[mlir][Linalg] Fix hoist padding through scf.for iter_arg

[nicolasvasilache]

Summary

Previously, hoisting through an iter_arg would mistakenly yield the unpadded value and
cast it to the padded value.

This was incorrect and resulted in out-of-bounds accesses.
The correct formulation is to yield the padded value and extract a smaller dynamic slice
out of it.

[qcolombet]

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

756

%0 => %arg0

760

I'd add "return tensor::ExtractSliceOp when no rewrite happened".

763

What does packedTensor hold with respect to the rewrite in the comment?

813

We used to assert in this case.
Is it something we can reach now?
And if so could you add a test case for it?

835

Isn't that guaranteed by the fact that they both feed and use the same iter_arg?

[nicolasvasilache]

Address comments.

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

813

asserting was actually fine

835

Unfortunately no, so I went ahead and implemented a poor man's analysis.
In the future, this will be easier with a tensor.unpad op.

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

 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/Interfaces/DestinationStyleOpInterface.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/Support/Debug.h"
 
 using llvm::dbgs;
 
 #define DEBUG_TYPE "hoist-padding"
 
 #define DBGS() (dbgs() << '[' << DEBUG_TYPE << "] ")

 /// Other cases succeed and will trigger hoisting of the pad op.
 struct HoistPaddingAnalysis {
   HoistPaddingAnalysis(tensor::PadOp padOp, int numLoops);
   HoistPaddingAnalysis(tensor::PadOp padOp, scf::ForOp outermostEnclosingForOp);
 
   bool isValid() { return valid.has_value() && valid.value(); }
   bool isInvalid() { return valid.has_value() && !valid.value(); }
 
-  /// Footprint of the packedTensor, computed from the packingLoops.
-  SmallVector<Value> getPackedTensorSizes(RewriterBase &rewriter,
+  /// Footprint of the hoistedPackedTensor, computed from the packingLoops.
+  SmallVector<Value> getHoistedPackedTensorSizes(RewriterBase &rewriter,
                                                  Location loc) const;
 
   /// Performs optional hoisting to enable hoist padding to occur. This may be
   /// necessary when `sliceOp` is not defined outside of the outermost enclosing
   /// loop we want to hoist above.
   ///
   /// Example:
   /// ```
   /// %source = linalg.fill(%cst, %arg0)
for (Operation *op : llvm::reverse(backwardSlice)) {
     if (!isa<arith::ConstantOp>(op))
       operationsToRemove.insert(op);
   }
   backwardSlice.set_subtract(operationsToRemove);
   return success();
 }
 
 SmallVector<Value>
-HoistPaddingAnalysis::getPackedTensorSizes(RewriterBase &rewriter,
+HoistPaddingAnalysis::getHoistedPackedTensorSizes(RewriterBase &rewriter,
                                                   Location loc) const {
   SmallVector<Value> dynamicTensorSizes;
 
   // Upper bound the packing loop lengths to size the packed tensor. Taking
   // upper bounds can make the sizes of the packed tensor independent of the
   // enclosing loops. This independence is a prerequisite for reusing the same
   // buffer for all enclosing loop iterations and hoisting its allocation out
   // of the enclosing loops.
   for (auto forOp : packingLoops) {
static Value buildLoopIterationCount(RewriterBase &rewriter, scf::ForOp outer,
   auto loc = forOp->getLoc();
   return rewriter.createOrFold<AffineApplyOp>(
       loc, (iv - lb).ceilDiv(step), ValueRange{ivVal, lbVal, stepVal});
 }
 
 // Build a packing loop nest by iteratively traversing the backward slice and
 // clone the operations, iteratively stepping into the loops that we encounter.
 // The implementation proceeds in a stack-like fashion:
-//   1. Iteratively clone and step into the loops, pushing the `packedTensor`
+//   1. Iteratively clone and step into the loops, pushing the
+//   `hoistedPackedTensor`
 //      deeper in the stack.
 //   2. At the innermost loop level, create a GenericOp if `transposeVector` is
 //      non-empty.
 //   3. At the innermost loop level, create a InsertSliceOp.
 //   4. Iteratively pop and yield the result of the InsertSliceOp across the
 //      cloned loops.
 static PackingResult buildPackingLoopNestImpl(
     RewriterBase &rewriter, IRMapping &bvm, tensor::PadOp opToHoist,
     ArrayRef<int64_t> transposeVector, RankedTensorType transposedTensorType,
     tensor::EmptyOp emptyOp, const HoistPaddingAnalysis &analysis) {
   SmallVector<OpFoldResult> offsets, sizes, strides;
-  SmallVector<Value> clonedLoopIvs, leadingPackedTensorIndexings;
+  SmallVector<Value> clonedLoopIvs, leadingHoistedPackedTensorIndexings;
 
   scf::ForOp outerLoop = analysis.outermostEnclosingForOp;
 
   Location loc = opToHoist->getLoc();
   RankedTensorType paddedTensorType = opToHoist.getResultType();
   int paddedRank = paddedTensorType.getRank();
 
   // Step 0. Populate bvm with opToHoist.getSource if relevant.
   BlockArgument bbArg = opToHoist.getSource().dyn_cast<BlockArgument>();
   while (bbArg) {
     auto forOp = dyn_cast<scf::ForOp>(bbArg.getOwner()->getParentOp());
     if (!forOp)
       break;
     if (forOp != outerLoop && !outerLoop->isAncestor(forOp))
       break;
     OpOperand &operand = forOp.getOpOperandForRegionIterArg(bbArg);
     bvm.map(bbArg, operand.get());
     bbArg = operand.get().dyn_cast<BlockArgument>();
   }
 
-  // Step 1. iteratively clone loops and push `packedTensor`.
-  Value packedTensor = emptyOp.getResult();
+  // Step 1. iteratively clone loops and push `hoistedPackedTensor`.
+  Value hoistedPackedTensor = emptyOp.getResult();
   OpBuilder::InsertionGuard g(rewriter);
   for (Operation *op : analysis.backwardSlice) {
-    // Specifically sit out in the extract_slice(packedTensor) case: this is
-    // the piece we seek to replace.
+    // Specifically sit out in the extract_slice(hoistedPackedTensor) case: this
+    // is the piece we seek to replace.
     if (auto sliceOp = dyn_cast<tensor::ExtractSliceOp>(op)) {
-      if (bvm.lookupOrDefault(sliceOp.getSource()) == packedTensor) {
+      if (bvm.lookupOrDefault(sliceOp.getSource()) == hoistedPackedTensor) {
         LLVM_DEBUG(DBGS() << "--Skip: " << sliceOp << "\n");
         continue;
       }
     }
 
     // Clone all operations except loops which require special handling.
     auto forOp = dyn_cast<scf::ForOp>(op);
     if (!forOp) {
       // We are at the right insertion point within the loop nest.
       rewriter.clone(*op, bvm);
       continue;
     }
 
-    // Create a packing loop that takes `packedTensor` as iteration argument.
+    // Create a packing loop that takes `hoistedPackedTensor` as iteration
+    // argument.
     auto clonedForOp = rewriter.create<scf::ForOp>(
         loc, bvm.lookupOrDefault(forOp.getLowerBound()),
         bvm.lookupOrDefault(forOp.getUpperBound()),
-        bvm.lookupOrDefault(forOp.getStep()), packedTensor);
+        bvm.lookupOrDefault(forOp.getStep()), hoistedPackedTensor);
 
     // Map the induction var, region args and results to the `clonedForOp`.
     bvm.map(forOp.getInductionVar(), clonedForOp.getInductionVar());
     bvm.map(forOp.getRegionIterArgs(), clonedForOp.getRegionIterArgs());
     bvm.map(forOp.getResults(), clonedForOp.getResults());
     assert(clonedForOp->getNumRegions() == 1);
     clonedLoopIvs.push_back(clonedForOp.getInductionVar());
 
     // Do not insert guard here, we get deeper into the loop nest.
     rewriter.setInsertionPointToStart(&clonedForOp->getRegion(0).front());
     Value loopIndependentIterationCount =
         buildLoopIterationCount(rewriter, outerLoop, clonedForOp);
 
     // Assert the loop-independent iteration count can be computed.
     if (!loopIndependentIterationCount)
       llvm_unreachable("loop independence prerequisite not met");
-    leadingPackedTensorIndexings.push_back(loopIndependentIterationCount);
-    packedTensor = clonedForOp.getRegionIterArgs().front();
+    leadingHoistedPackedTensorIndexings.push_back(
+        loopIndependentIterationCount);
+    hoistedPackedTensor = clonedForOp.getRegionIterArgs().front();
   }
 
   // Step 2. Construct offsets, sizes and strides for the innermost level of the
   // packing loop.
   int64_t nPackedLoops = clonedLoopIvs.size();
   // offsets = [clonedLoopIvs, 0 .. 0].
-  offsets = SmallVector<OpFoldResult>{leadingPackedTensorIndexings.begin(),
-                                      leadingPackedTensorIndexings.end()};
+  offsets =
+      SmallVector<OpFoldResult>{leadingHoistedPackedTensorIndexings.begin(),
+                                leadingHoistedPackedTensorIndexings.end()};
   offsets.append(paddedRank, rewriter.getIndexAttr(0));
   // sizes = [1 .. 1, transposedShape].
   sizes = SmallVector<OpFoldResult>(nPackedLoops, rewriter.getIndexAttr(1));
   for (int64_t sz : transposedTensorType.getShape()) {
     // TODO: go grab dims when needed, atm tensor::PadOp yields a static tensor.
     assert(!ShapedType::isDynamic(sz) && "padded tensor needs static sizes");
     sizes.push_back(rewriter.getIndexAttr(sz));
   }
   // strides = [1 .. 1].
   strides = SmallVector<OpFoldResult>(nPackedLoops + paddedRank,
                                       rewriter.getIndexAttr(1));
 
   // Step 3. Optionally transpose the padded tensor.
   GenericOp maybeTransposeOp;
   Value paddedTensor = bvm.lookup(opToHoist.getResult());
   if (!transposeVector.empty()) {
     Value outputTensor = rewriter.create<tensor::ExtractSliceOp>(
-        loc, transposedTensorType, packedTensor, offsets, sizes, strides);
+        loc, transposedTensorType, hoistedPackedTensor, offsets, sizes,
+        strides);
     maybeTransposeOp = makeTransposeOp(rewriter, loc, paddedTensor,
                                        outputTensor, transposeVector);
     paddedTensor = maybeTransposeOp.getResult(0);
   }
 
   // Innermost tensor.insert_slice and yields are optional / need loops.
   if (nPackedLoops > 0) {
     // Step 4. Create InsertSliceOp at the innermost loop level, inserting an
     // optionally transposed padded slice into the packed tensor.
     Value inserted = rewriter.create<tensor::InsertSliceOp>(
-        loc, paddedTensor, packedTensor, offsets, sizes, strides);
+        loc, paddedTensor, hoistedPackedTensor, offsets, sizes, strides);
 
     // Step 5. Iteratively pop the stack and propagate the yield.
     Value valueToYield = inserted;
     for (Value iv : llvm::reverse(clonedLoopIvs)) {
       auto forOp = scf::getForInductionVarOwner(iv);
       rewriter.setInsertionPointToEnd(&forOp.getRegion().front());
       rewriter.create<scf::YieldOp>(loc, valueToYield);
       valueToYield = forOp.getResult(0);
     }
   }
 
   return PackingResult{
       offsets,
       sizes,
       strides,
       clonedLoopIvs,
-      leadingPackedTensorIndexings,
+      leadingHoistedPackedTensorIndexings,
       maybeTransposeOp,
       cast<tensor::PadOp>(bvm.lookup(opToHoist.getResult()).getDefiningOp())};
 }
 
 /// Build the packing loop nest required to hoist `opToHoist` above
 /// `outermostEnclosingForOp`.
 /// The loop nest is built just before `outermostEnclosingForOp`.
 static FailureOr<PackingResult> buildPackingLoopNestImpl(
if (failed(transposedTensorType)) {
     LLVM_DEBUG(DBGS() << "--Could not compute transposed type -> Skip\n");
     return failure();
   }
 
   // Create the packed tensor<?x?x..? x transposedShape>.
   SmallVector<int64_t> packedShape(nPackedLoops, ShapedType::kDynamic);
   // TODO: go grab dims when needed, atm tensor::PadOp yields a static tensor.
   llvm::append_range(packedShape, transposedTensorType->getShape());
-  auto packedTensorType = RankedTensorType::get(
+  auto hoistedPackedTensorType = RankedTensorType::get(
       packedShape, transposedTensorType->getElementType());
 
   // Set the insertion point right before the outer loop and start packing.
   scf::ForOp outerLoop = analysis.outermostEnclosingForOp;
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(outerLoop);
   SmallVector<Value> dynamicTensorSizes =
-      analysis.getPackedTensorSizes(rewriter, loc);
+      analysis.getHoistedPackedTensorSizes(rewriter, loc);
   auto emptyOp = rewriter.create<tensor::EmptyOp>(
-      loc, packedTensorType.getShape(), packedTensorType.getElementType(),
-      dynamicTensorSizes);
+      loc, hoistedPackedTensorType.getShape(),
+      hoistedPackedTensorType.getElementType(), dynamicTensorSizes);
 
   return buildPackingLoopNestImpl(rewriter, bvm, opToHoist, transposeVector,
                                   *transposedTensorType, emptyOp, analysis);
 }
 
 /// Build the packing loop nest required to hoist `opToHoist` above
 /// `outermostEnclosingForOp`.
 /// The loop nest is built just before `outermostEnclosingForOp`.
FailureOr<PackingResult> mlir::linalg::detail::buildPackingLoopNest(
   return buildPackingLoopNestImpl(rewriter, bvm, opToHoist, transposeVector,
                                   analysis);
 }
 
 //===----------------------------------------------------------------------===//
 // hoistPaddingOnTensors Implementation.
 //===----------------------------------------------------------------------===//
 
-// If the original consumer of `sliceOp` was a `forOp` (i.e. through an iter
-// arg), propagate the `packedTensor` value through the same iter arg.
-// TODO: for multiple loops we need to track the use to the innermost loop.
-static Value padThroughLoopIterArg(RewriterBase &rewriter, Value packedTensor,
-                                   tensor::ExtractSliceOp sliceOp,
-                                   scf::ForOp forOp) {
+/// Return true if we can walk back the use-def chain from `extractSliceOp` to
+/// expectedSource going through DestinationStyleOpInterface inits only.
+/// This is a poor man's analysis that is sufficient to check the extractSliceOp
+/// the matches tensor.pad we want to hoist.
+/// In the future, it will be easier to ensure this with a matching symmetric
+/// tensor.unpad op.
+static bool tracesBackToExpectedValue(tensor::ExtractSliceOp extractSliceOp,
+                                      Value expectedSource) {
+  LLVM_DEBUG(DBGS() << "Start tracesBackToExpectedValue on: " << extractSliceOp
+                    << "\n");
+  LLVM_DEBUG(DBGS() << "--with extractSlice: " << extractSliceOp << "\n");
+  Value source = extractSliceOp.getSource();
+  LLVM_DEBUG(DBGS() << "--with starting source: " << source << "\n");
+  while (source && source != expectedSource) {
+    auto destOp =
+        dyn_cast_or_null<DestinationStyleOpInterface>(source.getDefiningOp());
+    if (!destOp)
+      break;
+    LLVM_DEBUG(DBGS() << "--step dest op: " << destOp << "\n");
+    source =
+        destOp.getDpsInitOperand(source.cast<OpResult>().getResultNumber())
+            ->get();
+  }
+  LLVM_DEBUG(DBGS() << "--final source: " << source << "\n");
+  LLVM_DEBUG(DBGS() << "--expected source: " << expectedSource << "\n");
+  return source == expectedSource;
+}
+
+/// If the original consumer of `outerSliceOp` was a `forOp` (i.e. through an
+/// iter arg), propagate the `hoistedPackedTensor` value through the same iter
+/// arg.
+/// TODO: for multiple loops we need to track the use to the innermost loop.
+///
+/// Match:
+/// ```
+///   %outerSliceOp = tensor.extract_slice ..
+///   %f = scf.for ... iter_args(%arg0 = %outerSliceOp) {
+///     %hoistedPackedTensor = tensor.pad %arg0
+///     %1 = compute %hoistedPackedTensor
+///     %2 = tensor.extract_slice %1
+///     scf.yield %2
+///   }
+/// ```
+///
+/// and rewrite as:
+/// ```
+///   %outerSliceOp = tensor.extract_slice ..
+///   %hoistedPackedTensor = tensor.pad %outerSliceOp
+///   %f = scf.for ... iter_args(%arg0 = %hoistedPackedTensor) {
+///     %1 = compute %arg0
+///     scf.yield %1
+///   }
+///   %2 = tensor.extract_slice %forOp
+/// ```
+///
+/// Return null when no rewrite happened.
+static tensor::ExtractSliceOp
+padThroughLoopIterArg(RewriterBase &rewriter, Value paddedValueBeforeHoisting,
+                      Value hoistedPackedTensor,
+                      tensor::ExtractSliceOp outerSliceOp, scf::ForOp forOp) {
+  LLVM_DEBUG(DBGS() << "Start padThroughLoopIterArg on: " << forOp << "\n");
+  LLVM_DEBUG(DBGS() << "--paddedValueBeforeHoisting: "
+                    << paddedValueBeforeHoisting << "\n");
   OpOperand *pUse = nullptr;
-  for (OpOperand &use : sliceOp->getUses()) {
+  for (OpOperand &use : outerSliceOp->getUses()) {
     if (use.getOwner() == forOp) {
       assert(!pUse && "Multiple slice uses in the for loop");
       pUse = &use;
     }
   }
   assert(pUse && "No slice use in the for loop");
   OpBuilder::InsertionGuard g(rewriter);
-  rewriter.setInsertionPointAfter(packedTensor.getDefiningOp());
-  Value casted = rewriter.create<tensor::CastOp>(
-      packedTensor.getLoc(), pUse->get().getType(), packedTensor);
+  rewriter.setInsertionPointAfter(hoistedPackedTensor.getDefiningOp());
 
-  std::optional<unsigned> operandNumber =
+  std::optional<unsigned> maybeOperandNumber =
       forOp.getIterArgNumberForOpOperand(*pUse);
-  assert(operandNumber.has_value() && "expected a proper iter arg number");
+  assert(maybeOperandNumber.has_value() && "expected a proper iter arg number");
+
+  int64_t operandNumber = maybeOperandNumber.value();
+  auto yieldOp = cast<scf::YieldOp>(forOp.getBody(0)->getTerminator());
+  auto yieldingExtractSliceOp = yieldOp->getOperand(operandNumber)
+                                    .getDefiningOp<tensor::ExtractSliceOp>();
+  if (!yieldingExtractSliceOp)
+    return tensor::ExtractSliceOp();
+
+  // Poor man's analysis sufficient to ensure extractSlice matches tensor.pad.
+  // In the future, it will be easier to ensure this with a matching symmetric
+  // tensor.unpad op.
+  if (!tracesBackToExpectedValue(yieldingExtractSliceOp,
+                                 paddedValueBeforeHoisting))
+    return tensor::ExtractSliceOp();
 
   SmallVector<Value> initArgs = forOp.getInitArgs();
-  initArgs[operandNumber.value()] = casted;
-  rewriter.startRootUpdate(forOp);
-  forOp.getInitArgsMutable().assign(initArgs);
-  rewriter.finalizeRootUpdate(forOp);
-  return forOp.getRegionIterArgForOpOperand(*pUse);
+  initArgs[operandNumber] = hoistedPackedTensor;
+  SmallVector<Value> yieldOperands = yieldOp.getOperands();
+  yieldOperands[operandNumber] = yieldingExtractSliceOp.getSource();
+
+  int64_t numOriginalForOpResults = initArgs.size();
+  LLVM_DEBUG(DBGS() << "numOriginalForOpResults: " << numOriginalForOpResults
+                    << "\n");
+  tensor::ExtractSliceOp extracted;
+  {
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPointAfter(forOp);
+    extracted = rewriter.create<tensor::ExtractSliceOp>(
+        hoistedPackedTensor.getLoc(), hoistedPackedTensor,
+        outerSliceOp.getMixedOffsets(), outerSliceOp.getMixedSizes(),
+        outerSliceOp.getMixedStrides());
+    rewriter.replaceAllUsesWith(forOp.getResult(operandNumber), extracted);
+  }
+  scf::ForOp newForOp =
+      replaceLoopWithNewYields(rewriter, forOp, initArgs, yieldOperands);
+
+  LLVM_DEBUG(DBGS() << "newForOp results: " << newForOp.getNumResults()
+                    << "\n");
+  LLVM_DEBUG(DBGS() << "replace source of: " << extracted << "\n");
+  LLVM_DEBUG(DBGS() << "with result #"
+                    << numOriginalForOpResults + operandNumber
+                    << " of forOp, giving us: " << extracted << "\n");
+  rewriter.startRootUpdate(extracted);
+  extracted.getSourceMutable().assign(
+      newForOp.getResult(numOriginalForOpResults + operandNumber));
+  rewriter.finalizeRootUpdate(extracted);
+
+  LLVM_DEBUG(DBGS() << "replace uses of: " << paddedValueBeforeHoisting
+                    << "\n");
+  LLVM_DEBUG(DBGS() << "with region iter arg #"
+                    << numOriginalForOpResults + operandNumber << "\n");
+  rewriter.replaceAllUsesWith(
+      paddedValueBeforeHoisting,
+      newForOp.getRegionIterArg(numOriginalForOpResults + operandNumber));
+
+  return extracted;
 }
 
 /// Produce a tensor extracted from the packingResult. This can be used as a
 /// replacement for `opToHoist` in callers.
 static Value replaceByPackingResult(RewriterBase &rewriter,
                                     const IRMapping &bvm,
                                     tensor::PadOp opToHoist,
                                     RankedTensorType transposedTensorType,
static Value replaceByPackingResult(RewriterBase &rewriter,
   int paddedRank = paddedTensorType.getRank();
 
   int64_t nPackedLoops = packingResult.clonedLoopIvs.size();
   LLVM_DEBUG(DBGS() << "nPackedLoops: " << nPackedLoops << " loops\n");
 
   scf::ForOp outerLoop = analysis.outermostEnclosingForOp;
   ArrayRef<scf::ForOp> packingLoops = analysis.packingLoops;
 
-  Value packedTensor;
+  Value hoistedPackedTensor;
   SmallVector<Value> loopIterationCounts;
   SmallVector<OpFoldResult> offsets(nPackedLoops + paddedRank,
                                     rewriter.getIndexAttr(0));
   if (nPackedLoops > 0) {
     loopIterationCounts =
         llvm::to_vector<4>(llvm::map_range(packingLoops, [&](Operation *loop) {
           return buildLoopIterationCount(rewriter, outerLoop,
                                          cast<scf::ForOp>(loop));
         }));
     // Assert all loop iteration counts can be computed.
     if (llvm ::any_of(loopIterationCounts, [](Value v) { return !v; }))
       llvm_unreachable("loop independence prerequisite not met");
 
     // offsets = [maybe_leading_ivs = originalLoopIvs, 0 .. 0].
     std::copy(loopIterationCounts.begin(), loopIterationCounts.end(),
               offsets.begin());
-    packedTensor =
+    hoistedPackedTensor =
         scf::getForInductionVarOwner(packingResult.clonedLoopIvs.front())
             ->getResult(0);
   } else {
     // If no loops were created, this is just hoisting without packing.
-    packedTensor = bvm.lookup(opToHoist.getResult());
+    hoistedPackedTensor = bvm.lookup(opToHoist.getResult());
   }
 
-  LLVM_DEBUG(DBGS() << "packedTensor: " << packedTensor << "\n");
+  LLVM_DEBUG(DBGS() << "hoistedPackedTensor: " << hoistedPackedTensor << "\n");
 
   // If the consumer of `padOp` was a `forOp`, propagate through iter args.
   scf::ForOp forOp = analysis.padConsumingForOp;
   if (forOp) {
-    packedTensor =
-        padThroughLoopIterArg(rewriter, packedTensor, analysis.sliceOp, forOp);
+    return padThroughLoopIterArg(rewriter, opToHoist, hoistedPackedTensor,
+                                 analysis.sliceOp, forOp);
   }
 
   // offsets = [maybe_leading_ivs, 0 .. 0].
   // sizes = [1 .. 1, transposedShape] (defined above).
   // strides = [1 .. 1] (defined above)
   return rewriter.create<tensor::ExtractSliceOp>(
-      loc, transposedTensorType, packedTensor, offsets, packingResult.sizes,
-      packingResult.strides);
+      loc, transposedTensorType, hoistedPackedTensor, offsets,
+      packingResult.sizes, packingResult.strides);
 }
 
 FailureOr<Value> mlir::linalg::hoistPaddingOnTensors(
     RewriterBase &rewriter, tensor::PadOp opToHoist, int64_t numLoops,
     ArrayRef<int64_t> transposeVector, tensor::PadOp &hoistedOp,
     SmallVectorImpl<GenericOp> &transposeOps) {
   LLVM_DEBUG(DBGS() << "\n"; DBGS() << " Try to hoist " << *(opToHoist) << "\n";
              DBGS() << " by " << numLoops << " loops\n");

mlir/test/Dialect/Linalg/transform-op-hoist-pad.mlir

 func.func @pad_and_hoist_init(
   %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)
      -> tensor<24x25xf32> 
 {
 
   //      CHECK: scf.for %{{.*}} -> (tensor<24x25xf32>) {
   //      CHECK:   %[[PADDED:.*]] = tensor.pad %{{.*}} 
   //      CHECK:     : tensor<?x25xf32> to tensor<5x25xf32>
-  //      CHECK:   scf.for %{{.*}} iter_args(%[[INNER_PADDED:[0-9a-zA-Z]*]] = %[[PADDED]]) -> (tensor<5x25xf32>)
+  //      CHECK:   %[[SCF_YIELD:.*]] = scf.for %{{.*}} iter_args(%[[INNER_PADDED:[0-9a-zA-Z]*]] = %[[PADDED]]) -> (tensor<5x25xf32>)
   //      CHECK:     %[[RES:.*]] = linalg.matmul {{.*}} outs(%[[INNER_PADDED]]
   // CHECK-SAME:       : tensor<5x25xf32>
   //      CHECK:     scf.yield %[[RES]] : tensor<5x25xf32>
-  //      CHECK:   %[[CAST:.*]] = tensor.cast %{{.*}} : tensor<5x25xf32> to tensor<?x25xf32>
-  //      CHECK:   tensor.insert_slice %[[CAST]] into %{{.*}}[%{{.*}}, 0] [%{{.*}}, 25] [1, 1]
+  //      CHECK:   %[[EXTRACTED:.*]] = tensor.extract_slice %[[SCF_YIELD]][%{{.*}}, 0] [%{{.*}}, 25] [1, 1]
+  // CHECK-SAME:     : tensor<5x25xf32> to tensor<?x25xf32>
+  //      CHECK:   tensor.insert_slice %[[EXTRACTED]] into %{{.*}}[%{{.*}}, 0] [%{{.*}}, 25] [1, 1]
   // CHECK-SAME:     : tensor<?x25xf32> into tensor<24x25xf32>
   %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
   func.return %0 : tensor<24x25xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1