From 5a27ef4b6a552b599b0a73e6dac8aa5a5b07c2a8 Mon Sep 17 00:00:00 2001 From: YeonwooSung Date: Mon, 28 Sep 2026 14:42:01 +0900 Subject: [PATCH 1/5] fix(cuda-device): run mbarrier init on one linear thread ManagedBarrier::init_by and inval_by selected the actor with threadIdx.x alone. In a 2D or 3D block that matches every thread with that x, so more than one thread executed mbarrier_init or mbarrier_inval. Select the thread by its x-major linear id instead. init() and inval() still pass 0, which is thread (0, 0, 0), and a 1D block still compares against threadIdx.x. Signed-off-by: YeonwooSung --- crates/cuda-device/src/barrier.rs | 174 ++++++++++++++++++++++++++++-- 1 file changed, 166 insertions(+), 8 deletions(-) diff --git a/crates/cuda-device/src/barrier.rs b/crates/cuda-device/src/barrier.rs index 703213caf3..e5e6dbc914 100644 --- a/crates/cuda-device/src/barrier.rs +++ b/crates/cuda-device/src/barrier.rs @@ -358,8 +358,8 @@ impl ManagedBarrier { /// Initialize the barrier with a specific thread performing initialization. /// - /// **All threads in the block should call this.** Only the thread with - /// `threadIdx.x == init_thread` performs the actual initialization; + /// **All threads in the block should call this.** Only the thread whose + /// linear id equals `init_thread` performs the actual initialization; /// all threads synchronize and receive a `Ready` handle. /// /// # Block-Scoped Barriers @@ -371,7 +371,9 @@ impl ManagedBarrier { /// # Parameters /// /// - `count`: Expected number of arrivals before barrier completes - /// - `init_thread`: Thread ID (threadIdx.x) that performs initialization + /// - `init_thread`: Linear thread id within the block that performs + /// initialization. `0` is thread `(0, 0, 0)`. In a 1D block this is + /// `threadIdx.x`. /// /// # Safety /// @@ -390,7 +392,19 @@ impl ManagedBarrier { /// ``` #[inline(always)] pub unsafe fn init_by(self, count: u32, init_thread: u32) -> ManagedBarrier { - if crate::thread::threadIdx_x() == init_thread { + if is_selected_linear_thread( + ( + crate::thread::threadIdx_x(), + crate::thread::threadIdx_y(), + crate::thread::threadIdx_z(), + ), + ( + crate::thread::blockDim_x(), + crate::thread::blockDim_y(), + crate::thread::blockDim_z(), + ), + init_thread, + ) { unsafe { mbarrier_init(self.ptr as *mut Barrier, count); fence_proxy_async_shared_cta(); @@ -509,15 +523,17 @@ impl ManagedBarrier { /// Invalidate the barrier with a specific thread performing invalidation. /// - /// **All threads in the block should call this.** Only the thread with - /// `threadIdx.x == inval_thread` performs the actual invalidation; + /// **All threads in the block should call this.** Only the thread whose + /// linear id equals `inval_thread` performs the actual invalidation; /// all threads synchronize before returning. /// /// Consumes the `Ready` barrier and returns an `Invalidated` barrier. /// /// # Parameters /// - /// - `inval_thread`: Thread ID (threadIdx.x) that performs invalidation + /// - `inval_thread`: Linear thread id within the block that performs + /// invalidation. `0` is thread `(0, 0, 0)`. In a 1D block this is + /// `threadIdx.x`. /// /// # Safety /// @@ -536,7 +552,19 @@ impl ManagedBarrier { // Ensure all threads are done with the barrier before invalidating crate::thread::sync_threads(); - if crate::thread::threadIdx_x() == inval_thread { + if is_selected_linear_thread( + ( + crate::thread::threadIdx_x(), + crate::thread::threadIdx_y(), + crate::thread::threadIdx_z(), + ), + ( + crate::thread::blockDim_x(), + crate::thread::blockDim_y(), + crate::thread::blockDim_z(), + ), + inval_thread, + ) { unsafe { mbarrier_inval(self.ptr as *mut Barrier) }; } @@ -566,3 +594,133 @@ pub type TmaBarrier0 = ManagedBarrier; /// Double-buffered TMA barrier #1 pub type TmaBarrier1 = ManagedBarrier; + +/// Whether this thread is the single block thread selected by a linear id. +/// +/// `tid = threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z)`. +/// `block_dim.z` is part of the launch shape but does not appear in the +/// x-major linearization. Exactly one `(thread_idx, block_dim)` pair matches +/// a given in-range `selected_linear_id`. +#[inline(always)] +fn is_selected_linear_thread( + thread_idx: (u32, u32, u32), + block_dim: (u32, u32, u32), + selected_linear_id: u32, +) -> bool { + let (tx, ty, tz) = thread_idx; + let (dx, dy, _dz) = block_dim; + let tid = tx + dx * (ty + dy * tz); + tid == selected_linear_id +} + +#[cfg(test)] +mod tests { + use super::is_selected_linear_thread; + + #[test] + fn one_d_block_matches_thread_idx_x() { + let block_dim = (128, 1, 1); + for x in 0..128 { + let thread_idx = (x, 0, 0); + assert!(is_selected_linear_thread(thread_idx, block_dim, x)); + if x != 0 { + assert!(!is_selected_linear_thread(thread_idx, block_dim, 0)); + } + if x != 31 { + assert!(!is_selected_linear_thread(thread_idx, block_dim, 31)); + } + } + } + + #[test] + fn two_d_block_selects_one_linear_thread() { + // blockDim = (32, 4, 1). Linear id is x-major: + // tid = x + 32 * y. + let block_dim = (32, 4, 1); + + assert!(is_selected_linear_thread((0, 0, 0), block_dim, 0)); + for y in 0..4 { + for x in 0..32 { + let selected = is_selected_linear_thread((x, y, 0), block_dim, 0); + assert_eq!(selected, x == 0 && y == 0); + } + } + + // Linear 32 is (0, 1), not every thread with threadIdx.x == 0. + assert!(is_selected_linear_thread((0, 1, 0), block_dim, 32)); + let mut matches = 0u32; + for y in 0..4 { + for x in 0..32 { + if is_selected_linear_thread((x, y, 0), block_dim, 32) { + matches += 1; + assert_eq!((x, y), (0, 1)); + } + } + } + assert_eq!(matches, 1); + + // The old threadIdx.x predicate would let every y with x == 0 through. + let x0_threads = [(0, 0, 0), (0, 1, 0), (0, 2, 0), (0, 3, 0)]; + let x0_hits = x0_threads + .iter() + .filter(|tid| is_selected_linear_thread(**tid, block_dim, 0)) + .count(); + assert_eq!(x0_hits, 1); + } + + #[test] + fn three_d_block_selects_one_linear_thread() { + // blockDim = (8, 4, 2). + // tid = x + 8 * (y + 4 * z). + let block_dim = (8, 4, 2); + + assert!(is_selected_linear_thread((0, 0, 0), block_dim, 0)); + // Linear 8 is (0, 1, 0); linear 32 is (0, 0, 1). + assert!(is_selected_linear_thread((0, 1, 0), block_dim, 8)); + assert!(is_selected_linear_thread((0, 0, 1), block_dim, 32)); + assert!(is_selected_linear_thread( + (3, 2, 1), + block_dim, + 3 + 8 * (2 + 4 * 1) + )); + + for selected in [0u32, 8, 31, 32] { + let mut matches = 0u32; + for z in 0..2 { + for y in 0..4 { + for x in 0..8 { + if is_selected_linear_thread((x, y, z), block_dim, selected) { + matches += 1; + } + } + } + } + assert_eq!(matches, 1, "selected linear id {selected}"); + } + } + + #[test] + fn two_threads_never_both_match_the_same_linear_id() { + let shapes = [(64, 1, 1), (32, 4, 1), (8, 4, 2), (16, 2, 2)]; + for (dx, dy, dz) in shapes { + let block_dim = (dx, dy, dz); + let n = dx * dy * dz; + for selected in 0..n { + let mut hits = 0u32; + for z in 0..dz { + for y in 0..dy { + for x in 0..dx { + if is_selected_linear_thread((x, y, z), block_dim, selected) { + hits += 1; + } + } + } + } + assert_eq!( + hits, 1, + "block {block_dim:?} selected {selected} matched {hits} threads" + ); + } + } + } +} From 4c1391e0e08b3a807d2bf054269d4b8c853da9d8 Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Mon, 19 Oct 2026 17:02:24 +0530 Subject: [PATCH 2/5] fix(cuda-device): elect barrier threads by the shared block-linear id init_by and inval_by now compare thread::block_linear_thread_id_here() with the selector instead of a private copy of the x-major linearization. The private predicate and its tests are gone; one property test next to the shared helper in thread.rs checks that every in-range id selects one thread across 1D, 2D, 3D and non-power-of-two blocks and that an id equal to the thread count selects nobody. This also removes the test literal that failed clippy::identity_op under -D warnings, and the init_by and inval_by safety rules now require a block-uniform selector below the block's thread count. Signed-off-by: nihalpasham --- cuda-oxide/crates/cuda-device/src/barrier.rs | 170 ++----------------- cuda-oxide/crates/cuda-device/src/thread.rs | 36 ++++ 2 files changed, 48 insertions(+), 158 deletions(-) diff --git a/cuda-oxide/crates/cuda-device/src/barrier.rs b/cuda-oxide/crates/cuda-device/src/barrier.rs index e5e6dbc914..6294f46a1d 100644 --- a/cuda-oxide/crates/cuda-device/src/barrier.rs +++ b/cuda-oxide/crates/cuda-device/src/barrier.rs @@ -379,7 +379,10 @@ impl ManagedBarrier { /// /// - Must be called before any arrive/wait operations /// - All participating threads in the block must call this together - /// - `init_thread` must be a valid thread ID within the block + /// - `init_thread` must be the same value in every thread and below + /// `blockDim.x * blockDim.y * blockDim.z`. A selector that differs between + /// threads elects zero or several threads; an out-of-range id selects no + /// thread, so the barrier is never initialized. /// /// # Example /// @@ -392,19 +395,8 @@ impl ManagedBarrier { /// ``` #[inline(always)] pub unsafe fn init_by(self, count: u32, init_thread: u32) -> ManagedBarrier { - if is_selected_linear_thread( - ( - crate::thread::threadIdx_x(), - crate::thread::threadIdx_y(), - crate::thread::threadIdx_z(), - ), - ( - crate::thread::blockDim_x(), - crate::thread::blockDim_y(), - crate::thread::blockDim_z(), - ), - init_thread, - ) { + // Block-linear id: x fastest, then y, then z; 0 is thread (0, 0, 0). + if crate::thread::block_linear_thread_id_here() == init_thread { unsafe { mbarrier_init(self.ptr as *mut Barrier, count); fence_proxy_async_shared_cta(); @@ -539,7 +531,10 @@ impl ManagedBarrier { /// /// - All threads must have completed their wait operations before calling /// - All participating threads in the block must call this together - /// - `inval_thread` must be a valid thread ID within the block + /// - `inval_thread` must be the same value in every thread and below + /// `blockDim.x * blockDim.y * blockDim.z`. A selector that differs between + /// threads elects zero or several threads; an out-of-range id selects no + /// thread, so the barrier is never invalidated. /// /// # Example /// @@ -552,19 +547,8 @@ impl ManagedBarrier { // Ensure all threads are done with the barrier before invalidating crate::thread::sync_threads(); - if is_selected_linear_thread( - ( - crate::thread::threadIdx_x(), - crate::thread::threadIdx_y(), - crate::thread::threadIdx_z(), - ), - ( - crate::thread::blockDim_x(), - crate::thread::blockDim_y(), - crate::thread::blockDim_z(), - ), - inval_thread, - ) { + // Block-linear id: x fastest, then y, then z; 0 is thread (0, 0, 0). + if crate::thread::block_linear_thread_id_here() == inval_thread { unsafe { mbarrier_inval(self.ptr as *mut Barrier) }; } @@ -594,133 +578,3 @@ pub type TmaBarrier0 = ManagedBarrier; /// Double-buffered TMA barrier #1 pub type TmaBarrier1 = ManagedBarrier; - -/// Whether this thread is the single block thread selected by a linear id. -/// -/// `tid = threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z)`. -/// `block_dim.z` is part of the launch shape but does not appear in the -/// x-major linearization. Exactly one `(thread_idx, block_dim)` pair matches -/// a given in-range `selected_linear_id`. -#[inline(always)] -fn is_selected_linear_thread( - thread_idx: (u32, u32, u32), - block_dim: (u32, u32, u32), - selected_linear_id: u32, -) -> bool { - let (tx, ty, tz) = thread_idx; - let (dx, dy, _dz) = block_dim; - let tid = tx + dx * (ty + dy * tz); - tid == selected_linear_id -} - -#[cfg(test)] -mod tests { - use super::is_selected_linear_thread; - - #[test] - fn one_d_block_matches_thread_idx_x() { - let block_dim = (128, 1, 1); - for x in 0..128 { - let thread_idx = (x, 0, 0); - assert!(is_selected_linear_thread(thread_idx, block_dim, x)); - if x != 0 { - assert!(!is_selected_linear_thread(thread_idx, block_dim, 0)); - } - if x != 31 { - assert!(!is_selected_linear_thread(thread_idx, block_dim, 31)); - } - } - } - - #[test] - fn two_d_block_selects_one_linear_thread() { - // blockDim = (32, 4, 1). Linear id is x-major: - // tid = x + 32 * y. - let block_dim = (32, 4, 1); - - assert!(is_selected_linear_thread((0, 0, 0), block_dim, 0)); - for y in 0..4 { - for x in 0..32 { - let selected = is_selected_linear_thread((x, y, 0), block_dim, 0); - assert_eq!(selected, x == 0 && y == 0); - } - } - - // Linear 32 is (0, 1), not every thread with threadIdx.x == 0. - assert!(is_selected_linear_thread((0, 1, 0), block_dim, 32)); - let mut matches = 0u32; - for y in 0..4 { - for x in 0..32 { - if is_selected_linear_thread((x, y, 0), block_dim, 32) { - matches += 1; - assert_eq!((x, y), (0, 1)); - } - } - } - assert_eq!(matches, 1); - - // The old threadIdx.x predicate would let every y with x == 0 through. - let x0_threads = [(0, 0, 0), (0, 1, 0), (0, 2, 0), (0, 3, 0)]; - let x0_hits = x0_threads - .iter() - .filter(|tid| is_selected_linear_thread(**tid, block_dim, 0)) - .count(); - assert_eq!(x0_hits, 1); - } - - #[test] - fn three_d_block_selects_one_linear_thread() { - // blockDim = (8, 4, 2). - // tid = x + 8 * (y + 4 * z). - let block_dim = (8, 4, 2); - - assert!(is_selected_linear_thread((0, 0, 0), block_dim, 0)); - // Linear 8 is (0, 1, 0); linear 32 is (0, 0, 1). - assert!(is_selected_linear_thread((0, 1, 0), block_dim, 8)); - assert!(is_selected_linear_thread((0, 0, 1), block_dim, 32)); - assert!(is_selected_linear_thread( - (3, 2, 1), - block_dim, - 3 + 8 * (2 + 4 * 1) - )); - - for selected in [0u32, 8, 31, 32] { - let mut matches = 0u32; - for z in 0..2 { - for y in 0..4 { - for x in 0..8 { - if is_selected_linear_thread((x, y, z), block_dim, selected) { - matches += 1; - } - } - } - } - assert_eq!(matches, 1, "selected linear id {selected}"); - } - } - - #[test] - fn two_threads_never_both_match_the_same_linear_id() { - let shapes = [(64, 1, 1), (32, 4, 1), (8, 4, 2), (16, 2, 2)]; - for (dx, dy, dz) in shapes { - let block_dim = (dx, dy, dz); - let n = dx * dy * dz; - for selected in 0..n { - let mut hits = 0u32; - for z in 0..dz { - for y in 0..dy { - for x in 0..dx { - if is_selected_linear_thread((x, y, z), block_dim, selected) { - hits += 1; - } - } - } - } - assert_eq!( - hits, 1, - "block {block_dim:?} selected {selected} matched {hits} threads" - ); - } - } - } -} diff --git a/cuda-oxide/crates/cuda-device/src/thread.rs b/cuda-oxide/crates/cuda-device/src/thread.rs index 5dd36cca53..0b3f33c435 100644 --- a/cuda-oxide/crates/cuda-device/src/thread.rs +++ b/cuda-oxide/crates/cuda-device/src/thread.rs @@ -1415,6 +1415,24 @@ const fn validate_unroll_factor(factor: u32) { mod tests { use super::block_linear_thread_id; + /// Number of threads in a `(dx, dy, dz)` block whose block-linear id is + /// `selected`. `ManagedBarrier::init_by` and `inval_by` make this + /// comparison on the device, one thread at a time. + fn threads_selected_by(block_dim: (u32, u32, u32), selected: u32) -> u32 { + let (dx, dy, dz) = block_dim; + let mut hits = 0; + for tz in 0..dz { + for ty in 0..dy { + for tx in 0..dx { + if block_linear_thread_id(tx, ty, tz, dx, dy) == selected { + hits += 1; + } + } + } + } + hits + } + #[test] fn a_1d_block_numbers_threads_by_x() { for tx in 0..256 { @@ -1452,4 +1470,22 @@ mod tests { // One step along z skips a whole x-y plane. assert_eq!(block_linear_thread_id(0, 0, 1, dx, dy), dx * dy); } + + #[test] + fn every_in_range_id_selects_one_thread() { + for block_dim in [(64, 1, 1), (32, 4, 1), (8, 4, 2), (16, 2, 2), (24, 3, 1)] { + let (dx, dy, dz) = block_dim; + let threads = dx * dy * dz; + for selected in 0..threads { + assert_eq!( + threads_selected_by(block_dim, selected), + 1, + "block {block_dim:?} id {selected}" + ); + } + // An id at or past the thread count selects nobody, which is why + // init_by and inval_by require the selector to be in range. + assert_eq!(threads_selected_by(block_dim, threads), 0); + } + } } From f8e41d1819f34d49bdaa0703e23112f498d91b5b Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Mon, 19 Oct 2026 17:03:09 +0530 Subject: [PATCH 3/5] docs(cuda-device): describe ManagedBarrier as an all-threads protocol The thread-requirement table, the struct example and the from_static docs still described a thread-0-only protocol. from_static is a plain constructor every thread calls; init() and inval() are all-threads calls in which linear thread 0, that is thread (0, 0, 0), acts, and both synchronize the block, so the old from_static example that called init() inside a thread-0 branch put a block-wide sync in divergent code. The struct example also repeated the fence and sync that init() already performs. Signed-off-by: nihalpasham --- cuda-oxide/crates/cuda-device/src/barrier.rs | 55 +++++++++++--------- 1 file changed, 30 insertions(+), 25 deletions(-) diff --git a/cuda-oxide/crates/cuda-device/src/barrier.rs b/cuda-oxide/crates/cuda-device/src/barrier.rs index 6294f46a1d..b21907e0c4 100644 --- a/cuda-oxide/crates/cuda-device/src/barrier.rs +++ b/cuda-oxide/crates/cuda-device/src/barrier.rs @@ -254,29 +254,29 @@ impl BarrierToken { /// /// # Thread Requirements /// -/// | Operation | Thread Requirement | -/// |-----------------|---------------------------| -/// | `from_static()` | Single thread (thread 0) | -/// | `init()` | Single thread (thread 0) | -/// | `arrive()` | All participating threads | -/// | `wait()` | All participating threads | -/// | `inval()` | Single thread (thread 0) | +/// | Operation | Thread Requirement | +/// |-----------------|----------------------------------------| +/// | `from_static()` | All threads call | +/// | `init()` | All threads call; linear thread 0 acts | +/// | `arrive()` | All participating threads | +/// | `wait()` | All participating threads | +/// | `inval()` | All threads call; linear thread 0 acts | /// /// # Example /// /// ```rust,ignore -/// // Thread 0: Create and initialize +/// // All threads: create the handle and initialize. Linear thread 0, that is +/// // thread (0, 0, 0), runs mbarrier.init; init() then fences and +/// // synchronizes the block, so every thread holds a Ready handle. /// let bar = ManagedBarrier::::from_static(&raw mut BAR); /// let bar = unsafe { bar.init(128) }; // Now Ready -/// fence_proxy_async_shared_cta(); -/// sync_threads(); /// /// // All threads: arrive and wait /// let token = bar.arrive(); /// bar.wait(token); /// -/// // Thread 0: Cleanup -/// sync_threads(); +/// // All threads: clean up. inval() synchronizes the block, then linear +/// // thread 0 runs mbarrier.inval. /// unsafe { bar.inval(); } /// ``` pub struct ManagedBarrier { @@ -301,7 +301,11 @@ impl ManagedBarrier { /// /// # Thread Requirements /// - /// **Single-thread**: Only call from the initialization thread (thread 0). + /// **All threads**: this is a plain constructor with no device effect. Every + /// thread wraps the same static, and `init()` then elects linear thread 0 to + /// run `mbarrier.init`. Do not wrap `init()` in a thread-0 branch: it + /// synchronizes the whole block, and a block-wide sync inside a branch that + /// only some threads take is undefined, in practice a hang. /// /// # Example /// @@ -309,13 +313,12 @@ impl ManagedBarrier { /// static mut BAR0: Barrier = Barrier::UNINIT; /// static mut BAR1: Barrier = Barrier::UNINIT; /// - /// if tid == 0 { - /// let bar0 = ManagedBarrier::::from_static(&raw mut BAR0); - /// let bar1 = ManagedBarrier::::from_static(&raw mut BAR1); + /// // All threads: wrap and initialize; init() elects linear thread 0. + /// let bar0 = ManagedBarrier::::from_static(&raw mut BAR0); + /// let bar1 = ManagedBarrier::::from_static(&raw mut BAR1); /// - /// let bar0 = unsafe { bar0.init(32) }; - /// let bar1 = unsafe { bar1.init(32) }; - /// } + /// let bar0 = unsafe { bar0.init(32) }; + /// let bar1 = unsafe { bar1.init(32) }; /// ``` pub fn from_static(ptr: *mut Barrier) -> Self { ManagedBarrier { @@ -327,8 +330,9 @@ impl ManagedBarrier { /// Initialize the barrier with an expected arrival count. /// - /// **All threads in the block should call this.** Only thread 0 performs - /// the actual initialization; all threads synchronize and receive a `Ready` handle. + /// **All threads in the block should call this.** Only linear thread 0, + /// that is thread (0, 0, 0), performs the actual initialization; all threads + /// synchronize and receive a `Ready` handle. /// /// This is a convenience wrapper for `init_by(count, 0)`. /// @@ -347,7 +351,7 @@ impl ManagedBarrier { /// ```rust,ignore /// static mut BAR: Barrier = Barrier::UNINIT; /// - /// // ALL threads call init - only thread 0 actually initializes + /// // ALL threads call init - only linear thread 0 actually initializes /// let bar = ManagedBarrier::::from_static(&raw mut BAR); /// let bar = unsafe { bar.init(128) }; // All threads get Ready handle /// ``` @@ -489,8 +493,9 @@ impl ManagedBarrier { /// Invalidate the barrier. /// - /// **All threads in the block should call this.** Only thread 0 performs - /// the actual invalidation; all threads synchronize before returning. + /// **All threads in the block should call this.** Only linear thread 0, + /// that is thread (0, 0, 0), performs the actual invalidation; all threads + /// synchronize before returning. /// /// This is a convenience wrapper for `inval_by(0)`. /// @@ -505,7 +510,7 @@ impl ManagedBarrier { /// # Example /// /// ```rust,ignore - /// // ALL threads call inval - only thread 0 actually invalidates + /// // ALL threads call inval - only linear thread 0 actually invalidates /// let _dead = unsafe { bar.inval() }; // Consumes Ready, returns Invalidated /// ``` #[inline(always)] From 1f65405d7038defabb2fe94249bcec256252df54 Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Mon, 19 Oct 2026 17:07:01 +0530 Subject: [PATCH 4/5] test(future_apis): run ManagedBarrier in a (32, 4, 1) block Add test_managed_barrier_2d: every thread calls init_by(128, 127) and inval_by(127), so block-linear thread 127, which is thread (31, 3, 0), runs mbarrier.init and mbarrier.inval; both coordinates of that thread are nonzero and no thread has threadIdx.x equal to the selector. Thread 0 fills the barrier word with all-ones first, a determinism aid that keeps the test independent of what shared memory held at launch. Each thread writes 42 when its row-major index equals its block rank, and the host returns an error on any other value. Against the previous election by threadIdx.x alone no thread matched the selector and the launch failed with error 719; with the block-linear election all 128 threads complete. Signed-off-by: nihalpasham --- .../examples/future_apis/README.md | 5 + .../examples/future_apis/src/main.rs | 102 ++++++++++++++++++ 2 files changed, 107 insertions(+) diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/README.md index 3b52278735..46757b69b6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/README.md @@ -16,6 +16,7 @@ Tests new type-safe abstractions: `CuSimd` for SIMD values and `ManagedBar 1. **test_managed_barrier**: Single barrier typestate flow 2. **test_multi_barrier**: Multiple barriers (TMA + General) 3. **test_double_buffered_barriers**: Ping-pong pattern with two barriers +4. **test_managed_barrier_2d**: One barrier in a (32, 4, 1) block; every thread calls `init_by` and `inval_by`, and linear thread 127, that is thread (31, 3, 0), runs the instructions ## Key Concepts Demonstrated @@ -153,6 +154,10 @@ cargo oxide run future_apis All 32 threads completed 3 iterations ✓ Double-buffered PASSED +--- Test 6: ManagedBarrier in a (32, 4, 1) block --- + All 128 threads wrote 42 in a (32, 4, 1) block +✓ ManagedBarrier 2D PASSED + === ALL FUTURE APIS TESTS COMPLETED === ``` diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/src/main.rs index d53c89ba93..55426c2f0a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/src/main.rs @@ -11,11 +11,20 @@ //! //! Run: cargo oxide run future_apis +use cuda_device::cooperative_groups::{ThreadGroup, this_thread_block}; use cuda_device::cusimd::CuSimd; use cuda_device::{Barrier, GeneralBarrier, ManagedBarrier, TmaBarrier, Uninit}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; +/// Block of the 2D barrier test: 4 rows of 32 threads, one warp per row. +const BLOCK_2D_WIDTH: usize = 32; +const BLOCK_2D_HEIGHT: usize = 4; +const BLOCK_2D_THREADS: u32 = (BLOCK_2D_WIDTH * BLOCK_2D_HEIGHT) as u32; +/// Block-linear id of thread (31, 3, 0), the last thread of the block, so the +/// x and y terms of the id both count. No thread has `threadIdx.x` equal to it. +const BARRIER_2D_THREAD: u32 = BLOCK_2D_THREADS - 1; + // ============================================================================= // CUSIMD KERNELS // ============================================================================= @@ -231,6 +240,63 @@ mod kernels { *output_elem = sum; // Should be 3 } } + + /// Test kernel for `ManagedBarrier` in a 2D block of (32, 4, 1) threads. + /// + /// Every thread calls `init_by` and `inval_by`, as their contract asks: + /// both synchronize the whole block. The selector is the block-linear id + /// of thread (31, 3, 0), the last thread of the block, so the x and y + /// terms of the id both count. No thread has `threadIdx.x` equal to it, + /// so an election by `threadIdx.x` alone initializes nothing and the + /// launch fails with error 719. + /// + /// Thread 0 first fills the barrier word with all-ones. The store is a + /// determinism aid, not the detector: a missing init then faults on this + /// word instead of depending on what shared memory held at launch. + /// + /// Each thread writes 42 at its row-major 2D index when that index equals + /// its block rank, so the host also checks that `thread_rank()` numbers + /// the block x fastest, the same numbering the barrier election uses. + #[kernel] + pub fn test_managed_barrier_2d( + mut output: DisjointSlice>, + ) { + static mut BAR: Barrier = Barrier::UNINIT; + + let rank = this_thread_block().thread_rank() as usize; + + // Determinism aid, not the detector: see the kernel docs above + if rank == 0 { + // SAFETY: BAR is this block's own shared memory, one 8-byte aligned + // u64 that nothing in this launch has made an mbarrier object yet; + // the sync below orders the store before init_by. + unsafe { *(&raw mut BAR as *mut u64) = u64::MAX }; + } + thread::sync_threads(); + + // ALL threads create the Uninit handle (wraps the same shared static) + let bar = ManagedBarrier::::from_static(&raw mut BAR); + + // ALL threads call init_by - only linear thread 127 runs mbarrier.init + // init_by includes: fence + sync_threads; all threads get a Ready handle + let bar = unsafe { bar.init_by(BLOCK_2D_THREADS, BARRIER_2D_THREAD) }; + + // All threads arrive and wait + let token = bar.arrive(); + bar.wait(token); + + // ALL threads call inval_by - it synchronizes the block, then only + // linear thread 127 runs mbarrier.inval, then synchronizes again + let _dead = unsafe { bar.inval_by(BARRIER_2D_THREAD) }; + + // Write results - 42 when the row-major index is the block rank + if let Some(idx) = thread::index_2d::() { + let value = if idx.get() == rank { 42 } else { 1 }; + if let Some(output_elem) = output.get_mut(idx) { + *output_elem = value; + } + } + } } // ============================================================================= @@ -433,6 +499,42 @@ fn main() -> Result<(), Box> { } } + // ==================================================================== + // Test 6: ManagedBarrier in a 2D block + // ==================================================================== + println!("--- Test 6: ManagedBarrier in a ({BLOCK_2D_WIDTH}, {BLOCK_2D_HEIGHT}, 1) block ---"); + { + const N: usize = BLOCK_2D_WIDTH * BLOCK_2D_HEIGHT; + let mut output_dev = DeviceBuffer::::zeroed(&stream, N)?; + + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (BLOCK_2D_WIDTH as u32, BLOCK_2D_HEIGHT as u32, 1), + shared_mem_bytes: 0, + }; + + // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. + unsafe { module.test_managed_barrier_2d((stream).as_ref(), cfg, &mut output_dev) }?; + stream.synchronize()?; + + let output: Vec = output_dev.to_host_vec(&stream)?; + + // Every thread writes 42 once it passed the barrier and its block + // rank matched its row-major index. Any other value means the + // election this block layout exists to test went wrong, so it is an + // error, not a note. + if let Some(thread) = output.iter().position(|&v| v != 42) { + let got = output[thread]; + println!(" Results: {:?}", &output[..8.min(output.len())]); + return Err(format!( + "ManagedBarrier 2D: thread {thread} wrote {got} instead of 42 in a ({BLOCK_2D_WIDTH}, {BLOCK_2D_HEIGHT}, 1) block" + ) + .into()); + } + println!(" All {N} threads wrote 42 in a ({BLOCK_2D_WIDTH}, {BLOCK_2D_HEIGHT}, 1) block"); + println!("✓ ManagedBarrier 2D PASSED\n"); + } + println!("=== ALL FUTURE APIS TESTS COMPLETED ==="); Ok(()) } From fd2bb390252d600ed2f19494e8bb794766b6d7e1 Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Mon, 19 Oct 2026 18:07:29 +0530 Subject: [PATCH 5/5] docs(book): show which thread runs mbarrier.init in a 2D block A one-page figure for the book and the pull request, in the style of the tile-libraries direct-linkage figure. Signed-off-by: nihalpasham --- .../images/mbarrier-thread-election.svg | 750 ++++++++++++++++++ .../advanced/tensor-memory-accelerator.md | 12 + 2 files changed, 762 insertions(+) create mode 100644 cuda-oxide/cuda-oxide-book/advanced/images/mbarrier-thread-election.svg diff --git a/cuda-oxide/cuda-oxide-book/advanced/images/mbarrier-thread-election.svg b/cuda-oxide/cuda-oxide-book/advanced/images/mbarrier-thread-election.svg new file mode 100644 index 0000000000..fc3a56b807 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/advanced/images/mbarrier-thread-election.svg @@ -0,0 +1,750 @@ + + + + + +ManagedBarrier: one thread runs mbarrier.init, chosen by its block-linear id +All threads call init_by(count, id) and inval_by(id); one thread acts. Main chose it by threadIdx.x alone, which names one thread only in a 1D block; #1349 uses the block-linear id. +THE BLOCK (future_apis, test_managed_barrier_2d) +MAIN (threadIdx_x() == selector) +FIXED (block_linear_thread_id_here() == selector) + +block_dim = (32, 4, 1): 128 threads, one warp per row +y +threadIdx.x = 0 +31 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +0 +1 +2 +3 + +(0, 0, 0): linear id 0; the thread init() and inval() select + +x == 0 in rows 1 to 3: threadIdx.x == 0 as well, like (0, 0, 0) + +(31, 3, 0): linear id 31 + 32 * 3 = 127; no thread has x == 127 + +future_apis/src/main.rs: test_managed_barrier_2d +// ALL threads call init_by - only linear thread 127 runs mbarrier.init +let bar = unsafe { bar.init_by(BLOCK_2D_THREADS, BARRIER_2D_THREAD) }; +let token = bar.arrive(); +bar.wait(token); +let _dead = unsafe { bar.inval_by(BARRIER_2D_THREAD) }; +BLOCK_2D_THREADS = 128; BARRIER_2D_THREAD = 127 = thread (31, 3, 0) +As landed, from_static(), init(), init_by, inval() and inval_by are +all-threads calls; one thread, named by the selector, acts. The +selector must be the same value in every thread and below the thread +count: an id at or past it names nobody; nothing is initialized. + +main: barrier.rs, init_by and inval_by +if crate::thread::threadIdx_x() == init_thread { +the x coordinate alone: one thread only in a 1D block +selector 0 (what init() and inval() pass) + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +four threads (x == 0 in each row) run mbarrier.init +on one word: undefined (PTX ISA), silent on sm_120 +selector 127 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +no thread has threadIdx.x == 127, so nothing runs +mbarrier.init; 128 arrive on an uninitialized word +result: the launch fails +Error: DriverError(719, "unspecified launch failure") +the same example source against main's cuda-device + +fixed: barrier.rs, init_by and inval_by +if crate::thread::block_linear_thread_id_here() + == init_thread { +id = x + 32 * (y + 4 * z): x fastest, then y, then z +selector 0 (what init() and inval() pass) + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +one thread, (0, 0, 0); 1D blocks keep the same thread +selector 127 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +one thread, (31, 3, 0): 31 + 32 * 3 = 127 +result: 128 of 128 threads correct +All 128 threads wrote 42 in a (32, 4, 1) block +memcheck, synccheck, racecheck: 0 errors, 0 hazards + +thread::block_linear_thread_id_here() +the shared helper #1351 landed; thread_rank() uses it +too, so the PR's private copy of the formula is gone + + +THE PTX ELECTION (the predicate before mbarrier.init; tabs shown as spaces; bra.uni lines, labels and the count mov.b32 left out) + +main: test_managed_barrier (the 1D test; init(), so selector 0) +mov.u32 %r1, %tid.x; +setp.ne.b32 %p1, %r1, 0; +@%p1 bra $L__BB4_2; +mbarrier.init.shared.b64 [__shared_mem_e69151d2fb141f63_4], %r2; + +fixed: test_managed_barrier_2d (selector 127) +mov.u32 %r2, %tid.x; +mov.u32 %r3, %tid.y; +mov.u32 %r4, %tid.z; +mov.u32 %r5, %ntid.x; +mov.u32 %r6, %ntid.y; +mad.lo.s32 %r7, %r6, %r4, %r3; +// y + ntid.y * z +mad.lo.s32 %r8, %r7, %r5, %r2; +// x + ntid.x * (y + ntid.y * z) +// 7 lines left out: thread 0 stores all-ones into the barrier word, bar.sync +setp.ne.b32 %p2, %r8, 127; +@%p2 bra $L__BB2_4; +mbarrier.init.shared.b64 [__shared_mem_ddf9d77ea1d0101f_0], %r9; +What this proves: on main, %p1 reads one register, %tid.x. With the fix, %p2 reads five +registers through two mad and compares tid.x + ntid.x * (tid.y + ntid.y * tid.z) with the +selector, so mbarrier.init runs on at most one thread per block for any in-range selector +in any block shape. In both trees the same predicate guards mbarrier.inval later. On main +it admits four threads for selector 0 in a (32, 4, 1) block and no thread for selector 127. + diff --git a/cuda-oxide/cuda-oxide-book/advanced/tensor-memory-accelerator.md b/cuda-oxide/cuda-oxide-book/advanced/tensor-memory-accelerator.md index a4b9da986a..add95f44be 100644 --- a/cuda-oxide/cuda-oxide-book/advanced/tensor-memory-accelerator.md +++ b/cuda-oxide/cuda-oxide-book/advanced/tensor-memory-accelerator.md @@ -188,6 +188,18 @@ The lifecycle is: `Uninit` → `init()` → `Ready` → `arrive()` / `wait()` `inval()` → `Invalidated`. The type system prevents calling `wait` on an uninitialized barrier or `arrive` on an invalidated one. +```{figure} images/mbarrier-thread-election.svg +:align: center +:width: 100% + +Which thread runs `mbarrier.init`. Every thread calls `init_by(count, id)` and `inval_by(id)`, +and the thread whose block-linear id (`x + blockDim.x * (y + blockDim.y * z)`, the numbering +`thread_rank()` uses) equals `id` runs the instruction; `init()` and `inval()` pass 0, thread +(0, 0, 0). Comparing `threadIdx.x` alone named one thread only in a 1D block: in a (32, 4, 1) +block selector 0 ran the instruction on four threads and selector 127 on none, so the launch +failed with error 719. +``` + ### Key barrier rules | Rule | Why |