153 lines · cpp
1//===- MemoryAllocation.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/Dialect/FIRDialect.h"10#include "flang/Optimizer/Dialect/FIROps.h"11#include "flang/Optimizer/Dialect/FIRType.h"12#include "flang/Optimizer/Transforms/MemoryUtils.h"13#include "flang/Optimizer/Transforms/Passes.h"14#include "mlir/Dialect/Func/IR/FuncOps.h"15#include "mlir/IR/Diagnostics.h"16#include "mlir/Pass/Pass.h"17#include "mlir/Transforms/DialectConversion.h"18#include "mlir/Transforms/Passes.h"19#include "llvm/ADT/TypeSwitch.h"20 21namespace fir {22#define GEN_PASS_DEF_MEMORYALLOCATIONOPT23#include "flang/Optimizer/Transforms/Passes.h.inc"24} // namespace fir25 26#define DEBUG_TYPE "flang-memory-allocation-opt"27 28// Number of elements in an array does not determine where it is allocated.29static constexpr std::size_t unlimitedArraySize = ~static_cast<std::size_t>(0);30 31/// Return `true` if this allocation is to remain on the stack (`fir.alloca`).32/// Otherwise the allocation should be moved to the heap (`fir.allocmem`).33static inline bool34keepStackAllocation(fir::AllocaOp alloca,35 const fir::MemoryAllocationOptOptions &options) {36 // Move all arrays and character with runtime determined size to the heap.37 if (options.dynamicArrayOnHeap && alloca.isDynamic())38 return false;39 // TODO: use data layout to reason in terms of byte size to cover all "big"40 // entities, which may be scalar derived types.41 if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(alloca.getInType())) {42 if (!fir::hasDynamicSize(seqTy)) {43 std::int64_t numberOfElements = 1;44 for (std::int64_t i : seqTy.getShape()) {45 numberOfElements *= i;46 // If the count is suspicious, then don't change anything here.47 if (numberOfElements <= 0)48 return true;49 }50 // If the number of elements exceeds the threshold, move the allocation to51 // the heap.52 if (static_cast<std::size_t>(numberOfElements) >53 options.maxStackArraySize) {54 return false;55 }56 }57 }58 return true;59}60 61static mlir::Value genAllocmem(mlir::OpBuilder &builder, fir::AllocaOp alloca,62 bool deallocPointsDominateAlloc) {63 mlir::Type varTy = alloca.getInType();64 auto unpackName = [](std::optional<llvm::StringRef> opt) -> llvm::StringRef {65 if (opt)66 return *opt;67 return {};68 };69 llvm::StringRef uniqName = unpackName(alloca.getUniqName());70 llvm::StringRef bindcName = unpackName(alloca.getBindcName());71 auto heap = fir::AllocMemOp::create(builder, alloca.getLoc(), varTy, uniqName,72 bindcName, alloca.getTypeparams(),73 alloca.getShape());74 LLVM_DEBUG(llvm::dbgs() << "memory allocation opt: replaced " << alloca75 << " with " << heap << '\n');76 return heap;77}78 79static void genFreemem(mlir::Location loc, mlir::OpBuilder &builder,80 mlir::Value allocmem) {81 [[maybe_unused]] auto free = fir::FreeMemOp::create(builder, loc, allocmem);82 LLVM_DEBUG(llvm::dbgs() << "memory allocation opt: add free " << free83 << " for " << allocmem << '\n');84}85 86/// This pass can reclassify memory allocations (fir.alloca, fir.allocmem) based87/// on heuristics and settings. The intention is to allow better performance and88/// workarounds for conditions such as environments with limited stack space.89///90/// Currently, implements two conversions from stack to heap allocation.91/// 1. If a stack allocation is an array larger than some threshold value92/// make it a heap allocation.93/// 2. If a stack allocation is an array with a runtime evaluated size make94/// it a heap allocation.95namespace {96class MemoryAllocationOpt97 : public fir::impl::MemoryAllocationOptBase<MemoryAllocationOpt> {98public:99 MemoryAllocationOpt() {100 // Set options with default values. (See Passes.td.) Note that the101 // command-line options, e.g. dynamicArrayOnHeap, are not set yet.102 options = {dynamicArrayOnHeap, maxStackArraySize};103 }104 105 MemoryAllocationOpt(bool dynOnHeap, std::size_t maxStackSize) {106 // Set options with default values. (See Passes.td.)107 options = {dynOnHeap, maxStackSize};108 }109 110 MemoryAllocationOpt(const fir::MemoryAllocationOptOptions &options)111 : options{options} {}112 113 /// Override `options` if command-line options have been set.114 inline void useCommandLineOptions() {115 if (dynamicArrayOnHeap)116 options.dynamicArrayOnHeap = dynamicArrayOnHeap;117 if (maxStackArraySize != unlimitedArraySize)118 options.maxStackArraySize = maxStackArraySize;119 }120 121 void runOnOperation() override {122 auto *context = &getContext();123 auto func = getOperation();124 mlir::RewritePatternSet patterns(context);125 mlir::ConversionTarget target(*context);126 127 useCommandLineOptions();128 LLVM_DEBUG(llvm::dbgs()129 << "dynamic arrays on heap: " << options.dynamicArrayOnHeap130 << "\nmaximum number of elements of array on stack: "131 << options.maxStackArraySize << '\n');132 133 // If func is a declaration, skip it.134 if (func.empty())135 return;136 auto tryReplacing = [&](fir::AllocaOp alloca) {137 bool res = !keepStackAllocation(alloca, options);138 if (res) {139 LLVM_DEBUG(llvm::dbgs()140 << "memory allocation opt: found " << alloca << '\n');141 }142 return res;143 };144 mlir::IRRewriter rewriter(context);145 fir::replaceAllocas(rewriter, func.getOperation(), tryReplacing,146 genAllocmem, genFreemem);147 }148 149private:150 fir::MemoryAllocationOptOptions options;151};152} // namespace153