[Mlir-commits] [mlir] [mlir][scf] fix canonicalization of unsigned scf.for/scf.parallel loops (PR #235272)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 26 07:54:27 PDT 2026
https://github.com/StarOne01 updated https://github.com/llvm/llvm-project/pull/235272
>From ea4aab208e97b1f20901804b1702efd4d9836a84 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sat, 24 Oct 2026 17:29:21 +0530
Subject: [PATCH 1/8] [mlir][scf] update constantTripCount call to use
getUnsignedCmp to check for signedness
update constantTripCount call to use getUnsignedCmp to check for signedness and only merge loops with the same signedness in `ParallelOpSingleOrZeroIterationDimsFolder`
---
mlir/lib/Dialect/SCF/IR/SCF.cpp | 7 +++++--
1 file changed, 5 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 99fd8a2c3433c..44597f621cc6c 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -3056,7 +3056,7 @@ struct ParallelOpSingleOrZeroIterationDimsFolder
llvm::zip(op.getLowerBound(), op.getUpperBound(), op.getStep(),
op.getInductionVars())) {
auto numIterations =
- constantTripCount(lb, ub, step, /*isSigned=*/true, computeUbMinusLb);
+ constantTripCount(lb, ub, step, !op.getUnsignedCmp(), computeUbMinusLb);
if (numIterations.has_value()) {
// Remove the loop if it performs zero iterations.
if (*numIterations == 0) {
@@ -3106,7 +3106,7 @@ struct ParallelOpSingleOrZeroIterationDimsFolder
// Replace the parallel loop by lower-dimensional parallel loop.
auto newOp =
ParallelOp::create(rewriter, op.getLoc(), newLowerBounds,
- newUpperBounds, newSteps, op.getInitVals(), nullptr);
+ 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
@@ -3130,6 +3130,9 @@ struct MergeNestedParallelLoops : public OpRewritePattern<ParallelOp> {
auto innerOp = dyn_cast<ParallelOp>(outerBody.front());
if (!innerOp)
return failure();
+
+ if(op.getUnsignedCmp() != innerOp.getUnsignedCmp())
+ return failure();
for (auto val : outerBody.getArguments())
if (llvm::is_contained(innerOp.getLowerBound(), val) ||
>From bccaf99521321330191cee1a469f33845722cb2e Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sat, 24 Oct 2026 16:18:11 +0530
Subject: [PATCH 2/8] [mlir][scf] fix unsigned trip count handling with
non-negative upper bounds and update tests
---
mlir/lib/Dialect/Utils/StaticValueUtils.cpp | 5 ++---
mlir/test/Dialect/SCF/trip_count.mlir | 4 ++--
2 files changed, 4 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index ee1086e70597c..d99aa47f4efe1 100644
--- a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
@@ -392,7 +392,7 @@ std::optional<APInt> constantTripCount(
/// Non-constant bound, let's try to compute the difference between the
/// upper and lower bound
- std::optional<llvm::APSInt> maybeDiff =
+ std::optional<llvm::APInt> maybeDiff =
computeUbMinusLb(cast<Value>(lb), cast<Value>(ub), isSigned);
if (!maybeDiff)
return std::nullopt;
@@ -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);
}
@@ -424,7 +424,6 @@ std::optional<APInt> constantTripCount(
llvm::APInt remainder = diff.urem(stepCst);
if (!remainder.isZero())
tripCount = tripCount + 1;
-
LDBG() << "constantTripCount found: " << tripCount;
return tripCount;
}
diff --git a/mlir/test/Dialect/SCF/trip_count.mlir b/mlir/test/Dialect/SCF/trip_count.mlir
index 7e74988b35019..732db0ead10f4 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
>From 00eb242fd169578aac7fc50ff409265cd80ca3f8 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sat, 24 Oct 2026 16:33:24 +0530
Subject: [PATCH 3/8] fix: change APInt to APSInt for trip count difference
computation
---
mlir/lib/Dialect/Utils/StaticValueUtils.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index d99aa47f4efe1..9d5161c5abb46 100644
--- a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
@@ -392,7 +392,7 @@ std::optional<APInt> constantTripCount(
/// Non-constant bound, let's try to compute the difference between the
/// upper and lower bound
- std::optional<llvm::APInt> maybeDiff =
+ std::optional<llvm::APSInt> maybeDiff =
computeUbMinusLb(cast<Value>(lb), cast<Value>(ub), isSigned);
if (!maybeDiff)
return std::nullopt;
>From b7ca0f2ebe4392041d88c97028916a0ae7c827ab Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sat, 24 Oct 2026 16:40:55 +0530
Subject: [PATCH 4/8] fix: add missing line
---
mlir/lib/Dialect/Utils/StaticValueUtils.cpp | 1 +
1 file changed, 1 insertion(+)
diff --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index 9d5161c5abb46..3be4d0166abe4 100644
--- a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
@@ -424,6 +424,7 @@ std::optional<APInt> constantTripCount(
llvm::APInt remainder = diff.urem(stepCst);
if (!remainder.isZero())
tripCount = tripCount + 1;
+
LDBG() << "constantTripCount found: " << tripCount;
return tripCount;
}
>From 2ac62460a1d15f0a3d6f82b31bdb31b9ce499a0b Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sat, 24 Oct 2026 17:53:38 +0530
Subject: [PATCH 5/8] add tests for unsigned trip count handling in
scf.parallel (part b)
---
mlir/test/Dialect/SCF/canonicalize.mlir | 83 +++++++++++++++++++++++++
1 file changed, 83 insertions(+)
diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir
index 4b4c4b1720a67..455b95e6e58f0 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/mlir/test/Dialect/SCF/canonicalize.mlir
@@ -2731,3 +2731,86 @@ func.func @select_chain(%p: i1, %a: i32, %b: i32, %c: i32) -> i32 {
return %y : i32
}
func.func private @consume_pair(i32, i32)
+
+// -----
+
+// B1: 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
+}
+
+// -----
+
+// B2: 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
+}
+
+// -----
+
+// B3: 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
+}
\ No newline at end of file
>From 0621dfd7bd482f4a4a26bd8b0a1606088fa6c3c5 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sat, 24 Oct 2026 17:56:29 +0530
Subject: [PATCH 6/8] add missing newline at end of canonicalize.mlir
---
mlir/test/Dialect/SCF/canonicalize.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir
index 455b95e6e58f0..5d64722cd9292 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/mlir/test/Dialect/SCF/canonicalize.mlir
@@ -2813,4 +2813,4 @@ func.func @mixed_signedness_parallel_not_merged(
scf.reduce
}
return
-}
\ No newline at end of file
+}
>From b307bfd8e6da6106a3f870eb4700094e00f67f85 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sat, 24 Oct 2026 17:57:43 +0530
Subject: [PATCH 7/8] update comments for unsigned trip count handling in
scf.parallel
---
mlir/test/Dialect/SCF/canonicalize.mlir | 7 +++----
1 file changed, 3 insertions(+), 4 deletions(-)
diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir
index 5d64722cd9292..e5b0f354e1808 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/mlir/test/Dialect/SCF/canonicalize.mlir
@@ -2734,7 +2734,7 @@ func.func private @consume_pair(i32, i32)
// -----
-// B1: bounds that look like ub < lb when read as signed must not erase an
+// 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>)
@@ -2761,8 +2761,7 @@ func.func @unsigned_parallel_not_erased(%arg0: memref<4xi64>) {
// -----
-// B2: removing the single-iteration dim must keep `unsigned` on the rebuilt
-// loop.
+// 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
@@ -2787,7 +2786,7 @@ func.func @unsigned_parallel_single_iter_dim_keeps_unsigned(
// -----
-// B3: a signed outer loop and an unsigned inner loop must not be merged.
+// 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
>From 1bc4f8c85cf64e810c97f842ebaba2c856f17dd5 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sat, 24 Oct 2026 18:03:13 +0530
Subject: [PATCH 8/8] improve formatting and readability
---
mlir/lib/Dialect/SCF/IR/SCF.cpp | 14 +++++++-------
1 file changed, 7 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 44597f621cc6c..e4e75ab3e4a0f 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, !op.getUnsignedCmp(), 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, op.getUnsignedCmp());
+ 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
@@ -3130,8 +3130,8 @@ struct MergeNestedParallelLoops : public OpRewritePattern<ParallelOp> {
auto innerOp = dyn_cast<ParallelOp>(outerBody.front());
if (!innerOp)
return failure();
-
- if(op.getUnsignedCmp() != innerOp.getUnsignedCmp())
+
+ if (op.getUnsignedCmp() != innerOp.getUnsignedCmp())
return failure();
for (auto val : outerBody.getArguments())
More information about the Mlir-commits
mailing list