[Mlir-commits] [mlir] [mlir][scf] fix unsigned trip count with nuw in `constantTripCount` (PR #235272)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Oct 24 05:24:49 PDT 2026
https://github.com/StarOne01 updated https://github.com/llvm/llvm-project/pull/235272
>From 9a43df0c42f3a7c52a570d02054aca23f45df3b9 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/5] [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 ae548cf1d2e10b..604665ffd2a482 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -3029,7 +3029,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) {
@@ -3079,7 +3079,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
@@ -3103,6 +3103,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 3e78eb3c33ea46df035eebba92f51334c01a5be9 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/5] [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 ee1086e70597c2..d99aa47f4efe14 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 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
>From 8ebf9056b50c1ad0d3860febf3f26953c2d548e4 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/5] 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 d99aa47f4efe14..9d5161c5abb463 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 2d0ec76d6e6489c3a615cba5e88940c0e4ae2134 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/5] 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 9d5161c5abb463..3be4d0166abe41 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 5eff4f4260a4f7ce2dbecd332279179d8df25dd1 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/5] 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 bd8c35ad882501..eba3b8ce7618b0 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/mlir/test/Dialect/SCF/canonicalize.mlir
@@ -2455,3 +2455,86 @@ func.func @promote_multiple_inserts(%a: tensor<2xf32>, %b: tensor<2xf32>, %init:
}
return %r : tensor<4xf32>
}
+
+// -----
+
+// 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
More information about the Mlir-commits
mailing list