diff --git a/crates/specialized/rtx-cfd/src/kernels/cuda/cg3.cu b/crates/specialized/rtx-cfd/src/kernels/cuda/cg3.cu new file mode 100644 index 0000000..9746b7e --- /dev/null +++ b/crates/specialized/rtx-cfd/src/kernels/cuda/cg3.cu @@ -0,0 +1,181 @@ +/** + * 3D Stage 1, gate 3: the f64 conjugate gradient's maps on the device over + * the active cells (index list), with FIXED-ORDER reductions: every block + * reduces its 256 lanes by the same shared-memory tree and writes one + * partial; `cg3_reduce` sums the partials in index order on one thread. + * Run-to-run bit-identical by construction (no atomics). + */ +#define NONE 0xFFFFFFFFu + +__device__ __forceinline__ double nb_sum3d( + int g, int nx, const double* ae, const double* aw, const double* an, const double* as_, + const double* at, const double* ab, const unsigned int* top, const unsigned int* bot, + const double* x) +{ + double s = 0.0; + double e = ae[g]; if (e != 0.0) s += e * x[g + 1]; + double w = aw[g]; if (w != 0.0) s += w * x[g - 1]; + double nn = an[g]; if (nn != 0.0) s += nn * x[g + nx]; + double ss = as_[g]; if (ss != 0.0) s += ss * x[g - nx]; + double t = at[g]; if (t != 0.0) s += t * x[top[g]]; + double b = ab[g]; if (b != 0.0) s += b * x[bot[g]]; + return s; +} + +__device__ __forceinline__ double block_reduce(double v) +{ + __shared__ double sh[256]; + int t = threadIdx.x; + sh[t] = v; + __syncthreads(); + for (int s = 128; s > 0; s >>= 1) { + if (t < s) sh[t] += sh[t + s]; + __syncthreads(); + } + return sh[0]; +} + +/* q = A d on the cells. */ +extern "C" __global__ void cg3_spmv( + int n_cells, const unsigned int* __restrict__ cells, + const double* __restrict__ ae, const double* __restrict__ aw, + const double* __restrict__ an, const double* __restrict__ as_, + const double* __restrict__ at, const double* __restrict__ ab, + const unsigned int* __restrict__ top, const unsigned int* __restrict__ bot, + const double* __restrict__ ap, const double* __restrict__ d, double* __restrict__ q, int nx) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= n_cells) return; + int g = cells[t]; + q[g] = ap[g] * d[g] - nb_sum3d(g, nx, ae, aw, an, as_, at, ab, top, bot, d); +} + +/* r = b − A p on the cells; partial[block] = Σ |r| over the block's cells. */ +extern "C" __global__ void cg3_residual( + int n_cells, const unsigned int* __restrict__ cells, + const double* __restrict__ ae, const double* __restrict__ aw, + const double* __restrict__ an, const double* __restrict__ as_, + const double* __restrict__ at, const double* __restrict__ ab, + const unsigned int* __restrict__ top, const unsigned int* __restrict__ bot, + const double* __restrict__ ap, const double* __restrict__ b, + const double* __restrict__ p, double* __restrict__ r, double* __restrict__ partial, int nx) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + double v = 0.0; + if (t < n_cells) { + int g = cells[t]; + v = b[g] - (ap[g] * p[g] - nb_sum3d(g, nx, ae, aw, an, as_, at, ab, top, bot, p)); + r[g] = v; + v = fabs(v); + } + double s = block_reduce(v); + if (threadIdx.x == 0) partial[blockIdx.x] = s; +} + +/* partial[block] = Σ a·b over the block's cells (b = a for a norm). */ +extern "C" __global__ void cg3_dot_partial( + int n_cells, const unsigned int* __restrict__ cells, + const double* __restrict__ a, const double* __restrict__ b, double* __restrict__ partial) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + double v = 0.0; + if (t < n_cells) { int g = cells[t]; v = a[g] * b[g]; } + double s = block_reduce(v); + if (threadIdx.x == 0) partial[blockIdx.x] = s; +} + +/* partial[block] = Σ |a| over the block's cells. */ +extern "C" __global__ void cg3_l1_partial( + int n_cells, const unsigned int* __restrict__ cells, + const double* __restrict__ a, double* __restrict__ partial) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + double v = 0.0; + if (t < n_cells) v = fabs(a[cells[t]]); + double s = block_reduce(v); + if (threadIdx.x == 0) partial[blockIdx.x] = s; +} + +/* partial[block] = Σ a over the block's cells. */ +extern "C" __global__ void cg3_sum_partial( + int n_cells, const unsigned int* __restrict__ cells, + const double* __restrict__ a, double* __restrict__ partial) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + double v = 0.0; + if (t < n_cells) v = a[cells[t]]; + double s = block_reduce(v); + if (threadIdx.x == 0) partial[blockIdx.x] = s; +} + +/* out[0] = Σ partial[0..n) in index order, one thread. */ +extern "C" __global__ void cg3_reduce(int n, const double* __restrict__ partial, double* __restrict__ out) +{ + if (blockIdx.x * blockDim.x + threadIdx.x != 0) return; + double s = 0.0; + for (int i = 0; i < n; ++i) s += partial[i]; + out[0] = s; +} + +/* p += alpha d; r -= alpha q. */ +extern "C" __global__ void cg3_axpy2( + int n_cells, const unsigned int* __restrict__ cells, double alpha, + const double* __restrict__ d, const double* __restrict__ q, + double* __restrict__ p, double* __restrict__ r) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= n_cells) return; + int g = cells[t]; + p[g] += alpha * d[g]; + r[g] -= alpha * q[g]; +} + +/* d = z + beta d. */ +extern "C" __global__ void cg3_xpay( + int n_cells, const unsigned int* __restrict__ cells, double beta, + const double* __restrict__ z, double* __restrict__ d) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= n_cells) return; + int g = cells[t]; + d[g] = z[g] + beta * d[g]; +} + +/* dst = src (copy over the cells). */ +extern "C" __global__ void cg3_copy( + int n_cells, const unsigned int* __restrict__ cells, + const double* __restrict__ src, double* __restrict__ dst) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= n_cells) return; + int g = cells[t]; + dst[g] = src[g]; +} + +/* v -= s over the cells. */ +extern "C" __global__ void cg3_shift( + int n_cells, const unsigned int* __restrict__ cells, double s, double* __restrict__ v) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= n_cells) return; + v[cells[t]] -= s; +} + +/* f32 b0 = (float) r over the cells; f64 z = (double) x0 over the cells. */ +extern "C" __global__ void cg3_gather_f32( + int n_cells, const unsigned int* __restrict__ cells, const double* __restrict__ r, float* __restrict__ b0) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= n_cells) return; + int g = cells[t]; + b0[g] = (float) r[g]; +} + +extern "C" __global__ void cg3_scatter_f64( + int n_cells, const unsigned int* __restrict__ cells, const float* __restrict__ x0, double* __restrict__ z) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= n_cells) return; + int g = cells[t]; + z[g] = (double) x0[g]; +} diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/poisson/device_cg.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/poisson/device_cg.rs new file mode 100644 index 0000000..a49fd21 --- /dev/null +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/poisson/device_cg.rs @@ -0,0 +1,646 @@ +//! Gate 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::{DeviceVcycle3, runtime3}; +use super::{Components3, Level3, OperatorKey3, PoissonProblem3D, export_hierarchy3}; +use crate::solvers::incompressible::poisson::{MultigridParameters, PoissonSolution}; +use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, LaunchConfig, PushKernelArg}; +use cudarc::nvrtc::{CompileOptions, compile_ptx_with_opts}; +use std::sync::{Arc, OnceLock}; + +const KERNELS: &str = include_str!("../../../../kernels/cuda/cg3.cu"); + +struct CgKernels { + _module: Arc, + 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 = OnceLock::new(); + +fn kernels() -> &'static CgKernels { + CG.get_or_init(|| { + let rt = runtime3(); + 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: cg3.cu"); + let module = rt.ctx.load_module(ptx).expect("cg3 module"); + let f = |name: &str| module.load_function(name).expect(name); + CgKernels { + spmv: f("cg3_spmv"), + residual: f("cg3_residual"), + dot_partial: f("cg3_dot_partial"), + l1_partial: f("cg3_l1_partial"), + sum_partial: f("cg3_sum_partial"), + reduce: f("cg3_reduce"), + axpy2: f("cg3_axpy2"), + xpay: f("cg3_xpay"), + copy: f("cg3_copy"), + shift: f("cg3_shift"), + gather_f32: f("cg3_gather_f32"), + scatter_f64: f("cg3_scatter_f64"), + _module: module, + } + }) +} + +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, + } +} + +/// The prepared operator on the device and the CG's work vectors. +pub struct DeviceCg3 { + key: OperatorKey3, + n: usize, + nx: i32, + n_cells: usize, + n_blocks: usize, + cells: CudaSlice, + top: CudaSlice, + bot: CudaSlice, + ae: CudaSlice, + aw: CudaSlice, + an: CudaSlice, + as_: CudaSlice, + at: CudaSlice, + ab: CudaSlice, + ap: CudaSlice, + b: CudaSlice, + r: CudaSlice, + z: CudaSlice, + d: CudaSlice, + q: CudaSlice, + partial: CudaSlice, + scalar: CudaSlice, + vcycle: DeviceVcycle3, + /// The singular component's members (all cells when singular; empty + /// otherwise) and whether the system is singular. + singular: bool, + active_host: Vec, + max_iterations: usize, + scalar_host: Vec, +} + +impl DeviceCg3 { + /// 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: &PoissonProblem3D, params: &MultigridParameters) -> Self { + let rt = runtime3(); + let fine = Level3::::new(problem.clone()); + let components = Components3::find(problem, &fine.cells); + let singular_count = components.singular.iter().filter(|&&s| s).count(); + assert!( + singular_count == 0 || components.members.len() == 1, + "DeviceCg3 (Stage 1): {} components with {} singular; one singular component is the scope", + components.members.len(), + singular_count + ); + let levels = export_hierarchy3(problem, params); + let vcycle = DeviceVcycle3::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::>() + }; + let up_u = |v: &[u32]| -> CudaSlice { + rt.stream + .memcpy_stod(if v.is_empty() { &[0u32][..] } else { v }) + .expect("upload") + }; + let up_f = |v: &[f64]| -> CudaSlice { rt.stream.memcpy_stod(v).expect("upload") }; + let zeros = || rt.stream.alloc_zeros::(n).expect("alloc"); + Self { + key: OperatorKey3::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::(n_blocks).expect("alloc"), + scalar: rt.stream.alloc_zeros::(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: &PoissonProblem3D, params: &MultigridParameters) -> bool { + self.key.matches(problem, params) + } + + pub fn n_cells(&self) -> usize { + self.n_cells + } + + fn reduce(&mut self) -> f64 { + let rt = runtime3(); + 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("cg3_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 = runtime3(); + 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("cg3_dot_partial"); + } + self.reduce() + } + + fn l1(&mut self, a: Which) -> f64 { + let rt = runtime3(); + 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("cg3_l1_partial"); + } + self.reduce() + } + + fn mean(&mut self, a: Which) -> f64 { + let rt = runtime3(); + 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("cg3_sum_partial"); + } + self.reduce() / self.n_cells as f64 + } + + fn shift(&mut self, a: Which, s: f64) { + let rt = runtime3(); + 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("cg3_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 { + let rt = runtime3(); + 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("cg3_residual"); + } + self.reduce() + } + + /// `z = M⁻¹ r` (gather to f32, the device V-cycle, scatter to f64). + fn precondition(&mut self) { + let rt = runtime3(); + 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("cg3_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("cg3_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, + p: &mut CudaSlice, + tolerance: f64, + anchor: Option, + setup_ns: u64, + ) -> PoissonSolution { + let rt = runtime3(); + 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("cg3_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, 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("cg3_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("cg3_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("cg3_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("cg3_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("cg3_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("cg3_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: &PoissonProblem3D, + p: &mut [f64], + tolerance: f64, + anchor: Option, + setup_ns: u64, + ) -> PoissonSolution { + let rt = runtime3(); + assert_eq!(p.len(), self.n); + let rhs_dev: CudaSlice = rt.stream.memcpy_stod(&problem.rhs).expect("rhs"); + let mut p_dev: CudaSlice = 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, + z: &'a CudaSlice, + d: &'a CudaSlice, + q: &'a CudaSlice, + b: &'a CudaSlice, +) -> &'a CudaSlice { + 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, + z: &'a mut CudaSlice, + d: &'a mut CudaSlice, + q: &'a mut CudaSlice, + b: &'a mut CudaSlice, +) -> &'a mut CudaSlice { + 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 DevicePcgCache3 { + slot: Option, +} + +/// The device counterpart of `solve_multigrid_pcg3_cached`. +pub fn solve_multigrid_pcg3_device_cached( + problem: &PoissonProblem3D, + p: &mut [f64], + params: &MultigridParameters, + tolerance: f64, + anchor: Option, + cache: &mut DevicePcgCache3, +) -> 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(DeviceCg3::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) +} diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/poisson/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/poisson/mod.rs index 5bcd731..b794573 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/poisson/mod.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/poisson/mod.rs @@ -13,6 +13,8 @@ use crate::solvers::incompressible::poisson::{ #[cfg(feature = "cuda")] pub mod device; +#[cfg(feature = "cuda")] +pub mod device_cg; pub mod export; pub use export::{LevelExport3, export_hierarchy3, vcycle_f32_reference3}; diff --git a/crates/specialized/rtx-cfd/tests/three_d_cg_device.rs b/crates/specialized/rtx-cfd/tests/three_d_cg_device.rs new file mode 100644 index 0000000..fa00525 --- /dev/null +++ b/crates/specialized/rtx-cfd/tests/three_d_cg_device.rs @@ -0,0 +1,179 @@ +//! 3D Stage 1, gate 3: the device-resident CG against the host PCG on the +//! same operator — converged, true residual below the tolerance, solutions +//! agreeing to the solve's accuracy, run-to-run bit-identical; the singular +//! closed box honours the anchor / the zero mean. +//! +//! `RTX_CUDA_ARCH=sm_120 cargo test --release -p rtx-cfd --features cuda --test three_d_cg_device -- --nocapture` +#![cfg(feature = "cuda")] + +use rtx_cfd::solvers::incompressible::three_d::poisson::device_cg::{ + DevicePcgCache3, solve_multigrid_pcg3_device_cached, +}; +use rtx_cfd::solvers::incompressible::three_d::poisson::{PoissonProblem3D, solve_multigrid_pcg3}; +use rtx_cfd::solvers::incompressible::{MgSmoother, MultigridParameters}; + +/// A channel on cubic cells with a z-cylinder hole and an outlet Dirichlet +/// (`dirichlet = true`), or a closed Neumann box (`false`), seeded rhs. +fn channel( + nx: usize, + ny: usize, + nz: usize, + periodic_z: bool, + dirichlet: bool, + seed: u64, +) -> PoissonProblem3D { + let mut p = PoissonProblem3D::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 dirichlet && 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 + }; + } + if !dirichlet { + // A compatible right-hand side: zero mean over the active cells. + let cells: Vec = (0..nx * ny * nz).filter(|&i| p.active[i]).collect(); + let mean = cells.iter().map(|&i| p.rhs[i]).sum::() / cells.len() as f64; + for &i in &cells { + p.rhs[i] -= mean; + } + } + p +} + +fn bits(v: &[f64]) -> Vec { + v.iter().map(|x| x.to_bits()).collect() +} + +#[test] +fn device_cg_matches_the_host_pcg() { + let params = MultigridParameters { + smoother: MgSmoother::RedBlack, + ..MultigridParameters::default() + }; + let tol = 1e-12; + let mut cache = DevicePcgCache3::default(); + for (nz, periodic, seed) in [(1usize, false, 3u64), (8, true, 5), (8, false, 7)] { + let p = channel(96, 40, nz, periodic, true, seed); + let n = p.nx * p.ny * p.nz; + let mut host = vec![0.0; n]; + let sh = solve_multigrid_pcg3(&p, &mut host, ¶ms, tol, None); + let mut dev = vec![0.0; n]; + let sd = solve_multigrid_pcg3_device_cached(&p, &mut dev, ¶ms, tol, None, &mut cache); + let mut dev2 = vec![0.0; n]; + let sd2 = solve_multigrid_pcg3_device_cached(&p, &mut dev2, ¶ms, tol, None, &mut cache); + assert!( + sh.converged && sd.converged, + "nz {nz}: host {} / device {}", + sh.converged, + sd.converged + ); + let res_dev = p.residual_l1(&dev); + assert!(res_dev < tol, "nz {nz}: device true residual {res_dev:.3e}"); + let scale = host.iter().fold(0.0_f64, |m, v| m.max(v.abs())); + let worst = host + .iter() + .zip(&dev) + .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs())); + println!( + " 96×40×{nz} periodic {periodic}: host {} it / device {} it (cache hit second solve: {} it); |Δp| {worst:.3e} of {scale:.3e}; device residual {res_dev:.3e}", + sh.iterations, sd.iterations, sd2.iterations + ); + assert!( + worst <= 1e-10 * scale, + "nz {nz}: device and host differ by {worst:.3e} of {scale:.3e}" + ); + assert_eq!( + bits(&dev), + bits(&dev2), + "nz {nz}: the device solve is not run-to-run identical" + ); + assert_eq!(sd.iterations, sd2.iterations); + } +} + +#[test] +fn the_singular_box_honours_the_anchor_and_the_mean() { + let params = MultigridParameters { + smoother: MgSmoother::RedBlack, + ..MultigridParameters::default() + }; + let tol = 1e-12; + let p = channel(48, 20, 8, false, false, 9); + let n = p.nx * p.ny * p.nz; + assert!(p.is_singular()); + let cells: Vec = (0..n).filter(|&i| p.active[i]).collect(); + let anchor = cells[cells.len() / 3]; + let mut cache = DevicePcgCache3::default(); + let mut host = vec![0.0; n]; + let sh = solve_multigrid_pcg3(&p, &mut host, ¶ms, tol, Some(anchor)); + let mut dev = vec![0.0; n]; + let sd = + solve_multigrid_pcg3_device_cached(&p, &mut dev, ¶ms, tol, Some(anchor), &mut cache); + assert!(sh.converged && sd.converged); + assert_eq!( + dev[anchor].to_bits(), + 0.0f64.to_bits(), + "anchor not at zero: {}", + dev[anchor] + ); + let scale = host.iter().fold(0.0_f64, |m, v| m.max(v.abs())); + let worst = host + .iter() + .zip(&dev) + .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs())); + println!( + " singular box with anchor: host {} it / device {} it; |Δp| {worst:.3e} of {scale:.3e}", + sh.iterations, sd.iterations + ); + assert!(worst <= 1e-10 * scale); + let mut dev0 = vec![0.0; n]; + let s0 = solve_multigrid_pcg3_device_cached(&p, &mut dev0, ¶ms, tol, None, &mut cache); + assert!(s0.converged); + let mean = cells.iter().map(|&i| dev0[i]).sum::() / cells.len() as f64; + println!(" singular box without anchor: mean {mean:.3e} of {scale:.3e}"); + assert!(mean.abs() <= 1e-13 * scale, "mean {mean:.3e}"); +}