[Util] Fix AssumeIntOp::inferResultRanges bug (iree-org#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]>
nod-ai · Nov 19, 2024 · 26ef79a · 26ef79a
1 parent 47432c6
commit 26ef79a
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/compiler/src/iree/compiler/Dialect/Util/IR/UtilOps.cpp b/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));
   }
 }
 

diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/test/optimize_int_arithmetic.mlir b/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
+}