//===- SampleProfile.cpp - Incorporate sample profiles into the IR --------===//
//
// 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 file implements the SampleProfileLoader transformation. This pass
// reads a profile file generated by a sampling profiler (e.g. Linux Perf -
// http://perf.wiki.kernel.org/) and generates IR metadata to reflect the
// profile information in the given profile.
//
// This pass generates branch weight annotations on the IR:
//
// - prof: Represents branch weights. This annotation is added to branches
//      to indicate the weights of each edge coming out of the branch.
//      The weight of each edge is the weight of the target block for
//      that edge. The weight of a block B is computed as the maximum
//      number of samples found in B.
//
//===----------------------------------------------------------------------===//

#include "llvm/Transforms/IPO/SampleProfile.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/PriorityQueue.h"
#include "llvm/ADT/SCCIterator.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/ADT/StringMap.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.h"
#include "llvm/Analysis/AssumptionCache.h"
#include "llvm/Analysis/BlockFrequencyInfoImpl.h"
#include "llvm/Analysis/InlineAdvisor.h"
#include "llvm/Analysis/InlineCost.h"
#include "llvm/Analysis/LazyCallGraph.h"
#include "llvm/Analysis/OptimizationRemarkEmitter.h"
#include "llvm/Analysis/ProfileSummaryInfo.h"
#include "llvm/Analysis/ReplayInlineAdvisor.h"
#include "llvm/Analysis/TargetLibraryInfo.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/DebugLoc.h"
#include "llvm/IR/DiagnosticInfo.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/GlobalValue.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/MDBuilder.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/PassManager.h"
#include "llvm/IR/ProfDataUtils.h"
#include "llvm/IR/PseudoProbe.h"
#include "llvm/IR/ValueSymbolTable.h"
#include "llvm/ProfileData/InstrProf.h"
#include "llvm/ProfileData/SampleProf.h"
#include "llvm/ProfileData/SampleProfReader.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/ErrorOr.h"
#include "llvm/Support/VirtualFileSystem.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/IPO.h"
#include "llvm/Transforms/IPO/ProfiledCallGraph.h"
#include "llvm/Transforms/IPO/SampleContextTracker.h"
#include "llvm/Transforms/IPO/SampleProfileProbe.h"
#include "llvm/Transforms/Instrumentation.h"
#include "llvm/Transforms/Utils/CallPromotionUtils.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/MisExpect.h"
#include "llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h"
#include "llvm/Transforms/Utils/SampleProfileLoaderBaseUtil.h"
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <functional>
#include <limits>
#include <map>
#include <memory>
#include <queue>
#include <string>
#include <system_error>
#include <utility>
#include <vector>

using namespace llvm;
using namespace sampleprof;
using namespace llvm::sampleprofutil;
using ProfileCount = Function::ProfileCount;
#define DEBUG_TYPE "sample-profile"
#define CSINLINE_DEBUG DEBUG_TYPE "-inline"

STATISTIC(NumCSInlined,
          "Number of functions inlined with context sensitive profile");
STATISTIC(NumCSNotInlined,
          "Number of functions not inlined with context sensitive profile");
STATISTIC(NumMismatchedProfile,
          "Number of functions with CFG mismatched profile");
STATISTIC(NumMatchedProfile, "Number of functions with CFG matched profile");
STATISTIC(NumDuplicatedInlinesite,
          "Number of inlined callsites with a partial distribution factor");

STATISTIC(NumCSInlinedHitMinLimit,
          "Number of functions with FDO inline stopped due to min size limit");
STATISTIC(NumCSInlinedHitMaxLimit,
          "Number of functions with FDO inline stopped due to max size limit");
STATISTIC(
    NumCSInlinedHitGrowthLimit,
    "Number of functions with FDO inline stopped due to growth size limit");

// Command line option to specify the file to read samples from. This is
// mainly used for debugging.
static cl::opt<std::string> SampleProfileFile(
    "sample-profile-file", cl::init(""), cl::value_desc("filename"),
    cl::desc("Profile file loaded by -sample-profile"), cl::Hidden);

// The named file contains a set of transformations that may have been applied
// to the symbol names between the program from which the sample data was
// collected and the current program's symbols.
static cl::opt<std::string> SampleProfileRemappingFile(
    "sample-profile-remapping-file", cl::init(""), cl::value_desc("filename"),
    cl::desc("Profile remapping file loaded by -sample-profile"), cl::Hidden);

static cl::opt<bool> SalvageStaleProfile(
    "salvage-stale-profile", cl::Hidden, cl::init(false),
    cl::desc("Salvage stale profile by fuzzy matching and use the remapped "
             "location for sample profile query."));

static cl::opt<bool> ReportProfileStaleness(
    "report-profile-staleness", cl::Hidden, cl::init(false),
    cl::desc("Compute and report stale profile statistical metrics."));

static cl::opt<bool> PersistProfileStaleness(
    "persist-profile-staleness", cl::Hidden, cl::init(false),
    cl::desc("Compute stale profile statistical metrics and write it into the "
             "native object file(.llvm_stats section)."));

static cl::opt<bool> ProfileSampleAccurate(
    "profile-sample-accurate", cl::Hidden, cl::init(false),
    cl::desc("If the sample profile is accurate, we will mark all un-sampled "
             "callsite and function as having 0 samples. Otherwise, treat "
             "un-sampled callsites and functions conservatively as unknown. "));

static cl::opt<bool> ProfileSampleBlockAccurate(
    "profile-sample-block-accurate", cl::Hidden, cl::init(false),
    cl::desc("If the sample profile is accurate, we will mark all un-sampled "
             "branches and calls as having 0 samples. Otherwise, treat "
             "them conservatively as unknown. "));

static cl::opt<bool> ProfileAccurateForSymsInList(
    "profile-accurate-for-symsinlist", cl::Hidden, cl::init(true),
    cl::desc("For symbols in profile symbol list, regard their profiles to "
             "be accurate. It may be overriden by profile-sample-accurate. "));

static cl::opt<bool> ProfileMergeInlinee(
    "sample-profile-merge-inlinee", cl::Hidden, cl::init(true),
    cl::desc("Merge past inlinee's profile to outline version if sample "
             "profile loader decided not to inline a call site. It will "
             "only be enabled when top-down order of profile loading is "
             "enabled. "));

static cl::opt<bool> ProfileTopDownLoad(
    "sample-profile-top-down-load", cl::Hidden, cl::init(true),
    cl::desc("Do profile annotation and inlining for functions in top-down "
             "order of call graph during sample profile loading. It only "
             "works for new pass manager. "));

static cl::opt<bool>
    UseProfiledCallGraph("use-profiled-call-graph", cl::init(true), cl::Hidden,
                         cl::desc("Process functions in a top-down order "
                                  "defined by the profiled call graph when "
                                  "-sample-profile-top-down-load is on."));

static cl::opt<bool> ProfileSizeInline(
    "sample-profile-inline-size", cl::Hidden, cl::init(false),
    cl::desc("Inline cold call sites in profile loader if it's beneficial "
             "for code size."));

// Since profiles are consumed by many passes, turning on this option has
// side effects. For instance, pre-link SCC inliner would see merged profiles
// and inline the hot functions (that are skipped in this pass).
static cl::opt<bool> DisableSampleLoaderInlining(
    "disable-sample-loader-inlining", cl::Hidden, cl::init(false),
    cl::desc("If true, artifically skip inline transformation in sample-loader "
             "pass, and merge (or scale) profiles (as configured by "
             "--sample-profile-merge-inlinee)."));

namespace llvm {
cl::opt<bool>
    SortProfiledSCC("sort-profiled-scc-member", cl::init(true), cl::Hidden,
                    cl::desc("Sort profiled recursion by edge weights."));

cl::opt<int> ProfileInlineGrowthLimit(
    "sample-profile-inline-growth-limit", cl::Hidden, cl::init(12),
    cl::desc("The size growth ratio limit for proirity-based sample profile "
             "loader inlining."));

cl::opt<int> ProfileInlineLimitMin(
    "sample-profile-inline-limit-min", cl::Hidden, cl::init(100),
    cl::desc("The lower bound of size growth limit for "
             "proirity-based sample profile loader inlining."));

cl::opt<int> ProfileInlineLimitMax(
    "sample-profile-inline-limit-max", cl::Hidden, cl::init(10000),
    cl::desc("The upper bound of size growth limit for "
             "proirity-based sample profile loader inlining."));

cl::opt<int> SampleHotCallSiteThreshold(
    "sample-profile-hot-inline-threshold", cl::Hidden, cl::init(3000),
    cl::desc("Hot callsite threshold for proirity-based sample profile loader "
             "inlining."));

cl::opt<int> SampleColdCallSiteThreshold(
    "sample-profile-cold-inline-threshold", cl::Hidden, cl::init(45),
    cl::desc("Threshold for inlining cold callsites"));
} // namespace llvm

static cl::opt<unsigned> ProfileICPRelativeHotness(
    "sample-profile-icp-relative-hotness", cl::Hidden, cl::init(25),
    cl::desc(
        "Relative hotness percentage threshold for indirect "
        "call promotion in proirity-based sample profile loader inlining."));

static cl::opt<unsigned> ProfileICPRelativeHotnessSkip(
    "sample-profile-icp-relative-hotness-skip", cl::Hidden, cl::init(1),
    cl::desc(
        "Skip relative hotness check for ICP up to given number of targets."));

static cl::opt<bool> CallsitePrioritizedInline(
    "sample-profile-prioritized-inline", cl::Hidden,

    cl::desc("Use call site prioritized inlining for sample profile loader."
             "Currently only CSSPGO is supported."));

static cl::opt<bool> UsePreInlinerDecision(
    "sample-profile-use-preinliner", cl::Hidden,

    cl::desc("Use the preinliner decisions stored in profile context."));

static cl::opt<bool> AllowRecursiveInline(
    "sample-profile-recursive-inline", cl::Hidden,

    cl::desc("Allow sample loader inliner to inline recursive calls."));

static cl::opt<std::string> ProfileInlineReplayFile(
    "sample-profile-inline-replay", cl::init(""), cl::value_desc("filename"),
    cl::desc(
        "Optimization remarks file containing inline remarks to be replayed "
        "by inlining from sample profile loader."),
    cl::Hidden);

static cl::opt<ReplayInlinerSettings::Scope> ProfileInlineReplayScope(
    "sample-profile-inline-replay-scope",
    cl::init(ReplayInlinerSettings::Scope::Function),
    cl::values(clEnumValN(ReplayInlinerSettings::Scope::Function, "Function",
                          "Replay on functions that have remarks associated "
                          "with them (default)"),
               clEnumValN(ReplayInlinerSettings::Scope::Module, "Module",
                          "Replay on the entire module")),
    cl::desc("Whether inline replay should be applied to the entire "
             "Module or just the Functions (default) that are present as "
             "callers in remarks during sample profile inlining."),
    cl::Hidden);

static cl::opt<ReplayInlinerSettings::Fallback> ProfileInlineReplayFallback(
    "sample-profile-inline-replay-fallback",
    cl::init(ReplayInlinerSettings::Fallback::Original),
    cl::values(
        clEnumValN(
            ReplayInlinerSettings::Fallback::Original, "Original",
            "All decisions not in replay send to original advisor (default)"),
        clEnumValN(ReplayInlinerSettings::Fallback::AlwaysInline,
                   "AlwaysInline", "All decisions not in replay are inlined"),
        clEnumValN(ReplayInlinerSettings::Fallback::NeverInline, "NeverInline",
                   "All decisions not in replay are not inlined")),
    cl::desc("How sample profile inline replay treats sites that don't come "
             "from the replay. Original: defers to original advisor, "
             "AlwaysInline: inline all sites not in replay, NeverInline: "
             "inline no sites not in replay"),
    cl::Hidden);

static cl::opt<CallSiteFormat::Format> ProfileInlineReplayFormat(
    "sample-profile-inline-replay-format",
    cl::init(CallSiteFormat::Format::LineColumnDiscriminator),
    cl::values(
        clEnumValN(CallSiteFormat::Format::Line, "Line", "<Line Number>"),
        clEnumValN(CallSiteFormat::Format::LineColumn, "LineColumn",
                   "<Line Number>:<Column Number>"),
        clEnumValN(CallSiteFormat::Format::LineDiscriminator,
                   "LineDiscriminator", "<Line Number>.<Discriminator>"),
        clEnumValN(CallSiteFormat::Format::LineColumnDiscriminator,
                   "LineColumnDiscriminator",
                   "<Line Number>:<Column Number>.<Discriminator> (default)")),
    cl::desc("How sample profile inline replay file is formatted"), cl::Hidden);

static cl::opt<unsigned>
    MaxNumPromotions("sample-profile-icp-max-prom", cl::init(3), cl::Hidden,
                     cl::desc("Max number of promotions for a single indirect "
                              "call callsite in sample profile loader"));

static cl::opt<bool> OverwriteExistingWeights(
    "overwrite-existing-weights", cl::Hidden, cl::init(false),
    cl::desc("Ignore existing branch weights on IR and always overwrite."));

static cl::opt<bool> AnnotateSampleProfileInlinePhase(
    "annotate-sample-profile-inline-phase", cl::Hidden, cl::init(false),
    cl::desc("Annotate LTO phase (prelink / postlink), or main (no LTO) for "
             "sample-profile inline pass name."));

namespace llvm {
extern cl::opt<bool> EnableExtTspBlockPlacement;
}

namespace {

using BlockWeightMap = DenseMap<const BasicBlock *, uint64_t>;
using EquivalenceClassMap = DenseMap<const BasicBlock *, const BasicBlock *>;
using Edge = std::pair<const BasicBlock *, const BasicBlock *>;
using EdgeWeightMap = DenseMap<Edge, uint64_t>;
using BlockEdgeMap =
    DenseMap<const BasicBlock *, SmallVector<const BasicBlock *, 8>>;

class GUIDToFuncNameMapper {
public:
  GUIDToFuncNameMapper(Module &M, SampleProfileReader &Reader,
                       DenseMap<uint64_t, StringRef> &GUIDToFuncNameMap)
      : CurrentReader(Reader), CurrentModule(M),
        CurrentGUIDToFuncNameMap(GUIDToFuncNameMap) {
    if (!CurrentReader.useMD5())
      return;

    for (const auto &F : CurrentModule) {
      StringRef OrigName = F.getName();
      CurrentGUIDToFuncNameMap.insert(
          {Function::getGUID(OrigName), OrigName});

      // Local to global var promotion used by optimization like thinlto
      // will rename the var and add suffix like ".llvm.xxx" to the
      // original local name. In sample profile, the suffixes of function
      // names are all stripped. Since it is possible that the mapper is
      // built in post-thin-link phase and var promotion has been done,
      // we need to add the substring of function name without the suffix
      // into the GUIDToFuncNameMap.
      StringRef CanonName = FunctionSamples::getCanonicalFnName(F);
      if (CanonName != OrigName)
        CurrentGUIDToFuncNameMap.insert(
            {Function::getGUID(CanonName), CanonName});
    }

    // Update GUIDToFuncNameMap for each function including inlinees.
    SetGUIDToFuncNameMapForAll(&CurrentGUIDToFuncNameMap);
  }

  ~GUIDToFuncNameMapper() {
    if (!CurrentReader.useMD5())
      return;

    CurrentGUIDToFuncNameMap.clear();

    // Reset GUIDToFuncNameMap for of each function as they're no
    // longer valid at this point.
    SetGUIDToFuncNameMapForAll(nullptr);
  }

private:
  void SetGUIDToFuncNameMapForAll(DenseMap<uint64_t, StringRef> *Map) {
    std::queue<FunctionSamples *> FSToUpdate;
    for (auto &IFS : CurrentReader.getProfiles()) {
      FSToUpdate.push(&IFS.second);
    }

    while (!FSToUpdate.empty()) {
      FunctionSamples *FS = FSToUpdate.front();
      FSToUpdate.pop();
      FS->GUIDToFuncNameMap = Map;
      for (const auto &ICS : FS->getCallsiteSamples()) {
        const FunctionSamplesMap &FSMap = ICS.second;
        for (const auto &IFS : FSMap) {
          FunctionSamples &FS = const_cast<FunctionSamples &>(IFS.second);
          FSToUpdate.push(&FS);
        }
      }
    }
  }

  SampleProfileReader &CurrentReader;
  Module &CurrentModule;
  DenseMap<uint64_t, StringRef> &CurrentGUIDToFuncNameMap;
};

// Inline candidate used by iterative callsite prioritized inliner
struct InlineCandidate {
  CallBase *CallInstr;
  const FunctionSamples *CalleeSamples;
  // Prorated callsite count, which will be used to guide inlining. For example,
  // if a callsite is duplicated in LTO prelink, then in LTO postlink the two
  // copies will get their own distribution factors and their prorated counts
  // will be used to decide if they should be inlined independently.
  uint64_t CallsiteCount;
  // Call site distribution factor to prorate the profile samples for a
  // duplicated callsite. Default value is 1.0.
  float CallsiteDistribution;
};

// Inline candidate comparer using call site weight
struct CandidateComparer {
  bool operator()(const InlineCandidate &LHS, const InlineCandidate &RHS) {
    if (LHS.CallsiteCount != RHS.CallsiteCount)
      return LHS.CallsiteCount < RHS.CallsiteCount;

    const FunctionSamples *LCS = LHS.CalleeSamples;
    const FunctionSamples *RCS = RHS.CalleeSamples;
    assert(LCS && RCS && "Expect non-null FunctionSamples");

    // Tie breaker using number of samples try to favor smaller functions first
    if (LCS->getBodySamples().size() != RCS->getBodySamples().size())
      return LCS->getBodySamples().size() > RCS->getBodySamples().size();

    // Tie breaker using GUID so we have stable/deterministic inlining order
    return LCS->getGUID() < RCS->getGUID();
  }
};

using CandidateQueue =
    PriorityQueue<InlineCandidate, std::vector<InlineCandidate>,
                  CandidateComparer>;

// Sample profile matching - fuzzy match.
class SampleProfileMatcher {
  Module &M;
  SampleProfileReader &Reader;
  const PseudoProbeManager *ProbeManager;
  SampleProfileMap FlattenedProfiles;
  // For each function, the matcher generates a map, of which each entry is a
  // mapping from the source location of current build to the source location in
  // the profile.
  StringMap<LocToLocMap> FuncMappings;

  // Profile mismatching statstics.
  uint64_t TotalProfiledCallsites = 0;
  uint64_t NumMismatchedCallsites = 0;
  uint64_t MismatchedCallsiteSamples = 0;
  uint64_t TotalCallsiteSamples = 0;
  uint64_t TotalProfiledFunc = 0;
  uint64_t NumMismatchedFuncHash = 0;
  uint64_t MismatchedFuncHashSamples = 0;
  uint64_t TotalFuncHashSamples = 0;

  // A dummy name for unknown indirect callee, used to differentiate from a
  // non-call instruction that also has an empty callee name.
  static constexpr const char *UnknownIndirectCallee =
      "unknown.indirect.callee";

public:
  SampleProfileMatcher(Module &M, SampleProfileReader &Reader,
                       const PseudoProbeManager *ProbeManager)
      : M(M), Reader(Reader), ProbeManager(ProbeManager){};
  void runOnModule();

private:
  FunctionSamples *getFlattenedSamplesFor(const Function &F) {
    StringRef CanonFName = FunctionSamples::getCanonicalFnName(F);
    auto It = FlattenedProfiles.find(FunctionId(CanonFName));
    if (It != FlattenedProfiles.end())
      return &It->second;
    return nullptr;
  }
  void runOnFunction(const Function &F);
  void findIRAnchors(const Function &F,
                     std::map<LineLocation, StringRef> &IRAnchors);
  void findProfileAnchors(
      const FunctionSamples &FS,
      std::map<LineLocation, std::unordered_set<FunctionId>>
          &ProfileAnchors);
  void countMismatchedSamples(const FunctionSamples &FS);
  void countProfileMismatches(
      const Function &F, const FunctionSamples &FS,
      const std::map<LineLocation, StringRef> &IRAnchors,
      const std::map<LineLocation, std::unordered_set<FunctionId>>
          &ProfileAnchors);
  void countProfileCallsiteMismatches(
      const FunctionSamples &FS,
      const std::map<LineLocation, StringRef> &IRAnchors,
      const std::map<LineLocation, std::unordered_set<FunctionId>>
          &ProfileAnchors,
      uint64_t &FuncMismatchedCallsites, uint64_t &FuncProfiledCallsites);
  LocToLocMap &getIRToProfileLocationMap(const Function &F) {
    auto Ret = FuncMappings.try_emplace(
        FunctionSamples::getCanonicalFnName(F.getName()), LocToLocMap());
    return Ret.first->second;
  }
  void distributeIRToProfileLocationMap();
  void distributeIRToProfileLocationMap(FunctionSamples &FS);
  void runStaleProfileMatching(
      const Function &F, const std::map<LineLocation, StringRef> &IRAnchors,
      const std::map<LineLocation, std::unordered_set<FunctionId>>
          &ProfileAnchors,
      LocToLocMap &IRToProfileLocationMap);
};

/// Sample profile pass.
///
/// This pass reads profile data from the file specified by
/// -sample-profile-file and annotates every affected function with the
/// profile information found in that file.
class SampleProfileLoader final : public SampleProfileLoaderBaseImpl<Function> {
public:
  SampleProfileLoader(
      StringRef Name, StringRef RemapName, ThinOrFullLTOPhase LTOPhase,
      IntrusiveRefCntPtr<vfs::FileSystem> FS,
      std::function<AssumptionCache &(Function &)> GetAssumptionCache,
      std::function<TargetTransformInfo &(Function &)> GetTargetTransformInfo,
      std::function<const TargetLibraryInfo &(Function &)> GetTLI)
      : SampleProfileLoaderBaseImpl(std::string(Name), std::string(RemapName),
                                    std::move(FS)),
        GetAC(std::move(GetAssumptionCache)),
        GetTTI(std::move(GetTargetTransformInfo)), GetTLI(std::move(GetTLI)),
        LTOPhase(LTOPhase),
        AnnotatedPassName(AnnotateSampleProfileInlinePhase
                              ? llvm::AnnotateInlinePassName(InlineContext{
                                    LTOPhase, InlinePass::SampleProfileInliner})
                              : CSINLINE_DEBUG) {}

  bool doInitialization(Module &M, FunctionAnalysisManager *FAM = nullptr);
  bool runOnModule(Module &M, ModuleAnalysisManager *AM,
                   ProfileSummaryInfo *_PSI, LazyCallGraph &CG);

protected:
  bool runOnFunction(Function &F, ModuleAnalysisManager *AM);
  bool emitAnnotations(Function &F);
  ErrorOr<uint64_t> getInstWeight(const Instruction &I) override;
  const FunctionSamples *findCalleeFunctionSamples(const CallBase &I) const;
  const FunctionSamples *
  findFunctionSamples(const Instruction &I) const override;
  std::vector<const FunctionSamples *>
  findIndirectCallFunctionSamples(const Instruction &I, uint64_t &Sum) const;
  void findExternalInlineCandidate(CallBase *CB, const FunctionSamples *Samples,
                                   DenseSet<GlobalValue::GUID> &InlinedGUIDs,
                                   uint64_t Threshold);
  // Attempt to promote indirect call and also inline the promoted call
  bool tryPromoteAndInlineCandidate(
      Function &F, InlineCandidate &Candidate, uint64_t SumOrigin,
      uint64_t &Sum, SmallVector<CallBase *, 8> *InlinedCallSites = nullptr);

  bool inlineHotFunctions(Function &F,
                          DenseSet<GlobalValue::GUID> &InlinedGUIDs);
  std::optional<InlineCost> getExternalInlineAdvisorCost(CallBase &CB);
  bool getExternalInlineAdvisorShouldInline(CallBase &CB);
  InlineCost shouldInlineCandidate(InlineCandidate &Candidate);
  bool getInlineCandidate(InlineCandidate *NewCandidate, CallBase *CB);
  bool
  tryInlineCandidate(InlineCandidate &Candidate,
                     SmallVector<CallBase *, 8> *InlinedCallSites = nullptr);
  bool
  inlineHotFunctionsWithPriority(Function &F,
                                 DenseSet<GlobalValue::GUID> &InlinedGUIDs);
  // Inline cold/small functions in addition to hot ones
  bool shouldInlineColdCallee(CallBase &CallInst);
  void emitOptimizationRemarksForInlineCandidates(
      const SmallVectorImpl<CallBase *> &Candidates, const Function &F,
      bool Hot);
  void promoteMergeNotInlinedContextSamples(
      MapVector<CallBase *, const FunctionSamples *> NonInlinedCallSites,
      const Function &F);
  std::vector<Function *> buildFunctionOrder(Module &M, LazyCallGraph &CG);
  std::unique_ptr<ProfiledCallGraph> buildProfiledCallGraph(Module &M);
  void generateMDProfMetadata(Function &F);

  /// Map from function name to Function *. Used to find the function from
  /// the function name. If the function name contains suffix, additional
  /// entry is added to map from the stripped name to the function if there
  /// is one-to-one mapping.
  HashKeyMap<std::unordered_map, FunctionId, Function *> SymbolMap;

  std::function<AssumptionCache &(Function &)> GetAC;
  std::function<TargetTransformInfo &(Function &)> GetTTI;
  std::function<const TargetLibraryInfo &(Function &)> GetTLI;

  /// Profile tracker for different context.
  std::unique_ptr<SampleContextTracker> ContextTracker;

  /// Flag indicating which LTO/ThinLTO phase the pass is invoked in.
  ///
  /// We need to know the LTO phase because for example in ThinLTOPrelink
  /// phase, in annotation, we should not promote indirect calls. Instead,
  /// we will mark GUIDs that needs to be annotated to the function.
  const ThinOrFullLTOPhase LTOPhase;
  const std::string AnnotatedPassName;

  /// Profle Symbol list tells whether a function name appears in the binary
  /// used to generate the current profile.
  std::unique_ptr<ProfileSymbolList> PSL;

  /// Total number of samples collected in this profile.
  ///
  /// This is the sum of all the samples collected in all the functions executed
  /// at runtime.
  uint64_t TotalCollectedSamples = 0;

  // Information recorded when we declined to inline a call site
  // because we have determined it is too cold is accumulated for
  // each callee function. Initially this is just the entry count.
  struct NotInlinedProfileInfo {
    uint64_t entryCount;
  };
  DenseMap<Function *, NotInlinedProfileInfo> notInlinedCallInfo;

  // GUIDToFuncNameMap saves the mapping from GUID to the symbol name, for
  // all the function symbols defined or declared in current module.
  DenseMap<uint64_t, StringRef> GUIDToFuncNameMap;

  // All the Names used in FunctionSamples including outline function
  // names, inline instance names and call target names.
  StringSet<> NamesInProfile;
  // MD5 version of NamesInProfile. Either NamesInProfile or GUIDsInProfile is
  // populated, depends on whether the profile uses MD5. Because the name table
  // generally contains several magnitude more entries than the number of
  // functions, we do not want to convert all names from one form to another.
  llvm::DenseSet<uint64_t> GUIDsInProfile;

  // For symbol in profile symbol list, whether to regard their profiles
  // to be accurate. It is mainly decided by existance of profile symbol
  // list and -profile-accurate-for-symsinlist flag, but it can be
  // overriden by -profile-sample-accurate or profile-sample-accurate
  // attribute.
  bool ProfAccForSymsInList;

  // External inline advisor used to replay inline decision from remarks.
  std::unique_ptr<InlineAdvisor> ExternalInlineAdvisor;

  // A helper to implement the sample profile matching algorithm.
  std::unique_ptr<SampleProfileMatcher> MatchingManager;

private:
  const char *getAnnotatedRemarkPassName() const {
    return AnnotatedPassName.c_str();
  }
};
} // end anonymous namespace

namespace llvm {
template <>
inline bool SampleProfileInference<Function>::isExit(const BasicBlock *BB) {
  return succ_empty(BB);
}

template <>
inline void SampleProfileInference<Function>::findUnlikelyJumps(
    const std::vector<const BasicBlockT *> &BasicBlocks,
    BlockEdgeMap &Successors, FlowFunction &Func) {
  for (auto &Jump : Func.Jumps) {
    const auto *BB = BasicBlocks[Jump.Source];
    const auto *Succ = BasicBlocks[Jump.Target];
    const Instruction *TI = BB->getTerminator();
    // Check if a block ends with InvokeInst and mark non-taken branch unlikely.
    // In that case block Succ should be a landing pad
    if (Successors[BB].size() == 2 && Successors[BB].back() == Succ) {
      if (isa<InvokeInst>(TI)) {
        Jump.IsUnlikely = true;
      }
    }
    const Instruction *SuccTI = Succ->getTerminator();
    // Check if the target block contains UnreachableInst and mark it unlikely
    if (SuccTI->getNumSuccessors() == 0) {
      if (isa<UnreachableInst>(SuccTI)) {
        Jump.IsUnlikely = true;
      }
    }
  }
}

template <>
void SampleProfileLoaderBaseImpl<Function>::computeDominanceAndLoopInfo(
    Function &F) {
  DT.reset(new DominatorTree);
  DT->recalculate(F);

  PDT.reset(new PostDominatorTree(F));

  LI.reset(new LoopInfo);
  LI->analyze(*DT);
}
} // namespace llvm

ErrorOr<uint64_t> SampleProfileLoader::getInstWeight(const Instruction &Inst) {
  if (FunctionSamples::ProfileIsProbeBased)
    return getProbeWeight(Inst);

  const DebugLoc &DLoc = Inst.getDebugLoc();
  if (!DLoc)
    return std::error_code();

  // Ignore all intrinsics, phinodes and branch instructions.
  // Branch and phinodes instruction usually contains debug info from sources
  // outside of the residing basic block, thus we ignore them during annotation.
  if (isa<BranchInst>(Inst) || isa<IntrinsicInst>(Inst) || isa<PHINode>(Inst))
    return std::error_code();

  // For non-CS profile, if a direct call/invoke instruction is inlined in
  // profile (findCalleeFunctionSamples returns non-empty result), but not
  // inlined here, it means that the inlined callsite has no sample, thus the
  // call instruction should have 0 count.
  // For CS profile, the callsite count of previously inlined callees is
  // populated with the entry count of the callees.
  if (!FunctionSamples::ProfileIsCS)
    if (const auto *CB = dyn_cast<CallBase>(&Inst))
      if (!CB->isIndirectCall() && findCalleeFunctionSamples(*CB))
        return 0;

  return getInstWeightImpl(Inst);
}

/// Get the FunctionSamples for a call instruction.
///
/// The FunctionSamples of a call/invoke instruction \p Inst is the inlined
/// instance in which that call instruction is calling to. It contains
/// all samples that resides in the inlined instance. We first find the
/// inlined instance in which the call instruction is from, then we
/// traverse its children to find the callsite with the matching
/// location.
///
/// \param Inst Call/Invoke instruction to query.
///
/// \returns The FunctionSamples pointer to the inlined instance.
const FunctionSamples *
SampleProfileLoader::findCalleeFunctionSamples(const CallBase &Inst) const {
  const DILocation *DIL = Inst.getDebugLoc();
  if (!DIL) {
    return nullptr;
  }

  StringRef CalleeName;
  if (Function *Callee = Inst.getCalledFunction())
    CalleeName = Callee->getName();

  if (FunctionSamples::ProfileIsCS)
    return ContextTracker->getCalleeContextSamplesFor(Inst, CalleeName);

  const FunctionSamples *FS = findFunctionSamples(Inst);
  if (FS == nullptr)
    return nullptr;

  return FS->findFunctionSamplesAt(FunctionSamples::getCallSiteIdentifier(DIL),
                                   CalleeName, Reader->getRemapper());
}

/// Returns a vector of FunctionSamples that are the indirect call targets
/// of \p Inst. The vector is sorted by the total number of samples. Stores
/// the total call count of the indirect call in \p Sum.
std::vector<const FunctionSamples *>
SampleProfileLoader::findIndirectCallFunctionSamples(
    const Instruction &Inst, uint64_t &Sum) const {
  const DILocation *DIL = Inst.getDebugLoc();
  std::vector<const FunctionSamples *> R;

  if (!DIL) {
    return R;
  }

  auto FSCompare = [](const FunctionSamples *L, const FunctionSamples *R) {
    assert(L && R && "Expect non-null FunctionSamples");
    if (L->getHeadSamplesEstimate() != R->getHeadSamplesEstimate())
      return L->getHeadSamplesEstimate() > R->getHeadSamplesEstimate();
    return L->getGUID() < R->getGUID();
  };

  if (FunctionSamples::ProfileIsCS) {
    auto CalleeSamples =
        ContextTracker->getIndirectCalleeContextSamplesFor(DIL);
    if (CalleeSamples.empty())
      return R;

    // For CSSPGO, we only use target context profile's entry count
    // as that already includes both inlined callee and non-inlined ones..
    Sum = 0;
    for (const auto *const FS : CalleeSamples) {
      Sum += FS->getHeadSamplesEstimate();
      R.push_back(FS);
    }
    llvm::sort(R, FSCompare);
    return R;
  }

  const FunctionSamples *FS = findFunctionSamples(Inst);
  if (FS == nullptr)
    return R;

  auto CallSite = FunctionSamples::getCallSiteIdentifier(DIL);
  auto T = FS->findCallTargetMapAt(CallSite);
  Sum = 0;
  if (T)
    for (const auto &T_C : T.get())
      Sum += T_C.second;
  if (const FunctionSamplesMap *M = FS->findFunctionSamplesMapAt(CallSite)) {
    if (M->empty())
      return R;
    for (const auto &NameFS : *M) {
      Sum += NameFS.second.getHeadSamplesEstimate();
      R.push_back(&NameFS.second);
    }
    llvm::sort(R, FSCompare);
  }
  return R;
}

const FunctionSamples *
SampleProfileLoader::findFunctionSamples(const Instruction &Inst) const {
  if (FunctionSamples::ProfileIsProbeBased) {
    std::optional<PseudoProbe> Probe = extractProbe(Inst);
    if (!Probe)
      return nullptr;
  }

  const DILocation *DIL = Inst.getDebugLoc();
  if (!DIL)
    return Samples;

  auto it = DILocation2SampleMap.try_emplace(DIL,nullptr);
  if (it.second) {
    if (FunctionSamples::ProfileIsCS)
      it.first->second = ContextTracker->getContextSamplesFor(DIL);
    else
      it.first->second =
          Samples->findFunctionSamples(DIL, Reader->getRemapper());
  }
  return it.first->second;
}

/// Check whether the indirect call promotion history of \p Inst allows
/// the promotion for \p Candidate.
/// If the profile count for the promotion candidate \p Candidate is
/// NOMORE_ICP_MAGICNUM, it means \p Candidate has already been promoted
/// for \p Inst. If we already have at least MaxNumPromotions
/// NOMORE_ICP_MAGICNUM count values in the value profile of \p Inst, we
/// cannot promote for \p Inst anymore.
static bool doesHistoryAllowICP(const Instruction &Inst, StringRef Candidate) {
  uint32_t NumVals = 0;
  uint64_t TotalCount = 0;
  std::unique_ptr<InstrProfValueData[]> ValueData =
      std::make_unique<InstrProfValueData[]>(MaxNumPromotions);
  bool Valid =
      getValueProfDataFromInst(Inst, IPVK_IndirectCallTarget, MaxNumPromotions,
                               ValueData.get(), NumVals, TotalCount, true);
  // No valid value profile so no promoted targets have been recorded
  // before. Ok to do ICP.
  if (!Valid)
    return true;

  unsigned NumPromoted = 0;
  for (uint32_t I = 0; I < NumVals; I++) {
    if (ValueData[I].Count != NOMORE_ICP_MAGICNUM)
      continue;

    // If the promotion candidate has NOMORE_ICP_MAGICNUM count in the
    // metadata, it means the candidate has been promoted for this
    // indirect call.
    if (ValueData[I].Value == Function::getGUID(Candidate))
      return false;
    NumPromoted++;
    // If already have MaxNumPromotions promotion, don't do it anymore.
    if (NumPromoted == MaxNumPromotions)
      return false;
  }
  return true;
}

/// Update indirect call target profile metadata for \p Inst.
/// Usually \p Sum is the sum of counts of all the targets for \p Inst.
/// If it is 0, it means updateIDTMetaData is used to mark a
/// certain target to be promoted already. If it is not zero,
/// we expect to use it to update the total count in the value profile.
static void
updateIDTMetaData(Instruction &Inst,
                  const SmallVectorImpl<InstrProfValueData> &CallTargets,
                  uint64_t Sum) {
  // Bail out early if MaxNumPromotions is zero.
  // This prevents allocating an array of zero length below.
  //
  // Note `updateIDTMetaData` is called in two places so check
  // `MaxNumPromotions` inside it.
  if (MaxNumPromotions == 0)
    return;
  uint32_t NumVals = 0;
  // OldSum is the existing total count in the value profile data.
  uint64_t OldSum = 0;
  std::unique_ptr<InstrProfValueData[]> ValueData =
      std::make_unique<InstrProfValueData[]>(MaxNumPromotions);
  bool Valid =
      getValueProfDataFromInst(Inst, IPVK_IndirectCallTarget, MaxNumPromotions,
                               ValueData.get(), NumVals, OldSum, true);

  DenseMap<uint64_t, uint64_t> ValueCountMap;
  if (Sum == 0) {
    assert((CallTargets.size() == 1 &&
            CallTargets[0].Count == NOMORE_ICP_MAGICNUM) &&
           "If sum is 0, assume only one element in CallTargets "
           "with count being NOMORE_ICP_MAGICNUM");
    // Initialize ValueCountMap with existing value profile data.
    if (Valid) {
      for (uint32_t I = 0; I < NumVals; I++)
        ValueCountMap[ValueData[I].Value] = ValueData[I].Count;
    }
    auto Pair =
        ValueCountMap.try_emplace(CallTargets[0].Value, CallTargets[0].Count);
    // If the target already exists in value profile, decrease the total
    // count OldSum and reset the target's count to NOMORE_ICP_MAGICNUM.
    if (!Pair.second) {
      OldSum -= Pair.first->second;
      Pair.first->second = NOMORE_ICP_MAGICNUM;
    }
    Sum = OldSum;
  } else {
    // Initialize ValueCountMap with existing NOMORE_ICP_MAGICNUM
    // counts in the value profile.
    if (Valid) {
      for (uint32_t I = 0; I < NumVals; I++) {
        if (ValueData[I].Count == NOMORE_ICP_MAGICNUM)
          ValueCountMap[ValueData[I].Value] = ValueData[I].Count;
      }
    }

    for (const auto &Data : CallTargets) {
      auto Pair = ValueCountMap.try_emplace(Data.Value, Data.Count);
      if (Pair.second)
        continue;
      // The target represented by Data.Value has already been promoted.
      // Keep the count as NOMORE_ICP_MAGICNUM in the profile and decrease
      // Sum by Data.Count.
      assert(Sum >= Data.Count && "Sum should never be less than Data.Count");
      Sum -= Data.Count;
    }
  }

  SmallVector<InstrProfValueData, 8> NewCallTargets;
  for (const auto &ValueCount : ValueCountMap) {
    NewCallTargets.emplace_back(
        InstrProfValueData{ValueCount.first, ValueCount.second});
  }

  llvm::sort(NewCallTargets,
             [](const InstrProfValueData &L, const InstrProfValueData &R) {
               if (L.Count != R.Count)
                 return L.Count > R.Count;
               return L.Value > R.Value;
             });

  uint32_t MaxMDCount =
      std::min(NewCallTargets.size(), static_cast<size_t>(MaxNumPromotions));
  annotateValueSite(*Inst.getParent()->getParent()->getParent(), Inst,
                    NewCallTargets, Sum, IPVK_IndirectCallTarget, MaxMDCount);
}

/// Attempt to promote indirect call and also inline the promoted call.
///
/// \param F  Caller function.
/// \param Candidate  ICP and inline candidate.
/// \param SumOrigin  Original sum of target counts for indirect call before
///                   promoting given candidate.
/// \param Sum        Prorated sum of remaining target counts for indirect call
///                   after promoting given candidate.
/// \param InlinedCallSite  Output vector for new call sites exposed after
/// inlining.
bool SampleProfileLoader::tryPromoteAndInlineCandidate(
    Function &F, InlineCandidate &Candidate, uint64_t SumOrigin, uint64_t &Sum,
    SmallVector<CallBase *, 8> *InlinedCallSite) {
  // Bail out early if sample-loader inliner is disabled.
  if (DisableSampleLoaderInlining)
    return false;

  // Bail out early if MaxNumPromotions is zero.
  // This prevents allocating an array of zero length in callees below.
  if (MaxNumPromotions == 0)
    return false;
  auto CalleeFunctionName = Candidate.CalleeSamples->getFunction();
  auto R = SymbolMap.find(CalleeFunctionName);
  if (R == SymbolMap.end() || !R->second)
    return false;

  auto &CI = *Candidate.CallInstr;
  if (!doesHistoryAllowICP(CI, R->second->getName()))
    return false;

  const char *Reason = "Callee function not available";
  // R->getValue() != &F is to prevent promoting a recursive call.
  // If it is a recursive call, we do not inline it as it could bloat
  // the code exponentially. There is way to better handle this, e.g.
  // clone the caller first, and inline the cloned caller if it is
  // recursive. As llvm does not inline recursive calls, we will
  // simply ignore it instead of handling it explicitly.
  if (!R->second->isDeclaration() && R->second->getSubprogram() &&
      R->second->hasFnAttribute("use-sample-profile") &&
      R->second != &F && isLegalToPromote(CI, R->second, &Reason)) {
    // For promoted target, set its value with NOMORE_ICP_MAGICNUM count
    // in the value profile metadata so the target won't be promoted again.
    SmallVector<InstrProfValueData, 1> SortedCallTargets = {InstrProfValueData{
        Function::getGUID(R->second->getName()), NOMORE_ICP_MAGICNUM}};
    updateIDTMetaData(CI, SortedCallTargets, 0);

    auto *DI = &pgo::promoteIndirectCall(
        CI, R->second, Candidate.CallsiteCount, Sum, false, ORE);
    if (DI) {
      Sum -= Candidate.CallsiteCount;
      // Do not prorate the indirect callsite distribution since the original
      // distribution will be used to scale down non-promoted profile target
      // counts later. By doing this we lose track of the real callsite count
      // for the leftover indirect callsite as a trade off for accurate call
      // target counts.
      // TODO: Ideally we would have two separate factors, one for call site
      // counts and one is used to prorate call target counts.
      // Do not update the promoted direct callsite distribution at this
      // point since the original distribution combined with the callee profile
      // will be used to prorate callsites from the callee if inlined. Once not
      // inlined, the direct callsite distribution should be prorated so that
      // the it will reflect the real callsite counts.
      Candidate.CallInstr = DI;
      if (isa<CallInst>(DI) || isa<InvokeInst>(DI)) {
        bool Inlined = tryInlineCandidate(Candidate, InlinedCallSite);
        if (!Inlined) {
          // Prorate the direct callsite distribution so that it reflects real
          // callsite counts.
          setProbeDistributionFactor(
              *DI, static_cast<float>(Candidate.CallsiteCount) / SumOrigin);
        }
        return Inlined;
      }
    }
  } else {
    LLVM_DEBUG(dbgs() << "\nFailed to promote indirect call to "
                      << FunctionSamples::getCanonicalFnName(
                             Candidate.CallInstr->getName())<< " because "
                      << Reason << "\n");
  }
  return false;
}

bool SampleProfileLoader::shouldInlineColdCallee(CallBase &CallInst) {
  if (!ProfileSizeInline)
    return false;

  Function *Callee = CallInst.getCalledFunction();
  if (Callee == nullptr)
    return false;

  InlineCost Cost = getInlineCost(CallInst, getInlineParams(), GetTTI(*Callee),
                                  GetAC, GetTLI);

  if (Cost.isNever())
    return false;

  if (Cost.isAlways())
    return true;

  return Cost.getCost() <= SampleColdCallSiteThreshold;
}

void SampleProfileLoader::emitOptimizationRemarksForInlineCandidates(
    const SmallVectorImpl<CallBase *> &Candidates, const Function &F,
    bool Hot) {
  for (auto *I : Candidates) {
    Function *CalledFunction = I->getCalledFunction();
    if (CalledFunction) {
      ORE->emit(OptimizationRemarkAnalysis(getAnnotatedRemarkPassName(),
                                           "InlineAttempt", I->getDebugLoc(),
                                           I->getParent())
                << "previous inlining reattempted for "
                << (Hot ? "hotness: '" : "size: '")
                << ore::NV("Callee", CalledFunction) << "' into '"
                << ore::NV("Caller", &F) << "'");
    }
  }
}

void SampleProfileLoader::findExternalInlineCandidate(
    CallBase *CB, const FunctionSamples *Samples,
    DenseSet<GlobalValue::GUID> &InlinedGUIDs, uint64_t Threshold) {

  // If ExternalInlineAdvisor(ReplayInlineAdvisor) wants to inline an external
  // function make sure it's imported
  if (CB && getExternalInlineAdvisorShouldInline(*CB)) {
    // Samples may not exist for replayed function, if so
    // just add the direct GUID and move on
    if (!Samples) {
      InlinedGUIDs.insert(
          Function::getGUID(CB->getCalledFunction()->getName()));
      return;
    }
    // Otherwise, drop the threshold to import everything that we can
    Threshold = 0;
  }

  // In some rare cases, call instruction could be changed after being pushed
  // into inline candidate queue, this is because earlier inlining may expose
  // constant propagation which can change indirect call to direct call. When
  // this happens, we may fail to find matching function samples for the
  // candidate later, even if a match was found when the candidate was enqueued.
  if (!Samples)
    return;

  // For AutoFDO profile, retrieve candidate profiles by walking over
  // the nested inlinee profiles.
  if (!FunctionSamples::ProfileIsCS) {
    Samples->findInlinedFunctions(InlinedGUIDs, SymbolMap, Threshold);
    return;
  }

  ContextTrieNode *Caller = ContextTracker->getContextNodeForProfile(Samples);
  std::queue<ContextTrieNode *> CalleeList;
  CalleeList.push(Caller);
  while (!CalleeList.empty()) {
    ContextTrieNode *Node = CalleeList.front();
    CalleeList.pop();
    FunctionSamples *CalleeSample = Node->getFunctionSamples();
    // For CSSPGO profile, retrieve candidate profile by walking over the
    // trie built for context profile. Note that also take call targets
    // even if callee doesn't have a corresponding context profile.
    if (!CalleeSample)
      continue;

    // If pre-inliner decision is used, honor that for importing as well.
    bool PreInline =
        UsePreInlinerDecision &&
        CalleeSample->getContext().hasAttribute(ContextShouldBeInlined);
    if (!PreInline && CalleeSample->getHeadSamplesEstimate() < Threshold)
      continue;
    
    Function *Func = SymbolMap.lookup(CalleeSample->getFunction());
    // Add to the import list only when it's defined out of module.
    if (!Func || Func->isDeclaration())
      InlinedGUIDs.insert(CalleeSample->getGUID());

    // Import hot CallTargets, which may not be available in IR because full
    // profile annotation cannot be done until backend compilation in ThinLTO.
    for (const auto &BS : CalleeSample->getBodySamples())
      for (const auto &TS : BS.second.getCallTargets())
        if (TS.second > Threshold) {
          const Function *Callee = SymbolMap.lookup(TS.first);
          if (!Callee || Callee->isDeclaration())
            InlinedGUIDs.insert(TS.first.getHashCode());
        }

    // Import hot child context profile associted with callees. Note that this
    // may have some overlap with the call target loop above, but doing this
    // based child context profile again effectively allow us to use the max of
    // entry count and call target count to determine importing.
    for (auto &Child : Node->getAllChildContext()) {
      ContextTrieNode *CalleeNode = &Child.second;
      CalleeList.push(CalleeNode);
    }
  }
}

/// Iteratively inline hot callsites of a function.
///
/// Iteratively traverse all callsites of the function \p F, so as to
/// find out callsites with corresponding inline instances.
///
/// For such callsites,
/// - If it is hot enough, inline the callsites and adds callsites of the callee
///   into the caller. If the call is an indirect call, first promote
///   it to direct call. Each indirect call is limited with a single target.
///
/// - If a callsite is not inlined, merge the its profile to the outline
///   version (if --sample-profile-merge-inlinee is true), or scale the
///   counters of standalone function based on the profile of inlined
///   instances (if --sample-profile-merge-inlinee is false).
///
///   Later passes may consume the updated profiles.
///
/// \param F function to perform iterative inlining.
/// \param InlinedGUIDs a set to be updated to include all GUIDs that are
///     inlined in the profiled binary.
///
/// \returns True if there is any inline happened.
bool SampleProfileLoader::inlineHotFunctions(
    Function &F, DenseSet<GlobalValue::GUID> &InlinedGUIDs) {
  // ProfAccForSymsInList is used in callsiteIsHot. The assertion makes sure
  // Profile symbol list is ignored when profile-sample-accurate is on.
  assert((!ProfAccForSymsInList ||
          (!ProfileSampleAccurate &&
           !F.hasFnAttribute("profile-sample-accurate"))) &&
         "ProfAccForSymsInList should be false when profile-sample-accurate "
         "is enabled");

  MapVector<CallBase *, const FunctionSamples *> LocalNotInlinedCallSites;
  bool Changed = false;
  bool LocalChanged = true;
  while (LocalChanged) {
    LocalChanged = false;
    SmallVector<CallBase *, 10> CIS;
    for (auto &BB : F) {
      bool Hot = false;
      SmallVector<CallBase *, 10> AllCandidates;
      SmallVector<CallBase *, 10> ColdCandidates;
      for (auto &I : BB) {
        const FunctionSamples *FS = nullptr;
        if (auto *CB = dyn_cast<CallBase>(&I)) {
          if (!isa<IntrinsicInst>(I)) {
            if ((FS = findCalleeFunctionSamples(*CB))) {
              assert((!FunctionSamples::UseMD5 || FS->GUIDToFuncNameMap) &&
                     "GUIDToFuncNameMap has to be populated");
              AllCandidates.push_back(CB);
              if (FS->getHeadSamplesEstimate() > 0 ||
                  FunctionSamples::ProfileIsCS)
                LocalNotInlinedCallSites.insert({CB, FS});
              if (callsiteIsHot(FS, PSI, ProfAccForSymsInList))
                Hot = true;
              else if (shouldInlineColdCallee(*CB))
                ColdCandidates.push_back(CB);
            } else if (getExternalInlineAdvisorShouldInline(*CB)) {
              AllCandidates.push_back(CB);
            }
          }
        }
      }
      if (Hot || ExternalInlineAdvisor) {
        CIS.insert(CIS.begin(), AllCandidates.begin(), AllCandidates.end());
        emitOptimizationRemarksForInlineCandidates(AllCandidates, F, true);
      } else {
        CIS.insert(CIS.begin(), ColdCandidates.begin(), ColdCandidates.end());
        emitOptimizationRemarksForInlineCandidates(ColdCandidates, F, false);
      }
    }
    for (CallBase *I : CIS) {
      Function *CalledFunction = I->getCalledFunction();
      InlineCandidate Candidate = {I, LocalNotInlinedCallSites.lookup(I),
                                   0 /* dummy count */,
                                   1.0 /* dummy distribution factor */};
      // Do not inline recursive calls.
      if (CalledFunction == &F)
        continue;
      if (I->isIndirectCall()) {
        uint64_t Sum;
        for (const auto *FS : findIndirectCallFunctionSamples(*I, Sum)) {
          uint64_t SumOrigin = Sum;
          if (LTOPhase == ThinOrFullLTOPhase::ThinLTOPreLink) {
            findExternalInlineCandidate(I, FS, InlinedGUIDs,
                                        PSI->getOrCompHotCountThreshold());
            continue;
          }
          if (!callsiteIsHot(FS, PSI, ProfAccForSymsInList))
            continue;

          Candidate = {I, FS, FS->getHeadSamplesEstimate(), 1.0};
          if (tryPromoteAndInlineCandidate(F, Candidate, SumOrigin, Sum)) {
            LocalNotInlinedCallSites.erase(I);
            LocalChanged = true;
          }
        }
      } else if (CalledFunction && CalledFunction->getSubprogram() &&
                 !CalledFunction->isDeclaration()) {
        if (tryInlineCandidate(Candidate)) {
          LocalNotInlinedCallSites.erase(I);
          LocalChanged = true;
        }
      } else if (LTOPhase == ThinOrFullLTOPhase::ThinLTOPreLink) {
        findExternalInlineCandidate(I, findCalleeFunctionSamples(*I),
                                    InlinedGUIDs,
                                    PSI->getOrCompHotCountThreshold());
      }
    }
    Changed |= LocalChanged;
  }

  // For CS profile, profile for not inlined context will be merged when
  // base profile is being retrieved.
  if (!FunctionSamples::ProfileIsCS)
    promoteMergeNotInlinedContextSamples(LocalNotInlinedCallSites, F);
  return Changed;
}

bool SampleProfileLoader::tryInlineCandidate(
    InlineCandidate &Candidate, SmallVector<CallBase *, 8> *InlinedCallSites) {
  // Do not attempt to inline a candidate if
  // --disable-sample-loader-inlining is true.
  if (DisableSampleLoaderInlining)
    return false;

  CallBase &CB = *Candidate.CallInstr;
  Function *CalledFunction = CB.getCalledFunction();
  assert(CalledFunction && "Expect a callee with definition");
  DebugLoc DLoc = CB.getDebugLoc();
  BasicBlock *BB = CB.getParent();

  InlineCost Cost = shouldInlineCandidate(Candidate);
  if (Cost.isNever()) {
    ORE->emit(OptimizationRemarkAnalysis(getAnnotatedRemarkPassName(),
                                         "InlineFail", DLoc, BB)
              << "incompatible inlining");
    return false;
  }

  if (!Cost)
    return false;

  InlineFunctionInfo IFI(GetAC);
  IFI.UpdateProfile = false;
  InlineResult IR = InlineFunction(CB, IFI,
                                   /*MergeAttributes=*/true);
  if (!IR.isSuccess())
    return false;

  // The call to InlineFunction erases I, so we can't pass it here.
  emitInlinedIntoBasedOnCost(*ORE, DLoc, BB, *CalledFunction, *BB->getParent(),
                             Cost, true, getAnnotatedRemarkPassName());

  // Now populate the list of newly exposed call sites.
  if (InlinedCallSites) {
    InlinedCallSites->clear();
    for (auto &I : IFI.InlinedCallSites)
      InlinedCallSites->push_back(I);
  }

  if (FunctionSamples::ProfileIsCS)
    ContextTracker->markContextSamplesInlined(Candidate.CalleeSamples);
  ++NumCSInlined;

  // Prorate inlined probes for a duplicated inlining callsite which probably
  // has a distribution less than 100%. Samples for an inlinee should be
  // distributed among the copies of the original callsite based on each
  // callsite's distribution factor for counts accuracy. Note that an inlined
  // probe may come with its own distribution factor if it has been duplicated
  // in the inlinee body. The two factor are multiplied to reflect the
  // aggregation of duplication.
  if (Candidate.CallsiteDistribution < 1) {
    for (auto &I : IFI.InlinedCallSites) {
      if (std::optional<PseudoProbe> Probe = extractProbe(*I))
        setProbeDistributionFactor(*I, Probe->Factor *
                                   Candidate.CallsiteDistribution);
    }
    NumDuplicatedInlinesite++;
  }

  return true;
}

bool SampleProfileLoader::getInlineCandidate(InlineCandidate *NewCandidate,
                                             CallBase *CB) {
  assert(CB && "Expect non-null call instruction");

  if (isa<IntrinsicInst>(CB))
    return false;

  // Find the callee's profile. For indirect call, find hottest target profile.
  const FunctionSamples *CalleeSamples = findCalleeFunctionSamples(*CB);
  // If ExternalInlineAdvisor wants to inline this site, do so even
  // if Samples are not present.
  if (!CalleeSamples && !getExternalInlineAdvisorShouldInline(*CB))
    return false;

  float Factor = 1.0;
  if (std::optional<PseudoProbe> Probe = extractProbe(*CB))
    Factor = Probe->Factor;

  uint64_t CallsiteCount =
      CalleeSamples ? CalleeSamples->getHeadSamplesEstimate() * Factor : 0;
  *NewCandidate = {CB, CalleeSamples, CallsiteCount, Factor};
  return true;
}

std::optional<InlineCost>
SampleProfileLoader::getExternalInlineAdvisorCost(CallBase &CB) {
  std::unique_ptr<InlineAdvice> Advice = nullptr;
  if (ExternalInlineAdvisor) {
    Advice = ExternalInlineAdvisor->getAdvice(CB);
    if (Advice) {
      if (!Advice->isInliningRecommended()) {
        Advice->recordUnattemptedInlining();
        return InlineCost::getNever("not previously inlined");
      }
      Advice->recordInlining();
      return InlineCost::getAlways("previously inlined");
    }
  }

  return {};
}

bool SampleProfileLoader::getExternalInlineAdvisorShouldInline(CallBase &CB) {
  std::optional<InlineCost> Cost = getExternalInlineAdvisorCost(CB);
  return Cost ? !!*Cost : false;
}

InlineCost
SampleProfileLoader::shouldInlineCandidate(InlineCandidate &Candidate) {
  if (std::optional<InlineCost> ReplayCost =
          getExternalInlineAdvisorCost(*Candidate.CallInstr))
    return *ReplayCost;
  // Adjust threshold based on call site hotness, only do this for callsite
  // prioritized inliner because otherwise cost-benefit check is done earlier.
  int SampleThreshold = SampleColdCallSiteThreshold;
  if (CallsitePrioritizedInline) {
    if (Candidate.CallsiteCount > PSI->getHotCountThreshold())
      SampleThreshold = SampleHotCallSiteThreshold;
    else if (!ProfileSizeInline)
      return InlineCost::getNever("cold callsite");
  }

  Function *Callee = Candidate.CallInstr->getCalledFunction();
  assert(Callee && "Expect a definition for inline candidate of direct call");

  InlineParams Params = getInlineParams();
  // We will ignore the threshold from inline cost, so always get full cost.
  Params.ComputeFullInlineCost = true;
  Params.AllowRecursiveCall = AllowRecursiveInline;
  // Checks if there is anything in the reachable portion of the callee at
  // this callsite that makes this inlining potentially illegal. Need to
  // set ComputeFullInlineCost, otherwise getInlineCost may return early
  // when cost exceeds threshold without checking all IRs in the callee.
  // The acutal cost does not matter because we only checks isNever() to
  // see if it is legal to inline the callsite.
  InlineCost Cost = getInlineCost(*Candidate.CallInstr, Callee, Params,
                                  GetTTI(*Callee), GetAC, GetTLI);

  // Honor always inline and never inline from call analyzer
  if (Cost.isNever() || Cost.isAlways())
    return Cost;

  // With CSSPGO, the preinliner in llvm-profgen can estimate global inline
  // decisions based on hotness as well as accurate function byte sizes for
  // given context using function/inlinee sizes from previous build. It
  // stores the decision in profile, and also adjust/merge context profile
  // aiming at better context-sensitive post-inline profile quality, assuming
  // all inline decision estimates are going to be honored by compiler. Here
  // we replay that inline decision under `sample-profile-use-preinliner`.
  // Note that we don't need to handle negative decision from preinliner as
  // context profile for not inlined calls are merged by preinliner already.
  if (UsePreInlinerDecision && Candidate.CalleeSamples) {
    // Once two node are merged due to promotion, we're losing some context
    // so the original context-sensitive preinliner decision should be ignored
    // for SyntheticContext.
    SampleContext &Context = Candidate.CalleeSamples->getContext();
    if (!Context.hasState(SyntheticContext) &&
        Context.hasAttribute(ContextShouldBeInlined))
      return InlineCost::getAlways("preinliner");
  }

  // For old FDO inliner, we inline the call site as long as cost is not
  // "Never". The cost-benefit check is done earlier.
  if (!CallsitePrioritizedInline) {
    return InlineCost::get(Cost.getCost(), INT_MAX);
  }

  // Otherwise only use the cost from call analyzer, but overwite threshold with
  // Sample PGO threshold.
  return InlineCost::get(Cost.getCost(), SampleThreshold);
}

bool SampleProfileLoader::inlineHotFunctionsWithPriority(
    Function &F, DenseSet<GlobalValue::GUID> &InlinedGUIDs) {
  // ProfAccForSymsInList is used in callsiteIsHot. The assertion makes sure
  // Profile symbol list is ignored when profile-sample-accurate is on.
  assert((!ProfAccForSymsInList ||
          (!ProfileSampleAccurate &&
           !F.hasFnAttribute("profile-sample-accurate"))) &&
         "ProfAccForSymsInList should be false when profile-sample-accurate "
         "is enabled");

  // Populating worklist with initial call sites from root inliner, along
  // with call site weights.
  CandidateQueue CQueue;
  InlineCandidate NewCandidate;
  for (auto &BB : F) {
    for (auto &I : BB) {
      auto *CB = dyn_cast<CallBase>(&I);
      if (!CB)
        continue;
      if (getInlineCandidate(&NewCandidate, CB))
        CQueue.push(NewCandidate);
    }
  }

  // Cap the size growth from profile guided inlining. This is needed even
  // though cost of each inline candidate already accounts for callee size,
  // because with top-down inlining, we can grow inliner size significantly
  // with large number of smaller inlinees each pass the cost check.
  assert(ProfileInlineLimitMax >= ProfileInlineLimitMin &&
         "Max inline size limit should not be smaller than min inline size "
         "limit.");
  unsigned SizeLimit = F.getInstructionCount() * ProfileInlineGrowthLimit;
  SizeLimit = std::min(SizeLimit, (unsigned)ProfileInlineLimitMax);
  SizeLimit = std::max(SizeLimit, (unsigned)ProfileInlineLimitMin);
  if (ExternalInlineAdvisor)
    SizeLimit = std::numeric_limits<unsigned>::max();

  MapVector<CallBase *, const FunctionSamples *> LocalNotInlinedCallSites;

  // Perform iterative BFS call site prioritized inlining
  bool Changed = false;
  while (!CQueue.empty() && F.getInstructionCount() < SizeLimit) {
    InlineCandidate Candidate = CQueue.top();
    CQueue.pop();
    CallBase *I = Candidate.CallInstr;
    Function *CalledFunction = I->getCalledFunction();

    if (CalledFunction == &F)
      continue;
    if (I->isIndirectCall()) {
      uint64_t Sum = 0;
      auto CalleeSamples = findIndirectCallFunctionSamples(*I, Sum);
      uint64_t SumOrigin = Sum;
      Sum *= Candidate.CallsiteDistribution;
      unsigned ICPCount = 0;
      for (const auto *FS : CalleeSamples) {
        // TODO: Consider disable pre-lTO ICP for MonoLTO as well
        if (LTOPhase == ThinOrFullLTOPhase::ThinLTOPreLink) {
          findExternalInlineCandidate(I, FS, InlinedGUIDs,
                                      PSI->getOrCompHotCountThreshold());
          continue;
        }
        uint64_t EntryCountDistributed =
            FS->getHeadSamplesEstimate() * Candidate.CallsiteDistribution;
        // In addition to regular inline cost check, we also need to make sure
        // ICP isn't introducing excessive speculative checks even if individual
        // target looks beneficial to promote and inline. That means we should
        // only do ICP when there's a small number dominant targets.
        if (ICPCount >= ProfileICPRelativeHotnessSkip &&
            EntryCountDistributed * 100 < SumOrigin * ProfileICPRelativeHotness)
          break;
        // TODO: Fix CallAnalyzer to handle all indirect calls.
        // For indirect call, we don't run CallAnalyzer to get InlineCost
        // before actual inlining. This is because we could see two different
        // types from the same definition, which makes CallAnalyzer choke as
        // it's expecting matching parameter type on both caller and callee
        // side. See example from PR18962 for the triggering cases (the bug was
        // fixed, but we generate different types).
        if (!PSI->isHotCount(EntryCountDistributed))
          break;
        SmallVector<CallBase *, 8> InlinedCallSites;
        // Attach function profile for promoted indirect callee, and update
        // call site count for the promoted inline candidate too.
        Candidate = {I, FS, EntryCountDistributed,
                     Candidate.CallsiteDistribution};
        if (tryPromoteAndInlineCandidate(F, Candidate, SumOrigin, Sum,
                                         &InlinedCallSites)) {
          for (auto *CB : InlinedCallSites) {
            if (getInlineCandidate(&NewCandidate, CB))
              CQueue.emplace(NewCandidate);
          }
          ICPCount++;
          Changed = true;
        } else if (!ContextTracker) {
          LocalNotInlinedCallSites.insert({I, FS});
        }
      }
    } else if (CalledFunction && CalledFunction->getSubprogram() &&
               !CalledFunction->isDeclaration()) {
      SmallVector<CallBase *, 8> InlinedCallSites;
      if (tryInlineCandidate(Candidate, &InlinedCallSites)) {
        for (auto *CB : InlinedCallSites) {
          if (getInlineCandidate(&NewCandidate, CB))
            CQueue.emplace(NewCandidate);
        }
        Changed = true;
      } else if (!ContextTracker) {
        LocalNotInlinedCallSites.insert({I, Candidate.CalleeSamples});
      }
    } else if (LTOPhase == ThinOrFullLTOPhase::ThinLTOPreLink) {
      findExternalInlineCandidate(I, findCalleeFunctionSamples(*I),
                                  InlinedGUIDs,
                                  PSI->getOrCompHotCountThreshold());
    }
  }

  if (!CQueue.empty()) {
    if (SizeLimit == (unsigned)ProfileInlineLimitMax)
      ++NumCSInlinedHitMaxLimit;
    else if (SizeLimit == (unsigned)ProfileInlineLimitMin)
      ++NumCSInlinedHitMinLimit;
    else
      ++NumCSInlinedHitGrowthLimit;
  }

  // For CS profile, profile for not inlined context will be merged when
  // base profile is being retrieved.
  if (!FunctionSamples::ProfileIsCS)
    promoteMergeNotInlinedContextSamples(LocalNotInlinedCallSites, F);
  return Changed;
}

void SampleProfileLoader::promoteMergeNotInlinedContextSamples(
    MapVector<CallBase *, const FunctionSamples *> NonInlinedCallSites,
    const Function &F) {
  // Accumulate not inlined callsite information into notInlinedSamples
  for (const auto &Pair : NonInlinedCallSites) {
    CallBase *I = Pair.first;
    Function *Callee = I->getCalledFunction();
    if (!Callee || Callee->isDeclaration())
      continue;

    ORE->emit(
        OptimizationRemarkAnalysis(getAnnotatedRemarkPassName(), "NotInline",
                                   I->getDebugLoc(), I->getParent())
        << "previous inlining not repeated: '" << ore::NV("Callee", Callee)
        << "' into '" << ore::NV("Caller", &F) << "'");

    ++NumCSNotInlined;
    const FunctionSamples *FS = Pair.second;
    if (FS->getTotalSamples() == 0 && FS->getHeadSamplesEstimate() == 0) {
      continue;
    }

    // Do not merge a context that is already duplicated into the base profile.
    if (FS->getContext().hasAttribute(sampleprof::ContextDuplicatedIntoBase))
      continue;

    if (ProfileMergeInlinee) {
      // A function call can be replicated by optimizations like callsite
      // splitting or jump threading and the replicates end up sharing the
      // sample nested callee profile instead of slicing the original
      // inlinee's profile. We want to do merge exactly once by filtering out
      // callee profiles with a non-zero head sample count.
      if (FS->getHeadSamples() == 0) {
        // Use entry samples as head samples during the merge, as inlinees
        // don't have head samples.
        const_cast<FunctionSamples *>(FS)->addHeadSamples(
            FS->getHeadSamplesEstimate());

        // Note that we have to do the merge right after processing function.
        // This allows OutlineFS's profile to be used for annotation during
        // top-down processing of functions' annotation.
        FunctionSamples *OutlineFS = Reader->getSamplesFor(*Callee);
        // If outlined function does not exist in the profile, add it to a
        // separate map so that it does not rehash the original profile.
        if (!OutlineFS)
          OutlineFS = &OutlineFunctionSamples[
              FunctionId(FunctionSamples::getCanonicalFnName(Callee->getName()))];
        OutlineFS->merge(*FS, 1);
        // Set outlined profile to be synthetic to not bias the inliner.
        OutlineFS->SetContextSynthetic();
      }
    } else {
      auto pair =
          notInlinedCallInfo.try_emplace(Callee, NotInlinedProfileInfo{0});
      pair.first->second.entryCount += FS->getHeadSamplesEstimate();
    }
  }
}

/// Returns the sorted CallTargetMap \p M by count in descending order.
static SmallVector<InstrProfValueData, 2>
GetSortedValueDataFromCallTargets(const SampleRecord::CallTargetMap &M) {
  SmallVector<InstrProfValueData, 2> R;
  for (const auto &I : SampleRecord::SortCallTargets(M)) {
    R.emplace_back(
        InstrProfValueData{I.first.getHashCode(), I.second});
  }
  return R;
}

// Generate MD_prof metadata for every branch instruction using the
// edge weights computed during propagation.
void SampleProfileLoader::generateMDProfMetadata(Function &F) {
  // Generate MD_prof metadata for every branch instruction using the
  // edge weights computed during propagation.
  LLVM_DEBUG(dbgs() << "\nPropagation complete. Setting branch weights\n");
  LLVMContext &Ctx = F.getContext();
  MDBuilder MDB(Ctx);
  for (auto &BI : F) {
    BasicBlock *BB = &BI;

    if (BlockWeights[BB]) {
      for (auto &I : *BB) {
        if (!isa<CallInst>(I) && !isa<InvokeInst>(I))
          continue;
        if (!cast<CallBase>(I).getCalledFunction()) {
          const DebugLoc &DLoc = I.getDebugLoc();
          if (!DLoc)
            continue;
          const DILocation *DIL = DLoc;
          const FunctionSamples *FS = findFunctionSamples(I);
          if (!FS)
            continue;
          auto CallSite = FunctionSamples::getCallSiteIdentifier(DIL);
          auto T = FS->findCallTargetMapAt(CallSite);
          if (!T || T.get().empty())
            continue;
          if (FunctionSamples::ProfileIsProbeBased) {
            // Prorate the callsite counts based on the pre-ICP distribution
            // factor to reflect what is already done to the callsite before
            // ICP, such as calliste cloning.
            if (std::optional<PseudoProbe> Probe = extractProbe(I)) {
              if (Probe->Factor < 1)
                T = SampleRecord::adjustCallTargets(T.get(), Probe->Factor);
            }
          }
          SmallVector<InstrProfValueData, 2> SortedCallTargets =
              GetSortedValueDataFromCallTargets(T.get());
          uint64_t Sum = 0;
          for (const auto &C : T.get())
            Sum += C.second;
          // With CSSPGO all indirect call targets are counted torwards the
          // original indirect call site in the profile, including both
          // inlined and non-inlined targets.
          if (!FunctionSamples::ProfileIsCS) {
            if (const FunctionSamplesMap *M =
                    FS->findFunctionSamplesMapAt(CallSite)) {
              for (const auto &NameFS : *M)
                Sum += NameFS.second.getHeadSamplesEstimate();
            }
          }
          if (Sum)
            updateIDTMetaData(I, SortedCallTargets, Sum);
          else if (OverwriteExistingWeights)
            I.setMetadata(LLVMContext::MD_prof, nullptr);
        } else if (!isa<IntrinsicInst>(&I)) {
          setBranchWeights(I, {static_cast<uint32_t>(BlockWeights[BB])});
        }
      }
    } else if (OverwriteExistingWeights || ProfileSampleBlockAccurate) {
      // Set profile metadata (possibly annotated by LTO prelink) to zero or
      // clear it for cold code.
      for (auto &I : *BB) {
        if (isa<CallInst>(I) || isa<InvokeInst>(I)) {
          if (cast<CallBase>(I).isIndirectCall()) {
            I.setMetadata(LLVMContext::MD_prof, nullptr);
          } else {
            setBranchWeights(I, {uint32_t(0)});
          }
        }
      }
    }

    Instruction *TI = BB->getTerminator();
    if (TI->getNumSuccessors() == 1)
      continue;
    if (!isa<BranchInst>(TI) && !isa<SwitchInst>(TI) &&
        !isa<IndirectBrInst>(TI))
      continue;

    DebugLoc BranchLoc = TI->getDebugLoc();
    LLVM_DEBUG(dbgs() << "\nGetting weights for branch at line "
                      << ((BranchLoc) ? Twine(BranchLoc.getLine())
                                      : Twine("<UNKNOWN LOCATION>"))
                      << ".\n");
    SmallVector<uint32_t, 4> Weights;
    uint32_t MaxWeight = 0;
    Instruction *MaxDestInst;
    // Since profi treats multiple edges (multiway branches) as a single edge,
    // we need to distribute the computed weight among the branches. We do
    // this by evenly splitting the edge weight among destinations.
    DenseMap<const BasicBlock *, uint64_t> EdgeMultiplicity;
    std::vector<uint64_t> EdgeIndex;
    if (SampleProfileUseProfi) {
      EdgeIndex.resize(TI->getNumSuccessors());
      for (unsigned I = 0; I < TI->getNumSuccessors(); ++I) {
        const BasicBlock *Succ = TI->getSuccessor(I);
        EdgeIndex[I] = EdgeMultiplicity[Succ];
        EdgeMultiplicity[Succ]++;
      }
    }
    for (unsigned I = 0; I < TI->getNumSuccessors(); ++I) {
      BasicBlock *Succ = TI->getSuccessor(I);
      Edge E = std::make_pair(BB, Succ);
      uint64_t Weight = EdgeWeights[E];
      LLVM_DEBUG(dbgs() << "\t"; printEdgeWeight(dbgs(), E));
      // Use uint32_t saturated arithmetic to adjust the incoming weights,
      // if needed. Sample counts in profiles are 64-bit unsigned values,
      // but internally branch weights are expressed as 32-bit values.
      if (Weight > std::numeric_limits<uint32_t>::max()) {
        LLVM_DEBUG(dbgs() << " (saturated due to uint32_t overflow)");
        Weight = std::numeric_limits<uint32_t>::max();
      }
      if (!SampleProfileUseProfi) {
        // Weight is added by one to avoid propagation errors introduced by
        // 0 weights.
        Weights.push_back(static_cast<uint32_t>(Weight + 1));
      } else {
        // Profi creates proper weights that do not require "+1" adjustments but
        // we evenly split the weight among branches with the same destination.
        uint64_t W = Weight / EdgeMultiplicity[Succ];
        // Rounding up, if needed, so that first branches are hotter.
        if (EdgeIndex[I] < Weight % EdgeMultiplicity[Succ])
          W++;
        Weights.push_back(static_cast<uint32_t>(W));
      }
      if (Weight != 0) {
        if (Weight > MaxWeight) {
          MaxWeight = Weight;
          MaxDestInst = Succ->getFirstNonPHIOrDbgOrLifetime();
        }
      }
    }

    misexpect::checkExpectAnnotations(*TI, Weights, /*IsFrontend=*/false);

    uint64_t TempWeight;
    // Only set weights if there is at least one non-zero weight.
    // In any other case, let the analyzer set weights.
    // Do not set weights if the weights are present unless under
    // OverwriteExistingWeights. In ThinLTO, the profile annotation is done
    // twice. If the first annotation already set the weights, the second pass
    // does not need to set it. With OverwriteExistingWeights, Blocks with zero
    // weight should have their existing metadata (possibly annotated by LTO
    // prelink) cleared.
    if (MaxWeight > 0 &&
        (!TI->extractProfTotalWeight(TempWeight) || OverwriteExistingWeights)) {
      LLVM_DEBUG(dbgs() << "SUCCESS. Found non-zero weights.\n");
      setBranchWeights(*TI, Weights);
      ORE->emit([&]() {
        return OptimizationRemark(DEBUG_TYPE, "PopularDest", MaxDestInst)
               << "most popular destination for conditional branches at "
               << ore::NV("CondBranchesLoc", BranchLoc);
      });
    } else {
      if (OverwriteExistingWeights) {
        TI->setMetadata(LLVMContext::MD_prof, nullptr);
        LLVM_DEBUG(dbgs() << "CLEARED. All branch weights are zero.\n");
      } else {
        LLVM_DEBUG(dbgs() << "SKIPPED. All branch weights are zero.\n");
      }
    }
  }
}

/// Once all the branch weights are computed, we emit the MD_prof
/// metadata on BB using the computed values for each of its branches.
///
/// \param F The function to query.
///
/// \returns true if \p F was modified. Returns false, otherwise.
bool SampleProfileLoader::emitAnnotations(Function &F) {
  bool Changed = false;

  if (FunctionSamples::ProfileIsProbeBased) {
    if (!ProbeManager->profileIsValid(F, *Samples)) {
      LLVM_DEBUG(
          dbgs() << "Profile is invalid due to CFG mismatch for Function "
                 << F.getName() << "\n");
      ++NumMismatchedProfile;
      if (!SalvageStaleProfile)
        return false;
    }
    ++NumMatchedProfile;
  } else {
    if (getFunctionLoc(F) == 0)
      return false;

    LLVM_DEBUG(dbgs() << "Line number for the first instruction in "
                      << F.getName() << ": " << getFunctionLoc(F) << "\n");
  }

  DenseSet<GlobalValue::GUID> InlinedGUIDs;
  if (CallsitePrioritizedInline)
    Changed |= inlineHotFunctionsWithPriority(F, InlinedGUIDs);
  else
    Changed |= inlineHotFunctions(F, InlinedGUIDs);

  Changed |= computeAndPropagateWeights(F, InlinedGUIDs);

  if (Changed)
    generateMDProfMetadata(F);

  emitCoverageRemarks(F);
  return Changed;
}

std::unique_ptr<ProfiledCallGraph>
SampleProfileLoader::buildProfiledCallGraph(Module &M) {
  std::unique_ptr<ProfiledCallGraph> ProfiledCG;
  if (FunctionSamples::ProfileIsCS)
    ProfiledCG = std::make_unique<ProfiledCallGraph>(*ContextTracker);
  else
    ProfiledCG = std::make_unique<ProfiledCallGraph>(Reader->getProfiles());

  // Add all functions into the profiled call graph even if they are not in
  // the profile. This makes sure functions missing from the profile still
  // gets a chance to be processed.
  for (Function &F : M) {
    if (F.isDeclaration() || !F.hasFnAttribute("use-sample-profile"))
      continue;
    ProfiledCG->addProfiledFunction(
          getRepInFormat(FunctionSamples::getCanonicalFnName(F)));
  }

  return ProfiledCG;
}

std::vector<Function *>
SampleProfileLoader::buildFunctionOrder(Module &M, LazyCallGraph &CG) {
  std::vector<Function *> FunctionOrderList;
  FunctionOrderList.reserve(M.size());

  if (!ProfileTopDownLoad && UseProfiledCallGraph)
    errs() << "WARNING: -use-profiled-call-graph ignored, should be used "
              "together with -sample-profile-top-down-load.\n";

  if (!ProfileTopDownLoad) {
    if (ProfileMergeInlinee) {
      // Disable ProfileMergeInlinee if profile is not loaded in top down order,
      // because the profile for a function may be used for the profile
      // annotation of its outline copy before the profile merging of its
      // non-inlined inline instances, and that is not the way how
      // ProfileMergeInlinee is supposed to work.
      ProfileMergeInlinee = false;
    }

    for (Function &F : M)
      if (!F.isDeclaration() && F.hasFnAttribute("use-sample-profile"))
        FunctionOrderList.push_back(&F);
    return FunctionOrderList;
  }

  if (UseProfiledCallGraph || (FunctionSamples::ProfileIsCS &&
                               !UseProfiledCallGraph.getNumOccurrences())) {
    // Use profiled call edges to augment the top-down order. There are cases
    // that the top-down order computed based on the static call graph doesn't
    // reflect real execution order. For example
    //
    // 1. Incomplete static call graph due to unknown indirect call targets.
    //    Adjusting the order by considering indirect call edges from the
    //    profile can enable the inlining of indirect call targets by allowing
    //    the caller processed before them.
    // 2. Mutual call edges in an SCC. The static processing order computed for
    //    an SCC may not reflect the call contexts in the context-sensitive
    //    profile, thus may cause potential inlining to be overlooked. The
    //    function order in one SCC is being adjusted to a top-down order based
    //    on the profile to favor more inlining. This is only a problem with CS
    //    profile.
    // 3. Transitive indirect call edges due to inlining. When a callee function
    //    (say B) is inlined into a caller function (say A) in LTO prelink,
    //    every call edge originated from the callee B will be transferred to
    //    the caller A. If any transferred edge (say A->C) is indirect, the
    //    original profiled indirect edge B->C, even if considered, would not
    //    enforce a top-down order from the caller A to the potential indirect
    //    call target C in LTO postlink since the inlined callee B is gone from
    //    the static call graph.
    // 4. #3 can happen even for direct call targets, due to functions defined
    //    in header files. A header function (say A), when included into source
    //    files, is defined multiple times but only one definition survives due
    //    to ODR. Therefore, the LTO prelink inlining done on those dropped
    //    definitions can be useless based on a local file scope. More
    //    importantly, the inlinee (say B), once fully inlined to a
    //    to-be-dropped A, will have no profile to consume when its outlined
    //    version is compiled. This can lead to a profile-less prelink
    //    compilation for the outlined version of B which may be called from
    //    external modules. while this isn't easy to fix, we rely on the
    //    postlink AutoFDO pipeline to optimize B. Since the survived copy of
    //    the A can be inlined in its local scope in prelink, it may not exist
    //    in the merged IR in postlink, and we'll need the profiled call edges
    //    to enforce a top-down order for the rest of the functions.
    //
    // Considering those cases, a profiled call graph completely independent of
    // the static call graph is constructed based on profile data, where
    // function objects are not even needed to handle case #3 and case 4.
    //
    // Note that static callgraph edges are completely ignored since they
    // can be conflicting with profiled edges for cyclic SCCs and may result in
    // an SCC order incompatible with profile-defined one. Using strictly
    // profile order ensures a maximum inlining experience. On the other hand,
    // static call edges are not so important when they don't correspond to a
    // context in the profile.

    std::unique_ptr<ProfiledCallGraph> ProfiledCG = buildProfiledCallGraph(M);
    scc_iterator<ProfiledCallGraph *> CGI = scc_begin(ProfiledCG.get());
    while (!CGI.isAtEnd()) {
      auto Range = *CGI;
      if (SortProfiledSCC) {
        // Sort nodes in one SCC based on callsite hotness.
        scc_member_iterator<ProfiledCallGraph *> SI(*CGI);
        Range = *SI;
      }
      for (auto *Node : Range) {
        Function *F = SymbolMap.lookup(Node->Name);
        if (F && !F->isDeclaration() && F->hasFnAttribute("use-sample-profile"))
          FunctionOrderList.push_back(F);
      }
      ++CGI;
    }
  } else {
    CG.buildRefSCCs();
    for (LazyCallGraph::RefSCC &RC : CG.postorder_ref_sccs()) {
      for (LazyCallGraph::SCC &C : RC) {
        for (LazyCallGraph::Node &N : C) {
          Function &F = N.getFunction();
          if (!F.isDeclaration() && F.hasFnAttribute("use-sample-profile"))
            FunctionOrderList.push_back(&F);
        }
      }
    }
  }

  std::reverse(FunctionOrderList.begin(), FunctionOrderList.end());

  LLVM_DEBUG({
    dbgs() << "Function processing order:\n";
    for (auto F : FunctionOrderList) {
      dbgs() << F->getName() << "\n";
    }
  });

  return FunctionOrderList;
}

bool SampleProfileLoader::doInitialization(Module &M,
                                           FunctionAnalysisManager *FAM) {
  auto &Ctx = M.getContext();

  auto ReaderOrErr = SampleProfileReader::create(
      Filename, Ctx, *FS, FSDiscriminatorPass::Base, RemappingFilename);
  if (std::error_code EC = ReaderOrErr.getError()) {
    std::string Msg = "Could not open profile: " + EC.message();
    Ctx.diagnose(DiagnosticInfoSampleProfile(Filename, Msg));
    return false;
  }
  Reader = std::move(ReaderOrErr.get());
  Reader->setSkipFlatProf(LTOPhase == ThinOrFullLTOPhase::ThinLTOPostLink);
  // set module before reading the profile so reader may be able to only
  // read the function profiles which are used by the current module.
  Reader->setModule(&M);
  if (std::error_code EC = Reader->read()) {
    std::string Msg = "profile reading failed: " + EC.message();
    Ctx.diagnose(DiagnosticInfoSampleProfile(Filename, Msg));
    return false;
  }

  PSL = Reader->getProfileSymbolList();

  // While profile-sample-accurate is on, ignore symbol list.
  ProfAccForSymsInList =
      ProfileAccurateForSymsInList && PSL && !ProfileSampleAccurate;
  if (ProfAccForSymsInList) {
    NamesInProfile.clear();
    GUIDsInProfile.clear();
    if (auto NameTable = Reader->getNameTable()) {
      if (FunctionSamples::UseMD5) {
        for (auto Name : *NameTable)
          GUIDsInProfile.insert(Name.getHashCode());
      } else {
        for (auto Name : *NameTable)
          NamesInProfile.insert(Name.stringRef());
      }
    }
    CoverageTracker.setProfAccForSymsInList(true);
  }

  if (FAM && !ProfileInlineReplayFile.empty()) {
    ExternalInlineAdvisor = getReplayInlineAdvisor(
        M, *FAM, Ctx, /*OriginalAdvisor=*/nullptr,
        ReplayInlinerSettings{ProfileInlineReplayFile,
                              ProfileInlineReplayScope,
                              ProfileInlineReplayFallback,
                              {ProfileInlineReplayFormat}},
        /*EmitRemarks=*/false, InlineContext{LTOPhase, InlinePass::ReplaySampleProfileInliner});
  }

  // Apply tweaks if context-sensitive or probe-based profile is available.
  if (Reader->profileIsCS() || Reader->profileIsPreInlined() ||
      Reader->profileIsProbeBased()) {
    if (!UseIterativeBFIInference.getNumOccurrences())
      UseIterativeBFIInference = true;
    if (!SampleProfileUseProfi.getNumOccurrences())
      SampleProfileUseProfi = true;
    if (!EnableExtTspBlockPlacement.getNumOccurrences())
      EnableExtTspBlockPlacement = true;
    // Enable priority-base inliner and size inline by default for CSSPGO.
    if (!ProfileSizeInline.getNumOccurrences())
      ProfileSizeInline = true;
    if (!CallsitePrioritizedInline.getNumOccurrences())
      CallsitePrioritizedInline = true;
    // For CSSPGO, we also allow recursive inline to best use context profile.
    if (!AllowRecursiveInline.getNumOccurrences())
      AllowRecursiveInline = true;

    if (Reader->profileIsPreInlined()) {
      if (!UsePreInlinerDecision.getNumOccurrences())
        UsePreInlinerDecision = true;
    }

    // Enable stale profile matching by default for probe-based profile.
    // Currently the matching relies on if the checksum mismatch is detected,
    // which is currently only available for pseudo-probe mode. Removing the
    // checksum check could cause regressions for some cases, so further tuning
    // might be needed if we want to enable it for all cases.
    if (Reader->profileIsProbeBased() &&
        !SalvageStaleProfile.getNumOccurrences()) {
      SalvageStaleProfile = true;
    }

    if (!Reader->profileIsCS()) {
      // Non-CS profile should be fine without a function size budget for the
      // inliner since the contexts in the profile are either all from inlining
      // in the prevoius build or pre-computed by the preinliner with a size
      // cap, thus they are bounded.
      if (!ProfileInlineLimitMin.getNumOccurrences())
        ProfileInlineLimitMin = std::numeric_limits<unsigned>::max();
      if (!ProfileInlineLimitMax.getNumOccurrences())
        ProfileInlineLimitMax = std::numeric_limits<unsigned>::max();
    }
  }

  if (Reader->profileIsCS()) {
    // Tracker for profiles under different context
    ContextTracker = std::make_unique<SampleContextTracker>(
        Reader->getProfiles(), &GUIDToFuncNameMap);
  }

  // Load pseudo probe descriptors for probe-based function samples.
  if (Reader->profileIsProbeBased()) {
    ProbeManager = std::make_unique<PseudoProbeManager>(M);
    if (!ProbeManager->moduleIsProbed(M)) {
      const char *Msg =
          "Pseudo-probe-based profile requires SampleProfileProbePass";
      Ctx.diagnose(DiagnosticInfoSampleProfile(M.getModuleIdentifier(), Msg,
                                               DS_Warning));
      return false;
    }
  }

  if (ReportProfileStaleness || PersistProfileStaleness ||
      SalvageStaleProfile) {
    MatchingManager =
        std::make_unique<SampleProfileMatcher>(M, *Reader, ProbeManager.get());
  }

  return true;
}

void SampleProfileMatcher::findIRAnchors(
    const Function &F, std::map<LineLocation, StringRef> &IRAnchors) {
  // For inlined code, recover the original callsite and callee by finding the
  // top-level inline frame. e.g. For frame stack "main:1 @ foo:2 @ bar:3", the
  // top-level frame is "main:1", the callsite is "1" and the callee is "foo".
  auto FindTopLevelInlinedCallsite = [](const DILocation *DIL) {
    assert((DIL && DIL->getInlinedAt()) && "No inlined callsite");
    const DILocation *PrevDIL = nullptr;
    do {
      PrevDIL = DIL;
      DIL = DIL->getInlinedAt();
    } while (DIL->getInlinedAt());

    LineLocation Callsite = FunctionSamples::getCallSiteIdentifier(DIL);
    StringRef CalleeName = PrevDIL->getSubprogramLinkageName();
    return std::make_pair(Callsite, CalleeName);
  };

  auto GetCanonicalCalleeName = [](const CallBase *CB) {
    StringRef CalleeName = UnknownIndirectCallee;
    if (Function *Callee = CB->getCalledFunction())
      CalleeName = FunctionSamples::getCanonicalFnName(Callee->getName());
    return CalleeName;
  };

  // Extract profile matching anchors in the IR.
  for (auto &BB : F) {
    for (auto &I : BB) {
      DILocation *DIL = I.getDebugLoc();
      if (!DIL)
        continue;

      if (FunctionSamples::ProfileIsProbeBased) {
        if (auto Probe = extractProbe(I)) {
          // Flatten inlined IR for the matching.
          if (DIL->getInlinedAt()) {
            IRAnchors.emplace(FindTopLevelInlinedCallsite(DIL));
          } else {
            // Use empty StringRef for basic block probe.
            StringRef CalleeName;
            if (const auto *CB = dyn_cast<CallBase>(&I)) {
              // Skip the probe inst whose callee name is "llvm.pseudoprobe".
              if (!isa<IntrinsicInst>(&I))
                CalleeName = GetCanonicalCalleeName(CB);
            }
            IRAnchors.emplace(LineLocation(Probe->Id, 0), CalleeName);
          }
        }
      } else {
        // TODO: For line-number based profile(AutoFDO), currently only support
        // find callsite anchors. In future, we need to parse all the non-call
        // instructions to extract the line locations for profile matching.
        if (!isa<CallBase>(&I) || isa<IntrinsicInst>(&I))
          continue;

        if (DIL->getInlinedAt()) {
          IRAnchors.emplace(FindTopLevelInlinedCallsite(DIL));
        } else {
          LineLocation Callsite = FunctionSamples::getCallSiteIdentifier(DIL);
          StringRef CalleeName = GetCanonicalCalleeName(dyn_cast<CallBase>(&I));
          IRAnchors.emplace(Callsite, CalleeName);
        }
      }
    }
  }
}

void SampleProfileMatcher::countMismatchedSamples(const FunctionSamples &FS) {
  const auto *FuncDesc = ProbeManager->getDesc(FS.getGUID());
  // Skip the function that is external or renamed.
  if (!FuncDesc)
    return;

  if (ProbeManager->profileIsHashMismatched(*FuncDesc, FS)) {
    MismatchedFuncHashSamples += FS.getTotalSamples();
    return;
  }
  for (const auto &I : FS.getCallsiteSamples())
    for (const auto &CS : I.second)
      countMismatchedSamples(CS.second);
}

void SampleProfileMatcher::countProfileMismatches(
    const Function &F, const FunctionSamples &FS,
    const std::map<LineLocation, StringRef> &IRAnchors,
    const std::map<LineLocation, std::unordered_set<FunctionId>>
        &ProfileAnchors) {
  [[maybe_unused]] bool IsFuncHashMismatch = false;
  if (FunctionSamples::ProfileIsProbeBased) {
    TotalFuncHashSamples += FS.getTotalSamples();
    TotalProfiledFunc++;
    const auto *FuncDesc = ProbeManager->getDesc(F);
    if (FuncDesc) {
      if (ProbeManager->profileIsHashMismatched(*FuncDesc, FS)) {
        NumMismatchedFuncHash++;
        IsFuncHashMismatch = true;
      }
      countMismatchedSamples(FS);
    }
  }

  uint64_t FuncMismatchedCallsites = 0;
  uint64_t FuncProfiledCallsites = 0;
  countProfileCallsiteMismatches(FS, IRAnchors, ProfileAnchors,
                                 FuncMismatchedCallsites,
                                 FuncProfiledCallsites);
  TotalProfiledCallsites += FuncProfiledCallsites;
  NumMismatchedCallsites += FuncMismatchedCallsites;
  LLVM_DEBUG({
    if (FunctionSamples::ProfileIsProbeBased && !IsFuncHashMismatch &&
        FuncMismatchedCallsites)
      dbgs() << "Function checksum is matched but there are "
             << FuncMismatchedCallsites << "/" << FuncProfiledCallsites
             << " mismatched callsites.\n";
  });
}

void SampleProfileMatcher::countProfileCallsiteMismatches(
    const FunctionSamples &FS,
    const std::map<LineLocation, StringRef> &IRAnchors,
    const std::map<LineLocation, std::unordered_set<FunctionId>>
        &ProfileAnchors,
    uint64_t &FuncMismatchedCallsites, uint64_t &FuncProfiledCallsites) {

  // Check if there are any callsites in the profile that does not match to any
  // IR callsites, those callsite samples will be discarded.
  for (const auto &I : ProfileAnchors) {
    const auto &Loc = I.first;
    const auto &Callees = I.second;
    assert(!Callees.empty() && "Callees should not be empty");

    StringRef IRCalleeName;
    const auto &IR = IRAnchors.find(Loc);
    if (IR != IRAnchors.end())
      IRCalleeName = IR->second;

    // Compute number of samples in the original profile.
    uint64_t CallsiteSamples = 0;
    auto CTM = FS.findCallTargetMapAt(Loc);
    if (CTM) {
      for (const auto &I : CTM.get())
        CallsiteSamples += I.second;
    }
    const auto *FSMap = FS.findFunctionSamplesMapAt(Loc);
    if (FSMap) {
      for (const auto &I : *FSMap)
        CallsiteSamples += I.second.getTotalSamples();
    }

    bool CallsiteIsMatched = false;
    // Since indirect call does not have CalleeName, check conservatively if
    // callsite in the profile is a callsite location. This is to reduce num of
    // false positive since otherwise all the indirect call samples will be
    // reported as mismatching.
    if (IRCalleeName == UnknownIndirectCallee)
      CallsiteIsMatched = true;
    else if (Callees.size() == 1 && Callees.count(getRepInFormat(IRCalleeName)))
      CallsiteIsMatched = true;

    FuncProfiledCallsites++;
    TotalCallsiteSamples += CallsiteSamples;
    if (!CallsiteIsMatched) {
      FuncMismatchedCallsites++;
      MismatchedCallsiteSamples += CallsiteSamples;
    }
  }
}

void SampleProfileMatcher::findProfileAnchors(const FunctionSamples &FS,
                                              std::map<LineLocation, std::unordered_set<FunctionId>> &ProfileAnchors) {
  auto isInvalidLineOffset = [](uint32_t LineOffset) {
    return LineOffset & 0x8000;
  };

  for (const auto &I : FS.getBodySamples()) {
    const LineLocation &Loc = I.first;
    if (isInvalidLineOffset(Loc.LineOffset))
      continue;
    for (const auto &I : I.second.getCallTargets()) {
      auto Ret = ProfileAnchors.try_emplace(Loc,
                                            std::unordered_set<FunctionId>());
      Ret.first->second.insert(I.first);
    }
  }

  for (const auto &I : FS.getCallsiteSamples()) {
    const LineLocation &Loc = I.first;
    if (isInvalidLineOffset(Loc.LineOffset))
      continue;
    const auto &CalleeMap = I.second;
    for (const auto &I : CalleeMap) {
      auto Ret = ProfileAnchors.try_emplace(Loc,
                                            std::unordered_set<FunctionId>());
      Ret.first->second.insert(I.first);
    }
  }
}

// Call target name anchor based profile fuzzy matching.
// Input:
// For IR locations, the anchor is the callee name of direct callsite; For
// profile locations, it's the call target name for BodySamples or inlinee's
// profile name for CallsiteSamples.
// Matching heuristic:
// First match all the anchors in lexical order, then split the non-anchor
// locations between the two anchors evenly, first half are matched based on the
// start anchor, second half are matched based on the end anchor.
// For example, given:
// IR locations:      [1, 2(foo), 3, 5, 6(bar), 7]
// Profile locations: [1, 2, 3(foo), 4, 7, 8(bar), 9]
// The matching gives:
//   [1,    2(foo), 3,  5,  6(bar), 7]
//    |     |       |   |     |     |
//   [1, 2, 3(foo), 4,  7,  8(bar), 9]
// The output mapping: [2->3, 3->4, 5->7, 6->8, 7->9].
void SampleProfileMatcher::runStaleProfileMatching(
    const Function &F,
    const std::map<LineLocation, StringRef> &IRAnchors,
    const std::map<LineLocation, std::unordered_set<FunctionId>>
        &ProfileAnchors,
    LocToLocMap &IRToProfileLocationMap) {
  LLVM_DEBUG(dbgs() << "Run stale profile matching for " << F.getName()
                    << "\n");
  assert(IRToProfileLocationMap.empty() &&
         "Run stale profile matching only once per function");

  std::unordered_map<FunctionId, std::set<LineLocation>>
      CalleeToCallsitesMap;
  for (const auto &I : ProfileAnchors) {
    const auto &Loc = I.first;
    const auto &Callees = I.second;
    // Filter out possible indirect calls, use direct callee name as anchor.
    if (Callees.size() == 1) {
      FunctionId CalleeName = *Callees.begin();
      const auto &Candidates = CalleeToCallsitesMap.try_emplace(
          CalleeName, std::set<LineLocation>());
      Candidates.first->second.insert(Loc);
    }
  }

  auto InsertMatching = [&](const LineLocation &From, const LineLocation &To) {
    // Skip the unchanged location mapping to save memory.
    if (From != To)
      IRToProfileLocationMap.insert({From, To});
  };

  // Use function's beginning location as the initial anchor.
  int32_t LocationDelta = 0;
  SmallVector<LineLocation> LastMatchedNonAnchors;

  for (const auto &IR : IRAnchors) {
    const auto &Loc = IR.first;
    auto CalleeName = IR.second;
    bool IsMatchedAnchor = false;
    // Match the anchor location in lexical order.
    if (!CalleeName.empty()) {
      auto CandidateAnchors = CalleeToCallsitesMap.find(
          getRepInFormat(CalleeName));
      if (CandidateAnchors != CalleeToCallsitesMap.end() &&
          !CandidateAnchors->second.empty()) {
        auto CI = CandidateAnchors->second.begin();
        const auto Candidate = *CI;
        CandidateAnchors->second.erase(CI);
        InsertMatching(Loc, Candidate);
        LLVM_DEBUG(dbgs() << "Callsite with callee:" << CalleeName
                          << " is matched from " << Loc << " to " << Candidate
                          << "\n");
        LocationDelta = Candidate.LineOffset - Loc.LineOffset;

        // Match backwards for non-anchor locations.
        // The locations in LastMatchedNonAnchors have been matched forwards
        // based on the previous anchor, spilt it evenly and overwrite the
        // second half based on the current anchor.
        for (size_t I = (LastMatchedNonAnchors.size() + 1) / 2;
             I < LastMatchedNonAnchors.size(); I++) {
          const auto &L = LastMatchedNonAnchors[I];
          uint32_t CandidateLineOffset = L.LineOffset + LocationDelta;
          LineLocation Candidate(CandidateLineOffset, L.Discriminator);
          InsertMatching(L, Candidate);
          LLVM_DEBUG(dbgs() << "Location is rematched backwards from " << L
                            << " to " << Candidate << "\n");
        }

        IsMatchedAnchor = true;
        LastMatchedNonAnchors.clear();
      }
    }

    // Match forwards for non-anchor locations.
    if (!IsMatchedAnchor) {
      uint32_t CandidateLineOffset = Loc.LineOffset + LocationDelta;
      LineLocation Candidate(CandidateLineOffset, Loc.Discriminator);
      InsertMatching(Loc, Candidate);
      LLVM_DEBUG(dbgs() << "Location is matched from " << Loc << " to "
                        << Candidate << "\n");
      LastMatchedNonAnchors.emplace_back(Loc);
    }
  }
}

void SampleProfileMatcher::runOnFunction(const Function &F) {
  // We need to use flattened function samples for matching.
  // Unlike IR, which includes all callsites from the source code, the callsites
  // in profile only show up when they are hit by samples, i,e. the profile
  // callsites in one context may differ from those in another context. To get
  // the maximum number of callsites, we merge the function profiles from all
  // contexts, aka, the flattened profile to find profile anchors.
  const auto *FSFlattened = getFlattenedSamplesFor(F);
  if (!FSFlattened)
    return;

  // Anchors for IR. It's a map from IR location to callee name, callee name is
  // empty for non-call instruction and use a dummy name(UnknownIndirectCallee)
  // for unknown indrect callee name.
  std::map<LineLocation, StringRef> IRAnchors;
  findIRAnchors(F, IRAnchors);
  // Anchors for profile. It's a map from callsite location to a set of callee
  // name.
  std::map<LineLocation, std::unordered_set<FunctionId>> ProfileAnchors;
  findProfileAnchors(*FSFlattened, ProfileAnchors);

  // Detect profile mismatch for profile staleness metrics report.
  // Skip reporting the metrics for imported functions.
  if (!GlobalValue::isAvailableExternallyLinkage(F.getLinkage()) &&
      (ReportProfileStaleness || PersistProfileStaleness)) {
    // Use top-level nested FS for counting profile mismatch metrics since
    // currently once a callsite is mismatched, all its children profiles are
    // dropped.
    if (const auto *FS = Reader.getSamplesFor(F))
      countProfileMismatches(F, *FS, IRAnchors, ProfileAnchors);
  }

  // Run profile matching for checksum mismatched profile, currently only
  // support for pseudo-probe.
  if (SalvageStaleProfile && FunctionSamples::ProfileIsProbeBased &&
      !ProbeManager->profileIsValid(F, *FSFlattened)) {
    // The matching result will be saved to IRToProfileLocationMap, create a new
    // map for each function.
    runStaleProfileMatching(F, IRAnchors, ProfileAnchors,
                            getIRToProfileLocationMap(F));
  }
}

void SampleProfileMatcher::runOnModule() {
  ProfileConverter::flattenProfile(Reader.getProfiles(), FlattenedProfiles,
                                   FunctionSamples::ProfileIsCS);
  for (auto &F : M) {
    if (F.isDeclaration() || !F.hasFnAttribute("use-sample-profile"))
      continue;
    runOnFunction(F);
  }
  if (SalvageStaleProfile)
    distributeIRToProfileLocationMap();

  if (ReportProfileStaleness) {
    if (FunctionSamples::ProfileIsProbeBased) {
      errs() << "(" << NumMismatchedFuncHash << "/" << TotalProfiledFunc << ")"
             << " of functions' profile are invalid and "
             << " (" << MismatchedFuncHashSamples << "/" << TotalFuncHashSamples
             << ")"
             << " of samples are discarded due to function hash mismatch.\n";
    }
    errs() << "(" << NumMismatchedCallsites << "/" << TotalProfiledCallsites
           << ")"
           << " of callsites' profile are invalid and "
           << "(" << MismatchedCallsiteSamples << "/" << TotalCallsiteSamples
           << ")"
           << " of samples are discarded due to callsite location mismatch.\n";
  }

  if (PersistProfileStaleness) {
    LLVMContext &Ctx = M.getContext();
    MDBuilder MDB(Ctx);

    SmallVector<std::pair<StringRef, uint64_t>> ProfStatsVec;
    if (FunctionSamples::ProfileIsProbeBased) {
      ProfStatsVec.emplace_back("NumMismatchedFuncHash", NumMismatchedFuncHash);
      ProfStatsVec.emplace_back("TotalProfiledFunc", TotalProfiledFunc);
      ProfStatsVec.emplace_back("MismatchedFuncHashSamples",
                                MismatchedFuncHashSamples);
      ProfStatsVec.emplace_back("TotalFuncHashSamples", TotalFuncHashSamples);
    }

    ProfStatsVec.emplace_back("NumMismatchedCallsites", NumMismatchedCallsites);
    ProfStatsVec.emplace_back("TotalProfiledCallsites", TotalProfiledCallsites);
    ProfStatsVec.emplace_back("MismatchedCallsiteSamples",
                              MismatchedCallsiteSamples);
    ProfStatsVec.emplace_back("TotalCallsiteSamples", TotalCallsiteSamples);

    auto *MD = MDB.createLLVMStats(ProfStatsVec);
    auto *NMD = M.getOrInsertNamedMetadata("llvm.stats");
    NMD->addOperand(MD);
  }
}

void SampleProfileMatcher::distributeIRToProfileLocationMap(
    FunctionSamples &FS) {
  const auto ProfileMappings = FuncMappings.find(FS.getFuncName());
  if (ProfileMappings != FuncMappings.end()) {
    FS.setIRToProfileLocationMap(&(ProfileMappings->second));
  }

  for (auto &Inlinees : FS.getCallsiteSamples()) {
    for (auto FS : Inlinees.second) {
      distributeIRToProfileLocationMap(FS.second);
    }
  }
}

// Use a central place to distribute the matching results. Outlined and inlined
// profile with the function name will be set to the same pointer.
void SampleProfileMatcher::distributeIRToProfileLocationMap() {
  for (auto &I : Reader.getProfiles()) {
    distributeIRToProfileLocationMap(I.second);
  }
}

bool SampleProfileLoader::runOnModule(Module &M, ModuleAnalysisManager *AM,
                                      ProfileSummaryInfo *_PSI,
                                      LazyCallGraph &CG) {
  GUIDToFuncNameMapper Mapper(M, *Reader, GUIDToFuncNameMap);

  PSI = _PSI;
  if (M.getProfileSummary(/* IsCS */ false) == nullptr) {
    M.setProfileSummary(Reader->getSummary().getMD(M.getContext()),
                        ProfileSummary::PSK_Sample);
    PSI->refresh();
  }
  // Compute the total number of samples collected in this profile.
  for (const auto &I : Reader->getProfiles())
    TotalCollectedSamples += I.second.getTotalSamples();

  auto Remapper = Reader->getRemapper();
  // Populate the symbol map.
  for (const auto &N_F : M.getValueSymbolTable()) {
    StringRef OrigName = N_F.getKey();
    Function *F = dyn_cast<Function>(N_F.getValue());
    if (F == nullptr || OrigName.empty())
      continue;
    SymbolMap[FunctionId(OrigName)] = F;
    StringRef NewName = FunctionSamples::getCanonicalFnName(*F);
    if (OrigName != NewName && !NewName.empty()) {
      auto r = SymbolMap.emplace(FunctionId(NewName), F);
      // Failiing to insert means there is already an entry in SymbolMap,
      // thus there are multiple functions that are mapped to the same
      // stripped name. In this case of name conflicting, set the value
      // to nullptr to avoid confusion.
      if (!r.second)
        r.first->second = nullptr;
      OrigName = NewName;
    }
    // Insert the remapped names into SymbolMap.
    if (Remapper) {
      if (auto MapName = Remapper->lookUpNameInProfile(OrigName)) {
        if (*MapName != OrigName && !MapName->empty())
          SymbolMap.emplace(FunctionId(*MapName), F);
      }
    }
  }
  assert(SymbolMap.count(FunctionId()) == 0 &&
         "No empty StringRef should be added in SymbolMap");

  if (ReportProfileStaleness || PersistProfileStaleness ||
      SalvageStaleProfile) {
    MatchingManager->runOnModule();
  }

  bool retval = false;
  for (auto *F : buildFunctionOrder(M, CG)) {
    assert(!F->isDeclaration());
    clearFunctionData();
    retval |= runOnFunction(*F, AM);
  }

  // Account for cold calls not inlined....
  if (!FunctionSamples::ProfileIsCS)
    for (const std::pair<Function *, NotInlinedProfileInfo> &pair :
         notInlinedCallInfo)
      updateProfileCallee(pair.first, pair.second.entryCount);

  return retval;
}

bool SampleProfileLoader::runOnFunction(Function &F, ModuleAnalysisManager *AM) {
  LLVM_DEBUG(dbgs() << "\n\nProcessing Function " << F.getName() << "\n");
  DILocation2SampleMap.clear();
  // By default the entry count is initialized to -1, which will be treated
  // conservatively by getEntryCount as the same as unknown (None). This is
  // to avoid newly added code to be treated as cold. If we have samples
  // this will be overwritten in emitAnnotations.
  uint64_t initialEntryCount = -1;

  ProfAccForSymsInList = ProfileAccurateForSymsInList && PSL;
  if (ProfileSampleAccurate || F.hasFnAttribute("profile-sample-accurate")) {
    // initialize all the function entry counts to 0. It means all the
    // functions without profile will be regarded as cold.
    initialEntryCount = 0;
    // profile-sample-accurate is a user assertion which has a higher precedence
    // than symbol list. When profile-sample-accurate is on, ignore symbol list.
    ProfAccForSymsInList = false;
  }
  CoverageTracker.setProfAccForSymsInList(ProfAccForSymsInList);

  // PSL -- profile symbol list include all the symbols in sampled binary.
  // If ProfileAccurateForSymsInList is enabled, PSL is used to treat
  // old functions without samples being cold, without having to worry
  // about new and hot functions being mistakenly treated as cold.
  if (ProfAccForSymsInList) {
    // Initialize the entry count to 0 for functions in the list.
    if (PSL->contains(F.getName()))
      initialEntryCount = 0;

    // Function in the symbol list but without sample will be regarded as
    // cold. To minimize the potential negative performance impact it could
    // have, we want to be a little conservative here saying if a function
    // shows up in the profile, no matter as outline function, inline instance
    // or call targets, treat the function as not being cold. This will handle
    // the cases such as most callsites of a function are inlined in sampled
    // binary but not inlined in current build (because of source code drift,
    // imprecise debug information, or the callsites are all cold individually
    // but not cold accumulatively...), so the outline function showing up as
    // cold in sampled binary will actually not be cold after current build.
    StringRef CanonName = FunctionSamples::getCanonicalFnName(F);
    if ((FunctionSamples::UseMD5 &&
         GUIDsInProfile.count(Function::getGUID(CanonName))) ||
        (!FunctionSamples::UseMD5 && NamesInProfile.count(CanonName)))
      initialEntryCount = -1;
  }

  // Initialize entry count when the function has no existing entry
  // count value.
  if (!F.getEntryCount())
    F.setEntryCount(ProfileCount(initialEntryCount, Function::PCT_Real));
  std::unique_ptr<OptimizationRemarkEmitter> OwnedORE;
  if (AM) {
    auto &FAM =
        AM->getResult<FunctionAnalysisManagerModuleProxy>(*F.getParent())
            .getManager();
    ORE = &FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
  } else {
    OwnedORE = std::make_unique<OptimizationRemarkEmitter>(&F);
    ORE = OwnedORE.get();
  }

  if (FunctionSamples::ProfileIsCS)
    Samples = ContextTracker->getBaseSamplesFor(F);
  else {
    Samples = Reader->getSamplesFor(F);
    // Try search in previously inlined functions that were split or duplicated
    // into base.
    if (!Samples) {
      StringRef CanonName = FunctionSamples::getCanonicalFnName(F);
      auto It = OutlineFunctionSamples.find(FunctionId(CanonName));
      if (It != OutlineFunctionSamples.end()) {
        Samples = &It->second;
      } else if (auto Remapper = Reader->getRemapper()) {
        if (auto RemppedName = Remapper->lookUpNameInProfile(CanonName)) {
          It = OutlineFunctionSamples.find(FunctionId(*RemppedName));
          if (It != OutlineFunctionSamples.end())
            Samples = &It->second;
        }
      }
    }
  }

  if (Samples && !Samples->empty())
    return emitAnnotations(F);
  return false;
}
SampleProfileLoaderPass::SampleProfileLoaderPass(
    std::string File, std::string RemappingFile, ThinOrFullLTOPhase LTOPhase,
    IntrusiveRefCntPtr<vfs::FileSystem> FS)
    : ProfileFileName(File), ProfileRemappingFileName(RemappingFile),
      LTOPhase(LTOPhase), FS(std::move(FS)) {}

PreservedAnalyses SampleProfileLoaderPass::run(Module &M,
                                               ModuleAnalysisManager &AM) {
  FunctionAnalysisManager &FAM =
      AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();

  auto GetAssumptionCache = [&](Function &F) -> AssumptionCache & {
    return FAM.getResult<AssumptionAnalysis>(F);
  };
  auto GetTTI = [&](Function &F) -> TargetTransformInfo & {
    return FAM.getResult<TargetIRAnalysis>(F);
  };
  auto GetTLI = [&](Function &F) -> const TargetLibraryInfo & {
    return FAM.getResult<TargetLibraryAnalysis>(F);
  };

  if (!FS)
    FS = vfs::getRealFileSystem();

  SampleProfileLoader SampleLoader(
      ProfileFileName.empty() ? SampleProfileFile : ProfileFileName,
      ProfileRemappingFileName.empty() ? SampleProfileRemappingFile
                                       : ProfileRemappingFileName,
      LTOPhase, FS, GetAssumptionCache, GetTTI, GetTLI);

  if (!SampleLoader.doInitialization(M, &FAM))
    return PreservedAnalyses::all();

  ProfileSummaryInfo *PSI = &AM.getResult<ProfileSummaryAnalysis>(M);
  LazyCallGraph &CG = AM.getResult<LazyCallGraphAnalysis>(M);
  if (!SampleLoader.runOnModule(M, &AM, PSI, CG))
    return PreservedAnalyses::all();

  return PreservedAnalyses::none();
}