1 #include <Kokkos_Core.hpp> 2 #include <petscvec_kokkos.hpp> 3 #include "ex18.h" 4 5 using DefaultMemorySpace = Kokkos::DefaultExecutionSpace::memory_space; 6 FillMatrixKokkosCOO(FEStruct * fe,Mat A)7PetscErrorCode 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