[MLIR] Keep cached symbol tables across buffer deallocation insertions #141956
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-mlir Author: Michele Scuttari (mscuttari) ChangesThe Full diff: https://github.com/llvm/llvm-project/pull/141956.diff 4 Files Affected:
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);
|
|
@llvm/pr-subscribers-mlir-bufferization Author: Michele Scuttari (mscuttari) ChangesThe Full diff: https://github.com/llvm/llvm-project/pull/141956.diff 4 Files Affected:
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);
|
matthias-springer
approved these changes
Jun 6, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
The
DeallocationStateclass has been modified to keep a reference to an externally ownedSymbolTableCollectionclass, to preserve the cached symbol tables across multiple insertions of deallocation instructions.