18 #define DEBUG_TYPE "linalg-padding"
23 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")
24 #define DBGSNL() (llvm::dbgs() << "\n")
29 PaddingInfo(int64_t padToMultipleOf = 1,
OpFoldResult size = {})
30 : padToMultipleOf(padToMultipleOf), size(size) {}
32 int64_t padToMultipleOf = 1;
42 FailureOr<bool> initialize(linalg::LinalgOp opToPad,
OpOperand *opOperand,
61 FailureOr<bool> PaddedShape::initialize(linalg::LinalgOp opToPad,
64 AffineMap indexingMap = opToPad.getMatchingIndexingMap(opOperand);
67 llvm::append_range(shape, opToPad.getShape(opOperand));
71 bool alreadyHasRequestedShape =
true;
74 if (en.value().isFunctionOfDim(dimEn.value())) {
75 PaddingInfo paddingInfo;
76 int64_t dimSize = shape[en.index()];
77 if (
options.padToMultipleOf.has_value()) {
78 paddingInfo.padToMultipleOf =
79 (*
options.padToMultipleOf)[dimEn.index()];
81 paddingInfo.padToMultipleOf = 1;
89 dimToInfo[en.index()] = paddingInfo;
90 if (ShapedType::isDynamic(dimSize) ||
91 dimSize % paddingInfo.padToMultipleOf != 0 ||
92 !paddingInfo.size.isNull()) {
93 alreadyHasRequestedShape =
false;
100 for (int64_t i = 0, e = shape.size(); i < e; ++i) {
101 LLVM_DEBUG(
DBGS() <<
"--computing un-padded size for dim " << i <<
"\n");
103 if (!dimToInfo.contains(i)) {
104 LLVM_DEBUG(
DBGS() <<
"----dim does not require padding, SKIP\n");
107 PaddingInfo &info = dimToInfo[i];
109 LLVM_DEBUG(
DBGS() <<
"----the user provided the size: " << info.size
114 FailureOr<int64_t> upperBound =
120 if (failed(upperBound)) {
122 DBGS() <<
"----could not compute a bounding box for padding\n");
127 LLVM_DEBUG(
DBGS() <<
"----new un-padded size: " << info.size <<
"\n");
129 return alreadyHasRequestedShape;
132 void PaddedShape::computePadding(
OpBuilder &builder,
Value operand) {
138 LLVM_DEBUG(
DBGS() <<
"--computing padded size for dim " << i <<
"\n");
141 PaddingInfo paddingInfo = dimToInfo.lookup(i);
144 if (paddingInfo.size.isNull()) {
145 LLVM_DEBUG(
DBGS() <<
"----dim does not require padding, SKIP\n");
149 if (ShapedType::isDynamic(dim)) {
157 AffineExpr szExpr = (sizeSym).ceilDiv(paddingInfo.padToMultipleOf) *
158 paddingInfo.padToMultipleOf;
160 builder, loc, szExpr, paddingInfo.size);
161 assert(paddedSize &&
"invalid arguments to affine apply");
163 if (
auto cstSzAttr = dyn_cast<Attribute>(paddedSize)) {
165 dim = cast<IntegerAttr>(cstSzAttr).getValue().getZExtValue();
168 dim = ShapedType::kDynamic;
169 dynDims.push_back(cast<Value>(paddedSize));
171 LLVM_DEBUG(
DBGS() <<
"----new dim size: " << paddedSize <<
"\n");
191 (!
options.padToMultipleOf.has_value() ||
192 options.padToMultipleOf->size() ==
options.paddingDimensions.size()) &&
193 "invalid number of elements in padToMultipleOf");
197 FailureOr<bool> alreadyHasRequestedShape =
198 shape.initialize(opToPad, opOperand,
options);
199 if (failed(alreadyHasRequestedShape)) {
201 "--failed to compute padded shape");
209 if (!nofold && *alreadyHasRequestedShape)
210 return opOperand->
get();
219 if (
auto complexTy = dyn_cast<ComplexType>(
221 auto complexAttr = cast<ArrayAttr>(paddingAttr);
222 paddingValue = rewriter.
create<complex::ConstantOp>(opToPad.getLoc(),
223 complexTy, complexAttr);
225 paddingValue = rewriter.
create<arith::ConstantOp>(
226 opToPad.getLoc(), cast<TypedAttr>(paddingAttr));
230 if (!*alreadyHasRequestedShape)
231 shape.computePadding(rewriter, opOperand->
get());
234 RankedTensorType paddedTensorType =
236 LLVM_DEBUG(
DBGS() <<
"--SUCCESS, makeComposedPadHighOp with type: "
237 << paddedTensorType);
239 opOperand->
get(), paddingValue, nofold,
248 LLVM_DEBUG(
DBGS() <<
"Start rewriteAsPaddedOp : " << opToPad <<
"\n");
254 if (
options.paddingValues.empty()) {
256 llvm::append_range(types, opToPad->getResultTypes());
257 for (
Type t : types) {
258 options.paddingValues.push_back(
264 if (!opToPad.hasPureTensorSemantics())
266 "expected operation on tensors");
274 newOperands.reserve(opToPad->getNumOperands());
275 for (
OpOperand &opOperand : opToPad->getOpOperands()) {
277 rewriter, opToPad, &opOperand,
options);
279 if (failed(paddedOperand)) {
280 LLVM_DEBUG(
DBGS() <<
"--operand cannot be bound statically : "
281 << opOperand.
get() <<
" -> FAIL\n");
283 "operand cannot be bound statically");
285 newOperands.push_back(*paddedOperand);
286 if (
auto padOp = paddedOperand->getDefiningOp<tensor::PadOp>())
287 padOps.push_back(padOp);
292 LLVM_DEBUG(
DBGS() <<
"--failed to reify result shapes -> FAIL\n");
294 "failed to reify result shapes");
296 assert(reifiedResultShapes.size() == opToPad->getNumResults() &&
297 "expected same number of results");
300 auto resultTensorTypes =
303 paddedOp =
clone(rewriter, opToPad, resultTensorTypes, newOperands);
304 LLVM_DEBUG(
DBGS() <<
"--cloned padded op: " << paddedOp <<
"\n");
309 paddedSubtensorResults.reserve(opToPad->getNumResults());
311 Value paddedResult = en.value();
312 int64_t resultNumber = en.index();
313 int64_t rank = cast<RankedTensorType>(paddedResult.
getType()).getRank();
316 paddedSubtensorResults.push_back(rewriter.
create<tensor::ExtractSliceOp>(
317 loc, paddedResult, offsets, reifiedResultShapes[resultNumber],
322 replacements = std::move(paddedSubtensorResults);
330 assert(
static_cast<int64_t
>(paddedSubtensorResults.size()) ==
331 opToPad.getNumDpsInits() &&
332 "expected matching number of results");
334 llvm::zip(paddedSubtensorResults, opToPad.getDpsInitsMutable())) {
336 replacements.push_back(rewriter
337 .create<linalg::CopyOp>(loc, std::get<0>(it),
338 std::get<1>(it).
get())
340 }
else if (
options.copyBackOp ==
342 BufferizationMaterializeInDestination) {
343 replacements.push_back(
345 .create<bufferization::MaterializeInDestinationOp>(
346 loc, std::get<0>(it), std::get<1>(it).
get())
349 llvm_unreachable(
"unsupported copy back op");
361 if (!linalgOp.hasPureTensorSemantics())
363 linalgOp,
"only applies to Linalg ops with tensor semantics");
370 newResults, padOps)))
372 "failed to rewrite as a padded op");
376 if (
static_cast<int64_t
>(en.index()) >= paddedOp->getNumOperands())
378 OpOperand &opOperand = paddedOp->getOpOperand(en.index());
380 if (!padOp || en.value() == 0) {
386 if (llvm::any_of(paddedOp.getShape(&opOperand), ShapedType::isDynamic)) {
388 "non static padding shape -- skip");
392 tensor::PadOp hoistedOp;
395 en.index() <
options.transposePaddings.size()
396 ?
options.transposePaddings[en.index()]
400 padOp, en.value(), transposeVector, hoistedOp, transposeOps);
401 if (failed(newResult)) {
403 "failed to apply hoistPadding");
410 rewriter.
replaceOp(linalgOp, newResults);
static FailureOr< Value > padOperandToSmallestStaticBoundingBox(RewriterBase &rewriter, linalg::LinalgOp opToPad, OpOperand *opOperand, const LinalgPaddingOptions &options)
Pad the opOperand in the "paddingDimensions" using the padding value and the nofold flag found in "pa...
static llvm::ManagedStatic< PassManagerOptions > options
Base type for affine expression.
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
ArrayRef< AffineExpr > getResults() const
Attributes are known-constant values of operations.
IntegerAttr getIndexAttr(int64_t value)
AffineExpr getAffineSymbolExpr(unsigned position)
TypedAttr getZeroAttr(Type type)
IRValueT get() const
Return the current value being used by this operand.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
RAII guard to reset the insertion point of the builder when destroyed.
This class helps build Operations.
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...
This class represents a single result from folding an operation.
This class represents an operand of an operation.
unsigned getOperandNumber()
Return which operand this is in the OpOperand list of the Operation.
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...
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
static FailureOr< int64_t > computeConstantBound(presburger::BoundType type, const Variable &var, StopConditionFn stopCondition=nullptr, bool closedUB=false)
Compute a constant bound for the given variable.
This class provides an abstraction over the different types of ranges over Values.
type_range getTypes() const
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.
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...
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
LogicalResult rewriteAsPaddedOp(RewriterBase &rewriter, LinalgOp opToPad, const LinalgPaddingOptions &options, LinalgOp &paddedOp, SmallVector< Value > &replacements, SmallVector< tensor::PadOp > &padOps)
Pad the iterator dimensions options.paddingDimensions of all opToPad operands to a static bounding bo...
FailureOr< Value > hoistPaddingOnTensors(RewriterBase &rewriter, tensor::PadOp opToHoist, int64_t numLoops, ArrayRef< int64_t > transposeVector, tensor::PadOp &hoistedOp, SmallVectorImpl< TransposeOp > &transposeOps)
Mechanically hoist padding operations on tensors by numLoops into a new, generally larger tensor.
FailureOr< LinalgOp > padAndHoistLinalgOp(RewriterBase &rewriter, LinalgOp linalgOp, const LinalgPaddingOptions &options)
Apply padding and hoisting to linalgOp according to the configuration specified in options.
Value makeComposedPadHighOp(OpBuilder &b, Location loc, RankedTensorType type, Value source, Value padding, bool nofold, ValueRange typeDynDims={})
Create a tensor::PadOp that pads source to the shape of type whose sizes are assumed to be greater th...
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).
Type getType(OpFoldResult ofr)
Returns the int type of the integer in ofr.
Type getElementTypeOrSelf(Type type)
Return the element type or return the type itself.
Operation * clone(OpBuilder &b, Operation *op, TypeRange newResultTypes, ValueRange newOperands)
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...