[mlir][TilingInterface] Modify `TilingInterface` methods to better return the state of the transformed IR.

[mravishankar]

Summary

Currently the getTiledImplementation and generateResultTileValue
return just SmallVector<Operation *> and FailureOr<Value>.

• For getTiledImplementation returning empty implies tiling wasnt done. There is also an implicit assumption that the tiled operation results correspond to the tiled values of the result of the original operation. This cannot handle cases where the tiled implementation might use multiple operations to compute the tiled value for the results of the untiled operation. Sometimes, the tiled operation might not directly give the tiled values, and might require casts, etc to get a replacement.
• For generateResultTileValue, it is assumed that the op defining the returned Value is the operation that represents the tiled computation. Again presence of casts, etc violate this.

Instead make these methods return

struct TilingResult {
  SmallVector<Operation *> tiledOps;
  SmallVector<Value> tiledValues;
};

The tiledOps represent the operations generated that are relevant
for subsequent transformations. The tiledValues represent the tiled
values for the results of the original operation. This better
transmits the state of the transformed IR.

As a consequence the following methods also return FailureOr<TilingResult>

• tensor::replaceExtractSliceWithTiledProducer
• tensor::bubbleUpPadSlice

[mravishankar]

Fix lit tests.

[pifon2a]

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

53–58

return tensor::bubbleUpPadSlice(b, cast<PadOp>(op), offsets, sizes);?

[pifon2a]

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

440–444

and here

[rengolin]

Overall looks good to me, thanks!

I have some inline comments but carrying out the information makes the code more readable and robust.

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

531–532

assert(tileValues.size() == 1)?

596–599

Is this tested with more than one tiled Ops?

The return below just takes the first value?

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

53–58

Then why not just call bubblePadOp directly? That cast<PadOp> is looking like it could crash on incorrect wrapper usage anyway.

416–417

This is a common pattern, could you have a constructor for TilingResult(Operation*) and avoid the repetition?

440–444

Are these just wrappers on top of wrappers? Do we need them all?

666

And here... :) etc.

[hanchung]

mlir/include/mlir/Interfaces/TilingInterface.h

26–34

Can we add a document about the order in tiledOps? In SCFTilingResult, we have the doc about the order:

/// Transformation information returned after tiling.
struct SCFTilingResult {
  /// Tiled operations that are generated during tiling. The order does not
  /// matter except the last op. The replacements are expected to be the results
  /// of the last op.
  SmallVector<Operation *> tiledOps;
  /// 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;
};

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

395–404

not related to this PR itself.. I've seen this code and feel that the assertion is not necessary. The casting to TilingInterface is not a requirement for the check. I can send out a fix for this part and the others if people think it's reasonable.

[mravishankar]

Rebase

[mravishankar]

Rebase again.

[ThomasRaoux]

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

480

This should be tileAndFuseResult->tiledValues[0] right?

It looks like we are missing a test where we fuse an op with multiple results used.
There is the following test but only the result 0 is used so this missed to catch this problems:
https://github.com/llvm/llvm-project/blob/ff11d6b6f6e27f5de389002b8f6102b6cf3ed474/mlir/test/Dialect/Linalg/transform-op-fuse-into-containing.mlir#L208

[mravishankar]

Rebase.

mlir/include/mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_TENSOR_IR_TENSORTILINGINTERFACEIMPL_H_
 #define MLIR_DIALECT_TENSOR_IR_TENSORTILINGINTERFACEIMPL_H_
 
 #include "mlir/IR/Dialect.h"
 
 namespace mlir {
+
+struct TilingResult;
+
 namespace tensor {
 
 class PadOp;
 
 /// Bubbles up a slice of this pad by taking the slice first and then performing
 /// the padding. `offsets` and `strides` specifies each dimension's start offset
 /// and size for the slice. The slice has unit strides along all dimensions.
 ///
 /// Specifically, this function converts:
 /// ```
 /// %0 = tensor.pad %source low[...] high[...] { linalg.yield %cst }
 /// %1 = <extract-slice> %0 offsets=[...], sizes[...]
 /// ```
 /// into
 /// ```
 /// %0 = tensor.extract_slice %source ...
 /// %0 = tensor.pad %0 low[...] high[...] { linalg.yield %cst }
 /// ```
 ///
 /// If `generateZeroSliceGuard` is true, the generated IR will contain logic
 /// to guard against the case that we might take a zero-sized slice from the
 /// original source. For such cases, we `tensor.generate` to generate the
 /// full tensor.
-Operation *bubbleUpPadSlice(OpBuilder &b, tensor::PadOp padOp,
+FailureOr<TilingResult> bubbleUpPadSlice(OpBuilder &b, tensor::PadOp padOp,
                                          ArrayRef<OpFoldResult> offsets,
                                          ArrayRef<OpFoldResult> sizes,
                                          bool generateZeroSliceGuard = true);
 
 /// Registers external models for Tiling interface for tensor ops.
 /// Currently, it registers:
 ///
 /// * TilingInterface for `tensor.pad`, `tensor.pack`, and `tensor.unpack`.
 ///
 /// Unfortunately, a "normal" internal registration is not possible at the
 /// moment, because of the dependency of the interface implementation for these

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

 //===- Transforms.h - Tensor Transformation Patterns ------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMS_H
 #define MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMS_H
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
+
+struct TilingResult;
+
 namespace tensor {
 
 /// Populates `patterns` with patterns to wrap a tensor.pad op with an scf.if op
 /// to separate the cases where we don't need padding (all pad sizes are
 /// actually zeros) and where we indeed need padding.
 void populateSplitPaddingPatterns(RewritePatternSet &patterns,
                                   PatternBenefit baseBenefit = 1);
 
 /// Pattern to swap an `tensor.extract_slice` with its producer when the
 /// producer implements the `TilingInterface`. The pattern itself does not
 /// provide a mechanism to control where the application happens. With use of
 /// transform dialect that control is done within the transform dialect. Other
 /// use cases can inherit from this pattern and add necessary controls.
-FailureOr<Value> replaceExtractSliceWithTiledProducer(
+FailureOr<TilingResult> replaceExtractSliceWithTiledProducer(
     OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producerOp);
 
 /// Collects patterns to merge consecutive tensor.insert_slice/extract_slice
 /// into one. These patterns are in in this separate entry point because the
 /// bufferization is sensitive over IR structure, particularly those
 /// tensor.extract_slice and tensor.insert_slice ops for creating the slices.
 void populateMergeConsecutiveInsertExtractSlicePatterns(
     RewritePatternSet &patterns);

mlir/include/mlir/Interfaces/TilingInterface.h

 
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "mlir/Support/LLVM.h"
 
+namespace mlir {
+
+/// Container for result values of tiling.
+/// - `tiledOps` contains operations created by the tiling implementation that
+/// are returned to the caller for further transformations.
+/// - `tiledValues` contains the tiled value corresponding to the result of the
+/// untiled operation.
+struct TilingResult {
+  SmallVector<Operation *> tiledOps;
+  SmallVector<Value> tiledValues;
+};
+
+} // namespace mlir
+
 /// Include the ODS generated interface header files.
 #include "mlir/Interfaces/TilingInterface.h.inc"
 
 #endif // MLIR_INTERFACES_TILINGINTERFACE_H_

mlir/include/mlir/Interfaces/TilingInterface.td

let methods = [
             dimension had non-zero offset, it must be included in the offset
             provided here (as opposed to zero-based offset "relative" to the
             iteration space).
           - `sizes` provides the size of the tile.
 
           The method returns the operation that is the tiled
           implementation.
         }],
-        /*retType=*/"SmallVector<Operation *>",
+        /*retType=*/"FailureOr<TilingResult>",
         /*methodName=*/"getTiledImplementation",
         /*args=*/(ins
             "OpBuilder &":$b,
             "ArrayRef<OpFoldResult> ":$offsets,
             "ArrayRef<OpFoldResult> ":$sizes),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return {};
let methods = [
 
           - `offsets` provides the offset of the tile in the coordinate system
             of the original iteration space, i.e., if an iteration space
             dimension had non-zero offset, it must be included in the offset
             provided here (as opposed to zero-based offset "relative" to the
             iteration space).
           - `sizes` provides the size of the tile.
         }],
-        /*retType=*/"FailureOr<Value>",
+        /*retType=*/"FailureOr<TilingResult>",
         /*methodName=*/"generateResultTileValue",
         /*args=*/(ins
           "OpBuilder &":$b,
           "unsigned":$resultNumber,
           "ArrayRef<OpFoldResult>":$offsets,
           "ArrayRef<OpFoldResult>":$sizes),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

void transform::FuseIntoContainingOp::build(OpBuilder &builder,
                                             Value containingOp) {
   result.addOperands({producerOp, containingOp});
   result.addTypes(pdl::OperationType::get(builder.getContext()));
 }
 
 /// Find the first "extract" user of `producerOp` and tile it right before its
 /// use. The tiled op is fused under the `containingOp`.
 /// Return this fused op on success or nullptr if anything fails.
-static Operation *tileAndFuseFirstExtractUse(RewriterBase &rewriter,
-                                             Diagnostic &diag,
-                                             Operation *producerOp,
-                                             Operation *containingOp) {
+static SmallVector<Operation *>
+tileAndFuseFirstExtractUse(RewriterBase &rewriter, Diagnostic &diag,
+                           Operation *producerOp, Operation *containingOp) {
   LLVM_DEBUG(DBGS() << "Try to fuse a direct extract use\n");
   auto tileableProducer = dyn_cast<TilingInterface>(producerOp);
   if (!tileableProducer) {
     diag.attachNote(producerOp->getLoc())
         << "producer is not a TileableInterface: " << *producerOp;
-    return nullptr;
+    return {};
   }
 
   // Search the producer slices accessed within the containing operation.
   // TODO: Generalize to more extract/insert/parallel_insert triples, maybe
   // evolve into an interface.
   auto it = llvm::find_if(tileableProducer->getUsers(), [&](Operation *user) {
     auto sliceOp = dyn_cast<tensor::ExtractSliceOp>(user);
     return sliceOp && containingOp->isProperAncestor(sliceOp);
   });
 
   // Find a fusion opportunity.
   if (it == tileableProducer->getUsers().end()) {
     diag.attachNote(tileableProducer->getLoc())
         << "could not find fusion opportunity for: " << *tileableProducer;
-    return nullptr;
+    return {};
   }
   auto sliceOpToTile = cast<tensor::ExtractSliceOp>(*it);
 
   // Try to fuse the producer in-place.
   OpBuilder::InsertionGuard guard(rewriter);
   rewriter.setInsertionPoint(sliceOpToTile);
 
   // Tile the producer.
   int64_t resultNumber =
       sliceOpToTile.getSource().cast<OpResult>().getResultNumber();
   LLVM_DEBUG(DBGS() << "resultNumber: " << resultNumber << "\n");
 
-  FailureOr<Value> tiledProducer = tileableProducer.generateResultTileValue(
-      rewriter, resultNumber, sliceOpToTile.getMixedOffsets(),
+  FailureOr<TilingResult> tileAndFuseResult =
+      tileableProducer.generateResultTileValue(rewriter, resultNumber,
+                                               sliceOpToTile.getMixedOffsets(),
                                                sliceOpToTile.getMixedSizes());
-  if (failed(tiledProducer)) {
+  if (failed(tileAndFuseResult)) {
     diag.attachNote(tileableProducer->getLoc())
         << "failed to tile producer op: " << *tileableProducer;
-    return nullptr;
+    return {};
+  }
+  for (auto tiledOp : tileAndFuseResult->tiledOps) {
+    LLVM_DEBUG(DBGS() << "tiledProducer: " << *tiledOp << "\n");
   }
-  LLVM_DEBUG(DBGS() << "tiledProducer: " << *tiledProducer << "\n");
 
   // Replace the extract op.
-  Operation *fusedOp = tiledProducer->getDefiningOp();
   auto maybeRankReduced = tensor::ExtractSliceOp::rankReduceIfNeeded(
-      rewriter, sliceOpToTile->getLoc(), fusedOp->getResult(resultNumber),
+      rewriter, sliceOpToTile->getLoc(),
+      tileAndFuseResult->tiledValues[resultNumber],
       sliceOpToTile->getResult(0)
           .getType()
           .cast<RankedTensorType>()
           .getShape());
   assert(succeeded(maybeRankReduced) && "unexpected shape");
   rewriter.replaceOp(sliceOpToTile, *maybeRankReduced);
-  return fusedOp;
+  return tileAndFuseResult->tiledOps;
 }
 
 /// First, find the first "scf::ForallOp" user of `producerOp` and ensure
 /// it is exactly the `containingOp`, otherwise bail.
 /// Then, find the first "extract" user of the tied block argument and tile it
 /// right before its "extract" use. The tiled op is fused under the
 /// `containingOp`.
 /// Return this fused op on success or nullptr if anything fails.
-static Operation *tileAndFuseFirstExtractUseThroughContainingOpBlockArgument(
+static SmallVector<Operation *>
+tileAndFuseFirstExtractUseThroughContainingOpBlockArgument(
     RewriterBase &rewriter, Diagnostic &diag, Operation *producerOp,
     Operation *containingOp) {
   LLVM_DEBUG(DBGS() << "Try to fuse an extract use through block argument\n");
 
   auto tileableProducer = dyn_cast<TilingInterface>(producerOp);
   if (!tileableProducer) {
     diag.attachNote(producerOp->getLoc())
         << "producer is not a TileableInterface: " << *producerOp;
-    return nullptr;
+    return {};
   }
 
   // Search the first use by a "scf::ForallOp" user.
   scf::ForallOp forallOp;
   auto itProducerUses =
       llvm::find_if(tileableProducer->getUses(), [&](OpOperand &use) {
         forallOp = dyn_cast<scf::ForallOp>(use.getOwner());
         return forallOp;
       });
   // If it's not from the containing op, return.
   if (!forallOp || forallOp != containingOp) {
     diag.attachNote(tileableProducer->getLoc())
         << "could not find a use by the containing op: " << *tileableProducer;
-    return nullptr;
+    return {};
   }
 
   // Search the producer slices accessed within the containing
   // operation.
   // TODO: Generalize to more extract/insert/parallel_insert triples.
   //   Maybe evolve into an interface.
   OpOperand *pUse = &(*itProducerUses);
   BlockArgument bbArg = forallOp.getTiedBlockArgument(pUse);
 
   // Search the producer slices accessed within the containing operation.
   // TODO: Generalize to more extract/insert/parallel_insert triples, maybe
   // evolve into an interface.
   auto itBBArgUsers = llvm::find_if(bbArg.getUsers(), [&](Operation *user) {
     auto sliceOp = dyn_cast<tensor::ExtractSliceOp>(user);
     return sliceOp && containingOp->isProperAncestor(sliceOp);
   });
 
   // Find a fusion opportunity.
   if (itBBArgUsers == bbArg.getUsers().end()) {
     diag.attachNote(containingOp->getLoc())
         << "could not find fusion opportunity for bbArg: " << bbArg;
-    return nullptr;
+    return {};
   }
   auto sliceOpToTile = cast<tensor::ExtractSliceOp>(*itBBArgUsers);
 
   // Try to fuse the producer in-place.
   OpBuilder::InsertionGuard guard(rewriter);
   rewriter.setInsertionPoint(sliceOpToTile);
 
   // Replace the use in the tileableProducer before tiling: clone, replace and
   // then tile.
   int64_t resultNumber = pUse->get().cast<OpResult>().getResultNumber();
   LLVM_DEBUG(DBGS() << "resultNumber: " << resultNumber << "\n");
 
   // Gather destination tensors.
   SmallVector<Value> destinationTensors;
   if (failed(tensor::getOrCreateDestinations(
           rewriter, tileableProducer->getLoc(), tileableProducer,
           destinationTensors))) {
     diag.attachNote(tileableProducer->getLoc())
         << "failed to get destination tensors for: " << *tileableProducer;
-    return nullptr;
+    return {};
   }
 
   IRMapping bvm;
   bvm.map(destinationTensors[resultNumber], bbArg);
   auto tileableProducerClone =
       cast<TilingInterface>(rewriter.clone(*tileableProducer, bvm));
   auto scopeGuard =
       llvm::make_scope_exit([&]() { rewriter.eraseOp(tileableProducerClone); });
 
   // Tile the producer.
-  FailureOr<Value> tiledProducer =
+  FailureOr<TilingResult> tileAndFuseResult =
       tileableProducerClone.generateResultTileValue(
           rewriter, resultNumber, sliceOpToTile.getMixedOffsets(),
           sliceOpToTile.getMixedSizes());
-  if (failed(tiledProducer)) {
+  if (failed(tileAndFuseResult)) {
     diag.attachNote(tileableProducer->getLoc())
         << "failed to tile producer op: " << *tileableProducer;
-    return nullptr;
+    return {};
   }
-  LLVM_DEBUG(DBGS() << "tiledProducer: " << *tiledProducer << "\n");
 
   // Replace the extract op.
-  Operation *fusedOp = tiledProducer->getDefiningOp();
   auto maybeRankReduced = tensor::ExtractSliceOp::rankReduceIfNeeded(
-      rewriter, sliceOpToTile->getLoc(), fusedOp->getResult(resultNumber),
+      rewriter, sliceOpToTile->getLoc(), tileAndFuseResult->tiledValues[0],
       sliceOpToTile->getResult(0)
           .getType()
           .cast<RankedTensorType>()
           .getShape());
   assert(succeeded(maybeRankReduced) && "unexpected shape");
   rewriter.replaceOp(sliceOpToTile, *maybeRankReduced);
 
   // Replace the use in containingOp.
   rewriter.updateRootInPlace(containingOp, [&]() {
     containingOp->setOperand(pUse->getOperandNumber(),
                              destinationTensors.front());
   });
 
-  return fusedOp;
+  return tileAndFuseResult->tiledOps;
 }
 
 static Operation *cloneAndFuseFirstUse(RewriterBase &rewriter, Diagnostic &diag,
                                        Operation *producerOp,
                                        Operation *containingOp) {
   LLVM_DEBUG(DBGS() << "Try to fuse an use by cloning\n");
 
   // Gather all uses inside the containing op.
while (!remainingProducers.empty()) {
     Diagnostic diag(producerOp->getLoc(), DiagnosticSeverity::Remark);
     diag << "could not fuse " << *producerOp << " into " << *containingOp;
 
     // TODO: If there are multiple uses of the producer in the containing op,
     // we currently tile/clone the op multiple times (once per use). In some
     // cases, we can tile/clone once and reuse the value for each use.
     // Futhermore, producers should then be traversed according to a
     // topological sorting.
-    Operation *tiled =
+    SmallVector<Operation *> tiledOps =
         tileAndFuseFirstExtractUse(rewriter, diag, producerOp, containingOp);
-    if (tiled) {
+    if (!tiledOps.empty()) {
       LLVM_DEBUG(DBGS() << "\nFused a direct extract use\n" << *containingOp);
-      fusedOps.push_back(tiled);
+      fusedOps.append(tiledOps);
       continue;
     }
 
-    Operation *tiledContainingOpOperand =
+    SmallVector<Operation *> tiledContainingOpOperand =
         tileAndFuseFirstExtractUseThroughContainingOpBlockArgument(
             rewriter, diag, producerOp, containingOp);
-    if (tiledContainingOpOperand) {
+    if (!tiledContainingOpOperand.empty()) {
       LLVM_DEBUG(DBGS() << "\nFused an extract use through block argument\n"
                         << *containingOp);
-      fusedOps.push_back(tiledContainingOpOperand);
+      fusedOps.append(tiledContainingOpOperand);
       continue;
     }
 
     Operation *cloned =
         cloneAndFuseFirstUse(rewriter, diag, producerOp, containingOp);
     if (cloned) {
       LLVM_DEBUG(DBGS() << "\nFused an use by cloning\n" << *containingOp);
       fusedOps.push_back(cloned);

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

createSplitPart(RewriterBase &b, Location loc, TilingInterface op,
                 SmallVectorImpl<Value> &results) {
   // Iteration space of the current part.
   SmallVector<OpFoldResult> sizesCopy = llvm::to_vector(sizes);
   SmallVector<OpFoldResult> offsetsCopy = llvm::to_vector(offsets);
   sizesCopy[dimension] = size;
   offsetsCopy[dimension] = offset;
 
   // Create the part as it it were a single tile.
-  SmallVector<Operation *> tiled =
+  FailureOr<TilingResult> tilingResult =
       op.getTiledImplementation(b, offsetsCopy, sizesCopy);
-  assert(tiled.size() == 1 && "expected a single result from tiling");
-  auto part = cast<TilingInterface>(tiled.front());
 
   // Insert the results back and populate the `results` list.
-  for (auto i : llvm::seq<unsigned>(0, part->getNumResults())) {
+  for (auto [index, result] : llvm::enumerate(tilingResult->tiledValues)) {
     SmallVector<OpFoldResult> resultOffsets, resultSizes;
-    if (failed(op.getResultTilePosition(b, i, offsetsCopy, sizesCopy,
+    if (failed(op.getResultTilePosition(b, index, offsetsCopy, sizesCopy,
                                         resultOffsets, resultSizes)))
       return nullptr;
     SmallVector<OpFoldResult> resultStrides(resultOffsets.size(),
                                             b.getIndexAttr(1));
     Value inserted = b.create<tensor::InsertSliceOp>(
-        loc, part->getResult(i), resultOperands[i], resultOffsets, resultSizes,
+        loc, result, resultOperands[index], resultOffsets, resultSizes,
         resultStrides);
     results.push_back(inserted);
   }
-
-  return part;
+  // TODO: this part can be generalized maybe to not expect a single op.
+  assert(tilingResult->tiledOps.size() == 1 &&
+         "expected split part to return a single tiled operation");
+  return cast<TilingInterface>(tilingResult->tiledOps[0]);
 }
 
 std::pair<TilingInterface, TilingInterface>
 linalg::splitOp(RewriterBase &rewriter, TilingInterface op, unsigned dimension,
                 OpFoldResult splitPoint) {
   // Compute the iteration space.
   SmallVector<Range> iterationSpace = op.getIterationDomain(rewriter);
 

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

if (destinationStyleOp) {
         if (it == dest.end())
           return op->emitOpError("must have \"tensor semantic\" for tiling");
         unsigned destNum = std::distance(dest.begin(), it);
         outOperand->set(destBbArgs[destNum]);
       }
     }
 
     // 4. Tile the cloned op and delete the clone.
-    SmallVector<Operation *> tiledOps =
+    FailureOr<TilingResult> tilingResult =
         cast<TilingInterface>(clonedOp).getTiledImplementation(b, tiledOffsets,
                                                                tiledSizes);
     b.eraseOp(clonedOp);
-    assert(tiledOps.size() == 1 && "expected a single produced tiled op");
-    tiledOp = tiledOps.front();
+    assert(tilingResult->tiledOps.size() == 1 &&
+           "expected a single produced tiled op");
+    tiledOp = tilingResult->tiledOps.front();
   }
 
   // 5. Parallel insert back into the result tensor.
   auto tilingInterfaceOp = dyn_cast<TilingInterface>(tiledOp);
   assert(tilingInterfaceOp && "Tiled op does not implement TilingInterface");
   for (auto it : llvm::zip(llvm::seq(unsigned(0), unsigned(dest.size())),
                            tilingInterfaceOp->getResults(), destBbArgs)) {
     // 5.a. Partial subset information is inserted just before the terminator.
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(forallOp.getTerminator());
 
     SmallVector<OpFoldResult> resultOffsets, resultSizes;
     if (failed(op.getResultTilePosition(b, std::get<0>(it), tiledOffsets,
                                         tiledSizes, resultOffsets,
                                         resultSizes)))
b.updateRootInPlace(clonedOp, [&]() {
                cast<DestinationStyleOpInterface>(clonedOp).getDpsInitOperands(),
                tiledDpsInitOperands)) {
         initOperandPtr->set(tiledInitValue);
       }
     });
 
     // 5. Tile the cloned op and delete the clone.
     if (tileSizes.empty()) {
-      SmallVector<Operation *> tiledOps =
+      FailureOr<TilingResult> tilingResult =
           cast<TilingInterface>(clonedOp).getTiledImplementation(
               b, tiledOffsets, tiledSizes);
-      assert(tiledOps.size() == 1 && "expected a single produced tiled op");
-      tiledOp = tiledOps.front();
-      tilingResults = tiledOp->getResults();
+      assert(tilingResult->tiledOps.size() == 1 &&
+             "expected a single produced tiled op");
+      tiledOp = tilingResult->tiledOps.front();
+      tilingResults = tilingResult->tiledValues;
     } else {
       LinalgTilingOptions options;
       FailureOr<TiledLinalgOp> maybeTiled = tileLinalgOpImpl<scf::ForOp>(
           b, cast<LinalgOp>(clonedOp), tileSizes, options);
       if (failed(maybeTiled))
         return b.notifyMatchFailure(op, "failed tileLinalgOpImpl");
 
       SmallVector<Value> ids = forallOp.getInductionVars();

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

return llvm::to_vector(
         llvm::map_range(map.getResults(), [&](AffineExpr loopExpr) {
           OpFoldResult ofr =
               makeComposedFoldedAffineApply(b, loc, loopExpr, allShapesSizes);
           return Range{b.getIndexAttr(0), ofr, b.getIndexAttr(1)};
         }));
   }
 
   // Instantiate the tiled implementation of the operation.
-  SmallVector<Operation *>
+  FailureOr<TilingResult>
   getTiledImplementation(Operation *op, OpBuilder &b,
                          ArrayRef<OpFoldResult> offsets,
                          ArrayRef<OpFoldResult> sizes) const {
     // Leave the `sizeBounds` value empty. That is only needed when the `sizes`
     // specified could lead to out of bounds accesses.
     Location loc = op->getLoc();
     LinalgOp linalgOp = cast<LinalgOp>(op);
     SmallVector<Value> valuesToTile = linalgOp->getOperands();
     SmallVector<Value, 4> tiledOperands = makeTiledShapes(
         b, loc, linalgOp, valuesToTile, offsets, sizes, {}, true);
 
     SmallVector<Type> resultTensorTypes =
         getTensorOutputTypes(linalgOp, tiledOperands);
 
     Operation *tiledOp = clone(b, linalgOp, resultTensorTypes, tiledOperands);
     offsetIndices(b, cast<LinalgOp>(tiledOp), offsets);
 
-    return {tiledOp};
+    return TilingResult{{tiledOp}, SmallVector<Value>(tiledOp->getResults())};
   }
 
   // Return the details of the output tile generated by the tiled
   // implementation.
   LogicalResult
   getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
                         ArrayRef<OpFoldResult> offsets,
                         ArrayRef<OpFoldResult> sizes,
SliceParameters sliceParams = computeSliceParameters(
         b, loc, outOperand->get(), sizes,
         linalgOp.getMatchingIndexingMap(outOperand), offsets,
         /*ubs*/ {}, subShapeSizes, true);
     resultOffsets = sliceParams.offsets;
     resultSizes = sliceParams.sizes;
     return success();
   }
 
-  FailureOr<Value> generateResultTileValue(Operation *op, OpBuilder &b,
-                                           unsigned resultNumber,
-                                           ArrayRef<OpFoldResult> offsets,
-                                           ArrayRef<OpFoldResult> sizes) const {
+  FailureOr<TilingResult>
+  generateResultTileValue(Operation *op, OpBuilder &b, unsigned resultNumber,
     auto linalgOp = cast<LinalgOp>(op);
 
     // Check that the indexing map used for the output is a projected
     // permutation. This could be relaxed with a more general approach that can
     // map the offsets and sizes from the result to iteration space tiles
     // (filling in full extent for dimensions not used to access the result).
     AffineMap indexingMap =
         linalgOp.getIndexingMapMatchingResult(op->getResult(resultNumber));
generateResultTileValue(Operation *op, OpBuilder &b, unsigned resultNumber,
     }
     for (const auto &resultExpr : llvm::enumerate(indexingMap.getResults())) {
       unsigned dimPosition =
           resultExpr.value().cast<AffineDimExpr>().getPosition();
       iterationTileOffsets[dimPosition] = offsets[resultExpr.index()];
       iterationTileSizes[dimPosition] = sizes[resultExpr.index()];
     }
 
-    SmallVector<Operation *> tiledOp = tilingInterfaceOp.getTiledImplementation(
-        b, iterationTileOffsets, iterationTileSizes);
-    if (tiledOp.size() != 1)
+    FailureOr<TilingResult> tilingResult =
+        tilingInterfaceOp.getTiledImplementation(b, iterationTileOffsets,
+                                                 iterationTileSizes);
+    if (tilingResult->tiledOps.size() != 1)
       return op->emitOpError("failed to generate tiled implementation");
 
-    return tiledOp[0]->getResult(resultNumber);
+    return TilingResult{
+        tilingResult->tiledOps,
+        SmallVector<Value>{tilingResult->tiledValues[resultNumber]}};
   }
 
   LogicalResult generateScalarImplementation(Operation *op, OpBuilder &builder,
                                              Location loc,
                                              ValueRange ivs) const {
     auto linalgOp = cast<LinalgOp>(op);
     if (!linalgOp.hasBufferSemantics())
       return op->emitOpError("expected operation to have buffer semantics");

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

LogicalResult ExtractSliceOfPadTensorSwapPattern::matchAndRewrite(
   bool zeroSliceGuard = true;
   if (controlFn) {
     if (std::optional<bool> control = controlFn(sliceOp))
       zeroSliceGuard = *control;
     else
       return failure();
   }
 
-  Operation *tiledPadOp =
+  FailureOr<TilingResult> tilingResult =
       tensor::bubbleUpPadSlice(rewriter, padOp, sliceOp.getMixedOffsets(),
                                sliceOp.getMixedSizes(), zeroSliceGuard);
+  if (failed(tilingResult))
+    return failure();
   // All shapes are static and the data source is actually used. Rewrite into
   // pad(extract_slice(x)).
-  rewriter.replaceOp(sliceOp, tiledPadOp->getResults());
+  rewriter.replaceOp(sliceOp, tilingResult->tiledValues);
   return success();
 }
 
 /// Returns a tensor.pad op if padding value is set. Otherwise, returns the
 /// source directly. The method assumes that the `packOp` has static shapes.
 static Value getPackOpSourceOrPaddedSource(OpBuilder &builder,
                                            tensor::PackOp packOp) {
   Value input = packOp.getSource();

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

updateDestinationOperandsForTiledOp(OpBuilder &builder,
   }
 }
 
 /// Helper method to yield the values of the tiled op, as well as
 /// update the destination operands of the tiled op, if it is
 /// a destination passing style op.
 static SmallVector<Value>
 yieldTiledValues(RewriterBase &rewriter, ArrayRef<Value> initValues,
-                 Operation *tiledOp,
+                 TilingResult tilingResult,
                  ArrayRef<SmallVector<OpFoldResult>> tileOffsetsList,
                  ArrayRef<SmallVector<OpFoldResult>> tileSizesList,
                  MutableArrayRef<scf::ForOp> loops) {
   SmallVector<Value> replacements =
-      yieldTiledValues(rewriter, initValues, tiledOp->getResults(),
+      yieldTiledValues(rewriter, initValues, tilingResult.tiledValues,
                        tileOffsetsList, tileSizesList, loops);
+  for (auto tiledOp : tilingResult.tiledOps) {
     if (auto dstOp = dyn_cast<DestinationStyleOpInterface>(tiledOp)) {
       auto innerMostLoop = loops.back();
       SmallVector<Value> tiledOpDestinationTensors = dstOp.getDpsInitOperands();
       updateDestinationOperandsForTiledOp(rewriter, tiledOpDestinationTensors,
                                           innerMostLoop.getRegionIterArgs());
     }
+  }
   return replacements;
 }
 
 /// Implementation of tiling transformation of `op` that implements the
 /// `TilingInterface` using `scf.for` to iterate over the tiles.
 FailureOr<scf::SCFTilingResult>
 mlir::scf::tileUsingSCFForOp(RewriterBase &rewriter, TilingInterface op,
                              const scf::SCFTilingOptions &options) {
if (!tilingResult.loops.empty()) {
       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 =
+  FailureOr<TilingResult> tiledImplementation =
       op.getTiledImplementation(rewriter, offsets, sizes);
-  tilingResult.tiledOps.append(tiledImplementation);
+  tilingResult.tiledOps.append(tiledImplementation->tiledOps);
   if (op->getNumResults() == 0) {
     // nothing more to do.
     return tilingResult;
   }
 
   // If loops are empty, the tiled op is used as the replacement for the untiled
   // op.
   if (tilingResult.loops.empty()) {
-    tilingResult.replacements = llvm::to_vector(
-        llvm::map_range(tiledImplementation[0]->getResults(),
-                        [](OpResult result) -> Value { return result; }));
+    tilingResult.replacements = tiledImplementation->tiledValues;
     return tilingResult;
   }
 
   // 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.
   int64_t numResults = op->getNumResults();
   SmallVector<SmallVector<OpFoldResult>> resultOffsetsList(numResults),
mlir::scf::tileUsingSCFForOp(RewriterBase &rewriter, TilingInterface op,
   }
 
   SmallVector<Value> destinationTensors;
   if (failed(tensor::getOrCreateDestinations(rewriter, op.getLoc(), op,
                                              destinationTensors)))
     return rewriter.notifyMatchFailure(op, "failed to get destinations");
 
   tilingResult.replacements = yieldTiledValues(
-      rewriter, destinationTensors, tilingResult.tiledOps.back(),
+      rewriter, destinationTensors, tiledImplementation.value(),
       resultOffsetsList, resultSizesList, tilingResult.loops);
 
   LLVM_DEBUG({
     if (!tilingResult.loops.empty()) {
       llvm::dbgs() << "After tiled implementation :\n";
       tilingResult.loops.front().dump();
       llvm::dbgs() << "\n";
     }
auto [fusableProducer, destinationIterArg] =
       getUntiledProducerFromSliceSource(&candidateSliceOp->getOpOperand(0),
                                         loops);
   if (!fusableProducer)
     return std::nullopt;
 
   // 2. Generate the tiled implementation of the producer of the source
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(candidateSliceOp);
-  FailureOr<Value> fusedProducerValue =
+  FailureOr<TilingResult> tileAndFuseResult =
       tensor::replaceExtractSliceWithTiledProducer(rewriter, candidateSliceOp,
                                                    fusableProducer);
-  if (failed(fusedProducerValue))
+  if (failed(tileAndFuseResult))
     return std::nullopt;
-  rewriter.replaceAllUsesWith(candidateSliceOp, fusedProducerValue.value());
+  rewriter.replaceAllUsesWith(candidateSliceOp,
+                              tileAndFuseResult->tiledValues[0]);
 
   // 3. 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) {
   //   %3 = scf.for .. iter_args(%arg1 = %arg0) {
mlir::scf::tileAndFuseProducerOfSlice(RewriterBase &rewriter,
   }
   if (iterArgNumber) {
     int64_t resultNumber = fusableProducer.getResultNumber();
     if (auto dstOp =
             dyn_cast<DestinationStyleOpInterface>(fusableProducer.getOwner())) {
       outerMostLoop.setIterArg(iterArgNumber.value(),
                                dstOp.getTiedOpOperand(fusableProducer)->get());
     }
-    if (auto dstOp = fusedProducerValue.value()
-                         .getDefiningOp<DestinationStyleOpInterface>()) {
+    for (auto tileAndFusedOp : tileAndFuseResult->tiledOps) {
+      auto dstOp = dyn_cast<DestinationStyleOpInterface>(tileAndFusedOp);
+      if (!dstOp)
+        continue;
       scf::ForOp innerMostLoop = loops.back();
       updateDestinationOperandsForTiledOp(
           rewriter, dstOp.getDpsInitOperand(resultNumber)->get(),
           innerMostLoop.getRegionIterArgs()[iterArgNumber.value()]);
     }
   }
   return scf::SCFFuseProducerOfSliceResult{fusableProducer,
-                                           fusedProducerValue.value()};
+                                           tileAndFuseResult->tiledValues[0]};
 }
 
 /// Reconstruct the fused producer from within the tiled-and-fused code.
 void mlir::scf::yieldReplacementForFusedProducer(
     RewriterBase &rewriter, tensor::ExtractSliceOp sliceOp,
     scf::SCFFuseProducerOfSliceResult fusedProducerInfo,
     MutableArrayRef<scf::ForOp> loops) {
   auto [fusableProducer, fusedProducerValue] = fusedProducerInfo;

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
     // Initialize all the ranges to {zero, one, one}. All the `ub`s are
     // overwritten.
     SmallVector<Range> loopRanges(reifiedShapes[0].size(), {zero, one, one});
     for (const auto &ub : enumerate(reifiedShapes[0]))
       loopRanges[ub.index()].size = ub.value();
     return loopRanges;
   }
 
-  SmallVector<Operation *>
+  FailureOr<TilingResult>
   getTiledImplementation(Operation *op, OpBuilder &b,
                          ArrayRef<OpFoldResult> offsets,
                          ArrayRef<OpFoldResult> sizes) const {
-    Operation *result =
+    FailureOr<TilingResult> result =
         tensor::bubbleUpPadSlice(b, cast<PadOp>(op), offsets, sizes);
-    if (!result)
-      return {};
-    return {result};
+    if (failed(result))
+      return failure();
+    return result.value();
   }
 
   LogicalResult
   getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
                         ArrayRef<OpFoldResult> offsets,
                         ArrayRef<OpFoldResult> sizes,
                         SmallVector<OpFoldResult> &resultOffsets,
                         SmallVector<OpFoldResult> &resultSizes) const {
     resultOffsets.assign(offsets.begin(), offsets.end());
SmallVector<utils::IteratorType> iteratorTypes(
         packOp.getSourceRank(), utils::IteratorType::parallel);
     return iteratorTypes;
   }
 
   SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
     return getPackUnPackIterationDomain<PackOp>(cast<PackOp>(op), b);
   }
 
-  SmallVector<Operation *>
+  FailureOr<TilingResult>
   getTiledImplementation(Operation *op, OpBuilder &b,
                          ArrayRef<OpFoldResult> offsets,
                          ArrayRef<OpFoldResult> sizes) const {
     auto packOp = cast<PackOp>(op);
     Location loc = packOp.getLoc();
 
     // The tiling is applied on interchanged dimensions. We have to undo the
     // interchange to map sizes and offsets to the original input.
getTiledImplementation(Operation *op, OpBuilder &b,
     if (auto val = packOp.getPaddingValue())
       tiledOperands.push_back(val);
     for (auto tile : packOp.getInnerTiles())
       tiledOperands.push_back(tile);
 
     Operation *tiledPackOp = b.create<PackOp>(
         loc, TypeRange{extractSlice.getType()}, tiledOperands, op->getAttrs());
 
-    return {tiledPackOp};
+    return TilingResult{{tiledPackOp},
+                        SmallVector<Value>(tiledPackOp->getResults())};
   }
 
   LogicalResult
   getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
                         ArrayRef<OpFoldResult> offsets,
                         ArrayRef<OpFoldResult> sizes,
                         SmallVector<OpFoldResult> &resultOffsets,
                         SmallVector<OpFoldResult> &resultSizes) const {
struct UnPackOpTiling
   /// Take Nn_to_N as an example. Say that N=32, n=8, and tiling_size=15. The
   /// coordinates of second tile (i.e., result[15..31]) are
   /// [(1, 7), (2, 0,), (2, 1) ... (3, 6), (3, 7)]. The first row and the last
   /// row are incomplete tiles. To represent the unpack op, we have to complete
   /// the rows. I.e., the input coordinates would start with (1, 0); end with
   /// (3, 7). In this context, the tiled unpack produces a (3 * n) elements
   /// because there are 3 rows in total. Follow by a tensor.extract_slice op, we
   /// can get the actual result.
-  SmallVector<Operation *>
+  FailureOr<TilingResult>
   getTiledImplementation(Operation *op, OpBuilder &b,
                          ArrayRef<OpFoldResult> offsets,
                          ArrayRef<OpFoldResult> sizes) const {
     auto unpackOp = cast<UnPackOp>(op);
     int64_t srcRank = unpackOp.getSourceRank();
     int64_t destRank = unpackOp.getDestRank();
     int64_t numInnerTiles = srcRank - destRank;
     Location loc = unpackOp.getLoc();
getTiledImplementation(Operation *op, OpBuilder &b,
     SmallVector<Value> tiledOperands = {sliceSource, sliceDest};
     for (auto tile : unpackOp.getInnerTiles())
       tiledOperands.push_back(tile);
 
     Operation *tiledUnpackOp = b.create<UnPackOp>(
         loc, TypeRange{sliceDest.getType()}, tiledOperands, op->getAttrs());
 
     if (isPerfectTilingCase)
-      return {tiledUnpackOp};
+      return TilingResult{{tiledUnpackOp},
+                          SmallVector<Value>(tiledUnpackOp->getResults())};
 
-    Operation *extractSlice =
+    auto extractSlice =
         b.create<ExtractSliceOp>(loc, tiledUnpackOp->getResult(0),
                                  resultOffsetsFromDest, sizes, destStrides);
-    return {tiledUnpackOp, extractSlice};
+    return TilingResult{{tiledUnpackOp}, {extractSlice.getResult()}};
   }
 
   LogicalResult
   getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
                         ArrayRef<OpFoldResult> offsets,
                         ArrayRef<OpFoldResult> sizes,
                         SmallVector<OpFoldResult> &resultOffsets,
                         SmallVector<OpFoldResult> &resultSizes) const {
     resultOffsets = llvm::to_vector(offsets);
     resultSizes = llvm::to_vector(sizes);
     return success();
   }
 
-  FailureOr<Value> generateResultTileValue(Operation *op, OpBuilder &b,
-                                           unsigned resultNumber,
-                                           ArrayRef<OpFoldResult> offsets,
-                                           ArrayRef<OpFoldResult> sizes) const {
-    return getTiledImplementation(op, b, offsets, sizes)
-        .back()
-        ->getResult(resultNumber);
+  FailureOr<TilingResult>
+  generateResultTileValue(Operation *op, OpBuilder &b, unsigned resultNumber,
+    FailureOr<TilingResult> tilingResult =
+        getTiledImplementation(op, b, offsets, sizes);
+    if (failed(tilingResult))
+      return failure();
+    return tilingResult.value();
   }
 };
 
 } // namespace
 
-Operation *tensor::bubbleUpPadSlice(OpBuilder &b, tensor::PadOp padOp,
+FailureOr<TilingResult> tensor::bubbleUpPadSlice(OpBuilder &b,
+                                                 tensor::PadOp padOp,
                                                  ArrayRef<OpFoldResult> offsets,
                                                  ArrayRef<OpFoldResult> sizes,
                                                  bool generateZeroSliceGuard) {
   // Only constant padding value supported.
   Value padValue = padOp.getConstantPaddingValue();
   if (!padValue)
-    return nullptr;
+    return failure();
 
   // Helper variables and functions for various arithmetic operations. These
   // are used extensively for computing new offset/length and padding values.
   Location loc = padOp->getLoc();
   AffineExpr dim0, dim1;
   bindDims(b.getContext(), dim0, dim1);
   // Add two integers.
   auto addMap = AffineMap::get(2, 0, {dim0 + dim1});
FailureOr<TilingResult> tensor::bubbleUpPadSlice(OpBuilder &b,
   // The shape of the result can be obtained from the sizes passed in.
   SmallVector<Value> dynDims;
   SmallVector<int64_t> shape;
   dispatchIndexOpFoldResults(sizes, dynDims, shape);
   RankedTensorType resultType =
       RankedTensorType::get(shape, padOp.getResultType().getElementType());
 
   // Insert cast to ensure that types match. (May be folded away.)
-  auto castResult = [&](Value val) -> Operation * {
+  auto castResult = [&](Value val) -> Value {
+    if (resultType == val.getType())
+      return val;
     return b.create<tensor::CastOp>(loc, resultType, val);
   };
 
   // In cases where the original data source is unused: Emit a GenerateOp and
   // do not generate a SliceOp. (The result shape of the SliceOp would
   // have a dimension of size 0, the semantics of which is unclear.)
   auto createGenerateOp = [&]() {
     // Create GenerateOp.
     auto generateOp = b.create<tensor::GenerateOp>(
         loc, resultType, dynDims,
         [&](OpBuilder &builder, Location gLoc, ValueRange indices) {
           builder.create<tensor::YieldOp>(gLoc, padValue);
         });
-    return castResult(generateOp);
+    return generateOp;
   };
 
   // Emit a SliceOp and a PadOp. Should not be used in cases where
   // the result shape of the new SliceOp has a zero dimension.
   auto createPadOfExtractSlice = [&]() {
     // Create pad(extract_slice(x)).
     auto newSliceOp = b.create<tensor::ExtractSliceOp>(
         loc, padOp.getSource(), newOffsets, newLengths, newStrides);
     auto newPadOp = b.create<PadOp>(loc, newSliceOp, staticNewLows,
                                     staticNewHighs, newLows, newHighs);
 
     // Copy region to new PadOp.
     IRMapping bvm;
     padOp.getRegion().cloneInto(&newPadOp.getRegion(), bvm);
 
     // Cast result and return.
-    return castResult(newPadOp);
+    return newPadOp;
   };
 
   // Rewrite extract_slice(pad(x)) into a GenerateOp it is statically known that
   // the original data source x is not used.
-  if (hasZeroLen)
-    return createGenerateOp();
+  if (hasZeroLen) {
+    Operation *generateOp = createGenerateOp();
+    return TilingResult{{generateOp}, {castResult(generateOp->getResult(0))}};
+  }
 
   // If there are dynamic dimensions: Generate an scf.if check to avoid
   // creating SliceOps with result dimensions of size 0 at runtime.
   if (generateZeroSliceGuard && dynHasZeroLenCond) {
+    Operation *thenOp;
+    Operation *elseOp;
     auto result = b.create<scf::IfOp>(
         loc, dynHasZeroLenCond,
         /*thenBuilder=*/
         [&](OpBuilder &b, Location loc) {
-          b.create<scf::YieldOp>(loc, createGenerateOp()->getResult(0));
+          thenOp = createGenerateOp();
+          b.create<scf::YieldOp>(loc, castResult(thenOp->getResult(0)));
         },
         /*elseBuilder=*/
         [&](OpBuilder &b, Location loc) {
-          b.create<scf::YieldOp>(loc, createPadOfExtractSlice()->getResult(0));
+          elseOp = createPadOfExtractSlice();
+          b.create<scf::YieldOp>(loc, castResult(elseOp->getResult(0)));
         });
-    return result;
+    return TilingResult{{result}, SmallVector<Value>(result->getResults())};
   }
-  return createPadOfExtractSlice();
+
+  Operation *newPadOp = createPadOfExtractSlice();
+  return TilingResult{{newPadOp}, {castResult(newPadOp->getResult(0))}};
 }
 
 void mlir::tensor::registerTilingInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, TensorDialect *dialect) {
     tensor::PadOp::attachInterface<PadOpTiling>(*ctx);
     tensor::PackOp::attachInterface<PackOpTiling>(*ctx);
     tensor::UnPackOp::attachInterface<UnPackOpTiling>(*ctx);

mlir/lib/Dialect/Tensor/Transforms/SwapExtractSliceWithProducerPatterns.cpp

 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Interfaces/TilingInterface.h"
 
 using namespace mlir;
 
-FailureOr<Value> tensor::replaceExtractSliceWithTiledProducer(
+FailureOr<TilingResult> tensor::replaceExtractSliceWithTiledProducer(
     OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producer) {
   auto producerOp = dyn_cast<TilingInterface>(producer.getOwner());
   if (!producerOp)
     return failure();
 
   // `TilingInterface` currently only supports strides being 1.
   if (llvm::any_of(sliceOp.getMixedStrides(), [](OpFoldResult ofr) {
         return !isConstantIntValue(ofr, 1);
       }))
     return failure();
 
-  FailureOr<Value> tiledResult = producerOp.generateResultTileValue(
+  FailureOr<TilingResult> tiledResult = producerOp.generateResultTileValue(
       builder, producer.getResultNumber(), sliceOp.getMixedOffsets(),
       sliceOp.getMixedSizes());
   if (failed(tiledResult))
     return failure();
 
   return *tiledResult;
 }