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

[Colin He]

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

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

>From a3edc3fb724120ea105ace218251d376c193d869 Mon Sep 17 00:00:00 2001
From: Colin He <hcolin88 at gmail.com>
Date: Wed, 4 Mar 2026 18:07:32 -0800
Subject: [PATCH] [mlir] Deterministic containers in OneShotModuleBufferize

---
 .../Bufferization/Transforms/OneShotModuleBufferize.cpp    | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

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