| //===- MemoryAllocation.cpp -----------------------------------------------===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "flang/Optimizer/Dialect/FIRDialect.h" |
| #include "flang/Optimizer/Dialect/FIROps.h" |
| #include "flang/Optimizer/Dialect/FIRType.h" |
| #include "flang/Optimizer/Transforms/Passes.h" |
| #include "mlir/Dialect/Func/IR/FuncOps.h" |
| #include "mlir/IR/Diagnostics.h" |
| #include "mlir/Pass/Pass.h" |
| #include "mlir/Transforms/DialectConversion.h" |
| #include "mlir/Transforms/Passes.h" |
| #include "llvm/ADT/TypeSwitch.h" |
| |
| namespace fir { |
| #define GEN_PASS_DEF_MEMORYALLOCATIONOPT |
| #include "flang/Optimizer/Transforms/Passes.h.inc" |
| } // namespace fir |
| |
| #define DEBUG_TYPE "flang-memory-allocation-opt" |
| |
| // Number of elements in an array does not determine where it is allocated. |
| static constexpr std::size_t unlimitedArraySize = ~static_cast<std::size_t>(0); |
| |
| namespace { |
| class ReturnAnalysis { |
| public: |
| MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ReturnAnalysis) |
| |
| ReturnAnalysis(mlir::Operation *op) { |
| if (auto func = mlir::dyn_cast<mlir::func::FuncOp>(op)) |
| for (mlir::Block &block : func) |
| for (mlir::Operation &i : block) |
| if (mlir::isa<mlir::func::ReturnOp>(i)) { |
| returnMap[op].push_back(&i); |
| break; |
| } |
| } |
| |
| llvm::SmallVector<mlir::Operation *> getReturns(mlir::Operation *func) const { |
| auto iter = returnMap.find(func); |
| if (iter != returnMap.end()) |
| return iter->second; |
| return {}; |
| } |
| |
| private: |
| llvm::DenseMap<mlir::Operation *, llvm::SmallVector<mlir::Operation *>> |
| returnMap; |
| }; |
| } // namespace |
| |
| /// Return `true` if this allocation is to remain on the stack (`fir.alloca`). |
| /// Otherwise the allocation should be moved to the heap (`fir.allocmem`). |
| static inline bool |
| keepStackAllocation(fir::AllocaOp alloca, mlir::Block *entry, |
| const fir::MemoryAllocationOptOptions &options) { |
| // Limitation: only arrays allocated on the stack in the entry block are |
| // considered for now. |
| // TODO: Generalize the algorithm and placement of the freemem nodes. |
| if (alloca->getBlock() != entry) |
| return true; |
| if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(alloca.getInType())) { |
| if (fir::hasDynamicSize(seqTy)) { |
| // Move all arrays with runtime determined size to the heap. |
| if (options.dynamicArrayOnHeap) |
| return false; |
| } else { |
| std::int64_t numberOfElements = 1; |
| for (std::int64_t i : seqTy.getShape()) { |
| numberOfElements *= i; |
| // If the count is suspicious, then don't change anything here. |
| if (numberOfElements <= 0) |
| return true; |
| } |
| // If the number of elements exceeds the threshold, move the allocation to |
| // the heap. |
| if (static_cast<std::size_t>(numberOfElements) > |
| options.maxStackArraySize) { |
| LLVM_DEBUG(llvm::dbgs() |
| << "memory allocation opt: found " << alloca << '\n'); |
| return false; |
| } |
| } |
| } |
| return true; |
| } |
| |
| namespace { |
| class AllocaOpConversion : public mlir::OpRewritePattern<fir::AllocaOp> { |
| public: |
| using OpRewritePattern::OpRewritePattern; |
| |
| AllocaOpConversion(mlir::MLIRContext *ctx, |
| llvm::ArrayRef<mlir::Operation *> rets) |
| : OpRewritePattern(ctx), returnOps(rets) {} |
| |
| mlir::LogicalResult |
| matchAndRewrite(fir::AllocaOp alloca, |
| mlir::PatternRewriter &rewriter) const override { |
| auto loc = alloca.getLoc(); |
| mlir::Type varTy = alloca.getInType(); |
| auto unpackName = |
| [](std::optional<llvm::StringRef> opt) -> llvm::StringRef { |
| if (opt) |
| return *opt; |
| return {}; |
| }; |
| auto uniqName = unpackName(alloca.getUniqName()); |
| auto bindcName = unpackName(alloca.getBindcName()); |
| auto heap = rewriter.create<fir::AllocMemOp>( |
| loc, varTy, uniqName, bindcName, alloca.getTypeparams(), |
| alloca.getShape()); |
| auto insPt = rewriter.saveInsertionPoint(); |
| for (mlir::Operation *retOp : returnOps) { |
| rewriter.setInsertionPoint(retOp); |
| [[maybe_unused]] auto free = rewriter.create<fir::FreeMemOp>(loc, heap); |
| LLVM_DEBUG(llvm::dbgs() << "memory allocation opt: add free " << free |
| << " for " << heap << '\n'); |
| } |
| rewriter.restoreInsertionPoint(insPt); |
| rewriter.replaceOpWithNewOp<fir::ConvertOp>( |
| alloca, fir::ReferenceType::get(varTy), heap); |
| LLVM_DEBUG(llvm::dbgs() << "memory allocation opt: replaced " << alloca |
| << " with " << heap << '\n'); |
| return mlir::success(); |
| } |
| |
| private: |
| llvm::ArrayRef<mlir::Operation *> returnOps; |
| }; |
| |
| /// This pass can reclassify memory allocations (fir.alloca, fir.allocmem) based |
| /// on heuristics and settings. The intention is to allow better performance and |
| /// workarounds for conditions such as environments with limited stack space. |
| /// |
| /// Currently, implements two conversions from stack to heap allocation. |
| /// 1. If a stack allocation is an array larger than some threshold value |
| /// make it a heap allocation. |
| /// 2. If a stack allocation is an array with a runtime evaluated size make |
| /// it a heap allocation. |
| class MemoryAllocationOpt |
| : public fir::impl::MemoryAllocationOptBase<MemoryAllocationOpt> { |
| public: |
| MemoryAllocationOpt() { |
| // Set options with default values. (See Passes.td.) Note that the |
| // command-line options, e.g. dynamicArrayOnHeap, are not set yet. |
| options = {dynamicArrayOnHeap, maxStackArraySize}; |
| } |
| |
| MemoryAllocationOpt(bool dynOnHeap, std::size_t maxStackSize) { |
| // Set options with default values. (See Passes.td.) |
| options = {dynOnHeap, maxStackSize}; |
| } |
| |
| MemoryAllocationOpt(const fir::MemoryAllocationOptOptions &options) |
| : options{options} {} |
| |
| /// Override `options` if command-line options have been set. |
| inline void useCommandLineOptions() { |
| if (dynamicArrayOnHeap) |
| options.dynamicArrayOnHeap = dynamicArrayOnHeap; |
| if (maxStackArraySize != unlimitedArraySize) |
| options.maxStackArraySize = maxStackArraySize; |
| } |
| |
| void runOnOperation() override { |
| auto *context = &getContext(); |
| auto func = getOperation(); |
| mlir::RewritePatternSet patterns(context); |
| mlir::ConversionTarget target(*context); |
| |
| useCommandLineOptions(); |
| LLVM_DEBUG(llvm::dbgs() |
| << "dynamic arrays on heap: " << options.dynamicArrayOnHeap |
| << "\nmaximum number of elements of array on stack: " |
| << options.maxStackArraySize << '\n'); |
| |
| // If func is a declaration, skip it. |
| if (func.empty()) |
| return; |
| |
| const auto &analysis = getAnalysis<ReturnAnalysis>(); |
| |
| target.addLegalDialect<fir::FIROpsDialect, mlir::arith::ArithDialect, |
| mlir::func::FuncDialect>(); |
| target.addDynamicallyLegalOp<fir::AllocaOp>([&](fir::AllocaOp alloca) { |
| return keepStackAllocation(alloca, &func.front(), options); |
| }); |
| |
| llvm::SmallVector<mlir::Operation *> returnOps = analysis.getReturns(func); |
| patterns.insert<AllocaOpConversion>(context, returnOps); |
| if (mlir::failed( |
| mlir::applyPartialConversion(func, target, std::move(patterns)))) { |
| mlir::emitError(func.getLoc(), |
| "error in memory allocation optimization\n"); |
| signalPassFailure(); |
| } |
| } |
| |
| private: |
| fir::MemoryAllocationOptOptions options; |
| }; |
| } // namespace |