Add MultiFab::sum(region) and sum_unique(region)

Add two new functions to MultiFab that return the sum of the given region.
AMReX-Codes · Apr 3, 2024 · 398fe4b · 398fe4b
1 parent 66bcf38
commit 398fe4b
Show file tree

Hide file tree

Showing 2 changed files with 111 additions and 0 deletions.
diff --git a/Src/Base/AMReX_MultiFab.H b/Src/Base/AMReX_MultiFab.H
@@ -238,6 +238,10 @@ public:
     * \brief Returns the sum of component "comp" over the MultiFab -- no ghost cells are included.
     */
     [[nodiscard]] Real sum (int comp = 0, bool local = false) const;
+    /**
+    * \brief Returns the sum of component "comp" in the given "region". -- no ghost cells are included.
+    */
+    [[nodiscard]] Real sum (Box const& region, int comp = 0, bool local = false) const;
 
     using FabArray<FArrayBox>::sum;
 
@@ -249,6 +253,13 @@ public:
                      bool local = false,
                      const Periodicity& period = Periodicity::NonPeriodic()) const;
     /**
+    * \brief Returns the unique sum of component "comp" in the gieven
+    * region. Non-unique points owned by multiple boxes in the MultiFab are
+    * skipped. No ghost cells are included. This function does not take
+    * periodicity into account in the determination of uniqueness of points.
+    */
+    [[nodiscard]] Real sum_unique (Box const& region, int comp = 0, bool local = false) const;
+    /**
     * \brief Adds the scalar value \p val to the value of each cell in the
     * specified subregion of the MultiFab.
     *

diff --git a/Src/Base/AMReX_MultiFab.cpp b/Src/Base/AMReX_MultiFab.cpp
@@ -1209,6 +1209,52 @@ MultiFab::sum (int comp, bool local) const
     return FabArray<FArrayBox>::sum(comp, IntVect(0), local);
 }
 
+Real
+MultiFab::sum (Box const& region, int comp, bool local) const
+{
+    BL_PROFILE("MultiFab::sum(region)");
+
+    auto sm = 0.0_rt;
+
+#ifdef AMREX_USE_GPU
+    if (Gpu::inLaunchRegion()) {
+        auto const& ma = this->const_arrays();
+        mx = ParReduce(TypeList<ReduceOpSum>{}, TypeList<Real>{}, *this, IntVect(0),
+        [=] AMREX_GPU_DEVICE (int box_no, int i, int j, int k) noexcept -> GpuTuple<Real>
+        {
+            if (region.contains(i,j,k)) {
+                return ma[box_no](i,j,k,comp);
+            } else {
+                return 0.0_rt;
+            }
+        });
+    } else
+#endif
+    {
+#ifdef AMREX_USE_OMP
+#pragma omp parallel if (!system::regtest_reduction) reduction(+:sm)
+#endif
+        for (MFIter mfi(*this,true); mfi.isValid(); ++mfi) {
+            Box const& bx = mfi.tilebox() & region;
+            if (bx.ok()) {
+                auto const& a = this->const_array(mfi);
+                auto tmp = 0.0_rt;
+                AMREX_LOOP_3D(bx, i, j, k,
+                {
+                    tmp += a(i,j,k,comp);
+                });
+                sm += tmp; // Do it this way so that it does not break regression tests.
+            }
+        }
+    }
+
+    if (!local) {
+        ParallelAllReduce::Sum(sm, ParallelContext::CommunicatorSub());
+    }
+
+    return sm;
+}
+
 Real
 MultiFab::sum_unique (int comp,
                       bool local,
@@ -1262,6 +1308,60 @@ MultiFab::sum_unique (int comp,
     return sm;
 }
 
+Real
+MultiFab::sum_unique (Box const& region, int comp, bool local) const
+{
+    BL_PROFILE("MultiFab::sum_unique(region)");
+
+    // no duplicatly distributed points if cell centered
+    if (ixType().cellCentered()) {
+        return this->sum(region, comp, local);
+    }
+
+    // Owner is the grid with the lowest grid number containing the data
+    std::unique_ptr<iMultiFab> owner_mask = OwnerMask();
+
+    auto sm = 0.0_rt;
+#ifdef AMREX_USE_GPU
+    if (Gpu::inLaunchRegion()) {
+        auto const& ma = this->const_arrays();
+        auto const& msk = owner_mask->const_arrays();
+        sm = ParReduce(TypeList<ReduceOpSum>{}, TypeList<Real>{}, *this, IntVect(0),
+        [=] AMREX_GPU_DEVICE (int box_no, int i, int j, int k) noexcept
+                       -> GpuTuple<Real>
+        {
+            return (region.contains(i,j,k) && msk[box_no](i,j,k))
+                ? ma[box_no](i,j,k,comp) : 0.0_rt;
+        });
+    } else
+#endif
+    {
+#ifdef AMREX_USE_OMP
+#pragma omp parallel if (!system::regtest_reduction) reduction(+:sm)
+#endif
+        for (MFIter mfi(*this,true); mfi.isValid(); ++mfi)
+        {
+            Box const& bx = mfi.tilebox() & region;
+            if (bx.ok()) {
+                Array4<Real const> const& a = this->const_array(mfi);
+                Array4<int const> const& msk = owner_mask->const_array(mfi);
+                auto tmp = 0.0_rt;
+                AMREX_LOOP_3D(bx, i, j, k,
+                {
+                    if (msk(i,j,k)) { tmp += a(i,j,k,comp); }
+                });
+                sm += tmp; // Do it this way so that it does not break regression tests.
+            }
+        }
+    }
+
+    if (!local) {
+        ParallelAllReduce::Sum(sm, ParallelContext::CommunicatorSub());
+    }
+
+    return sm;
+}
+
 void
 MultiFab::minus (const MultiFab& mf, int strt_comp, int num_comp, int nghost)
 {