#1934: Adjust buffers size in objs of the collection

DARMA-tasking · Nov 28, 2023 · f23d8e9 · f23d8e9
1 parent 771beb9
commit f23d8e9
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/src/vt/elm/elm_lb_data.h b/src/vt/elm/elm_lb_data.h
@@ -116,16 +116,13 @@ struct ElementLBData {
   static const constexpr SubphaseType no_subphase =
     std::numeric_limits<SubphaseType>::max();
 
-protected:
   /**
-   * \internal \brief Resize internal buffers
+   * \brief Resize internal buffers
    * 
    * \param[in] hist_lb_data_count the requested buffers capacity
    */
   void setHistoryCapacity(unsigned int hist_lb_data_count);
 
-  friend struct vrt::collection::balance::NodeLBData;
-
 protected:
   bool cur_time_started_ = false;
   TimeType cur_time_ = TimeType{0.0};

diff --git a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc
@@ -170,6 +170,14 @@ void LBManager::setLoadModel(std::shared_ptr<LoadModel> model) {
   model_ = model;
   model_->setLoads(nlb_data->getNodeLoad(),
                    nlb_data->getNodeComm());
+
+  // Adjust buffers size on the objs in collection
+  runInEpochCollective("LBManager: setLoadModel", [=] {
+    auto& objs = vt::objgroup::getObjs();
+    for (auto& obj : objs) {
+      obj.second->getLBData().setHistoryCapacity(min_hist_lb_data_);
+    }
+  });
 }
 
 template <typename LB>