[Util] Fix AssumeIntOp::inferResultRanges bug (#19195)

I noticed that the util-optimize-int-arithmetic pass, sometimes was
incorrectly optimizing away operations. I tracked the issue down to the
fact that `AssumeIntOp::inferResultRanges` doesn't always call
`setResultRange`. Looking at the docs for `InferIntRangeInterface` in
LLVM, it suggests that `setResultRange` must be called for each result,
and looking further it's clear that some of the `arith` op folders rely
on the fact that the range is always set (eg. `arith.select`).

This PR updates `AssumeIntOp::inferResultRanges` to call
`setResultRange` even when umin or umax are not set.

I added a test case that used to be optimized to an incorrect constant,
and now is not optimized out.

Signed-off-by: James Bartlett <[email protected]>

Author: JamesMBartlett

compiler/src/iree/compiler/Dialect/Util/IR/UtilOps.cpp

@@ -1217,11 +1217,14 @@ void AssumeIntOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
     else
       continue;
     auto [umin, umax] = getUnionedUnsignedRange(index);
-    if (umin && umax) {
-      APInt uminAp(bitWidth, *umin);
-      APInt umaxAp(bitWidth, *umax);
-      setResultRange(result, ConstantIntRanges::fromUnsigned(uminAp, umaxAp));
-    }
+    auto uminAp = APInt::getMinValue(bitWidth);
+    auto umaxAp = APInt::getMaxValue(bitWidth);
+    if (umin)
+      uminAp = APInt(bitWidth, *umin);
+    if (umax)
+      umaxAp = APInt(bitWidth, *umax);
+
+    setResultRange(result, ConstantIntRanges::fromUnsigned(uminAp, umaxAp));
   }
 }
 

compiler/src/iree/compiler/Dialect/Util/Transforms/test/optimize_int_arithmetic.mlir

@@ -523,3 +523,14 @@ util.func @hal_buffer_view_rank_min_max(%bv : !hal.buffer_view) -> (i1, i1, i1)
   // CHECK: util.return %[[FALSE]], %[[TRUE]], %[[FALSE]]
   util.return %1, %2, %3 : i1, i1, i1
 }
+
+// -----
+
+util.func @assume_int_with_single_bound(%bool: i1, %ind: index) -> index {
+  %c1 = arith.constant 1 : index
+  %0 = util.assume.int %ind<umin = 1> : index
+  %1 = arith.select %bool, %0, %c1 : index
+  // select should not be optimized away.
+  // CHECK: arith.select
+  util.return %1 : index
+}