diff --git a/doc/sphinx/src/sparse_packs.rst b/doc/sphinx/src/sparse_packs.rst index 8e378cbaccf28..20495d3540662 100644 --- a/doc/sphinx/src/sparse_packs.rst +++ b/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(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 ` + +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` 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 ` + .. [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. diff --git a/src/pack/subpack.hpp b/src/pack/subpack.hpp new file mode 100644 index 0000000000000..83983ad0ead14 --- /dev/null +++ b/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 +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 + KOKKOS_INLINE_FUNCTION Real &operator()(const Var_t &var) const { + return pack_(b_, var, k_, j_, i_); + } + + template + KOKKOS_INLINE_FUNCTION Real &operator()(TopologicalElement te, const Var_t &var) const { + return pack_(b_, te, var, k_, j_, i_); + } + + template + KOKKOS_INLINE_FUNCTION Real &flux(TopologicalElement te, const Var_t &var) const { + return pack_.flux(b_, te, var, k_, j_, i_); + } + + template + 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 +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 + 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 kji = kji_; + ([&]() { kji[static_cast(axes)] += idxs; }(), ...); + return pack_(b_, var, kji[0], kji[1], kji[2]); + } + + template + 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 kji = kji_; + ([&]() { kji[static_cast(axes)] += idxs; }(), ...); + return pack_(b_, te, var, kji[0], kji[1], kji[2]); + } + + template + 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 kji = kji_; + ([&]() { kji[static_cast(axes)] += idxs; }(), ...); + return pack_.flux(b_, te, var, kji[0], kji[1], kji[2]); + } + + template + KOKKOS_INLINE_FUNCTION std::size_t GetSize(const V &var) const { + return pack_.GetSize(b_, var); + } + + private: + const PackType &pack_; + const Kokkos::Array kji_; + const int b_; +}; + +template +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 + 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 kji = kji_; + ([&]() { kji[static_cast(axes)] += idxs; }(), ...); + return pack_(b_, var_, kji[0], kji[1], kji[2]); + } + + template + 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 kji = kji_; + ([&]() { kji[static_cast(axes)] += idxs; }(), ...); + return pack_(b_, te, var_, kji[0], kji[1], kji[2]); + } + + template + 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 kji = kji_; + ([&]() { kji[static_cast(axes)] += idxs; }(), ...); + return pack_.flux(b_, te, var_, kji[0], kji[1], kji[2]); + } + + private: + const PackType &pack_; + const Kokkos::Array kji_; + const Var_t var_; + const int b_; +}; + +template +using is_sparse_pack = is_specialization_of, SparsePack>; + +} // namespace impl + +template ::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(pack, b, var, k, j, i); +} + +template ::value)> +KOKKOS_INLINE_FUNCTION auto SubPack(PackType &pack, const int &b, const int &k, + const int &j, const int &i) { + return StencilSubPack_impl(pack, b, k, j, i); +} + +template ::value)> +KOKKOS_INLINE_FUNCTION auto SubPack(PackType &pack, const int &b, const int &k, + const int &j, const int &i) { + return SubPack_impl(pack, b, k, j, i); +} + +} // namespace parthenon +#endif // PACK_SUBPACK_HPP_ diff --git a/src/utils/concepts_lite.hpp b/src/utils/concepts_lite.hpp index b3b49d7561981..ae07748b87b57 100644 --- a/src/utils/concepts_lite.hpp +++ b/src/utils/concepts_lite.hpp @@ -39,6 +39,10 @@ struct is_specialization_of : public std::false_type {}; template