546 lines
24 KiB
C++
546 lines
24 KiB
C++
|
//===- LowerDeallocations.cpp - Bufferization Deallocs to MemRef pass -----===//
|
||
|
|
//
|
||
|
|
// 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 patterns to convert `bufferization.dealloc` operations
|
||
|
|
// to the MemRef dialect.
|
||
|
|
//
|
||
|
|
//===----------------------------------------------------------------------===//
|
||
|
|
|
||
|
|
#include "mlir/Dialect/Arith/IR/Arith.h"
|
||
|
|
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
|
||
|
|
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
|
||
|
|
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
||
|
|
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
||
|
|
#include "mlir/Dialect/SCF/IR/SCF.h"
|
||
|
|
#include "mlir/IR/BuiltinTypes.h"
|
||
|
|
#include "mlir/Pass/Pass.h"
|
||
|
|
#include "mlir/Support/LogicalResult.h"
|
||
|
|
#include "mlir/Transforms/DialectConversion.h"
|
||
|
|
|
||
|
|
namespace mlir {
|
||
|
|
namespace bufferization {
|
||
|
|
#define GEN_PASS_DEF_LOWERDEALLOCATIONS
|
||
|
|
#include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
|
||
|
|
} // namespace bufferization
|
||
|
|
} // namespace mlir
|
||
|
|
|
||
|
|
using namespace mlir;
|
||
|
|
|
||
|
|
namespace {
|
||
|
|
/// The DeallocOpConversion transforms all bufferization dealloc operations into
|
||
|
|
/// memref dealloc operations potentially guarded by scf if operations.
|
||
|
|
/// Additionally, memref extract_aligned_pointer_as_index and arith operations
|
||
|
|
/// are inserted to compute the guard conditions. We distinguish multiple cases
|
||
|
|
/// to provide an overall more efficient lowering. In the general case, a helper
|
||
|
|
/// func is created to avoid quadratic code size explosion (relative to the
|
||
|
|
/// number of operands of the dealloc operation). For examples of each case,
|
||
|
|
/// refer to the documentation of the member functions of this class.
|
||
|
|
class DeallocOpConversion
|
||
|
|
: public OpConversionPattern<bufferization::DeallocOp> {
|
||
|
|
|
||
|
|
/// Lower a simple case without any retained values and a single memref to
|
||
|
|
/// avoiding the helper function. Ideally, static analysis can provide enough
|
||
|
|
/// aliasing information to split the dealloc operations up into this simple
|
||
|
|
/// case as much as possible before running this pass.
|
||
|
|
///
|
||
|
|
/// Example:
|
||
|
|
/// ```
|
||
|
|
/// bufferization.dealloc (%arg0 : memref<2xf32>) if (%arg1)
|
||
|
|
/// ```
|
||
|
|
/// is lowered to
|
||
|
|
/// ```
|
||
|
|
/// scf.if %arg1 {
|
||
|
|
/// memref.dealloc %arg0 : memref<2xf32>
|
||
|
|
/// }
|
||
|
|
/// ```
|
||
|
|
LogicalResult
|
||
|
|
rewriteOneMemrefNoRetainCase(bufferization::DeallocOp op, OpAdaptor adaptor,
|
||
|
|
ConversionPatternRewriter &rewriter) const {
|
||
|
|
assert(adaptor.getMemrefs().size() == 1 && "expected only one memref");
|
||
|
|
assert(adaptor.getRetained().empty() && "expected no retained memrefs");
|
||
|
|
|
||
|
|
rewriter.replaceOpWithNewOp<scf::IfOp>(
|
||
|
|
op, adaptor.getConditions()[0], [&](OpBuilder &builder, Location loc) {
|
||
|
|
builder.create<memref::DeallocOp>(loc, adaptor.getMemrefs()[0]);
|
||
|
|
builder.create<scf::YieldOp>(loc);
|
||
|
|
});
|
||
|
|
return success();
|
||
|
|
}
|
||
|
|
|
||
|
|
/// A special case lowering for the deallocation operation with exactly one
|
||
|
|
/// memref, but arbitrary number of retained values. This avoids the helper
|
||
|
|
/// function that the general case needs and thus also avoids storing indices
|
||
|
|
/// to specifically allocated memrefs. The size of the code produced by this
|
||
|
|
/// lowering is linear to the number of retained values.
|
||
|
|
///
|
||
|
|
/// Example:
|
||
|
|
/// ```mlir
|
||
|
|
/// %0:2 = bufferization.dealloc (%m : memref<2xf32>) if (%cond)
|
||
|
|
// retain (%r0, %r1 : memref<1xf32>, memref<2xf32>)
|
||
|
|
/// return %0#0, %0#1 : i1, i1
|
||
|
|
/// ```
|
||
|
|
/// ```mlir
|
||
|
|
/// %m_base_pointer = memref.extract_aligned_pointer_as_index %m
|
||
|
|
/// %r0_base_pointer = memref.extract_aligned_pointer_as_index %r0
|
||
|
|
/// %r0_does_not_alias = arith.cmpi ne, %m_base_pointer, %r0_base_pointer
|
||
|
|
/// %r1_base_pointer = memref.extract_aligned_pointer_as_index %r1
|
||
|
|
/// %r1_does_not_alias = arith.cmpi ne, %m_base_pointer, %r1_base_pointer
|
||
|
|
/// %not_retained = arith.andi %r0_does_not_alias, %r1_does_not_alias : i1
|
||
|
|
/// %should_dealloc = arith.andi %not_retained, %cond : i1
|
||
|
|
/// scf.if %should_dealloc {
|
||
|
|
/// memref.dealloc %m : memref<2xf32>
|
||
|
|
/// }
|
||
|
|
/// %true = arith.constant true
|
||
|
|
/// %r0_does_alias = arith.xori %r0_does_not_alias, %true : i1
|
||
|
|
/// %r0_ownership = arith.andi %r0_does_alias, %cond : i1
|
||
|
|
/// %r1_does_alias = arith.xori %r1_does_not_alias, %true : i1
|
||
|
|
/// %r1_ownership = arith.andi %r1_does_alias, %cond : i1
|
||
|
|
/// return %r0_ownership, %r1_ownership : i1, i1
|
||
|
|
/// ```
|
||
|
|
LogicalResult rewriteOneMemrefMultipleRetainCase(
|
||
|
|
bufferization::DeallocOp op, OpAdaptor adaptor,
|
||
|
|
ConversionPatternRewriter &rewriter) const {
|
||
|
|
assert(adaptor.getMemrefs().size() == 1 && "expected only one memref");
|
||
|
|
|
||
|
|
// Compute the base pointer indices, compare all retained indices to the
|
||
|
|
// memref index to check if they alias.
|
||
|
|
SmallVector<Value> doesNotAliasList;
|
||
|
|
Value memrefAsIdx = rewriter.create<memref::ExtractAlignedPointerAsIndexOp>(
|
||
|
|
op->getLoc(), adaptor.getMemrefs()[0]);
|
||
|
|
for (Value retained : adaptor.getRetained()) {
|
||
|
|
Value retainedAsIdx =
|
||
|
|
rewriter.create<memref::ExtractAlignedPointerAsIndexOp>(op->getLoc(),
|
||
|
|
retained);
|
||
|
|
Value doesNotAlias = rewriter.create<arith::CmpIOp>(
|
||
|
|
op->getLoc(), arith::CmpIPredicate::ne, memrefAsIdx, retainedAsIdx);
|
||
|
|
doesNotAliasList.push_back(doesNotAlias);
|
||
|
|
}
|
||
|
|
|
||
|
|
// AND-reduce the list of booleans from above.
|
||
|
|
Value prev = doesNotAliasList.front();
|
||
|
|
for (Value doesNotAlias : ArrayRef(doesNotAliasList).drop_front())
|
||
|
|
prev = rewriter.create<arith::AndIOp>(op->getLoc(), prev, doesNotAlias);
|
||
|
|
|
||
|
|
// Also consider the condition given by the dealloc operation and perform a
|
||
|
|
// conditional deallocation guarded by that value.
|
||
|
|
Value shouldDealloc = rewriter.create<arith::AndIOp>(
|
||
|
|
op->getLoc(), prev, adaptor.getConditions()[0]);
|
||
|
|
|
||
|
|
rewriter.create<scf::IfOp>(
|
||
|
|
op.getLoc(), shouldDealloc, [&](OpBuilder &builder, Location loc) {
|
||
|
|
builder.create<memref::DeallocOp>(loc, adaptor.getMemrefs()[0]);
|
||
|
|
builder.create<scf::YieldOp>(loc);
|
||
|
|
});
|
||
|
|
|
||
|
|
// Compute the replacement values for the dealloc operation results. This
|
||
|
|
// inserts an already canonicalized form of
|
||
|
|
// `select(does_alias_with_memref(r), memref_cond, false)` for each retained
|
||
|
|
// value r.
|
||
|
|
SmallVector<Value> replacements;
|
||
|
|
Value trueVal = rewriter.create<arith::ConstantOp>(
|
||
|
|
op->getLoc(), rewriter.getBoolAttr(true));
|
||
|
|
for (Value doesNotAlias : doesNotAliasList) {
|
||
|
|
Value aliases =
|
||
|
|
rewriter.create<arith::XOrIOp>(op->getLoc(), doesNotAlias, trueVal);
|
||
|
|
Value result = rewriter.create<arith::AndIOp>(op->getLoc(), aliases,
|
||
|
|
adaptor.getConditions()[0]);
|
||
|
|
replacements.push_back(result);
|
||
|
|
}
|
||
|
|
|
||
|
|
rewriter.replaceOp(op, replacements);
|
||
|
|
|
||
|
|
return success();
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Lowering that supports all features the dealloc operation has to offer. It
|
||
|
|
/// computes the base pointer of each memref (as an index), stores it in a
|
||
|
|
/// new memref helper structure and passes it to the helper function generated
|
||
|
|
/// in 'buildDeallocationHelperFunction'. The results are stored in two lists
|
||
|
|
/// (represented as memrefs) of booleans passed as arguments. The first list
|
||
|
|
/// stores whether the corresponding condition should be deallocated, the
|
||
|
|
/// second list stores the ownership of the retained values which can be used
|
||
|
|
/// to replace the result values of the `bufferization.dealloc` operation.
|
||
|
|
///
|
||
|
|
/// Example:
|
||
|
|
/// ```
|
||
|
|
/// %0:2 = bufferization.dealloc (%m0, %m1 : memref<2xf32>, memref<5xf32>)
|
||
|
|
/// if (%cond0, %cond1)
|
||
|
|
/// retain (%r0, %r1 : memref<1xf32>, memref<2xf32>)
|
||
|
|
/// ```
|
||
|
|
/// lowers to (simplified):
|
||
|
|
/// ```
|
||
|
|
/// %c0 = arith.constant 0 : index
|
||
|
|
/// %c1 = arith.constant 1 : index
|
||
|
|
/// %dealloc_base_pointer_list = memref.alloc() : memref<2xindex>
|
||
|
|
/// %cond_list = memref.alloc() : memref<2xi1>
|
||
|
|
/// %retain_base_pointer_list = memref.alloc() : memref<2xindex>
|
||
|
|
/// %m0_base_pointer = memref.extract_aligned_pointer_as_index %m0
|
||
|
|
/// memref.store %m0_base_pointer, %dealloc_base_pointer_list[%c0]
|
||
|
|
/// %m1_base_pointer = memref.extract_aligned_pointer_as_index %m1
|
||
|
|
/// memref.store %m1_base_pointer, %dealloc_base_pointer_list[%c1]
|
||
|
|
/// memref.store %cond0, %cond_list[%c0]
|
||
|
|
/// memref.store %cond1, %cond_list[%c1]
|
||
|
|
/// %r0_base_pointer = memref.extract_aligned_pointer_as_index %r0
|
||
|
|
/// memref.store %r0_base_pointer, %retain_base_pointer_list[%c0]
|
||
|
|
/// %r1_base_pointer = memref.extract_aligned_pointer_as_index %r1
|
||
|
|
/// memref.store %r1_base_pointer, %retain_base_pointer_list[%c1]
|
||
|
|
/// %dyn_dealloc_base_pointer_list = memref.cast %dealloc_base_pointer_list :
|
||
|
|
/// memref<2xindex> to memref<?xindex>
|
||
|
|
/// %dyn_cond_list = memref.cast %cond_list : memref<2xi1> to memref<?xi1>
|
||
|
|
/// %dyn_retain_base_pointer_list = memref.cast %retain_base_pointer_list :
|
||
|
|
/// memref<2xindex> to memref<?xindex>
|
||
|
|
/// %dealloc_cond_out = memref.alloc() : memref<2xi1>
|
||
|
|
/// %ownership_out = memref.alloc() : memref<2xi1>
|
||
|
|
/// %dyn_dealloc_cond_out = memref.cast %dealloc_cond_out :
|
||
|
|
/// memref<2xi1> to memref<?xi1>
|
||
|
|
/// %dyn_ownership_out = memref.cast %ownership_out :
|
||
|
|
/// memref<2xi1> to memref<?xi1>
|
||
|
|
/// call @dealloc_helper(%dyn_dealloc_base_pointer_list,
|
||
|
|
/// %dyn_retain_base_pointer_list,
|
||
|
|
/// %dyn_cond_list,
|
||
|
|
/// %dyn_dealloc_cond_out,
|
||
|
|
/// %dyn_ownership_out) : (...)
|
||
|
|
/// %m0_dealloc_cond = memref.load %dyn_dealloc_cond_out[%c0] : memref<2xi1>
|
||
|
|
/// scf.if %m0_dealloc_cond {
|
||
|
|
/// memref.dealloc %m0 : memref<2xf32>
|
||
|
|
/// }
|
||
|
|
/// %m1_dealloc_cond = memref.load %dyn_dealloc_cond_out[%c1] : memref<2xi1>
|
||
|
|
/// scf.if %m1_dealloc_cond {
|
||
|
|
/// memref.dealloc %m1 : memref<5xf32>
|
||
|
|
/// }
|
||
|
|
/// %r0_ownership = memref.load %dyn_ownership_out[%c0] : memref<2xi1>
|
||
|
|
/// %r1_ownership = memref.load %dyn_ownership_out[%c1] : memref<2xi1>
|
||
|
|
/// memref.dealloc %dealloc_base_pointer_list : memref<2xindex>
|
||
|
|
/// memref.dealloc %retain_base_pointer_list : memref<2xindex>
|
||
|
|
/// memref.dealloc %cond_list : memref<2xi1>
|
||
|
|
/// memref.dealloc %dealloc_cond_out : memref<2xi1>
|
||
|
|
/// memref.dealloc %ownership_out : memref<2xi1>
|
||
|
|
/// // replace %0#0 with %r0_ownership
|
||
|
|
/// // replace %0#1 with %r1_ownership
|
||
|
|
/// ```
|
||
|
|
LogicalResult rewriteGeneralCase(bufferization::DeallocOp op,
|
||
|
|
OpAdaptor adaptor,
|
||
|
|
ConversionPatternRewriter &rewriter) const {
|
||
|
|
// Allocate two memrefs holding the base pointer indices of the list of
|
||
|
|
// memrefs to be deallocated and the ones to be retained. These can then be
|
||
|
|
// passed to the helper function and the for-loops can iterate over them.
|
||
|
|
// Without storing them to memrefs, we could not use for-loops but only a
|
||
|
|
// completely unrolled version of it, potentially leading to code-size
|
||
|
|
// blow-up.
|
||
|
|
Value toDeallocMemref = rewriter.create<memref::AllocOp>(
|
||
|
|
op.getLoc(), MemRefType::get({(int64_t)adaptor.getMemrefs().size()},
|
||
|
|
rewriter.getIndexType()));
|
||
|
|
Value conditionMemref = rewriter.create<memref::AllocOp>(
|
||
|
|
op.getLoc(), MemRefType::get({(int64_t)adaptor.getConditions().size()},
|
||
|
|
rewriter.getI1Type()));
|
||
|
|
Value toRetainMemref = rewriter.create<memref::AllocOp>(
|
||
|
|
op.getLoc(), MemRefType::get({(int64_t)adaptor.getRetained().size()},
|
||
|
|
rewriter.getIndexType()));
|
||
|
|
|
||
|
|
auto getConstValue = [&](uint64_t value) -> Value {
|
||
|
|
return rewriter.create<arith::ConstantOp>(op.getLoc(),
|
||
|
|
rewriter.getIndexAttr(value));
|
||
|
|
};
|
||
|
|
|
||
|
|
// Extract the base pointers of the memrefs as indices to check for aliasing
|
||
|
|
// at runtime.
|
||
|
|
for (auto [i, toDealloc] : llvm::enumerate(adaptor.getMemrefs())) {
|
||
|
|
Value memrefAsIdx =
|
||
|
|
rewriter.create<memref::ExtractAlignedPointerAsIndexOp>(op.getLoc(),
|
||
|
|
toDealloc);
|
||
|
|
rewriter.create<memref::StoreOp>(op.getLoc(), memrefAsIdx,
|
||
|
|
toDeallocMemref, getConstValue(i));
|
||
|
|
}
|
||
|
|
|
||
|
|
for (auto [i, cond] : llvm::enumerate(adaptor.getConditions()))
|
||
|
|
rewriter.create<memref::StoreOp>(op.getLoc(), cond, conditionMemref,
|
||
|
|
getConstValue(i));
|
||
|
|
|
||
|
|
for (auto [i, toRetain] : llvm::enumerate(adaptor.getRetained())) {
|
||
|
|
Value memrefAsIdx =
|
||
|
|
rewriter.create<memref::ExtractAlignedPointerAsIndexOp>(op.getLoc(),
|
||
|
|
toRetain);
|
||
|
|
rewriter.create<memref::StoreOp>(op.getLoc(), memrefAsIdx, toRetainMemref,
|
||
|
|
getConstValue(i));
|
||
|
|
}
|
||
|
|
|
||
|
|
// Cast the allocated memrefs to dynamic shape because we want only one
|
||
|
|
// helper function no matter how many operands the bufferization.dealloc
|
||
|
|
// has.
|
||
|
|
Value castedDeallocMemref = rewriter.create<memref::CastOp>(
|
||
|
|
op->getLoc(),
|
||
|
|
MemRefType::get({ShapedType::kDynamic}, rewriter.getIndexType()),
|
||
|
|
toDeallocMemref);
|
||
|
|
Value castedCondsMemref = rewriter.create<memref::CastOp>(
|
||
|
|
op->getLoc(),
|
||
|
|
MemRefType::get({ShapedType::kDynamic}, rewriter.getI1Type()),
|
||
|
|
conditionMemref);
|
||
|
|
Value castedRetainMemref = rewriter.create<memref::CastOp>(
|
||
|
|
op->getLoc(),
|
||
|
|
MemRefType::get({ShapedType::kDynamic}, rewriter.getIndexType()),
|
||
|
|
toRetainMemref);
|
||
|
|
|
||
|
|
Value deallocCondsMemref = rewriter.create<memref::AllocOp>(
|
||
|
|
op.getLoc(), MemRefType::get({(int64_t)adaptor.getMemrefs().size()},
|
||
|
|
rewriter.getI1Type()));
|
||
|
|
Value retainCondsMemref = rewriter.create<memref::AllocOp>(
|
||
|
|
op.getLoc(), MemRefType::get({(int64_t)adaptor.getRetained().size()},
|
||
|
|
rewriter.getI1Type()));
|
||
|
|
|
||
|
|
Value castedDeallocCondsMemref = rewriter.create<memref::CastOp>(
|
||
|
|
op->getLoc(),
|
||
|
|
MemRefType::get({ShapedType::kDynamic}, rewriter.getI1Type()),
|
||
|
|
deallocCondsMemref);
|
||
|
|
Value castedRetainCondsMemref = rewriter.create<memref::CastOp>(
|
||
|
|
op->getLoc(),
|
||
|
|
MemRefType::get({ShapedType::kDynamic}, rewriter.getI1Type()),
|
||
|
|
retainCondsMemref);
|
||
|
|
|
||
|
|
rewriter.create<func::CallOp>(
|
||
|
|
op.getLoc(), deallocHelperFunc,
|
||
|
|
SmallVector<Value>{castedDeallocMemref, castedRetainMemref,
|
||
|
|
castedCondsMemref, castedDeallocCondsMemref,
|
||
|
|
castedRetainCondsMemref});
|
||
|
|
|
||
|
|
for (unsigned i = 0, e = adaptor.getMemrefs().size(); i < e; ++i) {
|
||
|
|
Value idxValue = getConstValue(i);
|
||
|
|
Value shouldDealloc = rewriter.create<memref::LoadOp>(
|
||
|
|
op.getLoc(), deallocCondsMemref, idxValue);
|
||
|
|
rewriter.create<scf::IfOp>(
|
||
|
|
op.getLoc(), shouldDealloc, [&](OpBuilder &builder, Location loc) {
|
||
|
|
builder.create<memref::DeallocOp>(loc, adaptor.getMemrefs()[i]);
|
||
|
|
builder.create<scf::YieldOp>(loc);
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
SmallVector<Value> replacements;
|
||
|
|
for (unsigned i = 0, e = adaptor.getRetained().size(); i < e; ++i) {
|
||
|
|
Value idxValue = getConstValue(i);
|
||
|
|
Value ownership = rewriter.create<memref::LoadOp>(
|
||
|
|
op.getLoc(), retainCondsMemref, idxValue);
|
||
|
|
replacements.push_back(ownership);
|
||
|
|
}
|
||
|
|
|
||
|
|
// Deallocate above allocated memrefs again to avoid memory leaks.
|
||
|
|
// Deallocation will not be run on code after this stage.
|
||
|
|
rewriter.create<memref::DeallocOp>(op.getLoc(), toDeallocMemref);
|
||
|
|
rewriter.create<memref::DeallocOp>(op.getLoc(), toRetainMemref);
|
||
|
|
rewriter.create<memref::DeallocOp>(op.getLoc(), conditionMemref);
|
||
|
|
rewriter.create<memref::DeallocOp>(op.getLoc(), deallocCondsMemref);
|
||
|
|
rewriter.create<memref::DeallocOp>(op.getLoc(), retainCondsMemref);
|
||
|
|
|
||
|
|
rewriter.replaceOp(op, replacements);
|
||
|
|
return success();
|
||
|
|
}
|
||
|
|
|
||
|
|
public:
|
||
|
|
DeallocOpConversion(MLIRContext *context, func::FuncOp deallocHelperFunc)
|
||
|
|
: OpConversionPattern<bufferization::DeallocOp>(context),
|
||
|
|
deallocHelperFunc(deallocHelperFunc) {}
|
||
|
|
|
||
|
|
LogicalResult
|
||
|
|
matchAndRewrite(bufferization::DeallocOp op, OpAdaptor adaptor,
|
||
|
|
ConversionPatternRewriter &rewriter) const override {
|
||
|
|
// Lower the trivial case.
|
||
|
|
if (adaptor.getMemrefs().empty()) {
|
||
|
|
Value falseVal = rewriter.create<arith::ConstantOp>(
|
||
|
|
op.getLoc(), rewriter.getBoolAttr(false));
|
||
|
|
rewriter.replaceOp(
|
||
|
|
op, SmallVector<Value>(adaptor.getRetained().size(), falseVal));
|
||
|
|
return success();
|
||
|
|
}
|
||
|
|
|
||
|
|
if (adaptor.getMemrefs().size() == 1 && adaptor.getRetained().empty())
|
||
|
|
return rewriteOneMemrefNoRetainCase(op, adaptor, rewriter);
|
||
|
|
|
||
|
|
if (adaptor.getMemrefs().size() == 1)
|
||
|
|
return rewriteOneMemrefMultipleRetainCase(op, adaptor, rewriter);
|
||
|
|
|
||
|
|
if (!deallocHelperFunc)
|
||
|
|
return op->emitError(
|
||
|
|
"library function required for generic lowering, but cannot be "
|
||
|
|
"automatically inserted when operating on functions");
|
||
|
|
|
||
|
|
return rewriteGeneralCase(op, adaptor, rewriter);
|
||
|
|
}
|
||
|
|
|
||
|
|
private:
|
||
|
|
func::FuncOp deallocHelperFunc;
|
||
|
|
};
|
||
|
|
} // namespace
|
||
|
|
|
||
|
|
namespace {
|
||
|
|
struct LowerDeallocationsPass
|
||
|
|
: public bufferization::impl::LowerDeallocationsBase<
|
||
|
|
LowerDeallocationsPass> {
|
||
|
|
void runOnOperation() override {
|
||
|
|
if (!isa<ModuleOp, FunctionOpInterface>(getOperation())) {
|
||
|
|
emitError(getOperation()->getLoc(),
|
||
|
|
"root operation must be a builtin.module or a function");
|
||
|
|
signalPassFailure();
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
func::FuncOp helperFuncOp;
|
||
|
|
if (auto module = dyn_cast<ModuleOp>(getOperation())) {
|
||
|
|
OpBuilder builder =
|
||
|
|
OpBuilder::atBlockBegin(&module.getBodyRegion().front());
|
||
|
|
SymbolTable symbolTable(module);
|
||
|
|
|
||
|
|
// Build dealloc helper function if there are deallocs.
|
||
|
|
getOperation()->walk([&](bufferization::DeallocOp deallocOp) {
|
||
|
|
if (deallocOp.getMemrefs().size() > 1) {
|
||
|
|
helperFuncOp = bufferization::buildDeallocationLibraryFunction(
|
||
|
|
builder, getOperation()->getLoc(), symbolTable);
|
||
|
|
return WalkResult::interrupt();
|
||
|
|
}
|
||
|
|
return WalkResult::advance();
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
RewritePatternSet patterns(&getContext());
|
||
|
|
bufferization::populateBufferizationDeallocLoweringPattern(patterns,
|
||
|
|
helperFuncOp);
|
||
|
|
|
||
|
|
ConversionTarget target(getContext());
|
||
|
|
target.addLegalDialect<memref::MemRefDialect, arith::ArithDialect,
|
||
|
|
scf::SCFDialect, func::FuncDialect>();
|
||
|
|
target.addIllegalOp<bufferization::DeallocOp>();
|
||
|
|
|
||
|
|
if (failed(applyPartialConversion(getOperation(), target,
|
||
|
|
std::move(patterns))))
|
||
|
|
signalPassFailure();
|
||
|
|
}
|
||
|
|
};
|
||
|
|
} // namespace
|
||
|
|
|
||
|
|
func::FuncOp mlir::bufferization::buildDeallocationLibraryFunction(
|
||
|
|
OpBuilder &builder, Location loc, SymbolTable &symbolTable) {
|
||
|
|
Type indexMemrefType =
|
||
|
|
MemRefType::get({ShapedType::kDynamic}, builder.getIndexType());
|
||
|
|
Type boolMemrefType =
|
||
|
|
MemRefType::get({ShapedType::kDynamic}, builder.getI1Type());
|
||
|
|
SmallVector<Type> argTypes{indexMemrefType, indexMemrefType, boolMemrefType,
|
||
|
|
boolMemrefType, boolMemrefType};
|
||
|
|
builder.clearInsertionPoint();
|
||
|
|
|
||
|
|
// Generate the func operation itself.
|
||
|
|
auto helperFuncOp = func::FuncOp::create(
|
||
|
|
loc, "dealloc_helper", builder.getFunctionType(argTypes, {}));
|
||
|
|
helperFuncOp.setVisibility(SymbolTable::Visibility::Private);
|
||
|
|
symbolTable.insert(helperFuncOp);
|
||
|
|
auto &block = helperFuncOp.getFunctionBody().emplaceBlock();
|
||
|
|
block.addArguments(argTypes, SmallVector<Location>(argTypes.size(), loc));
|
||
|
|
|
||
|
|
builder.setInsertionPointToStart(&block);
|
||
|
|
Value toDeallocMemref = helperFuncOp.getArguments()[0];
|
||
|
|
Value toRetainMemref = helperFuncOp.getArguments()[1];
|
||
|
|
Value conditionMemref = helperFuncOp.getArguments()[2];
|
||
|
|
Value deallocCondsMemref = helperFuncOp.getArguments()[3];
|
||
|
|
Value retainCondsMemref = helperFuncOp.getArguments()[4];
|
||
|
|
|
||
|
|
// Insert some prerequisites.
|
||
|
|
Value c0 = builder.create<arith::ConstantOp>(loc, builder.getIndexAttr(0));
|
||
|
|
Value c1 = builder.create<arith::ConstantOp>(loc, builder.getIndexAttr(1));
|
||
|
|
Value trueValue =
|
||
|
|
builder.create<arith::ConstantOp>(loc, builder.getBoolAttr(true));
|
||
|
|
Value falseValue =
|
||
|
|
builder.create<arith::ConstantOp>(loc, builder.getBoolAttr(false));
|
||
|
|
Value toDeallocSize = builder.create<memref::DimOp>(loc, toDeallocMemref, c0);
|
||
|
|
Value toRetainSize = builder.create<memref::DimOp>(loc, toRetainMemref, c0);
|
||
|
|
|
||
|
|
builder.create<scf::ForOp>(
|
||
|
|
loc, c0, toRetainSize, c1, std::nullopt,
|
||
|
|
[&](OpBuilder &builder, Location loc, Value i, ValueRange iterArgs) {
|
||
|
|
builder.create<memref::StoreOp>(loc, falseValue, retainCondsMemref, i);
|
||
|
|
builder.create<scf::YieldOp>(loc);
|
||
|
|
});
|
||
|
|
|
||
|
|
builder.create<scf::ForOp>(
|
||
|
|
loc, c0, toDeallocSize, c1, std::nullopt,
|
||
|
|
[&](OpBuilder &builder, Location loc, Value outerIter,
|
||
|
|
ValueRange iterArgs) {
|
||
|
|
Value toDealloc =
|
||
|
|
builder.create<memref::LoadOp>(loc, toDeallocMemref, outerIter);
|
||
|
|
Value cond =
|
||
|
|
builder.create<memref::LoadOp>(loc, conditionMemref, outerIter);
|
||
|
|
|
||
|
|
// Build the first for loop that computes aliasing with retained
|
||
|
|
// memrefs.
|
||
|
|
Value noRetainAlias =
|
||
|
|
builder
|
||
|
|
.create<scf::ForOp>(
|
||
|
|
loc, c0, toRetainSize, c1, trueValue,
|
||
|
|
[&](OpBuilder &builder, Location loc, Value i,
|
||
|
|
ValueRange iterArgs) {
|
||
|
|
Value retainValue = builder.create<memref::LoadOp>(
|
||
|
|
loc, toRetainMemref, i);
|
||
|
|
Value doesAlias = builder.create<arith::CmpIOp>(
|
||
|
|
loc, arith::CmpIPredicate::eq, retainValue,
|
||
|
|
toDealloc);
|
||
|
|
builder.create<scf::IfOp>(
|
||
|
|
loc, doesAlias,
|
||
|
|
[&](OpBuilder &builder, Location loc) {
|
||
|
|
Value retainCondValue =
|
||
|
|
builder.create<memref::LoadOp>(
|
||
|
|
loc, retainCondsMemref, i);
|
||
|
|
Value aggregatedRetainCond =
|
||
|
|
builder.create<arith::OrIOp>(
|
||
|
|
loc, retainCondValue, cond);
|
||
|
|
builder.create<memref::StoreOp>(
|
||
|
|
loc, aggregatedRetainCond, retainCondsMemref,
|
||
|
|
i);
|
||
|
|
builder.create<scf::YieldOp>(loc);
|
||
|
|
});
|
||
|
|
Value doesntAlias = builder.create<arith::CmpIOp>(
|
||
|
|
loc, arith::CmpIPredicate::ne, retainValue,
|
||
|
|
toDealloc);
|
||
|
|
Value yieldValue = builder.create<arith::AndIOp>(
|
||
|
|
loc, iterArgs[0], doesntAlias);
|
||
|
|
builder.create<scf::YieldOp>(loc, yieldValue);
|
||
|
|
})
|
||
|
|
.getResult(0);
|
||
|
|
|
||
|
|
// Build the second for loop that adds aliasing with previously
|
||
|
|
// deallocated memrefs.
|
||
|
|
Value noAlias =
|
||
|
|
builder
|
||
|
|
.create<scf::ForOp>(
|
||
|
|
loc, c0, outerIter, c1, noRetainAlias,
|
||
|
|
[&](OpBuilder &builder, Location loc, Value i,
|
||
|
|
ValueRange iterArgs) {
|
||
|
|
Value prevDeallocValue = builder.create<memref::LoadOp>(
|
||
|
|
loc, toDeallocMemref, i);
|
||
|
|
Value doesntAlias = builder.create<arith::CmpIOp>(
|
||
|
|
loc, arith::CmpIPredicate::ne, prevDeallocValue,
|
||
|
|
toDealloc);
|
||
|
|
Value yieldValue = builder.create<arith::AndIOp>(
|
||
|
|
loc, iterArgs[0], doesntAlias);
|
||
|
|
builder.create<scf::YieldOp>(loc, yieldValue);
|
||
|
|
})
|
||
|
|
.getResult(0);
|
||
|
|
|
||
|
|
Value shouldDealoc = builder.create<arith::AndIOp>(loc, noAlias, cond);
|
||
|
|
builder.create<memref::StoreOp>(loc, shouldDealoc, deallocCondsMemref,
|
||
|
|
outerIter);
|
||
|
|
builder.create<scf::YieldOp>(loc);
|
||
|
|
});
|
||
|
|
|
||
|
|
builder.create<func::ReturnOp>(loc);
|
||
|
|
return helperFuncOp;
|
||
|
|
}
|
||
|
|
|
||
|
|
void mlir::bufferization::populateBufferizationDeallocLoweringPattern(
|
||
|
|
RewritePatternSet &patterns, func::FuncOp deallocLibraryFunc) {
|
||
|
|
patterns.add<DeallocOpConversion>(patterns.getContext(), deallocLibraryFunc);
|
||
|
|
}
|
||
|
|
|
||
|
|
std::unique_ptr<Pass> mlir::bufferization::createLowerDeallocationsPass() {
|
||
|
|
return std::make_unique<LowerDeallocationsPass>();
|
||
|
|
}
|