1: static char help[] = "Tests VecKokkosPlaceArray() and array-less vectors from VecCreate{Seq,MPI}KokkosWithArray().\n\n";

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

  6: int main(int argc, char **argv)
  7: {
  8:   PetscInt     n = 10;
  9:   PetscMPIInt  size;
 10:   Vec          x, y, z;
 11:   PetscReal    norm;
 12:   PetscScalar *array;

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

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

 28:     // Use kv's array to replace the device array in x
 29:     PetscCall(VecKokkosPlaceArray(x, kv.data())); // x = {2.0, 2.0, ...}
 30:     PetscCall(VecScale(x, 0.5));                  // x = {1.0, 1.0, ...}
 31:     PetscCall(VecGetArray(x, &array));            // must see the placed array, not the original one
 32:     PetscCheck(array[0] == 1.0, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "VecGetArray() did not return the array placed with VecKokkosPlaceArray()");
 33:     PetscCall(VecRestoreArray(x, &array));
 34:     PetscCall(VecKokkosResetArray(x)); // x = {4.0, 4.0, ...}, kv = {1,0, 1.0, ...}

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

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

 44:     PetscCall(VecDestroy(&x));
 45:     PetscCall(VecDestroy(&y));
 46:   }
 47:   {
 48:     // Create an array-less VecKokkos z: no host or device memory is allocated until an array is placed
 49:     if (size == 1) PetscCall(VecCreateSeqKokkosWithArray(PETSC_COMM_SELF, 1, n, NULL, &z));
 50:     else PetscCall(VecCreateMPIKokkosWithArray(PETSC_COMM_WORLD, 1, n, PETSC_DECIDE, NULL, &z));

 52:     // Place a device array, use z on device and on host, and take the array back
 53:     auto kz = Kokkos::View<PetscScalar *>("kz", n);
 54:     PetscCallCXX(Kokkos::deep_copy(kz, 2.0));
 55:     PetscCall(VecKokkosPlaceArray(z, kz.data())); // z = {2.0, 2.0, ...}
 56:     PetscCall(VecScale(z, 0.5));                  // z = {1.0, 1.0, ...}
 57:     PetscCall(VecNorm(z, NORM_1, &norm));
 58:     PetscCheck(PetscAbsReal(norm - (PetscReal)(n * size)) < PETSC_SMALL, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Test failed with VecKokkosPlaceArray() on an array-less vector");
 59:     PetscCall(VecGetArray(z, &array));
 60:     array[0] = 3.0; // z = {3.0, 1.0, ...}
 61:     PetscCall(VecRestoreArray(z, &array));
 62:     PetscCall(VecKokkosResetArray(z)); // kz = {3.0, 1.0, ...}
 63:     PetscCallCXX(Kokkos::fence());
 64:     {
 65:       auto kz_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), kz);

 67:       PetscCheck(kz_h(0) == 3.0 && kz_h(n - 1) == 1.0, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Test failed with VecKokkosResetArray() on an array-less vector");
 68:     }

 70:     // Place a host array, as MatDenseGetColumnVec() does, use z, and take the array back
 71:     PetscCall(PetscMalloc1(n, &array));
 72:     for (PetscInt i = 0; i < n; i++) array[i] = 1.0;
 73:     PetscCall(VecPlaceArray(z, array));
 74:     PetscCall(VecScale(z, 4.0)); // z = {4.0, 4.0, ...}
 75:     PetscCall(VecNorm(z, NORM_INFINITY, &norm));
 76:     PetscCheck(PetscAbsReal(norm - 4.0) < PETSC_SMALL, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Test failed with VecPlaceArray() on an array-less vector");
 77:     PetscCall(VecResetArray(z)); // array = {4.0, 4.0, ...}
 78:     PetscCheck(array[0] == 4.0 && array[n - 1] == 4.0, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Test failed with VecResetArray() on an array-less vector");
 79:     PetscCall(PetscFree(array));

 81:     // Place the device array once more; z must still work after the arrays were reset
 82:     PetscCall(VecKokkosPlaceArray(z, kz.data())); // z = {3.0, 1.0, ...}
 83:     PetscCall(VecNorm(z, NORM_INFINITY, &norm));
 84:     PetscCheck(PetscAbsReal(norm - 3.0) < PETSC_SMALL, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Test failed with a second VecKokkosPlaceArray() on an array-less vector");
 85:     PetscCall(VecKokkosResetArray(z));
 86:     PetscCall(VecDestroy(&z));
 87:   }
 88:   PetscCall(PetscFinalize());
 89:   return 0;
 90: }

 92: /*TEST

 94:     test:
 95:       requires: kokkos_kernels
 96:       nsize: {{1 2}}
 97:       output_file: output/empty.out

 99: TEST*/