From 0dc95a8de5eb08a44407d7f4de07979cd11d9b18 Mon Sep 17 00:00:00 2001 From: Omar Sobh Date: Wed, 16 Sep 2026 01:03:02 -0500 Subject: [PATCH] =?UTF-8?q?PERF-2=20P3:=20the=20batched=20device=20V-cycle?= =?UTF-8?q?=20microbenchmark=20=E2=80=94=20mg=5Fvcycle.cu=20(masked=20vari?= =?UTF-8?q?able-coefficient=20five-point=20red-black=20half-sweeps,=20resi?= =?UTF-8?q?dual,=20CSR=20restriction=20in=20a=20fixed=20order,=20prolongat?= =?UTF-8?q?ion,=20coarsest=20sweeps;=20[K][n]=20layout,=20blockIdx.y=20=3D?= =?UTF-8?q?=20march),=20LevelExport/export=5Fhierarchy=20and=20vcycle=5Ff3?= =?UTF-8?q?2=5Freference=20on=20the=20CPU=20side,=20and=20the=20ignored=20?= =?UTF-8?q?cuda-feature=20test=20that=20checks=20the=20K=20=3D=201=20devic?= =?UTF-8?q?e=20V-cycle=20against=20the=20CPU=20f32=20reference=20and=20tim?= =?UTF-8?q?es=20K=20=3D=201=20/=208=20/=2016=20per=20march?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01YJPeT6WA2e7YvAnS875AHL --- .../rtx-cfd/src/kernels/cuda/mg_vcycle.cu | 107 +++++ .../rtx-cfd/src/solvers/incompressible/mod.rs | 3 +- .../src/solvers/incompressible/poisson.rs | 93 +++++ .../rtx-cfd/tests/gpu_vcycle_bench.rs | 369 ++++++++++++++++++ 4 files changed, 571 insertions(+), 1 deletion(-) create mode 100644 crates/specialized/rtx-cfd/src/kernels/cuda/mg_vcycle.cu create mode 100644 crates/specialized/rtx-cfd/tests/gpu_vcycle_bench.rs diff --git a/crates/specialized/rtx-cfd/src/kernels/cuda/mg_vcycle.cu b/crates/specialized/rtx-cfd/src/kernels/cuda/mg_vcycle.cu new file mode 100644 index 0000000..601c872 --- /dev/null +++ b/crates/specialized/rtx-cfd/src/kernels/cuda/mg_vcycle.cu @@ -0,0 +1,107 @@ +/** + * PERF-2 P3 (`docs/perf2_campaign.md`): the multigrid V-cycle's maps for a + * MASKED, VARIABLE-COEFFICIENT five-point operator, batched over K marches. + * Layout: every per-cell array is [K][n] (march-major, n = nx*ny of the + * level); the index lists (cells, colours, children) are shared across the + * batch in this benchmark. blockIdx.y = the march. + */ +extern "C" __global__ void mg_rb_half( + int n_col, const unsigned int* __restrict__ col, int n, + const float* __restrict__ ae, const float* __restrict__ aw, + const float* __restrict__ an, const float* __restrict__ as_, + 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 base = blockIdx.y * n; + int idx = col[t]; + int g = base + idx; + 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]; + x[g] = (b[g] + s) / ap[g]; +} + +extern "C" __global__ void mg_residual( + int n_cells, const unsigned int* __restrict__ cells, int n, + const float* __restrict__ ae, const float* __restrict__ aw, + const float* __restrict__ an, const float* __restrict__ as_, + 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 base = blockIdx.y * n; + int g = base + cells[t]; + 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]; + 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 mg_restrict( + int n_coarse, const unsigned int* __restrict__ coarse_cells, + const unsigned int* __restrict__ children_ptr, const unsigned int* __restrict__ children_idx, + int n_f, int n_c, const float* __restrict__ r_f, float* __restrict__ b_c) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= n_coarse) return; + int k = blockIdx.y; + const float* rf = r_f + (size_t)k * n_f; + float s = 0.0f; + for (unsigned int p = children_ptr[t]; p < children_ptr[t + 1]; ++p) s += rf[children_idx[p]]; + b_c[(size_t)k * n_c + coarse_cells[t]] = s; +} + +/* x_f += 2 x_c[coarse_of[idx]] */ +extern "C" __global__ void mg_prolong( + int n_cells, const unsigned int* __restrict__ cells, const unsigned int* __restrict__ coarse_of, + int n_f, int n_c, float* __restrict__ x_f, const float* __restrict__ x_c) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= n_cells) return; + int k = blockIdx.y; + int idx = cells[t]; + x_f[(size_t)k * n_f + idx] += 2.0f * x_c[(size_t)k * n_c + coarse_of[idx]]; +} + +extern "C" __global__ void mg_zero(int n_cells, const unsigned int* __restrict__ cells, int n, float* __restrict__ x) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= n_cells) return; + x[(size_t)blockIdx.y * n + cells[t]] = 0.0f; +} + +/* The coarsest level: one thread per march, `sweeps` symmetric lexicographic sweeps over <= a few dozen cells. */ +extern "C" __global__ void mg_coarsest( + int K, int n_cells, const unsigned int* __restrict__ cells, int n, + const float* __restrict__ ae, const float* __restrict__ aw, + const float* __restrict__ an, const float* __restrict__ as_, + const float* __restrict__ ap, const float* __restrict__ b, + float* __restrict__ x, int nx, int sweeps) +{ + int k = blockIdx.x * blockDim.x + threadIdx.x; + if (k >= K) return; + int base = k * n; + for (int t = 0; t < n_cells; ++t) x[base + cells[t]] = 0.0f; + for (int sw = 0; sw < sweeps; ++sw) { + for (int pass = 0; pass < 2; ++pass) { + for (int q = 0; q < n_cells; ++q) { + int t = pass == 0 ? q : n_cells - 1 - q; + int g = base + cells[t]; + 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]; + x[g] = (b[g] + s) / ap[g]; + } + } + } +} diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs index fba8a8c..0d2a385 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs @@ -65,7 +65,8 @@ pub use piso::{PisoParameters, PisoResult, PisoSolver}; #[cfg(feature = "cuda")] pub use piso_gpu::PisoGpuSolver; pub use poisson::{ - MgPrecision, MgSmoother, MultigridParameters, PcgCache, PoissonProblem, PoissonSolution, + LevelExport, MgPrecision, MgSmoother, MultigridParameters, PcgCache, PoissonProblem, + PoissonSolution, export_hierarchy, vcycle_f32_reference, PoissonSolverKind, configure_threads, solve_multigrid_pcg, solve_multigrid_pcg_cached, }; pub use polygon_sdf::PolygonSdf; diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs index a545eed..bb5a31d 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs @@ -1001,6 +1001,99 @@ pub fn configure_threads(threads: usize) -> usize { n } +/// PERF-2 P3 (`docs/perf2_campaign.md`): one hierarchy level exported for a +/// device V-cycle — the sanitised f32 coefficients, the active-cell and +/// colour index lists, the parent map, and the coarse cells' children in +/// CSR form (so a restriction can sum in a fixed order on the device). +#[derive(Debug, Clone)] +pub struct LevelExport { + pub nx: usize, + pub ny: usize, + pub cells: Vec, + pub red: Vec, + pub black: Vec, + /// Fine cell → coarse cell (`u32::MAX` without an equation; empty on + /// the coarsest level). + pub coarse_of: Vec, + /// For the NEXT level's cells, in its `cells` order: the fine cells + /// restricting into each (CSR: `children_ptr[c]..children_ptr[c + 1]`). + pub children_ptr: Vec, + pub children_idx: Vec, + pub ae: Vec, + pub aw: Vec, + pub an: Vec, + pub as_: Vec, + pub ap: Vec, +} + +/// The f32 hierarchy of `problem` (red-black colour lists included), level +/// 0 fine, for a device implementation of [`Hierarchy::apply_preconditioner`]. +pub fn export_hierarchy(problem: &PoissonProblem, params: &MultigridParameters) -> Vec { + let hier = Hierarchy::::build(problem, params); + let depth = hier.levels.len(); + (0..depth) + .map(|l| { + let lv = &hier.levels[l]; + let to_u32 = |v: &[usize]| v.iter().map(|&i| i as u32).collect::>(); + let (children_ptr, children_idx) = if l + 1 < depth { + let coarse = &hier.levels[l + 1]; + let mut pos = vec![usize::MAX; coarse.problem.nx * coarse.problem.ny]; + for (k, &c) in coarse.cells.iter().enumerate() { + pos[c] = k; + } + let mut lists: Vec> = 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()) + }; + LevelExport { + nx: lv.problem.nx, + ny: lv.problem.ny, + cells: to_u32(&lv.cells), + red: to_u32(&lv.red), + black: to_u32(&lv.black), + coarse_of: lv + .coarse_of + .iter() + .map(|&c| if c == usize::MAX { u32::MAX } else { c as u32 }) + .collect(), + children_ptr, + children_idx, + ae: lv.ae.clone(), + aw: lv.aw.clone(), + an: lv.an.clone(), + as_: lv.as_.clone(), + ap: lv.ap.clone(), + } + }) + .collect() +} + +/// One f32 V-cycle of `problem`'s hierarchy on `r` (the CPU reference for a +/// device V-cycle): `z = M⁻¹ r` exactly as the preconditioner computes it. +pub fn vcycle_f32_reference( + problem: &PoissonProblem, + params: &MultigridParameters, + r: &[f64], + z: &mut [f64], +) { + let mut hier = Hierarchy::::build(problem, params); + hier.apply_preconditioner(r, z); +} + /// A reusable prepared solver per V-cycle precision (PERF-2 P1.1): the /// operator's hierarchy is rebuilt only when the operator changes. #[derive(Default)] diff --git a/crates/specialized/rtx-cfd/tests/gpu_vcycle_bench.rs b/crates/specialized/rtx-cfd/tests/gpu_vcycle_bench.rs new file mode 100644 index 0000000..819c962 --- /dev/null +++ b/crates/specialized/rtx-cfd/tests/gpu_vcycle_bench.rs @@ -0,0 +1,369 @@ +//! PERF-2 P3 go/no-go (`docs/perf2_campaign.md`): the batched f32 red-black +//! V-cycle on the device against the CPU one. Runs only with +//! `--features cuda` on a CUDA host, and only when asked (`--ignored`): +//! +//! `RTX_CUDA_ARCH=sm_120 cargo test --release -p rtx-cfd --features cuda --test gpu_vcycle_bench -- --ignored --nocapture` +//! +//! Checks first that the K = 1 device V-cycle reproduces the CPU f32 +//! red-black V-cycle (same algorithm; FMA contraction and summation order +//! differ, so to 1e-4 of the correction's scale), then times one V-cycle +//! per march at K = 1, 8, 16 (including the residual upload and the +//! correction download) against the CPU's serial f64 red-black V-cycle. +#![cfg(feature = "cuda")] + +use cudarc::driver::{CudaContext, CudaSlice, LaunchConfig, PushKernelArg}; +use cudarc::nvrtc::{CompileOptions, compile_ptx_with_opts}; +use rtx_cfd::solvers::incompressible::{ + LevelExport, MgSmoother, MultigridParameters, PoissonProblem, export_hierarchy, + vcycle_f32_reference, +}; +use std::sync::Arc; +use std::time::Instant; + +const KERNELS: &str = include_str!("../src/kernels/cuda/mg_vcycle.cu"); + +/// The anchor-sized masked operator (ny 62 over the 2.5 × 0.41 channel with +/// a cylinder-sized hole), the same construction as the Poisson pins. +fn problem(nx: usize, ny: usize, seed: u64) -> PoissonProblem { + let mut p = PoissonProblem::new(nx, ny); + let (dx, dy, dt) = (2.5 / nx as f64, 0.41 / ny as f64, 3.24e-4); + let (ae, an) = (dt * dy / dx, dt * dx / dy); + let hole = |i: usize, j: usize| { + let (x, y) = ((i as f64 + 0.5) * dx, (j as f64 + 0.5) * dy); + (x - 0.2).powi(2) + (y - 0.2).powi(2) < 0.08 * 0.08 + }; + for j in 0..ny { + for i in 0..nx { + let idx = j * nx + i; + if hole(i, j) { + p.active[idx] = false; + continue; + } + if i + 1 < nx && !hole(i + 1, j) { + p.ae[idx] = ae; + } + if i > 0 && !hole(i - 1, j) { + p.aw[idx] = ae; + } + if j + 1 < ny && !hole(i, j + 1) { + p.an[idx] = an; + } + if j > 0 && !hole(i, j - 1) { + p.as_[idx] = an; + } + if i + 1 == nx { + p.extra_diag[idx] = 2.0 * ae; + } + } + } + let mut state = seed | 1; + for idx in 0..nx * ny { + 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 + }; + } + p +} + +struct DeviceLevel { + n: usize, + nx: i32, + cells: CudaSlice, + red: CudaSlice, + black: CudaSlice, + coarse_of: CudaSlice, + children_ptr: CudaSlice, + children_idx: CudaSlice, + n_cells: usize, + n_red: usize, + n_black: usize, + ae: CudaSlice, + aw: CudaSlice, + an: CudaSlice, + as_: CudaSlice, + ap: CudaSlice, + b: CudaSlice, + x: CudaSlice, + r: CudaSlice, +} + +fn cfg(n_items: usize, k: usize) -> LaunchConfig { + LaunchConfig { + grid_dim: ((n_items as u32).div_ceil(256).max(1), k as u32, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + } +} + +#[test] +#[ignore] +fn batched_device_vcycle_go_no_go() { + let (nx, ny) = (372usize, 62usize); + let prob = problem(nx, ny, 3); + let params = MultigridParameters { + smoother: MgSmoother::RedBlack, + ..MultigridParameters::default() + }; + let levels: Vec = export_hierarchy(&prob, ¶ms); + let depth = levels.len(); + println!( + " hierarchy: {depth} levels, cells per level {:?}", + levels.iter().map(|l| l.cells.len()).collect::>() + ); + + // Device. + let ctx = CudaContext::new(0).expect("CUDA context"); + 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.clone().into_boxed_str())), + ..Default::default() + }, + ) + .expect("nvrtc"); + let module = ctx.load_module(ptx).expect("module"); + let f_half = module.load_function("mg_rb_half").unwrap(); + let f_res = module.load_function("mg_residual").unwrap(); + let f_restrict = module.load_function("mg_restrict").unwrap(); + let f_prolong = module.load_function("mg_prolong").unwrap(); + let f_zero = module.load_function("mg_zero").unwrap(); + let f_coarsest = module.load_function("mg_coarsest").unwrap(); + + // The CPU reference V-cycle on the fine right-hand side. + let n0 = nx * ny; + let r_host: Vec = prob.rhs.clone(); + let mut z_ref = vec![0.0; n0]; + vcycle_f32_reference(&prob, ¶ms, &r_host, &mut z_ref); + let z_scale = z_ref.iter().fold(0.0_f64, |m, v| m.max(v.abs())); + + // The CPU f64 serial red-black V-cycle's wall time (the comparison point). + let t0 = Instant::now(); + let reps = 20; + for _ in 0..reps { + let mut z = vec![0.0; n0]; + vcycle_f32_reference(&prob, ¶ms, &r_host, &mut z); + } + let cpu_ms = t0.elapsed().as_secs_f64() * 1e3 / reps as f64; + println!(" CPU f32 red-black V-cycle (incl. hierarchy build): {cpu_ms:.3} ms"); + + let sweeps = 2usize; + for &k in &[1usize, 8, 16] { + // Upload: coefficients replicated K times, index lists shared. + let mut dev: Vec = Vec::with_capacity(depth); + for l in &levels { + let n = l.nx * l.ny; + let rep = |v: &[f32]| -> Vec { + let mut out = Vec::with_capacity(v.len() * k); + for _ in 0..k { + out.extend_from_slice(v); + } + out + }; + let up_u = |v: &[u32]| stream.memcpy_stod(v).unwrap(); + let up_f = |v: &[f32]| stream.memcpy_stod(v).unwrap(); + dev.push(DeviceLevel { + n, + nx: l.nx as i32, + cells: up_u(&l.cells), + red: up_u(&l.red), + black: up_u(&l.black), + coarse_of: up_u(if l.coarse_of.is_empty() { &[0u32][..] } else { &l.coarse_of }), + children_ptr: up_u(if l.children_ptr.is_empty() { &[0u32][..] } else { &l.children_ptr }), + children_idx: up_u(if l.children_idx.is_empty() { &[0u32][..] } else { &l.children_idx }), + n_cells: l.cells.len(), + n_red: l.red.len(), + n_black: l.black.len(), + ae: up_f(&rep(&l.ae)), + aw: up_f(&rep(&l.aw)), + an: up_f(&rep(&l.an)), + as_: up_f(&rep(&l.as_)), + ap: up_f(&rep(&l.ap)), + b: stream.alloc_zeros::(n * k).unwrap(), + x: stream.alloc_zeros::(n * k).unwrap(), + r: stream.alloc_zeros::(n * k).unwrap(), + }); + } + stream.synchronize().unwrap(); + + let r_f32: Vec = { + let one: Vec = r_host.iter().map(|&v| v as f32).collect(); + let mut out = Vec::with_capacity(n0 * k); + for _ in 0..k { + out.extend_from_slice(&one); + } + out + }; + let mut z_out = vec![0.0f32; n0 * k]; + + let vcycle = |dev: &mut Vec, r_f32: &[f32], z_out: &mut [f32]| { + // Upload the residual into level 0's b. + stream.memcpy_htod(r_f32, &mut dev[0].b).unwrap(); + let half = |lv: &mut DeviceLevel, colour: u8| { + let (list, n_list) = if colour == 0 { (&lv.red, lv.n_red) } else { (&lv.black, lv.n_black) }; + let n_i = lv.n as i32; + let n_list_i = n_list as i32; + unsafe { + stream + .launch_builder(&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, k)) + .unwrap(); + } + }; + let smooth = |lv: &mut DeviceLevel| { + for _ in 0..sweeps { + half(lv, 0); + half(lv, 1); + half(lv, 1); + half(lv, 0); + } + }; + // Down. + for l in 0..depth - 1 { + { + let lv = &mut dev[l]; + let (n_i, n_cells_i) = (lv.n as i32, lv.n_cells as i32); + unsafe { + stream + .launch_builder(&f_zero) + .arg(&n_cells_i) + .arg(&lv.cells) + .arg(&n_i) + .arg(&mut lv.x) + .launch(cfg(lv.n_cells, k)) + .unwrap(); + } + smooth(lv); + unsafe { + stream + .launch_builder(&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, k)) + .unwrap(); + } + } + let (fine, coarse) = dev.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 { + stream + .launch_builder(&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, k)) + .unwrap(); + } + } + // Coarsest. + { + let lv = &mut dev[depth - 1]; + let (k_i, n_cells_i, n_i, sw_i) = (k as i32, lv.n_cells as i32, lv.n as i32, 50i32); + unsafe { + stream + .launch_builder(&f_coarsest) + .arg(&k_i) + .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(&mut lv.x) + .arg(&lv.nx) + .arg(&sw_i) + .launch(LaunchConfig { + grid_dim: ((k as u32).div_ceil(32), 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }) + .unwrap(); + } + } + // Up. + for l in (0..depth - 1).rev() { + let (fine, coarse) = dev.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 { + stream + .launch_builder(&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, k)) + .unwrap(); + } + smooth(lf); + } + stream.memcpy_dtoh(&dev[0].x, z_out).unwrap(); + stream.synchronize().unwrap(); + }; + + // Correctness at K = 1 (and every batch member at K > 1). + vcycle(&mut dev, &r_f32, &mut z_out); + let mut worst = 0.0_f64; + for m in 0..k { + for idx in 0..n0 { + let d = (z_out[m * n0 + idx] as f64 - z_ref[idx]).abs(); + worst = worst.max(d); + } + } + println!(" K = {k}: device V-cycle vs CPU f32 reference: max |Δz| {worst:.3e} on a scale of {z_scale:.3e}"); + assert!(worst < 1e-4 * z_scale, "device V-cycle disagrees with the CPU one"); + + // Timing: 50 V-cycles, per march. + let reps = 50; + vcycle(&mut dev, &r_f32, &mut z_out); + let t0 = Instant::now(); + for _ in 0..reps { + vcycle(&mut dev, &r_f32, &mut z_out); + } + let ms = t0.elapsed().as_secs_f64() * 1e3 / reps as f64; + println!( + " K = {k}: {ms:.3} ms per batched V-cycle = {:.3} ms per march (CPU serial {cpu_ms:.3} ms; ratio {:.2}×)", + ms / k as f64, + cpu_ms / (ms / k as f64) + ); + } +}