1: #include <Kokkos_Core.hpp>
  2: #include <petscvec_kokkos.hpp>
  3: #include "ex18.h"

  5: using DefaultMemorySpace = Kokkos::DefaultExecutionSpace::memory_space;

  7: PetscErrorCode FillMatrixKokkosCOO(FEStruct *fe, Mat A)
  8: {
  9:   Kokkos::View<PetscScalar *, DefaultMemorySpace> v("v", 3 * 3 * fe->Ne);

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