From 7810378f3c47516019b8b6f8c754be1a60662ee7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cezary=20Skrzy=C5=84ski?=
 <cezary.skrzynski@ng-analytics.com>
Date: Thu, 14 Apr 2022 20:18:16 +0200
Subject: [PATCH] #1672: send selected edges

---
 .../vrt/collection/balance/baselb/baselb.cc   | 10 ++++++
 src/vt/vrt/collection/balance/baselb/baselb.h | 21 +++++++++++-
 .../balance/lb_invoke/lb_manager.cc           | 32 ++-----------------
 3 files changed, 33 insertions(+), 30 deletions(-)

diff --git a/src/vt/vrt/collection/balance/baselb/baselb.cc b/src/vt/vrt/collection/balance/baselb/baselb.cc
index 07cc718ad9..1ad0323a71 100644
--- a/src/vt/vrt/collection/balance/baselb/baselb.cc
+++ b/src/vt/vrt/collection/balance/baselb/baselb.cc
@@ -252,6 +252,16 @@ void BaseLB::finalize(CountMsg* msg) {
   }
 }
 
+void BaseLB::recvSharedEdges(CommMsg* msg) {
+  auto& comm = msg->comm_;
+  for (auto&& elm : comm) {
+    vt_debug_print(
+      verbose, lb, "recv shared edge: from={}, to={}\n", elm.first.fromObj(),
+      elm.first.toObj()
+    );
+  }
+}
+
 }}}} /* end namespace vt::vrt::collection::lb */
 
 #endif /*INCLUDED_VT_VRT_COLLECTION_BALANCE_BASELB_BASELB_CC*/
diff --git a/src/vt/vrt/collection/balance/baselb/baselb.h b/src/vt/vrt/collection/balance/baselb/baselb.h
index cdbaa5f20d..22e52066bc 100644
--- a/src/vt/vrt/collection/balance/baselb/baselb.h
+++ b/src/vt/vrt/collection/balance/baselb/baselb.h
@@ -64,6 +64,8 @@ struct LoadModel;
 
 namespace lb {
 
+struct CommMsg;
+
 struct BaseLB {
   using ObjIDType        = balance::ElementIDStruct;
   using ElementLoadType  = std::unordered_map<ObjIDType,TimeType>;
@@ -142,6 +144,7 @@ struct BaseLB {
   TransferVecType& getTransfers() { return transfers_; }
 
   bool isCommAware() const { return comm_aware_; }
+  void recvSharedEdges(CommMsg* msg);
 
 protected:
   void getArgs(PhaseType phase);
@@ -173,6 +176,22 @@ struct BaseLB {
   std::shared_ptr<balance::Reassignment> pending_reassignment_ = nullptr;
 };
 
-}}}} /* end namespace vt::vrt::collection::lb */
+struct CommMsg : vt::Message {
+  using MessageParentType = vt::Message;
+  vt_msg_serialize_required();
+
+  CommMsg() = default;
+  explicit CommMsg(lb::BaseLB::ElementCommType in_comm) : comm_(in_comm) { }
+
+  lb::BaseLB::ElementCommType comm_;
+
+  template <typename SerializerT>
+  void serialize(SerializerT& s) {
+    MessageParentType::serialize(s);
+    s | comm_;
+  }
+};
+
+}}}} // namespace vt::vrt::collection::lb
 
 #endif /*INCLUDED_VT_VRT_COLLECTION_BALANCE_BASELB_BASELB_H*/
diff --git a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc
index a672d9c7c5..ad0f91fe63 100644
--- a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc
+++ b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc
@@ -532,32 +532,6 @@ bool LBManager::isCollectiveComm(elm::CommCategory cat) const {
   return is_collective;
 }
 
-struct CommMsg : vt::Message {
-  using MessageParentType = vt::Message;
-  vt_msg_serialize_required(); // comm_
-
-  CommMsg() = default;
-  explicit CommMsg(lb::BaseLB::ElementCommType in_comm) : comm_(in_comm) { }
-
-  lb::BaseLB::ElementCommType comm_;
-
-  template <typename SerializerT>
-  void serialize(SerializerT& s) {
-    MessageParentType::serialize(s);
-    s | comm_;
-  }
-};
-
-void recvSharedEdges(CommMsg* msg) {
-  auto& comm = msg->comm_;
-  for (auto&& elm : comm) {
-    vt_debug_print(
-      verbose, lb, "recv shared edge: from={}, to={}\n",
-      elm.first.fromObj(), elm.first.toObj()
-    );
-  }
-}
-
 void makeGraphSymmetric(
   PhaseType phase, objgroup::proxy::Proxy<lb::BaseLB> proxy
 ) {
@@ -605,9 +579,9 @@ void makeGraphSymmetric(
   }
 
   for (auto&& elm : shared_edges) {
-    (void)elm;
-    (void)proxy;
-    // proxy[elm.first].send<CommMsg, &recvSharedEdges>(elm.second);
+    proxy[elm.first].send<lb::CommMsg, &lb::BaseLB::recvSharedEdges>(
+      elm.second
+    );
   }
 }