[GPU] Enable GPU-GPU transfer through FMS

[edoyango]

• Get FMS2 working with nvfortran
• Figure out same MPI process transfers
• Figure out different MPI process transfers
• Implement gpu2gpu transfer in dynamic core
  • comes after loops are ported

[edoyango]

Things I know about transfers so far:

1. Transfers in MOM6 are setup using create_group_pass
   • create_group_pass in src/infra/MOM_domains.F90, which is a direct pass-through from src/infra/FMS{1,2}/MOM_domain_infra.F90
   • above calls mpp_create_group_update from src/FMS{1,2}/mpp/mpp_domains.F90
   • mpp_domains.F90 gets definition of mpp_create_group_update from include/mpp_domains_misc.inc, which gets it from include/mpp_group_update.h.
   • It looks like the creation of the group is setting up metadata and doesn't interact with MPI. Probably not important for the porting.
2. do_group_update leads to the same path as above (src/FMS{1,2}/mpp/include/mpp_group_update.h).
   • mpp_do_group_update calls mpp_send and mpp_recv which come from mpp.f90
   • definition of mpp_send/recv comes from a path of includes mpp_comm.inc -> mpp_comm_mpi.h -> mpp_transmit_mpi.h -> mpp_transmit.inc. There's also nocomm versions which we'll want to keep an eye on.
   • mpp_send/recv seem to be a thin wrapper around mpp_transmit which is defined in mpp_transmit_mpi/nocomm.h.
   • There's some logic in mpp_transmit which does a direct copy if the send/recv process is the same. We could port this part in the short term so single process case works without having to ship data to/from CPU.
   • The challenge might be the use of buffers both in mpp_do_group_update and mpp_transmit and the use of addresses (loc) and possibly pointers.

[edoyango]

Before calling mpp_send/recv, do_group_update packs data from the original array into buffer and the received data is also received in the buffer. The logic is in group_update_pack/unpack.inc.

The packing copies data from fieldx/y into buffer, and fieldx/y are local array variables, so I don't know when fieldx/y get populated with the original data to be sent.

[edoyango]

regarding how fieldx/y get their values, I think I've just misunderstood how cray pointers work. Prior to packing, ptr_fieldx/y get set based on the address that was collected from mpp_create_group_update which mean fieldx/y arrays point to the original arrays? Seems like

real :: field(nx, ny, nz)
pointer(ptr_field, field)

ptr_field = loc(arr)

Is equivalent to

real, pointer :: field(nx, ny, nz)
ptr => arr

I have no idea yet how NVIDIA's OpenMP implementation will behave with cray pointers...

[edoyango]

I have no idea yet how NVIDIA's OpenMP implementation will behave with cray pointers...

nvfortran seems pretty happy to work with cray compilers, but only with !$omp target teams distribute parallel do and not do concurrent or any !$omp target loop type situation. Toy example.

[edoyango]

I managed to swap out the do_group_pass subroutines in MOM_barotropic.F90 with a do_group_pass_omp which:

• allocates buffer on GPU
• Copies data from variables on GPU into buffer
• uses the device pointer to send/recv data in the corresponding FMS+MPI calls

here is my FMS branch with crude changes. For future reference, the changes I made were:

1. Copied MPP_DO_GROUP_UPDATE_ to MPP_DO_GROUP_UPDATE_OMP_ in mpp_group_update.h and a corresponding preprocessor declaration in mpp_domains_misc.inc.
2. MPP_DO_GROUP_UPDATE_OMP allocates buffer on GPU and includes group_update_{,un}pack_omp.inc which are copies of the non-_omp files. The pack/unpack code has only relevant loops annotated with the openmp directives so they run on the GPU
   • MPP_DO_GROUP_UPDATE_OMP also calls modified versions of mpp_send and mpp_recv
3. add a use_dev optional logical arg to mpp_send and mpp_recv.
4. The use_dev arg is passed to the mpp_transmit subroutine which does the sending/receiving.
5. Extra if branches were added for the openmp offload wrapper code around the MPI calls. I couldn't use the if functionality in the directive itself as I got segfaults when the condition was .false., indicating it wasn't turning itself off properly.

Performance wise (on my laptop, benchmark case 1 CPU, 1 GPU i.e., self-MPI transfers):

• before (using CPU-MPI transfers and target update): (Ocean barotropic mode stepping) 6.643228
• after (GPU-MPI transfers, no target update): (Ocean barotropic mode stepping) 4.446166

So I'm getting about 1/3rd performance improvement by removing the updates and doing gpu2gpu transfer (and results are identical!).

[edoyango]

There's some logic in mpp_transmit which does a direct copy if the send/recv process is the same. We could port this part in the short term so single process case works without having to ship data to/from CPU.

I should add that this comment is incorrect as the branch I saw is actually only for broadcasts and not for p2p which mpp_do_group_update uses.

[edoyango]

Can confirm my changes work across multiple GPUs

# on 1 node, 2 GPUs
mpiexec -n 1 -x CUDA_VISIBLE_DEVICES=0 ../build/MOM6 : -n 1 -x CUDA_VISIBLE_DEVICES=1 ../build/MOM6

With GPU2GPU MPI:

(Ocean BT pre-step halo updates)      0.359533      0.363895      0.361714      0.002181  0.003    41     0     1
...
(Ocean barotropic mode stepping)      3.987241      3.988508      3.987875      0.000633  0.037    31     0     1

CPU2CPU MPI:

(Ocean BT pre-step halo updates)      1.168880      1.172374      1.170627      0.001747  0.011    41     0     1
...
(Ocean barotropic mode stepping)      5.160809      5.162044      5.161426      0.000617  0.047    31     0     1

And results are bitwise identical.

[edoyango]

Get FMS2 working with nvfortran

Latest FMS works, but -stdpar=gpu needs to be excluded. Not a problem since do concurrent offload doesn't work with cray pointers anyway.

[edoyango]

In my FMS fork, I added the ability to use openmp-offloaded arrays in mpi comms for blocking case in latest. Corresponding mom6 branch