[Mlir-commits] [mlir] f849866 - [MLIR] Reduce complexity of searching circular function calls in bufferization (#142099)

[llvmlistbot at llvm.org]

Author: Michele Scuttari
Date: 2025-06-06T10:35:58+02:00
New Revision: f849866fc5e62091088e3e2b58214c614437ff68

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

LOG: [MLIR] Reduce complexity of searching circular function calls in bufferization (#142099)

The current algorithm searching for circular function calls scales quadratically due to the linear scan of the functions vector that is performed for each element of the vector itself. The PR replaces such algorithm with an O(V + E) version based on the Khan's algorithm for topological sorting, where V is the number of functions and E is the number of function calls.

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
index fc6424a25ac70..d1d106220a38c 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
@@ -341,15 +341,25 @@ static LogicalResult getFuncOpsOrderedByCalls(
 
   // Iteratively remove function operations that do not call any of the
   // functions remaining in the callCounter map and add them to ordered list.
-  while (!numberCallOpsContainedInFuncOp.empty()) {
-    auto it = llvm::find_if(numberCallOpsContainedInFuncOp,
-                            [](auto entry) { return entry.getSecond() == 0; });
-    if (it == numberCallOpsContainedInFuncOp.end())
-      break;
-    orderedFuncOps.push_back(it->getFirst());
-    for (auto callee : calledBy[it->getFirst()])
-      numberCallOpsContainedInFuncOp[callee]--;
-    numberCallOpsContainedInFuncOp.erase(it);
+  SmallVector<func::FuncOp> worklist;
+
+  for (const auto &entry : numberCallOpsContainedInFuncOp) {
+    if (entry.second == 0)
+      worklist.push_back(entry.first);
+  }
+
+  while (!worklist.empty()) {
+    func::FuncOp func = worklist.pop_back_val();
+    orderedFuncOps.push_back(func);
+
+    for (func::FuncOp caller : calledBy[func]) {
+      auto &count = numberCallOpsContainedInFuncOp[caller];
+
+      if (--count == 0)
+        worklist.push_back(caller);
+    }
+
+    numberCallOpsContainedInFuncOp.erase(func);
   }
 
   // Put all other functions in the list of remaining functions. These are