stage_target.hpp

Go to the documentation of this file.

// core — the TARGET of a coarse-level multigrid stage: onto which decomposition of a level's own
// grid the level moves when it can no longer coarsen in place (suite decision "Coarse-level
// redistribution lives in core; the hierarchies stay in the methods", amr/docs/
// amr_mg_core_boundary.md §4).
//
// A pure function of the current decomposition and the METHOD's lift predicate, replicated on every
// rank with no communication. Core does not know what makes a decomposition coarsenable — flow's
// rule is per-axis (every block even in origin and size on each axis whose global extent can still
// halve), amr's is all-axes — so the caller supplies it as `liftable(dec)`.
//
// The kinds:
//   InPlace      — no stage: the level coarsens on its current decomposition.
//   SiblingMerge — the ORB tree truncated (`BlockDecomposer::agglomerated(d)`): each target block
//                  is the union of a contiguous run of current blocks, owned by the lowest of their
//                  ranks. Moved by a group gather / scatter (stage_comm.hpp,
//                  redistribute_topology.hpp).
//   Repartition  — a fresh proportional ORB of the level grid on the first np_L parent ranks
//                  (repartitionTarget, §11.2 of the design). Moved point-to-point over the box
//                  intersections of the two decompositions (redistribute_topology.hpp).
//   Replicated   — the whole level as one block on EVERY rank. Moved by an Allgatherv.
//
// With `maxBlockCells == 0` the policy is flow's telescope search verbatim
// (flow/src/mac_cutcell_mg.hpp, CutcellMG::initMpi): so that flow's migration onto it is
// byte-identical. `maxBlockCells > 0` (design §11.1) bounds the block a stage may hand a rank.
#ifndef PECLET_CORE_DECOMP_STAGE_TARGET_HPP
#define PECLET_CORE_DECOMP_STAGE_TARGET_HPP
 
#include <cstddef>
#include <limits>
#include <optional>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>
 
#include "peclet/core/common/types.hpp"
#include "peclet/core/decomp/block_decomposer.hpp"
 
namespace peclet::core::decomp {
 
enum class StageKind {
  InPlace,       
  SiblingMerge,  
  Repartition,   
  Replicated     
};
 
template <int Dim>
struct StageTarget {
  StageKind kind = StageKind::InPlace;  
  BlockDecomposer<Dim> dec;
  std::vector<int> ownerOf;
  std::vector<int> groupOf;
};
 
template <int Dim>
inline Index minBlockExtent(const BlockDecomposer<Dim>& d) {
  Index m = std::numeric_limits<Index>::max();
  for (const auto& s : d.sizes())
    for (int k = 0; k < Dim; ++k)
      m = s[k] < m ? s[k] : m;
  return m;
}
 
template <int Dim>
inline Index largestBlockCells(const BlockDecomposer<Dim>& d) {
  Index m = 0;
  for (const auto& s : d.sizes()) {
    Index c = 1;
    for (int k = 0; k < Dim; ++k)
      c *= s[k];
    m = c > m ? c : m;
  }
  return m;
}
 
template <int Dim>
struct SiblingMergeChoice {
  int depth = -1;            
  BlockDecomposer<Dim> dec;  
  std::vector<int> groupOf;  
  std::vector<int> ownerOf;  
};
 
template <int Dim, class Liftable>
std::optional<SiblingMergeChoice<Dim>> shallowestLiftableMerge(const BlockDecomposer<Dim>& cur,
                                                               Liftable&& liftable, int minExtent) {
  for (int d = cur.treeDepth() - 1; d >= 0; --d) {
    SiblingMergeChoice<Dim> m;
    m.depth = d;
    m.dec = cur.agglomerated(d, &m.groupOf, &m.ownerOf);
    bool ok = liftable(m.dec);
    // Fat enough to STAY above the threshold after the halving that follows (a single block always
    // qualifies): merging to blocks of 4 that become 2 on the next level would merge again there.
    if (ok && minExtent > 0 && m.dec.numBlocks() > 1 &&
        minBlockExtent(m.dec) < 2 * static_cast<Index>(minExtent))
      ok = false;
    if (ok && m.dec.numBlocks() < cur.numBlocks())
      return m;
  }
  return std::nullopt;
}
 
template <int Dim, class Liftable>
std::optional<StageTarget<Dim>> repartitionTarget(const BlockDecomposer<Dim>& cur,
                                                  const std::type_identity_t<IVec<Dim>>& levelGrid,
                                                  Liftable&& liftable, int minExtent,
                                                  Index maxBlockCells) {
  if (maxBlockCells <= 0)
    throw std::invalid_argument("repartitionTarget: maxBlockCells must be positive");
  const Index np = static_cast<Index>(cur.numBlocks());
  Index cells = 1;
  for (int k = 0; k < Dim; ++k)
    cells *= levelGrid[k];
  Index npL = (cells + maxBlockCells - 1) / maxBlockCells;
  Index cap = np;  // the extent cap; no cap beyond the rank count at minExtent == 0
  if (minExtent > 0) {
    Index c = 1;
    for (int k = 0; k < Dim; ++k) {
      const Index f = levelGrid[k] / (2 * static_cast<Index>(minExtent));
      c *= f > 1 ? f : 1;
    }
    cap = c < np ? c : np;
    npL = npL < cap ? npL : cap;
  }
  Index up = 1;  // nextPow2(npL), stopping once it reaches np (the clamp below)
  while (up < npL && up < np)
    up *= 2;
  npL = up < np ? up : np;
  if (npL > cap) {  // the round-up overshot a cap that is not a power of two: the largest power
    Index p = 1;    // of two within the cap keeps the blocks fat enough (the cap's whole point)
    while (2 * p <= cap)
      p *= 2;
    npL = p;
  }
  npL = npL < 1 ? 1 : npL;
  Index p2 = 1;  // the largest power of two <= npL
  while (2 * p2 <= npL)
    p2 *= 2;
  std::vector<Index> tries;  // npL, then p2 (unless it is npL), then its halvings down to 1
  if (npL != p2)
    tries.push_back(npL);
  for (Index n = p2; n >= 1; n /= 2)
    tries.push_back(n);
  for (const Index n : tries) {
    BlockDecomposer<Dim> r(static_cast<std::size_t>(n), levelGrid);
    if (liftable(r)) {
      StageTarget<Dim> t;
      t.kind = StageKind::Repartition;
      t.dec = std::move(r);
      t.ownerOf.resize(static_cast<std::size_t>(n));
      for (Index b = 0; b < n; ++b)
        t.ownerOf[static_cast<std::size_t>(b)] = static_cast<int>(b);
      return t;
    }
  }
  return std::nullopt;
}
 
template <int Dim, class Liftable>
StageTarget<Dim> chooseStageTarget(const BlockDecomposer<Dim>& cur,
                                   const std::type_identity_t<IVec<Dim>>& levelGrid,
                                   Liftable&& liftable, int minExtent, Index maxBlockCells = 0) {
  if (levelGrid != cur.globalSize())
    throw std::invalid_argument(
        "chooseStageTarget: levelGrid differs from the current decomposition's global size");
  if (maxBlockCells < 0)
    throw std::invalid_argument("chooseStageTarget: maxBlockCells must be >= 0 (0 disables)");
 
  const std::size_t nb = cur.numBlocks();
  StageTarget<Dim> t;
 
  const bool tooSmall = minExtent > 0 && minBlockExtent(cur) < static_cast<Index>(minExtent);
  if (nb <= 1 || (liftable(cur) && !tooSmall)) {
    t.kind = StageKind::InPlace;
    t.dec = cur;
    t.ownerOf.resize(nb);
    for (std::size_t b = 0; b < nb; ++b)
      t.ownerOf[b] = static_cast<int>(b);
    return t;
  }
 
  auto m = shallowestLiftableMerge(cur, liftable, minExtent);
  auto merge = [&] {
    t.kind = StageKind::SiblingMerge;
    t.dec = std::move(m->dec);
    t.ownerOf = std::move(m->ownerOf);
    t.groupOf = std::move(m->groupOf);
    return std::move(t);
  };
  if (m && (maxBlockCells == 0 || largestBlockCells(m->dec) <= maxBlockCells))
    return merge();
 
  if (maxBlockCells > 0)
    if (auto r = repartitionTarget(cur, levelGrid, liftable, minExtent, maxBlockCells))
      return std::move(*r);
 
  if (m)
    return merge();
 
  t.kind = StageKind::Replicated;
  t.dec.init(1, levelGrid);
  t.ownerOf = {0};
  return t;
}
 
}  // namespace peclet::core::decomp
 
#endif  // PECLET_CORE_DECOMP_STAGE_TARGET_HPP