MLIR  19.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 
14 
26 #include "mlir/IR/AffineExpr.h"
27 #include "mlir/IR/AffineMap.h"
28 #include "mlir/IR/BuiltinOps.h"
29 #include "mlir/IR/ValueRange.h"
32 #include "llvm/ADT/STLExtras.h"
33 #include "llvm/Support/CommandLine.h"
34 #include <utility>
35 
36 namespace mlir {
37 #define GEN_PASS_DEF_LINALGTILINGPASS
38 #include "mlir/Dialect/Linalg/Passes.h.inc"
39 } // namespace mlir
40 
41 using namespace mlir;
42 using namespace mlir::affine;
43 using namespace mlir::linalg;
44 using namespace mlir::scf;
45 
46 #define DEBUG_TYPE "linalg-tiling"
47 
48 std::tuple<SmallVector<Range, 4>, LoopIndexToRangeIndexMap>
50  ArrayRef<OpFoldResult> allShapeSizes,
51  ArrayRef<OpFoldResult> allTileSizes) {
52  assert(allTileSizes.size() == map.getNumResults());
53  // Apply `map` to get shape sizes in loop order.
54  SmallVector<OpFoldResult> shapeSizes =
55  makeComposedFoldedMultiResultAffineApply(b, loc, map, allShapeSizes);
56  SmallVector<OpFoldResult> tileSizes(allTileSizes.begin(), allTileSizes.end());
57 
58  // Traverse the tile sizes, which are in loop order, erase zeros everywhere.
59  LoopIndexToRangeIndexMap loopIndexToRangeIndex;
60  for (int idx = 0, e = tileSizes.size(), zerosCount = 0; idx < e; ++idx) {
61  if (getConstantIntValue(tileSizes[idx - zerosCount]) ==
62  static_cast<int64_t>(0)) {
63  shapeSizes.erase(shapeSizes.begin() + idx - zerosCount);
64  tileSizes.erase(tileSizes.begin() + idx - zerosCount);
65  ++zerosCount;
66  continue;
67  }
68  loopIndexToRangeIndex[idx] = idx - zerosCount;
69  }
70 
71  // Create a new range with the applied tile sizes.
73  for (unsigned idx = 0, e = tileSizes.size(); idx < e; ++idx)
74  res.push_back(Range{b.getIndexAttr(0), shapeSizes[idx], tileSizes[idx]});
75  return std::make_tuple(res, loopIndexToRangeIndex);
76 }
77 
79  RewriterBase &b, LinalgOp op, SmallVectorImpl<Value> &ivs,
80  const LoopIndexToRangeIndexMap &loopIndexToRangeIndex) {
81  SmallVector<Value> allIvs(op.getNumLoops(), nullptr);
82  for (auto en : enumerate(allIvs)) {
83  auto rangeIndex = loopIndexToRangeIndex.find(en.index());
84  if (rangeIndex == loopIndexToRangeIndex.end())
85  continue;
86  en.value() = ivs[rangeIndex->second];
87  }
88  offsetIndices(b, op, getAsOpFoldResult(allIvs));
89 }
90 
91 /// Asserts that the given index-typed value is strictly positive. If the value
92 /// is an attribute, asserts at compile time, otherwise emits an assertion
93 /// checked at runtime.
95  OpFoldResult value) {
96  if (auto attr = llvm::dyn_cast_if_present<Attribute>(value)) {
97  assert(cast<IntegerAttr>(attr).getValue().isStrictlyPositive() &&
98  "expected strictly positive tile size and divisor");
99  return;
100  }
101 
102  Value zero = b.create<arith::ConstantIndexOp>(0);
103  Value condition = b.create<arith::CmpIOp>(arith::CmpIPredicate::sgt,
104  value.get<Value>(), zero);
105  b.create<cf::AssertOp>(
106  condition,
107  b.getStringAttr("expected strictly positive tile size and divisor"));
108 }
109 
111 mlir::linalg::computeStaticMultiTileSizes(LinalgOp op, unsigned dimension,
112  int64_t targetSize, int64_t divisor) {
113  assert(!op.hasDynamicShape() &&
114  "cannot compute static multi-tile sizes for an op with dynamic shape");
115  assert(targetSize > 0 && "target size must be non-negative");
116  assert(divisor > 0 && "divisor must be non-negative");
117  assert(dimension < op.getNumLoops() && "dimension overflow");
118 
120  int64_t tripCount = op.getStaticLoopRanges()[dimension];
121  int64_t a = tripCount / divisor;
122  int64_t t = (targetSize + divisor - 1) / divisor;
123  int64_t totalTripCount = (a + t - 1) / t;
124  spec.lowTileSize = (a / totalTripCount) * divisor;
125  spec.highTileSize = spec.lowTileSize + divisor;
126  spec.highTripCount = a % totalTripCount;
127  spec.lowTripCount = totalTripCount - spec.highTripCount;
128  if (spec.lowTileSize * spec.lowTripCount +
129  spec.highTileSize * spec.highTripCount !=
130  tripCount) {
131  return failure();
132  }
133  return spec;
134 }
135 
138  unsigned dimension, OpFoldResult targetSize,
139  OpFoldResult divisor, bool emitAssertions) {
140  // Bail out on dimension overflow.
141  if (dimension >= op.getNumLoops())
142  return failure();
143 
144  // The code below works only on values.
145  Location loc = op.getLoc();
146  ImplicitLocOpBuilder b(loc, builder);
147  if (emitAssertions) {
148  emitIsPositiveIndexAssertion(b, targetSize);
149  emitIsPositiveIndexAssertion(b, divisor);
150  }
151  Value targetSizeValue =
152  getValueOrCreateConstantIndexOp(builder, loc, targetSize);
153  Value divisorValue = getValueOrCreateConstantIndexOp(builder, loc, divisor);
154 
155  // Find the trip count of the iteration space dimension for which the tile
156  // sizes are computed.
157  SmallVector<OpFoldResult> allShapes =
158  op.createFlatListOfOperandDims(b, b.getLoc());
159  AffineMap shapesToLoops = op.getShapesToLoopsMap();
160  SmallVector<OpFoldResult> loopRanges =
161  makeComposedFoldedMultiResultAffineApply(b, op.getLoc(), shapesToLoops,
162  allShapes);
163  Value tripCount =
164  getValueOrCreateConstantIndexOp(b, op.getLoc(), loopRanges[dimension]);
165 
166  // Compute the tile sizes and the respective numbers of tiles.
170  auto apply = [&](AffineExpr expr, ArrayRef<OpFoldResult> ofrs) -> Value {
171  return affine::makeComposedAffineApply(b, b.getLoc(), expr, ofrs);
172  };
173  Value a = apply(s0.floorDiv(s1), {tripCount, divisorValue});
174  Value t = apply((s0 + s1 - 1).floorDiv(s1), {targetSizeValue, divisorValue});
175  Value d = apply((s0 + s1 - 1).floorDiv(s1), {a, t});
176  Value s = apply(s0.floorDiv(s1) * s2, {a, d, divisorValue});
177  Value v = apply(s0 % s1, {a, d});
178  Value u = apply(s0 - s1, {d, v});
179 
181  spec.lowTileSize = s;
182  spec.highTileSize = apply(s0 + s1, {s, divisorValue});
183  spec.lowTripCount = u;
184  spec.highTripCount = v;
185 
186  // If requested, emit the check that the tile sizes are computed correctly.
187  // For example, for iteration dimension size of 15 and the target size 8 it is
188  // impossible to find two tile sizes both divisible by 8 that fully cover the
189  // original space dimension.
190  if (emitAssertions) {
191  AffineExpr s3 = builder.getAffineSymbolExpr(3);
192  Value coveredSize =
193  apply(s0 * s1 + s2 * s3, {spec.lowTileSize, spec.lowTripCount,
194  spec.highTileSize, spec.highTripCount});
195  Value equals = b.create<arith::CmpIOp>(arith::CmpIPredicate::eq,
196  coveredSize, tripCount);
197  b.create<cf::AssertOp>(
198  equals, builder.getStringAttr(
199  "could not compute dynamic multi-size tile shapes"));
200  }
201 
202  return spec;
203 }
204 
205 /// Returns true if the maximum tile offset `tileSize * numThreads-1` is less
206 /// than `iterationSize`.
208  OpFoldResult numThreads,
209  OpFoldResult iterationSize) {
210  std::optional<int64_t> tileSizeConst = getConstantIntValue(tileSize);
211  std::optional<int64_t> numThreadsConst = getConstantIntValue(numThreads);
212  std::optional<int64_t> iterSizeConst = getConstantIntValue(iterationSize);
213  if (!tileSizeConst || !numThreadsConst || !iterSizeConst)
214  return false;
215  return *tileSizeConst * (*numThreadsConst - 1) < *iterSizeConst;
216 }
217 
218 /// Build an `affine_max` of all the `vals`.
220  ArrayRef<OpFoldResult> vals) {
222  b, loc, AffineMap::getMultiDimIdentityMap(vals.size(), loc.getContext()),
223  vals);
224 }
225 
226 /// Build an `affine_min` of all the `vals`.
228  ArrayRef<OpFoldResult> vals) {
230  b, loc, AffineMap::getMultiDimIdentityMap(vals.size(), loc.getContext()),
231  vals);
232 }
233 
234 /// Fill out the `tiledOffsets` and `tiledSizes` to be used to tile to a given
235 /// number of threads.
237  RewriterBase &b, Location loc, scf::ForallOp forallOp,
238  ArrayRef<OpFoldResult> numThreads, SmallVector<Range> loopRanges,
239  bool omitTileOffsetBoundsCheck,
240  std::optional<ArrayRef<OpFoldResult>> nominalTileSizes,
241  SmallVector<OpFoldResult> &tiledOffsets,
242  SmallVector<OpFoldResult> &tiledSizes) {
244  b.setInsertionPointToStart(forallOp.getBody(0));
245 
246  ValueRange threadIds = forallOp.getInductionVars();
247  SmallVector<OpFoldResult> nonZeroNumThreads =
248  llvm::to_vector(llvm::make_filter_range(numThreads, [](OpFoldResult ofr) {
249  return !isConstantIntValue(ofr, 0);
250  }));
251  int64_t nLoops = loopRanges.size();
252  tiledOffsets.reserve(nLoops);
253  tiledSizes.reserve(nLoops);
254  for (unsigned loopIdx = 0, threadIdIdx = 0; loopIdx < nLoops; ++loopIdx) {
255  bool overflow = loopIdx >= numThreads.size();
256  bool isZero = !overflow && isConstantIntValue(numThreads[loopIdx], 0);
257  // Degenerate case: take the whole domain.
258  if (overflow || isZero) {
259  tiledOffsets.push_back(loopRanges[loopIdx].offset);
260  tiledSizes.push_back(loopRanges[loopIdx].size);
261  continue;
262  }
263 
264  // Tiled case: compute the offset and size.
265  AffineExpr i, j, m, n, o;
266  bindDims(b.getContext(), i, j);
267  bindSymbols(b.getContext(), m, n, o);
268  OpFoldResult size = loopRanges[loopIdx].size;
269  OpFoldResult offset = loopRanges[loopIdx].offset;
270  OpFoldResult threadId = threadIds[threadIdIdx];
271  // Symbolic fixed max size per thread.
272  // TODO: floor + 0/1 depending on case for better load-balancing.
273  OpFoldResult tileSizePerThread =
274  nominalTileSizes.has_value()
275  ? (*nominalTileSizes)[loopIdx]
277  b, loc, m.ceilDiv(n),
278  ArrayRef<OpFoldResult>{size, nonZeroNumThreads[threadIdIdx]});
279 
280  // Dynamic offset shifted by threadId * maxSizePerThread.
282  b, loc, i + j * m, {offset, threadId, tileSizePerThread});
283  // Dynamic upper-bound depending on the threadId.
284  OpFoldResult residualTileSize = makeComposedFoldedAffineApply(
285  b, loc, i + j * m - n,
286  {offset, nonZeroNumThreads[threadIdIdx], tileSizePerThread, size});
287  if (!isConstantIntValue(residualTileSize, 0)) {
288  OpFoldResult sizeMinusOffsetPerThread = makeComposedFoldedAffineApply(
289  b, loc, -i + m, {offsetPerThread, size});
290  tileSizePerThread =
291  buildMin(b, loc, {sizeMinusOffsetPerThread, tileSizePerThread});
292  }
293 
294  tiledOffsets.push_back(offsetPerThread);
295  // TODO: if tileSizePerThread <= 0 early exit.
296  if (!omitTileOffsetBoundsCheck &&
297  !canOmitTileOffsetInBoundsCheck(tileSizePerThread,
298  nonZeroNumThreads[threadIdIdx], size))
299  tileSizePerThread =
300  buildMax(b, loc, {b.getIndexAttr(0), tileSizePerThread});
301 
302  tiledSizes.push_back(tileSizePerThread);
303  ++threadIdIdx;
304  }
305 }
306 
307 /// Rewrite a TilingInterface `op` to a tiled `scf.forall`. The
308 /// tiling is specified by the number of tiles/threads `numThreads` and the
309 /// optional nominal tile size `nominalTileSizes`. If `nominalTilSizes` is
310 /// not specified, then it is derived from `numThreads` as `ceilDiv(dimSize[i],
311 /// numThreads[i])`. If non-empty, the `mapping` is added as an
312 /// attribute to the resulting `scf.forall`. A zero tile sizes indicate
313 /// that the dimension is not tiled, and can be thought of as tiling by the full
314 /// size of data.
315 /// It is the user's responsibility to ensure that `numThreads` is a valid
316 /// tiling specification (i.e. that only tiles parallel dimensions, e.g. in the
317 /// Linalg case). If `omitTileOffsetBoundsCheck` is true, then the function will
318 /// assume that `tileSize[i] * (numThread[i] -1) <= dimSize[i]` holds.
320  RewriterBase &b, TilingInterface op, ArrayRef<OpFoldResult> numThreads,
321  std::optional<ArrayRef<OpFoldResult>> nominalTileSizes,
322  std::optional<ArrayAttr> mapping, bool omitTileOffsetBoundsCheck) {
323  Location loc = op->getLoc();
325 
326  SmallVector<Range> loopRanges = op.getIterationDomain(b);
327  if (loopRanges.empty())
328  return op->emitOpError("expected non-empty loop ranges");
329  auto hasStrideOne = [](Range r) { return !isConstantIntValue(r.stride, 1); };
330  if (llvm::any_of(loopRanges, hasStrideOne))
331  return op->emitOpError("only stride-1 supported atm");
332 
333  // Gather destination tensors.
334  SmallVector<Value> dest;
335  if (failed(tensor::getOrCreateDestinations(b, loc, op, dest)))
336  return op->emitOpError("failed to get destination tensors");
337 
338  SmallVector<OpFoldResult> nonZeroNumThreads =
339  llvm::to_vector(llvm::make_filter_range(numThreads, [](OpFoldResult ofr) {
340  return !isConstantIntValue(ofr, 0);
341  }));
342  SmallVector<Value> materializedNonZeroNumThreads =
343  llvm::to_vector(llvm::map_range(nonZeroNumThreads, [&](OpFoldResult ofr) {
344  return getValueOrCreateConstantIndexOp(b, loc, ofr);
345  }));
346 
347  // 1. Create the ForallOp. We don't use the lambda body-builder
348  // version because we require the use of RewriterBase in the body, so we
349  // manually move the insertion point to the body below.
350  scf::ForallOp forallOp = b.create<scf::ForallOp>(
351  loc, getAsOpFoldResult((materializedNonZeroNumThreads)), dest, mapping);
352 
353  // 2. Fill out the ForallOp body.
354  SmallVector<OpFoldResult> tiledOffsets, tiledSizes;
355  calculateTileOffsetsAndSizes(b, loc, forallOp, numThreads, loopRanges,
356  omitTileOffsetBoundsCheck, nominalTileSizes,
357  tiledOffsets, tiledSizes);
358 
359  // 3. Clone the tileable op and update its destination operands to use the
360  // output bbArgs of the ForallOp.
361  ArrayRef<BlockArgument> destBbArgs = forallOp.getRegionIterArgs();
362  Operation *tiledOp = nullptr;
363  SmallVector<Value> tiledValues;
364  {
365  // 3.a. RAII guard, inserting within forallOp, before terminator.
367  b.setInsertionPoint(forallOp.getTerminator());
368  Operation *clonedOp = b.clone(*op.getOperation());
369  auto destinationStyleOp = dyn_cast<DestinationStyleOpInterface>(clonedOp);
370  if (destinationStyleOp) {
371  for (OpOperand &outOperand : destinationStyleOp.getDpsInitsMutable()) {
372  // Swap tensor inits with the corresponding block argument of the
373  // scf.forall op. Memref inits remain as is.
374  if (outOperand.get().getType().isa<TensorType>()) {
375  auto *it = llvm::find(dest, outOperand.get());
376  assert(it != dest.end() && "could not find destination tensor");
377  unsigned destNum = std::distance(dest.begin(), it);
378  outOperand.set(destBbArgs[destNum]);
379  }
380  }
381  }
382 
383  // 4. Tile the cloned op and delete the clone.
384  FailureOr<TilingResult> tilingResult =
385  cast<TilingInterface>(clonedOp).getTiledImplementation(b, tiledOffsets,
386  tiledSizes);
387  if (failed(tilingResult))
388  return clonedOp->emitError("Failed to tile op: ");
389  if (tilingResult->tiledOps.size() != 1) {
390  return clonedOp->emitError("expected a single produced tiled op, got ")
391  << tilingResult->tiledOps.size();
392  }
393 
394  b.eraseOp(clonedOp);
395  tiledOp = tilingResult->tiledOps.front();
396  tiledValues = tilingResult->tiledValues;
397  }
398 
399  // 5. Parallel insert back into the result tensor.
400  for (auto it : llvm::zip(llvm::seq(unsigned(0), unsigned(dest.size())),
401  tiledValues, destBbArgs)) {
402  // 5.a. Partial subset information is inserted just before the terminator.
404  b.setInsertionPoint(forallOp.getTerminator());
405 
406  SmallVector<OpFoldResult> resultOffsets, resultSizes;
407  if (failed(op.getResultTilePosition(b, std::get<0>(it), tiledOffsets,
408  tiledSizes, resultOffsets,
409  resultSizes)))
410  return op->emitOpError("output offsets couldn't be calculated");
411  SmallVector<OpFoldResult> strides(resultSizes.size(), b.getIndexAttr(1));
412 
413  // 5.b. Parallel insertions are inserted at the end of the combining
414  // terminator.
415  b.setInsertionPointToEnd(forallOp.getTerminator().getBody());
416  b.create<tensor::ParallelInsertSliceOp>(loc, std::get<1>(it),
417  std::get<2>(it), resultOffsets,
418  resultSizes, strides);
419  }
420  return ForallTilingResult{forallOp, tiledOp};
421 }
422 
424 linalg::tileToForallOp(RewriterBase &b, TilingInterface op,
425  ArrayRef<OpFoldResult> numThreads,
426  std::optional<ArrayAttr> mapping) {
427  return tileToForallOpImpl(b, op, numThreads,
428  /*nominalTileSizes=*/std::nullopt, mapping,
429  /*omitTileOffsetBoundsCheck=*/false);
430 }
431 
434  ArrayRef<OpFoldResult> tileSizes,
435  std::optional<ArrayAttr> mapping) {
436  SmallVector<Range> loopRanges = op.getIterationDomain(b);
437  unsigned nLoops = loopRanges.size();
438  SmallVector<OpFoldResult> numThreads;
439  numThreads.reserve(nLoops);
440  AffineExpr s0, s1;
441  bindSymbols(b.getContext(), s0, s1);
442  AffineExpr divExpr = s0.ceilDiv(s1);
443  for (const auto &it : llvm::zip(tileSizes, loopRanges)) {
444  OpFoldResult numTiles = std::get<0>(it);
445  if (!isConstantIntValue(numTiles, 0))
447  b, op.getLoc(), divExpr, {std::get<1>(it).size, std::get<0>(it)});
448  numThreads.push_back(numTiles);
449  }
450  return tileToForallOpImpl(b, op, numThreads,
451  /*nominalTileSizes=*/tileSizes, mapping,
452  /*omitTileOffsetBoundsCheck=*/true);
453 }
454 
455 template <typename LoopTy>
458  const LinalgTilingOptions &options) {
460 
461  auto nLoops = op.getNumLoops();
462  // Initial tile sizes may be too big, only take the first nLoops.
463  tileSizes = tileSizes.take_front(nLoops);
464 
465  if (llvm::all_of(tileSizes, [](OpFoldResult ofr) {
466  return getConstantIntValue(ofr) == static_cast<int64_t>(0);
467  })) {
468  TiledLinalgOp tiledOp;
469  tiledOp.op = cast<LinalgOp>(b.clone(*op.getOperation()));
470  tiledOp.tensorResults.assign(tiledOp.op->result_begin(),
471  tiledOp.op->result_end());
472  return tiledOp;
473  }
474 
475  // 1. Build the tiled loop ranges.
476  SmallVector<OpFoldResult> allShapeSizes =
477  op.createFlatListOfOperandDims(b, op.getLoc());
478  AffineMap shapeSizesToLoopsMap = op.getShapesToLoopsMap();
479  if (!shapeSizesToLoopsMap)
480  return failure();
481 
482  auto [loopRanges, loopIndexToRangeIndex] = makeTiledLoopRanges(
483  b, op.getLoc(), shapeSizesToLoopsMap, allShapeSizes, tileSizes);
484 
486  for (const auto &attr : enumerate(op.getIteratorTypesArray())) {
487  if (loopIndexToRangeIndex.count(attr.index()))
488  iteratorTypes.push_back(attr.value());
489  }
490  // If interchangeVector is empty, use the identity. Build the permutation map
491  // otherwise.
492  auto invPermutationMap =
493  AffineMap::getMultiDimIdentityMap(tileSizes.size(), b.getContext());
494  if (!options.interchangeVector.empty()) {
495  // Based on the pruned iterations (due to zero tile size), recompute the
496  // interchange vector.
497  SmallVector<unsigned, 4> interchangeVector;
498  interchangeVector.reserve(options.interchangeVector.size());
499  for (auto pos : options.interchangeVector) {
500  auto it = loopIndexToRangeIndex.find(pos);
501  if (it == loopIndexToRangeIndex.end())
502  continue;
503  interchangeVector.push_back(it->second);
504  }
505  // Interchange vector is guaranteed to be a permutation,
506  // `inversePermutation` must succeed.
507  invPermutationMap = inversePermutation(
508  AffineMap::getPermutationMap(interchangeVector, b.getContext()));
509  assert(invPermutationMap);
510  SmallVector<int64_t> permutation(interchangeVector.begin(),
511  interchangeVector.end());
512  applyPermutationToVector(loopRanges, permutation);
513  applyPermutationToVector(iteratorTypes, permutation);
514  }
515 
516  // Handle distribution. Create a vector of the same size of loops that are to
517  // be tiled.
519  if (options.distribution) {
520  procInfo.resize(
521  iteratorTypes.size(),
522  linalg::ProcInfo{nullptr, nullptr, linalg::DistributionMethod::None});
523  // Collect loop ranges of tiled loops, loops that are parallel.
524  SmallVector<Range> parallelLoopRanges;
525  for (const auto &iteratorType : llvm::enumerate(iteratorTypes)) {
526  if (!isParallelIterator(iteratorType.value()))
527  break;
528  parallelLoopRanges.push_back(loopRanges[iteratorType.index()]);
529  }
530  auto returnedProcInfo =
531  options.distribution->procInfo(b, op.getLoc(), parallelLoopRanges);
532  unsigned procIdIdx = 0;
533  // Update the distribution information for the loops.
534  for (const auto &iteratorType : llvm::enumerate(iteratorTypes)) {
535  if (!isParallelIterator(iteratorType.value()))
536  break;
537  procInfo[iteratorType.index()] = returnedProcInfo[procIdIdx++];
538  }
539  }
540 
541  // 2. Create the tiled loops.
542  LinalgOp res = op;
543  SmallVector<Value, 4> ivs, tensorResults;
544  auto tiledLoopBodyBuilder =
545  [&](OpBuilder &builder, Location loc, ValueRange localIvs,
546  ValueRange operandValuesToUse) -> scf::ValueVector {
547  ivs.assign(localIvs.begin(), localIvs.end());
548 
549  // When an `interchangeVector` is present, it has been applied to the
550  // loop ranges and the iterator types. Apply its inverse to the
551  // resulting loop `ivs` to match the op definition.
552  SmallVector<Value, 4> interchangedIvs;
553  if (!options.interchangeVector.empty()) {
554  for (AffineExpr result : invPermutationMap.getResults())
555  interchangedIvs.push_back(
556  ivs[cast<AffineDimExpr>(result).getPosition()]);
557  } else {
558  interchangedIvs.assign(ivs.begin(), ivs.end());
559  }
560 
561  // Tile the `operandValuesToUse` that either match the `op` operands
562  // themselves or the tile loop arguments forwarding them.
563  assert(operandValuesToUse.size() ==
564  static_cast<size_t>(op->getNumOperands()) &&
565  "expect the number of operands and inputs and outputs to match");
566  SmallVector<Value> valuesToTile = operandValuesToUse;
567  SmallVector<OpFoldResult> sizeBounds =
568  makeComposedFoldedMultiResultAffineApply(b, loc, shapeSizesToLoopsMap,
569  allShapeSizes);
570  SmallVector<Value> tiledOperands = makeTiledShapes(
571  b, loc, op, valuesToTile, getAsOpFoldResult(interchangedIvs), tileSizes,
572  sizeBounds,
573  /*omitPartialTileCheck=*/false);
574 
575  SmallVector<Type> resultTensorTypes =
576  getTensorOutputTypes(op, tiledOperands);
577  res = clone(b, op, resultTensorTypes, tiledOperands);
578  tensorResults =
579  insertSlicesBack(builder, loc, op, tiledOperands, res->getResults());
580  return scf::ValueVector(tensorResults.begin(), tensorResults.end());
581  };
582  GenerateLoopNest<LoopTy>::doit(b, op.getLoc(), loopRanges, op, iteratorTypes,
583  tiledLoopBodyBuilder, procInfo);
584 
585  // 3. Transform IndexOp results w.r.t. the tiling.
586  transformIndexOps(b, res, ivs, loopIndexToRangeIndex);
587 
588  // 4. Gather the newly created loops and return them with the new op.
590  loops.reserve(ivs.size());
591  for (auto iv : ivs) {
592  if (isa<BlockArgument>(iv)) {
593  loops.push_back(cast<BlockArgument>(iv).getOwner()->getParentOp());
594  assert(loops.back() && "no owner found for induction variable!");
595  } else {
596  // TODO: Instead of doing this, try to recover the ops used instead of the
597  // loop.
598  loops.push_back(nullptr);
599  }
600  }
601 
602  // 5. Get the tensor results from the outermost loop if available. Otherwise
603  // use the previously captured `tensorResults`.
604  Operation *outermostLoop = nullptr;
605  for (Operation *loop : loops)
606  if ((outermostLoop = loop))
607  break;
608 
609  return TiledLinalgOp{
610  res, loops, outermostLoop ? outermostLoop->getResults() : tensorResults};
611 }
612 
614  RewriterBase &b, PartialReductionOpInterface op,
615  ArrayRef<OpFoldResult> numThreads, ArrayRef<OpFoldResult> tileSizes,
616  std::optional<ArrayAttr> mapping) {
617  Location loc = op.getLoc();
619 
620  // Ops implementing PartialReductionOpInterface are expected to implement
621  // TilingInterface.
622  // TODO: proper core mechanism to tie interfaces together.
623  auto tilingInterfaceOp = cast<TilingInterface>(op.getOperation());
624 
625  // Ops implementing PartialReductionOpInterface are not necessarily expected
626  // to implement TilingInterface.. This cast is unsafe atm.
627  // TODO: proper core mechanism to tie interfaces together.
628  // TODO: this function requires a pair of interfaces ..
629  auto destinationStyleOp =
630  dyn_cast<DestinationStyleOpInterface>(op.getOperation());
631  if (!destinationStyleOp)
632  return b.notifyMatchFailure(op, "not a destination style op");
633 
634  // Actually this only work for Linalg ops atm.
635  auto linalgOp = dyn_cast<linalg::LinalgOp>(op.getOperation());
636  if (!linalgOp)
637  return b.notifyMatchFailure(op, "not a linalg op");
638 
639  SmallVector<Range> iterationDomain = tilingInterfaceOp.getIterationDomain(b);
640  if (op->getNumResults() != 1)
641  return b.notifyMatchFailure(
642  op, "don't support ops with multiple results for now");
643 
645  tilingInterfaceOp.getLoopIteratorTypes();
646  SmallVector<unsigned> redDims;
647  linalgOp.getReductionDims(redDims);
648  if (redDims.size() != 1)
649  return b.notifyMatchFailure(
650  op, "only support ops with one reduction dimension.");
651  if (!tileSizes.empty() && tileSizes.size() != numThreads.size())
652  return b.notifyMatchFailure(op, "if tile sizes are present it must have as "
653  "many elements as number of threads");
654  int reductionDim = static_cast<int>(redDims.front());
655 
656  if (redDims.front() >= numThreads.size())
657  return b.notifyMatchFailure(
658  op, "reduction dimension must be mapped to threads");
659 
660  // 1. Create the inital tensor value.
661  FailureOr<Operation *> identityTensor =
662  op.generateInitialTensorForPartialReduction(b, loc, numThreads,
663  reductionDim);
664  if (failed(identityTensor))
665  return b.notifyMatchFailure(op,
666  "cannot create a tensor of identity value.");
667 
668  // Gather destination tensors.
669  SmallVector<Value> dest;
670  if (failed(tensor::getOrCreateDestinations(b, loc, op, dest)))
671  return b.notifyMatchFailure(op, "failed to get destination tensors");
672 
673  Operation *tiledOp = nullptr;
674 
675  SmallVector<OpFoldResult> nonZeroNumThreads =
676  llvm::to_vector(llvm::make_filter_range(numThreads, [](OpFoldResult ofr) {
677  return !isConstantIntValue(ofr, 0);
678  }));
679  SmallVector<Value> materializedNonZeroNumThreads =
680  getValueOrCreateConstantIndexOp(b, loc, nonZeroNumThreads);
681 
682  // 2. Create the ForallOp with an empty region.
683  scf::ForallOp forallOp = b.create<scf::ForallOp>(
684  loc, getAsOpFoldResult(materializedNonZeroNumThreads),
685  (*identityTensor)->getResults(), mapping);
686 
687  // 3. Calculate the tile offsets and sizes for the subsequent loop that will
688  // be nested under `forallOp`.
689  SmallVector<OpFoldResult> tiledOffsets, tiledSizes;
690  calculateTileOffsetsAndSizes(b, loc, forallOp, numThreads, iterationDomain,
691  /*omitTileOffsetBoundsCheck =*/false,
692  /*nominalTileSizes=*/std::nullopt, tiledOffsets,
693  tiledSizes);
694 
695  // 4. Clone the tileable op and update its destination operands to use the
696  // output bbArgs of the ForallOp.
697  SmallVector<Value> tilingResults;
698  ArrayRef<BlockArgument> destBbArgs = forallOp.getRegionIterArgs();
699  {
700  // 4.a. RAII guard, inserting within forallOp, before terminator.
702  b.setInsertionPoint(forallOp.getTerminator());
703 
704  SmallVector<Value> tiledDpsInitOperands;
705  for (Value initOperand : destinationStyleOp.getDpsInits()) {
706  auto *it = llvm::find(dest, initOperand);
707  assert(it != dest.end() && "dest operand not found in dest");
708  unsigned destNum = std::distance(dest.begin(), it);
709  SmallVector<OpFoldResult> strides(numThreads.size(), b.getIndexAttr(1));
710  SmallVector<OpFoldResult> outOffsets(numThreads.size(),
711  b.getIndexAttr(0));
712  SmallVector<OpFoldResult> sizes = tiledSizes;
713  sizes[reductionDim] = b.getIndexAttr(1);
714  outOffsets[reductionDim] = forallOp.getInductionVars().front();
715  // TODO: use SubsetExtractOpInterface once it is available.
716  tiledDpsInitOperands.push_back(b.create<tensor::ExtractSliceOp>(
717  loc, cast<RankedTensorType>(initOperand.getType()),
718  destBbArgs[destNum], outOffsets, sizes, strides));
719  }
720 
721  // 4.b. Clone the op and update init operands.
722  // We cannot use a IRMapping here because it can replace
723  // different OpOperands with the same value.
724  Operation *clonedOp = b.clone(*op.getOperation());
725  b.modifyOpInPlace(clonedOp, [&]() {
726  for (auto [initOperandPtr, tiledInitValue] : llvm::zip_equal(
727  cast<DestinationStyleOpInterface>(clonedOp).getDpsInitsMutable(),
728  tiledDpsInitOperands)) {
729  initOperandPtr.set(tiledInitValue);
730  }
731  });
732 
733  // 5. Tile the cloned op and delete the clone.
734  if (tileSizes.empty()) {
735  FailureOr<TilingResult> tilingResult =
736  cast<TilingInterface>(clonedOp).getTiledImplementation(
737  b, tiledOffsets, tiledSizes);
738  if (failed(tilingResult))
739  return clonedOp->emitError("Failed to tile op: ");
740  if (tilingResult->tiledOps.size() != 1) {
741  return clonedOp->emitError("expected a single produced tiled op, got ")
742  << tilingResult->tiledOps.size();
743  }
744  tiledOp = tilingResult->tiledOps.front();
745  tilingResults = tilingResult->tiledValues;
746  } else {
748  FailureOr<TiledLinalgOp> maybeTiled = tileLinalgOpImpl<scf::ForOp>(
749  b, cast<LinalgOp>(clonedOp), tileSizes, options);
750  if (failed(maybeTiled))
751  return b.notifyMatchFailure(op, "failed tileLinalgOpImpl");
752 
753  SmallVector<Value> ids = forallOp.getInductionVars();
754  mapLoopToProcessorIds(cast<scf::ForOp>(maybeTiled->loops.back()), ids,
755  materializedNonZeroNumThreads);
756  if (maybeTiled->loops.size() != 1) {
757  return clonedOp->emitError("expected a single produced loop");
758  }
759  tiledOp = maybeTiled->op;
760  tilingResults = maybeTiled->loops.front()->getResults();
761  }
762 
763  b.eraseOp(clonedOp);
764  }
765 
766  // 6. Insert the partial reductions back into a new tensor.
767  for (auto [index, result, bbArg] : llvm::zip(
768  llvm::seq<unsigned>(0, dest.size()), tilingResults, destBbArgs)) {
769  // 6.a. Partial subset information is inserted just before the terminator.
771  b.setInsertionPoint(forallOp.getTerminator());
772 
773  SmallVector<OpFoldResult> resultOffsets, resultSizes;
774  if (failed(tilingInterfaceOp.getResultTilePosition(
775  b, index, tiledOffsets, tiledSizes, resultOffsets, resultSizes)))
776  return op->emitOpError("output offsets couldn't be calculated");
777  SmallVector<OpFoldResult> resultOffsetsRank, resultSizesRank;
778  int64_t offIdx = 0;
779  int64_t sizeIdx = 0;
780  for (int64_t i = 0, e = numThreads.size(); i < e; ++i) {
781  if (i == reductionDim) {
782  resultOffsetsRank.push_back(forallOp.getInductionVars().front());
783  resultSizesRank.push_back(b.getIndexAttr(1));
784  continue;
785  }
786  resultOffsetsRank.push_back(resultOffsets[offIdx++]);
787  resultSizesRank.push_back(resultSizes[sizeIdx++]);
788  }
789  SmallVector<OpFoldResult> strides(resultSizesRank.size(),
790  b.getIndexAttr(1));
791 
792  // 6.b. Parallel insertions are inserted at the end of the combining
793  // terminator.
794  b.setInsertionPointToEnd(forallOp.getTerminator().getBody());
795  b.create<tensor::ParallelInsertSliceOp>(
796  loc, result, bbArg, resultOffsetsRank, resultSizesRank, strides);
797  }
798 
799  // 7. Merge the partial reductions.
800  b.setInsertionPointAfter(forallOp);
801  Operation *mergeOp =
802  op.mergeReductions(b, loc, forallOp->getResults(), reductionDim);
803  b.replaceOp(op, mergeOp->getResults());
804 
805  // 8. Return.
807  results.initialOp = *identityTensor;
808  results.loops = forallOp;
809  results.parallelTiledOp = tiledOp;
810  results.mergeOp = mergeOp;
811  return results;
812 }
813 
814 template <typename LoopTy>
816  RewriterBase &b, LinalgOp op, const LinalgTilingOptions &options) {
818  b.setInsertionPoint(op);
819 
820  if (!options.tileSizeComputationFunction)
821  return failure();
822 
823  // Enforce the convention that "tiling by zero" skips tiling a particular
824  // dimension. This convention is significantly simpler to handle instead of
825  // adjusting affine maps to account for missing dimensions.
826  auto nLoops = op.getNumLoops();
827  SmallVector<OpFoldResult> tileSizeVector =
828  getAsOpFoldResult(options.tileSizeComputationFunction(b, op));
829  if (tileSizeVector.size() < nLoops) {
830  tileSizeVector.append(nLoops - tileSizeVector.size(), b.getIndexAttr(0));
831  }
832 
833  return tileLinalgOpImpl<LoopTy>(b, op, tileSizeVector, options);
834 }
835 
838  const LinalgTilingOptions &options) {
839  switch (options.loopType) {
841  return tileLinalgOpImpl<scf::ForOp>(b, op, options);
842  case LinalgTilingLoopType::ParallelLoops:
843  return tileLinalgOpImpl<scf::ParallelOp>(b, op, options);
844  default:;
845  }
846  return failure();
847 }
848 
849 namespace {
850 /// Helper classes for type list expansion.
851 template <typename... OpTypes>
852 class CanonicalizationPatternList;
853 
854 template <>
855 class CanonicalizationPatternList<> {
856 public:
857  static void insert(RewritePatternSet &patterns) {}
858 };
859 
860 template <typename OpTy, typename... OpTypes>
861 class CanonicalizationPatternList<OpTy, OpTypes...> {
862 public:
863  static void insert(RewritePatternSet &patterns) {
864  OpTy::getCanonicalizationPatterns(patterns, patterns.getContext());
865  CanonicalizationPatternList<OpTypes...>::insert(patterns);
866  }
867 };
868 } // namespace
869 
872  RewritePatternSet patterns(ctx);
874  return patterns;
875 }
876 
878  RewritePatternSet &patterns) {
879  auto *ctx = patterns.getContext();
880  affine::AffineApplyOp::getCanonicalizationPatterns(patterns, ctx);
881  affine::AffineForOp::getCanonicalizationPatterns(patterns, ctx);
882  affine::AffineMinOp::getCanonicalizationPatterns(patterns, ctx);
883  affine::AffineMaxOp::getCanonicalizationPatterns(patterns, ctx);
884  arith::ConstantIndexOp::getCanonicalizationPatterns(patterns, ctx);
885 
886  memref::SubViewOp::getCanonicalizationPatterns(patterns, ctx);
887  memref::ViewOp::getCanonicalizationPatterns(patterns, ctx);
888 
889  scf::ForOp::getCanonicalizationPatterns(patterns, ctx);
890  scf::ParallelOp::getCanonicalizationPatterns(patterns, ctx);
891 
892  tensor::CastOp::getCanonicalizationPatterns(patterns, ctx);
893  tensor::EmptyOp::getCanonicalizationPatterns(patterns, ctx);
894  tensor::ExtractSliceOp::getCanonicalizationPatterns(patterns, ctx);
895  tensor::InsertSliceOp::getCanonicalizationPatterns(patterns, ctx);
896  tensor::PadOp::getCanonicalizationPatterns(patterns, ctx);
897  ctx->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(patterns);
898 
899  CanonicalizationPatternList<
900 #define GET_OP_LIST
901 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
902  >::insert(patterns);
903 }
DiagnosedSilenceableFailure doit(RewriterBase &rewriter, OpTy target, transform::ApplyToEachResultList &results, transform::TransformState &state)
static llvm::ManagedStatic< PassManagerOptions > options
static FailureOr< ForallTilingResult > tileToForallOpImpl(RewriterBase &b, TilingInterface op, ArrayRef< OpFoldResult > numThreads, std::optional< ArrayRef< OpFoldResult >> nominalTileSizes, std::optional< ArrayAttr > mapping, bool omitTileOffsetBoundsCheck)
Rewrite a TilingInterface op to a tiled scf.forall.
Definition: Tiling.cpp:319
static bool canOmitTileOffsetInBoundsCheck(OpFoldResult tileSize, OpFoldResult numThreads, OpFoldResult iterationSize)
Returns true if the maximum tile offset tileSize * numThreads-1 is less than iterationSize.
Definition: Tiling.cpp:207
static void emitIsPositiveIndexAssertion(ImplicitLocOpBuilder &b, OpFoldResult value)
Asserts that the given index-typed value is strictly positive.
Definition: Tiling.cpp:94
static OpFoldResult buildMax(OpBuilder &b, Location loc, ArrayRef< OpFoldResult > vals)
Build an affine_max of all the vals.
Definition: Tiling.cpp:219
static void calculateTileOffsetsAndSizes(RewriterBase &b, Location loc, scf::ForallOp forallOp, ArrayRef< OpFoldResult > numThreads, SmallVector< Range > loopRanges, bool omitTileOffsetBoundsCheck, std::optional< ArrayRef< OpFoldResult >> nominalTileSizes, SmallVector< OpFoldResult > &tiledOffsets, SmallVector< OpFoldResult > &tiledSizes)
Fill out the tiledOffsets and tiledSizes to be used to tile to a given number of threads.
Definition: Tiling.cpp:236
static FailureOr< TiledLinalgOp > tileLinalgOpImpl(RewriterBase &b, LinalgOp op, ArrayRef< OpFoldResult > tileSizes, const LinalgTilingOptions &options)
Definition: Tiling.cpp:457
static OpFoldResult buildMin(OpBuilder &b, Location loc, ArrayRef< OpFoldResult > vals)
Build an affine_min of all the vals.
Definition: Tiling.cpp:227
Base type for affine expression.
Definition: AffineExpr.h:69
AffineExpr floorDiv(uint64_t v) const
Definition: AffineExpr.cpp:882
AffineExpr ceilDiv(uint64_t v) const
Definition: AffineExpr.cpp:925
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:47
static AffineMap getMultiDimIdentityMap(unsigned numDims, MLIRContext *context)
Returns an AffineMap with 'numDims' identity result dim exprs.
Definition: AffineMap.cpp:318
unsigned getNumResults() const
Definition: AffineMap.cpp:388
static AffineMap getPermutationMap(ArrayRef< unsigned > permutation, MLIRContext *context)
Returns an AffineMap representing a permutation.
Definition: AffineMap.cpp:248
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:124
AffineExpr getAffineSymbolExpr(unsigned position)
Definition: Builders.cpp:375
StringAttr getStringAttr(const Twine &bytes)
Definition: Builders.cpp:269
MLIRContext * getContext() const
Definition: Builders.h:55
This class provides support for representing a failure result, or a valid value of type T.
Definition: LogicalResult.h:78
ImplicitLocOpBuilder maintains a 'current location', allowing use of the create<> method without spec...
Location getLoc() const
Accessors for the implied location.
OpTy create(Args &&...args)
Create an operation of specific op type at the current insertion point and location.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:63
MLIRContext * getContext() const
Return the context this location is uniqued in.
Definition: Location.h:73
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:350
This class helps build Operations.
Definition: Builders.h:209
Operation * clone(Operation &op, IRMapping &mapper)
Creates a deep copy of the specified operation, remapping any operands that use values outside of the...
Definition: Builders.cpp:553
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Definition: Builders.h:433
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:400
void setInsertionPointToEnd(Block *block)
Sets the insertion point to the end of the specified block.
Definition: Builders.h:438
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:464
void setInsertionPointAfter(Operation *op)
Sets the insertion point to the node after the specified operation, which will cause subsequent inser...
Definition: Builders.h:414
This class represents a single result from folding an operation.
Definition: OpDefinition.h:266
This class represents an operand of an operation.
Definition: Value.h:263
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
Location getLoc()
The source location the operation was defined or derived from.
Definition: Operation.h:223
unsigned getNumOperands()
Definition: Operation.h:341
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
Definition: Operation.cpp:268
result_range getResults()
Definition: Operation.h:410
InFlightDiagnostic emitOpError(const Twine &message={})
Emit an error with the op name prefixed, like "'dim' op " which is convenient for verifiers.
Definition: Operation.cpp:671
unsigned getNumResults()
Return the number of results held by this operation.
Definition: Operation.h:399
MLIRContext * getContext() const
Definition: PatternMatch.h:785
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:399
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the rewriter that the IR failed to be rewritten because of a match failure,...
Definition: PatternMatch.h:685
virtual void replaceOp(Operation *op, ValueRange newValues)
This method replaces the results of the operation with the specified list of values.
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
void modifyOpInPlace(Operation *root, CallableT &&callable)
This method is a utility wrapper around an in-place modification of an operation.
Definition: PatternMatch.h:631
Tensor types represent multi-dimensional arrays, and have two variants: RankedTensorType and Unranked...
Definition: BuiltinTypes.h:91
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:378
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
SmallVector< OpFoldResult > makeComposedFoldedMultiResultAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Variant of makeComposedFoldedAffineApply suitable for multi-result maps.
Definition: AffineOps.cpp:1232
AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Returns a composed AffineApplyOp by composing map and operands with other AffineApplyOps supplying th...
Definition: AffineOps.cpp:1135
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,...
Definition: AffineOps.cpp:1298
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,...
Definition: AffineOps.cpp:1291
OpFoldResult makeComposedFoldedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Constructs an AffineApplyOp that applies map to operands after composing the map with the maps of any...
Definition: AffineOps.cpp:1185
void mapLoopToProcessorIds(scf::ForOp forOp, ArrayRef< Value > processorId, ArrayRef< Value > numProcessors)
Maps forOp for execution on a parallel grid of virtual processorIds of size given by numProcessors.
Definition: LoopUtils.cpp:1761
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:285
FailureOr< ForallTilingResult > tileToForallOpUsingTileSizes(RewriterBase &builder, TilingInterface op, ArrayRef< OpFoldResult > tileSizes, std::optional< ArrayAttr > mapping)
Same as tileToForallOp, but calculate the number of threads required using the given tileSizes.
Definition: Tiling.cpp:433
SmallVector< Value > makeTiledShapes(OpBuilder &builder, Location loc, LinalgOp linalgOp, ValueRange valuesToTile, ArrayRef< OpFoldResult > ivs, ArrayRef< OpFoldResult > tileSizes, ArrayRef< OpFoldResult > sizeBounds, bool omitPartialTileCheck)
Creates extract_slice/subview ops for all valuesToTile of the given linalgOp with builder,...
Definition: Utils.cpp:829
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:78
bool isParallelIterator(utils::IteratorType iteratorType)
Check if iterator type has "parallel" semantics.
Definition: Utils.cpp:184
void populateLinalgTilingCanonicalizationPatterns(RewritePatternSet &patterns)
Definition: Tiling.cpp:877
FailureOr< ForallTilingResult > tileToForallOp(RewriterBase &builder, TilingInterface op, ArrayRef< OpFoldResult > numThreads, std::optional< ArrayAttr > mapping)
Definition: Tiling.cpp:424
SmallVector< Value > insertSlicesBack(OpBuilder &builder, Location loc, LinalgOp op, ValueRange operands, ValueRange results)
Creates insert_slice ops that insert results back into larger tensors they were originally extracted ...
Definition: Utils.cpp:749
std::tuple< SmallVector< Range, 4 >, LoopIndexToRangeIndexMap > makeTiledLoopRanges(RewriterBase &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > allShapeSizes, ArrayRef< OpFoldResult > allTileSizes)
Definition: Tiling.cpp:49
void offsetIndices(OpBuilder &b, LinalgOp linalgOp, ArrayRef< OpFoldResult > offests)
Add the specified offsets to any linalg.index ops contained in the given linalgOp.
Definition: Utils.cpp:850
FailureOr< StaticMultiSizeSpecification > computeStaticMultiTileSizes(LinalgOp op, unsigned dimension, int64_t targetSize, int64_t divisor)
Definition: Tiling.cpp:111
FailureOr< ForallReductionTilingResult > tileReductionUsingForall(RewriterBase &b, PartialReductionOpInterface op, ArrayRef< OpFoldResult > numThreads, ArrayRef< OpFoldResult > tileSizes={}, std::optional< ArrayAttr > mapping=std::nullopt)
Method to tile a reduction to parallel iterations computing partial reductions.
Definition: Tiling.cpp:613
FailureOr< TiledLinalgOp > tileLinalgOp(RewriterBase &b, LinalgOp op, const LinalgTilingOptions &options)
Definition: Tiling.cpp:837
RewritePatternSet getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx)
Canonicalization patterns relevant to apply after tiling patterns.
Definition: Tiling.cpp:871
SmallVector< Type > getTensorOutputTypes(LinalgOp op, ValueRange operands)
Returns the list of tensor output types produced when the given structured operation op is applied to...
Definition: Utils.cpp:740
FailureOr< MultiSizeSpecification > computeMultiTileSizes(OpBuilder &builder, LinalgOp op, unsigned dimension, OpFoldResult targetSize, OpFoldResult divisor, bool emitAssertions=true)
Emits the IR computing the multi-sized tiling specification with two tile sizes not exceeding targetS...
Definition: Tiling.cpp:137
SmallVector< Value > ValueVector
An owning vector of values, handy to return from functions.
Definition: SCF.h:70
LogicalResult getOrCreateDestinations(OpBuilder &b, Location loc, Operation *op, SmallVector< Value > &result)
This is a helper function for DestinationStyleOpInterface.
Definition: TensorOps.cpp:105
Include the generated interface declarations.
bool isConstantIntValue(OpFoldResult ofr, int64_t value)
Return true if ofr is constant integer equal to value.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
std::optional< int64_t > getConstantIntValue(OpFoldResult ofr)
If ofr is a constant integer or an IntegerAttr, return the integer.
int64_t floorDiv(int64_t lhs, int64_t rhs)
Returns the result of MLIR's floordiv operation on constants.
Definition: MathExtras.h:33
void bindDims(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to DimExpr at positions: [0 .
Definition: AffineExpr.h:349
AffineMap inversePermutation(AffineMap map)
Returns a map of codomain to domain dimensions such that the first codomain dimension for a particula...
Definition: AffineMap.cpp:755
void bindSymbols(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to SymbolExpr at positions: [0 .
Definition: AffineExpr.h:363
Value getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc, OpFoldResult ofr)
Converts an OpFoldResult to a Value.
Definition: Utils.cpp:41
Operation * clone(OpBuilder &b, Operation *op, TypeRange newResultTypes, ValueRange newOperands)
OpFoldResult getAsOpFoldResult(Value val)
Given a value, try to extract a constant Attribute.
SmallVector< scf::ForOp, 8 > Loops
Tile a nest of standard for loops rooted at rootForOp by finding such parametric tile sizes that the ...
Definition: Utils.h:121
void applyPermutationToVector(SmallVector< T, N > &inVec, ArrayRef< int64_t > permutation)
Apply the permutation defined by permutation to inVec.
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value.
Definition: LogicalResult.h:72
Represents a range (offset, size, and stride) where each element of the triple may be dynamic or stat...
Transformation information returned after reduction tiling.
Definition: Transforms.h:869
Operation * parallelTiledOp
The partial reduction tiled op generated.
Definition: Transforms.h:871
Operation * initialOp
The op initializing the tensor used for partial reductions.
Definition: Transforms.h:875
scf::ForallOp loops
The scf.forall operation that iterate over the tiles.
Definition: Transforms.h:877
Operation * mergeOp
The final reduction operation merging all the partial reductions.
Definition: Transforms.h:873
Rewrite a TilingInterface op to a tiled scf.forall, applying tiling by numThreads.
Definition: Transforms.h:852
A description of a multi-size tiling comprising tile sizes and numbers of tiles, expressed as Values ...
Definition: Transforms.h:805
Callback function type used to get processor ID, and number of processors used for distribution for a...
Definition: Utils.h:295
Perform standalone tiling of a single LinalgOp by tileSizes.
Definition: Transforms.h:663
SmallVector< Value, 4 > tensorResults
Definition: Transforms.h:666
T lowTripCount
Number of tiles associated with each size.
Definition: Transforms.h:797
Eliminates variable at the specified position using Fourier-Motzkin variable elimination.