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