rtx-cfd 3D Stage 1 item 2: mg3_vcycle.cu (seven-point red-black half-sweep, residual, coarsest; top/bot neighbour arrays carry the periodic wrap), LevelExport3/export_hierarchy3/vcycle_f32_reference3, DeviceVcycle3 (the K=1 sequence); gate 2 HELD: device = host f32 V-cycle to 4e-7 relative on 96×40×{1,8} (periodic and walls) and 378×62×62; 4.88 ms per V-cycle at 1.45 M cells incl. transfers

Co-Authored-By: Claude Fable 5.1 <[email protected]>

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

@@ -0,0 +1,116 @@
+/**
+ * 3D Stage 1 (omni-cortex `docs/three_d_stage1_campaign.md`, gate 2): the
+ * V-cycle's maps for a MASKED, VARIABLE-COEFFICIENT seven-point operator.
+ * Per-cell arrays are [n] (one march); the index lists (cells, colours,
+ * children) drive the maps; `top`/`bot` give the neighbour above/below per
+ * cell (UINT_MAX = none; a zero coefficient is never read), so a periodic
+ * z is data. Restriction, prolongation and zero are the 2D kernels (they
+ * never touch the stencil).
+ */
+#define NONE 0xFFFFFFFFu
+
+__device__ __forceinline__ float nb_sum3(
+    int g, int nx, const float* ae, const float* aw, const float* an, const float* as_,
+    const float* at, const float* ab, const unsigned int* top, const unsigned int* bot,
+    const float* x)
+{
+    float s = 0.0f;
+    float e = ae[g]; if (e != 0.0f) s += e * x[g + 1];
+    float w = aw[g]; if (w != 0.0f) s += w * x[g - 1];
+    float nn = an[g]; if (nn != 0.0f) s += nn * x[g + nx];
+    float ss = as_[g]; if (ss != 0.0f) s += ss * x[g - nx];
+    float t = at[g]; if (t != 0.0f) s += t * x[top[g]];
+    float b = ab[g]; if (b != 0.0f) s += b * x[bot[g]];
+    return s;
+}
+
+extern "C" __global__ void mg3_rb_half(
+    int n_col, const unsigned int* __restrict__ col,
+    const float* __restrict__ ae, const float* __restrict__ aw,
+    const float* __restrict__ an, const float* __restrict__ as_,
+    const float* __restrict__ at, const float* __restrict__ ab,
+    const unsigned int* __restrict__ top, const unsigned int* __restrict__ bot,
+    const float* __restrict__ ap, const float* __restrict__ b,
+    float* __restrict__ x, int nx)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_col) return;
+    int g = col[t];
+    float s = nb_sum3(g, nx, ae, aw, an, as_, at, ab, top, bot, x);
+    x[g] = (b[g] + s) / ap[g];
+}
+
+extern "C" __global__ void mg3_residual(
+    int n_cells, const unsigned int* __restrict__ cells,
+    const float* __restrict__ ae, const float* __restrict__ aw,
+    const float* __restrict__ an, const float* __restrict__ as_,
+    const float* __restrict__ at, const float* __restrict__ ab,
+    const unsigned int* __restrict__ top, const unsigned int* __restrict__ bot,
+    const float* __restrict__ ap, const float* __restrict__ b,
+    const float* __restrict__ x, float* __restrict__ r, int nx)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int g = cells[t];
+    float s = nb_sum3(g, nx, ae, aw, an, as_, at, ab, top, bot, x);
+    r[g] = b[g] - (ap[g] * x[g] - s);
+}
+
+/* b_c[c] = sum of r_f over the children of coarse cell c (fixed order). */
+extern "C" __global__ void mg3_restrict(
+    int n_coarse, const unsigned int* __restrict__ coarse_cells,
+    const unsigned int* __restrict__ children_ptr, const unsigned int* __restrict__ children_idx,
+    const float* __restrict__ r_f, float* __restrict__ b_c)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_coarse) return;
+    float s = 0.0f;
+    for (unsigned int p = children_ptr[t]; p < children_ptr[t + 1]; ++p) s += r_f[children_idx[p]];
+    b_c[coarse_cells[t]] = s;
+}
+
+/* x_f += 2 x_c[coarse_of[idx]] */
+extern "C" __global__ void mg3_prolong(
+    int n_cells, const unsigned int* __restrict__ cells, const unsigned int* __restrict__ coarse_of,
+    float* __restrict__ x_f, const float* __restrict__ x_c)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int idx = cells[t];
+    x_f[idx] += 2.0f * x_c[coarse_of[idx]];
+}
+
+extern "C" __global__ void mg3_zero(int n_cells, const unsigned int* __restrict__ cells, float* __restrict__ x)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    x[cells[t]] = 0.0f;
+}
+
+/* The coarsest level: one thread, `sweeps` symmetric RED-BLACK sweeps
+ * (red, black, black, red) from zero — the host's ordering. */
+extern "C" __global__ void mg3_coarsest(
+    int n_cells, const unsigned int* __restrict__ cells,
+    int n_red, const unsigned int* __restrict__ red,
+    int n_black, const unsigned int* __restrict__ black,
+    const float* __restrict__ ae, const float* __restrict__ aw,
+    const float* __restrict__ an, const float* __restrict__ as_,
+    const float* __restrict__ at, const float* __restrict__ ab,
+    const unsigned int* __restrict__ top, const unsigned int* __restrict__ bot,
+    const float* __restrict__ ap, const float* __restrict__ b,
+    float* __restrict__ x, int nx, int sweeps)
+{
+    if (blockIdx.x * blockDim.x + threadIdx.x != 0) return;
+    for (int t = 0; t < n_cells; ++t) x[cells[t]] = 0.0f;
+    for (int sw = 0; sw < sweeps; ++sw) {
+        for (int half = 0; half < 4; ++half) {
+            const unsigned int* list = (half == 0 || half == 3) ? red : black;
+            int n_list = (half == 0 || half == 3) ? n_red : n_black;
+            for (int t = 0; t < n_list; ++t) {
+                int g = list[t];
+                float s = nb_sum3(g, nx, ae, aw, an, as_, at, ab, top, bot, x);
+                x[g] = (b[g] + s) / ap[g];
+            }
+        }
+    }
+}

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

@@ -0,0 +1,328 @@
+//! Gate 2: the seven-point V-cycle on the CUDA device (`mg3_vcycle.cu`),
+//! one march, persistent buffers per operator; the 2D `poisson/device.rs`
+//! K = 1 sequence with the z terms. One runtime per process.
+
+use super::LevelExport3;
+use cudarc::driver::{
+    CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg,
+};
+use cudarc::nvrtc::{CompileOptions, compile_ptx_with_opts};
+use std::sync::{Arc, OnceLock};
+
+const KERNELS: &str = include_str!("../../../../kernels/cuda/mg3_vcycle.cu");
+
+pub(crate) struct Runtime3 {
+    pub(crate) ctx: Arc<CudaContext>,
+    pub(crate) stream: Arc<CudaStream>,
+    _module: Arc<CudaModule>,
+    f_half: CudaFunction,
+    f_res: CudaFunction,
+    f_restrict: CudaFunction,
+    f_prolong: CudaFunction,
+    f_zero: CudaFunction,
+    f_coarsest: CudaFunction,
+}
+
+static RUNTIME: OnceLock<Runtime3> = OnceLock::new();
+
+pub(crate) fn runtime3() -> &'static Runtime3 {
+    RUNTIME.get_or_init(|| {
+        let ctx = CudaContext::new(0).expect("CUDA context (device 0)");
+        let stream = ctx.default_stream();
+        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: mg3_vcycle.cu");
+        let module = ctx.load_module(ptx).expect("mg3_vcycle module");
+        let f = |name: &str| module.load_function(name).expect(name);
+        Runtime3 {
+            f_half: f("mg3_rb_half"),
+            f_res: f("mg3_residual"),
+            f_restrict: f("mg3_restrict"),
+            f_prolong: f("mg3_prolong"),
+            f_zero: f("mg3_zero"),
+            f_coarsest: f("mg3_coarsest"),
+            ctx,
+            stream,
+            _module: module,
+        }
+    })
+}
+
+pub(crate) struct DevLevel3 {
+    pub(crate) n: usize,
+    pub(crate) nx: i32,
+    pub(crate) n_cells: usize,
+    pub(crate) n_red: usize,
+    pub(crate) n_black: usize,
+    pub(crate) cells: CudaSlice<u32>,
+    pub(crate) red: CudaSlice<u32>,
+    pub(crate) black: CudaSlice<u32>,
+    pub(crate) top: CudaSlice<u32>,
+    pub(crate) bot: CudaSlice<u32>,
+    pub(crate) coarse_of: CudaSlice<u32>,
+    pub(crate) children_ptr: CudaSlice<u32>,
+    pub(crate) children_idx: CudaSlice<u32>,
+    pub(crate) ae: CudaSlice<f32>,
+    pub(crate) aw: CudaSlice<f32>,
+    pub(crate) an: CudaSlice<f32>,
+    pub(crate) as_: CudaSlice<f32>,
+    pub(crate) at: CudaSlice<f32>,
+    pub(crate) ab: CudaSlice<f32>,
+    pub(crate) ap: CudaSlice<f32>,
+    pub(crate) b: CudaSlice<f32>,
+    pub(crate) x: CudaSlice<f32>,
+    pub(crate) r: CudaSlice<f32>,
+}
+
+/// One operator's hierarchy on the device.
+pub struct DeviceVcycle3 {
+    pub(crate) levels: Vec<DevLevel3>,
+    sweeps: usize,
+    fine_cells: Vec<u32>,
+    r_f32: Vec<f32>,
+    z_f32: Vec<f32>,
+}
+
+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,
+    }
+}
+
+impl DeviceVcycle3 {
+    pub fn new(levels: &[LevelExport3], sweeps: usize) -> Self {
+        let rt = runtime3();
+        let up_u = |v: &[u32]| -> CudaSlice<u32> {
+            rt.stream
+                .memcpy_stod(if v.is_empty() { &[0u32][..] } else { v })
+                .expect("upload")
+        };
+        let up_f = |v: &[f32]| -> CudaSlice<f32> { rt.stream.memcpy_stod(v).expect("upload") };
+        let dev: Vec<DevLevel3> = levels
+            .iter()
+            .map(|l| {
+                let n = l.nx * l.ny * l.nz;
+                DevLevel3 {
+                    n,
+                    nx: l.nx as i32,
+                    n_cells: l.cells.len(),
+                    n_red: l.red.len(),
+                    n_black: l.black.len(),
+                    cells: up_u(&l.cells),
+                    red: up_u(&l.red),
+                    black: up_u(&l.black),
+                    top: up_u(&l.top),
+                    bot: up_u(&l.bot),
+                    coarse_of: up_u(&l.coarse_of),
+                    children_ptr: up_u(&l.children_ptr),
+                    children_idx: up_u(&l.children_idx),
+                    ae: up_f(&l.ae),
+                    aw: up_f(&l.aw),
+                    an: up_f(&l.an),
+                    as_: up_f(&l.as_),
+                    at: up_f(&l.at),
+                    ab: up_f(&l.ab),
+                    ap: up_f(&l.ap),
+                    b: rt.stream.alloc_zeros::<f32>(n).expect("alloc"),
+                    x: rt.stream.alloc_zeros::<f32>(n).expect("alloc"),
+                    r: rt.stream.alloc_zeros::<f32>(n).expect("alloc"),
+                }
+            })
+            .collect();
+        let n0 = dev[0].n;
+        Self {
+            levels: dev,
+            sweeps,
+            fine_cells: levels[0].cells.clone(),
+            r_f32: vec![0.0; n0],
+            z_f32: vec![0.0; n0],
+        }
+    }
+
+    pub fn depth(&self) -> usize {
+        self.levels.len()
+    }
+
+    fn half(&mut self, l: usize, colour: u8) {
+        let rt = runtime3();
+        let lv = &mut self.levels[l];
+        let (list, n_list) = if colour == 0 {
+            (&lv.red, lv.n_red)
+        } else {
+            (&lv.black, lv.n_black)
+        };
+        let n_list_i = n_list as i32;
+        unsafe {
+            rt.stream
+                .launch_builder(&rt.f_half)
+                .arg(&n_list_i)
+                .arg(list)
+                .arg(&lv.ae)
+                .arg(&lv.aw)
+                .arg(&lv.an)
+                .arg(&lv.as_)
+                .arg(&lv.at)
+                .arg(&lv.ab)
+                .arg(&lv.top)
+                .arg(&lv.bot)
+                .arg(&lv.ap)
+                .arg(&lv.b)
+                .arg(&mut lv.x)
+                .arg(&lv.nx)
+                .launch(cfg(n_list))
+                .expect("mg3_rb_half");
+        }
+    }
+
+    fn smooth(&mut self, l: usize) {
+        for _ in 0..self.sweeps {
+            self.half(l, 0);
+            self.half(l, 1);
+            self.half(l, 1);
+            self.half(l, 0);
+        }
+    }
+
+    /// The V-cycle on the device with `r` already in `levels[0].b`; the
+    /// correction is left in `levels[0].x`. No transfers.
+    pub(crate) fn vcycle_on_device(&mut self) {
+        let rt = runtime3();
+        let depth = self.levels.len();
+        for l in 0..depth - 1 {
+            {
+                let lv = &mut self.levels[l];
+                let n_cells_i = lv.n_cells as i32;
+                unsafe {
+                    rt.stream
+                        .launch_builder(&rt.f_zero)
+                        .arg(&n_cells_i)
+                        .arg(&lv.cells)
+                        .arg(&mut lv.x)
+                        .launch(cfg(lv.n_cells))
+                        .expect("mg3_zero");
+                }
+            }
+            self.smooth(l);
+            {
+                let lv = &mut self.levels[l];
+                let n_cells_i = lv.n_cells as i32;
+                unsafe {
+                    rt.stream
+                        .launch_builder(&rt.f_res)
+                        .arg(&n_cells_i)
+                        .arg(&lv.cells)
+                        .arg(&lv.ae)
+                        .arg(&lv.aw)
+                        .arg(&lv.an)
+                        .arg(&lv.as_)
+                        .arg(&lv.at)
+                        .arg(&lv.ab)
+                        .arg(&lv.top)
+                        .arg(&lv.bot)
+                        .arg(&lv.ap)
+                        .arg(&lv.b)
+                        .arg(&lv.x)
+                        .arg(&mut lv.r)
+                        .arg(&lv.nx)
+                        .launch(cfg(lv.n_cells))
+                        .expect("mg3_residual");
+                }
+            }
+            let (fine, coarse) = self.levels.split_at_mut(l + 1);
+            let (lf, lc) = (&fine[l], &mut coarse[0]);
+            let n_c_cells_i = lc.n_cells as i32;
+            unsafe {
+                rt.stream
+                    .launch_builder(&rt.f_restrict)
+                    .arg(&n_c_cells_i)
+                    .arg(&lc.cells)
+                    .arg(&lf.children_ptr)
+                    .arg(&lf.children_idx)
+                    .arg(&lf.r)
+                    .arg(&mut lc.b)
+                    .launch(cfg(lc.n_cells))
+                    .expect("mg3_restrict");
+            }
+        }
+        {
+            let lv = &mut self.levels[depth - 1];
+            let (n_cells_i, sw_i) = (lv.n_cells as i32, 50i32);
+            let (n_red_i, n_black_i) = (lv.n_red as i32, lv.n_black as i32);
+            unsafe {
+                rt.stream
+                    .launch_builder(&rt.f_coarsest)
+                    .arg(&n_cells_i)
+                    .arg(&lv.cells)
+                    .arg(&n_red_i)
+                    .arg(&lv.red)
+                    .arg(&n_black_i)
+                    .arg(&lv.black)
+                    .arg(&lv.ae)
+                    .arg(&lv.aw)
+                    .arg(&lv.an)
+                    .arg(&lv.as_)
+                    .arg(&lv.at)
+                    .arg(&lv.ab)
+                    .arg(&lv.top)
+                    .arg(&lv.bot)
+                    .arg(&lv.ap)
+                    .arg(&lv.b)
+                    .arg(&mut lv.x)
+                    .arg(&lv.nx)
+                    .arg(&sw_i)
+                    .launch(LaunchConfig {
+                        grid_dim: (1, 1, 1),
+                        block_dim: (32, 1, 1),
+                        shared_mem_bytes: 0,
+                    })
+                    .expect("mg3_coarsest");
+            }
+        }
+        for l in (0..depth - 1).rev() {
+            {
+                let (fine, coarse) = self.levels.split_at_mut(l + 1);
+                let (lf, lc) = (&mut fine[l], &coarse[0]);
+                let n_cells_i = lf.n_cells as i32;
+                unsafe {
+                    rt.stream
+                        .launch_builder(&rt.f_prolong)
+                        .arg(&n_cells_i)
+                        .arg(&lf.cells)
+                        .arg(&lf.coarse_of)
+                        .arg(&mut lf.x)
+                        .arg(&lc.x)
+                        .launch(cfg(lf.n_cells))
+                        .expect("mg3_prolong");
+                }
+            }
+            self.smooth(l);
+        }
+    }
+
+    /// `z = M⁻¹ r` on the active cells (upload, V-cycle, download).
+    pub fn apply(&mut self, r: &[f64], z: &mut [f64]) {
+        let rt = runtime3();
+        for (dst, &src) in self.r_f32.iter_mut().zip(r) {
+            *dst = src as f32;
+        }
+        rt.stream
+            .memcpy_htod(&self.r_f32, &mut self.levels[0].b)
+            .expect("upload r");
+        self.vcycle_on_device();
+        rt.stream
+            .memcpy_dtoh(&self.levels[0].x, &mut self.z_f32)
+            .expect("download z");
+        rt.stream.synchronize().expect("sync");
+        for &idx in &self.fine_cells {
+            z[idx as usize] = self.z_f32[idx as usize] as f64;
+        }
+    }
+}

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

@@ -0,0 +1,113 @@
+//! The f32 hierarchy exported for a device V-cycle (gate 2): per level the
+//! index lists (cells, colours, parent map, CSR children) and the seven
+//! coefficient arrays, plus explicit top/bottom neighbour indices so the
+//! periodic wrap is data, not arithmetic, on the device.
+
+use super::{Hierarchy3, PoissonProblem3D};
+use crate::solvers::incompressible::poisson::MultigridParameters;
+
+/// One exported level. `top`/`bot` hold the neighbour index above/below
+/// each cell (`u32::MAX` = none); a zero coefficient is never read.
+pub struct LevelExport3 {
+    pub nx: usize,
+    pub ny: usize,
+    pub nz: usize,
+    pub cells: Vec<u32>,
+    pub red: Vec<u32>,
+    pub black: Vec<u32>,
+    pub top: Vec<u32>,
+    pub bot: Vec<u32>,
+    /// Fine cell → coarse cell (`u32::MAX` without an equation; empty on
+    /// the coarsest level).
+    pub coarse_of: Vec<u32>,
+    /// For the NEXT level's cells in its `cells` order: the fine cells
+    /// restricting into each (CSR).
+    pub children_ptr: Vec<u32>,
+    pub children_idx: Vec<u32>,
+    pub ae: Vec<f32>,
+    pub aw: Vec<f32>,
+    pub an: Vec<f32>,
+    pub as_: Vec<f32>,
+    pub at: Vec<f32>,
+    pub ab: Vec<f32>,
+    pub ap: Vec<f32>,
+}
+
+/// The f32 hierarchy of `problem`, level 0 fine.
+pub fn export_hierarchy3(
+    problem: &PoissonProblem3D,
+    params: &MultigridParameters,
+) -> Vec<LevelExport3> {
+    export_levels3(&Hierarchy3::<f32>::build(problem, params))
+}
+
+pub(super) fn export_levels3(hier: &Hierarchy3<f32>) -> Vec<LevelExport3> {
+    let depth = hier.levels.len();
+    let to_u32 = |v: &[usize]| {
+        v.iter()
+            .map(|&i| if i == usize::MAX { u32::MAX } else { i as u32 })
+            .collect::<Vec<u32>>()
+    };
+    (0..depth)
+        .map(|l| {
+            let lv = &hier.levels[l];
+            let (children_ptr, children_idx) = if l + 1 < depth {
+                let coarse = &hier.levels[l + 1];
+                let nc = coarse.problem.nx * coarse.problem.ny * coarse.problem.nz;
+                let mut pos = vec![usize::MAX; nc];
+                for (k, &c) in coarse.cells.iter().enumerate() {
+                    pos[c] = k;
+                }
+                let mut lists: Vec<Vec<u32>> = vec![Vec::new(); coarse.cells.len()];
+                for &idx in &lv.cells {
+                    let c = lv.coarse_of[idx];
+                    if coarse.active[c] {
+                        lists[pos[c]].push(idx as u32);
+                    }
+                }
+                let mut ptr = Vec::with_capacity(lists.len() + 1);
+                let mut flat = Vec::new();
+                ptr.push(0u32);
+                for list in &lists {
+                    flat.extend_from_slice(list);
+                    ptr.push(flat.len() as u32);
+                }
+                (ptr, flat)
+            } else {
+                (Vec::new(), Vec::new())
+            };
+            LevelExport3 {
+                nx: lv.problem.nx,
+                ny: lv.problem.ny,
+                nz: lv.problem.nz,
+                cells: to_u32(&lv.cells),
+                red: to_u32(&lv.red),
+                black: to_u32(&lv.black),
+                top: to_u32(&lv.top),
+                bot: to_u32(&lv.bot),
+                coarse_of: to_u32(&lv.coarse_of),
+                children_ptr,
+                children_idx,
+                ae: lv.ae.clone(),
+                aw: lv.aw.clone(),
+                an: lv.an.clone(),
+                as_: lv.as_.clone(),
+                at: lv.at.clone(),
+                ab: lv.ab.clone(),
+                ap: lv.ap.clone(),
+            }
+        })
+        .collect()
+}
+
+/// `z = M⁻¹ r` by the host f32 V-cycle: the reference a device V-cycle is
+/// measured against.
+pub fn vcycle_f32_reference3(
+    problem: &PoissonProblem3D,
+    params: &MultigridParameters,
+    r: &[f64],
+    z: &mut [f64],
+) {
+    let mut hier = Hierarchy3::<f32>::build(problem, params);
+    hier.apply_preconditioner(r, z);
+}

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

@@ -11,6 +11,11 @@ use crate::solvers::incompressible::poisson::{
     MgPrecision, MgScalar, MgSmoother, MultigridParameters, PoissonSolution,
 };
 
+#[cfg(feature = "cuda")]
+pub mod device;
+pub mod export;
+pub use export::{LevelExport3, export_hierarchy3, vcycle_f32_reference3};
+
 /// Symmetric GS sweeps on the coarsest level (the 2D value).
 const COARSEST_SWEEPS: usize = 50;
 /// The 2D `COARSE_CORRECTION`, proved dimension-independent there.