Omar SobhandClaude Fable 5.1
95ffde9591
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 / 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 / Format Check (push) Failing after 5s
Documentation / Build API Documentation (push) Failing after 5s
CI / Build (ubuntu-latest) (push) Failing after 14s
CI / Clippy Check (push) Failing after 3m1s
Documentation / Build User Guide (push) Successful in 6s
Performance Benchmarks / Run Benchmarks (push) Failing after 11s
CI / Build CPU-Only (Explicit) (push) Failing after 2m33s
Co-Authored-By: Claude Fable 5.1 <[email protected]>
879 lines
29 KiB
Rust
879 lines
29 KiB
Rust
//! 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, DevicePtr, DeviceRepr,
|
||
LaunchConfig, PushKernelArg, ValidAsZeroBits,
|
||
};
|
||
use cudarc::nvrtc::{CompileOptions, compile_ptx_with_opts};
|
||
use std::cell::Cell;
|
||
use std::collections::BTreeMap;
|
||
use std::sync::{Arc, Condvar, Mutex, 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,
|
||
// The lane-table kernels (P3-iii).
|
||
l_half: CudaFunction,
|
||
l_res: CudaFunction,
|
||
l_restrict: CudaFunction,
|
||
l_prolong: CudaFunction,
|
||
l_zero: CudaFunction,
|
||
l_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"),
|
||
l_half: f("ml_rb_half"),
|
||
l_res: f("ml_residual"),
|
||
l_restrict: f("ml_restrict"),
|
||
l_prolong: f("ml_prolong"),
|
||
l_zero: f("ml_zero"),
|
||
l_coarsest: f("ml_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;
|
||
}
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// PERF-2 P3-iii: the `MarchPlane` — K lanes' V-cycles in ONE batched launch
|
||
// set, and the rendezvous that gathers the lanes' preconditioner calls.
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// One lane's level as the `ml_*` kernels read it: 14 device pointers then
|
||
/// 6 ints (`n, nx, n_cells, n_red, n_black, pad`), mirroring `struct
|
||
/// LaneLevel` in `mg_vcycle.cu`.
|
||
#[repr(C)]
|
||
#[derive(Clone, Copy, Default)]
|
||
struct LaneLevel {
|
||
ptrs: [u64; 14],
|
||
ints: [i32; 6],
|
||
}
|
||
unsafe impl DeviceRepr for LaneLevel {}
|
||
unsafe impl ValidAsZeroBits for LaneLevel {}
|
||
|
||
impl DeviceVcycle {
|
||
pub(super) fn depth(&self) -> usize {
|
||
self.levels.len()
|
||
}
|
||
|
||
fn lane_level(&self, l: usize, stream: &CudaStream) -> LaneLevel {
|
||
let lv = &self.levels[l];
|
||
let pf = |s: &CudaSlice<f32>| s.device_ptr(stream).0;
|
||
let pu = |s: &CudaSlice<u32>| s.device_ptr(stream).0;
|
||
LaneLevel {
|
||
ptrs: [
|
||
pf(&lv.ae),
|
||
pf(&lv.aw),
|
||
pf(&lv.an),
|
||
pf(&lv.as_),
|
||
pf(&lv.ap),
|
||
pf(&lv.b),
|
||
pf(&lv.x),
|
||
pf(&lv.r),
|
||
pu(&lv.cells),
|
||
pu(&lv.red),
|
||
pu(&lv.black),
|
||
pu(&lv.coarse_of),
|
||
pu(&lv.children_ptr),
|
||
pu(&lv.children_idx),
|
||
],
|
||
ints: [
|
||
lv.n as i32,
|
||
lv.nx,
|
||
lv.n_cells as i32,
|
||
lv.n_red as i32,
|
||
lv.n_black as i32,
|
||
0,
|
||
],
|
||
}
|
||
}
|
||
}
|
||
|
||
/// One lane's preconditioner call: its operator on the device, the residual
|
||
/// in, the correction out.
|
||
pub(super) struct LaneJob<'a> {
|
||
pub dev: &'a mut DeviceVcycle,
|
||
pub r: &'a [f64],
|
||
pub z: &'a mut [f64],
|
||
}
|
||
|
||
fn cfg_k(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,
|
||
}
|
||
}
|
||
|
||
/// `z_k = M_k⁻¹ r_k` for every job, the lanes of equal hierarchy depth in
|
||
/// one batched launch set (grid.y = the lane). Each lane's result is the
|
||
/// K = 1 [`DeviceVcycle::apply`]'s, bit for bit (the same per-cell
|
||
/// arithmetic in the same order; gate G-P0).
|
||
pub(super) fn apply_batch(jobs: &mut [LaneJob<'_>]) {
|
||
let rt = runtime();
|
||
let stream = &rt.stream;
|
||
let mut groups: BTreeMap<usize, Vec<usize>> = BTreeMap::new();
|
||
for (i, j) in jobs.iter().enumerate() {
|
||
groups.entry(j.dev.depth()).or_default().push(i);
|
||
}
|
||
for (depth, members) in groups {
|
||
let k = members.len();
|
||
// Residuals up.
|
||
for &i in &members {
|
||
let job = &mut jobs[i];
|
||
for (dst, &src) in job.dev.r_f32.iter_mut().zip(job.r) {
|
||
*dst = src as f32;
|
||
}
|
||
stream
|
||
.memcpy_htod(&job.dev.r_f32, &mut job.dev.levels[0].b)
|
||
.expect("upload r");
|
||
}
|
||
// The lane tables, one per level.
|
||
let tables: Vec<CudaSlice<LaneLevel>> = (0..depth)
|
||
.map(|l| {
|
||
let rows: Vec<LaneLevel> = members
|
||
.iter()
|
||
.map(|&i| jobs[i].dev.lane_level(l, stream))
|
||
.collect();
|
||
stream.memcpy_stod(&rows).expect("lane table")
|
||
})
|
||
.collect();
|
||
let max_of = |f: &dyn Fn(&DevLevel) -> usize, l: usize| -> usize {
|
||
members
|
||
.iter()
|
||
.map(|&i| f(&jobs[i].dev.levels[l]))
|
||
.max()
|
||
.unwrap_or(0)
|
||
};
|
||
let sweeps = jobs[members[0]].dev.sweeps;
|
||
let smooth = |l: usize| {
|
||
let n_red = max_of(&|lv| lv.n_red, l);
|
||
let n_black = max_of(&|lv| lv.n_black, l);
|
||
let half = |colour: i32| {
|
||
let n = if colour == 0 { n_red } else { n_black };
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&rt.l_half)
|
||
.arg(&tables[l])
|
||
.arg(&colour)
|
||
.launch(cfg_k(n, k))
|
||
.expect("ml_rb_half");
|
||
}
|
||
};
|
||
for _ in 0..sweeps {
|
||
half(0);
|
||
half(1);
|
||
half(1);
|
||
half(0);
|
||
}
|
||
};
|
||
// Down.
|
||
for l in 0..depth - 1 {
|
||
let n_cells = max_of(&|lv| lv.n_cells, l);
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&rt.l_zero)
|
||
.arg(&tables[l])
|
||
.launch(cfg_k(n_cells, k))
|
||
.expect("ml_zero");
|
||
}
|
||
smooth(l);
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&rt.l_res)
|
||
.arg(&tables[l])
|
||
.launch(cfg_k(n_cells, k))
|
||
.expect("ml_residual");
|
||
}
|
||
let n_coarse = max_of(&|lv| lv.n_cells, l + 1);
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&rt.l_restrict)
|
||
.arg(&tables[l])
|
||
.arg(&tables[l + 1])
|
||
.launch(cfg_k(n_coarse, k))
|
||
.expect("ml_restrict");
|
||
}
|
||
}
|
||
// Coarsest.
|
||
{
|
||
let (k_i, sw_i) = (k as i32, 50i32);
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&rt.l_coarsest)
|
||
.arg(&tables[depth - 1])
|
||
.arg(&k_i)
|
||
.arg(&sw_i)
|
||
.launch(LaunchConfig {
|
||
grid_dim: ((k as u32).div_ceil(32), 1, 1),
|
||
block_dim: (32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
})
|
||
.expect("ml_coarsest");
|
||
}
|
||
}
|
||
// Up.
|
||
for l in (0..depth - 1).rev() {
|
||
let n_cells = max_of(&|lv| lv.n_cells, l);
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&rt.l_prolong)
|
||
.arg(&tables[l])
|
||
.arg(&tables[l + 1])
|
||
.launch(cfg_k(n_cells, k))
|
||
.expect("ml_prolong");
|
||
}
|
||
smooth(l);
|
||
}
|
||
// Corrections down.
|
||
for &i in &members {
|
||
let job = &mut jobs[i];
|
||
stream
|
||
.memcpy_dtoh(&job.dev.levels[0].x, &mut job.dev.z_f32)
|
||
.expect("download z");
|
||
}
|
||
stream.synchronize().expect("sync");
|
||
for &i in &members {
|
||
let job = &mut jobs[i];
|
||
for &idx in &job.dev.fine_cells {
|
||
job.z[idx as usize] = job.dev.z_f32[idx as usize] as f64;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// --- The rendezvous ---------------------------------------------------------
|
||
//
|
||
// K lane threads each run an unchanged march. A lane inside a CG (between
|
||
// `cg_enter` and the guard's drop) that calls the preconditioner submits its
|
||
// job and blocks; when every lane currently inside a CG has submitted, the
|
||
// last arrival serves the whole batch and wakes the rest. A lane LEAVING its
|
||
// CG can complete a batch too (the others were waiting for it). Lanes
|
||
// outside a CG are not counted, so they never hold a batch up.
|
||
|
||
thread_local! {
|
||
static LANE: Cell<Option<usize>> = const { Cell::new(None) };
|
||
}
|
||
|
||
/// Mark this thread as plane lane `lane` (`None` = a solo march: the K = 1
|
||
/// path, unchanged).
|
||
pub fn set_lane(lane: Option<usize>) {
|
||
LANE.with(|l| l.set(lane));
|
||
}
|
||
|
||
pub fn lane() -> Option<usize> {
|
||
LANE.with(|l| l.get())
|
||
}
|
||
|
||
struct PendingJob {
|
||
dev: *mut DeviceVcycle,
|
||
r: *const f64,
|
||
r_len: usize,
|
||
z: *mut f64,
|
||
z_len: usize,
|
||
}
|
||
// The submitting thread blocks until its job is served, so the pointers
|
||
// outlive the batch.
|
||
unsafe impl Send for PendingJob {}
|
||
|
||
#[derive(Default)]
|
||
struct PlaneState {
|
||
in_cg: usize,
|
||
pending: Vec<PendingJob>,
|
||
generation: u64,
|
||
/// Batches served and lanes served (the plane's own bookkeeping).
|
||
batches: u64,
|
||
lanes_served: u64,
|
||
}
|
||
|
||
static PLANE: Mutex<PlaneState> = Mutex::new(PlaneState {
|
||
in_cg: 0,
|
||
pending: Vec::new(),
|
||
generation: 0,
|
||
batches: 0,
|
||
lanes_served: 0,
|
||
});
|
||
static PLANE_CV: Condvar = Condvar::new();
|
||
|
||
fn plane_lock() -> std::sync::MutexGuard<'static, PlaneState> {
|
||
PLANE.lock().unwrap_or_else(|e| e.into_inner())
|
||
}
|
||
|
||
fn serve(st: &mut PlaneState) {
|
||
let pending = std::mem::take(&mut st.pending);
|
||
let mut jobs: Vec<LaneJob<'_>> = pending
|
||
.iter()
|
||
.map(|p| unsafe {
|
||
LaneJob {
|
||
dev: &mut *p.dev,
|
||
r: std::slice::from_raw_parts(p.r, p.r_len),
|
||
z: std::slice::from_raw_parts_mut(p.z, p.z_len),
|
||
}
|
||
})
|
||
.collect();
|
||
apply_batch(&mut jobs);
|
||
st.batches += 1;
|
||
st.lanes_served += pending.len() as u64;
|
||
st.generation += 1;
|
||
PLANE_CV.notify_all();
|
||
}
|
||
|
||
/// Held by a lane for the duration of one CG solve.
|
||
pub(super) struct CgGuard(());
|
||
|
||
impl Drop for CgGuard {
|
||
fn drop(&mut self) {
|
||
let mut st = plane_lock();
|
||
st.in_cg -= 1;
|
||
if st.in_cg > 0 && st.pending.len() >= st.in_cg {
|
||
serve(&mut st);
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Enter a CG solve on this lane (`None` when this thread is not a lane).
|
||
pub(super) fn cg_enter() -> Option<CgGuard> {
|
||
lane()?;
|
||
plane_lock().in_cg += 1;
|
||
Some(CgGuard(()))
|
||
}
|
||
|
||
/// This lane's preconditioner call inside the plane: submit, and either
|
||
/// serve the batch (last arrival) or wait for it.
|
||
pub(super) fn apply_lane(dev: &mut DeviceVcycle, r: &[f64], z: &mut [f64]) {
|
||
let mut st = plane_lock();
|
||
let my_generation = st.generation;
|
||
st.pending.push(PendingJob {
|
||
dev: std::ptr::from_mut(dev),
|
||
r: r.as_ptr(),
|
||
r_len: r.len(),
|
||
z: z.as_mut_ptr(),
|
||
z_len: z.len(),
|
||
});
|
||
if st.pending.len() >= st.in_cg {
|
||
serve(&mut st);
|
||
return;
|
||
}
|
||
while st.generation == my_generation {
|
||
st = PLANE_CV.wait(st).unwrap_or_else(|e| e.into_inner());
|
||
}
|
||
}
|
||
|
||
/// `(batches served, lane calls served)` so far in this process.
|
||
pub fn plane_counters() -> (u64, u64) {
|
||
let st = plane_lock();
|
||
(st.batches, st.lanes_served)
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
//! G-P0 (`docs/perf2_campaign.md`, P3-iii): K distinct masked operators
|
||
//! through the batched path equal each one's K = 1 device apply bit for
|
||
//! bit, and the rendezvous serves uneven lanes exactly. GPU only:
|
||
//! `RTX_CUDA_ARCH=sm_120 cargo test --release -p rtx-cfd --features cuda --lib device::tests`.
|
||
use super::*;
|
||
use crate::solvers::incompressible::poisson::{
|
||
MgSmoother, MultigridParameters, PoissonProblem, export_hierarchy,
|
||
};
|
||
|
||
/// A masked channel operator with a hole at `(cx, cy)` and a seeded
|
||
/// right-hand side (the go/no-go benchmark's construction).
|
||
fn problem(nx: usize, ny: usize, cx: f64, cy: f64, 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 - cx).powi(2) + (y - cy).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
|
||
}
|
||
|
||
fn params() -> MultigridParameters {
|
||
MultigridParameters {
|
||
smoother: MgSmoother::RedBlack,
|
||
..MultigridParameters::default()
|
||
}
|
||
}
|
||
|
||
fn device_of(p: &PoissonProblem) -> DeviceVcycle {
|
||
DeviceVcycle::new(
|
||
&export_hierarchy(p, ¶ms()),
|
||
params().smoother_sweeps.max(1),
|
||
)
|
||
}
|
||
|
||
fn bits(v: &[f64]) -> Vec<u64> {
|
||
v.iter().map(|x| x.to_bits()).collect()
|
||
}
|
||
|
||
/// K operators with different holes: solo applies, then one batch.
|
||
#[test]
|
||
fn batched_lanes_equal_the_solo_apply_bit_for_bit() {
|
||
let (nx, ny) = (186usize, 31usize);
|
||
let probs: Vec<PoissonProblem> = [(0.2, 0.2, 3u64), (0.6, 0.15, 7), (1.1, 0.25, 11)]
|
||
.iter()
|
||
.map(|&(cx, cy, s)| problem(nx, ny, cx, cy, s))
|
||
.collect();
|
||
let n = nx * ny;
|
||
let mut solo: Vec<Vec<f64>> = Vec::new();
|
||
for p in &probs {
|
||
let mut d = device_of(p);
|
||
let mut z = vec![0.0; n];
|
||
d.apply(&p.rhs, &mut z);
|
||
solo.push(z);
|
||
}
|
||
assert!(
|
||
solo[0].iter().any(|v| *v != 0.0),
|
||
"the solo apply produced zeros"
|
||
);
|
||
// K = 1 through the batched path (the lane kernels vs the K = 1 kernels).
|
||
{
|
||
let mut d = device_of(&probs[0]);
|
||
let mut z = vec![0.0; n];
|
||
apply_batch(&mut [LaneJob {
|
||
dev: &mut d,
|
||
r: &probs[0].rhs,
|
||
z: &mut z,
|
||
}]);
|
||
assert_eq!(
|
||
bits(&z),
|
||
bits(&solo[0]),
|
||
"K = 1 through the lane kernels differs"
|
||
);
|
||
}
|
||
// K = 3.
|
||
let mut devs: Vec<DeviceVcycle> = probs.iter().map(device_of).collect();
|
||
let mut zs: Vec<Vec<f64>> = vec![vec![0.0; n]; 3];
|
||
{
|
||
let mut jobs: Vec<LaneJob<'_>> = devs
|
||
.iter_mut()
|
||
.zip(probs.iter())
|
||
.zip(zs.iter_mut())
|
||
.map(|((dev, p), z)| LaneJob { dev, r: &p.rhs, z })
|
||
.collect();
|
||
apply_batch(&mut jobs);
|
||
}
|
||
for k in 0..3 {
|
||
assert_eq!(
|
||
bits(&zs[k]),
|
||
bits(&solo[k]),
|
||
"lane {k} differs from its solo apply"
|
||
);
|
||
}
|
||
// A second batch on the same devices (persistent buffers) with the
|
||
// lanes in another order and one lane left out.
|
||
let mut z1 = vec![0.0; n];
|
||
let mut z2 = vec![0.0; n];
|
||
{
|
||
let (a, b) = devs.split_at_mut(2);
|
||
let mut jobs = [
|
||
LaneJob {
|
||
dev: &mut b[0],
|
||
r: &probs[2].rhs,
|
||
z: &mut z2,
|
||
},
|
||
LaneJob {
|
||
dev: &mut a[1],
|
||
r: &probs[1].rhs,
|
||
z: &mut z1,
|
||
},
|
||
];
|
||
apply_batch(&mut jobs);
|
||
}
|
||
assert_eq!(bits(&z2), bits(&solo[2]));
|
||
assert_eq!(bits(&z1), bits(&solo[1]));
|
||
}
|
||
|
||
/// Three lane threads with 2 / 4 / 5 preconditioner calls each inside
|
||
/// one CG guard: every call's correction equals the solo apply of the
|
||
/// same residual, the lanes leave at different times, and no thread
|
||
/// hangs.
|
||
#[test]
|
||
fn the_rendezvous_serves_uneven_lanes_exactly() {
|
||
let (nx, ny) = (124usize, 21usize);
|
||
let n = nx * ny;
|
||
let probs: Vec<PoissonProblem> = [(0.2, 0.2, 5u64), (0.7, 0.2, 9), (1.3, 0.22, 13)]
|
||
.iter()
|
||
.map(|&(cx, cy, s)| problem(nx, ny, cx, cy, s))
|
||
.collect();
|
||
let calls = [2usize, 4, 5];
|
||
// Expected: the solo apply of rhs × (1 + c) for call c.
|
||
let mut expected: Vec<Vec<Vec<u64>>> = Vec::new();
|
||
for (k, p) in probs.iter().enumerate() {
|
||
let mut d = device_of(p);
|
||
let mut per_call = Vec::new();
|
||
for c in 0..calls[k] {
|
||
let r: Vec<f64> = p.rhs.iter().map(|v| v * (1.0 + c as f64)).collect();
|
||
let mut z = vec![0.0; n];
|
||
d.apply(&r, &mut z);
|
||
per_call.push(bits(&z));
|
||
}
|
||
expected.push(per_call);
|
||
}
|
||
let (b0, l0) = plane_counters();
|
||
let handles: Vec<_> = probs
|
||
.into_iter()
|
||
.enumerate()
|
||
.map(|(k, p)| {
|
||
let calls_k = calls[k];
|
||
std::thread::spawn(move || {
|
||
set_lane(Some(k));
|
||
let mut d = device_of(&p);
|
||
let guard = cg_enter().expect("lane set");
|
||
let mut out = Vec::new();
|
||
for c in 0..calls_k {
|
||
let r: Vec<f64> = p.rhs.iter().map(|v| v * (1.0 + c as f64)).collect();
|
||
let mut z = vec![0.0; n];
|
||
apply_lane(&mut d, &r, &mut z);
|
||
out.push(bits(&z));
|
||
}
|
||
drop(guard);
|
||
out
|
||
})
|
||
})
|
||
.collect();
|
||
let results: Vec<Vec<Vec<u64>>> = handles
|
||
.into_iter()
|
||
.map(|h| h.join().expect("lane thread"))
|
||
.collect();
|
||
for k in 0..3 {
|
||
assert_eq!(results[k].len(), calls[k]);
|
||
for c in 0..calls[k] {
|
||
assert_eq!(results[k][c], expected[k][c], "lane {k} call {c} differs");
|
||
}
|
||
}
|
||
let (b1, l1) = plane_counters();
|
||
assert_eq!(l1 - l0, 11, "every lane call served once");
|
||
assert!(b1 - b0 >= 5 && b1 - b0 <= 11, "batches {}", b1 - b0);
|
||
}
|
||
}
|