[llvm] [mlir][bufferization] BufferDeallocation: support unstructured control flow loops (PR #66657)
via llvm-commits
llvm-commits at lists.llvm.org
Mon Sep 18 08:32:23 PDT 2023
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg
<details>
<summary>Changes</summary>
This commit adds support for any kind of unstructured control flow loops.
Depends on #<!-- -->66626
Only review the top commit.
---
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:
- (modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h (+11-2)
- (modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td (+2-2)
- (modified) mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h (+17)
- (modified) mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h (+5-3)
- (modified) mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td (+8)
- (modified) mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp (+3-3)
- (modified) mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp (+10-5)
- (modified) mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp (+13-1)
- (modified) mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt (+1)
- (modified) mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp (+58-133)
- (modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir (+7-1)
- (modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-memoryeffect-interface.mlir (+9-1)
- (modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir (+157-56)
- (added) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-unstructured-cf-loops.mlir (+404)
- (modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/invalid-buffer-deallocation.mlir (-66)
- (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-tensor-e2e.mlir (+1-1)
- (modified) utils/bazel/llvm-project-overlay/mlir/BUILD.bazel (+1)
``````````diff
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]
``````````
</details>
https://github.com/llvm/llvm-project/pull/66657
More information about the llvm-commits
mailing list