From 99d951764aadbe860b11f61de0e25ddc6a1e4594 Mon Sep 17 00:00:00 2001
From: Weiqun Zhang <weiqunzhang@lbl.gov>
Date: Thu, 2 May 2024 17:18:54 -0700
Subject: [PATCH] warpx.roundrobin_sfc: A runtime parameter to control
 distribution mapping

The default is false. If it's true, AMReX's RRSFS strategy will be used to
override the default SFC strategy used by amrex::AmrCore. The motivation for
this is that this might mitigate the load imbalance issue during
initialization by avoiding putting neighboring boxes on the same process.
---
 Source/WarpX.H   | 14 ++++++++++++++
 Source/WarpX.cpp | 16 ++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/Source/WarpX.H b/Source/WarpX.H
index d19843ca636..2f9471059a1 100644
--- a/Source/WarpX.H
+++ b/Source/WarpX.H
@@ -1212,6 +1212,20 @@ protected:
     //! This function is called in amrex::AmrCore::InitFromScratch.
     void PostProcessBaseGrids (amrex::BoxArray& ba0) const final;
 
+    //! If this is true, AMReX's RRSFC strategy is used to make
+    //! DistributionMapping in the virtual function MakeDistributionMap. The
+    //! default is false. Note that the DistributionMapping made by the
+    //! MakeDistributionMap function could still be overridden by load
+    //! balancing. In the RRSFC strategy, The Round robin method is used to
+    //! distribute Boxes orderd by the space filling curve. This might help
+    //! avoid some processes running out of memory due to having too many
+    //! particles during initialization.
+    static bool roundrobin_sfc;
+
+    //! Use this function to override how DistributionMapping is made.
+    [[nodiscard]] amrex::DistributionMapping
+    MakeDistributionMap (int lev, amrex::BoxArray const& ba) final;
+
     //! Make a new level from scratch using provided BoxArray and
     //! DistributionMapping.  Only used during initialization.  Called
     //! by AmrCoreInitFromScratch.
diff --git a/Source/WarpX.cpp b/Source/WarpX.cpp
index 9a29d7c2354..d80547fee47 100644
--- a/Source/WarpX.cpp
+++ b/Source/WarpX.cpp
@@ -219,6 +219,8 @@ bool WarpX::do_device_synchronize = true;
 bool WarpX::do_device_synchronize = false;
 #endif
 
+bool WarpX::roundrobin_sfc = false;
+
 WarpX* WarpX::m_instance = nullptr;
 
 void WarpX::MakeWarpX ()
@@ -1157,6 +1159,8 @@ WarpX::ReadParameters ()
             maxLevel() == 0 || !do_current_centering,
             "Finite-order centering of currents is not implemented with mesh refinement"
         );
+
+        pp_warpx.query("roundrobin_sfc", roundrobin_sfc);
     }
 
     {
@@ -3515,3 +3519,15 @@ WarpX::getField(FieldType field_type, const int lev, const int direction) const
 {
     return *getFieldPointer(field_type, lev, direction);
 }
+
+amrex::DistributionMapping
+WarpX::MakeDistributionMap (int lev, amrex::BoxArray const& ba)
+{
+    if (roundrobin_sfc) {
+        amrex::DistributionMapping dm;
+        dm.RRSFCProcessorMap(ba, amrex::ParallelContext::NProcsSub();
+        return dm;
+    } else {
+        return amrex::AmrCore::MakeDistributionMap(lev, ba);
+    }
+}