[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 ↗

(On Diff #266635)

Nit: newYieldedValues -> newIterOperands

mlir/lib/Transforms/Utils/LoopUtils.cpp

2534 ↗

(On Diff #266635)

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

2541 ↗

(On Diff #266635)

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

2546 ↗

(On Diff #266635)

Same here?

2551 ↗

(On Diff #266635)

How about using without_terminator()here?

2555 ↗

(On Diff #266635)

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

2530 ↗

(On Diff #266635)

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

2541 ↗

(On Diff #266635)

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

2555 ↗

(On Diff #266635)

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

2567 ↗

(On Diff #266635)

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

[bondhugula]

mlir/lib/Transforms/Utils/LoopUtils.cpp

2567 ↗

(On Diff #266635)

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

[bondhugula]

mlir/lib/Transforms/Utils/LoopUtils.cpp

2541–2543 ↗

(On Diff #266635)

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().

2560 ↗

(On Diff #266635)

Nit: Replace loop's results if ...

[nicolasvasilache]

mlir/lib/Transforms/Utils/LoopUtils.cpp

2534 ↗

(On Diff #266635)

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.

2567 ↗

(On Diff #266635)

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 ↗

(On Diff #267075)

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

2567 ↗

(On Diff #266635)

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

2567 ↗

(On Diff #266635)

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

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

+//===- Utils.h - SCF dialect utilities --------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines prototypes for various SCF utilities.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SCF_UTILS_H_
+#define MLIR_DIALECT_SCF_UTILS_H_
+
+namespace mlir {
+class OpBuilder;
+class ValueRange;
+
+namespace scf {
+class ForOp;
+class ParallelOp;
+} // end 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);
+
+} // end namespace mlir
+#endif // MLIR_DIALECT_SCF_UTILS_H_

mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRSCFTransforms
   ParallelLoopFusion.cpp
   ParallelLoopSpecialization.cpp
   ParallelLoopTiling.cpp
+  Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SCF
 
   DEPENDS
   MLIRSCFPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineOps
   MLIRIR
   MLIRPass
   MLIRSCF
   MLIRStandardOps
   MLIRSupport
   )

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

+//===- LoopUtils.cpp ---- Misc utilities for loop transformation ----------===//
+//
+// 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 miscellaneous loop transformation routines.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SCF/Utils.h"
+
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+
+using namespace mlir;
+
+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);
+
+  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;
+}

mlir/test/Transforms/loop-utils.mlir

+// RUN: mlir-opt -allow-unregistered-dialect -test-scf-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
   TestLiveness.cpp
   TestLoopMapping.cpp
   TestLoopParametricTiling.cpp
   TestLoopUnrolling.cpp
   TestOpaqueLoc.cpp
   TestMemRefBoundCheck.cpp
   TestMemRefDependenceCheck.cpp
   TestMemRefStrideCalculation.cpp
+  TestSCFUtils.cpp
   TestVectorTransforms.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   DEPENDS

mlir/test/lib/Transforms/TestSCFUtils.cpp

+//===- TestSCFUtils.cpp --- Pass to test independent SCF dialect 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 SCF dialect utils.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/SCF/Utils.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Passes.h"
+
+#include "llvm/ADT/SetVector.h"
+
+using namespace mlir;
+
+namespace {
+class TestSCFUtilsPass : public PassWrapper<TestSCFUtilsPass, FunctionPass> {
+public:
+  explicit TestSCFUtilsPass() {}
+
+  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 registerTestSCFUtilsPass() {
+  PassRegistration<TestSCFUtilsPass>("test-scf-utils", "test scf utils");
+}
+} // namespace mlir

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

 void registerTestLoopMappingPass();
 void registerTestLoopUnrollingPass();
 void registerTestMatchers();
 void registerTestMemRefDependenceCheck();
 void registerTestMemRefStrideCalculation();
 void registerTestOpaqueLoc();
 void registerTestParallelismDetection();
 void registerTestGpuParallelLoopMappingPass();
+void registerTestSCFUtilsPass();
 void registerTestVectorConversions();
 void registerVectorizerTestPass();
 } // namespace mlir
 
 static cl::opt<std::string>
     inputFilename(cl::Positional, cl::desc("<input file>"), cl::init("-"));
 
 static cl::opt<std::string> outputFilename("o", cl::desc("Output filename"),
#endif
   registerTestLoopMappingPass();
   registerTestLoopUnrollingPass();
   registerTestMatchers();
   registerTestMemRefDependenceCheck();
   registerTestMemRefStrideCalculation();
   registerTestOpaqueLoc();
   registerTestParallelismDetection();
   registerTestGpuParallelLoopMappingPass();
+  registerTestSCFUtilsPass();
   registerTestVectorConversions();
   registerVectorizerTestPass();
 }
 #endif
 
 static cl::opt<bool>
     showDialects("show-dialects",
                  cl::desc("Print the list of registered dialects"),