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 21 namespace 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. 29 static 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`). 33 static inline bool 34 keepStackAllocation(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 61 static 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 79 static 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. 95 namespace { 96 class MemoryAllocationOpt 97 : public fir::impl::MemoryAllocationOptBase<MemoryAllocationOpt> { 98 public: 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 149 private: 150 fir::MemoryAllocationOptOptions options; 151 }; 152 } // namespace 153