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]>

crates/specialized/rtx-cfd/src/kernels/cuda/cg3.cu

@@ -0,0 +1,181 @@
+/**
+ * 3D Stage 1, gate 3: the f64 conjugate gradient's maps on the device over
+ * the active cells (index list), with FIXED-ORDER reductions: every block
+ * reduces its 256 lanes by the same shared-memory tree and writes one
+ * partial; `cg3_reduce` sums the partials in index order on one thread.
+ * Run-to-run bit-identical by construction (no atomics).
+ */
+#define NONE 0xFFFFFFFFu
+
+__device__ __forceinline__ double nb_sum3d(
+    int g, int nx, const double* ae, const double* aw, const double* an, const double* as_,
+    const double* at, const double* ab, const unsigned int* top, const unsigned int* bot,
+    const double* x)
+{
+    double s = 0.0;
+    double e = ae[g]; if (e != 0.0) s += e * x[g + 1];
+    double w = aw[g]; if (w != 0.0) s += w * x[g - 1];
+    double nn = an[g]; if (nn != 0.0) s += nn * x[g + nx];
+    double ss = as_[g]; if (ss != 0.0) s += ss * x[g - nx];
+    double t = at[g]; if (t != 0.0) s += t * x[top[g]];
+    double b = ab[g]; if (b != 0.0) s += b * x[bot[g]];
+    return s;
+}
+
+__device__ __forceinline__ double block_reduce(double v)
+{
+    __shared__ double sh[256];
+    int t = threadIdx.x;
+    sh[t] = v;
+    __syncthreads();
+    for (int s = 128; s > 0; s >>= 1) {
+        if (t < s) sh[t] += sh[t + s];
+        __syncthreads();
+    }
+    return sh[0];
+}
+
+/* q = A d on the cells. */
+extern "C" __global__ void cg3_spmv(
+    int n_cells, const unsigned int* __restrict__ cells,
+    const double* __restrict__ ae, const double* __restrict__ aw,
+    const double* __restrict__ an, const double* __restrict__ as_,
+    const double* __restrict__ at, const double* __restrict__ ab,
+    const unsigned int* __restrict__ top, const unsigned int* __restrict__ bot,
+    const double* __restrict__ ap, const double* __restrict__ d, double* __restrict__ q, int nx)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int g = cells[t];
+    q[g] = ap[g] * d[g] - nb_sum3d(g, nx, ae, aw, an, as_, at, ab, top, bot, d);
+}
+
+/* r = b − A p on the cells; partial[block] = Σ |r| over the block's cells. */
+extern "C" __global__ void cg3_residual(
+    int n_cells, const unsigned int* __restrict__ cells,
+    const double* __restrict__ ae, const double* __restrict__ aw,
+    const double* __restrict__ an, const double* __restrict__ as_,
+    const double* __restrict__ at, const double* __restrict__ ab,
+    const unsigned int* __restrict__ top, const unsigned int* __restrict__ bot,
+    const double* __restrict__ ap, const double* __restrict__ b,
+    const double* __restrict__ p, double* __restrict__ r, double* __restrict__ partial, int nx)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    double v = 0.0;
+    if (t < n_cells) {
+        int g = cells[t];
+        v = b[g] - (ap[g] * p[g] - nb_sum3d(g, nx, ae, aw, an, as_, at, ab, top, bot, p));
+        r[g] = v;
+        v = fabs(v);
+    }
+    double s = block_reduce(v);
+    if (threadIdx.x == 0) partial[blockIdx.x] = s;
+}
+
+/* partial[block] = Σ a·b over the block's cells (b = a for a norm). */
+extern "C" __global__ void cg3_dot_partial(
+    int n_cells, const unsigned int* __restrict__ cells,
+    const double* __restrict__ a, const double* __restrict__ b, double* __restrict__ partial)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    double v = 0.0;
+    if (t < n_cells) { int g = cells[t]; v = a[g] * b[g]; }
+    double s = block_reduce(v);
+    if (threadIdx.x == 0) partial[blockIdx.x] = s;
+}
+
+/* partial[block] = Σ |a| over the block's cells. */
+extern "C" __global__ void cg3_l1_partial(
+    int n_cells, const unsigned int* __restrict__ cells,
+    const double* __restrict__ a, double* __restrict__ partial)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    double v = 0.0;
+    if (t < n_cells) v = fabs(a[cells[t]]);
+    double s = block_reduce(v);
+    if (threadIdx.x == 0) partial[blockIdx.x] = s;
+}
+
+/* partial[block] = Σ a over the block's cells. */
+extern "C" __global__ void cg3_sum_partial(
+    int n_cells, const unsigned int* __restrict__ cells,
+    const double* __restrict__ a, double* __restrict__ partial)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    double v = 0.0;
+    if (t < n_cells) v = a[cells[t]];
+    double s = block_reduce(v);
+    if (threadIdx.x == 0) partial[blockIdx.x] = s;
+}
+
+/* out[0] = Σ partial[0..n) in index order, one thread. */
+extern "C" __global__ void cg3_reduce(int n, const double* __restrict__ partial, double* __restrict__ out)
+{
+    if (blockIdx.x * blockDim.x + threadIdx.x != 0) return;
+    double s = 0.0;
+    for (int i = 0; i < n; ++i) s += partial[i];
+    out[0] = s;
+}
+
+/* p += alpha d; r -= alpha q. */
+extern "C" __global__ void cg3_axpy2(
+    int n_cells, const unsigned int* __restrict__ cells, double alpha,
+    const double* __restrict__ d, const double* __restrict__ q,
+    double* __restrict__ p, double* __restrict__ r)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int g = cells[t];
+    p[g] += alpha * d[g];
+    r[g] -= alpha * q[g];
+}
+
+/* d = z + beta d. */
+extern "C" __global__ void cg3_xpay(
+    int n_cells, const unsigned int* __restrict__ cells, double beta,
+    const double* __restrict__ z, double* __restrict__ d)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int g = cells[t];
+    d[g] = z[g] + beta * d[g];
+}
+
+/* dst = src (copy over the cells). */
+extern "C" __global__ void cg3_copy(
+    int n_cells, const unsigned int* __restrict__ cells,
+    const double* __restrict__ src, double* __restrict__ dst)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int g = cells[t];
+    dst[g] = src[g];
+}
+
+/* v -= s over the cells. */
+extern "C" __global__ void cg3_shift(
+    int n_cells, const unsigned int* __restrict__ cells, double s, double* __restrict__ v)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    v[cells[t]] -= s;
+}
+
+/* f32 b0 = (float) r over the cells; f64 z = (double) x0 over the cells. */
+extern "C" __global__ void cg3_gather_f32(
+    int n_cells, const unsigned int* __restrict__ cells, const double* __restrict__ r, float* __restrict__ b0)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int g = cells[t];
+    b0[g] = (float) r[g];
+}
+
+extern "C" __global__ void cg3_scatter_f64(
+    int n_cells, const unsigned int* __restrict__ cells, const float* __restrict__ x0, double* __restrict__ z)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int g = cells[t];
+    z[g] = (double) x0[g];
+}

crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/poisson/device_cg.rs

@@ -0,0 +1,646 @@
+//! Gate 3: the f64 conjugate gradient entirely on the device — the 2D
+//! `run_pcg`'s control flow (true-residual stop and resynchronisation,
+//! breakdown guard, mean projection on a singular system, anchor shift on
+//! exit) with the device V-cycle as its preconditioner; three scalars come
+//! back per iteration. Reductions are fixed-order (run-to-run identical).
+//! Stage 1 scope: at most one singular component (asserted at build).
+
+use super::device::{DeviceVcycle3, runtime3};
+use super::{Components3, Level3, OperatorKey3, PoissonProblem3D, export_hierarchy3};
+use crate::solvers::incompressible::poisson::{MultigridParameters, PoissonSolution};
+use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, LaunchConfig, PushKernelArg};
+use cudarc::nvrtc::{CompileOptions, compile_ptx_with_opts};
+use std::sync::{Arc, OnceLock};
+
+const KERNELS: &str = include_str!("../../../../kernels/cuda/cg3.cu");
+
+struct CgKernels {
+    _module: Arc<CudaModule>,
+    spmv: CudaFunction,
+    residual: CudaFunction,
+    dot_partial: CudaFunction,
+    l1_partial: CudaFunction,
+    sum_partial: CudaFunction,
+    reduce: CudaFunction,
+    axpy2: CudaFunction,
+    xpay: CudaFunction,
+    copy: CudaFunction,
+    shift: CudaFunction,
+    gather_f32: CudaFunction,
+    scatter_f64: CudaFunction,
+}
+
+static CG: OnceLock<CgKernels> = OnceLock::new();
+
+fn kernels() -> &'static CgKernels {
+    CG.get_or_init(|| {
+        let rt = runtime3();
+        let arch = std::env::var("RTX_CUDA_ARCH").unwrap_or_else(|_| "sm_120".to_string());
+        let ptx = compile_ptx_with_opts(
+            KERNELS,
+            CompileOptions {
+                arch: Some(Box::leak(arch.into_boxed_str())),
+                ..Default::default()
+            },
+        )
+        .expect("nvrtc: cg3.cu");
+        let module = rt.ctx.load_module(ptx).expect("cg3 module");
+        let f = |name: &str| module.load_function(name).expect(name);
+        CgKernels {
+            spmv: f("cg3_spmv"),
+            residual: f("cg3_residual"),
+            dot_partial: f("cg3_dot_partial"),
+            l1_partial: f("cg3_l1_partial"),
+            sum_partial: f("cg3_sum_partial"),
+            reduce: f("cg3_reduce"),
+            axpy2: f("cg3_axpy2"),
+            xpay: f("cg3_xpay"),
+            copy: f("cg3_copy"),
+            shift: f("cg3_shift"),
+            gather_f32: f("cg3_gather_f32"),
+            scatter_f64: f("cg3_scatter_f64"),
+            _module: module,
+        }
+    })
+}
+
+fn cfg(n_items: usize) -> LaunchConfig {
+    LaunchConfig {
+        grid_dim: ((n_items as u32).div_ceil(256).max(1), 1, 1),
+        block_dim: (256, 1, 1),
+        shared_mem_bytes: 0,
+    }
+}
+
+/// The prepared operator on the device and the CG's work vectors.
+pub struct DeviceCg3 {
+    key: OperatorKey3,
+    n: usize,
+    nx: i32,
+    n_cells: usize,
+    n_blocks: usize,
+    cells: CudaSlice<u32>,
+    top: CudaSlice<u32>,
+    bot: CudaSlice<u32>,
+    ae: CudaSlice<f64>,
+    aw: CudaSlice<f64>,
+    an: CudaSlice<f64>,
+    as_: CudaSlice<f64>,
+    at: CudaSlice<f64>,
+    ab: CudaSlice<f64>,
+    ap: CudaSlice<f64>,
+    b: CudaSlice<f64>,
+    r: CudaSlice<f64>,
+    z: CudaSlice<f64>,
+    d: CudaSlice<f64>,
+    q: CudaSlice<f64>,
+    partial: CudaSlice<f64>,
+    scalar: CudaSlice<f64>,
+    vcycle: DeviceVcycle3,
+    /// The singular component's members (all cells when singular; empty
+    /// otherwise) and whether the system is singular.
+    singular: bool,
+    active_host: Vec<bool>,
+    max_iterations: usize,
+    scalar_host: Vec<f64>,
+}
+
+impl DeviceCg3 {
+    /// Builds the device operator (a hierarchy build on the host, one
+    /// upload). Panics outside Stage 1's scope: more than one component
+    /// with a singular one among them.
+    pub fn new(problem: &PoissonProblem3D, params: &MultigridParameters) -> Self {
+        let rt = runtime3();
+        let fine = Level3::<f64>::new(problem.clone());
+        let components = Components3::find(problem, &fine.cells);
+        let singular_count = components.singular.iter().filter(|&&s| s).count();
+        assert!(
+            singular_count == 0 || components.members.len() == 1,
+            "DeviceCg3 (Stage 1): {} components with {} singular; one singular component is the scope",
+            components.members.len(),
+            singular_count
+        );
+        let levels = export_hierarchy3(problem, params);
+        let vcycle = DeviceVcycle3::new(&levels, params.smoother_sweeps.max(1));
+        let n = problem.nx * problem.ny * problem.nz;
+        let n_cells = fine.cells.len();
+        let n_blocks = n_cells.div_ceil(256).max(1);
+        let to_u32 = |v: &[usize]| {
+            v.iter()
+                .map(|&i| if i == usize::MAX { u32::MAX } else { i as u32 })
+                .collect::<Vec<u32>>()
+        };
+        let up_u = |v: &[u32]| -> CudaSlice<u32> {
+            rt.stream
+                .memcpy_stod(if v.is_empty() { &[0u32][..] } else { v })
+                .expect("upload")
+        };
+        let up_f = |v: &[f64]| -> CudaSlice<f64> { rt.stream.memcpy_stod(v).expect("upload") };
+        let zeros = || rt.stream.alloc_zeros::<f64>(n).expect("alloc");
+        Self {
+            key: OperatorKey3::of(problem, params),
+            n,
+            nx: problem.nx as i32,
+            n_cells,
+            n_blocks,
+            cells: up_u(&to_u32(&fine.cells)),
+            top: up_u(&to_u32(&fine.top)),
+            bot: up_u(&to_u32(&fine.bot)),
+            ae: up_f(&fine.ae),
+            aw: up_f(&fine.aw),
+            an: up_f(&fine.an),
+            as_: up_f(&fine.as_),
+            at: up_f(&fine.at),
+            ab: up_f(&fine.ab),
+            ap: up_f(&fine.ap),
+            b: zeros(),
+            r: zeros(),
+            z: zeros(),
+            d: zeros(),
+            q: zeros(),
+            partial: rt.stream.alloc_zeros::<f64>(n_blocks).expect("alloc"),
+            scalar: rt.stream.alloc_zeros::<f64>(1).expect("alloc"),
+            vcycle,
+            singular: singular_count > 0,
+            active_host: fine.active.clone(),
+            max_iterations: params.max_iterations,
+            scalar_host: vec![0.0],
+        }
+    }
+
+    /// Whether this prepared operator is the one for `problem` + `params`.
+    pub fn matches(&self, problem: &PoissonProblem3D, params: &MultigridParameters) -> bool {
+        self.key.matches(problem, params)
+    }
+
+    pub fn n_cells(&self) -> usize {
+        self.n_cells
+    }
+
+    fn reduce(&mut self) -> f64 {
+        let rt = runtime3();
+        let k = kernels();
+        let nb = self.n_blocks as i32;
+        unsafe {
+            rt.stream
+                .launch_builder(&k.reduce)
+                .arg(&nb)
+                .arg(&self.partial)
+                .arg(&mut self.scalar)
+                .launch(LaunchConfig {
+                    grid_dim: (1, 1, 1),
+                    block_dim: (32, 1, 1),
+                    shared_mem_bytes: 0,
+                })
+                .expect("cg3_reduce");
+        }
+        rt.stream
+            .memcpy_dtoh(&self.scalar, &mut self.scalar_host)
+            .expect("scalar");
+        rt.stream.synchronize().expect("sync");
+        self.scalar_host[0]
+    }
+
+    fn dot(&mut self, a: Which, b: Which) -> f64 {
+        let rt = runtime3();
+        let k = kernels();
+        let n_i = self.n_cells as i32;
+        let n_cells = self.n_cells;
+        let Self {
+            cells,
+            partial,
+            r,
+            z,
+            d,
+            q,
+            b: bb,
+            ..
+        } = self;
+        let pa = pick(a, r, z, d, q, bb);
+        let pb = pick(b, r, z, d, q, bb);
+        unsafe {
+            rt.stream
+                .launch_builder(&k.dot_partial)
+                .arg(&n_i)
+                .arg(&*cells)
+                .arg(pa)
+                .arg(pb)
+                .arg(&mut *partial)
+                .launch(cfg(n_cells))
+                .expect("cg3_dot_partial");
+        }
+        self.reduce()
+    }
+
+    fn l1(&mut self, a: Which) -> f64 {
+        let rt = runtime3();
+        let k = kernels();
+        let n_i = self.n_cells as i32;
+        let n_cells = self.n_cells;
+        let Self {
+            cells,
+            partial,
+            r,
+            z,
+            d,
+            q,
+            b: bb,
+            ..
+        } = self;
+        let pa = pick(a, r, z, d, q, bb);
+        unsafe {
+            rt.stream
+                .launch_builder(&k.l1_partial)
+                .arg(&n_i)
+                .arg(&*cells)
+                .arg(pa)
+                .arg(&mut *partial)
+                .launch(cfg(n_cells))
+                .expect("cg3_l1_partial");
+        }
+        self.reduce()
+    }
+
+    fn mean(&mut self, a: Which) -> f64 {
+        let rt = runtime3();
+        let k = kernels();
+        let n_i = self.n_cells as i32;
+        let n_cells = self.n_cells;
+        let Self {
+            cells,
+            partial,
+            r,
+            z,
+            d,
+            q,
+            b: bb,
+            ..
+        } = self;
+        let pa = pick(a, r, z, d, q, bb);
+        unsafe {
+            rt.stream
+                .launch_builder(&k.sum_partial)
+                .arg(&n_i)
+                .arg(&*cells)
+                .arg(pa)
+                .arg(&mut *partial)
+                .launch(cfg(n_cells))
+                .expect("cg3_sum_partial");
+        }
+        self.reduce() / self.n_cells as f64
+    }
+
+    fn shift(&mut self, a: Which, s: f64) {
+        let rt = runtime3();
+        let k = kernels();
+        let n_i = self.n_cells as i32;
+        let n_cells = self.n_cells;
+        let Self {
+            cells,
+            r,
+            z,
+            d,
+            q,
+            b: bb,
+            ..
+        } = self;
+        let pa = pick_mut(a, r, z, d, q, bb);
+        unsafe {
+            rt.stream
+                .launch_builder(&k.shift)
+                .arg(&n_i)
+                .arg(&*cells)
+                .arg(&s)
+                .arg(pa)
+                .launch(cfg(n_cells))
+                .expect("cg3_shift");
+        }
+    }
+
+    /// Project the mean out of `a` (the singular case).
+    fn project_mean(&mut self, a: Which) {
+        if self.singular {
+            let m = self.mean(a);
+            self.shift(a, m);
+        }
+    }
+
+    /// `r = b − A p`, returns `Σ |r|`.
+    fn true_residual(&mut self, p: &CudaSlice<f64>) -> f64 {
+        let rt = runtime3();
+        let k = kernels();
+        let n_i = self.n_cells as i32;
+        unsafe {
+            rt.stream
+                .launch_builder(&k.residual)
+                .arg(&n_i)
+                .arg(&self.cells)
+                .arg(&self.ae)
+                .arg(&self.aw)
+                .arg(&self.an)
+                .arg(&self.as_)
+                .arg(&self.at)
+                .arg(&self.ab)
+                .arg(&self.top)
+                .arg(&self.bot)
+                .arg(&self.ap)
+                .arg(&self.b)
+                .arg(p)
+                .arg(&mut self.r)
+                .arg(&mut self.partial)
+                .arg(&self.nx)
+                .launch(cfg(self.n_cells))
+                .expect("cg3_residual");
+        }
+        self.reduce()
+    }
+
+    /// `z = M⁻¹ r` (gather to f32, the device V-cycle, scatter to f64).
+    fn precondition(&mut self) {
+        let rt = runtime3();
+        let k = kernels();
+        let n_i = self.n_cells as i32;
+        unsafe {
+            rt.stream
+                .launch_builder(&k.gather_f32)
+                .arg(&n_i)
+                .arg(&self.cells)
+                .arg(&self.r)
+                .arg(&mut self.vcycle.levels[0].b)
+                .launch(cfg(self.n_cells))
+                .expect("cg3_gather_f32");
+        }
+        self.vcycle.vcycle_on_device();
+        unsafe {
+            rt.stream
+                .launch_builder(&k.scatter_f64)
+                .arg(&n_i)
+                .arg(&self.cells)
+                .arg(&self.vcycle.levels[0].x)
+                .arg(&mut self.z)
+                .launch(cfg(self.n_cells))
+                .expect("cg3_scatter_f64");
+        }
+    }
+
+    /// The CG on the device: `rhs` and `p` are device vectors of `n`
+    /// entries (inactive entries untouched). The 2D `run_pcg` sequence.
+    pub fn solve_device(
+        &mut self,
+        rhs: &CudaSlice<f64>,
+        p: &mut CudaSlice<f64>,
+        tolerance: f64,
+        anchor: Option<usize>,
+        setup_ns: u64,
+    ) -> PoissonSolution {
+        let rt = runtime3();
+        let k = kernels();
+        let n_i = self.n_cells as i32;
+        let t_iter = std::time::Instant::now();
+        if self.n_cells == 0 {
+            return PoissonSolution {
+                iterations: 0,
+                residual: 0.0,
+                converged: true,
+                setup_ns,
+                iterate_ns: 0,
+            };
+        }
+        // b = rhs on the cells, mean projected out when singular.
+        unsafe {
+            rt.stream
+                .launch_builder(&k.copy)
+                .arg(&n_i)
+                .arg(&self.cells)
+                .arg(rhs)
+                .arg(&mut self.b)
+                .launch(cfg(self.n_cells))
+                .expect("cg3_copy");
+        }
+        self.project_mean(Which::B);
+        let anchor = anchor.filter(|&a| a < self.n && self.active_host[a]);
+        let finish = |this: &mut Self, p: &mut CudaSlice<f64>, iterations: usize, residual: f64| {
+            if this.singular {
+                let shift = match anchor {
+                    Some(a) => {
+                        let mut one = vec![0.0f64];
+                        let view = p.slice(a..a + 1);
+                        rt.stream.memcpy_dtoh(&view, &mut one).expect("anchor");
+                        rt.stream.synchronize().expect("sync");
+                        one[0]
+                    }
+                    None => {
+                        // The mean of p over the cells.
+                        let n_i = this.n_cells as i32;
+                        unsafe {
+                            rt.stream
+                                .launch_builder(&k.sum_partial)
+                                .arg(&n_i)
+                                .arg(&this.cells)
+                                .arg(&*p)
+                                .arg(&mut this.partial)
+                                .launch(cfg(this.n_cells))
+                                .expect("cg3_sum_partial");
+                        }
+                        this.reduce() / this.n_cells as f64
+                    }
+                };
+                let n_i = this.n_cells as i32;
+                unsafe {
+                    rt.stream
+                        .launch_builder(&k.shift)
+                        .arg(&n_i)
+                        .arg(&this.cells)
+                        .arg(&shift)
+                        .arg(&mut *p)
+                        .launch(cfg(this.n_cells))
+                        .expect("cg3_shift");
+                }
+            }
+            PoissonSolution {
+                iterations,
+                residual,
+                converged: residual < tolerance,
+                setup_ns,
+                iterate_ns: t_iter.elapsed().as_nanos() as u64,
+            }
+        };
+
+        let mut res = self.true_residual(p);
+        if res < tolerance {
+            return finish(self, p, 0, res);
+        }
+        self.precondition();
+        self.project_mean(Which::Z);
+        unsafe {
+            rt.stream
+                .launch_builder(&k.copy)
+                .arg(&n_i)
+                .arg(&self.cells)
+                .arg(&self.z)
+                .arg(&mut self.d)
+                .launch(cfg(self.n_cells))
+                .expect("cg3_copy");
+        }
+        let mut rz = self.dot(Which::R, Which::Z);
+        let mut iterations = 0;
+        let mut last_true = res;
+        while iterations < self.max_iterations {
+            iterations += 1;
+            unsafe {
+                rt.stream
+                    .launch_builder(&k.spmv)
+                    .arg(&n_i)
+                    .arg(&self.cells)
+                    .arg(&self.ae)
+                    .arg(&self.aw)
+                    .arg(&self.an)
+                    .arg(&self.as_)
+                    .arg(&self.at)
+                    .arg(&self.ab)
+                    .arg(&self.top)
+                    .arg(&self.bot)
+                    .arg(&self.ap)
+                    .arg(&self.d)
+                    .arg(&mut self.q)
+                    .arg(&self.nx)
+                    .launch(cfg(self.n_cells))
+                    .expect("cg3_spmv");
+            }
+            let dq = self.dot(Which::D, Which::Q);
+            if !dq.is_finite() || dq <= 0.0 || !rz.is_finite() || rz <= 0.0 {
+                res = self.true_residual(p);
+                return finish(self, p, iterations, res);
+            }
+            let alpha = rz / dq;
+            unsafe {
+                rt.stream
+                    .launch_builder(&k.axpy2)
+                    .arg(&n_i)
+                    .arg(&self.cells)
+                    .arg(&alpha)
+                    .arg(&self.d)
+                    .arg(&self.q)
+                    .arg(&mut *p)
+                    .arg(&mut self.r)
+                    .launch(cfg(self.n_cells))
+                    .expect("cg3_axpy2");
+            }
+            if self.l1(Which::R) < tolerance {
+                res = self.true_residual(p);
+                if res < tolerance || res > 0.9 * last_true {
+                    return finish(self, p, iterations, res);
+                }
+                last_true = res;
+            }
+            self.precondition();
+            self.project_mean(Which::Z);
+            let rz_new = self.dot(Which::R, Which::Z);
+            let beta = rz_new / rz;
+            rz = rz_new;
+            unsafe {
+                rt.stream
+                    .launch_builder(&k.xpay)
+                    .arg(&n_i)
+                    .arg(&self.cells)
+                    .arg(&beta)
+                    .arg(&self.z)
+                    .arg(&mut self.d)
+                    .launch(cfg(self.n_cells))
+                    .expect("cg3_xpay");
+            }
+        }
+        res = self.true_residual(p);
+        finish(self, p, iterations, res)
+    }
+
+    /// Host-facing solve: uploads `problem.rhs` and `p`, solves, downloads `p`.
+    pub fn solve_host(
+        &mut self,
+        problem: &PoissonProblem3D,
+        p: &mut [f64],
+        tolerance: f64,
+        anchor: Option<usize>,
+        setup_ns: u64,
+    ) -> PoissonSolution {
+        let rt = runtime3();
+        assert_eq!(p.len(), self.n);
+        let rhs_dev: CudaSlice<f64> = rt.stream.memcpy_stod(&problem.rhs).expect("rhs");
+        let mut p_dev: CudaSlice<f64> = rt.stream.memcpy_stod(p).expect("p");
+        let sol = self.solve_device(&rhs_dev, &mut p_dev, tolerance, anchor, setup_ns);
+        rt.stream.memcpy_dtoh(&p_dev, p).expect("p down");
+        rt.stream.synchronize().expect("sync");
+        sol
+    }
+}
+
+fn pick<'a>(
+    w: Which,
+    r: &'a CudaSlice<f64>,
+    z: &'a CudaSlice<f64>,
+    d: &'a CudaSlice<f64>,
+    q: &'a CudaSlice<f64>,
+    b: &'a CudaSlice<f64>,
+) -> &'a CudaSlice<f64> {
+    match w {
+        Which::R => r,
+        Which::Z => z,
+        Which::D => d,
+        Which::Q => q,
+        Which::B => b,
+    }
+}
+
+fn pick_mut<'a>(
+    w: Which,
+    r: &'a mut CudaSlice<f64>,
+    z: &'a mut CudaSlice<f64>,
+    d: &'a mut CudaSlice<f64>,
+    q: &'a mut CudaSlice<f64>,
+    b: &'a mut CudaSlice<f64>,
+) -> &'a mut CudaSlice<f64> {
+    match w {
+        Which::R => r,
+        Which::Z => z,
+        Which::D => d,
+        Which::Q => q,
+        Which::B => b,
+    }
+}
+
+#[derive(Clone, Copy)]
+enum Which {
+    R,
+    Z,
+    D,
+    Q,
+    B,
+}
+
+/// One prepared device operator, reused while the operator is bit-identical.
+#[derive(Default)]
+pub struct DevicePcgCache3 {
+    slot: Option<DeviceCg3>,
+}
+
+/// The device counterpart of `solve_multigrid_pcg3_cached`.
+pub fn solve_multigrid_pcg3_device_cached(
+    problem: &PoissonProblem3D,
+    p: &mut [f64],
+    params: &MultigridParameters,
+    tolerance: f64,
+    anchor: Option<usize>,
+    cache: &mut DevicePcgCache3,
+) -> PoissonSolution {
+    let t_entry = std::time::Instant::now();
+    let hit = cache
+        .slot
+        .as_ref()
+        .is_some_and(|cg| cg.matches(problem, params));
+    if !hit {
+        cache.slot = Some(DeviceCg3::new(problem, params));
+    }
+    let setup_ns = t_entry.elapsed().as_nanos() as u64;
+    let cg = cache.slot.as_mut().expect("prepared");
+    cg.solve_host(problem, p, tolerance, anchor, setup_ns)
+}

crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/poisson/mod.rs

@@ -13,6 +13,8 @@ use crate::solvers::incompressible::poisson::{
 
 #[cfg(feature = "cuda")]
 pub mod device;
+#[cfg(feature = "cuda")]
+pub mod device_cg;
 pub mod export;
 pub use export::{LevelExport3, export_hierarchy3, vcycle_f32_reference3};
 

crates/specialized/rtx-cfd/tests/three_d_cg_device.rs

@@ -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}");
+}