872 lines · cpp
1//===- StackArrays.cpp ----------------------------------------------------===//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-exception6//7//===----------------------------------------------------------------------===//8 9#include "flang/Optimizer/Builder/FIRBuilder.h"10#include "flang/Optimizer/Builder/LowLevelIntrinsics.h"11#include "flang/Optimizer/Dialect/FIRAttr.h"12#include "flang/Optimizer/Dialect/FIRDialect.h"13#include "flang/Optimizer/Dialect/FIROps.h"14#include "flang/Optimizer/Dialect/FIRType.h"15#include "flang/Optimizer/Dialect/Support/FIRContext.h"16#include "flang/Optimizer/Support/DataLayout.h"17#include "flang/Optimizer/Transforms/Passes.h"18#include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"19#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"20#include "mlir/Analysis/DataFlow/DenseAnalysis.h"21#include "mlir/Analysis/DataFlowFramework.h"22#include "mlir/Dialect/DLTI/DLTI.h"23#include "mlir/Dialect/Func/IR/FuncOps.h"24#include "mlir/Dialect/LLVMIR/LLVMDialect.h"25#include "mlir/Dialect/OpenMP/OpenMPDialect.h"26#include "mlir/IR/Builders.h"27#include "mlir/IR/Diagnostics.h"28#include "mlir/IR/Value.h"29#include "mlir/Interfaces/LoopLikeInterface.h"30#include "mlir/Pass/Pass.h"31#include "mlir/Transforms/GreedyPatternRewriteDriver.h"32#include "mlir/Transforms/Passes.h"33#include "llvm/ADT/DenseMap.h"34#include "llvm/ADT/DenseSet.h"35#include "llvm/ADT/PointerUnion.h"36#include "llvm/Support/Casting.h"37#include "llvm/Support/raw_ostream.h"38#include <optional>39 40namespace fir {41#define GEN_PASS_DEF_STACKARRAYS42#include "flang/Optimizer/Transforms/Passes.h.inc"43} // namespace fir44 45#define DEBUG_TYPE "stack-arrays"46 47static llvm::cl::opt<std::size_t> maxAllocsPerFunc(48 "stack-arrays-max-allocs",49 llvm::cl::desc("The maximum number of heap allocations to consider in one "50 "function before skipping (to save compilation time). Set "51 "to 0 for no limit."),52 llvm::cl::init(1000), llvm::cl::Hidden);53 54static llvm::cl::opt<bool> emitLifetimeMarkers(55 "stack-arrays-lifetime",56 llvm::cl::desc("Add lifetime markers to generated constant size allocas"),57 llvm::cl::init(false), llvm::cl::Hidden);58 59namespace {60 61/// The state of an SSA value at each program point62enum class AllocationState {63 /// This means that the allocation state of a variable cannot be determined64 /// at this program point, e.g. because one route through a conditional freed65 /// the variable and the other route didn't.66 /// This asserts a known-unknown: different from the unknown-unknown of having67 /// no AllocationState stored for a particular SSA value68 Unknown,69 /// Means this SSA value was allocated on the heap in this function and has70 /// now been freed71 Freed,72 /// Means this SSA value was allocated on the heap in this function and is a73 /// candidate for moving to the stack74 Allocated,75};76 77/// Stores where an alloca should be inserted. If the PointerUnion is an78/// Operation the alloca should be inserted /after/ the operation. If it is a79/// block, the alloca can be placed anywhere in that block.80class InsertionPoint {81 llvm::PointerUnion<mlir::Operation *, mlir::Block *> location;82 bool saveRestoreStack;83 84 /// Get contained pointer type or nullptr85 template <class T>86 T *tryGetPtr() const {87 // Use llvm::dyn_cast_if_present because location may be null here.88 if (T *ptr = llvm::dyn_cast_if_present<T *>(location))89 return ptr;90 return nullptr;91 }92 93public:94 template <class T>95 InsertionPoint(T *ptr, bool saveRestoreStack = false)96 : location(ptr), saveRestoreStack{saveRestoreStack} {}97 InsertionPoint(std::nullptr_t null)98 : location(null), saveRestoreStack{false} {}99 100 /// Get contained operation, or nullptr101 mlir::Operation *tryGetOperation() const {102 return tryGetPtr<mlir::Operation>();103 }104 105 /// Get contained block, or nullptr106 mlir::Block *tryGetBlock() const { return tryGetPtr<mlir::Block>(); }107 108 /// Get whether the stack should be saved/restored. If yes, an llvm.stacksave109 /// intrinsic should be added before the alloca, and an llvm.stackrestore110 /// intrinsic should be added where the freemem is111 bool shouldSaveRestoreStack() const { return saveRestoreStack; }112 113 operator bool() const { return tryGetOperation() || tryGetBlock(); }114 115 bool operator==(const InsertionPoint &rhs) const {116 return (location == rhs.location) &&117 (saveRestoreStack == rhs.saveRestoreStack);118 }119 120 bool operator!=(const InsertionPoint &rhs) const { return !(*this == rhs); }121};122 123/// Maps SSA values to their AllocationState at a particular program point.124/// Also caches the insertion points for the new alloca operations125class LatticePoint : public mlir::dataflow::AbstractDenseLattice {126 // Maps all values we are interested in to states127 llvm::SmallDenseMap<mlir::Value, AllocationState, 1> stateMap;128 129public:130 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LatticePoint)131 using AbstractDenseLattice::AbstractDenseLattice;132 133 bool operator==(const LatticePoint &rhs) const {134 return stateMap == rhs.stateMap;135 }136 137 /// Join the lattice accross control-flow edges138 mlir::ChangeResult join(const AbstractDenseLattice &lattice) override;139 140 void print(llvm::raw_ostream &os) const override;141 142 /// Clear all modifications143 mlir::ChangeResult reset();144 145 /// Set the state of an SSA value146 mlir::ChangeResult set(mlir::Value value, AllocationState state);147 148 /// Get fir.allocmem ops which were allocated in this function and always149 /// freed before the function returns, plus whre to insert replacement150 /// fir.alloca ops151 void appendFreedValues(llvm::DenseSet<mlir::Value> &out) const;152 153 std::optional<AllocationState> get(mlir::Value val) const;154};155 156class AllocationAnalysis157 : public mlir::dataflow::DenseForwardDataFlowAnalysis<LatticePoint> {158public:159 using DenseForwardDataFlowAnalysis::DenseForwardDataFlowAnalysis;160 161 mlir::LogicalResult visitOperation(mlir::Operation *op,162 const LatticePoint &before,163 LatticePoint *after) override;164 165 /// At an entry point, the last modifications of all memory resources are166 /// yet to be determined167 void setToEntryState(LatticePoint *lattice) override;168 169protected:170 /// Visit control flow operations and decide whether to call visitOperation171 /// to apply the transfer function172 mlir::LogicalResult processOperation(mlir::Operation *op) override;173};174 175/// Drives analysis to find candidate fir.allocmem operations which could be176/// moved to the stack. Intended to be used with mlir::Pass::getAnalysis177class StackArraysAnalysisWrapper {178public:179 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(StackArraysAnalysisWrapper)180 181 // Maps fir.allocmem -> place to insert alloca182 using AllocMemMap = llvm::DenseMap<mlir::Operation *, InsertionPoint>;183 184 StackArraysAnalysisWrapper(mlir::Operation *op) {}185 186 // returns nullptr if analysis failed187 const AllocMemMap *getCandidateOps(mlir::Operation *func);188 189private:190 llvm::DenseMap<mlir::Operation *, AllocMemMap> funcMaps;191 192 llvm::LogicalResult analyseFunction(mlir::Operation *func);193};194 195/// Converts a fir.allocmem to a fir.alloca196class AllocMemConversion : public mlir::OpRewritePattern<fir::AllocMemOp> {197public:198 explicit AllocMemConversion(199 mlir::MLIRContext *ctx,200 const StackArraysAnalysisWrapper::AllocMemMap &candidateOps,201 std::optional<mlir::DataLayout> &dl,202 std::optional<fir::KindMapping> &kindMap)203 : OpRewritePattern(ctx), candidateOps{candidateOps}, dl{dl},204 kindMap{kindMap} {}205 206 llvm::LogicalResult207 matchAndRewrite(fir::AllocMemOp allocmem,208 mlir::PatternRewriter &rewriter) const override;209 210 /// Determine where to insert the alloca operation. The returned value should211 /// be checked to see if it is inside a loop212 static InsertionPoint213 findAllocaInsertionPoint(fir::AllocMemOp &oldAlloc,214 const llvm::SmallVector<mlir::Operation *> &freeOps);215 216private:217 /// Handle to the DFA (already run)218 const StackArraysAnalysisWrapper::AllocMemMap &candidateOps;219 220 const std::optional<mlir::DataLayout> &dl;221 const std::optional<fir::KindMapping> &kindMap;222 223 /// If we failed to find an insertion point not inside a loop, see if it would224 /// be safe to use an llvm.stacksave/llvm.stackrestore inside the loop225 static InsertionPoint findAllocaLoopInsertionPoint(226 fir::AllocMemOp &oldAlloc,227 const llvm::SmallVector<mlir::Operation *> &freeOps);228 229 /// Returns the alloca if it was successfully inserted, otherwise {}230 std::optional<fir::AllocaOp>231 insertAlloca(fir::AllocMemOp &oldAlloc,232 mlir::PatternRewriter &rewriter) const;233 234 /// Inserts a stacksave before oldAlloc and a stackrestore after each freemem235 void insertStackSaveRestore(fir::AllocMemOp oldAlloc,236 mlir::PatternRewriter &rewriter) const;237 /// Emit lifetime markers for newAlloc between oldAlloc and each freemem.238 /// If the allocation is dynamic, no life markers are emitted.239 void insertLifetimeMarkers(fir::AllocMemOp oldAlloc, fir::AllocaOp newAlloc,240 mlir::PatternRewriter &rewriter) const;241};242 243class StackArraysPass : public fir::impl::StackArraysBase<StackArraysPass> {244public:245 StackArraysPass() = default;246 StackArraysPass(const StackArraysPass &pass);247 248 llvm::StringRef getDescription() const override;249 250 void runOnOperation() override;251 252private:253 Statistic runCount{this, "stackArraysRunCount",254 "Number of heap allocations moved to the stack"};255};256 257} // namespace258 259static void print(llvm::raw_ostream &os, AllocationState state) {260 switch (state) {261 case AllocationState::Unknown:262 os << "Unknown";263 break;264 case AllocationState::Freed:265 os << "Freed";266 break;267 case AllocationState::Allocated:268 os << "Allocated";269 break;270 }271}272 273/// Join two AllocationStates for the same value coming from different CFG274/// blocks275static AllocationState join(AllocationState lhs, AllocationState rhs) {276 // | Allocated | Freed | Unknown277 // ========= | ========= | ========= | =========278 // Allocated | Allocated | Unknown | Unknown279 // Freed | Unknown | Freed | Unknown280 // Unknown | Unknown | Unknown | Unknown281 if (lhs == rhs)282 return lhs;283 return AllocationState::Unknown;284}285 286mlir::ChangeResult LatticePoint::join(const AbstractDenseLattice &lattice) {287 const auto &rhs = static_cast<const LatticePoint &>(lattice);288 mlir::ChangeResult changed = mlir::ChangeResult::NoChange;289 290 // add everything from rhs to map, handling cases where values are in both291 for (const auto &[value, rhsState] : rhs.stateMap) {292 auto it = stateMap.find(value);293 if (it != stateMap.end()) {294 // value is present in both maps295 AllocationState myState = it->second;296 AllocationState newState = ::join(myState, rhsState);297 if (newState != myState) {298 changed = mlir::ChangeResult::Change;299 it->getSecond() = newState;300 }301 } else {302 // value not present in current map: add it303 stateMap.insert({value, rhsState});304 changed = mlir::ChangeResult::Change;305 }306 }307 308 return changed;309}310 311void LatticePoint::print(llvm::raw_ostream &os) const {312 for (const auto &[value, state] : stateMap) {313 os << "\n * " << value << ": ";314 ::print(os, state);315 }316}317 318mlir::ChangeResult LatticePoint::reset() {319 if (stateMap.empty())320 return mlir::ChangeResult::NoChange;321 stateMap.clear();322 return mlir::ChangeResult::Change;323}324 325mlir::ChangeResult LatticePoint::set(mlir::Value value, AllocationState state) {326 if (stateMap.count(value)) {327 // already in map328 AllocationState &oldState = stateMap[value];329 if (oldState != state) {330 stateMap[value] = state;331 return mlir::ChangeResult::Change;332 }333 return mlir::ChangeResult::NoChange;334 }335 stateMap.insert({value, state});336 return mlir::ChangeResult::Change;337}338 339/// Get values which were allocated in this function and always freed before340/// the function returns341void LatticePoint::appendFreedValues(llvm::DenseSet<mlir::Value> &out) const {342 for (auto &[value, state] : stateMap) {343 if (state == AllocationState::Freed)344 out.insert(value);345 }346}347 348std::optional<AllocationState> LatticePoint::get(mlir::Value val) const {349 auto it = stateMap.find(val);350 if (it == stateMap.end())351 return {};352 return it->second;353}354 355static mlir::Value lookThroughDeclaresAndConverts(mlir::Value value) {356 while (mlir::Operation *op = value.getDefiningOp()) {357 if (auto declareOp = llvm::dyn_cast<fir::DeclareOp>(op))358 value = declareOp.getMemref();359 else if (auto convertOp = llvm::dyn_cast<fir::ConvertOp>(op))360 value = convertOp->getOperand(0);361 else362 return value;363 }364 return value;365}366 367mlir::LogicalResult AllocationAnalysis::visitOperation(368 mlir::Operation *op, const LatticePoint &before, LatticePoint *after) {369 LLVM_DEBUG(llvm::dbgs() << "StackArrays: Visiting operation: " << *op370 << "\n");371 LLVM_DEBUG(llvm::dbgs() << "--Lattice in: " << before << "\n");372 373 // propagate before -> after374 mlir::ChangeResult changed = after->join(before);375 376 if (auto allocmem = mlir::dyn_cast<fir::AllocMemOp>(op)) {377 assert(op->getNumResults() == 1 && "fir.allocmem has one result");378 auto attr = op->getAttrOfType<fir::MustBeHeapAttr>(379 fir::MustBeHeapAttr::getAttrName());380 if (attr && attr.getValue()) {381 LLVM_DEBUG(llvm::dbgs() << "--Found fir.must_be_heap: skipping\n");382 // skip allocation marked not to be moved383 return mlir::success();384 }385 386 auto retTy = allocmem.getAllocatedType();387 if (!mlir::isa<fir::SequenceType>(retTy)) {388 LLVM_DEBUG(llvm::dbgs()389 << "--Allocation is not for an array: skipping\n");390 return mlir::success();391 }392 393 mlir::Value result = op->getResult(0);394 changed |= after->set(result, AllocationState::Allocated);395 } else if (mlir::isa<fir::FreeMemOp>(op)) {396 assert(op->getNumOperands() == 1 && "fir.freemem has one operand");397 mlir::Value operand = op->getOperand(0);398 399 // Note: StackArrays is scheduled in the pass pipeline after lowering hlfir400 // to fir. Therefore, we only need to handle `fir::DeclareOp`s. Also look401 // past converts in case the pointer was changed between different pointer402 // types.403 operand = lookThroughDeclaresAndConverts(operand);404 405 std::optional<AllocationState> operandState = before.get(operand);406 if (operandState && *operandState == AllocationState::Allocated) {407 // don't tag things not allocated in this function as freed, so that we408 // don't think they are candidates for moving to the stack409 changed |= after->set(operand, AllocationState::Freed);410 }411 } else if (mlir::isa<fir::ResultOp>(op)) {412 mlir::Operation *parent = op->getParentOp();413 LatticePoint *parentLattice = getLattice(getProgramPointAfter(parent));414 assert(parentLattice);415 mlir::ChangeResult parentChanged = parentLattice->join(*after);416 propagateIfChanged(parentLattice, parentChanged);417 }418 419 // we pass lattices straight through fir.call because called functions should420 // not deallocate flang-generated array temporaries421 422 LLVM_DEBUG(llvm::dbgs() << "--Lattice out: " << *after << "\n");423 propagateIfChanged(after, changed);424 return mlir::success();425}426 427void AllocationAnalysis::setToEntryState(LatticePoint *lattice) {428 propagateIfChanged(lattice, lattice->reset());429}430 431/// Mostly a copy of AbstractDenseLattice::processOperation - the difference432/// being that call operations are passed through to the transfer function433mlir::LogicalResult AllocationAnalysis::processOperation(mlir::Operation *op) {434 mlir::ProgramPoint *point = getProgramPointAfter(op);435 // If the containing block is not executable, bail out.436 if (op->getBlock() != nullptr &&437 !getOrCreateFor<mlir::dataflow::Executable>(438 point, getProgramPointBefore(op->getBlock()))439 ->isLive())440 return mlir::success();441 442 // Get the dense lattice to update443 mlir::dataflow::AbstractDenseLattice *after = getLattice(point);444 445 // If this op implements region control-flow, then control-flow dictates its446 // transfer function.447 if (auto branch = mlir::dyn_cast<mlir::RegionBranchOpInterface>(op)) {448 visitRegionBranchOperation(point, branch, after);449 return mlir::success();450 }451 452 // pass call operations through to the transfer function453 454 // Get the dense state before the execution of the op.455 const mlir::dataflow::AbstractDenseLattice *before =456 getLatticeFor(point, getProgramPointBefore(op));457 458 /// Invoke the operation transfer function459 return visitOperationImpl(op, *before, after);460}461 462llvm::LogicalResult463StackArraysAnalysisWrapper::analyseFunction(mlir::Operation *func) {464 assert(mlir::isa<mlir::func::FuncOp>(func));465 size_t nAllocs = 0;466 func->walk([&nAllocs](fir::AllocMemOp) { nAllocs++; });467 // don't bother with the analysis if there are no heap allocations468 if (nAllocs == 0)469 return mlir::success();470 if ((maxAllocsPerFunc != 0) && (nAllocs > maxAllocsPerFunc)) {471 LLVM_DEBUG(llvm::dbgs() << "Skipping stack arrays for function with "472 << nAllocs << " heap allocations");473 return mlir::success();474 }475 476 mlir::DataFlowSolver solver;477 // constant propagation is required for dead code analysis, dead code analysis478 // is required to mark blocks live (required for mlir dense dfa)479 solver.load<mlir::dataflow::SparseConstantPropagation>();480 solver.load<mlir::dataflow::DeadCodeAnalysis>();481 482 auto [it, inserted] = funcMaps.try_emplace(func);483 AllocMemMap &candidateOps = it->second;484 485 solver.load<AllocationAnalysis>();486 if (failed(solver.initializeAndRun(func))) {487 llvm::errs() << "DataFlowSolver failed!";488 return mlir::failure();489 }490 491 LatticePoint point{solver.getProgramPointAfter(func)};492 auto joinOperationLattice = [&](mlir::Operation *op) {493 const LatticePoint *lattice =494 solver.lookupState<LatticePoint>(solver.getProgramPointAfter(op));495 // there will be no lattice for an unreachable block496 if (lattice)497 (void)point.join(*lattice);498 };499 500 func->walk([&](mlir::func::ReturnOp child) { joinOperationLattice(child); });501 func->walk([&](fir::UnreachableOp child) { joinOperationLattice(child); });502 func->walk(503 [&](mlir::omp::TerminatorOp child) { joinOperationLattice(child); });504 func->walk([&](mlir::omp::YieldOp child) { joinOperationLattice(child); });505 506 llvm::DenseSet<mlir::Value> freedValues;507 point.appendFreedValues(freedValues);508 509 // Find all fir.freemem operations corresponding to fir.allocmem510 // in freedValues. It is best to find the association going back511 // from fir.freemem to fir.allocmem through the def-use chains,512 // so that we can use lookThroughDeclaresAndConverts same way513 // the AllocationAnalysis is handling them.514 llvm::DenseMap<mlir::Operation *, llvm::SmallVector<mlir::Operation *>>515 allocToFreeMemMap;516 func->walk([&](fir::FreeMemOp freeOp) {517 mlir::Value memref = lookThroughDeclaresAndConverts(freeOp.getHeapref());518 if (!freedValues.count(memref))519 return;520 521 auto allocMem = memref.getDefiningOp<fir::AllocMemOp>();522 allocToFreeMemMap[allocMem].push_back(freeOp);523 });524 525 // We only replace allocations which are definately freed on all routes526 // through the function because otherwise the allocation may have an intende527 // lifetime longer than the current stack frame (e.g. a heap allocation which528 // is then freed by another function).529 for (mlir::Value freedValue : freedValues) {530 fir::AllocMemOp allocmem = freedValue.getDefiningOp<fir::AllocMemOp>();531 InsertionPoint insertionPoint =532 AllocMemConversion::findAllocaInsertionPoint(533 allocmem, allocToFreeMemMap[allocmem]);534 if (insertionPoint)535 candidateOps.insert({allocmem, insertionPoint});536 }537 538 LLVM_DEBUG(for (auto [allocMemOp, _]539 : candidateOps) {540 llvm::dbgs() << "StackArrays: Found candidate op: " << *allocMemOp << '\n';541 });542 return mlir::success();543}544 545const StackArraysAnalysisWrapper::AllocMemMap *546StackArraysAnalysisWrapper::getCandidateOps(mlir::Operation *func) {547 if (!funcMaps.contains(func))548 if (mlir::failed(analyseFunction(func)))549 return nullptr;550 return &funcMaps[func];551}552 553/// Restore the old allocation type exected by existing code554static mlir::Value convertAllocationType(mlir::PatternRewriter &rewriter,555 const mlir::Location &loc,556 mlir::Value heap, mlir::Value stack) {557 mlir::Type heapTy = heap.getType();558 mlir::Type stackTy = stack.getType();559 560 if (heapTy == stackTy)561 return stack;562 563 fir::HeapType firHeapTy = mlir::cast<fir::HeapType>(heapTy);564 [[maybe_unused]] fir::ReferenceType firRefTy =565 mlir::cast<fir::ReferenceType>(stackTy);566 assert(firHeapTy.getElementType() == firRefTy.getElementType() &&567 "Allocations must have the same type");568 569 auto insertionPoint = rewriter.saveInsertionPoint();570 rewriter.setInsertionPointAfter(stack.getDefiningOp());571 mlir::Value conv =572 fir::ConvertOp::create(rewriter, loc, firHeapTy, stack).getResult();573 rewriter.restoreInsertionPoint(insertionPoint);574 return conv;575}576 577llvm::LogicalResult578AllocMemConversion::matchAndRewrite(fir::AllocMemOp allocmem,579 mlir::PatternRewriter &rewriter) const {580 auto oldInsertionPt = rewriter.saveInsertionPoint();581 // add alloca operation582 std::optional<fir::AllocaOp> alloca = insertAlloca(allocmem, rewriter);583 rewriter.restoreInsertionPoint(oldInsertionPt);584 if (!alloca)585 return mlir::failure();586 587 // remove freemem operations588 llvm::SmallVector<mlir::Operation *> erases;589 mlir::Operation *parent = allocmem->getParentOp();590 // TODO: this shouldn't need to be re-calculated for every allocmem591 parent->walk([&](fir::FreeMemOp freeOp) {592 if (lookThroughDeclaresAndConverts(freeOp->getOperand(0)) == allocmem)593 erases.push_back(freeOp);594 });595 596 // now we are done iterating the users, it is safe to mutate them597 for (mlir::Operation *erase : erases)598 rewriter.eraseOp(erase);599 600 // replace references to heap allocation with references to stack allocation601 mlir::Value newValue = convertAllocationType(602 rewriter, allocmem.getLoc(), allocmem.getResult(), alloca->getResult());603 rewriter.replaceOp(allocmem, newValue);604 605 return mlir::success();606}607 608static bool isInLoop(mlir::Block *block) {609 return mlir::LoopLikeOpInterface::blockIsInLoop(block);610}611 612static bool isInLoop(mlir::Operation *op) {613 return isInLoop(op->getBlock()) ||614 op->getParentOfType<mlir::LoopLikeOpInterface>();615}616 617InsertionPoint AllocMemConversion::findAllocaInsertionPoint(618 fir::AllocMemOp &oldAlloc,619 const llvm::SmallVector<mlir::Operation *> &freeOps) {620 // Ideally the alloca should be inserted at the end of the function entry621 // block so that we do not allocate stack space in a loop. However,622 // the operands to the alloca may not be available that early, so insert it623 // after the last operand becomes available624 // If the old allocmem op was in an openmp region then it should not be moved625 // outside of that626 LLVM_DEBUG(llvm::dbgs() << "StackArrays: findAllocaInsertionPoint: "627 << oldAlloc << "\n");628 629 // check that an Operation or Block we are about to return is not in a loop630 auto checkReturn = [&](auto *point) -> InsertionPoint {631 if (isInLoop(point)) {632 mlir::Operation *oldAllocOp = oldAlloc.getOperation();633 if (isInLoop(oldAllocOp)) {634 // where we want to put it is in a loop, and even the old location is in635 // a loop. Give up.636 return findAllocaLoopInsertionPoint(oldAlloc, freeOps);637 }638 return {oldAllocOp};639 }640 return {point};641 };642 643 auto oldOmpRegion =644 oldAlloc->getParentOfType<mlir::omp::OutlineableOpenMPOpInterface>();645 646 // Find when the last operand value becomes available647 mlir::Block *operandsBlock = nullptr;648 mlir::Operation *lastOperand = nullptr;649 for (mlir::Value operand : oldAlloc.getOperands()) {650 LLVM_DEBUG(llvm::dbgs() << "--considering operand " << operand << "\n");651 mlir::Operation *op = operand.getDefiningOp();652 if (!op)653 return checkReturn(oldAlloc.getOperation());654 if (!operandsBlock)655 operandsBlock = op->getBlock();656 else if (operandsBlock != op->getBlock()) {657 LLVM_DEBUG(llvm::dbgs()658 << "----operand declared in a different block!\n");659 // Operation::isBeforeInBlock requires the operations to be in the same660 // block. The best we can do is the location of the allocmem.661 return checkReturn(oldAlloc.getOperation());662 }663 if (!lastOperand || lastOperand->isBeforeInBlock(op))664 lastOperand = op;665 }666 667 if (lastOperand) {668 // there were value operands to the allocmem so insert after the last one669 LLVM_DEBUG(llvm::dbgs()670 << "--Placing after last operand: " << *lastOperand << "\n");671 // check we aren't moving out of an omp region672 auto lastOpOmpRegion =673 lastOperand->getParentOfType<mlir::omp::OutlineableOpenMPOpInterface>();674 if (lastOpOmpRegion == oldOmpRegion)675 return checkReturn(lastOperand);676 // Presumably this happened because the operands became ready before the677 // start of this openmp region. (lastOpOmpRegion != oldOmpRegion) should678 // imply that oldOmpRegion comes after lastOpOmpRegion.679 return checkReturn(oldOmpRegion.getAllocaBlock());680 }681 682 // There were no value operands to the allocmem so we are safe to insert it683 // as early as we want684 685 // handle openmp case686 if (oldOmpRegion)687 return checkReturn(oldOmpRegion.getAllocaBlock());688 689 // fall back to the function entry block690 mlir::func::FuncOp func = oldAlloc->getParentOfType<mlir::func::FuncOp>();691 assert(func && "This analysis is run on func.func");692 mlir::Block &entryBlock = func.getBlocks().front();693 LLVM_DEBUG(llvm::dbgs() << "--Placing at the start of func entry block\n");694 return checkReturn(&entryBlock);695}696 697InsertionPoint AllocMemConversion::findAllocaLoopInsertionPoint(698 fir::AllocMemOp &oldAlloc,699 const llvm::SmallVector<mlir::Operation *> &freeOps) {700 mlir::Operation *oldAllocOp = oldAlloc;701 // This is only called as a last resort. We should try to insert at the702 // location of the old allocation, which is inside of a loop, using703 // llvm.stacksave/llvm.stackrestore704 705 assert(freeOps.size() && "DFA should only return freed memory");706 707 // Don't attempt to reason about a stacksave/stackrestore between different708 // blocks709 for (mlir::Operation *free : freeOps)710 if (free->getBlock() != oldAllocOp->getBlock())711 return {nullptr};712 713 // Check that there aren't any other stack allocations in between the714 // stack save and stack restore715 // note: for flang generated temporaries there should only be one free op716 for (mlir::Operation *free : freeOps) {717 for (mlir::Operation *op = oldAlloc; op && op != free;718 op = op->getNextNode()) {719 if (mlir::isa<fir::AllocaOp>(op))720 return {nullptr};721 }722 }723 724 return InsertionPoint{oldAllocOp, /*shouldStackSaveRestore=*/true};725}726 727std::optional<fir::AllocaOp>728AllocMemConversion::insertAlloca(fir::AllocMemOp &oldAlloc,729 mlir::PatternRewriter &rewriter) const {730 auto it = candidateOps.find(oldAlloc.getOperation());731 if (it == candidateOps.end())732 return {};733 InsertionPoint insertionPoint = it->second;734 if (!insertionPoint)735 return {};736 737 if (insertionPoint.shouldSaveRestoreStack())738 insertStackSaveRestore(oldAlloc, rewriter);739 740 mlir::Location loc = oldAlloc.getLoc();741 mlir::Type varTy = oldAlloc.getInType();742 if (mlir::Operation *op = insertionPoint.tryGetOperation()) {743 rewriter.setInsertionPointAfter(op);744 } else {745 mlir::Block *block = insertionPoint.tryGetBlock();746 assert(block && "There must be a valid insertion point");747 rewriter.setInsertionPointToStart(block);748 }749 750 auto unpackName = [](std::optional<llvm::StringRef> opt) -> llvm::StringRef {751 if (opt)752 return *opt;753 return {};754 };755 756 llvm::StringRef uniqName = unpackName(oldAlloc.getUniqName());757 llvm::StringRef bindcName = unpackName(oldAlloc.getBindcName());758 auto alloca =759 fir::AllocaOp::create(rewriter, loc, varTy, uniqName, bindcName,760 oldAlloc.getTypeparams(), oldAlloc.getShape());761 if (emitLifetimeMarkers)762 insertLifetimeMarkers(oldAlloc, alloca, rewriter);763 764 return alloca;765}766 767static void768visitFreeMemOp(fir::AllocMemOp oldAlloc,769 const std::function<void(mlir::Operation *)> &callBack) {770 for (mlir::Operation *user : oldAlloc->getUsers()) {771 if (auto declareOp = mlir::dyn_cast_if_present<fir::DeclareOp>(user)) {772 for (mlir::Operation *user : declareOp->getUsers()) {773 if (mlir::isa<fir::FreeMemOp>(user))774 callBack(user);775 }776 }777 778 if (mlir::isa<fir::FreeMemOp>(user))779 callBack(user);780 }781}782 783void AllocMemConversion::insertStackSaveRestore(784 fir::AllocMemOp oldAlloc, mlir::PatternRewriter &rewriter) const {785 mlir::OpBuilder::InsertionGuard insertGuard(rewriter);786 auto mod = oldAlloc->getParentOfType<mlir::ModuleOp>();787 fir::FirOpBuilder builder{rewriter, mod};788 789 builder.setInsertionPoint(oldAlloc);790 mlir::Value sp = builder.genStackSave(oldAlloc.getLoc());791 792 auto createStackRestoreCall = [&](mlir::Operation *user) {793 builder.setInsertionPoint(user);794 builder.genStackRestore(user->getLoc(), sp);795 };796 visitFreeMemOp(oldAlloc, createStackRestoreCall);797}798 799void AllocMemConversion::insertLifetimeMarkers(800 fir::AllocMemOp oldAlloc, fir::AllocaOp newAlloc,801 mlir::PatternRewriter &rewriter) const {802 if (!dl || !kindMap)803 return;804 llvm::StringRef attrName = fir::getHasLifetimeMarkerAttrName();805 // Do not add lifetime markers if the alloca already has any.806 if (newAlloc->hasAttr(attrName))807 return;808 if (std::optional<int64_t> size =809 fir::getAllocaByteSize(newAlloc, *dl, *kindMap)) {810 mlir::OpBuilder::InsertionGuard insertGuard(rewriter);811 rewriter.setInsertionPoint(oldAlloc);812 mlir::Value ptr = fir::factory::genLifetimeStart(813 rewriter, newAlloc.getLoc(), newAlloc, &*dl);814 visitFreeMemOp(oldAlloc, [&](mlir::Operation *op) {815 rewriter.setInsertionPoint(op);816 fir::factory::genLifetimeEnd(rewriter, op->getLoc(), ptr);817 });818 newAlloc->setAttr(attrName, rewriter.getUnitAttr());819 }820}821 822StackArraysPass::StackArraysPass(const StackArraysPass &pass)823 : fir::impl::StackArraysBase<StackArraysPass>(pass) {}824 825llvm::StringRef StackArraysPass::getDescription() const {826 return "Move heap allocated array temporaries to the stack";827}828 829void StackArraysPass::runOnOperation() {830 mlir::func::FuncOp func = getOperation();831 832 auto &analysis = getAnalysis<StackArraysAnalysisWrapper>();833 const StackArraysAnalysisWrapper::AllocMemMap *candidateOps =834 analysis.getCandidateOps(func);835 if (!candidateOps) {836 signalPassFailure();837 return;838 }839 840 if (candidateOps->empty())841 return;842 runCount += candidateOps->size();843 844 llvm::SmallVector<mlir::Operation *> opsToConvert;845 opsToConvert.reserve(candidateOps->size());846 for (auto [op, _] : *candidateOps)847 opsToConvert.push_back(op);848 849 mlir::MLIRContext &context = getContext();850 mlir::RewritePatternSet patterns(&context);851 mlir::GreedyRewriteConfig config;852 // prevent the pattern driver form merging blocks853 config.setRegionSimplificationLevel(854 mlir::GreedySimplifyRegionLevel::Disabled);855 856 auto module = func->getParentOfType<mlir::ModuleOp>();857 std::optional<mlir::DataLayout> dl =858 module ? fir::support::getOrSetMLIRDataLayout(859 module, /*allowDefaultLayout=*/false)860 : std::nullopt;861 std::optional<fir::KindMapping> kindMap;862 if (module)863 kindMap = fir::getKindMapping(module);864 865 patterns.insert<AllocMemConversion>(&context, *candidateOps, dl, kindMap);866 if (mlir::failed(mlir::applyOpPatternsGreedily(867 opsToConvert, std::move(patterns), config))) {868 mlir::emitError(func->getLoc(), "error in stack arrays optimization\n");869 signalPassFailure();870 }871}872