From 3407e877a7b219546c0ef12bef287845f5cb9f05 Mon Sep 17 00:00:00 2001 From: Weiqun Zhang Date: Wed, 20 Dec 2023 12:11:53 -0800 Subject: [PATCH] Add a few free functions for MLMG (#3680) These are useful when we use Array as the data type for MLMG. --- Src/Base/AMReX_FabArrayBase.H | 5 + Src/Base/AMReX_FabArrayBase.cpp | 20 ++++ Src/Base/AMReX_FabArrayUtility.H | 187 +++++++++++++++++++++++++++++++ Src/Base/AMReX_TypeTraits.H | 12 ++ 4 files changed, 224 insertions(+) diff --git a/Src/Base/AMReX_FabArrayBase.H b/Src/Base/AMReX_FabArrayBase.H index d8bc4411874..e2cf0ed9641 100644 --- a/Src/Base/AMReX_FabArrayBase.H +++ b/Src/Base/AMReX_FabArrayBase.H @@ -721,6 +721,11 @@ public: }; +[[nodiscard]] int nComp (FabArrayBase const& fa); +[[nodiscard]] IntVect nGrowVect (FabArrayBase const& fa); +[[nodiscard]] BoxArray const& boxArray (FabArrayBase const& fa); +[[nodiscard]] DistributionMapping const& DistributionMap (FabArrayBase const& fa); + #ifdef BL_USE_MPI bool CheckRcvStats (Vector& recv_stats, const Vector& recv_size, int tag); #endif diff --git a/Src/Base/AMReX_FabArrayBase.cpp b/Src/Base/AMReX_FabArrayBase.cpp index 8dd8275f66a..6997f3489dd 100644 --- a/Src/Base/AMReX_FabArrayBase.cpp +++ b/Src/Base/AMReX_FabArrayBase.cpp @@ -2699,4 +2699,24 @@ FabArrayBase::flushParForCache () #endif +int nComp (FabArrayBase const& fa) +{ + return fa.nComp(); +} + +IntVect nGrowVect (FabArrayBase const& fa) +{ + return fa.nGrowVect(); +} + +BoxArray const& boxArray (FabArrayBase const& fa) +{ + return fa.boxArray(); +} + +DistributionMapping const& DistributionMap (FabArrayBase const& fa) +{ + return fa.DistributionMap(); +} + } diff --git a/Src/Base/AMReX_FabArrayUtility.H b/Src/Base/AMReX_FabArrayUtility.H index ca80a070f45..0897c57ed4f 100644 --- a/Src/Base/AMReX_FabArrayUtility.H +++ b/Src/Base/AMReX_FabArrayUtility.H @@ -1602,6 +1602,193 @@ Dot (FabArray const& x, int xcomp, FabArray const& y, int ycomp, int n return sm; } +//! dst = val +template ,int> = 0> +void setVal (MF& dst, typename MF::value_type val) +{ + dst.setVal(val); +} + +//! dst = val in ghost cells. +template ,int> = 0> +void setBndry (MF& dst, typename MF::value_type val, int scomp, int ncomp) +{ + dst.setBndry(val, scomp, ncomp); +} + +//! dst = src +template && + IsMultiFabLike_v, int> = 0> +void LocalCopy (DMF& dst, SMF const& src, int scomp, int dcomp, + int ncomp, IntVect const& nghost) +{ + amrex::Copy(dst, src, scomp, dcomp, ncomp, nghost); +} + +//! dst += src +template ,int> = 0> +void LocalAdd (MF& dst, MF const& src, int scomp, int dcomp, + int ncomp, IntVect const& nghost) +{ + amrex::Add(dst, src, scomp, dcomp, ncomp, nghost); +} + +//! dst += a * src +template ,int> = 0> +void Saxpy (MF& dst, typename MF::value_type a, MF const& src, int scomp, int dcomp, + int ncomp, IntVect const& nghost) +{ + MF::Saxpy(dst, a, src, scomp, dcomp, ncomp, nghost); +} + +//! dst = src + a * dst +template ,int> = 0> +void Xpay (MF& dst, typename MF::value_type a, MF const& src, int scomp, int dcomp, + int ncomp, IntVect const& nghost) +{ + MF::Xpay(dst, a, src, scomp, dcomp, ncomp, nghost); +} + +//! dst = src w/ MPI communication +template , int> = 0> +void ParallelCopy (MF& dst, MF const& src, int scomp, int dcomp, int ncomp, + IntVect const& ng_src = IntVect(0), + IntVect const& ng_dst = IntVect(0), + Periodicity const& period = Periodicity::NonPeriodic()) +{ + dst.ParallelCopy(src, scomp, dcomp, ncomp, ng_src, ng_dst, period); +} + +template , int> = 0> +[[nodiscard]] typename MF::value_type +norminf (MF const& mf, int scomp, int ncomp, IntVect const& nghost, + bool local = false) +{ + return mf.norminf(scomp, ncomp, nghost, local); +} + +//! dst = val +template ,int> = 0> +void setVal (Array& dst, typename MF::value_type val) +{ + for (auto& mf: dst) { + mf.setVal(val); + } +} + +//! dst = val in ghost cells. +template ,int> = 0> +void setBndry (Array& dst, typename MF::value_type val, int scomp, int ncomp) +{ + for (auto& mf : dst) { + mf.setBndry(val, scomp, ncomp); + } +} + +//! dst = src +template && + IsMultiFabLike_v, int> = 0> +void LocalCopy (Array& dst, Array const& src, int scomp, int dcomp, + int ncomp, IntVect const& nghost) +{ + for (std::size_t i = 0; i < N; ++i) { + amrex::Copy(dst[i], src[i], scomp, dcomp, ncomp, nghost); + } +} + +//! dst += src +template ,int> = 0> +void LocalAdd (Array& dst, Array const& src, int scomp, int dcomp, + int ncomp, IntVect const& nghost) +{ + for (std::size_t i = 0; i < N; ++i) { + amrex::Add(dst[i], src[i], scomp, dcomp, ncomp, nghost); + } +} + +//! dst += a * src +template ,int> = 0> +void Saxpy (Array& dst, typename MF::value_type a, + Array const& src, int scomp, int dcomp, int ncomp, + IntVect const& nghost) +{ + for (std::size_t i = 0; i < N; ++i) { + MF::Saxpy(dst[i], a, src[i], scomp, dcomp, ncomp, nghost); + } +} + +//! dst = src + a * dst +template ,int> = 0> +void Xpay (Array& dst, typename MF::value_type a, + Array const& src, int scomp, int dcomp, int ncomp, + IntVect const& nghost) +{ + for (std::size_t i = 0; i < N; ++i) { + MF::Xpay(dst[i], a, src[i], scomp, dcomp, ncomp, nghost); + } +} + +//! dst = src w/ MPI communication +template , int> = 0> +void ParallelCopy (Array& dst, Array const& src, + int scomp, int dcomp, int ncomp, + IntVect const& ng_src = IntVect(0), + IntVect const& ng_dst = IntVect(0), + Periodicity const& period = Periodicity::NonPeriodic()) +{ + for (std::size_t i = 0; i < N; ++i) { + dst[i].ParallelCopy(src[i], scomp, dcomp, ncomp, ng_src, ng_dst, period); + } +} + +template , int> = 0> +[[nodiscard]] typename MF::value_type +norminf (Array const& mf, int scomp, int ncomp, IntVect const& nghost, + bool local = false) +{ + auto r = typename MF::value_type(0); + for (std::size_t i = 0; i < N; ++i) { + auto tmp = mf[i].norminf(scomp, ncomp, nghost, true); + r = std::max(r,tmp); + } + if (!local) { + ParallelAllReduce::Max(r, ParallelContext::CommunicatorSub()); + } + return r; +} + +template && (N > 0), + int> = 0> +[[nodiscard]] int nComp (Array const& mf) +{ + return mf[0].nComp(); +} + +template && (N > 0), + int> = 0> +[[nodiscard]] IntVect nGrowVect (Array const& mf) +{ + return mf[0].nGrowVect(); +} + +template && (N > 0), + int> = 0> +[[nodiscard]] BoxArray const& +boxArray (Array const& mf) +{ + return mf[0].boxArray(); +} + +template && (N > 0), + int> = 0> +[[nodiscard]] DistributionMapping const& +DistributionMap (Array const& mf) +{ + return mf[0].DistributionMap(); +} + } #endif diff --git a/Src/Base/AMReX_TypeTraits.H b/Src/Base/AMReX_TypeTraits.H index 222576f05f5..fbcb7a2c0e3 100644 --- a/Src/Base/AMReX_TypeTraits.H +++ b/Src/Base/AMReX_TypeTraits.H @@ -37,6 +37,18 @@ namespace amrex template inline constexpr bool IsFabArray_v = IsFabArray::value; + template + struct IsMultiFabLike : std::false_type {}; + // + template + struct IsMultiFabLike && + IsBaseFab_v > > + : std::true_type {}; + // + template + inline constexpr bool IsMultiFabLike_v = IsMultiFabLike::value; + + template using EnableIf_t = typename std::enable_if::type;