Simplified based on MLIR-side changes.

Author: htyu

lib/Dialect/Triton/Transforms/LoopUnroll.cpp

@@ -52,36 +52,6 @@ class LoopUnrollPass : public TritonLoopUnrollBase<LoopUnrollPass> {
   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;
@@ -95,17 +65,14 @@ class LoopUnrollPass : public TritonLoopUnrollBase<LoopUnrollPass> {
     for (unsigned i = 0; i < loops.size(); i++) {
       auto loop = loops[i];
       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));
       }
     }
   }