--- name: tirx-ptx-dialect description: Register, extend, or audit instructions in the table-driven T.ptx dialect (python/tvm/backend/cuda/ptx/table.py), and move the table to a newer PTX ISA version. Use when adding a PTX instruction or qualifier, widening an operand domain, fixing a ptxas certification failure, or checking the table's comments against the PTX ISA document. --- # TIRx PTX dialect (`T.ptx`) `T.ptx..(operands...)` is not hand-written per instruction. One data table describes every instruction family; one engine resolves calls, renders an inline-asm helper, and registers the TVM op; three generators derive the IDE stub, a coverage table and a helper dump from the same table. Adding an instruction means adding **data**, then proving it against ptxas. | file | role | |---|---| | `python/tvm/backend/cuda/ptx/table.py` | the table: `InstructionEntry` / `ModifierSlot` / `OperandSlot`, check functions, variant enumeration. tvm-free. **Read its three dataclass docstrings before editing.** | | `python/tvm/backend/cuda/ptx/render.py` | renders one variant to a `__forceinline__ __device__` helper (`asm` or `asm volatile`, according to `asm_volatile`). `CBinding` = how each TVM dtype binds a register; `BRIDGE` = three scoped cases covering the two register classes inline asm cannot bind directly (`.pred`, and the block-local `.b8` registers behind the `e2m1x2` and `st_async_b8reg` tokens) — not a general `.b8` mechanism. tvm-free. | | `python/tvm/backend/cuda/ptx/engine.py` | `register_table()` makes each entry a TVM op `tirx.ptx.` with a generic codegen; `PTXNamespace` resolves `T.ptx` attribute chains and the string form `T.ptx["..."]`; trace-time coercion and `check`/`imm_check`. | | `python/tvm/backend/cuda/ptx/gen_stubs.py` | `python -m tvm.backend.cuda.ptx.gen_stubs -o python/tvm/script/tirx.pyi` — the checked-in stub; a test diffs it against the generator. | | `python/tvm/backend/cuda/ptx/gen_helpers.py` | `python -m tvm.backend.cuda.ptx.gen_helpers ` — dump the exact helper source of every variant without compiling a kernel (it still imports `tvm`, so it needs the usual built worktree on `PYTHONPATH`). | | `python/tvm/backend/cuda/ptx/gen_coverage.py` | markdown coverage table of the whole dialect. | | `tests/python/tirx/codegen/test_ptx_dialect.py` | per-ISA-section dispatch tests, structural invariants, ptxas certification. | | `tests/python/tirx/codegen/test_ptx_cvt.py` | one case per registered `cvt` syntax line (`_FORM_CASES`); coverage of every cvt entry is asserted. | | `tests/python/tirx/codegen/test_ptx_addr.py` | `T.ptx.addr(base, byte_offset)` immediates. | The namespace itself is installed by `python/tvm/backend/cuda/__init__.py::script_namespaces()`; nothing there changes when the table grows. > **DO NOT introduce a new mechanism in an `InstructionEntry` without the > user's explicit approval.** Stop, explain why the existing entry model is > insufficient and what codegen or validation behavior the mechanism would > add, then wait for approval before implementing it. ## 1. Ground rules the table enforces The first four are tested mechanically — a new entry that breaks one fails the suite. The last one (what the comments claim) is audited by hand, §4. - **One entry = one syntax shape.** Split into sibling entries sharing the `mnemonic` when the register-group shape or the result structure changes (scalar vs vector destination, `mul` vs `mul.wide`), or when an optional operand is told apart only by call arity (`{, count}`, `state` vs the `_` sink); the engine then selects by arity and by which tokens are written. An operand whose presence is fixed by one written qualifier stays in the entry with a 0-or-1-register `lanes` function — `{, cache_policy}` with `.L2::cache_hint`, `{, ctaMask}` with `.multicast::cluster` use `lanes=_present_lanes("")`. Variants that only add or remove dotted qualifiers are optional slots of one entry. - **Single-instruction invariant**: every helper emits exactly one native PTX instruction. The only extra statements allowed are the framework `@p` wrapper and the `BRIDGE` boundary conversions. `raw_render` is the escape hatch and every user of it must be listed in `RAW_ENTRIES` inside `test_ptx_single_instruction_invariant`. - **Every legal closed variant assembles**: `renderings()` enumerates the product of modifiers × operand dtypes × closed immediates × `@p` (for entries without a destination) and full certification assembles all of it through ptxas at `cert_arch`. Four things get *representative* coverage instead of the full product, and each says so in its docstring: an OPEN immediate is certified at `imm_combos`'s samples (default `"0"`), which proves the shape, not the caller's constant; sink masks are walked once at one representative modifier combination per sink domain (`renderings`); address immediates are a small separate axis (`_addr_offset_samples`); and the destination-predicated `pred=` helpers (`*_pred_undef` / `*_pred_keep`) are not enumerated at all — cover those with a production-shaped test when a call site needs them. Conversely a `check()` must never reject a legal form to make certification pass — narrow with evidence. - **Dispatch must be unambiguous**: no two entries may accept the same call (`test_ptx_dispatch_unambiguous`); helper names must be unique (`test_ptx_all_variants_render_unique`). - **Comments cite the ISA**: the section number, the syntax line reproduced verbatim, the type/qualifier domains, `sm_XX` floors, and a `NOT REGISTERED:` note with the reason for every syntax line or token deliberately left out. A fact that comes from ptxas rather than the document is marked `MEASURED` and quotes the ptxas message. Section and table numbers follow the ISA version named in the module docstring (currently **PTX ISA 9.4**, the CUDA 13.4 developer-preview document, see §5). ## 2. Designing the entry Open the instruction's section in the ISA document and transcribe, in this order: 1. **Name / mnemonic / family.** `name` is the table key and must be a Python identifier. Dotted mnemonics are a family plus single-choice slots (`cvta.to.shared` → `cvta` + slots), or `mnemonic="st.bulk"` with `name="st_bulk"` when the dot is part of the instruction's identity. Several entries sharing a mnemonic (`mov`, `mbarrier`, `mma`) use distinct `name`s (`mov`, `mov_pack_2`, ...) and the same `mnemonic`. Entries that collide with an existing family on the same shape need a suffix (`add_int` vs `add`/`add_half`) and are told apart by their tokens. 2. **Slots** (`ModifierSlot(name, choices, optional)`), in asm render order — exactly the `{.qual}` positions of the syntax line, each with the ISA's `.qual = { ... }` list as `choices`. Tokens with `::` are written as-is (`"shared::cta"`); users type `shared__cta`. Python keywords get a trailing underscore only at the call site (`global_`). 3. **Operands** (`OperandSlot`), in PTX operand order, destinations first because PTX writes them first: - `rw`: `"w"` destination, `"rw"` accumulator (`"+"`), `"r"` input. A `"w"` operand binds `"="`; it does **not** block `pred=` — under a false predicate the value is undefined unless the caller passes `preserve_dst=True`, which switches the binding to `"+"` (see `test_ptx_predicated_destination_*`). Those two helpers are rendered on demand, not by `renderings()`. - `kind`: `"reg"` (default), `"addr"` (`[%k]`, space from `space=` or the entry's `space` slot; set `allow_imm_offset=True` only for an independent byte address), `"ptr"` (raw pointer value), `"imm"` (text operand: `literal=` fixed by the ISA, `choices=` closed caller set, neither = open immediate certified at samples; validate open ones with `imm_check`). - `dtype`: a `PTX_TYPE_DTYPES` key, a slot name, or a module-level pure function of the modifier map (`_wide_dtype`). `None` means the entry's `type` slot. `dtypes=` narrows/widens the TVM dtype domain independently (see the relaxed-carrier note in §6 before widening anything). - `lanes` (int or function of the modifier map) for brace-enclosed register groups; `vector=`, `bracket=`, `pipe=` for the few composite spellings; `sinkable=` only where the ISA lets a caller write `_` for a lane. 4. **`check`**: one pure module-level function `mod_map -> error | None` per entry with a one-line docstring (it is surfaced by the stub and the coverage table). Encode only what the ISA syntax block and notes state; every restriction that comes from ptxas instead is a separate `MEASURED` clause. Never a lambda: frozen dataclasses hash callables by identity. 5. **`cert_arch`**: the **maximum** `sm_XX` floor over the entry's variants (default is `PTX_ARCH`, `sm_90`). Certifying below a variant's floor makes ptxas report legal forms as illegal, and that verdict would then get baked into a `check()`. 6. `orders_memory=True` for fences/barriers/waits (no address operand but the `"memory"` clobber is needed); `asm_volatile` stays at its default. Worked example — the whole registration of `movmatrix` (commit a29a5e97ba, 4 files, 40 lines): ```python # movmatrix per PTX ISA 9.7.16.5.17 -- transpose one distributed m8n8 # matrix whose 16-bit elements are carried by one b32 register per lane. # # movmatrix.sync.aligned.m8n8.trans.b16 d, a; InstructionEntry( name="movmatrix", slots=( ModifierSlot("sync", ("sync",)), ModifierSlot("aligned", ("aligned",)), ModifierSlot("shape", ("m8n8",)), ModifierSlot("trans", ("trans",)), ModifierSlot("type", ("b16",)), ), cert_arch="sm_75", operands=( OperandSlot("d", rw="w", dtype="b32"), OperandSlot("a", dtype="b32"), ), ), ``` Put the entry under its ISA chapter banner in `_ENTRIES` (`# PTX ISA 9.7.x — ...`), next to the instructions it shares a mnemonic with. ## 3. Verify — every gate, in this order Run from the TVM repo root with the workspace `PYTHONPATH` (see `tir-test`). 1. **Table loads and enumerates.** A `check()` that filters out every combination, a duplicate name, or a bad `allow_imm_offset` slot fails at import. ```bash python -c "from tvm.backend.cuda.ptx.table import TABLE, variants; e=TABLE['']; print(len(variants(e)))" python -m tvm.backend.cuda.ptx.gen_helpers | head -60 # eyeball the asm ``` 2. **Add tests for the new surface** (before running the suite, so one run settles everything): - a call in the dispatch test of its ISA section (`test_ptx_
_dispatch`) asserting the exact emitted line, e.g. `assert "mul.wide.s32 %0, %1, %2;" in src`; - a new `cvt` line → a `_FORM_CASES` row in `test_ptx_cvt.py` (`test_cvt_cases_cover_every_registered_entry` fails otherwise); - `raw_render` → `RAW_ENTRIES`; - a destination-predicated (`pred=`) call site → a production-shaped codegen test, since certification does not enumerate those helpers; - only when the instruction belongs to the MegaMoE extracted-intrinsics surface → its kernel and mnemonic list in `test_codegen_cuda.py::test_megamoe_extracted_intrinsics_codegen` (that test covers one workload, not the dialect). 3. **Regenerate the stub** — any change to a family's tokens or shapes makes `test_ptx_stub_up_to_date` fail until you do: ```bash python -m tvm.backend.cuda.ptx.gen_stubs -o python/tvm/script/tirx.pyi ``` 4. **Fast suite** (no `PTX_CERT`): structural invariants + dispatch tests + a seeded ptxas sample. ```bash python -m pytest tests/python/tirx/codegen/test_ptx_dialect.py tests/python/tirx/codegen/test_ptx_cvt.py tests/python/tirx/codegen/test_ptx_addr.py -q ``` `test_ptx_all_variants_render_unique` ends with `assert total == `; its failure message prints the new total — update the constant and re-run. 5. **Full certification** — mandatory after any table change, and the only gate that proves the domain. It assembles every closed variant of every entry, split into 32 shards; pick 8, 16 or 32 workers to taste (each one drives its own nvcc; `-n auto` on a many-core box mostly burns memory). ~3.5 min at `-n 16`, ~6.5 min at `-n 8` on a B200-class host for the ~760k-variant PTX ISA 9.4 table (`test_ptx_all_variants_render_unique` pins the exact count). ```bash PTX_CERT=1 python -m pytest -n 16 -q \ tests/python/tirx/codegen/test_ptx_dialect.py::test_ptx_all_helpers_certify ``` A failure names ` batch ` and the ptxas message. To find the variant: render the failing family with `gen_helpers`, or assemble the batch yourself (`nvcc -arch= -ptx` then `ptxas`, map the error line back to the preceding `.entry`). Fix by narrowing the domain with a `MEASURED` clause, never by loosening the invariant. 6. `pre-commit run --files ` and commit as `feat(lower-tirx): support PTX ` (Conventional Commits, see the workspace `CLAUDE.md`). ## 4. Auditing the comments against the ISA The comments are load-bearing (they are what a reviewer checks the data against), so audit them the way the data is certified: download the ISA HTML for the version the module docstring names, convert to text, and verify every section number, quoted sentence, reproduced syntax line, type list and `sm_XX` note. Things that have gone wrong before and are worth a targeted pass: a `NOT REGISTERED` bullet naming something the table registers a few hundred lines later; a section number copied from a neighbouring instruction; a version number attached to the wrong syntax line; "the N-th syntax line" ordinals; `Table NN` numbers, which are global and shift between ISA versions. ## 5. Toolchain and ISA-version model - The dialect targets the **ptxas of the installed CUDA toolkit**, not the latest ISA document. CUDA 13.4's ptxas implements PTX ISA 9.4. Establish what yours implements from ptxas itself — `nvcc -ptx` only shows the version the front end chose to emit, and `-ptx` never validates inline asm: ```bash command -v nvcc ptxas && nvcc --version | tail -1 && ptxas --version | tail -1 printf '.version 9.4\n.target sm_107a\n.address_size 64\n.visible .entry k() { ret; }\n' \ | ptxas -arch=sm_107a - -o /dev/null # CUDA 13.4: accepted ``` A `.version` directive above what ptxas implements is refused outright ("Unsupported .version ...; current version is '9.4'"), which is the quickest way to read a toolkit's ceiling. Then let the certification suite (which compiles real inline-asm helpers to cubin) be the final word. - The table's citations follow that same version (module docstring of `table.py`: "Section and table numbers cite PTX ISA 9.4", URL `docs.nvidia.com/cuda/developer-preview/13.4/parallel-thread-execution/`; once CUDA 13.4 ships, the stable copy is `docs.nvidia.com/cuda/archive/13.4.0/parallel-thread-execution/`). The live `docs.nvidia.com/cuda/parallel-thread-execution/` page is whatever version is newest and its numbering differs; every ISA version stays available under `docs.nvidia.com/cuda/archive//`. - Every `MEASURED` clause in `table.py` / `render.py` names the toolkit it was taken on (currently CUDA 13.4). A toolkit bump re-measures all of them (§6 step 5); a clause that names an older toolkit is a bug. - Certification needs only `nvcc`/`ptxas`; the on-GPU round-trip tests need a driver that can load a cubin from that toolkit. A newer toolkit with an older driver can therefore certify a table but not run it. - `PTX_ARCH` (env, default `sm_90`) is the certification arch for entries without `cert_arch`. ## 6. Moving the table to a newer PTX ISA Do this as one commit series, in this order: 1. **Toolchain first.** Install the CUDA toolkit whose ptxas implements the target ISA and confirm via `.version` (§5). Until then, a new form cannot be registered: certification would fail, and a `check()` written against an older ptxas would silently delete the coverage later. 2. **Build the section map structurally, from both TOCs.** Dump both HTML documents to text, extract the two tables of contents (` `), and derive a `rule(old) -> new` function from where the new version inserted sections; cross-check it by asserting every old `9.7.*` entry maps to an identically titled new entry. Never trust title matching alone (duplicate titles such as "mov" x2 or "Async Proxy" x2 collide) and never shift chapters only — a chapter-only shift was the exact mistake made the first time this table changed versions. 3. **Retarget the citations.** Update the module docstring's version statement and URL, then run the renumbering over `table.py`, `engine.py`, `render.py`, `tests/python/tirx/codegen/test_ptx_*.py` and `test_codegen_cuda.py`, skipping any block that already carries the new numbering (`_PTX_94_ENTRIES` was such a block for 9.3→9.4). Hand-fix the forms a regex cannot: ranges that span an insertion (`9.7.x.a-b`), brace groups (`9.7.4.{1,2,3}`), sibling shorthand (`.12`, `/ .17`), global `Table NN` numbers, and `:line` offsets (re-derive against the new section text from the quoted sentence, or drop the offset and keep the quote). Then regenerate the stub (§3). 4. **Register the newly legal forms.** Read the new version's release notes (ISA chapter 13.1) and the per-section "introduced in PTX ISA version X.Y" notes; each one: read the section, apply §2, then §3 including full certification at the right `cert_arch`. Keep the new version's additions grouped (the `_PTX_94_ENTRIES` list is the 9.4 group). 5. **Re-measure every `MEASURED` clause on the new toolkit** and reword it to name that toolkit. Registered forms are re-proven by full certification; excluded forms need a targeted probe each (raw PTX through `ptxas -arch=<a>` for grammar claims, an exact force-inlined kernel through `nvcc` for crash claims — `test_ptx_dialect._certification_kernel` builds that shape). Quote the diagnostic the new ptxas actually prints; texts change between releases. If a gap has closed, either widen the domain with certification or say so in the clause — never leave a sentence that attributes the restriction to a toolkit you no longer use. ISA §9.4.1 Tables 27/28 are an upper bound and say so ("some combinations may still be invalid for a particular instruction"); the `_cvt_dst_dtypes` / `_cvt_src_dtypes` docstrings record the measured gaps. 6. Regenerate the stub, update the variant/address counts the tests pin, run the full suite, and run the §4 audit against the new document before merging. Worked example, PTX ISA 9.3 → 9.4: - 9.4 inserted chapter `9.7.6 Alternate Floating-Point Instructions`, so every `9.7.N` with N ≥ 6 became `9.7.(N+1)` (comparison `9.7.6→9.7.7`, data movement `9.7.9→9.7.10`, fabric `9.7.10→9.7.11`, sync `9.7.14→9.7.15`, mma `9.7.15→9.7.16`, wgmma `9.7.16→9.7.17`, tcgen05 `9.7.17→9.7.18`, misc `9.7.20→9.7.21`). - Inside data movement, `applypriority.async.bulk[.tensor]` were inserted as `9.7.10.18/19`, so `9.7.9.18..27 → 9.7.10.20..29` (+2), and "Overriding tensor property value" as `9.7.10.28.5.2`, so `9.7.9.26.5.2..4 → 9.7.10.28.5.3..5`. - Inside tcgen05, "Decompression of input matrices" became `9.7.18.10.8`, so `9.7.17.10.8.x → 9.7.18.10.9.x` and `9.7.17.10.9.x → 9.7.18.10.10.x`. - Global tables: tcgen05.ld/st register-count tables `52/53 → 59/61`, tensormap `new_val` validity `33 → 36`; the relaxed-typing tables `27/28` kept their numbers. - Two files still carried 9.2 numbering (`test_ptx_cvt.py`, `render.py`: cvt as `9.7.9.21`, which in 9.3 is cvta) — check every file's baseline before mapping it. - Toolchain facts that changed on 13.4: the `.L2::cache_hint` diagnostics on `.shared`/`.local`/`.volatile`, the bare `clusterlaunchcontrol.query_cancel` diagnostic, `multimem.st.async` accepting 16/32-bit sources for its byte forms, and the bf16 atom bit-bucket forms compiling again (still withheld pending certification). - 9.4-only sections worth registering later: `9.7.6.1-4` (alternate FP x4 arithmetic, sm_100a/sm_103a), `9.7.10.28.1.3` (report mechanisms), `9.7.18.10.7.2.6 / .3.6` (block16 K=128/256 scale layouts).