Merge pull request #1153 from parthenon-hpc-lab/lroberts36/fix-one-bl…

…ock-minimum-per-rank Allow for ranks to have zero blocks during initialization
parthenon-hpc-lab · Aug 23, 2024 · c7e4d23 · c7e4d23
2 parents 43380cd + b1e3d19
commit c7e4d23
Show file tree

Hide file tree

Showing 9 changed files with 123 additions and 94 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -30,6 +30,7 @@
 - [[PR 1019]](https://github.com/parthenon-hpc-lab/parthenon/pull/1019) Enable output for non-cell-centered variables
 
 ### Changed (changing behavior/API/variables/...)
+- [[PR 1153]](https://github.com/parthenon-hpc-lab/parthenon/pull/1153) Allow base grid with fewer blocks than ranks before initial AMR
 - [[PR 1105]](https://github.com/parthenon-hpc-lab/parthenon/pull/1105) Refactor parameter input for linear solvers
 - [[PR 1078]](https://github.com/parthenon-hpc-lab/parthenon/pull/1078) Add reduction fallback in 1D. Add IndexRange overload for 1D par loops
 - [[PR 1024]](https://github.com/parthenon-hpc-lab/parthenon/pull/1024) Add .outN. to history output filenames

diff --git a/src/bvals/comms/boundary_communication.cpp b/src/bvals/comms/boundary_communication.cpp
@@ -79,10 +79,12 @@ TaskStatus SendBoundBufs(std::shared_ptr<MeshData<Real>> &md) {
     }
   }
   // Restrict
-  auto pmb = md->GetBlockData(0)->GetBlockPointer();
-  StateDescriptor *resolved_packages = pmb->resolved_packages.get();
-  refinement::Restrict(resolved_packages, cache.prores_cache, pmb->cellbounds,
-                       pmb->c_cellbounds);
+  if (md->NumBlocks() > 0) {
+    auto pmb = md->GetBlockData(0)->GetBlockPointer();
+    StateDescriptor *resolved_packages = pmb->resolved_packages.get();
+    refinement::Restrict(resolved_packages, cache.prores_cache, pmb->cellbounds,
+                         pmb->c_cellbounds);
+  }
 
   // Load buffer data
   auto &bnd_info = cache.bnd_info;
@@ -335,7 +337,7 @@ TaskStatus SetBounds(std::shared_ptr<MeshData<Real>> &md) {
 #endif
   std::for_each(std::begin(cache.buf_vec), std::end(cache.buf_vec),
                 [](auto pbuf) { pbuf->Stale(); });
-  if (nbound > 0 && pmesh->multilevel) {
+  if (nbound > 0 && pmesh->multilevel && md->NumBlocks() > 0) {
     // Restrict
     auto pmb = md->GetBlockData(0)->GetBlockPointer();
     StateDescriptor *resolved_packages = pmb->resolved_packages.get();
@@ -377,7 +379,7 @@ TaskStatus ProlongateBounds(std::shared_ptr<MeshData<Real>> &md) {
     }
   }
 
-  if (nbound > 0 && pmesh->multilevel) {
+  if (nbound > 0 && pmesh->multilevel && md->NumBlocks() > 0) {
     auto pmb = md->GetBlockData(0)->GetBlockPointer();
     StateDescriptor *resolved_packages = pmb->resolved_packages.get();
 

diff --git a/src/coordinates/uniform_cartesian.hpp b/src/coordinates/uniform_cartesian.hpp
@@ -318,6 +318,7 @@ class UniformCartesian {
   const std::array<Real, 3> &GetXmin() const { return xmin_; }
   const std::array<int, 3> &GetStartIndex() const { return istart_; }
   const char *Name() const { return name_; }
+  static const char *StaticName() { return name_; }
 
  private:
   std::array<int, 3> istart_;

diff --git a/src/interface/mesh_data.hpp b/src/interface/mesh_data.hpp
@@ -519,6 +519,7 @@ class MeshData {
 
 template <typename T, typename... Args>
 std::vector<Uid_t> UidIntersection(MeshData<T> *md1, MeshData<T> *md2, Args &&...args) {
+  if (md1->NumBlocks() == 0 || md2->NumBlocks() == 0) return std::vector<Uid_t>();
   return UidIntersection(md1->GetBlockData(0).get(), md2->GetBlockData(0).get(),
                          std::forward<Args>(args)...);
 }

diff --git a/src/mesh/mesh-amr_loadbalance.cpp b/src/mesh/mesh-amr_loadbalance.cpp
@@ -388,18 +388,16 @@ void AssignBlocks(std::vector<double> const &costlist, std::vector<int> &ranklis
 
 void UpdateBlockList(std::vector<int> const &ranklist, std::vector<int> &nslist,
                      std::vector<int> &nblist) {
-  nslist.resize(Globals::nranks);
-  nblist.resize(Globals::nranks);
-
-  nslist[0] = 0;
-  int rank = 0;
-  for (int block_id = 1; block_id < ranklist.size(); block_id++) {
-    if (ranklist[block_id] != ranklist[block_id - 1]) {
-      nblist[rank] = block_id - nslist[rank];
-      nslist[++rank] = block_id;
-    }
-  }
-  nblist[rank] = ranklist.size() - nslist[rank];
+  // First count the number of blocks on each rank
+  nblist = std::vector<int>(Globals::nranks, 0);
+  for (int b = ranklist.size() - 1; b >= 0; --b)
+    nblist[ranklist[b]]++;
+
+  // Then find the starting gid of the blocks, assuming they
+  // are apportioned in increasing order
+  nslist = std::vector<int>(Globals::nranks, 0);
+  for (int b = 1; b < nslist.size(); ++b)
+    nslist[b] = nslist[b - 1] + nblist[b - 1];
 }
 } // namespace
 
@@ -735,30 +733,32 @@ void Mesh::RedistributeAndRefineMeshBlocks(ParameterInput *pin, ApplicationInput
   int nbe = nbs + nblist[Globals::my_rank] - 1;
 
   // Restrict fine to coarse buffers
-  ProResCache_t restriction_cache;
-  int nrestrict = 0;
-  for (int on = onbs; on <= onbe; on++) {
-    int nn = oldtonew[on];
-    auto pmb = FindMeshBlock(on);
-    if (newloc[nn].level() < loclist[on].level()) nrestrict += pmb->vars_cc_.size();
-  }
-  restriction_cache.Initialize(nrestrict, resolved_packages.get());
-  int irestrict = 0;
-  for (int on = onbs; on <= onbe; on++) {
-    int nn = oldtonew[on];
-    if (newloc[nn].level() < loclist[on].level()) {
+  if (block_list.size() > 0) {
+    ProResCache_t restriction_cache;
+    int nrestrict = 0;
+    for (int on = onbs; on <= onbe; on++) {
+      int nn = oldtonew[on];
       auto pmb = FindMeshBlock(on);
-      for (auto &var : pmb->vars_cc_) {
-        restriction_cache.RegisterRegionHost(
-            irestrict++, ProResInfo::GetInteriorRestrict(pmb.get(), NeighborBlock(), var),
-            var.get(), resolved_packages.get());
+      if (newloc[nn].level() < loclist[on].level()) nrestrict += pmb->vars_cc_.size();
+    }
+    restriction_cache.Initialize(nrestrict, resolved_packages.get());
+    int irestrict = 0;
+    for (int on = onbs; on <= onbe; on++) {
+      int nn = oldtonew[on];
+      if (newloc[nn].level() < loclist[on].level()) {
+        auto pmb = FindMeshBlock(on);
+        for (auto &var : pmb->vars_cc_) {
+          restriction_cache.RegisterRegionHost(
+              irestrict++,
+              ProResInfo::GetInteriorRestrict(pmb.get(), NeighborBlock(), var), var.get(),
+              resolved_packages.get());
+        }
       }
     }
+    restriction_cache.CopyToDevice();
+    refinement::Restrict(resolved_packages.get(), restriction_cache,
+                         block_list[0]->cellbounds, block_list[0]->c_cellbounds);
   }
-  restriction_cache.CopyToDevice();
-  refinement::Restrict(resolved_packages.get(), restriction_cache,
-                       block_list[0]->cellbounds, block_list[0]->c_cellbounds);
-
   Kokkos::fence();
 
 #ifdef MPI_PARALLEL
@@ -912,23 +912,26 @@ void Mesh::RedistributeAndRefineMeshBlocks(ParameterInput *pin, ApplicationInput
       auto pmb = FindMeshBlock(nn);
       if (newloc[nn].level() > loclist[on].level()) nprolong += pmb->vars_cc_.size();
     }
-    prolongation_cache.Initialize(nprolong, resolved_packages.get());
-    int iprolong = 0;
-    for (int nn = nbs; nn <= nbe; nn++) {
-      int on = newtoold[nn];
-      if (newloc[nn].level() > loclist[on].level()) {
-        auto pmb = FindMeshBlock(nn);
-        for (auto &var : pmb->vars_cc_) {
-          prolongation_cache.RegisterRegionHost(
-              iprolong++,
-              ProResInfo::GetInteriorProlongate(pmb.get(), NeighborBlock(), var),
-              var.get(), resolved_packages.get());
+    if (nprolong > 0) {
+      prolongation_cache.Initialize(nprolong, resolved_packages.get());
+      int iprolong = 0;
+      for (int nn = nbs; nn <= nbe; nn++) {
+        int on = newtoold[nn];
+        if (newloc[nn].level() > loclist[on].level()) {
+          auto pmb = FindMeshBlock(nn);
+          for (auto &var : pmb->vars_cc_) {
+            prolongation_cache.RegisterRegionHost(
+                iprolong++,
+                ProResInfo::GetInteriorProlongate(pmb.get(), NeighborBlock(), var),
+                var.get(), resolved_packages.get());
+          }
         }
+        prolongation_cache.CopyToDevice();
       }
-      prolongation_cache.CopyToDevice();
+      refinement::ProlongateShared(resolved_packages.get(), prolongation_cache,
+                                   block_list[0]->cellbounds,
+                                   block_list[0]->c_cellbounds);
     }
-    refinement::ProlongateShared(resolved_packages.get(), prolongation_cache,
-                                 block_list[0]->cellbounds, block_list[0]->c_cellbounds);
 
     // update the lists
     loclist = std::move(newloc);
@@ -978,9 +981,11 @@ void Mesh::RedistributeAndRefineMeshBlocks(ParameterInput *pin, ApplicationInput
                                     // ghosts of non-cell centered vars may get some junk
       // Now there is the correct data for prolongating on un-shared topological elements
       // on the new fine blocks
-      refinement::ProlongateInternal(resolved_packages.get(), prolongation_cache,
-                                     block_list[0]->cellbounds,
-                                     block_list[0]->c_cellbounds);
+      if (nprolong > 0) {
+        refinement::ProlongateInternal(resolved_packages.get(), prolongation_cache,
+                                       block_list[0]->cellbounds,
+                                       block_list[0]->c_cellbounds);
+      }
     }
 
     // Rebuild just the ownership model, this time weighting the "new" fine blocks just

diff --git a/src/mesh/mesh.cpp b/src/mesh/mesh.cpp
@@ -326,12 +326,7 @@ void Mesh::BuildBlockList(ParameterInput *pin, ApplicationInput *app_in,
 #ifdef MPI_PARALLEL
   // check if there are sufficient blocks
   if (nbtotal < Globals::nranks) {
-    if (mesh_test == 0) {
-      msg << "### FATAL ERROR in Mesh constructor" << std::endl
-          << "Too few mesh blocks: nbtotal (" << nbtotal << ") < nranks ("
-          << Globals::nranks << ")" << std::endl;
-      PARTHENON_FAIL(msg);
-    } else { // test
+    if (mesh_test != 0) {
       std::cout << "### Warning in Mesh constructor" << std::endl
                 << "Too few mesh blocks: nbtotal (" << nbtotal << ") < nranks ("
                 << Globals::nranks << ")" << std::endl;
@@ -356,6 +351,7 @@ void Mesh::BuildBlockList(ParameterInput *pin, ApplicationInput *app_in,
   // create MeshBlock list for this process
   int nbs = nslist[Globals::my_rank];
   int nbe = nbs + nblist[Globals::my_rank] - 1;
+
   // create MeshBlock list for this process
   block_list.clear();
   block_list.resize(nbe - nbs + 1);
@@ -398,9 +394,6 @@ void Mesh::BuildBlockPartitions(GridIdentifier grid) {
   auto partition_blocklists = partition::ToSizeN(
       grid.type == GridType::leaf ? block_list : gmg_block_lists[grid.logical_level],
       DefaultPackSize());
-  // Account for possibly empty block_list
-  if (partition_blocklists.size() == 0)
-    partition_blocklists = std::vector<BlockList_t>(1);
   std::vector<std::shared_ptr<BlockListPartition>> out;
   int id = 0;
   for (auto &part_bl : partition_blocklists)
@@ -725,6 +718,8 @@ void Mesh::PreCommFillDerived() {
   }
   for (auto &partition : GetDefaultBlockPartitions()) {
     auto &md = mesh_data.Add("base", partition);
+    PARTHENON_REQUIRE(partition->pmesh == this, "Bad partition mesh pointer");
+    PARTHENON_REQUIRE(md->GetParentPointer() == this, "Bad mesh pointer");
     Update::PreCommFillDerived(md.get());
   }
 }
@@ -847,12 +842,25 @@ void Mesh::Initialize(bool init_problem, ParameterInput *pin, ApplicationInput *
     }
   } while (!init_done);
 
+#ifdef MPI_PARALLEL
+  // check if there are sufficient blocks
+  if (nbtotal < Globals::nranks) {
+    std::stringstream msg;
+    msg << "### FATAL ERROR in Mesh Initialize" << std::endl
+        << "Too few mesh blocks after initialization: nbtotal (" << nbtotal
+        << ") < nranks (" << Globals::nranks << ")" << std::endl;
+    PARTHENON_FAIL(msg);
+  }
+#endif
+
   // Initialize the "base" MeshData object
   mesh_data.Get()->Initialize(block_list, this);
 }
 
 /// Finds location of a block with ID `tgid`.
 std::shared_ptr<MeshBlock> Mesh::FindMeshBlock(int tgid) const {
+  PARTHENON_REQUIRE(block_list.size() > 0,
+                    "Trying to call FindMeshBlock with empty block list");
   // Attempt to simply index into the block list.
   const int nbs = block_list[0]->gid;
   const int i = tgid - nbs;
@@ -878,24 +886,24 @@ bool Mesh::SetBlockSizeAndBoundaries(LogicalLocation loc, RegionSize &block_size
 }
 
 std::int64_t Mesh::GetTotalCells() {
-  auto &pmb = block_list.front();
-  return static_cast<std::int64_t>(nbtotal) * pmb->block_size.nx(X1DIR) *
-         pmb->block_size.nx(X2DIR) * pmb->block_size.nx(X3DIR);
+  return static_cast<std::int64_t>(nbtotal) * GetNumberOfMeshBlockCells();
 }
 
-// TODO(JMM): Move block_size into mesh.
 int Mesh::GetNumberOfMeshBlockCells() const {
-  return block_list.front()->GetNumberOfMeshBlockCells();
+  return base_block_size.nx(X1DIR) * base_block_size.nx(X2DIR) *
+         base_block_size.nx(X3DIR);
 }
 
-const IndexShape &Mesh::GetLeafBlockCellBounds(CellLevel level) const {
-  MeshBlock *pmb = block_list[0].get();
+const IndexShape Mesh::GetLeafBlockCellBounds(CellLevel level) const {
+  auto shapes = GetIndexShapes(
+      ndim > 0 ? base_block_size.nx(X1DIR) : 0, ndim > 1 ? base_block_size.nx(X2DIR) : 0,
+      ndim > 2 ? base_block_size.nx(X3DIR) : 0, multilevel, this);
   if (level == CellLevel::same) {
-    return pmb->cellbounds;
+    return shapes[0];
   } else if (level == CellLevel::fine) {
-    return pmb->f_cellbounds;
+    return shapes[1];
   } else { // if (level == CellLevel::coarse) {
-    return pmb->c_cellbounds;
+    return shapes[2];
   }
 }
 

diff --git a/src/mesh/mesh.hpp b/src/mesh/mesh.hpp
@@ -106,7 +106,7 @@ class Mesh {
   RegionSize GetBlockSize(const LogicalLocation &loc) const {
     return forest.GetBlockDomain(loc);
   }
-  const IndexShape &GetLeafBlockCellBounds(CellLevel level = CellLevel::same) const;
+  const IndexShape GetLeafBlockCellBounds(CellLevel level = CellLevel::same) const;
 
   const forest::Forest &Forest() const { return forest; }
 
@@ -143,7 +143,8 @@ class Mesh {
   void LoadBalancingAndAdaptiveMeshRefinement(ParameterInput *pin,
                                               ApplicationInput *app_in);
   int DefaultPackSize() {
-    return default_pack_size_ < 1 ? block_list.size() : default_pack_size_;
+    return default_pack_size_ < 1 ? std::max(static_cast<int>(block_list.size()), 1)
+                                  : default_pack_size_;
   }
   int DefaultNumPartitions() {
     return partition::partition_impl::IntCeil(block_list.size(), DefaultPackSize());

diff --git a/src/mesh/meshblock.cpp b/src/mesh/meshblock.cpp
@@ -201,27 +201,34 @@ void MeshBlock::Initialize(int igid, int ilid, LogicalLocation iloc,
 
 MeshBlock::~MeshBlock() = default;
 
+std::array<IndexShape, 3> GetIndexShapes(const int nx1, const int nx2, const int nx3,
+                                         bool multilevel, const Mesh *pmesh) {
+  IndexShape cellbounds(nx3, nx2, nx1, Globals::nghost);
+  IndexShape f_cellbounds(2 * nx3, 2 * nx2, 2 * nx1, Globals::nghost);
+  IndexShape c_cellbounds(nx3 / 2, nx2 / 2, nx1 / 2, 0);
+  if (multilevel) {
+    // Prevent the coarse bounds from going to zero
+    int cnx1 = nx1 / 2;
+    int cnx2 = nx2 / 2;
+    int cnx3 = nx3 / 2;
+    if (pmesh != nullptr) {
+      cnx1 = pmesh->mesh_size.symmetry(X1DIR) ? 0 : std::max(1, nx1 / 2);
+      cnx2 = pmesh->mesh_size.symmetry(X2DIR) ? 0 : std::max(1, nx2 / 2);
+      cnx3 = pmesh->mesh_size.symmetry(X3DIR) ? 0 : std::max(1, nx3 / 2);
+    }
+    c_cellbounds = IndexShape(cnx3, cnx2, cnx1, Globals::nghost);
+  }
+  return {cellbounds, f_cellbounds, c_cellbounds};
+}
+
 void MeshBlock::InitializeIndexShapesImpl(const int nx1, const int nx2, const int nx3,
                                           bool init_coarse, bool multilevel) {
-  cellbounds = IndexShape(nx3, nx2, nx1, Globals::nghost);
-  f_cellbounds = IndexShape(2 * nx3, 2 * nx2, 2 * nx1, Globals::nghost);
-
+  auto [cb, fcb, ccb] = GetIndexShapes(nx1, nx2, nx3, multilevel, pmy_mesh);
+  cellbounds = cb;
+  f_cellbounds = fcb;
   if (init_coarse) {
-    if (multilevel) {
-      // Prevent the coarse bounds from going to zero
-      int cnx1 = nx1 / 2;
-      int cnx2 = nx2 / 2;
-      int cnx3 = nx3 / 2;
-      if (pmy_mesh != nullptr) {
-        cnx1 = pmy_mesh->mesh_size.symmetry(X1DIR) ? 0 : std::max(1, nx1 / 2);
-        cnx2 = pmy_mesh->mesh_size.symmetry(X2DIR) ? 0 : std::max(1, nx2 / 2);
-        cnx3 = pmy_mesh->mesh_size.symmetry(X3DIR) ? 0 : std::max(1, nx3 / 2);
-      }
-      cnghost = (Globals::nghost + 1) / 2 + 1;
-      c_cellbounds = IndexShape(cnx3, cnx2, cnx1, Globals::nghost);
-    } else {
-      c_cellbounds = IndexShape(nx3 / 2, nx2 / 2, nx1 / 2, 0);
-    }
+    cnghost = (Globals::nghost + 1) / 2 + 1;
+    c_cellbounds = ccb;
   }
 }
 

diff --git a/src/mesh/meshblock.hpp b/src/mesh/meshblock.hpp
@@ -63,6 +63,9 @@ KOKKOS_FORCEINLINE_FUNCTION void par_for_inner(const team_mbr_t &team_member,
   parthenon::par_for_inner(DEFAULT_INNER_LOOP_PATTERN, team_member, il, iu, function);
 }
 
+std::array<IndexShape, 3> GetIndexShapes(const int nx1, const int nx2, const int nx3,
+                                         bool multilevel, const Mesh *pmesh);
+
 //----------------------------------------------------------------------------------------
 //! \class MeshBlock
 //! \brief data/functions associated with a single block