[mlir][Linalg] Improve the logic to perform tile and fuse with better dependence tracking.

[mravishankar]

Summary

This change does two main things

1. An operation might have multiple dependences to the same producer. Not tracking them correctly can result in incorrect code generation with fusion. To rectify this the dependence tracking needs to also have the operand number in the consumer.
2. Improve the logic used to find the fused loops making it easier to follow. The only constraint for fusion is that linalg ops (on buffers) have update semantics for the result. Fusion should be such that only one iteration of the fused loop (which is also a tiled loop) must touch only one (disjoint) tile of the output. This could be relaxed by allowing for recomputation that is the default when oeprands are tensors, or can be made legal with promotion of the fused view (in future).

[hanchung]

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

157–159

IIRC, "W" would be index + number of inputs, and "R" would be number of inputs. I feel we could miss any of one in a future change, so I think it would be good to have something like getLinalgOpViewW(int idx, LinalgOp op) and getLinalgOpViewR(int idx, LinalgOp op) and both of them will return a LinalgOpView.

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

654

Do we want tabs \t for indents?

[nicolasvasilache]

Can we split out the NFC part that extends support for operand + index in the dependence analysis and land that separately ?

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

535

It seems to me you are computing a simple projection of map.
Is a composition of map with AffineMap::getSubMap appropriate for this?

598

You probably want to document the case for which you have done all this:
A + A^T where A comes from a matmul ?

[nicolasvasilache]

Generally makes sense, thanks for fixing some of the underlying issues.
I left some things to address so I'll wait till next week for approval.
Note that in tensor world, recompute will become much more easy and side effects will become optimizations.
Would we good to think about potential refactorings of the fusable dependences towards using in cost models (not in this revision of course :) ).

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

699

The control flow under this loop is somewhat unclear to me.
Some conditions skip the current operand, some conditions exit the functions and the last one breaks out of the loop.

Can we restructure/fission this loop to first rule out the empty case ?
Processing the data 2 or 3 times seems minor compared to the expected gain in readability.

712

expose + reuse "areOpsInSameBlock" ?

730

typo

846

Please add some named accessors to make this nicer to read.

849

this seems to leak internal linalgop operand ordering .
Can we restructure and use input/output idx to hide such logic below the LinalgStructuredOpsInterface ?

[mravishankar]

Rebase

[mravishankar]

In D90579#2377499, @nicolasvasilache wrote:

Can we split out the NFC part that extends support for operand + index in the dependence analysis and land that separately ?

This seems a bit involved. I will have to first update the old code in Fusion.cpp to use the new dependence tracking and then change the code. Could we land it as one patch?

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

157–159

Changed it to use the getIndexOfShapedOperand method.

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

535

I couldnt work out how to use a compose to do this. Suggestions?

598

Right now this case doesnt really come here. It is treated as an unfusable dependence. When I add the promotion (or test it on tensors) I can add that example. It is actually easier since we just need to take the closest outer parallel loops.

654

Maybe, but its a simple check here. Not sure where to put that function.

699

Split the loop to first find the index of operand to fuse, and then do the fusion of those

846

I am not sure I understand what you mean by named accessors

[mravishankar]

Rebase and address comments

[antiagainst]

mlir/include/mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h

48

Nit: I'd suggest to use operandIndex. Number usually means total count to me.

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

433 ↗

(On Diff #303944)

SmallSet? I'd assume typically we don't have a large number of indices to fuse.

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

552

A fusable loop's iteration space only touches ...

564

The actual fusable loops are ..

[mravishankar]

Rebase

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

535

On second look, not with getSubMap (which is already a composition with a permutation) but I think something like this could work for you:

SmallVector<AffineExpr, 4> projection;
for (auto attr : llvm::enumerate(iteratorTypes))
  if (!isParallelIterator(attr.value()))
    projection.push_back(getAffineDimExpr(attr.index(), ctx));
unsigned projectionRank = projection.size();
projection.append(iteratorTypes.size() - projectionRank, getAffineConstant(0, ctx));
return AffineMap::get(projection.size(), 0, projectionRank, ctx).compose(map).getMajorSubMap(projectionRank);

The first n-1 lines could be refactored into a new AffineMap::getProjectionMap(ArrayRef<unsigned> positions, unsigned codomainRank).

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

535

actually 3 more things:

1. the proper expression is prob map.compose(AffineMap::get(...))
2. s/!ifParallelIterator/isParallelIterator
3. you may need to drop the 0 results of the output AffineMap for your use case.

[hanchung]

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td

578

s/then/the

581

s/retTyp/retTy

586–588

format this? (see above examples)

593

s/then/the

596

s/retTyp/retTy

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

157–159

nice!

249–254

[optional] Do we want to capture the value for consumer? it's only used in a later statement -- Value consumerView = ...

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

375

I think it's okay to pass without reference, because it is unsigned and we won't update the variable. It's a bit like passing an unsigned value without reference to a function.

[mravishankar]

Address comments

[mravishankar]

Addressing missed minor comments.

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

249–254

True, but the code is less readable IMO.

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

535

Nice! It actually helped me concretize something I was stuck at downstream of this change. Added the method as you asked for. PTAL

[hanchung]

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

801–803

nit: remove trivial braces

It's good to fix to me because you touched this condition. But it's okay to leave it this way because it doesn't really relate to the change...

mlir/include/mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h

 };
 
 /// Data structure for holding a dependence graph that operates on LinalgOp and
 /// views as SSA values.
 class LinalgDependenceGraph {
 public:
   struct LinalgOpView {
     Operation *op;
-    Value view;
+    unsigned operandNum;
   };
   struct LinalgDependenceGraphElem {
     // dependentOpView may be either:
     //   1. src in the case of dependencesIntoGraphs.
     //   2. dst in the case of dependencesFromDstGraphs.
     LinalgOpView dependentOpView;
     // View in the op that is used to index in the graph:
     //   1. src in the case of dependencesFromDstGraphs.
     //   2. dst in the case of dependencesIntoGraphs.
-    Value indexingView;
+    LinalgOpView indexingOpView;
   };
   using LinalgDependences = SmallVector<LinalgDependenceGraphElem, 8>;
   using DependenceGraph = DenseMap<Operation *, LinalgDependences>;
   using dependence_iterator = LinalgDependences::const_iterator;
   using dependence_range = iterator_range<dependence_iterator>;
 
   enum DependenceType { RAR = 0, RAW, WAR, WAW, NumTypes };
   static StringRef getDependenceTypeStr(DependenceType depType);

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td

InterfaceMethod<
         return llvm::to_vector<4>(
             llvm::map_range(inputOutputTypes, [](Type type) -> ShapedType {
               return type.cast<ShapedType>();
             }));
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
-        Return the position of the shaped operand in the operand list.
+        Return the first position of the shaped operand in the operand list.
       }],
       /*retTy=*/"Optional<unsigned>",
       /*methodName=*/"getIndexOfShapedOperand",
       /*args=*/(ins "Value":$value),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         Optional<unsigned> inputIndex = getIndexOfInput(value);
         if (inputIndex.hasValue()) return inputIndex.getValue();
         Optional<unsigned> outputIndex = getIndexOfOutputBuffer(value);
         if (outputIndex.hasValue())
           return $_op.getNumInputs() + outputIndex.getValue();
         Optional<unsigned> initTensorIndex = getIndexOfInitTensor(value);
         if (initTensorIndex.hasValue())
           return $_op.getNumInputs() + $_op.getNumOutputBuffers() + initTensorIndex.getValue();
         return llvm::None;
       }]
     >,
 
     //===------------------------------------------------------------------===//
     // Other interface methods.
     //===------------------------------------------------------------------===//
     InterfaceMethod<
       /*desc=*/[{
         Return the iterator types attribute within the current operation.
       }],
       /*retTy=*/"ArrayAttr",
       /*methodName=*/"iterator_types",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return $_op.iterator_types();
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return the indexing maps attribute within the current operation.
       }],
       /*retTy=*/"ArrayAttr",
       /*methodName=*/"indexing_maps"
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return the indexing maps within the current operation.
       }],
       /*retTy=*/"SmallVector<AffineMap, 4>",
       /*methodName=*/"getIndexingMaps",

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

 }
 
 void LinalgDependenceGraph::addDependenceElem(DependenceType dt,
                                               LinalgOpView indexingOpView,
                                               LinalgOpView dependentOpView) {
   LLVM_DEBUG(dbgs() << "\nAdd dep type " << getDependenceTypeStr(dt) << ":\t"
                     << *indexingOpView.op << " -> " << *dependentOpView.op);
   dependencesFromGraphs[dt][indexingOpView.op].push_back(
-      LinalgDependenceGraphElem{dependentOpView, indexingOpView.view});
+      LinalgDependenceGraphElem{dependentOpView, indexingOpView});
   dependencesIntoGraphs[dt][dependentOpView.op].push_back(
-      LinalgDependenceGraphElem{indexingOpView, dependentOpView.view});
+      LinalgDependenceGraphElem{indexingOpView, dependentOpView});
 }
 
 LinalgDependenceGraph::dependence_range
 LinalgDependenceGraph::getDependencesFrom(
     LinalgOp src, LinalgDependenceGraph::DependenceType dt) const {
   return getDependencesFrom(src.getOperation(), dt);
 }
 
LinalgDependenceGraph::getDependencesInto(
     Operation *dst, LinalgDependenceGraph::DependenceType dt) const {
   auto iter = dependencesIntoGraphs[dt].find(dst);
   if (iter == dependencesIntoGraphs[dt].end())
     return llvm::make_range(nullptr, nullptr);
   return llvm::make_range(iter->second.begin(), iter->second.end());
 }
 
 void LinalgDependenceGraph::addDependencesBetween(LinalgOp src, LinalgOp dst) {
-  for (auto srcView : src.getOutputBuffers()) { // W
+  for (auto srcView : llvm::enumerate(src.getOutputBuffers())) { // W
     // RAW graph
-    for (auto dstView : dst.getInputBuffers()) { // R
-      if (aliases.alias(srcView, dstView)) { // if alias, fill RAW
-        addDependenceElem(DependenceType::RAW,
-                          LinalgOpView{src.getOperation(), srcView},
-                          LinalgOpView{dst.getOperation(), dstView});
+    for (auto dstView : llvm::enumerate(dst.getInputBuffers())) { // R
+      if (aliases.alias(srcView.value(),
+                        dstView.value())) { // if alias, fill RAW
+        addDependenceElem(
+            DependenceType::RAW,
+            LinalgOpView{
+                src.getOperation(),
+                static_cast<unsigned>(srcView.index() + src.getNumInputs())},
+            LinalgOpView{dst.getOperation(),
+                         static_cast<unsigned>(dstView.index())});
       }
     }
     // WAW graph
-    for (auto dstView : dst.getOutputBuffers()) { // W
-      if (aliases.alias(srcView, dstView)) {      // if alias, fill WAW
-        addDependenceElem(DependenceType::WAW,
-                          LinalgOpView{src.getOperation(), srcView},
-                          LinalgOpView{dst.getOperation(), dstView});
+    for (auto dstView : llvm::enumerate(dst.getOutputBuffers())) { // W
+      if (aliases.alias(srcView.value(),
+                        dstView.value())) { // if alias, fill WAW
+        addDependenceElem(
+            DependenceType::WAW,
+            LinalgOpView{
+                src.getOperation(),
+                static_cast<unsigned>(srcView.index() + src.getNumInputs())},
+            LinalgOpView{
+                dst.getOperation(),
+                static_cast<unsigned>(dstView.index() + dst.getNumInputs())});
       }
     }
   }
-  for (auto srcView : src.getInputBuffers()) { // R
+  for (auto srcView : llvm::enumerate(src.getInputBuffers())) { // R
     // RAR graph
-    for (auto dstView : dst.getInputBuffers()) { // R
-      if (aliases.alias(srcView, dstView)) { // if alias, fill RAR
+    for (auto dstView : llvm::enumerate(dst.getInputBuffers())) { // R
+      if (aliases.alias(srcView.value(),
+                        dstView.value())) { // if alias, fill RAR
         addDependenceElem(DependenceType::RAR,
-                          LinalgOpView{src.getOperation(), srcView},
-                          LinalgOpView{dst.getOperation(), dstView});
+                          LinalgOpView{src.getOperation(),
+                                       static_cast<unsigned>(srcView.index())},
+                          LinalgOpView{dst.getOperation(),
+                                       static_cast<unsigned>(dstView.index())});
       }
     }
     // WAR graph
-    for (auto dstView : dst.getOutputBuffers()) { // W
-      if (aliases.alias(srcView, dstView)) {      // if alias, fill WAR
-        addDependenceElem(DependenceType::WAR,
-                          LinalgOpView{src.getOperation(), srcView},
-                          LinalgOpView{dst.getOperation(), dstView});
+    for (auto dstView : llvm::enumerate(dst.getOutputBuffers())) { // W
+      if (aliases.alias(srcView.value(),
+                        dstView.value())) { // if alias, fill WAR
+        addDependenceElem(
+            DependenceType::WAR,
+            LinalgOpView{src.getOperation(),
+                         static_cast<unsigned>(srcView.index())},
+            LinalgOpView{
+                dst.getOperation(),
+                static_cast<unsigned>(dstView.index() + dst.getNumInputs())});
       }
     }
   }
 }
 
 SmallVector<Operation *, 8>
 LinalgDependenceGraph::findCoveringDependences(LinalgOp srcLinalgOp,
                                                LinalgOp dstLinalgOp) const {
LinalgDependenceGraph::findOperationsWithCoveringDependences(
   // to an aliasing view on a src -> op -> dst path.
   // TODO: we are not considering paths yet, just interleaved positions.
   for (auto dt : types) {
     for (auto dependence : getDependencesFrom(src, dt)) {
       auto interimPos = linalgOpPositions.lookup(dependence.dependentOpView.op);
       // Skip if not interleaved.
       if (interimPos >= dstPos || interimPos <= srcPos)
         continue;
-      if (view && !aliases.alias(view, dependence.indexingView))
+      linalg::LinalgOp consumer =
+          cast<linalg::LinalgOp>(dependence.indexingOpView.op);
+      Value consumerView =
+          consumer.getShapedOperand(dependence.indexingOpView.operandNum);
+      if (view && !aliases.alias(view, consumerView))
         continue;
       auto *op = dependence.dependentOpView.op;
       LLVM_DEBUG(dbgs() << "\n***Found covering dependence of type "
                         << getDependenceTypeStr(dt) << ": " << *src << " -> "
-                        << *op << " on " << dependence.indexingView);
+                        << *op << " on " << consumerView);
       res.push_back(op);
     }
   }
   return res;
 }
 
 bool LinalgDependenceGraph::hasDependenceFrom(
     LinalgOp srcLinalgOp, LinalgOp dstLinalgOp,

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
+#include <set>
+
 #define DEBUG_TYPE "linalg-fusion"
 
 using namespace mlir;
 using namespace mlir::edsc;
 using namespace mlir::edsc::intrinsics;
 using namespace mlir::linalg;
 
 using folded_std_constant_index = FoldedValueBuilder<ConstantIndexOp>;
static LinalgOp cloneWithLoopRanges(OpBuilder &b, Location loc, LinalgOp op,
   SmallVector<Value, 8> clonedShapes;
   clonedShapes.reserve(op.getNumShapedOperands());
 
   // Iterate over the shape operands in order.
   // Extract the subranges from the linearized ranges.
   for (auto en : llvm::enumerate(op.getShapedOperands())) {
     unsigned shapedOperandIdx = en.index();
     AffineMap map = op.getIndexingMap(shapedOperandIdx);
-    LLVM_DEBUG(dbgs() << "shapedOperandIdx: " << shapedOperandIdx
+    LLVM_DEBUG(llvm::dbgs() << "shapedOperandIdx: " << shapedOperandIdx
                             << " with indexingMap: " << map << "\n");
     SmallVector<Value, 4> offsets, sizes, strides;
     inferShapeComponents(map, loopRanges, offsets, sizes, strides);
     Value shape = en.value();
     Value sub = shape.getType().isa<MemRefType>()
                     ? b.create<SubViewOp>(loc, shape, offsets, sizes, strides)
                           .getResult()
                     : b.create<SubTensorOp>(loc, shape, offsets, sizes, strides)
                           .getResult();
static ShapeDimension getShapeDefiningLoopRange(LinalgOp op,
                                                 unsigned loopDepth) {
   auto maps = op.indexing_maps();
   // Iterate over the inputs and outputs in order.
   // Extract the subranges from the linearized ranges.
   SmallVector<Value, 8> ios(op.getInputsAndOutputBuffers());
   for (auto en : llvm::enumerate(ios)) {
     unsigned idx = en.index();
     auto map = maps[idx].cast<AffineMapAttr>().getValue();
-    LLVM_DEBUG(dbgs() << "getShapeDefiningLoopRange I/O idx: " << idx << "\n");
-    LLVM_DEBUG(dbgs() << "getShapeDefiningLoopRange map: " << map << "\n");
+    LLVM_DEBUG(llvm::dbgs()
+               << "getShapeDefiningLoopRange I/O idx: " << idx << "\n");
+    LLVM_DEBUG(llvm::dbgs()
+               << "getShapeDefiningLoopRange map: " << map << "\n");
     Value shape = en.value();
     SmallVector<Value, 8> shapeRanges(map.getNumResults(), nullptr);
     for (auto en2 : llvm::enumerate(map.getResults())) {
       if (loopDepth == en2.value().cast<AffineDimExpr>().getPosition()) {
-        LLVM_DEBUG(dbgs() << "getShapeDefiningLoopRange loopDepth: "
+        LLVM_DEBUG(llvm::dbgs() << "getShapeDefiningLoopRange loopDepth: "
                                 << loopDepth << "\n");
-        LLVM_DEBUG(dbgs() << "getShapeDefiningLoopRange shape: " << shape
-                          << "\n");
+        LLVM_DEBUG(llvm::dbgs()
+                   << "getShapeDefiningLoopRange shape: " << shape << "\n");
         return ShapeDimension{shape, static_cast<unsigned>(en2.index())};
       }
     }
   }
   llvm_unreachable("Expect to be able to extract a shape defining loop range");
 }
 
 /// Fuses the producer of `producerIdx` into the loop immediately enclosing
assert((isa<SubViewOp>(shapeProducingOp) ||
           isa<SubTensorOp>(shapeProducingOp)) &&
          "SubviewOp or SubTensorOp expected");
 
   // loopToOperandRangesMaps are permutations-only by construction:
   //   we can always identify a data dimension with a (at least one) loop
   //   dimension.
   // TODO: extend this with range inference.
   AffineMap producerMap = producer.getOutputIndexingMap(producerIdx);
-  LLVM_DEBUG(dbgs() << "Producer Idx: " << producerIdx
+  LLVM_DEBUG(llvm::dbgs() << "Producer Idx: " << producerIdx
                           << ", producer map: " << producerMap << "\n");
 
   unsigned nPar = producer.getNumParallelLoops();
   unsigned nRed = producer.getNumReductionLoops();
   unsigned nWin = producer.getNumWindowLoops();
   SmallVector<Range, 8> loopRanges(nPar + nRed + nWin);
 
   // Iterate over dimensions identified by the producer map for `producerIdx`.
   // This defines a subset of the loop ranges that we need to complete later.

 // Some of these will be lifted in the future with better analysis.
 static bool isStructurallyFusableProducer(LinalgOp producer, Value consumedView,
                                           LinalgOp consumer) {
   assert(producer.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   assert(consumer.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   if (producer.getNumOutputs() != 1) {
-    LLVM_DEBUG(dbgs() << "\nNot structurally fusable (multi-output)");
+    LLVM_DEBUG(llvm::dbgs() << "\nNot structurally fusable (multi-output)");
     return false;
   }
   // Only fuse when the producer block dominates.
   DominanceInfo dom(producer.getOperation());
   if (!dom.dominates(producer.getOperation()->getBlock(),
                      consumer.getOperation()->getBlock())) {
     LLVM_DEBUG(
-        dbgs()
+        llvm::dbgs()
         << "\nNot structurally fusable (producer block does not dominate)");
     return false;
   }
   return true;
 }
 
 bool mlir::linalg::isProducerLastWriteOfView(const LinalgDependenceGraph &graph,
                                              LinalgOp consumer,
                                              Value consumedView,
                                              LinalgOp producer) {
   assert(producer.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   assert(consumer.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   // Make some simple structural checks that alleviate the need for more
   // complex analyses.
   if (!isStructurallyFusableProducer(producer, consumedView, consumer)) {
-    LLVM_DEBUG(dbgs() << "\n***Not static last write due to structure:\t"
+    LLVM_DEBUG(llvm::dbgs() << "\n***Not static last write due to structure:\t"
                             << *producer.getOperation());
     return false;
   }
   // Check for any interleaved write to consumedView.
   if (!graph.findCoveringWrites(producer, consumer, consumedView).empty()) {
-    LLVM_DEBUG(dbgs() << "\n***Not fusable due to interleaved write:\t"
+    LLVM_DEBUG(llvm::dbgs() << "\n***Not fusable due to interleaved write:\t"
                             << *producer.getOperation());
     return false;
   }
   return true;
 }
 
 bool mlir::linalg::isFusableInto(const LinalgDependenceGraph &graph,
                                  LinalgOp consumer, Value consumedView,
                                  LinalgOp producer) {
   assert(producer.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   assert(consumer.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   if (!isProducerLastWriteOfView(graph, consumer, consumedView, producer))
     return false;
   // Check for any fusion-preventing dependence to any shape read/written that
   // would violate dependences.
   if (!graph.findCoveringDependences(producer, consumer).empty()) {
-    LLVM_DEBUG(dbgs() << "\n***Not fusable due to an interleaved dependence:\t"
-                      << *producer.getOperation());
+    LLVM_DEBUG(llvm::dbgs()
+               << "\n***Not fusable due to an interleaved dependence:\t"
     return false;
   }
   if (auto convOp = dyn_cast<linalg::ConvOp>(producer.getOperation())) {
     // TODO: add a level of indirection to linalg.generic.
     if (convOp.padding())
       return false;
   }
   if (auto convOp = dyn_cast<linalg::ConvOp>(consumer.getOperation())) {

 static Optional<LinalgDependenceGraph::LinalgDependenceGraphElem>
 findFusableProducer(LinalgOp consumer, unsigned consumerIdx,
                     const LinalgDependenceGraph &dependenceGraph) {
   // Only consider RAW and WAW atm.
   for (auto depType : {
            LinalgDependenceGraph::DependenceType::RAW,
            LinalgDependenceGraph::DependenceType::WAW,
        }) {
-    for (auto dependence :
-         dependenceGraph.getDependencesInto(consumer, depType)) {
+    for (auto dependence : llvm::make_filter_range(
+             dependenceGraph.getDependencesInto(consumer, depType),
+             [&consumerIdx](
+                 LinalgDependenceGraph::LinalgDependenceGraphElem elem) {
+               return elem.indexingOpView.operandNum == consumerIdx;
+             })) {
       auto producer = cast<LinalgOp>(dependence.dependentOpView.op);
 
       // Check that the dependence is indeed on the input `consumerIdx` view.
-      auto consumedView = dependence.indexingView;
+      auto consumedView =
+          consumer.getBuffer(dependence.indexingOpView.operandNum);
       if (!isSameSubView(consumer.getBuffer(consumerIdx), consumedView))
         continue;
 
       // Consumer consumes this view, `isStructurallyFusableProducer` also
       // checks whether it is a strict subview of the producer view.
-      auto producedView = dependence.dependentOpView.view;
-      auto producerIdx =
-          producer.getIndexOfOutputBuffer(producedView).getValue();
-      // `consumerIdx` and `producerIdx` exist by construction.
-      LLVM_DEBUG(dbgs() << "\n"
-                        << LinalgDependenceGraph::getDependenceTypeStr(depType)
-                        << "producer: " << *producer.getOperation() << " view: "
-                        << producedView << " output index: " << producerIdx);
-      (void)producerIdx;
+      auto producedView =
+          producer.getBuffer(dependence.dependentOpView.operandNum);
+      LLVM_DEBUG(llvm::dbgs()
+                 << "\n"
+                 << "producer: " << *producer.getOperation()
+                 << " view: " << producedView << " output index: "
+                 << dependence.dependentOpView.operandNum -
+                        producer.getNumInputs()
+                 << "\n");
 
       // Simple fusability checks.
       if (!isFusableInto(dependenceGraph, consumer, consumedView, producer))
         continue;
 
       return dependence;
     }
   }
   return {};
 }
 
 Optional<FusionInfo> mlir::linalg::fuseProducerOfBuffer(
     OpBuilder &b, LinalgOp consumer, unsigned consumerIdx,
     const LinalgDependenceGraph &graph, OperationFolder *folder) {
   Optional<LinalgDependenceGraph::LinalgDependenceGraphElem> fusableDependence =
       findFusableProducer(consumer, consumerIdx, graph);
   if (!fusableDependence)
     return {};
 
   LinalgOp producerOp = cast<LinalgOp>(fusableDependence->dependentOpView.op);
-  Value producerView = fusableDependence->dependentOpView.view;
-  Value consumerView = fusableDependence->indexingView;
+  unsigned producerIdx =
+      fusableDependence->dependentOpView.operandNum - producerOp.getNumInputs();
+  Value consumerView = consumer.getShapedOperand(consumerIdx);
 
   // Must be a subview or a slice to guarantee there are loops we can fuse
   // into.
   auto subView = consumerView.getDefiningOp<SubViewOp>();
   auto slice = consumerView.getDefiningOp<SliceOp>();
   if (!subView && !slice) {
-    LLVM_DEBUG(dbgs() << "\nNot fusable (not a subview or slice)");
+    LLVM_DEBUG(llvm::dbgs() << "\nNot fusable (not a subview or slice)");
     return {};
   }
 
   // Fuse `producer` just before `consumer`.
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPoint(consumer.getOperation());
   ScopedContext scope(b, consumer.getLoc());
-  LLVM_DEBUG(dbgs() << "Fuse into consumer: " << *consumer << "\n");
-  Optional<unsigned> producerIdxOpt =
-      producerOp.getIndexOfOutputBuffer(producerView);
-  assert(producerIdxOpt.hasValue() && "incorrect operand index");
-  unsigned producerIdx = producerIdxOpt.getValue();
+  LLVM_DEBUG(llvm::dbgs() << "Fuse into consumer: " << *consumer << "\n");
 
   auto fusedProducer =
       fuse(b, producerOp, producerIdx, consumer, consumerIdx, folder);
   return FusionInfo{producerOp, fusedProducer};
 }
 
 /// Walk back use-def chain through scf::For yields.
 /// Sets `producer` and `outputIndex` if it finds a producer LinalgOp
mlir::linalg::fuseProducerOfTensor(OpBuilder &b, LinalgOp consumer,
   Value inputTensor = consumer.getInput(consumerIdx);
   LinalgOp producerOp;
   unsigned producerIdx;
   getProducerOfTensor(inputTensor, producerOp, producerIdx);
 
   // Must be a subtensor to guarantee there are loops we can fuse into.
   auto subTensor = inputTensor.getDefiningOp<SubTensorOp>();
   if (!subTensor || !producerOp) {
-    LLVM_DEBUG(dbgs() << "\nNot fusable (not a subtensor)");
+    LLVM_DEBUG(llvm::dbgs() << "\nNot fusable (not a subtensor)");
     return {};
   }
 
   // Insert fused `producer` just before `consumer`.
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPoint(consumer.getOperation());
   ScopedContext scope(b, consumer.getLoc());
-  LLVM_DEBUG(dbgs() << "Fuse into consumer: " << *consumer << "\n");
+  LLVM_DEBUG(llvm::dbgs() << "Fuse into consumer: " << *consumer << "\n");
   LinalgOp fusedProducer =
       fuse(b, producerOp, producerIdx, consumer, consumerIdx, folder);
 
   // Replace use.
   // Canonicalizations are not guaranteed to have happened before constructing
   // `fusedProducer`. In the tensor case this can result in temporary type
   // mismatches. Insert a `tensor_cast` op to propagate the transformation
   // invariant that types are compatible.
   Value def = fusedProducer.getOperation()->getResult(producerIdx);
   OpOperand &use = consumer.getOperation()->getOpOperand(consumerIdx);
   Type consumerType = use.get().getType();
   if (consumerType != def.getType())
     def = b.create<TensorCastOp>(fusedProducer.getLoc(), consumerType, def);
   use.set(def);
   return FusionInfo{producerOp, fusedProducer};
 }
 
+/// Given a map that is a projected permutation, remap dimensions according to
+/// the `remapping` provided. `remapping` is a map from original dimensions to
+/// new dimensions. Any of the original dimensions not in `remapping` are
+/// pruned.
+static AffineMap
+getSubMapBasedOnDims(AffineMap map,
+                     const DenseMap<unsigned, unsigned> &remapping) {
+  assert(map.isProjectedPermutation());
+  SmallVector<AffineExpr, 2> exprs;
+  for (AffineExpr e : map.getResults()) {
+    auto it = remapping.find(e.cast<AffineDimExpr>().getPosition());
+    if (it != remapping.end())
+      exprs.push_back(getAffineDimExpr(it->second, map.getContext()));
+  }
+  return AffineMap::get(remapping.size(), 0, exprs, map.getContext());
+}
+
+/// Prune all dimensions that are of reduction iterator type from `map`.
+static AffineMap pruneReductionDimsFromMap(ArrayRef<Attribute> iteratorTypes,
+                                           AffineMap map) {
+  DenseMap<unsigned, unsigned> remapping;
+  unsigned prunedDims = 0;
+  for (auto attr : llvm::enumerate(iteratorTypes)) {
+    if (!isParallelIterator(attr.value()))
+      continue;
+    remapping[attr.index()] = prunedDims++;
+  }
+  return getSubMapBasedOnDims(map, remapping);
+}
+
+using FusableOpDependencesTy = llvm::MapVector<
+    Operation *,
+    SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 1>>;
+
 /// Returns the positions of the loop in `op` that can be tiled based on the
 /// operations that are to be fused with it. For example, in a
 ///
 ///   linalg.matmul ins(%a, %b : ...) outs(%c : ...)
 ///
 /// if the producer of %a needs to be fused with this op, only the `i` loop of
 /// the matmul can be tiled while fusing. If producer of %a, and %b are to be
-/// fused, then no loops can be tiled while fusing.
-static DenseSet<unsigned> collectTileAndFuseLoops(
-    LinalgOp op, ArrayRef<LinalgDependenceGraph::LinalgDependenceGraphElem>
-                     fusableDependences) {
-  // 1. Only parallel loops can be used for tile + fuse. Find the number of
-  // common outer parallel loops between the op and its producers being fused.
+/// fused, then no loops can be tiled while fusing. The conditions used are:
+/// 1. Only parallel loops can be used for tile + fuse. Find the number of
+///    common outer parallel loops between the op and its producers being fused.
+/// 2. Of the parallel loops only some can be fused. Only those loops can be
+///    fused such where the fused loop iteration space only touches one tile of
+///    the fused operation. This is because the producer (which is writing the
+///    fused subview) has update semantics. To compute this,
+///    a. Find the mapping from iterations in the consumer that write to the
+///       same location as the iterations in the producer. To do so use
+///       - indexing map of the fused view in the consumer : consumerIndexMap
+///       - indexing map of the fused view in the producer : producerIndexMap
+///       consumerLoopToProducerLoop =
+///         inverse(producerIndexMap).compose(consumerIndexMap)
+///
+/// Since an inverse computation is needed, we need to consider the projection
+/// of the producerIndexMap w.r.t the parallel loops.  The actual fused loops
+/// are the dimensions of the consumerLoopToProducerLoop map that correspond to
+/// parallel loops and appear in the result of the map
+///
+/// Example 1:
+///   linalg.fill(%c, %cst)
+///   linalg.matmul ins(%a, %b) outs(%c)
+///     Number of parallel loops : 2
+///     producerIndexMap = affine_map<(i, j) ->(i , j)>
+///     consumerIndexMap = affine_map<(i, j, k) -> (i, j)>
+///     consumerLoopToProducerLoop = affine_map<(i, j, k) -> (i, j)>
+///     Fused dimensions : i, j
+///
+/// Example 2:
+///   linalg.matmul ins(%a, %b) outs(%c)
+///   linalg.generic {indexing_maps = [affine_map<(i, j) -> (j, i)>, ...
+///                   iterator_types = ["parallel", "parallel"]}
+///     ins(%c) ...
+///
+///     Number of parallel loops = 2:
+///     producerIndexMap (projected to parallel loops) =
+///       affine_map<(i, j) -> (i, j)>
+///     consumerLoopToProducerLoop2 = affine_map<(i, j) -> (j, i)>
+///     Fused dimensions : i, j
+///
+/// Example 3:
+///   linalg.copy(%s, %b)
+///   linalg.matmul ins(%a, %b) outs(%c)
+///
+///   Number of parallel loops = 2
+///   produceIndexMap : affine_map<(i, j) -> (i, j)>
+///   consumerLoopToProduceLoops = affine_map<(i, j, k) -> (k, j)>
+///     submap with only parallel loops = affine_map<(i, j) -> (j)>
+///   Fused dimensions : j
+static std::set<unsigned>
+collectTileAndFuseLoops(LinalgOp op,
+                        const FusableOpDependencesTy &fusableDependences) {
   auto getNumOuterParallelLoops = [](LinalgOp linalgOp) {
     return linalgOp.iterator_types()
         .getValue()
         .take_while([](Attribute attr) -> bool {
           return attr.cast<StringAttr>().getValue() ==
                  getParallelIteratorTypeName();
         })
         .size();
   };
 
+  LLVM_DEBUG({
+    llvm::dbgs() << "Op : ";
+    op.getOperation()->print(llvm::dbgs(), OpPrintingFlags().useLocalScope());
+    llvm::dbgs() << "\n";
+  });
+
   size_t numOuterParallelLoops = getNumOuterParallelLoops(op);
   for (auto dependence : fusableDependences) {
+    linalg::LinalgOp producer = cast<linalg::LinalgOp>(dependence.first);
     numOuterParallelLoops =
-        std::min(numOuterParallelLoops, getNumOuterParallelLoops(cast<LinalgOp>(
-                                            dependence.dependentOpView.op)));
+        std::min(numOuterParallelLoops, getNumOuterParallelLoops(producer));
   }
 
-  // Need to compute what tiled loops can be "fused". Given the precondition
-  // that all indexing map for the producer view is a projected permutation, we
-  // can assert that the producer iterates over the dimensions of the "fused
-  // view" only once. To be used a fused loop the producer should use this loop
-  // to access the fused view. For example, consider
-  //
-  // ```
-  //   linalg.add ins(%a, %b) outs(%c)
-  //   linalg.matmul ins(%d, %c) outs(%e)
-  // ```
-  //
-  // if `linalg.add` has the semantics of `c = a + b`, then the following
-  // tile+fuse code is correct.
-  //
-  // ```
-  // for j ... += TSj
-  //   %sa = subview %a[0, %j][...]
-  //   %sb = subview %b[0, %j][...]
-  //   %sc = subview %c[0, %j][...]
-  //   %sd = subview %d[0, 0][...]
-  //   %se = subview %e[0, %j][...]
-  //   linalg.add ins(%sa, %sb) outs(%sc)
-  //   linalg.matmul ins(%sd, %sc) outs(%se)
-  // ```
-  //
-  // On the other hand tiling along i would be incorrect
-  //
-  // ```
-  // for %i .. += TSi
-  //   %sa = subview %a[%i, 0][...]
-  //   %sb = subview %b[%i, 0][...]
-  //   %sc = subview %c[%i, 0][...]
-  //   %sc2 = subview %c[0, 0][...]
-  //   %sd = subview %d[%i, 0][...]
-  //   %se = subview %e[%i, 0][...]
-  //   linalg.add ins(%sa, %sb) outs(%sc)
-  //   linalg.matmul ins(%sd, %sc2) outs(%se)
-  // ```
-  //
-  // The write to the subview `%sc` in `linalg.add` is performed after the read
-  // from it using `%sc2` violating the RAW dependence of the original code. To
-  // find such loops indexing map of the fused view in the consumer op is
-  // used. For the above example, this indexing map is
-  //
-  //   affine_map<(d0, d1, d2) -> (d2, d1)>
-  //
-  // Since d0 is not in the result expressions of this map, it is not treated as
-  // tile + fuse loop, (but d1 is).
-  //
-  // TODO: The above is probably restrictive and there might be a generalization
-  // of these that might allow for more fusion opportunities. Explore based on
-  // needs.
-  SmallVector<DenseSet<unsigned>, 1> commonTilableLoops;
+  std::set<unsigned> fusableLoops;
+  auto range = llvm::seq<unsigned>(0, numOuterParallelLoops);
+  fusableLoops.insert(range.begin(), range.end());
   for (auto dependence : fusableDependences) {
-    unsigned consumerIdx =
-        op.getIndexOfShapedOperand(dependence.indexingView).getValue();
-    AffineMap consumerAccess = op.getIndexingMap(consumerIdx);
-    // Previously asserted that the consumerAccess map is a projected
-    // permutation, so all results are known to be AffineDimExprs. To remove
-    // this restriction walk the expression to find which dimensions of the
-    // consumer loop appear in the `consumerAccess`.
-    DenseSet<unsigned> positions;
-    for (auto expr : consumerAccess.getResults())
-      positions.insert(expr.cast<AffineDimExpr>().getPosition());
-    commonTilableLoops.emplace_back(std::move(positions));
-  }
-
-  // 2. Of the outer parallel loops, only those loops can be tiled + fused as
-  // computed above for all the fused dependences can be used to tile and fuse.
-  DenseSet<unsigned> tilableParallelLoops;
-  for (auto index : llvm::seq<unsigned>(0, numOuterParallelLoops)) {
-    if (llvm::all_of(commonTilableLoops,
-                     [&](const DenseSet<unsigned> &tilableLoops) {
-                       return tilableLoops.count(index);
-                     }))
-      tilableParallelLoops.insert(index);
+    LLVM_DEBUG({
+      llvm::dbgs() << "\t fusable :";
+      for (unsigned i : fusableLoops)
+        llvm::dbgs() << " " << i;
+      llvm::dbgs() << "\n";
+    });
+    linalg::LinalgOp producer = cast<linalg::LinalgOp>(dependence.first);
+
+    assert(!dependence.second.empty() &&
+           "unexpected producer but not dependences");
+    AffineMap producerIndexingMap = producer.getIndexingMap(
+        dependence.second.front().dependentOpView.operandNum);
+    AffineMap prunedProducerIndexingMap = pruneReductionDimsFromMap(
+        producer.iterator_types().getValue(), producerIndexingMap);
+    if (!prunedProducerIndexingMap.isPermutation())
+      return {};
+
+    AffineMap consumerIndexingMap =
+        op.getIndexingMap(dependence.second.front().indexingOpView.operandNum);
+    if (consumerIndexingMap.getNumResults() !=
+        prunedProducerIndexingMap.getNumResults())
+      return {};
+
+    LLVM_DEBUG({
+      llvm::dbgs() << "\t producerMap : ";
+      producerIndexingMap.print(llvm::dbgs());
+      llvm::dbgs() << "  pruned : ";
+      prunedProducerIndexingMap.print(llvm::dbgs());
+      llvm::dbgs() << "\n";
+      llvm::dbgs() << "\t consumerMap : ";
+      consumerIndexingMap.print(llvm::dbgs());
+      llvm::dbgs() << "\n";
+    });
+
+    AffineMap invProducerIndexMap =
+        inversePermutation(prunedProducerIndexingMap);
+    if (!invProducerIndexMap)
+      return {};
+
+    AffineMap consumerLoopToProducerLoop =
+        invProducerIndexMap.compose(consumerIndexingMap);
+    std::set<unsigned> candidates;
+    for (AffineExpr expr : consumerLoopToProducerLoop.getResults()) {
+      unsigned position = expr.cast<AffineDimExpr>().getPosition();
+      if (fusableLoops.count(position))
+        candidates.insert(position);
+    }
+    LLVM_DEBUG({
+      llvm::dbgs() << "\t candidates :";
+      for (unsigned i : candidates)
+        llvm::dbgs() << " " << i;
+      llvm::dbgs() << "\n";
+    });
+    if (candidates.empty())
+      return {};
+    std::swap(candidates, fusableLoops);
   }
-  return tilableParallelLoops;
+
+  return fusableLoops;
 }
 
 /// Find all dependences that are to be fusable.
-static Optional<
-    SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 1>>
+static FusableOpDependencesTy
 findAllFusableDependences(LinalgOp op,
                           const LinalgDependenceGraph &dependenceGraph,
                           const LinalgFusionOptions &fusionOptions) {
-  SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 1>
-      fusableDependences;
+  FusableOpDependencesTy fusableDependences;
+  // TODO: Currently fusion would not be legal if the fusable dependence is to
+  // the same producer but different indexing map in the consumer. Fix this, but
+  // in the meanwhile disallow such a fusion.
+  DenseMap<Operation *, AffineMap> fusedProducerIndexingMap;
   for (auto operand : llvm::enumerate(op.getInputsAndOutputBuffers())) {
     if (fusionOptions.indicesToFuse &&
         !fusionOptions.indicesToFuse->count(operand.index()))
       continue;
     Optional<LinalgDependenceGraph::LinalgDependenceGraphElem>
         fusableDependence =
             findFusableProducer(op, operand.index(), dependenceGraph);
     if (!fusableDependence)
       continue;
+    LinalgOp producerOp = cast<LinalgOp>(fusableDependence->dependentOpView.op);
+    // Do not fuse dependences that are to operations not in the same basic
+    // block. This avoid moving fused operations across loops that might
+    // themselves carry dependency making the fusion illegal.
+    if (producerOp.getOperation()->getBlock() != op.getOperation()->getBlock())
+      continue;
     // Make sure that the indexing map of the view used for fusion in the
     // producer is a projected permutation.
-    LinalgOp producerOp = cast<LinalgOp>(fusableDependence->dependentOpView.op);
-    Value producerView = fusableDependence->dependentOpView.view;
-    unsigned producerIdx =
-        producerOp.getIndexOfOutputBuffer(producerView).getValue();
-    AffineMap producerMap = producerOp.getOutputIndexingMap(producerIdx);
+    unsigned producerIdx = fusableDependence->dependentOpView.operandNum;
+    AffineMap producerMap = producerOp.getIndexingMap(producerIdx);
     if (!producerMap.isProjectedPermutation()) {
-      op.emitError("unhandled non permutation indexing map for fused view in "
+      op.emitRemark("unhandled non permutation indexing map for fused view in "
                     "producer for operand at index ")
           << operand.index();
-      return llvm::None;
+      return FusableOpDependencesTy{};
     }
-    Value consumerView = fusableDependence->indexingView;
-    unsigned consumerIdx = op.getIndexOfShapedOperand(consumerView).getValue();
-    if (!op.getIndexingMap(consumerIdx).isProjectedPermutation()) {
-      op.emitError(
+
+    unsigned consumerIdx = fusableDependence->indexingOpView.operandNum;
+    AffineMap consumerMap = op.getIndexingMap(consumerIdx);
+    if (!consumerMap.isProjectedPermutation()) {
+      op.emitRemark(
           "unhandled case where indexing map for fused view in the consumer is "
           "not a projected permuration while fusing at index ")
           << operand.index();
-      return llvm::None;
+      return FusableOpDependencesTy{};
+    }
+
+    // Check if the producer is already a fusion candidate. Cannot fuse this
+    // dependence if it has a different indexing map when used in the consumer.
+    if (fusedProducerIndexingMap.count(producerOp.getOperation()) &&
+        fusedProducerIndexingMap[producerOp.getOperation()] != consumerMap) {
+      op.emitRemark("unhandled fusion to the same producer but with different "
+                    "indexing maps");
+      return FusableOpDependencesTy{};
     }
-    fusableDependences.push_back(*fusableDependence);
+    fusedProducerIndexingMap[producerOp.getOperation()] = consumerMap;
+
+    fusableDependences[producerOp.getOperation()].push_back(*fusableDependence);
     if (!fusionOptions.indicesToFuse)
       break;
   }
   return fusableDependences;
 }
 
 static bool isZero(Value v) {
   if (auto cst = v.getDefiningOp<ConstantIndexOp>())
if (!tilingOptions.interchangeVector.empty()) {
     return llvm::None;
   }
 
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(op);
   ScopedContext scope(rewriter, op.getLoc());
 
   // Find all the producers.
-  Optional<SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 1>>
-      fusableDependencesOpt =
-          findAllFusableDependences(op, dependenceGraph, fusionOptions);
-  if (!fusableDependencesOpt)
+  FusableOpDependencesTy fusableDependences =
+  if (fusableDependences.empty())
     return llvm::None;
-  ArrayRef<LinalgDependenceGraph::LinalgDependenceGraphElem> fusableDependences(
-      *fusableDependencesOpt);
 
   // Enforce the convention that "tiling by zero" skips tiling a particular
   // dimension. This convention is significantly simpler to handle instead of
   // adjusting affine maps to account for missing dimensions.
   auto nLoops = op.getNumLoops();
   SmallVector<Value, 4> tileSizeVector =
       tilingOptions.tileSizeComputationFunction(rewriter, op);
   if (tileSizeVector.size() < nLoops) {
     auto zero = std_constant_index(0);
     tileSizeVector.append(nLoops - tileSizeVector.size(), zero);
   }
 
   TiledAndFusedLinalgOps ret;
 
   // Find the loops that can be tiled and fused.
-  DenseSet<unsigned> tileFuseLoops =
+  std::set<unsigned> tileFuseLoops =
       collectTileAndFuseLoops(op, fusableDependences);
 
   // If there are no fusable dependences or there are no tile+fusable loops,
   // just return.
-  if (fusableDependences.empty() || tileFuseLoops.empty()) {
+  if (tileFuseLoops.empty()) {
     return llvm::None;
   }
 
   // Get the tile sizes for the first and second tiling steps. For the first
   // step the tile size are set to zero for the loops that arent
   // fused. Similarly for the second step, the tile sizes are set to zero for
   // the loops that are fused. For example, if for the following input
   //
   // ```
   //   linalg.add ins(%a, %b) outs(%c)
Optional<TiledLinalgOp> firstTiledOp =
       tileLinalgOp(rewriter, op, firstTilingOptions);
   if (!firstTiledOp)
     return llvm::None;
   ret.op = firstTiledOp->op;
   ret.fusedLoops.assign(firstTiledOp->loops.begin(), firstTiledOp->loops.end());
 
   rewriter.setInsertionPoint(ret.op);
   // Fuse the operands.
-  for (auto producer : enumerate(fusableDependences)) {
-    LinalgOp producerOp = cast<LinalgOp>(producer.value().dependentOpView.op);
-    unsigned producerIdx =
-        producerOp.getIndexOfOutputBuffer(producer.value().dependentOpView.view)
-            .getValue();
-    unsigned consumerIdx =
-        op.getIndexOfShapedOperand(producer.value().indexingView).getValue();
+  for (auto dependence : fusableDependences) {
+    LinalgOp producerOp = cast<LinalgOp>(dependence.first);
+    unsigned producerIdx = dependence.second.front().dependentOpView.operandNum;
+    unsigned consumerIdx = dependence.second.front().indexingOpView.operandNum;
     LinalgOp fusedOp =
-        fuse(rewriter, producerOp, producerIdx, ret.op, consumerIdx);
+        fuse(rewriter, producerOp, producerIdx - producerOp.getNumInputs(),
+             ret.op, consumerIdx);
     ret.fusedProducers.push_back(fusedOp);
     ret.originalProducers.push_back(producerOp);
   }
 
   if (!llvm::all_of(tileSizes, isZero)) {
     // Tile the remaining loops of the root operation.
     LinalgTilingOptions secondTilingOptions = tilingOptions;
     // The distribution is done only for the tile+fused loops.
case LinalgTilingLoopType::ParallelLoops:
     return tileAndFuseLinalgOpsImpl<scf::ParallelOp>(
         rewriter, op, dependenceGraph, tilingOptions, fusionOptions);
   default:;
   }
   return llvm::None;
 }
 
 static void fuseLinalgOpsGreedily(FuncOp f) {
-  LLVM_DEBUG(f.print(dbgs() << "\nBefore linalg-fusion: \n"));
+  LLVM_DEBUG(f.print(llvm::dbgs() << "\nBefore linalg-fusion: \n"));
 
   OpBuilder b(f);
   OperationFolder folder(f.getContext());
   DenseSet<Operation *> eraseSet;
 
   // Save original Linalg ops, we only want to make a pass over those.
   SmallVector<Operation *, 8> linalgOps;
   f.walk([&](LinalgOp op) {
static void fuseLinalgOpsGreedily(FuncOp f) {
   // The `fuseProducerOfBuffer` function performs structural checks and in
   // particular that no covering read or write exist between the consumer and
   // the producer. As a consequence, the only fusions that may occur preserve
   // subsequent dependences and are guaranteed by construction to produce the
   // whole view. We may thus erase the producer once it is fused.
   for (auto *e : eraseSet)
     e->erase();
 
-  LLVM_DEBUG(f.print(dbgs() << "\nAfter linalg-fusion: \n"));
+  LLVM_DEBUG(f.print(llvm::dbgs() << "\nAfter linalg-fusion: \n"));
 }
 
 namespace {
 struct LinalgFusionPass : public LinalgFusionBase<LinalgFusionPass> {
   void runOnFunction() override { fuseLinalgOpsGreedily(getFunction()); }
 };
 } // namespace
 
 std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgFusionPass() {
   return std::make_unique<LinalgFusionPass>();
 }

mlir/test/Dialect/Linalg/fusion-pattern.mlir

 // CHECK-SAME:           ins(%[[SV6]], %[[SV7]]
 // CHECK-SAME:             : memref<?x?xf32, #[[MAP1]]>, memref<?x?xf32, #[[MAP1]]>)
 // CHECK-SAME:           outs(%[[SV8]] : memref<?x?xf32, #[[MAP1]]>)
 //      CHECK:       }
 //      CHECK:     }
 //      CHECK:   }
 //      CHECK:   linalg.matmul
 // CHECK-SAME:     __internal_linalg_transform__ = "after_lhs_fusion_original"
+
+// -----
+
+module {
+  func @matmul_plus_matmul(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>,
+                           %arg2: memref<?x?xf32>) {
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    %0 = dim %arg2, %c0 : memref<?x?xf32>
+    %1 = dim %arg2, %c1 : memref<?x?xf32>
+    %2 = alloc(%0, %1) : memref<?x?xf32>
+    linalg.matmul ins(%arg0, %arg1 : memref<?x?xf32>, memref<?x?xf32>)
+      outs(%2 : memref<?x?xf32>)
+    linalg.generic
+      {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                        affine_map<(d0, d1) -> (d0, d1)>,
+                        affine_map<(d0, d1) -> (d0, d1)>],
+       iterator_types = ["parallel", "parallel"],
+       __internal_linalg_transform__ = "transpose_fusion"}
+      ins(%2, %2 : memref<?x?xf32>, memref<?x?xf32>)
+      outs(%arg2 : memref<?x?xf32>) {
+      ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32) :
+        %3 = addf %arg3, %arg4 : f32
+        linalg.yield %3 : f32
+      }
+    return
+  }
+}
+//       CHECK: func @matmul_plus_matmul
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: memref<?x?xf32>
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: memref<?x?xf32>
+//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: memref<?x?xf32>
+//       CHECK:   %[[T2:.+]] = alloc(%{{.*}}, %{{.*}}) : memref<?x?xf32>
+//       CHECK:   linalg.matmul
+//  CHECK-SAME:     after_transpose_fusion_original
+//       CHECK:   scf.parallel (%[[ARG3:[a-zA-Z0-9_]+]], %[[ARG4:.[a-zA-Z0-9_]+]])
+//       CHECK:     %[[T5:.+]] = subview %[[T2]][%[[ARG3]], %[[ARG4]]]
+//       CHECK:     %[[T6:.+]] = subview %[[ARG2]][%[[ARG3]], %[[ARG4]]]
+//       CHECK:     %[[T8:.+]] = subview %[[ARG0]][%[[ARG3]], 0]
+//       CHECK:     %[[T9:.+]] = subview %[[ARG1]][0, %[[ARG4]]]
+//       CHECK:     linalg.matmul
+//  CHECK-SAME:       after_transpose_fusion_producer
+//  CHECK-SAME:       ins(%[[T8]], %[[T9]]
+//  CHECK-SAME:       outs(%[[T5]]
+//   CHECK-NOT:     linalg.matmul
+//       CHECK:     linalg.generic
+//  CHECK-SAME:       ins(%[[T5]], %[[T5]]
+//  CHECK-SAME:       outs(%[[T6]]
+//  CHECK-SAME:       after_transpose_fusion
+
+// -----
+
+module {
+  func @matmul_plus_transpose_matmul(%arg0: memref<?x?xf32>,
+                                     %arg1: memref<?x?xf32>,
+                                     %arg2: memref<?x?xf32>) {
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    %0 = dim %arg2, %c0 : memref<?x?xf32>
+    %1 = dim %arg2, %c1 : memref<?x?xf32>
+    %2 = alloc(%0, %1) : memref<?x?xf32>
+    linalg.matmul ins(%arg0, %arg1 : memref<?x?xf32>, memref<?x?xf32>)
+      outs(%2 : memref<?x?xf32>)
+    linalg.generic
+      {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                        affine_map<(d0, d1) -> (d1, d0)>,
+                        affine_map<(d0, d1) -> (d0, d1)>],
+       iterator_types = ["parallel", "parallel"],
+       __internal_linalg_transform__ = "transpose_fusion"}
+      ins(%2, %2 : memref<?x?xf32>, memref<?x?xf32>)
+      outs(%arg2 : memref<?x?xf32>) {
+      ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32) :
+        %3 = addf %arg3, %arg4 : f32
+        linalg.yield %3 : f32
+      }
+    return
+  }
+}
+
+//       CHECK: func @matmul_plus_transpose_matmul
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: memref<?x?xf32>
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: memref<?x?xf32>
+//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: memref<?x?xf32>
+//       CHECK:   %[[T2:.+]] = alloc(%{{.*}}, %{{.*}}) : memref<?x?xf32>
+//       CHECK:   linalg.matmul
+//  CHECK-SAME:     ins(%[[ARG0]], %[[ARG1]]
+//  CHECK-SAME:     outs(%[[T2]]
+//   CHECK-NOT:   linalg.matmul
+//       CHECK:   linalg.generic
+//  CHECK-SAME:     ins(%[[T2]], %[[T2]]
+//  CHECK-SAME:     outs(%[[ARG2]]
+
+// -----
+
+#map0 = affine_map<(d0)[s0] -> (32, -d0 + s0)>
+#map1 = affine_map<(d0)[s0] -> (64, -d0 + s0)>
+#map2 = affine_map<(d0)[s0] -> (16, -d0 + s0)>
+#map3 = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
+module {
+  func @basic_no_fusion(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>,
+                        %arg2: memref<?x?xf32>) {
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    %c2 = constant 2 : index
+    %c32 = constant 32 : index
+    %c64 = constant 64 : index
+    %c16 = constant 16 : index
+    %cst = constant 0.000000e+00 : f32
+    linalg.fill(%arg2, %cst) : memref<?x?xf32>, f32
+    %0 = dim %arg0, %c0 : memref<?x?xf32>
+    %1 = dim %arg1, %c1 : memref<?x?xf32>
+    %2 = dim %arg0, %c1 : memref<?x?xf32>
+    scf.parallel (%arg3, %arg4) = (%c0, %c0) to (%0, %1) step (%c32, %c64) {
+      scf.for %arg5 = %c0 to %2 step %c16 {
+        %3 = affine.min #map0(%arg3)[%0]
+        %4 = affine.min #map1(%arg4)[%1]
+        %5 = affine.min #map2(%arg5)[%2]
+        %6 = subview %arg0[%arg3, %arg5] [%3, %5] [1, 1] : memref<?x?xf32> to memref<?x?xf32, #map3>
+        %7 = subview %arg1[%arg5, %arg4] [%5, %4] [1, 1] : memref<?x?xf32> to memref<?x?xf32, #map3>
+        %8 = subview %arg2[%arg3, %arg4] [%3, %4] [1, 1] : memref<?x?xf32> to memref<?x?xf32, #map3>
+        linalg.matmul {__internal_linalg_transform__ = "basic_fusion"}
+          ins(%6, %7 : memref<?x?xf32, #map3>, memref<?x?xf32, #map3>)
+          outs(%8 : memref<?x?xf32, #map3>)
+      }
+      scf.yield
+    }
+    return
+  }
+}
+// CHECK-LABEL: func @basic_no_fusion
+//       CHECK:   linalg.fill
+//       CHECK:   scf.parallel
+//       CHECK:     scf.for
+//   CHECK-NOT:       linalg.fill
+//       CHECK:       linalg.matmul

mlir/test/lib/Transforms/TestLinalgFusionTransforms.cpp

patterns.insert<LinalgTileAndFusePattern<MatmulOp>>(
       LinalgMarker(Identifier::get("two_operand_fusion", context),
                    Identifier::get("after_two_operand_fusion", context)),
       LinalgMarker(
           ArrayRef<Identifier>(),
           Identifier::get("after_two_operand_fusion_producer", context)),
       LinalgMarker(
           ArrayRef<Identifier>(),
           Identifier::get("after_two_operand_fusion_original", context)));
+
+  patterns.insert<LinalgTileAndFusePattern<GenericOp>>(
+      context, dependenceGraph,
+      LinalgTilingOptions().setTileSizes({32, 64}).setLoopType(
+          LinalgTilingLoopType::ParallelLoops),
+      LinalgFusionOptions().setIndicesToFuse({0, 1}),
+      LinalgMarker(Identifier::get("transpose_fusion", context),
+                   Identifier::get("after_transpose_fusion", context)),
+      LinalgMarker(ArrayRef<Identifier>(),
+                   Identifier::get("after_transpose_fusion_producer", context)),
+      LinalgMarker(
+          ArrayRef<Identifier>(),
+          Identifier::get("after_transpose_fusion_original", context)));
 }
 
 static void applyFusionPatterns(MLIRContext *context, FuncOp funcOp) {
   OwningRewritePatternList fusionPatterns;
   Aliases alias;
   LinalgDependenceGraph dependenceGraph =
       LinalgDependenceGraph::buildDependenceGraph(alias, funcOp);
   fillFusionPatterns(context, dependenceGraph, fusionPatterns);