Actual source code: petscvec_kokkos.hpp

  1: #if !defined(PETSCVEC_KOKKOS_HPP)
  2: #define PETSCVEC_KOKKOS_HPP

  4: #include <petscconf.h>

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

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

 14: #include <petscvec.h>

 16: #if defined(PETSC_HAVE_KOKKOS)
 17:   #include <Kokkos_Core.hpp>

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

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

 27:    Logically Collective on Vec

 29:    Input Parameter:
 30: .  v - the vector in type of VECKOKKOS

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

 35:    Notes:
 36:    If the vector is not of type VECKOKKOS, an error will be raised.
 37:    The functions are similar to VecGetArrayRead() and VecGetArray() respectively. One can read-only or read/write the returned Kokkos View.
 38:    Note that passing in a const View enables read-only access.
 39:    One must return the View by a matching VecRestoreKokkosView() after finishing using the View. Currently, only two memory
 40:    spaces are supported: Kokkos::HostSpace and Kokkos::DefaultExecutionSpace::memory_space.
 41:    If needed, a memory copy will be internally called to copy the latest vector data to the specified memory space.

 43:    Level: beginner

 45: .seealso: VecRestoreKokkosView(), VecRestoreArray(), VecGetKokkosViewWrite(), VecGetArrayRead(), VecGetArrays(), VecGetArrayF90(), VecGetArrayReadF90(), VecPlaceArray(), VecGetArray2d(),
 46:           VecGetArrayPair(), VecRestoreArrayPair(), VecGetArrayWrite(), VecRestoreArrayWrite()
 47: @*/
 48:   template<class MemorySpace> PetscErrorCode VecGetKokkosView (Vec,Kokkos::View<const PetscScalar*,MemorySpace>*);
 49:   template<class MemorySpace> PetscErrorCode VecGetKokkosView (Vec,Kokkos::View<PetscScalar*,MemorySpace>*);

 51: /*@C
 52:    VecRestoreKokkosView - Returns a Kokkos View gotten by VecGetKokkosView().

 54:    Synopsis:
 55: #include <petscvec_kokkos.hpp>
 56:    PetscErrorCode VecRestoreKokkosView  (Vec v,Kokkos::View<const PetscScalar*,MemorySpace>* kv);
 57:    PetscErrorCode VecRestoreKokkosView  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);

 59:    Logically Collective on Vec

 61:    Input Parameters:
 62: +  v  - the vector in type of VECKOKKOS
 63: -  kv - the Kokkos View with a user-specified template parameter MemorySpace

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

 69:    Level: beginner

 71: .seealso: VecGetKokkosView(), VecRestoreKokkosViewWrite(), VecRestoreArray(), VecGetArrayRead(), VecGetArrays(), VecGetArrayF90(), VecGetArrayReadF90(), VecPlaceArray(), VecGetArray2d(),
 72:           VecGetArrayPair(), VecRestoreArrayPair(), VecGetArrayWrite(), VecRestoreArrayWrite()
 73: @*/
 74:   template<class MemorySpace> PetscErrorCode VecRestoreKokkosView(Vec,Kokkos::View<const PetscScalar*,MemorySpace>*){return 0;}
 75:   template<class MemorySpace> PetscErrorCode VecRestoreKokkosView(Vec,Kokkos::View<PetscScalar*,MemorySpace>*);


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

 81:    Synopsis:
 82: #include <petscvec_kokkos.hpp>
 83:    PetscErrorCode VecGetKokkosViewWrite  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);

 85:    Logically Collective on Vec

 87:    Input Parameter:
 88: .  v - the vector in type of VECKOKKOS

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

 93:    Notes:
 94:    If the vector is not of type VECKOKKOS, an error will be raised.
 95:    The functions is similar to VecGetArrayWrite(). The returned view might contain garbage data or stale data and one is not
 96:    expected to read data from the View. Instead, one is expected to overwrite all data in the View.
 97:    One must return the View by a matching VecRestoreKokkosViewWrite() after finishing using the View.
 98:   Currently, only two memory spaces are supported: Kokkos::HostSpace and Kokkos::DefaultExecutionSpace::memory_space.

100:    Level: beginner

102: .seealso: VecRestoreKokkosViewWrite(), VecRestoreKokkosView(), VecGetKokkosView(), VecRestoreArray(), VecGetArrayRead(), VecGetArrays(), VecGetArrayF90(), VecGetArrayReadF90(), VecPlaceArray(), VecGetArray2d(),
103:           VecGetArrayPair(), VecRestoreArrayPair(), VecGetArrayWrite(), VecRestoreArrayWrite()
104: @*/
105:   template<class MemorySpace> PetscErrorCode VecGetKokkosViewWrite    (Vec,Kokkos::View<PetscScalar*,MemorySpace>*);

107: /*@C
108:    VecRestoreKokkosViewWrite - Returns a Kokkos View gotten by VecGetKokkosViewWrite().

110:    Synopsis:
111: #include <petscvec_kokkos.hpp>
112:    PetscErrorCode VecRestoreKokkosViewWrite  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);

114:    Logically Collective on Vec

116:    Input Parameters:
117: +  v  - the vector in type of VECKOKKOS
118: -  kv - the Kokkos View with a user-specified template parameter MemorySpace

120:    Notes:
121:    If the vector is not of type VECKOKKOS, an error will be raised.
122:    The function is similar to VecRestoreArrayWrite(). It is the counterpart of VecGetKokkosViewWrite().

124:    Level: beginner

126: .seealso: VecGetKokkosViewWrite(), VecGetKokkosView(), VecGetKokkosView(), VecRestoreArray(), VecGetArrayRead(), VecGetArrays(), VecGetArrayF90(), VecGetArrayReadF90(), VecPlaceArray(), VecGetArray2d(),
127:           VecGetArrayPair(), VecRestoreArrayPair(), VecGetArrayWrite(), VecRestoreArrayWrite()
128: @*/
129:   template<class MemorySpace> PetscErrorCode VecRestoreKokkosViewWrite(Vec,Kokkos::View<PetscScalar*,MemorySpace>*);

131:  #if defined(PETSC_HAVE_COMPLEX) && defined(PETSC_USE_COMPLEX)
132:   static_assert(std::alignment_of<Kokkos::complex<PetscReal>>::value == std::alignment_of<std::complex<PetscReal>>::value,
133:    "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");
134:  #endif

136: #endif

138: #endif