09dfc5713d
The greedy rewriter is used in many different flows and it has a lot of
convenience (work list management, debugging actions, tracing, etc). But
it combines two kinds of greedy behavior 1) how ops are matched, 2)
folding wherever it can.
These are independent forms of greedy and leads to inefficiency. E.g.,
cases where one need to create different phases in lowering and is
required to applying patterns in specific order split across different
passes. Using the driver one ends up needlessly retrying folding/having
multiple rounds of folding attempts, where one final run would have
sufficed.
Of course folks can locally avoid this behavior by just building their
own, but this is also a common requested feature that folks keep on
working around locally in suboptimal ways.
For downstream users, there should be no behavioral change. Updating
from the deprecated should just be a find and replace (e.g., `find ./
-type f -exec sed -i
's|applyPatternsAndFoldGreedily|applyPatternsGreedily|g' {} \;` variety)
as the API arguments hasn't changed between the two.
485 lines
20 KiB
C++
485 lines
20 KiB
C++
//===- BufferDeallocationSimplification.cpp -------------------------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file implements logic for optimizing `bufferization.dealloc` operations
|
|
// that requires more analysis than what can be supported by regular
|
|
// canonicalization patterns.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
|
|
#include "mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h"
|
|
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
|
|
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
|
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
|
#include "mlir/IR/Matchers.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
|
|
namespace mlir {
|
|
namespace bufferization {
|
|
#define GEN_PASS_DEF_BUFFERDEALLOCATIONSIMPLIFICATION
|
|
#include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
|
|
} // namespace bufferization
|
|
} // namespace mlir
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::bufferization;
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Helpers
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
/// Given a memref value, return the "base" value by skipping over all
|
|
/// ViewLikeOpInterface ops (if any) in the reverse use-def chain.
|
|
static Value getViewBase(Value value) {
|
|
while (auto viewLikeOp = value.getDefiningOp<ViewLikeOpInterface>())
|
|
value = viewLikeOp.getViewSource();
|
|
return value;
|
|
}
|
|
|
|
static LogicalResult updateDeallocIfChanged(DeallocOp deallocOp,
|
|
ValueRange memrefs,
|
|
ValueRange conditions,
|
|
PatternRewriter &rewriter) {
|
|
if (deallocOp.getMemrefs() == memrefs &&
|
|
deallocOp.getConditions() == conditions)
|
|
return failure();
|
|
|
|
rewriter.modifyOpInPlace(deallocOp, [&]() {
|
|
deallocOp.getMemrefsMutable().assign(memrefs);
|
|
deallocOp.getConditionsMutable().assign(conditions);
|
|
});
|
|
return success();
|
|
}
|
|
|
|
/// Return "true" if the given values are guaranteed to be different (and
|
|
/// non-aliasing) allocations based on the fact that one value is the result
|
|
/// of an allocation and the other value is a block argument of a parent block.
|
|
/// Note: This is a best-effort analysis that will eventually be replaced by a
|
|
/// proper "is same allocation" analysis. This function may return "false" even
|
|
/// though the two values are distinct allocations.
|
|
static bool distinctAllocAndBlockArgument(Value v1, Value v2) {
|
|
Value v1Base = getViewBase(v1);
|
|
Value v2Base = getViewBase(v2);
|
|
auto areDistinct = [](Value v1, Value v2) {
|
|
if (Operation *op = v1.getDefiningOp())
|
|
if (hasEffect<MemoryEffects::Allocate>(op, v1))
|
|
if (auto bbArg = dyn_cast<BlockArgument>(v2))
|
|
if (bbArg.getOwner()->findAncestorOpInBlock(*op))
|
|
return true;
|
|
return false;
|
|
};
|
|
return areDistinct(v1Base, v2Base) || areDistinct(v2Base, v1Base);
|
|
}
|
|
|
|
/// Checks if `memref` may potentially alias a MemRef in `otherList`. It is
|
|
/// often a requirement of optimization patterns that there cannot be any
|
|
/// aliasing memref in order to perform the desired simplification.
|
|
static bool potentiallyAliasesMemref(BufferOriginAnalysis &analysis,
|
|
ValueRange otherList, Value memref) {
|
|
for (auto other : otherList) {
|
|
if (distinctAllocAndBlockArgument(other, memref))
|
|
continue;
|
|
std::optional<bool> analysisResult =
|
|
analysis.isSameAllocation(other, memref);
|
|
if (!analysisResult.has_value() || analysisResult == true)
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Patterns
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
namespace {
|
|
|
|
/// Remove values from the `memref` operand list that are also present in the
|
|
/// `retained` list (or a guaranteed alias of it) because they will never
|
|
/// actually be deallocated. However, we also need to be certain about which
|
|
/// other memrefs in the `retained` list can alias, i.e., there must not by any
|
|
/// may-aliasing memref. This is necessary because the `dealloc` operation is
|
|
/// defined to return one `i1` value per memref in the `retained` list which
|
|
/// represents the disjunction of the condition values corresponding to all
|
|
/// aliasing values in the `memref` list. In particular, this means that if
|
|
/// there is some value R in the `retained` list which aliases with a value M in
|
|
/// the `memref` list (but can only be staticaly determined to may-alias) and M
|
|
/// is also present in the `retained` list, then it would be illegal to remove M
|
|
/// because the result corresponding to R would be computed incorrectly
|
|
/// afterwards. Because we require an alias analysis, this pattern cannot be
|
|
/// applied as a regular canonicalization pattern.
|
|
///
|
|
/// Example:
|
|
/// ```mlir
|
|
/// %0:3 = bufferization.dealloc (%m0 : ...) if (%cond0)
|
|
/// retain (%m0, %r0, %r1 : ...)
|
|
/// ```
|
|
/// is canonicalized to
|
|
/// ```mlir
|
|
/// // bufferization.dealloc without memrefs and conditions returns %false for
|
|
/// // every retained value
|
|
/// %0:3 = bufferization.dealloc retain (%m0, %r0, %r1 : ...)
|
|
/// %1 = arith.ori %0#0, %cond0 : i1
|
|
/// // replace %0#0 with %1
|
|
/// ```
|
|
/// given that `%r0` and `%r1` may not alias with `%m0`.
|
|
struct RemoveDeallocMemrefsContainedInRetained
|
|
: public OpRewritePattern<DeallocOp> {
|
|
RemoveDeallocMemrefsContainedInRetained(MLIRContext *context,
|
|
BufferOriginAnalysis &analysis)
|
|
: OpRewritePattern<DeallocOp>(context), analysis(analysis) {}
|
|
|
|
/// The passed 'memref' must not have a may-alias relation to any retained
|
|
/// memref, and at least one must-alias relation. If there is no must-aliasing
|
|
/// memref in the retain list, we cannot simply remove the memref as there
|
|
/// could be situations in which it actually has to be deallocated. If it's
|
|
/// no-alias, then just proceed, if it's must-alias we need to update the
|
|
/// updated condition returned by the dealloc operation for that alias.
|
|
LogicalResult handleOneMemref(DeallocOp deallocOp, Value memref, Value cond,
|
|
PatternRewriter &rewriter) const {
|
|
rewriter.setInsertionPointAfter(deallocOp);
|
|
|
|
// Check that there is no may-aliasing memref and that at least one memref
|
|
// in the retain list aliases (because otherwise it might have to be
|
|
// deallocated in some situations and can thus not be dropped).
|
|
bool atLeastOneMustAlias = false;
|
|
for (Value retained : deallocOp.getRetained()) {
|
|
std::optional<bool> analysisResult =
|
|
analysis.isSameAllocation(retained, memref);
|
|
if (!analysisResult.has_value())
|
|
return failure();
|
|
if (analysisResult == true)
|
|
atLeastOneMustAlias = true;
|
|
}
|
|
if (!atLeastOneMustAlias)
|
|
return failure();
|
|
|
|
// Insert arith.ori operations to update the corresponding dealloc result
|
|
// values to incorporate the condition of the must-aliasing memref such that
|
|
// we can remove that operand later on.
|
|
for (auto [i, retained] : llvm::enumerate(deallocOp.getRetained())) {
|
|
Value updatedCondition = deallocOp.getUpdatedConditions()[i];
|
|
std::optional<bool> analysisResult =
|
|
analysis.isSameAllocation(retained, memref);
|
|
if (analysisResult == true) {
|
|
auto disjunction = rewriter.create<arith::OrIOp>(
|
|
deallocOp.getLoc(), updatedCondition, cond);
|
|
rewriter.replaceAllUsesExcept(updatedCondition, disjunction.getResult(),
|
|
disjunction);
|
|
}
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult matchAndRewrite(DeallocOp deallocOp,
|
|
PatternRewriter &rewriter) const override {
|
|
// There must not be any duplicates in the retain list anymore because we
|
|
// would miss updating one of the result values otherwise.
|
|
DenseSet<Value> retained(deallocOp.getRetained().begin(),
|
|
deallocOp.getRetained().end());
|
|
if (retained.size() != deallocOp.getRetained().size())
|
|
return failure();
|
|
|
|
SmallVector<Value> newMemrefs, newConditions;
|
|
for (auto [memref, cond] :
|
|
llvm::zip(deallocOp.getMemrefs(), deallocOp.getConditions())) {
|
|
|
|
if (succeeded(handleOneMemref(deallocOp, memref, cond, rewriter)))
|
|
continue;
|
|
|
|
if (auto extractOp =
|
|
memref.getDefiningOp<memref::ExtractStridedMetadataOp>())
|
|
if (succeeded(handleOneMemref(deallocOp, extractOp.getOperand(), cond,
|
|
rewriter)))
|
|
continue;
|
|
|
|
newMemrefs.push_back(memref);
|
|
newConditions.push_back(cond);
|
|
}
|
|
|
|
// Return failure if we don't change anything such that we don't run into an
|
|
// infinite loop of pattern applications.
|
|
return updateDeallocIfChanged(deallocOp, newMemrefs, newConditions,
|
|
rewriter);
|
|
}
|
|
|
|
private:
|
|
BufferOriginAnalysis &analysis;
|
|
};
|
|
|
|
/// Remove memrefs from the `retained` list which are guaranteed to not alias
|
|
/// any memref in the `memrefs` list. The corresponding result value can be
|
|
/// replaced with `false` in that case according to the operation description.
|
|
///
|
|
/// Example:
|
|
/// ```mlir
|
|
/// %0:2 = bufferization.dealloc (%m : memref<2xi32>) if (%cond)
|
|
/// retain (%r0, %r1 : memref<2xi32>, memref<2xi32>)
|
|
/// return %0#0, %0#1
|
|
/// ```
|
|
/// can be canonicalized to the following given that `%r0` and `%r1` do not
|
|
/// alias `%m`:
|
|
/// ```mlir
|
|
/// bufferization.dealloc (%m : memref<2xi32>) if (%cond)
|
|
/// return %false, %false
|
|
/// ```
|
|
struct RemoveRetainedMemrefsGuaranteedToNotAlias
|
|
: public OpRewritePattern<DeallocOp> {
|
|
RemoveRetainedMemrefsGuaranteedToNotAlias(MLIRContext *context,
|
|
BufferOriginAnalysis &analysis)
|
|
: OpRewritePattern<DeallocOp>(context), analysis(analysis) {}
|
|
|
|
LogicalResult matchAndRewrite(DeallocOp deallocOp,
|
|
PatternRewriter &rewriter) const override {
|
|
SmallVector<Value> newRetainedMemrefs, replacements;
|
|
|
|
for (auto retainedMemref : deallocOp.getRetained()) {
|
|
if (potentiallyAliasesMemref(analysis, deallocOp.getMemrefs(),
|
|
retainedMemref)) {
|
|
newRetainedMemrefs.push_back(retainedMemref);
|
|
replacements.push_back({});
|
|
continue;
|
|
}
|
|
|
|
replacements.push_back(rewriter.create<arith::ConstantOp>(
|
|
deallocOp.getLoc(), rewriter.getBoolAttr(false)));
|
|
}
|
|
|
|
if (newRetainedMemrefs.size() == deallocOp.getRetained().size())
|
|
return failure();
|
|
|
|
auto newDeallocOp = rewriter.create<DeallocOp>(
|
|
deallocOp.getLoc(), deallocOp.getMemrefs(), deallocOp.getConditions(),
|
|
newRetainedMemrefs);
|
|
int i = 0;
|
|
for (auto &repl : replacements) {
|
|
if (!repl)
|
|
repl = newDeallocOp.getUpdatedConditions()[i++];
|
|
}
|
|
|
|
rewriter.replaceOp(deallocOp, replacements);
|
|
return success();
|
|
}
|
|
|
|
private:
|
|
BufferOriginAnalysis &analysis;
|
|
};
|
|
|
|
/// Split off memrefs to separate dealloc operations to reduce the number of
|
|
/// runtime checks required and enable further canonicalization of the new and
|
|
/// simpler dealloc operations. A memref can be split off if it is guaranteed to
|
|
/// not alias with any other memref in the `memref` operand list. The results
|
|
/// of the old and the new dealloc operation have to be combined by computing
|
|
/// the element-wise disjunction of them.
|
|
///
|
|
/// Example:
|
|
/// ```mlir
|
|
/// %0:2 = bufferization.dealloc (%m0, %m1 : memref<2xi32>, memref<2xi32>)
|
|
/// if (%cond0, %cond1)
|
|
/// retain (%r0, %r1 : memref<2xi32>, memref<2xi32>)
|
|
/// return %0#0, %0#1
|
|
/// ```
|
|
/// Given that `%m0` is guaranteed to never alias with `%m1`, the above IR is
|
|
/// canonicalized to the following, thus reducing the number of runtime alias
|
|
/// checks by 1 and potentially enabling further canonicalization of the new
|
|
/// split-up dealloc operations.
|
|
/// ```mlir
|
|
/// %0:2 = bufferization.dealloc (%m0 : memref<2xi32>) if (%cond0)
|
|
/// retain (%r0, %r1 : memref<2xi32>, memref<2xi32>)
|
|
/// %1:2 = bufferization.dealloc (%m1 : memref<2xi32>) if (%cond1)
|
|
/// retain (%r0, %r1 : memref<2xi32>, memref<2xi32>)
|
|
/// %2 = arith.ori %0#0, %1#0
|
|
/// %3 = arith.ori %0#1, %1#1
|
|
/// return %2, %3
|
|
/// ```
|
|
struct SplitDeallocWhenNotAliasingAnyOther
|
|
: public OpRewritePattern<DeallocOp> {
|
|
SplitDeallocWhenNotAliasingAnyOther(MLIRContext *context,
|
|
BufferOriginAnalysis &analysis)
|
|
: OpRewritePattern<DeallocOp>(context), analysis(analysis) {}
|
|
|
|
LogicalResult matchAndRewrite(DeallocOp deallocOp,
|
|
PatternRewriter &rewriter) const override {
|
|
Location loc = deallocOp.getLoc();
|
|
if (deallocOp.getMemrefs().size() <= 1)
|
|
return failure();
|
|
|
|
SmallVector<Value> remainingMemrefs, remainingConditions;
|
|
SmallVector<SmallVector<Value>> updatedConditions;
|
|
for (int64_t i = 0, e = deallocOp.getMemrefs().size(); i < e; ++i) {
|
|
Value memref = deallocOp.getMemrefs()[i];
|
|
Value cond = deallocOp.getConditions()[i];
|
|
SmallVector<Value> otherMemrefs(deallocOp.getMemrefs());
|
|
otherMemrefs.erase(otherMemrefs.begin() + i);
|
|
// Check if `memref` can split off into a separate bufferization.dealloc.
|
|
if (potentiallyAliasesMemref(analysis, otherMemrefs, memref)) {
|
|
// `memref` alias with other memrefs, do not split off.
|
|
remainingMemrefs.push_back(memref);
|
|
remainingConditions.push_back(cond);
|
|
continue;
|
|
}
|
|
|
|
// Create new bufferization.dealloc op for `memref`.
|
|
auto newDeallocOp = rewriter.create<DeallocOp>(loc, memref, cond,
|
|
deallocOp.getRetained());
|
|
updatedConditions.push_back(
|
|
llvm::to_vector(ValueRange(newDeallocOp.getUpdatedConditions())));
|
|
}
|
|
|
|
// Fail if no memref was split off.
|
|
if (remainingMemrefs.size() == deallocOp.getMemrefs().size())
|
|
return failure();
|
|
|
|
// Create bufferization.dealloc op for all remaining memrefs.
|
|
auto newDeallocOp = rewriter.create<DeallocOp>(
|
|
loc, remainingMemrefs, remainingConditions, deallocOp.getRetained());
|
|
|
|
// Bit-or all conditions.
|
|
SmallVector<Value> replacements =
|
|
llvm::to_vector(ValueRange(newDeallocOp.getUpdatedConditions()));
|
|
for (auto additionalConditions : updatedConditions) {
|
|
assert(replacements.size() == additionalConditions.size() &&
|
|
"expected same number of updated conditions");
|
|
for (int64_t i = 0, e = replacements.size(); i < e; ++i) {
|
|
replacements[i] = rewriter.create<arith::OrIOp>(
|
|
loc, replacements[i], additionalConditions[i]);
|
|
}
|
|
}
|
|
rewriter.replaceOp(deallocOp, replacements);
|
|
return success();
|
|
}
|
|
|
|
private:
|
|
BufferOriginAnalysis &analysis;
|
|
};
|
|
|
|
/// Check for every retained memref if a must-aliasing memref exists in the
|
|
/// 'memref' operand list with constant 'true' condition. If so, we can replace
|
|
/// the operation result corresponding to that retained memref with 'true'. If
|
|
/// this condition holds for all retained memrefs we can also remove the
|
|
/// aliasing memrefs and their conditions since they will never be deallocated
|
|
/// due to the must-alias and we don't need them to compute the result value
|
|
/// anymore since it got replaced with 'true'.
|
|
///
|
|
/// Example:
|
|
/// ```mlir
|
|
/// %0:2 = bufferization.dealloc (%arg0, %arg1, %arg2 : ...)
|
|
/// if (%true, %true, %true)
|
|
/// retain (%arg0, %arg1 : memref<2xi32>, memref<2xi32>)
|
|
/// ```
|
|
/// becomes
|
|
/// ```mlir
|
|
/// %0:2 = bufferization.dealloc (%arg2 : memref<2xi32>) if (%true)
|
|
/// retain (%arg0, %arg1 : memref<2xi32>, memref<2xi32>)
|
|
/// // replace %0#0 with %true
|
|
/// // replace %0#1 with %true
|
|
/// ```
|
|
/// Note that the dealloc operation will still have the result values, but they
|
|
/// don't have uses anymore.
|
|
struct RetainedMemrefAliasingAlwaysDeallocatedMemref
|
|
: public OpRewritePattern<DeallocOp> {
|
|
RetainedMemrefAliasingAlwaysDeallocatedMemref(MLIRContext *context,
|
|
BufferOriginAnalysis &analysis)
|
|
: OpRewritePattern<DeallocOp>(context), analysis(analysis) {}
|
|
|
|
LogicalResult matchAndRewrite(DeallocOp deallocOp,
|
|
PatternRewriter &rewriter) const override {
|
|
BitVector aliasesWithConstTrueMemref(deallocOp.getRetained().size());
|
|
SmallVector<Value> newMemrefs, newConditions;
|
|
for (auto [memref, cond] :
|
|
llvm::zip(deallocOp.getMemrefs(), deallocOp.getConditions())) {
|
|
bool canDropMemref = false;
|
|
for (auto [i, retained, res] : llvm::enumerate(
|
|
deallocOp.getRetained(), deallocOp.getUpdatedConditions())) {
|
|
if (!matchPattern(cond, m_One()))
|
|
continue;
|
|
|
|
std::optional<bool> analysisResult =
|
|
analysis.isSameAllocation(retained, memref);
|
|
if (analysisResult == true) {
|
|
rewriter.replaceAllUsesWith(res, cond);
|
|
aliasesWithConstTrueMemref[i] = true;
|
|
canDropMemref = true;
|
|
continue;
|
|
}
|
|
|
|
// TODO: once our alias analysis is powerful enough we can remove the
|
|
// rest of this loop body
|
|
auto extractOp =
|
|
memref.getDefiningOp<memref::ExtractStridedMetadataOp>();
|
|
if (!extractOp)
|
|
continue;
|
|
|
|
std::optional<bool> extractAnalysisResult =
|
|
analysis.isSameAllocation(retained, extractOp.getOperand());
|
|
if (extractAnalysisResult == true) {
|
|
rewriter.replaceAllUsesWith(res, cond);
|
|
aliasesWithConstTrueMemref[i] = true;
|
|
canDropMemref = true;
|
|
}
|
|
}
|
|
|
|
if (!canDropMemref) {
|
|
newMemrefs.push_back(memref);
|
|
newConditions.push_back(cond);
|
|
}
|
|
}
|
|
if (!aliasesWithConstTrueMemref.all())
|
|
return failure();
|
|
|
|
return updateDeallocIfChanged(deallocOp, newMemrefs, newConditions,
|
|
rewriter);
|
|
}
|
|
|
|
private:
|
|
BufferOriginAnalysis &analysis;
|
|
};
|
|
|
|
} // namespace
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// BufferDeallocationSimplificationPass
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
namespace {
|
|
|
|
/// The actual buffer deallocation pass that inserts and moves dealloc nodes
|
|
/// into the right positions. Furthermore, it inserts additional clones if
|
|
/// necessary. It uses the algorithm described at the top of the file.
|
|
struct BufferDeallocationSimplificationPass
|
|
: public bufferization::impl::BufferDeallocationSimplificationBase<
|
|
BufferDeallocationSimplificationPass> {
|
|
void runOnOperation() override {
|
|
BufferOriginAnalysis analysis(getOperation());
|
|
RewritePatternSet patterns(&getContext());
|
|
patterns.add<RemoveDeallocMemrefsContainedInRetained,
|
|
RemoveRetainedMemrefsGuaranteedToNotAlias,
|
|
SplitDeallocWhenNotAliasingAnyOther,
|
|
RetainedMemrefAliasingAlwaysDeallocatedMemref>(&getContext(),
|
|
analysis);
|
|
// We don't want that the block structure changes invalidating the
|
|
// `BufferOriginAnalysis` so we apply the rewrites witha `Normal` level of
|
|
// region simplification
|
|
GreedyRewriteConfig config;
|
|
config.enableRegionSimplification = GreedySimplifyRegionLevel::Normal;
|
|
populateDeallocOpCanonicalizationPatterns(patterns, &getContext());
|
|
|
|
if (failed(
|
|
applyPatternsGreedily(getOperation(), std::move(patterns), config)))
|
|
signalPassFailure();
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
std::unique_ptr<Pass>
|
|
mlir::bufferization::createBufferDeallocationSimplificationPass() {
|
|
return std::make_unique<BufferDeallocationSimplificationPass>();
|
|
}
|