[Mlir-commits] [mlir] [MLIR][IntRange] Materialize known constant values only when the value type is integer-like (PR #196133)

[Jeremy Kun]

https://github.com/j2kun updated https://github.com/llvm/llvm-project/pull/196133

>From 2005ad12f7d7cdc491e9adae8365c0fd319c9952 Mon Sep 17 00:00:00 2001
From: Jeremy Kun <jkun at google.com>
Date: Wed, 6 May 2026 10:27:25 -0700
Subject: [PATCH 1/3] add reproducing crash

---
 .../Dialect/Arith/int-range-opts-crash.mlir   | 13 ++++++++++
 mlir/test/lib/Dialect/Test/TestOpDefs.cpp     | 24 +++++++++++++++++++
 mlir/test/lib/Dialect/Test/TestOps.td         | 11 +++++++++
 3 files changed, 48 insertions(+)
 create mode 100644 mlir/test/Dialect/Arith/int-range-opts-crash.mlir

diff --git a/mlir/test/Dialect/Arith/int-range-opts-crash.mlir b/mlir/test/Dialect/Arith/int-range-opts-crash.mlir
new file mode 100644
index 0000000000000..fa763c163160d
--- /dev/null
+++ b/mlir/test/Dialect/Arith/int-range-opts-crash.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-opt -int-range-optimizations %s | FileCheck %s
+
+// CHECK-LABEL: func.func @repro_crash() -> !test.i32 {
+func.func @repro_crash() -> !test.i32 {
+  %cst = arith.constant 1 : i32
+  // CHECK: %[[RES:.*]] = test.region_type_changer
+  %0 = "test.region_type_changer"(%cst) ({
+  ^bb0(%arg0: i32):
+    "test.types_compat_yield"(%arg0) : (i32) -> ()
+  }) : (i32) -> !test.i32
+  // CHECK: return %[[RES]] : !test.i32
+  return %0 : !test.i32
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 340b44b14dd96..73eed41acc3e8 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -1546,6 +1546,30 @@ bool TestLoopTypesCompatOp::areTypesCompatible(Type lhs, Type rhs) {
   return lhs == rhs || (isa<IntegerType>(lhs) && isa<IntegerType>(rhs));
 }
 
+void TestRegionTypeChangerOp::getSuccessorRegions(
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
+  if (point.isParent())
+    regions.emplace_back(&getBody());
+  else
+    regions.push_back(RegionSuccessor::parent());
+}
+
+OperandRange
+TestRegionTypeChangerOp::getEntrySuccessorOperands(RegionSuccessor) {
+  return getEntries();
+}
+
+ValueRange
+TestRegionTypeChangerOp::getSuccessorInputs(RegionSuccessor successor) {
+  if (successor.isParent())
+    return getResults();
+  return getBody().getArguments();
+}
+
+bool TestRegionTypeChangerOp::areTypesCompatible(Type lhs, Type rhs) {
+  return true;
+}
+
 //===----------------------------------------------------------------------===//
 // TestVersionedOpA
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 348ff5d7f4ea0..90cda01f2704b 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2799,6 +2799,17 @@ def TestLoopTypesCompatOp : TEST_Op<"loop_types_compat",
   let regions = (region SizedRegion<1>:$body);
 }
 
+def TestRegionTypeChangerOp : TEST_Op<"region_type_changer",
+    [DeclareOpInterfaceMethods<RegionBranchOpInterface,
+       ["getEntrySuccessorOperands", "getSuccessorInputs",
+        "areTypesCompatible"]>,
+     RecursiveMemoryEffects]> {
+  let arguments = (ins Variadic<AnyType>:$entries);
+  let results = (outs Variadic<AnyType>:$results);
+  let regions = (region SizedRegion<1>:$body);
+  let assemblyFormat = "`(` $entries `)` $body `:` functional-type($entries, $results) attr-dict";
+}
+
 //===----------------------------------------------------------------------===//
 // Test TableGen generated build() methods
 //===----------------------------------------------------------------------===//

>From 890a8b73960f2d317f863ea7e8d7d63b49e4d226 Mon Sep 17 00:00:00 2001
From: Jeremy Kun <jkun at google.com>
Date: Wed, 6 May 2026 10:39:00 -0700
Subject: [PATCH 2/3] Add a guard for attempts to materialize non-integer types

---
 .../Arith/Transforms/IntRangeOptimizations.cpp       | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
index 85578c22799c6..559073951c3e3 100644
--- a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
@@ -71,6 +71,10 @@ LogicalResult maybeReplaceWithConstant(DataFlowSolver &solver,
     return failure();
 
   Type type = value.getType();
+  // If the type or element type is non-integral, the attribute constructor
+  // will crash, so eagerly check for an integer type to avoid this.
+  if (!getElementTypeOrSelf(type).isIntOrIndex())
+    return failure();
   Location loc = value.getLoc();
   Operation *maybeDefiningOp = value.getDefiningOp();
   Dialect *valueDialect =
@@ -133,8 +137,14 @@ 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.
     auto needsReplacing = [&](Value v) {
-      return getMaybeConstantValue(solver, v).has_value() && !v.use_empty();
+      return getElementTypeOrSelf(v.getType()).isIntOrIndex() &&
+             getMaybeConstantValue(solver, v).has_value() && !v.use_empty();
     };
     bool hasConstantResults = llvm::any_of(op->getResults(), needsReplacing);
     if (op->getNumRegions() == 0)

>From c6e44b3c983be0b2ba14717fdaab78f8e64b71da Mon Sep 17 00:00:00 2001
From: Jeremy Kun <jkun at google.com>
Date: Wed, 6 May 2026 14:54:31 -0700
Subject: [PATCH 3/3] format

---
 .../Dialect/Arith/Transforms/IntRangeOptimizations.cpp | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
index 559073951c3e3..9fcda39089b2c 100644
--- a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
@@ -137,11 +137,11 @@ 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() 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.
     auto needsReplacing = [&](Value v) {
       return getElementTypeOrSelf(v.getType()).isIntOrIndex() &&
              getMaybeConstantValue(solver, v).has_value() && !v.use_empty();