Merge pull request #998 from parthenon-hpc-lab/mg/tensor_indices

tensor indices added to sparse pack

Author: Yurlungur

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

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 <bool REGEX, int... NCOMP>
 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 <typename... Args, REQUIRES(all_implement<integral(Args...)>::value),
+            REQUIRES(sizeof...(Args) == sizeof...(NCOMP))>
+  KOKKOS_INLINE_FUNCTION explicit base_t(Args... args)
+      : idx(GetIndex_(std::forward<Args>(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 <int idx>
+  static constexpr auto GetDim() {
+    return std::get<sizeof...(NCOMP) - idx>(std::make_tuple(NCOMP...));
+  }
+  static std::vector<int> GetShape() { return std::vector<int>{NCOMP...}; }
   KOKKOS_INLINE_FUNCTION
   static bool regex() { return REGEX; }
   KOKKOS_INLINE_FUNCTION
@@ -114,6 +135,23 @@ struct base_t {
   static int size() { return multiply<NCOMP...>::value; }
 
   const int idx;
+
+ private:
+  template <typename... Tail, REQUIRES(all_implement<integral(Tail...)>::value)>
+  static constexpr bool CheckArgs_(int head, Tail... tail) {
+    return (... && (tail > 0));
+  }
+  template <class... Args>
+  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

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<false> {
   static std::string name() { return "v5"; }
 };
 
+using parthenon::variable_names::ANYDIM;
+struct v7 : public parthenon::variable_names::base_t<false, ANYDIM, 3> {
+  template <class... Ts>
+  KOKKOS_INLINE_FUNCTION v7(Ts &&...args)
+      : parthenon::variable_names::base_t<false, ANYDIM, 3>(std::forward<Ts>(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<int> tensor_shape{N, N, N, 3, 3};
+    Metadata m_tensor({Metadata::Independent}, tensor_shape);
+    auto pkg = std::make_shared<StateDescriptor>("Test package");
+    pkg->AddField<v7>(m_tensor);
+    BlockList_t block_list = MakeBlockList(pkg, NBLOCKS, N, NDIM);
+
+    MeshData<Real> 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<v7>(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 &ltot) {
+              // 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<int> scalar_shape{N, N, N};
     const std::vector<int> 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");