[Mlir-commits] [mlir] [MLIR] Keep cached symbol tables across buffer deallocation insertions (PR #141956)

[Michele Scuttari]

https://github.com/mscuttari updated https://github.com/llvm/llvm-project/pull/141956

>From f3c8ae7a113643e5ecc57616654f0a40ca0c3c60 Mon Sep 17 00:00:00 2001
From: Michele Scuttari <michele.scuttari at outlook.com>
Date: Thu, 29 May 2025 16:25:49 +0200
Subject: [PATCH 1/2] Keep cached symbol tables across buffer deallocation
 insertions

---
 .../IR/BufferDeallocationOpInterface.h          |  4 ++--
 .../Dialect/Bufferization/Transforms/Passes.h   |  6 ++++--
 .../IR/BufferDeallocationOpInterface.cpp        |  4 +++-
 .../OwnershipBasedBufferDeallocation.cpp        | 17 ++++++++++-------
 4 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
index 752a4a2c6f42a..5d33817c7d9fc 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
@@ -104,7 +104,7 @@ struct DeallocationOptions {
 /// BufferDeallocation pass.
 class DeallocationState {
 public:
-  DeallocationState(Operation *op);
+  DeallocationState(Operation *op, SymbolTableCollection &symbolTables);
 
   // The state should always be passed by reference.
   DeallocationState(const DeallocationState &) = delete;
@@ -189,7 +189,7 @@ class DeallocationState {
 private:
   // Symbol cache to lookup functions from call operations to check attributes
   // on the function operation.
-  SymbolTableCollection symbolTable;
+  SymbolTableCollection &symbolTable;
 
   // Mapping from each SSA value with MemRef type to the associated ownership in
   // each block.
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
index a8cc37a103f04..596c470ef6d23 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
@@ -120,8 +120,10 @@ func::FuncOp buildDeallocationLibraryFunction(OpBuilder &builder, Location loc,
                                               SymbolTable &symbolTable);
 
 /// Run the ownership-based buffer deallocation.
-LogicalResult deallocateBuffersOwnershipBased(FunctionOpInterface op,
-                                              DeallocationOptions options);
+LogicalResult
+deallocateBuffersOwnershipBased(FunctionOpInterface op,
+                                DeallocationOptions options,
+                                SymbolTableCollection &symbolTables);
 
 // Options struct for BufferResultsToOutParams pass.
 // Note: defined only here, not in tablegen.
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
index eed7a56fff8af..dee234c6314da 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
@@ -94,7 +94,9 @@ void Ownership::combine(Ownership other) { *this = getCombined(other); }
 // DeallocationState
 //===----------------------------------------------------------------------===//
 
-DeallocationState::DeallocationState(Operation *op) : liveness(op) {}
+DeallocationState::DeallocationState(Operation *op,
+                                     SymbolTableCollection &symbolTables)
+    : liveness(op), symbolTable(symbolTables) {}
 
 void DeallocationState::updateOwnership(Value memref, Ownership ownership,
                                         Block *block) {
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index c5b2f3020712e..1eeafc4df8cf1 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
@@ -166,8 +166,9 @@ namespace {
 /// program have a corresponding de-allocation.
 class BufferDeallocation {
 public:
-  BufferDeallocation(Operation *op, DeallocationOptions options)
-      : state(op), options(options) {}
+  BufferDeallocation(Operation *op, DeallocationOptions options,
+                     SymbolTableCollection &symbolTables)
+      : state(op, symbolTables), options(options) {}
 
   /// Performs the actual placement/creation of all dealloc operations.
   LogicalResult deallocate(FunctionOpInterface op);
@@ -1027,11 +1028,13 @@ struct OwnershipBasedBufferDeallocationPass
     DeallocationOptions options;
     options.privateFuncDynamicOwnership = privateFuncDynamicOwnership;
 
+    mlir::SymbolTableCollection symbolTables;
+
     auto status = getOperation()->walk([&](func::FuncOp func) {
       if (func.isExternal())
         return WalkResult::skip();
 
-      if (failed(deallocateBuffersOwnershipBased(func, options)))
+      if (failed(deallocateBuffersOwnershipBased(func, options, symbolTables)))
         return WalkResult::interrupt();
 
       return WalkResult::advance();
@@ -1047,11 +1050,11 @@ struct OwnershipBasedBufferDeallocationPass
 // Implement bufferization API
 //===----------------------------------------------------------------------===//
 
-LogicalResult
-bufferization::deallocateBuffersOwnershipBased(FunctionOpInterface op,
-                                               DeallocationOptions options) {
+LogicalResult bufferization::deallocateBuffersOwnershipBased(
+    FunctionOpInterface op, DeallocationOptions options,
+    SymbolTableCollection &symbolTables) {
   // Gather all required allocation nodes and prepare the deallocation phase.
-  BufferDeallocation deallocation(op, options);
+  BufferDeallocation deallocation(op, options, symbolTables);
 
   // Place all required temporary clone and dealloc nodes.
   return deallocation.deallocate(op);

>From 3bde109f45978f89e814cf0e31e990780be3e4a1 Mon Sep 17 00:00:00 2001
From: Michele Scuttari <michele.scuttari at outlook.com>
Date: Thu, 5 Jun 2025 12:41:43 +0200
Subject: [PATCH 2/2] Inizialize SymbolTableCollection before Liveness

---
 .../Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
index dee234c6314da..ca914df8b7890 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
@@ -96,7 +96,7 @@ void Ownership::combine(Ownership other) { *this = getCombined(other); }
 
 DeallocationState::DeallocationState(Operation *op,
                                      SymbolTableCollection &symbolTables)
-    : liveness(op), symbolTable(symbolTables) {}
+    : symbolTable(symbolTables), liveness(op) {}
 
 void DeallocationState::updateOwnership(Value memref, Ownership ownership,
                                         Block *block) {