[Mlir-commits] [mlir] [mlir][arith] Fix SelectOp int range unsafe inference with unfinalized case ranges (PR #173716)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-arith

Author: Longsheng Du (LongshengDu)

<details>
<summary>Changes</summary>

This PR fixes a bug in `arith::SelectOp::inferResultRangesFromOptional` where uninitialized SelectOp branch int ranges were incorrectly joined with initialized int ranges during dataflow analysis, leading to incorrect folding in `-int-range-optimizations`.

**The Issue:**
When a `arith.select` branch has an uninitialized range (e.g., from an op like `nvvm.read.ptx.sreg.cluster.ctaid.x`, `scf.switch`, `llvm.call`, etc. that lacks range inference), the analysis computed `IntegerValueRange::join(Uninitialized, Constant) = Constant`. This caused the `arith.select` to be replaced with the constant, ignoring the dynamic branch.

**Example:**
```mlir
// The bug before fix: -int-range-optimizations replaces %1 with %c32
// led to incorrect results and unsafe behaviours
%0 = nvvm.read.ptx.sreg.cluster.ctaid.x : i32 // Uninitialized int range
%c32 = arith.constant 32 : i32
%1 = arith.select %cond, %0, %c32 : i32
```

**The Fix:**
Explicitly ensure `inferResultRangesFromOptional` all select cases have initialized ranges before combining them. If any case is uninitialized, the result is now treated as uninitialized. Also added `test.without_bounds` op to simulate and test uninitialized ranges.


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


4 Files Affected:

- (modified) mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp (+5-1) 
- (modified) mlir/test/Dialect/Arith/int-range-interface.mlir (+18) 
- (modified) mlir/test/lib/Dialect/Test/TestOpDefs.cpp (+10) 
- (modified) mlir/test/lib/Dialect/Test/TestOps.td (+17) 


``````````diff
diff --git a/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp b/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp
index 7673185487eef..4d20430b39f91 100644
--- a/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp
+++ b/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp
@@ -329,7 +329,11 @@ void arith::SelectOp::inferResultRangesFromOptional(
       setResultRange(getResult(), trueCase);
     return;
   }
-  setResultRange(getResult(), IntegerValueRange::join(trueCase, falseCase));
+
+  if (trueCase.isUninitialized() || falseCase.isUninitialized())
+    setResultRange(getResult(), IntegerValueRange{});
+  else
+    setResultRange(getResult(), IntegerValueRange::join(trueCase, falseCase));
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Arith/int-range-interface.mlir b/mlir/test/Dialect/Arith/int-range-interface.mlir
index 130782ba9f525..30b7128dab42c 100644
--- a/mlir/test/Dialect/Arith/int-range-interface.mlir
+++ b/mlir/test/Dialect/Arith/int-range-interface.mlir
@@ -663,6 +663,24 @@ func.func @select_union(%arg0 : index, %arg1 : i1) -> i1 {
     func.return %5 : i1
 }
 
+// CHECK-LABEL: func @select_undefined_union
+// CHECK-COUNT-2: arith.select
+// CHECK: %[[ret:.*]] = arith.cmpi eq
+// CHECK: return %[[ret]]
+
+func.func @select_undefined_union(%arg0: i1) -> i1 {
+  %c32 = arith.constant 32 : index
+  %c64 = arith.constant 64 : index
+  %0 = test.without_bounds : index
+  %1 = arith.select %arg0, %0, %c64 : index
+  %2 = arith.cmpi eq, %1, %c64 : index
+  %3 = test.without_bounds : index
+  %4 = arith.select %2, %c32, %3 : index
+  %5 = arith.cmpi eq, %4, %c32 : index
+
+  return %5 : i1
+}
+
 // CHECK-LABEL: func @if_union
 // CHECK: %[[true:.*]] = arith.constant true
 // CHECK: return %[[true]]
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 868926520af05..110f83c75ef00 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -863,6 +863,16 @@ void TestWithBoundsOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
   setResultRanges(getResult(), {getUmin(), getUmax(), getSmin(), getSmax()});
 }
 
+//===----------------------------------------------------------------------===//
+// TestWithoutBoundsOp
+//===----------------------------------------------------------------------===//
+
+void TestWithoutBoundsOp::inferResultRangesFromOptional(
+    ArrayRef<IntegerValueRange> argRanges, SetIntLatticeFn setResultRanges) {
+  // mimic ops with uninitialized range
+  setResultRanges(getResult(), IntegerValueRange{});
+}
+
 //===----------------------------------------------------------------------===//
 // TestWithBoundsRegionOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 5417ae94f00d7..bceb49ebe17f6 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -3182,6 +3182,23 @@ def TestWithBoundsOp : TEST_Op<"with_bounds",
   let assemblyFormat = "attr-dict `:` type($fakeVal)";
 }
 
+def TestWithoutBoundsOp : TEST_Op<"without_bounds",
+                         [DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRangesFromOptional"]>,
+                         NoMemoryEffect]> {
+  let description = [{
+    Creates a value with uninitialized range for integer range analysis tests.
+
+    Example:
+
+    ```mlir
+    %0 = test.without_bounds : index
+    ```
+  }];
+  let results = (outs InferIntRangeType:$result);
+
+  let assemblyFormat = "attr-dict `:` type($result)";
+}
+
 def TestWithBoundsRegionOp : TEST_Op<"with_bounds_region",
                           [DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>,
                            SingleBlock, NoTerminator]> {

``````````

</details>


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