PERF-2 P3-ii: the device V-cycle as the CG's preconditioner — poisson/device.rs (one CUDA runtime per process, persistent per-operator buffers, the mg_vcycle.cu kernels at K = 1; upload r, run the V-cycle, download z; the f64 CG unchanged), MultigridParameters::device, Prepared holds the device hierarchy, the CG driver destructures the prepared operator instead of cloning it; EmbeddedPisoSolver::set_poisson_device, overset pass-through, harness knob RTX_FSI2O_MG_DEVICE=1; export_levels factored out; the quarantine's dangling cfg attribute fixed

Co-Authored-By: Claude Fable 5.1 <[email protected]>
Claude-Session: https://claude.ai/code/session_01YJPeT6WA2e7YvAnS875AHL

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

@@ -0,0 +1,312 @@
+//! PERF-2 P3-ii (`docs/perf2_campaign.md`): the multigrid V-cycle on the
+//! CUDA device as the CG's preconditioner — the `mg_vcycle.cu` kernels of
+//! the go/no-go benchmark, K = 1, with persistent device buffers per
+//! operator. One CUDA context, stream and module per process (built on
+//! first use); the CG, its residual and its stop stay f64 on the CPU.
+
+use super::LevelExport;
+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/mg_vcycle.cu");
+
+struct Runtime {
+    _ctx: Arc<CudaContext>,
+    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<Runtime> = OnceLock::new();
+
+fn runtime() -> &'static Runtime {
+    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: mg_vcycle.cu");
+        let module = ctx.load_module(ptx).expect("mg_vcycle module");
+        let f = |name: &str| module.load_function(name).expect(name);
+        Runtime {
+            f_half: f("mg_rb_half"),
+            f_res: f("mg_residual"),
+            f_restrict: f("mg_restrict"),
+            f_prolong: f("mg_prolong"),
+            f_zero: f("mg_zero"),
+            f_coarsest: f("mg_coarsest"),
+            _ctx: ctx,
+            stream,
+            _module: module,
+        }
+    })
+}
+
+struct DevLevel {
+    n: usize,
+    nx: i32,
+    n_cells: usize,
+    n_red: usize,
+    n_black: usize,
+    cells: CudaSlice<u32>,
+    red: CudaSlice<u32>,
+    black: CudaSlice<u32>,
+    coarse_of: CudaSlice<u32>,
+    children_ptr: CudaSlice<u32>,
+    children_idx: CudaSlice<u32>,
+    ae: CudaSlice<f32>,
+    aw: CudaSlice<f32>,
+    an: CudaSlice<f32>,
+    as_: CudaSlice<f32>,
+    ap: CudaSlice<f32>,
+    b: CudaSlice<f32>,
+    x: CudaSlice<f32>,
+    r: CudaSlice<f32>,
+}
+
+/// One operator's hierarchy on the device (K = 1).
+pub(super) struct DeviceVcycle {
+    levels: Vec<DevLevel>,
+    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 DeviceVcycle {
+    pub(super) fn new(levels: &[LevelExport], sweeps: usize) -> Self {
+        let rt = runtime();
+        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<DevLevel> = levels
+            .iter()
+            .map(|l| {
+                let n = l.nx * l.ny;
+                DevLevel {
+                    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),
+                    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_),
+                    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],
+        }
+    }
+
+    fn half(&mut self, l: usize, colour: u8) {
+        let rt = runtime();
+        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_i, n_list_i) = (lv.n as i32, n_list as i32);
+        unsafe {
+            rt.stream
+                .launch_builder(&rt.f_half)
+                .arg(&n_list_i)
+                .arg(list)
+                .arg(&n_i)
+                .arg(&lv.ae)
+                .arg(&lv.aw)
+                .arg(&lv.an)
+                .arg(&lv.as_)
+                .arg(&lv.ap)
+                .arg(&lv.b)
+                .arg(&mut lv.x)
+                .arg(&lv.nx)
+                .launch(cfg(n_list))
+                .expect("mg_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);
+        }
+    }
+
+    /// `z = M⁻¹ r` on the active cells (the same V-cycle as
+    /// `Hierarchy::apply_preconditioner`, in f32 on the device).
+    pub(super) fn apply(&mut self, r: &[f64], z: &mut [f64]) {
+        let rt = runtime();
+        let depth = self.levels.len();
+        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");
+        // Down.
+        for l in 0..depth - 1 {
+            {
+                let lv = &mut self.levels[l];
+                let (n_i, n_cells_i) = (lv.n as i32, lv.n_cells as i32);
+                unsafe {
+                    rt.stream
+                        .launch_builder(&rt.f_zero)
+                        .arg(&n_cells_i)
+                        .arg(&lv.cells)
+                        .arg(&n_i)
+                        .arg(&mut lv.x)
+                        .launch(cfg(lv.n_cells))
+                        .expect("mg_zero");
+                }
+            }
+            self.smooth(l);
+            {
+                let lv = &mut self.levels[l];
+                let (n_i, n_cells_i) = (lv.n as i32, lv.n_cells as i32);
+                unsafe {
+                    rt.stream
+                        .launch_builder(&rt.f_res)
+                        .arg(&n_cells_i)
+                        .arg(&lv.cells)
+                        .arg(&n_i)
+                        .arg(&lv.ae)
+                        .arg(&lv.aw)
+                        .arg(&lv.an)
+                        .arg(&lv.as_)
+                        .arg(&lv.ap)
+                        .arg(&lv.b)
+                        .arg(&lv.x)
+                        .arg(&mut lv.r)
+                        .arg(&lv.nx)
+                        .launch(cfg(lv.n_cells))
+                        .expect("mg_residual");
+                }
+            }
+            let (fine, coarse) = self.levels.split_at_mut(l + 1);
+            let (lf, lc) = (&fine[l], &mut coarse[0]);
+            let (n_c_cells_i, n_f_i, n_c_i) = (lc.n_cells as i32, lf.n as i32, lc.n 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(&n_f_i)
+                    .arg(&n_c_i)
+                    .arg(&lf.r)
+                    .arg(&mut lc.b)
+                    .launch(cfg(lc.n_cells))
+                    .expect("mg_restrict");
+            }
+        }
+        // Coarsest.
+        {
+            let lv = &mut self.levels[depth - 1];
+            let (k_i, n_cells_i, n_i, sw_i) = (1i32, lv.n_cells as i32, lv.n 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(&k_i)
+                    .arg(&n_cells_i)
+                    .arg(&lv.cells)
+                    .arg(&n_red_i)
+                    .arg(&lv.red)
+                    .arg(&n_black_i)
+                    .arg(&lv.black)
+                    .arg(&n_i)
+                    .arg(&lv.ae)
+                    .arg(&lv.aw)
+                    .arg(&lv.an)
+                    .arg(&lv.as_)
+                    .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("mg_coarsest");
+            }
+        }
+        // Up.
+        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, n_f_i, n_c_i) = (lf.n_cells as i32, lf.n as i32, lc.n as i32);
+                unsafe {
+                    rt.stream
+                        .launch_builder(&rt.f_prolong)
+                        .arg(&n_cells_i)
+                        .arg(&lf.cells)
+                        .arg(&lf.coarse_of)
+                        .arg(&n_f_i)
+                        .arg(&n_c_i)
+                        .arg(&mut lf.x)
+                        .arg(&lc.x)
+                        .launch(cfg(lf.n_cells))
+                        .expect("mg_prolong");
+                }
+            }
+            self.smooth(l);
+        }
+        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;
+        }
+    }
+}