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
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:
co-authored by
Claude Fable 5.1
parent
845e0ae01a
commit
d4cd7d9545
@@ -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();
|
||||||
|
|||||||
Reference in New Issue
Block a user