CBC_FLUDS uses memory pool allocator to minimize memory usage for cell angular fluxes during sweeps

modules/linear_boltzmann_solvers/discrete_ordinates_problem/discrete_ordinates_problem.cc

@@ -1022,16 +1022,18 @@ DiscreteOrdinatesProblem::InitFluxDataStructures(LBSGroupset& groupset)
       }
       else if (sweep_type_ == "CBC")
       {
-        OpenSnLogicalErrorIf(not options_.save_angular_flux,
-                             "When using sweep_type \"CBC\" then "
-                             "\"save_angular_flux\" must be true.");
+        const size_t num_local_cells = grid_->local_cells.size();
+        const auto cbc_sweep_ordering = std::static_pointer_cast<CBC_SPDS>(sweep_ordering);
+        const auto& min_num_pool_allocator_slots =
+          cbc_sweep_ordering->GetMinNumPoolAllocatorSlots();
+
         std::shared_ptr<FLUDS> fluds =
           std::make_shared<CBC_FLUDS>(gs_num_grps,
                                       angle_indices.size(),
                                       dynamic_cast<const CBC_FLUDSCommonData&>(fluds_common_data),
-                                      psi_new_local_[groupset.id],
-                                      groupset.psi_uk_man_,
-                                      *discretization_);
+                                      num_local_cells,
+                                      max_cell_dof_count_,
+                                      min_num_pool_allocator_slots);
 
         auto angle_set = std::make_shared<CBC_AngleSet>(angle_set_id++,
                                                         gs_num_grps,

modules/linear_boltzmann_solvers/discrete_ordinates_problem/sweep/angle_set/cbc_angle_set.cc

@@ -24,7 +24,8 @@ CBC_AngleSet::CBC_AngleSet(size_t id,
                            bool use_gpu)
   : AngleSet(id, num_groups, spds, fluds, angle_indices, boundaries, use_gpu),
     cbc_spds_(dynamic_cast<const CBC_SPDS&>(spds_)),
-    async_comm_(id, *fluds, comm_set)
+    async_comm_(id, *fluds, comm_set),
+    cbc_fluds_(dynamic_cast<CBC_FLUDS&>(*fluds_))
 {
 }
 
@@ -70,15 +71,31 @@ CBC_AngleSet::AngleSetAdvance(SweepChunk& sweep_chunk, AngleSetStatus permission
         all_tasks_completed = false;
       if (cell_task.num_dependencies == 0 and not cell_task.completed)
       {
+        cbc_fluds_.Allocate(cell_task.cell_ptr->local_id);
+
         sweep_chunk.SetCell(cell_task.cell_ptr, *this);
         sweep_chunk.Sweep(*this);
 
-        for (uint64_t local_task_num : cell_task.successors)
+        for (uint64_t local_task_num : cell_task.local_successors)
           --current_task_list_[local_task_num].num_dependencies;
 
         cell_task.completed = true;
         a_task_executed = true;
         async_comm_.SendData();
+
+        // Update predecessor dependency consumption counts
+        for (uint64_t local_task_num : cell_task.local_predecessors)
+        {
+          ++current_task_list_[local_task_num].num_satisfied_downwind_deps;
+
+          if (current_task_list_[local_task_num].num_satisfied_downwind_deps >=
+              current_task_list_[local_task_num].local_successors.size())
+            cbc_fluds_.Deallocate(current_task_list_[local_task_num].cell_ptr->local_id);
+        }
+
+        // Deallocate if cell has no local successors
+        if (cell_task.local_successors.empty())
+          cbc_fluds_.Deallocate(cell_task.cell_ptr->local_id);
       }
     } // for cell_task
     async_comm_.SendData();

modules/linear_boltzmann_solvers/discrete_ordinates_problem/sweep/angle_set/cbc_angle_set.h

@@ -18,6 +18,7 @@ class CBC_AngleSet : public AngleSet
   const CBC_SPDS& cbc_spds_;
   std::vector<Task> current_task_list_;
   CBC_ASynchronousCommunicator async_comm_;
+  CBC_FLUDS& cbc_fluds_;
 
 public:
   CBC_AngleSet(size_t id,

modules/linear_boltzmann_solvers/discrete_ordinates_problem/sweep/fluds/cbc_fluds.cc

@@ -5,22 +5,27 @@
 #include "modules/linear_boltzmann_solvers/discrete_ordinates_problem/sweep/spds/spds.h"
 #include "framework/math/spatial_discretization/spatial_discretization.h"
 #include "framework/mesh/mesh_continuum/mesh_continuum.h"
+#include "caliper/cali.h"
+#include <cstddef>
 
 namespace opensn
 {
 
 CBC_FLUDS::CBC_FLUDS(size_t num_groups,
                      size_t num_angles,
                      const CBC_FLUDSCommonData& common_data,
-                     std::vector<double>& local_psi_data,
-                     const UnknownManager& psi_uk_man,
-                     const SpatialDiscretization& sdm)
+                     size_t num_local_cells,
+                     size_t max_cell_dof_count,
+                     size_t min_num_pool_allocator_slots)
   : FLUDS(num_groups, num_angles, common_data.GetSPDS()),
     common_data_(common_data),
-    local_psi_data_(local_psi_data),
-    psi_uk_man_(psi_uk_man),
-    sdm_(sdm)
+    slot_size_(max_cell_dof_count * num_groups_and_angles_),
+    cell_local_ID_to_psi_map_(num_local_cells, nullptr),
+    local_psi_data_backing_buffer_(min_num_pool_allocator_slots * slot_size_)
 {
+  local_psi_data_.add_block(local_psi_data_backing_buffer_.data(),
+                            (min_num_pool_allocator_slots * slot_size_) * sizeof(double),
+                            slot_size_ * sizeof(double));
 }
 
 const FLUDSCommonData&
@@ -29,32 +34,60 @@ CBC_FLUDS::GetCommonData() const
   return common_data_;
 }
 
-const std::vector<double>&
-CBC_FLUDS::GetLocalUpwindDataBlock() const
+void
+CBC_FLUDS::Allocate(uint64_t cell_local_ID)
 {
-  return local_psi_data_;
+  assert(cell_local_ID_to_psi_map_[cell_local_ID] == nullptr);
+  void* cell_block_ptr = local_psi_data_.malloc();
+  cell_local_ID_to_psi_map_[cell_local_ID] = static_cast<double*>(cell_block_ptr);
 }
 
-const double*
-CBC_FLUDS::GetLocalCellUpwindPsi(const std::vector<double>& psi_data_block, const Cell& cell)
+void
+CBC_FLUDS::Deallocate(uint64_t cell_local_ID)
 {
-  const auto dof_map = sdm_.MapDOFLocal(cell, 0, psi_uk_man_, 0, 0);
-  return &psi_data_block[dof_map];
+  assert(cell_local_ID_to_psi_map_[cell_local_ID] != nullptr);
+  local_psi_data_.free(cell_local_ID_to_psi_map_[cell_local_ID]);
+  cell_local_ID_to_psi_map_[cell_local_ID] = nullptr;
 }
 
-const std::vector<double>&
-CBC_FLUDS::GetNonLocalUpwindData(uint64_t cell_global_id, unsigned int face_id) const
+double*
+CBC_FLUDS::UpwindPsi(uint64_t cell_local_id, unsigned int adj_cell_node, size_t as_ss_idx)
 {
-  return deplocs_outgoing_messages_.at({cell_global_id, face_id});
+  assert(cell_local_ID_to_psi_map_[cell_local_id] != nullptr);
+  const size_t addr_offset = adj_cell_node * num_groups_and_angles_ + as_ss_idx * num_groups_;
+  return cell_local_ID_to_psi_map_[cell_local_id] + addr_offset;
 }
 
-const double*
-CBC_FLUDS::GetNonLocalUpwindPsi(const std::vector<double>& psi_data,
-                                unsigned int face_node_mapped,
-                                unsigned int angle_set_index)
+double*
+CBC_FLUDS::OutgoingPsi(uint64_t cell_local_ID, unsigned int cell_node, size_t as_ss_idx)
 {
-  const size_t dof_map = face_node_mapped * num_groups_and_angles_ + angle_set_index * num_groups_;
-  return &psi_data[dof_map];
+  assert(cell_local_ID_to_psi_map_[cell_local_ID] != nullptr);
+  const size_t addr_offset = cell_node * num_groups_and_angles_ + as_ss_idx * num_groups_;
+  return cell_local_ID_to_psi_map_[cell_local_ID] + addr_offset;
+}
+
+double*
+CBC_FLUDS::NLUpwindPsi(uint64_t cell_global_id,
+                       unsigned int face_id,
+                       unsigned int face_node_mapped,
+                       size_t as_ss_idx)
+{
+  std::vector<double>& psi = deplocs_outgoing_messages_.at({cell_global_id, face_id});
+  const size_t dof_map =
+    face_node_mapped * num_groups_and_angles_ + //  Offset to start of data for face_node_mapped
+    as_ss_idx * num_groups_;                    // Offset to start of data for angle_set_index
+  assert((dof_map >= 0) and (dof_map < psi.size()));
+  return &psi[dof_map];
+}
+
+double*
+CBC_FLUDS::NLOutgoingPsi(std::vector<double>* psi_nonlocal_outgoing,
+                         size_t face_node,
+                         size_t as_ss_idx)
+{
+  assert(psi_nonlocal_outgoing != nullptr);
+  const size_t addr_offset = face_node * num_groups_and_angles_ + as_ss_idx * num_groups_;
+  return &(*psi_nonlocal_outgoing)[addr_offset];
 }
 
 } // namespace opensn

modules/linear_boltzmann_solvers/discrete_ordinates_problem/sweep/fluds/cbc_fluds.h

@@ -5,6 +5,8 @@
 
 #include "modules/linear_boltzmann_solvers/discrete_ordinates_problem/sweep/fluds/cbc_fluds_common_data.h"
 #include "modules/linear_boltzmann_solvers/discrete_ordinates_problem/sweep/fluds/fluds.h"
+#include "boost/pool/simple_segregated_storage.hpp"
+#include <cstddef>
 #include <map>
 #include <functional>
 
@@ -15,28 +17,73 @@ class UnknownManager;
 class SpatialDiscretization;
 class Cell;
 
+/**
+ * Flux data structures (FLUDS) specific to the cell-by-cell (CBC) sweep algorithm
+ *
+ * This class manages the storage and access of angular flux data during a CBC sweep
+ *
+ * It provides methods to access:
+ * - Upwind angular flux data from local neighbor cells
+ * - Storage locations for downwind angular flux data for the current cell
+ * - Upwind angular flux data received from remote MPI ranks
+ */
 class CBC_FLUDS : public FLUDS
 {
 public:
   CBC_FLUDS(size_t num_groups,
             size_t num_angles,
             const CBC_FLUDSCommonData& common_data,
-            std::vector<double>& local_psi_data,
-            const UnknownManager& psi_uk_man,
-            const SpatialDiscretization& sdm);
+            size_t num_local_cells,
+            size_t max_cell_dof_count,
+            size_t min_num_pool_allocator_slots);
 
   const FLUDSCommonData& GetCommonData() const;
 
-  const std::vector<double>& GetLocalUpwindDataBlock() const;
-
-  const double* GetLocalCellUpwindPsi(const std::vector<double>& psi_data_block, const Cell& cell);
-
-  const std::vector<double>& GetNonLocalUpwindData(uint64_t cell_global_id,
-                                                   unsigned int face_id) const;
-
-  const double* GetNonLocalUpwindPsi(const std::vector<double>& psi_data,
-                                     unsigned int face_node_mapped,
-                                     unsigned int angle_set_index);
+  /**
+   * Given a cell's local ID, retrieve a free slot in the pool allocator for storing
+   * the cell's angular flux data and map the cell's local ID to the slot
+   */
+  void Allocate(uint64_t cell_local_ID);
+
+  /**
+   * Given a cell's local ID, deallocate the memory used for the cell's angular flux data
+   * and remove the mapping from the cell's local ID to the slot
+   */
+  void Deallocate(uint64_t cell_local_ID);
+
+  /**
+   * Given a local upwind neighbor cell local cell ID, a node index on this cell, and an
+   * angleset subset index, this function returns a pointer to
+   * the start of the group data for the specified node and angle.
+   */
+  double* UpwindPsi(uint64_t cell_local_id, unsigned int adj_cell_node, size_t as_ss_idx);
+
+  /**
+   * Given a local cell's local ID, a node index on this cell, and an angleset subset index,
+   * this function returns a pointer to the start of the group data for the specified
+   * node and angle for writing its just solved angular fluxes.
+   */
+  double* OutgoingPsi(uint64_t cell_local_id, unsigned int cell_node, size_t as_ss_idx);
+
+  /**
+   * Given a remote upwind cell's global ID, a face ID on this cell,
+   * a node index on this face, and an angleset subset index,
+   * this function returns a pointer to the start of the group data for the specified
+   * face node and angle.
+   */
+  double* NLUpwindPsi(uint64_t cell_global_id,
+                      unsigned int face_id,
+                      unsigned int face_node_mapped,
+                      size_t as_ss_idx);
+
+  /**
+   * Given a pointer to a vector holding the non-local outgoing psi data for a face,
+   * a node index on this face, and an angleset subset index,
+   * this function returns a pointer to the start of the group data for the specified
+   * face node and angle.
+   */
+  double*
+  NLOutgoingPsi(std::vector<double>* psi_nonlocal_outgoing, size_t face_node, size_t as_ss_idx);
 
   void ClearLocalAndReceivePsi() override { deplocs_outgoing_messages_.clear(); }
   void ClearSendPsi() override {}
@@ -57,9 +104,10 @@ class CBC_FLUDS : public FLUDS
 
 private:
   const CBC_FLUDSCommonData& common_data_;
-  std::reference_wrapper<std::vector<double>> local_psi_data_;
-  const UnknownManager& psi_uk_man_;
-  const SpatialDiscretization& sdm_;
+  size_t slot_size_;
+  std::vector<double*> cell_local_ID_to_psi_map_;
+  std::vector<double> local_psi_data_backing_buffer_;
+  boost::simple_segregated_storage<size_t> local_psi_data_;
 
   std::vector<std::vector<double>> boundryI_incoming_psi_;