[mlir][TilingInterface]NFC Refactor of tile and fuse using `TilingInterface`.

[mravishankar]

Summary

This patch refactors the tiling and tile + fuse implementation using
TilingInterface. Primarily, it exposes the functionality as simple
utility functions instead of as a Pattern to allow calling it from a
pattern as it is done in the test today or from within the transform
dialect (in the future). This is a step towards deprecating similar
methods in Linalg dialect.

• The utility methods do not erase the root operations.
• The return value provides the values to use for replacements.
• Along the way also simplify the implementation and remove unnecessary methods. In particular, it is not legal to replace all uses of values made iter_args within the loop body. That replacement needs to happen in a more targeted manner only when fusion happens through operands that are destination operands. Gaurd this change in the utility method to avoid doing that during the tile (+ fuse) transformation.

[nicolasvasilache]

This is getting more reviewable but it is still trying to sneak a load-bearing behavior through an NFC commit :)

Please drop anything related to destinationIterArg to make this truly NFC then we can evaluate that part separately.

Since this is the 3rd time we go through this, if this is getting too annoying, please lmk and I'll commander.

mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h

76–86

tiling is by itself oblivious to the existent of a consumer, please drop this from the name as it is confusing to bring in fusion-related concepts here

95

guaranteed

mlir/include/mlir/IR/OpDefinition.h

243 ↗

(On Diff #461094)

please drop this change, this PR does should not touch core IR

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

284

This feels a lot like premature optimization where the APIs is made significantly more complex for efficiency reasons that I'd like to see measured.
I strongly prefer we get the simplest possible API first, really understand the problem in its simplest form, then measure and then invest in removing overhead.

287

It would be better and more future-proof for this yieldTiledValues to take a SubsetExtractOpInterface and return an instance of the symmetrical SubsetInsertOpInterface.
Depending on the op, we can then perform the getMixedOffsets() inside that.
This will make it easy to support other extract / insert op pairs and in particular parallel_insert ops.

You can approximate this by just passing ArrayRef<Operation*> sliceOps for now and switch to dyn_cast<tensor::ExtractSliceOp> + add a TODO.

[mravishankar]

Added https://reviews.llvm.org/D134411 to separate out the functional changes from this patch. Once that lands, Ill rebase on top of that to make this a true NFC.

[mravishankar]

Rebase on top of D134411 to make this truly NFC.

[mravishankar]

@nicolasvasilache this change is now just NFC.

mlir/include/mlir/IR/OpDefinition.h

243 ↗

(On Diff #461094)

This helps with debugging. All dump methods are const. This seems to be an error. Ok to add this as part of this (now NFC) change?

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

287

Agreed! That would be clean. Added a TODO to move to using the interface that will come as part of the overall change to move TilingInterface to use SubsetExtractOpInterface.

[mravishankar]

Address comments.

[mravishankar]

Drop changes to OpDefinition.h.

[nicolasvasilache]

Thanks @mravishankar !
Minor comment but this is otherwise GTG!

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

391

we usually dump debug to dbgs, any reason to go to errs here ?

mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h

struct SCFTilingOptions {
   /// The interchange vector to reorder the tiled loops.
   SmallVector<unsigned> interchangeVector = {};
   SCFTilingOptions &setInterchange(ArrayRef<unsigned> interchange) {
     interchangeVector = llvm::to_vector(interchange);
     return *this;
   }
 };
 
+/// Transformation information returned after tiling.
 struct SCFTilingResult {
+  /// The tiled operation generated.
   Operation *tiledOp;
+  /// The `scf.for` operations that iterate over the tiles.
   SmallVector<scf::ForOp> loops;
+  /// Values to use as replacements for the untiled op. Is the same size as the
+  /// number of results of the untiled op.
+  SmallVector<Value> replacements;
 };
 
-/// Pattern to tile an op that implements the `TilingInterface` using
+/// Method to tile an op that implements the `TilingInterface` using
 /// `scf.for` for iterating over the tiles.
-struct TileUsingSCFForOp : public OpInterfaceRewritePattern<TilingInterface> {
-  /// Construct a generic pattern applied to all TilingInterface ops.
-  TileUsingSCFForOp(MLIRContext *context, SCFTilingOptions options,
-                    PatternBenefit benefit = 1);
-
-  /// Construct a generic pattern applied to `opName`.
-  TileUsingSCFForOp(StringRef opName, MLIRContext *context,
-                    SCFTilingOptions options, PatternBenefit benefit = 1);
-
-  /// `matchAndRewrite` implementation that returns the significant transformed
-  /// pieces of IR.
-  FailureOr<SCFTilingResult>
-  returningMatchAndRewrite(TilingInterface op, PatternRewriter &rewriter) const;
-
-  LogicalResult matchAndRewrite(TilingInterface op,
-                                PatternRewriter &rewriter) const override {
-    return returningMatchAndRewrite(op, rewriter);
+FailureOr<SCFTilingResult> tileUsingSCFForOp(RewriterBase &rewriter,
+                                             TilingInterface op,
+                                             SCFTilingOptions options);
+
+/// Options used to control tile + fuse.
+struct SCFTileAndFuseOptions {
+  /// The tiling options used to control the tiling of the consumer.
+  SCFTilingOptions tilingOptions;
+  SCFTileAndFuseOptions &setTilingOptions(SCFTilingOptions options) {
+    tilingOptions = options;
+    return *this;
   }
+};
 
-private:
-  /// Options to control tiling;
-  SCFTilingOptions options;
+/// Transformation information returned after tile and fuse.
+struct SCFTileAndFuseResult {
+  /// List of untiled operations that were fused with the tiled consumer.
+  llvm::SetVector<Operation *> fusedProducers;
+  /// List of tiled and fused operations generated. The first one in this list
+  /// is guaranteed to be the tiled operations generated during tiling of the
+  /// generated operation.
+  llvm::SetVector<Operation *> tiledAndFusedOps;
+  /// The `scf.for` operations that iterate over the tiles.
+  SmallVector<scf::ForOp> loops;
+  /// The replacement values to use for the tiled and fused operations.
+  llvm::DenseMap<Value, Value> replacements;
 };
 
-/// Pattern to tile and fuse a sequence of operations, by tiling the consumer
+/// Method to tile and fuse a sequence of operations, by tiling the consumer
 /// and fusing its producers. Note that this assumes that it is valid to
 /// tile+fuse the producer into the innermost tiled loop. Its up to the caller
 /// to ensure that the tile sizes provided make this fusion valid.
 ///
 /// For example, for the following sequence
 ///
 /// ```mlir
-/// %0 = linalg.fill ...
-/// %1 = linalg.matmul ... outs(%0 : ...) ...
+/// %0 =
+/// %1 = linalg.fill ... outs(%0 : ... )
+/// %2 = linalg.matmul ... outs(%1 : ...) ...
 /// ```
 ///
 /// it is legal to fuse the fill with the matmul only if the matmul is tiled
 /// along the parallel dimensions and not the reduction dimension, i.e. the tile
-/// size for the reduction dimension should be 0.
-struct SCFTileAndFuseResult {
-  SmallVector<Operation *> tiledAndFusedOps;
-  SmallVector<scf::ForOp> loops;
-};
-struct TileConsumerAndFuseProducersUsingSCFForOp
-    : public OpInterfaceRewritePattern<TilingInterface> {
-
-  /// Construct a generic pattern applied to all TilingInterface ops.
-  TileConsumerAndFuseProducersUsingSCFForOp(MLIRContext *context,
-                                            SCFTilingOptions options,
-                                            PatternBenefit benefit = 1);
-
-  /// Construct a generic pattern applied to `opName`.
-  TileConsumerAndFuseProducersUsingSCFForOp(StringRef opName,
-                                            MLIRContext *context,
-                                            SCFTilingOptions options,
-                                            PatternBenefit benefit = 1);
-
-  /// `matchAndRewrite` implementation that returns the significant transformed
-  /// pieces of IR.
-  FailureOr<SCFTileAndFuseResult>
-  returningMatchAndRewrite(TilingInterface op, PatternRewriter &rewriter) const;
-
-  LogicalResult matchAndRewrite(TilingInterface op,
-                                PatternRewriter &rewriter) const override {
-    return returningMatchAndRewrite(op, rewriter);
-  }
-
-private:
-  /// This pattern uses the tiling pattern. Instead of using inheritance, use
-  /// the patterns as private object that is instantiated at the same time as
-  /// this pattern.
-  TileUsingSCFForOp tilingPattern;
-};
-
-/// Pattern to lower operations that implement the `TilingInterface` to
-/// loops/scalar IR using `scf.for`.
-struct LowerToLoopsUsingSCFForOp
-    : public OpInterfaceRewritePattern<TilingInterface> {
-  using OpInterfaceRewritePattern<TilingInterface>::OpInterfaceRewritePattern;
+/// size for the reduction dimension should be 0. The resulting fused
+/// transformation is
+///
+/// ```mlir
+/// %1 = scf.for ... iter_args(%arg0 = %0)
+///   %2 = tensor.extract_slice %arg0
+///   %3 = linalg.fill .. outs(%2 : ... )
+///   %4 = linalg.matmul .. outs(%3 : ...)
+/// }
+/// ```
+tileConsumerAndFuseProducerGreedilyUsingSCFForOp(RewriterBase &rewriter,
+                                                 TilingInterface consumer,
+                                                 SCFTileAndFuseOptions options);
 
-  /// `matchAndRewrite` implementation that returns the significant transformed
-  /// pieces of IR.
-  FailureOr<SmallVector<scf::ForOp>>
-  returningMatchAndRewrite(TilingInterface op, PatternRewriter &rewriter) const;
-
-  LogicalResult matchAndRewrite(TilingInterface op,
-                                PatternRewriter &rewriter) const override {
-    return returningMatchAndRewrite(op, rewriter);
-  }
-};
+/// Method to lower an `op` that implements the `TilingInterface` to
+/// loops/scalars.
+lowerToLoopsUsingSCFForOp(RewriterBase &rewriter, TilingInterface op);
 
 } // namespace scf
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_TRANSFORMS_TILEUSINGINTERFACE_H

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

for (auto val : interchange) {
     if (seenVals.count(val))
       return false;
     seenVals.insert(val);
   }
   return seenVals.size() == interchange.size();
 }
 
 //===----------------------------------------------------------------------===//
-// TileUsingSCFForOp pattern implementation.
+// tileUsingSCFForOp implementation.
 //===----------------------------------------------------------------------===//
 
 // Check if `stride` evenly divides the trip count `size - offset`.
 static bool tileDividesIterationDomain(Range loopRange) {
   Optional<int64_t> offsetAsInt = getConstantIntValue(loopRange.offset);
   if (!offsetAsInt)
     return false;
   Optional<int64_t> sizeAsInt = getConstantIntValue(loopRange.size);
auto loop = builder.create<scf::ForOp>(
         });
     offsets[loopRange.index()] = loop.getInductionVar();
     loops.push_back(loop);
     builder.setInsertionPoint(loop.getBody()->getTerminator());
   }
   return loops;
 }
 
-/// If the tiled operation is in destination passing style, update the
+/// For a value to be yielded (`yieldedValue`) from within a loop nest `loops`,
+/// construct the destructive update pattern that inserts the yielded
+/// value into a destination tensor provided by `initValue` at offset
+/// `tileOffsets` and size `tileSizes`. For example,
+///
+/// ```mlir
+/// scf.for %iv0 = ... {
+///   %0 = tiled_op
+/// }
+/// ```
+///
+/// is transformed to
+///
+/// ```mlir
+/// scf.for %iv0 = ... iter_args(%arg = %0) {
+///   %1 = tensor.extract_slice %arg
+///   %2 = tiled_op
+///   %3 = tensor.insert_slice %2 into %arg
+///   scf.yield %3
+/// }
+/// ```
+/// TODO: This API can be cleaned up by using `SubsetExtractOpInterface`.
+static FailureOr<SmallVector<Value>>
+yieldTiledValues(RewriterBase &rewriter, ValueRange initValues,
+                 ValueRange yieldedValues,
+                 ArrayRef<SmallVector<OpFoldResult>> tileOffsetsList,
+                 ArrayRef<SmallVector<OpFoldResult>> tileSizesList,
+                 MutableArrayRef<scf::ForOp> loops) {
+  NewYieldValueFn yieldValueFn =
+      [&](OpBuilder &b, Location loc,
+          ArrayRef<BlockArgument> newBBArgs) -> SmallVector<Value> {
+    SmallVector<Value> inserts;
+    for (auto yieldedValue : llvm::enumerate(yieldedValues)) {
+      ArrayRef<OpFoldResult> tileOffsets =
+          tileOffsetsList[yieldedValue.index()];
+      ArrayRef<OpFoldResult> tileSizes = tileSizesList[yieldedValue.index()];
+      SmallVector<OpFoldResult> tileStrides(tileOffsets.size(),
+                                            b.getIndexAttr(1));
+      Value insert = b.create<tensor::InsertSliceOp>(
+          loc, yieldedValue.value(), newBBArgs[yieldedValue.index()],
+          tileOffsets, tileSizes, tileStrides);
+      inserts.push_back(insert);
+    }
+    return inserts;
+  };
+
+  SmallVector<scf::ForOp> newLoops =
+      replaceLoopNestWithNewYields(rewriter, loops, initValues, yieldValueFn,
+                                   /*replaceIterOperandsUsesInLoop =*/false);
+  for (const auto &loop : llvm::enumerate(loops)) {
+    rewriter.eraseOp(loop.value());
+    loops[loop.index()] = newLoops[loop.index()];
+  }
+  return llvm::to_vector(llvm::map_range(
+      loops.front().getResults().take_back(yieldedValues.size()),
+      [](OpResult r) -> Value { return r; }));
+}
+
+/// If the tiled operation is destination passing style, update the
 /// slice of the destination used (which refers to the untiled destination)
 /// to use the corresponding region argument of the innermost loop.
 ///
 /// ```mlir
 /// %0 =
 /// scf.for %iv0 = ... iter_args(%arg = %0) {
 ///   %1 = tensor.extract_slice %0
 ///   %2 = tiled_op
 ///   %3 = tensor.insert_slice %2 into %arg
 ///   scf.yield %3
 /// }
 /// ```
 ///
 /// is transformed to
 ///
 /// ```mlir
 /// scf.for %iv0 = ... iter_args(%arg = %0) {
 ///   %1 = tensor.extract_slice %arg
 ///   %2 = tiled_op
 ///   %3 = tensor.insert_slice %2 into %arg
 ///   scf.yield %3
 /// }
 /// ```
-/// TODO: This can be made much cleaner when `DestinationStyleOp` interface is
-/// available generally.
 static void
 updateDestinationOperandsForTiledOp(OpBuilder &builder,
                                     ValueRange tiledOpDestinationValues,
                                     ValueRange bbArgsList) {
   for (auto destValue : llvm::enumerate(tiledOpDestinationValues)) {
     auto sliceOp = destValue.value().getDefiningOp<tensor::ExtractSliceOp>();
     if (!sliceOp)
       continue;
     sliceOp.setOperand(0, bbArgsList[destValue.index()]);
   }
 }
 
-scf::TileUsingSCFForOp::TileUsingSCFForOp(MLIRContext *context,
-                                          scf::SCFTilingOptions options,
-                                          PatternBenefit benefit)
-    : OpInterfaceRewritePattern<TilingInterface>(context, benefit),
-      options(std::move(options)) {}
-
-scf::TileUsingSCFForOp::TileUsingSCFForOp(StringRef opName,
-                                          MLIRContext *context,
-                                          scf::SCFTilingOptions options,
-                                          PatternBenefit benefit)
-    : OpInterfaceRewritePattern<TilingInterface>(context, benefit),
-      options(std::move(options)) {}
-
+/// Implementation of tiling transformation of `op` that implements the
+/// `TilingInterface` using `scf.for` to iterate over the tiles.
 FailureOr<scf::SCFTilingResult>
-scf::TileUsingSCFForOp::returningMatchAndRewrite(
-    TilingInterface op, PatternRewriter &rewriter) const {
+mlir::scf::tileUsingSCFForOp(RewriterBase &rewriter, TilingInterface op,
+                             scf::SCFTilingOptions options) {
   OpBuilder::InsertionGuard guard(rewriter);
   rewriter.setInsertionPointAfter(op);
 
   if (!options.tileSizeComputationFunction) {
     return rewriter.notifyMatchFailure(
         op, "missing tile size computation function");
   }
 
   // 1. Get the range of the loops that are represented by the operation.
   SmallVector<Range> iterationDomain = op.getIterationDomain(rewriter);
   size_t numLoops = iterationDomain.size();
   if (numLoops == 0) {
     return rewriter.notifyMatchFailure(
         op, "unable to tile op with no iteration domain");
   }
 
   // 2. Materialize the tile sizes. Enforce the convention that "tiling by zero"
   // skips tiling a particular dimension. This convention is significantly
   // simpler to handle instead of adjusting affine maps to account for missing
   // dimensions.
   SmallVector<Value> tileSizeVector =
       options.tileSizeComputationFunction(rewriter, op);
   if (tileSizeVector.size() < iterationDomain.size()) {
     auto zero = rewriter.create<arith::ConstantIndexOp>(op.getLoc(), 0);
     tileSizeVector.append(numLoops - tileSizeVector.size(), zero);
   }
 
SmallVector<OpFoldResult> offsets, sizes;
     tilingResult.loops = generateTileLoopNest(
         rewriter, op.getLoc(), iterationDomain, tileSizeVector, offsets, sizes);
 
     if (!interchangeVector.empty()) {
       auto inversePermutation = invertPermutationVector(interchangeVector);
       offsets = applyPermutationToVector(offsets, inversePermutation);
       sizes = applyPermutationToVector(sizes, inversePermutation);
     }
+  }
 
-    LLVM_DEBUG({
-      if (!tilingResult.loops.empty()) {
-        llvm::errs() << "LoopNest shell :\n";
-        tilingResult.loops.front().dump();
-        llvm::errs() << "\n";
-      }
-    });
+      llvm::dbgs() << "LoopNest shell :\n";
+      llvm::dbgs() << "\n";
 
-    // 4. Generate the tiled implementation within the inner most loop.
-    if (!tilingResult.loops.empty())
-      rewriter.setInsertionPoint(
-          tilingResult.loops.back().getBody()->getTerminator());
-    SmallVector<Operation *> tiledImplementation =
-        op.getTiledImplementation(rewriter, offsets, sizes);
-    if (tiledImplementation.size() != 1) {
-      return rewriter.notifyMatchFailure(
-          op, "expected tiled implementation to return a single op");
-    }
-    tilingResult.tiledOp = tiledImplementation[0];
-
-    LLVM_DEBUG({
-      if (!tilingResult.loops.empty()) {
-        llvm::errs() << "After tiled implementation :\n";
-        tilingResult.loops.front().dump();
-        llvm::errs() << "\n";
-      }
-    });
-  }
-
   if (op->getNumResults() == 0) {
-    rewriter.eraseOp(op);
+    // nothing more to do.
     return tilingResult;
   }
 
-  // 5. If the original operations has results, modify the loop nest to yield
-  // the replacement values.
+  // If loops are empty, the tiled op is used as the replacement for the untiled
+  // op.
   if (tilingResult.loops.empty()) {
-    // 5a. If there were no loops, the tiled implementation results are the
-    // replacements.
-    rewriter.replaceOp(op, tilingResult.tiledOp->getResults());
+    tilingResult.replacements = llvm::to_vector(
+        llvm::map_range(tiledImplementation[0]->getResults(),
+                        [](OpResult result) -> Value { return result; }));
     return tilingResult;
   }
 
-  // 6. Yield the results of the tiled operation from the loop nest as
-  //    replacements for the original untiled ops.
-  if (tilingResult.tiledOp->getNumResults() != op->getNumResults()) {
+  // 5. Yield all the results of the tiled operation. The surrounding loop
+  //    nest is modified to insert a destructive update pattern to yield
+  //    from the loop nest values to replace the untiled op with.
+  unsigned numResults = op->getNumResults();
+  SmallVector<SmallVector<OpFoldResult>> resultOffsetsList(numResults),
+      resultSizesList(numResults);
+  for (auto result : llvm::enumerate(op->getResults())) {
+    if (failed(op.getResultTilePosition(rewriter, result.index(), offsets,
+                                        sizes,
+                                        resultOffsetsList[result.index()],
+                                        resultSizesList[result.index()]))) {
       return rewriter.notifyMatchFailure(
-        tilingResult.tiledOp,
-        "expected tiled op to have as many results as the untiled operation");
+          op, "failed to get slice of result produced");
+    }
   }
 
-  // `scf.for` with tensor semantics requires the loop nest to yield the
-  // replacement values using destructive updates. Use the `TilingInterface`
-  // to get the position of the result tiles and use that to generate the
-  // destructive update pattern, i.e.,
-  //
-  // ```mlir
-  // scf.for %iv0 = ... {
-  //   %0 = tiled_op
-  // }
-  // ```
-  //
-  // is transformed to
-  //
-  // ```mlir
-  // %result = scf.for %iv0 = ... iter_args(%arg = %init) -> .. {
-  //   %0 = tiled_op
-  //   %1 = tensor.insert_slice %0 into %arg[..] [..] [..]
-  //   scf.yield %1
-  // }
-  // ```
-  NewYieldValueFn yieldValueFn =
-      [&](OpBuilder &b, Location loc,
-          ArrayRef<BlockArgument> newBBArgs) -> SmallVector<Value> {
-    SmallVector<Value> yieldedValues;
-    Attribute one = b.getIndexAttr(1);
-    for (auto resultNum : llvm::seq<unsigned>(0, op->getNumResults())) {
-      SmallVector<OpFoldResult> resultTileOffsets, resultTileSizes;
-      if (failed(op.getResultTilePosition(b, resultNum, offsets, sizes,
-                                          resultTileOffsets,
-                                          resultTileSizes))) {
-        op.emitOpError("unable to get position of result ")
-            << resultNum << " of the tiled implementation";
-        return {};
-      }
-      SmallVector<OpFoldResult> resultTileStrides(resultTileOffsets.size(),
-                                                  one);
-      Value yieldedValue = b.create<tensor::InsertSliceOp>(
-          op->getLoc(), tilingResult.tiledOp->getResult(resultNum),
-          newBBArgs[resultNum], resultTileOffsets, resultTileSizes,
-          resultTileStrides);
-      yieldedValues.push_back(yieldedValue);
-    }
-    return yieldedValues;
-  };
-  SmallVector<scf::ForOp> newLoops = replaceLoopNestWithNewYields(
-      rewriter, tilingResult.loops, op.getDestinationOperands(rewriter),
-      yieldValueFn, /*replaceIterOperandsUsesInLoops =*/false);
-  for (const auto &loop : llvm::enumerate(tilingResult.loops)) {
-    rewriter.eraseOp(loop.value());
-    tilingResult.loops[loop.index()] = newLoops[loop.index()];
+  FailureOr<SmallVector<Value>> replacementOr =
+      yieldTiledValues(rewriter, op.getDestinationOperands(rewriter),
+                       tilingResult.tiledOp->getResults(), resultOffsetsList,
+                       resultSizesList, tilingResult.loops);
+  if (failed(replacementOr))
+    return rewriter.notifyMatchFailure(op, "failed to yield replacement");
+  if (auto tiledInterfaceOp = dyn_cast<TilingInterface>(tilingResult.tiledOp)) {
+    auto innerMostLoop = tilingResult.loops.back();
+    updateDestinationOperandsForTiledOp(
+        rewriter, tiledInterfaceOp.getDestinationOperands(rewriter),
+        innerMostLoop.getRegionIterArgs());
+
+  tilingResult.replacements = replacementOr.value();
+
+  LLVM_DEBUG({
+    if (!tilingResult.loops.empty()) {
+      llvm::dbgs() << "After tiled implementation :\n";
+      tilingResult.loops.front().dump();
+      llvm::dbgs() << "\n";
     }
-  rewriter.replaceOp(op, tilingResult.loops.front().getResults());
+  });
   return tilingResult;
 }
 
 //===----------------------------------------------------------------------===//
-// TileConsumerAndFuseProducersUsingSCFForOp pattern implementation.
+// tileConsumerAndFuseProducerGreedilyUsingSCFForOp implementation.
 //===----------------------------------------------------------------------===//
 
-scf::TileConsumerAndFuseProducersUsingSCFForOp::
-    TileConsumerAndFuseProducersUsingSCFForOp(MLIRContext *context,
-                                              scf::SCFTilingOptions options,
-                                              PatternBenefit benefit)
-    : OpInterfaceRewritePattern<TilingInterface>(context, benefit),
-      tilingPattern(context, std::move(options)) {}
-
-scf::TileConsumerAndFuseProducersUsingSCFForOp::
-    TileConsumerAndFuseProducersUsingSCFForOp(StringRef opName,
-                                              MLIRContext *context,
-                                              scf::SCFTilingOptions options,
-                                              PatternBenefit benefit)
-    : OpInterfaceRewritePattern<TilingInterface>(context, benefit),
-      tilingPattern(context, std::move(options)) {}
-
 /// Return the untiled producer whose slice is used in a tiled consumer. The
 /// method traverses the tile loop nest (`loops`) if needed, and returns the
 /// `iter_args` of the outer most that is encountered. Traversing the iter_args
 /// indicates that this is a destination operand of the consumer. If there was
 /// no loop traversal needed, the second value of the returned tuple is empty.
 static std::tuple<OpResult, Optional<OpOperand *>>
 getUntiledProducerFromSliceSource(OpOperand *source,
                                   ArrayRef<scf::ForOp> loops) {
   Optional<OpOperand *> destinationIterArg;
   auto loopIt = loops.rbegin();
   while (auto iterArg = source->get().dyn_cast<BlockArgument>()) {
     scf::ForOp loop = *loopIt;
     if (iterArg.getOwner()->getParentOp() != loop)
       break;
     source = &loop.getOpOperandForRegionIterArg(iterArg);
     loopIt++;
   }
   if (loopIt == loops.rend())
     destinationIterArg = source;
   return {source->get().dyn_cast<OpResult>(), destinationIterArg};
 }
 
+/// Implementation of tile consumer and fuse producer greedily.
 FailureOr<scf::SCFTileAndFuseResult>
-scf::TileConsumerAndFuseProducersUsingSCFForOp::returningMatchAndRewrite(
-    TilingInterface op, PatternRewriter &rewriter) const {
+mlir::scf::tileConsumerAndFuseProducerGreedilyUsingSCFForOp(
+    RewriterBase &rewriter, TilingInterface consumer,
+    scf::SCFTileAndFuseOptions options) {
   // This transformation is only valid for ops that return values (i.e. not
   // valid to use with operations that have memref operands).
-  if (!op->getNumResults()) {
+  if (!consumer->getNumResults()) {
     return rewriter.notifyMatchFailure(
-        op, "invalid pattern for op with no results");
+        consumer, "invalid pattern for op with no results");
   }
 
   // 1. First tile the consumer.
-  SCFTileAndFuseResult tileAndFuseResult;
+  scf::SCFTileAndFuseResult tileAndFuseResult;
+  llvm::SmallDenseMap<Value, unsigned> yieldedValueToResultNumber;
   {
-    FailureOr<SCFTilingResult> tilingResult =
-        tilingPattern.returningMatchAndRewrite(op, rewriter);
-    if (failed(tilingResult)) {
-      return failure();
-    }
-    tileAndFuseResult.tiledAndFusedOps.push_back(tilingResult->tiledOp);
+    FailureOr<scf::SCFTilingResult> tilingResult =
+        tileUsingSCFForOp(rewriter, consumer, options.tilingOptions);
+    if (failed(tilingResult))
+      return rewriter.notifyMatchFailure(consumer, "failed to tile consumer");
+    tileAndFuseResult.tiledAndFusedOps.insert(tilingResult->tiledOp);
     tileAndFuseResult.loops = std::move(tilingResult->loops);
+    for (auto result : llvm::enumerate(
+             llvm::zip(consumer->getResults(), tilingResult->replacements))) {
+      tileAndFuseResult.replacements[std::get<0>(result.value())] =
+          std::get<1>(result.value());
+      yieldedValueToResultNumber[tilingResult->tiledOp->getResult(
+          result.index())] = result.index();
     }
+  }
+
+  // If there are no loops generated, fusion is immaterial.
+  if (tileAndFuseResult.loops.empty())
+    return tileAndFuseResult;
 
   // 2. Typically, the operands of the tiled operation are slices of the
   //    operands of the untiled operation. These are expressed in IR using
   //    `tensor.extract_slice` operations with source being the operands of the
   //    untiled operation. Create a worklist of these `tensor.extract_slice`
   //    operations. If the producers of the source of the `tensor.extract_slice`
   //    can be tiled such that the tiled value is generated in-place, that
   //    effectively tiles + fuses the operations.
while (!candidates.empty()) {
     if (failed(fusedProducerValue))
       continue;
     rewriter.replaceOp(candidateSliceOp, fusedProducerValue.value());
 
     // 2d. The operands of the fused producer might themselved be slices of
     //     values produced by operations that implement the `TilingInterface`.
     //     Add these operations to the worklist.
     Operation *fusedProducer = fusedProducerValue->getDefiningOp();
-    tileAndFuseResult.tiledAndFusedOps.push_back(fusedProducer);
+    tileAndFuseResult.tiledAndFusedOps.insert(fusedProducer);
     addCandidateSlices(fusedProducer, candidates);
 
     // 2e. If the slice is for a destination operand, for example,
     //
     // ```mlir
     // %0 = linalg.init
     // %1 = linalg.fill .. outs(%0 : )
     // %2 = scf.for .. iter_args(%arg0 = %1) {
if (iterArgNumber) {
             innerMostLoop.getRegionIterArgs()[iterArgNumber.value()]);
       }
     }
   }
   return tileAndFuseResult;
 }
 
 //===----------------------------------------------------------------------===//
-// LowerToLoopsUsingSCFForOp
+// lowerToLoopsUsingSCFForOp implementation.
 //===----------------------------------------------------------------------===//
 
 FailureOr<SmallVector<scf::ForOp>>
-scf::LowerToLoopsUsingSCFForOp::returningMatchAndRewrite(
-    TilingInterface op, PatternRewriter &rewriter) const {
-  SmallVector<Range> domain = op.getIterationDomain(rewriter);
-
+mlir::scf::lowerToLoopsUsingSCFForOp(RewriterBase &rewriter,
+                                     TilingInterface op) {
   // TODO: Handle cases where the op has results if needed.
   if (op->getNumResults() > 0) {
     return rewriter.notifyMatchFailure(
         op, "unable to lower to loops operations with return values");
   }
 
+  SmallVector<Range> domain = op.getIterationDomain(rewriter);
   SmallVector<Value> ivs;
   SmallVector<scf::ForOp> loops;
   Location loc = op.getLoc();
   for (auto loopRange : domain) {
     Value offsetVal =
         getValueOrCreateConstantIndexOp(rewriter, loc, loopRange.offset);
     Value sizeVal =
         getValueOrCreateConstantIndexOp(rewriter, loc, loopRange.size);
     Value strideVal =
         getValueOrCreateConstantIndexOp(rewriter, loc, loopRange.stride);
     auto loop = rewriter.create<scf::ForOp>(op.getLoc(), offsetVal, sizeVal,
                                             strideVal, ValueRange{});
     loops.push_back(loop);
     ivs.push_back(loop.getInductionVar());
     rewriter.setInsertionPoint(loop.getBody()->getTerminator());
   }
   if (failed(op.generateScalarImplementation(rewriter, op.getLoc(), ivs))) {
     return failure();
   }
-  rewriter.eraseOp(op);
   return loops;
 }

mlir/test/lib/Interfaces/TilingInterface/TestTilingInterface.cpp

 
 using namespace mlir;
 
 namespace {
 
 /// Pattern for testing `TileUsingSCFForOp` pattern (that tiles operations using
 /// the `TilingInterface` with `scf.for` ops for iterating over the tiles) while
 /// using a `filter` to avoid recursive application.
-struct TestTileUsingSCFForOpWithFilter : public scf::TileUsingSCFForOp {
-  TestTileUsingSCFForOpWithFilter(MLIRContext *context,
-                                  scf::SCFTilingOptions options,
-                                  linalg::LinalgTransformationFilter filter =
-                                      linalg::LinalgTransformationFilter(),
-                                  PatternBenefit benefit = 1)
-      : scf::TileUsingSCFForOp(context, std::move(options), benefit),
-        filter(std::move(filter)) {}
+struct TestTileUsingSCFForOp
+    : public OpInterfaceRewritePattern<TilingInterface> {
+  TestTileUsingSCFForOp(MLIRContext *context, scf::SCFTilingOptions options,
+      : OpInterfaceRewritePattern<TilingInterface>(context, benefit),
+        options(std::move(options)), filter(std::move(filter)) {}
 
   /// Construct a generic pattern applied to `opName`.
-  TestTileUsingSCFForOpWithFilter(StringRef opName, MLIRContext *context,
-                                  scf::SCFTilingOptions options,
-                                  linalg::LinalgTransformationFilter filter =
-                                      linalg::LinalgTransformationFilter(),
-                                  PatternBenefit benefit = 1)
-      : scf::TileUsingSCFForOp(context, std::move(options), benefit),
-        filter(std::move(filter)) {}
+  TestTileUsingSCFForOp(StringRef opName, MLIRContext *context,
+      : OpInterfaceRewritePattern<TilingInterface>(context, benefit),
+        options(std::move(options)), filter(std::move(filter)) {}
 
   LogicalResult matchAndRewrite(TilingInterface op,
                                 PatternRewriter &rewriter) const override {
     if (failed(filter.checkAndNotify(rewriter, op)))
       return failure();
 
-    auto tilingResult = returningMatchAndRewrite(op, rewriter);
-    if (failed(tilingResult)) {
-      return failure();
+    FailureOr<scf::SCFTilingResult> tilingResult =
+        scf::tileUsingSCFForOp(rewriter, op, options);
+    if (failed(tilingResult))
+      return rewriter.notifyMatchFailure(op, "failed to tile operation");
+
+    if (op->getNumResults()) {
+      rewriter.replaceOp(op, tilingResult->replacements);
+    } else {
+      rewriter.eraseOp(op);
     }
+
     filter.replaceLinalgTransformationFilter(rewriter, tilingResult->tiledOp);
     return success();
   }
 
 private:
+  scf::SCFTilingOptions options;
   linalg::LinalgTransformationFilter filter;
 };
 
 /// Pattern for testing `TileConsumerAndFuseProducersUsingSCFForOp` pattern
 /// (that tiles and fuses operations using the `TilingInterface` with `scf.for`
 /// ops for iterating over the tiles) while using a `filter` to avoid recursive
 /// application.
-struct TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter
-    : public scf::TileConsumerAndFuseProducersUsingSCFForOp {
-  TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter(
-      MLIRContext *context, scf::SCFTilingOptions options,
+struct TestTileConsumerAndFuseProducersGreedilyUsingSCFForOp
+    : public OpInterfaceRewritePattern<TilingInterface> {
+  TestTileConsumerAndFuseProducersGreedilyUsingSCFForOp(
+      MLIRContext *context, scf::SCFTileAndFuseOptions options,
       linalg::LinalgTransformationFilter filter =
           linalg::LinalgTransformationFilter(),
       PatternBenefit benefit = 1)
-      : scf::TileConsumerAndFuseProducersUsingSCFForOp(
-            context, std::move(options), benefit),
-        filter(std::move(filter)) {}
+      : OpInterfaceRewritePattern<TilingInterface>(context, benefit),
+        options(std::move(options)), filter(std::move(filter)) {}
 
   /// Construct a generic pattern applied to `opName`.
-  TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter(
-      StringRef opName, MLIRContext *context, scf::SCFTilingOptions options,
+  TestTileConsumerAndFuseProducersGreedilyUsingSCFForOp(
+      StringRef opName, MLIRContext *context,
+      scf::SCFTileAndFuseOptions options,
       linalg::LinalgTransformationFilter filter =
           linalg::LinalgTransformationFilter(),
       PatternBenefit benefit = 1)
-      : scf::TileConsumerAndFuseProducersUsingSCFForOp(
-            context, std::move(options), benefit),
-        filter(std::move(filter)) {}
+      : OpInterfaceRewritePattern<TilingInterface>(context, benefit),
+        options(std::move(options)), filter(std::move(filter)) {}
 
   LogicalResult matchAndRewrite(TilingInterface op,
                                 PatternRewriter &rewriter) const override {
     if (failed(filter.checkAndNotify(rewriter, op)))
       return failure();
 
-    auto tileAndFuseResult = returningMatchAndRewrite(op, rewriter);
+    FailureOr<scf::SCFTileAndFuseResult> tileAndFuseResult =
+        scf::tileConsumerAndFuseProducerGreedilyUsingSCFForOp(rewriter, op,
+                                                              options);
     if (failed(tileAndFuseResult)) {
       return failure();
     }
+    // Replace the tiled op with replacements.
+    SmallVector<Value> replacements(op->getNumResults());
+    for (auto result : llvm::enumerate(op->getResults())) {
+      replacements[result.index()] =
+          tileAndFuseResult->replacements.lookup(result.value());
+    }
+    rewriter.replaceOp(op, replacements);
+
     filter.replaceLinalgTransformationFilter(
         rewriter, tileAndFuseResult->tiledAndFusedOps.front());
     return success();
   }
 
 private:
+  scf::SCFTileAndFuseOptions options;
   linalg::LinalgTransformationFilter filter;
 };
 
+/// Pattern to lower operations that implement the `TilingInterface` to
+/// loops/scalar IR using `scf.for`.
+struct LowerToLoopsUsingSCFForOp
+    : public OpInterfaceRewritePattern<TilingInterface> {
+  using OpInterfaceRewritePattern<TilingInterface>::OpInterfaceRewritePattern;
+
+  /// `matchAndRewrite` implementation that returns the significant transformed
+  /// pieces of IR.
+  LogicalResult matchAndRewrite(TilingInterface op,
+                                PatternRewriter &rewriter) const override {
+    FailureOr<SmallVector<scf::ForOp>> loops =
+        scf::lowerToLoopsUsingSCFForOp(rewriter, op);
+    if (failed(loops))
+      return rewriter.notifyMatchFailure(op, "failed to lower to loops");
+    rewriter.eraseOp(op);
+    return loops;
+  }
+};
+
 /// Test pass for testing the use of `TilingInterface`.
 struct TestTilingInterfacePass
     : public PassWrapper<TestTilingInterfacePass, OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTilingInterfacePass)
 
   TestTilingInterfacePass() = default;
   TestTilingInterfacePass(const TestTilingInterfacePass &pass)
       : PassWrapper(pass) {}
struct TestTilingInterfacePass
 
   void runOnOperation() override;
 
 private:
   void addTestPatterns(MLIRContext *context, RewritePatternSet &patterns);
 };
 } // namespace
 
-template <class Pattern>
-static void
-addPatternForTiling(MLIRContext *context, RewritePatternSet &patterns,
-                    StringRef filterName, ArrayRef<int64_t> tileSizes,
+static void addPatternForTiling(MLIRContext *context,
+                                RewritePatternSet &patterns,
+                                StringRef filterName,
+                                ArrayRef<int64_t> tileSizes,
                                 ArrayRef<unsigned> interchange = {}) {
   scf::SCFTilingOptions tilingOptions;
   tilingOptions.setTileSizes(tileSizes).setInterchange(interchange);
   linalg::LinalgTransformationFilter filter(
       StringAttr::get(context, filterName), StringAttr::get(context, "tiled"));
-  patterns.add<Pattern>(context, tilingOptions, filter);
+  patterns.add<TestTileUsingSCFForOp>(context, tilingOptions, filter);
+}
+
+static void addPatternForTileAndFuse(MLIRContext *context,
+                                     RewritePatternSet &patterns,
+                                     StringRef filterName,
+                                     ArrayRef<int64_t> tileSizes,
+                                     ArrayRef<unsigned> interchange = {}) {
+  scf::SCFTileAndFuseOptions tileAndFuseOptions;
+  tileAndFuseOptions.tilingOptions.setTileSizes(tileSizes).setInterchange(
+      interchange);
+  linalg::LinalgTransformationFilter filter(
+      StringAttr::get(context, filterName), StringAttr::get(context, "tiled"));
+  patterns.add<TestTileConsumerAndFuseProducersGreedilyUsingSCFForOp>(
+      context, tileAndFuseOptions, filter);
 }
 
 void TestTilingInterfacePass::addTestPatterns(MLIRContext *context,
                                               RewritePatternSet &patterns) {
   if (testTiling) {
     // 1. Tiling M and N dims of `linalg.matmul` on tensors.
-    addPatternForTiling<TestTileUsingSCFForOpWithFilter>(
-        context, patterns, "simple_gemm", {10, 20});
+    addPatternForTiling(context, patterns, "simple_gemm", {10, 20});
     // 2. Tiling M, N and K of `linalg.matmul` on buffers.
-    addPatternForTiling<TestTileUsingSCFForOpWithFilter>(
-        context, patterns, "simple_gemm_memref", {10, 20, 30});
+    addPatternForTiling(context, patterns, "simple_gemm_memref", {10, 20, 30});
     // 3. Tiling 3D parallel generic op which implements a transpose
-    addPatternForTiling<TestTileUsingSCFForOpWithFilter>(
-        context, patterns, "parallel_generic_transpose", {10, 0, 20});
+    addPatternForTiling(context, patterns, "parallel_generic_transpose",
+                        {10, 0, 20});
     // 4. Tiling 2D conv op.
-    addPatternForTiling<TestTileUsingSCFForOpWithFilter>(
-        context, patterns, "simple_conv", {0, 0, 0, 0, 10, 20, 30});
+    addPatternForTiling(context, patterns, "simple_conv",
+                        {0, 0, 0, 0, 10, 20, 30});
     // 5. Tiling a simple op with `linalg.index` inside.
-    addPatternForTiling<TestTileUsingSCFForOpWithFilter>(
-        context, patterns, "indexed_semantics", {10, 20});
+    addPatternForTiling(context, patterns, "indexed_semantics", {10, 20});
     // 6. Tiling + interchange of an operation
-    addPatternForTiling<TestTileUsingSCFForOpWithFilter>(
-        context, patterns, "gemm_interchange", {10, 20, 30}, {1, 2, 0});
+    addPatternForTiling(context, patterns, "gemm_interchange", {10, 20, 30},
+                        {1, 2, 0});
     // 7. Tiling for 2D pad tensor operations.
-    addPatternForTiling<TestTileUsingSCFForOpWithFilter>(
-        context, patterns, "pad_2dtiling", {2, 3});
+    addPatternForTiling(context, patterns, "pad_2dtiling", {2, 3});
     // 8. Tiling inner dimension of 2d pad tensor operations.
-    addPatternForTiling<TestTileUsingSCFForOpWithFilter>(
-        context, patterns, "pad_inner_tiling", {0, 3});
+    addPatternForTiling(context, patterns, "pad_inner_tiling", {0, 3});
     // 9. Tiling inner dimension of 2d pad tensor operations.
-    addPatternForTiling<TestTileUsingSCFForOpWithFilter>(
-        context, patterns, "pad_outer_tiling", {2, 3});
+    addPatternForTiling(context, patterns, "pad_outer_tiling", {2, 3});
 
     return;
   }
   if (testTileConsumerAndFuseProducer) {
-    // 1. Tile and fuse of gemm with bias-add operation.
-    addPatternForTiling<
-        TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter>(
-        context, patterns, "fusion", {10, 20});
-    addPatternForTiling<
-        TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter>(
-        context, patterns, "gemm_fusion", {10});
-    addPatternForTiling<
-        TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter>(
-        context, patterns, "gemm_interchange_fusion", {10, 20}, {1, 0});
-    addPatternForTiling<
-        TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter>(
-        context, patterns, "gemm_plus_gemm_fusion", {10, 20});
-    addPatternForTiling<
-        TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter>(
-        context, patterns, "gemm_sequence_fusion", {10});
+    // 1. Tile and fuse of gemm with fill producer and bias-add consumer.
+    addPatternForTileAndFuse(context, patterns, "fusion", {10, 20});
+    // 2. Tile and fuse sequence of GEMMs, by fusing only along M.
+    addPatternForTileAndFuse(context, patterns, "gemm_fusion", {10});
+    // 3. Tile and fuse gemm with consumer + interchange of tiled loops.
+    addPatternForTileAndFuse(context, patterns, "gemm_interchange_fusion",
+                             {10, 20}, {1, 0});
+    // 4. Tile and fuse matmul + transpose(matmul). Will introduce redundant
+    // computations.
+    addPatternForTileAndFuse(context, patterns, "gemm_plus_gemm_fusion",
+                             {10, 20});
+    // 5. Tile and fuse a sequence of GEMMs by tiling and fusing only along M
+    // dimension.
+    addPatternForTileAndFuse(context, patterns, "gemm_sequence_fusion", {10});
     return;
   }
   if (testLoweringToScalar) {
-    patterns.add<scf::LowerToLoopsUsingSCFForOp>(context);
+    patterns.add<LowerToLoopsUsingSCFForOp>(context);
   }
 }
 
 void TestTilingInterfacePass::runOnOperation() {
   MLIRContext *context = &getContext();
 
   RewritePatternSet tilingPatterns(context);
   addTestPatterns(context, tilingPatterns);