AllocationPlacementPolicy.h

//===-- Optimizer/Transforms/AllocationPlacementPolicy.h --------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This header defines the policy used by the allocation-placement pass to
// decide whether an array allocation should live on the stack (fir.alloca) or
// on the heap (fir.allocmem). The policy is expressed with tunable thresholds
// and a per-function stack budget, and can be customized through a hook (e.g.
// with different thresholds inside device routines or parallel regions).
//
//===----------------------------------------------------------------------===//
 
#ifndef FORTRAN_OPTIMIZER_TRANSFORMS_ALLOCATIONPLACEMENTPOLICY_H
#define FORTRAN_OPTIMIZER_TRANSFORMS_ALLOCATIONPLACEMENTPOLICY_H
 
#include <cstddef>
#include <cstdint>
#include <functional>
#include <optional>
 
namespace mlir {
class Operation;
} // namespace mlir
 
namespace fir {

enum class AllocationPlacement {
  Stack,
  Heap,
  Leave,
};

struct AllocationPlacementThresholds {
  bool stackArrays = false;
  std::size_t smallArrayThresholdBytes = 64;
  std::size_t totalStackLimitBytes = 4ull * 1024 * 1024;
};

struct AllocationInfo {
  mlir::Operation *op = nullptr;
  bool isCurrentlyOnStack = false;
  bool isTemporary = false;
  bool isDynamic = false;
  std::optional<std::int64_t> byteSize;
};

AllocationPlacement
decideAllocationPlacement(const AllocationInfo &info,
                          const AllocationPlacementThresholds &thresholds,
                          std::size_t stackBytesUsed);

using AllocationPlacementHook = std::function<AllocationPlacement(
    const AllocationInfo & /*info*/,
    const AllocationPlacementThresholds & /*thresholds*/,
    std::size_t /*stackBytesUsed*/)>;
 
} // namespace fir
 
#endif // FORTRAN_OPTIMIZER_TRANSFORMS_ALLOCATIONPLACEMENTPOLICY_H