Simplified based on MLIR-side changes.

Author: htyu

lib/Dialect/Triton/Transforms/LoopUnroll.cpp

@@ -1,7 +1,6 @@
 #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"
@@ -36,52 +35,12 @@ class LoopUnrollPass : public TritonLoopUnrollBase<LoopUnrollPass> {
     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);
-    });
-
-    assert(loops.size() <= 2 && "Expect at most 2 loops, one for the main loop "
-                                "and one for the prolog/epilog");
-    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);
-    }
-    if (loops.size() == 2) {
-      // 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.
-      // 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;
@@ -92,20 +51,16 @@ class LoopUnrollPass : public TritonLoopUnrollBase<LoopUnrollPass> {
     });
 
     auto ctx = getOperation()->getContext();
-    for (unsigned i = 0; i < loops.size(); i++) {
-      auto loop = loops[i];
+    for (auto loop : loops) {
       auto unrollFactor = getUnrollFactorOrDefault(loop);
-      loop->setAttr(unrolledLoopIdAttrName,
-                    mlir::IntegerAttr::get(IntegerType::get(ctx, 32), i + 1));
+      loop->removeAttr(loopUnrollFactorAttrName);
       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));
+      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));
       }
     }
   }