25 #include "llvm/ADT/STLExtras.h"
35 OperandRange::iterator &elementIt,
37 if (dim ==
static_cast<int>(shape.size()) - 1) {
38 for (
int i = 0; i < shape.back(); ++i) {
39 indices.back() = constants[i];
40 destination = rewriter.
create<tensor::InsertOp>(loc, *elementIt,
41 destination, indices);
46 for (
int i = 0; i < shape[dim]; ++i) {
47 indices[dim] = constants[i];
48 destination =
createInserts(rewriter, loc, dim + 1, destination, shape,
49 constants, elementIt, indices);
59 auto tensorType = dyn_cast<RankedTensorType>(tensorSource.
getType());
60 assert(tensorType &&
"expected ranked tensor");
61 assert(isa<MemRefType>(memrefDest.
getType()) &&
"expected ranked memref");
68 auto materializeOp = b.
create<bufferization::MaterializeInDestinationOp>(
69 loc, tensorSource, memrefDest);
70 materializeOp.setWritable(
true);
76 Value toBuffer = b.
create<bufferization::ToBufferOp>(
79 b.
create<memref::CopyOp>(loc, toBuffer, memrefDest);
85 Value toBuffer = b.
create<bufferization::ToBufferOp>(
88 b.
create<linalg::CopyOp>(loc, toBuffer, memrefDest);
97 RankedTensorType resultType = padOp.getResultType();
102 cast<tensor::YieldOp>(padOp.getBody()->getTerminator()).getValue();
106 isa<BlockArgument>(yieldedValue) &&
107 cast<BlockArgument>(yieldedValue).getOwner()->getParentOp() !=
108 padOp.getOperation();
110 bool outsideOpResult =
111 isa<OpResult>(yieldedValue) &&
113 bool invariantYieldedValue = outsideBbArg || outsideOpResult;
128 if (invariantYieldedValue) {
130 auto fillOp = rewriter.
create<linalg::FillOp>(loc,
ValueRange(yieldedValue),
137 utils::IteratorType::parallel);
140 auto genericOp = rewriter.
create<linalg::GenericOp>(
143 indexingMaps, iteratorTypes);
145 resultType.getElementType(), loc);
148 for (int64_t i = 0; i < resultType.getRank(); ++i)
149 bbArgReplacements.push_back(rewriter.
create<linalg::IndexOp>(loc, i));
150 rewriter.
mergeBlocks(padOp.getBody(), body, bbArgReplacements);
153 auto yieldOp = cast<tensor::YieldOp>(body->getTerminator());
160 auto tensorType = cast<RankedTensorType>(value.
getType());
161 if (tensorType.hasStaticShape())
166 if (isa<OpResult>(value) &&
169 for (int64_t i = 0; i < tensorType.getRank(); ++i) {
170 if (tensorType.isDynamicDim(i))
171 dynSizes.push_back(cast<Value>(
172 reifiedShape[cast<OpResult>(value).getResultNumber()][i]));
179 for (int64_t i = 0; i < tensorType.getRank(); ++i) {
180 if (tensorType.isDynamicDim(i))
193 auto tensorType = cast<RankedTensorType>(value.
getType());
198 tensorType, memorySpace));
204 alloc = rewriter.
create<memref::AllocOp>(loc, memrefType, dynamicSizes);
208 rewriter.
create<memref::DeallocOp>(loc, alloc);
212 alloc = rewriter.
create<memref::AllocaOp>(loc, memrefType, dynamicSizes);
223 assert(!
options.bufferizeDestinationOnly &&
"invalid options");
234 if (!padOp.hasZeroLowPad() || !padOp.hasZeroHighPad()) {
246 Value subview = rewriter.
create<memref::SubViewOp>(
247 loc, alloc, padOp.getMixedLowPad(), sizes, strides);
252 Value toTensorOp = rewriter.
create<bufferization::ToTensorOp>(
253 loc, padOp.getResult().getType(), alloc,
true,
262 assert(llvm::range_size(maskOp.getMaskBlock()->without_terminator()) == 1 &&
263 "expected single masked op");
270 Operation *yieldOp = maskOp.getMaskRegion().front().getTerminator();
271 assert(isa<vector::YieldOp>(yieldOp) &&
"expected yield op terminator");
276 rewriter,
options, maskOp.getMaskableOp(), memorySpace,
277 insertionPoint ? insertionPoint : maskOp);
279 if (
options.bufferizeDestinationOnly)
284 if (failed(cast<bufferization::BufferizableOpInterface>(yieldOp).bufferize(
285 rewriter, bufferizationOptions, bufferizationState)))
292 maskOp.walk([&](bufferization::ToTensorOp toTensorOp) {
293 if (toTensorOp->getUses().empty())
294 toTensorOps.push_back(toTensorOp.getOperation());
301 for (
Value result : maskOp.getResults())
302 if (isa<TensorType>(result.getType()))
304 resultUses.push_back(&use);
307 cast<bufferization::BufferizableOpInterface>(maskOp.getOperation())
308 .bufferize(rewriter, bufferizationOptions, bufferizationState)))
313 for (
OpOperand *resultUse : resultUses) {
315 resultUse->get().getDefiningOp<bufferization::ToTensorOp>();
316 assert(toTensorOp &&
"expected to_tensor op");
318 toTensorOp.setRestrict(
true);
319 toTensorOp.setWritable(
true);
328 bufferization::AllocTensorOp allocTensorOp,
Attribute memorySpace,
330 Location loc = allocTensorOp.getLoc();
337 rewriter, loc, allocTensorOp.getResult(),
options, memorySpace);
341 Value toTensorOp = rewriter.
create<bufferization::ToTensorOp>(
342 loc, allocTensorOp.getResult().getType(), alloc,
true,
344 rewriter.
replaceOp(allocTensorOp, toTensorOp);
350 RewriterBase &rewriter, tensor::FromElementsOp fromElementsOp) {
351 Location loc = fromElementsOp.getLoc();
352 RankedTensorType tensorType =
353 cast<RankedTensorType>(fromElementsOp.getType());
354 auto shape = tensorType.getShape();
362 fromElementsOp, fromElementsOp.getElements().front(),
368 auto maxDim = *llvm::max_element(shape);
370 constants.reserve(maxDim);
371 for (
int i = 0; i < maxDim; ++i)
375 auto elementIt = fromElementsOp.getElements().begin();
378 shape, constants, elementIt, indices);
381 rewriter.
replaceOp(fromElementsOp, result);
386 FailureOr<Operation *>
388 tensor::GenerateOp generateOp) {
390 if (!generateOp.getBody().hasOneBlock())
394 RankedTensorType tensorType = cast<RankedTensorType>(generateOp.getType());
398 rewriter.
create<EmptyOp>(loc, tensorType, generateOp.getDynamicExtents());
402 utils::IteratorType::parallel);
405 auto genericOp = rewriter.
create<linalg::GenericOp>(
408 indexingMaps, iteratorTypes);
410 tensorType.getElementType(), loc);
413 for (int64_t i = 0; i < tensorType.getRank(); ++i)
414 bbArgReplacements.push_back(rewriter.
create<linalg::IndexOp>(loc, i));
415 rewriter.
mergeBlocks(&generateOp.getBody().front(), body, bbArgReplacements);
418 auto yieldOp = cast<tensor::YieldOp>(body->getTerminator());
422 rewriter.
replaceOp(generateOp, genericOp->getResult(0));
423 return genericOp.getOperation();
427 FailureOr<Operation *>
429 tensor::PadOp padOp) {
431 if (!padOp.getBodyRegion().hasOneBlock())
436 RankedTensorType resultType = padOp.getResultType();
440 padOp,
"failed to reify tensor.pad op result shape");
442 for (int64_t i = 0; i < resultType.getRank(); ++i)
443 if (resultType.isDynamicDim(i))
444 dynamicSizes.push_back(cast<Value>(reifiedShape[0][i]));
448 if (padOp.getNofoldAttr() &&
451 using bufferization::AllocTensorOp;
453 rewriter.
create<AllocTensorOp>(loc, resultType, dynamicSizes);
455 padOp, padOp.getSource(), allocated);
456 return copyOp.getOperation();
459 Value empty = rewriter.
create<EmptyOp>(loc, resultType, dynamicSizes);
470 padOp, padOp.getSource(), fillOp->getResult(0),
471 padOp.getMixedLowPad(), sliceSizes, sliceStrides);
472 return insertSliceOp.getOperation();
478 using namespace bufferization;
481 if (
auto padOp = dyn_cast<tensor::PadOp>(op))
483 if (
auto maskOp = dyn_cast<vector::MaskOp>(op))
485 if (
auto allocTensorOp = dyn_cast<bufferization::AllocTensorOp>(op))
489 auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
494 BufferizationOptions bufferizationOptions;
496 BufferizationState bufferizationState;
499 if (!
options.bufferizeDestinationOnly) {
506 [](
Value v) { return isa<TensorType>(v.getType()); }))
507 llvm_unreachable(
"ops with nested tensor ops are not supported yet");
509 [](
Value v) { return isa<TensorType>(v.getType()); }))
510 llvm_unreachable(
"ops with nested tensor ops are not supported yet");
518 if (!isa<TensorType>(result.getType()))
521 if (!isa<RankedTensorType>(result.getType()))
524 if (bufferizableOp.bufferizesToAllocation(result))
526 tensorResults.push_back(result);
532 auto addOutOfPlaceOperand = [&](
OpOperand *operand) {
533 if (!llvm::is_contained(outOfPlaceOperands, operand))
534 outOfPlaceOperands.push_back(operand);
536 for (
OpResult result : tensorResults) {
538 analysisState.getAliasingOpOperands(result);
539 for (
const AliasingOpOperand &operand : aliasingOperands) {
540 addOutOfPlaceOperand(operand.opOperand);
541 for (
OpOperand &resultUse : result.getUses())
542 resultUses.push_back(&resultUse);
546 if (!analysisState.bufferizesToMemoryWrite(operand))
548 if (!isa<RankedTensorType>(operand.get().getType()))
550 addOutOfPlaceOperand(&operand);
553 if (outOfPlaceOperands.size() != 1)
560 for (
OpOperand *operand : outOfPlaceOperands) {
562 rewriter, op->
getLoc(), operand->get(),
options, memorySpace);
563 allocs.push_back(alloc);
564 if (!analysisState.findDefinitions(operand).empty()) {
570 auto toTensorOp = rewriter.
create<ToTensorOp>(
571 op->
getLoc(), operand->get().getType(), alloc);
572 operand->set(toTensorOp);
573 if (
options.bufferizeDestinationOnly) {
575 toTensorOp.setRestrict(
true);
576 toTensorOp.setWritable(
true);
582 if (
options.bufferizeDestinationOnly)
583 return allocs.front();
587 if (failed(bufferizableOp.bufferize(rewriter, bufferizationOptions,
588 bufferizationState)))
593 for (
OpOperand *resultUse : resultUses) {
594 auto toTensorOp = resultUse->get().
getDefiningOp<ToTensorOp>();
595 assert(toTensorOp &&
"expected to_tensor op");
597 toTensorOp.setRestrict(
true);
598 toTensorOp.setWritable(
true);
601 return allocs.front();
606 template <
typename OpTy>
607 LogicalResult rewriteOpInDestinationPassingStyle(OpTy op,
616 patterns.add(rewriteOpInDestinationPassingStyle<tensor::FromElementsOp>);
617 patterns.add(rewriteOpInDestinationPassingStyle<tensor::GenerateOp>);
618 patterns.add(rewriteOpInDestinationPassingStyle<tensor::PadOp>);
static Operation * movePaddingToFillOrGenericOp(RewriterBase &rewriter, Location loc, PadOp padOp, Value dest)
static Value createAllocationForTensor(RewriterBase &rewriter, Location loc, Value value, const linalg::BufferizeToAllocationOptions &options, Attribute memorySpace={})
static void createMemcpy(OpBuilder &b, Location loc, Value tensorSource, Value memrefDest, const linalg::BufferizeToAllocationOptions &options)
Create a memcpy from the given source tensor to the given destination memref.
static SmallVector< Value > reifyOrComputeDynamicSizes(OpBuilder &b, Value value)
static Value createInserts(RewriterBase &rewriter, Location loc, int dim, Value destination, ArrayRef< int64_t > shape, ArrayRef< Value > constants, OperandRange::iterator &elementIt, SmallVectorImpl< Value > &indices)
static llvm::ManagedStatic< PassManagerOptions > options
Base class for generic analysis states.
Attributes are known-constant values of operations.
Block represents an ordered list of Operations.
Operation * getTerminator()
Get the terminator operation of this block.
IntegerAttr getIndexAttr(int64_t value)
AffineMap getMultiDimIdentityMap(unsigned rank)
MLIRContext * getContext() const
Dialects are groups of MLIR operations, types and attributes, as well as behavior associated with the...
virtual Operation * materializeConstant(OpBuilder &builder, Attribute value, Type type, Location loc)
Registered hook to materialize a single constant operation from a given attribute value with the desi...
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Dialect * getLoadedDialect(StringRef name)
Get a registered IR dialect with the given namespace.
RAII guard to reset the insertion point of the builder when destroyed.
This class helps build Operations.
Block * createBlock(Region *parent, Region::iterator insertPt={}, TypeRange argTypes={}, ArrayRef< Location > locs={})
Add new block with 'argTypes' arguments and set the insertion point to the end of it.
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...
Block * getInsertionBlock() const
Return the block the current insertion point belongs to.
This class represents an operand of an operation.
This is a value defined by a result of an operation.
Operation is the basic unit of execution within MLIR.
std::enable_if_t< llvm::function_traits< std::decay_t< FnT > >::num_args==1, RetT > walk(FnT &&callback)
Walk the operation by calling the callback for each nested operation (including this one),...
Location getLoc()
The source location the operation was defined or derived from.
Operation * getParentOp()
Returns the closest surrounding operation that contains this operation or nullptr if this is a top-le...
MutableArrayRef< OpOperand > getOpOperands()
operand_range getOperands()
Returns an iterator on the underlying Value's.
result_range getResults()
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the listener that the IR failed to be rewritten because of a match failure,...
virtual void replaceOp(Operation *op, ValueRange newValues)
Replace the results of the given (original) operation with the specified list of values (replacements...
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
void mergeBlocks(Block *source, Block *dest, ValueRange argValues={})
Inline the operations of block 'source' into the end of block 'dest'.
void modifyOpInPlace(Operation *root, CallableT &&callable)
This method is a utility wrapper around an in-place modification of an operation.
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
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.
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.
Specialization of arith.constant op that returns an integer of index type.
BufferizationState provides information about the state of the IR during the bufferization process.
BaseMemRefType getMemRefTypeWithStaticIdentityLayout(TensorType tensorType, Attribute memorySpace=nullptr)
Return a MemRef type with a static identity layout (i.e., no layout map).
AliasList< AliasingOpOperand > AliasingOpOperandList
A list of possible aliasing OpOperands.
BaseMemRefType getMemRefTypeWithFullyDynamicLayout(TensorType tensorType, Attribute memorySpace=nullptr)
Return a MemRef type with fully dynamic layout.
Value bufferizeToAllocation(RewriterBase &rewriter, const BufferizeToAllocationOptions &options, tensor::PadOp padOp, Attribute memorySpace={}, Operation *insertionPoint=nullptr)
Materialize a buffer allocation for the given tensor.pad op and lower the op to linalg....
FailureOr< Operation * > rewriteInDestinationPassingStyle(RewriterBase &rewriter, tensor::FromElementsOp fromElementsOp)
Rewrite tensor.from_elements to linalg.generic.
void populateConvertToDestinationStylePatterns(RewritePatternSet &patterns)
Populate patterns that convert non-destination-style ops to destination style ops.
SmallVector< OpFoldResult > getMixedSizes(OpBuilder &builder, Location loc, Value value)
Return the dimensions of the given tensor value.
Include the generated interface declarations.
bool matchPattern(Value value, const Pattern &pattern)
Entry point for matching a pattern over a Value.
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).
const FrozenRewritePatternSet & patterns
bool isZeroInteger(OpFoldResult v)
Return true if v is an IntegerAttr with value 0.
detail::constant_op_matcher m_Constant()
Matches a constant foldable operation.
Options for BufferizableOpInterface-based bufferization.
@ MaterializeInDestination