26 #include "llvm/ADT/STLExtras.h"
27 #include "llvm/Support/Debug.h"
37 OperandRange::iterator &elementIt,
39 if (dim ==
static_cast<int>(shape.size()) - 1) {
40 for (
int i = 0; i < shape.back(); ++i) {
41 indices.back() = constants[i];
42 destination = rewriter.
create<tensor::InsertOp>(loc, *elementIt,
43 destination, indices);
48 for (
int i = 0; i < shape[dim]; ++i) {
49 indices[dim] = constants[i];
50 destination =
createInserts(rewriter, loc, dim + 1, destination, shape,
51 constants, elementIt, indices);
61 auto tensorType = dyn_cast<RankedTensorType>(tensorSource.
getType());
62 assert(tensorType &&
"expected ranked tensor");
63 assert(isa<MemRefType>(memrefDest.
getType()) &&
"expected ranked memref");
70 auto materializeOp = b.
create<bufferization::MaterializeInDestinationOp>(
71 loc, tensorSource, memrefDest);
72 materializeOp.setWritable(
true);
78 Value toBuffer = b.
create<bufferization::ToBufferOp>(
81 b.
create<memref::CopyOp>(loc, toBuffer, memrefDest);
87 Value toBuffer = b.
create<bufferization::ToBufferOp>(
90 b.
create<linalg::CopyOp>(loc, toBuffer, memrefDest);
99 RankedTensorType resultType = padOp.getResultType();
104 cast<tensor::YieldOp>(padOp.getBody()->getTerminator()).getValue();
108 isa<BlockArgument>(yieldedValue) &&
109 cast<BlockArgument>(yieldedValue).getOwner()->getParentOp() !=
110 padOp.getOperation();
112 bool outsideOpResult =
113 isa<OpResult>(yieldedValue) &&
115 bool invariantYieldedValue = outsideBbArg || outsideOpResult;
130 if (invariantYieldedValue) {
132 auto fillOp = rewriter.
create<linalg::FillOp>(loc,
ValueRange(yieldedValue),
139 utils::IteratorType::parallel);
142 auto genericOp = rewriter.
create<linalg::GenericOp>(
145 indexingMaps, iteratorTypes);
147 resultType.getElementType(), loc);
150 for (int64_t i = 0; i < resultType.getRank(); ++i)
151 bbArgReplacements.push_back(rewriter.
create<linalg::IndexOp>(loc, i));
152 rewriter.
mergeBlocks(padOp.getBody(), body, bbArgReplacements);
155 auto yieldOp = cast<tensor::YieldOp>(body->getTerminator());
162 auto tensorType = cast<RankedTensorType>(value.
getType());
163 if (tensorType.hasStaticShape())
168 if (isa<OpResult>(value) &&
171 for (int64_t i = 0; i < tensorType.getRank(); ++i) {
172 if (tensorType.isDynamicDim(i))
173 dynSizes.push_back(cast<Value>(
174 reifiedShape[cast<OpResult>(value).getResultNumber()][i]));
181 for (int64_t i = 0; i < tensorType.getRank(); ++i) {
182 if (tensorType.isDynamicDim(i))
195 auto tensorType = cast<RankedTensorType>(value.
getType());
200 tensorType, memorySpace));
206 alloc = rewriter.
create<memref::AllocOp>(loc, memrefType, dynamicSizes);
210 rewriter.
create<memref::DeallocOp>(loc, alloc);
214 alloc = rewriter.
create<memref::AllocaOp>(loc, memrefType, dynamicSizes);
225 assert(!
options.bufferizeDestinationOnly &&
"invalid options");
236 if (!padOp.hasZeroLowPad() || !padOp.hasZeroHighPad()) {
248 Value subview = rewriter.
create<memref::SubViewOp>(
249 loc, alloc, padOp.getMixedLowPad(), sizes, strides);
254 Value toTensorOp = rewriter.
create<bufferization::ToTensorOp>(
255 loc, alloc,
true,
true);
263 assert(llvm::range_size(maskOp.getMaskBlock()->without_terminator()) == 1 &&
264 "expected single masked op");
271 Operation *yieldOp = maskOp.getMaskRegion().front().getTerminator();
272 assert(isa<vector::YieldOp>(yieldOp) &&
"expected yield op terminator");
277 rewriter,
options, maskOp.getMaskableOp(), memorySpace,
278 insertionPoint ? insertionPoint : maskOp);
280 if (
options.bufferizeDestinationOnly)
285 if (failed(cast<bufferization::BufferizableOpInterface>(yieldOp).bufferize(
286 rewriter, bufferizationOptions, bufferizationState)))
293 maskOp.walk([&](bufferization::ToTensorOp toTensorOp) {
294 if (toTensorOp->getUses().empty())
295 toTensorOps.push_back(toTensorOp.getOperation());
302 for (
Value result : maskOp.getResults())
303 if (isa<TensorType>(result.getType()))
305 resultUses.push_back(&use);
308 cast<bufferization::BufferizableOpInterface>(maskOp.getOperation())
309 .bufferize(rewriter, bufferizationOptions, bufferizationState)))
314 for (
OpOperand *resultUse : resultUses) {
316 resultUse->get().getDefiningOp<bufferization::ToTensorOp>();
317 assert(toTensorOp &&
"expected to_tensor op");
319 toTensorOp.setRestrict(
true);
320 toTensorOp.setWritable(
true);
329 bufferization::AllocTensorOp allocTensorOp,
Attribute memorySpace,
331 Location loc = allocTensorOp.getLoc();
338 rewriter, loc, allocTensorOp.getResult(),
options, memorySpace);
342 Value toTensorOp = rewriter.
create<bufferization::ToTensorOp>(
343 loc, alloc,
true,
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>(op->
getLoc(), alloc);
571 operand->set(toTensorOp);
572 if (
options.bufferizeDestinationOnly) {
574 toTensorOp.setRestrict(
true);
575 toTensorOp.setWritable(
true);
581 if (
options.bufferizeDestinationOnly)
582 return allocs.front();
586 if (failed(bufferizableOp.bufferize(rewriter, bufferizationOptions,
587 bufferizationState)))
592 for (
OpOperand *resultUse : resultUses) {
593 auto toTensorOp = resultUse->get().
getDefiningOp<ToTensorOp>();
594 assert(toTensorOp &&
"expected to_tensor op");
596 toTensorOp.setRestrict(
true);
597 toTensorOp.setWritable(
true);
600 return allocs.front();
605 template <
typename OpTy>
606 LogicalResult rewriteOpInDestinationPassingStyle(OpTy op,
615 patterns.add(rewriteOpInDestinationPassingStyle<tensor::FromElementsOp>);
616 patterns.add(rewriteOpInDestinationPassingStyle<tensor::GenerateOp>);
617 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.
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.
Block * createBlock(Region *parent, Region::iterator insertPt={}, TypeRange argTypes=std::nullopt, ArrayRef< Location > locs=std::nullopt)
Add new block with 'argTypes' arguments and set the insertion point to the end of it.
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...
void mergeBlocks(Block *source, Block *dest, ValueRange argValues=std::nullopt)
Inline the operations of block 'source' into the end of block 'dest'.
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
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