warpx.roundrobin_sfc: A runtime parameter to control distribution map…

…ping

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

@@ -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.

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,17 @@ 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) {
+        auto old_strategy = amrex::DistributionMapping::strategy();
+        amrex::DistributionMapping::strategy(amrex::DistributionMapping::RRSFC);
+        amrex::DistributionMapping dm(ba);
+        amrex::DistributionMapping::strategy(old_strategy);
+        return dm;
+    } else {
+        return amrex::AmrCore::MakeDistributionMap(lev, ba);
+    }
+}