embedded3 PERF-3 P2: the Poisson initial guess by projection onto the last K corrections of the same corrector index (GuessBasis, DeviceCg::project_guess: Gram matrix rebuilt per solve, Cholesky with pivot dropping; e3_cg_axpy) behind RTX_E3_POISSON_PROJECT=K (default 0, digit-identical); static DFG 2D ny 122: CG 9.5 -> 2.5 per step, loads identical to 4 digits; the moving flag: -4 % only (consecutive corrections nearly uncorrelated); RTX_E3_POISSON_PROJECT_DEBUG prints the residual before / after
CI / Build (macos-latest) (push) Waiting to run
CI / Test (macos-latest) (push) Blocked by required conditions
CI / Test (ubuntu-latest) (push) Blocked by required conditions
CI / Python Bindings (maturin) (macos-latest) (push) Blocked by required conditions
CI / Python Bindings (maturin) (ubuntu-latest) (push) Blocked by required conditions
CI / WASM Build + Size Check (push) Blocked by required conditions
CI / Distributed Training Tests (push) Blocked by required conditions
CI / CI Success (push) Blocked by required conditions
CI / Build CPU-Only (Explicit) (push) Failing after 5s
Documentation / Build API Documentation (push) Failing after 5s
Documentation / Build User Guide (push) Successful in 5s
CI / Format Check (push) Failing after 15s
CI / Build (ubuntu-latest) (push) Failing after 2m0s
CI / Clippy Check (push) Failing after 2m19s
Performance Benchmarks / Run Benchmarks (push) Successful in 2m57s

Co-Authored-By: Claude Fable 5.1 <[email protected]>
This commit is contained in:
Omar Sobh
2026-09-19 18:57:50 -05:00
co-authored by Claude Fable 5.1
parent 845e0ae01a
commit d4cd7d9545
5 changed files with 309 additions and 1 deletions
@@ -138,6 +138,17 @@ extern "C" __global__ void e3_cg_axpy2(
r[g] -= alpha * q[g]; 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. */ /* d = z + beta d. */
extern "C" __global__ void e3_cg_xpay( extern "C" __global__ void e3_cg_xpay(
int n_cells, const unsigned int* __restrict__ cells, double beta, int n_cells, const unsigned int* __restrict__ cells, double beta,
@@ -23,6 +23,7 @@ struct CgKernels {
reduce: CudaFunction, reduce: CudaFunction,
axpy2: CudaFunction, axpy2: CudaFunction,
xpay: CudaFunction, xpay: CudaFunction,
axpy: CudaFunction,
copy: CudaFunction, copy: CudaFunction,
shift: CudaFunction, shift: CudaFunction,
gather_f32: CudaFunction, gather_f32: CudaFunction,
@@ -44,6 +45,7 @@ fn kernels() -> &'static CgKernels {
reduce: f("e3_cg_reduce"), reduce: f("e3_cg_reduce"),
axpy2: f("e3_cg_axpy2"), axpy2: f("e3_cg_axpy2"),
xpay: f("e3_cg_xpay"), xpay: f("e3_cg_xpay"),
axpy: f("e3_cg_axpy"),
copy: f("e3_cg_copy"), copy: f("e3_cg_copy"),
shift: f("e3_cg_shift"), shift: f("e3_cg_shift"),
gather_f32: f("e3_cg_gather_f32"), 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<CudaSlice<f64>>,
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<f64>) {
if self.cap == 0 {
return;
}
let rt = runtime();
if self.vecs.len() < self.cap {
let mut v = rt.stream.alloc_zeros::<f64>(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<f64>,
p: &mut CudaSlice<f64>,
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::<Vec<_>>()
);
}
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<f64>, p: &CudaSlice<f64>) -> 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<u32>,
partial: &mut CudaSlice<f64>,
scalar: &mut CudaSlice<f64>,
scalar_host: &mut Vec<f64>,
n_blocks: usize,
n_i: i32,
a: &CudaSlice<f64>,
b: &CudaSlice<f64>,
) -> 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<f64> {
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>( fn pick<'a>(
w: Which, w: Which,
r: &'a CudaSlice<f64>, r: &'a CudaSlice<f64>,
@@ -13,6 +13,7 @@ mod hierarchy;
mod pcg; mod pcg;
mod problem; mod problem;
pub use device_cg::GuessBasis;
pub use export::{LevelExport, export_hierarchy, vcycle_f32_reference}; pub use export::{LevelExport, export_hierarchy, vcycle_f32_reference};
pub use hierarchy::Hierarchy; pub use hierarchy::Hierarchy;
pub use pcg::{PcgCache, solve_pcg, solve_pcg_cached}; pub use pcg::{PcgCache, solve_pcg, solve_pcg_cached};
@@ -151,6 +151,8 @@ pub struct DeviceStep {
n_blocks: usize, n_blocks: usize,
cg: Option<DeviceCg>, cg: Option<DeviceCg>,
cg_dt: f64, cg_dt: f64,
/// PERF-3 P2: per corrector index, the basis of the projected initial guess.
pub(super) guess: Vec<crate::solvers::incompressible::embedded3::poisson::GuessBasis>,
timers: Option<StepTimers>, timers: Option<StepTimers>,
initialized: bool, initialized: bool,
/// A static cut-cell mask's tables (item 9b), when the solver has one. /// A static cut-cell mask's tables (item 9b), when the solver has one.
@@ -213,6 +215,7 @@ impl DeviceStep {
n_blocks, n_blocks,
cg: None, cg: None,
cg_dt: 0.0, cg_dt: 0.0,
guess: Vec::new(),
timers, timers,
initialized: false, initialized: false,
cut, cut,
@@ -7,6 +7,7 @@
use super::{DeviceStep, E3Params, E3Ptrs, StepResult}; use super::{DeviceStep, E3Params, E3Ptrs, StepResult};
use crate::solvers::incompressible::embedded3::Grid; use crate::solvers::incompressible::embedded3::Grid;
use crate::solvers::incompressible::embedded3::field::Field; 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::{cfg, load_module, runtime};
use crate::solvers::incompressible::embedded3::poisson::device_cg::DeviceCg; use crate::solvers::incompressible::embedded3::poisson::device_cg::DeviceCg;
use crate::solvers::incompressible::embedded3::step::Solver; use crate::solvers::incompressible::embedded3::step::Solver;
@@ -515,9 +516,18 @@ impl DeviceStep {
if corrector > 0 { if corrector > 0 {
rt.stream.memset_zeros(&mut self.p_prime).expect("p' = 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 sol = {
let cg = self.cg.as_mut().expect("cg"); 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; cg_iterations += sol.iterations;
t_poisson += tp.elapsed(); t_poisson += tp.elapsed();