21 #include "llvm/Support/Debug.h"
28 namespace bufferization {
30 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.cpp.inc"
37 #define DEBUG_TYPE "bufferizable-op-interface"
38 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
39 #define LDBG(X) LLVM_DEBUG(DBGS() << (X))
42 using namespace bufferization;
47 if (
auto bufferizableOp =
options.dynCastBufferizableOp(op))
95 return opResult.getDefiningOp();
101 auto bufferizableOp = opResult.
getDefiningOp<BufferizableOpInterface>();
102 assert(bufferizableOp && bufferizableOp.bufferizesToAllocation(opResult) &&
103 "expected op that bufferizes to an allocation");
108 if (!op->
hasAttr(BufferizationDialect::kEscapeAttrName))
111 op->
getAttrOfType<ArrayAttr>(BufferizationDialect::kEscapeAttrName);
122 if (shapedValue.
getType().
isa<RankedTensorType>()) {
123 tensor = shapedValue;
124 }
else if (shapedValue.
getType().
isa<MemRefType>()) {
125 tensor = b.
create<ToTensorOp>(loc, shapedValue);
126 }
else if (shapedValue.
getType().
isa<UnrankedTensorType>() ||
129 ->
emitError(
"copying of unranked tensors is not implemented");
131 llvm_unreachable(
"expected RankedTensorType or MemRefType");
133 RankedTensorType tensorType = tensor.
getType().
cast<RankedTensorType>();
138 bool reifiedShapes =
false;
139 if (shapedValue.
getType().
isa<RankedTensorType>() &&
144 reifiedShapes =
true;
146 resultDims[shapedValue.
cast<
OpResult>().getResultNumber()];
147 for (
const auto &dim :
enumerate(tensorType.getShape()))
148 if (ShapedType::isDynamic(dim.value()))
149 dynamicSizes.push_back(shape[dim.index()].get<
Value>());
159 auto allocTensorOp = b.
create<AllocTensorOp>(loc, tensorType, dynamicSizes,
161 allocTensorOp->
setAttr(BufferizationDialect::kEscapeAttrName,
166 return allocTensorOp.getResult();
168 if (
failed(copyBufferType))
170 Attribute memorySpace = copyBufferType->getMemorySpace();
173 allocTensorOp.setMemorySpaceAttr(memorySpace);
174 return allocTensorOp.getResult();
177 LogicalResult BufferizableOpInterface::resolveTensorOpOperandConflicts(
190 Type operandType = opOperand.get().getType();
193 if (state.isInPlace(opOperand))
195 if (operandType.
isa<UnrankedTensorType>())
196 return op->
emitError(
"copying of unranked tensors is not implemented");
199 state.getAliasingOpResults(opOperand);
203 bool escape = !state.getOptions().createDeallocs ||
205 return state.isTensorYielded(a.
opResult);
209 !state.bufferizesToMemoryWrite(opOperand) &&
210 state.getAliasingOpOperands(aliasingOpResults.
getAliases()[0].opResult)
211 .getNumAliases() == 1 &&
214 .isa<UnrankedTensorType>()) {
222 outOfPlaceOpResults.push_back(opResult);
223 if (!state.canOmitTensorCopy(opOperand))
224 copiedOpResults.insert(opResult);
226 escapingOpResultCopies.insert(opResult);
229 outOfPlaceOpOperands.push_back(&opOperand);
230 if (!state.canOmitTensorCopy(opOperand))
231 copiedOpOperands.insert(&opOperand);
233 escapingOpOperandCopies.insert(&opOperand);
239 for (
OpOperand *opOperand : outOfPlaceOpOperands) {
241 rewriter, op->
getLoc(), opOperand->get(),
242 escapingOpOperandCopies.contains(opOperand), state.getOptions(),
243 copiedOpOperands.contains(opOperand));
251 for (
OpResult opResult : outOfPlaceOpResults) {
253 rewriter, op->
getLoc(), opResult,
254 escapingOpResultCopies.contains(opResult), state.getOptions(),
255 copiedOpResults.count(opResult));
262 if (use->getOwner() ==
copy->getDefiningOp())
266 if (isa<tensor::DimOp>(use->getOwner()))
278 assert(
options.dynCastBufferizableOp(op).bufferizesToAllocation(opResult) &&
279 "expected that op allocates");
282 if (op->
hasAttr(BufferizationDialect::kEscapeAttrName)) {
284 ArrayAttr escapeAttr =
285 op->
getAttr(BufferizationDialect::kEscapeAttrName).
cast<ArrayAttr>();
286 return !escapeAttr[0].cast<
BoolAttr>().getValue();
304 bool isAllowed = !hasAllowRule();
305 for (
const Entry &entry : entries) {
306 bool filterResult = entry.fn(op);
307 switch (entry.type) {
309 isAllowed |= filterResult;
340 bool isFuncBoundaryOp = isa_and_nonnull<func::FuncDialect>(op->
getDialect());
347 BufferizableOpInterface
349 auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
354 return bufferizableOp;
357 BufferizableOpInterface
359 if (
auto bufferizableOp = value.
getDefiningOp<BufferizableOpInterface>())
361 return bufferizableOp;
383 if (
auto bufferizableOp =
getOptions().dynCastBufferizableOp(op))
384 return bufferizableOp.getAliasingOpOperands(opResult, *
this);
394 if (
auto bufferizableOp =
396 return bufferizableOp.getAliasingOpResults(opOperand, *
this);
405 if (
auto bufferizableOp =
407 return bufferizableOp.bufferizesToMemoryRead(opOperand, *
this);
417 if (
auto bufferizableOp =
419 return bufferizableOp.bufferizesToMemoryWrite(opOperand, *
this);
429 if (
auto bufferizableOp =
431 return bufferizableOp.bufferizesToAliasOnly(opOperand, *
this);
445 return bufferizableOp.resultBufferizesToMemoryWrite(opResult, *
this);
455 workingSet.push_back(&use);
457 while (!workingSet.empty()) {
458 OpOperand *uMaybeReading = workingSet.pop_back_val();
462 for (
OpOperand &use : alias.opResult.getUses())
463 workingSet.push_back(&use);
477 bool followEquivalentOnly,
bool alwaysIncludeLeaves)
const {
479 workingSet.insert(value);
481 while (!workingSet.empty()) {
482 Value value = workingSet.pop_back_val();
483 if (condition(value)) {
484 result.insert(value);
489 if (alwaysIncludeLeaves)
490 result.insert(value);
495 BufferizableOpInterface bufferizableOp =
503 if (alwaysIncludeLeaves)
504 result.insert(value);
512 if (alwaysIncludeLeaves)
513 result.insert(value);
515 workingSet.insert(a.opOperand->get());
563 if (isa<ToMemrefOp>(opOperand.
getOwner()))
598 worklist.push_back(&use);
600 while (!worklist.empty()) {
601 OpOperand *operand = worklist.pop_back_val();
611 if (isa<ToMemrefOp>(op))
623 for (
OpOperand &use : alias.opResult.getUses())
624 worklist.push_back(&use);
635 assert((!rankedTensorType || memrefType.
cast<MemRefType>().getRank() ==
636 rankedTensorType.getRank()) &&
637 "to_memref would be invalid: mismatching ranks");
645 assert(tensorType &&
"unexpected non-tensor type");
649 if (
auto toTensorOp = value.
getDefiningOp<bufferization::ToTensorOp>())
650 return toTensorOp.getMemref();
660 .
create<bufferization::ToMemrefOp>(value.
getLoc(), *memrefType, value)
678 const auto &it = fixedTypes.find(value);
679 if (it != fixedTypes.end())
684 auto bufferizableOp =
options.dynCastBufferizableOp(op);
686 return bufferizableOp.getBufferType(value,
options, fixedTypes);
689 if (!
options.defaultMemorySpace.has_value())
690 return op->
emitError(
"could not infer memory space");
700 "expected one value per OpResult");
710 assert((replacement.
getType().
isa<MemRefType>() ||
712 "tensor op result should be replaced with a memref value");
717 replacement = rewriter.
create<bufferization::ToTensorOp>(
718 replacement.
getLoc(), replacement);
720 replacements.push_back(replacement);
740 .
create<memref::AllocOp>(loc, type, dynShape,
743 return b.
create<memref::AllocOp>(loc, type, dynShape).getResult();
749 Value allocatedBuffer)
const {
754 b.
create<memref::DeallocOp>(loc, allocatedBuffer);
762 return (*
memCpyFn)(b, loc, from, to);
764 b.
create<memref::CopyOp>(loc, from, to);
776 return isa<func::FuncOp>(bbArg.getOwner()->getParentOp());
781 MemRefLayoutAttrInterface layout,
786 if (
auto unrankedTensorType = tensorType.dyn_cast<UnrankedTensorType>()) {
787 assert(!layout &&
"UnrankedTensorType cannot have a layout map");
788 return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
793 auto rankedTensorType = tensorType.cast<RankedTensorType>();
795 return MemRefType::get(rankedTensorType.getShape(),
796 rankedTensorType.getElementType(), layout,
800 return options.unknownTypeConverterFn(value, memorySpace,
options);
807 if (
auto unrankedTensorType = tensorType.
dyn_cast<UnrankedTensorType>()) {
808 return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
813 auto rankedTensorType = tensorType.
cast<RankedTensorType>();
814 int64_t dynamicOffset = ShapedType::kDynamic;
816 ShapedType::kDynamic);
817 auto stridedLayout = StridedLayoutAttr::get(tensorType.
getContext(),
818 dynamicOffset, dynamicStrides);
819 return MemRefType::get(rankedTensorType.getShape(),
820 rankedTensorType.getElementType(), stridedLayout,
830 if (
auto unrankedTensorType = tensorType.
dyn_cast<UnrankedTensorType>()) {
831 return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
836 auto rankedTensorType = tensorType.
cast<RankedTensorType>();
837 MemRefLayoutAttrInterface layout = {};
838 return MemRefType::get(rankedTensorType.getShape(),
839 rankedTensorType.getElementType(), layout,
849 auto bufferizableOp = cast<BufferizableOpInterface>(opResult.
getDefiningOp());
851 bufferizableOp.getAliasingOpOperands(opResult, state);
894 auto isMemoryWriteInsideOp = [&](
Value v) {
903 isMemoryWriteInsideOp,
919 if (!opOperand.get().getType().isa<
TensorType>())
923 for (
const auto &it : aliasingOpResults)
924 if (it.opResult == opResult)
925 result.emplace_back(&opOperand, it.relation, it.isDefinite);
948 Value equivalentOperand = aliases.
getAliases().front().opOperand->get();
954 if (!
options.defaultMemorySpace.has_value())
955 return op->
emitError(
"could not infer memory space");
962 BufferizableOpInterface bufferizableOp,
unsigned index) {
963 assert(index < bufferizableOp->getNumRegions() &&
"invalid region index");
964 auto regionInterface =
965 dyn_cast<RegionBranchOpInterface>(bufferizableOp.getOperation());
966 if (!regionInterface)
968 return regionInterface.isRepetitiveRegion(index);
976 if (operand.get().getType().isa<
TensorType>())
static void ensureToMemrefOpIsValid(Value tensor, Type memrefType)
static void setInsertionPointAfter(OpBuilder &b, Value value)
static BaseMemRefType defaultUnknownTypeConverter(Value value, Attribute memorySpace, const BufferizationOptions &options)
Default unknown type converter: Use a fully dynamic layout map.
static bool isRepetitiveRegion(Region *region, const BufferizationOptions &options)
static void copy(Location loc, Value dst, Value src, Value size, OpBuilder &builder)
Copies the given number of bytes from src to dst pointers.
static llvm::ManagedStatic< PassManagerOptions > options
#define MLIR_DEFINE_EXPLICIT_TYPE_ID(CLASS_NAME)
Base class for generic analysis states.
Attributes are known-constant values of operations.
This class provides a shared interface for ranked and unranked memref types.
This class represents an argument of a Block.
Block represents an ordered list of Operations.
Region * getParent() const
Provide a 'getParent' method for ilist_node_with_parent methods.
Special case of IntegerAttr to represent boolean integers, i.e., signless i1 integers.
IntegerAttr getI64IntegerAttr(int64_t value)
ArrayAttr getBoolArrayAttr(ArrayRef< bool > values)
This class provides support for representing a failure result, or a valid value of type T.
IRValueT get() const
Return the current value being used by this operand.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
RAII guard to reset the insertion point of the builder when destroyed.
This class helps build Operations.
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
void setInsertionPointAfter(Operation *op)
Sets the insertion point to the node after the specified operation, which will cause subsequent inser...
This class represents an operand of an operation.
This is a value defined by a result of an operation.
Operation * getOwner() const
Returns the operation that owns this result.
unsigned getResultNumber() const
Returns the number of this result.
Operation is the basic unit of execution within MLIR.
Dialect * getDialect()
Return the dialect this operation is associated with, or nullptr if the associated dialect is not loa...
AttrClass getAttrOfType(StringAttr name)
Attribute getAttr(StringAttr name)
Return the specified attribute if present, null otherwise.
bool hasAttr(StringAttr name)
Return true if the operation has an attribute with the provided name, false otherwise.
Location getLoc()
The source location the operation was defined or derived from.
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
Block * getBlock()
Returns the operation block that contains this operation.
void setAttr(StringAttr name, Attribute value)
If the an attribute exists with the specified name, change it to the new value.
MutableArrayRef< OpOperand > getOpOperands()
bool isAncestor(Operation *other)
Return true if this operation is an ancestor of the other operation.
result_range getOpResults()
Region * getParentRegion()
Returns the region to which the instruction belongs.
unsigned getNumResults()
Return the number of results held by this operation.
This class contains a list of basic blocks and a link to the parent operation it is attached to.
Region * getParentRegion()
Return the region containing this region or nullptr if the region is attached to a top-level operatio...
unsigned getRegionNumber()
Return the number of this region in the parent operation.
Operation * getParentOp()
Return the parent operation this region is attached to.
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
virtual void replaceOp(Operation *op, ValueRange newValues)
This method replaces the results of the operation with the specified list of values.
void updateRootInPlace(Operation *root, CallableT &&callable)
This method is a utility wrapper around a root update of an operation.
Tensor types represent multi-dimensional arrays, and have two variants: RankedTensorType and Unranked...
This class provides an efficient unique identifier for a specific C++ type.
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
MLIRContext * getContext() const
Return the MLIRContext in which this type was uniqued.
This class provides an abstraction over the different types of ranges over Values.
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Type getType() const
Return the type of this value.
use_range getUses() const
Returns a range of all uses, which is useful for iterating over all uses.
Location getLoc() const
Return the location of this value.
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Region * getParentRegion()
Return the Region in which this Value is defined.
size_t getNumAliases() const
ArrayRef< T > getAliases() const
AnalysisState provides a variety of helper functions for dealing with tensor values.
bool isValueRead(Value value) const
Return true if the given value is read by an op that bufferizes to a memory read.
AliasingOpOperandList getAliasingOpOperands(OpResult result) const
Determine which OpOperand* will alias with result if the op is bufferized in place.
virtual bool areAliasingBufferizedValues(Value v1, Value v2) const
Return true if v1 and v2 may bufferize to aliasing buffers.
virtual bool hasUndefinedContents(OpOperand *opOperand) const
Return true if the given tensor has undefined contents.
bool canOmitTensorCopy(OpOperand &opOperand) const
Return true if a copy can always be avoided when allocating a new tensor for the given OpOperand.
bool bufferizesToMemoryWrite(OpOperand &opOperand) const
Return true if opOperand bufferizes to a memory write.
virtual bool isInPlace(OpOperand &opOperand) const
Return true if the given OpResult has been decided to bufferize inplace.
bool bufferizesToAliasOnly(OpOperand &opOperand) const
Return true if opOperand does neither read nor write but bufferizes to an alias.
SetVector< Value > findValueInReverseUseDefChain(Value value, llvm::function_ref< bool(Value)> condition, bool followEquivalentOnly=false, bool alwaysIncludeLeaves=true) const
Starting from value, follow the use-def chain in reverse, always selecting the aliasing OpOperands.
AnalysisState(const BufferizationOptions &options)
const BufferizationOptions & getOptions() const
Return a reference to the BufferizationOptions.
virtual bool isTensorYielded(Value tensor) const
Return true if the given tensor (or an aliasing tensor) is yielded from the containing block.
bool bufferizesToMemoryRead(OpOperand &opOperand) const
Return true if opOperand bufferizes to a memory read.
AliasingOpResultList getAliasingOpResults(OpOperand &opOperand) const
Determine which OpResult will alias with opOperand if the op is bufferized in place.
SetVector< Value > findDefinitions(Value value) const
Find the values that may define the contents of the given value at runtime.
virtual bool areEquivalentBufferizedValues(Value v1, Value v2) const
Return true if v1 and v2 bufferize to equivalent buffers.
bool isOpAllowed(Operation *op) const
Return whether the op is allowed or not.
Operation * getOwner() const
Return the owner of this operand.
FailureOr< BaseMemRefType > defaultGetBufferType(Value value, const BufferizationOptions &options, const DenseMap< Value, BaseMemRefType > &fixedTypes)
This is the default implementation of BufferizableOpInterface::getBufferType.
AliasingOpOperandList defaultGetAliasingOpOperands(OpResult opResult, const AnalysisState &state)
This is the default implementation of BufferizableOpInterface::getAliasingOpOperands.
bool defaultResultBufferizesToMemoryWrite(OpResult opResult, const AnalysisState &state)
This is the default implementation of BufferizableOpInterface::resultBufferizesToMemoryWrite.
bool defaultIsRepetitiveRegion(BufferizableOpInterface bufferizableOp, unsigned index)
This is the default implementation of BufferizableOpInterface::isRepetitiveRegion.
AliasingOpResultList unknownGetAliasingOpResults(OpOperand &opOperand)
This is the default implementation of getAliasingOpResults in case the owner op does not implement th...
AliasingOpOperandList unknownGetAliasingOpOperands(OpResult opResult)
This is the default implementation of getAliasingOpOperands in case the defining op does not implemen...
void replaceOpWithBufferizedValues(RewriterBase &rewriter, Operation *op, ValueRange values)
Replace an op with replacement values.
BaseMemRefType getMemRefTypeWithStaticIdentityLayout(TensorType tensorType, Attribute memorySpace=nullptr)
Return a MemRef type with a static identity layout (i.e., no layout map).
Operation * getOwnerOfValue(Value value)
Return the owner of the given value.
bool allocationDoesNotEscape(OpResult opResult)
Return true if the allocation of the given op is guaranteed to not escape the containing block.
BaseMemRefType getMemRefType(Value value, const BufferizationOptions &options, MemRefLayoutAttrInterface layout={}, Attribute memorySpace=nullptr)
Return a MemRefType to which the type of the given value can be bufferized.
bool shouldDeallocateOpResult(OpResult opResult, const BufferizationOptions &options)
Return true if the buffer of given OpResult should be deallocated.
Region * getEnclosingRepetitiveRegion(Operation *op, const BufferizationOptions &options)
Return the closest enclosing repetitive region around the given op.
FailureOr< Value > allocateTensorForShapedValue(OpBuilder &b, Location loc, Value shapedValue, bool escape, const BufferizationOptions &options, bool copy=true)
Create an AllocTensorOp for the given shaped value (memref or tensor).
Region * getNextEnclosingRepetitiveRegion(Region *region, const BufferizationOptions &options)
Assuming that the given region is repetitive, find the next enclosing repetitive region.
AliasList< AliasingOpOperand > AliasingOpOperandList
A list of possible aliasing OpOperands.
bool isFunctionArgument(Value value)
Return true if the given value is a BlockArgument of a func::FuncOp.
FailureOr< BaseMemRefType > getBufferType(Value value, const BufferizationOptions &options)
Return the buffer type for a given Value (tensor) after bufferization without bufferizing any IR.
FailureOr< Value > getBuffer(RewriterBase &rewriter, Value value, const BufferizationOptions &options)
Lookup the buffer for the given value.
BaseMemRefType getMemRefTypeWithFullyDynamicLayout(TensorType tensorType, Attribute memorySpace=nullptr)
Return a MemRef type with fully dynamic layout.
void populateDynamicDimSizes(OpBuilder &b, Location loc, Value shapedValue, SmallVector< Value > &dynamicDims)
Populate dynamicDims with tensor::DimOp / memref::DimOp results for all dynamic dimensions of the giv...
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
This header declares functions that assit transformations in the MemRef dialect.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
LogicalResult reifyResultShapes(OpBuilder &b, Operation *op, ReifiedRankedShapedTypeDims &reifiedReturnShapes)
Reify the shape of the result of an operation (typically in terms of the shape of its operands).
bool succeeded(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a success value.
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
bool isRegionReturnLike(Operation *operation)
Returns true if the given operation is either annotated with the ReturnLike trait or implements the R...
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value.
This class represents an efficient way to signal success or failure.
A maybe aliasing OpOperand.
A maybe aliasing OpResult.
Options for BufferizableOpInterface-based bufferization.
std::function< void(AnalysisState &)> AnalysisStateInitFn
Initializer function for analysis state.
BufferizableOpInterface dynCastBufferizableOp(Operation *op) const
Try to cast the given op to BufferizableOpInterface if the op is allow listed.
unsigned int bufferAlignment
Buffer alignment for new memory allocations.
std::optional< AllocationFn > allocationFn
Helper functions for allocation, deallocation, memory copying.
OpFilter opFilter
A filter that specifies which ops should be bufferized and which ops should be ignored.
bool isOpAllowed(Operation *op) const
Return true if the given op should be bufferized.
std::optional< DeallocationFn > deallocationFn
std::optional< MemCpyFn > memCpyFn
bool bufferizeFunctionBoundaries
Specifies whether function boundaries (ops in the func dialect) should be bufferized or not.
LogicalResult createDealloc(OpBuilder &b, Location loc, Value allocatedBuffer) const
Creates a memref deallocation.
FailureOr< Value > createAlloc(OpBuilder &b, Location loc, MemRefType type, ValueRange dynShape) const
Create a memref allocation with the given type and dynamic extents.
LogicalResult createMemCpy(OpBuilder &b, Location loc, Value from, Value to) const
Creates a memcpy between two given buffers.
SmallVector< AnalysisStateInitFn > stateInitializers
Initializer functions for analysis state.