xref: /petsc/include/petscvec_kokkos.hpp (revision cb5db2414029547a5ccf00a1710ee072432a08af)
1 #pragma once
2 
3 #include <petscconf.h>
4 
5 /* SUBMANSEC = Vec */
6 
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
11 
12   #define PETSC_DESIRE_KOKKOS_COMPLEX 1 /* To control the definition of petsccomplexlib in petscsystypes.h */
13 #endif
14 
15 #include <petscvec.h>
16 
17 #if defined(PETSC_HAVE_KOKKOS)
18   #include <Kokkos_Core.hpp>
19 
20 /*@C
21      VecGetKokkosView - Returns a constant Kokkos View that contains up-to-date data of a vector in the specified memory space.
22 
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);
27 
28    Logically Collective
29 
30    Input Parameter:
31 .  v - the vector in type of `VECKOKKOS`
32 
33    Output Parameter:
34 .  kv - the Kokkos View with a user-specified template parameter MemorySpace
35 
36    Notes:
37    If the vector is not of type `VECKOKKOS`, an error will be raised.
38 
39    The functions are similar to `VecGetArrayRead()` and `VecGetArray()` respectively. One can read-only or read/write the returned Kokkos View.
40 
41    Passing in a const View enables read-only access.
42 
43    One must return the View by a matching `VecRestoreKokkosView()` after finishing using the View. Currently, only two memory
44    spaces are supported: Kokkos::HostSpace and Kokkos::DefaultExecutionSpace::memory_space.
45    If needed, a memory copy will be internally called to copy the latest vector data to the specified memory space.
46 
47    Level: beginner
48 
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> *);
56 
57 /*@C
58    VecRestoreKokkosView - Returns a Kokkos View gotten by `VecGetKokkosView()`.
59 
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);
64 
65    Logically Collective
66 
67    Input Parameters:
68 +  v  - the vector in type of `VECKOKKOS`
69 -  kv - the Kokkos View with a user-specified template parameter MemorySpace
70 
71    Notes:
72    If the vector is not of type `VECKOKKOS`, an error will be raised.
73    The functions are similar to `VecRestoreArrayRead()` and `VecRestoreArray()` respectively. They are the counterpart of `VecGetKokkosView()`.
74 
75    Level: beginner
76 
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> *);
87 
88 /*@C
89    VecGetKokkosViewWrite - Returns a Kokkos View that contains the array of a vector in the specified memory space.
90 
91    Synopsis:
92    #include <petscvec_kokkos.hpp>
93    PetscErrorCode VecGetKokkosViewWrite  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);
94 
95    Logically Collective
96 
97    Input Parameter:
98 .  v - the vector in type of `VECKOKKOS`
99 
100    Output Parameter:
101 .  kv - the Kokkos View with a user-specified template parameter MemorySpace
102 
103    Notes:
104    If the vector is not of type `VECKOKKOS`, an error will be raised.
105 
106    The functions is similar to `VecGetArrayWrite()`. The returned view might contain garbage data or stale data and one is not
107    expected to read data from the View. Instead, one is expected to overwrite all data in the View.
108    One must return the View by a matching `VecRestoreKokkosViewWrite()` after finishing using the View.
109 
110    Currently, only two memory spaces are supported: Kokkos::HostSpace and Kokkos::DefaultExecutionSpace::memory_space.
111 
112    Level: beginner
113 
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> *);
119 
120 /*@C
121    VecRestoreKokkosViewWrite - Returns a Kokkos View gotten with `VecGetKokkosViewWrite()`.
122 
123    Synopsis:
124    #include <petscvec_kokkos.hpp>
125    PetscErrorCode VecRestoreKokkosViewWrite  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);
126 
127    Logically Collective
128 
129    Input Parameters:
130 +  v  - the vector in type of `VECKOKKOS`
131 -  kv - the Kokkos View with a user-specified template parameter MemorySpace
132 
133    Notes:
134    If the vector is not of type `VECKOKKOS`, an error will be raised.
135 
136    The function is similar to `VecRestoreArrayWrite()`. It is the counterpart of `VecGetKokkosViewWrite()`.
137 
138    Level: beginner
139 
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> *);
145 
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
150 
151 #endif
152