From 79c18def329cd48dc740da08a2168e685feef897 Mon Sep 17 00:00:00 2001 From: Omar Sobh Date: Tue, 15 Sep 2026 23:30:53 -0500 Subject: [PATCH] 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 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01YJPeT6WA2e7YvAnS875AHL --- .../solvers/incompressible/embedded/mod.rs | 4 + .../incompressible/embedded/projection.rs | 5 +- .../rtx-cfd/src/solvers/incompressible/mod.rs | 3 +- .../src/solvers/incompressible/poisson.rs | 180 ++++++++++++++++-- .../rtx-cfd/tests/poisson_cache_exact.rs | 117 ++++++++++++ 5 files changed, 294 insertions(+), 15 deletions(-) create mode 100644 crates/specialized/rtx-cfd/tests/poisson_cache_exact.rs diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded/mod.rs index 63b581d..3c12b2a 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded/mod.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded/mod.rs @@ -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, 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, diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded/projection.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded/projection.rs index 77914fc..9eb44c4 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded/projection.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded/projection.rs @@ -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(( diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs index fee8162..4cccedb 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs @@ -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}; diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs index dbc41d9..cd2e78c 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs @@ -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, + coefficients: Vec, + 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 { + key: OperatorKey, + hier: Hierarchy, + fine: Level, + cells: Vec, + components: Components, +} + +impl Prepared { + fn build(problem: &PoissonProblem, params: &MultigridParameters) -> Self { + let hier = Hierarchy::::build(problem, params); + let fine = Level::::new(problem.clone()); + let cells: Vec = 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>, + f32: Option>, +} + +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, + cache: &mut PcgCache, +) -> PoissonSolution { + match params.precision { + MgPrecision::F64 => { + solve_cached_with::(problem, p, params, tolerance, anchor, &mut cache.f64) + } + MgPrecision::F32 => { + solve_cached_with::(problem, p, params, tolerance, anchor, &mut cache.f32) + } + } +} + +fn solve_cached_with( + problem: &PoissonProblem, + p: &mut [f64], + params: &MultigridParameters, + tolerance: f64, + anchor: Option, + slot: &mut Option>, +) -> 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::::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( problem: &PoissonProblem, p: &mut [f64], params: &MultigridParameters, tolerance: f64, anchor: Option, +) -> PoissonSolution { + let t_entry = std::time::Instant::now(); + let mut prep = Prepared::::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( + prep: &mut Prepared, + problem: &PoissonProblem, + p: &mut [f64], + params: &MultigridParameters, + tolerance: f64, + anchor: Option, + 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( "invalid PoissonProblem: {:?}", problem.validate() ); - let t_entry = std::time::Instant::now(); - - let mut hier = Hierarchy::::build(problem, params); - let fine = Level::::new(problem.clone()); - let cells: Vec = 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( // 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( // 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( |p: &[f64], r: &mut [f64], fine: &Level| -> 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( 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( 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( 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]; } } diff --git a/crates/specialized/rtx-cfd/tests/poisson_cache_exact.rs b/crates/specialized/rtx-cfd/tests/poisson_cache_exact.rs new file mode 100644 index 0000000..c970544 --- /dev/null +++ b/crates/specialized/rtx-cfd/tests/poisson_cache_exact.rs @@ -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, ¶ms, 1e-12, None); + let s1 = solve_multigrid_pcg_cached(&prob, &mut p1, ¶ms, 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, ¶ms, 1e-12, None); + let s1 = solve_multigrid_pcg_cached(&changed, &mut p1, ¶ms, 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, ¶ms, 1e-12, None); + let s1 = solve_multigrid_pcg_cached(&base, &mut p1, ¶ms, 1e-12, None, &mut cache); + assert!(same(&p0, &p1)); + assert_eq!(s0.iterations, s1.iterations); + } +}