[BACKEND] Set 2CTA mode as a global flag (#8653)

We do so by looking at the flags of the `tcgen05.mma` dots and we make
sure they all agree. Once we support this mode in `dot_scaled` we'll
check these as well.

Author: lezcano

include/triton/Dialect/TritonNvidiaGPU/IR/Dialect.h

@@ -26,6 +26,7 @@
 
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
@@ -51,6 +52,17 @@ LogicalResult verifyMMAv5Op(Operation *op);
 
 namespace mlir::triton::nvidia_gpu {
 
+constexpr static char AttrTwoCTAsName[] = "ttng.two-ctas";
+
+inline bool getModuleTwoCTAs(ModuleOp mod) {
+  auto attr = mod->getAttrOfType<BoolAttr>(AttrTwoCTAsName);
+  return attr ? attr.getValue() : false;
+}
+
+inline bool getModuleTwoCTAs(Operation *op) {
+  return getModuleTwoCTAs(op->getParentOfType<ModuleOp>());
+}
+
 struct TensorMemory : public SideEffects::Resource::Base<TensorMemory> {
   StringRef getName() final { return "<TensorMemory>"; }
 };

include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.td

@@ -174,4 +174,14 @@ def TritonNvidiaGPURemoveTMEMTokensPass : Pass<"triton-nvidia-gpu-remove-tmem-to
   }];
 }
 
+def TritonNvidiaGPUCheckMatmulTwoCTAPass : Pass<"triton-nvidia-check-matmul-two-cta", "mlir::ModuleOp"> {
+  let summary = "Verify consistent two_ctas usage across matmuls";
+
+  let description = [{
+    Inspect all matmul operations and ensure they agree on the `two_ctas`
+    setting. Propagate the chosen value to the module so later lowering steps
+    can access it. Compilation fails if mixed configurations are detected.
+  }];
+}
+
 #endif

lib/Dialect/TritonNvidiaGPU/Transforms/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_triton_library(TritonNvidiaGPUTransforms
+  CheckMatmulTwoCTAs.cpp
   FenceInsertion.cpp
   InterleaveTMem.cpp
   MMALowering.cpp

lib/Dialect/TritonNvidiaGPU/Transforms/CheckMatmulTwoCTAs.cpp

@@ -0,0 +1,63 @@
+#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
+#include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h"
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Visitors.h"
+
+namespace ttng = mlir::triton::nvidia_gpu;
+
+namespace mlir::triton::nvidia_gpu {
+
+#define GEN_PASS_DEF_TRITONNVIDIAGPUCHECKMATMULTWOCTAPASS
+#include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h.inc"
+
+namespace {
+
+class TritonNvidiaGPUCheckMatmulTwoCTAPass
+    : public impl::TritonNvidiaGPUCheckMatmulTwoCTAPassBase<
+          TritonNvidiaGPUCheckMatmulTwoCTAPass> {
+public:
+  using impl::TritonNvidiaGPUCheckMatmulTwoCTAPassBase<
+      TritonNvidiaGPUCheckMatmulTwoCTAPass>::
+      TritonNvidiaGPUCheckMatmulTwoCTAPassBase;
+
+  void runOnOperation() override {
+    ModuleOp mod = getOperation();
+    Operation *firstMatmul = nullptr;
+    bool firstTwoCTA = false;
+
+    WalkResult result = mod.walk([&](ttng::TCGen5MMAOp op) {
+      bool currentTwoCTA = op.getTwoCtas();
+      if (!firstMatmul) {
+        firstMatmul = op;
+        firstTwoCTA = currentTwoCTA;
+        return WalkResult::advance();
+      }
+      if (currentTwoCTA != firstTwoCTA) {
+        auto diag = op.emitError()
+                    << "inconsistent two_ctas setting across matmuls; "
+                       "expected all matmuls to "
+                    << (firstTwoCTA ? "enable" : "disable") << " two_ctas.";
+        diag.attachNote(firstMatmul->getLoc())
+            << "first matmul here has two_ctas="
+            << (firstTwoCTA ? "true" : "false") << ".";
+        return WalkResult::interrupt();
+      }
+      return WalkResult::advance();
+    });
+
+    if (result.wasInterrupted()) {
+      signalPassFailure();
+      return;
+    }
+
+    bool twoCTAValue = firstMatmul ? firstTwoCTA : false;
+    mod->setAttr(AttrTwoCTAsName, BoolAttr::get(mod.getContext(), twoCTAValue));
+  }
+};
+
+} // namespace
+
+} // namespace mlir::triton::nvidia_gpu

test/Conversion/tritongpu_to_llvm_blackwell.mlir

@@ -270,7 +270,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
 #shared1 = #ttg.nvmma_shared<{swizzlingByteWidth = 64, transposed = false, elementBitWidth = 16, CTAsPerCGA = [1, 2], CTASplitNum = [1, 2], CTAOrder = [1, 0]}>
 #shared2 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0], CTAsPerCGA = [2], CTASplitNum = [1], CTAOrder = [0]}>
 #tmem = #ttng.tensor_memory_encoding<blockM = 128, blockN = 128, colStride = 1, CTASplitM = 2>
-module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32} {
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32, "ttng.two-ctas" = true} {
   // CHECK-LABEL: @tc_gen5_mma_2ctas
   tt.func @tc_gen5_mma_2ctas(%a: !ttg.memdesc<256x32xf16, #shared, #ttg.shared_memory>,
                        %b: !ttg.memdesc<32x128xf16, #shared1, #ttg.shared_memory>,

third_party/nvidia/backend/compiler.py

@@ -348,6 +348,7 @@ def make_llir(self, src, metadata, options, capability):
         passes.gluon.add_inliner(pm)
         nvidia.passes.ttgpuir.add_allocate_shared_memory_nv(pm, capability, ptx_version)
         nvidia.passes.ttnvgpuir.add_allocate_tensor_memory(pm)
+        nvidia.passes.ttnvgpuir.add_check_matmul_two_cta(pm)
         if knobs.compilation.instrumentation_mode == "consan":
             # Call ConcurrencySanitizerPass here, before allocating global scratch memory but after allocating tensor and shared
             passes.ttgpuir.add_concurrency_sanitizer(pm)

third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp

@@ -595,10 +595,7 @@ static Value initTensorMemory(LLVM::LLVMFuncOp func) {
     return LLVM::UndefOp::create(rewriter, loc, ptr_ty(ctx, 6));
   }
 
-  int numCTAs = triton::gpu::TritonGPUDialect::getNumCTAs(mod);
-  // Assume that 2CTAs is used if we have two CTAs this is pessimistic but
-  // should be fine for now.
-  bool useTwoCTAs = numCTAs == 2;
+  bool useTwoCTAs = mlir::triton::nvidia_gpu::getModuleTwoCTAs(mod);
   // This code is only executed by the default warp group.
   Value threadId = NVVM::ThreadIdXOp::create(rewriter, loc, i32_ty);
   Value pred = b.icmp_ult(threadId, b.i32_val(32));

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv5.cpp

@@ -264,17 +264,17 @@ static void createScaledGen5MMA(ConversionPatternRewriter &rewriter,
                                 MemDescOperand a, Value b, MemDescOperand d,
                                 Value scaleA, Value scaleB, Value pred,
                                 Value instDescriptor, Value useInitAcc,
-                                bool aInTmem, mxfpKind mxfpInstKind) {
+                                bool aInTmem, mxfpKind mxfpInstKind,
+                                bool twoCTAs) {
   PTXBuilder ptxBuilder;
-  std::string opcode;
+  std::string opcode =
+      "tcgen05.mma.cta_group::" + std::to_string(twoCTAs ? 2 : 1) + ".kind::";
   if (mxfpInstKind == mxfpKind::mxf8f6f4) {
-    opcode =
-        "tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X";
+    opcode += "mxf8f6f4.block_scale.scale_vec::1X";
   } else if (mxfpInstKind == mxfpKind::mxf4) {
-    opcode = "tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X";
+    opcode += "mxf4.block_scale.scale_vec::2X";
   } else if (mxfpInstKind == mxfpKind::mxf4nvf4) {
-    opcode =
-        "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X";
+    opcode += "mxf4nvf4.block_scale.scale_vec::4X";
   } else {
     assert(0 && "Unsupported mxfp kind.");
   }
@@ -312,7 +312,9 @@ static void createMMACommit(ConversionPatternRewriter &rewriter, Location loc,
              "tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::"
              "cluster.multicast::cluster.b64 [$1], $2;";
   } else {
-    opcode = "@$0 tcgen05.commit.cta_group::1.mbarrier::arrive::one.b64 [$1];";
+    opcode =
+        "@$0 tcgen05.commit.cta_group::" + std::to_string(twoCTAs ? 2 : 1) +
+        ".mbarrier::arrive::one.b64 [$1];";
   }
   auto &barrierOp = *ptxBuilder.create(opcode);
   barrierOp(ptxOperands, /*onlyAttachMLIRArgs=*/true);
@@ -486,7 +488,8 @@ void convertDot(const LLVMTypeConverter &typeConverter,
   MemDescType bTensorTy = op.getB().getType();
   MemDescType dTensorTy = op.getD().getType();
   auto dLayout = cast<ttng::TensorMemoryEncodingAttr>(dTensorTy.getEncoding());
-  bool twoCTAs = op.getTwoCtas();
+  bool twoCTAs = ttng::getModuleTwoCTAs(op);
+  assert(twoCTAs == op.getTwoCtas());
 
   DotConversion dot;
 
@@ -595,6 +598,7 @@ void convertScaledDot(const LLVMTypeConverter &typeConverter,
   Value baseD = tb.ptrtoint(i32_ty, adaptor.getD());
   Value baseScaleA = tb.ptrtoint(i32_ty, adaptor.getAScale());
   Value baseScaleB = tb.ptrtoint(i32_ty, adaptor.getBScale());
+  bool twoCTAs = ttng::getModuleTwoCTAs(op);
 
   int numRows = 128;
   int colSizeInBits = 32;
@@ -634,14 +638,13 @@ void convertScaledDot(const LLVMTypeConverter &typeConverter,
         subWordIdx, subWordIdx, mxfpInstKind);
     createScaledGen5MMA(rewriter, loc, op, a, b, accAddress, scaleA, scaleB,
                         pred, instDescriptor, useInitAcc, desc.aInTmem,
-                        mxfpInstKind);
+                        mxfpInstKind, twoCTAs);
   };
 
   convertDotImpl(typeConverter, rewriter, loc, op.getA(), op.getB(),
                  adaptor.getA(), adaptor.getB(), dTensorTy, adaptor.getUseD(),
                  adaptor.getPred(), adaptor.getBarriers(),
-                 adaptor.getBarrierPreds(), /*twoCTAs=*/false, opKindIsMXFP4,
-                 dot);
+                 adaptor.getBarrierPreds(), twoCTAs, opKindIsMXFP4, dot);
 }
 
 //===----------------------------------------------------------------------===//
@@ -699,7 +702,7 @@ struct TCGen5CommitOpConversion
       pred = b.and_(adaptor.getPred(), pred);
 
     createMMACommit(rewriter, op.getLoc(), smemObj.getBase(), pred,
-                    op.getTwoCtas());
+                    ttng::getModuleTwoCTAs(op));
     rewriter.eraseOp(op);
     return success();
   }

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TensorMemoryToLLVM.cpp

@@ -524,18 +524,20 @@ struct TensorMemoryAllocOpConversion
 };
 
 static void createCommit(ConversionPatternRewriter &rewriter, Location loc,
-                         Value barrier, Value pred) {
+                         Value barrier, Value pred, bool twoCTAs) {
   PTXBuilder ptxBuilder;
   auto *barrierOperand = ptxBuilder.newAddrOperand(barrier, "r");
-  std::string opcode = "tcgen05.commit.cta_group::1.mbarrier::arrive::one.b64";
+  std::string opcode =
+      "tcgen05.commit.cta_group::" + std::to_string(twoCTAs ? 2 : 1) +
+      ".mbarrier::arrive::one.b64";
   auto &barrierOp = *ptxBuilder.create(opcode);
   barrierOp(barrierOperand).predicate(pred);
   ptxBuilder.launch(rewriter, loc, void_ty(rewriter.getContext()));
 }
 
 static void createTcgen05Cp(ConversionPatternRewriter &rewriter, Location loc,
                             Value tmem_address, Value src_desc, Value pred,
-                            TMemCopyAtom atom) {
+                            TMemCopyAtom atom, bool twoCTAs) {
   PTXBuilder ptxBuilder;
   auto dst = ptxBuilder.newAddrOperand(tmem_address, "r");
   auto src = ptxBuilder.newOperand(src_desc, "l");
@@ -547,9 +549,9 @@ static void createTcgen05Cp(ConversionPatternRewriter &rewriter, Location loc,
   } else if (atom.multicast == 3) {
     warp = ".warpx4";
   }
-  std::string opcode = "tcgen05.cp.cta_group::1" + warp + "." +
-                       std::to_string(atom.nRow) + "x" +
-                       std::to_string(atom.bCol) + "b";
+  std::string opcode =
+      "tcgen05.cp.cta_group::" + std::to_string(twoCTAs ? 2 : 1) + warp + "." +
+      std::to_string(atom.nRow) + "x" + std::to_string(atom.bCol) + "b";
   auto &op = *ptxBuilder.create(opcode);
   op({dst, src}).predicate(pred);
   ptxBuilder.launch(rewriter, loc, void_ty(rewriter.getContext()));
@@ -592,6 +594,7 @@ static void copySharedToTmem(ConversionPatternRewriter &rewriter, Location loc,
   auto loader = DotOpMmaSmemLoader::build(loc, rewriter, cvtWarp, bitwidth,
                                           smemBase, instrShape, 0, 5);
   assert(!loader.getDescriptor().transposed);
+  bool twoCTAs = getModuleTwoCTAs(op);
   // Check correct lbo/sbo along the multicast
   auto strideRow = cvt.getBasis(kRow, llvm::Log2_32(8), kOffset);
   if ((atom.multicast & 1) == 0) {
@@ -608,7 +611,7 @@ static void copySharedToTmem(ConversionPatternRewriter &rewriter, Location loc,
     auto tmemAddr =
         b.or_(b.ptrtoint(i32_ty, baseDst), b.i32_val(col * bitwidth / 32),
               /*disjoint=*/true);
-    createTcgen05Cp(rewriter, loc, tmemAddr, desc, pred, atom);
+    createTcgen05Cp(rewriter, loc, tmemAddr, desc, pred, atom, twoCTAs);
   }
 }
 
@@ -622,13 +625,14 @@ struct TensorMemoryCopyOpConversion
     assert(lookupNumCTAs(rewriter) == 1 && "NYI");
     Location loc = op->getLoc();
     Value pred = LLVM::NVIDIA::createElectPredicateWarp0(loc, rewriter);
+    bool twoCTAs = getModuleTwoCTAs(op);
     copySharedToTmem(rewriter, loc, typeConverter, op, adaptor.getSrc(),
                      adaptor.getDst(), pred);
 
     if (op.getBarrier()) {
       auto barrier = LLVM::getSharedMemoryObjectFromStruct(
           op.getLoc(), adaptor.getBarrier(), i64_ty, rewriter);
-      createCommit(rewriter, loc, barrier.getBase(), pred);
+      createCommit(rewriter, loc, barrier.getBase(), pred, twoCTAs);
     }
 
     rewriter.eraseOp(op);

third_party/nvidia/triton_nvidia.cc

@@ -60,6 +60,8 @@ void init_triton_nvidia_passes_ttnvgpuir(py::module &&m) {
                      ttng::createTritonNvidiaGPUPromoteLHSToTMemPass);
   ADD_PASS_WRAPPER_0("add_remove_tmem_tokens",
                      ttng::createTritonNvidiaGPURemoveTMEMTokensPass);
+  ADD_PASS_WRAPPER_0("add_check_matmul_two_cta",
+                     ttng::createTritonNvidiaGPUCheckMatmulTwoCTAPass);
   ADD_PASS_WRAPPER_0("add_nvgpu_to_llvm",
                      mlir::triton::createConvertNVGPUToLLVM);
   ADD_PASS_WRAPPER_0("add_warp_specialize_to_llvm",