#1672: use theNodeLBData where necessary

DARMA-tasking · May 30, 2022 · 8be2e96 · 8be2e96
1 parent 730d050
commit 8be2e96
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/src/vt/vrt/collection/balance/baselb/baselb.cc b/src/vt/vrt/collection/balance/baselb/baselb.cc
@@ -263,7 +263,7 @@ void BaseLB::finalize(CountMsg* msg) {
 
 void BaseLB::recvSharedEdges(CommMsg* msg) {
   auto phase = thePhase()->getCurrentPhase();
-  auto comm_map = theNodeStats()->getNodeComm(phase);
+  auto comm_map = theNodeLBData()->getNodeComm(phase);
 
   if (comm_map != nullptr) {
     auto& comm = msg->comm_;

diff --git a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc
@@ -732,8 +732,8 @@ void makeGraphSymmetric(
   PhaseType phase, objgroup::proxy::Proxy<lb::BaseLB> proxy
 ) {
   auto const this_node = theContext()->getNode();
-  auto iter = theNodeStats()->getNodeComm()->find(phase);
-  if (iter == theNodeStats()->getNodeComm()->end()) {
+  auto iter = theNodeLBData()->getNodeComm()->find(phase);
+  if (iter == theNodeLBData()->getNodeComm()->end()) {
     return;
   }
 

diff --git a/tests/unit/collection/test_lb.extended.cc b/tests/unit/collection/test_lb.extended.cc
@@ -49,6 +49,7 @@
 
 #include "vt/vrt/collection/manager.h"
 #include "vt/vrt/collection/balance/lb_data_holder.h"
+#include "vt/vrt/collection/balance/node_lb_data.h"
 #include "vt/vrt/collection/balance/lb_invoke/lb_manager.h"
 #include "vt/utils/json/json_reader.h"
 #include "vt/utils/json/json_appender.h"
@@ -163,8 +164,8 @@ TEST_F(TestLoadBalancerOther, test_make_graph_symmetric) {
   // );
 
   auto phase = num_phases - 1;
-  auto iter = theNodeStats()->getNodeComm()->find(phase);
-  ASSERT_NE(iter, theNodeStats()->getNodeComm()->end());
+  auto iter = theNodeLBData()->getNodeComm()->find(phase);
+  ASSERT_NE(iter, theNodeLBData()->getNodeComm()->end());
 
   elm::CommMapType const& comm_data = iter->second;