[Mlir-commits] [mlir] f9f340d - [mlir][affine] Clamp negative static trip counts to zero (#227134)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 8 11:54:33 PDT 2026
Author: Slava Zakharin
Date: 2026-10-08T11:54:26-07:00
New Revision: f9f340dc111d85d0a8c8dd7c7c753d9a5d60cc38
URL: https://github.com/llvm/llvm-project/commit/f9f340dc111d85d0a8c8dd7c7c753d9a5d60cc38
DIFF: https://github.com/llvm/llvm-project/commit/f9f340dc111d85d0a8c8dd7c7c753d9a5d60cc38.diff
LOG: [mlir][affine] Clamp negative static trip counts to zero (#227134)
When the bounds of an affine.for are not both constant,
getStaticTripCount subtracts the lower bound from each upper bound
result, divides by the step, and returns the minimum if every result is
constant. A result that folds to a negative constant means the loop does
not execute, but it was converted to uint64_t and wrapped around. A
single-result bound such as `%n to %n - 2` reported a trip count of 2^64
- 2, and a multi-result bound such as
`max (%n) to min (%n - 4, %n + 4)` reported 4, because the wrapped value
lost the minimum to the other result.
Callers acted on these counts. Full unrolling turned the multi-result
loop into four unconditional copies of its body and hit an assertion on
the single-result one. Unrolling the single-result loop by four added a
cleanup loop that ran two iterations, and LICM treated both loops as
executing at least once and hoisted side-effecting ops out of them.
Clamp negative results to zero, as the constant-bound path already does.
Assisted-by: Cursor
Added:
Modified:
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
mlir/test/Dialect/Affine/trip-count.mlir
mlir/test/Dialect/Affine/unroll.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 0b14c7de070585..ad9845d8935fa3 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2860,7 +2860,7 @@ std::optional<APInt> AffineForOp::getStaticTripCount() {
for (unsigned i = 0, e = tripCountValueMap.getNumResults(); i < e; ++i) {
AffineExpr expr = tripCountValueMap.getResult(i).ceilDiv(step);
if (auto constExpr = llvm::dyn_cast<AffineConstantExpr>(expr)) {
- uint64_t value = constExpr.getValue();
+ uint64_t value = std::max<int64_t>(constExpr.getValue(), 0);
if (tripCount.has_value())
tripCount = std::min(*tripCount, value);
else
diff --git a/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir b/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
index 2cb3826d8f6dbc..2666958ce52b57 100644
--- a/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
+++ b/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
@@ -965,6 +965,30 @@ func.func @no_hoist_from_zero_trip_count_loop(%x: i32) -> i32 {
// -----
+// The bounds are not constant, but the upper bound is below the lower bound, so
+// this loop never executes either.
+
+// CHECK-LABEL: func @no_hoist_from_zero_trip_count_loop_nonconstant_bounds
+// CHECK-SAME: (%[[X:.*]]: i32,
+func.func @no_hoist_from_zero_trip_count_loop_nonconstant_bounds(%x: i32, %n: index) -> i32 {
+ %alloc = memref.alloc() : memref<1xi32>
+ %c42 = arith.constant 42 : i32
+ affine.store %c42, %alloc[0] : memref<1xi32>
+ affine.for %i = affine_map<(d0) -> (d0)>(%n) to affine_map<(d0) -> (d0 - 2)>(%n) {
+ affine.store %x, %alloc[0] : memref<1xi32>
+ }
+ // CHECK: %[[ALLOC:.*]] = memref.alloc()
+ // CHECK: %[[C42:.*]] = arith.constant 42 : i32
+ // CHECK: affine.store %[[C42]], %[[ALLOC]][0]
+ // CHECK: affine.for
+ // CHECK-NEXT: affine.store %[[X]], %[[ALLOC]][0]
+ // CHECK-NEXT: }
+ %r = affine.load %alloc[0] : memref<1xi32>
+ return %r : i32
+}
+
+// -----
+
// Side-effectful ops must not be hoisted from loops with an unknown (dynamic)
// trip count, because such a loop may not execute at all (e.g. when the upper
// bound is zero at runtime). Only pure ops may be hoisted.
diff --git a/mlir/test/Dialect/Affine/trip-count.mlir b/mlir/test/Dialect/Affine/trip-count.mlir
index e28e410fd21123..1c53675b44b996 100644
--- a/mlir/test/Dialect/Affine/trip-count.mlir
+++ b/mlir/test/Dialect/Affine/trip-count.mlir
@@ -36,3 +36,27 @@ func.func @affine_symbolic_loops(%N : index) {
return
}
+
+// -----
+
+// The upper bound is below the lower bound, so none of these loops executes.
+
+// CHECK-LABEL: func.func @affine_empty_loops
+func.func @affine_empty_loops(%N : index) {
+ // CHECK: "test.trip-count" = 0
+ affine.for %i = affine_map<(d0) -> (d0)>(%N) to affine_map<(d0) -> (d0 - 2)>(%N) {
+ affine.yield
+ }
+
+ // CHECK: "test.trip-count" = 0
+ affine.for %i = max affine_map<(d0) -> (d0)>(%N) to min affine_map<(d0) -> (d0 - 4, d0 + 4)>(%N) {
+ affine.yield
+ }
+
+ // CHECK: "test.trip-count" = 0
+ affine.for %i = 5 to min affine_map<() -> (3, 10)>() {
+ affine.yield
+ }
+
+ return
+}
diff --git a/mlir/test/Dialect/Affine/unroll.mlir b/mlir/test/Dialect/Affine/unroll.mlir
index 5dec623f684120..4d4c980090cfa9 100644
--- a/mlir/test/Dialect/Affine/unroll.mlir
+++ b/mlir/test/Dialect/Affine/unroll.mlir
@@ -733,6 +733,24 @@ func.func @unroll_zero_trip_count_case() {
return
}
+// The upper bounds are below the lower bounds, so neither loop executes.
+// UNROLL-FULL-LABEL: func @unroll_zero_trip_count_nonconstant_bounds
+func.func @unroll_zero_trip_count_nonconstant_bounds(%N: index) {
+ // UNROLL-FULL: affine.for
+ // UNROLL-FULL-NEXT: "foo"() : () -> ()
+ // UNROLL-FULL-NEXT: }
+ affine.for %i = affine_map<(d0) -> (d0)>(%N) to affine_map<(d0) -> (d0 - 2)>(%N) {
+ "foo"() : () -> ()
+ }
+ // UNROLL-FULL: affine.for
+ // UNROLL-FULL-NEXT: "foo"() : () -> ()
+ // UNROLL-FULL-NEXT: }
+ affine.for %i = max affine_map<(d0) -> (d0)>(%N) to min affine_map<(d0) -> (d0 - 4, d0 + 4)>(%N) {
+ "foo"() : () -> ()
+ }
+ return
+}
+
// UNROLL-CLEANUP-LOOP-LABEL: func @unroll_cleanup_loop_with_larger_unroll_factor()
func.func @unroll_cleanup_loop_with_larger_unroll_factor() {
affine.for %i = 0 to 3 {
More information about the Mlir-commits
mailing list