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