xref: /petsc/src/mat/tutorials/ex18k.kokkos.cxx (revision a336c15037c72f93cd561f5a5e11e93175f2efd9)
1 #include <Kokkos_Core.hpp>
2 #include <petscvec_kokkos.hpp>
3 #include "ex18.h"
4 
5 using DefaultMemorySpace = Kokkos::DefaultExecutionSpace::memory_space;
6 
7 PetscErrorCode FillMatrixKokkosCOO(FEStruct *fe, Mat A)
8 {
9   Kokkos::View<PetscScalar *, DefaultMemorySpace> v("v", 3 * 3 * fe->Ne);
10 
11   PetscFunctionBeginUser;
12   // Simulation of GPU based finite assembly process with COO
13   Kokkos::parallel_for(
14     "AssembleElementMatrices", fe->Ne, KOKKOS_LAMBDA(PetscInt i) {
15       PetscScalar *s = &v(3 * 3 * i);
16       for (PetscInt vi = 0; vi < 3; vi++) {
17         for (PetscInt vj = 0; vj < 3; vj++) s[vi * 3 + vj] = vi + 2 * vj;
18       }
19     });
20   PetscCall(MatSetValuesCOO(A, v.data(), INSERT_VALUES));
21   PetscFunctionReturn(PETSC_SUCCESS);
22 }
23