WIP: Generalize Sparse Indices

CHANGELOG.md

@@ -41,6 +41,7 @@
 
 
 ### Incompatibilities (i.e. breaking changes)
+- [[PR 1398]](https://github.com/parthenon-hpc-lab/parthenon/pull/1398) Generalize sparse identifiers and add sparse control groups. Downstream sparse-state access may need to change from `sparse_idx` to `sparse_idx()`.
 - [[PR 1351]](https://github.com/parthenon-hpc-lab/parthenon/pull/1351) Bump Kokkos 5 & C++20
 - [[PR 1377]](https://github.com/parthenon-hpc-lab/parthenon/pull/1377) Extend Initialization Hierarchy
 - [[PR 1376]](https://github.com/parthenon-hpc-lab/parthenon/pull/1376) Refactor SwarmPacks

dev/plan_histories/1398.md

@@ -0,0 +1,53 @@
+<!-- This file was made in part with generative AI. -->
+
+## Goal
+
+Generalize sparse identifiers so Parthenon can represent more than a single integer
+index per sparse field, while keeping the existing scalar behavior and restart behavior
+unchanged for current codes. At the same time, introduce a cleaner control-group model
+for sparse fields so allocation and deallocation can operate on groups of controlling
+fields rather than only a single controller label.
+
+## Functional Plan
+
+- Introduce a dedicated `SparseID` value type and thread it through the sparse field
+  model instead of storing raw integers directly in the core interfaces.
+- Keep the sparse label format backward compatible for current scalar IDs so existing
+  sparse names continue to look the same.
+- Add support for multi-component sparse IDs in a way that remains device-visible and
+  inexpensive enough to be used in hot code paths.
+- Mark the sparse-ID comparison operators as device-available so pack selection and
+  other device code can reason about sparse ordering without special handling.
+- Keep the host-facing sparse APIs compatible with existing scalar call sites by
+  allowing integer inputs to be converted to `SparseID` at the API boundary.
+- Update sparse pools so they store and operate on typed sparse IDs internally rather
+  than treating sparse IDs as anonymous integers.
+- Add a sparse-ID projection rule for sparse pools so a pool can group fields by a
+  controller key derived from the sparse index space.
+- Extend `StateDescriptor` so it resolves sparse-pool control groups into the same
+  controller-reverse-map machinery used by ordinary fields.
+- Treat control groups as the source of truth for grouped sparse control, while keeping
+  the legacy single-controller accessor available as a compatibility shim.
+- Keep per-field sparse deallocation counters intact, but evaluate deallocation at the
+  group level by requiring every control member to be ready before the group is removed.
+- Make sparse deallocation ignore dense-only controller groups so dense fields do not
+  enter the sparse dealloc path.
+- Preserve restart exactness by keeping the sparse deallocation metadata and sparse
+  allocation behavior compatible with existing restart files.
+- Keep the existing sparse-advection example and restart flow working as the primary
+  exactness check while the sparse-control machinery changes underneath it.
+- Add focused unit coverage for the new sparse-ID type and for grouped sparse control.
+- Update sparse documentation so the split between sparse IDs, sparse pools, and
+  controlling groups is described clearly.
+
+## Scope Boundaries
+
+- Do not change the public downstream package API unless it is required for backward
+  compatibility with existing scalar sparse call sites.
+- Keep dense fields and sparse fields conceptually separate; only unify their internal
+  registration and control bookkeeping where that reduces duplication.
+- Do not change the numerical update kernels in the example problems unless a restart or
+  sparse-control bug forces a targeted fix.
+- Preserve the current sparse-advection restart gold output as the exactness oracle for
+  this PR.
+

example/advection/advection_package.cpp

@@ -11,6 +11,8 @@
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
 
+// This file was made in part with generative AI.
+
 #include <algorithm>
 #include <cmath>
 #include <iostream>
@@ -390,8 +392,10 @@ void PostFill(MeshBlockData<Real> *rc) {
     IndexRange kb = pmb->cellbounds.GetBoundsK(IndexDomain::entire);
 
     // check that we have the sparse indices we want
-    pmb->AllocSparseID("one_minus_sqrt_one_minus_advected_sq", 12);
-    pmb->AllocSparseID("one_minus_sqrt_one_minus_advected_sq", 37);
+    pmb->AllocSparseID("one_minus_sqrt_one_minus_advected_sq",
+                       parthenon::SparseID::Scalar(12));
+    pmb->AllocSparseID("one_minus_sqrt_one_minus_advected_sq",
+                       parthenon::SparseID::Scalar(37));
 
     // packing in principle unnecessary/convoluted here and just done for demonstration
     std::vector<std::string> vars(
@@ -402,7 +406,9 @@ void PostFill(MeshBlockData<Real> *rc) {
     const int in = imap.get("one_minus_advected_sq").first;
     // we can get sparse fields either by specifying base name and sparse id, or the full
     // name
-    const int out12 = imap.get("one_minus_sqrt_one_minus_advected_sq", 12).first;
+    const int out12 =
+        imap.get("one_minus_sqrt_one_minus_advected_sq", parthenon::SparseID::Scalar(12))
+            .first;
     const int out37 = imap.get("one_minus_sqrt_one_minus_advected_sq_37").first;
     const auto num_vars = rc->Get("advected").data.GetDim(4);
     pmb->par_for(

example/calculate_pi/calculate_pi.cpp

@@ -11,6 +11,8 @@
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
 
+// This file was made in part with generative AI.
+
 // Standard Includes
 #include <iostream>
 #include <memory>
@@ -72,8 +74,8 @@ void SetInOrOut(MeshBlockData<Real> *rc) {
       return;
     }
 
-    pmb->AllocSparseID("in_or_out", 0);
-    v = rc->Get("in_or_out", 0).data;
+    pmb->AllocSparseID("in_or_out", parthenon::SparseID::Scalar(0));
+    v = rc->Get("in_or_out", parthenon::SparseID::Scalar(0)).data;
   } else {
     v = rc->Get("in_or_out").data;
   }

example/fine_advection/parthenon_app_inputs.cpp

@@ -11,6 +11,8 @@
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
 
+// This file was made in part with generative AI.
+
 #include <memory>
 #include <sstream>
 #include <string>
@@ -65,7 +67,7 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
   if (do_regular_advection) {
     const int sparse_size = pkg->Param<int>("sparse_size");
     for (int s = 0; s < sparse_size; ++s)
-      pmb->AllocSparseID(phi::name(), s);
+      pmb->AllocSparseID(phi::name(), SparseID::Scalar(s));
   }
   static auto desc = parthenon::MakePackDescriptor<phi, phi_fine, C, D>(data.get());
   auto pack = desc.GetPack(data.get());

example/sparse_advection/parthenon_app_inputs.cpp

@@ -10,6 +10,8 @@
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
 
+// This file was made in part with generative AI.
+
 #include <algorithm>
 #include <cstdio>
 #include <limits>
@@ -89,15 +91,16 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
       VariablePack<Real> v;
 
       if (restart_test) {
-        pmb->AllocSparseID("shape_shift", 1);
-        pmb->AllocSparseID("shape_shift", 3);
-        pmb->AllocSparseID("shape_shift", 4);
+        pmb->AllocSparseID("shape_shift", parthenon::SparseID::Scalar(1));
+        pmb->AllocSparseID("shape_shift", parthenon::SparseID::Scalar(3));
+        pmb->AllocSparseID("shape_shift", parthenon::SparseID::Scalar(4));
 
         v = data->PackVariables(
             std::vector<std::string>{"dense_A", "dense_B", "shape_shift"});
       } else {
-        pmb->AllocSparseID("sparse", f);
-        v = data->PackVariables(std::vector<std::string>{MakeVarLabel("sparse", f)});
+        pmb->AllocSparseID("sparse", parthenon::SparseID::Scalar(f));
+        v = data->PackVariables(std::vector<std::string>{
+            parthenon::MakeVarLabel("sparse", parthenon::SparseID::Scalar(f))});
       }
 
       pmb->par_for(
@@ -127,7 +130,7 @@ void PostStepDiagnosticsInLoop(Mesh *mesh, ParameterInput *pin, const SimTime &t
   for (auto &pmb : mesh->block_list) {
     auto rc = pmb->meshblock_data.Get(); // get base container
     for (int i = 0; i < n; ++i) {
-      if (rc->IsAllocated("sparse", i)) {
+      if (rc->IsAllocated("sparse", parthenon::SparseID::Scalar(i))) {
         num_allocated[i] += 1;
       }
     }

example/sparse_advection/sparse_advection_package.cpp

@@ -11,6 +11,8 @@
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
 
+// This file was made in part with generative AI.
+
 #include <algorithm>
 #include <cmath>
 #include <limits>
@@ -76,7 +78,7 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
     SparsePool pool("sparse", m);
 
     for (int sid = 0; sid < NUM_FIELDS; ++sid) {
-      pool.Add(sid);
+      pool.Add(parthenon::SparseID::Scalar(sid));
     }
     pkg->AddSparsePool(pool);
   }
@@ -93,10 +95,11 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
                        Metadata::FillGhost, Metadata::Sparse});
 
     SparsePool pool("shape_shift", m_sparse);
-    pool.Add(1, std::vector<int>{1}, std::vector<std::string>{"scalar"});
-    pool.Add(3, std::vector<int>{3}, Metadata::Vector,
+    pool.Add(parthenon::SparseID::Scalar(1), std::vector<int>{1},
+             std::vector<std::string>{"scalar"});
+    pool.Add(parthenon::SparseID::Scalar(3), std::vector<int>{3}, Metadata::Vector,
              std::vector<std::string>{"vec_x", "vec_y", "vec_z"});
-    pool.Add(4, std::vector<int>{4}, Metadata::Vector);
+    pool.Add(parthenon::SparseID::Scalar(4), std::vector<int>{4}, Metadata::Vector);
 
     pkg->AddSparsePool(pool);
   }
@@ -216,7 +219,7 @@ TaskStatus CalculateFluxes(std::shared_ptr<MeshBlockData<Real>> &rc) {
 
         for (int n = 0; n < nvar; n++) {
           if (!v.IsAllocated(n)) continue;
-          const auto this_v = vx[v(n).sparse_id % NUM_FIELDS];
+          const auto this_v = vx[v(n).sparse_id() % NUM_FIELDS];
           par_for_inner(member, ib.s, ib.e + 1, [&](const int i) {
             v.flux(X1DIR, n, k, j, i) = (this_v > 0.0 ? ql(n, i) : qr(n, i)) * this_v;
           });
@@ -246,7 +249,7 @@ TaskStatus CalculateFluxes(std::shared_ptr<MeshBlockData<Real>> &rc) {
 
           for (int n = 0; n < nvar; n++) {
             if (!v.IsAllocated(n)) continue;
-            const auto this_v = vy[v(n).sparse_id % NUM_FIELDS];
+            const auto this_v = vy[v(n).sparse_id() % NUM_FIELDS];
             par_for_inner(member, ib.s, ib.e, [&](const int i) {
               v.flux(X2DIR, n, k, j, i) = (this_v > 0.0 ? ql(n, i) : qr(n, i)) * this_v;
             });

src/interface/meshblock_data.cpp

@@ -11,6 +11,8 @@
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
 
+// This file was made in part with generative AI.
+
 #include "interface/meshblock_data.hpp"
 
 #include <algorithm>
@@ -44,7 +46,7 @@ namespace parthenon {
 /// @param sparse_id the sparse id of the variable
 template <typename T>
 void MeshBlockData<T>::AddField(const std::string &base_name, const Metadata &metadata,
-                                int sparse_id) {
+                                SparseID sparse_id) {
   auto pvar = std::make_shared<Variable<T>>(base_name, metadata, sparse_id, pmy_block);
   Add(pvar);
 

src/interface/meshblock_data.hpp

@@ -13,7 +13,7 @@
 #ifndef INTERFACE_MESHBLOCK_DATA_HPP_
 #define INTERFACE_MESHBLOCK_DATA_HPP_
 
-// This file was made in part with generative AI
+// This file was made in part with generative AI.
 
 #include <algorithm>
 #include <limits>
@@ -266,7 +266,8 @@ class MeshBlockData {
 
   const auto &GetUidMap() const { return varUidMap_; }
 
-  Variable<T> &Get(const std::string &base_name, int sparse_id = InvalidSparseID) const {
+  Variable<T> &Get(const std::string &base_name,
+                   SparseID sparse_id = InvalidSparseID) const {
     return *GetVarPtr(MakeVarLabel(base_name, sparse_id));
   }
   Variable<T> &Get(const Uid_t &uid) const { return *(varUidMap_.at(uid)); }
@@ -286,7 +287,8 @@ class MeshBlockData {
     return it->second->IsAllocated();
   }
 
-  inline bool IsAllocated(std::string const &base_name, int sparse_id) const noexcept {
+  inline bool IsAllocated(std::string const &base_name,
+                          SparseID sparse_id) const noexcept {
     return IsAllocated(MakeVarLabel(base_name, sparse_id));
   }
 
@@ -585,14 +587,14 @@ class MeshBlockData {
 
  private:
   void AddField(const std::string &base_name, const Metadata &metadata,
-                int sparse_id = InvalidSparseID);
+                SparseID sparse_id);
 
   void Add(std::shared_ptr<Variable<T>> var) noexcept;
 
   std::shared_ptr<Variable<T>> AllocateSparse(std::string const &label,
                                               bool flag_uninitialized = false);
   std::shared_ptr<Variable<T>> AllocSparseID(std::string const &base_name,
-                                             const int sparse_id) {
+                                             const SparseID sparse_id) {
     return AllocateSparse(MakeVarLabel(base_name, sparse_id));
   }
   void DeallocateSparse(std::string const &label);

src/interface/sparse_pool.cpp

@@ -11,6 +11,8 @@
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
 
+// This file was made in part with generative AI.
+
 #include <memory>
 #include <string>
 #include <vector>
@@ -23,7 +25,7 @@
 namespace parthenon {
 
 SparsePool::SparsePool(const std::string &base_name, const Metadata &metadata,
-                       const std::vector<int> &sparse_ids,
+                       const std::vector<SparseID> &sparse_ids,
                        const std::vector<std::vector<int>> &shapes,
                        const std::vector<MetadataFlag> &vector_tensor_flags,
                        const std::vector<std::vector<std::string>> &component_labels,
@@ -92,7 +94,7 @@ MakeSparseVarMetadataImpl(Metadata *in, const std::vector<int> &shape,
   return this_metadata;
 }
 
-const Metadata &SparsePool::AddImpl(int sparse_id, const std::vector<int> &shape,
+const Metadata &SparsePool::AddImpl(SparseID sparse_id, const std::vector<int> &shape,
                                     const MetadataFlag *vector_tensor,
                                     const std::vector<std::string> &component_labels) {
   PARTHENON_REQUIRE_THROWS(sparse_id != InvalidSparseID,
@@ -107,22 +109,22 @@ const Metadata &SparsePool::AddImpl(int sparse_id, const std::vector<int> &shape
   }
 
   const auto ins = pool_.insert({sparse_id, *this_metadata});
-  PARTHENON_REQUIRE_THROWS(ins.second, "Tried to add sparse ID " +
-                                           std::to_string(sparse_id) +
-                                           " to sparse pool '" + base_name_ +
+  PARTHENON_REQUIRE_THROWS(ins.second, "Tried to add sparse field '" +
+                                           MakeVarLabel(base_name_, sparse_id) +
+                                           "' to sparse pool '" + base_name_ +
                                            "', but this sparse ID already exists");
 
   return ins.first->second;
 }
 
-const Metadata &SparsePool::Add(int sparse_id, const Metadata &md) {
+const Metadata &SparsePool::Add(SparseID sparse_id, const Metadata &md) {
   PARTHENON_REQUIRE_THROWS(sparse_id != InvalidSparseID,
                            "Tried to add InvalidSparseID to sparse pool " + base_name_);
 
   const auto ins = pool_.insert({sparse_id, md});
-  PARTHENON_REQUIRE_THROWS(ins.second, "Tried to add sparse ID " +
-                                           std::to_string(sparse_id) +
-                                           " to sparse pool '" + base_name_ +
+  PARTHENON_REQUIRE_THROWS(ins.second, "Tried to add sparse field '" +
+                                           MakeVarLabel(base_name_, sparse_id) +
+                                           "' to sparse pool '" + base_name_ +
                                            "', but this sparse ID already exists");
 
   return ins.first->second;