Fix ci breakage for missing protection to is_exclusive

python/cudnn/gemm/frost/kernel_templates/_tile_helpers.py

@@ -129,3 +129,47 @@ def copy_tensormap_to_workspace(src_desc_ptr, dst_i64_ptr) -> None:
     src_words = cute.make_ptr(cutlass.Int64, src_desc_ptr.toint(), mem_space=cute.AddressSpace.generic)
     for i in cutlass.range_constexpr(TENSOR_MAP_QWORDS):
         dst_i64_ptr.subview(i).store((src_words + i).load())
+
+
+def tcgen05_alloc(tmem_ptr, num_cols, *, is_exclusive=False, group=None):
+    if is_exclusive:
+        nvvm.tcgen05_alloc(tmem_ptr, num_cols, is_exclusive=True, group=group)
+    else:
+        nvvm.tcgen05_alloc(tmem_ptr, num_cols, group=group)
+
+
+def tcgen05_dealloc(tmem_ptr, num_cols, *, is_exclusive=False, group=None):
+    if is_exclusive:
+        nvvm.tcgen05_dealloc(tmem_ptr, num_cols, is_exclusive=True, group=group)
+    else:
+        nvvm.tcgen05_dealloc(tmem_ptr, num_cols, group=group)
+
+
+def tcgen05_mma_block_scale(mma_kind, cta_group, d, a, b, idesc, *, enable_input_d, scale_a, scale_b, scale_vec_size, b_collector_op=None):
+    if b_collector_op is None:
+        nvvm.tcgen05_mma_block_scale(
+            mma_kind,
+            cta_group,
+            d,
+            a,
+            b,
+            idesc,
+            enable_input_d=enable_input_d,
+            scale_a=scale_a,
+            scale_b=scale_b,
+            scale_vec_size=scale_vec_size,
+        )
+    else:
+        nvvm.tcgen05_mma_block_scale(
+            mma_kind,
+            cta_group,
+            d,
+            a,
+            b,
+            idesc,
+            enable_input_d=enable_input_d,
+            scale_a=scale_a,
+            scale_b=scale_b,
+            scale_vec_size=scale_vec_size,
+            b_collector_op=b_collector_op,
+        )

python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma.py

@@ -25,6 +25,9 @@
 import cutlass.experimental.primitives as nvvm
 from cudnn.gemm.frost.kernel_templates._tile_helpers import (
     l2_swizzle_tile as _l2_swizzle_tile,
+    tcgen05_alloc as _tcgen05_alloc,
+    tcgen05_dealloc as _tcgen05_dealloc,
+    tcgen05_mma_block_scale as _tcgen05_mma_block_scale,
 )
 import cutlass.experimental.cuda.tensor_map as _tma
 import cutlass._mlir_helpers.vector as _cvec
@@ -608,7 +611,7 @@ def _kernel(
 
     if warp_idx == mma_warp_id:
         nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE)
-        nvvm.tcgen05_alloc(
+        _tcgen05_alloc(
             tmem_ptr_i32,
             cutlass.Int32(num_tmem_alloc_cols),
             is_exclusive=tmem_alloc_exclusive,
@@ -762,7 +765,7 @@ def _kernel(
                                 # shared; A's SF word block follows the M block.
                                 desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi)
                                 if elect_one:
-                                    nvvm.tcgen05_mma_block_scale(
+                                    _tcgen05_mma_block_scale(
                                         mma_block_scale_kind,
                                         nvvm.CTAGroup.CTA_1,
                                         acc_tmem_ptrs[g][mi],
@@ -831,7 +834,7 @@ def _kernel(
 
         nvvm.bar_warp_sync(0xFFFFFFFF)
         alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32)
-        nvvm.tcgen05_dealloc(
+        _tcgen05_dealloc(
             alloc_ptr,
             cutlass.Int32(num_tmem_alloc_cols),
             is_exclusive=tmem_alloc_exclusive,

python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma_static.py

@@ -27,6 +27,9 @@
 import cutlass.experimental.primitives as nvvm
 from cudnn.gemm.frost.kernel_templates._tile_helpers import (
     l2_swizzle_tile as _l2_swizzle_tile,
+    tcgen05_alloc as _tcgen05_alloc,
+    tcgen05_dealloc as _tcgen05_dealloc,
+    tcgen05_mma_block_scale as _tcgen05_mma_block_scale,
 )
 import cutlass.experimental.cuda.tensor_map as _tma
 import cutlass._mlir_helpers.vector as _cvec
@@ -513,7 +516,7 @@ def _kernel(
 
     if warp_idx == mma_warp_id:
         nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE)
-        nvvm.tcgen05_alloc(
+        _tcgen05_alloc(
             tmem_ptr_i32,
             cutlass.Int32(num_tmem_alloc_cols),
             is_exclusive=tmem_alloc_exclusive,
@@ -668,7 +671,7 @@ def _kernel(
                                 # shared; A's SF word block follows the M block.
                                 desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi)
                                 if elect_one:
-                                    nvvm.tcgen05_mma_block_scale(
+                                    _tcgen05_mma_block_scale(
                                         mma_block_scale_kind,
                                         nvvm.CTAGroup.CTA_1,
                                         acc_tmem_ptrs[g][mi],
@@ -721,7 +724,7 @@ def _kernel(
 
         nvvm.bar_warp_sync(0xFFFFFFFF)
         alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32)
-        nvvm.tcgen05_dealloc(
+        _tcgen05_dealloc(
             alloc_ptr,
             cutlass.Int32(num_tmem_alloc_cols),
             is_exclusive=tmem_alloc_exclusive,

python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma.py

@@ -24,6 +24,9 @@
 import cutlass.experimental.primitives as nvvm
 from cudnn.gemm.frost.kernel_templates._tile_helpers import (
     l2_swizzle_tile as _l2_swizzle_tile,
+    tcgen05_alloc as _tcgen05_alloc,
+    tcgen05_dealloc as _tcgen05_dealloc,
+    tcgen05_mma_block_scale as _tcgen05_mma_block_scale,
 )
 import cutlass.experimental.cuda.tensor_map as _tma
 import cutlass._mlir_helpers.vector as _cvec
@@ -635,7 +638,7 @@ def _kernel(
     ab_empty_arrive_mask = cutlass.Int16(a_part | b_part)
     if warp_idx == mma_warp_id:
         nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE)
-        nvvm.tcgen05_alloc(
+        _tcgen05_alloc(
             tmem_ptr_i32,
             cutlass.Int32(num_tmem_alloc_cols),
             is_exclusive=tmem_alloc_exclusive,
@@ -797,7 +800,7 @@ def _kernel(
                                     # are shared; A's SF word block follows the M block.
                                     desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi)
                                     if elect_one:
-                                        nvvm.tcgen05_mma_block_scale(
+                                        _tcgen05_mma_block_scale(
                                             mma_block_scale_kind,
                                             nvvm.CTAGroup.CTA_2,
                                             acc_tmem_ptrs[g][mi],
@@ -874,7 +877,7 @@ def _kernel(
             if cutlass.const_expr(not use_acc_overlap):
                 nvvm.mbarrier_arrive(peer_mbar, scope=nvvm.MemScope.CLUSTER, relaxed=True)
             alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32)
-            nvvm.tcgen05_dealloc(
+            _tcgen05_dealloc(
                 alloc_ptr,
                 cutlass.Int32(num_tmem_alloc_cols),
                 is_exclusive=tmem_alloc_exclusive,
@@ -914,7 +917,7 @@ def _kernel(
             while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000):
                 pass
             alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32)
-            nvvm.tcgen05_dealloc(
+            _tcgen05_dealloc(
                 alloc_ptr,
                 cutlass.Int32(num_tmem_alloc_cols),
                 is_exclusive=tmem_alloc_exclusive,

python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma_static.py

@@ -25,6 +25,9 @@
 import cutlass.experimental.primitives as nvvm
 from cudnn.gemm.frost.kernel_templates._tile_helpers import (
     l2_swizzle_tile as _l2_swizzle_tile,
+    tcgen05_alloc as _tcgen05_alloc,
+    tcgen05_dealloc as _tcgen05_dealloc,
+    tcgen05_mma_block_scale as _tcgen05_mma_block_scale,
 )
 import cutlass.experimental.cuda.tensor_map as _tma
 import cutlass._mlir_helpers.vector as _cvec
@@ -535,7 +538,7 @@ def _kernel(
     ab_empty_arrive_mask = cutlass.Int16(a_part | b_part)
     if warp_idx == mma_warp_id:
         nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE)
-        nvvm.tcgen05_alloc(
+        _tcgen05_alloc(
             tmem_ptr_i32,
             cutlass.Int32(num_tmem_alloc_cols),
             is_exclusive=tmem_alloc_exclusive,
@@ -697,7 +700,7 @@ def _kernel(
                                     # are shared; A's SF word block follows the M block.
                                     desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi)
                                     if elect_one:
-                                        nvvm.tcgen05_mma_block_scale(
+                                        _tcgen05_mma_block_scale(
                                             mma_block_scale_kind,
                                             nvvm.CTAGroup.CTA_2,
                                             acc_tmem_ptrs[g][mi],
@@ -759,7 +762,7 @@ def _kernel(
             if cutlass.const_expr(not use_acc_overlap):
                 nvvm.mbarrier_arrive(peer_mbar, scope=nvvm.MemScope.CLUSTER, relaxed=True)
             alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32)
-            nvvm.tcgen05_dealloc(
+            _tcgen05_dealloc(
                 alloc_ptr,
                 cutlass.Int32(num_tmem_alloc_cols),
                 is_exclusive=tmem_alloc_exclusive,
@@ -777,7 +780,7 @@ def _kernel(
             while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000):
                 pass
             alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32)
-            nvvm.tcgen05_dealloc(
+            _tcgen05_dealloc(
                 alloc_ptr,
                 cutlass.Int32(num_tmem_alloc_cols),
                 is_exclusive=tmem_alloc_exclusive,

python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py

@@ -28,6 +28,8 @@
 from cudnn.gemm.frost.kernel_templates._tile_helpers import (
     epi_subtile_spans as _epi_subtile_spans,
     l2_swizzle_tile as _l2_swizzle_tile,
+    tcgen05_alloc as _tcgen05_alloc,
+    tcgen05_dealloc as _tcgen05_dealloc,
 )
 import cutlass.experimental.cuda.tensor_map as _tma
 import cutlass._mlir_helpers.vector as _cvec
@@ -543,7 +545,7 @@ def _kernel(
 
     if warp_idx == mma_warp_id:
         nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE)
-        nvvm.tcgen05_alloc(
+        _tcgen05_alloc(
             tmem_ptr_i32,
             cutlass.Int32(num_tmem_alloc_cols),
             is_exclusive=tmem_alloc_exclusive,
@@ -681,7 +683,7 @@ def _kernel(
         nvvm.bar_warp_sync(0xFFFFFFFF)
         nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1)
         alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32)
-        nvvm.tcgen05_dealloc(
+        _tcgen05_dealloc(
             alloc_ptr,
             cutlass.Int32(num_tmem_alloc_cols),
             is_exclusive=tmem_alloc_exclusive,

python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma_static.py

@@ -26,6 +26,8 @@
 from cudnn.gemm.frost.kernel_templates._tile_helpers import (
     epi_subtile_spans as _epi_subtile_spans,
     l2_swizzle_tile as _l2_swizzle_tile,
+    tcgen05_alloc as _tcgen05_alloc,
+    tcgen05_dealloc as _tcgen05_dealloc,
 )
 import cutlass.experimental.cuda.tensor_map as _tma
 import cutlass._mlir_helpers.vector as _cvec
@@ -440,7 +442,7 @@ def _kernel(
 
     if warp_idx == mma_warp_id:
         nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE)
-        nvvm.tcgen05_alloc(
+        _tcgen05_alloc(
             tmem_ptr_i32,
             cutlass.Int32(num_tmem_alloc_cols),
             is_exclusive=tmem_alloc_exclusive,
@@ -559,7 +561,7 @@ def _kernel(
         nvvm.bar_warp_sync(0xFFFFFFFF)
         nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1)
         alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32)
-        nvvm.tcgen05_dealloc(
+        _tcgen05_dealloc(
             alloc_ptr,
             cutlass.Int32(num_tmem_alloc_cols),
             is_exclusive=tmem_alloc_exclusive,

python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py

@@ -25,6 +25,8 @@
 from cudnn.gemm.frost.kernel_templates._tile_helpers import (
     epi_subtile_spans as _epi_subtile_spans,
     l2_swizzle_tile as _l2_swizzle_tile,
+    tcgen05_alloc as _tcgen05_alloc,
+    tcgen05_dealloc as _tcgen05_dealloc,
 )
 import cutlass.experimental.cuda.tensor_map as _tma
 import cutlass._mlir_helpers.vector as _cvec
@@ -563,7 +565,7 @@ def _kernel(
     ab_empty_arrive_mask = cutlass.Int16(a_part | b_part)
     if warp_idx == mma_warp_id:
         nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE)
-        nvvm.tcgen05_alloc(
+        _tcgen05_alloc(
             tmem_ptr_i32,
             cutlass.Int32(num_tmem_alloc_cols),
             is_exclusive=tmem_alloc_exclusive,
@@ -717,7 +719,7 @@ def _kernel(
                 pass
             nvvm.mbarrier_arrive(peer_mbar, scope=nvvm.MemScope.CLUSTER, relaxed=True)
             alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32)
-            nvvm.tcgen05_dealloc(
+            _tcgen05_dealloc(
                 alloc_ptr,
                 cutlass.Int32(num_tmem_alloc_cols),
                 is_exclusive=tmem_alloc_exclusive,
@@ -756,7 +758,7 @@ def _kernel(
             while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000):
                 pass
             alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32)
-            nvvm.tcgen05_dealloc(
+            _tcgen05_dealloc(
                 alloc_ptr,
                 cutlass.Int32(num_tmem_alloc_cols),
                 is_exclusive=tmem_alloc_exclusive,

python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma_static.py

@@ -25,6 +25,8 @@
 from cudnn.gemm.frost.kernel_templates._tile_helpers import (
     epi_subtile_spans as _epi_subtile_spans,
     l2_swizzle_tile as _l2_swizzle_tile,
+    tcgen05_alloc as _tcgen05_alloc,
+    tcgen05_dealloc as _tcgen05_dealloc,
 )
 import cutlass.experimental.cuda.tensor_map as _tma
 import cutlass._mlir_helpers.vector as _cvec
@@ -461,7 +463,7 @@ def _kernel(
     ab_empty_arrive_mask = cutlass.Int16(a_part | b_part)
     if warp_idx == mma_warp_id:
         nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE)
-        nvvm.tcgen05_alloc(
+        _tcgen05_alloc(
             tmem_ptr_i32,
             cutlass.Int32(num_tmem_alloc_cols),
             is_exclusive=tmem_alloc_exclusive,
@@ -599,7 +601,7 @@ def _kernel(
                 pass
             nvvm.mbarrier_arrive(peer_mbar, scope=nvvm.MemScope.CLUSTER, relaxed=True)
             alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32)
-            nvvm.tcgen05_dealloc(
+            _tcgen05_dealloc(
                 alloc_ptr,
                 cutlass.Int32(num_tmem_alloc_cols),
                 is_exclusive=tmem_alloc_exclusive,
@@ -616,7 +618,7 @@ def _kernel(
             while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000):
                 pass
             alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32)
-            nvvm.tcgen05_dealloc(
+            _tcgen05_dealloc(
                 alloc_ptr,
                 cutlass.Int32(num_tmem_alloc_cols),
                 is_exclusive=tmem_alloc_exclusive,

python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_1ctamma.py

@@ -26,6 +26,8 @@
 from cudnn.gemm.frost.kernel_templates._tile_helpers import (
     epi_subtile_spans as _epi_subtile_spans,
     l2_swizzle_tile as _l2_swizzle_tile,
+    tcgen05_alloc as _tcgen05_alloc,
+    tcgen05_dealloc as _tcgen05_dealloc,
 )
 import cutlass.experimental.cuda.tensor_map as _tma
 import cutlass._mlir_helpers.vector as _cvec
@@ -556,7 +558,7 @@ def _kernel(
 
     if warp_idx == mma_warp_id:
         nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE)
-        nvvm.tcgen05_alloc(
+        _tcgen05_alloc(
             tmem_ptr_i32,
             cutlass.Int32(num_tmem_alloc_cols),
             is_exclusive=tmem_alloc_exclusive,
@@ -707,7 +709,7 @@ def _kernel(
         nvvm.bar_warp_sync(0xFFFFFFFF)
         nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1)
         alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32)
-        nvvm.tcgen05_dealloc(
+        _tcgen05_dealloc(
             alloc_ptr,
             cutlass.Int32(num_tmem_alloc_cols),
             is_exclusive=tmem_alloc_exclusive,