From f3c8ae7a113643e5ecc57616654f0a40ca0c3c60 Mon Sep 17 00:00:00 2001 From: Michele Scuttari 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 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) {