Generalize constraints on MPI rank counts for tests
[alexxy/gromacs.git] / src / gromacs / domdec / tests / haloexchange_mpi.cpp
index c763db0a9e58c57772a3f8a7450058d5ba0f451a..fac6cc09abc64ab40da09c11e394436caf7ac87f 100644 (file)
@@ -1,7 +1,7 @@
 /*
  * This file is part of the GROMACS molecular simulation package.
  *
- * Copyright (c) 2020, by the GROMACS development team, led by
+ * Copyright (c) 2020,2021, by the GROMACS development team, led by
  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
  * and including many others, as listed in the AUTHORS file in the
  * top-level source directory and at http://www.gromacs.org.
@@ -63,8 +63,8 @@
 #if GMX_GPU_CUDA
 #    include "gromacs/gpu_utils/device_stream.h"
 #    include "gromacs/gpu_utils/devicebuffer.h"
-#    include "gromacs/gpu_utils/gpueventsynchronizer.cuh"
 #endif
+#include "gromacs/gpu_utils/gpueventsynchronizer.h"
 #include "gromacs/gpu_utils/hostallocator.h"
 #include "gromacs/mdtypes/inputrec.h"
 
@@ -150,8 +150,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,6 +164,7 @@ void gpuHalo(gmx_domdec_t* dd, matrix box, HostVector<RVec>* h_x, int numAtomsTo
             gpuHaloExchange[d][pulse].communicateHaloCoordinates(box, &coordinatesReadyOnDeviceEvent);
         }
     }
+    MPI_Barrier(MPI_COMM_WORLD);
 
     GpuEventSynchronizer haloCompletedEvent;
     haloCompletedEvent.markEvent(deviceStream);
@@ -191,8 +192,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 +511,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 +567,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 +624,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 +680,7 @@ TEST(HaloExchangeTest, Coordinates2dHaloWith1PulseInEachDim)
 
 TEST(HaloExchangeTest, Coordinates2dHaloWith2PulsesInDim1)
 {
-    GMX_MPI_TEST(4);
+    GMX_MPI_TEST(RequireRankCount<4>);
 
     // Set up atom data
     const int        numHomeAtoms  = 10;