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

[Matthias Springer via flang-commits]

================
@@ -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:

Generally speaking, for whatever modification you make to the loop body, you have to be sure that the loop body is still computing the same thing as before. Just blanket replacing all occurrences of the old step size (even if it's just in `affine.min/max` ops) with the new step size may change the semantics of the loop.

The only safe way of canonicalizing the loop body that we have at the moment is `rewritePeeledMinMaxOp`. This function will look for `affine.min/max` ops and select one of the provided options if it can prove that doing so would always be correct in the peeled or partial loop.



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