MPI gpu-direct: make it optional and have fallback solution with device-host explicit transfers (#389)

* gpu-direct: handle communications when MPI GPU DIRECT is not there/available.

---------

Co-authored-by: Geoffroy Lesur <geoffroy.lesur@univ-grenoble-alpes.fr>

Author: svalat

CMakeLists.txt

@@ -16,6 +16,9 @@ option(Idefix_DEBUG "Enable Idefix debug features (makes the code very slow)" OF
 option(Idefix_RUNTIME_CHECKS "Enable runtime sanity checks" OFF)
 option(Idefix_WERROR "Treat compiler warnings as errors" OFF)
 option(Idefix_PYTHON "Enable python bindings (requires pybind11)" OFF)
+option(Idefix_MPI_GPU_DIRECT "Enable usage of GPU direct to avoid CPU-GPU copies for MPI messages." ON)
+option(Idefix_MPI_GPU_FORCE_COPY "When GPU_DIRECT is disabled, force using transfers for validation purpose." OFF)
+mark_as_advanced(Idefix_MPI_GPU_FORCE_TRANSFERS)
 option(Idefix_SUPPRESS_FMA "Disable FMA (fused multiply-add) contraction/codegen. Useful for code validation across architectures." OFF)
 set(Idefix_PROBLEM_DIR "${CMAKE_BINARY_DIR}" CACHE STRING "Problem directory to build for.")
 set(Idefix_CXX_FLAGS "" CACHE STRING "Additional compiler/linker flag")
@@ -36,6 +39,10 @@ set_property(CACHE Idefix_PRECISION PROPERTY STRINGS Double Single)
 set(Idefix_LOOP_PATTERN "Default" CACHE STRING "Loop pattern for idefix_for")
 set_property(CACHE Idefix_LOOP_PATTERN PROPERTY STRINGS Default SIMD Range MDRange TeamPolicy TeamPolicyInnerVector)
 
+set(Idefix_MPI_MODE "Persistent" CACHE STRING "MPI communication mode to use.")
+set_property(CACHE Idefix_MPI_MODE PROPERTY STRINGS Persistent Blocking NonBlocking)
+mark_as_advanced(Idefix_MPI_MODE)
+
 # load git revision tools
 list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/")
 include(GetGitRevisionDescription)
@@ -109,6 +116,10 @@ if(Idefix_MPI)
   add_subdirectory(src/mpi)
 endif()
 
+if (Idefix_MPI_GPU_DIRECT)
+  add_compile_definitions("WITH_MPI_GPU_DIRECT")
+endif()
+
 if(Idefix_HDF5)
   add_compile_definitions("WITH_HDF5")
   if(Idefix_MPI)
@@ -245,6 +256,20 @@ if(${Idefix_PRECISION} STREQUAL "Single")
   add_compile_definitions("SINGLE_PRECISION")
 endif()
 
+# MPI mode
+if(${Idefix_MPI_MODE} STREQUAL "Persistent")
+  add_compile_definitions("COMMUNICATION_MODE_PERSISTENT")
+elseif(${Idefix_MPI_MODE} STREQUAL "NonBlocking")
+  add_compile_definitions("COMMUNICATION_MODE_NON_BLOCKING")
+else(NOT ${Idefix_MPI_MODE} STREQUAL "Blocking")
+  message(ERROR "Unknown MPI communication mode : Idefix_MPI_MODE=${Idefix_MPI_MODE}")
+endif()
+
+# GPU direct disabled and force transfers
+if (${Idefix_MPI_GPU_FORCE_COPY})
+  add_compile_definitions("WITH_GPU_FORCE_COPY")
+endif()
+
 target_include_directories(idefix PUBLIC
                            "${Idefix_PROBLEM_DIR_ABS}"
                            )
@@ -297,6 +322,7 @@ endif()
 message(STATUS "    MPI:  ${Idefix_MPI}")
 message(STATUS "    HDF5: ${Idefix_HDF5}")
 message(STATUS "    Python: ${Idefix_PYTHON}")
+message(STATUS "    GPU direct: ${Idefix_MPI_GPU_DIRECT}")
 message(STATUS "    Reconstruction: ${Idefix_RECONSTRUCTION}")
 message(STATUS "    Precision: ${Idefix_PRECISION}")
 message(STATUS "    Version: ${Idefix_VERSION}")

doc/source/faq.rst

@@ -27,6 +27,9 @@ I have a complex setup, and have written some functions in separate .cpp files.
 I want to run on the GPUs of xxx machine, how do I proceed?
   Check the examples in :ref:`setupExamples`
 
+I don't have a CUDA or HIP aware MPI with GPU DIRECT support, how I can run *Idefix* ?
+  You can look at the compile option `-DIdefix_MPI_GPU_DIRECT=OFF` in :ref:`makefile` to disable GPU DIRECT usage. Note that this will decrease performances.
+
 Compilation
 -----------
 

doc/source/reference/makefile.rst

@@ -1,3 +1,5 @@
+.. _makefile:
+
 Code configuration with Cmake
 =============================
 
@@ -33,6 +35,10 @@ Several options can be enabled from the command line (or are accessible with ``c
 ``-D Idefix_MPI=ON``
     Enable MPI parallelisation. Requires an MPI library. When used in conjonction with CUDA (Nvidia GPUs), a CUDA-aware MPI library is required by *Idefix*.
 
+``-D Idefix_MPI_GPU_DIRECT=OFF``
+    Disable the usage of MPI GPU direct for inter-process communication. In this case it uses a copy between GPU and GPU before and after making communications.
+    It is usefull if your MPI is not GPU-aware. Note that this option decreases performances.
+
 ``-D Idefix_DEFS=foo.hpp``
     Specify a particular filename to be used in place of the default problem file ``definitions.hpp``
 

pytools/idfx_test.py

@@ -92,6 +92,20 @@ def __init__(self, current_test_file, name=""):
 
         parser.add_argument("-mpi", help="Enable MPI", action="store_true")
 
+        parser.add_argument(
+            "-mpiMode",
+            default="Persistent",
+            choices=["Default", "Persistent", "NonBlocking", "Blocking"],
+            help="MPI communication mode to use (Default)",
+            type=str,
+        )
+
+        parser.add_argument(
+            "-mpiGpuForceCopy",
+            action="store_true",
+            help="Disable MPI GPU direct and force memory transfers for validation purpose.",
+        )
+
         parser.add_argument(
             "-all",
             help="Do all test suite (otherwise, just do the test with the current configuration)",
@@ -266,6 +280,13 @@ def _genCmakeCommand(self, definitionFile=""):
         else:
             comm.append("-DIdefix_MPI=OFF")
 
+        if self.mpiMode != "Default":
+            comm.append(f"-DIdefix_MPI_MODE={self.mpiMode}")
+
+        if self.mpiGpuForceCopy:
+            comm.append("-DIdefix_MPI_GPU_DIRECT=OFF")
+            comm.append("-DIdefix_MPI_GPU_FORCE_COPY=ON")
+
         if self.reconstruction == 2:
             comm.append("-DIdefix_RECONSTRUCTION=Linear")
         elif self.reconstruction == 3:

src/fluid/boundary/axis.cpp

@@ -25,7 +25,7 @@ void Axis::ShowConfig() {
 void Axis::SymmetrizeEx1Side(int jref, IdefixArray3D<real> Ex1) {
 #if DIMENSIONS == 3
 
-  IdefixArray1D<real> Ex1Avg = this->Ex1Avg;
+  auto Ex1Avg = this->Ex1Avg.deviceView();
 
   idefix_for("Ex1_ini",0,data->np_tot[IDIR],
       KOKKOS_LAMBDA(int i) {
@@ -40,7 +40,7 @@ void Axis::SymmetrizeEx1Side(int jref, IdefixArray3D<real> Ex1) {
     #ifdef WITH_MPI
       Kokkos::fence();
       // sum along all of the processes on the same r
-      MPI_Allreduce(MPI_IN_PLACE, Ex1Avg.data(), data->np_tot[IDIR], realMPI,
+      idfx::MPI_Allreduce(MPI_IN_PLACE, this->Ex1Avg, data->np_tot[IDIR], realMPI,
                     MPI_SUM, data->mygrid->AxisComm);
     #endif
   }
@@ -88,7 +88,8 @@ void Axis::RegularizeCurrentSide(int side) {
       jc = data->end[JDIR]-1;
       sign = -1;
     }
-    IdefixArray1D<real> BAvg = this->Ex1Avg;
+    auto BAvg = this->Ex1Avg.deviceView();
+    auto BAvgComm = this->Ex1Avg;
     IdefixArray1D<real> x1 = data->x[IDIR];
     IdefixArray1D<real> dx3 = data->dx[KDIR];
     IdefixArray1D<real> dx2 = data->dx[JDIR];
@@ -107,7 +108,7 @@ void Axis::RegularizeCurrentSide(int side) {
       #ifdef WITH_MPI
         Kokkos::fence();
         // sum along all of the processes on the same r
-        MPI_Allreduce(MPI_IN_PLACE, BAvg.data(), data->np_tot[IDIR], realMPI,
+        MPI_Allreduce(MPI_IN_PLACE, BAvgComm, data->np_tot[IDIR], realMPI,
                       MPI_SUM, data->mygrid->AxisComm);
       #endif
     }
@@ -169,7 +170,7 @@ void Axis::FixBx2sAxis(int side) {
   // Compute the values of Bx and By that are consistent with BX2 along the axis
   #if DIMENSIONS == 3
     IdefixArray4D<real> Vs = this->Vs;
-    IdefixArray2D<real> BAvg = this->BAvg;
+    auto BAvg = this->BAvg.deviceView();
     IdefixArray1D<real> phi = data->x[KDIR];
 
     int jin = 0;
@@ -208,7 +209,7 @@ void Axis::FixBx2sAxis(int side) {
       Kokkos::fence();
       #ifdef WITH_MPI
         // sum along all of the processes on the same r
-        MPI_Allreduce(MPI_IN_PLACE, BAvg.data(), 2*data->np_tot[IDIR], realMPI,
+        idfx::MPI_Allreduce(MPI_IN_PLACE, this->BAvg, 2*data->np_tot[IDIR], realMPI,
                       MPI_SUM, data->mygrid->AxisComm);
       #endif
     }
@@ -392,7 +393,6 @@ void Axis::ExchangeMPI(int side) {
   idfx::pushRegion("Axis::ExchangeMPI");
   #ifdef WITH_MPI
   // Load  the buffers with data
-  [[maybe_unused]] int ibeg,iend,jbeg,jend,kbeg,kend;
   int offset;
   int ny;
   Buffer bufferSend = this->bufferSend;
@@ -409,7 +409,7 @@ void Axis::ExchangeMPI(int side) {
   MPI_Status recvStatus;
 
   double tStart = MPI_Wtime();
-  MPI_SAFE_CALL(MPI_Start(&recvRequest));
+  MPI_SAFE_CALL(idfx::MPI_Start(&recvRequest));
   idfx::mpiCallsTimer += MPI_Wtime() - tStart;
 
   // Coordinates of the ghost region which needs to be transfered
@@ -471,8 +471,8 @@ void Axis::ExchangeMPI(int side) {
   Kokkos::fence();
 
   tStart = MPI_Wtime();
-  MPI_SAFE_CALL(MPI_Start(&sendRequest));
-  MPI_Wait(&recvRequest,&recvStatus);
+  MPI_SAFE_CALL(idfx::MPI_Start(&sendRequest));
+  idfx::MPI_Wait(&recvRequest,&recvStatus);
   idfx::mpiCallsTimer += MPI_Wtime() - tStart;
 
   // Unpack
@@ -526,8 +526,7 @@ void Axis::ExchangeMPI(int side) {
     } // MHD
   }
 
-  MPI_Wait(&sendRequest, &sendStatus);
-
+  idfx::MPI_Wait(&sendRequest, &sendStatus);
   idfx::mpiCallsTimer += MPI_Wtime() - tStart;
 
 
@@ -594,11 +593,11 @@ void Axis::InitMPI() {
   MPI_SAFE_CALL(MPI_Cart_shift(data->mygrid->AxisComm,0,data->mygrid->nproc[KDIR]/2,
                                &procRecv,&procSend ));
 
-  MPI_SAFE_CALL(MPI_Send_init(bufferSend.data(), bufferSend.Size(), realMPI, procSend,
-                650, data->mygrid->AxisComm, &sendRequest));
+  MPI_SAFE_CALL(idfx::MPI_Send_init(bufferSend.commView(), bufferSend.Size(),
+                realMPI, procSend, 650, data->mygrid->AxisComm, &sendRequest));
 
-  MPI_SAFE_CALL(MPI_Recv_init(bufferRecv.data(), bufferRecv.Size(), realMPI, procRecv,
-                650, data->mygrid->AxisComm, &recvRequest));
+  MPI_SAFE_CALL(idfx::MPI_Recv_init(bufferRecv.commView(), bufferRecv.Size(),
+                realMPI, procRecv, 650, data->mygrid->AxisComm, &recvRequest));
 
   #endif
   idfx::popRegion();

src/fluid/boundary/axis.hpp

@@ -57,8 +57,8 @@ class Axis {
 
   enum {faceTop, faceBot};
 #ifdef WITH_MPI
-  MPI_Request sendRequest;
-  MPI_Request recvRequest;
+  idfx::MPI_Request_1D<real> sendRequest;
+  idfx::MPI_Request_1D<real> recvRequest;
 
   Buffer bufferSend;
   Buffer bufferRecv;
@@ -71,8 +71,8 @@ class Axis {
 #endif
   void InitMPI();
 
-  IdefixArray1D<real> Ex1Avg;
-  IdefixArray2D<real> BAvg;
+  idfx::IdefixCommArray1D<real> Ex1Avg;
+  idfx::IdefixCommArray2D<real> BAvg;
   bool haveCurrent;
   IdefixArray2D<real> JAvg;
   IdefixArray1D<int> symmetryVc;
@@ -155,8 +155,8 @@ Axis::Axis(Boundary<Phys> *boundary) {
     }
     Kokkos::deep_copy(symmetryVs, symmetryVsHost);
 
-    this->Ex1Avg = IdefixArray1D<real>("Axis:Ex1Avg",data->np_tot[IDIR]);
-    this->BAvg = IdefixArray2D<real>("Axis:BxAvg",data->np_tot[IDIR],2);
+    this->Ex1Avg = idfx::IdefixCommArray1D<real>("Axis:Ex1Avg",data->np_tot[IDIR]);
+    this->BAvg = idfx::IdefixCommArray2D<real>("Axis:BxAvg",data->np_tot[IDIR],2);
     if(haveCurrent) {
       this->JAvg = IdefixArray2D<real>("Axis:JAvg",data->np_tot[IDIR],3);
     }

src/fluid/constrainedTransport/EMFexchange.hpp

@@ -41,7 +41,7 @@ void ConstrainedTransport<Phys>::ExchangeX1(IdefixArray3D<real> ey, IdefixArray3
   MPI_Status recvStatus[2];
 
   double tStart = MPI_Wtime();
-  MPI_SAFE_CALL(MPI_Startall(2, recvRequestX1));
+  MPI_SAFE_CALL(idfx::MPI_Startall(2, recvRequestX1));
   idfx::mpiCallsTimer += MPI_Wtime() - tStart;
 
   BoundaryType lbound = data->lbound[IDIR];
@@ -82,11 +82,10 @@ void ConstrainedTransport<Phys>::ExchangeX1(IdefixArray3D<real> ey, IdefixArray3
   Kokkos::fence();
 
   tStart = MPI_Wtime();
-  MPI_SAFE_CALL(MPI_Startall(2, sendRequestX1));
+  MPI_SAFE_CALL(idfx::MPI_Startall(2, sendRequestX1));
   // Wait for buffers to be received
-
-  MPI_Waitall(2,recvRequestX1,recvStatus);
-  MPI_Waitall(2, sendRequestX1, sendStatus);
+  idfx::MPI_Waitall(2,recvRequestX1,recvStatus);
+  idfx::MPI_Waitall(2, sendRequestX1, sendStatus);
   idfx::mpiCallsTimer += MPI_Wtime() - tStart;
 
   // Unpack
@@ -130,7 +129,7 @@ void ConstrainedTransport<Phys>::ExchangeX2(IdefixArray3D<real> ex, IdefixArray3
   double tStart = MPI_Wtime();
   MPI_Status sendStatus[2];
   MPI_Status recvStatus[2];
-  MPI_SAFE_CALL(MPI_Startall(2, recvRequestX2));
+  MPI_SAFE_CALL(idfx::MPI_Startall(2, recvRequestX2));
   idfx::mpiCallsTimer += MPI_Wtime() - tStart;
 
   BoundaryType lbound = data->lbound[JDIR];
@@ -169,10 +168,10 @@ void ConstrainedTransport<Phys>::ExchangeX2(IdefixArray3D<real> ex, IdefixArray3
   Kokkos::fence();
 
   tStart = MPI_Wtime();
-  MPI_SAFE_CALL(MPI_Startall(2, sendRequestX2));
+  MPI_SAFE_CALL(idfx::MPI_Startall(2, sendRequestX2));
   // Wait for buffers to be received
-  MPI_Waitall(2, recvRequestX2, recvStatus);
-  MPI_Waitall(2, sendRequestX2, sendStatus);
+  idfx::MPI_Waitall(2, recvRequestX2, recvStatus);
+  idfx::MPI_Waitall(2, sendRequestX2, sendStatus);
   idfx::mpiCallsTimer += MPI_Wtime() - tStart;
 
   // Unpack
@@ -216,7 +215,7 @@ void ConstrainedTransport<Phys>::ExchangeX3(IdefixArray3D<real> ex, IdefixArray3
   double tStart = MPI_Wtime();
   MPI_Status sendStatus[2];
   MPI_Status recvStatus[2];
-  MPI_SAFE_CALL(MPI_Startall(2, recvRequestX3));
+  MPI_SAFE_CALL(idfx::MPI_Startall(2, recvRequestX3));
   idfx::mpiCallsTimer += MPI_Wtime() - tStart;
 
   BoundaryType lbound = data->lbound[KDIR];
@@ -256,10 +255,10 @@ void ConstrainedTransport<Phys>::ExchangeX3(IdefixArray3D<real> ex, IdefixArray3
   Kokkos::fence();
 
   tStart = MPI_Wtime();
-  MPI_SAFE_CALL(MPI_Startall(2, sendRequestX3));
+  MPI_SAFE_CALL(idfx::MPI_Startall(2, sendRequestX3));
   // Wait for buffers to be received
-  MPI_Waitall(2, recvRequestX3, recvStatus);
-  MPI_Waitall(2, sendRequestX3, sendStatus);
+  idfx::MPI_Waitall(2, recvRequestX3, recvStatus);
+  idfx::MPI_Waitall(2, sendRequestX3, sendStatus);
   idfx::mpiCallsTimer += MPI_Wtime() - tStart;
 
   // Unpack