feat(tma): add CTA-local G2S tile intrinsics

[lucifer1004]

What this adds

Two things that land together: SharedPtr<T>, a pointer into the issuing CTA's shared memory whose address space is part of its type, and cp_async_bulk_tensor_{1..5}d_g2s_cta, TMA tile loads into that memory through the PTX .shared::cta destination form (PTX 8.6, sm_90 and newer). The existing cp_async_bulk_tensor_Nd_g2s family and its cluster contract are unchanged.

A *mut T in device code is a generic pointer: once it crosses a helper boundary nothing says it points at this CTA's shared memory. SharedPtr<T> keeps that fact in the type and lowers to ptr addrspace(3) in parameters, returns, locals, struct fields, generics and closures, so the .shared::cta copy never has to guess, and ptxas on sm_120a no longer pays the runtime CTA check that made it drop setmaxnreg.

before (cluster form, sm_120a):
  cvta.to.shared::cluster.u64 ...
  cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes ...
  ptxas: C7506 setmaxnreg ignored, 24 registers, 8-byte stack
  SASS:  CALL.ABS.NOINC, UTMALDG.2D, no USETMAXREG

after (CTA form, sm_120a):
  cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::complete_tx::bytes [TILE], [tmap, {x, y}], [BAR]
  ptxas: no C7506, 168 registers, 0 stack
  SASS:  USETMAXREG.DEALLOC.CTAPOOL 0x28, UTMALDG.2D, USETMAXREG.TRY_ALLOC.CTAPOOL 0xe8, no CALL

static mut TILE: SharedArray<f32, 4096, 128> = SharedArray::UNINIT;
static mut BAR: Barrier = Barrier::UNINIT;

#[inline(never)]
unsafe fn load(dst: SharedPtr<u8>, tmap: &TmaDescriptor, x: i32, y: i32, bar: *mut Barrier) {
    unsafe { cp_async_bulk_tensor_2d_g2s_cta(dst, tmap, x, y, bar) } // dst is still addrspace(3) here
}

unsafe {
    let dst = SharedArray::shared_ptr(&raw mut TILE).cast::<u8>();
    load(dst, tensor_map, tile_x, tile_y, &raw mut BAR);
}

SharedPtr<T>
unsafe SharedArray::shared_ptr(&raw mut S)	the array's shared address, no conversion
DynamicSharedArray::<T>::shared_ptr() / shared_ptr_at(bytes)	as get() / offset(bytes)
unsafe SharedPtr::from_raw(p)	the explicit locality assertion, like __cvta_generic_to_shared or CuTe make_smem_ptr
add, offset, byte_add, cast	pointer arithmetic that keeps the type
as_ptr()	the generic *mut T, for loads, stores and raw-pointer APIs
addr_u32()	the 32-bit .shared::cta address, for the _shared_u32 forms

What we fixed

• On sm_120a, CUDA 13.3 and 13.4 ptxas lower every non-multicast .shared::cluster tensor load to a runtime CTA-id check with a syscall fallback. That call path makes ptxas ignore setmaxnreg (C7506), so a warp-specialized kernel runs with 24 registers instead of 168. Only a statically emitted .shared::cta form avoids it; sm_100a is native either way.
• A *mut u8 from as_mut_ptr() or get() loses its shared-memory origin at the first helper boundary. SharedPtr keeps that origin in the type, so the proof the .shared::cta form needs does not depend on a caller promise.
• A SharedPtr nested in another place, such as region.base after rustc inlines cast, was read as the address of an aggregate. The place walkers now drop the identity projection onto the newtype's field.
• A kernel parameter that holds a shared or local pointer inside a struct or array compiled and applied cvta.shared to host-supplied bytes. The exporter now refuses it the way it refuses a bare one. Pre-existing with *mut Barrier, easier to reach with SharedPtr.
• SharedArray::shared_ptr is unsafe, like as_raw_mut_ptr: the constructors establish the locality that addr_u32 relies on.

Contract

• dst: SharedPtr<u8> points into the issuing CTA's shared memory by type. barrier: *mut Barrier must point into the issuing CTA's shared memory, as for every mbarrier operation.
• A SharedPtr<T> comes from shared memory the CTA owns, or from a generic pointer through unsafe from_raw. For an address outside the shared window, addr_u32 yields an unspecified value.
• Rejected at compile time: a SharedPtr kernel parameter, bare or inside a struct or array (the host has no shared address to pass), and a barrier passed straight from map_shared_rank_mut into the _cta form (a located error that names the cluster form).
• Not rejected: SharedPtr::from_raw and pointer casts applied to a cluster pointer. Those are the caller's unsafe contract; the PTX ISA makes a .shared::cta copy into another CTA's memory undefined behavior.
• The raw layer (cuda_intrinsics::tma::cp_async_bulk_tensor_Nd_g2s_cta, ABI i1034 to i1038) keeps *mut u8 and converts a generic pointer with cvta.to.shared; its safety text states CTA locality.
• Unions and MaybeUninit<SharedPtr<T>> fail closed under the union storage guard; repr(packed) carriers follow the packed-carrier limits. On the modern libNVVM route (p3:32) a struct that mixes a SharedPtr with a narrower field is not yet supported in arrays or shared memory; see Caveats.
• .shared::cta needs PTX 8.6; the module floor is raised when the typed intrinsic is present. The typed intrinsic needs LLVM 22 or newer; the pinned sysroot llc is LLVM 23.

Implementation

• SharedSlot<T>(T) is a hidden repr(transparent) pointee marker; SharedPtr<T> { ptr: *mut SharedSlot<T> } is repr(transparent) with manual Clone and Copy. The importer types *mut SharedSlot<T> as a shared pointer in the pointee-keyed arm next to SharedArray and Barrier, and lowers SharedPtr<T> itself as a scalar newtype next to ThreadIndex; strip_shared_ptr_field_projections drops the Field(0) identity projection at the top of both place walkers.
• The methods are plain #[inline(always)] Rust: casts between *mut SharedSlot<_> stay in addrspace(3), as_ptr is the 3 to 0 cast, from_raw the 0 to 3 cast, addr_u32 composes with cvta_generic_to_shared_u32.
• Catalog: five [[tma.variant]] records with ABI entries i1034 to i1038 (appended, no renumbering), per-backend evidence profiles rust-llvm-23.1.3-aab11ee0-tma-cta and cuda-13.3-libnvvm-13.3.73-tma-cta, overlay shard schema 65 with a TMA_CTA_SHARD_SCHEMA gate; the two new shard fields are Option<String> with serde defaults and are required only when a shard lists g2s_cta variants.
• Generator: recipes without the i16 multicast mask; imported properties ImmArg, IntrConvergent, ReadOnly<arg2>, WriteOnly<arg0>; PTX floor 8.6; dst: SharedPtr<u8> rendered for the device wrapper only; the raw-layer safety text carries the CTA-locality clause for the five records.
• Importer: emit_tma_g2s_cta reuses the G2S emitter with a cta flag. Lowering: convert_g2s_cta beside the cluster lowering; LLVM NVPTX declares llvm.nvvm.cp.async.bulk.tensor.g2s.cta.tile.Nd(p3, p3, p0, i32 x N, i64, i1), libNVVM emits the .shared::cta inline PTX with ~{memory}, and an addrspace(7) operand is a located error. PTX-version detection recognizes the typed intrinsic (8.6 floor).
• Exporter: the kernel-parameter check walks struct fields and array elements and names the offending field path.
• Examples: tma_copy gains tma_copy_2d_cta_test and a generated-code check; shared_memory_helper gains the SharedPtr probe through a generic non-inlined helper, the two-field struct and returned-pointer regressions, and a code check that reads the unoptimized build; shared_ptr_api covers the remaining surface and runs through the libNVVM verifier in CI; error_kernel_shared_ptr_param is the nested-parameter negative.
• Book: the SharedPtr section with its limits, the TMA chapter entry with the figure, the API quick reference, the supported-features note and the catalog stamp.

Verification

Checked at fbe27a694 on an RTX 5090 (sm_120a) with CUDA 13.4 and nightly-2026-10-11:

• ✅ Tests of every touched crate (cuda-device, cuda-intrinsics-gen, cuda-intrinsics, cuda-oxide-codegen, mir-importer, mir-lower, dialect-nvvm, llvm-export, rustc-codegen-cuda), clippy and rustdoc with -D warnings, cargo oxide fmt --check, the 18 guards, the example lock check, cargo deny, the book with -W.
• ✅ cuda-intrinsics-gen check, probe --all --skip-terminal --per-target, check-abi-history (1033 preserved, 1038 total), a second generate with no diff, the book stamp guard.
• ✅ Evidence: the five CTA records' LLVM stages reproduce byte for byte from the pinned llc (LLVM 23.1.3) and ptxas 13.3 V13.3.73; the libNVVM file is the author's.
• ✅ GPU: tma_copy on LLVM NVPTX and on libNVVM (--emit-nvvm-ir --arch sm_120), 4096 of 4096 values for the _cta kernel next to the unchanged kernels; shared_memory_helper on both routes, in the normal and the unoptimized build; shared_ptr_api on both routes; memcheck and synccheck 0 errors; the generated-code checks pass.
• ✅ Ranks 1, 3, 4 and 5 executed at review time with real tensor maps through every SharedPtr provenance (static array, dynamic array, byte offset, from_raw, non-inlined helper): 1088 values on both backends, sanitizer clean; in tree they are covered by probes, cubin assembly and conversion tests.
• ✅ The TMA G2S lowering makes ptxas ignore setmaxnreg via an extern compatibility call #1238 kernel (384 threads, setmaxnreg_dec::<40> / setmaxnreg_inc::<232>) at sm_120a with ptxas 13.4 and nvdisasm: the _cta form, direct and from an #[inline(never)] helper taking SharedPtr<u8>, gives 168 registers, 0 stack, USETMAXREG.DEALLOC/TRY_ALLOC, UTMALDG.2D, no CALL; the unchanged cluster form gives C7506, 24 registers, CALL.ABS.NOINC.
• ✅ Negatives: a barrier from map_shared_rank_mut into _g2s_cta is a located error naming cp_async_bulk_tensor_2d_g2s; error_kernel_shared_ptr_param (Args { p: SharedPtr<f32>, n: u32 }) fails to compile with the located nested-path error; the llvm-export unit tests refuse { ptr addrspace(3), i32 } and [2 x ptr addrspace(3)] kernel parameters and keep accepting a struct of generic and global pointers.
• ✅ PTX floor: tma_copy --arch sm_90a emits .version 8.6 and ptxas sm_90a assembles all three entries; the five probe PTX files assemble under ptxas 13.4; the pinned llc selects the typed intrinsic at +ptx86 and rejects it at +ptx85.
• ✅ Baseline: the PTX of tma_copy_2d_test and tma_pipeline_test matches main's backend apart from symbol numbering; 8 unchanged shared-memory examples produce identical per-kernel PTX.

Caveats

• Part of TMA G2S lowering makes ptxas ignore setmaxnreg via an extern compatibility call #1238, not a fix: the reproducer calls cp_async_bulk_tensor_2d_g2s, and the 39 in-tree cluster-form call sites keep .shared::cluster. Migration is tracked in Typed CTA-shared pointers so TMA G2S can use .shared::cta (SM12x setmaxnreg loss, #1238) #1337.
• SharedPtr::from_raw and pointer casts can launder a map_shared_rank pointer into the typed form; the compiler rejects only an uncast addrspace(7) operand. A located importer error for that cast is a follow-up.
• Modern libNVVM (p3:32): struct { p: SharedPtr<u8>, x: u32 } in a local array or a SharedArray traps with misaligned address because mir-lower sizes every pointer at 8 bytes when it picks a struct layout; the same happens with *mut Barrier on main. NVPTX is unaffected, and so are Region { SharedPtr, usize } and arrays of pointers. Tracked separately.
• The shared_memory_helper code check reads the unoptimized build, the house mechanism for checks that need a real call boundary.
• In-tree runtime coverage of the _cta family is 2D (tma_copy); ranks 1 and 3 to 5 ran in a scratch example at review time.
• Tile mode only: no im2col, no L2 cache hint, no cta_group forms, matching the cluster family.
• An llc older than LLVM 22 emits .extern .func declarations for the typed intrinsic without a diagnostic; the pinned sysroot llc is LLVM 23.

Supersedes #1341, whose commit lands here unchanged. Part of #1337.

[lucifer1004]

Reworked per the plan in #1337. _g2s_cta now takes dst: SharedPtr<u8>, the typed CTA-shared pointer added in #1341, so locality comes from the type instead of an unsafe promise. The raw layer and ABI entries i1030–i1034 are unchanged. The branch is now #1341 plus one commit; review #1341 first. The description is updated.

[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 fbe27a6

[nihalpasham]

Thanks, this is the explicit local API I asked for in #1249, with the cluster contract left alone. This PR carries the SharedPtr commit from #1341, so both land here. I replayed both commits onto current main and layered on what the landing needed:

• Re-recorded the CTA LLVM evidence under the pinned LLVM (rust-llvm-23.1.3-aab11ee0) and regenerated; the five records sit at ABI ids i1034 to i1038 because feat(intrinsics): admit bf16x2 to FP8x2 pack conversions #1354 landed first.
• Gated the two new shard fields behind a schema bump, as feat(intrinsics): add TMA tensor reduction intrinsics #770 did for the reduction forms, so older shards still parse.
• Stated CTA locality in the raw-layer safety text, which still takes a generic *mut u8 and converts it with cvta.to.shared.
• Made SharedArray::shared_ptr unsafe, like as_raw_mut_ptr: the constructors are what make the type's locality claim true.
• The exporter now refuses a kernel parameter that carries a shared or local pointer inside a struct or array, not only a bare one. struct { p: SharedPtr<f32>, n: u32 } compiled and applied cvta.shared to host bytes. New fixture and unit tests.
• Added shared_ptr_api, the generic-helper regressions from fix(tma): select CTA-local G2S on SM120/SM121 #1249, plus the two-field struct and returned-SharedPtr shapes in shared_memory_helper, whose code check now reads the unoptimized build so the call boundary is real.
• Book: the limits paragraph, two stale comments from the pre-SharedPtr design, and a one-page figure.

On the 5090, tma_copy, shared_memory_helper and shared_ptr_api pass on both backends with memcheck and synccheck clean, and the #1238 kernel keeps setmaxnreg with the _cta form: 168 registers, UTMALDG.2D, no CALL. The cluster path is unchanged in PTX.

At fbe27a694, the full local checks and every CI check passed, apart from the CodeQL analysis, still queued behind today's GitHub Actions incident.

This is part of #1238; the reproducer path and the in-tree cluster-form callers stay on #1337.