Functionality for Indexer masking (#901)

* Start on in one AMR

* Fix MPI bug

* Make things work w/o MPI

* small

* update changelog

* Make AMR helper functions free to alleviate compilation issues on GPU

* Fix MPI compilation error

* Actually compile with MPI...

* Fix indexing bug

* Remove unused vars

* Fix indexing bug

* Make things work for tensor variables

* CellVariable to Variable

* Respond to Philipp's comments

* Maybe fix compiler issues

* Second half of fix

* Add fence back

* Update copyright dates

* Pass dealloc count when sending same to same

* Almost working, but differs on step 301 after remesh

* Add deletion check

* Fix a bunch of bugs with a couple changes

* Format, lint, changelog

* Remove debugging MPI Barrier

* Add maximum number of iterations for deletion check

* Remove commented lines and clean up

* Fix MPI AMR bug related to not passing dereference count when nothing is allocated on a block

* format and lint

* Fix bug when sparse is disabled

* overloads to create packs with meshdata

* Always apply fine boundary conditions

* Split pack descriptor into separate header

* Actually include meshdata and Meshblockdata

* more includes

* more includes

* Reworking neighbor finding

* Start on adding morton numbers

* Make stuff private

* format and lint

* Add real comparison operators

* Split out morton number

* Fix bug when shift is larger than bit size

* Add some more functionality to logical location

* Start on unit test

* Actually compare Morton numbers

* Add neighbor check, untested

* Fix TE neighbor finding and add tests

* Update copyrights

* Add logical location to NeighborBlock

* add interleave constant test and fix bug

* Add bit interleave test

* update changelog

* Explicitly start at zero

* Add routines for calculating ownership and tests

* Remove comments

* Add another ownership test

* Format and lint

* Add mask and start testing

* switch to bits

* Separate logical location header

* Switch to class

* reserve

* Seemingly working masking code

* Somewhat more extensive tests

* Add ownership to neighbor blocks

* Update to mask and run over all indices, still some bug were some neighbor block ownership is not being set

* Ownership model passing for cell centered vars

* format and lint

* Fix errors hidden by MPI ifdef

* Remove printf

* Add sparse seed nans flag to output

* changelog

* change signaling NaN to quiet NaN

* format

* format and lint

* Deal with coarse to fine corner issue

* changelog

* Format and lint

* Remove commented out code

* Mask prolongation and restriction

* Add MakePackDescriptor overload

* Add possible neighbors on periodic boundaries

* Deal with periodic boundaries

* Explicitly deal with periodic vs non-periodic directions

* format and lint

* Actually pass base grid information

* Act on Philipp's comments

* Fix linter error

---------

Co-authored-by: Philipp Grete <[email protected]>
Co-authored-by: Jonah Maxwell Miller <[email protected]>

Author: 3 people

CHANGELOG.md

@@ -24,6 +24,7 @@
 - [[PR 890]](https://github.com/parthenon-hpc-lab/parthenon/pull/890) Fix bugs in sparse communication and prolongation
 
 ### Infrastructure (changes irrelevant to downstream codes)
+- [[PR 901]](https://github.com/parthenon-hpc-lab/parthenon/pull/901) Implement shared element ownership model
 
 ### Removed (removing behavior/API/varaibles/...)
 

src/CMakeLists.txt

@@ -135,6 +135,7 @@ add_library(parthenon
   interface/swarm.hpp
   interface/swarm_boundaries.hpp
   interface/swarm_device_context.hpp
+  interface/make_pack_descriptor.hpp
   interface/metadata.cpp
   interface/metadata.hpp
   interface/packages.hpp

src/bvals/boundary_conditions_generic.hpp

@@ -21,6 +21,7 @@
 #include <vector>
 
 #include "basic_types.hpp"
+#include "interface/make_pack_descriptor.hpp"
 #include "interface/meshblock_data.hpp"
 #include "interface/sparse_pack.hpp"
 #include "mesh/domain.hpp"

src/bvals/bvals.hpp

@@ -90,6 +90,8 @@ class BoundaryBase {
   RegionSize block_size_;
   ParArrayND<Real> sarea_[2];
 
+  void SetNeighborOwnership();
+
  private:
   // calculate 3x shared static data members when constructing only the 1st class instance
   // int maxneighbor_=BufferID() computes ni[] and then calls bufid[]=CreateBufferID()

src/bvals/bvals_base.cpp

@@ -368,6 +368,7 @@ void BoundaryBase::SearchAndSetNeighbors(MeshBlockTree &tree, int *ranklist,
     }
   }
   if (block_size_.nx2 == 1) {
+    SetNeighborOwnership();
     Kokkos::Profiling::popRegion(); // SearchAndSetNeighbors
     return;
   }
@@ -502,6 +503,7 @@ void BoundaryBase::SearchAndSetNeighbors(MeshBlockTree &tree, int *ranklist,
   }
 
   if (block_size_.nx3 == 1) {
+    SetNeighborOwnership();
     Kokkos::Profiling::popRegion(); // SearchAndSetNeighbors
     return;
   }
@@ -623,7 +625,25 @@ void BoundaryBase::SearchAndSetNeighbors(MeshBlockTree &tree, int *ranklist,
       }
     }
   }
+
+  SetNeighborOwnership();
   Kokkos::Profiling::popRegion(); // SearchAndSetNeighbors
 }
 
+void BoundaryBase::SetNeighborOwnership() {
+  // Set neighbor block ownership
+  std::set<LogicalLocation> allowed_neighbors;
+  allowed_neighbors.insert(loc); // Insert the location of this block
+  for (int n = 0; n < nneighbor; ++n)
+    allowed_neighbors.insert(neighbor[n].loc);
+  // Although the neighbor blocks abut more blocks than are contained in this
+  // list, the unaccounted for blocks cannot impact the ownership of elements
+  // that are shared with *this
+  RootGridInfo rg_info = pmy_mesh_->GetRootGridInfo();
+  for (int n = 0; n < nneighbor; ++n) {
+    neighbor[n].ownership =
+        DetermineOwnership(neighbor[n].loc, allowed_neighbors, rg_info);
+    neighbor[n].ownership.initialized = true;
+  }
+}
 } // namespace parthenon

src/bvals/bvals_interfaces.hpp

@@ -143,6 +143,8 @@ struct NeighborBlock { // aggregate and POD type. Inheritance breaks standard-la
   int bufid, eid, targetid;
   BoundaryFace fid;
   LogicalLocation loc;
+  block_ownership_t ownership;
+
   void SetNeighbor(LogicalLocation inloc, int irank, int ilevel, int igid, int ilid,
                    int iox1, int iox2, int iox3, NeighborConnect itype, int ibid,
                    int itargetid, int ifi1 = 0, int ifi2 = 0);

src/bvals/comms/bnd_info.cpp

@@ -41,9 +41,9 @@ enum class InterfaceType { SameToSame, CoarseToFine, FineToCoarse };
 namespace parthenon {
 
 template <InterfaceType INTERFACE>
-Indexer6D CalcLoadIndices(const NeighborIndexes &ni, TopologicalElement el,
-                          std::array<int, 3> tensor_shape,
-                          const parthenon::IndexShape &shape) {
+SpatiallyMaskedIndexer6D CalcLoadIndices(const NeighborIndexes &ni, TopologicalElement el,
+                                         std::array<int, 3> tensor_shape,
+                                         const parthenon::IndexShape &shape) {
   IndexDomain interior = IndexDomain::interior;
   std::array<IndexRange, 3> bounds{shape.GetBoundsI(interior, el),
                                    shape.GetBoundsJ(interior, el),
@@ -82,25 +82,30 @@ Indexer6D CalcLoadIndices(const NeighborIndexes &ni, TopologicalElement el,
                  // offset in some direction
     } else if (block_offset[dir] > 0) {
       s[dir] = bounds[dir].e - Globals::nghost + 1 - top_offset[dir];
-      e[dir] = bounds[dir].e - top_offset[dir];
+      e[dir] = bounds[dir].e;
     } else {
-      s[dir] = bounds[dir].s + top_offset[dir];
+      s[dir] = bounds[dir].s;
       e[dir] = bounds[dir].s + Globals::nghost - 1 + top_offset[dir];
     }
   }
-  return Indexer6D({0, tensor_shape[0] - 1}, {0, tensor_shape[1] - 1},
-                   {0, tensor_shape[2] - 1}, {s[2], e[2]}, {s[1], e[1]}, {s[0], e[0]});
+  block_ownership_t owns(true);
+  return SpatiallyMaskedIndexer6D(owns, {0, tensor_shape[0] - 1},
+                                  {0, tensor_shape[1] - 1}, {0, tensor_shape[2] - 1},
+                                  {s[2], e[2]}, {s[1], e[1]}, {s[0], e[0]});
 }
 
 template <InterfaceType INTERFACE, bool PROLONGATEORRESTRICT = false>
-Indexer6D CalcSetIndices(const NeighborIndexes &ni, LogicalLocation loc,
-                         TopologicalElement el, std::array<int, 3> tensor_shape,
-                         const parthenon::IndexShape &shape) {
+SpatiallyMaskedIndexer6D
+CalcSetIndices(const NeighborBlock &nb, LogicalLocation loc, TopologicalElement el,
+               std::array<int, 3> tensor_shape, const parthenon::IndexShape &shape) {
+  const auto &ni = nb.ni;
   IndexDomain interior = IndexDomain::interior;
   std::array<IndexRange, 3> bounds{shape.GetBoundsI(interior, el),
                                    shape.GetBoundsJ(interior, el),
                                    shape.GetBoundsK(interior, el)};
 
+  std::array<int, 3> top_offset{TopologicalOffsetI(el), TopologicalOffsetJ(el),
+                                TopologicalOffsetK(el)};
   std::array<int, 3> block_offset{ni.ox1, ni.ox2, ni.ox3};
   // This is gross, but the face offsets do not contain the correct
   // information for going from coarse to fine and the neighbor block
@@ -128,15 +133,28 @@ Indexer6D CalcSetIndices(const NeighborIndexes &ni, LogicalLocation loc,
       }
       ++off_idx;
     } else if (block_offset[dir] > 0) {
-      s[dir] = bounds[dir].e + 1;
+      s[dir] = bounds[dir].e + 1 - top_offset[dir];
       e[dir] = bounds[dir].e + ghosts;
     } else {
       s[dir] = bounds[dir].s - ghosts;
-      e[dir] = bounds[dir].s - 1;
+      e[dir] = bounds[dir].s - 1 + top_offset[dir];
     }
   }
-  return Indexer6D({0, tensor_shape[0] - 1}, {0, tensor_shape[1] - 1},
-                   {0, tensor_shape[2] - 1}, {s[2], e[2]}, {s[1], e[1]}, {s[0], e[0]});
+  int sox1 = -ni.ox1;
+  int sox2 = -ni.ox2;
+  int sox3 = -ni.ox3;
+  if (INTERFACE == InterfaceType::CoarseToFine) {
+    // For coarse to fine interfaces, we are passing zones from only an
+    // interior corner of the cell, never an entire face or edge
+    if (sox1 == 0) sox1 = logic_loc[0] % 2 == 1 ? 1 : -1;
+    if (sox2 == 0) sox2 = logic_loc[1] % 2 == 1 ? 1 : -1;
+    if (sox3 == 0) sox3 = logic_loc[2] % 2 == 1 ? 1 : -1;
+  }
+  block_ownership_t owns =
+      GetIndexRangeMaskFromOwnership(el, nb.ownership, sox1, sox2, sox3);
+  return SpatiallyMaskedIndexer6D(owns, {0, tensor_shape[0] - 1},
+                                  {0, tensor_shape[1] - 1}, {0, tensor_shape[2] - 1},
+                                  {s[2], e[2]}, {s[1], e[1]}, {s[0], e[0]});
 }
 
 int GetBufferSize(std::shared_ptr<MeshBlock> pmb, const NeighborBlock &nb,
@@ -152,9 +170,9 @@ int GetBufferSize(std::shared_ptr<MeshBlock> pmb, const NeighborBlock &nb,
   const int isize = cb.ie(in) - cb.is(in) + 2;
   const int jsize = cb.je(in) - cb.js(in) + 2;
   const int ksize = cb.ke(in) - cb.ks(in) + 2;
-  return (nb.ni.ox1 == 0 ? isize : Globals::nghost) *
-         (nb.ni.ox2 == 0 ? jsize : Globals::nghost) *
-         (nb.ni.ox3 == 0 ? ksize : Globals::nghost) * v->GetDim(6) * v->GetDim(5) *
+  return (nb.ni.ox1 == 0 ? isize : Globals::nghost + 1) *
+         (nb.ni.ox2 == 0 ? jsize : Globals::nghost + 1) *
+         (nb.ni.ox3 == 0 ? ksize : Globals::nghost + 1) * v->GetDim(6) * v->GetDim(5) *
          v->GetDim(4) * topo_comp;
 }
 
@@ -262,29 +280,29 @@ BndInfo BndInfo::GetSetBndInfo(std::shared_ptr<MeshBlock> pmb, const NeighborBlo
     if (nb.snb.level == mylevel) {
       out.var = v->data.Get();
       out.idxer[idx] = CalcSetIndices<InterfaceType::SameToSame>(
-          nb.ni, pmb->loc, el, {Nt, Nu, Nv}, pmb->cellbounds);
+          nb, pmb->loc, el, {Nt, Nu, Nv}, pmb->cellbounds);
       if (restricted) {
         out.refinement_op = RefinementOp_t::Restriction;
         out.prores_idxer[static_cast<int>(el)] =
             CalcSetIndices<InterfaceType::SameToSame, true>(
-                nb.ni, pmb->loc, el, {Nt, Nu, Nv}, pmb->c_cellbounds);
+                nb, pmb->loc, el, {Nt, Nu, Nv}, pmb->c_cellbounds);
       }
     } else if (nb.snb.level < mylevel) {
       out.idxer[idx] = CalcSetIndices<InterfaceType::CoarseToFine>(
-          nb.ni, pmb->loc, el, {Nt, Nu, Nv}, pmb->c_cellbounds);
+          nb, pmb->loc, el, {Nt, Nu, Nv}, pmb->c_cellbounds);
       out.var = v->coarse_s.Get();
       out.refinement_op = RefinementOp_t::Prolongation;
       out.prores_idxer[idx] = CalcSetIndices<InterfaceType::CoarseToFine, true>(
-          nb.ni, pmb->loc, el, {Nt, Nu, Nv}, pmb->c_cellbounds);
+          nb, pmb->loc, el, {Nt, Nu, Nv}, pmb->c_cellbounds);
     } else {
       out.var = v->data.Get();
       out.idxer[idx] = CalcSetIndices<InterfaceType::FineToCoarse>(
-          nb.ni, pmb->loc, el, {Nt, Nu, Nv}, pmb->cellbounds);
+          nb, pmb->loc, el, {Nt, Nu, Nv}, pmb->cellbounds);
       if (restricted) {
         out.refinement_op = RefinementOp_t::Restriction;
         out.prores_idxer[static_cast<int>(el)] =
             CalcSetIndices<InterfaceType::FineToCoarse, true>(
-                nb.ni, pmb->loc, el, {Nt, Nu, Nv}, pmb->c_cellbounds);
+                nb, pmb->loc, el, {Nt, Nu, Nv}, pmb->c_cellbounds);
       }
     }
   }
@@ -302,7 +320,7 @@ BndInfo BndInfo::GetSetBndInfo(std::shared_ptr<MeshBlock> pmb, const NeighborBlo
     for (auto el : {TE::CC, TE::F1, TE::F2, TE::F3, TE::E1, TE::E2, TE::E3, TE::NN}) {
       out.prores_idxer[static_cast<int>(el)] =
           CalcSetIndices<InterfaceType::CoarseToFine, true>(
-              nb.ni, pmb->loc, el, {Nt, Nu, Nv}, pmb->c_cellbounds);
+              nb, pmb->loc, el, {Nt, Nu, Nv}, pmb->c_cellbounds);
     }
   }
 
@@ -359,8 +377,10 @@ BndInfo BndInfo::GetSendCCFluxCor(std::shared_ptr<MeshBlock> pmb, const Neighbor
 
   out.var = v->flux[out.dir];
   out.coords = pmb->coords;
-  out.idxer[0] = Indexer6D({0, out.var.GetDim(6) - 1}, {0, out.var.GetDim(5) - 1},
-                           {0, out.var.GetDim(4) - 1}, {sk, ek}, {sj, ej}, {si, ei});
+  block_ownership_t owns(true);
+  out.idxer[0] = SpatiallyMaskedIndexer6D(
+      owns, {0, out.var.GetDim(6) - 1}, {0, out.var.GetDim(5) - 1},
+      {0, out.var.GetDim(4) - 1}, {sk, ek}, {sj, ej}, {si, ei});
   return out;
 }
 
@@ -435,8 +455,10 @@ BndInfo BndInfo::GetSetCCFluxCor(std::shared_ptr<MeshBlock> pmb, const NeighborB
   out.var = v->flux[out.dir];
 
   out.coords = pmb->coords;
-  out.idxer[0] = Indexer6D({0, out.var.GetDim(6) - 1}, {0, out.var.GetDim(5) - 1},
-                           {0, out.var.GetDim(4) - 1}, {sk, ek}, {sj, ej}, {si, ei});
+  block_ownership_t owns(true);
+  out.idxer[0] = SpatiallyMaskedIndexer6D(
+      owns, {0, out.var.GetDim(6) - 1}, {0, out.var.GetDim(5) - 1},
+      {0, out.var.GetDim(4) - 1}, {sk, ek}, {sj, ej}, {si, ei});
   return out;
 }
 

src/bvals/comms/bnd_info.hpp

@@ -42,9 +42,10 @@ class Variable;
 
 struct BndInfo {
   int ntopological_elements = 1;
-  Indexer6D idxer[3];
-  Indexer6D prores_idxer[10]; // Has to be large enough to allow for maximum integer
-                              // conversion of TopologicalElements
+  SpatiallyMaskedIndexer6D idxer[3];
+  SpatiallyMaskedIndexer6D
+      prores_idxer[10]; // Has to be large enough to allow for maximum integer
+                        // conversion of TopologicalElements
 
   CoordinateDirection dir;
   bool allocated = true;

src/bvals/comms/boundary_communication.cpp

@@ -228,8 +228,9 @@ TaskStatus SetBounds(std::shared_ptr<MeshData<Real>> &md) {
             Kokkos::parallel_for(Kokkos::TeamThreadRange<>(team_member, idxer.size()),
                                  [&](const int idx) {
                                    const auto [t, u, v, k, j, i] = idxer(idx);
-                                   bnd_info(b).var(iel, t, u, v, k, j, i) =
-                                       bnd_info(b).buf(idx + idx_offset);
+                                   if (idxer.IsActive(k, j, i))
+                                     bnd_info(b).var(iel, t, u, v, k, j, i) =
+                                         bnd_info(b).buf(idx + idx_offset);
                                  });
           } else if (bnd_info(b).allocated) {
             const Real default_val = bnd_info(b).var.sparse_default_val;

src/interface/make_pack_descriptor.hpp

@@ -0,0 +1,121 @@
+//========================================================================================
+// (C) (or copyright) 2020-2023. 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
+// for the U.S. Department of Energy/National Nuclear Security Administration. All rights
+// in the program are reserved by Triad National Security, LLC, and the U.S. Department
+// of Energy/National Nuclear Security Administration. The Government is granted for
+// itself and others acting on its behalf a nonexclusive, paid-up, irrevocable worldwide
+// license in this material to reproduce, prepare derivative works, distribute copies to
+// the public, perform publicly and display publicly, and to permit others to do so.
+//========================================================================================
+#ifndef INTERFACE_MAKE_PACK_DESCRIPTOR_HPP_
+#define INTERFACE_MAKE_PACK_DESCRIPTOR_HPP_
+
+#include <algorithm>
+#include <functional>
+#include <limits>
+#include <map>
+#include <memory>
+#include <regex>
+#include <set>
+#include <string>
+#include <tuple>
+#include <type_traits>
+#include <utility>
+#include <vector>
+
+#include "interface/mesh_data.hpp"
+#include "interface/meshblock_data.hpp"
+#include "interface/metadata.hpp"
+#include "interface/sparse_pack.hpp"
+#include "interface/state_descriptor.hpp"
+#include "mesh/mesh.hpp"
+
+namespace parthenon {
+
+inline auto MakePackDescriptor(StateDescriptor *psd, const std::vector<std::string> &vars,
+                               const std::vector<bool> &use_regex,
+                               const std::vector<MetadataFlag> &flags = {},
+                               const std::set<PDOpt> &options = {}) {
+  PARTHENON_REQUIRE(vars.size() == use_regex.size(),
+                    "Vargroup names and use_regex need to be the same size.");
+  auto selector = [&](int vidx, const VarID &id, const Metadata &md) {
+    if (flags.size() > 0) {
+      for (const auto &flag : flags) {
+        if (!md.IsSet(flag)) return false;
+      }
+    }
+
+    if (use_regex[vidx]) {
+      if (std::regex_match(std::string(id.label()), std::regex(vars[vidx]))) return true;
+    } else {
+      if (vars[vidx] == id.label()) return true;
+      if (vars[vidx] == id.base_name && id.sparse_id != InvalidSparseID) return true;
+    }
+    return false;
+  };
+
+  impl::PackDescriptor base_desc(psd, vars, selector, options);
+  return typename SparsePack<>::Descriptor(base_desc);
+}
+
+template <class... Ts>
+inline auto MakePackDescriptor(StateDescriptor *psd,
+                               const std::vector<MetadataFlag> &flags = {},
+                               const std::set<PDOpt> &options = {}) {
+  static_assert(sizeof...(Ts) > 0, "Must have at least one variable type for type pack");
+
+  std::vector<std::string> vars{Ts::name()...};
+  std::vector<bool> use_regex{Ts::regex()...};
+
+  return typename SparsePack<Ts...>::Descriptor(static_cast<impl::PackDescriptor>(
+      MakePackDescriptor(psd, vars, use_regex, flags, options)));
+}
+
+inline auto MakePackDescriptor(StateDescriptor *psd, const std::vector<std::string> &vars,
+                               const std::vector<MetadataFlag> &flags = {},
+                               const std::set<PDOpt> &options = {}) {
+  return MakePackDescriptor(psd, vars, std::vector<bool>(vars.size(), false), flags,
+                            options);
+}
+
+template <class... Ts>
+inline auto MakePackDescriptor(MeshBlockData<Real> *pmbd,
+                               const std::vector<MetadataFlag> &flags = {},
+                               const std::set<PDOpt> &options = {}) {
+  return MakePackDescriptor<Ts...>(
+      pmbd->GetBlockPointer()->pmy_mesh->resolved_packages.get(), flags, options);
+}
+
+template <class... Ts>
+inline auto MakePackDescriptor(MeshData<Real> *pmd,
+                               const std::vector<MetadataFlag> &flags = {},
+                               const std::set<PDOpt> &options = {}) {
+  return MakePackDescriptor<Ts...>(pmd->GetMeshPointer()->resolved_packages.get(), flags,
+                                   options);
+}
+
+template <class... Ts>
+inline auto MakePackDescriptor(SparsePack<Ts...> pack, StateDescriptor *psd,
+                               const std::vector<MetadataFlag> &flags = {},
+                               const std::set<PDOpt> &options = {}) {
+  return parthenon::MakePackDescriptor<Ts...>(psd, flags, options);
+}
+
+inline auto MakePackDescriptor(
+    StateDescriptor *psd, const std::vector<std::pair<std::string, bool>> &var_regexes,
+    const std::vector<MetadataFlag> &flags = {}, const std::set<PDOpt> &options = {}) {
+  std::vector<std::string> vars;
+  std::vector<bool> use_regex;
+  for (const auto &[v, r] : var_regexes) {
+    vars.push_back(v);
+    use_regex.push_back(r);
+  }
+  return MakePackDescriptor(psd, vars, use_regex, flags, options);
+}
+
+} // namespace parthenon
+
+#endif // INTERFACE_MAKE_PACK_DESCRIPTOR_HPP_