#2201: implement memory aware TemperedLB in vt

src/vt/configs/arguments/app_config.h

@@ -158,6 +158,7 @@ struct AppConfig {
   bool vt_lb_self_migration      = false;
   bool vt_lb_spec                = false;
   std::string vt_lb_spec_file    = "";
+  bool vt_lb_run_lb_first_phase = false;
 
 
   bool vt_no_detect_hang       = false;

src/vt/configs/arguments/args.cc

@@ -481,6 +481,7 @@ void addLbArgs(CLI::App& app, AppConfig& appConfig) {
   auto lb_self_migration = "Allow load balancer to migrate objects to the same node";
   auto lb_spec      = "Enable LB spec file (defines which phases output LB data)";
   auto lb_spec_file = "File containing LB spec; --vt_lb_spec to enable";
+  auto lb_first_phase_info = "Force LB to run on the first phase (phase 0)";
   auto s  = app.add_flag("--vt_lb", appConfig.vt_lb, lb);
   auto t1 = app.add_flag("--vt_lb_quiet", appConfig.vt_lb_quiet, lb_quiet);
   auto u  = app.add_option("--vt_lb_file_name", appConfig.vt_lb_file_name, lb_file_name)->capture_default_str()->check(CLI::ExistingFile);
@@ -503,6 +504,7 @@ void addLbArgs(CLI::App& app, AppConfig& appConfig) {
   auto lbasm = app.add_flag("--vt_lb_self_migration",   appConfig.vt_lb_self_migration,      lb_self_migration);
   auto lbspec = app.add_flag("--vt_lb_spec",            appConfig.vt_lb_spec,                lb_spec);
   auto lbspecfile = app.add_option("--vt_lb_spec_file", appConfig.vt_lb_spec_file,           lb_spec_file)->capture_default_str()->check(CLI::ExistingFile);
+  auto lb_first_phase = app.add_flag("--vt_lb_run_lb_first_phase", appConfig.vt_lb_run_lb_first_phase, lb_first_phase_info);
 
   // --vt_lb_name excludes --vt_lb_file_name, and vice versa
   v->excludes(u);
@@ -531,6 +533,7 @@ void addLbArgs(CLI::App& app, AppConfig& appConfig) {
   lbasm->group(debugLB);
   lbspec->group(debugLB);
   lbspecfile->group(debugLB);
+  lb_first_phase->group(debugLB);
 
   // help options deliberately omitted from the debugLB group above so that
   // they appear grouped with --vt_help when --vt_help is used

src/vt/elm/elm_comm.h

@@ -58,7 +58,9 @@ enum struct CommCategory : int8_t {
   CollectionToNodeBcast = 5,
   NodeToCollectionBcast = 6,
   CollectiveToCollectionBcast = 7,
-  LocalInvoke = 8
+  LocalInvoke = 8,
+  WriteShared = 9,
+  ReadOnlyShared = 10
 };
 
 inline NodeType objGetNode(ElementIDStruct const id) {
@@ -71,6 +73,8 @@ struct CommKey {
   struct CollectionTag { };
   struct CollectionToNodeTag { };
   struct NodeToCollectionTag { };
+  struct WriteSharedTag { };
+  struct ReadOnlySharedTag { };
 
   CommKey() = default;
   CommKey(CommKey const&) = default;
@@ -107,12 +111,25 @@ struct CommKey {
       cat_(bcast ? CommCategory::NodeToCollectionBcast : CommCategory::NodeToCollection)
   { }
 
+  CommKey(
+    WriteSharedTag,
+    NodeType in_home, int in_shared_id
+  ) : nto_(in_home), shared_id_(in_shared_id), cat_(CommCategory::WriteShared)
+  { }
+
+  CommKey(
+    ReadOnlySharedTag,
+    NodeType in_home, int in_shared_id
+  ) : nto_(in_home), shared_id_(in_shared_id), cat_(CommCategory::ReadOnlyShared)
+  { }
+
   ElementIDStruct from_ = {};
   ElementIDStruct to_   = {};
 
   ElementIDStruct edge_id_ = {};
   NodeType nfrom_          = uninitialized_destination;
   NodeType nto_            = uninitialized_destination;
+  int shared_id_            = -1;
   CommCategory  cat_       = CommCategory::SendRecv;
 
   ElementIDStruct fromObj()    const { return from_; }
@@ -121,6 +138,7 @@ struct CommKey {
   ElementIDType toNode()       const { return nto_; }
   ElementIDStruct edgeID()     const { return edge_id_; }
   CommCategory commCategory()  const { return cat_; }
+  int sharedID()               const { return shared_id_; }
 
   bool selfEdge() const { return cat_ == CommCategory::SendRecv and from_ == to_; }
   bool offNode() const {
@@ -140,12 +158,12 @@ struct CommKey {
     return
       k.from_  ==  from_ and k.to_  ==  to_ and
       k.nfrom_ == nfrom_ and k.nto_ == nto_ and
-      k.cat_   == cat_;
+      k.cat_   == cat_   and k.shared_id_ == shared_id_;
   }
 
   template <typename SerializerT>
   void serialize(SerializerT& s) {
-    s | from_ | to_ | nfrom_ | nto_ | cat_ | edge_id_;
+    s | from_ | to_ | nfrom_ | nto_ | cat_ | edge_id_ | shared_id_;
   }
 };
 
@@ -189,7 +207,8 @@ struct hash<vt::elm::CommKey> {
   size_t operator()(vt::elm::CommKey const& in) const {
     return std::hash<uint64_t>()(
       std::hash<vt::elm::ElementIDStruct>()(in.from_) ^
-      std::hash<vt::elm::ElementIDStruct>()(in.to_) ^ in.nfrom_ ^ in.nto_
+      std::hash<vt::elm::ElementIDStruct>()(in.to_) ^ in.nfrom_ ^ in.nto_ ^
+      in.shared_id_
     );
   }
 };

src/vt/elm/elm_id.cc

@@ -41,6 +41,7 @@
 //@HEADER
 */
 
+#include "vt/context/context.h"
 #include "vt/elm/elm_id.h"
 #include "vt/elm/elm_id_bits.h"
 
@@ -58,4 +59,8 @@ NodeType ElementIDStruct::getCurrNode() const {
   return curr_node;
 }
 
+bool ElementIDStruct::isLocatedOnThisNode() const {
+  return theContext()->getNode() == curr_node and not isMigratable();
+}
+
 }} /* end namespace vt::elm */

src/vt/elm/elm_id.h

@@ -78,6 +78,7 @@ struct ElementIDStruct {
   bool isMigratable() const;
   NodeType getHomeNode() const;
   NodeType getCurrNode() const;
+  bool isLocatedOnThisNode() const;
 };
 
 

src/vt/elm/elm_lb_data.cc

@@ -86,6 +86,24 @@ void ElementLBData::sendToEntity(
   sendComm(key, bytes);
 }
 
+void ElementLBData::addWritableSharedID(
+  NodeType home, int shared_id, double bytes
+) {
+  elm::CommKey key(elm::CommKey::WriteSharedTag{}, home, shared_id);
+  phase_comm_[cur_phase_][key].sendMsg(bytes);
+  subphase_comm_[cur_phase_].resize(cur_subphase_ + 1);
+  subphase_comm_[cur_phase_].at(cur_subphase_)[key].sendMsg(bytes);
+}
+
+void ElementLBData::addReadOnlySharedID(
+  NodeType home, int shared_id, double bytes
+) {
+  elm::CommKey key(elm::CommKey::ReadOnlySharedTag{}, home, shared_id);
+  phase_comm_[cur_phase_][key].sendMsg(bytes);
+  subphase_comm_[cur_phase_].resize(cur_subphase_ + 1);
+  subphase_comm_[cur_phase_].at(cur_subphase_)[key].sendMsg(bytes);
+}
+
 void ElementLBData::sendComm(elm::CommKey key, double bytes) {
   phase_comm_[cur_phase_][key].sendMsg(bytes);
   subphase_comm_[cur_phase_].resize(cur_subphase_ + 1);

src/vt/elm/elm_lb_data.h

@@ -68,6 +68,9 @@ struct ElementLBData {
   void sendToEntity(ElementIDStruct to, ElementIDStruct from, double bytes);
   void sendComm(elm::CommKey key, double bytes);
 
+  void addWritableSharedID(NodeType home, int shared_id, double bytes);
+  void addReadOnlySharedID(NodeType home, int shared_id, double bytes);
+
   void recvComm(elm::CommKey key, double bytes);
   void recvObjData(
     ElementIDStruct to_perm,

src/vt/messaging/active.h

@@ -1722,6 +1722,19 @@ struct ActiveMessenger : runtime::component::PollableComponent<ActiveMessenger>
     MsgSizeType const msg_size
   );
 
+public:
+  /**
+   * \brief Get the rank-based LB data along with element ID for rank-based work
+   *
+   * \return tuple with pointers to each one
+   */
+  auto getRankLBData() {
+    return std::make_tuple(
+      &bare_handler_dummy_elm_id_for_lb_data_,
+      &bare_handler_lb_data_
+    );
+  }
+
 private:
 # if vt_check_enabled(trace_enabled)
   trace::UserEventIDType trace_irecv             = trace::no_user_event_id;

src/vt/vrt/collection/balance/lb_data_holder.cc

@@ -209,6 +209,22 @@ std::unique_ptr<nlohmann::json> LBDataHolder::toJson(PhaseType phase) const {
         outputEntity(j["communications"][i]["from"], key.fromObj());
         break;
       }
+      case elm::CommCategory::ReadOnlyShared: {
+        j["communications"][i]["type"] = "ReadOnlyShared";
+        j["communications"][i]["to"]["type"] = "node";
+        j["communications"][i]["to"]["id"] = key.toNode();
+        j["communications"][i]["from"]["type"] = "shared_id";
+        j["communications"][i]["from"]["id"] = key.sharedID();
+        break;
+      }
+      case elm::CommCategory::WriteShared: {
+        j["communications"][i]["type"] = "WriteShared";
+        j["communications"][i]["to"]["type"] = "node";
+        j["communications"][i]["to"]["id"] = key.toNode();
+        j["communications"][i]["from"]["type"] = "shared_id";
+        j["communications"][i]["from"]["id"] = key.sharedID();
+        break;
+      }
       case elm::CommCategory::LocalInvoke:
       case elm::CommCategory::CollectiveToCollectionBcast:
         // not currently supported
@@ -379,6 +395,34 @@ LBDataHolder::LBDataHolder(nlohmann::json const& j)
               );
               CommVolume vol{bytes, messages};
               this->node_comm_[id][key] = vol;
+            } else if (
+              type == "ReadOnlyShared" or type == "WriteShared"
+            ) {
+              vtAssertExpr(comm["from"]["type"] == "shared_id");
+              vtAssertExpr(comm["to"]["type"] == "node");
+
+              CommVolume vol{bytes, messages};
+              auto to_node = comm["to"]["id"];
+              vtAssertExpr(to_node.is_number());
+
+              auto from_shared_id = comm["from"]["id"];
+              vtAssertExpr(from_shared_id.is_number());
+
+              if (type == "ReadOnlyShared") {
+                CommKey key(
+                  CommKey::ReadOnlySharedTag{},
+                  static_cast<NodeType>(to_node),
+                  static_cast<int>(from_shared_id)
+                );
+                this->node_comm_[id][key] = vol;
+              } else {
+                CommKey key(
+                  CommKey::WriteSharedTag{},
+                  static_cast<NodeType>(to_node),
+                  static_cast<int>(from_shared_id)
+                );
+                this->node_comm_[id][key] = vol;
+              }
             }
           }
         }

src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc

@@ -130,7 +130,11 @@ LBType LBManager::decideLBToRun(PhaseType phase, bool try_file) {
   } else {
     auto interval = theConfig()->vt_lb_interval;
     vtAssert(interval != 0, "LB Interval must not be 0");
-    if (phase % interval == 1 || (interval == 1 && phase != 0)) {
+    vt::PhaseType offset = theConfig()->vt_lb_run_lb_first_phase ? 0 : 1;
+    if (
+      phase % interval == offset ||
+      (interval == 1 && phase != 0)
+    ) {
       bool name_match = false;
       for (auto&& elm : get_lb_names()) {
         if (elm.second == theConfig()->vt_lb_name) {

src/vt/vrt/collection/balance/model/composed_model.cc

@@ -102,4 +102,8 @@ int ComposedModel::getNumSubphases() const {
   return base_->getNumSubphases();
 }
 
+CommMapType ComposedModel::getComm(PhaseOffset when) const {
+  return base_->getComm(when);
+}
+
 }}}}

src/vt/vrt/collection/balance/model/composed_model.h

@@ -77,6 +77,7 @@ class ComposedModel : public LoadModel
   bool hasUserData() const override;
   ElmUserDataType getUserData(ElementIDStruct object, PhaseOffset when) const override;
   unsigned int getNumPastPhasesNeeded(unsigned int look_back) const override;
+  CommMapType getComm(PhaseOffset offset) const override;
 
   ObjectIterator begin() const override;
 

src/vt/vrt/collection/balance/model/load_model.h

@@ -215,6 +215,17 @@ struct LoadModel
    */
   virtual void updateLoads(PhaseType last_completed_phase) = 0;
 
+  /**
+   * \brief Provide all the comm info for a given phase
+   *
+   * \param[in] when the interval in which comm is desired
+   *
+   * \return the comm info
+   */
+  virtual CommMapType getComm(PhaseOffset when) const {
+    return CommMapType{};
+  }
+
   /**
    * \brief Provide an estimate of the given object's load during a specified interval
    *

src/vt/vrt/collection/balance/model/raw_data.cc

@@ -130,6 +130,15 @@ ElmUserDataType RawData::getUserData(ElementIDStruct object, PhaseOffset offset)
   }
 }
 
+CommMapType RawData::getComm(PhaseOffset offset) const {
+  auto phase = getNumCompletedPhases() + offset.phases;
+  if (auto it = proc_comm_->find(phase); it != proc_comm_->end()) {
+    return it->second;
+  } else {
+    return CommMapType{};
+  }
+}
+
 unsigned int RawData::getNumPastPhasesNeeded(unsigned int look_back) const
 {
   return look_back;

src/vt/vrt/collection/balance/model/raw_data.h

@@ -64,6 +64,7 @@ struct RawData : public LoadModel {
   LoadType getRawLoad(ElementIDStruct object, PhaseOffset when) const override;
   bool hasUserData() const override { return user_data_ != nullptr; }
   ElmUserDataType getUserData(ElementIDStruct object, PhaseOffset when) const override;
+  CommMapType getComm(PhaseOffset when) const override;
 
   void setLoads(std::unordered_map<PhaseType, LoadMapType> const* proc_load,
                 std::unordered_map<PhaseType, CommMapType> const* proc_comm,

src/vt/vrt/collection/balance/node_lb_data.cc

@@ -261,6 +261,9 @@ getRecvSendDirection(elm::CommKeyType const& comm) {
   case elm::CommCategory::CollectiveToCollectionBcast:
   case elm::CommCategory::LocalInvoke:
     return std::make_pair(ElementIDType{}, ElementIDType{});
+  case elm::CommCategory::WriteShared:
+  case elm::CommCategory::ReadOnlyShared:
+    return std::make_pair(ElementIDType{}, ElementIDType{});
   }
 
   vtAssert(false, "Invalid balance::CommCategory enum value");

src/vt/vrt/collection/balance/temperedlb/tempered_enums.h

@@ -71,6 +71,36 @@ enum struct InformTypeEnum : uint8_t {
   AsyncInform = 1
 };
 
+/// Enum for the strategy to be used in transfer stage
+enum struct TransferTypeEnum : uint8_t {
+  /**
+   * \brief Original strategy
+   *
+   * Transfer one object per transfer as in original Grapevine approach.
+   */
+  Original   = 0,
+  /**
+   * \brief Original strategy improved by recursion
+   *
+   * When single object transfer is rejected, attempt to recurse in order to
+   * pull more objects into the transfer and hereby minimize work added by
+   * said transfer.
+   * This is especially useful when communication is taken into account, as
+   * object transfers typically disrupt local vs. global communication edges.
+   */
+  Recursive  = 1,
+  /**
+   * \brief Form object clusters and attempt to perform swaps.
+   *
+   * Object can be clustered including to arbitrary definition, and swaps
+   * of entire clusters, including the nullset, between ranks are attempted.
+   * This is especially useful when shared memory constraints are present,
+   * as breaking shared memory clusters results in higher overall memory
+   * footprint, in constrast with whole cluster swaps.
+   */
+  SwapClusters = 2,
+};
+
 /// Enum for the order in which local objects are considered for transfer
 enum struct ObjectOrderEnum : uint8_t {
   Arbitrary = 0, //< Arbitrary order: iterate as defined by the unordered_map