#1838: lb: only use modeled work in TemperedWMin

DARMA-tasking · Aug 5, 2022 · efc2016 · efc2016
1 parent bb120ca
commit efc2016
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 5 deletions.
diff --git a/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc b/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc
@@ -499,11 +499,7 @@ void TemperedLB::doLBStages(TimeType start_imb) {
         cur_objs_.clear();
         for (auto obj : *load_model_) {
           if (obj.isMigratable()) {
-            cur_objs_[obj] = load_model_->getModeledWork(
-              obj,
-              {balance::PhaseOffset::NEXT_PHASE,
-               balance::PhaseOffset::WHOLE_PHASE}
-            );
+            cur_objs_[obj] = getModeledValue(obj);
           }
         }
         this_new_load_ = this_load;
@@ -1365,4 +1361,10 @@ void TemperedLB::migrate() {
   vtAssertExpr(false);
 }
 
+TimeType TemperedLB::getModeledValue(const elm::ElementIDStruct& obj) {
+  return load_model_->getModeledLoad(
+    obj, {balance::PhaseOffset::NEXT_PHASE, balance::PhaseOffset::WHOLE_PHASE}
+  );
+}
+
 }}}} /* end namespace vt::vrt::collection::lb */
diff --git a/src/vt/vrt/collection/balance/temperedlb/temperedlb.h b/src/vt/vrt/collection/balance/temperedlb/temperedlb.h
@@ -111,6 +111,7 @@ struct TemperedLB : BaseLB {
   ElementLoadType::iterator selectObject(
     LoadType size, ElementLoadType& load, std::set<ObjIDType> const& available
   );
+  virtual TimeType getModeledValue(const elm::ElementIDStruct& obj);
 
   void lazyMigrateObjsTo(EpochType epoch, NodeType node, ObjsType const& objs);
   void inLazyMigrations(balance::LazyMigrationMsg* msg);

diff --git a/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.cc b/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.cc
@@ -43,6 +43,8 @@
 
 #include "vt/vrt/collection/balance/temperedwmin/temperedwmin.h"
 
+#include "vt/vrt/collection/balance/model/load_model.h"
+
 namespace vt { namespace vrt { namespace collection { namespace lb {
 
 void TemperedWMin::init(objgroup::proxy::Proxy<TemperedWMin> in_proxy) {
@@ -52,4 +54,11 @@ void TemperedWMin::init(objgroup::proxy::Proxy<TemperedWMin> in_proxy) {
   strat->init(proxy);
 }
 
+TimeType TemperedWMin::getModeledValue(const elm::ElementIDStruct& obj) {
+  balance::PhaseOffset when =
+      {balance::PhaseOffset::NEXT_PHASE, balance::PhaseOffset::WHOLE_PHASE};
+
+  return load_model_->getModeledWork(obj, when);
+}
+
 }}}} // namespace vt::vrt::collection::lb
diff --git a/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h b/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h
@@ -56,6 +56,9 @@ struct TemperedWMin : TemperedLB {
 
 public:
   void init(objgroup::proxy::Proxy<TemperedWMin> in_proxy);
+
+protected:
+  TimeType getModeledValue(const elm::ElementIDStruct& obj) override;
 };
 
 }}}} /* end namespace vt::vrt::collection::lb */