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 
110 FailureOr<StaticContinuousTileSizeSpecification>
112  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 = getValueOrCreateConstantIndexOp(b, loc,
186  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 =
379  llvm::to_vector(llvm::make_filter_range(numThreads, [](OpFoldResult ofr) {
380  return !isConstantIntValue(ofr, 0);
381  }));
382  int64_t nLoops = loopRanges.size();
383  tiledOffsets.reserve(nLoops);
384  tiledSizes.reserve(nLoops);
385  for (unsigned loopIdx = 0, threadIdIdx = 0; loopIdx < nLoops; ++loopIdx) {
386  bool overflow = loopIdx >= numThreads.size();
387  bool isZero = !overflow && isConstantIntValue(numThreads[loopIdx], 0);
388  // Degenerate case: take the whole domain.
389  if (overflow || isZero) {
390  tiledOffsets.push_back(loopRanges[loopIdx].offset);
391  tiledSizes.push_back(loopRanges[loopIdx].size);
392  continue;
393  }
394 
395  // Tiled case: compute the offset and size.
396  AffineExpr i, j, m, n, o;
397  bindDims(b.getContext(), i, j);
398  bindSymbols(b.getContext(), m, n, o);
399  OpFoldResult size = loopRanges[loopIdx].size;
400  OpFoldResult offset = loopRanges[loopIdx].offset;
401  OpFoldResult threadId = threadIds[threadIdIdx];
402  // Symbolic fixed max size per thread.
403  // TODO: floor + 0/1 depending on case for better load-balancing.
404  OpFoldResult tileSizePerThread =
405  nominalTileSizes.has_value()
406  ? (*nominalTileSizes)[loopIdx]
408  b, loc, m.ceilDiv(n),
409  ArrayRef<OpFoldResult>{size, nonZeroNumThreads[threadIdIdx]});
410 
411  // Dynamic offset shifted by threadId * maxSizePerThread.
413  b, loc, i + j * m, {offset, threadId, tileSizePerThread});
414  // Dynamic upper-bound depending on the threadId.
415  OpFoldResult residualTileSize = makeComposedFoldedAffineApply(
416  b, loc, i + j * m - n,
417  {offset, nonZeroNumThreads[threadIdIdx], tileSizePerThread, size});
418  if (!isConstantIntValue(residualTileSize, 0)) {
419  OpFoldResult sizeMinusOffsetPerThread = makeComposedFoldedAffineApply(
420  b, loc, -i + m, {offsetPerThread, size});
421  tileSizePerThread =
422  buildMin(b, loc, {sizeMinusOffsetPerThread, tileSizePerThread});
423  }
424 
425  tiledOffsets.push_back(offsetPerThread);
426  // TODO: if tileSizePerThread <= 0 early exit.
427  if (!omitTileOffsetBoundsCheck &&
428  !canOmitTileOffsetInBoundsCheck(tileSizePerThread,
429  nonZeroNumThreads[threadIdIdx], size))
430  tileSizePerThread =
431  buildMax(b, loc, {b.getIndexAttr(0), tileSizePerThread});
432 
433  tiledSizes.push_back(tileSizePerThread);
434  ++threadIdIdx;
435  }
436 }
437 
438 /// Returns a vector of bools representing if, for each axis, `op` can be tiled
439 /// without incurring in a race condition and thus it is thread-safe to do the
440 /// tiling. This is checked by iterating over numThreads and ensuring that the
441 /// corresponding iterator type is "parallel". If it is not, then we know that
442 /// such dimension is unsafe to tile.
444  ArrayRef<OpFoldResult> numThreads) {
445  auto iterators = linalgOp.getIteratorTypesArray();
446  SmallVector<bool> safeToTile(numThreads.size(), true);
447 
448  for (unsigned i = 0, e = numThreads.size(); i != e; i++) {
449  if (auto attr = llvm::dyn_cast_if_present<Attribute>(numThreads[i])) {
450  if (cast<IntegerAttr>(attr).getValue().getSExtValue() > 1) {
451  safeToTile[i] = iterators[i] == utils::IteratorType::parallel;
452  }
453  } else {
454  safeToTile[i] = iterators[i] == utils::IteratorType::parallel;
455  }
456  }
457  return safeToTile;
458 }
459 
460 /// Rewrite a TilingInterface `op` to a tiled `scf.forall`. The
461 /// tiling is specified by the number of tiles/threads `numThreads` and the
462 /// optional nominal tile size `nominalTileSizes`. If `nominalTilSizes` is
463 /// not specified, then it is derived from `numThreads` as `ceilDiv(dimSize[i],
464 /// numThreads[i])`. If non-empty, the `mapping` is added as an
465 /// attribute to the resulting `scf.forall`. A zero tile sizes indicate
466 /// that the dimension is not tiled, and can be thought of as tiling by the full
467 /// size of data.
468 /// It is the user's responsibility to ensure that `numThreads` is a valid
469 /// tiling specification (i.e. that only tiles parallel dimensions, e.g. in the
470 /// Linalg case). If the dimension is not parallelizable, a warning is issued to
471 /// notify the user that the generated code is not safe to parallelize. If
472 /// `omitTileOffsetBoundsCheck` is true, then the function will assume that
473 /// `tileSize[i] * (numThread[i] -1) <= dimSize[i]` holds.
474 static FailureOr<ForallTilingResult> tileToForallOpImpl(
475  RewriterBase &b, TilingInterface op, ArrayRef<OpFoldResult> numThreads,
476  std::optional<ArrayRef<OpFoldResult>> nominalTileSizes,
477  std::optional<ArrayAttr> mapping, bool omitTileOffsetBoundsCheck) {
478  Location loc = op->getLoc();
480 
481  SmallVector<Range> loopRanges = op.getIterationDomain(b);
482  if (loopRanges.empty())
483  return op->emitOpError("expected non-empty loop ranges");
484  auto hasStrideOne = [](Range r) { return !isConstantIntValue(r.stride, 1); };
485  if (llvm::any_of(loopRanges, hasStrideOne))
486  return op->emitOpError("only stride-1 supported atm");
487 
488  // Gather destination tensors.
489  SmallVector<Value> dest;
490  if (failed(tensor::getOrCreateDestinations(b, loc, op, dest)))
491  return op->emitOpError("failed to get destination tensors");
492 
493  SmallVector<OpFoldResult> nonZeroNumThreads =
494  llvm::to_vector(llvm::make_filter_range(numThreads, [](OpFoldResult ofr) {
495  return !isConstantIntValue(ofr, 0);
496  }));
497  SmallVector<Value> materializedNonZeroNumThreads =
498  llvm::to_vector(llvm::map_range(nonZeroNumThreads, [&](OpFoldResult ofr) {
499  return getValueOrCreateConstantIndexOp(b, loc, ofr);
500  }));
501 
502  LinalgOp linalgOp = dyn_cast<LinalgOp>(op.getOperation());
503  if (linalgOp) {
504  // Check if tiling is thread safe and print a warning if not.
505  SmallVector<bool> tilingSafety =
506  safeToTileToForall(b.getContext(), linalgOp, numThreads);
507  for (size_t i = 0; i < tilingSafety.size(); i++)
508  if (!tilingSafety[i])
509  op.emitWarning() << "tiling is not thread safe at axis #" << i;
510  }
511 
512  // 1. Create the ForallOp. We don't use the lambda body-builder
513  // version because we require the use of RewriterBase in the body, so we
514  // manually move the insertion point to the body below.
515  scf::ForallOp forallOp = b.create<scf::ForallOp>(
516  loc, getAsOpFoldResult((materializedNonZeroNumThreads)), dest, mapping);
517 
518  // 2. Fill out the ForallOp body.
519  SmallVector<OpFoldResult> tiledOffsets, tiledSizes;
520  calculateTileOffsetsAndSizes(b, loc, forallOp, numThreads, loopRanges,
521  omitTileOffsetBoundsCheck, nominalTileSizes,
522  tiledOffsets, tiledSizes);
523 
524  // 3. Clone the tileable op and update its destination operands to use the
525  // output bbArgs of the ForallOp.
526  ArrayRef<BlockArgument> destBbArgs = forallOp.getRegionIterArgs();
527  Operation *tiledOp = nullptr;
528  SmallVector<Value> tiledValues;
529  {
530  // 3.a. RAII guard, inserting within forallOp, before terminator.
532  b.setInsertionPoint(forallOp.getTerminator());
533  Operation *clonedOp = b.clone(*op.getOperation());
534  auto destinationStyleOp = dyn_cast<DestinationStyleOpInterface>(clonedOp);
535  if (destinationStyleOp) {
536  for (OpOperand &outOperand : destinationStyleOp.getDpsInitsMutable()) {
537  // Swap tensor inits with the corresponding block argument of the
538  // scf.forall op. Memref inits remain as is.
539  if (isa<TensorType>(outOperand.get().getType())) {
540  auto *it = llvm::find(dest, outOperand.get());
541  assert(it != dest.end() && "could not find destination tensor");
542  unsigned destNum = std::distance(dest.begin(), it);
543  outOperand.set(destBbArgs[destNum]);
544  }
545  }
546  }
547 
548  // 4. Tile the cloned op and delete the clone.
549  FailureOr<TilingResult> tilingResult =
550  cast<TilingInterface>(clonedOp).getTiledImplementation(b, tiledOffsets,
551  tiledSizes);
552  if (failed(tilingResult))
553  return clonedOp->emitError("Failed to tile op: ");
554  if (tilingResult->tiledOps.size() != 1) {
555  return clonedOp->emitError("expected a single produced tiled op, got ")
556  << tilingResult->tiledOps.size();
557  }
558 
559  b.eraseOp(clonedOp);
560  tiledOp = tilingResult->tiledOps.front();
561  tiledValues = tilingResult->tiledValues;
562  }
563 
564  // 5. Parallel insert back into the result tensor.
565  for (auto it : llvm::zip(llvm::seq(unsigned(0), unsigned(dest.size())),
566  tiledValues, destBbArgs)) {
567  // 5.a. Partial subset information is inserted just before the terminator.
569  b.setInsertionPoint(forallOp.getTerminator());
570 
571  SmallVector<OpFoldResult> resultOffsets, resultSizes;
572  if (failed(op.getResultTilePosition(b, std::get<0>(it), tiledOffsets,
573  tiledSizes, resultOffsets,
574  resultSizes)))
575  return op->emitOpError("output offsets couldn't be calculated");
576  SmallVector<OpFoldResult> strides(resultSizes.size(), b.getIndexAttr(1));
577 
578  // 5.b. Parallel insertions are inserted at the end of the combining
579  // terminator.
580  b.setInsertionPointToEnd(forallOp.getTerminator().getBody());
581  b.create<tensor::ParallelInsertSliceOp>(loc, std::get<1>(it),
582  std::get<2>(it), resultOffsets,
583  resultSizes, strides);
584  }
585  return ForallTilingResult{forallOp, tiledOp};
586 }
587 
588 FailureOr<ForallTilingResult>
589 linalg::tileToForallOp(RewriterBase &b, TilingInterface op,
590  ArrayRef<OpFoldResult> numThreads,
591  std::optional<ArrayAttr> mapping) {
592  return tileToForallOpImpl(b, op, numThreads,
593  /*nominalTileSizes=*/std::nullopt, mapping,
594  /*omitTileOffsetBoundsCheck=*/false);
595 }
596 
597 FailureOr<ForallTilingResult>
599  ArrayRef<OpFoldResult> tileSizes,
600  std::optional<ArrayAttr> mapping) {
601  SmallVector<Range> loopRanges = op.getIterationDomain(b);
602  unsigned nLoops = loopRanges.size();
603  SmallVector<OpFoldResult> numThreads;
604  numThreads.reserve(nLoops);
605  AffineExpr s0, s1;
606  bindSymbols(b.getContext(), s0, s1);
607  AffineExpr divExpr = s0.ceilDiv(s1);
608  for (const auto &it : llvm::zip(tileSizes, loopRanges)) {
609  OpFoldResult numTiles = std::get<0>(it);
610  if (!isConstantIntValue(numTiles, 0))
612  b, op.getLoc(), divExpr, {std::get<1>(it).size, std::get<0>(it)});
613  numThreads.push_back(numTiles);
614  }
615  return tileToForallOpImpl(b, op, numThreads,
616  /*nominalTileSizes=*/tileSizes, mapping,
617  /*omitTileOffsetBoundsCheck=*/true);
618 }
619 
620 template <typename LoopTy>
621 static FailureOr<TiledLinalgOp>
623  const LinalgTilingOptions &options) {
625 
626  auto nLoops = op.getNumLoops();
627  // Initial tile sizes may be too big, only take the first nLoops.
628  tileSizes = tileSizes.take_front(nLoops);
629 
630  if (llvm::all_of(tileSizes, [](OpFoldResult ofr) {
631  return getConstantIntValue(ofr) == static_cast<int64_t>(0);
632  })) {
633  TiledLinalgOp tiledOp;
634  tiledOp.op = cast<LinalgOp>(b.clone(*op.getOperation()));
635  tiledOp.tensorResults.assign(tiledOp.op->result_begin(),
636  tiledOp.op->result_end());
637  return tiledOp;
638  }
639 
640  // 1. Build the tiled loop ranges.
641  SmallVector<OpFoldResult> allShapeSizes =
642  op.createFlatListOfOperandDims(b, op.getLoc());
643  AffineMap shapeSizesToLoopsMap = op.getShapesToLoopsMap();
644  if (!shapeSizesToLoopsMap)
645  return failure();
646 
647  auto [loopRanges, loopIndexToRangeIndex] = makeTiledLoopRanges(
648  b, op.getLoc(), shapeSizesToLoopsMap, allShapeSizes, tileSizes);
649 
651  for (const auto &attr : enumerate(op.getIteratorTypesArray())) {
652  if (loopIndexToRangeIndex.count(attr.index()))
653  iteratorTypes.push_back(attr.value());
654  }
655  // If interchangeVector is empty, use the identity. Build the permutation map
656  // otherwise.
657  auto invPermutationMap =
658  AffineMap::getMultiDimIdentityMap(tileSizes.size(), b.getContext());
659  if (!options.interchangeVector.empty()) {
660  // Based on the pruned iterations (due to zero tile size), recompute the
661  // interchange vector.
662  SmallVector<unsigned, 4> interchangeVector;
663  interchangeVector.reserve(options.interchangeVector.size());
664  for (auto pos : options.interchangeVector) {
665  auto it = loopIndexToRangeIndex.find(pos);
666  if (it == loopIndexToRangeIndex.end())
667  continue;
668  interchangeVector.push_back(it->second);
669  }
670  // Interchange vector is guaranteed to be a permutation,
671  // `inversePermutation` must succeed.
672  invPermutationMap = inversePermutation(
673  AffineMap::getPermutationMap(interchangeVector, b.getContext()));
674  assert(invPermutationMap);
675  SmallVector<int64_t> permutation(interchangeVector.begin(),
676  interchangeVector.end());
677  applyPermutationToVector(loopRanges, permutation);
678  applyPermutationToVector(iteratorTypes, permutation);
679  }
680 
681  // Handle distribution. Create a vector of the same size of loops that are to
682  // be tiled.
684  if (options.distribution) {
685  procInfo.resize(
686  iteratorTypes.size(),
687  linalg::ProcInfo{nullptr, nullptr, linalg::DistributionMethod::None});
688  // Collect loop ranges of tiled loops, loops that are parallel.
689  SmallVector<Range> parallelLoopRanges;
690  for (const auto &iteratorType : llvm::enumerate(iteratorTypes)) {
691  if (!isParallelIterator(iteratorType.value()))
692  break;
693  parallelLoopRanges.push_back(loopRanges[iteratorType.index()]);
694  }
695  auto returnedProcInfo =
696  options.distribution->procInfo(b, op.getLoc(), parallelLoopRanges);
697  unsigned procIdIdx = 0;
698  // Update the distribution information for the loops.
699  for (const auto &iteratorType : llvm::enumerate(iteratorTypes)) {
700  if (!isParallelIterator(iteratorType.value()))
701  break;
702  procInfo[iteratorType.index()] = returnedProcInfo[procIdIdx++];
703  }
704  }
705 
706  // 2. Create the tiled loops.
707  LinalgOp res = op;
708  SmallVector<Value, 4> ivs, tensorResults;
709  auto tiledLoopBodyBuilder =
710  [&](OpBuilder &builder, Location loc, ValueRange localIvs,
711  ValueRange operandValuesToUse) -> scf::ValueVector {
712  ivs.assign(localIvs.begin(), localIvs.end());
713 
714  // When an `interchangeVector` is present, it has been applied to the
715  // loop ranges and the iterator types. Apply its inverse to the
716  // resulting loop `ivs` to match the op definition.
717  SmallVector<Value, 4> interchangedIvs;
718  if (!options.interchangeVector.empty()) {
719  for (AffineExpr result : invPermutationMap.getResults())
720  interchangedIvs.push_back(
721  ivs[cast<AffineDimExpr>(result).getPosition()]);
722  } else {
723  interchangedIvs.assign(ivs.begin(), ivs.end());
724  }
725 
726  // Tile the `operandValuesToUse` that either match the `op` operands
727  // themselves or the tile loop arguments forwarding them.
728  assert(operandValuesToUse.size() ==
729  static_cast<size_t>(op->getNumOperands()) &&
730  "expect the number of operands and inputs and outputs to match");
731  SmallVector<Value> valuesToTile = operandValuesToUse;
732  SmallVector<OpFoldResult> sizeBounds =
733  makeComposedFoldedMultiResultAffineApply(b, loc, shapeSizesToLoopsMap,
734  allShapeSizes);
735  SmallVector<Value> tiledOperands = makeTiledShapes(
736  b, loc, op, valuesToTile, getAsOpFoldResult(interchangedIvs), tileSizes,
737  sizeBounds,
738  /*omitPartialTileCheck=*/false);
739 
740  SmallVector<Type> resultTensorTypes =
741  getTensorOutputTypes(op, tiledOperands);
742  res = clone(b, op, resultTensorTypes, tiledOperands);
743  tensorResults =
744  insertSlicesBack(builder, loc, op, tiledOperands, res->getResults());
745  return scf::ValueVector(tensorResults.begin(), tensorResults.end());
746  };
747  GenerateLoopNest<LoopTy>::doit(b, op.getLoc(), loopRanges, op, iteratorTypes,
748  tiledLoopBodyBuilder, procInfo);
749 
750  // 3. Transform IndexOp results w.r.t. the tiling.
751  transformIndexOps(b, res, ivs, loopIndexToRangeIndex);
752 
753  // 4. Gather the newly created loops and return them with the new op.
755  loops.reserve(ivs.size());
756  for (auto iv : ivs) {
757  if (isa<BlockArgument>(iv)) {
758  loops.push_back(cast<BlockArgument>(iv).getOwner()->getParentOp());
759  assert(loops.back() && "no owner found for induction variable!");
760  } else {
761  // TODO: Instead of doing this, try to recover the ops used instead of the
762  // loop.
763  loops.push_back(nullptr);
764  }
765  }
766 
767  // 5. Get the tensor results from the outermost loop if available. Otherwise
768  // use the previously captured `tensorResults`.
769  Operation *outermostLoop = nullptr;
770  for (Operation *loop : loops)
771  if ((outermostLoop = loop))
772  break;
773 
774  return TiledLinalgOp{
775  res, loops, outermostLoop ? outermostLoop->getResults() : tensorResults};
776 }
777 
778 FailureOr<linalg::ForallReductionTilingResult> linalg::tileReductionUsingForall(
779  RewriterBase &b, PartialReductionOpInterface op,
780  ArrayRef<OpFoldResult> numThreads, ArrayRef<OpFoldResult> tileSizes,
781  std::optional<ArrayAttr> mapping) {
782  Location loc = op.getLoc();
784 
785  // Ops implementing PartialReductionOpInterface are expected to implement
786  // TilingInterface.
787  // TODO: proper core mechanism to tie interfaces together.
788  auto tilingInterfaceOp = cast<TilingInterface>(op.getOperation());
789 
790  // Ops implementing PartialReductionOpInterface are not necessarily expected
791  // to implement TilingInterface.. This cast is unsafe atm.
792  // TODO: proper core mechanism to tie interfaces together.
793  // TODO: this function requires a pair of interfaces ..
794  auto destinationStyleOp =
795  dyn_cast<DestinationStyleOpInterface>(op.getOperation());
796  if (!destinationStyleOp)
797  return b.notifyMatchFailure(op, "not a destination style op");
798 
799  // Actually this only work for Linalg ops atm.
800  auto linalgOp = dyn_cast<linalg::LinalgOp>(op.getOperation());
801  if (!linalgOp)
802  return b.notifyMatchFailure(op, "not a linalg op");
803 
804  SmallVector<Range> iterationDomain = tilingInterfaceOp.getIterationDomain(b);
805  if (op->getNumResults() != 1)
806  return b.notifyMatchFailure(
807  op, "don't support ops with multiple results for now");
808 
810  tilingInterfaceOp.getLoopIteratorTypes();
811  SmallVector<unsigned> redDims;
812  linalgOp.getReductionDims(redDims);
813  if (redDims.size() != 1)
814  return b.notifyMatchFailure(
815  op, "only support ops with one reduction dimension.");
816  if (!tileSizes.empty() && tileSizes.size() != numThreads.size())
817  return b.notifyMatchFailure(op, "if tile sizes are present it must have as "
818  "many elements as number of threads");
819  int reductionDim = static_cast<int>(redDims.front());
820 
821  if (redDims.front() >= numThreads.size())
822  return b.notifyMatchFailure(
823  op, "reduction dimension must be mapped to threads");
824 
825  // 1. Create the inital tensor value.
826  FailureOr<SmallVector<Value>> maybeInitTensors =
827  op.generateInitialTensorForPartialReduction(b, loc, numThreads,
828  reductionDim);
829  if (failed(maybeInitTensors))
830  return b.notifyMatchFailure(
831  op, "Failed to create inital tensors for partial reduction");
832  SmallVector<Value> &initTensors = maybeInitTensors.value();
833 
834  // Gather destination tensors.
835  SmallVector<Value> dest;
836  if (failed(tensor::getOrCreateDestinations(b, loc, op, dest)))
837  return b.notifyMatchFailure(op, "failed to get destination tensors");
838 
839  Operation *tiledOp = nullptr;
840 
841  SmallVector<OpFoldResult> nonZeroNumThreads =
842  llvm::to_vector(llvm::make_filter_range(numThreads, [](OpFoldResult ofr) {
843  return !isConstantIntValue(ofr, 0);
844  }));
845  SmallVector<Value> materializedNonZeroNumThreads =
846  getValueOrCreateConstantIndexOp(b, loc, nonZeroNumThreads);
847 
848  // 2. Create the ForallOp with an empty region.
849  scf::ForallOp forallOp = b.create<scf::ForallOp>(
850  loc, getAsOpFoldResult(materializedNonZeroNumThreads), initTensors,
851  mapping);
852 
853  // 3. Calculate the tile offsets and sizes for the subsequent loop that will
854  // be nested under `forallOp`.
855  SmallVector<OpFoldResult> tiledOffsets, tiledSizes;
856  calculateTileOffsetsAndSizes(b, loc, forallOp, numThreads, iterationDomain,
857  /*omitTileOffsetBoundsCheck =*/false,
858  /*nominalTileSizes=*/std::nullopt, tiledOffsets,
859  tiledSizes);
860 
861  // 4b. Clone the tileable op and update its destination operands to use the
862  // output bbArgs of the ForallOp.
863  SmallVector<Value> tilingResults;
864  ArrayRef<BlockArgument> destBbArgs = forallOp.getRegionIterArgs();
865  {
866  // 4.a. RAII guard, inserting within forallOp, before terminator.
868  b.setInsertionPoint(forallOp.getTerminator());
869 
870  SmallVector<Value> tiledDpsInitOperands;
871  for (Value initOperand : destinationStyleOp.getDpsInits()) {
872  auto *it = llvm::find(dest, initOperand);
873  assert(it != dest.end() && "dest operand not found in dest");
874  unsigned destNum = std::distance(dest.begin(), it);
875  SmallVector<OpFoldResult> strides(numThreads.size(), b.getIndexAttr(1));
876  SmallVector<OpFoldResult> outOffsets(numThreads.size(),
877  b.getIndexAttr(0));
878  SmallVector<OpFoldResult> sizes = tiledSizes;
879  sizes[reductionDim] = b.getIndexAttr(1);
880  outOffsets[reductionDim] = forallOp.getInductionVars()[0];
881  // TODO: use SubsetExtractOpInterface once it is available.
882  tiledDpsInitOperands.push_back(b.create<tensor::ExtractSliceOp>(
883  loc, cast<RankedTensorType>(initOperand.getType()),
884  destBbArgs[destNum], outOffsets, sizes, strides));
885  }
886 
887  // 4.b. Clone the op and update init operands.
888  // We cannot use a IRMapping here because it can replace
889  // different OpOperands with the same value.
890  Operation *clonedOp = b.clone(*op.getOperation());
891  b.modifyOpInPlace(clonedOp, [&]() {
892  for (auto [initOperandPtr, tiledInitValue] : llvm::zip_equal(
893  cast<DestinationStyleOpInterface>(clonedOp).getDpsInitsMutable(),
894  tiledDpsInitOperands)) {
895  initOperandPtr.set(tiledInitValue);
896  }
897  });
898 
899  // 5. Tile the cloned op and delete the clone.
900  if (tileSizes.empty()) {
901  FailureOr<TilingResult> tilingResult =
902  cast<TilingInterface>(clonedOp).getTiledImplementation(
903  b, tiledOffsets, tiledSizes);
904  if (failed(tilingResult))
905  return clonedOp->emitError("Failed to tile op: ");
906  if (tilingResult->tiledOps.size() != 1) {
907  return clonedOp->emitError("expected a single produced tiled op, got ")
908  << tilingResult->tiledOps.size();
909  }
910  tiledOp = tilingResult->tiledOps.front();
911  tilingResults = tilingResult->tiledValues;
912  } else {
914  FailureOr<TiledLinalgOp> maybeTiled = tileLinalgOpImpl<scf::ForOp>(
915  b, cast<LinalgOp>(clonedOp), tileSizes, options);
916  if (failed(maybeTiled))
917  return b.notifyMatchFailure(op, "failed tileLinalgOpImpl");
918 
919  SmallVector<Value> ids = forallOp.getInductionVars();
920  mapLoopToProcessorIds(cast<scf::ForOp>(maybeTiled->loops.back()), ids,
921  materializedNonZeroNumThreads);
922  if (maybeTiled->loops.size() != 1) {
923  return clonedOp->emitError("expected a single produced loop");
924  }
925  tiledOp = maybeTiled->op;
926  tilingResults = maybeTiled->loops.front()->getResults();
927  }
928 
929  b.eraseOp(clonedOp);
930  }
931 
932  // 6. Insert the partial reductions back into a new tensor.
933  for (auto [index, result, bbArg] : llvm::zip(
934  llvm::seq<unsigned>(0, dest.size()), tilingResults, destBbArgs)) {
935  // 6.a. Partial subset information is inserted just before the terminator.
937  b.setInsertionPoint(forallOp.getTerminator());
938 
939  SmallVector<OpFoldResult> resultOffsets, resultSizes;
940  if (failed(tilingInterfaceOp.getResultTilePosition(
941  b, index, tiledOffsets, tiledSizes, resultOffsets, resultSizes)))
942  return op->emitOpError("output offsets couldn't be calculated");
943  SmallVector<OpFoldResult> resultOffsetsRank, resultSizesRank;
944  int64_t offIdx = 0;
945  int64_t sizeIdx = 0;
946  for (int64_t i = 0, e = numThreads.size(); i < e; ++i) {
947  if (i == reductionDim) {
948  resultOffsetsRank.push_back(forallOp.getInductionVars()[0]);
949  resultSizesRank.push_back(b.getIndexAttr(1));
950  continue;
951  }
952  resultOffsetsRank.push_back(resultOffsets[offIdx++]);
953  resultSizesRank.push_back(resultSizes[sizeIdx++]);
954  }
955  SmallVector<OpFoldResult> strides(resultSizesRank.size(),
956  b.getIndexAttr(1));
957 
958  // 6.b. Parallel insertions are inserted at the end of the combining
959  // terminator.
960  b.setInsertionPointToEnd(forallOp.getTerminator().getBody());
961  b.create<tensor::ParallelInsertSliceOp>(
962  loc, result, bbArg, resultOffsetsRank, resultSizesRank, strides);
963  }
964 
965  // 7. Merge the partial reductions.
966  b.setInsertionPointAfter(forallOp);
967  FailureOr<MergeResult> mergeResult =
968  op.mergeReductions(b, loc, forallOp->getResults(), reductionDim);
969  if (failed(mergeResult)) {
970  return failure();
971  }
972  b.replaceOp(op, mergeResult->replacements);
973 
974  // 8. Return.
976  results.initialValues = initTensors;
977  results.loops = forallOp;
978  results.parallelTiledOps.push_back(tiledOp);
979  results.mergeOps.append(mergeResult->mergeOps);
980  return results;
981 }
982 
983 template <typename LoopTy>
984 FailureOr<TiledLinalgOp> static tileLinalgOpImpl(
985  RewriterBase &b, LinalgOp op, const LinalgTilingOptions &options) {
987  b.setInsertionPoint(op);
988 
989  if (!options.tileSizeComputationFunction)
990  return failure();
991 
992  // Enforce the convention that "tiling by zero" skips tiling a particular
993  // dimension. This convention is significantly simpler to handle instead of
994  // adjusting affine maps to account for missing dimensions.
995  auto nLoops = op.getNumLoops();
996  SmallVector<OpFoldResult> tileSizeVector =
997  getAsOpFoldResult(options.tileSizeComputationFunction(b, op));
998  if (tileSizeVector.size() < nLoops) {
999  tileSizeVector.append(nLoops - tileSizeVector.size(), b.getIndexAttr(0));
1000  }
1001 
1002  return tileLinalgOpImpl<LoopTy>(b, op, tileSizeVector, options);
1003 }
1004 
1005 FailureOr<TiledLinalgOp>
1007  const LinalgTilingOptions &options) {
1008  switch (options.loopType) {
1010  return tileLinalgOpImpl<scf::ForOp>(b, op, options);
1011  case LinalgTilingLoopType::ParallelLoops:
1012  return tileLinalgOpImpl<scf::ParallelOp>(b, op, options);
1013  default:;
1014  }
1015  return failure();
1016 }
1017 
1018 namespace {
1019 /// Helper classes for type list expansion.
1020 template <typename... OpTypes>
1021 class CanonicalizationPatternList;
1022 
1023 template <>
1024 class CanonicalizationPatternList<> {
1025 public:
1026  static void insert(RewritePatternSet &patterns) {}
1027 };
1028 
1029 template <typename OpTy, typename... OpTypes>
1030 class CanonicalizationPatternList<OpTy, OpTypes...> {
1031 public:
1032  static void insert(RewritePatternSet &patterns) {
1033  OpTy::getCanonicalizationPatterns(patterns, patterns.getContext());
1034  CanonicalizationPatternList<OpTypes...>::insert(patterns);
1035  }
1036 };
1037 } // namespace
1038 
1041  RewritePatternSet patterns(ctx);
1043  return patterns;
1044 }
1045 
1047  RewritePatternSet &patterns) {
1048  auto *ctx = patterns.getContext();
1049  affine::AffineApplyOp::getCanonicalizationPatterns(patterns, ctx);
1050  affine::AffineForOp::getCanonicalizationPatterns(patterns, ctx);
1051  affine::AffineMinOp::getCanonicalizationPatterns(patterns, ctx);
1052  affine::AffineMaxOp::getCanonicalizationPatterns(patterns, ctx);
1053  arith::ConstantIndexOp::getCanonicalizationPatterns(patterns, ctx);
1054 
1055  memref::SubViewOp::getCanonicalizationPatterns(patterns, ctx);
1056  memref::ViewOp::getCanonicalizationPatterns(patterns, ctx);
1057 
1058  scf::ForOp::getCanonicalizationPatterns(patterns, ctx);
1059  scf::ParallelOp::getCanonicalizationPatterns(patterns, ctx);
1060 
1061  tensor::CastOp::getCanonicalizationPatterns(patterns, ctx);
1062  tensor::EmptyOp::getCanonicalizationPatterns(patterns, ctx);
1063  tensor::ExtractSliceOp::getCanonicalizationPatterns(patterns, ctx);
1064  tensor::InsertSliceOp::getCanonicalizationPatterns(patterns, ctx);
1065  tensor::PadOp::getCanonicalizationPatterns(patterns, ctx);
1066  ctx->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(patterns);
1067 
1068  CanonicalizationPatternList<
1069 #define GET_OP_LIST
1070 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
1071  >::insert(patterns);
1072 }
DiagnosedSilenceableFailure doit(RewriterBase &rewriter, OpTy target, transform::ApplyToEachResultList &results, transform::TransformState &state)
static llvm::ManagedStatic< PassManagerOptions > options
SmallVector< bool > safeToTileToForall(mlir::MLIRContext *ctx, LinalgOp linalgOp, ArrayRef< OpFoldResult > numThreads)
Returns a vector of bools representing if, for each axis, op can be tiled without incurring in a race...
Definition: Tiling.cpp:443
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:474
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:94
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:622
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:901
AffineExpr ceilDiv(uint64_t v) const
Definition: AffineExpr.cpp:948
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:321
unsigned getNumResults() const
Definition: AffineMap.cpp:389
static AffineMap getPermutationMap(ArrayRef< unsigned > permutation, MLIRContext *context)
Returns an AffineMap representing a permutation.
Definition: AffineMap.cpp:251
Attributes are known-constant values of operations.
Definition: Attributes.h:25
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
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:555
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
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:522
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:268
This class represents an operand of an operation.
Definition: Value.h:267
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
InFlightDiagnostic emitWarning(const Twine &message={})
Emit a warning about this operation, reporting up to any diagnostic handlers that may be listening.
Definition: Operation.cpp:280
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:823
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:400
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:718
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:630
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:381
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:92
SmallVector< OpFoldResult > makeComposedFoldedMultiResultAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Variant of makeComposedFoldedAffineApply suitable for multi-result maps.
Definition: AffineOps.cpp:1239
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:1142
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:1305
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:1298
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:1192
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:1722
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:598
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:1046
FailureOr< ForallTilingResult > tileToForallOp(RewriterBase &builder, TilingInterface op, ArrayRef< OpFoldResult > numThreads, std::optional< ArrayAttr > mapping)
Definition: Tiling.cpp:589
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: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: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:778
FailureOr< TiledLinalgOp > tileLinalgOp(RewriterBase &b, LinalgOp op, const LinalgTilingOptions &options)
Definition: Tiling.cpp:1006
RewritePatternSet getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx)
Canonicalization patterns relevant to apply after tiling patterns.
Definition: Tiling.cpp:1040
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:268
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:109
Include the generated interface declarations.
bool isConstantIntValue(OpFoldResult ofr, int64_t value)
Return true if ofr is constant integer equal to value.
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:348
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:362
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:144
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:894
SmallVector< Operation * > mergeOps
The final reduction operation merging all the partial reductions.
Definition: Transforms.h:898
SmallVector< Value > initialValues
Initial values used for partial reductions.
Definition: Transforms.h:900
scf::ForallOp loops
The scf.forall operation that iterate over the tiles.
Definition: Transforms.h:902
SmallVector< Operation * > parallelTiledOps
The partial reduction tiled op generated.
Definition: Transforms.h:896
Rewrite a TilingInterface op to a tiled scf.forall, applying tiling by numThreads.
Definition: Transforms.h:877
A description of a multi-size tiling comprising tile sizes and numbers of tiles, expressed as Values ...
Definition: Transforms.h:818
Callback function type used to get processor ID, and number of processors used for distribution for a...
Definition: Utils.h:293
Perform standalone tiling of a single LinalgOp by tileSizes.
Definition: Transforms.h:667
SmallVector< Value, 4 > tensorResults
Definition: Transforms.h:670
SmallVector< T > tripCounts
Number of tiles associated with each size.
Definition: Transforms.h:809
T lowTripCount
Number of tiles associated with each size.
Definition: Transforms.h:801
Eliminates variable at the specified position using Fourier-Motzkin variable elimination.