1: #pragma once

  3: #include <petscconf.h>

  5: /* SUBMANSEC = Vec */

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

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

 15: #include <petscvec.h>

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

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

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

 28:    Logically Collective

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

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

 36:    Level: beginner

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

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

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

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

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

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

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

 65:    Logically Collective

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

 71:    Level: beginner

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

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

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

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

 95:    Logically Collective

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

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

103:    Level: beginner

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

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

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

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

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

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

127:    Logically Collective

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

133:    Level: beginner

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

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

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

146:   #if defined(PETSC_HAVE_COMPLEX) && defined(PETSC_USE_COMPLEX)
147: static_assert(std::alignment_of<Kokkos::complex<PetscReal>>::value == std::alignment_of<std::complex<PetscReal>>::value,
148:               "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");
149:   #endif

151: #endif