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