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 <[email protected]>
Claude-Session: https://claude.ai/code/session_01YJPeT6WA2e7YvAnS875AHL

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<bool>,
+    coefficients: Vec<u64>,
+    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<T: MgScalar> {
+    key: OperatorKey,
+    hier: Hierarchy<T>,
+    fine: Level<f64>,
+    cells: Vec<usize>,
+    components: Components,
+}
+
+impl<T: MgScalar> Prepared<T> {
+    fn build(problem: &PoissonProblem, params: &MultigridParameters) -> Self {
+        let hier = Hierarchy::<T>::build(problem, params);
+        let fine = Level::<f64>::new(problem.clone());
+        let cells: Vec<usize> = 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<Prepared<f64>>,
+    f32: Option<Prepared<f32>>,
+}
+
+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<usize>,
+    cache: &mut PcgCache,
+) -> PoissonSolution {
+    match params.precision {
+        MgPrecision::F64 => {
+            solve_cached_with::<f64>(problem, p, params, tolerance, anchor, &mut cache.f64)
+        }
+        MgPrecision::F32 => {
+            solve_cached_with::<f32>(problem, p, params, tolerance, anchor, &mut cache.f32)
+        }
+    }
+}
+
+fn solve_cached_with<T: MgScalar>(
+    problem: &PoissonProblem,
+    p: &mut [f64],
+    params: &MultigridParameters,
+    tolerance: f64,
+    anchor: Option<usize>,
+    slot: &mut Option<Prepared<T>>,
+) -> 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::<T>::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<T: MgScalar>(
     problem: &PoissonProblem,
     p: &mut [f64],
     params: &MultigridParameters,
     tolerance: f64,
     anchor: Option<usize>,
+) -> PoissonSolution {
+    let t_entry = std::time::Instant::now();
+    let mut prep = Prepared::<T>::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<T: MgScalar>(
+    prep: &mut Prepared<T>,
+    problem: &PoissonProblem,
+    p: &mut [f64],
+    params: &MultigridParameters,
+    tolerance: f64,
+    anchor: Option<usize>,
+    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<T: MgScalar>(
         "invalid PoissonProblem: {:?}",
         problem.validate()
     );
-    let t_entry = std::time::Instant::now();
-
-    let mut hier = Hierarchy::<T>::build(problem, params);
-    let fine = Level::<f64>::new(problem.clone());
-    let cells: Vec<usize> = 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<T: MgScalar>(
     // 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<T: MgScalar>(
 
     // 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<T: MgScalar>(
         |p: &[f64], r: &mut [f64], fine: &Level<f64>| -> 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<T: MgScalar>(
     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<T: MgScalar>(
             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<T: MgScalar>(
         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];
         }
     }