[MLIR] Keep cached symbol tables across buffer deallocation insertions (#141956)

The `DeallocationState` class has been modified to keep a reference to an externally owned `SymbolTableCollection` class, to preserve the cached symbol tables across multiple insertions of deallocation instructions.

Author: mscuttari

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.

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.

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)
+    : symbolTable(symbolTables), liveness(op) {}
 
 void DeallocationState::updateOwnership(Value memref, Ownership ownership,
                                         Block *block) {

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