Tirx Ptx Dialect

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.

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