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 <[email protected]>

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<T: MgScalar> Level<T> {
         let n = nx * ny * nz;
         let ap: Vec<f64> = problem.diagonals();
         let active: Vec<bool> = (0..n)
+            .into_par_iter()
             .map(|idx| problem.active[idx] && ap[idx] > 0.0)
             .collect();
-        let mut top = vec![usize::MAX; n];
+        // Couplings to inactive neighbours are zeroed and the z links recorded.
-        let mut bot = vec![usize::MAX; n];
+        // Every entry depends on the activity flags alone, so the pass runs in
-        for k in 0..nz {
+        // parallel with the serial loop's numbers (PERF-3 P1-1: this constructor
-            for j in 0..ny {
+        // runs on every step of a moving body).
-                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<usize> = (0..n).filter(|&idx| active[idx]).collect();
         let nxy = nx * ny;
+        let mask = |v: &[f64], ok: &(dyn Fn(usize) -> bool + Sync)| -> Vec<f64> {
+            (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<usize> {
+            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<usize> {
+            (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<usize> = (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<usize> = cells
-            .iter()
+            .par_iter()
             .copied()
             .filter(|&idx| parity(idx) == 0)
             .collect();
         let black: Vec<usize> = cells
-            .iter()
+            .par_iter()
             .copied()
             .filter(|&idx| parity(idx) == 1)
             .collect();
-        let cast = |v: &[f64]| v.iter().map(|&x| T::from_f64(x)).collect::<Vec<T>>();
+        let cast = |v: &[f64]| v.par_iter().map(|&x| T::from_f64(x)).collect::<Vec<T>>();
         let links: Vec<Vec<(usize, T)>> = if problem.links.is_empty() {
             Vec::new()
         } else {

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]