Fix for constant-bound loop where no epilog will be generated.

Author: htyu

lib/Dialect/Triton/Transforms/LoopUnroll.cpp

@@ -62,20 +62,23 @@ class LoopUnrollPass : public TritonLoopUnrollBase<LoopUnrollPass> {
     // 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);
-    }
+    assert(loops.size() <= 2 && "Expect at most 2 loops, one for the main loop "
+                                "and one for the prolog/epilog");
+    if (loops.size() == 2) {
+      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]);
+      // 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;
   }