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