[mlir][bufferization] Improve performance of DropEquivalentBufferResultsPass

[CoTinker]

By using DenseMap to minimize the traveral time of callOps, and the efficiency of running this pass has been greatly improved.

[llvmbot]

@llvm/pr-subscribers-mlir-bufferization

@llvm/pr-subscribers-mlir

Author: Longsheng Mou (CoTinker)

Changes

By using DenseMap to minimize the traveral time of callOps, and the efficiency of running this pass has been greatly improved.

---

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

1 Files Affected:

• (modified) mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp (+10-6)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp b/mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp
index 016ec2be62dce..d86bdb20a66bb 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp
@@ -71,6 +71,14 @@ LogicalResult
 mlir::bufferization::dropEquivalentBufferResults(ModuleOp module) {
   IRRewriter rewriter(module.getContext());
 
+  DenseMap<func::FuncOp, DenseSet<func::CallOp>> callerMap;
+  // Collect the mapping of functions to their call sites.
+  module.walk([&](func::CallOp callOp) {
+    if (func::FuncOp calledFunc = getCalledFunction(callOp)) {
+      callerMap[calledFunc].insert(callOp);
+    }
+  });
+
   for (auto funcOp : module.getOps<func::FuncOp>()) {
     if (funcOp.isExternal())
       continue;
@@ -109,10 +117,7 @@ mlir::bufferization::dropEquivalentBufferResults(ModuleOp module) {
     returnOp.getOperandsMutable().assign(newReturnValues);
 
     // Update function calls.
-    module.walk([&](func::CallOp callOp) {
-      if (getCalledFunction(callOp) != funcOp)
-        return WalkResult::skip();
-
+    for (func::CallOp callOp : callerMap[funcOp]) {
       rewriter.setInsertionPoint(callOp);
       auto newCallOp = rewriter.create<func::CallOp>(callOp.getLoc(), funcOp,
                                                      callOp.getOperands());
@@ -136,8 +141,7 @@ mlir::bufferization::dropEquivalentBufferResults(ModuleOp module) {
         newResults.push_back(replacement);
       }
       rewriter.replaceOp(callOp, newResults);
-      return WalkResult::advance();
-    });
+    }
   }
 
   return success();

[CoTinker]

Related Discussions: https://discourse.llvm.org/t/mlir-bufferization-improve-performance-of-dropequivalentbufferresults-pass/80453
Currently, the pass execution time has been reduced from 1 hours to less than 5s. @matthias-springer