Fix MPI Transfer for ib_markers on GPUs

Author: Anand

src/simulation/m_mpi_proxy.fpp

@@ -345,12 +345,40 @@ contains
         #:endfor
         call nvtxEndRange ! Packbuf
 
-        call nvtxStartRange("IB-MARKER-SENDRECV")
-        call MPI_SENDRECV( &
-            ib_buff_send, buffer_count, MPI_INTEGER, dst_proc, send_tag, &
-            ib_buff_recv, buffer_count, MPI_INTEGER, src_proc, recv_tag, &
-            MPI_COMM_WORLD, MPI_STATUS_IGNORE, ierr)
-        call nvtxEndRange ! RHS-MPI-SENDRECV-(NO)-RDMA
+        #:for rdma_mpi in [False, True]
+            if (rdma_mpi .eqv. ${'.true.' if rdma_mpi else '.false.'}$) then
+                #:if rdma_mpi
+                    #:call GPU_HOST_DATA(use_device='[ib_buff_send, ib_buff_recv]')
+                        call nvtxStartRange("IB-MARKER-SENDRECV-RDMA")
+
+                        call MPI_SENDRECV( &
+                            ib_buff_send, buffer_count, MPI_INTEGER, dst_proc, send_tag, &
+                            ib_buff_recv, buffer_count, MPI_INTEGER, src_proc, recv_tag, &
+                            MPI_COMM_WORLD, MPI_STATUS_IGNORE, ierr)
+
+                        call nvtxEndRange ! RHS-MPI-SENDRECV-(NO)-RDMA
+
+                    #:endcall GPU_HOST_DATA
+                    $:GPU_WAIT()
+                #:else
+                    call nvtxStartRange("IB-MARKER-DEV2HOST")
+                    $:GPU_UPDATE(host='[ib_buff_send]')
+                    call nvtxEndRange
+                    call nvtxStartRange("IB-MARKER-SENDRECV-NO-RMDA")
+
+                    call MPI_SENDRECV( &
+                        ib_buff_send, buffer_count, MPI_INTEGER, dst_proc, send_tag, &
+                        ib_buff_recv, buffer_count, MPI_INTEGER, src_proc, recv_tag, &
+                        MPI_COMM_WORLD, MPI_STATUS_IGNORE, ierr)
+
+                    call nvtxEndRange ! RHS-MPI-SENDRECV-(NO)-RDMA
+
+                    call nvtxStartRange("IB-MARKER-HOST2DEV")
+                    $:GPU_UPDATE(device='[ib_buff_recv]')
+                    call nvtxEndRange
+                #:endif
+            end if
+        #:endfor
 
         ! Unpack Received Buffer
         call nvtxStartRange("IB-MARKER-COMM-UNPACKBUF")