xref: /petsc/include/petscvec_kokkos.hpp (revision 030f984af8d8bb4c203755d35bded3c05b3d83ce)
1 #if !defined(PETSCVEC_KOKKOS_HPP)
2 #define PETSCVEC_KOKKOS_HPP
3 
4 #include <petscconf.h>
5 
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
10 
11   #define PETSC_DESIRE_KOKKOS_COMPLEX   1   /* To control the definition of petsccomplexlib in petscsystypes.h */
12 #endif
13 
14 #include <petscvec.h>
15 
16 #if defined(PETSC_HAVE_KOKKOS)
17   #include <Kokkos_Core.hpp>
18 
19   #if defined(PETSC_HAVE_CUDA)
20     #define WaitForKokkos() PetscCUDASynchronize ? (Kokkos::fence(),0) : 0
21   #elif defined(PETSC_HAVE_HIP)
22     #define WaitForKokkos() PetscHIPSynchronize ? (Kokkos::fence(),0) : 0
23   #else
24     #define WaitForKokkos() 0
25   #endif
26 
27   /* Routines to get/restore Kokkos Views from PETSc vectors */
28 
29   /* Like VecGetArrayRead() */
30   template<class MemorySpace> PetscErrorCode VecGetKokkosView    (Vec,Kokkos::View<const PetscScalar*,MemorySpace>*);
31   template<class MemorySpace> PetscErrorCode VecRestoreKokkosView(Vec,Kokkos::View<const PetscScalar*,MemorySpace>*){return 0;}
32 
33   /* Like VecGetArray() */
34   template<class MemorySpace> PetscErrorCode VecGetKokkosView    (Vec,Kokkos::View<PetscScalar*,MemorySpace>*);
35   template<class MemorySpace> PetscErrorCode VecRestoreKokkosView(Vec,Kokkos::View<PetscScalar*,MemorySpace>*);
36 
37   /* Like VecGetArrayWrite() */
38   template<class MemorySpace> PetscErrorCode VecGetKokkosViewWrite    (Vec,Kokkos::View<PetscScalar*,MemorySpace>*);
39   template<class MemorySpace> PetscErrorCode VecRestoreKokkosViewWrite(Vec,Kokkos::View<PetscScalar*,MemorySpace>*);
40 
41 #endif
42 
43 #endif
44