1: static char help[] = "Tests VecCUDAGetArray() and friends along with OpenMP target offload\n\n";

  3: #include <petscvec.h>
  4: #include <petscdevice.h>

  6: int main(int argc, char **argv)
  7: {
  8:   const PetscInt     n = 16;
  9:   PetscScalar       *a;
 10:   const PetscScalar *ar;
 11:   Vec                x;
 12:   PetscDeviceContext dctx;

 14:   PetscFunctionBeginUser;
 15:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
 16:   PetscCall(PetscDeviceContextGetCurrentContext(&dctx));

 18:   PetscCall(VecCreate(PETSC_COMM_WORLD, &x));
 19:   PetscCall(VecSetSizes(x, PETSC_DECIDE, n));
 20:   PetscCall(VecSetType(x, VECCUDA));

 22:   // A PETSc write queued on the current-context stream; the first OpenMP kernel overwrites the
 23:   // same buffer, so the two race (write-after-write) unless PETSc's stream is synchronized first.
 24:   PetscCall(VecSet(x, 100.0));

 26:   /*
 27:      This test interleaves PETSc Vec operations with OpenMP target offload to exercise their
 28:      stream interaction. PETSc queues work on the current device context's stream, and the
 29:      VecCUDAGetArray*() calls return without synchronizing; the omp target regions run on a
 30:      separate, runtime-managed stream with no implicit ordering relative to PETSc's. So whenever
 31:      an OpenMP kernel touches a buffer a preceding PETSc operation is still writing on PETSc's
 32:      stream (e.g. VecSet() or VecScale() here), the current device context must be synchronized
 33:      first, or the two streams race. The omp target regions have no nowait, so they block the
 34:      host on return, which orders the subsequent PETSc calls after the kernel.
 35:   */

 37:   // Write access: overwrite with x[i] = i + 1.
 38:   PetscCall(VecCUDAGetArrayWrite(x, &a));
 39:   PetscCall(PetscDeviceContextSynchronize(dctx));
 40: #pragma omp target teams distribute parallel for is_device_ptr(a)
 41:   for (PetscInt i = 0; i < n; i++) a[i] = i + 1.0;
 42:   PetscCall(VecCUDARestoreArrayWrite(x, &a));

 44:   // A PETSc operation the next OpenMP kernel depends on: double every entry, x[i] = 2(i + 1).
 45:   PetscCall(VecScale(x, 2.0));

 47:   // Read-write access: add to what VecScale() produced, so sync PETSc's stream first, x[i] = 2(i + 1) + 3.
 48:   PetscCall(VecCUDAGetArray(x, &a));
 49:   PetscCall(PetscDeviceContextSynchronize(dctx));
 50: #pragma omp target teams distribute parallel for is_device_ptr(a)
 51:   for (PetscInt i = 0; i < n; i++) a[i] = a[i] + 3.0;
 52:   PetscCall(VecCUDARestoreArray(x, &a));

 54:   // A PETSc operation consuming the OpenMP result; the synchronous target above ordered it, x[i] = 2(i + 1) + 4.
 55:   PetscCall(VecShift(x, 1.0));

 57:   // Read-only access round-trip.
 58:   PetscCall(VecCUDAGetArrayRead(x, &ar));
 59:   PetscCall(VecCUDARestoreArrayRead(x, &ar));

 61:   PetscCall(PetscObjectSetName((PetscObject)x, "x"));
 62:   PetscCall(VecView(x, PETSC_VIEWER_STDOUT_WORLD));

 64:   PetscCall(VecDestroy(&x));
 65:   PetscCall(PetscFinalize());
 66:   return 0;
 67: }

 69: /*TEST

 71:   test:
 72:     requires: cuda defined(PETSC_HAVE_OPENMP_TARGET_OFFLOAD_CC)
 73:     output_file: output/ex67_1.out

 75: TEST*/