From 73e43d597d532dbfdeafdef14d5eac16ca04124a Mon Sep 17 00:00:00 2001 From: Omar Sobh Date: Sat, 19 Sep 2026 16:41:28 -0500 Subject: [PATCH] embedded3 PERF-3 P1-1: the level constructor's masking pass, cell lists and the diagonal sum as parallel per-entry maps (level 0.90 -> 0.60 s per step at 11.6 M cells), digit-identical (slab flag CSV byte-identical; host Poisson / MMS / identity and device cg / moving / vcycle tests green) Co-Authored-By: Claude Fable 5.1 --- .../embedded3/poisson/hierarchy.rs | 99 ++++++++++--------- .../embedded3/poisson/problem.rs | 2 + 2 files changed, 57 insertions(+), 44 deletions(-) diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/hierarchy.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/hierarchy.rs index 4a4049d..0153475 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/hierarchy.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/hierarchy.rs @@ -6,6 +6,7 @@ use super::{COARSE_CORRECTION, COARSEST_SWEEPS, MAX_LEVELS, Problem}; use crate::solvers::incompressible::poisson::{MgScalar, MgSmoother, MultigridParameters}; +use rayon::prelude::*; /// One level: the problem, its coefficients in the V-cycle scalar, the /// active cells, the two colours, the neighbour arrays, the parent map. @@ -43,62 +44,72 @@ impl Level { let n = nx * ny * nz; let ap: Vec = problem.diagonals(); let active: Vec = (0..n) + .into_par_iter() .map(|idx| problem.active[idx] && ap[idx] > 0.0) .collect(); - let mut top = vec![usize::MAX; n]; - let mut bot = vec![usize::MAX; n]; - for k in 0..nz { - for j in 0..ny { - for i in 0..nx { - let idx = problem.index(k, j, i); - let t = problem.top(idx, k); - let b = problem.bottom(idx, k); - if !active[idx] { - problem.ae[idx] = 0.0; - problem.aw[idx] = 0.0; - problem.an[idx] = 0.0; - problem.as_[idx] = 0.0; - problem.at[idx] = 0.0; - problem.ab[idx] = 0.0; - continue; - } - if !(i + 1 < nx && active[idx + 1]) { - problem.ae[idx] = 0.0; - } - if !(i > 0 && active[idx - 1]) { - problem.aw[idx] = 0.0; - } - if !(j + 1 < ny && active[idx + nx]) { - problem.an[idx] = 0.0; - } - if !(j > 0 && active[idx - nx]) { - problem.as_[idx] = 0.0; - } - match t { - Some(t) if active[t] => top[idx] = t, - _ => problem.at[idx] = 0.0, - } - match b { - Some(b) if active[b] => bot[idx] = b, - _ => problem.ab[idx] = 0.0, - } - } - } - } - let cells: Vec = (0..n).filter(|&idx| active[idx]).collect(); + // Couplings to inactive neighbours are zeroed and the z links recorded. + // Every entry depends on the activity flags alone, so the pass runs in + // parallel with the serial loop's numbers (PERF-3 P1-1: this constructor + // runs on every step of a moving body). let nxy = nx * ny; + let mask = |v: &[f64], ok: &(dyn Fn(usize) -> bool + Sync)| -> Vec { + (0..n) + .into_par_iter() + .map(|idx| if active[idx] && ok(idx) { v[idx] } else { 0.0 }) + .collect() + }; + let link = |idx: usize, up: bool| -> Option { + let k = idx / nxy; + let o = if up { + problem.top(idx, k) + } else { + problem.bottom(idx, k) + }; + o.filter(|&t| active[t]) + }; + let ae = mask(&problem.ae, &|idx| idx % nx + 1 < nx && active[idx + 1]); + let aw = mask(&problem.aw, &|idx| idx % nx > 0 && active[idx - 1]); + let an = mask(&problem.an, &|idx| { + (idx % nxy) / nx + 1 < ny && active[idx + nx] + }); + let as_ = mask(&problem.as_, &|idx| { + (idx % nxy) / nx > 0 && active[idx - nx] + }); + let at = mask(&problem.at, &|idx| link(idx, true).is_some()); + let ab = mask(&problem.ab, &|idx| link(idx, false).is_some()); + let z_link = |up: bool| -> Vec { + (0..n) + .into_par_iter() + .map(|idx| { + if active[idx] { + link(idx, up).unwrap_or(usize::MAX) + } else { + usize::MAX + } + }) + .collect() + }; + let (top, bot) = (z_link(true), z_link(false)); + problem.ae = ae; + problem.aw = aw; + problem.an = an; + problem.as_ = as_; + problem.at = at; + problem.ab = ab; + // rayon's collect keeps the index order, so the lists equal the serial ones. + let cells: Vec = (0..n).into_par_iter().filter(|&idx| active[idx]).collect(); let parity = |idx: usize| (idx % nx + (idx % nxy) / nx + idx / nxy) % 2; let red: Vec = cells - .iter() + .par_iter() .copied() .filter(|&idx| parity(idx) == 0) .collect(); let black: Vec = cells - .iter() + .par_iter() .copied() .filter(|&idx| parity(idx) == 1) .collect(); - let cast = |v: &[f64]| v.iter().map(|&x| T::from_f64(x)).collect::>(); + let cast = |v: &[f64]| v.par_iter().map(|&x| T::from_f64(x)).collect::>(); let links: Vec> = if problem.links.is_empty() { Vec::new() } else { diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/problem.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/problem.rs index f8d616a..a9e9276 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/problem.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/problem.rs @@ -3,6 +3,7 @@ //! periodic; the neighbour above/below a cell is computed HERE, once, and //! stored as data by the hierarchy — no stencil branches on it. +use rayon::prelude::*; #[derive(Debug, Clone)] pub struct Problem { pub nx: usize, @@ -141,6 +142,7 @@ impl Problem { link_sum[b] += c; } (0..n) + .into_par_iter() .map(|idx| { let stencil = self.ae[idx] + self.aw[idx]