fix(cuda-device): run mbarrier init on one linear thread

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::<Uninit, TmaBarrier>::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<State, Kind, const ID: usize = 0> {
@@ -301,21 +301,24 @@ impl<Kind, const ID: usize> ManagedBarrier<Uninit, Kind, ID> {
     ///
     /// # 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
     ///
     /// ```rust,ignore
     /// static mut BAR0: Barrier = Barrier::UNINIT;
     /// static mut BAR1: Barrier = Barrier::UNINIT;
     ///
-    /// if tid == 0 {
-    ///     let bar0 = ManagedBarrier::<Uninit, TmaBarrier>::from_static(&raw mut BAR0);
-    ///     let bar1 = ManagedBarrier::<Uninit, GeneralBarrier>::from_static(&raw mut BAR1);
+    /// // All threads: wrap and initialize; init() elects linear thread 0.
+    /// let bar0 = ManagedBarrier::<Uninit, TmaBarrier>::from_static(&raw mut BAR0);
+    /// let bar1 = ManagedBarrier::<Uninit, GeneralBarrier>::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<Kind, const ID: usize> ManagedBarrier<Uninit, Kind, ID> {
 
     /// 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<Kind, const ID: usize> ManagedBarrier<Uninit, Kind, ID> {
     /// ```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::<Uninit, TmaBarrier>::from_static(&raw mut BAR);
     /// let bar = unsafe { bar.init(128) };  // All threads get Ready handle
     /// ```
@@ -358,8 +362,8 @@ impl<Kind, const ID: usize> ManagedBarrier<Uninit, Kind, ID> {
 
     /// 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,13 +375,18 @@ impl<Kind, const ID: usize> ManagedBarrier<Uninit, Kind, ID> {
     /// # 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
     ///
     /// - 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
     ///
@@ -390,7 +399,8 @@ impl<Kind, const ID: usize> ManagedBarrier<Uninit, Kind, ID> {
     /// ```
     #[inline(always)]
     pub unsafe fn init_by(self, count: u32, init_thread: u32) -> ManagedBarrier<Ready, Kind, ID> {
-        if crate::thread::threadIdx_x() == 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();
@@ -483,8 +493,9 @@ impl<Kind, const ID: usize> ManagedBarrier<Ready, Kind, ID> {
 
     /// 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)`.
     ///
@@ -499,7 +510,7 @@ impl<Kind, const ID: usize> ManagedBarrier<Ready, Kind, ID> {
     /// # 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)]
@@ -509,21 +520,26 @@ impl<Kind, const ID: usize> ManagedBarrier<Ready, Kind, ID> {
 
     /// 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
     ///
     /// - 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
     ///
@@ -536,7 +552,8 @@ impl<Kind, const ID: usize> ManagedBarrier<Ready, Kind, ID> {
         // Ensure all threads are done with the barrier before invalidating
         crate::thread::sync_threads();
 
-        if crate::thread::threadIdx_x() == 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) };
         }
 

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

cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/README.md

@@ -16,6 +16,7 @@ Tests new type-safe abstractions: `CuSimd<T, N>` 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 ===
 ```
 

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<u32, thread::Index2D<BLOCK_2D_WIDTH>>,
+    ) {
+        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::<Uninit, GeneralBarrier>::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::<BLOCK_2D_WIDTH>() {
+            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<dyn std::error::Error>> {
         }
     }
 
+    // ====================================================================
+    // 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::<u32>::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<u32> = 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(())
 }