assembly.hpp

Go to the documentation of this file.

// core — device assembly of the FV (pressure) operator (D2), built on the S1 CSR primitive.
//
// AmrPoisson::assembleFv walks `forEachFaceNeighbor` on the host (serial) to produce the weight-CSR
// (invVol / start / nbr / coef=openness·A_f/d_f / bcDiag) that the device MG then applies. For a
// STATIC geometry that host walk runs once; for a DYNAMIC geometry (moving SDF / adapt) it would
// run every step and force a host round-trip. This header reproduces that walk as a device functor
// + a per-cell diagonal kernel, so the FvOp is assembled entirely on the device — no host
// round-trip — feeding the same shared face_csr.hpp apply kernels.
//
// Bit-exactness: the device functor emits faces in the exact same axis/dir/sub-face order as the
// host `forEachFaceNeighbor`, each cell fills its own CSR slice (S1, atomic-free), and every
// coefficient uses the identical double arithmetic. So on the OpenMP backend the device-assembled
// FvOp is bit-for-bit equal to the host `assembleFv` / `buildFaceCsr` (the cross-backend anti-drift
// lock in tests/test_amr_assembly_kokkos). GPU is tolerance-not-bit-exact (FMA), per the
// convention.
//
// KOKKOS_FUNCTION).
#ifndef PECLET_AMR_ASSEMBLY_HPP
#define PECLET_AMR_ASSEMBLY_HPP
 
#include <array>
 
#include "peclet/amr/block_octree_view.hpp"
#include "peclet/amr/common.hpp"
#include "peclet/amr/csr.hpp"
#include "peclet/amr/fv_op.hpp"
#include "peclet/amr/poisson.hpp"
#include "peclet/core/common/view.hpp"
 
namespace peclet::amr {
 
template <int Dim, unsigned Bits = (Dim == 2 ? 32u : (Dim == 3 ? 21u : 16u))>
struct FvFaceEmit {
  using View_t = BlockOctreeView<Dim, Bits>;
  using M = typename View_t::M;
  using Code = typename View_t::Code;
  using Coord = typename View_t::Coord;
 
  View_t ov;                 
  View<const double> alpha;  
  Coord fineExt[Dim] = {};   
  Real h0[Dim] = {};         
  bool hasOpen = false;
  bool periodic = true;
  bool immersedWall = false;
 
  static constexpr int kFaces = 2 * Dim;
  KOKKOS_INLINE_FUNCTION static int faceIndex(int axis, int dir) {
    return 2 * axis + (dir > 0 ? 0 : 1);
  }
  KOKKOS_INLINE_FUNCTION double openness(Index leaf, int axis, int dir) const {
    if (!hasOpen)
      return 1.0;
    return alpha(static_cast<std::size_t>(leaf) * kFaces + faceIndex(axis, dir));
  }
  KOKKOS_INLINE_FUNCTION Real areaOf(Coord s, int axis) const {
    Real a = 1;
    for (int d = 0; d < Dim; ++d)
      if (d != axis)
        a *= static_cast<Real>(s) * h0[d];
    return a;
  }
  KOKKOS_INLINE_FUNCTION Real coeff(Coord si, Coord sj, int axis) const {
    const Real dist = 0.5 * (static_cast<Real>(si) + static_cast<Real>(sj)) * h0[axis];
    return areaOf(si < sj ? si : sj, axis) / dist;
  }
  KOKKOS_INLINE_FUNCTION Coord wrap(long c, int axis) const {
    const long e = static_cast<long>(fineExt[axis]);
    return static_cast<Coord>(((c % e) + e) % e);
  }
  KOKKOS_INLINE_FUNCTION Real cellWidth(Index i, int axis) const {
    return h0[axis] * static_cast<Real>(Index(1) << ov.levels(i));
  }
  KOKKOS_INLINE_FUNCTION Real cellWidth(Index i) const { return cellWidth(i, 0); }
  KOKKOS_INLINE_FUNCTION Real cellVolume(Index i) const {
    const Real f = static_cast<Real>(Index(1) << ov.levels(i));
    Real v = 1;
    for (int d = 0; d < Dim; ++d)
      v *= h0[d] * f;
    return v;
  }
 
  KOKKOS_INLINE_FUNCTION double bcDiag(Index i) const {
    if (periodic && !immersedWall)
      return 0.0;
    std::array<Coord, Dim> lo = M::from_code(ov.codes(i)).decode();
    const Coord si = Coord(Coord(1) << ov.levels(i));
    double s = 0.0;
    for (int axis = 0; axis < Dim; ++axis) {
      const double wall = areaOf(si, axis) / (0.5 * static_cast<Real>(si) * h0[axis]);
      for (int dir = -1; dir <= 1; dir += 2) {
        const long pc = (dir > 0) ? static_cast<long>(lo[axis]) + static_cast<long>(si)
                                  : static_cast<long>(lo[axis]) - 1;
        const bool domainBoundary = !periodic && (pc < 0 || pc >= static_cast<long>(fineExt[axis]));
        if (domainBoundary)
          s += openness(i, axis, dir) * wall;
        else if (immersedWall)
          s += (1.0 - openness(i, axis, dir)) * wall;
      }
    }
    return s;
  }
 
  template <class Sink>
  KOKKOS_INLINE_FUNCTION void operator()(Index i, Sink& sink) const {
    std::array<Coord, Dim> lo = M::from_code(ov.codes(i)).decode();
    const unsigned Li = ov.levels(i);
    const Coord si = Coord(Coord(1) << Li);
    for (int axis = 0; axis < Dim; ++axis)
      for (int dir = -1; dir <= 1; dir += 2) {
        const long pc = (dir > 0) ? static_cast<long>(lo[axis]) + static_cast<long>(si)
                                  : static_cast<long>(lo[axis]) - 1;
        if (!periodic && (pc < 0 || pc >= static_cast<long>(fineExt[axis])))
          continue;
        std::array<Coord, Dim> p = lo;
        p[axis] = wrap(pc, axis);
        const Index j = ov.locate(M::encode(p).code());
        const unsigned Lj = ov.levels(j);
        if (Lj >= Li) {
          sink(j, openness(i, axis, dir) * coeff(si, Coord(Coord(1) << Lj), axis));
        } else {
          const Coord sj = Coord(si >> 1);
          const int nsub = 1 << (Dim - 1);
          for (int k = 0; k < nsub; ++k) {
            std::array<Coord, Dim> q = lo;
            q[axis] = wrap(pc, axis);
            int bit = 0;
            for (int t = 0; t < Dim; ++t) {
              if (t == axis)
                continue;
              const Coord off = ((k >> bit) & 1) ? sj : Coord(0);
              q[t] = wrap(static_cast<long>(lo[t]) + static_cast<long>(off), t);
              ++bit;
            }
            const Index jj = ov.locate(M::encode(q).code());
            sink(jj, openness(jj, axis, -dir) * coeff(si, sj, axis));
          }
        }
      }
  }
};
 
template <int Dim, unsigned Bits>
FvOp assembleFv(const AmrPoisson<Dim, Bits>& ap, const BlockOctreeView<Dim, Bits>& ov) {
  FvFaceEmit<Dim, Bits> emit;
  emit.ov = ov;
  for (int d = 0; d < Dim; ++d)
    emit.h0[d] = ap.h0()[d];
  emit.periodic = ap.periodic();
  emit.immersedWall = ap.immersedWall();
  emit.hasOpen = ap.hasOpenness();
  for (int d = 0; d < Dim; ++d)
    emit.fineExt[d] = ap.fineExt()[d];
  if (ap.hasOpenness())
    emit.alpha = toDevice(ap.opennessRaw(), "amr::alpha");
 
  Csr csr = buildFaceCsr(ov.numLeaves(), emit);
 
  FvOp op;
  op.n = ov.numLeaves();
  op.faceStart = csr.start;
  op.faceNbr = csr.nbr;
  op.faceW = csr.coef;
  op.invVol = View<double>(Kokkos::view_alloc("amr::invVol", Kokkos::WithoutInitializing),
                           static_cast<std::size_t>(op.n));
  op.bcDiag = View<double>(Kokkos::view_alloc("amr::bcDiag", Kokkos::WithoutInitializing),
                           static_cast<std::size_t>(op.n));
  View<double> invVol = op.invVol;
  View<double> bcDiag = op.bcDiag;
  const FvFaceEmit<Dim, Bits> e = emit;
  Kokkos::parallel_for(
      "amr::fv_diag", op.n, KOKKOS_LAMBDA(const Index i) {
        invVol(i) = 1.0 / e.cellVolume(i);
        bcDiag(i) = e.bcDiag(i);
      });
  return op;
}
 
}  // namespace peclet::amr
 
#endif  // PECLET_AMR_ASSEMBLY_HPP