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-exception
6//
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_MEMORYALLOCATIONOPT
23#include "flang/Optimizer/Transforms/Passes.h.inc"
24} // namespace fir
25
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 bool
34keepStackAllocation(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 to
51 // 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 = builder.create<fir::AllocMemOp>(alloca.getLoc(), varTy, uniqName,
72 bindcName, alloca.getTypeparams(),
73 alloca.getShape());
74 LLVM_DEBUG(llvm::dbgs() << "memory allocation opt: replaced " << alloca
75 << " 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 = builder.create<fir::FreeMemOp>(loc, allocmem);
82 LLVM_DEBUG(llvm::dbgs() << "memory allocation opt: add free " << free
83 << " for " << allocmem << '\n');
84}
85
86/// This pass can reclassify memory allocations (fir.alloca, fir.allocmem) based
87/// on heuristics and settings. The intention is to allow better performance and
88/// 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 value
92/// make it a heap allocation.
93/// 2. If a stack allocation is an array with a runtime evaluated size make
94/// it a heap allocation.
95namespace {
96class MemoryAllocationOpt
97 : public fir::impl::MemoryAllocationOptBase<MemoryAllocationOpt> {
98public:
99 MemoryAllocationOpt() {
100 // Set options with default values. (See Passes.td.) Note that the
101 // 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.dynamicArrayOnHeap
130 << "\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} // namespace
153

source code of flang/lib/Optimizer/Transforms/MemoryAllocation.cpp