1 files changed, 54 insertions, 7 deletions

diff --git a/Source/Parallelization/WarpXRegrid.cpp b/Source/Parallelization/WarpXRegrid.cpp
index 36a1ed776..43a03f5ed 100644
--- a/Source/Parallelization/WarpXRegrid.cpp
+++ b/Source/Parallelization/WarpXRegrid.cpp
@@ -30,22 +30,33 @@ WarpX::LoadBalance ()
     for (int lev = 0; lev <= nLevels; ++lev)
     {
 #ifdef AMREX_USE_MPI
-        // Parallel reduce the costs_heurisitc
+        // Parallel reduce the costs
         amrex::Vector<Real>::iterator it = (*costs[lev]).begin();
         amrex::Real* itAddr = &(*it);
-        ParallelAllReduce::Sum(itAddr,
-                               costs[lev]->size(),
-                               ParallelContext::CommunicatorSub());
+        ParallelAllReduce::Sum(itAddr, costs[lev]->size(), ParallelContext::CommunicatorSub());
 #endif
+        // Compute efficiency for the current distribution mapping
+        const DistributionMapping& currentdm = DistributionMap(lev);
+        amrex::Real currentEfficiency = 0.0;
+        if (load_balance_efficiency_ratio_threshold > 0.0)
+        {
+            ComputeDistributionMappingEfficiency(currentdm, *costs[lev],
+                                                 currentEfficiency);
+        }
+
         const amrex::Real nboxes = costs[lev]->size();
         const amrex::Real nprocs = ParallelContext::NProcsSub();
         const int nmax = static_cast<int>(std::ceil(nboxes/nprocs*load_balance_knapsack_factor));
 
+        amrex::Real proposedEfficiency = 0.0;
         const DistributionMapping newdm = (load_balance_with_sfc)
-            ? DistributionMapping::makeSFC(*costs[lev], boxArray(lev), false)
-            : DistributionMapping::makeKnapSack(*costs[lev], nmax);
+            ? DistributionMapping::makeSFC(*costs[lev], boxArray(lev), proposedEfficiency, false)
+            : DistributionMapping::makeKnapSack(*costs[lev], proposedEfficiency, nmax);
 
-        RemakeLevel(lev, t_new[lev], boxArray(lev), newdm);
+        if (proposedEfficiency > load_balance_efficiency_ratio_threshold*currentEfficiency)
+        {
+            RemakeLevel(lev, t_new[lev], boxArray(lev), newdm);
+        }
     }
     mypc->Redistribute();
 }
@@ -284,3 +295,39 @@ WarpX::ResetCosts ()
         costs[lev]->assign((*costs[lev]).size(), 0.0);
     }
 }
+
+void
+WarpX::ComputeDistributionMappingEfficiency (const DistributionMapping& dm,
+                                             const Vector<Real>& cost,
+                                             Real& efficiency)
+{
+    const Real nprocs = ParallelDescriptor::NProcs();
+
+    // Collect costs per fab corresponding to each rank, then collapse into vector
+    // of total cost per proc
+
+    // This will store mapping from (proc) --> ([cost_FAB_1, cost_FAB_2, ... ])
+    // for each proc
+    std::map<int, Vector<Real>> rankToCosts;
+
+    for (int i=0; i<cost.size(); ++i)
+    {
+        rankToCosts[dm[i]].push_back(cost[i]);
+    }
+
+    Real maxCost = -1.0;
+
+    // This will store mapping from (proc) --> (sum of cost) for each proc
+    Vector<Real> rankToCost = {0.0};
+    rankToCost.resize(nprocs);
+    for (int i=0; i<nprocs; ++i) {
+        const Real rwSum = std::accumulate(rankToCosts[i].begin(),
+                                           rankToCosts[i].end(), 0.0);
+        rankToCost[i] = rwSum;
+        maxCost = std::max(maxCost, rwSum);
+    }
+
+    // `efficiency` is mean cost per proc
+    efficiency = (std::accumulate(rankToCost.begin(),
+                  rankToCost.end(), 0.0) / (nprocs*maxCost));
+}