From a649bef271ad21a7a6552f3d4a9e73d848ee9c9e Mon Sep 17 00:00:00 2001
From: Tobias Meyer Andersen <tobiasmeyer.andersen@sintef.no>
Date: Wed, 17 Apr 2024 11:37:25 +0200
Subject: [PATCH] update GPU direct support check

---
 .../linalg/cuistl/SolverAdapter.hpp           | 27 ++++++++++++-------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/opm/simulators/linalg/cuistl/SolverAdapter.hpp b/opm/simulators/linalg/cuistl/SolverAdapter.hpp
index 322d2e7e9..9d2c1db5c 100644
--- a/opm/simulators/linalg/cuistl/SolverAdapter.hpp
+++ b/opm/simulators/linalg/cuistl/SolverAdapter.hpp
@@ -33,6 +33,12 @@
 #include <opm/simulators/linalg/cuistl/PreconditionerAdapter.hpp>
 #include <opm/simulators/linalg/cuistl/detail/has_function.hpp>
 
+#ifdef OPEN_MPI
+#if OPEN_MPI
+#include "mpi-ext.h"
+#endif
+#endif
+
 namespace Opm::cuistl
 {
 //! @brief Wraps a CUDA solver to work with CPU data.
@@ -162,20 +168,23 @@ private:
             const auto& communication = m_opOnCPUWithMatrix.getCommunication();
 
             // Temporary solution use the GPU Direct communication solely based on these prepcrosessor statements
-            bool mpiMightBeSupportedDuringCompilation = true;
-            bool mpiMightBeSupportedDuringRuntime = true;
+            bool mpiSUpportsCudaAwareAtCompileTime = false;
+            bool mpiSupportsCudaAwareAtRunTime = false;
 
-            #if defined(MPIX_CUDA_AWARE_SUPPORT) && !MPIX_CUDA_AWARE_SUPPORT
-                mpiMightBeSupportedDuringCompilation = false;
-            #endif /* MPIX_CUDA_AWARE_SUPPORT */
+#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 defined(MPIX_CUDA_AWARE_SUPPORT) && !MPIX_Query_cuda_support
-                mpiMightBeSupportedDuringRuntime = false;
-            #endif /* MPIX_CUDA_AWARE_SUPPORT */
 
             // TODO add typename Operator communication type as a named type with using
             std::shared_ptr<Opm::cuistl::GPUSender<real_type, typename Operator::communication_type>> gpuComm;
-            if (mpiMightBeSupportedDuringCompilation && mpiMightBeSupportedDuringRuntime){
+            if (mpiSUpportsCudaAwareAtCompileTime && mpiSupportsCudaAwareAtRunTime){
                 gpuComm = std::make_shared<Opm::cuistl::GPUAwareMPISender<real_type, block_size, typename Operator::communication_type>>(communication);
             }
             else{