PERF-2 P1.1: the multigrid-PCG's prepared operator (hierarchy, f64 fine level, active cells, components) cached on the embedded solver and reused while the operator is bit-identical (an exact key over the coefficient bit patterns, the mask and the hierarchy parameters); solve_multigrid_pcg_cached gives the uncached solve's answer bit for bit (pin poisson_cache_exact.rs: five right-hand sides on one operator, a one-coefficient miss, both precisions); the CG driver split into Prepared::build + run_pcg
CI / Build (macos-latest) (push) Waiting to run
CI / CI Success (push) Blocked by required conditions
CI / Test (macos-latest) (push) Blocked by required conditions
CI / Test (ubuntu-latest) (push) Blocked by required conditions
CI / Python Bindings (maturin) (macos-latest) (push) Blocked by required conditions
CI / Python Bindings (maturin) (ubuntu-latest) (push) Blocked by required conditions
CI / WASM Build + Size Check (push) Blocked by required conditions
CI / Distributed Training Tests (push) Blocked by required conditions
CI / Build CPU-Only (Explicit) (push) Failing after 6s
Documentation / Build API Documentation (push) Failing after 6s
Documentation / Build User Guide (push) Successful in 6s
CI / Format Check (push) Failing after 16s
CI / Build (ubuntu-latest) (push) Failing after 2m13s
CI / Clippy Check (push) Failing after 2m41s
Performance Benchmarks / Run Benchmarks (push) Successful in 3m7s

Co-Authored-By: Claude Fable 5.1 <[email protected]>
Claude-Session: https://claude.ai/code/session_01YJPeT6WA2e7YvAnS875AHL
This commit is contained in:
Omar Sobh
2026-09-15 23:30:53 -05:00
co-authored by Claude Fable 5.1
parent 01830e5c8d
commit 79c18def32
5 changed files with 294 additions and 15 deletions
@@ -171,6 +171,9 @@ pub struct EmbeddedPisoSolver {
/// PERF-2 P0: Poisson `(setup ns, iterate ns, calls, CG iterations)`
/// summed over every multigrid-PCG solve (`docs/perf2_campaign.md`).
poisson_profile: std::cell::Cell<(u64, u64, u64, u64)>,
/// PERF-2 P1.1: the prepared multigrid operator, reused across rounds
/// and steps while the operator's coefficients and mask are unchanged.
pcg_cache: std::cell::RefCell<super::poisson::PcgCache>,
moving: bool,
/// Mask hysteresis band in multiples of the min cell size (0 = off).
mask_hysteresis: f64,
@@ -198,6 +201,7 @@ impl EmbeddedPisoSolver {
fringe_correction: Vec::new(),
inner_stop_factor: 1e-2,
poisson_profile: std::cell::Cell::new((0, 0, 0, 0)),
pcg_cache: std::cell::RefCell::new(super::poisson::PcgCache::default()),
moving: false,
mask_hysteresis: 0.0,
time: 0.0,
@@ -7,7 +7,7 @@ use super::EmbeddedPisoSolver;
use crate::CfdResult;
use crate::solvers::incompressible::ale::SideBoundary;
use crate::solvers::incompressible::poisson::{
MultigridParameters, PoissonProblem, PoissonSolverKind, solve_multigrid_pcg,
MultigridParameters, PoissonProblem, PoissonSolverKind, solve_multigrid_pcg_cached,
};
use crate::solvers::incompressible::{EmbeddedMask, FlowField};
@@ -294,7 +294,7 @@ impl EmbeddedPisoSolver {
}
let anchor_cell =
(!any_outlet && !self.has_fringe()).then_some(anchor.0 * nx + anchor.1);
let solution = solve_multigrid_pcg(
let solution = solve_multigrid_pcg_cached(
&problem,
&mut p_prime,
&MultigridParameters {
@@ -303,6 +303,7 @@ impl EmbeddedPisoSolver {
},
inner_stop,
anchor_cell,
&mut self.pcg_cache.borrow_mut(),
);
let (s0, i0, c0, k0) = self.poisson_profile.get();
self.poisson_profile.set((
@@ -65,7 +65,8 @@ pub use piso::{PisoParameters, PisoResult, PisoSolver};
#[cfg(feature = "cuda")]
pub use piso_gpu::PisoGpuSolver;
pub use poisson::{
MgPrecision, MultigridParameters, PoissonProblem, PoissonSolution, PoissonSolverKind,
MgPrecision, MultigridParameters, PcgCache, PoissonProblem, PoissonSolution, PoissonSolverKind,
solve_multigrid_pcg, solve_multigrid_pcg_cached,
};
pub use polygon_sdf::PolygonSdf;
pub use simple::{ConvectionScheme, SimpleParameters, SimpleResult, SimpleSolver};
@@ -751,12 +751,171 @@ pub fn solve_multigrid_pcg(
/// level used for the CG's own products and true residual is a separate
/// `f64` level, so the precision of the preconditioner never enters the
/// stopping rule or the reported residual.
/// The operator part of a [`PoissonProblem`] (everything but the
/// right-hand side) plus the hierarchy parameters, kept to decide whether
/// a prepared solver can be reused (PERF-2 P1.1, `docs/perf2_campaign.md`).
/// The comparison is exact (bit patterns), so a reuse changes nothing.
struct OperatorKey {
nx: usize,
ny: usize,
active: Vec<bool>,
coefficients: Vec<u64>,
smoother_sweeps: usize,
coarsest_cells: usize,
}
impl OperatorKey {
fn of(problem: &PoissonProblem, params: &MultigridParameters) -> Self {
let coefficients = problem
.ae
.iter()
.chain(&problem.aw)
.chain(&problem.an)
.chain(&problem.as_)
.chain(&problem.extra_diag)
.map(|v| v.to_bits())
.collect();
Self {
nx: problem.nx,
ny: problem.ny,
active: problem.active.clone(),
coefficients,
smoother_sweeps: params.smoother_sweeps,
coarsest_cells: params.coarsest_cells,
}
}
fn matches(&self, problem: &PoissonProblem, params: &MultigridParameters) -> bool {
self.nx == problem.nx
&& self.ny == problem.ny
&& self.smoother_sweeps == params.smoother_sweeps
&& self.coarsest_cells == params.coarsest_cells
&& self.active == problem.active
&& self.coefficients.iter().copied().eq(problem
.ae
.iter()
.chain(&problem.aw)
.chain(&problem.an)
.chain(&problem.as_)
.chain(&problem.extra_diag)
.map(|v| v.to_bits()))
}
}
/// Everything the CG driver derives from the OPERATOR: the hierarchy, the
/// `f64` fine level, the active cells and the connected components. A
/// deterministic function of the operator; the work vectors inside the
/// hierarchy are re-initialised on the active set at every use, so a
/// prepared solver reused for another right-hand side gives the same
/// answer as a fresh one, bit for bit.
struct Prepared<T: MgScalar> {
key: OperatorKey,
hier: Hierarchy<T>,
fine: Level<f64>,
cells: Vec<usize>,
components: Components,
}
impl<T: MgScalar> Prepared<T> {
fn build(problem: &PoissonProblem, params: &MultigridParameters) -> Self {
let hier = Hierarchy::<T>::build(problem, params);
let fine = Level::<f64>::new(problem.clone());
let cells: Vec<usize> = fine.cells.clone();
let components = Components::find(problem, &cells);
Self {
key: OperatorKey::of(problem, params),
hier,
fine,
cells,
components,
}
}
}
/// A reusable prepared solver per V-cycle precision (PERF-2 P1.1): the
/// operator's hierarchy is rebuilt only when the operator changes.
#[derive(Default)]
pub struct PcgCache {
f64: Option<Prepared<f64>>,
f32: Option<Prepared<f32>>,
}
impl PcgCache {
/// Number of prepared operators held (0, 1 or 2).
pub fn len(&self) -> usize {
usize::from(self.f64.is_some()) + usize::from(self.f32.is_some())
}
/// Whether nothing is cached yet.
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
/// [`solve_multigrid_pcg`] with the operator's hierarchy taken from
/// `cache` when the operator (coefficients, active mask, hierarchy
/// parameters) is bit-identical to the cached one, rebuilt into it
/// otherwise. The answer is that of the uncached solve, bit for bit.
pub fn solve_multigrid_pcg_cached(
problem: &PoissonProblem,
p: &mut [f64],
params: &MultigridParameters,
tolerance: f64,
anchor: Option<usize>,
cache: &mut PcgCache,
) -> PoissonSolution {
match params.precision {
MgPrecision::F64 => {
solve_cached_with::<f64>(problem, p, params, tolerance, anchor, &mut cache.f64)
}
MgPrecision::F32 => {
solve_cached_with::<f32>(problem, p, params, tolerance, anchor, &mut cache.f32)
}
}
}
fn solve_cached_with<T: MgScalar>(
problem: &PoissonProblem,
p: &mut [f64],
params: &MultigridParameters,
tolerance: f64,
anchor: Option<usize>,
slot: &mut Option<Prepared<T>>,
) -> PoissonSolution {
let t_entry = std::time::Instant::now();
let hit = slot
.as_ref()
.is_some_and(|prep| prep.key.matches(problem, params));
if !hit {
*slot = Some(Prepared::<T>::build(problem, params));
}
let setup_ns = t_entry.elapsed().as_nanos() as u64;
let prep = slot.as_mut().expect("prepared");
run_pcg(prep, problem, p, params, tolerance, anchor, setup_ns)
}
fn solve_pcg_with<T: MgScalar>(
problem: &PoissonProblem,
p: &mut [f64],
params: &MultigridParameters,
tolerance: f64,
anchor: Option<usize>,
) -> PoissonSolution {
let t_entry = std::time::Instant::now();
let mut prep = Prepared::<T>::build(problem, params);
let setup_ns = t_entry.elapsed().as_nanos() as u64;
run_pcg(&mut prep, problem, p, params, tolerance, anchor, setup_ns)
}
/// The CG loop on a prepared operator (see [`Prepared`]).
fn run_pcg<T: MgScalar>(
prep: &mut Prepared<T>,
problem: &PoissonProblem,
p: &mut [f64],
params: &MultigridParameters,
tolerance: f64,
anchor: Option<usize>,
setup_ns: u64,
) -> PoissonSolution {
let n = problem.nx * problem.ny;
assert_eq!(p.len(), n, "p must have nx*ny entries");
@@ -765,18 +924,17 @@ fn solve_pcg_with<T: MgScalar>(
"invalid PoissonProblem: {:?}",
problem.validate()
);
let t_entry = std::time::Instant::now();
let mut hier = Hierarchy::<T>::build(problem, params);
let fine = Level::<f64>::new(problem.clone());
let cells: Vec<usize> = fine.cells.clone();
let hier = &mut prep.hier;
let fine = &prep.fine;
let cells: &[usize] = &prep.cells;
let components = &prep.components;
let active_n = cells.len();
if active_n == 0 {
return PoissonSolution {
iterations: 0,
residual: 0.0,
converged: true,
setup_ns: t_entry.elapsed().as_nanos() as u64,
setup_ns,
iterate_ns: 0,
};
}
@@ -791,7 +949,6 @@ fn solve_pcg_with<T: MgScalar>(
// even a well-posed Dirichlet component (found in review, 1e-8
// relative was enough). So the mean is projected per singular
// component, and the exit shift is applied per singular component.
let components = Components::find(problem, &cells);
let singular_any = components.singular.iter().any(|&s| s);
let project_mean = |v: &mut [f64]| {
@@ -810,7 +967,7 @@ fn solve_pcg_with<T: MgScalar>(
// Right-hand side (per-component mean projected out where singular).
let mut b = vec![0.0; n];
for &idx in &cells {
for &idx in cells {
b[idx] = problem.rhs[idx];
}
project_mean(&mut b);
@@ -825,7 +982,6 @@ fn solve_pcg_with<T: MgScalar>(
|p: &[f64], r: &mut [f64], fine: &Level<f64>| -> f64 { fine.residual(&b, p, r) };
let anchor = anchor.filter(|&a| a < n && fine.active[a]);
let setup_ns = t_entry.elapsed().as_nanos() as u64;
let t_iter = std::time::Instant::now();
let finish = |p: &mut [f64], iterations: usize, residual: f64| {
// Level of each singular component: the anchor's component is
@@ -862,7 +1018,7 @@ fn solve_pcg_with<T: MgScalar>(
if singular {
project_mean(&mut z);
}
for &idx in &cells {
for &idx in cells {
d[idx] = z[idx];
}
let mut rz = dot(&r, &z);
@@ -884,7 +1040,7 @@ fn solve_pcg_with<T: MgScalar>(
return finish(p, iterations, res);
}
let alpha = rz / dq;
for &idx in &cells {
for &idx in cells {
p[idx] += alpha * d[idx];
r[idx] -= alpha * q[idx];
}
@@ -904,7 +1060,7 @@ fn solve_pcg_with<T: MgScalar>(
let rz_new = dot(&r, &z);
let beta = rz_new / rz;
rz = rz_new;
for &idx in &cells {
for &idx in cells {
d[idx] = z[idx] + beta * d[idx];
}
}
@@ -0,0 +1,117 @@
//! PERF-2 P1.1 (`docs/perf2_campaign.md`): the cached multigrid-PCG (the
//! operator's hierarchy, fine level and components kept across solves)
//! must give the uncached solve's answer BIT FOR BIT — a masked,
//! outlet-anchored problem like the overset background's, five right-hand
//! sides in a row on one operator, then a one-coefficient change that
//! must miss the cache and still match.
use rtx_cfd::solvers::incompressible::{
MgPrecision, MultigridParameters, PcgCache, PoissonProblem, solve_multigrid_pcg,
solve_multigrid_pcg_cached,
};
fn problem(nx: usize, ny: usize, seed: u64) -> PoissonProblem {
let mut p = PoissonProblem::new(nx, ny);
let (dx, dy, dt) = (1.0 / nx as f64, 0.41 / ny as f64, 1e-3);
let (ae, an) = (dt * dy / dx, dt * dx / dy);
let hole = |i: usize, j: usize| {
let (x, y) = ((i as f64 + 0.5) * dx, (j as f64 + 0.5) * dy);
(x - 0.2).powi(2) + (y - 0.2).powi(2) < 0.05 * 0.05
};
for j in 0..ny {
for i in 0..nx {
let idx = j * nx + i;
if hole(i, j) {
p.active[idx] = false;
continue;
}
if i + 1 < nx && !hole(i + 1, j) {
p.ae[idx] = ae;
}
if i > 0 && !hole(i - 1, j) {
p.aw[idx] = ae;
}
if j + 1 < ny && !hole(i, j + 1) {
p.an[idx] = an;
}
if j > 0 && !hole(i, j - 1) {
p.as_[idx] = an;
}
if i + 1 == nx {
p.extra_diag[idx] = 2.0 * ae; // the outlet
}
}
}
let mut state = seed | 1;
for idx in 0..nx * ny {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
p.rhs[idx] = if p.active[idx] {
1e-6 * ((state >> 11) as f64 / (1u64 << 53) as f64 - 0.5)
} else {
0.0
};
}
p
}
fn same(a: &[f64], b: &[f64]) -> bool {
a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
}
#[test]
fn cached_pcg_reproduces_the_uncached_solve_bit_for_bit() {
let (nx, ny) = (96, 40);
for precision in [MgPrecision::F64, MgPrecision::F32] {
let params = MultigridParameters {
precision,
..MultigridParameters::default()
};
let mut cache = PcgCache::default();
let base = problem(nx, ny, 7);
for k in 0..5u64 {
// Same operator, a new right-hand side each time.
let mut prob = base.clone();
prob.rhs = problem(nx, ny, 11 + k).rhs;
let (mut p0, mut p1) = (vec![0.0; nx * ny], vec![0.0; nx * ny]);
let s0 = solve_multigrid_pcg(&prob, &mut p0, &params, 1e-12, None);
let s1 = solve_multigrid_pcg_cached(&prob, &mut p1, &params, 1e-12, None, &mut cache);
assert_eq!(cache.len(), 1);
assert!(same(&p0, &p1), "{precision:?} rhs {k}: solutions differ");
assert_eq!(s0.iterations, s1.iterations);
assert_eq!(s0.residual.to_bits(), s1.residual.to_bits());
if k > 0 {
assert!(
s1.setup_ns < s0.setup_ns / 4,
"{precision:?} rhs {k}: cache hit should skip the setup ({} vs {} ns)",
s1.setup_ns,
s0.setup_ns
);
}
println!(
" {precision:?} rhs {k}: {} iterations, residual {:.3e}, setup {} vs {} ns",
s0.iterations, s0.residual, s0.setup_ns, s1.setup_ns
);
}
// One coefficient changes: a miss, still exact.
let mut changed = base.clone();
let idx = (ny / 2) * nx + nx / 3;
changed.ae[idx] *= 1.5;
changed.aw[idx + 1] *= 1.5;
let (mut p0, mut p1) = (vec![0.0; nx * ny], vec![0.0; nx * ny]);
let s0 = solve_multigrid_pcg(&changed, &mut p0, &params, 1e-12, None);
let s1 = solve_multigrid_pcg_cached(&changed, &mut p1, &params, 1e-12, None, &mut cache);
assert!(
same(&p0, &p1),
"{precision:?}: the changed operator's solutions differ"
);
assert_eq!(s0.iterations, s1.iterations);
// And back to the base operator: a miss again (the cache holds one), still exact.
let (mut p0, mut p1) = (vec![0.0; nx * ny], vec![0.0; nx * ny]);
let s0 = solve_multigrid_pcg(&base, &mut p0, &params, 1e-12, None);
let s1 = solve_multigrid_pcg_cached(&base, &mut p1, &params, 1e-12, None, &mut cache);
assert!(same(&p0, &p1));
assert_eq!(s0.iterations, s1.iterations);
}
}