[Mlir-commits] [mlir] [mlir][test] Fix crash in TestWithBoundsOp with mismatched attribute width (PR #184093)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/184093

`TestWithBoundsOp::inferResultRanges` was passing the stored APInt bound values directly to `setResultRanges` without checking whether their bit width matched the result type. When bounds were stored using a wider type (e.g., `10 : i64` for an `i8` result), the `ConstantIntRanges` invariant was violated: the APInts had width 64 while the integer type had width 8.

Downstream consumers such as `TestReflectBoundsOp::inferResultRanges` would then call `Builder::getIntegerAttr(i8Type, 64-bit-APInt)`, which asserts that the bit widths match, causing a crash.

Fix by truncating or extending the stored APInts to match the result type's storage width in `TestWithBoundsOp::inferResultRanges`.

Fixes #120882

>From 82e6bcdbf42ecd263f449a62144f9cf2fa5581d2 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Sun, 1 Mar 2026 16:54:47 -0800
Subject: [PATCH] [mlir][test] Fix crash in TestWithBoundsOp with mismatched
 attribute width

`TestWithBoundsOp::inferResultRanges` was passing the stored APInt bound
values directly to `setResultRanges` without checking whether their bit
width matched the result type. When bounds were stored using a wider type
(e.g., `10 : i64` for an `i8` result), the `ConstantIntRanges` invariant
was violated: the APInts had width 64 while the integer type had width 8.

Downstream consumers such as `TestReflectBoundsOp::inferResultRanges`
would then call `Builder::getIntegerAttr(i8Type, 64-bit-APInt)`, which
asserts that the bit widths match, causing a crash.

Fix by truncating or extending the stored APInts to match the result
type's storage width in `TestWithBoundsOp::inferResultRanges`.

Fixes #120882
---
 .../infer-int-range-test-ops.mlir             | 12 +++++++++
 mlir/test/lib/Dialect/Test/TestOpDefs.cpp     | 25 ++++++++++++++++++-
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/mlir/test/Interfaces/InferIntRangeInterface/infer-int-range-test-ops.mlir b/mlir/test/Interfaces/InferIntRangeInterface/infer-int-range-test-ops.mlir
index c6344447d9f74..90386052f7107 100644
--- a/mlir/test/Interfaces/InferIntRangeInterface/infer-int-range-test-ops.mlir
+++ b/mlir/test/Interfaces/InferIntRangeInterface/infer-int-range-test-ops.mlir
@@ -200,3 +200,15 @@ func.func @multiple_loop_ivs(%arg0: memref<?x64xi32>) {
   }
   return
 }
+
+// Verify that test.with_bounds with mismatched attribute width (e.g., i64
+// bounds for an i8 result) does not crash but instead truncates the bounds
+// to the result type's width. See: https://github.com/llvm/llvm-project/issues/120882
+// CHECK-LABEL: func @with_bounds_mismatched_width
+// CHECK: test.reflect_bounds {smax = 15 : si8, smin = 10 : si8, umax = 15 : ui8, umin = 10 : ui8}
+func.func @with_bounds_mismatched_width() -> i8 {
+  %0 = test.with_bounds { umin = 10 : i64, umax = 15 : i64,
+                           smin = 10 : i64, smax = 15 : i64 } : i8
+  %1 = test.reflect_bounds %0 : i8
+  return %1 : i8
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index d2826cc42b270..7cf608b66c657 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -878,7 +878,30 @@ LogicalResult TestVerifiersOp::verifyRegions() {
 
 void TestWithBoundsOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                          SetIntRangeFn setResultRanges) {
-  setResultRanges(getResult(), {getUmin(), getUmax(), getSmin(), getSmax()});
+  APInt umin = getUmin(), umax = getUmax(), smin = getSmin(), smax = getSmax();
+  // Ensure the APInt widths match the result type's expected storage width.
+  // The stored attributes may use a different width (e.g., i64 bounds for an
+  // i8 result), which would violate the ConstantIntRanges invariant.
+  Type type = getElementTypeOrSelf(getResult().getType());
+  unsigned targetWidth = 0;
+  if (type.isIndex())
+    targetWidth = IndexType::kInternalStorageBitWidth;
+  else if (auto intTy = llvm::dyn_cast<IntegerType>(type))
+    targetWidth = intTy.getWidth();
+  if (targetWidth != 0 && umin.getBitWidth() != targetWidth) {
+    if (umin.getBitWidth() > targetWidth) {
+      umin = umin.trunc(targetWidth);
+      umax = umax.trunc(targetWidth);
+      smin = smin.trunc(targetWidth);
+      smax = smax.trunc(targetWidth);
+    } else {
+      umin = umin.zext(targetWidth);
+      umax = umax.zext(targetWidth);
+      smin = smin.sext(targetWidth);
+      smax = smax.sext(targetWidth);
+    }
+  }
+  setResultRanges(getResult(), {umin, umax, smin, smax});
 }
 
 //===----------------------------------------------------------------------===//