Set up event consumption counter in DomDecMpiTests

[alexxy/gromacs.git] / src / gromacs / domdec / tests / haloexchange_mpi.cpp

diff --git a/src/gromacs/domdec/tests/haloexchange_mpi.cpp b/src/gromacs/domdec/tests/haloexchange_mpi.cpp
index 1f1138fb36d5e41f910df13e0fa3fb63a43058d9..93aad8c288d62356af6124ba40995e6075bc5f67 100644 (file)
--- a/src/gromacs/domdec/tests/haloexchange_mpi.cpp
+++ b/src/gromacs/domdec/tests/haloexchange_mpi.cpp
@@ -53,6 +53,7 @@
 #include "config.h"
 
 #include <array>
+#include <numeric>
 #include <vector>
 
 #include <gtest/gtest.h>
@@ -140,7 +141,14 @@ void gpuHalo(gmx_domdec_t* dd, matrix box, HostVector<RVec>* h_x, int numAtomsTo
 
     copyToDeviceBuffer(&d_x, h_x->data(), 0, numAtomsTotal, deviceStream, GpuApiCallBehavior::Sync, nullptr);
 
-    GpuEventSynchronizer coordinatesReadyOnDeviceEvent;
+    const int numPulses = std::accumulate(
+            dd->comm->cd.begin(), dd->comm->cd.end(), 0, [](const int a, const auto& b) {
+                return a + b.numPulses();
+            });
+    const int numExtraConsumptions = GMX_THREAD_MPI ? 1 : 0;
+    // Will be consumed once for each pulse, and, with tMPI, once more for dim=0,pulse=0 case
+    GpuEventSynchronizer coordinatesReadyOnDeviceEvent(numPulses + numExtraConsumptions,
+                                                       numPulses + numExtraConsumptions);
     coordinatesReadyOnDeviceEvent.markEvent(deviceStream);
 
     std::array<std::vector<GpuHaloExchange>, DIM> gpuHaloExchange;
@@ -150,8 +158,8 @@ void gpuHalo(gmx_domdec_t* dd, matrix box, HostVector<RVec>* h_x, int numAtomsTo
     {
         for (int pulse = 0; pulse < dd->comm->cd[d].numPulses(); pulse++)
         {
-            gpuHaloExchange[d].push_back(GpuHaloExchange(
-                    dd, d, MPI_COMM_WORLD, deviceContext, deviceStream, deviceStream, pulse, nullptr));
+            gpuHaloExchange[d].push_back(
+                    GpuHaloExchange(dd, d, MPI_COMM_WORLD, deviceContext, pulse, nullptr));
         }
     }
 
@@ -164,10 +172,10 @@ void gpuHalo(gmx_domdec_t* dd, matrix box, HostVector<RVec>* h_x, int numAtomsTo
             gpuHaloExchange[d][pulse].communicateHaloCoordinates(box, &coordinatesReadyOnDeviceEvent);
         }
     }
+    // Barrier is needed to avoid other threads using events after its owner has exited and destroyed the context.
+    MPI_Barrier(MPI_COMM_WORLD);
 
-    GpuEventSynchronizer haloCompletedEvent;
-    haloCompletedEvent.markEvent(deviceStream);
-    haloCompletedEvent.waitForEvent();
+    deviceStream.synchronize();
 
     // Copy results back to host
     copyFromDeviceBuffer(
@@ -191,8 +199,9 @@ void define1dRankTopology(gmx_domdec_t* dd)
     int rank;
     MPI_Comm_rank(MPI_COMM_WORLD, &rank);
 
-    dd->neighbor[0][0] = (rank + 1) % 4;
-    dd->neighbor[0][1] = (rank == 0) ? 3 : rank - 1;
+    const int numRanks = getNumberOfTestMpiRanks();
+    dd->neighbor[0][0] = (rank + 1) % numRanks;
+    dd->neighbor[0][1] = (rank == 0) ? (numRanks - 1) : rank - 1;
 }
 
 /*! \brief Define 2D rank topology with 4 MPI tasks
@@ -509,7 +518,7 @@ void checkResults2dHaloWith2PulsesInDim1(const RVec* x, const gmx_domdec_t* dd,
 
 TEST(HaloExchangeTest, Coordinates1dHaloWith1Pulse)
 {
-    GMX_MPI_TEST(4);
+    GMX_MPI_TEST(RequireRankCount<4>);
 
     // Set up atom data
     const int        numHomeAtoms  = 10;
@@ -565,7 +574,7 @@ TEST(HaloExchangeTest, Coordinates1dHaloWith1Pulse)
 
 TEST(HaloExchangeTest, Coordinates1dHaloWith2Pulses)
 {
-    GMX_MPI_TEST(4);
+    GMX_MPI_TEST(RequireRankCount<4>);
 
     // Set up atom data
     const int        numHomeAtoms  = 10;
@@ -622,7 +631,7 @@ TEST(HaloExchangeTest, Coordinates1dHaloWith2Pulses)
 
 TEST(HaloExchangeTest, Coordinates2dHaloWith1PulseInEachDim)
 {
-    GMX_MPI_TEST(4);
+    GMX_MPI_TEST(RequireRankCount<4>);
 
     // Set up atom data
     const int        numHomeAtoms  = 10;
@@ -678,7 +687,7 @@ TEST(HaloExchangeTest, Coordinates2dHaloWith1PulseInEachDim)
 
 TEST(HaloExchangeTest, Coordinates2dHaloWith2PulsesInDim1)
 {
-    GMX_MPI_TEST(4);
+    GMX_MPI_TEST(RequireRankCount<4>);
 
     // Set up atom data
     const int        numHomeAtoms  = 10;