PERF-2 P2: the red-black colour maps, the residual and the V-cycle's per-level maps on rayon threads (MultigridParameters::threads; RTX_THREADS in the harness; rtx_cfd::configure_threads) — each colour's values computed from the unchanged other colour into a scratch and written back, the L1 sum in the serial order: bit-identical to the serial red-black (pin: 4 right-hand sides at 4 threads); the coarsest level stays serial; no effect on the lexicographic regime
CI / Build (macos-latest) (push) Waiting to run
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 / CI Success (push) Blocked by required conditions
CI / Format Check (push) Failing after 6s
Performance Benchmarks / Run Benchmarks (push) Failing after 7s
CI / Clippy Check (push) Failing after 7s
CI / Build (ubuntu-latest) (push) Failing after 7s
Documentation / Build User Guide (push) Successful in 7s
CI / Build CPU-Only (Explicit) (push) Failing after 1m19s
Documentation / Build API Documentation (push) Failing after 1m22s

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:59:29 -05:00
co-authored by Claude Fable 5.1
parent 364766a54a
commit 7f144e0005
7 changed files with 164 additions and 13 deletions
@@ -178,6 +178,8 @@ pub struct EmbeddedPisoSolver {
/// PERF-2 P1.1: the prepared multigrid operator, reused across rounds /// PERF-2 P1.1: the prepared multigrid operator, reused across rounds
/// and steps while the operator's coefficients and mask are unchanged. /// and steps while the operator's coefficients and mask are unchanged.
pcg_cache: std::cell::RefCell<super::poisson::PcgCache>, pcg_cache: std::cell::RefCell<super::poisson::PcgCache>,
/// PERF-2 P2: threads for the multigrid's red-black maps (default 1).
poisson_threads: usize,
moving: bool, moving: bool,
/// Mask hysteresis band in multiples of the min cell size (0 = off). /// Mask hysteresis band in multiples of the min cell size (0 = off).
mask_hysteresis: f64, mask_hysteresis: f64,
@@ -206,6 +208,7 @@ impl EmbeddedPisoSolver {
inner_stop_factor: 1e-2, inner_stop_factor: 1e-2,
poisson_profile: std::cell::Cell::new((0, 0, 0, 0)), poisson_profile: std::cell::Cell::new((0, 0, 0, 0)),
pcg_cache: std::cell::RefCell::new(super::poisson::PcgCache::default()), pcg_cache: std::cell::RefCell::new(super::poisson::PcgCache::default()),
poisson_threads: 1,
moving: false, moving: false,
mask_hysteresis: 0.0, mask_hysteresis: 0.0,
time: 0.0, time: 0.0,
@@ -225,6 +228,12 @@ impl EmbeddedPisoSolver {
self.parameters.poisson_smoother = smoother; self.parameters.poisson_smoother = smoother;
} }
/// Threads for the multigrid's red-black colour maps (PERF-2 P2; the
/// caller configures rayon's global pool). Bit-identical to 1.
pub fn set_poisson_threads(&mut self, threads: usize) {
self.poisson_threads = threads.max(1);
}
/// Mask hysteresis for the moving-body rebuild, as a fraction of the /// Mask hysteresis for the moving-body rebuild, as a fraction of the
/// min cell size (default 0, exactly the plain rebuild). With a band, /// min cell size (default 0, exactly the plain rebuild). With a band,
/// a cell within `band * h_min` of the surface keeps the /// a cell within `band * h_min` of the surface keeps the
@@ -300,6 +300,7 @@ impl EmbeddedPisoSolver {
&MultigridParameters { &MultigridParameters {
precision: self.parameters.poisson_precision, precision: self.parameters.poisson_precision,
smoother: self.parameters.poisson_smoother, smoother: self.parameters.poisson_smoother,
threads: self.poisson_threads,
..MultigridParameters::default() ..MultigridParameters::default()
}, },
inner_stop, inner_stop,
@@ -66,7 +66,7 @@ pub use piso::{PisoParameters, PisoResult, PisoSolver};
pub use piso_gpu::PisoGpuSolver; pub use piso_gpu::PisoGpuSolver;
pub use poisson::{ pub use poisson::{
MgPrecision, MgSmoother, MultigridParameters, PcgCache, PoissonProblem, PoissonSolution, MgPrecision, MgSmoother, MultigridParameters, PcgCache, PoissonProblem, PoissonSolution,
PoissonSolverKind, solve_multigrid_pcg, solve_multigrid_pcg_cached, PoissonSolverKind, configure_threads, solve_multigrid_pcg, solve_multigrid_pcg_cached,
}; };
pub use polygon_sdf::PolygonSdf; pub use polygon_sdf::PolygonSdf;
pub use simple::{ConvectionScheme, SimpleParameters, SimpleResult, SimpleSolver}; pub use simple::{ConvectionScheme, SimpleParameters, SimpleResult, SimpleSolver};
@@ -329,6 +329,11 @@ impl OversetPisoSolver {
}) })
} }
/// PERF-2 P2: threads for the background multigrid's red-black maps.
pub fn set_poisson_threads(&mut self, threads: usize) {
self.background.set_poisson_threads(threads);
}
/// The step timers, when profiling (`RTX_PROFILE`). /// The step timers, when profiling (`RTX_PROFILE`).
pub fn timers(&self) -> Option<&StepTimers> { pub fn timers(&self) -> Option<&StepTimers> {
self.timers.as_deref() self.timers.as_deref()
@@ -289,6 +289,12 @@ pub struct MultigridParameters {
pub precision: MgPrecision, pub precision: MgPrecision,
/// Smoother ordering (see [`MgSmoother`]). /// Smoother ordering (see [`MgSmoother`]).
pub smoother: MgSmoother, pub smoother: MgSmoother,
/// Threads for the red-black colour maps, the residual and the matvec
/// (rayon; default 1 = the serial code path). The parallel path computes
/// the same per-cell values from the same inputs and sums in the same
/// order, so it is bit-identical to the serial one. No effect on the
/// lexicographic smoother (a dependency chain).
pub threads: usize,
/// Symmetric GaussSeidel sweeps before AND after the coarse correction /// Symmetric GaussSeidel sweeps before AND after the coarse correction
/// (default 2; a value of 0 is treated as 1). One count for both on /// (default 2; a value of 0 is treated as 1). One count for both on
/// purpose: unequal pre/post counts make the V-cycle non-symmetric and /// purpose: unequal pre/post counts make the V-cycle non-symmetric and
@@ -307,6 +313,7 @@ impl Default for MultigridParameters {
Self { Self {
precision: MgPrecision::F64, precision: MgPrecision::F64,
smoother: MgSmoother::Lexicographic, smoother: MgSmoother::Lexicographic,
threads: 1,
smoother_sweeps: 2, smoother_sweeps: 2,
coarsest_cells: 32, coarsest_cells: 32,
max_iterations: 500, max_iterations: 500,
@@ -380,6 +387,8 @@ const MAX_LEVELS: usize = 64;
/// `f32` is the mixed-precision probe. /// `f32` is the mixed-precision probe.
pub trait MgScalar: pub trait MgScalar:
Copy Copy
+ Send
+ Sync
+ PartialEq + PartialEq
+ PartialOrd + PartialOrd
+ std::ops::Add<Output = Self> + std::ops::Add<Output = Self>
@@ -464,6 +473,8 @@ struct Work<T: MgScalar> {
x: Vec<T>, x: Vec<T>,
/// Residual. /// Residual.
r: Vec<T>, r: Vec<T>,
/// Scratch for the parallel maps (one entry per active cell).
tmp: Vec<T>,
} }
impl<T: MgScalar> Work<T> { impl<T: MgScalar> Work<T> {
@@ -472,6 +483,7 @@ impl<T: MgScalar> Work<T> {
b: vec![T::ZERO; n], b: vec![T::ZERO; n],
x: vec![T::ZERO; n], x: vec![T::ZERO; n],
r: vec![T::ZERO; n], r: vec![T::ZERO; n],
tmp: vec![T::ZERO; n],
} }
} }
} }
@@ -583,6 +595,56 @@ impl<T: MgScalar> Level<T> {
l1 l1
} }
/// [`Self::residual`] with the per-cell values computed in parallel into
/// `tmp` (one per active cell, in `cells` order) and the L1 sum taken
/// serially in the same order: bit-identical.
fn residual_par(&self, b: &[T], x: &[T], r: &mut [T], tmp: &mut [T]) -> T {
use rayon::prelude::*;
let n = self.cells.len();
self.cells
.par_iter()
.zip(tmp[..n].par_iter_mut())
.for_each(|(&idx, t)| {
*t = b[idx] - (self.ap[idx] * x[idx] - self.neighbour_sum(x, idx));
});
let mut l1 = T::ZERO;
for (k, &idx) in self.cells.iter().enumerate() {
let v = tmp[k];
r[idx] = v;
l1 += v.abs();
}
l1
}
/// One red-black half-sweep over `colour` in parallel: every cell of a
/// colour reads only the other colour, so the values are computed into
/// `tmp` from the unchanged `x` and written back — the serial
/// half-sweep's values exactly.
fn half_sweep_par(&self, colour: &[usize], b: &[T], x: &mut [T], tmp: &mut [T]) {
use rayon::prelude::*;
let n = colour.len();
{
let x_ro: &[T] = x;
colour
.par_iter()
.zip(tmp[..n].par_iter_mut())
.for_each(|(&idx, t)| {
*t = (b[idx] + self.neighbour_sum(x_ro, idx)) / self.ap[idx];
});
}
for (k, &idx) in colour.iter().enumerate() {
x[idx] = tmp[k];
}
}
/// [`Self::symmetric_gs_rb`] on threads (bit-identical).
fn symmetric_gs_rb_par(&self, b: &[T], x: &mut [T], tmp: &mut [T]) {
self.half_sweep_par(&self.red, b, x, tmp);
self.half_sweep_par(&self.black, b, x, tmp);
self.half_sweep_par(&self.black, b, x, tmp);
self.half_sweep_par(&self.red, b, x, tmp);
}
/// One symmetric GaussSeidel sweep (forward then backward) on `A x = b`. /// One symmetric GaussSeidel sweep (forward then backward) on `A x = b`.
fn symmetric_gs(&self, b: &[T], x: &mut [T]) { fn symmetric_gs(&self, b: &[T], x: &mut [T]) {
for &idx in &self.cells { for &idx in &self.cells {
@@ -612,11 +674,13 @@ impl<T: MgScalar> Level<T> {
} }
} }
/// One symmetric sweep in the chosen ordering. /// One symmetric sweep in the chosen ordering (`tmp`: the parallel
fn smooth(&self, b: &[T], x: &mut [T], smoother: MgSmoother) { /// path's scratch, used only with red-black on threads).
match smoother { fn smooth(&self, b: &[T], x: &mut [T], tmp: &mut [T], smoother: MgSmoother, parallel: bool) {
MgSmoother::Lexicographic => self.symmetric_gs(b, x), match (smoother, parallel) {
MgSmoother::RedBlack => self.symmetric_gs_rb(b, x), (MgSmoother::Lexicographic, _) => self.symmetric_gs(b, x),
(MgSmoother::RedBlack, false) => self.symmetric_gs_rb(b, x),
(MgSmoother::RedBlack, true) => self.symmetric_gs_rb_par(b, x, tmp),
} }
} }
@@ -663,6 +727,7 @@ pub(crate) struct Hierarchy<T: MgScalar = f64> {
work: Vec<Work<T>>, work: Vec<Work<T>>,
sweeps: usize, sweeps: usize,
smoother: MgSmoother, smoother: MgSmoother,
parallel: bool,
} }
impl<T: MgScalar> Hierarchy<T> { impl<T: MgScalar> Hierarchy<T> {
@@ -693,6 +758,7 @@ impl<T: MgScalar> Hierarchy<T> {
work, work,
sweeps: params.smoother_sweeps.max(1), sweeps: params.smoother_sweeps.max(1),
smoother: params.smoother, smoother: params.smoother,
parallel: params.threads > 1,
} }
} }
@@ -727,6 +793,7 @@ impl<T: MgScalar> Hierarchy<T> {
self.work[0].b[idx] = T::from_f64(r[idx]); self.work[0].b[idx] = T::from_f64(r[idx]);
} }
// Down: smooth from zero, restrict the residual. // Down: smooth from zero, restrict the residual.
let (smoother, parallel) = (self.smoother, self.parallel);
for l in 0..depth - 1 { for l in 0..depth - 1 {
let (fine, coarse) = (&levels[l], &levels[l + 1]); let (fine, coarse) = (&levels[l], &levels[l + 1]);
let (head, tail) = self.work.split_at_mut(l + 1); let (head, tail) = self.work.split_at_mut(l + 1);
@@ -734,10 +801,15 @@ impl<T: MgScalar> Hierarchy<T> {
for &idx in &fine.cells { for &idx in &fine.cells {
wf.x[idx] = T::ZERO; wf.x[idx] = T::ZERO;
} }
let Work { b, x, r, tmp } = wf;
for _ in 0..self.sweeps { for _ in 0..self.sweeps {
fine.smooth(&wf.b, &mut wf.x, self.smoother); fine.smooth(b, x, tmp, smoother, parallel);
}
if parallel {
fine.residual_par(b, x, r, tmp);
} else {
fine.residual(b, x, r);
} }
fine.residual(&wf.b, &wf.x, &mut wf.r);
for &idx in &coarse.cells { for &idx in &coarse.cells {
wc.b[idx] = T::ZERO; wc.b[idx] = T::ZERO;
} }
@@ -746,7 +818,7 @@ impl<T: MgScalar> Hierarchy<T> {
// A coarse cell without an equation (a whole component // A coarse cell without an equation (a whole component
// inside one aggregate) receives the component's zero sum. // inside one aggregate) receives the component's zero sum.
if coarse.active[c] { if coarse.active[c] {
wc.b[c] += wf.r[idx]; wc.b[c] += r[idx];
} }
} }
} }
@@ -757,8 +829,10 @@ impl<T: MgScalar> Hierarchy<T> {
for &idx in &bottom.cells { for &idx in &bottom.cells {
wb.x[idx] = T::ZERO; wb.x[idx] = T::ZERO;
} }
let Work { b, x, tmp, .. } = wb;
// The coarsest level is tiny (≤ 32 cells): always serial.
for _ in 0..COARSEST_SWEEPS { for _ in 0..COARSEST_SWEEPS {
bottom.smooth(&wb.b, &mut wb.x, self.smoother); bottom.smooth(b, x, tmp, smoother, false);
} }
} }
// Up: prolongate, smooth. // Up: prolongate, smooth.
@@ -769,8 +843,9 @@ impl<T: MgScalar> Hierarchy<T> {
for &idx in &fine.cells { for &idx in &fine.cells {
wf.x[idx] += T::from_f64(COARSE_CORRECTION) * wc.x[fine.coarse_of[idx]]; wf.x[idx] += T::from_f64(COARSE_CORRECTION) * wc.x[fine.coarse_of[idx]];
} }
let Work { b, x, tmp, .. } = wf;
for _ in 0..self.sweeps { for _ in 0..self.sweeps {
fine.smooth(&wf.b, &mut wf.x, self.smoother); fine.smooth(b, x, tmp, smoother, parallel);
} }
} }
for &idx in &levels[0].cells { for &idx in &levels[0].cells {
@@ -826,6 +901,7 @@ struct OperatorKey {
smoother_sweeps: usize, smoother_sweeps: usize,
coarsest_cells: usize, coarsest_cells: usize,
smoother: MgSmoother, smoother: MgSmoother,
threads: usize,
} }
impl OperatorKey { impl OperatorKey {
@@ -847,6 +923,7 @@ impl OperatorKey {
smoother_sweeps: params.smoother_sweeps, smoother_sweeps: params.smoother_sweeps,
coarsest_cells: params.coarsest_cells, coarsest_cells: params.coarsest_cells,
smoother: params.smoother, smoother: params.smoother,
threads: params.threads,
} }
} }
@@ -856,6 +933,7 @@ impl OperatorKey {
&& self.smoother_sweeps == params.smoother_sweeps && self.smoother_sweeps == params.smoother_sweeps
&& self.coarsest_cells == params.coarsest_cells && self.coarsest_cells == params.coarsest_cells
&& self.smoother == params.smoother && self.smoother == params.smoother
&& self.threads == params.threads
&& self.active == problem.active && self.active == problem.active
&& self.coefficients.iter().copied().eq(problem && self.coefficients.iter().copied().eq(problem
.ae .ae
@@ -898,6 +976,19 @@ impl<T: MgScalar> Prepared<T> {
} }
} }
/// Configure rayon's global thread pool for the multigrid's parallel maps
/// (PERF-2 P2). Idempotent: a pool that already exists is kept (rayon
/// refuses a second global pool) — the first caller decides.
pub fn configure_threads(threads: usize) -> usize {
let n = threads.max(1);
if n > 1 {
let _ = rayon::ThreadPoolBuilder::new()
.num_threads(n)
.build_global();
}
n
}
/// A reusable prepared solver per V-cycle precision (PERF-2 P1.1): the /// A reusable prepared solver per V-cycle precision (PERF-2 P1.1): the
/// operator's hierarchy is rebuilt only when the operator changes. /// operator's hierarchy is rebuilt only when the operator changes.
#[derive(Default)] #[derive(Default)]
@@ -6,8 +6,8 @@
//! red-black solve (the cache keys on the smoother). //! red-black solve (the cache keys on the smoother).
use rtx_cfd::solvers::incompressible::{ use rtx_cfd::solvers::incompressible::{
MgSmoother, MultigridParameters, PcgCache, PoissonProblem, solve_multigrid_pcg, MgSmoother, MultigridParameters, PcgCache, PoissonProblem, configure_threads,
solve_multigrid_pcg_cached, solve_multigrid_pcg, solve_multigrid_pcg_cached,
}; };
fn problem(nx: usize, ny: usize, seed: u64) -> PoissonProblem { fn problem(nx: usize, ny: usize, seed: u64) -> PoissonProblem {
@@ -106,3 +106,35 @@ fn red_black_solves_the_masked_problem_and_caches_exactly() {
assert!(a.iter().zip(&b).all(|(x, y)| x.to_bits() == y.to_bits())); assert!(a.iter().zip(&b).all(|(x, y)| x.to_bits() == y.to_bits()));
assert_eq!(sa.iterations, sb.iterations); assert_eq!(sa.iterations, sb.iterations);
} }
/// The threaded red-black V-cycle (colour maps, residual on threads) is the
/// serial red-black one bit for bit — the same per-cell arithmetic from the
/// same inputs, the sums in the same order.
#[test]
fn threaded_red_black_is_the_serial_red_black_bit_for_bit() {
let (nx, ny) = (96, 40);
configure_threads(4);
for k in 0..4u64 {
let prob = problem(nx, ny, 31 + k);
let serial = MultigridParameters {
smoother: MgSmoother::RedBlack,
threads: 1,
..MultigridParameters::default()
};
let threaded = MultigridParameters {
smoother: MgSmoother::RedBlack,
threads: 4,
..MultigridParameters::default()
};
let (mut a, mut b) = (vec![0.0; nx * ny], vec![0.0; nx * ny]);
let sa = solve_multigrid_pcg(&prob, &mut a, &serial, 1e-12, None);
let sb = solve_multigrid_pcg(&prob, &mut b, &threaded, 1e-12, None);
assert!(sa.converged && sb.converged);
assert_eq!(sa.iterations, sb.iterations, "rhs {k}");
assert!(
a.iter().zip(&b).all(|(x, y)| x.to_bits() == y.to_bits()),
"rhs {k}: threaded red-black differs from serial"
);
}
println!(" threaded red-black: 4 right-hand sides bit-identical to serial at 4 threads");
}
@@ -313,6 +313,19 @@ impl OversetFluid {
convection_scheme: bg_convection(), convection_scheme: bg_convection(),
}, },
)?; )?;
// PERF-2 P2: `RTX_THREADS=n` — rayon's global pool and the multigrid's
// red-black maps on n threads (bit-identical to 1; no effect on the
// lexicographic smoother).
let threads: usize = std::env::var("RTX_THREADS")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(1)
.max(1);
if threads > 1 {
rtx_cfd::solvers::incompressible::configure_threads(threads);
background.set_poisson_threads(threads);
println!(" multigrid threads: {threads} (RTX_THREADS)");
}
background.set_boundary_velocity(move |x, y, t| { background.set_boundary_velocity(move |x, y, t| {
if x <= 0.0 { if x <= 0.0 {
(inflow_for(u_mean, y, t), 0.0) (inflow_for(u_mean, y, t), 0.0)