[mlir][Linalg] Extend tile+fuse to work on Linalg operation on tensors.

[mravishankar]

Summary

Depends On D93077

[nicolasvasilache]

Looks good!

Making a note here that I locally had to do 2 things for fusion on tensors:

1. change some logic in getProducerOfTensor to properly handle scf.for iter_args:

static void getProducerOfTensor(Value tensor, LinalgOp &producer,
                                unsigned &outputIndex) {

  while (true) {
    if (!tensor.getType().isa<RankedTensorType>())
      return;
    if (auto linalgOp = tensor.getDefiningOp<LinalgOp>()) {
      producer = linalgOp;
      outputIndex = tensor.cast<OpResult>().getResultNumber();
      return;
    }
    if (auto subTensorOp = tensor.getDefiningOp<SubTensorOp>()) {
      tensor = subTensorOp.source();
      continue;
    }
    if (auto blockArg = tensor.dyn_cast<BlockArgument>()) {
      if (auto forOp = dyn_cast<scf::ForOp>(blockArg.getOwner()->getParentOp())) {
        tensor = *(forOp.getIterOperands().begin() +
                   (blockArg.getArgNumber() - 1)); // account for iv
        continue;
      }
    }
    return;
  }
}

2. fix an indexing error at the beginning of fuseProducerOfTensor

Optional<FusionInfo> mlir::linalg::fuseProducerOfTensor(OpBuilder &b,
                                                        LinalgOp consumer,
                                                        unsigned consumerIdx) {
  Value inputTensor =
      consumerIdx < consumer.getNumInputs()
          ? consumer.getInput(consumerIdx)
          : consumer.getInitTensor(consumer.getNumInputs() - consumerIdx);
  LinalgOp producerOp;
  unsigned producerIdx;
  getProducerOfTensor(inputTensor, producerOp, producerIdx);
  if (!producerOp)
    return llvm::None;

I don't have a good setting to flush these changes at this time but signal boosting them here so you have them in case you need it in the future.
If not I'll prob hit the case again in a few more weeks and will deal with it then.

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

741

Now would also be a good time to refactor op + idx into just OpOperand which is a more idiomatic Value that also tracks its operand index.
v.isa<BlockArgument>() || v.isa<OpOperand>() should always be true (unless I missed a case..).

[mravishankar]

Rebase

[mravishankar]

Addressing comments.

[mravishankar]

In D93086#2448323, @nicolasvasilache wrote:

Looks good!

Making a note here that I locally had to do 2 things for fusion on tensors:

1. change some logic in getProducerOfTensor to properly handle scf.for iter_args:

static void getProducerOfTensor(Value tensor, LinalgOp &producer,
                                unsigned &outputIndex) {

  while (true) {
    if (!tensor.getType().isa<RankedTensorType>())
      return;
    if (auto linalgOp = tensor.getDefiningOp<LinalgOp>()) {
      producer = linalgOp;
      outputIndex = tensor.cast<OpResult>().getResultNumber();
      return;
    }
    if (auto subTensorOp = tensor.getDefiningOp<SubTensorOp>()) {
      tensor = subTensorOp.source();
      continue;
    }
    if (auto blockArg = tensor.dyn_cast<BlockArgument>()) {
      if (auto forOp = dyn_cast<scf::ForOp>(blockArg.getOwner()->getParentOp())) {
        tensor = *(forOp.getIterOperands().begin() +
                   (blockArg.getArgNumber() - 1)); // account for iv
        continue;
      }
    }
    return;
  }
}

2. fix an indexing error at the beginning of fuseProducerOfTensor

Optional<FusionInfo> mlir::linalg::fuseProducerOfTensor(OpBuilder &b,
                                                        LinalgOp consumer,
                                                        unsigned consumerIdx) {
  Value inputTensor =
      consumerIdx < consumer.getNumInputs()
          ? consumer.getInput(consumerIdx)
          : consumer.getInitTensor(consumer.getNumInputs() - consumerIdx);
  LinalgOp producerOp;
  unsigned producerIdx;
  getProducerOfTensor(inputTensor, producerOp, producerIdx);
  if (!producerOp)
    return llvm::None;

I don't have a good setting to flush these changes at this time but signal boosting them here so you have them in case you need it in the future.
If not I'll prob hit the case again in a few more weeks and will deal with it then.

Added this to this change itself.

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

741

Good point. I will look into this since the dependent changes might need some time to resolve.

[mravishankar]

Rebase on chnages to linalg on tensors

[mravishankar]

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

741

The latest changes uses OpOperand * to track dependencies. Wasnt sure how to adapt that to acsount for OpResult as well. Couple of things that I was thinking off but wasnt sure if those make sense

1. PointerPair<OpOperand *, OpResult *> : but I am not sure it would be valid to take OpResult * These dont seemed to be "owned" by the underlying operation but are just wrappers around Value. This would actually be a nice interface if that was not the case.

2)std::pair<OpOperand *, OpResult> this would work, but it probably needs a reimplementation of LinalgDependenceGraphElem to have nice interface to get the operation and operand/result number.

Please let me know which would be better. Happy to do either.

[mravishankar]

Rebase

[hanchung]

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

353

s/For the now/For now

354–358

If this is for consumer with buffer semantics only, we should keep the assertion.

355

s/LinalgDependenceGraphElem/dependenceGraph

I think we typically use the variable name instead of the type. (it would confuse if there are two variables in the same type.)

742–743

I think this comment is for later statements, ie Value value = consumer.getOperation()->getOperand(consumerIdx); ....? If so, let's move it two lines below.

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

61

let's add three more spaces to align it if you intend to modify this...

[mravishankar]

Rebase

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

742–743

With the rebase, the comment is at the right place.

[mravishankar]

Address comments.

[hanchung]

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

357–360

If we are going to check consumerOp, the above code should be dyn_cast. Otherwise, it results in a segfault.

744–746

ditto, should be dyn_cast.

Another point here is, we do assertion for buffers, but return llvm::None for tensors. I think this should be consistent. Let's make them both assert or return llvm::None?

[mravishankar]

Rebase on top of change to LinalgDependenceGraph that accounts for
tracking dependence through result values.

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

private:
   /// Analyses the aliasing views between `src` and `dst` and inserts the proper
   /// dependences in the graph.
   void addDependencesBetween(LinalgOp src, LinalgOp dst);
 
   // Adds an new dependence unit in the proper graph.
   // Uses std::pair to keep operations and view together and avoid usage errors
   // related to src/dst and producer/consumer terminology in the context of
   // dependences.
-  void addDependenceElem(DependenceType dt, OpOperand *indexingOpView,
-                         OpOperand *dependentOpView);
+  void addDependenceElem(DependenceType dt,
+                         LinalgDependenceGraphElem::OpView indexingOpView,
+                         LinalgDependenceGraphElem::OpView dependentOpView);
 
   /// Implementation detail for findCoveringxxx.
   SmallVector<Operation *, 8>
   findOperationsWithCoveringDependences(LinalgOp srcLinalgOp,
                                         LinalgOp dstLinalgOp, Value view,
                                         ArrayRef<DependenceType> types) const;
 
   Aliases &aliases;
   SmallVector<LinalgOp, 8> linalgOps;
   DenseMap<Operation *, unsigned> linalgOpPositions;
 };
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_ANALYSIS_DEPENDENCEANALYSIS_H_

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

LinalgDependenceGraph::LinalgDependenceGraph(Aliases &aliases,
   }
   for (unsigned i = 0, e = ops.size(); i < e; ++i) {
     for (unsigned j = i + 1; j < e; ++j) {
       addDependencesBetween(ops[i], ops[j]);
     }
   }
 }
 
-void LinalgDependenceGraph::addDependenceElem(DependenceType dt,
-                                              OpOperand *indexingOpView,
-                                              OpOperand *dependentOpView) {
+void LinalgDependenceGraph::addDependenceElem(
+    DependenceType dt, LinalgDependenceGraphElem::OpView indexingOpView,
+    LinalgDependenceGraphElem::OpView dependentOpView) {
   LLVM_DEBUG(dbgs() << "\nAdd dep type " << getDependenceTypeStr(dt) << ":\t ("
-                    << indexingOpView->get() << " @"
-                    << indexingOpView->getOperandNumber() << ") -> \n\t\t("
-                    << dependentOpView->get() << " @"
-                    << dependentOpView->getOperandNumber() << ")");
-  dependencesFromGraphs[dt][indexingOpView->getOwner()].push_back(
+                    << LinalgDependenceGraphElem::getValue(indexingOpView)
+                    << " @) -> \n\t\t("
+                    << LinalgDependenceGraphElem::getValue(dependentOpView)
+                    << " @)");
+  dependencesFromGraphs[dt][LinalgDependenceGraphElem::getOwner(indexingOpView)]
+      .push_back(
           LinalgDependenceGraphElem{dependentOpView, indexingOpView, dt});
-  dependencesIntoGraphs[dt][dependentOpView->getOwner()].push_back(
-      LinalgDependenceGraphElem{indexingOpView, dependentOpView, dt});
+  dependencesIntoGraphs[dt]
+                       [LinalgDependenceGraphElem::getOwner(dependentOpView)]
+                           .push_back(LinalgDependenceGraphElem{
+                               indexingOpView, dependentOpView, dt});
 }
 
 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) {
+  if (src.hasTensorSemantics() && dst.hasTensorSemantics()) {
+    for (OpOperand &dstOpOperand : dst.getInputOpOperands()) {
+      // Check if the operand is defined by the src.
+      auto definingOp = dstOpOperand.get().getDefiningOp<LinalgOp>();
+      if (definingOp && definingOp == src)
+        addDependenceElem(DependenceType::RAW, dstOpOperand.get(),
+                          &dstOpOperand);
+    }
+    return;
+  }
+  assert(src.hasBufferSemantics() && dst.hasBufferSemantics() &&
+         "unhandled dependence tracking for mixed buffer/tensor operations");
   for (OpOperand *srcOpOperand : src.getOutputBuffersOpOperands()) { // W
     // RAW graph
     for (OpOperand *dstOpOperand : dst.getInputBuffersOpOperands()) // R
       if (aliases.alias(srcOpOperand->get(), dstOpOperand->get()))  // RAW alias
         addDependenceElem(DependenceType::RAW, srcOpOperand, dstOpOperand);
     // WAW graph
     for (OpOperand *dstOpOperand : dst.getOutputBuffersOpOperands()) // W
       if (aliases.alias(srcOpOperand->get(), dstOpOperand->get())) // WAW alias

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

 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
 #include <set>
 
 #define DEBUG_TYPE "linalg-fusion"
 
bool mlir::linalg::isFusableInto(const LinalgDependenceGraph &graph,
   if (auto convOp = dyn_cast<linalg::ConvOp>(consumer.getOperation())) {
     // TODO: add a level of indirection to linalg.generic.
     if (convOp.padding())
       return false;
   }
   return true;
 }
 
+/// For `consumer` with buffer semantics, find the Linalg operation on buffers
+/// that is the last writer of `consumerOpOperand`. For now the fusable
+/// dependence is returned as an instance of the `dependenceGraph`.
 static Optional<LinalgDependenceGraph::LinalgDependenceGraphElem>
 findFusableProducer(OpOperand &consumerOpOperand,
                     const LinalgDependenceGraph &dependenceGraph) {
-  LinalgOp consumerOp = cast<LinalgOp>(consumerOpOperand.getOwner());
-  // Note that buffer semantics implies that the dependence will only be from
-  // OpOperand -> OpOperand.
-  assert(consumerOp.hasBufferSemantics() && "revisit usage of shaped operand");
+  LinalgOp consumerOp = dyn_cast<LinalgOp>(consumerOpOperand.getOwner());
+  if (!consumerOp)
+    return {};
 
   // Only consider RAW and WAW atm.
   for (auto depType : {
            LinalgDependenceGraph::DependenceType::RAW,
            LinalgDependenceGraph::DependenceType::WAW,
        }) {
     for (auto dependence : llvm::make_filter_range(
              dependenceGraph.getDependencesInto(consumerOp, depType),
              [&](LinalgDependenceGraph::LinalgDependenceGraphElem elem) {
                Value v = elem.getIndexingValue();
                Optional<unsigned> operandNum =
                    elem.getIndexingOpViewOperandNum();
                return isa<LinalgOp>(elem.getDependentOp()) &&
                       v == consumerOpOperand.get() && operandNum &&
                       operandNum.getValue() ==
                           consumerOpOperand.getOperandNumber();
              })) {
       // Consumer consumes this view, `isStructurallyFusableProducer` also
       // checks whether it is a strict subview of the producer view.
       auto producer = cast<LinalgOp>(dependence.getDependentOp());
       LLVM_DEBUG(llvm::dbgs()
                  << "\n"
                  << LinalgDependenceGraph::getDependenceTypeStr(depType)
-                 << "producer: " << *dependence.getDependentOp() << " view: "
-                 << dependence.getDependentValue() << " output index: "
-                 << (dependence.getDependentOpViewOperandNum().getValue() -
-                     producer.getNumInputs())
-                 << "\n");
+                 << "producer: " << *dependence.getDependentOp()
+                 << " view: " << dependence.getDependentValue() << "\n");
 
-      // Simple fusability checks.
-      if (!isFusableInto(dependenceGraph, consumerOp, consumerOpOperand.get(),
-                         producer))
-        continue;
-
+      // If the producer and consumer have tensor semantics, the only dependence
+      // between them is through a RAW dependence and they are fusable by
+      // construction. For buffer semantics need additional checks.
+      if (producer.hasBufferSemantics() && consumerOp.hasBufferSemantics() &&
+          isFusableInto(dependenceGraph, consumerOp, consumerOpOperand.get(),
         return dependence;
+      if (producer.hasTensorSemantics() && consumerOp.hasTensorSemantics()) {
+        assert(dependence.dependenceType ==
+               LinalgDependenceGraph::DependenceType::RAW);
+        return dependence;
+      }
     }
   }
   return {};
 }
 
 Optional<FusionInfo>
 mlir::linalg::fuseProducerOfBuffer(OpBuilder &b, OpOperand &consumerOpOperand,
                                    const LinalgDependenceGraph &graph) {
LLVM_DEBUG(llvm::dbgs() << "Fuse into consumer: "
                           << *consumerOpOperand.getOwner() << "\n");
 
   auto fusedProducer = fuse(b, producerOp, *producerMap, consumerOpOperand);
   return FusionInfo{producerOp, fusedProducer};
 }
 
 /// Walk back use-def chain through scf::For yields.
 /// Sets `producer` and `outputIndex` if it finds a producer LinalgOp
+
+// TODO(ravishankarm, ntv): This can be moved into the dependence graphs
+// dependence tracking since the dependence tracking is similar to what is done
+// w.r.t to buffers.
 static void getProducerOfTensor(Value tensor, OpResult &opResult) {
   if (!tensor.getType().isa<RankedTensorType>())
     return;
 
   while (true) {
     LLVM_DEBUG(llvm::dbgs() << "\ngetProducerOfTensor: " << tensor);
     if (auto linalgOp = tensor.getDefiningOp<LinalgOp>()) {
       opResult = tensor.cast<OpResult>();
for (auto dependence : fusableDependences.lookup(op)) {
         return {};
       std::swap(candidates, fusableLoops);
     }
   }
 
   return fusableLoops;
 }
 
+// /// For `consumer` with tensor semantics, find the Linalg operation on
+// tensors
+// /// producer the operand at position `consumerIdx`. This is a simple use-def
+// /// chain using the SSA value, but returned as an element of the
+// /// `LinalgDependenceGraphElem` to use the same analysis for both tensors and
+// /// buffers.
+// static Optional<LinalgDependenceGraph::LinalgDependenceGraphElem>
+// findFusableProducerForTensorOp(OpOperand &consumerOpOperand) {
+//   // For now only looking for cases where the operand is produced by another
+//   // Linalg structured operation.
+//   LinalgOp consumer = cast<LinalgOp>(consumerOpOperand.getOwner());
+//   if (!consumer || !consumer.hasTensorSemantics())
+//     return llvm::None;
+//   unsigned consumerIdx = consumerOpOperand.getOperandNumber();
+//   Value value = consumerOpOperand.get();
+//   if (auto linalgOp = value.getDefiningOp<LinalgOp>()) {
+//     return LinalgDependenceGraph::LinalgDependenceGraphElem{
+//         &(linalgOp
+//               .getOutputOpOperands()[value.cast<OpResult>().getResultNumber()]),
+//         &(consumer.getInputOpOperands()[consumerIdx]),
+//         LinalgDependenceGraph::DependenceType::RAW};
+//   }
+//   return llvm::None;
+// }
+
+// static Optional<LinalgDependenceGraph::LinalgDependenceGraphElem>
+// findFusableProducer(OpOperand &consumerOpOperand,
+//                     const LinalgDependenceGraph &dependenceGraph) {
+//   LinalgOp consumer = cast<LinalgOp>(consumerOpOperand.getOwner());
+//   if (!consumer)
+//     return llvm::None;
+//   if (consumer.hasBufferSemantics())
+//     return findFusableProducerForBufferOp(consumerOpOperand,
+//     dependenceGraph);
+//   if (consumer.hasTensorSemantics())
+//     return findFusableProducerForTensorOp(consumerOpOperand);
+//   return llvm::None;
+// }
+
 /// Find all dependences that are fusable.
 FusableOpDependencesTy mlir::linalg::findAllFusableDependences(
     ArrayRef<LinalgOp> ops, const LinalgDependenceGraph &dependenceGraph) {
   FusableOpDependencesTy fusableDependences;
   DenseMap<Operation *, SmallVector<AffineMap, 1>> fusedProducerIndexingMap;
   for (LinalgOp op : reverse(ops)) {
     for (OpOperand &opOperand : op.getShapedOpOperands()) {
       Optional<LinalgDependenceGraph::LinalgDependenceGraphElem>
static Optional<TiledLinalgOp> tileRootOperation(
   return tileLinalgOp(builder, op, tileFusedLoopsOptions);
 }
 
 /// Fuse the operations in `fusionCandidates` with `tiledOp`. Latter is expected
 /// to be a tiled operation such that it is valid to fuse all operations in
 /// `fusionCandidates`, i.e. move the operation within the inter-tile loops of
 /// `tiledOp`.
 static SmallVector<LinalgOp, 1>
-fuseOperations(OpBuilder &builder, LinalgOp tiledOp,
+fuseOperations(OpBuilder &builder, LinalgOp rootOp, LinalgOp tiledOp,
                ArrayRef<LinalgOp> fusionCandidates,
                const FusableOpDependencesTy &fusableDependences,
                const std::set<unsigned> &fusedLoops) {
   OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPoint(tiledOp);
   DenseMap<unsigned, Range> fusedLoopsAndRanges;
   for (unsigned loop : fusedLoops) {
     ShapeDimension shapeDim = getShapeDefiningLoopRange(tiledOp, loop, true);
     fusedLoopsAndRanges[loop] = getRangeFromOperandShape(
         builder, tiledOp.getLoc(), shapeDim.shape, shapeDim.dimension);
   }
 
   SmallVector<LinalgOp, 1> fusedOps(fusionCandidates.size());
+  DenseMap<Operation *, LinalgOp> origOpToFusedOp;
+  origOpToFusedOp[rootOp.getOperation()] = tiledOp;
   for (auto candidate : enumerate(llvm::reverse(fusionCandidates))) {
-    LinalgOp fusedOp = fuse(builder, candidate.value(), fusedLoopsAndRanges);
+    LinalgOp origOp = candidate.value();
+    LinalgOp fusedOp = fuse(builder, origOp, fusedLoopsAndRanges);
+    origOpToFusedOp[origOp.getOperation()] = fusedOp;
     fusedOps[fusionCandidates.size() - candidate.index() - 1] = fusedOp;
+    // If the producer consumer operations are linalg operations on tensors, the
+    // dependence is due to value produced (as a return tensor) by the producer
+    // and used in the consumer. The returned value of the fused op needs to be
+    // made the operand of the tiled/fused consumer operation. By construction
+    // the value returned by the producer is the value used by the consumer.
+    for (auto &dependence : fusableDependences.lookup(origOp.getOperation())) {
+      if (origOp.hasTensorSemantics() &&
+          dependence.dependenceType ==
+              LinalgDependenceGraph::DependenceType::RAW) {
+        unsigned resultIndex =
+            dependence.getDependentOpViewResultNum().getValue();
+        LinalgOp consumer = origOpToFusedOp.lookup(dependence.getIndexingOp());
+        if (!consumer)
+          continue;
+        Value replacementValue = fusedOp.getOperation()->getResult(resultIndex);
+        consumer.getOperation()->setOperand(
+            dependence.getIndexingOpViewOperandNum().getValue(),
+            replacementValue);
+      }
+    }
     builder.setInsertionPoint(fusedOp);
   }
   return fusedOps;
 }
 
 template <typename LoopType>
 static Optional<TiledAndFusedLinalgOps>
 tileAndFuseLinalgOpsImpl(OpBuilder &builder, ArrayRef<LinalgOp> ops,
                          const LinalgDependenceGraph &dependenceGraph,
                          const LinalgTilingOptions &tilingOptions) {
   if (ops.size() < 2)
     return llvm::None;
   LinalgOp rootOp = ops.back();
-  for (auto op : enumerate(ops)) {
-    // TODO: Nothing in the fusion of sequence of ops is specific to
-    // buffers. This check can be removed after it is tested on tensors.
-    LinalgOp linalgOp = op.value();
-    if (!linalgOp.hasBufferSemantics()) {
-      linalgOp.emitRemark("tile and fuse only tested for buffer operation");
-      return llvm::None;
-    }
-  }
+  if (!llvm::all_of(
+          ops,
+          [](LinalgOp linalgOp) { return linalgOp.hasBufferSemantics(); }) &&
+      !llvm::all_of(ops, [](LinalgOp linalgOp) {
+        return linalgOp.hasTensorSemantics();
+      })) {
+    rootOp.emitError(
+        "unable to fuse operations that have tensor semantics with operations "
+        "that have buffer semantics and viceversa.");
   // TODO: Support interchange with tile + fuse. This might actually help do
   // better fusion.
   if (!tilingOptions.interchangeVector.empty()) {
     rootOp.emitRemark("unable to handle tile and fuse with interchange");
     return llvm::None;
   }
 
   OpBuilder::InsertionGuard guard(builder);
tileAndFuseLinalgOpsImpl(OpBuilder &builder, ArrayRef<LinalgOp> ops,
   if (!tiledRootOp) {
     rootOp.emitRemark("failed to tile the fused loops");
     return llvm::None;
   }
   ret.op = tiledRootOp->op;
   ret.fusedLoops.assign(tiledRootOp->loops.begin(), tiledRootOp->loops.end());
 
   // Fuse the other operations into the fused inter-tile loops produced above.
-  ret.fusedProducers = fuseOperations(builder, ret.op, ops.drop_back(),
+  ret.fusedProducers = fuseOperations(builder, rootOp, ret.op, ops.drop_back(),
                                       fusableDependences, ret.fusedLoopDims);
+
   return ret;
 }
 
 Optional<TiledAndFusedLinalgOps>
 mlir::linalg::tileAndFuseLinalgOps(OpBuilder &builder, ArrayRef<LinalgOp> ops,
                                    const LinalgDependenceGraph &dependenceGraph,
                                    const LinalgTilingOptions &tilingOptions) {
   switch (tilingOptions.loopType) {

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

 //       CHECK:     scf.yield
 //       CHECK:   }
 
 // -----
 
 module {
   func @sequence_of_matmul(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>,
                            %arg2: memref<?x?xf32>, %arg3: memref<?x?xf32>,
-         %arg4: memref<?x?xf32>) {
+                           %arg4: memref<?x?xf32>) {
     %cst = constant 0.000000e+00 : f32
     %c0 = constant 0 : index
     %c1 = constant 1 : index
     %m = dim %arg0, %c0 : memref<?x?xf32>
     %n1 = dim %arg1, %c1 : memref<?x?xf32>
     %n2 = dim %arg2, %c1 : memref<?x?xf32>
     %n3 = dim %arg3, %c1 : memref<?x?xf32>
     %0 = alloc(%m, %n1) : memref<?x?xf32>

 //  CHECK-SAME:        outs(%[[SV_ALLOC2]] : memref<?x?xf32, #[[MAP1]]>)
 //       CHECK:     linalg.fill(%[[SV_ARG4_2]], %{{.+}})
 //       CHECK:     linalg.matmul ins(%[[SV_ALLOC2]], %[[ARG3]]
 //  CHECK-SAME:        : memref<?x?xf32, #[[MAP1]]>, memref<?x?xf32>)
 //  CHECK-SAME:        outs(%[[SV_ARG4]] : memref<?x?xf32, #[[MAP1]]>)
 //       CHECK:     scf.yield
 //       CHECK:   }
 
+// -----
+
+module {
+  func @tensor_op_fusion(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
+                         %arg2: tensor<?x?xf32>, %arg3: tensor<?xf32>)
+    -> tensor<?x?xf32> {
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    %0 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+    %1 = dim %0, %c0 : tensor<?x?xf32>
+    %2 = dim %0, %c1 : tensor<?x?xf32>
+    %3 = linalg.init_tensor [%1, %2] : tensor<?x?xf32>
+    %4 = linalg.generic
+      {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                        affine_map<(d0, d1) -> (d0)>,
+                        affine_map<(d0, d1) -> (d0, d1)>],
+       iterator_types = ["parallel", "parallel"]}
+      ins(%0, %arg3 : tensor<?x?xf32>, tensor<?xf32>)
+      outs(%3 : tensor<?x?xf32>) {
+      ^bb0(%arg4: f32, %arg5: f32, %arg6: f32):
+        %5 = addf %arg4, %arg5 : f32
+        linalg.yield %5 : f32
+      } -> tensor<?x?xf32>
+    return %4 : tensor<?x?xf32>
+  }
+}
+// CHECK-LABEL: func @tensor_op_fusion
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG3:[a-zA-Z0-9_]+]]: tensor<?xf32>
+//       CHECK:   %[[INIT:.+]] = linalg.init_tensor
+//       CHECK:   %[[R0:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[ARG5:.+]] = %[[INIT]]) -> (tensor<?x?xf32>) {
+//       CHECK:     %[[R1:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[ARG7:.+]] = %[[ARG5]]) -> (tensor<?x?xf32>) {
+//   CHECK-DAG:       %[[STARG3:.+]] = subtensor %[[ARG3]]
+//   CHECK-DAG:       %[[STARG7:.+]] = subtensor %[[ARG7]]
+//   CHECK-DAG:       %[[STARG0:.+]] = subtensor %[[ARG0]]
+//   CHECK-DAG:       %[[STARG1:.+]] = subtensor %[[ARG1]]
+//   CHECK-DAG:       %[[STARG2:.+]] = subtensor %[[ARG2]]
+//       CHECK:       %[[T0:.+]] = linalg.matmul
+//  CHECK-SAME:         ins(%[[STARG0]], %[[STARG1]] : tensor<?x?xf32>, tensor<?x?xf32>)
+//  CHECK-SAME:         outs(%[[STARG2]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+//       CHECK:       %[[T1:.+]] = linalg.generic
+//  CHECK-SAME:         ins(%[[T0:.+]], %[[STARG3]] : tensor<?x?xf32>, tensor<?xf32>)
+//  CHECK-SAME:         outs(%[[STARG7]] : tensor<?x?xf32>)
+//       CHECK:       %[[RESULT:.+]] = subtensor_insert %[[T1]] into %[[ARG7]]
+//       CHECK:       scf.yield %[[RESULT]]
+//       CHECK:     }
+//       CHECK:     scf.yield %[[R1]]
+//       CHECK:   }
+//       CHECK:   return %[[R0]]
+
+// -----
+
+module {
+  func @tensor_matmul_fusion(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
+                             %arg2: tensor<?x?xf32>, %arg3: tensor<?x?xf32>,
+           %arg4: tensor<?x?xf32>, %arg5: tensor<?x?xf32>,
+           %arg6: tensor<?x?xf32>) -> tensor<?x?xf32> {
+    %0 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> // [M, N0] * [N0, N1]
+    %1 = linalg.matmul ins(%0, %arg3 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%arg4 : tensor<?x?xf32>) -> tensor<?x?xf32> // [M, N1] * [N1, N2]
+    %2 = linalg.matmul ins(%1, %arg5 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%arg6 : tensor<?x?xf32>) -> tensor<?x?xf32> // [M, N2] * [N2, N3]
+    return %2 : tensor<?x?xf32>
+  }
+}
+// CHECK-LABEL: func @tensor_matmul_fusion(
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG3:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG4:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG5:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG6:[a-zA-Z0-9_]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
+//   CHECK-DAG:   %[[C0:.+]] = constant 0 : index
+//   CHECK-DAG:   %[[C1:.+]] = constant 1 : index
+//       CHECK:   %[[R0:.+]] = scf.for %[[IV0:[a-zA-Z0-9_]+]] =
+//  CHECK-SAME:     iter_args(%[[ARG8:.+]] = %[[ARG6]]) -> (tensor<?x?xf32>) {
+//       CHECK:       %[[N3:.+]] = dim %[[ARG8]], %[[C1]]
+//       CHECK:       %[[STARG6:.+]] = subtensor %[[ARG8]][%[[IV0]], 0]
+//  CHECK-SAME:         [%{{[a-zA-Z0-9_]+}}, %[[N3]]]
+//       CHECK:       %[[N2:.+]] = dim %[[ARG3]], %[[C1]]
+//       CHECK:       %[[N1:.+]] = dim %[[ARG1]], %[[C1]]
+//       CHECK:       %[[STARG3:.+]] = subtensor %[[ARG3]][0, 0]
+//  CHECK-SAME:         [%[[N1]], %[[N2]]]
+//       CHECK:       %[[STARG4:.+]] = subtensor %[[ARG4]][%[[IV0]], 0]
+//  CHECK-SAME:         [%{{[a-zA-Z0-9_]+}}, %[[N2]]]
+//       CHECK:       %[[N0:.+]] = dim %[[ARG0]], %[[C1]]
+//       CHECK:       %[[STARG0:.+]] = subtensor %[[ARG0]][%[[IV0]], 0]
+//  CHECK-SAME:         [%{{[a-zA-Z0-9_]+}}, %[[N0]]]
+//       CHECK:       %[[STARG1:.+]] = subtensor %[[ARG1]][0, 0]
+//  CHECK-SAME:         [%[[N0]], %[[N1]]]
+//       CHECK:       %[[STARG2:.+]] = subtensor %[[ARG2]][%[[IV0]], 0]
+//  CHECK-SAME:         [%{{[a-zA-Z0-9_]+}}, %[[N1]]]
+//       CHECK:       %[[T0:.+]] = linalg.matmul
+//  CHECK-SAME:         ins(%[[STARG0]], %[[STARG1]]
+//  CHECK-SAME:         ) outs(%[[STARG2]] : tensor<?x?xf32>)
+//       CHECK:       %[[T1:.+]] = linalg.matmul
+//  CHECK-SAME:         ins(%[[T0]], %[[STARG3]]
+//  CHECK-SAME:         ) outs(%[[STARG4]] : tensor<?x?xf32>)
+//       CHECK:       %[[T2:.+]] = linalg.matmul
+//  CHECK-SAME:         ins(%[[T1]], %[[ARG5]]
+//  CHECK-SAME:         ) outs(%[[STARG6]] : tensor<?x?xf32>)
+//       CHECK:       %[[R1:.+]] = subtensor_insert %[[T2]]
+//  CHECK-SAME:         into %[[ARG8]][%[[IV0]], %[[C0]]]
+//       CHECK:       scf.yield %[[R1]]
+//       CHECK:     }
+//       CHECK:     return %[[R0]]
+//       CHECK:   }

mlir/test/lib/Transforms/TestLinalgFusionTransforms.cpp

void runOnFunction() override {
     if (!llvm::hasSingleElement(blocks)) {
       return;
     }
     SmallVector<LinalgOp, 2> linalgOps =
         llvm::to_vector<2>(blocks.front().getOps<LinalgOp>());
     Aliases aliases;
     LinalgDependenceGraph dependenceGraph(aliases, linalgOps);
     OpBuilder builder(funcOp.getContext());
+    linalg::LinalgTilingLoopType loopType = LinalgTilingLoopType::ParallelLoops;
+    if (llvm::all_of(linalgOps, [](LinalgOp linalgOp) {
+          return linalgOp.hasTensorSemantics();
+        }))
+      loopType = LinalgTilingLoopType::Loops;
     Optional<TiledAndFusedLinalgOps> tileAndFuseOps = tileAndFuseLinalgOps(
         builder, linalgOps, dependenceGraph,
-        LinalgTilingOptions().setTileSizes(tileSizes).setLoopType(
-            LinalgTilingLoopType::ParallelLoops));
+        LinalgTilingOptions().setTileSizes(tileSizes).setLoopType(loopType));
     if (!tileAndFuseOps)
       return signalPassFailure();
+    if (linalgOps.back().hasTensorSemantics()) {
+      linalgOps.back().getOperation()->replaceAllUsesWith(
+          tileAndFuseOps->fusedLoops.front());
+    }
     for (auto op : linalgOps)
+      if (op.hasBufferSemantics())
         op.erase();
   }
 };
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestLinalgFusionTransforms() {
   PassRegistration<TestLinalgFusionTransforms> testFusionTransformsPass(