[llvm][NFC] Use CallBase instead of Instruction in ProfileSummaryInfo

[mtrofin]

Summary

getProfileCount requires the parameter be a valid CallBase, and its uses
reflect that.

[dblaikie]

Looks good!

llvm/include/llvm/Analysis/ProfileSummaryInfo.h

 #include "llvm/IR/ProfileSummary.h"
 #include "llvm/Pass.h"
 #include <memory>
 
 namespace llvm {
 class BasicBlock;
 class BlockFrequencyInfo;
 class CallBase;
-class Instruction;
 class ProfileSummary;
 /// Analysis providing profile information.
 ///
 /// This is an immutable analysis pass that provides ability to query global
 /// (program-level) profile information. The main APIs are isHotCount and
 /// isColdCount that tells whether a given profile count is considered hot/cold
 /// based on the profile summary. This also provides convenience methods to
 /// check whether a function is hot or cold.
public:
   /// summary is immutable after it is annotated on the module, we return false
   /// here.
   bool invalidate(Module &, const PreservedAnalyses &,
                   ModuleAnalysisManager::Invalidator &) {
     return false;
   }
 
   /// Returns the profile count for \p CallInst.
-  Optional<uint64_t> getProfileCount(const Instruction *CallInst,
+  Optional<uint64_t> getProfileCount(const CallBase &CallInst,
                                      BlockFrequencyInfo *BFI,
                                      bool AllowSynthetic = false);
   /// Returns true if the working set size of the code is considered huge.
   bool hasHugeWorkingSetSize();
   /// Returns true if the working set size of the code is considered large.
   bool hasLargeWorkingSetSize();
   /// Returns true if \p F has hot function entry.
   bool isFunctionEntryHot(const Function *F);

llvm/lib/Analysis/ModuleSummaryAnalysis.cpp

for (const Instruction &I : BB) {
         if (CI && CalledFunction->isIntrinsic()) {
           addIntrinsicToSummary(
               CI, TypeTests, TypeTestAssumeVCalls, TypeCheckedLoadVCalls,
               TypeTestAssumeConstVCalls, TypeCheckedLoadConstVCalls, DT);
           continue;
         }
         // We should have named any anonymous globals
         assert(CalledFunction->hasName());
-        auto ScaledCount = PSI->getProfileCount(&I, BFI);
+        auto ScaledCount = PSI->getProfileCount(*CB, BFI);
         auto Hotness = ScaledCount ? getHotness(ScaledCount.getValue(), PSI)
                                    : CalleeInfo::HotnessType::Unknown;
         if (ForceSummaryEdgesCold != FunctionSummary::FSHT_None)
           Hotness = CalleeInfo::HotnessType::Cold;
 
         // Use the original CalledValue, in case it was an alias. We want
         // to record the call edge to the alias in that case. Eventually
         // an alias summary will be created to associate the alias and

llvm/lib/Analysis/ProfileSummaryInfo.cpp

bool ProfileSummaryInfo::computeSummary() {
   // This will actually return PSK_Instr or PSK_Sample summary.
   SummaryMD = M.getProfileSummary(/* IsCS */ false);
   if (!SummaryMD)
     return false;
   Summary.reset(ProfileSummary::getFromMD(SummaryMD));
   return true;
 }
 
-// FIXME(CallSite): the parameter should be a CallBase.
-Optional<uint64_t>
-ProfileSummaryInfo::getProfileCount(const Instruction *Inst,
+Optional<uint64_t> ProfileSummaryInfo::getProfileCount(const CallBase &Call,
                                                        BlockFrequencyInfo *BFI,
                                                        bool AllowSynthetic) {
-  if (!Inst)
-    return None;
-  assert((isa<CallInst>(Inst) || isa<InvokeInst>(Inst)) &&
+  assert((isa<CallInst>(Call) || isa<InvokeInst>(Call)) &&
          "We can only get profile count for call/invoke instruction.");
   if (hasSampleProfile()) {
     // In sample PGO mode, check if there is a profile metadata on the
     // instruction. If it is present, determine hotness solely based on that,
     // since the sampled entry count may not be accurate. If there is no
     // annotated on the instruction, return None.
     uint64_t TotalCount;
-    if (Inst->extractProfTotalWeight(TotalCount))
+    if (Call.extractProfTotalWeight(TotalCount))
       return TotalCount;
     return None;
   }
   if (BFI)
-    return BFI->getBlockProfileCount(Inst->getParent(), AllowSynthetic);
+    return BFI->getBlockProfileCount(Call.getParent(), AllowSynthetic);
   return None;
 }
 
 /// Returns true if the function's entry is hot. If it returns false, it
 /// either means it is not hot or it is unknown whether it is hot or not (for
 /// example, no profile data is available).
 bool ProfileSummaryInfo::isFunctionEntryHot(const Function *F) {
   if (!F || !computeSummary())
if (auto FunctionCount = F->getEntryCount())
     if (isHotCount(FunctionCount.getCount()))
       return true;
 
   if (hasSampleProfile()) {
     uint64_t TotalCallCount = 0;
     for (const auto &BB : *F)
       for (const auto &I : BB)
         if (isa<CallInst>(I) || isa<InvokeInst>(I))
-          if (auto CallCount = getProfileCount(&I, nullptr))
+          if (auto CallCount = getProfileCount(cast<CallBase>(I), nullptr))
             TotalCallCount += CallCount.getValue();
     if (isHotCount(TotalCallCount))
       return true;
   }
   for (const auto &BB : *F)
     if (isHotBlock(&BB, &BFI))
       return true;
   return false;
if (auto FunctionCount = F->getEntryCount())
     if (!isColdCount(FunctionCount.getCount()))
       return false;
 
   if (hasSampleProfile()) {
     uint64_t TotalCallCount = 0;
     for (const auto &BB : *F)
       for (const auto &I : BB)
         if (isa<CallInst>(I) || isa<InvokeInst>(I))
-          if (auto CallCount = getProfileCount(&I, nullptr))
+          if (auto CallCount = getProfileCount(cast<CallBase>(I), nullptr))
             TotalCallCount += CallCount.getValue();
     if (!isColdCount(TotalCallCount))
       return false;
   }
   for (const auto &BB : *F)
     if (!isColdBlock(&BB, &BFI))
       return false;
   return true;
if (!isHot &&
         !isColdCountNthPercentile(PercentileCutoff, FunctionCount.getCount()))
       return false;
   }
   if (hasSampleProfile()) {
     uint64_t TotalCallCount = 0;
     for (const auto &BB : *F)
       for (const auto &I : BB)
         if (isa<CallInst>(I) || isa<InvokeInst>(I))
-          if (auto CallCount = getProfileCount(&I, nullptr))
+          if (auto CallCount = getProfileCount(cast<CallBase>(I), nullptr))
             TotalCallCount += CallCount.getValue();
     if (isHot && isHotCountNthPercentile(PercentileCutoff, TotalCallCount))
       return true;
     if (!isHot && !isColdCountNthPercentile(PercentileCutoff, TotalCallCount))
       return false;
   }
   for (const auto &BB : *F) {
     if (isHot && isHotBlockNthPercentile(PercentileCutoff, &BB, &BFI))

 bool ProfileSummaryInfo::isColdBlockNthPercentile(int PercentileCutoff,
                                                   const BasicBlock *BB,
                                                   BlockFrequencyInfo *BFI) {
   return isHotOrColdBlockNthPercentile<false>(PercentileCutoff, BB, BFI);
 }
 
 bool ProfileSummaryInfo::isHotCallSite(const CallBase &CB,
                                        BlockFrequencyInfo *BFI) {
-  auto C = getProfileCount(&CB, BFI);
+  auto C = getProfileCount(CB, BFI);
   return C && isHotCount(*C);
 }
 
 bool ProfileSummaryInfo::isColdCallSite(const CallBase &CB,
                                         BlockFrequencyInfo *BFI) {
-  auto C = getProfileCount(&CB, BFI);
+  auto C = getProfileCount(CB, BFI);
   if (C)
     return isColdCount(*C);
 
   // In SamplePGO, if the caller has been sampled, and there is no profile
   // annotated on the callsite, we consider the callsite as cold.
   return hasSampleProfile() && CB.getCaller()->hasProfileData();
 }
 

llvm/lib/Transforms/Utils/InlineFunction.cpp

static void updateCallerBFI(BasicBlock *CallSiteBlock,
   CallerBFI->setBlockFreqAndScale(
       EntryClone, CallerBFI->getBlockFreq(CallSiteBlock).getFrequency(),
       ClonedBBs);
 }
 
 /// Update the branch metadata for cloned call instructions.
 static void updateCallProfile(Function *Callee, const ValueToValueMapTy &VMap,
                               const ProfileCount &CalleeEntryCount,
-                              const Instruction *TheCall,
-                              ProfileSummaryInfo *PSI,
+                              const CallBase &TheCall, ProfileSummaryInfo *PSI,
                               BlockFrequencyInfo *CallerBFI) {
   if (!CalleeEntryCount.hasValue() || CalleeEntryCount.isSynthetic() ||
       CalleeEntryCount.getCount() < 1)
     return;
   auto CallSiteCount = PSI ? PSI->getProfileCount(TheCall, CallerBFI) : None;
   int64_t CallCount =
       std::min(CallSiteCount.hasValue() ? CallSiteCount.getValue() : 0,
                CalleeEntryCount.getCount());
Function::iterator FirstNewBlock;
     // Remember the first block that is newly cloned over.
     FirstNewBlock = LastBlock; ++FirstNewBlock;
 
     if (IFI.CallerBFI != nullptr && IFI.CalleeBFI != nullptr)
       // Update the BFI of blocks cloned into the caller.
       updateCallerBFI(OrigBB, VMap, IFI.CallerBFI, IFI.CalleeBFI,
                       CalledFunc->front());
 
-    updateCallProfile(CalledFunc, VMap, CalledFunc->getEntryCount(), &CB,
+    updateCallProfile(CalledFunc, VMap, CalledFunc->getEntryCount(), CB,
                       IFI.PSI, IFI.CallerBFI);
 
     // Inject byval arguments initialization.
     for (std::pair<Value*, Value*> &Init : ByValInit)
       HandleByValArgumentInit(Init.first, Init.second, Caller->getParent(),
                               &*FirstNewBlock, IFI);
 
     Optional<OperandBundleUse> ParentDeopt =