xref: /llvm-project/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp (revision f9c8febc522c2d26a44d4881f015e0e11e4f9167)
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