1 files changed, 89 insertions, 7 deletions

diff --git a/Source/Parallelization/WarpXRegrid.cpp b/Source/Parallelization/WarpXRegrid.cpp
index 5868642f8..a86256079 100644
--- a/Source/Parallelization/WarpXRegrid.cpp
+++ b/Source/Parallelization/WarpXRegrid.cpp
@@ -7,6 +7,7 @@
  * License: BSD-3-Clause-LBNL
  */
 #include <WarpX.H>
+#include <WarpXAlgorithmSelection.H>
 #include <AMReX_BLProfiler.H>
 
 using namespace amrex;
@@ -17,6 +18,18 @@ WarpX::LoadBalance ()
     WARPX_PROFILE_REGION("LoadBalance");
     WARPX_PROFILE("WarpX::LoadBalance()");
 
+    if (WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Timers)
+    {
+        LoadBalanceTimers();
+    } else if (WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Heuristic)
+    {
+        LoadBalanceHeuristic();
+    }
+}
+
+void
+WarpX::LoadBalanceTimers ()
+{
     AMREX_ALWAYS_ASSERT(costs[0] != nullptr);
 
     const int nLevels = finestLevel();
@@ -30,7 +43,36 @@ WarpX::LoadBalance ()
             : DistributionMapping::makeKnapSack(*costs[lev], nmax);
         RemakeLevel(lev, t_new[lev], boxArray(lev), newdm);
     }
+    mypc->Redistribute();
+}
+
+void
+WarpX::LoadBalanceHeuristic ()
+{
+    AMREX_ALWAYS_ASSERT(costs_heuristic[0] != nullptr);
+    WarpX::ComputeCostsHeuristic();
+
+    const int nLevels = finestLevel();
+    for (int lev = 0; lev <= nLevels; ++lev)
+    {
+#ifdef AMREX_USE_MPI
+        // Parallel reduce the costs_heurisitc
+        amrex::Vector<Real>::iterator it = (*costs_heuristic[lev]).begin();
+        amrex::Real* itAddr = &(*it);
+        ParallelAllReduce::Sum(itAddr,
+                               costs_heuristic[lev]->size(),
+                               ParallelContext::CommunicatorSub());
+#endif
+        const amrex::Real nboxes = costs_heuristic[lev]->size();
+        const amrex::Real nprocs = ParallelContext::NProcsSub();
+        const int nmax = static_cast<int>(std::ceil(nboxes/nprocs*load_balance_knapsack_factor));
 
+        const DistributionMapping newdm = (load_balance_with_sfc)
+            ? DistributionMapping::makeSFC(*costs_heuristic[lev], boxArray(lev), false)
+            : DistributionMapping::makeKnapSack(*costs_heuristic[lev], nmax);
+
+        RemakeLevel(lev, t_new[lev], boxArray(lev), newdm);
+    }
     mypc->Redistribute();
 }
 
@@ -47,7 +89,6 @@ WarpX::RemakeLevel (int lev, Real /*time*/, const BoxArray& ba, const Distributi
 #endif // WARPX_DO_ELECTROSTATIC
 
         // Fine patch
-
         const auto& period = Geom(lev).periodicity();
         for (int idim=0; idim < 3; ++idim)
         {
@@ -98,7 +139,6 @@ WarpX::RemakeLevel (int lev, Real /*time*/, const BoxArray& ba, const Distributi
         }
 
         // Aux patch
-
         if (lev == 0 && Bfield_aux[0][0]->ixType() == Bfield_fp[0][0]->ixType())
         {
             for (int idim = 0; idim < 3; ++idim) {
@@ -223,15 +263,57 @@ WarpX::RemakeLevel (int lev, Real /*time*/, const BoxArray& ba, const Distributi
             }
         }
 
-        if (costs[lev] != nullptr) {
-            costs[lev].reset(new MultiFab(costs[lev]->boxArray(), dm, 1, 0));
-            costs[lev]->setVal(0.0);
+        if (WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Timers)
+        {
+            if (costs[lev] != nullptr)
+            {
+                costs[lev].reset(new MultiFab(costs[lev]->boxArray(), dm, 1, 0));
+                costs[lev]->setVal(0.0);
+            }
+        } else if (WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Heuristic)
+        {
+            if (costs_heuristic[lev] != nullptr)
+            {
+                costs_heuristic[lev].reset(new amrex::Vector<Real>);
+                const int nboxes = Efield_fp[lev][0].get()->size();
+                costs_heuristic[lev]->resize(nboxes, 0.0); // Initializes to 0.0?
+            }
         }
 
         SetDistributionMap(lev, dm);
-    }
-    else
+
+    } else
     {
         amrex::Abort("RemakeLevel: to be implemented");
     }
 }
+
+void
+WarpX::ComputeCostsHeuristic ()
+{
+    for (int lev = 0; lev <= finest_level; ++lev)
+    {
+        auto & mypc = WarpX::GetInstance().GetPartContainer();
+        auto nSpecies = mypc.nSpecies();
+
+        // Species loop
+        for (int i_s = 0; i_s < nSpecies; ++i_s)
+        {
+            auto & myspc = mypc.GetParticleContainer(i_s);
+
+            // Particle loop
+            for (WarpXParIter pti(myspc, lev); pti.isValid(); ++pti)
+            {
+                (*costs_heuristic[lev])[pti.index()] += costs_heuristic_particles_wt*pti.numParticles();
+            }
+        }
+
+        //Cell loop
+        MultiFab* Ex = Efield_fp[lev][0].get();
+        for (MFIter mfi(*Ex, false); mfi.isValid(); ++mfi)
+        {
+            const Box& gbx = mfi.growntilebox();
+            (*costs_heuristic[lev])[mfi.index()] += costs_heuristic_cells_wt*gbx.numPts();
+        }
+    } // for (int lev ...)
+} // WarpX::ComputeCostsHeuristic