MLIR  21.0.0git
Padding.cpp
Go to the documentation of this file.
1 //===- Padding.cpp - Padding of Linalg ops --------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
10 
17 
18 #define DEBUG_TYPE "linalg-padding"
19 
20 using namespace mlir;
21 using namespace mlir::linalg;
22 
23 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")
24 #define DBGSNL() (llvm::dbgs() << "\n")
25 
26 namespace {
27 /// Helper class for storing padding information.
28 struct PaddingInfo {
29  PaddingInfo(int64_t padToMultipleOf = 1, OpFoldResult size = {})
30  : padToMultipleOf(padToMultipleOf), size(size) {}
31  /// Pad the tensor to a multiple of.
32  int64_t padToMultipleOf = 1;
33  /// The size used for padding.
34  OpFoldResult size = {};
35 };
36 
37 /// Helper class for storing and computing the padded shape.
38 struct PaddedShape {
39  /// Initializes the shape information and on success it returns whether the
40  /// shape of the operand will change. Returns failure if the operand cannot be
41  /// padded.
42  FailureOr<bool> initialize(linalg::LinalgOp opToPad, OpOperand *opOperand,
44 
45  /// Computs the padded shape.
46  void computePadding(OpBuilder &builder, Value operand);
47 
48  /// Returns the new tensor type.
49  RankedTensorType getType(Type elemTy) {
50  return RankedTensorType::get(shape, elemTy);
51  }
52 
53  SmallVector<Value> dynDims;
54 
55 private:
58 };
59 } // namespace
60 
61 FailureOr<bool> PaddedShape::initialize(linalg::LinalgOp opToPad,
62  OpOperand *opOperand,
64  AffineMap indexingMap = opToPad.getMatchingIndexingMap(opOperand);
65 
66  // Initialize the padded shape.
67  llvm::append_range(shape, opToPad.getShape(opOperand));
68 
69  // Collect the shape dimensions that are a function of "paddingDimensions",
70  // along with the multiple that they should be padded to ("1" if none).
71  bool alreadyHasRequestedShape = true;
72  for (const auto &dimEn : enumerate(options.paddingDimensions)) {
73  for (const auto &en : enumerate(indexingMap.getResults())) {
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()];
80  } else {
81  paddingInfo.padToMultipleOf = 1;
82  }
83 
84  // Check if the user provided a size in the options.
85  paddingInfo.size =
86  options.getSizeToPadTo(opOperand->getOperandNumber(), en.index());
87 
88  // Set the padding info.
89  dimToInfo[en.index()] = paddingInfo;
90  if (ShapedType::isDynamic(dimSize) ||
91  dimSize % paddingInfo.padToMultipleOf != 0 ||
92  !paddingInfo.size.isNull()) {
93  alreadyHasRequestedShape = false;
94  }
95  }
96  }
97  }
98 
99  // Upper bound the sizes to obtain a static bounding box.
100  for (int64_t i = 0, e = shape.size(); i < e; ++i) {
101  LLVM_DEBUG(DBGS() << "--computing un-padded size for dim " << i << "\n");
102  // Skip dimensions that do not require padding.
103  if (!dimToInfo.contains(i)) {
104  LLVM_DEBUG(DBGS() << "----dim does not require padding, SKIP\n");
105  continue;
106  }
107  PaddingInfo &info = dimToInfo[i];
108  if (info.size) {
109  LLVM_DEBUG(DBGS() << "----the user provided the size: " << info.size
110  << "\n");
111  continue;
112  }
113  // Otherwise, try to compute a constant upper bound for the size value.
114  FailureOr<int64_t> upperBound =
117  {opOperand->get(),
118  /*dim=*/i},
119  /*stopCondition=*/nullptr, /*closedUB=*/true);
120  if (failed(upperBound)) {
121  LLVM_DEBUG(
122  DBGS() << "----could not compute a bounding box for padding\n");
123  return failure();
124  }
125  info.size =
126  IntegerAttr::get(IndexType::get(opToPad.getContext()), *upperBound);
127  LLVM_DEBUG(DBGS() << "----new un-padded size: " << info.size << "\n");
128  }
129  return alreadyHasRequestedShape;
130 }
131 
132 void PaddedShape::computePadding(OpBuilder &builder, Value operand) {
133  Location loc = operand.getLoc();
134  AffineExpr sizeSym = builder.getAffineSymbolExpr(0);
135 
136  // Compute the padding for each dimension.
137  for (auto &&[i, dim] : llvm::enumerate(shape)) {
138  LLVM_DEBUG(DBGS() << "--computing padded size for dim " << i << "\n");
139 
140  // Get the padding info or default info for the shape dimension.
141  PaddingInfo paddingInfo = dimToInfo.lookup(i);
142 
143  // Skip dimensions that do not require padding.
144  if (paddingInfo.size.isNull()) {
145  LLVM_DEBUG(DBGS() << "----dim does not require padding, SKIP\n");
146 
147  // We still need to push the size as `makeComposedPadHighOp` expects a
148  // range with all the dynamic sizes, whether they're being padded or not.
149  if (ShapedType::isDynamic(dim)) {
150  dynDims.push_back(
151  cast<Value>(tensor::getMixedSize(builder, loc, operand, i)));
152  }
153  continue;
154  }
155 
156  // Compute the padded size to be a multiple of `padToMultipleOf`.
157  AffineExpr szExpr = (sizeSym).ceilDiv(paddingInfo.padToMultipleOf) *
158  paddingInfo.padToMultipleOf;
160  builder, loc, szExpr, paddingInfo.size);
161  assert(paddedSize && "invalid arguments to affine apply");
162 
163  if (auto cstSzAttr = dyn_cast<Attribute>(paddedSize)) {
164  // Update the shape as the size is static.
165  dim = cast<IntegerAttr>(cstSzAttr).getValue().getZExtValue();
166  } else {
167  // Add a dynamic dimension.
168  dim = ShapedType::kDynamic;
169  dynDims.push_back(cast<Value>(paddedSize));
170  }
171  LLVM_DEBUG(DBGS() << "----new dim size: " << paddedSize << "\n");
172  }
173 }
174 
175 /// Pad the `opOperand` in the "paddingDimensions" using the padding value and
176 /// the nofold flag found in "paddingValues" and "nofoldFlags", respectively.
177 ///
178 /// Exit early and return the `opOperand` value if it already has the requested
179 /// shape. i.e.:
180 /// - static shape
181 /// - nofold is not set
182 /// - dim sizes are multiples of "padToMultipleOf"
183 ///
184 /// Otherwise, try to pad the shape dimensions that match the iterator
185 /// dimensions "paddingDimensions" and return the tensor::PadOp result if
186 /// padding succeeds or failure otherwise.
188  RewriterBase &rewriter, linalg::LinalgOp opToPad, OpOperand *opOperand,
189  const LinalgPaddingOptions &options) {
190  assert(
191  (!options.padToMultipleOf.has_value() ||
192  options.padToMultipleOf->size() == options.paddingDimensions.size()) &&
193  "invalid number of elements in padToMultipleOf");
194 
195  // Initialize the padded shape and get whether it requires padding.
196  PaddedShape shape;
197  FailureOr<bool> alreadyHasRequestedShape =
198  shape.initialize(opToPad, opOperand, options);
199  if (failed(alreadyHasRequestedShape)) {
200  return rewriter.notifyMatchFailure(opToPad,
201  "--failed to compute padded shape");
202  }
203 
204  // Return the un-padded operand if padding to a static shape is not needed and
205  // if the nofold flag is not set.
206  bool nofold = opOperand->getOperandNumber() < options.nofoldFlags.size()
207  ? bool(options.nofoldFlags[opOperand->getOperandNumber()])
208  : false;
209  if (!nofold && *alreadyHasRequestedShape)
210  return opOperand->get();
211 
212  // Fail if `paddingValues` specifies no padding value.
213  if (opOperand->getOperandNumber() >= options.paddingValues.size()) {
214  return rewriter.notifyMatchFailure(opToPad, "--no padding value specified");
215  }
216  Attribute paddingAttr = options.paddingValues[opOperand->getOperandNumber()];
217 
218  Value paddingValue;
219  if (auto complexTy = dyn_cast<ComplexType>(
220  getElementTypeOrSelf(opOperand->get().getType()))) {
221  auto complexAttr = cast<ArrayAttr>(paddingAttr);
222  paddingValue = rewriter.create<complex::ConstantOp>(opToPad.getLoc(),
223  complexTy, complexAttr);
224  } else {
225  paddingValue = rewriter.create<arith::ConstantOp>(
226  opToPad.getLoc(), cast<TypedAttr>(paddingAttr));
227  }
228 
229  // Computes the padded shape.
230  if (!*alreadyHasRequestedShape)
231  shape.computePadding(rewriter, opOperand->get());
232 
233  // Pad the operand to the bounding box defined by `paddedShape`.
234  RankedTensorType paddedTensorType =
235  shape.getType(getElementTypeOrSelf(opOperand->get()));
236  LLVM_DEBUG(DBGS() << "--SUCCESS, makeComposedPadHighOp with type: "
237  << paddedTensorType);
238  return makeComposedPadHighOp(rewriter, opToPad->getLoc(), paddedTensorType,
239  opOperand->get(), paddingValue, nofold,
240  shape.dynDims);
241 }
242 
243 LogicalResult
244 linalg::rewriteAsPaddedOp(RewriterBase &rewriter, LinalgOp opToPad,
245  const LinalgPaddingOptions &constOptions,
246  LinalgOp &paddedOp, SmallVector<Value> &replacements,
247  SmallVector<tensor::PadOp> &padOps) {
248  LLVM_DEBUG(DBGS() << "Start rewriteAsPaddedOp : " << opToPad << "\n");
249  Location loc = opToPad->getLoc();
250 
251  LinalgPaddingOptions options(constOptions);
252  // Allow inference of pad values if they are not explicitly specified.
253  // TODO: be mindful about the value depending on the actual operation.
254  if (options.paddingValues.empty()) {
255  SmallVector<Type> types(opToPad->getOperandTypes());
256  llvm::append_range(types, opToPad->getResultTypes());
257  for (Type t : types) {
258  options.paddingValues.push_back(
259  rewriter.getZeroAttr(getElementTypeOrSelf(t)));
260  }
261  }
262 
263  // TODO: there are cases where we may still want to pad to larger sizes.
264  if (!opToPad.hasPureTensorSemantics())
265  return rewriter.notifyMatchFailure(opToPad,
266  "expected operation on tensors");
267 
268  OpBuilder::InsertionGuard g(rewriter);
269  // Set IP after op because we also take the dims of the original output.
270  rewriter.setInsertionPointAfter(opToPad);
271 
272  // Make a copy of the shaped operands and update it.
273  SmallVector<Value> newOperands;
274  newOperands.reserve(opToPad->getNumOperands());
275  for (OpOperand &opOperand : opToPad->getOpOperands()) {
276  FailureOr<Value> paddedOperand = padOperandToSmallestStaticBoundingBox(
277  rewriter, opToPad, &opOperand, options);
278  // Exit if `paddingDimensions` cannot be bounded statically.
279  if (failed(paddedOperand)) {
280  LLVM_DEBUG(DBGS() << "--operand cannot be bound statically : "
281  << opOperand.get() << " -> FAIL\n");
282  return rewriter.notifyMatchFailure(opToPad,
283  "operand cannot be bound statically");
284  }
285  newOperands.push_back(*paddedOperand);
286  if (auto padOp = paddedOperand->getDefiningOp<tensor::PadOp>())
287  padOps.push_back(padOp);
288  }
289 
290  ReifiedRankedShapedTypeDims reifiedResultShapes;
291  if (failed(reifyResultShapes(rewriter, opToPad, reifiedResultShapes))) {
292  LLVM_DEBUG(DBGS() << "--failed to reify result shapes -> FAIL\n");
293  return rewriter.notifyMatchFailure(opToPad,
294  "failed to reify result shapes");
295  }
296  assert(reifiedResultShapes.size() == opToPad->getNumResults() &&
297  "expected same number of results");
298 
299  // Clone `opToPad` to operate on the statically padded shapes.
300  auto resultTensorTypes =
301  ValueRange(newOperands).take_back(opToPad.getNumDpsInits()).getTypes();
302  // clone **should** properly notify the rewriter.
303  paddedOp = clone(rewriter, opToPad, resultTensorTypes, newOperands);
304  LLVM_DEBUG(DBGS() << "--cloned padded op: " << paddedOp << "\n");
305 
306  // Recover the slice out of the new static results. This keeps the original
307  // linalg op around because it uses the dims of the original results.
308  SmallVector<Value> paddedSubtensorResults;
309  paddedSubtensorResults.reserve(opToPad->getNumResults());
310  for (const auto &en : llvm::enumerate(paddedOp->getResults())) {
311  Value paddedResult = en.value();
312  int64_t resultNumber = en.index();
313  int64_t rank = cast<RankedTensorType>(paddedResult.getType()).getRank();
314  SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
315  SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
316  paddedSubtensorResults.push_back(rewriter.create<tensor::ExtractSliceOp>(
317  loc, paddedResult, offsets, reifiedResultShapes[resultNumber],
318  strides));
319  }
320 
322  replacements = std::move(paddedSubtensorResults);
323  return success();
324  }
325 
326  // Copy back unpadded results to the original destination (i.e., inits of the
327  // linalg op), so that the destination buffer of the computation does not
328  // change. If the padding folds away, this will materialize as a memcpy
329  // between two identical buffers, which will then also fold away.
330  assert(static_cast<int64_t>(paddedSubtensorResults.size()) ==
331  opToPad.getNumDpsInits() &&
332  "expected matching number of results");
333  for (auto it :
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())
339  .getResult(0));
340  } else if (options.copyBackOp ==
342  BufferizationMaterializeInDestination) {
343  replacements.push_back(
344  rewriter
345  .create<bufferization::MaterializeInDestinationOp>(
346  loc, std::get<0>(it), std::get<1>(it).get())
347  ->getResult(0));
348  } else {
349  llvm_unreachable("unsupported copy back op");
350  }
351  }
352  return success();
353 }
354 
355 FailureOr<LinalgOp>
356 mlir::linalg::padAndHoistLinalgOp(RewriterBase &rewriter, LinalgOp linalgOp,
357  const LinalgPaddingOptions &options) {
358  assert(options.copyBackOp == LinalgPaddingOptions::CopyBackOp::None &&
359  "invalid options");
360 
361  if (!linalgOp.hasPureTensorSemantics())
362  return rewriter.notifyMatchFailure(
363  linalgOp, "only applies to Linalg ops with tensor semantics");
364 
365  // Pad the operation.
366  LinalgOp paddedOp;
367  SmallVector<Value> newResults;
369  if (failed(rewriteAsPaddedOp(rewriter, linalgOp, options, paddedOp,
370  newResults, padOps)))
371  return rewriter.notifyMatchFailure(linalgOp,
372  "failed to rewrite as a padded op");
373 
374  // Hoist the padding.
375  for (const auto &en : enumerate(options.hoistPaddings)) {
376  if (static_cast<int64_t>(en.index()) >= paddedOp->getNumOperands())
377  break;
378  OpOperand &opOperand = paddedOp->getOpOperand(en.index());
379  auto padOp = opOperand.get().getDefiningOp<tensor::PadOp>();
380  if (!padOp || en.value() == 0) {
381  (void)rewriter.notifyMatchFailure(linalgOp, "not a tensor.pad -- skip");
382  continue;
383  }
384 
385  // Fail hoisting if the operand shape is not fully static.
386  if (llvm::any_of(paddedOp.getShape(&opOperand), ShapedType::isDynamic)) {
387  (void)rewriter.notifyMatchFailure(linalgOp,
388  "non static padding shape -- skip");
389  continue;
390  }
391 
392  tensor::PadOp hoistedOp;
393  SmallVector<TransposeOp> transposeOps;
394  SmallVector<int64_t> transposeVector =
395  en.index() < options.transposePaddings.size()
396  ? options.transposePaddings[en.index()]
398 
399  FailureOr<Value> newResult = hoistPaddingOnTensors(
400  padOp, en.value(), transposeVector, hoistedOp, transposeOps);
401  if (failed(newResult)) {
402  (void)rewriter.notifyMatchFailure(linalgOp,
403  "failed to apply hoistPadding");
404  continue;
405  }
406  rewriter.replaceOp(padOp, *newResult);
407  }
408 
409  // Replace the original operation to pad.
410  rewriter.replaceOp(linalgOp, newResults);
411 
412  return paddedOp;
413 }
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...
Definition: Padding.cpp:187
#define DBGS()
Definition: Padding.cpp:23
static llvm::ManagedStatic< PassManagerOptions > options
Base type for affine expression.
Definition: AffineExpr.h:68
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:46
ArrayRef< AffineExpr > getResults() const
Definition: AffineMap.cpp:403
Attributes are known-constant values of operations.
Definition: Attributes.h:25
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:103
AffineExpr getAffineSymbolExpr(unsigned position)
Definition: Builders.cpp:363
TypedAttr getZeroAttr(Type type)
Definition: Builders.cpp:319
IRValueT get() const
Return the current value being used by this operand.
Definition: UseDefLists.h:160
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:76
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:346
This class helps build Operations.
Definition: Builders.h:205
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:452
void setInsertionPointAfter(Operation *op)
Sets the insertion point to the node after the specified operation, which will cause subsequent inser...
Definition: Builders.h:410
This class represents a single result from folding an operation.
Definition: OpDefinition.h:271
This class represents an operand of an operation.
Definition: Value.h:257
unsigned getOperandNumber()
Return which operand this is in the OpOperand list of the Operation.
Definition: Value.cpp:228
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:358
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,...
Definition: PatternMatch.h:681
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...
Definition: Types.h:74
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.
Definition: ValueRange.h:387
type_range getTypes() const
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Type getType() const
Return the type of this value.
Definition: Value.h:105
Location getLoc() const
Return the location of this value.
Definition: Value.cpp:26
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Definition: Value.cpp:20
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...
Definition: AffineOps.cpp:1331
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:344
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...
Definition: Padding.cpp:244
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.
Definition: Padding.cpp:356
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...
Definition: Utils.cpp:243
OpFoldResult getMixedSize(OpBuilder &builder, Location loc, Value value, int64_t dim)
Return the dimension of the given tensor value.
Definition: TensorOps.cpp:61
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.
Definition: Utils.cpp:305
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...