[Mlir-commits] [mlir] Improving MinUI and MaxUI implementation in valuebounds (PR #207701)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Hagai Lev Hacohen (HagaiLevHacohen)

<details>
<summary>Changes</summary>

Previously, the valuebounds implementation for MinUI and MaxUI was too conservative.
We can already make correct assumptions even when we know only one operand is provably not negative.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp (+21-13) 
- (modified) mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir (+19-7) 


``````````diff
diff --git a/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
index bc15d3b22f1b3..074f7efffc81b 100644
--- a/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -354,20 +354,24 @@ struct MinUIOpInterface
     assert(value == minOp.getResult() && "invalid value");
 
     // ValueBoundsConstraintSet models values as signed integers (e.g. an i8
-    // 0xff is treated as -1, not 255).So, we can only derive bounds for minui
-    // if both operands are provably non-negative.
+    // 0xff is treated as -1, not 255). For an unsigned minimum it is enough
+    // that a single operand is provably non-negative: minui(x, y) is in
+    // [0, y] whenever y >= 0 (and symmetrically for x).
     bool lhsNonNegative =
         ValueBoundsConstraintSet::isProvablyNonNegative(minOp.getLhs(), cstr);
     bool rhsNonNegative =
         ValueBoundsConstraintSet::isProvablyNonNegative(minOp.getRhs(), cstr);
-    if (!lhsNonNegative || !rhsNonNegative)
+    if (!lhsNonNegative && !rhsNonNegative)
       return;
 
-    cstr.bound(value) >= 0;
-    AffineExpr lhs = cstr.getExpr(minOp.getLhs());
-    AffineExpr rhs = cstr.getExpr(minOp.getRhs());
-    cstr.bound(value) <= lhs;
-    cstr.bound(value) <= rhs;
+    if (lhsNonNegative) {
+      AffineExpr lhs = cstr.getExpr(minOp.getLhs());
+      cstr.bound(value) <= lhs;
+    }
+    if (rhsNonNegative) {
+      AffineExpr rhs = cstr.getExpr(minOp.getRhs());
+      cstr.bound(value) <= rhs;
+    }
   }
 };
 
@@ -384,13 +388,17 @@ struct MaxUIOpInterface
         ValueBoundsConstraintSet::isProvablyNonNegative(maxOp.getLhs(), cstr);
     bool rhsNonNegative =
         ValueBoundsConstraintSet::isProvablyNonNegative(maxOp.getRhs(), cstr);
-    if (!lhsNonNegative || !rhsNonNegative)
+    if (!lhsNonNegative && !rhsNonNegative)
       return;
 
-    AffineExpr lhs = cstr.getExpr(maxOp.getLhs());
-    AffineExpr rhs = cstr.getExpr(maxOp.getRhs());
-    cstr.bound(value) >= lhs;
-    cstr.bound(value) >= rhs;
+    if (lhsNonNegative) {
+      AffineExpr lhs = cstr.getExpr(maxOp.getLhs());
+      cstr.bound(value) >= lhs;
+    }
+    if (rhsNonNegative) {
+      AffineExpr rhs = cstr.getExpr(maxOp.getRhs());
+      cstr.bound(value) >= rhs;
+    }
   }
 };
 } // namespace
diff --git a/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
index 41f7f18baefec..ceb8256003b1e 100644
--- a/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
@@ -486,32 +486,38 @@ func.func @arith_maxui() -> index {
 
 // -----
 
-func.func @arith_maxui_unknown_sign(%a: index) -> index {
+// CHECK-LABEL: func @arith_maxui_unknown_operand(
+//       CHECK:   %[[lb:.*]] = arith.constant 4 : index
+//       CHECK:   return %[[lb]]
+func.func @arith_maxui_unknown_operand(%a: index) -> index {
   %c4 = arith.constant 4 : index
   %0 = arith.maxui %a, %c4 : index
-  // expected-error @below{{could not reify bound}}
-  %1 = "test.reify_bound"(%0) {type = "LB"} : (index) -> (index)
+  %1 = "test.reify_bound"(%0) {type = "LB", constant} : (index) -> (index)
   return %1 : index
 }
 
 // -----
 
+// CHECK-LABEL: func @arith_minui_wraparound(
+//       CHECK:   %[[ub:.*]] = arith.constant 11 : index
+//       CHECK:   return %[[ub]]
 func.func @arith_minui_wraparound() -> index {
   %c255 = arith.constant 0xFF : i8
   %c10 = arith.constant 10 : i8
   %0 = arith.minui %c255, %c10 : i8
-  // expected-error @below{{could not reify bound}}
   %1 = "test.reify_bound"(%0) {type = "UB", allow_integer_type} : (i8) -> (index)
   return %1 : index
 }
 
 // -----
 
+// CHECK-LABEL: func @arith_maxui_wraparound(
+//       CHECK:   %[[lb:.*]] = arith.constant 10 : index
+//       CHECK:   return %[[lb]]
 func.func @arith_maxui_wraparound() -> index {
   %c255 = arith.constant 0xFF : i8
   %c10 = arith.constant 10 : i8
   %0 = arith.maxui %c255, %c10 : i8
-  // expected-error @below{{could not reify bound}}
   %1 = "test.reify_bound"(%0) {type = "LB", allow_integer_type} : (i8) -> (index)
   return %1 : index
 }
@@ -547,12 +553,15 @@ func.func @arith_minui_nonneg_symbolic(%a: index) -> index {
 
 // -----
 
+// CHECK-LABEL: func @arith_minui_negative_symbolic(
+//  CHECK-SAME:     %[[a:.*]]: index
+//       CHECK:   %[[ub:.*]] = arith.constant 5 : index
+//       CHECK:   return %[[ub]]
 func.func @arith_minui_negative_symbolic(%a: index) -> index {
   %cm1 = arith.constant -1 : index
   %c4 = arith.constant 4 : index
   %neg = arith.minsi %a, %cm1 : index
   %0 = arith.minui %neg, %c4 : index
-  // expected-error @below{{could not reify bound}}
   %1 = "test.reify_bound"(%0) {type = "UB", constant} : (index) -> (index)
   return %1 : index
 }
@@ -574,12 +583,15 @@ func.func @arith_maxui_nonneg_symbolic(%a: index) -> index {
 
 // -----
 
+// CHECK-LABEL: func @arith_maxui_negative_symbolic(
+//  CHECK-SAME:     %[[a:.*]]: index
+//       CHECK:   %[[lb:.*]] = arith.constant 4 : index
+//       CHECK:   return %[[lb]]
 func.func @arith_maxui_negative_symbolic(%a: index) -> index {
   %cm1 = arith.constant -1 : index
   %c4 = arith.constant 4 : index
   %neg = arith.minsi %a, %cm1 : index
   %0 = arith.maxui %neg, %c4 : index
-  // expected-error @below{{could not reify bound}}
   %1 = "test.reify_bound"(%0) {type = "LB", constant} : (index) -> (index)
   return %1 : index
 }

``````````

</details>


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