Add flat sparse packs

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Current develop
 
 ### Added (new features/APIs/variables/...)
+- [[PR 877]](https://github.com/parthenon-hpc-lab/parthenon/pull/877) Add flat sparse packs
 - [[PR 830]](https://github.com/parthenon-hpc-lab/parthenon/pull/830) Add particle output
 - [[PR 840]](https://github.com/parthenon-hpc-lab/parthenon/pull/840) Generalized integrators infrastructure in a backwards compatible way
 - [[PR 810]](https://github.com/parthenon-hpc-lab/parthenon/pull/810) Add suport for Ascent in-situ visualization

doc/sphinx/src/interface/sparse.rst

@@ -137,6 +137,16 @@ simulation. The set of possible operations for a sparse field are:
    ``lroberts36/merge-sparse-with-jdolence-sparse`` that is being used
    in ``Riot``. This should probably be brought into ``develop``, since
    the “sparse sparse pack” access pattern is probably not desirable.
+- *Flatten outer indices:* Most common packs over ``MeshData`` produce
+  a 5-dimensional data structure, where the slowest two moving indices
+  are over ``MeshBlock``s then tensor/vector components of
+  ``Variable``s. However, it is sometimes desirable to construct a
+  4-dimensional data structure where the slowest moving index ranges
+  over all variables and all blocks. So for example if we had a system
+  with 9 blocks and 5 variables per block, the slowest moving index
+  would be of size 45. This can be enabled ``SparsePack``s with the
+  ``GetFlat`` series of factory functions or by passing the optional
+  ``flat`` boolean into the constructor.
 
 In comparison to a sparse field, a dense field only requires the
 operation *Access*.

src/interface/sparse_pack.hpp

@@ -125,10 +125,10 @@ class SparsePack : public SparsePackBase {
   // The pack will be created and accessible on the device
   template <class T>
   static SparsePack Get(T *pmd, const std::vector<MetadataFlag> &flags = {},
-                        bool fluxes = false, bool coarse = false) {
+                        bool fluxes = false, bool coarse = false, bool flatten = false) {
     const impl::PackDescriptor desc(std::vector<std::string>{Ts::name()...},
                                     std::vector<bool>{Ts::regex()...}, flags, fluxes,
-                                    coarse);
+                                    coarse, flatten);
     return SparsePack(SparsePackBase::GetPack(pmd, desc));
   }
 
@@ -143,9 +143,9 @@ class SparsePack : public SparsePackBase {
   template <class T, class VAR_VEC>
   static std::tuple<SparsePack, SparsePackIdxMap>
   Get(T *pmd, const VAR_VEC &vars, const std::vector<MetadataFlag> &flags = {},
-      bool fluxes = false, bool coarse = false) {
+      bool fluxes = false, bool coarse = false, bool flatten = false) {
     static_assert(sizeof...(Ts) == 0, "Cannot create a string/type hybrid pack");
-    impl::PackDescriptor desc(vars, flags, fluxes, coarse);
+    impl::PackDescriptor desc(vars, flags, fluxes, coarse, flatten);
     return {SparsePack(SparsePackBase::GetPack(pmd, desc)),
             SparsePackBase::GetIdxMap(desc)};
   }
@@ -183,13 +183,65 @@ class SparsePack : public SparsePackBase {
     Get(pmd, vars, flags, fluxes, coarse);
   }
 
+  template <class T>
+  static SparsePack GetFlat(T *pmd, const std::vector<MetadataFlag> &flags = {}) {
+    const bool coarse = false;
+    const bool fluxes = false;
+    const bool flatten = true;
+    return Get(pmd, flags, fluxes, coarse, flatten);
+  }
+
+  template <class T, class VAR_VEC>
+  static std::tuple<SparsePack, SparsePackIdxMap>
+  GetFlat(T *pmd, const VAR_VEC &vars, const std::vector<MetadataFlag> &flags = {}) {
+    const bool coarse = false;
+    const bool fluxes = false;
+    const bool flatten = true;
+    return Get(pmd, vars, flags, fluxes, coarse, flatten);
+  }
+
+  template <class T>
+  static SparsePack GetFlatWithFluxes(T *pmd,
+                                      const std::vector<MetadataFlag> &flags = {}) {
+    const bool coarse = false;
+    const bool fluxes = true;
+    const bool flatten = true;
+    return Get(pmd, flags, fluxes, coarse, flatten);
+  }
+
+  template <class T, class VAR_VEC>
+  static std::tuple<SparsePack, SparsePackIdxMap>
+  GetFlatWithFluxes(T *pmd, const VAR_VEC &vars,
+                    const std::vector<MetadataFlag> &flags = {}) {
+    const bool coarse = false;
+    const bool fluxes = true;
+    const bool flatten = true;
+    return Get(pmd, vars, flags, fluxes, coarse, flatten);
+  }
+
+  template <class T>
+  static SparsePack GetFlatWithCoarse(T *pmd,
+                                      const std::vector<MetadataFlag> &flags = {}) {
+    const bool coarse = true;
+    const bool fluxes = false;
+    const bool flatten = true;
+    return Get(pmd, flags, fluxes, coarse, flatten);
+  }
+
+  template <class T, class VAR_VEC>
+  static std::tuple<SparsePack, SparsePackIdxMap>
+  GetFlatWithCoarse(T *pmd, const VAR_VEC &vars,
+                    const std::vector<MetadataFlag> &flags = {}) {
+    const bool coarse = true;
+    const bool fluxes = false;
+    const bool flatten = true;
+    return Get(pmd, vars, flags, fluxes, coarse, flatten);
+  }
+
   // Methods for getting parts of the shape of the pack
   KOKKOS_FORCEINLINE_FUNCTION
   int GetNBlocks() const { return nblocks_; }
 
-  KOKKOS_FORCEINLINE_FUNCTION
-  int GetNDim() const { return ndim_; }
-
   KOKKOS_FORCEINLINE_FUNCTION
   int GetMaxNumberOfVars() const { return pack_.extent_int(2); }
 
@@ -200,27 +252,32 @@ class SparsePack : public SparsePackBase {
   KOKKOS_INLINE_FUNCTION int GetLowerBound(const int b) const { return 0; }
 
   KOKKOS_INLINE_FUNCTION int GetUpperBound(const int b) const {
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Bounds not valid for flat packs");
     return bounds_(1, b, nvar_);
   }
 
   KOKKOS_INLINE_FUNCTION int GetLowerBound(const int b, PackIdx idx) const {
     static_assert(sizeof...(Ts) == 0, "Cannot create a string/type hybrid pack");
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Bounds not valid for flat packs");
     return bounds_(0, b, idx.VariableIdx());
   }
 
   KOKKOS_INLINE_FUNCTION int GetUpperBound(const int b, PackIdx idx) const {
     static_assert(sizeof...(Ts) == 0, "Cannot create a string/type hybrid pack");
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Bounds not valid for flat packs");
     return bounds_(1, b, idx.VariableIdx());
   }
 
   template <class TIn, REQUIRES(IncludesType<TIn, Ts...>::value)>
   KOKKOS_INLINE_FUNCTION int GetLowerBound(const int b, const TIn &) const {
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Bounds not valid for flat packs");
     const int vidx = GetTypeIdx<TIn, Ts...>::value;
     return bounds_(0, b, vidx);
   }
 
   template <class TIn, REQUIRES(IncludesType<TIn, Ts...>::value)>
   KOKKOS_INLINE_FUNCTION int GetUpperBound(const int b, const TIn &) const {
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Bounds not valid for flat packs");
     const int vidx = GetTypeIdx<TIn, Ts...>::value;
     return bounds_(1, b, vidx);
   }
@@ -229,27 +286,32 @@ class SparsePack : public SparsePackBase {
   KOKKOS_INLINE_FUNCTION int GetLowerBoundHost(const int b) const { return 0; }
 
   KOKKOS_INLINE_FUNCTION int GetUpperBoundHost(const int b) const {
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Bounds not valid for flat packs");
     return bounds_h_(1, b, nvar_);
   }
 
   KOKKOS_INLINE_FUNCTION int GetLowerBoundHost(const int b, PackIdx idx) const {
     static_assert(sizeof...(Ts) == 0);
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Bounds not valid for flat packs");
     return bounds_h_(0, b, idx.VariableIdx());
   }
 
   KOKKOS_INLINE_FUNCTION int GetUpperBoundHost(const int b, PackIdx idx) const {
     static_assert(sizeof...(Ts) == 0);
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Bounds not valid for flat packs");
     return bounds_h_(1, b, idx.VariableIdx());
   }
 
   template <class TIn, REQUIRES(IncludesType<TIn, Ts...>::value)>
   KOKKOS_INLINE_FUNCTION int GetLowerBoundHost(const int b, const TIn &) const {
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Bounds not valid for flat packs");
     const int vidx = GetTypeIdx<TIn, Ts...>::value;
     return bounds_h_(0, b, vidx);
   }
 
   template <class TIn, REQUIRES(IncludesType<TIn, Ts...>::value)>
   KOKKOS_INLINE_FUNCTION int GetUpperBoundHost(const int b, const TIn &) const {
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Bounds not valid for flat packs");
     const int vidx = GetTypeIdx<TIn, Ts...>::value;
     return bounds_h_(1, b, vidx);
   }
@@ -261,55 +323,75 @@ class SparsePack : public SparsePackBase {
   KOKKOS_INLINE_FUNCTION
   auto &operator()(const int b, PackIdx idx) const {
     static_assert(sizeof...(Ts) == 0);
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Accessor cannot be used for flat packs");
     const int n = bounds_(0, b, idx.VariableIdx()) + idx.Offset();
     return pack_(0, b, n);
   }
 
   template <class TIn, REQUIRES(IncludesType<TIn, Ts...>::value)>
   KOKKOS_INLINE_FUNCTION auto &operator()(const int b, const TIn &t) const {
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Accessor cannot be used for flat packs");
     const int vidx = GetLowerBound(b, t) + t.idx;
     return pack_(0, b, vidx);
   }
 
   KOKKOS_INLINE_FUNCTION
   Real &operator()(const int b, const int idx, const int k, const int j,
                    const int i) const {
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Accessor cannot be used for flat packs");
     return pack_(0, b, idx)(k, j, i);
   }
+  KOKKOS_INLINE_FUNCTION
+  Real &operator()(int idx, const int k, const int j, const int i) const {
+    PARTHENON_DEBUG_REQUIRE(flat_, "Accessor valid only for flat packs");
+    return pack_(0, 0, idx)(k, j, i);
+  }
 
   KOKKOS_INLINE_FUNCTION
   Real &operator()(const int b, PackIdx idx, const int k, const int j,
                    const int i) const {
     static_assert(sizeof...(Ts) == 0, "Cannot create a string/type hybrid pack");
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Accessor cannot be used for flat packs");
     const int n = bounds_(0, b, idx.VariableIdx()) + idx.Offset();
     return pack_(0, b, n)(k, j, i);
   }
 
   template <class TIn, REQUIRES(IncludesType<TIn, Ts...>::value)>
   KOKKOS_INLINE_FUNCTION Real &operator()(const int b, const TIn &t, const int k,
                                           const int j, const int i) const {
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Accessor cannot be used for flat packs");
     const int vidx = GetLowerBound(b, t) + t.idx;
     return pack_(0, b, vidx)(k, j, i);
   }
 
   // flux() overloads
   KOKKOS_INLINE_FUNCTION
   auto &flux(const int b, const int dir, const int idx) const {
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Accessor cannot be used for flat packs");
     PARTHENON_DEBUG_REQUIRE(dir > 0 && dir < 4 && with_fluxes_, "Bad input to flux call");
     return pack_(dir, b, idx);
   }
 
   KOKKOS_INLINE_FUNCTION
   Real &flux(const int b, const int dir, const int idx, const int k, const int j,
              const int i) const {
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Accessor cannot be used for flat packs");
     PARTHENON_DEBUG_REQUIRE(dir > 0 && dir < 4 && with_fluxes_, "Bad input to flux call");
     return pack_(dir, b, idx)(k, j, i);
   }
 
+  KOKKOS_INLINE_FUNCTION
+  Real &flux(const int dir, const int idx, const int k, const int j, const int i) const {
+    PARTHENON_DEBUG_REQUIRE(flat_, "Accessor must only be used for flat packs");
+    PARTHENON_DEBUG_REQUIRE(dir > 0 && dir < 4 && with_fluxes_, "Bad input to flux call");
+    return pack_(dir, 0, idx)(k, j, i);
+  }
+
   KOKKOS_INLINE_FUNCTION
   Real &flux(const int b, const int dir, PackIdx idx, const int k, const int j,
              const int i) const {
     static_assert(sizeof...(Ts) == 0, "Cannot create a string/type hybrid pack");
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Accessor cannot be used for flat packs");
     PARTHENON_DEBUG_REQUIRE(dir > 0 && dir < 4 && with_fluxes_, "Bad input to flux call");
     const int n = bounds_(0, b, idx.VariableIdx()) + idx.Offset();
     return pack_(dir, b, n)(k, j, i);
@@ -318,6 +400,7 @@ class SparsePack : public SparsePackBase {
   template <class TIn, REQUIRES(IncludesType<TIn, Ts...>::value)>
   KOKKOS_INLINE_FUNCTION Real &flux(const int b, const int dir, const TIn &t, const int k,
                                     const int j, const int i) const {
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Accessor cannot be used for flat packs");
     PARTHENON_DEBUG_REQUIRE(dir > 0 && dir < 4 && with_fluxes_, "Bad input to flux call");
     const int vidx = GetLowerBound(b, t) + t.idx;
     return pack_(dir, b, vidx)(k, j, i);

src/interface/sparse_pack_base.cpp

@@ -85,21 +85,22 @@ SparsePackBase SparsePackBase::Build(T *pmd, const PackDescriptor &desc) {
   pack.with_fluxes_ = desc.with_fluxes;
   pack.coarse_ = desc.coarse;
   pack.nvar_ = desc.vars.size();
+  pack.flat_ = desc.flat;
 
   // Count up the size of the array that is required
   int max_size = 0;
-  int nblocks = 0;
-  int ndim = 3;
+  int nblocks = desc.flat;
+  int size = 0;
   ForEachBlock(pmd, [&](int b, mbd_t *pmbd) {
-    int size = 0;
-    nblocks++;
+    if (!desc.flat) {
+      size = 0;
+      nblocks++;
+    }
     for (auto &pv : pmbd->GetCellVariableVector()) {
       for (int i = 0; i < nvar; ++i) {
         if (desc.IncludeVariable(i, pv)) {
           if (pv->IsAllocated()) {
             size += pv->GetDim(6) * pv->GetDim(5) * pv->GetDim(4);
-            ndim = (pv->GetDim(1) > 1 ? 1 : 0) + (pv->GetDim(2) > 1 ? 1 : 0) +
-                   (pv->GetDim(3) > 1 ? 1 : 0);
           }
         }
       }
@@ -122,12 +123,19 @@ SparsePackBase SparsePackBase::Build(T *pmd, const PackDescriptor &desc) {
   auto coords_h = Kokkos::create_mirror_view(pack.coords_);
 
   // Fill the views
-  ForEachBlock(pmd, [&](int b, mbd_t *pmbd) {
-    int idx = 0;
-    coords_h(b) = pmbd->GetBlockPointer()->coords_device;
+  int idx = 0;
+  ForEachBlock(pmd, [&](int block, mbd_t *pmbd) {
+    int b = 0;
+    if (!desc.flat) {
+      idx = 0;
+      b = block;
+      coords_h(b) = pmbd->GetBlockPointer()->coords_device;
+    }
 
     for (int i = 0; i < nvar; ++i) {
-      pack.bounds_h_(0, b, i) = idx;
+      if (!desc.flat) {
+        pack.bounds_h_(0, b, i) = idx;
+      }
 
       for (auto &pv : pmbd->GetCellVariableVector()) {
         if (desc.IncludeVariable(i, pv)) {
@@ -148,18 +156,8 @@ SparsePackBase SparsePackBase::Build(T *pmd, const PackDescriptor &desc) {
                     PARTHENON_REQUIRE(pack_h(1, b, idx).size() ==
                                           pack_h(0, b, idx).size(),
                                       "Different size fluxes.");
-                    if (ndim > 1) {
-                      pack_h(2, b, idx) = pv->flux[2].Get(t, u, v);
-                      PARTHENON_REQUIRE(pack_h(2, b, idx).size() ==
-                                            pack_h(0, b, idx).size(),
-                                        "Different size fluxes.");
-                    }
-                    if (ndim > 2) {
-                      pack_h(3, b, idx) = pv->flux[3].Get(t, u, v);
-                      PARTHENON_REQUIRE(pack_h(3, b, idx).size() ==
-                                            pack_h(0, b, idx).size(),
-                                        "Different size fluxes.");
-                    }
+                    pack_h(2, b, idx) = pv->flux[2].Get(t, u, v);
+                    pack_h(3, b, idx) = pv->flux[3].Get(t, u, v);
                   }
                   idx++;
                 }
@@ -169,23 +167,25 @@ SparsePackBase SparsePackBase::Build(T *pmd, const PackDescriptor &desc) {
         }
       }
 
-      pack.bounds_h_(1, b, i) = idx - 1;
-
-      if (pack.bounds_h_(1, b, i) < pack.bounds_h_(0, b, i)) {
-        // Did not find any allocated variables meeting our criteria
-        pack.bounds_h_(0, b, i) = -1;
-        // Make the upper bound more negative so a for loop won't iterate once
-        pack.bounds_h_(1, b, i) = -2;
+      if (!desc.flat) {
+        pack.bounds_h_(1, b, i) = idx - 1;
+        if (pack.bounds_h_(1, b, i) < pack.bounds_h_(0, b, i)) {
+          // Did not find any allocated variables meeting our criteria
+          pack.bounds_h_(0, b, i) = -1;
+          // Make the upper bound more negative so a for loop won't iterate once
+          pack.bounds_h_(1, b, i) = -2;
+        }
       }
     }
     // Record the maximum for easy access
-    pack.bounds_h_(1, b, nvar) = idx - 1;
+    if (!desc.flat) {
+      pack.bounds_h_(1, b, nvar) = idx - 1;
+    }
   });
 
   Kokkos::deep_copy(pack.pack_, pack_h);
   Kokkos::deep_copy(pack.bounds_, pack.bounds_h_);
   Kokkos::deep_copy(pack.coords_, coords_h);
-  pack.ndim_ = ndim;
   pack.dims_[1] = pack.nblocks_;
   pack.dims_[2] = -1; // Not allowed to ask for the ragged dimension anyway
   pack.dims_[3] = pack_h(0, 0, 0).extent_int(0);
@@ -206,8 +206,6 @@ SparsePackBase &SparsePackCache::Get(T *pmd, const PackDescriptor &desc) {
   std::string ident = GetIdentifier(desc);
   if (pack_map.count(ident) > 0) {
     auto &pack = pack_map[ident].first;
-    if (desc.with_fluxes != pack.with_fluxes_) return BuildAndAdd(pmd, desc, ident);
-    if (desc.coarse != pack.coarse_) return BuildAndAdd(pmd, desc, ident);
     auto alloc_status_in = SparsePackBase::GetAllocStatus(pmd, desc);
     auto &alloc_status = pack_map[ident].second;
     if (alloc_status.size() != alloc_status_in.size())
@@ -246,6 +244,10 @@ std::string SparsePackCache::GetIdentifier(const PackDescriptor &desc) const {
   identifier += "____";
   for (int i = 0; i < desc.vars.size(); ++i)
     identifier += desc.vars[i] + std::to_string(desc.use_regex[i]);
+  identifier += "____";
+  identifier += std::to_string(desc.with_fluxes);
+  identifier += std::to_string(desc.coarse);
+  identifier += std::to_string(desc.flat);
   return identifier;
 }