[InstrProf] Factor out getRecord() and use NamedInstrProfRecord (#145417)

Factor out code in `populateCounters()` and `populateCoverage()` used to
grab the record into `PGOUseFunc::getRecord()` to reduce code
duplication. And return `NamedInstrProfRecord` in `getInstrProfRecord()`
to avoid an unnecessary cast. No functional change is intented.

Author: ellishg

clang/lib/CodeGen/CodeGenPGO.cpp

@@ -1423,8 +1423,7 @@ void CodeGenPGO::loadRegionCounts(llvm::IndexedInstrProfReader *PGOReader,
                                   bool IsInMainFile) {
   CGM.getPGOStats().addVisited(IsInMainFile);
   RegionCounts.clear();
-  llvm::Expected<llvm::InstrProfRecord> RecordExpected =
-      PGOReader->getInstrProfRecord(FuncName, FunctionHash);
+  auto RecordExpected = PGOReader->getInstrProfRecord(FuncName, FunctionHash);
   if (auto E = RecordExpected.takeError()) {
     auto IPE = std::get<0>(llvm::InstrProfError::take(std::move(E)));
     if (IPE == llvm::instrprof_error::unknown_function)

llvm/include/llvm/ProfileData/InstrProfReader.h

@@ -824,7 +824,7 @@ class LLVM_ABI IndexedInstrProfReader : public InstrProfReader {
   /// functions, MismatchedFuncSum returns the maximum. If \c FuncName is not
   /// found, try to lookup \c DeprecatedFuncName to handle profiles built by
   /// older compilers.
-  Expected<InstrProfRecord>
+  Expected<NamedInstrProfRecord>
   getInstrProfRecord(StringRef FuncName, uint64_t FuncHash,
                      StringRef DeprecatedFuncName = "",
                      uint64_t *MismatchedFuncSum = nullptr);

llvm/lib/ProfileData/InstrProfReader.cpp

@@ -1369,7 +1369,7 @@ InstrProfSymtab &IndexedInstrProfReader::getSymtab() {
   return *Symtab;
 }
 
-Expected<InstrProfRecord> IndexedInstrProfReader::getInstrProfRecord(
+Expected<NamedInstrProfRecord> IndexedInstrProfReader::getInstrProfRecord(
     StringRef FuncName, uint64_t FuncHash, StringRef DeprecatedFuncName,
     uint64_t *MismatchedFuncSum) {
   ArrayRef<NamedInstrProfRecord> Data;
@@ -1572,7 +1572,7 @@ memprof::AllMemProfData IndexedMemProfReader::getAllMemProfData() const {
 Error IndexedInstrProfReader::getFunctionCounts(StringRef FuncName,
                                                 uint64_t FuncHash,
                                                 std::vector<uint64_t> &Counts) {
-  Expected<InstrProfRecord> Record = getInstrProfRecord(FuncName, FuncHash);
+  auto Record = getInstrProfRecord(FuncName, FuncHash);
   if (Error E = Record.takeError())
     return error(std::move(E));
 
@@ -1583,7 +1583,7 @@ Error IndexedInstrProfReader::getFunctionCounts(StringRef FuncName,
 Error IndexedInstrProfReader::getFunctionBitmap(StringRef FuncName,
                                                 uint64_t FuncHash,
                                                 BitVector &Bitmap) {
-  Expected<InstrProfRecord> Record = getInstrProfRecord(FuncName, FuncHash);
+  auto Record = getInstrProfRecord(FuncName, FuncHash);
   if (Error E = Record.takeError())
     return error(std::move(E));
 

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

@@ -1176,15 +1176,20 @@ class PGOUseFunc {
 
   void handleInstrProfError(Error Err, uint64_t MismatchedFuncSum);
 
+  /// Get the profile record, assign it to \p ProfileRecord, handle errors if
+  /// necessary, and assign \p ProgramMaxCount. \returns true if there are no
+  /// errors.
+  bool getRecord(IndexedInstrProfReader *PGOReader);
+
   // Read counts for the instrumented BB from profile.
-  bool readCounters(IndexedInstrProfReader *PGOReader, bool &AllZeros,
+  bool readCounters(bool &AllZeros,
                     InstrProfRecord::CountPseudoKind &PseudoKind);
 
   // Populate the counts for all BBs.
   void populateCounters();
 
   // Set block coverage based on profile coverage values.
-  void populateCoverage(IndexedInstrProfReader *PGOReader);
+  void populateCoverage();
 
   // Set the branch weights based on the count values.
   void setBranchWeights();
@@ -1208,7 +1213,7 @@ class PGOUseFunc {
   uint64_t getFuncHash() const { return FuncInfo.FunctionHash; }
 
   // Return the profile record for this function;
-  InstrProfRecord &getProfileRecord() { return ProfileRecord; }
+  NamedInstrProfRecord &getProfileRecord() { return ProfileRecord; }
 
   // Return the auxiliary BB information.
   PGOUseBBInfo &getBBInfo(const BasicBlock *BB) const {
@@ -1246,7 +1251,7 @@ class PGOUseFunc {
   uint32_t ProfileCountSize = 0;
 
   // ProfileRecord for this function.
-  InstrProfRecord ProfileRecord;
+  NamedInstrProfRecord ProfileRecord;
 
   // Function hotness info derived from profile.
   FuncFreqAttr FreqAttr;
@@ -1441,21 +1446,26 @@ void PGOUseFunc::handleInstrProfError(Error Err, uint64_t MismatchedFuncSum) {
   });
 }
 
-// Read the profile from ProfileFileName and assign the value to the
-// instrumented BB and the edges. This function also updates ProgramMaxCount.
-// Return true if the profile are successfully read, and false on errors.
-bool PGOUseFunc::readCounters(IndexedInstrProfReader *PGOReader, bool &AllZeros,
-                              InstrProfRecord::CountPseudoKind &PseudoKind) {
-  auto &Ctx = M->getContext();
+bool PGOUseFunc::getRecord(IndexedInstrProfReader *PGOReader) {
   uint64_t MismatchedFuncSum = 0;
-  Expected<InstrProfRecord> Result = PGOReader->getInstrProfRecord(
+  auto Result = PGOReader->getInstrProfRecord(
       FuncInfo.FuncName, FuncInfo.FunctionHash, FuncInfo.DeprecatedFuncName,
       &MismatchedFuncSum);
   if (Error E = Result.takeError()) {
     handleInstrProfError(std::move(E), MismatchedFuncSum);
     return false;
   }
   ProfileRecord = std::move(Result.get());
+  ProgramMaxCount = PGOReader->getMaximumFunctionCount(IsCS);
+  return true;
+}
+
+// Read the profile from ProfileFileName and assign the value to the
+// instrumented BB and the edges. Return true if the profile are successfully
+// read, and false on errors.
+bool PGOUseFunc::readCounters(bool &AllZeros,
+                              InstrProfRecord::CountPseudoKind &PseudoKind) {
+  auto &Ctx = M->getContext();
   PseudoKind = ProfileRecord.getCountPseudoKind();
   if (PseudoKind != InstrProfRecord::NotPseudo) {
     return true;
@@ -1488,22 +1498,13 @@ bool PGOUseFunc::readCounters(IndexedInstrProfReader *PGOReader, bool &AllZeros,
         DS_Warning));
     return false;
   }
-  ProgramMaxCount = PGOReader->getMaximumFunctionCount(IsCS);
   return true;
 }
 
-void PGOUseFunc::populateCoverage(IndexedInstrProfReader *PGOReader) {
-  uint64_t MismatchedFuncSum = 0;
-  Expected<InstrProfRecord> Result = PGOReader->getInstrProfRecord(
-      FuncInfo.FuncName, FuncInfo.FunctionHash, FuncInfo.DeprecatedFuncName,
-      &MismatchedFuncSum);
-  if (auto Err = Result.takeError()) {
-    handleInstrProfError(std::move(Err), MismatchedFuncSum);
-    return;
-  }
+void PGOUseFunc::populateCoverage() {
   IsCS ? NumOfCSPGOFunc++ : NumOfPGOFunc++;
 
-  std::vector<uint64_t> &CountsFromProfile = Result.get().Counts;
+  ArrayRef<uint64_t> CountsFromProfile = ProfileRecord.Counts;
   DenseMap<const BasicBlock *, bool> Coverage;
   unsigned Index = 0;
   for (auto &BB : F)
@@ -2243,8 +2244,10 @@ static bool annotateAllFunctions(
     PGOUseFunc Func(F, &M, TLI, ComdatMembers, BPI, BFI, LI, PSI, IsCS,
                     InstrumentFuncEntry, InstrumentLoopEntries,
                     HasSingleByteCoverage);
+    if (!Func.getRecord(PGOReader.get()))
+      continue;
     if (HasSingleByteCoverage) {
-      Func.populateCoverage(PGOReader.get());
+      Func.populateCoverage();
       continue;
     }
     // When PseudoKind is set to a vaule other than InstrProfRecord::NotPseudo,
@@ -2253,7 +2256,7 @@ static bool annotateAllFunctions(
     // attribute and drop all the profile counters.
     InstrProfRecord::CountPseudoKind PseudoKind = InstrProfRecord::NotPseudo;
     bool AllZeros = false;
-    if (!Func.readCounters(PGOReader.get(), AllZeros, PseudoKind))
+    if (!Func.readCounters(AllZeros, PseudoKind))
       continue;
     if (AllZeros) {
       F.setEntryCount(ProfileCount(0, Function::PCT_Real));

llvm/unittests/ProfileData/InstrProfTest.cpp

@@ -135,7 +135,7 @@ TEST_P(MaybeSparseInstrProfTest, get_instr_prof_record) {
   auto Profile = Writer.writeBuffer();
   readProfile(std::move(Profile));
 
-  Expected<InstrProfRecord> R = Reader->getInstrProfRecord("foo", 0x1234);
+  auto R = Reader->getInstrProfRecord("foo", 0x1234);
   EXPECT_THAT_ERROR(R.takeError(), Succeeded());
   ASSERT_EQ(2U, R->Counts.size());
   ASSERT_EQ(1U, R->Counts[0]);
@@ -251,7 +251,7 @@ TEST_F(InstrProfTest, test_writer_merge) {
   auto Profile = Writer.writeBuffer();
   readProfile(std::move(Profile));
 
-  Expected<InstrProfRecord> R = Reader->getInstrProfRecord("func1", 0x1234);
+  auto R = Reader->getInstrProfRecord("func1", 0x1234);
   EXPECT_THAT_ERROR(R.takeError(), Succeeded());
   ASSERT_EQ(1U, R->Counts.size());
   ASSERT_EQ(42U, R->Counts[0]);
@@ -600,7 +600,7 @@ TEST_F(InstrProfTest, test_memprof_merge) {
   auto Profile = Writer.writeBuffer();
   readProfile(std::move(Profile));
 
-  Expected<InstrProfRecord> R = Reader->getInstrProfRecord("func1", 0x1234);
+  auto R = Reader->getInstrProfRecord("func1", 0x1234);
   EXPECT_THAT_ERROR(R.takeError(), Succeeded());
   ASSERT_EQ(1U, R->Counts.size());
   ASSERT_EQ(42U, R->Counts[0]);
@@ -800,7 +800,7 @@ TEST_P(InstrProfReaderWriterTest, icall_and_vtable_data_read_write) {
   // Set reader value prof data endianness.
   Reader->setValueProfDataEndianness(getEndianness());
 
-  Expected<InstrProfRecord> R = Reader->getInstrProfRecord("caller", 0x1234);
+  auto R = Reader->getInstrProfRecord("caller", 0x1234);
   ASSERT_THAT_ERROR(R.takeError(), Succeeded());
 
   // Test the number of instrumented indirect call sites and the number of
@@ -874,7 +874,7 @@ TEST_P(MaybeSparseInstrProfTest, annotate_vp_data) {
   Writer.addRecord(std::move(Record), Err);
   auto Profile = Writer.writeBuffer();
   readProfile(std::move(Profile));
-  Expected<InstrProfRecord> R = Reader->getInstrProfRecord("caller", 0x1234);
+  auto R = Reader->getInstrProfRecord("caller", 0x1234);
   EXPECT_THAT_ERROR(R.takeError(), Succeeded());
 
   LLVMContext Ctx;
@@ -1051,7 +1051,7 @@ TEST_P(MaybeSparseInstrProfTest, icall_and_vtable_data_merge) {
 
   // Test the number of instrumented value sites and the number of profiled
   // values for each site.
-  Expected<InstrProfRecord> R = Reader->getInstrProfRecord("caller", 0x1234);
+  auto R = Reader->getInstrProfRecord("caller", 0x1234);
   EXPECT_THAT_ERROR(R.takeError(), Succeeded());
   // For indirect calls.
   ASSERT_EQ(5U, R->getNumValueSites(IPVK_IndirectCallTarget));
@@ -1190,13 +1190,11 @@ TEST_P(ValueProfileMergeEdgeCaseTest, value_profile_data_merge_saturation) {
   readProfile(std::move(Profile));
 
   // Verify saturation of counts.
-  Expected<InstrProfRecord> ReadRecord1 =
-      Reader->getInstrProfRecord("foo", 0x1234);
+  auto ReadRecord1 = Reader->getInstrProfRecord("foo", 0x1234);
   ASSERT_THAT_ERROR(ReadRecord1.takeError(), Succeeded());
   EXPECT_EQ(MaxEdgeCount, ReadRecord1->Counts[0]);
 
-  Expected<InstrProfRecord> ReadRecord2 =
-      Reader->getInstrProfRecord("baz", 0x5678);
+  auto ReadRecord2 = Reader->getInstrProfRecord("baz", 0x5678);
   ASSERT_TRUE(bool(ReadRecord2));
   ASSERT_EQ(1U, ReadRecord2->getNumValueSites(ValueKind));
   auto VD = ReadRecord2->getValueArrayForSite(ValueKind, 0);
@@ -1241,7 +1239,7 @@ TEST_P(ValueProfileMergeEdgeCaseTest, value_profile_data_merge_site_trunc) {
   auto Profile = Writer.writeBuffer();
   readProfile(std::move(Profile));
 
-  Expected<InstrProfRecord> R = Reader->getInstrProfRecord("caller", 0x1234);
+  auto R = Reader->getInstrProfRecord("caller", 0x1234);
   ASSERT_THAT_ERROR(R.takeError(), Succeeded());
   ASSERT_EQ(2U, R->getNumValueSites(ValueKind));
   auto VD = R->getValueArrayForSite(ValueKind, 0);