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/Arithmetic/IR/Arithmetic.h" 16 #include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h" 17 #include "mlir/Dialect/Linalg/IR/Linalg.h" 18 #include "mlir/Dialect/Linalg/Passes.h" 19 #include "mlir/Dialect/Linalg/Transforms/Transforms.h" 20 #include "mlir/Dialect/Linalg/Utils/Utils.h" 21 #include "mlir/Dialect/MemRef/IR/MemRef.h" 22 #include "mlir/Dialect/Tensor/IR/Tensor.h" 23 #include "mlir/IR/AffineExpr.h" 24 #include "mlir/IR/AffineMap.h" 25 #include "mlir/IR/Dominance.h" 26 #include "mlir/Support/LLVM.h" 27 #include "mlir/Transforms/GreedyPatternRewriteDriver.h" 28 #include "mlir/Transforms/RegionUtils.h" 29 #include "llvm/ADT/MapVector.h" 30 #include "llvm/ADT/ScopeExit.h" 31 #include "llvm/Support/CommandLine.h" 32 #include "llvm/Support/Debug.h" 33 34 #include <set> 35 36 #define DEBUG_TYPE "linalg-fusion" 37 38 using namespace mlir; 39 using namespace mlir::linalg; 40 41 /// Implements a simple high-level fusion pass on linalg structured operations. 42 /// 43 /// In each block, linalg ops are processed in reverse textual order. 44 /// Given a linalg op `O`, fusion occurs by: 45 /// 1. inspecting the linalg ops that write into the views read by `O`. There 46 /// are 2 cases: 47 /// a) buffer case: use the SSA value of the views and a simple alias 48 /// analysis on subview ops to determine producer-consumer dependences; 49 /// b) tensor case: use SSA use-def chains on extract_slice ops; 50 /// 2. greedily fuse the linalg ops that produce the subview/extract_slice. 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 struct ShapeDimension { 58 Value shape; 59 unsigned dimension; 60 }; 61 62 // Given an `op`, returns the first (`shape`, `dimension`) pair that identifies 63 // the loop range at `loopDepth`. The semantics of the loopToOperandRangesMaps 64 // guarantees at least one such dimension is found. If multiple candidates exist 65 // they must agree by construction (i.e. have the same size) and we just return 66 // the first one. 67 static ShapeDimension 68 getShapeDefiningLoopRange(LinalgOp op, unsigned loopDepth, 69 bool fromSubViewOpOnly = false) { 70 // Iterate over the inputs and outputs in order. 71 // Extract the subranges from the linearized ranges. 72 for (OpOperand *opOperand : op.getInputAndOutputOperands()) { 73 // The method `getRangeFromOperandShape` requires using SubViewOp or 74 // ExtractSliceOps. If the value isn't defined from there continue. 75 // todo: The method should be adapted to get the values from 76 // `ViewInterface`. The interface needs a `getOrCreateRanges` method which 77 // currently returns a `linalg.range`. The fix here is to move this op to 78 // `std` dialect and add the method to `ViewInterface`. 79 if (fromSubViewOpOnly && 80 !isa_and_nonnull<memref::SubViewOp, tensor::ExtractSliceOp>( 81 opOperand->get().getDefiningOp())) 82 continue; 83 84 AffineMap map = op.getTiedIndexingMap(opOperand); 85 LLVM_DEBUG(llvm::dbgs() << "getShapeDefiningLoopRange I/O idx: " 86 << opOperand->getOperandNumber() << "\n"); 87 LLVM_DEBUG(llvm::dbgs() 88 << "getShapeDefiningLoopRange map: " << map << "\n"); 89 SmallVector<Value, 8> shapeRanges(map.getNumResults(), nullptr); 90 for (const auto &en : llvm::enumerate(map.getResults())) { 91 auto dimExpr = en.value().dyn_cast<AffineDimExpr>(); 92 if (!dimExpr) 93 continue; 94 if (loopDepth == en.value().cast<AffineDimExpr>().getPosition()) { 95 LLVM_DEBUG(llvm::dbgs() << "getShapeDefiningLoopRange loopDepth: " 96 << loopDepth << "\n"); 97 LLVM_DEBUG(llvm::dbgs() << "getShapeDefiningLoopRange shape: " 98 << opOperand->get() << "\n"); 99 return ShapeDimension{opOperand->get(), 100 static_cast<unsigned>(en.index())}; 101 } 102 } 103 } 104 llvm_unreachable("Expect to be able to extract a shape defining loop range"); 105 } 106 107 static SmallVector<Value> getTiledOperands(LinalgOp producer) { 108 return producer.getInputAndOutputOperands(); 109 } 110 111 /// Fuses the producer by cloning the `producer`. The `fusedLoopsAndRanges` 112 /// provides the loop range information for the fused loops. The rest are 113 /// obtained from the producer itself, since they are not tiled + fused. 114 static LinalgOp fuse(OpBuilder &b, LinalgOp producer, 115 const DenseMap<unsigned, Range> &fusedLoopsAndRanges) { 116 SmallVector<Value, 8> ivs, tileSizes, sizeBounds; 117 SmallVector<Range, 8> loopRanges; 118 Location loc = producer.getLoc(); 119 auto zero = b.create<arith::ConstantIndexOp>(loc, 0); 120 auto one = b.create<arith::ConstantIndexOp>(loc, 1); 121 122 for (unsigned i = 0, e = producer.getNumLoops(); i < e; ++i) { 123 auto shapeDim = getShapeDefiningLoopRange(producer, i); 124 Value dim = createOrFoldDimOp(b, loc, shapeDim.shape, shapeDim.dimension); 125 sizeBounds.push_back(dim); 126 auto it = fusedLoopsAndRanges.find(i); 127 if (it != fusedLoopsAndRanges.end()) { 128 ivs.push_back(it->second.offset); 129 tileSizes.push_back(it->second.size); 130 loopRanges.push_back(it->second); 131 LLVM_DEBUG(llvm::dbgs() << "tiled loop#" << i << " with LoopRange " 132 << loopRanges.back() << "\n"); 133 } else { 134 tileSizes.push_back(zero); 135 loopRanges.push_back(Range{zero, dim, one}); 136 LLVM_DEBUG(llvm::dbgs() << "full loop#" << i << " with LoopRange " 137 << loopRanges.back() << "\n"); 138 } 139 } 140 141 SmallVector<Value, 8> clonedShapes; 142 clonedShapes.reserve(producer.getNumInputsAndOutputs()); 143 144 // Compute subranges for all tensor input/output operands. 145 clonedShapes.append(makeTiledShapes( 146 b, loc, producer, getTiledOperands(producer), ivs, tileSizes, sizeBounds, 147 /**omitPartialTileCheck=*/false)); 148 149 // Iterate over the results in order. 150 // Extract the subtensor type from the linearized range. 151 // Since we do not enforce any canonicalizations on the fly, this is always 152 // fully dynamic at construction time. 153 SmallVector<Type, 4> resultTypes; 154 resultTypes.reserve(producer->getNumResults()); 155 for (RankedTensorType t : producer.getOutputTensorTypes()) { 156 unsigned rank = t.getRank(); 157 SmallVector<int64_t, 4> staticOffsetsVector( 158 rank, ShapedType::kDynamicStrideOrOffset); 159 SmallVector<int64_t, 4> staticSizesVector(rank, ShapedType::kDynamicSize); 160 SmallVector<int64_t, 4> staticStridesVector( 161 rank, ShapedType::kDynamicStrideOrOffset); 162 resultTypes.push_back(tensor::ExtractSliceOp::inferResultType( 163 t.cast<RankedTensorType>(), staticOffsetsVector, staticSizesVector, 164 staticStridesVector)); 165 } 166 167 Operation *clonedOp = producer.clone(b, loc, resultTypes, clonedShapes); 168 169 // Shift all IndexOp results by the tile offset. 170 SmallVector<Value> allIvs; 171 llvm::transform(loopRanges, std::back_inserter(allIvs), 172 [](Range range) { return range.offset; }); 173 offsetIndices(b, clonedOp, allIvs); 174 175 return clonedOp; 176 } 177 178 /// Get the loop range for a dimension `dim` based on the `shapedOperand`. It is 179 /// expected to be defined by a subview op or an extract_slice op. 180 static Range getRangeFromOperandShape(OpBuilder &b, Location loc, 181 Value shapedOperand, unsigned dim) { 182 Operation *shapeProducingOp = shapedOperand.getDefiningOp(); 183 if (auto subViewOp = dyn_cast<memref::SubViewOp>(shapeProducingOp)) 184 return subViewOp.getOrCreateRanges(b, loc)[dim]; 185 if (auto sliceOp = dyn_cast<tensor::ExtractSliceOp>(shapeProducingOp)) 186 return sliceOp.getOrCreateRanges(b, loc)[dim]; 187 llvm_unreachable("SubviewOp or ExtractSliceOp expected"); 188 } 189 190 /// Fuses the producer into the loop immediately enclosing the consumer. 191 /// This is achieved by "recomputing" the producer at the time it 192 /// is needed just before the consumer. 193 static LinalgOp fuse(OpBuilder &b, LinalgOp producerOp, AffineMap producerMap, 194 OpOperand &consumerOpOperand) { 195 LLVM_DEBUG(llvm::dbgs() << "Producer map: " << producerMap << "\n"); 196 DenseMap<unsigned, Range> fusedLoopsAndRanges; 197 Value shapedOperand = consumerOpOperand.get(); 198 for (const auto &en : llvm::enumerate(producerMap.getResults())) { 199 unsigned posInProducerLoop = en.value().cast<AffineDimExpr>().getPosition(); 200 fusedLoopsAndRanges[posInProducerLoop] = getRangeFromOperandShape( 201 b, consumerOpOperand.getOwner()->getLoc(), shapedOperand, en.index()); 202 } 203 return fuse(b, producerOp, fusedLoopsAndRanges); 204 } 205 206 // Encode structural fusion safety preconditions. 207 // Some of these will be lifted in the future with better analysis. 208 static bool isStructurallyFusableProducer(LinalgOp producer, Value consumedView, 209 LinalgOp consumer) { 210 assert(producer.hasBufferSemantics() && 211 "expected linalg op with buffer semantics"); 212 assert(consumer.hasBufferSemantics() && 213 "expected linalg op with buffer semantics"); 214 if (producer.getNumOutputs() != 1) { 215 LLVM_DEBUG(llvm::dbgs() << "\nNot structurally fusable (multi-output)"); 216 return false; 217 } 218 // Only fuse when the producer block dominates. 219 DominanceInfo dom(producer.getOperation()); 220 if (!dom.dominates(producer->getBlock(), consumer->getBlock())) { 221 LLVM_DEBUG( 222 llvm::dbgs() 223 << "\nNot structurally fusable (producer block does not dominate)"); 224 return false; 225 } 226 return true; 227 } 228 229 bool mlir::linalg::isProducerLastWriteOfView(const LinalgDependenceGraph &graph, 230 LinalgOp consumer, 231 Value consumedView, 232 LinalgOp producer) { 233 assert(producer.hasBufferSemantics() && 234 "expected linalg op with buffer semantics"); 235 assert(consumer.hasBufferSemantics() && 236 "expected linalg op with buffer semantics"); 237 // Make some simple structural checks that alleviate the need for more 238 // complex analyses. 239 if (!isStructurallyFusableProducer(producer, consumedView, consumer)) { 240 LLVM_DEBUG(llvm::dbgs() << "\n***Not static last write due to structure:\t" 241 << *producer.getOperation()); 242 return false; 243 } 244 // Check for any interleaved write to consumedView. 245 if (!graph.findCoveringWrites(producer, consumer, consumedView).empty()) { 246 LLVM_DEBUG(llvm::dbgs() << "\n***Not fusable due to interleaved write:\t" 247 << *producer.getOperation()); 248 return false; 249 } 250 return true; 251 } 252 253 bool mlir::linalg::isFusableInto(const LinalgDependenceGraph &graph, 254 LinalgOp consumer, Value consumedView, 255 LinalgOp producer) { 256 assert(producer.hasBufferSemantics() && 257 "expected linalg op with buffer semantics"); 258 assert(consumer.hasBufferSemantics() && 259 "expected linalg op with buffer semantics"); 260 if (!isProducerLastWriteOfView(graph, consumer, consumedView, producer)) 261 return false; 262 // Check for any fusion-preventing dependence to any shape read/written that 263 // would violate dependences. 264 if (!graph.findCoveringDependences(producer, consumer).empty()) { 265 LLVM_DEBUG(llvm::dbgs() 266 << "\n***Not fusable due to an interleaved dependence:\t" 267 << *producer.getOperation()); 268 return false; 269 } 270 return true; 271 } 272 273 /// For `consumer` with buffer semantics, find the Linalg operation on buffers 274 /// that is the last writer of `consumerOpOperand`. For now the fusable 275 /// dependence is returned as an instance of the `dependenceGraph`. 276 static FailureOr<LinalgDependenceGraph::LinalgDependenceGraphElem> 277 findFusableProducer(OpOperand &consumerOpOperand, 278 const LinalgDependenceGraph &dependenceGraph) { 279 LLVM_DEBUG(llvm::dbgs() << "findFusableProducer for: " 280 << consumerOpOperand.get() << " @" 281 << consumerOpOperand.getOperandNumber() << " in " 282 << *consumerOpOperand.getOwner() << "\n"); 283 LinalgOp consumerOp = dyn_cast<LinalgOp>(consumerOpOperand.getOwner()); 284 if (!consumerOp) 285 return failure(); 286 287 // Only consider RAW and WAW atm. 288 for (auto depType : { 289 LinalgDependenceGraph::DependenceType::RAW, 290 LinalgDependenceGraph::DependenceType::WAW, 291 }) { 292 LLVM_DEBUG(llvm::dbgs() 293 << "Dependencies into: " << *consumerOp.getOperation() << "\n"); 294 for (auto dependence : llvm::make_filter_range( 295 dependenceGraph.getDependencesInto(consumerOp, depType), 296 [&](LinalgDependenceGraph::LinalgDependenceGraphElem elem) { 297 LLVM_DEBUG(llvm::dbgs() << "Inspect dependence btw: " 298 << elem.getIndexingValue() << " and " 299 << elem.getDependentValue() << "\n"); 300 Value v = elem.getIndexingValue(); 301 Optional<unsigned> operandNum = 302 elem.getIndexingOpViewOperandNum(); 303 return isa<LinalgOp>(elem.getDependentOp()) && 304 v == consumerOpOperand.get() && operandNum && 305 *operandNum == consumerOpOperand.getOperandNumber(); 306 })) { 307 // Consumer consumes this view, `isStructurallyFusableProducer` also 308 // checks whether it is a strict subview of the producer view. 309 auto producer = cast<LinalgOp>(dependence.getDependentOp()); 310 LLVM_DEBUG(llvm::dbgs() 311 << "\n" 312 << LinalgDependenceGraph::getDependenceTypeStr(depType) 313 << "producer: " << *dependence.getDependentOp() 314 << " view: " << dependence.getDependentValue() << "\n"); 315 316 // If the producer and consumer have tensor semantics, the only dependence 317 // between them is through a RAW dependence and they are fusable by 318 // construction. For buffer semantics need additional checks. 319 if (producer.hasBufferSemantics() && consumerOp.hasBufferSemantics() && 320 isFusableInto(dependenceGraph, consumerOp, consumerOpOperand.get(), 321 producer)) 322 return dependence; 323 if (producer.hasTensorSemantics() && consumerOp.hasTensorSemantics()) { 324 assert(dependence.dependenceType == 325 LinalgDependenceGraph::DependenceType::RAW); 326 return dependence; 327 } 328 } 329 } 330 return failure(); 331 } 332 333 FailureOr<FusionInfo> 334 mlir::linalg::fuseProducerOfBuffer(OpBuilder &b, OpOperand &consumerOpOperand, 335 const LinalgDependenceGraph &graph) { 336 Optional<LinalgDependenceGraph::LinalgDependenceGraphElem> fusableDependence = 337 findFusableProducer(consumerOpOperand, graph); 338 if (!fusableDependence) 339 return failure(); 340 341 LinalgOp producerOp = dyn_cast<LinalgOp>(fusableDependence->getDependentOp()); 342 if (!producerOp) 343 return failure(); 344 345 // If producer is already in the same block as consumer, we are done. 346 if (consumerOpOperand.get().getParentBlock() == 347 fusableDependence->getDependentValue().getParentBlock()) 348 return failure(); 349 350 Optional<AffineMap> producerMap = 351 fusableDependence->getDependentOpViewIndexingMap(); 352 if (!producerMap) 353 return failure(); 354 355 // Must be a subview or an extract_slice to guarantee there are loops we can 356 // fuse into. 357 auto subView = consumerOpOperand.get().getDefiningOp<memref::SubViewOp>(); 358 if (!subView) { 359 LLVM_DEBUG(llvm::dbgs() << "\nNot fusable (not a subview)"); 360 return failure(); 361 } 362 363 // Fuse `producer` just before `consumer`. 364 OpBuilder::InsertionGuard g(b); 365 b.setInsertionPoint(consumerOpOperand.getOwner()); 366 LLVM_DEBUG(llvm::dbgs() << "Fuse into consumer: " 367 << *consumerOpOperand.getOwner() << "\n"); 368 369 auto fusedProducer = fuse(b, producerOp, *producerMap, consumerOpOperand); 370 return FusionInfo{producerOp, fusedProducer}; 371 } 372 373 /// Walk back use-def chain through scf::For yields. 374 /// Sets `producer` and `outputIndex` if it finds a producer LinalgOp 375 376 // TODO(ravishankarm, ntv): This can be moved into the dependence graphs 377 // dependence tracking since the dependence tracking is similar to what is done 378 // w.r.t to buffers. 379 static void getProducerOfTensor(Value tensor, OpResult &opResult) { 380 if (!tensor.getType().isa<RankedTensorType>()) 381 return; 382 383 while (true) { 384 LLVM_DEBUG(llvm::dbgs() << "\ngetProducerOfTensor: " << tensor); 385 if (auto linalgOp = tensor.getDefiningOp<LinalgOp>()) { 386 opResult = tensor.cast<OpResult>(); 387 return; 388 } 389 if (auto sliceOp = tensor.getDefiningOp<tensor::ExtractSliceOp>()) { 390 tensor = sliceOp.getSource(); 391 continue; 392 } 393 if (auto blockArg = tensor.dyn_cast<BlockArgument>()) { 394 if (auto forOp = blockArg.getDefiningOp<scf::ForOp>()) { 395 tensor = *(forOp.getIterOperands().begin() + blockArg.getArgNumber()); 396 continue; 397 } 398 } 399 return; 400 } 401 } 402 403 FailureOr<FusionInfo> 404 mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpOperand &consumerOpOperand) { 405 Value inputTensor = consumerOpOperand.get(); 406 OpResult producerOpResult; 407 getProducerOfTensor(inputTensor, producerOpResult); 408 if (!producerOpResult) { 409 LLVM_DEBUG(llvm::dbgs() << "\nUnable to find producer"); 410 return failure(); 411 } 412 return fuseProducerOfTensor(b, producerOpResult, consumerOpOperand); 413 } 414 415 FailureOr<FusionInfo> 416 mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult, 417 OpOperand &consumerOpOperand) { 418 auto producerOp = dyn_cast<LinalgOp>(producerOpResult.getOwner()); 419 if (!producerOp) 420 return failure(); 421 422 LinalgOp consumerOp = dyn_cast<LinalgOp>(consumerOpOperand.getOwner()); 423 if (!consumerOp) 424 return failure(); 425 426 Value inputTensor = consumerOpOperand.get(); 427 428 // Must be an extract_slice op to guarantee there are loops we can fuse into. 429 auto sliceOp = inputTensor.getDefiningOp<tensor::ExtractSliceOp>(); 430 if (!sliceOp) { 431 LLVM_DEBUG(llvm::dbgs() 432 << "\nNot fusable, not an extract_slice op: " << inputTensor); 433 return failure(); 434 } 435 436 // If producer is already in the same block as consumer, we are done. 437 if (consumerOpOperand.get().getParentBlock() == 438 producerOpResult.getParentBlock()) 439 return failure(); 440 441 // Insert fused `producer` just before `consumer`. 442 OpBuilder::InsertionGuard g(b); 443 b.setInsertionPoint(consumerOp); 444 LLVM_DEBUG(llvm::dbgs() << "Fuse into consumer: " << *consumerOp << "\n"); 445 OpOperand *opOperand = 446 producerOp.getOutputOperand(producerOpResult.getResultNumber()); 447 LinalgOp fusedProducer = 448 fuse(b, producerOp, producerOp.getTiedIndexingMap(opOperand), 449 consumerOpOperand); 450 451 // Replace use. 452 // Canonicalizations are not guaranteed to have happened before constructing 453 // `fusedProducer`. In the tensor case this can result in temporary type 454 // mismatches. Insert a `tensor.cast` op to propagate the transformation 455 // invariant that types are compatible. 456 Value def = fusedProducer->getResult(producerOpResult.getResultNumber()); 457 Type consumerType = consumerOpOperand.get().getType(); 458 if (consumerType != def.getType()) 459 def = b.create<tensor::CastOp>(fusedProducer.getLoc(), consumerType, def); 460 consumerOpOperand.set(def); 461 return FusionInfo{cast<LinalgOp>(producerOpResult.getOwner()), fusedProducer}; 462 } 463