From f73126a826e4eefab85b78188348184915a41dea Mon Sep 17 00:00:00 2001
From: Tyson Jones <tyson.jones.input@gmail.com>
Date: Sun, 31 May 2026 03:23:50 -0400
Subject: [PATCH] Add GPU-aware MPICH detection

supersedes #724

Co-Authored-By: JPRichings <james.richings@ed.ac.uk>
---
 quest/include/environment.h    |  1 +
 quest/src/api/environment.cpp  |  8 +++++++-
 quest/src/comm/comm_config.cpp | 15 ++++++++++++---
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/quest/include/environment.h b/quest/include/environment.h
index 608829912..15a6ac5e2 100644
--- a/quest/include/environment.h
+++ b/quest/include/environment.h
@@ -45,6 +45,7 @@ typedef struct {
 
     // deployment configurations which can be changed via environment variables
     int isGpuSharingEnabled;
+    int isMpiGpuAware;
 
     // distributed configuration
     int rank;
diff --git a/quest/src/api/environment.cpp b/quest/src/api/environment.cpp
index 6c72b14c0..b0eb38c77 100644
--- a/quest/src/api/environment.cpp
+++ b/quest/src/api/environment.cpp
@@ -130,6 +130,10 @@ void validateAndInitCustomQuESTEnv(int useDistrib, bool userOwnsMpi, int useGpuA
         gpu_initCuQuantum();
     }
 
+    // MPI GPU-awareness detection is platform specific; sometimes it is
+    // known at compile-time, other times according to env-vars
+    bool isMpiGpuAware = comm_isMpiGpuAware();
+
     // initialise RNG, used by measurements and random-state generation
     rand_setSeedsToDefault();
 
@@ -147,6 +151,7 @@ void validateAndInitCustomQuESTEnv(int useDistrib, bool userOwnsMpi, int useGpuA
     globalEnvPtr->userOwnsMpi         = userOwnsMpi;
     globalEnvPtr->isCuQuantumEnabled  = useCuQuantum;
     globalEnvPtr->isGpuSharingEnabled = permitGpuSharing;
+    globalEnvPtr->isMpiGpuAware       = isMpiGpuAware;
 
     // bind distributed info
     globalEnvPtr->rank     = (useDistrib)? comm_getRank()     : 0;
@@ -211,6 +216,7 @@ void printDeploymentInfo() {
         {"isOmpEnabled",        globalEnvPtr->isMultithreaded},
         {"isCuQuantumEnabled",  globalEnvPtr->isCuQuantumEnabled},
         {"isGpuSharingEnabled", globalEnvPtr->isGpuSharingEnabled},
+        {"isMpiGpuAware",       globalEnvPtr->isMpiGpuAware},
     });
 }
 
@@ -269,7 +275,7 @@ void printDistributionInfo() {
 
     print_table(
         "distribution", {
-        {"isMpiGpuAware", (comm_isMpiCompiled())? printer_toStr(comm_isMpiGpuAware()) : na},
+        {"isMpiGpuAware", comm_isInit()? printer_toStr(globalEnvPtr->isMpiGpuAware) : na},
         {"numMpiNodes",   printer_toStr(globalEnvPtr->numNodes)},
     });
 }
diff --git a/quest/src/comm/comm_config.cpp b/quest/src/comm/comm_config.cpp
index 9da8f34e1..67fd53231 100644
--- a/quest/src/comm/comm_config.cpp
+++ b/quest/src/comm/comm_config.cpp
@@ -18,6 +18,8 @@
 #include "quest/src/comm/comm_config.hpp"
 #include "quest/src/core/errors.hpp"
 
+#include <string>
+
 #if QUEST_COMPILE_MPI
     #include <mpi.h>
 
@@ -62,15 +64,17 @@ bool comm_isMpiCompiled() {
     return (bool) QUEST_COMPILE_MPI;
 }
 
+
 bool comm_isMpiSubCommunicatorCompiled() {
     return (bool) QUEST_COMPILE_SUBCOMM;
 }
 
+
 bool comm_isMpiGpuAware() {
 
-    /// @todo these checks may be OpenMPI specific, so that
-    /// non-OpenMPI MPI compilers are always dismissed as
-    /// not being CUDA-aware. Check e.g. MPICH method!
+    // well duh
+    if (!comm_isMpiCompiled())
+        return false;
 
     // definitely not GPU-aware if compiler declares it is not
     #if defined(MPIX_CUDA_AWARE_SUPPORT) && ! MPIX_CUDA_AWARE_SUPPORT
@@ -82,6 +86,11 @@ bool comm_isMpiGpuAware() {
         return (bool) MPIX_Query_cuda_support();
     #endif
 
+    // check whether an MPICH env-var indicates support (we assume it never lies!)
+    static const auto var = std::getenv("MPICH_GPU_SUPPORT_ENABLED");
+    if (var && std::string(var) == "1") // ill-formed vars = 0
+        return true;
+
     // if we can't ascertain CUDA-awareness, just assume no to avoid seg-fault
     return false;
 }