rtx-cfd/rtx-fsi PERF-2 P3-iii K-stream mode: each lane thread's device operator lives on its own CUDA stream (no rendezvous; set_plane_streams, RTX_FSI2O_PLANE_MODE=streams default | batch); stream-mode unit test digit-identical to the solo apply

Co-Authored-By: Claude Fable 5.1 <[email protected]>

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};

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<usize>) {}
 #[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)]

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<CudaStream>,
     levels: Vec<DevLevel>,
     sweeps: usize,
     fine_cells: Vec<u32>,
@@ -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<u32> {
-            rt.stream
+            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 up_f = |v: &[f32]| -> CudaSlice<f32> { stream.memcpy_stod(v).expect("upload") };
         let dev: Vec<DevLevel> = 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::<f32>(n).expect("alloc"),
-                    x: rt.stream.alloc_zeros::<f32>(n).expect("alloc"),
-                    r: rt.stream.alloc_zeros::<f32>(n).expect("alloc"),
+                    b: stream.alloc_zeros::<f32>(n).expect("alloc"),
+                    x: stream.alloc_zeros::<f32>(n).expect("alloc"),
+                    r: stream.alloc_zeros::<f32>(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<Option<Arc<CudaStream>>> = 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<CudaStream> {
+    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<CgGuard> {
     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<PoissonProblem> = [(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<Vec<Vec<u64>>> = probs
+            .iter()
+            .map(|p| {
+                let mut d = device_of(p);
+                (0..4)
+                    .map(|c| {
+                        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);
+                        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<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);
+                            bits(&z)
+                        })
+                        .collect::<Vec<_>>()
+                })
+            })
+            .collect();
+        let results: Vec<Vec<Vec<u64>>> = 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