Typed CTA-shared pointers so TMA G2S can use .shared::cta (SM12x setmaxnreg loss, #1238) #1337

Problem

#1238: a kernel that combines setmaxnreg register donation with cp_async_bulk_tensor_*_g2s loses both USETMAXREG instructions on SM12x (ptxas C7506).

Root cause, from hand-written PTX assembled with CUDA 13.3 ptxas and inspected in SASS:

On SM12x a unicast TMA into another CTA's shared memory is emulated by a syscall. ptxas keeps that path whenever it cannot prove that the destination is CTA-local. Only a statically emitted .shared::cta form avoids it; runtime checks and address tricks do not.

Options considered

• Select .shared::cta by target (fix(tma): select CTA-local G2S on SM120/SM121 #1249). Unsound: SM12x does support peer destinations, through the emulated path.
• An explicit _g2s_cta API with an unsafe locality promise (feat(tma): add CTA-local G2S tile intrinsics #1336). Works, but the caller carries the proof.
• Infer locality from pointer provenance, by propagating address spaces through casts in the importer. Pointers passed to helpers stay generic, since helpers are not inlined, and changing pointer types tree-wide affects unrelated inline-PTX lowerings, such as tcgen05.alloc's cvta.to.shared.
• Carry the address space in the type, like CuTe's smem_ptr. Locality is known statically, survives helper calls, and needs no dataflow. The importer already makes *mut Barrier and *mut SharedArray<..> addrspace(3) by type; this extends that to arbitrary CTA-shared pointers.

Plan

• feat(cuda-device): add SharedPtr, a CTA-shared pointer typed addrspace(3) #1341 SharedPtr<T>: a CTA-shared pointer that is addrspace(3) by type.
  • Built safely from SharedArray / DynamicSharedArray, or with unsafe SharedPtr::from_raw.
  • add / byte_add / cast keep it SharedPtr, and as_ptr widens it to a generic pointer.
• feat(tma): add CTA-local G2S tile intrinsics #1336 (on top of feat(cuda-device): add SharedPtr, a CTA-shared pointer typed addrspace(3) #1341): cp_async_bulk_tensor_{1..5}d_g2s_cta(dst: SharedPtr<u8>, …, barrier: *mut Barrier), lowered to .shared::cta, so locality comes from the type rather than the safety contract.
• Migrate the TMA G2S lowering makes ptxas ignore setmaxnreg via an extern compatibility call #1238 reproducer and the in-tree SM12x TMA users (tma_copy, gemm_sol, tcgen05_matmul).
• Follow-ups:
  • a ClusterSharedPtr from map_shared_rank;
  • SharedPtr forms for cp.async, ldmatrix/stmatrix, and tcgen05 and wgmma descriptors.

The existing cp_async_bulk_tensor_*_g2s keeps its cluster contract unchanged.

🤖 Generated with Claude Code