Adding support to memory collection of Particles and MultiFabs within a grid

Src/Base/AMReX_FabArray.H

@@ -817,6 +817,10 @@ public:
     template <class F=FAB, std::enable_if_t<IsBaseFab<F>::value,int> = 0>
     void setDomainBndry (value_type val, int strt_comp, int ncomp, const Geometry& geom);
 
+    //! Get capacity of the FabArray
+    template <typename I, class F=FAB, std::enable_if_t<IsBaseFab<F>::value && std::is_integral_v<I> && (sizeof(I) >= sizeof(Long)), int> = 0>
+    void capacityOfFabs (LayoutData<I>& mem) const;
+
     /**
     * \brief Returns the sum of component "comp"
     *
@@ -2506,6 +2510,17 @@ FabArray<FAB>::setDomainBndry (value_type val,
     }
 }
 
+
+template <class FAB>
+template <typename I, class F, std::enable_if_t<IsBaseFab<F>::value && std::is_integral_v<I> && (sizeof(I) >= sizeof(Long)), int> FOO>
+void
+FabArray<FAB>::capacityOfFabs (LayoutData<I>& mem) const {
+    AMREX_ASSERT(amrex::isMFIterSafe(*this, mem));
+    for (MFIter mfi(*this, MFItInfo{}.DisableDeviceSync()); mfi.isValid(); ++mfi) {
+        mem[mfi] += static_cast<I>((*this)[mfi].nBytesOwned());
+    }
+}
+
 template <class FAB>
 template <class F, std::enable_if_t<IsBaseFab<F>::value,int> FOO>
 typename F::value_type

Src/Particle/AMReX_ParIter.H

@@ -93,6 +93,8 @@ public:
 
     [[nodiscard]] int numNeighborParticles () const { return GetParticleTile().numNeighborParticles(); }
 
+    [[nodiscard]] Long capacity () const { return GetParticleTile().capacity(); }
+
     [[nodiscard]] int GetLevel () const { return m_level; }
 
     [[nodiscard]] std::pair<int, int> GetPairIndex () const { return std::make_pair(this->index(), this->LocalTileIndex()); }

Src/Particle/AMReX_ParticleContainer.H

@@ -578,6 +578,12 @@ public:
 
     Vector<Long> NumberOfParticlesInGrid  (int level, bool only_valid = true, bool only_local = false) const;
 
+    /**
+     * \brief Return capacity of memory for particles at specific grid
+     */
+    template <typename I, std::enable_if_t<std::is_integral_v<I> && (sizeof(I) >= sizeof(Long)), int> = 0>
+    void CapacityOfParticlesInGrid (LayoutData<I>& mem, int lev) const;
+
     /**
     * \brief Returns # of particles at all levels
     *

Src/Particle/AMReX_ParticleContainerI.H

@@ -586,6 +586,25 @@
     return nparticles;
 }
 
+template <typename ParticleType, int NArrayReal, int NArrayInt,
+          template<class> class Allocator, class CellAssignor>
+template <typename I, std::enable_if_t<std::is_integral_v<I> && (sizeof(I) >= sizeof(Long)), int> FOO>
+void
+ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssignor>::CapacityOfParticlesInGrid (LayoutData<I>& mem, int lev) const
+{
+    AMREX_ASSERT(lev >= 0 && lev < int(m_particles.size()));
+
+    AMREX_ASSERT(BoxArray::SameRefs(mem.boxArray(), ParticleBoxArray(lev)) &&
+                 mem.DistributionMap() == ParticleDistributionMap(lev));
+
+    [[maybe_unused]] Gpu::NoSyncRegion no_sync{};
+    for (ParConstIterType pti(*this, lev); pti.isValid(); ++pti)
+    {
+        int gid = pti.index();
+        mem[gid] += static_cast<I>(pti.capacity());
+    }
+}
+
 //
 // This includes both valid and invalid particles since invalid particles still take up space.
 //

Tests/CMakeLists.txt

@@ -125,9 +125,9 @@ else()
    #
    # List of subdirectories to search for CMakeLists.
    #
-   set( AMREX_TESTS_SUBDIRS Amr AsyncOut CallNoinline CLZ CommType CTOParFor DeviceGlobal Enum
-                            MultiBlock MultiPeriod ParmParse Parser Parser2 ParserUserFn Reinit
-                            RoundoffDomain SmallMatrix)
+   set( AMREX_TESTS_SUBDIRS Amr AsyncOut CallNoinline CLZ CommType CTOParFor DeviceGlobal
+                            Enum HeatEquation MultiBlock MultiPeriod ParmParse Parser
+                            Parser2 ParserUserFn Reinit RoundoffDomain SmallMatrix)
 
    if (AMReX_PARTICLES)
      list(APPEND AMREX_TESTS_SUBDIRS Particles)

Tests/HeatEquation/GNUmakefile

@@ -0,0 +1,21 @@
+# AMREX_HOME defines the directory in which we will find all the AMReX code.
+AMREX_HOME ?= ../../../../amrex
+
+USE_ASSERTION = TRUE
+
+DEBUG        = FALSE
+USE_MPI      = FALSE
+USE_OMP      = FALSE
+PRECISION	 = FLOAT
+COMP         = gnu
+DIM          = 3
+
+include $(AMREX_HOME)/Tools/GNUMake/Make.defs
+
+include Make.package
+# VPATH_LOCATIONS  += ../Source
+# INCLUDE_LOCATIONS += ../Source
+
+include $(AMREX_HOME)/Src/Base/Make.package
+
+include $(AMREX_HOME)/Tools/GNUMake/Make.rules

Tests/HeatEquation/Make.package

@@ -0,0 +1,2 @@
+CEXE_sources += main.cpp
+CEXE_headers += myfunc.H

Tests/HeatEquation/inputs

@@ -0,0 +1,7 @@
+n_cell = 32
+max_grid_size = 16
+
+nsteps = 1000
+plot_int = 0
+
+dt = 1.e-5

Tests/HeatEquation/main.cpp

@@ -0,0 +1,208 @@
+#include <AMReX_PlotFileUtil.H>
+#include <AMReX_ParmParse.H>
+#include <AMReX_MultiFabUtil.H>
+// #include <AMReX.H>
+// #include <AMReX_FabArray.H>
+
+#include "myfunc.H"
+
+using namespace amrex;
+
+int main (int argc, char* argv[])
+{
+    amrex::Initialize(argc,argv);
+
+    main_main();
+
+    amrex::Finalize();
+    return 0;
+}
+
+void main_main ()
+{
+
+    // **********************************
+    // SIMULATION PARAMETERS
+
+    // number of cells on each side of the domain
+    int n_cell;
+
+    // size of each box (or grid)
+    int max_grid_size;
+
+    // total steps in simulation
+    int nsteps;
+
+    // how often to write a plotfile
+    int plot_int;
+
+    // time step
+    Real dt;
+
+    // inputs parameters
+    {
+        // ParmParse is way of reading inputs from the inputs file
+        // pp.get means we require the inputs file to have it
+        // pp.query means we optionally need the inputs file to have it - but we must supply a default here
+        ParmParse pp;
+
+        // We need to get n_cell from the inputs file - this is the number of cells on each side of
+        //   a square (or cubic) domain.
+        pp.get("n_cell",n_cell);
+
+        // The domain is broken into boxes of size max_grid_size
+        pp.get("max_grid_size",max_grid_size);
+
+        // Default nsteps to 10, allow us to set it to something else in the inputs file
+        nsteps = 10;
+        pp.query("nsteps",nsteps);
+
+        // Default plot_int to -1, allow us to set it to something else in the inputs file
+        //  If plot_int < 0 then no plot files will be written
+        plot_int = -1;
+        pp.query("plot_int",plot_int);
+
+        // time step
+        pp.get("dt",dt);
+    }
+
+    // **********************************
+    // SIMULATION SETUP
+
+    // make BoxArray and Geometry
+    // ba will contain a list of boxes that cover the domain
+    // geom contains information such as the physical domain size,
+    //               number of points in the domain, and periodicity
+    BoxArray ba;
+    Geometry geom;
+
+    // AMREX_D_DECL means "do the first X of these, where X is the dimensionality of the simulation"
+    IntVect dom_lo(AMREX_D_DECL(       0,        0,        0));
+    IntVect dom_hi(AMREX_D_DECL(n_cell-1, n_cell-1, n_cell-1));
+
+    // Make a single box that is the entire domain
+    Box domain(dom_lo, dom_hi);
+
+    // Initialize the boxarray "ba" from the single box "domain"
+    ba.define(domain);
+
+    // Break up boxarray "ba" into chunks no larger than "max_grid_size" along a direction
+    ba.maxSize(max_grid_size);
+
+    // This defines the physical box, [0,1] in each direction.
+    RealBox real_box({AMREX_D_DECL( 0., 0., 0.)},
+                     {AMREX_D_DECL( 1., 1., 1.)});
+
+    // periodic in all direction
+    Array<int,AMREX_SPACEDIM> is_periodic{AMREX_D_DECL(1,1,1)};
+
+    // This defines a Geometry object
+    geom.define(domain, real_box, CoordSys::cartesian, is_periodic);
+
+    // extract dx from the geometry object
+    GpuArray<Real,AMREX_SPACEDIM> dx = geom.CellSizeArray();
+
+    // Nghost = number of ghost cells for each array
+    int Nghost = 1;
+
+    // Ncomp = number of components for each array
+    int Ncomp = 1;
+
+    // How Boxes are distributed among MPI processes
+    DistributionMapping dm(ba);
+
+    // we allocate two phi multifabs; one will store the old state, the other the new.
+    MultiFab phi_old(ba, dm, Ncomp, Nghost);
+    MultiFab phi_new(ba, dm, Ncomp, Nghost);
+
+    // time = starting time in the simulation
+    Real time = 0.0;
+
+    // **********************************
+    // INITIALIZE DATA
+
+    // loop over boxes
+    for (MFIter mfi(phi_old); mfi.isValid(); ++mfi)
+    {
+        const Box& bx = mfi.validbox();
+
+        const Array4<Real>& phiOld = phi_old.array(mfi);
+
+        // set phi = 1 + e^(-(r-0.5)^2)
+        amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k)
+        {
+            Real x = (i+0.5) * dx[0];
+            Real y = (j+0.5) * dx[1];
+#if (AMREX_SPACEDIM == 2)
+            Real rsquared = ((x-0.5)*(x-0.5)+(y-0.5)*(y-0.5))/0.01;
+#elif (AMREX_SPACEDIM == 3)
+            Real z= (k+0.5) * dx[2];
+            Real rsquared = ((x-0.5)*(x-0.5)+(y-0.5)*(y-0.5)+(z-0.5)*(z-0.5))/0.01;
+#endif
+            phiOld(i,j,k) = 1. + std::exp(-rsquared);
+        });
+    }
+
+    // Write a plotfile of the initial data if plot_int > 0
+    if (plot_int > 0)
+    {
+        int step = 0;
+        const std::string& pltfile = amrex::Concatenate("plt",step,5);
+        WriteSingleLevelPlotfile(pltfile, phi_old, {"phi"}, geom, time, 0);
+    }
+
+    for (int step = 1; step <= nsteps; ++step)
+    {
+        // fill periodic ghost cells
+        phi_old.FillBoundary(geom.periodicity());
+
+        // new_phi = old_phi + dt * Laplacian(old_phi)
+        // loop over boxes
+        for ( MFIter mfi(phi_old); mfi.isValid(); ++mfi )
+        {
+            const Box& bx = mfi.validbox();
+
+            const Array4<Real>& phiOld = phi_old.array(mfi);
+            const Array4<Real>& phiNew = phi_new.array(mfi);
+
+            // advance the data by dt
+            amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
+            {
+                phiNew(i,j,k) = phiOld(i,j,k) + dt *
+                    ( (phiOld(i+1,j,k) - 2.*phiOld(i,j,k) + phiOld(i-1,j,k)) / (dx[0]*dx[0])
+                     +(phiOld(i,j+1,k) - 2.*phiOld(i,j,k) + phiOld(i,j-1,k)) / (dx[1]*dx[1])
+#if (AMREX_SPACEDIM == 3)
+                     +(phiOld(i,j,k+1) - 2.*phiOld(i,j,k) + phiOld(i,j,k-1)) / (dx[2]*dx[2])
+#endif
+                        );
+            });
+        }
+
+        // update time
+        time = time + dt;
+
+        // copy new solution into old solution
+        MultiFab::Copy(phi_old, phi_new, 0, 0, 1, 0);
+
+        // Tell the I/O Processor to write out which step we're doing
+        amrex::Print() << "Advanced step " << step << "\n";
+
+        // Write a plotfile of the current data (plot_int was defined in the inputs file)
+        if (plot_int > 0 && step%plot_int == 0)
+        {
+            const std::string& pltfile = amrex::Concatenate("plt",step,5);
+            WriteSingleLevelPlotfile(pltfile, phi_new, {"phi"}, geom, time, step);
+        }
+    }
+
+    // Testing capacity function
+    amrex::LayoutData<std::size_t> mem(ba, dm);
+    phi_old.capacityOfFabs(mem);
+    phi_new.capacityOfFabs(mem);
+
+    long total_bytes = 0;
+    for (amrex::MFIter mfi(mem); mfi.isValid(); ++mfi) {
+        amrex::Print() << "Box " << mfi.index() << " uses "
+                     << mem[mfi] << " bytes\n";
+    }
+}

Tests/HeatEquation/myfunc.H

@@ -0,0 +1,6 @@
+#ifndef MYFUNC_H_
+#define MYFUNC_H_
+
+void main_main ();
+
+#endif

Tests/Particles/NeighborList/main.cpp

@@ -138,4 +138,13 @@ void testNeighborList ()
         auto& nlist2 = neighbors_type_2[0][std::make_pair(gid, tid)];
         nlist2.print();
     }
+
+    LayoutData<std::size_t> mem(ba, dm);
+    pc1.CapacityOfParticlesInGrid(mem, 0);
+
+    Vector<Long> particles_in_grid = pc1.NumberOfParticlesInGrid(0, false,true);
+    for (amrex::MFIter mfi(mem); mfi.isValid(); ++mfi) {
+        amrex::Print() << "Box " << mfi.index() << " uses "
+                     << mem[mfi] << " bytes and have " << particles_in_grid[mfi.LocalIndex()] << " particles\n";
+    }
 }

build_docs.sh

@@ -7,6 +7,7 @@ cd Docs/Doxygen
 doxygen doxygen.conf &> doxygen.out
 if grep -q "warning:" doxygen.out; then
     echo "Doxygen warnings detected! Failing..."
+    cat doxygen.out
     exit 1
 fi
 cd ../..