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

[Longsheng Du]

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

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.


>From 65f5b35ed24b14bb60b2ed67de731bbe03c39ea5 Mon Sep 17 00:00:00 2001
From: Longsheng Du <longshengd at nvidia.com>
Date: Sat, 27 Dec 2025 03:43:39 -0800
Subject: [PATCH] fix select range

---
 .../Arith/IR/InferIntRangeInterfaceImpls.cpp   |  6 +++++-
 .../Dialect/Arith/int-range-interface.mlir     | 18 ++++++++++++++++++
 mlir/test/lib/Dialect/Test/TestOpDefs.cpp      | 10 ++++++++++
 mlir/test/lib/Dialect/Test/TestOps.td          | 17 +++++++++++++++++
 4 files changed, 50 insertions(+), 1 deletion(-)

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]> {