[BACKEND] Do not pipeline epilog loops generated by loop unrolling

Author: htyu

lib/Dialect/Triton/Transforms/LoopUnroll.cpp

@@ -1,6 +1,7 @@
 #include <memory>
 
 #include "mlir/Dialect/SCF/Utils/Utils.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
@@ -22,24 +23,63 @@
 
 namespace mlir::triton {
 
-static const char *loopUnrollFactorAttrName = "tt.loop_unroll_factor";
-
 namespace {
 
 class LoopUnrollPass : public TritonLoopUnrollBase<LoopUnrollPass> {
 
   int getUnrollFactorOrDefault(scf::ForOp forOp) {
     // Use the attribute attached to the loop if it exists otherwise set the
     // factor to 1 to suppress the unrolling.
-    if (auto factor = forOp->getAttrOfType<IntegerAttr>(
-            mlir::triton::loopUnrollFactorAttrName))
+    if (auto factor =
+            forOp->getAttrOfType<IntegerAttr>(loopUnrollFactorAttrName))
       return factor.getInt();
     return 1;
   }
 
+  int getUnrollIdOrDefault(scf::ForOp forOp) {
+    // Use the attribute attached to the loop if it exists otherwise set the
+    // factor to 1 to suppress the unrolling.
+    if (auto factor = forOp->getAttrOfType<IntegerAttr>(unrolledLoopIdAttrName))
+      return factor.getInt();
+    return 0;
+  }
+
+  const char *loopUnrollFactorAttrName = "tt.loop_unroll_factor";
+  const char *unrolledLoopIdAttrName = "tt.unrolled_loop_id";
+  const char *pipelineStagesAttrName = "tt.num_stages";
+
 public:
   LoopUnrollPass() = default;
   LoopUnrollPass(const LoopUnrollPass &) {}
+
+  SmallVector<scf::ForOp, 2> getUnrolledLoopsAndClearAttrs(unsigned loopId) {
+    SmallVector<scf::ForOp, 2> loops;
+    getOperation()->walk([&](scf::ForOp forOp) {
+      if (getUnrollIdOrDefault(forOp) == loopId)
+        loops.push_back(forOp);
+    });
+
+    // check which one is the unrolled loop and which one is the prolog/epilog
+    // loop. A simple heuristic is to check the number of instructions in the
+    // loop. The unrolled main loop should have the most instructions.
+    assert(loops.size() == 2 && "only support unrolling one loop at a time");
+    SmallVector<int, 2> loopInstructionCounts;
+    for (auto loop : loops) {
+      loop->removeAttr(loopUnrollFactorAttrName);
+      loop->removeAttr(unrolledLoopIdAttrName);
+      int count = 0;
+      loop->walk([&](Operation *op) { count++; });
+      loopInstructionCounts.push_back(count);
+    }
+
+    // sort the loops by the number of instructions. The unrolled main loop
+    // should go first.
+    if (loopInstructionCounts[0] < loopInstructionCounts[1])
+      std::swap(loops[0], loops[1]);
+
+    return loops;
+  }
+
   void runOnOperation() override {
     LDBG("Loop unroll pass");
     SmallVector<scf::ForOp, 4> loops;
@@ -49,11 +89,22 @@ class LoopUnrollPass : public TritonLoopUnrollBase<LoopUnrollPass> {
         loops.push_back(forOp);
     });
 
-    for (auto loop : loops) {
+    auto ctx = getOperation()->getContext();
+    for (unsigned i = 0; i < loops.size(); i++) {
+      auto loop = loops[i];
       auto unrollFactor = getUnrollFactorOrDefault(loop);
-      loop->removeAttr(mlir::triton::loopUnrollFactorAttrName);
+      loop->setAttr(unrolledLoopIdAttrName,
+                    mlir::IntegerAttr::get(IntegerType::get(ctx, 32), i + 1));
       LDBG("Unrolling loop by " << unrollFactor << " times\n" << loop);
       (void)loopUnrollByFactor(loop, unrollFactor);
+      auto unrolledLoops = getUnrolledLoopsAndClearAttrs(i + 1);
+      // Do not pipeline the prolog/epilog loop.
+      if (unrolledLoops.size() == 2) {
+        auto prologEpilogLoop = unrolledLoops[1];
+        prologEpilogLoop->setAttr(
+            pipelineStagesAttrName,
+            mlir::IntegerAttr::get(IntegerType::get(ctx, 32), 1));
+      }
     }
   }
 };

test/Triton/loop-unroll.mlir

@@ -13,6 +13,7 @@ tt.func @add_kernel_unroll(%arg0: tensor<256x!tt.ptr<f32>>, %arg1: i32) {
   // CHECK: scf.for
   // CHECK: tt.load
   // CHECK-NOT: tt.load
+  // CHECK: tt.num_stages = 1 : i32
   %2:2 = scf.for %arg3 = %c1_i32 to %arg1 step %c1_i32 iter_args(%arg4 = %1, %arg5 = %arg0) -> (tensor<256xf32>, tensor<256x!tt.ptr<f32>>)  : i32 {
       %3 = tt.load %arg5 : tensor<256x!tt.ptr<f32>>
     %4 = arith.addf %arg4, %3 : tensor<256xf32>