[Blackwell] Add rewrite pattern to merge explicit commit ops into MMAv5 (triton-lang#8026)

An explicit commit op and the barrier arguments to MMAv5 express the
same synchronization patterns. In practice, however, emitting
`tcgen05.mma` and `tcgen05.commit` together into the same basic block,
as implemented in `MMAv5.cpp` for MMAv5 ops with the barrier arguments,
can yield slightly more efficient SASS sometimes.

This PR adds a simple rewrite pattern that folds explicit commit ops
into a preceding MMAv5 op, so that use cases involving explicit commit
can also benefit from the low-level optimization. More discussions that
motivated this work are in
triton-lang#7984.

Author: masahi

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

@@ -53,6 +53,9 @@ def MMAv5OpInterface : OpInterface<"MMAv5OpInterface"> {
                     "void",
                     "setIsAsync",
                     (ins "bool":$isAsync)>,
+    InterfaceMethod<"Return true if this MMA op executes asynchronously.",
+                    "bool",
+                    "isAsync">
   ];
 
   let verify = [{

lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp

@@ -400,6 +400,8 @@ void TCGen5MMAOp::build(OpBuilder &builder, OperationState &state, Type token,
         useTwoCTAs ? builder.getUnitAttr() : UnitAttr());
 }
 
+bool TCGen5MMAOp::isAsync() { return getIsAsync(); }
+
 // -- TCGen5MMAScaledOp --
 LogicalResult TCGen5MMAScaledOp::verify() {
   if (!getIsAsync() && !getBarriers().empty()) {
@@ -573,6 +575,8 @@ void TCGen5MMAScaledOp::build(OpBuilder &builder, OperationState &state,
         barrierPreds, isAsync ? builder.getUnitAttr() : UnitAttr());
 }
 
+bool TCGen5MMAScaledOp::isAsync() { return getIsAsync(); }
+
 // -- TMEMStoreOp --
 static LogicalResult verifyTMEMOperand(Operation *op, RankedTensorType type,
                                        MemDescType memdesc, StringRef regName) {

lib/Dialect/TritonNvidiaGPU/Transforms/MMALowering.cpp

@@ -1,3 +1,4 @@
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
@@ -16,15 +17,14 @@ namespace nvidia_gpu {
 
 namespace {
 
-template <typename TCGen5MMAOpTy>
-class SyncMMALowering : public OpRewritePattern<TCGen5MMAOpTy> {
+class SyncMMALowering : public OpInterfaceRewritePattern<MMAv5OpInterface> {
 public:
-  using OpRewritePattern<TCGen5MMAOpTy>::OpRewritePattern;
+  using OpInterfaceRewritePattern<MMAv5OpInterface>::OpInterfaceRewritePattern;
 
-  LogicalResult matchAndRewrite(TCGen5MMAOpTy op,
+  LogicalResult matchAndRewrite(MMAv5OpInterface op,
                                 PatternRewriter &rewriter) const override {
     // If the op doesn't have synchronous semantic skip the pattern.
-    if (op.getIsAsync())
+    if (op.isAsync())
       return failure();
     MLIRContext *ctx = op.getContext();
     Location loc = op.getLoc();
@@ -46,7 +46,7 @@ class SyncMMALowering : public OpRewritePattern<TCGen5MMAOpTy> {
 
     rewriter.setInsertionPointAfter(op);
     Value phase = rewriter.create<arith::ConstantIntOp>(loc, 0, 32);
-    rewriter.create<WaitBarrierOp>(loc, barrierAlloc, phase, op.getPred());
+    rewriter.create<WaitBarrierOp>(loc, barrierAlloc, phase, op.getPredicate());
     rewriter.create<InvalBarrierOp>(loc, barrierAlloc);
     return success();
   }
@@ -101,6 +101,103 @@ struct TCGen5MMAScaleSharedToTmemConversion
   }
 };
 
+std::pair<SmallVector<TCGen5CommitOp>, SmallVector<Value>>
+collectCommitOpsAfter(MMAv5OpInterface mmaOp) {
+  auto isConstTrue = [](Value v) {
+    if (auto constOp = v.getDefiningOp<arith::ConstantOp>()) {
+      if (auto attr = dyn_cast<BoolAttr>(constOp.getValueAttr())) {
+        return attr.getValue();
+      }
+    }
+    return false;
+  };
+
+  SmallVector<TCGen5CommitOp> commitOps;
+  SmallVector<Value> commitPredicates;
+  auto mmaPred = mmaOp.getPredicate();
+  Operation *nextOp = mmaOp->getNextNode();
+
+  while (nextOp) {
+    if (auto commit = dyn_cast<TCGen5CommitOp>(nextOp)) {
+      // If the mma predicate is true, or mma and commit ops use the same
+      // predicate, it is safe to merge them
+      if (isConstTrue(mmaPred) || mmaPred == commit.getPred()) {
+        commitOps.push_back(commit);
+        commitPredicates.push_back(commit.getPred());
+      }
+    } else if (!isPure(nextOp)) {
+      // Only move commits across pure ops. We also bail here when encountering
+      // another MMAv5 op.
+      break;
+    }
+    nextOp = nextOp->getNextNode();
+  }
+
+  return {commitOps, commitPredicates};
+}
+
+// Return false if defining ops cannot be moved above the target op
+bool moveDefiningOpsBefore(Value val, Operation *target) {
+  SetVector<Operation *> toMove;
+
+  std::function<bool(Value)> collectOpsToMove = [&](Value val) {
+    if (auto defOp = val.getDefiningOp()) {
+      if (defOp->getBlock() == target->getBlock() &&
+          target->isBeforeInBlock(defOp)) {
+        if (!isPure(defOp)) {
+          // This defOp needs to move above the target op, but it is unsafe due
+          // to impurity.
+          return false;
+        }
+        for (Value operand : defOp->getOperands()) {
+          if (!collectOpsToMove(operand)) {
+            return false;
+          }
+        }
+        toMove.insert(defOp);
+      }
+    }
+    return true;
+  };
+
+  if (!collectOpsToMove(val)) {
+    return false;
+  }
+
+  for (Operation *op : toMove) {
+    op->moveBefore(target);
+  }
+
+  return true;
+}
+
+class MergeCommitIntoMMA : public OpInterfaceRewritePattern<MMAv5OpInterface> {
+public:
+  using OpInterfaceRewritePattern<MMAv5OpInterface>::OpInterfaceRewritePattern;
+
+  LogicalResult matchAndRewrite(MMAv5OpInterface op,
+                                PatternRewriter &rewriter) const override {
+    auto [commitOps, predicates] = collectCommitOpsAfter(op);
+    if (commitOps.size() == 0) {
+      return llvm::failure();
+    }
+    for (auto [commit, pred] : llvm::zip(commitOps, predicates)) {
+      if (!pred) {
+        pred = rewriter.create<arith::ConstantIntOp>(op.getLoc(), true, 1);
+      }
+      if (!moveDefiningOpsBefore(commit.getBarrier(), op) ||
+          !moveDefiningOpsBefore(pred, op)) {
+        // Give up merging a commit if its defining ops cannot be moved above
+        // the mma op.
+        continue;
+      }
+      op.addCompletionBarrier(commit.getBarrier(), pred);
+      rewriter.eraseOp(commit);
+    }
+    return success();
+  }
+};
+
 } // anonymous namespace
 
 class TritonNvidiaGPUMMALoweringPass
@@ -112,9 +209,9 @@ class TritonNvidiaGPUMMALoweringPass
     ModuleOp m = getOperation();
 
     mlir::RewritePatternSet patterns(context);
-    patterns
-        .add<SyncMMALowering<TCGen5MMAOp>, SyncMMALowering<TCGen5MMAScaledOp>,
-             TCGen5MMAScaleSharedToTmemConversion>(context);
+    patterns.add<SyncMMALowering, TCGen5MMAScaleSharedToTmemConversion,
+                 MergeCommitIntoMMA>(context);
+
     if (applyPatternsGreedily(m, std::move(patterns)).failed())
       signalPassFailure();
   }

test/TritonNvidiaGPU/mma_lowering.mlir

@@ -58,3 +58,72 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.targ
     tt.return
   }
 }
+
+// -----
+#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 32, transposed = false, elementBitWidth = 8}>
+#shared1 = #ttg.nvmma_shared<{swizzlingByteWidth = 32, transposed = true, elementBitWidth = 8}>
+#shared2 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0], CTAsPerCGA = [1], CTASplitNum = [1], CTAOrder = [0]}>
+#tmem = #ttng.tensor_memory_encoding<blockM = 128, blockN = 64, unpacked = true>
+#smem = #ttg.shared_memory
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+  // CHECK-LABEL: tcgen5_with_commit
+  tt.func @tcgen5_with_commit(
+    // CHECK: [[BARRIER1:%.*]]: !ttg.memdesc<1xi64, #shared
+    %barrier: !ttg.memdesc<1xi64, #shared2, #ttg.shared_memory, mutable>,
+    // CHECK: [[BARRIER_PRED:%.*]]: i1,
+    %barrierPred: i1,
+    // CHECK: [[A_SMEM:%.*]]: !ttg.memdesc<128x128xf8E5M2
+    %a: !ttg.memdesc<128x128xf8E5M2, #shared, #ttg.shared_memory>,
+    %b: !ttg.memdesc<128x256xf8E5M2, #shared1, #ttg.shared_memory>,
+    %c: !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>) {
+    %barrier2 = ttg.local_alloc : () -> !ttg.memdesc<2x1xi64, #shared2, #smem, mutable>
+    %c0_i32 = arith.constant 0 : i32
+    // CHECK: [[TRUE:%.*]] = arith.constant true
+    // CHECK: [[BARRIER_SLICE:%.*]] = ttg.memdesc_index
+    // CHECK: ttng.tc_gen5_mma {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, [[BARRIER1]][[[BARRIER_PRED]]], [[BARRIER_SLICE]][[[TRUE]]]
+    %accUse = arith.constant false
+    %pred = arith.constant true
+    ttng.tc_gen5_mma %a, %b, %c, %accUse, %pred {is_async} :
+       !ttg.memdesc<128x128xf8E5M2, #shared, #ttg.shared_memory>,
+       !ttg.memdesc<128x256xf8E5M2, #shared1, #ttg.shared_memory>,
+       !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>
+    ttng.tc_gen5_commit %barrier, %barrierPred : !ttg.memdesc<1xi64, #shared2, #ttg.shared_memory, mutable>
+    %barrier_slice = ttg.memdesc_index %barrier2[%c0_i32] : !ttg.memdesc<2x1xi64, #shared2, #smem, mutable> -> !ttg.memdesc<1xi64, #shared2, #smem, mutable, 2x1>
+    ttng.tc_gen5_commit %barrier_slice : !ttg.memdesc<1xi64, #shared2, #ttg.shared_memory, mutable, 2x1>
+
+    ttng.tc_gen5_mma %a, %b, %c, %accUse, %pred {is_async} :
+       !ttg.memdesc<128x128xf8E5M2, #shared, #ttg.shared_memory>,
+       !ttg.memdesc<128x256xf8E5M2, #shared1, #ttg.shared_memory>,
+       !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>
+
+    %random_pred = arith.cmpi eq, %barrierPred, %pred : i1
+    scf.if %random_pred {
+      ttng.tc_gen5_mma %a, %b, %c, %accUse, %pred {is_async} :
+       !ttg.memdesc<128x128xf8E5M2, #shared, #ttg.shared_memory>,
+       !ttg.memdesc<128x256xf8E5M2, #shared1, #ttg.shared_memory>,
+       !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>
+    }
+    // This commit should not be merged into any of two mma ops above
+    // CHECK: tc_gen5_commit
+    ttng.tc_gen5_commit %barrier, %barrierPred : !ttg.memdesc<1xi64, #shared2, #ttg.shared_memory, mutable>
+
+    // The mma predicate is not a constant true. The commit op should not be merged
+    // CHECK: tc_gen5_commit
+    ttng.tc_gen5_mma %a, %b, %c, %accUse, %random_pred {is_async} :
+       !ttg.memdesc<128x128xf8E5M2, #shared, #ttg.shared_memory>,
+       !ttg.memdesc<128x256xf8E5M2, #shared1, #ttg.shared_memory>,
+       !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>
+    ttng.tc_gen5_commit %barrier : !ttg.memdesc<1xi64, #shared2, #ttg.shared_memory, mutable>
+
+    // There is an impure op between mma and commit ops. Do not allow merging in such cases.
+    // CHECK: tc_gen5_commit
+    ttng.tc_gen5_mma %a, %b, %c, %accUse, %pred {is_async} :
+       !ttg.memdesc<128x128xf8E5M2, #shared, #ttg.shared_memory>,
+       !ttg.memdesc<128x256xf8E5M2, #shared1, #ttg.shared_memory>,
+       !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>
+    ttng.wait_barrier %barrier, %c0_i32 : !ttg.memdesc<1xi64, #shared2, #ttg.shared_memory, mutable>
+    ttng.tc_gen5_commit %barrier : !ttg.memdesc<1xi64, #shared2, #ttg.shared_memory, mutable>
+
+    tt.return
+  }
+}