rtx-cfd 3D Stage 1 item 3: cg3.cu + DeviceCg3 — the f64 CG entirely on the device (run_pcg's control flow: true-residual stop and resync, breakdown guard, mean projection, anchor shift; fixed-order block reductions, no atomics; the device V-cycle as preconditioner; DevicePcgCache3); gate 3 HELD: iterations = host, |Δp| ≤ 2e-14 relative, run-to-run bit-identical, singular box anchor exact / mean 2e-15

Co-Authored-By: Claude Fable 5.1 <[email protected]>
2026-09-17 11:18:52 -05:00
parent 05ed96f383
commit 2f476a38d5
4 changed files with 1008 additions and 0 deletions
@@ -0,0 +1,179 @@
+//! 3D Stage 1, gate 3: the device-resident CG against the host PCG on the
+//! same operator — converged, true residual below the tolerance, solutions
+//! agreeing to the solve's accuracy, run-to-run bit-identical; the singular
+//! closed box honours the anchor / the zero mean.
+//!
+//! `RTX_CUDA_ARCH=sm_120 cargo test --release -p rtx-cfd --features cuda --test three_d_cg_device -- --nocapture`
+#![cfg(feature = "cuda")]
+
+use rtx_cfd::solvers::incompressible::three_d::poisson::device_cg::{
+    DevicePcgCache3, solve_multigrid_pcg3_device_cached,
+};
+use rtx_cfd::solvers::incompressible::three_d::poisson::{PoissonProblem3D, solve_multigrid_pcg3};
+use rtx_cfd::solvers::incompressible::{MgSmoother, MultigridParameters};
+
+/// A channel on cubic cells with a z-cylinder hole and an outlet Dirichlet
+/// (`dirichlet = true`), or a closed Neumann box (`false`), seeded rhs.
+fn channel(
+    nx: usize,
+    ny: usize,
+    nz: usize,
+    periodic_z: bool,
+    dirichlet: bool,
+    seed: u64,
+) -> PoissonProblem3D {
+    let mut p = PoissonProblem3D::new(nx, ny, nz);
+    p.periodic_z = periodic_z;
+    let h = 0.41 / ny as f64;
+    let a = 3.24e-4 * h;
+    let hole = |i: usize, j: usize| {
+        let (x, y) = ((i as f64 + 0.5) * h, (j as f64 + 0.5) * h);
+        (x - 0.2).powi(2) + (y - 0.2).powi(2) < 0.05 * 0.05
+    };
+    for k in 0..nz {
+        for j in 0..ny {
+            for i in 0..nx {
+                let idx = p.index(k, j, i);
+                if hole(i, j) {
+                    p.active[idx] = false;
+                    continue;
+                }
+                if i + 1 < nx && !hole(i + 1, j) {
+                    p.ae[idx] = a;
+                }
+                if i > 0 && !hole(i - 1, j) {
+                    p.aw[idx] = a;
+                }
+                if j + 1 < ny && !hole(i, j + 1) {
+                    p.an[idx] = a;
+                }
+                if j > 0 && !hole(i, j - 1) {
+                    p.as_[idx] = a;
+                }
+                if k + 1 < nz || periodic_z {
+                    p.at[idx] = a;
+                }
+                if k > 0 || periodic_z {
+                    p.ab[idx] = a;
+                }
+                if dirichlet && i + 1 == nx {
+                    p.extra_diag[idx] = 2.0 * a;
+                }
+            }
+        }
+    }
+    let mut state = seed | 1;
+    for idx in 0..nx * ny * nz {
+        state ^= state << 13;
+        state ^= state >> 7;
+        state ^= state << 17;
+        p.rhs[idx] = if p.active[idx] {
+            1e-3 * ((state >> 11) as f64 / (1u64 << 53) as f64 - 0.5)
+        } else {
+            0.0
+        };
+    }
+    if !dirichlet {
+        // A compatible right-hand side: zero mean over the active cells.
+        let cells: Vec<usize> = (0..nx * ny * nz).filter(|&i| p.active[i]).collect();
+        let mean = cells.iter().map(|&i| p.rhs[i]).sum::<f64>() / cells.len() as f64;
+        for &i in &cells {
+            p.rhs[i] -= mean;
+        }
+    }
+    p
+}
+
+fn bits(v: &[f64]) -> Vec<u64> {
+    v.iter().map(|x| x.to_bits()).collect()
+}
+
+#[test]
+fn device_cg_matches_the_host_pcg() {
+    let params = MultigridParameters {
+        smoother: MgSmoother::RedBlack,
+        ..MultigridParameters::default()
+    };
+    let tol = 1e-12;
+    let mut cache = DevicePcgCache3::default();
+    for (nz, periodic, seed) in [(1usize, false, 3u64), (8, true, 5), (8, false, 7)] {
+        let p = channel(96, 40, nz, periodic, true, seed);
+        let n = p.nx * p.ny * p.nz;
+        let mut host = vec![0.0; n];
+        let sh = solve_multigrid_pcg3(&p, &mut host, &params, tol, None);
+        let mut dev = vec![0.0; n];
+        let sd = solve_multigrid_pcg3_device_cached(&p, &mut dev, &params, tol, None, &mut cache);
+        let mut dev2 = vec![0.0; n];
+        let sd2 = solve_multigrid_pcg3_device_cached(&p, &mut dev2, &params, tol, None, &mut cache);
+        assert!(
+            sh.converged && sd.converged,
+            "nz {nz}: host {} / device {}",
+            sh.converged,
+            sd.converged
+        );
+        let res_dev = p.residual_l1(&dev);
+        assert!(res_dev < tol, "nz {nz}: device true residual {res_dev:.3e}");
+        let scale = host.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
+        let worst = host
+            .iter()
+            .zip(&dev)
+            .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
+        println!(
+            "  96×40×{nz} periodic {periodic}: host {} it / device {} it (cache hit second solve: {} it); |Δp| {worst:.3e} of {scale:.3e}; device residual {res_dev:.3e}",
+            sh.iterations, sd.iterations, sd2.iterations
+        );
+        assert!(
+            worst <= 1e-10 * scale,
+            "nz {nz}: device and host differ by {worst:.3e} of {scale:.3e}"
+        );
+        assert_eq!(
+            bits(&dev),
+            bits(&dev2),
+            "nz {nz}: the device solve is not run-to-run identical"
+        );
+        assert_eq!(sd.iterations, sd2.iterations);
+    }
+}
+
+#[test]
+fn the_singular_box_honours_the_anchor_and_the_mean() {
+    let params = MultigridParameters {
+        smoother: MgSmoother::RedBlack,
+        ..MultigridParameters::default()
+    };
+    let tol = 1e-12;
+    let p = channel(48, 20, 8, false, false, 9);
+    let n = p.nx * p.ny * p.nz;
+    assert!(p.is_singular());
+    let cells: Vec<usize> = (0..n).filter(|&i| p.active[i]).collect();
+    let anchor = cells[cells.len() / 3];
+    let mut cache = DevicePcgCache3::default();
+    let mut host = vec![0.0; n];
+    let sh = solve_multigrid_pcg3(&p, &mut host, &params, tol, Some(anchor));
+    let mut dev = vec![0.0; n];
+    let sd =
+        solve_multigrid_pcg3_device_cached(&p, &mut dev, &params, tol, Some(anchor), &mut cache);
+    assert!(sh.converged && sd.converged);
+    assert_eq!(
+        dev[anchor].to_bits(),
+        0.0f64.to_bits(),
+        "anchor not at zero: {}",
+        dev[anchor]
+    );
+    let scale = host.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
+    let worst = host
+        .iter()
+        .zip(&dev)
+        .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
+    println!(
+        "  singular box with anchor: host {} it / device {} it; |Δp| {worst:.3e} of {scale:.3e}",
+        sh.iterations, sd.iterations
+    );
+    assert!(worst <= 1e-10 * scale);
+    let mut dev0 = vec![0.0; n];
+    let s0 = solve_multigrid_pcg3_device_cached(&p, &mut dev0, &params, tol, None, &mut cache);
+    assert!(s0.converged);
+    let mean = cells.iter().map(|&i| dev0[i]).sum::<f64>() / cells.len() as f64;
+    println!("  singular box without anchor: mean {mean:.3e} of {scale:.3e}");
+    assert!(mean.abs() <= 1e-13 * scale, "mean {mean:.3e}");
+}