Implements PGOBBAddrMap in Object and ObjectYAML with enough support for tests.

Removes public and adds check for number of bits.

Author: red1bluelost

llvm/include/llvm/BinaryFormat/ELF.h

@@ -1037,6 +1037,7 @@ 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

@@ -416,6 +416,12 @@ class ELFFile {
   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;
+
   /// Returns a map from every section matching \p IsMatch to its relocation
   /// section, or \p nullptr if it has no relocation section. This function
   /// returns an error if any of the \p IsMatch calls fail or if it fails to

llvm/include/llvm/Object/ELFObjectFile.h

@@ -113,6 +113,9 @@ class ELFObjectFileBase : public ObjectFile {
   // corresponding to the section with that index.
   Expected<std::vector<BBAddrMap>>
   readBBAddrMap(std::optional<unsigned> TextSectionIndex = std::nullopt) const;
+
+  Expected<std::vector<PGOBBAddrMap>> readPGOBBAddrMap(
+      std::optional<unsigned> TextSectionIndex = std::nullopt) const;
 };
 
 class ELFSectionRef : public SectionRef {

llvm/include/llvm/Object/ELFTypes.h

@@ -10,9 +10,12 @@
 #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"
+#include "llvm/Support/BlockFrequency.h"
+#include "llvm/Support/BranchProbability.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MathExtras.h"
@@ -805,6 +808,11 @@ struct BBAddrMap {
       bool HasIndirectBranch : 1; // If this block ends with an indirect branch
                                   // (branch via a register).
 
+      // Number of bits used when encoding Metadata, that way an extension
+      // can pack into the extra space if possible. This must be updated when
+      // new bits are added here.
+      static constexpr uint32_t NumberOfBits = 5;
+
       bool operator==(const Metadata &Other) const {
         return HasReturn == Other.HasReturn &&
                HasTailCall == Other.HasTailCall && IsEHPad == Other.IsEHPad &&
@@ -875,6 +883,106 @@ struct BBAddrMap {
   std::vector<BBEntry> BBEntries; // Basic block entries for this function.
 };
 
+/// An extension of BBAddrMap that holds information relevant to PGO.
+struct PGOBBAddrMap {
+  /// 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 {
+    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
+      /// probability.
+      None = 0,
+      /// Single successor blocks are not present in the successor entries.
+      One = 1,
+      /// Common case for conditional branches to avoid encoding size.
+      Two = 2,
+      /// Uncommon case which needs successor size to be encoded.
+      Multiple = 3,
+    };
+
+    /// Single successor of a given basic block that contains the tag and branch
+    /// probability associated with it.
+    struct SuccessorEntry {
+      /// Unique ID of this successor basic block.
+      uint32_t ID;
+      /// Branch Probability of the edge to this successor taken from MBPI
+      BranchProbability Prob;
+
+      bool operator==(const SuccessorEntry &Other) const {
+        return std::tie(ID, Prob) == std::tie(Other.ID, Other.Prob);
+      }
+    };
+
+    /// 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
+    llvm::SmallVector<SuccessorEntry, 2> Successors;
+
+    /// Converts number of successors into a SuccessorsType.
+    static SuccessorsType getSuccessorsType(unsigned SuccessorsCount) {
+      return SuccessorsCount == 0   ? SuccessorsType::None
+             : SuccessorsCount == 1 ? SuccessorsType::One
+             : SuccessorsCount == 2 ? SuccessorsType::Two
+                                    : SuccessorsType::Multiple;
+    }
+
+    /// Encodes extra information in the free bits of the base metadata
+    static uint32_t encodeMD(BaseMetadata MD, SuccessorsType SuccType) {
+      return MD.encode() | static_cast<uint32_t>(SuccType)
+                               << BaseMetadata::NumberOfBits;
+    }
+
+    /// Extracts successors type then defers all errors to the base metadata
+    static Expected<std::pair<BaseMetadata, SuccessorsType>>
+    decodeMD(uint32_t V) {
+      auto SuccType = SuccessorsType((V >> BaseMetadata::NumberOfBits) & 0b11);
+      V &= ~(0b11 << BaseMetadata::NumberOfBits); // Clear extra bits
+      BaseMetadata MD;
+      if (llvm::Error E = BaseMetadata::decode(V).moveInto(MD))
+        return std::move(E);
+      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);
+    }
+  };
+  // 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
+
+  // 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,
+                    BBFreqEnabled, BBSuccProbEnabled) ==
+           std::tie(Other.Addr, Other.FuncEntryCount, Other.BBEntries,
+                    Other.FuncEntryCountEnabled, Other.BBFreqEnabled,
+                    Other.BBSuccProbEnabled);
+  }
+};
+
 } // end namespace object.
 } // end namespace llvm.
 

llvm/include/llvm/ObjectYAML/ELFYAML.h

@@ -156,17 +156,36 @@ 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;
   };
-  uint8_t Version;
-  llvm::yaml::Hex8 Feature;
-  llvm::yaml::Hex64 Address;
-  std::optional<uint64_t> NumBlocks;
+  BBAddrMapCommonBase Common;
+  std::optional<std::vector<BBEntry>> BBEntries;
+};
+
+struct PGOBBAddrMapEntry {
+  struct BBEntry {
+    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;
 };
 
@@ -204,6 +223,7 @@ struct Chunk {
     DependentLibraries,
     CallGraphProfile,
     BBAddrMap,
+    PGOBBAddrMap,
 
     // Special chunks.
     SpecialChunksStart,
@@ -329,6 +349,20 @@ 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;
 
@@ -737,6 +771,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::PGOBBAddrMapEntry::BBEntry::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)
@@ -905,6 +943,20 @@ 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::PGOBBAddrMapEntry::BBEntry> {
+  static void mapping(IO &IO, ELFYAML::PGOBBAddrMapEntry::BBEntry &Rel);
+};
+
+template <>
+struct MappingTraits<ELFYAML::PGOBBAddrMapEntry::BBEntry::SuccessorEntry> {
+  static void mapping(IO &IO,
+                      ELFYAML::PGOBBAddrMapEntry::BBEntry::SuccessorEntry &Rel);
+};
+
 template <> struct MappingTraits<ELFYAML::GnuHashHeader> {
   static void mapping(IO &IO, ELFYAML::GnuHashHeader &Rel);
 };

llvm/lib/MC/MCParser/ELFAsmParser.cpp

@@ -673,6 +673,8 @@ 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,6 +170,8 @@ 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)