[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.

Depends D125147

[mravishankar]

Add patch dependency.

[nicolasvasilache]

A first round of comments to improve the doc and help digest the complex implementation of the function.
I suspect it will need to be broken into smaller more composable pieces.

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

550

type: Interchange

568

typo: construct

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

349–358

Please rephrase, clarify and use better variable names.

Assuming `tiledOp` is the tiled version of `originalOp` and `mixedXXX` are YYY.
Assuming `tiledOp` is fused (or do we want to fuse it at this point?) into the current loop nest.
Update `tileLoopOps` with (one?) additional iter_args and (one?) additional yield operand such that ZZZ.

Example:
IR

352

the name isn't descriptive enough for this seemingly complicated process.

357

Better description plz.

Returns the OpResult of the outermost loop that corresponds to the original operation before it has been tiled for fusion.
(Any single result assumption here?).

404

You have access to: 1. the top-level loop, 2. fusedOps and 3. replacements.
IIUC 3. can be recovered from 1+2 and 2 can be recovered from 1+3.

Can we make one or the other an accessor that computes the information rather than storing it?
This will reduce opportunities for errors with them going out of sync.

410

I am concerned that we are making this already very complex function even more complex ..

Side note, do you still need tileDistribution in this function or could we get rid of it (in a followup).

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

571–576

You can rewrite as a 1-liner: fuseProducersGreedily(..., returnFusedOpValues)

[mravishankar]

Rebase and address comments.

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

349–358

Added IR example near the implementation.

352

Happy to rename to any thing that is more descriptive. This is essentially what this method does. (Also, its the bare minimum I could think of to make fusion + scf.for + tensors work.)

357

No single result assumption here.

404

I dont think so. You need to know both 2, and 3 to be able to map which result maps to a replacement of which of the fused op.

410

For a fusion that is generating scf.for with tensors, I cant see a simpler implementation than this. If we move away from using scf.for in general, then this would be much easier.

I do need the tileDistribution here. Essentially there is a custom tileConsumerAndFuseProducer method in IREE that I am trying to replace. That needs distribution as well (with the replacement for the fused ops where possible, without which fusion is useless since the original op ends up in the generated executable anyway). When IREE moves away from using scf.for + tensors, we can drop the distribution aspect of this. Still think the change in this patch is required for completeness (unless we drop support of fusion + scf.for + tensors completely).

[nicolasvasilache]

I think this is now subsumed by the TilingInterface ?
Unsubscribing, plz readd me if I missed something.

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

568

typo: to to

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

658

spurious change

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.source().dyn_cast<OpResult>();
   if (auto bbArg = sliceOp.source().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());
     }
   };
 
   // Tile the outer parallel loops and fuse the output operands.
   SmallVector<int64_t> outerTileSizes;
   outerTileSizes.append(tileSizes.begin(), tileSizes.begin() + split);
   outerTileSizes.append(tileSizes.size() - split, 0);
   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.
   SmallVector<int64_t> innerTileSizes;
   innerTileSizes.append(split, 0);
   innerTileSizes.append(tileSizes.begin() + split, tileSizes.end());
   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-tensor-fusion-transform-patterns -resolve-shaped-type-result-dims -canonicalize -cse --split-input-file | FileCheck %s
+// RUN: mlir-opt %s  -test-linalg-tile-and-fuse-on-tensors -cse --split-input-file | FileCheck %s
 
-module {
-  func.func @matmul_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> {
-    %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 %ABC : tensor<?x?xf32>
-  }
-}
-//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 32)>
-//  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 16)>
-//  CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 64)>
-//  CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0)[s0, s1] -> (-d0 + s1, -d0 + s0, 32)>
-
-//      CHECK: func @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-DAG:   %[[C0:.+]] = arith.constant 0 : index
-//  CHECK-DAG:   %[[C1:.+]] = arith.constant 1 : index
-//  CHECK-DAG:   %[[C32:.+]] = arith.constant 32 : index
-//  CHECK-DAG:   %[[C64:.+]] = arith.constant 64 : index
-//  CHECK-DAG:   %[[C16:.+]] = arith.constant 16 : index
-//  CHECK-DAG:   %[[M:.+]] = tensor.dim %[[ARG0]], %[[C0]]
-//      CHECK:   %[[RESULT:.+]] = scf.for %[[IV0:[a-zA-Z0-9]+]] =
-// CHECK-SAME:     %[[C0]] to %[[M]] step %[[C32]]
-// CHECK-SAME:     iter_args(%[[ARG6:.+]] = %[[ARG4]]) -> (tensor<?x?xf32>) {
-//      CHECK:     %[[TILE_M_2:.+]] = affine.min #[[MAP1]](%[[IV0]])[%[[M]]]
-//      CHECK:     %[[N3:.+]] = tensor.dim %[[ARG6]], %[[C1]]
-//      CHECK:     %[[ST_ARG6:.+]] = tensor.extract_slice %[[ARG6]][%[[IV0]], 0]
-// CHECK-SAME:       [%[[TILE_M_2]], %[[N3]]]
-//      CHECK:     %[[TILE_M_3:.+]] = affine.min #[[MAP5]](%[[IV0]])[%[[M]], %[[M]]]
-//      CHECK:     %[[N1:.+]] = tensor.dim %[[ARG0]], %[[C1]]
-//      CHECK:     %[[ST_ARG0:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0]
-// CHECK-SAME:       [%[[TILE_M_3]], %[[N1]]]
-//      CHECK:     %[[N2_2:.+]] = tensor.dim %[[ARG2]], %[[C1]]
-//      CHECK:     %[[ST_ARG2:.+]] = tensor.extract_slice %[[ARG2]][%[[IV0]], 0]
-// CHECK-SAME:       [%[[TILE_M_3]], %[[N2_2]]]
-//      CHECK:     %[[LHS:.+]] = linalg.matmul
-// CHECK-SAME:       __internal_linalg_transform__ = "after_lhs_fusion_producer"
-// CHECK-SAME:       ins(%[[ST_ARG0]], %[[ARG1]] : tensor<?x?xf32>, tensor<?x?xf32>)
-// CHECK-SAME:       outs(%[[ST_ARG2]] : tensor<?x?xf32>)
-//      CHECK:     %[[N2:.+]] = tensor.dim %[[ARG1]], %[[C1]]
-//      CHECK:     %[[N3_2:.+]] = tensor.dim %[[ARG3]], %[[C1]]
-//      CHECK:     %[[YIELD0:.+]] = scf.for %[[IV1:[a-zA-Z0-9]+]] =
-// CHECK-SAME:       %[[C0]] to %[[N3_2]] step %[[C64]]
-// CHECK-SAME:       iter_args(%[[ARG8:.+]] = %[[ST_ARG6]]) -> (tensor<?x?xf32>) {
-//      CHECK:       %[[YIELD1:.+]] = scf.for %[[IV2:[a-zA-Z0-9]+]] =
-// CHECK-SAME:         %[[C0]] to %[[N2]] step %[[C16]]
-// CHECK-SAME:         iter_args(%[[ARG10:.+]] = %[[ARG8]]) -> (tensor<?x?xf32>) {
-//      CHECK:         %[[TILE_N2:.+]] = affine.min #[[MAP2]](%[[IV2]])[%[[N2]]]
-//      CHECK:         %[[ST_LHS:.+]] = tensor.extract_slice %[[LHS]][0, %[[IV2]]]
-// CHECK-SAME:           [%[[TILE_M_3]], %[[TILE_N2]]]
-//      CHECK:         %[[TILE_N3:.+]] = affine.min #[[MAP3]](%[[IV1]])[%[[N3_2]]]
-//      CHECK:         %[[ST_ARG3:.+]] = tensor.extract_slice %[[ARG3]][%[[IV2]], %[[IV1]]]
-// CHECK-SAME:           [%[[TILE_N2]], %[[TILE_N3]]]
-//      CHECK:         %[[M_4:.+]] = tensor.dim %[[ARG10]], %[[C0]]
-//      CHECK:         %[[ST_ARG4:.+]] = tensor.extract_slice %[[ARG10]][0, %[[IV1]]]
-// CHECK-SAME:           [%[[M_4]], %[[TILE_N3]]]
-//      CHECK:         %[[ST_RESULT:.+]] = linalg.matmul
-// CHECK-SAME:           __internal_linalg_transform__ = "after_lhs_fusion"
-// CHECK-SAME:           ins(%[[ST_LHS]], %[[ST_ARG3]]
-// CHECK-SAME:             : tensor<?x?xf32>, tensor<?x?xf32>)
-// CHECK-SAME:           outs(%[[ST_ARG4]] : tensor<?x?xf32>)
-//      CHECK:         %[[UPDATE1:.+]] = tensor.insert_slice %[[ST_RESULT]]
-// CHECK-SAME:           into %[[ARG10]][0, %[[IV1]]] [%[[M_4]], %[[TILE_N3]]]
-//      CHECK:         scf.yield %[[UPDATE1]]
-//      CHECK:       }
-//      CHECK:       scf.yield %[[YIELD1]]
-//      CHECK:     }
-//      CHECK:     %[[UPDATE0:.+]] = tensor.insert_slice %[[YIELD0]] into
-// CHECK-SAME:       %[[ARG6]][%[[IV0]], 0] [%[[TILE_M_2]], %[[N3]]]
-//      CHECK:     scf.yield %[[UPDATE0]]
-//      CHECK:   }
-//      CHECK:   return %[[RESULT]]
+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>) {
+  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
 
 // -----
 
-module {
-  func.func @matmul_plus_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
-                           %arg2: tensor<?x?xf32>) -> tensor<?x?xf32>{
-    %c0 = arith.constant 0 : index
-    %c1 = arith.constant 1 : index
-    %0 = tensor.dim %arg2, %c0 : tensor<?x?xf32>
-    %1 = tensor.dim %arg2, %c1 : tensor<?x?xf32>
-    %2 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
-      outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
-    %3 = tensor.dim %2, %c0 : tensor<?x?xf32>
-    %4 = tensor.dim %2, %c1 : tensor<?x?xf32>
-    %5 = linalg.init_tensor [%3, %4] : tensor<?x?xf32>
-    %6 = linalg.generic
-      {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
-                        affine_map<(d0, d1) -> (d0, d1)>,
-                        affine_map<(d0, d1) -> (d0, d1)>],
-       iterator_types = ["parallel", "parallel"],
-       __internal_linalg_transform__ = "transpose_fusion"}
-      ins(%2, %2 : tensor<?x?xf32>, tensor<?x?xf32>)
-      outs(%5 : tensor<?x?xf32>) {
-      ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32) :
-        %7 = arith.addf %arg3, %arg4 : f32
-        linalg.yield %7 : f32
-      } -> tensor<?x?xf32>
-    return %6 : tensor<?x?xf32>
-  }
+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: func @matmul_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<?x?xf32>
-//       CHECK:   %[[RESULT:.+]] = scf.for %[[IV0:[a-zA-Z0-9_]+]]
-//  CHECK-SAME:     iter_args(%[[ARG4:.+]] = %{{[a-zA-Z0-9_]+}})
-//       CHECK:     %[[YIELD:.+]] = scf.for %[[IV1:[a-zA-Z0-9_]+]]
-//  CHECK-SAME:       iter_args(%[[ARG6:.+]] = %[[ARG4]])
-//       CHECK:       %[[ST_ARG6:.+]] = tensor.extract_slice %[[ARG6]][%[[IV0]], %[[IV1]]]
-//       CHECK:       %[[ST_ARG0:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0]
-//       CHECK:       %[[ST_ARG1:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]]
-//       CHECK:       %[[ST_ARG2:.+]] = tensor.extract_slice %[[ARG2]][%[[IV0]], %[[IV1]]]
-//       CHECK:       %[[LHS:.+]] = linalg.matmul
-//  CHECK-SAME:         ins(%[[ST_ARG0]], %[[ST_ARG1]]
-//  CHECK-SAME:           : tensor<?x?xf32>, tensor<?x?xf32>)
-//  CHECK-SAME:         outs(%[[ST_ARG2]] : tensor<?x?xf32>)
-//       CHECK:       %[[ST_RESULT:.+]] = linalg.generic
-//  CHECK-SAME:         ins(%[[LHS]] : tensor<?x?xf32>)
-//  CHECK-SAME:         outs(%[[ST_ARG6]] : tensor<?x?xf32>)
-//       CHECK:       %[[UPDATE:.+]] = tensor.insert_slice %[[ST_RESULT]]
-//  CHECK-SAME:         into %[[ARG6]][%[[IV0]], %[[IV1]]]
-//       CHECK:       scf.yield %[[UPDATE]]
-//       CHECK:     scf.yield %[[YIELD]]
-//       CHECK:   return %[[RESULT]]
+// 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
 
 // -----
 
-module {
-  func.func @matmul_out_fusion(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
-                      %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
-    %c0 = arith.constant 0.0 : f32
-    %0 = linalg.fill ins(%c0 : f32) outs(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32>
-    %1 = linalg.matmul {__internal_linalg_transform__ = "out_fusion"}
-      ins(%arg1, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>)
-      outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
-    return %1 : 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-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
-//       CHECK: %[[C0:.*]] = arith.constant 0.0{{.*}} : f32
-//   CHECK-NOT: fill
-//       CHECK: scf.for %[[I:.*]]{{.*}}iter_args(%{{.*}} = %[[ARG0]]) -> (tensor<?x?xf32>) {
-//       CHECK:   scf.for %[[J:.*]]
-//       CHECK:     %[[ST:.*]] = tensor.extract_slice %[[ARG0]]
-//       CHECK:     %[[ST_FILL:.*]] = linalg.fill
-//  CHECK-SAME:       {__internal_linalg_transform__ = "after_out_fusion_producer"}
-//  CHECK-SAME:       ins(%[[C0]] : f32) outs(%[[ST]] : tensor<?x?xf32>) -> tensor<?x?xf32>
-//       CHECK:     %[[ST_MM_RES:.*]] = scf.for %[[K:.*]]{{.*}}iter_args(%[[BB:.*]] = %[[ST_FILL]]) -> (tensor<?x?xf32>) {
-//   CHECK-NOT:       fill
-//       CHECK:       %[[ST_FILL_SUB:.*]] = tensor.extract_slice %[[BB]][0, 0]
-//       CHECK:       %[[ST_MM_SUB:.*]] = linalg.matmul {__internal_linalg_transform__ = "after_out_fusion"} ins(%{{.*}}, %{{.*}} : tensor<?x?xf32>, tensor<?x?xf32>) outs(%[[ST_FILL_SUB]] : tensor<?x?xf32>) -> tensor<?x?xf32>
-//       CHECK:       %[[ST_MM:.*]] = tensor.insert_slice %[[ST_MM_SUB]] into %[[BB]]
-//       CHECK:       scf.yield %[[ST_MM]] : tensor<?x?xf32>
-//       CHECK:     %[[MM:.*]] = tensor.insert_slice %[[ST_MM_RES]] into {{.*}}
-//       CHECK:     scf.yield %[[MM]] : tensor<?x?xf32>
+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]]
 
 // -----
 
-module {
-  func.func @generic_plus_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
-                      %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
-    %c0 = arith.constant 0.0 : f32
-    %0 = linalg.generic {
-      indexing_maps = [affine_map<(m, n) -> ()>, affine_map<(m, n) -> (m, n)>],
+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(%c0 : f32)
-    outs(%arg0: tensor<?x?xf32>) {
-      ^bb(%0: f32, %1: f32) :
-        linalg.yield %0 : f32
+      ins(%arg2 : tensor<?xf32>) outs(%init: tensor<?x?xf32>) {
+    ^bb(%b0: f32, %b1: f32) :
+      linalg.yield %b0 : f32
     } -> tensor<?x?xf32>
-    %1 = linalg.matmul {__internal_linalg_transform__ = "out_fusion"}
-      ins(%arg1, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>)
-      outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
-    return %1 : 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
   TestLinalgCodegenStrategy.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 registerTestMultiBuffering();
 void registerTestGpuParallelLoopMappingPass();
 void registerTestIRVisitorsPass();
 void registerTestGenericIRVisitorsPass();
 void registerTestGenericIRVisitorsInterruptPass();
 void registerTestInterfaces();
 void registerTestLinalgCodegenStrategy();
 void registerTestLinalgElementwiseFusion();
+void registerTestLinalgFusionOnTensorsPass();
 void registerTestLinalgFusionTransforms();
 void registerTestLinalgTensorFusionTransforms();
 void registerTestLinalgTiledLoopFusionTransforms();
 void registerTestLinalgGreedyFusion();
 void registerTestLinalgHoisting();
 void registerTestLinalgTileAndFuseSequencePass();
 void registerTestLinalgTransforms();
 void registerTestLivenessPass();
#endif
   mlir::test::registerTestComposeSubView();
   mlir::test::registerTestMultiBuffering();
   mlir::test::registerTestGpuParallelLoopMappingPass();
   mlir::test::registerTestIRVisitorsPass();
   mlir::test::registerTestGenericIRVisitorsPass();
   mlir::test::registerTestInterfaces();
   mlir::test::registerTestLinalgCodegenStrategy();
   mlir::test::registerTestLinalgElementwiseFusion();
+  mlir::test::registerTestLinalgFusionOnTensorsPass();
   mlir::test::registerTestLinalgFusionTransforms();
   mlir::test::registerTestLinalgTensorFusionTransforms();
   mlir::test::registerTestLinalgTiledLoopFusionTransforms();
   mlir::test::registerTestLinalgGreedyFusion();
   mlir::test::registerTestLinalgHoisting();
   mlir::test::registerTestLinalgTileAndFuseSequencePass();
   mlir::test::registerTestLinalgTransforms();
   mlir::test::registerTestLivenessPass();