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

@@ -12,6 +12,7 @@
 #include <iosfwd>
 #include <cstddef>
 #include <map>
+#include <memory>
 #include <unordered_map>
 
 namespace amrex
@@ -850,6 +851,10 @@ public:
 
     [[nodiscard]] BATransformer const& transformer () const;
 
+    [[nodiscard]] std::weak_ptr<BARef> getWeakRef () const;
+    [[nodiscard]] std::shared_ptr<BARef> const& getSharedRef () const;
+    std::shared_ptr<BARef>& getSharedRef ();
+
     friend class AmrMesh;
     friend class FabArrayBase;
 

Src/Base/AMReX_BoxArray.cpp

@@ -1623,6 +1623,24 @@ BoxArray::transformer () const
     return m_bat;
 }
 
+std::weak_ptr<BARef>
+BoxArray::getWeakRef () const
+{
+    return std::weak_ptr<BARef>{m_ref};
+}
+
+std::shared_ptr<BARef> const&
+BoxArray::getSharedRef () const
+{
+    return m_ref;
+}
+
+std::shared_ptr<BARef>&
+BoxArray::getSharedRef ()
+{
+    return m_ref;
+}
+
 std::ostream&
 operator<< (std::ostream&   os,
             const BoxArray& ba)

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<int>& pmap) : m_pmap(pmap) {}
+
+        explicit Ref (Vector<int>&& 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<int> m_pmap; //!< index array for all boxes
+        Vector<int> m_index_array;  //!< index array for local boxes owned by the team
+        std::vector<bool> 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<Ref> 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<T>& cost,
                                                       Real* efficiency);
 
+    [[nodiscard]] std::weak_ptr<Ref> getWeakRef () const;
+
 private:
 
     const Vector<int>& 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<int>& pmap) : m_pmap(pmap) {}
-
-        explicit Ref (Vector<int>&& 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<int> m_pmap; //!< index array for all boxes
-        Vector<int> m_index_array;  //!< index array for local boxes owned by the team
-        std::vector<bool> m_ownership; //!< true ownership
-    };
     //
     //! The data -- a reference-counted pointer to a Ref.
     std::shared_ptr<Ref> 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);
 

Src/Base/AMReX_DistributionMapping.cpp

@@ -326,6 +326,11 @@ DistributionMapping::DistributionMapping (const BoxArray& boxes,
     define(boxes,nprocs);
 }
 
+DistributionMapping::DistributionMapping (std::shared_ptr<Ref> 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<DMRef>
+DistributionMapping::getWeakRef () const
+{
+    return std::weak_ptr<DMRef>{m_ref};
+}
+
 std::ostream&
 operator<< (std::ostream&              os,
             const DistributionMapping& pmap)

Src/Base/AMReX_VisMF.cpp

@@ -9,6 +9,8 @@
 #include <cerrno>
 #include <cstdio>
 #include <limits>
+#include <vector>
+#include <utility>
 
 namespace amrex {
 
@@ -60,6 +62,45 @@ namespace
         }
     }
 #endif
+
+    std::vector<std::pair<std::weak_ptr<BARef>, std::weak_ptr<DMRef>>> s_layout_cache;
+
+    DistributionMapping vismf_make_dm (BoxArray& ba)
+    {
+        auto& ba_ref = ba.getSharedRef();
+
+        std::shared_ptr<DMRef> 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<FArrayBox> &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()));