From c451ec030d2a82fd3838495ff1f338b865ab4aba Mon Sep 17 00:00:00 2001 From: Weiqun Zhang Date: Sat, 16 Nov 2024 20:00:08 -0800 Subject: [PATCH] VisMF::Read: Cache BoxArray and DistributionMapping (#4233) The cache ensures consistent DistributionMapping for the same BoxArray. It is now safe to call VisMF::Read twice with empty MultiFabs. Previously, even if the two MulitFabs had the same BoxArray, their DistributionMapping could be different, which was a potential source of bugs. --- Src/Base/AMReX_BoxArray.H | 5 +++ Src/Base/AMReX_BoxArray.cpp | 18 ++++++++++ Src/Base/AMReX_DistributionMapping.H | 48 +++++++++++++++----------- Src/Base/AMReX_DistributionMapping.cpp | 11 ++++++ Src/Base/AMReX_VisMF.cpp | 44 ++++++++++++++++++++++- 5 files changed, 104 insertions(+), 22 deletions(-) diff --git a/Src/Base/AMReX_BoxArray.H b/Src/Base/AMReX_BoxArray.H index 066e3b073f7..02ab23c2a26 100644 --- a/Src/Base/AMReX_BoxArray.H +++ b/Src/Base/AMReX_BoxArray.H @@ -12,6 +12,7 @@ #include #include #include +#include #include namespace amrex @@ -850,6 +851,10 @@ public: [[nodiscard]] BATransformer const& transformer () const; + [[nodiscard]] std::weak_ptr getWeakRef () const; + [[nodiscard]] std::shared_ptr const& getSharedRef () const; + std::shared_ptr& getSharedRef (); + friend class AmrMesh; friend class FabArrayBase; diff --git a/Src/Base/AMReX_BoxArray.cpp b/Src/Base/AMReX_BoxArray.cpp index fbb137f5e87..287c60d2ebe 100644 --- a/Src/Base/AMReX_BoxArray.cpp +++ b/Src/Base/AMReX_BoxArray.cpp @@ -1623,6 +1623,24 @@ BoxArray::transformer () const return m_bat; } +std::weak_ptr +BoxArray::getWeakRef () const +{ + return std::weak_ptr{m_ref}; +} + +std::shared_ptr const& +BoxArray::getSharedRef () const +{ + return m_ref; +} + +std::shared_ptr& +BoxArray::getSharedRef () +{ + return m_ref; +} + std::ostream& operator<< (std::ostream& os, const BoxArray& ba) diff --git a/Src/Base/AMReX_DistributionMapping.H b/Src/Base/AMReX_DistributionMapping.H index e9aa82f16a2..bc8d54e6851 100644 --- a/Src/Base/AMReX_DistributionMapping.H +++ b/Src/Base/AMReX_DistributionMapping.H @@ -47,6 +47,26 @@ class DistributionMapping //! The distribution strategies enum Strategy { UNDEFINED = -1, ROUNDROBIN, KNAPSACK, SFC, RRSFC }; + struct Ref + { + //! Constructors to match those in DistributionMapping .... + Ref () = default; + + explicit Ref (int len) : m_pmap(len) {} + + explicit Ref (const Vector& pmap) : m_pmap(pmap) {} + + explicit Ref (Vector&& pmap) noexcept : m_pmap(std::move(pmap)) {} + + //! dtor, copy-ctor, copy-op=, move-ctor, and move-op= are compiler generated. + + void clear () { m_pmap.clear(); m_index_array.clear(); m_ownership.clear(); } + + Vector m_pmap; //!< index array for all boxes + Vector m_index_array; //!< index array for local boxes owned by the team + std::vector m_ownership; //!< true ownership + }; + //! The default constructor. DistributionMapping () noexcept; @@ -73,6 +93,9 @@ class DistributionMapping //! Build mapping out of BoxArray over nprocs processors. explicit DistributionMapping (const BoxArray& boxes, int nprocs = ParallelDescriptor::NProcs()); + + explicit DistributionMapping (std::shared_ptr a_ref); + /** * \brief This is a very specialized distribution map. * Do NOT use it unless you really understand what it does. @@ -259,6 +282,8 @@ class DistributionMapping const std::vector& cost, Real* efficiency); + [[nodiscard]] std::weak_ptr getWeakRef () const; + private: const Vector& getIndexArray (); @@ -332,27 +357,6 @@ private: */ static PVMF m_BuildMap; - struct Ref - { - friend class DistributionMapping; - - //! Constructors to match those in DistributionMapping .... - Ref () = default; - - explicit Ref (int len) : m_pmap(len) {} - - explicit Ref (const Vector& pmap) : m_pmap(pmap) {} - - explicit Ref (Vector&& pmap) noexcept : m_pmap(std::move(pmap)) {} - - //! dtor, copy-ctor, copy-op=, move-ctor, and move-op= are compiler generated. - - void clear () { m_pmap.clear(); m_index_array.clear(); m_ownership.clear(); } - - Vector m_pmap; //!< index array for all boxes - Vector m_index_array; //!< index array for local boxes owned by the team - std::vector m_ownership; //!< true ownership - }; // //! The data -- a reference-counted pointer to a Ref. std::shared_ptr m_ref; @@ -375,6 +379,8 @@ public: [[nodiscard]] RefID getRefID () const noexcept { return RefID { m_ref.get() }; } }; +using DMRef = DistributionMapping::Ref; + //! Our output operator. std::ostream& operator<< (std::ostream& os, const DistributionMapping& pmap); diff --git a/Src/Base/AMReX_DistributionMapping.cpp b/Src/Base/AMReX_DistributionMapping.cpp index 45304cf5e04..7e524d0a7d2 100644 --- a/Src/Base/AMReX_DistributionMapping.cpp +++ b/Src/Base/AMReX_DistributionMapping.cpp @@ -326,6 +326,11 @@ DistributionMapping::DistributionMapping (const BoxArray& boxes, define(boxes,nprocs); } +DistributionMapping::DistributionMapping (std::shared_ptr a_ref) + : m_ref(std::move(a_ref)) +{ +} + DistributionMapping::DistributionMapping (const DistributionMapping& d1, const DistributionMapping& d2) : @@ -1958,6 +1963,12 @@ DistributionMapping::getOwnerShip () return m_ref->m_ownership; } +std::weak_ptr +DistributionMapping::getWeakRef () const +{ + return std::weak_ptr{m_ref}; +} + std::ostream& operator<< (std::ostream& os, const DistributionMapping& pmap) diff --git a/Src/Base/AMReX_VisMF.cpp b/Src/Base/AMReX_VisMF.cpp index f5f1430d59e..6749c626c3e 100644 --- a/Src/Base/AMReX_VisMF.cpp +++ b/Src/Base/AMReX_VisMF.cpp @@ -9,6 +9,8 @@ #include #include #include +#include +#include namespace amrex { @@ -60,6 +62,45 @@ namespace } } #endif + + std::vector, std::weak_ptr>> s_layout_cache; + + DistributionMapping vismf_make_dm (BoxArray& ba) + { + auto& ba_ref = ba.getSharedRef(); + + std::shared_ptr dm_ref; + int ielem = -1; + for (auto it = s_layout_cache.begin(); it != s_layout_cache.end(); ) { + auto cached_ba = it->first.lock(); + if (cached_ba) { + if (cached_ba->m_abox == ba_ref->m_abox) { + ba_ref = cached_ba; + dm_ref = it->second.lock(); + ielem = int(std::distance(s_layout_cache.begin(), it)); + break; + } + ++it; + } else { // expired + s_layout_cache.erase(it); + } + } + + auto have_dm_in_cache = bool(dm_ref); + ParallelDescriptor::ReduceBoolAnd(have_dm_in_cache); + if (have_dm_in_cache) { + return DistributionMapping{std::move(dm_ref)}; + } else { + DistributionMapping dm{ba}; + if (ielem >= 0) { + s_layout_cache[ielem].first = ba.getWeakRef(); + s_layout_cache[ielem].second = dm.getWeakRef(); + } else { + s_layout_cache.emplace_back(ba.getWeakRef(), dm.getWeakRef()); + } + return dm; + } + } } void @@ -104,6 +145,7 @@ void VisMF::Finalize () { initialized = false; + s_layout_cache.clear(); } void @@ -1542,7 +1584,7 @@ VisMF::Read (FabArray &mf, } if (mf.empty()) { - DistributionMapping dm(hdr.m_ba); + DistributionMapping dm = vismf_make_dm(hdr.m_ba); mf.define(hdr.m_ba, dm, hdr.m_ncomp, hdr.m_ngrow, MFInfo(), FArrayBoxFactory()); } else { BL_ASSERT(amrex::match(hdr.m_ba,mf.boxArray()));