mg_predict.hpp

Go to the documentation of this file.

// amr — what pressure-multigrid ladder a given grid, tree depth and rank count will produce.
//
// The prediction tool of docs/amr_mg_depth.md §6.7. It is a PURE HOST FUNCTION (no MPI, no Kokkos,
// no octree): it re-runs the §6.2 ladder rule on the ORB the distributed octree would build
// (peclet::core::decomp::BlockDecomposer on the global ROOT grid), so a test can assert the
// hierarchy the solver built against the hierarchy the design says it should build, instead of
// against a literal that goes stale.
//
// The rule, once:
//   * levels above the root brick are the octree's own coarsenings — `lmax + 1` of them for a mesh
//     refined to level 0 (the case §6.2's table describes; see the `lmax` note on the function);
//   * below the root the ladder LIFTS while the global grid halves into a cube level AND every
//     rank's block stays even in origin and size (§6.2); the lift is lockstep, so one rank's odd
//     block stops the ladder for all;
//   * when the coarsest in-place grid is still larger than `bottomExtent` on some axis and there
//     is more than one rank, the tail gathers it and continues redundantly with the single-rank
//     ladder (§6.5);
//   * the bottom is exact damped Jacobi while the final extent is <= `bottomExtent`, and the
//     agglomerated GraphAMG-PCG solve otherwise (§6.6).
//
// The prediction was written before the tail (WO4) and the GraphAMG bottom (WO5) were built, as
// the ladder of the finished design; both have since landed and match it, which is what makes it
// the specification the built ladder is checked against (`amr_mg_predict`, `amr_mg_tail`,
// `amr_mg_bottom{,_dist}`). It predicts the default bottom selection, `auto`, on the ORB a
// freshly initialised DistributedOctree builds — not a weighted or rebalanced partition.
#ifndef PECLET_AMR_MG_PREDICT_HPP
#define PECLET_AMR_MG_PREDICT_HPP
 
#include <algorithm>
#include <cstddef>
#include <string>
#include <vector>
 
#include "peclet/amr/common.hpp"
#include "peclet/core/decomp/block_decomposer.hpp"
#include "peclet/core/decomp/stage_target.hpp"
 
namespace peclet::amr {
 
struct PressureStagePolicy {
  bool enabled = true;  
  int minExtent = 4;
  Index maxBlockCells = -1;
};
 
template <int Dim>
bool amrStageLiftable(const core::decomp::BlockDecomposer<Dim>& dec) {
  const IVec<Dim>& G = dec.globalSize();
  for (int a = 0; a < Dim; ++a)
    if ((G[a] % 2) != 0 || (G[a] / 2) < 2)
      return false;
  for (std::size_t b = 0; b < dec.numBlocks(); ++b)
    for (int a = 0; a < Dim; ++a)
      if ((dec.origins()[b][a] % 2) != 0 || (dec.sizes()[b][a] % 2) != 0)
        return false;
  return true;
}
 
enum class MgLevelKind {
  Octree,       
  Lifted,       
  Tail,         
  Sibling,      
  Repartition,  
};
 
enum class MgBottomKind {
  Jacobi,  
  Amg      
};
 
inline const char* toString(MgLevelKind k) {
  switch (k) {
    case MgLevelKind::Octree:
      return "octree";
    case MgLevelKind::Lifted:
      return "lifted";
    case MgLevelKind::Tail:
      return "tail";
    case MgLevelKind::Sibling:
      return "sibling";
    case MgLevelKind::Repartition:
      return "repartition";
  }
  return "?";
}
inline const char* toString(MgBottomKind k) {
  return k == MgBottomKind::Jacobi ? "jacobi" : "amg";
}
 
template <int Dim>
struct MgLevelPrediction {
  IVec<Dim> extent{};   
  long long cells = 0;  
  MgLevelKind kind = MgLevelKind::Octree;  
};
 
template <int Dim>
struct MgLadderPrediction {
  std::vector<MgLevelPrediction<Dim>> levels;
  std::vector<MgLevelKind> stages;
  bool tail = false;                           
  IVec<Dim> tailFrom{};                        
  MgBottomKind bottom = MgBottomKind::Jacobi;  
 
  bool hasStage() const { return !stages.empty(); }
  std::size_t numInPlace() const {
    std::size_t k = 0;
    while (k < levels.size() &&
           (levels[k].kind == MgLevelKind::Octree || levels[k].kind == MgLevelKind::Lifted))
      ++k;
    return k;
  }
  std::string bottomName() const {
    std::string s = toString(bottom);
    for (auto it = stages.rbegin(); it != stages.rend(); ++it)
      s += std::string("+") + toString(*it);
    return s;
  }
};
 
namespace detail {
 
template <int Dim>
Index maxExtent(const IVec<Dim>& G) {
  Index mx = 0;
  for (int d = 0; d < Dim; ++d)
    mx = std::max(mx, G[d]);
  return mx;
}
 
template <int Dim>
bool gridCanHalve(const IVec<Dim>& G, Index bottomExtent) {
  if (maxExtent<Dim>(G) <= bottomExtent)
    return false;
  for (int d = 0; d < Dim; ++d)
    if ((G[d] % 2) != 0 || (G[d] / 2) < 2)
      return false;
  return true;
}
 
template <int Dim>
long long cellsOf(const IVec<Dim>& G) {
  long long c = 1;
  for (int d = 0; d < Dim; ++d)
    c *= static_cast<long long>(G[d]);
  return c;
}
 
template <int Dim>
void pushLevel(MgLadderPrediction<Dim>& p, const IVec<Dim>& G, MgLevelKind kind) {
  MgLevelPrediction<Dim> lv;
  lv.extent = G;
  lv.cells = cellsOf<Dim>(G);
  lv.kind = kind;
  p.levels.push_back(lv);
}
 
}  // namespace detail
 
template <int Dim>
MgLadderPrediction<Dim> predictPressureLadder(IVec<Dim> G, unsigned lmax, int numRanks,
                                              Index bottomExtent = 4,
                                              const PressureStagePolicy& policy = {}) {
  MgLadderPrediction<Dim> p;
  if (numRanks < 1)
    numRanks = 1;
 
  // --- the octree levels: the finest grid coarsened down to the root brick.
  for (unsigned j = 0; j <= lmax; ++j) {
    IVec<Dim> e{};
    for (int d = 0; d < Dim; ++d)
      e[d] = G[d] << (lmax - j);
    detail::pushLevel<Dim>(p, e, MgLevelKind::Octree);
  }
 
  // --- the in-place lifted levels: lockstep while the grid halves and every block stays even —
  //     then, where they block above `bottomExtent`, a STAGE, whose continued ladder lifts again on
  //     its own decomposition and may stage again (the same loop, on the target decomposition).
  namespace cd = core::decomp;
  cd::BlockDecomposer<Dim> dec(static_cast<std::size_t>(numRanks), G);
  Index maxBlockCells = policy.maxBlockCells;
  if (maxBlockCells < 0) {  // the leaf count of the largest finest block, uniform mesh (§9.6)
    maxBlockCells = static_cast<Index>(cd::largestBlockCells(dec));
    for (unsigned j = 0; j < lmax; ++j)
      maxBlockCells <<= Dim;
  }
  IVec<Dim> cur = G;
  MgLevelKind liftKind = MgLevelKind::Lifted;
  for (;;) {
    while (detail::gridCanHalve<Dim>(cur, bottomExtent) && amrStageLiftable<Dim>(dec)) {
      IVec<Dim> two{};
      for (int d = 0; d < Dim; ++d)
        two[d] = 2;
      dec = dec.coarsened(two);
      cur = dec.globalSize();
      detail::pushLevel<Dim>(p, cur, liftKind);
    }
    if (dec.numBlocks() <= 1 || detail::maxExtent<Dim>(cur) <= bottomExtent)
      break;  // one rank: its bottom solves the level whole; or the Jacobi bottom is exact here
    if (!p.hasStage())
      p.tailFrom = cur;
    if (policy.enabled) {
      const cd::StageTarget<Dim> t = cd::chooseStageTarget(
          dec, cur, [](const cd::BlockDecomposer<Dim>& c) { return amrStageLiftable<Dim>(c); },
          policy.minExtent, maxBlockCells);
      if (t.kind == cd::StageKind::SiblingMerge || t.kind == cd::StageKind::Repartition) {
        liftKind =
            t.kind == cd::StageKind::SiblingMerge ? MgLevelKind::Sibling : MgLevelKind::Repartition;
        p.stages.push_back(liftKind);
        detail::pushLevel<Dim>(p, cur, liftKind);  // the continued ladder's level 0: moved
        dec = t.dec;
        continue;
      }
    }
    // --- the redundant tail: the gathered grid, continued on one rank.
    p.tail = true;
    p.stages.push_back(MgLevelKind::Tail);
    detail::pushLevel<Dim>(p, cur,
                           MgLevelKind::Tail);  // the tail's own level 0 == the gathered grid
    while (detail::gridCanHalve<Dim>(cur, bottomExtent)) {
      for (int d = 0; d < Dim; ++d)
        cur[d] /= 2;
      detail::pushLevel<Dim>(p, cur, MgLevelKind::Tail);
    }
    break;
  }
 
  p.bottom = detail::maxExtent<Dim>(cur) > bottomExtent ? MgBottomKind::Amg : MgBottomKind::Jacobi;
  return p;
}
 
}  // namespace peclet::amr
 
#endif  // PECLET_AMR_MG_PREDICT_HPP