[Mlir-commits] [mlir] [mlir] Add splitForOpAtBound utility and use it in loop unrolling (PR #215108)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Aug 13 04:10:00 PDT 2026
https://github.com/davidlerner96 updated https://github.com/llvm/llvm-project/pull/215108
>From 4dc3beadcf252bfee5f2fdaed78f1d4c93918c44 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 ++++++++++++++-----
.../Dialect/SCF/split-for-op-at-point.mlir | 53 +++++++++++++
mlir/test/lib/Dialect/SCF/CMakeLists.txt | 1 +
.../lib/Dialect/SCF/TestSplitForOpAtPoint.cpp | 72 ++++++++++++++++++
mlir/tools/mlir-opt/mlir-opt.cpp | 2 +
6 files changed, 197 insertions(+), 17 deletions(-)
create mode 100644 mlir/test/Dialect/SCF/split-for-op-at-point.mlir
create mode 100644 mlir/test/lib/Dialect/SCF/TestSplitForOpAtPoint.cpp
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/Dialect/SCF/split-for-op-at-point.mlir b/mlir/test/Dialect/SCF/split-for-op-at-point.mlir
new file mode 100644
index 0000000000000..ca6e6cc0139a7
--- /dev/null
+++ b/mlir/test/Dialect/SCF/split-for-op-at-point.mlir
@@ -0,0 +1,53 @@
+// RUN: mlir-opt %s -test-split-for-op-at-point -split-input-file -verify-diagnostics | FileCheck %s
+
+// Split [0, 10) at 9 into [0, 9) and [9, 10).
+func.func @basic_split(%mem: memref<?xf32>) {
+ %cst = arith.constant 0.0 : f32
+ %c0 = arith.constant 0 : index
+ %c10 = arith.constant 10 : index
+ %c1 = arith.constant 1 : index
+ scf.for %i = %c0 to %c10 step %c1 {
+ memref.store %cst, %mem[%i] : memref<?xf32>
+ } {test.split_at = 9 : index}
+ return
+}
+// CHECK-LABEL: func @basic_split
+// CHECK: scf.for %{{.*}} = %c0 to %c9 step %c1
+// CHECK-NEXT: memref.store
+// CHECK: scf.for %{{.*}} = %c9 to %c10 step %c1
+// CHECK-NEXT: memref.store
+
+// -----
+
+// Chain loop-carried values across the split.
+func.func @split_with_iter_args() -> i32 {
+ %c0 = arith.constant 0 : index
+ %c10 = arith.constant 10 : index
+ %c1 = arith.constant 1 : index
+ %c0_i32 = arith.constant 0 : i32
+ %r = scf.for %i = %c0 to %c10 step %c1 iter_args(%arg = %c0_i32) -> i32 {
+ %one = arith.constant 1 : i32
+ %add = arith.addi %arg, %one : i32
+ scf.yield %add : i32
+ } {test.split_at = 9 : index}
+ return %r : i32
+}
+// CHECK-LABEL: func @split_with_iter_args
+// CHECK: %[[FIRST:.*]] = scf.for %{{.*}} = %c0 to %c9 step %c1 iter_args(%{{.*}} = %c0_i32) -> (i32)
+// CHECK: %[[RESULT:.*]] = scf.for %{{.*}} = %c9 to %c10 step %c1 iter_args(%{{.*}} = %[[FIRST]]) -> (i32)
+// CHECK: return %[[RESULT]] : i32
+
+// -----
+
+// Invalid split point is rejected when bounds are constant.
+func.func @invalid_split(%mem: memref<?xf32>) {
+ %cst = arith.constant 0.0 : f32
+ %c0 = arith.constant 0 : index
+ %c10 = arith.constant 10 : index
+ %c1 = arith.constant 1 : index
+ // expected-error @+1 {{failed to split scf.for}}
+ scf.for %i = %c0 to %c10 step %c1 {
+ memref.store %cst, %mem[%i] : memref<?xf32>
+ } {test.split_at = 10 : index}
+ return
+}
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/test/lib/Dialect/SCF/TestSplitForOpAtPoint.cpp b/mlir/test/lib/Dialect/SCF/TestSplitForOpAtPoint.cpp
new file mode 100644
index 0000000000000..3251230d895c1
--- /dev/null
+++ b/mlir/test/lib/Dialect/SCF/TestSplitForOpAtPoint.cpp
@@ -0,0 +1,72 @@
+//===- TestSplitForOpAtPoint.cpp - Pass to test splitForOpAtPoint ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to test splitForOpAtPoint.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SCF/Utils/Utils.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+
+static constexpr StringLiteral kSplitAtAttr = "test.split_at";
+
+struct TestSplitForOpAtPointPass
+ : public PassWrapper<TestSplitForOpAtPointPass,
+ OperationPass<func::FuncOp>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSplitForOpAtPointPass)
+
+ StringRef getArgument() const final { return "test-split-for-op-at-point"; }
+
+ StringRef getDescription() const final { return "test splitForOpAtPoint"; }
+
+ TestSplitForOpAtPointPass() = default;
+ TestSplitForOpAtPointPass(const TestSplitForOpAtPointPass &) {}
+
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<arith::ArithDialect>();
+ }
+
+ void runOnOperation() override {
+ func::FuncOp func = getOperation();
+ SmallVector<scf::ForOp> loopsToSplit;
+ func.walk([&](scf::ForOp forOp) {
+ if (forOp->hasAttr(kSplitAtAttr))
+ loopsToSplit.push_back(forOp);
+ });
+
+ for (scf::ForOp forOp : loopsToSplit) {
+ auto splitAttr = forOp->getAttrOfType<IntegerAttr>(kSplitAtAttr);
+ OpBuilder builder(forOp);
+ Value splitPoint =
+ arith::ConstantOp::create(builder, forOp.getLoc(), splitAttr);
+ forOp->removeAttr(kSplitAtAttr);
+ if (failed(splitForOpAtPoint(forOp, splitPoint))) {
+ emitError(forOp.getLoc(), "failed to split scf.for");
+ return signalPassFailure();
+ }
+ }
+ }
+};
+
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestSplitForOpAtPointPass() {
+ PassRegistration<TestSplitForOpAtPointPass>();
+}
+} // namespace test
+} // namespace mlir
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