From 5e2b56597177c16ae51ca88098c757aebbb807df Mon Sep 17 00:00:00 2001 From: Omar Sobh Date: Thu, 17 Sep 2026 14:52:18 -0500 Subject: [PATCH] =?UTF-8?q?rtx-cfd=20embedded3=20item=202:=20e3=5Fmg.cu=20?= =?UTF-8?q?+=20poisson::{export,=20device}=20(one=20shared=20CUDA=20runtim?= =?UTF-8?q?e=20and=20module=20loader=20for=20embedded3);=20gate=202=20HELD?= =?UTF-8?q?:=20device=20V-cycle=20=3D=20host=20f32=20to=204e-7=20relative?= =?UTF-8?q?=20on=2096=C3=9740=C3=97{1,8}=20and=20378=C3=9762=C3=9762,=204.?= =?UTF-8?q?97=20ms=20per=20V-cycle=20at=201.45=20M=20cells=20incl.=20trans?= =?UTF-8?q?fers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5.1 --- .../rtx-cfd/src/kernels/cuda/e3_mg.cu | 116 ++++++ .../embedded3/poisson/device.rs | 350 ++++++++++++++++++ .../embedded3/poisson/export.rs | 103 ++++++ .../incompressible/embedded3/poisson/mod.rs | 4 + .../rtx-cfd/tests/embedded3_gpu_vcycle.rs | 118 ++++++ 5 files changed, 691 insertions(+) create mode 100644 crates/specialized/rtx-cfd/src/kernels/cuda/e3_mg.cu create mode 100644 crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device.rs create mode 100644 crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/export.rs create mode 100644 crates/specialized/rtx-cfd/tests/embedded3_gpu_vcycle.rs diff --git a/crates/specialized/rtx-cfd/src/kernels/cuda/e3_mg.cu b/crates/specialized/rtx-cfd/src/kernels/cuda/e3_mg.cu new file mode 100644 index 0000000..f7c3b4b --- /dev/null +++ b/crates/specialized/rtx-cfd/src/kernels/cuda/e3_mg.cu @@ -0,0 +1,116 @@ +/** + * embedded3 item 2 (omni-cortex `docs/embedded3_campaign.md`): 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 e3_mg_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 e3_mg_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 e3_mg_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 e3_mg_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 e3_mg_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 e3_mg_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]; + } + } + } +} diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device.rs new file mode 100644 index 0000000..3e02738 --- /dev/null +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device.rs @@ -0,0 +1,350 @@ +//! The seven-point V-cycle on the CUDA device (`e3_mg.cu`), one march, +//! persistent buffers per operator. One CUDA runtime per process, shared +//! by every device module of `embedded3`. + +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/e3_mg.cu"); + +/// The process's CUDA context and default stream. +pub struct Runtime { + pub ctx: Arc, + pub stream: Arc, +} + +static RUNTIME: OnceLock = OnceLock::new(); + +pub fn runtime() -> &'static Runtime { + RUNTIME.get_or_init(|| { + let ctx = CudaContext::new(0).expect("CUDA context (device 0)"); + let stream = ctx.default_stream(); + Runtime { ctx, stream } + }) +} + +/// Compile a kernel source for the device's architecture (`RTX_CUDA_ARCH`, +/// default sm_120) and load it. `fmad_off` = the host's f64 arithmetic. +pub fn load_module(source: &str, name: &str, fmad_off: bool) -> Arc { + let rt = runtime(); + let arch = std::env::var("RTX_CUDA_ARCH").unwrap_or_else(|_| "sm_120".to_string()); + let mut options = CompileOptions { + arch: Some(Box::leak(arch.into_boxed_str())), + ..Default::default() + }; + if fmad_off { + options.options.push("--fmad=false".to_string()); + } + let ptx = + compile_ptx_with_opts(source, options).unwrap_or_else(|e| panic!("nvrtc: {name}: {e:?}")); + rt.ctx + .load_module(ptx) + .unwrap_or_else(|e| panic!("module {name}: {e:?}")) +} + +struct MgKernels { + _module: Arc, + half: CudaFunction, + residual: CudaFunction, + restrict: CudaFunction, + prolong: CudaFunction, + zero: CudaFunction, + coarsest: CudaFunction, +} + +static MG: OnceLock = OnceLock::new(); + +fn kernels() -> &'static MgKernels { + MG.get_or_init(|| { + let module = load_module(KERNELS, "e3_mg.cu", false); + let f = |name: &str| module.load_function(name).expect(name); + MgKernels { + half: f("e3_mg_rb_half"), + residual: f("e3_mg_residual"), + restrict: f("e3_mg_restrict"), + prolong: f("e3_mg_prolong"), + zero: f("e3_mg_zero"), + coarsest: f("e3_mg_coarsest"), + _module: module, + } + }) +} + +pub(crate) struct DevLevel { + 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, + pub(crate) red: CudaSlice, + pub(crate) black: CudaSlice, + pub(crate) top: CudaSlice, + pub(crate) bot: CudaSlice, + pub(crate) coarse_of: CudaSlice, + pub(crate) children_ptr: CudaSlice, + pub(crate) children_idx: CudaSlice, + pub(crate) ae: CudaSlice, + pub(crate) aw: CudaSlice, + pub(crate) an: CudaSlice, + pub(crate) as_: CudaSlice, + pub(crate) at: CudaSlice, + pub(crate) ab: CudaSlice, + pub(crate) ap: CudaSlice, + pub(crate) b: CudaSlice, + pub(crate) x: CudaSlice, + pub(crate) r: CudaSlice, +} + +/// One operator's hierarchy on the device. +pub struct DeviceVcycle { + pub(crate) levels: Vec, + sweeps: usize, + fine_cells: Vec, + r_f32: Vec, + z_f32: Vec, +} + +pub(crate) 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 fn new(levels: &[LevelExport], sweeps: usize) -> Self { + let rt = runtime(); + let up_u = |v: &[u32]| -> CudaSlice { + rt.stream + .memcpy_stod(if v.is_empty() { &[0u32][..] } else { v }) + .expect("upload") + }; + let up_f = |v: &[f32]| -> CudaSlice { rt.stream.memcpy_stod(v).expect("upload") }; + let dev: Vec = levels + .iter() + .map(|l| { + let n = l.nx * l.ny * l.nz; + 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), + 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::(n).expect("alloc"), + x: rt.stream.alloc_zeros::(n).expect("alloc"), + r: rt.stream.alloc_zeros::(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 = runtime(); + let k = kernels(); + 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(&k.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("e3_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); + } + } + + /// The V-cycle 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 = runtime(); + let k = kernels(); + 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(&k.zero) + .arg(&n_cells_i) + .arg(&lv.cells) + .arg(&mut lv.x) + .launch(cfg(lv.n_cells)) + .expect("e3_mg_zero"); + } + } + self.smooth(l); + { + let lv = &mut self.levels[l]; + let n_cells_i = lv.n_cells as i32; + unsafe { + rt.stream + .launch_builder(&k.residual) + .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("e3_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 = lc.n_cells as i32; + unsafe { + rt.stream + .launch_builder(&k.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("e3_mg_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(&k.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("e3_mg_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(&k.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("e3_mg_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 = runtime(); + 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; + } + } +} diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/export.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/export.rs new file mode 100644 index 0000000..f6b6f67 --- /dev/null +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/export.rs @@ -0,0 +1,103 @@ +//! The f32 hierarchy exported for the device V-cycle: per level the index +//! lists (cells, colours, the neighbour arrays, the parent map, the CSR +//! children) and the seven coefficient arrays. + +use super::{Hierarchy, Problem}; +use crate::solvers::incompressible::poisson::MultigridParameters; + +/// One exported level. `top`/`bot`: the neighbour above/below each cell +/// (`u32::MAX` = none; a zero coefficient is never read). +pub struct LevelExport { + pub nx: usize, + pub ny: usize, + pub nz: usize, + pub cells: Vec, + pub red: Vec, + pub black: Vec, + pub top: Vec, + pub bot: 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). + pub children_ptr: Vec, + pub children_idx: Vec, + pub ae: Vec, + pub aw: Vec, + pub an: Vec, + pub as_: Vec, + pub at: Vec, + pub ab: Vec, + pub ap: Vec, +} + +/// The f32 hierarchy of `problem`, level 0 fine. +pub fn export_hierarchy(problem: &Problem, params: &MultigridParameters) -> Vec { + let hier = Hierarchy::::build(problem, params); + 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::>() + }; + (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![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, + 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_reference( + problem: &Problem, + params: &MultigridParameters, + r: &[f64], + z: &mut [f64], +) { + let mut hier = Hierarchy::::build(problem, params); + hier.apply_preconditioner(r, z); +} diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/mod.rs index 8e1da7f..624ef7d 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/mod.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/mod.rs @@ -4,10 +4,14 @@ //! rule for rule. At `nz = 1` with zero z-coefficients the arithmetic is the //! 2D solver's in the same order (gate 1: bit identity). +#[cfg(feature = "cuda")] +pub mod device; +mod export; mod hierarchy; mod pcg; mod problem; +pub use export::{LevelExport, export_hierarchy, vcycle_f32_reference}; pub use hierarchy::Hierarchy; pub use pcg::{PcgCache, solve_pcg, solve_pcg_cached}; pub use problem::Problem; diff --git a/crates/specialized/rtx-cfd/tests/embedded3_gpu_vcycle.rs b/crates/specialized/rtx-cfd/tests/embedded3_gpu_vcycle.rs new file mode 100644 index 0000000..eb37036 --- /dev/null +++ b/crates/specialized/rtx-cfd/tests/embedded3_gpu_vcycle.rs @@ -0,0 +1,118 @@ +//! embedded3 gate 2: the device seven-point V-cycle against the host f32 +//! reference. The cheap pin runs on a small channel with a z-cylinder hole; +//! `bench_anchor_size` (ignored) is 378 × 62 × 62 and prints ms per V-cycle. +//! +//! `RTX_CUDA_ARCH=sm_120 cargo test --release -p rtx-cfd --features cuda --test embedded3_gpu_vcycle -- --nocapture [--ignored]` +#![cfg(feature = "cuda")] + +use rtx_cfd::solvers::incompressible::embedded3::poisson::device::DeviceVcycle; +use rtx_cfd::solvers::incompressible::embedded3::poisson::{ + Problem, export_hierarchy, vcycle_f32_reference, +}; +use rtx_cfd::solvers::incompressible::{MgSmoother, MultigridParameters}; +use std::time::Instant; + +fn channel(nx: usize, ny: usize, nz: usize, periodic_z: bool, seed: u64) -> Problem { + let mut p = Problem::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 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 + }; + } + p +} + +fn compare(p: &Problem, label: &str) -> (f64, f64) { + let params = MultigridParameters { + smoother: MgSmoother::RedBlack, + ..MultigridParameters::default() + }; + let n = p.nx * p.ny * p.nz; + let mut z_ref = vec![0.0; n]; + vcycle_f32_reference(p, ¶ms, &p.rhs, &mut z_ref); + let levels = export_hierarchy(p, ¶ms); + let mut dev = DeviceVcycle::new(&levels, params.smoother_sweeps.max(1)); + let mut z = vec![0.0; n]; + dev.apply(&p.rhs, &mut z); + let scale = z_ref.iter().fold(0.0_f64, |m, v| m.max(v.abs())); + let worst = z + .iter() + .zip(&z_ref) + .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs())); + let reps = 20; + dev.apply(&p.rhs, &mut z); + let t0 = Instant::now(); + for _ in 0..reps { + dev.apply(&p.rhs, &mut z); + } + let ms = t0.elapsed().as_secs_f64() * 1e3 / reps as f64; + println!( + " {label}: {} levels {:?}; device vs host f32 max |Δz| {worst:.3e} on {scale:.3e}; {ms:.3} ms per V-cycle incl. transfers ({n} cells)", + levels.len(), + levels.iter().map(|l| l.cells.len()).collect::>() + ); + (worst, scale) +} + +#[test] +fn device_vcycle_matches_the_host_reference() { + for (nz, periodic) in [(1usize, false), (8, true), (8, false)] { + let (worst, scale) = compare( + &channel(96, 40, nz, periodic, 3), + &format!("96×40×{nz} periodic {periodic}"), + ); + assert!(worst < 1e-4 * scale, "{worst:.3e} of {scale:.3e}"); + } +} + +#[test] +#[ignore = "the anchor-size bench (378 × 62 × 62): prints ms per V-cycle"] +fn bench_anchor_size() { + let (worst, scale) = compare(&channel(378, 62, 62, false, 11), "378×62×62 walls"); + assert!(worst < 1e-4 * scale); +}