Simplify IFP message passing (#3719)

Co-authored-by: Paul Romano <[email protected]>

Author: GuySten

include/openmc/ifp.h

@@ -6,6 +6,8 @@
 #include "openmc/particle_data.h"
 #include "openmc/settings.h"
 
+#include <algorithm> // for copy
+
 namespace openmc {
 
 //! Check the value of the IFP parameter for beta effective or both.
@@ -113,14 +115,25 @@ void broadcast_ifp_n_generation(int& n_generation,
 //! \param[in] n_generation Number of generations
 //! \param[in] neighbor Index of the neighboring processor
 //! \param[in] requests MPI requests
-//! \param[in] delayed_groups List of delayed group numbers lists
-//! \param[out] send_delayed_groups Delayed group numbers buffer
-//! \param[in] lifetimes List of lifetimes lists
-//! \param[out] send_lifetimes Lifetimes buffer
+//! \param[in] data List of data lists
+//! \param[out] send_data data buffer
+template<typename T>
 void send_ifp_info(int64_t idx, int64_t n, int n_generation, int neighbor,
-  vector<MPI_Request>& requests, const vector<vector<int>>& delayed_groups,
-  vector<int>& send_delayed_groups, const vector<vector<double>>& lifetimes,
-  vector<double>& send_lifetimes);
+  vector<MPI_Request>& requests, const vector<vector<T>>& data,
+  vector<T>& send_data)
+{
+  // Copy data in buffer
+  for (int i = idx; i < idx + n; i++) {
+    std::copy(
+      data[i].begin(), data[i].end(), send_data.begin() + i * n_generation);
+  }
+
+  // Send data
+  requests.emplace_back();
+  MPI_Datatype datatype = mpi::MPITypeMap<T>::mpi_type;
+  MPI_Isend(&send_data[n_generation * idx], n_generation * static_cast<int>(n),
+    datatype, neighbor, mpi::rank, mpi::intracomm, &requests.back());
+}
 
 //! Receive IFP data using MPI.
 //!
@@ -129,12 +142,22 @@ void send_ifp_info(int64_t idx, int64_t n, int n_generation, int neighbor,
 //! \param[in] n_generation Number of generations
 //! \param[in] neighbor Index of the neighboring processor
 //! \param[in] requests MPI requests
-//! \param[in] delayed_groups List of delayed group numbers
-//! \param[in] lifetimes List of lifetimes
+//! \param[in] data data buffer
 //! \param[out] deserialization Information to deserialize the received data
+template<typename T>
 void receive_ifp_data(int64_t idx, int64_t n, int n_generation, int neighbor,
-  vector<MPI_Request>& requests, vector<int>& delayed_groups,
-  vector<double>& lifetimes, vector<DeserializationInfo>& deserialization);
+  vector<MPI_Request>& requests, vector<T>& data,
+  vector<DeserializationInfo>& deserialization)
+{
+  requests.emplace_back();
+  MPI_Datatype datatype = mpi::MPITypeMap<T>::mpi_type;
+  MPI_Irecv(&data[n_generation * idx], n_generation * static_cast<int>(n),
+    datatype, neighbor, neighbor, mpi::intracomm, &requests.back());
+
+  // Deserialization info to reconstruct data later
+  DeserializationInfo info = {idx, n};
+  deserialization.push_back(info);
+}
 
 //! Copy partial IFP data from local lists to source banks.
 //!
@@ -151,12 +174,24 @@ void copy_partial_ifp_data_to_source_banks(int64_t idx, int n, int64_t i_bank,
 //! the IFP source banks.
 //!
 //! \param[in] n_generation Number of generations
+//! \param[in] data data to deserialize
+//! \param[in] bank bank to store data
 //! \param[out] deserialization Information to deserialize the received data
-//! \param[in] delayed_groups List of delayed group numbers
-//! \param[in] lifetimes List of lifetimes
-void deserialize_ifp_info(int n_generation,
-  const vector<DeserializationInfo>& deserialization,
-  const vector<int>& delayed_groups, const vector<double>& lifetimes);
+template<typename T>
+void deserialize_ifp_info(int n_generation, const vector<T>& data,
+  vector<vector<T>>& bank, const vector<DeserializationInfo>& deserialization)
+{
+  for (auto info : deserialization) {
+    int64_t index_local = info.index_local;
+    int64_t n = info.n;
+
+    for (int i = index_local; i < index_local + n; i++) {
+      vector<T> data_received(
+        data.begin() + n_generation * i, data.begin() + n_generation * (i + 1));
+      bank[i] = data_received;
+    }
+  }
+}
 
 #endif
 

include/openmc/message_passing.h

@@ -22,6 +22,19 @@ extern MPI_Datatype collision_track_site;
 extern MPI_Comm intracomm;
 #endif
 
+//==============================================================================
+// Template struct used to map types to MPI datatypes
+// By having a single static data member, the template can
+// be specialized for each type we know of. The specializations appear in the
+// .cpp file since they are definitions.
+//==============================================================================
+#ifdef OPENMC_MPI
+template<typename T>
+struct MPITypeMap {
+  static const MPI_Datatype mpi_type;
+};
+#endif
+
 // Calculates global indices of the bank particles
 // across all ranks using a parallel scan. This is used to write
 // the surface source file in parallel runs. It will probably

src/eigenvalue.cpp

@@ -248,9 +248,12 @@ void synchronize_bank()
 
         if (settings::ifp_on) {
           // Send IFP data
-          send_ifp_info(index_local, n, ifp_n_generation, neighbor, requests,
-            temp_delayed_groups, send_delayed_groups, temp_lifetimes,
-            send_lifetimes);
+          if (is_beta_effective_or_both())
+            send_ifp_info(index_local, n, ifp_n_generation, neighbor, requests,
+              temp_delayed_groups, send_delayed_groups);
+          if (is_generation_time_or_both())
+            send_ifp_info(index_local, n, ifp_n_generation, neighbor, requests,
+              temp_lifetimes, send_lifetimes);
         }
       }
 
@@ -316,8 +319,12 @@ void synchronize_bank()
 
       if (settings::ifp_on) {
         // Receive IFP data
-        receive_ifp_data(index_local, n, ifp_n_generation, neighbor, requests,
-          recv_delayed_groups, recv_lifetimes, deserialization_info);
+        if (is_beta_effective_or_both())
+          receive_ifp_data(index_local, n, ifp_n_generation, neighbor, requests,
+            recv_delayed_groups, deserialization_info);
+        if (is_generation_time_or_both())
+          receive_ifp_data(index_local, n, ifp_n_generation, neighbor, requests,
+            recv_lifetimes, deserialization_info);
       }
 
     } else {
@@ -348,8 +355,12 @@ void synchronize_bank()
   MPI_Waitall(n_request, requests.data(), MPI_STATUSES_IGNORE);
 
   if (settings::ifp_on) {
-    deserialize_ifp_info(ifp_n_generation, deserialization_info,
-      recv_delayed_groups, recv_lifetimes);
+    if (is_beta_effective_or_both())
+      deserialize_ifp_info(ifp_n_generation, recv_delayed_groups,
+        simulation::ifp_source_delayed_group_bank, deserialization_info);
+    if (is_generation_time_or_both())
+      deserialize_ifp_info(ifp_n_generation, recv_lifetimes,
+        simulation::ifp_source_lifetime_bank, deserialization_info);
   }
 
 #else

src/ifp.cpp

@@ -63,7 +63,6 @@ void copy_ifp_data_from_fission_banks(
 }
 
 #ifdef OPENMC_MPI
-
 void broadcast_ifp_n_generation(int& n_generation,
   const vector<vector<int>>& delayed_groups,
   const vector<vector<double>>& lifetimes)
@@ -78,61 +77,6 @@ void broadcast_ifp_n_generation(int& n_generation,
   MPI_Bcast(&n_generation, 1, MPI_INT, 0, mpi::intracomm);
 }
 
-void send_ifp_info(int64_t idx, int64_t n, int n_generation, int neighbor,
-  vector<MPI_Request>& requests, const vector<vector<int>>& delayed_groups,
-  vector<int>& send_delayed_groups, const vector<vector<double>>& lifetimes,
-  vector<double>& send_lifetimes)
-{
-  // Copy data in send buffers
-  for (int i = idx; i < idx + n; i++) {
-    if (is_beta_effective_or_both()) {
-      std::copy(delayed_groups[i].begin(), delayed_groups[i].end(),
-        send_delayed_groups.begin() + i * n_generation);
-    }
-    if (is_generation_time_or_both()) {
-      std::copy(lifetimes[i].begin(), lifetimes[i].end(),
-        send_lifetimes.begin() + i * n_generation);
-    }
-  }
-  // Send delayed groups
-  if (is_beta_effective_or_both()) {
-    requests.emplace_back();
-    MPI_Isend(&send_delayed_groups[n_generation * idx],
-      n_generation * static_cast<int>(n), MPI_INT, neighbor, mpi::rank,
-      mpi::intracomm, &requests.back());
-  }
-  // Send lifetimes
-  if (is_generation_time_or_both()) {
-    requests.emplace_back();
-    MPI_Isend(&send_lifetimes[n_generation * idx],
-      n_generation * static_cast<int>(n), MPI_DOUBLE, neighbor, mpi::rank,
-      mpi::intracomm, &requests.back());
-  }
-}
-
-void receive_ifp_data(int64_t idx, int64_t n, int n_generation, int neighbor,
-  vector<MPI_Request>& requests, vector<int>& delayed_groups,
-  vector<double>& lifetimes, vector<DeserializationInfo>& deserialization)
-{
-  // Receive delayed groups
-  if (is_beta_effective_or_both()) {
-    requests.emplace_back();
-    MPI_Irecv(&delayed_groups[n_generation * idx],
-      n_generation * static_cast<int>(n), MPI_INT, neighbor, neighbor,
-      mpi::intracomm, &requests.back());
-  }
-  // Receive lifetimes
-  if (is_generation_time_or_both()) {
-    requests.emplace_back();
-    MPI_Irecv(&lifetimes[n_generation * idx],
-      n_generation * static_cast<int>(n), MPI_DOUBLE, neighbor, neighbor,
-      mpi::intracomm, &requests.back());
-  }
-  // Deserialization info to reconstruct data later
-  DeserializationInfo info = {idx, n};
-  deserialization.push_back(info);
-}
-
 void copy_partial_ifp_data_to_source_banks(int64_t idx, int n, int64_t i_bank,
   const vector<vector<int>>& delayed_groups,
   const vector<vector<double>>& lifetimes)
@@ -146,31 +90,6 @@ void copy_partial_ifp_data_to_source_banks(int64_t idx, int n, int64_t i_bank,
       &simulation::ifp_source_lifetime_bank[i_bank]);
   }
 }
-
-void deserialize_ifp_info(int n_generation,
-  const vector<DeserializationInfo>& deserialization,
-  const vector<int>& delayed_groups, const vector<double>& lifetimes)
-{
-  for (auto info : deserialization) {
-    int64_t index_local = info.index_local;
-    int64_t n = info.n;
-
-    for (int i = index_local; i < index_local + n; i++) {
-      if (is_beta_effective_or_both()) {
-        vector<int> delayed_groups_received(
-          delayed_groups.begin() + n_generation * i,
-          delayed_groups.begin() + n_generation * (i + 1));
-        simulation::ifp_source_delayed_group_bank[i] = delayed_groups_received;
-      }
-      if (is_generation_time_or_both()) {
-        vector<double> lifetimes_received(lifetimes.begin() + n_generation * i,
-          lifetimes.begin() + n_generation * (i + 1));
-        simulation::ifp_source_lifetime_bank[i] = lifetimes_received;
-      }
-    }
-  }
-}
-
 #endif
 
 void copy_complete_ifp_data_to_source_banks(

src/message_passing.cpp

@@ -39,6 +39,14 @@ vector<int64_t> calculate_parallel_index_vector(int64_t size)
   return result;
 }
 
+#ifdef OPENMC_MPI
+// Specializations of the MPITypeMap template struct
+template<>
+const MPI_Datatype MPITypeMap<int>::mpi_type = MPI_INT;
+template<>
+const MPI_Datatype MPITypeMap<double>::mpi_type = MPI_DOUBLE;
+#endif
+
 } // namespace mpi
 
 } // namespace openmc