[mlir][SCF] Add utility to clone an scf.ForOp while appending new yield values.

[nicolasvasilache]

Summary

This utility factors out the machinery required to add iterArgs and yield values to an scf.ForOp.

[nicolasvasilache]

Drop spurious changes to unrelated parts of the code.

[nicolasvasilache]

Revert spurious changes.

[herhut]

mlir/include/mlir/Transforms/LoopUtils.h

297

Nit: newYieldedValues -> newIterOperands

mlir/lib/Transforms/Utils/LoopUtils.cpp

2532

Why not create this in the right block to start with?

2539

Would bvm.map(loop.region().front().getArguments(), newLoop.region().front().getArguments().take_front(loop.region().front().getNumArguments())) work?

2544

Same here?

2549

How about using without_terminator()here?

2553

bvm.map(o.getResults(), cloned->getResults())?

[ftynse]

It feels like this should live in SCF rather than LoopUtils, it's not anyhow generic.

mlir/lib/Transforms/Utils/LoopUtils.cpp

2528

Would loop.body().getTerminator() work instead?

2539

Could we make it a bit shorter by, e.g., caching the result of loop.region().front() in a variable loopBody?

2553

I think this is already done inside clone, no need to do it manually

2565

This is not allowed if OpBuilder is a ConversionRewriter, use b.eraseOp instead

[bondhugula]

mlir/lib/Transforms/Utils/LoopUtils.cpp

2565

Hmm... but there is no OpBuilder::eraseOp; it only exists in PatternRewriters.

[bondhugula]

mlir/lib/Transforms/Utils/LoopUtils.cpp

2539–2541

Since there's only one region and one block, region().front() -> getBody() / body(). I'm surprised there isn't an accessor. For eg. spirv::LoopOp has a body().

2558

Nit: Replace loop's results if ...

[nicolasvasilache]

mlir/lib/Transforms/Utils/LoopUtils.cpp

2532

Because I would need to manually replace the values, so I just let the clone do it then erase the op.
If you have strong feelings about this I can change to the other form.

2565

added one, let's see what @rriddle says

[nicolasvasilache]

Address review comments.

[ftynse]

mlir/include/mlir/IR/Builders.h

465 ↗

(On Diff #267075)

OpBuilder was refactored to use the Listener pattern instead of virtual methods. If we agree that we want to have eraseOp on OpBuilder, it should also go through the listener.

mlir/include/mlir/Transforms/LoopUtils.h

323

My commit-level comment might have gone unnoticed, so reiterating here: "It feels like this should live in SCF rather than LoopUtils, it's not anyhow generic."

mlir/lib/Transforms/Utils/LoopUtils.cpp

2565

Good point, @bondhugula !

@nicolasvasilache do you intend to use this in a place where no PatternRewriter is available? If not, just take a PatternRewriter& instead of OpBuilder& and be done with it.

[nicolasvasilache]

Address review.

[nicolasvasilache]

mlir/include/mlir/IR/Builders.h

465 ↗

(On Diff #267075)

reverted this for now, it can be done separately.

mlir/lib/Transforms/Utils/LoopUtils.cpp

2565

For now I reverted to a comment, when I need it in a PatternRewriter I'll update.
Thanks!

mlir/include/mlir/Transforms/LoopUtils.h

 ///      // full_tile
 ///    else
 ///      // partial tile
 ///
 LogicalResult
 separateFullTiles(MutableArrayRef<AffineForOp> nest,
                   SmallVectorImpl<AffineForOp> *fullTileNest = nullptr);
 
+/// Creates a clone of `loop` with `newYieldedValues` 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.
+///
+/// Returns 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);
+
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_LOOP_UTILS_H

mlir/lib/Transforms/Utils/LoopUtils.cpp

stripmineSink(AffineForOp forOp, uint64_t factor,
   auto ubMap = forOp.getUpperBoundMap();
   SmallVector<Value, 4> ubOperands(forOp.getUpperBoundOperands());
   augmentMapAndBounds(b, forOp.getInductionVar(), &ubMap, &ubOperands,
                       /*offset=*/scaledStep);
 
   auto iv = forOp.getInductionVar();
   SmallVector<AffineForOp, 8> innerLoops;
   for (auto t : targets) {
-    // Insert newForOp before the terminator of `t`.
+    // Insert newLoop before the terminator of `t`.
     auto b = OpBuilder::atBlockTerminator(t.getBody());
-    auto newForOp = b.create<AffineForOp>(t.getLoc(), lbOperands, lbMap,
-                                          ubOperands, ubMap, originalStep);
+    auto newLoop = b.create<AffineForOp>(t.getLoc(), lbOperands, lbMap,
     auto begin = t.getBody()->begin();
-    // Skip terminator and `newForOp` which is just before the terminator.
+    // Skip terminator and `newLoop` which is just before the terminator.
     auto nOps = t.getBody()->getOperations().size() - 2;
-    newForOp.getBody()->getOperations().splice(
-        newForOp.getBody()->getOperations().begin(),
+    newLoop.getBody()->getOperations().splice(
+        newLoop.getBody()->getOperations().begin(),
         t.getBody()->getOperations(), begin, std::next(begin, nOps));
-    replaceAllUsesInRegionWith(iv, newForOp.getInductionVar(),
-                               newForOp.region());
-    innerLoops.push_back(newForOp);
+    replaceAllUsesInRegionWith(iv, newLoop.getInductionVar(), newLoop.region());
+    innerLoops.push_back(newLoop);
   }
 
   return innerLoops;
 }
 
 static Loops stripmineSink(scf::ForOp forOp, Value factor,
                            ArrayRef<scf::ForOp> targets) {
   auto originalStep = forOp.step();
   auto iv = forOp.getInductionVar();
 
   OpBuilder b(forOp);
   forOp.setStep(b.create<MulIOp>(forOp.getLoc(), originalStep, factor));
 
   Loops innerLoops;
   for (auto t : targets) {
     // Save information for splicing ops out of t when done
     auto begin = t.getBody()->begin();
     auto nOps = t.getBody()->getOperations().size();
 
-    // Insert newForOp before the terminator of `t`.
+    // Insert newLoop before the terminator of `t`.
     auto b = OpBuilder::atBlockTerminator((t.getBody()));
     Value stepped = b.create<AddIOp>(t.getLoc(), iv, forOp.step());
     Value less = b.create<CmpIOp>(t.getLoc(), CmpIPredicate::slt,
                                   forOp.upperBound(), stepped);
     Value ub =
         b.create<SelectOp>(t.getLoc(), less, forOp.upperBound(), stepped);
 
-    // Splice [begin, begin + nOps - 1) into `newForOp` and replace uses.
-    auto newForOp = b.create<scf::ForOp>(t.getLoc(), iv, ub, originalStep);
-    newForOp.getBody()->getOperations().splice(
-        newForOp.getBody()->getOperations().begin(),
+    // Splice [begin, begin + nOps - 1) into `newLoop` and replace uses.
+    auto newLoop = b.create<scf::ForOp>(t.getLoc(), iv, ub, originalStep);
+    newLoop.getBody()->getOperations().splice(
+        newLoop.getBody()->getOperations().begin(),
         t.getBody()->getOperations(), begin, std::next(begin, nOps - 1));
-    replaceAllUsesInRegionWith(iv, newForOp.getInductionVar(),
-                               newForOp.region());
+    replaceAllUsesInRegionWith(iv, newLoop.getInductionVar(), newLoop.region());
 
-    innerLoops.push_back(newForOp);
+    innerLoops.push_back(newLoop);
   }
 
   return innerLoops;
 }
 
 // Stripmines a `forOp` by `factor` and sinks it under a single `target`.
 // Returns the new AffineForOps, nested immediately under `target`.
 template <typename ForType, typename SizeType>
elseBlock->getOperations().splice(
       firstLoop.getOperation()->getBlock()->getOperations(),
       Block::iterator(firstLoop));
 
   if (fullTileNest)
     *fullTileNest = std::move(fullTileLoops);
 
   return success();
 }
+
+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.lowerBound(), loop.upperBound(),
+                           loop.step(), operands);
+
+  // 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>(loop.region().front().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.
+  for (auto it : llvm::zip(loop.region().front().getArguments(),
+                           newLoop.region().front().getArguments().take_front(
+                               loop.region().front().getNumArguments())))
+    bvm.map(std::get<0>(it), std::get<1>(it));
+  // c. remap
+  for (auto it :
+       llvm::zip(newIterOperands,
+                 newLoop.getRegionIterArgs().take_back(newIterOperands.size())))
+    bvm.map(std::get<0>(it), std::get<1>(it));
+  b.setInsertionPointToStart(&newLoop.region().front());
+  for (auto &o : loop.region().front().getOperations()) {
+    if (&o != yield.getOperation()) {
+      // Skip the original yield terminator which does not have enough operands.
+      auto *cloned = b.clone(o, bvm);
+      for (unsigned idx = 0, e = cloned->getNumResults(); idx < e; ++idx)
+        bvm.map(o.getResult(idx), cloned->getResult(idx));
+    }
+  }
+
+  // Perform `loop` results replacement 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));
+  }
+
+  newYield.erase();
+  return newLoop;
+}

mlir/test/Transforms/loop-utils.mlir

+// RUN: mlir-opt -allow-unregistered-dialect -test-loop-utils -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 @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 @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/Transforms/CMakeLists.txt

add_mlir_library(MLIRTestTransforms
   TestGpuMemoryPromotion.cpp
   TestGpuParallelLoopMapping.cpp
   TestInlining.cpp
   TestLinalgTransforms.cpp
   TestLiveness.cpp
   TestLoopMapping.cpp
   TestLoopParametricTiling.cpp
   TestLoopUnrolling.cpp
+  TestLoopUntils.cpp
   TestOpaqueLoc.cpp
   TestMemRefBoundCheck.cpp
   TestMemRefDependenceCheck.cpp
   TestMemRefStrideCalculation.cpp
   TestVectorToSCFConversion.cpp
   TestVectorTransforms.cpp
 
   EXCLUDE_FROM_LIBMLIR

mlir/test/lib/Transforms/TestLoopUtils.cpp

+//===- TestLoopUtils.cpp --- Pass to test independent loop utils ----------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to test loop utils.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/LoopUtils.h"
+#include "mlir/Transforms/Passes.h"
+
+#include "llvm/ADT/SetVector.h"
+
+using namespace mlir;
+
+namespace {
+class TestLoopUtilsPass : public PassWrapper<TestLoopUtilsPass, FunctionPass> {
+public:
+  explicit TestLoopUtilsPass() {}
+
+  void runOnFunction() override {
+    FuncOp func = getFunction();
+    SmallVector<scf::ForOp, 4> toErase;
+
+    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();
+  }
+};
+} // end namespace
+
+namespace mlir {
+void registerTestLoopUtilsPass() {
+  PassRegistration<TestLoopUtilsPass>("test-loop-utils", "test loop utils");
+}
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestDominancePass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestLinalgTransforms();
 void registerTestLivenessPass();
 void registerTestLoopFusion();
 void registerTestLoopMappingPass();
 void registerTestLoopUnrollingPass();
+void registerTestLoopUtilsPass();
 void registerTestMatchers();
 void registerTestMemRefDependenceCheck();
 void registerTestMemRefStrideCalculation();
 void registerTestOpaqueLoc();
 void registerTestParallelismDetection();
 void registerTestGpuParallelLoopMappingPass();
 void registerTestVectorConversions();
 void registerTestVectorToSCFPass();
#endif
   registerTestDominancePass();
   registerTestFunc();
   registerTestGpuMemoryPromotionPass();
   registerTestLinalgTransforms();
   registerTestLivenessPass();
   registerTestLoopFusion();
   registerTestLoopMappingPass();
   registerTestLoopUnrollingPass();
+  registerTestLoopUtilsPass();
   registerTestMatchers();
   registerTestMemRefDependenceCheck();
   registerTestMemRefStrideCalculation();
   registerTestOpaqueLoc();
   registerTestParallelismDetection();
   registerTestGpuParallelLoopMappingPass();
   registerTestVectorConversions();
   registerTestVectorToSCFPass();