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.uniq_name());
119     auto bindcName = unpackName(alloca.bindc_name());
120     auto heap = rewriter.create<fir::AllocMemOp>(
121         loc, varTy, uniqName, bindcName, alloca.typeparams(), alloca.shape());
122     auto insPt = rewriter.saveInsertionPoint();
123     for (mlir::Operation *retOp : returnOps) {
124       rewriter.setInsertionPoint(retOp);
125       [[maybe_unused]] auto free = rewriter.create<fir::FreeMemOp>(loc, heap);
126       LLVM_DEBUG(llvm::dbgs() << "memory allocation opt: add free " << free
127                               << " for " << heap << '\n');
128     }
129     rewriter.restoreInsertionPoint(insPt);
130     rewriter.replaceOpWithNewOp<fir::ConvertOp>(
131         alloca, fir::ReferenceType::get(varTy), heap);
132     LLVM_DEBUG(llvm::dbgs() << "memory allocation opt: replaced " << alloca
133                             << " with " << heap << '\n');
134     return mlir::success();
135   }
136 
137 private:
138   llvm::ArrayRef<mlir::Operation *> returnOps;
139 };
140 
141 /// This pass can reclassify memory allocations (fir.alloca, fir.allocmem) based
142 /// on heuristics and settings. The intention is to allow better performance and
143 /// workarounds for conditions such as environments with limited stack space.
144 ///
145 /// Currently, implements two conversions from stack to heap allocation.
146 ///   1. If a stack allocation is an array larger than some threshold value
147 ///      make it a heap allocation.
148 ///   2. If a stack allocation is an array with a runtime evaluated size make
149 ///      it a heap allocation.
150 class MemoryAllocationOpt
151     : public fir::MemoryAllocationOptBase<MemoryAllocationOpt> {
152 public:
153   MemoryAllocationOpt() {
154     // Set options with default values. (See Passes.td.) Note that the
155     // command-line options, e.g. dynamicArrayOnHeap,  are not set yet.
156     options = {dynamicArrayOnHeap, maxStackArraySize};
157   }
158 
159   MemoryAllocationOpt(bool dynOnHeap, std::size_t maxStackSize) {
160     // Set options with default values. (See Passes.td.)
161     options = {dynOnHeap, maxStackSize};
162   }
163 
164   /// Override `options` if command-line options have been set.
165   inline void useCommandLineOptions() {
166     if (dynamicArrayOnHeap)
167       options.dynamicArrayOnHeap = dynamicArrayOnHeap;
168     if (maxStackArraySize != unlimitedArraySize)
169       options.maxStackArraySize = maxStackArraySize;
170   }
171 
172   void runOnOperation() override {
173     auto *context = &getContext();
174     auto func = getOperation();
175     mlir::RewritePatternSet patterns(context);
176     mlir::ConversionTarget target(*context);
177 
178     useCommandLineOptions();
179     LLVM_DEBUG(llvm::dbgs()
180                << "dynamic arrays on heap: " << options.dynamicArrayOnHeap
181                << "\nmaximum number of elements of array on stack: "
182                << options.maxStackArraySize << '\n');
183 
184     // If func is a declaration, skip it.
185     if (func.empty())
186       return;
187 
188     const auto &analysis = getAnalysis<ReturnAnalysis>();
189 
190     target.addLegalDialect<fir::FIROpsDialect, mlir::arith::ArithmeticDialect,
191                            mlir::StandardOpsDialect>();
192     target.addDynamicallyLegalOp<fir::AllocaOp>([&](fir::AllocaOp alloca) {
193       return keepStackAllocation(alloca, &func.front(), options);
194     });
195 
196     patterns.insert<AllocaOpConversion>(context, analysis.getReturns(func));
197     if (mlir::failed(
198             mlir::applyPartialConversion(func, target, std::move(patterns)))) {
199       mlir::emitError(func.getLoc(),
200                       "error in memory allocation optimization\n");
201       signalPassFailure();
202     }
203   }
204 
205 private:
206   MemoryAllocationOptions options;
207 };
208 } // namespace
209 
210 std::unique_ptr<mlir::Pass> fir::createMemoryAllocationPass() {
211   return std::make_unique<MemoryAllocationOpt>();
212 }
213 
214 std::unique_ptr<mlir::Pass>
215 fir::createMemoryAllocationPass(bool dynOnHeap, std::size_t maxStackSize) {
216   return std::make_unique<MemoryAllocationOpt>(dynOnHeap, maxStackSize);
217 }
218