xref: /petsc/include/petscvec_kokkos.hpp (revision ee12ae39415b2e672d944cdca066227dadbf8b14)
1 #if !defined(PETSCVEC_KOKKOS_HPP)
2 #define PETSCVEC_KOKKOS_HPP
3 
4 #include <petscvec.h>
5 
6 #if defined(PETSC_HAVE_KOKKOS)
7   #include <Kokkos_Core.hpp>
8 
9   #if defined(PETSC_HAVE_CUDA)
10     #define WaitForKokkos() PetscCUDASynchronize ? (Kokkos::fence(),0) : 0
11   #elif defined(PETSC_HAVE_HIP)
12     #define WaitForKokkos() PetscHIPSynchronize ? (Kokkos::fence(),0) : 0
13   #else
14     #define WaitForKokkos() 0
15   #endif
16 
17   /* Routines to get/restore Kokkos Views from PETSc vectors */
18 
19   /* Like VecGetArrayRead() */
20   template<class MemorySpace> PetscErrorCode VecGetKokkosView    (Vec,Kokkos::View<const PetscScalar*,MemorySpace>*);
21   template<class MemorySpace> PetscErrorCode VecRestoreKokkosView(Vec,Kokkos::View<const PetscScalar*,MemorySpace>*){return 0;}
22 
23   /* Like VecGetArray() */
24   template<class MemorySpace> PetscErrorCode VecGetKokkosView    (Vec,Kokkos::View<PetscScalar*,MemorySpace>*);
25   template<class MemorySpace> PetscErrorCode VecRestoreKokkosView(Vec,Kokkos::View<PetscScalar*,MemorySpace>*);
26 
27   /* Like VecGetArrayWrite() */
28   template<class MemorySpace> PetscErrorCode VecGetKokkosViewWrite    (Vec,Kokkos::View<PetscScalar*,MemorySpace>*);
29   template<class MemorySpace> PetscErrorCode VecRestoreKokkosViewWrite(Vec,Kokkos::View<PetscScalar*,MemorySpace>*);
30 
31 #endif
32 
33 #endif
34