TIRx PTX dialect (T.ptx)

T.ptx.<mnemonic>.<qualifiers>(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.<name> 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 <family> — 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("<slot>"). 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 names (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):

    # 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.

   python -c "from tvm.backend.cuda.ptx.table import TABLE, variants; e=TABLE['<name>']; print(len(variants(e)))"
   python -m tvm.backend.cuda.ptx.gen_helpers <name> | 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_<section>_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:

   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.

   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 == <N>; 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).

   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 <arch> batch <k> and the ptxas message. To find the variant: render the failing family with gen_helpers, or assemble the batch yourself (nvcc -arch=<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 <changed files> and commit as feat(lower-tirx): support PTX <instruction> (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:

  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/<cuda version>/.
• 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 (<num> <title>), 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).