1: #include <petsc_kokkos.hpp>
  2: #include <petsc/private/petschypre.h>
  3: #include <Kokkos_Core.hpp>
  4: #include <../src/mat/impls/hypre/mhypre.h>
  5: #include <petsc/private/kokkosimpl.hpp>

  7: PetscErrorCode MatZeroRows_Kokkos(PetscInt n, const PetscInt rows[], const HYPRE_Int i[], const HYPRE_Int j[], HYPRE_Complex a[], HYPRE_Complex diag)
  8: {
  9:   PetscFunctionBegin;
 10:   if (!n) PetscFunctionReturn(PETSC_SUCCESS);
 11:   PetscCall(PetscKokkosInitializeCheck()); // As we might have not created any petsc/kokkos object yet
 12:   Kokkos::parallel_for(
 13:     Kokkos::TeamPolicy<>(n, Kokkos::AUTO()), KOKKOS_LAMBDA(const Kokkos::TeamPolicy<>::member_type &t) {
 14:       PetscInt r = rows[t.league_rank()]; // row r
 15:       Kokkos::parallel_for(Kokkos::TeamThreadRange(t, i[r + 1] - i[r]), [&](PetscInt c) {
 16:         if (r == j[i[r] + c]) a[i[r] + c] = diag;
 17:         else a[i[r] + c] = 0.0;
 18:       });
 19:     });
 20:   PetscFunctionReturn(PETSC_SUCCESS);
 21: }

 23: PetscErrorCode PetscHypreIntCastArray_Kokkos(PetscInt n, const PetscInt *a, HYPRE_Int *b)
 24: {
 25:   PetscFunctionBegin;
 26:   if (n) PetscCallCXX(Kokkos::parallel_for(Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, n), KOKKOS_LAMBDA(const size_t i) { b[i] = static_cast<HYPRE_Int>(a[i]); }));
 27:   PetscFunctionReturn(PETSC_SUCCESS);
 28: }

 30: PetscErrorCode MatHypreDeviceMalloc_Kokkos(size_t size, void **ptr)
 31: {
 32:   PetscFunctionBegin;
 33:   if (size) PetscCallCXX(*ptr = Kokkos::kokkos_malloc<DefaultMemorySpace>(size));
 34:   else *ptr = NULL;
 35:   PetscFunctionReturn(PETSC_SUCCESS);
 36: }

 38: PetscErrorCode MatHypreDeviceFree_Kokkos(void *a)
 39: {
 40:   PetscFunctionBegin;
 41:   PetscCallCXX(Kokkos::kokkos_free<DefaultMemorySpace>(a));
 42:   PetscFunctionReturn(PETSC_SUCCESS);
 43: }