Omar Sobh
3b7f8fb362
CI / CI Success (push) Blocked by required conditions
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 / Format Check (push) Failing after 5s
CI / Build (ubuntu-latest) (push) Failing after 6s
CI / Clippy Check (push) Failing after 7s
Documentation / Build User Guide (push) Successful in 6s
Performance Benchmarks / Run Benchmarks (push) Failing after 14s
Documentation / Build API Documentation (push) Failing after 30s
CI / Build CPU-Only (Explicit) (push) Failing after 1m31s
402 lines
15 KiB
Rust
402 lines
15 KiB
Rust
//! PERF-2 P3 go/no-go (`docs/perf2_campaign.md`): the batched f32 red-black
|
||
//! V-cycle on the device against the CPU one. Runs only with
|
||
//! `--features cuda` on a CUDA host, and only when asked (`--ignored`):
|
||
//!
|
||
//! `RTX_CUDA_ARCH=sm_120 cargo test --release -p rtx-cfd --features cuda --test gpu_vcycle_bench -- --ignored --nocapture`
|
||
//!
|
||
//! Checks first that the K = 1 device V-cycle reproduces the CPU f32
|
||
//! red-black V-cycle (same algorithm; FMA contraction and summation order
|
||
//! differ, so to 1e-4 of the correction's scale), then times one V-cycle
|
||
//! per march at K = 1, 8, 16 (including the residual upload and the
|
||
//! correction download) against the CPU's serial f64 red-black V-cycle.
|
||
#![cfg(feature = "cuda")]
|
||
|
||
use cudarc::driver::{CudaContext, CudaSlice, LaunchConfig, PushKernelArg};
|
||
use cudarc::nvrtc::{CompileOptions, compile_ptx_with_opts};
|
||
use rtx_cfd::solvers::incompressible::{
|
||
LevelExport, MgSmoother, MultigridParameters, PoissonProblem, export_hierarchy,
|
||
vcycle_f32_reference, vcycle_f32_work,
|
||
};
|
||
use std::sync::Arc;
|
||
use std::time::Instant;
|
||
|
||
const KERNELS: &str = include_str!("../src/kernels/cuda/mg_vcycle.cu");
|
||
|
||
/// The anchor-sized masked operator (ny 62 over the 2.5 × 0.41 channel with
|
||
/// a cylinder-sized hole), the same construction as the Poisson pins.
|
||
fn problem(nx: usize, ny: usize, seed: u64) -> PoissonProblem {
|
||
let mut p = PoissonProblem::new(nx, ny);
|
||
let (dx, dy, dt) = (2.5 / nx as f64, 0.41 / ny as f64, 3.24e-4);
|
||
let (ae, an) = (dt * dy / dx, dt * dx / dy);
|
||
let hole = |i: usize, j: usize| {
|
||
let (x, y) = ((i as f64 + 0.5) * dx, (j as f64 + 0.5) * dy);
|
||
(x - 0.2).powi(2) + (y - 0.2).powi(2) < 0.08 * 0.08
|
||
};
|
||
for j in 0..ny {
|
||
for i in 0..nx {
|
||
let idx = j * nx + i;
|
||
if hole(i, j) {
|
||
p.active[idx] = false;
|
||
continue;
|
||
}
|
||
if i + 1 < nx && !hole(i + 1, j) {
|
||
p.ae[idx] = ae;
|
||
}
|
||
if i > 0 && !hole(i - 1, j) {
|
||
p.aw[idx] = ae;
|
||
}
|
||
if j + 1 < ny && !hole(i, j + 1) {
|
||
p.an[idx] = an;
|
||
}
|
||
if j > 0 && !hole(i, j - 1) {
|
||
p.as_[idx] = an;
|
||
}
|
||
if i + 1 == nx {
|
||
p.extra_diag[idx] = 2.0 * ae;
|
||
}
|
||
}
|
||
}
|
||
let mut state = seed | 1;
|
||
for idx in 0..nx * ny {
|
||
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
|
||
};
|
||
}
|
||
p
|
||
}
|
||
|
||
struct DeviceLevel {
|
||
n: usize,
|
||
nx: i32,
|
||
cells: CudaSlice<u32>,
|
||
red: CudaSlice<u32>,
|
||
black: CudaSlice<u32>,
|
||
coarse_of: CudaSlice<u32>,
|
||
children_ptr: CudaSlice<u32>,
|
||
children_idx: CudaSlice<u32>,
|
||
n_cells: usize,
|
||
n_red: usize,
|
||
n_black: usize,
|
||
ae: CudaSlice<f32>,
|
||
aw: CudaSlice<f32>,
|
||
an: CudaSlice<f32>,
|
||
as_: CudaSlice<f32>,
|
||
ap: CudaSlice<f32>,
|
||
b: CudaSlice<f32>,
|
||
x: CudaSlice<f32>,
|
||
r: CudaSlice<f32>,
|
||
}
|
||
|
||
fn cfg(n_items: usize, k: usize) -> LaunchConfig {
|
||
LaunchConfig {
|
||
grid_dim: ((n_items as u32).div_ceil(256).max(1), k as u32, 1),
|
||
block_dim: (256, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
#[ignore]
|
||
fn batched_device_vcycle_go_no_go() {
|
||
let (nx, ny) = (372usize, 62usize);
|
||
let prob = problem(nx, ny, 3);
|
||
let params = MultigridParameters {
|
||
smoother: MgSmoother::RedBlack,
|
||
..MultigridParameters::default()
|
||
};
|
||
let levels: Vec<LevelExport> = export_hierarchy(&prob, ¶ms);
|
||
let depth = levels.len();
|
||
println!(
|
||
" hierarchy: {depth} levels, cells per level {:?}",
|
||
levels.iter().map(|l| l.cells.len()).collect::<Vec<_>>()
|
||
);
|
||
|
||
// Device.
|
||
let ctx = CudaContext::new(0).expect("CUDA context");
|
||
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.clone().into_boxed_str())),
|
||
..Default::default()
|
||
},
|
||
)
|
||
.expect("nvrtc");
|
||
let module = ctx.load_module(ptx).expect("module");
|
||
let f_half = module.load_function("mg_rb_half").unwrap();
|
||
let f_res = module.load_function("mg_residual").unwrap();
|
||
let f_restrict = module.load_function("mg_restrict").unwrap();
|
||
let f_prolong = module.load_function("mg_prolong").unwrap();
|
||
let f_zero = module.load_function("mg_zero").unwrap();
|
||
let f_coarsest = module.load_function("mg_coarsest").unwrap();
|
||
|
||
// The CPU reference V-cycle on the fine right-hand side.
|
||
let n0 = nx * ny;
|
||
let r_host: Vec<f64> = prob.rhs.clone();
|
||
let mut z_ref = vec![0.0; n0];
|
||
vcycle_f32_reference(&prob, ¶ms, &r_host, &mut z_ref);
|
||
let z_scale = z_ref.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
|
||
|
||
// The CPU f64 serial red-black V-cycle's wall time (the comparison point).
|
||
let t0 = Instant::now();
|
||
let reps = 20;
|
||
for _ in 0..reps {
|
||
let mut z = vec![0.0; n0];
|
||
vcycle_f32_reference(&prob, ¶ms, &r_host, &mut z);
|
||
}
|
||
let cpu_ms = t0.elapsed().as_secs_f64() * 1e3 / reps as f64;
|
||
println!(" CPU f32 red-black V-cycle (incl. hierarchy build): {cpu_ms:.3} ms");
|
||
|
||
let sweeps = 2usize;
|
||
for &k in &[1usize, 8, 16] {
|
||
// Upload: coefficients replicated K times, index lists shared.
|
||
let mut dev: Vec<DeviceLevel> = Vec::with_capacity(depth);
|
||
for l in &levels {
|
||
let n = l.nx * l.ny;
|
||
let rep = |v: &[f32]| -> Vec<f32> {
|
||
let mut out = Vec::with_capacity(v.len() * k);
|
||
for _ in 0..k {
|
||
out.extend_from_slice(v);
|
||
}
|
||
out
|
||
};
|
||
let up_u = |v: &[u32]| stream.memcpy_stod(v).unwrap();
|
||
let up_f = |v: &[f32]| stream.memcpy_stod(v).unwrap();
|
||
dev.push(DeviceLevel {
|
||
n,
|
||
nx: l.nx as i32,
|
||
cells: up_u(&l.cells),
|
||
red: up_u(&l.red),
|
||
black: up_u(&l.black),
|
||
coarse_of: up_u(if l.coarse_of.is_empty() { &[0u32][..] } else { &l.coarse_of }),
|
||
children_ptr: up_u(if l.children_ptr.is_empty() { &[0u32][..] } else { &l.children_ptr }),
|
||
children_idx: up_u(if l.children_idx.is_empty() { &[0u32][..] } else { &l.children_idx }),
|
||
n_cells: l.cells.len(),
|
||
n_red: l.red.len(),
|
||
n_black: l.black.len(),
|
||
ae: up_f(&rep(&l.ae)),
|
||
aw: up_f(&rep(&l.aw)),
|
||
an: up_f(&rep(&l.an)),
|
||
as_: up_f(&rep(&l.as_)),
|
||
ap: up_f(&rep(&l.ap)),
|
||
b: stream.alloc_zeros::<f32>(n * k).unwrap(),
|
||
x: stream.alloc_zeros::<f32>(n * k).unwrap(),
|
||
r: stream.alloc_zeros::<f32>(n * k).unwrap(),
|
||
});
|
||
}
|
||
stream.synchronize().unwrap();
|
||
|
||
let r_f32: Vec<f32> = {
|
||
let one: Vec<f32> = r_host.iter().map(|&v| v as f32).collect();
|
||
let mut out = Vec::with_capacity(n0 * k);
|
||
for _ in 0..k {
|
||
out.extend_from_slice(&one);
|
||
}
|
||
out
|
||
};
|
||
let mut z_out = vec![0.0f32; n0 * k];
|
||
|
||
let vcycle = |dev: &mut Vec<DeviceLevel>, r_f32: &[f32], z_out: &mut [f32]| {
|
||
// Upload the residual into level 0's b.
|
||
stream.memcpy_htod(r_f32, &mut dev[0].b).unwrap();
|
||
let half = |lv: &mut DeviceLevel, colour: u8| {
|
||
let (list, n_list) = if colour == 0 { (&lv.red, lv.n_red) } else { (&lv.black, lv.n_black) };
|
||
let n_i = lv.n as i32;
|
||
let n_list_i = n_list as i32;
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&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, k))
|
||
.unwrap();
|
||
}
|
||
};
|
||
let smooth = |lv: &mut DeviceLevel| {
|
||
for _ in 0..sweeps {
|
||
half(lv, 0);
|
||
half(lv, 1);
|
||
half(lv, 1);
|
||
half(lv, 0);
|
||
}
|
||
};
|
||
// Down.
|
||
for l in 0..depth - 1 {
|
||
{
|
||
let lv = &mut dev[l];
|
||
let (n_i, n_cells_i) = (lv.n as i32, lv.n_cells as i32);
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&f_zero)
|
||
.arg(&n_cells_i)
|
||
.arg(&lv.cells)
|
||
.arg(&n_i)
|
||
.arg(&mut lv.x)
|
||
.launch(cfg(lv.n_cells, k))
|
||
.unwrap();
|
||
}
|
||
smooth(lv);
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&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, k))
|
||
.unwrap();
|
||
}
|
||
}
|
||
let (fine, coarse) = dev.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 {
|
||
stream
|
||
.launch_builder(&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, k))
|
||
.unwrap();
|
||
}
|
||
}
|
||
// Coarsest.
|
||
{
|
||
let lv = &mut dev[depth - 1];
|
||
let (k_i, n_cells_i, n_i, sw_i) = (k as i32, 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 {
|
||
stream
|
||
.launch_builder(&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: ((k as u32).div_ceil(32), 1, 1),
|
||
block_dim: (32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
})
|
||
.unwrap();
|
||
}
|
||
}
|
||
// Up.
|
||
for l in (0..depth - 1).rev() {
|
||
let (fine, coarse) = dev.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 {
|
||
stream
|
||
.launch_builder(&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, k))
|
||
.unwrap();
|
||
}
|
||
smooth(lf);
|
||
}
|
||
stream.memcpy_dtoh(&dev[0].x, z_out).unwrap();
|
||
stream.synchronize().unwrap();
|
||
};
|
||
|
||
// Correctness at K = 1 (and every batch member at K > 1).
|
||
vcycle(&mut dev, &r_f32, &mut z_out);
|
||
if k == 1 {
|
||
// Level by level against the CPU's work vectors (march 0).
|
||
let work = vcycle_f32_work(&prob, ¶ms, &r_host);
|
||
for (l, (b_ref, x_ref, r_ref)) in work.iter().enumerate() {
|
||
let n = dev[l].n;
|
||
let mut b = vec![0.0f32; n * k];
|
||
let mut x = vec![0.0f32; n * k];
|
||
let mut r = vec![0.0f32; n * k];
|
||
stream.memcpy_dtoh(&dev[l].b, &mut b).unwrap();
|
||
stream.memcpy_dtoh(&dev[l].x, &mut x).unwrap();
|
||
stream.memcpy_dtoh(&dev[l].r, &mut r).unwrap();
|
||
stream.synchronize().unwrap();
|
||
let cmp = |a: &[f32], c: &[f32]| {
|
||
let (mut worst, mut scale) = (0.0f32, 0.0f32);
|
||
for &idx in &levels[l].cells {
|
||
let i = idx as usize;
|
||
worst = worst.max((a[i] - c[i]).abs());
|
||
scale = scale.max(c[i].abs());
|
||
}
|
||
(worst, scale)
|
||
};
|
||
let (db, sb) = cmp(&b[..n], b_ref);
|
||
let (dx, sx) = cmp(&x[..n], x_ref);
|
||
let (dr, sr) = cmp(&r[..n], r_ref);
|
||
println!(" level {l} ({} cells): |Δb| {db:.3e} / {sb:.3e}, |Δx| {dx:.3e} / {sx:.3e}, |Δr| {dr:.3e} / {sr:.3e}", levels[l].cells.len());
|
||
}
|
||
}
|
||
let mut worst = 0.0_f64;
|
||
for m in 0..k {
|
||
for idx in 0..n0 {
|
||
let d = (z_out[m * n0 + idx] as f64 - z_ref[idx]).abs();
|
||
worst = worst.max(d);
|
||
}
|
||
}
|
||
println!(" K = {k}: device V-cycle vs CPU f32 reference: max |Δz| {worst:.3e} on a scale of {z_scale:.3e}");
|
||
assert!(worst < 1e-4 * z_scale, "device V-cycle disagrees with the CPU one");
|
||
|
||
// Timing: 50 V-cycles, per march.
|
||
let reps = 50;
|
||
vcycle(&mut dev, &r_f32, &mut z_out);
|
||
let t0 = Instant::now();
|
||
for _ in 0..reps {
|
||
vcycle(&mut dev, &r_f32, &mut z_out);
|
||
}
|
||
let ms = t0.elapsed().as_secs_f64() * 1e3 / reps as f64;
|
||
println!(
|
||
" K = {k}: {ms:.3} ms per batched V-cycle = {:.3} ms per march (CPU serial {cpu_ms:.3} ms; ratio {:.2}×)",
|
||
ms / k as f64,
|
||
cpu_ms / (ms / k as f64)
|
||
);
|
||
}
|
||
}
|