momentum_assembly.hpp

Go to the documentation of this file.

// core — device assembly of the cut-cell MOMENTUM operator (D3), on the S1 CSR primitive.
//
// AmrCutCell::build (host) samples the SDF per cell and runs buildCutStencil to produce the
// ξ-overlay stencil (AC_/off_/cut_/rscale_); assembleOperator (host) then merges that stencil, the
// regular-fluid C/F-aware ∇², and the implicit-FOU advection into one diag + face-CSR via a
// std::vector<vector<pair>>. For a moving boundary that whole walk reruns every step on the host
// and round-trips the operator.
//
// This header moves the two device-portable halves onto the device:
//   * the per-cell ξ stencil rebuild — buildCutStencil is pure math (now MORTON_HD), fed the cell +
//   6
//     neighbour SDF samples (staged like D2 stages openness; the SDF *sampling* itself, like
//     openness sampling, stays host until a device SDF exists); and
//   * the operator merge — the three per-cell branches (solid identity / ξ-overlay / regular ∇²·μ)
//   plus
//     the advection fold, emitted as a face-CSR through buildFaceCsr (S1) + a per-cell diagonal.
// The regular-fluid faces reuse the D2 FvFaceEmit traversal (α=1) via a scaling sink adapter, so
// the 2:1 enumeration and coeff are shared, not re-derived.
//
// Bit-exactness: emit order, branch logic and arithmetic mirror assembleOperator exactly, each cell
// fills its own CSR slice (S1, atomic-free), so on OpenMP the device MomentumOp == host
// assembleOperator bit-for-bit (test_amr_momentum). GPU is tolerance-not-bit-exact (FMA), per the
// convention.
//
#ifndef PECLET_AMR_MOMENTUM_ASSEMBLY_HPP
#define PECLET_AMR_MOMENTUM_ASSEMBLY_HPP
 
#include "peclet/amr/assembly.hpp"  // FvFaceEmit (the α=1 ∇² geometry traversal)
#include "peclet/amr/block_octree_view.hpp"
#include "peclet/amr/common.hpp"
#include "peclet/amr/csr.hpp"
#include "peclet/amr/cut_cell.hpp"
#include "peclet/amr/momentum.hpp"
#include "peclet/core/common/view.hpp"
 
namespace peclet::amr {
 
template <class Sink>
struct MomScaleSink {
  Sink* sink;
  double factor;
  KOKKOS_INLINE_FUNCTION void operator()(Index j, double c) const { (*sink)(j, factor* c); }
};
 
struct MomSumSink {
  double sum = 0.0;
  KOKKOS_INLINE_FUNCTION void operator()(Index /*j*/, double c) { sum += c; }
};
 
template <unsigned Bits = 21u>
struct MomFaceEmit {
  FvFaceEmit<3, Bits> geom;  
  View<const char> fluid, cut;
  View<const double> AC, off, rscale;  
  View<const Index> nbr6;              
  View<const double> advDiag, advCoef;
  View<const Index> advStart, advNbr;
  double idiag = 0.0, mu = 1.0;
  bool hasAdv = false, scaleAdv = false;
 
  KOKKOS_INLINE_FUNCTION double cellVolume(Index i) const { return geom.cellVolume(i); }
 
  template <class Sink>
  KOKKOS_INLINE_FUNCTION void operator()(Index i, Sink& sink) const {
    if (!fluid(i))
      return;  // solid: identity row, no off-diagonals
    const std::size_t s = static_cast<std::size_t>(i);
    if (cut(i)) {
      for (int k = 0; k < 6; ++k) {
        const double a = off(s * 6 + k);
        if (a == 0.0)
          continue;
        const Index j = nbr6(s * 6 + k);
        if (j >= 0)
          sink(j, a);
      }
    } else {
      const double invV = 1.0 / geom.cellVolume(i);
      MomScaleSink<Sink> adapter{&sink, -mu * invV};
      geom(i, adapter);  // emits (j, -mu·invV·(a·c)) per C/F face (a=1)
    }
    if (hasAdv) {
      const double as = scaleAdv ? rscale(i) : 1.0;
      for (Index p = advStart(i); p < advStart(i + 1); ++p)
        sink(advNbr(p), as * advCoef(p));
    }
  }
 
  KOKKOS_INLINE_FUNCTION double diag(Index i) const {
    if (!fluid(i))
      return 1.0;
    double d;
    if (cut(i)) {
      d = AC(i);
    } else {
      const double invV = 1.0 / geom.cellVolume(i);
      MomSumSink ss;
      geom(i, ss);
      d = idiag + mu * invV * ss.sum;
    }
    if (hasAdv) {
      const double as = scaleAdv ? rscale(i) : 1.0;
      d += as * advDiag(i);
    }
    return d;
  }
};
 
template <unsigned Bits>
void rebuildCutStencil(Index n, const double betaIn[3], double idiag, const View<double>& sdfC,
                       const View<Index>& nbr6, const View<char>& fluid,
                       const View<std::uint8_t>& levels, View<double> AC, View<double> off,
                       View<char> cut, View<double> rscale) {
  // Phase 3: capture the three axis coefficients by value into the device lambda. `betaIn` is
  // mu/h_finest^2, so it must be divided by 4^level per leaf — the device twin of the host
  // AmrCutCell::build Pass 2 (see the long note there; bit-for-bit the old expressions at level 0).
  const double b0 = betaIn[0], b1 = betaIn[1], b2 = betaIn[2], id0 = idiag;
  Kokkos::parallel_for(
      "amr::cut_stencil", n, KOKKOS_LAMBDA(const Index i) {
        const double f = static_cast<double>(Index(1) << levels(i));  // h_leaf / h_finest
        const double inv = 1.0 / (f * f);
        const double beta[3] = {b0 * inv, b1 * inv, b2 * inv};
        const double AC0 = id0 + ((2.0 * beta[0] + 2.0 * beta[1]) + 2.0 * beta[2]);
        const std::size_t s = static_cast<std::size_t>(i);
        if (!fluid(i)) {  // solid: identity row
          AC(i) = 1.0;
          for (int k = 0; k < 6; ++k)
            off(s * 6 + k) = 0.0;
          cut(i) = 0;
          rscale(i) = 1.0;
          return;
        }
        double sdf_n[6];
        bool anyGhost = false;
        for (int k = 0; k < 6; ++k) {
          const Index j = nbr6(s * 6 + k);
          sdf_n[k] = (j >= 0) ? sdfC(j) : -1.0;  // missing neighbour => solid
          if (sdf_n[k] < 0.0)
            anyGhost = true;
        }
        cut(i) = anyGhost ? 1 : 0;
        double AClocal = AC0, offl[6];
        for (int k = 0; k < 6; ++k)
          offl[k] = -beta[k / 2];
        double rsl = 1.0, inhoml = 0.0;
        if (anyGhost)
          AmrCutCell<Bits>::buildCutStencil(sdfC(i), sdf_n, beta, AC0, AClocal, offl, rsl, inhoml);
        AC(i) = AClocal;
        for (int k = 0; k < 6; ++k)
          off(s * 6 + k) = offl[k];
        rscale(i) = rsl;
      });
}
 
template <unsigned Bits>
MomentumOp assembleMomentum(const AmrCutCell<Bits>& ccop, const BlockOctreeView<3, Bits>& ov,
                            bool scaleAdvByRscale = false) {
  const Index n = ov.numLeaves();
  // Phase 3: beta per axis, and the diagonal as one parenthesised sum (Rule B).
  const Vec<3> betaV = ccop.beta();
  const double beta[3] = {betaV[0], betaV[1], betaV[2]};
 
  // Stage build inputs + rebuild the ξ stencil on device.
  View<double> sdfC = toDevice(ccop.sdfCRaw(), "mom::sdfC");
  View<Index> nbr6 = toDevice(ccop.nbRaw(), "mom::nb");
  View<char> fluid = toDevice(ccop.fluidRaw(), "mom::fluid");
  View<double> AC(Kokkos::view_alloc("mom::AC", Kokkos::WithoutInitializing),
                  static_cast<std::size_t>(n));
  View<double> off(Kokkos::view_alloc("mom::off", Kokkos::WithoutInitializing),
                   static_cast<std::size_t>(n) * 6);
  View<char> cut(Kokkos::view_alloc("mom::cut", Kokkos::WithoutInitializing),
                 static_cast<std::size_t>(n));
  View<double> rscale(Kokkos::view_alloc("mom::rscale", Kokkos::WithoutInitializing),
                      static_cast<std::size_t>(n));
  rebuildCutStencil<Bits>(n, beta, ccop.idiag(), sdfC, nbr6, fluid, ov.levels, AC, off, cut,
                          rscale);
 
  // The α=1 ∇² geometry for regular fluid cells (reuses the D2 traversal with openness off).
  FvFaceEmit<3, Bits> geom;
  geom.ov = ov;
  for (int d = 0; d < 3; ++d)
    geom.h0[d] = ccop.lap().h0()[d];
  geom.periodic = ccop.lap().periodic();
  geom.hasOpen = false;
  for (int d = 0; d < 3; ++d)
    geom.fineExt[d] = ccop.lap().fineExt()[d];
 
  MomFaceEmit<Bits> emit;
  emit.geom = geom;
  emit.fluid = fluid;
  emit.cut = cut;
  emit.AC = AC;
  emit.off = off;
  emit.rscale = rscale;
  emit.nbr6 = nbr6;
  emit.idiag = ccop.idiag();
  emit.mu = ccop.mu();
  emit.hasAdv = ccop.hasAdv();
  emit.scaleAdv = scaleAdvByRscale;
  if (ccop.hasAdv()) {
    emit.advDiag = toDevice(ccop.advDiagRaw(), "mom::advDiag");
    emit.advCoef = toDevice(ccop.advCoefRaw(), "mom::advCoef");
    emit.advStart = toDevice(ccop.advStartRaw(), "mom::advStart");
    emit.advNbr = toDevice(ccop.advNbrRaw(), "mom::advNbr");
  }
 
  Csr csr = buildFaceCsr(n, emit);
 
  MomentumOp op;
  op.n = n;
  op.faceStart = csr.start;
  op.faceNbr = csr.nbr;
  op.faceCoef = csr.coef;
  op.hasAdv = false;  // advection folded into the single CSR (as host assembleOperator does)
  op.diag = View<double>(Kokkos::view_alloc("mom::diag", Kokkos::WithoutInitializing),
                         static_cast<std::size_t>(n));
  View<double> diagV = op.diag;
  const MomFaceEmit<Bits> e = emit;
  Kokkos::parallel_for("amr::mom_diag", n, KOKKOS_LAMBDA(const Index i) { diagV(i) = e.diag(i); });
  return op;
}
 
}  // namespace peclet::amr
 
#endif  // PECLET_AMR_MOMENTUM_ASSEMBLY_HPP