From 208fc5c9ed9feca471191f64d6a4dcb7f863d00b Mon Sep 17 00:00:00 2001
From: Kjetil Olsen Lye <kjetil.olsen.lye@sintef.no>
Date: Tue, 24 Jun 2025 14:12:17 +0200
Subject: [PATCH] Added CLI argument for GPU aware MPI

---
 .../linalg/FlowLinearSolverParameters.cpp     |  9 ++++
 .../linalg/FlowLinearSolverParameters.hpp     |  2 +
 .../linalg/gpuistl/GpuOwnerOverlapCopy.hpp    | 53 +++++++++++++++++++
 .../linalg/gpuistl/SolverAdapter.hpp          |  5 --
 4 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/opm/simulators/linalg/FlowLinearSolverParameters.cpp b/opm/simulators/linalg/FlowLinearSolverParameters.cpp
index d90e871e4b..dc5ff467c6 100644
--- a/opm/simulators/linalg/FlowLinearSolverParameters.cpp
+++ b/opm/simulators/linalg/FlowLinearSolverParameters.cpp
@@ -47,6 +47,7 @@ void FlowLinearSolverParameters::init(bool cprRequestedInDataFile)
     linear_solver_print_json_definition_ = Parameters::Get<Parameters::LinearSolverPrintJsonDefinition>();
     cpr_reuse_setup_  = Parameters::Get<Parameters::CprReuseSetup>();
     cpr_reuse_interval_  = Parameters::Get<Parameters::CprReuseInterval>();
+    gpu_aware_mpi_ = Parameters::Get<Parameters::GPUAwareMPI>();
 
     if (!Parameters::IsSet<Parameters::LinearSolver>() && cprRequestedInDataFile) {
         linsolver_ = "cpr";
@@ -162,6 +163,13 @@ void FlowLinearSolverParameters::registerParameters()
     Parameters::Register<Parameters::LinearSolverAccelerator>
         ("Choose the backend for the linear solver, usage: "
          "'--linear-solver-accelerator=[cpu|gpu]'.");
+    Parameters::Register<Parameters::GPUAwareMPI>
+        ("MPI communication use GPU aware MPI in the sense that "
+            "it will use GPU direct communication. Setting this to true "
+            " will require that the MPI implementation "
+            "supports GPU direct communication. "
+            "If you are unsure, set this to false. "
+            "Usage: --gpu-aware-mpi=[true|false]. ");
 
     Parameters::SetDefault<Parameters::LinearSolverVerbosity>(0);
 }
@@ -191,6 +199,7 @@ void FlowLinearSolverParameters::reset()
     opencl_platform_id_       = 0;
     opencl_ilu_parallel_      = true;
     linear_solver_accelerator_ = Parameters::LinearSolverAcceleratorType::CPU;
+    gpu_aware_mpi_              = false;
 }
 
 } // namespace Opm
diff --git a/opm/simulators/linalg/FlowLinearSolverParameters.hpp b/opm/simulators/linalg/FlowLinearSolverParameters.hpp
index 82b0cd428e..ea4e6ffa43 100644
--- a/opm/simulators/linalg/FlowLinearSolverParameters.hpp
+++ b/opm/simulators/linalg/FlowLinearSolverParameters.hpp
@@ -86,6 +86,7 @@ struct AcceleratorMode { static constexpr auto value = "none"; };
 struct GpuDeviceId { static constexpr int value = 0; };
 struct OpenclPlatformId { static constexpr int value = 0; };
 struct OpenclIluParallel { static constexpr bool value = true; }; // note: false should only be used in debug
+struct GPUAwareMPI { static constexpr bool value = false; };
 } // namespace Opm::Parameters
 
 namespace Opm {
@@ -116,6 +117,7 @@ struct FlowLinearSolverParameters
     int opencl_platform_id_;
     bool opencl_ilu_parallel_;
     Parameters::LinearSolverAcceleratorType linear_solver_accelerator_;
+    bool gpu_aware_mpi_;
 
     FlowLinearSolverParameters() { reset(); }
 
diff --git a/opm/simulators/linalg/gpuistl/GpuOwnerOverlapCopy.hpp b/opm/simulators/linalg/gpuistl/GpuOwnerOverlapCopy.hpp
index fd5a7613a7..2ffccd2492 100644
--- a/opm/simulators/linalg/gpuistl/GpuOwnerOverlapCopy.hpp
+++ b/opm/simulators/linalg/gpuistl/GpuOwnerOverlapCopy.hpp
@@ -22,6 +22,7 @@
 #include <dune/istl/owneroverlapcopy.hh>
 
 #include <opm/simulators/linalg/gpuistl/GpuVector.hpp>
+#include <opm/simulators/linalg/FlowLinearSolverParameters.hpp>
 
 #include <mpi.h>
 
@@ -29,6 +30,13 @@
 #include <mutex>
 #include <vector>
 
+
+#ifdef OPEN_MPI
+#if OPEN_MPI
+#include "mpi-ext.h"
+#endif
+#endif
+
 namespace Opm::gpuistl {
 
 /**
@@ -424,6 +432,51 @@ private:
     std::shared_ptr<GPUSender<field_type, OwnerOverlapCopyCommunicationType>> m_sender;
 };
 
+template<class field_type, int block_size, class OwnerOverlapCopyCommunicationType>
+std::shared_ptr<GpuOwnerOverlapCopy<field_type, block_size, OwnerOverlapCopyCommunicationType>>
+makeGpuOwnerOverlapCopy(
+    const OwnerOverlapCopyCommunicationType& cpuOwnerOverlapCopy)
+{
+
+    const auto useGPUAwareMPI = Opm::Parameters::Get<Opm::Parameters::GPUAwareMPI>();
+    std::shared_ptr<Opm::gpuistl::GPUSender<field_type, OwnerOverlapCopyCommunicationType>> gpuComm;
+
+    if (useGPUAwareMPI) {
+        // Temporary solution use the GPU Direct communication solely based on these prepcrosessor statements
+        bool mpiSupportsCudaAwareAtCompileTime = false;
+        bool mpiSupportsCudaAwareAtRunTime = false;
+
+    #if defined(MPIX_CUDA_AWARE_SUPPORT) && MPIX_CUDA_AWARE_SUPPORT
+        mpiSupportsCudaAwareAtCompileTime = true;
+    #endif /* MPIX_CUDA_AWARE_SUPPORT */
+
+    #if defined(MPIX_CUDA_AWARE_SUPPORT)
+        if (1 == MPIX_Query_cuda_support()) {
+            mpiSupportsCudaAwareAtRunTime = true;
+        }
+    #endif /* MPIX_CUDA_AWARE_SUPPORT */
+        
+        if (!mpiSupportsCudaAwareAtCompileTime || !mpiSupportsCudaAwareAtRunTime) {
+            OPM_THROW(std::runtime_error,
+                      "The GPU-aware MPI support is not available. "
+                      "Please check your MPI installation and the OPM configuration "
+                      "or run with --gpu-aware-mpi=false");
+            }
+            gpuComm = std::make_shared<
+                Opm::gpuistl::GPUAwareMPISender<field_type, block_size, OwnerOverlapCopyCommunicationType>>(
+                cpuOwnerOverlapCopy);
+
+    } else {
+        gpuComm = std::make_shared<
+            Opm::gpuistl::GPUObliviousMPISender<field_type, block_size, OwnerOverlapCopyCommunicationType>>(
+            cpuOwnerOverlapCopy);
+    }
+    
+    using CudaCommunication = GpuOwnerOverlapCopy<field_type, block_size, OwnerOverlapCopyCommunicationType>;
+
+    return std::make_shared<CudaCommunication>(gpuComm);
+
+}
 } // namespace Opm::gpuistl
 
 #endif
diff --git a/opm/simulators/linalg/gpuistl/SolverAdapter.hpp b/opm/simulators/linalg/gpuistl/SolverAdapter.hpp
index 1d0fbdc790..b4cf832529 100644
--- a/opm/simulators/linalg/gpuistl/SolverAdapter.hpp
+++ b/opm/simulators/linalg/gpuistl/SolverAdapter.hpp
@@ -37,11 +37,6 @@
 #include <opm/simulators/linalg/gpuistl/GpuOwnerOverlapCopy.hpp>
 #endif
 
-#ifdef OPEN_MPI
-#if OPEN_MPI
-#include "mpi-ext.h"
-#endif
-#endif
 
 #include <memory>