[Mlir-commits] [clang] [flang] [llvm] [mlir] [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:

The original transform was intended for applying tiling repeatedly by lowering the step size. This doesn't work as generic loop transform and the right place to implement this is in Linalg/TransformOps. The patch will be moved there and submitted in a separate PR.

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