43 LogicalResult matchAndRewrite(tensor::PadOp padOp,
46 Value padValue = padOp.getConstantPaddingValue();
52 Value source = padOp.getSource();
56 padOp,
"expected source to be linalg.generic op");
59 if (linalgOp.getNumLoops() != linalgOp.getNumParallelLoops()) {
61 padOp,
"only supported for ops with all parallel iterator types");
65 resultShape.size() != 1) {
67 padOp,
"failed to get shape of pad op result");
73 RankedTensorType padResultType = padOp.getResultType();
74 auto resultSizes = resultShape[0];
75 auto emptyTensor = tensor::EmptyOp::create(rewriter, loc, resultSizes,
76 padResultType.getElementType());
81 auto fillTensor = linalg::FillOp::create(rewriter, loc, padValue,
82 emptyTensor.getResult());
88 unsigned resultNumber = cast<OpResult>(source).getResultNumber();
91 sizes.reserve(offsets.size());
92 for (
const auto &shape :
94 if (ShapedType::isDynamic(shape.value())) {
96 tensor::DimOp::create(rewriter, loc, source, shape.index())
103 auto slice = tensor::ExtractSliceOp::create(
104 rewriter, loc, fillTensor.getResult(0), offsets, sizes, strides);
108 cast<linalg::GenericOp>(rewriter.
clone(*linalgOp.getOperation()));
109 clonedOp.setDpsInitOperand(resultNumber, slice.getResult());
113 padOp, clonedOp.getResult(resultNumber), fillTensor.getResult(0),
114 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...
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...