Refactor the PSI to extract getCallSiteCount and remove checks for profile type.

[danielcdh]

Summary

There is no need to check profile count as only CallInst will have metadata attached.

[danielcdh]

update

[eraman]

include/llvm/Analysis/ProfileSummaryInfo.h

57

May be rename it to getProfileCount()?

lib/Analysis/ProfileSummaryInfo.cpp

73

Add an assert to check this is CallInst or InvokeInst.

[danielcdh]

update

[eraman]

Please run clang-format to fix the formatting issue before committing this patch.

include/llvm/Analysis/ProfileSummaryInfo.h

59

Not formatted correctly.

[danielcdh]

update

include/llvm/Analysis/ProfileSummaryInfo.h

 class ProfileSummaryInfo {
 private:
   Module &M;
   std::unique_ptr<ProfileSummary> Summary;
   bool computeSummary();
   void computeThresholds();
   // Count thresholds to answer isHotCount and isColdCount queries.
   Optional<uint64_t> HotCountThreshold, ColdCountThreshold;
-  bool extractProfTotalWeight(const Instruction *TI, uint64_t &TotalCount);
 
 public:
   ProfileSummaryInfo(Module &M) : M(M) {}
   ProfileSummaryInfo(ProfileSummaryInfo &&Arg)
       : M(Arg.M), Summary(std::move(Arg.Summary)) {}
+  /// Returns the profile count for \p CallInst.
+  static Optional<uint64_t> getProfileCount(const Instruction *CallInst,
+                                            BlockFrequencyInfo *BFI);
   /// \brief Returns true if \p F has hot function entry.
   bool isFunctionEntryHot(const Function *F);
   /// \brief Returns true if \p F has cold function entry.
   bool isFunctionEntryCold(const Function *F);
   /// \brief Returns true if \p F is a hot function.
   bool isHotCount(uint64_t C);
   /// \brief Returns true if count \p C is considered cold.
   bool isColdCount(uint64_t C);

lib/Analysis/ProfileSummaryInfo.cpp

if (Summary)
     return true;
   auto *SummaryMD = M.getProfileSummary();
   if (!SummaryMD)
     return false;
   Summary.reset(ProfileSummary::getFromMD(SummaryMD));
   return true;
 }
 
+Optional<uint64_t>
+ProfileSummaryInfo::getProfileCount(const Instruction *Inst,
+                                    BlockFrequencyInfo *BFI) {
+  if (!Inst)
+    return None;
+  assert((isa<CallInst>(Inst) || isa<InvokeInst>(Inst)) &&
+         "We can only get profile count for call/invoke instruction.");
+  // Check if there is a profile metadata on the instruction. If it is present,
+  // determine hotness solely based on that.
+  uint64_t TotalCount;
+  if (Inst->extractProfTotalWeight(TotalCount))
+    return TotalCount;
+  if (BFI)
+    return BFI->getBlockProfileCount(Inst->getParent());
+  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())
     return false;
   auto FunctionCount = F->getEntryCount();
   // FIXME: The heuristic used below for determining hotness is based on

 }
 
 bool ProfileSummaryInfo::isColdBB(const BasicBlock *B,
                                   BlockFrequencyInfo *BFI) {
   auto Count = BFI->getBlockProfileCount(B);
   return Count && isColdCount(*Count);
 }
 
-bool ProfileSummaryInfo::extractProfTotalWeight(const Instruction *I,
-                                                uint64_t &TotalCount) {
-  if (!computeSummary())
-    return false;
-  // Use profile weight on metadata only for sample profiling where block counts
-  // could differ from the count of an instruction within the block.
-  if (Summary.get()->getKind() != ProfileSummary::PSK_Sample)
-    return false;
-
-  return (isa<CallInst>(I) ||
-          (isa<TerminatorInst>(I) && !isa<ReturnInst>(I))) &&
-         I->extractProfTotalWeight(TotalCount);
-}
-
 bool ProfileSummaryInfo::isHotCallSite(const CallSite &CS,
                                        BlockFrequencyInfo *BFI) {
-  auto *CallInst = CS.getInstruction();
-  if (!CS)
-    return false;
-  // Check if there is a profile metadata on the instruction. If it is present,
-  // determine hotness solely based on that.
-  uint64_t TotalCount;
-  if (extractProfTotalWeight(CallInst, TotalCount))
-    return isHotCount(TotalCount);
-  return BFI && isHotBB(CallInst->getParent(), BFI);
+  auto C = getProfileCount(CS.getInstruction(), BFI);
+  return C && isHotCount(*C);
 }
 
 bool ProfileSummaryInfo::isColdCallSite(const CallSite &CS,
                                         BlockFrequencyInfo *BFI) {
-  auto *CallInst = CS.getInstruction();
-  if (!CS)
-    return false;
-  // Check if there is a profile metadata on the instruction. If it is present,
-  // and tells that the callsite is not cold, then return false;
-  uint64_t TotalCount;
-  if (extractProfTotalWeight(CallInst, TotalCount) && !isColdCount(TotalCount))
-    return false;
-  return BFI && isColdBB(CallInst->getParent(), BFI);
+  auto C = getProfileCount(CS.getInstruction(), BFI);
+  return C && isColdCount(*C);
 }
 
 INITIALIZE_PASS(ProfileSummaryInfoWrapperPass, "profile-summary-info",
                 "Profile summary info", false, true)
 
 ProfileSummaryInfoWrapperPass::ProfileSummaryInfoWrapperPass()
     : ImmutablePass(ID) {
   initializeProfileSummaryInfoWrapperPassPass(*PassRegistry::getPassRegistry());

unittests/Analysis/ProfileSummaryInfoTest.cpp

TEST_F(ProfileSummaryInfoTest, InstrProf) {
   EXPECT_TRUE(PSI.isHotBB(BB3, &BFI));
 
   CallSite CS1(BB1->getFirstNonPHI());
   auto *CI2 = BB2->getFirstNonPHI();
   CallSite CS2(CI2);
 
   EXPECT_TRUE(PSI.isHotCallSite(CS1, &BFI));
   EXPECT_FALSE(PSI.isHotCallSite(CS2, &BFI));
-
-  // Test that adding an MD_prof metadata with a hot count on CS2 does not
-  // change itas hotness as it has no effect in instrumented profiling.
-  MDBuilder MDB(M->getContext());
-  CI2->setMetadata(llvm::LLVMContext::MD_prof, MDB.createBranchWeights({400}));
-  EXPECT_FALSE(PSI.isHotCallSite(CS2, &BFI));
 }
 
 TEST_F(ProfileSummaryInfoTest, SampleProf) {
   auto M = makeLLVMModule("SampleProfile");
   Function *F = M->getFunction("f");
   ProfileSummaryInfo PSI = buildPSI(M.get());
 
   BasicBlock &BB0 = F->getEntryBlock();