MPI refactor complete

paulromano · Feb 11, 2025 · Jan 29, 2025 · Jan 29, 2025 · Jan 30, 2025 · Jan 30, 2025
commit c8aad402142f1e9b2747e2a2aec0e672aaf28159
diff --git a/include/openmc/random_ray/flat_source_domain.h b/include/openmc/random_ray/flat_source_domain.h
@@ -37,7 +37,7 @@ class FlatSourceDomain {
   void random_ray_tally();
   virtual void accumulate_iteration_flux();
   void output_to_vtk() const;
-  virtual void all_reduce_replicated_source_regions();
+  void all_reduce_replicated_source_regions();
   void convert_external_sources();
   void count_external_source_regions();
   void set_adjoint_sources(const vector<double>& forward_flux);

diff --git a/include/openmc/random_ray/linear_source_domain.h b/include/openmc/random_ray/linear_source_domain.h
@@ -27,7 +27,6 @@ class LinearSourceDomain : public FlatSourceDomain {
   void batch_reset() override;
   void accumulate_iteration_flux() override;
   void output_to_vtk() const;
-  void all_reduce_replicated_source_regions() override;
   double evaluate_flux_at_point(Position r, int64_t sr, int g) const override;
 
 protected:

diff --git a/include/openmc/random_ray/source_region.h b/include/openmc/random_ray/source_region.h
@@ -318,6 +318,7 @@ class SourceRegionContainer {
   void push_back(const SourceRegion& sr);
   void assign(int n_source_regions, const SourceRegion& source_region);
   void flux_swap();
+  void mpi_sync_ranks(bool reduce_position);
 
   //----------------------------------------------------------------------------
   // Public Data Members

diff --git a/src/random_ray/flat_source_domain.cpp b/src/random_ray/flat_source_domain.cpp
@@ -658,7 +658,6 @@ void FlatSourceDomain::random_ray_tally()
 void FlatSourceDomain::all_reduce_replicated_source_regions()
 {
 #ifdef OPENMC_MPI
-
   // If we only have 1 MPI rank, no need
   // to reduce anything.
   if (mpi::n_procs <= 1)
@@ -671,74 +670,10 @@ void FlatSourceDomain::all_reduce_replicated_source_regions()
   int mapped_all_tallies_i = static_cast<int>(mapped_all_tallies_);
   MPI_Bcast(&mapped_all_tallies_i, 1, MPI_INT, 0, mpi::intracomm);
 
-  // Perform separate set of all reductions for each source region
-  for (int sr = 0; sr < n_source_regions_; sr++) {
-    SourceRegion& region = source_regions_[sr];
-
-    // The "position_recorded" variable needs to be allreduced (and maxed),
-    // as whether or not a cell was hit will affect some decisions in how the
-    // source is calculated in the next iteration so as to avoid dividing
-    // by zero. We take the max rather than the sum as the hit values are
-    // expected to be zero or 1.
-    MPI_Allreduce(MPI_IN_PLACE, &region.position_recorded_, 1, MPI_INT, MPI_MAX,
-      mpi::intracomm);
-
-    // The position variable is more complicated to reduce than the others,
-    // as we do not want the sum of all positions in each cell, rather, we
-    // want to just pick any single valid position. Thus, we perform a gather
-    // and then pick the first valid position we find for all source regions
-    // that have had a position recorded. This operation does not need to
-    // be broadcast back to other ranks, as this value is only used for the
-    // tally conversion operation, which is only performed on the master rank.
-    // While this is expensive, it only needs to be done for active batches,
-    // and only if we have not mapped all the tallies yet. Once tallies are
-    // fully mapped, then the position vector is fully populated, so this
-    // operation can be skipped.
-
-    // Then, we perform the gather of position data, if needed
-    if (simulation::current_batch > settings::n_inactive &&
-        !mapped_all_tallies_i) {
-
-      // Master rank will gather results and pick valid positions
-      if (mpi::master) {
-        // Initialize temporary vector for receiving positions
-        vector<Position> all_position;
-        all_position.resize(mpi::n_procs);
-
-        // Copy master rank data into gathered vector for convenience
-        all_position[0] = region.position_;
-
-        // Receive all data into gather vector
-        for (int i = 1; i < mpi::n_procs; i++) {
-          MPI_Recv(&all_position[i], 3, MPI_DOUBLE, i, 0, mpi::intracomm,
-            MPI_STATUS_IGNORE);
-        }
-
-        // Scan through gathered data and pick first valid cell posiiton
-        if (region.position_recorded_ == 1) {
-          for (int i = 0; i < mpi::n_procs; i++) {
-            if (all_position[i].x != 0.0 || all_position[i].y != 0.0 ||
-                all_position[i].z != 0.0) {
-              region.position_ = all_position[i];
-              break;
-            }
-          }
-        }
-      } else {
-        // Other ranks just send in their data
-        MPI_Send(&region.position_, 3, MPI_DOUBLE, 0, 0, mpi::intracomm);
-      }
-    }
-
-    // For the rest of the source region data, we simply perform an all reduce,
-    // as these values will be needed on all ranks for transport during the
-    // next iteration.
-    MPI_Allreduce(
-      MPI_IN_PLACE, &region.volume_, 1, MPI_DOUBLE, MPI_SUM, mpi::intracomm);
+  bool reduce_position =
+    simulation::current_batch > settings::n_inactive && !mapped_all_tallies_i;
 
-    MPI_Allreduce(MPI_IN_PLACE, region.scalar_flux_new_.data(), negroups_,
-      MPI_DOUBLE, MPI_SUM, mpi::intracomm);
-  }
+  source_regions_.mpi_sync_ranks(reduce_position);
 
   simulation::time_bank_sendrecv.stop();
 #endif

diff --git a/src/random_ray/linear_source_domain.cpp b/src/random_ray/linear_source_domain.cpp
@@ -158,40 +158,6 @@ void LinearSourceDomain::accumulate_iteration_flux()
   }
 }
 
-void LinearSourceDomain::all_reduce_replicated_source_regions()
-{
-#ifdef OPENMC_MPI
-  FlatSourceDomain::all_reduce_replicated_source_regions();
-  simulation::time_bank_sendrecv.start();
-
-  // We are going to assume we can safely cast Position, MomentArray,
-  // and MomentMatrix to contiguous arrays of doubles for the MPI
-  // allreduce operation. This is a safe assumption as typically
-  // compilers will at most pad to 8 byte boundaries. If a new FP32 MomentArray
-  // type is introduced, then there will likely be padding, in which case this
-  // function will need to become more complex.
-  if (sizeof(MomentArray) != 3 * sizeof(double) ||
-      sizeof(MomentMatrix) != 6 * sizeof(double)) {
-    fatal_error("Unexpected buffer padding in linear source domain reduction.");
-  }
-
-  // Perform separate all reductions for each source region
-  for (int sr = 0; sr < n_source_regions_; sr++) {
-    SourceRegion& region = source_regions_[sr];
-
-    MPI_Allreduce(MPI_IN_PLACE,
-      static_cast<void*>(region.flux_moments_new_.data()), negroups_ * 3,
-      MPI_DOUBLE, MPI_SUM, mpi::intracomm);
-    MPI_Allreduce(MPI_IN_PLACE, static_cast<void*>(&region.mom_matrix_), 6,
-      MPI_DOUBLE, MPI_SUM, mpi::intracomm);
-    MPI_Allreduce(MPI_IN_PLACE, static_cast<void*>(&region.centroid_iteration_),
-      3, MPI_DOUBLE, MPI_SUM, mpi::intracomm);
-  }
-
-  simulation::time_bank_sendrecv.stop();
-#endif
-}
-
 double LinearSourceDomain::evaluate_flux_at_point(
   Position r, int64_t sr, int g) const
 {

diff --git a/src/random_ray/source_region.cpp b/src/random_ray/source_region.cpp
@@ -1,5 +1,7 @@
 #include "openmc/random_ray/source_region.h"
 
+#include "openmc/message_passing.h"
+
 namespace openmc {
 
 //==============================================================================
@@ -131,4 +133,101 @@ void SourceRegionContainer::flux_swap()
   }
 }
 
+void SourceRegionContainer::mpi_sync_ranks(bool reduce_position)
+{
+#ifdef OPENMC_MPI
+
+  // The "position_recorded" variable needs to be allreduced (and maxed),
+  // as whether or not a cell was hit will affect some decisions in how the
+  // source is calculated in the next iteration so as to avoid dividing
+  // by zero. We take the max rather than the sum as the hit values are
+  // expected to be zero or 1.
+  MPI_Allreduce(MPI_IN_PLACE, position_recorded_.data(), 1, MPI_INT, MPI_MAX,
+    mpi::intracomm);
+
+  // The position variable is more complicated to reduce than the others,
+  // as we do not want the sum of all positions in each cell, rather, we
+  // want to just pick any single valid position. Thus, we perform a gather
+  // and then pick the first valid position we find for all source regions
+  // that have had a position recorded. This operation does not need to
+  // be broadcast back to other ranks, as this value is only used for the
+  // tally conversion operation, which is only performed on the master rank.
+  // While this is expensive, it only needs to be done for active batches,
+  // and only if we have not mapped all the tallies yet. Once tallies are
+  // fully mapped, then the position vector is fully populated, so this
+  // operation can be skipped.
+
+  // Then, we perform the gather of position data, if needed
+  if (reduce_position) {
+
+    // Master rank will gather results and pick valid positions
+    if (mpi::master) {
+      // Initialize temporary vector for receiving positions
+      vector<vector<Position>> all_position;
+      all_position.resize(mpi::n_procs);
+      for (int i = 0; i < mpi::n_procs; i++) {
+        all_position[i].resize(num_regions_);
+      }
+
+      // Copy master rank data into gathered vector for convenience
+      all_position[0] = position_;
+
+      // Receive all data into gather vector
+      for (int i = 1; i < mpi::n_procs; i++) {
+        MPI_Recv(all_position[i].data(), n_source_regions_ * 3, MPI_DOUBLE, i,
+          0, mpi::intracomm, MPI_STATUS_IGNORE);
+      }
+
+      // Scan through gathered data and pick first valid cell posiiton
+      for (int sr = 0; sr < num_regions_; sr++) {
+        if (position_recorded_[sr] == 1) {
+          for (int i = 0; i < mpi::n_procs; i++) {
+            if (all_position[i][sr].x != 0.0 || all_position[i][sr].y != 0.0 ||
+                all_position[i][sr].z != 0.0) {
+              position_[sr] = all_position[i][sr];
+              break;
+            }
+          }
+        }
+      }
+    } else {
+      // Other ranks just send in their data
+      MPI_Send(
+        position_.data(), num_regions_ * 3, MPI_DOUBLE, 0, 0, mpi::intracomm);
+    }
+  }
+
+  // For the rest of the source region data, we simply perform an all reduce,
+  // as these values will be needed on all ranks for transport during the
+  // next iteration.
+  MPI_Allreduce(MPI_IN_PLACE, volume_.data(), num_regions_, MPI_DOUBLE, MPI_SUM,
+    mpi::intracomm);
+
+  MPI_Allreduce(MPI_IN_PLACE, scalar_flux_new_.data(), num_regions_ * negroups_,
+    MPI_DOUBLE, MPI_SUM, mpi::intracomm);
+
+  if (is_linear_) {
+    // We are going to assume we can safely cast Position, MomentArray,
+    // and MomentMatrix to contiguous arrays of doubles for the MPI
+    // allreduce operation. This is a safe assumption as typically
+    // compilers will at most pad to 8 byte boundaries. If a new FP32
+    // MomentArray type is introduced, then there will likely be padding, in
+    // which case this function will need to become more complex.
+    if (sizeof(MomentArray) != 3 * sizeof(double) ||
+        sizeof(MomentMatrix) != 6 * sizeof(double)) {
+      fatal_error(
+        "Unexpected buffer padding in linear source domain reduction.");
+    }
+
+    MPI_Allreduce(MPI_IN_PLACE, static_cast<void*>(flux_moments_new_.data()),
+      n_source_elements_ * 3, MPI_DOUBLE, MPI_SUM, mpi::intracomm);
+    MPI_Allreduce(MPI_IN_PLACE, static_cast<void*>(mom_matrix_.data()),
+      n_source_regions_ * 6, MPI_DOUBLE, MPI_SUM, mpi::intracomm);
+    MPI_Allreduce(MPI_IN_PLACE, static_cast<void*>(centroid_iteration_.data()),
+      n_source_regions_ * 3, MPI_DOUBLE, MPI_SUM, mpi::intracomm);
+  }
+
+#endif
+}
+
 } // namespace openmc