#1934: Adjust buffers size in objs of the collection

DARMA-tasking · Dec 11, 2023 · 5aefbe3 · 5aefbe3
1 parent d0a4f72
commit 5aefbe3
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 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
@@ -169,8 +169,16 @@ void LBManager::setLoadModel(std::shared_ptr<LoadModel> model) {
 
   model_ = model;
   model_->setLoads(nlb_data->getNodeLoad(),
-                   nlb_data->getNodeComm(),
                    nlb_data->getUserData());
+                   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>