[mlir][SCF] Retire cloneWithNewYields helper function.

[nicolasvasilache]

Summary

This is now subsumed by replaceLoopWithNewYields.

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

 } // namespace func
 
 namespace scf {
 class IfOp;
 class ForOp;
 class ParallelOp;
 } // namespace scf
 
-/// Create a clone of `loop` with `newIterOperands` added as new initialization
-/// values and `newYieldedValues` added as new yielded values. The returned
-/// ForOp has `newYieldedValues.size()` new result values.  The `loop` induction
-/// variable and `newIterOperands` are remapped to the new induction variable
-/// and the new entry block arguments respectively.
-///
-/// Additionally, if `replaceLoopResults` is true, all uses of
-/// `loop.getResults()` are replaced with the first `loop.getNumResults()`
-/// return values respectively. This additional replacement is provided as a
-/// convenience to update the consumers of `loop`, in the case e.g. when `loop`
-/// is soon to be deleted.
-///
-/// Return the cloned loop.
-///
-/// This convenience function is useful to factorize common mechanisms related
-/// to hoisting roundtrips to memory into yields. It does not perform any
-/// legality checks.
-///
-/// Prerequisite: `newYieldedValues.size() == newYieldedValues.size()`.
-scf::ForOp cloneWithNewYields(OpBuilder &b, scf::ForOp loop,
-                              ValueRange newIterOperands,
-                              ValueRange newYieldedValues,
-                              bool replaceLoopResults = true);
-
 /// Replace the `loop` with `newIterOperands` added as new initialization
 /// values. `newYieldValuesFn` is a callback that can be used to specify
 /// the additional values to be yielded by the loop. The number of
 /// 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.
-/// TODO: This method could be used instead of `cloneWithNewYields`. Making
-/// this change though hits assertions in the walk mechanism that is unrelated
-/// to this method itself.
 using NewYieldValueFn = std::function<SmallVector<Value>(
     OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs)>;
 scf::ForOp replaceLoopWithNewYields(OpBuilder &builder, scf::ForOp loop,
                                     ValueRange newIterOperands,
                                     NewYieldValueFn newYieldValuesFn);
 
 /// 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

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

static void hoistReadWrite(HoistableRead read, HoistableWrite write,
   auto yieldOp = cast<scf::YieldOp>(forOp.getRegion().front().getTerminator());
   if (write.insertSliceOp)
     yieldOp->setOperand(initArgNumber, write.insertSliceOp.dest());
   else
     yieldOp->setOperand(initArgNumber, write.transferWriteOp.getSource());
 
   // Rewrite `loop` with additional new yields.
   OpBuilder b(read.transferReadOp);
-  auto newForOp = cloneWithNewYields(b, forOp, read.transferReadOp.getVector(),
-                                     write.transferWriteOp.getVector());
+  NewYieldValueFn yieldFn = [&](OpBuilder &b, Location loc,
+                                ArrayRef<BlockArgument> newBBArgs) {
+    return SmallVector<Value>{write.transferWriteOp.getVector()};
+  };
+  auto newForOp = replaceLoopWithNewYields(
+      b, forOp, read.transferReadOp.getVector(), yieldFn);
+
   // Transfer write has been hoisted, need to update the vector and tensor
   // source. Replace the result of the loop to use the new tensor created
   // outside the loop.
   // Depending on whether a insert_slice is present or not, it carries the
   // update on the tensor operands.
   if (write.insertSliceOp) {
     newForOp.getResult(initArgNumber)
         .replaceAllUsesWith(write.insertSliceOp.getResult());
void mlir::linalg::hoistRedundantVectorTransfersOnTensor(func::FuncOp func) {
   }
 }
 
 void mlir::linalg::hoistRedundantVectorTransfers(func::FuncOp func) {
   bool changed = true;
   while (changed) {
     changed = false;
     // First move loop invariant ops outside of their loop. This needs to be
-    // done before as we cannot move ops without interputing the function walk.
-    func.walk([&](LoopLikeOpInterface loopLike) {
-      moveLoopInvariantCode(loopLike);
-    });
+    // done before as we cannot move ops without interrupting the function walk.
+    func.walk(
+        [&](LoopLikeOpInterface loopLike) { moveLoopInvariantCode(loopLike); });
 
     func.walk([&](vector::TransferReadOp transferRead) {
       if (!transferRead.getShapedType().isa<MemRefType>())
         return WalkResult::advance();
 
       LLVM_DEBUG(DBGS() << "Candidate for hoisting: "
                         << *transferRead.getOperation() << "\n");
       auto loop = dyn_cast<scf::ForOp>(transferRead->getParentOp());
func.walk([&](vector::TransferReadOp transferRead) {
       // Hoist read before.
       loop.moveOutOfLoop(transferRead);
 
       // Hoist write after.
       transferWrite->moveAfter(loop);
 
       // Rewrite `loop` with new yields by cloning and erase the original loop.
       OpBuilder b(transferRead);
-      auto newForOp = cloneWithNewYields(b, loop, transferRead.getVector(),
-                                         transferWrite.getVector());
+      NewYieldValueFn yieldFn = [&](OpBuilder &b, Location loc,
+                                    ArrayRef<BlockArgument> newBBArgs) {
+        return SmallVector<Value>{transferWrite.getVector()};
+      };
+      auto newForOp =
+          replaceLoopWithNewYields(b, loop, transferRead.getVector(), yieldFn);
 
-      // Transfer write has been hoisted, need to update the written value to
+      // Transfer write has been hoisted, need to update the written vector by
       // the value yielded by the newForOp.
-      transferWrite.getVector().replaceAllUsesWith(
-          newForOp.getResults().take_back()[0]);
+      transferWrite.getVectorMutable().assign(newForOp.getResults().back());
 
       changed = true;
       loop.erase();
       // Need to interrupt and restart because erasing the loop messes up the
       // walk.
       return WalkResult::interrupt();
     });
   }
 }

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

 // confusing the order.
 struct LoopParams {
   Value lowerBound;
   Value upperBound;
   Value step;
 };
 } // namespace
 
-scf::ForOp mlir::cloneWithNewYields(OpBuilder &b, scf::ForOp loop,
-                                    ValueRange newIterOperands,
-                                    ValueRange newYieldedValues,
-                                    bool replaceLoopResults) {
-  assert(newIterOperands.size() == newYieldedValues.size() &&
-         "newIterOperands must be of the same size as newYieldedValues");
-
-  // Create a new loop before the existing one, with the extra operands.
-  OpBuilder::InsertionGuard g(b);
-  b.setInsertionPoint(loop);
-  auto operands = llvm::to_vector<4>(loop.getIterOperands());
-  operands.append(newIterOperands.begin(), newIterOperands.end());
-  scf::ForOp newLoop =
-      b.create<scf::ForOp>(loop.getLoc(), loop.getLowerBound(),
-                           loop.getUpperBound(), loop.getStep(), operands);
-
-  auto &loopBody = *loop.getBody();
-  auto &newLoopBody = *newLoop.getBody();
-  // Clone / erase the yield inside the original loop to both:
-  //   1. augment its operands with the newYieldedValues.
-  //   2. automatically apply the BlockAndValueMapping on its operand
-  auto yield = cast<scf::YieldOp>(loopBody.getTerminator());
-  b.setInsertionPoint(yield);
-  auto yieldOperands = llvm::to_vector<4>(yield.getOperands());
-  yieldOperands.append(newYieldedValues.begin(), newYieldedValues.end());
-  auto newYield = b.create<scf::YieldOp>(yield.getLoc(), yieldOperands);
-
-  // Clone the loop body with remaps.
-  BlockAndValueMapping bvm;
-  // a. remap the induction variable.
-  bvm.map(loop.getInductionVar(), newLoop.getInductionVar());
-  // b. remap the BB args.
-  bvm.map(loopBody.getArguments(),
-          newLoopBody.getArguments().take_front(loopBody.getNumArguments()));
-  // c. remap the iter args.
-  bvm.map(newIterOperands,
-          newLoop.getRegionIterArgs().take_back(newIterOperands.size()));
-  b.setInsertionPointToStart(&newLoopBody);
-  // Skip the original yield terminator which does not have enough operands.
-  for (auto &o : loopBody.without_terminator())
-    b.clone(o, bvm);
-
-  // Replace `loop`'s results if requested.
-  if (replaceLoopResults) {
-    for (auto it : llvm::zip(loop.getResults(), newLoop.getResults().take_front(
-                                                    loop.getNumResults())))
-      std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
-  }
-
-  // TODO: this is unsafe in the context of a PatternRewrite.
-  newYield.erase();
-
-  return newLoop;
-}
-
 scf::ForOp mlir::replaceLoopWithNewYields(OpBuilder &builder, scf::ForOp loop,
                                           ValueRange newIterOperands,
                                           NewYieldValueFn newYieldValuesFn) {
   // 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) {});
 
   Block *loopBody = loop.getBody();
   Block *newLoopBody = newLoop.getBody();
 
   // Move the body of the original loop to the new loop.
   newLoopBody->getOperations().splice(newLoopBody->end(),
                                       loopBody->getOperations());
 
-  // Generate the new yield values to use by using the callback and ppend the
+  // Generate the new yield values to use by using the callback and append the
   // yield values to the scf.yield operation.
   auto yield = cast<scf::YieldOp>(newLoopBody->getTerminator());
   ArrayRef<BlockArgument> newBBArgs =
       newLoopBody->getArguments().take_back(newIterOperands.size());
   {
     OpBuilder::InsertionGuard g(builder);
     builder.setInsertionPoint(yield);
     SmallVector<Value> newYieldedValues =

mlir/test/Transforms/scf-loop-utils.mlir

-// RUN: mlir-opt -allow-unregistered-dialect -test-scf-for-utils=test-clone-with-new-yields -mlir-disable-threading %s | FileCheck %s
-
-// CHECK-LABEL: @hoist
-//  CHECK-SAME: %[[lb:[a-zA-Z0-9]*]]: index,
-//  CHECK-SAME: %[[ub:[a-zA-Z0-9]*]]: index,
-//  CHECK-SAME: %[[step:[a-zA-Z0-9]*]]: index
-func.func @hoist(%lb: index, %ub: index, %step: index) {
-  // CHECK: %[[A:.*]] = "fake_read"() : () -> index
-  // CHECK: %[[RES:.*]] = scf.for %[[I:.*]] = %[[lb]] to %[[ub]] step %[[step]] iter_args(%[[VAL:.*]] = %[[A]]) -> (index)
-  // CHECK:   %[[YIELD:.*]] = "fake_compute"(%[[VAL]]) : (index) -> index
-  // CHECK:   scf.yield %[[YIELD]] : index
-  // CHECK: "fake_write"(%[[RES]]) : (index) -> ()
-  scf.for %i = %lb to %ub step %step {
-    %0 = "fake_read"() : () -> (index)
-    %1 = "fake_compute"(%0) : (index) -> (index)
-    "fake_write"(%1) : (index) -> ()
-  }
-  return
-}
-
-// CHECK-LABEL: @hoist2
-//  CHECK-SAME: %[[lb:[a-zA-Z0-9]*]]: index,
-//  CHECK-SAME: %[[ub:[a-zA-Z0-9]*]]: index,
-//  CHECK-SAME: %[[step:[a-zA-Z0-9]*]]: index
-//  CHECK-SAME: %[[extra_arg:[a-zA-Z0-9]*]]: f32
-func.func @hoist2(%lb: index, %ub: index, %step: index, %extra_arg: f32) -> f32 {
-  // CHECK: %[[A:.*]] = "fake_read"() : () -> index
-  // CHECK: %[[RES:.*]]:2 = scf.for %[[I:.*]] = %[[lb]] to %[[ub]] step %[[step]] iter_args(%[[VAL0:.*]] = %[[extra_arg]], %[[VAL1:.*]] = %[[A]]) -> (f32, index)
-  // CHECK:   %[[YIELD:.*]] = "fake_compute"(%[[VAL1]]) : (index) -> index
-  // CHECK:   scf.yield %[[VAL0]], %[[YIELD]] : f32, index
-  // CHECK: "fake_write"(%[[RES]]#1) : (index) -> ()
-  // CHECK: return %[[RES]]#0 : f32
-  %0 = scf.for %i = %lb to %ub step %step iter_args(%iter = %extra_arg) -> (f32) {
-    %0 = "fake_read"() : () -> (index)
-    %1 = "fake_compute"(%0) : (index) -> (index)
-    "fake_write"(%1) : (index) -> ()
-    scf.yield %iter: f32
-  }
-  return %0: f32
-}

mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp

struct TestSCFForUtilsPass
     : public PassWrapper<TestSCFForUtilsPass, OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSCFForUtilsPass)
 
   StringRef getArgument() const final { return "test-scf-for-utils"; }
   StringRef getDescription() const final { return "test scf.for utils"; }
   explicit TestSCFForUtilsPass() = default;
   TestSCFForUtilsPass(const TestSCFForUtilsPass &pass) : PassWrapper(pass) {}
 
-  Option<bool> testCloneWithNewYields{
-      *this, "test-clone-with-new-yields",
-      llvm::cl::desc(
-          "Test cloning of a loop while returning additional yield values"),
-      llvm::cl::init(false)};
-
   Option<bool> testReplaceWithNewYields{
       *this, "test-replace-with-new-yields",
       llvm::cl::desc("Test replacing a loop with a new loop that returns new "
                      "additional yeild values"),
       llvm::cl::init(false)};
 
   void runOnOperation() override {
     func::FuncOp func = getOperation();
     SmallVector<scf::ForOp, 4> toErase;
 
-    if (testCloneWithNewYields) {
-      func.walk([&](Operation *fakeRead) {
-        if (fakeRead->getName().getStringRef() != "fake_read")
-          return;
-        auto *fakeCompute = fakeRead->getResult(0).use_begin()->getOwner();
-        auto *fakeWrite = fakeCompute->getResult(0).use_begin()->getOwner();
-        auto loop = fakeRead->getParentOfType<scf::ForOp>();
-
-        OpBuilder b(loop);
-        loop.moveOutOfLoop(fakeRead);
-        fakeWrite->moveAfter(loop);
-        auto newLoop = cloneWithNewYields(b, loop, fakeRead->getResult(0),
-                                          fakeCompute->getResult(0));
-        fakeCompute->getResult(0).replaceAllUsesWith(
-            newLoop.getResults().take_back()[0]);
-        toErase.push_back(loop);
-      });
-      for (auto loop : llvm::reverse(toErase))
-        loop.erase();
-    }
-
     if (testReplaceWithNewYields) {
       func.walk([&](scf::ForOp forOp) {
         if (forOp.getNumResults() == 0)
           return;
         auto newInitValues = forOp.getInitArgs();
         if (newInitValues.empty())
           return;
         NewYieldValueFn fn = [&](OpBuilder &b, Location loc,