[MultiCTA] Implement broadcast semantics in mbarrier.arrive (#9475)

This was missing and it is necessary for multiCTA warp-specialised
kernels.

We also fix `mbarrier.init` to rescale the number of arrivals on the
leader CTA to follow this pattern.

In `mbarrier.expect` we emit `mbarrier.expect` from the leader CTA
(which actas as an `mbarrier.arrive bar, 1`),
and we emit `mbarrier.arrive bar, 1` from the non leader CTAs to go with
the semantics above.
This has as a nice corollary now `expect` also has release semantics,
which is nice. We hope this
should not be a perf issue in real kernels really.

Finally, we implement these semantics in two helper functions and use
them across the codebase.

Author: lezcano

test/Conversion/tritonnvidiagpu_to_llvm.mlir

@@ -13,6 +13,20 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
 
 // -----
 
+#shared0 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0], CGALayout = [[0]]}>
+#smem = #ttg.shared_memory
+module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32} {
+  // CHECK-LABEL: init_barrier_cluster_broadcast
+  tt.func @init_barrier_cluster_broadcast(%alloc: !ttg.memdesc<1xi64, #shared0, #smem>) {
+    // CHECK: nvg.cluster_id
+    // CHECK: @$0 mbarrier.init.shared::cta.b64 [$1], 2;
+    ttng.init_barrier %alloc, 1 : !ttg.memdesc<1xi64, #shared0, #smem>
+    tt.return
+  }
+}
+
+// -----
+
 #shared0 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}>
 #smem = #ttg.shared_memory
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
@@ -45,30 +59,55 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
 
   // CHECK-LABEL: arrive_barrier
   tt.func @arrive_barrier(%alloc: !ttg.memdesc<1xi64, #shared0, #smem>) {
+    // CHECK-NEXT: [[BASE:%.*]] = llvm.extractvalue %arg0[0] : !llvm.struct<(ptr<3>, i32)>
+    // CHECK-NEXT: llvm.extractvalue %arg0[1] : !llvm.struct<(ptr<3>, i32)>
+    // CHECK-NEXT: nvvm.barrier0
     // CHECK-NEXT: [[TID:%.*]] = nvvm.read.ptx.sreg.tid.x
     // CHECK-NEXT: [[C127:%.*]] = llvm.mlir.constant(127 : i32)
     // CHECK-NEXT: [[RTID:%.*]] = llvm.and [[TID]], [[C127]]
     // CHECK-NEXT: [[C0:%.*]] = llvm.mlir.constant(0 : i32)
     // CHECK-NEXT: [[IS_ZERO:%.*]] = llvm.icmp "eq" [[RTID]], [[C0]]
-    // CHECK-NEXT: "@$0 mbarrier.arrive.shared::cta.b64 _, [$1], 2;", "b,r" [[IS_ZERO]], %arg0
+    // CHECK-NEXT: "@$0 mbarrier.arrive.shared::cta.b64 _, [$1], 2;", "b,r" [[IS_ZERO]], [[BASE]]
     ttng.arrive_barrier %alloc, 2 : !ttg.memdesc<1xi64, #shared0, #smem>
     tt.return
   }
 
   // CHECK-LABEL: arrive_barrier_pred
   tt.func @arrive_barrier_pred(%alloc: !ttg.memdesc<1xi64, #shared0, #smem>, %pred: i1) {
+    // CHECK-NEXT: [[BASE:%.*]] = llvm.extractvalue %arg0[0] : !llvm.struct<(ptr<3>, i32)>
+    // CHECK-NEXT: llvm.extractvalue %arg0[1] : !llvm.struct<(ptr<3>, i32)>
+    // CHECK-NEXT: nvvm.barrier0
     // CHECK-NEXT: [[TID:%.*]] = nvvm.read.ptx.sreg.tid.x
     // CHECK-NEXT: [[C127:%.*]] = llvm.mlir.constant(127 : i32)
     // CHECK-NEXT: [[RTID:%.*]] = llvm.and [[TID]], [[C127]]
     // CHECK-NEXT: [[C0:%.*]] = llvm.mlir.constant(0 : i32)
     // CHECK-NEXT: [[IS_ZERO:%.*]] = llvm.icmp "eq" [[RTID]], [[C0]]
     // CHECK-NEXT: [[PRED:%.*]] = llvm.and [[IS_ZERO]], %arg1
-    // CHECK-NEXT: "@$0 mbarrier.arrive.shared::cta.b64 _, [$1], 2;", "b,r" [[PRED]], %arg0
+    // CHECK-NEXT: "@$0 mbarrier.arrive.shared::cta.b64 _, [$1], 2;", "b,r" [[PRED]], [[BASE]]
     ttng.arrive_barrier %alloc, 2, %pred : !ttg.memdesc<1xi64, #shared0, #smem>
     tt.return
   }
 }
 
+// -----
+
+#shared0 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0], CGALayout = [[0]]}>
+#smem = #ttg.shared_memory
+module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32} {
+  // CHECK-LABEL: arrive_barrier_cluster_broadcast
+  tt.func @arrive_barrier_cluster_broadcast(%alloc: !ttg.memdesc<1xi64, #shared0, #smem>) {
+    // CHECK: nvvm.barrier0
+    // CHECK: nvg.cluster_id
+    // CHECK: llvm.ptrtoint
+    // CHECK: llvm.and
+    // CHECK: llvm.inttoptr
+    // CHECK: mbarrier.arrive.shared::cluster.b64
+    // CHECK-NOT: mbarrier.arrive.shared::cta.b64
+    ttng.arrive_barrier %alloc, 1 : !ttg.memdesc<1xi64, #shared0, #smem>
+    tt.return
+  }
+}
+
 
 // -----
 
@@ -218,6 +257,24 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
 
 // -----
 
+#shared0 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0], CGALayout = [[0]]}>
+#smem = #ttg.shared_memory
+module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32} {
+  // CHECK-LABEL: expect_barrier_cluster_broadcast
+  // CHECK: nvg.cluster_id
+  // CHECK: llvm.ptrtoint
+  // CHECK: llvm.and
+  // CHECK: llvm.inttoptr
+  // CHECK: @$0 mbarrier.arrive.expect_tx.shared::cta.b64 _, [$1], 32768;
+  // CHECK: @$0 mbarrier.arrive.shared::cluster.b64 _, [$1], 1;
+  tt.func @expect_barrier_cluster_broadcast(%barrier: !ttg.memdesc<1xi64, #shared0, #smem, mutable>, %pred: i1) {
+    ttng.barrier_expect %barrier, 16384, %pred : !ttg.memdesc<1xi64, #shared0, #smem, mutable>
+    tt.return
+  }
+}
+
+// -----
+
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
   // CHECK-LABEL: byval_tma_desc
   // CHECK: llvm.align = 64

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/BarrierOpToLLVM.cpp

@@ -111,17 +111,28 @@ struct InitBarrierOpConversion
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
     auto b = TritonLLVMOpBuilder(loc, rewriter);
+    auto barrierTy = op.getAlloc().getType();
     auto smemObj = LLVM::getSharedMemoryObjectFromStruct(
         loc, adaptor.getAlloc(),
-        typeConverter->convertType(op.getAlloc().getType().getElementType()),
-        rewriter);
+        typeConverter->convertType(barrierTy.getElementType()), rewriter);
 
     // We use an elect predicate to tell ptxas that the operation is uniform,
     // which results in better codegen.
     Value pred = getElectWarp0OrThread0(*targetInfo, b);
+
+    if (auto leaderPred =
+            LLVM::NVIDIA::getLeaderCTAPredicate(loc, rewriter, barrierTy))
+      pred = b.and_(pred, *leaderPred);
+
+    auto numCTAs = triton::gpu::lookupNumCTAs(op);
+    auto initCount = op.getCount();
+    // The lead barrier accounts for all arrives from CTAs that broadcast into
+    // the same barrier.
+    initCount *= numCTAs / barrierTy.getNumElements();
+
     ::mlir::triton::PTXBuilder ptxBuilder;
     const std::string ptx = "@$0 mbarrier.init.shared::cta.b64 [$1], " +
-                            std::to_string(op.getCount()) + ";";
+                            std::to_string(initCount) + ";";
     auto &barSyncOp = *ptxBuilder.create(ptx);
     barSyncOp({ptxBuilder.newOperand(pred, "b"),
                ptxBuilder.newOperand(smemObj.getBase(), "r")},
@@ -188,31 +199,40 @@ struct BarrierExpectConversion
     auto expectedBytes = op.getSize() * (numCTAs / barrierTy.getNumElements());
 
     auto id = getThreadId(rewriter, loc);
-    Value pred = b.icmp_eq(id, b.i32_val(0));
-    pred = b.and_(pred, adaptor.getPred());
-
-    auto kBlock = StringAttr::get(op->getContext(), "block");
-    auto maskCGABroadcast =
-        toLinearLayout(barrierTy).getFreeVariableMasks().lookup(kBlock);
-    if (maskCGABroadcast) {
-      // If several CTAs cast to the same barrier, as when we do a TMA into a
-      // tcgen05.mma 2CTA, we just register the expect in the lead barrier, as
-      // it is the only one that will receive the mbarrier signals
-      auto ctaId = nvgpu::ClusterCTAIdOp::create(rewriter, loc);
-      auto ctaIdInGroup = b.and_(ctaId, b.i32_val(maskCGABroadcast));
-      pred = b.and_(pred, b.icmp_eq(ctaIdInGroup, b.i32_val(0)));
-    }
-
-    ::mlir::triton::PTXBuilder ptxBuilder;
-    const std::string ptx =
+    Value basePred = b.icmp_eq(id, b.i32_val(0));
+    basePred = b.and_(basePred, adaptor.getPred());
+    auto leaderCTAPred =
+        LLVM::NVIDIA::getLeaderCTAPredicate(loc, rewriter, barrierTy);
+    bool crossCluster = leaderCTAPred.has_value();
+    Value leaderPred =
+        leaderCTAPred ? b.and_(basePred, *leaderCTAPred) : basePred;
+    Value leaderBarrierPtr = LLVM::NVIDIA::getLeaderAddress(
+        loc, rewriter, smemObj.getBase(), barrierTy);
+
+    ::mlir::triton::PTXBuilder expectPtxBuilder;
+    const std::string expectPtx =
         "@$0 mbarrier.arrive.expect_tx.shared::cta.b64 _, [$1], " +
         std::to_string(expectedBytes) + ";";
-    auto &barSyncOp = *ptxBuilder.create(ptx);
-    barSyncOp({ptxBuilder.newOperand(pred, "b"),
-               ptxBuilder.newOperand(smemObj.getBase(), "r")},
-              /*onlyAttachMLIRArgs=*/true);
+    auto &expectOp = *expectPtxBuilder.create(expectPtx);
+    expectOp({expectPtxBuilder.newOperand(leaderPred, "b"),
+              expectPtxBuilder.newOperand(leaderBarrierPtr, "r")},
+             /*onlyAttachMLIRArgs=*/true);
     auto voidTy = void_ty(op->getContext());
-    ptxBuilder.launch(rewriter, loc, voidTy);
+    expectPtxBuilder.launch(rewriter, loc, voidTy);
+
+    if (crossCluster) {
+      // Non-leader CTAs still contribute one arrival to the lead CTA barrier.
+      auto nonLeaderPred = b.and_(basePred, b.xor_(leaderPred, b.true_val()));
+      ::mlir::triton::PTXBuilder arrivePtxBuilder;
+      const std::string arrivePtx =
+          "@$0 mbarrier.arrive.shared::cluster.b64 _, [$1], 1;";
+      auto &arriveOp = *arrivePtxBuilder.create(arrivePtx);
+      arriveOp({arrivePtxBuilder.newOperand(nonLeaderPred, "b"),
+                arrivePtxBuilder.newOperand(leaderBarrierPtr, "r")},
+               /*onlyAttachMLIRArgs=*/true);
+      arrivePtxBuilder.launch(rewriter, loc, voidTy);
+    }
+
     rewriter.eraseOp(op);
     return success();
   }
@@ -238,19 +258,9 @@ struct WaitBarrierOpConversion
     auto loc = op.getLoc();
     auto b = TritonLLVMOpBuilder(loc, rewriter);
     auto pred = adaptor.getPred();
-
-    auto kBlock = StringAttr::get(ctx, "block");
-    auto maskCGABroadcast =
-        toLinearLayout(barrierTy).getFreeVariableMasks().lookup(kBlock);
-    if (maskCGABroadcast) {
-      // If several CTAs cast to the same barrier, as when we do a TMA into a
-      // tcgen05.mma 2CTA, we send all the signals to the lead CTA, so even if
-      // this barrier is waiting for zero bytes, no one will arrive on it. As
-      // such, we predicate it out
-      auto ctaId = nvgpu::ClusterCTAIdOp::create(rewriter, loc);
-      auto ctaIdInGroup = b.and_(ctaId, b.i32_val(maskCGABroadcast));
-      pred = b.and_(pred, b.icmp_eq(ctaIdInGroup, b.i32_val(0)));
-    }
+    if (auto leaderPred =
+            LLVM::NVIDIA::getLeaderCTAPredicate(loc, rewriter, barrierTy))
+      pred = b.and_(pred, *leaderPred);
 
     bool predicated = pred && !matchPattern(pred, m_NonZero());
     std::string ptx;
@@ -323,29 +333,48 @@ struct ArriveBarrierOpConversion
   LogicalResult
   matchAndRewrite(triton::nvidia_gpu::ArriveBarrierOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    auto b = TritonLLVMOpBuilder(loc, rewriter);
+    auto barrierTy = op.getAlloc().getType();
+    auto smemObj = LLVM::getSharedMemoryObjectFromStruct(
+        loc, adaptor.getAlloc(),
+        typeConverter->convertType(barrierTy.getElementType()), rewriter);
+
+    // Arrive has block-level semantics, so we must synchronize
+    // Technically, this should be MemBar's job but it can include TMEM
+    // accesses which doesn't have a MemBar equivalent :/
+    ttg::BarrierOp::create(rewriter, loc, ttg::AddrSpace::Local);
+
+    Value id = getThreadId(rewriter, loc);
+    Value pred = b.icmp_eq(id, b.i32_val(0));
+    if (op.getPred())
+      pred = b.and_(pred, adaptor.getPred());
+
+    bool isCrossCluster =
+        LLVM::NVIDIA::getLeaderCTAPredicate(loc, rewriter, barrierTy)
+            .has_value();
+
+    Value barrierPtr = LLVM::NVIDIA::getLeaderAddress(
+        loc, rewriter, smemObj.getBase(), barrierTy);
     // TODO: Add phase result as needed.
     std::stringstream ptxAsm;
-    ptxAsm << "@$0 mbarrier.arrive.shared::cta.b64 _, [$1]";
+    ptxAsm << "@$0 mbarrier.arrive."
+           << (isCrossCluster ? "shared::cluster" : "shared::cta")
+           << ".b64 _, [$1]";
     if (op.getCount() > 1) {
       ptxAsm << ", " << op.getCount();
     }
     ptxAsm << ";";
 
-    TritonLLVMOpBuilder b(op.getLoc(), rewriter);
-    Value id = getThreadId(rewriter, op.getLoc());
-    Value pred = b.icmp_eq(id, b.i32_val(0));
-    if (op.getPred())
-      pred = b.and_(pred, adaptor.getPred());
-
     PTXBuilder ptxBuilder;
     SmallVector<PTXBuilder::Operand *, 2> operands = {
         ptxBuilder.newOperand(pred, "b"),
-        ptxBuilder.newOperand(adaptor.getAlloc(), "r")};
+        ptxBuilder.newOperand(barrierPtr, "r")};
 
     auto arriveOp = *ptxBuilder.create(ptxAsm.str());
     arriveOp(operands, /*onlyAttachMLIRArgs=*/true);
     auto voidTy = void_ty(getContext());
-    ptxBuilder.launch(rewriter, op.getLoc(), voidTy);
+    ptxBuilder.launch(rewriter, loc, voidTy);
 
     rewriter.eraseOp(op);
     return success();

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp

@@ -1379,14 +1379,8 @@ struct AsyncTMACopyGlobalToLocalOpConversion
     // out)
     bool clusterBarrier = barrierMask & ~maskCGABroadcast;
     if (clusterBarrier) {
-      // This part is to support TMA into tcgen05.mma 2CTA mostly, i.e.,
-      // barrierMask == 1
-      // Mask with ones on the bits where the CTA broadcasts.
-      // This is a trick from cutlass to implement a faster `mapa`.
-      uint32_t fullMask = ~(barrierMask << 24);
-      Value barrierInt = b.ptrtoint(i32_ty, barrierPtr);
-      barrierInt = b.and_(barrierInt, b.i32_val(fullMask));
-      barrierPtr = b.inttoptr(barrierPtr.getType(), barrierInt);
+      barrierPtr =
+          LLVM::NVIDIA::getLeaderAddress(loc, rewriter, barrierPtr, barrierTy);
     }
 
     // Don't set cta_group::1 as it doesn't exist pre-Blackwell

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/Utility.cpp

@@ -155,6 +155,39 @@ Value createTMAMulticastMask(Location loc, ConversionPatternRewriter &rewriter,
   return b.shl(b.i32_val(pattern), base);
 }
 
+static uint32_t getCGABroadcastMask(mlir::triton::gpu::MemDescType barrierTy) {
+  auto kBlock = StringAttr::get(barrierTy.getContext(), "block");
+  return toLinearLayout(barrierTy).getFreeVariableMasks().lookup(kBlock);
+}
+
+std::optional<Value>
+getLeaderCTAPredicate(Location loc, ConversionPatternRewriter &rewriter,
+                      mlir::triton::gpu::MemDescType barrierTy) {
+  auto b = TritonLLVMOpBuilder(loc, rewriter);
+  uint32_t maskCGABroadcast = getCGABroadcastMask(barrierTy);
+  if (!maskCGABroadcast)
+    return std::nullopt;
+
+  Value ctaId = nvgpu::ClusterCTAIdOp::create(rewriter, loc);
+  Value ctaIdInGroup = b.and_(ctaId, b.i32_val(maskCGABroadcast));
+  return std::optional<Value>(b.icmp_eq(ctaIdInGroup, b.i32_val(0)));
+}
+
+Value getLeaderAddress(Location loc, ConversionPatternRewriter &rewriter,
+                       Value barrierPtr,
+                       mlir::triton::gpu::MemDescType barrierTy) {
+  uint32_t barrierMask = getCGABroadcastMask(barrierTy);
+  if (!barrierMask)
+    return barrierPtr;
+
+  // Trick from cutlass to implement a faster `mapa` via a single and
+  auto b = TritonLLVMOpBuilder(loc, rewriter);
+  uint32_t fullMask = ~(barrierMask << 24);
+  Value barrierInt = b.ptrtoint(i32_ty, barrierPtr);
+  barrierInt = b.and_(barrierInt, b.i32_val(fullMask));
+  return b.inttoptr(barrierPtr.getType(), barrierInt);
+}
+
 LogicalResult lowerLdStMatrix(
     Location loc, LinearLayout cvt, bool transpose,
     SmallVector<Value> &vals, // Input for stmatrix, output for ldmatrix

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/Utility.h

@@ -2,6 +2,7 @@
 #define TRITON_CONVERSION_TRITONNVIDIAGPU_TO_LLVM_UTILITY_H
 
 #include <cstdint>
+#include <optional>
 
 #include "nvidia/include/TritonNVIDIAGPUToLLVM/PTXAsmFormat.h"
 
@@ -63,6 +64,18 @@ LogicalResult lowerLdStMatrix(
 // group
 Value createTMAMulticastMask(Location loc, ConversionPatternRewriter &rewriter,
                              uint16_t broadcastBits);
+
+// Returns the lead CTA predicate for this barrier layout when lowering through
+// cluster scope. Returns std::nullopt for CTA-local lowering.
+std::optional<Value>
+getLeaderCTAPredicate(Location loc, ConversionPatternRewriter &rewriter,
+                      mlir::triton::gpu::MemDescType barrierTy);
+
+// Returns the lead CTA barrier address for this layout. If there is no
+// cross-cluster lowering, returns barrierPtr unchanged.
+Value getLeaderAddress(Location loc, ConversionPatternRewriter &rewriter,
+                       Value barrierPtr,
+                       mlir::triton::gpu::MemDescType barrierTy);
 } // namespace NVIDIA
 } // namespace LLVM