/*========================== begin_copyright_notice ============================

Copyright (C) 2017-2023 Intel Corporation

SPDX-License-Identifier: MIT

============================= end_copyright_notice ===========================*/

#pragma once

#include "common/LLVMWarningsPush.hpp"
#include "llvm/ADT/DenseMap.h"
#include "llvm/Pass.h"
#include <llvm/IR/InstVisitor.h>
#include <llvm/ADT/StringRef.h>
#include "llvm/Support/ScaledNumber.h"
#include "common/LLVMWarningsPop.hpp"
#include "Probe/Assertion.h"
#include <cstddef>
#include <deque>
#include <functional>
#include <unordered_map>
#include <unordered_set>

namespace llvm {
class Loop;
class LoopInfo;
class BlockFrequencyInfo;
class ScalarEvolution;
} // namespace llvm

namespace IGC {
namespace IGCMD {
class MetaDataUtils;
}
class CodeGenContext;
} // namespace IGC

namespace IGC {

/// \brief Estimate function size after complete inlining.
///
/// This pass visits the call graph and estimates the number of llvm IR
/// instructions after complete inlining.
class EstimateFunctionSize : public llvm::ModulePass, public llvm::InstVisitor<EstimateFunctionSize> {
public:
  static char ID;

  enum AnalysisLevel { AL_Module, AL_Kernel };

  explicit EstimateFunctionSize(AnalysisLevel = AL_Module, bool = false);
  ~EstimateFunctionSize();
  EstimateFunctionSize(const EstimateFunctionSize &) = delete;
  EstimateFunctionSize &operator=(const EstimateFunctionSize &) = delete;
  llvm::StringRef getPassName() const override { return "Estimate Function Sizes"; }
  void getAnalysisUsage(llvm::AnalysisUsage &AU) const override;
  bool runOnModule(llvm::Module &M) override;

  // Per-function analyses (LoopInfo/BlockFrequencyInfo/ScalarEvolution) and the MetaDataUtils are
  // injected by the caller so the engine works under both pass managers. The legacy runOnModule
  // wires the getters to getAnalysis<...>(F); the new pass manager consumer calls computeInline()
  // with getters backed by the function analysis manager. This avoids renaming the pass, which is
  // required because it is also consumed by codegen passes (GenCodeGenModule, CallMergerPass).
  using LoopInfoGetter = std::function<llvm::LoopInfo &(llvm::Function &)>;
  using BFIGetter = std::function<llvm::BlockFrequencyInfo &(llvm::Function &)>;
  using SEGetter = std::function<llvm::ScalarEvolution &(llvm::Function &)>;

  // Compute the analysis on a stack-constructed instance (outside any pass manager).
  bool computeInline(llvm::Module &M, LoopInfoGetter LIG, BFIGetter BFIG, SEGetter SEG,
                     IGC::IGCMD::MetaDataUtils *MdUtils, IGC::CodeGenContext *Ctx);

  /// \brief Return the estimated maximal function size after complete inlining.
  std::size_t getMaxExpandedSize() const;

  /// \brief Return the estimated function size after complete inlining.
  std::size_t getExpandedSize(const llvm::Function *F) const;

  bool onlyCalledOnce(const llvm::Function *F, const llvm::Function *CallerF);

  bool hasRecursion() const { return HasRecursion; }

  bool shouldEnableSubroutine() const { return EnableSubroutine; }

  bool isTrimmedFunction(llvm::Function *F);

  void visitCallInst(llvm::CallInst &CI);

  bool isStackCallAssigned(llvm::Function *F);

  bool isLargeKernelThresholdExceeded() const;

private:
  struct TotalTrimmingPool {
    llvm::SmallVector<void *, 64> trimming_pool;
    llvm::SmallVector<void *, 64> tiny_fn_trimming_pool;
    llvm::SmallVector<void *, 64> implicit_arg_trimming_pool;
  };

  void analyze();
  void checkSubroutine();
  void clear();
  void reduceKernelSize();

  /// \brief Return the associated opaque data.
  template <typename T> T *get(llvm::Function *F) {
    IGC_ASSERT(ECG.count(F));
    return static_cast<T *>(ECG[F]);
  }
  void performImplArgsAnalysis();
  void initializeTopologicalVisit(llvm::Function *root, std::unordered_map<void *, uint32_t> &FunctionsInKernel,
                                  std::deque<void *> &BottomUpQueue, bool ignoreStackCallBoundary);
  uint32_t updateExpandedUnitSize(llvm::Function *F, bool ignoreStackCallBoundary);
  void UpdateSizeAfterCollapsing(std::deque<void *> &leafNodes, std::unordered_set<void *> &funcsInKernel);
  void updateInlineCnt(llvm::Function *root);
  llvm::ScaledNumber<uint64_t> calculateTotalWeight(llvm::Function *root);
  uint32_t bottomUpHeuristic(llvm::Function *F, uint32_t &stackCall_cnt);
  void partitionKernel();
  void runStaticAnalysis();
  void reduceCompilationUnitSize();
  void trimCompilationUnit(llvm::SmallVector<void *, 64> &unitHeads, uint32_t threshold, bool ignoreStackCallBoundary);
  void performTrimming(llvm::Function *head, llvm::SmallVector<void *, 64> &functions_to_trim, uint32_t threshold,
                       bool ignoreStackCallBoundary);
  void performGreedyTrimming(llvm::Function *head, llvm::SmallVector<void *, 64> &functions_to_trim, uint32_t threshold,
                             bool ignoreStackCallBoundary);
  uint32_t getMaxUnitSize();
  void getFunctionsToTrim(llvm::Function *root, TotalTrimmingPool &pools, bool ignoreStackCallBoundary,
                          uint32_t &func_cnt);
  void updateStaticFuncFreq();
  void estimateTotalLoopIteration(llvm::Function &F, llvm::LoopInfo *LI);

  /// \brief Run the core analysis once the getters/MetaDataUtils are set up.
  bool runAnalysis(llvm::Module &Mod);

  /// \brief The module being analyzed.
  llvm::Module *M;

  /// \brief Injected per-function analysis getters and MetaDataUtils (see computeInline).
  LoopInfoGetter m_getLI;
  BFIGetter m_getBFI;
  SEGetter m_getSE;
  IGC::IGCMD::MetaDataUtils *m_pMdUtils = nullptr;
  IGC::CodeGenContext *m_pCtx = nullptr;

  /// \brief The analysis level to be performed.
  AnalysisLevel AL;

  bool tmpHasImplicitArg;
  bool matchImplicitArg(llvm::CallInst &CI);

  const llvm::StringRef GET_LOCAL_ID_X = "__builtin_IB_get_local_id_x";
  const llvm::StringRef GET_LOCAL_ID_Y = "__builtin_IB_get_local_id_y";
  const llvm::StringRef GET_LOCAL_ID_Z = "__builtin_IB_get_local_id_z";
  const llvm::StringRef GET_GROUP_ID = "__builtin_IB_get_group_id";
  const llvm::StringRef GET_LOCAL_THREAD_ID = "__builtin_IB_get_local_thread_id";
  const llvm::StringRef GET_GLOBAL_SIZE = "__builtin_IB_get_global_size";
  const llvm::StringRef GET_LOCAL_SIZE = "__builtin_IB_get_local_size";
  const llvm::StringRef GET_GLOBAL_OFFSET = "__builtin_IB_get_global_offset";
  const llvm::StringRef GET_WORK_DIM = "__builtin_IB_get_work_dim";
  const llvm::StringRef GET_NUM_GROUPS = "__builtin_IB_get_num_groups";
  const llvm::StringRef GET_ENQUEUED_LOCAL_SIZE = "__builtin_IB_get_enqueued_local_size";
  const llvm::StringRef GET_STAGE_IN_GRID_ORIGIN = "__builtin_IB_get_stage_in_grid_origin";
  const llvm::StringRef GET_STAGE_IN_GRID_SIZE = "__builtin_IB_get_stage_in_grid_size";
  const llvm::StringRef GET_SYNC_BUFFER = "__builtin_IB_get_sync_buffer";
  const llvm::StringRef GET_ASSERT_BUFFER = "__builtin_IB_get_assert_buffer";
  const llvm::StringRef GET_REGION_GROUP_SIZE = "__builtin_IB_get_region_group_size";
  const llvm::StringRef GET_REGION_GROUP_WG_COUNT = "__builtin_IB_get_region_group_wg_count";
  const llvm::StringRef GET_REGION_GROUP_BARRIER_BUFFER = "__builtin_IB_get_region_group_barrier_buffer";
  bool HasRecursion;
  bool EnableSubroutine;

  /// Internal data structure for the analysis which is approximately an
  /// extended call graph.
  llvm::SmallDenseMap<llvm::Function *, void *> ECG;
  // Kernel entries
  llvm::SmallVector<void *, 64> kernelEntries;
  // Functions that are assigned stackcalls
  llvm::SmallVector<void *, 64> stackCallFuncs;
  llvm::SmallVector<void *, 64> addressTakenFuncs;
  llvm::ScaledNumber<uint64_t> threshold_func_freq;
  llvm::ScaledNumber<uint64_t> thresholdForTrimming;
  std::unordered_map<llvm::Loop *, llvm::ScaledNumber<uint64_t>> LoopIterCnts;

  // Flags for Kernel trimming
  bool ControlKernelTotalSize;
  bool ControlUnitSize;
  unsigned ControlInlineTinySize;
  unsigned UnitSizeThreshold;

  // Flags for Static Profile-guided trimming
  bool StaticProfileGuidedTrimming;
  bool UseFrequencyInfoForSPGT;
  bool BlockFrequencySampling;
  bool EnableLeafCollapsing;
  bool EnableSizeContributionOptimization;
  bool LoopCountAwareTrimming;
  bool EnableGreedyTrimming;
  unsigned SizeWeightForSPGT;
  unsigned FrequencyWeightForSPGT;
  unsigned MetricForKernelSizeReduction;
  unsigned ParameterForColdFuncThreshold;
  unsigned ControlInlineTinySizeForSPGT;
  unsigned MaxUnrollCountForFunctionSizeAnalysis;
  unsigned SkipTrimmingOneCopyFunction;
  std::string SelectiveTrimming;

  // Flags for Partitioning
  bool PartitionUnit;
  bool StaticProfileGuidedPartitioning;

  // Flags for implicit arguments and external functions
  bool ForceInlineExternalFunctions;
  bool ForceInlineStackCallWithImplArg;
  bool ControlInlineImplicitArgs;
  bool TrimImplicitArgFunctionsForLargeKernels;
  unsigned SubroutineThreshold;
  unsigned KernelTotalSizeThreshold;
  unsigned ExpandedUnitSizeThreshold;

  // Flags for large kernel handling
  unsigned LargeKernelSmallFunctionLimit;
  unsigned LargeKernelThresholdMultiplier;
};

llvm::ModulePass *createEstimateFunctionSizePass();
llvm::ModulePass *createEstimateFunctionSizePass(bool);
llvm::ModulePass *createEstimateFunctionSizePass(EstimateFunctionSize::AnalysisLevel);

} // namespace IGC
