[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 26 01:27:01 PDT 2026
https://github.com/davidlerner96 updated https://github.com/llvm/llvm-project/pull/215108
>From bcbb17b62b8805152c5695bc028ddd158671b257 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 | 11 +
.../mlir/Dialect/Utils/StaticValueUtils.h | 3 +
mlir/lib/Dialect/SCF/Utils/CMakeLists.txt | 1 +
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 129 +++++-
mlir/lib/Dialect/Utils/StaticValueUtils.cpp | 9 +
mlir/test/Dialect/SCF/loop-unroll.mlir | 40 ++
.../Dialect/SCF/split-for-op-at-point.mlir | 381 ++++++++++++++++++
mlir/test/Transforms/scf-loop-unroll.mlir | 4 +
mlir/test/lib/Dialect/SCF/CMakeLists.txt | 1 +
.../lib/Dialect/SCF/TestLoopUnrolling.cpp | 3 +-
mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp | 57 +++
11 files changed, 623 insertions(+), 16 deletions(-)
create mode 100644 mlir/test/Dialect/SCF/split-for-op-at-point.mlir
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index dd394f7b355b6..47357557f4026 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;
@@ -108,6 +109,16 @@ struct UnrolledLoopInfo {
std::optional<scf::ForOp> epilogueLoopOp = std::nullopt;
};
+/// Splits `forOp` into two consecutive loops at `splitPoint`:
+/// first: [lowerBound, splitPoint)
+/// second: [splitPoint, upperBound)
+///
+/// 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
/// returns failure if the loop cannot be unrolled either due to restrictions or
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 2350f705a7ed4..01210a0cfd744 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"
@@ -363,6 +364,112 @@ void mlir::generateUnrolledLoop(
loopBodyBlock->getTerminator()->setOperands(lastYielded);
}
+/// Splits `forOp` into two consecutive loops at `splitPoint`.
+FailureOr<std::pair<scf::ForOp, scf::ForOp>>
+mlir::splitForOpAtPoint(scf::ForOp forOp, Value splitPoint) {
+ if (splitPoint.getType() != forOp.getLowerBound().getType())
+ 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);
+ 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 `firstForOp`'s results.
+ // - Users of `forOp`'s results are redirected to `secondForOp`'s results,
+ // so downstream code observes the final carried values.
+ secondForOp->setOperands(secondForOp.getNumControlOperands(),
+ secondForOp.getInitArgs().size(),
+ firstForOp.getResults());
+ forOp->replaceAllUsesWith(secondForOp.getResults());
+ forOp.erase();
+
+ return std::pair<scf::ForOp, scf::ForOp>{firstForOp, secondForOp};
+}
+
/// Unrolls 'forOp' by 'unrollFactor', returns the unrolled main loop and the
/// epilogue loop, if the loop is unrolled.
FailureOr<UnrolledLoopInfo> mlir::loopUnrollByFactor(
@@ -484,27 +591,19 @@ 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);
+ auto splitLoops = splitForOpAtPoint(forOp, upperBoundUnrolled);
+ if (failed(splitLoops))
+ return failure();
+ forOp = splitLoops->first;
+ scf::ForOp epilogueForOp = splitLoops->second;
if (!shouldPromoteIfSingleIteration ||
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/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 89d86b09cddfb..c5e2e68d65dbe 100644
--- a/mlir/test/Dialect/SCF/loop-unroll.mlir
+++ b/mlir/test/Dialect/SCF/loop-unroll.mlir
@@ -38,6 +38,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
@@ -71,6 +82,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
@@ -699,7 +721,25 @@ func.func @static_loop_unroll_by_3_no_promote_epilogue(%arg0 : memref<?xf32>) {
// PROMOTE-BY-3-NOT: scf.for
// PROMOTE-BY-3: memref.store
+// -----
+// 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
new file mode 100644
index 0000000000000..41055429507e1
--- /dev/null
+++ b/mlir/test/Dialect/SCF/split-for-op-at-point.mlir
@@ -0,0 +1,381 @@
+// 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
+}
+
+// -----
+
+// 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/Transforms/scf-loop-unroll.mlir b/mlir/test/Transforms/scf-loop-unroll.mlir
index db96c659c49fb..f07104ab68377 100644
--- a/mlir/test/Transforms/scf-loop-unroll.mlir
+++ b/mlir/test/Transforms/scf-loop-unroll.mlir
@@ -38,6 +38,10 @@ func.func @scf_loop_unroll_double_symbolic_ub(%arg0 : f32, %arg1 : f32, %n : ind
// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
// CHECK-NEXT: %[[REM:.*]] = arith.remsi %[[N]], %[[C3]]
// CHECK-NEXT: %[[UB:.*]] = arith.subi %[[N]], %[[REM]]
+ // CHECK-NEXT: %[[LB_OK:.*]] = arith.cmpi sge, %[[UB]], %[[C0]]
+ // CHECK-NEXT: cf.assert %[[LB_OK]]
+ // CHECK-NEXT: %[[UB_OK:.*]] = arith.cmpi slt, %[[UB]], %[[N]]
+ // CHECK-NEXT: cf.assert %[[UB_OK]]
// CHECK-NEXT: %[[SUM:.*]]:2 = scf.for {{.*}} = %[[C0]] to %[[UB]] step %[[C3]] iter_args
// CHECK: }
// CHECK-NEXT: %[[SUM1:.*]]:2 = scf.for {{.*}} = %[[UB]] to %[[N]] step %[[C1]] iter_args(%[[V1:.*]] = %[[SUM]]#0, %[[V2:.*]] = %[[SUM]]#1)
diff --git a/mlir/test/lib/Dialect/SCF/CMakeLists.txt b/mlir/test/lib/Dialect/SCF/CMakeLists.txt
index d2f97e816cc14..c9097f1d5a901 100644
--- a/mlir/test/lib/Dialect/SCF/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/SCF/CMakeLists.txt
@@ -11,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 bbeae9d39db8d..4e91364fe7832 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"
@@ -51,7 +52,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
>From 16dfa240f0ee2885a6c80909faebabcdc334fa61 Mon Sep 17 00:00:00 2001
From: David Lerner <davidlerner96 at gmail.com>
Date: Wed, 26 Aug 2026 11:14:38 +0300
Subject: [PATCH 2/2] [mlir][scf] Generalize splitForOpAtPoint
Make the utility safe for rewrite infrastructure by using a caller-provided rewriter, documenting dynamic preconditions, and rejecting statically invalid split points.
Co-authored-by: Cursor <cursoragent at cursor.com>
---
mlir/include/mlir/Dialect/SCF/Utils/Utils.h | 14 ++-
mlir/lib/Dialect/SCF/Utils/CMakeLists.txt | 1 -
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 102 ++++--------------
mlir/test/Dialect/SCF/loop-unroll.mlir | 22 ----
.../Dialect/SCF/split-for-op-at-point.mlir | 56 ++--------
mlir/test/Transforms/scf-loop-unroll.mlir | 4 -
mlir/test/lib/Dialect/SCF/CMakeLists.txt | 1 -
.../lib/Dialect/SCF/TestLoopUnrolling.cpp | 3 +-
mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp | 16 +--
9 files changed, 51 insertions(+), 168 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index 47357557f4026..4cc7e62e6fd58 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -113,11 +113,17 @@ struct UnrolledLoopInfo {
/// first: [lowerBound, splitPoint)
/// second: [splitPoint, upperBound)
///
-/// 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.
+/// Uses `rewriter` to replace `forOp` and returns the two new loops. Iter-args
+/// are chained from the first loop to the second.
+///
+/// The caller must ensure that `splitPoint` has the same type as the loop
+/// bounds, that the step is positive, and that
+/// `lowerBound <= splitPoint < upperBound`. The split point must also lie on
+/// the loop's iteration lattice: `splitPoint == lowerBound + k * step` for
+/// some non-negative integer `k`. Statically known violations cause failure;
+/// dynamic values are assumed to satisfy these preconditions.
FailureOr<std::pair<scf::ForOp, scf::ForOp>>
-splitForOpAtPoint(scf::ForOp forOp, Value splitPoint);
+splitForOpAtPoint(RewriterBase &rewriter, 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/lib/Dialect/SCF/Utils/CMakeLists.txt b/mlir/lib/Dialect/SCF/Utils/CMakeLists.txt
index f690300edce77..5ac6ed50659e7 100644
--- a/mlir/lib/Dialect/SCF/Utils/CMakeLists.txt
+++ b/mlir/lib/Dialect/SCF/Utils/CMakeLists.txt
@@ -10,7 +10,6 @@ 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 01210a0cfd744..490972a837e96 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -15,7 +15,6 @@
#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"
@@ -366,106 +365,51 @@ void mlir::generateUnrolledLoop(
/// Splits `forOp` into two consecutive loops at `splitPoint`.
FailureOr<std::pair<scf::ForOp, scf::ForOp>>
-mlir::splitForOpAtPoint(scf::ForOp forOp, Value splitPoint) {
+mlir::splitForOpAtPoint(RewriterBase &rewriter, scf::ForOp forOp,
+ Value splitPoint) {
if (splitPoint.getType() != forOp.getLowerBound().getType())
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();
+ // Reject statically known violations of the split preconditions.
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");
- }
+ if ((lb && split && *lb > *split) || (split && ub && *split >= *ub) ||
+ (step && *step <= 0))
+ return failure();
+ if (lb && step && split && (*split - *lb) % *step != 0)
+ return failure();
return success();
};
if (failed(isUnsigned ? checkSplitPoint(getConstantUIntValue)
: checkSplitPoint(getConstantIntValue)))
return failure();
- OpBuilder builder(forOp->getContext());
- builder.setInsertionPointAfter(forOp);
- auto firstForOp = cast<scf::ForOp>(builder.clone(*forOp));
- auto secondForOp = cast<scf::ForOp>(builder.clone(*forOp));
- firstForOp.setUpperBound(splitPoint);
- secondForOp.setLowerBound(splitPoint);
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPointAfter(forOp);
+ auto firstForOp = cast<scf::ForOp>(rewriter.clone(*forOp));
+ auto secondForOp = cast<scf::ForOp>(rewriter.clone(*forOp));
+ rewriter.modifyOpInPlace(firstForOp,
+ [&] { firstForOp.setUpperBound(splitPoint); });
+ rewriter.modifyOpInPlace(secondForOp,
+ [&] { secondForOp.setLowerBound(splitPoint); });
// Chain iter-args across the split:
// - `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.
- secondForOp->setOperands(secondForOp.getNumControlOperands(),
- secondForOp.getInitArgs().size(),
- firstForOp.getResults());
- forOp->replaceAllUsesWith(secondForOp.getResults());
- forOp.erase();
+ rewriter.modifyOpInPlace(secondForOp, [&] {
+ secondForOp->setOperands(secondForOp.getNumControlOperands(),
+ secondForOp.getInitArgs().size(),
+ firstForOp.getResults());
+ });
+ rewriter.replaceOp(forOp, secondForOp.getResults());
return std::pair<scf::ForOp, scf::ForOp>{firstForOp, secondForOp};
}
@@ -591,7 +535,7 @@ FailureOr<UnrolledLoopInfo> mlir::loopUnrollByFactor(
// Create epilogue clean up loop starting at 'upperBoundUnrolled'.
if (generateEpilogueLoop) {
- auto splitLoops = splitForOpAtPoint(forOp, upperBoundUnrolled);
+ auto splitLoops = splitForOpAtPoint(rewriter, forOp, upperBoundUnrolled);
if (failed(splitLoops))
return failure();
forOp = splitLoops->first;
diff --git a/mlir/test/Dialect/SCF/loop-unroll.mlir b/mlir/test/Dialect/SCF/loop-unroll.mlir
index c5e2e68d65dbe..91decbde075c7 100644
--- a/mlir/test/Dialect/SCF/loop-unroll.mlir
+++ b/mlir/test/Dialect/SCF/loop-unroll.mlir
@@ -38,17 +38,6 @@ 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
@@ -82,17 +71,6 @@ 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
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 41055429507e1..0deb8b78c965b 100644
--- a/mlir/test/Dialect/SCF/split-for-op-at-point.mlir
+++ b/mlir/test/Dialect/SCF/split-for-op-at-point.mlir
@@ -69,7 +69,7 @@ func.func @invalid_split_not_multiple(%mem: memref<?xf32>) {
// -----
-// Dynamic upper bound: split happens and `split < ub` is checked at runtime.
+// Dynamic upper bound.
func.func @dynamic_ub_split(%mem: memref<?xf32>, %ub: index) {
%cst = arith.constant 0.0 : f32
%c0 = arith.constant 0 : index
@@ -82,8 +82,6 @@ func.func @dynamic_ub_split(%mem: memref<?xf32>, %ub: index) {
// 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
@@ -91,7 +89,7 @@ func.func @dynamic_ub_split(%mem: memref<?xf32>, %ub: index) {
// -----
-// Dynamic step: split happens and lattice alignment is checked at runtime.
+// Dynamic step.
func.func @dynamic_step_split(%mem: memref<?xf32>, %step: index) {
%cst = arith.constant 0.0 : f32
%c0 = arith.constant 0 : index
@@ -104,12 +102,6 @@ func.func @dynamic_step_split(%mem: memref<?xf32>, %step: index) {
// 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]]
@@ -199,7 +191,7 @@ func.func @split_integer_iv() -> i32 {
// -----
-// Unsigned i32 split uses unsigned compares when a bound is dynamic.
+// Unsigned i32 split.
func.func @split_unsigned_i32() -> i32 {
%c0 = arith.constant 0 : i32
%c10 = arith.constant 10 : i32
@@ -242,7 +234,7 @@ func.func @split_unsigned_i3() -> i32 {
// -----
-// Dynamic unsigned upper bound: `split < ub` is an unsigned compare.
+// Dynamic unsigned upper bound.
func.func @unsigned_dynamic_ub(%ub: i32) -> i32 {
%c0 = arith.constant 0 : i32
%c1 = arith.constant 1 : i32
@@ -258,14 +250,12 @@ func.func @unsigned_dynamic_ub(%ub: i32) -> 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.
+// Dynamic unsigned step.
func.func @unsigned_dynamic_step(%step: i32) -> i32 {
%c0 = arith.constant 0 : i32
%c12 = arith.constant 12 : i32
@@ -281,18 +271,12 @@ func.func @unsigned_dynamic_step(%step: i32) -> 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.
+// Dynamic lower bound.
func.func @dynamic_lb_split(%mem: memref<?xf32>, %lb: index) {
%cst = arith.constant 0.0 : f32
%c10 = arith.constant 10 : index
@@ -305,18 +289,12 @@ func.func @dynamic_lb_split(%mem: memref<?xf32>, %lb: index) {
// 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.
+// Fully dynamic bounds.
func.func @dynamic_all_split(%mem: memref<?xf32>, %lb: index, %ub: index,
%step: index) {
%cst = arith.constant 0.0 : f32
@@ -328,22 +306,12 @@ func.func @dynamic_all_split(%mem: memref<?xf32>, %lb: index, %ub: index,
// 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.
+// Dynamic split point (function argument).
func.func @dynamic_split_point(%mem: memref<?xf32>, %split: index) {
%cst = arith.constant 0.0 : f32
%c0 = arith.constant 0 : index
@@ -356,14 +324,6 @@ func.func @dynamic_split_point(%mem: memref<?xf32>, %split: index) {
}
// 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
diff --git a/mlir/test/Transforms/scf-loop-unroll.mlir b/mlir/test/Transforms/scf-loop-unroll.mlir
index f07104ab68377..db96c659c49fb 100644
--- a/mlir/test/Transforms/scf-loop-unroll.mlir
+++ b/mlir/test/Transforms/scf-loop-unroll.mlir
@@ -38,10 +38,6 @@ func.func @scf_loop_unroll_double_symbolic_ub(%arg0 : f32, %arg1 : f32, %n : ind
// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
// CHECK-NEXT: %[[REM:.*]] = arith.remsi %[[N]], %[[C3]]
// CHECK-NEXT: %[[UB:.*]] = arith.subi %[[N]], %[[REM]]
- // CHECK-NEXT: %[[LB_OK:.*]] = arith.cmpi sge, %[[UB]], %[[C0]]
- // CHECK-NEXT: cf.assert %[[LB_OK]]
- // CHECK-NEXT: %[[UB_OK:.*]] = arith.cmpi slt, %[[UB]], %[[N]]
- // CHECK-NEXT: cf.assert %[[UB_OK]]
// CHECK-NEXT: %[[SUM:.*]]:2 = scf.for {{.*}} = %[[C0]] to %[[UB]] step %[[C3]] iter_args
// CHECK: }
// CHECK-NEXT: %[[SUM1:.*]]:2 = scf.for {{.*}} = %[[UB]] to %[[N]] step %[[C1]] iter_args(%[[V1:.*]] = %[[SUM]]#0, %[[V2:.*]] = %[[SUM]]#1)
diff --git a/mlir/test/lib/Dialect/SCF/CMakeLists.txt b/mlir/test/lib/Dialect/SCF/CMakeLists.txt
index c9097f1d5a901..d2f97e816cc14 100644
--- a/mlir/test/lib/Dialect/SCF/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/SCF/CMakeLists.txt
@@ -11,7 +11,6 @@ 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 4e91364fe7832..bbeae9d39db8d 100644
--- a/mlir/test/lib/Dialect/SCF/TestLoopUnrolling.cpp
+++ b/mlir/test/lib/Dialect/SCF/TestLoopUnrolling.cpp
@@ -11,7 +11,6 @@
//===----------------------------------------------------------------------===//
#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"
@@ -52,7 +51,7 @@ struct TestLoopUnrollingPass
}
void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<arith::ArithDialect, cf::ControlFlowDialect>();
+ registry.insert<arith::ArithDialect>();
}
void runOnOperation() override {
diff --git a/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
index 45a93d0fc5f37..c6de295b6c1bc 100644
--- a/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
+++ b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
@@ -11,7 +11,6 @@
//===----------------------------------------------------------------------===//
#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"
@@ -285,7 +284,7 @@ struct TestSplitForOpAtPointPass
TestSplitForOpAtPointPass(const TestSplitForOpAtPointPass &) {}
void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<arith::ArithDialect, cf::ControlFlowDialect>();
+ registry.insert<arith::ArithDialect>();
}
void runOnOperation() override {
@@ -296,13 +295,15 @@ struct TestSplitForOpAtPointPass
loopsToSplit.push_back(forOp);
});
+ IRRewriter rewriter(func.getContext());
for (scf::ForOp forOp : loopsToSplit) {
Value splitPoint;
if (auto splitAttr = forOp->getAttrOfType<IntegerAttr>(kSplitAtAttr)) {
- OpBuilder builder(forOp);
+ rewriter.setInsertionPoint(forOp);
splitPoint =
- arith::ConstantOp::create(builder, forOp.getLoc(), splitAttr);
- forOp->removeAttr(kSplitAtAttr);
+ arith::ConstantOp::create(rewriter, forOp.getLoc(), splitAttr);
+ rewriter.modifyOpInPlace(forOp,
+ [&] { forOp->removeAttr(kSplitAtAttr); });
} else if (auto argAttr =
forOp->getAttrOfType<IntegerAttr>(kSplitArgAttr)) {
int64_t argNo = argAttr.getInt();
@@ -312,11 +313,12 @@ struct TestSplitForOpAtPointPass
return signalPassFailure();
}
splitPoint = func.getArgument(argNo);
- forOp->removeAttr(kSplitArgAttr);
+ rewriter.modifyOpInPlace(forOp,
+ [&] { forOp->removeAttr(kSplitArgAttr); });
} else {
continue;
}
- if (failed(splitForOpAtPoint(forOp, splitPoint))) {
+ if (failed(splitForOpAtPoint(rewriter, forOp, splitPoint))) {
emitError(forOp.getLoc(), "failed to split scf.for");
return signalPassFailure();
}
More information about the Mlir-commits
mailing list