2CTA Block Scale MMA with tcgen05.cp (#9460)

# 2CTA Block Scale MMA with tcgen05.cp

* **New Features**
* Functional 2-CTA Block Scale MMA with tcgen05.cp: full path (TMA → cp
→ MMA → commit) for two CTAs

* **Bug Fixes**
* TMA barrier (2CTA hang): arrive on lead CTA barrier and use .dst
shared::cluster
* TCGen05.cp (2CTA hang): use Lead CTA predicate for tcgen05.cp
instruction
  * Scaled MMA 2CTA: per-CTA M/N; double M in the scale descriptor
  * TCGen05.cp: only copy per-CTA scales (block=0)

* **Tests**
* test_mma_scaled_tcgen05_copy: 48 cases, parameters: num_ctas,
multicast, block size, dtype
* tritongpu_to_llvm_blackwell.mlir: tmem_copy cta_group::2, lead-CTA
predicate

# New contributor declaration
- [x] I am not making a trivial change, such as fixing a typo in a
comment.

- [x] I have written a PR description following these
  [rules](https://cbea.ms/git-commit/#why-not-how).

- [x] I have run `pre-commit run --from-ref origin/main --to-ref HEAD`.

- Select one of the following.
  - [x] I have added tests.
    - `/test` for `lit` tests
    - `/unittest` for C++ tests
    - `/python/test` for end-to-end tests
  - [ ] This PR does not need a test because `FILL THIS IN`.

- Select one of the following.
  - [ ] I have not added any `lit` tests.
- [x] The `lit` tests I have added follow these [best
practices](https://mlir.llvm.org/getting_started/TestingGuide/#filecheck-best-practices),
including the "tests should be minimal" section. (Usually running Python
code
    and using the instructions it generates is not minimal.)

Author: adstraw

include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOpInterfaces.td

@@ -62,6 +62,9 @@ def MMAv5OpInterface : OpInterface<"MMAv5OpInterface"> {
                     "void",
                     "setIsAsync",
                     (ins "bool":$isAsync)>,
+    InterfaceMethod<"Return true if this MMA op uses two CTAs.",
+                    "bool",
+                    "getTwoCtas">,
     InterfaceMethod<"Return true if this MMA op executes asynchronously.",
                     "bool",
                     "isAsync">

include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td

@@ -633,7 +633,9 @@ def TTNG_TCGen5MMAScaledOp : TTNG_Op<"tc_gen5_mma_scaled", [
 
   let description = [{
     $d += matrix_multiply(scale($lhs, $lhs_scale), scale(rlhs, $rhs_scale))
-    if is_async is false, the op executes synchronously. The barrier operands must not be present in that case.
+    If $two_ctas is set the op will execute a matmul across two contiguous CTAs, it will read the data distributed across the two CTAs
+    and synchronize both CTAs if the op is synchronous.
+    If is_async is false, the op executes synchronously. The barrier operands must not be present in that case.
     Otherwise, if a barrier is given, the op will trigger a commit/arrive on it.
     The result will be safe to read after a barrier wait.
 
@@ -655,6 +657,7 @@ def TTNG_TCGen5MMAScaledOp : TTNG_Op<"tc_gen5_mma_scaled", [
     I1:$pred,
     Variadic<TTG_MemDescType>:$barriers,
     Variadic<I1>:$barrier_preds,
+    UnitAttr:$two_ctas,
     UnitAttr:$is_async
   );
   let results = (outs Optional<TTG_AsyncToken>:$token);
@@ -676,6 +679,7 @@ def TTNG_TCGen5MMAScaledOp : TTNG_Op<"tc_gen5_mma_scaled", [
       "::mlir::Value":$useD, "::mlir::Value":$pred,
       CArg<"::mlir::ValueRange", "{}">:$barriers,
       CArg<"::mlir::ValueRange", "{}">:$barrier_preds,
+      CArg<"bool", "false">:$two_ctas,
       CArg<"bool", "false">:$is_async)>
   ];
 

lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp

@@ -894,15 +894,17 @@ void TCGen5MMAScaledOp::build(OpBuilder &builder, OperationState &state,
                               Value accDep, Value aScale, Value bScale,
                               ScaleDotElemType aType, ScaleDotElemType bType,
                               Value useD, Value pred, ValueRange barriers,
-                              ValueRange barrierPreds, bool isAsync) {
+                              ValueRange barrierPreds, bool twoCTAs,
+                              bool isAsync) {
   MLIRContext *ctx = builder.getContext();
   if (!barriers.empty()) {
     isAsync = true;
   }
   build(builder, state, token, a, b, d, accDep, aScale, bScale,
         ScaleDotElemTypeAttr::get(ctx, aType),
         ScaleDotElemTypeAttr::get(ctx, bType), useD, pred, barriers,
-        barrierPreds, isAsync ? builder.getUnitAttr() : UnitAttr());
+        barrierPreds, twoCTAs ? builder.getUnitAttr() : UnitAttr(),
+        isAsync ? builder.getUnitAttr() : UnitAttr());
 }
 
 bool TCGen5MMAScaledOp::isAsync() { return getIsAsync(); }
@@ -1065,11 +1067,6 @@ LogicalResult TMEMCopyOp::verify() {
                        "representable in a matrix descriptor.");
   }
 
-  auto mod = getOperation()->getParentOfType<ModuleOp>();
-  unsigned numCTAs = triton::gpu::TritonGPUDialect::getNumCTAs(mod);
-  if (numCTAs != 1)
-    return emitOpError("NYI: Only one CTA is supported for now.");
-
   // Fp4 we could lift if we needed
   auto nvmmaEnc =
       dyn_cast<triton::gpu::NVMMASharedEncodingAttr>(srcTy.getEncoding());

lib/Dialect/TritonNvidiaGPU/Transforms/CheckMatmulTwoCTAs.cpp

@@ -28,7 +28,8 @@ class TritonNvidiaGPUCheckMatmulTwoCTAPass
     Operation *firstMatmul = nullptr;
     bool firstTwoCTA = false;
 
-    WalkResult result = mod.walk([&](ttng::TCGen5MMAOp op) {
+    // Walk all MMAv5 ops using the interface
+    WalkResult result = mod.walk([&](ttng::MMAv5OpInterface op) -> WalkResult {
       bool currentTwoCTA = op.getTwoCtas();
       if (!firstMatmul) {
         firstMatmul = op;

lib/Dialect/TritonNvidiaGPU/Transforms/ClusterBarrierInsertion.cpp

@@ -40,10 +40,7 @@ static bool isDistributedMultiCTAOp(Operation *op, bool isRead) {
   if (auto mma = dyn_cast<ttng::TCGen5MMAOp>(op)) {
     return mma.getTwoCtas();
   } else if (auto mmaScaled = dyn_cast<ttng::TCGen5MMAScaledOp>(op)) {
-    // TODO: Change when we support scaled MMA with 2CTAs
-    assert(!ttng::getModuleTwoCTAs(op->getParentOfType<ModuleOp>()) &&
-           "Scaled MMA with 2CTAs not supported");
-    return false;
+    return mmaScaled.getTwoCtas();
   } else if (auto tma = dyn_cast<ttng::AsyncTMACopyGlobalToLocalOp>(op)) {
     return tma.getMulticast();
   }

python/src/gluon_ir.cc

@@ -904,13 +904,13 @@ void init_gluon_ir(py::module &&m) {
            [](GluonOpBuilder &self, Value a, Value b, Value acc, Value aScale,
               Value bScale, tt::ScaleDotElemType aType,
               tt::ScaleDotElemType bType, Value useAcc, Value pred,
-              std::vector<Value> &mbarriers,
-              std::vector<Value> &mbarrier_preds) {
+              std::vector<Value> &mbarriers, std::vector<Value> &mbarrier_preds,
+              bool two_ctas) {
              Value accDep;
              auto tokType = self.getBuilder().getType<ttg::AsyncTokenType>();
              self.create<ttng::TCGen5MMAScaledOp>(
                  tokType, a, b, acc, accDep, aScale, bScale, aType, bType,
-                 useAcc, pred, mbarriers, mbarrier_preds);
+                 useAcc, pred, mbarriers, mbarrier_preds, two_ctas);
            })
       .def("create_tcgen05_commit",
            [](GluonOpBuilder &self, Value &barrier, Value &pred,