1 | //===- StackReclaim.cpp -- Insert stacksave/stackrestore in region --------===// |
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 | #include "flang/Optimizer/Builder/FIRBuilder.h" |
10 | #include "flang/Optimizer/Dialect/FIRDialect.h" |
11 | #include "flang/Optimizer/Dialect/FIROps.h" |
12 | #include "flang/Optimizer/Transforms/Passes.h" |
13 | #include "flang/Support/Fortran.h" |
14 | #include "mlir/Dialect/LLVMIR/LLVMDialect.h" |
15 | #include "mlir/IR/Matchers.h" |
16 | #include "mlir/Pass/Pass.h" |
17 | |
18 | namespace fir { |
19 | #define GEN_PASS_DEF_STACKRECLAIM |
20 | #include "flang/Optimizer/Transforms/Passes.h.inc" |
21 | } // namespace fir |
22 | |
23 | using namespace mlir; |
24 | |
25 | namespace { |
26 | |
27 | class StackReclaimPass : public fir::impl::StackReclaimBase<StackReclaimPass> { |
28 | public: |
29 | using StackReclaimBase<StackReclaimPass>::StackReclaimBase; |
30 | |
31 | void runOnOperation() override; |
32 | }; |
33 | } // namespace |
34 | |
35 | void StackReclaimPass::runOnOperation() { |
36 | auto *op = getOperation(); |
37 | fir::FirOpBuilder builder(op, fir::getKindMapping(op)); |
38 | |
39 | op->walk([&](fir::DoLoopOp loopOp) { |
40 | mlir::Location loc = loopOp.getLoc(); |
41 | |
42 | if (!loopOp.getRegion().getOps<fir::AllocaOp>().empty()) { |
43 | builder.setInsertionPointToStart(&loopOp.getRegion().front()); |
44 | mlir::Value sp = builder.genStackSave(loc); |
45 | |
46 | auto *terminator = loopOp.getRegion().back().getTerminator(); |
47 | builder.setInsertionPoint(terminator); |
48 | builder.genStackRestore(loc, sp); |
49 | } |
50 | }); |
51 | } |
52 | |