diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs index 02eb116..0074de5 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs @@ -67,8 +67,8 @@ pub use piso_gpu::PisoGpuSolver; pub use poisson::{ LevelExport, MgPrecision, MgSmoother, MultigridParameters, PcgCache, PoissonProblem, PoissonSolution, PoissonSolverKind, configure_threads, export_hierarchy, plane_counters, - set_plane_lane, solve_multigrid_pcg, solve_multigrid_pcg_cached, vcycle_f32_reference, - vcycle_f32_work, + plane_streams, set_plane_lane, set_plane_streams, solve_multigrid_pcg, + solve_multigrid_pcg_cached, vcycle_f32_reference, vcycle_f32_work, }; pub use polygon_sdf::PolygonSdf; pub use simple::{ConvectionScheme, SimpleParameters, SimpleResult, SimpleSolver}; diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs index 73040b7..94722d1 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson.rs @@ -289,13 +289,19 @@ mod device; /// solo march). On a lane, every device V-cycle of a CG goes through the /// plane's rendezvous. Without the `cuda` feature this is a no-op. #[cfg(feature = "cuda")] -pub use device::{plane_counters, set_lane as set_plane_lane}; +pub use device::{plane_counters, plane_streams, set_lane as set_plane_lane, set_plane_streams}; #[cfg(not(feature = "cuda"))] pub fn set_plane_lane(_lane: Option) {} #[cfg(not(feature = "cuda"))] pub fn plane_counters() -> (u64, u64) { (0, 0) } +#[cfg(not(feature = "cuda"))] +pub fn set_plane_streams(_on: bool) {} +#[cfg(not(feature = "cuda"))] +pub fn plane_streams() -> bool { + false +} /// Multigrid preconditioner parameters. #[derive(Debug, Clone)] diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson/device.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson/device.rs index b6e5421..85f995f 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson/device.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/poisson/device.rs @@ -12,6 +12,7 @@ use cudarc::driver::{ use cudarc::nvrtc::{CompileOptions, compile_ptx_with_opts}; use std::cell::Cell; use std::collections::BTreeMap; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Condvar, Mutex, OnceLock}; const KERNELS: &str = include_str!("../../../kernels/cuda/mg_vcycle.cu"); @@ -96,6 +97,10 @@ struct DevLevel { /// One operator's hierarchy on the device (K = 1). pub(super) struct DeviceVcycle { + /// The stream this operator's buffers and launches live on: the + /// process's default stream for a solo march or a batch-mode lane, the + /// lane thread's own stream in stream mode (P3-iii K-stream). + stream: Arc, levels: Vec, sweeps: usize, fine_cells: Vec, @@ -113,13 +118,13 @@ fn cfg(n_items: usize) -> LaunchConfig { impl DeviceVcycle { pub(super) fn new(levels: &[LevelExport], sweeps: usize) -> Self { - let rt = runtime(); + let stream = stream_for_this_thread(); let up_u = |v: &[u32]| -> CudaSlice { - rt.stream + stream .memcpy_stod(if v.is_empty() { &[0u32][..] } else { v }) .expect("upload") }; - let up_f = |v: &[f32]| -> CudaSlice { rt.stream.memcpy_stod(v).expect("upload") }; + let up_f = |v: &[f32]| -> CudaSlice { stream.memcpy_stod(v).expect("upload") }; let dev: Vec = levels .iter() .map(|l| { @@ -141,14 +146,15 @@ impl DeviceVcycle { an: up_f(&l.an), as_: up_f(&l.as_), ap: up_f(&l.ap), - b: rt.stream.alloc_zeros::(n).expect("alloc"), - x: rt.stream.alloc_zeros::(n).expect("alloc"), - r: rt.stream.alloc_zeros::(n).expect("alloc"), + b: stream.alloc_zeros::(n).expect("alloc"), + x: stream.alloc_zeros::(n).expect("alloc"), + r: stream.alloc_zeros::(n).expect("alloc"), } }) .collect(); let n0 = dev[0].n; Self { + stream, levels: dev, sweeps, fine_cells: levels[0].cells.clone(), @@ -159,6 +165,7 @@ impl DeviceVcycle { fn half(&mut self, l: usize, colour: u8) { let rt = runtime(); + let stream = Arc::clone(&self.stream); let lv = &mut self.levels[l]; let (list, n_list) = if colour == 0 { (&lv.red, lv.n_red) @@ -167,7 +174,7 @@ impl DeviceVcycle { }; let (n_i, n_list_i) = (lv.n as i32, n_list as i32); unsafe { - rt.stream + stream .launch_builder(&rt.f_half) .arg(&n_list_i) .arg(list) @@ -198,11 +205,12 @@ impl DeviceVcycle { /// `Hierarchy::apply_preconditioner`, in f32 on the device). pub(super) fn apply(&mut self, r: &[f64], z: &mut [f64]) { let rt = runtime(); + let stream = Arc::clone(&self.stream); let depth = self.levels.len(); for (dst, &src) in self.r_f32.iter_mut().zip(r) { *dst = src as f32; } - rt.stream + stream .memcpy_htod(&self.r_f32, &mut self.levels[0].b) .expect("upload r"); // Down. @@ -211,7 +219,7 @@ impl DeviceVcycle { let lv = &mut self.levels[l]; let (n_i, n_cells_i) = (lv.n as i32, lv.n_cells as i32); unsafe { - rt.stream + stream .launch_builder(&rt.f_zero) .arg(&n_cells_i) .arg(&lv.cells) @@ -226,7 +234,7 @@ impl DeviceVcycle { let lv = &mut self.levels[l]; let (n_i, n_cells_i) = (lv.n as i32, lv.n_cells as i32); unsafe { - rt.stream + stream .launch_builder(&rt.f_res) .arg(&n_cells_i) .arg(&lv.cells) @@ -248,7 +256,7 @@ impl DeviceVcycle { 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 + stream .launch_builder(&rt.f_restrict) .arg(&n_c_cells_i) .arg(&lc.cells) @@ -268,7 +276,7 @@ impl DeviceVcycle { 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 + stream .launch_builder(&rt.f_coarsest) .arg(&k_i) .arg(&n_cells_i) @@ -302,7 +310,7 @@ impl DeviceVcycle { 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 + stream .launch_builder(&rt.f_prolong) .arg(&n_cells_i) .arg(&lf.cells) @@ -317,10 +325,10 @@ impl DeviceVcycle { } self.smooth(l); } - rt.stream + stream .memcpy_dtoh(&self.levels[0].x, &mut self.z_f32) .expect("download z"); - rt.stream.synchronize().expect("sync"); + stream.synchronize().expect("sync"); for &idx in &self.fine_cells { z[idx as usize] = self.z_f32[idx as usize] as f64; } @@ -332,6 +340,42 @@ impl DeviceVcycle { // set, and the rendezvous that gathers the lanes' preconditioner calls. // --------------------------------------------------------------------------- +/// Stream mode (P3-iii K-stream): a lane thread launches its own V-cycles +/// on its own stream and never waits for other lanes; the GPU overlaps the +/// streams inside one process. Off = batch mode (the rendezvous). +static STREAM_MODE: AtomicBool = AtomicBool::new(false); + +/// Select the plane's mode for this process: `true` = one stream per lane +/// thread (no rendezvous), `false` = the batched rendezvous. +pub fn set_plane_streams(on: bool) { + STREAM_MODE.store(on, Ordering::SeqCst); +} + +pub fn plane_streams() -> bool { + STREAM_MODE.load(Ordering::SeqCst) +} + +thread_local! { + static LANE_STREAM: std::cell::RefCell>> = const { std::cell::RefCell::new(None) }; +} + +/// The stream a device operator built on this thread lives on: the lane's +/// own stream in stream mode on a lane thread, the default stream otherwise. +fn stream_for_this_thread() -> Arc { + let rt = runtime(); + if lane().is_some() && plane_streams() { + LANE_STREAM.with(|c| { + let mut slot = c.borrow_mut(); + if slot.is_none() { + *slot = Some(rt._ctx.new_stream().expect("lane stream")); + } + Arc::clone(slot.as_ref().expect("lane stream")) + }) + } else { + Arc::clone(&rt.stream) + } +} + /// 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`. @@ -627,6 +671,9 @@ impl Drop for CgGuard { /// Enter a CG solve on this lane (`None` when this thread is not a lane). pub(super) fn cg_enter() -> Option { lane()?; + if plane_streams() { + return None; + } plane_lock().in_cg += 1; Some(CgGuard(())) } @@ -813,6 +860,60 @@ mod tests { assert_eq!(bits(&z1), bits(&solo[1])); } + /// Stream mode: three lane threads, each with its own stream and its own + /// operator, applying concurrently with no rendezvous — every result + /// equals the solo apply on the default stream, bit for bit. + #[test] + fn stream_mode_lanes_equal_the_solo_apply_bit_for_bit() { + let (nx, ny) = (124usize, 21usize); + let n = nx * ny; + let probs: Vec = [(0.2, 0.2, 21u64), (0.7, 0.2, 23), (1.3, 0.22, 29)] + .iter() + .map(|&(cx, cy, s)| problem(nx, ny, cx, cy, s)) + .collect(); + let expected: Vec>> = probs + .iter() + .map(|p| { + let mut d = device_of(p); + (0..4) + .map(|c| { + let r: Vec = p.rhs.iter().map(|v| v * (1.0 + c as f64)).collect(); + let mut z = vec![0.0; n]; + d.apply(&r, &mut z); + bits(&z) + }) + .collect() + }) + .collect(); + let handles: Vec<_> = probs + .into_iter() + .enumerate() + .map(|(k, p)| { + std::thread::spawn(move || { + set_lane(Some(k)); + set_plane_streams(true); + assert!(cg_enter().is_none(), "stream mode must not rendezvous"); + let mut d = device_of(&p); + assert!(!Arc::ptr_eq(&d.stream, &runtime().stream), "lane stream"); + (0..4) + .map(|c| { + let r: Vec = p.rhs.iter().map(|v| v * (1.0 + c as f64)).collect(); + let mut z = vec![0.0; n]; + d.apply(&r, &mut z); + bits(&z) + }) + .collect::>() + }) + }) + .collect(); + let results: Vec>> = handles + .into_iter() + .map(|h| h.join().expect("lane thread")) + .collect(); + set_plane_streams(false); + assert_eq!(results, expected); + } + /// 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 diff --git a/crates/specialized/rtx-fsi/tests/fsi2_overset_plane.rs b/crates/specialized/rtx-fsi/tests/fsi2_overset_plane.rs index ba3fac8..9e9615e 100644 --- a/crates/specialized/rtx-fsi/tests/fsi2_overset_plane.rs +++ b/crates/specialized/rtx-fsi/tests/fsi2_overset_plane.rs @@ -10,6 +10,9 @@ //! (the `p7_points.txt` format; `#` comments and blank lines skipped). //! Without it the test returns. //! - `RTX_FSI2O_PLANE_DIR`: each lane's CSV goes to `/.csv`. +//! - `RTX_FSI2O_PLANE_MODE`: `streams` (default; one CUDA stream per lane +//! thread, no rendezvous — the K-stream design) or `batch` (the +//! rendezvous: every lane inside a CG gathered into one batched launch). //! - `RTX_FSI2O_PLANE_STACK_MB`: the lane threads' stack (default 64). //! Requires `RTX_FSI2O_MG_RB=1 RTX_FSI2O_MG_DEVICE=1` on a //! `--features rtx-cfd/cuda` build; the test fails loudly otherwise (a @@ -21,7 +24,7 @@ mod fsi2_harness; use fsi2_harness::overset_march::{OversetMarchConfig, run_march_overset}; use fsi2_harness::{BenchmarkCase, FSI2, FSI3}; -use rtx_cfd::solvers::incompressible::{plane_counters, set_plane_lane}; +use rtx_cfd::solvers::incompressible::{plane_counters, set_plane_lane, set_plane_streams}; use std::time::Instant; struct Point { @@ -80,6 +83,13 @@ fn fsi2_overset_plane() { device_on, "the plane needs RTX_FSI2O_MG_RB=1 RTX_FSI2O_MG_DEVICE=1 on a cuda build" ); + let mode = std::env::var("RTX_FSI2O_PLANE_MODE").unwrap_or_else(|_| "streams".into()); + let streams = match mode.as_str() { + "streams" => true, + "batch" => false, + other => panic!("RTX_FSI2O_PLANE_MODE must be streams or batch, got {other:?}"), + }; + set_plane_streams(streams); let points = read_points(&points_path); let k = points.len(); assert!(k >= 1, "no plane points in {points_path}"); @@ -113,7 +123,7 @@ fn fsi2_overset_plane() { FSI2 }; println!( - " MARCH PLANE: {k} lanes in one process ({} base, ny {}, t_end {}, coupler {}, s = {}); points from {points_path}, CSVs under {dir}", + " MARCH PLANE: {k} lanes in one process, mode {mode} ({} base, ny {}, t_end {}, coupler {}, s = {}); points from {points_path}, CSVs under {dir}", base.name, config.ny, config.t_end, config.coupler, config.subcycle ); for (i, p) in points.iter().enumerate() { @@ -176,7 +186,7 @@ fn fsi2_overset_plane() { let plane_wall = t_plane.elapsed().as_secs_f64(); let (batches, served) = plane_counters(); println!( - " MARCH PLANE done: {k} lanes in {plane_wall:.0} s wall; {batches} batched V-cycle launches served {served} lane calls ({:.2} lanes per batch)", + " MARCH PLANE done: {k} lanes in {plane_wall:.0} s wall, mode {mode}; {batches} batched V-cycle launches served {served} lane calls ({:.2} lanes per batch)", if batches > 0 { served as f64 / batches as f64 } else {