#2201: temperedlb: add computation for cluster/memory summary

DARMA-tasking · Nov 29, 2023 · 690134a · 690134a
1 parent 0470b12
commit 690134a
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 7 deletions.
diff --git a/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc b/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc
@@ -566,14 +566,28 @@ void TemperedLB::readClustersMemoryData() {
         obj, shared_id, shared_bytes
       );
 
+      has_memory_data_ = true;
       obj_shared_block_[obj] = shared_id;
       obj_working_bytes_[obj] = working_bytes;
       shared_block_size_[shared_id] = shared_bytes;
-      has_memory_data_ = true;
     }
   }
 }
 
+void TemperedLB::computeClusterSummary() {
+  for (auto const& [shared_id, shared_bytes] : shared_block_size_) {
+    LoadType cluster_load = 0;
+    for (auto const& [obj_id, obj_load] : cur_objs_) {
+      if (auto iter = obj_shared_block_.find(obj_id); iter != obj_shared_block_.end()) {
+        if (iter->second == shared_id) {
+          cluster_load += obj_load;
+        }
+      }
+    }
+    cur_blocks_[shared_id] = std::make_tuple(shared_bytes, cluster_load);
+  }
+}
+
 TemperedLB::BytesType TemperedLB::computeMemoryUsage() const {
   // Compute bytes used by shared blocks mapped here based on object mapping
   auto const blocks_here = getSharedBlocksHere();
@@ -659,12 +673,25 @@ void TemperedLB::doLBStages(LoadType start_imb) {
         LoadType(this_new_load_)
       );
 
-      vt_print(
-        temperedlb,
-        "Current memory info: total memory usage={}, shared blocks here={}, "
-        "memory_threshold={}\n", computeMemoryUsage(), getSharedBlocksHere().size(),
-        mem_thresh_
-      );
+      if (has_memory_data_) {
+        vt_print(
+          temperedlb,
+          "Current memory info: total memory usage={}, shared blocks here={}, "
+          "memory_threshold={}\n", computeMemoryUsage(),
+          getSharedBlocksHere().size(), mem_thresh_
+        );
+
+        computeClusterSummary();
+
+        for (auto const& [shared_id, value] : cur_blocks_) {
+          auto const& [shared_bytes, cluster_load] = value;
+          vt_print(
+            temperedlb,
+            "Cluster: id={}, bytes={}, load={}\n",
+            shared_id, shared_bytes, cluster_load
+          );
+        }
+      }
 
       if (isOverloaded(this_new_load_)) {
         is_overloaded_ = true;

diff --git a/src/vt/vrt/collection/balance/temperedlb/temperedlb.h b/src/vt/vrt/collection/balance/temperedlb/temperedlb.h
@@ -143,6 +143,11 @@ struct TemperedLB : BaseLB {
    */
   std::set<SharedIDType> getSharedBlocksHere() const;
 
+  /**
+   * \brief Compute the current cluster assignment summary for this rank
+   */
+  void computeClusterSummary();
+
 private:
   uint16_t f_                                       = 0;
   uint8_t k_max_                                    = 0;
@@ -223,6 +228,8 @@ struct TemperedLB : BaseLB {
   std::unordered_map<SharedIDType, BytesType> shared_block_size_;
   /// Working bytes for each object
   std::unordered_map<ObjIDType, BytesType> obj_working_bytes_;
+  /// Current assignment memory/load summary
+  std::unordered_map<SharedIDType, std::tuple<BytesType, LoadType>> cur_blocks_;
   /// User-defined memory threshold
   BytesType mem_thresh_ = 0;
 };