xref: /llvm-project/flang/lib/Optimizer/Transforms/StackReclaim.cpp (revision 5aaf384b1614fcef5504d0b16d3e5063f72943c1)
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/Common/Fortran.h"
10 #include "flang/Optimizer/Builder/FIRBuilder.h"
11 #include "flang/Optimizer/Dialect/FIRDialect.h"
12 #include "flang/Optimizer/Dialect/FIROps.h"
13 #include "flang/Optimizer/Transforms/Passes.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