diff --git a/CHANGELOG.md b/CHANGELOG.md index a4d9fa2847b0..07fa749a5531 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ ## Current develop ### Added (new features/APIs/variables/...) +- [[PR 998]](https://github.com/parthenon-hpc-lab/parthenon/pull/998) tensor indices added to sparse pack - [[PR 999]](https://github.com/parthenon-hpc-lab/parthenon/pull/999) Add a post-initialization hook - [[PR 987]](https://github.com/parthenon-hpc-lab/parthenon/pull/987) New tasking infrastructure and capabilities - [[PR 969]](https://github.com/parthenon-hpc-lab/parthenon/pull/969) New macro-based auto-naming of profiling regions and kernels diff --git a/src/interface/sparse_pack.hpp b/src/interface/sparse_pack.hpp index d1befba64b93..9753c401561b 100644 --- a/src/interface/sparse_pack.hpp +++ b/src/interface/sparse_pack.hpp @@ -1,5 +1,5 @@ //======================================================================================== -// (C) (or copyright) 2020-2023. Triad National Security, LLC. All rights reserved. +// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights reserved. // // This program was produced under U.S. Government contract 89233218CNA000001 for Los // Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC @@ -90,6 +90,7 @@ KOKKOS_INLINE_FUNCTION PackIdx operator+(T offset, PackIdx idx) { // device namespace variable_names { // Struct that all variable_name types should inherit from +constexpr int ANYDIM = -1234; // ANYDIM must be a slowest-moving index template struct base_t { KOKKOS_INLINE_FUNCTION @@ -98,6 +99,21 @@ struct base_t { KOKKOS_INLINE_FUNCTION explicit base_t(int idx1) : idx(idx1) {} + /* + for 2D:, (M, N), + idx(m, n) = N*m + n + for 3D: (L, M, N) + idx(l, m, n) = (M*l + m)*N + n + = l*M*N + m*N + n + */ + template ::value), + REQUIRES(sizeof...(Args) == sizeof...(NCOMP))> + KOKKOS_INLINE_FUNCTION explicit base_t(Args... args) + : idx(GetIndex_(std::forward(args)...)) { + static_assert(CheckArgs_(NCOMP...), + "All dimensions must be strictly positive, " + "except the first (slowest), which may be ANYDIM."); + } virtual ~base_t() = default; // All of these are just static methods so that there is no @@ -106,6 +122,11 @@ struct base_t { PARTHENON_FAIL("Need to implement your own name method."); return "error"; } + template + static constexpr auto GetDim() { + return std::get(std::make_tuple(NCOMP...)); + } + static std::vector GetShape() { return std::vector{NCOMP...}; } KOKKOS_INLINE_FUNCTION static bool regex() { return REGEX; } KOKKOS_INLINE_FUNCTION @@ -114,6 +135,23 @@ struct base_t { static int size() { return multiply::value; } const int idx; + + private: + template ::value)> + static constexpr bool CheckArgs_(int head, Tail... tail) { + return (... && (tail > 0)); + } + template + KOKKOS_INLINE_FUNCTION static auto GetIndex_(Args... args) { + int idx = 0; + ( + [&] { + idx *= NCOMP; + idx += args; + }(), + ...); + return idx; + } }; // An example variable name type that selects all variables available diff --git a/tst/unit/test_sparse_pack.cpp b/tst/unit/test_sparse_pack.cpp index 83ed9b0244f4..4efea18a5054 100644 --- a/tst/unit/test_sparse_pack.cpp +++ b/tst/unit/test_sparse_pack.cpp @@ -1,5 +1,5 @@ //======================================================================================== -// (C) (or copyright) 2020-2023. Triad National Security, LLC. All rights reserved. +// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights reserved. // // This program was produced under U.S. Government contract 89233218CNA000001 for Los // Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC @@ -77,13 +77,80 @@ struct v5 : public parthenon::variable_names::base_t { static std::string name() { return "v5"; } }; +using parthenon::variable_names::ANYDIM; +struct v7 : public parthenon::variable_names::base_t { + template + KOKKOS_INLINE_FUNCTION v7(Ts &&...args) + : parthenon::variable_names::base_t(std::forward(args)...) {} + static std::string name() { return "v7"; } +}; + } // namespace TEST_CASE("Test behavior of sparse packs", "[SparsePack]") { + constexpr int N = 6; + constexpr int NDIM = 3; + constexpr int NBLOCKS = 9; + + GIVEN("A tensor variable on a mesh") { + const std::vector tensor_shape{N, N, N, 3, 3}; + Metadata m_tensor({Metadata::Independent}, tensor_shape); + auto pkg = std::make_shared("Test package"); + pkg->AddField(m_tensor); + BlockList_t block_list = MakeBlockList(pkg, NBLOCKS, N, NDIM); + + MeshData mesh_data("base"); + mesh_data.Set(block_list, nullptr); + + WHEN("We initialize the independent variables by hand and deallocate one") { + auto ib = block_list[0]->cellbounds.GetBoundsI(IndexDomain::entire); + auto jb = block_list[0]->cellbounds.GetBoundsJ(IndexDomain::entire); + auto kb = block_list[0]->cellbounds.GetBoundsK(IndexDomain::entire); + for (int b = 0; b < NBLOCKS; ++b) { + auto &pmb = block_list[b]; + auto &pmbd = pmb->meshblock_data.Get(); + auto var = pmbd->Get("v7"); + auto var5 = var.data.Get<5>(); + int slower_rank = var5.GetDim(5); + int faster_rank = var5.GetDim(4); + par_for( + loop_pattern_mdrange_tag, "initializev7", DevExecSpace(), kb.s, kb.e, jb.s, + jb.e, ib.s, ib.e, KOKKOS_LAMBDA(int k, int j, int i) { + for (int l = 0; l < slower_rank; ++l) { + for (int m = 0; m < faster_rank; ++m) { + Real n = m + 1e1 * l; + var5(l, m, k, j, i) = n; + } + } + }); + } + THEN("A sparse pack can correctly index into tensor types") { + auto desc = parthenon::MakePackDescriptor(pkg.get()); + auto sparse_pack = desc.GetPack(&mesh_data); + int nwrong = 0; + int nl = tensor_shape[4]; + int nm = tensor_shape[3]; + par_reduce( + loop_pattern_mdrange_tag, "check vector", DevExecSpace(), 0, + sparse_pack.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, + KOKKOS_LAMBDA(int b, int k, int j, int i, int <ot) { + // 0-th is ANYDIM, 1st is 3. + for (int l = 0; l < nl; ++l) { + for (int m = 0; m < nm; ++m) { + Real n = m + 1e1 * l; + if (sparse_pack(b, v7(l, m), k, j, i) != n) { + ltot += 1; + } + } + } + }, + nwrong); + REQUIRE(nwrong == 0); + } + } + } + GIVEN("A set of meshblocks and meshblock and mesh data") { - constexpr int N = 6; - constexpr int NDIM = 3; - constexpr int NBLOCKS = 9; const std::vector scalar_shape{N, N, N}; const std::vector vector_shape{N, N, N, 3}; @@ -122,6 +189,7 @@ TEST_CASE("Test behavior of sparse packs", "[SparsePack]") { }); } } + // Deallocate a variable on an arbitrary block block_list[2]->DeallocateSparse("v3");