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

[llvmlistbot at llvm.org]

Author: Jeremy Kun
Date: 2026-05-06T18:54:39-07:00
New Revision: b07aeece4d8b3af5578c58a2a2d4092070b92773

URL: https://github.com/llvm/llvm-project/commit/b07aeece4d8b3af5578c58a2a2d4092070b92773
DIFF: https://github.com/llvm/llvm-project/commit/b07aeece4d8b3af5578c58a2a2d4092070b92773.diff

LOG: [MLIR][IntRange] Materialize known constant values only when the value type is integer-like (#196133)

This patch came out of
https://discourse.llvm.org/t/request-for-advice-in-updating-intrange-optimziations/90685

In short, my out-of-tree dialect has region-bearing operations whose
return-like terminators semantically wrap an integer type in a
non-integer type when returning to the parent for control flow (also
this dialect that has no constant materializer). In this case,
`MaterializeKnownConstantValues` will attempt to construct integer
attributes using those non-integer types, which causes a stack trace and
prevents us from using int range optimizations on IR that contains these
ops.

This patch adds guards to the materialization step, so that if the value
that is inferred to be constant is not an integer (or shaped container
with an integer element type), no materialization is attempted.

Added: 
    mlir/test/Dialect/Arith/int-range-opts-crash.mlir

Modified: 
    mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
    mlir/test/lib/Dialect/Test/TestOpDefs.cpp
    mlir/test/lib/Dialect/Test/TestOps.td

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
index 85578c22799c6..9fcda39089b2c 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)

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 b5a6cd535d9f2..edd120de3b667 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2808,6 +2808,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
 //===----------------------------------------------------------------------===//