[LoopUnroll] Do not pipeline epilog loops generated by loop unrolling (…

…triton-lang#5027)

The epilog loop created by the loop unroller may not be run if the main
unrolled loop covers all original loop iterations, thus pipelining it
non-speculatively may not be beneficial. It can also cause some
correctness issue when combined with the downstream PTXAS optimizer.

lib/Dialect/Triton/Transforms/LoopUnroll.cpp

@@ -22,21 +22,22 @@
 
 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;
   }
 
+  const char *loopUnrollFactorAttrName = "tt.loop_unroll_factor";
+  const char *pipelineStagesAttrName = "tt.num_stages";
+
 public:
   LoopUnrollPass() = default;
   LoopUnrollPass(const LoopUnrollPass &) {}
@@ -49,11 +50,18 @@ class LoopUnrollPass : public TritonLoopUnrollBase<LoopUnrollPass> {
         loops.push_back(forOp);
     });
 
+    auto ctx = getOperation()->getContext();
     for (auto loop : loops) {
       auto unrollFactor = getUnrollFactorOrDefault(loop);
-      loop->removeAttr(mlir::triton::loopUnrollFactorAttrName);
+      loop->removeAttr(loopUnrollFactorAttrName);
       LDBG("Unrolling loop by " << unrollFactor << " times\n" << loop);
-      (void)loopUnrollByFactor(loop, unrollFactor);
+      auto resultLoops = loopUnrollByFactor(loop, unrollFactor);
+      // Do not pipeline the epilog loop.
+      if (succeeded(resultLoops) && resultLoops->epilogueLoopOp) {
+        (*resultLoops->epilogueLoopOp)
+            ->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>