Safe writes for block-owned rows: a block-strided column witness for DisjointSlice #1692

What problem does this solve?
Block-per-row kernels (LayerNorm, softmax, row reductions, one block per row and thread t taking columns t, t + blockDim.x, …) have no safe output path. DisjointSlice offers per-thread witnesses (index_1d[_u32], element_thread32, tile_thread32, thread_run32, grid_stride_runs32, tile_2d32; view.rs:1135-1435), one per-warp space (thread::warp_index, thread.rs:1006), and, in open PR #1585, one scalar per block, but nothing for "block b owns row b, thread t owns its block-strided columns". These kernels end in unsafe { out.get_unchecked_mut(row * cols + c) } with a SAFETY comment. An in-place variant, normalizing a buffer over itself, can't take the buffer as both &[T] and DisjointSlice<T> (aliasing), so it reads and writes through as_mut_ptr().

What a kernel writes today

Both examples assume a 1D launch (grid and block in x only), which nothing checks.

use cuda_device::thread::{blockDim_x, blockIdx_x, threadIdx_x};
use cuda_device::{DisjointSlice, kernel};

#[kernel]
pub fn scale_rows(mut out: DisjointSlice<f32>, x: &[f32], cols: u32, s: f32) {
    let row = blockIdx_x() as usize;
    let cols = cols as usize;
    let end = (row + 1) * cols;
    if end > out.len() || end > x.len() {
        return;
    }
    let mut c = threadIdx_x() as usize;
    while c < cols {
        let i = row * cols + c;
        // SAFETY: rows are disjoint per block and columns per thread (c = t mod blockDim.x),
        // and i < end <= out.len()
        unsafe { *out.get_unchecked_mut(i) = x[i] * s };
        c += blockDim_x() as usize;
    }
}

/// In place: the buffer can't be both `&[f32]` and `DisjointSlice<f32>`
#[kernel]
pub fn scale_rows_inplace(mut xy: DisjointSlice<f32>, cols: u32, s: f32) {
    let row = blockIdx_x() as usize;
    let cols = cols as usize;
    if (row + 1) * cols > xy.len() {
        return;
    }
    let p = xy.as_mut_ptr();
    let mut c = threadIdx_x() as usize;
    while c < cols {
        // SAFETY: as above; this thread is the only one touching this element
        unsafe {
            let e = p.add(row * cols + c);
            *e = *e * s;
        }
        c += blockDim_x() as usize;
    }
}

Both compile at 88e52d117. Each disjointness argument lives in a SAFETY comment that the type system doesn't check.

Proposed API / behavior

A separate block-row index space, following the RuntimeRowMajorTiles pattern:

#[kernel(launch_context = lc)]
#[launch_contract(domain = 1, coordinates = u32)]
pub fn layernorm(mut out: DisjointSlice<f32, BlockRows>, x: &[f32], eps: f32) {
    // the row width is the one the host bound into the launch packet with the slice;
    // None, uniformly per block, when row blockIdx.x is outside `out`
    let Some(mut row) = out.block_row32(lc) else { return };
    // yields InBoundsMut32 for columns threadIdx.x, threadIdx.x + blockDim.x, ...
    for mut y in row.thread_columns() {
        let c = y.column();
        y.write(/* ... */);
    }
}

// in place: read and write the same element through the one witness
for mut y in row.thread_columns() {
    let v = y.read();
    y.write((v - mean) * rstd);
}

The row width can't be a per-call argument. Two blocks passing different widths, or one thread-dependent choice between two widths, describe two grids over one buffer, and their "disjoint" rows collide: with 2 threads per block, block 0 at width 4 reaches element 2 through thread 0's second column, and so does block 1 at width 2 through thread 0's first. disjoint.rs:61-70 gives the same argument for tiles, and #515 and #584 discussed it. So the width is bound to the slice once, on the host, as SpaceLayout::Data already allows.

The space also can't accept flat Index1D witnesses: at width 512 with 256 threads, block 0's thread 1 owns row element 257 through its second column, and block 1's thread 1 would reach flat element 257 through element_thread32. With those two rules, a 1D launch gives each block one row (blockIdx.x) and each thread its residue class mod blockDim.x. The bounds result is uniform within a block when every thread reaches the call, which keeps it safe ahead of sync_threads.

Alternatives considered

• DisjointSlice<T, BlockIndex> from open PR Add checked per-block scalar output witnesses #1585 gives thread (0, 0, 0) of each block one scalar. That covers block reductions, but not a row that every thread of the block writes. This request is the row-shaped sibling of that space.
• get_unchecked_mut with a SAFETY comment, which is what sam3rs ships today in 4 kernels.
• WarpIndex covers only one value per warp.
• One thread per element doesn't fit, because the row needs a block reduction first.

Additional context
Found porting sam3rs's ~50 CUDA C++ kernels to cuda-oxide (main 2d47e43, rechecked at 88e52d117): layernorm, layernorm_bf16, layernorm_inplace, softmax. The in-place LayerNorm also showed the aliasing problem: host code that called layernorm(out = x, x) was UB under the &[T] + DisjointSlice signature, so the kernel had to be split.