Enable user to take ownership of MPI

CMakeLists.txt

@@ -145,6 +145,12 @@ option(
 )
 message(STATUS "Distribution is turned ${ENABLE_DISTRIBUTION}. Set ENABLE_DISTRIBUTION to modify.")
 
+option(
+  ENABLE_SUBCOMM
+  "Whether QuEST will be built with support for restricting it to a user-defined MPI communicator. Turned OFF by default."
+  OFF
+)
+message(STATUS "Custom communicator support is turned ${ENABLE_SUBCOMM}. Set ENABLE_SUBCOMM to modify.")
 
 # GPU Acceleration
 option(
@@ -207,6 +213,9 @@ if (ENABLE_CUQUANTUM AND NOT ENABLE_CUDA)
   message(FATAL_ERROR "Use of cuQuantum requires CUDA.")
 endif()
 
+if (ENABLE_SUBCOMM AND NOT ENABLE_DISTRIBUTION)
+  message(FATAL_ERROR "Distribution must be enabled to make use of a user-defined communicator for QuEST.")
+endif()
 
 if(WIN32)
 
@@ -402,8 +411,11 @@ endif()
 # MPI
 if (ENABLE_DISTRIBUTION)
   find_package(MPI REQUIRED
+    # Component CXX is the C api usable from C++
+    # NOT the deprecated C++ API
     COMPONENTS CXX
   )
+
   target_link_libraries(QuEST
     PRIVATE
     MPI::MPI_CXX
@@ -476,6 +488,7 @@ endif()
 # set vars which will be written to config.h.in (auto-converted to 0 or 1)
 set(COMPILE_OPENMP ${ENABLE_MULTITHREADING})
 set(COMPILE_MPI ${ENABLE_DISTRIBUTION})
+set(COMPILE_SUBCOMM ${ENABLE_SUBCOMM})
 set(COMPILE_CUQUANTUM ${ENABLE_CUQUANTUM})
 set(INCLUDE_DEPRECATED_FUNCTIONS ${ENABLE_DEPRECATED_API})
 
@@ -530,6 +543,10 @@ add_executable(min_example
 )
 target_link_libraries(min_example PRIVATE QuEST::QuEST)
 
+if (ENABLE_DISTRIBUTION AND ENABLE_SUBCOMM)
+  target_link_libraries(min_example PRIVATE MPI::MPI_CXX)
+endif()
+
 if (INSTALL_BINARIES)
   install(TARGETS min_example
     RUNTIME

examples/CMakeLists.txt

@@ -20,6 +20,10 @@ function(add_example direc in_fn)
   add_executable(${target} ${in_fn})
   target_link_libraries(${target} PUBLIC QuEST)
 
+  if (ENABLE_DISTRIBUTION AND ENABLE_SUBCOMM)
+    target_link_libraries(${target} PRIVATE MPI::MPI_CXX)
+  endif()
+
   if (INSTALL_BINARIES)
     install(
       TARGETS ${target}

quest/include/config.h.in

@@ -37,6 +37,7 @@
 #if defined(FLOAT_PRECISION)              || \
     defined(COMPILE_OPENMP)               || \
     defined(COMPILE_MPI)                  || \
+    defined(COMPILE_SUBCOMM)              || \
     defined(COMPILE_CUDA)                 || \
     defined(COMPILE_HIP)                  || \
     defined(COMPILE_CUQUANTUM)            || \
@@ -79,6 +80,7 @@
 // crucial to QuEST source (informs external library usage)
 #cmakedefine01 COMPILE_OPENMP
 #cmakedefine01 COMPILE_MPI
+#cmakedefine01 COMPILE_SUBCOMM
 #cmakedefine01 COMPILE_CUDA
 #cmakedefine01 COMPILE_CUQUANTUM
 
@@ -118,6 +120,7 @@
 #if ! defined(FLOAT_PRECISION)              || \
     ! defined(COMPILE_OPENMP)               || \
     ! defined(COMPILE_MPI)                  || \
+    ! defined(COMPILE_SUBCOMM)              || \
     ! defined(COMPILE_CUDA)                 || \
     ! defined(COMPILE_HIP)                  || \
     ! defined(COMPILE_CUQUANTUM)            || \
@@ -144,6 +147,7 @@
 
 #if ! (COMPILE_OPENMP               == 0 || COMPILE_OPENMP               == 1) || \
     ! (COMPILE_MPI                  == 0 || COMPILE_MPI                  == 1) || \
+    ! (COMPILE_SUBCOMM             == 0  || COMPILE_SUBCOMM              == 1) || \
     ! (COMPILE_CUDA                 == 0 || COMPILE_CUDA                 == 1) || \
     ! (COMPILE_HIP                  == 0 || COMPILE_HIP                  == 1) || \
     ! (COMPILE_CUQUANTUM            == 0 || COMPILE_CUQUANTUM            == 1) || \

quest/include/environment.h

@@ -14,6 +14,8 @@
 #ifndef ENVIRONMENT_H
 #define ENVIRONMENT_H
 
+#include <stdbool.h>
+
 // enable invocation by both C and C++ binaries
 #ifdef __cplusplus
 extern "C" {
@@ -36,6 +38,7 @@ typedef struct {
     int isMultithreaded;
     int isGpuAccelerated;
     int isDistributed;
+    bool userOwnsMpi;
 
     // deployment modes which cannot be directly changed after compilation
     int isCuQuantumEnabled;
@@ -61,6 +64,8 @@ void initQuESTEnv();
  */
 void initCustomQuESTEnv(int useDistrib, int useGpuAccel, int useMultithread);
 
+void initCustomMpiQuESTEnv(int useDistrib, bool userOwnsMpi, int useGpuAccel, int useMultithread);
+
 /// @notyetdoced
 void finalizeQuESTEnv();
 

quest/include/quest.h

@@ -45,6 +45,7 @@
 #include "quest/include/operations.h"
 #include "quest/include/paulis.h"
 #include "quest/include/qureg.h"
+#include "quest/include/subcommunicator.h"
 #include "quest/include/matrices.h"
 #include "quest/include/wrappers.h"
 

quest/include/subcommunicator.h

@@ -0,0 +1,22 @@
+#ifndef SUBCOMMUNICATOR_H
+#define SUBCOMMUNICATOR_H
+
+#include "quest/include/config.h" 
+
+#if COMPILE_MPI && COMPILE_SUBCOMM
+
+#include <mpi.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+void initCustomMpiCommQuESTEnv(MPI_Comm questComm, int useGpuAccel, int useMultithread);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif
+
+#endif

quest/src/api/CMakeLists.txt

@@ -12,6 +12,7 @@ target_sources(QuEST
   operations.cpp
   paulis.cpp
   qureg.cpp
+  subcommunicator.cpp
   trotterisation.cpp
   types.cpp
-)
+)

quest/src/api/environment.cpp

@@ -71,7 +71,7 @@ static bool hasEnvBeenFinalized = false;
  */
 
 
-void validateAndInitCustomQuESTEnv(int useDistrib, int useGpuAccel, int useMultithread, const char* caller) {
+void validateAndInitCustomQuESTEnv(int useDistrib, bool userOwnsMpi, int useGpuAccel, int useMultithread, const char* caller) {
 
     // ensure that we are never re-initialising QuEST (even after finalize) because
     // this leads to undefined behaviour in distributed mode, as per the MPI
@@ -94,7 +94,7 @@ void validateAndInitCustomQuESTEnv(int useDistrib, int useGpuAccel, int useMulti
     // perform that specifically upon the MPI-process-bound GPU(s). Further,
     // we can make sure validation errors are reported only by the root node.
     if (useDistrib)
-        comm_init();
+        comm_init(userOwnsMpi);
 
     validate_newEnvDistributedBetweenPower2Nodes(caller);
 
@@ -145,6 +145,7 @@ void validateAndInitCustomQuESTEnv(int useDistrib, int useGpuAccel, int useMulti
     globalEnvPtr->isMultithreaded     = useMultithread;
     globalEnvPtr->isGpuAccelerated    = useGpuAccel;
     globalEnvPtr->isDistributed       = useDistrib;
+    globalEnvPtr->userOwnsMpi         = userOwnsMpi;
     globalEnvPtr->isCuQuantumEnabled  = useCuQuantum;
     globalEnvPtr->isGpuSharingEnabled = permitGpuSharing;
 
@@ -153,6 +154,11 @@ void validateAndInitCustomQuESTEnv(int useDistrib, int useGpuAccel, int useMulti
     globalEnvPtr->numNodes = (useDistrib)? comm_getNumNodes() : 1;
 }
 
+void updateQuESTEnvDistInfo() {
+    globalEnvPtr->rank     = (globalEnvPtr->isDistributed)? comm_getRank()     : 0;
+    globalEnvPtr->numNodes = (globalEnvPtr->isDistributed)? comm_getNumNodes() : 1;
+    return;
+}
 
 
 /*
@@ -187,10 +193,11 @@ void printCompilationInfo() {
 
     print_table(
         "compilation", {
-        {"isMpiCompiled",       comm_isMpiCompiled()},
-        {"isGpuCompiled",       gpu_isGpuCompiled()},
-        {"isOmpCompiled",       cpu_isOpenmpCompiled()},
-        {"isCuQuantumCompiled", gpu_isCuQuantumCompiled()},
+        {"isMpiCompiled",                comm_isMpiCompiled()},
+        {"isMpiSubCommunicatorCompiled", comm_isMpiSubCommunicatorCompiled()},
+        {"isGpuCompiled",                gpu_isGpuCompiled()},
+        {"isOmpCompiled",                cpu_isOpenmpCompiled()},
+        {"isCuQuantumCompiled",          gpu_isCuQuantumCompiled()},
     });
 }
 
@@ -200,6 +207,7 @@ void printDeploymentInfo() {
     print_table(
         "deployment", {
         {"isMpiEnabled",        globalEnvPtr->isDistributed},
+        {"doesUserOwnMpi",      globalEnvPtr->userOwnsMpi},
         {"isGpuEnabled",        globalEnvPtr->isGpuAccelerated},
         {"isOmpEnabled",        globalEnvPtr->isMultithreaded},
         {"isCuQuantumEnabled",  globalEnvPtr->isCuQuantumEnabled},
@@ -397,13 +405,19 @@ extern "C" {
 
 void initCustomQuESTEnv(int useDistrib, int useGpuAccel, int useMultithread) {
 
-    validateAndInitCustomQuESTEnv(useDistrib, useGpuAccel, useMultithread, __func__);
+    const int USER_OWNS_MPI = 0;
+    validateAndInitCustomQuESTEnv(useDistrib, USER_OWNS_MPI, useGpuAccel, useMultithread, __func__);
 }
 
 
+void initCustomMpiQuESTEnv(int useDistrib, bool userOwnsMpi, int useGpuAccel, int useMultithread) {
+    validateAndInitCustomQuESTEnv(useDistrib, userOwnsMpi, useGpuAccel, useMultithread, __func__);
+}
+
 void initQuESTEnv() {
 
-    validateAndInitCustomQuESTEnv(modeflag::USE_AUTO, modeflag::USE_AUTO, modeflag::USE_AUTO, __func__);
+    const int USER_OWNS_MPI = 0;
+    validateAndInitCustomQuESTEnv(modeflag::USE_AUTO, USER_OWNS_MPI, modeflag::USE_AUTO, modeflag::USE_AUTO, __func__);
 }
 
 
@@ -436,7 +450,7 @@ void finalizeQuESTEnv() {
 
     if (globalEnvPtr->isDistributed) {
         comm_sync();
-        comm_end();
+        comm_end(globalEnvPtr->userOwnsMpi);
     }
 
     // free global env's heap memory and flag it as unallocated
@@ -454,8 +468,12 @@ void syncQuESTEnv() {
     if (globalEnvPtr->isGpuAccelerated)
         gpu_sync();
 
-    if (globalEnvPtr->isDistributed)
+    if (globalEnvPtr->isDistributed) {
         comm_sync();
+        #if COMPILE_SUBCOMM
+            updateQuESTEnvDistInfo();
+        #endif
+    }
 }
 
 
@@ -498,10 +516,11 @@ void getEnvironmentString(char str[200]) {
     int cuQuantum = env.isGpuAccelerated && gpu_isCuQuantumCompiled();
     int gpuDirect = env.isGpuAccelerated && gpu_isDirectGpuCommPossible();
 
-    snprintf(str, 200, "CUDA=%d OpenMP=%d MPI=%d threads=%d ranks=%d cuQuantum=%d gpuDirect=%d",
+    snprintf(str, 200, "CUDA=%d OpenMP=%d MPI=%d userOwnsMPI=%d threads=%d ranks=%d cuQuantum=%d gpuDirect=%d",
         env.isGpuAccelerated,
         env.isMultithreaded,
         env.isDistributed,
+        env.userOwnsMpi,
         numThreads,
         env.numNodes,
         cuQuantum,

quest/src/api/subcommunicator.cpp

@@ -0,0 +1,25 @@
+#include "quest/include/config.h"
+#include "quest/include/environment.h"
+#include "quest/include/subcommunicator.h"
+
+#include "quest/src/comm/comm_config.hpp"
+
+#if COMPILE_MPI && COMPILE_SUBCOMM
+
+#include <mpi.h>
+
+void initCustomMpiCommQuESTEnv(MPI_Comm userQuestComm, int useGpuAccel, int useMultithread) {
+    // useDistrib and userOwnsMpi are implied by the user of this initialiser
+    const int useDistrib = 1;
+    const int userOwnsMpi = 1;
+
+    // set mpiCommQuest to user provided communicator
+    comm_setMpiComm(userQuestComm);
+
+    // initialise QuEST around that communicator
+    initCustomMpiQuESTEnv(useDistrib, userOwnsMpi, useGpuAccel, useMultithread);
+
+    return;
+}
+
+#endif