[WIP] Add sub packs & scratchpad/array type list indexers

doc/sphinx/src/sparse_packs.rst

@@ -123,6 +123,53 @@ A given sparse field may or may not be allocated on each block within a pack. To
      return TaskStatus::complete;
    }
 
+Slicing into ``SparsePack``\ s with ``SubPack``\ s
+--------------------------------------------------
+
+A `SubPack` provdies a view into a slice of a `SparsePack` along a given dimension(s).
+`SubPack`\ s are constructed with a block + `kji` index to give a slice into the 
+fields at the meshblock + cell index. When a `SubPack` is constructed with the
+`Axis` template parameters then the `SubPack` also providies a view into slices
+of the `SparsePack` along the provided axes offset from the provided `kji` indices.
+
+.. code:: c++
+
+  const int ni = ib.e - ib.s + 1;
+  const int ic = ib.s + ni / 2;
+  par_for(
+      PARTHENON_AUTO_LABEL, 0, sparse_pack.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e,
+      KOKKOS_LAMBDA(int b, int k, int j) {
+        int ltot = 0;
+        int lo = sparse_pack.GetLowerBound(b, v3());
+        int hi = sparse_pack.GetUpperBound(b, v3());
+        auto sub_pack = parthenon::SubPack<Axis::I>(sparse_pack, b, k, j, ic);
+
+        for (int i = ib.s - ni / 2; i <= ib.e - ni / 2; i++) {
+          for (int c = 0; c <= hi - lo; ++c) {
+            Real n = i + ic + 1e1 * j + 1e2 * k + 1e4 * c + 1e5 * v + 1e3 * b;
+            // indexes into sparse_pack(b, v3(c), k, j, ic + i)
+            if (n != sub_pack(v3(c), i)) ltot += 1;
+          }
+        }
+      });
+
+For example usage see the `unit test <https://github.com/parthenon-hpc-lab/parthenon/blob/develop/tst/unit/test_sparse_pack.cpp>`
+
+Type based indexing arrays
+--------------------------
+
+It is often convenient when working with type-based packs to also be able to index into
+an array or scratchpads using the same type-based indexing used for the packs. Parthenon
+provides some objects to wrap either a `Kokkos::Array` or a `ScratchPad2D<Real>` and
+to index into these with types with the `TypeListArray` or `ScratchPack` functions.
+Users only need to provide an interface that conforms to the `PackLike` interface 
+that provides a compile time value for number of variable components to be indexed into
+and a method to determine an integer index from a type. Interfaces are provided by Parthenon
+for `SparsePacks`, `TypeList`s & `VarList`s for variables that are defined with a 
+`ncomp` property.
+
+For example usage see the `unit test <https://github.com/parthenon-hpc-lab/parthenon/blob/develop/tst/unit/test_sparse_pack.cpp>`
+
 .. [1] In practice, there are ways of selecting subsets of fields for inclusion in a given ``MeshBlockData`` instance.
 .. [2] ``ParArray``\ s are lite wrappers around ``Kokkos::View``\ s and therefore obey reference semantics.
 .. [3] Additionally, because many ``std::`` library containers don't work on device, the chain from field name through ``MeshBlockData`` to ``Variable`` to the underlying ``ParArray`` cannot be legally followed within a kernel.

src/pack/subpack.hpp

@@ -0,0 +1,156 @@
+#ifndef PACK_SUBPACK_HPP_
+#define PACK_SUBPACK_HPP_
+
+#include "pack/sparse_pack.hpp"
+#include "utils/concepts_lite.hpp"
+
+namespace parthenon {
+
+namespace impl {
+enum class Axis { K = 0, J = 1, I = 2 };
+
+template <typename PackType, Axis... axes>
+struct SubPack_impl {
+  KOKKOS_INLINE_FUNCTION SubPack_impl(PackType &pack, const int &b, const int &k,
+                                      const int &j, const int &i)
+      : pack_(pack), b_(b), k_(k), j_(j), i_(i) {}
+
+  template <typename Var_t>
+  KOKKOS_INLINE_FUNCTION Real &operator()(const Var_t &var) const {
+    return pack_(b_, var, k_, j_, i_);
+  }
+
+  template <typename Var_t>
+  KOKKOS_INLINE_FUNCTION Real &operator()(TopologicalElement te, const Var_t &var) const {
+    return pack_(b_, te, var, k_, j_, i_);
+  }
+
+  template <typename Var_t>
+  KOKKOS_INLINE_FUNCTION Real &flux(TopologicalElement te, const Var_t &var) const {
+    return pack_.flux(b_, te, var, k_, j_, i_);
+  }
+
+  template <typename V>
+  KOKKOS_INLINE_FUNCTION std::size_t GetSize(const V &var) const {
+    return pack_.GetSize(b_, var);
+  }
+
+ private:
+  PackType &pack_;
+  const int b_, k_, j_, i_;
+};
+
+template <typename PackType, Axis... axes>
+struct StencilSubPack_impl {
+  KOKKOS_INLINE_FUNCTION StencilSubPack_impl(PackType &pack, const int &b, const int &k,
+                                             const int &j, const int &i)
+      : pack_(pack), b_(b), kji_({k, j, i}) {}
+
+  template <typename Var_t, typename... Is>
+  KOKKOS_INLINE_FUNCTION Real &operator()(const Var_t &var, Is &&...idxs) {
+    static_assert(sizeof...(Is) == sizeof...(axes),
+                  "number of indices passed to sub pack must match number of axes.");
+    Kokkos::Array<int, 3> kji = kji_;
+    ([&]() { kji[static_cast<int>(axes)] += idxs; }(), ...);
+    return pack_(b_, var, kji[0], kji[1], kji[2]);
+  }
+
+  template <typename Var_t, typename... Is>
+  KOKKOS_INLINE_FUNCTION Real &operator()(TopologicalElement te, const Var_t &var,
+                                          Is &&...idxs) {
+    static_assert(sizeof...(Is) == sizeof...(axes),
+                  "number of indices passed to sub pack must match number of axes.");
+    Kokkos::Array<int, 3> kji = kji_;
+    ([&]() { kji[static_cast<int>(axes)] += idxs; }(), ...);
+    return pack_(b_, te, var, kji[0], kji[1], kji[2]);
+  }
+
+  template <typename Var_t, typename... Is>
+  KOKKOS_INLINE_FUNCTION Real &flux(TopologicalElement te, const Var_t &var,
+                                    Is &&...idxs) {
+    static_assert(sizeof...(Is) == sizeof...(axes),
+                  "number of indices passed to sub pack must match number of axes.");
+    Kokkos::Array<int, 3> kji = kji_;
+    ([&]() { kji[static_cast<int>(axes)] += idxs; }(), ...);
+    return pack_.flux(b_, te, var, kji[0], kji[1], kji[2]);
+  }
+
+  template <typename V>
+  KOKKOS_INLINE_FUNCTION std::size_t GetSize(const V &var) const {
+    return pack_.GetSize(b_, var);
+  }
+
+ private:
+  const PackType &pack_;
+  const Kokkos::Array<int, 3> kji_;
+  const int b_;
+};
+
+template <typename Var_t, typename PackType, Axis... axes>
+struct VarStencilSubPack_impl {
+  KOKKOS_INLINE_FUNCTION VarStencilSubPack_impl(PackType &pack, const int &b,
+                                                const Var_t &var, const int &k,
+                                                const int &j, const int &i)
+      : pack_(pack), b_(b), var_(var), kji_({k, j, i}) {}
+
+  template <typename... Is>
+  KOKKOS_INLINE_FUNCTION Real &operator()(Is &&...idxs) {
+    static_assert(sizeof...(Is) == sizeof...(axes),
+                  "number of indices passed to sub pack must match number of axes.");
+    Kokkos::Array<int, 3> kji = kji_;
+    ([&]() { kji[static_cast<int>(axes)] += idxs; }(), ...);
+    return pack_(b_, var_, kji[0], kji[1], kji[2]);
+  }
+
+  template <typename... Is>
+  KOKKOS_INLINE_FUNCTION Real &operator()(TopologicalElement te, Is &&...idxs) {
+    static_assert(sizeof...(Is) == sizeof...(axes),
+                  "number of indices passed to sub pack must match number of axes.");
+    Kokkos::Array<int, 3> kji = kji_;
+    ([&]() { kji[static_cast<int>(axes)] += idxs; }(), ...);
+    return pack_(b_, te, var_, kji[0], kji[1], kji[2]);
+  }
+
+  template <typename... Is>
+  KOKKOS_INLINE_FUNCTION Real &flux(TopologicalElement te, Is &&...idxs) {
+    static_assert(sizeof...(Is) == sizeof...(axes),
+                  "number of indices passed to sub pack must match number of axes.");
+    Kokkos::Array<int, 3> kji = kji_;
+    ([&]() { kji[static_cast<int>(axes)] += idxs; }(), ...);
+    return pack_.flux(b_, te, var_, kji[0], kji[1], kji[2]);
+  }
+
+ private:
+  const PackType &pack_;
+  const Kokkos::Array<int, 3> kji_;
+  const Var_t var_;
+  const int b_;
+};
+
+template <typename PackType>
+using is_sparse_pack = is_specialization_of<base_type<PackType>, SparsePack>;
+
+} // namespace impl
+
+template <Axis axis, Axis... axes, typename Var_t, typename PackType,
+          REQUIRES(is_sparse_pack<PackType>::value)>
+KOKKOS_INLINE_FUNCTION auto SubPack(PackType &pack, const int &b, const Var_t &var,
+                                    const int &k, const int &j, const int &i) {
+  return VarStencilSubPack_impl<Var_t, PackType, axis, axes...>(pack, b, var, k, j, i);
+}
+
+template <Axis axis, Axis... axes, typename PackType,
+          REQUIRES(is_sparse_pack<PackType>::value)>
+KOKKOS_INLINE_FUNCTION auto SubPack(PackType &pack, const int &b, const int &k,
+                                    const int &j, const int &i) {
+  return StencilSubPack_impl<PackType, axis, axes...>(pack, b, k, j, i);
+}
+
+template <typename PackType, REQUIRES(is_sparse_pack<PackType>::value)>
+KOKKOS_INLINE_FUNCTION auto SubPack(PackType &pack, const int &b, const int &k,
+                                    const int &j, const int &i) {
+  return SubPack_impl<PackType>(pack, b, k, j, i);
+}
+
+} // namespace parthenon
+#endif // PACK_SUBPACK_HPP_

src/utils/concepts_lite.hpp

@@ -39,6 +39,10 @@ struct is_specialization_of : public std::false_type {};
 template <template <class...> class TEMPL, class... TPARAMS>
 struct is_specialization_of<TEMPL<TPARAMS...>, TEMPL> : public std::true_type {};
 
+// this is in c++20 to remove any const & ref from a given type
+template <typename T>
+using base_type = typename std::remove_cv_t<typename std::remove_reference_t<T>>;
+
 // This is a variadic template class that accepts any set of types
 // and is always equal to void as long as the types are well formed.
 // Although it seems simple, it is the basis of the SFINAE "void_t

src/utils/type_arrays.hpp

@@ -0,0 +1,174 @@
+#ifndef UTILS_TYPE_ARRAY_HPP_
+#define UTILS_TYPE_ARRAY_HPP_
+
+#include "Kokkos_Macros.hpp"
+#include "pack/sparse_pack.hpp"
+#include "utils/concepts_lite.hpp"
+#include "utils/type_list.hpp"
+#include <utility>
+
+namespace parthenon {
+template <typename>
+struct SparsePackList {};
+
+template <typename... Ts>
+struct SparsePackList<SparsePack<Ts...>> {
+  using type = SparsePack<Ts...>;
+  // this doesn't actually reflect the size of the types packed, as
+  // that can not be guaranteed at runtime
+  static constexpr std::size_t ncomp = sizeof...(Ts);
+
+  KOKKOS_INLINE_FUNCTION SparsePackList(const type &pack_in, const int &b_in)
+      : pack(pack_in), b(b_in) {}
+
+  template <typename T>
+  KOKKOS_INLINE_FUNCTION std::size_t GetIndex(const T &t) const {
+    return pack.GetIndex(b, t);
+  }
+
+ private:
+  const type &pack;
+  const int b;
+};
+
+template <typename... Vars>
+struct VarList {
+  template <typename... Ts>
+  using TypeList = parthenon::TypeList<Ts...>;
+
+  static constexpr std::size_t GetSize() {
+    std::size_t size = 0;
+    ([&] { size += Vars::ncomp; }(), ...);
+    return size;
+  }
+
+  static constexpr std::size_t ncomp = GetSize();
+  template <typename V>
+  KOKKOS_INLINE_FUNCTION std::size_t GetIndex(const V &var) const {
+    return GetIndex_(TypeList<Vars...>(), var);
+  }
+
+ private:
+  template <typename V, typename... Vs>
+  KOKKOS_INLINE_FUNCTION std::size_t GetIndex_(TypeList<V, Vs...>, const V &var) const {
+    return var.idx;
+  }
+
+  template <typename V, typename U, typename... Us>
+  KOKKOS_INLINE_FUNCTION std::size_t GetIndex_(TypeList<U, Us...>, const V &var) const {
+    return U::ncomp + GetIndex_(TypeList<Us...>(), var);
+  }
+};
+
+namespace impl {
+template <typename>
+struct TypeListArray {};
+
+template <template <typename...> typename PackType, typename... Ts>
+struct TypeListArray<PackType<Ts...>> {
+  using type = PackType<Ts...>;
+  using Arr_t = Kokkos::Array<Real, type::ncomp>;
+
+  KOKKOS_INLINE_FUNCTION TypeListArray(const type &pack_in) : pack(pack_in) {}
+  KOKKOS_INLINE_FUNCTION TypeListArray(const type &pack_in, const Real &value)
+      : TypeListArray(pack_in) {
+    for (int idx = 0; idx < type::ncomp; idx++) {
+      data[idx] = value;
+    }
+  }
+  KOKKOS_INLINE_FUNCTION TypeListArray(const type &pack_in, Arr_t data_in)
+      : TypeListArray(pack_in), data(data_in) {}
+
+  template <typename V, REQUIRES(IncludesType<V, Ts...>::value)>
+  KOKKOS_INLINE_FUNCTION Real &operator()(const V &var) {
+    return data[pack.GetIndex(var)];
+  }
+
+  KOKKOS_INLINE_FUNCTION Real &operator[](const std::size_t &idx) { return data[idx]; }
+
+ private:
+  Arr_t data;
+  const type &pack;
+};
+
+template <typename, typename, typename>
+struct ScratchPack_impl {};
+
+template <typename ScratchPad, template <typename...> typename PackType, typename... Ts,
+          int... Is>
+struct ScratchPack_impl<ScratchPad, PackType<Ts...>, std::integer_sequence<int, Is...>> {
+  using type = PackType<Ts...>;
+
+  template <typename... Args>
+  KOKKOS_INLINE_FUNCTION ScratchPack_impl(const type pack_, ScratchPad scratch_,
+                                          Args &&...idxs)
+      : pack(pack_), scratch(scratch_), kji({idxs...}) {}
+
+  template <typename V, REQUIRES(IncludesType<V, Ts...>::value)>
+  KOKKOS_INLINE_FUNCTION Real &operator()(const V &var) const {
+    return scratch(pack.GetIndex(var), kji[Is]...);
+  }
+
+  template <typename V, typename... Args, REQUIRES(IncludesType<V, Ts...>::value)>
+  KOKKOS_INLINE_FUNCTION Real &operator()(const V &var, Args &&...idxs) const {
+    static_assert(sizeof...(Is) == sizeof...(Args),
+                  "Must provide number of indices equal to dimension of the underlying "
+                  "ScratchPad.");
+    return scratch(pack.GetIndex(var), kji[Is] + idxs...);
+  }
+
+  template <typename... Args>
+  KOKKOS_INLINE_FUNCTION Real &operator()(const int &var, Args &&...idxs) const {
+    return scratch(var, kji[Is] + idxs...);
+  }
+
+  KOKKOS_INLINE_FUNCTION Real &operator()(const int &var) const {
+    return scratch(var, kji[Is]...);
+  }
+
+ private:
+  const type pack;
+  ScratchPad scratch;
+  const Kokkos::Array<int, sizeof...(Is)> kji;
+};
+
+// TypeList containers that can be used to index into an integer array need
+// to provide
+//    * a static constexpr std::size_t ncomp
+//       that declares the size of the array to index into
+//       note that this is not used by the ScratchPack, as it assumes
+//       that the scratch memory is already allocated
+//    * and an int GetIndex() method templated on the types in the list
+template <typename... Ts>
+struct PackLike {
+  template <typename T, REQUIRES(implements<integral(decltype(T::ncomp))>::value)>
+  auto requires_(T) -> void_t<decltype(T::ncomp), decltype(T().GetIndex(Ts()))...>;
+};
+
+} // namespace impl
+
+template <template <typename...> typename PackType, typename... Ts, typename... Args,
+          REQUIRES(implements<PackLike<Ts...>(PackType<Ts...>)>::value &&
+                   !is_specialization_of<PackType<Ts...>, SparsePackList>::value)>
+KOKKOS_INLINE_FUNCTION auto TypeListArray(const PackType<Ts...> &pack, Args &&...args) {
+  return impl::TypeListArray<PackType<Ts...>>(pack, std::forward<Args>(args)...);
+}
+
+template <typename ScratchPad, template <typename...> typename PackType, typename... Ts,
+          typename... Args, REQUIRES(implements<PackLike<Ts...>(PackType<Ts...>)>::value)>
+KOKKOS_INLINE_FUNCTION auto ScratchPack(const PackType<Ts...> &pack, ScratchPad &scratch,
+                                        Args &&...args) {
+  return ScratchPack_impl<ScratchPad, PackType<Ts...>,
+                          std::make_integer_sequence<int, sizeof...(Args)>>(
+      pack, scratch, std::forward<Args>(args)...);
+}
+
+template <typename ScratchPad, typename... Ts>
+KOKKOS_INLINE_FUNCTION auto ScratchPack(const SparsePack<Ts...> &pack,
+                                        ScratchPad &scratch, const int &b, const int &i) {
+  auto spl = SparsePackList(pack, b);
+  return ScratchPack(spl, scratch, i);
+}
+
+} // namespace parthenon
+#endif // UTILS_TYPE_ARRAY_HPP_