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

[Matthias Springer]

================
@@ -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);
----------------
matthias-springer wrote:

I think this won't work. The SSA value of `forOp.getStep()` could be used in different ways inside of the loop and you don't want to change that.

E.g.:
```mlir
scf.for ... step %c16 {
  // This op should not be changed as part of loop peeling.
  "test.foo"(%c16) : () -> ()
}
```

What's the purpose of this `map.map`? Is it meant to canonicalize `affine.min/max` ops, taking into account the fact that the loop was peeled?


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