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