block_decomposer.hpp

Go to the documentation of this file.

// core — orthogonal recursive bisection (ORB) domain decomposition.
//
// Ported from block_decomposer/src/BlockDecomposer.hpp (pbs::BlockDecomposer), modernized into the
// peclet::core namespace. The global cell grid is split recursively along its largest axis into
// `numBlocks` rank-owned blocks. Adds ownerOf() (a tree walk) for halo topology construction.
#ifndef PECLET_CORE_DECOMP_BLOCK_DECOMPOSER_HPP
#define PECLET_CORE_DECOMP_BLOCK_DECOMPOSER_HPP
 
#include <array>
#include <cassert>
#include <cmath>
#include <cstddef>
#include <limits>
#include <stack>
#include <vector>
 
#include "peclet/core/common/types.hpp"
 
namespace peclet::core::decomp {
 
template <int Dim>
struct Block {
  IVec<Dim> origin{};  
  IVec<Dim> size{};    
};
 
template <int Dim>
class BlockDecomposer {
 public:
  BlockDecomposer() = default;
  BlockDecomposer(std::size_t numBlocks, IVec<Dim> globalSize) { init(numBlocks, globalSize); }
  BlockDecomposer(std::size_t numBlocks, IVec<Dim> globalSize, const std::vector<Real>& weights) {
    init(numBlocks, globalSize, weights);
  }
  BlockDecomposer(std::size_t numBlocks, IVec<Dim> globalSize, const IVec<Dim>& align) {
    init(numBlocks, globalSize, align);
  }
 
  void init(std::size_t numBlocks, IVec<Dim> globalSize) {
    align_ = IVec<Dim>{};
    for (int i = 0; i < Dim; ++i)
      align_[i] = 1;
    initImpl(numBlocks, globalSize, nullptr);
  }
 
  void init(std::size_t numBlocks, IVec<Dim> globalSize, const IVec<Dim>& align) {
    align_ = align;
    for (int i = 0; i < Dim; ++i)
      if (align_[i] < 1)
        align_[i] = 1;
    initImpl(numBlocks, globalSize, nullptr);
  }
 
  void init(std::size_t numBlocks, IVec<Dim> globalSize, const IVec<Dim>& align,
            const IVec<Dim>& cellExtent) {
    cellExtent_ = cellExtent;
    init(numBlocks, globalSize, align);
  }
 
  void init(std::size_t numBlocks, IVec<Dim> globalSize, const std::vector<Real>& weights) {
    assert(weights.size() == static_cast<std::size_t>([&] {
             Index v = 1;
             for (int i = 0; i < Dim; ++i)
               v *= globalSize[i];
             return v;
           }()) &&
           "weights array must cover the global grid (x-fastest)");
    align_ = IVec<Dim>{};
    for (int i = 0; i < Dim; ++i)
      align_[i] = 1;
    initImpl(numBlocks, globalSize, &weights);
  }
 
  void init(std::size_t numBlocks, IVec<Dim> globalSize, const std::vector<Real>& weights,
            const IVec<Dim>& align) {
    IVec<Dim> gc{};
    Index nc = 1;
    for (int k = 0; k < Dim; ++k) {
      assert(align[k] >= 1 && "aligned weighted init(): align must be >= 1");
      assert(globalSize[k] % align[k] == 0 &&
             "aligned weighted init(): globalSize must be a multiple of align");
      gc[k] = globalSize[k] / align[k];
      nc *= gc[k];
    }
    // The coarse weight field. The fine grid is walked x-fastest, so every coarse cell receives
    // its align-box in the box's own x-fastest order; the first cell of a box ASSIGNS (a one-term
    // sum is then a bitwise copy), the rest accumulate. Deterministic, hence replicated.
    std::vector<Real> wc(static_cast<std::size_t>(nc), 0.0);
    BlockDecomposer<Dim> fineIndex;  // x-fastest indexing on the fine grid only
    fineIndex.globalSize_ = globalSize;
    BlockDecomposer<Dim> coarse;
    coarse.globalSize_ = gc;
    assert(weights.size() == static_cast<std::size_t>([&] {
             Index v = 1;
             for (int i = 0; i < Dim; ++i)
               v *= globalSize[i];
             return v;
           }()) &&
           "weights array must cover the global grid (x-fastest)");
    forEachInBox<Dim>(IVec<Dim>{}, globalSize, [&](const IVec<Dim>& g) {
      IVec<Dim> c{};
      bool first = true;
      for (int k = 0; k < Dim; ++k) {
        c[k] = g[k] / align[k];
        if (g[k] % align[k] != 0)
          first = false;
      }
      const std::size_t ic = static_cast<std::size_t>(coarse.linearGlobal(c));
      const Real wv = weights[static_cast<std::size_t>(fineIndex.linearGlobal(g))];
      if (first)
        wc[ic] = wv;
      else
        wc[ic] += wv;
    });
    coarse.cellExtent_ = align;
    coarse.init(numBlocks, gc, wc);  // the existing weighted ORB (sets align_ = 1)
    *this = coarse.refined(align);   // the existing exact inverse of coarsened(); align_ = align
  }
 
  std::size_t numBlocks() const { return origins_.size(); }
  const IVec<Dim>& globalSize() const { return globalSize_; }
  const std::vector<IVec<Dim>>& origins() const { return origins_; }
  const std::vector<IVec<Dim>>& sizes() const { return sizes_; }
 
  Block<Dim> block(std::size_t b) const { return {origins_[b], sizes_[b]}; }
 
  int ownerOf(const IVec<Dim>& g) const {
    Index node = 0;
    while (tree_[node].splitDim != -1) {
      node = (g[tree_[node].splitDim] < tree_[node].splitValue) ? 2 * node + 1 : 2 * node + 2;
    }
    return static_cast<int>(tree_[node].splitValue);
  }
 
  void flattenTree(std::vector<int>& splitDim, std::vector<Index>& splitVal) const {
    splitDim.resize(tree_.size());
    splitVal.resize(tree_.size());
    for (std::size_t i = 0; i < tree_.size(); ++i) {
      splitDim[i] = tree_[i].splitDim;
      splitVal[i] = tree_[i].splitValue;
    }
  }
 
  Index linearGlobal(const IVec<Dim>& g) const {
    Index idx = g[Dim - 1];
    for (int i = Dim - 2; i >= 0; --i) {
      idx *= globalSize_[i];
      idx += g[i];
    }
    return idx;
  }
 
  IVec<Dim> multiGlobal(Index lin) const {
    IVec<Dim> g{};
    for (int i = 0; i < Dim; ++i) {
      g[i] = lin % globalSize_[i];
      lin /= globalSize_[i];
    }
    return g;
  }
 
  BlockDecomposer<Dim> coarsened(const IVec<Dim>& ratio) const {
    BlockDecomposer<Dim> c;
    for (int k = 0; k < Dim; ++k) {
      assert(ratio[k] >= 1 && globalSize_[k] % ratio[k] == 0 &&
             "coarsened(): global size not divisible by ratio (build the decomposition aligned)");
      c.globalSize_[k] = globalSize_[k] / ratio[k];
      c.align_[k] = align_[k] / ratio[k] >= 1 ? align_[k] / ratio[k] : 1;
    }
    c.tree_ = tree_;
    for (auto& nd : c.tree_)
      if (nd.splitDim != -1) {
        assert(nd.splitValue % ratio[nd.splitDim] == 0 &&
               "coarsened(): split not divisible by ratio (build the decomposition aligned)");
        nd.splitValue /= ratio[nd.splitDim];
      }
    c.origins_.resize(origins_.size());
    c.sizes_.resize(sizes_.size());
    for (std::size_t b = 0; b < origins_.size(); ++b)
      for (int k = 0; k < Dim; ++k) {
        assert(origins_[b][k] % ratio[k] == 0 && sizes_[b][k] % ratio[k] == 0 &&
               "coarsened(): block origin/size not divisible by ratio (build aligned)");
        c.origins_[b][k] = origins_[b][k] / ratio[k];
        c.sizes_[b][k] = sizes_[b][k] / ratio[k];
      }
    return c;
  }
 
  BlockDecomposer<Dim> refined(const IVec<Dim>& ratio) const {
    BlockDecomposer<Dim> f;
    for (int k = 0; k < Dim; ++k) {
      assert(ratio[k] >= 1 && "refined(): ratio must be >= 1");
      f.globalSize_[k] = globalSize_[k] * ratio[k];
      f.align_[k] = align_[k] * ratio[k];
      f.cellExtent_[k] = 1;  // the refined grid's cells are the fine cells
    }
    f.tree_ = tree_;
    for (auto& nd : f.tree_)
      if (nd.splitDim != -1)
        nd.splitValue *= ratio[nd.splitDim];
    f.origins_.resize(origins_.size());
    f.sizes_.resize(sizes_.size());
    for (std::size_t b = 0; b < origins_.size(); ++b)
      for (int k = 0; k < Dim; ++k) {
        f.origins_[b][k] = origins_[b][k] * ratio[k];
        f.sizes_[b][k] = sizes_[b][k] * ratio[k];
      }
    return f;
  }
 
 private:
  struct TreeNode {
    int splitDim = -1;     
    Index splitValue = 0;  
  };
 
 public:
  int treeDepth() const {
    int best = 0;
    depthWalk(0, 0, best);
    return best;
  }
 
  BlockDecomposer<Dim> agglomerated(int depth, std::vector<int>* groupOf = nullptr,
                                    std::vector<int>* rootOf = nullptr) const {
    assert(depth >= 0 && "agglomerated(): depth must be >= 0");
    BlockDecomposer<Dim> a;
    a.globalSize_ = globalSize_;
    a.align_ = align_;
    a.cellExtent_ = cellExtent_;
    a.tree_.assign(tree_.size(), TreeNode{});  // everything a leaf until visited
    std::vector<int> gmap(origins_.size(), -1);
    std::vector<int> roots;
    agglomWalk(0, 0, depth, IVec<Dim>{}, globalSize_, a, gmap, roots);
    a.tree_.resize(prunedSize(a.tree_));
    if (groupOf)
      *groupOf = gmap;
    if (rootOf)
      *rootOf = roots;
    return a;
  }
 
 private:
  void depthWalk(Index node, int d, int& best) const {
    if (node >= static_cast<Index>(tree_.size()))
      return;
    if (tree_[node].splitDim == -1) {
      if (d > best)
        best = d;
      return;
    }
    depthWalk(2 * node + 1, d + 1, best);
    depthWalk(2 * node + 2, d + 1, best);
  }
  // Collect the (in-order) original leaf indices under `node`.
  void leavesUnder(Index node, std::vector<int>& out) const {
    if (tree_[node].splitDim == -1) {
      out.push_back(static_cast<int>(tree_[node].splitValue));
      return;
    }
    leavesUnder(2 * node + 1, out);
    leavesUnder(2 * node + 2, out);
  }
  void agglomWalk(Index node, int d, int depth, IVec<Dim> origin, IVec<Dim> size,
                  BlockDecomposer<Dim>& a, std::vector<int>& gmap, std::vector<int>& roots) const {
    const TreeNode& nd = tree_[node];
    if (nd.splitDim == -1 || d == depth) {
      // becomes a block of the agglomerated decomposition
      const int newIdx = static_cast<int>(a.origins_.size());
      a.tree_[node] = TreeNode{-1, newIdx};
      a.origins_.push_back(origin);
      a.sizes_.push_back(size);
      std::vector<int> leaves;
      leavesUnder(node, leaves);
      int root = leaves.front();
      for (int b : leaves) {
        gmap[b] = newIdx;
        if (b < root)
          root = b;
      }
      roots.push_back(root);
      return;
    }
    a.tree_[node] = nd;  // kept internal node: same split
    IVec<Dim> lsize = size, rorigin = origin, rsize = size;
    lsize[nd.splitDim] = nd.splitValue - origin[nd.splitDim];
    rorigin[nd.splitDim] = nd.splitValue;
    rsize[nd.splitDim] = origin[nd.splitDim] + size[nd.splitDim] - nd.splitValue;
    agglomWalk(2 * node + 1, d + 1, depth, origin, lsize, a, gmap, roots);
    agglomWalk(2 * node + 2, d + 1, depth, rorigin, rsize, a, gmap, roots);
  }
  static std::size_t prunedSize(const std::vector<TreeNode>& t) {
    // keep up to the last node that is reachable (an internal node or a numbered leaf); leaves
    // default to {-1, 0} which is indistinguishable from block 0, so keep by reachability instead
    std::size_t last = 0;
    std::vector<Index> stack{0};
    while (!stack.empty()) {
      const Index n = stack.back();
      stack.pop_back();
      if (n >= static_cast<Index>(t.size()))
        continue;
      if (static_cast<std::size_t>(n) > last)
        last = static_cast<std::size_t>(n);
      if (t[n].splitDim != -1) {
        stack.push_back(2 * n + 1);
        stack.push_back(2 * n + 2);
      }
    }
    return last + 1;
  }
 
 private:
  void initImpl(std::size_t numBlocks, IVec<Dim> globalSize, const std::vector<Real>* weights);
 
  Index splitPosition(const IVec<Dim>& origin, const IVec<Dim>& size, int kLargest,
                      std::size_t numSub, std::size_t numTotal,
                      const std::vector<Real>* weights) const;
 
  IVec<Dim> globalSize_{};
  IVec<Dim>
      align_{};  
  IVec<Dim> cellExtent_{};  
  std::vector<IVec<Dim>> origins_;
  std::vector<IVec<Dim>> sizes_;
  std::vector<TreeNode> tree_;
};
 
template <int Dim>
void BlockDecomposer<Dim>::initImpl(std::size_t numBlocks, IVec<Dim> globalSize,
                                    const std::vector<Real>* weights) {
  globalSize_ = globalSize;
  for (int k = 0; k < Dim; ++k)
    if (cellExtent_[k] <= 0)
      cellExtent_[k] = 1;
  origins_.clear();
  sizes_.clear();
  tree_.clear();
 
  struct StackBlock {
    IVec<Dim> origin;
    IVec<Dim> size;
    std::size_t numSub;
    Index treeIndx;
  };
  std::stack<StackBlock> stack;
  stack.push(StackBlock{IVec<Dim>{}, globalSize, numBlocks, 0});
 
  Index leafIndex = 0;
  while (!stack.empty()) {
    StackBlock cur = stack.top();
    stack.pop();
 
    if (static_cast<std::size_t>(cur.treeIndx) >= tree_.size()) {
      tree_.resize(cur.treeIndx + 1);
    }
 
    if (cur.numSub > 1) {
      // Split along the largest axis. The split position balances either the sub-block cell count
      // (unweighted) or the cumulative weight (weighted) of the two children.
      // Compare PHYSICAL extents: size*cellExtent. cellExtent is all-ones for an ordinary grid, so
      // this is the plain largest-cell-count test unless an anisotropically coarsened grid says
      // otherwise (see the cellExtent init).
      int kLargest = 0;
      for (int k = 1; k < Dim; ++k) {
        if (cur.size[k] * cellExtent_[k] > cur.size[kLargest] * cellExtent_[kLargest])
          kLargest = k;
      }
      std::size_t numSub = cur.numSub / 2;
      Index szSub = splitPosition(cur.origin, cur.size, kLargest, numSub, cur.numSub, weights);
 
      // Snap the split to a multiple of align_[kLargest] so the global split position
      // (cur.origin[kLargest] + szSub, with cur.origin already aligned) is a multiple too — the
      // precondition for coarsened() to divide cleanly. Skip if the box is too small to split
      // aligned.
      const Index a = align_[kLargest];
      if (a > 1 && cur.size[kLargest] >= 2 * a) {
        szSub = ((szSub + a / 2) / a) * a;
        if (szSub < a)
          szSub = a;
        if (szSub > cur.size[kLargest] - a)
          szSub = cur.size[kLargest] - a;
      }
 
      StackBlock left = cur;
      left.numSub = numSub;
      left.size[kLargest] = szSub;
      left.treeIndx = 2 * cur.treeIndx + 1;
 
      StackBlock right = cur;
      right.origin[kLargest] += szSub;
      right.size[kLargest] -= szSub;
      right.numSub = cur.numSub - numSub;
      right.treeIndx = 2 * cur.treeIndx + 2;
 
      tree_[cur.treeIndx] = TreeNode{kLargest, cur.origin[kLargest] + szSub};
      stack.push(right);
      stack.push(left);
    } else {
      tree_[cur.treeIndx] = TreeNode{-1, leafIndex++};
      origins_.push_back(cur.origin);
      sizes_.push_back(cur.size);
    }
  }
}
 
template <int Dim>
Index BlockDecomposer<Dim>::splitPosition(const IVec<Dim>& origin, const IVec<Dim>& size,
                                          int kLargest, std::size_t numSub, std::size_t numTotal,
                                          const std::vector<Real>* weights) const {
  const Index n = size[kLargest];
 
  // Unweighted (and the degenerate non-splittable axis): the classic proportional-to-count split.
  // Kept as the exact same expression so the unweighted API is byte-for-byte unchanged.
  if (weights == nullptr || n <= 1) {
    return static_cast<Index>(std::round(static_cast<double>(n) * static_cast<double>(numSub) /
                                         static_cast<double>(numTotal)));
  }
 
  // Weighted: accumulate the weight of each slab (fixed kLargest-coordinate) within the box, then
  // pick the boundary whose cumulative weight is closest to the target fraction of the total.
  std::vector<double> slab(static_cast<std::size_t>(n), 0.0);
  IVec<Dim> bgn = origin, end{};
  for (int i = 0; i < Dim; ++i)
    end[i] = origin[i] + size[i];
  forEachInBox<Dim>(bgn, end, [&](const IVec<Dim>& g) {
    slab[static_cast<std::size_t>(g[kLargest] - origin[kLargest])] += (*weights)[linearGlobal(g)];
  });
 
  double total = 0.0;
  for (double w : slab)
    total += w;
  const double target = total * static_cast<double>(numSub) / static_cast<double>(numTotal);
 
  // Search boundaries in [1, n-1] (non-empty children). Ties resolve to the larger boundary, which
  // mirrors std::round's half-away-from-zero rule so equal weights reproduce the unweighted split.
  double cum = 0.0;
  double bestDist = std::numeric_limits<double>::max();
  Index best = 1;
  for (Index s = 1; s < n; ++s) {
    cum += slab[static_cast<std::size_t>(s - 1)];
    const double dist = std::abs(cum - target);
    if (dist <= bestDist) {
      bestDist = dist;
      best = s;
    }
  }
  return best;
}
 
template <int Dim>
Real weightImbalance(const BlockDecomposer<Dim>& dec, const std::vector<Real>& weights) {
  const std::size_t nb = dec.numBlocks();
  double total = 0.0, hi = 0.0;
  for (std::size_t b = 0; b < nb; ++b) {
    IVec<Dim> bgn = dec.origins()[b], end{};
    for (int k = 0; k < Dim; ++k)
      end[k] = bgn[k] + dec.sizes()[b][k];
    double w = 0.0;
    forEachInBox<Dim>(bgn, end, [&](const IVec<Dim>& g) {
      w += weights[static_cast<std::size_t>(dec.linearGlobal(g))];
    });
    total += w;
    if (b == 0 || w > hi)
      hi = w;
  }
  if (!(total > 0.0) || nb == 0)
    return std::numeric_limits<Real>::infinity();
  return hi / (total / static_cast<double>(nb));
}
 
template <int Dim>
struct AlignedWeightedChoice {
  BlockDecomposer<Dim> dec;  
  int a = 0;                 
  Real imbalance = 0.0;      
};
 
template <std::size_t N, int Dim = static_cast<int>(N)>
AlignedWeightedChoice<Dim> chooseAlignedWeighted(std::size_t numBlocks,
                                                 const std::array<Index, N>& G,
                                                 const std::vector<Real>& weights,
                                                 Real budget = 1.05,
                                                 int aMax = std::numeric_limits<int>::max()) {
  int cap = aMax;
  for (int k = 0; k < Dim; ++k) {
    int tz = 0;
    for (Index g = G[k]; g > 0 && g % 2 == 0; g /= 2)
      ++tz;
    cap = tz < cap ? tz : cap;
  }
  int top = 0;  // the largest a <= cap whose coarse grid still holds numBlocks, >= 2 cells per axis
  while (top < cap) {
    const int a = top + 1;
    bool ok = true;
    double cells = 1.0;
    for (int k = 0; k < Dim; ++k) {
      const Index gk = G[k] >> a;
      if (gk < 2)
        ok = false;
      cells *= static_cast<double>(gk);
    }
    if (!ok || cells < static_cast<double>(numBlocks))
      break;
    top = a;
  }
  for (int a = top; a >= 1; --a) {
    IVec<Dim> align{};
    for (int k = 0; k < Dim; ++k)
      align[k] = Index(1) << a;
    AlignedWeightedChoice<Dim> c;
    c.dec.init(numBlocks, G, weights, align);
    bool empty = false;
    for (const auto& sz : c.dec.sizes())
      for (int k = 0; k < Dim; ++k)
        if (sz[k] <= 0)
          empty = true;
    if (empty)
      continue;
    c.imbalance = weightImbalance(c.dec, weights);
    if (c.imbalance <= budget) {
      c.a = a;
      return c;
    }
  }
  AlignedWeightedChoice<Dim> c;
  c.dec.init(numBlocks, G, weights);
  c.a = 0;
  c.imbalance = weightImbalance(c.dec, weights);
  return c;
}
 
}  // namespace peclet::core::decomp
 
#endif  // PECLET_CORE_DECOMP_BLOCK_DECOMPOSER_HPP