[TIRx][CUDA] Add a table-driven PTX dialect and retire tirx.ptx.* - #20103
Open
spectrometerHBH wants to merge 17 commits into
Open
[TIRx][CUDA] Add a table-driven PTX dialect and retire tirx.ptx.*#20103spectrometerHBH wants to merge 17 commits into
spectrometerHBH wants to merge 17 commits into
Conversation
* feat(lower-tirx): ldstmatrix XOR address emission and consecutive dst registers - _swizzle_iter: add the XOR emit path (try_recognize_xor, emit_base, emit_xor_offset, xor_delta) built on the GF(2)-linearity of the swizzle family — per-iter addresses become (base + D_high) ^ sigma(D_low) with compile-time constants, no signed_strides buffer, no (C2) requirement. Also skip extent-1 iters before the zero-stride guard so m_outer==1 copies are recognized instead of falling back to full applies. - ld_stmatrix: order the num-atom (seg 4) iters by R m-stride descending so ldmatrix/stmatrix dst registers come out consecutive with the R fragment word order — ptxas no longer inserts register-shuffle MOVs. - ld_stmatrix: region-min constant split — lift compile-time region-min constants out of the swizzle apply input and emit (apply(X) + D_high) ^ sigma(D_low), so chains of copies emit textually identical applies that nvcc CSE merges into one base + one XOR per copy, matching the handwritten CUDA's shape. * refactor(lower-tirx): abolish additive signed-strides in favor of the XOR swizzle emit The additive signed_strides encoding (emit_init / emit_iter_offset / per-thread sign computations + local stride buffer) covers nothing the XOR form does not: every consumer's iter index is either compile-time (emit_xor_offset, one XOR immediate) or a TIR var (emit_xor_offset_var, per-bit XOR with compile-time sigma constants), the runtime +-1 signs only cover cases (C1) already excludes, and (C2) support-disjointness is unneeded since sigma is GF(2)-linear and inner-outer iter pairs cancel exactly. Removes the whole additive path with no compatibility aliases: - _swizzle_iter: delete emit_init/emit_iter_offset/signed_strides/C2; single recognizer (C1+distinctness); add emit_xor_offset_var (Var k, floordiv/floormod bit extraction that folds downstream). - consumers (reg, gmem_smem, vec_auto_reg, vec_auto_gmem_smem, ldgsts, ld_stmatrix): emit_base + emit_xor_offset(_var); signed_strides state removed everywhere. - tests: test_swizzle_iter rewritten for the XOR formula (incl. inner-outer-pair acceptance and mask-source-bit toggling); fast-path fingerprints in test_reg/test_gmem_smem/test_ld_stmatrix updated to the XOR form with additive-form negative guards. * feat(lower-tirx): support predicated tcgen05 barrier arrival * chore(lower-tirx): normalize swizzle notation * perf(layout): materialize swizzle fallback once * refactor(layout): migrate CUDA swizzles to ComposeLayout * perf(op-dispatch): hoist linear register copy base offsets
…#47) * chore(tvmscript): remove dead tmem_pool re-export shim The module only re-exported TMEMPool from tvm.tirx.lang.alloc_pool and had no importers; builder/__init__.py and script/__init__.py already import from the canonical location. * feat(op): add table-driven PTX dialect prototype (T.ptxd) One instruction table + one generic engine + thin generators, quarantined under the ptxd namespace (tirx.ptxd.* ops) alongside the existing T.ptx: - table.py: pure-data instruction entries (modifier slots, per-entry check function for cross-slot PTX grammar rules, per-operand space/dtype overrides for mixed-space instructions like cp.async.bulk) - engine.py: registers ops + one codegen closure per family at backend load; resolves T.ptxd attribute chains and the exact-text string form; coerces addresses at trace time (shared pointer -> explicit cvta node, raw u32 passthrough) and supports framework-level @p predication - render.py: tvm-free asm-helper rendering shared byte-identical between the codegen and the offline gen_helpers dump - generators: gen_stubs (tvm/script/tirx.pyi for Pyright completion, with a freshness unit test), gen_coverage, gen_helpers - tests: golden helper sources, coercion IR forms, parser round-trip, full-variant nvcc soundness gate, GPU ld/st roundtrip Demo table covers prefetch/ld/st/red/cvta/cp.async.bulk (103 variants). * feat(op): transcribe the full ld family into ptxd and retire ptx.ld_acquire Transcribe PTX ISA 9.7.9.8 (ld) and 9.7.9.9 (ld.global.nc) into the ptxd table: 9 modifier slots (mmio/sem/scope/ss/cop/nc/l1ev/prefetch/type) and a check() carrying the grammar rules with per-rule doc citations. 9,548 legal variants, all assembled by ptxas at sm_90. Migrate the six tirx.ptx.ld_acquire call sites to T.ptxd and delete the op from every registration point. Supporting fixes found while doing it: - The certification gate never validated anything: nvcc -ptx silently drops unreferenced __forceinline__ device functions, so a bogus instruction passed. Strip __forceinline__ and assemble with -rdc=true -cubin, and add a test that the gate rejects a known-bad instruction. - Split the suite into a fast tier (sampled ptxas smoke, ~8s) and a PTXD_CERT=1 certification tier sharded 32 ways for xdist, so the full 9.6k-variant sweep does not run on every invocation. - Widen PTX_TYPES to all 14 scalar types with an explicit asm carrier, so 8/16-bit values ride in 16-bit "h" registers and narrow on return. - Support per-operand space/dtype overrides for instructions whose operands live in different state spaces (cp.async.bulk), and accept the printed round-trip form so script() output re-parses. - Reject raw uint64 addresses: the helper parameter is const void*, so an integer address is only meaningful as a uint32 shared-window offset. - Generate the .pyi stub through ruff format and with an ASF header so it is stable under pre-commit, and allow the pyi extension in check_file_type. * test(op): enforce the one-call-one-instruction invariant mechanically The ptxd dialect's defining constraint is that a call emits exactly one native PTX instruction, with the @p wrapper and cvta coercion as the only sanctioned framework mechanisms. That was a convention; make it a test that runs over every variant of every family, so it cannot erode as the table grows to cover the ISA. The probe is falsified against the three shapes it must reject — a chained bundle, a spin loop with a label and a branch, and a prologue+instruction pair — since a guard never shown to fail proves nothing. It also rejects a cvta hidden inside any helper other than the cvta family's own. Make the certification architecture overridable via PTXD_ARCH. It was pinned to sm_90, which would misreport every legal variant of a Blackwell-only family (tcgen05, clusterlaunchcontrol) as illegal, and those verdicts would then get baked into a check() and silently delete real coverage. * refactor(op): retire 13 unreferenced legacy ptx ops These ops had no call site anywhere in the workspace, verified with five grep forms (T.ptx.<n>, Tx.ptx.<n>, flat ptx_<n>(, the "tirx.ptx.<n>" string, and the tvm_builtin_ptx_<n> helper name that tests assert on), all matched with trailing word boundaries. The boundaries matter more than usual here: 60 pairs of these op names are prefixes of one another, so ptx_sub_f32 matches inside ptx_sub_f32x2 and ptx_tcgen05_mma_sp inside ptx_tcgen05_mma_sp_block_scale. A naive grep both over-reports references and under-deletes registrations. Removed: add_f64, sub_f32, sub_f64, mul_f64, fma_f64, ld_mmio, ld_relaxed, st_mmio, st_relaxed, st_release, st_volatile, tcgen05_mma_sp, tcgen05_mma_sp_block_scale — from the C++ op table, the op.py wrappers, the PTXNamespace members (including the Tcgen05MmaSpNamespace class that existed only for the two sparse MMA ops), the codegen registrations, and one test enumeration. 111 registered ptx ops -> 98. Deliberately kept: cp_async_bulk_s2s_cluster (a test still asserts on its helper) and map_shared_rank (the op registration is dead, but the user-facing T.ptx.map_shared_rank name has 26 live call sites and moves with the mapa family). * feat(op): migrate the ex2/rcp family to ptxd and retire the legacy ops Transcribe ex2 (PTX ISA 9.7.3.21) and rcp (9.7.3.13) into the table and move all 26 call sites in tirx-kernels, then delete tirx.ptx.exp2 and tirx.ptx.rcp. 98 registered ptx ops -> 96. exp2 is renamed to its actual mnemonic: the instruction is ex2, and the legacy name never matched the ISA. The generated PTX is byte-identical, verified by compiling flash_attention4 before and after and diffing the assembled PTX with virtual registers canonically renumbered: 10383 lines each, zero differing lines, and the opcode histogram matches at 4501 instructions. Getting there required fixing two conflations in the renderer, both found by that diff rather than by review: - The effect kind was driving the asm `volatile` qualifier. They answer different questions - whether the IR may share a let-bound result, versus whether nvcc may delete or reorder the emitted asm. Marking ex2/rcp opaque (to get volatile) stopped the IR from inlining a let-bound reciprocal at its 35 use sites; marking them pure (matching the legacy registration, and correct - they are pure math) dropped volatile and let nvcc common 256 redundant reciprocals down to 2. Neither matched. Always emit volatile: the caller asked for a specific hardware instruction, so C must not delete or duplicate it, and every hand-written helper already does this. - The memory clobber was also tied to the effect kind. An instruction with no memory operand cannot clobber memory, and claiming it does is a needless optimization barrier around pure-register instructions like ex2. It is now tied to whether the instruction has an address operand, which leaves every existing family's output unchanged. Certification covers 9663 variants. bench_suite --filter flash_attention4 exits 3, but so does a re-run with no code change at all (4 rows below -1%), and the two rows flagged here move by 2.5 points between identical runs or reproduce unchanged without the migration - the pinned baseline predates many commits, so that gate is measuring drift and box noise, not this change. * feat(op): migrate the fns family to ptxd and retire the legacy op fns (PTX ISA 9.7.1.18, `fns.b32 d, mask, base, offset;`) has one form and three differently-typed operands, so each operand declares its own dtype rather than sharing the entry's type slot. 97 registered ptx ops -> 96. Add an explicit asm_volatile knob, because the previous "always volatile" rule was wrong for this family. Whether the emitted inline asm carries `volatile` is a C-level optimization barrier: it never changes which PTX instruction is emitted, only whether nvcc may common up identical calls. It is genuinely independent of the IR's effect kind, and the two disagree in opposite directions in the code being migrated - ex2/rcp are pure yet carry the barrier, fns is pure and does not. Deriving one from the other cannot reproduce both, so entries state it when it differs from the default. Measured rather than assumed: with volatile forced on, a kernel issuing four identical fns calls assembles to 4 instructions; without it, nvcc commons them to 3, which is what the legacy helper does. The migrated form now assembles to the same 18-instruction PTX as before. Verified: full /tir-test green (2641 passed, 0 failed) and bench_suite --filter mega_moe exits 0 with no row outside +/-1%. * feat(op): migrate the red/atom family to ptxd and retire the legacy ops Widen red and add atom, both from the scalar .op syntax line (PTX ISA 9.7.14.6 and 9.7.14.5): sem/scope/space all optional, the full 8-operation and 8-type sets, and atom returning through its type slot. 4200 certified variants. 97 registered ptx ops -> 95. The op x type pairing came from the oracle rather than the prose: the ISA lists the union of types across all operations, so a first pass that only excluded the obvious bit/arithmetic mismatch left 1000 combinations ptxas rejects. ptxas names the exact per-operation type set in its diagnostics, so the check now encodes those, and the variant count drops by exactly the 1000 that failed. Deliberately excluded, each needing a mechanism this shape lacks: {.level::cache_hint} with its trailing cache_policy operand, the vector forms, the half-precision .add.noftz forms, and atom's .cas and .exch (both are separate syntax shapes with their own operand lists). Emitted PTX is unchanged: a probe issuing all six forms used by mega_moe assembles to identical PTX before and after, 54 lines each with virtual registers canonically renumbered. Verified: full /tir-test green (2641 passed) and bench_suite --filter mega_moe exits 0 with no row outside +/-1%. * feat(op): migrate prefetch.tensormap to ptxd and retire the legacy op Widen the prefetch entry to three of the four syntax lines in PTX ISA 9.7.9.16 - {.space}.level, .global.level::eviction_priority, and {.tensormap_space}.tensormap - with a check that exactly one target qualifier is present and that .const/.param only pair with .tensormap. 11 variants, all certified. 94 registered ptx ops -> 93. (prefetchu is a different mnemonic and would be its own entry.) Accept a u64 address handle on address operands. A tensormap address arrives as T.address_of(...), which is a 64-bit value rather than a typed pointer - the legacy helper declared `unsigned long long` for exactly this reason, and PTX binds either to the same "l" register. The conversion is now an explicit reinterpret node in the IR rather than a type pun inside the helper, matching how cvta coercion is handled. This also fixes a latent bug in the legacy op: passing it an actual pointer produced a C type error. Emitted PTX is unchanged for both affected kernels: 1855 and 1880 lines, zero differing lines with virtual registers canonically renumbered. Found by /tir-test, not by review: the first pass missed the call site in tile_primitive/copy_async/tma.py because the search output was truncated, which broke 61 copy_async dispatch tests. The census greps must not be piped through head. Verified: full /tir-test green (2642 passed) and bench_suite --filter mega_moe exits 0 with no row outside +/-1%. The paged_mqa workloads cannot bench in this environment - a reference impl needs sglang, which is not installed - so they rest on the PTX equivalence and their correctness tests. * feat(op): migrate the max family to ptxd and retire the legacy op max per PTX ISA 9.7.3.12, two-source form: {.ftz}{.NaN} on .f32 and the bare .f64 line, 5 certified variants. 92 registered ptx ops -> 91. Deliberately excluded: {.xorsign.abs} (a paired qualifier the slot model cannot express as one token), the half-precision forms of 9.7.4.8, and the three-source `max{.ftz}{.NaN}{.abs}.f32 d, a, b, c` line, which is a different operand shape and so needs its own entry. That last exclusion corrects an entry in the inventory: reduce3_max_f32 and reduce3_min_f32 were listed as having no PTX equivalent, but the ISA does define the three-source form. They are 1:1 after all and can migrate once there is an entry for that shape, rather than moving to T.cuda as planned. Emitted asm is byte-identical to the legacy helper, including its volatile barrier on an otherwise pure instruction. Verified: full /tir-test green (2643 passed). No kernel uses max_f32 - its only call site is a tirx-base test, which is outside /tir-test's scope and was run explicitly. * feat(op): migrate st.bulk to ptxd and retire the legacy op st.bulk per PTX ISA 9.7.9.14, 4 certified variants. 92 registered ptx ops -> 91. Three small mechanisms, each needed by the instruction's own syntax: - An "imm" operand role for a value the ISA fixes: st.bulk's initval "must be zero", so it belongs in the instruction text and takes neither a C parameter nor a call argument. - A `mnemonic` field, because the table key has to be a Python identifier while the PTX name contains a dot. st.bulk is the first instruction where the two cannot be the same string. - A `min_arch` field. st.bulk requires sm_100 and certification defaults to sm_90, so it was reported as illegal at every variant. Encoding that into a check() would have been the wrong fix - it would delete real coverage - so certification now groups variants by their family's arch floor and assembles each group at that arch. This is the failure mode PTXD_ARCH was added to anticipate, hit for real. Also drop the storage-scope gate on shared-space address operands. It rejected correct code: a shared buffer's ptr_to() reports scope 'global', so the check is stricter than the metadata can support, and the legacy helpers converted unconditionally anyway. Verified: full /tir-test green (2642 passed) and bench_suite --filter mega_moe exits 0 with no row outside +/-1%. * refactor(op): give ptxd a uniform allocate-then-pass surface PTX has no defining form: a register is declared first and instructions then write into it. ptxd modelled this two ways instead -- some families returned a value, others were void statements -- so the call shape depended on which family you reached for, and the returned form had a second problem: a kPure result bound to T.let gets substituted at every use by the analyzer, turning one source-level call into N instructions (flash_attention4 non-causal emits 258 rcp.approx.ftz.f32 where the author wrote one). Model destinations as ordinary operands instead. Every entry lists its destination as a leading OperandSlot(role="dst") in PTX operand order, every helper is void with the destination taken by reference, and every call is a statement: val = T.local_scalar("uint32") # .reg .b32 val; T.ptxd.ld.acquire.gpu.global_.b32(val, p) # ld.acquire.gpu.global.b32 val, [p]; This collapses `returns` and `effect` out of the schema. `returns` was a second encoding of "this instruction has a destination"; `effect` is now constant kOpaque, which a void call needs to survive RemoveNoOp and which also states the honest contract for a hand-written PTX call. asm_volatile stops deriving from effect and becomes a plain per-entry bool that keeps each instruction's established barrier. The emitted PTX is unchanged: across all 13,871 variants the asm instruction text and volatile flag are byte-identical to before, and the five families without destinations render byte-identically including their @p twins. Only the C wrapper shape moved, which __forceinline__ erases at -O3. Also: - @p is now gated on has_dst rather than on the return convention. It was already unavailable for these families; the new gate states why (a false predicate leaves the destination unwritten, and "=" tells nvcc its prior value is dead). - 8-bit destinations narrow through a carrier local, mirroring the old narrowing return, so the 2,046 8-bit ld variants keep compiling. - dst operands get a real coercion: they must be a writable lvalue of the right dtype, which rejects the T.let binding that caused the 258x expansion above. - fix an arity bug in the round-trip path: the trailing-pred check counted ISA-fixed `imm` operands, so a predicated st.bulk could not re-parse. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(op): register the floating-point add/sub/mul/fma family in ptxd Transcribed from a live fetch of PTX ISA 9.7.3.{2,3,4,7} (same-precision) and 9.7.5.{1,2,3} (mixed-precision), covering every syntax line: add{.rnd}{.ftz}{.sat}.f32 d, a, b; add{.rnd}{.ftz}.f32x2 d, a, b; add{.rnd}.f64 d, a, b; add{.rnd}{.sat}.f32.atype d, a, c; fma.rnd{.ftz}{.sat}.f32 d, a, b, c; fma.rnd{.sat}.f32.abtype d, a, b, c; 189 new variants, all certified through ptxas with zero errors. `mul` has no mixed-precision line, and fma's .rnd is mandatory on every line, so both fall out of the table rather than out of a special case. Two mechanisms this needs, both native syntax: - `.f32x2` and the 16-bit mixed sources join PTX_TYPES. Per the ISA, .f32x2 operands "have .b64 type" and mixed sources sit in a `.reg .b16`, so both ride bit-container carriers rather than float ones. - An operand may name a modifier slot for its type, and an omitted optional slot falls back to `type`. That is exactly what the mixed lines need: `add.rn.f32.bf16`'s `a` is the .bf16 source while plain `add.rn.f32`'s `a` is just .f32 — one rule, no per-form branching. `operand_type()` replaces the `slot.dtype or mod_map["type"]` that render and engine each open-coded. min_arch is sm_100a for these entries: .f32/.f64 assemble everywhere but .f32x2 and the mixed lines need sm_100, and certification has to run somewhere every legal variant is legal. Also pin _LD_TYPES to its 14 scalar tokens instead of deriving it from PTX_TYPES, which now also carries the packed and mixed tokens that are not ld types. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(op): retire the scalar legacy T.ptx arithmetic ops add_f32, mul_f32, fma_f32 and add_rn_f32_bf16 are now covered by the table-driven entries, so all four go away at every registration point: the PTXNamespace bindings, the op.py wrappers, the codegen registrations in math.py, and the kOpaque entries in the C++ kDeviceIntrinsics table. Equivalence, measured per kernel rather than assumed: for the six affected kernels the set of emitted asm instruction texts is unchanged and ptxas allocates the same number of registers; only the helper name moves from tvm_builtin_ptx_* to tvm_builtin_ptxd_*. flash_attention4 is byte-identical. The one real instruction change is add_rn_f32_bf16, whose legacy helper tied its accumulator with "+f" and emitted `add.rn.f32.bf16 %0, %1, %0`. PTX has a distinct `c` operand, so ptxd emits the honest three-address `%0, %1, %2` and the call site passes the accumulator twice. Measured on a B200: identical register allocation (mega_moe: 128 both ways) and bench_suite -0.0% -- ptxas coalesces the tie completely, so the tie was only ever a PTX-level artifact. The scalar registration loop also carried five orphans -- ptx_sub_f32, ptx_add_f64, ptx_sub_f64, ptx_mul_f64, ptx_fma_f64 -- whose C++ ops were deleted earlier but whose codegen registrations survived, unreachable because register_codegen does not check that the op exists. They go with the loop. The .f32x2 ops stay for now. Their destination is a b64 view over two adjacent float32 buffer elements, which the ptxd destination coercion cannot express: it requires a writable lvalue whose dtype matches the operand. Migrating them needs that mechanism, not a use-site workaround. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(op): add register groups and register the mov pack/unpack family PTX writes an instruction's operand *shape* in the operand list, not in the dotted modifier text: `mov.b64 d, {lo, hi}` and `mov.b64 {lo, hi}, a` are the same opcode. ptxd assumed shape was always determined by the mnemonic plus modifier tokens -- true for every family so far only because they all spell their group size in the text (`ld.v4`, `ldmatrix.x2`, `mma.m16n8k16`). `mov` is the first family that does not, so it needs two things. Register groups: `OperandSlot.lanes` makes an operand a brace-enclosed vector expression. It renders as `{%k, %k+1, ...}` and takes one call argument per lane; the C parameter list stays flat, since inline asm puts the braces in the asm text and binds each lane as an ordinary operand (a grouped destination is simply N `"="` outputs). `call_slots()` is now the single definition of the flat argument layout, replacing the three places that each open-coded it. Shape-directed dispatch: a chain carries candidate entries instead of one, each modifier token narrows them, and the call's argument count and operand dtypes pick the survivor -- the same information ptxas resolves them by. So the spelling stays the PTX text: T.ptxd.mov.b64(packed, lo, hi) # mov.b64 packed, {lo, hi}; T.ptxd.mov.b64(lo, hi, packed) # mov.b64 {lo, hi}, packed; 16 mov entries share mnemonic "mov", covering all five reachable (type, lanes, lane type) forms; all certify through ptxas at sm_90 and sm_100a. The helper name now keys off the table name rather than the mnemonic so those entries do not collide -- byte-identical for every existing family (st.bulk already normalized to st_bulk), verified as 0 diffs across all 15,735 existing renders. Two of the ISA's seven vector forms are deliberately absent, and the table says so: `mov.b16 {a,b}` and `mov.b32 {a,b,c,d}` have 8-bit lanes, and inline asm has no 8-bit register constraint, so the lanes ride a 16-bit carrier and ptxas rejects the widths. Both are legal in hand-written PTX and unreachable via make_uchar4 as well. This is the first recorded case of legal PTX that CUDA C cannot express in one instruction, which is the real bound on this dialect's coverage. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(op): retire the legacy packed f32x2 arithmetic ops add/sub/mul/fma_f32x2 were the last four legacy T.ptx arithmetic ops. They survived the earlier batch because their destination is eight bytes written through a `void*` into a float array -- an operand shape the ptxd destination coercion cannot express, since it requires a writable lvalue whose dtype matches. With `mov` registered that shape is no longer needed: the packing that the legacy helper hid becomes two visible instructions. T.ptxd.mov.b64(acc, local_sum[2 * j], local_sum[2 * j + 1]) T.ptxd.add.rn.ftz.f32x2(acc, acc, rhs) T.ptxd.mov.b64(local_sum[2 * j], local_sum[2 * j + 1], acc) Those movs are not new work. nvcc already emitted exactly them for the old make_float2/float2_x glue -- 12 mov.b64 for four packed adds in the reduction primitive -- and ptxas folds them away entirely, because a 64-bit register is an aligned pair of 32-bit registers and the pack is virtual bookkeeping. Measured: identical ptxas register counts on all nine affected kernel configurations, and asm instruction texts gain only the two mov forms with none lost. bench_suite is in the noise on flash_attention4 (32 configs) and sparse_flashmla_prefill_head64_phase1. The VecImpl contract gains one thing: `emit` may return None after emitting its own statements. A packed op bracketed by pack/unpack is several statements and cannot be a single expression, which is what the old `-> Expr` wrapped in one T.evaluate assumed. Two fixes fall out: - `_coerce_operand` unwraps the T.local_scalar wrapper for every operand role, not just destinations. A scalar used as a *source* previously reached the FFI as an opaque object. - TScriptPrinterName is now the surface path (the mnemonic) rather than the table key. They agreed for every single-shape family, but the `mov_*` entries all answer to `T.ptxd.mov`, and test_op_namespace_cleanup rightly requires a printer name to be a path that actually resolves from T. The round-trip test now covers mov so the shared-mnemonic dispatch stays honest through print/reparse. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(op): name ld's state-space slot "space" so shared addressing engages An `addr` operand picks its C carrier from the entry's `"space"` modifier slot: a shared state space takes a 32-bit shared-window address bound with "r", anything else takes a generic pointer bound with "l". `ld` named that slot `"ss"`, so the lookup missed and all 882 `ld.shared*` variants rendered a 64-bit generic pointer instead: ld.shared::cta.b32 (uint32_t& __d, const void* __addr) "l" <- wrong st.shared::cta.b32 (uint32_t __addr, uint32_t __value) "r" <- right Nothing caught it. The address coercion never took the shared path either, so a shared pointer was rejected as "generic address but got a shared-scope pointer" while a raw uint32 window address was rejected outright -- and ptxas assembles the wrong form happily, since PTX permits a 64-bit address for a 32-bit state space and simply truncates it (ISA 6.4.1). The result would have been a truncated *generic* address used as a shared offset. No in-tree caller uses ld.shared, so this is a landmine rather than a live bug. The trap generalizes though, which is why the slot now carries a comment: name the space slot anything else in a future family and shared addressing silently stops engaging, with no diagnostic anywhere. Variant count is unchanged; all 14,076 re-certify through ptxas, the 882 renamed helpers now with a 32-bit "r" address. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(op): bring st to parity with ld (84 -> 2912 variants) `st` was the largest conformance gap in the table -- 84 of a modelled 2912 forms, with no comment recording what was left out or why. Six axes were missing or wrong against PTX ISA 9.7.9.11: .scope was ("cta","gpu","sys"); ISA: "{ .cta, .cluster, .gpu, .sys }" .ss was mandatory and 2-token; ISA has 5 spaces plus "If no state space is given, perform the store using Generic Addressing" .cop absent; ISA: "{ .wb, .cg, .cs, .wt }" .level1::eviction_priority absent (syntax line 2) st.mmio.sem.sys{.global} absent (syntax line 6) .type 6 of the ISA's 15 _check_st is now the mirror of _check_ld, each rule carrying the ISA prose it comes from, and one rule is deleted outright: "st.shared::cta is CTA-local; scope must be cta or omitted" has no basis in the ISA -- the only state-space statement for these forms says nothing about scope, and ptxas accepts st.release.sys.shared::cta.b32 and st.release.cluster.shared::cta.b32. We were denying legal PTX. The entry now also records what stays out and why, the way `ld` does: .vec/.b128 multi-register sources, .level::cache_hint with its trailing cache_policy operand, .level2::eviction_priority (vector-only), and .param::func (kernel-parameter addresses cannot cross the helper ABI). All 16,904 variants certify through ptxas, the 2,828 new ones included. _LDST_TYPES and _SCOPES go away -- st was their only consumer. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * refactor(op): consolidate the ptxd machinery before the dtype axis lands Today's churn -- destinations became operands, register groups, entries sharing a mnemonic -- left three copies of the operand emitter and several rules written down twice. The next change touches exactly this code, so collapse it first. Emitted PTX is unchanged: all 21,407 renders (every variant of every entry, with the @p twin where one exists) are byte-identical to HEAD. The consolidations, in the order they matter: - render.py's operand loop was one lanes-generic body wearing three copies: a grouped-dst branch, a grouped-value branch, and the original scalar chain. The grouped-dst copy had also dropped the carrier, so an 8-bit lane would have emitted "=h" bound to a uint8_t&. One loop now, -21 lines. - operand_space() joins operand_type() in table.py. It was the same line in engine.py and render.py, and it is the line that made `ld` naming its slot "ss" a silent wrong-address bug; with one definition that failure mode cannot come back unnoticed. - call_slots() moves to table.py (it never needed tvm) and gen_stubs uses it instead of re-deriving the layout without the imm filter. That was a real bug: the checked-in stub declared st_bulk(addr, size, initval) while the runtime rejects a third argument. Plus the smaller ones: `_fill` returns None so `_narrow` owns the single error message and `_accepts` (a second full scan of the same slots) goes away; the unread `_InstrChain._entry`; error messages stop being reassembled with `split(": ", 1)[1]`, which produced "mov_pack_b32x2: mov_pack_b32x2: ..." on the aggregate path; `except (ValueError, TypeError)` narrows to ValueError so a table typo surfaces instead of being reported as a bad call; `entry.family` and `pred_forms()` replace idioms open-coded at five sites; _EFFECT_OPAQUE reads CallEffectKind.Opaque.value rather than hand-copying 3. `min_arch` becomes `cert_arch`. Its docstring said "lowest -arch that assembles this family", but for add/sub/mul/fma it is the *maximum* floor over the family's variants -- the file already admitted the contradiction in a comment, and the misreading produced a wrong recommendation in the audit. One CWD-dependent test failure fixed: gen_stubs passed a bare relative --stdin-filename to ruff, so running pytest from the workspace root resolved no pyproject.toml and fell back to ruff's 88-column default, failing test_ptxd_stub_up_to_date. It now passes the absolute STUB_PATH, and the test reads that same constant. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(op): certify at sm_100, and write down what each entry leaves out Three ISA-conformance items from the audit, none of them changing a variant count. cert_arch drops the arch-specific suffix on the three entries that carried it. The ISA Target Notes say "requires sm_100 or higher" with no `a`, and add.rn.f32x2, add.rn.sat.f32.bf16 and st.bulk.weak.shared::cta all assemble at sm_100 and sm_120 under nvcc 13.2. Certifying at sm_100a was a strictly weaker gate: it would not have caught a variant that stopped being portable. 193 variants now certify at sm_100. _check_rcp told callers "rcp.approx is only defined for .f32", which is a statement about this entry dressed up as a statement about PTX -- rcp.approx.ftz.f64 is ISA 9.7.3.14, an existing syntax line we simply have not registered. The message now says so. The rest is documentation, held to the table's own convention that an absence must be recorded with its reason. cvta and cp had no comment at all; cvta registers 1 of 32 legal combinations and cp 1 of 8 syntax lines. Also recorded: st_bulk's 32-bit size form (an operand-shape axis, like mov's) and its two value constraints that check() structurally cannot see; mov's sink symbol and scalar mov; the integer/extended-precision/half-precision lines absent from the arithmetic group; the integer max family; and why prefetch keeps .level::eviction_priority bound to .global (its syntax line writes .global in, where ld's writes {.ss}) -- an audit flagged that as a gap once already. Corrected rather than added: ld's .b128 reason (it has a "q" carrier now; what it needs is __int128 in the host compiler), ex2's stale carrier reason, _check_atomic's provenance (ISA Tables 35/36 give the per-op pairing; the Syntax block's .type line is only the union, which is why it cannot be transcribed directly), and every ISA section number on the arithmetic entries, which pointed at copysign and mad. ld and atom now note that the ISA permits @p on them while ptxd does not, with the mechanism reason on InstructionEntry.has_dst. st does not get that note -- it has no destination and @p st works today. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(lower-tirx): accept any same-width dtype on a .bN ptxd operand PTX ISA 5.2: "The bit-size type is compatible with any fundamental type having the same size." A `.b32` operand therefore accepts float32, int32 or uint32 -- all three name the same 32 bits, and PTX draws no distinction between them. The generated helper is C, which does draw one. Routing a float through the single canonical `uint32_t` parameter is a numeric conversion, and nvcc emits `cvt.rzi.u32.f32` for it: the wrong bits, silently. So each dtype gets its own helper, and an instruction with several bit-typed operands gets their product. Mechanics: - `BIT_DTYPES` maps each bit type to the dtypes of its width; `DTYPE_C` gives every dtype its C parameter type, asm constraint and, where the value cannot bind a register class directly (8-bit, __half, __nv_bfloat16), the carrier and the bit-cast on each boundary. These are C-expression-level casts: the asm block still holds exactly one instruction. - The helper name gains a positional dtype suffix, but only when the choice is non-canonical -- so no existing helper is renamed except the six f32x2 mov shapes, which the axis subsumes. Naming only the changed operands would collide the moment two operands swap which is non-canonical (`atom`'s d and b do exactly that). - Dropping those six entries (mov 16 -> 10) is required, not incidental: with the axis they accept the same calls as the b32x2/b64x2/b32x4 shapes, and shared-mnemonic dispatch would be ambiguous. - `_coerce_operand` widens from exact-dtype to any dtype of the width; the codegen closure derives the actual dtypes from the call arguments. 31,024 variants (up from 16,904) assemble under ptxas, the certification tiers walk the dtype axis, and the nine kernels using ptxd today have byte-identical helper bodies and call sites -- only the two f32x2 mov helpers change name. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * test(lower-tirx): cover register groups, and reject two silent conversions The register-group and dtype-axis mechanisms had no tests of their own, and writing them turned up two ways a call could reach the wrong C parameter type -- the same silent-conversion failure the dtype axis exists to prevent: - Lanes of one group disagreeing on dtype. Each lane alone is a legal dtype for the operand's bit type, so per-lane coercion accepted it, and the group was then typed from its first lane -- binding the odd lane to a parameter of another type. A register group is ONE operand spanning N registers (ISA 6.4.3: "similarly typed"), so it has one dtype; `_check_lane_dtypes` now says so at trace time. - A bare Python float literal. It carries no dtype, so `1.5` on a `.b32` operand became `T.uint32(1)`: neither the float's bits nor a diagnostic. Integer literals stay accepted (their bits are their bits); float ones now ask for an explicit constant. Tests: golden helpers for mov pack/unpack, the b128 "q" carrier and a dtype-suffixed variant; an end-to-end assertion on the brace text `mov.b64 %0, {%1, %2};`, which nothing checked before; and negatives for wrong arity, a non-lvalue lane, mismatched lanes and the bare literal. `tokens_for(entry, **by_name)` replaces the positional token tuples in the goldens -- `ld` took nine, so inserting a slot would have shifted them all silently. It caught a wrong slot name while being written. All nine ptxd kernels are byte-identical. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * style(test): declare the f32x2 scalars with T.local_scalar The bare `prod: T.uint64` annotation form does not bind a name in plain Python, so ruff reported F821 on every use. Every other ptxd test already declares its registers with T.local_scalar; this one now matches. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * refactor(op): give the ptxd dtype axis one owner per layer Four review passes over the dtype-axis work agreed on the same shape: the facts it introduced were spread across layers, and its enumeration was copied into every consumer. Nothing here changes what is generated -- all 38,091 renders are byte-identical and the nine ptxd kernels recompile unchanged. Type tables, one per layer. `PTX_TYPES` had four columns of which three were reproduced exactly by `DTYPE_C`, and after the axis landed the only surviving read was column 0. It collapses into `PTX_TYPE_DTYPES` (PTX token -> the dtypes it accepts), absorbing `BIT_DTYPES`, so `table.py` states the ISA 5.2 rule and nothing about C. The C facts move to `render.py` as one `CBinding` NamedTuple per dtype, folding in `DTYPE_SUFFIX`: three dicts over one key set, each needing a matching edit, became one row. Coercion per operand, not per argument. `_check_lane_dtypes` was a repair pass for information the per-argument walk destroyed -- its own docstring said as much. `_coerce_operand` now receives a whole operand's lanes, so choosing the dtype once *is* the lane-agreement rule and the separate check is gone. A literal lane now takes its dtype from the group rather than being typed alone, and the float-literal complaint keys off "does this operand have a dtype axis" instead of probing whether the canonical member happens to be a float. One enumerator. `variants x dtype_combos x pred_forms` was written in four tests and, crucially, *not* in `gen_helpers` -- the tool whose job is to dump what each variant compiles to had been emitting only canonical-dtype helpers since the axis landed. `renderings(entry)` defines the product once; the next axis lands in one place. Also: `call_slots`/`typed_operands` become cached properties (they ran twice per render and twice per dispatch candidate), `PTXDNamespace` groups by family once instead of rescanning the table on every attribute access, `canonical_dtypes` replaces indexing element 0 of a cartesian product, and `tokens_for` now runs the entry's `check` so a test cannot pin a variant the dialect refuses to emit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(op): register fence, bar/barrier, mapa, min/max and clc in ptxd Transcribes six instruction families from the PTX ISA (fetched from the live doc, ISA 9.3) and retires the fourteen legacy ops they cover. A new `orders_memory` flag carries the `"memory"` clobber for instructions that name no address yet constrain everyone else's memory order. The clobber was derived purely from "does this entry have an `addr` operand", which is right for instructions that touch memory but blind to fences: `asm volatile` alone pins the asm block without stopping the compiler moving ordinary loads and stores across it. Every legacy barrier helper carried `::: "memory"` with no address operand, so without this the whole family would have silently lost its barrier. Default false, and all 38,078 existing renders are byte-identical. Registered, each syntax line transcribed and certified under ptxas at sm_90 / sm_100 / sm_100a: - fence (5 lines, incl. the tensormap acquire form whose size operand the ISA fixes at the literal 128) and griddepcontrol - bar / barrier / barrier.cluster - mapa, split by `.type` because it fixes the width of both d and a -- a `.u32` mapa takes a 32-bit shared-window address, not a pointer. The legacy helper bound a 64-bit pointer for both and only ever used `.u64`, so ptxas never saw the mismatch. - min and max, complete: integer, single/double and half-precision lines plus the three-source `d, a, b, c` form, which retires reduce3_{max,min}_f32 - clusterlaunchcontrol.try_cancel Not registered, both for mechanism reasons rather than ISA ones, and noted where the reader will look for them: the no-count `bar.sync a` / `barrier.sync a` lines, whose arity collides with the framework's positional `pred` form, and `min`/`max` `.u8x4`/`.s8x4`, which need sm_120f. Call sites pass the count explicitly now: the ISA types bar/barrier's operands `.u32`, and a concrete PTX type accepts exactly its dtype, so the int32 expressions the legacy C signature swallowed are now written as `T.uint32(...)`. Also fixes seven dead call sites in rmsnorm that named `T.ptx.fence.proxy`, which FenceNamespace never had -- they raised AttributeError at trace time and no test covered them. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(op): register the tcgen05 alloc and sync instructions in ptxd Transcribes PTX ISA 9.7.16's allocation and synchronisation lines and retires the seven legacy ops they cover: alloc, dealloc, relinquish_alloc_permit, wait::{ld,st} and fence::{before,after} _thread_sync. All 32 renders assemble under ptxas at sm_100a. Every entry sets `orders_memory=True`: the legacy helpers all carried `::: "memory"`, and the waits and fences name no address at all. `tcgen05.dealloc`'s taddr is a plain register operand in the ISA, not a bracketed address, so it is a `value` slot -- which is also why it needs no tmem address space while `shift` does. Not registered, and noted where a reader will look for them: - `tcgen05.shift.cta_group.down [taddr]`. The operand is bracketed, so it needs the `addr` role, but `addr` only picks the 32-bit carrier when the state space starts with "shared". A tmem address is neither shared nor generic, and labelling it shared to get the right carrier would put a falsehood in the table. - `tcgen05.commit`'s `{.multicast}{, ctaMask}` form, the same collision `bar.sync` hit: an optional trailing operand makes a one-operand and a two-operand entry both accept a two-operand call, because the framework's positional `pred` consumes the difference. Call sites do pass `cta_mask=3`, so `tcgen05.commit` stays on the legacy op. /tir-test: 2646 passed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(op): move mapa onto ptxd and retire the legacy op mapa returned a value; ptxd writes into a destination the caller declares, as PTX itself does. Call sites therefore gain a one-element local scratch buffer, whose element is both a writable lvalue and an ordinary Expr -- `T.local_scalar` is not, it hands back a wrapper that `T.reinterpret` rejects. Two things this shook out: - `MBarrier.remote_view` is a plain Python method, not TVMScript, so a bare ptxd call there is discarded rather than emitted. The blackwell IR test caught it (zero mapa calls in the body); it needs an explicit `T.evaluate`. The dsmem emitter sits inside a `@T.prim_func` and does not. - The two gemm_comm kernels give `Tx` opposite meanings -- `gemm_reduce_scatter` imports the whole tirx namespace as `Tx`, while `allgather_gemm` imports only the tile submodule -- so the same helper cannot be shared verbatim between them. The blackwell assertions move with the shape: `arrive` now reads the mapped address back out of mapa's destination instead of re-deriving it, so the two are no longer structurally equal, and mapa's rank carries the `.u32` cast the ISA asks for. /tir-test: 2646 passed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(op): register the cp.async completion-tracking instructions Adds `cp.async.mbarrier.arrive{.noinc}{.shared{::cta}}.b64` and `cp.async.bulk.commit_group` (PTX ISA 9.7.9.13 / 9.7.9.15) and retires three legacy ops -- the third being `ptx_cp_async_mbarrier_arrive_noinc`, a wrapper that only forwarded to the one being deleted. Both entries take `mnemonic="cp"`: the surface splits a family name at its first dot, so the remaining tokens have to be modifier slots. The existing `cp` entry already models `async` and `bulk` that way. Watch the default state space when reading the diff: the legacy plain wrapper defaulted to `.shared` while the `_noinc` one defaulted to `.shared::cta`, so the two call shapes migrate to different chains. Not registered: `cp.async.bulk.wait_group N`, whose count lands in the instruction text rather than a register, and the `cp.async` ca/cg lines, which carry an optional ignore-src operand. /tir-test: 2646 passed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(op): retire the reduce3 min/max ops Their single call site moved to ptxd's three-source `max`/`min` line in an earlier commit; these registrations were left behind. Deleting them needed an AST-span removal rather than "cut to the next def": the two wrappers have the `_PTX_CVT_*` tables sitting between them and the following function, and a def-to-def cut swallows those. The guard that compares top-level names before and after caught it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(op): register the mbarrier family in ptxd Transcribes PTX ISA 9.7.14.5-9.7.14.12: init, inval, arrive, arrive.expect_tx, expect_tx and complete_tx. All 140 renders assemble under ptxas at sm_90. Two shapes the ISA offers and this cannot take: - The `state, [addr]` arrive forms. `_, [addr]` is the sink spelling, so `_` is an ISA-fixed immediate the way `st.bulk`'s initval is, and the instruction is left without a destination -- which is what lets `pred=` work at all. - The no-count arrive line, for the collision `bar.sync` already hit: a one-operand and a two-operand entry cannot share a mnemonic while the framework's positional `pred` can swallow the count. The ISA defines the omitted count as 1, so call sites write it. The addr slot deliberately fixes no state space. With `.space` omitted the ISA means a *generic* address, 64-bit on sm_90+, so pinning the operand to shared binds a 32-bit register there and ptxas rejects the whole file with "32-Bit ABI is not supported on sm_90 or higher". 36 of the 140 variants failed that way before the fix. Letting `operand_space` read the modifier picks the carrier per variant, as ld/st do. /tir-test: 2646 passed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(op): register wgmma.fence and wgmma.commit_group, retire the ops The last two zero-operand lines from PTX ISA 9.7.15.4. Both set `orders_memory=True`, matching the `"memory"` clobber their legacy helpers carried. `wgmma.wait_group.sync.aligned N` stays behind: the group count lands in the instruction text, not a register. So do the `wgmma.mma_async` lines, whose accumulator is a register group up to 128 wide. /tir-test: 2645 passed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(lower-tirx): add explicit pred marker and caller-chosen immediates to ptxd Two mechanisms, both needed to transcribe instruction families whose ISA syntax lines the table could not express before. Explicit pred marker. The Call layout was [operands..., pred?] [slot tokens], so a predicate could only be recovered by counting arguments. Where two syntax lines differ by one optional trailing operand, that is ambiguous: `mbarrier.arrive [addr]` with pred= and `mbarrier.arrive [addr], count` without one have identical argument counts, so printing and reparsing silently turned one into the other. The layout now ends with a "pred"/"" marker and the positional fallback is gone. Optional trailing operands can therefore be two entries sharing a mnemonic, told apart by arity -- which is how bar.sync/barrier.sync gain their no-count lines, mbarrier.arrive its implicit-count line, and tcgen05.commit its multicast form. Caller-chosen immediates. Some operands have no register form in the ISA at all: the N in `cp.async.wait_group N` and setmaxnreg's register budget only exist as integer literals in the instruction text. OperandSlot.choices names the closed set of legal values; the value is validated at trace time, travels as an IntImm, is not forwarded to the helper, and is baked into both the asm text and the helper name -- one helper per value. The closed set is what makes every generated helper certifiable: certification walks the product of choices through ptxas. Deleted 6 legacy ops: tcgen05_commit, cp_async_commit_group, cp_async_wait_group, cp_async_bulk_wait_group, setmaxnreg, wgmma_wait_group. The software-pipeline attribute lowering in codegen_cuda.cc now builds ptxd calls directly. One behavior note: a ptxd immediate must be constant at trace time, where the legacy wrappers deferred to codegen (after unrolling). The one site that relied on that -- flash_attention4's wait_group(1 - i_q) over an unrolled loop -- is now written out per stage. * refactor(lower-tirx): move the remote mbarrier arrive path to ptxd The legacy wrappers fused mapa and arrive into one helper so the mapped address never left the asm block. Split into two ptxd calls, the emitted instructions are the same: cvta, mapa.shared::cluster.u32, then the arrive on that window address. The address operand already accepts a raw uint32 shared-window address, so nothing has to round-trip back through a pointer -- mapa.u64 would have cost a 64-bit register to carry one. MBarrier.remote_view now maps once instead of once per arrive: the view's buffer is already the remote address, so arrives on it use it directly rather than re-deriving it the way the fused helper did. Also let ptxd's pred take a bare Python bool. It names no dtype but is unambiguous, and the alternative was every call site spelling T.bool(True). The blocker recorded for this path -- "T.alloc_local does not bind inside a class method" -- was wrong. remote_view had been doing exactly that all along; the requirement is only that a plain-Python method hand its calls to the frame with T.evaluate(), or they are discarded. Deletes ptx_mbarrier_arrive and ptx_mbarrier_arrive_expect_tx. * refactor(lower-tirx): drop the mbarrier init and complete_tx legacy ops Both had ptxd entries already; only tvm's own tests still named the legacy ops. The remote complete_tx forms become a mapa plus a complete_tx on the mapped window address, matching what the fused wrapper emitted inside one asm block, so the pinned instruction texts are unchanged. No kernel path changes -- every remaining caller was a test. * feat(lower-tirx): let a register group's lanes be a function of the modifiers Some instructions state their operand's register count as a formula over the modifiers -- ldmatrix's destination is "1, 2, or 4 32-bit registers as per the value of .num". OperandSlot.lanes now also accepts a callable over the modifier map, following the schema's existing precedent (`check`). The modifier set is closed, so every resulting length is still enumerable and certifiable. The per-entry cached call layout becomes operand_layout(entry, mod_map), memoized per token combination; trace-time resolves modifiers before arity (the attribute chain parses before the call), and codegen parses the token tail (whose length is static) before slicing operands. All 38,752 existing renders are byte-identical. First family on the mechanism: ldmatrix, transcribed from ISA 9.7.15.5.15 as three entries and one four-line lanes function, 162 variants certified (sm_90 + sm_100a). Certification caught two transcription errors on the spot -- .m16n16 doubles the register count per matrix, and a vector operand keeps its braces even at length one -- and ptxas ruled that .m16n16 requires .trans, all now encoded in the table. Call sites move to the ptxd spelling (tile_primitive ld_stmatrix and the numeric codegen test); destinations land through a uint32 view of the b16 fragment buffer, two elements per word, the tcgen05_ldst pattern. The legacy op stays: lower_warp_memory.cc recognizes tirx.ptx.ldmatrix calls to rewrite warp-buffer indices on the s_tir tensorize path, so deleting it rides with the C++ pass rewrite batch. * feat(lower-tirx): transcribe stmatrix, the store mirror of ldmatrix One ISA syntax line (9.7.15.5.16), two entries: .m8n8 pairs with .b16, and .m16n8 exists only as .b8 with .trans mandatory on sm_100a. The register group is `.num` registers flat -- no shape term, unlike ldmatrix's .m16n16 doubling -- so both entries share a four-line lanes function. 162 variants certified across sm_90 and sm_100a. Operand order is the reverse of ldmatrix's: the address first, the register group second. Call sites move to the ptxd spelling. The b16 fragment buffers reach the b32 registers through a uint32 view, two elements per word, which is what the ld side already does since the ldmatrix migration -- both directions now share one view and one index expression. One trap worth recording: a chain string bound to a name inside the traced body (`trans_seg = ".trans" if trans else ""`) makes the whole dispatch silently fall back to an element-wise copy. Chains are built before the prim_func now. Deletes ptx_stmatrix. ptxd derives a "memory" clobber from the address operand, which the legacy helper omitted; the instruction stores to shared memory, so the clobber is the honest reading. Measured: identical register allocation, no spills. * feat(lower-tirx): transcribe the ld/st vector lines Four entries, split by target floor the way ldmatrix's are: the .v2/.v4 lines up to 128 bits, and the two 256-bit lines (.v8 with a 32-bit type, .v4 with a 64-bit type) that ptxas gates behind sm_100 ("Feature '256 bit wide load/store' requires .target sm_100 or higher"). The group length is `.vec`, one lanes function for all four. .level2::eviction_priority lives only on the 256-bit entries, which is where the ISA spells it. Certification ruled on two things the ISA text does not: the sm_100 floor above, and .L2::evict_unchanged, which ptxas rejects on ld/st outright. Both are now in the table. 57,828 variants assemble. The vector entries are separate from the scalar ones rather than an optional slot on them, because a vector operand keeps its braces even at one register -- vector-ness has to be a property of the entry. Call sites move to ptxd across the copy library, the swizzle permute path, and the ld/st tests. copy_ptx_form is now copy_ptxd_form, returning the chain tail, the lane count, and the register container dtype; a scratch is allocated in the container type (these copies move bits, not values) and a register buffer is read through a container view with a scaled index. The `nc` three-way branch collapses into the chain string. Deletes ptx_ld and ptx_ld_global_nc. .b128 is NOT transcribed, and the reason is worth recording: the rendered helper is correct and assembles, but nothing can call it. The operand would have to be a 128-bit TVM value and the CUDA codegen has no uint128 type. The legacy helper did the 128-bit load inside itself, so the wide value never entered the IR. ptx_st therefore stays for the one b128 store site. * feat(lower-tirx): give ptxd a tmem address space, and transcribe tcgen05.ld/st A tmem address is a packed (row << 16 | col) 32-bit value -- neither a shared window address nor a generic pointer. The table said so in a comment and had nowhere to put it; `space="tmem"` is that place. It renders the same 32-bit "r" carrier a shared address does and is written [%N], but nothing is converted: the operand must already be a uint32, and a pointer is refused. tcgen05.ld / .st follow, per ISA 9.7.17.8.3/8.4. The register vector is `.num` scaled by the shape width (Table 52/53), up to 128 registers, which is one lanes function; the two instructions mirror each other's operand order. 58,176 variants certified at sm_100a. NOT REGISTERED, with reasons: the .16x32bx2 shape, whose immHalfSplitoff is an instruction-text immediate the ISA gives no value domain for (the choices mechanism needs a closed set) and which no call site uses; and tcgen05.ld.red, which is sm_101a-only and therefore cannot be certified here. Since ptxd is one instruction per call, the row/col packing the legacy helper did internally moves to the call site as T.cuda.get_tmem_addr -- and is skipped entirely when both are zero, which is most sites. Deletes ptx_tcgen05_ld, ptx_tcgen05_st and ptx_tcgen05_shift. Two fixes fell out. The stub generator emitted shape tokens like `16x64b` as attribute names, which is not valid Python -- those variants are reachable only through the string form, so the stub now skips them. And a codegen test asserted a helper name was absent when it meant no call was made; a definition pulled in by another site is not a use. * feat(lower-tirx): transcribe mma, the last family the lanes mechanism unlocks Four entries covering the certifiable syntax lines of ISA 9.7.15.5.14: the .f32- and .f16-accumulator floating-point lines, the integer / sub-byte / single-bit lines, and double precision. Four operand groups (d, a, b, c), each a register vector whose length follows the Matrix Fragments tables, so four callable `lanes` off one shared rule. d and c are separate operands, as the ISA lists them and as the legacy helper bound them ("=" and "r"); a caller passing the same registers for both is accumulating in place, which needs no read-modify-write constraint. Certification earned its keep here -- four transcription errors, none of them visible in the ISA text: * .m8n8k4 runs four independent 8-thread MMAs, so a thread holds an eighth of the tile, not a thirty-second: d=4, a=2, b=2 for .f16. * The ISA spells ".m8n84" among the double-precision shapes. It is a typo, and not for .m8n8k4 -- ptxas rejects that shape outright, so the f64 shapes are the three m16n8 ones. * .bitOp.popc belongs to the single-bit line alone. * The register groups carry fragments, not elements: A and B are packed into .b32 whatever the element format, and only an .f32 accumulator binds "f". Leaving that to the dtype axis offered 2,678 combinations ptxas refuses. NOT REGISTERED: the .kind::/.block_scale lines and the .e3m2/.e2m3/.e2m1 types, which need sm_120a and so cannot be certified here; and mma.sp, a separate instruction with a metadata operand. The op stays: lower_warp_memory.cc and codegen_cuda.cc both recognize tirx.ptx.mma, so the call sites migrate with the C++ pass rewrite rather than ahead of it. * feat(lower-tirx): print 128-bit integers, and migrate the mma call sites Two things the ptxd migration had been blocked on. The CUDA codegen had no 128-bit integer type, so a .b128 operand could not be written at all: the legacy helper did the wide load inside itself, keeping the value out of the IR entirely. PrintType now spells __uint128_t / __int128_t, ahead of the unsigned "u" prefix -- the name is not "u" plus a signed one. .b128 rejoins the scalar ld/st type axis, its two call sites move to ptxd, and ptx_st is deleted. mma's call sites move to ptxd as well: the tile-primitive GEMM inner loop, the Ampere numeric tests, the script assertions, and nymph's emission. Fragments reach the instruction as register lvalues rather than pointers, and the packed 16-bit multiplicands ride a uint32 view -- two elements per b32, which is what the fragment tables describe. The legacy "omit c" convenience fed literal zeros; ptxd takes the accumulator as the operand it is, so a caller wanting beta=0 zeroes the registers itself. ptx_mma itself stays. lower_warp_memory.cc and codegen_cuda.cc both recognize it, and until those two are rewritten the op has to remain -- only the call sites move. * feat(lower-tirx): let a lanes function return zero, and declare vector-ness Two small extensions of the lanes-as-function mechanism, for the syntax lines that bracket their last operand as optional ({, cache_policy}, {, ctaMask}, {, src-size}): A length function may now return 0, which makes the operand disappear -- no argument, no text -- when its modifier is absent. One entry then covers both spellings of such a line, instead of a copy per optional tail. Vector-ness becomes declarable. The brace rule used to be "callable lanes means a vector group", which ptxas wants for real register groups even at length one -- but an operand whose length varies between 0 and 1 is a bracketed-optional scalar, not a vector. `vector=False` says so; every existing entry keeps the derived default. All 73,315 existing renders are byte-identical. * feat(lower-tirx): add the acc operand role, transcribe wgmma.mma_async and mbarrier.arrive.noComplete An acc operand is a register group the instruction reads and writes in place -- the "+" asm constraint. Unlike a dst it does not block pred=, because "+" keeps the old value live under a false predicate. All 73,315 existing renders are byte-identical. wgmma.mma_async (PTX ISA 9.7.16.5.2, sm_90a) lands as 16 entries: six type groups by ss/rs form, split per accumulator register type so every operand dtype stays pinned, 19,600 renders all ptxas-certified. Two ptxas verdicts encoded: the scale-d predicate position accepts the literals 0/1 (so it is a choices immediate, no setp needed), and imm-scale-a/b take the ISA's documented {-1, 1} (the legacy helper's 0 was outside that domain). Helper-name mangling now maps "-" to "m" for the -1 immediates. mbarrier.arrive.noComplete has no sink form, so its state result rides role=acc ("+l", pinned u64 -- ptxas rejects an .f64 register there), which keeps pred= legal on the write. Delete the ptx_wgmma_mma_async_ss/_rs and ptx_mbarrier_arrive_no_complete legacy ops and migrate their sites (hopper ss/rs tests, noComplete codegen tests, printer assertions). * feat(lower-tirx): add composite bracket operands, transcribe the TMA tensor family A bracket names one composite memory operand: adjacent slots naming the same bracket render inside a single pair of square brackets, each member keeping its own registers and constraints. That is how TMA spells its tensor address -- `[tensorMap, {c0, c1}]` is one PTX operand holding a 64-bit pointer and an .s32 coordinate vector. Zero disturbance: all 92,927 existing renders are byte-identical. cp.async.bulk.tensor / cp.async.bulk.prefetch.tensor / cp.reduce.async.bulk.tensor (PTX ISA 9.7.9.26.5.2-4) land as five entries. The coordinate count follows .dim except under tile::gather4/scatter4 (fixed five, 2d only); ctaMask and cache_policy are trailing operands that exist exactly when .multicast::cluster / .L2::cache_hint are written -- zero-lanes functions each. All 798 renders are ptxas-certified at sm_100a. The im2col load modes stay NOT REGISTERED (no call site ever used them). This is the first family whose optional operands go through the lanes-zero path at codegen time, which exposed a latent bug in the codegen closure: a vanished operand still owns a dtype slot (the render aligns dtypes with every typed operand) but has no argument to read, so it now reports its canonical dtype instead of indexing past the call's arguments. Delete the five legacy ops (g2s_cluster/g2s_cta/prefetch/ shared_to_global/shared_to_global_reduce) and migrate every site: the tma.py dispatch library (which now spells the multicast/cta_group/ cache-hint decisions in the instruction text it always meant), hopper and codegen tests, printer/roundtrip/verifier tests. The .s32 coordinates and .u16 ctaMask conversions the legacy helpers performed implicitly in C are now explicit casts at the call sites. * feat(lower-tirx): add the .pred boundary roles, transcribe tcgen05.mma and the parity waits Inline asm has no constraint letter for predicate registers, so a .pred crossing the asm boundary converts inside the block -- the same exception @p's own setp already established. role="pred_dst" names the predicate the instruction writes and materializes it with a trailing selp.b32 into a "=r" uint32 reference (it gates @p like any dst); role="pred_src" converts a "r"-bound uint32 in with a leading setp.ne.b32. All 93,731 prior renders are byte-identical. mbarrier.test_wait.parity / try_wait.parity (ISA 9.7.14.12) land with waitComplete as a pred_dst; try_wait registers its timeHint arity only. tcgen05.mma (ISA 9.7.17.10.9.1-3, sm_100a) lands as six entries -- the dense, .ws and .block_scale lines, each split ss/ts on where A lives -- with enable-input-d as a pred_src, the disable-output-lane vector sized by .cta_group (4 or 8), .ws's zero-column-mask descriptor operand, and block_scale's scale_vec::NX domain checked per kind. 110 new renders, all ptxas-certified; .sp, collector/ashift, scale-input-d and the phase_type qualifiers stay NOT REGISTERED (no call sites). Restore legacy parity on the derived "memory" clobber: a tmem address is not C-visible memory, and the legacy tcgen05 ld/st/mma helpers never claimed the clobber -- for mma, issued per K-step in the hottest loop, the needless barrier is a real optimization fence. tcgen05.ld/st lose theirs too (their legacy helpers had none either); ordering against tmem consumers belongs to tcgen05.fence/commit/wait. The gemm_async dispatch now spells the full instruction (kind from the trace-time dtypes, ws/block_scale/scale_vec in the chain) and passes the uint32 conversions the legacy helper parameters performed implicitly as explicit casts. Delete ptx_tcgen05_mma, ptx_tcgen05_mma_block_scale, ptx_mbarrier_try_wait_once and ptx_mbarrier_test_wait_parity, and migrate the blackwell tests (the @p + pred_src composition keeps the predicated single-SASS form). * feat(lower-tirx): transcribe cp.async and the non-tensor cp.async.bulk directions cp.async (PTX ISA 9.7.9.26.3.1) lands as four entries: ca/cg, each in its…
* refactor(op): rename the PTX dialect namespace from ptxd to ptx The table-driven dialect was quarantined under `ptxd` because the name it wanted was still taken by the hand-written surface it replaced. That surface is gone -- `tirx.ptx.*` has been empty since the dialect landed -- so the dialect takes the namespace it was always meant to have. One property generates all 174 op names, so the rename is small at the source: `InstructionEntry.op_name`, the printer/namespace attrs and codegen key in `register_table`, and the script-namespace dict key in `backend/cuda/__init__.py`. Everything else is the call sites: `T.ptxd` -> `T.ptx` across the dispatch layer, tests and docs, the two `Op::Get` literals the CUDA codegen holds for cp.async group tracking, and the regenerated `tirx.pyi` stub. Also renamed with it: `PTXDNamespace` -> `PTXNamespace`, the `_PTXD` stub class, `copy_ptxd_form` / `copy_ptxd_ld_chain` / `_ptxd_call_parts`, the `test_ptxd_dialect.py` file, and the `PTXD_ARCH` / `PTXD_CERT` environment variables that gate the certification shards (now `PTX_ARCH` / `PTX_CERT` -- external runbooks setting the old names need updating). The generated CUDA helper prefix `tvm_builtin_ptxd_` is left alone here; it is an independent name and moves in its own commit. * refactor(op): rename generated helper prefix tvm_builtin_ptxd_ to tvm_builtin_ptx_ The dialect's emitted CUDA helpers were the last thing still spelling the quarantine name. The prefix is generated in exactly one place, so this is `render.py` plus the golden strings that assert on the emitted source. No collision with the four legacy `tvm_builtin_ptx_*` helpers that survive on `T.cuda.*` code paths (fetch_register, barrier_cluster_arrive/wait, tcgen05_encode_matrix_descriptor): the audit enumerated all 225428 generated helper names against them and the intersection is empty. A dialect name always repeats the mnemonic after the table key -- barrier.cluster.arrive renders as `..._barrier_cluster_arrive_cluster_arrive` -- so the shapes cannot coincide. One golden also needed its two `cp.async.wait_group` helper *definitions* swapped. The CUDA codegen holds them in an `unordered_map` keyed by helper name (codegen_cuda.h), so the preamble order follows the string hash; the set of definitions is unchanged. * docs(op): fix stale examples and comments left from the legacy surface retirement Around thirty documentation snippets still showed the retired hand-written surface's call forms -- kwargs like `tcgen05.alloc(n_cols=512, cta_group=1)`, positional `mma(shape_str, "row", "col", ...)`, `T.ptx.add_f32x2(...)` for a family that does not exist -- so the namespace rename would have made wrong examples look freshly correct. Rewritten against the actual table entries and dispatch call sites: chain or bracket-string spellings, allocate-then-pass operand order (`tcgen05.ld` takes destinations first and the address last, `tcgen05.st` the reverse), explicit `T.uint32` operands, and `wait::ld` / `wait::st` spelled `wait__ld` / `wait__st`. The "Generated CUDA" blocks now carry the helper names and signatures the renderer actually emits. Also corrected: two docstrings and a C++ comment naming `tirx.ptx.cp_async_raw` (it moved to `tirx.s_tir.cp_async_raw`), a comment naming `T.ptx.ldmatrix_legacy` (now `T.ptx_legacy.ldmatrix`), and the prototype/quarantine framing in the dialect's module docstrings -- the dialect shipped and owns the namespace now. * test(op-dispatch): drop dead legacy helper-name fallbacks and a vacuous guard The packed f32x2 assertions each accepted a second spelling, `tvm_builtin_ptx_<op>_packed_*`, from the hand-written surface. No emitter produces that shape any more -- no table entry or modifier token contains "packed" -- so the alternative could only ever hide a regression in the asm-text assertion that is doing the real work. `test_codegen_cuda.py` also asserted that `tirx.ptx.cp_async_mbarrier_arrive_noinc` was absent from the generated CUDA. That op was deleted with the legacy surface, and op names never reach the emitted source in the first place, so the check was vacuous in both directions.
* fix(infra): give the NVSHMEM objects CUDA's cccl include directory
`tvm_nvshmem_objs` compiles both .cc and .cu sources but was only handed
`${NVSHMEM_INCLUDE_DIR}`. CUDA 13 moved libcu++ from `${CTK}/include/cuda/std`
to `${CTK}/include/cccl/cuda/std`, and NVSHMEM's own headers include
`<cuda/std/tuple>` -- nvcc adds that directory implicitly, so the .cu sources
kept building while `init.cc` and `memory_allocator.cc` failed a from-scratch
build with:
nvshmem_tensor.h:37:10: fatal error: cuda/std/tuple: No such file or directory
This is the relocation NVIDIA documents for CUDA 13.0: nvcc needs no action,
host-compiled translation units must name the directory themselves. The global
`include_directories(SYSTEM ${CUDA_INCLUDE_DIRS})` in cmake/modules/CUDA.cmake
cannot cover it because `CUDA_INCLUDE_DIRS` is `${CUDA_TOOLKIT_ROOT_DIR}/include`
and cccl is a subdirectory of it.
Guarded by `if(EXISTS)`, so CUDA 12 toolkits -- where the headers still sit
directly under `include/` and are already covered -- are unaffected.
* style(infra): apply the two pre-commit fixes main has been carrying
`pre-commit run --all-files` -- what CI runs -- reformats two files that have
been out of compliance since #47 introduced the dialect: a string concatenation
in `intrinsics/sync.py` that ruff-format joins onto one line, and an import pair
in `contrib/hexagon/tools.py` that ruff-check reorders. Neither is reachable
from this branch's change; both reproduce on a clean checkout of `main` and of
`209a4a2239` itself.
They went unnoticed because contributors run pre-commit over changed files
only, and the last PR reported no checks at all. Since CI lints the whole tree,
they fail every pull request against this repo until they land.
Purely mechanical -- the hook output is taken verbatim, no hand edits. The
adjacent string literals in sync.py still concatenate before the `+`, so the
emitted helper text is unchanged; the codegen suites that assert on
`tvm_builtin_cuda_mbarrier_wait` pass (281 tests).
* style(infra): satisfy the pinned ruff's UP038 in the ONNX backend test
The remaining `pre-commit run --all-files` failure: two `isinstance` calls in
`tests/python/relax/test_frontend_onnx_backend.py` still use the tuple form,
which the pinned ruff (v0.12.3) flags as UP038 and cannot fix automatically.
Same provenance as the other two -- last touched by #47, reproduces on a clean
`main`.
Worth noting for anyone reaching for a local `ruff`: UP038 was removed in later
ruff releases, so a system ruff (0.15 here) refuses to select the rule and
reports the file clean. Only the hook's pinned version sees it, which is what
CI runs.
`X | Y` in `isinstance` needs 3.10 and the project requires 3.10 already
(`target-version = "py310"`). Behavior is identical -- checked both branches of
the union against the tuple form -- and the file's 389 tests pass.
`pre-commit run --all-files` is now clean end to end.
tirx-kernels moved the pinned bench sweep from a single workloads.yaml to one config file per kernel, dropping DEFAULT_WORKLOADS. Read the sweep through load_config_dir() so collection no longer fails.
The registry correctness test reads the pinned bench sweep out of a sibling tirx-kernels checkout. Newer checkouts assemble it from one config file per kernel via load_config_dir; older ones keep a single workloads.yaml behind load_workloads(DEFAULT_WORKLOADS). Pick whichever the checkout provides, so the module keeps collecting either way.
The table's law is the ISA: an entry models one syntax group, and no two
entries may model the same one. Two did. PTX ISA 9.7.9.26.4.1 gives the
global -> shared::cta direction as a single line whose {.sem},
{.level::cache_hint} and {.ignore_oob} are optional qualifiers on it, and
both a bare `cp` entry and `cp_async_bulk_g2s_cta` modelled it. They
rendered byte-identical assembly and accepted the same four operands, so
any call to that form resolved to two entries and died with 'ambiguous
ptx table' — the form was unreachable unless a cache hint was written,
purely to break the tie. The comment above the redundant entry recorded
the overlap ('also subsumes this entry's rendering') instead of removing
it.
Delete it, and add the checks that would have caught it at authoring
time rather than years later at a call site:
- `len(TABLE) == len(_ENTRIES)`, since keying by name silently drops a
duplicate and a dropped entry is an unreachable ISA line.
- test_ptx_no_instruction_registered_twice: strip the helper and
parameter names off every rendering and assert no two entries produce
the same instruction text and constraints.
- test_ptx_dispatch_unambiguous: model what the engine actually resolves
by — written tokens, operand count, and each position's acceptance
class, with declared spaces that _coerce_address treats alike
collapsed. Stricter than the rendering check: two entries can emit
different assembly and still leave a call with nothing to choose by.
Both fail on the pair they were written for and pass on the fixed table.
The ambiguity assertion also named the surviving candidates rather than
the entries that accepted, which pointed at the wrong row.
The table's order carried no information: 9.7.9 was split across five runs, 9.7.14 across four, and the warp-matrix chapters interleaved, so finding an instruction meant grepping and adding one meant guessing. Since the ISA is what the table models, let the ISA order it. Every top-level element is now placed by the section it cites (or, where it cites none, by matching its spelling against the ISA table of contents), and each 9.7.N chapter gets one banner. Nine chapters, nine runs, in document order. The moves are pure block permutation: entries are structurally identical before and after, only their order differs. The scattered horizontal rules are gone, since the chapter banners now carry that structure; the prose they headed describes the entry it sits on and travels with it.
The generator piped its output through `ruff format` and silently returned the unformatted text when ruff was missing or failed. The checked-in stub is the formatted text, so the freshness test compared formatted against unformatted wherever ruff is absent — green on a dev box, red on the CPU CI image, which is exactly where it failed. Emit the two shapes ruff was fixing (a docstring that fits on one line closes on that line, and one blank line after each class body) and drop the formatter pass. `generate()` is now pure: with ruff removed from PATH the checked-in stub still matches it byte for byte. A future construct ruff would rewrite now fails the freshness test instead of being papered over, which is the signal we want.
The package models the PTX ISA and exposes exactly one namespace, `T.ptx`; the `_dialect` suffix named the mechanism rather than the contents and made every import path longer than the thing it points at. Pure rename. The module path `tvm.backend.cuda.ptx_dialect` becomes `tvm.backend.cuda.ptx` in its 33 references, both lazy-submodule lists keep their alphabetical order, and the generated stub is regenerated for its new `python -m` line. No entry name, op name, helper symbol or rendered instruction changes, so nothing on the codegen or tirx-kernels side moves.
…nto a dtype `role` fused two independent axes. Data direction decides the whole C-boundary story -- constraint letter, pass-by, lvalue requirement, @p legality -- while "is this a register at all" decides something else entirely. Fusing them left the pred+accumulator cell inexpressible and made `has_dst` carry a special case for a value that is simply a destination. Split into `rw` ("r"/"w"/"rw") and `kind` ("reg"/"addr"/"ptr"/"imm"). The non-register kinds are inherently read-only -- no ISA line writes an address register or a text immediate -- so they leave `rw` at its default and the flat enum loses nothing. `has_dst` is now just "any register written". With direction factored out, `.pred` stops being a role and becomes what the ISA calls it (5.2: a fundamental type, declared `.reg .pred p` exactly as `.reg .b32 r` is). The setp/selp conversions move out of hand-written render branches into a BRIDGE row keyed by dtype, with `rw` picking which fire; the emitted text is unchanged. That, in turn, closes the one place where dispatch had run out of axes. PTX tells a predicate operand from an integer one by the declared register class, but both cross the C boundary in a uint32, so the class was erased before the engine saw it -- which is why the two `cp.async` syntax lines differing only in `{, src-size}` vs `{, ignore-src}` could not both be registered. `T.ptx.pred(x)` puts that declaration back at the call, where PTX writes it. It is a tag, not a conversion: the value reaches the helper untouched and the asm is byte-identical. A bool-typed expression needs no tag, since its dtype already names the class; an untagged integer is now refused, by a message naming the fix. Proven behaviour-identical by rendering all 143288 variants before and after. One name moves: clusterlaunchcontrol.query_cancel.is_canceled gains a `u32` segment when its .b128 operand takes the non-canonical dtype, because the discriminator names every typed operand positionally and the predicate is now one. Nothing calls that variant and no golden pins it.
`.e2m1x2`'s operand is typed .b8 by the ISA, a register class inline asm cannot bind, so it is staged through a block-local `.reg .b8` with a conversion at each boundary. That was a hand-written helper body behind `raw_render`, exempted by name from the single-instruction invariant. It is the same shape as `.pred`: a register class the constraint alphabet cannot express, a wider carrier that can, and a measured conversion pair between them. So it becomes a `render.BRIDGE` row, and the four cvt entries become ordinary entries. `_cvt_f4x2_raw` and the exemption set are gone; `raw_render` itself stays for a family that is genuinely irregular, with no users today. The invariant's sanctioned prefixes grow to match, and its falsification twin had to change: its "prologue" case was literally the b8 shape, which is now legal, so it is replaced by a prologue that computes (`shl`) rather than converts, plus a conversion on a class that has no bridge row -- and the b8 staging joins the positive cases. Rendering is untouched: all 143288 variants are byte-identical to before.
ISA 9.7.9.26.3.1 gives cp.async four syntax lines, two of which take
`{, ignore-src}` where the others take `{, src-size}`. Only the src-size half
was in the table: both add one operand at the same position, so arity cannot
separate them, and before `.pred` was a dtype neither could the acceptance
class -- a predicate and a byte count arrive as the same uint32. Registering
both would have created two entries with identical discriminators, one of them
permanently unreachable.
The ISA separates them by register class, in as many words: "The optional and
non-immediate predicate argument ignore-src" against "a 32-bit integer operand
src-size". The table now separates them the same way, and the caller writes
`T.ptx.pred(...)` where PTX writes `%p`.
The falsification twin is the point of the commit as much as the entries are:
collapse `.pred` back into u32 in the dispatch model and exactly these two
pairs collide -- shown, not asserted. All 128 new variants assemble under
ptxas at sm_90.
`{.level::cache_hint}` with its trailing `{, cache_policy}` was documented in
four places as needing "a mechanism this shape lacks". It does not: an operand
whose presence is a function of the modifiers is what `_present_lanes` has
always been, and cp.async and the TMA family already spell exactly this pair.
The comments were stale; the four families now carry the axis.
Its grammar is measured, not assumed. ptxas 13.2 at sm_90 answers "Modifier
'.L2::cache_hint' cannot be ..." for `.local`, `.shared`, the `.volatile` line
(whose ISA syntax carries no cache_hint at all) and `.mmio` -- so the qualifier
is confined to `.global` and generic addressing, and the mmio rule grew to
name it. It is accepted alongside .cop, .nc, both eviction priorities,
.level::prefetch_size and the scoped acquire/relaxed lines, all probed.
59246 cache_hint variants assemble.
Two latent bugs surfaced and are fixed here because the axis is what exposed
them:
- The helper-name discriminator counted operands that are not there. A
bracketed-optional operand resolves to zero lanes and contributes no C
parameter, so counting it renamed every non-canonical helper of an entry
the moment such an operand was added. It now names only present operands,
which leaves every existing name untouched.
- gen_stubs emitted two `*` parameters for a single-entry family with
modifier-dependent lanes -- `atom` became the first such family, and the
stub stopped parsing. The catch-all is now recognised as already covering
the round-trip arguments.
The stale NOT REGISTERED notes are rewritten to say what actually keeps the
remaining forms out: .unified/.param/.const are symbols rather than values and
need a rendering model that emits where the symbol is in scope, and red/atom's
vector and half-type lines are separate syntax shapes awaiting a caller.
PTX writes a discarded destination element as `_`, at the operand position:
`mov.b64 {_, %0}, %1;`. The table now writes it there too. `T.ptx.SINK` at a
lane of a `sinkable` operand renders the symbol, and that lane loses its C
parameter and its constraint -- which makes the chosen mask part of the
variant, since a different mask is a different helper signature. The empty
mask keeps every name exactly as it was.
Unlike the other axes this one is the caller's choice rather than a property
of the instruction, so it is enumerated per call site: `sink_combos` walks the
subsets of the sinkable lanes, minus the all-sunk one. ISA 9.7.9.4 states that
exclusion for mov ("provided that at least one element is a scalar register")
and it is the conservative reading elsewhere -- an instruction whose every
destination is discarded has nothing left to do.
Registered where the ISA sanctions it and the domain is bounded: mov's unpack
destinations, and clusterlaunchcontrol's .v4, whose own ISA example is
`{xctaid, _, _, _}` and whose fourth element is documented as unspecified.
Sinking a lane of an accumulator drops both halves of its read-modify-write,
which is exactly what that example does, so `rw="rw"` operands take it too.
474 sink variants assemble under ptxas.
The 256-bit ld/st lines stay out, and the note now says why with numbers
rather than blaming a missing mechanism: 2**8 masks over their 25392 and 11952
renderings is 3453312 and 1625472 helpers against 161108 for the whole table,
and each mask is a signature the certification tier would have to prove.
The round trip needed the marker again, and more than pred did: a sunk lane
leaves no argument at all, so the printed call is shorter than the one written
and would re-parse at the wrong arity. The marker names the positions and they
are re-inserted before dispatch.
…ink axis
The 256-bit ld/st sink was held out on a reason that did not survive being
questioned. The claim was that its domain is unaffordable: 2**8 masks over
25392 modifier combinations is 3.45 million helpers. The arithmetic is right
and the conclusion was wrong, for two reasons.
Nothing is instantiated ahead of time. A helper is rendered from the mask the
call site actually wrote -- three ptx calls in a kernel produce three helpers,
sunk lane included. The table holds 175 entries; `renderings()` is a
verification-time enumeration, not a build artifact. So registering the lines
costs nothing at compile time.
And the multiplication bought nothing. Whether ptxas takes `_` at a given lane
does not depend on `.cop`, `.scope` or the eviction priorities, so the product
re-proved one fact 25392 times. The sink axis is now *added* rather than
multiplied: the full product with nothing sunk, plus the mask domain once per
distinct domain -- `.v4` and `.v8` sink different numbers of lanes, so one
representative each, not one per entry. The table grows 161108 -> 162716
instead of to millions, and 2412 new sink variants assemble under ptxas.
Registering them also falsified the rule I had written down. `_` is not a
destination spelling: ISA 9.7.9.11 puts it in st's "vector expression b", the
data being stored, and ptxas assembles
`st.global.v8.b32 [%4], {%0,%1,%2,%3,_,_,_,_};`. So sinkability is a per-slot
fact read off the syntax line, and it now says what each direction means --
not written (ld), not stored (st), neither (clusterlaunchcontrol's accumulator).
`sinkable` may be a function of the modifiers, like `lanes`, because the ISA
gates these two lines on the vector width and element type. ptxas is looser
than the ISA there -- it takes `_` on .v4 with a 32-bit type as well -- and
that spelling stays out: toolchain evidence narrows what the ISA permits, it
never widens it.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation and context
tirx.ptx.*grew as a hand-written intrinsic per PTX instruction: every newinstruction meant a Python wrapper, a codegen entry, and a hand-rolled
asm volatiletemplate, with the modifier grammar encoded implicitly in keywordarguments. Adding a modifier meant editing three places, and nothing checked the
result against the ISA's own grammar.
This replaces that surface with a table-driven dialect. Instruction shape lives
in one table (mnemonic, modifier slots with their legal tokens, operand roles and
dtypes); the surface, the codegen, and the type stubs are all derived from it.
T.ptx.fence.proxy.async_.shared__cta()andT.ptx["cp.async.bulk.tensor.2d. shared::cluster.global.mbarrier::complete_tx::bytes"](...)are the sameinstruction reached two ways, and an illegal modifier is rejected at parse time
with the open slots listed.
Changes
python/tvm/backend/cuda/ptx_dialect/): 174instruction entries with modifier slots, operand roles/dtypes, and per-entry
checks; a renderer that builds the
asmtemplate from the resolved chain; andgenerators for the coverage report and the
T.ptxtype stubs(
python/tvm/script/tirx.pyi).tirx.ptx.*and its per-instructionintrinsics are retired; all in-repo call sites move to the new surface.
dialect needed (cvt/mma/ld-st families).
replaces the additive signed-strides path,
ldmatrix/stmatrixdestinationregisters are ordered to match the fragment word order, and register-copy base
offsets are hoisted out of the serial loop.
ccclinclude directory, fixing their build.workload layout the sibling
tirx-kernelscheckout provides.Testing
tests/python/tirx/suite on B200 (sm_100a): 2679 passed, 73 skipped,3 xpassed, 0 failed, with both
tirx_kernelsimport gates green.PTX_CERT=1assembles the generated forms throughptxasat each entry'sown ISA floor;
test_ptx_registrationchecks every table entry is registeredin a
USE_CUDA=OFFbuild andtest_ptx_stub_up_to_datebyte-compares thechecked-in stub against the generator.
docs/tirxpages.pre-commit run --all-filesclean at the pinned hook versions.