llvm-project/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
Andrei Golubev 128caa1ba3
[mlir][bufferization] Refine tensor-buffer compatibility checks (#167705)
Generally, to_tensor and to_buffer already perform sufficient
verification. However, there are some unnecessarily strict constraints:
* builtin tensor requires its buffer counterpart to always be memref
* to_buffer on ranked tensor requires to always return memref

These checks are assertions (i.e. preconditions), however, they actually
prevent an apparently useful bufferization where builtin tensors could
become custom buffers. Lift these assertions, maintaining the
verification procedure unchanged, to allow builtin -> custom
bufferizations at operation boundary level.
2025-11-18 11:18:53 +01:00

1073 lines
41 KiB
C++

//===- BufferizableOpInterface.cpp - Bufferizable Ops ---=----------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/AsmState.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/IR/Value.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "llvm/ADT/ScopeExit.h"
//===----------------------------------------------------------------------===//
// BufferizableOpInterface
//===----------------------------------------------------------------------===//
namespace mlir {
namespace bufferization {
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.cpp.inc"
} // namespace bufferization
} // namespace mlir
MLIR_DEFINE_EXPLICIT_TYPE_ID(mlir::bufferization::AnalysisState)
#define DEBUG_TYPE "bufferizable-op-interface"
using namespace mlir;
using namespace bufferization;
static bool isRepetitiveRegion(Region *region,
const BufferizationOptions &options) {
Operation *op = region->getParentOp();
if (auto bufferizableOp = options.dynCastBufferizableOp(op))
if (bufferizableOp.isRepetitiveRegion(region->getRegionNumber()))
return true;
return false;
}
Region *AnalysisState::getEnclosingRepetitiveRegion(
Operation *op, const BufferizationOptions &options) {
if (!op->getBlock())
return nullptr;
if (auto iter = enclosingRepetitiveRegionCache.find_as(op);
iter != enclosingRepetitiveRegionCache.end())
return iter->second;
return enclosingRepetitiveRegionCache[op] =
getEnclosingRepetitiveRegion(op->getBlock(), options);
}
Region *AnalysisState::getEnclosingRepetitiveRegion(
Value value, const BufferizationOptions &options) {
if (auto iter = enclosingRepetitiveRegionCache.find_as(value);
iter != enclosingRepetitiveRegionCache.end())
return iter->second;
Region *region = value.getParentRegion();
// Collect all visited regions since we only know the repetitive region we
// want to map it to later on
SmallVector<Region *> visitedRegions;
while (region) {
visitedRegions.push_back(region);
if (isRepetitiveRegion(region, options))
break;
region = region->getParentRegion();
}
enclosingRepetitiveRegionCache[value] = region;
for (Region *r : visitedRegions)
enclosingRepetitiveRegionCache[r] = region;
return region;
}
Region *AnalysisState::getEnclosingRepetitiveRegion(
Block *block, const BufferizationOptions &options) {
if (auto iter = enclosingRepetitiveRegionCache.find_as(block);
iter != enclosingRepetitiveRegionCache.end())
return iter->second;
Region *region = block->getParent();
Operation *op = nullptr;
// Collect all visited regions since we only know the repetitive region we
// want to map it to later on
SmallVector<Region *> visitedRegions;
do {
op = region->getParentOp();
if (isRepetitiveRegion(region, options))
break;
} while ((region = op->getParentRegion()));
enclosingRepetitiveRegionCache[block] = region;
for (Region *r : visitedRegions)
enclosingRepetitiveRegionCache[r] = region;
return region;
}
bool AnalysisState::insideMutuallyExclusiveRegions(Operation *op0,
Operation *op1) {
auto key = std::make_pair(op0, op1);
if (auto iter = insideMutuallyExclusiveRegionsCache.find(key);
iter != insideMutuallyExclusiveRegionsCache.end())
return iter->second;
bool result = ::mlir::insideMutuallyExclusiveRegions(op0, op1);
// Populate results for both orderings of the ops.
insideMutuallyExclusiveRegionsCache[key] = result;
insideMutuallyExclusiveRegionsCache[std::make_pair(op1, op0)] = result;
return result;
}
void AnalysisState::resetCache() {
enclosingRepetitiveRegionCache.clear();
insideMutuallyExclusiveRegionsCache.clear();
}
SymbolTableCollection &BufferizationState::getSymbolTables() {
return symbolTables;
}
Region *bufferization::getNextEnclosingRepetitiveRegion(
Region *region, const BufferizationOptions &options) {
assert(isRepetitiveRegion(region, options) && "expected repetitive region");
while ((region = region->getParentRegion())) {
if (isRepetitiveRegion(region, options))
break;
}
return region;
}
Region *bufferization::getParallelRegion(Region *region,
const BufferizationOptions &options) {
while (region) {
auto bufferizableOp = options.dynCastBufferizableOp(region->getParentOp());
if (bufferizableOp &&
bufferizableOp.isParallelRegion(region->getRegionNumber())) {
assert(isRepetitiveRegion(region, options) &&
"expected that all parallel regions are also repetitive regions");
return region;
}
region = region->getParentRegion();
}
return nullptr;
}
Operation *bufferization::getOwnerOfValue(Value value) {
if (auto opResult = llvm::dyn_cast<OpResult>(value))
return opResult.getDefiningOp();
return llvm::cast<BlockArgument>(value).getOwner()->getParentOp();
}
/// Create an AllocTensorOp for the given shaped value. If `copy` is set, the
/// shaped value is copied. Otherwise, a tensor with undefined contents is
/// allocated.
FailureOr<Value> bufferization::allocateTensorForShapedValue(
OpBuilder &b, Location loc, Value shapedValue,
const BufferizationOptions &options, const BufferizationState &state,
bool copy) {
Value tensor;
if (llvm::isa<RankedTensorType>(shapedValue.getType())) {
tensor = shapedValue;
} else if (llvm::isa<MemRefType>(shapedValue.getType())) {
tensor = ToTensorOp::create(
b, loc, memref::getTensorTypeFromMemRefType(shapedValue.getType()),
shapedValue);
} else if (llvm::isa<UnrankedTensorType>(shapedValue.getType()) ||
llvm::isa<UnrankedMemRefType>(shapedValue.getType())) {
return getOwnerOfValue(shapedValue)
->emitError("copying of unranked tensors is not implemented");
} else {
llvm_unreachable("expected RankedTensorType or MemRefType");
}
RankedTensorType tensorType = llvm::cast<RankedTensorType>(tensor.getType());
SmallVector<Value> dynamicSizes;
if (!copy) {
// Compute the dynamic part of the shape.
// First try to query the shape via ReifyRankedShapedTypeOpInterface.
bool reifiedShapes = false;
if (llvm::isa<RankedTensorType>(shapedValue.getType()) &&
llvm::isa<OpResult>(shapedValue)) {
ReifiedRankedShapedTypeDims resultDims;
if (succeeded(
reifyResultShapes(b, shapedValue.getDefiningOp(), resultDims))) {
reifiedShapes = true;
auto &shape =
resultDims[llvm::cast<OpResult>(shapedValue).getResultNumber()];
for (const auto &dim : enumerate(tensorType.getShape())) {
if (ShapedType::isDynamic(dim.value())) {
dynamicSizes.push_back(
getValueOrCreateConstantIndexOp(b, loc, shape[dim.index()]));
}
}
}
}
// If the shape could not be reified, create DimOps.
if (!reifiedShapes)
populateDynamicDimSizes(b, loc, tensor, dynamicSizes);
}
// Create AllocTensorOp.
auto allocTensorOp = AllocTensorOp::create(b, loc, tensorType, dynamicSizes,
copy ? tensor : Value());
// Add 'memory_space' attribute. Not needed if 'copy' operand is specified.
if (copy)
return allocTensorOp.getResult();
auto copyBufferType =
detail::asMemRefType(getBufferType(tensor, options, state));
if (failed(copyBufferType))
return failure();
std::optional<Attribute> memorySpace = copyBufferType->getMemorySpace();
if (!memorySpace)
memorySpace = options.defaultMemorySpaceFn(tensorType);
if (memorySpace.has_value())
allocTensorOp.setMemorySpaceAttr(memorySpace.value());
return allocTensorOp.getResult();
}
LogicalResult BufferizableOpInterface::resolveTensorOpOperandConflicts(
RewriterBase &rewriter, const AnalysisState &analysisState,
const BufferizationState &bufferizationState) {
OpBuilder::InsertionGuard g(rewriter);
Operation *op = getOperation();
SmallVector<OpOperand *> outOfPlaceOpOperands;
DenseSet<OpOperand *> copiedOpOperands;
SmallVector<Value> outOfPlaceValues;
DenseSet<Value> copiedOpValues;
// Find all out-of-place OpOperands.
for (OpOperand &opOperand : op->getOpOperands()) {
Type operandType = opOperand.get().getType();
if (!llvm::isa<TensorType>(operandType))
continue;
if (analysisState.isInPlace(opOperand))
continue;
if (llvm::isa<UnrankedTensorType>(operandType))
return op->emitError("copying of unranked tensors is not implemented");
AliasingValueList aliasingValues =
analysisState.getAliasingValues(opOperand);
if (aliasingValues.getNumAliases() == 1 &&
isa<OpResult>(aliasingValues.getAliases()[0].value) &&
!analysisState.bufferizesToMemoryWrite(opOperand) &&
analysisState
.getAliasingOpOperands(aliasingValues.getAliases()[0].value)
.getNumAliases() == 1 &&
!isa<UnrankedTensorType>(
aliasingValues.getAliases()[0].value.getType())) {
// The op itself does not write but may create exactly one alias. Instead
// of copying the OpOperand, copy the OpResult. The OpResult can sometimes
// be smaller than the OpOperand (e.g., in the case of an extract_slice,
// where the result is usually a smaller part of the source). Do not apply
// this optimization if the OpResult is an unranked tensor (because those
// cannot be copied at the moment).
Value value = aliasingValues.getAliases()[0].value;
outOfPlaceValues.push_back(value);
if (!analysisState.canOmitTensorCopy(opOperand))
copiedOpValues.insert(value);
} else {
// In all other cases, make a copy of the OpOperand.
outOfPlaceOpOperands.push_back(&opOperand);
if (!analysisState.canOmitTensorCopy(opOperand))
copiedOpOperands.insert(&opOperand);
}
}
// Insert copies of OpOperands.
rewriter.setInsertionPoint(op);
for (OpOperand *opOperand : outOfPlaceOpOperands) {
FailureOr<Value> copy = allocateTensorForShapedValue(
rewriter, op->getLoc(), opOperand->get(), analysisState.getOptions(),
bufferizationState, copiedOpOperands.contains(opOperand));
if (failed(copy))
return failure();
rewriter.modifyOpInPlace(op, [&]() { opOperand->set(*copy); });
}
// Insert copies of Values.
rewriter.setInsertionPointAfter(op);
for (Value value : outOfPlaceValues) {
FailureOr<Value> copy = allocateTensorForShapedValue(
rewriter, op->getLoc(), value, analysisState.getOptions(),
bufferizationState, copiedOpValues.count(value));
if (failed(copy))
return failure();
SmallVector<OpOperand *> uses = llvm::to_vector(
llvm::map_range(value.getUses(), [](OpOperand &use) { return &use; }));
for (OpOperand *use : uses) {
// Do not update the alloc_tensor op that we just created.
if (use->getOwner() == copy->getDefiningOp())
continue;
// tensor.dim ops may have been created to be used as alloc_tensor op
// dynamic extents. Do not update these either.
if (isa<tensor::DimOp>(use->getOwner()))
continue;
rewriter.modifyOpInPlace(use->getOwner(), [&]() { use->set(*copy); });
}
}
return success();
}
//===----------------------------------------------------------------------===//
// OpFilter
//===----------------------------------------------------------------------===//
bool OpFilter::isOpAllowed(Operation *op) const {
// All other ops: Allow/disallow according to filter.
bool isAllowed = !hasAllowRule();
for (const Entry &entry : entries) {
bool filterResult = entry.fn(op);
switch (entry.type) {
case Entry::ALLOW:
isAllowed |= filterResult;
break;
case Entry::DENY:
if (filterResult)
// DENY filter matches. This op is no allowed. (Even if other ALLOW
// filters may match.)
return false;
};
}
return isAllowed;
}
//===----------------------------------------------------------------------===//
// BufferizationOptions
//===----------------------------------------------------------------------===//
namespace {
/// Default function arg type converter: Use a fully dynamic layout map.
BufferLikeType
defaultFunctionArgTypeConverter(TensorLikeType type, Attribute memorySpace,
func::FuncOp funcOp,
const BufferizationOptions &options) {
if (auto tensorType = mlir::dyn_cast<TensorType>(type)) {
return cast<BufferLikeType>(
getMemRefTypeWithFullyDynamicLayout(tensorType, memorySpace));
}
// If not builtin, fallback to TensorLikeType::getBufferType()
auto bufferType =
type.getBufferType(options, [&]() { return funcOp->emitError(); });
assert(succeeded(bufferType) &&
"a valid buffer is always expected at function boundary");
return *bufferType;
}
/// Default unknown type converter: Use a fully dynamic layout map.
BaseMemRefType
defaultUnknownTypeConverter(TensorType tensorType, Attribute memorySpace,
const BufferizationOptions &options) {
return getMemRefTypeWithFullyDynamicLayout(tensorType, memorySpace);
}
} // namespace
// Default constructor for BufferizationOptions.
BufferizationOptions::BufferizationOptions()
: functionArgTypeConverterFn(defaultFunctionArgTypeConverter),
unknownTypeConverterFn(defaultUnknownTypeConverter) {}
bool BufferizationOptions::isOpAllowed(Operation *op) const {
// Special case: If function boundary bufferization is deactivated, do not
// allow ops that belong to the `func` dialect.
bool isFuncBoundaryOp = isa_and_nonnull<func::FuncDialect>(op->getDialect());
if (!bufferizeFunctionBoundaries && isFuncBoundaryOp)
return false;
return opFilter.isOpAllowed(op);
}
BufferizableOpInterface
BufferizationOptions::dynCastBufferizableOp(Operation *op) const {
if (!isOpAllowed(op))
return nullptr;
auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
if (!bufferizableOp)
return nullptr;
return bufferizableOp;
}
BufferizableOpInterface
BufferizationOptions::dynCastBufferizableOp(Value value) const {
return dynCastBufferizableOp(getOwnerOfValue(value));
}
void BufferizationOptions::setFunctionBoundaryTypeConversion(
LayoutMapOption layoutMapOption) {
functionArgTypeConverterFn = [=](TensorLikeType type, Attribute memorySpace,
func::FuncOp funcOp,
const BufferizationOptions &options) {
if (auto tensorType = mlir::dyn_cast<TensorType>(type)) {
if (layoutMapOption == LayoutMapOption::IdentityLayoutMap)
return cast<BufferLikeType>(
bufferization::getMemRefTypeWithStaticIdentityLayout(tensorType,
memorySpace));
return cast<BufferLikeType>(
bufferization::getMemRefTypeWithFullyDynamicLayout(tensorType,
memorySpace));
}
// If not builtin, fallback to TensorLikeType::getBufferType()
auto bufferType =
type.getBufferType(options, [&]() { return funcOp->emitError(); });
assert(succeeded(bufferType) &&
"a valid buffer is always expected at function boundary");
return *bufferType;
};
inferFunctionResultLayout =
layoutMapOption == LayoutMapOption::InferLayoutMap;
}
//===----------------------------------------------------------------------===//
// Helper functions for BufferizableOpInterface
//===----------------------------------------------------------------------===//
static void setInsertionPointAfter(OpBuilder &b, Value value) {
if (auto bbArg = llvm::dyn_cast<BlockArgument>(value)) {
b.setInsertionPointToStart(bbArg.getOwner());
} else {
b.setInsertionPointAfter(value.getDefiningOp());
}
}
/// Determine which OpOperand* will alias with `value` if the op is bufferized
/// in place. Return all tensor OpOperand* if the op is not bufferizable.
AliasingOpOperandList AnalysisState::getAliasingOpOperands(Value value) const {
if (Operation *op = getOwnerOfValue(value))
if (auto bufferizableOp = getOptions().dynCastBufferizableOp(op))
return bufferizableOp.getAliasingOpOperands(value, *this);
// The op is not bufferizable.
return detail::unknownGetAliasingOpOperands(value);
}
/// Determine which Values will alias with `opOperand` if the op is bufferized
/// in place. Return all tensor Values if the op is not bufferizable.
AliasingValueList AnalysisState::getAliasingValues(OpOperand &opOperand) const {
if (auto bufferizableOp =
getOptions().dynCastBufferizableOp(opOperand.getOwner()))
return bufferizableOp.getAliasingValues(opOperand, *this);
// The op is not bufferizable.
return detail::unknownGetAliasingValues(opOperand);
}
/// Return true if `opOperand` bufferizes to a memory read. Return `true` if the
/// op is not bufferizable.
bool AnalysisState::bufferizesToMemoryRead(OpOperand &opOperand) const {
if (auto bufferizableOp =
getOptions().dynCastBufferizableOp(opOperand.getOwner()))
return bufferizableOp.bufferizesToMemoryRead(opOperand, *this);
// Unknown op that returns a tensor. The inplace analysis does not support it.
// Conservatively return true.
return true;
}
/// Return true if `opOperand` bufferizes to a memory write. Return
/// `true` if the op is not bufferizable.
bool AnalysisState::bufferizesToMemoryWrite(OpOperand &opOperand) const {
if (auto bufferizableOp =
getOptions().dynCastBufferizableOp(opOperand.getOwner()))
return bufferizableOp.bufferizesToMemoryWrite(opOperand, *this);
// Unknown op that returns a tensor. The inplace analysis does not support it.
// Conservatively return true.
return true;
}
/// Return true if `opOperand` does neither read nor write but bufferizes to an
/// alias. Return false if the op is not bufferizable.
bool AnalysisState::bufferizesToAliasOnly(OpOperand &opOperand) const {
if (auto bufferizableOp =
getOptions().dynCastBufferizableOp(opOperand.getOwner()))
return bufferizableOp.bufferizesToAliasOnly(opOperand, *this);
// Unknown op that returns a tensor. The inplace analysis does not support it.
// Conservatively return false.
return false;
}
bool AnalysisState::bufferizesToMemoryWrite(Value value) const {
auto opResult = llvm::dyn_cast<OpResult>(value);
if (!opResult)
return true;
auto bufferizableOp = getOptions().dynCastBufferizableOp(value);
if (!bufferizableOp)
return true;
return bufferizableOp.resultBufferizesToMemoryWrite(opResult, *this);
}
/// Return true if the given value is read by an op that bufferizes to a memory
/// read. Also takes into account ops that create an alias but do not read by
/// themselves (e.g., ExtractSliceOp).
bool AnalysisState::isValueRead(Value value) const {
assert(llvm::isa<TensorType>(value.getType()) && "expected TensorType");
SmallVector<OpOperand *> workingSet;
DenseSet<OpOperand *> visited;
for (OpOperand &use : value.getUses())
workingSet.push_back(&use);
while (!workingSet.empty()) {
OpOperand *uMaybeReading = workingSet.pop_back_val();
if (!visited.insert(uMaybeReading).second)
continue;
// Skip over all ops that neither read nor write (but create an alias).
if (bufferizesToAliasOnly(*uMaybeReading))
for (AliasingValue alias : getAliasingValues(*uMaybeReading))
for (OpOperand &use : alias.value.getUses())
workingSet.push_back(&use);
if (bufferizesToMemoryRead(*uMaybeReading))
return true;
}
return false;
}
// Starting from `opOperand`, follow the use-def chain in reverse, always
// selecting the aliasing OpOperands. Find and return Values for which
// `condition` evaluates to true. Uses of such matching Values are not
// traversed any further, the visited aliasing opOperands will be preserved
// through `visitedOpOperands`.
llvm::SetVector<Value> AnalysisState::findValueInReverseUseDefChain(
OpOperand *opOperand, llvm::function_ref<bool(Value)> condition,
TraversalConfig config,
llvm::DenseSet<OpOperand *> *visitedOpOperands) const {
llvm::DenseSet<Value> visited;
llvm::SetVector<Value> result, workingSet;
workingSet.insert(opOperand->get());
if (visitedOpOperands)
visitedOpOperands->insert(opOperand);
while (!workingSet.empty()) {
Value value = workingSet.pop_back_val();
if (!config.revisitAlreadyVisitedValues && visited.contains(value)) {
// Stop traversal if value was already visited.
if (config.alwaysIncludeLeaves)
result.insert(value);
continue;
}
visited.insert(value);
if (condition(value)) {
result.insert(value);
continue;
}
if (!config.followUnknownOps && !options.dynCastBufferizableOp(value)) {
// Stop iterating if `followUnknownOps` is unset and the op is either
// not bufferizable or excluded in the OpFilter.
if (config.alwaysIncludeLeaves)
result.insert(value);
continue;
}
AliasingOpOperandList aliases = getAliasingOpOperands(value);
if (aliases.getNumAliases() == 0) {
// The traversal ends naturally if there are no more OpOperands that
// could be followed.
if (config.alwaysIncludeLeaves)
result.insert(value);
continue;
}
for (AliasingOpOperand a : aliases) {
if (config.followEquivalentOnly &&
a.relation != BufferRelation::Equivalent) {
// Stop iterating if `followEquivalentOnly` is set but the alias is not
// equivalent.
if (config.alwaysIncludeLeaves)
result.insert(value);
continue;
}
if (config.followInPlaceOnly && !isInPlace(*a.opOperand)) {
// Stop iterating if `followInPlaceOnly` is set but the alias is
// out-of-place.
if (config.alwaysIncludeLeaves)
result.insert(value);
continue;
}
if (config.followSameTypeOrCastsOnly &&
a.opOperand->get().getType() != value.getType() &&
!value.getDefiningOp<CastOpInterface>()) {
// Stop iterating if `followSameTypeOrCastsOnly` is set but the alias is
// has a different type and the op is not a cast.
if (config.alwaysIncludeLeaves)
result.insert(value);
continue;
}
workingSet.insert(a.opOperand->get());
if (visitedOpOperands)
visitedOpOperands->insert(a.opOperand);
}
}
return result;
}
// Find the values that define the contents of the given operand's value.
llvm::SetVector<Value>
AnalysisState::findDefinitions(OpOperand *opOperand) const {
TraversalConfig config;
config.alwaysIncludeLeaves = false;
return findValueInReverseUseDefChain(
opOperand, [&](Value v) { return this->bufferizesToMemoryWrite(v); },
config);
}
AnalysisState::AnalysisState(const BufferizationOptions &options)
: AnalysisState(options, TypeID::get<AnalysisState>()) {}
AnalysisState::AnalysisState(const BufferizationOptions &options, TypeID type)
: options(options), type(type) {
for (const BufferizationOptions::AnalysisStateInitFn &fn :
options.stateInitializers)
fn(*this);
}
bool AnalysisState::canOmitTensorCopy(OpOperand &opOperand) const {
// Do not copy if the tensor has undefined contents.
if (hasUndefinedContents(&opOperand))
return true;
// Do not copy if the buffer of the tensor is entirely overwritten (with
// values that do not depend on the old tensor).
if (bufferizesToMemoryWrite(opOperand) && !bufferizesToMemoryRead(opOperand))
return true;
// Do not copy if the tensor is never read.
AliasingValueList aliases = getAliasingValues(opOperand);
if (!bufferizesToMemoryRead(opOperand) &&
llvm::none_of(aliases,
[&](AliasingValue a) { return isValueRead(a.value); }))
return true;
// Default: Cannot omit the copy.
return false;
}
bool AnalysisState::isInPlace(OpOperand &opOperand) const {
// ToBufferOps are always in-place.
if (isa<ToBufferOp>(opOperand.getOwner()))
return true;
// In the absence of analysis information, OpOperands that bufferize to a
// memory write are out-of-place, i.e., an alloc and copy is inserted.
return !bufferizesToMemoryWrite(opOperand);
}
bool AnalysisState::areEquivalentBufferizedValues(Value v1, Value v2) const {
// In the absence of analysis information, we do not know if the values are
// equivalent. The conservative answer is "false".
return false;
}
bool AnalysisState::areAliasingBufferizedValues(Value v1, Value v2) const {
// In the absence of analysis information, we do not know if the values may be
// aliasing. The conservative answer is "true".
return true;
}
bool AnalysisState::hasUndefinedContents(OpOperand *opOperand) const {
// In the absence of analysis information, the conservative answer is "false".
return false;
}
FailureOr<Value> bufferization::getBuffer(RewriterBase &rewriter, Value value,
const BufferizationOptions &options,
const BufferizationState &state) {
#ifndef NDEBUG
auto tensorType = llvm::dyn_cast<TensorLikeType>(value.getType());
assert(tensorType && "unexpected non-tensor type");
#endif // NDEBUG
// Replace "%t = to_tensor %m" with %m.
if (auto toTensorOp = value.getDefiningOp<bufferization::ToTensorOp>())
return toTensorOp.getBuffer();
// Insert to_buffer op.
OpBuilder::InsertionGuard g(rewriter);
setInsertionPointAfter(rewriter, value);
FailureOr<BufferLikeType> bufferType = getBufferType(value, options, state);
if (failed(bufferType))
return failure();
return bufferization::ToBufferOp::create(rewriter, value.getLoc(),
*bufferType, value)
.getResult();
}
/// Return the buffer type for a given Value (tensor) after bufferization.
FailureOr<BufferLikeType>
bufferization::getBufferType(Value value, const BufferizationOptions &options,
const BufferizationState &state) {
SmallVector<Value> invocationStack;
return getBufferType(value, options, state, invocationStack);
}
/// Return the buffer type for a given Value (tensor) after bufferization.
FailureOr<BufferLikeType>
bufferization::getBufferType(Value value, const BufferizationOptions &options,
const BufferizationState &state,
SmallVector<Value> &invocationStack) {
assert(llvm::isa<TensorLikeType>(value.getType()) &&
"unexpected non-tensor type");
invocationStack.push_back(value);
auto popFromStack =
llvm::make_scope_exit([&]() { invocationStack.pop_back(); });
// Try querying BufferizableOpInterface.
Operation *op = getOwnerOfValue(value);
auto bufferizableOp = options.dynCastBufferizableOp(op);
if (bufferizableOp)
return bufferizableOp.getBufferType(value, options, state, invocationStack);
// Op is not bufferizable.
return cast<TensorLikeType>(value.getType()).getBufferType(options, [&]() {
return op->emitError();
});
}
bool bufferization::hasTensorSemantics(Operation *op) {
if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
return bufferizableOp.hasTensorSemantics();
return detail::defaultHasTensorSemantics(op);
}
void bufferization::replaceOpWithBufferizedValues(RewriterBase &rewriter,
Operation *op,
ValueRange values) {
assert(values.size() == op->getNumResults() &&
"expected one value per OpResult");
OpBuilder::InsertionGuard g(rewriter);
// Replace all OpResults with the given values.
SmallVector<Value> replacements;
for (OpResult opResult : op->getOpResults()) {
Value replacement = values[opResult.getResultNumber()];
if (llvm::isa<TensorLikeType>(opResult.getType())) {
// The OpResult is a tensor. Such values are replaced with memrefs during
// bufferization.
assert(llvm::isa<BufferLikeType>(replacement.getType()) &&
"tensor op result should be replaced with a buffer value");
// The existing uses of the OpResult still expect a tensor. Insert a
// ToTensorOp. Throughout bufferization, this ToTensorOp will gradually
// loose all of its users and eventually DCE away.
rewriter.setInsertionPointAfter(op);
replacement = bufferization::ToTensorOp::create(
rewriter, replacement.getLoc(), opResult.getType(), replacement);
}
replacements.push_back(replacement);
}
rewriter.replaceOp(op, replacements);
}
//===----------------------------------------------------------------------===//
// Bufferization-specific scoped alloc insertion support.
//===----------------------------------------------------------------------===//
/// Create a memref allocation with the given type and dynamic extents.
FailureOr<Value> BufferizationOptions::createAlloc(OpBuilder &b, Location loc,
MemRefType type,
ValueRange dynShape) const {
if (allocationFn)
return (*allocationFn)(b, loc, type, dynShape, bufferAlignment);
// Default bufferallocation via AllocOp.
if (bufferAlignment != 0)
return memref::AllocOp::create(b, loc, type, dynShape,
b.getI64IntegerAttr(bufferAlignment))
.getResult();
return memref::AllocOp::create(b, loc, type, dynShape).getResult();
}
/// Create a memory copy between two memref buffers.
LogicalResult BufferizationOptions::createMemCpy(OpBuilder &b, Location loc,
Value from, Value to) const {
if (memCpyFn)
return (*memCpyFn)(b, loc, from, to);
memref::CopyOp::create(b, loc, from, to);
return success();
}
//===----------------------------------------------------------------------===//
// Bufferization-specific IRMapping support with debugging.
//===----------------------------------------------------------------------===//
BaseMemRefType bufferization::getMemRefType(TensorType tensorType,
const BufferizationOptions &options,
MemRefLayoutAttrInterface layout,
Attribute memorySpace) {
// Case 1: Unranked memref type.
if (auto unrankedTensorType =
llvm::dyn_cast<UnrankedTensorType>(tensorType)) {
assert(!layout && "UnrankedTensorType cannot have a layout map");
return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
memorySpace);
}
// Case 2: Ranked memref type with specified layout.
auto rankedTensorType = llvm::cast<RankedTensorType>(tensorType);
if (layout) {
return MemRefType::get(rankedTensorType.getShape(),
rankedTensorType.getElementType(), layout,
memorySpace);
}
return options.unknownTypeConverterFn(tensorType, memorySpace, options);
}
BaseMemRefType
bufferization::getMemRefTypeWithFullyDynamicLayout(TensorType tensorType,
Attribute memorySpace) {
// Case 1: Unranked memref type.
if (auto unrankedTensorType =
llvm::dyn_cast<UnrankedTensorType>(tensorType)) {
return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
memorySpace);
}
// Case 2: Ranked memref type.
auto rankedTensorType = llvm::cast<RankedTensorType>(tensorType);
int64_t dynamicOffset = ShapedType::kDynamic;
SmallVector<int64_t> dynamicStrides(rankedTensorType.getRank(),
ShapedType::kDynamic);
auto stridedLayout = StridedLayoutAttr::get(tensorType.getContext(),
dynamicOffset, dynamicStrides);
return MemRefType::get(rankedTensorType.getShape(),
rankedTensorType.getElementType(), stridedLayout,
memorySpace);
}
/// Return a MemRef type with a static identity layout (i.e., no layout map). If
/// the given tensor type is unranked, return an unranked MemRef type.
BaseMemRefType
bufferization::getMemRefTypeWithStaticIdentityLayout(TensorType tensorType,
Attribute memorySpace) {
// Case 1: Unranked memref type.
if (auto unrankedTensorType =
llvm::dyn_cast<UnrankedTensorType>(tensorType)) {
return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
memorySpace);
}
// Case 2: Ranked memref type.
auto rankedTensorType = llvm::cast<RankedTensorType>(tensorType);
MemRefLayoutAttrInterface layout = {};
return MemRefType::get(rankedTensorType.getShape(),
rankedTensorType.getElementType(), layout,
memorySpace);
}
//===----------------------------------------------------------------------===//
// Default implementations of interface methods
//===----------------------------------------------------------------------===//
bool bufferization::detail::defaultResultBufferizesToMemoryWrite(
OpResult opResult, const AnalysisState &state) {
auto bufferizableOp = cast<BufferizableOpInterface>(opResult.getDefiningOp());
AliasingOpOperandList opOperands =
bufferizableOp.getAliasingOpOperands(opResult, state);
// Case 1: OpResults that have no aliasing OpOperand usually bufferize to
// memory writes.
if (opOperands.getAliases().empty())
return true;
// Case 2: If an aliasing OpOperand bufferizes to a memory write, the OpResult
// may bufferize to a memory write.
if (llvm::any_of(opOperands, [&](AliasingOpOperand alias) {
return state.bufferizesToMemoryWrite(*alias.opOperand);
}))
return true;
// Case 3: Check if a nested aliasing OpOperand value bufferizes to a memory
// write. (Or: The reverse SSA use-def chain ends inside the reigon.) In that
// case, the OpResult bufferizes to a memory write. E.g.:
//
// %0 = "some_writing_op" : tensor<?xf32>
// %r = scf.if ... -> tensor<?xf32> {
// scf.yield %0 : tensor<?xf32>
// } else {
// %1 = "another_writing_op"(%0) : tensor<?xf32>
// scf.yield %1 : tensor<?xf32>
// }
// "some_reading_op"(%r)
//
// %r bufferizes to a memory write because an aliasing OpOperand value (%1)
// bufferizes to a memory write and the defining op is inside the scf.if.
//
// Note: This treatment of surrouding ops is useful for ops that have a
// region but no OpOperand such as scf.if or scf.execute_region. It simplifies
// the analysis considerably.
//
// "another_writing_op" in the above example should be able to bufferize
// inplace in the absence of another read of %0. However, if the scf.if op
// would not be considered a "write", the analysis would detect the
// following conflict:
//
// * read = some_reading_op
// * lastWrite = %0 (Note: The last write of %r would be a set: {%0, %1}.)
// * conflictingWrite = %1
//
auto isMemoryWriteInsideOp = [&](Value v) {
Operation *op = getOwnerOfValue(v);
if (!opResult.getDefiningOp()->isAncestor(op))
return false;
return state.bufferizesToMemoryWrite(v);
};
TraversalConfig config;
config.alwaysIncludeLeaves = false;
for (AliasingOpOperand alias : opOperands) {
if (!state
.findValueInReverseUseDefChain(alias.opOperand,
isMemoryWriteInsideOp, config)
.empty())
return true;
}
return false;
}
// Compute the AliasingOpOperandList for a given Value based on
// getAliasingValues.
AliasingOpOperandList bufferization::detail::defaultGetAliasingOpOperands(
Value value, const AnalysisState &state) {
Operation *op = getOwnerOfValue(value);
SmallVector<AliasingOpOperand> result;
for (OpOperand &opOperand : op->getOpOperands()) {
if (!llvm::isa<TensorType>(opOperand.get().getType()))
continue;
AliasingValueList aliasingValues = state.getAliasingValues(opOperand);
for (const auto &it : aliasingValues)
if (it.value == value)
result.emplace_back(&opOperand, it.relation, it.isDefinite);
}
return AliasingOpOperandList(std::move(result));
}
FailureOr<BufferLikeType> bufferization::detail::defaultGetBufferType(
Value value, const BufferizationOptions &options,
const BufferizationState &bufferizationState,
SmallVector<Value> &invocationStack) {
assert(llvm::isa<TensorType>(value.getType()) && "expected tensor type");
auto tensorType = cast<TensorType>(value.getType());
// No further analysis is possible for a block argument.
if (llvm::isa<BlockArgument>(value)) {
return cast<BufferLikeType>(
bufferization::getMemRefType(tensorType, options));
}
// Value is an OpResult.
Operation *op = getOwnerOfValue(value);
auto opResult = llvm::cast<OpResult>(value);
AnalysisState analysisState(options);
AliasingOpOperandList aliases = analysisState.getAliasingOpOperands(opResult);
if (aliases.getNumAliases() > 0 &&
aliases.getAliases()[0].relation == BufferRelation::Equivalent) {
// If the OpResult has an equivalent OpOperand, both OpResult and
// OpOperand bufferize to the exact same buffer type.
Value equivalentOperand = aliases.getAliases().front().opOperand->get();
return getBufferType(equivalentOperand, options, bufferizationState,
invocationStack);
}
// If we do not know the memory space and there is no default memory space,
// report a failure.
auto memSpace =
options.defaultMemorySpaceFn(cast<TensorType>(value.getType()));
if (!memSpace.has_value())
return op->emitError("could not infer memory space");
return cast<BufferLikeType>(
getMemRefType(tensorType, options, /*layout=*/{}, *memSpace));
}
bool bufferization::detail::defaultIsRepetitiveRegion(
BufferizableOpInterface bufferizableOp, unsigned index) {
assert(index < bufferizableOp->getNumRegions() && "invalid region index");
auto regionInterface =
dyn_cast<RegionBranchOpInterface>(bufferizableOp.getOperation());
if (!regionInterface)
return false;
return regionInterface.isRepetitiveRegion(index);
}
AliasingOpOperandList
bufferization::detail::unknownGetAliasingOpOperands(Value value) {
// TODO: Take into account successor blocks.
// No aliasing in case of non-entry blocks.
if (auto bbArg = dyn_cast<BlockArgument>(value))
if (bbArg.getOwner() != &bbArg.getOwner()->getParent()->getBlocks().front())
return {};
// Unknown op: Conservatively assume that each OpResult may alias with every
// OpOperand. In addition, each block argument of an entry block may alias
// with every OpOperand.
AliasingOpOperandList r;
for (OpOperand &operand : value.getDefiningOp()->getOpOperands())
if (isa<TensorType>(operand.get().getType()))
r.addAlias({&operand, BufferRelation::Unknown, /*isDefinite=*/false});
return r;
}
AliasingValueList
bufferization::detail::unknownGetAliasingValues(OpOperand &opOperand) {
// TODO: Take into account successor blocks.
// Unknown op: Conservatively assume that each OpResult may alias with every
// OpOperand. In addition, each block argument of an entry block may alias
// with every OpOperand.
AliasingValueList r;
for (OpResult result : opOperand.getOwner()->getOpResults())
if (llvm::isa<TensorType>(result.getType()))
r.addAlias({result, BufferRelation::Unknown, /*isDefinite=*/false});
for (Region &region : opOperand.getOwner()->getRegions())
if (!region.getBlocks().empty())
for (BlockArgument bbArg : region.getBlocks().front().getArguments())
if (isa<TensorType>(bbArg.getType()))
r.addAlias({bbArg, BufferRelation::Unknown, /*isDefinite=*/false});
return r;
}
bool bufferization::detail::defaultHasTensorSemantics(Operation *op) {
auto isaTensor = [](Type t) { return isa<TensorLikeType>(t); };
bool hasTensorBlockArgument = any_of(op->getRegions(), [&](Region &r) {
return any_of(r.getBlocks(), [&](Block &b) {
return any_of(b.getArguments(), [&](BlockArgument bbArg) {
return isaTensor(bbArg.getType());
});
});
});
if (hasTensorBlockArgument)
return true;
if (any_of(op->getResultTypes(), isaTensor))
return true;
return any_of(op->getOperandTypes(), isaTensor);
}
FailureOr<BaseMemRefType>
bufferization::detail::asMemRefType(FailureOr<BufferLikeType> bufferType) {
if (failed(bufferType))
return failure();
return cast<BaseMemRefType>(*bufferType);
}
bool bufferization::detail::typesMatchAfterBufferization(Operation &op,
Value tensor,
Value buffer) {
return mlir::succeeded(
cast<TensorLikeType>(tensor.getType())
.verifyCompatibleBufferType(cast<BufferLikeType>(buffer.getType()),
[&]() { return op.emitError(); }));
}