diff --git a/src/vt/vrt/collection/balance/elm_stats.cc b/src/vt/vrt/collection/balance/elm_stats.cc index 01ad535f12..68e5d8647b 100644 --- a/src/vt/vrt/collection/balance/elm_stats.cc +++ b/src/vt/vrt/collection/balance/elm_stats.cc @@ -202,6 +202,15 @@ SubphaseType ElementStats::getSubPhase() const { return cur_subphase_; } +void ElementStats::startIterCleanup(PhaseType phase, unsigned int look_back) { + if (phase >= look_back) { + phase_timings_.erase(phase - look_back); + subphase_timings_.erase(phase - look_back); + comm_.erase(phase - look_back); + subphase_comm_.erase(phase - look_back); + } +} + std::size_t ElementStats::getLoadPhaseCount() const { return phase_timings_.size(); } diff --git a/src/vt/vrt/collection/balance/elm_stats.h b/src/vt/vrt/collection/balance/elm_stats.h index 793570e505..d33c4c27d0 100644 --- a/src/vt/vrt/collection/balance/elm_stats.h +++ b/src/vt/vrt/collection/balance/elm_stats.h @@ -92,6 +92,11 @@ struct ElementStats { void setSubPhase(SubphaseType subphase); SubphaseType getSubPhase() const; + /** + * \internal \brief Cleanup after LB runs + */ + void startIterCleanup(PhaseType phase, unsigned int look_back); + // these are just for unit testing std::size_t getLoadPhaseCount() const; std::size_t getCommPhaseCount() const; diff --git a/src/vt/vrt/collection/balance/elm_stats.impl.h b/src/vt/vrt/collection/balance/elm_stats.impl.h index 4e15927962..11e3366b28 100644 --- a/src/vt/vrt/collection/balance/elm_stats.impl.h +++ b/src/vt/vrt/collection/balance/elm_stats.impl.h @@ -95,6 +95,9 @@ void ElementStats::syncNextPhase(CollectStatsMsg* msg, ColT* col) { theNodeStats()->addNodeStats( col, cur_phase, total_load, subphase_loads, comm, subphase_comm ); + + auto model = theLBManager()->getLoadModel(); + stats.startIterCleanup(cur_phase, model->getNumPastPhasesNeeded()); } }}}} /* end namespace vt::vrt::collection::balance */