1 //===- Fusion.cpp - Implementation of linalg Fusion -----------------------===// 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 Fusion pass. 10 // 11 //===----------------------------------------------------------------------===// 12 13 #include "PassDetail.h" 14 #include "mlir/Dialect/Affine/IR/AffineOps.h" 15 #include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h" 16 #include "mlir/Dialect/Linalg/EDSC/FoldedIntrinsics.h" 17 #include "mlir/Dialect/Linalg/IR/LinalgOps.h" 18 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h" 19 #include "mlir/Dialect/Linalg/Passes.h" 20 #include "mlir/Dialect/Linalg/Utils/Utils.h" 21 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h" 22 #include "mlir/IR/AffineExpr.h" 23 #include "mlir/IR/AffineMap.h" 24 #include "mlir/IR/Dominance.h" 25 #include "mlir/IR/PatternMatch.h" 26 #include "mlir/Support/LLVM.h" 27 #include "mlir/Transforms/FoldUtils.h" 28 #include "llvm/ADT/SetVector.h" 29 #include "llvm/Support/CommandLine.h" 30 #include "llvm/Support/Debug.h" 31 32 #define DEBUG_TYPE "linalg-fusion" 33 34 using namespace mlir; 35 using namespace mlir::edsc; 36 using namespace mlir::edsc::intrinsics; 37 using namespace mlir::linalg; 38 39 using folded_std_constant_index = FoldedValueBuilder<ConstantIndexOp>; 40 41 using llvm::dbgs; 42 43 /// Implements a simple high-level fusion pass of linalg library operations. 44 /// 45 /// In each block, linalg ops are processed in reverse textual order. 46 /// Given a linalg op `O`, fusion occurs by: 47 /// 1. inspecting the linalg ops that write into the views read by `O`. This 48 /// uses the SSA value of the views and a simple subview/slice analysis to 49 /// determine producer-consumer dependences; 50 /// 2. greedily fuse the linalg ops that produce subview 51 /// 3. inspect the fused ops and determine whether they have other remaining 52 /// LinalgOp uses. If not, then erase the original producing linalg op. 53 /// 54 /// More advanced use cases, analyses as well as profitability heuristics are 55 /// left for future work. 56 57 // Return a cloned version of `op` that operates on `loopRanges`, assumed to be 58 // a subset of the original loop ranges of `op`. 59 // This is achieved by applying the `loopToOperandRangesMaps` permutation maps 60 // to the `loopRanges` in order to obtain view ranges. 61 static LinalgOp cloneWithLoopRanges(OpBuilder &b, Location loc, LinalgOp op, 62 ArrayRef<SubViewOp::Range> loopRanges) { 63 assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics"); 64 auto maps = op.indexing_maps(); 65 SmallVector<Value, 8> clonedViews; 66 clonedViews.reserve(op.getNumInputsAndOutputs()); 67 // Iterate over the inputs and outputs in order. 68 // Extract the subranges from the linearized ranges. 69 SmallVector<Value, 8> ios(op.getInputsAndOutputBuffers()); 70 for (auto en : llvm::enumerate(ios)) { 71 unsigned idx = en.index(); 72 auto map = maps[idx].cast<AffineMapAttr>().getValue(); 73 LLVM_DEBUG(dbgs() << "map: " << map << "\n"); 74 Value view = en.value(); 75 SmallVector<SubViewOp::Range, 4> viewRanges(map.getNumResults()); 76 for (auto en2 : llvm::enumerate(map.getResults())) { 77 unsigned d = en2.index(); 78 // loopToOperandRangesMaps are permutations-only. 79 unsigned loopPos = en2.value().cast<AffineDimExpr>().getPosition(); 80 viewRanges[d] = loopRanges[loopPos]; 81 LLVM_DEBUG(dbgs() << "\ni,j: " << en.index() << ", " << en2.index() 82 << "\t" 83 << "loopPos: " << loopPos << "\t" << viewRanges[d]); 84 } 85 // Construct a new subview for the tile. 86 unsigned rank = viewRanges.size(); 87 SmallVector<Value, 4> offsets, sizes, strides; 88 offsets.reserve(rank); 89 sizes.reserve(rank); 90 strides.reserve(rank); 91 for (auto r : viewRanges) { 92 offsets.push_back(r.offset); 93 sizes.push_back(r.size); 94 strides.push_back(r.stride); 95 } 96 clonedViews.push_back( 97 b.create<SubViewOp>(loc, view, offsets, sizes, strides)); 98 } 99 auto operands = getAssumedNonViewOperands(op); 100 clonedViews.append(operands.begin(), operands.end()); 101 102 Operation *clonedOp = op.clone(b, loc, clonedViews); 103 // When the producer is an IndexedGenercOp, we have to transform its block 104 // IV arguments according to the tiling of the consumer, i.e. offset them by 105 // the values computed in `loopRanges`. 106 if (auto indexedGenericOp = dyn_cast<IndexedGenericOp>(clonedOp)) { 107 auto &block = indexedGenericOp.region().front(); 108 109 OpBuilder::InsertionGuard g(b); 110 b.setInsertionPointToStart(&block); 111 for (unsigned i = 0, e = indexedGenericOp.getNumLoops(); i < e; ++i) { 112 Value oldIndex = block.getArgument(i); 113 AddIOp newIndex = b.create<AddIOp>(indexedGenericOp.getLoc(), oldIndex, 114 loopRanges[i].offset); 115 oldIndex.replaceAllUsesExcept(newIndex, 116 SmallPtrSet<Operation *, 1>{newIndex}); 117 } 118 } 119 return clonedOp; 120 } 121 122 struct ViewDimension { 123 Value view; 124 unsigned dimension; 125 }; 126 127 // Given an `op`, returns the first (`view`, `dimension`) pair that identifies 128 // the loop range at `loopDepth`. The semantics of the loopToOperandRangesMaps 129 // guarantees at least one such dimension is found. If multiple candidates exist 130 // they must agree by construction (i.e. have the same size) and we just return 131 // the first one. 132 static ViewDimension getViewDefiningLoopRange(LinalgOp op, unsigned loopDepth) { 133 assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics"); 134 auto maps = op.indexing_maps(); 135 // Iterate over the inputs and outputs in order. 136 // Extract the subranges from the linearized ranges. 137 SmallVector<Value, 8> ios(op.getInputsAndOutputBuffers()); 138 for (auto en : llvm::enumerate(ios)) { 139 unsigned idx = en.index(); 140 auto map = maps[idx].cast<AffineMapAttr>().getValue(); 141 LLVM_DEBUG(dbgs() << "getViewDefiningLoopRange I/O idx: " << idx << "\n"); 142 LLVM_DEBUG(dbgs() << "getViewDefiningLoopRange map: " << map << "\n"); 143 Value view = en.value(); 144 SmallVector<Value, 8> viewRanges(map.getNumResults(), nullptr); 145 for (auto en2 : llvm::enumerate(map.getResults())) { 146 if (loopDepth == en2.value().cast<AffineDimExpr>().getPosition()) { 147 LLVM_DEBUG(dbgs() << "getViewDefiningLoopRange loopDepth: " << loopDepth 148 << "\n"); 149 LLVM_DEBUG(dbgs() << "getViewDefiningLoopRange view: " << view << "\n"); 150 return ViewDimension{view, static_cast<unsigned>(en2.index())}; 151 } 152 } 153 } 154 llvm_unreachable("Expect to be able to extract a view defining loop range"); 155 } 156 157 static LinalgOp fuse(Value producedView, LinalgOp producer, LinalgOp consumer, 158 unsigned consumerIdx, unsigned producerIdx, 159 OperationFolder *folder) { 160 assert(producer.hasBufferSemantics() && 161 "expected linalg op with buffer semantics"); 162 assert(consumer.hasBufferSemantics() && 163 "expected linalg op with buffer semantics"); 164 165 auto subView = dyn_cast_or_null<SubViewOp>( 166 consumer.getBuffer(consumerIdx).getDefiningOp()); 167 auto slice = dyn_cast_or_null<SliceOp>( 168 consumer.getBuffer(consumerIdx).getDefiningOp()); 169 assert(subView || slice); 170 (void)subView; 171 (void)slice; 172 173 // loopToOperandRangesMaps are permutations-only by construction: 174 // we can always identify a data dimension with a (at least one) loop 175 // dimension. 176 AffineMap producerMap = 177 producer.indexing_maps()[producer.getNumInputs() + producerIdx] 178 .cast<AffineMapAttr>() 179 .getValue(); 180 LLVM_DEBUG(dbgs() << "Producer Idx: " << producerIdx 181 << ", producer map: " << producerMap << "\n"); 182 183 unsigned nPar = producer.getNumParallelLoops(); 184 unsigned nRed = producer.getNumReductionLoops(); 185 unsigned nWin = producer.getNumWindowLoops(); 186 SmallVector<SubViewOp::Range, 8> loopRanges(nPar + nRed + nWin); 187 188 OpBuilder b(consumer.getOperation()); 189 auto loc = consumer.getLoc(); 190 // Iterate over dimensions identified by the producer map for `producerIdx`. 191 // This defines a subset of the loop ranges that we need to complete later. 192 for (auto en : llvm::enumerate(producerMap.getResults())) { 193 unsigned posInProducerLoop = en.value().cast<AffineDimExpr>().getPosition(); 194 loopRanges[posInProducerLoop] = 195 subView.getOrCreateRanges(b, loc)[en.index()]; 196 } 197 198 // Iterate over all dimensions. For the dimensions not identified by the 199 // producer map for `producerIdx`, we need to explicitly compute the view that 200 // defines the loop ranges using the `producer`. 201 for (unsigned i = 0, nLoops = loopRanges.size(); i < nLoops; ++i) { 202 if (loopRanges[i].offset) 203 LLVM_DEBUG(llvm::dbgs() 204 << "existing LoopRange: " << loopRanges[i] << "\n"); 205 else { 206 auto viewDim = getViewDefiningLoopRange(producer, i); 207 loopRanges[i] = SubViewOp::Range{folded_std_constant_index(folder, 0), 208 std_dim(viewDim.view, viewDim.dimension), 209 folded_std_constant_index(folder, 1)}; 210 LLVM_DEBUG(llvm::dbgs() << "new LoopRange: " << loopRanges[i] << "\n"); 211 } 212 } 213 214 return cloneWithLoopRanges(b, loc, producer, loopRanges); 215 } 216 217 // Encode structural fusion safety preconditions. 218 // Some of these will be lifted in the future with better analysis. 219 static bool isStructurallyFusableProducer(LinalgOp producer, Value consumedView, 220 LinalgOp consumer) { 221 assert(producer.hasBufferSemantics() && 222 "expected linalg op with buffer semantics"); 223 assert(consumer.hasBufferSemantics() && 224 "expected linalg op with buffer semantics"); 225 if (producer.getNumOutputs() != 1) { 226 LLVM_DEBUG(dbgs() << "\nNot structurally fusable (multi-output)"); 227 return false; 228 } 229 // Only fuse when the producer block dominates. 230 DominanceInfo dom(producer.getOperation()); 231 if (!dom.dominates(producer.getOperation()->getBlock(), 232 consumer.getOperation()->getBlock())) { 233 LLVM_DEBUG( 234 dbgs() 235 << "\nNot structurally fusable (producer block does not dominate)"); 236 return false; 237 } 238 return true; 239 } 240 241 bool mlir::linalg::isProducerLastWriteOfView(const LinalgDependenceGraph &graph, 242 LinalgOp consumer, 243 Value consumedView, 244 LinalgOp producer) { 245 assert(producer.hasBufferSemantics() && 246 "expected linalg op with buffer semantics"); 247 assert(consumer.hasBufferSemantics() && 248 "expected linalg op with buffer semantics"); 249 // Make some simple structural checks that alleviate the need for more 250 // complex analyses. 251 if (!isStructurallyFusableProducer(producer, consumedView, consumer)) { 252 LLVM_DEBUG(dbgs() << "\n***Not static last write due to structure:\t" 253 << *producer.getOperation()); 254 return false; 255 } 256 // Check for any interleaved write to consumedView. 257 if (!graph.findCoveringWrites(producer, consumer, consumedView).empty()) { 258 LLVM_DEBUG(dbgs() << "\n***Not fusable due to interleaved write:\t" 259 << *producer.getOperation()); 260 return false; 261 } 262 return true; 263 } 264 265 bool mlir::linalg::isFusableInto(const LinalgDependenceGraph &graph, 266 LinalgOp consumer, Value consumedView, 267 LinalgOp producer) { 268 assert(producer.hasBufferSemantics() && 269 "expected linalg op with buffer semantics"); 270 assert(consumer.hasBufferSemantics() && 271 "expected linalg op with buffer semantics"); 272 if (!isProducerLastWriteOfView(graph, consumer, consumedView, producer)) 273 return false; 274 // Check for any fusion-preventing dependence to any view read/written that 275 // would violate dependences. 276 if (!graph.findCoveringDependences(producer, consumer).empty()) { 277 LLVM_DEBUG(dbgs() << "\n***Not fusable due to an interleaved dependence:\t" 278 << *producer.getOperation()); 279 return false; 280 } 281 if (auto convOp = dyn_cast<linalg::ConvOp>(producer.getOperation())) { 282 // TODO: add a level of indirection to linalg.generic. 283 if (convOp.padding()) 284 return false; 285 } 286 if (auto convOp = dyn_cast<linalg::ConvOp>(consumer.getOperation())) { 287 // TODO: add a level of indirection to linalg.generic. 288 if (convOp.padding()) 289 return false; 290 } 291 return true; 292 } 293 294 static bool isSameSubView(Value a, Value b) { 295 if (a == b) 296 return true; 297 auto sva = a.getDefiningOp<SubViewOp>(); 298 auto svb = b.getDefiningOp<SubViewOp>(); 299 if (!sva || !svb) 300 return false; 301 if (!isSameSubView(sva.getViewSource(), svb.getViewSource())) 302 return false; 303 if (sva.getType() != svb.getType()) 304 return false; 305 if (sva.getRank() != svb.getRank()) 306 return false; 307 if (sva.getNumOperands() != svb.getNumOperands()) 308 return false; 309 if (sva.static_offsets() != svb.static_offsets()) 310 return false; 311 if (sva.static_sizes() != svb.static_sizes()) 312 return false; 313 if (sva.static_strides() != svb.static_strides()) 314 return false; 315 /// Skip the "viewSource" operand. 316 for (unsigned idx = 1, e = sva.getNumOperands(); idx != e; ++idx) 317 if (sva.getOperand(idx) != svb.getOperand(idx)) 318 return false; 319 return true; 320 } 321 322 static Optional<FusionInfo> 323 fuseProducerOfDep(OpBuilder &b, LinalgOp consumer, unsigned consumerIdx, 324 const LinalgDependenceGraph &graph, OperationFolder *folder, 325 LinalgDependenceGraph::DependenceType depType) { 326 assert(consumer.hasBufferSemantics() && 327 "expected linalg op with buffer semantics"); 328 LLVM_DEBUG(dbgs() << "\nStart examining consumer: " 329 << *consumer.getOperation()); 330 for (auto dependence : graph.getDependencesInto(consumer, depType)) { 331 LLVM_DEBUG(dbgs() << "\n***Consider producer:\t" 332 << *dependence.dependentOpView.op << "\n"); 333 auto producer = cast<LinalgOp>(dependence.dependentOpView.op); 334 335 // Check that the dependence is indeed on the input `consumerIdx` view. 336 auto consumedView = dependence.indexingView; 337 if (!isSameSubView(consumer.getBuffer(consumerIdx), consumedView)) 338 continue; 339 340 // Consumer consumes this view, `isStructurallyFusableProducer` also checks 341 // whether it is a strict subview of the producer view. 342 auto producedView = dependence.dependentOpView.view; 343 auto producerIdx = producer.getIndexOfOutputBuffer(producedView).getValue(); 344 // `consumerIdx` and `producerIdx` exist by construction. 345 LLVM_DEBUG(dbgs() << "\n" 346 << LinalgDependenceGraph::getDependenceTypeStr(depType) 347 << "producer: " << *producer.getOperation() << " view: " 348 << producedView << " output index: " << producerIdx); 349 350 // Must be a subview or a slice to guarantee there are loops we can fuse 351 // into. 352 auto subView = consumedView.getDefiningOp<SubViewOp>(); 353 auto slice = consumedView.getDefiningOp<SliceOp>(); 354 if (!subView && !slice) { 355 LLVM_DEBUG(dbgs() << "\nNot fusable (not a subview or slice)"); 356 continue; 357 } 358 359 // Simple fusability checks. 360 if (!isFusableInto(graph, consumer, consumedView, producer)) 361 continue; 362 363 // Fuse `producer` just before `consumer`. 364 OpBuilder::InsertionGuard g(b); 365 b.setInsertionPoint(consumer.getOperation()); 366 ScopedContext scope(b, consumer.getLoc()); 367 LLVM_DEBUG(dbgs() << "Fuse into consumer: " << *consumer << "\n"); 368 auto fusedProducer = fuse(producedView, producer, consumer, consumerIdx, 369 producerIdx, folder); 370 371 return FusionInfo{producer, fusedProducer}; 372 } 373 return llvm::None; 374 } 375 376 // Only consider RAW and WAW atm. 377 Optional<FusionInfo> mlir::linalg::fuseProducerOf( 378 OpBuilder &b, LinalgOp consumer, unsigned consumerIdx, 379 const LinalgDependenceGraph &graph, OperationFolder *folder) { 380 SmallVector<LinalgDependenceGraph::DependenceType, 4> deps = { 381 LinalgDependenceGraph::DependenceType::RAW, 382 LinalgDependenceGraph::DependenceType::WAW, 383 }; 384 for (auto dep : deps) { 385 if (auto res = 386 fuseProducerOfDep(b, consumer, consumerIdx, graph, folder, dep)) 387 return res; 388 } 389 return llvm::None; 390 } 391 392 static void fuseLinalgOpsGreedily(FuncOp f) { 393 LLVM_DEBUG(f.print(dbgs() << "\nBefore linalg-fusion: \n")); 394 395 OpBuilder b(f); 396 OperationFolder folder(f.getContext()); 397 DenseSet<Operation *> eraseSet; 398 399 // Save original Linalg ops, we only want to make a pass over those. 400 SmallVector<Operation *, 8> linalgOps; 401 f.walk([&](LinalgOp op) { 402 if (op.hasBufferSemantics()) 403 linalgOps.push_back(op); 404 }); 405 406 // TODO: LinalgDependenceGraph should be able to update itself. 407 // The current naive and expensive reconstruction of the graph should be 408 // removed. 409 for (auto *op : llvm::reverse(linalgOps)) { 410 for (unsigned id = 0, e = LinalgOp(op).getNumInputsAndOutputBuffers(); 411 id < e; ++id) { 412 linalg::Aliases aliases; 413 linalg::LinalgDependenceGraph graph(aliases, linalgOps); 414 if (auto info = fuseProducerOf(b, op, id, graph, &folder)) { 415 auto *originalOp = info->originalProducer.getOperation(); 416 eraseSet.insert(originalOp); 417 auto *originalOpInLinalgOpsVector = 418 std::find(linalgOps.begin(), linalgOps.end(), originalOp); 419 *originalOpInLinalgOpsVector = info->fusedProducer.getOperation(); 420 } 421 } 422 } 423 // The `fuseProducerOf` function performs structural checks and in particular 424 // that no covering read or write exist between the consumer and the producer. 425 // As a consequence, the only fusions that may occur preserve subsequent 426 // dependences and are guaranteed by construction to produce the whole view. 427 // We may thus erase the producer once it is fused. 428 for (auto *e : eraseSet) 429 e->erase(); 430 LLVM_DEBUG(f.print(dbgs() << "\nAfter linalg-fusion: \n")); 431 } 432 433 //====---------------------------------------------------------------------===// 434 // Fusion on Tensor operation. 435 //====---------------------------------------------------------------------===// 436 437 namespace { 438 439 /// Implementation of fusion of generic ops and indexed_generic ops. 440 struct FuseGenericOpsOnTensors { 441 static bool isFusible(LinalgOp producer, LinalgOp consumer, 442 unsigned consumerIdx) { 443 // Producer and consumer must have tensor semantics. 444 if (!producer.hasTensorSemantics() || !consumer.hasTensorSemantics()) 445 return false; 446 447 // Verify that 448 // - the producer has all "parallel" iterator type. 449 if (producer.getNumParallelLoops() != producer.getNumLoops()) 450 return false; 451 452 // Get the consumer index map. The number of results of the consumer index 453 // map must match the number of loops of the producer. 454 AffineMap consumerIndexMap = consumer.getIndexingMap(consumerIdx); 455 if (consumerIndexMap.getNumResults() != producer.getNumLoops()) 456 return false; 457 458 // Finally the index_map for the result must be invertible. For now just 459 // verify it is a permutation. 460 AffineMap producerResultIndexMap = producer.getOutputIndexingMap(0); 461 return producerResultIndexMap.isPermutation(); 462 } 463 464 static LinalgOp fuse(LinalgOp producer, LinalgOp consumer, 465 unsigned consumerIdx, PatternRewriter &rewriter, 466 OperationFolder *folder = nullptr) { 467 if (!isFusible(producer, consumer, consumerIdx)) 468 return nullptr; 469 470 unsigned numFusedOperands = producer.getOperation()->getNumOperands() + 471 consumer.getOperation()->getNumOperands() - 1; 472 473 // Compute the fused operands list, 474 SmallVector<Value, 2> fusedOperands; 475 fusedOperands.reserve(numFusedOperands); 476 auto consumerOperands = consumer.getOperation()->getOperands(); 477 auto producerOperands = producer.getOperation()->getOperands(); 478 fusedOperands.assign(consumerOperands.begin(), 479 std::next(consumerOperands.begin(), consumerIdx)); 480 fusedOperands.append(producerOperands.begin(), producerOperands.end()); 481 fusedOperands.append(std::next(consumerOperands.begin(), consumerIdx + 1), 482 consumerOperands.end()); 483 484 // Compute indexing_maps for the fused operation. The indexing_maps for the 485 // operands of the consumers that arent fused are the same. The 486 // indexing_maps for the producers need to be computed based on the 487 // indexing_map of the operand at consumerIdx in the consumer. 488 SmallVector<Attribute, 4> fusedIndexMaps; 489 auto consumerIndexMaps = consumer.indexing_maps(); 490 fusedIndexMaps.reserve(fusedOperands.size() + 491 consumer.getOperation()->getNumResults()); 492 fusedIndexMaps.assign(consumerIndexMaps.begin(), 493 std::next(consumerIndexMaps.begin(), consumerIdx)); 494 // Compute indexing maps for the producer args in the fused operation. 495 computeProducerOperandIndex( 496 producer, consumer.getInputIndexingMap(consumerIdx), fusedIndexMaps); 497 498 // Append the indexing maps for the remaining consumer operands. 499 fusedIndexMaps.append(std::next(consumerIndexMaps.begin(), consumerIdx + 1), 500 consumerIndexMaps.end()); 501 502 // Generate the fused op. 503 LinalgOp fusedOp; 504 if (isa<GenericOp>(producer.getOperation()) && 505 isa<GenericOp>(consumer.getOperation())) { 506 fusedOp = 507 rewriter 508 .create<GenericOp>( 509 rewriter.getUnknownLoc(), 510 consumer.getOperation()->getResultTypes(), fusedOperands, 511 rewriter.getI64IntegerAttr(fusedOperands.size()), 512 rewriter.getI64IntegerAttr( 513 consumer.getOperation()->getNumResults()), 514 rewriter.getArrayAttr(fusedIndexMaps), 515 consumer.iterator_types(), 516 /*doc=*/nullptr, 517 /*library_call=*/nullptr, 518 /*symbol_source=*/nullptr) 519 .getOperation(); 520 } else { 521 fusedOp = 522 rewriter 523 .create<IndexedGenericOp>( 524 rewriter.getUnknownLoc(), 525 consumer.getOperation()->getResultTypes(), fusedOperands, 526 rewriter.getI64IntegerAttr(fusedOperands.size()), 527 rewriter.getI64IntegerAttr( 528 consumer.getOperation()->getNumResults()), 529 rewriter.getArrayAttr(fusedIndexMaps), 530 consumer.iterator_types(), 531 /*doc=*/nullptr, 532 /*library_call=*/nullptr, 533 /*symbol_source=*/nullptr) 534 .getOperation(); 535 } 536 537 // Construct an AffineMap from consumer loops to producer loops. 538 // consumer loop -> tensor index 539 AffineMap consumerResultIndexMap = 540 consumer.getInputIndexingMap(consumerIdx); 541 // producer loop -> tensor index 542 AffineMap producerResultIndexMap = producer.getOutputIndexingMap(0); 543 // tensor index -> producer loop 544 AffineMap invProducerResultIndexMap = 545 inversePermutation(producerResultIndexMap); 546 assert(invProducerResultIndexMap && 547 "expected producer result indexig map to be invertible"); 548 // consumer loop -> producer loop 549 AffineMap consumerToProducerLoopsMap = 550 invProducerResultIndexMap.compose(consumerResultIndexMap); 551 552 generateFusedRegion(rewriter, fusedOp, producer, consumer, 553 consumerToProducerLoopsMap, consumerIdx, 554 consumer.getNumLoops()); 555 return fusedOp; 556 } 557 558 private: 559 /// Append to `fusedOpIndexingMapAttrs` the indexing maps for the operands of 560 /// the `producer` to use in the fused operation given the indexing map of the 561 /// result of the producer in the consumer. 562 static void computeProducerOperandIndex( 563 LinalgOp producer, AffineMap fusedConsumerArgIndexMap, 564 SmallVectorImpl<Attribute> &fusedOpIndexingMapAttrs) { 565 // The indexing map in the consumer op (fusedConsumerArgIndexMap) is a map 566 // from consumer loop -> consumer arg tensor index/producer result tensor 567 // index. The fused loop is same as the consumer loop. For each producer arg 568 // the indexing map to be computed is a map from consumer loop -> producer 569 // arg tensor index. 570 571 AffineMap producerResultIndexMap = producer.getOutputIndexingMap(0); 572 // producerResultIndexMap is a map from producer loop -> tensor index. 573 // Compute the inverse to get map from tensor index -> producer loop. 574 // The inverse is a map from producer result tensor index -> producer loop. 575 AffineMap invProducerResultIndexMap = 576 inversePermutation(producerResultIndexMap); 577 assert(invProducerResultIndexMap && 578 "expected producer result indexig map to be invertible"); 579 for (unsigned argNum : llvm::seq<unsigned>(0, producer.getNumInputs())) { 580 // argMap is a map from producer loop -> producer arg tensor index. 581 AffineMap argMap = producer.getInputIndexingMap(argNum); 582 583 // Compose argMap with invProducerResultIndexMap to get a map from 584 // producer result tensor index -> producer arg tensor index. 585 AffineMap t1 = argMap.compose(invProducerResultIndexMap); 586 587 // Compose t1 with fusedConsumerArgIndexMap gives an indexing map from 588 // consumer loop/ fused loop -> producer arg tensor index. 589 AffineMap indexingMap = t1.compose(fusedConsumerArgIndexMap); 590 fusedOpIndexingMapAttrs.push_back(AffineMapAttr::get(indexingMap)); 591 } 592 } 593 594 /// Generate the region of the fused operation. The region of the fused op 595 /// must be empty. 596 static void generateFusedRegion(PatternRewriter &rewriter, Operation *fusedOp, 597 LinalgOp producer, LinalgOp consumer, 598 AffineMap consumerToProducerLoopsMap, 599 unsigned consumerIdx, unsigned nloops) { 600 // Build the region of the fused op. 601 Block &producerBlock = producer.getOperation()->getRegion(0).front(); 602 Block &consumerBlock = consumer.getOperation()->getRegion(0).front(); 603 Block *fusedBlock = new Block(); 604 fusedOp->getRegion(0).push_back(fusedBlock); 605 BlockAndValueMapping mapper; 606 OpBuilder::InsertionGuard guard(rewriter); 607 rewriter.setInsertionPointToStart(fusedBlock); 608 609 // The block arguments are 610 // [index_0, index_1, ... , 611 // consumer_operand_0, ... , consumer_operand_(`consumerIdx`-1), 612 // producer_operand_0, ... , producer_operand_(n-1)], 613 // consumer_operand_(`consumerIdx`), .. consumer_operand_(m-1)] 614 // , where n is the number of producer's operand and m is the number 615 // consumer's operand. 616 // If both `numProducerIndices` and `numConsumerIndices` are zero, this is a 617 // generic op. In this case, there are no indices in block arguments. 618 unsigned numProducerIndices = 619 isa<IndexedGenericOp>(producer.getOperation()) ? nloops : 0; 620 unsigned numConsumerIndices = 621 isa<IndexedGenericOp>(consumer.getOperation()) ? nloops : 0; 622 // Firstly, add all the indices to the block arguments. 623 for (unsigned i = 0, e = std::max(numProducerIndices, numConsumerIndices); 624 i < e; ++i) 625 fusedBlock->addArgument(rewriter.getIndexType()); 626 // Map the arguments for the unmodified args from the consumer. 627 for (auto consumerArg : llvm::enumerate(consumerBlock.getArguments())) { 628 if (consumerArg.index() == consumerIdx + numConsumerIndices) { 629 // Map the arguments for the args from the producer. 630 for (auto producerArg : llvm::enumerate(producerBlock.getArguments())) { 631 // If producer is an indexed_generic op, map the indices from consumer 632 // loop to producer loop (because the fusedOp is built based on 633 // consumer's perspective). 634 if (producerArg.index() < numProducerIndices) { 635 auto newIndex = rewriter.create<mlir::AffineApplyOp>( 636 producer.getLoc(), 637 consumerToProducerLoopsMap.getSubMap(producerArg.index()), 638 fusedBlock->getArguments().take_front(nloops)); 639 mapper.map(producerArg.value(), newIndex); 640 } else { 641 mapper.map(producerArg.value(), 642 fusedBlock->addArgument(producerArg.value().getType())); 643 } 644 } 645 continue; 646 } 647 648 // If consumer is an indexed_generic op, map the indices to the block 649 // arguments directly. Otherwise, add the same type of arugment and map to 650 // it. 651 if (consumerArg.index() < numConsumerIndices) { 652 mapper.map(consumerArg.value(), 653 fusedBlock->getArgument(consumerArg.index())); 654 } else { 655 mapper.map(consumerArg.value(), 656 fusedBlock->addArgument(consumerArg.value().getType())); 657 } 658 } 659 660 // Add operations from producer (except the yield operation) to the fused 661 // op. 662 for (auto &op : producerBlock.getOperations()) { 663 if (auto yieldOp = dyn_cast<YieldOp>(op)) { 664 // Lookup the value the yield operation is mapped to. 665 Value yieldVal = yieldOp.getOperand(0); 666 if (Value clonedVal = mapper.lookupOrNull(yieldVal)) 667 mapper.map( 668 consumerBlock.getArgument(consumerIdx + numConsumerIndices), 669 clonedVal); 670 continue; 671 } 672 rewriter.clone(op, mapper); 673 } 674 for (auto &op : consumerBlock.getOperations()) 675 rewriter.clone(op, mapper); 676 } 677 }; 678 } // namespace 679 680 /// Linearize the expressions in `sourceMap` based on the `reassociationMaps` 681 /// provided, given the shape of the source tensor that corresponds to the 682 /// `sourceMap`. Note that this implicitly assumes that the tensors dimensions 683 /// are "row-major" ordered logically. 684 /// 685 /// For example: 686 /// 687 /// %0 = op ... : tensor<?x?x4x5xf32> 688 /// with output index_map `affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>` 689 /// 690 /// and reshape: 691 /// %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>, 692 /// affine_map<(i, j, k, l) -> (j, k, l)>] : 693 /// tensor<?x?x4x5xf32> into tensor<?x?xf32> 694 /// 695 /// would be rewritten into: 696 /// %0 = op ... : tensor<?x?x4x5xf32> 697 /// with output index_map 698 /// `affine_map<(d0, d1, d2, d3) -> (d0, d1 * 20 + d2 * 5 + d3)>` 699 static AffineMap linearizeCollapsedDims(AffineMap sourceMap, 700 ArrayRef<int64_t> sourceShape, 701 ArrayRef<AffineMap> reassociationMaps) { 702 SmallVector<AffineExpr, 4> resultExprs; 703 resultExprs.reserve(reassociationMaps.size()); 704 ArrayRef<AffineExpr> sourceExprs = sourceMap.getResults(); 705 MLIRContext *context = sourceMap.getContext(); 706 707 // Compute the result exprs based on the reassociation maps. 708 for (AffineMap map : reassociationMaps) { 709 ArrayRef<AffineExpr> collapsedDims = map.getResults(); 710 // Assume that they are in-order and contiguous (already checked in 711 // verifier). 712 assert(!collapsedDims.empty()); 713 unsigned startDim = 714 collapsedDims.front().cast<AffineDimExpr>().getPosition(); 715 AffineExpr linearizedExpr = makeCanonicalStridedLayoutExpr( 716 sourceShape.slice(startDim, collapsedDims.size()), 717 sourceExprs.slice(startDim, collapsedDims.size()), context); 718 resultExprs.push_back(linearizedExpr); 719 } 720 return AffineMap::get(sourceMap.getNumDims(), sourceMap.getNumSymbols(), 721 resultExprs, context); 722 } 723 724 /// Checks if the `reshapeOp` can be fused with it consumer (if `asProducer` is 725 /// true) or its producer (if `asProducer` is false) given the indexing map at 726 /// its use. 727 static bool isTensorReshapeOpFusible(TensorReshapeOp reshapeOp, 728 AffineMap useIndexMap, bool asProducer) { 729 RankedTensorType returnType = reshapeOp.getResultType(); 730 RankedTensorType operandType = reshapeOp.getSrcType(); 731 // Reshape is fusible with its consumer (i.e. reshape as a producer) when its 732 // operand is of lesser rank than the result. Fusing when operand has higher 733 // rank will require use of mods and divs in the indexing maps of the fused op 734 // which would make it non-invertible. Similarly reshape is fused with its 735 // producer (i.e. reshape as consumer) only if the return type has lesser 736 // rank. 737 if ((asProducer && returnType.getRank() < operandType.getRank()) || 738 (!asProducer && operandType.getRank() < returnType.getRank())) 739 return false; 740 return useIndexMap.isIdentity(); 741 } 742 743 /// Based on the type of `op` create a linalg op of the same type, i.e. if `op` 744 /// is a linalg.generic operation, the create a `linalg.generic` operation with 745 /// the given `args`. Expects `op` to be `linalg.generic` or 746 /// `linalg.indexed_generic`. 747 template <typename... Args> 748 static LinalgOp createLinalgOpOfSameType(LinalgOp op, PatternRewriter &rewriter, 749 Args... args) { 750 if (isa<GenericOp>(op.getOperation())) 751 return cast<LinalgOp>(rewriter.create<GenericOp>(args...).getOperation()); 752 if (isa<IndexedGenericOp>(op.getOperation())) 753 return cast<LinalgOp>( 754 rewriter.create<IndexedGenericOp>(args...).getOperation()); 755 llvm_unreachable( 756 "expected only linalg.generic or linalg.indexed_generic ops"); 757 return nullptr; 758 } 759 760 namespace { 761 762 /// Implementation of fusion on tensor ops when producer is a TensorReshapeOp. 763 struct FuseTensorReshapeOpAsProducer { 764 static bool isFusible(TensorReshapeOp producer, LinalgOp consumer, 765 unsigned consumerIdx) { 766 return isa<GenericOp, IndexedGenericOp>(consumer.getOperation()) && 767 consumer.hasTensorSemantics() && 768 isTensorReshapeOpFusible(producer, 769 consumer.getInputIndexingMap(consumerIdx), 770 /*asProducer=*/true); 771 } 772 773 static LinalgOp fuse(TensorReshapeOp producer, LinalgOp consumer, 774 unsigned consumerIdx, PatternRewriter &rewriter, 775 OperationFolder *folder = nullptr) { 776 if (!isFusible(producer, consumer, consumerIdx)) 777 return nullptr; 778 779 // Compute the fused operands list, 780 Operation *consumerOp = consumer.getOperation(); 781 SmallVector<Value, 2> fusedOperands(consumerOp->getOperands()); 782 fusedOperands[consumerIdx] = producer.src(); 783 784 // Compute indexing_maps for the fused operation. The indexing_maps for the 785 // operands of the consumers that arent fused are the same. 786 SmallVector<AffineMap, 4> fusedIndexMaps = 787 llvm::to_vector<4>(llvm::map_range( 788 consumer.indexing_maps(), [](Attribute attr) -> AffineMap { 789 return attr.cast<AffineMapAttr>().getValue(); 790 })); 791 792 // Compute the indexing map to use for the operand of the producer. 793 AffineMap modifiedMap = linearizeCollapsedDims( 794 fusedIndexMaps[consumerIdx], producer.getResultType().getShape(), 795 producer.getReassociationMaps()); 796 for (AffineExpr expr : modifiedMap.getResults()) { 797 if (!expr.isPureAffine()) 798 return nullptr; 799 } 800 fusedIndexMaps[consumerIdx] = modifiedMap; 801 802 // Further check that the resulting index maps can be fused and 803 // inverted. Without this the resultant op is not legal. 804 if (!inversePermutation(concatAffineMaps(fusedIndexMaps))) 805 return nullptr; 806 807 SmallVector<Attribute, 4> indexMapAttrs = llvm::to_vector<4>( 808 llvm::map_range(fusedIndexMaps, [](AffineMap map) -> Attribute { 809 return AffineMapAttr::get(map); 810 })); 811 LinalgOp fusedOp = createLinalgOpOfSameType( 812 consumer, rewriter, rewriter.getUnknownLoc(), 813 consumerOp->getResultTypes(), fusedOperands, 814 rewriter.getI64IntegerAttr(fusedOperands.size()), 815 rewriter.getI64IntegerAttr(consumerOp->getNumResults()), 816 rewriter.getArrayAttr(indexMapAttrs), consumer.iterator_types(), 817 /*doc=*/nullptr, 818 /*library_call=*/nullptr, 819 /*symbol_source=*/nullptr); 820 auto &fusedRegion = fusedOp.getOperation()->getRegion(0); 821 rewriter.cloneRegionBefore(consumerOp->getRegion(0), fusedRegion, 822 fusedRegion.begin()); 823 return fusedOp; 824 } 825 }; 826 827 /// Implementation of fusion on tensor ops when consumer is a TensorReshapeOp. 828 struct FuseTensorReshapeOpAsConsumer { 829 static bool isFusible(LinalgOp producer, TensorReshapeOp consumer, 830 unsigned consumerIdx) { 831 return isa<GenericOp, IndexedGenericOp>(producer.getOperation()) && 832 producer.hasTensorSemantics() && 833 isTensorReshapeOpFusible(consumer, producer.getOutputIndexingMap(0), 834 /*asProducer=*/false); 835 } 836 837 static LinalgOp fuse(LinalgOp producer, TensorReshapeOp consumer, 838 unsigned consumerIdx, PatternRewriter &rewriter, 839 OperationFolder *folder = nullptr) { 840 if (!isFusible(producer, consumer, consumerIdx)) 841 return nullptr; 842 843 // The indexing_maps for the operands of the fused operation are same as 844 // those for the operands of the producer. 845 SmallVector<AffineMap, 4> fusedIndexMaps = 846 llvm::to_vector<4>(llvm::map_range( 847 producer.indexing_maps(), [](Attribute attr) -> AffineMap { 848 return attr.cast<AffineMapAttr>().getValue(); 849 })); 850 // Compute the indexing map to use for the operand of the producer. 851 AffineMap modifiedMap = linearizeCollapsedDims( 852 producer.getOutputIndexingMap(0), consumer.getSrcType().getShape(), 853 consumer.getReassociationMaps()); 854 for (AffineExpr expr : modifiedMap.getResults()) { 855 if (!expr.isPureAffine()) 856 return nullptr; 857 } 858 fusedIndexMaps.back() = modifiedMap; 859 860 // Further check that the resulting index maps can be fused and 861 // inverted. Without this the resultant op is not legal. 862 if (!inversePermutation(concatAffineMaps(fusedIndexMaps))) 863 return nullptr; 864 865 SmallVector<Attribute, 4> indexMapAttrs = llvm::to_vector<4>( 866 llvm::map_range(fusedIndexMaps, [](AffineMap map) -> Attribute { 867 return AffineMapAttr::get(map); 868 })); 869 870 Operation *producerOp = producer.getOperation(); 871 LinalgOp fusedOp = createLinalgOpOfSameType( 872 producer, rewriter, rewriter.getUnknownLoc(), consumer.getResultType(), 873 producerOp->getOperands(), 874 rewriter.getI64IntegerAttr(producerOp->getNumOperands()), 875 rewriter.getI64IntegerAttr(1), rewriter.getArrayAttr(indexMapAttrs), 876 producer.iterator_types(), 877 /*doc=*/nullptr, 878 /*library_call=*/nullptr, 879 /*symbol_source=*/nullptr); 880 auto &fusedRegion = fusedOp.getOperation()->getRegion(0); 881 rewriter.cloneRegionBefore(producerOp->getRegion(0), fusedRegion, 882 fusedRegion.begin()); 883 return fusedOp; 884 } 885 }; 886 887 /// Implementation of fusion on tensor ops when producer is a splat constant. 888 struct FuseConstantOpAsProducer { 889 static bool isFusible(ConstantOp producer, LinalgOp consumer, 890 unsigned consumerIdx) { 891 return isa<GenericOp, IndexedGenericOp>(consumer.getOperation()) && 892 consumer.hasTensorSemantics() && 893 producer.getResult().getType().isa<RankedTensorType>() && 894 producer.value().cast<DenseElementsAttr>().isSplat(); 895 } 896 897 static LinalgOp fuse(ConstantOp producer, LinalgOp consumer, 898 unsigned consumerIdx, PatternRewriter &rewriter, 899 OperationFolder *folder = nullptr) { 900 if (!isFusible(producer, consumer, consumerIdx)) 901 return nullptr; 902 903 // The indexing_maps for the operands of the fused operation are same as 904 // those for the operands of the consumer without the indexing map at 905 // consumerIdx 906 SmallVector<AffineMap, 4> fusedIndexMaps = 907 llvm::to_vector<4>(llvm::map_range( 908 consumer.indexing_maps(), [](Attribute attr) -> AffineMap { 909 return attr.cast<AffineMapAttr>().getValue(); 910 })); 911 fusedIndexMaps.erase(std::next(fusedIndexMaps.begin(), consumerIdx)); 912 913 // The operands list is same as the consumer with the argument for constant 914 // index dropped. 915 Operation *consumerOp = consumer.getOperation(); 916 SmallVector<Value, 4> fusedOperands(consumerOp->getOperands()); 917 fusedOperands.erase(std::next(fusedOperands.begin(), consumerIdx)); 918 919 // Create a constant scalar value from the splat constant. 920 Value scalarConstant = rewriter.create<ConstantOp>( 921 producer.getLoc(), 922 producer.value().cast<DenseElementsAttr>().getSplatValue()); 923 924 LinalgOp fusedOp = createLinalgOpOfSameType( 925 consumer, rewriter, rewriter.getUnknownLoc(), 926 consumerOp->getResultTypes(), fusedOperands, 927 rewriter.getI64IntegerAttr(consumerOp->getNumOperands() - 1), 928 rewriter.getI64IntegerAttr(consumerOp->getNumResults()), 929 rewriter.getAffineMapArrayAttr(fusedIndexMaps), 930 consumer.iterator_types(), 931 /*doc=*/nullptr, 932 /*library_call=*/nullptr, 933 /*symbol_source=*/nullptr); 934 935 // Map the block argument corresponding to the replaced argument with the 936 // scalar constant. 937 Region &consumerRegion = consumerOp->getRegion(0); 938 Block &entryBlock = *consumerRegion.begin(); 939 unsigned argIndex = entryBlock.getNumArguments() - 940 consumerOp->getNumOperands() + consumerIdx; 941 BlockAndValueMapping mapping; 942 mapping.map(entryBlock.getArgument(argIndex), scalarConstant); 943 Region &fusedRegion = fusedOp.getOperation()->getRegion(0); 944 rewriter.cloneRegionBefore(consumerRegion, fusedRegion, fusedRegion.begin(), 945 mapping); 946 return fusedOp; 947 } 948 }; 949 } // namespace 950 951 Operation *mlir::linalg::fuseTensorOps(PatternRewriter &rewriter, 952 Operation *consumer, 953 unsigned consumerIdx, 954 OperationFolder *folder) { 955 if (consumerIdx >= consumer->getNumOperands()) 956 return nullptr; 957 Operation *producer = consumer->getOperand(consumerIdx).getDefiningOp(); 958 if (!producer || producer->getNumResults() != 1) 959 return nullptr; 960 961 // Fuse when consumer is GenericOp or IndexedGenericOp. 962 if (isa<GenericOp, IndexedGenericOp>(consumer)) { 963 if (isa<GenericOp, IndexedGenericOp>(producer)) 964 return FuseGenericOpsOnTensors::fuse(cast<LinalgOp>(producer), 965 cast<LinalgOp>(consumer), 966 consumerIdx, rewriter, folder); 967 if (auto reshapeOpProducer = dyn_cast<TensorReshapeOp>(producer)) 968 return FuseTensorReshapeOpAsProducer::fuse(reshapeOpProducer, 969 cast<LinalgOp>(consumer), 970 consumerIdx, rewriter, folder); 971 if (auto constantOpProducer = dyn_cast<ConstantOp>(producer)) 972 return FuseConstantOpAsProducer::fuse(constantOpProducer, 973 cast<LinalgOp>(consumer), 974 consumerIdx, rewriter, folder); 975 return nullptr; 976 } 977 978 if (isa<GenericOp, IndexedGenericOp>(producer)) { 979 // Fuse when consumer is a TensorReshapeOp. 980 if (TensorReshapeOp reshapeOp = dyn_cast<TensorReshapeOp>(consumer)) { 981 return FuseTensorReshapeOpAsConsumer::fuse( 982 cast<LinalgOp>(producer), reshapeOp, consumerIdx, rewriter, folder); 983 } 984 } 985 986 return nullptr; 987 } 988 989 namespace { 990 /// Patterns to fuse a generic op, with the producer of its operands. 991 template <typename LinalgOpTy> 992 struct FuseTensorOps : public OpRewritePattern<LinalgOpTy> { 993 using OpRewritePattern<LinalgOpTy>::OpRewritePattern; 994 995 LogicalResult matchAndRewrite(LinalgOpTy op, 996 PatternRewriter &rewriter) const override { 997 // Find the first operand that is defined by another generic op on tensors. 998 for (auto operandNum : 999 llvm::seq<unsigned>(0, op.getOperation()->getNumOperands())) { 1000 Operation *producer = 1001 op.getOperation()->getOperand(operandNum).getDefiningOp(); 1002 if (Operation *fusedOp = fuseTensorOps(rewriter, op, operandNum)) { 1003 rewriter.replaceOp(op, fusedOp->getResults()); 1004 if (producer && llvm::all_of(producer->getResults(), 1005 [](Value val) { return val.use_empty(); })) 1006 rewriter.eraseOp(producer); 1007 return success(); 1008 } 1009 } 1010 return failure(); 1011 } 1012 }; 1013 1014 /// Pass that fuses generic ops on tensors. Used only for testing. 1015 struct FusionOfTensorOpsPass 1016 : public LinalgFusionOfTensorOpsBase<FusionOfTensorOpsPass> { 1017 void runOnOperation() override { 1018 OwningRewritePatternList patterns; 1019 Operation *op = getOperation(); 1020 populateLinalgTensorOpsFusionPatterns(op->getContext(), patterns); 1021 applyPatternsAndFoldGreedily(op->getRegions(), patterns); 1022 }; 1023 }; 1024 1025 struct LinalgFusionPass : public LinalgFusionBase<LinalgFusionPass> { 1026 void runOnFunction() override { fuseLinalgOpsGreedily(getFunction()); } 1027 }; 1028 } // namespace 1029 1030 void mlir::populateLinalgTensorOpsFusionPatterns( 1031 MLIRContext *context, OwningRewritePatternList &patterns) { 1032 patterns.insert<FuseTensorOps<GenericOp>, FuseTensorOps<IndexedGenericOp>, 1033 FuseTensorOps<TensorReshapeOp>>(context); 1034 } 1035 1036 std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgFusionPass() { 1037 return std::make_unique<LinalgFusionPass>(); 1038 } 1039 1040 std::unique_ptr<Pass> mlir::createLinalgFusionOfTensorOpsPass() { 1041 return std::make_unique<FusionOfTensorOpsPass>(); 1042 } 1043