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