Merge pull request #973 from parthenon-hpc-lab/lroberts36/mg-riot-upd…

…ates

MG performance upgrades

CHANGELOG.md

@@ -15,6 +15,7 @@
 - [[PR 996]](https://github.com/parthenon-hpc-lab/parthenon/pull/996) Remove dynamic allocations from swarm particle creation
 
 ### Changed (changing behavior/API/variables/...)
+- [[PR 973]](https://github.com/parthenon-hpc-lab/parthenon/pull/973) Multigrid performance upgrades
 
 ### Fixed (not changing behavior/API/variables/...)
 - [[PR1023]](https://github.com/parthenon-hpc-lab/parthenon/pull/1023) Fix broken param of a scalar bool

example/poisson_gmg/main.cpp

@@ -51,6 +51,8 @@ int main(int argc, char *argv[]) {
     if (driver_status != parthenon::DriverStatus::complete ||
         driver.final_rms_residual > 1.e-10 || driver.final_rms_error > 1.e-12)
       success = false;
+    if (driver.final_rms_residual != driver.final_rms_residual) success = false;
+    if (driver.final_rms_error != driver.final_rms_error) success = false;
   }
   // call MPI_Finalize and Kokkos::finalize if necessary
   pman.ParthenonFinalize();

example/poisson_gmg/poisson_driver.cpp

@@ -79,7 +79,6 @@ TaskCollection PoissonDriver::MakeTaskCollection(BlockList_t &blocks) {
 
   const int num_partitions = pmesh->DefaultNumPartitions();
   TaskRegion &region = tc.AddRegion(num_partitions);
-  int reg_dep_id = 0;
   for (int i = 0; i < num_partitions; ++i) {
     TaskList &tl = region[i];
     auto &md = pmesh->mesh_data.GetOrAdd("base", i);
@@ -100,9 +99,11 @@ TaskCollection PoissonDriver::MakeTaskCollection(BlockList_t &blocks) {
 
     auto solve = zero_u;
     if (solver == "BiCGSTAB") {
-      solve = bicgstab_solver->AddTasks(tl, zero_u, pmesh, i);
+      auto setup = bicgstab_solver->AddSetupTasks(tl, zero_u, i, pmesh);
+      solve = bicgstab_solver->AddTasks(tl, setup, pmesh, i);
     } else if (solver == "MG") {
-      solve = mg_solver->AddTasks(tl, zero_u, pmesh, i);
+      auto setup = mg_solver->AddSetupTasks(tl, zero_u, i, pmesh);
+      solve = mg_solver->AddTasks(tl, setup, pmesh, i);
     } else {
       PARTHENON_FAIL("Unknown solver type.");
     }

example/poisson_gmg/poisson_driver.hpp

@@ -1,5 +1,5 @@
 //========================================================================================
-// (C) (or copyright) 2021-2023. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2021-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

example/poisson_gmg/poisson_package.cpp

@@ -1,5 +1,5 @@
 //========================================================================================
-// (C) (or copyright) 2021-2023. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2021-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
@@ -147,6 +147,7 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
   bicgstab_params.max_iters = max_poisson_iterations;
   bicgstab_params.residual_tolerance = res_tol;
   bicgstab_params.precondition = precondition;
+  bicgstab_params.print_per_step = true;
   parthenon::solvers::BiCGSTABSolver<u, rhs, PoissonEquation> bicg_solver(
       pkg.get(), bicgstab_params, eq);
   pkg->AddParam<>("MGBiCGSTABsolver", bicg_solver,

src/bvals/boundary_conditions_generic.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
@@ -18,6 +18,9 @@
 #include <memory>
 #include <set>
 #include <string>
+#include <tuple>
+#include <unordered_map>
+#include <utility>
 #include <vector>
 
 #include "basic_types.hpp"
@@ -34,6 +37,35 @@ namespace BoundaryFunction {
 enum class BCSide { Inner, Outer };
 enum class BCType { Outflow, Reflect, ConstantDeriv, Fixed, FixedFace };
 
+namespace impl {
+using desc_key_t = std::tuple<bool, TopologicalType>;
+template <class... var_ts>
+using map_bc_pack_descriptor_t =
+    std::unordered_map<desc_key_t, typename SparsePack<var_ts...>::Descriptor,
+                       tuple_hash<desc_key_t>>;
+
+template <class... var_ts>
+map_bc_pack_descriptor_t<var_ts...>
+GetPackDescriptorMap(std::shared_ptr<MeshBlockData<Real>> &rc) {
+  std::vector<std::pair<TopologicalType, MetadataFlag>> elements{
+      {TopologicalType::Cell, Metadata::Cell},
+      {TopologicalType::Face, Metadata::Face},
+      {TopologicalType::Edge, Metadata::Edge},
+      {TopologicalType::Node, Metadata::Node}};
+  map_bc_pack_descriptor_t<var_ts...> my_map;
+  for (auto [tt, md] : elements) {
+    std::vector<MetadataFlag> flags{Metadata::FillGhost};
+    flags.push_back(md);
+    std::set<PDOpt> opts{PDOpt::Coarse};
+    my_map.emplace(std::make_pair(desc_key_t{true, tt},
+                                  MakePackDescriptor<var_ts...>(rc.get(), flags, opts)));
+    my_map.emplace(std::make_pair(desc_key_t{false, tt},
+                                  MakePackDescriptor<var_ts...>(rc.get(), flags)));
+  }
+  return my_map;
+}
+} // namespace impl
+
 template <CoordinateDirection DIR, BCSide SIDE, BCType TYPE, class... var_ts>
 void GenericBC(std::shared_ptr<MeshBlockData<Real>> &rc, bool coarse,
                TopologicalElement el, Real val) {
@@ -46,17 +78,9 @@ void GenericBC(std::shared_ptr<MeshBlockData<Real>> &rc, bool coarse,
   constexpr bool X3 = (DIR == X3DIR);
   constexpr bool INNER = (SIDE == BCSide::Inner);
 
-  std::vector<MetadataFlag> flags{Metadata::FillGhost};
-  if (GetTopologicalType(el) == TopologicalType::Cell) flags.push_back(Metadata::Cell);
-  if (GetTopologicalType(el) == TopologicalType::Face) flags.push_back(Metadata::Face);
-  if (GetTopologicalType(el) == TopologicalType::Edge) flags.push_back(Metadata::Edge);
-  if (GetTopologicalType(el) == TopologicalType::Node) flags.push_back(Metadata::Node);
-
-  std::set<PDOpt> opts;
-  if (coarse) opts = {PDOpt::Coarse};
-  auto desc = MakePackDescriptor<var_ts...>(
-      rc->GetBlockPointer()->pmy_mesh->resolved_packages.get(), flags, opts);
-  auto q = desc.GetPack(rc.get());
+  static auto descriptors = impl::GetPackDescriptorMap<var_ts...>(rc);
+  auto q =
+      descriptors[impl::desc_key_t{coarse, GetTopologicalType(el)}].GetPack(rc.get());
   const int b = 0;
   const int lstart = q.GetLowerBoundHost(b);
   const int lend = q.GetUpperBoundHost(b);

src/interface/mesh_data.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
@@ -25,9 +25,19 @@ void MeshData<T>::Initialize(const MeshData<T> *src,
   pmy_mesh_ = src->GetParentPointer();
   const int nblocks = src->NumBlocks();
   block_data_.resize(nblocks);
-  for (int i = 0; i < nblocks; i++) {
-    block_data_[i] = pmy_mesh_->block_list[i]->meshblock_data.Add(
-        stage_name_, src->GetBlockData(i), names, shallow);
+
+  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);
+    }
+  } else {
+    for (int i = 0; i < nblocks; i++) {
+      block_data_[i] = pmy_mesh_->block_list[i]->meshblock_data.Add(
+          stage_name_, src->GetBlockData(i), names, shallow);
+    }
   }
 }
 

src/interface/sparse_pack.hpp

@@ -181,17 +181,16 @@ class SparsePack : public SparsePackBase {
     // accessed on device via instance of types in the type list Ts...
     // The pack will be created and accessible on the device
     template <class T>
-    SparsePack GetPack(T *pmd, std::vector<bool> include_block = {},
+    SparsePack GetPack(T *pmd, std::vector<bool> &include_block,
                        bool only_fine_two_level_composite_blocks = true) const {
       // If this is a composite grid MeshData object and if
       // only_fine_two_level_composite_blocks is true, only
       // include blocks on the finer level
       if constexpr (std::is_same<T, MeshData<Real>>::value) {
         if (pmd->grid.type == GridType::two_level_composite &&
             only_fine_two_level_composite_blocks) {
-          if (include_block.size() != pmd->NumBlocks()) {
-            include_block = std::vector<bool>(pmd->NumBlocks(), true);
-          }
+          PARTHENON_REQUIRE(include_block.size() == pmd->NumBlocks(),
+                            "Passed wrong size block include list.");
           int fine_level = pmd->grid.logical_level;
           for (int b = 0; b < pmd->NumBlocks(); ++b)
             include_block[b] =
@@ -202,6 +201,27 @@ class SparsePack : public SparsePackBase {
       return SparsePack(SparsePackBase::GetPack(pmd, *this, include_block));
     }
 
+    template <class T>
+    SparsePack GetPack(T *pmd, bool only_fine_two_level_composite_blocks = true) const {
+      // If this is a composite grid MeshData object, only include blocks on
+      // the finer level
+      if constexpr (std::is_same<T, MeshData<Real>>::value) {
+        if (pmd->grid.type == GridType::two_level_composite &&
+            only_fine_two_level_composite_blocks) {
+          auto include_block = std::vector<bool>(pmd->NumBlocks(), true);
+          int fine_level = pmd->grid.logical_level;
+          for (int b = 0; b < pmd->NumBlocks(); ++b)
+            include_block[b] =
+                include_block[b] &&
+                (fine_level == pmd->GetBlockData(b)->GetBlockPointer()->loc.level());
+          return SparsePack(SparsePackBase::GetPack(pmd, *this, include_block));
+        } else {
+          return SparsePack(SparsePackBase::GetPack(pmd, *this, std::vector<bool>{}));
+        }
+      }
+      return SparsePack(SparsePackBase::GetPack(pmd, *this, std::vector<bool>{}));
+    }
+
     SparsePackIdxMap GetMap() const {
       PARTHENON_REQUIRE(sizeof...(Ts) == 0,
                         "Should not be getting an IdxMap for a type based pack");