Reworks the design to use same ELF section and not superset data structure.

Author: red1bluelost

llvm/include/llvm/BinaryFormat/ELF.h

@@ -1037,7 +1037,6 @@ enum : unsigned {
   SHT_LLVM_BB_ADDR_MAP = 0x6fff4c0a,        // LLVM Basic Block Address Map.
   SHT_LLVM_OFFLOADING = 0x6fff4c0b,         // LLVM device offloading data.
   SHT_LLVM_LTO = 0x6fff4c0c,                // .llvm.lto for fat LTO.
-  SHT_LLVM_PGO_BB_ADDR_MAP = 0x6fff4c0d,    // LLVM PGO extended BB Addr Map.
   // Android's experimental support for SHT_RELR sections.
   // https://android.googlesource.com/platform/bionic/+/b7feec74547f84559a1467aca02708ff61346d2a/libc/include/elf.h#512
   SHT_ANDROID_RELR = 0x6fffff00,   // Relocation entries; only offsets.

llvm/include/llvm/Object/ELF.h

@@ -413,14 +413,12 @@ class ELFFile {
   /// within the text section that the SHT_LLVM_BB_ADDR_MAP section \p Sec
   /// is associated with. If the current ELFFile is relocatable, a corresponding
   /// \p RelaSec must be passed in as an argument.
+  /// Optional out variable to all collect PGO Analyses. New elements are only
+  /// added if no error occurs. If not provided, the PGO Analyses are decoded
+  /// then ignored.
   Expected<std::vector<BBAddrMap>>
-  decodeBBAddrMap(const Elf_Shdr &Sec, const Elf_Shdr *RelaSec = nullptr) const;
-
-  /// Decodes same as decodeBBAddrMap but also decodes extra information from
-  /// features which are enabled.
-  Expected<std::vector<PGOBBAddrMap>>
-  decodePGOBBAddrMap(const Elf_Shdr &Sec,
-                     const Elf_Shdr *RelaSec = nullptr) const;
+  decodeBBAddrMap(const Elf_Shdr &Sec, const Elf_Shdr *RelaSec = nullptr,
+                  std::vector<PGOAnalysisMap> *PGOAnalyses = nullptr) const;
 
   /// Returns a map from every section matching \p IsMatch to its relocation
   /// section, or \p nullptr if it has no relocation section. This function

llvm/include/llvm/Object/ELFObjectFile.h

@@ -110,12 +110,13 @@ class ELFObjectFileBase : public ObjectFile {
 
   /// Returns a vector of all BB address maps in the object file. When
   // `TextSectionIndex` is specified, only returns the BB address maps
-  // corresponding to the section with that index.
+  // corresponding to the section with that index. When `PGOAnalyses`is
+  // specified, the vector is cleared then filled with extra PGO data if the
+  // feature when enabled in the ELF section. `PGOAnalyses` will always be the
+  // same length as the return value on success, otherwise it is empty.
   Expected<std::vector<BBAddrMap>>
-  readBBAddrMap(std::optional<unsigned> TextSectionIndex = std::nullopt) const;
-
-  Expected<std::vector<PGOBBAddrMap>> readPGOBBAddrMap(
-      std::optional<unsigned> TextSectionIndex = std::nullopt) const;
+  readBBAddrMap(std::optional<unsigned> TextSectionIndex = std::nullopt,
+                std::vector<PGOAnalysisMap> *PGOAnalyses = nullptr) const;
 };
 
 class ELFSectionRef : public SectionRef {

llvm/include/llvm/Object/ELFTypes.h

@@ -10,7 +10,6 @@
 #define LLVM_OBJECT_ELFTYPES_H
 
 #include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/BitmaskEnum.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/Object/Error.h"
@@ -883,26 +882,24 @@ struct BBAddrMap {
   std::vector<BBEntry> BBEntries; // Basic block entries for this function.
 };
 
-/// An extension of BBAddrMap that holds information relevant to PGO.
-struct PGOBBAddrMap {
+/// A feature extension of BBAddrMap that holds information relevant to PGO.
+struct PGOAnalysisMap {
   /// Bitmask of optional features to include in the PGO extended map.
   enum class Features {
-    None = 0,
     FuncEntryCnt = (1 << 0),
     BBFreq = (1 << 1),
     BrProb = (1 << 2),
-    LLVM_MARK_AS_BITMASK_ENUM(/*LargestValue=*/BrProb),
   };
 
   /// Super-set of BBAddrMap::BBEntry with additional fields for block frequency
   /// and branch probability.
-  struct BBEntry {
+  struct PGOBBEntry {
     using BaseMetadata = BBAddrMap::BBEntry::Metadata;
 
     /// Enum indicating the how many successors a block has. This enum must fit
     /// into two bits.
     enum class SuccessorsType {
-      /// None should be present if PGOBBAddrMap has disabled branch
+      /// None should be present if BBAddrMap.feature has disabled branch
       /// probability.
       None = 0,
       /// Single successor blocks are not present in the successor entries.
@@ -926,8 +923,6 @@ struct PGOBBAddrMap {
       }
     };
 
-    /// Reuse of the fields provided by regular BBAddrMap
-    BBAddrMap::BBEntry Base;
     /// Block frequency taken from MBFI
     BlockFrequency BlockFreq;
     /// List of successors of the current block
@@ -958,26 +953,25 @@ struct PGOBBAddrMap {
       return std::make_pair(MD, SuccType);
     }
 
-    bool operator==(const BBEntry &Other) const {
-      return std::tie(Base, BlockFreq, Successors) ==
-             std::tie(Other.Base, Other.BlockFreq, Other.Successors);
+    bool operator==(const PGOBBEntry &Other) const {
+      return std::tie(BlockFreq, Successors) ==
+             std::tie(Other.BlockFreq, Other.Successors);
     }
   };
   // This field is duplicated from BBAddrMap since this class needs a different
   // type for the vector of entries.
-  uint64_t Addr;                  // Function address
-  std::vector<BBEntry> BBEntries; // Extended basic block entries
-  uint64_t FuncEntryCount;        // Prof count from IR function
+  uint64_t FuncEntryCount;           // Prof count from IR function
+  std::vector<PGOBBEntry> BBEntries; // Extended basic block entries
 
   // Flags to indicate if each PGO related info was enabled in this function
   bool FuncEntryCountEnabled : 1;
   bool BBFreqEnabled : 1;
   bool BBSuccProbEnabled : 1;
 
-  bool operator==(const PGOBBAddrMap &Other) const {
-    return std::tie(Addr, FuncEntryCount, BBEntries, FuncEntryCountEnabled,
+  bool operator==(const PGOAnalysisMap &Other) const {
+    return std::tie(FuncEntryCount, BBEntries, FuncEntryCountEnabled,
                     BBFreqEnabled, BBSuccProbEnabled) ==
-           std::tie(Other.Addr, Other.FuncEntryCount, Other.BBEntries,
+           std::tie(Other.FuncEntryCount, Other.BBEntries,
                     Other.FuncEntryCountEnabled, Other.BBFreqEnabled,
                     Other.BBSuccProbEnabled);
   }

llvm/include/llvm/ObjectYAML/ELFYAML.h

@@ -156,37 +156,31 @@ struct DynamicEntry {
   llvm::yaml::Hex64 Val;
 };
 
-struct BBAddrMapCommonBase {
-  uint8_t Version;
-  llvm::yaml::Hex8 Feature;
-  llvm::yaml::Hex64 Address;
-  std::optional<uint64_t> NumBlocks;
-};
-
 struct BBAddrMapEntry {
   struct BBEntry {
     uint32_t ID;
     llvm::yaml::Hex64 AddressOffset;
     llvm::yaml::Hex64 Size;
     llvm::yaml::Hex64 Metadata;
   };
-  BBAddrMapCommonBase Common;
+  uint8_t Version;
+  llvm::yaml::Hex8 Feature;
+  llvm::yaml::Hex64 Address;
+  std::optional<uint64_t> NumBlocks;
   std::optional<std::vector<BBEntry>> BBEntries;
 };
 
-struct PGOBBAddrMapEntry {
-  struct BBEntry {
+struct PGOAnalysisMapEntry {
+  struct PGOBBEntry {
     struct SuccessorEntry {
       uint32_t ID;
       llvm::yaml::Hex32 BrProb;
     };
-    BBAddrMapEntry::BBEntry Base;
     std::optional<uint64_t> BBFreq;
     std::optional<std::vector<SuccessorEntry>> Successors;
   };
-  BBAddrMapCommonBase Common;
   std::optional<uint64_t> FuncEntryCount;
-  std::optional<std::vector<BBEntry>> BBEntries;
+  std::optional<std::vector<PGOBBEntry>> PGOBBEntries;
 };
 
 struct StackSizeEntry {
@@ -223,7 +217,6 @@ struct Chunk {
     DependentLibraries,
     CallGraphProfile,
     BBAddrMap,
-    PGOBBAddrMap,
 
     // Special chunks.
     SpecialChunksStart,
@@ -337,6 +330,7 @@ struct SectionHeaderTable : Chunk {
 
 struct BBAddrMapSection : Section {
   std::optional<std::vector<BBAddrMapEntry>> Entries;
+  std::optional<std::vector<PGOAnalysisMapEntry>> PGOAnalyses;
 
   BBAddrMapSection() : Section(ChunkKind::BBAddrMap) {}
 
@@ -349,20 +343,6 @@ struct BBAddrMapSection : Section {
   }
 };
 
-struct PGOBBAddrMapSection : Section {
-  std::optional<std::vector<PGOBBAddrMapEntry>> Entries;
-
-  PGOBBAddrMapSection() : Section(ChunkKind::PGOBBAddrMap) {}
-
-  std::vector<std::pair<StringRef, bool>> getEntries() const override {
-    return {{"Entries", Entries.has_value()}};
-  };
-
-  static bool classof(const Chunk *S) {
-    return S->Kind == ChunkKind::PGOBBAddrMap;
-  }
-};
-
 struct StackSizesSection : Section {
   std::optional<std::vector<StackSizeEntry>> Entries;
 
@@ -771,10 +751,10 @@ bool shouldAllocateFileSpace(ArrayRef<ProgramHeader> Phdrs,
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::ELFYAML::StackSizeEntry)
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::ELFYAML::BBAddrMapEntry)
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::ELFYAML::BBAddrMapEntry::BBEntry)
-LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::ELFYAML::PGOBBAddrMapEntry)
-LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::ELFYAML::PGOBBAddrMapEntry::BBEntry)
+LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::ELFYAML::PGOAnalysisMapEntry)
+LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::ELFYAML::PGOAnalysisMapEntry::PGOBBEntry)
 LLVM_YAML_IS_SEQUENCE_VECTOR(
-    llvm::ELFYAML::PGOBBAddrMapEntry::BBEntry::SuccessorEntry)
+    llvm::ELFYAML::PGOAnalysisMapEntry::PGOBBEntry::SuccessorEntry)
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::ELFYAML::DynamicEntry)
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::ELFYAML::LinkerOption)
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::ELFYAML::CallGraphEntryWeight)
@@ -943,18 +923,19 @@ template <> struct MappingTraits<ELFYAML::BBAddrMapEntry::BBEntry> {
   static void mapping(IO &IO, ELFYAML::BBAddrMapEntry::BBEntry &Rel);
 };
 
-template <> struct MappingTraits<ELFYAML::PGOBBAddrMapEntry> {
-  static void mapping(IO &IO, ELFYAML::PGOBBAddrMapEntry &Rel);
+template <> struct MappingTraits<ELFYAML::PGOAnalysisMapEntry> {
+  static void mapping(IO &IO, ELFYAML::PGOAnalysisMapEntry &Rel);
 };
 
-template <> struct MappingTraits<ELFYAML::PGOBBAddrMapEntry::BBEntry> {
-  static void mapping(IO &IO, ELFYAML::PGOBBAddrMapEntry::BBEntry &Rel);
+template <> struct MappingTraits<ELFYAML::PGOAnalysisMapEntry::PGOBBEntry> {
+  static void mapping(IO &IO, ELFYAML::PGOAnalysisMapEntry::PGOBBEntry &Rel);
 };
 
 template <>
-struct MappingTraits<ELFYAML::PGOBBAddrMapEntry::BBEntry::SuccessorEntry> {
-  static void mapping(IO &IO,
-                      ELFYAML::PGOBBAddrMapEntry::BBEntry::SuccessorEntry &Rel);
+struct MappingTraits<ELFYAML::PGOAnalysisMapEntry::PGOBBEntry::SuccessorEntry> {
+  static void
+  mapping(IO &IO,
+          ELFYAML::PGOAnalysisMapEntry::PGOBBEntry::SuccessorEntry &Rel);
 };
 
 template <> struct MappingTraits<ELFYAML::GnuHashHeader> {

llvm/lib/MC/MCParser/ELFAsmParser.cpp

@@ -673,8 +673,6 @@ bool ELFAsmParser::ParseSectionArguments(bool IsPush, SMLoc loc) {
       Type = ELF::SHT_LLVM_SYMPART;
     else if (TypeName == "llvm_bb_addr_map")
       Type = ELF::SHT_LLVM_BB_ADDR_MAP;
-    else if (TypeName == "llvm_pgo_bb_addr_map")
-      Type = ELF::SHT_LLVM_PGO_BB_ADDR_MAP;
     else if (TypeName == "llvm_offloading")
       Type = ELF::SHT_LLVM_OFFLOADING;
     else if (TypeName == "llvm_lto")

llvm/lib/MC/MCSectionELF.cpp

@@ -170,8 +170,6 @@ void MCSectionELF::printSwitchToSection(const MCAsmInfo &MAI, const Triple &T,
     OS << "llvm_bb_addr_map";
   else if (Type == ELF::SHT_LLVM_BB_ADDR_MAP_V0)
     OS << "llvm_bb_addr_map_v0";
-  else if (Type == ELF::SHT_LLVM_PGO_BB_ADDR_MAP)
-    OS << "llvm_pgo_bb_addr_map";
   else if (Type == ELF::SHT_LLVM_OFFLOADING)
     OS << "llvm_offloading";
   else if (Type == ELF::SHT_LLVM_LTO)