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