[CUDA][TIRx] Allow collector-A-only SM107 block-scale MMA and open cp.async wait_group immediates

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

@@ -5328,7 +5328,7 @@ def _check_set_packed(m):
                 ),
                 ModifierSlot("block_scale", ("block_scale",)),
                 ModifierSlot("collector_a", _TCGEN05_COLLECTOR_A),
-                ModifierSlot("collector_b", _TCGEN05_COLLECTOR_B),
+                ModifierSlot("collector_b", _TCGEN05_COLLECTOR_B, optional=True),
             ),
             cert_arch="sm_107f",
             operands=(
@@ -8516,6 +8516,12 @@ def _check_set_packed(m):
             OperandSlot("b", dtype="u32"),
         ),
     ),
+    # PTX ISA 9.7.10.28.3.3 / 9.7.10.28.6.2 defines N only as an integer
+    # constant and declares no value domain. Keep it OPEN: callers may use any
+    # compile-time integer, while certification samples the instruction shape.
+    # MEASURED on CUDA 13.4 ptxas at sm_107a: ordinary and bulk forms accept
+    # values beyond 7 (8, 9, 16, 255), and the bulk `.read` form also accepts
+    # 2147483647 and -1.
     *[
         InstructionEntry(  # cp.async{.bulk}.wait_group{.read} N;
             name=f"cp_async{'_bulk' if bulk else ''}_wait_group",
@@ -8527,7 +8533,7 @@ def _check_set_packed(m):
                 *((ModifierSlot("read", ("read",), optional=True),) if bulk else ()),
             ),
             orders_memory=True,
-            operands=(OperandSlot("group", kind="imm", choices=tuple(str(n) for n in range(8))),),
+            operands=(OperandSlot("group", kind="imm"),),
         )
         for bulk in (False, True)
     ],
@@ -10179,10 +10185,10 @@ def _check_set_packed(m):
     # and cp.async.wait_all per 9.7.10.28.3.3, cp.async.bulk.commit_group /
     # .wait_group per 9.7.10.28.6.1 / 9.7.10.28.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`
-    # set is what makes every one of them certifiable. 0..7 covers every call
-    # site (pipeline depths); widen the tuple if a deeper pipeline appears.
+    # The wait_group counts are caller-chosen OPEN immediates: the ISA gives N
+    # no register form or value domain. Each call-site constant becomes its own
+    # helper; enumeration and full-table certification sample the open operand
+    # at 0 and therefore certify the instruction shape rather than every value.
     #
     # (The `cp.async` ca/cg copy lines this note once excluded are registered
     # in the 9.7.10 group above, ignore-src operand and all.)

tests/python/tirx/codegen/test_ptx_dialect.py

@@ -1977,6 +1977,10 @@ def kernel(a_ptr: T.handle):
         T.ptx.applypriority.global_.L2__evict_normal(A.ptr_to([6]))
         T.ptx.discard.global_.L2(A.ptr_to([7]))
         T.ptx.prefetchu.L1(A.ptr_to([0]))
+        T.ptx.cp.async_.wait_group(255)
+        T.ptx.cp.async_.bulk.wait_group(255)
+        T.ptx.cp.async_.bulk.wait_group.read(8)
+        T.ptx.cp.async_.bulk.wait_group.read(-1)
         T.ptx.multimem_ld_reduce.add.u32(v, A.ptr_to([0]))
         T.ptx.multimem_red.relaxed.gpu.add.u32(A.ptr_to([0]), v)
         smem[tx % 4] = d + p + v
@@ -2000,6 +2004,10 @@ def kernel(a_ptr: T.handle):
         "applypriority.global.L2::evict_normal [%0], 128;",
         "discard.global.L2 [%0], 128;",
         "prefetchu.L1 [%0];",
+        "cp.async.wait_group 255;",
+        "cp.async.bulk.wait_group 255;",
+        "cp.async.bulk.wait_group.read 8;",
+        "cp.async.bulk.wait_group.read -1;",
         "multimem.ld_reduce.add.u32 %0, [%1];",
         "multimem.red.relaxed.gpu.add.u32 [%0], %1;",
     ):
@@ -2402,6 +2410,32 @@ def invalid_mxf4_block16():
             )
 
 
+@requires_nvcc
+def test_ptx_tcgen05_mma_block_scale_collector_a_without_block_size():
+    """SM107 activation-stationary FP8 accepts collector A without `.block*`."""
+
+    @T.prim_func
+    def kernel(a_ptr: T.handle):
+        A = T.match_buffer(a_ptr, (32,), "uint32")
+        T.device_entry()
+        T.cta_id([1])
+        tx = T.thread_id([32])
+        if tx == 0:
+            tmem = T.local_scalar("uint32")
+            desc = T.local_scalar("uint64")
+            idesc = T.local_scalar("uint32")
+            flag = T.local_scalar("uint32")
+            T.ptx["tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.collector::a::discard"](
+                tmem, desc, desc, idesc, tmem, tmem, T.ptx.pred(flag)
+            )
+        A[tx] = A[tx]
+
+    src = _cuda_source(kernel)
+    opcode = "tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.collector::a::discard"
+    assert opcode in src
+    _assert_ptxas_ok(src, arch="sm_107a")
+
+
 def test_ptx_tcgen05_mma_block_size_collector_legality():
     from tvm.backend.cuda.ptx.table import TABLE, tokens_for
 
@@ -4097,7 +4131,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 == 762023  # update when the table grows or a ptxas gap narrows it
+    assert total == 762050  # update when the table grows or a ptxas gap narrows it
 
 
 def test_ptx_no_instruction_registered_twice():