1 //===- Tiling.cpp - Implementation of linalg Tiling -----------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements the linalg dialect Tiling pass.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include <utility>
14 
15 #include "PassDetail.h"
16 #include "mlir/Dialect/Linalg/IR/Linalg.h"
17 #include "mlir/Dialect/Linalg/Passes.h"
18 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
19 #include "mlir/Dialect/Linalg/Utils/Utils.h"
20 #include "mlir/Dialect/MemRef/IR/MemRef.h"
21 #include "mlir/Dialect/SCF/Transforms.h"
22 #include "mlir/Dialect/Tensor/IR/Tensor.h"
23 #include "mlir/IR/AffineExpr.h"
24 #include "mlir/IR/AffineMap.h"
25 #include "mlir/Transforms/FoldUtils.h"
26 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
27 
28 #include "llvm/Support/CommandLine.h"
29 
30 using namespace mlir;
31 using namespace mlir::linalg;
32 using namespace mlir::scf;
33 
34 #define DEBUG_TYPE "linalg-tiling"
35 
36 static bool isZero(Value v) {
37   if (auto cst = v.getDefiningOp<arith::ConstantIndexOp>())
38     return cst.value() == 0;
39   return false;
40 }
41 
42 std::tuple<SmallVector<Range, 4>, LoopIndexToRangeIndexMap>
43 mlir::linalg::makeTiledLoopRanges(RewriterBase &b, Location loc, AffineMap map,
44                                   ValueRange allShapeSizes,
45                                   ValueRange allTileSizes) {
46   assert(allTileSizes.size() == map.getNumResults());
47   // Apply `map` to get shape sizes in loop order.
48   auto shapeSizes = applyMapToValues(b, loc, map, allShapeSizes);
49   SmallVector<Value, 4> tileSizes(allTileSizes.begin(), allTileSizes.end());
50 
51   // Traverse the tile sizes, which are in loop order, erase zeros everywhere.
52   LoopIndexToRangeIndexMap loopIndexToRangeIndex;
53   for (int idx = 0, e = tileSizes.size(), zerosCount = 0; idx < e; ++idx) {
54     if (isZero(tileSizes[idx - zerosCount])) {
55       shapeSizes.erase(shapeSizes.begin() + idx - zerosCount);
56       tileSizes.erase(tileSizes.begin() + idx - zerosCount);
57       ++zerosCount;
58       continue;
59     }
60     loopIndexToRangeIndex[idx] = idx - zerosCount;
61   }
62 
63   // Create a new range with the applied tile sizes.
64   SmallVector<Range, 4> res;
65   for (unsigned idx = 0, e = tileSizes.size(); idx < e; ++idx)
66     res.push_back(Range{b.create<arith::ConstantIndexOp>(loc, 0),
67                         shapeSizes[idx], tileSizes[idx]});
68   return std::make_tuple(res, loopIndexToRangeIndex);
69 }
70 
71 void mlir::linalg::transformIndexOps(
72     RewriterBase &b, LinalgOp op, SmallVectorImpl<Value> &ivs,
73     const LoopIndexToRangeIndexMap &loopIndexToRangeIndex) {
74   SmallVector<Value> allIvs(op.getNumLoops(), nullptr);
75   for (auto &en : enumerate(allIvs)) {
76     auto rangeIndex = loopIndexToRangeIndex.find(en.index());
77     if (rangeIndex == loopIndexToRangeIndex.end())
78       continue;
79     en.value() = ivs[rangeIndex->second];
80   }
81   addTileLoopIvsToIndexOpResults(b, op, allIvs);
82 }
83 
84 // Insert a tile `source` into the destination tensor `dest`. The position at
85 // which the tile is inserted (as well as size of tile) is taken from a given
86 // ExtractSliceOp `sliceOp`.
87 static Value insertSliceIntoTensor(RewriterBase &b, Location loc,
88                                    tensor::ExtractSliceOp sliceOp, Value source,
89                                    Value dest) {
90   return b.create<tensor::InsertSliceOp>(
91       loc, sliceOp.source().getType(), source, dest, sliceOp.offsets(),
92       sliceOp.sizes(), sliceOp.strides(), sliceOp.static_offsets(),
93       sliceOp.static_sizes(), sliceOp.static_strides());
94 }
95 
96 template <typename LoopTy>
97 static FailureOr<TiledLinalgOp>
98 tileLinalgOpImpl(RewriterBase &b, LinalgOp op, ValueRange tileSizes,
99                  const LinalgTilingOptions &options) {
100   auto nLoops = op.getNumLoops();
101   // Initial tile sizes may be too big, only take the first nLoops.
102   tileSizes = tileSizes.take_front(nLoops);
103 
104   if (llvm::all_of(tileSizes, isZero)) {
105     TiledLinalgOp tiledOp;
106     tiledOp.op = cast<LinalgOp>(b.clone(*op.getOperation()));
107     tiledOp.tensorResults.assign(tiledOp.op->result_begin(),
108                                  tiledOp.op->result_end());
109     return tiledOp;
110   }
111 
112   // 1. Build the tiled loop ranges.
113   auto allShapeSizes = op.createFlatListOfOperandDims(b, op.getLoc());
114   AffineMap shapeSizesToLoopsMap = op.getShapesToLoopsMap();
115   if (!shapeSizesToLoopsMap)
116     return failure();
117 
118   SmallVector<Range, 4> loopRanges;
119   LoopIndexToRangeIndexMap loopIndexToRangeIndex;
120   std::tie(loopRanges, loopIndexToRangeIndex) = makeTiledLoopRanges(
121       b, op.getLoc(), shapeSizesToLoopsMap, allShapeSizes, tileSizes);
122 
123   SmallVector<Attribute, 4> iteratorTypes;
124   for (const auto &attr :
125        enumerate(op.iterator_types().cast<ArrayAttr>().getValue())) {
126     if (loopIndexToRangeIndex.count(attr.index()))
127       iteratorTypes.push_back(attr.value());
128   }
129   // If interchangeVector is empty, use the identity. Build the permutation map
130   // otherwise.
131   auto invPermutationMap =
132       AffineMap::getMultiDimIdentityMap(tileSizes.size(), b.getContext());
133   if (!options.interchangeVector.empty()) {
134     // Based on the pruned iterations (due to zero tile size), recompute the
135     // interchange vector.
136     SmallVector<unsigned, 4> interchangeVector;
137     interchangeVector.reserve(options.interchangeVector.size());
138     for (auto pos : options.interchangeVector) {
139       auto it = loopIndexToRangeIndex.find(pos);
140       if (it == loopIndexToRangeIndex.end())
141         continue;
142       interchangeVector.push_back(it->second);
143     }
144     // Interchange vector is guaranteed to be a permutation,
145     // `inversePermutation` must succeed.
146     invPermutationMap = inversePermutation(
147         AffineMap::getPermutationMap(interchangeVector, b.getContext()));
148     assert(invPermutationMap);
149     SmallVector<int64_t> permutation(interchangeVector.begin(),
150                                      interchangeVector.end());
151     applyPermutationToVector(loopRanges, permutation);
152     applyPermutationToVector(iteratorTypes, permutation);
153   }
154 
155   // 2. Create the tiled loops.
156   LinalgOp res = op;
157   SmallVector<Value, 4> ivs, tensorResults;
158   auto tiledLoopBodyBuilder =
159       [&](OpBuilder &builder, Location loc, ValueRange localIvs,
160           ValueRange operandValuesToUse) -> scf::ValueVector {
161     ivs.assign(localIvs.begin(), localIvs.end());
162 
163     // When an `interchangeVector` is present, it has been applied to the
164     // loop ranges and the iterator types. Apply its inverse to the
165     // resulting loop `ivs` to match the op definition.
166     SmallVector<Value, 4> interchangedIvs;
167     if (!options.interchangeVector.empty())
168       interchangedIvs = applyMapToValues(b, loc, invPermutationMap, ivs);
169     else
170       interchangedIvs.assign(ivs.begin(), ivs.end());
171 
172     // Tile the `operandValuesToUse` that either match the `op` operands
173     // themselves or the tile loop arguments forwarding them.
174     assert(operandValuesToUse.size() ==
175                static_cast<size_t>(op.getNumInputsAndOutputs()) &&
176            "expect the number of operands and inputs and outputs to match");
177     SmallVector<Value> valuesToTile = operandValuesToUse;
178     auto sizeBounds =
179         applyMapToValues(b, loc, shapeSizesToLoopsMap, allShapeSizes);
180     SmallVector<Value, 4> tiledOperands = makeTiledShapes(
181         b, loc, op, valuesToTile, interchangedIvs, tileSizes, sizeBounds);
182 
183     // TODO: use an interface/adaptor to avoid leaking position in
184     // `tiledOperands`.
185     SmallVector<Type, 4> resultTensorTypes;
186     for (OpOperand *opOperand : op.getOutputTensorOperands())
187       resultTensorTypes.push_back(
188           tiledOperands[opOperand->getOperandNumber()].getType());
189 
190     res = op.clone(b, loc, resultTensorTypes, tiledOperands);
191 
192     // Insert a insert_slice for each output tensor.
193     unsigned resultIdx = 0;
194     for (OpOperand *opOperand : op.getOutputTensorOperands()) {
195       // TODO: use an interface/adaptor to avoid leaking position in
196       // `tiledOperands`.
197       Value outputTensor = tiledOperands[opOperand->getOperandNumber()];
198       // TODO: Propagate RewriterBase everywhere.
199       IRRewriter rewriter(b);
200       if (auto sliceOp = outputTensor.getDefiningOp<tensor::ExtractSliceOp>()) {
201         tensorResults.push_back(insertSliceIntoTensor(rewriter, loc, sliceOp,
202                                                       res->getResult(resultIdx),
203                                                       sliceOp.source()));
204       } else {
205         tensorResults.push_back(res->getResult(resultIdx));
206       }
207       ++resultIdx;
208     }
209     return scf::ValueVector(tensorResults.begin(), tensorResults.end());
210   };
211   GenerateLoopNest<LoopTy>::doit(b, op.getLoc(), loopRanges, op, iteratorTypes,
212                                  tiledLoopBodyBuilder, options.distribution,
213                                  options.distributionTypes);
214 
215   // 3. Transform IndexOp results w.r.t. the tiling.
216   transformIndexOps(b, res, ivs, loopIndexToRangeIndex);
217 
218   // 4. Gather the newly created loops and return them with the new op.
219   SmallVector<Operation *, 8> loops;
220   loops.reserve(ivs.size());
221   for (auto iv : ivs) {
222     if (iv.isa<BlockArgument>()) {
223       loops.push_back(iv.cast<BlockArgument>().getOwner()->getParentOp());
224       assert(loops.back() && "no owner found for induction variable!");
225     } else {
226       // TODO: Instead of doing this, try to recover the ops used instead of the
227       // loop.
228       loops.push_back(nullptr);
229     }
230   }
231 
232   // 5. Get the tensor results from the outermost loop if available. Otherwise
233   // use the previously captured `tensorResults`.
234   Operation *outermostLoop = nullptr;
235   for (Operation *loop : loops)
236     if ((outermostLoop = loop))
237       break;
238 
239   return TiledLinalgOp{
240       res, loops, outermostLoop ? outermostLoop->getResults() : tensorResults};
241 }
242 
243 template <typename LoopTy>
244 FailureOr<TiledLinalgOp> static tileLinalgOpImpl(
245     RewriterBase &b, LinalgOp op, const LinalgTilingOptions &options) {
246   OpBuilder::InsertionGuard g(b);
247   b.setInsertionPoint(op);
248 
249   if (!options.tileSizeComputationFunction)
250     return failure();
251 
252   // Enforce the convention that "tiling by zero" skips tiling a particular
253   // dimension. This convention is significantly simpler to handle instead of
254   // adjusting affine maps to account for missing dimensions.
255   auto nLoops = op.getNumLoops();
256   SmallVector<Value, 4> tileSizeVector =
257       options.tileSizeComputationFunction(b, op);
258   if (tileSizeVector.size() < nLoops) {
259     auto zero = b.create<arith::ConstantIndexOp>(op.getLoc(), 0);
260     tileSizeVector.append(nLoops - tileSizeVector.size(), zero);
261   }
262 
263   return tileLinalgOpImpl<LoopTy>(b, op, tileSizeVector, options);
264 }
265 
266 FailureOr<TiledLinalgOp>
267 mlir::linalg::tileLinalgOp(RewriterBase &b, LinalgOp op,
268                            const LinalgTilingOptions &options) {
269   switch (options.loopType) {
270   case LinalgTilingLoopType::Loops:
271     return tileLinalgOpImpl<scf::ForOp>(b, op, options);
272   case LinalgTilingLoopType::ParallelLoops:
273     return tileLinalgOpImpl<scf::ParallelOp>(b, op, options);
274   case LinalgTilingLoopType::TiledLoops:
275     return tileLinalgOpImpl<linalg::TiledLoopOp>(b, op, options);
276   default:;
277   }
278   return failure();
279 }
280 
281 /// Generate a loop nest around a given tensor::PadOp (for tiling). `newPadOp`
282 /// and `loopNest` are output parameters that return the new (tiled)
283 /// tensor::PadOp and the loop nest.
284 static LogicalResult tilePadOp(RewriterBase &builder, tensor::PadOp op,
285                                tensor::PadOp &newPadOp, LoopNest &loopNest,
286                                const LinalgTilingOptions &options) {
287   Location loc = op.getLoc();
288   OpBuilder::InsertionGuard g(builder);
289   builder.setInsertionPoint(op);
290 
291   // Clone tensor::PadOp so that the existing op can be replaced more easily.
292   newPadOp = cast<tensor::PadOp>(builder.clone(*op.getOperation()));
293   // Get rank and tile sizes.
294   int64_t rank = op.getResultType().getRank();
295   SmallVector<Value> tileSizes =
296       options.tileSizeComputationFunction(builder, op);
297   // Normalize untiled padding dimensions to 0.
298   Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
299   tileSizes.append(rank - tileSizes.size(), zero);
300   // Compute lower and upper bounds of the loop nest.
301   TilingInterface tilingInterface =
302       dyn_cast<TilingInterface>(op.getOperation());
303   SmallVector<Range> ranges = tilingInterface.getIterationDomain(builder);
304   SmallVector<Value> lbs, dims, allDims, steps;
305   for (int64_t i = 0; i < rank; ++i) {
306     allDims.push_back(ranges[i].size);
307     if (!isZero(tileSizes[i])) {
308       lbs.push_back(ranges[i].offset);
309       dims.push_back(ranges[i].size);
310       steps.push_back(tileSizes[i]);
311     }
312   }
313   // Generate loop nest: One loop per dimension.
314   SmallVector<Value> destOperand =
315       tilingInterface.getDestinationOperands(builder);
316   loopNest = mlir::scf::buildLoopNest(
317       builder, loc, lbs, /*ubs=*/dims, steps, ValueRange(destOperand),
318       [&](OpBuilder &b, Location loc, ValueRange localIvs,
319           ValueRange iterArgs) -> scf::ValueVector {
320         // Compute offsets and sizes of ExtractSliceOp.
321         SmallVector<Value> offsets =
322             computeTileOffsets(b, loc, localIvs, tileSizes);
323         SmallVector<Value> sizes =
324             computeTileSizes(b, loc, localIvs, tileSizes, allDims);
325         // Create ExtractSliceOp: Extract a tile from the tensor::PadOp.
326         // Note: The tensor::PadOp is located outside of the loop nest. It is
327         // later moved inside by ExtractSliceOfPadTensorSwapPattern.
328         auto map = AffineMap::getMultiDimIdentityMap(rank, b.getContext());
329         Value tiledOutput =
330             makeTiledShape(b, loc, newPadOp->getResult(0), tileSizes, map,
331                            offsets, allDims, sizes);
332         auto sliceOp = tiledOutput.getDefiningOp<tensor::ExtractSliceOp>();
333         assert(sliceOp && "expected ExtractSliceOp");
334         // Insert the tile into the output tensor.
335         // TODO: Propagate RewriterBase everywhere.
336         IRRewriter rewriter(b);
337         Value yieldValue =
338             insertSliceIntoTensor(rewriter, loc, sliceOp, sliceOp, iterArgs[0]);
339         return scf::ValueVector({yieldValue});
340       });
341   return success();
342 }
343 
344 namespace {
345 struct PadOpTilingPattern : public OpRewritePattern<tensor::PadOp> {
346   PadOpTilingPattern(MLIRContext *ctx, LinalgTilingOptions opt)
347       : OpRewritePattern<tensor::PadOp>(ctx), options(std::move(opt)) {}
348 
349   LogicalResult matchAndRewrite(tensor::PadOp op,
350                                 PatternRewriter &rewriter) const override {
351     if (op->hasAttr(LinalgTransforms::kLinalgTransformMarker))
352       return failure();
353     tensor::PadOp newPadOp;
354     LoopNest loopNest;
355     if (failed(tilePadOp(rewriter, op, newPadOp, loopNest, options)))
356       return failure();
357     newPadOp->setAttr(LinalgTransforms::kLinalgTransformMarker,
358                       rewriter.getUnitAttr());
359     // Replace all uses of the original tensor::PadOp.
360     rewriter.replaceOp(op, loopNest.getResults()[0]);
361     return success();
362   }
363 
364   LinalgTilingOptions options;
365 };
366 } // namespace
367 
368 namespace {
369 /// Helper classes for type list expansion.
370 template <typename... OpTypes>
371 class CanonicalizationPatternList;
372 
373 template <>
374 class CanonicalizationPatternList<> {
375 public:
376   static void insert(RewritePatternSet &patterns) {}
377 };
378 
379 template <typename OpTy, typename... OpTypes>
380 class CanonicalizationPatternList<OpTy, OpTypes...> {
381 public:
382   static void insert(RewritePatternSet &patterns) {
383     OpTy::getCanonicalizationPatterns(patterns, patterns.getContext());
384     CanonicalizationPatternList<OpTypes...>::insert(patterns);
385   }
386 };
387 } // namespace
388 
389 RewritePatternSet
390 mlir::linalg::getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx) {
391   RewritePatternSet patterns(ctx);
392   populateLinalgTilingCanonicalizationPatterns(patterns);
393   return patterns;
394 }
395 
396 void mlir::linalg::populateLinalgTilingCanonicalizationPatterns(
397     RewritePatternSet &patterns) {
398   auto *ctx = patterns.getContext();
399   AffineApplyOp::getCanonicalizationPatterns(patterns, ctx);
400   AffineForOp::getCanonicalizationPatterns(patterns, ctx);
401   AffineMinOp::getCanonicalizationPatterns(patterns, ctx);
402   AffineMaxOp::getCanonicalizationPatterns(patterns, ctx);
403   arith::ConstantIndexOp::getCanonicalizationPatterns(patterns, ctx);
404 
405   memref::SubViewOp::getCanonicalizationPatterns(patterns, ctx);
406   memref::ViewOp::getCanonicalizationPatterns(patterns, ctx);
407 
408   scf::ForOp::getCanonicalizationPatterns(patterns, ctx);
409   scf::ParallelOp::getCanonicalizationPatterns(patterns, ctx);
410 
411   tensor::CastOp::getCanonicalizationPatterns(patterns, ctx);
412   tensor::ExtractSliceOp::getCanonicalizationPatterns(patterns, ctx);
413   tensor::InsertSliceOp::getCanonicalizationPatterns(patterns, ctx);
414 
415   InitTensorOp::getCanonicalizationPatterns(patterns, ctx);
416   tensor::PadOp::getCanonicalizationPatterns(patterns, ctx);
417   ctx->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(patterns);
418 
419   CanonicalizationPatternList<
420 #define GET_OP_LIST
421 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
422       >::insert(patterns);
423 }
424 
425 /// Populate the given list with patterns that apply Linalg tiling.
426 static void insertTilingPatterns(RewritePatternSet &patterns,
427                                  const LinalgTilingOptions &options) {
428   auto *ctx = patterns.getContext();
429   LinalgTransformationFilter f(ArrayRef<StringAttr>{},
430                                StringAttr::get(ctx, "tiled"));
431   TilingPatterns<GenericOp,
432 #define GET_OP_LIST
433 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
434                  >::insert(patterns, options, f);
435   patterns.add<PadOpTilingPattern>(ctx, options);
436 }
437 
438 void mlir::linalg::populatePadTensorTilingPatterns(
439     RewritePatternSet &patterns, const LinalgTilingOptions &options) {
440   auto *ctx = patterns.getContext();
441   patterns.add<PadOpTilingPattern>(ctx, options);
442 }
443 
444 static void applyExtractSliceOfPadTensorSwapPattern(FuncOp funcOp) {
445   MLIRContext *ctx = funcOp.getContext();
446   RewritePatternSet patterns(ctx);
447   patterns.add<ExtractSliceOfPadTensorSwapPattern>(patterns.getContext());
448   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
449   (void)applyPatternsAndFoldGreedily(
450       funcOp, getLinalgTilingCanonicalizationPatterns(ctx));
451 }
452 
453 namespace {
454 struct LinalgTilingPass : public LinalgTilingBase<LinalgTilingPass> {
455   LinalgTilingPass() = default;
456   LinalgTilingPass(ArrayRef<int64_t> tileSizes, LinalgTilingLoopType loopType,
457                    ArrayRef<StringRef> distributionTypes) {
458     this->tileSizes = tileSizes;
459     this->loopType = "";
460     this->loopTypeEnum = loopType;
461     this->distributionTypes = llvm::to_vector<2>(llvm::map_range(
462         distributionTypes, [](StringRef ref) { return ref.str(); }));
463   }
464 
465   void runOnOperation() override {
466     FuncOp funcOp = getOperation();
467     LinalgTilingLoopType type =
468         llvm::StringSwitch<LinalgTilingLoopType>(loopType)
469             .Case("for", LinalgTilingLoopType::Loops)
470             .Case("affine", LinalgTilingLoopType::AffineLoops)
471             .Case("parallel", LinalgTilingLoopType::ParallelLoops)
472             .Case("tiled_loop", LinalgTilingLoopType::TiledLoops)
473             .Default(loopTypeEnum);
474     auto distTypes = llvm::to_vector<2>(llvm::map_range(
475         distributionTypes, [](std::string &str) { return StringRef(str); }));
476     auto options = LinalgTilingOptions()
477                        .setTileSizes(tileSizes)
478                        .setLoopType(type)
479                        .setDistributionTypes(distTypes);
480     MLIRContext *ctx = funcOp.getContext();
481     RewritePatternSet patterns(ctx);
482     insertTilingPatterns(patterns, options);
483     scf::populateSCFForLoopCanonicalizationPatterns(patterns);
484     (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
485     (void)applyPatternsAndFoldGreedily(
486         funcOp, getLinalgTilingCanonicalizationPatterns(ctx));
487     // Drop the marker.
488     funcOp.walk([](LinalgOp op) {
489       op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
490     });
491 
492     // Apply swap pattern after generating loop nest and running
493     // canonicalizations.
494     applyExtractSliceOfPadTensorSwapPattern(funcOp);
495   }
496 
497   LinalgTilingLoopType loopTypeEnum;
498 };
499 
500 } // namespace
501 
502 std::unique_ptr<OperationPass<FuncOp>>
503 mlir::createLinalgTilingPass(ArrayRef<int64_t> tileSizes,
504                              linalg::LinalgTilingLoopType loopType,
505                              ArrayRef<StringRef> distributionTypes) {
506   return std::make_unique<LinalgTilingPass>(tileSizes, loopType,
507                                             distributionTypes);
508 }
509