[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 19 02:03:48 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 1/2] [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();
>From 056df405262ae0c1657c8b2deacd39396b770031 Mon Sep 17 00:00:00 2001
From: David Lerner <davidlerner96 at gmail.com>
Date: Wed, 19 Aug 2026 11:59:55 +0300
Subject: [PATCH 2/2] [mlir][scf] Check split points and return both loops from
splitForOpAtPoint
- Clone two new loops instead of mutating the original scf.for; return
{first, second}, chain iter-args, replace uses, and erase the input.
- Validate lb <= split < ub, step > 0, and lattice alignment
(split == lb + k * step). Fail statically when values are constant;
otherwise emit arith.cmpi / cf.assert.
- Use signed or unsigned compares according to unsignedCmp.
- Add getConstantUIntValue for zero-extended constant integers.
- Link MLIRControlFlowDialect from SCF Utils.
- Update loopUnrollByFactor to unroll the first returned loop and keep
the second as the epilogue.
- Fold the split test pass into TestSCFUtils; support test.split_at and
test.split_arg; register ControlFlow on the SCF test passes.
- Remove the standalone TestSplitForOpAtPoint pass and its mlir-opt
registration.
- Expand split-for-op-at-point.mlir for dynamic bounds, unsigned and
integer IVs, invalid range/step, and split at the lower bound.
- Update loop-unroll.mlir for runtime asserts on dynamic unroll and a
constant zero-step case.
Co-authored-by: Cursor <cursoragent at cursor.com>
---
mlir/include/mlir/Dialect/SCF/Utils/Utils.h | 11 +-
.../mlir/Dialect/Utils/StaticValueUtils.h | 3 +
mlir/lib/Dialect/SCF/Utils/CMakeLists.txt | 1 +
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 135 ++++---
mlir/lib/Dialect/Utils/StaticValueUtils.cpp | 9 +
mlir/test/Dialect/SCF/loop-unroll.mlir | 44 +++
.../Dialect/SCF/split-for-op-at-point.mlir | 328 ++++++++++++++++++
mlir/test/lib/Dialect/SCF/CMakeLists.txt | 2 +-
.../lib/Dialect/SCF/TestLoopUnrolling.cpp | 3 +-
mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp | 57 +++
.../lib/Dialect/SCF/TestSplitForOpAtPoint.cpp | 72 ----
mlir/tools/mlir-opt/mlir-opt.cpp | 2 -
12 files changed, 547 insertions(+), 120 deletions(-)
delete 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 c5c6871bd758e..809cd77c03ef6 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -19,6 +19,7 @@
#include "llvm/ADT/STLExtras.h"
#include <optional>
#include <tuple>
+#include <utility>
namespace mlir {
class Location;
@@ -112,11 +113,11 @@ struct UnrolledLoopInfo {
/// 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);
+/// Returns the two new loops and replaces `forOp`. Iter-args are chained
+/// from the first loop to the second. The split point is checked statically
+/// when bounds are constant, and with runtime asserts otherwise.
+FailureOr<std::pair<scf::ForOp, 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
diff --git a/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h b/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
index 511e9c5c2c76a..a0ceee6e543eb 100644
--- a/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
@@ -117,6 +117,9 @@ SmallVector<OpFoldResult> getAsIndexOpFoldResult(MLIRContext *ctx,
std::optional<std::pair<APInt, bool>> getConstantAPIntValue(OpFoldResult ofr);
/// If ofr is a constant integer or an IntegerAttr, return the integer.
std::optional<int64_t> getConstantIntValue(OpFoldResult ofr);
+/// If ofr is a constant integer or an IntegerAttr, return the integer
+/// zero-extended to 64 bits.
+std::optional<uint64_t> getConstantUIntValue(OpFoldResult ofr);
/// If all ofrs are constant integers or IntegerAttrs, return the integers.
std::optional<SmallVector<int64_t>>
getConstantIntValues(ArrayRef<OpFoldResult> ofrs);
diff --git a/mlir/lib/Dialect/SCF/Utils/CMakeLists.txt b/mlir/lib/Dialect/SCF/Utils/CMakeLists.txt
index 5ac6ed50659e7..f690300edce77 100644
--- a/mlir/lib/Dialect/SCF/Utils/CMakeLists.txt
+++ b/mlir/lib/Dialect/SCF/Utils/CMakeLists.txt
@@ -10,6 +10,7 @@ add_mlir_dialect_library(MLIRSCFUtils
MLIRAffineAnalysis
MLIRAnalysis
MLIRArithDialect
+ MLIRControlFlowDialect
MLIRDialectUtils
MLIRFuncDialect
MLIRIR
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index c266a2de0e54d..f46e1163aeede 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -15,6 +15,7 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/IRMapping.h"
@@ -364,55 +365,109 @@ void mlir::generateUnrolledLoop(
}
/// Splits `forOp` into two consecutive loops at `splitPoint`.
-FailureOr<scf::ForOp> mlir::splitForOpAtPoint(scf::ForOp forOp,
- Value splitPoint) {
+FailureOr<std::pair<scf::ForOp, 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();
- }
+ // When a bound is constant, require a valid lattice-aligned split point.
+ // When it is dynamic, emit the same checks as runtime asserts.
+ OpBuilder runtimeBuilder(forOp);
+ Location loc = forOp.getLoc();
+ bool isUnsigned = forOp.getUnsignedCmp();
+ Value lbVal = forOp.getLowerBound();
+ Value ubVal = forOp.getUpperBound();
+ Value stepVal = forOp.getStep();
+ // Fail at runtime if `cond` is false.
+ auto emitAssert = [&](Value cond, StringRef msg) {
+ forOp->getContext()->getOrLoadDialect<cf::ControlFlowDialect>();
+ cf::AssertOp::create(runtimeBuilder, loc, cond,
+ runtimeBuilder.getStringAttr(msg));
+ };
+ // Compare according to the loop's signedness.
+ auto emitICmp = [&](arith::CmpIPredicate signedPred,
+ arith::CmpIPredicate unsignedPred, Value lhs, Value rhs) {
+ return arith::CmpIOp::create(
+ runtimeBuilder, loc, isUnsigned ? unsignedPred : signedPred, lhs, rhs);
+ };
+ // Check the split point statically when constant, otherwise at runtime.
+ auto checkSplitPoint = [&](auto getBound) -> LogicalResult {
+ auto lb = getBound(lbVal);
+ auto ub = getBound(ubVal);
+ auto step = getBound(stepVal);
+ auto split = getBound(splitPoint);
+ // lowerBound <= splitPoint
+ if (lb && split) {
+ if (*lb > *split)
+ return failure();
+ } else {
+ emitAssert(emitICmp(arith::CmpIPredicate::sle, arith::CmpIPredicate::ule,
+ lbVal, splitPoint),
+ "splitForOpAtPoint: split point below lower bound");
+ }
+ // splitPoint < upperBound
+ if (split && ub) {
+ if (*split >= *ub)
+ return failure();
+ } else {
+ emitAssert(emitICmp(arith::CmpIPredicate::slt, arith::CmpIPredicate::ult,
+ splitPoint, ubVal),
+ "splitForOpAtPoint: split point not below upper bound");
+ }
+ // step > 0
+ if (step) {
+ if (*step <= 0)
+ return failure();
+ } else {
+ Value zero = arith::ConstantOp::create(
+ runtimeBuilder, loc,
+ runtimeBuilder.getIntegerAttr(stepVal.getType(), 0));
+ emitAssert(emitICmp(arith::CmpIPredicate::sgt, arith::CmpIPredicate::ugt,
+ stepVal, zero),
+ "splitForOpAtPoint: step must be positive");
+ }
+ // splitPoint == lowerBound + k * step
+ if (lb && step && split) {
+ if ((*split - *lb) % *step != 0)
+ return failure();
+ } else {
+ Value zero = arith::ConstantOp::create(
+ runtimeBuilder, loc,
+ runtimeBuilder.getIntegerAttr(stepVal.getType(), 0));
+ Value diff =
+ arith::SubIOp::create(runtimeBuilder, loc, splitPoint, lbVal);
+ Value rem = isUnsigned ? static_cast<Value>(arith::RemUIOp::create(
+ runtimeBuilder, loc, diff, stepVal))
+ : static_cast<Value>(arith::RemSIOp::create(
+ runtimeBuilder, loc, diff, stepVal));
+ emitAssert(arith::CmpIOp::create(runtimeBuilder, loc,
+ arith::CmpIPredicate::eq, rem, zero),
+ "splitForOpAtPoint: split point is not lattice-aligned");
+ }
+ return success();
+ };
+ if (failed(isUnsigned ? checkSplitPoint(getConstantUIntValue)
+ : checkSplitPoint(getConstantIntValue)))
+ return failure();
OpBuilder builder(forOp->getContext());
builder.setInsertionPointAfter(forOp);
- // Clone before mutating `forOp` so the second loop inherits the original
- // body and bounds.
+ auto firstForOp = cast<scf::ForOp>(builder.clone(*forOp));
auto secondForOp = cast<scf::ForOp>(builder.clone(*forOp));
+ firstForOp.setUpperBound(splitPoint);
secondForOp.setLowerBound(splitPoint);
// Chain iter-args across the split:
- // - `secondForOp` is initialized from `forOp`'s results.
+ // - `secondForOp` is initialized from `firstForOp`'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);
+ secondForOp.getInitArgs().size(),
+ firstForOp.getResults());
+ forOp->replaceAllUsesWith(secondForOp.getResults());
+ forOp.erase();
- return secondForOp;
+ return std::pair<scf::ForOp, scf::ForOp>{firstForOp, secondForOp};
}
/// Unrolls 'forOp' by 'unrollFactor', returns the unrolled main loop and the
@@ -525,11 +580,13 @@ FailureOr<UnrolledLoopInfo> mlir::loopUnrollByFactor(
// Create epilogue clean up loop starting at 'upperBoundUnrolled'.
if (generateEpilogueLoop) {
- auto epilogueForOp = splitForOpAtPoint(forOp, upperBoundUnrolled);
- if (failed(epilogueForOp))
+ auto splitLoops = splitForOpAtPoint(forOp, upperBoundUnrolled);
+ if (failed(splitLoops))
return failure();
- if (epilogueForOp->promoteIfSingleIteration(rewriter).failed())
- resultLoops.epilogueLoopOp = *epilogueForOp;
+ forOp = splitLoops->first;
+ scf::ForOp epilogueForOp = splitLoops->second;
+ if (epilogueForOp.promoteIfSingleIteration(rewriter).failed())
+ resultLoops.epilogueLoopOp = epilogueForOp;
} else {
forOp.setUpperBound(upperBoundUnrolled);
}
diff --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index 4e30693353440..3f66d86e48991 100644
--- a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
@@ -152,6 +152,15 @@ std::optional<int64_t> getConstantIntValue(OpFoldResult ofr) {
return apInt->first.getSExtValue();
}
+/// If ofr is a constant integer or an IntegerAttr, return the integer
+/// zero-extended to 64 bits.
+std::optional<uint64_t> getConstantUIntValue(OpFoldResult ofr) {
+ std::optional<std::pair<APInt, bool>> apInt = getConstantAPIntValue(ofr);
+ if (!apInt)
+ return std::nullopt;
+ return apInt->first.getZExtValue();
+}
+
std::optional<SmallVector<int64_t>>
getConstantIntValues(ArrayRef<OpFoldResult> ofrs) {
SmallVector<int64_t> res;
diff --git a/mlir/test/Dialect/SCF/loop-unroll.mlir b/mlir/test/Dialect/SCF/loop-unroll.mlir
index f764013ed50f9..8000161509911 100644
--- a/mlir/test/Dialect/SCF/loop-unroll.mlir
+++ b/mlir/test/Dialect/SCF/loop-unroll.mlir
@@ -36,6 +36,17 @@ func.func @dynamic_loop_unroll(%arg0 : index, %arg1 : index, %arg2 : index,
// UNROLL-BY-2-DAG: %[[V7:.*]] = arith.addi %[[LB]], %[[V6]] : index
// Compute step of unrolled loop in V8.
// UNROLL-BY-2-DAG: %[[V8:.*]] = arith.muli %[[STEP]], %[[C2]] : index
+// Runtime checks from splitForOpAtPoint (dynamic split = lb + evenMult * step).
+// UNROLL-BY-2: %[[LB_OK:.*]] = arith.cmpi sle, %[[LB]], %[[V7]]
+// UNROLL-BY-2: cf.assert %[[LB_OK]]
+// UNROLL-BY-2: %[[UB_OK:.*]] = arith.cmpi slt, %[[V7]], %[[UB]]
+// UNROLL-BY-2: cf.assert %[[UB_OK]]
+// UNROLL-BY-2: arith.cmpi sgt, %[[STEP]]
+// UNROLL-BY-2: cf.assert
+// UNROLL-BY-2: %[[SPLIT_DIFF:.*]] = arith.subi %[[V7]], %[[LB]]
+// UNROLL-BY-2: %[[SPLIT_REM:.*]] = arith.remsi %[[SPLIT_DIFF]], %[[STEP]]
+// UNROLL-BY-2: %[[ALIGNED:.*]] = arith.cmpi eq, %[[SPLIT_REM]]
+// UNROLL-BY-2: cf.assert %[[ALIGNED]]
// UNROLL-BY-2: scf.for %[[IV:.*]] = %[[LB]] to %[[V7]] step %[[V8]] {
// UNROLL-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>
// UNROLL-BY-2-NEXT: %[[C1_IV:.*]] = arith.constant 1 : index
@@ -69,6 +80,17 @@ func.func @dynamic_loop_unroll(%arg0 : index, %arg1 : index, %arg2 : index,
// UNROLL-BY-3-DAG: %[[V7:.*]] = arith.addi %[[LB]], %[[V6]] : index
// Compute step of unrolled loop in V8.
// UNROLL-BY-3-DAG: %[[V8:.*]] = arith.muli %[[STEP]], %[[C3]] : index
+// Runtime checks from splitForOpAtPoint (dynamic split = lb + evenMult * step).
+// UNROLL-BY-3: %[[LB_OK:.*]] = arith.cmpi sle, %[[LB]], %[[V7]]
+// UNROLL-BY-3: cf.assert %[[LB_OK]]
+// UNROLL-BY-3: %[[UB_OK:.*]] = arith.cmpi slt, %[[V7]], %[[UB]]
+// UNROLL-BY-3: cf.assert %[[UB_OK]]
+// UNROLL-BY-3: arith.cmpi sgt, %[[STEP]]
+// UNROLL-BY-3: cf.assert
+// UNROLL-BY-3: %[[SPLIT_DIFF:.*]] = arith.subi %[[V7]], %[[LB]]
+// UNROLL-BY-3: %[[SPLIT_REM:.*]] = arith.remsi %[[SPLIT_DIFF]], %[[STEP]]
+// UNROLL-BY-3: %[[ALIGNED:.*]] = arith.cmpi eq, %[[SPLIT_REM]]
+// UNROLL-BY-3: cf.assert %[[ALIGNED]]
// UNROLL-BY-3: scf.for %[[IV:.*]] = %[[LB]] to %[[V7]] step %[[V8]] {
// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>
// UNROLL-BY-3-NEXT: %[[C1_IV:.*]] = arith.constant 1 : index
@@ -173,6 +195,7 @@ func.func @static_loop_unroll_by_2(%arg0 : memref<?xf32>) {
// UNROLL-BY-2-DAG: %[[C1:.*]] = arith.constant 1 : index
// UNROLL-BY-2-DAG: %[[C20:.*]] = arith.constant 20 : index
// UNROLL-BY-2-DAG: %[[C2:.*]] = arith.constant 2 : index
+// UNROLL-BY-2-NOT: cf.assert
// UNROLL-BY-2: scf.for %[[IV:.*]] = %[[C0]] to %[[C20]] step %[[C2]] {
// UNROLL-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>
// UNROLL-BY-2-NEXT: %[[C1_IV:.*]] = arith.constant 1 : index
@@ -237,6 +260,7 @@ func.func @static_loop_unroll_by_3(%arg0 : memref<?xf32>) {
// UNROLL-BY-3-DAG: %[[C20:.*]] = arith.constant 20 : index
// UNROLL-BY-3-DAG: %[[C18:.*]] = arith.constant 18 : index
// UNROLL-BY-3-DAG: %[[C3:.*]] = arith.constant 3 : index
+// UNROLL-BY-3-NOT: cf.assert
// UNROLL-BY-3: scf.for %[[IV:.*]] = %[[C0]] to %[[C18]] step %[[C3]] {
// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>
// UNROLL-BY-3-NEXT: %[[C1_IV:.*]] = arith.constant 1 : index
@@ -660,3 +684,23 @@ func.func @unroll_unsigned_i2_step2_bug2() -> (i32, i32) {
// UNROLL-BY-2: arith.addi
// UNROLL-BY-2: arith.muli
// UNROLL-BY-2: return
+
+// -----
+
+// Dynamic bounds with a constant zero step: splitForOpAtPoint fails the static
+// step check, so unrolling does not rewrite the loop.
+func.func @dynamic_unroll_zero_step(%lb: index, %ub: index,
+ %mem: memref<?xf32>) {
+ %0 = arith.constant 7.0 : f32
+ %step = arith.constant 0 : index
+ scf.for %i0 = %lb to %ub step %step {
+ memref.store %0, %mem[%i0] : memref<?xf32>
+ }
+ return
+}
+// UNROLL-BY-2-LABEL: func @dynamic_unroll_zero_step
+// UNROLL-BY-2-SAME: %[[LB:.*]]: index, %[[UB:.*]]: index
+// UNROLL-BY-2: scf.for %{{.*}} = %[[LB]] to %[[UB]] step %{{.*}}
+// UNROLL-BY-2-NEXT: memref.store
+// UNROLL-BY-2-NEXT: }
+// UNROLL-BY-2-NEXT: return
diff --git a/mlir/test/Dialect/SCF/split-for-op-at-point.mlir b/mlir/test/Dialect/SCF/split-for-op-at-point.mlir
index ca6e6cc0139a7..41055429507e1 100644
--- a/mlir/test/Dialect/SCF/split-for-op-at-point.mlir
+++ b/mlir/test/Dialect/SCF/split-for-op-at-point.mlir
@@ -51,3 +51,331 @@ func.func @invalid_split(%mem: memref<?xf32>) {
} {test.split_at = 10 : index}
return
}
+
+// -----
+
+// Split point must be lowerBound + k * step.
+func.func @invalid_split_not_multiple(%mem: memref<?xf32>) {
+ %cst = arith.constant 0.0 : f32
+ %c0 = arith.constant 0 : index
+ %c10 = arith.constant 10 : index
+ %c3 = arith.constant 3 : index
+ // expected-error @+1 {{failed to split scf.for}}
+ scf.for %i = %c0 to %c10 step %c3 {
+ memref.store %cst, %mem[%i] : memref<?xf32>
+ } {test.split_at = 8 : index}
+ return
+}
+
+// -----
+
+// Dynamic upper bound: split happens and `split < ub` is checked at runtime.
+func.func @dynamic_ub_split(%mem: memref<?xf32>, %ub: index) {
+ %cst = arith.constant 0.0 : f32
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ scf.for %i = %c0 to %ub step %c1 {
+ memref.store %cst, %mem[%i] : memref<?xf32>
+ } {test.split_at = 9 : index}
+ return
+}
+// CHECK-LABEL: func @dynamic_ub_split
+// CHECK-SAME: %[[MEM:.*]]: memref<?xf32>, %[[UB:.*]]: index
+// CHECK: %[[C9:.*]] = arith.constant 9 : index
+// CHECK: %[[CMP:.*]] = arith.cmpi slt, %[[C9]], %[[UB]]
+// CHECK: cf.assert %[[CMP]]
+// CHECK: scf.for %{{.*}} = %c0 to %[[C9]] step %c1
+// CHECK-NEXT: memref.store
+// CHECK: scf.for %{{.*}} = %[[C9]] to %[[UB]] step %c1
+// CHECK-NEXT: memref.store
+
+// -----
+
+// Dynamic step: split happens and lattice alignment is checked at runtime.
+func.func @dynamic_step_split(%mem: memref<?xf32>, %step: index) {
+ %cst = arith.constant 0.0 : f32
+ %c0 = arith.constant 0 : index
+ %c12 = arith.constant 12 : index
+ scf.for %i = %c0 to %c12 step %step {
+ memref.store %cst, %mem[%i] : memref<?xf32>
+ } {test.split_at = 6 : index}
+ return
+}
+// CHECK-LABEL: func @dynamic_step_split
+// CHECK-SAME: %[[MEM:.*]]: memref<?xf32>, %[[STEP:.*]]: index
+// CHECK: %[[C6:.*]] = arith.constant 6 : index
+// CHECK: arith.cmpi sgt, %[[STEP]]
+// CHECK: cf.assert
+// CHECK: %[[DIFF:.*]] = arith.subi %[[C6]], %c0
+// CHECK: %[[REM:.*]] = arith.remsi %[[DIFF]], %[[STEP]]
+// CHECK: %[[ALIGNED:.*]] = arith.cmpi eq, %[[REM]]
+// CHECK: cf.assert %[[ALIGNED]]
+// CHECK: scf.for %{{.*}} = %c0 to %[[C6]] step %[[STEP]]
+// CHECK: scf.for %{{.*}} = %[[C6]] to %c12 step %[[STEP]]
+
+// -----
+
+// Split at the lower bound (k = 0): first loop is empty, second keeps the range.
+func.func @split_at_lower_bound(%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 = 0 : index}
+ return
+}
+// CHECK-LABEL: func @split_at_lower_bound
+// CHECK: %[[LB:.*]] = arith.constant 0 : index
+// CHECK: %[[UB:.*]] = arith.constant 10 : index
+// CHECK: %[[STEP:.*]] = arith.constant 1 : index
+// CHECK: %[[SPLIT:.*]] = arith.constant 0 : index
+// CHECK: scf.for %{{.*}} = %[[LB]] to %[[SPLIT]] step %[[STEP]]
+// CHECK: scf.for %{{.*}} = %[[SPLIT]] to %[[UB]] step %[[STEP]]
+
+// -----
+
+// Constant split point below the lower bound is rejected.
+func.func @invalid_split_below_lb(%mem: memref<?xf32>) {
+ %cst = arith.constant 0.0 : f32
+ %c5 = arith.constant 5 : index
+ %c10 = arith.constant 10 : index
+ %c1 = arith.constant 1 : index
+ // expected-error @+1 {{failed to split scf.for}}
+ scf.for %i = %c5 to %c10 step %c1 {
+ memref.store %cst, %mem[%i] : memref<?xf32>
+ } {test.split_at = 3 : index}
+ return
+}
+
+// -----
+
+// Constant non-positive step is rejected.
+func.func @invalid_zero_step(%mem: memref<?xf32>) {
+ %cst = arith.constant 0.0 : f32
+ %c0 = arith.constant 0 : index
+ %c10 = arith.constant 10 : index
+ // expected-error @+1 {{failed to split scf.for}}
+ scf.for %i = %c0 to %c10 step %c0 {
+ memref.store %cst, %mem[%i] : memref<?xf32>
+ } {test.split_at = 0 : index}
+ return
+}
+
+// -----
+
+func.func @invalid_negative_step(%mem: memref<?xf32>) {
+ %cst = arith.constant 0.0 : f32
+ %c0 = arith.constant 0 : index
+ %c10 = arith.constant 10 : index
+ %cm1 = arith.constant -1 : index
+ // expected-error @+1 {{failed to split scf.for}}
+ scf.for %i = %c0 to %c10 step %cm1 {
+ memref.store %cst, %mem[%i] : memref<?xf32>
+ } {test.split_at = 0 : index}
+ return
+}
+
+// -----
+
+// Integer (non-index) induction variable.
+func.func @split_integer_iv() -> i32 {
+ %c0 = arith.constant 0 : i32
+ %c10 = arith.constant 10 : i32
+ %c1 = arith.constant 1 : i32
+ %init = arith.constant 0 : i32
+ %r = scf.for %i = %c0 to %c10 step %c1 iter_args(%acc = %init) -> i32 : i32 {
+ %one = arith.constant 1 : i32
+ %add = arith.addi %acc, %one : i32
+ scf.yield %add : i32
+ } {test.split_at = 9 : i32}
+ return %r : i32
+}
+// CHECK-LABEL: func @split_integer_iv
+// CHECK: %[[FIRST:.*]] = scf.for %{{.*}} = %{{.*}} to %c9_i32 step %{{.*}} iter_args(%{{.*}} = %{{.*}}) -> (i32)
+// CHECK: %[[RESULT:.*]] = scf.for %{{.*}} = %c9_i32 to %c10_i32 step %{{.*}} iter_args(%{{.*}} = %[[FIRST]]) -> (i32)
+// CHECK: return %[[RESULT]] : i32
+
+// -----
+
+// Unsigned i32 split uses unsigned compares when a bound is dynamic.
+func.func @split_unsigned_i32() -> i32 {
+ %c0 = arith.constant 0 : i32
+ %c10 = arith.constant 10 : i32
+ %c1 = arith.constant 1 : i32
+ %init = arith.constant 0 : i32
+ %r = scf.for unsigned %i = %c0 to %c10 step %c1
+ iter_args(%acc = %init) -> i32 : i32 {
+ %one = arith.constant 1 : i32
+ %add = arith.addi %acc, %one : i32
+ scf.yield %add : i32
+ } {test.split_at = 9 : i32}
+ return %r : i32
+}
+// CHECK-LABEL: func @split_unsigned_i32
+// CHECK: %[[FIRST:.*]] = scf.for unsigned %{{.*}} = %{{.*}} to %c9_i32 step %{{.*}}
+// CHECK: %[[RESULT:.*]] = scf.for unsigned %{{.*}} = %c9_i32 to %c10_i32 step %{{.*}}
+// CHECK: return %[[RESULT]] : i32
+
+// -----
+
+// Narrow unsigned IV: split at 4 in [0, 5) as i3. Sign-extending 4:i3 is -4
+// and would reject this split; zero-extension keeps it in range.
+func.func @split_unsigned_i3() -> i32 {
+ %c0 = arith.constant 0 : i3
+ %c5 = arith.constant 5 : i3
+ %c1 = arith.constant 1 : i3
+ %init = arith.constant 0 : i32
+ %r = scf.for unsigned %i = %c0 to %c5 step %c1
+ iter_args(%acc = %init) -> i32 : i3 {
+ %one = arith.constant 1 : i32
+ %add = arith.addi %acc, %one : i32
+ scf.yield %add : i32
+ } {test.split_at = 4 : i3}
+ return %r : i32
+}
+// CHECK-LABEL: func @split_unsigned_i3
+// CHECK: %[[FIRST:.*]] = scf.for unsigned %{{.*}} = %{{.*}} to %c-4_i3 step %{{.*}}
+// CHECK: %[[RESULT:.*]] = scf.for unsigned %{{.*}} = %c-4_i3 to %{{.*}} step %{{.*}}
+// CHECK: return %[[RESULT]] : i32
+
+// -----
+
+// Dynamic unsigned upper bound: `split < ub` is an unsigned compare.
+func.func @unsigned_dynamic_ub(%ub: i32) -> i32 {
+ %c0 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : i32
+ %init = arith.constant 0 : i32
+ %r = scf.for unsigned %i = %c0 to %ub step %c1
+ iter_args(%acc = %init) -> i32 : i32 {
+ %one = arith.constant 1 : i32
+ %add = arith.addi %acc, %one : i32
+ scf.yield %add : i32
+ } {test.split_at = 9 : i32}
+ return %r : i32
+}
+// CHECK-LABEL: func @unsigned_dynamic_ub
+// CHECK-SAME: %[[UB:.*]]: i32
+// CHECK: %[[C9:.*]] = arith.constant 9 : i32
+// CHECK: %[[CMP:.*]] = arith.cmpi ult, %[[C9]], %[[UB]]
+// CHECK: cf.assert %[[CMP]]
+// CHECK: scf.for unsigned %{{.*}} = %{{.*}} to %[[C9]]
+// CHECK: scf.for unsigned %{{.*}} = %[[C9]] to %[[UB]]
+
+// -----
+
+// Dynamic unsigned step: `step > 0` and lattice use unsigned ops.
+func.func @unsigned_dynamic_step(%step: i32) -> i32 {
+ %c0 = arith.constant 0 : i32
+ %c12 = arith.constant 12 : i32
+ %init = arith.constant 0 : i32
+ %r = scf.for unsigned %i = %c0 to %c12 step %step
+ iter_args(%acc = %init) -> i32 : i32 {
+ %one = arith.constant 1 : i32
+ %add = arith.addi %acc, %one : i32
+ scf.yield %add : i32
+ } {test.split_at = 6 : i32}
+ return %r : i32
+}
+// CHECK-LABEL: func @unsigned_dynamic_step
+// CHECK-SAME: %[[STEP:.*]]: i32
+// CHECK: %[[C6:.*]] = arith.constant 6 : i32
+// CHECK: arith.cmpi ugt, %[[STEP]]
+// CHECK: cf.assert
+// CHECK: %[[DIFF:.*]] = arith.subi %[[C6]], %c0_i32
+// CHECK: %[[REM:.*]] = arith.remui %[[DIFF]], %[[STEP]]
+// CHECK: %[[ALIGNED:.*]] = arith.cmpi eq, %[[REM]]
+// CHECK: cf.assert %[[ALIGNED]]
+// CHECK: scf.for unsigned %{{.*}} = %{{.*}} to %[[C6]] step %[[STEP]]
+// CHECK: scf.for unsigned %{{.*}} = %[[C6]] to %{{.*}} step %[[STEP]]
+
+// -----
+
+// Dynamic lower bound: `lb <= split` and lattice alignment are runtime checks.
+func.func @dynamic_lb_split(%mem: memref<?xf32>, %lb: index) {
+ %cst = arith.constant 0.0 : f32
+ %c10 = arith.constant 10 : index
+ %c1 = arith.constant 1 : index
+ scf.for %i = %lb to %c10 step %c1 {
+ memref.store %cst, %mem[%i] : memref<?xf32>
+ } {test.split_at = 9 : index}
+ return
+}
+// CHECK-LABEL: func @dynamic_lb_split
+// CHECK-SAME: %[[MEM:.*]]: memref<?xf32>, %[[LB:.*]]: index
+// CHECK: %[[C9:.*]] = arith.constant 9 : index
+// CHECK: %[[CMP:.*]] = arith.cmpi sle, %[[LB]], %[[C9]]
+// CHECK: cf.assert %[[CMP]]
+// CHECK: %[[DIFF:.*]] = arith.subi %[[C9]], %[[LB]]
+// CHECK: %[[REM:.*]] = arith.remsi %[[DIFF]], %c1
+// CHECK: %[[ALIGNED:.*]] = arith.cmpi eq, %[[REM]]
+// CHECK: cf.assert %[[ALIGNED]]
+// CHECK: scf.for %{{.*}} = %[[LB]] to %[[C9]] step %c1
+// CHECK: scf.for %{{.*}} = %[[C9]] to %c10 step %c1
+
+// -----
+
+// Fully dynamic bounds: every check is a runtime assert.
+func.func @dynamic_all_split(%mem: memref<?xf32>, %lb: index, %ub: index,
+ %step: index) {
+ %cst = arith.constant 0.0 : f32
+ scf.for %i = %lb to %ub step %step {
+ memref.store %cst, %mem[%i] : memref<?xf32>
+ } {test.split_at = 9 : index}
+ return
+}
+// CHECK-LABEL: func @dynamic_all_split
+// CHECK-SAME: %[[MEM:.*]]: memref<?xf32>, %[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index
+// CHECK: %[[C9:.*]] = arith.constant 9 : index
+// CHECK: arith.cmpi sle, %[[LB]], %[[C9]]
+// CHECK: cf.assert
+// CHECK: arith.cmpi slt, %[[C9]], %[[UB]]
+// CHECK: cf.assert
+// CHECK: arith.cmpi sgt, %[[STEP]]
+// CHECK: cf.assert
+// CHECK: %[[DIFF:.*]] = arith.subi %[[C9]], %[[LB]]
+// CHECK: %[[REM:.*]] = arith.remsi %[[DIFF]], %[[STEP]]
+// CHECK: %[[ALIGNED:.*]] = arith.cmpi eq, %[[REM]]
+// CHECK: cf.assert %[[ALIGNED]]
+// CHECK: scf.for %{{.*}} = %[[LB]] to %[[C9]] step %[[STEP]]
+// CHECK: scf.for %{{.*}} = %[[C9]] to %[[UB]] step %[[STEP]]
+
+// -----
+
+// Dynamic split point (function argument): range and lattice are runtime.
+func.func @dynamic_split_point(%mem: memref<?xf32>, %split: index) {
+ %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_arg = 1 : i64}
+ return
+}
+// CHECK-LABEL: func @dynamic_split_point
+// CHECK-SAME: %[[MEM:.*]]: memref<?xf32>, %[[SPLIT:.*]]: index
+// CHECK: arith.cmpi sle, %c0, %[[SPLIT]]
+// CHECK: cf.assert
+// CHECK: arith.cmpi slt, %[[SPLIT]], %c10
+// CHECK: cf.assert
+// CHECK: %[[DIFF:.*]] = arith.subi %[[SPLIT]], %c0
+// CHECK: %[[REM:.*]] = arith.remsi %[[DIFF]], %c1
+// CHECK: %[[ALIGNED:.*]] = arith.cmpi eq, %[[REM]]
+// CHECK: cf.assert %[[ALIGNED]]
+// CHECK: scf.for %{{.*}} = %c0 to %[[SPLIT]] step %c1
+// CHECK: scf.for %{{.*}} = %[[SPLIT]] to %c10 step %c1
+
+// -----
+
+// Dynamic ub with a constant zero step: the static step check fails.
+func.func @dynamic_ub_zero_step(%mem: memref<?xf32>, %ub: index) {
+ %cst = arith.constant 0.0 : f32
+ %c0 = arith.constant 0 : index
+ // expected-error @+1 {{failed to split scf.for}}
+ scf.for %i = %c0 to %ub step %c0 {
+ memref.store %cst, %mem[%i] : memref<?xf32>
+ } {test.split_at = 0 : index}
+ return
+}
diff --git a/mlir/test/lib/Dialect/SCF/CMakeLists.txt b/mlir/test/lib/Dialect/SCF/CMakeLists.txt
index 80b8268768d46..c9097f1d5a901 100644
--- a/mlir/test/lib/Dialect/SCF/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/SCF/CMakeLists.txt
@@ -4,7 +4,6 @@ add_mlir_library(MLIRSCFTestPasses
TestLoopUnrolling.cpp
TestParallelLoopUnrolling.cpp
TestSCFUtils.cpp
- TestSplitForOpAtPoint.cpp
TestSCFWrapInZeroTripCheck.cpp
TestUpliftWhileToFor.cpp
TestWhileOpBuilder.cpp
@@ -12,6 +11,7 @@ add_mlir_library(MLIRSCFTestPasses
EXCLUDE_FROM_LIBMLIR
)
mlir_target_link_libraries(MLIRSCFTestPasses PUBLIC
+ MLIRControlFlowDialect
MLIRMemRefDialect
MLIRPass
MLIRSCFDialect
diff --git a/mlir/test/lib/Dialect/SCF/TestLoopUnrolling.cpp b/mlir/test/lib/Dialect/SCF/TestLoopUnrolling.cpp
index e2d6996e435bb..5e887d07e393f 100644
--- a/mlir/test/lib/Dialect/SCF/TestLoopUnrolling.cpp
+++ b/mlir/test/lib/Dialect/SCF/TestLoopUnrolling.cpp
@@ -11,6 +11,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/IR/Builders.h"
@@ -49,7 +50,7 @@ struct TestLoopUnrollingPass
}
void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<arith::ArithDialect>();
+ registry.insert<arith::ArithDialect, cf::ControlFlowDialect>();
}
void runOnOperation() override {
diff --git a/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
index fafa03b6c089f..45a93d0fc5f37 100644
--- a/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
+++ b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
@@ -11,6 +11,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
@@ -267,6 +268,61 @@ struct TestSCFPipeliningPass
});
}
};
+
+static constexpr StringLiteral kSplitAtAttr = "test.split_at";
+static constexpr StringLiteral kSplitArgAttr = "test.split_arg";
+
+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, cf::ControlFlowDialect>();
+ }
+
+ void runOnOperation() override {
+ func::FuncOp func = getOperation();
+ SmallVector<scf::ForOp> loopsToSplit;
+ func.walk([&](scf::ForOp forOp) {
+ if (forOp->hasAttr(kSplitAtAttr) || forOp->hasAttr(kSplitArgAttr))
+ loopsToSplit.push_back(forOp);
+ });
+
+ for (scf::ForOp forOp : loopsToSplit) {
+ Value splitPoint;
+ if (auto splitAttr = forOp->getAttrOfType<IntegerAttr>(kSplitAtAttr)) {
+ OpBuilder builder(forOp);
+ splitPoint =
+ arith::ConstantOp::create(builder, forOp.getLoc(), splitAttr);
+ forOp->removeAttr(kSplitAtAttr);
+ } else if (auto argAttr =
+ forOp->getAttrOfType<IntegerAttr>(kSplitArgAttr)) {
+ int64_t argNo = argAttr.getInt();
+ if (argNo < 0 ||
+ static_cast<unsigned>(argNo) >= func.getNumArguments()) {
+ emitError(forOp.getLoc(), "test.split_arg is out of range");
+ return signalPassFailure();
+ }
+ splitPoint = func.getArgument(argNo);
+ forOp->removeAttr(kSplitArgAttr);
+ } else {
+ continue;
+ }
+ if (failed(splitForOpAtPoint(forOp, splitPoint))) {
+ emitError(forOp.getLoc(), "failed to split scf.for");
+ return signalPassFailure();
+ }
+ }
+ }
+};
} // namespace
namespace mlir {
@@ -275,6 +331,7 @@ void registerTestSCFUtilsPass() {
PassRegistration<TestSCFForUtilsPass>();
PassRegistration<TestSCFIfUtilsPass>();
PassRegistration<TestSCFPipeliningPass>();
+ PassRegistration<TestSplitForOpAtPointPass>();
}
} // namespace test
} // namespace mlir
diff --git a/mlir/test/lib/Dialect/SCF/TestSplitForOpAtPoint.cpp b/mlir/test/lib/Dialect/SCF/TestSplitForOpAtPoint.cpp
deleted file mode 100644
index 3251230d895c1..0000000000000
--- a/mlir/test/lib/Dialect/SCF/TestSplitForOpAtPoint.cpp
+++ /dev/null
@@ -1,72 +0,0 @@
-//===- 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 bad235043a096..47dffb1b2d28b 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -143,7 +143,6 @@ void registerTestParallelLoopUnrollingPass();
void registerTestRecursiveTypesPass();
void registerTestSCFUpliftWhileToFor();
void registerTestSCFUtilsPass();
-void registerTestSplitForOpAtPointPass();
void registerTestSCFWhileOpBuilderPass();
void registerTestSCFWrapInZeroTripCheckPasses();
void registerTestShapeMappingPass();
@@ -293,7 +292,6 @@ 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