MLIR  15.0.0git
Tiling.cpp
Go to the documentation of this file.
1 //===- Tiling.cpp - Implementation of linalg Tiling -----------------------===//
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 //
9 // This file implements the linalg dialect Tiling pass.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include <utility>
14 
15 #include "PassDetail.h"
24 #include "mlir/IR/AffineExpr.h"
25 #include "mlir/IR/AffineMap.h"
28 
29 #include "llvm/Support/CommandLine.h"
30 
31 using namespace mlir;
32 using namespace mlir::linalg;
33 using namespace mlir::scf;
34 
35 #define DEBUG_TYPE "linalg-tiling"
36 
37 static bool isZero(Value v) {
38  if (auto cst = v.getDefiningOp<arith::ConstantIndexOp>())
39  return cst.value() == 0;
40  return false;
41 }
42 
43 std::tuple<SmallVector<Range, 4>, LoopIndexToRangeIndexMap>
45  ValueRange allShapeSizes,
46  ValueRange allTileSizes) {
47  assert(allTileSizes.size() == map.getNumResults());
48  // Apply `map` to get shape sizes in loop order.
49  auto shapeSizes = applyMapToValues(b, loc, map, allShapeSizes);
50  SmallVector<Value, 4> tileSizes(allTileSizes.begin(), allTileSizes.end());
51 
52  // Traverse the tile sizes, which are in loop order, erase zeros everywhere.
53  LoopIndexToRangeIndexMap loopIndexToRangeIndex;
54  for (int idx = 0, e = tileSizes.size(), zerosCount = 0; idx < e; ++idx) {
55  if (isZero(tileSizes[idx - zerosCount])) {
56  shapeSizes.erase(shapeSizes.begin() + idx - zerosCount);
57  tileSizes.erase(tileSizes.begin() + idx - zerosCount);
58  ++zerosCount;
59  continue;
60  }
61  loopIndexToRangeIndex[idx] = idx - zerosCount;
62  }
63 
64  // Create a new range with the applied tile sizes.
66  for (unsigned idx = 0, e = tileSizes.size(); idx < e; ++idx)
67  res.push_back(Range{b.create<arith::ConstantIndexOp>(loc, 0),
68  shapeSizes[idx], tileSizes[idx]});
69  return std::make_tuple(res, loopIndexToRangeIndex);
70 }
71 
73  RewriterBase &b, LinalgOp op, SmallVectorImpl<Value> &ivs,
74  const LoopIndexToRangeIndexMap &loopIndexToRangeIndex) {
75  SmallVector<Value> allIvs(op.getNumLoops(), nullptr);
76  for (auto &en : enumerate(allIvs)) {
77  auto rangeIndex = loopIndexToRangeIndex.find(en.index());
78  if (rangeIndex == loopIndexToRangeIndex.end())
79  continue;
80  en.value() = ivs[rangeIndex->second];
81  }
82  addTileLoopIvsToIndexOpResults(b, op, allIvs);
83 }
84 
85 // Insert a tile `source` into the destination tensor `dest`. The position at
86 // which the tile is inserted (as well as size of tile) is taken from a given
87 // ExtractSliceOp `sliceOp`.
89  tensor::ExtractSliceOp sliceOp, Value source,
90  Value dest) {
91  return b.create<tensor::InsertSliceOp>(
92  loc, sliceOp.source().getType(), source, dest, sliceOp.offsets(),
93  sliceOp.sizes(), sliceOp.strides(), sliceOp.static_offsets(),
94  sliceOp.static_sizes(), sliceOp.static_strides());
95 }
96 
97 template <typename LoopTy>
99 tileLinalgOpImpl(RewriterBase &b, LinalgOp op, ValueRange tileSizes,
100  const LinalgTilingOptions &options) {
101  auto nLoops = op.getNumLoops();
102  // Initial tile sizes may be too big, only take the first nLoops.
103  tileSizes = tileSizes.take_front(nLoops);
104 
105  if (llvm::all_of(tileSizes, isZero)) {
106  TiledLinalgOp tiledOp;
107  tiledOp.op = cast<LinalgOp>(b.clone(*op.getOperation()));
108  tiledOp.tensorResults.assign(tiledOp.op->result_begin(),
109  tiledOp.op->result_end());
110  return tiledOp;
111  }
112 
113  // 1. Build the tiled loop ranges.
114  auto allShapeSizes = op.createFlatListOfOperandDims(b, op.getLoc());
115  AffineMap shapeSizesToLoopsMap = op.getShapesToLoopsMap();
116  if (!shapeSizesToLoopsMap)
117  return failure();
118 
119  SmallVector<Range, 4> loopRanges;
120  LoopIndexToRangeIndexMap loopIndexToRangeIndex;
121  std::tie(loopRanges, loopIndexToRangeIndex) = makeTiledLoopRanges(
122  b, op.getLoc(), shapeSizesToLoopsMap, allShapeSizes, tileSizes);
123 
124  SmallVector<Attribute, 4> iteratorTypes;
125  for (const auto &attr :
126  enumerate(op.iterator_types().cast<ArrayAttr>().getValue())) {
127  if (loopIndexToRangeIndex.count(attr.index()))
128  iteratorTypes.push_back(attr.value());
129  }
130  // If interchangeVector is empty, use the identity. Build the permutation map
131  // otherwise.
132  auto invPermutationMap =
133  AffineMap::getMultiDimIdentityMap(tileSizes.size(), b.getContext());
134  if (!options.interchangeVector.empty()) {
135  // Based on the pruned iterations (due to zero tile size), recompute the
136  // interchange vector.
137  SmallVector<unsigned, 4> interchangeVector;
138  interchangeVector.reserve(options.interchangeVector.size());
139  for (auto pos : options.interchangeVector) {
140  auto it = loopIndexToRangeIndex.find(pos);
141  if (it == loopIndexToRangeIndex.end())
142  continue;
143  interchangeVector.push_back(it->second);
144  }
145  // Interchange vector is guaranteed to be a permutation,
146  // `inversePermutation` must succeed.
147  invPermutationMap = inversePermutation(
148  AffineMap::getPermutationMap(interchangeVector, b.getContext()));
149  assert(invPermutationMap);
150  SmallVector<int64_t> permutation(interchangeVector.begin(),
151  interchangeVector.end());
152  applyPermutationToVector(loopRanges, permutation);
153  applyPermutationToVector(iteratorTypes, permutation);
154  }
155 
156  // 2. Create the tiled loops.
157  LinalgOp res = op;
158  SmallVector<Value, 4> ivs, tensorResults;
159  auto tiledLoopBodyBuilder =
160  [&](OpBuilder &builder, Location loc, ValueRange localIvs,
161  ValueRange operandValuesToUse) -> scf::ValueVector {
162  ivs.assign(localIvs.begin(), localIvs.end());
163 
164  // When an `interchangeVector` is present, it has been applied to the
165  // loop ranges and the iterator types. Apply its inverse to the
166  // resulting loop `ivs` to match the op definition.
167  SmallVector<Value, 4> interchangedIvs;
168  if (!options.interchangeVector.empty())
169  interchangedIvs = applyMapToValues(b, loc, invPermutationMap, ivs);
170  else
171  interchangedIvs.assign(ivs.begin(), ivs.end());
172 
173  // Tile the `operandValuesToUse` that either match the `op` operands
174  // themselves or the tile loop arguments forwarding them.
175  assert(operandValuesToUse.size() ==
176  static_cast<size_t>(op.getNumInputsAndOutputs()) &&
177  "expect the number of operands and inputs and outputs to match");
178  SmallVector<Value> valuesToTile = operandValuesToUse;
179  auto sizeBounds =
180  applyMapToValues(b, loc, shapeSizesToLoopsMap, allShapeSizes);
181  SmallVector<Value, 4> tiledOperands =
182  makeTiledShapes(b, loc, op, valuesToTile, interchangedIvs, tileSizes,
183  sizeBounds, /*omitPartialTileCheck=*/false);
184 
185  // TODO: use an interface/adaptor to avoid leaking position in
186  // `tiledOperands`.
187  SmallVector<Type, 4> resultTensorTypes;
188  for (OpOperand *opOperand : op.getOutputTensorOperands())
189  resultTensorTypes.push_back(
190  tiledOperands[opOperand->getOperandNumber()].getType());
191 
192  res = op.clone(b, loc, resultTensorTypes, tiledOperands);
193 
194  // Insert a insert_slice for each output tensor.
195  unsigned resultIdx = 0;
196  for (OpOperand *opOperand : op.getOutputTensorOperands()) {
197  // TODO: use an interface/adaptor to avoid leaking position in
198  // `tiledOperands`.
199  Value outputTensor = tiledOperands[opOperand->getOperandNumber()];
200  // TODO: Propagate RewriterBase everywhere.
201  IRRewriter rewriter(b);
202  if (auto sliceOp = outputTensor.getDefiningOp<tensor::ExtractSliceOp>()) {
203  tensorResults.push_back(insertSliceIntoTensor(rewriter, loc, sliceOp,
204  res->getResult(resultIdx),
205  sliceOp.source()));
206  } else {
207  tensorResults.push_back(res->getResult(resultIdx));
208  }
209  ++resultIdx;
210  }
211  return scf::ValueVector(tensorResults.begin(), tensorResults.end());
212  };
213  GenerateLoopNest<LoopTy>::doit(b, op.getLoc(), loopRanges, op, iteratorTypes,
214  tiledLoopBodyBuilder, options.distribution,
215  options.distributionTypes);
216 
217  // 3. Transform IndexOp results w.r.t. the tiling.
218  transformIndexOps(b, res, ivs, loopIndexToRangeIndex);
219 
220  // 4. Gather the newly created loops and return them with the new op.
222  loops.reserve(ivs.size());
223  for (auto iv : ivs) {
224  if (iv.isa<BlockArgument>()) {
225  loops.push_back(iv.cast<BlockArgument>().getOwner()->getParentOp());
226  assert(loops.back() && "no owner found for induction variable!");
227  } else {
228  // TODO: Instead of doing this, try to recover the ops used instead of the
229  // loop.
230  loops.push_back(nullptr);
231  }
232  }
233 
234  // 5. Get the tensor results from the outermost loop if available. Otherwise
235  // use the previously captured `tensorResults`.
236  Operation *outermostLoop = nullptr;
237  for (Operation *loop : loops)
238  if ((outermostLoop = loop))
239  break;
240 
241  return TiledLinalgOp{
242  res, loops, outermostLoop ? outermostLoop->getResults() : tensorResults};
243 }
244 
245 template <typename LoopTy>
247  RewriterBase &b, LinalgOp op, const LinalgTilingOptions &options) {
249  b.setInsertionPoint(op);
250 
251  if (!options.tileSizeComputationFunction)
252  return failure();
253 
254  // Enforce the convention that "tiling by zero" skips tiling a particular
255  // dimension. This convention is significantly simpler to handle instead of
256  // adjusting affine maps to account for missing dimensions.
257  auto nLoops = op.getNumLoops();
258  SmallVector<Value, 4> tileSizeVector =
259  options.tileSizeComputationFunction(b, op);
260  if (tileSizeVector.size() < nLoops) {
261  auto zero = b.create<arith::ConstantIndexOp>(op.getLoc(), 0);
262  tileSizeVector.append(nLoops - tileSizeVector.size(), zero);
263  }
264 
265  return tileLinalgOpImpl<LoopTy>(b, op, tileSizeVector, options);
266 }
267 
270  const LinalgTilingOptions &options) {
271  switch (options.loopType) {
273  return tileLinalgOpImpl<scf::ForOp>(b, op, options);
275  return tileLinalgOpImpl<scf::ParallelOp>(b, op, options);
276  default:;
277  }
278  return failure();
279 }
280 
281 /// Generate a loop nest around a given tensor::PadOp (for tiling). `newPadOp`
282 /// and `loopNest` are output parameters that return the new (tiled)
283 /// tensor::PadOp and the loop nest.
284 static LogicalResult tilePadOp(RewriterBase &builder, tensor::PadOp op,
285  tensor::PadOp &newPadOp, LoopNest &loopNest,
286  const LinalgTilingOptions &options) {
287  Location loc = op.getLoc();
288  OpBuilder::InsertionGuard g(builder);
289  builder.setInsertionPoint(op);
290 
291  // Clone tensor::PadOp so that the existing op can be replaced more easily.
292  newPadOp = cast<tensor::PadOp>(builder.clone(*op.getOperation()));
293  // Get rank and tile sizes.
294  int64_t rank = op.getResultType().getRank();
295  SmallVector<Value> tileSizes =
296  options.tileSizeComputationFunction(builder, op);
297  // Normalize untiled padding dimensions to 0.
298  Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
299  tileSizes.append(rank - tileSizes.size(), zero);
300  // Compute lower and upper bounds of the loop nest.
301  TilingInterface tilingInterface =
302  dyn_cast<TilingInterface>(op.getOperation());
303  SmallVector<Range> ranges = tilingInterface.getIterationDomain(builder);
304  SmallVector<Value> lbs, dims, allDims, steps;
305  for (int64_t i = 0; i < rank; ++i) {
306  allDims.push_back(ranges[i].size);
307  if (!isZero(tileSizes[i])) {
308  lbs.push_back(ranges[i].offset);
309  dims.push_back(ranges[i].size);
310  steps.push_back(tileSizes[i]);
311  }
312  }
313  // Generate loop nest: One loop per dimension.
314  SmallVector<Value> destOperand =
315  tilingInterface.getDestinationOperands(builder);
316  loopNest = mlir::scf::buildLoopNest(
317  builder, loc, lbs, /*ubs=*/dims, steps, ValueRange(destOperand),
318  [&](OpBuilder &b, Location loc, ValueRange localIvs,
319  ValueRange iterArgs) -> scf::ValueVector {
320  // Compute offsets and sizes of ExtractSliceOp.
321  SmallVector<Value> offsets =
322  computeTileOffsets(b, loc, localIvs, tileSizes);
323  SmallVector<Value> sizes =
324  computeTileSizes(b, loc, localIvs, tileSizes, allDims);
325  // Create ExtractSliceOp: Extract a tile from the tensor::PadOp.
326  // Note: The tensor::PadOp is located outside of the loop nest. It is
327  // later moved inside by ExtractSliceOfPadTensorSwapPattern.
328  auto map = AffineMap::getMultiDimIdentityMap(rank, b.getContext());
329  Value tiledOutput = makeTiledShape(
330  b, loc, newPadOp->getResult(0), tileSizes, map, offsets, allDims,
331  sizes, /*omitPartialTileCheck=*/false);
332  auto sliceOp = tiledOutput.getDefiningOp<tensor::ExtractSliceOp>();
333  assert(sliceOp && "expected ExtractSliceOp");
334  // Insert the tile into the output tensor.
335  // TODO: Propagate RewriterBase everywhere.
336  IRRewriter rewriter(b);
337  Value yieldValue =
338  insertSliceIntoTensor(rewriter, loc, sliceOp, sliceOp, iterArgs[0]);
339  return scf::ValueVector({yieldValue});
340  });
341  return success();
342 }
343 
344 namespace {
345 struct PadOpTilingPattern : public OpRewritePattern<tensor::PadOp> {
346  PadOpTilingPattern(MLIRContext *ctx, LinalgTilingOptions opt)
347  : OpRewritePattern<tensor::PadOp>(ctx), options(std::move(opt)) {}
348 
349  LogicalResult matchAndRewrite(tensor::PadOp op,
350  PatternRewriter &rewriter) const override {
352  return failure();
353  tensor::PadOp newPadOp;
354  LoopNest loopNest;
355  if (failed(tilePadOp(rewriter, op, newPadOp, loopNest, options)))
356  return failure();
358  rewriter.getUnitAttr());
359  // Replace all uses of the original tensor::PadOp.
360  rewriter.replaceOp(op, loopNest.getResults()[0]);
361  return success();
362  }
363 
365 };
366 } // namespace
367 
368 namespace {
369 /// Helper classes for type list expansion.
370 template <typename... OpTypes>
371 class CanonicalizationPatternList;
372 
373 template <>
374 class CanonicalizationPatternList<> {
375 public:
376  static void insert(RewritePatternSet &patterns) {}
377 };
378 
379 template <typename OpTy, typename... OpTypes>
380 class CanonicalizationPatternList<OpTy, OpTypes...> {
381 public:
382  static void insert(RewritePatternSet &patterns) {
383  OpTy::getCanonicalizationPatterns(patterns, patterns.getContext());
384  CanonicalizationPatternList<OpTypes...>::insert(patterns);
385  }
386 };
387 } // namespace
388 
391  RewritePatternSet patterns(ctx);
393  return patterns;
394 }
395 
397  RewritePatternSet &patterns) {
398  auto *ctx = patterns.getContext();
399  AffineApplyOp::getCanonicalizationPatterns(patterns, ctx);
400  AffineForOp::getCanonicalizationPatterns(patterns, ctx);
401  AffineMinOp::getCanonicalizationPatterns(patterns, ctx);
402  AffineMaxOp::getCanonicalizationPatterns(patterns, ctx);
403  arith::ConstantIndexOp::getCanonicalizationPatterns(patterns, ctx);
404 
405  memref::SubViewOp::getCanonicalizationPatterns(patterns, ctx);
406  memref::ViewOp::getCanonicalizationPatterns(patterns, ctx);
407 
408  scf::ForOp::getCanonicalizationPatterns(patterns, ctx);
409  scf::ParallelOp::getCanonicalizationPatterns(patterns, ctx);
410 
411  tensor::CastOp::getCanonicalizationPatterns(patterns, ctx);
412  tensor::ExtractSliceOp::getCanonicalizationPatterns(patterns, ctx);
413  tensor::InsertSliceOp::getCanonicalizationPatterns(patterns, ctx);
414 
415  InitTensorOp::getCanonicalizationPatterns(patterns, ctx);
416  tensor::PadOp::getCanonicalizationPatterns(patterns, ctx);
417  ctx->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(patterns);
418 
419  CanonicalizationPatternList<
420 #define GET_OP_LIST
421 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
422  >::insert(patterns);
423 }
424 
425 /// Populate the given list with patterns that apply Linalg tiling.
427  const LinalgTilingOptions &options) {
428  auto *ctx = patterns.getContext();
430  StringAttr::get(ctx, "tiled"));
431  TilingPatterns<GenericOp,
432 #define GET_OP_LIST
433 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
434  >::insert(patterns, options, f);
435  patterns.add<PadOpTilingPattern>(ctx, options);
436 }
437 
439  RewritePatternSet &patterns, const LinalgTilingOptions &options) {
440  auto *ctx = patterns.getContext();
441  patterns.add<PadOpTilingPattern>(ctx, options);
442 }
443 
444 static void applyExtractSliceOfPadTensorSwapPattern(func::FuncOp funcOp) {
445  MLIRContext *ctx = funcOp.getContext();
446  RewritePatternSet patterns(ctx);
447  patterns.add<ExtractSliceOfPadTensorSwapPattern>(patterns.getContext());
448  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
451 }
452 
453 namespace {
454 struct LinalgTilingPass : public LinalgTilingBase<LinalgTilingPass> {
455  LinalgTilingPass() = default;
456  LinalgTilingPass(ArrayRef<int64_t> tileSizes, LinalgTilingLoopType loopType) {
457  this->tileSizes = tileSizes;
458  this->loopType = "";
459  this->loopTypeEnum = loopType;
460  }
461 
462  void runOnOperation() override {
463  func::FuncOp funcOp = getOperation();
464  LinalgTilingLoopType type =
466  .Case("for", LinalgTilingLoopType::Loops)
467  .Case("affine", LinalgTilingLoopType::AffineLoops)
468  .Case("parallel", LinalgTilingLoopType::ParallelLoops)
469  .Default(loopTypeEnum);
470  auto options =
471  LinalgTilingOptions().setTileSizes(tileSizes).setLoopType(type);
472  MLIRContext *ctx = funcOp.getContext();
473  RewritePatternSet patterns(ctx);
474  insertTilingPatterns(patterns, options);
476  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
479  // Drop the marker.
480  funcOp.walk([](LinalgOp op) {
482  });
483 
484  // Apply swap pattern after generating loop nest and running
485  // canonicalizations.
487  }
488 
489  LinalgTilingLoopType loopTypeEnum;
490 };
491 
492 } // namespace
493 
494 std::unique_ptr<OperationPass<func::FuncOp>>
496  linalg::LinalgTilingLoopType loopType) {
497  return std::make_unique<LinalgTilingPass>(tileSizes, loopType);
498 }
Include the generated interface declarations.
Helper class to control application of linalg transformation patterns.
Definition: Transforms.h:382
AffineMap inversePermutation(AffineMap map)
Returns a map of codomain to domain dimensions such that the first codomain dimension for a particula...
Definition: AffineMap.cpp:658
MLIRContext * getContext() const
Definition: Builders.h:54
SmallVector< Value > computeTileSizes(OpBuilder &b, Location loc, ValueRange ivs, ValueRange tileSizes, ArrayRef< Value > sizeBounds)
Compute tile sizes, given a list of loop ivs, tileSizes and dimension sizes (sizeBounds).
Definition: Utils.cpp:896
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:600
Operation is a basic unit of execution within MLIR.
Definition: Operation.h:28
SmallVector< StringRef, 2 > distributionTypes
Specification markers of how to distribute the linalg.tiled_loop.
Definition: Transforms.h:613
void applyPermutationToVector(SmallVector< T, N > &inVec, ArrayRef< int64_t > permutation)
Apply the permutation defined by permutation to inVec.
Definition: IndexingUtils.h:38
Operation * getParentOp()
Returns the closest surrounding operation that contains this block.
Definition: Block.cpp:30
SmallVector< Value, 4 > tensorResults
Definition: Transforms.h:127
TileSizeComputationFunction tileSizeComputationFunction
Computation function that returns the tile sizes for each operation.
Definition: Transforms.h:563
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:329
Operation * clone(Operation &op, BlockAndValueMapping &mapper)
Creates a deep copy of the specified operation, remapping any operands that use values outside of the...
Definition: Builders.cpp:468
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value...
Definition: LogicalResult.h:72
static AffineMap getPermutationMap(ArrayRef< unsigned > permutation, MLIRContext *context)
Returns an AffineMap representing a permutation.
Definition: AffineMap.cpp:205
std::vector< Value > ValueVector
An owning vector of values, handy to return from functions.
Definition: SCF.h:55
Rewrite extract_slice(tensor.pad(x)) into tensor.pad(extract_slice(x)).
Definition: Transforms.h:1294
std::tuple< SmallVector< Range, 4 >, LoopIndexToRangeIndexMap > makeTiledLoopRanges(RewriterBase &b, Location loc, AffineMap map, ValueRange allShapeSizes, ValueRange allTileSizes)
Definition: Tiling.cpp:44
void populatePadTensorTilingPatterns(RewritePatternSet &patterns, const LinalgTilingOptions &options)
Definition: Tiling.cpp:438
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:48
Auxiliary range data structure to unpack the offset, size and stride operands into a list of triples...
LinalgTilingLoopType loopType
The type of tile loops to generate.
Definition: Transforms.h:595
Block * getOwner() const
Returns the block that owns this argument.
Definition: Value.h:309
std::unique_ptr< OperationPass< func::FuncOp > > createLinalgTilingPass(ArrayRef< int64_t > tileSizes={}, linalg::LinalgTilingLoopType loopType=linalg::LinalgTilingLoopType::Loops)
Definition: Tiling.cpp:495
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:380
Value makeTiledShape(OpBuilder &builder, Location loc, Value valueToTile, ValueRange tileSizes, AffineMap map, ValueRange lbs, ValueRange ubs, ValueRange subShapeSizes, bool omitPartialTileCheck)
Creates an extract_slice/subview op for a single valueToTile with builder.
Definition: Utils.cpp:764
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
virtual void replaceOp(Operation *op, ValueRange newValues)
This method replaces the results of the operation with the specified list of values.
Optional< LinalgLoopDistributionOptions > distribution
When specified, specifies distribution of generated tile loops to processors.
Definition: Transforms.h:604
static LogicalResult tilePadOp(RewriterBase &builder, tensor::PadOp op, tensor::PadOp &newPadOp, LoopNest &loopNest, const LinalgTilingOptions &options)
Generate a loop nest around a given tensor::PadOp (for tiling).
Definition: Tiling.cpp:284
This class provides support for representing a failure result, or a valid value of type T...
Definition: LogicalResult.h:77
static Value insertSliceIntoTensor(RewriterBase &b, Location loc, tensor::ExtractSliceOp sliceOp, Value source, Value dest)
Definition: Tiling.cpp:88
UnitAttr getUnitAttr()
Definition: Builders.cpp:85
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:234
static FailureOr< TiledLinalgOp > tileLinalgOpImpl(RewriterBase &b, LinalgOp op, ValueRange tileSizes, const LinalgTilingOptions &options)
Definition: Tiling.cpp:99
ResultRange getResults()
Definition: SCF.h:58
SmallVector< Value, 4 > applyMapToValues(OpBuilder &b, Location loc, AffineMap map, ValueRange values)
Returns the values obtained by applying map to the list of values.
Definition: AffineOps.cpp:736
unsigned getNumResults() const
Definition: AffineMap.cpp:302
SmallVector< unsigned, 4 > interchangeVector
The interchange vector to reorder the tiled loops.
Definition: Transforms.h:587
A multi-dimensional affine map Affine map&#39;s are immutable like Type&#39;s, and they are uniqued...
Definition: AffineMap.h:41
static void applyExtractSliceOfPadTensorSwapPattern(func::FuncOp funcOp)
Definition: Tiling.cpp:444
FailureOr< TiledLinalgOp > tileLinalgOp(RewriterBase &b, LinalgOp op, const LinalgTilingOptions &options)
Definition: Tiling.cpp:269
This class represents an argument of a Block.
Definition: Value.h:300
SmallVector< Value > computeTileOffsets(OpBuilder &b, Location loc, ValueRange ivs, ValueRange tileSizes)
Compute tile offsets, given a list of loop ivs and tileSizes.
Definition: Utils.cpp:881
void populateLinalgTilingCanonicalizationPatterns(RewritePatternSet &patterns)
Definition: Tiling.cpp:396
static void insertTilingPatterns(RewritePatternSet &patterns, const LinalgTilingOptions &options)
Populate the given list with patterns that apply Linalg tiling.
Definition: Tiling.cpp:426
void populateSCFForLoopCanonicalizationPatterns(RewritePatternSet &patterns)
Populate patterns for canonicalizing operations inside SCF loop bodies.
This class coordinates rewriting a piece of IR outside of a pattern rewrite, providing a way to keep ...
Definition: PatternMatch.h:584
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:85
LoopNest buildLoopNest(OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs, ValueRange steps, ValueRange iterArgs, function_ref< ValueVector(OpBuilder &, Location, ValueRange, ValueRange)> bodyBuilder=nullptr)
Creates a perfect nest of "for" loops, i.e.
Definition: SCF.cpp:504
static bool isZero(Value v)
Definition: Tiling.cpp:37
void transformIndexOps(RewriterBase &b, LinalgOp op, SmallVectorImpl< Value > &ivs, const LoopIndexToRangeIndexMap &loopIndexToRangeIndex)
All indices returned by IndexOp should be invariant with respect to tiling.
Definition: Tiling.cpp:72
static llvm::ManagedStatic< PassManagerOptions > options
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:355
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:279
LinalgTilingLoopType
The type of loops to be generated during tiling.
Definition: Utils.h:142
SmallVector< Value, 4 > makeTiledShapes(OpBuilder &builder, Location loc, LinalgOp linalgOp, ArrayRef< Value > valuesToTile, ValueRange ivs, ValueRange tileSizes, ArrayRef< Value > sizeBounds, bool omitPartialTileCheck)
Creates extract_slice/subview ops for all valuesToTile of the given linalgOp with builder...
Definition: Utils.cpp:911
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&... args)
Add an instance of each of the pattern types &#39;Ts&#39; to the pattern list with the given arguments...
static const StringLiteral kLinalgTransformMarker
Definition: Transforms.h:372
static void doit(OpBuilder &b, Location loc, ArrayRef< Range > loopRanges, LinalgOp linalgOp, ArrayRef< Attribute > iteratorTypes, function_ref< scf::ValueVector(OpBuilder &, Location, ValueRange, ValueRange)> bodyBuilderFn, Optional< LinalgLoopDistributionOptions >=None, ArrayRef< StringRef > distributionTypes={})
LinalgTilingOptions & setLoopType(LinalgTilingLoopType lt)
Definition: Transforms.h:597
Specialization of arith.constant op that returns an integer of index type.
Definition: Arithmetic.h:79
Perform standalone tiling of a single LinalgOp by tileSizes.
Definition: Transforms.h:124
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Definition: Value.cpp:20
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:55
This class represents an operand of an operation.
Definition: Value.h:251
static AffineMap getMultiDimIdentityMap(unsigned numDims, MLIRContext *context)
Returns an AffineMap with &#39;numDims&#39; identity result dim exprs.
Definition: AffineMap.cpp:244
RewritePatternSet getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx)
Canonicalization patterns relevant to apply after tiling patterns.
Definition: Tiling.cpp:390
LogicalResult applyPatternsAndFoldGreedily(MutableArrayRef< Region > regions, const FrozenRewritePatternSet &patterns, GreedyRewriteConfig config=GreedyRewriteConfig())
Rewrite the regions of the specified operation, which must be isolated from above, by repeatedly applying the highest benefit patterns in a greedy work-list driven manner.
result_range getResults()
Definition: Operation.h:339
This class helps build Operations.
Definition: Builders.h:177
This class provides an abstraction over the different types of ranges over Values.
MLIRContext * getContext() const
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:398
void addTileLoopIvsToIndexOpResults(OpBuilder &b, LinalgOp tiledOp, ArrayRef< Value > ivs)
Add the tile loop induction variables ivs to the IndexOp results found in the body of the tiledOp to ...
Definition: Utils.cpp:958
LinalgTilingOptions & setTileSizes(const SmallVector< Value, 4 > &ts)
Set the tileSizeComputationFunction to return the values ts.
Definition: Transforms.h:573