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, padOp.getResult().getType(), alloc,
true,
264 assert(llvm::range_size(maskOp.getMaskBlock()->without_terminator()) == 1 &&
265 "expected single masked op");
272 Operation *yieldOp = maskOp.getMaskRegion().front().getTerminator();
273 assert(isa<vector::YieldOp>(yieldOp) &&
"expected yield op terminator");
278 rewriter,
options, maskOp.getMaskableOp(), memorySpace,
279 insertionPoint ? insertionPoint : maskOp);
281 if (
options.bufferizeDestinationOnly)
286 if (failed(cast<bufferization::BufferizableOpInterface>(yieldOp).bufferize(
287 rewriter, bufferizationOptions, bufferizationState)))
294 maskOp.walk([&](bufferization::ToTensorOp toTensorOp) {
295 if (toTensorOp->getUses().empty())
296 toTensorOps.push_back(toTensorOp.getOperation());
303 for (
Value result : maskOp.getResults())
304 if (isa<TensorType>(result.getType()))
306 resultUses.push_back(&use);
309 cast<bufferization::BufferizableOpInterface>(maskOp.getOperation())
310 .bufferize(rewriter, bufferizationOptions, bufferizationState)))
315 for (
OpOperand *resultUse : resultUses) {
317 resultUse->get().getDefiningOp<bufferization::ToTensorOp>();
318 assert(toTensorOp &&
"expected to_tensor op");
320 toTensorOp.setRestrict(
true);
321 toTensorOp.setWritable(
true);
330 bufferization::AllocTensorOp allocTensorOp,
Attribute memorySpace,
332 Location loc = allocTensorOp.getLoc();
339 rewriter, loc, allocTensorOp.getResult(),
options, memorySpace);
343 Value toTensorOp = rewriter.
create<bufferization::ToTensorOp>(
344 loc, allocTensorOp.getResult().getType(), alloc,
true,
346 rewriter.
replaceOp(allocTensorOp, toTensorOp);
352 RewriterBase &rewriter, tensor::FromElementsOp fromElementsOp) {
353 Location loc = fromElementsOp.getLoc();
354 RankedTensorType tensorType =
355 cast<RankedTensorType>(fromElementsOp.getType());
356 auto shape = tensorType.getShape();
364 fromElementsOp, fromElementsOp.getElements().front(),
370 auto maxDim = *llvm::max_element(shape);
372 constants.reserve(maxDim);
373 for (
int i = 0; i < maxDim; ++i)
377 auto elementIt = fromElementsOp.getElements().begin();
380 shape, constants, elementIt, indices);
383 rewriter.
replaceOp(fromElementsOp, result);
388 FailureOr<Operation *>
390 tensor::GenerateOp generateOp) {
392 if (!generateOp.getBody().hasOneBlock())
396 RankedTensorType tensorType = cast<RankedTensorType>(generateOp.getType());
400 rewriter.
create<EmptyOp>(loc, tensorType, generateOp.getDynamicExtents());
404 utils::IteratorType::parallel);
407 auto genericOp = rewriter.
create<linalg::GenericOp>(
410 indexingMaps, iteratorTypes);
412 tensorType.getElementType(), loc);
415 for (int64_t i = 0; i < tensorType.getRank(); ++i)
416 bbArgReplacements.push_back(rewriter.
create<linalg::IndexOp>(loc, i));
417 rewriter.
mergeBlocks(&generateOp.getBody().front(), body, bbArgReplacements);
420 auto yieldOp = cast<tensor::YieldOp>(body->getTerminator());
424 rewriter.
replaceOp(generateOp, genericOp->getResult(0));
425 return genericOp.getOperation();
429 FailureOr<Operation *>
431 tensor::PadOp padOp) {
433 if (!padOp.getBodyRegion().hasOneBlock())
438 RankedTensorType resultType = padOp.getResultType();
442 padOp,
"failed to reify tensor.pad op result shape");
444 for (int64_t i = 0; i < resultType.getRank(); ++i)
445 if (resultType.isDynamicDim(i))
446 dynamicSizes.push_back(cast<Value>(reifiedShape[0][i]));
450 if (padOp.getNofoldAttr() &&
453 using bufferization::AllocTensorOp;
455 rewriter.
create<AllocTensorOp>(loc, resultType, dynamicSizes);
457 padOp, padOp.getSource(), allocated);
458 return copyOp.getOperation();
461 Value empty = rewriter.
create<EmptyOp>(loc, resultType, dynamicSizes);
472 padOp, padOp.getSource(), fillOp->getResult(0),
473 padOp.getMixedLowPad(), sliceSizes, sliceStrides);
474 return insertSliceOp.getOperation();
480 using namespace bufferization;
483 if (
auto padOp = dyn_cast<tensor::PadOp>(op))
485 if (
auto maskOp = dyn_cast<vector::MaskOp>(op))
487 if (
auto allocTensorOp = dyn_cast<bufferization::AllocTensorOp>(op))
491 auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
496 BufferizationOptions bufferizationOptions;
498 BufferizationState bufferizationState;
501 if (!
options.bufferizeDestinationOnly) {
508 [](
Value v) { return isa<TensorType>(v.getType()); }))
509 llvm_unreachable(
"ops with nested tensor ops are not supported yet");
511 [](
Value v) { return isa<TensorType>(v.getType()); }))
512 llvm_unreachable(
"ops with nested tensor ops are not supported yet");
520 if (!isa<TensorType>(result.getType()))
523 if (!isa<RankedTensorType>(result.getType()))
526 if (bufferizableOp.bufferizesToAllocation(result))
528 tensorResults.push_back(result);
534 auto addOutOfPlaceOperand = [&](
OpOperand *operand) {
535 if (!llvm::is_contained(outOfPlaceOperands, operand))
536 outOfPlaceOperands.push_back(operand);
538 for (
OpResult result : tensorResults) {
540 analysisState.getAliasingOpOperands(result);
541 for (
const AliasingOpOperand &operand : aliasingOperands) {
542 addOutOfPlaceOperand(operand.opOperand);
543 for (
OpOperand &resultUse : result.getUses())
544 resultUses.push_back(&resultUse);
548 if (!analysisState.bufferizesToMemoryWrite(operand))
550 if (!isa<RankedTensorType>(operand.get().getType()))
552 addOutOfPlaceOperand(&operand);
555 if (outOfPlaceOperands.size() != 1)
562 for (
OpOperand *operand : outOfPlaceOperands) {
564 rewriter, op->
getLoc(), operand->get(),
options, memorySpace);
565 allocs.push_back(alloc);
566 if (!analysisState.findDefinitions(operand).empty()) {
572 auto toTensorOp = rewriter.
create<ToTensorOp>(
573 op->
getLoc(), operand->get().getType(), alloc);
574 operand->set(toTensorOp);
575 if (
options.bufferizeDestinationOnly) {
577 toTensorOp.setRestrict(
true);
578 toTensorOp.setWritable(
true);
584 if (
options.bufferizeDestinationOnly)
585 return allocs.front();
589 if (failed(bufferizableOp.bufferize(rewriter, bufferizationOptions,
590 bufferizationState)))
595 for (
OpOperand *resultUse : resultUses) {
596 auto toTensorOp = resultUse->get().
getDefiningOp<ToTensorOp>();
597 assert(toTensorOp &&
"expected to_tensor op");
599 toTensorOp.setRestrict(
true);
600 toTensorOp.setWritable(
true);
603 return allocs.front();
608 template <
typename OpTy>
609 LogicalResult rewriteOpInDestinationPassingStyle(OpTy op,
618 patterns.add(rewriteOpInDestinationPassingStyle<tensor::FromElementsOp>);
619 patterns.add(rewriteOpInDestinationPassingStyle<tensor::GenerateOp>);
620 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