[Mlir-commits] [mlir] [mlir][bufferization] Add option to LowerDeallocations to choose the kind of dealloc op to build (PR #67565)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Sep 27 07:54:36 PDT 2023
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg
<details>
<summary>Changes</summary>
* Add option to LowerDeallocations to choose the kind of dealloc op to build
* Use alloca instead of alloc in the generic `bufferization.dealloc` op lowering.
Depends on #<!-- -->67556
Only review top commit
---
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:
- (modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h (+64-7)
- (modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td (+32-2)
- (modified) mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h (+23)
- (modified) mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h (+19-5)
- (modified) mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td (+16-1)
- (modified) mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp (+7-6)
- (modified) mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp (+73-24)
- (modified) mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp (+14-1)
- (modified) mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt (+1)
- (modified) mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt (+1)
- (modified) mlir/lib/Dialect/Bufferization/Transforms/LowerDeallocations.cpp (+30-24)
- (modified) mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp (+63-72)
- (modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir (+1-1)
- (modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir (+1-1)
- (modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir (+9-17)
- (modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-memoryeffect-interface.mlir (+2-19)
- (added) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-mixed-allocations.mlir (+33)
- (modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir (+14-57)
- (added) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-runtime-verification.mlir (+13)
- (modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-subviews.mlir (+1-1)
- (added) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-unknown-ops.mlir (+29)
- (modified) mlir/test/Dialect/Bufferization/Transforms/lower-deallocations.mlir (+5-10)
- (modified) mlir/test/Integration/Dialect/Linalg/CPU/test-collapse-tensor.mlir (+1-1)
- (modified) mlir/test/Integration/Dialect/Linalg/CPU/test-expand-tensor.mlir (+1-1)
- (modified) mlir/test/Integration/Dialect/Linalg/CPU/test-one-shot-bufferize.mlir (+1-1)
- (modified) mlir/test/Integration/Dialect/Linalg/CPU/test-tensor-e2e.mlir (+1-1)
- (modified) mlir/test/lib/Dialect/Bufferization/CMakeLists.txt (+4)
- (added) mlir/test/lib/Dialect/Bufferization/TestOwnershipBasedBufferDeallocation.cpp (+139)
- (modified) mlir/tools/mlir-opt/mlir-opt.cpp (+2)
- (modified) utils/bazel/llvm-project-overlay/mlir/BUILD.bazel (+1)
- (modified) utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel (+6-1)
``````````diff
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]
``````````
</details>
https://github.com/llvm/llvm-project/pull/67565
More information about the Mlir-commits
mailing list