[Mlir-commits] [mlir] 95acd8d - [mlir][scf] fix canonicalization of unsigned scf.for/scf.parallel loops (#235272)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 27 06:53:52 PDT 2026
Author: Prashanth
Date: 2026-10-27T06:53:47-07:00
New Revision: 95acd8de58cdebbcf4246c15f8f7c358f2ef2727
URL: https://github.com/llvm/llvm-project/commit/95acd8de58cdebbcf4246c15f8f7c358f2ef2727
DIFF: https://github.com/llvm/llvm-project/commit/95acd8de58cdebbcf4246c15f8f7c358f2ef2727.diff
LOG: [mlir][scf] fix canonicalization of unsigned scf.for/scf.parallel loops (#235272)
- `constantTripCount` didn't check the signedness of `diff` before
checking whether it is negative, so it returned the trip count as 0 for
unsigned loops with "non-constant bound" - problem A according to the
issue
- update `constantTripCount` call to use `getUnsignedCmp` to check for
signedness and only merge loops with the same signedness in
`ParallelOpSingleOrZeroIterationDimsFolder` - problem B* according to
the issue
fixes #233223
claude helped me with the tests
Added:
Modified:
mlir/lib/Dialect/SCF/IR/SCF.cpp
mlir/lib/Dialect/Utils/StaticValueUtils.cpp
mlir/test/Dialect/SCF/canonicalize.mlir
mlir/test/Dialect/SCF/trip_count.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 99fd8a2c3433c1..e4e75ab3e4a0f8 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -3055,8 +3055,8 @@ struct ParallelOpSingleOrZeroIterationDimsFolder
for (auto [lb, ub, step, iv] :
llvm::zip(op.getLowerBound(), op.getUpperBound(), op.getStep(),
op.getInductionVars())) {
- auto numIterations =
- constantTripCount(lb, ub, step, /*isSigned=*/true, computeUbMinusLb);
+ auto numIterations = constantTripCount(lb, ub, step, !op.getUnsignedCmp(),
+ computeUbMinusLb);
if (numIterations.has_value()) {
// Remove the loop if it performs zero iterations.
if (*numIterations == 0) {
@@ -3104,9 +3104,9 @@ struct ParallelOpSingleOrZeroIterationDimsFolder
return success();
}
// Replace the parallel loop by lower-dimensional parallel loop.
- auto newOp =
- ParallelOp::create(rewriter, op.getLoc(), newLowerBounds,
- newUpperBounds, newSteps, op.getInitVals(), nullptr);
+ auto newOp = ParallelOp::create(rewriter, op.getLoc(), newLowerBounds,
+ newUpperBounds, newSteps, op.getInitVals(),
+ nullptr, op.getUnsignedCmp());
// Erase the empty block that was inserted by the builder.
rewriter.eraseBlock(newOp.getBody());
// Clone the loop body and remap the block arguments of the collapsed loops
@@ -3131,6 +3131,9 @@ struct MergeNestedParallelLoops : public OpRewritePattern<ParallelOp> {
if (!innerOp)
return failure();
+ if (op.getUnsignedCmp() != innerOp.getUnsignedCmp())
+ return failure();
+
for (auto val : outerBody.getArguments())
if (llvm::is_contained(innerOp.getLowerBound(), val) ||
llvm::is_contained(innerOp.getUpperBound(), val) ||
diff --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index ee1086e70597c2..3be4d0166abe41 100644
--- a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
@@ -400,7 +400,7 @@ std::optional<APInt> constantTripCount(
}
LDBG() << "constantTripCount: " << (isSigned ? "isSigned" : "isUnsigned")
<< ", ub-lb: " <<
diff << "(" <<
diff .getBitWidth() << "b)";
- if (
diff .isNegative()) {
+ if (isSigned &&
diff .isNegative()) {
LDBG() << "constantTripCount is 0 because ub-lb
diff is negative";
return APInt(bitwidth, 0);
}
diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir
index 4b4c4b1720a67d..e5b0f354e1808d 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/mlir/test/Dialect/SCF/canonicalize.mlir
@@ -2731,3 +2731,85 @@ func.func @select_chain(%p: i1, %a: i32, %b: i32, %c: i32) -> i32 {
return %y : i32
}
func.func private @consume_pair(i32, i32)
+
+// -----
+
+// bounds that look like ub < lb when read as signed must not erase an
+// unsigned scf.parallel. The unsigned trip count is 4.
+// CHECK-LABEL: func @unsigned_parallel_not_erased
+// CHECK-SAME: (%[[A:.*]]: memref<4xi64>)
+// CHECK-DAG: %[[LB:.*]] = arith.constant 9223372036854775806 : index
+// CHECK-DAG: %[[UB:.*]] = arith.constant -9223372036854775806 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: scf.parallel unsigned (%[[I:.*]]) = (%[[LB]]) to (%[[UB]]) step (%[[C1]]) {
+// CHECK: %[[OFF:.*]] = arith.subi %[[I]], %[[LB]] : index
+// CHECK: %[[V:.*]] = arith.index_cast %[[OFF]] : index to i64
+// CHECK: memref.store %[[V]], %[[A]][%[[OFF]]] : memref<4xi64>
+// CHECK: scf.reduce
+func.func @unsigned_parallel_not_erased(%arg0: memref<4xi64>) {
+ %lb = arith.constant 9223372036854775806 : index
+ %ub = arith.constant -9223372036854775806 : index
+ %c1 = arith.constant 1 : index
+ scf.parallel unsigned (%i) = (%lb) to (%ub) step (%c1) {
+ %off = arith.subi %i, %lb : index
+ %v = arith.index_cast %off : index to i64
+ memref.store %v, %arg0[%off] : memref<4xi64>
+ scf.reduce
+ }
+ return
+}
+
+// -----
+
+// removing the single-iteration dim must keep `unsigned` on the rebuilt loop.
+// CHECK-LABEL: func @unsigned_parallel_single_iter_dim_keeps_unsigned
+// CHECK-SAME: (%[[A:.*]]: memref<4xi64>, %[[LB:.*]]: index, %[[UB:.*]]: index)
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-NOT: scf.parallel (%{{.*}}, %{{.*}})
+// CHECK: scf.parallel unsigned (%[[I:.*]]) = (%[[LB]]) to (%[[UB]]) step (%[[C1]]) {
+// CHECK: %[[OFF:.*]] = arith.subi %[[I]], %[[LB]] : index
+// CHECK: %[[V:.*]] = arith.index_cast %[[OFF]] : index to i64
+// CHECK: memref.store %[[V]], %[[A]][%[[OFF]]] : memref<4xi64>
+func.func @unsigned_parallel_single_iter_dim_keeps_unsigned(
+ %arg0: memref<4xi64>, %lb: index, %ub: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ scf.parallel unsigned (%k, %i) = (%c0, %lb) to (%c1, %ub) step (%c1, %c1) {
+ %off0 = arith.subi %i, %lb : index
+ %off = arith.addi %off0, %k : index
+ %v = arith.index_cast %off : index to i64
+ memref.store %v, %arg0[%off] : memref<4xi64>
+ scf.reduce
+ }
+ return
+}
+
+// -----
+
+// a signed outer loop and an unsigned inner loop must not be merged.
+// CHECK-LABEL: func @mixed_signedness_parallel_not_merged
+// CHECK-SAME: (%[[A:.*]]: memref<2x4xi64>, %[[LB:.*]]: index, %[[UB:.*]]: index)
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: scf.parallel (%[[K:.*]]) = (%[[C0]]) to (%[[C2]]) step (%[[C1]]) {
+// CHECK: scf.parallel unsigned (%[[I:.*]]) = (%[[LB]]) to (%[[UB]]) step (%[[C1]]) {
+// CHECK: %[[OFF:.*]] = arith.subi %[[I]], %[[LB]] : index
+// CHECK: %[[V:.*]] = arith.index_cast %[[OFF]] : index to i64
+// CHECK: memref.store %[[V]], %[[A]][%[[K]], %[[OFF]]] : memref<2x4xi64>
+func.func @mixed_signedness_parallel_not_merged(
+ %arg0: memref<2x4xi64>, %lb: index, %ub: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ scf.parallel (%k) = (%c0) to (%c2) step (%c1) {
+ scf.parallel unsigned (%i) = (%lb) to (%ub) step (%c1) {
+ %off = arith.subi %i, %lb : index
+ %v = arith.index_cast %off : index to i64
+ memref.store %v, %arg0[%k, %off] : memref<2x4xi64>
+ scf.reduce
+ }
+ scf.reduce
+ }
+ return
+}
diff --git a/mlir/test/Dialect/SCF/trip_count.mlir b/mlir/test/Dialect/SCF/trip_count.mlir
index 7e74988b350197..732db0ead10f42 100644
--- a/mlir/test/Dialect/SCF/trip_count.mlir
+++ b/mlir/test/Dialect/SCF/trip_count.mlir
@@ -663,7 +663,7 @@ func.func @trip_count_arith_add_negative_nuw_loop_unsigned(%lb : i32) -> i32 {
%c0_i32 = arith.constant 0 : i32
%c4_i32 = arith.constant 4 : i32
%c-16_i32 = arith.constant -16 : i32
- // CHECK: "test.trip-count" = 0
+ // CHECK: "test.trip-count" = 1073741820
%ub = arith.addi %lb, %c-16_i32 overflow<nuw> : i32
%1 = scf.for unsigned %arg0 = %lb to %ub step %c4_i32 iter_args(%arg1 = %c0_i32) -> (i32) : i32 {
scf.yield %arg0 : i32
@@ -677,7 +677,7 @@ func.func @trip_count_arith_add_negative_nuw_loop_unsigned(%lb : i32) -> i32 {
func.func @trip_count_arith_add_negative_nuw_loop_unsigned_step_dyn(%lb : i32, %step : i32) -> i32 {
%c0_i32 = arith.constant 0 : i32
%c-16_i32 = arith.constant -16 : i32
- // CHECK: "test.trip-count" = 0
+ // CHECK: "test.trip-count" = "none"
%ub = arith.addi %lb, %c-16_i32 overflow<nuw> : i32
%1 = scf.for unsigned %arg0 = %lb to %ub step %step iter_args(%arg1 = %c0_i32) -> (i32) : i32 {
scf.yield %arg0 : i32
More information about the Mlir-commits
mailing list