From 8a54bf30b07e1f0903b569963a9ab112ccb1efd4 Mon Sep 17 00:00:00 2001 From: Bohan Hou Date: Fri, 4 Sep 2026 19:56:52 -0400 Subject: [PATCH] [CUDA][TIRx] Allow collector-A-only SM107 block-scale MMA and open cp.async wait_group immediates - tcgen05.mma block-scale SM107 entries: collector B is optional, so the collector-A-only form is accepted alongside the A+B form. - cp.async.wait_group / cp.async.bulk.wait_group N is an OPEN immediate. PTX ISA 9.7.10.28.3.3 / 9.7.10.28.6.2 define N only as an integer constant with no value domain, and CUDA 13.4 ptxas at sm_107a accepts values beyond 7 (8, 9, 16, 255; the bulk .read form also 2147483647 and -1). Enumeration and full-table certification sample the operand at 0 and certify the instruction shape. Squash of: - [CUDA][TIRx] Allow collector-A-only SM107 block-scale MMA - feat(lower-tirx): open cp.async wait-group immediates --- python/tvm/backend/cuda/ptx/table.py | 18 ++++++---- tests/python/tirx/codegen/test_ptx_dialect.py | 36 ++++++++++++++++++- 2 files changed, 47 insertions(+), 7 deletions(-) diff --git a/python/tvm/backend/cuda/ptx/table.py b/python/tvm/backend/cuda/ptx/table.py index 1aeb7dda50f8..94e0121f829d 100644 --- a/python/tvm/backend/cuda/ptx/table.py +++ b/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.) diff --git a/tests/python/tirx/codegen/test_ptx_dialect.py b/tests/python/tirx/codegen/test_ptx_dialect.py index 2304824933b0..b599c06e6e1d 100644 --- a/tests/python/tirx/codegen/test_ptx_dialect.py +++ b/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():