rtx-cfd embedded3 item 3: e3_cg.cu (FMA off) + poisson::device_cg::{DeviceCg, DevicePcgCache, solve_pcg_device_cached}; gate 3 HELD: device CG = host PCG (9 / 9 iterations, |Δp| ≤ 1.5e-14 relative, run-to-run bit-identical, anchor exact, mean 9e-16)

Co-Authored-By: Claude Fable 5.1 <[email protected]>
2026-09-17 14:53:20 -05:00
parent 5e2b565971
commit 765ba6d3f6
4 changed files with 988 additions and 0 deletions
@@ -0,0 +1,627 @@
+//! embedded3 item 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::{DeviceVcycle, cfg, load_module, runtime};
+use super::{Components, Level, OperatorKey, Problem, export_hierarchy};
+use crate::solvers::incompressible::poisson::{MultigridParameters, PoissonSolution};
+use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, LaunchConfig, PushKernelArg};
+use std::sync::{Arc, OnceLock};
+
+const KERNELS: &str = include_str!("../../../../kernels/cuda/e3_cg.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 module = load_module(KERNELS, "e3_cg.cu", true);
+        let f = |name: &str| module.load_function(name).expect(name);
+        CgKernels {
+            spmv: f("e3_cg_spmv"),
+            residual: f("e3_cg_residual"),
+            dot_partial: f("e3_cg_dot_partial"),
+            l1_partial: f("e3_cg_l1_partial"),
+            sum_partial: f("e3_cg_sum_partial"),
+            reduce: f("e3_cg_reduce"),
+            axpy2: f("e3_cg_axpy2"),
+            xpay: f("e3_cg_xpay"),
+            copy: f("e3_cg_copy"),
+            shift: f("e3_cg_shift"),
+            gather_f32: f("e3_cg_gather_f32"),
+            scatter_f64: f("e3_cg_scatter_f64"),
+            _module: module,
+        }
+    })
+}
+
+/// The prepared operator on the device and the CG's work vectors.
+pub struct DeviceCg {
+    key: OperatorKey,
+    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: DeviceVcycle,
+    /// 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 DeviceCg {
+    /// 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: &Problem, params: &MultigridParameters) -> Self {
+        let rt = runtime();
+        let fine = Level::<f64>::new(problem.clone());
+        let components = Components::find(problem, &fine.cells);
+        let singular_count = components.singular.iter().filter(|&&s| s).count();
+        assert!(
+            singular_count == 0 || components.members.len() == 1,
+            "DeviceCg: {} components with {} singular; one singular component is the scope",
+            components.members.len(),
+            singular_count
+        );
+        let levels = export_hierarchy(problem, params);
+        let vcycle = DeviceVcycle::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: OperatorKey::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: &Problem, params: &MultigridParameters) -> bool {
+        self.key.matches(problem, params)
+    }
+
+    pub fn n_cells(&self) -> usize {
+        self.n_cells
+    }
+
+    fn reduce(&mut self) -> f64 {
+        let rt = runtime();
+        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("e3_cg_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 = runtime();
+        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("e3_cg_dot_partial");
+        }
+        self.reduce()
+    }
+
+    fn l1(&mut self, a: Which) -> f64 {
+        let rt = runtime();
+        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("e3_cg_l1_partial");
+        }
+        self.reduce()
+    }
+
+    fn mean(&mut self, a: Which) -> f64 {
+        let rt = runtime();
+        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("e3_cg_sum_partial");
+        }
+        self.reduce() / self.n_cells as f64
+    }
+
+    fn shift(&mut self, a: Which, s: f64) {
+        let rt = runtime();
+        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("e3_cg_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 = runtime();
+        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("e3_cg_residual");
+        }
+        self.reduce()
+    }
+
+    /// `z = M⁻¹ r` (gather to f32, the device V-cycle, scatter to f64).
+    fn precondition(&mut self) {
+        let rt = runtime();
+        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("e3_cg_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("e3_cg_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 = runtime();
+        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("e3_cg_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("e3_cg_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("e3_cg_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("e3_cg_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("e3_cg_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("e3_cg_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("e3_cg_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: &Problem,
+        p: &mut [f64],
+        tolerance: f64,
+        anchor: Option<usize>,
+        setup_ns: u64,
+    ) -> PoissonSolution {
+        let rt = runtime();
+        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 DevicePcgCache {
+    slot: Option<DeviceCg>,
+}
+
+/// The device counterpart of `solve_pcg_cached`.
+pub fn solve_pcg_device_cached(
+    problem: &Problem,
+    p: &mut [f64],
+    params: &MultigridParameters,
+    tolerance: f64,
+    anchor: Option<usize>,
+    cache: &mut DevicePcgCache,
+) -> 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(DeviceCg::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)
+}
@@ -6,6 +6,8 @@

 #[cfg(feature = "cuda")]
 pub mod device;
+#[cfg(feature = "cuda")]
+pub mod device_cg;
 mod export;
 mod hierarchy;
 mod pcg;