[BACKEND][EZ] Tighten TMA multicta verifier (#9941)

lezcano · web-flow · commit 76fbef9dbbcf · 2026-04-10T08:34:32.000Z
As per the PTX docs, TMAs have a very specific behaviour when executed
in a 2CTA kernel:

&gt;.cta_group::1 : The mbarrier signal is also multicasted to the same
offset as mbar in the shared memory of the destination CTA.
.cta_group::2 : The mbarrier signal is multicasted either to all the odd
numbered CTAs or the even numbered CTAs within the corresponding
CTA-Pair. For each destination CTA specified in the ctaMask, the
mbarrier signal is sent either to the destination CTA or its peer-CTA
based on CTAs %cluster_ctarank parity of shared memory where the
mbarrier object mbar resides.

As such, we require these CTA layouts in TMA barriers.
diff --git a/lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp b/lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp
@@ -290,6 +290,62 @@ LogicalResult ClusterBarrierOp::verify() {
 }
 
 // -- TMA operation verifiers --
+static std::string formatCGALayout(CGAEncodingAttr cgaLayout) {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  auto kBlock = StringAttr::get(cgaLayout.getContext(), "block");
+  os << "[";
+  llvm::interleaveComma(cgaLayout.getLinearLayout().getBases().lookup(kBlock),
+                        os, [&](const auto &basis) {
+                          os << "[";
+                          llvm::interleaveComma(basis, os);
+                          os << "]";
+                        });
+  os << "]";
+  return os.str();
+}
+
+static LogicalResult verifyBarrierCGALayout(Operation *op, Value barrier,
+                                            CGAEncodingAttr expectedCGALayout,
+                                            StringRef barrierName) {
+  auto barrierTy = cast<MemDescType>(barrier.getType());
+  auto actualCGALayout = getCGALayout(barrierTy.getEncoding());
+  if (actualCGALayout != expectedCGALayout)
+    return op->emitOpError() << barrierName << " cga_layout must be "
+                             << formatCGALayout(expectedCGALayout) << ", got "
+                             << formatCGALayout(actualCGALayout);
+  return success();
+}
+
+static LogicalResult verifyCompletionBarrierLayout(Operation *op,
+                                                   Value barrier) {
+  auto expectedCGALayout =
+      CGAEncodingAttr::get1DLayout(op->getContext(), gpu::lookupNumCTAs(op));
+  return verifyBarrierCGALayout(op, barrier, expectedCGALayout,
+                                "completion barrier");
+}
+
+static LogicalResult verifyTMABarrierLayout(Operation *op, Value barrier) {
+  auto twoCTAsAttr =
+      op->getParentOfType<ModuleOp>()->getAttrOfType<BoolAttr>(AttrTwoCTAsName);
+  if (!twoCTAsAttr)
+    return success();
+
+  auto ctx = op->getContext();
+  int numCTAs = gpu::lookupNumCTAs(op);
+  CGAEncodingAttr expectedCGALayout;
+  if (twoCTAsAttr.getValue()) {
+    auto kBlock = StringAttr::get(ctx, "block");
+    auto dim = standardOutDimNames(ctx, /*rank=*/1)[0];
+    auto layout = LinearLayout::zeros1D(2, kBlock, dim) *
+                  LinearLayout::identity1D(numCTAs / 2, kBlock, dim);
+    expectedCGALayout = CGAEncodingAttr::get(ctx, std::move(layout));
+  } else {
+    expectedCGALayout = CGAEncodingAttr::get1DLayout(ctx, numCTAs);
+  }
+  return verifyBarrierCGALayout(op, barrier, expectedCGALayout, "TMA barrier");
+}
+
 static LogicalResult verifyTMAEncoding(Operation *op, TensorDescInterface desc,
                                        Attribute enc) {
   auto nvmma = dyn_cast<NVMMASharedEncodingAttr>(enc);
@@ -318,6 +374,8 @@ static LogicalResult verifyAsyncTMALoadOp(Operation *op,
                                           MemDescType resultType) {
   if (failed(verifyBarrierType(op, barrier.getType())))
     return failure();
+  if (failed(verifyTMABarrierLayout(op, barrier)))
+    return failure();
   if (!resultType.getMutableMemory())
     return op->emitOpError("cannot store into immutable memory");
   if (failed(verifyTMAEncoding(op, desc, resultType.getEncoding())))
@@ -599,34 +657,6 @@ static LogicalResult verifyMMADType(Operation *op, Type a, Type b, Type d) {
   return success();
 }
 
-static std::string formatCGALayout(CGAEncodingAttr cgaLayout) {
-  std::string str;
-  llvm::raw_string_ostream os(str);
-  auto kBlock = StringAttr::get(cgaLayout.getContext(), "block");
-  os << "[";
-  llvm::interleaveComma(cgaLayout.getLinearLayout().getBases().lookup(kBlock),
-                        os, [&](const auto &basis) {
-                          os << "[";
-                          llvm::interleaveComma(basis, os);
-                          os << "]";
-                        });
-  os << "]";
-  return os.str();
-}
-
-static LogicalResult verifyCompletionBarrierLayout(Operation *op,
-                                                   Value barrier) {
-  auto barrierTy = cast<MemDescType>(barrier.getType());
-  auto expectedCGALayout =
-      CGAEncodingAttr::get1DLayout(op->getContext(), gpu::lookupNumCTAs(op));
-  auto actualCGALayout = getCGALayout(barrierTy.getEncoding());
-  if (actualCGALayout != expectedCGALayout)
-    return op->emitOpError("completion barrier cga_layout must be ")
-           << formatCGALayout(expectedCGALayout) << ", got "
-           << formatCGALayout(actualCGALayout);
-  return success();
-}
-
 LogicalResult TCGen5MMAOp::verify() {
   if (!getIsAsync() && !getBarriers().empty()) {
     return emitOpError("The op is synchronous but a barrier is present.");
diff --git a/python/test/gluon/test_consan.py b/python/test/gluon/test_consan.py
@@ -158,13 +158,10 @@ def test_consan_uses_profile_scratch(device, fresh_knobs, num_ctas):
 
 @pytest.mark.skipif(not is_cuda() or torch.cuda.get_device_capability()[0] < 9, reason="Requires hopper or newer")
 @pytest.mark.parametrize("FAILURE", [True, False])
-@pytest.mark.parametrize("TWO_CTA_BARRIER", [False, True])
-def test_async_tma_kernel(FAILURE, TWO_CTA_BARRIER, device, run_wrapper, monkeypatch, num_ctas):
-    if TWO_CTA_BARRIER and num_ctas == 1:
-        pytest.skip("Need at least 2 CTAs for a two-CTA barrier")
+def test_async_tma_kernel(FAILURE, device, run_wrapper, monkeypatch, num_ctas):
     if run_wrapper:
-        result = run_in_process(test_async_tma_kernel, (FAILURE, TWO_CTA_BARRIER, device, False, monkeypatch, num_ctas))
-        if FAILURE or TWO_CTA_BARRIER:
+        result = run_in_process(test_async_tma_kernel, (FAILURE, device, False, monkeypatch, num_ctas))
+        if FAILURE:
             assert_expected_cuda_failure(result.exc)
             assert "Buffer being accessed has outstanding writes" in result.driver_stderr_output
         else:
@@ -177,13 +174,13 @@ def test_async_tma_kernel(FAILURE, TWO_CTA_BARRIER, device, run_wrapper, monkeyp
     knobs.refresh_knobs()
 
     @gluon.jit
-    def kernel(input_desc, out, FAILURE: ttgl.constexpr, TWO_CTA_BARRIER: ttgl.constexpr):
+    def kernel(input_desc, out, FAILURE: ttgl.constexpr):
         block_m: ttgl.constexpr = XBLOCK * ttgl.num_ctas()
         cga_layout: ttgl.constexpr = default_cga_layout(ttgl.num_ctas(), 2)
         blocked_layout: ttgl.constexpr = ttgl.BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[32, 1],
                                                             warps_per_cta=[4, 1], order=[0, 1], cga_layout=cga_layout)
         smem = ttgl.allocate_shared_memory(ttgl.float16, [block_m, XBLOCK], input_desc.layout)
-        bar = mbarrier.allocate_mbarrier(two_ctas=TWO_CTA_BARRIER)
+        bar = mbarrier.allocate_mbarrier()
         mbarrier.init(bar, count=1)
         mbarrier.expect(bar, input_desc.nbytes_per_cta)
         tma.async_copy_global_to_shared(input_desc, [0, 0], bar, smem)
@@ -203,19 +200,17 @@ def kernel(input_desc, out, FAILURE: ttgl.constexpr, TWO_CTA_BARRIER: ttgl.const
     shared_layout = ttgl.NVMMASharedLayout(swizzle_byte_width=128, element_bitwidth=16, rank=2,
                                            cga_layout=default_cga_layout(num_ctas, 2))
     input_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(input, [block_m, XBLOCK.value], shared_layout)
-    kernel[(1, )](input_desc, output, FAILURE=FAILURE, TWO_CTA_BARRIER=TWO_CTA_BARRIER, num_warps=4, num_ctas=num_ctas)
+    kernel[(1, )](input_desc, output, FAILURE=FAILURE, num_warps=4, num_ctas=num_ctas)
 
 
 @pytest.mark.skipif(not is_cuda() or torch.cuda.get_device_capability()[0] < 9, reason="Requires hopper or newer")
 @pytest.mark.parametrize("FAILURE", [True, False])
-@pytest.mark.parametrize("TWO_CTA_BARRIER", [False, True])
-def test_async_tma_multicast_kernel(FAILURE, TWO_CTA_BARRIER, device, run_wrapper, monkeypatch, num_ctas):
+def test_async_tma_multicast_kernel(FAILURE, device, run_wrapper, monkeypatch, num_ctas):
     if num_ctas == 1:
         pytest.skip("Need at least 2 CTAs for multicast in this test")
     if run_wrapper:
-        result = run_in_process(test_async_tma_multicast_kernel,
-                                (FAILURE, TWO_CTA_BARRIER, device, False, monkeypatch, num_ctas))
-        if FAILURE or TWO_CTA_BARRIER:
+        result = run_in_process(test_async_tma_multicast_kernel, (FAILURE, device, False, monkeypatch, num_ctas))
+        if FAILURE:
             assert_expected_cuda_failure(result.exc)
             assert "Buffer being accessed has outstanding writes" in result.driver_stderr_output
         else:
@@ -228,12 +223,12 @@ def test_async_tma_multicast_kernel(FAILURE, TWO_CTA_BARRIER, device, run_wrappe
     knobs.refresh_knobs()
 
     @gluon.jit
-    def kernel(input_desc, out, FAILURE: ttgl.constexpr, TWO_CTA_BARRIER: ttgl.constexpr):
+    def kernel(input_desc, out, FAILURE: ttgl.constexpr):
         cga_layout: ttgl.constexpr = multicast_cga_layout(ttgl.num_ctas(), 2)
         blocked_layout: ttgl.constexpr = ttgl.BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[32, 1],
                                                             warps_per_cta=[4, 1], order=[0, 1], cga_layout=cga_layout)
         smem = ttgl.allocate_shared_memory(ttgl.float16, [XBLOCK, XBLOCK], input_desc.layout)
-        bar = mbarrier.allocate_mbarrier(two_ctas=TWO_CTA_BARRIER)
+        bar = mbarrier.allocate_mbarrier()
         mbarrier.init(bar, count=1)
         mbarrier.expect(bar, input_desc.nbytes_per_cta)
         tma.async_copy_global_to_shared(input_desc, [0, 0], bar, smem, multicast=True)
@@ -252,54 +247,7 @@ def kernel(input_desc, out, FAILURE: ttgl.constexpr, TWO_CTA_BARRIER: ttgl.const
     shared_layout = ttgl.NVMMASharedLayout(swizzle_byte_width=128, element_bitwidth=16, rank=2,
                                            cga_layout=multicast_cga_layout(num_ctas, 2))
     input_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(input, [XBLOCK.value, XBLOCK.value], shared_layout)
-    kernel[(1, )](input_desc, output, FAILURE=FAILURE, TWO_CTA_BARRIER=TWO_CTA_BARRIER, num_warps=4, num_ctas=num_ctas)
-
-
-@pytest.mark.skipif(not is_cuda() or torch.cuda.get_device_capability()[0] < 9, reason="Requires hopper or newer")
-@pytest.mark.parametrize("TWO_CTA_BARRIER", [False, True])
-def test_async_tma_multicast_kernel_two_cta_barrier(TWO_CTA_BARRIER, device, run_wrapper, monkeypatch, num_ctas):
-    if num_ctas != 2:
-        pytest.skip("This test covers a single 2-CTA multicast group")
-    if run_wrapper:
-        result = run_in_process(test_async_tma_multicast_kernel_two_cta_barrier,
-                                (TWO_CTA_BARRIER, device, False, monkeypatch, num_ctas))
-        if TWO_CTA_BARRIER:
-            assert_expected_cuda_failure(result.exc)
-            assert "Buffer being accessed has outstanding writes" in result.driver_stderr_output
-        else:
-            assert result.exc is None
-            assert result.driver_stderr_output == ""
-        return
-
-    monkeypatch.setenv("TRITON_INSTRUMENTATION_MODE", "consan")
-    monkeypatch.setenv("CUDA_LAUNCH_BLOCKING", "1")
-    knobs.refresh_knobs()
-
-    @gluon.jit
-    def kernel(input_desc, out, TWO_CTA_BARRIER: ttgl.constexpr):
-        cga_layout: ttgl.constexpr = multicast_cga_layout(ttgl.num_ctas(), 2)
-        blocked_layout: ttgl.constexpr = ttgl.BlockedLayout(size_per_thread=[1, 1], threads_per_warp=[32, 1],
-                                                            warps_per_cta=[4, 1], order=[0, 1], cga_layout=cga_layout)
-        smem = ttgl.allocate_shared_memory(ttgl.float16, [XBLOCK, XBLOCK], input_desc.layout)
-        bar = mbarrier.allocate_mbarrier(two_ctas=TWO_CTA_BARRIER)
-        mbarrier.init(bar, count=1)
-        mbarrier.expect(bar, input_desc.nbytes_per_cta)
-        tma.async_copy_global_to_shared(input_desc, [0, 0], bar, smem, multicast=True)
-        mbarrier.wait(bar, 0, deps=[smem])
-        val = smem.load(blocked_layout)
-        mbarrier.invalidate(bar)
-
-        out_m = ttgl.arange(0, XBLOCK, ttgl.SliceLayout(1, blocked_layout))[:, None]
-        out_n = ttgl.arange(0, XBLOCK, ttgl.SliceLayout(0, blocked_layout))[None, :]
-        out_ptr = out + out_m * XBLOCK + out_n
-        ttgl.store(out_ptr, val)
-
-    input = torch.randn((XBLOCK.value, XBLOCK.value), device=device, dtype=torch.float16)
-    output = torch.empty((XBLOCK.value, XBLOCK.value), device=device, dtype=torch.float16)
-    shared_layout = ttgl.NVMMASharedLayout(swizzle_byte_width=128, element_bitwidth=16, rank=2,
-                                           cga_layout=multicast_cga_layout(num_ctas, 2))
-    input_desc = gluon.nvidia.hopper.TensorDescriptor.from_tensor(input, [XBLOCK.value, XBLOCK.value], shared_layout)
-    kernel[(1, )](input_desc, output, TWO_CTA_BARRIER=TWO_CTA_BARRIER, num_warps=4, num_ctas=num_ctas)
+    kernel[(1, )](input_desc, output, FAILURE=FAILURE, num_warps=4, num_ctas=num_ctas)
 
 
 @pytest.mark.skipif(not is_cuda() or torch.cuda.get_device_capability()[0] < 9, reason="Requires hopper or newer")
@@ -741,21 +689,24 @@ def kernel(input_desc, output_desc, FAILURE: ttgl.constexpr, MEM_ACCESS_KIND: tt
             ttgl.NVMMASharedLayout.get_default_for([XBLOCK, block_n], ttgl.float16,
                                                    cga_layout=mma_cga_layout(ttgl.num_ctas(), 1, TWO_CTAS)),
         )
-        bar = mbarrier.allocate_mbarrier(batch=2)
+        mma_bar = mbarrier.allocate_mbarrier()
         acc = blackwell.allocate_tensor_memory(ttgl.float32, [block_m, block_n], acc_layout)
-        mbarrier.init(bar.index(0), count=1)
-        mbarrier.init(bar.index(1), count=1)
+        mbarrier.init(mma_bar, count=1)
+        if MEM_ACCESS_KIND == "tma_cp":
+            tma_bar = mbarrier.allocate_mbarrier(two_ctas=TWO_CTAS)
+            mbarrier.init(tma_bar, count=1)
 
         blackwell.tcgen05_mma(smemA, smemB, acc)
-        blackwell.tcgen05_commit(bar.index(0))
+        blackwell.tcgen05_commit(mma_bar)
 
         if not FAILURE:
-            mbarrier.wait(bar.index(0), 0)
+            mbarrier.wait(mma_bar, 0)
 
         if MEM_ACCESS_KIND == "tma_cp":
-            mbarrier.expect(bar.index(1), input_desc.nbytes_per_cta)
-            tma.async_copy_global_to_shared(input_desc, [0, 0], bar.index(1), smemA)
-            mbarrier.wait(bar.index(1), 0)
+            mbarrier.expect(tma_bar, input_desc.nbytes_per_cta)
+            tma.async_copy_global_to_shared(input_desc, [0, 0], tma_bar, smemA)
+            mbarrier.wait(tma_bar, 0)
+            mbarrier.invalidate(tma_bar)
         elif MEM_ACCESS_KIND == "local_store":
             smemA.store(ttgl.full([block_m, XBLOCK], 42, ttgl.float16, smem_a_blocked_layout))
         elif MEM_ACCESS_KIND == "tmem_load":
@@ -769,8 +720,7 @@ def kernel(input_desc, output_desc, FAILURE: ttgl.constexpr, MEM_ACCESS_KIND: tt
         elif MEM_ACCESS_KIND == "tmem_store":
             acc.store(ttgl.full([block_m, block_n], 42, ttgl.float32, acc_blocked_layout))
 
-        mbarrier.invalidate(bar.index(0))
-        mbarrier.invalidate(bar.index(1))
+        mbarrier.invalidate(mma_bar)
 
     block_m = mma_block_m(num_ctas)
     block_n = mma_block_n(num_ctas)
diff --git a/python/tutorials/gluon/14-multicta.py b/python/tutorials/gluon/14-multicta.py
@@ -477,8 +477,9 @@ def broadcast(b):
 # every CTA in the multicast group atomically, so the wait side does not need a
 # different API.
 #
-# The only new ingredient is the layout. The TMA destination must use a
-# broadcast `cga_layout`, so that both CTAs view the same shared-memory tile.
+# The TMA destination must use a broadcast `cga_layout`, so that both CTAs
+# receive the same shared-memory tile. The barrier stays a regular 1D TMA
+# barrier unless the kernel is in 2CTA mode.
 #
 # The example below keeps things intentionally simple: it multicasts one tile
 # into shared memory and then materializes that same tile back to global memory.
@@ -489,6 +490,7 @@ def tma_multicast_copy_kernel(in_desc, out_desc):
     gl.static_assert(gl.num_ctas() == 2)
 
     smem = gl.allocate_shared_memory(in_desc.dtype, in_desc.block_shape, in_desc.layout)
+    # This kernel is not in 2CTA mode, so the TMA barrier is per-CTA.
     bar = mbarrier.allocate_mbarrier()
     mbarrier.init(bar, count=1)
 
diff --git a/test/Conversion/tritonnvidiagpu_to_llvm.mlir b/test/Conversion/tritonnvidiagpu_to_llvm.mlir
@@ -167,14 +167,14 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
 #shared0_cluster = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0], CGALayout = [[0]]}>
 #shared1_cga = #ttg.nvmma_shared<{swizzlingByteWidth = 0, transposed = false, elementBitWidth = 8, CGALayout = [[0, 0]]}>
 #smem = #ttg.shared_memory
-module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32} {
+module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.two-ctas" = true} {
   // CHECK-LABEL: tma_copy_barrier_mask_nonzero
   // Barrier pointer is modified when barrier mask != 0
   // CHECK: llvm.ptrtoint
   // CHECK: llvm.and
   // CHECK: llvm.inttoptr
   // TMA uses shared::cluster when barrier mask is non-zero
-  // CHECK: cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes
+  // CHECK: cp.async.bulk.tensor.2d.cta_group::2.shared::cluster.global.mbarrier::complete_tx::bytes
   // CHECK-NOT: cp.async.bulk.tensor.2d.shared::cta.global.mbarrier
   tt.func @tma_copy_barrier_mask_nonzero(%tma: !tt.tensordesc<128x128xf32, #shared1_cga>, %alloc: !ttg.memdesc<128x128xf32, #shared1_cga, #smem, mutable>, %x: i32, %barrier: !ttg.memdesc<1xi64, #shared0_cluster, #smem>, %pred: i1) {
     ttng.async_tma_copy_global_to_local %tma[%x, %x] %alloc, %barrier, %pred : !tt.tensordesc<128x128xf32, #shared1_cga>, !ttg.memdesc<1xi64, #shared0_cluster, #smem> -> !ttg.memdesc<128x128xf32, #shared1_cga, #smem, mutable>
diff --git a/test/TritonNvidiaGPU/invalid.mlir b/test/TritonNvidiaGPU/invalid.mlir
diff --git a/test/TritonNvidiaGPU/membar-cluster.mlir b/test/TritonNvidiaGPU/membar-cluster.mlir
