[Mlir-commits] [mlir] [mlir] Deterministic containers in OneShotModuleBufferize (PR #184722)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-bufferization

Author: Colin He (CPlusMinus2000)

<details>
<summary>Changes</summary>

Iteration over funcOps in `getFuncOpsOrderedByCalls` is non-deterministic as a result of using Dense containers. Replacing with Vector-backed containers restores deterministic behaviour.

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


1 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp (+4-3) 


``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
index 1cced539b5318..d29150a7403f9 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
@@ -70,6 +70,7 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Operation.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallVectorExtras.h"
 
 using namespace mlir;
@@ -316,9 +317,9 @@ static LogicalResult getFuncOpsOrderedByCalls(
     SymbolTableCollection &symbolTables) {
   // For each FuncOp, the set of functions called by it (i.e. the union of
   // symbols of all nested func::CallOp).
-  DenseMap<func::FuncOp, DenseSet<func::FuncOp>> calledBy;
+  DenseMap<func::FuncOp, SetVector<func::FuncOp>> calledBy;
   // For each FuncOp, the number of func::CallOp it contains.
-  DenseMap<func::FuncOp, unsigned> numberCallOpsContainedInFuncOp;
+  llvm::MapVector<func::FuncOp, unsigned> numberCallOpsContainedInFuncOp;
   for (mlir::Region &region : moduleOp->getRegions()) {
     for (mlir::Block &block : region.getBlocks()) {
       for (func::FuncOp funcOp : block.getOps<func::FuncOp>()) {
@@ -333,7 +334,7 @@ static LogicalResult getFuncOpsOrderedByCalls(
             return WalkResult::skip();
 
           callerMap[calledFunction].insert(callOp);
-          if (calledBy[calledFunction].insert(funcOp).second) {
+          if (calledBy[calledFunction].insert(funcOp)) {
             numberCallOpsContainedInFuncOp[funcOp]++;
           }
           return WalkResult::advance();

``````````

</details>


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