[Mlir-commits] [mlir] [mlir][func] Fix multiple bugs in `DuplicateFunctionElimination` (PR #109571)

[Christopher Bate]

================
@@ -97,14 +101,14 @@ struct DuplicateFunctionEliminationPass
       }
     });
 
-    // Update call ops to call unique func op representants.
-    module.walk([&](func::CallOp callOp) {
-      func::FuncOp callee = getRepresentant[callOp.getCalleeAttr().getAttr()];
-      callOp.setCallee(callee.getSymName());
-    });
-
-    // Erase redundant func ops.
+    // Update all symbol uses to reference unique func op
+    // representants and erase redundant func ops.
     for (auto it : toBeErased) {
+      auto oldSymbol = it.getSymNameAttr();
+      auto newSymbol = getRepresentant[oldSymbol].getSymNameAttr();
+      if (failed(
+              SymbolTable::replaceAllSymbolUses(oldSymbol, newSymbol, module)))
----------------
christopherbate wrote:

IIUC, the issue for performance looks to be the `module` in ` SymbolTable::replaceAllSymbolUses(oldSymbol, newSymbol, module)))`. That means it needs to traverse all the ops in `module` and look for uses of `oldSymbol` that can be replaced. 

The class `SymbolUserMap` pre-computes this information for just this use case. Create the `SymbolUserMap` and then call `SymbolUserMap::replaceAllUsesWith` in the loop. That will iterate over the relevant users and update them rather than iterating over the whole module.



```
SymbolTableCollection collection;
SymolUserMap userMap(collection, module);

for(....) {
    ....
    userMap.replace(old, new);
}
```







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