#5: LB: add work model calculations

Author: lifflander

src/vt-lb/algo/temperedlb/temperedlb.h

@@ -55,6 +55,7 @@
 #include <random>
 #include <ostream>
 #include <fstream>
+#include <cassert>
 
 #include <mpi.h>
 
@@ -70,6 +71,17 @@ struct WorkModel {
   /// @brief  Coefficient for shared-memory communication component
   double delta = 0.0;
 
+  /// @brief Whether memory information is available
+  bool has_memory_info = true;
+  /// @brief Has task serialized memory info
+  bool has_task_serialized_memory_info = true;
+  /// @brief Has task working memory info
+  bool has_task_working_memory_info = true;
+  /// @brief Has task footprint memory info
+  bool has_task_footprint_memory_info = true;
+  /// @brief Has shared block memory info
+  bool has_shared_block_memory_info = true;
+
   double applyWorkFormula(
     double compute, double inter_comm_bytes, double intra_comm_bytes,
     double shared_comm_bytes
@@ -90,6 +102,20 @@ struct Configuration {
     k_max_ = std::ceil(std::sqrt(std::log(num_ranks)/std::log(2.0)));
   }
 
+  bool hasMemoryInfo() const { return work_model_.has_memory_info; }
+  bool hasTaskSerializedMemoryInfo() const {
+    return hasMemoryInfo() && work_model_.has_task_serialized_memory_info;
+  }
+  bool hasTaskWorkingMemoryInfo() const {
+    return hasMemoryInfo() && work_model_.has_task_working_memory_info;
+  }
+  bool hasTaskFootprintMemoryInfo() const {
+    return hasMemoryInfo() && work_model_.has_task_footprint_memory_info;
+  }
+  bool hasSharedBlockMemoryInfo() const {
+    return hasMemoryInfo() && work_model_.has_shared_block_memory_info;
+  }
+
   /// @brief  Number of trials to perform
   int num_trials_ = 1;
   /// @brief  Number of iterations per trial
@@ -249,6 +275,31 @@ struct TaskClusterSummaryInfo {
   model::BytesType max_object_serialized_bytes = 0;
   model::BytesType max_object_serialized_bytes_outside = 0;
   model::BytesType cluster_footprint = 0;
+
+  template <typename SerializerT>
+  void serializer(SerializerT& s) {
+    s | cluster_id;
+    s | num_tasks_;
+    s | cluster_load;
+    s | cluster_intra_send_bytes;
+    s | cluster_intra_recv_bytes;
+    s | inter_edges_;
+    s | shared_block_bytes_;
+    s | max_object_working_bytes;
+    s | max_object_working_bytes_outside;
+    s | max_object_serialized_bytes;
+    s | max_object_serialized_bytes_outside;
+    s | cluster_footprint;
+  }
+};
+
+struct WorkBreakdown {
+  double compute = 0.0;
+  double inter_node_recv_comm = 0.0;
+  double inter_node_send_comm = 0.0;
+  double intra_node_recv_comm = 0.0;
+  double intra_node_send_comm = 0.0;
+  double shared_mem_comm = 0.0;
 };
 
 template <typename CommT>
@@ -347,6 +398,102 @@ struct TemperedLB : baselb::BaseLB {
 
   Clusterer const* getClusterer() const { return clusterer_.get(); }
 
+private:
+  WorkBreakdown computeWorkBreakdown() const {
+    WorkBreakdown breakdown;
+    std::unordered_set<model::SharedBlockType> shared_blocks_here;
+
+    // Rank-alpha term
+    for (auto const& [id, task] : this->getPhaseData().getTasksMap()) {
+      breakdown.compute += task.getLoad();
+      for (auto const& sb : task.getSharedBlocks()) {
+        shared_blocks_here.insert(sb);
+      }
+    }
+
+    // Communication terms
+    for (auto const& e : this->getPhaseData().getCommunications()) {
+      assert(
+        (e.getFromRank() == comm_.getRank() || e.getToRank() == comm_.getRank()) &&
+        "Edge does not belong to this rank"
+      );
+      if (e.getFromRank() != e.getToRank()) {
+        if (e.getToRank() == comm_.getRank()) {
+          breakdown.inter_node_recv_comm += e.getVolume();
+        } else {
+          breakdown.inter_node_send_comm += e.getVolume();
+        }
+      } else {
+        if (e.getToRank() == comm_.getRank()) {
+          breakdown.intra_node_recv_comm += e.getVolume();
+        } else {
+          breakdown.intra_node_send_comm += e.getVolume();
+        }
+      }
+    }
+
+    // Shared-memory communication term
+    for (auto const& sb : shared_blocks_here) {
+      assert(getPhaseData().hasSharedBlock(sb) && "Shared block information missing");
+      auto info = getPhaseData().getSharedBlock(sb);
+      if (info->getHome() != comm_.getRank()) {
+        breakdown.shared_mem_comm += info->getSize();
+      }
+    }
+
+    return breakdown;
+  }
+
+  double computeWork(WorkBreakdown breakdown) const {
+    return config_.work_model_.applyWorkFormula(
+      breakdown.compute,
+      std::max(breakdown.inter_node_recv_comm, breakdown.inter_node_send_comm),
+      std::max(breakdown.intra_node_recv_comm, breakdown.intra_node_send_comm),
+      breakdown.shared_mem_comm
+    );
+  }
+
+  double computeMemoryUsage() const {
+    if (!config_.hasMemoryInfo()) {
+      return 0.0;
+    }
+
+    double task_footprint_bytes_ = 0.0;
+    double task_max_working_bytes_ = 0.0;
+    double task_max_serialized_bytes_ = 0.0;
+    double shared_blocks_bytes_ = 0.0;
+    std::unordered_set<model::SharedBlockType> shared_blocks_here;
+    for (auto const& [id, task] : this->getPhaseData().getTasksMap()) {
+      if (config_.hasTaskFootprintMemoryInfo()) {
+        task_footprint_bytes_ += task.getMemory().footprint_bytes;
+      }
+      if (config_.hasTaskWorkingMemoryInfo()) {
+        task_max_working_bytes_ = std::max(
+          task_max_working_bytes_, task.getMemory().working_bytes
+        );
+      }
+      if (config_.hasTaskSerializedMemoryInfo()) {
+        task_max_serialized_bytes_ = std::max(
+          task_max_serialized_bytes_, task.getMemory().serialized_bytes
+        );
+      }
+      if (config_.hasSharedBlockMemoryInfo()) {
+        for (auto const& sb : task.getSharedBlocks()) {
+          shared_blocks_here.insert(sb);
+        }
+      }
+    }
+    for (auto const& sb : shared_blocks_here) {
+      assert(getPhaseData().hasSharedBlock(sb) && "Shared block information missing");
+      auto info = getPhaseData().getSharedBlock(sb);
+      shared_blocks_bytes_ += info->getSize();
+    }
+    return this->getPhaseData().getRankFootprintBytes() +
+      task_footprint_bytes_ +
+      task_max_working_bytes_ +
+      shared_blocks_bytes_;
+  }
+
 private:
   void computeGlobalMaxClusters() {
     // compute max number of clusters on any rank

src/vt-lb/model/PhaseData.h

@@ -82,6 +82,12 @@ struct PhaseData {
   std::vector<Edge> const& getCommunications() const { return communications_; }
   std::unordered_map<SharedBlockType, SharedBlock> const& getSharedBlocksMap() const { return shared_blocks_; }
 
+  BytesType getRankFootprintBytes() const { return rank_footprint_bytes_; }
+  void setRankFootprintBytes(BytesType bytes) { rank_footprint_bytes_ = bytes; }
+
+  BytesType getRankMaxMemoryAvailable() const { return rank_max_memory_available_; }
+  void setRankMaxMemoryAvailable(BytesType bytes) { rank_max_memory_available_ = bytes; }
+
   void clear() {
     tasks_.clear();
     communications_.clear();
@@ -94,13 +100,17 @@ struct PhaseData {
     s | tasks_;
     s | communications_;
     s | shared_blocks_;
+    s | rank_footprint_bytes_;
+    s | rank_max_memory_available_;
   }
 
 private:
   RankType rank_ = invalid_node;
   std::unordered_map<TaskType, Task> tasks_;
   std::vector<Edge> communications_;
   std::unordered_map<SharedBlockType, SharedBlock> shared_blocks_;
+  BytesType rank_footprint_bytes_ = 0.0;
+  BytesType rank_max_memory_available_ = 0.0;
 };
 
 } /* end namespace vt_lb::model */

src/vt-lb/model/Task.h

@@ -51,74 +51,74 @@
 namespace vt_lb::model {
 
 struct TaskMemory {
-    TaskMemory() = default;
-    TaskMemory(BytesType working, BytesType footprint, BytesType serialized)
-      : working_(working), footprint_(footprint), serialized_(serialized)
-    {}
-
-    BytesType getWorking() const { return working_; }
-    BytesType getFootprint() const { return footprint_; }
-    BytesType getSerialized() const { return serialized_; }
-
-    template <typename Serializer>
-    void serialize(Serializer& s) {
-      s | working_;
-      s | footprint_;
-      s | serialized_;
-    }
+  TaskMemory() = default;
+  TaskMemory(BytesType working, BytesType footprint, BytesType serialized)
+    : working_(working), footprint_(footprint), serialized_(serialized)
+  {}
+
+  BytesType getWorking() const { return working_; }
+  BytesType getFootprint() const { return footprint_; }
+  BytesType getSerialized() const { return serialized_; }
+
+  template <typename Serializer>
+  void serialize(Serializer& s) {
+    s | working_;
+    s | footprint_;
+    s | serialized_;
+  }
 
 private:
-    BytesType working_ = 0.0;
-    BytesType footprint_ = 0.0;
-    BytesType serialized_ = 0.0;
+  BytesType working_ = 0.0;
+  BytesType footprint_ = 0.0;
+  BytesType serialized_ = 0.0;
 };
 
 struct Task {
-    Task() = default;
-    Task(TaskType id, RankType home, RankType current, bool migratable,
-         TaskMemory const& memory, LoadType load)
-      : id_(id),
-        home_(home),
-        current_(current),
-        migratable_(migratable),
-        memory_(memory),
-        load_(load)
-    {}
-
-    TaskType getId() const { return id_; }
-    RankType getHome() const { return home_; }
-    RankType getCurrent() const { return current_; }
-    bool isMigratable() const { return migratable_; }
-    TaskMemory const& getMemory() const { return memory_; }
-    LoadType getLoad() const { return load_; }
-
-    // Add accessors for shared blocks
-    void addSharedBlock(SharedBlockType sb) { shared_blocks_.insert(sb); }
-    std::unordered_set<SharedBlockType> const& getSharedBlocks() const { return shared_blocks_; }
-
-    template <typename Serializer>
-    void serialize(Serializer& s) {
-      s | id_;
-      s | home_;
-      s | current_;
-      s | migratable_;
-      s | memory_;
-      s | load_;
-      s | shared_blocks_;
-    }
+  Task() = default;
+  Task(TaskType id, RankType home, RankType current, bool migratable,
+        TaskMemory const& memory, LoadType load)
+    : id_(id),
+      home_(home),
+      current_(current),
+      migratable_(migratable),
+      memory_(memory),
+      load_(load)
+  {}
+
+  TaskType getId() const { return id_; }
+  RankType getHome() const { return home_; }
+  RankType getCurrent() const { return current_; }
+  bool isMigratable() const { return migratable_; }
+  TaskMemory const& getMemory() const { return memory_; }
+  LoadType getLoad() const { return load_; }
+
+  // Add accessors for shared blocks
+  void addSharedBlock(SharedBlockType sb) { shared_blocks_.insert(sb); }
+  std::unordered_set<SharedBlockType> const& getSharedBlocks() const { return shared_blocks_; }
+
+  template <typename Serializer>
+  void serialize(Serializer& s) {
+    s | id_;
+    s | home_;
+    s | current_;
+    s | migratable_;
+    s | memory_;
+    s | load_;
+    s | shared_blocks_;
+  }
 
 private:
-    TaskType id_ = invalid_task;
-    int home_ = invalid_node;
-    int current_ = invalid_node;
-    bool migratable_ = true;
-    TaskMemory memory_;
-    LoadType load_ = 0.0;
-    std::unordered_set<SharedBlockType> shared_blocks_;
+  TaskType id_ = invalid_task;
+  int home_ = invalid_node;
+  int current_ = invalid_node;
+  bool migratable_ = true;
+  TaskMemory memory_;
+  LoadType load_ = 0.0;
+  std::unordered_set<SharedBlockType> shared_blocks_;
 
 public:
-    bool operator==(const Task& other) const { return id_ == other.id_; }
-    bool operator!=(const Task& other) const { return !(*this == other); }
+  bool operator==(const Task& other) const { return id_ == other.id_; }
+  bool operator!=(const Task& other) const { return !(*this == other); }
 };
 
 } /* end namespace vt_lb::model */