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 "PassDetail.h"
10 #include "flang/Optimizer/Dialect/FIRDialect.h"
11 #include "flang/Optimizer/Dialect/FIROps.h"
12 #include "flang/Optimizer/Dialect/FIRType.h"
13 #include "flang/Optimizer/Transforms/Passes.h"
14 #include "mlir/Dialect/StandardOps/IR/Ops.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 
21 #define DEBUG_TYPE "flang-memory-allocation-opt"
22 
23 // Number of elements in an array does not determine where it is allocated.
24 static constexpr std::size_t unlimitedArraySize = ~static_cast<std::size_t>(0);
25 
26 namespace {
27 struct MemoryAllocationOptions {
28   // Always move dynamic array allocations to the heap. This may result in more
29   // heap fragmentation, so may impact performance negatively.
30   bool dynamicArrayOnHeap = false;
31 
32   // Number of elements in array threshold for moving to heap. In environments
33   // with limited stack size, moving large arrays to the heap can avoid running
34   // out of stack space.
35   std::size_t maxStackArraySize = unlimitedArraySize;
36 };
37 
38 class ReturnAnalysis {
39 public:
40   ReturnAnalysis(mlir::Operation *op) {
41     if (auto func = mlir::dyn_cast<mlir::FuncOp>(op))
42       for (mlir::Block &block : func)
43         for (mlir::Operation &i : block)
44           if (mlir::isa<mlir::ReturnOp>(i)) {
45             returnMap[op].push_back(&i);
46             break;
47           }
48   }
49 
50   llvm::SmallVector<mlir::Operation *> getReturns(mlir::Operation *func) const {
51     auto iter = returnMap.find(func);
52     if (iter != returnMap.end())
53       return iter->second;
54     return {};
55   }
56 
57 private:
58   llvm::DenseMap<mlir::Operation *, llvm::SmallVector<mlir::Operation *>>
59       returnMap;
60 };
61 } // namespace
62 
63 /// Return `true` if this allocation is to remain on the stack (`fir.alloca`).
64 /// Otherwise the allocation should be moved to the heap (`fir.allocmem`).
65 static inline bool keepStackAllocation(fir::AllocaOp alloca, mlir::Block *entry,
66                                        const MemoryAllocationOptions &options) {
67   // Limitation: only arrays allocated on the stack in the entry block are
68   // considered for now.
69   // TODO: Generalize the algorithm and placement of the freemem nodes.
70   if (alloca->getBlock() != entry)
71     return true;
72   if (auto seqTy = alloca.getInType().dyn_cast<fir::SequenceType>()) {
73     if (fir::hasDynamicSize(seqTy)) {
74       // Move all arrays with runtime determined size to the heap.
75       if (options.dynamicArrayOnHeap)
76         return false;
77     } else {
78       std::int64_t numberOfElements = 1;
79       for (std::int64_t i : seqTy.getShape()) {
80         numberOfElements *= i;
81         // If the count is suspicious, then don't change anything here.
82         if (numberOfElements <= 0)
83           return true;
84       }
85       // If the number of elements exceeds the threshold, move the allocation to
86       // the heap.
87       if (static_cast<std::size_t>(numberOfElements) >
88           options.maxStackArraySize) {
89         LLVM_DEBUG(llvm::dbgs()
90                    << "memory allocation opt: found " << alloca << '\n');
91         return false;
92       }
93     }
94   }
95   return true;
96 }
97 
98 namespace {
99 class AllocaOpConversion : public mlir::OpRewritePattern<fir::AllocaOp> {
100 public:
101   using OpRewritePattern::OpRewritePattern;
102 
103   AllocaOpConversion(mlir::MLIRContext *ctx,
104                      llvm::ArrayRef<mlir::Operation *> rets)
105       : OpRewritePattern(ctx), returnOps(rets) {}
106 
107   mlir::LogicalResult
108   matchAndRewrite(fir::AllocaOp alloca,
109                   mlir::PatternRewriter &rewriter) const override {
110     auto loc = alloca.getLoc();
111     mlir::Type varTy = alloca.getInType();
112     auto unpackName =
113         [](llvm::Optional<llvm::StringRef> opt) -> llvm::StringRef {
114       if (opt)
115         return *opt;
116       return {};
117     };
118     auto uniqName = unpackName(alloca.getUniqName());
119     auto bindcName = unpackName(alloca.getBindcName());
120     auto heap = rewriter.create<fir::AllocMemOp>(
121         loc, varTy, uniqName, bindcName, alloca.getTypeparams(),
122         alloca.getShape());
123     auto insPt = rewriter.saveInsertionPoint();
124     for (mlir::Operation *retOp : returnOps) {
125       rewriter.setInsertionPoint(retOp);
126       [[maybe_unused]] auto free = rewriter.create<fir::FreeMemOp>(loc, heap);
127       LLVM_DEBUG(llvm::dbgs() << "memory allocation opt: add free " << free
128                               << " for " << heap << '\n');
129     }
130     rewriter.restoreInsertionPoint(insPt);
131     rewriter.replaceOpWithNewOp<fir::ConvertOp>(
132         alloca, fir::ReferenceType::get(varTy), heap);
133     LLVM_DEBUG(llvm::dbgs() << "memory allocation opt: replaced " << alloca
134                             << " with " << heap << '\n');
135     return mlir::success();
136   }
137 
138 private:
139   llvm::ArrayRef<mlir::Operation *> returnOps;
140 };
141 
142 /// This pass can reclassify memory allocations (fir.alloca, fir.allocmem) based
143 /// on heuristics and settings. The intention is to allow better performance and
144 /// workarounds for conditions such as environments with limited stack space.
145 ///
146 /// Currently, implements two conversions from stack to heap allocation.
147 ///   1. If a stack allocation is an array larger than some threshold value
148 ///      make it a heap allocation.
149 ///   2. If a stack allocation is an array with a runtime evaluated size make
150 ///      it a heap allocation.
151 class MemoryAllocationOpt
152     : public fir::MemoryAllocationOptBase<MemoryAllocationOpt> {
153 public:
154   MemoryAllocationOpt() {
155     // Set options with default values. (See Passes.td.) Note that the
156     // command-line options, e.g. dynamicArrayOnHeap,  are not set yet.
157     options = {dynamicArrayOnHeap, maxStackArraySize};
158   }
159 
160   MemoryAllocationOpt(bool dynOnHeap, std::size_t maxStackSize) {
161     // Set options with default values. (See Passes.td.)
162     options = {dynOnHeap, maxStackSize};
163   }
164 
165   /// Override `options` if command-line options have been set.
166   inline void useCommandLineOptions() {
167     if (dynamicArrayOnHeap)
168       options.dynamicArrayOnHeap = dynamicArrayOnHeap;
169     if (maxStackArraySize != unlimitedArraySize)
170       options.maxStackArraySize = maxStackArraySize;
171   }
172 
173   void runOnOperation() override {
174     auto *context = &getContext();
175     auto func = getOperation();
176     mlir::RewritePatternSet patterns(context);
177     mlir::ConversionTarget target(*context);
178 
179     useCommandLineOptions();
180     LLVM_DEBUG(llvm::dbgs()
181                << "dynamic arrays on heap: " << options.dynamicArrayOnHeap
182                << "\nmaximum number of elements of array on stack: "
183                << options.maxStackArraySize << '\n');
184 
185     // If func is a declaration, skip it.
186     if (func.empty())
187       return;
188 
189     const auto &analysis = getAnalysis<ReturnAnalysis>();
190 
191     target.addLegalDialect<fir::FIROpsDialect, mlir::arith::ArithmeticDialect,
192                            mlir::StandardOpsDialect>();
193     target.addDynamicallyLegalOp<fir::AllocaOp>([&](fir::AllocaOp alloca) {
194       return keepStackAllocation(alloca, &func.front(), options);
195     });
196 
197     patterns.insert<AllocaOpConversion>(context, analysis.getReturns(func));
198     if (mlir::failed(
199             mlir::applyPartialConversion(func, target, std::move(patterns)))) {
200       mlir::emitError(func.getLoc(),
201                       "error in memory allocation optimization\n");
202       signalPassFailure();
203     }
204   }
205 
206 private:
207   MemoryAllocationOptions options;
208 };
209 } // namespace
210 
211 std::unique_ptr<mlir::Pass> fir::createMemoryAllocationPass() {
212   return std::make_unique<MemoryAllocationOpt>();
213 }
214 
215 std::unique_ptr<mlir::Pass>
216 fir::createMemoryAllocationPass(bool dynOnHeap, std::size_t maxStackSize) {
217   return std::make_unique<MemoryAllocationOpt>(dynOnHeap, maxStackSize);
218 }
219