AllocationPolicy.h

//===-- Optimizer/Support/AllocationPolicy.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
//
//===----------------------------------------------------------------------===//
//
// Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
//
//===----------------------------------------------------------------------===//
//
// The policy controlling where array allocations should live: on the stack
// (fir.alloca) or on the heap (fir.allocmem). Lowering records it on the
// module, and on the functions that need a narrower one, so that policy-aware
// passes can make consistent decisions and dumped IR replays with the policy it
// was compiled with.
//
//===----------------------------------------------------------------------===//
 
#ifndef FORTRAN_OPTIMIZER_SUPPORT_ALLOCATIONPOLICY_H
#define FORTRAN_OPTIMIZER_SUPPORT_ALLOCATIONPOLICY_H
 
#include "flang/Optimizer/Dialect/Support/FIRContext.h"
#include "flang/Optimizer/Dialect/Support/KindMapping.h"
#include "flang/Optimizer/Support/DataLayout.h"
#include <cstddef>
#include <cstdint>
#include <functional>
#include <optional>
 
namespace mlir {
class Location;
class ModuleOp;
class Operation;
class Type;
} // namespace mlir
 
namespace fir {

struct AllocationPolicy {
  static constexpr bool stackArraysDefault = false;
  static constexpr std::uint64_t smallArrayThresholdBytesDefault = 1024;
  static constexpr std::uint64_t totalStackLimitBytesDefault =
      4ull * 1024 * 1024;

  bool stackArrays = stackArraysDefault;
  std::uint64_t smallArrayThresholdBytes = smallArrayThresholdBytesDefault;
  std::uint64_t totalStackLimitBytes = totalStackLimitBytesDefault;
};

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

struct PendingAllocationInfo {
  bool isTemporary = false;
  bool isDynamic = false;
  std::optional<std::int64_t> byteSize;
  std::optional<std::int64_t> alignment;
  mlir::Operation *context = nullptr;
};

std::optional<std::size_t> stackBytesNeeded(const PendingAllocationInfo &info);

struct AllocationSizeContext {
  std::optional<mlir::DataLayout> dataLayout;
  std::optional<fir::KindMapping> kindMap;
};

AllocationSizeContext getAllocationSizeContext(mlir::Operation *op);

bool shouldUseStackForCopyin(
    mlir::Location loc, mlir::Type sequenceType, const AllocationPolicy &policy,
    const AllocationSizeContext &sizeContext,
    std::optional<std::int64_t> alignment = std::nullopt);

struct AllocationInfo : PendingAllocationInfo {
  mlir::Operation *op = nullptr;
  bool isCurrentlyOnStack = false;
};

bool shouldAllocateOnStack(const PendingAllocationInfo &info,
                           const AllocationPolicy &policy,
                           std::size_t stackBytesUsed);

AllocationPlacement decideAllocationPlacement(const AllocationInfo &info,
                                              const AllocationPolicy &policy,
                                              std::size_t stackBytesUsed);

template <typename FieldT, typename OptionT>
void overrideIfExplicitlySet(FieldT &field, const OptionT &option) {
  if (option.hasValue())
    field = static_cast<FieldT>(option);
}

using AllocationPlacementHook = std::function<AllocationPlacement(
    const AllocationInfo & /*info*/, const AllocationPolicy & /*policy*/,
    std::size_t /*stackBytesUsed*/)>;

AllocationPolicy getCommandLineAllocationPolicy(bool stackArrays);

void setAllocationPolicy(mlir::ModuleOp mod, const AllocationPolicy &policy);

void setAllocationPolicy(mlir::Operation *op, const AllocationPolicy &policy);

AllocationPolicy getAllocationPolicy(mlir::ModuleOp mod);

std::optional<AllocationPolicy> getLocalAllocationPolicy(mlir::Operation *op);

AllocationPolicy getAllocationPolicy(mlir::Operation *op);
 
} // namespace fir
 
#endif // FORTRAN_OPTIMIZER_SUPPORT_ALLOCATIONPOLICY_H