#2074: Update finding next phase for OfflineLB

DARMA-tasking · May 16, 2023 · ca02e10 · ca02e10
1 parent 8984e5a
commit ca02e10
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 14 deletions.
diff --git a/src/vt/vrt/collection/balance/lb_data_restart_reader.h b/src/vt/vrt/collection/balance/lb_data_restart_reader.h
@@ -123,6 +123,24 @@ struct LBDataRestartReader : runtime::component::Component<LBDataRestartReader>
       | num_phases_;
   }
 
+  /**
+   * \brief Find the next phase
+   *
+   * \param phase the current phase
+   *
+   * \return the next phase
+   */
+  PhaseType findNextPhase(PhaseType phase) const {
+    auto next = phase + 1;
+    for(; next < num_phases_; ++next) {
+      if(history_.find(next) != history_.end()) {
+        return next;
+      }
+    }
+    vtAssert(history_.find(next) != history_.end(), "Must have a valid phase");
+    return next;
+  }
+
   /**
    * \brief Get the elements assigned for a given phase
    *
@@ -131,12 +149,6 @@ struct LBDataRestartReader : runtime::component::Component<LBDataRestartReader>
    * \return element assigned to this node
    */
   std::set<ElementIDStruct> const& getDistro(PhaseType phase) {
-    for(; phase < num_phases_; ++phase) {
-      if(history_.find(phase) != history_.end()) {
-        break;
-      }
-    }
-
     auto iter = history_.find(phase);
     vtAssert(iter != history_.end(), "Must have a valid phase");
     return iter->second;
@@ -148,12 +160,6 @@ struct LBDataRestartReader : runtime::component::Component<LBDataRestartReader>
    * \param[in] phase the phase to clear
    */
   void clearDistro(PhaseType phase) {
-    for(; phase < num_phases_; ++phase) {
-      if(history_.find(phase) != history_.end()) {
-        break;
-      }
-    }
-
     auto iter = history_.find(phase);
     if (iter != history_.end()) {
       history_.erase(iter);

diff --git a/src/vt/vrt/collection/balance/offlinelb/offlinelb.cc b/src/vt/vrt/collection/balance/offlinelb/offlinelb.cc
@@ -55,11 +55,12 @@ void OfflineLB::init(objgroup::proxy::Proxy<OfflineLB> in_proxy) {
 }
 
 void OfflineLB::runLB(TimeType) {
-  auto const& distro = theLBDataReader()->getDistro(phase_ + 1);
+  auto nextPhase = theLBDataReader()->findNextPhase(phase_);
+  auto const& distro = theLBDataReader()->getDistro(nextPhase);
   for (auto&& elm : distro) {
     migrateObjectTo(elm, theContext()->getNode());
   }
-  theLBDataReader()->clearDistro(phase_ + 1);
+  theLBDataReader()->clearDistro(nextPhase);
 }
 
 }}}} /* end namespace vt::vrt::collection::lb */