xref: /petsc/include/petscvec_kokkos.hpp (revision 6ffe77eaecce1557e50d00ca5347a7f48e598865)
1 #if !defined(PETSCVEC_KOKKOS_HPP)
2 #define PETSCVEC_KOKKOS_HPP
3 
4 #include <petscconf.h>
5 
6 /* SUBMANSEC = Vec */
7 
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
12 
13   #define PETSC_DESIRE_KOKKOS_COMPLEX   1   /* To control the definition of petsccomplexlib in petscsystypes.h */
14 #endif
15 
16 #include <petscvec.h>
17 
18 #if defined(PETSC_HAVE_KOKKOS)
19   #include <Kokkos_Core.hpp>
20 
21 /*@C
22      VecGetKokkosView - Returns a constant Kokkos View that contains up-to-date data of a vector in the specified memory space.
23 
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);
28 
29    Logically Collective on Vec
30 
31    Input Parameter:
32 .  v - the vector in type of VECKOKKOS
33 
34    Output Parameter:
35 .  kv - the Kokkos View with a user-specified template parameter MemorySpace
36 
37    Notes:
38    If the vector is not of type VECKOKKOS, an error will be raised.
39    The functions are similar to VecGetArrayRead() and VecGetArray() respectively. One can read-only or read/write the returned Kokkos View.
40    Note that passing in a const View enables read-only access.
41    One must return the View by a matching VecRestoreKokkosView() after finishing using the View. Currently, only two memory
42    spaces are supported: Kokkos::HostSpace and Kokkos::DefaultExecutionSpace::memory_space.
43    If needed, a memory copy will be internally called to copy the latest vector data to the specified memory space.
44 
45    Level: beginner
46 
47 .seealso: `VecRestoreKokkosView()`, `VecRestoreArray()`, `VecGetKokkosViewWrite()`, `VecGetArrayRead()`, `VecGetArrays()`, `VecGetArrayF90()`, `VecGetArrayReadF90()`, `VecPlaceArray()`, `VecGetArray2d()`,
48           `VecGetArrayPair()`, `VecRestoreArrayPair()`, `VecGetArrayWrite()`, `VecRestoreArrayWrite()`
49 @*/
50   template<class MemorySpace> PetscErrorCode VecGetKokkosView (Vec,Kokkos::View<const PetscScalar*,MemorySpace>*);
51   template<class MemorySpace> PetscErrorCode VecGetKokkosView (Vec,Kokkos::View<PetscScalar*,MemorySpace>*);
52 
53 /*@C
54    VecRestoreKokkosView - Returns a Kokkos View gotten by VecGetKokkosView().
55 
56    Synopsis:
57    #include <petscvec_kokkos.hpp>
58    PetscErrorCode VecRestoreKokkosView  (Vec v,Kokkos::View<const PetscScalar*,MemorySpace>* kv);
59    PetscErrorCode VecRestoreKokkosView  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);
60 
61    Logically Collective on Vec
62 
63    Input Parameters:
64 +  v  - the vector in type of VECKOKKOS
65 -  kv - the Kokkos View with a user-specified template parameter MemorySpace
66 
67    Notes:
68    If the vector is not of type VECKOKKOS, an error will be raised.
69    The functions are similar to VecRestoreArrayRead() and VecRestoreArray() respectively. They are the counterpart of VecGetKokkosView().
70 
71    Level: beginner
72 
73 .seealso: `VecGetKokkosView()`, `VecRestoreKokkosViewWrite()`, `VecRestoreArray()`, `VecGetArrayRead()`, `VecGetArrays()`, `VecGetArrayF90()`, `VecGetArrayReadF90()`, `VecPlaceArray()`, `VecGetArray2d()`,
74           `VecGetArrayPair()`, `VecRestoreArrayPair()`, `VecGetArrayWrite()`, `VecRestoreArrayWrite()`
75 @*/
76   template<class MemorySpace> PetscErrorCode VecRestoreKokkosView(Vec,Kokkos::View<const PetscScalar*,MemorySpace>*){return 0;}
77   template<class MemorySpace> PetscErrorCode VecRestoreKokkosView(Vec,Kokkos::View<PetscScalar*,MemorySpace>*);
78 
79 
80 /*@C
81    VecGetKokkosViewWrite - Returns a Kokkos View that contains the array of a vector in the specified memory space.
82 
83    Synopsis:
84    #include <petscvec_kokkos.hpp>
85    PetscErrorCode VecGetKokkosViewWrite  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);
86 
87    Logically Collective on Vec
88 
89    Input Parameter:
90 .  v - the vector in type of VECKOKKOS
91 
92    Output Parameter:
93 .  kv - the Kokkos View with a user-specified template parameter MemorySpace
94 
95    Notes:
96    If the vector is not of type VECKOKKOS, an error will be raised.
97    The functions is similar to VecGetArrayWrite(). The returned view might contain garbage data or stale data and one is not
98    expected to read data from the View. Instead, one is expected to overwrite all data in the View.
99    One must return the View by a matching VecRestoreKokkosViewWrite() after finishing using the View.
100   Currently, only two memory spaces are supported: Kokkos::HostSpace and Kokkos::DefaultExecutionSpace::memory_space.
101 
102    Level: beginner
103 
104 .seealso: `VecRestoreKokkosViewWrite()`, `VecRestoreKokkosView()`, `VecGetKokkosView()`, `VecRestoreArray()`, `VecGetArrayRead()`, `VecGetArrays()`, `VecGetArrayF90()`, `VecGetArrayReadF90()`, `VecPlaceArray()`, `VecGetArray2d()`,
105           `VecGetArrayPair()`, `VecRestoreArrayPair()`, `VecGetArrayWrite()`, `VecRestoreArrayWrite()`
106 @*/
107   template<class MemorySpace> PetscErrorCode VecGetKokkosViewWrite    (Vec,Kokkos::View<PetscScalar*,MemorySpace>*);
108 
109 /*@C
110    VecRestoreKokkosViewWrite - Returns a Kokkos View gotten by VecGetKokkosViewWrite().
111 
112    Synopsis:
113    #include <petscvec_kokkos.hpp>
114    PetscErrorCode VecRestoreKokkosViewWrite  (Vec v,Kokkos::View<PetscScalar*,MemorySpace>* kv);
115 
116    Logically Collective on Vec
117 
118    Input Parameters:
119 +  v  - the vector in type of VECKOKKOS
120 -  kv - the Kokkos View with a user-specified template parameter MemorySpace
121 
122    Notes:
123    If the vector is not of type VECKOKKOS, an error will be raised.
124    The function is similar to VecRestoreArrayWrite(). It is the counterpart of VecGetKokkosViewWrite().
125 
126    Level: beginner
127 
128 .seealso: `VecGetKokkosViewWrite()`, `VecGetKokkosView()`, `VecGetKokkosView()`, `VecRestoreArray()`, `VecGetArrayRead()`, `VecGetArrays()`, `VecGetArrayF90()`, `VecGetArrayReadF90()`, `VecPlaceArray()`, `VecGetArray2d()`,
129           `VecGetArrayPair()`, `VecRestoreArrayPair()`, `VecGetArrayWrite()`, `VecRestoreArrayWrite()`
130 @*/
131   template<class MemorySpace> PetscErrorCode VecRestoreKokkosViewWrite(Vec,Kokkos::View<PetscScalar*,MemorySpace>*);
132 
133  #if defined(PETSC_HAVE_COMPLEX) && defined(PETSC_USE_COMPLEX)
134   static_assert(std::alignment_of<Kokkos::complex<PetscReal>>::value == std::alignment_of<std::complex<PetscReal>>::value,
135    "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");
136  #endif
137 
138 #endif
139 
140 #endif
141