[Mlir-commits] [mlir] 6867324 - [mlir][bufferization] Improve performance of DropEquivalentBufferResultsPass (#101281)

[llvmlistbot at llvm.org]

Author: Longsheng Mou
Date: 2024-08-02T09:22:20+08:00
New Revision: 6867324eeec7c4f297c2f787d9c7b4d751a384c7

URL: https://github.com/llvm/llvm-project/commit/6867324eeec7c4f297c2f787d9c7b4d751a384c7
DIFF: https://github.com/llvm/llvm-project/commit/6867324eeec7c4f297c2f787d9c7b4d751a384c7.diff

LOG: [mlir][bufferization] Improve performance of DropEquivalentBufferResultsPass (#101281)

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

Added: 
    

Modified: 
    mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp

Removed: 
    


################################################################################
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();