[Mlir-commits] [mlir] [mlir] Add splitForOpAtBound utility and use it in loop unrolling (PR #215108)

[llvmlistbot at llvm.org]

https://github.com/davidlerner96 updated https://github.com/llvm/llvm-project/pull/215108

>From 5222c55626778b1cc4385118ed4c7e74e8df7ae5 Mon Sep 17 00:00:00 2001
From: David Lerner <davidlerner96 at gmail.com>
Date: Sun, 9 Aug 2026 18:20:20 +0300
Subject: [PATCH] [mlir] Add splitForOpAtBound utility and use it in loop
 unrolling

Factor scf.for loop splitting into a new splitForOpAtBound helper. The utility
reuses the input loop for [lowerBound, splitPoint), clones a second loop for
[splitPoint, upperBound), and returns the clone. Loop-carried values are
chained across the split by initializing the second loop from the first loop's
results and redirecting users of the first loop's results to the second.

When all bounds are constant, the utility fails unless lowerBound <= splitPoint
< upperBound, using signed or unsigned comparison per the loop's unsignedCmp
attribute. Non-constant bounds are accepted without this check.

Use splitForOpAtBound in loopUnrollByFactor when building epilogue loops,
replacing the duplicated clone-and-rewire logic.
---
 mlir/include/mlir/Dialect/SCF/Utils/Utils.h | 10 +++
 mlir/lib/Dialect/SCF/Utils/Utils.cpp        | 76 ++++++++++++++++-----
 2 files changed, 69 insertions(+), 17 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index a758032ef69b4..f7dce3964156c 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -117,6 +117,16 @@ struct UnrolledLoopInfo {
   std::optional<scf::ForOp> epilogueLoopOp = std::nullopt;
 };
 
+/// Splits `forOp` into two consecutive loops at `splitPoint`:
+///   first:  [lowerBound, splitPoint)
+///   second: [splitPoint, upperBound)
+///
+/// Reuses `forOp` as the first loop and clones a second loop. Returns the
+/// second loop. When all bounds are constant, fails unless splitPoint is within
+/// the loop bounds. Chains loop-carried values from the first loop's results to
+/// the second.
+FailureOr<scf::ForOp> splitForOpAtBound(scf::ForOp forOp, Value splitPoint);
+
 /// Unrolls this for operation by the specified unroll factor. Returns the
 /// unrolled main loop and the epilogue loop, if the loop is unrolled. Otherwise
 /// returns failure if the loop cannot be unrolled either due to restrictions or
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index c789b4c8904d3..237a3cd12b9fe 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -363,6 +363,58 @@ void mlir::generateUnrolledLoop(
   loopBodyBlock->getTerminator()->setOperands(lastYielded);
 }
 
+/// Splits `forOp` into two consecutive loops at `splitPoint`.
+FailureOr<scf::ForOp> mlir::splitForOpAtBound(scf::ForOp forOp,
+                                              Value splitPoint) {
+  if (splitPoint.getType() != forOp.getLowerBound().getType())
+    return failure();
+
+  // When all bounds are constant, require a valid split point.
+  if (forOp.getUnsignedCmp()) { // Unsigned comparison
+    auto getUnsignedBound = [](Value v) -> std::optional<uint64_t> {
+      auto apInt = getConstantAPIntValue(v);
+      if (!apInt)
+        return std::nullopt;
+      return apInt->first.getZExtValue();
+    };
+    auto lb = getUnsignedBound(forOp.getLowerBound());
+    auto ub = getUnsignedBound(forOp.getUpperBound());
+    auto split = getUnsignedBound(splitPoint);
+    if (lb && ub && split && (*lb > *split || *split >= *ub))
+      return failure();
+  } else { // Signed comparison
+    auto lbCst = getConstantIntValue(forOp.getLowerBound());
+    auto ubCst = getConstantIntValue(forOp.getUpperBound());
+    auto splitCst = getConstantIntValue(splitPoint);
+    if (lbCst && ubCst && splitCst &&
+        (*lbCst > *splitCst || *splitCst >= *ubCst))
+      return failure();
+  }
+
+  OpBuilder builder(forOp->getContext());
+  builder.setInsertionPointAfter(forOp);
+  // Clone before mutating `forOp` so the second loop inherits the original
+  // body and bounds.
+  auto secondForOp = cast<scf::ForOp>(builder.clone(*forOp));
+  secondForOp.setLowerBound(splitPoint);
+
+  // Chain iter-args across the split:
+  //   - `secondForOp` is initialized from `forOp`'s results.
+  //   - Users of `forOp`'s results are redirected to `secondForOp`'s results,
+  //     so downstream code observes the final carried values.
+  auto results = forOp.getResults();
+  auto secondResults = secondForOp.getResults();
+  for (auto [origResult, secondResult] : llvm::zip(results, secondResults))
+    origResult.replaceAllUsesWith(secondResult);
+  secondForOp->setOperands(secondForOp.getNumControlOperands(),
+                           secondForOp.getInitArgs().size(), results);
+
+  // Truncate the first loop to [lowerBound, splitPoint).
+  forOp.setUpperBound(splitPoint);
+
+  return secondForOp;
+}
+
 /// Unrolls 'forOp' by 'unrollFactor', returns the unrolled main loop and the
 /// epilogue loop, if the loop is unrolled.
 FailureOr<UnrolledLoopInfo> mlir::loopUnrollByFactor(
@@ -473,26 +525,16 @@ FailureOr<UnrolledLoopInfo> mlir::loopUnrollByFactor(
 
   // Create epilogue clean up loop starting at 'upperBoundUnrolled'.
   if (generateEpilogueLoop) {
-    OpBuilder epilogueBuilder(forOp->getContext());
-    epilogueBuilder.setInsertionPointAfter(forOp);
-    auto epilogueForOp = cast<scf::ForOp>(epilogueBuilder.clone(*forOp));
-    epilogueForOp.setLowerBound(upperBoundUnrolled);
-
-    // Update uses of loop results.
-    auto results = forOp.getResults();
-    auto epilogueResults = epilogueForOp.getResults();
-
-    for (auto e : llvm::zip(results, epilogueResults)) {
-      std::get<0>(e).replaceAllUsesWith(std::get<1>(e));
-    }
-    epilogueForOp->setOperands(epilogueForOp.getNumControlOperands(),
-                               epilogueForOp.getInitArgs().size(), results);
-    if (epilogueForOp.promoteIfSingleIteration(rewriter).failed())
-      resultLoops.epilogueLoopOp = epilogueForOp;
+    auto epilogueForOp = splitForOpAtBound(forOp, upperBoundUnrolled);
+    if (failed(epilogueForOp))
+      return failure();
+    if (epilogueForOp->promoteIfSingleIteration(rewriter).failed())
+      resultLoops.epilogueLoopOp = *epilogueForOp;
+  } else {
+    forOp.setUpperBound(upperBoundUnrolled);
   }
 
   // Create unrolled loop.
-  forOp.setUpperBound(upperBoundUnrolled);
   forOp.setStep(stepUnrolled);
 
   auto iterArgs = ValueRange(forOp.getRegionIterArgs());