[Mlir-commits] [mlir] d31eaf1 - [mlir] Add splitForOpAtBound utility and use it in loop unrolling (#215108)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Aug 26 01:43:32 PDT 2026
Author: davidlerner96
Date: 2026-08-26T08:43:27Z
New Revision: d31eaf1864d7d74997c31d03ecd2fb5d8b5e3194
URL: https://github.com/llvm/llvm-project/commit/d31eaf1864d7d74997c31d03ecd2fb5d8b5e3194
DIFF: https://github.com/llvm/llvm-project/commit/d31eaf1864d7d74997c31d03ecd2fb5d8b5e3194.diff
LOG: [mlir] Add splitForOpAtBound utility and use it in loop unrolling (#215108)
## Summary
Adds `splitForOpAtPoint` to split an `scf.for` at a given split point:
- first loop: `[lowerBound, splitPoint)`
- second loop: `[splitPoint, upperBound)`
The original loop is not mutated in place. Two new loops are cloned,
iter-args
are chained from the first to the second, uses of the original results
are
rewired, and the original op is erased. The helper returns `{first,
second}`.
The split point is validated as:
- `lowerBound <= splitPoint < upperBound`
- `step > 0`
- `splitPoint == lowerBound + k * step`
When the needed values are constant, an invalid split fails at compile
time.
When any of them is dynamic, the same conditions are emitted as runtime
`arith.cmpi` / `cf.assert`. Signed vs unsigned follows the loop's
`unsignedCmp` attribute (`getConstantIntValue` /
`getConstantUIntValue`).
`loopUnrollByFactor` uses this helper for the epilogue instead of
inlining
clone-and-rewire logic. The first returned loop is the main loop to
unroll;
the second is the epilogue.
Also adds `getConstantUIntValue` (zero-extend to `uint64_t`) next to
`getConstantIntValue`, and links the ControlFlow dialect from SCF Utils
so
runtime asserts can be created.
---------
Co-authored-by: Cursor <cursoragent at cursor.com>
Added:
mlir/test/Dialect/SCF/split-for-op-at-point.mlir
Modified:
mlir/include/mlir/Dialect/SCF/Utils/Utils.h
mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
mlir/lib/Dialect/SCF/Utils/Utils.cpp
mlir/lib/Dialect/Utils/StaticValueUtils.cpp
mlir/test/Dialect/SCF/loop-unroll.mlir
mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index dd394f7b355b6..4cc7e62e6fd58 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,22 @@ struct UnrolledLoopInfo {
std::optional<scf::ForOp> epilogueLoopOp = std::nullopt;
};
+/// Splits `forOp` into two consecutive loops at `splitPoint`:
+/// first: [lowerBound, splitPoint)
+/// second: [splitPoint, upperBound)
+///
+/// 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(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
/// 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/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 2350f705a7ed4..490972a837e96 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -363,6 +363,57 @@ void mlir::generateUnrolledLoop(
loopBodyBlock->getTerminator()->setOperands(lastYielded);
}
+/// Splits `forOp` into two consecutive loops at `splitPoint`.
+FailureOr<std::pair<scf::ForOp, scf::ForOp>>
+mlir::splitForOpAtPoint(RewriterBase &rewriter, scf::ForOp forOp,
+ Value splitPoint) {
+ if (splitPoint.getType() != forOp.getLowerBound().getType())
+ return failure();
+
+ // Reject statically known violations of the split preconditions.
+ bool isUnsigned = forOp.getUnsignedCmp();
+ Value lbVal = forOp.getLowerBound();
+ Value ubVal = forOp.getUpperBound();
+ Value stepVal = forOp.getStep();
+ auto checkSplitPoint = [&](auto getBound) -> LogicalResult {
+ auto lb = getBound(lbVal);
+ auto ub = getBound(ubVal);
+ auto step = getBound(stepVal);
+ auto split = getBound(splitPoint);
+ 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::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.
+ 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};
+}
+
/// Unrolls 'forOp' by 'unrollFactor', returns the unrolled main loop and the
/// epilogue loop, if the loop is unrolled.
FailureOr<UnrolledLoopInfo> mlir::loopUnrollByFactor(
@@ -484,27 +535,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(rewriter, 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..91decbde075c7 100644
--- a/mlir/test/Dialect/SCF/loop-unroll.mlir
+++ b/mlir/test/Dialect/SCF/loop-unroll.mlir
@@ -699,7 +699,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..0deb8b78c965b
--- /dev/null
+++ b/mlir/test/Dialect/SCF/split-for-op-at-point.mlir
@@ -0,0 +1,341 @@
+// 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.
+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: 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.
+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: 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.
+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.
+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: scf.for unsigned %{{.*}} = %{{.*}} to %[[C9]]
+// CHECK: scf.for unsigned %{{.*}} = %[[C9]] to %[[UB]]
+
+// -----
+
+// Dynamic unsigned step.
+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: scf.for unsigned %{{.*}} = %{{.*}} to %[[C6]] step %[[STEP]]
+// CHECK: scf.for unsigned %{{.*}} = %[[C6]] to %{{.*}} step %[[STEP]]
+
+// -----
+
+// Dynamic lower bound.
+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: scf.for %{{.*}} = %[[LB]] to %[[C9]] step %c1
+// CHECK: scf.for %{{.*}} = %[[C9]] to %c10 step %c1
+
+// -----
+
+// Fully dynamic bounds.
+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: scf.for %{{.*}} = %[[LB]] to %[[C9]] step %[[STEP]]
+// CHECK: scf.for %{{.*}} = %[[C9]] to %[[UB]] step %[[STEP]]
+
+// -----
+
+// 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
+ %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: 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/TestSCFUtils.cpp b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
index fafa03b6c089f..c6de295b6c1bc 100644
--- a/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
+++ b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
@@ -267,6 +267,64 @@ 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>();
+ }
+
+ 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);
+ });
+
+ IRRewriter rewriter(func.getContext());
+ for (scf::ForOp forOp : loopsToSplit) {
+ Value splitPoint;
+ if (auto splitAttr = forOp->getAttrOfType<IntegerAttr>(kSplitAtAttr)) {
+ rewriter.setInsertionPoint(forOp);
+ splitPoint =
+ 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();
+ 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);
+ rewriter.modifyOpInPlace(forOp,
+ [&] { forOp->removeAttr(kSplitArgAttr); });
+ } else {
+ continue;
+ }
+ if (failed(splitForOpAtPoint(rewriter, forOp, splitPoint))) {
+ emitError(forOp.getLoc(), "failed to split scf.for");
+ return signalPassFailure();
+ }
+ }
+ }
+};
} // namespace
namespace mlir {
@@ -275,6 +333,7 @@ void registerTestSCFUtilsPass() {
PassRegistration<TestSCFForUtilsPass>();
PassRegistration<TestSCFIfUtilsPass>();
PassRegistration<TestSCFPipeliningPass>();
+ PassRegistration<TestSplitForOpAtPointPass>();
}
} // namespace test
} // namespace mlir
More information about the Mlir-commits
mailing list