[Mlir-commits] [clang] [mlir] [llvm] [flang] [MLIR][LLVM] Add Continuous Loop Peeling transform to SCF (PR #77328)

[llvmlistbot at llvm.org]

================
@@ -105,6 +106,161 @@ static void specializeForLoopForUnrolling(ForOp op) {
   op.erase();
 }
 
+/// Create a new for loop for the remaining iterations (partialIteration)
+/// after a for loop has been peeled. This is followed by correcting the
+/// loop bounds for both loops given the index (splitBound) where the
+/// iteration space is to be split up. Returns failure if the loop can not
+/// be split and no new partialIteration is created.
+static LogicalResult splitLoopHelper(RewriterBase &b, scf::ForOp forOp,
+                                     scf::ForOp &partialIteration,
+                                     Value splitBound) {
+  RewriterBase::InsertionGuard guard(b);
+  auto lbInt = getConstantIntValue(forOp.getLowerBound());
+  auto ubInt = getConstantIntValue(forOp.getUpperBound());
+  auto stepInt = getConstantIntValue(forOp.getStep());
+
+  // No specialization necessary if step already divides upper bound evenly.
+  if (lbInt && ubInt && stepInt && (*ubInt - *lbInt) % *stepInt == 0)
+    return failure();
+  // No specialization necessary if step size is 1.
+  if (stepInt == static_cast<int64_t>(1))
+    return failure();
+
+  // Create ForOp for partial iteration.
+  b.setInsertionPointAfter(forOp);
+  IRMapping map;
+  auto constStepOp =
+      b.create<arith::ConstantIndexOp>(forOp.getLoc(), *stepInt / 2);
+  // The new for loop for the remaining iterations has half the step size
+  // as continuous peeling requires the step size to diminish exponentially
+  // across subsequent loops.
+  map.map(forOp.getStep(), constStepOp);
----------------
muneebkhan85 wrote:

Yes, the purpose is to canonicalize min/max ops inside the new loop body. This requires the affine.min/max ops to have the correct step given the step size has changed (halved) for the partialIteraton loop.

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