#1838: TemperedWMin: compose total work model on top of existing one

DARMA-tasking · Nov 8, 2022 · 0143287 · 0143287
1 parent df36c63
commit 0143287
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.cc b/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.cc
@@ -43,7 +43,9 @@
 
 #include "vt/vrt/collection/balance/temperedwmin/temperedwmin.h"
 
+#include "vt/vrt/collection/balance/lb_invoke/lb_manager.h"
 #include "vt/vrt/collection/balance/model/load_model.h"
+#include "vt/vrt/collection/balance/model/weighted_communication_volume.h"
 
 namespace vt { namespace vrt { namespace collection { namespace lb {
 
@@ -93,13 +95,17 @@ void TemperedWMin::inputParams(balance::SpecEntry* spec) {
     "TemperedWMin::inputParams: alpha={}, beta={}, gamma={}\n",
     alpha_, beta_, gamma_
   );
+
+  total_work_model_ = std::make_unique<balance::WeightedCommunicationVolume>(
+    theLBManager()->getLoadModel(), alpha_, beta_, gamma_
+  );
 }
 
 TimeType TemperedWMin::getModeledValue(const elm::ElementIDStruct& obj) {
   balance::PhaseOffset when =
       {balance::PhaseOffset::NEXT_PHASE, balance::PhaseOffset::WHOLE_PHASE};
 
-  return load_model_->getModeledWork(obj, when);
+  return total_work_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
@@ -44,8 +44,11 @@
 #if !defined INCLUDED_VT_VRT_COLLECTION_BALANCE_TEMPEREDWMIN_TEMPEREDWMIN_H
 #define INCLUDED_VT_VRT_COLLECTION_BALANCE_TEMPEREDWMIN_TEMPEREDWMIN_H
 
+#include "vt/vrt/collection/balance/model/load_model.h"
 #include "vt/vrt/collection/balance/temperedlb/temperedlb.h"
 
+#include <memory>
+
 namespace vt { namespace vrt { namespace collection { namespace lb {
 
 struct TemperedWMin : TemperedLB {
@@ -64,6 +67,8 @@ struct TemperedWMin : TemperedLB {
   TimeType getModeledValue(const elm::ElementIDStruct& obj) override;
 
 private:
+  std::unique_ptr<balance::LoadModel> total_work_model_ = nullptr;
+
   double alpha_ = 1.0;
   double beta_  = 0.0;
   double gamma_ = 0.0;