Date: 2025-03-29T05:58:39+00:00

Actual source code: ex18kok.kokkos.cxx

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