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