[Mlir-commits] [mlir] [mlir][scf] Fix scf.for value bounds for empty and unsigned loops (PR #226761)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Dhairyashil R G (dhairyashilRG)

<details>
<summary>Changes</summary>

Value bounds assumes that an `scf.for` result is `init + ceildiv(ub - lb, step) * (yield - iter_arg)`. This is only true if the loop runs at least once. When the loop does not run, `ceildiv(ub - lb, step)` is negative instead of 0. For unsigned loops, `ub - lb` is computed as if the bounds were signed, which also gives a wrong trip count. The induction variable bounds of unsigned loops have the same signed problem.

Example:
```mlir
%r = scf.for %iv = %c5 to %c2 step %c1 iter_args(%arg = %c7) -> index {
  %n = arith.addi %arg, %c1 : index
  scf.yield %n : index
}
%m = affine.min affine_map<(d0) -> (d0, 5)>(%r)
```
The loop does not run, so `%r` is 7 and `%m` should be 5. Value bounds thinks `%r` is 4, so `affine-simplify-min-max` replaces `%m` with `%r` and the result becomes 7.

Canonicalization usually removes a constant empty loop like this one, but the analysis should not depend on that. The unsigned case is not removed by canonicalization and gives wrong results too.

With this change, the result formula is used only for signed loops that are known to run at least once, and unsigned loops get no induction variable bounds.

Assisted-by: Claude

---
Full diff: https://github.com/llvm/llvm-project/pull/226761.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp (+13) 
- (modified) mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir (+54) 


``````````diff
diff --git a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
index 496a7b036e65d..171ab3f19e09a 100644
--- a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -101,6 +101,16 @@ struct ForOpInterface
     if (dim.has_value() || isa<BlockArgument>(value))
       return;
 
+    // The closed form below holds only if the loop runs at least once (the
+    // trip count is not clamped at 0) and the bounds are compared as signed
+    // integers.
+    if (forOp.getUnsignedCmp() ||
+        !cstr.populateAndCompare(
+            forOp.getUpperBound(),
+            ValueBoundsConstraintSet::ComparisonOperator::GT,
+            forOp.getLowerBound()))
+      return;
+
     // `value` is result of `forOp`, we can prove that:
     // %result == %init_arg + trip_count * (%yielded_value - %iter_arg).
     // Where trip_count is (ub - lb) / step.
@@ -117,6 +127,9 @@ struct ForOpInterface
     auto forOp = cast<ForOp>(op);
 
     if (value == forOp.getInductionVar()) {
+      // The IV bounds are signed; unsigned loops are not modeled.
+      if (forOp.getUnsignedCmp())
+        return;
       return populateIVBounds(forOp.getLowerBound(), forOp.getUpperBound(),
                               forOp.getStep(), value, cstr);
     }
diff --git a/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
index 60fe96d52d20b..ab3e0b3989cb1 100644
--- a/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
@@ -388,3 +388,57 @@ func.func @scf_forall_computed_upper_bound(%x: index) {
   }
   return
 }
+
+// -----
+
+// The loop does not run (5 to 2), so %0 is 3. ceildiv(ub - lb, step) is
+// negative and must not be used as the trip count: %0 is not 6.
+func.func @scf_for_result_empty() {
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %c3 = arith.constant 3 : index
+  %c5 = arith.constant 5 : index
+  %c6 = arith.constant 6 : index
+  %0 = scf.for %iv = %c5 to %c2 step %c1 iter_args(%arg = %c3) -> index {
+    %1 = arith.subi %arg, %c1 : index
+    scf.yield %1 : index
+  }
+  // expected-error @below{{unknown}}
+  "test.compare"(%0, %c6) {cmp = "EQ"} : (index, index) -> ()
+  return
+}
+
+// -----
+
+// Unsigned loop from 2^64 - 1 to 5: the loop does not run, so %0 is 3. The
+// signed difference of the bounds must not be used as the trip count: %0 is
+// not 9.
+func.func @scf_for_unsigned_result() {
+  %c1 = arith.constant 1 : index
+  %c3 = arith.constant 3 : index
+  %c5 = arith.constant 5 : index
+  %c9 = arith.constant 9 : index
+  %lb = arith.constant -1 : index
+  %0 = scf.for unsigned %iv = %lb to %c5 step %c1 iter_args(%arg = %c3) -> index {
+    %1 = arith.addi %arg, %c1 : index
+    scf.yield %1 : index
+  }
+  // expected-error @below{{unknown}}
+  "test.compare"(%0, %c9) {cmp = "EQ"} : (index, index) -> ()
+  return
+}
+
+// -----
+
+// Unsigned loop: %iv is 1, 2^62 + 1 and 2^63 + 1. Signed bounds do not hold.
+func.func @scf_for_unsigned_induction_var() {
+  %c1 = arith.constant 1 : index
+  %ub = arith.constant -4611686018427387904 : index
+  %step = arith.constant 4611686018427387904 : index
+  scf.for unsigned %iv = %c1 to %ub step %step {
+    // expected-error @below{{could not reify bound}}
+    %0 = "test.reify_bound"(%iv) {type = "UB", constant} : (index) -> (index)
+    "test.some_use"(%0) : (index) -> ()
+  }
+  return
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/226761