Merge commit '8de17d2525531c9cb9690405add1ed615c0410cd'

Author: whitneywhtsang

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;
 }