[YeonwooSung]

What this adds

TmemGuard::alloc_by and dealloc_by select the warp that runs the warp-collective tcgen05.alloc and tcgen05.dealloc. They used warp::warp_id(), which is threadIdx.x / 32 and names a hardware warp only in a 1D block. In a 2D or 3D block every row with the same threadIdx.x range matched, so several warps ran the allocation. The two methods now select by block-linear warp id, the numbering the hardware uses to form warps.

linear_tid = threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z)
warp       = linear_tid / 32

blockDim = (32, 4, 1), alloc_by(0):
  before: threadIdx.x / 32 == 0   selects 128 threads, four full warps, four tcgen05.alloc
  after:  linear_tid / 32 == 0    selects 32 threads, the y == 0 row, one tcgen05.alloc
blockDim = (256, 1, 1): both rules select threads 0..=31, unchanged

static mut TMEM_SLOT: SharedArray<u32, 1, 4> = SharedArray::UNINIT;

// Every thread of the (32, 4, 1) block calls both; only block-linear warp 0 issues the instructions.
let tmem = TmemGuard::<TmemUninit, 128>::from_static(&raw mut TMEM_SLOT as *mut u32);
let tmem = unsafe { tmem.alloc() };
// ... MMA through tmem.address() ...
let _done = unsafe { tmem.dealloc() };

What we fixed

• alloc_by and dealloc_by read threadIdx.{x,y,z} and blockDim.{x,y} and compare the block-linear warp id, so one hardware warp executes each .sync.aligned instruction in every block layout.
• One crate-private linearization helper now serves TmemGuard, the cooperative-groups block reductions and thread_rank(); the private copy in cooperative_groups.rs is gone.
• The tcgen05 usage pattern, the warp_id() rustdoc and the book no longer present threadIdx.x / 32 as adequate for multi-dimensional blocks.

Contract

• alloc_by(n) and dealloc_by(n) take a block-linear warp id: warp 0 is linear tids 0..=31. In a 1D block this equals threadIdx.x / 32, so existing callers are unchanged.
• The selected warp must be a full 32-thread warp: the block must hold at least 32 * (n + 1) threads.
• Any one full warp of the CTA may deallocate once all TMEM accesses are complete; using the allocating warp is the convention. A kernel that allocates more than once must issue every allocation from the same warp.
• warp::warp_id() is unchanged and remains documented as 1D-only.

Implementation

• thread::block_linear_thread_id(tx, ty, tz, dx, dy) is a pure const fn, host-testable; block_linear_thread_id_here() reads the special registers through the generated stubs.
• warp::block_linear_warp_id divides it by 32; block_linear_warp_id_here() is what alloc_by and dealloc_by call.
• blockDim.z is not needed: z is the slowest axis.
• Host tests enumerate every thread of 256x1x1, 32x4x1 and 8x4x2 blocks and check warp membership, plus a control that the old rule selects 128 threads of a 32x4 block and straddles two warps at blockDim.x = 48.
• The tcgen05 example gains two 32x4 TmemGuard kernels behind its existing sm_100 gate, one TMEM allocation each (warp 0 and warp 3).

Verification

Checked at 5528df9fd (the PR head merged with main e84e5b3d7):

• ✅ cargo test -p cuda-device --all-targets: 111 lib tests, integration, compile-fail and doctests pass.
• ✅ cargo clippy -p cuda-device --all-targets -- -D warnings, RUSTDOCFLAGS='-D warnings' cargo doc -p cuda-device, cargo oxide fmt --check, all 18 scripts/check-*.sh guards, scripts/sync-example-locks.sh --check and just book pass. All 62 CI checks passed.
• ✅ sm_100a PTX of the two 32x4 TmemGuard kernels: one tcgen05.alloc and one tcgen05.dealloc each, predicated on the block-linear thread id (%tid.y, %tid.z, %ntid.x, %ntid.y, two mad.lo); a kernel built with the old rule compares %tid.x alone. ptxas -arch=sm_100a accepts the PTX; the libNVVM route (--emit-nvvm-ir --arch sm_100a) builds the same selection.
• ✅ Hardware census on an RTX 5090 (sm_120): in a 32x4 block warp_id() == 0 holds for 128 threads in four full warps, the linear rule for 32; with blockDim.x = 16 the old rule selects half-warps; in a 256x1 block both rules agree. memcheck, synccheck and racecheck clean.
• ✅ tcgen05, tcgen05_matmul and gemm_sol PTX unchanged (identical per-kernel opcode histograms); they use raw tcgen05_alloc on 1D blocks.

Caveats

• tcgen05.alloc and tcgen05.dealloc cannot execute on sm_120, so the 2D path is verified by PTX, ptxas and the census of the selection rule, not by a TMEM allocation on hardware; the example's new kernels run only on an sm_100 machine.
• Users of the manual tcgen05_alloc path still select their warp themselves; this_thread_block().thread_rank() / 32 is the public block-linear id.

---

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[nihalpasham]

/ok to test 5528df9

[nihalpasham]

Good catch, and the same bug sat in your mbarrier PR, so I gave both one shared helper:

• The block-linear thread id now lives once in thread.rs; TmemGuard, the cooperative-groups reductions and init_by/inval_by (fix(cuda-device): run mbarrier init on one linear thread #1349) all use it instead of carrying their own formula.
• The tcgen05 usage notes, the warp_id() docs and two book lines still presented threadIdx.x / 32 as fine for 2D blocks; fixed.
• alloc_by/dealloc_by now state the rules: a full 32-thread warp, any one full warp may deallocate once TMEM accesses are done, and repeated allocations must come from the same warp.
• A control test for the old rule (128 threads selected in a 32x4 block, two half warps at blockDim.x = 48) and two 32x4 TmemGuard kernels in the tcgen05 example, one allocation each, so CI compiles the path.

On the 5090 the old rule picks four full warps of a 32x4 block and the new one picks one; the sm_100a PTX gates one tcgen05.alloc and one tcgen05.dealloc per kernel on the linear id, and ptxas accepts it. Running tcgen05.alloc in a 2D block for real needs sm_100 hardware, which I do not have here. 111 cuda-device tests, clippy, rustdoc, the book, guards and all 62 CI checks green at 5528df9fd.