[mlir][bufferization] Add option to LowerDeallocations to choose the kind of dealloc op to build #67565
Conversation
This file contains 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
…ction boundary ABI Inserting clones requires a lot of assumptions to hold on the input IR, e.g., all writes to a buffer need to dominate all reads. This is not guaranteed by one-shot bufferization and isn't easy to verify, thus it could quickly lead to incorrect results that are hard to debug. This commit changes the mechanism of how an ownership indicator is materialized when there is not already a unique ownership present. Additionally, we don't create copies of returned memrefs anymore when we don't have ownership. Instead, we insert assert operations to make sure we have ownership at runtime, or otherwise report to the user that correctness could not be guaranteed.
…loc operations, add option to specify the kind of alloc operations to consider
…kind of dealloc op to build
llvmbot
added
mlir:core
|
@llvm/pr-subscribers-mlir-core @llvm/pr-subscribers-mlir-linalg Changes
Depends on #67556 Patch is 79.66 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/67565.diff 31 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
index 752a4a2c6f42a2f..d2e718ac93045ce 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
@@ -10,6 +10,8 @@
#define MLIR_DIALECT_BUFFERIZATION_IR_BUFFERDEALLOCATIONOPINTERFACE_H_
#include "mlir/Analysis/Liveness.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Support/LLVM.h"
@@ -92,10 +94,65 @@ class Ownership {
/// Options for BufferDeallocationOpInterface-based buffer deallocation.
struct DeallocationOptions {
+ using DetectionFn = std::function<bool(Operation *)>;
+ using ReplaceDeallocFn = std::function<FailureOr<ValueRange>(Operation *)>;
+
// A pass option indicating whether private functions should be modified to
// pass the ownership of MemRef values instead of adhering to the function
// boundary ABI.
- bool privateFuncDynamicOwnership = false;
+ bool privateFuncDynamicOwnership = true;
+
+ /// Inserts `cf.assert` operations to verify the function boundary ABI at
+ /// runtime. Currently, it is only checked that the ownership of returned
+ /// MemRefs is 'true'. This makes sure that ownership is yielded and the
+ /// returned MemRef does not originate from the same allocation as a function
+ /// argument. TODO: check that returned MemRefs don't alias each other.
+ /// If it can be determined statically that the ABI is not adhered
+ /// to, an error will already be emitted at compile time. This cannot be
+ /// changed with this option.
+ bool verifyFunctionBoundaryABI = true;
+
+ /// Given an allocation side-effect on the passed operation, determine whether
+ /// this allocation operation is of relevance (i.e., should assign ownership
+ /// to the allocated value). If it is determined to not be relevant,
+ /// ownership will be set to 'false', i.e., it will be leaked. This is useful
+ /// to support deallocation of multiple different kinds of allocation ops.
+ DetectionFn isRelevantAllocOp = [](Operation *op) {
+ return isa<memref::MemRefDialect, bufferization::BufferizationDialect>(
+ op->getDialect());
+ };
+
+ /// Given a free side-effect on the passed operation, determine whether this
+ /// deallocation operation is of relevance (i.e., should be removed if the
+ /// `removeExistingDeallocations` option is enabled or otherwise an error
+ /// should be emitted because existing deallocation operations are not
+ /// supported without that flag). If it is determined to not be relevant,
+ /// the operation will be ignored. This is useful to support deallocation of
+ /// multiple different kinds of allocation ops where deallocations for some of
+ /// them are already present in the IR.
+ DetectionFn isRelevantDeallocOp = [](Operation *op) {
+ return isa<memref::MemRefDialect, bufferization::BufferizationDialect>(
+ op->getDialect());
+ };
+
+ /// When enabled, remove deallocation operations determined to be relevant
+ /// according to `isRelevantDeallocOp`. If the operation has result values,
+ /// `getDeallocReplacement` will be called to determine the SSA values that
+ /// should be used as replacements.
+ bool removeExistingDeallocations = false;
+
+ /// Provides SSA values for deallocation operations when
+ /// `removeExistingDeallocations` is enabled. May return a failure when the
+ /// given deallocation operation is not supported (e.g., because no
+ /// replacement for a result value can be determined). A failure will directly
+ /// lead to a failure emitted by the deallocation pass.
+ ReplaceDeallocFn getDeallocReplacement =
+ [](Operation *op) -> FailureOr<ValueRange> {
+ if (isa<memref::DeallocOp>(op))
+ return ValueRange{};
+ // ReallocOp has to be expanded before running the dealloc pass.
+ return failure();
+ };
};
/// This class collects all the state that we need to perform the buffer
@@ -138,12 +195,12 @@ class DeallocationState {
void getLiveMemrefsIn(Block *block, SmallVectorImpl<Value> &memrefs);
/// Given an SSA value of MemRef type, this function queries the ownership and
- /// if it is not already in the 'Unique' state, potentially inserts IR to get
- /// 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);
+ /// if it is not already in the 'Unique' state, potentially inserts IR to
+ /// determine the ownership (which might involve expensive aliasing checks at
+ /// runtime).
+ 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..8fe143eea5bf8ec 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td
@@ -55,8 +55,38 @@ def BufferDeallocationOpInterface :
ownership indicator when needed, it should be implemented using this
method (which is especially important if operations are created that
cannot be easily canonicalized away anymore).
+ Ownership indicators have to be materialized when
+ * needed for the condition operands of a `bufferization.dealloc` op
+ * passed along MemRefs to successor blocks via additional forwarded
+ operands of terminator ops
+ * passing them as additional operands to nested regions (e.g.,
+ init_args of `scf.for`)
+ * passing them as additional operands to a call operation when
+ `private-function-dynamic-ownership` is enabled
+ * a copy is made conditionally on the current ownership, etc.
+
+ In the following example, the deallocation pass would add an
+ additional block argument to `^bb1` for passing the ownership of `%0`
+ along and thus the ownership indicator has to be materialized before
+ the `cf.br` operation and added as a forwarded operand.
+ ```mlir
+ %0 = arith.select %cond, %m1, %m2 : memref<f32>
+ cf.br ^bb1(%0 : memref<f32>)
+ ^bb1(%arg0: memref<f32>)
+ ...
+ ```
+ The `arith.select` operation could implement this interface method to
+ materialize another `arith.select` operation to select the
+ corresponding ownership indicator.
+ ```mlir
+ %0 = arith.select %cond, %m1, %m2 : memref<f32>
+ %0_ownership = arith.select %cond, %m1_ownership, %m2_ownership : i1
+ cf.br ^bb1(%0, %0_ownership : memref<f32>, i1)
+ ^bb1(%arg0: memref<f32>, %arg1: i1)
+ ...
+ ```
}],
- /*retType=*/"std::pair<Value, Value>",
+ /*retType=*/"Value",
/*methodName=*/"materializeUniqueOwnershipForMemref",
/*args=*/(ins "DeallocationState &":$state,
"const DeallocationOptions &":$options,
@@ -65,7 +95,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..38883ff4588d2ee 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
@@ -27,7 +28,29 @@ struct BufferDeallocationPipelineOptions
"Allows to add additional arguments to private functions to "
"dynamically pass ownership of memrefs to callees. This can enable "
"earlier deallocations."),
+ llvm::cl::init(true)};
+ PassOptions::Option<bool> verifyFunctionBoundaryABI{
+ *this, "verify-function-boundary-abi",
+ llvm::cl::desc(
+ "Inserts `cf.assert` operations to verify the function boundary ABI "
+ "at runtime. Currently, it is only checked that the ownership of "
+ "returned MemRefs is 'true'. This makes sure that ownership is "
+ "yielded and the returned MemRef does not originate from the same "
+ "allocation as a function argument. If it can be determined "
+ "statically that the ABI is not adhered to, an error will already be "
+ "emitted at compile time. This cannot be changed with this option."),
+ llvm::cl::init(true)};
+ PassOptions::Option<bool> removeExistingDeallocations{
+ *this, "remove-existing-deallocations",
+ llvm::cl::desc("Removes all pre-existing memref.dealloc operations and "
+ "insert all deallocations according to the buffer "
+ "deallocation rules."),
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 a6f668b26aa10e4..59b3ab5c84e80de 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 {
@@ -17,6 +18,16 @@ class FuncOp;
namespace bufferization {
struct OneShotBufferizationOptions;
+/// Options for the LowerDeallocation pass and rewrite patterns.
+struct LowerDeallocationOptions {
+ /// Given a MemRef value, build the operation(s) necessary to properly
+ /// deallocate the value.
+ std::function<void(OpBuilder &, Location, Value)> buildDeallocOp =
+ [](OpBuilder &builder, Location loc, Value memref) {
+ builder.create<memref::DeallocOp>(loc, memref);
+ };
+};
+
//===----------------------------------------------------------------------===//
// Passes
//===----------------------------------------------------------------------===//
@@ -31,7 +42,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
@@ -40,12 +51,14 @@ std::unique_ptr<Pass> createBufferDeallocationSimplificationPass();
/// Creates an instance of the LowerDeallocations pass to lower
/// `bufferization.dealloc` operations to the `memref` dialect.
-std::unique_ptr<Pass> createLowerDeallocationsPass();
+std::unique_ptr<Pass> createLowerDeallocationsPass(
+ const LowerDeallocationOptions &options = LowerDeallocationOptions());
/// Adds the conversion pattern of the `bufferization.dealloc` operation to the
/// given pattern set for use in other transformation passes.
void populateBufferizationDeallocLoweringPattern(
- RewritePatternSet &patterns, func::FuncOp deallocLibraryFunc);
+ RewritePatternSet &patterns, func::FuncOp deallocLibraryFunc,
+ const LowerDeallocationOptions &options = LowerDeallocationOptions());
/// Construct the library function needed for the fully generic
/// `bufferization.dealloc` lowering implemented in the LowerDeallocations pass.
@@ -134,8 +147,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 e01f36b8daa18d6..e182ae6548b95bf 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -219,10 +219,25 @@ def OwnershipBasedBufferDeallocation : Pass<
}];
let options = [
Option<"privateFuncDynamicOwnership", "private-function-dynamic-ownership",
- "bool", /*default=*/"false",
+ "bool", /*default=*/"true",
"Allows to add additional arguments to private functions to "
"dynamically pass ownership of memrefs to callees. This can enable "
"earlier deallocations.">,
+ Option<"verifyFunctionBoundaryABI", "verify-function-boundary-abi",
+ "bool", /*default=*/"true",
+ "Inserts `cf.assert` operations to verify the function boundary ABI "
+ "at runtime. Currently, it is only checked that the ownership of "
+ "returned MemRefs is 'true'. This makes sure that ownership is "
+ "yielded and the returned MemRef does not originate from the same "
+ "allocation as a function argument. "
+ "If it can be determined statically that the ABI is not adhered "
+ "to, an error will already be emitted at compile time. This cannot "
+ "be changed with this option.">,
+ Option<"removeExistingDeallocations", "remove-existing-deallocations",
+ "bool", /*default=*/"false",
+ "Remove already existing MemRef deallocation operations and let the "
+ "deallocation pass insert the deallocation operations according to "
+ "its rules.">,
];
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..29e6ba9bfcf5891 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp
@@ -53,10 +53,11 @@ struct SelectOpInterface
return op; // nothing to do
}
- std::pair<Value, Value>
- materializeUniqueOwnershipForMemref(Operation *op, DeallocationState &state,
- const DeallocationOptions &options,
- OpBuilder &builder, Value value) const {
+ Value materializeUniqueOwnershipForMemref(Operation *op,
+ DeallocationState &state,
+ const DeallocationOptions &options,
+ OpBuilder &builder,
+ Value value) const {
auto selectOp = cast<arith::SelectOp>(op);
assert(value == selectOp.getResult() &&
"Value not defined by this operation");
@@ -64,14 +65,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 ownership;
}
};
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
index 8d21446f1eb777e..50def5f45538bb5 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
@@ -132,30 +132,79 @@ void DeallocationState::getLiveMemrefsIn(Block *block,
memrefs.append(liveMemrefs);
}
-std::pair<Value, Value>
-DeallocationState::getMemrefWithUniqueOwnership(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()};
-
- // Instead of inserting a clone operation we could also insert a dealloc
- // operation earlier in the block and use the updated ownerships returned by
- // the op for the retained values. Alternatively, we could insert code to
- // check aliasing at runtime and use this information to combine two unique
- // ownerships more intelligently to not end up with an 'Unknown' ownership in
- // the first place.
- auto cloneOp =
- builder.create<bufferization::CloneOp>(memref.getLoc(), memref);
- Value condition = buildBoolValue(builder, memref.getLoc(), true);
- Value newMemref = cloneOp.getResult();
- updateOwnership(newMemref, condition);
- memrefsToDeallocatePerBlock[newMemref.getParentBlock()].push_back(newMemref);
- return {newMemref, condition};
+Value DeallocationState::getMemrefWithUniqueOwnership(
+ const DeallocationOptions &options, OpBuilder &builder, Value memref,
+ Block *block) {
+ // NOTE: * if none of the operands have the same allocated pointer, a new
+ // memref was allocated and thus the operation should have the allocate
+ // side-effect defined on that result value and thus the correct unique
+ // ownership is pre-populated by the ownership pass (unless an interface
+ // implementation is incorrect).
+ // * if exactly one operand has the same allocated pointer, this retunes
+ // the ownership of exactly that operand
+ // * if multiple operands match the allocated pointer of the result, the
+ // ownership indicators of all of them always have to evaluate to the
+ // same value because no dealloc operations may be present and because
+ // of the rules they are passed to nested regions and successor blocks.
+ // This could be verified at runtime by inserting `cf.assert`
+ // operations, but would require O(|operands|^2) additional operations
+ // to check and is thus not implemented yet (would need to insert a
+ // library function to avoid code-size explosion which would make the
+ // deallo...
[truncated]
|
matthias-springer
approved these changes
Sep 28, 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.
bufferization.deallocop lowering.Depends on #67556
Only review top commit