flang/lib/Optimizer/Transforms/MemoryAllocation.cpp

   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