[MLIR][SCF] Add for-to-while loop transformation pass

[mortbopet]

Summary

This pass transforms SCF.ForOp operations to SCF.WhileOp. The For loop condition is placed in the 'before' region of the while operation, and indctuion variable incrementation + the loop body in the 'after' region. The loop carried values of the while op are the induction variable (IV) of the for-loop + any iter_args specified for the for-loop.
Any 'yield' ops in the for-loop are rewritten to additionally yield the (incremented) induction variable.

This transformation is useful for passes where we want to consider structured control flow solely on the basis of a loop body and the computation of a loop condition. As an example, when doing high-level synthesis in CIRCT, the incrementation of an IV in a for-loop is "just another part" of a circuit datapath, and what we really care about is the distinction between our datapath and our control logic (the condition variable).

[mortbopet]

Fix clang-tidy nits

[mortbopet]

• Moved induction variable incrementation to start of while body
• updated tests

[mortbopet]

scf.execute_region ops are no longer necessary.

[nicolasvasilache]

Could you add some context as to when this is expected to be useful?
Is there also a plan to replace scf -> block lowering by scf -> while -> blocks ?

[mortbopet]

I've added a rationale to the PR summary.

Could you elaborate on what you mean by "scf -> block lowering by scf -> while -> blocks"?

[nicolasvasilache]

In D108454#3006185, @mortbopet wrote:

I've added a rationale to the PR summary.

Could you elaborate on what you mean by "scf -> block lowering by scf -> while -> blocks"?

thanks!

[mortbopet]

mising newline

[stephenneuendorffer]

Thanks, Morten!

mlir/include/mlir/Dialect/SCF/Passes.td

69

Even though this is pretty simple, can you add a description here for the generated documentation?

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

1

Name wrong?

41

I didn't think this was necessary? Isn't the insertionPoint set to the op being lowered already?

89

"erase the it"?

mlir/test/Dialect/SCF/for-loop-to-while-loop.mlir

6

The preference is not to assume anything about naming in tests and to use the regular expression matching capabilities of FileCheck.

[mortbopet]

Address review comments

[mehdi_amini]

I think you had tested this on an old revision but it failed CI at HEAD, can you rebase and fix the test before recommitting?

[mortbopet]

In D108454#3009869, @mehdi_amini wrote:

I think you had tested this on an old revision but it failed CI at HEAD, can you rebase and fix the test before recommitting?

Sorry about that; rebased and tested locally, and the tests seemed to run successfully. I'll update this PR and rerun CI - thank you for notifying me.

mlir/include/mlir/Dialect/SCF/Passes.h

 std::unique_ptr<Pass>
 createParallelLoopTilingPass(llvm::ArrayRef<int64_t> tileSize = {},
                              bool noMinMaxBounds = false);
 
 /// Creates a pass which folds arith ops on induction variable into
 /// loop range.
 std::unique_ptr<Pass> createForLoopRangeFoldingPass();
 
+// Creates a pass which lowers for loops into while loops.
+std::unique_ptr<Pass> createForToWhileLoopPass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/SCF/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_PASSES_H_

mlir/include/mlir/Dialect/SCF/Passes.td

 }
 
 def SCFForLoopRangeFolding
     : Pass<"for-loop-range-folding"> {
   let summary = "Fold add/mul ops into loop range";
   let constructor = "mlir::createForLoopRangeFoldingPass()";
 }
 
+def SCFForToWhileLoop
+    : FunctionPass<"scf-for-to-while"> {
+  let summary = "Convert SCF for loops to SCF while loops";
+  let constructor = "mlir::createForToWhileLoopPass()";
+}
+
 #endif // MLIR_DIALECT_SCF_PASSES

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

 add_mlir_dialect_library(MLIRSCFTransforms
   Bufferize.cpp
+  ForToWhile.cpp
   LoopPipelining.cpp
   LoopRangeFolding.cpp
   LoopSpecialization.cpp
   ParallelLoopFusion.cpp
   ParallelLoopTiling.cpp
   StructuralTypeConversions.cpp
   Utils.cpp
 

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

+//===- LoopSpecialization.cpp - scf.parallel/SCR.for specialization -------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Transforms SCF.ForOps into SCF.WhileOps.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/SCF/Passes.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/SCF/Transforms.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace llvm;
+using namespace mlir;
+using scf::ForOp;
+using scf::WhileOp;
+
+namespace {
+
+struct ForLoopLoweringPattern : public OpRewritePattern<ForOp> {
+  using OpRewritePattern<ForOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ForOp forOp,
+                                PatternRewriter &rewriter) const override {
+    // Generate type signature for the loop-carried values. The induction
+    // variable is placed first, followed by the forOp.iterArgs.
+    SmallVector<Type, 8> lcvTypes;
+    lcvTypes.push_back(forOp.getInductionVar().getType());
+    llvm::transform(forOp.initArgs(), std::back_inserter(lcvTypes),
+                    [&](auto v) { return v.getType(); });
+
+    // Build scf.WhileOp
+    rewriter.setInsertionPoint(forOp);
+    SmallVector<Value> initArgs;
+    initArgs.push_back(forOp.lowerBound());
+    llvm::append_range(initArgs, forOp.initArgs());
+    auto whileOp = rewriter.create<WhileOp>(forOp.getLoc(), lcvTypes, initArgs,
+                                            forOp->getAttrs());
+
+    // 'before' region contains the loop condition and forwarding of iteration
+    // arguments to the 'after' region.
+    auto *beforeBlock = rewriter.createBlock(
+        &whileOp.before(), whileOp.before().begin(), lcvTypes, {});
+    rewriter.setInsertionPointToStart(&whileOp.before().front());
+    auto cmpOp = rewriter.create<CmpIOp>(whileOp.getLoc(), CmpIPredicate::slt,
+                                         beforeBlock->getArgument(0),
+                                         forOp.upperBound());
+    rewriter.create<scf::ConditionOp>(whileOp.getLoc(), cmpOp.getResult(),
+                                      beforeBlock->getArguments());
+
+    // Inline for-loop body into an executeRegion operation in the "after"
+    // region. The return type of the execRegionOp does not contain the
+    // iv - yields in the source for-loop contain only iterArgs.
+    auto *afterBlock = rewriter.createBlock(
+        &whileOp.after(), whileOp.after().begin(), lcvTypes, {});
+
+    // Add induction variable incrementation
+    rewriter.setInsertionPointToEnd(afterBlock);
+    auto ivIncOp = rewriter.create<AddIOp>(
+        whileOp.getLoc(), afterBlock->getArgument(0), forOp.step());
+
+    // Rewrite uses of the for-loop block arguments to the new while-loop
+    // "after" arguments
+    for (auto barg : enumerate(forOp.getBody(0)->getArguments()))
+      barg.value().replaceAllUsesWith(afterBlock->getArgument(barg.index()));
+
+    // Inline for-loop body operations into 'after' region.
+    for (auto &arg : llvm::make_early_inc_range(*forOp.getBody()))
+      arg.moveBefore(afterBlock, afterBlock->end());
+
+    // Add incremented IV to yield operations
+    for (auto yieldOp : afterBlock->getOps<scf::YieldOp>()) {
+      SmallVector<Value> yieldOperands = yieldOp.getOperands();
+      yieldOperands.insert(yieldOperands.begin(), ivIncOp.getResult());
+      yieldOp->setOperands(yieldOperands);
+    }
+
+    // We cannot do a direct replacement of the forOp since the while op returns
+    // one more value (the induction variable escapes the loop through being
+    // carried in the set of iterargs). Instead, rewrite uses of the forOp
+    // results, and erase the it.
+    for (auto arg : llvm::enumerate(forOp.getResults()))
+      arg.value().replaceAllUsesWith(whileOp.getResult(arg.index() + 1));
+
+    rewriter.eraseOp(forOp);
+    return success();
+  }
+};
+
+struct ForToWhileLoop : public SCFForToWhileLoopBase<ForToWhileLoop> {
+  void runOnFunction() override {
+    FuncOp funcOp = getFunction();
+    MLIRContext *ctx = funcOp.getContext();
+    RewritePatternSet patterns(ctx);
+    patterns.add<ForLoopLoweringPattern>(ctx);
+    (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::createForToWhileLoopPass() {
+  return std::make_unique<ForToWhileLoop>();
+}

mlir/test/Dialect/SCF/for-loop-to-while-loop.mlir

+// RUN: mlir-opt %s -pass-pipeline='builtin.func(scf-for-to-while)' -split-input-file | FileCheck %s
+// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
+
+// CHECK:      builtin.module  {
+// CHECK-NEXT:   builtin.func @single_loop(%arg0: memref<?xi32>, %arg1: index, %arg2: i32) {
+// CHECK-NEXT:     %c0 = constant 0 : index
+// CHECK-NEXT:     %c1 = constant 1 : index
+// CHECK-NEXT:     %0 = scf.while (%arg3 = %c0) : (index) -> index {
+// CHECK-NEXT:       %1 = cmpi slt, %arg3, %arg1 : index
+// CHECK-NEXT:       scf.condition(%1) %arg3 : index
+// CHECK-NEXT:     } do {
+// CHECK-NEXT:     ^bb0(%arg3: index):  // no predecessors
+// CHECK-NEXT:       %1 = addi %arg3, %c1 : index
+// CHECK-NEXT:       %2 = addi %arg2, %arg2 : i32
+// CHECK-NEXT:       memref.store %2, %arg0[%arg3] : memref<?xi32>
+// CHECK-NEXT:       scf.yield %1 : index
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+func @single_loop(%arg0: memref<?xi32>, %arg1: index, %arg2: i32) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  scf.for %i = %c0 to %arg1 step %c1 {
+    %0 = addi %arg2, %arg2 : i32
+    memref.store %0, %arg0[%i] : memref<?xi32>
+  }
+  return
+}
+
+// -----
+
+// CHECK:      builtin.module  {
+// CHECK-NEXT:   builtin.func @nested_loop(%arg0: memref<?xi32>, %arg1: index, %arg2: i32) {
+// CHECK-NEXT:     %c0 = constant 0 : index
+// CHECK-NEXT:     %c1 = constant 1 : index
+// CHECK-NEXT:     %0 = scf.while (%arg3 = %c0) : (index) -> index {
+// CHECK-NEXT:       %1 = cmpi slt, %arg3, %arg1 : index
+// CHECK-NEXT:       scf.condition(%1) %arg3 : index
+// CHECK-NEXT:     } do {
+// CHECK-NEXT:     ^bb0(%arg3: index):  // no predecessors
+// CHECK-NEXT:       %1 = addi %arg3, %c1 : index
+// CHECK-NEXT:       %2 = scf.while (%arg4 = %c0) : (index) -> index {
+// CHECK-NEXT:         %3 = cmpi slt, %arg4, %arg1 : index
+// CHECK-NEXT:         scf.condition(%3) %arg4 : index
+// CHECK-NEXT:       } do {
+// CHECK-NEXT:       ^bb0(%arg4: index):  // no predecessors
+// CHECK-NEXT:         %3 = addi %arg4, %c1 : index
+// CHECK-NEXT:         %4 = addi %arg2, %arg2 : i32
+// CHECK-NEXT:         memref.store %4, %arg0[%arg3] : memref<?xi32>
+// CHECK-NEXT:         memref.store %4, %arg0[%arg4] : memref<?xi32>
+// CHECK-NEXT:         scf.yield %3 : index
+// CHECK-NEXT:       }
+// CHECK-NEXT:       scf.yield %1 : index
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+func @nested_loop(%arg0: memref<?xi32>, %arg1: index, %arg2: i32) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  scf.for %i = %c0 to %arg1 step %c1 {
+    scf.for %j = %c0 to %arg1 step %c1 {
+      %0 = addi %arg2, %arg2 : i32
+      memref.store %0, %arg0[%i] : memref<?xi32>
+      memref.store %0, %arg0[%j] : memref<?xi32>
+    }
+  }
+  return
+}
+
+// -----
+
+// CHECK:      builtin.module  {
+// CHECK-NEXT:   builtin.func @for_iter_args(%arg0: index, %arg1: index, %arg2: index) -> f32 {
+// CHECK-NEXT:     %cst = constant 0.000000e+00 : f32
+// CHECK-NEXT:     %0:3 = scf.while (%arg3 = %arg0, %arg4 = %cst, %arg5 = %cst) : (index, f32, f32) -> (index, f32, f32) {
+// CHECK-NEXT:       %1 = cmpi slt, %arg3, %arg1 : index
+// CHECK-NEXT:       scf.condition(%1) %arg3, %arg4, %arg5 : index, f32, f32
+// CHECK-NEXT:     } do {
+// CHECK-NEXT:     ^bb0(%arg3: index, %arg4: f32, %arg5: f32):  // no predecessors
+// CHECK-NEXT:       %1 = addi %arg3, %arg2 : index
+// CHECK-NEXT:       %2 = addf %arg4, %arg5 : f32
+// CHECK-NEXT:       scf.yield %1, %2, %2 : index, f32, f32
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return %0#2 : f32
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+func @for_iter_args(%arg0 : index, %arg1: index, %arg2: index) -> f32 {
+  %s0 = constant 0.0 : f32
+  %result:2 = scf.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%iarg0 = %s0, %iarg1 = %s0) -> (f32, f32) {
+    %sn = addf %iarg0, %iarg1 : f32
+    scf.yield %sn, %sn : f32, f32
+  }
+  return %result#1 : f32
+}
+
+// -----
+
+// CHECK:      builtin.module  {
+// CHECK-NEXT:   builtin.func @exec_region_multiple_yields(%arg0: i32, %arg1: index, %arg2: i32) -> i32 {
+// CHECK-NEXT:     %c0 = constant 0 : index
+// CHECK-NEXT:     %c1 = constant 1 : index
+// CHECK-NEXT:     %0:2 = scf.while (%arg3 = %c0, %arg4 = %arg0) : (index, i32) -> (index, i32) {
+// CHECK-NEXT:       %1 = cmpi slt, %arg3, %arg1 : index
+// CHECK-NEXT:       scf.condition(%1) %arg3, %arg4 : index, i32
+// CHECK-NEXT:     } do {
+// CHECK-NEXT:     ^bb0(%arg3: index, %arg4: i32):  // no predecessors
+// CHECK-NEXT:       %1 = addi %arg3, %c1 : index
+// CHECK-NEXT:       %2 = scf.execute_region -> i32 {
+// CHECK-NEXT:         %3 = cmpi slt, %arg3, %c1 : index
+// CHECK-NEXT:         cond_br %3, ^bb1, ^bb2
+// CHECK-NEXT:       ^bb1:  // pred: ^bb0
+// CHECK-NEXT:         %4 = subi %arg4, %arg0 : i32
+// CHECK-NEXT:         scf.yield %4 : i32
+// CHECK-NEXT:       ^bb2:  // pred: ^bb0
+// CHECK-NEXT:         %5 = muli %arg4, %arg2 : i32
+// CHECK-NEXT:         scf.yield %5 : i32
+// CHECK-NEXT:       }
+// CHECK-NEXT:       scf.yield %1, %2 : index, i32
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return %0#1 : i32
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+func @exec_region_multiple_yields(%arg0: i32, %arg1: index, %arg2: i32) -> i32 {
+  %c1_i32 = constant 1 : i32
+  %c2_i32 = constant 2 : i32
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c5 = constant 5 : index
+  %0 = scf.for %i = %c0 to %arg1 step %c1 iter_args(%iarg0 = %arg0) -> i32 {
+    %2 = scf.execute_region -> i32 {
+      %1 = cmpi slt, %i, %c1 : index
+      cond_br %1, ^bb1, ^bb2
+    ^bb1:
+      %2 = subi %iarg0, %arg0 : i32
+      scf.yield %2 : i32
+    ^bb2:
+      %3 = muli %iarg0, %arg2 : i32
+      scf.yield %3 : i32
+    }
+    scf.yield %2 : i32
+  }
+  return %0 : i32
+}
 No newline at end of file