Performance optimizations for shared secondary bank (#4011)

include/openmc/shared_array.h

@@ -114,6 +114,32 @@ public:
     data_[size_++] = value;
   }
 
+  //! Increase the size of the container by count elements without assigning
+  //! values to the new elements. Existing elements are preserved if the
+  //! container needs to grow. This does not perform any thread safety checks.
+  //
+  //! \param count The number of elements to append
+  //! \return The starting index of the appended range
+  int64_t extend_uninitialized(int64_t count)
+  {
+    int64_t offset = size_;
+    int64_t new_size = size_ + count;
+    if (new_size > capacity_) {
+      int64_t new_capacity = capacity_ == 0 ? 8 : capacity_;
+      while (new_capacity < new_size) {
+        new_capacity *= 2;
+      }
+      unique_ptr<T[]> new_data = make_unique<T[]>(new_capacity);
+      if (size_ > 0) {
+        std::copy_n(data_.get(), size_, new_data.get());
+      }
+      data_ = std::move(new_data);
+      capacity_ = new_capacity;
+    }
+    size_ = new_size;
+    return offset;
+  }
+
   //! Return the number of elements in the container
   int64_t size() { return size_; }
   int64_t size() const { return size_; }

src/simulation.cpp

@@ -12,6 +12,7 @@
 #include "openmc/material.h"
 #include "openmc/message_passing.h"
 #include "openmc/nuclide.h"
+#include "openmc/openmp_interface.h"
 #include "openmc/output.h"
 #include "openmc/particle.h"
 #include "openmc/photon.h"
@@ -41,6 +42,7 @@
 
 #include <algorithm>
 #include <cmath>
+#include <numeric>
 #include <string>
 
 //==============================================================================
@@ -354,6 +356,94 @@ int64_t simulation_tracks_completed {0};
 
 } // namespace simulation
 
+namespace {
+
+//! Collect thread-local secondary banks into the shared secondary bank in
+//! sorted order.
+//!
+//! \param thread_banks  Secondary banks produced by each OpenMP thread
+void collect_sorted_history_secondary_banks(
+  vector<vector<SourceSite>>& thread_banks)
+{
+  // Count the total number of all secondary sites produced
+  int64_t n_collected = 0;
+  for (const auto& bank : thread_banks) {
+    n_collected += bank.size();
+  }
+
+  // Count the expected number of progeny from per-parent progeny counts
+  int64_t n_progeny = 0;
+  for (int64_t count : simulation::progeny_per_particle) {
+    n_progeny += count;
+  }
+
+  if (n_collected != n_progeny) {
+    fatal_error("Mismatch detected between sum of all particle progeny and "
+                "secondary bank size during collection.");
+  }
+
+  // Convert per-parent progeny counts to offsets into the sorted bank
+  std::exclusive_scan(simulation::progeny_per_particle.begin(),
+    simulation::progeny_per_particle.end(),
+    simulation::progeny_per_particle.begin(), 0);
+
+  // Allocate the shared bank once for the complete generation
+  simulation::shared_secondary_bank_write.resize(0);
+  simulation::shared_secondary_bank_write.extend_uninitialized(n_progeny);
+
+  // Place each secondary according to its parent and progeny identifiers
+  for (const auto& bank : thread_banks) {
+    for (const auto& site : bank) {
+      if (site.parent_id < 0 ||
+          site.parent_id >=
+            static_cast<int64_t>(simulation::progeny_per_particle.size())) {
+        fatal_error(fmt::format("Invalid parent_id {} for banked site "
+                                "(expected range [0, {})).",
+          site.parent_id, simulation::progeny_per_particle.size()));
+      }
+      int64_t idx =
+        simulation::progeny_per_particle[site.parent_id] + site.progeny_id;
+      if (idx < 0 || idx >= n_progeny) {
+        fatal_error("Mismatch detected between sum of all particle progeny and "
+                    "secondary bank size during collection.");
+      }
+      simulation::shared_secondary_bank_write[idx] = site;
+    }
+  }
+}
+
+//! Collect particle-local secondary banks into the shared secondary bank.
+//!
+//! \param n_particles  Number of particles in the active event-based buffer
+void collect_event_secondary_banks(int64_t n_particles)
+{
+  // Compute offsets for each particle's local secondary bank.
+  vector<int64_t> offsets(n_particles);
+  int64_t total = 0;
+  for (int64_t i = 0; i < n_particles; ++i) {
+    offsets[i] = total;
+    total += simulation::particles[i].local_secondary_bank().size();
+  }
+
+  // Extend the shared bank once for all collected secondaries
+  int64_t bank_offset =
+    simulation::shared_secondary_bank_write.extend_uninitialized(total);
+
+  // Copy each local bank into its assigned range and clear the local storage
+#pragma omp parallel for schedule(static)
+  for (int64_t i = 0; i < n_particles; ++i) {
+    auto& local_bank = simulation::particles[i].local_secondary_bank();
+    if (!local_bank.empty()) {
+      std::copy(local_bank.cbegin(), local_bank.cend(),
+        simulation::shared_secondary_bank_write.data() + bank_offset +
+          offsets[i]);
+      local_bank.clear();
+    }
+  }
+}
+
+} // namespace
+
 //==============================================================================
 // Non-member functions
 //==============================================================================
@@ -921,11 +1011,13 @@ void transport_history_based_shared_secondary()
   std::fill(simulation::progeny_per_particle.begin(),
     simulation::progeny_per_particle.end(), 0);
 
+  vector<vector<SourceSite>> thread_banks(num_threads());
+
   // Phase 1: Transport primary particles and deposit first generation of
   // secondaries in the shared secondary bank
 #pragma omp parallel
   {
-    vector<SourceSite> thread_bank;
+    auto& thread_bank = thread_banks[thread_num()];
     Particle p;
 
 #pragma omp for schedule(runtime)
@@ -937,15 +1029,9 @@ void transport_history_based_shared_secondary()
       }
       p.local_secondary_bank().clear();
     }
-
-    // Drain thread-local bank into the shared secondary bank (once per thread)
-#pragma omp critical(SharedSecondaryBank)
-    {
-      for (auto& site : thread_bank) {
-        simulation::shared_secondary_bank_write.thread_unsafe_append(site);
-      }
-    }
   }
+  collect_sorted_history_secondary_banks(thread_banks);
+  thread_banks.clear();
 
   simulation::simulation_tracks_completed += settings::n_particles;
 
@@ -955,10 +1041,6 @@ void transport_history_based_shared_secondary()
   int64_t alive_secondary = 1;
   while (alive_secondary) {
 
-    // Sort the shared secondary bank by parent ID then progeny ID to
-    // ensure reproducibility.
-    sort_bank(simulation::shared_secondary_bank_write, false);
-
     // Synchronize the shared secondary bank amongst all MPI ranks, such
     // that each MPI rank has an approximately equal number of secondary
     // tracks. Also reports the total number of secondaries alive across
@@ -987,11 +1069,12 @@ void transport_history_based_shared_secondary()
       simulation::shared_secondary_bank_read.size());
     std::fill(simulation::progeny_per_particle.begin(),
       simulation::progeny_per_particle.end(), 0);
+    thread_banks.resize(num_threads());
 
     // Transport all secondary tracks from the shared secondary bank
 #pragma omp parallel
     {
-      vector<SourceSite> thread_bank;
+      auto& thread_bank = thread_banks[thread_num()];
       Particle p;
 
 #pragma omp for schedule(runtime)
@@ -1006,17 +1089,12 @@ void transport_history_based_shared_secondary()
         }
         p.local_secondary_bank().clear();
       }
-
-      // Drain thread-local bank into the shared secondary bank (once per
-      // thread)
-#pragma omp critical(SharedSecondaryBank)
-      {
-        for (auto& secondary_site : thread_bank) {
-          simulation::shared_secondary_bank_write.thread_unsafe_append(
-            secondary_site);
-        }
-      }
     } // End of transport loop over tracks in shared secondary bank
+    simulation::shared_secondary_bank_write =
+      std::move(simulation::shared_secondary_bank_read);
+    simulation::shared_secondary_bank_read = SharedArray<SourceSite>();
+    collect_sorted_history_secondary_banks(thread_banks);
+    thread_banks.clear();
     n_generation_depth++;
     simulation::simulation_tracks_completed += alive_secondary;
   } // End of loop over secondary generations
@@ -1080,13 +1158,7 @@ void transport_event_based_shared_secondary()
     process_transport_events();
     process_death_events(n_particles);
 
-    // Collect secondaries from all particle buffers into shared bank
-    for (int64_t i = 0; i < n_particles; i++) {
-      for (auto& site : simulation::particles[i].local_secondary_bank()) {
-        simulation::shared_secondary_bank_write.thread_unsafe_append(site);
-      }
-      simulation::particles[i].local_secondary_bank().clear();
-    }
+    collect_event_secondary_banks(n_particles);
 
     remaining_work -= n_particles;
     source_offset += n_particles;
@@ -1150,13 +1222,7 @@ void transport_event_based_shared_secondary()
       process_transport_events();
       process_death_events(n_particles);
 
-      // Collect secondaries from all particle buffers into shared bank
-      for (int64_t i = 0; i < n_particles; i++) {
-        for (auto& site : simulation::particles[i].local_secondary_bank()) {
-          simulation::shared_secondary_bank_write.thread_unsafe_append(site);
-        }
-        simulation::particles[i].local_secondary_bank().clear();
-      }
+      collect_event_secondary_banks(n_particles);
 
       sec_remaining -= n_particles;
       sec_offset += n_particles;