Actual source code: petscvec_kokkos.hpp

  1: #ifndef PETSCVEC_KOKKOS_HPP
  2: #define PETSCVEC_KOKKOS_HPP

  4: #include <petscconf.h>

  6: /* SUBMANSEC = Vec */

  8: #if defined(PETSC_HAVE_KOKKOS)
  9:   #if defined(petsccomplexlib)
 10:     #error "Error: You must include petscvec_kokkos.hpp before other petsc headers in this C++ file to use petsc complex with Kokkos"
 11:   #endif

 13:   #define PETSC_DESIRE_KOKKOS_COMPLEX 1 /* To control the definition of petsccomplexlib in petscsystypes.h */
 14: #endif

 16: #include <petscvec.h>

 18: #if defined(PETSC_HAVE_KOKKOS)
 19:   #include <Kokkos_Core.hpp>

 21: /*@C
 22:      VecGetKokkosView - Returns a constant Kokkos View that contains up-to-date data of a vector in the specified memory space.

 24:    Synopsis:
 25: #include <petscvec_kokkos.hpp>
 26:    PetscErrorCode VecGetKokkosView  (Vec v,Kokkos::View<const PetscScalar*,MemorySpace>* kv);
 27:    PetscErrorCode VecGetKokkosView  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);

 29:    Logically Collective

 31:    Input Parameter:
 32: .  v - the vector in type of `VECKOKKOS`

 34:    Output Parameter:
 35: .  kv - the Kokkos View with a user-specified template parameter MemorySpace

 37:    Notes:
 38:    If the vector is not of type `VECKOKKOS`, an error will be raised.

 40:    The functions are similar to `VecGetArrayRead()` and `VecGetArray()` respectively. One can read-only or read/write the returned Kokkos View.

 42:    Passing in a const View enables read-only access.

 44:    One must return the View by a matching `VecRestoreKokkosView()` after finishing using the View. Currently, only two memory
 45:    spaces are supported: Kokkos::HostSpace and Kokkos::DefaultExecutionSpace::memory_space.
 46:    If needed, a memory copy will be internally called to copy the latest vector data to the specified memory space.

 48:    Level: beginner

 50: .seealso: `VecRestoreKokkosView()`, `VecRestoreArray()`, `VecGetKokkosViewWrite()`, `VecGetArrayRead()`, `VecGetArrays()`, `VecGetArrayF90()`, `VecGetArrayReadF90()`, `VecPlaceArray()`, `VecGetArray2d()`,
 51:           `VecGetArrayPair()`, `VecRestoreArrayPair()`, `VecGetArrayWrite()`, `VecRestoreArrayWrite()`
 52: @*/
 53: template <class MemorySpace>
 54: PetscErrorCode VecGetKokkosView(Vec, Kokkos::View<const PetscScalar *, MemorySpace> *);
 55: template <class MemorySpace>
 56: PetscErrorCode VecGetKokkosView(Vec, Kokkos::View<PetscScalar *, MemorySpace> *);

 58: /*@C
 59:    VecRestoreKokkosView - Returns a Kokkos View gotten by `VecGetKokkosView()`.

 61:    Synopsis:
 62: #include <petscvec_kokkos.hpp>
 63:    PetscErrorCode VecRestoreKokkosView  (Vec v,Kokkos::View<const PetscScalar*,MemorySpace>* kv);
 64:    PetscErrorCode VecRestoreKokkosView  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);

 66:    Logically Collective

 68:    Input Parameters:
 69: +  v  - the vector in type of `VECKOKKOS`
 70: -  kv - the Kokkos View with a user-specified template parameter MemorySpace

 72:    Notes:
 73:    If the vector is not of type `VECKOKKOS`, an error will be raised.
 74:    The functions are similar to `VecRestoreArrayRead()` and `VecRestoreArray()` respectively. They are the counterpart of `VecGetKokkosView()`.

 76:    Level: beginner

 78: .seealso: `VecGetKokkosView()`, `VecRestoreKokkosViewWrite()`, `VecRestoreArray()`, `VecGetArrayRead()`, `VecGetArrays()`, `VecGetArrayF90()`, `VecGetArrayReadF90()`, `VecPlaceArray()`, `VecGetArray2d()`,
 79:           `VecGetArrayPair()`, `VecRestoreArrayPair()`, `VecGetArrayWrite()`, `VecRestoreArrayWrite()`
 80: @*/
 81: template <class MemorySpace>
 82: PetscErrorCode VecRestoreKokkosView(Vec, Kokkos::View<const PetscScalar *, MemorySpace> *)
 83: {
 84:   return 0;
 85: }
 86: template <class MemorySpace>
 87: PetscErrorCode VecRestoreKokkosView(Vec, Kokkos::View<PetscScalar *, MemorySpace> *);

 89: /*@C
 90:    VecGetKokkosViewWrite - Returns a Kokkos View that contains the array of a vector in the specified memory space.

 92:    Synopsis:
 93: #include <petscvec_kokkos.hpp>
 94:    PetscErrorCode VecGetKokkosViewWrite  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);

 96:    Logically Collective

 98:    Input Parameter:
 99: .  v - the vector in type of `VECKOKKOS`

101:    Output Parameter:
102: .  kv - the Kokkos View with a user-specified template parameter MemorySpace

104:    Notes:
105:    If the vector is not of type `VECKOKKOS`, an error will be raised.

107:    The functions is similar to `VecGetArrayWrite()`. The returned view might contain garbage data or stale data and one is not
108:    expected to read data from the View. Instead, one is expected to overwrite all data in the View.
109:    One must return the View by a matching `VecRestoreKokkosViewWrite()` after finishing using the View.

111:    Currently, only two memory spaces are supported: Kokkos::HostSpace and Kokkos::DefaultExecutionSpace::memory_space.

113:    Level: beginner

115: .seealso: `VecRestoreKokkosViewWrite()`, `VecRestoreKokkosView()`, `VecGetKokkosView()`, `VecRestoreArray()`, `VecGetArrayRead()`, `VecGetArrays()`, `VecGetArrayF90()`, `VecGetArrayReadF90()`, `VecPlaceArray()`, `VecGetArray2d()`,
116:           `VecGetArrayPair()`, `VecRestoreArrayPair()`, `VecGetArrayWrite()`, `VecRestoreArrayWrite()`
117: @*/
118: template <class MemorySpace>
119: PetscErrorCode VecGetKokkosViewWrite(Vec, Kokkos::View<PetscScalar *, MemorySpace> *);

121: /*@C
122:    VecRestoreKokkosViewWrite - Returns a Kokkos View gotten with `VecGetKokkosViewWrite()`.

124:    Synopsis:
125: #include <petscvec_kokkos.hpp>
126:    PetscErrorCode VecRestoreKokkosViewWrite  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);

128:    Logically Collective

130:    Input Parameters:
131: +  v  - the vector in type of `VECKOKKOS`
132: -  kv - the Kokkos View with a user-specified template parameter MemorySpace

134:    Notes:
135:    If the vector is not of type `VECKOKKOS`, an error will be raised.

137:    The function is similar to `VecRestoreArrayWrite()`. It is the counterpart of `VecGetKokkosViewWrite()`.

139:    Level: beginner

141: .seealso: `VecGetKokkosViewWrite()`, `VecGetKokkosView()`, `VecGetKokkosView()`, `VecRestoreArray()`, `VecGetArrayRead()`, `VecGetArrays()`, `VecGetArrayF90()`, `VecGetArrayReadF90()`, `VecPlaceArray()`, `VecGetArray2d()`,
142:           `VecGetArrayPair()`, `VecRestoreArrayPair()`, `VecGetArrayWrite()`, `VecRestoreArrayWrite()`
143: @*/
144: template <class MemorySpace>
145: PetscErrorCode VecRestoreKokkosViewWrite(Vec, Kokkos::View<PetscScalar *, MemorySpace> *);

147:   #if defined(PETSC_HAVE_COMPLEX) && defined(PETSC_USE_COMPLEX)
148: static_assert(std::alignment_of<Kokkos::complex<PetscReal>>::value == std::alignment_of<std::complex<PetscReal>>::value,
149:               "Alignment of Kokkos::complex<PetscReal> and std::complex<PetscReal> mismatch. Reconfigure your Kokkos with -DKOKKOS_ENABLE_COMPLEX_ALIGN=OFF, or let PETSc install Kokkos for you with --download-kokkos --download-kokkos-kernels");
150:   #endif

152: #endif

154: #endif