[PIPELINER] Reintroduce epilogue peeling (#6962)

Reintroduce "[PIPELINE] Peel single epilogue iteration after loop
expansion #6893"
Fixed the issue that was exposing IMA by using poison values as buffer
indices used by speculatively executed async copy.

Author: pawelszczerbuk

bin/CMakeLists.txt

@@ -12,6 +12,7 @@ target_link_libraries(triton-opt PRIVATE
   ${triton_libs}
   # tests
   TritonTestAnalysis
+  TritonTestDialect
   TritonAMDGPUTestAnalysis
   # MLIR core
   MLIROptLib
@@ -31,6 +32,7 @@ target_link_libraries(triton-reduce PRIVATE
   ${triton_libs}
   # tests
   TritonTestAnalysis
+  TritonTestDialect
   TritonAMDGPUTestAnalysis
   # MLIR core
   MLIRReduceLib
@@ -49,6 +51,7 @@ target_link_libraries(triton-lsp PRIVATE
   ${triton_libs}
   # tests
   TritonTestAnalysis
+  TritonTestDialect
   TritonAMDGPUTestAnalysis
   # MLIR core
   MLIRLspServerLib
@@ -85,5 +88,6 @@ target_link_libraries(triton-tensor-layout PRIVATE
   ${conversion_libs}
   ${dialect_libs}
   TritonTestAnalysis
+  TritonTestDialect
   TritonAMDGPUTestAnalysis
   )

bin/RegisterTritonDialects.h

@@ -37,6 +37,7 @@ void registerTestAllocationPass();
 void registerTestMembarPass();
 void registerTestAMDGPUMembarPass();
 void registerTestTritonAMDGPURangeAnalysis();
+void registerTestLoopPeelingPass();
 } // namespace test
 } // namespace mlir
 
@@ -49,6 +50,7 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::test::registerTestAlignmentPass();
   mlir::test::registerTestAllocationPass();
   mlir::test::registerTestMembarPass();
+  mlir::test::registerTestLoopPeelingPass();
   mlir::test::registerTestAMDGPUMembarPass();
   mlir::test::registerTestTritonAMDGPURangeAnalysis();
   mlir::triton::registerConvertTritonToTritonGPUPass();

include/triton/Dialect/Triton/IR/Utility.h

@@ -177,6 +177,9 @@ template <typename T> auto seq(T start, T end, T step) {
 Value getPredMask(RewriterBase &rewriter, Type typeLike, Value currentMask,
                   Value pred);
 
+// Get the value of the induction variable at the end of the loop.
+Value getLastInductionValue(OpBuilder &b, scf::ForOp loop);
+
 MakeTensorPtrOp getMakeTensorPtrOp(Value v);
 
 } // namespace triton

include/triton/Dialect/Triton/Transforms/LoopPeeling.h

@@ -0,0 +1,18 @@
+#ifndef TRITON_DIALECT_TRITON_TRANSFORMS_LOOP_PEELING_H_
+#define TRITON_DIALECT_TRITON_TRANSFORMS_LOOP_PEELING_H_
+
+#include "mlir/Dialect/SCF/IR/SCF.h"
+
+namespace mlir {
+namespace triton {
+
+// Peel the single last iteration of the loop.
+void peelLoopEpilogue(
+    scf::ForOp forOp,
+    function_ref<Operation *(RewriterBase &, Operation *, bool)>
+        processPeeledOp = nullptr);
+
+} // namespace triton
+} // namespace mlir
+
+#endif // TRITON_DIALECT_TRITON_TRANSFORMS_LOOP_PEELING_H_

include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td

@@ -337,6 +337,24 @@ def TTG_PredicateStageOp: TTG_Op<"predicate_stage",
   let assemblyFormat = "$iv `,` $ub `,` $step `maxStage` $maxStage `stage` $stage attr-dict `:` type($iv) `->` type($result)";
 }
 
+def TTG_MaskOp: TTG_Op<"mask",
+                       [SingleBlock]> {
+    let summary = "mask op for pipelining";
+    let arguments = (ins I1:$pred);
+    let results = (outs Variadic<AnyType>:$result);
+    let regions = (region SizedRegion<1>:$region);
+    let builders = [
+        OpBuilder<(ins "Value":$pred)>,
+    ];
+}
+
+def TTG_MaskReturnOp: TTG_Op<"mask.return",
+                             [HasParent<"MaskOp">, Pure, Terminator, ReturnLike]> {
+    let summary = "terminator for mask operator";
+    let arguments = (ins Variadic<AnyType>:$result);
+    let assemblyFormat = "$result attr-dict `:` type($result)";
+}
+
 def TTG_Fp4ToFpOp : TTG_Op<"fp4_to_fp", [Pure]> {
   let summary = "Upcast fp4 (e2m1) to fp";
 

lib/Dialect/Triton/IR/Utility.cpp

@@ -1,5 +1,6 @@
 #include "triton/Dialect/Triton/IR/Utility.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
 
 using namespace mlir;
@@ -90,3 +91,15 @@ tt::MakeTensorPtrOp tt::getMakeTensorPtrOp(Value v) {
   }
   llvm_unreachable("Unable to getMakeTensorPtr()");
 }
+
+Value tt::getLastInductionValue(OpBuilder &b, scf::ForOp loop) {
+  Location loc = loop.getLoc();
+  // (ub - lb -1) // step * step + lb
+  Value diff =
+      b.create<arith::SubIOp>(loc, loop.getUpperBound(), loop.getLowerBound());
+  diff = b.create<arith::SubIOp>(
+      loc, diff, b.create<arith::ConstantOp>(loc, b.getI32IntegerAttr(1)));
+  Value ceilStep = b.create<arith::MulIOp>(
+      loc, b.create<arith::DivSIOp>(loc, diff, loop.getStep()), loop.getStep());
+  return b.create<arith::AddIOp>(loc, ceilStep, loop.getLowerBound());
+}

lib/Dialect/Triton/Transforms/CMakeLists.txt

@@ -6,6 +6,7 @@ add_triton_library(TritonTransforms
   Combine.cpp
   LoopAwareCSE.cpp
   LoopInvariantCodeMotion.cpp
+  LoopPeeling.cpp
   LoopUnroll.cpp
   ReorderBroadcast.cpp
   RewriteTensorPointer.cpp
@@ -20,5 +21,7 @@ add_triton_library(TritonTransforms
   LINK_LIBS PUBLIC
   MLIRPass
   MLIRTransformUtils
+  MLIRTransforms
+  MLIRSCFToControlFlow
   TritonIR
 )

lib/Dialect/Triton/Transforms/LoopPeeling.cpp

@@ -0,0 +1,68 @@
+#include "triton/Dialect/Triton/Transforms/LoopPeeling.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Pass/Pass.h"
+#include "triton/Dialect/Triton/IR/Utility.h"
+
+using namespace mlir;
+
+namespace mlir {
+namespace triton {
+
+void peelLoopEpilogue(
+    scf::ForOp forOp,
+    function_ref<Operation *(RewriterBase &, Operation *, bool)>
+        processPeeledOp) {
+  SmallVector<Operation *> loopBodyOps;
+  IRRewriter rewriter(forOp);
+  Location loc = forOp.getLoc();
+  Type type = forOp.getStep().getType();
+
+  // Fetch loop bounds and step
+  Value lowerBound = forOp.getLowerBound();
+  Value upperBound = forOp.getUpperBound();
+  Value step = forOp.getStep();
+  Value newUpperBound = rewriter.create<arith::SubIOp>(loc, upperBound, step);
+
+  rewriter.setInsertionPointAfter(forOp);
+  Value lastIV = getLastInductionValue(rewriter, forOp);
+
+  auto cond = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt,
+                                             lowerBound, upperBound);
+
+  // Create an if op to execute the peeled iteration
+  IRMapping map;
+  map.map(forOp.getRegionIterArgs(), forOp.getResults());
+  map.map(forOp.getInductionVar(), lastIV);
+  auto ifOp = rewriter.create<scf::IfOp>(loc, forOp.getResultTypes(), cond,
+                                         /*hasElse=*/true);
+  ifOp.getThenRegion().front().erase();
+  forOp.getBodyRegion().cloneInto(&ifOp.getThenRegion(), map);
+  rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front());
+  rewriter.create<scf::YieldOp>(loc, forOp.getResults());
+
+  forOp->replaceUsesWithIf(ifOp, [&](OpOperand &operand) {
+    return !ifOp->isAncestor(operand.getOwner());
+  });
+
+  forOp.getUpperBoundMutable().assign(newUpperBound);
+
+  if (processPeeledOp) {
+    for (auto &op :
+         llvm::make_early_inc_range(forOp.getBody()->without_terminator())) {
+      Operation *newOp = processPeeledOp(rewriter, &op, /*isEpilogue=*/false);
+      if (newOp && newOp != &op) {
+        op.replaceAllUsesWith(newOp);
+      }
+    }
+    for (auto &op : llvm::make_early_inc_range(
+             ifOp.getThenRegion().front().without_terminator())) {
+      Operation *newOp = processPeeledOp(rewriter, &op, /*isEpilogue=*/true);
+      if (newOp && newOp != &op) {
+        op.replaceAllUsesWith(newOp);
+      }
+    }
+  }
+}
+
+} // namespace triton
+} // namespace mlir

lib/Dialect/TritonGPU/Transforms/CMakeLists.txt

@@ -44,6 +44,7 @@ add_triton_library(TritonGPUTransforms
   MLIRTransformUtils
   TritonAnalysis
   TritonIR
+  TritonTransforms
   TritonGPUIR
   TritonNvidiaGPUIR
   TritonToTritonGPU

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

@@ -275,10 +275,10 @@ class AssignMMALatencies {
         // place the wait right before the loads.
 
         if (hasSyncDots(forOp)) {
-          // Skip pipelining MMA in the loops where sync dots are used. This is
-          // dirty heuristic for performance drops in kernels where we would
-          // rather want to have last iteration peeled instead of having a full
-          // iteration of masked operations only to execute single wait.
+          // Skip pipelining MMA in the loops where sync dots are used. This
+          // is a dirty heuristic for performance drops in kernels where we
+          // would rather want to have last iteration peeled instead of having a
+          // full iteration of masked operations only to execute single wait.
           continue;
         }
         auto pipeHelper = ttng::MMAv5PipelineableOperandsHelper(