flang/lib/Optimizer/Transforms/StackReclaim.cpp

   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 }