[Mlir-commits] [mlir] [mlir] Add splitForOpAtBound utility and use it in loop unrolling (PR #215108)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Aug 12 04:04:14 PDT 2026
https://github.com/davidlerner96 updated https://github.com/llvm/llvm-project/pull/215108
>From 38c7934c40043007616f7f268cfa109c5263f3a9 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 ++++++++++++++++-----
mlir/test/lib/Dialect/SCF/CMakeLists.txt | 1 +
mlir/tools/mlir-opt/mlir-opt.cpp | 2 +
4 files changed, 72 insertions(+), 17 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index 316bf179e22c3..c5c6871bd758e 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -108,6 +108,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> splitForOpAtPoint(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..c266a2de0e54d 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::splitForOpAtPoint(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 = splitForOpAtPoint(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());
diff --git a/mlir/test/lib/Dialect/SCF/CMakeLists.txt b/mlir/test/lib/Dialect/SCF/CMakeLists.txt
index d2f97e816cc14..80b8268768d46 100644
--- a/mlir/test/lib/Dialect/SCF/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/SCF/CMakeLists.txt
@@ -4,6 +4,7 @@ add_mlir_library(MLIRSCFTestPasses
TestLoopUnrolling.cpp
TestParallelLoopUnrolling.cpp
TestSCFUtils.cpp
+ TestSplitForOpAtPoint.cpp
TestSCFWrapInZeroTripCheck.cpp
TestUpliftWhileToFor.cpp
TestWhileOpBuilder.cpp
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index 47dffb1b2d28b..bad235043a096 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -143,6 +143,7 @@ void registerTestParallelLoopUnrollingPass();
void registerTestRecursiveTypesPass();
void registerTestSCFUpliftWhileToFor();
void registerTestSCFUtilsPass();
+void registerTestSplitForOpAtPointPass();
void registerTestSCFWhileOpBuilderPass();
void registerTestSCFWrapInZeroTripCheckPasses();
void registerTestShapeMappingPass();
@@ -292,6 +293,7 @@ static void registerTestPasses() {
mlir::test::registerTestRecursiveTypesPass();
mlir::test::registerTestSCFUpliftWhileToFor();
mlir::test::registerTestSCFUtilsPass();
+ mlir::test::registerTestSplitForOpAtPointPass();
mlir::test::registerTestSCFWhileOpBuilderPass();
mlir::test::registerTestSCFWrapInZeroTripCheckPasses();
mlir::test::registerTestShapeMappingPass();
More information about the Mlir-commits
mailing list