You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by bo...@apache.org on 2026/08/29 06:22:41 UTC
(tvm) branch main updated: [TIRx][CUDA] Align the PTX dialect with PTX ISA 9.2 (#20224)
This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 6a05afaf88 [TIRx][CUDA] Align the PTX dialect with PTX ISA 9.2 (#20224)
6a05afaf88 is described below
commit 6a05afaf888e14db4ef9e8a7047c544cd318880d
Author: Bohan Hou <bo...@andrew.cmu.edu>
AuthorDate: Sat Aug 29 02:22:34 2026 -0400
[TIRx][CUDA] Align the PTX dialect with PTX ISA 9.2 (#20224)
## Summary
This PR audits the table-driven `T.ptx` dialect against PTX ISA 9.2, the
ISA version shipped by CUDA 13.2, and brings the registered surface,
operand contracts, generated stubs, documentation references, and
certification tests into alignment.
The main changes are:
- add missing PTX 9.2 forms for vector `red`, `cp.reduce.async.bulk`,
atomic bit-bucket destinations, mbarrier bit buckets, `mapa` carrier
shapes, and the `lop3` predicate-destination sibling;
- model relaxed operand carriers for `ld`, `st`, `ldu`, scalar `cvt`,
`slct`, and `st.bulk` while preserving existing canonical helper names
and signatures;
- correct `wgmma.mma_async` `scale-d` from a text immediate to a runtime
predicate operand;
- distinguish generic and explicit shared address carriers for bare
versus `.shared::*` `tcgen05.alloc` and `tcgen05.commit` forms;
- add an instruction-local `lop3 immLut` check for the documented 0..255
domain;
- add a private byte-register bridge for the PTX 9.2
`st.async.release.{b8,u8,s8}` forms, without changing the existing
16/32/64-bit helpers;
- keep CUDA 13.2 toolchain gaps local to the affected operand callbacks,
including scalar `cvt` and `st.v2.f64`, without adding a generic
operand-range or renderer mechanism;
- correct PTX 9.2 section and table references and regenerate the
checked-in stub;
- extend generated-stub, invariant, round-trip, exact-rendering, and
production-shaped nvcc/ptxas certification coverage.
## CUDA 13.2 boundaries
The table exposes callable inline-assembly helpers, so documented PTX
forms that CUDA 13.2 ptxas does not compile are reported explicitly
instead of being advertised as usable helpers. These restrictions are
operand-local and measured with exact force-inlined callers. They cover
bf16 atomic bit-bucket forms, specific widened scalar `cvt`
combinations, `st.v2.f64` with 128-bit source carriers, and widened
carriers on 256-bit `ld`/`st`. Supported wide-carrier forms remain
registered and certified.
## Testing
- `pre-commit run --files` on all changed files
- `python -m pytest -q -ra tests/python/tirx/codegen/test_ptx_dialect.py
tests/python/tirx/codegen/test_ptx_cvt.py
tests/python/tirx/codegen/test_ptx_addr.py
tests/python/tirx/codegen/test_codegen_hopper.py -k 'not
test_ptx_all_helpers_certify'` (`215 passed, 2 skipped, 32 deselected`)
- `PTX_CERT=1 python -m pytest -n 16 -q -ra
tests/python/tirx/codegen/test_ptx_dialect.py::test_ptx_all_helpers_certify`
(`32 passed`)
- `cmake --build build --parallel`
- `python -m tirx_kernels.bench_suite --check-imports` (80 kernels
imported)
- `python -m tirx_kernels.registry --cc 10 --strict`
- full TIRX suite on three otherwise-idle B200 GPUs: `2888 passed, 91
skipped, 3 xpassed`; one external cuDNN/CUTLASS reference failed because
its DLPack tensor was on CUDA device 1 while the CUTLASS runtime current
device remained 0, and the same case passed in a fresh single-GPU
process (`1 passed`)
- five large registry cases that hit OOM on GPUs occupied by unrelated
processes in the initial all-GPU run each passed in fresh single-GPU
processes
---
.agents/skills/tirx-ptx-dialect/SKILL.md | 308 +++++
python/tvm/backend/cuda/ptx/engine.py | 29 +-
python/tvm/backend/cuda/ptx/gen_stubs.py | 2 +
python/tvm/backend/cuda/ptx/render.py | 20 +-
python/tvm/backend/cuda/ptx/table.py | 1470 +++++++++++++++-------
python/tvm/script/tirx.pyi | 40 +-
tests/python/tirx/codegen/test_codegen_hopper.py | 4 +-
tests/python/tirx/codegen/test_ptx_addr.py | 4 +-
tests/python/tirx/codegen/test_ptx_cvt.py | 229 +++-
tests/python/tirx/codegen/test_ptx_dialect.py | 1169 ++++++++++++++++-
10 files changed, 2756 insertions(+), 519 deletions(-)
diff --git a/.agents/skills/tirx-ptx-dialect/SKILL.md b/.agents/skills/tirx-ptx-dialect/SKILL.md
new file mode 100644
index 0000000000..983350d149
--- /dev/null
+++ b/.agents/skills/tirx-ptx-dialect/SKILL.md
@@ -0,0 +1,308 @@
+---
+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 (9.3, 9.4, ...). 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.<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.2**, 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.14.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['<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:
+ ```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 == <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 the current B200 host for the
+ ~800k-variant PTX ISA 9.2 table.
+ ```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 `<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.2's ptxas implements PTX ISA 9.2. 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.3\n.target sm_90\n.address_size 64\n.visible .entry k() { ret; }\n' \
+ | ptxas -arch=sm_90 - -o /dev/null # CUDA 13.2: "Unsupported .version 9.3; current version is '9.2'"
+ ```
+ 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.2", archive URL
+ `docs.nvidia.com/cuda/archive/13.2.0/parallel-thread-execution/`). The
+ live `docs.nvidia.com/cuda/parallel-thread-execution/` page is the newest
+ version and its numbering differs; every ISA version stays available
+ under `docs.nvidia.com/cuda/archive/<cuda version>/`.
+- 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 PTX ISA 9.3 (and later)
+
+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`. Until then, a 9.3 form cannot be
+ registered: certification would fail, and a `check()` written against an
+ older ptxas would silently delete the coverage later.
+2. **Retarget the citations.** Update the module docstring's version
+ statement and archive URL, then renumber. Between 9.2 and 9.3 the shifts
+ are: chapter `9.7.10 Fabric Instructions` inserted (everything from
+ Texture onward moves +1: sync `9.7.13→9.7.14`, mma `9.7.14→9.7.15`, wgmma
+ `9.7.15→9.7.16`, tcgen05 `9.7.16→9.7.17`, misc `9.7.19→9.7.20`);
+ `9.7.1.5 clmad` inserted (integer instructions from `mul24` on move +1);
+ `9.7.9.13 multimem.st.async` inserted (data movement from `st.bulk` on
+ move +1); `9.7.14.8 multimem.red.async` inserted (sync instructions from
+ `vote` on move +1); the mbarrier chapter gains three descriptive
+ subsections (Layouts, Tracking successful completion, Report-on), so its
+ instruction subsections move +3 (`mbarrier.arrive` `9.7.13.15.13 →
+ 9.7.14.16.16`) and `mbarrier.check_layout` is appended; global table
+ numbers shift (the tcgen05.ld/st register-count tables are 49/50 in 9.2
+ and 52/53 in 9.3).
+ Renumber intra-chapter subsections too, not just chapters, and re-run the
+ audit of §4 afterwards — a chapter-only shift was the exact mistake made
+ the first time this table changed versions.
+3. **Register the newly legal forms.** Per the 9.3 release notes (ISA
+ chapter 13.1) and the per-section "introduced in PTX ISA version 9.3"
+ notes, the forms this table currently leaves out because they are 9.3 are:
+ - `ld.mmio` with `.acquire`, `st.mmio` with `.release` (`_check_ld` /
+ `_check_st`: "PTX ISA 9.2 spells only relaxed");
+ - `clmad` (new integer family, 9.7.1 banner note);
+ - `multimem.st.async`, `multimem.red.async` (new families, 9.7.9 / sync
+ chapter);
+ - `.sem`/`.scope` on `cp.async.bulk`, `cp.reduce.async.bulk`,
+ `multimem.cp.async.bulk`, `multimem.cp.reduce.async.bulk` (new slots
+ on the cp.async.bulk entries);
+ - mbarrier: `.phase_type::*` on `test_wait`/`try_wait`, the
+ `waitComplete|reportPredicate{, reportValue}` report forms (new
+ entries — a second destination is a new shape), the `.layout` qualifier
+ on `mbarrier.init`/`pending_count`, and `mbarrier.check_layout` (new
+ family);
+ - `fence.proxy.to_proxykind::from_proxykind_fabric.alias.sem_fabric.sys`
+ (two new slots on the fence family; sm_100);
+ - the `fabric.*` instructions (a whole new chapter; new families).
+ Each one: read the section, apply §2, then §3 including full
+ certification at the right `cert_arch`.
+4. **Do not widen operand carriers from the ISA's relaxed-typing tables
+ without certifying.** ISA §9.4.1 Tables 27/28 are an upper bound and say
+ so ("some combinations may still be invalid for a particular
+ instruction"); ptxas rejects, for example, any `cvt` that pairs a `.bf16`
+ operand with a wider register on the other side, and 128-bit carriers on
+ several `.ftz` conversions. The `_cvt_dst_dtypes` / `_cvt_src_dtypes`
+ docstrings in `table.py` record the measured gaps; a version bump must
+ re-run certification and re-measure them, because gaps can close.
+5. Regenerate the stub, update the variant count, run the full suite, and
+ run the §4 audit against the new document before merging.
diff --git a/python/tvm/backend/cuda/ptx/engine.py b/python/tvm/backend/cuda/ptx/engine.py
index 2b5069491f..74042b62c8 100644
--- a/python/tvm/backend/cuda/ptx/engine.py
+++ b/python/tvm/backend/cuda/ptx/engine.py
@@ -228,12 +228,15 @@ def _make_codegen(entry: InstructionEntry):
# this point, then are baked into the instruction text and NOT
# forwarded to the helper (which has no parameter for them).
imm_at = {i for slot, i, _ in layout if slot.kind == "imm"}
- imm_values = [rest[at[i]] for i in sorted(imm_at)]
+ imm_layout = [(slot, i) for slot, i, _ in layout if slot.kind == "imm"]
+ imm_values = [rest[at[i]] for _, i in imm_layout]
if any(not isinstance(value, IntImm) for value in imm_values):
raise ValueError(
f"{entry.name}: immediate operands must become compile-time constants "
"before CUDA codegen (use an explicitly-unrolled loop)"
)
+ for (slot, _), value in zip(imm_layout, imm_values, strict=True):
+ _validate_imm(entry, slot, int(value), mod_map)
imms = tuple(str(int(value)) for value in imm_values)
# ``tirx.ptx.addr`` is an expression only in the outer PTX call's IR.
# The helper still receives the coerced base, while the signed byte
@@ -413,7 +416,7 @@ def _coerce_operand(entry, slot, values, mod_map):
if is_pred:
return _coerce_pred_operand(entry, slot, values)
if slot.kind == "imm":
- return [_coerce_imm(entry, slot, v) for v in values]
+ return [_coerce_imm(entry, slot, v, mod_map) for v in values]
if slot.kind in ("addr", "ptr"):
return [_coerce_address(entry, slot, v, mod_map) for v in values]
return _coerce_typed(entry, slot, values, mod_map)
@@ -576,7 +579,15 @@ def _coerce_address(entry, slot, value, mod_map):
raise ValueError(f"{entry.name}: operand '{slot.name}' must be a pointer or uint64 handle")
-def _coerce_imm(entry, slot, value):
+def _validate_imm(entry, slot, value, mod_map):
+ if entry.imm_check is None:
+ return
+ error = entry.imm_check(mod_map, slot.name, value)
+ if error:
+ raise ValueError(f"{entry.name}: {error}")
+
+
+def _coerce_imm(entry, slot, value, mod_map):
"""A caller-passed immediate: a compile-time constant.
The value lands in the instruction *text* -- the ISA gives these operands
@@ -584,8 +595,9 @@ def _coerce_imm(entry, slot, value):
mismatch" to a register there) -- so a runtime expression has nothing to
lower to and is rejected outright rather than silently materialized.
With `choices` the value is additionally checked against the declared set;
- an open slot declares no domain because the ISA declares none, so any
- constant passes.
+ an open slot declares no enumerable domain. An instruction may still attach
+ an entry-local validator for a compact semantic constraint such as lop3's
+ one-byte LUT.
"""
if isinstance(value, IntImm):
value = value.value
@@ -608,6 +620,7 @@ def _coerce_imm(entry, slot, value):
f"{entry.name}: operand '{slot.name}' is an immediate in the instruction "
f"text; it needs a compile-time integer constant, got {type(value).__name__}"
)
+ _validate_imm(entry, slot, value, mod_map)
if slot.choices is not None and str(value) not in slot.choices:
raise ValueError(
f"{entry.name}: operand '{slot.name}' must be one of "
@@ -670,8 +683,10 @@ def _emit(entry, filled, operands, pred=None, preserve_dst=False):
sunk.add(i + lane)
if lanes and all(i + lane in sunk for lane in range(lanes)):
# ISA 9.7.9.4 states it for mov ("provided that at least one
- # element is a scalar register"); an instruction whose every
- # destination is discarded has nothing left to do anyway.
+ # element is a scalar register"). This API is deliberately the
+ # partial-lane sink facility; whole-operand bit buckets such as
+ # atom's `_` keep their memory side effect and are registered as
+ # fixed-literal siblings selected by their shorter arity.
raise ValueError(
f"{entry.name}: at least one lane of '{slot.name}' must be a real register"
)
diff --git a/python/tvm/backend/cuda/ptx/gen_stubs.py b/python/tvm/backend/cuda/ptx/gen_stubs.py
index 85b2526033..90865aec86 100644
--- a/python/tvm/backend/cuda/ptx/gen_stubs.py
+++ b/python/tvm/backend/cuda/ptx/gen_stubs.py
@@ -170,9 +170,11 @@ def generate() -> str:
out.append(_chain_class(family, families[family]))
out.append("")
out.append("class _PTX:")
+ out.append(" SINK: Any")
for family in sorted(families):
out.append(f" {escape_token(family)}: _Chain_{family}")
out.append(" def addr(self, base: Any, byte_offset: Any) -> Any: ...")
+ out.append(" def pred(self, value: Any) -> Any: ...")
out.append(" def __getitem__(self, text: str) -> Any: ...")
out.append("")
out.append("ptx: _PTX")
diff --git a/python/tvm/backend/cuda/ptx/render.py b/python/tvm/backend/cuda/ptx/render.py
index b2b2f70846..e1a71087e1 100644
--- a/python/tvm/backend/cuda/ptx/render.py
+++ b/python/tvm/backend/cuda/ptx/render.py
@@ -94,12 +94,12 @@ class Bridge(NamedTuple):
"""How one PTX register class that inline asm cannot bind is reached.
``C_BINDING`` above answers "what C type carries this value"; this answers
- the other half, "what does the instruction actually name". Two ISA types
- need it, for the same reason from opposite ends of the constraint
- alphabet: ``.pred`` has no letter at all, and ``.b8`` is below the
- narrowest one (``"h"``). Both are reached the way hand-written PTX reaches
- them -- declare a register of the real class inside the asm block, and
- convert between it and the bound carrier.
+ the other half, "what does the instruction actually name". The bridged
+ classes sit outside the constraint alphabet: ``.pred`` has no letter at
+ all, and the byte registers are below the narrowest constraint (``"h"``).
+ They are reached the way hand-written PTX reaches them -- declare a
+ register of the real class inside the asm block, and convert between it
+ and the bound carrier.
``read``/``write`` are the local register's name (``{slot}`` = the operand
name, ``{n}`` = how many of this direction the block already declared).
@@ -126,7 +126,7 @@ BRIDGE = {
"pred": Bridge(
".pred", "ps{n}", "pd{n}", "setp.ne.b32 {reg}, %{idx}, 0;", "selp.b32 %{idx}, 1, 0, {reg};"
),
- # `.e2m1x2`: the ISA types this operand .b8 (9.7.9.22:92 for the
+ # `.e2m1x2`: the ISA types this operand .b8 (9.7.9.21:92 for the
# destination, :101 for the source). Its general prose says a wider
# register may be used and names no exception for e2m1x2 (:476-486), but
# the toolchain disagrees, so the width here is measured, not read:
@@ -144,6 +144,12 @@ BRIDGE = {
"e2m1x2": Bridge(
".b8", "raw_{slot}", "raw_{slot}", "cvt.u8.u16 {reg}, %{idx};", "cvt.u16.u8 %{idx}, {reg};"
),
+ # st.async.release's b8/u8/s8 sources share one private .b8 staging
+ # register. Stores preserve the low eight bits; the instruction suffix
+ # still supplies the public signed/unsigned interpretation.
+ "st_async_b8reg": Bridge(
+ ".b8", "raw_{slot}", "raw_{slot}", "cvt.u8.u16 {reg}, %{idx};", "cvt.u16.u8 %{idx}, {reg};"
+ ),
}
diff --git a/python/tvm/backend/cuda/ptx/table.py b/python/tvm/backend/cuda/ptx/table.py
index fb15a4ca16..b799f71636 100644
--- a/python/tvm/backend/cuda/ptx/table.py
+++ b/python/tvm/backend/cuda/ptx/table.py
@@ -20,6 +20,12 @@ Pure data + pure functions: this module deliberately imports nothing from
``tvm`` so the thin generators (``gen_stubs``, ``gen_coverage``,
``gen_helpers``) can load it standalone.
+Section and table numbers cite PTX ISA 9.2:
+``https://docs.nvidia.com/cuda/archive/13.2.0/parallel-thread-execution/index.html``.
+That is the version accepted by CUDA 13.2 ptxas. PTX ISA 9.3 added or
+reorganized several sections, including Fabric Instructions, ``clmad``,
+``multimem.st.async``, and ``multimem.red.async``, so its numbering differs.
+
The converged :class:`InstructionEntry` design:
- ``name`` is a single identifier-safe token. Multi-token PTX mnemonics
@@ -55,7 +61,7 @@ import itertools
import keyword
import re
from collections.abc import Callable
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
# PTX operand type token -> the TVM dtypes it accepts, canonical first.
#
@@ -103,7 +109,7 @@ PTX_TYPE_DTYPES = {
# in one 64-bit register. It names a packed layout whose container happens
# to be .b64, not a bit container, so it is not widened.
"f32x2": ("uint64",),
- # cvt's narrow packed formats, per ISA 9.7.9.22's operand table. Each names
+ # cvt's narrow packed formats, per ISA 9.7.9.21's operand table. Each names
# a lane layout, not a container: two (or four) sub-byte or 8-bit elements
# ride one integer register, exactly as `.f16x2` does. The carrier width is
# what the register must be, so it is the dtype the operand binds.
@@ -118,6 +124,8 @@ PTX_TYPE_DTYPES = {
"e5m2x4": ("uint32",),
"e2m3x4": ("uint32",),
"e3m2x4": ("uint32",),
+ # `.tf32` is not packed: ISA 5.2.3 makes it a single 32-bit format whose
+ # register variable must be declared with `.b32` type.
"tf32": ("uint32",),
# `.e2m1x2` is the one that does not fit: the ISA types it .b8, and ptxas
# rejects a wider register in that position, so it can only be reached
@@ -125,6 +133,11 @@ PTX_TYPE_DTYPES = {
# is here so those entries can name it; the uint8 is the C-boundary
# contract, not the register the instruction sees.
"e2m1x2": ("uint8",),
+ # Private carrier for st.async.release's byte source. Inline asm has no
+ # 8-bit constraint, so render.BRIDGE stages the C-boundary value through a
+ # local .b8 register. The entry's dtypes callback preserves the public
+ # b8/u8/s8 domains; this token never reaches the opcode or another family.
+ "st_async_b8reg": ("uint8", "int8"),
# Mixed-precision sources live in a plain 16-bit register -- the ISA's own
# example declares them `.reg .b16`.
"f16": ("uint16",),
@@ -220,9 +233,9 @@ class OperandSlot:
one helper generated per value (the way `cp.async.wait_group N` or
`setmaxnreg`'s nreg exist only as integer literals in the ISA).
- ``dtype`` names the ISA type, and two of those name a register class the
+ ``dtype`` names the ISA type, and a few of those name a register class the
inline-asm constraint alphabet cannot bind (``.pred`` has no letter at
- all, ``.b8`` is below ``"h"``). Those dtypes carry a *bridge* in
+ all, byte registers are below ``"h"``). Those dtypes carry a *bridge* in
:data:`render.BRIDGE`: the value crosses the C boundary in a wider
carrier, and a conversion instruction inside the asm block moves it
between the carrier and a block-local register of the real class. ``rw``
@@ -234,7 +247,7 @@ class OperandSlot:
operands the ISA types by formula rather than by a token in the
instruction text: ``mul.wide``'s destination ("d is twice as wide as a and
b", 9.7.1.3) and ``dp4a``'s accumulator ("c has type .u32 if both .atype
- and .btype are .u32 else .s32", 9.7.1.24). Neither type is spellable as a
+ and .btype are .u32 else .s32", 9.7.1.23). Neither type is spellable as a
slot reference, because neither is written in the instruction -- and a
slot could not hold it either: every written token is rendered into the
opcode, so a `.s64` slot would emit `mul.wide.s32.s64`.
@@ -327,10 +340,12 @@ class OperandSlot:
# the condition per syntax line, and for the 256-bit ld/st it depends on
# the vector width and the element type.
#
- # An always-sunk operand is NOT this: it is a fixed value in the
- # instruction text, i.e. a `kind="imm"` with `literal="_"`, which is how
- # mbarrier.arrive's `state` is registered. This field is for the case where
- # the caller chooses.
+ # A whole-operand bit bucket is NOT this per-lane facility: it is a fixed
+ # value in the instruction text, i.e. a `kind="imm"` with `literal="_"`.
+ # That is how mbarrier.arrive's `state` is registered, and how atom's
+ # caller-selected bit-bucket spelling is represented by a sibling entry
+ # whose shorter arity selects it. This field is for choosing individual
+ # lanes inside a register-group operand.
#
# Unlike every other field this one grants a *per-call* choice rather than
# describing the instruction, so the chosen mask is part of the variant:
@@ -341,6 +356,7 @@ class OperandSlot:
CheckFn = Callable[[dict], str | None]
+ImmCheckFn = Callable[[dict, str, int], str | None]
@dataclass(frozen=True)
@@ -359,6 +375,12 @@ class InstructionEntry:
operands: tuple[OperandSlot, ...]
slots: tuple[ModifierSlot, ...] = ()
check: CheckFn | None = None # cross-slot validation, mod_map -> error | None
+ # Optional validation for caller immediates after they become concrete.
+ # It runs at trace time for literal constants and again after unrolling /
+ # simplification in CUDA codegen. Keeping it on the entry lets an
+ # instruction validate an operand without turning every OperandSlot into
+ # a range-policy schema.
+ imm_check: ImmCheckFn | None = None
# Whether the emitted inline asm carries `volatile`. This is purely a
# C-level optimization barrier — it never changes *which* PTX instruction
# is emitted, only whether nvcc may common up or drop identical calls.
@@ -573,8 +595,10 @@ def sink_combos(entry: InstructionEntry, tokens) -> tuple[frozenset, ...]:
A sink is spelled per element, so the domain is 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"); it is the conservative reading everywhere else, and an
- instruction whose every destination is discarded has nothing left to do.
+ register"), and this per-lane facility keeps the same conservative rule
+ for every register group. A whole-operand bit bucket can still have a side
+ effect -- atom continues to update memory -- so those spellings are fixed
+ ``literal="_"`` sibling entries rather than an all-sunk lane mask.
Empty set first, so the un-sunk variant keeps the helper name it had
before the slot became sinkable.
@@ -683,7 +707,8 @@ def variants(entry: InstructionEntry) -> tuple:
def _check_cache_hint(m):
"""`{.level::cache_hint}` and its `{, cache_policy}` operand, per ISA
- 9.7.9.8/9.7.9.11/9.7.12.x and MEASURED on ptxas 13.2 at -arch=sm_90.
+ 9.7.9.8/9.7.9.9/9.7.9.11/9.7.13.5/9.7.13.6 and MEASURED on ptxas 13.2
+ at -arch=sm_90.
The qualifier is an L2 policy, so it is spelled only on lines that can
reach L2: `.global` or generic addressing. `.local`, `.shared` and the
@@ -727,9 +752,7 @@ def _check_ld(m):
return "only ld.relaxed/ld.acquire take a scope"
if mmio:
# "ld.mmio.sem.sys{.global}": "Only .sys thread scope is valid";
- # global or generic addressing only. The ISA also allows .acquire
- # (PTX ISA 9.3+), but the current toolchain assembles 9.2 and ptxas
- # rejects it — widen when the toolchain catches up.
+ # global or generic addressing only. PTX ISA 9.2 spells only relaxed.
if sem != "relaxed":
return "ld.mmio requires .relaxed"
if scope != "sys":
@@ -892,9 +915,7 @@ def _check_st(m):
return "only st.relaxed/st.release take a scope"
if mmio:
# "st.mmio.sem.sys{.global}": "Only .sys thread scope is valid for the
- # st.mmio operation." .release with .mmio arrives in PTX ISA 9.3; the
- # toolchain here assembles 9.2 and ptxas rejects it, so keep .relaxed
- # only and widen when the toolchain catches up.
+ # st.mmio operation." PTX ISA 9.2 spells only relaxed.
if sem != "relaxed":
return "st.mmio requires .relaxed"
if scope != "sys":
@@ -982,7 +1003,7 @@ def _check_div_f(m):
# The .type line shared by most of PTX ISA 9.7.1: mul/mad (9.7.1.3/4), sad
-# (9.7.1.8), div (9.7.1.9) and rem (9.7.1.10) all spell exactly this set.
+# (9.7.1.7), div (9.7.1.8) and rem (9.7.1.9) all spell exactly this set.
_INT_TYPES = ("u16", "u32", "u64", "s16", "s32", "s64")
# `.wide`'s result type (ISA 9.7.1.3: "If .wide is specified, then d is twice
@@ -998,15 +1019,127 @@ def _wide_dtype(m):
return _WIDE_RESULT[m["type"]]
-def _ld_dst_dtypes(m):
- """Scalar ``ld.s32`` may sign-extend into a wider destination register."""
- if m["type"] == "s32":
- return ("int32", "int64")
- return PTX_TYPE_DTYPES[m["type"]]
+_RELAXED_MEM_DTYPES = {
+ # PTX ISA 9.4.1, Tables 27/28. `ld`, `st`, and `ldu` accept wider data
+ # registers. Bit-size instructions accept any fundamental register class;
+ # integer instructions accept bit/integer classes; floating instructions
+ # accept their native class or a bit carrier. Canonical dtypes stay first
+ # so existing helper names and signatures remain stable.
+ "b8": (
+ "uint8",
+ "int8",
+ "uint16",
+ "int16",
+ "float16",
+ "bfloat16",
+ "uint32",
+ "int32",
+ "float32",
+ "uint64",
+ "int64",
+ "float64",
+ "uint128",
+ "int128",
+ ),
+ "b16": (
+ "uint16",
+ "int16",
+ "float16",
+ "bfloat16",
+ "uint32",
+ "int32",
+ "float32",
+ "uint64",
+ "int64",
+ "float64",
+ "uint128",
+ "int128",
+ ),
+ "b32": (
+ "uint32",
+ "int32",
+ "float32",
+ "uint64",
+ "int64",
+ "float64",
+ "uint128",
+ "int128",
+ ),
+ "b64": ("uint64", "int64", "float64", "uint128", "int128"),
+ "b128": ("uint128", "int128"),
+ "u8": (
+ "uint8",
+ "int8",
+ "uint16",
+ "int16",
+ "uint32",
+ "int32",
+ "uint64",
+ "int64",
+ "uint128",
+ "int128",
+ ),
+ "s8": (
+ "int8",
+ "uint8",
+ "int16",
+ "uint16",
+ "int32",
+ "uint32",
+ "int64",
+ "uint64",
+ "int128",
+ "uint128",
+ ),
+ "u16": (
+ "uint16",
+ "int16",
+ "uint32",
+ "int32",
+ "uint64",
+ "int64",
+ "uint128",
+ "int128",
+ ),
+ "s16": (
+ "int16",
+ "uint16",
+ "int32",
+ "uint32",
+ "int64",
+ "uint64",
+ "int128",
+ "uint128",
+ ),
+ "u32": ("uint32", "int32", "uint64", "int64", "uint128", "int128"),
+ "s32": ("int32", "uint32", "int64", "uint64", "int128", "uint128"),
+ "u64": ("uint64", "int64", "uint128", "int128"),
+ "s64": ("int64", "uint64", "int128", "uint128"),
+ "f32": ("float32", "uint32", "int32", "uint64", "int64", "uint128", "int128"),
+ "f64": ("float64", "uint64", "int64", "uint128", "int128"),
+}
+
+
+def _relaxed_mem_dtypes(m):
+ """Relaxed carriers for scalar and at-most-128-bit ``ld``, ``st``, and ``ldu``."""
+ return _RELAXED_MEM_DTYPES[m["type"]]
+
+
+def _st_vec_dtypes(m):
+ """Relaxed carriers accepted by CUDA 13.2 for the <=128-bit vector ``st`` lines."""
+ dtypes = _RELAXED_MEM_DTYPES[m["type"]]
+ # PTX ISA 9.2 section 9.4.1 permits a wider bit register, but CUDA 13.2
+ # ptxas reports C7907 for an otherwise minimal `st.v2.f64` whose source
+ # lanes use 128-bit carriers. The same carriers assemble for scalar st,
+ # ld/ldu and their <=128-bit vector forms, so keep the toolchain exception
+ # local to this one source-operand shape.
+ if m["vec"] == "v2" and m["type"] == "f64":
+ return tuple(dtype for dtype in dtypes if dtype not in ("uint128", "int128"))
+ return dtypes
def _dp_acc_dtype(m):
- """dp4a/dp2a's accumulator type (ISA 9.7.1.24/9.7.1.25).
+ """dp4a/dp2a's accumulator type (ISA 9.7.1.23/9.7.1.24).
"Operand c has type .u32 if both .atype and .btype are .u32 else operand c
has type .s32" -- and d accumulates into the same type. It is a function of
@@ -1015,7 +1148,7 @@ def _dp_acc_dtype(m):
return "u32" if (m["atype"], m["btype"]) == ("u32", "u32") else "s32"
-def _check_int_addsub(m):
+def _check_int_addsub(m, mnemonic):
"""Which types the integer add/sub lines take `.sat` on (ISA 9.7.1.1/9.7.1.2).
add.type1 d, a, b; .type1 = {.u16, .u64, .s16, .s64}
@@ -1023,25 +1156,39 @@ def _check_int_addsub(m):
sub.type1 d, a, b; .type1 = {.u16, .u32, .u64, .s16, .s64}
sub{.sat}.type2 d, a, b; .type2 = {.s32, .u8x4, .s8x4}
- Of the saturating forms only `.s32` is registered: `add.sat` on
+ Of the saturating forms only `.s32` is registered. `add.sat` on
.u32/.u16x2/.s16x2 arrived in PTX ISA 9.2 as "sm_120f or higher in the same
- family", and the `.u8x4`/`.s8x4` types are excluded outright for the same
- reason (see the `_MINMAX_INT_PLAIN` note below). `add.sat.s32` is the PTX
- 1.0 form and is supported everywhere.
+ family". The remaining types in these entries belong to syntax lines that
+ never carry `.sat`: .u16/.u64/.s16/.s64 for add, and
+ .u16/.u32/.u64/.s16/.s64 for sub. The `.u8x4`/`.s8x4` types are excluded
+ outright because they are sm_120f-only (see `_MINMAX_INT_PLAIN` below).
+ `add.sat.s32` and `sub.sat.s32` are the PTX 1.0 forms.
"""
- if m["sat"] and m["type"] != "s32":
+ if not m["sat"] or m["type"] == "s32":
+ return None
+ if mnemonic == "add" and m["type"] in ("u32", "u16x2", "s16x2"):
return (
- f".sat is registered on .s32 only; .sat.{m['type']} is an sm_120f-only form "
- f"(PTX ISA 9.2)"
+ f"add.sat.{m['type']} is an sm_120f-only form (PTX ISA 9.2); "
+ ".sat is registered on .s32 only"
)
- return None
+ return f".sat is not on the {mnemonic}.{m['type']} syntax line"
+
+
+def _check_int_add(m):
+ """Apply the integer add syntax-line split documented by `_check_int_addsub`."""
+ return _check_int_addsub(m, "add")
+
+
+def _check_int_sub(m):
+ """Apply the integer sub syntax-line split documented by `_check_int_addsub`."""
+ return _check_int_addsub(m, "sub")
def _check_int_mad(m):
"""`.sat` on the multiply-add lines: `.hi` mode, `.s32` type, nothing else.
Both lines spell it as a syntax line of its own -- `mad.hi.sat.s32 d, a, b,
- c;` (ISA 9.7.1.4) and `mad24.hi.sat.s32 d, a, b, c;` (9.7.1.7) -- with the
+ c;` (ISA 9.7.1.4) and `mad24.hi.sat.s32 d, a, b, c;` (9.7.1.6) -- with the
Notes repeating "Applies only to .s32 type in .hi mode".
"""
if m["sat"] and (m["mode"], m["type"]) != ("hi", "s32"):
@@ -1049,7 +1196,7 @@ def _check_int_mad(m):
return None
-# The .type tokens of the integer min/max lines (PTX ISA 9.7.1.13/9.7.1.14).
+# The .type tokens of the integer min/max lines (PTX ISA 9.7.1.12/9.7.1.13).
# `.relu` is only on the signed line (type2); the unsigned/other line takes none.
# NOT REGISTERED: `.u8x4` and `{.relu}.s8x4`, which the ISA supports only on
# "sm_120f or higher in the same family" -- outside the architectures this
@@ -1064,8 +1211,8 @@ def _check_mbarrier_sem_scope(m):
"""`.sem` and `.scope` are one qualifier pair: both or neither.
Stated in so many words by every mbarrier section that has the pair --
- "Qualifiers .sem and .scope must be specified together." (ISA 9.7.14.16.14
- expect_tx, .15 complete_tx, .16 arrive, .17 arrive_drop, .19 test_wait /
+ "Qualifiers .sem and .scope must be specified together." (ISA 9.7.13.15.11
+ expect_tx, .12 complete_tx, .13 arrive, .14 arrive_drop, .16 test_wait /
try_wait). The rule is what the sections say, not how they spell it: some
lines write the pair as one `{.sem.scope}` group and others as two
`{.sem}{.scope}` groups, and `mbarrier.arrive.noComplete` fixes it to
@@ -1080,17 +1227,55 @@ def _check_fence_proxy(m):
"""`.proxykind = { .alias, .async, .async.global, .async.shared::{cta,cluster} }`.
Modelled as proxykind + an optional space so the surface can walk it one
- attribute at a time; only `.async` carries a state space (ISA 9.7.14.4).
+ attribute at a time; only `.async` carries a state space (ISA 9.7.13.4).
"""
if m["space"] and m["proxykind"] != "async":
return f"fence.proxy.{m['proxykind']} takes no state space"
return None
+_CP_REDUCE_BULK_TYPES = {
+ "shared::cluster": {
+ "add": ("u32", "s32", "u64"),
+ "min": ("u32", "s32"),
+ "max": ("u32", "s32"),
+ "inc": ("u32",),
+ "dec": ("u32",),
+ "and": ("b32",),
+ "or": ("b32",),
+ "xor": ("b32",),
+ },
+ "global": {
+ "add": ("u32", "s32", "u64", "f32", "f64", "f16", "bf16"),
+ "min": ("u32", "s32", "u64", "s64", "f16", "bf16"),
+ "max": ("u32", "s32", "u64", "s64", "f16", "bf16"),
+ "inc": ("u32",),
+ "dec": ("u32",),
+ "and": ("b32", "b64"),
+ "or": ("b32", "b64"),
+ "xor": ("b32", "b64"),
+ },
+}
+
+
+def _check_cp_reduce_async_bulk(m):
+ """Exact redOp/type grid of non-tensor cp.reduce.async.bulk (ISA 9.7.9.25.4.2)."""
+ dst, op, ty = m["dst"], m["redop"], m["type"]
+ if ty not in _CP_REDUCE_BULK_TYPES[dst][op]:
+ allowed = " / ".join("." + item for item in _CP_REDUCE_BULK_TYPES[dst][op])
+ return f".{op} to .{dst} takes {allowed}"
+ needs_noftz = dst == "global" and op == "add" and ty in ("f16", "bf16")
+ if bool(m.get("noftz", "")) != needs_noftz:
+ if needs_noftz:
+ return f"cp.reduce.async.bulk.add.{ty} requires .noftz"
+ return ".noftz applies only to .add.f16 and .add.bf16 into .global"
+ return None
+
+
def _check_minmax(m):
"""Which qualifiers each min/max syntax line allows.
- Integer lines (9.7.1.13/14): op.type1 d,a,b | op{.relu}.type2 d,a,b
+ Integer lines (9.7.1.12/13): op.type1 d,a,b | op{.relu}.type2 d,a,b
Floating lines (9.7.3.11/12): op{.ftz}{.NaN}{.xorsign.abs}.f32 d,a,b
op.f64 d,a,b
Half lines (9.7.4.7/8): op{.ftz}{.NaN}{.xorsign.abs}.f16{x2} d,a,b
@@ -1292,8 +1477,8 @@ _LOGIC_TYPES = ("pred", *_BIT_TYPES)
# bit-size types are required, we include the arithmetic types for the
# programmer's convenience".
_MOV_TYPES = ("pred", "b16", "b32", "b64", "u16", "u32", "u64", "s16", "s32", "s64", "f32", "f64")
-# The state spaces cvta converts between and isspacep queries (ISA 9.7.9.21,
-# 9.7.9.20) -- the same eight, spelled the same way, in both instructions.
+# The state spaces cvta converts between and isspacep queries (ISA 9.7.9.20,
+# 9.7.9.19) -- the same eight, spelled the same way, in both instructions.
_CVTA_SPACES = (
"const", "global", "local", "shared", "shared::cta", "shared::cluster", "param", "param::entry",
) # fmt: skip
@@ -1348,7 +1533,7 @@ def _check_half_set(m):
return f"a .{dt} destination takes no .ftz"
if st not in _FTZ_SET_STYPES:
return f".ftz flushes subnormal inputs, so a .{st} source does not take it"
- if st.startswith("b"):
+ if st in _BIT_TYPES:
if op not in ("eq", "ne"):
return f"the bit-size source .{st} compares only with eq/ne"
elif st[0] in "us":
@@ -1370,24 +1555,36 @@ def _check_half_setp(m):
return None
-# --- PTX ISA 9.7.9.12 (st.async) and 9.7.9.15 (multimem) --------------------
+# --- PTX ISA 9.7.9.12 (st.async) and 9.7.9.14 (multimem) --------------------
# The vector lines of st.async: "`.v2` is supported with .b32, .b64, .s32,
# .s64, .u32, .u64, .f32 and .f64 types. `.v4` qualifier is supported with
# .b32, .s32, .u32 and .f32 types."
_ST_ASYNC_TYPES = ("b32", "b64", "b128", "u32", "u64", "s32", "s64", "f32", "f64")
_ST_ASYNC_V2 = ("b32", "b64", "u32", "u64", "s32", "s64", "f32", "f64")
_ST_ASYNC_V4 = ("b32", "u32", "s32", "f32")
-# The release line's type list (9.7.9.12, second syntax block), which reaches
-# below the 32-bit floor of the mbarrier line.
-# MEASURED: the ISA writes `.b8`, `.u8` and `.s8` into that list, but ptxas
-# 13.2 rejects all three at sm_100 -- with and without `.mmio`, at either scope
-# -- while every 16-bit and wider type assembles. So the line starts at 16 bits
-# here, and the byte forms wait for a toolchain that accepts them.
+# The release line's complete type list (9.7.9.12, second syntax block), which
+# reaches below the 32-bit floor of the mbarrier line. The byte forms need an
+# exact local .b8 source register: inline asm has no 8-bit constraint,
+# and binding their C values directly through the usual 16-bit carrier makes
+# ptxas reject the instruction. The private bridge is selected only here;
+# the separate dtypes callback retains each suffix's public type domain.
_ST_ASYNC_REL_TYPES = (
- "b16", "b32", "b64", "u16", "u32", "u64", "s16", "s32", "s64", "f32", "f64",
+ "b8", "b16", "b32", "b64", "u8", "u16", "u32", "u64",
+ "s8", "s16", "s32", "s64", "f32", "f64",
) # fmt: skip
-# createpolicy's `.level::primary_priority` line (ISA 9.7.9.19). Wider than
+
+def _st_async_rel_operand_type(m):
+ """Private byte-register bridge for ``st.async.release``'s 8-bit source."""
+ return "st_async_b8reg" if m["type"] in ("b8", "u8", "s8") else m["type"]
+
+
+def _st_async_rel_operand_dtypes(m):
+ """Public dtype domain remains the one named by the instruction suffix."""
+ return PTX_TYPE_DTYPES[m["type"]]
+
+
+# createpolicy's `.level::primary_priority` line (ISA 9.7.9.18). Wider than
# the `_L2_EVICT` set the ld/st eviction hints use: this instruction is where
# a priority is *created*, so it spells all four.
_CACHE_PRIORITIES = (
@@ -1418,17 +1615,16 @@ _MM_VEC_TYPES = {
def _check_multimem_sem(m):
- """How `.sem` and `.scope` go together on every multimem line (ISA 9.7.9.15).
-
- The ISA writes two syntax lines per mnemonic -- `{.sem}{.scope}` and a
- `.weak` one with no scope position -- which reads as "either may be
- omitted". ptxas is stricter, and its rule is the one that makes the two
- lines distinct: everywhere but `.weak`, the pair is the same both-or-
- neither rule the mbarrier sections state in words, so that check is what
- decides it here. ptxas says both halves -- "Modifier '.relaxed' requires
- scope with 'multimem.st' instruction" and "Modifier '.cta' requires order
- with 'multimem.st' instruction". What multimem adds is `.weak`: the line
- with no scope position at all.
+ """How `.sem` and `.scope` go together on every multimem line (ISA 9.7.9.14).
+
+ The ISA writes two syntax lines for `multimem.ld_reduce` and `multimem.st`:
+ `{.sem}{.scope}` and a `.weak` line with no scope position. `multimem.red`
+ has one line whose `.redsem` is `{.relaxed, .release}`, with no `.weak`
+ form. ptxas makes the non-weak semantic/scope pair the same both-or-neither
+ rule the mbarrier sections state in words, so that check decides it here.
+ It diagnoses both halves: "Modifier '.relaxed' requires scope with
+ 'multimem.st' instruction" and "Modifier '.cta' requires order with
+ 'multimem.st' instruction".
"""
if m.get("sem", "") == "weak":
return ".weak is a syntax line of its own and takes no scope" if m["scope"] else None
@@ -1436,7 +1632,7 @@ def _check_multimem_sem(m):
def _check_multimem_int(m):
- """op x type on the integer multimem lines (ISA 9.7.9.15).
+ """op x type on the integer multimem lines (ISA 9.7.9.14).
The `.type` line in the Syntax block is the union over ops; the pairing is
the "valid combinations of .op and base type" table, and `.add` is the row
@@ -1512,19 +1708,20 @@ def _check_st_async_rel(m):
return None
-# --- PTX ISA 9.7.14 warp-level and async reductions -------------------------
-# red.async's four mbarrier lines (9.7.14.7), one op group each.
+# --- PTX ISA 9.7.13 warp-level and async reductions -------------------------
+# red.async's four mbarrier lines (9.7.13.7), one op group each.
_RED_ASYNC_OPS = {"inc": ("u32",), "dec": ("u32",), "min": ("u32", "s32"),
"max": ("u32", "s32"), "and": ("b32",), "or": ("b32",), "xor": ("b32",),
"add": ("u32", "s32", "u64")} # fmt: skip
-# atom's two vector lines (9.7.14.5). A half element rides `.v2`/`.v4`/`.v8`;
-# a packed pair, being twice as wide, stops at `.v4`; and `.f32` has only the
-# `.add` line. Every cell probed.
+# atom/red's three vector lines (9.7.13.5 / 9.7.13.6): `.f32`, a half-word,
+# and a packed pair. The latter two share one table entry per mnemonic; a half
+# element rides `.v2`/`.v4`/`.v8`, while a packed pair stops at `.v4`.
+# `.f32` has only the `.add` line. Every cell probed.
_ATOM_VEC_HALF = ("f16", "bf16")
def _check_red_async(m):
- """op x type on red.async's mbarrier lines (ISA 9.7.14.7).
+ """op x type on red.async's mbarrier lines (ISA 9.7.13.7).
The ISA writes one syntax line per op group rather than a table, so the
grid is read off those four lines: increment/decrement on the unsigned
@@ -1537,8 +1734,8 @@ def _check_red_async(m):
return None
-def _check_atom_vec(m):
- """`.vec` x type on atom's vector lines (ISA 9.7.14.5).
+def _check_atomic_vec(m):
+ """`.vec` x type on atom/red's vector lines (ISA 9.7.13.5 / 9.7.13.6).
atom{.sem}{.scope}{.global}.add{.cache}.vec_32_bit.f32 d, [a], b;
atom{...}.op.noftz{.cache}.vec_16_bit.half_word_type d, [a], b;
@@ -1555,6 +1752,32 @@ def _check_atom_vec(m):
return None
+def _check_atom_bitbucket_bf16(m):
+ """Withhold the documented bf16 atom bit-bucket forms from this backend.
+
+ PTX ISA 9.2 permits ``_`` as the destination of a simple ``atom``
+ reduction, including the scalar and vector bf16 lines. CUDA 13.2 ptxas
+ does not compile those forms: an exact force-inlined kernel probe
+ containing ``atom...bf16 _`` terminates ptxas with a segmentation fault.
+ The corresponding returned-value bf16 form and f16 bit-bucket form both
+ compile, so this is deliberately limited to bf16/bf16x2 bit buckets.
+ """
+ if m["type"].startswith("bf16"):
+ return (
+ "PTX ISA 9.2 supports a bf16 atom bit-bucket destination, but CUDA 13.2 "
+ "ptxas does not compile it"
+ )
+ return None
+
+
+def _check_atom_half_bitbucket(m):
+ return _check_cache_hint(m) or _check_atom_bitbucket_bf16(m)
+
+
+def _check_atom_vec_half_bitbucket(m):
+ return _check_atomic_vec(m) or _check_atom_bitbucket_bf16(m)
+
+
def _check_slct(m):
"""slct's two lines (ISA 9.7.6.4), which differ only in the selector type.
@@ -1569,6 +1792,30 @@ def _check_slct(m):
return None
+_SLCT_VALUE_DTYPES = {
+ # d/a/b are bit values of the first instruction type's width, not numeric
+ # operands of that type. Keep its native TVM dtype first so the canonical
+ # helper names and signatures remain stable; the rest are the carrier
+ # classes accepted by ptxas 13.2 over the full independent d x a x b grid.
+ "b16": ("uint16", "int16", "float16", "bfloat16"),
+ "u16": ("uint16", "int16", "float16", "bfloat16"),
+ "s16": ("int16", "uint16", "float16", "bfloat16"),
+ "b32": ("uint32", "int32", "float32"),
+ "u32": ("uint32", "int32"),
+ "s32": ("int32", "uint32"),
+ "f32": ("float32", "uint32", "int32"),
+ "b64": ("uint64", "int64", "float64"),
+ "u64": ("uint64", "int64"),
+ "s64": ("int64", "uint64"),
+ "f64": ("float64", "uint64", "int64"),
+}
+
+
+def _slct_value_dtypes(m):
+ """TVM carrier domain for slct's independently bit-typed d/a/b operands."""
+ return _SLCT_VALUE_DTYPES[m["dtype"]]
+
+
def _check_absneg(m):
"""The one-source sign lines, `op{.ftz}.f32 d, a;` and `op.f64 d, a;`.
@@ -1606,7 +1853,7 @@ def _check_farith(m):
def _check_prefetch(m):
- """Each prefetch syntax line names exactly one target (PTX ISA 9.7.9.16).
+ """Each prefetch syntax line names exactly one target (PTX ISA 9.7.9.15).
`.level::eviction_priority` stays bound to `.global` on purpose: its syntax
line is `prefetch.global.level::eviction_priority`, with `.global` written
@@ -1639,12 +1886,12 @@ _ATOM_TYPES = ("b32", "b64", "u32", "u64", "s32", "s64", "f32", "f64")
def _check_atomic(m):
- """op x type pairings for atom/red (PTX ISA 9.7.14.5 / 9.7.14.6).
+ """op x type pairings for atom/red (PTX ISA 9.7.13.5 / 9.7.13.6).
- Normative source: ISA Table 35 (atom) and Table 36 (red), which give the
- pairing cell by cell. The `.type = {...}` line in the Syntax block is only
- the union across ops, which is why it cannot be transcribed directly. Half-precision
- types appear in ptxas' message but are excluded from this entry (they need
+ The op/type rules in those sections give the pairings cell by cell. The
+ `.type = {...}` line in the Syntax block is only the union across ops,
+ which is why it cannot be transcribed directly. Half-precision types
+ appear in ptxas' message but are excluded from this entry (they need
.noftz and a half carrier type).
"""
hint = _check_cache_hint(m)
@@ -1716,14 +1963,14 @@ _FRND = ("rn", "rz", "rm", "rp") # .rnd on the floating-point arithmetic lines
def _matrix_num_lanes(m):
- # ISA 9.7.15.5.16 (stmatrix): "a brace-enclosed vector expression consisting
+ # ISA 9.7.14.5.16 (stmatrix): "a brace-enclosed vector expression consisting
# of 1, 2, or 4 32-bit registers as per the value of .num" -- no shape term,
# unlike ldmatrix's .m16n16 doubling.
return int(m["num"][1:])
def _ldmatrix_lanes(m):
- # ISA 9.7.15.5.15: "a brace-enclosed vector expression consisting of 1, 2,
+ # ISA 9.7.14.5.15: "a brace-enclosed vector expression consisting of 1, 2,
# or 4 32-bit registers as per the value of .num" -- and, for shape 16x16,
# "two destination registers r0 and r1 of type .b32 must be specified" per
# matrix, so .m16n16 doubles the count (ptxas: "Vector of size 2 is
@@ -1745,21 +1992,23 @@ def _check_ldmatrix_b8fmt(m):
def _tcgen05_ldst_lanes(m):
- # ISA 9.7.17.8.3 Table 52 / 9.7.17.8.4 Table 53: the register vector holds
- # `.num` x (shape width / 32b) registers, capped at 128.
+ # ISA 9.7.16.8.3 Table 49 / 9.7.16.8.4 Table 50: the register vector holds
+ # `.num` x (rows x shape width / 1024b) registers -- 1 per `.num` for
+ # .16x32bx2/.16x64b/.32x32b, 2 for .16x128b, and 4 for .16x256b -- capped
+ # at 128.
per_num = {"16x64b": 1, "32x32b": 1, "16x128b": 2, "16x256b": 4, "16x32bx2": 1}
return int(m["num"][1:]) * per_num[m["shape"]]
def _check_tcgen05_ldst(m):
- """The Table 52/53 rows marked NA -- the products that exceed 128 registers."""
+ """The Table 49/50 rows marked NA -- the products that exceed 128 registers."""
if _tcgen05_ldst_lanes(m) > 128:
return f"shape {m['shape']} caps .num where the vector would exceed 128 registers"
return None
def _check_tcgen05_cp(m):
- """The shape <-> multicast pairings ISA 9.7.17.9.2 states, and the fmt pair.
+ """The shape <-> multicast pairings ISA 9.7.16.9.2 states, and the fmt pair.
".64x128b requires .warpx2::02_13 or .warpx2::01_23" and ".32x128b
requires .warpx4"; the wider shapes copy to all warps and take no
@@ -1780,15 +2029,16 @@ def _check_tcgen05_cp(m):
return None
-# mma fragment sizes, per the Matrix Fragments tables of ISA 9.7.15.5.1-13.
+# mma fragment sizes, per the Matrix Fragments tables of ISA 9.7.14.5.1-13.
# Each is `rows * cols * bits / threads / 32`, the register count a thread holds
# of an MxN (or MxK / KxN) tile -- the ISA states the tables, this states the
# rule they follow. .m8n8k4 with .f16 multiplicands is the one shape a warp
-# runs as four independent 8-thread MMAs, so its A/B fragments divide by 8.
+# runs as four independent 8-thread MMAs, so all four of its fragments -- A, B
+# and C/D alike -- divide by 8 rather than 32.
_MMA_BITS = {
"f16": 16, "bf16": 16, "tf32": 32, "f32": 32, "f64": 64, "s32": 32,
"u8": 8, "s8": 8, "u4": 4, "s4": 4, "b1": 1,
- "e4m3": 8, "e5m2": 8, "e3m2": 6, "e2m3": 6, "e2m1": 4,
+ "e4m3": 8, "e5m2": 8,
} # fmt: skip
@@ -1804,7 +2054,7 @@ def _mma_regs(dtype, rows, cols, threads, reg_bits=32):
def _mma_threads(m):
"""Threads sharing one tile: 8 on the .f16 .m8n8k4 line, 32 everywhere else.
- ISA 9.7.15.5.14: "A warp executing mma.sync.m8n8k4 instruction computes 4
+ ISA 9.7.14.5.14: "A warp executing mma.sync.m8n8k4 instruction computes 4
matrix multiply and accumulate operations. Rest of the mma.sync operations
compute a single matrix mutliply and accumulate operation per warp." Four
operations to a warp is 8 threads each, so a thread's fragment of that line
@@ -1812,7 +2062,7 @@ def _mma_threads(m):
(d=4, a=2, b=2 for .f16.f16.f16.f16; anything else is "Arguments mismatch").
The division by 8 belongs to that line alone. The .f64 .m8n8k4 line has its
- own fragment section, ISA 9.7.15.5.2, which opens "A warp executing
+ own fragment section, ISA 9.7.14.5.2, which opens "A warp executing
mma.m8n8k4 with .f64 floating point type will compute an MMA operation of
shape .m8n8k4" -- one operation, the whole warp -- and tabulates A and B as
"A vector expression containing a single .f64 register" and C/D as "A
@@ -1851,10 +2101,13 @@ def _mma_lanes(which):
def _mma_sp_lanes(which):
"""Like `_mma_lanes`, but A holds half of K -- the structured-sparse half.
- ISA 9.7.15.6: "For an MxNxK sparse mma.sp{::ordered_metadata} operation,
- the MxK matrix A is packed into MxK/2 elements" -- two of every four along
- K, so the A fragment of an MxK sparse line is the dense fragment of MxK/2.
- The other three groups are unchanged.
+ ISA 9.7.14.6: "For an MxNxK sparse mma.sp{::ordered_metadata} operation,
+ the MxK matrix A is packed into MxK/2 elements" -- 50% zeros per row at a
+ shape- and kind-specific granularity: e2m1 is 2:4 in the
+ f8f6f4/mxf8f6f4 lines and 4:8 only in mxf4/mxf4nvf4; the other forms have
+ the ratios stated by their own syntax line. Thus the stored A fragment of
+ an MxK sparse line is the dense fragment of MxK/2; the other three groups
+ are unchanged.
"""
def lanes(m):
@@ -1884,12 +2137,171 @@ _CVT_INT_BITS = {"u8": 8, "s8": 8, "u16": 16, "s16": 16, "u32": 32, "s32": 32, "
_CVT_IRND = ("rni", "rzi", "rmi", "rpi")
_CVT_FRND = ("rn", "rz", "rm", "rp")
+# The generic scalar cvt line follows ISA 9.4.1's relaxed type checking: an
+# integer instruction type may use a same-width or wider integer/bit register,
+# while a floating instruction type may use its native register or a same-width
+# or wider bit register. These are the base TVM carrier spellings, canonical
+# first so the pre-existing helper names remain stable. The cvt section
+# explicitly excludes `.bf16` from widening its own operand, so it stays pinned
+# to its exact 16-bit carrier. Packed/narrow cvt entries do not use these
+# callbacks; their operand prose gives exact carriers and several are stricter
+# than the generic relaxed rule in the current toolchain.
+_CVT_RELAXED_DTYPES = {
+ "u8": (
+ "uint8",
+ "int8",
+ "uint16",
+ "int16",
+ "uint32",
+ "int32",
+ "uint64",
+ "int64",
+ "uint128",
+ "int128",
+ ),
+ "s8": (
+ "int8",
+ "uint8",
+ "int16",
+ "uint16",
+ "int32",
+ "uint32",
+ "int64",
+ "uint64",
+ "int128",
+ "uint128",
+ ),
+ "u16": (
+ "uint16",
+ "int16",
+ "uint32",
+ "int32",
+ "uint64",
+ "int64",
+ "uint128",
+ "int128",
+ ),
+ "s16": (
+ "int16",
+ "uint16",
+ "int32",
+ "uint32",
+ "int64",
+ "uint64",
+ "int128",
+ "uint128",
+ ),
+ "u32": ("uint32", "int32", "uint64", "int64", "uint128", "int128"),
+ "s32": ("int32", "uint32", "int64", "uint64", "int128", "uint128"),
+ "u64": ("uint64", "int64", "uint128", "int128"),
+ "s64": ("int64", "uint64", "int128", "uint128"),
+ "f16": (
+ "uint16",
+ "int16",
+ "float16",
+ "bfloat16",
+ "uint32",
+ "int32",
+ "uint64",
+ "int64",
+ "uint128",
+ "int128",
+ ),
+ "bf16": ("uint16",),
+ "f32": ("float32", "uint32", "int32", "uint64", "int64", "uint128", "int128"),
+ "f64": ("float64", "uint64", "int64", "uint128", "int128"),
+}
+
+# The same-width subset of each relaxed domain. This is intentionally broader
+# than the canonical dtype: for example, a `.f16` operand may ride any 16-bit
+# register class. CUDA 13.2 ptxas needs this subset for a few cvt-specific
+# gaps in the general relaxed type-checking rules, documented at the callbacks
+# below. Packed/narrow cvt entries do not use this table.
+_CVT_SAME_WIDTH_DTYPES = {
+ "u8": ("uint8", "int8"),
+ "s8": ("int8", "uint8"),
+ "u16": ("uint16", "int16"),
+ "s16": ("int16", "uint16"),
+ "u32": ("uint32", "int32"),
+ "s32": ("int32", "uint32"),
+ "u64": ("uint64", "int64"),
+ "s64": ("int64", "uint64"),
+ "f16": ("uint16", "int16", "float16", "bfloat16"),
+ "bf16": ("uint16",),
+ "f32": ("float32", "uint32", "int32"),
+ "f64": ("float64", "uint64", "int64"),
+}
+
+# PTX ISA 9.2, 9.4.1 Table 27 and cvt's own notes permit a floating source
+# operand to use a wider bit register, with the instruction consuming its low
+# bits. CUDA 13.2 ptxas accepts `.version 9.2` and the exact-width form, but a
+# direct probe such as
+#
+# .reg .b16 d;
+# .reg .b64 a;
+# cvt.rn.ftz.bf16.f32 d, a;
+#
+# fails with "Arguments mismatch for instruction 'cvt'". These entries expose
+# callable nvcc inline-asm helpers, not abstract PTX grammar, so do not advertise
+# floating source carriers the supported assembler cannot compile. Destination
+# widening and integer source widening continue to use the documented relaxed
+# domains above; only the ptxas-rejected floating source direction is narrowed.
+_CVT_FLOAT_SRC_DTYPES = {
+ "f16": ("uint16", "int16", "float16", "bfloat16"),
+ "bf16": ("uint16",),
+ "f32": ("float32", "uint32", "int32"),
+ "f64": ("float64", "uint64", "int64"),
+}
+
+
+def _cvt_dst_dtypes(m):
+ """Destination carriers supported by CUDA 13.2 for scalar ``cvt``.
+
+ PTX ISA 9.2 Tables 27/28 describe a relaxed upper bound, but also warn
+ that a particular instruction may reject combinations from that grid.
+ Exact force-inlined CUDA 13.2 probes expose three operand-local gaps:
+
+ - when the source instruction type is ``.bf16``, the destination carrier
+ cannot be wider than its own instruction type;
+ - ``.ftz`` forms ``{u64,s64}.f32``, ``f64.f32``, and ``f32.f64`` reject a
+ 128-bit destination carrier;
+ - ``{rz,rm,rp}.ftz.f32.f64`` also rejects a 64-bit destination carrier,
+ while the corresponding ``.rn`` forms accept it.
+
+ These are ptxas restrictions, not additional PTX ISA rules. In
+ particular, no-ftz 128-bit destinations and the accepted ``.rn`` 64-bit
+ forms remain in the callable domain.
+ """
+ dtypes = _CVT_RELAXED_DTYPES[m["dtype"]]
+ if m["atype"] == "bf16":
+ dtypes = _CVT_SAME_WIDTH_DTYPES[m["dtype"]]
+ if m["ftz"] and (
+ (m["dtype"] in ("u64", "s64", "f64") and m["atype"] == "f32")
+ or (m["dtype"], m["atype"]) == ("f32", "f64")
+ ):
+ dtypes = tuple(dtype for dtype in dtypes if dtype not in ("uint128", "int128"))
+ if m["ftz"] and (m["dtype"], m["atype"]) == ("f32", "f64") and m["rnd"] in ("rz", "rm", "rp"):
+ dtypes = tuple(dtype for dtype in dtypes if dtype not in ("uint64", "int64"))
+ return dtypes
+
+
+def _cvt_src_dtypes(m):
+ """Source carriers supported by ptxas for the generic scalar ``cvt`` line."""
+ atype = m["atype"]
+ dtypes = _CVT_FLOAT_SRC_DTYPES.get(atype, _CVT_RELAXED_DTYPES[atype])
+ # CUDA 13.2 rejects a wider carrier on the opposite operand whenever the
+ # destination instruction type is `.bf16`. Same-width bit/integer carrier
+ # spellings still assemble and remain exposed.
+ if m["dtype"] == "bf16":
+ dtypes = _CVT_SAME_WIDTH_DTYPES[atype]
+ return dtypes
+
def _present_lanes(slot: str) -> LanesFn:
"""A bracketed-optional operand: one register when its qualifier is written.
The ISA spells several of these `{, operand}` against a `{.qualifier}`, and
- says so in the same words each time -- ISA 9.7.9.22:180-182 for cvt's
+ says so in the same words each time -- ISA 9.7.9.21:180-182 for cvt's
scale-factor is typical: "Operand scale-factor and qualifier
.scaled::n2::ue8m0 must be used together." `lanes=0` is what makes the
operand vanish from the helper signature (see render.operand_layout), so
@@ -1899,14 +2311,14 @@ def _present_lanes(slot: str) -> LanesFn:
# `{, scale-factor}` exists exactly when `.scaled::n2::ue8m0` is written. ISA
-# 9.7.9.22:180-182: "Optional qualifier .scaled::n2::ue8m0 specifies that the
+# 9.7.9.21:180-182: "Optional qualifier .scaled::n2::ue8m0 specifies that the
# instruction uses packed scale-factor with 2 scale values of ue8m0 type.
# Operand scale-factor and qualifier .scaled::n2::ue8m0 must be used together."
_cvt_scale_lanes = _present_lanes("scaled")
def _check_cvt_tf32(m):
- """The two `.tf32` lines, ISA 9.7.9.22:18-19 --
+ """The two `.tf32` lines, ISA 9.7.9.21:18-19 --
cvt.rna{.satfinite}.tf32.f32 d, a;
cvt.frnd2{.satfinite}{.relu}.tf32.f32 d, a;
@@ -1944,7 +2356,7 @@ def _cvt_int_covers(d: str, a: str) -> bool:
def _check_cvt_same_type(t, rnd):
- """The `.dtype == .atype` sub-grid of the generic scalar line, ISA 9.7.9.22.
+ """The `.dtype == .atype` sub-grid of the generic scalar line, ISA 9.7.9.21.
A same-type cvt is a real instruction, not a move: an *integer* rounding
mode may ride a float-to-float conversion here. The ISA licenses that for
@@ -1952,10 +2364,10 @@ def _check_cvt_same_type(t, rnd):
six pairs, not four -- `.f16` and `.bf16` are both 16-bit. This toolchain
assembles only the same-type four; see the third ptxas-only restriction in
`_check_cvt_scalar`, which is where the cross-type pairs are ruled on. ISA
- 9.7.9.22:329-331 (Integer Notes): "Integer rounding is required for
+ 9.7.9.21:329-331 (Integer Notes): "Integer rounding is required for
float-to-integer conversions, and for same-size float-to-float conversions
where the value is rounded to an integer. Integer rounding is illegal in
- all other instances." And 9.7.9.22:424-426: "A floating-point value may be
+ all other instances." And 9.7.9.21:424-426: "A floating-point value may be
rounded to an integral value using the integer rounding modifiers (see
Integer Notes). The operands must be of the same size. The result is an
integral value, stored in floating-point format." The section's Examples
@@ -1979,7 +2391,7 @@ def _check_cvt_same_type(t, rnd):
two types being equal, so the caller's shared tail applies them to this
sub-grid too. (For an integer `t` that tail rejects `.sat` because
`_cvt_int_covers(t, t)` holds -- the destination range is the source range
- -- and it rejects `.sat` on .bf16 as a toolchain limit.)
+ -- and for `.bf16` it applies the ISA's destination-type exclusion.)
Together with that tail this admits 53 of the 432 same-type spellings in this entry's slot grid,
which is exactly the set ptxas assembles (measured over the full grid, nvcc
@@ -1993,7 +2405,7 @@ def _check_cvt_same_type(t, rnd):
def _check_cvt_frnd2_scalar(d, a, rnd, ftz, sat):
- """The frnd2 *scalar* lines' sub-grid, ISA 9.7.9.22:10 and :14 --
+ """The frnd2 *scalar* lines' sub-grid, ISA 9.7.9.21:10 and :14 --
cvt.frnd2{.relu}{.satfinite}.f16.f32 d, a;
cvt.frnd2{.relu}{.satfinite}.bf16.f32 d, a;
@@ -2033,7 +2445,7 @@ def _check_cvt_frnd2_scalar(d, a, rnd, ftz, sat):
def _check_cvt_scalar(m):
- """The generic scalar line's rules, quoting ISA 9.7.9.22.
+ """The generic scalar line's rules, quoting ISA 9.7.9.21.
Rounding: "Integer rounding is required for float-to-integer conversions,
and for same-size float-to-float conversions where the value is rounded to
@@ -2058,14 +2470,18 @@ def _check_cvt_scalar(m):
Three restrictions are ptxas's rather than the ISA's, and are recorded here
because the certification pass would otherwise report legal-looking forms
as illegal: this toolchain assembles no conversion between .bf16 and an
- 8-bit integer, takes no `.sat` on any .bf16 operand, and refuses an integer
- rounding mode on the two same-size cross-type float pairs -- `cvt.rni.bf16.f16`
+ 8-bit integer, takes no `.sat` when .bf16 is the source, and refuses an
+ integer rounding mode on the two same-size cross-type float pairs -- `cvt.rni.bf16.f16`
and `cvt.rni.f16.bf16` are "Illegal rounding modifier for instruction 'cvt'"
at ptxas 13.2 / sm_90 even though the Integer Notes clause quoted above
requires integer rounding for exactly those conversions. The `.dtype !=
.atype` branch below therefore rejects them, which is a toolchain verdict,
not the ISA's.
+ A `.bf16` destination is a different case: the ISA's own Floating Point
+ Notes limit `.sat` destinations to `.f16`, `.f32`, and `.f64`, so rejecting
+ `.sat.bf16` is an ISA rule rather than a toolchain deviation.
+
``.relu``/``.satfinite`` belong to the two frnd2 scalar lines, which share
this entry's shape and type list; `_check_cvt_frnd2_scalar` is their
sub-grid, and it runs whenever either token is written.
@@ -2101,8 +2517,10 @@ def _check_cvt_scalar(m):
return ".ftz applies only where .f32 is one of the types"
if sat and d_int and a_int and _cvt_int_covers(d, a):
return f"{d} already covers {a}, so saturation is not possible"
- if sat and "bf16" in (d, a):
- return "this toolchain assembles no .sat on a .bf16 operand"
+ if sat and d == "bf16":
+ return "the ISA limits floating-point .sat destinations to .f16, .f32, and .f64"
+ if sat and a == "bf16":
+ return "this toolchain assembles no .sat when .bf16 is the source"
if "bf16" in (d, a) and (d in ("u8", "s8") or a in ("u8", "s8")):
return "this toolchain assembles no .bf16 <-> 8-bit-integer conversion"
return None
@@ -2110,7 +2528,7 @@ def _check_cvt_scalar(m):
# Which multiplicand-type *set* each mma / mma.sp syntax line draws BOTH of its
# type positions from. A line never pairs a type from one set with a type from
-# another -- ISA 9.7.15.5.14 spells the integer lines as
+# another -- ISA 9.7.14.5.14 spells the integer lines as
#
# mma.sync.aligned.shape.row.col{.satfinite}.s32.atype.btype.s32 d, a, b, c;
# .atype = {.u8, .s8};
@@ -2125,12 +2543,16 @@ def _check_cvt_scalar(m):
# and nowhere else. (`mma_sp_all` is the one entry whose slot domain does the
# job on its own: its .atype/.btype are exactly `.f8type = {.e4m3, .e5m2}`.)
_MMA_INT_LINE = {"u8": "i8", "s8": "i8", "u4": "i4", "s4": "i4", "b1": "b1"}
-# The one floating-point line that quantifies its two multiplicand positions
-# independently: "mma.sync.aligned.shape.row.col.dtype.f8type.f8type.ctype"
-# with ".f8type = {.e4m3, .e5m2};" (ISA 9.7.15.5.14:23,28). Every other fp line
-# writes a literal token in both positions (.f16.f16, .bf16.bf16, .tf32.tf32)
-# or, at .m16n8k8, names .atype/.btype separately and then requires them equal
-# (:129) -- so outside .f8type the two must simply match.
+# The one floating-point line registered here that quantifies its two
+# multiplicand positions independently is
+# "mma.sync.aligned.shape.row.col.dtype.f8type.f8type.ctype", with
+# ".f8type = {.e4m3, .e5m2};" (ISA 9.7.14.5.14:23,28). The unregistered
+# `.kind::f8f6f4` and `.block_scale` lines do the same over
+# `.f8f6f4type = {.e4m3, .e5m2, .e3m2, .e2m3, .e2m1}`. Every other fp line
+# registered here writes a literal token in both positions (.f16.f16,
+# .bf16.bf16, .tf32.tf32) or, at .m16n8k8, names .atype/.btype separately and
+# then requires them equal (:129). Thus among the registered lines, the two
+# types must match outside `.f8type`.
_MMA_F8 = ("e4m3", "e5m2")
@@ -2143,7 +2565,7 @@ def _check_mma_fp_pair(a: str, b: str) -> str | None:
def _check_mma_sp_fp_types(m):
"""The sparse floating-point lines spell one literal token in both
- multiplicand positions, per ISA 9.7.15.6.3:8-21 --
+ multiplicand positions, per ISA 9.7.14.6.3:8-21 --
mma.spvariant.sync.aligned.m16n8k16.row.col.dtype.f16.f16.ctype d, a, b, c, e, f;
mma.spvariant.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 d, a, b, c, e, f;
@@ -2157,12 +2579,12 @@ def _check_mma_sp_fp_types(m):
data-type of the elements in the matrices scale_A and scale_B. In case of
shapes .m16n8k16, .m16n8k32 and .m16n8k64, .dtype must be the same as
.ctype." -- .dtype/.ctype only; grepping "must be the same" over the whole
- 9.7.15.6.3 slice returns that line and nothing else, so the section states
+ 9.7.14.6.3 slice returns that line and nothing else, so the section states
no .atype/.btype equality rule.
The blanket "the multiplicand types must match" this check used to apply to
every sparse entry came from the **wmma.mma** section, a different
- instruction: ISA 9.7.15.4.5:77-78, "For integer wmma, .ctype and .dtype must
+ instruction: ISA 9.7.14.4.5:77-78, "For integer wmma, .ctype and .dtype must
be specified as .s32. Also, the values for .atype and .btype must be the
same, i.e., either both are .s8 or both are .u8." The sparse lines that do
name two independent type variables -- ".s32.atype.btype.s32" and
@@ -2179,7 +2601,7 @@ def _check_mma_sp_fp_types(m):
def _check_mma_sp_int_types(m):
- """The sparse integer lines pair by width class, per ISA 9.7.15.6.3:60-71 --
+ """The sparse integer lines pair by width class, per ISA 9.7.14.6.3:60-71 --
mma.spvariant.sync.aligned.shape.row.col{.satfinite}.s32.atype.btype.s32 d, a, b, c, e, f;
.atype = {.u8, .s8};
@@ -2190,9 +2612,9 @@ def _check_mma_sp_int_types(m):
mixed signedness (.u8 x .s8, .u4 x .s4) is in the grammar; only crossing
the 8-bit and 4-bit lines is not.
- 9.7.15.6.3 states no .atype/.btype equality rule (see
+ 9.7.14.6.3 states no .atype/.btype equality rule (see
`_check_mma_sp_fp_types` for the section's one type-equality sentence);
- the equality this check used to apply is wmma.mma's, ISA 9.7.15.4.5:77-78.
+ the equality this check used to apply is wmma.mma's, ISA 9.7.14.4.5:77-78.
"""
a, b = m["atype"], m["btype"]
if _MMA_INT_LINE[a] != _MMA_INT_LINE[b]:
@@ -2203,7 +2625,7 @@ def _check_mma_sp_int_types(m):
def _check_mma_sp_fp_thread(m):
"""The floating-point lines whose selector names one thread of four.
- ISA 9.7.15.6.1: .f16/.bf16 at .m16n8k16 and .tf32 at .m16n8k8. Sparse
+ ISA 9.7.14.6.1: .f16/.bf16 at .m16n8k16 and .tf32 at .m16n8k8. Sparse
doubles K against the dense line it mirrors, so these shape lists differ
from the dense `_check_mma_fp_f32`'s.
"""
@@ -2248,7 +2670,7 @@ def _check_mma_sp_int_all(m):
def _check_mma_fp_f16(m):
"""The .f16-accumulator forms: the ISA's f16 and f8 lines.
- The half-precision lines (ISA 9.7.15.5.14:8-10, with ".ctype = {.f16,
+ The half-precision lines (ISA 9.7.14.5.14:8-10, with ".ctype = {.f16,
.f32};" / ".dtype = {.f16, .f32};" at :14-15) spell the token literally in
both multiplicand positions --
@@ -2256,7 +2678,7 @@ def _check_mma_fp_f16(m):
mma.sync.aligned.m16n8k8.row.col.dtype.f16.f16.ctype d, a, b, c;
mma.sync.aligned.m16n8k16.row.col.dtype.f16.f16.ctype d, a, b, c;
- -- while the f8 line, ISA 9.7.15.5.14:23 with ".f8type = {.e4m3,
+ -- while the f8 line, ISA 9.7.14.5.14:23 with ".f8type = {.e4m3,
.e5m2};" (:28) and ".shape = {.m16n8k16, .m16n8k32};" (:32),
mma.sync.aligned.shape.row.col.dtype.f8type.f8type.ctype d, a, b, c;
@@ -2268,10 +2690,10 @@ def _check_mma_fp_f16(m):
This check used to reject every `.atype != .btype`. That blanket rule is a
sentence from the **wmma.mma** section, a different instruction: ISA
- 9.7.15.4.5:77-78, "For integer wmma, .ctype and .dtype must be specified as
+ 9.7.14.4.5:77-78, "For integer wmma, .ctype and .dtype must be specified as
.s32. Also, the values for .atype and .btype must be the same, i.e., either
both are .s8 or both are .u8." mma's own restriction block
- (9.7.15.5.14:122-135) scopes ".atype must be the same as .btype." to
+ (9.7.14.5.14:122-135) scopes ".atype must be the same as .btype." to
.m16n8k8 alone, and this entry reaches .m16n8k8 only through the .f16.f16
line above -- the f8 line's shapes are .m16n8k16 / .m16n8k32 -- so that rule
holds by construction here.
@@ -2293,7 +2715,7 @@ def _check_mma_fp_f16(m):
def _check_mma_fp_f32(m):
- """One check per floating-point syntax line of ISA 9.7.15.5.14.
+ """One check per floating-point syntax line of ISA 9.7.14.5.14.
The lines differ in which shapes pair with which operand types, and only
the .m8n8k4 line leaves the layouts free -- every other line spells
@@ -2301,10 +2723,10 @@ def _check_mma_fp_f32(m):
Multiplicand types are NOT equal in general. This check used to reject
every `.atype != .btype`; that blanket rule is a sentence from the
- **wmma.mma** section, a different instruction -- ISA 9.7.15.4.5:77-78, "For
+ **wmma.mma** section, a different instruction -- ISA 9.7.14.4.5:77-78, "For
integer wmma, .ctype and .dtype must be specified as .s32. Also, the values
for .atype and .btype must be the same, i.e., either both are .s8 or both
- are .u8." mma's own restriction block, ISA 9.7.15.5.14:122-135, reads:
+ are .u8." mma's own restriction block, ISA 9.7.14.5.14:122-135, reads:
Specific shapes have type restrictions :
.m8n8k4 : When .ctype is .f32, .dtype must also be .f32.
@@ -2316,7 +2738,7 @@ def _check_mma_fp_f32(m):
so `.atype == .btype` binds at .m16n8k8 alone. That is exactly the one
alternate-fp line whose two type positions are separate variables over one
- set (9.7.15.5.14:21,26-27):
+ set (9.7.14.5.14:21,26-27):
mma.sync.aligned.m16n8k8.row.col.f32.atype.btype.f32 d, a, b, c;
.atype = {.bf16, .tf32};
@@ -2364,18 +2786,18 @@ def _check_mma_fp_f32(m):
def _check_mma_int(m):
- """The integer / sub-byte / single-bit lines of ISA 9.7.15.5.14.
+ """The integer / sub-byte / single-bit lines of ISA 9.7.14.5.14.
Mixed signedness is legal. This check used to reject every `.atype !=
.btype` under the comment "the values for .atype and .btype must be the
same" -- that sentence is the **wmma.mma** section's, a different
- instruction: ISA 9.7.15.4.5:77-78, "For integer wmma, .ctype and .dtype
+ instruction: ISA 9.7.14.4.5:77-78, "For integer wmma, .ctype and .dtype
must be specified as .s32. Also, the values for .atype and .btype must be
the same, i.e., either both are .s8 or both are .u8." mma's own restriction
- block (9.7.15.5.14:122-135) states no rule for the integer lines at all,
+ block (9.7.14.5.14:122-135) states no rule for the integer lines at all,
and scopes ".atype must be the same as .btype." to .m16n8k8 -- a shape no
integer or single-bit line lists. The integer lines quantify their two
- positions independently (9.7.15.5.14:67-77):
+ positions independently (9.7.14.5.14:67-77):
mma.sync.aligned.shape.row.col{.satfinite}.s32.atype.btype.s32 d, a, b, c;
.shape = {.m8n8k16, .m16n8k16, .m16n8k32}
@@ -2418,17 +2840,19 @@ def _check_mma_int(m):
return None
-# wgmma.mma_async register fragments, per ISA 9.7.16.5.1.1 (Register
+# wgmma.mma_async register fragments, per ISA 9.7.15.5.1.1.1-.4 (Matrix
# Fragments): across the 128-thread warpgroup the accumulator D holds
# M*N/128 = N/2 registers per thread (.f32 and .s32), and M*N/256 = N/4
# when .dtype is .f16 (two halves per register). The A fragment of the rs
# form works out to M*K/128/(32/bits) = 4 registers for every (K, type)
# pairing the ISA defines, so it is a plain `lanes=4`, not a function.
#
-# The N domains, straight from each syntax line's `.shape =` set: the
-# floating-point and fp8 lines take every multiple of 8 up to 256; the
-# s8/u8 line drops 40, 56, ... (the odd multiples of 8 above 32) and stops
-# at 224; the single-bit line adds 240 and 256 back.
+# The N domains follow ISA 9.7.15.2 (Matrix Shape): the floating-point and
+# fp8 lines take every multiple of 8 up to 256; the integer and single-bit
+# lines drop 40, 56, ... (the odd multiples of 8 above 32). The integer
+# syntax enumeration in 9.7.15.5.2 stops at 224 even though the Matrix Shape
+# table includes 240 and 256. The public surface follows the instruction's
+# concrete syntax line; the single-bit syntax separately includes both shapes.
_WGMMA_N_FULL = tuple(str(8 * i) for i in range(1, 33))
_WGMMA_N_S8 = ("8", "16", "24", "32", "48", "64", "80", "96", "112", "128",
"144", "160", "176", "192", "208", "224") # fmt: skip
@@ -2440,7 +2864,7 @@ def _wgmma_acc_lanes(m):
return n // 4 if m["dtype"] == "f16" else n // 2
-# cp.async.bulk.tensor coordinate vectors, per ISA 9.7.9.26.5.2: "Vector of
+# cp.async.bulk.tensor coordinate vectors, per ISA 9.7.9.25.5.2: "Vector of
# n elements where n = .dim" -- except the gather4/scatter4 load modes, whose
# tensorCoords is a "fixed length vector of size 5" (one column index plus
# four row indices) whatever the dimension.
@@ -2466,7 +2890,7 @@ def _check_tma_gather4(m):
return None
-# tcgen05.mma disable-output-lane, per ISA 9.7.17.10.9.1: "The size of the
+# tcgen05.mma disable-output-lane, per ISA 9.7.16.10.9.1: "The size of the
# vector is as follows: .cta_group::1 -> 4, .cta_group::2 -> 8".
def _tcgen05_mma_mask_lanes(m):
return 8 if m["cta_group"] == "cta_group::2" else 4
@@ -2507,15 +2931,23 @@ _cp_mask_lanes = _present_lanes("cp_mask")
_ignore_oob_lanes = _present_lanes("ignore_oob")
+def _check_lop3_imm(_m, operand, value):
+ """PTX ISA 9.7.8.6 constrains immLut to one unsigned byte."""
+ if operand == "immLut" and not 0 <= value <= 255:
+ return f"operand 'immLut' must be in the inclusive range 0..255, got {value}"
+ return None
+
+
_ENTRIES = [
# ------------------------------------------------------------------
# PTX ISA 9.7.1 — Integer Arithmetic Instructions
# ------------------------------------------------------------------
- # The section is registered in full, with two exceptions, each stated at
- # the entry it belongs to: the packed `.u8x4`/`.s8x4` types (and the
- # saturating forms that arrived with them), which the ISA gives only to
- # "sm_120f or higher in the same family"; and `clmad` (9.7.1.5), which is
- # PTX ISA 9.3 and does not assemble on this toolchain.
+ # The existing public families cover this section with two exceptions,
+ # each stated at the entry it belongs to: the packed `.u8x4`/`.s8x4` types
+ # (and the saturating forms that arrived with them), which the ISA gives
+ # only to "sm_120f or higher in the same family"; and the distinct `clmad`
+ # family, which was introduced in PTX ISA 9.3 and is not in the current
+ # PTX ISA 9.2 toolchain.
#
# Several mnemonics have a floating-point line further down the table. The
# integer lines are separate entries because they share no qualifier with
@@ -2523,7 +2955,7 @@ _ENTRIES = [
# dispatch resolves them apart on their type tokens, exactly as it already
# does for `add`/`add_half`.
#
- # min / max per PTX ISA 9.7.1.13, 9.7.1.14 (integer), 9.7.3.11, 9.7.3.12
+ # min / max per PTX ISA 9.7.1.12, 9.7.1.13 (integer), 9.7.3.11, 9.7.3.12
# (single/double), 9.7.4.7, 9.7.4.8 (half).
# Each mnemonic gets two entries because the ISA gives it two operand
# shapes: the usual `d, a, b` and the three-source `d, a, b, c` line.
@@ -2574,7 +3006,7 @@ _ENTRIES = [
)
for name in ("max", "min")
],
- # fns per PTX ISA 9.7.1.18: `fns.b32 d, mask, base, offset;` — one form.
+ # fns per PTX ISA 9.7.1.17: `fns.b32 d, mask, base, offset;` — one form.
# The operands carry three different types (mask .b32, base .b32/.u32/.s32,
# offset .s32), so each declares its own dtype rather than sharing the
# entry's type slot.
@@ -2616,7 +3048,7 @@ _ENTRIES = [
ModifierSlot("sat", ("sat",), optional=True),
ModifierSlot("type", ("u16", "u32", "u64", "u16x2", "s16", "s32", "s64", "s16x2")),
),
- check=_check_int_addsub,
+ check=_check_int_add,
# `.u16x2`/`.s16x2` require sm_90 (ISA Target Notes); cert_arch is the
# max floor over the entry's variants.
cert_arch="sm_90",
@@ -2633,7 +3065,7 @@ _ENTRIES = [
ModifierSlot("sat", ("sat",), optional=True),
ModifierSlot("type", _INT_TYPES),
),
- check=_check_int_addsub,
+ check=_check_int_sub,
operands=(
OperandSlot("d", rw="w"),
OperandSlot("a"),
@@ -2706,13 +3138,7 @@ _ENTRIES = [
OperandSlot("c", dtype=_wide_dtype),
),
),
- # NOT REGISTERED: clmad (PTX ISA 9.7.1.5, `clmad.mode.u64 d, a, b, c;`).
- # It was introduced in PTX ISA 9.3; this toolchain assembles 9.2, so ptxas
- # rejects every form and certification could not prove a single variant.
- # Register it when the toolchain catches up -- it is an ordinary
- # mode + fixed-type entry with four .u64 operands.
- #
- # mul24 / mad24 per PTX ISA 9.7.1.6, 9.7.1.7: a 24x24-bit multiply held in
+ # mul24 / mad24 per PTX ISA 9.7.1.5, 9.7.1.6: a 24x24-bit multiply held in
# 32-bit registers, so there is no `.wide` and the type line is 32-bit only.
# mul24.mode.type d, a, b; mad24.mode.type d, a, b, c;
# .mode = {.hi, .lo} mad24.hi.sat.s32 d, a, b, c;
@@ -2743,7 +3169,7 @@ _ENTRIES = [
OperandSlot("c"),
),
),
- # sad per PTX ISA 9.7.1.8: `sad.type d, a, b, c;` -- d = c + |a - b|.
+ # sad per PTX ISA 9.7.1.7: `sad.type d, a, b, c;` -- d = c + |a - b|.
InstructionEntry(
name="sad",
slots=(ModifierSlot("type", _INT_TYPES),),
@@ -2754,7 +3180,7 @@ _ENTRIES = [
OperandSlot("c"),
),
),
- # div / rem per PTX ISA 9.7.1.9, 9.7.1.10: `div.type d, a, b;`. These are
+ # div / rem per PTX ISA 9.7.1.8, 9.7.1.9: `div.type d, a, b;`. These are
# the integer lines; the floating-point `div` (9.7.3.8) is the `div_f`
# entry below, sharing this mnemonic.
*[
@@ -2769,7 +3195,7 @@ _ENTRIES = [
)
for name in ("div", "rem")
],
- # abs / neg per PTX ISA 9.7.1.11, 9.7.1.12 -- signed integers only. Their
+ # abs / neg per PTX ISA 9.7.1.10, 9.7.1.11 -- signed integers only. Their
# floating-point lines (9.7.3.9, 9.7.3.10) are the `abs_f` and `neg`
# entries below; which of each pair carries a suffix is just which arrived
# second. NOT REGISTERED: `neg.s8x4`, sm_120f-only as above.
@@ -2790,7 +3216,7 @@ _ENTRIES = [
OperandSlot("a"),
),
),
- # popc / clz per PTX ISA 9.7.1.15, 9.7.1.16: `popc.type d, a;`. Both count
+ # popc / clz per PTX ISA 9.7.1.14, 9.7.1.15: `popc.type d, a;`. Both count
# bits of a `.b32`/`.b64` source into a `.u32` destination whatever the
# source width -- "destination d has type .u32" -- so d is typed outright
# while a takes the instruction type.
@@ -2805,7 +3231,7 @@ _ENTRIES = [
)
for name in ("popc", "clz")
],
- # bfind per PTX ISA 9.7.1.17: `bfind{.shiftamt}.type d, a;`, d again .u32.
+ # bfind per PTX ISA 9.7.1.16: `bfind{.shiftamt}.type d, a;`, d again .u32.
InstructionEntry(
name="bfind",
slots=(
@@ -2817,7 +3243,7 @@ _ENTRIES = [
OperandSlot("a"),
),
),
- # brev per PTX ISA 9.7.1.19: `brev.type d, a;`.
+ # brev per PTX ISA 9.7.1.18: `brev.type d, a;`.
InstructionEntry(
name="brev",
slots=(ModifierSlot("type", ("b32", "b64")),),
@@ -2826,7 +3252,7 @@ _ENTRIES = [
OperandSlot("a"),
),
),
- # bfe / bfi per PTX ISA 9.7.1.20, 9.7.1.21. Both take their position and
+ # bfe / bfi per PTX ISA 9.7.1.19, 9.7.1.20. Both take their position and
# length operands as `.u32` regardless of the instruction type ("Operands b
# and c are type .u32, but are restricted to the 8-bit value range 0..255").
# bfe.type d, a, b, c; .type = {.u32, .u64, .s32, .s64}
@@ -2845,7 +3271,7 @@ _ENTRIES = [
name="bfi",
# The ISA's own operand names, kept as they are: the destination is `f`
# and `d` is the *length input*, the one family where `d` is not the
- # result. Renaming them would make the helper unreadable against 9.7.1.21.
+ # result. Renaming them would make the helper unreadable against 9.7.1.20.
slots=(ModifierSlot("type", ("b32", "b64")),),
operands=(
OperandSlot("f", rw="w"),
@@ -2855,7 +3281,7 @@ _ENTRIES = [
OperandSlot("d", dtype="u32"), # field length
),
),
- # szext / bmsk per PTX ISA 9.7.1.22, 9.7.1.23 -- both `.mode = {.clamp,
+ # szext / bmsk per PTX ISA 9.7.1.21, 9.7.1.22 -- both `.mode = {.clamp,
# .wrap}`, and both take their width operand as an unsigned 32-bit value.
# szext.mode.type d, a, b; .type = {.u32, .s32}
# bmsk.mode.b32 d, a, b; -- a (position) and b (width) are .u32
@@ -2889,7 +3315,7 @@ _ENTRIES = [
OperandSlot("b", dtype="b32i"), # mask width
),
),
- # dp4a / dp2a per PTX ISA 9.7.1.24, 9.7.1.25:
+ # dp4a / dp2a per PTX ISA 9.7.1.23, 9.7.1.24:
# dp4a.atype.btype d, a, b, c; .atype = .btype = {.u32, .s32}
# dp2a.mode.atype.btype d, a, b, c; .mode = {.lo, .hi}
# The only families in the table with no `type` slot at all: every operand
@@ -3048,7 +3474,7 @@ _ENTRIES = [
),
# abs / neg per PTX ISA 9.7.3.9, 9.7.3.10 -- one shape, one check:
# op{.ftz}.f32 d, a; op.f64 d, a;
- # `abs_f` carries the suffix because the integer `abs` (9.7.1.11) took the
+ # `abs_f` carries the suffix because the integer `abs` (9.7.1.10) took the
# bare name first; `neg` keeps it because the integer one came second.
*[
InstructionEntry(
@@ -3377,6 +3803,13 @@ _ENTRIES = [
# and the number whose sign selects it are independent: "operands d, a and
# b are treated as a bitsize type of the same width as the first
# instruction type; operand c must match the second".
+ #
+ # "Treated as a bitsize type" is load-bearing for the helper ABI: d, a,
+ # and b independently accept the same-width carrier classes ptxas accepts,
+ # with the instruction type's native dtype first. Exhaustively probing the
+ # product gives 996 helpers (up from the former exact-type domain's 378).
+ # c is different: its slot deliberately has no `dtypes` override, so it
+ # remains an exact .s32/.f32 operand through PTX_TYPE_DTYPES.
InstructionEntry(
name="slct",
slots=(
@@ -3386,9 +3819,9 @@ _ENTRIES = [
),
check=_check_slct,
operands=(
- OperandSlot("d", rw="w", dtype="dtype"),
- OperandSlot("a", dtype="dtype"),
- OperandSlot("b", dtype="dtype"),
+ OperandSlot("d", rw="w", dtype="dtype", dtypes=_slct_value_dtypes),
+ OperandSlot("a", dtype="dtype", dtypes=_slct_value_dtypes),
+ OperandSlot("b", dtype="dtype", dtypes=_slct_value_dtypes),
OperandSlot("c", dtype="ctype"),
),
),
@@ -3477,12 +3910,12 @@ _ENTRIES = [
# `escape_token`/`unescape_token` carry them across the boundary. The PTX
# text is unaffected; only the attribute a program types is.
#
- # and / or / xor (PTX ISA 9.7.8.1-9.7.8.3) and not (9.7.8.4): `op.type d,
- # a, b;` over `.pred` and the three bit widths. The predicate line is a
- # real one here -- "Allowed types include predicate registers" -- so with
- # `.pred` the sources want `T.ptx.pred(x)` (or a bool expression) and the
- # destination a uint32 lvalue, and the helper carries the setp/selp bridges
- # around the single instruction.
+ # and / or / xor (PTX ISA 9.7.8.1-9.7.8.3): `op.type d, a, b;`, and not
+ # (9.7.8.4): `not.type d, a;` -- all four over `.pred` and the three bit
+ # widths. The predicate line is real here -- "Allowed types include
+ # predicate registers" -- so with `.pred` the sources want `T.ptx.pred(x)`
+ # (or a bool expression) and the destination a uint32 lvalue, and the
+ # helper carries the setp/selp bridges around the single instruction.
*[
InstructionEntry(
name=name,
@@ -3519,14 +3952,14 @@ _ENTRIES = [
# named by a look-up table byte rather than by an opcode.
# lop3.b32 d, a, b, c, immLut;
# lop3.BoolOp.b32 d|p, a, b, c, immLut, q;
- # `immLut` is an OPEN immediate: it lives in the instruction text, and the
- # ISA gives it a range (0..255) rather than a list of meaningful values, so
- # enumerating 256 helpers per dtype combination would be certifying an
- # arithmetic identity 256 times. Certification samples it, exactly as it
- # does for tcgen05's immHalfSplitoff.
+ # `immLut` is an OPEN immediate because explicitly-unrolled expressions
+ # must survive tracing until they specialize. The entry-level imm_check
+ # enforces the ISA's inclusive 0..255 range both for direct constants and
+ # after specialization without multiplying the rendering domain by 256.
InstructionEntry(
name="lop3",
slots=(ModifierSlot("type", ("b32",)),),
+ imm_check=_check_lop3_imm,
operands=(
OperandSlot("d", rw="w"),
OperandSlot("a"),
@@ -3541,9 +3974,7 @@ _ENTRIES = [
# register classes -- the mechanism only groups the text, so each half
# keeps its own constraint and its own bridge.
# NOT REGISTERED: `.xor` (ptxas: "Illegal operation '.xor' for instruction
- # 'lop3'" -- the ISA's `.BoolOp` line stops at `.or`/`.and`), and `_` in
- # place of `d`, which the ISA allows here but which the engine's per-slot
- # sink guard would read as "every lane sunk"; pass a scratch lvalue.
+ # 'lop3'" -- the ISA's `.BoolOp` line stops at `.or`/`.and`).
InstructionEntry(
name="lop3_bool",
mnemonic="lop3",
@@ -3551,6 +3982,7 @@ _ENTRIES = [
ModifierSlot("boolop", ("or", "and")),
ModifierSlot("type", ("b32",)),
),
+ imm_check=_check_lop3_imm,
operands=(
OperandSlot("d", rw="w", pipe="dp"),
OperandSlot("p", rw="w", dtype="pred", pipe="dp"),
@@ -3561,6 +3993,28 @@ _ENTRIES = [
OperandSlot("q", dtype="pred"),
),
),
+ # The same BoolOp line also permits the bit-size result to be discarded:
+ # `_|p`. This is a separate fixed operand shape, not a caller-selectable
+ # sink lane. A literal immediate models the ISA-owned `_` directly and
+ # leaves the global all-sunk rule untouched; arity selects this sibling.
+ InstructionEntry(
+ name="lop3_bool_sink",
+ mnemonic="lop3",
+ slots=(
+ ModifierSlot("boolop", ("or", "and")),
+ ModifierSlot("type", ("b32",)),
+ ),
+ imm_check=_check_lop3_imm,
+ operands=(
+ OperandSlot("d", kind="imm", literal="_", pipe="dp"),
+ OperandSlot("p", rw="w", dtype="pred", pipe="dp"),
+ OperandSlot("a"),
+ OperandSlot("b"),
+ OperandSlot("c"),
+ OperandSlot("immLut", kind="imm"),
+ OperandSlot("q", dtype="pred"),
+ ),
+ ),
# shf per PTX ISA 9.7.8.7: the funnel shift. a and b are the low and high
# halves of one 64-bit source ("Operand b holds bits 63:32 and operand a
# holds bits 31:0"), and the direction picks which 32 bits of the shifted
@@ -3605,15 +4059,13 @@ _ENTRIES = [
# ------------------------------------------------------------------
# PTX ISA 9.7.9 — Data Movement and Conversion Instructions
# ------------------------------------------------------------------
- # The largest section of the ISA, and registered across all 27 of its
+ # The largest section of the ISA, and registered across all 26 of its
# subsections. What is left out is listed at the entry it belongs to;
# the exclusions that are about the toolchain or the architecture rather
# than about this table are:
# - `shfl` without `.sync` (9.7.9.5), removed by the ISA for sm_70+;
- # - `multimem.st.async` (9.7.9.13), PTX ISA 9.3, which ptxas 13.2
- # (ISA 9.2) cannot assemble;
- # - the FP8 multimem types and `.acc::f16` (9.7.9.15), and
- # `tensormap.replace.swizzle_atomicity` (9.7.9.27), all scoped to
+ # - the FP8 multimem types and `.acc::f16` (9.7.9.14), and
+ # `tensormap.replace.swizzle_atomicity` (9.7.9.26), all scoped to
# architectures above the one their instruction is certified at.
# Everything else absent is an operand the helper ABI cannot carry -- a
# symbol rather than a value -- and says so where it is excluded.
@@ -3647,8 +4099,9 @@ _ENTRIES = [
# symbol '_' may be used for one or more elements provided that at least
# one element is a scalar register." That proviso is `sink_combos`' rule.
#
- # Also unregistered: scalar mov (9.7.9.3), a different instruction that
- # shares the mnemonic.
+ # Registered separately below: scalar mov (9.7.9.3), a different
+ # instruction that shares the mnemonic; dispatch tells them apart by
+ # arity.
*[
InstructionEntry(
name=f"mov_{direction}_{lane_dtype}x{lanes}",
@@ -3692,9 +4145,9 @@ _ENTRIES = [
# because it is a *symbol* rather than a value --
# - .unified (asserts its address names a variable declared with that
# attribute: a declaration property, invisible from a register)
- # - .param/.const spaces (kernel-parameter and const addresses are named,
- # not computed; taking `¶m` in CUDA C already materializes a generic
- # pointer, so no C expression can deliver one)
+ # - .param/.const spaces (PTX can hold either state-space address in a
+ # register, but the helper's CUDA-C ABI cannot originate one: taking an
+ # address in C materializes a generic pointer instead)
# Registering these needs a second rendering model that emits at the site
# where the symbol is in scope, not a slot field.
# The ISA permits @p on this instruction; ptx does not, because it writes a
@@ -3725,7 +4178,7 @@ _ENTRIES = [
),
check=_check_ld,
operands=(
- OperandSlot("d", rw="w", dtypes=_ld_dst_dtypes),
+ OperandSlot("d", rw="w", dtypes=_relaxed_mem_dtypes),
OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("cache_policy", dtype="u64", lanes=_present_lanes("cache"), vector=False),
),
@@ -3774,7 +4227,7 @@ _ENTRIES = [
),
check=_check_ld_vec,
operands=(
- OperandSlot("d", rw="w", lanes=_vec_lanes),
+ OperandSlot("d", rw="w", dtypes=_relaxed_mem_dtypes, lanes=_vec_lanes),
OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("cache_policy", dtype="u64", lanes=_present_lanes("cache"), vector=False),
),
@@ -3793,6 +4246,7 @@ _ENTRIES = [
ModifierSlot("nc", ("nc",), optional=True),
ModifierSlot("l1ev", _L1_EVICT, optional=True),
ModifierSlot("l2ev", _L2_EVICT, optional=True),
+ ModifierSlot("cache", ("L2::cache_hint",), optional=True),
ModifierSlot("prefetch", ("L2::64B", "L2::128B", "L2::256B"), optional=True),
ModifierSlot("vec", ("v4", "v8")),
ModifierSlot("type", _BITS32 + _BITS64),
@@ -3800,15 +4254,24 @@ _ENTRIES = [
cert_arch="sm_100",
check=_check_ld_vec256,
operands=(
+ # PTX ISA 9.2 section 9.4.1 permits a data register wider than the
+ # instruction type, but CUDA 13.2 ptxas does not compile that
+ # carrier axis reliably on the 256-bit ld lines: exact inlined
+ # callers produce C7907, segfaults, or stack-smashing aborts. Keep
+ # these lanes at PTX_TYPE_DTYPES' exact width. Scalar and <=128-bit
+ # vector ld retain the documented relaxed carriers above.
+ #
# `_` means this element is not read from memory.
OperandSlot("d", rw="w", lanes=_vec_lanes, sinkable=_sink256),
OperandSlot("addr", kind="addr", allow_imm_offset=True),
+ OperandSlot("cache_policy", dtype="u64", lanes=_present_lanes("cache"), vector=False),
),
),
# Complete scalar `st` per PTX ISA 9.7.9.11, at parity with `ld`.
- # NOT REGISTERED: .param::func -- a kernel-parameter address is a symbol,
- # not a value, so it cannot flow through the helper-function ABI (see the
- # note on `ld` above).
+ # NOT REGISTERED: .param::func -- a device-function parameter address
+ # cannot originate on the helper's CUDA-C boundary (see the note on `ld`
+ # above). Bare `.param` has the same meaning here because it defaults to
+ # `.param::func`.
InstructionEntry(
name="st",
slots=(
@@ -3828,7 +4291,7 @@ _ENTRIES = [
check=_check_st,
operands=(
OperandSlot("addr", kind="addr", allow_imm_offset=True),
- OperandSlot("value"),
+ OperandSlot("value", dtypes=_relaxed_mem_dtypes),
OperandSlot("cache_policy", dtype="u64", lanes=_present_lanes("cache"), vector=False),
),
),
@@ -3852,7 +4315,7 @@ _ENTRIES = [
check=_check_st_vec,
operands=(
OperandSlot("addr", kind="addr", allow_imm_offset=True),
- OperandSlot("value", lanes=_vec_lanes),
+ OperandSlot("value", dtypes=_st_vec_dtypes, lanes=_vec_lanes),
OperandSlot("cache_policy", dtype="u64", lanes=_present_lanes("cache"), vector=False),
),
),
@@ -3866,6 +4329,7 @@ _ENTRIES = [
ModifierSlot("cop", ("wb", "cg", "cs", "wt"), optional=True),
ModifierSlot("l1ev", _L1_EVICT, optional=True),
ModifierSlot("l2ev", _L2_EVICT, optional=True),
+ ModifierSlot("cache", ("L2::cache_hint",), optional=True),
ModifierSlot("vec", ("v4", "v8")),
ModifierSlot("type", _BITS32 + _BITS64),
),
@@ -3873,20 +4337,25 @@ _ENTRIES = [
check=_check_st_vec256,
operands=(
OperandSlot("addr", kind="addr", allow_imm_offset=True),
+ # PTX ISA 9.2 section 9.4.1 permits wider data registers, but CUDA
+ # 13.2 ptxas does not compile that carrier axis reliably for these
+ # 256-bit st lines (C7907/segfault/stack-smashing on exact inlined
+ # callers). Register only exact-width carriers, matching ld_vec256.
+ # Scalar and <=128-bit vector st keep the relaxed domain.
+ #
# ISA 9.7.9.11 puts the sink in "vector expression b" -- the data
# being stored -- so here `_` means this element is not written to
# memory. Sink is not a destination-only spelling.
OperandSlot("value", lanes=_vec_lanes, sinkable=_sink256),
+ OperandSlot("cache_policy", dtype="u64", lanes=_present_lanes("cache"), vector=False),
),
),
- # st.bulk per PTX ISA 9.7.9.14:
+ # st.bulk per PTX ISA 9.7.9.13:
# st.bulk{.weak}{.shared::cta} [a], size, initval; // initval must be zero
- # NOT REGISTERED: the 32-bit `size` form (ISA: "The 32-bit or 64-bit integer
- # operand size ..."), because no modifier token distinguishes the two -- it
- # would be an operand-shape axis, like mov's. Two ISA constraints are also
- # unenforceable here, being properties of a value rather than of the
- # modifier map that check() sees: "size must be a multiple of 8" and "The
- # maximum value of size operand can be 16777216".
+ # `size` accepts the ISA's 32- or 64-bit integer register forms. Its two
+ # numeric conditions remain caller preconditions because they depend on a
+ # runtime value rather than the modifier map: it must be a multiple of 8
+ # and at most 16777216, otherwise behavior is undefined.
InstructionEntry(
name="st_bulk",
mnemonic="st.bulk",
@@ -3897,11 +4366,11 @@ _ENTRIES = [
),
operands=(
OperandSlot("addr", kind="addr", allow_imm_offset=True),
- OperandSlot("size", dtype="u64"),
+ OperandSlot("size", dtype="u64", dtypes=("uint64", "int64", "uint32", "int32")),
OperandSlot("initval", kind="imm", literal="0"),
),
),
- # prefetch per PTX ISA 9.7.9.16, covering three of its four syntax lines:
+ # prefetch per PTX ISA 9.7.9.15, covering three of its four syntax lines:
# prefetch{.space}.level [a]
# prefetch.global.level::eviction_priority [a]
# prefetch{.tensormap_space}.tensormap [a]
@@ -3962,10 +4431,14 @@ _ENTRIES = [
cert_arch="sm_100",
operands=(
OperandSlot("addr", kind="addr", allow_imm_offset=True),
- OperandSlot("b"),
+ OperandSlot(
+ "b",
+ dtype=_st_async_rel_operand_type,
+ dtypes=_st_async_rel_operand_dtypes,
+ ),
),
),
- # multimem per PTX ISA 9.7.9.15: operations on a multimem address, which
+ # multimem per PTX ISA 9.7.9.14: operations on a multimem address, which
# names one location on each of several GPUs at once. Three mnemonics --
# ld_reduce reduces across the copies into a register, st writes all of
# them, red reduces into all of them -- and each splits into an integer and
@@ -3973,9 +4446,7 @@ _ENTRIES = [
# `.acc::f32`, and their op x type tables are different.
# NOT REGISTERED: the FP8 types (.e4m3*/.e5m2*) and `.acc::f16`, whose
# Target Notes scope them to the sm_100a/sm_120a family architectures
- # rather than to sm_90 where the rest of this subsection lives; and
- # `multimem.st.async` (9.7.9.13), which is PTX ISA 9.3 and does not
- # assemble on this toolchain.
+ # rather than to sm_90 where the rest of this subsection lives.
InstructionEntry(
name="multimem_ld_reduce",
mnemonic="multimem.ld_reduce",
@@ -4083,7 +4554,7 @@ _ENTRIES = [
)
for vec in (False, True)
],
- # createpolicy per PTX ISA 9.7.9.19: build the opaque 64-bit cache-eviction
+ # createpolicy per PTX ISA 9.7.9.18: build the opaque 64-bit cache-eviction
# policy that the `.level::cache_hint` operand of ld/st/cp takes. Three
# syntax lines, three shapes:
# createpolicy.range{.global}.pri{.sec}.b64 policy, [a], primary, total;
@@ -4138,7 +4609,7 @@ _ENTRIES = [
OperandSlot("access_property", dtype="u64"),
),
),
- # cp.async.bulk.prefetch per PTX ISA 9.7.9.26.4.3 -- the non-tensor bulk
+ # cp.async.bulk.prefetch per PTX ISA 9.7.9.25.4.3 -- the non-tensor bulk
# prefetch, beside the tensor one further down:
# cp.async.bulk.prefetch.L2.src{.level::cache_hint} [srcMem], size{, policy};
InstructionEntry(
@@ -4159,7 +4630,7 @@ _ENTRIES = [
OperandSlot("cache_policy", dtype="u64", lanes=_present_lanes("cache"), vector=False),
),
),
- # tensormap.replace per PTX ISA 9.7.9.27: overwrite one field of a 1024-bit
+ # tensormap.replace per PTX ISA 9.7.9.26: overwrite one field of a 1024-bit
# tensor-map object in place. Three field groups, and they differ in shape
# rather than in qualifier, so they are three entries:
# .field1 = {global_address, rank} [addr], new_val
@@ -4290,7 +4761,7 @@ _ENTRIES = [
ModifierSlot("type", _LD_TYPES),
),
operands=(
- OperandSlot("d", rw="w"),
+ OperandSlot("d", rw="w", dtypes=_relaxed_mem_dtypes),
OperandSlot("addr", kind="addr", allow_imm_offset=True),
),
),
@@ -4308,19 +4779,19 @@ _ENTRIES = [
# to send those to, so here the check is the whole rule.
check=_check_vec128,
operands=(
- OperandSlot("d", rw="w", lanes=_vec_lanes),
+ OperandSlot("d", rw="w", dtypes=_relaxed_mem_dtypes, lanes=_vec_lanes),
OperandSlot("addr", kind="addr", allow_imm_offset=True),
),
),
- # prefetchu per PTX ISA 9.7.9.16, the fourth line of that subsection: the
- # uniform cache rather than the data cache, so a different mnemonic and no
- # state space (it "requires a generic address").
+ # prefetchu per PTX ISA 9.7.9.15, the third syntax form of that subsection.
+ # It targets the uniform cache rather than the data cache, so it has a
+ # different mnemonic and no state space (it "requires a generic address").
InstructionEntry(
name="prefetchu",
slots=(ModifierSlot("level", ("L1",)),),
operands=(OperandSlot("addr", kind="addr", allow_imm_offset=True),),
),
- # applypriority / discard per PTX ISA 9.7.9.17, 9.7.9.18. Same shape: an
+ # applypriority / discard per PTX ISA 9.7.9.16, 9.7.9.17. Same shape: an
# address range and a cache level, one hinting how to evict it and the
# other saying it need not be written back at all.
# applypriority{.global}.L2::evict_normal [a], size;
@@ -4342,7 +4813,7 @@ _ENTRIES = [
)
for name, level in (("applypriority", "L2::evict_normal"), ("discard", "L2"))
],
- # isspacep per PTX ISA 9.7.9.20: does this generic address fall inside the
+ # isspacep per PTX ISA 9.7.9.19: does this generic address fall inside the
# window of a given state space? A `.pred` result, so it rides the bridge.
# The address is a plain pointer rather than an `addr` operand: the
# instruction reads the value, it does not dereference it, so there is
@@ -4356,7 +4827,7 @@ _ENTRIES = [
),
asm_volatile=False, # a pure query of an address value
),
- # getctarank per PTX ISA 9.7.9.25: which CTA of the cluster owns this
+ # getctarank per PTX ISA 9.7.9.24: which CTA of the cluster owns this
# address. Two lines by whether the address is a shared one or a generic
# one; `.type` describes the *source*, while "destination d is always a
# 32-bit register".
@@ -4377,7 +4848,7 @@ _ENTRIES = [
)
for name, shared in (("getctarank", True), ("getctarank_generic", False))
],
- # cvt.pack per PTX ISA 9.7.9.23: convert two .s32 values with saturation
+ # cvt.pack per PTX ISA 9.7.9.22: convert two .s32 values with saturation
# and pack them into one register. Two syntax lines, and the operand count
# is what separates them -- the sub-byte line needs `c` to supply the bits
# the pair does not fill, and the 16-bit line has none left over.
@@ -4434,7 +4905,7 @@ _ENTRIES = [
)
for pq in (False, True)
],
- # cvta per PTX ISA 9.7.9.21, both directions over all eight state spaces.
+ # cvta per PTX ISA 9.7.9.20, both directions over all eight state spaces.
# The `.to` direction converts a generic address into a space-specific one
# and the bare direction does the reverse; they are two entries because the
# bare one takes a plain value where `.to` takes a pointer.
@@ -4469,7 +4940,7 @@ _ENTRIES = [
),
asm_volatile=False,
),
- # cvt per PTX ISA 9.7.9.22.
+ # cvt per PTX ISA 9.7.9.21.
#
# The generic scalar line first: `cvt{.irnd|.frnd}{.ftz}{.sat}.dtype.atype`
# over the ISA's twelve-type dtype/atype product. The two spellings the ISA
@@ -4525,13 +4996,13 @@ _ENTRIES = [
# {.f16x2, .bf16, .bf16x2, .tf32} destination formats require sm_80 or
# higher."
operands=(
- OperandSlot("d", rw="w", dtype="dtype"),
- OperandSlot("a", dtype="atype"),
+ OperandSlot("d", rw="w", dtype="dtype", dtypes=_cvt_dst_dtypes),
+ OperandSlot("a", dtype="atype", dtypes=_cvt_src_dtypes),
),
),
# The two frnd2 lines that pack *two* .f32 sources into one register are a
# different shape (`d, a, b`), so they are their own entries. ISA
- # 9.7.9.22:65-68: "For .f16x2 and .bf16x2 instruction type, two inputs a and
+ # 9.7.9.21:65-68: "For .f16x2 and .bf16x2 instruction type, two inputs a and
# b of .f32 type are converted into .f16 or .bf16 type and the converted
# values are packed in the destination register d, such that the value
# converted from input a is stored in the upper half of d and the value
@@ -4717,9 +5188,9 @@ _ENTRIES = [
ModifierSlot("dtype", ("e4m3x2", "e5m2x2")),
ModifierSlot("atype", ("f16x2", "bf16x2")),
),
- # The .bf16x2 source is the PTX ISA 9.2 line: "cvt.rn.satfinite{.relu}
+ # The .bf16x2 source is the PTX ISA 9.1 line: "cvt.rn.satfinite{.relu}
# {.e5m2x2/.e4m3x2}{.bf16x2} is supported on following family-specific
- # architectures:" (9.7.9.22), listing sm_100f, sm_110f and sm_120f.
+ # architectures:" (9.7.9.21), listing sm_100f, sm_110f and sm_120f.
cert_arch="sm_100a",
operands=(
OperandSlot("d", rw="w", dtype="dtype"),
@@ -5122,13 +5593,15 @@ _ENTRIES = [
),
),
),
- # mapa per PTX ISA 9.7.9.24: map a shared address into another CTA of the
- # cluster. All four syntax lines differ only in how `a` is spelled at the
- # PTX level (register / variable / variable+imm); through a C helper the
- # operand is always a register, so they collapse to one entry.
- # `.type` fixes the width of BOTH d and a, so the two type tokens are two
- # entries: `.u64` maps a generic address (a pointer), `.u32` maps a 32-bit
- # shared-window address (a plain register, not bracketed).
+ # mapa per PTX ISA 9.7.9.23: map a shared address into another CTA of the
+ # cluster. Syntax lines that differ only in how `a` is spelled at the PTX
+ # level (register / variable / variable+imm) collapse to a register operand
+ # through the C helper, while generic versus shared carrier semantics stay
+ # separate.
+ # `.type` fixes the width of BOTH d and a. `.u64` accepts either a pointer
+ # (with or without explicit `.shared::cluster`) or an already-materialized
+ # raw address register. The distinct carrier shapes are sibling entries;
+ # no space-aware pointer mechanism is needed.
InstructionEntry(
name="mapa",
slots=(
@@ -5142,15 +5615,39 @@ _ENTRIES = [
OperandSlot("b", dtype="u32"),
),
),
+ InstructionEntry(
+ name="mapa_u64_raw",
+ mnemonic="mapa",
+ slots=(ModifierSlot("type", ("u64",)),),
+ asm_volatile=False,
+ operands=(
+ OperandSlot("d", rw="w"),
+ OperandSlot("a", dtype="u64"),
+ OperandSlot("b", dtype="u32"),
+ ),
+ ),
+ InstructionEntry(
+ name="mapa_u64_shared",
+ mnemonic="mapa",
+ slots=(
+ ModifierSlot("space", ("shared::cluster",)),
+ ModifierSlot("type", ("u64",)),
+ ),
+ asm_volatile=False,
+ operands=(
+ OperandSlot("d", rw="w"),
+ OperandSlot("a", dtype="u64"),
+ OperandSlot("b", dtype="u32"),
+ ),
+ ),
InstructionEntry(
name="mapa_u32",
mnemonic="mapa",
slots=(
- # Not optional here, unlike the .u64 entry: the ISA says that with
- # `.space` omitted "both a and d are registers containing generic
- # addresses", and a generic address does not fit 32 bits on this
- # target. ptxas tolerates the bare `mapa.u32`, but nothing can
- # legitimately call it.
+ # The ISA says that with `.space` omitted "both a and d are
+ # registers containing generic addresses", and a generic address
+ # does not fit 32 bits on this target. ptxas tolerates the bare
+ # `mapa.u32`, but nothing can legitimately call it.
ModifierSlot("space", ("shared::cluster",)),
ModifierSlot("type", ("u32",)),
),
@@ -5176,7 +5673,7 @@ _ENTRIES = [
)
for bulk in (False, True)
],
- # cp.async per PTX ISA 9.7.9.26.3.1: the non-bulk asynchronous copy.
+ # cp.async per PTX ISA 9.7.9.25.3.1: the non-bulk asynchronous copy.
# cp-size is an integer constant the ISA closes to {4, 8, 16} (and to 16
# alone under .cg) -- a choices immediate. The src-size arity zero-fills
# the destination tail; it is a separate entry told apart by arity, the
@@ -5237,7 +5734,7 @@ _ENTRIES = [
operands=(),
),
InstructionEntry( # cp.async.wait_all ;
- # The ISA's second syntax line of 9.7.9.26.3.3, and the one with no
+ # The ISA's second syntax line of 9.7.9.25.3.3, and the one with no
# operand: "cp.async.wait_all is equivalent to :
#
# cp.async.commit_group;
@@ -5251,7 +5748,7 @@ _ENTRIES = [
# it nvcc may hoist the loads of the copied data above the wait.
# "Writes performed by cp.async operations are made visible to the
# executing thread only after: The completion of cp.async.wait_all"
- # (9.7.9.26.3.3) is exactly what that clobber has to protect.
+ # (9.7.9.25.3.3) is exactly what that clobber has to protect.
name="cp_async_wait_all",
mnemonic="cp",
slots=(
@@ -5261,29 +5758,17 @@ _ENTRIES = [
orders_memory=True,
operands=(),
),
- # cp.async.bulk (non-tensor), per PTX ISA 9.7.9.26.4.1: four directions.
- #
- # NOT REGISTERED:
- # - the `.sem.scope`/`.type` lines: "Support for .weak and .relaxed
- # semantics, .scope and .type qualifiers are introduced in PTX ISA version
- # 9.3", which ptxas 13.2 in this toolchain (9.2) cannot assemble, and no
- # call site uses them.
- # - the `{.sem}` (`.weak`) position on the plain lines, for the same reason:
- # omitting an optional qualifier renders the same instruction, and the
- # token itself is 9.3. Same situation `_check_ld` records for
- # ld.mmio.acquire.
+ # cp.async.bulk (non-tensor), per PTX ISA 9.7.9.25.4.1: four directions.
#
# `.ignore_oob` is registered, on the one entry below whose direction the
# ISA gives it: "The qualifier .ignore_oob is only available for the global
- # to .shared::cta copy direction." (9.7.9.26.4.1). It is a PTX ISA 9.2
- # feature -- "Support for .ignore_oob qualifier introduced in PTX ISA
- # version 9.2." -- so unlike the .sem.scope/.type lines above it is inside
- # what this toolchain assembles; ptxas takes it at sm_90 and sm_100.
+ # to .shared::cta copy direction." (9.7.9.25.4.1). It is a PTX ISA 9.2
+ # feature and ptxas 13.2 takes it at sm_90 and sm_100.
InstructionEntry( # global -> shared::cta
# The `{.ignore_oob}` position and its two operands, from the syntax
- # line (9.7.9.26.4.1), wrapped but otherwise verbatim:
+ # line (9.7.9.25.4.1), wrapped but otherwise verbatim:
#
- # cp.async.bulk{.sem}.dst.src.completion_mechanism{.level::cache_hint}
+ # cp.async.bulk.dst.src.completion_mechanism{.level::cache_hint}
# {.ignore_oob}
# [dstMem], [srcMem], size{, ignoreBytesLeft, ignoreBytesRight},
# [mbar] {, cache_policy};
@@ -5381,9 +5866,9 @@ _ENTRIES = [
# .cp_mask masks bytes within each 16-byte chunk: "The i-th bit in
# the 16-bit wide byteMask operand specifies whether the i-th byte
# of each 16-byte wide chunk of source data is copied to the
- # destination." (ISA 9.7.9.26.4.1:181-182). It is independent of
+ # destination." (ISA 9.7.9.25.4.1:181-182). It is independent of
# .L2::cache_hint -- the syntax line brace-marks the two separately
- # ("cp.async.bulk{.sem}.dst.src.completion_mechanism{.level::cache_hint}{.cp_mask}",
+ # ("cp.async.bulk.dst.src.completion_mechanism{.level::cache_hint}{.cp_mask}",
# :72) and the section's only couplings are ":173 When the optional
# argument cache_policy is specified, the qualifier
# .level::cache_hint is required." and ":180 When the optional
@@ -5410,7 +5895,57 @@ _ENTRIES = [
OperandSlot("byte_mask", dtype="u16", lanes=_cp_mask_lanes, vector=False),
),
),
- # cp.async.bulk.tensor (TMA), per ISA 9.7.9.26.5.2-4. The tensor address
+ # cp.reduce.async.bulk (non-tensor), per PTX ISA 9.7.9.25.4.2.
+ # PTX 9.2 exposes the original forms without explicit sem/scope tokens.
+ InstructionEntry(
+ name="cp_reduce_async_bulk_s2c",
+ mnemonic="cp",
+ slots=(
+ ModifierSlot("op", ("reduce",)),
+ ModifierSlot("api", ("async",)),
+ ModifierSlot("kind", ("bulk",)),
+ ModifierSlot("dst", ("shared::cluster",)),
+ ModifierSlot("src", ("shared::cta",)),
+ ModifierSlot("completion", ("mbarrier::complete_tx::bytes",)),
+ ModifierSlot("redop", tuple(_CP_REDUCE_BULK_TYPES["shared::cluster"])),
+ ModifierSlot("type", ("b32", "u32", "s32", "u64")),
+ ),
+ check=_check_cp_reduce_async_bulk,
+ cert_arch="sm_90",
+ operands=(
+ OperandSlot("dst_mem", kind="addr", allow_imm_offset=True, space="shared::cluster"),
+ OperandSlot("src_mem", kind="addr", allow_imm_offset=True, space="shared::cta"),
+ OperandSlot("size", dtype="u32"),
+ OperandSlot("mbar", kind="addr", allow_imm_offset=True, space="shared"),
+ ),
+ ),
+ InstructionEntry(
+ name="cp_reduce_async_bulk_s2g",
+ mnemonic="cp",
+ slots=(
+ ModifierSlot("op", ("reduce",)),
+ ModifierSlot("api", ("async",)),
+ ModifierSlot("kind", ("bulk",)),
+ ModifierSlot("dst", ("global",)),
+ ModifierSlot("src", ("shared::cta",)),
+ ModifierSlot("completion", ("bulk_group",)),
+ ModifierSlot("cache", ("L2::cache_hint",), optional=True),
+ ModifierSlot("redop", tuple(_CP_REDUCE_BULK_TYPES["global"])),
+ ModifierSlot("noftz", ("noftz",), optional=True),
+ ModifierSlot(
+ "type", ("f16", "bf16", "b32", "u32", "s32", "b64", "u64", "s64", "f32", "f64")
+ ),
+ ),
+ check=_check_cp_reduce_async_bulk,
+ cert_arch="sm_90",
+ operands=(
+ OperandSlot("dst_mem", kind="addr", allow_imm_offset=True, space="global"),
+ OperandSlot("src_mem", kind="addr", allow_imm_offset=True, space="shared::cta"),
+ OperandSlot("size", dtype="u32"),
+ OperandSlot("cache_policy", dtype="u64", lanes=_tma_cache_lanes, vector=False),
+ ),
+ ),
+ # cp.async.bulk.tensor (TMA), per ISA 9.7.9.25.5.2-4. The tensor address
# is the composite `[tensorMap, tensorCoords]` -- one PTX operand holding
# a 64-bit tensor-map pointer plus an .s32 coordinate vector -- which is
# what OperandSlot.bracket transcribes. The coordinate count follows .dim
@@ -5558,10 +6093,10 @@ _ENTRIES = [
operands=(),
),
# ------------------------------------------------------------------
- # PTX ISA 9.7.14 — Parallel Synchronization and Communication Instructions
+ # PTX ISA 9.7.13 — Parallel Synchronization and Communication Instructions
# ------------------------------------------------------------------
- # bar / barrier per PTX ISA 9.7.14.1, bar.warp.sync per 9.7.14.2,
- # barrier.cluster per 9.7.14.3.
+ # bar / barrier per PTX ISA 9.7.13.1, bar.warp.sync per 9.7.13.2,
+ # barrier.cluster per 9.7.13.3.
#
# barrier{.cta}.sync{.aligned} a{, b}; bar{.cta}.sync a{, b};
# barrier{.cta}.arrive{.aligned} a, b; bar{.cta}.arrive a, b;
@@ -5644,7 +6179,7 @@ _ENTRIES = [
*[
InstructionEntry( # barrier.cluster.arrive{.sem}{.aligned} / .wait{.acquire}{.aligned}
name=f"barrier_cluster_{act}",
- # Shares the `barrier` mnemonic with 9.7.14.1 so the surface reads
+ # Shares the `barrier` mnemonic with 9.7.13.1 so the surface reads
# `T.ptx.barrier.cluster.arrive(...)`; `cluster` is the token that
# tells the two ISA sections apart.
mnemonic="barrier",
@@ -5664,7 +6199,7 @@ _ENTRIES = [
for act in ("arrive", "wait")
],
InstructionEntry( # bar.warp.sync membermask;
- # The `bar` mnemonic's warp-level line, ISA 9.7.14.2 -- a different
+ # The `bar` mnemonic's warp-level line, ISA 9.7.13.2 -- a different
# section from the CTA barriers above and a different operand: not a
# barrier id but a lane mask. "Operand membermask specifies a 32-bit
# integer which is a mask indicating threads participating in barrier
@@ -5689,7 +6224,7 @@ _ENTRIES = [
orders_memory=True,
operands=(OperandSlot("membermask", dtype="u32"),),
),
- # The reduction forms of bar / barrier (PTX ISA 9.7.14.1), which combine a
+ # The reduction forms of bar / barrier (PTX ISA 9.7.13.1), which combine a
# predicate across the barrier's threads instead of only waiting:
# bar{.cta}.red.popc.u32 d, a{, b}, {!}c;
# bar{.cta}.red.op.pred p, a{, b}, {!}c; .op = {.and, .or}
@@ -5726,10 +6261,10 @@ _ENTRIES = [
for kind in ("popc", "pred")
for count in (False, True)
],
- # vote.sync per PTX ISA 9.7.14.10: reduce a predicate across a warp. The
+ # vote.sync per PTX ISA 9.7.13.9: reduce a predicate across a warp. The
# `.ballot` line returns one bit per lane instead of a single answer, so it
# is a second entry rather than a fourth `.mode` token.
- # NOT REGISTERED: `vote` without `.sync` (9.7.14.9), deprecated in PTX ISA
+ # NOT REGISTERED: `vote` without `.sync` (9.7.13.8), deprecated in PTX ISA
# 6.0 and rejected outright by ptxas at sm_70 and higher, exactly as the
# non-sync `shfl` is; and the `{!}a` negation, per the usual policy.
*[
@@ -5751,7 +6286,7 @@ _ENTRIES = [
("vote_sync_ballot", ("ballot",), "b32"),
)
],
- # match.sync per PTX ISA 9.7.14.11: which lanes of the warp hold the same
+ # match.sync per PTX ISA 9.7.13.10: which lanes of the warp hold the same
# value. `.all` optionally reports, through a second predicate, whether
# every lane agreed -- `.any` has no such answer to give, and ptxas says so
# ("Predicate output not allowed for instruction 'match.any'"), so the
@@ -5787,14 +6322,14 @@ _ENTRIES = [
("match_all_sync_p", "all", True),
)
],
- # activemask per PTX ISA 9.7.14.12: the one instruction in the table that
+ # activemask per PTX ISA 9.7.13.11: the one instruction in the table that
# writes a destination and reads nothing at all.
InstructionEntry(
name="activemask",
slots=(ModifierSlot("type", ("b32",)),),
operands=(OperandSlot("d", rw="w"),),
),
- # elect.sync per PTX ISA 9.7.14.15: pick one leader lane out of the member
+ # elect.sync per PTX ISA 9.7.13.14: pick one leader lane out of the member
# mask. Its `d|p` is not optional the way match's is -- ptxas answers
# "Predicate output expected for instruction 'elect'" to a bare
# destination -- so this family has exactly one shape, and it is the pipe.
@@ -5811,7 +6346,7 @@ _ENTRIES = [
OperandSlot("membermask", dtype="u32"),
),
),
- # redux.sync per PTX ISA 9.7.14.13: reduce a value across the warp.
+ # redux.sync per PTX ISA 9.7.13.12: reduce a value across the warp.
*[
InstructionEntry(
name=name,
@@ -5857,16 +6392,16 @@ _ENTRIES = [
OperandSlot("membermask", dtype="u32"),
),
),
- # fence / membar per PTX ISA 9.7.14.4, griddepcontrol per 9.7.14.14.
+ # fence / membar per PTX ISA 9.7.13.4, griddepcontrol per 9.7.13.13.
#
# These name no address yet constrain everyone else's memory order, so they
# set `orders_memory=True` -- see InstructionEntry for why `asm volatile`
# alone is not enough. Each syntax line whose token sequence differs is its
# own entry, all sharing the `fence` mnemonic, so the surface stays
# `T.ptx.fence...` and the chain narrows on the tokens themselves.
- # NOT REGISTERED: the `.sync_restrict` lines and the fabric-proxy line
- # (fixed token sequences with no users yet), and the deprecated `membar`
- # spellings, which the ISA itself marks as the old style for `fence`.
+ # NOT REGISTERED: the `.sync_restrict` lines (fixed token sequences with no
+ # users yet), and the deprecated `membar` spellings, which the ISA itself
+ # marks as the old style for `fence`.
InstructionEntry( # fence{.sem}.scope;
name="fence",
slots=(
@@ -5924,13 +6459,13 @@ _ENTRIES = [
operands=(
OperandSlot("addr", kind="addr", allow_imm_offset=True),
# "The only supported value for the size operand is 128, which must
- # be a constant integer literal" -- ISA 9.7.14.4.
+ # be a constant integer literal" -- ISA 9.7.13.4.
OperandSlot("size", kind="imm", literal="128"),
),
),
- # The ISA permits @p on this instruction; ptx does not, because it writes a
- # destination -- see InstructionEntry.has_dst for why that needs a "+"
- # constraint first.
+ # The ISA permits @p. A returned-value call may use `preserve_dst=True` when
+ # the old destination must survive the false path; the bit-bucket sibling
+ # has no destination and therefore needs no read-write binding.
InstructionEntry(
name="atom",
slots=(
@@ -5949,7 +6484,7 @@ _ENTRIES = [
OperandSlot("cache_policy", dtype="u64", lanes=_present_lanes("cache"), vector=False),
),
),
- # atom's other syntax lines (PTX ISA 9.7.14.5), each its own entry because
+ # atom's other syntax lines (PTX ISA 9.7.13.5), each its own entry because
# each differs in shape or in type domain from the `.op` line above.
#
# `.cas` compares and swaps, so it takes a fourth value; it is also the one
@@ -6021,10 +6556,11 @@ _ENTRIES = [
)
for mnem in ("atom", "red")
],
- # atom's two vector lines. Both operands are register groups, so this is a
- # different shape again; `_check_atom_vec` holds which widths pair with
- # which element. `.f32` has no `.noftz` (there is no half to flush), which
- # is why the two lines are two entries rather than one.
+ # atom's three vector syntax lines become two entries: the half-word and
+ # packed lines share one schema, and `_check_atomic_vec` withholds `.v8`
+ # from the packed types. Both data operands are register groups, so this is
+ # a different shape again. `.f32` has no `.noftz` (there is no half to
+ # flush), which is why it remains a separate entry.
# The ISA's own examples spell these tokens in a different order than its
# syntax line (`atom.global.v8.f16.max.noftz` against
# `atom{...}.op.noftz{.cache}.vec.type`); ptxas accepts both, and the
@@ -6043,7 +6579,7 @@ _ENTRIES = [
ModifierSlot("vec", ("v2", "v4", "v8") if noftz else ("v2", "v4")),
ModifierSlot("type", types),
),
- check=_check_atom_vec,
+ check=_check_atomic_vec,
operands=(
OperandSlot("d", rw="w", lanes=_vec_lanes),
OperandSlot("addr", kind="addr", allow_imm_offset=True),
@@ -6058,7 +6594,39 @@ _ENTRIES = [
("atom_vec_f32", ("add",), False, ("f32",)),
)
],
- # red / atom scalar `.op` forms per PTX ISA 9.7.14.6 and 9.7.14.5.
+ # red's same three vector syntax lines, represented by the same two shapes
+ # as atom but without the returned-value destination. Vector red is global
+ # or generic-only; the cache-policy operand is present exactly with its
+ # `.L2::cache_hint` qualifier.
+ *[
+ InstructionEntry(
+ name=name,
+ mnemonic="red",
+ slots=(
+ ModifierSlot("sem", _RED_SEM, optional=True),
+ ModifierSlot("scope", _ATOM_SCOPES, optional=True),
+ ModifierSlot("space", ("global",), optional=True),
+ ModifierSlot("op", ops),
+ *((ModifierSlot("noftz", ("noftz",)),) if noftz else ()),
+ ModifierSlot("cache", ("L2::cache_hint",), optional=True),
+ ModifierSlot("vec", ("v2", "v4", "v8") if noftz else ("v2", "v4")),
+ ModifierSlot("type", types),
+ ),
+ check=_check_atomic_vec,
+ operands=(
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
+ OperandSlot("value", lanes=_vec_lanes),
+ OperandSlot(
+ "cache_policy", dtype="u64", lanes=_present_lanes("cache"), vector=False
+ ),
+ ),
+ )
+ for name, ops, noftz, types in (
+ ("red_vec_half", ("add", "min", "max"), True, (*_ATOM_VEC_HALF, *_HALF_X2)),
+ ("red_vec_f32", ("add",), False, ("f32",)),
+ )
+ ],
+ # red / atom scalar `.op` forms per PTX ISA 9.7.13.6 and 9.7.13.5.
InstructionEntry(
name="red",
slots=(
@@ -6082,47 +6650,18 @@ _ENTRIES = [
orders_memory=True,
operands=(),
),
- # mbarrier, per the PTX ISA 9.7.14.16 chapter: init 9.7.14.16.12, inval
- # 9.7.14.16.13, expect_tx 9.7.14.16.14, complete_tx 9.7.14.16.15, arrive
- # 9.7.14.16.16, arrive_drop 9.7.14.16.17, test_wait / try_wait 9.7.14.16.19,
- # pending_count 9.7.14.16.20.
- #
- # The arrive lines come in a `state`-returning form and a sink form
- # (`_, [addr]`); the sink is what a void helper can express, so `_` is an
- # ISA-fixed immediate the way `st.bulk`'s initval is. That also leaves the
- # instruction without a destination, so `pred=` works.
+ # mbarrier, per PTX ISA 9.7.13.15: the PTX 9.2 operation instructions from
+ # init through pending_count are registered below.
#
- # This is the same ISA facility `OperandSlot.sinkable` models elsewhere,
- # spelled differently on purpose: here the sink is not a choice. The
- # state-returning form needs a destination nothing reads today, so only the
- # sunk spelling is registered, and an operand that is ALWAYS `_` is a fixed
- # instruction-text value -- which is what a literal imm is. `sinkable`
- # would let a caller ask for the other form, and there is no other form
- # here. Unify the two the day the state-returning line is registered. arrive_drop's five
- # syntax lines (9.7.14.16.17) are the same five shapes, so they are
- # registered as the same four entries with the action token swapped.
+ # A syntax line with an optional operand is represented by one entry per
+ # arity. Likewise, a destination written as `state` or as the sink `_` is
+ # represented by sibling entries: the sink is an ISA-fixed text operand,
+ # `kind="imm", literal="_"`, rather than a caller-selectable sink mask.
+ # This covers the optional sink on local-CTA arrive/arrive_drop, including
+ # their noComplete forms, while the shared::cluster lines use the same
+ # always-sunk entries because they cannot return a state.
#
- # NOT REGISTERED:
- # - the `state, [addr]` forms: the destination is the barrier's pre-arrival
- # state, which nothing reads today. (`arrive.noComplete` has no sink form
- # at all, so it is registered below with a real state result.)
- # - the non-parity try_wait / test_wait lines (their `state` operand is the
- # arrive-returned token nothing reads today), the `.phase_type` qualifiers
- # (no call site), and try_wait's timeHint-less arity (every caller passes
- # the tick budget).
- # - the mbarrier layout facility, all of it: `mbarrier.init{.layout}`
- # (9.7.14.16.12), `mbarrier.pending_count{.layout}` (9.7.14.16.20) and
- # the whole `mbarrier.check_layout.layout{.ss}.b64 p, [addr];`
- # instruction (9.7.14.16.21). Every one is a PTX ISA 9.3 feature --
- # "Support for .layout qualifier introduced in PTX ISA version 9.3."
- # (9.7.14.16.12, 9.7.14.16.20) and "Introduced in PTX ISA version 9.3."
- # (9.7.14.16.21) -- and ptxas 13.2 in this toolchain tops out at PTX ISA
- # 9.2, so it cannot assemble them at any -arch: measured, sm_90 through
- # sm_120a, "Unknown modifier '.layout::v1'" and "Not a name of any known
- # instruction: 'mbarrier.check_layout'". Same reason the cp.async.bulk
- # `.sem.scope`/`.type` lines above are unregistered. Register them when
- # the toolchain moves to PTX ISA 9.3; the bare `mbarrier.pending_count`
- # line, which predates the qualifier, is registered below.
+ # try_wait's optional timeHint is split by arity in every applicable shape.
#
# The addr slot takes no fixed space: with `.space` omitted the ISA means
# a generic address, which is 64-bit on sm_90+, so pinning the operand to
@@ -6142,7 +6681,7 @@ _ENTRIES = [
OperandSlot("count", dtype="u32"),
),
),
- # mbarrier.arrive (9.7.14.16.16) and mbarrier.arrive_drop (9.7.14.16.17)
+ # mbarrier.arrive (9.7.13.15.13) and mbarrier.arrive_drop (9.7.13.15.14)
# are the same five syntax lines under two action tokens -- compare
#
# mbarrier.arrive{.sem.scope}{.shared{::cta}}.b64 state, [addr]{, count};
@@ -6151,7 +6690,7 @@ _ENTRIES = [
# mbarrier.arrive.expect_tx{.sem.scope}{.shared::cluster}.b64 _, [addr], txCount;
# mbarrier.arrive.noComplete{.release.cta}{.shared{::cta}}.b64 state, [addr], count;
#
- # with 9.7.14.16.17's five, which differ only in the mnemonic's action and
+ # with 9.7.13.15.14's five, which differ only in the mnemonic's action and
# in arrive_drop also decrementing the expected count ("Decrements the
# expected arrival count of the mbarrier object by the value specified by
# the 32-bit integer operand count"). So each shape below is one
@@ -6159,7 +6698,7 @@ _ENTRIES = [
#
# Both sections state the pairing rule `_check_mbarrier_sem_scope` enforces
# in the same words: "Qualifiers .sem and .scope must be specified
- # together." (9.7.14.16.16, 9.7.14.16.17).
+ # together." (9.7.13.15.13, 9.7.13.15.14).
*[
InstructionEntry( # mbarrier.<act>{.sem.scope}{.space}.b64 _, [addr];
# The `{, count}` optionality is one ISA line; here it is two entries
@@ -6222,14 +6761,14 @@ _ENTRIES = [
)
for act in ("arrive", "arrive_drop")
],
- # mbarrier.<act>.noComplete{.release.cta}{.shared{::cta}}.b64 state, [addr], count;
+ # mbarrier.<act>.noComplete{.release.cta}{.shared{::cta}}.b64 state|_, [addr], count;
*[
InstructionEntry(
- # This line has no sink form, so `state` is a real register result.
- # It rides rw="rw" rather than "w": "+" keeps the old value live
- # under a false predicate, which is what lets `pred=` remain legal on
- # an instruction that writes a register. Its qualifier pair is fixed
- # (.release.cta only) and the space domain has no ::cluster.
+ # The state-result sibling rides rw="rw" rather than "w": "+"
+ # keeps the old value live under a false predicate, which is what
+ # lets `pred=` remain legal on an instruction that writes a
+ # register. Its qualifier pair is fixed (.release.cta only) and
+ # the space domain has no ::cluster.
name=f"mbarrier_{act}_no_complete",
mnemonic="mbarrier",
slots=(
@@ -6252,6 +6791,29 @@ _ENTRIES = [
)
for act in ("arrive", "arrive_drop")
],
+ *[
+ InstructionEntry(
+ # The optional sink spelling is its own fixed-text sibling. It has
+ # no destination and therefore accepts ordinary framework pred=.
+ name=f"mbarrier_{act}_no_complete_sink",
+ mnemonic="mbarrier",
+ slots=(
+ ModifierSlot("action", (act,)),
+ ModifierSlot("nocomplete", ("noComplete",)),
+ ModifierSlot("sem", ("release",), optional=True),
+ ModifierSlot("scope", ("cta",), optional=True),
+ ModifierSlot("space", ("shared", "shared::cta"), optional=True),
+ ModifierSlot("type", ("b64",)),
+ ),
+ check=_check_mbarrier_sem_scope,
+ operands=(
+ OperandSlot("state", kind="imm", literal="_"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
+ OperandSlot("count", dtype="u32"),
+ ),
+ )
+ for act in ("arrive", "arrive_drop")
+ ],
# mbarrier.test_wait.parity{.sem.scope}{.ss}.b64 waitComplete, [addr], phaseParity;
# mbarrier.try_wait.parity{.sem.scope}{.ss}.b64 waitComplete, [addr], phaseParity, timeHint;
#
@@ -6270,6 +6832,7 @@ _ENTRIES = [
ModifierSlot("type", ("b64",)),
),
check=_check_mbarrier_sem_scope,
+ cert_arch="sm_90",
operands=(
OperandSlot("wait_complete", rw="w", dtype="pred"),
OperandSlot("addr", kind="addr", allow_imm_offset=True),
@@ -6291,6 +6854,7 @@ _ENTRIES = [
ModifierSlot("type", ("b64",)),
),
check=_check_mbarrier_sem_scope,
+ cert_arch="sm_90",
operands=(
OperandSlot("wait_complete", rw="w", dtype="pred"),
OperandSlot("addr", kind="addr", allow_imm_offset=True),
@@ -6326,7 +6890,7 @@ _ENTRIES = [
),
operands=(OperandSlot("addr", kind="addr", allow_imm_offset=True),),
),
- # The state-returning arrive lines (PTX ISA 9.7.14.16.16 / .17). The
+ # The state-returning arrive lines (PTX ISA 9.7.13.15.13 / .14). The
# entries above bake `_` into the text, which is the only spelling the
# `.shared::cluster` lines offer; these return the phase state instead, and
# the ISA gives them `.shared{::cta}` alone.
@@ -6354,17 +6918,13 @@ _ENTRIES = [
)
for act in ("arrive", "arrive_drop")
# Three shapes per action, mirroring the `_` family above: bare, with a
- # thread count, and with a transaction count. The noComplete line needs
- # no twin -- it has no `_` spelling, so `mbarrier_*_no_complete` above
- # is already the state-returning entry.
+ # thread count, and with a transaction count. noComplete always has a
+ # count and its state/sink siblings are registered separately above.
for suffix, count in (("", False), ("_count", True), ("_expect_tx", True))
],
- # The non-parity waits (PTX ISA 9.7.14.16.19), which take the state token
- # an arrive returned rather than a phase parity. try_wait may also carry a
- # sleep hint, which is a separate syntax line and so a separate entry.
- # NOT REGISTERED: the `.phase_type::*` qualifiers and the
- # `waitComplete|reportPredicate{, reportValue}` report forms -- both are
- # PTX ISA 9.3, and ptxas 13.2 (9.2) rejects the token outright.
+ # The non-parity waits (PTX ISA 9.7.13.15.16), which take the state token
+ # an arrive returned rather than a phase parity. try_wait's optional
+ # timeHint is represented by two arities of the same ISA syntax line.
*[
InstructionEntry(
name=f"mbarrier_{act}{'_hint' if hint else ''}",
@@ -6377,6 +6937,7 @@ _ENTRIES = [
ModifierSlot("type", ("b64",)),
),
check=_check_mbarrier_sem_scope,
+ cert_arch="sm_90",
operands=(
OperandSlot("wait_complete", rw="w", dtype="pred"),
OperandSlot("addr", kind="addr", allow_imm_offset=True),
@@ -6386,7 +6947,7 @@ _ENTRIES = [
)
for act, hint in (("test_wait", False), ("try_wait", False), ("try_wait", True))
],
- # tensormap.cp_fenceproxy per PTX ISA 9.7.14.17: copy a tensor-map object
+ # tensormap.cp_fenceproxy per PTX ISA 9.7.13.16: copy a tensor-map object
# from shared to global and fence the tensormap proxy over it, in one
# instruction. `size` is the literal 128 -- ptxas both refuses a register
# there ("Arguments mismatch") and names the only legal value ("unexpected
@@ -6410,7 +6971,7 @@ _ENTRIES = [
OperandSlot("size", kind="imm", literal="128"),
),
),
- # red.async per PTX ISA 9.7.14.7: an asynchronous reduction whose
+ # red.async per PTX ISA 9.7.13.7: an asynchronous reduction whose
# completion is signalled through an mbarrier. The ISA writes one syntax
# line per op group; `_check_red_async` is that grid.
InstructionEntry(
@@ -6435,8 +6996,6 @@ _ENTRIES = [
# The release line of the same subsection, which reduces straight into
# global memory and needs no mbarrier. Like st.async's release line it is
# sm_100, and `.mmio` is system-scoped only.
- # NOT REGISTERED: `multimem.red.async` (9.7.14.8) -- ptxas 13.2 rejects the
- # mnemonic at sm_90 and sm_100 alike, as it does `multimem.st.async`.
InstructionEntry(
name="red_async_release",
mnemonic="red.async",
@@ -6456,9 +7015,9 @@ _ENTRIES = [
),
),
# cp.async completion tracking: cp.async.mbarrier.arrive per PTX ISA
- # 9.7.14.16.18, cp.async.commit_group per 9.7.9.26.3.2, cp.async.wait_group
- # and cp.async.wait_all per 9.7.9.26.3.3, cp.async.bulk.commit_group /
- # .wait_group per 9.7.9.26.6.1 / 9.7.9.26.6.2.
+ # 9.7.13.15.15, cp.async.commit_group per 9.7.9.25.3.2, cp.async.wait_group
+ # and cp.async.wait_all per 9.7.9.25.3.3, cp.async.bulk.commit_group /
+ # .wait_group per 9.7.9.25.6.1 / 9.7.9.25.6.2.
#
# The wait_group counts are caller-chosen immediates: the ISA gives N no
# register form, so each value is its own helper, and the closed `choices`
@@ -6480,17 +7039,17 @@ _ENTRIES = [
),
# No fixed space on addr, for the reason the mbarrier family states at
# length below: "If no state space is specified then Generic Addressing
- # is used." (ISA 9.7.14.16.18), and a generic address is 64-bit on
+ # is used." (ISA 9.7.13.15.15), and a generic address is 64-bit on
# sm_90+, so pinning the operand to shared would bind a 32-bit register
# under the space-omitted spelling. `operand_space` reads the entry's
# `space` slot instead, so the carrier follows the spelling.
operands=(OperandSlot("addr", kind="addr", allow_imm_offset=True),),
),
InstructionEntry( # mbarrier.pending_count.b64 count, state;
- # The reader of the `state` result the two `.noComplete` entries above
- # produce: "The state operand is a 64-bit register that must be the
- # result of a prior mbarrier.arrive.noComplete or
- # mbarrier.arrive_drop.noComplete instruction." (ISA 9.7.14.16.20).
+ # The reader of the `state` result the state-returning `.noComplete`
+ # entries above produce: "The state operand is a 64-bit register that
+ # must be the result of a prior mbarrier.arrive.noComplete or
+ # mbarrier.arrive_drop.noComplete instruction." (ISA 9.7.13.15.17).
#
# No address and no state space -- the instruction reads a register,
# not the mbarrier object -- so there is no `space` slot here and
@@ -6499,10 +7058,9 @@ _ENTRIES = [
# mbarrier object prior to the arrive-on operation from which the state
# register was obtained."
#
- # `state` is pinned u64 for the reason its producer is: the bit-type
- # dtype axis would otherwise offer an .f64 carrier the instruction has
- # no use for. The `{.layout}` position is unregistered -- see the
- # region note above.
+ # `state` uses the integer-only bit carrier for the reason its producer
+ # does: the ordinary b64 dtype axis would also offer an f64 register,
+ # which ptxas rejects in this position.
name="mbarrier_pending_count",
mnemonic="mbarrier",
slots=(
@@ -6511,10 +7069,10 @@ _ENTRIES = [
),
operands=(
OperandSlot("count", rw="w", dtype="u32"),
- OperandSlot("state", dtype="u64"),
+ OperandSlot("state", dtype="b64i"),
),
),
- # clusterlaunchcontrol.try_cancel per PTX ISA 9.7.14.18.
+ # clusterlaunchcontrol.try_cancel per PTX ISA 9.7.13.17.
InstructionEntry(
name="clusterlaunchcontrol_try_cancel",
mnemonic="clusterlaunchcontrol",
@@ -6527,12 +7085,15 @@ _ENTRIES = [
ModifierSlot("type", ("b128",)),
),
# The instruction itself needs sm_100, but `.multicast::cluster::all`
- # is only on the arch-specific targets (ISA: sm_100a / sm_101a / sm_120a).
+ # is restricted: ISA 9.7.13.17 lists sm_100a, sm_101a and sm_120a,
+ # plus the family-specific sm_100f / sm_101f / sm_110f / sm_120f "or
+ # higher in the same family" from PTX ISA 8.8. sm_100a is the
+ # certification target.
cert_arch="sm_100a",
# Neither address operand takes a fixed space: `.space` is optional on
# the syntax line, and "The .space qualifier is specified, both operands
# addr and mbar must be in the .shared::cta state space. Otherwise,
- # generic addressing will be assumed for both." (ISA 9.7.14.18). A
+ # generic addressing will be assumed for both." (ISA 9.7.13.17). A
# generic address is 64-bit on sm_100, so pinning both to shared would
# bind 32-bit registers under the space-omitted spelling. Letting
# `operand_space` read the entry's `space` slot gives each variant the
@@ -6542,9 +7103,9 @@ _ENTRIES = [
OperandSlot("mbar", kind="addr", allow_imm_offset=True),
),
),
- # clusterlaunchcontrol.query_cancel per PTX ISA 9.7.14.19: decode the
+ # clusterlaunchcontrol.query_cancel per PTX ISA 9.7.13.18: decode the
# opaque b128 response a try_cancel wrote. Three syntax shapes, so three
- # entries -- the ISA writes them as separate lines (9.7.14.19, wrapped but
+ # entries -- the ISA writes them as separate lines (9.7.13.18, wrapped but
# otherwise verbatim):
#
# clusterlaunchcontrol.query_cancel.is_canceled.pred.b128
@@ -6636,9 +7197,9 @@ _ENTRIES = [
),
),
# ------------------------------------------------------------------
- # PTX ISA 9.7.15 — Warp Level Matrix Multiply-Accumulate Instructions
+ # PTX ISA 9.7.14 — Warp Level Matrix Multiply-Accumulate Instructions
# ------------------------------------------------------------------
- # mma per PTX ISA 9.7.15.5.14. Four operand groups (d, a, b, c), each a
+ # mma per PTX ISA 9.7.14.5.14. Four operand groups (d, a, b, c), each a
# register vector whose length follows the Matrix Fragments tables -- four
# callable `lanes`, one per group. `d` and `c` are separate operands: the
# ISA lists them separately and the legacy helper bound them to separate
@@ -6648,7 +7209,8 @@ _ENTRIES = [
# NOT REGISTERED:
# - The `.kind::`/`.block_scale` lines and the .e3m2/.e2m3/.e2m1 types,
# which require sm_120a -- outside the architectures this table certifies.
- # - `mma.sp`, a separate instruction with a metadata operand.
+ # (`mma.sp` / `mma.sp::ordered_metadata` per 9.7.14.6.3 is registered in
+ # the `mma_sp*` entries below.)
InstructionEntry( # half precision, and the alternate formats that share it
name="mma",
slots=(
@@ -6699,7 +7261,7 @@ _ENTRIES = [
),
InstructionEntry( # mma.sync.aligned.m8n8k4.alayout.blayout.f32.f16.f16.f16
# The one dense mma line whose .dtype and .ctype differ. The
- # half-precision syntax line of 9.7.15.5.14 quantifies both ends
+ # half-precision syntax line of 9.7.14.5.14 quantifies both ends
# independently --
#
# mma.sync.aligned.m8n8k4.alayout.blayout.dtype.f16.f16.ctype d, a, b, c;
@@ -6802,7 +7364,7 @@ _ENTRIES = [
# floating point type mma operation with .m8n8k4 shape introduced
# in PTX ISA version 7.0.", ".f64 floating point type mma operation
# with .m8n8k4 shape requires sm_80 or higher.", and a fragment
- # section of its own, 9.7.15.5.2 "Matrix Fragments for mma.m8n8k4
+ # section of its own, 9.7.14.5.2 "Matrix Fragments for mma.m8n8k4
# with .f64 floating point type".
#
# This slot used to omit the shape under a note claiming ptxas
@@ -6810,7 +7372,7 @@ _ENTRIES = [
# evidence were both wrong: that probe fed the operand vectors
# `_mma_threads` produced while it still divided every .m8n8k4 by
# 8, and ptxas was objecting to the vectors, not to the shape. With
- # 9.7.15.5.2's counts (a = b = 1, c = d = 2) ptxas 13.2 assembles
+ # 9.7.14.5.2's counts (a = b = 1, c = d = 2) ptxas 13.2 assembles
# the line at sm_80, sm_90 and sm_90a, and its sm_90a SASS is
# `DMMA.8x8x4 R4, R4, R6, R8` -- one real hardware instruction,
# unlike the .m8n8k4 .bf16 / .tf32 spellings ptxas takes and then
@@ -6833,7 +7395,7 @@ _ENTRIES = [
# are : .rn : mantissa LSB rounds to nearest even. This is the
# default. .rz : mantissa LSB rounds towards zero. .rm : mantissa
# LSB rounds towards negative infinity. .rp : mantissa LSB rounds
- # towards positive infinity." (9.7.15.5.14)
+ # towards positive infinity." (9.7.14.5.14)
#
# The syntax line leaves the position out; the section's own f64
# examples spell it, last, after .ctype -- they write the operands
@@ -6921,7 +7483,7 @@ _ENTRIES = [
("_all", ("m16n8k64", "m16n8k128"), ("0",), _check_mma_sp_int_all),
)
],
- # ldmatrix per PTX ISA 9.7.15.5.15 -- warp-level matrix load. Three syntax
+ # ldmatrix per PTX ISA 9.7.14.5.15 -- warp-level matrix load. Three syntax
# lines; the destination is a brace-enclosed vector of 1/2/4 32-bit
# registers "as per the value of .num", which is what a callable `lanes`
# transcribes. The first line's two shapes split into two entries because
@@ -6991,7 +7553,7 @@ _ENTRIES = [
OperandSlot("p", kind="addr", allow_imm_offset=True),
),
),
- # movmatrix per PTX ISA 9.7.15.5.14 -- transpose one distributed m8n8
+ # movmatrix per PTX ISA 9.7.14.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;
@@ -7010,7 +7572,7 @@ _ENTRIES = [
OperandSlot("a", dtype="b32"),
),
),
- # stmatrix per PTX ISA 9.7.15.5.16 -- the store mirror of ldmatrix. One
+ # stmatrix per PTX ISA 9.7.14.5.16 -- the store mirror of ldmatrix. One
# syntax line; the two shapes split into two entries because their type and
# target floors differ.
#
@@ -7056,14 +7618,14 @@ _ENTRIES = [
OperandSlot("r", dtype="b32", lanes=_matrix_num_lanes),
),
),
- # mma.sp / mma.sp::ordered_metadata per PTX ISA 9.7.15.6.3: the same
+ # mma.sp / mma.sp::ordered_metadata per PTX ISA 9.7.14.6.3: the same
# multiply-accumulate with a structured-sparse A. A holds half of K (the
# other half is implied), so only its group shrinks; `e` carries the
# sparsity metadata and `f` selects which threads contributed it -- the ISA
# calls it "a 32-bit integer constant with values in the range 0..3", an
# instruction-text immediate rather than a register.
#
- # That domain is not 0..3 everywhere: ISA 9.7.15.6.1 states it per shape and
+ # That domain is not 0..3 everywhere: ISA 9.7.14.6.1 states it per shape and
# type as "one thread within a group of four" (0..3), "a thread-pair" (0 or
# 1), or "all threads ... must be 0". A `check` sees only the modifier map,
# never the immediate axis, so each selector domain is its own entry -- the
@@ -7122,7 +7684,7 @@ _ENTRIES = [
),
# "All threads within a group of four ... must be 0."
#
- # No check: this entry is one syntax line, ISA 9.7.15.6.3:22,
+ # No check: this entry is one syntax line, ISA 9.7.14.6.3:22,
# "mma.spvariant.sync.aligned.m16n8k64.row.col.f32.f8type.f8type.f32
# d, a, b, c, e, f;" with ".f8type = {.e4m3, .e5m2};", and the
# slots spell it exactly -- one shape, and .atype/.btype domains
@@ -7130,28 +7692,25 @@ _ENTRIES = [
# so .e4m3 x .e5m2 is in the grammar (the section's own example at
# :249 is "mma.sp.sync.aligned.m16n8k64.row.col.f32.e5m2.e4m3.f32").
# The same-type rule this entry used to carry was wmma.mma's
- # (9.7.15.4.5:77-78), not mma.sp's -- see `_check_mma_sp_fp_types`.
+ # (9.7.14.4.5:77-78), not mma.sp's -- see `_check_mma_sp_fp_types`.
("_all", ("m16n8k64",), ("e4m3", "e5m2"), ("0",), None),
)
],
# ------------------------------------------------------------------
- # PTX ISA 9.7.16 — Asynchronous Warpgroup Level Matrix Multiply-Accumulate Instructions
+ # PTX ISA 9.7.15 — Asynchronous Warpgroup Level Matrix Multiply-Accumulate Instructions
# ------------------------------------------------------------------
- # wgmma.mma_async per PTX ISA 9.7.16.5.2 (sm_90a). Six type groups, each
+ # wgmma.mma_async per PTX ISA 9.7.15.5.2 (sm_90a). Six type groups, each
# with an ss line (both A and B from shared memory, named by 64-bit matrix
# descriptors) and an rs line (A from registers -- always four .b32, see
# the fragment note above -- which also drops imm-trans-a). Like mma, one
# entry per accumulator register type so every operand dtype is pinned.
#
# The accumulator is read and written in place: D = A*B + D, so `d` is
- # rw="rw", the "+" constraint. The trailing arguments all live in the
- # instruction text as caller immediates:
- # - scale-d: "the operation of the form D = A*B is issued when the input
- # predicate argument scale-d is false". The syntax calls it a
- # predicate, but ptxas accepts the literals 0 and 1 in that position
- # (certified), and every call site passes a compile-time constant --
- # so it is a choices immediate, not a runtime operand. (The legacy
- # helper burned a setp + predicate register on it per call.)
+ # rw="rw", the "+" constraint. The trailing arguments are:
+ # - scale-d: a predicate operand. False issues D = A*B; true issues
+ # D = A*B+D. It is a runtime operand and uses the ordinary `.pred`
+ # bridge. Because the operation is warpgroup-collective, callers must
+ # ensure every thread in the warpgroup supplies the same predicate.
# - imm-scale-a/b: "the valid values ... are -1 and 1". The legacy
# helper emitted 0 for "no negate", outside the ISA's domain; ptx
# transcribes the documented set.
@@ -7160,7 +7719,7 @@ _ENTRIES = [
#
# NOT REGISTERED: `wgmma.mma_async.sp` (sparse A, a separate instruction
# with a metadata operand) and the `wgmma.fence`/`commit_group`/
- # `wait_group` companions (registered above).
+ # `wait_group` companions (registered below).
*[
InstructionEntry( # .f16 inputs, .f32 accumulator
name=f"wgmma_f16_{form}",
@@ -7183,7 +7742,7 @@ _ENTRIES = [
else (OperandSlot("a", dtype="u32", lanes=4),)
),
OperandSlot("b_desc", dtype="u64"),
- OperandSlot("scale_d", kind="imm", choices=("0", "1")),
+ OperandSlot("scale_d", dtype="pred"),
OperandSlot("scale_a", kind="imm", choices=("-1", "1")),
OperandSlot("scale_b", kind="imm", choices=("-1", "1")),
*(
@@ -7218,7 +7777,7 @@ _ENTRIES = [
else (OperandSlot("a", dtype="u32", lanes=4),)
),
OperandSlot("b_desc", dtype="u64"),
- OperandSlot("scale_d", kind="imm", choices=("0", "1")),
+ OperandSlot("scale_d", dtype="pred"),
OperandSlot("scale_a", kind="imm", choices=("-1", "1")),
OperandSlot("scale_b", kind="imm", choices=("-1", "1")),
*(
@@ -7253,7 +7812,7 @@ _ENTRIES = [
else (OperandSlot("a", dtype="u32", lanes=4),)
),
OperandSlot("b_desc", dtype="u64"),
- OperandSlot("scale_d", kind="imm", choices=("0", "1")),
+ OperandSlot("scale_d", dtype="pred"),
OperandSlot("scale_a", kind="imm", choices=("-1", "1")),
OperandSlot("scale_b", kind="imm", choices=("-1", "1")),
*(
@@ -7288,7 +7847,7 @@ _ENTRIES = [
else (OperandSlot("a", dtype="u32", lanes=4),)
),
OperandSlot("b_desc", dtype="u64"),
- OperandSlot("scale_d", kind="imm", choices=("0", "1")),
+ OperandSlot("scale_d", dtype="pred"),
OperandSlot("scale_a", kind="imm", choices=("-1", "1")),
OperandSlot("scale_b", kind="imm", choices=("-1", "1")),
),
@@ -7317,7 +7876,7 @@ _ENTRIES = [
else (OperandSlot("a", dtype="u32", lanes=4),)
),
OperandSlot("b_desc", dtype="u64"),
- OperandSlot("scale_d", kind="imm", choices=("0", "1")),
+ OperandSlot("scale_d", dtype="pred"),
OperandSlot("scale_a", kind="imm", choices=("-1", "1")),
OperandSlot("scale_b", kind="imm", choices=("-1", "1")),
),
@@ -7346,7 +7905,7 @@ _ENTRIES = [
else (OperandSlot("a", dtype="u32", lanes=4),)
),
OperandSlot("b_desc", dtype="u64"),
- OperandSlot("scale_d", kind="imm", choices=("0", "1")),
+ OperandSlot("scale_d", dtype="pred"),
OperandSlot("scale_a", kind="imm", choices=("-1", "1")),
OperandSlot("scale_b", kind="imm", choices=("-1", "1")),
),
@@ -7378,7 +7937,7 @@ _ENTRIES = [
else (OperandSlot("a", dtype="u32", lanes=4),)
),
OperandSlot("b_desc", dtype="u64"),
- OperandSlot("scale_d", kind="imm", choices=("0", "1")),
+ OperandSlot("scale_d", dtype="pred"),
),
)
for form in ("ss", "rs")
@@ -7407,15 +7966,15 @@ _ENTRIES = [
else (OperandSlot("a", dtype="u32", lanes=4),)
),
OperandSlot("b_desc", dtype="u64"),
- OperandSlot("scale_d", kind="imm", choices=("0", "1")),
+ OperandSlot("scale_d", dtype="pred"),
),
)
for form in ("ss", "rs")
],
- # wgmma group synchronisation, per PTX ISA 9.7.16.7: fence 9.7.16.7.1,
- # commit_group 9.7.16.7.2, wait_group 9.7.16.7.3. (wait_group's textual
- # group count is the `choices` immediate registered above; the mma_async
- # lines are the acc-role entries further down.)
+ # wgmma group synchronisation, per PTX ISA 9.7.15.7: fence 9.7.15.7.1,
+ # commit_group 9.7.15.7.2, wait_group 9.7.15.7.3. (wait_group's textual
+ # group count is the `choices` immediate registered below; the mma_async
+ # lines are the accumulator entries above.)
*[
InstructionEntry(
name=f"wgmma_{act}",
@@ -7431,11 +7990,11 @@ _ENTRIES = [
)
for act in ("fence", "commit_group")
],
- # wgmma.wait_group per PTX ISA 9.7.16.7.3, same caller-immediate shape.
+ # wgmma.wait_group per PTX ISA 9.7.15.7.3, same caller-immediate shape.
#
# The ISA's domain for N is open: the whole section says only "Operand N is
- # an integer constant." (9.7.16.7.3:14), with no upper bound anywhere --
- # unlike setmaxnreg above, whose [24, 256] the ISA closes itself. The 0..7
+ # an integer constant." (9.7.15.7.3:14), with no upper bound anywhere --
+ # unlike setmaxnreg, whose [24, 256] the ISA closes itself. The 0..7
# below is therefore this table's closure, not an ISA rule: `choices` needs
# a finite set to enumerate and certify, and every variant of a registered
# entry must assemble.
@@ -7461,11 +8020,11 @@ _ENTRIES = [
operands=(OperandSlot("group", kind="imm", choices=tuple(str(n) for n in range(8))),),
),
# ------------------------------------------------------------------
- # PTX ISA 9.7.17 — TensorCore 5th Generation Family Instructions
+ # PTX ISA 9.7.16 — TensorCore 5th Generation Family Instructions
# ------------------------------------------------------------------
- # tcgen05 memory-allocation and synchronisation, per PTX ISA 9.7.17:
- # alloc / dealloc / relinquish_alloc_permit 9.7.17.7.1, wait 9.7.17.8.5,
- # fence 9.7.17.11.1, commit 9.7.17.12.1.
+ # tcgen05 memory-allocation and synchronisation, per PTX ISA 9.7.16:
+ # alloc / dealloc / relinquish_alloc_permit 9.7.16.7.1, wait 9.7.16.8.5,
+ # fence 9.7.16.11.1, commit 9.7.16.12.1.
#
# Every one of these carries `orders_memory=True`: the legacy helpers all
# had `::: "memory"`, and the waits/fences name no address at all.
@@ -7488,7 +8047,7 @@ _ENTRIES = [
cert_arch="sm_100a",
orders_memory=True,
operands=(
- OperandSlot("dst", kind="addr", allow_imm_offset=True, space="shared::cta"),
+ OperandSlot("dst", kind="addr", allow_imm_offset=True),
OperandSlot("ncols", dtype="u32"),
),
),
@@ -7523,9 +8082,10 @@ _ENTRIES = [
orders_memory=True,
operands=(),
),
- # tcgen05.ld / .st per PTX ISA 9.7.17.8.3 / 9.7.17.8.4. The register vector
- # length is `.num` scaled by the shape width, which is what a callable
- # `lanes` transcribes; `taddr` is a tmem address, a packed 32-bit
+ # tcgen05.ld / .st per PTX ISA 9.7.16.8.3 / 9.7.16.8.4. Per Table 49 / 50,
+ # the register vector is `.num` lanes for .16x32bx2/.16x64b/.32x32b,
+ # 2x`.num` for .16x128b, and 4x`.num` for .16x256b; that table is what a
+ # callable `lanes` transcribes. `taddr` is a tmem address, a packed 32-bit
# (row << 16 | col) value rather than a pointer.
#
# tcgen05.ld.sync.aligned.shape.num{.pack::16b}.b32 r, [taddr];
@@ -7558,7 +8118,7 @@ _ENTRIES = [
OperandSlot("taddr", kind="addr", space="tmem"),
),
),
- # The `.16x32bx2` lines, ISA 9.7.17.8.3:11 / 9.7.17.8.4:9 --
+ # The `.16x32bx2` lines, ISA 9.7.16.8.3:11 / 9.7.16.8.4:9 --
#
# tcgen05.ld.sync.aligned.16x32bx2.num{.pack}.b32 r, [taddr], immHalfSplitoff;
# tcgen05.st.sync.aligned.16x32bx2.num{.unpack}.b32 [taddr], immHalfSplitoff, r;
@@ -7639,7 +8199,7 @@ _ENTRIES = [
orders_memory=True,
operands=(),
),
- # tcgen05.cp per PTX ISA 9.7.17.9.2: an async shared -> tmem copy of one
+ # tcgen05.cp per PTX ISA 9.7.16.9.2: an async shared -> tmem copy of one
# shape, optionally multicast across the warps of a warpgroup and
# optionally decompressing fp4/fp6 into fp8 on the way. `s-desc` is the
# same 64-bit shared-memory matrix descriptor tcgen05.mma takes.
@@ -7665,7 +8225,7 @@ _ENTRIES = [
OperandSlot("s_desc", dtype="b64"),
),
),
- # tcgen05.mma per PTX ISA 9.7.17.10.9.1 (sm_100a): D = A*B + D where D
+ # tcgen05.mma per PTX ISA 9.7.16.10.9.1 (sm_100a): D = A*B + D where D
# lives in Tensor Memory (an address, not registers -- so the instruction
# has no dst and @p stays available). Two entries per family split on the
# A operand's home: a 64-bit shared-memory descriptor ("ss") or a tmem
@@ -7818,9 +8378,7 @@ _ENTRIES = [
),
cert_arch="sm_100a",
orders_memory=True,
- operands=(
- OperandSlot("mbar", kind="addr", allow_imm_offset=True, space="shared::cluster"),
- ),
+ operands=(OperandSlot("mbar", kind="addr", allow_imm_offset=True),),
),
# tcgen05.commit...{.shared::cluster}.multicast::cluster.b64 [mbar], ctaMask;
# The multicast form: `pred` is keyword-only, so the trailing mask
@@ -7840,14 +8398,14 @@ _ENTRIES = [
cert_arch="sm_100a",
orders_memory=True,
operands=(
- OperandSlot("mbar", kind="addr", allow_imm_offset=True, space="shared::cluster"),
+ OperandSlot("mbar", kind="addr", allow_imm_offset=True),
OperandSlot("mask", dtype="u16"),
),
),
# ------------------------------------------------------------------
- # PTX ISA 9.7.20 — Miscellaneous Instructions
+ # PTX ISA 9.7.19 — Miscellaneous Instructions
# ------------------------------------------------------------------
- # setmaxnreg per PTX ISA 9.7.20.5: the register count is an immediate the
+ # setmaxnreg per PTX ISA 9.7.19.5: the register count is an immediate the
# ISA bounds to [24, 256] in steps of 8 -- exactly a closed choices set.
InstructionEntry(
name="setmaxnreg",
@@ -7866,6 +8424,54 @@ _ENTRIES = [
]
+# ISA 9.7.13.5: "Simple reductions may be specified by using the bit bucket
+# destination operand `_`." This is a whole PTX operand, including on vector
+# atom, rather than a brace group of per-lane sinks. Derive one fixed-literal
+# sibling from every returned-value atom shape; the shorter Python-call arity
+# selects the bit bucket.
+#
+# One documented subset cannot stay modifier-for-modifier identical: CUDA 13.2
+# ptxas does not compile scalar or vector bf16 atom bit buckets (it segfaults on
+# exact force-inlined kernel probes), although PTX ISA 9.2 explicitly supports
+# them and the returned-value bf16 forms compile. The two half-entry checks
+# below therefore withhold only bf16/bf16x2 bit buckets. This is a registration
+# restriction for the current callable toolchain, not an ISA restriction.
+_ATOM_BITBUCKET_BASES = (
+ "atom",
+ "atom_cas",
+ "atom_exch",
+ "atom_half",
+ "atom_vec_half",
+ "atom_vec_f32",
+)
+
+
+def _atom_bitbucket_sibling(entry: InstructionEntry) -> InstructionEntry:
+ destination, *tail = entry.operands
+ if not (
+ entry.ptx_name == "atom"
+ and destination.name == "d"
+ and destination.kind == "reg"
+ and destination.rw == "w"
+ ):
+ raise ValueError(f"{entry.name}: atom bit-bucket base must start with a written register d")
+ bitbucket_checks = {
+ "atom_half": _check_atom_half_bitbucket,
+ "atom_vec_half": _check_atom_vec_half_bitbucket,
+ }
+ return replace(
+ entry,
+ name=f"{entry.name}_bitbucket",
+ mnemonic="atom",
+ check=bitbucket_checks.get(entry.name, entry.check),
+ operands=(OperandSlot("d", kind="imm", literal="_"), *tail),
+ )
+
+
+_entries_by_name = {entry.name: entry for entry in _ENTRIES}
+_ENTRIES.extend(_atom_bitbucket_sibling(_entries_by_name[name]) for name in _ATOM_BITBUCKET_BASES)
+
+
def _validate_imm_offset_slots(entries) -> None:
"""Reject capability bits on operand classes that cannot spell ``[addr+imm]``."""
errors = []
diff --git a/python/tvm/script/tirx.pyi b/python/tvm/script/tirx.pyi
index 111a77ff17..eb88764ac0 100644
--- a/python/tvm/script/tirx.pyi
+++ b/python/tvm/script/tirx.pyi
@@ -104,8 +104,9 @@ class _Chain_applypriority:
def __call__(self, addr: Any, *args: Any, pred: Any = None) -> None: ...
class _Chain_atom:
- """`atom` — 6 entries sharing this mnemonic; PTX puts their difference in the operand list,
- so the call selects one. Shapes: (*__operands); (d, addr, compare, value)
+ """`atom` — 12 entries sharing this mnemonic; PTX puts their difference in the operand
+ list, so the call selects one. Shapes: (*__operands); (d, addr, compare, value); (addr,
+ compare, value)
"""
L2__cache_hint: _Chain_atom
@@ -355,7 +356,7 @@ class _Chain_cos:
) -> None: ...
class _Chain_cp:
- """`cp` — 22 entries sharing this mnemonic; PTX puts their difference in the operand list,
+ """`cp` — 24 entries sharing this mnemonic; PTX puts their difference in the operand list,
so the call selects one. Shapes: (*__operands); (group); (); (dst_mem, src_mem, size,
mbar); (addr)
"""
@@ -369,7 +370,9 @@ class _Chain_cp:
and_: _Chain_cp
arrive: _Chain_cp
async_: _Chain_cp
+ b32: _Chain_cp
b64: _Chain_cp
+ bf16: _Chain_cp
bulk: _Chain_cp
bulk_group: _Chain_cp
ca: _Chain_cp
@@ -379,6 +382,9 @@ class _Chain_cp:
cta_group__1: _Chain_cp
cta_group__2: _Chain_cp
dec: _Chain_cp
+ f16: _Chain_cp
+ f32: _Chain_cp
+ f64: _Chain_cp
global_: _Chain_cp
ignore_oob: _Chain_cp
inc: _Chain_cp
@@ -387,11 +393,14 @@ class _Chain_cp:
mbarrier__complete_tx__bytes: _Chain_cp
min: _Chain_cp
multicast__cluster: _Chain_cp
+ noftz: _Chain_cp
noinc: _Chain_cp
or_: _Chain_cp
prefetch: _Chain_cp
read: _Chain_cp
reduce: _Chain_cp
+ s32: _Chain_cp
+ s64: _Chain_cp
shared: _Chain_cp
shared__cluster: _Chain_cp
shared__cta: _Chain_cp
@@ -399,6 +408,8 @@ class _Chain_cp:
tile: _Chain_cp
tile__gather4: _Chain_cp
tile__scatter4: _Chain_cp
+ u32: _Chain_cp
+ u64: _Chain_cp
wait_all: _Chain_cp
wait_group: _Chain_cp
xor: _Chain_cp
@@ -817,8 +828,9 @@ class _Chain_lg2:
) -> None: ...
class _Chain_lop3:
- """`lop3` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list,
- so the call selects one. Shapes: (d, a, b, c, immLut); (d, p, a, b, c, immLut, q)
+ """`lop3` — 3 entries sharing this mnemonic; PTX puts their difference in the operand list,
+ so the call selects one. Shapes: (d, a, b, c, immLut); (d, p, a, b, c, immLut, q); (p,
+ a, b, c, immLut, q)
"""
and_: _Chain_lop3
@@ -854,7 +866,7 @@ class _Chain_mad24:
"""`mad24` — mode∈{hi,lo}; sat∈{sat} (opt); type∈{u32,s32} — `.sat` on the multiply-add
lines: `.hi` mode, `.s32` type, nothing else. Both lines spell it as a syntax line of
its own -- `mad.hi.sat.s32 d, a, b, c;` (ISA 9.7.1.4) and `mad24.hi.sat.s32 d, a, b, c;`
- (9.7.1.7) -- with the Notes repeating "Applies only to .s32 type in .hi mode".
+ (9.7.1.6) -- with the Notes repeating "Applies only to .s32 type in .hi mode".
"""
hi: _Chain_mad24
@@ -874,7 +886,7 @@ class _Chain_mad24:
) -> None: ...
class _Chain_mapa:
- """`mapa` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list,
+ """`mapa` — 4 entries sharing this mnemonic; PTX puts their difference in the operand list,
so the call selects one. Shapes: (d, a, b)
"""
@@ -922,7 +934,7 @@ class _Chain_max:
def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ...
class _Chain_mbarrier:
- """`mbarrier` — 25 entries sharing this mnemonic; PTX puts their difference in the operand
+ """`mbarrier` — 27 entries sharing this mnemonic; PTX puts their difference in the operand
list, so the call selects one. Shapes: (addr, count); (addr); (addr, tx_count); (state,
addr, count); (wait_complete, addr, phase); (wait_complete, addr, phase, time_hint);
(state, addr); (wait_complete, addr, state); (wait_complete, addr, state, time_hint);
@@ -1275,7 +1287,7 @@ class _Chain_popc:
class _Chain_prefetch:
"""`prefetch` — space∈{global,local,const,param} (opt); level∈{L1,L2} (opt);
evict∈{L2::evict_last,L2::evict_normal} (opt); tensormap∈{tensormap} (opt) — Each
- prefetch syntax line names exactly one target (PTX ISA 9.7.9.16).
+ prefetch syntax line names exactly one target (PTX ISA 9.7.9.15).
`.level::eviction_priority` stays bound to `.global` on purpose: its syntax line is
`prefetch.global.level::eviction_priority`, with `.global` written in rather than the
`{.ss}` that the `ld` lines carry. Generic addressing is not offered there, so neither
@@ -1349,7 +1361,7 @@ class _Chain_rcp:
) -> None: ...
class _Chain_red:
- """`red` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list,
+ """`red` — 4 entries sharing this mnemonic; PTX puts their difference in the operand list,
so the call selects one. Shapes: (*__operands)
"""
@@ -1384,6 +1396,9 @@ class _Chain_red:
sys: _Chain_red
u32: _Chain_red
u64: _Chain_red
+ v2: _Chain_red
+ v4: _Chain_red
+ v8: _Chain_red
xor: _Chain_red
def __call__(self, *args: Any, pred: Any = None) -> None: ...
@@ -1814,6 +1829,7 @@ class _Chain_st_async:
b16: _Chain_st_async
b32: _Chain_st_async
b64: _Chain_st_async
+ b8: _Chain_st_async
f32: _Chain_st_async
f64: _Chain_st_async
global_: _Chain_st_async
@@ -1824,11 +1840,13 @@ class _Chain_st_async:
s16: _Chain_st_async
s32: _Chain_st_async
s64: _Chain_st_async
+ s8: _Chain_st_async
shared__cluster: _Chain_st_async
sys: _Chain_st_async
u16: _Chain_st_async
u32: _Chain_st_async
u64: _Chain_st_async
+ u8: _Chain_st_async
v2: _Chain_st_async
v4: _Chain_st_async
weak: _Chain_st_async
@@ -2210,6 +2228,7 @@ class _Chain_xor:
) -> None: ...
class _PTX:
+ SINK: Any
abs: _Chain_abs
activemask: _Chain_activemask
add: _Chain_add
@@ -2305,6 +2324,7 @@ class _PTX:
wgmma: _Chain_wgmma
xor: _Chain_xor
def addr(self, base: Any, byte_offset: Any) -> Any: ...
+ def pred(self, value: Any) -> Any: ...
def __getitem__(self, text: str) -> Any: ...
ptx: _PTX
diff --git a/tests/python/tirx/codegen/test_codegen_hopper.py b/tests/python/tirx/codegen/test_codegen_hopper.py
index 1ac6c43cd8..8bfc6f56fa 100644
--- a/tests/python/tirx/codegen/test_codegen_hopper.py
+++ b/tests/python/tirx/codegen/test_codegen_hopper.py
@@ -1029,7 +1029,7 @@ def test_wgmma_ss_nt():
T.cuda.wgmma.encode_matrix_descriptor(T.address_of(descB), B_smem.data, *B_encode_args) # noqa: F821
T.ptx.wgmma.fence.sync.aligned()
T.ptx[mma_chain](*get_accum_list(C_local, C_elems), descA, descB, # noqa: F821
- 0, 1, 1, int(transA), int(transB))
+ False, 1, 1, int(transA), int(transB))
T.ptx.wgmma.commit_group.sync.aligned()
T.ptx.wgmma.wait_group.sync.aligned(0)
@@ -1194,7 +1194,7 @@ def test_wgmma_rs_nt():
T.cuda.wgmma.encode_matrix_descriptor(T.address_of(descB), B_smem.data, *B_encode_args) # noqa: F821
T.ptx.wgmma.fence.sync.aligned()
T.ptx[mma_chain](*get_accum_list(C_local, C_elems), *get_A_list(A_local_b32, A_elems_b32), # noqa: E501
- descB, 0, 1, 1, int(transB)) # noqa: F821
+ descB, False, 1, 1, int(transB)) # noqa: F821
T.ptx.wgmma.commit_group.sync.aligned()
T.ptx.wgmma.wait_group.sync.aligned(0)
diff --git a/tests/python/tirx/codegen/test_ptx_addr.py b/tests/python/tirx/codegen/test_ptx_addr.py
index fa4dac30b6..8f8aed7e9e 100644
--- a/tests/python/tirx/codegen/test_ptx_addr.py
+++ b/tests/python/tirx/codegen/test_ptx_addr.py
@@ -54,8 +54,8 @@ def test_ptx_addr_registration_and_table_capabilities():
assert "tirx.ptx.addr" in CODEGEN_REGISTRY
addresses = [slot for entry in TABLE.values() for slot in entry.operands if slot.kind == "addr"]
- assert len(addresses) == 145
- assert sum(slot.allow_imm_offset for slot in addresses) == 115
+ assert len(addresses) == 160
+ assert sum(slot.allow_imm_offset for slot in addresses) == 130
assert sum(slot.bracket is not None and not slot.allow_imm_offset for slot in addresses) == 5
assert sum(slot.space == "tmem" and not slot.allow_imm_offset for slot in addresses) == 25
diff --git a/tests/python/tirx/codegen/test_ptx_cvt.py b/tests/python/tirx/codegen/test_ptx_cvt.py
index eaf868d5e9..21cf6b6beb 100644
--- a/tests/python/tirx/codegen/test_ptx_cvt.py
+++ b/tests/python/tirx/codegen/test_ptx_cvt.py
@@ -14,12 +14,12 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""The ptx cvt entries: one case per registered syntax line of ISA 9.7.9.22."""
+"""The ptx cvt entries: one case per registered syntax line of ISA 9.7.9.21."""
import pytest
from tvm.backend.cuda.ptx.render import render_variant
-from tvm.backend.cuda.ptx.table import TABLE, renderings, tokens_for
+from tvm.backend.cuda.ptx.table import TABLE, mods, operand_dtypes, renderings, tokens_for
# (entry name, the modifier slots to write, the instruction that combination emits).
# Slots are named, not positional: `tokens_for` shifts nothing when a slot is
@@ -235,8 +235,8 @@ _FORM_CASES = [
# The generic scalar line is one entry named plain "cvt"; the packed lines are
# the "cvt_*" family.
-# Every entry of the `cvt` instruction (ISA 9.7.9.22). Keyed off the mnemonic
-# rather than the table name: `cvt.pack` (9.7.9.23) is a different instruction
+# Every entry of the `cvt` instruction (ISA 9.7.9.21). Keyed off the mnemonic
+# rather than the table name: `cvt.pack` (9.7.9.22) is a different instruction
# that happens to sort under the same prefix, and its forms are not points on
# this conversion grid.
_CVT_ENTRIES = {name for name, entry in TABLE.items() if entry.ptx_name == "cvt"}
@@ -278,10 +278,221 @@ def test_cvt_packed_operands_bind_their_carrier():
assert "uint16_t __a" in source
+_CVT_RELAXED_DTYPE_CASES = [
+ (
+ "u8",
+ (
+ "uint8",
+ "int8",
+ "uint16",
+ "int16",
+ "uint32",
+ "int32",
+ "uint64",
+ "int64",
+ "uint128",
+ "int128",
+ ),
+ None,
+ ),
+ (
+ "s8",
+ (
+ "int8",
+ "uint8",
+ "int16",
+ "uint16",
+ "int32",
+ "uint32",
+ "int64",
+ "uint64",
+ "int128",
+ "uint128",
+ ),
+ None,
+ ),
+ (
+ "u16",
+ ("uint16", "int16", "uint32", "int32", "uint64", "int64", "uint128", "int128"),
+ None,
+ ),
+ (
+ "s16",
+ ("int16", "uint16", "int32", "uint32", "int64", "uint64", "int128", "uint128"),
+ None,
+ ),
+ ("u32", ("uint32", "int32", "uint64", "int64", "uint128", "int128"), None),
+ ("s32", ("int32", "uint32", "int64", "uint64", "int128", "uint128"), None),
+ ("u64", ("uint64", "int64", "uint128", "int128"), None),
+ ("s64", ("int64", "uint64", "int128", "uint128"), None),
+ (
+ "f16",
+ (
+ "uint16",
+ "int16",
+ "float16",
+ "bfloat16",
+ "uint32",
+ "int32",
+ "uint64",
+ "int64",
+ "uint128",
+ "int128",
+ ),
+ ("uint16", "int16", "float16", "bfloat16"),
+ ),
+ ("bf16", ("uint16",), ("uint16",)),
+ (
+ "f32",
+ ("float32", "uint32", "int32", "uint64", "int64", "uint128", "int128"),
+ ("float32", "uint32", "int32"),
+ ),
+ (
+ "f64",
+ ("float64", "uint64", "int64", "uint128", "int128"),
+ ("float64", "uint64", "int64"),
+ ),
+]
+
+
+@pytest.mark.parametrize("ptx_type,dst_expected,src_expected", _CVT_RELAXED_DTYPE_CASES)
+def test_cvt_generic_scalar_relaxed_carriers(ptx_type, dst_expected, src_expected):
+ """Keep documented relaxation except where ptxas rejects a floating source.
+
+ The native spelling remains first to preserve canonical helper names.
+ Integer instruction types never admit a floating register, and the cvt
+ section explicitly exempts `.bf16` from widening in both directions. PTX
+ 9.2 permits wider bit registers for the other floating sources, but ptxas
+ rejects them, so they are not callable interfaces in this backend.
+ """
+ entry = TABLE["cvt"]
+ mod_map = mods(entry, tokens_for(entry, dtype=ptx_type, atype=ptx_type))
+ assert operand_dtypes(entry.operands[0], mod_map) == dst_expected
+ assert operand_dtypes(entry.operands[1], mod_map) == (src_expected or dst_expected)
+
+
+def test_cvt_generic_scalar_cuda_13_2_bf16_opposite_operand_gap():
+ """CUDA 13.2 also rejects widening the operand opposite `.bf16`."""
+ entry = TABLE["cvt"]
+ to_bf16 = mods(entry, tokens_for(entry, rnd="rn", dtype="bf16", atype="u32"))
+ assert operand_dtypes(entry.operands[0], to_bf16) == ("uint16",)
+ assert operand_dtypes(entry.operands[1], to_bf16) == ("uint32", "int32")
+
+ from_bf16 = mods(entry, tokens_for(entry, rnd="rni", dtype="u32", atype="bf16"))
+ assert operand_dtypes(entry.operands[0], from_bf16) == ("uint32", "int32")
+ assert operand_dtypes(entry.operands[1], from_bf16) == ("uint16",)
+
+ # The gap is about width, not signedness or register interpretation.
+ from_bf16_f16 = mods(entry, tokens_for(entry, rnd="rn", dtype="f16", atype="bf16"))
+ assert operand_dtypes(entry.operands[0], from_bf16_f16) == (
+ "uint16",
+ "int16",
+ "float16",
+ "bfloat16",
+ )
+
+
+def test_cvt_generic_scalar_cuda_13_2_ftz_destination_gaps():
+ """Pin the two destination-only `.ftz` gaps measured on CUDA 13.2."""
+ entry = TABLE["cvt"]
+
+ u64_from_f32 = mods(entry, tokens_for(entry, rnd="rni", ftz="ftz", dtype="u64", atype="f32"))
+ assert operand_dtypes(entry.operands[0], u64_from_f32) == ("uint64", "int64")
+
+ f64_from_f32 = mods(entry, tokens_for(entry, ftz="ftz", dtype="f64", atype="f32"))
+ assert operand_dtypes(entry.operands[0], f64_from_f32) == (
+ "float64",
+ "uint64",
+ "int64",
+ )
+
+ f32_from_f64_rz = mods(entry, tokens_for(entry, rnd="rz", ftz="ftz", dtype="f32", atype="f64"))
+ assert operand_dtypes(entry.operands[0], f32_from_f64_rz) == (
+ "float32",
+ "uint32",
+ "int32",
+ )
+
+ # ptxas accepts the 64-bit destination on `.rn`, but still rejects q.
+ f32_from_f64_rn = mods(entry, tokens_for(entry, rnd="rn", ftz="ftz", dtype="f32", atype="f64"))
+ assert operand_dtypes(entry.operands[0], f32_from_f64_rn) == (
+ "float32",
+ "uint32",
+ "int32",
+ "uint64",
+ "int64",
+ )
+
+
+def test_cvt_generic_scalar_keeps_supported_wide_carriers():
+ """Do not turn CUDA 13.2's narrow gaps into a blanket cvt restriction."""
+ entry = TABLE["cvt"]
+
+ no_ftz = mods(entry, tokens_for(entry, rnd="rn", dtype="f32", atype="f64"))
+ assert operand_dtypes(entry.operands[0], no_ftz)[-2:] == ("uint128", "int128")
+
+ narrow_ftz = mods(entry, tokens_for(entry, rnd="rn", ftz="ftz", dtype="f16", atype="f32"))
+ assert operand_dtypes(entry.operands[0], narrow_ftz)[-2:] == ("uint128", "int128")
+
+ wide_integer_source = mods(
+ entry, tokens_for(entry, rnd="rn", ftz="ftz", dtype="f32", atype="u64")
+ )
+ assert operand_dtypes(entry.operands[1], wide_integer_source)[-2:] == (
+ "uint128",
+ "int128",
+ )
+
+
+def test_cvt_generic_scalar_relaxed_carriers_render_exact_instruction():
+ entry = TABLE["cvt"]
+
+ tokens = tokens_for(entry, dtype="s16", atype="u16")
+ opcode, helper, source = render_variant(entry, tokens, dtypes=("int64", "uint64"))
+ assert opcode == "cvt.s16.u16"
+ assert helper == "tvm_builtin_ptx_cvt_s16_u16_s64_u64"
+ assert "(int64_t& __d, uint64_t __a)" in source
+ assert '"cvt.s16.u16 %0, %1;" : "=l"(__d) : "l"(__a)' in source
+
+ # Destination widening remains supported, while the floating source is
+ # exact-width because ptxas rejects the wider source form permitted by PTX
+ # 9.2 Table 27.
+ tokens = tokens_for(entry, rnd="rn", dtype="f16", atype="f32")
+ opcode, helper, source = render_variant(entry, tokens, dtypes=("uint128", "uint32"))
+ assert opcode == "cvt.rn.f16.f32"
+ assert helper == "tvm_builtin_ptx_cvt_rn_f16_f32_u128_u32"
+ assert "(__uint128_t& __d, uint32_t __a)" in source
+ assert '"cvt.rn.f16.f32 %0, %1;" : "=q"(__d) : "r"(__a)' in source
+
+
+def test_cvt_packed_and_bf16_carriers_are_not_blanket_widened():
+ generic_bf16 = TABLE["cvt"]
+ bf16_map = mods(generic_bf16, tokens_for(generic_bf16, dtype="bf16", atype="bf16"))
+ assert operand_dtypes(generic_bf16.operands[0], bf16_map) == ("uint16",)
+ assert operand_dtypes(generic_bf16.operands[1], bf16_map) == ("uint16",)
+
+ packed = TABLE["cvt_f16x2_f32"]
+ packed_map = mods(packed, tokens_for(packed, rnd="rn", dtype="f16x2", atype="f32"))
+ assert operand_dtypes(packed.operands[0], packed_map) == ("uint32",)
+ assert operand_dtypes(packed.operands[1], packed_map) == ("float32",)
+
+ tf32 = TABLE["cvt_tf32_f32"]
+ tf32_map = mods(tf32, tokens_for(tf32, rnd="rna", dtype="tf32", atype="f32"))
+ assert operand_dtypes(tf32.operands[0], tf32_map) == ("uint32",)
+ assert operand_dtypes(tf32.operands[1], tf32_map) == ("float32",)
+
+
+def test_cvt_bf16_sat_diagnostics_distinguish_isa_and_toolchain():
+ entry = TABLE["cvt"]
+ with pytest.raises(ValueError, match=r"ISA limits floating-point \.sat destinations"):
+ tokens_for(entry, rnd="rn", sat="sat", dtype="bf16", atype="f32")
+ with pytest.raises(ValueError, match=r"toolchain assembles no \.sat when \.bf16 is the source"):
+ tokens_for(entry, sat="sat", dtype="f32", atype="bf16")
+
+
def test_cvt_e2m1x2_stages_its_b8_operand():
"""`.e2m1x2` is the one cvt format with no register of its own width.
- ISA 9.7.9.22:92 "When converting to .e2m1x2 data formats, the destination
+ ISA 9.7.9.21:92 "When converting to .e2m1x2 data formats, the destination
operand d has .b8 type." and :101 "When converting from .e2m1x2 to
.f16x2/.bf16x2, source operand a has .b8 type." Inline asm has no 8-bit
constraint letter, so both directions declare the register inside the block
@@ -311,7 +522,7 @@ def test_cvt_e2m1x2_stages_its_b8_operand():
assert "cvt.rn.scaled::n2::ue8m0.bf16x2.e2m1x2 %0, raw_a, %2;" in source
-# The cvt type tokens ISA 9.7.9.22's Target ISA Notes list architecture by
+# The cvt type tokens ISA 9.7.9.21's Target ISA Notes list architecture by
# architecture (:527-639), plus the .rs rounding mode (":602 .rs rounding mode
# is supported on following architectures:", listing sm_100a and sm_103a).
_CVT_BLACKWELL_TOKENS = {
@@ -349,7 +560,7 @@ def test_cvt_blackwell_lines_carry_their_arch_floor():
the sm_90 default would report legal forms as illegal.
The floor rides the narrow format, not `.bf16x2` on its own: ISA
- 9.7.9.22:517-518 puts `.bf16x2` as a destination format at "sm_80 or
+ 9.7.9.21:517-518 puts `.bf16x2` as a destination format at "sm_80 or
higher", which is where cvt.frnd2{.relu}{.satfinite}.bf16x2.f32 lives,
while :634-639 restrict `.bf16x2` *from* an fp8/fp6/fp4 format to
family-specific architectures.
@@ -361,7 +572,7 @@ def test_cvt_blackwell_lines_carry_their_arch_floor():
def test_cvt_tf32_satfinite_carries_its_sm_100_floor():
- """ISA 9.7.9.22:526 "cvt.{rn/rz}.satfinite.tf32.f32 requires sm_100 or
+ """ISA 9.7.9.21:526 "cvt.{rn/rz}.satfinite.tf32.f32 requires sm_100 or
higher." -- the maximum floor over the entry, whose other spellings sit at
sm_80/sm_90."""
entry = TABLE["cvt_tf32_f32"]
@@ -374,7 +585,7 @@ def test_cvt_rs_and_scale_factor_shapes():
"""The two operand shapes this family added: the .rs lines' trailing rbits
(with a grouped ``{a, b, e, f}`` source on the x4 forms), and the
scale-factor operand that exists exactly when .scaled::n2::ue8m0 is
- written (ISA 9.7.9.22:180-182 "Operand scale-factor and qualifier
+ written (ISA 9.7.9.21:180-182 "Operand scale-factor and qualifier
.scaled::n2::ue8m0 must be used together.")."""
_, _, source = render_variant(TABLE["cvt_rs_f16x2_f32"], ("rs", "", "", "f16x2", "f32"))
assert "uint32_t& __d, float __a, float __b, uint32_t __rbits" in source
diff --git a/tests/python/tirx/codegen/test_ptx_dialect.py b/tests/python/tirx/codegen/test_ptx_dialect.py
index f6902c2331..010fb9ed92 100644
--- a/tests/python/tirx/codegen/test_ptx_dialect.py
+++ b/tests/python/tirx/codegen/test_ptx_dialect.py
@@ -202,6 +202,91 @@ def test_ptx_predication_codegen():
assert "setp.ne.b32 p, %2, 0; @p red.relaxed.gpu.global.add.u32 [%0], %1;" in src
+def test_ptx_red_vector_codegen_and_roundtrip():
+ @T.prim_func
+ def kernel(a_ptr: T.handle):
+ A = T.match_buffer(a_ptr, (16,), "uint32")
+ T.device_entry()
+ T.cta_id([1])
+ tx = T.thread_id([32])
+ half = T.local_scalar("uint16")
+ packed = T.local_scalar("uint32")
+ value = T.local_scalar("float32")
+ policy = T.local_scalar("uint64")
+ T.ptx.red.relaxed.gpu.global_.add.noftz.v8.f16(
+ A.ptr_to([0]), half, half, half, half, half, half, half, half
+ )
+ T.ptx.red.release.sys.global_.max.noftz.L2__cache_hint.v4.f16x2(
+ A.ptr_to([1]), packed, packed, packed, packed, policy
+ )
+ T.ptx.red.global_.add.v4.f32(A.ptr_to([2]), value, value, value, value)
+ A[tx % 16] = A[tx % 16]
+
+ src = _cuda_source(kernel)
+ assert (
+ "red.relaxed.gpu.global.add.noftz.v8.f16 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};"
+ ) in src
+ assert (
+ "red.release.sys.global.max.noftz.L2::cache_hint.v4.f16x2 [%0], {%1, %2, %3, %4}, %5;"
+ ) in src
+ assert "red.global.add.v4.f32 [%0], {%1, %2, %3, %4};" in src
+
+ reparsed = tvm.script.from_source(kernel.script())
+ tvm.ir.assert_structural_equal(kernel, reparsed)
+
+ with pytest.raises((ValueError, tvm.error.DiagnosticError), match="already a 32-bit pair"):
+
+ @T.prim_func
+ def packed_v8(a_ptr: T.handle):
+ A = T.match_buffer(a_ptr, (1,), "uint32")
+ T.device_entry()
+ v = T.local_scalar("uint32")
+ T.ptx.red.global_.add.noftz.v8.f16x2(A.ptr_to([0]), v, v, v, v, v, v, v, v)
+
+ with pytest.raises((AttributeError, tvm.error.DiagnosticError), match="not a valid modifier"):
+
+ @T.prim_func
+ def f32_max(a_ptr: T.handle):
+ A = T.match_buffer(a_ptr, (1,), "float32")
+ T.device_entry()
+ v = T.local_scalar("float32")
+ T.ptx.red.global_.max.v2.f32(A.ptr_to([0]), v, v)
+
+
+def test_ptx_atom_bitbucket_codegen_and_roundtrip():
+ @T.prim_func
+ def kernel(a_ptr: T.handle):
+ A = T.match_buffer(a_ptr, (16,), "uint32")
+ T.device_entry()
+ T.cta_id([1])
+ tx = T.thread_id([32])
+ word = T.local_scalar("uint32")
+ half = T.local_scalar("uint16")
+ value = T.local_scalar("float32")
+ T.ptx.atom.global_.add.u32(A.ptr_to([0]), word)
+ T.ptx.atom.global_.cas.b32(A.ptr_to([1]), word, word)
+ T.ptx.atom.global_.exch.b32(A.ptr_to([2]), word)
+ T.ptx.atom.global_.add.noftz.f16(A.ptr_to([3]), half)
+ T.ptx.atom.global_.add.noftz.v2.f16(A.ptr_to([4]), half, half)
+ T.ptx.atom.global_.add.v2.f32(A.ptr_to([5]), value, value, pred=tx)
+ A[tx % 16] = A[tx % 16]
+
+ src = _cuda_source(kernel)
+ for text in (
+ "atom.global.add.u32 _, [%0], %1;",
+ "atom.global.cas.b32 _, [%0], %1, %2;",
+ "atom.global.exch.b32 _, [%0], %1;",
+ "atom.global.add.noftz.f16 _, [%0], %1;",
+ "atom.global.add.noftz.v2.f16 _, [%0], {%1, %2};",
+ "@p atom.global.add.v2.f32 _, [%0], {%1, %2};",
+ ):
+ assert text in src, text
+ assert "atom.global.add.v2.f32 {_," not in src
+
+ reparsed = tvm.script.from_source(kernel.script())
+ tvm.ir.assert_structural_equal(kernel, reparsed)
+
+
@requires_nvcc
def test_ptx_predicated_destination_preserves_old_value():
@T.prim_func
@@ -266,6 +351,139 @@ def test_ptx_string_form_matches_chain():
tvm.ir.assert_structural_equal(make(chain_call), make(string_call))
+def test_ptx_92_cp_bulk_exact_renderings():
+ """PTX 9.2 bulk-copy and non-tensor reduction forms render exactly."""
+ from tvm.backend.cuda.ptx.render import render_variant
+ from tvm.backend.cuda.ptx.table import TABLE, tokens_for
+
+ cases = (
+ (
+ "cp_async_bulk_g2s_cta",
+ dict(
+ api="async",
+ kind="bulk",
+ dst="shared::cta",
+ src="global",
+ completion="mbarrier::complete_tx::bytes",
+ cache="L2::cache_hint",
+ ignore_oob="ignore_oob",
+ ),
+ "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes."
+ "L2::cache_hint.ignore_oob [%0], [%1], %2, %3, %4, [%5], %6;",
+ ),
+ (
+ "cp_reduce_async_bulk_s2c",
+ dict(
+ op="reduce",
+ api="async",
+ kind="bulk",
+ dst="shared::cluster",
+ src="shared::cta",
+ completion="mbarrier::complete_tx::bytes",
+ redop="xor",
+ type="b32",
+ ),
+ "cp.reduce.async.bulk.shared::cluster.shared::cta."
+ "mbarrier::complete_tx::bytes.xor.b32 [%0], [%1], %2, [%3];",
+ ),
+ (
+ "cp_reduce_async_bulk_s2g",
+ dict(
+ op="reduce",
+ api="async",
+ kind="bulk",
+ dst="global",
+ src="shared::cta",
+ completion="bulk_group",
+ cache="L2::cache_hint",
+ redop="add",
+ noftz="noftz",
+ type="f16",
+ ),
+ "cp.reduce.async.bulk.global.shared::cta.bulk_group.L2::cache_hint."
+ "add.noftz.f16 [%0], [%1], %2, %3;",
+ ),
+ )
+ for name, modifiers, expected in cases:
+ entry = TABLE[name]
+ _, _, source = render_variant(entry, tokens_for(entry, **modifiers))
+ assert f'asm volatile("{expected}"' in source, source
+
+
+def test_ptx_92_cp_reduce_negative_grids():
+ from tvm.backend.cuda.ptx.table import TABLE, tokens_for
+ from tvm.ir.type import PointerType, PrimType
+
+ shared_reduce = TABLE["cp_reduce_async_bulk_s2c"]
+ with pytest.raises(ValueError, match=r"\.add to \.shared::cluster takes"):
+ tokens_for(
+ shared_reduce,
+ op="reduce",
+ api="async",
+ kind="bulk",
+ dst="shared::cluster",
+ src="shared::cta",
+ completion="mbarrier::complete_tx::bytes",
+ redop="add",
+ type="b32",
+ )
+
+ global_reduce = TABLE["cp_reduce_async_bulk_s2g"]
+ with pytest.raises(ValueError, match=r"add\.f16 requires \.noftz"):
+ tokens_for(
+ global_reduce,
+ op="reduce",
+ api="async",
+ kind="bulk",
+ dst="global",
+ src="shared::cta",
+ completion="bulk_group",
+ redop="add",
+ type="f16",
+ )
+
+ shared_ptr = tvm.tirx.Var("shared_ptr", PointerType(PrimType("uint32"), "shared"))
+ global_ptr = tvm.tirx.Var("global_ptr", PointerType(PrimType("uint32"), "global"))
+ size = tvm.tirx.Var("size", "uint32")
+ with pytest.raises(ValueError, match=r"\.add to \.shared::cluster takes"):
+ T.ptx[
+ "cp.reduce.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes.add.b32"
+ ](shared_ptr, shared_ptr, size, shared_ptr)
+ with pytest.raises(ValueError, match=r"add\.f16 requires \.noftz"):
+ T.ptx["cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f16"](
+ global_ptr, shared_ptr, size
+ )
+
+
+def test_ptx_92_cp_bulk_roundtrip():
+ @T.prim_func
+ def kernel(src: T.Buffer((64,), "uint32"), dst: T.Buffer((64,), "uint32")):
+ T.device_entry()
+ smem = T.alloc_buffer((64,), "uint32", scope="shared")
+ mbar = T.alloc_buffer((2,), "uint64", scope="shared")
+ T.ptx[
+ "cp.async.bulk.shared::cta.global."
+ "mbarrier::complete_tx::bytes.L2::cache_hint.ignore_oob"
+ ](
+ smem.ptr_to([0]),
+ src.ptr_to([0]),
+ T.uint32(16),
+ T.uint32(0),
+ T.uint32(0),
+ mbar.ptr_to([0]),
+ T.uint64(0),
+ )
+ T.ptx[
+ "cp.reduce.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes.xor.b32"
+ ](smem.ptr_to([16]), smem.ptr_to([0]), T.uint32(16), mbar.ptr_to([0]))
+ T.ptx["cp.reduce.async.bulk.global.shared::cta.bulk_group.add.noftz.f16"](
+ dst.ptr_to([0]), smem.ptr_to([0]), T.uint32(16)
+ )
+
+ reparsed = tvm.script.from_source(kernel.script())
+ tvm.ir.assert_structural_equal(kernel, reparsed)
+
+
def test_ptx_trace_time_errors():
# Global ld fed a raw uint32 address.
with pytest.raises((ValueError, tvm.error.DiagnosticError), match="shared state space"):
@@ -283,15 +501,16 @@ def test_ptx_trace_time_errors():
T.device_entry()
T.ptx.ld.global_.bogus.b32(out[0], T.uint32(0))
- # Value dtype of the wrong *width*. A bit type accepts any dtype of its own
- # width (see test_ptx_bit_width_axis), so the rejection is about size.
+ # A floating register is not compatible with an integer instruction type,
+ # even when it is wider (the relaxed ld/st rule admits a bit carrier, not a
+ # differently typed floating register).
with pytest.raises((ValueError, tvm.error.DiagnosticError), match="must have dtype"):
@T.prim_func
def bad_value_dtype(a_ptr: T.handle):
A = T.match_buffer(a_ptr, (1,), "uint32")
T.device_entry()
- T.ptx.st.global_.b32(A.ptr_to([0]), T.float64(1.0))
+ T.ptx.st.global_.u32(A.ptr_to([0]), T.float64(1.0))
# Missing required modifier (no type token).
with pytest.raises((ValueError, tvm.error.DiagnosticError), match="missing required modifier"):
@@ -323,14 +542,15 @@ def test_ptx_trace_time_errors():
def test_ptx_destination_errors():
"""A destination is a register the caller declared: it must be a writable lvalue."""
- # Destination of the wrong width (a .b32 load into a 64-bit slot).
+ # A floating destination is incompatible with an integer instruction type;
+ # relaxed widening does not numerically convert between the two.
with pytest.raises((ValueError, tvm.error.DiagnosticError), match="must have dtype"):
@T.prim_func
def wrong_dst_dtype(out: T.Buffer((1,), "float64"), a_ptr: T.handle):
A = T.match_buffer(a_ptr, (1,), "uint32")
T.device_entry()
- T.ptx.ld.global_.b32(out[0], A.ptr_to([0]))
+ T.ptx.ld.global_.u32(out[0], A.ptr_to([0]))
# A T.let binding is immutable, so it cannot be written into. This is the
# gate that keeps the analyzer from re-expanding one call into N.
@@ -488,6 +708,85 @@ def test_ptx_optional_operand_arity_dispatch():
tvm.ir.assert_structural_equal(kernel, reparsed)
+def test_ptx_mbarrier_92_shapes_render_and_roundtrip():
+ """PTX 9.2 noComplete sink/state and wait shapes render exactly."""
+
+ @T.prim_func
+ def kernel(out_ptr: T.handle):
+ out = T.match_buffer(out_ptr, (4,), "uint32")
+ T.device_entry()
+ T.cta_id([1])
+ tx = T.thread_id([32])
+ bar = T.alloc_buffer((2,), "uint64", scope="shared")
+ state = T.local_scalar("uint64")
+ pending = T.local_scalar("uint32")
+ wait_complete = T.local_scalar("uint32")
+
+ T.ptx.mbarrier.init.shared.b64(bar.ptr_to([0]), T.uint32(1))
+ T.ptx.mbarrier.arrive.noComplete.shared.b64(bar.ptr_to([0]), T.uint32(1))
+ T.ptx.mbarrier.arrive_drop.noComplete.shared__cta.b64(state, bar.ptr_to([1]), T.uint32(1))
+ T.ptx.mbarrier.pending_count.b64(pending, state)
+ T.ptx.mbarrier.test_wait.shared__cta.b64(wait_complete, bar.ptr_to([0]), state)
+ T.ptx.mbarrier.try_wait.parity.relaxed.cluster.shared__cta.b64(
+ wait_complete, bar.ptr_to([1]), T.uint32(0), T.uint32(20)
+ )
+ out[tx % 4] = pending + wait_complete
+
+ from tvm.backend.cuda.ptx.render import render_variant
+ from tvm.backend.cuda.ptx.table import TABLE
+
+ selected = (
+ ("mbarrier_init", ("init", "shared", "b64"), False),
+ (
+ "mbarrier_arrive_no_complete_sink",
+ ("arrive", "noComplete", "", "", "shared", "b64"),
+ False,
+ ),
+ (
+ "mbarrier_arrive_drop_no_complete",
+ ("arrive_drop", "noComplete", "", "", "shared::cta", "b64"),
+ False,
+ ),
+ ("mbarrier_pending_count", ("pending_count", "b64"), False),
+ ("mbarrier_test_wait", ("test_wait", "", "", "shared::cta", "b64"), False),
+ (
+ "mbarrier_try_wait_parity",
+ ("try_wait", "parity", "relaxed", "cluster", "shared::cta", "b64"),
+ False,
+ ),
+ )
+ src = "\n".join(
+ render_variant(TABLE[name], tokens, predicated)[2] for name, tokens, predicated in selected
+ )
+ for text in (
+ "mbarrier.init.shared.b64 [%0], %1;",
+ "mbarrier.arrive.noComplete.shared.b64 _, [%0], %1;",
+ "mbarrier.arrive_drop.noComplete.shared::cta.b64 %0, [%1], %2;",
+ "mbarrier.pending_count.b64 %0, %1;",
+ "mbarrier.test_wait.shared::cta.b64 pd0, [%1], %2;",
+ "mbarrier.try_wait.parity.relaxed.cluster.shared::cta.b64 pd0, [%1], %2, %3;",
+ ):
+ assert text in src, text
+
+ reparsed = tvm.script.from_source(kernel.script())
+ tvm.ir.assert_structural_equal(kernel, reparsed)
+
+
+def test_ptx_mbarrier_92_state_carriers():
+ from tvm.backend.cuda.ptx.table import TABLE, mods, operand_dtypes, variants
+
+ pending = TABLE["mbarrier_pending_count"]
+ for tokens in variants(pending):
+ mod_map = mods(pending, tokens)
+ assert operand_dtypes(pending.operands[1], mod_map) == ("uint64", "int64")
+
+ for name in ("mbarrier_arrive_no_complete", "mbarrier_arrive_drop_no_complete"):
+ entry = TABLE[name]
+ for tokens in variants(entry):
+ mod_map = mods(entry, tokens)
+ assert operand_dtypes(entry.operands[0], mod_map) == ("uint64", "int64")
+
+
def test_ptx_bit_width_axis():
"""A `.bN` operand takes any dtype of that width, each with its own helper.
@@ -531,6 +830,173 @@ def test_ptx_bit_width_axis():
assert render_variant(ld, tokens)[1] == "tvm_builtin_ptx_ld_global_b32"
+def test_ptx_relaxed_load_store_typing():
+ """Scalar/vector ld, st and ldu accept ISA 9.4.1's wider register carriers."""
+
+ @T.prim_func
+ def kernel(a_ptr: T.handle):
+ A = T.match_buffer(a_ptr, (8,), "uint64")
+ T.device_entry()
+ T.cta_id([1])
+ tx = T.thread_id([32])
+ wide = T.local_scalar("uint64")
+ signed0 = T.local_scalar("int64")
+ signed1 = T.local_scalar("int64")
+ T.ptx.ld.global_.b8(wide, A.ptr_to([0]))
+ T.ptx.ld.global_.v2.s8(signed0, signed1, A.ptr_to([1]))
+ T.ptx.st.global_.b8(A.ptr_to([2]), wide)
+ T.ptx.st.global_.v2.u16(A.ptr_to([3]), wide, wide)
+ T.ptx.ldu.global_.u16(wide, A.ptr_to([4]))
+ T.ptx.ldu.global_.v2.s8(signed0, signed1, A.ptr_to([5]))
+ A[tx % 8] = wide + T.uint64(signed0 + signed1)
+
+ src = _cuda_source(kernel)
+ for text in (
+ "ld.global.b8 %0, [%1];",
+ "ld.global.v2.s8 {%0, %1}, [%2];",
+ "st.global.b8 [%0], %1;",
+ "st.global.v2.u16 [%0], {%1, %2};",
+ "ldu.global.u16 %0, [%1];",
+ "ldu.global.v2.s8 {%0, %1}, [%2];",
+ ):
+ assert text in src, text
+ # All data operands above use 64-bit registers; widening is expressed by
+ # the inline-asm constraint and does not insert a numeric conversion.
+ assert '"=l"(__d)' in src or '"=l"(__d0)' in src
+ assert '"l"(__value)' in src or '"l"(__value0)' in src
+ assert "cvt." not in src
+
+ with pytest.raises((ValueError, tvm.error.DiagnosticError), match="must have dtype"):
+
+ @T.prim_func
+ def floating_source_for_integer_type(a_ptr: T.handle):
+ A = T.match_buffer(a_ptr, (1,), "uint16")
+ T.device_entry()
+ T.ptx.st.global_.u16(A.ptr_to([0]), T.float32(1))
+
+
+def test_ptx_st_vec_f64_cuda_13_2_carrier_gap():
+ """Keep CUDA 13.2's st.v2.f64 gap local to that source-operand shape."""
+ from tvm.backend.cuda.ptx.table import TABLE, dtype_combos, tokens_for
+
+ st_vec_tokens = tokens_for(TABLE["st_vec"], vec="v2", type="f64")
+ assert dtype_combos(TABLE["st_vec"], st_vec_tokens) == (
+ ("float64", "uint64"),
+ ("uint64", "uint64"),
+ ("int64", "uint64"),
+ )
+
+ for name, tokens in (
+ ("st", tokens_for(TABLE["st"], type="f64")),
+ ("ld_vec", tokens_for(TABLE["ld_vec"], vec="v2", type="f64")),
+ ("ldu_vec", tokens_for(TABLE["ldu_vec"], vec="v2", type="f64")),
+ ):
+ combos = dtype_combos(TABLE[name], tokens)
+ assert any(combo[0] == "uint128" for combo in combos), name
+ assert any(combo[0] == "int128" for combo in combos), name
+
+
+def test_ptx_vec256_cache_policy():
+ """Pin the PTX 9.2 256-bit cache-policy arity and exact-width carriers."""
+ from tvm.backend.cuda.ptx.render import render_variant
+ from tvm.backend.cuda.ptx.table import TABLE, mods, operand_dtypes, tokens_for
+
+ ld256 = TABLE["ld_vec256"]
+ ld_tokens = tokens_for(ld256, space="global", cache="L2::cache_hint", vec="v8", type="b32")
+ ld_src = render_variant(ld256, ld_tokens)[2]
+ assert "uint64_t __cache_policy" in ld_src
+ assert "ld.global.L2::cache_hint.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8], %9;" in ld_src
+
+ st256 = TABLE["st_vec256"]
+ st_tokens = tokens_for(st256, space="global", cache="L2::cache_hint", vec="v8", type="b32")
+ st_src = render_variant(st256, st_tokens)[2]
+ assert "uint64_t __cache_policy" in st_src
+ assert "st.global.L2::cache_hint.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8}, %9;" in st_src
+ assert operand_dtypes(st256.operands[1], mods(st256, st_tokens)) == (
+ "uint32",
+ "int32",
+ "float32",
+ )
+
+ @T.prim_func
+ def vec256_calls(a_ptr: T.handle):
+ A = T.match_buffer(a_ptr, (8,), "uint32")
+ T.device_entry()
+ policy = T.local_scalar("uint64")
+ x0 = T.local_scalar("uint32")
+ x1 = T.local_scalar("uint32")
+ x2 = T.local_scalar("uint32")
+ x3 = T.local_scalar("uint32")
+ x4 = T.local_scalar("uint32")
+ x5 = T.local_scalar("uint32")
+ x6 = T.local_scalar("uint32")
+ x7 = T.local_scalar("uint32")
+ T.ptx.ld.global_.L2__cache_hint.v8.b32(
+ x0, x1, x2, x3, x4, x5, x6, x7, A.ptr_to([0]), policy
+ )
+ T.ptx.st.global_.L2__cache_hint.v8.b32(
+ A.ptr_to([0]), x0, x1, x2, x3, x4, x5, x6, x7, policy
+ )
+
+ reparsed = tvm.script.from_source(vec256_calls.script())
+ tvm.ir.assert_structural_equal(vec256_calls, reparsed)
+
+
+@requires_nvcc
+def test_ptx_st_bulk_size_carriers_and_st_async_byte_bridge():
+ """st.bulk takes 32/64-bit sizes; st.async stages one private .b8 register."""
+ from tvm.backend.cuda.ptx.render import render_variant
+ from tvm.backend.cuda.ptx.table import TABLE, dtype_combos, tokens_for
+
+ bulk = TABLE["st_bulk"]
+ bulk_tokens = tokens_for(bulk, space="shared::cta")
+ assert dtype_combos(bulk, bulk_tokens) == (
+ ("uint64",),
+ ("int64",),
+ ("uint32",),
+ ("int32",),
+ )
+ assert "uint32_t __size" in render_variant(bulk, bulk_tokens, dtypes=("uint32",))[2]
+
+ release = TABLE["st_async_release"]
+ for ty, dtype in (
+ ("b8", "uint8"),
+ ("u8", "uint8"),
+ ("s8", "int8"),
+ ):
+ tokens = tokens_for(release, sem="release", scope="sys", space="global", type=ty)
+ source = render_variant(release, tokens, dtypes=(dtype,))[2]
+ assert ".reg .b8 raw_b;" in source
+ assert "cvt.u8.u16 raw_b, %1;" in source
+ assert f"st.async.release.sys.global.{ty} [%0], raw_b;" in source
+ asm_text = _ASM_RE.findall(source)[0]
+ assert _sole_instruction(asm_text) == f"st.async.release.sys.global.{ty} [%0], raw_b;"
+
+ # The C-boundary casts preserve all byte patterns; the bridge conversion
+ # consumes their low 8 bits in the exact register class named above.
+ for bits in (0x00, 0x7F, 0x80, 0xFF):
+ signed_value = bits if bits < 0x80 else bits - 0x100
+ assert (bits & 0xFF) == bits
+ assert (signed_value & 0xFF) == bits
+
+ @T.prim_func
+ def carrier_calls(a_ptr: T.handle):
+ A = T.match_buffer(a_ptr, (8,), "uint8")
+ T.device_entry()
+ T.cta_id([1])
+ T.thread_id([1])
+ T.ptx.st_async.release.sys.global_.b8(A.ptr_to([0]), T.uint8(0x00))
+ T.ptx.st_async.release.sys.global_.b8(A.ptr_to([1]), T.uint8(0x7F))
+ T.ptx.st_async.release.sys.global_.b8(A.ptr_to([2]), T.uint8(0x80))
+ T.ptx.st_async.release.sys.global_.b8(A.ptr_to([3]), T.uint8(0xFF))
+ T.ptx.st_async.release.sys.global_.u8(A.ptr_to([4]), T.uint8(0xFF))
+ T.ptx.st_async.release.sys.global_.s8(A.ptr_to([5]), T.int8(-1))
+
+ reparsed = tvm.script.from_source(carrier_calls.script())
+ tvm.ir.assert_structural_equal(carrier_calls, reparsed)
+ _assert_ptxas_ok(_cuda_source(carrier_calls), arch="sm_100")
+
+
def test_ptx_u64_address_handle_is_reinterpreted():
"""A 64-bit address handle is accepted via an explicit, visible cast.
@@ -609,6 +1075,31 @@ def test_ptx_integer_arithmetic_dispatch():
T.device_entry()
T.ptx.mad.lo.sat.s32(out[0], T.int32(2), T.int32(3), T.int32(4))
+ # The unavailable add.sat forms added in PTX 9.2 are distinct from the
+ # integer lines on which `.sat` has never been legal. Keep both reasons
+ # visible so an invalid spelling is not misreported as an architecture
+ # gate. In particular, sub.sat.u32 is not the counterpart of add.sat.u32.
+ with pytest.raises((ValueError, tvm.error.DiagnosticError), match="sm_120f-only"):
+
+ @T.prim_func
+ def add_sat_sm120(out: T.Buffer((1,), "uint32")):
+ T.device_entry()
+ T.ptx.add.sat.u32(out[0], T.uint32(2), T.uint32(3))
+
+ with pytest.raises((ValueError, tvm.error.DiagnosticError), match=r"not on the add\.u64"):
+
+ @T.prim_func
+ def add_sat_no_syntax_line(out: T.Buffer((1,), "uint64")):
+ T.device_entry()
+ T.ptx.add.sat.u64(out[0], T.uint64(2), T.uint64(3))
+
+ with pytest.raises((ValueError, tvm.error.DiagnosticError), match=r"not on the sub\.u32"):
+
+ @T.prim_func
+ def sub_sat_no_syntax_line(out: T.Buffer((1,), "uint32")):
+ T.device_entry()
+ T.ptx.sub.sat.u32(out[0], T.uint32(2), T.uint32(3))
+
def test_ptx_floating_point_dispatch():
"""ISA 9.7.3's lines, including three mnemonics shared with the integer group.
@@ -872,6 +1363,9 @@ def test_ptx_comparison_selection_dispatch():
T.ptx.set.lt.u32.f32(d, f, f) # writes a value, not a predicate
T.ptx.selp.b32(d, d, p, T.ptx.pred(q)) # predicate selects
T.ptx.slct.ftz.b32.f32(d, d, p, f) # a sign selects
+ # slct treats d/a/b independently as bit-size values. This mixes all
+ # three 32-bit carrier classes while c remains exactly .s32.
+ T.ptx.slct.f32.s32(d, f, A[0], A[1])
A[tx % 4] = T.int32(p + q + d)
src = _cuda_source(kernel)
@@ -883,6 +1377,7 @@ def test_ptx_comparison_selection_dispatch():
"set.lt.u32.f32 %0, %1, %2;",
"selp.b32 %0, %1, %2, ps0;",
"slct.ftz.b32.f32 %0, %1, %2, %3;",
+ "slct.f32.s32 %0, %1, %2, %3;",
):
assert text in src, text
@@ -931,6 +1426,56 @@ def test_ptx_comparison_selection_dispatch():
T.device_entry()
T.ptx.slct.ftz.b32.s32(out[0], T.uint32(1), T.uint32(2), T.int32(3))
+ # Only d/a/b receive slct's relaxed bit-size typing. The selector must
+ # still match the second instruction type exactly, even when another
+ # same-constraint integer carrier would fit through inline asm.
+ with pytest.raises(
+ (ValueError, tvm.error.DiagnosticError), match=r"operand 'c'.*int32.*uint32"
+ ):
+
+ @T.prim_func
+ def relaxed_selector(out: T.Buffer((1,), "uint32")):
+ T.device_entry()
+ T.ptx.slct.f32.s32(out[0], T.float32(1), T.int32(2), T.uint32(3))
+
+
+def test_ptx_slct_relaxed_value_dtype_domain():
+ """The ptxas-proven d/a/b carrier product is complete and native-first."""
+ from tvm.backend.cuda.ptx.table import (
+ PTX_TYPE_DTYPES,
+ TABLE,
+ dtype_combos,
+ mods,
+ operand_dtypes,
+ tokens_for,
+ variants,
+ )
+
+ expected = {
+ "b16": ("uint16", "int16", "float16", "bfloat16"),
+ "u16": ("uint16", "int16", "float16", "bfloat16"),
+ "s16": ("int16", "uint16", "float16", "bfloat16"),
+ "b32": ("uint32", "int32", "float32"),
+ "u32": ("uint32", "int32"),
+ "s32": ("int32", "uint32"),
+ "f32": ("float32", "uint32", "int32"),
+ "b64": ("uint64", "int64", "float64"),
+ "u64": ("uint64", "int64"),
+ "s64": ("int64", "uint64"),
+ "f64": ("float64", "uint64", "int64"),
+ }
+ entry = TABLE["slct"]
+ for dtype, value_dtypes in expected.items():
+ for ctype in ("s32", "f32"):
+ tokens = tokens_for(entry, dtype=dtype, ctype=ctype)
+ mod_map = mods(entry, tokens)
+ domains = tuple(operand_dtypes(slot, mod_map) for slot in entry.typed_operands)
+ assert domains == (value_dtypes, value_dtypes, value_dtypes, PTX_TYPE_DTYPES[ctype])
+
+ # The three value operands form an independent product. Across the legal
+ # (dtype, ctype, ftz) variants this is the exhaustively certified 996.
+ assert sum(len(dtype_combos(entry, tokens)) for tokens in variants(entry)) == 996
+
def test_ptx_half_comparison_dispatch():
"""ISA 9.7.7, the half twin of 9.7.6 -- same two mnemonics, different grids.
@@ -959,6 +1504,8 @@ def test_ptx_half_comparison_dispatch():
T.ptx.setp.eq.and_.bf16(p, h, h, T.ptx.pred(q))
T.ptx.setp.gt.or_.bf16x2(p, q, x2, x2, T.ptx.pred(d32))
T.ptx.set.lt.u32.f16(d32, h, h) # integer answer to a half compare
+ T.ptx.set.lt.u32.bf16(d32, h, h) # bf16 is floating, not a .b* bit type
+ T.ptx.set.nan.u32.bf16x2(d32, x2, x2) # nor is packed bf16x2
T.ptx.set.gt.f16.f32(h, T.float32(1.0), T.float32(2.0)) # ... and back
T.ptx.set.eq.f16x2.f16x2(x2, x2, x2) # packed both sides
A[tx % 4] = p + q + d32 + T.uint32(h) + x2
@@ -970,6 +1517,8 @@ def test_ptx_half_comparison_dispatch():
"setp.eq.and.bf16 pd0, %1, %2, ps0;",
"setp.gt.or.bf16x2 pd0|pd1, %2, %3, ps0;",
"set.lt.u32.f16 %0, %1, %2;",
+ "set.lt.u32.bf16 %0, %1, %2;",
+ "set.nan.u32.bf16x2 %0, %1, %2;",
"set.gt.f16.f32 %0, %1, %2;",
"set.eq.f16x2.f16x2 %0, %1, %2;",
):
@@ -1018,6 +1567,15 @@ def test_ptx_half_comparison_dispatch():
T.device_entry()
T.ptx.set.equ.f16.s32(out[0], T.int32(0), T.int32(0))
+ # The classification fix is narrow: a real .b16 source remains a bit-size
+ # value and therefore still permits equality comparisons only.
+ with pytest.raises((ValueError, tvm.error.DiagnosticError), match="bit-size source"):
+
+ @T.prim_func
+ def ordered_on_bitsize(out: T.Buffer((1,), "uint16")):
+ T.device_entry()
+ T.ptx.set.lt.f16.b16(out[0], T.uint16(0), T.uint16(0))
+
# 9.7.7 spells no unsigned alternates at all, so `lo` never resolves here.
with pytest.raises((AttributeError, tvm.error.DiagnosticError), match="not a valid modifier"):
@@ -1059,6 +1617,7 @@ def test_ptx_logic_shift_dispatch():
T.ptx.cnot.b32(d, d)
T.ptx.lop3.b32(d, d, d, d, 0x80) # a & b & c, by look-up table
T.ptx.lop3.or_.b32(d, p, d, d, d, 0xFE, T.ptx.pred(q)) # d|p pair
+ T.ptx.lop3.and_.b32(p, d, d, d, 0x80, T.ptx.pred(q)) # fixed _|p sibling
T.ptx.shf.l.clamp.b32(d, d, d, T.uint32(4)) # funnel shift
T.ptx.shl.b32(d, d, T.uint32(2))
T.ptx.shr.s32(s, s, T.uint32(1)) # signed: fills with the sign bit
@@ -1075,6 +1634,7 @@ def test_ptx_logic_shift_dispatch():
"cnot.b32 %0, %1;",
"lop3.b32 %0, %1, %2, %3, 128;",
"lop3.or.b32 %0|pd0, %2, %3, %4, 254, ps0;",
+ "lop3.and.b32 _|pd0, %1, %2, %3, 128, ps0;",
"shf.l.clamp.b32 %0, %1, %2, %3;",
"shl.b32 %0, %1, %2;",
"shr.s32 %0, %1, %2;",
@@ -1142,6 +1702,55 @@ def test_ptx_logic_shift_dispatch():
assert "lop3.b32 %0, %1, %2, %3, 0;" in unrolled_src
assert "lop3.b32 %0, %1, %2, %3, 128;" in unrolled_src
+ @T.prim_func
+ def lut_boundaries(a_ptr: T.handle):
+ A = T.match_buffer(a_ptr, (1,), "uint32")
+ T.device_entry()
+ T.cta_id([1])
+ T.thread_id([1])
+ d = T.local_scalar("uint32")
+ T.ptx.lop3.b32(d, A[0], A[0], A[0], 0)
+ T.ptx.lop3.b32(d, A[0], A[0], A[0], 255)
+
+ boundary_src = _cuda_source(lut_boundaries)
+ assert "lop3.b32 %0, %1, %2, %3, 0;" in boundary_src
+ assert "lop3.b32 %0, %1, %2, %3, 255;" in boundary_src
+
+ with pytest.raises(
+ (ValueError, tvm.error.DiagnosticError), match=r"inclusive range 0\.\.255, got -1"
+ ):
+
+ @T.prim_func
+ def lut_below_range(a_ptr: T.handle):
+ A = T.match_buffer(a_ptr, (1,), "uint32")
+ T.device_entry()
+ d = T.local_scalar("uint32")
+ T.ptx.lop3.b32(d, A[0], A[0], A[0], -1)
+
+ with pytest.raises(
+ (ValueError, tvm.error.DiagnosticError), match=r"inclusive range 0\.\.255, got 256"
+ ):
+
+ @T.prim_func
+ def lut_above_range(a_ptr: T.handle):
+ A = T.match_buffer(a_ptr, (1,), "uint32")
+ T.device_entry()
+ d = T.local_scalar("uint32")
+ T.ptx.lop3.b32(d, A[0], A[0], A[0], 256)
+
+ @T.prim_func
+ def lut_unrolled_out_of_range(a_ptr: T.handle):
+ A = T.match_buffer(a_ptr, (1,), "uint32")
+ T.device_entry()
+ T.cta_id([1])
+ T.thread_id([1])
+ d = T.local_scalar("uint32")
+ for i in T.unroll(2):
+ T.ptx.lop3.b32(d, A[0], A[0], A[0], i * 256)
+
+ with pytest.raises(ValueError, match=r"inclusive range 0\.\.255, got 256"):
+ _cuda_source(lut_unrolled_out_of_range)
+
def test_ptx_data_movement_dispatch():
"""ISA 9.7.9's newly registered instructions, end to end.
@@ -1274,13 +1883,13 @@ def test_ptx_data_movement_dispatch():
def test_ptx_parallel_sync_dispatch():
- """ISA 9.7.14's warp-level primitives and the atom shapes beside `.op`.
+ """ISA 9.7.13's warp-level primitives and the atom shapes beside `.op`.
The section is where predicates are most load-bearing: they are sources
(bar.red's `c`), destinations (vote, elect), and both at once. It is also
- where one mnemonic carries the most shapes -- `atom` now has the `.op`
- line, `.cas`, `.exch`, the half-precision adds and two vector lines, all
- resolved by tokens and arity alone.
+ where one mnemonic carries the most shapes -- `atom` has the `.op` line,
+ `.cas`, `.exch`, the half-precision adds and three vector syntax lines
+ represented by two entries, all resolved by tokens and arity alone.
"""
@T.prim_func
@@ -1431,6 +2040,89 @@ def test_ptx_pred_operand_roundtrip():
tvm.ir.assert_structural_equal(kernel, reparsed)
+def test_ptx_wgmma_scale_d_runtime_predicate_roundtrip():
+ """WGMMA scale-d is a runtime predicate, not a 0/1 text immediate."""
+
+ @T.prim_func
+ def kernel(out_ptr: T.handle):
+ Out = T.match_buffer(out_ptr, (128,), "uint32")
+ T.device_entry()
+ T.cta_id([1])
+ tx = T.thread_id([128])
+ scale_d: T.uint32 = T.cast(tx < 128, "uint32")
+ d0: T.uint32 = T.uint32(0)
+ d1: T.uint32 = T.uint32(0)
+ d2: T.uint32 = T.uint32(0)
+ d3: T.uint32 = T.uint32(0)
+ T.ptx.wgmma.mma_async.sync.aligned.m64n8k32.s32.s8.s8(
+ d0, d1, d2, d3, T.uint64(0), T.uint64(0), T.ptx.pred(scale_d)
+ )
+ Out[tx] = d0 + d1 + d2 + d3
+
+ target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"})
+ with target:
+ mod = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx")
+ src = mod.mod.imports[0].inspect_source("cuda")
+ assert "wgmma.mma_async.sync.aligned.m64n8k32.s32.s8.s8" in src
+ assert ".reg .pred ps0;" in src
+ assert "setp.ne.b32 ps0," in src
+ assert "}, %4, %5, ps0;" in src
+
+ reparsed = tvm.script.from_source(kernel.script())
+ tvm.ir.assert_structural_equal(kernel, reparsed)
+
+ # There is deliberately no compatibility overload for the old integer
+ # immediate spelling. A bool already identifies a predicate; a runtime
+ # integer must carry the explicit T.ptx.pred(...) register-class marker.
+ with pytest.raises((ValueError, tvm.error.DiagnosticError), match=r"T\.ptx\.pred"):
+
+ @T.prim_func
+ def bare_integer_scale_d():
+ T.device_entry()
+ d0: T.uint32 = T.uint32(0)
+ d1: T.uint32 = T.uint32(0)
+ d2: T.uint32 = T.uint32(0)
+ d3: T.uint32 = T.uint32(0)
+ T.ptx.wgmma.mma_async.sync.aligned.m64n8k32.s32.s8.s8(
+ d0, d1, d2, d3, T.uint64(0), T.uint64(0), 0
+ )
+
+
+def test_ptx_wgmma_integer_shape_domains_follow_concrete_syntax():
+ """s8/u8 stop at N=224, while the b1 syntax includes N=240/256."""
+ from tvm.backend.cuda.ptx.render import render_variant
+ from tvm.backend.cuda.ptx.table import TABLE, lanes_of, mods, tokens_for
+
+ for form in ("ss", "rs"):
+ entry = TABLE[f"wgmma_int_{form}"]
+ shape_slot = next(slot for slot in entry.slots if slot.name == "shape")
+ assert "m64n224k32" in shape_slot.choices
+ assert "m64n240k32" not in shape_slot.choices
+ assert "m64n256k32" not in shape_slot.choices
+ assert entry.operands[-1].dtype == "pred"
+
+ tokens = tokens_for(
+ entry,
+ action="mma_async",
+ sync="sync",
+ aligned="aligned",
+ shape="m64n224k32",
+ dtype="s32",
+ atype="s8",
+ btype="s8",
+ )
+ assert lanes_of(entry.operands[0], mods(entry, tokens)) == 112
+ opcode, _, source = render_variant(entry, tokens)
+ assert opcode == "wgmma.mma_async.sync.aligned.m64n224k32.s32.s8.s8"
+ assert f"{opcode} {{" in source
+ assert ", ps0;" in source
+
+ b1_entry = TABLE[f"wgmma_b1_{form}"]
+ b1_shapes = next(slot for slot in b1_entry.slots if slot.name == "shape").choices
+ assert "m64n240k256" in b1_shapes
+ assert "m64n256k256" in b1_shapes
+
+
@requires_nvcc
def test_ptx_tcgen05_mma_block_size_form():
@T.prim_func
@@ -1585,6 +2277,160 @@ def test_ptx_printer_form():
# ---------------------------------------------------------------------------
+def test_ptx_tcgen05_mapa_address_rendering():
+ """Optional state spaces select generic versus shared address carriers."""
+ from tvm.backend.cuda.ptx.render import render_variant
+ from tvm.backend.cuda.ptx.table import TABLE, tokens_for, variants
+
+ def render(name, **by_name):
+ entry = TABLE[name]
+ return render_variant(entry, tokens_for(entry, **by_name))[2]
+
+ bare_alloc = render(
+ "tcgen05_alloc",
+ action="alloc",
+ cta_group="cta_group::1",
+ sync="sync",
+ aligned="aligned",
+ type="b32",
+ )
+ shared_alloc = render(
+ "tcgen05_alloc",
+ action="alloc",
+ cta_group="cta_group::1",
+ sync="sync",
+ aligned="aligned",
+ space="shared::cta",
+ type="b32",
+ )
+ assert "(const void* __dst, uint32_t __ncols)" in bare_alloc
+ assert '"l"(__dst), "r"(__ncols)' in bare_alloc
+ assert "tcgen05.alloc.cta_group::1.sync.aligned.b32 [%0], %1;" in bare_alloc
+ assert "(uint32_t __dst, uint32_t __ncols)" in shared_alloc
+ assert '"r"(__dst), "r"(__ncols)' in shared_alloc
+ assert "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" in shared_alloc
+
+ for name, extra in (
+ ("tcgen05_commit", {}),
+ ("tcgen05_commit_multicast", {"multicast": "multicast::cluster"}),
+ ):
+ required = {
+ "action": "commit",
+ "cta_group": "cta_group::2",
+ "completion": "mbarrier::arrive::one",
+ "type": "b64",
+ **extra,
+ }
+ bare = render(name, **required)
+ shared = render(name, space="shared::cluster", **required)
+ assert "(const void* __mbar" in bare
+ assert '"l"(__mbar)' in bare
+ assert "(uint32_t __mbar" in shared
+ assert '"r"(__mbar)' in shared
+
+ generic_mapa = render("mapa", type="u64")
+ shared_mapa_ptr = render("mapa", space="shared::cluster", type="u64")
+ raw_mapa = render("mapa_u64_raw", type="u64")
+ shared_mapa_u64 = render("mapa_u64_shared", space="shared::cluster", type="u64")
+ shared_mapa_u32 = render("mapa_u32", space="shared::cluster", type="u32")
+ assert "(uint64_t& __d, const void* __a, uint32_t __b)" in generic_mapa
+ assert '"=l"(__d) : "l"(__a), "r"(__b)' in generic_mapa
+ assert "mapa.u64 %0, %1, %2;" in generic_mapa
+ assert "(uint64_t& __d, const void* __a, uint32_t __b)" in shared_mapa_ptr
+ assert "mapa.shared::cluster.u64 %0, %1, %2;" in shared_mapa_ptr
+ assert "(uint64_t& __d, uint64_t __a, uint32_t __b)" in raw_mapa
+ assert '"=l"(__d) : "l"(__a), "r"(__b)' in raw_mapa
+ assert "mapa.u64 %0, %1, %2;" in raw_mapa
+ assert "(uint64_t& __d, uint64_t __a, uint32_t __b)" in shared_mapa_u64
+ assert '"=l"(__d) : "l"(__a), "r"(__b)' in shared_mapa_u64
+ assert "mapa.shared::cluster.u64 %0, %1, %2;" in shared_mapa_u64
+ assert "(uint32_t& __d, uint32_t __a, uint32_t __b)" in shared_mapa_u32
+
+ # Two entries share bare mapa.u64's ISA spelling but accept disjoint
+ # pointer/register call shapes; explicit shared has both carrier shapes.
+ assert (
+ sum(
+ len(variants(TABLE[name]))
+ for name in ("mapa", "mapa_u64_raw", "mapa_u64_shared", "mapa_u32")
+ )
+ == 5
+ )
+
+
+def test_ptx_tcgen05_mapa_address_coercion():
+ """Bare forms preserve generic pointers; explicit shared forms coerce only addresses."""
+ from tvm.ir.type import PointerType, PrimType
+
+ generic_ptr = tvm.tirx.Var("g", PointerType(PrimType("uint64"), "global"))
+ shared_ptr = tvm.tirx.Var("s", PointerType(PrimType("uint64"), "shared"))
+ raw_u32 = tvm.tirx.Var("a32", "uint32")
+ raw_u64 = tvm.tirx.Var("a64", "uint64")
+ ncols = tvm.tirx.Var("n", "uint32")
+ mask = tvm.tirx.Var("mask", "uint16")
+ out32 = tvm.tirx.decl_buffer((1,), "uint32", name="out32", scope="local")
+ out64 = tvm.tirx.decl_buffer((1,), "uint64", name="out64", scope="local")
+
+ bare_alloc = T.ptx.tcgen05.alloc.cta_group__1.sync.aligned.b32(generic_ptr, ncols)
+ assert bare_alloc.args[0].same_as(generic_ptr)
+ shared_alloc = T.ptx.tcgen05.alloc.cta_group__1.sync.aligned.shared__cta.b32(shared_ptr, ncols)
+ assert shared_alloc.args[0].op.name == "tirx.cuda.cvta_generic_to_shared"
+ assert shared_alloc.args[0].args[0].same_as(shared_ptr)
+
+ bare_commit = T.ptx.tcgen05.commit.cta_group__1.mbarrier__arrive__one.b64(generic_ptr)
+ assert bare_commit.args[0].same_as(generic_ptr)
+ shared_commit = T.ptx[
+ "tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64"
+ ](raw_u32, mask)
+ assert shared_commit.args[0].same_as(raw_u32)
+
+ generic_mapa = T.ptx.mapa.u64(out64[0], generic_ptr, ncols)
+ assert generic_mapa.args[1].same_as(generic_ptr)
+ raw_mapa = T.ptx.mapa.u64(out64[0], raw_u64, ncols)
+ assert raw_mapa.args[1].same_as(raw_u64)
+ shared_mapa_u64 = T.ptx.mapa.shared__cluster.u64(out64[0], raw_u64, ncols)
+ assert shared_mapa_u64.args[1].same_as(raw_u64)
+ shared_mapa_u32 = T.ptx.mapa.shared__cluster.u32(out32[0], raw_u32, ncols)
+ assert shared_mapa_u32.args[1].same_as(raw_u32)
+
+ shared_mapa_ptr = T.ptx.mapa.shared__cluster.u64(out64[0], shared_ptr, ncols)
+ assert shared_mapa_ptr.args[1].same_as(shared_ptr)
+
+
+def test_ptx_tcgen05_mapa_address_roundtrip():
+ """The generic/shared split remains exact through TVMScript print and parse."""
+
+ @T.prim_func
+ def kernel(ptr: T.handle):
+ generic = T.match_buffer(ptr, (4,), "uint64")
+ T.device_entry()
+ mapped32 = T.local_scalar("uint32")
+ mapped64 = T.local_scalar("uint64")
+ window64 = T.local_scalar("uint64")
+ T.ptx.tcgen05.alloc.cta_group__1.sync.aligned.b32(generic.ptr_to([0]), T.uint32(32))
+ T.ptx.tcgen05.alloc.cta_group__1.sync.aligned.shared__cta.b32(
+ generic.ptr_to([0]), T.uint32(32)
+ )
+ T.ptx.tcgen05.commit.cta_group__1.mbarrier__arrive__one.b64(generic.ptr_to([0]))
+ T.ptx.tcgen05.commit.cta_group__1.mbarrier__arrive__one.shared__cluster.b64(
+ generic.ptr_to([0])
+ )
+ T.ptx.tcgen05.commit.cta_group__2.mbarrier__arrive__one.multicast__cluster.b64(
+ generic.ptr_to([0]), T.uint16(3)
+ )
+ T.ptx[
+ "tcgen05.commit.cta_group::2.mbarrier::arrive::one"
+ ".shared::cluster.multicast::cluster.b64"
+ ](generic.ptr_to([0]), T.uint16(3))
+ T.ptx.mapa.u64(mapped64, generic.ptr_to([0]), T.uint32(1))
+ T.ptx.cvta.to.shared__cluster.u64(window64, generic.ptr_to([0]))
+ T.ptx.mapa.u64(mapped64, window64, T.uint32(1))
+ T.ptx.mapa.shared__cluster.u64(mapped64, window64, T.uint32(1))
+ T.ptx.mapa.shared__cluster.u32(mapped32, T.uint32(0), T.uint32(1))
+
+ reparsed = tvm.script.from_source(kernel.script())
+ tvm.ir.assert_structural_equal(kernel, reparsed)
+
+
def test_ptx_helper_source_golden():
"""Exact generated helper source, one per family (executable documentation)."""
from tvm.backend.cuda.ptx.render import render_variant
@@ -1722,7 +2568,7 @@ def test_ptx_helper_source_golden():
)
# Two written type tokens, and an accumulator derived from both: c and d
# are .u32 only when atype and btype both are, so this mixed pair makes
- # them .s32 (ISA 9.7.1.24).
+ # them .s32 (ISA 9.7.1.23).
assert render("dp4a", atype="u32", btype="s32") == (
"__forceinline__ __device__ void tvm_builtin_ptx_dp4a_u32_s32"
"(int32_t& __d, uint32_t __a, int32_t __b, int32_t __c) {\n"
@@ -1867,6 +2713,21 @@ def test_ptx_helper_source_golden():
'"r"(__b), "f"(__c));\n'
"}\n"
)
+ # The first type controls the bit width, but does not force d/a/b to share
+ # an interpretation or a carrier class. A noncanonical helper names all
+ # four operand dtypes so independently mixed products cannot collide.
+ assert render(
+ "slct",
+ dtypes=("uint32", "float32", "int32", "int32"),
+ dtype="f32",
+ ctype="s32",
+ ) == (
+ "__forceinline__ __device__ void tvm_builtin_ptx_slct_f32_s32_u32_f32_s32_s32"
+ "(uint32_t& __d, float __a, int32_t __b, int32_t __c) {\n"
+ ' asm volatile("slct.f32.s32 %0, %1, %2, %3;" : "=r"(__d) : "f"(__a), '
+ '"r"(__b), "r"(__c));\n'
+ "}\n"
+ )
# Half-precision comparison (ISA 9.7.7). The packed setp reuses the pipe
# pair, but its two halves mean something else than 9.7.6's: p and q are
@@ -1921,6 +2782,16 @@ def test_ptx_helper_source_golden():
': "=r"(__d), "=r"(__p) : "r"(__a), "r"(__b), "r"(__c), "r"(__q));\n'
"}\n"
)
+ # Discarding d is a fixed syntax sibling, so `_` is table-owned text and
+ # neither a C parameter nor a use of the caller-selectable sink mechanism.
+ assert render("lop3_bool_sink", boolop="and", type="b32", imms=("128",)) == (
+ "__forceinline__ __device__ void tvm_builtin_ptx_lop3_bool_sink_and_b32_128"
+ "(uint32_t& __p, uint32_t __a, uint32_t __b, uint32_t __c, uint32_t __q) {\n"
+ ' asm volatile("{ .reg .pred pd0; .reg .pred ps0; setp.ne.b32 ps0, %4, 0; '
+ 'lop3.and.b32 _|pd0, %1, %2, %3, 128, ps0; selp.b32 %0, 1, 0, pd0; }" '
+ ': "=r"(__p) : "r"(__a), "r"(__b), "r"(__c), "r"(__q));\n'
+ "}\n"
+ )
# The shift amount is a 32-bit value "regardless of the instruction type",
# so a 16-bit shl still takes a uint32 there.
assert render("shl", type="b16") == (
@@ -1981,7 +2852,7 @@ def test_ptx_helper_source_golden():
"}\n"
)
- # Parallel synchronization (ISA 9.7.14). elect.sync's `d|p` is the one
+ # Parallel synchronization (ISA 9.7.13). elect.sync's `d|p` is the one
# pipe pair the ISA makes mandatory -- ptxas rejects a bare destination --
# and its halves are two different register classes.
assert render("elect_sync") == (
@@ -2001,6 +2872,61 @@ def test_ptx_helper_source_golden():
'"l"(__addr), "q"(__compare), "q"(__value) : "memory");\n'
"}\n"
)
+ # The bit-bucket overload is a fixed `_` operand, selected by omitting the
+ # returned-value destination from the call. It has no output constraint.
+ assert render("atom_bitbucket", space="global", op="add", type="u32") == (
+ "__forceinline__ __device__ void tvm_builtin_ptx_atom_bitbucket_global_add_u32"
+ "(const void* __addr, uint32_t __value) {\n"
+ ' asm volatile("atom.global.add.u32 _, [%0], %1;" : : "l"(__addr), '
+ '"r"(__value) : "memory");\n'
+ "}\n"
+ )
+ # Vector atom also uses one whole bit bucket, not a brace group of lane
+ # sinks. The value remains a normal register group.
+ assert render("atom_vec_f32_bitbucket", space="global", op="add", vec="v2", type="f32") == (
+ "__forceinline__ __device__ void "
+ "tvm_builtin_ptx_atom_vec_f32_bitbucket_global_add_v2_f32"
+ "(const void* __addr, float __value0, float __value1) {\n"
+ ' asm volatile("atom.global.add.v2.f32 _, [%0], {%1, %2};" : : '
+ '"l"(__addr), "f"(__value0), "f"(__value1) : "memory");\n'
+ "}\n"
+ )
+ # PTX ISA 9.2 documents the same bit-bucket spelling for bf16 atom, but
+ # CUDA 13.2 ptxas does not compile it (an exact inlined caller segfaults).
+ # Returned-value bf16 atom and f16 bit buckets remain registered.
+ for entry_name, by_name in (
+ ("atom_half_bitbucket", {"op": "add", "noftz": "noftz", "type": "bf16"}),
+ (
+ "atom_vec_half_bitbucket",
+ {"op": "add", "noftz": "noftz", "vec": "v2", "type": "bf16"},
+ ),
+ ):
+ with pytest.raises(ValueError, match="PTX ISA 9.2 supports.*ptxas does not compile"):
+ tokens_for(TABLE[entry_name], **by_name)
+ # red's half-word and packed vector syntax lines share this entry; the
+ # conditional policy operand follows the value group.
+ assert render(
+ "red_vec_half",
+ sem="release",
+ scope="sys",
+ space="global",
+ op="max",
+ noftz="noftz",
+ cache="L2::cache_hint",
+ vec="v4",
+ type="f16x2",
+ ) == (
+ "__forceinline__ __device__ void "
+ "tvm_builtin_ptx_red_vec_half_release_sys_global_max_noftz_"
+ "L2__cache_hint_v4_f16x2"
+ "(const void* __addr, uint32_t __value0, uint32_t __value1, "
+ "uint32_t __value2, uint32_t __value3, uint64_t __cache_policy) {\n"
+ ' asm volatile("red.release.sys.global.max.noftz.L2::cache_hint.v4.f16x2 '
+ '[%0], {%1, %2, %3, %4}, %5;" : : "l"(__addr), "r"(__value0), '
+ '"r"(__value1), "r"(__value2), "r"(__value3), "l"(__cache_policy) '
+ ': "memory");\n'
+ "}\n"
+ )
# bar.red reduces a predicate across a barrier, so a predicate crosses the
# boundary in BOTH directions around one instruction: setp on the way in,
# selp on the way out.
@@ -2111,6 +3037,129 @@ def test_ptx_coercion_ir_forms():
# fp16/bf16 dtypes bring in __half / __nv_bfloat16 and their bit-cast helpers.
_CERT_PRELUDE = "#include <cstdint>\n#include <cuda_fp16.h>\n#include <cuda_bf16.h>"
+# A certification unit must be small enough that ptxas is checking helpers,
+# rather than being stress-tested by one translation unit containing every
+# sampled table variant. The limit is on generated source (not helper count),
+# because helper sizes vary substantially across instruction families.
+_CERT_MAX_SOURCE_CHARS = 256 * 1024
+_CERT_HELPER_SIGNATURE_RE = re.compile(
+ r"\A__forceinline__ __device__ void (?P<helper>[A-Za-z_]\w*)\((?P<params>[^)]*)\) \{"
+)
+_CERT_PARAM_RE = re.compile(r"(?P<type>.+?)(?P<reference>&)?\s+(?P<name>__[A-Za-z_]\w*)\Z")
+
+
+def _certification_kernel(helper, helper_source, kernel_index):
+ """Keep a helper inline and call it from a retained, production-shaped kernel.
+
+ Merely placing an unreferenced ``__forceinline__`` helper in a translation
+ unit does not certify it: nvcc may discard it before ptxas sees its body.
+ The old workaround stripped ``__forceinline__`` and retained the resulting
+ device function with RDC. That changes the device ABI and is not an
+ equivalent compilation shape for wide inline-asm operands.
+
+ A global caller is retained without RDC and forces the original helper to
+ inline, which is the shape emitted code uses. Pointer arguments come from
+ the runtime sink so address instructions are not specialized to null.
+ Register inputs are locals, avoiding an artificial device ABI boundary;
+ ordinary values and raw 32-bit addresses are loaded through the sink so
+ they keep the runtime data flow of generated calls. Predicate and b128
+ inputs remain local because crossing either through the certification ABI
+ creates a tool-only compilation shape. Written values are loaded and
+ stored through the sink so non-volatile output instructions cannot be
+ optimized away before ptxas.
+ """
+ match = _CERT_HELPER_SIGNATURE_RE.match(helper_source)
+ assert match is not None, f"cannot parse certification signature for {helper}"
+ assert match.group("helper") == helper
+
+ declarations = []
+ arguments = []
+ stores = []
+ params = match.group("params")
+ parsed_params = params.split(", ") if params else ()
+ for param_index, param in enumerate(parsed_params):
+ param_match = _CERT_PARAM_RE.fullmatch(param)
+ assert param_match is not None, f"cannot parse {helper} parameter: {param}"
+ c_type = param_match.group("type")
+ param_name = param_match.group("name")
+ local = f"__cert_arg{param_index}"
+ input_offset = 16 * param_index
+ if c_type.endswith("*"):
+ declarations.append(f" {c_type} {local} = __cert_sink + {input_offset};")
+ elif param_match.group("reference"):
+ declarations.append(
+ f" {c_type} {local} = "
+ f"*reinterpret_cast<const {c_type}*>(__cert_sink + {input_offset});"
+ )
+ elif c_type in ("__int128_t", "__uint128_t") or param_name == "__pred":
+ declarations.append(f" {c_type} {local}{{}};")
+ else:
+ declarations.append(
+ f" {c_type} {local} = "
+ f"*reinterpret_cast<const {c_type}*>(__cert_sink + {input_offset});"
+ )
+ arguments.append(local)
+ if param_match.group("reference"):
+ # Every C binding is at most 16 bytes. Distinct, aligned locations
+ # keep all outputs observable without imposing a constraint letter
+ # of our own on the value being certified.
+ stores.append(
+ f" *reinterpret_cast<{c_type}*>(__cert_sink + "
+ f"{16 * (len(parsed_params) + len(stores))}) = {local};"
+ )
+
+ kernel = [
+ f'extern "C" __global__ void __ptx_cert_{kernel_index}(char* __cert_sink) {{',
+ *declarations,
+ f" {helper}({', '.join(arguments)});",
+ *stores,
+ "}",
+ ]
+ return "\n".join((helper_source, *kernel))
+
+
+def _append_certification(by_arch, arch, helper, helper_source):
+ items = by_arch.setdefault(arch, [])
+ items.append((helper, _certification_kernel(helper, helper_source, len(items))))
+
+
+def _certification_batches(items):
+ """Yield deterministic, source-size-bounded groups of certification kernels."""
+ batch = []
+ size = len(_CERT_PRELUDE) + 1
+ for item in items:
+ item_size = len(item[1]) + 1
+ assert item_size + len(_CERT_PRELUDE) <= _CERT_MAX_SOURCE_CHARS, (
+ f"single certification helper exceeds source limit: {item[0]}"
+ )
+ if batch and size + item_size > _CERT_MAX_SOURCE_CHARS:
+ yield batch
+ batch = []
+ size = len(_CERT_PRELUDE) + 1
+ batch.append(item)
+ size += item_size
+ if batch:
+ yield batch
+
+
+def _assert_certifications_ok(by_arch):
+ for arch, items in by_arch.items():
+ for batch_index, batch in enumerate(_certification_batches(items)):
+ names, sources = zip(*batch, strict=True)
+ try:
+ _assert_ptxas_ok("\n".join((_CERT_PRELUDE, *sources)), arch=arch)
+ except Exception as err:
+ message = str(err)
+ marker = "Compilation error:"
+ diagnostic = (
+ message[message.rfind(marker) :] if marker in message else message[-4000:]
+ )
+ raise AssertionError(
+ f"PTX certification failed for {arch} batch {batch_index} "
+ f"({len(names)} helpers, {names[0]} .. {names[-1]}):\n{diagnostic}"
+ ) from None
+
+
_ASM_RE = re.compile(r'asm(?: volatile)?\("(.*?)"\s*:', re.S)
_BLOCK_RE = re.compile(r"^\{ (?P<body>.*) \}$")
# The asm block's sanctioned non-instructions: `render.BRIDGE`'s register
@@ -2131,19 +3180,18 @@ _BOUNDARY_RE = re.compile(
r"|".join(
(
r"\.reg \.pred (?:p|ps\d+|pd\d+);", # @p guard / pred bridge declarations
- r"\.reg \.b8 raw_\w+;", # b8 bridge declaration
+ r"\.reg \.b8 raw_\w+;", # st.async byte-register bridge declaration
r"setp\.ne\.b32 (?:p|ps\d+), %\d+, 0;", # @p guard, pred_src conversion in
r"selp\.b32 %\d+, 1, 0, pd\d+;", # pred_dst materialization out
- r"cvt\.u8\.u16 raw_\w+, %\d+;", # b8 conversion in
- r"cvt\.u16\.u8 %\d+, raw_\w+;", # b8 conversion out
+ r"cvt\.u8\.u16 raw_\w+, %\d+;", # st.async byte conversion in
+ r"cvt\.u16\.u8 %\d+, raw_\w+;", # st.async byte conversion out
)
)
)
def _as_render_args(rendering):
- """`renderings` yields (tokens, dtypes, predicated, imms); render_variant
- takes (tokens, predicated, dtypes, imms)."""
+ """Reorder a five-field ``renderings`` item for ``render_variant``."""
tokens, dtypes, predicated, imms, sinks = rendering
return tokens, predicated, dtypes, imms, sinks
@@ -2208,13 +3256,11 @@ def test_ptx_single_instruction_invariant():
rather than adding one. cvta coercion is a separate IR node and must never
appear inside a helper body.
- ``RAW_ENTRIES`` below is the one exemption, and it is a list of names, not
- a predicate: the entries whose helper body the table cannot derive because
- one operand is typed ``.b8`` and inline asm has no 8-bit constraint, so the
- value has to be staged through a block-local ``.reg .b8``. Naming them here
- and asserting the set equals the table's own ``raw_render`` entries means a
- new one can never be added without editing this test. Every other assertion
- still applies to them.
+ Framework boundary conversions such as predicate materialization and the
+ byte-register bridges are peeled before counting. ``RAW_ENTRIES`` is the
+ separate closed exemption list for a genuinely hand-written helper body;
+ asserting that it equals the table's ``raw_render`` entries means a new one
+ cannot be added without editing this test.
"""
from tvm.backend.cuda.ptx.render import render_variant
from tvm.backend.cuda.ptx.table import TABLE, renderings
@@ -2238,9 +3284,9 @@ def test_ptx_single_instruction_invariant():
asm_blocks = asm_re.findall(source)
assert len(asm_blocks) == 1, f"{opcode}: {len(asm_blocks)} asm blocks, expected 1"
if raw:
- # The `.reg .b8` prologue is several statements, which is why
- # this entry is exempt. The instruction must still be in there
- # as its own statement. The helper name needs no assertion:
+ # A raw helper may contain several statements, which is why it
+ # is exempt. The instruction must still be in there as its own
+ # statement. The helper name needs no assertion:
# `render_variant` passes the name it derived into
# `raw_render`, so a raw body cannot declare a different one.
assert f"; {opcode} " in asm_blocks[0], (
@@ -2329,17 +3375,16 @@ def test_ptx_all_variants_render_unique():
_, helper, _ = render_variant(entry, *args, addr_offsets=addr_offsets)
assert helper not in names, f"address-offset helper name collision: {helper}"
names.add(helper)
- assert total == 200027 # update when the table grows
+ assert total == 660867 # update when the table grows or a ptxas gap narrows it
def test_ptx_no_instruction_registered_twice():
- """One PTX instruction, one entry.
+ """No two entries register the same PTX instruction and C-call shape.
- The table's law is the ISA: an entry models one syntax group, and no two
- entries may model the same one. Strip the helper name and the parameter
- names off a rendering and what is left is the instruction itself plus its
- operand constraints — if two entries ever produce the same one, the ISA
- line has been registered twice and calls to it are unresolvable.
+ Sibling entries may intentionally share an opcode when their public call
+ shapes are disjoint (bare mapa.u64 accepts a pointer or a raw uint64).
+ Strip only helper and parameter *names*; the remaining C types, asm text,
+ and constraints form the identity whose duplication would be ambiguous.
"""
from tvm.backend.cuda.ptx.render import render_variant
from tvm.backend.cuda.ptx.table import TABLE, renderings
@@ -2467,14 +3512,34 @@ def test_ptx_stub_up_to_date():
@requires_nvcc
def test_ptxas_gate_rejects_invalid():
"""Honesty check: the gate path must actually reject bad instructions."""
- bogus = '__device__ void f(unsigned x) { asm volatile("totally.bogus.instr %0;" : : "r"(x)); }'
+ bogus = (
+ "__forceinline__ __device__ void f(uint32_t __x) {\n"
+ ' asm volatile("totally.bogus.instr %0;" : : "r"(__x));\n'
+ "}\n"
+ )
+ source = "\n".join((_CERT_PRELUDE, _certification_kernel("f", bogus, 0)))
with pytest.raises(Exception, match="bogus|error"):
- _assert_ptxas_ok(bogus, rdc=True)
+ _assert_ptxas_ok(source)
+
+
+def test_ptx_vec256_wide_carriers_not_registered():
+ """PTX documents wider carriers, but CUDA 13.2 ptxas cannot compile this axis."""
+ from tvm.backend.cuda.ptx.table import TABLE, dtype_combos, tokens_for
+
+ expected = (
+ ("uint32", "uint64"),
+ ("int32", "uint64"),
+ ("float32", "uint64"),
+ )
+ for name in ("ld_vec256", "st_vec256"):
+ entry = TABLE[name]
+ tokens = tokens_for(entry, vec="v8", type="b32")
+ assert dtype_combos(entry, tokens) == expected
@requires_nvcc
def test_ptx_sampled_helpers_assemble():
- """Fast tier: a seeded sample of every family's variants assembles."""
+ """Fast tier: production-shaped calls to a seeded sample assemble."""
import random
from tvm.backend.cuda.ptx.render import render_variant
@@ -2484,15 +3549,18 @@ def test_ptx_sampled_helpers_assemble():
by_arch = {}
for entry in TABLE.values():
arch = entry.cert_arch or PTX_ARCH
- rendered = list(renderings(entry))
+ # Some dtype domains originate in sets/frozensets. Sort the complete
+ # rendering tuple so the seeded certification sample is reproducible
+ # across Python hash seeds and can be isolated from a reported helper
+ # name alone.
+ rendered = sorted(renderings(entry), key=repr)
for i in rng.sample(range(len(rendered)), min(48, len(rendered))):
- _, _, source = render_variant(entry, *_as_render_args(rendered[i]))
- by_arch.setdefault(arch, []).append(source.replace("__forceinline__ ", ""))
+ _, helper, source = render_variant(entry, *_as_render_args(rendered[i]))
+ _append_certification(by_arch, arch, helper, source)
for args, addr_offsets in _addr_offset_samples(entry):
- _, _, source = render_variant(entry, *args, addr_offsets=addr_offsets)
- by_arch.setdefault(arch, []).append(source.replace("__forceinline__ ", ""))
- for arch, sources in by_arch.items():
- _assert_ptxas_ok("\n".join([_CERT_PRELUDE, *sources]), rdc=True, arch=arch)
+ _, helper, source = render_variant(entry, *args, addr_offsets=addr_offsets)
+ _append_certification(by_arch, arch, helper, source)
+ _assert_certifications_ok(by_arch)
_CERT_SHARDS = 32
@@ -2521,9 +3589,11 @@ def test_ptx_all_helpers_certify(shard):
Stride slicing keeps shards balanced (ld dominates the variant count).
Variants are grouped by their family's arch floor and each group is
assembled at that arch: below an instruction's floor ptxas rejects legal
- variants, and believing it would delete real coverage.
- __forceinline__ must be stripped and -rdc used: unreferenced inline
- device functions are silently dropped before ptxas ever sees them.
+ variants, and believing it would delete real coverage. Every helper keeps
+ ``__forceinline__`` and is called by a retained certification kernel, so
+ ptxas sees the same inline-asm shape as a real generated caller. Each
+ pytest shard is further split by source size to avoid oversized compiler
+ translation units.
"""
from tvm.backend.cuda.ptx.render import render_variant
from tvm.backend.cuda.ptx.table import TABLE, renderings
@@ -2544,11 +3614,10 @@ def test_ptx_all_helpers_certify(shard):
if index % _CERT_SHARDS == shard:
covered += 1
arch = entry.cert_arch or PTX_ARCH
- _, _, src = render_variant(entry, *args, addr_offsets=addr_offsets)
- by_arch.setdefault(arch, []).append(src.replace("__forceinline__ ", ""))
+ _, helper, src = render_variant(entry, *args, addr_offsets=addr_offsets)
+ _append_certification(by_arch, arch, helper, src)
assert covered, "empty shard: lower _CERT_SHARDS"
- for arch, sources in by_arch.items():
- _assert_ptxas_ok("\n".join([_CERT_PRELUDE, *sources]), rdc=True, arch=arch)
+ _assert_certifications_ok(by_arch)
@requires_nvcc