[Trivial] Reorganize LogicalLocation code

src/CMakeLists.txt

@@ -158,15 +158,17 @@ add_library(parthenon
 
   mesh/amr_loadbalance.cpp
   mesh/domain.hpp
+  mesh/forest/block_ownership.cpp
+  mesh/forest/block_ownership.hpp
   mesh/forest/cell_center_offsets.hpp
   mesh/forest/forest.cpp
   mesh/forest/forest.hpp
   mesh/forest/relative_orientation.hpp
   mesh/forest/relative_orientation.cpp
   mesh/forest/tree.hpp
   mesh/forest/tree.cpp
-  mesh/logical_location.cpp
-  mesh/logical_location.hpp
+  mesh/forest/logical_location.cpp
+  mesh/forest/logical_location.hpp
   mesh/mesh_refinement.cpp
   mesh/mesh_refinement.hpp
   mesh/mesh-gmg.cpp

src/bvals/neighbor_block.cpp

@@ -30,7 +30,7 @@
 #include <unordered_set>
 
 #include "globals.hpp"
-#include "mesh/logical_location.hpp"
+#include "mesh/forest/logical_location.hpp"
 #include "mesh/mesh.hpp"
 #include "utils/buffer_utils.hpp"
 #include "utils/error_checking.hpp"

src/bvals/neighbor_block.hpp

@@ -29,6 +29,8 @@
 #include "parthenon_mpi.hpp"
 
 #include "defs.hpp"
+#include "mesh/forest/block_ownership.hpp"
+#include "mesh/forest/logical_location.hpp"
 #include "parthenon_arrays.hpp"
 #include "utils/error_checking.hpp"
 

src/defs.hpp

@@ -25,7 +25,7 @@
 
 #include "basic_types.hpp"
 #include "config.hpp"
-#include "mesh/logical_location.hpp"
+#include "mesh/forest/logical_location.hpp"
 #include "parthenon_arrays.hpp"
 
 namespace parthenon {

src/mesh/forest/block_ownership.cpp

@@ -0,0 +1,183 @@
+//========================================================================================
+// Athena++ astrophysical MHD code
+// Copyright(C) 2014 James M. Stone <jmstone@princeton.edu> and other code contributors
+// Licensed under the 3-clause BSD License, see LICENSE file for details
+//========================================================================================
+// Parthenon performance portable AMR framework
+// Copyright(C) 2020-2023 The Parthenon collaboration
+// Licensed under the 3-clause BSD License, see LICENSE file for details
+//========================================================================================
+// (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.
+//========================================================================================
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <set>
+#include <string>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "mesh/forest/block_ownership.hpp"
+#include "mesh/forest/logical_location.hpp"
+#include "utils/error_checking.hpp"
+#include "utils/morton_number.hpp"
+
+namespace parthenon {
+
+// TODO(LFR): Remove this
+block_ownership_t
+DetermineOwnership(const LogicalLocation &main_block,
+                   const std::unordered_set<LogicalLocation> &allowed_neighbors,
+                   const RootGridInfo &rg_info,
+                   const std::unordered_set<LogicalLocation> &newly_refined) {
+  block_ownership_t main_owns;
+
+  auto ownership_level = [&](const LogicalLocation &a) {
+    // Newly-refined blocks are treated as higher-level than blocks at their
+    // parent level, but lower-level than previously-refined blocks at their
+    // current level.
+    if (newly_refined.count(a)) return 2 * a.level() - 1;
+    return 2 * a.level();
+  };
+
+  auto ownership_less_than = [ownership_level](const LogicalLocation &a,
+                                               const LogicalLocation &b) {
+    // Ownership is first determined by block with the highest level, then by maximum
+    // Morton number this is reversed in precedence from the normal comparators where
+    // Morton number takes precedence
+    if (ownership_level(a) == ownership_level(b)) return a.morton() < b.morton();
+    return ownership_level(a) < ownership_level(b);
+  };
+
+  for (int ox1 : {-1, 0, 1}) {
+    for (int ox2 : {-1, 0, 1}) {
+      for (int ox3 : {-1, 0, 1}) {
+        main_owns(ox1, ox2, ox3) = true;
+        for (auto &n : allowed_neighbors) {
+          if (ownership_less_than(main_block, n) &&
+              main_block.IsNeighborOfTE(n, ox1, ox2, ox3, rg_info)) {
+            main_owns(ox1, ox2, ox3) = false;
+            break;
+          }
+        }
+      }
+    }
+  }
+  return main_owns;
+}
+
+block_ownership_t
+DetermineOwnershipForest(const LogicalLocation &main_block,
+                         const std::vector<NeighborLocation> &allowed_neighbors,
+                         const std::unordered_set<LogicalLocation> &newly_refined) {
+  block_ownership_t main_owns;
+
+  auto ownership_level = [&](const LogicalLocation &a) {
+    // Newly-refined blocks are treated as higher-level than blocks at their
+    // parent level, but lower-level than previously-refined blocks at their
+    // current level.
+    if (newly_refined.count(a)) return 2 * a.level() - 1;
+    return 2 * a.level();
+  };
+
+  auto ownership_less_than = [ownership_level](const LogicalLocation &a,
+                                               const LogicalLocation &b) {
+    // Ownership is first determined by block with the highest level, then by maximum
+    // (tree, Morton) number this is reversed in precedence from the normal comparators
+    // where (tree, Morton) number takes precedence
+    if (ownership_level(a) != ownership_level(b))
+      return ownership_level(a) < ownership_level(b);
+    if (a.tree() != b.tree()) return a.tree() < b.tree();
+    return a.morton() < b.morton();
+  };
+
+  for (int ox1 : {-1, 0, 1}) {
+    for (int ox2 : {-1, 0, 1}) {
+      for (int ox3 : {-1, 0, 1}) {
+        main_owns(ox1, ox2, ox3) = true;
+        for (const auto &n : allowed_neighbors) {
+          if (ownership_less_than(main_block, n.global_loc) &&
+              main_block.IsNeighborOfTEForest(n.origin_loc, {ox1, ox2, ox3})) {
+            main_owns(ox1, ox2, ox3) = false;
+            break;
+          }
+        }
+      }
+    }
+  }
+  return main_owns;
+}
+
+// Given a topological element, ownership array of the sending block, and offset indices
+// defining the location of an index region within the block (i.e. the ghost zones passed
+// across the x-face or the ghost zones passed across the z-edge), return the index range
+// masking array required for masking out unowned regions of the index space. ox? defines
+// buffer location on the owner block
+block_ownership_t
+GetIndexRangeMaskFromOwnership(TopologicalElement el,
+                               const block_ownership_t &sender_ownership, int ox1,
+                               int ox2, int ox3) {
+  using vp_t = std::vector<std::pair<int, int>>;
+
+  // Transform general block ownership to element ownership over entire block. For
+  // instance, x-faces only care about block ownership in the x-direction First index of
+  // the pair is the element index and the second index is the block index that is copied
+  // to that element index
+  block_ownership_t element_ownership = sender_ownership;
+  auto x1_idxs = TopologicalOffsetI(el) ? vp_t{{-1, -1}, {0, 0}, {1, 1}}
+                                        : vp_t{{-1, 0}, {0, 0}, {1, 0}};
+  auto x2_idxs = TopologicalOffsetJ(el) ? vp_t{{-1, -1}, {0, 0}, {1, 1}}
+                                        : vp_t{{-1, 0}, {0, 0}, {1, 0}};
+  auto x3_idxs = TopologicalOffsetK(el) ? vp_t{{-1, -1}, {0, 0}, {1, 1}}
+                                        : vp_t{{-1, 0}, {0, 0}, {1, 0}};
+  for (auto [iel, ibl] : x1_idxs) {
+    for (auto [jel, jbl] : x2_idxs) {
+      for (auto [kel, kbl] : x3_idxs) {
+        element_ownership(iel, jel, kel) = sender_ownership(ibl, jbl, kbl);
+      }
+    }
+  }
+
+  // Now, the ownership status is correct for the entire interior index range of the
+  // block, but the offsets ox? define a subset of these indices (e.g. one edge of the
+  // interior). Therefore, we need to set the index ownership to true for edges of the
+  // index range that are contained in the interior of the sending block
+  if (ox1 != 0) {
+    for (auto j : {-1, 0, 1}) {
+      for (auto k : {-1, 0, 1}) {
+        element_ownership(-ox1, j, k) = element_ownership(0, j, k);
+      }
+    }
+  }
+  if (ox2 != 0) {
+    for (auto i : {-1, 0, 1}) {
+      for (auto k : {-1, 0, 1}) {
+        element_ownership(i, -ox2, k) = element_ownership(i, 0, k);
+      }
+    }
+  }
+  if (ox3 != 0) {
+    for (auto i : {-1, 0, 1}) {
+      for (auto j : {-1, 0, 1}) {
+        element_ownership(i, j, -ox3) = element_ownership(i, j, 0);
+      }
+    }
+  }
+
+  return element_ownership;
+}
+
+} // namespace parthenon

src/mesh/forest/block_ownership.hpp

@@ -0,0 +1,108 @@
+//========================================================================================
+// Athena++ astrophysical MHD code
+// Copyright(C) 2014 James M. Stone <jmstone@princeton.edu> and other code contributors
+// Licensed under the 3-clause BSD License, see LICENSE file for details
+//========================================================================================
+// Parthenon performance portable AMR framework
+// Copyright(C) 2020-2024 The Parthenon collaboration
+// Licensed under the 3-clause BSD License, see LICENSE file for details
+//========================================================================================
+// (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
+// 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 MESH_FOREST_BLOCK_OWNERSHIP_HPP_
+#define MESH_FOREST_BLOCK_OWNERSHIP_HPP_
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <set>
+#include <string>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "basic_types.hpp"
+#include "mesh/forest/logical_location.hpp"
+#include "utils/error_checking.hpp"
+#include "utils/morton_number.hpp"
+
+namespace parthenon {
+
+struct block_ownership_t {
+ public:
+  KOKKOS_FORCEINLINE_FUNCTION
+  const bool &operator()(int ox1, int ox2, int ox3) const {
+    return ownership[ox1 + 1][ox2 + 1][ox3 + 1];
+  }
+  KOKKOS_FORCEINLINE_FUNCTION
+  bool &operator()(int ox1, int ox2, int ox3) {
+    return ownership[ox1 + 1][ox2 + 1][ox3 + 1];
+  }
+
+  KOKKOS_FORCEINLINE_FUNCTION
+  block_ownership_t() : block_ownership_t(false) {}
+
+  KOKKOS_FORCEINLINE_FUNCTION
+  explicit block_ownership_t(bool value) : initialized(false) {
+    for (int i = 0; i < 3; ++i) {
+      for (int j = 0; j < 3; ++j) {
+        for (int k = 0; k < 3; ++k) {
+          ownership[i][j][k] = value;
+        }
+      }
+    }
+  }
+
+  bool initialized;
+
+  bool operator==(const block_ownership_t &rhs) const {
+    bool same = initialized == rhs.initialized;
+    for (int i = 0; i < 3; ++i) {
+      for (int j = 0; j < 3; ++j) {
+        for (int k = 0; k < 3; ++k) {
+          same = same && (ownership[i][j][k] == rhs.ownership[i][j][k]);
+        }
+      }
+    }
+    return same;
+  }
+
+ private:
+  bool ownership[3][3][3];
+};
+
+block_ownership_t
+DetermineOwnership(const LogicalLocation &main_block,
+                   const std::unordered_set<LogicalLocation> &allowed_neighbors,
+                   const RootGridInfo &rg_info = RootGridInfo(),
+                   const std::unordered_set<LogicalLocation> &newly_refined = {});
+
+block_ownership_t
+DetermineOwnershipForest(const LogicalLocation &main_block,
+                         const std::vector<NeighborLocation> &allowed_neighbors,
+                         const std::unordered_set<LogicalLocation> &newly_refined = {});
+
+// Given a topological element, ownership array of the sending block, and offset indices
+// defining the location of an index region within the block (i.e. the ghost zones passed
+// across the x-face or the ghost zones passed across the z-edge), return the index range
+// masking array required for masking out unowned regions of the index space. ox? defines
+// buffer location on the owner block
+block_ownership_t
+GetIndexRangeMaskFromOwnership(TopologicalElement el,
+                               const block_ownership_t &sender_ownership, int ox1,
+                               int ox2, int ox3);
+
+} // namespace parthenon
+
+#endif // MESH_FOREST_BLOCK_OWNERSHIP_HPP_

src/mesh/forest/cell_center_offsets.hpp

@@ -25,7 +25,7 @@
 
 #include "basic_types.hpp"
 #include "defs.hpp"
-#include "mesh/logical_location.hpp"
+#include "mesh/forest/logical_location.hpp"
 #include "utils/bit_hacks.hpp"
 #include "utils/indexer.hpp"
 

src/mesh/forest/forest.cpp

@@ -27,9 +27,9 @@
 #include "basic_types.hpp"
 #include "defs.hpp"
 #include "mesh/forest/forest.hpp"
+#include "mesh/forest/logical_location.hpp"
 #include "mesh/forest/relative_orientation.hpp"
 #include "mesh/forest/tree.hpp"
-#include "mesh/logical_location.hpp"
 #include "utils/bit_hacks.hpp"
 #include "utils/indexer.hpp"
 

src/mesh/forest/forest.hpp

@@ -25,8 +25,8 @@
 
 #include "basic_types.hpp"
 #include "defs.hpp"
+#include "mesh/forest/logical_location.hpp"
 #include "mesh/forest/tree.hpp"
-#include "mesh/logical_location.hpp"
 #include "utils/bit_hacks.hpp"
 #include "utils/indexer.hpp"
 
@@ -145,6 +145,14 @@ class Forest {
                          std::array<BoundaryFlag, BOUNDARY_NFACES> mesh_bcs);
 };
 
+struct NeighborLocation {
+  NeighborLocation(const LogicalLocation &g, const LogicalLocation &o)
+      : global_loc(g), origin_loc(o) {}
+  LogicalLocation global_loc; // Global location of neighboring block
+  LogicalLocation
+      origin_loc; // Logical location of neighboring block in index space of origin block
+};
+
 } // namespace forest
 } // namespace parthenon