[mlir][Linalg] Add basic support for TileAndFuse on Linalg on tensors.

[nicolasvasilache]

Summary

This revision allows the fusion of the producer of input tensors in the consumer under a tiling transformation (which produces subtensors).
Many pieces are still missing (e.g. support init_tensors, better refactor LinalgStructuredOp interface support, try to merge implementations and reuse code) but this still allows getting started.

The greedy pass itself is just for testing purposes and will be extracted in a separate test pass.

[mravishankar]

This is really cool! Minimal change to get this working on tensors! Some minor comments.

Do you want to try using the TileAndFusePattern to test the fusion on tensors. I think it should just work.

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

92

This can be confusing with the fuseTensorOps method below. They are doing different things. maybe some more comments on that would be helpful.

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

225

I am not sure why you need to do this. Cant you just use the resultTypes of the producer directly. That should already be the the tensorType that you have in the consumer right?

440

Looking at this, the fuse method is already able to handle tensor or buffer ops. The only difference between this and fuseProduceOfBuffer method is the logic to get the producer. Do we want to even have a separate function for this. Based on the type of the operand at consumerIdx we can get the produer using the buffer approach or tensor approach and just call fuse. That seems like it should work.

[nicolasvasilache]

Address review.

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

92

Ah yes, thanks!

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

225

I don't think this is true, see the test:

  %t0 = linalg.matmul ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
                     init(%arg2: tensor<?x?xf32>)
    -> tensor<?x?xf32>
  ...
  %3 = scf.for %arg3 = %c0 to %0 step %c2 iter_args(%arg4 = %arg2) -> (tensor<?x?xf32>) {
    %4 = scf.for %arg5 = %c0 to %2 step %c3 iter_args(%arg6 = %arg4) -> (tensor<?x?xf32>) {
      %5 = scf.for %arg7 = %c0 to %1 step %c4 iter_args(%arg8 = %arg6) -> (tensor<?x?xf32>) {
        ...
        %10 = subtensor %t0[%arg3, %arg7] [%7, %9] [1, 1]  : tensor<?x?xf32> to tensor<?x?xf32>
        ...
        %21 = linalg.matmul ins(%10, %15 : tensor<?x?xf32>, tensor<?x?xf32>) init(%20 : tensor<?x?xf32>)  -> tensor<?x?xf32>
...

That subtensor %t0 may easily change the type if we have some static size and the "tile and fused producer" needs to produce that type.
So I am just taking it from the consumer input.
Updated the test to use a mix of static and dynamic to better illustrate what is happening problem.

Also adding an assertion to guard against multi-outputs: this will require more thinking as to where to get the output shape for other returned tensors.

440

Yes, I am planning on doing tighter integration in subsequent commits.
Originally, I had split fuseLinalgOpsGreedily into 2 functions but then saw it was ok to dispatch per operand.
Still, I wanted to keep the chances of messing up input/output_buffers/input_tensors operands + types low.

I can wait until all steps are ready and squash them together later if you prefer?
Tentatively marking as done for now, please reopen if you feel this is not enough.

[mravishankar]

Thanks Nicolas! If there is no hurry, I'd like to take a couple of days and do a deep dive into the transformation here and get a better idea of the changes. So if it isnt blocking, then I will circle back to this a couple of days.

In the meanwhile, some tests seem to be failing. Maybe its misattribute...

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

225

Just not sure why this needs to be done for tensor, when this is not done for the buffer. AFAICS, the subview computed gets the right shape of the operands for the cloned op. Same should be true for the subtensor op.
I might be missing something here. Its worth me going into details of this. Sorry for the delay. Ill do a deep dive and update here.

Btw, I dont see the updated example. Seems like it needs to be uploaded?

[mravishankar]

Seems to be some failing tests here. Will circle back after that is fixed.

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

101

This is the part that I am referring to below. The subview calculation above takes the view, offset, sizes, strides and creates a memref that exactly matches (at least as far as I know), the type of the memref of the operand that is the fused view. In the tensor equivalent the fused view is by construction the result of the subtensor operation. It must be creating the same tensor type that is the operand in the consumer. If this isnt the case then the fusion is wrong AFAICS.

[nicolasvasilache]

Refactor and address review.

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

101

The subview calculation above takes the view, offset, sizes, strides and creates a memref that exactly matches (at least as far as I know), the type of the memref of the operand that is the fused view

This does not hold at the moment of construction: the type of the producing subview is always built dynamically.
Even when passing constant sizes, they are constant index ops and the subview builder returns a fully dynamic type.
On the other hand, the consumer may already have been canonicalized (or constructed this way) and can be "more static".
In other words, at the time of construction, the types may not match, but after enough canonicalizations+foldings kick in, the static types are consistent.
The invariants of the transformation guarantee that the types are compatible.
We can debate on whether or not this is reasonable and whether we should write more code to detect static sizes coming from SSA values.
For now this is left to canonicalization + folding which may pull quite a few patterns (e.g. affine_min + scf.for canonicalization for "statically known to be full" tiles).

In the buffer case, there is no use-def chain whose type needs to be consistent so this mismatch did not pose problems.
We have 2 aliasing views that after canonicalization + cse become the same view object and trigger further fusions.

In the tensor case, the IR needs to be valid and this needs a cast.
This is similar to other transformations that introduce memref_cast and tensor_cast ops.

I refactored this to separate make it more explicit what happens, in particular output tensor types are now inferred (always dynamic).

It must be creating the same tensor type that is the operand in the consumer. If this isn't the case then the fusion is wrong AFAICS.

It does but this is only apparent after canonicalization + cse, I added an example in the test to show this.
In the following patch I'll refactor the pass to become a purely test pass that uses pattern rewrites mixed with canonicalization, so I'll be able to remove that special from the RUN command.

225

Just not sure why this needs to be done for tensor, when this is not done for the buffer. AFAICS, the subview computed gets the right shape of the operands for the cloned op. Same should be true for the subtensor op.

See explanation above. In the buffer case, the type mismatch was hidden by aliasing. In the tensor case we need to be type consistent.

Btw, I dont see the updated example. Seems like it needs to be uploaded?

Seems I missed something when uploading before leaving for 1 week, sorry about that!
Should be fixed now.

[mravishankar]

Thanks for the refactoring. Make sense now. Exciting!

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

110

THanks for the comments and refactoring. I understand now how this works.

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h

 #include "mlir/IR/StandardTypes.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/CopyOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "mlir/Support/LLVM.h"
 
+#include "llvm/ADT/STLExtras.h"
+
 namespace mlir {
 namespace linalg {
 
 class ConvOp;
 class PoolingMaxOp;
 class PoolingMinOp;
 class PoolingSumOp;
 

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

InterfaceMethod<
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         auto range = this->getOperation()->getOperands();
         return {range.begin(), range.begin() + $_op.getNumInputs()};
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
+        Return the range over the input operands that are of buffer type.
+      }],
+      /*retTy=*/"SmallVector<Value, 4>",
+      /*methodName=*/"getInputBuffers",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return llvm::to_vector<4>(llvm::make_filter_range(
+          getInputs(), [](Value in){ return in.getType().isa<MemRefType>(); }));
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
         Return the subset of input operands that are of ranked tensor type.
       }],
       /*retTy=*/"SmallVector<RankedTensorType, 4>",
       /*methodName=*/"getInputTensorTypes" ,
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         SmallVector<RankedTensorType, 4> res;
InterfaceMethod<
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return getNumInputsAndOutputBuffers() + $_op.getNumInitTensors();
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
+        Return the `i`-th shaped operand value, which can be an arbitrary input
+        tensor/buffer, init tensor or output buffer.
+      }],
+      /*retTy=*/"Value",
+      /*methodName=*/"getShapedOperand",
+      /*args=*/(ins "unsigned":$i),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        assert(i < $_op.getNumShapedOperands());
+        return this->getOperation()->getOperand(i);
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
         Return the range over inputs, output buffers and init tensors.
       }],
       /*retTy=*/"Operation::operand_range",
       /*methodName=*/"getShapedOperands",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         auto range = this->getOperation()->getOperands();

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 bool isFusableInto(const LinalgDependenceGraph &graph, LinalgOp consumer,
                    Value consumedView, LinalgOp producer);
 
 /// Fuses producer into consumer if the producer is structurally feasible and
 /// the fusion would not violate dependencies.
 /// When non-null, the optional pointer `folder` is used to call into the
 /// `createAndFold` builder method. If `folder` is null, the regular `create`
 /// method is called.
-Optional<FusionInfo> fuseProducerOf(OpBuilder &b, LinalgOp consumer,
+Optional<FusionInfo> fuseProducerOfBuffer(OpBuilder &b, LinalgOp consumer,
                                           unsigned consumerIdx,
                                           const LinalgDependenceGraph &graph,
                                           OperationFolder *folder = nullptr);
+/// Tensor counterpart of `fuseProducerOfBuffer`.
+Optional<FusionInfo> fuseProducerOfTensor(OpBuilder &b, LinalgOp consumer,
+                                          unsigned consumerIdx,
+                                          OperationFolder *folder);
 
 /// Fuse linalg operation on tensors, with the producer of the operand at
 /// position `consumerIdx` of the consumer.
 Operation *fuseTensorOps(PatternRewriter &rewriter, Operation *consumer,
                          unsigned consumerIdx,
                          OperationFolder *folder = nullptr);
 
 /// Returns the linearized list of all shape dimensions in a `linalgOp`.

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

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) {
-  assert(src.hasBufferSemantics() &&
-         "expected linalg op with buffer semantics");
-  assert(dst.hasBufferSemantics() &&
-         "expected linalg op with buffer semantics");
   for (auto srcView : src.getOutputBuffers()) { // W
     // RAW graph
-    for (auto dstView : dst.getInputs()) {   // R
+    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});
       }
     }
     // 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 srcView : src.getInputs()) { // R
+  for (auto srcView : src.getInputBuffers()) { // R
     // RAR graph
-    for (auto dstView : dst.getInputs()) {   // R
+    for (auto dstView : dst.getInputBuffers()) { // R
       if (aliases.alias(srcView, dstView)) { // if alias, fill RAR
         addDependenceElem(DependenceType::RAR,
                           LinalgOpView{src.getOperation(), srcView},
                           LinalgOpView{dst.getOperation(), dstView});
       }
     }
     // WAR graph
     for (auto dstView : dst.getOutputBuffers()) { // W

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

 ///
 /// More advanced use cases, analyses as well as profitability heuristics are
 /// left for future work.
 
 // Return a cloned version of `op` that operates on `loopRanges`, assumed to be
 // a subset of the original loop ranges of `op`.
 // This is achieved by applying the `loopToOperandRangesMaps` permutation maps
 // to the `loopRanges` in order to obtain view ranges.
-static LinalgOp cloneWithLoopRanges(OpBuilder &b, Location loc, LinalgOp op,
-                                    ArrayRef<Range> loopRanges) {
-  assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics");
+static LinalgOp cloneWithLoopRangesAndTypes(OpBuilder &b, Location loc,
+                                            LinalgOp op,
+                                            ArrayRef<Range> loopRanges,
+                                            TypeRange resultTypes) {
   auto maps = op.indexing_maps();
   SmallVector<Value, 8> clonedViews;
   clonedViews.reserve(op.getNumInputsAndOutputs());
   // 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)) {
+  for (auto en : llvm::enumerate(op.getShapedOperands())) {
     unsigned idx = en.index();
     auto map = maps[idx].cast<AffineMapAttr>().getValue();
     LLVM_DEBUG(dbgs() << "map: " << map << "\n");
     Value view = en.value();
     SmallVector<Range, 4> viewRanges(map.getNumResults());
     for (auto en2 : llvm::enumerate(map.getResults())) {
       unsigned d = en2.index();
       // loopToOperandRangesMaps are permutations-only.
for (auto en : llvm::enumerate(op.getShapedOperands())) {
     offsets.reserve(rank);
     sizes.reserve(rank);
     strides.reserve(rank);
     for (auto r : viewRanges) {
       offsets.push_back(r.offset);
       sizes.push_back(r.size);
       strides.push_back(r.stride);
     }
+    if (view.getType().isa<MemRefType>())
       clonedViews.push_back(
           b.create<SubViewOp>(loc, view, offsets, sizes, strides));
+    else
+      clonedViews.push_back(
+          b.create<SubTensorOp>(loc, view, offsets, sizes, strides));
   }
   auto operands = op.getAssumedNonShapedOperands();
   clonedViews.append(operands.begin(), operands.end());
 
-  Operation *clonedOp = op.clone(b, loc, /*resultTypes*/ {}, clonedViews);
+  Operation *clonedOp = op.clone(b, loc, resultTypes, clonedViews);
   // When the producer is an IndexedGenercOp, we have to transform its block
   // IV arguments according to the tiling of the consumer, i.e. offset them by
   // the values computed in `loopRanges`.
   if (auto indexedGenericOp = dyn_cast<IndexedGenericOp>(clonedOp)) {
     auto &block = indexedGenericOp.region().front();
 
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPointToStart(&block);
     for (unsigned i = 0, e = indexedGenericOp.getNumLoops(); i < e; ++i) {
       Value oldIndex = block.getArgument(i);
       AddIOp newIndex = b.create<AddIOp>(indexedGenericOp.getLoc(), oldIndex,

 };
 
 // Given an `op`, returns the first (`view`, `dimension`) pair that identifies
 // the loop range at `loopDepth`. The semantics of the loopToOperandRangesMaps
 // guarantees at least one such dimension is found. If multiple candidates exist
 // they must agree by construction (i.e. have the same size) and we just return
 // the first one.
 static ViewDimension getViewDefiningLoopRange(LinalgOp op, unsigned loopDepth) {
-  assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics");
   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() << "getViewDefiningLoopRange I/O idx: " << idx << "\n");
     LLVM_DEBUG(dbgs() << "getViewDefiningLoopRange map: " << map << "\n");
     Value view = en.value();
     SmallVector<Value, 8> viewRanges(map.getNumResults(), nullptr);
     for (auto en2 : llvm::enumerate(map.getResults())) {
       if (loopDepth == en2.value().cast<AffineDimExpr>().getPosition()) {
         LLVM_DEBUG(dbgs() << "getViewDefiningLoopRange loopDepth: " << loopDepth
                           << "\n");
         LLVM_DEBUG(dbgs() << "getViewDefiningLoopRange view: " << view << "\n");
         return ViewDimension{view, static_cast<unsigned>(en2.index())};
       }
     }
   }
   llvm_unreachable("Expect to be able to extract a view defining loop range");
 }
 
+/// Fuses the producer of `producerIdx` into the loop immediately enclosing
+/// `consumer`. This is achieved by "recomputing" the `producer` at the time it
+/// is needed just before the `consumer.
+///
+/// Depending on the type of `consumer.getShapedOperand(consumerIdx)`, there are
+/// 2 cases:
+///   1. Buffer case: `producerIdx` is the index of the buffer in
+///      `producer.getShapedOperands()`.
+///   2. Tensor case: `producerIdx` is the index of the tensor in
+///      `producer.getResults()`.
 static LinalgOp fuse(OpBuilder &b, LinalgOp producer, unsigned producerIdx,
                      LinalgOp consumer, unsigned consumerIdx,
                      OperationFolder *folder = nullptr) {
-  assert(producer.hasBufferSemantics() &&
-         "expected linalg op with buffer semantics");
-  assert(consumer.hasBufferSemantics() &&
-         "expected linalg op with buffer semantics");
-
-  auto subView = dyn_cast_or_null<SubViewOp>(
-      consumer.getBuffer(consumerIdx).getDefiningOp());
-  auto slice = dyn_cast_or_null<SliceOp>(
-      consumer.getBuffer(consumerIdx).getDefiningOp());
-  assert(subView || slice);
-  (void)subView;
-  (void)slice;
+  Operation *shapeProducingOp =
+      consumer.getShapedOperand(consumerIdx).getDefiningOp();
+  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.indexing_maps()[producerIdx].cast<AffineMapAttr>().getValue();
   LLVM_DEBUG(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.
   auto loc = consumer.getLoc();
   for (auto en : llvm::enumerate(producerMap.getResults())) {
     unsigned posInProducerLoop = en.value().cast<AffineDimExpr>().getPosition();
     loopRanges[posInProducerLoop] =
-        subView.getOrCreateRanges(b, loc)[en.index()];
+        isa<SubViewOp>(shapeProducingOp)
+            ? cast<SubViewOp>(shapeProducingOp)
+                  .getOrCreateRanges(b, loc)[en.index()]
+            : cast<SubTensorOp>(shapeProducingOp)
+                  .getOrCreateRanges(b, loc)[en.index()];
   }
 
   // Iterate over all dimensions. For the dimensions not identified by the
   // producer map for `producerIdx`, we need to explicitly compute the view that
   // defines the loop ranges using the `producer`.
   for (unsigned i = 0, nLoops = loopRanges.size(); i < nLoops; ++i) {
     if (loopRanges[i].offset)
       LLVM_DEBUG(llvm::dbgs()
                  << "existing LoopRange: " << loopRanges[i] << "\n");
     else {
       auto viewDim = getViewDefiningLoopRange(producer, i);
       loopRanges[i] = Range{folded_std_constant_index(folder, 0),
                             std_dim(viewDim.view, viewDim.dimension),
                             folded_std_constant_index(folder, 1)};
       LLVM_DEBUG(llvm::dbgs() << "new LoopRange: " << loopRanges[i] << "\n");
     }
   }
 
-  return cloneWithLoopRanges(b, loc, producer, loopRanges);
+  // Pass updated resultTypes to support the tensor case: just copy the
+  // corresponding tensor type from the consumer if appropriate.
+  SmallVector<Type, 4> resultTypes =
+      llvm::to_vector<4>(producer.getOperation()->getResultTypes());
+  RankedTensorType tensorType =
+      consumer.getShapedType(consumerIdx).dyn_cast<RankedTensorType>();
+  if (tensorType)
+    resultTypes[producerIdx] = tensorType;
+  return cloneWithLoopRangesAndTypes(b, loc, producer, loopRanges, resultTypes);
 }
 
 // Encode structural fusion safety preconditions.
 // 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");
for (auto dependence :
         continue;
 
       return dependence;
     }
   }
   return {};
 }
 
-Optional<FusionInfo> mlir::linalg::fuseProducerOf(
+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);
Optional<FusionInfo> mlir::linalg::fuseProducerOfBuffer(
   assert(producerIdxOpt.hasValue() && "incorrect operand index");
   unsigned producerIdx = producerIdxOpt.getValue();
 
   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
+static void getProducerOfTensor(Value tensor, LinalgOp &producer,
+                                unsigned &outputIndex) {
+  if (!tensor.getType().isa<RankedTensorType>())
+    return;
+
+  while (true) {
+    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 = blockArg.getDefiningOp<scf::ForOp>()) {
+        tensor = forOp.getResult(blockArg.getArgNumber());
+        continue;
+      }
+    }
+    return;
+  }
+}
+
+Optional<FusionInfo>
+mlir::linalg::fuseProducerOfTensor(OpBuilder &b, LinalgOp consumer,
+                                   unsigned consumerIdx,
+                                   OperationFolder *folder) {
+  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)");
+    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");
+  auto fusedProducer =
+      fuse(b, producerOp, producerIdx, consumer, consumerIdx, folder);
+  // Replace use.
+  consumer.getOperation()
+      ->getOpOperand(consumerIdx)
+      .set(fusedProducer.getOperation()->getResult(producerIdx));
+  return FusionInfo{producerOp, fusedProducer};
+}
+
 /// 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 void fuseLinalgOpsGreedily(FuncOp f) {
   LLVM_DEBUG(f.print(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) {
-    if (op.hasBufferSemantics())
-      linalgOps.push_back(op);
-  });
+  f.walk([&](LinalgOp op) { linalgOps.push_back(op); });
 
+  // Tile and Fuse for tensors inputs (TODO: all tensor operands).
+  for (auto *op : llvm::reverse(linalgOps)) {
+    LinalgOp linalgOp = cast<LinalgOp>(op);
+    for (auto en : llvm::enumerate(linalgOp.getShapedOperands())) {
+      if (en.value().getType().isa<MemRefType>()) {
         // TODO: LinalgDependenceGraph should be able to update itself.
         // The current naive and expensive reconstruction of the graph should be
         // removed.
-  for (auto *op : llvm::reverse(linalgOps)) {
-    for (unsigned id = 0, e = LinalgOp(op).getNumInputsAndOutputBuffers();
-         id < e; ++id) {
         linalg::Aliases aliases;
         linalg::LinalgDependenceGraph graph(aliases, linalgOps);
-      if (auto info = fuseProducerOf(b, op, id, graph, &folder)) {
+        if (auto info =
+                fuseProducerOfBuffer(b, op, en.index(), graph, &folder)) {
           auto *originalOp = info->originalProducer.getOperation();
           eraseSet.insert(originalOp);
           auto *originalOpInLinalgOpsVector =
               std::find(linalgOps.begin(), linalgOps.end(), originalOp);
           *originalOpInLinalgOpsVector = info->fusedProducer.getOperation();
         }
+      } else {
+        assert(en.value().getType().isa<RankedTensorType>());
+        // Tile and Fuse tensor input (TODO: init_tensors too).
+        if (en.index() >= linalgOp.getNumInputs())
+          continue;
+        if (auto info = fuseProducerOfTensor(b, op, en.index(), &folder)) {
+          auto *originalOp = info->originalProducer.getOperation();
+          auto *originalOpInLinalgOpsVector =
+              std::find(linalgOps.begin(), linalgOps.end(), originalOp);
+          *originalOpInLinalgOpsVector = info->fusedProducer.getOperation();
+        }
       }
     }
-  // The `fuseProducerOf` 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.
+  }
+  // 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"));
 }
 
 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/tile-and-fuse-tensors.mlir

+// RUN: mlir-opt %s -linalg-fusion -split-input-file | FileCheck %s
+
+#map0 = affine_map<(d0)[s0] -> (2, -d0 + s0)>
+#map1 = affine_map<(d0)[s0] -> (4, -d0 + s0)>
+#map2 = affine_map<(d0)[s0] -> (3, -d0 + s0)>
+#map3 = affine_map<(d0, d1) -> (2, d0 - d1)>
+#map4 = affine_map<(d0, d1) -> (3, d0 - d1)>
+
+func @matmul_tensors(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %t0 = linalg.matmul ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
+                     init(%arg2: tensor<?x?xf32>)
+    -> tensor<?x?xf32>
+
+  %c4 = constant 4 : index
+  %c2 = constant 2 : index
+  %c0 = constant 0 : index
+  %c3 = constant 3 : index
+  %c1 = constant 1 : index
+  %0 = dim %t0, %c0 : tensor<?x?xf32>
+  %1 = dim %t0, %c1 : tensor<?x?xf32>
+  %2 = dim %arg1, %c1 : tensor<?x?xf32>
+  %3 = scf.for %arg3 = %c0 to %0 step %c2 iter_args(%arg4 = %arg2) -> (tensor<?x?xf32>) {
+    %4 = scf.for %arg5 = %c0 to %2 step %c3 iter_args(%arg6 = %arg4) -> (tensor<?x?xf32>) {
+      %5 = scf.for %arg7 = %c0 to %1 step %c4 iter_args(%arg8 = %arg6) -> (tensor<?x?xf32>) {
+        %6 = dim %t0, %c0 : tensor<?x?xf32>
+        %7 = affine.min #map0(%arg3)[%6]
+        %8 = dim %t0, %c1 : tensor<?x?xf32>
+        %9 = affine.min #map1(%arg7)[%8]
+        %10 = subtensor %t0[%arg3, %arg7] [%7, %9] [1, 1]  : tensor<?x?xf32> to tensor<?x?xf32>
+        %11 = dim %arg1, %c0 : tensor<?x?xf32>
+        %12 = affine.min #map1(%arg7)[%11]
+        %13 = dim %arg1, %c1 : tensor<?x?xf32>
+        %14 = affine.min #map2(%arg5)[%13]
+        %15 = subtensor %arg1[%arg7, %arg5] [%12, %14] [1, 1]  : tensor<?x?xf32> to tensor<?x?xf32>
+        %16 = dim %arg8, %c0 : tensor<?x?xf32>
+        %17 = affine.min #map3(%16, %arg3)
+        %18 = dim %arg8, %c1 : tensor<?x?xf32>
+        %19 = affine.min #map4(%18, %arg5)
+        %20 = subtensor %arg8[%arg3, %arg5] [%17, %19] [1, 1]  : tensor<?x?xf32> to tensor<?x?xf32>
+        %21 = linalg.matmul ins(%10, %15 : tensor<?x?xf32>, tensor<?x?xf32>) init(%20 : tensor<?x?xf32>)  -> tensor<?x?xf32>
+        %22 = subtensor_insert %21 into %arg8[%arg3, %arg5] [%17, %19] [%c1, %c1]  : tensor<?x?xf32> into tensor<?x?xf32>
+        scf.yield %22 : tensor<?x?xf32>
+      }
+      scf.yield %5 : tensor<?x?xf32>
+    }
+    scf.yield %4 : tensor<?x?xf32>
+  }
+  return %3 : tensor<?x?xf32>
+}
+
+// CHECK-LABEL: func @matmul_tensors(
+ // CHECK-SAME: %[[A:[0-9a-z]*]]: tensor<?x?xf32>
+ // CHECK-SAME: %[[B:[0-9a-z]*]]: tensor<?x?xf32>
+ // CHECK-SAME: %[[C:[0-9a-z]*]]: tensor<?x?xf32>
+//       CHECK: %[[C0:.*]] = constant 0 : index
+//       CHECK: scf.for %[[I:[0-9a-z]*]]
+//  CHECK-NEXT:   scf.for %[[J:[0-9a-z]*]]
+//  CHECK-NEXT:     scf.for %[[K:[0-9a-z]*]]
+//
+// subtensor of the original program, first one refers to the unfused matmul and becomes a dead SSA value.
+//       CHECK:     subtensor %{{.*}}[%[[I]], %[[K]]] {{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
+//       CHECK:     %[[stB1:.*]] = subtensor %[[B]][%[[K]], %[[J]]] {{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
+//       CHECK:     %[[stF:.*]] = subtensor %{{.*}}[%[[I]], %[[J]]] {{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
+//
+// subtensors of the producing matmul.
+//       CHECK:     %[[stA:.*]] = subtensor %[[A]][%[[I]], %[[K]]] {{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
+//  CHECK-NEXT:     %[[stB2:.*]] = subtensor %[[B]][%[[K]], %[[C0]]] {{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
+//  CHECK-NEXT:     %[[stC:.*]] = subtensor %[[C]][%[[I]], %[[C0]]] {{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
+//  CHECK-NEXT:     %[[stD:.*]] = linalg.matmul ins(%[[stA]], %[[stB2]] : tensor<?x?xf32>, tensor<?x?xf32>) init(%[[stC]] : tensor<?x?xf32>)  -> tensor<?x?xf32>
+//  CHECK-NEXT:     %[[stG:.*]] = linalg.matmul ins(%[[stD]], %[[stB1]] : tensor<?x?xf32>, tensor<?x?xf32>) init(%[[stF]] : tensor<?x?xf32>)  -> tensor<?x?xf32>
+//  CHECK-NEXT:     subtensor_insert %[[stG]]