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 ®ion = 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