added O(n) fission bank sorting implementation

include/openmc/bank.h

@@ -20,6 +20,7 @@ extern std::vector<Particle::Bank> source_bank;
 extern Particle::Bank* fission_bank;
 extern int64_t fission_bank_length;
 extern int64_t fission_bank_max;
+extern std::vector<int64_t> progeny_per_particle;
 
 } // namespace simulation
 
@@ -27,6 +28,8 @@ extern int64_t fission_bank_max;
 // Non-member functions
 //==============================================================================
 
+void sort_fission_bank();
+
 void free_memory_bank();
 
 void init_fission_bank(int64_t max);

include/openmc/particle.h

@@ -171,6 +171,7 @@ public:
     int delayed_group;
     Type particle;
     int64_t parent_id;
+    int64_t progeny_id;
     bool operator<(const Bank& bank) const { return (parent_id < bank.parent_id); }
   };
   
@@ -359,6 +360,8 @@ public:
   moab::DagMC::RayHistory history_;
   Direction last_dir_;
   #endif
+
+  int64_t n_progeny_ {0}; // Number of progeny produced by this particle
 };
 
 } // namespace openmc

src/bank.cpp

@@ -1,6 +1,7 @@
 #include "openmc/bank.h"
 #include "openmc/capi.h"
 #include "openmc/error.h"
+#include "openmc/simulation.h"
 
 #include <cstdint>
 
@@ -19,6 +20,11 @@ Particle::Bank* fission_bank {nullptr};
 int64_t fission_bank_length {0};
 int64_t fission_bank_max;
 
+// Each entry in this vector corresponds to the number of progeny produced
+// this generation for the particle located at that index. This vector is
+// used to efficiently sort the fission bank after each iteration.
+std::vector<int64_t> progeny_per_particle;
+
 } // namespace simulation
 
 //==============================================================================
@@ -39,6 +45,38 @@ void init_fission_bank(int64_t max)
   simulation::fission_bank_max = max;
   simulation::fission_bank = new Particle::Bank[max];
   simulation::fission_bank_length = 0;
+  simulation::progeny_per_particle.resize(simulation::work_per_rank);
+}
+
+void sort_fission_bank()
+{
+  if( simulation::progeny_per_particle.size() == 0 )
+    return;
+
+  // Perform exclusive scan summation to determine starting indices in fission
+  // bank for each parent particle id
+  int64_t tmp = simulation::progeny_per_particle[0];
+  simulation::progeny_per_particle[0] = 0;
+  for (int64_t i = 1; i < simulation::progeny_per_particle.size(); i++) {
+    int64_t value = simulation::progeny_per_particle[i-1] + tmp;
+    tmp = simulation::progeny_per_particle[i];
+    simulation::progeny_per_particle[i] = value;
+  }
+
+  // Create temporary scratch vector for permutation
+  std::vector<Particle::Bank> sorted_bank(simulation::fission_bank_length);
+
+  // Use parent and progeny indices to sort fission bank
+  for (int64_t i = 0; i < simulation::fission_bank_length; i++) {
+    Particle::Bank& site = simulation::fission_bank[i];
+    int64_t idx = simulation::progeny_per_particle[site.parent_id-1] + site.progeny_id;
+    sorted_bank[idx] = site;
+  }
+
+  // Copy sorted bank into the fission bank
+  for (int64_t i = 0; i < simulation::fission_bank_length; i++) {
+    simulation::fission_bank[i] = sorted_bank[i];
+  }
 }
 
 //==============================================================================

src/particle.cpp

@@ -391,6 +391,10 @@ Particle::event_death()
   keff_tally_collision_   = 0.0;
   keff_tally_tracklength_ = 0.0;
   keff_tally_leakage_     = 0.0;
+
+  // Record the number of progeny created by this particle.
+  // This data will be used to efficiently sort the fission bank.
+  simulation::progeny_per_particle[id_-1] = n_progeny_;
 }
 
 

src/physics.cpp

@@ -203,6 +203,7 @@ create_fission_sites(Particle* p, int i_nuclide, const Reaction* rx)
     site->particle = Particle::Type::neutron;
     site->wgt = 1. / weight;
     site->parent_id = p->id_;
+    site->progeny_id = p->n_progeny_++;
 
     // Sample delayed group and angle/energy for fission reaction
     sample_fission_neutron(i_nuclide, rx, p->E_, site, p->current_seed());

src/physics_mg.cpp

@@ -151,6 +151,7 @@ create_fission_sites(Particle* p)
     site->particle = Particle::Type::neutron;
     site->wgt = 1. / weight;
     site->parent_id = p->id_;
+    site->progeny_id = p->n_progeny_++;
 
     // Sample the cosine of the angle, assuming fission neutrons are emitted
     // isotropically

src/simulation.cpp

@@ -398,8 +398,9 @@ void finalize_generation()
     // so as to allow for reproducibility regardless of which order particles
     // are run in.
     #ifdef _OPENMP
-    std::stable_sort(simulation::fission_bank,
-        simulation::fission_bank + simulation::fission_bank_length);
+    //std::stable_sort(simulation::fission_bank,
+    //    simulation::fission_bank + simulation::fission_bank_length);
+    sort_fission_bank();
     #endif
 
     // Distribute fission bank across processors evenly
@@ -434,6 +435,9 @@ void initialize_history(Particle* p, int64_t index_source)
   // set identifier for particle
   p->id_ = simulation::work_index[mpi::rank] + index_source;
 
+  // set progeny count to zero
+  p->n_progeny_ = 0;
+
   // set random number seed
   int64_t particle_seed = (simulation::total_gen + overall_generation() - 1)
     * settings::n_particles + p->id_;