PERF-2 P3-ii: the device V-cycle as the CG's preconditioner — poisson/device.rs (one CUDA runtime per process, persistent per-operator buffers, the mg_vcycle.cu kernels at K = 1; upload r, run the V-cycle, download z; the f64 CG unchanged), MultigridParameters::device, Prepared holds the device hierarchy, the CG driver destructures the prepared operator instead of cloning it; EmbeddedPisoSolver::set_poisson_device, overset pass-through, harness knob RTX_FSI2O_MG_DEVICE=1; export_levels factored out; the quarantine's dangling cfg attribute fixed
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 (macos-latest) (push) Waiting to run
CI / Build CPU-Only (Explicit) (push) Failing after 4s
Documentation / Build API Documentation (push) Failing after 5s
Documentation / Build User Guide (push) Successful in 8s
CI / Format Check (push) Failing after 13s
CI / Build (ubuntu-latest) (push) Failing after 2m10s
CI / Clippy Check (push) Failing after 2m28s
Performance Benchmarks / Run Benchmarks (push) Successful in 3m18s

Co-Authored-By: Claude Fable 5.1 <[email protected]>
Claude-Session: https://claude.ai/code/session_01YJPeT6WA2e7YvAnS875AHL
This commit is contained in:
Omar Sobh
2026-09-16 01:16:55 -05:00
co-authored by Claude Fable 5.1
parent 3b7f8fb362
commit fc556f8a88
9 changed files with 410 additions and 20 deletions
@@ -180,6 +180,9 @@ pub struct EmbeddedPisoSolver {
pcg_cache: std::cell::RefCell<super::poisson::PcgCache>,
/// PERF-2 P2: threads for the multigrid's red-black maps (default 1).
poisson_threads: usize,
/// PERF-2 P3-ii: the V-cycle on the CUDA device (needs the `cuda`
/// feature and the red-black smoother; default off).
poisson_device: bool,
moving: bool,
/// Mask hysteresis band in multiples of the min cell size (0 = off).
mask_hysteresis: f64,
@@ -209,6 +212,7 @@ impl EmbeddedPisoSolver {
poisson_profile: std::cell::Cell::new((0, 0, 0, 0)),
pcg_cache: std::cell::RefCell::new(super::poisson::PcgCache::default()),
poisson_threads: 1,
poisson_device: false,
moving: false,
mask_hysteresis: 0.0,
time: 0.0,
@@ -234,6 +238,11 @@ impl EmbeddedPisoSolver {
self.poisson_threads = threads.max(1);
}
/// The V-cycle on the CUDA device (PERF-2 P3-ii; a regime, band-gated).
pub fn set_poisson_device(&mut self, on: bool) {
self.poisson_device = on;
}
/// Mask hysteresis for the moving-body rebuild, as a fraction of the
/// min cell size (default 0, exactly the plain rebuild). With a band,
/// a cell within `band * h_min` of the surface keeps the
@@ -301,6 +301,7 @@ impl EmbeddedPisoSolver {
precision: self.parameters.poisson_precision,
smoother: self.parameters.poisson_smoother,
threads: self.poisson_threads,
device: self.poisson_device,
..MultigridParameters::default()
},
inner_stop,
@@ -1,12 +1,12 @@
// PERF-2 (2026-09-16): the legacy GPU modules have not compiled since the
// initial commit (cudarc API drift); kept out of the `cuda` build until
// someone needs them — the batched V-cycle benchmark uses cudarc directly.
//! Incompressible flow solvers
//!
//! This module implements pressure-velocity coupling algorithms for incompressible flows:
//! - SIMPLE (Semi-Implicit Method for Pressure Linked Equations)
//! - PISO (Pressure-Implicit with Splitting of Operators)
//! - SIMPLER (SIMPLE Revised)
// PERF-2 (2026-09-16): the legacy GPU modules have not compiled since the
// initial commit (cudarc API drift); kept out of the `cuda` build until
// someone needs them — the batched V-cycle benchmark uses cudarc directly.
use crate::{CfdConfig, CfdError, CfdResult};
// use nalgebra::{DMatrix, DVector};
@@ -342,6 +342,11 @@ impl OversetPisoSolver {
self.background.set_poisson_threads(threads);
}
/// PERF-2 P3-ii: the background multigrid's V-cycle on the CUDA device.
pub fn set_poisson_device(&mut self, on: bool) {
self.background.set_poisson_device(on);
}
/// The step timers, when profiling (`RTX_PROFILE`).
pub fn timers(&self) -> Option<&StepTimers> {
self.timers.as_deref()
@@ -282,6 +282,9 @@ pub enum MgSmoother {
RedBlack,
}
#[cfg(feature = "cuda")]
mod device;
/// Multigrid preconditioner parameters.
#[derive(Debug, Clone)]
pub struct MultigridParameters {
@@ -295,6 +298,11 @@ pub struct MultigridParameters {
/// order, so it is bit-identical to the serial one. No effect on the
/// lexicographic smoother (a dependency chain).
pub threads: usize,
/// PERF-2 P3-ii: run the V-cycle on the CUDA device (f32, red-black,
/// the `mg_vcycle.cu` kernels) as the CG's preconditioner; the CG, its
/// residual and its stop stay f64 on the CPU (the M1 contract). Needs
/// the `cuda` feature and `MgSmoother::RedBlack`; ignored otherwise.
pub device: bool,
/// Symmetric GaussSeidel sweeps before AND after the coarse correction
/// (default 2; a value of 0 is treated as 1). One count for both on
/// purpose: unequal pre/post counts make the V-cycle non-symmetric and
@@ -314,6 +322,7 @@ impl Default for MultigridParameters {
precision: MgPrecision::F64,
smoother: MgSmoother::Lexicographic,
threads: 1,
device: false,
smoother_sweeps: 2,
coarsest_cells: 32,
max_iterations: 500,
@@ -446,6 +455,7 @@ impl MgScalar for f32 {
/// and inspection), its coefficients and diagonal in the V-cycle scalar,
/// the list of active cells that carry an equation (row-major), and the
/// parent map into the next coarser level.
#[derive(Clone)]
struct Level<T: MgScalar> {
problem: PoissonProblem,
/// `problem.active && ap > 0`: cells with an equation.
@@ -906,6 +916,7 @@ struct OperatorKey {
coarsest_cells: usize,
smoother: MgSmoother,
threads: usize,
device: bool,
}
impl OperatorKey {
@@ -928,6 +939,7 @@ impl OperatorKey {
coarsest_cells: params.coarsest_cells,
smoother: params.smoother,
threads: params.threads,
device: params.device,
}
}
@@ -938,6 +950,7 @@ impl OperatorKey {
&& self.coarsest_cells == params.coarsest_cells
&& self.smoother == params.smoother
&& self.threads == params.threads
&& self.device == params.device
&& self.active == problem.active
&& self.coefficients.iter().copied().eq(problem
.ae
@@ -962,6 +975,9 @@ struct Prepared<T: MgScalar> {
fine: Level<f64>,
cells: Vec<usize>,
components: Components,
/// The device V-cycle for this operator (PERF-2 P3-ii), when asked for.
#[cfg(feature = "cuda")]
device: Option<device::DeviceVcycle>,
}
impl<T: MgScalar> Prepared<T> {
@@ -970,12 +986,21 @@ impl<T: MgScalar> Prepared<T> {
let fine = Level::<f64>::new(problem.clone());
let cells: Vec<usize> = fine.cells.clone();
let components = Components::find(problem, &cells);
#[cfg(feature = "cuda")]
let device = (params.device && params.smoother == MgSmoother::RedBlack).then(|| {
device::DeviceVcycle::new(
&export_levels(&Hierarchy::<f32>::build(problem, params)),
params.smoother_sweeps.max(1),
)
});
Self {
key: OperatorKey::of(problem, params),
hier,
fine,
cells,
components,
#[cfg(feature = "cuda")]
device,
}
}
}
@@ -1028,8 +1053,15 @@ pub struct LevelExport {
/// The f32 hierarchy of `problem` (red-black colour lists included), level
/// 0 fine, for a device implementation of [`Hierarchy::apply_preconditioner`].
pub fn export_hierarchy(problem: &PoissonProblem, params: &MultigridParameters) -> Vec<LevelExport> {
let hier = Hierarchy::<f32>::build(problem, params);
pub fn export_hierarchy(
problem: &PoissonProblem,
params: &MultigridParameters,
) -> Vec<LevelExport> {
export_levels(&Hierarchy::<f32>::build(problem, params))
}
/// The levels of a built f32 hierarchy (see [`export_hierarchy`]).
fn export_levels(hier: &Hierarchy<f32>) -> Vec<LevelExport> {
let depth = hier.levels.len();
(0..depth)
.map(|l| {
@@ -1204,10 +1236,34 @@ fn run_pcg<T: MgScalar>(
"invalid PoissonProblem: {:?}",
problem.validate()
);
let hier = &mut prep.hier;
let fine = &prep.fine;
let cells: &[usize] = &prep.cells;
let components = &prep.components;
#[cfg(feature = "cuda")]
let Prepared {
hier,
fine,
cells,
components,
device,
..
} = prep;
#[cfg(not(feature = "cuda"))]
let Prepared {
hier,
fine,
cells,
components,
..
} = prep;
let fine: &Level<f64> = fine;
let cells: &[usize] = cells;
let components: &Components = components;
let mut precond = |r: &[f64], z: &mut [f64]| {
#[cfg(feature = "cuda")]
if let Some(d) = device.as_mut() {
d.apply(r, z);
return;
}
hier.apply_preconditioner(r, z);
};
let active_n = cells.len();
if active_n == 0 {
return PoissonSolution {
@@ -1294,7 +1350,7 @@ fn run_pcg<T: MgScalar>(
return finish(p, 0, res);
}
hier.apply_preconditioner(&r, &mut z);
precond(&r, &mut z);
if singular {
project_mean(&mut z);
}
@@ -1333,7 +1389,7 @@ fn run_pcg<T: MgScalar>(
}
last_true = res;
}
hier.apply_preconditioner(&r, &mut z);
precond(&r, &mut z);
if singular {
project_mean(&mut z);
}
@@ -1351,6 +1407,7 @@ fn run_pcg<T: MgScalar>(
/// Connected components of the active cells of a [`PoissonProblem`],
/// connected through faces carrying a non-zero coefficient, and whether
/// each component is singular (carries no Dirichlet contribution).
#[derive(Clone)]
struct Components {
/// Component id per cell (`usize::MAX` for inactive cells).
id: Vec<usize>,
@@ -0,0 +1,312 @@
//! PERF-2 P3-ii (`docs/perf2_campaign.md`): the multigrid V-cycle on the
//! CUDA device as the CG's preconditioner — the `mg_vcycle.cu` kernels of
//! the go/no-go benchmark, K = 1, with persistent device buffers per
//! operator. One CUDA context, stream and module per process (built on
//! first use); the CG, its residual and its stop stay f64 on the CPU.
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/mg_vcycle.cu");
struct Runtime {
_ctx: Arc<CudaContext>,
stream: Arc<CudaStream>,
_module: Arc<CudaModule>,
f_half: CudaFunction,
f_res: CudaFunction,
f_restrict: CudaFunction,
f_prolong: CudaFunction,
f_zero: CudaFunction,
f_coarsest: CudaFunction,
}
static RUNTIME: OnceLock<Runtime> = OnceLock::new();
fn runtime() -> &'static Runtime {
RUNTIME.get_or_init(|| {
let ctx = CudaContext::new(0).expect("CUDA context (device 0)");
let stream = ctx.default_stream();
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: mg_vcycle.cu");
let module = ctx.load_module(ptx).expect("mg_vcycle module");
let f = |name: &str| module.load_function(name).expect(name);
Runtime {
f_half: f("mg_rb_half"),
f_res: f("mg_residual"),
f_restrict: f("mg_restrict"),
f_prolong: f("mg_prolong"),
f_zero: f("mg_zero"),
f_coarsest: f("mg_coarsest"),
_ctx: ctx,
stream,
_module: module,
}
})
}
struct DevLevel {
n: usize,
nx: i32,
n_cells: usize,
n_red: usize,
n_black: usize,
cells: CudaSlice<u32>,
red: CudaSlice<u32>,
black: CudaSlice<u32>,
coarse_of: CudaSlice<u32>,
children_ptr: CudaSlice<u32>,
children_idx: CudaSlice<u32>,
ae: CudaSlice<f32>,
aw: CudaSlice<f32>,
an: CudaSlice<f32>,
as_: CudaSlice<f32>,
ap: CudaSlice<f32>,
b: CudaSlice<f32>,
x: CudaSlice<f32>,
r: CudaSlice<f32>,
}
/// One operator's hierarchy on the device (K = 1).
pub(super) struct DeviceVcycle {
levels: Vec<DevLevel>,
sweeps: usize,
fine_cells: Vec<u32>,
r_f32: Vec<f32>,
z_f32: Vec<f32>,
}
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(super) fn new(levels: &[LevelExport], sweeps: usize) -> Self {
let rt = runtime();
let up_u = |v: &[u32]| -> CudaSlice<u32> {
rt.stream
.memcpy_stod(if v.is_empty() { &[0u32][..] } else { v })
.expect("upload")
};
let up_f = |v: &[f32]| -> CudaSlice<f32> { rt.stream.memcpy_stod(v).expect("upload") };
let dev: Vec<DevLevel> = levels
.iter()
.map(|l| {
let n = l.nx * l.ny;
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),
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_),
ap: up_f(&l.ap),
b: rt.stream.alloc_zeros::<f32>(n).expect("alloc"),
x: rt.stream.alloc_zeros::<f32>(n).expect("alloc"),
r: rt.stream.alloc_zeros::<f32>(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],
}
}
fn half(&mut self, l: usize, colour: u8) {
let rt = runtime();
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_i, n_list_i) = (lv.n as i32, n_list as i32);
unsafe {
rt.stream
.launch_builder(&rt.f_half)
.arg(&n_list_i)
.arg(list)
.arg(&n_i)
.arg(&lv.ae)
.arg(&lv.aw)
.arg(&lv.an)
.arg(&lv.as_)
.arg(&lv.ap)
.arg(&lv.b)
.arg(&mut lv.x)
.arg(&lv.nx)
.launch(cfg(n_list))
.expect("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);
}
}
/// `z = M⁻¹ r` on the active cells (the same V-cycle as
/// `Hierarchy::apply_preconditioner`, in f32 on the device).
pub(super) fn apply(&mut self, r: &[f64], z: &mut [f64]) {
let rt = runtime();
let depth = self.levels.len();
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");
// Down.
for l in 0..depth - 1 {
{
let lv = &mut self.levels[l];
let (n_i, n_cells_i) = (lv.n as i32, lv.n_cells as i32);
unsafe {
rt.stream
.launch_builder(&rt.f_zero)
.arg(&n_cells_i)
.arg(&lv.cells)
.arg(&n_i)
.arg(&mut lv.x)
.launch(cfg(lv.n_cells))
.expect("mg_zero");
}
}
self.smooth(l);
{
let lv = &mut self.levels[l];
let (n_i, n_cells_i) = (lv.n as i32, lv.n_cells as i32);
unsafe {
rt.stream
.launch_builder(&rt.f_res)
.arg(&n_cells_i)
.arg(&lv.cells)
.arg(&n_i)
.arg(&lv.ae)
.arg(&lv.aw)
.arg(&lv.an)
.arg(&lv.as_)
.arg(&lv.ap)
.arg(&lv.b)
.arg(&lv.x)
.arg(&mut lv.r)
.arg(&lv.nx)
.launch(cfg(lv.n_cells))
.expect("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, n_f_i, n_c_i) = (lc.n_cells as i32, lf.n as i32, lc.n as i32);
unsafe {
rt.stream
.launch_builder(&rt.f_restrict)
.arg(&n_c_cells_i)
.arg(&lc.cells)
.arg(&lf.children_ptr)
.arg(&lf.children_idx)
.arg(&n_f_i)
.arg(&n_c_i)
.arg(&lf.r)
.arg(&mut lc.b)
.launch(cfg(lc.n_cells))
.expect("mg_restrict");
}
}
// Coarsest.
{
let lv = &mut self.levels[depth - 1];
let (k_i, n_cells_i, n_i, sw_i) = (1i32, lv.n_cells as i32, lv.n 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(&rt.f_coarsest)
.arg(&k_i)
.arg(&n_cells_i)
.arg(&lv.cells)
.arg(&n_red_i)
.arg(&lv.red)
.arg(&n_black_i)
.arg(&lv.black)
.arg(&n_i)
.arg(&lv.ae)
.arg(&lv.aw)
.arg(&lv.an)
.arg(&lv.as_)
.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("mg_coarsest");
}
}
// Up.
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, n_f_i, n_c_i) = (lf.n_cells as i32, lf.n as i32, lc.n as i32);
unsafe {
rt.stream
.launch_builder(&rt.f_prolong)
.arg(&n_cells_i)
.arg(&lf.cells)
.arg(&lf.coarse_of)
.arg(&n_f_i)
.arg(&n_c_i)
.arg(&mut lf.x)
.arg(&lc.x)
.launch(cfg(lf.n_cells))
.expect("mg_prolong");
}
}
self.smooth(l);
}
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;
}
}
}
@@ -1,6 +1,3 @@
// PERF-2 (2026-09-16): the legacy GPU modules have not compiled since the
// initial commit (cudarc API drift); kept out of the `cuda` build until
// someone needs them — the batched V-cycle benchmark uses cudarc directly.
//! Lattice Boltzmann Method (LBM) solvers
//!
//! This module implements various LBM models for fluid flow simulation:
@@ -8,6 +5,9 @@
//! - D3Q19: 3D model with 19 discrete velocities
//! - Boundary conditions: bounce-back, Zou-He
//! - Collision operators: BGK, MRT
// PERF-2 (2026-09-16): the legacy GPU modules have not compiled since the
// initial commit (cudarc API drift); kept out of the `cuda` build until
// someone needs them — the batched V-cycle benchmark uses cudarc directly.
pub mod boundary;
pub mod d2q9;
@@ -1,6 +1,3 @@
// PERF-2 (2026-09-16): the legacy GPU modules have not compiled since the
// initial commit (cudarc API drift); kept out of the `cuda` build until
// someone needs them — the batched V-cycle benchmark uses cudarc directly.
//! Turbulence modeling for CFD
//!
//! This module provides various turbulence models for simulating turbulent flows:
@@ -10,9 +7,10 @@
//! - **Wall Functions**: Log-law, enhanced wall treatment
//! - **Transition Models**: γ-Reθ, k-kL-ω
// PERF-2 (2026-09-16): the legacy GPU modules have not compiled since the
// initial commit (cudarc API drift); kept out of the `cuda` build until
// someone needs them — the batched V-cycle benchmark uses cudarc directly.
pub mod k_epsilon;
/// GPU-accelerated turbulence models
#[cfg(feature = "cuda")]
// PERF-2 (2026-09-16): the legacy k-epsilon GPU model no longer matches
// `KEpsilonConstants` and has not compiled since the initial commit; kept out
// of the `cuda` build until someone needs it.
@@ -24,7 +22,7 @@ pub mod transition;
pub mod wall_functions;
pub use k_epsilon::{KEpsilonConstants, KEpsilonModel, KEpsilonVariant};
#[cfg(feature = "cuda")]
// #[cfg(feature = "cuda")]
// pub use k_epsilon_gpu::KEpsilonGpuModel;
pub use smagorinsky::{SmagorinskyConstants, SmagorinskyModel};
pub use wall_functions::{EnhancedWallTreatment, LogLawWallFunction, WallFunction};