Actual source code: hypre3.kokkos.cxx

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

      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: }