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 "PassDetail.h"
14 #include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
15 #include "mlir/Dialect/Linalg/EDSC/FoldedIntrinsics.h"
16 #include "mlir/Dialect/Linalg/IR/LinalgTypes.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/SCF/EDSC/Builders.h"
21 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
22 #include "mlir/IR/AffineExpr.h"
23 #include "mlir/IR/AffineExprVisitor.h"
24 #include "mlir/IR/AffineMap.h"
25 #include "mlir/Support/LLVM.h"
26 #include "mlir/Transforms/FoldUtils.h"
27 
28 #include "llvm/Support/CommandLine.h"
29 
30 using namespace mlir;
31 using namespace mlir::edsc;
32 using namespace mlir::edsc::intrinsics;
33 using namespace mlir::linalg;
34 using namespace mlir::scf;
35 
36 using folded_affine_min = FoldedValueBuilder<AffineMinOp>;
37 
38 #define DEBUG_TYPE "linalg-tiling"
39 
40 static bool isZero(Value v) {
41   if (auto cst = v.getDefiningOp<ConstantIndexOp>())
42     return cst.getValue() == 0;
43   return false;
44 }
45 
46 using LoopIndexToRangeIndexMap = DenseMap<int, int>;
47 
48 // Creates a number of ranges equal to the number of non-zero in `tileSizes`.
49 // One for each loop of the LinalgOp that is tiled. The `tileSizes` argument has
50 // one entry per surrounding loop. It uses zero as the convention that a
51 // particular loop is not tiled. This convention simplifies implementations by
52 // avoiding affine map manipulations.
53 // The returned ranges correspond to the loop ranges, in the proper order, that
54 // are tiled and for which new loops will be created. Also the function returns
55 // a map from loop indices of the LinalgOp to the corresponding non-empty range
56 // indices of newly created loops.
57 static std::tuple<SmallVector<Range, 4>, LoopIndexToRangeIndexMap>
58 makeTiledLoopRanges(OpBuilder &b, Location loc, AffineMap map,
59                     ArrayRef<Value> allViewSizes,
60                     ArrayRef<Value> allTileSizes) {
61   assert(allTileSizes.size() == map.getNumResults());
62   // Apply `map` to get view sizes in loop order.
63   auto viewSizes = applyMapToValues(b, loc, map, allViewSizes);
64   SmallVector<Value, 4> tileSizes(allTileSizes.begin(), allTileSizes.end());
65 
66   // Traverse the tile sizes, which are in loop order, erase zeros everywhere.
67   LoopIndexToRangeIndexMap loopIndexToRangeIndex;
68   for (int idx = 0, e = tileSizes.size(), zerosCount = 0; idx < e; ++idx) {
69     if (isZero(tileSizes[idx - zerosCount])) {
70       viewSizes.erase(viewSizes.begin() + idx - zerosCount);
71       tileSizes.erase(tileSizes.begin() + idx - zerosCount);
72       ++zerosCount;
73       continue;
74     }
75     loopIndexToRangeIndex[idx] = idx - zerosCount;
76   }
77 
78   // Create a new range with the applied tile sizes.
79   SmallVector<Range, 4> res;
80   for (unsigned idx = 0, e = tileSizes.size(); idx < e; ++idx)
81     res.push_back(Range{std_constant_index(0), viewSizes[idx], tileSizes[idx]});
82   return std::make_tuple(res, loopIndexToRangeIndex);
83 }
84 
85 namespace {
86 
87 // Helper visitor to determine whether an AffineExpr is tiled.
88 // This is achieved by traversing every AffineDimExpr with position `pos` and
89 // checking whether the corresponding `tileSizes[pos]` is non-zero.
90 // This also enforces only positive coefficients occur in multiplications.
91 //
92 // Example:
93 //   `d0 + 2 * d1 + d3` is tiled by [0, 0, 0, 2] but not by [0, 0, 2, 0]
94 //
95 struct TileCheck : public AffineExprVisitor<TileCheck> {
96   TileCheck(ArrayRef<Value> tileSizes) : isTiled(false), tileSizes(tileSizes) {}
97 
98   void visitDimExpr(AffineDimExpr expr) {
99     isTiled |= !isZero(tileSizes[expr.getPosition()]);
100   }
101   void visitAffineBinaryOpExpr(AffineBinaryOpExpr expr) {
102     visit(expr.getLHS());
103     visit(expr.getRHS());
104     if (expr.getKind() == mlir::AffineExprKind::Mul)
105       assert(expr.getRHS().cast<AffineConstantExpr>().getValue() > 0 &&
106              "nonpositive multiplying coefficient");
107   }
108   bool isTiled;
109   ArrayRef<Value> tileSizes;
110 };
111 
112 } // namespace
113 
114 // IndexedGenericOp explicitly uses induction variables in the loop body. The
115 // values of the indices that are used in the loop body for any given access of
116 // input/output memref before `subview` op was applied should be invariant with
117 // respect to tiling.
118 //
119 // Therefore, if the operation is tiled, we have to transform the indices
120 // accordingly, i.e. offset them by the values of the corresponding induction
121 // variables that are captured implicitly in the body of the op.
122 //
123 // Example. `linalg.indexed_generic` before tiling:
124 //
125 // #id_2d = (i, j) -> (i, j)
126 // #pointwise_2d_trait = {
127 //   indexing_maps = [#id_2d, #id_2d],
128 //   iterator_types = ["parallel", "parallel"],
129 //   n_views = [1, 1]
130 // }
131 // linalg.indexed_generic #pointwise_2d_trait %operand, %result {
132 //   ^bb0(%i: index, %j: index, %operand_in: f32, %result_in: f32):
133 //     <some operations that use %i, %j>
134 // }: memref<50x100xf32>, memref<50x100xf32>
135 //
136 // After tiling pass with tiles sizes 10 and 25:
137 //
138 // #strided = (i, j)[s0, s1, s2] -> (i * s1 + s0 + j * s2)
139 //
140 // %c1 = constant 1 : index
141 // %c0 = constant 0 : index
142 // %c25 = constant 25 : index
143 // %c10 = constant 10 : index
144 // operand_dim_0 = dim %operand, 0 : memref<50x100xf32>
145 // operand_dim_1 = dim %operand, 1 : memref<50x100xf32>
146 // scf.for %k = %c0 to operand_dim_0 step %c10 {
147 //   scf.for %l = %c0 to operand_dim_1 step %c25 {
148 //     %4 = std.subview %operand[%k, %l][%c10, %c25][%c1, %c1]
149 //       : memref<50x100xf32> to memref<?x?xf32, #strided>
150 //     %5 = std.subview %result[%k, %l][%c10, %c25][%c1, %c1]
151 //       : memref<50x100xf32> to memref<?x?xf32, #strided>
152 //     linalg.indexed_generic pointwise_2d_trait %4, %5 {
153 //     ^bb0(%i: index, %j: index, %operand_in: f32, %result_in: f32):
154 //       // Indices `k` and `l` are implicitly captured in the body.
155 //       %transformed_i = addi %i, %k : index // index `i` is offset by %k
156 //       %transformed_j = addi %j, %l : index // index `j` is offset by %l
157 //       // Every use of %i, %j is replaced with %transformed_i, %transformed_j
158 //       <some operations that use %transformed_i, %transformed_j>
159 //     }: memref<?x?xf32, #strided>, memref<?x?xf32, #strided>
160 //   }
161 // }
162 //
163 // TODO: Investigate whether mixing implicit and explicit indices
164 // does not lead to losing information.
165 static void transformIndexedGenericOpIndices(
166     OpBuilder &b, LinalgOp op, SmallVectorImpl<Value> &ivs,
167     const LoopIndexToRangeIndexMap &loopIndexToRangeIndex) {
168   assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics");
169   auto indexedGenericOp = dyn_cast<IndexedGenericOp>(op.getOperation());
170   if (!indexedGenericOp)
171     return;
172 
173   // `linalg.indexed_generic` comes in two flavours. One has a region with a
174   // single block that defines the loop body. The other has a `fun` attribute
175   // that refers to an existing function symbol. The `fun` function call will be
176   // inserted in the loop body in that case.
177   //
178   // TODO: Add support for `linalg.indexed_generic` with `fun` attribute.
179   auto &region = indexedGenericOp.region();
180   if (region.empty()) {
181     indexedGenericOp.emitOpError("expected a region");
182     return;
183   }
184   auto &block = region.front();
185 
186   OpBuilder::InsertionGuard g(b);
187   b.setInsertionPointToStart(&block);
188   for (unsigned i = 0; i < indexedGenericOp.getNumLoops(); ++i) {
189     auto rangeIndex = loopIndexToRangeIndex.find(i);
190     if (rangeIndex == loopIndexToRangeIndex.end())
191       continue;
192     Value oldIndex = block.getArgument(i);
193     // Offset the index argument `i` by the value of the corresponding induction
194     // variable and replace all uses of the previous value.
195     Value newIndex = b.create<AddIOp>(indexedGenericOp.getLoc(), oldIndex,
196                                       ivs[rangeIndex->second]);
197     for (auto &use : oldIndex.getUses()) {
198       if (use.getOwner() == newIndex.getDefiningOp())
199         continue;
200       use.set(newIndex);
201     }
202   }
203 }
204 
205 static bool isTiled(AffineExpr expr, ArrayRef<Value> tileSizes) {
206   if (!expr)
207     return false;
208   TileCheck t(tileSizes);
209   t.visit(expr);
210   return t.isTiled;
211 }
212 
213 // Checks whether the view with index `viewIndex` within `linalgOp` varies with
214 // respect to a non-zero `tileSize`.
215 static bool isTiled(AffineMap map, ArrayRef<Value> tileSizes) {
216   if (!map)
217     return false;
218   for (unsigned r = 0; r < map.getNumResults(); ++r)
219     if (isTiled(map.getResult(r), tileSizes))
220       return true;
221   return false;
222 }
223 
224 static SmallVector<Value, 4> makeTiledViews(OpBuilder &b, Location loc,
225                                             LinalgOp linalgOp, AffineMap map,
226                                             ArrayRef<Value> ivs,
227                                             ArrayRef<Value> tileSizes,
228                                             ArrayRef<Value> allViewSizes) {
229   assert(linalgOp.hasBufferSemantics() &&
230          "expected linalg op with buffer semantics");
231   assert(ivs.size() == static_cast<size_t>(llvm::count_if(
232                            llvm::make_range(tileSizes.begin(), tileSizes.end()),
233                            [](Value v) { return !isZero(v); })) &&
234          "expected as many ivs as non-zero sizes");
235 
236   using namespace edsc::op;
237 
238   auto viewSizes = applyMapToValues(b, loc, map, allViewSizes);
239   // Construct (potentially temporary) mins and maxes on which to apply maps
240   // that define tile subviews.
241   SmallVector<Value, 8> lbs, subViewSizes;
242   for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) {
243     bool isTiled = !isZero(tileSizes[idx]);
244     lbs.push_back(isTiled ? ivs[idxIvs++] : (Value)std_constant_index(0));
245     // Before composing, we need to make range a closed interval.
246     Value size = isTiled ? tileSizes[idx] : viewSizes[idx];
247     subViewSizes.push_back(size - std_constant_index(1));
248   }
249 
250   auto *op = linalgOp.getOperation();
251 
252   SmallVector<Value, 4> res;
253   res.reserve(op->getNumOperands());
254   auto viewIteratorBegin = linalgOp.getInputsAndOutputBuffers().begin();
255   for (unsigned viewIndex = 0; viewIndex < linalgOp.getNumInputsAndOutputs();
256        ++viewIndex) {
257     Value view = *(viewIteratorBegin + viewIndex);
258     auto viewType = view.getType().cast<MemRefType>();
259     unsigned rank = viewType.getRank();
260     auto mapAttr = linalgOp.indexing_maps()[viewIndex];
261     auto map = mapAttr.cast<AffineMapAttr>().getValue();
262     // If the view is not tiled, we can use it as is.
263     if (!isTiled(map, tileSizes)) {
264       res.push_back(view);
265       continue;
266     }
267 
268     // Construct a new subview for the tile.
269     SmallVector<Value, 4> offsets, sizes, strides;
270     offsets.reserve(rank);
271     sizes.reserve(rank);
272     strides.reserve(rank);
273     for (unsigned r = 0; r < rank; ++r) {
274       if (!isTiled(map.getSubMap({r}), tileSizes)) {
275         offsets.push_back(std_constant_index(0));
276         sizes.push_back(std_dim(view, r));
277         strides.push_back(std_constant_index(1));
278         continue;
279       }
280 
281       // Tiling creates a new slice at the proper index, the slice step is 1
282       // (i.e. the slice view does not subsample, stepping occurs in the loop).
283       auto m = map.getSubMap({r});
284       auto offset = applyMapToValues(b, loc, m, lbs).front();
285       offsets.push_back(offset);
286       auto closedIntSize = applyMapToValues(b, loc, m, subViewSizes).front();
287       // Resulting size needs to be made half open interval again.
288       auto size = closedIntSize + std_constant_index(1);
289 
290       // The size of the subview should be trimmed to avoid out-of-bounds
291       // accesses, unless we statically know the subview size divides the view
292       // size evenly.
293       int64_t viewSize = viewType.getDimSize(r);
294       auto sizeCst = size.getDefiningOp<ConstantIndexOp>();
295       if (ShapedType::isDynamic(viewSize) || !sizeCst ||
296           (viewSize % sizeCst.getValue()) != 0) {
297         // Compute min(size, dim - offset) to avoid out-of-bounds accesses.
298         auto minMap = AffineMap::get(
299             /*dimCount=*/3, /*symbolCount=*/0,
300             {getAffineDimExpr(/*position=*/0, b.getContext()),
301              getAffineDimExpr(/*position=*/1, b.getContext()) -
302                  getAffineDimExpr(/*position=*/2, b.getContext())},
303             b.getContext());
304         auto d = std_dim(view, r);
305         size =
306             affine_min(b.getIndexType(), minMap, ValueRange{size, d, offset});
307       }
308 
309       sizes.push_back(size);
310       strides.push_back(std_constant_index(1));
311     }
312 
313     res.push_back(b.create<SubViewOp>(loc, view, offsets, sizes, strides));
314   }
315 
316   return res;
317 }
318 
319 template <typename LoopTy>
320 static Optional<TiledLinalgOp>
321 tileLinalgOpImpl(OpBuilder &b, LinalgOp op, ArrayRef<Value> tileSizes,
322                  const LinalgTilingOptions &options) {
323   auto nLoops = op.getNumLoops();
324   // Initial tile sizes may be too big, only take the first nLoops.
325   tileSizes = tileSizes.take_front(nLoops);
326 
327   if (llvm::all_of(tileSizes, isZero))
328     return llvm::None;
329 
330   if (auto convOp = dyn_cast<linalg::ConvOp>(op.getOperation())) {
331     // For conv op only support tiling along batch dimension (which is the first
332     // loop).
333     if (convOp.padding() && !llvm::all_of(tileSizes.drop_front(), isZero))
334       return llvm::None;
335   }
336 
337   // 1. Build the tiled loop ranges.
338   auto allViewSizes = getViewSizes(b, op);
339   // The flattened loopToOperandRangesMaps is expected to be an invertible
340   // permutation map (asserted in the inverse calculation).
341   auto mapsRange = op.indexing_maps().getAsRange<AffineMapAttr>();
342   auto maps = llvm::to_vector<8>(
343       llvm::map_range(mapsRange, [](AffineMapAttr a) { return a.getValue(); }));
344   auto viewSizesToLoopsMap = inversePermutation(concatAffineMaps(maps));
345   if (!viewSizesToLoopsMap)
346     return llvm::None;
347 
348   SmallVector<Range, 4> loopRanges;
349   LoopIndexToRangeIndexMap loopIndexToRangeIndex;
350   std::tie(loopRanges, loopIndexToRangeIndex) = makeTiledLoopRanges(
351       b, op.getLoc(), viewSizesToLoopsMap, allViewSizes, tileSizes);
352   SmallVector<Attribute, 4> iteratorTypes;
353   for (auto attr :
354        enumerate(op.iterator_types().cast<ArrayAttr>().getValue())) {
355     if (loopIndexToRangeIndex.count(attr.index()))
356       iteratorTypes.push_back(attr.value());
357   }
358   // If interchangeVector is empty, use the identity. Build the permutation map
359   // otherwise.
360   auto invPermutationMap =
361       AffineMap::getMultiDimIdentityMap(tileSizes.size(), b.getContext());
362   if (!options.interchangeVector.empty()) {
363     // Based on the pruned iterations (due to zero tile size), recompute the
364     // interchange vector.
365     SmallVector<unsigned, 4> interchangeVector;
366     interchangeVector.reserve(options.interchangeVector.size());
367     for (auto pos : options.interchangeVector) {
368       auto it = loopIndexToRangeIndex.find(pos);
369       if (it == loopIndexToRangeIndex.end())
370         continue;
371       interchangeVector.push_back(it->second);
372     }
373     invPermutationMap = inversePermutation(
374         AffineMap::getPermutationMap(interchangeVector, b.getContext()));
375     if (!invPermutationMap)
376       return llvm::None;
377     applyPermutationToVector(loopRanges, interchangeVector);
378     applyPermutationToVector(iteratorTypes, interchangeVector);
379   }
380 
381   // 2. Create the tiled loops.
382   LinalgOp res = op;
383   SmallVector<Value, 4> ivs;
384   GenerateLoopNest<LoopTy>::doit(
385       loopRanges, /*iterArgInitValues*/ {}, iteratorTypes,
386       [&](ValueRange localIvs, ValueRange iterArgs) -> scf::ValueVector {
387         auto &b = ScopedContext::getBuilderRef();
388         auto loc = ScopedContext::getLocation();
389         ivs.assign(localIvs.begin(), localIvs.end());
390         SmallVector<Value, 4> ivValues(ivs.begin(), ivs.end());
391 
392         // If we have to apply a permutation to the tiled loop nest, we have to
393         // reorder the induction variables This permutation is the right one
394         // assuming that loopRanges have previously been permuted by
395         // (i,j,k)->(k,i,j) So this permutation should be the inversePermutation
396         // of that one: (d0,d1,d2)->(d2,d0,d1)
397         if (!options.interchangeVector.empty())
398           ivValues = applyMapToValues(b, loc, invPermutationMap, ivValues);
399 
400         auto views = makeTiledViews(b, loc, op, viewSizesToLoopsMap, ivValues,
401                                     tileSizes, allViewSizes);
402         auto operands = getAssumedNonViewOperands(op);
403         views.append(operands.begin(), operands.end());
404         res = op.clone(b, loc, /*resultTypes*/ {}, views);
405         return scf::ValueVector{};
406       },
407       options.distribution);
408 
409   // 3. Transforms index arguments of `linalg.generic` w.r.t. to the tiling.
410   transformIndexedGenericOpIndices(b, res, ivs, loopIndexToRangeIndex);
411 
412   // 4. Gather the newly created loops and return them with the new op.
413   SmallVector<Operation *, 8> loops;
414   loops.reserve(ivs.size());
415   for (auto iv : ivs) {
416     if (iv.isa<BlockArgument>()) {
417       loops.push_back(iv.cast<BlockArgument>().getOwner()->getParentOp());
418       assert(loops.back() && "no owner found for induction variable!");
419     } else {
420       // TODO: Instead of doing this, try to recover the ops used instead of the
421       // loop.
422       loops.push_back(nullptr);
423     }
424   }
425   return TiledLinalgOp{res, loops};
426 }
427 
428 template <typename LoopTy>
429 Optional<TiledLinalgOp> static tileLinalgOpImpl(
430     OpBuilder &b, LinalgOp op, const LinalgTilingOptions &options) {
431   OpBuilder::InsertionGuard g(b);
432   b.setInsertionPoint(op);
433   ScopedContext scope(b, op.getLoc());
434 
435   assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics");
436   // Enforce the convention that "tiling by zero" skips tiling a particular
437   // dimension. This convention is significantly simpler to handle instead of
438   // adjusting affine maps to account for missing dimensions.
439   auto nLoops = op.getNumLoops();
440   SmallVector<Value, 4> tileSizeVector =
441       options.tileSizeComputationFunction(b, op);
442   if (tileSizeVector.size() < nLoops) {
443     auto zero = std_constant_index(0);
444     tileSizeVector.append(nLoops - tileSizeVector.size(), zero);
445   }
446 
447   return tileLinalgOpImpl<LoopTy>(b, op, tileSizeVector, options);
448 }
449 
450 Optional<TiledLinalgOp>
451 mlir::linalg::tileLinalgOp(OpBuilder &b, LinalgOp op,
452                            const LinalgTilingOptions &options) {
453   switch (options.loopType) {
454   case LinalgTilingLoopType::Loops:
455     return tileLinalgOpImpl<scf::ForOp>(b, op, options);
456   case LinalgTilingLoopType::ParallelLoops:
457     return tileLinalgOpImpl<scf::ParallelOp>(b, op, options);
458   default:;
459   }
460   return llvm::None;
461 }
462 
463 namespace {
464 /// Helper classes for type list expansion.
465 template <typename... OpTypes>
466 class CanonicalizationPatternList;
467 
468 template <>
469 class CanonicalizationPatternList<> {
470 public:
471   static void insert(OwningRewritePatternList &patterns, MLIRContext *ctx) {}
472 };
473 
474 template <typename OpTy, typename... OpTypes>
475 class CanonicalizationPatternList<OpTy, OpTypes...> {
476 public:
477   static void insert(OwningRewritePatternList &patterns, MLIRContext *ctx) {
478     OpTy::getCanonicalizationPatterns(patterns, ctx);
479     CanonicalizationPatternList<OpTypes...>::insert(patterns, ctx);
480   }
481 };
482 
483 /// Helper classes for type list expansion.
484 template <typename... OpTypes>
485 class RewritePatternList;
486 
487 template <>
488 class RewritePatternList<> {
489 public:
490   static void insert(OwningRewritePatternList &patterns,
491                      const LinalgTilingOptions &options, MLIRContext *ctx) {}
492 };
493 
494 template <typename OpTy, typename... OpTypes>
495 class RewritePatternList<OpTy, OpTypes...> {
496 public:
497   static void insert(OwningRewritePatternList &patterns,
498                      const LinalgTilingOptions &options, MLIRContext *ctx) {
499     patterns.insert<LinalgTilingPattern<OpTy>>(
500         ctx, options, LinalgMarker({}, Identifier::get("tiled", ctx)));
501     RewritePatternList<OpTypes...>::insert(patterns, options, ctx);
502   }
503 };
504 } // namespace
505 
506 OwningRewritePatternList
507 mlir::linalg::getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx) {
508   OwningRewritePatternList patterns;
509   AffineApplyOp::getCanonicalizationPatterns(patterns, ctx);
510   AffineForOp::getCanonicalizationPatterns(patterns, ctx);
511   AffineMinOp::getCanonicalizationPatterns(patterns, ctx);
512   AffineMaxOp::getCanonicalizationPatterns(patterns, ctx);
513   scf::ForOp::getCanonicalizationPatterns(patterns, ctx);
514   scf::ParallelOp::getCanonicalizationPatterns(patterns, ctx);
515   ConstantIndexOp::getCanonicalizationPatterns(patterns, ctx);
516   SubViewOp::getCanonicalizationPatterns(patterns, ctx);
517   ViewOp::getCanonicalizationPatterns(patterns, ctx);
518   CanonicalizationPatternList<
519 #define GET_OP_LIST
520 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
521       >::insert(patterns, ctx);
522   return patterns;
523 }
524 
525 /// Populate the given list with patterns that apply Linalg tiling.
526 static void insertTilingPatterns(OwningRewritePatternList &patterns,
527                                  const LinalgTilingOptions &options,
528                                  MLIRContext *ctx) {
529   RewritePatternList<
530 #define GET_OP_LIST
531 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
532       >::insert(patterns, options, ctx);
533 }
534 
535 static void applyTilingToLoopPatterns(LinalgTilingLoopType loopType,
536                                       FuncOp funcOp,
537                                       ArrayRef<int64_t> tileSizes) {
538   auto options =
539       LinalgTilingOptions().setTileSizes(tileSizes).setLoopType(loopType);
540   MLIRContext *ctx = funcOp.getContext();
541   OwningRewritePatternList patterns;
542   insertTilingPatterns(patterns, options, ctx);
543   applyPatternsAndFoldGreedily(funcOp, patterns);
544   applyPatternsAndFoldGreedily(funcOp,
545                                getLinalgTilingCanonicalizationPatterns(ctx));
546   // Drop the marker.
547   funcOp.walk([](LinalgOp op) {
548     op.removeAttr(LinalgTransforms::kLinalgTransformMarker);
549   });
550 }
551 
552 namespace {
553 struct LinalgTilingPass : public LinalgTilingBase<LinalgTilingPass> {
554   LinalgTilingPass() = default;
555   LinalgTilingPass(ArrayRef<int64_t> sizes) { tileSizes = sizes; }
556 
557   void runOnFunction() override {
558     applyTilingToLoopPatterns(LinalgTilingLoopType::Loops, getFunction(),
559                               tileSizes);
560   }
561 };
562 
563 struct LinalgTilingToParallelLoopsPass
564     : public LinalgTilingToParallelLoopsBase<LinalgTilingToParallelLoopsPass> {
565   LinalgTilingToParallelLoopsPass() = default;
566   LinalgTilingToParallelLoopsPass(ArrayRef<int64_t> sizes) {
567     tileSizes = sizes;
568   }
569 
570   void runOnFunction() override {
571     applyTilingToLoopPatterns(LinalgTilingLoopType::ParallelLoops,
572                               getFunction(), tileSizes);
573   }
574 };
575 
576 } // namespace
577 
578 std::unique_ptr<OperationPass<FuncOp>>
579 mlir::createLinalgTilingPass(ArrayRef<int64_t> tileSizes) {
580   return std::make_unique<LinalgTilingPass>(tileSizes);
581 }
582 
583 std::unique_ptr<OperationPass<FuncOp>>
584 mlir::createLinalgTilingToParallelLoopsPass(ArrayRef<int64_t> tileSizes) {
585   return std::make_unique<LinalgTilingToParallelLoopsPass>(tileSizes);
586 }
587