1 | //===- ForallToFor.cpp - scf.forall to scf.for loop conversion ------------===// |
2 | // |
3 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
4 | // See https://llvm.org/LICENSE.txt for license information. |
5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | // |
9 | // Transforms SCF.ForallOp's into SCF.ForOp's. |
10 | // |
11 | //===----------------------------------------------------------------------===// |
12 | |
13 | #include "mlir/Dialect/SCF/Transforms/Passes.h" |
14 | |
15 | #include "mlir/Dialect/SCF/IR/SCF.h" |
16 | #include "mlir/Dialect/SCF/Transforms/Transforms.h" |
17 | #include "mlir/IR/PatternMatch.h" |
18 | |
19 | namespace mlir { |
20 | #define GEN_PASS_DEF_SCFFORALLTOFORLOOP |
21 | #include "mlir/Dialect/SCF/Transforms/Passes.h.inc" |
22 | } // namespace mlir |
23 | |
24 | using namespace llvm; |
25 | using namespace mlir; |
26 | using scf::ForallOp; |
27 | using scf::ForOp; |
28 | using scf::LoopNest; |
29 | |
30 | LogicalResult |
31 | mlir::scf::forallToForLoop(RewriterBase &rewriter, scf::ForallOp forallOp, |
32 | SmallVectorImpl<Operation *> *results) { |
33 | OpBuilder::InsertionGuard guard(rewriter); |
34 | rewriter.setInsertionPoint(forallOp); |
35 | |
36 | Location loc = forallOp.getLoc(); |
37 | SmallVector<Value> lbs = getValueOrCreateConstantIndexOp( |
38 | rewriter, loc, forallOp.getMixedLowerBound()); |
39 | SmallVector<Value> ubs = getValueOrCreateConstantIndexOp( |
40 | rewriter, loc, forallOp.getMixedUpperBound()); |
41 | SmallVector<Value> steps = |
42 | getValueOrCreateConstantIndexOp(rewriter, loc, forallOp.getMixedStep()); |
43 | LoopNest loopNest = scf::buildLoopNest(rewriter, loc, lbs, ubs, steps); |
44 | |
45 | SmallVector<Value> ivs = llvm::map_to_vector( |
46 | loopNest.loops, [](scf::ForOp loop) { return loop.getInductionVar(); }); |
47 | |
48 | Block *innermostBlock = loopNest.loops.back().getBody(); |
49 | rewriter.eraseOp(op: forallOp.getBody()->getTerminator()); |
50 | rewriter.inlineBlockBefore(forallOp.getBody(), innermostBlock, |
51 | innermostBlock->getTerminator()->getIterator(), |
52 | ivs); |
53 | rewriter.eraseOp(op: forallOp); |
54 | |
55 | if (results) { |
56 | llvm::move(loopNest.loops, std::back_inserter(*results)); |
57 | } |
58 | |
59 | return success(); |
60 | } |
61 | |
62 | namespace { |
63 | struct ForallToForLoop : public impl::SCFForallToForLoopBase<ForallToForLoop> { |
64 | void runOnOperation() override { |
65 | Operation *parentOp = getOperation(); |
66 | IRRewriter rewriter(parentOp->getContext()); |
67 | |
68 | parentOp->walk([&](scf::ForallOp forallOp) { |
69 | if (failed(scf::forallToForLoop(rewriter, forallOp: forallOp))) { |
70 | return signalPassFailure(); |
71 | } |
72 | }); |
73 | } |
74 | }; |
75 | } // namespace |
76 | |
77 | std::unique_ptr<Pass> mlir::createForallToForLoopPass() { |
78 | return std::make_unique<ForallToForLoop>(); |
79 | } |
80 | |