#1934: Refactor trimLBDataHistory to have better performance

DARMA-tasking · Jan 24, 2023 · d074668 · d074668
1 parent ff1b4ee
commit d074668
Showing 1 changed file with 11 additions and 6 deletions.
diff --git a/src/vt/vrt/collection/balance/node_lb_data.cc b/src/vt/vrt/collection/balance/node_lb_data.cc
@@ -131,12 +131,17 @@ void NodeLBData::startIterCleanup(PhaseType phase) {
 
 void NodeLBData::trimLBDataHistory(PhaseType phase) {
   if (phase != no_lb_phase && phase > min_hist_lb_data_) {
-    // Clean phases which are no longer needed
-    for(auto i = 0ull; i < phase - min_hist_lb_data_; i++) {
-      lb_data_->node_data_.erase(i);
-      lb_data_->node_comm_.erase(i);
-      lb_data_->node_subphase_comm_.erase(i);
-    }
+    auto phaseToLookFor = phase - min_hist_lb_data_ - 1;
+
+    // Trim all maps to have 'min_hist_lb_data_ + 1' phases of data
+    auto fromData = lb_data_->node_data_.find(phaseToLookFor);
+    lb_data_->node_data_.erase(fromData, lb_data_->node_data_.end());
+
+    auto fromComm = lb_data_->node_comm_.find(phaseToLookFor);
+    lb_data_->node_comm_.erase(fromComm, lb_data_->node_comm_.end());
+
+    auto fromSub = lb_data_->node_subphase_comm_.find(phaseToLookFor);
+    lb_data_->node_subphase_comm_.erase(fromSub, lb_data_->node_subphase_comm_.end());
 
     NodeLBData::node_migrate_.clear();
     node_collection_lookup_.clear();