44 LogicalResult matchAndRewrite(tensor::PadOp padOp,
47 Value padValue = padOp.getConstantPaddingValue();
53 Value source = padOp.getSource();
57 padOp,
"expected source to be linalg.generic op");
60 if (linalgOp.getNumLoops() != linalgOp.getNumParallelLoops()) {
62 padOp,
"only supported for ops with all parallel iterator types");
66 resultShape.size() != 1) {
68 padOp,
"failed to get shape of pad op result");
74 RankedTensorType padResultType = padOp.getResultType();
75 auto resultSizes = resultShape[0];
76 auto emptyTensor = rewriter.
create<tensor::EmptyOp>(
77 loc, resultSizes, padResultType.getElementType());
83 rewriter.
create<linalg::FillOp>(loc, padValue, emptyTensor.getResult());
89 unsigned resultNumber = cast<OpResult>(source).getResultNumber();
92 sizes.reserve(offsets.size());
93 for (
const auto &shape :
95 if (ShapedType::isDynamic(shape.value())) {
97 rewriter.
create<tensor::DimOp>(loc, source, shape.index())
104 auto slice = rewriter.
create<tensor::ExtractSliceOp>(
105 loc, fillTensor.getResult(0), offsets, sizes, strides);
109 cast<linalg::GenericOp>(rewriter.
clone(*linalgOp.getOperation()));
110 clonedOp.setDpsInitOperand(resultNumber, slice.getResult());
114 padOp, clonedOp.getResult(resultNumber), fillTensor.getResult(0),
115 offsets, sizes, strides);
IntegerAttr getIndexAttr(int64_t value)
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Operation * clone(Operation &op, IRMapping &mapper)
Creates a deep copy of the specified operation, remapping any operands that use values outside of the...
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
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,...
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 represents an instance of an SSA value in the MLIR system, representing a computable value...
Type getType() const
Return the type of this value.
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
void populateFuseTensorPadWithProducerLinalgOpPatterns(RewritePatternSet &patterns)
Pattern to fuse a tensor.pad operation with the producer of its source, if the producer is a linalg o...
Include the generated interface declarations.
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
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...