NgenSimulation MPI interfacing adjustments

include/core/NgenSimulation.hpp

@@ -3,6 +3,10 @@
 
 #include <NGenConfig.h>
 
+#if NGEN_WITH_MPI
+#include <mpi.h>
+#endif
+
 #include <Simulation_Time.hpp>
 #include <Layer.hpp>
 
@@ -24,6 +28,30 @@ namespace routing_py_adapter {
 class NgenSimulation
 {
 public:
+#if NGEN_WITH_MPI
+    /**
+     * Construct from an MPI communicator, deriving the rank and process
+     * count from it.
+     *
+     * The communicator is duplicated (MPI_Comm_dup), giving the simulation a
+     * private communication context for its collective operations that is
+     * insulated from any other use of the caller's communicator. The
+     * duplicate is freed when the simulation is destroyed.
+     */
+    NgenSimulation(
+                   Simulation_Time const& sim_time,
+                   std::vector<std::shared_ptr<ngen::Layer>> layers,
+                   std::unordered_map<std::string, int> catchment_indexes,
+                   std::unordered_map<std::string, int> nexus_indexes,
+                   MPI_Comm comm
+                   );
+
+    /**
+     * Construct with an explicit MPI rank and process count, without owning a
+     * communicator (the internal communicator is left MPI_COMM_NULL). Routing
+     * collective operations are unavailable through such an instance; this is
+     * intended for single-process and test use where no communicator is needed.
+     */
     NgenSimulation(
                    Simulation_Time const& sim_time,
                    std::vector<std::shared_ptr<ngen::Layer>> layers,
@@ -32,6 +60,19 @@ class NgenSimulation
                    int mpi_rank,
                    int mpi_num_procs
                    );
+#else
+    /**
+     * Construct for a serial (non-MPI) run; the simulation behaves as rank 0
+     * of a single process.
+     */
+    NgenSimulation(
+                   Simulation_Time const& sim_time,
+                   std::vector<std::shared_ptr<ngen::Layer>> layers,
+                   std::unordered_map<std::string, int> catchment_indexes,
+                   std::unordered_map<std::string, int> nexus_indexes
+                   );
+#endif // NGEN_WITH_MPI
+
     NgenSimulation() = delete;
 
     ~NgenSimulation();
@@ -87,6 +128,12 @@ class NgenSimulation
 
     int mpi_rank_;
     int mpi_num_procs_;
+#if NGEN_WITH_MPI
+    // Private, duplicated communicator owned by this simulation. Set to
+    // MPI_COMM_NULL when constructed without a communicator (rank/size given
+    // explicitly), in which case the destructor leaves it alone.
+    MPI_Comm mpi_comm_;
+#endif
 
 #if NGEN_WITH_ROUTING
     std::unique_ptr<routing_py_adapter::Routing_Py_Adapter> router_;

src/NGen.cpp

@@ -132,7 +132,9 @@ int main(int argc, char* argv[]) {
     std::string PARTITION_PATH = "";
 
     // This default value should lead to behavior matching the single-process case in the standalone or non-MPI case
-    int mpi_num_procs = 1;
+    // Read only in MPI-guarded code; in non-MPI builds the serial NgenSimulation
+    // constructor takes no rank/size, so mark this to avoid an unused warning there.
+    [[maybe_unused]] int mpi_num_procs = 1;
     // Define in the non-MPI case so that we don't need to conditionally compile `if (mpi_rank == 0)`
     int mpi_rank = 0;
 
@@ -609,12 +611,18 @@ int main(int argc, char* argv[]) {
     }
 #endif // NGEN_WITH_ROUTING && NGEN_WITH_ROUTING_TROUTE_BMI
 
+#if NGEN_WITH_MPI
     auto simulation = std::make_unique<NgenSimulation>(*sim_time,
                                                        layers,
                                                        std::move(catchment_indexes),
                                                        std::move(nexus_indexes),
-                                                       mpi_rank,
-                                                       mpi_num_procs);
+                                                       MPI_COMM_WORLD);
+#else
+    auto simulation = std::make_unique<NgenSimulation>(*sim_time,
+                                                       layers,
+                                                       std::move(catchment_indexes),
+                                                       std::move(nexus_indexes));
+#endif
 
     auto time_done_init                             = std::chrono::steady_clock::now();
     std::chrono::duration<double> time_elapsed_init = time_done_init - time_start;
@@ -665,6 +673,9 @@ int main(int argc, char* argv[]) {
   manager->finalize();
 
 #if NGEN_WITH_MPI
+    // Destroy the simulation, freeing its duplicated MPI communicator, while
+    // MPI is still active.
+    simulation.reset();
     MPI_Finalize();
 #endif
 

src/core/NgenSimulation.cpp

@@ -14,6 +14,30 @@
 
 #include "parallel_utils.h"
 
+#if NGEN_WITH_MPI
+NgenSimulation::NgenSimulation(
+    Simulation_Time const& sim_time,
+    std::vector<std::shared_ptr<ngen::Layer>> layers,
+    std::unordered_map<std::string, int> catchment_indexes,
+    std::unordered_map<std::string, int> nexus_indexes,
+    MPI_Comm comm
+                               )
+    : simulation_step_(0)
+    , sim_time_(std::make_shared<Simulation_Time>(sim_time))
+    , layers_(std::move(layers))
+    , catchment_indexes_(std::move(catchment_indexes))
+    , nexus_indexes_(std::move(nexus_indexes))
+{
+    MPI_Comm_dup(comm, &mpi_comm_);
+    MPI_Comm_rank(mpi_comm_, &mpi_rank_);
+    MPI_Comm_size(mpi_comm_, &mpi_num_procs_);
+
+#if NGEN_WITH_ROUTING && NGEN_WITH_ROUTING_TROUTE_BMI
+    catchment_outflows_.reserve(catchment_indexes_.size() * get_num_output_times());
+    nexus_downstream_flows_.reserve(nexus_indexes_.size() * get_num_output_times());
+#endif
+}
+
 NgenSimulation::NgenSimulation(
     Simulation_Time const& sim_time,
     std::vector<std::shared_ptr<ngen::Layer>> layers,
@@ -30,13 +54,44 @@ NgenSimulation::NgenSimulation(
     , mpi_rank_(mpi_rank)
     , mpi_num_procs_(mpi_num_procs)
 {
+    // This constructor does not own an MPI communicator; mark it so the
+    // destructor leaves it alone.
+    mpi_comm_ = MPI_COMM_NULL;
 #if NGEN_WITH_ROUTING && NGEN_WITH_ROUTING_TROUTE_BMI
     catchment_outflows_.reserve(catchment_indexes_.size() * get_num_output_times());
     nexus_downstream_flows_.reserve(nexus_indexes_.size() * get_num_output_times());
 #endif
 }
+#else
+NgenSimulation::NgenSimulation(
+    Simulation_Time const& sim_time,
+    std::vector<std::shared_ptr<ngen::Layer>> layers,
+    std::unordered_map<std::string, int> catchment_indexes,
+    std::unordered_map<std::string, int> nexus_indexes
+                               )
+    : simulation_step_(0)
+    , sim_time_(std::make_shared<Simulation_Time>(sim_time))
+    , layers_(std::move(layers))
+    , catchment_indexes_(std::move(catchment_indexes))
+    , nexus_indexes_(std::move(nexus_indexes))
+    , mpi_rank_(0)
+    , mpi_num_procs_(1)
+{
+#if NGEN_WITH_ROUTING && NGEN_WITH_ROUTING_TROUTE_BMI
+    catchment_outflows_.reserve(catchment_indexes_.size() * get_num_output_times());
+    nexus_downstream_flows_.reserve(nexus_indexes_.size() * get_num_output_times());
+#endif
+}
+#endif // NGEN_WITH_MPI
 
-NgenSimulation::~NgenSimulation() = default;
+NgenSimulation::~NgenSimulation()
+{
+#if NGEN_WITH_MPI
+    if (mpi_comm_ != MPI_COMM_NULL) {
+        MPI_Comm_free(&mpi_comm_);
+    }
+#endif
+}
 
 void NgenSimulation::run_catchments()
 {
@@ -169,7 +224,7 @@ void NgenSimulation::run_routing_bmi(NgenSimulation::hy_features_t &features, st
             local_nexus_ids.push_back(nexus.first);
         }
         // MPI_Gather all nexus IDs into a single vector
-        std::vector<std::string> all_nexus_ids = parallel::gather_strings(local_nexus_ids, MPI_COMM_WORLD);
+        std::vector<std::string> all_nexus_ids = parallel::gather_strings(local_nexus_ids, mpi_comm_);
         if (mpi_rank_ == 0) {
             // filter to only the unique IDs
             std::sort(all_nexus_ids.begin(), all_nexus_ids.end());
@@ -179,7 +234,7 @@ void NgenSimulation::run_routing_bmi(NgenSimulation::hy_features_t &features, st
             );
         }
         // MPI_Broadcast so all processes share the nexus IDs
-        all_nexus_ids = std::move(parallel::broadcast_strings(all_nexus_ids, MPI_COMM_WORLD));
+        all_nexus_ids = std::move(parallel::broadcast_strings(all_nexus_ids, mpi_comm_));
 
         // MPI_Reduce to collect the results from processes
         if (mpi_rank_ == 0) {
@@ -200,7 +255,7 @@ void NgenSimulation::run_routing_bmi(NgenSimulation::hy_features_t &features, st
                 // if this process does not have the id, fill with 0 to make sure it doesn't affect reduce sum
                 std::fill(local_buffer.begin(), local_buffer.end(), 0.0);
             }
-            MPI_Reduce(local_buffer.data(), receive_buffer.data(), number_of_timesteps, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
+            MPI_Reduce(local_buffer.data(), receive_buffer.data(), number_of_timesteps, MPI_DOUBLE, MPI_SUM, 0, mpi_comm_);
             if (mpi_rank_ == 0) {
                 // copy reduce values to a combined downflows vector
                 all_nexus_indexes[nexus_id] = i;

test/core/NgenSimulationTests.cpp

@@ -10,9 +10,12 @@ TEST(NgenSimulation_Test, Construction)
     std::vector<std::shared_ptr<ngen::Layer>> layers;
     std::unordered_map<std::string, int> catchment_indexes;
     std::unordered_map<std::string, int> nexus_indexes;
+#if NGEN_WITH_MPI
     int mpi_rank = 0;
     int mpi_num_procs = 1;
 
+    // Use the explicit rank/process-count constructor, which owns no
+    // communicator and so needs no MPI runtime initialization.
     std::unique_ptr<NgenSimulation> simulation{
         new NgenSimulation(
                            sim_time,
@@ -22,4 +25,13 @@ TEST(NgenSimulation_Test, Construction)
                            mpi_rank,
                            mpi_num_procs
                            )};
+#else
+    std::unique_ptr<NgenSimulation> simulation{
+        new NgenSimulation(
+                           sim_time,
+                           layers,
+                           catchment_indexes,
+                           nexus_indexes
+                           )};
+#endif
 }