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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Jeremy Kun (j2kun)

<details>
<summary>Changes</summary>

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 that wrap an integer type in a non-integer type (for a dialect that has no constant materializer), and `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.

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.

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


4 Files Affected:

- (modified) mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp (+11-1) 
- (added) mlir/test/Dialect/Arith/int-range-opts-crash.mlir (+13) 
- (modified) mlir/test/lib/Dialect/Test/TestOpDefs.cpp (+24) 
- (modified) mlir/test/lib/Dialect/Test/TestOps.td (+11) 


``````````diff
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)
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
 //===----------------------------------------------------------------------===//

``````````

</details>


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