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