profile format change

mingmingl-llvm · mingmingl-llvm · commit 2664d700b0f2 · 2025-05-26T15:38:00.000-07:00
diff --git a/llvm/include/llvm/ProfileData/SampleProf.h b/llvm/include/llvm/ProfileData/SampleProf.h
@@ -201,6 +201,9 @@ enum class SecProfSummaryFlags : uint32_t {
   /// SecFlagIsPreInlined means this profile contains ShouldBeInlined
   /// contexts thus this is CS preinliner computed.
   SecFlagIsPreInlined = (1 << 4),
+
+  /// SecFlagHasVTableTypeProf means this profile contains vtable type profiles.
+  SecFlagHasVTableTypeProf = (1 << 5),
 };
 
 enum class SecFuncMetadataFlags : uint32_t {
@@ -312,16 +315,19 @@ struct LineLocationHash {
 
 raw_ostream &operator<<(raw_ostream &OS, const LineLocation &Loc);
 
+using TypeMap = std::map<FunctionId, uint64_t>;
+
 /// Representation of a single sample record.
 ///
 /// A sample record is represented by a positive integer value, which
 /// indicates how frequently was the associated line location executed.
 ///
 /// Additionally, if the associated location contains a function call,
-/// the record will hold a list of all the possible called targets. For
-/// direct calls, this will be the exact function being invoked. For
-/// indirect calls (function pointers, virtual table dispatch), this
-/// will be a list of one or more functions.
+/// the record will hold a list of all the possible called targets and the types
+/// for virtual table dispatches. For direct calls, this will be the exact
+/// function being invoked. For indirect calls (function pointers, virtual table
+/// dispatch), this will be a list of one or more functions. For virtual table
+/// dispatches, this record will also hold the type of the object.
 class SampleRecord {
 public:
   using CallTarget = std::pair<FunctionId, uint64_t>;
@@ -336,6 +342,7 @@ class SampleRecord {
 
   using SortedCallTargetSet = std::set<CallTarget, CallTargetComparator>;
   using CallTargetMap = std::unordered_map<FunctionId, uint64_t>;
+
   SampleRecord() = default;
 
   /// Increment the number of samples for this record by \p S.
@@ -374,6 +381,14 @@ class SampleRecord {
                       : sampleprof_error::success;
   }
 
+  sampleprof_error addTypeCount(FunctionId F, uint64_t S, uint64_t Weight = 1) {
+    uint64_t &Samples = TypeCounts[F];
+    bool Overflowed;
+    Samples = SaturatingMultiplyAdd(S, Weight, Samples, &Overflowed);
+    return Overflowed ? sampleprof_error::counter_overflow
+                      : sampleprof_error::success;
+  }
+
   /// Remove called function from the call target map. Return the target sample
   /// count of the called function.
   uint64_t removeCalledTarget(FunctionId F) {
@@ -391,6 +406,7 @@ class SampleRecord {
 
   uint64_t getSamples() const { return NumSamples; }
   const CallTargetMap &getCallTargets() const { return CallTargets; }
+  const TypeMap &getTypes() const { return TypeCounts; }
   const SortedCallTargetSet getSortedCallTargets() const {
     return sortCallTargets(CallTargets);
   }
@@ -439,6 +455,7 @@ class SampleRecord {
 private:
   uint64_t NumSamples = 0;
   CallTargetMap CallTargets;
+  TypeMap TypeCounts;
 };
 
 raw_ostream &operator<<(raw_ostream &OS, const SampleRecord &Sample);
@@ -734,6 +751,7 @@ using BodySampleMap = std::map<LineLocation, SampleRecord>;
 // memory, which is *very* significant for large profiles.
 using FunctionSamplesMap = std::map<FunctionId, FunctionSamples>;
 using CallsiteSampleMap = std::map<LineLocation, FunctionSamplesMap>;
+using CallsiteTypeMap = std::map<LineLocation, TypeMap>;
 using LocToLocMap =
     std::unordered_map<LineLocation, LineLocation, LineLocationHash>;
 
@@ -791,6 +809,11 @@ class FunctionSamples {
         Func, Num, Weight);
   }
 
+  sampleprof_error addTypeSamples(const LineLocation &Loc, FunctionId Func,
+                                  uint64_t Num, uint64_t Weight = 1) {
+    return BodySamples[Loc].addTypeCount(Func, Num, Weight);
+  }
+
   sampleprof_error addSampleRecord(LineLocation Location,
                                    const SampleRecord &SampleRecord,
                                    uint64_t Weight = 1) {
@@ -916,6 +939,13 @@ class FunctionSamples {
     return &Iter->second;
   }
 
+  const TypeMap *findTypeSamplesAt(const LineLocation &Loc) const {
+    auto Iter = VirtualCallsiteTypes.find(mapIRLocToProfileLoc(Loc));
+    if (Iter == VirtualCallsiteTypes.end())
+      return nullptr;
+    return &Iter->second;
+  }
+
   /// Returns a pointer to FunctionSamples at the given callsite location
   /// \p Loc with callee \p CalleeName. If no callsite can be found, relax
   /// the restriction to return the FunctionSamples at callsite location
@@ -977,6 +1007,14 @@ class FunctionSamples {
     return CallsiteSamples;
   }
 
+  const CallsiteTypeMap &getCallsiteTypes() const {
+    return VirtualCallsiteTypes;
+  }
+
+  TypeMap& getTypeSamplesAt(const LineLocation &Loc) {
+    return VirtualCallsiteTypes[mapIRLocToProfileLoc(Loc)];
+  }
+
   /// Return the maximum of sample counts in a function body. When SkipCallSite
   /// is false, which is the default, the return count includes samples in the
   /// inlined functions. When SkipCallSite is true, the return count only
@@ -1274,6 +1312,8 @@ class FunctionSamples {
   /// collected in the call to baz() at line offset 8.
   CallsiteSampleMap CallsiteSamples;
 
+  CallsiteTypeMap VirtualCallsiteTypes;
+
   /// IR to profile location map generated by stale profile matching.
   ///
   /// Each entry is a mapping from the location on current build to the matched
diff --git a/llvm/include/llvm/ProfileData/SampleProfReader.h b/llvm/include/llvm/ProfileData/SampleProfReader.h
@@ -701,6 +701,17 @@ class SampleProfileReaderBinary : public SampleProfileReader {
   /// otherwise same as readStringFromTable, also return its hash value.
   ErrorOr<std::pair<SampleContext, uint64_t>> readSampleContextFromTable();
 
+  /// Overridden by SampleProfileReaderExtBinary to read the vtable profile.
+  virtual std::error_code readVTableProf(const LineLocation &Loc,
+                                         FunctionSamples &FProfile) {
+    return sampleprof_error::success;
+  }
+
+  virtual std::error_code readCallsiteVTableProf(FunctionSamples &FProfile) {
+    return sampleprof_error::success;
+  }
+
+
   /// Points to the current location in the buffer.
   const uint8_t *Data = nullptr;
 
@@ -814,6 +825,8 @@ class SampleProfileReaderExtBinaryBase : public SampleProfileReaderBinary {
   /// The set containing the functions to use when compiling a module.
   DenseSet<StringRef> FuncsToUse;
 
+  bool ReadVTableProf = false;
+
 public:
   SampleProfileReaderExtBinaryBase(std::unique_ptr<MemoryBuffer> B,
                                    LLVMContext &C, SampleProfileFormat Format)
@@ -854,6 +867,12 @@ class SampleProfileReaderExtBinary : public SampleProfileReaderExtBinaryBase {
     return sampleprof_error::success;
   };
 
+  std::error_code readVTableProf(const LineLocation &Loc,
+                                 FunctionSamples &FProfile) override;
+
+  std::error_code readCallsiteVTableProf(FunctionSamples &FProfile) override;
+
+   std::error_code readTypeMap(TypeMap& M);
 public:
   SampleProfileReaderExtBinary(std::unique_ptr<MemoryBuffer> B, LLVMContext &C,
                                SampleProfileFormat Format = SPF_Ext_Binary)
diff --git a/llvm/include/llvm/ProfileData/SampleProfWriter.h b/llvm/include/llvm/ProfileData/SampleProfWriter.h
@@ -209,6 +209,17 @@ class SampleProfileWriterBinary : public SampleProfileWriter {
   virtual MapVector<FunctionId, uint32_t> &getNameTable() { return NameTable; }
   virtual std::error_code writeMagicIdent(SampleProfileFormat Format);
   virtual std::error_code writeNameTable();
+  virtual std::error_code
+  writeSampleRecordVTableProf(const SampleRecord &Record, raw_ostream &OS) {
+    // TODO: This is not virtual because SampleProfWriter may create objects of
+    // type SampleProfileWriterRawBinary.
+    return sampleprof_error::success;
+  }
+  virtual std::error_code
+  writeCallsiteType(const FunctionSamples &FunctionSample, raw_ostream &OS) {
+    return sampleprof_error::success;
+  }
+  virtual void addTypeNames(const TypeMap &M) {}
   std::error_code writeHeader(const SampleProfileMap &ProfileMap) override;
   std::error_code writeSummary();
   virtual std::error_code writeContextIdx(const SampleContext &Context);
@@ -218,8 +229,9 @@ class SampleProfileWriterBinary : public SampleProfileWriter {
                                 std::set<FunctionId> &V);
   
   MapVector<FunctionId, uint32_t> NameTable;
-  
+
   void addName(FunctionId FName);
+  void addTypeName(FunctionId TypeName);
   virtual void addContext(const SampleContext &Context);
   void addNames(const FunctionSamples &S);
 
@@ -409,8 +421,23 @@ class SampleProfileWriterExtBinaryBase : public SampleProfileWriterBinary {
 
 class SampleProfileWriterExtBinary : public SampleProfileWriterExtBinaryBase {
 public:
-  SampleProfileWriterExtBinary(std::unique_ptr<raw_ostream> &OS)
-      : SampleProfileWriterExtBinaryBase(OS) {}
+  SampleProfileWriterExtBinary(std::unique_ptr<raw_ostream> &OS);
+
+protected:
+  std::error_code writeSampleRecordVTableProf(const SampleRecord &Record,
+                                              raw_ostream &OS) override;
+
+  std::error_code writeCallsiteType(const FunctionSamples &FunctionSample,
+                                    raw_ostream &OS) override;
+
+  void addTypeNames(const TypeMap &M) override {
+    if (WriteVTableProf)
+      return;
+    // Add type name to TypeNameTable.
+    for (const auto &[Type, Cnt] : M) {
+      addName(Type);
+    }
+  }
 
 private:
   std::error_code writeDefaultLayout(const SampleProfileMap &ProfileMap);
@@ -426,6 +453,11 @@ class SampleProfileWriterExtBinary : public SampleProfileWriterExtBinaryBase {
     assert((SL == DefaultLayout || SL == CtxSplitLayout) &&
            "Unsupported layout");
   }
+
+  std::error_code writeTypeMap(const TypeMap &Map, raw_ostream& OS);
+
+  // TODO:This should be configurable by flag.
+  bool WriteVTableProf = false;
 };
 
 } // end namespace sampleprof
diff --git a/llvm/lib/ProfileData/SampleProfReader.cpp b/llvm/lib/ProfileData/SampleProfReader.cpp
@@ -594,6 +594,67 @@ SampleProfileReaderBinary::readSampleContextFromTable() {
   return std::make_pair(Context, Hash);
 }
 
+std::error_code SampleProfileReaderExtBinary::readTypeMap(TypeMap &M) {
+  auto NumVTableTypes = readNumber<uint32_t>();
+  if (std::error_code EC = NumVTableTypes.getError())
+    return EC;
+
+  for (uint32_t I = 0; I < *NumVTableTypes; ++I) {
+    auto VTableType(readStringFromTable());
+    if (std::error_code EC = VTableType.getError())
+      return EC;
+
+    auto VTableSamples = readNumber<uint64_t>();
+    if (std::error_code EC = VTableSamples.getError())
+      return EC;
+
+    M.insert(std::make_pair(*VTableType, *VTableSamples));
+  }
+  return sampleprof_error::success;
+}
+
+std::error_code
+SampleProfileReaderExtBinary::readVTableProf(const LineLocation &Loc,
+                                             FunctionSamples &FProfile) {
+  if (!ReadVTableProf)
+    return sampleprof_error::success;
+
+  return readTypeMap(FProfile.getTypeSamplesAt(Loc));
+}
+
+std::error_code SampleProfileReaderExtBinary::readCallsiteVTableProf(
+    FunctionSamples &FProfile) {
+  if (!ReadVTableProf)
+    return sampleprof_error::success;
+
+  // Read the vtable type profile for the callsite.
+  auto NumCallsites = readNumber<uint32_t>();
+  if (std::error_code EC = NumCallsites.getError())
+    return EC;
+
+  for (uint32_t I = 0; I < *NumCallsites; ++I) {
+    auto LineOffset = readNumber<uint64_t>();
+    if (std::error_code EC = LineOffset.getError())
+      return EC;
+
+    auto Discriminator = readNumber<uint64_t>();
+    if (std::error_code EC = Discriminator.getError())
+      return EC;
+
+    // Here we handle FS discriminators:
+    uint32_t DiscriminatorVal = (*Discriminator) & getDiscriminatorMask();
+
+    if (!isOffsetLegal(*LineOffset)) {
+      return std::error_code();
+    }
+
+    if (std::error_code EC = readVTableProf(
+            LineLocation(*LineOffset, DiscriminatorVal), FProfile))
+      return EC;
+  }
+  return sampleprof_error::success;
+}
+
 std::error_code
 SampleProfileReaderBinary::readProfile(FunctionSamples &FProfile) {
   auto NumSamples = readNumber<uint64_t>();
@@ -643,6 +704,11 @@ SampleProfileReaderBinary::readProfile(FunctionSamples &FProfile) {
                                       *CalledFunction, *CalledFunctionSamples);
     }
 
+    // read vtable type profiles.
+    if (std::error_code EC = readVTableProf(
+            LineLocation(*LineOffset, DiscriminatorVal), FProfile))
+      return EC;
+
     FProfile.addBodySamples(*LineOffset, DiscriminatorVal, *NumSamples);
   }
 
@@ -674,7 +740,7 @@ SampleProfileReaderBinary::readProfile(FunctionSamples &FProfile) {
       return EC;
   }
 
-  return sampleprof_error::success;
+  return readCallsiteVTableProf(FProfile);
 }
 
 std::error_code
@@ -736,6 +802,8 @@ std::error_code SampleProfileReaderExtBinaryBase::readOneSection(
       FunctionSamples::ProfileIsPreInlined = ProfileIsPreInlined = true;
     if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagFSDiscriminator))
       FunctionSamples::ProfileIsFS = ProfileIsFS = true;
+    if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagHasVTableTypeProf))
+      ReadVTableProf = true;
     break;
   case SecNameTable: {
     bool FixedLengthMD5 =
diff --git a/llvm/lib/ProfileData/SampleProfWriter.cpp b/llvm/lib/ProfileData/SampleProfWriter.cpp
