[Mlir-commits] [mlir] [mlir][SCF] Add support for peeling the first iteration out of the loop (PR #74015)

[llvmlistbot at llvm.org]

================
@@ -205,32 +222,84 @@ LogicalResult mlir::scf::peelForLoopAndSimplifyBounds(RewriterBase &rewriter,
   return success();
 }
 
+LogicalResult mlir::scf::peelFirstIterationForLoop(RewriterBase &b, ForOp forOp,
+                                                   ForOp &firstIteration) {
+  RewriterBase::InsertionGuard guard(b);
+  auto lbInt = getConstantIntValue(forOp.getLowerBound());
+  auto ubInt = getConstantIntValue(forOp.getUpperBound());
+  auto stepInt = getConstantIntValue(forOp.getStep());
+
+  // Peeling is not needed if there is one or less iteration.
+  if (lbInt && ubInt && stepInt && (*ubInt - *lbInt) / *stepInt <= 1)
+    return success();
----------------
yzhang93 wrote:

I've changed this to `return failure`. The reason I wanted to return success is because I don't want the transform dialect operation `transform.loop.peel` failed when the for loop only has one iteration. I'm not sure how to make a `if else` condition in transform dialect like the regular C++ does.

https://github.com/llvm/llvm-project/pull/74015