24 struct PadOpTiling :
public TilingInterface::ExternalModel<PadOpTiling, PadOp> {
27 auto padOp = cast<PadOp>(op);
29 padOp.getResultType().getRank(), utils::IteratorType::parallel);
41 for (
const auto &ub :
enumerate(reifiedShapes[0]))
42 loopRanges[ub.index()].size = ub.value();
46 FailureOr<TilingResult>
50 FailureOr<TilingResult> result =
54 return result.value();
63 resultOffsets.assign(offsets.begin(), offsets.end());
64 resultSizes.assign(sizes.begin(), sizes.end());
68 LogicalResult getIterationDomainTileFromResultTile(
73 iterDomainOffsets.assign(offsets.begin(), offsets.end());
74 iterDomainSizes.assign(sizes.begin(), sizes.end());
78 FailureOr<TilingResult>
92 bool generateZeroSliceGuard) {
94 Value padValue = padOp.getConstantPaddingValue();
124 bool hasZeroLen =
false;
127 Value dynHasZeroLenCond;
129 int64_t rank = padOp.getSourceType().getRank();
132 for (
unsigned dim = 0; dim < rank; ++dim) {
133 auto low = padOp.getMixedLowPad()[dim];
135 auto high = padOp.getMixedHighPad()[dim];
137 auto offset = offsets[dim];
138 auto length = sizes[dim];
141 if (!hasLowPad && !hasHighPad) {
142 newOffsets.push_back(offset);
143 newLengths.push_back(length);
144 newLows.push_back(low);
145 newHighs.push_back(high);
157 newLows.push_back(newLow);
172 ?
min(
max(sub(offset, low), zero), srcSize)
173 :
min(offset, srcSize);
174 newOffsets.push_back(newOffset);
198 OpFoldResult newLength =
min(sub(srcSize, newOffset), sub(length, newLow));
202 newLength =
max(newLength, zero);
203 newLengths.push_back(newLength);
209 }
else if (!hasZeroLen) {
211 loc, arith::CmpIPredicate::eq,
216 ? b.
create<arith::OrIOp>(loc, check, dynHasZeroLenCond)
225 hasHighPad ? sub(sub(length, newLength), newLow) : zero;
226 newHighs.push_back(newHigh);
233 RankedTensorType resultType =
238 if (resultType == val.getType())
240 return b.
create<tensor::CastOp>(loc, resultType, val);
246 auto createGenerateOp = [&]() {
248 auto generateOp = b.
create<tensor::GenerateOp>(
249 loc, resultType, dynDims,
251 builder.create<tensor::YieldOp>(gLoc, padValue);
258 auto createPadOfExtractSlice = [&]() {
260 auto newSliceOp = b.
create<tensor::ExtractSliceOp>(
261 loc, padOp.getSource(), newOffsets, newLengths, newStrides);
262 auto newPadOp = b.
create<PadOp>(
263 loc,
Type(), newSliceOp, newLows, newHighs,
269 padOp.getRegion().cloneInto(&newPadOp.getRegion(), bvm);
272 return std::make_tuple(newPadOp, newSliceOp);
278 Operation *generateOp = createGenerateOp();
286 if (generateZeroSliceGuard && dynHasZeroLenCond) {
290 auto result = b.
create<scf::IfOp>(
291 loc, dynHasZeroLenCond,
294 thenOp = createGenerateOp();
295 b.create<scf::YieldOp>(loc, castResult(thenOp->
getResult(0)));
299 std::tie(elseOp, sliceOp) = createPadOfExtractSlice();
300 b.create<scf::YieldOp>(loc, castResult(elseOp->
getResult(0)));
306 auto [newPadOp, sliceOp] = createPadOfExtractSlice();
308 {newPadOp}, {castResult(newPadOp->getResult(0))}, {sliceOp}};
314 tensor::PadOp::attachInterface<PadOpTiling>(*ctx);
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
static Value min(ImplicitLocOpBuilder &builder, Value value, Value bound)
static LogicalResult getResultTilePosition(RewriterBase &rewriter, ReductionTilingStrategy reductionStrategy, int64_t index, Value tiledResult, TilingInterface op, ArrayRef< OpFoldResult > offsets, ArrayRef< OpFoldResult > sizes, ValueRange ivs, ArrayRef< OpFoldResult > numThreads, ArrayRef< OpFoldResult > tileSizes, const SetVector< unsigned > &reductionDims, SmallVector< OpFoldResult > &resultOffset, SmallVector< OpFoldResult > &resultSize)
static FailureOr< TilingResult > getTiledImplementation(RewriterBase &rewriter, TilingInterface op, ReductionTilingStrategy reductionStrategy, ValueRange regionIterArg, ArrayRef< OpFoldResult > offsets, ArrayRef< OpFoldResult > sizes, ValueRange ivs, ArrayRef< OpFoldResult > numThreads, ArrayRef< OpFoldResult > tileSizes, const SetVector< unsigned > &reductionDims)
Base type for affine expression.
static AffineMap getMultiDimIdentityMap(unsigned numDims, MLIRContext *context)
Returns an AffineMap with 'numDims' identity result dim exprs.
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
IntegerAttr getIndexAttr(int64_t value)
MLIRContext * getContext() const
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
bool addExtension(TypeID extensionID, std::unique_ptr< DialectExtensionBase > extension)
Add the given extension to the registry.
This is a utility class for mapping one set of IR entities to another.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
MLIRContext is the top-level object for a collection of MLIR operations.
This class helps build Operations.
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
This class represents a single result from folding an operation.
Operation is the basic unit of execution within MLIR.
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
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...
OpFoldResult makeComposedFoldedAffineMax(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Constructs an AffineMinOp that computes a maximum across the results of applying map to operands,...
OpFoldResult makeComposedFoldedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands, bool composeAffineMin=false)
Constructs an AffineApplyOp that applies map to operands after composing the map with the maps of any...
OpFoldResult makeComposedFoldedAffineMin(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Constructs an AffineMinOp that computes a minimum across the results of applying map to operands,...
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
FailureOr< TilingResult > bubbleUpPadSlice(OpBuilder &b, tensor::PadOp padOp, ArrayRef< OpFoldResult > offsets, ArrayRef< OpFoldResult > sizes, bool generateZeroSliceGuard=true)
Bubbles up a slice of this pad by taking the slice first and then performing the padding.
void registerTilingInterfaceExternalModels(mlir::DialectRegistry ®istry)
Registers external models for Tiling interface for tensor ops.
OpFoldResult getMixedSize(OpBuilder &builder, Location loc, Value value, int64_t dim)
Return the dimension of the given tensor value.
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).
void bindDims(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to DimExpr at positions: [0 .
bool isZeroInteger(OpFoldResult v)
Return true if v is an IntegerAttr with value 0.
void dispatchIndexOpFoldResults(ArrayRef< OpFoldResult > ofrs, SmallVectorImpl< Value > &dynamicVec, SmallVectorImpl< int64_t > &staticVec)
Helper function to dispatch multiple OpFoldResults according to the behavior of dispatchIndexOpFoldRe...
Value getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc, OpFoldResult ofr)
Converts an OpFoldResult to a Value.
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
SmallVector< NamedAttribute > getPrunedAttributeList(Operation *op, ArrayRef< StringRef > elidedAttrs)
Container for result values of tiling.