[Mlir-commits] [mlir] [mlir][Builtin] Fix crash folding self-referential unrealized_conversion_cast in graph regions (PR #207185)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Anutosh Bhat (anutosh491)

<details>
<summary>Changes</summary>

Fix #<!-- -->205064

Graph regions permit use-def cycles. When an `unrealized_conversion_cast`uses its own result as an operand and the input/output types match, the fold returned success with the operand (= the result itself) as the replacement value. `replaceOp` then called `replaceAllUsesWith(%0, %0)` (a no-op) followed by `eraseOp`, which asserted `use_empty()` because the self-use was still live.

This can be fixed by bailing out of the fold when any operand is produced by this op itself. The greedy rewriter then correctly DCEs the op by dropping the self-use before erasing. As a result we have the following 
```
 anutosh491@<!-- -->Anutoshs-MacBook-Air mlir-build % cat bin/a.mlir                                     
module {
  test.graph_region {
    %0 = builtin.unrealized_conversion_cast %0 : i32 to i32
    "test.return"() : () -> ()
  }
}%
                                                                                                                                          
 anutosh491@<!-- -->Anutoshs-MacBook-Air mlir-build % ./bin/mlir-opt -canonicalize bin/a.mlir            
module {
  test.graph_region {
    "test.return"() : () -> ()
  }
}

 anutosh491@<!-- -->Anutoshs-MacBook-Air mlir-build % ./bin/llvm-lit ../mlir/test/Dialect/Builtin/canonicalize.mlir -v
-- Testing: 1 tests, 1 workers --
PASS: MLIR :: Dialect/Builtin/canonicalize.mlir (1 of 1)

Testing Time: 0.12s

Total Discovered Tests: 1
  Passed: 1 (100.00%)
```

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


2 Files Affected:

- (modified) mlir/lib/IR/BuiltinDialect.cpp (+5) 
- (modified) mlir/test/Dialect/Builtin/canonicalize.mlir (+11) 


``````````diff
diff --git a/mlir/lib/IR/BuiltinDialect.cpp b/mlir/lib/IR/BuiltinDialect.cpp
index c88b328282275..00cad06122dd8 100644
--- a/mlir/lib/IR/BuiltinDialect.cpp
+++ b/mlir/lib/IR/BuiltinDialect.cpp
@@ -201,6 +201,11 @@ UnrealizedConversionCastOp::fold(FoldAdaptor adaptor,
   ResultRange results = getOutputs();
 
   if (operands.getType() == results.getType()) {
+    // Skip if any operand is this op's own result (self-referential cycle).
+    if (llvm::any_of(operands, [this](Value v) {
+          return v.getDefiningOp() == getOperation();
+        }))
+      return failure();
     foldResults.append(operands.begin(), operands.end());
     return success();
   }
diff --git a/mlir/test/Dialect/Builtin/canonicalize.mlir b/mlir/test/Dialect/Builtin/canonicalize.mlir
index 2e36b7ee371c3..537868e990d78 100644
--- a/mlir/test/Dialect/Builtin/canonicalize.mlir
+++ b/mlir/test/Dialect/Builtin/canonicalize.mlir
@@ -23,3 +23,14 @@ func.func @multiple_conversion_casts_failure(%arg0: i32, %arg1: i32, %arg2: i64)
   %outputs:2 = builtin.unrealized_conversion_cast %arg2, %inputs#1 : i64, i64 to i32, i32
   return %outputs#0, %outputs#1 : i32, i32
 }
+
+// Test a self-referential cast in a graph region must not crash.
+// CHECK-LABEL: func @self_referential_cast
+// CHECK-NOT: unrealized_conversion_cast
+func.func @self_referential_cast() {
+  test.graph_region {
+    %0 = builtin.unrealized_conversion_cast %0 : i32 to i32
+    "test.return"() : () -> ()
+  }
+  return
+}

``````````

</details>


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