Fix Multigrid Solver in Forest

.github/workflows/ci-extended.yml

@@ -60,15 +60,15 @@ jobs:
           cd build
           # Pick GPU with most available memory
           export CUDA_VISIBLE_DEVICES=$(nvidia-smi --query-gpu=memory.free,index --format=csv,nounits,noheader | sort -nr | head -1 | awk '{ print $NF }')
-          ctest -L performance -LE perf-reg -E gmg
+          ctest -L performance -LE perf-reg
 
       # run regression tests
       - name: Regression tests
         run: |
           cd build
           # Pick GPU with most available memory
           export CUDA_VISIBLE_DEVICES=$(nvidia-smi --query-gpu=memory.free,index --format=csv,nounits,noheader | sort -nr | head -1 | awk '{ print $NF }')
-          ctest -L regression -L ${{ matrix.parallel }} -LE perf-reg -E gmg --timeout 3600
+          ctest -L regression -L ${{ matrix.parallel }} -LE perf-reg --timeout 3600
 
       # Test Ascent integration (only most complex setup with MPI and on device)
       - name: Ascent tests

CHANGELOG.md

@@ -16,6 +16,7 @@
 - [[PR 1031]](https://github.com/parthenon-hpc-lab/parthenon/pull/1031) Fix bug in non-cell centered AMR
 
 ### Infrastructure (changes irrelevant to downstream codes)
+- [[PR 1035]](https://github.com/parthenon-hpc-lab/parthenon/pull/1035) Fix multigrid infrastructure to work with forest
 - [[PR 1048]](https://github.com/parthenon-hpc-lab/parthenon/pull/1048) Tiny fixes to custom coords logic
 - [[PR 1028]](https://github.com/parthenon-hpc-lab/parthenon/pull/1028) Internal reorganization of LogicalLocation files
 - [[PR 1009]](https://github.com/parthenon-hpc-lab/parthenon/pull/1009) Move from a single octree to a forest of octrees

doc/sphinx/src/mesh/mesh.rst

@@ -71,7 +71,10 @@ level that are neighbors to finer blocks, which implies that below the root grid
 level the blocks may not cover the entire mesh. For levels above the root grid, 
 blocks may change shape so that they only cover the domain of the root grid. Note 
 that leaf blocks may be contained in multiple blocklists, and the lists all point
-to the same block (not a separate copy). To be explicit, when ``parthenon/mesh/multigrid`` is set to ``true`` blocks corresponding to *all* internal nodes of the refinement tree are created, in addition to the leaf node blocks that are normally created.
+to the same block (not a separate copy). To be explicit, when 
+``parthenon/mesh/multigrid`` is set to ``true`` blocks corresponding to *all* 
+internal nodes of the refinement tree are created, in addition to the leaf node blocks 
+that are normally created.
 
 *GMG Implementation Note:*
 The reason for including two levels in the GMG block lists is for dealing with 

src/basic_types.hpp

@@ -67,6 +67,17 @@ enum class BoundaryType : int {
   gmg_prolongate_recv
 };
 
+enum class GridType { none, leaf, two_level_composite, single_level_with_internal };
+struct GridIdentifier {
+  GridType type = GridType::none;
+  int logical_level = 0;
+
+  static GridIdentifier leaf() { return GridIdentifier{GridType::leaf, 0}; }
+  static GridIdentifier two_level_composite(int level) {
+    return GridIdentifier{GridType::two_level_composite, level};
+  }
+};
+
 constexpr bool IsSender(BoundaryType btype) {
   if (btype == BoundaryType::flxcor_recv) return false;
   if (btype == BoundaryType::gmg_restrict_recv) return false;

src/driver/driver.cpp

@@ -1,5 +1,5 @@
 //========================================================================================
-// (C) (or copyright) 2020-2021. 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
@@ -191,8 +191,8 @@ void EvolutionDriver::InitializeBlockTimeStepsAndBoundaries() {
     auto &mbase = pmesh->mesh_data.GetOrAdd("base", i);
     Update::EstimateTimestep(mbase.get());
     BuildBoundaryBuffers(mbase);
-    for (int gmg_level = 0; gmg_level < pmesh->gmg_mesh_data.size(); ++gmg_level) {
-      auto &mdg = pmesh->gmg_mesh_data[gmg_level].GetOrAdd(gmg_level, "base", i);
+    for (auto &[gmg_level, mdc] : pmesh->gmg_mesh_data) {
+      auto &mdg = mdc.GetOrAdd(gmg_level, "base", i);
       BuildBoundaryBuffers(mdg);
       BuildGMGBoundaryBuffers(mdg);
     }

src/interface/data_collection.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
@@ -58,9 +58,9 @@ std::shared_ptr<MeshData<Real>> &
 GetOrAdd_impl(Mesh *pmy_mesh_,
               std::map<std::string, std::shared_ptr<MeshData<Real>>> &containers_,
               BlockList_t &block_list, const std::string &mbd_label,
-              const int &partition_id, const int gmg_level) {
+              const int &partition_id, const std::optional<int> gmg_level) {
   std::string label = mbd_label + "_part-" + std::to_string(partition_id);
-  if (gmg_level >= 0) label = label + "_gmg-" + std::to_string(gmg_level);
+  if (gmg_level) label = label + "_gmg-" + std::to_string(*gmg_level);
   auto it = containers_.find(label);
   if (it == containers_.end()) {
     // TODO(someone) add caching of partitions to Mesh at some point
@@ -70,15 +70,13 @@ GetOrAdd_impl(Mesh *pmy_mesh_,
     if (partitions.size() == 0) partitions = std::vector<BlockList_t>(1);
     for (auto i = 0; i < partitions.size(); i++) {
       std::string md_label = mbd_label + "_part-" + std::to_string(i);
-      if (gmg_level >= 0) md_label = md_label + "_gmg-" + std::to_string(gmg_level);
+      if (gmg_level) md_label = md_label + "_gmg-" + std::to_string(*gmg_level);
       containers_[md_label] = std::make_shared<MeshData<Real>>(mbd_label);
       containers_[md_label]->Set(partitions[i], pmy_mesh_);
-      if (gmg_level >= 0) {
-        int min_gmg_logical_level = pmy_mesh_->GetGMGMinLogicalLevel();
-        containers_[md_label]->grid = GridIdentifier{GridType::two_level_composite,
-                                                     gmg_level + min_gmg_logical_level};
+      if (gmg_level) {
+        containers_[md_label]->grid = GridIdentifier::two_level_composite(*gmg_level);
       } else {
-        containers_[md_label]->grid = GridIdentifier{GridType::leaf, 0};
+        containers_[md_label]->grid = GridIdentifier::leaf();
       }
     }
   }
@@ -90,7 +88,7 @@ std::shared_ptr<MeshData<Real>> &
 DataCollection<MeshData<Real>>::GetOrAdd(const std::string &mbd_label,
                                          const int &partition_id) {
   return GetOrAdd_impl(pmy_mesh_, containers_, pmy_mesh_->block_list, mbd_label,
-                       partition_id, -1);
+                       partition_id, {});
 }
 
 template <>

src/interface/mesh_data.cpp

@@ -28,10 +28,10 @@ void MeshData<T>::Initialize(const MeshData<T> *src,
 
   grid = src->grid;
   if (grid.type == GridType::two_level_composite) {
-    int gmg_level = src->grid.logical_level - pmy_mesh_->GetGMGMinLogicalLevel();
     for (int i = 0; i < nblocks; i++) {
-      block_data_[i] = pmy_mesh_->gmg_block_lists[gmg_level][i]->meshblock_data.Add(
-          stage_name_, src->GetBlockData(i), names, shallow);
+      block_data_[i] =
+          pmy_mesh_->gmg_block_lists[src->grid.logical_level][i]->meshblock_data.Add(
+              stage_name_, src->GetBlockData(i), names, shallow);
     }
   } else {
     for (int i = 0; i < nblocks; i++) {

src/interface/mesh_data.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
@@ -181,12 +181,6 @@ const MeshBlockPack<P> &PackOnMesh(M &map, BlockDataList_t<Real> &block_data_,
 
 } // namespace pack_on_mesh_impl
 
-enum class GridType { none, leaf, two_level_composite, single_level_with_internal };
-struct GridIdentifier {
-  GridType type = GridType::none;
-  int logical_level = 0;
-};
-
 /// The MeshData class is a container for cached MeshBlockPacks, i.e., it
 /// contains both the pointers to the MeshBlockData of the MeshBlocks contained
 /// in the object as well as maps to the cached MeshBlockPacks of VariablePacks or

src/mesh/forest/block_ownership.cpp

@@ -4,10 +4,10 @@
 // Licensed under the 3-clause BSD License, see LICENSE file for details
 //========================================================================================
 // Parthenon performance portable AMR framework
-// Copyright(C) 2020-2023 The Parthenon collaboration
+// Copyright(C) 2020-2024 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.
+// (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
@@ -37,11 +37,9 @@
 
 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::vector<NeighborLocation> &allowed_neighbors,
                    const std::unordered_set<LogicalLocation> &newly_refined) {
   block_ownership_t main_owns;
 
@@ -53,46 +51,6 @@ DetermineOwnership(const LogicalLocation &main_block,
     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
@@ -110,7 +68,7 @@ DetermineOwnershipForest(const LogicalLocation &main_block,
         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_block.IsNeighborOfTE(n.origin_loc, {ox1, ox2, ox3})) {
             main_owns(ox1, ox2, ox3) = false;
             break;
           }

src/mesh/forest/block_ownership.hpp

@@ -84,15 +84,9 @@ struct block_ownership_t {
 
 block_ownership_t
 DetermineOwnership(const LogicalLocation &main_block,
-                   const std::unordered_set<LogicalLocation> &allowed_neighbors,
-                   const RootGridInfo &rg_info = RootGridInfo(),
+                   const std::vector<NeighborLocation> &allowed_neighbors,
                    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

src/mesh/forest/forest.cpp

@@ -41,13 +41,22 @@ std::vector<LogicalLocation> Forest::GetMeshBlockListAndResolveGids() {
   std::uint64_t gid{0};
   for (auto &[id, tree] : trees) {
     std::size_t start = mb_list.size();
-    auto tree_mbs = tree->GetMeshBlockList();
+    auto tree_mbs = tree->GetSortedMeshBlockList();
     mb_list.insert(mb_list.end(), std::make_move_iterator(tree_mbs.begin()),
                    std::make_move_iterator(tree_mbs.end()));
     std::size_t end = mb_list.size();
     for (int i = start; i < end; ++i)
       tree->InsertGid(mb_list[i], gid++);
   }
+
+  // Assign gids to the internal nodes
+  for (auto &[id, tree] : trees) {
+    std::size_t start = mb_list.size();
+    auto tree_int_locs = tree->GetSortedInternalNodeList();
+    for (auto &loc : tree_int_locs)
+      tree->InsertGid(loc, gid++);
+  }
+
   // The index of blocks in this list corresponds to their gid
   gids_resolved = true;
   return mb_list;

src/mesh/forest/forest.hpp

@@ -1,5 +1,5 @@
 //========================================================================================
-// (C) (or copyright) 2023. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2023-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
@@ -84,8 +84,10 @@ class Forest {
     return trees.at(loc.tree())->FindNeighbors(loc, ox1, ox2, ox3);
   }
 
-  std::vector<NeighborLocation> FindNeighbors(const LogicalLocation &loc) const {
-    return trees.at(loc.tree())->FindNeighbors(loc);
+  std::vector<NeighborLocation>
+  FindNeighbors(const LogicalLocation &loc,
+                GridIdentifier grid_id = GridIdentifier::leaf()) const {
+    return trees.at(loc.tree())->FindNeighbors(loc, grid_id);
   }
   std::size_t CountMeshBlock() const {
     std::size_t count{0};
@@ -134,6 +136,13 @@ class Forest {
     return trees.at(loc.tree())->GetGid(loc);
   }
 
+  // Get the gid of the leaf block with the same Morton number
+  // as loc (on the same tree)
+  std::int64_t GetLeafGid(const LogicalLocation &loc) const {
+    PARTHENON_REQUIRE(gids_resolved, "Asking for GID in invalid state.");
+    return trees.at(loc.tree())->GetLeafGid(loc);
+  }
+
   std::int64_t GetOldGid(const LogicalLocation &loc) const {
     PARTHENON_REQUIRE(gids_resolved, "Asking for GID in invalid state.");
     return trees.at(loc.tree())->GetOldGid(loc);