This is an automated email from the ASF dual-hosted git repository.
tqchen 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 ed5e2fed3f [FIX][TIRx][CUDA] Support SM100 weight-stationary B
collectors (#20329)
ed5e2fed3f is described below
commit ed5e2fed3f32030d9407e72751550b26653c5fac
Author: Hongyi Jin <[email protected]>
AuthorDate: Sun Sep 13 09:44:50 2026 -0400
[FIX][TIRx][CUDA] Support SM100 weight-stationary B collectors (#20329)
The SM100 `tcgen05.mma.ws` SS/TS entries omit the collector-B modifier
slot, preventing callers from expressing weight-stationary B reuse
through the PTX namespace. Add the existing optional B0–B3
fill/use/lastuse/discard domain to both entries, preserving calls that
omit the qualifier.
Also reject modifier-domain and operand-count mismatches during code
generation, before stale table/codegen layouts can emit malformed CUDA
helpers. Add dispatch and parser round-trip regressions, SM100
certification for all WS variants, and stale-layout diagnostics. Make
the test source helper accept an explicit target so the existing SM107
collector test compiles at its required architecture.
Validation with CUDA 13.4 NVCC/NVRTC and a matching worktree FFI:
- PTX dialect, conversion, and address suites: 179 non-skipped cases
passed across the full run and targeted reruns.
- Full-table NVCC/ptxas certification: **32/32 shards passed**,
including predicated helpers and address-offset samples at their
declared architectures.
- Regenerated IDE stubs: no diff needed; the tokens already exist in the
ti16 families.
- Changed-file pre-commit hooks and `git diff --check` passed.
Closes #20328.
---
python/tvm/backend/cuda/ptx/engine.py | 12 ++++
python/tvm/backend/cuda/ptx/table.py | 2 +
tests/python/tirx/codegen/test_ptx_dialect.py | 86 +++++++++++++++++++++++++--
3 files changed, 95 insertions(+), 5 deletions(-)
diff --git a/python/tvm/backend/cuda/ptx/engine.py
b/python/tvm/backend/cuda/ptx/engine.py
index ae1b2b9ad7..0b7de8a736 100644
--- a/python/tvm/backend/cuda/ptx/engine.py
+++ b/python/tvm/backend/cuda/ptx/engine.py
@@ -186,6 +186,12 @@ def _make_codegen(entry: InstructionEntry):
preserve_dst = "keep" in flags
tokens = [parse_str(a) for a in args[len(args) - n_slots - 1 : -1]]
rest = list(args[: len(args) - n_slots - 1]) # operands, plus pred
when present
+ for slot, token in zip(entry.slots, tokens, strict=True):
+ if token not in slot.choices and not (slot.optional and token ==
""):
+ raise ValueError(
+ f"{entry.name}: invalid codegen modifier {token!r} for
slot {slot.name!r}; "
+ "the PTX call and registered table entry may be
inconsistent"
+ )
mod_map = mods(entry, tokens)
layout = operand_layout(entry, mod_map)
n_operands = sum(n for _, _, n in layout)
@@ -199,6 +205,12 @@ def _make_codegen(entry: InstructionEntry):
at[i] = pos
pos += 1
n_present = pos
+ expected_args = n_present + int(predicated)
+ if len(rest) != expected_args:
+ raise ValueError(
+ f"{entry.name}: expected {expected_args} codegen operand(s),
got {len(rest)}; "
+ "the PTX call and registered table entry may be inconsistent"
+ )
sinks = frozenset(
(slot.name, lane)
for slot, i, lanes in layout
diff --git a/python/tvm/backend/cuda/ptx/table.py
b/python/tvm/backend/cuda/ptx/table.py
index 94e0121f82..ac283dca18 100644
--- a/python/tvm/backend/cuda/ptx/table.py
+++ b/python/tvm/backend/cuda/ptx/table.py
@@ -11529,6 +11529,7 @@ _ENTRIES = [
)
for form in ("ss", "ts")
],
+ # PTX ISA 9.7.18.10.10.3: optional
collector::b{0,1,2,3}::{fill,use,lastuse,discard}.
*[
InstructionEntry( # weight-stationary: no mask vector, a zero-column
desc
name=f"tcgen05_mma_ws_{form}",
@@ -11538,6 +11539,7 @@ _ENTRIES = [
ModifierSlot("ws", ("ws",)),
ModifierSlot("cta_group", ("cta_group::1",)),
ModifierSlot("kind", ("kind::f16", "kind::tf32",
"kind::f8f6f4", "kind::i8")),
+ ModifierSlot("collector_b", _TCGEN05_WS_COLLECTOR_B,
optional=True),
),
cert_arch="sm_100a",
operands=(
diff --git a/tests/python/tirx/codegen/test_ptx_dialect.py
b/tests/python/tirx/codegen/test_ptx_dialect.py
index 0e4737292f..9db4666189 100644
--- a/tests/python/tirx/codegen/test_ptx_dialect.py
+++ b/tests/python/tirx/codegen/test_ptx_dialect.py
@@ -34,9 +34,9 @@ TARGET = tvm.target.Target("cuda")
requires_nvcc = pytest.mark.skipif(shutil.which("nvcc") is None, reason="nvcc
not available")
-def _cuda_source(func) -> str:
- with TARGET:
- mod = tvm.compile(tvm.IRModule({"main": func}), target=TARGET,
tir_pipeline="tirx")
+def _cuda_source(func, target=TARGET) -> str:
+ with target:
+ mod = tvm.compile(tvm.IRModule({"main": func}), target=target,
tir_pipeline="tirx")
return mod.mod.imports[0].inspect_source("cuda")
@@ -2340,6 +2340,80 @@ def
test_ptx_wgmma_integer_shape_domains_follow_concrete_syntax():
assert "m64n256k256" in b1_shapes
[email protected]("form", ("ss", "ts"))
[email protected]("collector", ("", "b0::fill", "b1::use",
"b2::lastuse", "b3::discard"))
+def test_ptx_tcgen05_mma_ws_collector_dispatch(form, collector):
+ opcode = "tcgen05.mma.ws.cta_group::1.kind::f16"
+ if collector:
+ opcode += f".collector::{collector}"
+ a_dtype = "uint64" if form == "ss" else "uint32"
+
+ @T.prim_func
+ def kernel():
+ T.device_entry()
+ T.cta_id([1])
+ T.thread_id([32])
+ T.ptx[opcode](
+ T.uint32(0),
+ T.cast(0, a_dtype),
+ T.uint64(0),
+ T.uint32(0),
+ T.ptx.pred(T.uint32(0)),
+ T.uint64(0),
+ )
+
+ src = _cuda_source(kernel, tvm.target.Target({"kind": "cuda", "arch":
"sm_100a"}))
+ a_operand = "%1" if form == "ss" else "[%1]"
+ assert f"{opcode} [%0], {a_operand}, %2, %3, ps0, %5;" in src
+ tvm.ir.assert_structural_equal(kernel,
tvm.script.from_source(kernel.script()))
+
+
+@requires_nvcc
+def test_ptx_tcgen05_mma_ws_collectors_certify_sm100a():
+ """All WS kinds, B buffers/operations, and predication assemble on
SM100a."""
+ from tvm.backend.cuda.ptx.render import render_variant
+ from tvm.backend.cuda.ptx.table import TABLE, renderings
+
+ by_arch = {}
+ for form in ("ss", "ts"):
+ entry = TABLE[f"tcgen05_mma_ws_{form}"]
+ for rendering in renderings(entry):
+ _, helper, source = render_variant(entry,
*_as_render_args(rendering))
+ _append_certification(by_arch, "sm_100a", helper, source)
+ _assert_certifications_ok(by_arch)
+
+
[email protected]("mismatch", ("modifier", "operand"))
+def test_ptx_codegen_rejects_stale_table_layout(mismatch):
+ """A traced call must not silently feed modifier strings to an old
helper."""
+ from dataclasses import replace
+
+ from tvm.backend.cuda.ptx.engine import PTXNamespace, _make_codegen
+ from tvm.backend.cuda.ptx.table import TABLE, ModifierSlot, OperandSlot
+
+ entry = TABLE["tcgen05_mma_ws_ss"]
+ if mismatch == "modifier":
+ changed = replace(entry, slots=(*entry.slots, ModifierSlot("extra",
("extra",))))
+ opcode = "tcgen05.mma.ws.cta_group::1.kind::f16.extra"
+ extra = ()
+ else:
+ changed = replace(entry, operands=(*entry.operands,
OperandSlot("extra", dtype="u64")))
+ opcode = "tcgen05.mma.ws.cta_group::1.kind::f16"
+ extra = (T.uint64(0),)
+ namespace = PTXNamespace({changed.name: changed})
+ call = namespace[opcode](
+ T.uint32(0),
+ T.uint64(0),
+ T.uint64(0),
+ T.uint32(0),
+ namespace.pred(T.uint32(0)),
+ T.uint64(0),
+ *extra,
+ )
+ with pytest.raises(ValueError, match="PTX call and registered table entry
may be inconsistent"):
+ _make_codegen(entry)(*call.args)
+
+
@requires_nvcc
def test_ptx_tcgen05_mma_block_size_form():
@T.prim_func
@@ -2405,7 +2479,9 @@ def test_ptx_tcgen05_mma_block_size_collector_form():
](tmem, tmem, desc, idesc, tmem, tmem, T.ptx.pred(flag))
A[tx] = A[tx]
- collector_src = _cuda_source(sm107_collector_kernel)
+ collector_src = _cuda_source(
+ sm107_collector_kernel, tvm.target.Target({"kind": "cuda", "arch":
"sm_107f"})
+ )
ss_collector_opcode = (
"tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16"
".collector::a::discard.collector::b::fill"
@@ -4144,7 +4220,7 @@ 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 == 762050 # update when the table grows or a ptxas gap
narrows it
+ assert total == 762178 # update when the table grows or a ptxas gap
narrows it
def test_ptx_no_instruction_registered_twice():
