[mlir][Linalg] Allow `tileConsumerAndFuseProducers` to return the values of the fused operations.

[mravishankar]

Summary

Current tileConsumerAndFuseProducers method allows for only
replacing the uses of the tiled consumer, and not the fused producers.
Based on how fusion works currently, it is not always possible/easy to
return a tensor value from the tiled loop nest to replace the uses of
the fused producers since the same tile of the fused producer might be
recomputed while computing multiple tiles of the consumer op. This
patch adds an option that will allow the transformation to return the
tensor replacements for the fused producer ops, but expects the caller
to control the tile sizes so that the tiles of the fused producer are
computed only once, by controlling which loops are tiled. There might
be an analysis that can determine this, but for now this left to
control from the caller.

[nicolasvasilache]

Discussed this offline, the plan is to start by building small independent IR transform utils that can be tested independently and reevaluate how they compose once we have a better visibility of the parts.

[mravishankar]

Abondoning this in favor of https://reviews.llvm.org/D133474

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

 
 struct LinalgTilingAndFusionOptions {
   /// Tile sizes used to tile the root operation.
   SmallVector<int64_t> tileSizes;
   LinalgTilingAndFusionOptions &setTileSizes(ArrayRef<int64_t> ts) {
     tileSizes.assign(ts.begin(), ts.end());
     return *this;
   }
+
   /// Tile interchange used to permute the tile loops.
   SmallVector<int64_t> tileInterchange;
+  LinalgTilingAndFusionOptions &setInterchane(ArrayRef<unsigned> interchange) {
+    tileInterchange.assign(interchange.begin(), interchange.end());
+    return *this;
+  }
+
   /// When specified, specifies distribution of generated tile loops to
   /// processors.
   Optional<LinalgLoopDistributionOptions> tileDistribution = None;
   LinalgTilingAndFusionOptions &
   setDistributionOptions(LinalgLoopDistributionOptions distributionOptions) {
     tileDistribution = std::move(distributionOptions);
     return *this;
   }
+
+  /// When specified, this returns from the tiled loops, the tensor values
+  /// that can be used as replacements for the operations that are fused into
+  /// the consumer tiled loops. Note that fusion can lead to the same tile of
+  /// the producer being recomputed for producing different tiles of the
+  /// consumer. In such situations, it is tricky to to contruct the return value
+  /// for the fused ops within the loop. So this option is to be used with care
+  /// only when the caller knows that the fused op is not recomputed. For
+  /// example,
+  ///
+  /// ```mlir
+  ///   %C = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
+  ///       outs (%init1 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  ///   %E = linalg.matmul ins(%C, %D : tensor<?x?xf32>, tensor<?x?xf32>)
+  ///       outs (%init2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  /// ```
+  ///
+  /// depending the tiling chosen for the consumer `linalg.matmul` can be tiled
+  /// either along two outer loop dimensions (`i` and `j` loops), or just the
+  /// first outer loop dimension (`i`) loop could be tiled. Fusing the producer
+  /// results in either
+  ///
+  /// ```mlir
+  /// %l0 = scf.for %iv0 = ... iter_args(%arg0 = %init2) {
+  ///   %l1 = scf.for %iv1 = ... iter_args(%arg1 = %arg0) {
+  ///     %C = linalg.matmul
+  ///     %E = linalg.matmul ins(%C, ...) ...
+  ///     scf.yield %E
+  ///   }
+  /// }
+  /// ```
+  ///
+  /// or
+  ///
+  /// ```mlir
+  /// %l0 = scf.for %iv0 = ... iter_args(%arg0 = %init2) {
+  ///     %C = linalg.matmul ...
+  ///     %E = linalg.matmul ins(%C, ...) ...
+  ///     scf.yield %E
+  ///   }
+  /// }
+  /// ```
+  ///
+  /// In the first case, the same producer is recomputed along all iterations of
+  /// the inner loop. In the second case, there is no recomputation. In such
+  /// cases the value of the producer computed within the tile loop can be
+  /// returned as well so that it could replace the uses of the untiled producer
+  /// op. Setting the option to `true` assumes that the producer tiles are not
+  /// recomputed after fusion and the caller is expected to ensure this
+  /// assumption holds.
+  bool returnFusedOpValues = false;
+  LinalgTilingAndFusionOptions &setReturnFusedOpValues(bool v) {
+    returnFusedOpValues = v;
+    return *this;
+  }
 };
 
 struct LinalgTilingOptions {
   /// Computation function that returns the tile sizes for each operation.
   /// Delayed construction of constant tile sizes should occur to interoperate
   /// with folding.
   TileSizeComputationFunction tileSizeComputationFunction = nullptr;
 

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

public:
   /// and `tileDistribution`.
   LogicalResult
   tileRootOp(OpBuilder &b, ArrayRef<int64_t> tileSizes,
              ArrayRef<int64_t> tileInterchange,
              Optional<LinalgLoopDistributionOptions> tileDistribution);
 
   /// Fuse the producer of `consumerOpOperand` into the tile loop nest. Returns
   /// the fused producer or fails if fusion is not possible.
-  FailureOr<LinalgOp> fuseProducer(OpBuilder &b, OpOperand *consumerOpOperand);
+  FailureOr<LinalgOp> fuseProducer(OpBuilder &b, OpOperand *consumerOpOperand,
+                                   bool replaceFusedOpExternalUses);
 
-  /// Returns the replacement results for the original untiled root operation.
-  ValueRange getRootOpReplacementResults();
+  /// For the `untiledOp` that is fused with the root op, `clonedOp` being the
+  /// fused operation, modify the tiled loop nest to yield a value that can
+  /// be used to replace `untiledOp`.
+  void yieldFusedValues(OpBuilder &b, LinalgOp untiledOp, LinalgOp clonedOp,
+                        ArrayRef<OpFoldResult> mixedOffsets,
+                        ArrayRef<OpFoldResult> mixedSizes,
+                        ArrayRef<OpFoldResult> mixedStrides);
+
+  /// Returns the replacement results for the original untiled operation.
+  Value getUntiledOpResultReplacement(Value v);
 
   /// Returns the tiled root operation.
   LinalgOp getRootOp() { return rootOp; }
 
   /// Returns the tiled root operation and the fused producers.
   SmallVector<LinalgOp> getAllTiledAndFusedOps();
 
+  /// Returns the ops that were tiled and fused into the loop (returns the
+  /// original ops)
+  ArrayRef<Operation *> getFusedOps() { return fusedOps; }
+
   /// Returns the loop ops generated from tiling.
   ArrayRef<scf::ForOp> getLoopOps() { return tileLoopOps; }
 
   /// Returns true if the tile loop nest has no tile loops.
   bool isEmpty();
 
 private:
   /// Returns true if the tile loop nest invariants are satisfied:
private:
   /// Returns true if `bbArg` has other used than `sliceOp` and its
   /// dependencies. Only if there are no other uses, the producer output
   /// iteration argument may reused to pass the producer result after fusion.
   bool hasOtherUses(BlockArgument bbArg, tensor::ExtractSliceOp sliceOp);
 
   LinalgOp rootOp;
   SmallVector<scf::ForOp> tileLoopOps;
   DenseMap<Operation *, SmallVector<int64_t>> tiledRootAndFusedOpsLoops;
+
+  SmallVector<Operation *> fusedOps;
+  DenseMap<Value, Value> replacements;
 };
 
 /// Tiles `consumerOp` and fuses its dependencies if possible. Uses the
 /// `tileSizes`, `tileInterchange`, and `tileDistribution` parameters to control
 /// the tiling.
 FailureOr<TileLoopNest> tileConsumerAndFuseProducers(
     OpBuilder &b, LinalgOp consumerOp, ArrayRef<int64_t> tileSizes,
     ArrayRef<int64_t> tileInterchange,
-    const Optional<LinalgLoopDistributionOptions> &tileDistribution);
+    const Optional<LinalgLoopDistributionOptions> &tileDistribution,
+    bool returnFusedOpValues = false);
 
 //===----------------------------------------------------------------------===//
 // Generic op region utilities
 //===----------------------------------------------------------------------===//
 
 /// A struct containing common matchers over linalg op's region.
 struct RegionMatcher {
   enum class BinaryOpKind {

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

LogicalResult TileLoopNest::tileRootOp(
   if (failed(tiledRootOp))
     return failure();
 
   // Replace all uses of the root operation if it has been tiled before. All
   // uses of the original untiled root operation are updated by the calling pass
   // or pattern.
   if (!isEmpty())
     rootOp->replaceAllUsesWith(tiledRootOp->tensorResults);
+  fusedOps.push_back(rootOp);
+  for (auto result : llvm::enumerate(rootOp->getResults()))
+    replacements[result.value()] = tiledRootOp->tensorResults[result.index()];
 
   // Transfer the stored `rootOp` loop dimensions if it has been tiled before.
   if (tiledRootAndFusedOpsLoops.count(rootOp) != 0) {
     tiledRootAndFusedOpsLoops[tiledRootOp->op] =
         tiledRootAndFusedOpsLoops[rootOp];
   }
 
   // Update the root operation and append the loops and tile loop dimensions.
   rootOp = tiledRootOp->op;
   tileLoopOps.append(tiledRootOp->loops.begin(), tiledRootOp->loops.end());
   for (const auto &en : enumerate(tileSizes)) {
     // Copy only the tiled loop dimensions with non-zero tile size.
     if (en.value() == 0)
       continue;
     tiledRootAndFusedOpsLoops[rootOp].push_back(tileInterchange[en.index()]);
   }
   assert(isValid() && "expect tile loop nest to be valid after tiling");
   return success();
 }
 
+void TileLoopNest::yieldFusedValues(OpBuilder &b, LinalgOp untiledOp,
+                                    LinalgOp tiledOp,
+                                    ArrayRef<OpFoldResult> offsets,
+                                    ArrayRef<OpFoldResult> sizes,
+                                    ArrayRef<OpFoldResult> strides) {
+  if (tileLoopOps.empty())
+    return;
+
+  SmallVector<Value> newIterOperands;
+  SmallVector<Value> yieldedResults;
+  SmallVector<Value> returnedFusedOpResults;
+  for (auto result : llvm::enumerate(untiledOp->getResults())) {
+    // Check if any of the uses are outside the tile loop nest. Dim
+    // uses are not treated as "real" uses, further the dim ops could be used
+    // to compute the bounds of the loop itself. They are expected to be
+    // resolved using `ReifyRankedShapedTypeOpInterface`.
+    if (llvm::all_of(result.value().getUsers(), [&](Operation *user) {
+          return tileLoopOps.front()->isAncestor(user) ||
+                 isa<tensor::DimOp>(user) || llvm::is_contained(fusedOps, user);
+          ;
+        }))
+      continue;
+
+    unsigned resultNumber = result.value().getResultNumber();
+    // 1. Replace all existing uses of untiled op in the iter_args of the tile
+    //    loops with the `outs` of the untiled operation.
+    OpOperand *untiledOpOutOperand = untiledOp.getOutputOperand(resultNumber);
+    tileLoopOps.front()->replaceUsesOfWith(result.value(),
+                                           untiledOpOutOperand->get());
+
+    // 2. The `outs` ops of the untiled op are the "new" iter operands to be
+    //    added.
+    newIterOperands.push_back(untiledOpOutOperand->get());
+    yieldedResults.push_back(tiledOp->getResult(resultNumber));
+    returnedFusedOpResults.push_back(result.value());
+  }
+
+  // The inner most loop body needs to be modified to add `tensor.insert_slice`s
+  // to reconstruct the full tensor using the values produced from the tiled
+  // ops.
+  scf::ForOp outerMostLoop = tileLoopOps.front();
+  scf::ForOp innerMostLoop = tileLoopOps.back();
+  NewYieldValueFn innerFn = [&](OpBuilder &innerBuilder, Location innerLoc,
+                                ArrayRef<BlockArgument> innerNewBBArgs) {
+    SmallVector<Value> newYieldVals;
+    for (auto yieldedResult : llvm::enumerate(yieldedResults)) {
+      newYieldVals.push_back(innerBuilder.create<tensor::InsertSliceOp>(
+          tiledOp->getLoc(), yieldedResult.value(),
+          innerNewBBArgs[yieldedResult.index()], offsets, sizes, strides));
+    }
+    return newYieldVals;
+  };
+  scf::ForOp newInnerMostLoop =
+      replaceLoopWithNewYields(b, innerMostLoop, newIterOperands, innerFn);
+  tileLoopOps.back() = newInnerMostLoop;
+
+  // Propagate the newly yielded values back up the loop nest.
+  for (unsigned loopDepth :
+       llvm::reverse(llvm::seq<unsigned>(0, tileLoopOps.size() - 1))) {
+    NewYieldValueFn fn = [&](OpBuilder &innerBuilder, Location innerLoc,
+                             ArrayRef<BlockArgument> innerNewBBArgs) {
+      return llvm::to_vector(
+          llvm::map_range(tileLoopOps[loopDepth + 1].getResults().take_back(
+                              newIterOperands.size()),
+                          [](OpResult r) -> Value { return r; }));
+    };
+    tileLoopOps[loopDepth] = replaceLoopWithNewYields(b, tileLoopOps[loopDepth],
+                                                      newIterOperands, fn);
+  }
+
+  // The `replacements` uses values returned by the tile loop nest. Update these
+  // to use values from the new loop nests created here. Build a map from result
+  // of the previous outermost loop to the new outermost loop.
+  llvm::SmallDenseMap<Value, Value> resultValRemap;
+  scf::ForOp newOuterMostLoop = tileLoopOps.front();
+  for (auto origResult : llvm::enumerate(outerMostLoop->getResults())) {
+    resultValRemap[origResult.value()] =
+        newOuterMostLoop->getResult(origResult.index());
+  }
+  for (auto &replacement : replacements) {
+    auto newReplacement = resultValRemap.lookup(replacement.second);
+    if (newReplacement)
+      replacement.second = newReplacement;
+  }
+  for (auto fusedOpResult : llvm::enumerate(returnedFusedOpResults)) {
+    replacements[fusedOpResult.value()] = newOuterMostLoop.getResult(
+        outerMostLoop.getNumResults() + fusedOpResult.index());
+  }
+}
+
 FailureOr<LinalgOp> TileLoopNest::fuseProducer(OpBuilder &b,
-                                               OpOperand *consumerOpOperand) {
+                                               OpOperand *consumerOpOperand,
+                                               bool returnFusedOpValues) {
   // Check if the consumer has been tiled before. For example, it may not have
   // been tiled if the outermost tile loop is a reduction loop.
   if (tiledRootAndFusedOpsLoops.count(consumerOpOperand->getOwner()) == 0)
     return failure();
 
   assert(this->isValid() &&
          "expect the tile loop nest to satisfy all invariants");
 
FailureOr<LinalgOp> TileLoopNest::fuseProducer(OpBuilder &b,
   BlockArgument bbArg = consumerOp.getTiedBlockArgument(consumerOpOperand);
   if (bbArg.getUses().empty())
     return failure();
 
   // Check if the producer is a LinalgOp possibly passed by iteration argument.
   OpOperand *iterArg = nullptr;
   auto producerResult = sliceOp.getSource().dyn_cast<OpResult>();
   if (auto bbArg = sliceOp.getSource().dyn_cast<BlockArgument>()) {
+    if (bbArg.getParentBlock() != sliceOp->getBlock())
+      return failure();
     iterArg = getTiedIterArg(bbArg);
     // Check the iteration argument may be used to pass in the producer output.
     if (!iterArg || hasOtherUses(bbArg, sliceOp))
       return failure();
     producerResult = iterArg->get().dyn_cast<OpResult>();
   }
-  if (!producerResult || !isa<LinalgOp>(producerResult.getOwner()))
+  if (!producerResult)
+    return failure();
+  LinalgOp producer = dyn_cast<LinalgOp>(producerResult.getOwner());
+  if (!producer)
     return failure();
 
   // Compute the tiled producer slice dimensions given the tiled consumer loops.
   SmallVector<int64_t> tiledSliceDimIndices = getTiledSliceDims(
       consumerOpOperand, tiledRootAndFusedOpsLoops[consumerOp]);
   if (tiledSliceDimIndices.empty())
     return failure();
 
   // Compute the tiled producer loop indices.
   SmallVector<int64_t> tiledProducerLoopIndices =
       getTiledProducerLoops(producerResult, tiledSliceDimIndices);
 
   // Tile the producer operands and clone the producer in place of `sliceOp`.
   LinalgOp clonedOp =
       getTiledProducer(b, producerResult, sliceOp, tiledSliceDimIndices,
                        tiledProducerLoopIndices, iterArg);
   tiledRootAndFusedOpsLoops[clonedOp] = tiledProducerLoopIndices;
+  OpBuilder::InsertionGuard guard(b);
+  b.setInsertionPointAfter(clonedOp);
 
   // Cast the `clonedOp` result to gap type mismatches before canonicalization.
   Type consumerOperandType = consumerOpOperand->get().getType();
   Value newResult = clonedOp->getResult(producerResult.getResultNumber());
   if (newResult.getType() != consumerOperandType) {
-    OpBuilder::InsertionGuard guard(b);
-    b.setInsertionPointAfter(clonedOp);
     newResult = b.create<tensor::CastOp>(producerResult.getLoc(),
                                          consumerOperandType, newResult);
   }
 
   // Replace the `sliceOp` uses except for the `clonedOp` output uses.
   sliceOp.getResult().replaceAllUsesExcept(newResult, clonedOp);
+
+  if (returnFusedOpValues) {
+    fusedOps.push_back(producer);
+    yieldFusedValues(b, producer, clonedOp, sliceOp.getMixedOffsets(),
+                     sliceOp.getMixedSizes(), sliceOp.getMixedStrides());
+  }
   return clonedOp;
 }
 
-ValueRange TileLoopNest::getRootOpReplacementResults() {
-  assert(!isEmpty() && "expect tile loop nest to be non-empty");
-  return tileLoopOps.front()->getOpResults();
+Value TileLoopNest::getUntiledOpResultReplacement(Value v) {
+  return replacements.lookup(v);
 }
 
 SmallVector<LinalgOp> TileLoopNest::getAllTiledAndFusedOps() {
   SmallVector<LinalgOp> result;
   for (const auto &kvp : tiledRootAndFusedOpsLoops) {
     auto linalgOp = dyn_cast<LinalgOp>(kvp.getFirst());
     assert(linalgOp &&
            "expect all tiled and fused operations are linalg operations");
     result.push_back(linalgOp);
   }
   return result;
 }
 
 //===----------------------------------------------------------------------===//
 // Tile and fuse entry-points.
 //===----------------------------------------------------------------------===//
 
 FailureOr<TileLoopNest> mlir::linalg::tileConsumerAndFuseProducers(
     OpBuilder &b, LinalgOp consumerOp, ArrayRef<int64_t> tileSizes,
     ArrayRef<int64_t> tileInterchange,
-    const Optional<LinalgLoopDistributionOptions> &tileDistribution) {
+    const Optional<LinalgLoopDistributionOptions> &tileDistribution,
+    bool returnFusedOpValues) {
   assert(tileSizes.size() == tileInterchange.size() &&
          "expect the number of tile sizes and interchange dims to match");
   assert(isPermutation(tileInterchange) &&
          "expect tile interchange is a permutation");
 
   // Create an empty tile loop nest.
   TileLoopNest tileLoopNest(consumerOp);
 
   // Search the number of outer parallel loops to separate them from possible
   // inner reduction dimensions.
   SmallVector<StringAttr> iterTypes =
       llvm::to_vector<6>(consumerOp.iterator_types().getAsRange<StringAttr>());
   applyPermutationToVector(iterTypes, tileInterchange);
   auto *it = find_if(iterTypes, [&](StringAttr iterType) {
     return !isParallelIterator(iterType);
   });
   int64_t split = std::distance(iterTypes.begin(), it);
 
   // Helper to fuse the producers greedily using a queue of fusion candidates.
-  auto fuseProducersGreedily = [&](ArrayRef<OpOperand *> operands) {
+  auto fuseProducersGreedily = [&](ArrayRef<OpOperand *> operands,
+                                   bool returnFusedOpValues) {
     SmallVector<OpOperand *> candidates(operands.begin(), operands.end());
     while (!candidates.empty()) {
-      FailureOr<LinalgOp> fusedProducer =
-          tileLoopNest.fuseProducer(b, candidates.pop_back_val());
+      FailureOr<LinalgOp> fusedProducer = tileLoopNest.fuseProducer(
+          b, candidates.pop_back_val(), returnFusedOpValues);
       if (failed(fusedProducer))
         continue;
       candidates.append(fusedProducer->getInputAndOutputOperands());
     }
   };
 
   // Perform tiling and fusion in two steps. We need to respect the loop
   // interchange here; filter parellel dimensions based on their order *after*
   // permutation but pass in the original configuration *before* permuation,
   // given the tiling and interchange happen together.
   SmallVector<int64_t> outerTileSizes(tileSizes.size(), 0);
   SmallVector<int64_t> innerTileSizes(tileSizes.size(), 0);
   for (int64_t i : tileInterchange.take_front(split))
     outerTileSizes[i] = tileSizes[i];
   for (int64_t i : tileInterchange.drop_front(split))
     innerTileSizes[i] = tileSizes[i];
 
   // Tile the outer parallel loops and fuse the output operands.
   if (failed(tileLoopNest.tileRootOp(b, outerTileSizes, tileInterchange,
                                      tileDistribution)))
     return failure();
-  fuseProducersGreedily(tileLoopNest.getRootOp().getOutputOperands());
+  if (returnFusedOpValues) {
+    fuseProducersGreedily(tileLoopNest.getRootOp().getInputAndOutputOperands(),
+                          true);
+  } else {
+    fuseProducersGreedily(tileLoopNest.getRootOp().getOutputOperands(), false);
+  }
 
   // Tile the remaining loops and fuse the input operands.
   if (failed(tileLoopNest.tileRootOp(b, innerTileSizes, tileInterchange,
                                      tileDistribution)))
     return failure();
-  fuseProducersGreedily(tileLoopNest.getRootOp().getInputOperands());
+  fuseProducersGreedily(tileLoopNest.getRootOp().getInputOperands(), false);
 
   // Exit if the tile loop nest is empty since all tile sizes are zero.
   if (tileLoopNest.isEmpty())
     return failure();
 
   return tileLoopNest;
 }

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

return rewriter.notifyMatchFailure(
         op, "expect the number of tile sizes and interchange dims to match");
 
   // Copy the `tileSizes` and `tileInterchange` prefixes needed for `rootOp`.
   SmallVector<int64_t> rootTileSizes(options.tileSizes.begin(),
                                      options.tileSizes.begin() +
                                          rootOp.getNumLoops());
   SmallVector<int64_t> rootInterchange =
       options.tileInterchange.empty()
-          ? llvm::to_vector<6>(llvm::seq<int64_t>(0, rootOp.getNumLoops()))
+          ? llvm::to_vector(llvm::seq<int64_t>(0, rootOp.getNumLoops()))
           : SmallVector<int64_t>(options.tileInterchange.begin(),
                                  options.tileInterchange.begin() +
                                      rootOp.getNumLoops());
 
   // Check `rootTileSizes` contains non-zero tile sizes.
   if (llvm::count(rootTileSizes, 0) == static_cast<long>(rootTileSizes.size()))
     return rewriter.notifyMatchFailure(
         op, "expect at least one non-zero tile size");
 
   // Check `rootInterchange` is a permutation of the `rootOp` loop dimensions.
   // It has to be a permutation since the tiling cannot tile the same loop
   // dimension multiple times.
   if (!isPermutation(rootInterchange))
     return rewriter.notifyMatchFailure(
         op, "expect the tile interchange permutes the root loops");
 
   // Tile `rootOp` and fuse its producers.
-  FailureOr<TileLoopNest> tileLoopNest =
-      tileConsumerAndFuseProducers(rewriter, rootOp, rootTileSizes,
-                                   rootInterchange, options.tileDistribution);
+  FailureOr<TileLoopNest> tileLoopNest = tileConsumerAndFuseProducers(
+      rewriter, rootOp, rootTileSizes, rootInterchange,
+      options.tileDistribution, options.returnFusedOpValues);
+
   if (failed(tileLoopNest))
     return rewriter.notifyMatchFailure(
         op, "tileConsumerAndFuseProducers failed unexpectedly");
 
-  // Replace all uses of the tiled loop operation.
-  rootOp->replaceAllUsesWith(tileLoopNest->getRootOpReplacementResults());
+  // Replace uses of the tiled operation(s) which are
+  // - Not within the loop itself (that creates an SSA use-def violation)
+  // - Ignore `tensor.dim` uses. These `tensor.dim` operations might be
+  //   used in computation of the loop bounds themselves. Replacing these uses
+  //   causes use-def violations. These uses should be handled using
+  //   `ReifyRankedShapedTypeOpInterface`.
+  scf::ForOp outerMostLoop = tileLoopNest->getLoopOps().front();
+  auto useOutsideLoopFn = [&outerMostLoop](OpOperand &use) {
+    Operation *user = use.getOwner();
+    return !isa<tensor::DimOp>(user) &&
+           !outerMostLoop->isAncestor(use.getOwner());
+  };
+  for (Value rootResult : rootOp->getResults()) {
+    if (Value replacement =
+            tileLoopNest->getUntiledOpResultReplacement(rootResult))
+      rootResult.replaceUsesWithIf(replacement, useOutsideLoopFn);
+  }
+  if (options.returnFusedOpValues) {
+    for (auto fusedOp : tileLoopNest->getFusedOps()) {
+      for (auto fusedOpResult : fusedOp->getResults()) {
+        if (Value replacement =
+                tileLoopNest->getUntiledOpResultReplacement(fusedOpResult))
+          fusedOpResult.replaceUsesWithIf(replacement, useOutsideLoopFn);
+      }
+    }
+  }
 
   // Apply the filter if specified.
-  for (LinalgOp linalgOp : tileLoopNest->getAllTiledAndFusedOps())
+  for (LinalgOp linalgOp : tileLoopNest->getAllTiledAndFusedOps()) {
     filter.replaceLinalgTransformationFilter(rewriter, linalgOp);
+  }
   return tileLoopNest;
 }
 
 /// Linalg generic interchange pattern.
 mlir::linalg::GenericOpInterchangePattern::GenericOpInterchangePattern(
     MLIRContext *context, ArrayRef<unsigned> interchangeVector,
     LinalgTransformationFilter f, PatternBenefit benefit)
     : OpRewritePattern(context, benefit), filter(std::move(f)),

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

+// RUN: mlir-opt %s  -test-linalg-tile-and-fuse-on-tensors -cse --split-input-file | FileCheck %s
+
+func.func @matmul_bias_add(%lhs : tensor<?x?xf32>, %rhs : tensor<?x?xf32>, %bias : tensor<?xf32>)
+    -> (tensor<?x?xf32>, tensor<?x?xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %cst = arith.constant 0.0 : f32
+  %d0 = tensor.dim %lhs, %c0 : tensor<?x?xf32>
+  %d1 = tensor.dim %rhs, %c1 : tensor<?x?xf32>
+  %init = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32>
+  %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<?x?xf32>) -> tensor<?x?xf32>
+  %matmul = linalg.matmul ins(%lhs, %rhs : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%fill : tensor<?x?xf32>) -> tensor<?x?xf32>
+  %bias_add = linalg.generic {
+      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                       affine_map<(d0, d1) -> (d1)>,
+                       affine_map<(d0, d1) -> (d0, d1)>],
+      iterator_types = ["parallel", "parallel"],
+      __internal_linalg_transform__ = "return_fused_values"}
+      ins(%matmul, %bias : tensor<?x?xf32>, tensor<?xf32>)
+      outs(%init : tensor<?x?xf32>) {
+    ^bb0(%b0 : f32, %b1 : f32, %b3 : f32):
+      %add = arith.addf %b0, %b1 : f32
+      linalg.yield %add : f32
+    } -> tensor<?x?xf32>
+  return %matmul, %bias_add : tensor<?x?xf32>, tensor<?x?xf32>
+}
+// CHECK-LABEL: func @matmul_bias_add
+//  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<?xf32>
+//   CHECK-DAG:   %[[INIT:.+]] = linalg.init_tensor
+//       CHECK:   %[[OUTER:.+]]:2 = scf.for %[[IV0:.+]] = %{{.+}} to %{{.+}} step
+//  CHECK-SAME:       iter_args(%[[OUTER_ITER0:[a-zA-Z0-9]+]] = %[[INIT]],
+//  CHECK-SAME:       %[[OUTER_ITER1:[a-zA-Z0-9]+]] = %[[INIT]])
+//       CHECK:     %[[INNER:.+]]:2 = scf.for %[[IV1:.+]] = %{{.+}} to %[[UB1:.+]] step
+//  CHECK-SAME:       iter_args(%[[INNER_ITER0:[a-zA-Z0-9]+]] = %[[OUTER_ITER0]],
+//  CHECK-SAME:       %[[INNER_ITER1:[a-zA-Z0-9]+]] = %[[OUTER_ITER1]])
+//   CHECK-DAG:       %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0]
+//   CHECK-DAG:       %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]]
+//   CHECK-DAG:       %[[OUTS_TILE:.+]] = tensor.extract_slice %[[INNER_ITER1]][%[[IV0]], %[[IV1]]]
+//   CHECK-DAG:       %[[FILL_TILE:.+]] = linalg.fill
+//  CHECK-SAME:           outs(%[[OUTS_TILE]] :
+//   CHECK-DAG:       %[[MATMUL_TILE:.+]] = linalg.matmul
+//  CHECK-SAME:           ins(%[[LHS_TILE]], %[[RHS_TILE]] :
+//  CHECK-SAME:           outs(%[[FILL_TILE:.+]] :
+//   CHECK-DAG:       %[[BIAS_TILE:.+]] = tensor.extract_slice %[[ARG2]][%[[IV1]]]
+//   CHECK-DAG:       %[[OUTS2_TILE:.+]] = tensor.extract_slice %[[INNER_ITER0]][%[[IV0]], %[[IV1]]]
+//       CHECK:       %[[ROOT_TILE:.+]] = linalg.generic
+//  CHECK-SAME:           ins(%[[MATMUL_TILE]], %[[BIAS_TILE]] :
+//  CHECK-SAME:           outs(%[[OUTS2_TILE]] :
+//   CHECK-DAG:       %[[ROOT_INSERT:.+]] = tensor.insert_slice %[[ROOT_TILE]] into %[[INNER_ITER0]][%[[IV0]], %[[IV1]]]
+//   CHECK-DAG:       %[[MATMUL_INSERT:.+]] = tensor.insert_slice %[[MATMUL_TILE]] into %[[INNER_ITER1]][%[[IV0]], %[[IV1]]]
+//       CHECK:       scf.yield %[[ROOT_INSERT]], %[[MATMUL_INSERT]]
+//       CHECK:     scf.yield %[[INNER]]#0, %[[INNER]]#1
+//       CHECK:   return %[[OUTER]]#1, %[[OUTER]]#0
+
+// -----
+
+func.func @matmul_lhs_fusion(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
+    %AB_init: tensor<?x?xf32>, %C: tensor<?x?xf32>, %ABC_init: tensor<?x?xf32>)
+     -> (tensor<?x?xf32>, tensor<?x?xf32>) {
+  %AB = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
+    outs(%AB_init : tensor<?x?xf32>) -> tensor<?x?xf32>   // <MxN1> <N1xN2>
+  %ABC = linalg.matmul {__internal_linalg_transform__ = "lhs_fusion"}
+    ins(%AB, %C : tensor<?x?xf32>, tensor<?x?xf32>)
+    outs(%ABC_init : tensor<?x?xf32>) -> tensor<?x?xf32>   // <MxN2> <N2xN3>
+  return %AB, %ABC : tensor<?x?xf32>, tensor<?x?xf32>
+}
+// CHECK-LABEL: func @matmul_lhs_fusion
+//  CHECK-SAME:     %[[A:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//  CHECK-SAME:     %[[B:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//  CHECK-SAME:     %[[AB_INIT:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//  CHECK-SAME:     %[[C:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//  CHECK-SAME:     %[[ABC_INIT:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//       CHECK:   %[[LOOP:.+]]:2 = scf.for %[[IV0:[a-zA-Z0-9]+]]
+//  CHECK-SAME:       iter_args(%[[ITER0:[a-zA-Z0-9]+]] = %[[ABC_INIT]],
+//  CHECK-SAME:       %[[ITER1:[a-zA-Z0-9]+]] = %[[AB_INIT]])
+//   CHECK-DAG:     %[[LHS_TILE:.+]] = tensor.extract_slice %[[A]][%[[IV0]], 0]
+//   CHECK-DAG:     %[[OUTS0_TILE:.+]] = tensor.extract_slice %[[ITER1]][%[[IV0]], 0]
+//   CHECK-DAG:     %[[GEMM0_TILE:.+]] = linalg.matmul
+//  CHECK-SAME:         ins(%[[LHS_TILE]], %[[B]] :
+//  CHECK-SAME:         outs(%[[OUTS0_TILE]] :
+//   CHECK-DAG:     %[[OUTS1_TILE:.+]] = tensor.extract_slice %[[ITER0]][%[[IV0]], 0]
+//   CHECK-DAG:     %[[GEMM1_TILE:.+]] = linalg.matmul
+//  CHECK-SAME:         ins(%[[GEMM0_TILE]], %[[C]] :
+//  CHECK-SAME:         outs(%[[OUTS1_TILE:.+]] :
+//   CHECK-DAG:     %[[GEMM1_INSERT:.+]] = tensor.insert_slice %[[GEMM1_TILE]] into %[[ITER0]][%[[IV0]], 0]
+//   CHECK-DAG:     %[[GEMM0_INSERT:.+]] = tensor.insert_slice %[[GEMM0_TILE]] into %[[ITER1]][%[[IV0]], 0]
+//       CHECK:     scf.yield %[[GEMM1_INSERT]], %[[GEMM0_INSERT]]
+//       CHECK:   return %[[LOOP]]#1, %[[LOOP]]#0
+
+// -----
+
+func.func @matmul_rhs_fusion(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
+    %AB_init: tensor<?x?xf32>, %C: tensor<?x?xf32>, %ABC_init: tensor<?x?xf32>)
+    -> (tensor<?x?xf32>, tensor<?x?xf32>) {
+  %AB = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
+    outs(%AB_init : tensor<?x?xf32>) -> tensor<?x?xf32>   // <MxN1> <N1xN2>
+  %ABC = linalg.matmul {__internal_linalg_transform__ = "rhs_fusion"}
+    ins(%C, %AB : tensor<?x?xf32>, tensor<?x?xf32>)
+    outs(%ABC_init : tensor<?x?xf32>) -> tensor<?x?xf32>   // <MxN2> <N2xN3>
+  return %AB, %ABC : tensor<?x?xf32>, tensor<?x?xf32>
+}
+// CHECK-LABEL: func @matmul_rhs_fusion
+//  CHECK-SAME:     %[[A:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//  CHECK-SAME:     %[[B:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//  CHECK-SAME:     %[[AB_INIT:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//  CHECK-SAME:     %[[C:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//  CHECK-SAME:     %[[ABC_INIT:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//       CHECK:   %[[LOOP:.+]]:2 = scf.for %[[IV0:[a-zA-Z0-9]+]]
+//  CHECK-SAME:       iter_args(%[[ITER0:[a-zA-Z0-9]+]] = %[[ABC_INIT]],
+//  CHECK-SAME:       %[[ITER1:[a-zA-Z0-9]+]] = %[[AB_INIT]])
+//   CHECK-DAG:     %[[RHS_TILE:.+]] = tensor.extract_slice %[[B]][0, %[[IV0]]]
+//   CHECK-DAG:     %[[OUTS0_TILE:.+]] = tensor.extract_slice %[[ITER1]][0, %[[IV0]]]
+//   CHECK-DAG:     %[[GEMM0_TILE:.+]] = linalg.matmul
+//  CHECK-SAME:         ins(%[[A]], %[[RHS_TILE]] :
+//  CHECK-SAME:         outs(%[[OUTS0_TILE]] :
+//   CHECK-DAG:     %[[OUTS1_TILE:.+]] = tensor.extract_slice %[[ITER0]][0, %[[IV0]]]
+//   CHECK-DAG:     %[[GEMM1_TILE:.+]] = linalg.matmul
+//  CHECK-SAME:         ins(%[[C]], %[[GEMM0_TILE]] :
+//  CHECK-SAME:         outs(%[[OUTS1_TILE:.+]] :
+//   CHECK-DAG:     %[[GEMM1_INSERT:.+]] = tensor.insert_slice %[[GEMM1_TILE]] into %[[ITER0]][0, %[[IV0]]]
+//   CHECK-DAG:     %[[GEMM0_INSERT:.+]] = tensor.insert_slice %[[GEMM0_TILE]] into %[[ITER1]][0, %[[IV0]]]
+//       CHECK:     scf.yield %[[GEMM1_INSERT]], %[[GEMM0_INSERT]]
+//       CHECK:   return %[[LOOP]]#1, %[[LOOP]]#0
+
+// -----
+
+func.func @matmul_out_fusion(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %cst = arith.constant 0.0 : f32
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+  %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
+  %init = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32>
+  %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<?x?xf32>) -> tensor<?x?xf32>
+  %gemm = linalg.matmul {__internal_linalg_transform__ = "matmul_outs_fusion"}
+      ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%fill : tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %gemm : tensor<?x?xf32>
+}
+// CHECK-LABEL: func @matmul_out_fusion
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//   CHECK-DAG:   %[[INIT:.+]] = linalg.init_tensor
+//       CHECK:   %[[OUTER:.+]] = scf.for %[[IV0:.+]] = %{{.+}} to %{{.+}} step
+//  CHECK-SAME:       iter_args(%[[OUTER_ITER:[a-zA-Z0-9]+]] = %[[INIT]])
+//       CHECK:     %[[INNER:.+]] = scf.for %[[IV1:.+]] = %{{.+}} to %[[UB1:.+]] step
+//  CHECK-SAME:       iter_args(%[[INNER_ITER:[a-zA-Z0-9]+]] = %[[OUTER_ITER]])
+//   CHECK-DAG:       %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0]
+//   CHECK-DAG:       %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]]
+//   CHECK-DAG:       %[[OUTS_TILE:.+]] = tensor.extract_slice %[[INNER_ITER]][%[[IV0]], %[[IV1]]]
+//   CHECK-DAG:       %[[FILL_TILE:.+]] = linalg.fill
+//  CHECK-SAME:           outs(%[[OUTS_TILE]] :
+//   CHECK-DAG:       %[[MATMUL_TILE:.+]] = linalg.matmul
+//  CHECK-SAME:           ins(%[[LHS_TILE]], %[[RHS_TILE]] :
+//  CHECK-SAME:           outs(%[[FILL_TILE:.+]] :
+//   CHECK-DAG:       %[[MATMUL_INSERT:.+]] = tensor.insert_slice %[[MATMUL_TILE]] into %[[INNER_ITER]][%[[IV0]], %[[IV1]]]
+//       CHECK:       scf.yield %[[MATMUL_INSERT]]
+//       CHECK:     scf.yield %[[INNER]]
+//       CHECK:   return %[[OUTER]]
+
+// -----
+
+func.func @generic_plus_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2 : tensor<?xf32>)
+    -> (tensor<?x?xf32>, tensor<?x?xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+  %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
+  %init = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32>
+  %bias = linalg.generic {
+      indexing_maps = [affine_map<(m, n) -> (n)>, affine_map<(m, n) -> (m, n)>],
+      iterator_types = ["parallel", "parallel"]}
+      ins(%arg2 : tensor<?xf32>) outs(%init: tensor<?x?xf32>) {
+    ^bb(%b0: f32, %b1: f32) :
+      linalg.yield %b0 : f32
+    } -> tensor<?x?xf32>
+  %gemm = linalg.matmul {__internal_linalg_transform__ = "matmul_outs_fusion"}
+    ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+    outs(%bias : tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %bias, %gemm : tensor<?x?xf32>, tensor<?x?xf32>
+}
+// CHECK-LABEL: func @generic_plus_matmul
+//  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<?xf32>
+//   CHECK-DAG:   %[[INIT:.+]] = linalg.init_tensor
+//       CHECK:   %[[OUTER:.+]]:2 = scf.for %[[IV0:.+]] = %{{.+}} to %{{.+}} step
+//  CHECK-SAME:       iter_args(%[[OUTER_ITER0:[a-zA-Z0-9]+]] = %[[INIT]],
+//  CHECK-SAME:       %[[OUTER_ITER1:[a-zA-Z0-9]+]] = %[[INIT]])
+//       CHECK:     %[[INNER:.+]]:2 = scf.for %[[IV1:.+]] = %{{.+}} to %[[UB1:.+]] step
+//  CHECK-SAME:       iter_args(%[[INNER_ITER0:[a-zA-Z0-9]+]] = %[[OUTER_ITER0]],
+//  CHECK-SAME:       %[[INNER_ITER1:[a-zA-Z0-9]+]] = %[[OUTER_ITER1]])
+//   CHECK-DAG:       %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0]
+//   CHECK-DAG:       %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]]
+//   CHECK-DAG:       %[[OUTS_TILE:.+]] = tensor.extract_slice %[[INNER_ITER0]][%[[IV0]], %[[IV1]]]
+//   CHECK-DAG:       %[[BIAS_TILE:.+]] = tensor.extract_slice %[[ARG2]][%[[IV1]]]
+//       CHECK:       %[[BCAST_TILE:.+]] = linalg.generic
+//  CHECK-SAME:           ins(%[[BIAS_TILE]] :
+//  CHECK-SAME:           outs(%[[OUTS_TILE]] :
+//   CHECK-DAG:       %[[MATMUL_TILE:.+]] = linalg.matmul
+//  CHECK-SAME:           ins(%[[LHS_TILE]], %[[RHS_TILE]] :
+//  CHECK-SAME:           outs(%[[BCAST_TILE:.+]] :
+//   CHECK-DAG:       %[[MATMUL_INSERT:.+]] = tensor.insert_slice %[[MATMUL_TILE]] into %[[INNER_ITER0]][%[[IV0]], %[[IV1]]]
+//   CHECK-DAG:       %[[BCAST_INSERT:.+]] = tensor.insert_slice %[[BCAST_TILE]] into %[[INNER_ITER1]][%[[IV0]], %[[IV1]]]
+//       CHECK:       scf.yield %[[MATMUL_INSERT]], %[[BCAST_INSERT]]
+//       CHECK:     scf.yield %[[INNER]]#0, %[[INNER]]#1
+//       CHECK:   return %[[OUTER]]#1, %[[OUTER]]#0

mlir/test/lib/Dialect/Linalg/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRLinalgTestPasses
   TestLinalgDecomposeOps.cpp
   TestLinalgElementwiseFusion.cpp
+  TestLinalgFusionOnTensors.cpp
   TestLinalgFusionTransforms.cpp
   TestLinalgHoisting.cpp
   TestLinalgTransforms.cpp
   TestPadFusion.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC

mlir/test/lib/Dialect/Linalg/TestLinalgFusionOnTensors.cpp

+//===- TestLinalgFusionOnTensorsPass.cpp - Test Linalg fusion patterns ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements logic for testing Linalg fusion patterns on tensors.
+// Specifically the `tileConsumerAndFuseProducers` method.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/SCF/Transforms.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/Passes.h"
+
+using namespace mlir;
+using namespace mlir::linalg;
+
+namespace {
+struct TestLinalgFusionOnTensorsPass
+    : public PassWrapper<TestLinalgFusionOnTensorsPass,
+                         OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgFusionOnTensorsPass)
+
+  StringRef getArgument() const final {
+    return "test-linalg-tile-and-fuse-on-tensors";
+  }
+
+  StringRef getDescription() const final {
+    return "Test Linalg tiling and fusion on tensor operations.";
+  }
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AffineDialect, linalg::LinalgDialect, tensor::TensorDialect,
+                    scf::SCFDialect>();
+  }
+  TestLinalgFusionOnTensorsPass() = default;
+  TestLinalgFusionOnTensorsPass(const TestLinalgFusionOnTensorsPass &pass) {}
+
+  void runOnOperation() override {
+    MLIRContext *context = &this->getContext();
+    func::FuncOp funcOp = this->getOperation();
+
+    RewritePatternSet fusionPatterns(context);
+
+    fusionPatterns.insert<LinalgTileAndFuseTensorOpsPattern>(
+        context,
+        LinalgTilingAndFusionOptions()
+            .setTileSizes({10, 20})
+            .setReturnFusedOpValues(true),
+        LinalgTransformationFilter(
+            StringAttr::get(context, "return_fused_values"),
+            StringAttr::get(context, "fused_ops")));
+
+    fusionPatterns.insert<LinalgTileAndFuseTensorOpsPattern>(
+        context,
+        LinalgTilingAndFusionOptions()
+            .setTileSizes({10, 0, 0})
+            .setReturnFusedOpValues(true),
+        LinalgTransformationFilter(StringAttr::get(context, "lhs_fusion"),
+                                   StringAttr::get(context, "fused_ops")));
+
+    fusionPatterns.insert<LinalgTileAndFuseTensorOpsPattern>(
+        context,
+        LinalgTilingAndFusionOptions()
+            .setTileSizes({0, 20, 0})
+            .setReturnFusedOpValues(true),
+        LinalgTransformationFilter(StringAttr::get(context, "rhs_fusion"),
+                                   StringAttr::get(context, "fused_ops")));
+
+    fusionPatterns.insert<LinalgTileAndFuseTensorOpsPattern>(
+        context,
+        LinalgTilingAndFusionOptions()
+            .setTileSizes({10, 20, 0})
+            .setReturnFusedOpValues(true),
+        LinalgTransformationFilter(
+            StringAttr::get(context, "matmul_outs_fusion"),
+            StringAttr::get(context, "fused_ops")));
+
+    (void)applyPatternsAndFoldGreedily(funcOp, std::move(fusionPatterns));
+  }
+};
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestLinalgFusionOnTensorsPass() {
+  PassRegistration<TestLinalgFusionOnTensorsPass>();
+}
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestIntRangeInference();
 void registerTestIRVisitorsPass();
 void registerTestGenericIRVisitorsPass();
 void registerTestGenericIRVisitorsInterruptPass();
 void registerTestInterfaces();
 void registerTestLastModifiedPass();
 void registerTestLinalgDecomposeOps();
 void registerTestLinalgElementwiseFusion();
+void registerTestLinalgFusionOnTensorsPass();
 void registerTestLinalgGreedyFusion();
 void registerTestLinalgHoisting();
 void registerTestLinalgTransforms();
 void registerTestLivenessPass();
 void registerTestLoopFusion();
 void registerTestLoopMappingPass();
 void registerTestLoopUnrollingPass();
 void registerTestMatchReductionPass();
#endif
   mlir::test::registerTestMultiBuffering();
   mlir::test::registerTestIntRangeInference();
   mlir::test::registerTestIRVisitorsPass();
   mlir::test::registerTestGenericIRVisitorsPass();
   mlir::test::registerTestInterfaces();
   mlir::test::registerTestLastModifiedPass();
   mlir::test::registerTestLinalgDecomposeOps();
   mlir::test::registerTestLinalgElementwiseFusion();
+  mlir::test::registerTestLinalgFusionOnTensorsPass();
   mlir::test::registerTestLinalgGreedyFusion();
   mlir::test::registerTestLinalgHoisting();
   mlir::test::registerTestLinalgTransforms();
   mlir::test::registerTestLivenessPass();
   mlir::test::registerTestLoopFusion();
   mlir::test::registerTestLoopMappingPass();
   mlir::test::registerTestLoopUnrollingPass();
   mlir::test::registerTestMatchReductionPass();