Actual source code: ex3k.kokkos.cxx

  1: static char help[] = "Tests VecKokkosPlaceArray().\n\n";

  3: #include <petscvec.h>
  4: #include <Kokkos_Core.hpp>

  6: int main(int argc, char **argv)
  7: {
  8:   PetscInt  n = 10;
  9:   Vec       x, y;
 10:   PetscReal norm;

 12:   PetscFunctionBeginUser;
 13:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
 14:   {
 15:     // Create a VecKokkos x and init it
 16:     PetscCall(VecCreate(PETSC_COMM_WORLD, &x));
 17:     PetscCall(VecSetSizes(x, n, PETSC_DECIDE));
 18:     PetscCall(VecSetType(x, VECKOKKOS));
 19:     PetscCall(VecSet(x, 4.0));

 21:     // Allocate a Kokkos View kv and init it with a different value
 22:     auto kv = Kokkos::View<PetscScalar *>("kv", n);
 23:     PetscCallCXX(Kokkos::deep_copy(kv, 2.0));

 25:     // Use kv's array to replace the device array in x
 26:     PetscCall(VecKokkosPlaceArray(x, kv.data())); // x = {2.0, 2.0, ...}
 27:     PetscCall(VecScale(x, 0.5));                  // x = {1.0, 1.0, ...}
 28:     PetscCall(VecKokkosResetArray(x));            // x = {4.0, 4.0, ...}, kv = {1,0, 1.0, ...}

 30:     // Create a vector y with kv
 31:     PetscCall(VecCreateMPIKokkosWithArray(PETSC_COMM_WORLD, 1, n, PETSC_DECIDE, kv.data(), &y));

 33:     // Check both x and y have correct values
 34:     PetscCall(VecAXPY(x, -4.0, y)); // x -= 4 * y
 35:     PetscCall(VecNorm(x, NORM_2, &norm));
 36:     PetscCheck(norm < PETSC_SMALL, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Test failed with VecKokkosPlaceArray");

 38:     PetscCall(VecDestroy(&x));
 39:     PetscCall(VecDestroy(&y));
 40:   }
 41:   PetscCall(PetscFinalize());
 42:   return 0;
 43: }

 45: /*TEST

 47:     test:
 48:       requires: kokkos_kernels
 49:       nsize: {{1 2}}
 50:       output_file: output/empty.out

 52: TEST*/