From 04c94501888fe7af94147a2ceb5beceab63a265d Mon Sep 17 00:00:00 2001
From: Alan Gray <alangraygerrit@gmail.com>
Date: Sat, 2 Nov 2019 12:24:43 -0700
Subject: [PATCH] Fix race condition in GPU X halo exchange with GPU update

When the GPU Update featured is enabled on multi-GPU with the GPU halo
exchange feature, it was possible that the GPU receiving the halo data
would still be working on the previous timestep and writing to the
same coordinate buffer in the GPU update, causing a race
condition. This change replaces the pre-existing synchronization
mechanism of a MPI Send/Recv of a boolean flag with the same but
exchanging the events that are recorded when the coordinates are ready
on the device. The remote event is then enqueued to the stream that
the GPU uses for pushing the data, satisfying the dependency of the
push on the remote update being completed on the GPU.

Change-Id: Ie16a93dcc3064d73da1e560d82b4e76870d03e52
---
 src/gromacs/domdec/gpuhaloexchange_impl.cu  | 20 +++++++++++++-------
 src/gromacs/domdec/gpuhaloexchange_impl.cuh |  5 ++++-
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/src/gromacs/domdec/gpuhaloexchange_impl.cu b/src/gromacs/domdec/gpuhaloexchange_impl.cu
index d917be3f39..09061493b3 100644
--- a/src/gromacs/domdec/gpuhaloexchange_impl.cu
+++ b/src/gromacs/domdec/gpuhaloexchange_impl.cu
@@ -235,7 +235,7 @@ void GpuHaloExchange::Impl::communicateHaloCoordinates(const matrix          box
         launchGpuKernel(kernelFn, config, nullptr, "Domdec GPU Apply X Halo Exchange", kernelArgs);
     }
 
-    communicateHaloData(d_x_, HaloQuantity::HaloCoordinates);
+    communicateHaloData(d_x_, HaloQuantity::HaloCoordinates, coordinatesReadyOnDeviceEvent);
 
     return;
 }
@@ -246,7 +246,7 @@ void GpuHaloExchange::Impl::communicateHaloForces(bool accumulateForces)
 {
 
     // Communicate halo data (in non-local stream)
-    communicateHaloData(d_f_, HaloQuantity::HaloForces);
+    communicateHaloData(d_f_, HaloQuantity::HaloForces, nullptr);
 
     float3* d_f = d_f_;
 
@@ -292,7 +292,9 @@ void GpuHaloExchange::Impl::communicateHaloForces(bool accumulateForces)
 }
 
 
-void GpuHaloExchange::Impl::communicateHaloData(float3* d_ptr, HaloQuantity haloQuantity)
+void GpuHaloExchange::Impl::communicateHaloData(float3*               d_ptr,
+                                                HaloQuantity          haloQuantity,
+                                                GpuEventSynchronizer* coordinatesReadyOnDeviceEvent)
 {
 
     void* sendPtr;
@@ -310,10 +312,14 @@ void GpuHaloExchange::Impl::communicateHaloData(float3* d_ptr, HaloQuantity halo
         recvRank  = recvRankX_;
 
 #if GMX_MPI
-        // Wait for signal from receiving task that it is ready, and similarly send signal to task that will push data to this task
-        char thisTaskIsReady, remoteTaskIsReady;
-        MPI_Sendrecv(&thisTaskIsReady, sizeof(char), MPI_BYTE, recvRank, 0, &remoteTaskIsReady,
-                     sizeof(char), MPI_BYTE, sendRank, 0, mpi_comm_mysim_, MPI_STATUS_IGNORE);
+        // Wait for event from receiving task that remote coordinates are ready, and enqueue that event to stream used
+        // for subsequent data push. This avoids a race condition with the remote data being written in the previous timestep.
+        // Similarly send event to task that will push data to this task.
+        GpuEventSynchronizer* remoteCoordinatesReadyOnDeviceEvent;
+        MPI_Sendrecv(&coordinatesReadyOnDeviceEvent, sizeof(GpuEventSynchronizer*), MPI_BYTE,
+                     recvRank, 0, &remoteCoordinatesReadyOnDeviceEvent, sizeof(GpuEventSynchronizer*),
+                     MPI_BYTE, sendRank, 0, mpi_comm_mysim_, MPI_STATUS_IGNORE);
+        remoteCoordinatesReadyOnDeviceEvent->enqueueWaitEvent(nonLocalStream_);
 #endif
     }
     else
diff --git a/src/gromacs/domdec/gpuhaloexchange_impl.cuh b/src/gromacs/domdec/gpuhaloexchange_impl.cuh
index 626810b04a..9f4bf8cdd8 100644
--- a/src/gromacs/domdec/gpuhaloexchange_impl.cuh
+++ b/src/gromacs/domdec/gpuhaloexchange_impl.cuh
@@ -100,8 +100,11 @@ private:
     /*! \brief Data transfer wrapper for GPU halo exchange
      * \param [inout] d_ptr      pointer to coordinates or force buffer in GPU memory
      * \param [in] haloQuantity  switch on whether X or F halo exchange is being performed
+     * \param [in] coordinatesReadyOnDeviceEvent event recorded when coordinates have been copied to device
      */
-    void communicateHaloData(float3* d_ptr, HaloQuantity haloQuantity);
+    void communicateHaloData(float3*               d_ptr,
+                             HaloQuantity          haloQuantity,
+                             GpuEventSynchronizer* coordinatesReadyOnDeviceEvent);
 
     /*! \brief Data transfer for GPU halo exchange using CUDA memcopies
      * \param [inout] sendPtr    address to send data from
-- 
2.11.4.GIT