1: !
  2: ! Description: Tests VecCUDAGetArray() and friends along with OpenMP target offload.
  3: ! It is the Fortran version of ex67.c.
  4: !
  5: #include <petsc/finclude/petscvec.h>
  6: program main
  7:   use petscvec
  8:   implicit none

 10:   PetscInt, parameter :: n = 16
 11:   PetscScalar, parameter :: one = 1.0, two = 2.0, three = 3.0, hundred = 100.0
 12:   PetscScalar, pointer, dimension(:) :: a
 13:   PetscInt :: i
 14:   PetscErrorCode :: ierr
 15:   Vec :: x
 16:   PetscDeviceContext :: dctx

 18:   PetscCallA(PetscInitialize(ierr))
 19:   PetscCallA(PetscDeviceContextGetCurrentContext(dctx, ierr))

 21:   PetscCallA(VecCreate(PETSC_COMM_WORLD, x, ierr))
 22:   PetscCallA(VecSetSizes(x, PETSC_DECIDE, n, ierr))
 23:   PetscCallA(VecSetType(x, 'cuda', ierr))

 25: !  A PETSc write queued on the current-context stream; the first OpenMP kernel overwrites the
 26: !  same buffer, so the two race (write-after-write) unless PETSc's stream is synchronized first.

 28:   PetscCallA(VecSet(x, hundred, ierr))

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

 39: !  Write access: overwrite with x(i) = i.

 41:   PetscCallA(VecCUDAGetArrayWrite(x, a, ierr))
 42:   PetscCallA(PetscDeviceContextSynchronize(dctx, ierr))
 43:   !$omp target teams distribute parallel do is_device_ptr(a)
 44:   do i = 1, n
 45:     a(i) = i
 46:   end do
 47:   !$omp end target teams distribute parallel do
 48:   PetscCallA(VecCUDARestoreArrayWrite(x, a, ierr))

 50: !  A PETSc operation the next OpenMP kernel depends on: double every entry, x(i) = 2*i.

 52:   PetscCallA(VecScale(x, two, ierr))

 54: !  Read-write access: add to what VecScale() produced, so sync PETSc's stream first, x(i) = 2*i + 3.

 56:   PetscCallA(VecCUDAGetArray(x, a, ierr))
 57:   PetscCallA(PetscDeviceContextSynchronize(dctx, ierr))
 58:   !$omp target teams distribute parallel do is_device_ptr(a)
 59:   do i = 1, n
 60:     a(i) = a(i) + three
 61:   end do
 62:   !$omp end target teams distribute parallel do
 63:   PetscCallA(VecCUDARestoreArray(x, a, ierr))

 65: !  A PETSc operation consuming the OpenMP result; the synchronous target above ordered it, x(i) = 2*i + 4.

 67:   PetscCallA(VecShift(x, one, ierr))

 69: !  Read-only access round-trip.

 71:   PetscCallA(VecCUDAGetArrayRead(x, a, ierr))
 72:   PetscCallA(VecCUDARestoreArrayRead(x, a, ierr))

 74:   PetscCallA(PetscObjectSetName(x, 'x', ierr))
 75:   PetscCallA(VecView(x, PETSC_VIEWER_STDOUT_WORLD, ierr))

 77:   PetscCallA(VecDestroy(x, ierr))
 78:   PetscCallA(PetscFinalize(ierr))
 79: end

 81: !
 82: !/*TEST
 83: !
 84: ! test:
 85: !    requires: cuda defined(PETSC_HAVE_OPENMP_TARGET_OFFLOAD_FC)
 86: !    output_file: output/ex67_1.out
 87: !
 88: !TEST*/