[mlir][bufferization] OwnershipBasedBufferDeallocation: support unstructured control flow loops #66657
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
…n to forbid clones Adds a pass option to the `ownership-based-buffer-deallocation` pass to forbid insertion of clone operations. This is necessary to support IR that does not have the property that every buffer write dominates every buffer read to the same buffer. Instead of silently producing invalid IR, the pass would then emit an error. This is a restriction in the old `buffer-deallocation` pass, but the new function boundary ABI was not enforced in this old pass. Having this option allows easier migration from the old to the new deallocation pass because enabling this option allows the new deallocation pass to fix IR that does not adhere to the function boundary ABI (in some situations).
…l flow loops This commit adds support for any kind of unstructured control flow loops.
llvmbot
added
mlir:linalg
mlir
mlir:bufferization
|
@llvm/pr-subscribers-mlir-bufferization @llvm/pr-subscribers-mlir-linalg ChangesThis commit adds support for any kind of unstructured control flow loops. Depends on #66626 Patch is 77.17 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/66657.diff 19 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
index 7ac4592de7875fb..3aa61fae8c6caee 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
@@ -96,6 +96,14 @@ struct DeallocationOptions {
// pass the ownership of MemRef values instead of adhering to the function
// boundary ABI.
bool privateFuncDynamicOwnership = false;
+
+ // Allows the pass to insert `bufferization.clone` operations. This is useful
+ // for supporting IR that does not adhere to the function boundary ABI
+ // initially (excl. external functions) and to support operations with results
+ // with 'Unknown' ownership. However, it requires that all buffer writes
+ // dominate all buffer reads (i.e., only enable this option if your IR is
+ // guaranteed to have this property).
+ bool allowCloning = false;
};
/// This class collects all the state that we need to perform the buffer
@@ -142,8 +150,9 @@ class DeallocationState {
/// a new SSA value, returned as the first element of the pair, which has
/// 'Unique' ownership and can be used instead of the passed Value with the
/// the ownership indicator returned as the second element of the pair.
- std::pair<Value, Value>
- getMemrefWithUniqueOwnership(OpBuilder &builder, Value memref, Block *block);
+ FailureOr<std::pair<Value, Value>>
+ getMemrefWithUniqueOwnership(const DeallocationOptions &options,
+ OpBuilder &builder, Value memref, Block *block);
/// Given two basic blocks and the values passed via block arguments to the
/// destination block, compute the list of MemRefs that have to be retained in
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td
index 3e11432c65c5f08..3b9a9c3f4fef667 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td
@@ -56,7 +56,7 @@ def BufferDeallocationOpInterface :
method (which is especially important if operations are created that
cannot be easily canonicalized away anymore).
}],
- /*retType=*/"std::pair<Value, Value>",
+ /*retType=*/"FailureOr<std::pair<Value, Value>>",
/*methodName=*/"materializeUniqueOwnershipForMemref",
/*args=*/(ins "DeallocationState &":$state,
"const DeallocationOptions &":$options,
@@ -65,7 +65,7 @@ def BufferDeallocationOpInterface :
/*methodBody=*/[{}],
/*defaultImplementation=*/[{
return state.getMemrefWithUniqueOwnership(
- builder, memref, memref.getParentBlock());
+ options, builder, memref, memref.getParentBlock());
}]>,
];
}
diff --git a/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h
index 7acacb763cd2c18..7500257ed95eac8 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h
@@ -17,6 +17,7 @@
namespace mlir {
namespace bufferization {
+struct DeallocationOptions;
/// Options for the buffer deallocation pipeline.
struct BufferDeallocationPipelineOptions
@@ -28,6 +29,22 @@ struct BufferDeallocationPipelineOptions
"dynamically pass ownership of memrefs to callees. This can enable "
"earlier deallocations."),
llvm::cl::init(false)};
+ PassOptions::Option<bool> allowCloning{
+ *this, "allow-cloning",
+ llvm::cl::desc(
+ "Allows the pass to insert `bufferization.clone` operations. This is "
+ "useful for supporting IR that does not adhere to the function "
+ "boundary ABI initially (excl. external functions) and to support "
+ "operations with results with 'Unknown' ownership. However, it "
+ "requires that all buffer writes dominate all buffer reads (i.e., "
+ "only enable this option if your IR is guaranteed to have this "
+ "property)."),
+ llvm::cl::init(false)};
+
+ /// Convert this BufferDeallocationPipelineOptions struct to a
+ /// DeallocationOptions struct to be passed to the
+ /// OwnershipBasedBufferDeallocationPass.
+ DeallocationOptions asDeallocationOptions() const;
};
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
index 92520eb13da6875..37a3942f7bac6c5 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
@@ -1,6 +1,7 @@
#ifndef MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H
#define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H
+#include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h"
#include "mlir/Pass/Pass.h"
namespace mlir {
@@ -31,7 +32,7 @@ std::unique_ptr<Pass> createBufferDeallocationPass();
/// Creates an instance of the OwnershipBasedBufferDeallocation pass to free all
/// allocated buffers.
std::unique_ptr<Pass> createOwnershipBasedBufferDeallocationPass(
- bool privateFuncDynamicOwnership = false);
+ const DeallocationOptions &options = DeallocationOptions());
/// Creates a pass that optimizes `bufferization.dealloc` operations. For
/// example, it reduces the number of alias checks needed at runtime using
@@ -134,8 +135,9 @@ func::FuncOp buildDeallocationLibraryFunction(OpBuilder &builder, Location loc,
LogicalResult deallocateBuffers(Operation *op);
/// Run ownership basedbuffer deallocation.
-LogicalResult deallocateBuffersOwnershipBased(FunctionOpInterface op,
- bool privateFuncDynamicOwnership);
+LogicalResult deallocateBuffersOwnershipBased(
+ FunctionOpInterface op,
+ const DeallocationOptions &options = DeallocationOptions());
/// Creates a pass that moves allocations upwards to reduce the number of
/// required copies that are inserted during the BufferDeallocation pass.
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 62383e376f6f7a3..5b8af7a975c34b5 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -223,6 +223,14 @@ def OwnershipBasedBufferDeallocation : Pass<
"Allows to add additional arguments to private functions to "
"dynamically pass ownership of memrefs to callees. This can enable "
"earlier deallocations.">,
+ Option<"allowCloning", "allow-cloning", "bool", /*default=*/"false",
+ "Allows the pass to insert `bufferization.clone` operations. This "
+ "is useful for supporting IR that does not adhere to the function "
+ "boundary ABI initially (excl. external functions) and to support "
+ "operations with results with 'Unknown' ownership. However, it "
+ "requires that all buffer writes dominate all buffer reads (i.e., "
+ "only enable this option if your IR is guaranteed to have this "
+ "property).">,
];
let constructor = "mlir::bufferization::createOwnershipBasedBufferDeallocationPass()";
diff --git a/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp
index f2e7732e8ea4aa3..8ab4717739a7643 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp
@@ -53,7 +53,7 @@ struct SelectOpInterface
return op; // nothing to do
}
- std::pair<Value, Value>
+ FailureOr<std::pair<Value, Value>>
materializeUniqueOwnershipForMemref(Operation *op, DeallocationState &state,
const DeallocationOptions &options,
OpBuilder &builder, Value value) const {
@@ -64,14 +64,14 @@ struct SelectOpInterface
Block *block = value.getParentBlock();
if (!state.getOwnership(selectOp.getTrueValue(), block).isUnique() ||
!state.getOwnership(selectOp.getFalseValue(), block).isUnique())
- return state.getMemrefWithUniqueOwnership(builder, value,
+ return state.getMemrefWithUniqueOwnership(options, builder, value,
value.getParentBlock());
Value ownership = builder.create<arith::SelectOp>(
op->getLoc(), selectOp.getCondition(),
state.getOwnership(selectOp.getTrueValue(), block).getIndicator(),
state.getOwnership(selectOp.getFalseValue(), block).getIndicator());
- return {selectOp.getResult(), ownership};
+ return std::make_pair(selectOp.getResult(), ownership);
}
};
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
index 407d75e2426e9f9..9ac2e09dec385aa 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
@@ -132,16 +132,21 @@ void DeallocationState::getLiveMemrefsIn(Block *block,
memrefs.append(liveMemrefs);
}
-std::pair<Value, Value>
-DeallocationState::getMemrefWithUniqueOwnership(OpBuilder &builder,
- Value memref, Block *block) {
+FailureOr<std::pair<Value, Value>>
+DeallocationState::getMemrefWithUniqueOwnership(
+ const DeallocationOptions &options, OpBuilder &builder, Value memref,
+ Block *block) {
auto iter = ownershipMap.find({memref, block});
assert(iter != ownershipMap.end() &&
"Value must already have been registered in the ownership map");
Ownership ownership = iter->second;
if (ownership.isUnique())
- return {memref, ownership.getIndicator()};
+ return std::make_pair(memref, ownership.getIndicator());
+
+ if (!options.allowCloning)
+ return emitError(memref.getLoc(),
+ "MemRef value does not have valid ownership");
// Instead of inserting a clone operation we could also insert a dealloc
// operation earlier in the block and use the updated ownerships returned by
@@ -155,7 +160,7 @@ DeallocationState::getMemrefWithUniqueOwnership(OpBuilder &builder,
Value newMemref = cloneOp.getResult();
updateOwnership(newMemref, condition);
memrefsToDeallocatePerBlock[newMemref.getParentBlock()].push_back(newMemref);
- return {newMemref, condition};
+ return std::make_pair(newMemref, condition);
}
void DeallocationState::getMemrefsToRetain(
diff --git a/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp b/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp
index b2a60feb9a7f011..f08de33345ce605 100644
--- a/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp
+++ b/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp
@@ -8,23 +8,35 @@
#include "mlir/Dialect/Bufferization/Pipelines/Passes.h"
+#include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h"
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/Passes.h"
+using namespace mlir;
+using namespace bufferization;
+
//===----------------------------------------------------------------------===//
// Pipeline implementation.
//===----------------------------------------------------------------------===//
+DeallocationOptions
+BufferDeallocationPipelineOptions::asDeallocationOptions() const {
+ DeallocationOptions opts;
+ opts.privateFuncDynamicOwnership = privateFunctionDynamicOwnership.getValue();
+ opts.allowCloning = allowCloning.getValue();
+ return opts;
+}
+
void mlir::bufferization::buildBufferDeallocationPipeline(
OpPassManager &pm, const BufferDeallocationPipelineOptions &options) {
pm.addNestedPass<func::FuncOp>(
memref::createExpandReallocPass(/*emitDeallocs=*/false));
pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
pm.addNestedPass<func::FuncOp>(createOwnershipBasedBufferDeallocationPass(
- options.privateFunctionDynamicOwnership.getValue()));
+ options.asDeallocationOptions()));
pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
pm.addNestedPass<func::FuncOp>(createBufferDeallocationSimplificationPass());
pm.addPass(createLowerDeallocationsPass());
diff --git a/mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt b/mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt
index 6e8dab64ba6b935..d67b28b308fa10e 100644
--- a/mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt
+++ b/mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt
@@ -5,6 +5,7 @@ add_mlir_dialect_library(MLIRBufferizationPipelines
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization
LINK_LIBS PUBLIC
+ MLIRBufferizationDialect
MLIRBufferizationTransforms
MLIRMemRefTransforms
MLIRFuncDialect
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index 43ba11cf132cb92..c022d47199200ed 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
@@ -47,89 +47,6 @@ static Value buildBoolValue(OpBuilder &builder, Location loc, bool value) {
static bool isMemref(Value v) { return v.getType().isa<BaseMemRefType>(); }
-//===----------------------------------------------------------------------===//
-// Backedges analysis
-//===----------------------------------------------------------------------===//
-
-namespace {
-
-/// A straight-forward program analysis which detects loop backedges induced by
-/// explicit control flow.
-class Backedges {
-public:
- using BlockSetT = SmallPtrSet<Block *, 16>;
- using BackedgeSetT = llvm::DenseSet<std::pair<Block *, Block *>>;
-
-public:
- /// Constructs a new backedges analysis using the op provided.
- Backedges(Operation *op) { recurse(op); }
-
- /// Returns the number of backedges formed by explicit control flow.
- size_t size() const { return edgeSet.size(); }
-
- /// Returns the start iterator to loop over all backedges.
- BackedgeSetT::const_iterator begin() const { return edgeSet.begin(); }
-
- /// Returns the end iterator to loop over all backedges.
- BackedgeSetT::const_iterator end() const { return edgeSet.end(); }
-
-private:
- /// Enters the current block and inserts a backedge into the `edgeSet` if we
- /// have already visited the current block. The inserted edge links the given
- /// `predecessor` with the `current` block.
- bool enter(Block ¤t, Block *predecessor) {
- bool inserted = visited.insert(¤t).second;
- if (!inserted)
- edgeSet.insert(std::make_pair(predecessor, ¤t));
- return inserted;
- }
-
- /// Leaves the current block.
- void exit(Block ¤t) { visited.erase(¤t); }
-
- /// Recurses into the given operation while taking all attached regions into
- /// account.
- void recurse(Operation *op) {
- Block *current = op->getBlock();
- // If the current op implements the `BranchOpInterface`, there can be
- // cycles in the scope of all successor blocks.
- if (isa<BranchOpInterface>(op)) {
- for (Block *succ : current->getSuccessors())
- recurse(*succ, current);
- }
- // Recurse into all distinct regions and check for explicit control-flow
- // loops.
- for (Region ®ion : op->getRegions()) {
- if (!region.empty())
- recurse(region.front(), current);
- }
- }
-
- /// Recurses into explicit control-flow structures that are given by
- /// the successor relation defined on the block level.
- void recurse(Block &block, Block *predecessor) {
- // Try to enter the current block. If this is not possible, we are
- // currently processing this block and can safely return here.
- if (!enter(block, predecessor))
- return;
-
- // Recurse into all operations and successor blocks.
- for (Operation &op : block.getOperations())
- recurse(&op);
-
- // Leave the current block.
- exit(block);
- }
-
- /// Stores all blocks that are currently visited and on the processing stack.
- BlockSetT visited;
-
- /// Stores all backedges in the format (source, target).
- BackedgeSetT edgeSet;
-};
-
-} // namespace
-
//===----------------------------------------------------------------------===//
// BufferDeallocation
//===----------------------------------------------------------------------===//
@@ -139,10 +56,8 @@ namespace {
/// program have a corresponding de-allocation.
class BufferDeallocation {
public:
- BufferDeallocation(Operation *op, bool privateFuncDynamicOwnership)
- : state(op) {
- options.privateFuncDynamicOwnership = privateFuncDynamicOwnership;
- }
+ BufferDeallocation(Operation *op, DeallocationOptions options)
+ : state(op), options(options) {}
/// Performs the actual placement/creation of all dealloc operations.
LogicalResult deallocate(FunctionOpInterface op);
@@ -376,8 +291,9 @@ class BufferDeallocation {
/// Given an SSA value of MemRef type, returns the same of a new SSA value
/// which has 'Unique' ownership where the ownership indicator is guaranteed
/// to be always 'true'.
- Value materializeMemrefWithGuaranteedOwnership(OpBuilder &builder,
- Value memref, Block *block);
+ FailureOr<Value> materializeMemrefWithGuaranteedOwnership(OpBuilder &builder,
+ Value memref,
+ Block *block);
/// Returns whether the given operation implements FunctionOpInterface, has
/// private visibility, and the private-function-dynamic-ownership pass option
@@ -391,15 +307,9 @@ class BufferDeallocation {
/// is requested does not match the block in which 'memref' is defined, the
/// default implementation in
/// `DeallocationState::getMemrefWithUniqueOwnership` is queried instead.
- std::pair<Value, Value>
+ FailureOr<std::pair<Value, Value>>
materializeUniqueOwnership(OpBuilder &builder, Value memref, Block *block);
- /// Checks all the preconditions for operations implementing the
- /// FunctionOpInterface that have to hold for the deallocation to be
- /// applicable:
- /// (1) Checks that there are not explicit control flow loops.
- static LogicalResult verifyFunctionPreconditions(FunctionOpInterface op);
-
/// Checks all the preconditions for operations inside the region of
/// operations implementing the FunctionOpInterface that have to hold for the
/// deallocation to be applicable:
@@ -430,7 +340,7 @@ class BufferDeallocation {
DeallocationState state;
/// Collects all pass options in a single place.
- DeallocationOptions options;
+ const DeallocationOptions options;
};
} // namespace
@@ -439,13 +349,13 @@ class BufferDeallocation {
// BufferDeallocation Implementation
//===----------------------------------------------------------------------===//
-std::pair<Value, Value>
+FailureOr<std::pair<Value, Value>>
BufferDeallocation::materializeUniqueOwnership(OpBuilder &builder, Value memref,
Block *block) {
// The interface can only materialize ownership indicators in the same block
// as the defining op.
if (memref.getParentBlock() != block)
- return state.getMemrefWithUniqueOwnership(builder, memref, block);
+ return state.getMemrefWithUniqueOwnership(options, builder, memref, block);
Operation *owner = memref.getDefiningOp();
if (!owner)
@@ -458,7 +368,7 @@ BufferDeallocation::materializeUniqueOwnership(OpBuilder &builder, Value memref,
state, options, builder, memr...
[truncated]
|
maerhart
changed the title
matthias-springer
approved these changes
Sep 26, 2023
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.
This commit adds support for any kind of unstructured control flow loops.
Depends on #66626
Only review the top commit.