mg_stage.hpp

Go to the documentation of this file.

// amr — a multigrid STAGE: move a level, at its own resolution, onto a new decomposition of that
// level's grid, and continue the ladder there. docs/amr_mg_depth.md §5.6 / §6.5, and
// suite/docs/archive/MG_TELESCOPING_PLAN.md §4.0, whose primitive this is.
//
// The in-place ladder (§6.2) lifts the root while every rank's ORB block stays even in origin and
// size. It stops somewhere — grid-limited on a badly factored brick mesh, decomposition-limited as
// soon as a block turns odd, which under a WEIGHTED (brick-granular, DEM-shared) partition is
// typically the root brick level itself. Below that point the hierarchy does not stop: the level
// MOVES onto a coarser decomposition of its own grid and the ladder continues there. §5.6's three
// instantiations, in the order the trigger tries them:
//
//   * sibling merge — target `BlockDecomposer::agglomerated(d)`, group gather; proportional trees;
//   * repartition   — target a fresh proportional ORB on np_L <= np ranks, moved point-to-point
//                     over the box intersections; weighted trees, where no d > 0 is liftable;
//   * REPLICATED    — target one block on every rank, moved by `Allgatherv`; the bottom on the
//                     last sub-communicator, and the fallback. THIS FILE implements this one.
//
// All three take their target, communicators and movement from core (`peclet::core::decomp`:
// `chooseStageTarget`, `makeStageComm`, `RedistributeTopology`; docs/amr_mg_core_boundary.md). The
// other two are WO4b's `DistributedStage` (distributed_flow_mg.hpp, since its continued ladder is
// a `DistributedFlowMultigrid`). They are why this is an interface rather than a function: what
// differs between them is only the target decomposition and the two movement steps (`moveUp` /
// `moveDown`); the continued ladder, the per-level communicator and the way `vcycle` calls the
// stage are shared. `DistributedFlowMultigrid::vcycle` never sees a message — it calls `apply` at
// the stage point.
#ifndef PECLET_AMR_MG_STAGE_HPP
#define PECLET_AMR_MG_STAGE_HPP
 
#include <algorithm>
#include <array>
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>
 
#include "peclet/amr/block_octree.hpp"
#include "peclet/amr/common.hpp"
#include "peclet/amr/distributed_octree.hpp"
#include "peclet/amr/multigrid.hpp"
#include "peclet/amr/poisson.hpp"
#include "peclet/core/common/mpi.hpp"
#include "peclet/core/common/view.hpp"
#include "peclet/core/decomp/block_decomposer.hpp"
#include "peclet/core/decomp/redistribute_topology.hpp"
#include "peclet/core/decomp/stage_comm.hpp"
#include "peclet/core/decomp/stage_target.hpp"
 
namespace peclet::amr {
 
template <int Dim, unsigned Bits = (Dim == 2 ? 32u : (Dim == 3 ? 21u : 16u))>
class MgStage {
 public:
  using Octree = BlockOctree<Dim, Bits>;                     
  using Decomposition = core::decomp::BlockDecomposer<Dim>;  
 
  virtual ~MgStage() = default;
 
  // ---- what the stage is ---------------------------------------------------------------------
  virtual const char* kind() const = 0;
  virtual const char* diagnosticSuffix() const = 0;
  virtual const Decomposition& targetDecomposition() const = 0;
  virtual MPI_Comm targetComm() const = 0;
  virtual bool active() const = 0;
  virtual const Octree& targetOctree() const = 0;
 
  // ---- the continued ladder (its level 0 IS the moved level) ---------------------------------
  virtual std::size_t numLevels() const = 0;
  virtual Index numLeaves(std::size_t L) const = 0;
  virtual std::string bottomName() const = 0;
  virtual void setRemoveMean(bool on) = 0;
  virtual void setBottom(typename Multigrid<Dim, Bits>::Bottom b) = 0;
  virtual View<double> targetX() = 0;
 
  // ---- the movement: the ONLY part the three instantiations differ in ------------------------
  virtual void moveUp(View<const double> src, Index nSrc) = 0;
  virtual void cycle(int pre, int post, int bottom, double omega) = 0;
  virtual void moveDown(View<double> dst, Index nDst) = 0;
 
  void apply(View<const double> residual, View<double> correction, Index n, int pre, int post,
             int bottom, double omega) {
    moveUp(residual, n);
    if (active())
      cycle(pre, post, bottom, omega);
    moveDown(correction, n);
  }
};
 
template <int Dim, unsigned Bits = (Dim == 2 ? 32u : (Dim == 3 ? 21u : 16u))>
class ReplicatedTailStage final : public MgStage<Dim, Bits> {
 public:
  using Base = MgStage<Dim, Bits>;                     
  using Octree = typename Base::Octree;                
  using Decomposition = typename Base::Decomposition;  
  using DO = DistributedOctree<Dim, Bits>;             
  using Poisson = AmrPoisson<Dim, Bits>;               
 
  void build(const DO& d, const Vec<Dim>& h0, const Poisson& ap, Index n, Index nExt,
             Index bottomExtent, MPI_Comm comm) {
    namespace cd = core::decomp;
    G_ = d.globalRootSize();
    nt_ = 1;
    for (int a = 0; a < Dim; ++a)
      nt_ *= G_[a];
    for (int a = 0; a < Dim; ++a)
      if (!d.periodic()[a])
        throw std::runtime_error(
            "amr::ReplicatedTailStage: periodic-only (as the whole distributed pressure path is); "
            "a non-periodic bottom must be gated the way flow gates its anchored path");
    Index blockCells = 1;
    for (int a = 0; a < Dim; ++a)
      blockCells *= d.blockBrick()[a];
    if (n != blockCells || n != d.local().numLeaves())
      throw std::runtime_error(
          "amr::ReplicatedTailStage: the moved level does not tile its block with root cells (the "
          "lift is not nested)");
    target_.kind = cd::StageKind::Replicated;  // ONE block, the whole level, on every rank
    target_.dec.init(1, G_);
    target_.ownerOf = {0};
    target_.groupOf.clear();
    oct_.init(G_, d.lmax(), IVec<Dim>{});
    const Index W = d.rootSpan();
 
    // The two index maps core's topology calls once per cell: a global cell of the moved level's
    // grid → this rank's leaf holding it, and → the tail octree's row. Every leaf of the moved
    // level is a root cell of the lifted root grid, which is what nesting buys.
    auto srcIndex = [&](const IVec<Dim>& g) -> Index {
      const Index i = d.findGlobalRoot(g);
      if (i < 0 || d.local().level(i) != d.lmax())
        throw std::runtime_error(
            "amr::ReplicatedTailStage: a cell of this rank's block is not one of its root leaves "
            "(the lift is not nested)");
      return i;
    };
    auto dstIndex = [&](const IVec<Dim>& g) -> Index {
      std::array<typename Octree::Coord, Dim> f{};
      for (int a = 0; a < Dim; ++a)
        f[a] = static_cast<typename Octree::Coord>(g[a] * W);
      return oct_.find(f);
    };
    comm_ = cd::makeStageComm<Dim>(comm, target_);
    topo_.build(d.decomposition(), target_, comm_, srcIndex, dstIndex);
 
    // The moved level's α rows, one field per face, gathered into target leaf order.
    const int F = 2 * Dim;
    const std::vector<double>& aloc = ap.opennessRaw();
    std::vector<std::vector<double>> aCol(static_cast<std::size_t>(F),
                                          std::vector<double>(static_cast<std::size_t>(n), 1.0));
    if (!aloc.empty())
      for (Index i = 0; i < n; ++i)
        for (int f = 0; f < F; ++f)
          aCol[static_cast<std::size_t>(f)][static_cast<std::size_t>(i)] =
              aloc[static_cast<std::size_t>(i) * F + static_cast<std::size_t>(f)];
    std::vector<std::vector<double>> tCol(static_cast<std::size_t>(F),
                                          std::vector<double>(static_cast<std::size_t>(nt_), 1.0));
    std::vector<const double*> aSrc(static_cast<std::size_t>(F));
    std::vector<double*> aDst(static_cast<std::size_t>(F));
    for (int f = 0; f < F; ++f) {
      aSrc[static_cast<std::size_t>(f)] = aCol[static_cast<std::size_t>(f)].data();
      aDst[static_cast<std::size_t>(f)] = tCol[static_cast<std::size_t>(f)].data();
    }
    topo_.forward(aSrc, aDst);
    std::vector<double> aTail(static_cast<std::size_t>(nt_) * F);
    for (Index r = 0; r < nt_; ++r)
      for (int f = 0; f < F; ++f)
        aTail[static_cast<std::size_t>(r) * F + static_cast<std::size_t>(f)] =
            tCol[static_cast<std::size_t>(f)][static_cast<std::size_t>(r)];
 
    mg_ = std::make_unique<Multigrid<Dim, Bits>>();
    mg_->buildRaw(oct_, h0, std::move(aTail), /*periodic=*/true, /*liftRoot=*/true, bottomExtent);
 
    (void)nExt;
    // Sized by the level's OWN cell count, and staged through device scratch, so the caller may
    // hand in a longer view (the distributed levels carry a ghost tail) without a subview.
    dSrc_ = View<double>("stage_dsrc", static_cast<std::size_t>(n));
    dDst_ = View<double>("stage_ddst", static_cast<std::size_t>(n));
    srcMirror_ = Kokkos::View<double*, Kokkos::HostSpace>("stage_src", static_cast<std::size_t>(n));
    dstMirror_ = Kokkos::View<double*, Kokkos::HostSpace>("stage_dst", static_cast<std::size_t>(n));
    tbMirror_ = Kokkos::View<double*, Kokkos::HostSpace>("stage_tb", static_cast<std::size_t>(nt_));
    txMirror_ = Kokkos::View<double*, Kokkos::HostSpace>("stage_tx", static_cast<std::size_t>(nt_));
  }
 
  const char* kind() const override { return "replicated"; }
  const char* diagnosticSuffix() const override { return "+tail"; }
  const Decomposition& targetDecomposition() const override { return target_.dec; }
  MPI_Comm targetComm() const override { return comm_.sub; }
  bool active() const override { return true; }  // replicated: every rank owns the whole level
  const Octree& targetOctree() const override { return oct_; }
 
  std::size_t numLevels() const override { return mg_->numLevels(); }
  Index numLeaves(std::size_t L) const override { return mg_->numLeaves(L); }
  std::string bottomName() const override { return mg_->bottomName(); }
  void setRemoveMean(bool on) override { mg_->setRemoveMean(on); }
  void setBottom(typename Multigrid<Dim, Bits>::Bottom b) override { mg_->setBottom(b); }
  View<double> targetX() override { return mg_->x(0); }
 
  Multigrid<Dim, Bits>& multigrid() { return *mg_; }
 
  void moveUp(View<const double> src, Index nSrc) override {
    auto ds = dSrc_;
    Kokkos::parallel_for("amr::stage_pack", nSrc, KOKKOS_LAMBDA(const Index i) { ds(i) = src(i); });
    Kokkos::deep_copy(srcMirror_, dSrc_);
    topo_.forward({srcMirror_.data()}, {tbMirror_.data()});
    Kokkos::deep_copy(mg_->b(0), tbMirror_);
    Kokkos::deep_copy(mg_->x(0), 0.0);
  }
 
  void cycle(int pre, int post, int bottom, double omega) override {
    mg_->vcycle(pre, post, bottom, omega);
  }
 
  void moveDown(View<double> dst, Index nDst) override {
    Kokkos::deep_copy(txMirror_, mg_->x(0));
    topo_.backward({txMirror_.data()}, {dstMirror_.data()});  // a local pick: no communication
    Kokkos::deep_copy(dDst_, dstMirror_);
    auto dd = dDst_;
    Kokkos::parallel_for(
        "amr::stage_scatter", nDst, KOKKOS_LAMBDA(const Index i) { dst(i) = dd(i); });
  }
 
 private:
  IVec<Dim> G_{};
  Index nt_ = 0;
  core::decomp::StageTarget<Dim> target_;
  core::decomp::StageComm comm_;  // owns its communicators; must outlive topo_'s use of them
  core::decomp::RedistributeTopology<Dim, double> topo_;
  Octree oct_;
  std::unique_ptr<Multigrid<Dim, Bits>> mg_;
  View<double> dSrc_, dDst_;
  Kokkos::View<double*, Kokkos::HostSpace> srcMirror_, dstMirror_, tbMirror_, txMirror_;
};
 
}  // namespace peclet::amr
 
#endif  // PECLET_AMR_MG_STAGE_HPP