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 // Verify that 444 // - the producer has all "parallel" iterator type. 445 if (producer.getNumParallelLoops() != producer.getNumLoops()) 446 return false; 447 448 // Get the consumer index map. The number of results of the consumer index 449 // map must match the number of loops of the producer. 450 AffineMap consumerIndexMap = consumer.getIndexingMap(consumerIdx); 451 if (consumerIndexMap.getNumResults() != producer.getNumLoops()) 452 return false; 453 454 // Finally the index_map for the result must be invertible. For now just 455 // verify it is a permutation. 456 AffineMap producerResultIndexMap = producer.getOutputIndexingMap(0); 457 return producerResultIndexMap.isPermutation(); 458 } 459 460 static Operation *fuse(LinalgOp producer, LinalgOp consumer, 461 unsigned consumerIdx, PatternRewriter &rewriter, 462 OperationFolder *folder = nullptr) { 463 if (!isFusible(producer, consumer, consumerIdx)) 464 return nullptr; 465 466 unsigned numFusedOperands = producer.getOperation()->getNumOperands() + 467 consumer.getOperation()->getNumOperands() - 1; 468 469 // Compute the fused operands list, 470 SmallVector<Value, 2> fusedOperands; 471 fusedOperands.reserve(numFusedOperands); 472 auto consumerOperands = consumer.getOperation()->getOperands(); 473 auto producerOperands = producer.getOperation()->getOperands(); 474 fusedOperands.assign(consumerOperands.begin(), 475 std::next(consumerOperands.begin(), consumerIdx)); 476 fusedOperands.append(producerOperands.begin(), producerOperands.end()); 477 fusedOperands.append(std::next(consumerOperands.begin(), consumerIdx + 1), 478 consumerOperands.end()); 479 480 // Compute indexing_maps for the fused operation. The indexing_maps for the 481 // operands of the consumers that arent fused are the same. The 482 // indexing_maps for the producers need to be computed based on the 483 // indexing_map of the operand at consumerIdx in the consumer. 484 SmallVector<Attribute, 4> fusedIndexMaps; 485 auto consumerIndexMaps = consumer.indexing_maps(); 486 fusedIndexMaps.reserve(fusedOperands.size() + 487 consumer.getOperation()->getNumResults()); 488 fusedIndexMaps.assign(consumerIndexMaps.begin(), 489 std::next(consumerIndexMaps.begin(), consumerIdx)); 490 // Compute indexing maps for the producer args in the fused operation. 491 computeProducerOperandIndex( 492 producer, consumer.getInputIndexingMap(consumerIdx), fusedIndexMaps); 493 494 // Append the indexing maps for the remaining consumer operands. 495 fusedIndexMaps.append(std::next(consumerIndexMaps.begin(), consumerIdx + 1), 496 consumerIndexMaps.end()); 497 498 // Generate the fused op. 499 LinalgOp fusedOp; 500 if (isa<GenericOp>(producer.getOperation()) && 501 isa<GenericOp>(consumer.getOperation())) { 502 fusedOp = 503 rewriter 504 .create<GenericOp>( 505 rewriter.getUnknownLoc(), 506 consumer.getOperation()->getResultTypes(), fusedOperands, 507 rewriter.getI64IntegerAttr(fusedOperands.size()), 508 rewriter.getI64IntegerAttr( 509 consumer.getOperation()->getNumResults()), 510 rewriter.getArrayAttr(fusedIndexMaps), 511 consumer.iterator_types(), 512 /*doc=*/nullptr, 513 /*library_call=*/nullptr, 514 /*symbol_source=*/nullptr) 515 .getOperation(); 516 } else { 517 fusedOp = 518 rewriter 519 .create<IndexedGenericOp>( 520 rewriter.getUnknownLoc(), 521 consumer.getOperation()->getResultTypes(), fusedOperands, 522 rewriter.getI64IntegerAttr(fusedOperands.size()), 523 rewriter.getI64IntegerAttr( 524 consumer.getOperation()->getNumResults()), 525 rewriter.getArrayAttr(fusedIndexMaps), 526 consumer.iterator_types(), 527 /*doc=*/nullptr, 528 /*library_call=*/nullptr, 529 /*symbol_source=*/nullptr) 530 .getOperation(); 531 } 532 533 // Construct an AffineMap from consumer loops to producer loops. 534 // consumer loop -> tensor index 535 AffineMap consumerResultIndexMap = 536 consumer.getInputIndexingMap(consumerIdx); 537 // producer loop -> tensor index 538 AffineMap producerResultIndexMap = producer.getOutputIndexingMap(0); 539 // tensor index -> producer loop 540 AffineMap invProducerResultIndexMap = 541 inversePermutation(producerResultIndexMap); 542 assert(invProducerResultIndexMap && 543 "expected producer result indexig map to be invertible"); 544 // consumer loop -> producer loop 545 AffineMap consumerToProducerLoopsMap = 546 invProducerResultIndexMap.compose(consumerResultIndexMap); 547 548 generateFusedRegion(rewriter, fusedOp, producer, consumer, 549 consumerToProducerLoopsMap, consumerIdx, 550 consumer.getNumLoops()); 551 return fusedOp; 552 } 553 554 private: 555 /// Append to `fusedOpIndexingMapAttrs` the indexing maps for the operands of 556 /// the `producer` to use in the fused operation given the indexing map of the 557 /// result of the producer in the consumer. 558 static void computeProducerOperandIndex( 559 LinalgOp producer, AffineMap fusedConsumerArgIndexMap, 560 SmallVectorImpl<Attribute> &fusedOpIndexingMapAttrs) { 561 // The indexing map in the consumer op (fusedConsumerArgIndexMap) is a map 562 // from consumer loop -> consumer arg tensor index/producer result tensor 563 // index. The fused loop is same as the consumer loop. For each producer arg 564 // the indexing map to be computed is a map from consumer loop -> producer 565 // arg tensor index. 566 567 AffineMap producerResultIndexMap = producer.getOutputIndexingMap(0); 568 // producerResultIndexMap is a map from producer loop -> tensor index. 569 // Compute the inverse to get map from tensor index -> producer loop. 570 // The inverse is a map from producer result tensor index -> producer loop. 571 AffineMap invProducerResultIndexMap = 572 inversePermutation(producerResultIndexMap); 573 assert(invProducerResultIndexMap && 574 "expected producer result indexig map to be invertible"); 575 for (unsigned argNum : llvm::seq<unsigned>(0, producer.getNumInputs())) { 576 // argMap is a map from producer loop -> producer arg tensor index. 577 AffineMap argMap = producer.getInputIndexingMap(argNum); 578 579 // Compose argMap with invProducerResultIndexMap to get a map from 580 // producer result tensor index -> producer arg tensor index. 581 AffineMap t1 = argMap.compose(invProducerResultIndexMap); 582 583 // Compose t1 with fusedConsumerArgIndexMap gives an indexing map from 584 // consumer loop/ fused loop -> producer arg tensor index. 585 AffineMap indexingMap = t1.compose(fusedConsumerArgIndexMap); 586 fusedOpIndexingMapAttrs.push_back(AffineMapAttr::get(indexingMap)); 587 } 588 } 589 590 /// Generate the region of the fused operation. The region of the fused op 591 /// must be empty. 592 static void generateFusedRegion(PatternRewriter &rewriter, Operation *fusedOp, 593 LinalgOp producer, LinalgOp consumer, 594 AffineMap consumerToProducerLoopsMap, 595 unsigned consumerIdx, unsigned nloops) { 596 // Build the region of the fused op. 597 Block &producerBlock = producer.getOperation()->getRegion(0).front(); 598 Block &consumerBlock = consumer.getOperation()->getRegion(0).front(); 599 Block *fusedBlock = new Block(); 600 fusedOp->getRegion(0).push_back(fusedBlock); 601 BlockAndValueMapping mapper; 602 OpBuilder::InsertionGuard guard(rewriter); 603 rewriter.setInsertionPointToStart(fusedBlock); 604 605 // The block arguments are 606 // [index_0, index_1, ... , 607 // consumer_operand_0, ... , consumer_operand_(`consumerIdx`-1), 608 // producer_operand_0, ... , producer_operand_(n-1)], 609 // consumer_operand_(`consumerIdx`), .. consumer_operand_(m-1)] 610 // , where n is the number of producer's operand and m is the number 611 // consumer's operand. 612 // If both `numProducerIndices` and `numConsumerIndices` are zero, this is a 613 // generic op. In this case, there are no indices in block arguments. 614 unsigned numProducerIndices = 615 isa<IndexedGenericOp>(producer.getOperation()) ? nloops : 0; 616 unsigned numConsumerIndices = 617 isa<IndexedGenericOp>(consumer.getOperation()) ? nloops : 0; 618 // Firstly, add all the indices to the block arguments. 619 for (unsigned i = 0, e = std::max(numProducerIndices, numConsumerIndices); 620 i < e; ++i) 621 fusedBlock->addArgument(rewriter.getIndexType()); 622 // Map the arguments for the unmodified args from the consumer. 623 for (auto consumerArg : llvm::enumerate(consumerBlock.getArguments())) { 624 if (consumerArg.index() == consumerIdx + numConsumerIndices) { 625 // Map the arguments for the args from the producer. 626 for (auto producerArg : llvm::enumerate(producerBlock.getArguments())) { 627 // If producer is an indexed_generic op, map the indices from consumer 628 // loop to producer loop (because the fusedOp is built based on 629 // consumer's perspective). 630 if (producerArg.index() < numProducerIndices) { 631 auto newIndex = rewriter.create<mlir::AffineApplyOp>( 632 producer.getLoc(), 633 consumerToProducerLoopsMap.getSubMap(producerArg.index()), 634 fusedBlock->getArguments().take_front(nloops)); 635 mapper.map(producerArg.value(), newIndex); 636 } else { 637 mapper.map(producerArg.value(), 638 fusedBlock->addArgument(producerArg.value().getType())); 639 } 640 } 641 continue; 642 } 643 644 // If consumer is an indexed_generic op, map the indices to the block 645 // arguments directly. Otherwise, add the same type of arugment and map to 646 // it. 647 if (consumerArg.index() < numConsumerIndices) { 648 mapper.map(consumerArg.value(), 649 fusedBlock->getArgument(consumerArg.index())); 650 } else { 651 mapper.map(consumerArg.value(), 652 fusedBlock->addArgument(consumerArg.value().getType())); 653 } 654 } 655 656 // Add operations from producer (except the yield operation) to the fused 657 // op. 658 for (auto &op : producerBlock.getOperations()) { 659 if (auto yieldOp = dyn_cast<YieldOp>(op)) { 660 // Lookup the value the yield operation is mapped to. 661 Value yieldVal = yieldOp.getOperand(0); 662 if (Value clonedVal = mapper.lookupOrNull(yieldVal)) 663 mapper.map( 664 consumerBlock.getArgument(consumerIdx + numConsumerIndices), 665 clonedVal); 666 continue; 667 } 668 rewriter.clone(op, mapper); 669 } 670 for (auto &op : consumerBlock.getOperations()) 671 rewriter.clone(op, mapper); 672 } 673 }; 674 } // namespace 675 676 /// Linearize the expressions in `sourceMap` based on the `reassociationMaps` 677 /// provided, given the shape of the source tensor that corresponds to the 678 /// `sourceMap`. Note that this implicitly assumes that the tensors dimensions 679 /// are "row-major" ordered logically. 680 /// 681 /// For example: 682 /// 683 /// %0 = op ... : tensor<?x?x4x5xf32> 684 /// with output index_map `affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>` 685 /// 686 /// and reshape: 687 /// %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>, 688 /// affine_map<(i, j, k, l) -> (j, k, l)>] : 689 /// tensor<?x?x4x5xf32> into tensor<?x?xf32> 690 /// 691 /// would be rewritten into: 692 /// %0 = op ... : tensor<?x?x4x5xf32> 693 /// with output index_map 694 /// `affine_map<(d0, d1, d2, d3) -> (d0, d1 * 20 + d2 * 5 + d3)>` 695 static AffineMap linearizeCollapsedDims(AffineMap sourceMap, 696 ArrayRef<int64_t> sourceShape, 697 ArrayRef<AffineMap> reassociationMaps) { 698 SmallVector<AffineExpr, 4> resultExprs; 699 resultExprs.reserve(reassociationMaps.size()); 700 ArrayRef<AffineExpr> sourceExprs = sourceMap.getResults(); 701 MLIRContext *context = sourceMap.getContext(); 702 703 // Compute the result exprs based on the reassociation maps. 704 for (AffineMap map : reassociationMaps) { 705 ArrayRef<AffineExpr> collapsedDims = map.getResults(); 706 // Assume that they are in-order and contiguous (already checked in 707 // verifier). 708 assert(!collapsedDims.empty()); 709 unsigned startDim = 710 collapsedDims.front().cast<AffineDimExpr>().getPosition(); 711 AffineExpr linearizedExpr = makeCanonicalStridedLayoutExpr( 712 sourceShape.slice(startDim, collapsedDims.size()), 713 sourceExprs.slice(startDim, collapsedDims.size()), context); 714 resultExprs.push_back(linearizedExpr); 715 } 716 return AffineMap::get(sourceMap.getNumDims(), sourceMap.getNumSymbols(), 717 resultExprs, context); 718 } 719 720 /// Checks if the `reshapeOp` can be fused with it consumer (if `asProducer` is 721 /// true) or its producer (if `asProducer` is false) given the indexing map at 722 /// its use. 723 static bool isTensorReshapeOpFusible(TensorReshapeOp reshapeOp, 724 AffineMap useIndexMap, bool asProducer) { 725 RankedTensorType returnType = reshapeOp.getResultType(); 726 RankedTensorType operandType = reshapeOp.getSrcType(); 727 // Reshape is fusible with its consumer (i.e. reshape as a producer) when its 728 // operand is of lesser rank than the result. Fusing when operand has higher 729 // rank will require use of mods and divs in the indexing maps of the fused op 730 // which would make it non-invertible. Similarly reshape is fused with its 731 // producer (i.e. reshape as consumer) only if the return type has lesser 732 // rank. 733 if ((asProducer && returnType.getRank() < operandType.getRank()) || 734 (!asProducer && operandType.getRank() < returnType.getRank())) 735 return false; 736 return useIndexMap.isIdentity(); 737 } 738 739 namespace { 740 /// Implementation of fusion on tensor ops when producer is a TensorReshapeOp. 741 template <typename LinalgOpTy> struct FuseTensorReshapeOpAsProducer { 742 static bool isFusible(TensorReshapeOp producer, LinalgOpTy consumer, 743 unsigned consumerIdx) { 744 return isTensorReshapeOpFusible( 745 producer, consumer.getInputIndexingMap(consumerIdx), true); 746 } 747 748 static Operation *fuse(TensorReshapeOp producer, LinalgOpTy consumer, 749 unsigned consumerIdx, PatternRewriter &rewriter, 750 OperationFolder *folder = nullptr) { 751 if (!isFusible(producer, consumer, consumerIdx)) 752 return nullptr; 753 754 // Compute the fused operands list, 755 SmallVector<Value, 2> fusedOperands(consumer.operand_begin(), 756 consumer.operand_end()); 757 fusedOperands[consumerIdx] = producer.src(); 758 759 // Compute indexing_maps for the fused operation. The indexing_maps for the 760 // operands of the consumers that arent fused are the same. 761 SmallVector<AffineMap, 4> fusedIndexMaps = 762 llvm::to_vector<4>(llvm::map_range( 763 consumer.indexing_maps(), [](Attribute attr) -> AffineMap { 764 return attr.cast<AffineMapAttr>().getValue(); 765 })); 766 767 // Compute the indexing map to use for the operand of the producer. 768 AffineMap modifiedMap = linearizeCollapsedDims( 769 fusedIndexMaps[consumerIdx], producer.getResultType().getShape(), 770 producer.getReassociationMaps()); 771 for (AffineExpr expr : modifiedMap.getResults()) { 772 if (!expr.isPureAffine()) 773 return nullptr; 774 } 775 fusedIndexMaps[consumerIdx] = modifiedMap; 776 777 // Further check that the resulting index maps can be fused and 778 // inverted. Without this the resultant op is not legal. 779 if (!inversePermutation(concatAffineMaps(fusedIndexMaps))) 780 return nullptr; 781 782 SmallVector<Attribute, 4> indexMapAttrs = llvm::to_vector<4>( 783 llvm::map_range(fusedIndexMaps, [](AffineMap map) -> Attribute { 784 return AffineMapAttr::get(map); 785 })); 786 auto fusedOp = rewriter.create<LinalgOpTy>( 787 rewriter.getUnknownLoc(), consumer.getResultTypes(), fusedOperands, 788 rewriter.getI64IntegerAttr(fusedOperands.size()), 789 rewriter.getI64IntegerAttr(consumer.getNumResults()), 790 rewriter.getArrayAttr(indexMapAttrs), consumer.iterator_types(), 791 /*doc=*/nullptr, 792 /*library_call=*/nullptr, 793 /*symbol_source=*/nullptr); 794 auto &fusedRegion = fusedOp.region(); 795 rewriter.cloneRegionBefore(consumer.region(), fusedRegion, 796 fusedRegion.begin()); 797 return fusedOp; 798 } 799 }; 800 801 /// Implementation of fusion on tensor ops when consumer is a TensorReshapeOp. 802 template <typename LinalgOpTy> struct FuseTensorReshapeOpAsConsumer { 803 static bool isFusible(LinalgOpTy producer, TensorReshapeOp consumer, 804 unsigned consumerIdx) { 805 return isTensorReshapeOpFusible(consumer, producer.getOutputIndexingMap(0), 806 false); 807 } 808 809 static Operation *fuse(LinalgOpTy producer, TensorReshapeOp consumer, 810 unsigned consumerIdx, PatternRewriter &rewriter, 811 OperationFolder *folder = nullptr) { 812 if (!isFusible(producer, consumer, consumerIdx)) 813 return nullptr; 814 815 // The indexing_maps for the operands of the fused operation are same as 816 // those for the operands of the producer. 817 SmallVector<AffineMap, 4> fusedIndexMaps = 818 llvm::to_vector<4>(llvm::map_range( 819 producer.indexing_maps(), [](Attribute attr) -> AffineMap { 820 return attr.cast<AffineMapAttr>().getValue(); 821 })); 822 // Compute the indexing map to use for the operand of the producer. 823 AffineMap modifiedMap = linearizeCollapsedDims( 824 producer.getOutputIndexingMap(0), consumer.getSrcType().getShape(), 825 consumer.getReassociationMaps()); 826 for (AffineExpr expr : modifiedMap.getResults()) { 827 if (!expr.isPureAffine()) 828 return nullptr; 829 } 830 fusedIndexMaps.back() = modifiedMap; 831 832 // Further check that the resulting index maps can be fused and 833 // inverted. Without this the resultant op is not legal. 834 if (!inversePermutation(concatAffineMaps(fusedIndexMaps))) 835 return nullptr; 836 837 SmallVector<Attribute, 4> indexMapAttrs = llvm::to_vector<4>( 838 llvm::map_range(fusedIndexMaps, [](AffineMap map) -> Attribute { 839 return AffineMapAttr::get(map); 840 })); 841 842 auto fusedOp = rewriter.create<LinalgOpTy>( 843 rewriter.getUnknownLoc(), consumer.getResultType(), 844 producer.getOperands(), 845 rewriter.getI64IntegerAttr(producer.getNumOperands()), 846 rewriter.getI64IntegerAttr(1), rewriter.getArrayAttr(indexMapAttrs), 847 producer.iterator_types(), 848 /*doc=*/nullptr, 849 /*library_call=*/nullptr, 850 /*symbol_source=*/nullptr); 851 auto &fusedRegion = fusedOp.region(); 852 rewriter.cloneRegionBefore(producer.region(), fusedRegion, 853 fusedRegion.begin()); 854 return fusedOp; 855 } 856 }; 857 858 /// Implementation of fusion on tensor ops when producer is a splat constant. 859 template <typename LinalgOpTy> struct FuseConstantOpAsProducer { 860 static bool isFusible(ConstantOp producer, LinalgOpTy consumer, 861 unsigned consumerIdx) { 862 return producer.getResult().getType().isa<RankedTensorType>() && 863 producer.value().template cast<DenseElementsAttr>().isSplat(); 864 } 865 866 static Operation *fuse(ConstantOp producer, LinalgOpTy consumer, 867 unsigned consumerIdx, PatternRewriter &rewriter, 868 OperationFolder *folder = nullptr) { 869 if (!isFusible(producer, consumer, consumerIdx)) 870 return nullptr; 871 872 // The indexing_maps for the operands of the fused operation are same as 873 // those for the operands of the consumer without the indexing map at 874 // consumerIdx 875 SmallVector<AffineMap, 4> fusedIndexMaps = 876 llvm::to_vector<4>(llvm::map_range( 877 consumer.indexing_maps(), [](Attribute attr) -> AffineMap { 878 return attr.cast<AffineMapAttr>().getValue(); 879 })); 880 fusedIndexMaps.erase(std::next(fusedIndexMaps.begin(), consumerIdx)); 881 882 // The operands list is same as the consumer with the argument for constant 883 // index dropped. 884 SmallVector<Value, 4> fusedOperands(consumer.operand_begin(), 885 consumer.operand_end()); 886 fusedOperands.erase(std::next(fusedOperands.begin(), consumerIdx)); 887 888 // Create a constant scalar value from the splat constant. 889 Value scalarConstant = rewriter.create<ConstantOp>( 890 producer.getLoc(), 891 producer.value().template cast<DenseElementsAttr>().getSplatValue()); 892 893 auto fusedOp = rewriter.create<LinalgOpTy>( 894 rewriter.getUnknownLoc(), consumer.getResultTypes(), fusedOperands, 895 rewriter.getI64IntegerAttr(consumer.getNumOperands() - 1), 896 rewriter.getI64IntegerAttr(consumer.getNumResults()), 897 rewriter.getAffineMapArrayAttr(fusedIndexMaps), 898 consumer.iterator_types(), 899 /*doc=*/nullptr, 900 /*library_call=*/nullptr, 901 /*symbol_source=*/nullptr); 902 903 // Map the block argument corresponding to the replaced argument with the 904 // scalar constant. 905 Region &consumerRegion = consumer.region(); 906 Block &entryBlock = *consumerRegion.begin(); 907 unsigned argIndex = 908 entryBlock.getNumArguments() - consumer.getNumOperands() + consumerIdx; 909 BlockAndValueMapping mapping; 910 mapping.map(entryBlock.getArgument(argIndex), scalarConstant); 911 Region &fusedRegion = fusedOp.region(); 912 rewriter.cloneRegionBefore(consumerRegion, fusedRegion, fusedRegion.begin(), 913 mapping); 914 return fusedOp; 915 } 916 }; 917 918 } // namespace 919 920 Operation *mlir::linalg::fuseTensorOps(PatternRewriter &rewriter, 921 Operation *consumer, 922 unsigned consumerIdx, 923 OperationFolder *folder) { 924 if (consumerIdx >= consumer->getNumOperands()) 925 return nullptr; 926 Operation *producer = consumer->getOperand(consumerIdx).getDefiningOp(); 927 if (!producer || producer->getNumResults() != 1) 928 return nullptr; 929 930 // Fuse when consumer is GenericOp or IndexedGenericOp. 931 if (isa<GenericOp, IndexedGenericOp>(consumer)) { 932 auto linalgOpConsumer = cast<LinalgOp>(consumer); 933 if (!linalgOpConsumer.hasTensorSemantics()) 934 return nullptr; 935 if (isa<GenericOp, IndexedGenericOp>(producer)) { 936 auto linalgOpProducer = cast<LinalgOp>(producer); 937 if (linalgOpProducer.hasTensorSemantics()) 938 return FuseGenericOpsOnTensors::fuse(linalgOpProducer, linalgOpConsumer, 939 consumerIdx, rewriter, folder); 940 } else if (auto reshapeOpProducer = dyn_cast<TensorReshapeOp>(producer)) { 941 if (auto genericOpConsumer = dyn_cast<GenericOp>(consumer)) { 942 return FuseTensorReshapeOpAsProducer<GenericOp>::fuse( 943 reshapeOpProducer, genericOpConsumer, consumerIdx, rewriter, 944 folder); 945 } else if (auto indexedGenericOpConsumer = 946 dyn_cast<IndexedGenericOp>(consumer)) { 947 return FuseTensorReshapeOpAsProducer<IndexedGenericOp>::fuse( 948 reshapeOpProducer, indexedGenericOpConsumer, consumerIdx, rewriter, 949 folder); 950 } 951 } else if (auto constantOpProducer = dyn_cast<ConstantOp>(producer)) { 952 if (auto genericOpConsumer = dyn_cast<GenericOp>(consumer)) { 953 return FuseConstantOpAsProducer<GenericOp>::fuse( 954 constantOpProducer, genericOpConsumer, consumerIdx, rewriter, 955 folder); 956 } 957 } 958 return nullptr; 959 } 960 961 // Fuse when consumer is a TensorReshapeOp. 962 if (TensorReshapeOp reshapeOp = dyn_cast<TensorReshapeOp>(consumer)) { 963 if (auto genericOpProducer = dyn_cast<GenericOp>(producer)) { 964 if (genericOpProducer.hasTensorSemantics()) 965 return FuseTensorReshapeOpAsConsumer<GenericOp>::fuse( 966 genericOpProducer, reshapeOp, consumerIdx, rewriter, folder); 967 } else if (auto indexedGenericOpProducer = 968 dyn_cast<IndexedGenericOp>(producer)) { 969 if (indexedGenericOpProducer.hasTensorSemantics()) 970 return FuseTensorReshapeOpAsConsumer<IndexedGenericOp>::fuse( 971 indexedGenericOpProducer, reshapeOp, consumerIdx, rewriter, folder); 972 } 973 return nullptr; 974 } 975 976 return nullptr; 977 } 978 979 namespace { 980 /// Patterns to fuse a generic op, with the producer of its operands. 981 template <typename LinalgOpTy> 982 struct FuseTensorOps : public OpRewritePattern<LinalgOpTy> { 983 using OpRewritePattern<LinalgOpTy>::OpRewritePattern; 984 985 LogicalResult matchAndRewrite(LinalgOpTy op, 986 PatternRewriter &rewriter) const override { 987 // Find the first operand that is defined by another generic op on tensors. 988 for (auto operandNum : 989 llvm::seq<unsigned>(0, op.getOperation()->getNumOperands())) { 990 Operation *producer = 991 op.getOperation()->getOperand(operandNum).getDefiningOp(); 992 if (Operation *fusedOp = fuseTensorOps(rewriter, op, operandNum)) { 993 rewriter.replaceOp(op, fusedOp->getResults()); 994 if (producer && llvm::all_of(producer->getResults(), 995 [](Value val) { return val.use_empty(); })) 996 rewriter.eraseOp(producer); 997 return success(); 998 } 999 } 1000 return failure(); 1001 } 1002 }; 1003 1004 /// Pass that fuses generic ops on tensors. Used only for testing. 1005 struct FusionOfTensorOpsPass 1006 : public LinalgFusionOfTensorOpsBase<FusionOfTensorOpsPass> { 1007 void runOnOperation() override { 1008 OwningRewritePatternList patterns; 1009 Operation *op = getOperation(); 1010 populateLinalgTensorOpsFusionPatterns(op->getContext(), patterns); 1011 applyPatternsAndFoldGreedily(op->getRegions(), patterns); 1012 }; 1013 }; 1014 1015 struct LinalgFusionPass : public LinalgFusionBase<LinalgFusionPass> { 1016 void runOnFunction() override { fuseLinalgOpsGreedily(getFunction()); } 1017 }; 1018 } // namespace 1019 1020 void mlir::populateLinalgTensorOpsFusionPatterns( 1021 MLIRContext *context, OwningRewritePatternList &patterns) { 1022 patterns.insert<FuseTensorOps<GenericOp>, FuseTensorOps<IndexedGenericOp>, 1023 FuseTensorOps<TensorReshapeOp>>(context); 1024 } 1025 1026 std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgFusionPass() { 1027 return std::make_unique<LinalgFusionPass>(); 1028 } 1029 1030 std::unique_ptr<Pass> mlir::createLinalgFusionOfTensorOpsPass() { 1031 return std::make_unique<FusionOfTensorOpsPass>(); 1032 } 1033