[mlir][TilingInterface] Fix `iter_args` handling in tile (and fuse).

[mravishankar]

Summary

The current approach for handling iter_args was to replace all uses
of the value that is used as init value with the corresponding
region block argument within the scf.for. This is not always
correct. Instead a more deliberate approach needs to be taken to
handle these. If the slice being fused represents a slice of the
destination operand of the untiled op, then

• Make the destination of the fused producer the init value of the loop nest
• For the tiled and fused producer op created, replace the slice of the destination operand with a slice of the corresponding region iter arg of the innermost loop of the generated loop nest

[mravishankar]

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

191

@nicolasvasilache responding to some of the comments from https://reviews.llvm.org/D134144 here.

1. This is just moved here from below.
2. I forgot that all the results of the tiled operation need to be yielded the same time for things to work correctly. The reason being that to get the tile offset and size of the result computed you need to use the offset and size of the iteration space. If you just yield one result of the tiled operation at a time, after the first yield, the offset of the iteration space (which depends on the induction variable of the loop) becomes invalid, but is still referenced in the offset computation. So you need to first compute the resultTileOffsets and resultTileSizes` of all the results before you yield them. Note that this is only an issue for returning the tiled results for the consumer (or the first tiled op). After that, if you fuse a producer and need to yield the value of the tiled producer, you have the resultTileOffset and resultTileSizes from the slice op that is used. Since after the yield the IR is in a consistent state, processing other slices doesnt have an issue. So this "batch processing of yields" only needs to happen for the tiled operation.. In any case, this method is just "move" of code from below. This is what was always done (I just had forgotten this).
3. You left a comment about front-loading the use of SubsetExtractOpInterface. I think that is a larger change that has to happen as a generalization of the whole TilingInterface. That probably needs to be done later. Also note that for just the tiling, you dont have slices. So we need a different container for the offsets and sizes that does not need the op itself. So there are things to be handled here w.r.t to SubsetExtractOpInterface` that are sort of unrelated to this change. This is really fixing a bug.

[nicolasvasilache]

Much better commit message with an IR example that points the error and the solution.
This is all too abstract and in your head right now, this is not a "simple fix" by any stretch of the definition.

Please address the comments first, name things better, add significant doc and remove all unnecessary complexity / code motion that you still can.

Once done I'll spend another hour trying to review this.

Feel free to add another reviewer that may have a better ability to grok where this is going.

mlir/include/mlir/Dialect/SCF/Utils/Utils.h

69

incomplete sentence.

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

191

1. this move creates significant reviewcomplexity, can it be undone?

Is the objective to just add /*replaceIterOperandsUsesInLoop =*/false ?
If so, it will be much easier to just do without moving the code and reducing the review complexity.

Then you can refactor as part of the NFC.

I am unable to tell whether

for (auto tiledResult : llvm::enumerate(tilingResult.tiledOp->getResults())) {
  if (failed(op.getResultTilePosition(
          rewriter, tiledResult.index(), offsets, sizes,
          resultTileOffsetsList[tiledResult.index()],
          resultTileSizesList[tiledResult.index()]))) {
    return rewriter.notifyMatchFailure(
        op, "unable to get insertion position of tiled result");
  }
}

+
changes in the function is NFC compared to:

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 {};
      }

If the answer is yes, please make it easier to accept your PR by avoiding complex refactorings mixed with non-trivial functional changes: leave the lambda where it was and just add /*replaceIterOperandsUsesInLoop =*/false.

If the answer is no, we need significantly more comments here to understand why not.

390

Treating all the code above this line as unnecessary change and assuming that adding`/*replaceIterOperandsUsesInLoop =*/false.` is the only useful change.

415

this functions needs a significantly better name that is descriptive of what it does.
walkBackXAndReturnY or something

478

you can now auto [x, y] = ...

498–499

At this point, if the slice is for a destination operand, the IR is ill-formed because XXX (with a clear example with multiple loops and clear pointing out where in the IR the problem is and why).

Problematic IR Example

To correct this, we need to update Y and Z as follows:
  - iter_arg ...
  - destination of the slice ...

Corrected IR example.

mlir/lib/Dialect/SCF/Utils/Utils.cpp

131

Comments and examples needed here, I have no idea why this post-hoc patchup is happening here and why.
Can this be achieved in a more idiomatic fashion, without post-hoc patchup, by e.g. splitting replaceLoopWithNewYields in multiple parts?

[mravishankar]

Remove code movements that are NFC to reduce diff.

[mravishankar]

Address comments.

[mravishankar]

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

191

Inlined the function. The comment doesnt point to where the code is now, but it is just that change.

415

Went with getUntiledProducerFromSliceSource.

478

TIL....

mlir/lib/Dialect/SCF/Utils/Utils.cpp

131

I think this can be simplified. Instead of going bottom up from inner to outer, we can try making this go top down, from outer to inner... I can do this after the fact. Maybe not add to the patch right now?

[nicolasvasilache]

Thanks Mahesh, much nicer an easier to follow!
Approving conditioned on the top-down/bottom-up/middle-out cleanup arriving very soon.

mlir/include/mlir/Dialect/SCF/Utils/Utils.h

48

nit: double are replaced

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

170

nit: is in

191

Great, thanks for removing that cognitive barrier!

198

Plz add some TODO about SubsetExtractOpInterface

415

Thanks!

520

Could you please highlight

%arg0 = %1 /*incorrect value*/)

and

tensor.extract_slice %0 /* incorrect value */

?

520

Could you please highlight

%arg0 = %0 /*now this is correct*/)
and

tensor.extract_slice %arg1 /* now this is correct */
?

547

can we use (or add) some helper function to avoid magic indexings in clients of this op (there should already be something available).

mlir/lib/Dialect/SCF/Utils/Utils.cpp

131

ok, but let's please do it as part of this list of patches.
Before or after the NFC is fine, but please don't let this linger, approving conditioned on that cleanup arriving very soon.

[mravishankar]

Rebase and address comments.

[mravishankar]

Rebase and address comments.

mlir/include/mlir/Dialect/SCF/Utils/Utils.h

 /// values returned by the callback should match the number of new
 /// initialization values. This function
 /// - Moves (i.e. doesnt clone) operations from the `loop` to the newly created
 ///   loop
 /// - Replaces the uses of `loop` with the new loop.
 /// - `loop` isnt erased, but is left in a "no-op" state where the body of the
 ///   loop just yields the basic block arguments that correspond to the
 ///   initialization values of a loop. The loop is dead after this method.
-/// - All uses of the `newIterOperands` within the generated new loop
-///   are replaced with the corresponding `BlockArgument` in the loop body.
+/// - If `replaceIterOperandsUsesInLoop` is true, all uses of the
+///   `newIterOperands` within the generated new loop
 using NewYieldValueFn = std::function<SmallVector<Value>(
     OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs)>;
 scf::ForOp replaceLoopWithNewYields(OpBuilder &builder, scf::ForOp loop,
                                     ValueRange newIterOperands,
-                                    const NewYieldValueFn &newYieldValuesFn);
+                                    const NewYieldValueFn &newYieldValuesFn,
+                                    bool replaceIterOperandsUsesInLoop = true);
 
 /// Update a perfectly nested loop nest to yield new values from the innermost
 /// loop and propagating it up through the loop nest. This function
 /// - Expects `loopNest` to be a perfectly nested loop with outer most loop
 ///   first and innermost loop last.
 /// - `newIterOperands` are the initialization values to be used for the
 ///    outermost loop
 /// - `newYielValueFn` is the callback that generates the new values to be
 ///   yielded from within the innermost loop.
 /// - The original loops are not erased,  but are left in a "no-op" state where
 ///   the body of the loop just yields the basic block arguments that correspond
 ///   to the initialization values of a loop. The original loops are dead after
 ///   this method.
-/// - All uses of the `newIterOperands` within the generated new loop
-///   are replaced with the corresponding `BlockArgument` in the loop body.
+/// - If `replaceIterOperandsUsesInLoop` is true, all uses of the
+///   `newIterOperands` within the generated new loop
 SmallVector<scf::ForOp>
 replaceLoopNestWithNewYields(OpBuilder &builder, ArrayRef<scf::ForOp> loopNest,
                              ValueRange newIterOperands,
-                             const NewYieldValueFn &newYieldValueFn);
+                             const NewYieldValueFn &newYieldValueFn,
+                             bool replaceIterOperandsUsesInLoop = true);
 
 /// Outline a region with a single block into a new FuncOp.
 /// Assumes the FuncOp result types is the type of the yielded operands of the
 /// single block. This constraint makes it easy to determine the result.
 /// This method also clones the `arith::ConstantIndexOp` at the start of
 /// `outlinedFuncBody` to alloc simple canonicalizations.
 /// Creates a new FuncOp and thus cannot be used in a FuncOp pass.
 /// The client is responsible for providing a unique `funcName` that will not

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

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 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
+/// }
+/// ```
+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::TileUsingSCFForOp::returningMatchAndRewrite(
 
   if (op->getNumResults() == 0) {
     rewriter.eraseOp(op);
     return tilingResult;
   }
 
   // 5. If the original operations has results, modify the loop nest to yield
   // the replacement values.
-  SmallVector<Value> replacements;
   if (tilingResult.loops.empty()) {
     // 5a. If there were no loops, the tiled implementation results are the
     // replacements.
     rewriter.replaceOp(op, tilingResult.tiledOp->getResults());
     return tilingResult;
   }
 
-  // 5b. `scf.for` with tensor semantics requires the loop nest to yield the
+  // 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()) {
+    return rewriter.notifyMatchFailure(
+        tilingResult.tiledOp,
+        "expected tiled op to have as many results as the untiled operation");
+  }
+
+  // `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
   // }
for (auto resultNum : llvm::seq<unsigned>(0, op->getNumResults())) {
           newBBArgs[resultNum], resultTileOffsets, resultTileSizes,
           resultTileStrides);
       yieldedValues.push_back(yieldedValue);
     }
     return yieldedValues;
   };
   SmallVector<scf::ForOp> newLoops = replaceLoopNestWithNewYields(
       rewriter, tilingResult.loops, op.getDestinationOperands(rewriter),
-      yieldValueFn);
+      yieldValueFn, /*replaceIterOperandsUsesInLoops =*/false);
   for (const auto &loop : llvm::enumerate(tilingResult.loops)) {
     rewriter.eraseOp(loop.value());
     tilingResult.loops[loop.index()] = newLoops[loop.index()];
   }
   rewriter.replaceOp(op, tilingResult.loops.front().getResults());
   return tilingResult;
 }
 
 //===----------------------------------------------------------------------===//
 // TileConsumerAndFuseProducersUsingSCFForOp pattern 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 `Value` that is defined by an operation that implements
-/// the `TilingInterface`. Looks through `iter_args` of scf.for nest
-/// if required.
-static Optional<OpResult> getFusableProducer(Value v) {
-  while (auto blockArg = v.dyn_cast<BlockArgument>()) {
-    auto loopOp = dyn_cast<scf::ForOp>(blockArg.getOwner()->getParentOp());
-    if (!loopOp)
-      return llvm::None;
-    v = loopOp.getOpOperandForRegionIterArg(blockArg).get();
-  }
-  if (!isa_and_nonnull<TilingInterface>(v.getDefiningOp()))
-    return llvm::None;
-  return v.cast<OpResult>();
-}
-
-// Replace iter args of the outer most loop with region args of the inner most
-// one.
-static void replaceIterArgs(scf::ForOp outerFor, scf::ForOp innerFor,
-                            PatternRewriter &rewriter) {
-  assert(outerFor.getNumIterOperands() == innerFor.getNumIterOperands() &&
-         "expect same number of iter args");
-  Block *block = &(*innerFor.getRegion().begin());
-  for (auto it :
-       llvm::zip(outerFor.getIterOperands(), innerFor.getRegionIterArgs())) {
-    Value source = std::get<0>(it);
-    Value target = std::get<1>(it);
-    source.replaceUsesWithIf(target, [&](OpOperand &use) {
-      return use.getOwner()->getBlock() == block;
-    });
-  }
+/// 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 *>>
+getProducerOp(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};
 }
 
 FailureOr<scf::SCFTileAndFuseResult>
 scf::TileConsumerAndFuseProducersUsingSCFForOp::returningMatchAndRewrite(
     TilingInterface op, PatternRewriter &rewriter) const {
   // 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()) {
scf::TileConsumerAndFuseProducersUsingSCFForOp::returningMatchAndRewrite(
   OpBuilder::InsertionGuard g(rewriter);
   while (!candidates.empty()) {
     // 2a. Traverse the slices in BFS fashion.
     tensor::ExtractSliceOp candidateSliceOp = candidates.front();
     candidates.pop_front();
 
     // 2b. Get the producer of the source (potentially walking through
     // `iter_args` of nested `scf.for`)
-    Optional<OpResult> fusableProducer =
-        getFusableProducer(candidateSliceOp.getSource());
+    OpResult fusableProducer;
+    Optional<OpOperand *> destinationIterArg;
+    std::tie(fusableProducer, destinationIterArg) = getProducerOp(
+        &candidateSliceOp->getOpOperand(0), tileAndFuseResult.loops);
     if (!fusableProducer)
       continue;
 
     // 2c. Generate the tiled implementation of the producer of the source
     rewriter.setInsertionPoint(candidateSliceOp);
     FailureOr<Value> fusedProducerValue =
         tensor::replaceExtractSliceWithTiledProducer(rewriter, candidateSliceOp,
-                                                     fusableProducer.value());
+                                                     fusableProducer);
     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);
     addCandidateSlices(fusedProducer, candidates);
 
-    // 2e. If the operation being fused creates a value that is used as `outs`
-    //     in the tiled operation, the result of the unfused operation will be
-    //     used in the `iter_args` of the tiled loop generated. When the
-    //     operation is fused, this use in `iter_args` needs to be modified to
-    //     use the destination of the fused operation. For example, starting
-    //     with
+    // 2e. If the slice is for a destination operand, then
+    //     - Update the iter_arg of the outer most loop to use the destination
+    //       of the untiled producer.
+    //     - Update the destination of the slice of the tiled producer generated
+    //       to use the same basic block argument as the slice that was used to
+    //       producer the tiled implementation.
+    // For example
     //
-    //     ```mlir
-    //     %0 = linalg.init_tensor ...
-    //     %1 = linalg.fill ... outs(%0:...)...
-    //     %2 = linalg.matmul ... outs(%1:...)....
-    //     ```
-    //
-    //     First the `linalg.matmul` gets tiled
-    //
-    //     ```mlir
-    //     %0 = linalg.init_tensor
-    //     %1 = linalg.fill
-    //     %2 = scf.for .... iter_args(%arg0 = %1)...
-    //        ...
-    //        ... = linalg.matmul ...
-    //
-    //     ```
+    // ```mlir
+    // %0 = linalg.init
+    // %1 = linalg.fill .. outs(%0 : )
+    // %2 = scf.for .. iter_args(%arg0 = %1) {
+    //   %3 = tensor.extract_slice %arg0
+    //   .. = linalg.matmul .. outs(%3 : )
+    // }
+    // ```
     //
-    //     When the `linalg.fill` gets fused, the `iter_args` needs to be
-    //     modified
+    // is transformed to
     //
-    //     ```mlir
-    //     %0 = linalg.init_tensor
-    //     %1 = scf.for ... iter_args(%arg0 = %0)...
-    //        ...
-    //        %2 = linalg.fill ...
-    //        %3 = linalg.matmul ... outs(%2: ...)...
-    //     ```
-    TilingInterface unfusedProducerOp =
-        cast<TilingInterface>(fusableProducer->getOwner());
-    scf::ForOp outerMostTiledLoop = tileAndFuseResult.loops.front();
-    SmallVector<Value> unfusedProducerOpDestValues =
-        unfusedProducerOp.getDestinationOperands(rewriter);
-    for (OpOperand &uses : unfusedProducerOp->getUses()) {
-      if (uses.getOwner() == outerMostTiledLoop.getOperation()) {
-        unsigned resultNumber = uses.get().cast<OpResult>().getResultNumber();
-        unsigned operandNumber = uses.getOperandNumber();
-        outerMostTiledLoop->setOperand(
-            operandNumber, unfusedProducerOpDestValues[resultNumber]);
+    // ```mlir
+    // %0 = linalg.init
+    // %1 = scf.for .. iter_args(%arg0 = %0) {
+    //   %3 = tensor.extract_slice %arg0
+    //   %4 = linalg.fill .. outs(%3 : )
+    //   .. = linalg.matmul .. outs(%4 : )
+    // }
+    // ```
+    if (destinationIterArg) {
+      unsigned iterArgNumber = destinationIterArg.value()->getOperandNumber();
+      unsigned resultNumber = fusableProducer.getResultNumber();
+      if (auto producerOp =
+              dyn_cast<TilingInterface>(fusableProducer.getOwner())) {
+        scf::ForOp outerMostLoop = tileAndFuseResult.loops.front();
+        SmallVector<Value> destination =
+            producerOp.getDestinationOperands(rewriter);
+        outerMostLoop.setOperand(iterArgNumber, destination[resultNumber]);
+      }
+      if (auto tiledAndFusedInterfaceOp =
+              fusedProducerValue.value().getDefiningOp<TilingInterface>()) {
+        scf::ForOp innerMostLoop = tileAndFuseResult.loops.back();
+        SmallVector<Value> destination =
+            tiledAndFusedInterfaceOp.getDestinationOperands(rewriter);
+        updateDestinationOperandsForTiledOp(
+            rewriter, destination[resultNumber],
+            innerMostLoop
+                .getRegionIterArgs()[iterArgNumber -
+                                     innerMostLoop.getNumControlOperands()]);
       }
     }
   }
-  replaceIterArgs(tileAndFuseResult.loops.front(),
-                  tileAndFuseResult.loops.back(), rewriter);
   return tileAndFuseResult;
 }
 
 //===----------------------------------------------------------------------===//
 // LowerToLoopsUsingSCFForOp
 //===----------------------------------------------------------------------===//
 
 FailureOr<SmallVector<scf::ForOp>>

mlir/lib/Dialect/SCF/Utils/Utils.cpp

struct LoopParams {
   Value upperBound;
   Value step;
 };
 } // namespace
 
 scf::ForOp
 mlir::replaceLoopWithNewYields(OpBuilder &builder, scf::ForOp loop,
                                ValueRange newIterOperands,
-                               const NewYieldValueFn &newYieldValuesFn) {
+                               const NewYieldValueFn &newYieldValuesFn,
+                               bool replaceIterOperandsUsesInLoop) {
   // Create a new loop before the existing one, with the extra operands.
   OpBuilder::InsertionGuard g(builder);
   builder.setInsertionPoint(loop);
   auto operands = llvm::to_vector(loop.getIterOperands());
   operands.append(newIterOperands.begin(), newIterOperands.end());
   scf::ForOp newLoop = builder.create<scf::ForOp>(
       loop.getLoc(), loop.getLowerBound(), loop.getUpperBound(), loop.getStep(),
       operands, [](OpBuilder &, Location, Value, ValueRange) {});
mlir::replaceLoopWithNewYields(OpBuilder &builder, scf::ForOp loop,
 
   // Remap the BlockArguments from the original loop to the new loop
   // BlockArguments.
   ArrayRef<BlockArgument> bbArgs = loopBody->getArguments();
   for (auto it :
        llvm::zip(bbArgs, newLoopBody->getArguments().take_front(bbArgs.size())))
     std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
 
+  if (replaceIterOperandsUsesInLoop) {
     // Replace all uses of `newIterOperands` with the corresponding basic block
     // arguments.
     for (auto it : llvm::zip(newIterOperands, newBBArgs)) {
       std::get<0>(it).replaceUsesWithIf(std::get<1>(it), [&](OpOperand &use) {
         Operation *user = use.getOwner();
         return newLoop->isProperAncestor(user);
       });
     }
+  }
 
   // Replace all uses of the original loop with corresponding values from the
   // new loop.
   loop.replaceAllUsesWith(
       newLoop.getResults().take_front(loop.getNumResults()));
 
   // Add a fake yield to the original loop body that just returns the
   // BlockArguments corresponding to the iter_args. This makes it a no-op loop.
   // The loop is dead. The caller is expected to erase it.
   builder.setInsertionPointToEnd(loopBody);
   builder.create<scf::YieldOp>(loop->getLoc(), loop.getRegionIterArgs());
 
   return newLoop;
 }
 
 SmallVector<scf::ForOp> mlir::replaceLoopNestWithNewYields(
     OpBuilder &builder, ArrayRef<scf::ForOp> loopNest,
-    ValueRange newIterOperands, const NewYieldValueFn &newYieldValueFn) {
+    ValueRange newIterOperands, const NewYieldValueFn &newYieldValueFn,
+    bool replaceIterOperandsUsesInLoop) {
   if (loopNest.empty())
     return {};
   SmallVector<scf::ForOp> newLoopNest(loopNest.size());
 
   newLoopNest.back() = replaceLoopWithNewYields(
       builder, loopNest.back(), newIterOperands, newYieldValueFn);
 
   for (unsigned loopDepth :
        llvm::reverse(llvm::seq<unsigned>(0, loopNest.size() - 1))) {
     NewYieldValueFn fn = [&](OpBuilder &innerBuilder, Location loc,
                              ArrayRef<BlockArgument> innerNewBBArgs) {
       SmallVector<Value> newYields(
           newLoopNest[loopDepth + 1]->getResults().take_back(
               newIterOperands.size()));
       return newYields;
     };
-    newLoopNest[loopDepth] = replaceLoopWithNewYields(
-        builder, loopNest[loopDepth], newIterOperands, fn);
+    newLoopNest[loopDepth] =
+        replaceLoopWithNewYields(builder, loopNest[loopDepth], newIterOperands,
+                                 fn, replaceIterOperandsUsesInLoop);
+    if (!replaceIterOperandsUsesInLoop) {
+      /// The yield is expected to producer the following structure
+      /// ```
+      /// %0 = scf.for ... iter_args(%arg0 = %init) {
+      ///   %1 = scf.for ... iter_args(%arg1 = %arg0) {
+      ///     scf.yield %yield
+      ///   }
+      /// }
+      /// ```
+      ///
+      /// since the yield is propagated from inside out, after the inner
+      /// loop is processed the IR is in this form
+      ///
+      /// ```
+      /// scf.for ... iter_args {
+      ///   %1 = scf.for ... iter_args(%arg1 = %init) {
+      ///     scf.yield %yield
+      ///   }
+      /// ```
+      ///
+      /// If `replaceIterOperandUsesInLoops` is true, there is nothing to do.
+      /// `%init` will be replaced with `%arg0` when it is created for the
+      /// outer loop. But without that this has to be done explicitly.
+      unsigned subLen = newIterOperands.size();
+      unsigned subStart =
+          newLoopNest[loopDepth + 1].getNumIterOperands() - subLen;
+      auto resetOperands =
+          newLoopNest[loopDepth + 1].getInitArgsMutable().slice(subStart,
+                                                                subLen);
+      resetOperands.assign(
+          newLoopNest[loopDepth].getRegionIterArgs().take_back(subLen));
+    }
   }
   return newLoopNest;
 }
 
 /// Outline a region with a single block into a new FuncOp.
 /// Assumes the FuncOp result types is the type of the yielded operands of the
 /// single block. This constraint makes it easy to determine the result.
 /// This method also clones the `arith::ConstantIndexOp` at the start of

mlir/test/Interfaces/TilingInterface/tile-and-fuse-using-interface.mlir

 //  CHECK-DAG:       %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]]
 //  CHECK-DAG:       %[[INIT_TILE:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV0]], %[[IV1]]]
 //      CHECK:       %[[FILL_TILE:.+]] = linalg.fill
 // CHECK-SAME:           outs(%[[INIT_TILE]] :
 //      CHECK:       %[[GEMM_TILE:.+]] = linalg.matmul
 // CHECK-SAME:           ins(%[[LHS_TILE]], %[[RHS_TILE]] :
 // CHECK-SAME:           outs(%[[FILL_TILE]] :
 //      CHECK:       %[[INSERT:.+]] = tensor.insert_slice %[[GEMM_TILE]] into %[[ITERARG1]][%[[IV0]], %[[IV1]]]
-//      CHECK        scf.yield %[[INSERT]]
+//      CHECK:       scf.yield %[[INSERT]]
 
 // -----
 
 func.func @gemm_generic_fusion(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
     %arg2 : tensor<?xf32>) -> tensor<?x?xf32> {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %cst = arith.constant 0.0 : f32

 // CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: tensor<?xf32>)
 //      CHECK:   %[[INIT:.+]] = linalg.init_tensor
 //      CHECK:   scf.for %[[IV0:[a-zA-Z0-9]+]] =
 // CHECK-SAME:       iter_args(%[[ITERARG0:.+]] = %[[INIT]])
 //      CHECK:     scf.for %[[IV1:[a-zA-Z0-9]+]] =
 // CHECK-SAME:         iter_args(%[[ITERARG1:.+]] = %[[ITERARG0]])
 //  CHECK-DAG:       %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0]
 //  CHECK-DAG:       %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]]
-//  CHECK-DAG:       %[[INIT_TILE:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV0]], %[[IV1]]]
+//  CHECK-DAG:       %[[INIT_TILE:.+]] = tensor.extract_slice %[[INIT]][%[[IV0]], %[[IV1]]]
 //      CHECK:       %[[FILL_TILE:.+]] = linalg.fill
 // CHECK-SAME:           outs(%[[INIT_TILE]] :
 //      CHECK:       %[[GEMM_TILE:.+]] = linalg.matmul
 // CHECK-SAME:           ins(%[[LHS_TILE]], %[[RHS_TILE]] :
 // CHECK-SAME:           outs(%[[FILL_TILE]] :
 //  CHECK-DAG:       %[[BIAS_TILE:.+]] = tensor.extract_slice %[[ARG2]][%[[IV1]]]
 //  CHECK-DAG:       %[[OUTS_TILE:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV0]], %[[IV1]]]
 //      CHECK:       %[[GENERIC_TILE:.+]] = linalg.generic
 // CHECK-SAME:           ins(%[[GEMM_TILE]], %[[BIAS_TILE]] :
 // CHECK-SAME:           outs(%[[OUTS_TILE]] :
 //      CHECK:       %[[INSERT:.+]] = tensor.insert_slice %[[GENERIC_TILE]] into %[[ITERARG1]][%[[IV0]], %[[IV1]]]
-//      CHECK        scf.yield %[[INSERT]]
+//      CHECK:       scf.yield %[[INSERT]]
 
 // -----
 
 func.func @gemm_gemm_fusion(%lhs0 : tensor<?x?xf32>, %rhs0 : tensor<?x?xf32>, %rhs1 : tensor<?x?xf32>) -> tensor<?x?xf32> {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %cst = arith.constant 0.0 : f32
   %d0 = tensor.dim %lhs0, %c0 : tensor<?x?xf32>

 //  CHECK-DAG:     %[[RHS1_TILE:.+]] = tensor.extract_slice %[[RHS1]][0, 0]
 //  CHECK-DAG:     %[[INIT1_TILE:.+]] = tensor.extract_slice %[[ITERARG]][%[[IV]], 0]
 //      CHECK:     %[[FILL1_TILE:.+]] = linalg.fill
 // CHECK-SAME:         outs(%[[INIT1_TILE]] :
 //      CHECK:     %[[GEMM1_TILE:.+]] = linalg.matmul
 // CHECK-SAME:         ins(%[[GEMM0_TILE]], %[[RHS1_TILE]] :
 // CHECK-SAME:         outs(%[[FILL1_TILE]] :
 //      CHECK:     %[[INSERT:.+]] = tensor.insert_slice %[[GEMM1_TILE]] into %[[ITERARG]][%[[IV]], 0]
-//      CHECK      scf.yield %[[INSERT]]
+//      CHECK:     scf.yield %[[INSERT]]
 
 // -----
 
 func.func @gemm_transpose_fusion(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32> {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %cst = arith.constant 0.0 : f32
   %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>

 //      CHECK:       %[[GEMM_TILE:.+]] = linalg.matmul
 // CHECK-SAME:           ins(%[[LHS_TILE]], %[[RHS_TILE]] :
 // CHECK-SAME:           outs(%[[FILL_TILE]] :
 //  CHECK-DAG:       %[[OUTS_TILE:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV1]], %[[IV0]]]
 //      CHECK:       %[[GENERIC_TILE:.+]] = linalg.generic
 // CHECK-SAME:           ins(%[[GEMM_TILE]] :
 // CHECK-SAME:           outs(%[[OUTS_TILE]] :
 //      CHECK:       %[[INSERT:.+]] = tensor.insert_slice %[[GENERIC_TILE]] into %[[ITERARG1]][%[[IV1]], %[[IV0]]]
-//      CHECK        scf.yield %[[INSERT]]
+//      CHECK:       scf.yield %[[INSERT]]
 
 // -----
 
 func.func @interchange_matmul_fusion(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32> {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
   %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>

 // CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>)
 //      CHECK:   %[[INIT:.+]] = linalg.init_tensor
 //      CHECK:   scf.for %[[IV0:[a-zA-Z0-9]+]] =
 // CHECK-SAME:       iter_args(%[[ITERARG0:.+]] = %[[INIT]])
 //      CHECK:     scf.for %[[IV1:[a-zA-Z0-9]+]] =
 // CHECK-SAME:         iter_args(%[[ITERARG1:.+]] = %[[ITERARG0]])
 //  CHECK-DAG:       %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV1]], 0]
 //  CHECK-DAG:       %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV0]]]
-//  CHECK-DAG:       %[[INIT_TILE:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV1]], %[[IV0]]]
+//  CHECK-DAG:       %[[INIT_TILE:.+]] = tensor.extract_slice %[[INIT]][%[[IV1]], %[[IV0]]]
 //      CHECK:       %[[FILL_TILE:.+]] = linalg.fill
 // CHECK-SAME:           outs(%[[INIT_TILE]] :
 //      CHECK:       %[[GEMM_TILE:.+]] = linalg.matmul
 // CHECK-SAME:           ins(%[[LHS_TILE]], %[[RHS_TILE]] :
 // CHECK-SAME:           outs(%[[FILL_TILE]] :
 //      CHECK:       %[[INIT_TILE_2:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV1]], %[[IV0]]]
 //      CHECK:       %[[GENERIC_TILE:.+]] = linalg.generic
 // CHECK-SAME:           ins(%[[GEMM_TILE]] :
 // CHECK-SAME:           outs(%[[INIT_TILE_2]] :
 //      CHECK:       %[[INSERT:.+]] = tensor.insert_slice %[[GENERIC_TILE]] into %[[ITERARG1]][%[[IV1]], %[[IV0]]]
-//      CHECK        scf.yield %[[INSERT]]
+//      CHECK:       scf.yield %[[INSERT]]
 
 // -----
 
 func.func @matmul_plus_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
                          %arg2: tensor<?x?xf32>) -> tensor<?x?xf32>{
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %0 = tensor.dim %arg2, %c0 : tensor<?x?xf32>