[Mlir-commits] [mlir] [mlir][func] Fix a crash in DuplicateFunctionEliminationPass (PR #209667)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Longsheng Mou (CoTinker)

<details>
<summary>Changes</summary>

Previously, we used `SymbolUserMap::replaceAllUsesWith` to replace symbols, but this could not update the symbol table cached by `SymbolUserMap` during traversal, leading to a crash. This PR switches to `SymbolTable::replaceAllSymbolUses`, which always operates on the latest symbol table and avoids the crash. Fixes #<!-- -->209648.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp (+3-3) 
- (modified) mlir/test/Dialect/Func/duplicate-function-elimination.mlir (+26) 


``````````diff
diff --git a/mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp b/mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp
index a2186e433f544..c480e469fdcc8 100644
--- a/mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp
+++ b/mlir/lib/Dialect/Func/Transforms/DuplicateFunctionElimination.cpp
@@ -101,12 +101,12 @@ struct DuplicateFunctionEliminationPass
 
     // Update all symbol uses to reference unique func op
     // representants and erase redundant func ops.
-    SymbolTableCollection symbolTable;
-    SymbolUserMap userMap(symbolTable, module);
     for (auto it : toBeErased) {
       StringAttr oldSymbol = it.getSymNameAttr();
       StringAttr newSymbol = getRepresentant[oldSymbol].getSymNameAttr();
-      userMap.replaceAllUsesWith(it, newSymbol);
+      if (failed(
+              SymbolTable::replaceAllSymbolUses(oldSymbol, newSymbol, module)))
+        return signalPassFailure();
       it.erase();
     }
   }
diff --git a/mlir/test/Dialect/Func/duplicate-function-elimination.mlir b/mlir/test/Dialect/Func/duplicate-function-elimination.mlir
index 4d00d8a954d17..b6588b39c5a87 100644
--- a/mlir/test/Dialect/Func/duplicate-function-elimination.mlir
+++ b/mlir/test/Dialect/Func/duplicate-function-elimination.mlir
@@ -412,3 +412,29 @@ func.func @user(%arg0: tensor<f32>) -> tensor<f32> {
 // CHECK:     @user
 // CHECK-2:     constant @identity
 // CHECK:       call @identity
+
+// -----
+
+func.func @callee0() {
+  return
+}
+
+func.func @caller0() {
+  call @callee1() : () -> ()
+  return
+}
+
+func.func @caller1() {
+  call @callee1() : () -> ()
+  return
+}
+
+func.func @callee1() {
+  return
+}
+
+// CHECK:     @callee0
+// CHECK:     @caller0
+// CHECK:       call @callee0()
+// CHECK-NOT: @caller1
+// CHECK-NOT: @callee1

``````````

</details>


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