fix loop pattern, reset lb vars

parthenon-hpc-lab · Nov 3, 2023 · 2c78030 · 2c78030
1 parent 64a023e
commit 2c78030
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/src/mesh/amr_loadbalance.cpp b/src/mesh/amr_loadbalance.cpp
@@ -368,6 +368,7 @@ Real DistributeTrial(std::vector<Real> const &cost, std::vector<int> &start,
 double CalculateNewBalance(std::vector<double> const &cost, std::vector<int> &start,
                            std::vector<int> &nb, const double avg_cost,
                            const double max_block_cost) {
+  PARTHENON_INSTRUMENT
   const int nblocks = cost.size();
   const int max_rank = std::min(nblocks, Globals::nranks);
 
@@ -689,7 +690,10 @@ bool Mesh::RedistributeAndRefineMeshBlocks(ParameterInput *pin, ApplicationInput
       double new_max =
           CalculateNewBalance(costlist, start_trial, nb_trial, avg_cost, max_block_cost);
       // if the improvement isn't large enough, just return because we're done
-      if ((max_rank_cost - new_max) / max_rank_cost < lb_tolerance_) return false;
+      if ((max_rank_cost - new_max) / max_rank_cost < lb_tolerance_) {
+        ResetLoadBalanceVariables();
+        return false;
+      }
       newrank.resize(ntot);
       AssignBlocks(start_trial, nb_trial, newrank);
       nslist = std::move(start_trial);

diff --git a/src/mesh/mesh.cpp b/src/mesh/mesh.cpp
@@ -1147,7 +1147,7 @@ void Mesh::Initialize(bool init_problem, ParameterInput *pin, ApplicationInput *
         block_cost_host[i] = 1.0;
 #ifdef ENABLE_LB_TIMERS
       parthenon::par_for(
-          DEFAULT_LOOP_PATTERN, PARTHENON_AUTO_LABEL, DevExecSpace(), 0, nmb - 1,
+          loop_pattern_flatrange_tag, PARTHENON_AUTO_LABEL, DevExecSpace(), 0, nmb - 1,
           KOKKOS_LAMBDA(const int b) { block_cost(b) = 1.0; });
 #endif
       LoadBalancingAndAdaptiveMeshRefinement(pin, app_in);