[PIPELINE] Add support for pipelining mma loops with modified accumulator (#6393)

This change improves mmav5 pipelining so mma with RMW accumulator are
pipelined by putting the tmem load before the mma.
The approach is not entirely clean, as to avoid having dist-1
dependencies on the wait_barrier we allow only pipelining of cases where
the load and the mma are in the same stage, and we skip adding buffer
dependencies to the wait.
Second dot pipelining is still disabled by default as even with this
optimization there are some performance problems when we enable it.
With second dot pipelining disabled however, this change improves some
internal attention kernels perf by up to 30% by reducing spilling

Author: pawelszczerbuk

include/triton/Dialect/TritonGPU/Transforms/Schedule.h

@@ -28,9 +28,6 @@ void lowerLoops(ModuleOp moduleOp);
 /// Pipeline the TMA stores in the loop.
 bool pipelineTMAStores(scf::ForOp forOp);
 
-/// Simple pipelining for the MMA ops which accumulator is modified in the loop.
-scf::ForOp pipelineMMAWithScaledAcc(scf::ForOp forOp);
-
 /// This does post-processing on the pipelined loop to try to pipeline wgmma
 /// ops.
 // TODO: this should be included as part of the pipeline but currently the wgmma
@@ -75,6 +72,8 @@ class CoarseSchedule {
     }
 
     bool isBefore(iterator a, iterator b) const {
+      if (a == b)
+        return false;
       for (auto it = begin(); it != end(); ++it) {
         if (it == a)
           return true;

lib/Dialect/TritonGPU/Transforms/CMakeLists.txt

@@ -22,7 +22,6 @@ add_triton_library(TritonGPUTransforms
   Pipeliner/SoftwarePipeliner.cpp
   Pipeliner/TMAStoresPipeline.cpp
   Pipeliner/MMAv5PipelineUtility.cpp
-  Pipeliner/ModifiedAccMMAPipeline.cpp
   Pipeliner/Partition.cpp
   Pipeliner/PipeliningUtility.cpp
   Pipeliner/Schedule.cpp

lib/Dialect/TritonGPU/Transforms/HoistTMEMAlloc.cpp

@@ -229,6 +229,53 @@ class RotateTMEMStoreInLoop : public OpRewritePattern<ttng::TMEMStoreOp> {
   }
 };
 
+class RotateTMEMLoadInLoop : public OpRewritePattern<ttng::TMEMLoadOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ttng::TMEMLoadOp load,
+                                PatternRewriter &rewriter) const override {
+    scf::ForOp forOp = dyn_cast<scf::ForOp>(load->getParentOp());
+    if (!forOp) {
+      return failure();
+    }
+    if (!forOp.isDefinedOutsideOfLoop(load.getSrc())) {
+      return failure();
+    }
+    if (!load.getResult().hasOneUse()) {
+      return failure();
+    }
+    OpOperand &use = *load.getResult().getUses().begin();
+    auto yield = dyn_cast<scf::YieldOp>(use.getOwner());
+    if (!yield) {
+      return failure();
+    }
+    // Create two copies of the store: one before the loop, storing the initial
+    // value, and one before the yield, storing the value carried by the loop
+    // arg.
+    int argNo = use.getOperandNumber();
+    Value initVal = forOp.getInitArgs()[argNo];
+    rewriter.setInsertionPoint(forOp);
+    auto vTrue = rewriter.create<arith::ConstantIntOp>(load.getLoc(), 1, 1);
+    rewriter.create<ttng::TMEMStoreOp>(load.getLoc(), load.getSrc(), initVal,
+                                       vTrue);
+    auto attributes = load->getAttrDictionary();
+    rewriter.moveOpBefore(load, &forOp.getBody()->front());
+    forOp.getRegionIterArg(argNo).replaceAllUsesWith(load.getResult());
+    load->setAttrs(attributes);
+
+    // Load from the tmem after the loop, and use it instead of the loop carried
+    // value.
+    rewriter.setInsertionPointAfter(forOp);
+    auto loadAfterLoop = rewriter.create<ttng::TMEMLoadOp>(
+        load.getLoc(), load.getResult().getType(), load.getSrc());
+    forOp->getResult(argNo).replaceAllUsesWith(loadAfterLoop.getResult());
+    // Loop carried value is no longer used, short-circuit it.
+    yield.setOperand(argNo, forOp.getRegionIterArg(argNo));
+    return success();
+  }
+};
+
 ttng::TMEMAllocOp hoistTMEMAlloc(ttng::TMEMAllocOp alloc, scf::ForOp forOp) {
   OpBuilder builder(alloc);
   builder.setInsertionPoint(forOp);
@@ -300,8 +347,10 @@ struct HoistTMEMAlloc
     }
 
     mlir::RewritePatternSet patterns(&getContext());
-    patterns.add<RotateTMEMStoreInLoop, CombineTMEMLoadAndStore,
-                 CombineTMEMStoreAndSelect, SinkTMEMLoad>(&getContext());
+    patterns
+        .add<RotateTMEMStoreInLoop, RotateTMEMLoadInLoop,
+             CombineTMEMLoadAndStore, CombineTMEMStoreAndSelect, SinkTMEMLoad>(
+            &getContext());
     scf::ForOp::getCanonicalizationPatterns(patterns, &getContext());
     if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
       llvm_unreachable("Failed to hoist tmem_store");

lib/Dialect/TritonGPU/Transforms/Pipeliner/LowerLoops.cpp

@@ -118,7 +118,9 @@ Operation *getFirstUseOfPipelinedOp(SmallVector<Operation *> ops,
       auto [_firstUserStage, _firstUserCluster] = schedule[firstUser];
       if (_useStage < _firstUserStage ||
           (_useStage == _firstUserStage &&
-           schedule.clusters.isBefore(_useCluster, _firstUserCluster))) {
+           schedule.clusters.isBefore(_useCluster, _firstUserCluster)) ||
+          (_useStage == _firstUserStage && _useCluster == _firstUserCluster &&
+           topLevelUser->isBeforeInBlock(firstUser))) {
         firstUser = topLevelUser;
       }
     }
@@ -214,6 +216,8 @@ Value createIncrementModulo(BuilderT &builder, Location loc, Value counter,
 
 void replaceAllUsesDominatedBy(Operation *domOp, Value newValue, Value oldValue,
                                DominanceInfo &domInfo) {
+  if (newValue == oldValue)
+    return;
   oldValue.replaceUsesWithIf(newValue, [&](OpOperand &use) {
     return domInfo.properlyDominates(domOp, use.getOwner());
   });
@@ -232,15 +236,6 @@ static Value createAlloc(scf::ForOp &forOp, Operation *loadOp,
       loadOp->getLoc(), sharedEnc, distance);
 }
 
-template <typename BuilderT, typename... Args>
-Operation *createWithStage(BuilderT &builder, Location loc, int stage,
-                           CoarseSchedule::Cluster cluster, Args &&...args) {
-  Operation *op = builder.template create<ttg::AsyncCopyGlobalToLocalOp>(
-      loc, std::forward<Args>(args)...);
-
-  return op;
-}
-
 void createAsyncCopy(scf::ForOp forOp, tt::LoadOp loadOp, Value alloc,
                      Value insertIdx, Value extractIdx,
                      CoarseSchedule &schedule) {
@@ -874,12 +869,36 @@ std::pair<int, int> getTmemUseStageBounds(ttng::TMEMAllocOp alloc,
   return bounds;
 }
 
-void createBarrierAndWaitOps(scf::ForOp forOp, CoarseSchedule &schedule,
-                             ttng::MMAv5OpInterface mma,
-                             ttng::TMEMAllocOp alloc, Value &phase,
-                             Value &barrierIdx, int numStages) {
+// Create a predicate argument for the dist-1wait
+scf::ForOp prepLoopForDist1Wait(scf::ForOp forOp, CoarseSchedule &schedule,
+                                ttng::MMAv5OpInterface mma) {
+  OpBuilderForStage builder(forOp, schedule);
+  Location loc = mma.getLoc();
+  Value vFalse = builder.create<arith::ConstantIntOp>(loc, 0, 1);
+
+  // Create a predicate for the wait (start with false and change to true on the
+  // first mma execution)
+  scf::ForOp newForOp = replaceForOpWithNewSignature(builder, forOp, {vFalse});
+  forOp.erase();
+  forOp = newForOp;
+
+  builder.setInsertionPointAfter(mma);
+  builder.setStageCluster(schedule[mma]);
+  Value vTrue = builder.create<arith::ConstantIntOp>(loc, 1, 1);
+
+  auto yieldOp = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
+  yieldOp.getResultsMutable().append(vTrue); // predicate
+  return forOp;
+}
+
+scf::ForOp createBarrierAndWaitOps(scf::ForOp forOp, CoarseSchedule &schedule,
+                                   ttng::MMAv5OpInterface mma,
+                                   ttng::TMEMAllocOp alloc, int phaseArgIdx,
+                                   int barrierIdxArgIdx, int numStages) {
   OpBuilderForStage builder(forOp, schedule);
   DominanceInfo domInfo(forOp);
+  Value phase = forOp.getRegionIterArg(phaseArgIdx);
+  Value barrierIdx = forOp.getRegionIterArg(barrierIdxArgIdx);
   Value zero = builder.create<arith::ConstantIntOp>(forOp.getLoc(), 0, 32);
   Value one = builder.create<arith::ConstantIntOp>(forOp.getLoc(), 1, 32);
   Value numStagesVal =
@@ -889,8 +908,11 @@ void createBarrierAndWaitOps(scf::ForOp forOp, CoarseSchedule &schedule,
   builder.setStageCluster(schedule[mma]);
   builder.setInsertionPoint(mma);
   Location loc = mma->getLoc();
-  Value barrierSlice =
-      triton::createSingleBufferView(builder, barrierAlloc, barrierIdx);
+  Value barrierSlice = barrierAlloc;
+  if (numStages > 1) {
+    barrierSlice =
+        triton::createSingleBufferView(builder, barrierAlloc, barrierIdx);
+  }
   mma.setBarrier(barrierSlice);
 
   // List of buffers that may be used until wait completes
@@ -953,27 +975,85 @@ void createBarrierAndWaitOps(scf::ForOp forOp, CoarseSchedule &schedule,
     }
   }
 
+  std::optional<int> predicateArgIdx;
   for (auto wp : waitPoints) {
     builder.setStageCluster({wp.stage, wp.cluster});
     builder.setInsertionPoint(wp.op);
-    builder.create<ttng::WaitBarrierOp>(loc, barrierSlice, phase, waitBuffers);
+    SmallVector<Value> currWaitBuffers = waitBuffers;
+    Value pred = nullptr;
+    if (!domInfo.properlyDominates(mma, wp.op)) {
+      // Waits before the mma are special: they are dist-1 to the mma which
+      // means two things:
+      // 1. The wait should not be executed before the mma executes for the
+      //    first time
+      // 2. The waitBuffers do not dominate this wait. We either need to hoist
+      //    the local allocs out of the loop, or have dist-1 buffer references.
+      if (!predicateArgIdx) {
+        forOp = prepLoopForDist1Wait(forOp, schedule, mma);
+        predicateArgIdx = forOp.getInitArgs().size() - 1;
+        // HACK: Clear the wait buffers. This is meant to avoid having to have
+        // dist-1 buffer references, or hoisting the local allocs out of the
+        // loop. This works as long as the wait is in the same stage as the mma,
+        // and is the reason for which we prevent pipelining the mma if there
+        // are loads from the accumulator before the mma in a different stage.
+        currWaitBuffers.clear();
+      }
+      pred = forOp.getRegionIterArg(*predicateArgIdx);
+    }
+    Value barrierSlice = barrierAlloc;
+    if (numStages > 1) {
+      barrierSlice =
+          triton::createSingleBufferView(builder, barrierAlloc, barrierIdx);
+    }
+    builder.create<ttng::WaitBarrierOp>(loc, barrierSlice, phase, pred,
+                                        currWaitBuffers);
   }
 
   builder.setStageCluster(schedule[mma]);
   builder.setInsertionPoint(forOp.getBody()->getTerminator());
-  Value barWrap;
-  barrierIdx = createIncrementModulo(builder, loc, barrierIdx, numStagesVal,
-                                     zero, one, &barWrap);
-  phase = builder.create<arith::SelectOp>(
-      loc, phase.getType(), barWrap,
-      builder.create<arith::XOrIOp>(loc, phase, one), phase);
+  Value newPhase = builder.create<arith::XOrIOp>(loc, phase, one);
+  Value newBarrierIdx = barrierIdx;
+  if (numStages > 1) {
+    Value barWrap;
+    newBarrierIdx = createIncrementModulo(builder, loc, barrierIdx,
+                                          numStagesVal, zero, one, &barWrap);
+    newPhase = builder.create<arith::SelectOp>(loc, phase.getType(), barWrap,
+                                               newPhase, phase);
+  }
+  if (predicateArgIdx) {
+    // If there is a wait predicate, we need to select the phase and barrierIdx
+    // based on the predicate
+    Value pred = forOp.getRegionIterArg(*predicateArgIdx);
+    newPhase = builder.create<arith::SelectOp>(loc, phase.getType(), pred,
+                                               newPhase, phase);
+    newBarrierIdx = builder.create<arith::SelectOp>(loc, phase.getType(), pred,
+                                                    newBarrierIdx, barrierIdx);
+  }
+  replaceAllUsesDominatedBy(newPhase.getDefiningOp(), newPhase, phase, domInfo);
+  replaceAllUsesDominatedBy(newBarrierIdx.getDefiningOp(), newBarrierIdx,
+                            barrierIdx, domInfo);
+
+  // If there is a dist-1 wait, we need to add a wait after the loop
+  if (predicateArgIdx) {
+    builder.setInsertionPointAfter(forOp);
+    Value barrierSlice = barrierAlloc;
+    Value barrierIdx = forOp.getResult(barrierIdxArgIdx);
+    Value phase = forOp.getResult(phaseArgIdx);
+    if (numStages > 1) {
+      barrierSlice =
+          triton::createSingleBufferView(builder, barrierAlloc, barrierIdx);
+    }
+    builder.create<ttng::WaitBarrierOp>(loc, barrierSlice, phase,
+                                        forOp.getResult(*predicateArgIdx));
+  }
+  return forOp;
 }
 
 void multibufferTensorMemory(scf::ForOp forOp, CoarseSchedule &schedule,
-                             ttng::TMEMAllocOp alloc, Value &bufIdx,
-                             int bufIdxArgIdx, int tmemUseNumStages) {
+                             ttng::TMEMAllocOp alloc, int bufIdxArgIdx,
+                             int tmemUseNumStages) {
   DominanceInfo domInfo(forOp);
-
+  Value bufIdx = forOp.getRegionIterArg(bufIdxArgIdx);
   SmallVector<std::pair<Operation *, Value>> bufIdxDefs;
   auto getCurrBufIdx = [&](Operation *op) {
     for (auto [_op, _val] : llvm::reverse(bufIdxDefs)) {
@@ -1077,17 +1157,16 @@ void multibufferTensorMemory(scf::ForOp forOp, CoarseSchedule &schedule,
         "accumulator, and the mma uses the accumulator all the time.");
   }
   alloc->erase();
-  bufIdx = bufIdxDefs.back().second;
+  Value newBufIdx = bufIdxDefs.back().second;
+  replaceAllUsesDominatedBy(newBufIdx.getDefiningOp(), newBufIdx, bufIdx,
+                            domInfo);
 }
 
 scf::ForOp lowerMMA(ttng::MMAv5OpInterface mma, scf::ForOp forOp,
                     CoarseSchedule &schedule) {
   auto isLoadPipelineable = [&](Operation *op) {
     return schedule[mma].first > schedule[op].first;
   };
-  if (!mmaHasPipelineableOperands(mma, forOp, isLoadPipelineable)) {
-    return forOp;
-  }
   auto alloc = mma.getAccumulator().getDefiningOp<ttng::TMEMAllocOp>();
   if (!alloc) {
     return forOp;
@@ -1099,7 +1178,8 @@ scf::ForOp lowerMMA(ttng::MMAv5OpInterface mma, scf::ForOp forOp,
   int tmemUseNumStages =
       tmemUseStageBounds.second - tmemUseStageBounds.first + 1;
   int waitNumStages = tmemUseStageBounds.second - schedule[mma].first + 1;
-  if (waitNumStages == 1 && !hasAccReadModifyWrite(mma, forOp)) {
+  if (waitNumStages == 1 && !hasAccReadModifyWrite(mma, forOp) &&
+      mmaHasPipelineableOperands(mma, forOp, isLoadPipelineable)) {
     // Overlap the mma with itself, even if there is no use of the accumulator
     // after the mma
     waitNumStages = 2;
@@ -1112,35 +1192,40 @@ scf::ForOp lowerMMA(ttng::MMAv5OpInterface mma, scf::ForOp forOp,
   // Add arguments to the forOp
   unsigned newOperandIndex = forOp.getInitArgs().size();
   SmallVector<Value> newOperands = {
-      zero,     // phase
-      zero,     // barrierIdx
-      minusOne, // bufIdx
+      zero, // phase
+      zero, // barrierIdx
   };
+  if (tmemUseNumStages > 1) {
+    newOperands.push_back(minusOne); // bufIdx
+  }
   scf::ForOp newForOp =
       replaceForOpWithNewSignature(builder, forOp, newOperands);
   forOp.erase();
   forOp = newForOp;
 
-  Value phase = forOp.getRegionIterArg(newOperandIndex + 0);
-  Value barrierIdx = forOp.getRegionIterArg(newOperandIndex + 1);
-  Value bufIdx = forOp.getRegionIterArg(newOperandIndex + 2);
+  int phaseArgIdx = newOperandIndex + 0;
+  int barrierIdxArgIdx = newOperandIndex + 1;
+  int bufIdxArgIdx = newOperandIndex + 2;
+  Value phase = forOp.getRegionIterArg(phaseArgIdx);
+  Value barrierIdx = forOp.getRegionIterArg(barrierIdxArgIdx);
 
-  if (waitNumStages > 1) {
-    createBarrierAndWaitOps(forOp, schedule, mma, alloc, phase, barrierIdx,
-                            waitNumStages);
+  SmallVector<Value> newYieldOperands = {phase, barrierIdx};
+  if (tmemUseNumStages > 1) {
+    Value bufIdx = forOp.getRegionIterArg(bufIdxArgIdx);
+    newYieldOperands.push_back(bufIdx);
   }
+  cast<scf::YieldOp>(forOp.getBody()->getTerminator())
+      .getResultsMutable()
+      .append(newYieldOperands);
+
+  forOp = createBarrierAndWaitOps(forOp, schedule, mma, alloc, phaseArgIdx,
+                                  barrierIdxArgIdx, waitNumStages);
 
   if (tmemUseNumStages > 1) {
-    multibufferTensorMemory(forOp, schedule, alloc, bufIdx, newOperandIndex + 2,
+    multibufferTensorMemory(forOp, schedule, alloc, bufIdxArgIdx,
                             tmemUseNumStages);
   }
 
-  SmallVector<Value> newYieldOperands;
-  newYieldOperands.push_back(phase);
-  newYieldOperands.push_back(barrierIdx);
-  newYieldOperands.push_back(bufIdx);
-  appendToForOpYield(forOp, newYieldOperands);
-
   return forOp;
 }