From 7f144e00056fe4cc3fcf26b420c713add475353e Mon Sep 17 00:00:00 2001 From: Omar Sobh Date: Tue, 15 Sep 2026 23:59:29 -0500 Subject: [PATCH] =?UTF-8?q?PERF-2=20P2:=20the=20red-black=20colour=20maps,?= =?UTF-8?q?=20the=20residual=20and=20the=20V-cycle's=20per-level=20maps=20?= =?UTF-8?q?on=20rayon=20threads=20(MultigridParameters::threads;=20RTX=5FT?= =?UTF-8?q?HREADS=20in=20the=20harness;=20rtx=5Fcfd::configure=5Fthreads)?= =?UTF-8?q?=20=E2=80=94=20each=20colour's=20values=20computed=20from=20the?= =?UTF-8?q?=20unchanged=20other=20colour=20into=20a=20scratch=20and=20writ?= =?UTF-8?q?ten=20back,=20the=20L1=20sum=20in=20the=20serial=20order:=20bit?= =?UTF-8?q?-identical=20to=20the=20serial=20red-black=20(pin:=204=20right-?= =?UTF-8?q?hand=20sides=20at=204=20threads);=20the=20coarsest=20level=20st?= =?UTF-8?q?ays=20serial;=20no=20effect=20on=20the=20lexicographic=20regime?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01YJPeT6WA2e7YvAnS875AHL --- .../solvers/incompressible/embedded/mod.rs | 9 ++ .../incompressible/embedded/projection.rs | 1 + .../rtx-cfd/src/solvers/incompressible/mod.rs | 2 +- .../src/solvers/incompressible/overset/mod.rs | 5 + .../src/solvers/incompressible/poisson.rs | 111 ++++++++++++++++-- .../rtx-cfd/tests/poisson_redblack.rs | 36 +++++- .../rtx-fsi/tests/fsi2_harness/overset.rs | 13 ++ 7 files changed, 164 insertions(+), 13 deletions(-) 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 ef67563..1c165a3 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded/mod.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded/mod.rs @@ -178,6 +178,8 @@ pub struct EmbeddedPisoSolver { /// 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, + /// PERF-2 P2: threads for the multigrid's red-black maps (default 1). + poisson_threads: usize, moving: bool, /// Mask hysteresis band in multiples of the min cell size (0 = off). mask_hysteresis: f64, @@ -206,6 +208,7 @@ impl EmbeddedPisoSolver { 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()), + poisson_threads: 1, moving: false, mask_hysteresis: 0.0, time: 0.0, @@ -225,6 +228,12 @@ impl EmbeddedPisoSolver { 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 /// min cell size (default 0, exactly the plain rebuild). With a band, /// a cell within `band * h_min` of the surface keeps the 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 8ec7fd7..6572407 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded/projection.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded/projection.rs @@ -300,6 +300,7 @@ impl EmbeddedPisoSolver { &MultigridParameters { precision: self.parameters.poisson_precision, smoother: self.parameters.poisson_smoother, + threads: self.poisson_threads, ..MultigridParameters::default() }, inner_stop, diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs index d901b6f..fba8a8c 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs @@ -66,7 +66,7 @@ pub use piso::{PisoParameters, PisoResult, PisoSolver}; pub use piso_gpu::PisoGpuSolver; pub use poisson::{ 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 simple::{ConvectionScheme, SimpleParameters, SimpleResult, SimpleSolver}; diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/overset/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/overset/mod.rs index e2486b8..ca44415 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/overset/mod.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/overset/mod.rs @@ -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`). pub fn timers(&self) -> Option<&StepTimers> { self.timers.as_deref() diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs index 91d0b0c..5798fb4 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs @@ -289,6 +289,12 @@ pub struct MultigridParameters { pub precision: MgPrecision, /// Smoother ordering (see [`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 Gauss–Seidel sweeps before AND after the coarse correction /// (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 @@ -307,6 +313,7 @@ impl Default for MultigridParameters { Self { precision: MgPrecision::F64, smoother: MgSmoother::Lexicographic, + threads: 1, smoother_sweeps: 2, coarsest_cells: 32, max_iterations: 500, @@ -380,6 +387,8 @@ const MAX_LEVELS: usize = 64; /// `f32` is the mixed-precision probe. pub trait MgScalar: Copy + + Send + + Sync + PartialEq + PartialOrd + std::ops::Add @@ -464,6 +473,8 @@ struct Work { x: Vec, /// Residual. r: Vec, + /// Scratch for the parallel maps (one entry per active cell). + tmp: Vec, } impl Work { @@ -472,6 +483,7 @@ impl Work { b: vec![T::ZERO; n], x: vec![T::ZERO; n], r: vec![T::ZERO; n], + tmp: vec![T::ZERO; n], } } } @@ -583,6 +595,56 @@ impl Level { 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 Gauss–Seidel sweep (forward then backward) on `A x = b`. fn symmetric_gs(&self, b: &[T], x: &mut [T]) { for &idx in &self.cells { @@ -612,11 +674,13 @@ impl Level { } } - /// One symmetric sweep in the chosen ordering. - fn smooth(&self, b: &[T], x: &mut [T], smoother: MgSmoother) { - match smoother { - MgSmoother::Lexicographic => self.symmetric_gs(b, x), - MgSmoother::RedBlack => self.symmetric_gs_rb(b, x), + /// One symmetric sweep in the chosen ordering (`tmp`: the parallel + /// path's scratch, used only with red-black on threads). + fn smooth(&self, b: &[T], x: &mut [T], tmp: &mut [T], smoother: MgSmoother, parallel: bool) { + match (smoother, parallel) { + (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 { work: Vec>, sweeps: usize, smoother: MgSmoother, + parallel: bool, } impl Hierarchy { @@ -693,6 +758,7 @@ impl Hierarchy { work, sweeps: params.smoother_sweeps.max(1), smoother: params.smoother, + parallel: params.threads > 1, } } @@ -727,6 +793,7 @@ impl Hierarchy { self.work[0].b[idx] = T::from_f64(r[idx]); } // Down: smooth from zero, restrict the residual. + let (smoother, parallel) = (self.smoother, self.parallel); for l in 0..depth - 1 { let (fine, coarse) = (&levels[l], &levels[l + 1]); let (head, tail) = self.work.split_at_mut(l + 1); @@ -734,10 +801,15 @@ impl Hierarchy { for &idx in &fine.cells { wf.x[idx] = T::ZERO; } + let Work { b, x, r, tmp } = wf; 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 { wc.b[idx] = T::ZERO; } @@ -746,7 +818,7 @@ impl Hierarchy { // A coarse cell without an equation (a whole component // inside one aggregate) receives the component's zero sum. if coarse.active[c] { - wc.b[c] += wf.r[idx]; + wc.b[c] += r[idx]; } } } @@ -757,8 +829,10 @@ impl Hierarchy { for &idx in &bottom.cells { 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 { - bottom.smooth(&wb.b, &mut wb.x, self.smoother); + bottom.smooth(b, x, tmp, smoother, false); } } // Up: prolongate, smooth. @@ -769,8 +843,9 @@ impl Hierarchy { for &idx in &fine.cells { 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 { - fine.smooth(&wf.b, &mut wf.x, self.smoother); + fine.smooth(b, x, tmp, smoother, parallel); } } for &idx in &levels[0].cells { @@ -826,6 +901,7 @@ struct OperatorKey { smoother_sweeps: usize, coarsest_cells: usize, smoother: MgSmoother, + threads: usize, } impl OperatorKey { @@ -847,6 +923,7 @@ impl OperatorKey { smoother_sweeps: params.smoother_sweeps, coarsest_cells: params.coarsest_cells, smoother: params.smoother, + threads: params.threads, } } @@ -856,6 +933,7 @@ impl OperatorKey { && self.smoother_sweeps == params.smoother_sweeps && self.coarsest_cells == params.coarsest_cells && self.smoother == params.smoother + && self.threads == params.threads && self.active == problem.active && self.coefficients.iter().copied().eq(problem .ae @@ -898,6 +976,19 @@ impl Prepared { } } +/// 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 /// operator's hierarchy is rebuilt only when the operator changes. #[derive(Default)] diff --git a/crates/specialized/rtx-cfd/tests/poisson_redblack.rs b/crates/specialized/rtx-cfd/tests/poisson_redblack.rs index 482f4f8..c645496 100644 --- a/crates/specialized/rtx-cfd/tests/poisson_redblack.rs +++ b/crates/specialized/rtx-cfd/tests/poisson_redblack.rs @@ -6,8 +6,8 @@ //! red-black solve (the cache keys on the smoother). use rtx_cfd::solvers::incompressible::{ - MgSmoother, MultigridParameters, PcgCache, PoissonProblem, solve_multigrid_pcg, - solve_multigrid_pcg_cached, + MgSmoother, MultigridParameters, PcgCache, PoissonProblem, configure_threads, + solve_multigrid_pcg, solve_multigrid_pcg_cached, }; 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_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"); +} diff --git a/crates/specialized/rtx-fsi/tests/fsi2_harness/overset.rs b/crates/specialized/rtx-fsi/tests/fsi2_harness/overset.rs index e44fa3c..4347470 100644 --- a/crates/specialized/rtx-fsi/tests/fsi2_harness/overset.rs +++ b/crates/specialized/rtx-fsi/tests/fsi2_harness/overset.rs @@ -313,6 +313,19 @@ impl OversetFluid { 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| { if x <= 0.0 { (inflow_for(u_mean, y, t), 0.0)