rtx-cfd 3D Stage 1 item 5: piso3.cu + three_d::piso_device::Piso3Device — the PISO step device-resident (predictors compiled without FMA contraction = the host's arithmetic, per-side boundary tables, divergence/correct/add_p kernels with fixed-order partials, the device CG); poisson_operator/boundary_tables/source_tables/anchor_cell/inner_stop split out of the host solver; StepTimers3 (RTX_PROFILE). Gates 4–6 on the device HELD: MMS 9e-16 (upwind) / 7e-16 (TVD), Beltrami 3.3e-12, Poiseuille 2e-15 with periodic planes equal to 2e-16 (tight tolerances; default-tolerance differences = the projection's inner stop); 55 ms per step at 378×62×62 (predictor 22, poisson 32, apply 2)

Co-Authored-By: Claude Fable 5.1 <[email protected]>
2026-09-17 11:44:04 -05:00
parent e0cbb99343
commit d9fd849239
5 changed files with 1740 additions and 14 deletions
@@ -0,0 +1,561 @@
+//! 3D Stage 1, gates 4–6 on the DEVICE: the device-resident step against
+//! the host step from the same start — the manufactured problem (upwind
+//! and TVD), Beltrami with time-dependent boundary data, and Poiseuille
+//! at nz = 1 and on the periodic extrusion (plane agreement). The
+//! predictors are compiled without FMA contraction, so the only difference
+//! from the host is the CG's reduction order: agreement to 1e-12 relative.
+//!
+//! `RTX_CUDA_ARCH=sm_120 cargo test --release -p rtx-cfd --features cuda --test three_d_device_step -- --nocapture`
+#![cfg(feature = "cuda")]
+
+use rtx_cfd::solvers::incompressible::three_d::piso_device::Piso3Device;
+use rtx_cfd::solvers::incompressible::three_d::{
+    Boundaries3, FlowField3D, Fluid3, Grid3, Piso3Parameters, Piso3Solver, SideBoundary3,
+};
+use rtx_cfd::solvers::incompressible::{ConvectionScheme, MgSmoother};
+use std::f64::consts::PI;
+
+const RHO: f64 = 1.0;
+
+fn fluid(mu: f64) -> Fluid3 {
+    Fluid3 {
+        density: RHO,
+        viscosity: mu,
+        reference_velocity: 1.0,
+        reference_length: 1.0,
+    }
+}
+
+/// `tight` = the identity configuration: inner stop 1e-6 of the source
+/// scale AND mass tolerance 1e-12 (its 0.1× floor on the inner stop is
+/// what dominates near a steady state); otherwise the defaults.
+fn params(scheme: ConvectionScheme, z: SideBoundary3, tight: bool) -> Piso3Parameters {
+    Piso3Parameters {
+        corrector_steps: 2,
+        tolerance: if tight { 1e-12 } else { 1e-8 },
+        inner_stop_factor: if tight { 1e-6 } else { 1e-2 },
+        boundaries: Boundaries3 {
+            z0: z,
+            z1: z,
+            ..Boundaries3::default()
+        },
+        // The device V-cycle is red-black; the host uses the same so the
+        // preconditioners match.
+        poisson_smoother: MgSmoother::RedBlack,
+        convection_scheme: scheme,
+        ..Piso3Parameters::default()
+    }
+}
+
+/// Max |Δ| over u, v, w between two fields on the velocity scale, and max
+/// |Δp| on the pressure scale floored at the dynamic pressure `ρ U²` (a
+/// flat pressure field must not inflate a rounding-level difference).
+fn compare(a: &FlowField3D, b: &FlowField3D) -> (f64, f64) {
+    let mut worst = 0.0_f64;
+    let mut scale = 0.0_f64;
+    for (x, y) in
+        a.u.iter()
+            .zip(&b.u)
+            .chain(a.v.iter().zip(&b.v))
+            .chain(a.w.iter().zip(&b.w))
+    {
+        worst = worst.max((x - y).abs());
+        scale = scale.max(x.abs());
+    }
+    let mut worst_p = 0.0_f64;
+    let mut scale_p = 0.0_f64;
+    for (x, y) in a.p.iter().zip(&b.p) {
+        worst_p = worst_p.max((x - y).abs());
+        scale_p = scale_p.max(x.abs());
+    }
+    let scale_p = scale_p.max(RHO * scale * scale);
+    // One number: the larger of the two relative differences, on the velocity scale.
+    (worst.max(worst_p / scale_p * scale), scale)
+}
+
+// ---- the manufactured solution of three_d_mms.rs ----
+fn u3(x: f64, y: f64, z: f64) -> f64 {
+    (PI * x).sin() * (PI * y).cos() * (PI * z).cos()
+}
+fn v3(x: f64, y: f64, z: f64) -> f64 {
+    (PI * x).cos() * (PI * y).sin() * (PI * z).cos()
+}
+fn w3(x: f64, y: f64, z: f64) -> f64 {
+    -2.0 * (PI * x).cos() * (PI * y).cos() * (PI * z).sin()
+}
+fn source3(mu: f64, x: f64, y: f64, z: f64) -> (f64, f64, f64) {
+    let (sx, cx) = (PI * x).sin_cos();
+    let (sy, cy) = (PI * y).sin_cos();
+    let (sz, cz) = (PI * z).sin_cos();
+    let (u, v, w) = (u3(x, y, z), v3(x, y, z), w3(x, y, z));
+    let (ux, uy, uz) = (PI * cx * cy * cz, -PI * sx * sy * cz, -PI * sx * cy * sz);
+    let (vx, vy, vz) = (-PI * sx * sy * cz, PI * cx * cy * cz, -PI * cx * sy * sz);
+    let (wx, wy, wz) = (
+        2.0 * PI * sx * cy * sz,
+        2.0 * PI * cx * sy * sz,
+        -2.0 * PI * cx * cy * cz,
+    );
+    let (px, py, pz) = (PI * cx * sy * sz, PI * sx * cy * sz, PI * sx * sy * cz);
+    let lap = -3.0 * PI * PI;
+    (
+        RHO * (u * ux + v * uy + w * uz) + px - mu * lap * u,
+        RHO * (u * vx + v * vy + w * vz) + py - mu * lap * v,
+        RHO * (u * wx + v * wy + w * wz) + pz - mu * lap * w,
+    )
+}
+fn boundary3(x: f64, y: f64, z: f64) -> (f64, f64, f64) {
+    let u = if x <= 0.0 || x >= 1.0 {
+        0.0
+    } else {
+        u3(x, y, z)
+    };
+    let v = if y <= 0.0 || y >= 1.0 {
+        0.0
+    } else {
+        v3(x, y, z)
+    };
+    let w = if z <= 0.0 || z >= 1.0 {
+        0.0
+    } else {
+        w3(x, y, z)
+    };
+    (u, v, w)
+}
+
+fn mms_pair(
+    n: usize,
+    scheme: ConvectionScheme,
+    isf: bool,
+) -> (Piso3Solver, Piso3Solver, Grid3, f64) {
+    let mu = 0.05;
+    let h = 1.0 / n as f64;
+    let dt = 0.4 * (h * h / (4.0 * mu / RHO)).min(h);
+    let mk = || {
+        let mut s = Piso3Solver::new(fluid(mu), params(scheme, SideBoundary3::Velocity, isf));
+        s.set_momentum_source(move |x, y, z, _t| source3(mu, x, y, z));
+        s.set_boundary_velocity(|x, y, z, _t| boundary3(x, y, z));
+        s
+    };
+    (
+        mk(),
+        mk(),
+        Grid3 {
+            nx: n,
+            ny: n,
+            nz: n,
+            dx: h,
+            dy: h,
+            dz: h,
+        },
+        dt,
+    )
+}
+
+fn run_pair(
+    host: Piso3Solver,
+    dev: Piso3Solver,
+    g: Grid3,
+    dt: f64,
+    steps: usize,
+    label: &str,
+) -> (f64, f64) {
+    let mut host = host;
+    let mut fh = FlowField3D::new(g);
+    let mut device = Piso3Device::new(dev, g);
+    device.upload(&fh);
+    let mut mismatched = 0usize;
+    for _ in 0..steps {
+        let rh = host.advance(&mut fh, dt);
+        let rd = device.advance(dt);
+        if rh.poisson_iterations != rd.poisson_iterations {
+            mismatched += 1;
+        }
+    }
+    let mut fd = FlowField3D::new(g);
+    device.download(&mut fd);
+    let (worst, scale) = compare(&fh, &fd);
+    println!(
+        "  {label}: {steps} steps, host vs device max |Δ| {worst:.3e} on a scale of {scale:.3e}; CG iteration counts differ on {mismatched} steps"
+    );
+    (worst, scale)
+}
+
+/// The CG's stop is a threshold on a reduction; host and device reduce in
+/// different orders, so on a marginal step one side takes one more
+/// iteration and the answers differ by the inner tolerance. The gate is
+/// therefore taken at a TIGHT inner stop (1e-6 of the source scale), where
+/// that flip cannot show above 1e-12; the default stop's difference is
+/// reported alongside.
+#[test]
+fn device_step_matches_the_host_step_on_the_manufactured_problem() {
+    for scheme in [ConvectionScheme::Upwind, ConvectionScheme::TvdVanAlbada] {
+        let (h, d, g, dt) = mms_pair(12, scheme, false);
+        run_pair(
+            h,
+            d,
+            g,
+            dt,
+            100,
+            &format!("MMS n 12 {scheme:?} (default tolerances)"),
+        );
+        let (h, d, g, dt) = mms_pair(12, scheme, true);
+        let (worst, scale) = run_pair(h, d, g, dt, 100, &format!("MMS n 12 {scheme:?} (tight)"));
+        assert!(
+            worst <= 1e-11 * scale,
+            "{scheme:?}: {worst:.3e} of {scale:.3e}"
+        );
+    }
+}
+
+// ---- Beltrami (three_d_beltrami.rs) ----
+const NU: f64 = 0.02;
+const A: f64 = PI / 4.0;
+const D: f64 = PI / 2.0;
+fn exact(x: f64, y: f64, z: f64, t: f64) -> (f64, f64, f64) {
+    let decay = (-D * D * NU * t).exp();
+    (
+        -A * ((A * x).exp() * (A * y + D * z).sin() + (A * z).exp() * (A * x + D * y).cos())
+            * decay,
+        -A * ((A * y).exp() * (A * z + D * x).sin() + (A * x).exp() * (A * y + D * z).cos())
+            * decay,
+        -A * ((A * z).exp() * (A * x + D * y).sin() + (A * y).exp() * (A * z + D * x).cos())
+            * decay,
+    )
+}
+
+#[test]
+fn device_step_matches_the_host_step_on_beltrami() {
+    let n = 16;
+    let h = 1.0 / n as f64;
+    let dt = 0.25 * h * h / (4.0 * NU);
+    let g = Grid3 {
+        nx: n,
+        ny: n,
+        nz: n,
+        dx: h,
+        dy: h,
+        dz: h,
+    };
+    for tight in [false, true] {
+        let mk = || {
+            let mut s = Piso3Solver::new(
+                fluid(NU * RHO),
+                params(ConvectionScheme::Upwind, SideBoundary3::Velocity, tight),
+            );
+            s.set_boundary_velocity(|x, y, z, t| exact(x, y, z, t));
+            s
+        };
+        let mut host = mk();
+        let mut fh = FlowField3D::new(g);
+        for k in 0..n {
+            for j in 0..n {
+                for i in 0..=n {
+                    fh.u[g.uface(k, j, i)] = exact(
+                        i as f64 * h,
+                        (j as f64 + 0.5) * h,
+                        (k as f64 + 0.5) * h,
+                        0.0,
+                    )
+                    .0;
+                }
+            }
+        }
+        for k in 0..n {
+            for j in 0..=n {
+                for i in 0..n {
+                    fh.v[g.vface(k, j, i)] = exact(
+                        (i as f64 + 0.5) * h,
+                        j as f64 * h,
+                        (k as f64 + 0.5) * h,
+                        0.0,
+                    )
+                    .1;
+                }
+            }
+        }
+        for k in 0..=n {
+            for j in 0..n {
+                for i in 0..n {
+                    fh.w[g.wface(k, j, i)] = exact(
+                        (i as f64 + 0.5) * h,
+                        (j as f64 + 0.5) * h,
+                        k as f64 * h,
+                        0.0,
+                    )
+                    .2;
+                }
+            }
+        }
+        let mut device = Piso3Device::new(mk(), g);
+        device.upload(&fh);
+        let mut mismatched = 0;
+        for _ in 0..40 {
+            let rh = host.advance(&mut fh, dt);
+            let rd = device.advance(dt);
+            if rh.poisson_iterations != rd.poisson_iterations {
+                mismatched += 1;
+            }
+        }
+        let mut fd = FlowField3D::new(g);
+        device.download(&mut fd);
+        let (worst, scale) = compare(&fh, &fd);
+        println!(
+            "  Beltrami n 16 (tight {tight}): 40 steps (time-dependent tables), host vs device max |Δ| {worst:.3e} on {scale:.3e}; CG counts differ on {mismatched} steps"
+        );
+        if tight {
+            assert!(worst <= 1e-11 * scale, "{worst:.3e} of {scale:.3e}");
+        }
+    }
+}
+
+// ---- Poiseuille (three_d_poiseuille_identity.rs) ----
+const MU_P: f64 = 0.1;
+const G: f64 = 0.8;
+fn discrete_profile(n: usize) -> Vec<f64> {
+    let h = 1.0 / n as f64;
+    let rhs_value = -G * h * h / MU_P;
+    let mut diag = vec![-2.0; n];
+    diag[0] = -3.0;
+    diag[n - 1] = -3.0;
+    let mut rhs = vec![rhs_value; n];
+    let upper = vec![1.0; n];
+    for j in 1..n {
+        let factor = 1.0 / diag[j - 1];
+        diag[j] -= factor * upper[j - 1];
+        rhs[j] -= factor * rhs[j - 1];
+    }
+    let mut u = vec![0.0; n];
+    u[n - 1] = rhs[n - 1] / diag[n - 1];
+    for j in (0..n - 1).rev() {
+        u[j] = (rhs[j] - upper[j] * u[j + 1]) / diag[j];
+    }
+    u
+}
+
+#[test]
+fn device_step_matches_the_host_step_on_poiseuille_and_is_z_invariant() {
+    let n = 16;
+    let h = 1.0 / n as f64;
+    let dt = 0.4 * (h * h / (4.0 * MU_P)).min(h);
+    for (nz, dz, z, tight) in [
+        (1usize, 1.0, SideBoundary3::SlipWall, false),
+        (1, 1.0, SideBoundary3::SlipWall, true),
+        (4, h, SideBoundary3::Periodic, true),
+    ] {
+        let g = Grid3 {
+            nx: n,
+            ny: n,
+            nz,
+            dx: h,
+            dy: h,
+            dz,
+        };
+        let mk = || {
+            let mut s = Piso3Solver::new(fluid(MU_P), params(ConvectionScheme::Upwind, z, tight));
+            s.set_momentum_source(|_x, _y, _z, _t| (G, 0.0, 0.0));
+            let u_hat = discrete_profile(n);
+            s.set_boundary_velocity(move |x, y, _z, _t| {
+                if x <= 0.0 || x >= 1.0 {
+                    let j = ((y / h - 0.5).round().max(0.0) as usize).min(n - 1);
+                    (u_hat[j], 0.0, 0.0)
+                } else {
+                    (0.0, 0.0, 0.0)
+                }
+            });
+            s
+        };
+        let mut host = mk();
+        let mut fh = FlowField3D::new(g);
+        let u_hat = discrete_profile(n);
+        for k in 0..nz {
+            for (j, &uj) in u_hat.iter().enumerate() {
+                fh.u[g.uface(k, j, 0)] = uj;
+                fh.u[g.uface(k, j, n)] = uj;
+            }
+        }
+        let mut device = Piso3Device::new(mk(), g);
+        device.upload(&fh);
+        let mut mismatched = 0;
+        for _ in 0..300 {
+            let rh = host.advance(&mut fh, dt);
+            let rd = device.advance(dt);
+            if rh.poisson_iterations != rd.poisson_iterations {
+                mismatched += 1;
+            }
+        }
+        let mut fd = FlowField3D::new(g);
+        device.download(&mut fd);
+        let (worst, scale) = compare(&fh, &fd);
+        println!(
+            "  Poiseuille n 16 nz {nz} (tight {tight}): 300 steps, host vs device max |Δ| {worst:.3e} on {scale:.3e}; CG counts differ on {mismatched} steps"
+        );
+        if tight {
+            assert!(
+                worst <= 1e-11 * scale,
+                "nz {nz}: {worst:.3e} of {scale:.3e}"
+            );
+        }
+        if nz > 1 {
+            let plane =
+                |f: &FlowField3D, k: usize| f.u[k * n * (n + 1)..(k + 1) * n * (n + 1)].to_vec();
+            let p0 = plane(&fd, 0);
+            let mut worst_plane = 0.0_f64;
+            for k in 1..nz {
+                for (a, b) in plane(&fd, k).iter().zip(&p0) {
+                    worst_plane = worst_plane.max((a - b).abs());
+                }
+            }
+            println!("  Poiseuille nz {nz}: device planes within {worst_plane:.3e} of {scale:.3e}");
+            assert!(worst_plane <= 1e-11 * scale);
+        }
+    }
+}
+
+/// Diagnostic: where and when the Poiseuille host/device difference enters.
+#[test]
+#[ignore = "diagnostic: per-component host/device differences on Poiseuille variants"]
+fn poiseuille_difference_diagnostic() {
+    let n = 16;
+    let h = 1.0 / n as f64;
+    let dt = 0.4 * (h * h / (4.0 * MU_P)).min(h);
+    for (label, dz, inlet, source) in [
+        ("as is (dz 1, inlet profile, source G)", 1.0, true, true),
+        ("dz = h", h, true, true),
+        ("closed box (no inlet), source G", 1.0, false, true),
+        ("inlet profile, no source", 1.0, true, false),
+    ] {
+        let g = Grid3 {
+            nx: n,
+            ny: n,
+            nz: 1,
+            dx: h,
+            dy: h,
+            dz,
+        };
+        let mk = || {
+            let mut s = Piso3Solver::new(
+                fluid(MU_P),
+                params(ConvectionScheme::Upwind, SideBoundary3::SlipWall, true),
+            );
+            if source {
+                s.set_momentum_source(|_x, _y, _z, _t| (G, 0.0, 0.0));
+            }
+            let u_hat = discrete_profile(n);
+            s.set_boundary_velocity(move |x, y, _z, _t| {
+                if inlet && (x <= 0.0 || x >= 1.0) {
+                    let j = ((y / h - 0.5).round().max(0.0) as usize).min(n - 1);
+                    (u_hat[j], 0.0, 0.0)
+                } else {
+                    (0.0, 0.0, 0.0)
+                }
+            });
+            s
+        };
+        let mut host = mk();
+        let mut fh = FlowField3D::new(g);
+        if inlet {
+            let u_hat = discrete_profile(n);
+            for (j, &uj) in u_hat.iter().enumerate() {
+                fh.u[g.uface(0, j, 0)] = uj;
+                fh.u[g.uface(0, j, n)] = uj;
+            }
+        }
+        let mut device = Piso3Device::new(mk(), g);
+        device.upload(&fh);
+        let mut fd = FlowField3D::new(g);
+        for step in 1..=300 {
+            let rh = host.advance(&mut fh, dt);
+            let rd = device.advance(dt);
+            if [1, 2, 10, 100, 300].contains(&step) {
+                device.download(&mut fd);
+                let du =
+                    fh.u.iter()
+                        .zip(&fd.u)
+                        .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
+                let dv =
+                    fh.v.iter()
+                        .zip(&fd.v)
+                        .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
+                let dp =
+                    fh.p.iter()
+                        .zip(&fd.p)
+                        .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
+                let dpp = fh
+                    .p_prime
+                    .iter()
+                    .zip(&fd.p_prime)
+                    .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
+                let dsp = fh
+                    .sp
+                    .iter()
+                    .zip(&fd.sp)
+                    .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
+                println!(
+                    "  {label} step {step}: Δu {du:.2e} Δv {dv:.2e} Δp {dp:.2e} Δp' {dpp:.2e} Δsp {dsp:.2e}; CG it host {} / dev {}; correctors {} / {}; residual {:.2e} / {:.2e}",
+                    rh.poisson_iterations,
+                    rd.poisson_iterations,
+                    rh.corrector_steps_performed,
+                    rd.corrector_steps_performed,
+                    rh.final_residual,
+                    rd.final_residual
+                );
+            }
+        }
+    }
+}
+
+/// The device step's cost at the anchor size (378 × 62 × 62, cubic cells,
+/// the manufactured source, red-black + device CG), `RTX_PROFILE=1` for the
+/// phase split. Recorded, not asserted.
+#[test]
+#[ignore = "bench: ms per device step at 378×62×62"]
+fn bench_device_step_anchor_size() {
+    let (nx, ny, nz) = (378usize, 62usize, 62usize);
+    let h = 0.41 / ny as f64;
+    let mu = 1.0e-3;
+    let g = Grid3 {
+        nx,
+        ny,
+        nz,
+        dx: h,
+        dy: h,
+        dz: h,
+    };
+    let dt = 0.4 * (h * h / (4.0 * mu / RHO)).min(h / 2.0);
+    let mut s = Piso3Solver::new(
+        fluid(mu),
+        params(
+            ConvectionScheme::TvdVanAlbada,
+            SideBoundary3::Velocity,
+            false,
+        ),
+    );
+    s.set_momentum_source(move |x, y, z, _t| source3(mu, x / 2.5, y / 0.41, z / 0.41));
+    s.set_boundary_velocity(|x, y, z, _t| boundary3(x / 2.5, y / 0.41, z / 0.41));
+    let mut device = Piso3Device::new(s, g);
+    let f = FlowField3D::new(g);
+    device.upload(&f);
+    device.advance(dt);
+    let t0 = std::time::Instant::now();
+    let steps = 20;
+    let mut it = 0;
+    for _ in 0..steps {
+        it += device.advance(dt).poisson_iterations;
+    }
+    let ms = t0.elapsed().as_secs_f64() * 1e3 / steps as f64;
+    println!(
+        "  device step at {nx}×{ny}×{nz} ({} cells): {ms:.1} ms per step, {:.1} CG iterations per step",
+        g.cells(),
+        it as f64 / steps as f64
+    );
+    if let Some(t) = device.timers() {
+        let n = t.steps.max(1) as f64;
+        println!(
+            "  split per step: predictor {:.1} ms, poisson {:.1} ms, apply {:.1} ms ({} steps timed)",
+            t.predictor_ns as f64 / n / 1e6,
+            t.poisson_ns as f64 / n / 1e6,
+            t.apply_ns as f64 / n / 1e6,
+            t.steps
+        );
+    }
+}