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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-arith

@llvm/pr-subscribers-mlir

Author: Anutosh Bhat (anutosh491)

<details>
<summary>Changes</summary>

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.



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


2 Files Affected:

- (modified) mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp (+22-7) 
- (modified) mlir/test/Dialect/Arith/int-range-opts-crash.mlir (+14) 


``````````diff
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

``````````

</details>


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