Enabling GPU direct communication for Cray-MPICH

quest/include/environment.h

@@ -36,6 +36,7 @@ typedef struct {
     int isMultithreaded;
     int isGpuAccelerated;
     int isDistributed;
+    int isMpiGpuAware;
 
     // deployment modes which cannot be directly changed after compilation
     int isCuQuantumEnabled;

quest/src/api/environment.cpp

@@ -147,6 +147,7 @@ void validateAndInitCustomQuESTEnv(int useDistrib, int useGpuAccel, int useMulti
     globalEnvPtr->isDistributed       = useDistrib;
     globalEnvPtr->isCuQuantumEnabled  = useCuQuantum;
     globalEnvPtr->isGpuSharingEnabled = permitGpuSharing;
+    globalEnvPtr->isMPIGPUAware = comm_set_isMpiGpuAware();
 
     // bind distributed info
     globalEnvPtr->rank     = (useDistrib)? comm_getRank()     : 0;

quest/src/comm/comm_config.cpp

@@ -12,12 +12,17 @@
  * @author Tyson Jones
  */
 
+#include "quest/include/environment.h"
 #include "quest/include/config.h"
 #include "quest/include/types.h"
 
 #include "quest/src/comm/comm_config.hpp"
 #include "quest/src/core/errors.hpp"
 
+#include <cstdlib>
+#include <string>
+#include <iostream>
+
 #if COMPILE_MPI
     #include <mpi.h>
 #endif
@@ -60,23 +65,71 @@ bool comm_isMpiCompiled() {
     return (bool) COMPILE_MPI;
 }
 
+enum Mpi_version {NONE, OPENMPI, CRAYMPICH};
 
-bool comm_isMpiGpuAware() {
+int comm_whichMpi() {
 
-    /// @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!
+    #ifdef COMPILE_MPI
+        char version_string[MPI_MAX_LIBRARY_VERSION_STRING];
+    #else
+        char version_string[];
+    #endif
 
-    // definitely not GPU-aware if compiler declares it is not
-    #if defined(MPIX_CUDA_AWARE_SUPPORT) && ! MPIX_CUDA_AWARE_SUPPORT
-        return false;
+    int resultlen[] = {0};
+
+    #ifdef COMPILE_MPI
+        MPI_Get_library_version(version_string, resultlen);
     #endif
+    enum Mpi_version version = NONE;
 
-    // check CUDA-awareness at run-time if we know it's principally supported
-    #if defined(MPIX_CUDA_AWARE_SUPPORT)
-        return (bool) MPIX_Query_cuda_support();
+    // Check if Openmpi used
+    #ifdef OPEN_MPI
+        version = OPENMPI;
     #endif
 
+    // Check if Cray MPI used 
+
+    const char* cray_string = "CRAY MPICH";
+
+    std::string v_string = version_string;
+
+    if (v_string.find(cray_string) != string::npos) {
+        version = CRAYMPICH;
+    }
+
+    return version;
+
+}
+
+
+
+bool comm_isMpiGpuAware() {
+	return getQuESTEnv().isMPIGPUAware;
+}
+
+bool comm_set_isMpiGpuAware() {
+
+   int mpi_lib = comm_whichMpi();
+
+    if(OPENMPI==mpi_lib) {
+        // definitely not GPU-aware if compiler declares it is not
+        #if defined(MPIX_CUDA_AWARE_SUPPORT) && ! MPIX_CUDA_AWARE_SUPPORT
+            return false;
+        #endif
+
+        // check CUDA-awareness at run-time if we know it's principally supported
+        #if defined(MPIX_CUDA_AWARE_SUPPORT)
+            return (bool) MPIX_Query_cuda_support();
+	#endif
+    }
+
+    if(CRAYMPICH==mpi_lib) { 
+
+       const char* var = std::getenv("MPICH_GPU_SUPPORT_ENABLED");
+
+       return (bool) var;
+    }
+
     // if we can't ascertain CUDA-awareness, just assume no to avoid seg-fault
     return false;
 }

quest/src/comm/comm_config.hpp

@@ -13,8 +13,11 @@
 
 constexpr int ROOT_RANK = 0;
 
+int comm_whichMpi();
+
 bool comm_isMpiCompiled();
 bool comm_isMpiGpuAware();
+bool comm_set_isMpiGpuAware();
 
 void comm_init();
 void comm_end();
@@ -28,4 +31,4 @@ bool comm_isRootNode();
 bool comm_isRootNode(int rank);
 
 
-#endif // COMM_CONFIG_HPP
+#endif // COMM_CONFIG_HPP