Adding bittree interface to improve regridding performance in octree …

…mode (AMReX-Codes#3555) ## Summary This PR introduces dependency on [Bittree library](https://github.com/Flash-X/Bittree) to improve regridding performance in octree mode. ## Additional background Testing and development of this feature is done in sync with Flash-X and is recorded in this [reproducibility capsule](https://github.com/Lab-Notebooks/AMReX-Bittree-Performance). This PR is primarily created to consolidate development work and avoid creating multiple branches. At present using AMReX+Bittree improves regridding performance [by a factor of 2](https://github.com/Lab-Notebooks/AMReX-Bittree-Performance/blob/14faa2212c4e5dba7fd99a6526c6937414f9c109/analysis/Performance.ipynb) at > 20000 ranks. We hope to improve performance further using a bittree-based distribution mapping and therefore adding a new function `AmrMesh::MakeDistributionMap` as a place-holder. Continuation of AMReX-Codes#2893 and AMReX-Codes#3547 ## Checklist The proposed changes: - [ ] fix a bug or incorrect behavior in AMReX - [x] add new capabilities to AMReX - [ ] changes answers in the test suite to more than roundoff level - [ ] are likely to significantly affect the results of downstream AMReX users - [ ] include documentation in the code and/or rst files, if appropriate --------- Co-authored-by: Weiqun Zhang <[email protected]> Co-authored-by: Tom Klosterman <https://github.com/tklos96>
ajnonaka · Oct 1, 2023 · 8515ea8 · 8515ea8
1 parent 703a204
commit 8515ea8
Show file tree

Hide file tree

Showing 14 changed files with 737 additions and 7 deletions.
diff --git a/GNUmakefile.in b/GNUmakefile.in
@@ -45,6 +45,9 @@ endif
 ifeq ($(USE_SUNDIALS),TRUE)
     Pdirs += Extern/SUNDIALS
 endif
+ifeq ($(USE_BITTREE),TRUE)
+	Pdirs += Extern/Bittree
+endif
 Ppack := $(foreach dir, $(Pdirs), $(AMREX_HOME)/Src/$(dir)/Make.package)
 include $(Ppack)
 

diff --git a/Src/AmrCore/AMReX_AmrCore.cpp b/Src/AmrCore/AMReX_AmrCore.cpp
@@ -104,7 +104,7 @@ AmrCore::regrid (int lbase, Real time, bool)
                 DistributionMapping level_dmap = dmap[lev];
                 if (ba_changed) {
                     level_grids = new_grids[lev];
-                    level_dmap = DistributionMapping(level_grids);
+                    level_dmap = MakeDistributionMap(lev, level_grids);
                 }
                 const auto old_num_setdm = num_setdm;
                 RemakeLevel(lev, time, level_grids, level_dmap);
@@ -117,7 +117,7 @@ AmrCore::regrid (int lbase, Real time, bool)
         }
         else  // a new level
         {
-            DistributionMapping new_dmap(new_grids[lev]);
+            DistributionMapping new_dmap = MakeDistributionMap(lev, new_grids[lev]);
             const auto old_num_setdm = num_setdm;
             MakeNewLevelFromCoarse(lev, time, new_grids[lev], new_dmap);
             SetBoxArray(lev, new_grids[lev]);

diff --git a/Src/AmrCore/AMReX_AmrMesh.H b/Src/AmrCore/AMReX_AmrMesh.H
@@ -11,6 +11,10 @@
 #include <AMReX_BoxArray.H>
 #include <AMReX_TagBox.H>
 
+#ifdef AMREX_USE_BITTREE
+#include <Bittree_BittreeAmr.h>
+#endif
+
 namespace amrex {
 
 struct AmrInfo {
@@ -253,13 +257,20 @@ public:
 
     [[nodiscard]] long CountCells (int lev) noexcept;
 
+    [[nodiscard]] virtual DistributionMapping MakeDistributionMap (int lev, BoxArray const& ba);
+
 protected:
 
     int finest_level;    //!< Current finest level.
     Vector<Geometry>            geom;
     Vector<DistributionMapping> dmap;
     Vector<BoxArray>            grids;
 
+#ifdef AMREX_USE_BITTREE
+    bool use_bittree = false;
+    std::unique_ptr<bittree::BittreeAmr> btmesh;
+#endif
+
     unsigned int num_setdm = 0;
     unsigned int num_setba = 0;
 

diff --git a/Src/AmrCore/AMReX_AmrMesh.cpp b/Src/AmrCore/AMReX_AmrMesh.cpp
@@ -5,6 +5,13 @@
 #include <AMReX_ParmParse.H>
 #include <AMReX_ParallelDescriptor.H>
 #include <AMReX_Print.H>
+#include <AMReX_Vector.H>
+
+#ifdef AMREX_USE_BITTREE
+#include <AMReX_Bittree.H>
+#endif
+
+#include <memory>
 
 namespace amrex {
 
@@ -376,6 +383,10 @@ AmrMesh::InitAmrMesh (int max_level_in, const Vector<int>& n_cell_in,
 
     finest_level = -1;
 
+#ifdef AMREX_USE_BITTREE
+    pp.queryAdd("use_bittree",use_bittree);
+#endif
+
     if (check_input) { checkInput(); }
 }
 
@@ -437,6 +448,26 @@ AmrMesh::LevelDefined (int lev) noexcept
     return lev <= max_level && !grids[lev].empty() && !dmap[lev].empty();
 }
 
+DistributionMapping
+AmrMesh::MakeDistributionMap (int lev, BoxArray const& ba)
+{
+
+    BL_PROFILE("AmrMesh::MakeDistributionMap()");
+
+    if (verbose) {
+        amrex::Print() << "Creating new distribution map on level: " << lev << "\n";
+    }
+
+#ifdef AMREX_USE_BITTREE
+    // if (use_bittree) {
+    //     return DistributionMapping(ba);
+    // } else
+#endif
+    {
+        return DistributionMapping(ba);
+    }
+}
+
 void
 AmrMesh::ChopGrids (int lev, BoxArray& ba, int target_size) const
 {
@@ -514,6 +545,10 @@ AmrMesh::MakeNewGrids (int lbase, Real time, int& new_finest, Vector<BoxArray>&
 
     if (new_grids.size() < max_crse+2) { new_grids.resize(max_crse+2); }
 
+#ifdef AMREX_USE_BITTREE
+    if(!use_bittree) {
+#endif
+
     //
     // Construct problem domain at each level.
     //
@@ -774,6 +809,72 @@ AmrMesh::MakeNewGrids (int lbase, Real time, int& new_finest, Vector<BoxArray>&
             }
         }
     }
+
+#ifdef AMREX_USE_BITTREE
+    }
+#endif
+
+#ifdef AMREX_USE_BITTREE
+    // Bittree version
+    if(use_bittree) {
+        // Initialize BT refinement
+        btmesh->refine_init();
+
+        // -------------------------------------------------------------------
+        // Use tagging data to mark BT for refinement, then use the new bitmap
+        // to calculate the new grids.
+        auto tree0 = btmesh->getTree();
+
+        // [1] Error Estimation and tagging
+        // btTags is indexed by bitid, Bittree's internal indexing scheme.
+        // For any id, btTags = 1 if should be parent, -1 if should not be parent (or not exist).
+        std::vector<int> btTags(tree0->id_upper_bound(),0);
+
+        for (int lev=max_crse; lev>=lbase; --lev) {
+
+            TagBoxArray tags(grids[lev],dmap[lev], n_error_buf[lev]);
+            ErrorEst(lev, tags, time, 0);
+            tags.buffer(n_error_buf[lev]);
+
+            for (MFIter mfi(tags); mfi.isValid(); ++mfi) {
+                auto const& tagbox = tags.const_array(mfi);
+                bool has_set_tags = amrex::Reduce::AnyOf(mfi.validbox(),
+                                                         [=] AMREX_GPU_DEVICE (int i, int j, int k)
+                                                         {
+                                                              return tagbox(i,j,k)!=TagBox::CLEAR;
+                                                         });
+
+                // Set the values of btTags.
+                int bitid = btUnit::getBitid(btmesh.get(),false,lev,mfi.index());
+                // TODO Check lev == tree0->block_level(bitid)
+                if(has_set_tags) {
+                    btTags[bitid] = 1;
+                }
+                else {
+                    btTags[bitid] = -1;
+                }
+            }
+        }
+
+        // [2] btRefine - check for proper octree nesting and update bitmap
+        MPI_Comm comm = ParallelContext::CommunicatorSub();
+        int changed = btUnit::btRefine(btmesh.get(), btTags, max_crse, lbase, grids, dmap, comm);
+
+        // [3] btCalculateGrids - use new bitmap to generate new grids
+        if (changed>0) {
+            btUnit::btCalculateGrids(btmesh.get(),lbase,new_finest,new_grids,max_grid_size);
+        } else {
+            new_finest = finest_level;
+            for(int i=0; i<=finest_level; ++i) {
+                new_grids[i] = grids[i];
+            }
+        }
+
+        // Finalize BT refinement
+        btmesh->refine_apply();
+    }
+#endif
+
 }
 
 void
@@ -783,11 +884,48 @@ AmrMesh::MakeNewGrids (Real time)
     {
         finest_level = 0;
 
-        const BoxArray& ba = MakeBaseGrids();
-        DistributionMapping dm(ba);
+        BoxArray ba;
+        DistributionMapping dm;
         const auto old_num_setdm = num_setdm;
         const auto old_num_setba = num_setba;
 
+#ifdef AMREX_USE_BITTREE
+        if(!use_bittree) {
+#endif
+            ba = MakeBaseGrids();
+            dm = MakeDistributionMap(0, ba);
+
+#ifdef AMREX_USE_BITTREE
+        }
+        else {
+            //Initialize Bittree
+
+            // top = number of grids on coarsest level in each direction
+            std::vector<int> top(AMREX_SPACEDIM,0);
+            IntVect ncells = geom[0].Domain().length();
+            for(int i=0; i<AMREX_SPACEDIM; ++i) {
+                top[i] = ncells[i] / max_grid_size[0][i];
+            }
+
+            // includes = boolean to check each coarsest level grid exists
+            // (Bittree supports having "holes" in the mesh)
+            int ngrids = AMREX_D_TERM(top[0],*top[1],*top[2]);
+            std::vector<int> includes(ngrids,1);
+
+            btmesh = std::make_unique<bittree::BittreeAmr>(top.data(),includes.data());
+
+            // Set BCs
+            for(int d=0; d<AMREX_SPACEDIM; ++d) {
+                btUnit::bcPeriodic[d] = geom[0].isPeriodic(d);
+            }
+
+
+            // Use Bittree to make coarsest level (don't need MakeBaseGrids)
+            // Need to use Bittree, so the indices of grids[lev] will be compatible with BT.
+            btUnit::btCalculateLevel(btmesh.get(),0,ba,max_grid_size[0]);
+            dm = MakeDistributionMap(0, ba);
+        }
+#endif
         MakeNewLevelFromScratch(0, time, ba, dm);
 
         if (old_num_setba == num_setba) {
@@ -812,7 +950,7 @@ AmrMesh::MakeNewGrids (Real time)
             if (new_finest <= finest_level) { break; }
             finest_level = new_finest;
 
-            DistributionMapping dm(new_grids[new_finest]);
+            DistributionMapping dm = MakeDistributionMap(new_finest, new_grids[new_finest]);
             const auto old_num_setdm = num_setdm;
 
             MakeNewLevelFromScratch(new_finest, time, new_grids[finest_level], dm);
@@ -843,7 +981,7 @@ AmrMesh::MakeNewGrids (Real time)
                 for (int lev = 1; lev <= new_finest; ++lev) {
                     if (new_grids[lev] != grids[lev]) {
                         grids_the_same = false;
-                        DistributionMapping dm(new_grids[lev]);
+                        DistributionMapping dm = MakeDistributionMap(lev, new_grids[lev]);
                         const auto old_num_setdm = num_setdm;
 
                         MakeNewLevelFromScratch(lev, time, new_grids[lev], dm);

diff --git a/Src/AmrCore/Make.package b/Src/AmrCore/Make.package
@@ -1,5 +1,6 @@
 
-CEXE_headers += AMReX_AmrCore.H AMReX_Cluster.H AMReX_ErrorList.H AMReX_FillPatchUtil.H AMReX_FillPatchUtil_I.H AMReX_FluxRegister.H \
+CEXE_headers += AMReX_AmrCore.H AMReX_Cluster.H AMReX_ErrorList.H AMReX_FillPatchUtil.H \
+                AMReX_FillPatchUtil_I.H AMReX_FluxRegister.H \
                 AMReX_Interpolater.H AMReX_MFInterpolater.H AMReX_TagBox.H AMReX_AmrMesh.H \
                 AMReX_InterpBase.H
 CEXE_sources += AMReX_AmrCore.cpp AMReX_Cluster.cpp AMReX_ErrorList.cpp AMReX_FillPatchUtil.cpp AMReX_FluxRegister.cpp \

diff --git a/Src/Extern/Bittree/AMReX_Bittree.H b/Src/Extern/Bittree/AMReX_Bittree.H
@@ -0,0 +1,75 @@
+#ifndef BL_BITTREE_H_
+#define BL_BITTREE_H_
+
+#include <AMReX_IntVect.H>
+#include <AMReX_BoxArray.H>
+#include <AMReX_DistributionMapping.H>
+#include <Bittree_BittreeAmr.h>
+
+namespace amrex {
+
+/*
+Include in Make.local:
+BITTREE_PATH = /path/to/bittree/installation
+INCLUDE_LOCATIONS += $(BITTREE_PATH)/include
+LIBRARY_LOCATIONS += $(BITTREE_PATH)/lib
+LIBRARIES += -lbittree
+
+Include in inputs:
+amr.use_bittree = true
+*/
+
+class btUnit {
+  // Functions used in AmrMesh
+  public:
+    static int  btRefine(bittree::BittreeAmr* const mesh,
+                         std::vector<int>& btTags,
+                         int max_crse, int lbase,
+                         Vector<BoxArray>& grids, Vector<DistributionMapping>& dmap,
+                         MPI_Comm comm);
+    static void btCalculateGrids(bittree::BittreeAmr* const mesh,
+                                 int lbase,
+                                 int& new_finest,
+                                 Vector<BoxArray>& new_grids,
+                                 Vector<IntVect> const& max_grid_size);
+    static void btCalculateLevel(bittree::BittreeAmr* const mesh,
+                                 int lev,
+                                 BoxArray& ba,
+                                 IntVect const& max_grid_size);
+  // Utils
+  public:
+    static int getBitid(bittree::BittreeAmr* const mesh, bool updated,
+                        int lev, int idx_on_lev);
+    static int getIndex(bittree::BittreeAmr* const mesh, bool updated,
+                        int lev, int bitid);
+
+  // Functions to implement strict octree logic
+  private:
+    static void btCheckRefine(bittree::BittreeAmr* const mesh,
+                              std::vector<int>& btTags,
+                              int max_crse, int lbase,
+                              Vector<BoxArray>& grids, Vector<DistributionMapping>& dmap,
+                              MPI_Comm comm);
+
+    static void btCheckDerefine(bittree::BittreeAmr* const mesh,
+                                std::vector<int>& btTags,
+                                int max_crse, int lbase,
+                                Vector<BoxArray>& grids, Vector<DistributionMapping>& dmap,
+                                MPI_Comm comm);
+
+    // Utility Functions
+    static bool checkNeighborsRefine(bittree::BittreeAmr* const mesh,
+                                     bittree::MortonTree::Block b);
+    static std::vector<int> neighIntCoords(bittree::BittreeAmr* const mesh,
+                                           unsigned lev, unsigned* lcoord,
+                                           int* gCell);
+
+  public:
+    // Represents whether domain has periodic BC in each direction
+    // true = Periodic, false = Non-Periodic
+    static bool bcPeriodic[AMREX_SPACEDIM];
+};
+
+
+}
+#endif