diff --git a/mlir/include/mlir/Dialect/SCF/Passes.h b/mlir/include/mlir/Dialect/SCF/Passes.h --- a/mlir/include/mlir/Dialect/SCF/Passes.h +++ b/mlir/include/mlir/Dialect/SCF/Passes.h @@ -48,6 +48,9 @@ /// loop range. std::unique_ptr createForLoopRangeFoldingPass(); +// Creates a pass which lowers for loops into while loops. +std::unique_ptr createForToWhileLoopPass(); + //===----------------------------------------------------------------------===// // Registration //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Dialect/SCF/Passes.td b/mlir/include/mlir/Dialect/SCF/Passes.td --- a/mlir/include/mlir/Dialect/SCF/Passes.td +++ b/mlir/include/mlir/Dialect/SCF/Passes.td @@ -62,4 +62,10 @@ 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 diff --git a/mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt b/mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt --- a/mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt +++ b/mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt @@ -1,5 +1,6 @@ add_mlir_dialect_library(MLIRSCFTransforms Bufferize.cpp + ForToWhile.cpp LoopPipelining.cpp LoopRangeFolding.cpp LoopSpecialization.cpp diff --git a/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp b/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp @@ -0,0 +1,111 @@ +//===- 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 { + using OpRewritePattern::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 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 initArgs; + initArgs.push_back(forOp.lowerBound()); + llvm::append_range(initArgs, forOp.initArgs()); + auto whileOp = rewriter.create(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(whileOp.getLoc(), CmpIPredicate::slt, + beforeBlock->getArgument(0), + forOp.upperBound()); + rewriter.create(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( + 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()) { + SmallVector 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 { + void runOnFunction() override { + FuncOp funcOp = getFunction(); + MLIRContext *ctx = funcOp.getContext(); + RewritePatternSet patterns(ctx); + patterns.add(ctx); + (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns)); + } +}; +} // namespace + +std::unique_ptr mlir::createForToWhileLoopPass() { + return std::make_unique(); +} diff --git a/mlir/test/Dialect/SCF/for-loop-to-while-loop.mlir b/mlir/test/Dialect/SCF/for-loop-to-while-loop.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/SCF/for-loop-to-while-loop.mlir @@ -0,0 +1,145 @@ +// 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, %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 +// CHECK-NEXT: scf.yield %1 : index +// CHECK-NEXT: } +// CHECK-NEXT: return +// CHECK-NEXT: } +// CHECK-NEXT: } +func @single_loop(%arg0: memref, %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 + } + return +} + +// ----- + +// CHECK: builtin.module { +// CHECK-NEXT: builtin.func @nested_loop(%arg0: memref, %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 +// CHECK-NEXT: memref.store %4, %arg0[%arg4] : memref +// 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, %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 + memref.store %0, %arg0[%j] : memref + } + } + 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