[Mlir-commits] [mlir] [mlir][arith] Fix APInt bitwidth mismatch crash in int-range-optimizations (PR #205110)

[Anutosh Bhat]

https://github.com/anutosh491 created https://github.com/llvm/llvm-project/pull/205110

Fixes https://github.com/llvm/llvm-project/issues/204909

When an op's `areTypesCompatible()` hook accepts integers of different widths across a region boundary, the range analysis can propagate a constant range whose APInt bitwidth does not match the IR type of the destination value.
This caused `IntegerAttr::get` to `assert` in `maybeReplaceWithConstant`.

Fix by bailing out in `maybeReplaceWithConstant` when the bitwidths mismatch, and adding the same check to the needsReplacing lambda in matchAndRewrite.

The second guard is necessary to mirror the existing isIntOrIndex() guard — without it the pattern claims success without changing the IR, causing the greedy rewrite driver to loop.



>From ac2e6554e07e3c3e9c847f361b14f0f65a66d385 Mon Sep 17 00:00:00 2001
From: anutosh491 <andersonbhat491 at gmail.com>
Date: Mon, 22 Jun 2026 19:23:38 +0530
Subject: [PATCH] [mlir][arith] Fix APInt bitwidth mismatch crash in
 int-range-optimizations

---
 .../Transforms/IntRangeOptimizations.cpp      | 29 ++++++++++++++-----
 .../Dialect/Arith/int-range-opts-crash.mlir   | 14 +++++++++
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
index 9fcda39089b2c..298c0dc2f3bda 100644
--- a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
@@ -75,6 +75,13 @@ LogicalResult maybeReplaceWithConstant(DataFlowSolver &solver,
   // will crash, so eagerly check for an integer type to avoid this.
   if (!getElementTypeOrSelf(type).isIntOrIndex())
     return failure();
+
+  // Bail out if the inferred APInt bitwidth does not match the storage width
+  // of the IR type; IntegerAttr::get would assert otherwise.
+  unsigned storageWidth = ConstantIntRanges::getStorageBitwidth(type);
+  if (storageWidth != 0 && maybeConstValue->getBitWidth() != storageWidth)
+    return failure();
+
   Location loc = value.getLoc();
   Operation *maybeDefiningOp = value.getDefiningOp();
   Dialect *valueDialect =
@@ -137,14 +144,22 @@ struct MaterializeKnownConstantValues : public RewritePattern {
     if (matchPattern(op, m_Constant()))
       return failure();
 
-    // We need to check isIntOrIndex() here as well to avoid infinite loops in
-    // the greedy pattern rewriter. If we only check it in
-    // maybeReplaceWithConstant, this lambda might still return true for
-    // non-integral types, causing the pattern to match and claim success
-    // without making any changes, leading to non-convergence.
+    // We need to check isIntOrIndex() and APInt bitwidth compatibility here
+    // as well to avoid infinite loops in the greedy pattern rewriter. If we
+    // only check in maybeReplaceWithConstant, this lambda might still return
+    // true for values that cannot be materialized, causing the pattern to
+    // match and claim success without making any changes, leading to
+    // non-convergence.
     auto needsReplacing = [&](Value v) {
-      return getElementTypeOrSelf(v.getType()).isIntOrIndex() &&
-             getMaybeConstantValue(solver, v).has_value() && !v.use_empty();
+      if (!getElementTypeOrSelf(v.getType()).isIntOrIndex())
+        return false;
+      std::optional<APInt> maybeConstValue = getMaybeConstantValue(solver, v);
+      if (!maybeConstValue.has_value() || v.use_empty())
+        return false;
+      unsigned storageWidth =
+          ConstantIntRanges::getStorageBitwidth(v.getType());
+      return storageWidth == 0 ||
+             maybeConstValue->getBitWidth() == storageWidth;
     };
     bool hasConstantResults = llvm::any_of(op->getResults(), needsReplacing);
     if (op->getNumRegions() == 0)
diff --git a/mlir/test/Dialect/Arith/int-range-opts-crash.mlir b/mlir/test/Dialect/Arith/int-range-opts-crash.mlir
index fa763c163160d..99aa03714edb9 100644
--- a/mlir/test/Dialect/Arith/int-range-opts-crash.mlir
+++ b/mlir/test/Dialect/Arith/int-range-opts-crash.mlir
@@ -1,5 +1,19 @@
 // RUN: mlir-opt -int-range-optimizations %s | FileCheck %s
 
+// CHECK-LABEL: func.func @repro_bitwidth_mismatch
+func.func @repro_bitwidth_mismatch() -> i32 {
+  %c0_i32 = arith.constant 0 : i32
+  // CHECK: test.region_types_compat
+  %0 = "test.region_types_compat"(%c0_i32) ({
+  ^bb0(%arg0: i64):
+    %c1_i64 = arith.constant 1 : i64
+    test.types_compat_yield %c1_i64 : i64
+  }) : (i32) -> i32
+  return %0 : i32
+}
+
+// -----
+
 // CHECK-LABEL: func.func @repro_crash() -> !test.i32 {
 func.func @repro_crash() -> !test.i32 {
   %cst = arith.constant 1 : i32