Actual source code: petscmat_kokkos.hpp

  1: #pragma once

  3: #include <petscconf.h>

  5: /* SUBMANSEC = Mat */

  7: #if defined(PETSC_HAVE_KOKKOS)

  9:   #include <Kokkos_Core.hpp>
 10: #include <petscmat.h>

 12: /*@C
 13:    MatCreateSeqAIJKokkosWithKokkosViews - Creates a MATSEQAIJKOKKOS matrix with Kokkos views of the aij data

 15:    Synopsis:
 16: #include <petscmat_kokkos.hpp>
 17:    PetscErrorCode MatCreateSeqAIJKokkosWithKokkosViews  (MPI_Comm comm, PetscInt m, PetscInt n, Kokkos::View<PetscInt *>& i_d, Kokkos::View<PetscInt *>& j_d, Kokkos::View<PetscScalar *>& a_d, Mat *A);

 19:    Logically Collective, No Fortran Support

 21:    Input Parameter:
 22: +  comm  - the MPI communicator
 23: -  m     - row size
 24: -  n     - the column size
 25: -  i     - the Kokkos view of row data (in Kokkos::DefaultExecutionSpace)
 26: -  j     - the Kokkos view of the column data (in Kokkos::DefaultExecutionSpace)
 27: -  a     - the Kokkos view of the values (in Kokkos::DefaultExecutionSpace)

 29:    Output Parameter:
 30: .  A  - the `MATSEQAIJKOKKOS` matrix

 32:    Level: intermediate

 34:    Notes:
 35:    Creates a Mat given the csr data input as Kokkos views. This routine allows a Mat
 36:    to be built without involving the host. Don't modify entries in the views after this routine.
 37:    There should be no outstanding asynchronous operations on the views (ie this routine does not call fence()
 38:    before using the views)

 40: .seealso:
 41: @*/
 42: PetscErrorCode MatCreateSeqAIJKokkosWithKokkosViews(MPI_Comm, PetscInt, PetscInt, Kokkos::View<PetscInt *> &, Kokkos::View<PetscInt *> &, Kokkos::View<PetscScalar *> &, Mat *);

 44: /*@C
 45:    MatSeqAIJGetKokkosView - Returns a Kokkos View of the matrix nonzero values on the device, with up-to-date data.

 47:    Not Collective, No Fortran Support

 49:    Input Parameter:
 50: .  A - the matrix in type of `MATSEQAIJKOKKOS`

 52:    Output Parameter:
 53: .  kv - the Kokkos View over the matrix nonzero values on the device

 55:    Level: intermediate

 57:    Notes:
 58:    If the matrix is not of type `MATSEQAIJKOKKOS`, an error will be raised.

 60:    Passing in a const View enables read-only access; passing in a non-const View enables read-write access. In the
 61:    read-write case, the matching `MatSeqAIJRestoreKokkosView()` marks the device side as modified so subsequent
 62:    operations see the new values.

 64:    One must return the View with a matching `MatSeqAIJRestoreKokkosView()` after finishing using the View.

 66: .seealso: `MatSeqAIJRestoreKokkosView()`, `MatSeqAIJGetKokkosViewWrite()`, `MatSeqAIJGetCSRAndMemType()`, `MatCreateSeqAIJKokkosWithKokkosViews()`, `VecGetKokkosView()`
 67: @*/
 68: PetscErrorCode MatSeqAIJGetKokkosView(Mat, Kokkos::View<const PetscScalar *> *);
 69: PetscErrorCode MatSeqAIJGetKokkosView(Mat, Kokkos::View<PetscScalar *> *);

 71: /*@C
 72:    MatSeqAIJRestoreKokkosView - Returns a Kokkos View obtained with `MatSeqAIJGetKokkosView()`.

 74:    Not Collective, No Fortran Support

 76:    Input Parameters:
 77: +  A  - the matrix in type of `MATSEQAIJKOKKOS`
 78: -  kv - the Kokkos View previously obtained with `MatSeqAIJGetKokkosView()`

 80:    Level: intermediate

 82:    Notes:
 83:    If the matrix is not of type `MATSEQAIJKOKKOS`, an error will be raised.

 85:    When the matching `MatSeqAIJGetKokkosView()` returned a non-const View, the device side of the matrix is marked
 86:    modified so subsequent PETSc operations see the new values. The const overload is a no-op.

 88: .seealso: `MatSeqAIJGetKokkosView()`, `MatSeqAIJRestoreKokkosViewWrite()`, `MatSeqAIJGetKokkosViewWrite()`
 89: @*/
 90: PetscErrorCode MatSeqAIJRestoreKokkosView(Mat, Kokkos::View<const PetscScalar *> *);
 91: PetscErrorCode MatSeqAIJRestoreKokkosView(Mat, Kokkos::View<PetscScalar *> *);

 93: /*@C
 94:    MatSeqAIJGetKokkosViewWrite - Returns a Kokkos View of the matrix nonzero values on the device for write-only access.

 96:    Not Collective, No Fortran Support

 98:    Input Parameter:
 99: .  A - the matrix in type of `MATSEQAIJKOKKOS`

101:    Output Parameter:
102: .  kv - the Kokkos View over the matrix nonzero values on the device

104:    Level: intermediate

106:    Notes:
107:    If the matrix is not of type `MATSEQAIJKOKKOS`, an error will be raised.

109:    This routine does not synchronize the device with the host first. The caller is expected to overwrite all entries
110:    in the View; reading from it may return stale or garbage data.

112:    One must return the View with a matching `MatSeqAIJRestoreKokkosViewWrite()` after finishing using the View.

114: .seealso: `MatSeqAIJRestoreKokkosViewWrite()`, `MatSeqAIJGetKokkosView()`, `MatSeqAIJRestoreKokkosView()`
115: @*/
116: PetscErrorCode MatSeqAIJGetKokkosViewWrite(Mat, Kokkos::View<PetscScalar *> *);

118: /*@C
119:    MatSeqAIJRestoreKokkosViewWrite - Returns a Kokkos View obtained with `MatSeqAIJGetKokkosViewWrite()`.

121:    Not Collective, No Fortran Support

123:    Input Parameters:
124: +  A  - the matrix in type of `MATSEQAIJKOKKOS`
125: -  kv - the Kokkos View previously obtained with `MatSeqAIJGetKokkosViewWrite()`

127:    Level: intermediate

129:    Notes:
130:    If the matrix is not of type `MATSEQAIJKOKKOS`, an error will be raised.

132:    The device side of the matrix is marked modified so subsequent PETSc operations see the new values.

134: .seealso: `MatSeqAIJGetKokkosViewWrite()`, `MatSeqAIJGetKokkosView()`, `MatSeqAIJRestoreKokkosView()`
135: @*/
136: PetscErrorCode MatSeqAIJRestoreKokkosViewWrite(Mat, Kokkos::View<PetscScalar *> *);

138: #endif