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

Author: jinhongyii

python/tvm/backend/cuda/ptx/engine.py

@@ -186,6 +186,12 @@ def codegen(*args):
         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 codegen(*args):
                 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

python/tvm/backend/cuda/ptx/table.py

@@ -11529,6 +11529,7 @@ def _check_set_packed(m):
         )
         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 @@ def _check_set_packed(m):
                 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=(

tests/python/tirx/codegen/test_ptx_dialect.py

@@ -34,9 +34,9 @@
 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
 
 
+@pytest.mark.parametrize("form", ("ss", "ts"))
+@pytest.mark.parametrize("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)
+
+
+@pytest.mark.parametrize("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 sm107_collector_kernel(a_ptr: T.handle):
             ](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():