diff --git a/crates/specialized/rtx-cfd/src/kernels/cuda/e3_cg.cu b/crates/specialized/rtx-cfd/src/kernels/cuda/e3_cg.cu index c72e41b..f92ea06 100644 --- a/crates/specialized/rtx-cfd/src/kernels/cuda/e3_cg.cu +++ b/crates/specialized/rtx-cfd/src/kernels/cuda/e3_cg.cu @@ -138,6 +138,17 @@ extern "C" __global__ void e3_cg_axpy2( r[g] -= alpha * q[g]; } +/* y += alpha x over the cells (PERF-3 P2: the projected initial guess). */ +extern "C" __global__ void e3_cg_axpy( + int n_cells, const unsigned int* __restrict__ cells, double alpha, + const double* __restrict__ x, double* __restrict__ y) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= n_cells) return; + int g = cells[t]; + y[g] += alpha * x[g]; +} + /* d = z + beta d. */ extern "C" __global__ void e3_cg_xpay( int n_cells, const unsigned int* __restrict__ cells, double beta, diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device_cg.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device_cg.rs index a7e4ea2..f53deb6 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device_cg.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device_cg.rs @@ -23,6 +23,7 @@ struct CgKernels { reduce: CudaFunction, axpy2: CudaFunction, xpay: CudaFunction, + axpy: CudaFunction, copy: CudaFunction, shift: CudaFunction, gather_f32: CudaFunction, @@ -44,6 +45,7 @@ fn kernels() -> &'static CgKernels { reduce: f("e3_cg_reduce"), axpy2: f("e3_cg_axpy2"), xpay: f("e3_cg_xpay"), + axpy: f("e3_cg_axpy"), copy: f("e3_cg_copy"), shift: f("e3_cg_shift"), gather_f32: f("e3_cg_gather_f32"), @@ -686,6 +688,287 @@ impl DeviceCg { } } +/// PERF-3 P2: the last K solutions of one corrector index, the basis of the +/// projected initial guess (`RTX_E3_POISSON_PROJECT=K`). +pub struct GuessBasis { + vecs: Vec>, + next: usize, + cap: usize, + n: usize, +} + +impl GuessBasis { + /// `K` from `RTX_E3_POISSON_PROJECT` (0 = off). + #[must_use] + pub fn from_env(n: usize) -> Self { + let cap = std::env::var("RTX_E3_POISSON_PROJECT") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(0); + Self { + vecs: Vec::new(), + next: 0, + cap, + n, + } + } + + pub fn is_on(&self) -> bool { + self.cap > 0 + } + + /// Keep `p` (the oldest goes when the basis is full). + pub fn push(&mut self, p: &CudaSlice) { + if self.cap == 0 { + return; + } + let rt = runtime(); + if self.vecs.len() < self.cap { + let mut v = rt.stream.alloc_zeros::(self.n).expect("basis"); + rt.stream.memcpy_dtod(p, &mut v).expect("basis copy"); + self.vecs.push(v); + } else { + let slot = self.next; + rt.stream + .memcpy_dtod(p, &mut self.vecs[slot]) + .expect("basis copy"); + self.next = (slot + 1) % self.cap; + } + } +} + +impl DeviceCg { + /// `p := Σ α_j x_j`, the A-norm best combination of the basis for the + /// right-hand side on the current cells: `G α = r` with + /// `G_jk = ⟨x_j, A x_k⟩`, `r_j = ⟨x_j, b⟩` (K spmv + K² dots on one + /// scratch vector; the operator changes every step on a moving body, so + /// the Gram matrix is rebuilt, not stored), Cholesky on the host, a + /// pivot under 1e-12 of the diagonal's maximum drops its vector. + /// Returns the number of vectors used. + pub fn project_guess( + &mut self, + rhs: &CudaSlice, + p: &mut CudaSlice, + basis: &GuessBasis, + ) -> usize { + let m = basis.vecs.len(); + if m == 0 || self.n_cells == 0 { + return 0; + } + let rt = runtime(); + let k = kernels(); + let n_i = self.n_cells as i32; + let mut gram = vec![0.0f64; m * m]; + let mut rhs_dots = vec![0.0f64; m]; + for (col, xk) in basis.vecs.iter().enumerate() { + // q = A x_k + 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.link_ptr) + .arg(&self.link_idx) + .arg(&self.link_coef) + .arg(&self.ap) + .arg(xk) + .arg(&mut self.q) + .arg(&self.nx) + .launch(cfg(self.n_cells)) + .expect("e3_cg_spmv"); + } + let Self { + cells, + partial, + scalar, + scalar_host, + q, + n_blocks, + .. + } = self; + for (row, xj) in basis.vecs.iter().enumerate() { + gram[row * m + col] = + dot_raw(cells, partial, scalar, scalar_host, *n_blocks, n_i, xj, q); + } + rhs_dots[col] = dot_raw(cells, partial, scalar, scalar_host, *n_blocks, n_i, xk, rhs); + } + // Symmetrise (the two orders of a dot differ at round-off) and solve. + for j in 0..m { + for c in j + 1..m { + let v = 0.5 * (gram[j * m + c] + gram[c * m + j]); + gram[j * m + c] = v; + gram[c * m + j] = v; + } + } + let alpha = cholesky_solve(&gram, &rhs_dots, m); + let used = alpha.iter().filter(|a| **a != 0.0).count(); + if used == 0 { + return 0; + } + let debug = std::env::var("RTX_E3_POISSON_PROJECT_DEBUG").is_ok(); + let before = if debug { self.residual_of(rhs, p) } else { 0.0 }; + rt.stream.memset_zeros(p).expect("p = 0"); + for (a, x) in alpha.iter().zip(&basis.vecs) { + if *a == 0.0 { + continue; + } + unsafe { + rt.stream + .launch_builder(&k.axpy) + .arg(&n_i) + .arg(&self.cells) + .arg(a) + .arg(x) + .arg(&mut *p) + .launch(cfg(self.n_cells)) + .expect("e3_cg_axpy"); + } + } + if debug { + let after = self.residual_of(rhs, p); + eprintln!( + " projection: {used} of {m} vectors, |r| {before:.3e} -> {after:.3e} (alpha {:?})", + alpha.iter().map(|a| format!("{a:.3}")).collect::>() + ); + } + used + } + + /// `Σ |b − A p|` over the cells (the CG's own residual measure), without + /// touching the solve's state beyond `r`. + fn residual_of(&mut self, rhs: &CudaSlice, p: &CudaSlice) -> f64 { + let rt = runtime(); + let k = kernels(); + let n_i = self.n_cells as i32; + 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("e3_cg_copy"); + } + self.true_residual(p) + } +} + +/// `⟨a, b⟩` over the cells with the solver's partial-sum buffers (the same +/// kernels and reduction order as `DeviceCg::dot`). +#[allow(clippy::too_many_arguments)] +fn dot_raw( + cells: &CudaSlice, + partial: &mut CudaSlice, + scalar: &mut CudaSlice, + scalar_host: &mut Vec, + n_blocks: usize, + n_i: i32, + a: &CudaSlice, + b: &CudaSlice, +) -> f64 { + let rt = runtime(); + let k = kernels(); + let n_cells = n_i as usize; + unsafe { + rt.stream + .launch_builder(&k.dot_partial) + .arg(&n_i) + .arg(cells) + .arg(a) + .arg(b) + .arg(&mut *partial) + .launch(cfg(n_cells)) + .expect("e3_cg_dot_partial"); + let nb = n_blocks as i32; + rt.stream + .launch_builder(&k.reduce) + .arg(&nb) + .arg(&*partial) + .arg(&mut *scalar) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }) + .expect("e3_cg_reduce"); + } + rt.stream + .memcpy_dtoh(&*scalar, scalar_host) + .expect("scalar"); + rt.stream.synchronize().expect("sync"); + scalar_host[0] +} + +/// Cholesky solve of a small SPD system with pivot dropping: a row whose +/// pivot falls under 1e-12 × the largest diagonal gets α = 0 (its vector is +/// nearly dependent on the ones before it). +fn cholesky_solve(g: &[f64], r: &[f64], m: usize) -> Vec { + let gmax = (0..m).map(|j| g[j * m + j]).fold(0.0f64, f64::max); + let floor = 1e-12 * gmax; + let mut l = vec![0.0f64; m * m]; + let mut keep = vec![true; m]; + for j in 0..m { + let mut d = g[j * m + j]; + for c in 0..j { + if keep[c] { + d -= l[j * m + c] * l[j * m + c]; + } + } + if !(d > floor) || !d.is_finite() { + keep[j] = false; + continue; + } + let dj = d.sqrt(); + l[j * m + j] = dj; + for i in j + 1..m { + let mut v = g[i * m + j]; + for c in 0..j { + if keep[c] { + v -= l[i * m + c] * l[j * m + c]; + } + } + l[i * m + j] = v / dj; + } + } + // L y = r, Lᵀ α = y over the kept rows. + let mut y = vec![0.0f64; m]; + for i in 0..m { + if !keep[i] { + continue; + } + let mut v = r[i]; + for c in 0..i { + if keep[c] { + v -= l[i * m + c] * y[c]; + } + } + y[i] = v / l[i * m + i]; + } + let mut alpha = vec![0.0f64; m]; + for i in (0..m).rev() { + if !keep[i] { + continue; + } + let mut v = y[i]; + for c in i + 1..m { + if keep[c] { + v -= l[c * m + i] * alpha[c]; + } + } + alpha[i] = v / l[i * m + i]; + } + alpha +} + fn pick<'a>( w: Which, r: &'a CudaSlice, 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 4d66ad2..e91763b 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 @@ -13,6 +13,7 @@ mod hierarchy; mod pcg; mod problem; +pub use device_cg::GuessBasis; pub use export::{LevelExport, export_hierarchy, vcycle_f32_reference}; pub use hierarchy::Hierarchy; pub use pcg::{PcgCache, solve_pcg, solve_pcg_cached}; diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device.rs index d82f2a3..72ee676 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device.rs @@ -151,6 +151,8 @@ pub struct DeviceStep { n_blocks: usize, cg: Option, cg_dt: f64, + /// PERF-3 P2: per corrector index, the basis of the projected initial guess. + pub(super) guess: Vec, timers: Option, initialized: bool, /// A static cut-cell mask's tables (item 9b), when the solver has one. @@ -213,6 +215,7 @@ impl DeviceStep { n_blocks, cg: None, cg_dt: 0.0, + guess: Vec::new(), timers, initialized: false, cut, diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device/cut.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device/cut.rs index 9b84161..4014371 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device/cut.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device/cut.rs @@ -7,6 +7,7 @@ use super::{DeviceStep, E3Params, E3Ptrs, StepResult}; use crate::solvers::incompressible::embedded3::Grid; use crate::solvers::incompressible::embedded3::field::Field; +use crate::solvers::incompressible::embedded3::poisson::GuessBasis; use crate::solvers::incompressible::embedded3::poisson::device::{cfg, load_module, runtime}; use crate::solvers::incompressible::embedded3::poisson::device_cg::DeviceCg; use crate::solvers::incompressible::embedded3::step::Solver; @@ -515,9 +516,18 @@ impl DeviceStep { if corrector > 0 { rt.stream.memset_zeros(&mut self.p_prime).expect("p' = 0"); } + // PERF-3 P2: the projected initial guess from this corrector's previous solutions. + while self.guess.len() <= corrector { + self.guess.push(GuessBasis::from_env(g.cells())); + } let sol = { let cg = self.cg.as_mut().expect("cg"); - cg.solve_device(&self.sp, &mut self.p_prime, inner_stop, anchor, 0) + if self.guess[corrector].is_on() { + cg.project_guess(&self.sp, &mut self.p_prime, &self.guess[corrector]); + } + let sol = cg.solve_device(&self.sp, &mut self.p_prime, inner_stop, anchor, 0); + self.guess[corrector].push(&self.p_prime); + sol }; cg_iterations += sol.iterations; t_poisson += tp.elapsed();