[llvm] Implements PGOBBAddrMap in Object and ObjectYAML with tests [1/5] (PR #71750)

Micah Weston via llvm-commits llvm-commits at lists.llvm.org
Wed Nov 22 10:55:16 PST 2023


https://github.com/red1bluelost updated https://github.com/llvm/llvm-project/pull/71750

>From 6374da590be4a1f4a1ab99c5c483da9490fee6b8 Mon Sep 17 00:00:00 2001
From: Micah Weston <micahsweston at gmail.com>
Date: Tue, 24 Oct 2023 22:31:15 -0400
Subject: [PATCH 1/8] Implements PGOBBAddrMap in Object and ObjectYAML with
 enough support for tests.

Removes public and adds check for number of bits.
---
 llvm/include/llvm/BinaryFormat/ELF.h        |   1 +
 llvm/include/llvm/Object/ELF.h              |   6 +
 llvm/include/llvm/Object/ELFObjectFile.h    |   3 +
 llvm/include/llvm/Object/ELFTypes.h         | 108 +++++
 llvm/include/llvm/ObjectYAML/ELFYAML.h      |  60 ++-
 llvm/lib/MC/MCParser/ELFAsmParser.cpp       |   2 +
 llvm/lib/MC/MCSectionELF.cpp                |   2 +
 llvm/lib/Object/ELF.cpp                     | 263 +++++++++---
 llvm/lib/Object/ELFObjectFile.cpp           |  52 ++-
 llvm/lib/ObjectYAML/ELFEmitter.cpp          | 130 +++++-
 llvm/lib/ObjectYAML/ELFYAML.cpp             |  45 +-
 llvm/test/MC/AsmParser/llvm_section_types.s |   4 +
 llvm/tools/obj2yaml/elf2yaml.cpp            |   2 +-
 llvm/unittests/Object/ELFObjectFileTest.cpp | 440 ++++++++++++++++++++
 llvm/unittests/Object/ELFTypesTest.cpp      |  53 ++-
 15 files changed, 1080 insertions(+), 91 deletions(-)

diff --git a/llvm/include/llvm/BinaryFormat/ELF.h b/llvm/include/llvm/BinaryFormat/ELF.h
index 3596174f74dde80..299d7e8265317b3 100644
--- a/llvm/include/llvm/BinaryFormat/ELF.h
+++ b/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.
diff --git a/llvm/include/llvm/Object/ELF.h b/llvm/include/llvm/Object/ELF.h
index 927deeea2cd6aef..2efcdf2cfe714ad 100644
--- a/llvm/include/llvm/Object/ELF.h
+++ b/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
diff --git a/llvm/include/llvm/Object/ELFObjectFile.h b/llvm/include/llvm/Object/ELFObjectFile.h
index 8e16fc148a3c78d..7cf2226038b2fc9 100644
--- a/llvm/include/llvm/Object/ELFObjectFile.h
+++ b/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 {
diff --git a/llvm/include/llvm/Object/ELFTypes.h b/llvm/include/llvm/Object/ELFTypes.h
index 4670abc867de63c..7ab34c8caaf5c33 100644
--- a/llvm/include/llvm/Object/ELFTypes.h
+++ b/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.
 
diff --git a/llvm/include/llvm/ObjectYAML/ELFYAML.h b/llvm/include/llvm/ObjectYAML/ELFYAML.h
index 1ba41232f552e3a..3f18d189b7dc870 100644
--- a/llvm/include/llvm/ObjectYAML/ELFYAML.h
+++ b/llvm/include/llvm/ObjectYAML/ELFYAML.h
@@ -156,6 +156,13 @@ 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;
@@ -163,10 +170,22 @@ struct BBAddrMapEntry {
     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);
 };
diff --git a/llvm/lib/MC/MCParser/ELFAsmParser.cpp b/llvm/lib/MC/MCParser/ELFAsmParser.cpp
index dbfe0d83e1b29b1..e4c1cc7763f7804 100644
--- a/llvm/lib/MC/MCParser/ELFAsmParser.cpp
+++ b/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")
diff --git a/llvm/lib/MC/MCSectionELF.cpp b/llvm/lib/MC/MCSectionELF.cpp
index 95fdf33522076e4..02e63f7b03919a8 100644
--- a/llvm/lib/MC/MCSectionELF.cpp
+++ b/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)
diff --git a/llvm/lib/Object/ELF.cpp b/llvm/lib/Object/ELF.cpp
index 1d73a6ffa73f5f9..e5b2b38048b9382 100644
--- a/llvm/lib/Object/ELF.cpp
+++ b/llvm/lib/Object/ELF.cpp
@@ -312,6 +312,7 @@ StringRef llvm::object::getELFSectionTypeName(uint32_t Machine, unsigned Type) {
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_PART_PHDR);
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_BB_ADDR_MAP_V0);
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_BB_ADDR_MAP);
+    STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_PGO_BB_ADDR_MAP);
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_OFFLOADING);
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_LTO);
     STRINGIFY_ENUM_CASE(ELF, SHT_GNU_ATTRIBUTES);
@@ -645,11 +646,179 @@ ELFFile<ELFT>::toMappedAddr(uint64_t VAddr, WarningHandler WarnHandler) const {
   return base() + Offset;
 }
 
-template <class ELFT>
-Expected<std::vector<BBAddrMap>>
-ELFFile<ELFT>::decodeBBAddrMap(const Elf_Shdr &Sec,
-                               const Elf_Shdr *RelaSec) const {
-  bool IsRelocatable = getHeader().e_type == ELF::ET_REL;
+// Helper to extract and decode the next ULEB128 value as unsigned int.
+// Returns zero and sets ULEBSizeErr if the ULEB128 value exceeds the unsigned
+// int limit.
+// Also returns zero if ULEBSizeErr is already in an error state.
+// ULEBSizeErr is an out variable if an error occurs.
+template <typename IntTy>
+static IntTy readULEB128As(DataExtractor &Data, DataExtractor::Cursor &Cur,
+                           Error &ULEBSizeErr) {
+  static_assert(std::is_unsigned_v<IntTy> &&
+                    (std::numeric_limits<IntTy>::radix == 2),
+                "only use unsigned radix 2");
+  // Bail out and do not extract data if ULEBSizeErr is already set.
+  if (ULEBSizeErr)
+    return 0;
+  uint64_t Offset = Cur.tell();
+  uint64_t Value = Data.getULEB128(Cur);
+  if (Value > std::numeric_limits<IntTy>::max()) {
+    ULEBSizeErr = createError("ULEB128 value at offset 0x" +
+                              Twine::utohexstr(Offset) + " exceeds UINT" +
+                              Twine(std::numeric_limits<IntTy>::digits) +
+                              "_MAX (0x" + Twine::utohexstr(Value) + ")");
+    return 0;
+  }
+  return static_cast<IntTy>(Value);
+}
+
+namespace {
+struct BasicAddrMapBBEntry {
+  uint32_t ID;
+  uint32_t Offset;
+  uint32_t Size;
+  uint32_t MD;
+};
+
+template <typename ELFT, typename AddrMap> struct AddrMapDecodeTrait;
+
+template <typename ELFT> struct AddrMapDecodeTrait<ELFT, BBAddrMap> {
+  // Base addr map has no extra data
+  struct ExtraData {};
+
+  static constexpr unsigned SectionID = ELF::SHT_LLVM_BB_ADDR_MAP;
+
+  DataExtractor &Data;
+  DataExtractor::Cursor &Cur;
+  Error &ULEBSizeErr;
+  Error &MetadataDecodeErr;
+
+  uint32_t PrevBBEndOffset = 0;
+
+  AddrMapDecodeTrait(DataExtractor &Data, DataExtractor::Cursor &Cur,
+                     Error &ULEBSizeErr, Error &MetadataDecodeErr)
+      : Data(Data), Cur(Cur), ULEBSizeErr(ULEBSizeErr),
+        MetadataDecodeErr(MetadataDecodeErr) {}
+
+  ExtraData decodeExtraDataHeader(uint8_t, uint8_t) {
+    PrevBBEndOffset = 0;
+    return {};
+  }
+
+  // This helper method avoids decoding the metadata so other AddrMaps can use
+  BasicAddrMapBBEntry decodeBasicInfo(uint8_t Version, uint8_t Feature,
+                                      uint32_t BlockIndex) {
+    uint32_t ID = Version >= 2 ? readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr)
+                               : BlockIndex;
+    uint32_t Offset = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
+    uint32_t Size = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
+    uint32_t MD = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
+    if (Version >= 1) {
+      // Offset is calculated relative to the end of the previous BB.
+      Offset += PrevBBEndOffset;
+      PrevBBEndOffset = Offset + Size;
+    }
+    return {ID, Offset, Size, MD};
+  }
+
+  BBAddrMap::BBEntry decodeBBEntry(uint8_t Version, uint8_t Feature,
+                                   uint32_t BlockIndex) {
+    auto [ID, Offset, Size, MD] = decodeBasicInfo(Version, Feature, BlockIndex);
+    auto MetadataOrErr = BBAddrMap::BBEntry::Metadata::decode(MD);
+    if (Error E = MetadataOrErr.takeError()) {
+      MetadataDecodeErr = std::move(E);
+      return {{}, {}, {}, {}};
+    }
+    return {ID, Offset, Size, *MetadataOrErr};
+  }
+
+  BBAddrMap makeAddrMap(ExtraData, uint8_t Feature, uint64_t Address,
+                        std::vector<BBAddrMap::BBEntry> BBEntries) {
+    return {Address, std::move(BBEntries)};
+  }
+};
+
+template <typename ELFT> struct AddrMapDecodeTrait<ELFT, PGOBBAddrMap> {
+  using Features = PGOBBAddrMap::Features;
+
+  struct ExtraData {
+    uint64_t FuncEntryCount;
+  };
+
+  static constexpr unsigned SectionID = ELF::SHT_LLVM_PGO_BB_ADDR_MAP;
+
+  AddrMapDecodeTrait<ELFT, BBAddrMap> Base;
+
+  AddrMapDecodeTrait(DataExtractor &Data, DataExtractor::Cursor &Cur,
+                     Error &ULEBSizeErr, Error &MetadataDecodeErr)
+      : Base(Data, Cur, ULEBSizeErr, MetadataDecodeErr) {}
+
+  ExtraData decodeExtraDataHeader(uint8_t Version, uint8_t Feature) {
+    Base.decodeExtraDataHeader(Version, Feature);
+    if (Version < 2) {
+      // hijack the metadata error if version is too low
+      Base.MetadataDecodeErr =
+          createError("unsupported SHT_LLVM_PGO_BB_ADDR_MAP version: " +
+                      Twine(static_cast<int>(Version)));
+      return {};
+    }
+    return {Feature & uint8_t(PGOBBAddrMap::Features::FuncEntryCnt)
+                ? readULEB128As<uint64_t>(Base.Data, Base.Cur, Base.ULEBSizeErr)
+                : 0};
+  }
+
+  PGOBBAddrMap::BBEntry decodeBBEntry(uint8_t Version, uint8_t Feature,
+                                      uint32_t BlockIndex) {
+    auto [ID, Offset, Size, MD] =
+        Base.decodeBasicInfo(Version, Feature, BlockIndex);
+    auto MetadataOrErr = PGOBBAddrMap::BBEntry::decodeMD(MD);
+    if (Error E = MetadataOrErr.takeError()) {
+      Base.MetadataDecodeErr = std::move(E);
+      return {{{}, {}, {}, {}}, {}, {}};
+    }
+    auto [MetaData, SuccsType] = *MetadataOrErr;
+
+    uint64_t BBF =
+        Feature & uint8_t(PGOBBAddrMap::Features::BBFreq)
+            ? readULEB128As<uint64_t>(Base.Data, Base.Cur, Base.ULEBSizeErr)
+            : 0;
+
+    llvm::SmallVector<PGOBBAddrMap::BBEntry::SuccessorEntry, 2> Successors;
+    if (Feature & uint8_t(PGOBBAddrMap::Features::BrProb)) {
+      auto SuccCount =
+          SuccsType == PGOBBAddrMap::BBEntry::SuccessorsType::Multiple
+              ? readULEB128As<uint64_t>(Base.Data, Base.Cur, Base.ULEBSizeErr)
+              : uint64_t(SuccsType);
+      for (uint64_t I = 0; I < SuccCount; ++I) {
+        uint32_t BBID =
+            readULEB128As<uint32_t>(Base.Data, Base.Cur, Base.ULEBSizeErr);
+        uint32_t BrProb =
+            readULEB128As<uint32_t>(Base.Data, Base.Cur, Base.ULEBSizeErr);
+        Successors.push_back({BBID, BranchProbability::getRaw(BrProb)});
+      }
+    }
+    return PGOBBAddrMap::BBEntry{
+        {ID, Offset, Size, MetaData}, BlockFrequency(BBF), Successors};
+  }
+
+  PGOBBAddrMap makeAddrMap(ExtraData ED, uint8_t Feature, uint64_t Address,
+                           std::vector<PGOBBAddrMap::BBEntry> BBEntries) {
+    return {Address,
+            std::move(BBEntries),
+            ED.FuncEntryCount,
+            bool(Feature & uint8_t(Features::FuncEntryCnt)),
+            bool(Feature & uint8_t(Features::BBFreq)),
+            bool(Feature & uint8_t(Features::BrProb))};
+  }
+};
+} // namespace
+
+template <typename AddrMap, typename ELFT>
+static Expected<std::vector<AddrMap>>
+decodeBBAddrMapCommonImpl(const ELFFile<ELFT> &EF,
+                          const typename ELFFile<ELFT>::Elf_Shdr &Sec,
+                          const typename ELFFile<ELFT>::Elf_Shdr *RelaSec) {
+  bool IsRelocatable = EF.getHeader().e_type == ELF::ET_REL;
 
   // This DenseMap maps the offset of each function (the location of the
   // reference to the function in the SHT_LLVM_BB_ADDR_MAP section) to the
@@ -659,57 +828,44 @@ ELFFile<ELFT>::decodeBBAddrMap(const Elf_Shdr &Sec,
     assert(RelaSec &&
            "Can't read a SHT_LLVM_BB_ADDR_MAP section in a relocatable "
            "object file without providing a relocation section.");
-    Expected<Elf_Rela_Range> Relas = this->relas(*RelaSec);
+    Expected<typename ELFFile<ELFT>::Elf_Rela_Range> Relas = EF.relas(*RelaSec);
     if (!Relas)
       return createError("unable to read relocations for section " +
-                         describe(*this, Sec) + ": " +
+                         describe(EF, Sec) + ": " +
                          toString(Relas.takeError()));
-    for (Elf_Rela Rela : *Relas)
+    for (typename ELFFile<ELFT>::Elf_Rela Rela : *Relas)
       FunctionOffsetTranslations[Rela.r_offset] = Rela.r_addend;
   }
-  Expected<ArrayRef<uint8_t>> ContentsOrErr = getSectionContents(Sec);
+  Expected<ArrayRef<uint8_t>> ContentsOrErr = EF.getSectionContents(Sec);
   if (!ContentsOrErr)
     return ContentsOrErr.takeError();
   ArrayRef<uint8_t> Content = *ContentsOrErr;
-  DataExtractor Data(Content, isLE(), ELFT::Is64Bits ? 8 : 4);
-  std::vector<BBAddrMap> FunctionEntries;
+  DataExtractor Data(Content, EF.isLE(), ELFT::Is64Bits ? 8 : 4);
+  std::vector<AddrMap> FunctionEntries;
 
   DataExtractor::Cursor Cur(0);
   Error ULEBSizeErr = Error::success();
   Error MetadataDecodeErr = Error::success();
-  // Helper to extract and decode the next ULEB128 value as uint32_t.
-  // Returns zero and sets ULEBSizeErr if the ULEB128 value exceeds the uint32_t
-  // limit.
-  // Also returns zero if ULEBSizeErr is already in an error state.
-  auto ReadULEB128AsUInt32 = [&Data, &Cur, &ULEBSizeErr]() -> uint32_t {
-    // Bail out and do not extract data if ULEBSizeErr is already set.
-    if (ULEBSizeErr)
-      return 0;
-    uint64_t Offset = Cur.tell();
-    uint64_t Value = Data.getULEB128(Cur);
-    if (Value > UINT32_MAX) {
-      ULEBSizeErr = createError(
-          "ULEB128 value at offset 0x" + Twine::utohexstr(Offset) +
-          " exceeds UINT32_MAX (0x" + Twine::utohexstr(Value) + ")");
-      return 0;
-    }
-    return static_cast<uint32_t>(Value);
-  };
+
+  using DecodeTrait = AddrMapDecodeTrait<ELFT, AddrMap>;
+  DecodeTrait DT(Data, Cur, ULEBSizeErr, MetadataDecodeErr);
 
   uint8_t Version = 0;
+  uint8_t Feature = 0;
   while (!ULEBSizeErr && !MetadataDecodeErr && Cur &&
          Cur.tell() < Content.size()) {
-    if (Sec.sh_type == ELF::SHT_LLVM_BB_ADDR_MAP) {
+    if (Sec.sh_type == DecodeTrait::SectionID) {
       Version = Data.getU8(Cur);
       if (!Cur)
         break;
       if (Version > 2)
         return createError("unsupported SHT_LLVM_BB_ADDR_MAP version: " +
                            Twine(static_cast<int>(Version)));
-      Data.getU8(Cur); // Feature byte
+      Feature = Data.getU8(Cur); // Feature byte
     }
     uint64_t SectionOffset = Cur.tell();
-    uintX_t Address = static_cast<uintX_t>(Data.getAddress(Cur));
+    auto Address =
+        static_cast<typename ELFFile<ELFT>::uintX_t>(Data.getAddress(Cur));
     if (!Cur)
       return Cur.takeError();
     if (IsRelocatable) {
@@ -718,34 +874,23 @@ ELFFile<ELFT>::decodeBBAddrMap(const Elf_Shdr &Sec,
       if (FOTIterator == FunctionOffsetTranslations.end()) {
         return createError("failed to get relocation data for offset: " +
                            Twine::utohexstr(SectionOffset) + " in section " +
-                           describe(*this, Sec));
+                           describe(EF, Sec));
       }
       Address = FOTIterator->second;
     }
-    uint32_t NumBlocks = ReadULEB128AsUInt32();
-    std::vector<BBAddrMap::BBEntry> BBEntries;
-    uint32_t PrevBBEndOffset = 0;
+    uint32_t NumBlocks = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
+    auto ExtraData = DT.decodeExtraDataHeader(Version, Feature);
+    std::vector<typename AddrMap::BBEntry> BBEntries;
     for (uint32_t BlockIndex = 0;
          !MetadataDecodeErr && !ULEBSizeErr && Cur && (BlockIndex < NumBlocks);
          ++BlockIndex) {
-      uint32_t ID = Version >= 2 ? ReadULEB128AsUInt32() : BlockIndex;
-      uint32_t Offset = ReadULEB128AsUInt32();
-      uint32_t Size = ReadULEB128AsUInt32();
-      uint32_t MD = ReadULEB128AsUInt32();
-      if (Version >= 1) {
-        // Offset is calculated relative to the end of the previous BB.
-        Offset += PrevBBEndOffset;
-        PrevBBEndOffset = Offset + Size;
-      }
-      Expected<BBAddrMap::BBEntry::Metadata> MetadataOrErr =
-          BBAddrMap::BBEntry::Metadata::decode(MD);
-      if (!MetadataOrErr) {
-        MetadataDecodeErr = MetadataOrErr.takeError();
+      auto Entry = DT.decodeBBEntry(Version, Feature, BlockIndex);
+      if (MetadataDecodeErr)
         break;
-      }
-      BBEntries.push_back({ID, Offset, Size, *MetadataOrErr});
+      BBEntries.push_back(std::move(Entry));
     }
-    FunctionEntries.emplace_back(Address, std::move(BBEntries));
+    FunctionEntries.push_back(
+        DT.makeAddrMap(ExtraData, Feature, Address, std::move(BBEntries)));
   }
   // Either Cur is in the error state, or we have an error in ULEBSizeErr or
   // MetadataDecodeErr (but not both), but we join all errors here to be safe.
@@ -755,6 +900,20 @@ ELFFile<ELFT>::decodeBBAddrMap(const Elf_Shdr &Sec,
   return FunctionEntries;
 }
 
+template <class ELFT>
+Expected<std::vector<BBAddrMap>>
+ELFFile<ELFT>::decodeBBAddrMap(const Elf_Shdr &Sec,
+                               const Elf_Shdr *RelaSec) const {
+  return decodeBBAddrMapCommonImpl<BBAddrMap>(*this, Sec, RelaSec);
+}
+
+template <class ELFT>
+Expected<std::vector<PGOBBAddrMap>>
+ELFFile<ELFT>::decodePGOBBAddrMap(const Elf_Shdr &Sec,
+                                  const Elf_Shdr *RelaSec) const {
+  return decodeBBAddrMapCommonImpl<PGOBBAddrMap>(*this, Sec, RelaSec);
+}
+
 template <class ELFT>
 Expected<
     MapVector<const typename ELFT::Shdr *, const typename ELFT::Shdr *>>
diff --git a/llvm/lib/Object/ELFObjectFile.cpp b/llvm/lib/Object/ELFObjectFile.cpp
index 143f9d37849d238..0c200f090802269 100644
--- a/llvm/lib/Object/ELFObjectFile.cpp
+++ b/llvm/lib/Object/ELFObjectFile.cpp
@@ -708,17 +708,18 @@ std::vector<ELFPltEntry> ELFObjectFileBase::getPltEntries() const {
   return Result;
 }
 
-template <class ELFT>
-Expected<std::vector<BBAddrMap>> static readBBAddrMapImpl(
-    const ELFFile<ELFT> &EF, std::optional<unsigned> TextSectionIndex) {
+template <typename AddrMap, typename ELFT, typename DecodeAddrMapFn>
+Expected<std::vector<AddrMap>> static readBBAddrMapCommonImpl(
+    const ELFFile<ELFT> &EF, std::optional<unsigned> TextSectionIndex,
+    ArrayRef<unsigned> SectionIDs, DecodeAddrMapFn DecodeAddrMap) {
   using Elf_Shdr = typename ELFT::Shdr;
   bool IsRelocatable = EF.getHeader().e_type == ELF::ET_REL;
-  std::vector<BBAddrMap> BBAddrMaps;
+  std::vector<AddrMap> BBAddrMaps;
 
   const auto &Sections = cantFail(EF.sections());
   auto IsMatch = [&](const Elf_Shdr &Sec) -> Expected<bool> {
-    if (Sec.sh_type != ELF::SHT_LLVM_BB_ADDR_MAP &&
-        Sec.sh_type != ELF::SHT_LLVM_BB_ADDR_MAP_V0)
+    if (llvm::none_of(SectionIDs,
+                      [&](unsigned ID) { return ID == Sec.sh_type; }))
       return false;
     if (!TextSectionIndex)
       return true;
@@ -741,8 +742,8 @@ Expected<std::vector<BBAddrMap>> static readBBAddrMapImpl(
     if (IsRelocatable && !RelocSec)
       return createError("unable to get relocation section for " +
                          describe(EF, *Sec));
-    Expected<std::vector<BBAddrMap>> BBAddrMapOrErr =
-        EF.decodeBBAddrMap(*Sec, RelocSec);
+    Expected<std::vector<AddrMap>> BBAddrMapOrErr =
+        DecodeAddrMap(EF, *Sec, RelocSec);
     if (!BBAddrMapOrErr)
       return createError("unable to read " + describe(EF, *Sec) + ": " +
                          toString(BBAddrMapOrErr.takeError()));
@@ -752,6 +753,29 @@ Expected<std::vector<BBAddrMap>> static readBBAddrMapImpl(
   return BBAddrMaps;
 }
 
+template <class ELFT>
+Expected<std::vector<BBAddrMap>> static readBBAddrMapImpl(
+    const ELFFile<ELFT> &EF, std::optional<unsigned> TextSectionIndex) {
+  return readBBAddrMapCommonImpl<BBAddrMap>(
+      EF, TextSectionIndex,
+      {ELF::SHT_LLVM_BB_ADDR_MAP, ELF::SHT_LLVM_BB_ADDR_MAP_V0},
+      [](const ELFFile<ELFT> &EF, const typename ELFFile<ELFT>::Elf_Shdr &Sec,
+         const typename ELFFile<ELFT>::Elf_Shdr *RelaSec) {
+        return EF.decodeBBAddrMap(Sec, RelaSec);
+      });
+}
+
+template <class ELFT>
+Expected<std::vector<PGOBBAddrMap>> static readPGOBBAddrMapImpl(
+    const ELFFile<ELFT> &EF, std::optional<unsigned> TextSectionIndex) {
+  return readBBAddrMapCommonImpl<PGOBBAddrMap>(
+      EF, TextSectionIndex, {ELF::SHT_LLVM_PGO_BB_ADDR_MAP},
+      [](const ELFFile<ELFT> &EF, const typename ELFFile<ELFT>::Elf_Shdr &Sec,
+         const typename ELFFile<ELFT>::Elf_Shdr *RelaSec) {
+        return EF.decodePGOBBAddrMap(Sec, RelaSec);
+      });
+}
+
 template <class ELFT>
 static Expected<std::vector<VersionEntry>>
 readDynsymVersionsImpl(const ELFFile<ELFT> &EF,
@@ -832,3 +856,15 @@ Expected<std::vector<BBAddrMap>> ELFObjectFileBase::readBBAddrMap(
   return readBBAddrMapImpl(cast<ELF64BEObjectFile>(this)->getELFFile(),
                            TextSectionIndex);
 }
+
+Expected<std::vector<PGOBBAddrMap>> ELFObjectFileBase::readPGOBBAddrMap(
+    std::optional<unsigned> TextSectionIndex) const {
+  if (const auto *Obj = dyn_cast<ELF32LEObjectFile>(this))
+    return readPGOBBAddrMapImpl(Obj->getELFFile(), TextSectionIndex);
+  if (const auto *Obj = dyn_cast<ELF64LEObjectFile>(this))
+    return readPGOBBAddrMapImpl(Obj->getELFFile(), TextSectionIndex);
+  if (const auto *Obj = dyn_cast<ELF32BEObjectFile>(this))
+    return readPGOBBAddrMapImpl(Obj->getELFFile(), TextSectionIndex);
+  return readPGOBBAddrMapImpl(cast<ELF64BEObjectFile>(this)->getELFFile(),
+                              TextSectionIndex);
+}
diff --git a/llvm/lib/ObjectYAML/ELFEmitter.cpp b/llvm/lib/ObjectYAML/ELFEmitter.cpp
index 40f81f867efa423..4265c4bfd226eac 100644
--- a/llvm/lib/ObjectYAML/ELFEmitter.cpp
+++ b/llvm/lib/ObjectYAML/ELFEmitter.cpp
@@ -175,8 +175,9 @@ struct Fragment {
 /// TODO: This class still has a ways to go before it is truly a "single
 /// point of truth".
 template <class ELFT> class ELFState {
+public:
   LLVM_ELF_IMPORT_TYPES_ELFT(ELFT)
-
+private:
   enum class SymtabType { Static, Dynamic };
 
   /// The future symbol table string section.
@@ -283,6 +284,9 @@ template <class ELFT> class ELFState {
   void writeSectionContent(Elf_Shdr &SHeader,
                            const ELFYAML::BBAddrMapSection &Section,
                            ContiguousBlobAccumulator &CBA);
+  void writeSectionContent(Elf_Shdr &SHeader,
+                           const ELFYAML::PGOBBAddrMapSection &Section,
+                           ContiguousBlobAccumulator &CBA);
   void writeSectionContent(Elf_Shdr &SHeader,
                            const ELFYAML::HashSection &Section,
                            ContiguousBlobAccumulator &CBA);
@@ -894,6 +898,8 @@ void ELFState<ELFT>::initSectionHeaders(std::vector<Elf_Shdr> &SHeaders,
       writeSectionContent(SHeader, *S, CBA);
     } else if (auto S = dyn_cast<ELFYAML::BBAddrMapSection>(Sec)) {
       writeSectionContent(SHeader, *S, CBA);
+    } else if (auto S = dyn_cast<ELFYAML::PGOBBAddrMapSection>(Sec)) {
+      writeSectionContent(SHeader, *S, CBA);
     } else {
       llvm_unreachable("Unknown section type");
     }
@@ -1386,44 +1392,124 @@ void ELFState<ELFT>::writeSectionContent(
   }
 }
 
-template <class ELFT>
-void ELFState<ELFT>::writeSectionContent(
-    Elf_Shdr &SHeader, const ELFYAML::BBAddrMapSection &Section,
-    ContiguousBlobAccumulator &CBA) {
-  if (!Section.Entries)
-    return;
-
-  for (const ELFYAML::BBAddrMapEntry &E : *Section.Entries) {
+namespace {
+template <typename ELFT, typename AddrMap> struct BBAddrMapWriteTrait;
+
+template <typename ELFT>
+struct BBAddrMapWriteTrait<ELFT, ELFYAML::BBAddrMapSection> {
+  typename ELFState<ELFT>::Elf_Shdr &SHeader;
+  const ELFYAML::Section &Section;
+  ContiguousBlobAccumulator &CBA;
+  const unsigned UpToDateSectionID;
+
+  BBAddrMapWriteTrait(typename ELFState<ELFT>::Elf_Shdr &SHeader,
+                      const ELFYAML::Section &Section,
+                      ContiguousBlobAccumulator &CBA,
+                      unsigned UpToDateSectionID = ELF::SHT_LLVM_BB_ADDR_MAP)
+      : SHeader(SHeader), Section(Section), CBA(CBA),
+        UpToDateSectionID(UpToDateSectionID) {}
+
+  template <typename EntryTy> void writeEntry(const EntryTy &E) {
+    using uintX_t = typename ELFState<ELFT>::uintX_t;
     // Write version and feature values.
-    if (Section.Type == llvm::ELF::SHT_LLVM_BB_ADDR_MAP) {
-      if (E.Version > 2)
+    if (Section.Type == UpToDateSectionID) {
+      if (E.Common.Version > 2)
         WithColor::warning() << "unsupported SHT_LLVM_BB_ADDR_MAP version: "
-                             << static_cast<int>(E.Version)
+                             << static_cast<int>(E.Common.Version)
                              << "; encoding using the most recent version";
-      CBA.write(E.Version);
-      CBA.write(E.Feature);
+      CBA.write(E.Common.Version);
+      CBA.write(E.Common.Feature);
       SHeader.sh_size += 2;
     }
     // Write the address of the function.
-    CBA.write<uintX_t>(E.Address, ELFT::TargetEndianness);
+    CBA.write<uintX_t>(E.Common.Address, ELFT::TargetEndianness);
     // Write number of BBEntries (number of basic blocks in the function). This
     // is overridden by the 'NumBlocks' YAML field when specified.
     uint64_t NumBlocks =
-        E.NumBlocks.value_or(E.BBEntries ? E.BBEntries->size() : 0);
+        E.Common.NumBlocks.value_or(E.BBEntries ? E.BBEntries->size() : 0);
     SHeader.sh_size += sizeof(uintX_t) + CBA.writeULEB128(NumBlocks);
+  }
+
+  void writeBBEntry(const ELFYAML::BBAddrMapEntry &E,
+                    const ELFYAML::BBAddrMapEntry::BBEntry &BBE) {
+    if (Section.Type == UpToDateSectionID && E.Common.Version > 1)
+      SHeader.sh_size += CBA.writeULEB128(BBE.ID);
+    SHeader.sh_size += CBA.writeULEB128(BBE.AddressOffset) +
+                       CBA.writeULEB128(BBE.Size) +
+                       CBA.writeULEB128(BBE.Metadata);
+  }
+};
+
+template <typename ELFT>
+struct BBAddrMapWriteTrait<ELFT, ELFYAML::PGOBBAddrMapSection> {
+  using Features = object::PGOBBAddrMap::Features;
+
+  BBAddrMapWriteTrait<ELFT, ELFYAML::BBAddrMapSection> Base;
+
+  BBAddrMapWriteTrait(typename ELFState<ELFT>::Elf_Shdr &SHeader,
+                      const ELFYAML::PGOBBAddrMapSection &Section,
+                      ContiguousBlobAccumulator &CBA)
+      : Base(SHeader, Section, CBA, ELF::SHT_LLVM_PGO_BB_ADDR_MAP) {}
+
+  void writeEntry(const ELFYAML::PGOBBAddrMapEntry &E) {
+    if (E.Common.Version < 2)
+      WithColor::warning() << "unsupported SHT_LLVM_PGO_BB_ADDR_MAP version: "
+                           << static_cast<int>(E.Common.Version)
+                           << "; must use version >= 2";
+    Base.writeEntry(E);
+    if (E.FuncEntryCount)
+      Base.SHeader.sh_size += Base.CBA.writeULEB128(*E.FuncEntryCount);
+  }
+
+  void writeBBEntry(const ELFYAML::PGOBBAddrMapEntry &E,
+                    const ELFYAML::PGOBBAddrMapEntry::BBEntry &BBE) {
+    Base.writeBBEntry(ELFYAML::BBAddrMapEntry{E.Common, {}}, BBE.Base);
+    if (BBE.BBFreq)
+      Base.SHeader.sh_size += Base.CBA.writeULEB128(*BBE.BBFreq);
+    if (BBE.Successors) {
+      if (BBE.Successors->size() > 2)
+        Base.SHeader.sh_size += Base.CBA.writeULEB128(BBE.Successors->size());
+      for (const auto &Succ : *BBE.Successors)
+        Base.SHeader.sh_size +=
+            Base.CBA.writeULEB128(Succ.ID) + Base.CBA.writeULEB128(Succ.BrProb);
+    }
+  }
+};
+} // namespace
+
+template <typename ELFT, typename SectionTy>
+static void writeAddrMapSectionContent(
+    ELFState<ELFT> &ES, typename ELFState<ELFT>::Elf_Shdr &SHeader,
+    const SectionTy &Section, ContiguousBlobAccumulator &CBA) {
+  if (!Section.Entries)
+    return;
+
+  BBAddrMapWriteTrait<ELFT, SectionTy> AddrTrait(SHeader, Section, CBA);
+  for (const auto &E : *Section.Entries) {
+    AddrTrait.writeEntry(E);
     // Write all BBEntries.
     if (!E.BBEntries)
       continue;
-    for (const ELFYAML::BBAddrMapEntry::BBEntry &BBE : *E.BBEntries) {
-      if (Section.Type == llvm::ELF::SHT_LLVM_BB_ADDR_MAP && E.Version > 1)
-        SHeader.sh_size += CBA.writeULEB128(BBE.ID);
-      SHeader.sh_size += CBA.writeULEB128(BBE.AddressOffset) +
-                         CBA.writeULEB128(BBE.Size) +
-                         CBA.writeULEB128(BBE.Metadata);
+    for (const auto &BBE : *E.BBEntries) {
+      AddrTrait.writeBBEntry(E, BBE);
     }
   }
 }
 
+template <class ELFT>
+void ELFState<ELFT>::writeSectionContent(
+    Elf_Shdr &SHeader, const ELFYAML::BBAddrMapSection &Section,
+    ContiguousBlobAccumulator &CBA) {
+  writeAddrMapSectionContent(*this, SHeader, Section, CBA);
+}
+
+template <class ELFT>
+void ELFState<ELFT>::writeSectionContent(
+    Elf_Shdr &SHeader, const ELFYAML::PGOBBAddrMapSection &Section,
+    ContiguousBlobAccumulator &CBA) {
+  writeAddrMapSectionContent(*this, SHeader, Section, CBA);
+}
+
 template <class ELFT>
 void ELFState<ELFT>::writeSectionContent(
     Elf_Shdr &SHeader, const ELFYAML::LinkerOptionsSection &Section,
diff --git a/llvm/lib/ObjectYAML/ELFYAML.cpp b/llvm/lib/ObjectYAML/ELFYAML.cpp
index 872b89420a9e7a7..8aa79b6724e0a07 100644
--- a/llvm/lib/ObjectYAML/ELFYAML.cpp
+++ b/llvm/lib/ObjectYAML/ELFYAML.cpp
@@ -683,6 +683,7 @@ void ScalarEnumerationTraits<ELFYAML::ELF_SHT>::enumeration(
   ECase(SHT_LLVM_PART_PHDR);
   ECase(SHT_LLVM_BB_ADDR_MAP_V0);
   ECase(SHT_LLVM_BB_ADDR_MAP);
+  ECase(SHT_LLVM_PGO_BB_ADDR_MAP);
   ECase(SHT_LLVM_OFFLOADING);
   ECase(SHT_LLVM_LTO);
   ECase(SHT_GNU_ATTRIBUTES);
@@ -1389,6 +1390,12 @@ static void sectionMapping(IO &IO, ELFYAML::BBAddrMapSection &Section) {
   IO.mapOptional("Entries", Section.Entries);
 }
 
+static void sectionMapping(IO &IO, ELFYAML::PGOBBAddrMapSection &Section) {
+  commonSectionMapping(IO, Section);
+  IO.mapOptional("Content", Section.Content);
+  IO.mapOptional("Entries", Section.Entries);
+}
+
 static void sectionMapping(IO &IO, ELFYAML::StackSizesSection &Section) {
   commonSectionMapping(IO, Section);
   IO.mapOptional("Entries", Section.Entries);
@@ -1682,6 +1689,11 @@ void MappingTraits<std::unique_ptr<ELFYAML::Chunk>>::mapping(
       Section.reset(new ELFYAML::BBAddrMapSection());
     sectionMapping(IO, *cast<ELFYAML::BBAddrMapSection>(Section.get()));
     break;
+  case ELF::SHT_LLVM_PGO_BB_ADDR_MAP:
+    if (!IO.outputting())
+      Section.reset(new ELFYAML::PGOBBAddrMapSection());
+    sectionMapping(IO, *cast<ELFYAML::PGOBBAddrMapSection>(Section.get()));
+    break;
   default:
     if (!IO.outputting()) {
       StringRef Name;
@@ -1803,13 +1815,17 @@ void MappingTraits<ELFYAML::StackSizeEntry>::mapping(
   IO.mapRequired("Size", E.Size);
 }
 
-void MappingTraits<ELFYAML::BBAddrMapEntry>::mapping(
-    IO &IO, ELFYAML::BBAddrMapEntry &E) {
-  assert(IO.getContext() && "The IO context is not initialized");
+static void mapBBAddrMapCommonBase(IO &IO, ELFYAML::BBAddrMapCommonBase &E) {
   IO.mapRequired("Version", E.Version);
   IO.mapOptional("Feature", E.Feature, Hex8(0));
   IO.mapOptional("Address", E.Address, Hex64(0));
   IO.mapOptional("NumBlocks", E.NumBlocks);
+}
+
+void MappingTraits<ELFYAML::BBAddrMapEntry>::mapping(
+    IO &IO, ELFYAML::BBAddrMapEntry &E) {
+  assert(IO.getContext() && "The IO context is not initialized");
+  mapBBAddrMapCommonBase(IO, E.Common);
   IO.mapOptional("BBEntries", E.BBEntries);
 }
 
@@ -1822,6 +1838,29 @@ void MappingTraits<ELFYAML::BBAddrMapEntry::BBEntry>::mapping(
   IO.mapRequired("Metadata", E.Metadata);
 }
 
+void MappingTraits<ELFYAML::PGOBBAddrMapEntry>::mapping(
+    IO &IO, ELFYAML::PGOBBAddrMapEntry &E) {
+  assert(IO.getContext() && "The IO context is not initialized");
+  mapBBAddrMapCommonBase(IO, E.Common);
+  IO.mapOptional("FuncEntryCount", E.FuncEntryCount);
+  IO.mapOptional("BBEntries", E.BBEntries);
+}
+
+void MappingTraits<ELFYAML::PGOBBAddrMapEntry::BBEntry>::mapping(
+    IO &IO, ELFYAML::PGOBBAddrMapEntry::BBEntry &E) {
+  assert(IO.getContext() && "The IO context is not initialized");
+  MappingTraits<ELFYAML::BBAddrMapEntry::BBEntry>::mapping(IO, E.Base);
+  IO.mapOptional("BBFreq", E.BBFreq);
+  IO.mapOptional("Successors", E.Successors);
+}
+
+void MappingTraits<ELFYAML::PGOBBAddrMapEntry::BBEntry::SuccessorEntry>::
+    mapping(IO &IO, ELFYAML::PGOBBAddrMapEntry::BBEntry::SuccessorEntry &E) {
+  assert(IO.getContext() && "The IO context is not initialized");
+  IO.mapRequired("ID", E.ID);
+  IO.mapRequired("BrProb", E.BrProb);
+}
+
 void MappingTraits<ELFYAML::GnuHashHeader>::mapping(IO &IO,
                                                     ELFYAML::GnuHashHeader &E) {
   assert(IO.getContext() && "The IO context is not initialized");
diff --git a/llvm/test/MC/AsmParser/llvm_section_types.s b/llvm/test/MC/AsmParser/llvm_section_types.s
index 147b1499d2b8883..2d7b650e10432ae 100644
--- a/llvm/test/MC/AsmParser/llvm_section_types.s
+++ b/llvm/test/MC/AsmParser/llvm_section_types.s
@@ -17,6 +17,8 @@
 .byte 1
 .section    .section8,"", at llvm_lto
 .byte 1
+.section    .section9,"", at llvm_pgo_bb_addr_map
+.byte 1
 
 # CHECK:        Name: .section1
 # CHECK-NEXT:   Type: SHT_LLVM_BB_ADDR_MAP
@@ -34,3 +36,5 @@
 # CHECK-NEXT:   Type: SHT_LLVM_OFFLOADING
 # CHECK:        Name: .section8
 # CHECK-NEXT:   Type: SHT_LLVM_LTO
+# CHECK:        Name: .section9
+# CHECK-NEXT:   Type: SHT_LLVM_PGO_BB_ADDR_MAP
diff --git a/llvm/tools/obj2yaml/elf2yaml.cpp b/llvm/tools/obj2yaml/elf2yaml.cpp
index bec4a10a5d35f2a..653e97959795014 100644
--- a/llvm/tools/obj2yaml/elf2yaml.cpp
+++ b/llvm/tools/obj2yaml/elf2yaml.cpp
@@ -915,7 +915,7 @@ ELFDumper<ELFT>::dumpBBAddrMapSection(const Elf_Shdr *Shdr) {
       BBEntries.push_back({ID, Offset, Size, Metadata});
     }
     Entries.push_back(
-        {Version, Feature, Address, /*NumBlocks=*/{}, std::move(BBEntries)});
+        {{Version, Feature, Address, /*NumBlocks=*/{}}, std::move(BBEntries)});
   }
 
   if (!Cur) {
diff --git a/llvm/unittests/Object/ELFObjectFileTest.cpp b/llvm/unittests/Object/ELFObjectFileTest.cpp
index 17402f39a5df834..7ae9a2055040bdb 100644
--- a/llvm/unittests/Object/ELFObjectFileTest.cpp
+++ b/llvm/unittests/Object/ELFObjectFileTest.cpp
@@ -744,6 +744,446 @@ TEST(ELFObjectFileTest, ReadBBAddrMap) {
                   Section1BBAddrMaps);
 }
 
+// Tests for error paths of the ELFFile::decodePGOBBAddrMap API.
+TEST(ELFObjectFileTest, InvalidDecodePGOBBAddrMap) {
+  if (IsHostWindows())
+    GTEST_SKIP();
+  StringRef CommonYamlString(R"(
+--- !ELF
+FileHeader:
+  Class: ELFCLASS64
+  Data:  ELFDATA2LSB
+  Type:  ET_EXEC
+Sections:
+  - Type: SHT_LLVM_PGO_BB_ADDR_MAP
+    Name: .llvm_pgo_bb_addr_map
+    Entries:
+      - Address: 0x11111
+)");
+
+  auto DoCheck = [&](StringRef YamlString, const char *ErrMsg) {
+    SmallString<0> Storage;
+    Expected<ELFObjectFile<ELF64LE>> ElfOrErr =
+        toBinary<ELF64LE>(Storage, YamlString);
+    ASSERT_THAT_EXPECTED(ElfOrErr, Succeeded());
+    const ELFFile<ELF64LE> &Elf = ElfOrErr->getELFFile();
+
+    Expected<const typename ELF64LE::Shdr *> PGOBBAddrMapSecOrErr =
+        Elf.getSection(1);
+    ASSERT_THAT_EXPECTED(PGOBBAddrMapSecOrErr, Succeeded());
+    EXPECT_THAT_ERROR(
+        Elf.decodePGOBBAddrMap(**PGOBBAddrMapSecOrErr).takeError(),
+        FailedWithMessage(ErrMsg));
+  };
+
+  // Check that we can detect unsupported versions.
+  SmallString<128> UnsupportedVersionYamlString(CommonYamlString);
+  UnsupportedVersionYamlString += R"(
+        Version: 3
+        BBEntries:
+          - AddressOffset: 0x0
+            Size:          0x1
+            Metadata:      0x2
+)";
+
+  DoCheck(UnsupportedVersionYamlString,
+          "unsupported SHT_LLVM_BB_ADDR_MAP version: 3");
+
+  // Check that we can detect unsupported versions that is too low
+  SmallString<128> UnsupportedLowVersionYamlString(CommonYamlString);
+  UnsupportedLowVersionYamlString += R"(
+        Version: 1
+        BBEntries:
+          - AddressOffset: 0x0
+            Size:          0x1
+            Metadata:      0x2
+)";
+
+  DoCheck(UnsupportedLowVersionYamlString,
+          "unsupported SHT_LLVM_PGO_BB_ADDR_MAP version: 1");
+
+  SmallString<128> CommonVersionedYamlString(CommonYamlString);
+  CommonVersionedYamlString += R"(
+        Version: 2
+        BBEntries:
+          - ID:            1
+            AddressOffset: 0x0
+            Size:          0x1
+            Metadata:      0x2
+)";
+
+  // Check that we can detect the malformed encoding when the section is
+  // truncated.
+  SmallString<128> TruncatedYamlString(CommonVersionedYamlString);
+  TruncatedYamlString += R"(
+    ShSize: 0xb
+)";
+  DoCheck(TruncatedYamlString, "unable to decode LEB128 at offset 0x0000000b: "
+                               "malformed uleb128, extends past end");
+
+  // Check that we can detect when the encoded BB entry fields exceed the UINT32
+  // limit.
+  SmallVector<SmallString<128>, 3> OverInt32LimitYamlStrings(
+      3, CommonVersionedYamlString);
+  OverInt32LimitYamlStrings[0] += R"(
+          - ID:            1
+            AddressOffset: 0x100000000
+            Size:          0xFFFFFFFF
+            Metadata:      0xFFFFFFFF
+)";
+
+  OverInt32LimitYamlStrings[1] += R"(
+          - ID:            2
+            AddressOffset: 0xFFFFFFFF
+            Size:          0x100000000
+            Metadata:      0xFFFFFFFF
+)";
+
+  OverInt32LimitYamlStrings[2] += R"(
+          - ID:            3
+            AddressOffset: 0xFFFFFFFF
+            Size:          0xFFFFFFFF
+            Metadata:      0x100000000
+)";
+
+  DoCheck(OverInt32LimitYamlStrings[0],
+          "ULEB128 value at offset 0x10 exceeds UINT32_MAX (0x100000000)");
+  DoCheck(OverInt32LimitYamlStrings[1],
+          "ULEB128 value at offset 0x15 exceeds UINT32_MAX (0x100000000)");
+  DoCheck(OverInt32LimitYamlStrings[2],
+          "ULEB128 value at offset 0x1a exceeds UINT32_MAX (0x100000000)");
+
+  // Check the proper error handling when the section has fields exceeding
+  // UINT32 and is also truncated. This is for checking that we don't generate
+  // unhandled errors.
+  SmallVector<SmallString<128>, 3> OverInt32LimitAndTruncated(
+      3, OverInt32LimitYamlStrings[1]);
+  // Truncate before the end of the 5-byte field.
+  OverInt32LimitAndTruncated[0] += R"(
+    ShSize: 0x19
+)";
+  // Truncate at the end of the 5-byte field.
+  OverInt32LimitAndTruncated[1] += R"(
+    ShSize: 0x1a
+)";
+  // Truncate after the end of the 5-byte field.
+  OverInt32LimitAndTruncated[2] += R"(
+    ShSize: 0x1b
+)";
+
+  DoCheck(OverInt32LimitAndTruncated[0],
+          "unable to decode LEB128 at offset 0x00000015: malformed uleb128, "
+          "extends past end");
+  DoCheck(OverInt32LimitAndTruncated[1],
+          "ULEB128 value at offset 0x15 exceeds UINT32_MAX (0x100000000)");
+  DoCheck(OverInt32LimitAndTruncated[2],
+          "ULEB128 value at offset 0x15 exceeds UINT32_MAX (0x100000000)");
+
+  // Check for proper error handling when the 'NumBlocks' field is overridden
+  // with an out-of-range value.
+  SmallString<128> OverLimitNumBlocks(CommonVersionedYamlString);
+  OverLimitNumBlocks += R"(
+        NumBlocks: 0x100000000
+)";
+
+  DoCheck(OverLimitNumBlocks,
+          "ULEB128 value at offset 0xa exceeds UINT32_MAX (0x100000000)");
+
+  // Check that we fail when function entry count is enabled but not provided.
+  SmallString<128> MissingFuncEntryCount(CommonYamlString);
+  MissingFuncEntryCount += R"(
+        Version: 2
+        Feature: 0x01
+)";
+
+  DoCheck(MissingFuncEntryCount,
+          "unable to decode LEB128 at offset 0x0000000b: malformed uleb128, "
+          "extends past end");
+
+  // Check that we fail when basic block frequency is enabled but not provided.
+  SmallString<128> MissingBBFreq(CommonYamlString);
+  MissingBBFreq += R"(
+        Version: 2
+        Feature: 0x02
+        BBEntries:
+          - ID:            1
+            AddressOffset: 0x0
+            Size:          0x1
+            Metadata:      0x2
+)";
+
+  DoCheck(MissingBBFreq, "unable to decode LEB128 at offset 0x0000000f: "
+                         "malformed uleb128, extends past end");
+
+  // Check that we fail when branch probability is enabled but not provided.
+  SmallString<128> MissingBrProb(CommonYamlString);
+  MissingBrProb += R"(
+        Version: 2
+        Feature: 0x02
+        BBEntries:
+          - ID:            1
+            AddressOffset: 0x0
+            Size:          0x1
+            Metadata:      0x60
+          - ID:            2
+            AddressOffset: 0x1
+            Size:          0x1
+            Metadata:      0x2
+          - ID:            3
+            AddressOffset: 0x2
+            Size:          0x1
+            Metadata:      0x2
+          - ID:            4
+            AddressOffset: 0x3
+            Size:          0x1
+            Metadata:      0x2
+)";
+
+  DoCheck(MissingBrProb, "unable to decode LEB128 at offset 0x0000001b: "
+                         "malformed uleb128, extends past end");
+}
+
+// Test for the ELFObjectFile::readPGOBBAddrMap API.
+TEST(ELFObjectFileTest, ReadPGOBBAddrMap) {
+  if (IsHostWindows())
+    GTEST_SKIP();
+  StringRef CommonYamlString(R"(
+--- !ELF
+FileHeader:
+  Class: ELFCLASS64
+  Data:  ELFDATA2LSB
+  Type:  ET_EXEC
+Sections:
+  - Name: .llvm_pgo_bb_addr_map_1
+    Type: SHT_LLVM_PGO_BB_ADDR_MAP
+    Link: 1
+    Entries:
+      - Version: 2
+        Address: 0x11111
+        Feature: 0x1
+        FuncEntryCount: 892
+        BBEntries:
+          - ID:            1
+            AddressOffset: 0x0
+            Size:          0x1
+            Metadata:      0x2
+  - Name: .llvm_pgo_bb_addr_map_2
+    Type: SHT_LLVM_PGO_BB_ADDR_MAP
+    Link: 1
+    Entries:
+      - Version: 2
+        Address: 0x22222
+        Feature: 0x2
+        BBEntries:
+          - ID:            2
+            AddressOffset: 0x0
+            Size:          0x2
+            Metadata:      0x4
+            BBFreq:        343
+  - Name: .llvm_bb_addr_map_3
+    Type: SHT_LLVM_PGO_BB_ADDR_MAP
+    Link: 2
+    Entries:
+      - Version: 2
+        Address: 0x33333
+        Feature: 0x4
+        BBEntries:
+          - ID:            0
+            AddressOffset: 0x0
+            Size:          0x3
+            Metadata:      0x46
+            Successors:
+            - ID:          1
+              BrProb:      0x11111111
+            - ID:          2
+              BrProb:      0xeeeeeeee
+          - ID:            1
+            AddressOffset: 0x0
+            Size:          0x3
+            Metadata:      0x24
+            Successors:
+            - ID:          2
+              BrProb:      0xffffffff
+          - ID:            2
+            AddressOffset: 0x0
+            Size:          0x3
+            Metadata:      0x00
+            Successors:    []
+  - Name: .llvm_pgo_bb_addr_map_4
+    Type: SHT_LLVM_PGO_BB_ADDR_MAP
+  # Link: 0 (by default, can be overriden)
+    Entries:
+      - Version: 2
+        Address: 0x44444
+        Feature: 0x7
+        FuncEntryCount: 1000
+        BBEntries:
+          - ID:            0
+            AddressOffset: 0x0
+            Size:          0x4
+            Metadata:      0x78
+            BBFreq:        1000
+            Successors:
+            - ID:          1
+              BrProb:      0x22222222
+            - ID:          2
+              BrProb:      0x33333333
+            - ID:          3
+              BrProb:      0xaaaaaaaa
+          - ID:            1
+            AddressOffset: 0x0
+            Size:          0x4
+            Metadata:      0x40
+            BBFreq:        133
+            Successors:
+            - ID:          2
+              BrProb:      0x11111111
+            - ID:          3
+              BrProb:      0xeeeeeeee
+          - ID:            2
+            AddressOffset: 0x0
+            Size:          0x4
+            Metadata:      0x20
+            BBFreq:        18
+            Successors:
+            - ID:          3
+              BrProb:      0xffffffff
+          - ID:            3
+            AddressOffset: 0x0
+            Size:          0x4
+            Metadata:      0x00
+            BBFreq:        1000
+            Successors:    []
+)");
+
+  PGOBBAddrMap E1 = {
+      0x11111, {{{1, 0x0, 0x1, {false, true, false, false, false}}, {}, {}}},
+      892,     true,
+      false,   false};
+  PGOBBAddrMap E2 = {0x22222,
+                     {{{2, 0x0, 0x2, {false, false, true, false, false}},
+                       BlockFrequency(343),
+                       {}}},
+                     0,
+                     false,
+                     true,
+                     false};
+  PGOBBAddrMap E3 = {
+      0x33333,
+      {{{0, 0x0, 0x3, {false, true, true, false, false}},
+        {},
+        {{1, BranchProbability::getRaw(0x1111'1111)},
+         {2, BranchProbability::getRaw(0xeeee'eeee)}}},
+       {{1, 0x3, 0x3, {false, false, true, false, false}},
+        {},
+        {{2, BranchProbability::getRaw(0xffff'ffff)}}},
+       {{2, 0x6, 0x3, {false, false, false, false, false}}, {}, {}}},
+      0,
+      false,
+      false,
+      true};
+  PGOBBAddrMap E4 = {0x44444,
+                     {{{0, 0x0, 0x4, {false, false, false, true, true}},
+                       BlockFrequency(1000),
+                       {{1, BranchProbability::getRaw(0x2222'2222)},
+                        {2, BranchProbability::getRaw(0x3333'3333)},
+                        {3, BranchProbability::getRaw(0xaaaa'aaaa)}}},
+                      {{1, 0x4, 0x4, {false, false, false, false, false}},
+                       BlockFrequency(133),
+                       {{2, BranchProbability::getRaw(0x1111'1111)},
+                        {3, BranchProbability::getRaw(0xeeee'eeee)}}},
+                      {{2, 0x8, 0x4, {false, false, false, false, false}},
+                       BlockFrequency(18),
+                       {{3, BranchProbability::getRaw(0xffff'ffff)}}},
+                      {{3, 0xc, 0x4, {false, false, false, false, false}},
+                       BlockFrequency(1000),
+                       {}}},
+                     1000,
+                     true,
+                     true,
+                     true};
+
+  std::vector<PGOBBAddrMap> Section0PGOBBAddrMaps = {E4};
+  std::vector<PGOBBAddrMap> Section1PGOBBAddrMaps = {E3};
+  std::vector<PGOBBAddrMap> Section2PGOBBAddrMaps = {E1, E2};
+  std::vector<PGOBBAddrMap> AllPGOBBAddrMaps = {E1, E2, E3, E4};
+
+  auto DoCheckSucceeds = [&](StringRef YamlString,
+                             std::optional<unsigned> TextSectionIndex,
+                             std::vector<PGOBBAddrMap> ExpectedResult) {
+    SmallString<0> Storage;
+    Expected<ELFObjectFile<ELF64LE>> ElfOrErr =
+        toBinary<ELF64LE>(Storage, YamlString);
+    ASSERT_THAT_EXPECTED(ElfOrErr, Succeeded());
+
+    Expected<const typename ELF64LE::Shdr *> PGOBBAddrMapSecOrErr =
+        ElfOrErr->getELFFile().getSection(1);
+    ASSERT_THAT_EXPECTED(PGOBBAddrMapSecOrErr, Succeeded());
+    auto PGOBBAddrMaps = ElfOrErr->readPGOBBAddrMap(TextSectionIndex);
+    EXPECT_THAT_EXPECTED(PGOBBAddrMaps, Succeeded());
+    EXPECT_EQ(*PGOBBAddrMaps, ExpectedResult);
+  };
+
+  auto DoCheckFails = [&](StringRef YamlString,
+                          std::optional<unsigned> TextSectionIndex,
+                          const char *ErrMsg) {
+    SmallString<0> Storage;
+    Expected<ELFObjectFile<ELF64LE>> ElfOrErr =
+        toBinary<ELF64LE>(Storage, YamlString);
+    ASSERT_THAT_EXPECTED(ElfOrErr, Succeeded());
+
+    Expected<const typename ELF64LE::Shdr *> BBAddrMapSecOrErr =
+        ElfOrErr->getELFFile().getSection(1);
+    ASSERT_THAT_EXPECTED(BBAddrMapSecOrErr, Succeeded());
+    EXPECT_THAT_ERROR(ElfOrErr->readPGOBBAddrMap(TextSectionIndex).takeError(),
+                      FailedWithMessage(ErrMsg));
+  };
+
+  // Check that we can retrieve the data in the normal case.
+  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/std::nullopt,
+                  AllPGOBBAddrMaps);
+  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/0,
+                  Section0PGOBBAddrMaps);
+  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/2,
+                  Section1PGOBBAddrMaps);
+  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/1,
+                  Section2PGOBBAddrMaps);
+  // Check that when no bb-address-map section is found for a text section,
+  // we return an empty result.
+  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/3, {});
+
+  // Check that we detect when a bb-addr-map section is linked to an invalid
+  // (not present) section.
+  SmallString<128> InvalidLinkedYamlString(CommonYamlString);
+  InvalidLinkedYamlString += R"(
+    Link: 10
+)";
+
+  DoCheckFails(InvalidLinkedYamlString, /*TextSectionIndex=*/4,
+               "unable to get the linked-to section for "
+               "SHT_LLVM_PGO_BB_ADDR_MAP section with index 4: invalid section "
+               "index: 10");
+  // Linked sections are not checked when we don't target a specific text
+  // section.
+  DoCheckSucceeds(InvalidLinkedYamlString, /*TextSectionIndex=*/std::nullopt,
+                  AllPGOBBAddrMaps);
+
+  // Check that we can detect when bb-address-map decoding fails.
+  SmallString<128> TruncatedYamlString(CommonYamlString);
+  TruncatedYamlString += R"(
+    ShSize: 0xa
+)";
+
+  DoCheckFails(TruncatedYamlString, /*TextSectionIndex=*/std::nullopt,
+               "unable to read SHT_LLVM_PGO_BB_ADDR_MAP section with index 4: "
+               "unable to decode LEB128 at offset 0x0000000a: malformed "
+               "uleb128, extends past end");
+  // Check that we can read the other section's bb-address-maps which are
+  // valid.
+  DoCheckSucceeds(TruncatedYamlString, /*TextSectionIndex=*/2,
+                  Section1PGOBBAddrMaps);
+}
+
 // Test for ObjectFile::getRelocatedSection: check that it returns a relocated
 // section for executable and relocatable files.
 TEST(ELFObjectFileTest, ExecutableWithRelocs) {
diff --git a/llvm/unittests/Object/ELFTypesTest.cpp b/llvm/unittests/Object/ELFTypesTest.cpp
index 2c9e8b7aebac293..8689f77af805073 100644
--- a/llvm/unittests/Object/ELFTypesTest.cpp
+++ b/llvm/unittests/Object/ELFTypesTest.cpp
@@ -73,8 +73,12 @@ TEST(ELFTypesTest, BBEntryMetadataEncodingTest) {
        {false, false, false, false, true},
        {true, true, true, true, true}}};
   const std::array<uint32_t, 7> Encoded = {{0, 1, 2, 4, 8, 16, 31}};
-  for (size_t i = 0; i < Decoded.size(); ++i)
+  for (size_t i = 0; i < Decoded.size(); ++i) {
     EXPECT_EQ(Decoded[i].encode(), Encoded[i]);
+    EXPECT_LT(Decoded[i].encode(),
+              uint32_t{1} << BBAddrMap::BBEntry::Metadata::NumberOfBits)
+        << "If a new bit was added, please update NumberOfBits.";
+  }
   for (size_t i = 0; i < Encoded.size(); ++i) {
     Expected<BBAddrMap::BBEntry::Metadata> MetadataOrError =
         BBAddrMap::BBEntry::Metadata::decode(Encoded[i]);
@@ -94,3 +98,50 @@ TEST(ELFTypesTest, BBEntryMetadataInvalidEncodingTest) {
         FailedWithMessage(Errors[i]));
   }
 }
+
+static_assert(
+    std::is_same_v<decltype(PGOBBAddrMap::BBEntry::SuccessorEntry::ID),
+                   decltype(BBAddrMap::BBEntry::ID)>,
+    "PGOBBAddrMap should use the same type for basic block ID as BBAddrMap");
+static_assert(BBAddrMap::BBEntry::Metadata::NumberOfBits <
+                  (sizeof(uint32_t) * 8) - 2,
+              "currently PGOBBAddrMap relies on having two bits of space to "
+              "encode number of successors, to add more we need increase the "
+              "encoded size of Metadata");
+
+TEST(ELFTypesTest, PGOBBEntryMetadataEncodingTest) {
+  using ST = PGOBBAddrMap::BBEntry::SuccessorsType;
+  const std::array<std::pair<BBAddrMap::BBEntry::Metadata, ST>, 7> Decoded = {
+      {{{false, false, false, false, false}, ST::None},
+       {{true, false, false, false, false}, ST::One},
+       {{false, true, false, false, false}, ST::Two},
+       {{false, false, true, false, false}, ST::Multiple},
+       {{false, false, false, true, false}, ST::One},
+       {{false, false, false, false, true}, ST::Two},
+       {{true, true, true, true, true}, ST::Multiple}}};
+  const std::array<uint32_t, 7> Encoded = {{0b00'00000, 0b01'00001, 0b10'00010,
+                                            0b11'00100, 0b01'01000, 0b10'10000,
+                                            0b11'11111}};
+  for (auto [Enc, Dec] : llvm::zip(Encoded, Decoded)) {
+    auto [MD, SuccType] = Dec;
+    EXPECT_EQ(PGOBBAddrMap::BBEntry::encodeMD(MD, SuccType), Enc);
+  }
+  for (auto [Enc, Dec] : llvm::zip(Encoded, Decoded)) {
+    Expected<std::pair<BBAddrMap::BBEntry::Metadata, ST>> MetadataOrError =
+        PGOBBAddrMap::BBEntry::decodeMD(Enc);
+    ASSERT_THAT_EXPECTED(MetadataOrError, Succeeded());
+    EXPECT_EQ(*MetadataOrError, Dec);
+  }
+}
+
+TEST(ELFTypesTest, PGOBBEntryMetadataInvalidEncodingTest) {
+  const std::array<std::string, 3> Errors = {
+      "invalid encoding for BBEntry::Metadata: 0xff9f",
+      "invalid encoding for BBEntry::Metadata: 0x100001",
+      "invalid encoding for BBEntry::Metadata: 0x80"};
+  const std::array<uint32_t, 3> Values = {0xFFFF, 0x100001, 0x00c0};
+  for (auto [Val, Err] : llvm::zip(Values, Errors)) {
+    EXPECT_THAT_ERROR(PGOBBAddrMap::BBEntry::decodeMD(Val).takeError(),
+                      FailedWithMessage(Err));
+  }
+}

>From 989c0843e0f049eb39be3288caeb075076a9292e Mon Sep 17 00:00:00 2001
From: Micah Weston <micahsweston at gmail.com>
Date: Mon, 20 Nov 2023 19:25:23 -0500
Subject: [PATCH 2/8] Reworks the design to use same ELF section and not
 superset data structure.

---
 llvm/include/llvm/BinaryFormat/ELF.h        |   1 -
 llvm/include/llvm/Object/ELF.h              |  12 +-
 llvm/include/llvm/Object/ELFObjectFile.h    |  11 +-
 llvm/include/llvm/Object/ELFTypes.h         |  30 +-
 llvm/include/llvm/ObjectYAML/ELFYAML.h      |  57 +--
 llvm/lib/MC/MCParser/ELFAsmParser.cpp       |   2 -
 llvm/lib/MC/MCSectionELF.cpp                |   2 -
 llvm/lib/Object/ELF.cpp                     | 266 +++++---------
 llvm/lib/Object/ELFObjectFile.cpp           |  70 ++--
 llvm/lib/ObjectYAML/ELFEmitter.cpp          | 174 ++++------
 llvm/lib/ObjectYAML/ELFYAML.cpp             |  40 +--
 llvm/test/MC/AsmParser/llvm_section_types.s |   4 -
 llvm/tools/obj2yaml/elf2yaml.cpp            |   2 +-
 llvm/unittests/Object/ELFObjectFileTest.cpp | 365 +++++++++-----------
 llvm/unittests/Object/ELFTypesTest.cpp      |  14 +-
 15 files changed, 404 insertions(+), 646 deletions(-)

diff --git a/llvm/include/llvm/BinaryFormat/ELF.h b/llvm/include/llvm/BinaryFormat/ELF.h
index 299d7e8265317b3..3596174f74dde80 100644
--- a/llvm/include/llvm/BinaryFormat/ELF.h
+++ b/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.
diff --git a/llvm/include/llvm/Object/ELF.h b/llvm/include/llvm/Object/ELF.h
index 2efcdf2cfe714ad..cdc902559fb3bb3 100644
--- a/llvm/include/llvm/Object/ELF.h
+++ b/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
diff --git a/llvm/include/llvm/Object/ELFObjectFile.h b/llvm/include/llvm/Object/ELFObjectFile.h
index 7cf2226038b2fc9..66171f1c355cc4b 100644
--- a/llvm/include/llvm/Object/ELFObjectFile.h
+++ b/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 {
diff --git a/llvm/include/llvm/Object/ELFTypes.h b/llvm/include/llvm/Object/ELFTypes.h
index 7ab34c8caaf5c33..a89cdd58cf78afe 100644
--- a/llvm/include/llvm/Object/ELFTypes.h
+++ b/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);
   }
diff --git a/llvm/include/llvm/ObjectYAML/ELFYAML.h b/llvm/include/llvm/ObjectYAML/ELFYAML.h
index 3f18d189b7dc870..12b47c271da2cd6 100644
--- a/llvm/include/llvm/ObjectYAML/ELFYAML.h
+++ b/llvm/include/llvm/ObjectYAML/ELFYAML.h
@@ -156,13 +156,6 @@ 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;
@@ -170,23 +163,24 @@ struct BBAddrMapEntry {
     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> {
diff --git a/llvm/lib/MC/MCParser/ELFAsmParser.cpp b/llvm/lib/MC/MCParser/ELFAsmParser.cpp
index e4c1cc7763f7804..dbfe0d83e1b29b1 100644
--- a/llvm/lib/MC/MCParser/ELFAsmParser.cpp
+++ b/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")
diff --git a/llvm/lib/MC/MCSectionELF.cpp b/llvm/lib/MC/MCSectionELF.cpp
index 02e63f7b03919a8..95fdf33522076e4 100644
--- a/llvm/lib/MC/MCSectionELF.cpp
+++ b/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)
diff --git a/llvm/lib/Object/ELF.cpp b/llvm/lib/Object/ELF.cpp
index e5b2b38048b9382..6517ce3ead092ff 100644
--- a/llvm/lib/Object/ELF.cpp
+++ b/llvm/lib/Object/ELF.cpp
@@ -10,6 +10,8 @@
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/Support/DataExtractor.h"
+#include <type_traits>
+#include <vector>
 
 using namespace llvm;
 using namespace object;
@@ -312,7 +314,6 @@ StringRef llvm::object::getELFSectionTypeName(uint32_t Machine, unsigned Type) {
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_PART_PHDR);
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_BB_ADDR_MAP_V0);
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_BB_ADDR_MAP);
-    STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_PGO_BB_ADDR_MAP);
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_OFFLOADING);
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_LTO);
     STRINGIFY_ENUM_CASE(ELF, SHT_GNU_ATTRIBUTES);
@@ -672,152 +673,12 @@ static IntTy readULEB128As(DataExtractor &Data, DataExtractor::Cursor &Cur,
   return static_cast<IntTy>(Value);
 }
 
-namespace {
-struct BasicAddrMapBBEntry {
-  uint32_t ID;
-  uint32_t Offset;
-  uint32_t Size;
-  uint32_t MD;
-};
-
-template <typename ELFT, typename AddrMap> struct AddrMapDecodeTrait;
-
-template <typename ELFT> struct AddrMapDecodeTrait<ELFT, BBAddrMap> {
-  // Base addr map has no extra data
-  struct ExtraData {};
-
-  static constexpr unsigned SectionID = ELF::SHT_LLVM_BB_ADDR_MAP;
-
-  DataExtractor &Data;
-  DataExtractor::Cursor &Cur;
-  Error &ULEBSizeErr;
-  Error &MetadataDecodeErr;
-
-  uint32_t PrevBBEndOffset = 0;
-
-  AddrMapDecodeTrait(DataExtractor &Data, DataExtractor::Cursor &Cur,
-                     Error &ULEBSizeErr, Error &MetadataDecodeErr)
-      : Data(Data), Cur(Cur), ULEBSizeErr(ULEBSizeErr),
-        MetadataDecodeErr(MetadataDecodeErr) {}
-
-  ExtraData decodeExtraDataHeader(uint8_t, uint8_t) {
-    PrevBBEndOffset = 0;
-    return {};
-  }
-
-  // This helper method avoids decoding the metadata so other AddrMaps can use
-  BasicAddrMapBBEntry decodeBasicInfo(uint8_t Version, uint8_t Feature,
-                                      uint32_t BlockIndex) {
-    uint32_t ID = Version >= 2 ? readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr)
-                               : BlockIndex;
-    uint32_t Offset = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-    uint32_t Size = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-    uint32_t MD = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-    if (Version >= 1) {
-      // Offset is calculated relative to the end of the previous BB.
-      Offset += PrevBBEndOffset;
-      PrevBBEndOffset = Offset + Size;
-    }
-    return {ID, Offset, Size, MD};
-  }
-
-  BBAddrMap::BBEntry decodeBBEntry(uint8_t Version, uint8_t Feature,
-                                   uint32_t BlockIndex) {
-    auto [ID, Offset, Size, MD] = decodeBasicInfo(Version, Feature, BlockIndex);
-    auto MetadataOrErr = BBAddrMap::BBEntry::Metadata::decode(MD);
-    if (Error E = MetadataOrErr.takeError()) {
-      MetadataDecodeErr = std::move(E);
-      return {{}, {}, {}, {}};
-    }
-    return {ID, Offset, Size, *MetadataOrErr};
-  }
-
-  BBAddrMap makeAddrMap(ExtraData, uint8_t Feature, uint64_t Address,
-                        std::vector<BBAddrMap::BBEntry> BBEntries) {
-    return {Address, std::move(BBEntries)};
-  }
-};
-
-template <typename ELFT> struct AddrMapDecodeTrait<ELFT, PGOBBAddrMap> {
-  using Features = PGOBBAddrMap::Features;
-
-  struct ExtraData {
-    uint64_t FuncEntryCount;
-  };
-
-  static constexpr unsigned SectionID = ELF::SHT_LLVM_PGO_BB_ADDR_MAP;
-
-  AddrMapDecodeTrait<ELFT, BBAddrMap> Base;
-
-  AddrMapDecodeTrait(DataExtractor &Data, DataExtractor::Cursor &Cur,
-                     Error &ULEBSizeErr, Error &MetadataDecodeErr)
-      : Base(Data, Cur, ULEBSizeErr, MetadataDecodeErr) {}
-
-  ExtraData decodeExtraDataHeader(uint8_t Version, uint8_t Feature) {
-    Base.decodeExtraDataHeader(Version, Feature);
-    if (Version < 2) {
-      // hijack the metadata error if version is too low
-      Base.MetadataDecodeErr =
-          createError("unsupported SHT_LLVM_PGO_BB_ADDR_MAP version: " +
-                      Twine(static_cast<int>(Version)));
-      return {};
-    }
-    return {Feature & uint8_t(PGOBBAddrMap::Features::FuncEntryCnt)
-                ? readULEB128As<uint64_t>(Base.Data, Base.Cur, Base.ULEBSizeErr)
-                : 0};
-  }
-
-  PGOBBAddrMap::BBEntry decodeBBEntry(uint8_t Version, uint8_t Feature,
-                                      uint32_t BlockIndex) {
-    auto [ID, Offset, Size, MD] =
-        Base.decodeBasicInfo(Version, Feature, BlockIndex);
-    auto MetadataOrErr = PGOBBAddrMap::BBEntry::decodeMD(MD);
-    if (Error E = MetadataOrErr.takeError()) {
-      Base.MetadataDecodeErr = std::move(E);
-      return {{{}, {}, {}, {}}, {}, {}};
-    }
-    auto [MetaData, SuccsType] = *MetadataOrErr;
-
-    uint64_t BBF =
-        Feature & uint8_t(PGOBBAddrMap::Features::BBFreq)
-            ? readULEB128As<uint64_t>(Base.Data, Base.Cur, Base.ULEBSizeErr)
-            : 0;
-
-    llvm::SmallVector<PGOBBAddrMap::BBEntry::SuccessorEntry, 2> Successors;
-    if (Feature & uint8_t(PGOBBAddrMap::Features::BrProb)) {
-      auto SuccCount =
-          SuccsType == PGOBBAddrMap::BBEntry::SuccessorsType::Multiple
-              ? readULEB128As<uint64_t>(Base.Data, Base.Cur, Base.ULEBSizeErr)
-              : uint64_t(SuccsType);
-      for (uint64_t I = 0; I < SuccCount; ++I) {
-        uint32_t BBID =
-            readULEB128As<uint32_t>(Base.Data, Base.Cur, Base.ULEBSizeErr);
-        uint32_t BrProb =
-            readULEB128As<uint32_t>(Base.Data, Base.Cur, Base.ULEBSizeErr);
-        Successors.push_back({BBID, BranchProbability::getRaw(BrProb)});
-      }
-    }
-    return PGOBBAddrMap::BBEntry{
-        {ID, Offset, Size, MetaData}, BlockFrequency(BBF), Successors};
-  }
-
-  PGOBBAddrMap makeAddrMap(ExtraData ED, uint8_t Feature, uint64_t Address,
-                           std::vector<PGOBBAddrMap::BBEntry> BBEntries) {
-    return {Address,
-            std::move(BBEntries),
-            ED.FuncEntryCount,
-            bool(Feature & uint8_t(Features::FuncEntryCnt)),
-            bool(Feature & uint8_t(Features::BBFreq)),
-            bool(Feature & uint8_t(Features::BrProb))};
-  }
-};
-} // namespace
-
-template <typename AddrMap, typename ELFT>
-static Expected<std::vector<AddrMap>>
-decodeBBAddrMapCommonImpl(const ELFFile<ELFT> &EF,
-                          const typename ELFFile<ELFT>::Elf_Shdr &Sec,
-                          const typename ELFFile<ELFT>::Elf_Shdr *RelaSec) {
+template <typename ELFT>
+static Expected<std::vector<BBAddrMap>>
+decodeBBAddrMapImpl(const ELFFile<ELFT> &EF,
+                    const typename ELFFile<ELFT>::Elf_Shdr &Sec,
+                    const typename ELFFile<ELFT>::Elf_Shdr *RelaSec,
+                    std::vector<PGOAnalysisMap> *PGOAnalyses) {
   bool IsRelocatable = EF.getHeader().e_type == ELF::ET_REL;
 
   // This DenseMap maps the offset of each function (the location of the
@@ -841,20 +702,20 @@ decodeBBAddrMapCommonImpl(const ELFFile<ELFT> &EF,
     return ContentsOrErr.takeError();
   ArrayRef<uint8_t> Content = *ContentsOrErr;
   DataExtractor Data(Content, EF.isLE(), ELFT::Is64Bits ? 8 : 4);
-  std::vector<AddrMap> FunctionEntries;
+  std::vector<BBAddrMap> FunctionEntries;
 
   DataExtractor::Cursor Cur(0);
   Error ULEBSizeErr = Error::success();
   Error MetadataDecodeErr = Error::success();
 
-  using DecodeTrait = AddrMapDecodeTrait<ELFT, AddrMap>;
-  DecodeTrait DT(Data, Cur, ULEBSizeErr, MetadataDecodeErr);
-
   uint8_t Version = 0;
   uint8_t Feature = 0;
+  bool FuncEntryCountEnabled = false;
+  bool BBFreqEnabled = false;
+  bool BrProbEnabled = false;
   while (!ULEBSizeErr && !MetadataDecodeErr && Cur &&
          Cur.tell() < Content.size()) {
-    if (Sec.sh_type == DecodeTrait::SectionID) {
+    if (Sec.sh_type == ELF::SHT_LLVM_BB_ADDR_MAP) {
       Version = Data.getU8(Cur);
       if (!Cur)
         break;
@@ -862,6 +723,10 @@ decodeBBAddrMapCommonImpl(const ELFFile<ELFT> &EF,
         return createError("unsupported SHT_LLVM_BB_ADDR_MAP version: " +
                            Twine(static_cast<int>(Version)));
       Feature = Data.getU8(Cur); // Feature byte
+      FuncEntryCountEnabled =
+          Feature & uint8_t(PGOAnalysisMap::Features::FuncEntryCnt);
+      BBFreqEnabled = Feature & uint8_t(PGOAnalysisMap::Features::BBFreq);
+      BrProbEnabled = Feature & uint8_t(PGOAnalysisMap::Features::BrProb);
     }
     uint64_t SectionOffset = Cur.tell();
     auto Address =
@@ -879,18 +744,81 @@ decodeBBAddrMapCommonImpl(const ELFFile<ELFT> &EF,
       Address = FOTIterator->second;
     }
     uint32_t NumBlocks = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-    auto ExtraData = DT.decodeExtraDataHeader(Version, Feature);
-    std::vector<typename AddrMap::BBEntry> BBEntries;
+
+    if (Feature != 0 && Version < 2 && Cur)
+      return createError("version should be >= 2 for SHT_LLVM_BB_ADDR_MAP when "
+                         "PGO features are enabled: version = " +
+                         Twine(static_cast<int>(Version)) +
+                         " feature = " + Twine(static_cast<int>(Feature)));
+
+    uint64_t FuncEntryCount =
+        FuncEntryCountEnabled ? readULEB128As<uint64_t>(Data, Cur, ULEBSizeErr)
+                              : 0;
+
+    std::vector<BBAddrMap::BBEntry> BBEntries;
+    std::vector<PGOAnalysisMap::PGOBBEntry> PGOBBEntries;
+    uint32_t PrevBBEndOffset = 0;
     for (uint32_t BlockIndex = 0;
          !MetadataDecodeErr && !ULEBSizeErr && Cur && (BlockIndex < NumBlocks);
          ++BlockIndex) {
-      auto Entry = DT.decodeBBEntry(Version, Feature, BlockIndex);
-      if (MetadataDecodeErr)
-        break;
-      BBEntries.push_back(std::move(Entry));
+      uint32_t ID = Version >= 2
+                        ? readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr)
+                        : BlockIndex;
+      uint32_t Offset = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
+      uint32_t Size = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
+      uint32_t MD = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
+      if (Version >= 1) {
+        // Offset is calculated relative to the end of the previous BB.
+        Offset += PrevBBEndOffset;
+        PrevBBEndOffset = Offset + Size;
+      }
+
+      BBAddrMap::BBEntry::Metadata Metadata;
+      PGOAnalysisMap::PGOBBEntry::SuccessorsType ST{};
+      if (BrProbEnabled) {
+        auto MetadataOrErr = PGOAnalysisMap::PGOBBEntry::decodeMD(MD);
+        if (Error E = MetadataOrErr.takeError()) {
+          MetadataDecodeErr = std::move(E);
+          break;
+        }
+        std::tie(Metadata, ST) = *MetadataOrErr;
+      } else {
+        auto MetadataOrErr = BBAddrMap::BBEntry::Metadata::decode(MD);
+        if (Error E = MetadataOrErr.takeError()) {
+          MetadataDecodeErr = std::move(E);
+          break;
+        }
+        Metadata = *MetadataOrErr;
+      }
+
+      uint64_t BBF =
+          BBFreqEnabled ? readULEB128As<uint64_t>(Data, Cur, ULEBSizeErr) : 0;
+
+      llvm::SmallVector<PGOAnalysisMap::PGOBBEntry::SuccessorEntry, 2>
+          Successors;
+      if (BrProbEnabled) {
+        auto SuccCount =
+            ST == PGOAnalysisMap::PGOBBEntry::SuccessorsType::Multiple
+                ? readULEB128As<uint64_t>(Data, Cur, ULEBSizeErr)
+                : uint64_t(ST);
+        for (uint64_t I = 0; I < SuccCount; ++I) {
+          uint32_t BBID = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
+          uint32_t BrProb = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
+          if (PGOAnalyses)
+            Successors.push_back({BBID, BranchProbability::getRaw(BrProb)});
+        }
+      }
+
+      BBEntries.push_back({ID, Offset, Size, Metadata});
+      if (PGOAnalyses)
+        PGOBBEntries.push_back({BlockFrequency(BBF), std::move(Successors)});
     }
-    FunctionEntries.push_back(
-        DT.makeAddrMap(ExtraData, Feature, Address, std::move(BBEntries)));
+
+    if (PGOAnalyses)
+      PGOAnalyses->push_back({FuncEntryCount, std::move(PGOBBEntries),
+                              FuncEntryCountEnabled, BBFreqEnabled,
+                              BrProbEnabled});
+    FunctionEntries.push_back({Address, std::move(BBEntries)});
   }
   // Either Cur is in the error state, or we have an error in ULEBSizeErr or
   // MetadataDecodeErr (but not both), but we join all errors here to be safe.
@@ -902,16 +830,14 @@ decodeBBAddrMapCommonImpl(const ELFFile<ELFT> &EF,
 
 template <class ELFT>
 Expected<std::vector<BBAddrMap>>
-ELFFile<ELFT>::decodeBBAddrMap(const Elf_Shdr &Sec,
-                               const Elf_Shdr *RelaSec) const {
-  return decodeBBAddrMapCommonImpl<BBAddrMap>(*this, Sec, RelaSec);
-}
-
-template <class ELFT>
-Expected<std::vector<PGOBBAddrMap>>
-ELFFile<ELFT>::decodePGOBBAddrMap(const Elf_Shdr &Sec,
-                                  const Elf_Shdr *RelaSec) const {
-  return decodeBBAddrMapCommonImpl<PGOBBAddrMap>(*this, Sec, RelaSec);
+ELFFile<ELFT>::decodeBBAddrMap(const Elf_Shdr &Sec, const Elf_Shdr *RelaSec,
+                               std::vector<PGOAnalysisMap> *PGOAnalyses) const {
+  size_t OriginalPGOSize = PGOAnalyses ? PGOAnalyses->size() : 0;
+  auto AddrMapsOrErr = decodeBBAddrMapImpl(*this, Sec, RelaSec, PGOAnalyses);
+  // remove new analyses when an error occurs
+  if (!AddrMapsOrErr && PGOAnalyses)
+    PGOAnalyses->resize(OriginalPGOSize);
+  return std::move(AddrMapsOrErr);
 }
 
 template <class ELFT>
diff --git a/llvm/lib/Object/ELFObjectFile.cpp b/llvm/lib/Object/ELFObjectFile.cpp
index 0c200f090802269..630141b20a9ae0f 100644
--- a/llvm/lib/Object/ELFObjectFile.cpp
+++ b/llvm/lib/Object/ELFObjectFile.cpp
@@ -33,6 +33,7 @@
 #include <optional>
 #include <string>
 #include <utility>
+#include <vector>
 
 using namespace llvm;
 using namespace object;
@@ -708,18 +709,20 @@ std::vector<ELFPltEntry> ELFObjectFileBase::getPltEntries() const {
   return Result;
 }
 
-template <typename AddrMap, typename ELFT, typename DecodeAddrMapFn>
-Expected<std::vector<AddrMap>> static readBBAddrMapCommonImpl(
+template <class ELFT>
+Expected<std::vector<BBAddrMap>> static readBBAddrMapImpl(
     const ELFFile<ELFT> &EF, std::optional<unsigned> TextSectionIndex,
-    ArrayRef<unsigned> SectionIDs, DecodeAddrMapFn DecodeAddrMap) {
+    std::vector<PGOAnalysisMap> *PGOAnalyses) {
   using Elf_Shdr = typename ELFT::Shdr;
   bool IsRelocatable = EF.getHeader().e_type == ELF::ET_REL;
-  std::vector<AddrMap> BBAddrMaps;
+  std::vector<BBAddrMap> BBAddrMaps;
+  if (PGOAnalyses)
+    PGOAnalyses->clear();
 
   const auto &Sections = cantFail(EF.sections());
   auto IsMatch = [&](const Elf_Shdr &Sec) -> Expected<bool> {
-    if (llvm::none_of(SectionIDs,
-                      [&](unsigned ID) { return ID == Sec.sh_type; }))
+    if (Sec.sh_type != ELF::SHT_LLVM_BB_ADDR_MAP &&
+        Sec.sh_type != ELF::SHT_LLVM_BB_ADDR_MAP_V0)
       return false;
     if (!TextSectionIndex)
       return true;
@@ -742,40 +745,20 @@ Expected<std::vector<AddrMap>> static readBBAddrMapCommonImpl(
     if (IsRelocatable && !RelocSec)
       return createError("unable to get relocation section for " +
                          describe(EF, *Sec));
-    Expected<std::vector<AddrMap>> BBAddrMapOrErr =
-        DecodeAddrMap(EF, *Sec, RelocSec);
-    if (!BBAddrMapOrErr)
+    Expected<std::vector<BBAddrMap>> BBAddrMapOrErr =
+        EF.decodeBBAddrMap(*Sec, RelocSec, PGOAnalyses);
+    if (!BBAddrMapOrErr) {
+      if (PGOAnalyses)
+        PGOAnalyses->clear();
       return createError("unable to read " + describe(EF, *Sec) + ": " +
                          toString(BBAddrMapOrErr.takeError()));
+    }
     std::move(BBAddrMapOrErr->begin(), BBAddrMapOrErr->end(),
               std::back_inserter(BBAddrMaps));
   }
   return BBAddrMaps;
 }
 
-template <class ELFT>
-Expected<std::vector<BBAddrMap>> static readBBAddrMapImpl(
-    const ELFFile<ELFT> &EF, std::optional<unsigned> TextSectionIndex) {
-  return readBBAddrMapCommonImpl<BBAddrMap>(
-      EF, TextSectionIndex,
-      {ELF::SHT_LLVM_BB_ADDR_MAP, ELF::SHT_LLVM_BB_ADDR_MAP_V0},
-      [](const ELFFile<ELFT> &EF, const typename ELFFile<ELFT>::Elf_Shdr &Sec,
-         const typename ELFFile<ELFT>::Elf_Shdr *RelaSec) {
-        return EF.decodeBBAddrMap(Sec, RelaSec);
-      });
-}
-
-template <class ELFT>
-Expected<std::vector<PGOBBAddrMap>> static readPGOBBAddrMapImpl(
-    const ELFFile<ELFT> &EF, std::optional<unsigned> TextSectionIndex) {
-  return readBBAddrMapCommonImpl<PGOBBAddrMap>(
-      EF, TextSectionIndex, {ELF::SHT_LLVM_PGO_BB_ADDR_MAP},
-      [](const ELFFile<ELFT> &EF, const typename ELFFile<ELFT>::Elf_Shdr &Sec,
-         const typename ELFFile<ELFT>::Elf_Shdr *RelaSec) {
-        return EF.decodePGOBBAddrMap(Sec, RelaSec);
-      });
-}
-
 template <class ELFT>
 static Expected<std::vector<VersionEntry>>
 readDynsymVersionsImpl(const ELFFile<ELFT> &EF,
@@ -846,25 +829,14 @@ ELFObjectFileBase::readDynsymVersions() const {
 }
 
 Expected<std::vector<BBAddrMap>> ELFObjectFileBase::readBBAddrMap(
-    std::optional<unsigned> TextSectionIndex) const {
+    std::optional<unsigned> TextSectionIndex,
+    std::vector<PGOAnalysisMap> *PGOAnalyses) const {
   if (const auto *Obj = dyn_cast<ELF32LEObjectFile>(this))
-    return readBBAddrMapImpl(Obj->getELFFile(), TextSectionIndex);
+    return readBBAddrMapImpl(Obj->getELFFile(), TextSectionIndex, PGOAnalyses);
   if (const auto *Obj = dyn_cast<ELF64LEObjectFile>(this))
-    return readBBAddrMapImpl(Obj->getELFFile(), TextSectionIndex);
+    return readBBAddrMapImpl(Obj->getELFFile(), TextSectionIndex, PGOAnalyses);
   if (const auto *Obj = dyn_cast<ELF32BEObjectFile>(this))
-    return readBBAddrMapImpl(Obj->getELFFile(), TextSectionIndex);
+    return readBBAddrMapImpl(Obj->getELFFile(), TextSectionIndex, PGOAnalyses);
   return readBBAddrMapImpl(cast<ELF64BEObjectFile>(this)->getELFFile(),
-                           TextSectionIndex);
-}
-
-Expected<std::vector<PGOBBAddrMap>> ELFObjectFileBase::readPGOBBAddrMap(
-    std::optional<unsigned> TextSectionIndex) const {
-  if (const auto *Obj = dyn_cast<ELF32LEObjectFile>(this))
-    return readPGOBBAddrMapImpl(Obj->getELFFile(), TextSectionIndex);
-  if (const auto *Obj = dyn_cast<ELF64LEObjectFile>(this))
-    return readPGOBBAddrMapImpl(Obj->getELFFile(), TextSectionIndex);
-  if (const auto *Obj = dyn_cast<ELF32BEObjectFile>(this))
-    return readPGOBBAddrMapImpl(Obj->getELFFile(), TextSectionIndex);
-  return readPGOBBAddrMapImpl(cast<ELF64BEObjectFile>(this)->getELFFile(),
-                              TextSectionIndex);
+                           TextSectionIndex, PGOAnalyses);
 }
diff --git a/llvm/lib/ObjectYAML/ELFEmitter.cpp b/llvm/lib/ObjectYAML/ELFEmitter.cpp
index 4265c4bfd226eac..7df4ca786e21b1d 100644
--- a/llvm/lib/ObjectYAML/ELFEmitter.cpp
+++ b/llvm/lib/ObjectYAML/ELFEmitter.cpp
@@ -32,6 +32,7 @@
 #include "llvm/Support/YAMLTraits.h"
 #include "llvm/Support/raw_ostream.h"
 #include <optional>
+#include <variant>
 
 using namespace llvm;
 
@@ -284,9 +285,6 @@ template <class ELFT> class ELFState {
   void writeSectionContent(Elf_Shdr &SHeader,
                            const ELFYAML::BBAddrMapSection &Section,
                            ContiguousBlobAccumulator &CBA);
-  void writeSectionContent(Elf_Shdr &SHeader,
-                           const ELFYAML::PGOBBAddrMapSection &Section,
-                           ContiguousBlobAccumulator &CBA);
   void writeSectionContent(Elf_Shdr &SHeader,
                            const ELFYAML::HashSection &Section,
                            ContiguousBlobAccumulator &CBA);
@@ -898,8 +896,6 @@ void ELFState<ELFT>::initSectionHeaders(std::vector<Elf_Shdr> &SHeaders,
       writeSectionContent(SHeader, *S, CBA);
     } else if (auto S = dyn_cast<ELFYAML::BBAddrMapSection>(Sec)) {
       writeSectionContent(SHeader, *S, CBA);
-    } else if (auto S = dyn_cast<ELFYAML::PGOBBAddrMapSection>(Sec)) {
-      writeSectionContent(SHeader, *S, CBA);
     } else {
       llvm_unreachable("Unknown section type");
     }
@@ -1392,124 +1388,96 @@ void ELFState<ELFT>::writeSectionContent(
   }
 }
 
-namespace {
-template <typename ELFT, typename AddrMap> struct BBAddrMapWriteTrait;
-
-template <typename ELFT>
-struct BBAddrMapWriteTrait<ELFT, ELFYAML::BBAddrMapSection> {
-  typename ELFState<ELFT>::Elf_Shdr &SHeader;
-  const ELFYAML::Section &Section;
-  ContiguousBlobAccumulator &CBA;
-  const unsigned UpToDateSectionID;
-
-  BBAddrMapWriteTrait(typename ELFState<ELFT>::Elf_Shdr &SHeader,
-                      const ELFYAML::Section &Section,
-                      ContiguousBlobAccumulator &CBA,
-                      unsigned UpToDateSectionID = ELF::SHT_LLVM_BB_ADDR_MAP)
-      : SHeader(SHeader), Section(Section), CBA(CBA),
-        UpToDateSectionID(UpToDateSectionID) {}
-
-  template <typename EntryTy> void writeEntry(const EntryTy &E) {
-    using uintX_t = typename ELFState<ELFT>::uintX_t;
+template <class ELFT>
+void ELFState<ELFT>::writeSectionContent(
+    Elf_Shdr &SHeader, const ELFYAML::BBAddrMapSection &Section,
+    ContiguousBlobAccumulator &CBA) {
+  if (!Section.Entries) {
+    if (Section.PGOAnalyses)
+      WithColor::warning()
+          << "PGOAnalyses should not exist in SHT_LLVM_BB_ADDR_MAP when "
+             "Entries does not exist";
+    return;
+  }
+
+  const std::vector<ELFYAML::PGOAnalysisMapEntry> *PGOAnalyses = nullptr;
+  if (Section.PGOAnalyses) {
+    if (Section.Entries->size() != Section.PGOAnalyses->size())
+      WithColor::warning() << "PGOAnalyses must be the same length as Entries "
+                              "in SHT_LLVM_BB_ADDR_MAP";
+    else
+      PGOAnalyses = &Section.PGOAnalyses.value();
+  }
+
+  for (const auto &[Idx, E] : llvm::enumerate(*Section.Entries)) {
+    const ELFYAML::PGOAnalysisMapEntry *PGOEntry =
+        PGOAnalyses ? &PGOAnalyses->at(Idx) : nullptr;
     // Write version and feature values.
-    if (Section.Type == UpToDateSectionID) {
-      if (E.Common.Version > 2)
+    if (Section.Type == llvm::ELF::SHT_LLVM_BB_ADDR_MAP) {
+      if (E.Version > 2)
         WithColor::warning() << "unsupported SHT_LLVM_BB_ADDR_MAP version: "
-                             << static_cast<int>(E.Common.Version)
+                             << static_cast<int>(E.Version)
                              << "; encoding using the most recent version";
-      CBA.write(E.Common.Version);
-      CBA.write(E.Common.Feature);
+      CBA.write(E.Version);
+      CBA.write(E.Feature);
       SHeader.sh_size += 2;
     }
+
+    if (Section.PGOAnalyses) {
+      if (E.Version < 2)
+        WithColor::warning()
+            << "unsupported SHT_LLVM_BB_ADDR_MAP version when using PGO: "
+            << static_cast<int>(E.Version) << "; must use version >= 2";
+    }
+
     // Write the address of the function.
-    CBA.write<uintX_t>(E.Common.Address, ELFT::TargetEndianness);
+    CBA.write<uintX_t>(E.Address, ELFT::TargetEndianness);
     // Write number of BBEntries (number of basic blocks in the function). This
     // is overridden by the 'NumBlocks' YAML field when specified.
     uint64_t NumBlocks =
-        E.Common.NumBlocks.value_or(E.BBEntries ? E.BBEntries->size() : 0);
+        E.NumBlocks.value_or(E.BBEntries ? E.BBEntries->size() : 0);
     SHeader.sh_size += sizeof(uintX_t) + CBA.writeULEB128(NumBlocks);
-  }
-
-  void writeBBEntry(const ELFYAML::BBAddrMapEntry &E,
-                    const ELFYAML::BBAddrMapEntry::BBEntry &BBE) {
-    if (Section.Type == UpToDateSectionID && E.Common.Version > 1)
-      SHeader.sh_size += CBA.writeULEB128(BBE.ID);
-    SHeader.sh_size += CBA.writeULEB128(BBE.AddressOffset) +
-                       CBA.writeULEB128(BBE.Size) +
-                       CBA.writeULEB128(BBE.Metadata);
-  }
-};
 
-template <typename ELFT>
-struct BBAddrMapWriteTrait<ELFT, ELFYAML::PGOBBAddrMapSection> {
-  using Features = object::PGOBBAddrMap::Features;
-
-  BBAddrMapWriteTrait<ELFT, ELFYAML::BBAddrMapSection> Base;
-
-  BBAddrMapWriteTrait(typename ELFState<ELFT>::Elf_Shdr &SHeader,
-                      const ELFYAML::PGOBBAddrMapSection &Section,
-                      ContiguousBlobAccumulator &CBA)
-      : Base(SHeader, Section, CBA, ELF::SHT_LLVM_PGO_BB_ADDR_MAP) {}
-
-  void writeEntry(const ELFYAML::PGOBBAddrMapEntry &E) {
-    if (E.Common.Version < 2)
-      WithColor::warning() << "unsupported SHT_LLVM_PGO_BB_ADDR_MAP version: "
-                           << static_cast<int>(E.Common.Version)
-                           << "; must use version >= 2";
-    Base.writeEntry(E);
-    if (E.FuncEntryCount)
-      Base.SHeader.sh_size += Base.CBA.writeULEB128(*E.FuncEntryCount);
-  }
-
-  void writeBBEntry(const ELFYAML::PGOBBAddrMapEntry &E,
-                    const ELFYAML::PGOBBAddrMapEntry::BBEntry &BBE) {
-    Base.writeBBEntry(ELFYAML::BBAddrMapEntry{E.Common, {}}, BBE.Base);
-    if (BBE.BBFreq)
-      Base.SHeader.sh_size += Base.CBA.writeULEB128(*BBE.BBFreq);
-    if (BBE.Successors) {
-      if (BBE.Successors->size() > 2)
-        Base.SHeader.sh_size += Base.CBA.writeULEB128(BBE.Successors->size());
-      for (const auto &Succ : *BBE.Successors)
-        Base.SHeader.sh_size +=
-            Base.CBA.writeULEB128(Succ.ID) + Base.CBA.writeULEB128(Succ.BrProb);
+    if (PGOEntry && PGOEntry->FuncEntryCount)
+      SHeader.sh_size += CBA.writeULEB128(*PGOEntry->FuncEntryCount);
+
+    const auto *PGOBBEntries = PGOEntry && PGOEntry->PGOBBEntries
+                                   ? &PGOEntry->PGOBBEntries.value()
+                                   : nullptr;
+    if (PGOBBEntries && E.BBEntries &&
+        E.BBEntries->size() != PGOBBEntries->size()) {
+      PGOBBEntries = nullptr;
+      WithColor::warning() << "PBOBBEntries must be the same length as "
+                              "BBEntries in SHT_LLVM_BB_ADDR_MAP.\n"
+                           << "Mismatch on function with address: "
+                           << E.Address;
     }
-  }
-};
-} // namespace
 
-template <typename ELFT, typename SectionTy>
-static void writeAddrMapSectionContent(
-    ELFState<ELFT> &ES, typename ELFState<ELFT>::Elf_Shdr &SHeader,
-    const SectionTy &Section, ContiguousBlobAccumulator &CBA) {
-  if (!Section.Entries)
-    return;
-
-  BBAddrMapWriteTrait<ELFT, SectionTy> AddrTrait(SHeader, Section, CBA);
-  for (const auto &E : *Section.Entries) {
-    AddrTrait.writeEntry(E);
     // Write all BBEntries.
     if (!E.BBEntries)
       continue;
-    for (const auto &BBE : *E.BBEntries) {
-      AddrTrait.writeBBEntry(E, BBE);
+    for (const auto &[BBIdx, BBE] : llvm::enumerate(*E.BBEntries)) {
+      const auto *PGOBBE = PGOBBEntries ? &PGOBBEntries->at(BBIdx) : nullptr;
+      if (Section.Type == llvm::ELF::SHT_LLVM_BB_ADDR_MAP && E.Version > 1)
+        SHeader.sh_size += CBA.writeULEB128(BBE.ID);
+      SHeader.sh_size += CBA.writeULEB128(BBE.AddressOffset) +
+                         CBA.writeULEB128(BBE.Size) +
+                         CBA.writeULEB128(BBE.Metadata);
+
+      if (PGOBBE) {
+        if (PGOBBE->BBFreq)
+          SHeader.sh_size += CBA.writeULEB128(*PGOBBE->BBFreq);
+        if (PGOBBE->Successors) {
+          if (PGOBBE->Successors->size() > 2)
+            SHeader.sh_size += CBA.writeULEB128(PGOBBE->Successors->size());
+          for (const auto &[ID, BrProb] : *PGOBBE->Successors)
+            SHeader.sh_size += CBA.writeULEB128(ID) + CBA.writeULEB128(BrProb);
+        }
+      }
     }
   }
 }
 
-template <class ELFT>
-void ELFState<ELFT>::writeSectionContent(
-    Elf_Shdr &SHeader, const ELFYAML::BBAddrMapSection &Section,
-    ContiguousBlobAccumulator &CBA) {
-  writeAddrMapSectionContent(*this, SHeader, Section, CBA);
-}
-
-template <class ELFT>
-void ELFState<ELFT>::writeSectionContent(
-    Elf_Shdr &SHeader, const ELFYAML::PGOBBAddrMapSection &Section,
-    ContiguousBlobAccumulator &CBA) {
-  writeAddrMapSectionContent(*this, SHeader, Section, CBA);
-}
-
 template <class ELFT>
 void ELFState<ELFT>::writeSectionContent(
     Elf_Shdr &SHeader, const ELFYAML::LinkerOptionsSection &Section,
diff --git a/llvm/lib/ObjectYAML/ELFYAML.cpp b/llvm/lib/ObjectYAML/ELFYAML.cpp
index 8aa79b6724e0a07..3783dbdaa4a5c21 100644
--- a/llvm/lib/ObjectYAML/ELFYAML.cpp
+++ b/llvm/lib/ObjectYAML/ELFYAML.cpp
@@ -683,7 +683,6 @@ void ScalarEnumerationTraits<ELFYAML::ELF_SHT>::enumeration(
   ECase(SHT_LLVM_PART_PHDR);
   ECase(SHT_LLVM_BB_ADDR_MAP_V0);
   ECase(SHT_LLVM_BB_ADDR_MAP);
-  ECase(SHT_LLVM_PGO_BB_ADDR_MAP);
   ECase(SHT_LLVM_OFFLOADING);
   ECase(SHT_LLVM_LTO);
   ECase(SHT_GNU_ATTRIBUTES);
@@ -1388,12 +1387,7 @@ static void sectionMapping(IO &IO, ELFYAML::BBAddrMapSection &Section) {
   commonSectionMapping(IO, Section);
   IO.mapOptional("Content", Section.Content);
   IO.mapOptional("Entries", Section.Entries);
-}
-
-static void sectionMapping(IO &IO, ELFYAML::PGOBBAddrMapSection &Section) {
-  commonSectionMapping(IO, Section);
-  IO.mapOptional("Content", Section.Content);
-  IO.mapOptional("Entries", Section.Entries);
+  IO.mapOptional("PGOAnalyses", Section.PGOAnalyses);
 }
 
 static void sectionMapping(IO &IO, ELFYAML::StackSizesSection &Section) {
@@ -1689,11 +1683,6 @@ void MappingTraits<std::unique_ptr<ELFYAML::Chunk>>::mapping(
       Section.reset(new ELFYAML::BBAddrMapSection());
     sectionMapping(IO, *cast<ELFYAML::BBAddrMapSection>(Section.get()));
     break;
-  case ELF::SHT_LLVM_PGO_BB_ADDR_MAP:
-    if (!IO.outputting())
-      Section.reset(new ELFYAML::PGOBBAddrMapSection());
-    sectionMapping(IO, *cast<ELFYAML::PGOBBAddrMapSection>(Section.get()));
-    break;
   default:
     if (!IO.outputting()) {
       StringRef Name;
@@ -1815,17 +1804,13 @@ void MappingTraits<ELFYAML::StackSizeEntry>::mapping(
   IO.mapRequired("Size", E.Size);
 }
 
-static void mapBBAddrMapCommonBase(IO &IO, ELFYAML::BBAddrMapCommonBase &E) {
+void MappingTraits<ELFYAML::BBAddrMapEntry>::mapping(
+    IO &IO, ELFYAML::BBAddrMapEntry &E) {
+  assert(IO.getContext() && "The IO context is not initialized");
   IO.mapRequired("Version", E.Version);
   IO.mapOptional("Feature", E.Feature, Hex8(0));
   IO.mapOptional("Address", E.Address, Hex64(0));
   IO.mapOptional("NumBlocks", E.NumBlocks);
-}
-
-void MappingTraits<ELFYAML::BBAddrMapEntry>::mapping(
-    IO &IO, ELFYAML::BBAddrMapEntry &E) {
-  assert(IO.getContext() && "The IO context is not initialized");
-  mapBBAddrMapCommonBase(IO, E.Common);
   IO.mapOptional("BBEntries", E.BBEntries);
 }
 
@@ -1838,24 +1823,23 @@ void MappingTraits<ELFYAML::BBAddrMapEntry::BBEntry>::mapping(
   IO.mapRequired("Metadata", E.Metadata);
 }
 
-void MappingTraits<ELFYAML::PGOBBAddrMapEntry>::mapping(
-    IO &IO, ELFYAML::PGOBBAddrMapEntry &E) {
+void MappingTraits<ELFYAML::PGOAnalysisMapEntry>::mapping(
+    IO &IO, ELFYAML::PGOAnalysisMapEntry &E) {
   assert(IO.getContext() && "The IO context is not initialized");
-  mapBBAddrMapCommonBase(IO, E.Common);
   IO.mapOptional("FuncEntryCount", E.FuncEntryCount);
-  IO.mapOptional("BBEntries", E.BBEntries);
+  IO.mapOptional("PGOBBEntries", E.PGOBBEntries);
 }
 
-void MappingTraits<ELFYAML::PGOBBAddrMapEntry::BBEntry>::mapping(
-    IO &IO, ELFYAML::PGOBBAddrMapEntry::BBEntry &E) {
+void MappingTraits<ELFYAML::PGOAnalysisMapEntry::PGOBBEntry>::mapping(
+    IO &IO, ELFYAML::PGOAnalysisMapEntry::PGOBBEntry &E) {
   assert(IO.getContext() && "The IO context is not initialized");
-  MappingTraits<ELFYAML::BBAddrMapEntry::BBEntry>::mapping(IO, E.Base);
   IO.mapOptional("BBFreq", E.BBFreq);
   IO.mapOptional("Successors", E.Successors);
 }
 
-void MappingTraits<ELFYAML::PGOBBAddrMapEntry::BBEntry::SuccessorEntry>::
-    mapping(IO &IO, ELFYAML::PGOBBAddrMapEntry::BBEntry::SuccessorEntry &E) {
+void MappingTraits<ELFYAML::PGOAnalysisMapEntry::PGOBBEntry::SuccessorEntry>::
+    mapping(IO &IO,
+            ELFYAML::PGOAnalysisMapEntry::PGOBBEntry::SuccessorEntry &E) {
   assert(IO.getContext() && "The IO context is not initialized");
   IO.mapRequired("ID", E.ID);
   IO.mapRequired("BrProb", E.BrProb);
diff --git a/llvm/test/MC/AsmParser/llvm_section_types.s b/llvm/test/MC/AsmParser/llvm_section_types.s
index 2d7b650e10432ae..147b1499d2b8883 100644
--- a/llvm/test/MC/AsmParser/llvm_section_types.s
+++ b/llvm/test/MC/AsmParser/llvm_section_types.s
@@ -17,8 +17,6 @@
 .byte 1
 .section    .section8,"", at llvm_lto
 .byte 1
-.section    .section9,"", at llvm_pgo_bb_addr_map
-.byte 1
 
 # CHECK:        Name: .section1
 # CHECK-NEXT:   Type: SHT_LLVM_BB_ADDR_MAP
@@ -36,5 +34,3 @@
 # CHECK-NEXT:   Type: SHT_LLVM_OFFLOADING
 # CHECK:        Name: .section8
 # CHECK-NEXT:   Type: SHT_LLVM_LTO
-# CHECK:        Name: .section9
-# CHECK-NEXT:   Type: SHT_LLVM_PGO_BB_ADDR_MAP
diff --git a/llvm/tools/obj2yaml/elf2yaml.cpp b/llvm/tools/obj2yaml/elf2yaml.cpp
index 653e97959795014..bec4a10a5d35f2a 100644
--- a/llvm/tools/obj2yaml/elf2yaml.cpp
+++ b/llvm/tools/obj2yaml/elf2yaml.cpp
@@ -915,7 +915,7 @@ ELFDumper<ELFT>::dumpBBAddrMapSection(const Elf_Shdr *Shdr) {
       BBEntries.push_back({ID, Offset, Size, Metadata});
     }
     Entries.push_back(
-        {{Version, Feature, Address, /*NumBlocks=*/{}}, std::move(BBEntries)});
+        {Version, Feature, Address, /*NumBlocks=*/{}, std::move(BBEntries)});
   }
 
   if (!Cur) {
diff --git a/llvm/unittests/Object/ELFObjectFileTest.cpp b/llvm/unittests/Object/ELFObjectFileTest.cpp
index 7ae9a2055040bdb..da8af2f69195dea 100644
--- a/llvm/unittests/Object/ELFObjectFileTest.cpp
+++ b/llvm/unittests/Object/ELFObjectFileTest.cpp
@@ -9,6 +9,7 @@
 #include "llvm/Object/ELFObjectFile.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ObjectYAML/yaml2obj.h"
+#include "llvm/Support/BlockFrequency.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/YAMLTraits.h"
 #include "llvm/Testing/Support/Error.h"
@@ -744,7 +745,7 @@ TEST(ELFObjectFileTest, ReadBBAddrMap) {
                   Section1BBAddrMaps);
 }
 
-// Tests for error paths of the ELFFile::decodePGOBBAddrMap API.
+// Tests for error paths of the ELFFile::decodeBBAddrMap with PGOBBAddrMap API.
 TEST(ELFObjectFileTest, InvalidDecodePGOBBAddrMap) {
   if (IsHostWindows())
     GTEST_SKIP();
@@ -755,8 +756,8 @@ TEST(ELFObjectFileTest, InvalidDecodePGOBBAddrMap) {
   Data:  ELFDATA2LSB
   Type:  ET_EXEC
 Sections:
-  - Type: SHT_LLVM_PGO_BB_ADDR_MAP
-    Name: .llvm_pgo_bb_addr_map
+  - Type: SHT_LLVM_BB_ADDR_MAP
+    Name: .llvm_bb_addr_map
     Entries:
       - Address: 0x11111
 )");
@@ -768,31 +769,22 @@ TEST(ELFObjectFileTest, InvalidDecodePGOBBAddrMap) {
     ASSERT_THAT_EXPECTED(ElfOrErr, Succeeded());
     const ELFFile<ELF64LE> &Elf = ElfOrErr->getELFFile();
 
-    Expected<const typename ELF64LE::Shdr *> PGOBBAddrMapSecOrErr =
+    Expected<const typename ELF64LE::Shdr *> BBAddrMapSecOrErr =
         Elf.getSection(1);
-    ASSERT_THAT_EXPECTED(PGOBBAddrMapSecOrErr, Succeeded());
+    ASSERT_THAT_EXPECTED(BBAddrMapSecOrErr, Succeeded());
+
+    std::vector<PGOAnalysisMap> PGOAnalyses;
     EXPECT_THAT_ERROR(
-        Elf.decodePGOBBAddrMap(**PGOBBAddrMapSecOrErr).takeError(),
+        Elf.decodeBBAddrMap(**BBAddrMapSecOrErr, nullptr, &PGOAnalyses)
+            .takeError(),
         FailedWithMessage(ErrMsg));
   };
 
-  // Check that we can detect unsupported versions.
-  SmallString<128> UnsupportedVersionYamlString(CommonYamlString);
-  UnsupportedVersionYamlString += R"(
-        Version: 3
-        BBEntries:
-          - AddressOffset: 0x0
-            Size:          0x1
-            Metadata:      0x2
-)";
-
-  DoCheck(UnsupportedVersionYamlString,
-          "unsupported SHT_LLVM_BB_ADDR_MAP version: 3");
-
   // Check that we can detect unsupported versions that is too low
   SmallString<128> UnsupportedLowVersionYamlString(CommonYamlString);
   UnsupportedLowVersionYamlString += R"(
         Version: 1
+        Feature: 0x4
         BBEntries:
           - AddressOffset: 0x0
             Size:          0x1
@@ -800,7 +792,8 @@ TEST(ELFObjectFileTest, InvalidDecodePGOBBAddrMap) {
 )";
 
   DoCheck(UnsupportedLowVersionYamlString,
-          "unsupported SHT_LLVM_PGO_BB_ADDR_MAP version: 1");
+          "version should be >= 2 for SHT_LLVM_BB_ADDR_MAP when PGO features "
+          "are enabled: version = 1 feature = 4");
 
   SmallString<128> CommonVersionedYamlString(CommonYamlString);
   CommonVersionedYamlString += R"(
@@ -812,83 +805,6 @@ TEST(ELFObjectFileTest, InvalidDecodePGOBBAddrMap) {
             Metadata:      0x2
 )";
 
-  // Check that we can detect the malformed encoding when the section is
-  // truncated.
-  SmallString<128> TruncatedYamlString(CommonVersionedYamlString);
-  TruncatedYamlString += R"(
-    ShSize: 0xb
-)";
-  DoCheck(TruncatedYamlString, "unable to decode LEB128 at offset 0x0000000b: "
-                               "malformed uleb128, extends past end");
-
-  // Check that we can detect when the encoded BB entry fields exceed the UINT32
-  // limit.
-  SmallVector<SmallString<128>, 3> OverInt32LimitYamlStrings(
-      3, CommonVersionedYamlString);
-  OverInt32LimitYamlStrings[0] += R"(
-          - ID:            1
-            AddressOffset: 0x100000000
-            Size:          0xFFFFFFFF
-            Metadata:      0xFFFFFFFF
-)";
-
-  OverInt32LimitYamlStrings[1] += R"(
-          - ID:            2
-            AddressOffset: 0xFFFFFFFF
-            Size:          0x100000000
-            Metadata:      0xFFFFFFFF
-)";
-
-  OverInt32LimitYamlStrings[2] += R"(
-          - ID:            3
-            AddressOffset: 0xFFFFFFFF
-            Size:          0xFFFFFFFF
-            Metadata:      0x100000000
-)";
-
-  DoCheck(OverInt32LimitYamlStrings[0],
-          "ULEB128 value at offset 0x10 exceeds UINT32_MAX (0x100000000)");
-  DoCheck(OverInt32LimitYamlStrings[1],
-          "ULEB128 value at offset 0x15 exceeds UINT32_MAX (0x100000000)");
-  DoCheck(OverInt32LimitYamlStrings[2],
-          "ULEB128 value at offset 0x1a exceeds UINT32_MAX (0x100000000)");
-
-  // Check the proper error handling when the section has fields exceeding
-  // UINT32 and is also truncated. This is for checking that we don't generate
-  // unhandled errors.
-  SmallVector<SmallString<128>, 3> OverInt32LimitAndTruncated(
-      3, OverInt32LimitYamlStrings[1]);
-  // Truncate before the end of the 5-byte field.
-  OverInt32LimitAndTruncated[0] += R"(
-    ShSize: 0x19
-)";
-  // Truncate at the end of the 5-byte field.
-  OverInt32LimitAndTruncated[1] += R"(
-    ShSize: 0x1a
-)";
-  // Truncate after the end of the 5-byte field.
-  OverInt32LimitAndTruncated[2] += R"(
-    ShSize: 0x1b
-)";
-
-  DoCheck(OverInt32LimitAndTruncated[0],
-          "unable to decode LEB128 at offset 0x00000015: malformed uleb128, "
-          "extends past end");
-  DoCheck(OverInt32LimitAndTruncated[1],
-          "ULEB128 value at offset 0x15 exceeds UINT32_MAX (0x100000000)");
-  DoCheck(OverInt32LimitAndTruncated[2],
-          "ULEB128 value at offset 0x15 exceeds UINT32_MAX (0x100000000)");
-
-  // Check for proper error handling when the 'NumBlocks' field is overridden
-  // with an out-of-range value.
-  SmallString<128> OverLimitNumBlocks(CommonVersionedYamlString);
-  OverLimitNumBlocks += R"(
-        NumBlocks: 0x100000000
-)";
-
-  DoCheck(OverLimitNumBlocks,
-          "ULEB128 value at offset 0xa exceeds UINT32_MAX (0x100000000)");
-
   // Check that we fail when function entry count is enabled but not provided.
   SmallString<128> MissingFuncEntryCount(CommonYamlString);
   MissingFuncEntryCount += R"(
@@ -924,7 +840,7 @@ TEST(ELFObjectFileTest, InvalidDecodePGOBBAddrMap) {
           - ID:            1
             AddressOffset: 0x0
             Size:          0x1
-            Metadata:      0x60
+            Metadata:      0x6
           - ID:            2
             AddressOffset: 0x1
             Size:          0x1
@@ -954,21 +870,22 @@ TEST(ELFObjectFileTest, ReadPGOBBAddrMap) {
   Data:  ELFDATA2LSB
   Type:  ET_EXEC
 Sections:
-  - Name: .llvm_pgo_bb_addr_map_1
-    Type: SHT_LLVM_PGO_BB_ADDR_MAP
+  - Name: .llvm_bb_addr_map_1
+    Type: SHT_LLVM_BB_ADDR_MAP
     Link: 1
     Entries:
       - Version: 2
         Address: 0x11111
         Feature: 0x1
-        FuncEntryCount: 892
         BBEntries:
           - ID:            1
             AddressOffset: 0x0
             Size:          0x1
             Metadata:      0x2
-  - Name: .llvm_pgo_bb_addr_map_2
-    Type: SHT_LLVM_PGO_BB_ADDR_MAP
+    PGOAnalyses:
+      - FuncEntryCount: 892
+  - Name: .llvm_bb_addr_map_2
+    Type: SHT_LLVM_BB_ADDR_MAP
     Link: 1
     Entries:
       - Version: 2
@@ -979,9 +896,11 @@ TEST(ELFObjectFileTest, ReadPGOBBAddrMap) {
             AddressOffset: 0x0
             Size:          0x2
             Metadata:      0x4
-            BBFreq:        343
+    PGOAnalyses:
+      - PGOBBEntries:
+         - BBFreq:         343
   - Name: .llvm_bb_addr_map_3
-    Type: SHT_LLVM_PGO_BB_ADDR_MAP
+    Type: SHT_LLVM_BB_ADDR_MAP
     Link: 2
     Entries:
       - Version: 2
@@ -992,37 +911,53 @@ TEST(ELFObjectFileTest, ReadPGOBBAddrMap) {
             AddressOffset: 0x0
             Size:          0x3
             Metadata:      0x46
-            Successors:
-            - ID:          1
-              BrProb:      0x11111111
-            - ID:          2
-              BrProb:      0xeeeeeeee
           - ID:            1
             AddressOffset: 0x0
             Size:          0x3
             Metadata:      0x24
-            Successors:
-            - ID:          2
-              BrProb:      0xffffffff
           - ID:            2
             AddressOffset: 0x0
             Size:          0x3
             Metadata:      0x00
-            Successors:    []
-  - Name: .llvm_pgo_bb_addr_map_4
-    Type: SHT_LLVM_PGO_BB_ADDR_MAP
+    PGOAnalyses:
+      - PGOBBEntries:
+         - Successors:
+            - ID:          1
+              BrProb:      0x11111111
+            - ID:          2
+              BrProb:      0xeeeeeeee
+         - Successors:
+            - ID:          2
+              BrProb:      0xffffffff
+         - Successors:     []
+  - Name: .llvm_bb_addr_map_4
+    Type: SHT_LLVM_BB_ADDR_MAP
   # Link: 0 (by default, can be overriden)
     Entries:
       - Version: 2
         Address: 0x44444
         Feature: 0x7
-        FuncEntryCount: 1000
         BBEntries:
           - ID:            0
             AddressOffset: 0x0
             Size:          0x4
             Metadata:      0x78
-            BBFreq:        1000
+          - ID:            1
+            AddressOffset: 0x0
+            Size:          0x4
+            Metadata:      0x40
+          - ID:            2
+            AddressOffset: 0x0
+            Size:          0x4
+            Metadata:      0x20
+          - ID:            3
+            AddressOffset: 0x0
+            Size:          0x4
+            Metadata:      0x00
+    PGOAnalyses:
+      - FuncEntryCount: 1000
+        PGOBBEntries:
+          - BBFreq:         1000
             Successors:
             - ID:          1
               BrProb:      0x22222222
@@ -1030,99 +965,90 @@ TEST(ELFObjectFileTest, ReadPGOBBAddrMap) {
               BrProb:      0x33333333
             - ID:          3
               BrProb:      0xaaaaaaaa
-          - ID:            1
-            AddressOffset: 0x0
-            Size:          0x4
-            Metadata:      0x40
-            BBFreq:        133
+          - BBFreq:         133
             Successors:
             - ID:          2
               BrProb:      0x11111111
             - ID:          3
               BrProb:      0xeeeeeeee
-          - ID:            2
-            AddressOffset: 0x0
-            Size:          0x4
-            Metadata:      0x20
-            BBFreq:        18
+          - BBFreq:         18
             Successors:
             - ID:          3
               BrProb:      0xffffffff
-          - ID:            3
-            AddressOffset: 0x0
-            Size:          0x4
-            Metadata:      0x00
-            BBFreq:        1000
+          - BBFreq:         1000
             Successors:    []
 )");
 
-  PGOBBAddrMap E1 = {
-      0x11111, {{{1, 0x0, 0x1, {false, true, false, false, false}}, {}, {}}},
-      892,     true,
-      false,   false};
-  PGOBBAddrMap E2 = {0x22222,
-                     {{{2, 0x0, 0x2, {false, false, true, false, false}},
-                       BlockFrequency(343),
-                       {}}},
-                     0,
-                     false,
-                     true,
-                     false};
-  PGOBBAddrMap E3 = {
-      0x33333,
-      {{{0, 0x0, 0x3, {false, true, true, false, false}},
-        {},
-        {{1, BranchProbability::getRaw(0x1111'1111)},
-         {2, BranchProbability::getRaw(0xeeee'eeee)}}},
-       {{1, 0x3, 0x3, {false, false, true, false, false}},
-        {},
-        {{2, BranchProbability::getRaw(0xffff'ffff)}}},
-       {{2, 0x6, 0x3, {false, false, false, false, false}}, {}, {}}},
-      0,
-      false,
-      false,
+  BBAddrMap E1 = {0x11111, {{1, 0x0, 0x1, {false, true, false, false, false}}}};
+  PGOAnalysisMap P1 = {892, {{}}, true, false, false};
+  BBAddrMap E2 = {0x22222, {{2, 0x0, 0x2, {false, false, true, false, false}}}};
+  PGOAnalysisMap P2 = {{}, {{BlockFrequency(343), {}}}, false, true, false};
+  BBAddrMap E3 = {0x33333,
+                  {{0, 0x0, 0x3, {false, true, true, false, false}},
+                   {1, 0x3, 0x3, {false, false, true, false, false}},
+                   {2, 0x6, 0x3, {false, false, false, false, false}}}};
+  PGOAnalysisMap P3 = {{},
+                       {{{},
+                         {{1, BranchProbability::getRaw(0x1111'1111)},
+                          {2, BranchProbability::getRaw(0xeeee'eeee)}}},
+                        {{}, {{2, BranchProbability::getRaw(0xffff'ffff)}}},
+                        {{}, {}}},
+                       false,
+                       false,
+                       true};
+  BBAddrMap E4 = {0x44444,
+                  {{0, 0x0, 0x4, {false, false, false, true, true}},
+                   {1, 0x4, 0x4, {false, false, false, false, false}},
+                   {2, 0x8, 0x4, {false, false, false, false, false}},
+                   {3, 0xc, 0x4, {false, false, false, false, false}}}};
+  PGOAnalysisMap P4 = {
+      1000,
+      {{BlockFrequency(1000),
+        {{1, BranchProbability::getRaw(0x2222'2222)},
+         {2, BranchProbability::getRaw(0x3333'3333)},
+         {3, BranchProbability::getRaw(0xaaaa'aaaa)}}},
+       {BlockFrequency(133),
+        {{2, BranchProbability::getRaw(0x1111'1111)},
+         {3, BranchProbability::getRaw(0xeeee'eeee)}}},
+       {BlockFrequency(18), {{3, BranchProbability::getRaw(0xffff'ffff)}}},
+       {BlockFrequency(1000), {}}},
+      true,
+      true,
       true};
-  PGOBBAddrMap E4 = {0x44444,
-                     {{{0, 0x0, 0x4, {false, false, false, true, true}},
-                       BlockFrequency(1000),
-                       {{1, BranchProbability::getRaw(0x2222'2222)},
-                        {2, BranchProbability::getRaw(0x3333'3333)},
-                        {3, BranchProbability::getRaw(0xaaaa'aaaa)}}},
-                      {{1, 0x4, 0x4, {false, false, false, false, false}},
-                       BlockFrequency(133),
-                       {{2, BranchProbability::getRaw(0x1111'1111)},
-                        {3, BranchProbability::getRaw(0xeeee'eeee)}}},
-                      {{2, 0x8, 0x4, {false, false, false, false, false}},
-                       BlockFrequency(18),
-                       {{3, BranchProbability::getRaw(0xffff'ffff)}}},
-                      {{3, 0xc, 0x4, {false, false, false, false, false}},
-                       BlockFrequency(1000),
-                       {}}},
-                     1000,
-                     true,
-                     true,
-                     true};
-
-  std::vector<PGOBBAddrMap> Section0PGOBBAddrMaps = {E4};
-  std::vector<PGOBBAddrMap> Section1PGOBBAddrMaps = {E3};
-  std::vector<PGOBBAddrMap> Section2PGOBBAddrMaps = {E1, E2};
-  std::vector<PGOBBAddrMap> AllPGOBBAddrMaps = {E1, E2, E3, E4};
 
-  auto DoCheckSucceeds = [&](StringRef YamlString,
-                             std::optional<unsigned> TextSectionIndex,
-                             std::vector<PGOBBAddrMap> ExpectedResult) {
-    SmallString<0> Storage;
-    Expected<ELFObjectFile<ELF64LE>> ElfOrErr =
-        toBinary<ELF64LE>(Storage, YamlString);
-    ASSERT_THAT_EXPECTED(ElfOrErr, Succeeded());
+  std::vector<BBAddrMap> Section0BBAddrMaps = {E4};
+  std::vector<BBAddrMap> Section1BBAddrMaps = {E3};
+  std::vector<BBAddrMap> Section2BBAddrMaps = {E1, E2};
+  std::vector<BBAddrMap> AllBBAddrMaps = {E1, E2, E3, E4};
 
-    Expected<const typename ELF64LE::Shdr *> PGOBBAddrMapSecOrErr =
-        ElfOrErr->getELFFile().getSection(1);
-    ASSERT_THAT_EXPECTED(PGOBBAddrMapSecOrErr, Succeeded());
-    auto PGOBBAddrMaps = ElfOrErr->readPGOBBAddrMap(TextSectionIndex);
-    EXPECT_THAT_EXPECTED(PGOBBAddrMaps, Succeeded());
-    EXPECT_EQ(*PGOBBAddrMaps, ExpectedResult);
-  };
+  std::vector<PGOAnalysisMap> Section0PGOAnalysisMaps = {P4};
+  std::vector<PGOAnalysisMap> Section1PGOAnalysisMaps = {P3};
+  std::vector<PGOAnalysisMap> Section2PGOAnalysisMaps = {P1, P2};
+  std::vector<PGOAnalysisMap> AllPGOAnalysisMaps = {P1, P2, P3, P4};
+
+  auto DoCheckSucceeds =
+      [&](StringRef YamlString, std::optional<unsigned> TextSectionIndex,
+          std::vector<BBAddrMap> ExpectedResult,
+          std::optional<std::vector<PGOAnalysisMap>> ExpectedPGO) {
+        SmallString<0> Storage;
+        Expected<ELFObjectFile<ELF64LE>> ElfOrErr =
+            toBinary<ELF64LE>(Storage, YamlString);
+        ASSERT_THAT_EXPECTED(ElfOrErr, Succeeded());
+
+        Expected<const typename ELF64LE::Shdr *> BBAddrMapSecOrErr =
+            ElfOrErr->getELFFile().getSection(1);
+        ASSERT_THAT_EXPECTED(BBAddrMapSecOrErr, Succeeded());
+
+        std::vector<PGOAnalysisMap> PGOAnalyses;
+        auto BBAddrMaps = ElfOrErr->readBBAddrMap(
+            TextSectionIndex, ExpectedPGO ? &PGOAnalyses : nullptr);
+        EXPECT_THAT_EXPECTED(BBAddrMaps, Succeeded());
+        EXPECT_EQ(*BBAddrMaps, ExpectedResult);
+        if (ExpectedPGO) {
+          EXPECT_EQ(BBAddrMaps->size(), PGOAnalyses.size());
+          EXPECT_EQ(PGOAnalyses, *ExpectedPGO);
+        }
+      };
 
   auto DoCheckFails = [&](StringRef YamlString,
                           std::optional<unsigned> TextSectionIndex,
@@ -1135,22 +1061,35 @@ TEST(ELFObjectFileTest, ReadPGOBBAddrMap) {
     Expected<const typename ELF64LE::Shdr *> BBAddrMapSecOrErr =
         ElfOrErr->getELFFile().getSection(1);
     ASSERT_THAT_EXPECTED(BBAddrMapSecOrErr, Succeeded());
-    EXPECT_THAT_ERROR(ElfOrErr->readPGOBBAddrMap(TextSectionIndex).takeError(),
-                      FailedWithMessage(ErrMsg));
+    std::vector<PGOAnalysisMap> PGOAnalyses;
+    EXPECT_THAT_ERROR(
+        ElfOrErr->readBBAddrMap(TextSectionIndex, &PGOAnalyses).takeError(),
+        FailedWithMessage(ErrMsg));
   };
 
   // Check that we can retrieve the data in the normal case.
   DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/std::nullopt,
-                  AllPGOBBAddrMaps);
-  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/0,
-                  Section0PGOBBAddrMaps);
-  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/2,
-                  Section1PGOBBAddrMaps);
-  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/1,
-                  Section2PGOBBAddrMaps);
+                  AllBBAddrMaps, std::nullopt);
+  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/0, Section0BBAddrMaps,
+                  std::nullopt);
+  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/2, Section1BBAddrMaps,
+                  std::nullopt);
+  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/1, Section2BBAddrMaps,
+                  std::nullopt);
+
+  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/std::nullopt,
+                  AllBBAddrMaps, AllPGOAnalysisMaps);
+  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/0, Section0BBAddrMaps,
+                  Section0PGOAnalysisMaps);
+  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/2, Section1BBAddrMaps,
+                  Section1PGOAnalysisMaps);
+  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/1, Section2BBAddrMaps,
+                  Section2PGOAnalysisMaps);
   // Check that when no bb-address-map section is found for a text section,
   // we return an empty result.
-  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/3, {});
+  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/3, {}, std::nullopt);
+  DoCheckSucceeds(CommonYamlString, /*TextSectionIndex=*/3, {},
+                  std::vector<PGOAnalysisMap>{});
 
   // Check that we detect when a bb-addr-map section is linked to an invalid
   // (not present) section.
@@ -1161,12 +1100,14 @@ TEST(ELFObjectFileTest, ReadPGOBBAddrMap) {
 
   DoCheckFails(InvalidLinkedYamlString, /*TextSectionIndex=*/4,
                "unable to get the linked-to section for "
-               "SHT_LLVM_PGO_BB_ADDR_MAP section with index 4: invalid section "
+               "SHT_LLVM_BB_ADDR_MAP section with index 4: invalid section "
                "index: 10");
   // Linked sections are not checked when we don't target a specific text
   // section.
   DoCheckSucceeds(InvalidLinkedYamlString, /*TextSectionIndex=*/std::nullopt,
-                  AllPGOBBAddrMaps);
+                  AllBBAddrMaps, std::nullopt);
+  DoCheckSucceeds(InvalidLinkedYamlString, /*TextSectionIndex=*/std::nullopt,
+                  AllBBAddrMaps, AllPGOAnalysisMaps);
 
   // Check that we can detect when bb-address-map decoding fails.
   SmallString<128> TruncatedYamlString(CommonYamlString);
@@ -1175,13 +1116,15 @@ TEST(ELFObjectFileTest, ReadPGOBBAddrMap) {
 )";
 
   DoCheckFails(TruncatedYamlString, /*TextSectionIndex=*/std::nullopt,
-               "unable to read SHT_LLVM_PGO_BB_ADDR_MAP section with index 4: "
+               "unable to read SHT_LLVM_BB_ADDR_MAP section with index 4: "
                "unable to decode LEB128 at offset 0x0000000a: malformed "
                "uleb128, extends past end");
   // Check that we can read the other section's bb-address-maps which are
   // valid.
   DoCheckSucceeds(TruncatedYamlString, /*TextSectionIndex=*/2,
-                  Section1PGOBBAddrMaps);
+                  Section1BBAddrMaps, std::nullopt);
+  DoCheckSucceeds(TruncatedYamlString, /*TextSectionIndex=*/2,
+                  Section1BBAddrMaps, Section1PGOAnalysisMaps);
 }
 
 // Test for ObjectFile::getRelocatedSection: check that it returns a relocated
diff --git a/llvm/unittests/Object/ELFTypesTest.cpp b/llvm/unittests/Object/ELFTypesTest.cpp
index 8689f77af805073..d354b5e4570c352 100644
--- a/llvm/unittests/Object/ELFTypesTest.cpp
+++ b/llvm/unittests/Object/ELFTypesTest.cpp
@@ -100,17 +100,17 @@ TEST(ELFTypesTest, BBEntryMetadataInvalidEncodingTest) {
 }
 
 static_assert(
-    std::is_same_v<decltype(PGOBBAddrMap::BBEntry::SuccessorEntry::ID),
+    std::is_same_v<decltype(PGOAnalysisMap::PGOBBEntry::SuccessorEntry::ID),
                    decltype(BBAddrMap::BBEntry::ID)>,
-    "PGOBBAddrMap should use the same type for basic block ID as BBAddrMap");
+    "PGOAnalysisMap should use the same type for basic block ID as BBAddrMap");
 static_assert(BBAddrMap::BBEntry::Metadata::NumberOfBits <
                   (sizeof(uint32_t) * 8) - 2,
-              "currently PGOBBAddrMap relies on having two bits of space to "
+              "currently PGOAnalysisMap relies on having two bits of space to "
               "encode number of successors, to add more we need increase the "
               "encoded size of Metadata");
 
 TEST(ELFTypesTest, PGOBBEntryMetadataEncodingTest) {
-  using ST = PGOBBAddrMap::BBEntry::SuccessorsType;
+  using ST = PGOAnalysisMap::PGOBBEntry::SuccessorsType;
   const std::array<std::pair<BBAddrMap::BBEntry::Metadata, ST>, 7> Decoded = {
       {{{false, false, false, false, false}, ST::None},
        {{true, false, false, false, false}, ST::One},
@@ -124,11 +124,11 @@ TEST(ELFTypesTest, PGOBBEntryMetadataEncodingTest) {
                                             0b11'11111}};
   for (auto [Enc, Dec] : llvm::zip(Encoded, Decoded)) {
     auto [MD, SuccType] = Dec;
-    EXPECT_EQ(PGOBBAddrMap::BBEntry::encodeMD(MD, SuccType), Enc);
+    EXPECT_EQ(PGOAnalysisMap::PGOBBEntry::encodeMD(MD, SuccType), Enc);
   }
   for (auto [Enc, Dec] : llvm::zip(Encoded, Decoded)) {
     Expected<std::pair<BBAddrMap::BBEntry::Metadata, ST>> MetadataOrError =
-        PGOBBAddrMap::BBEntry::decodeMD(Enc);
+        PGOAnalysisMap::PGOBBEntry::decodeMD(Enc);
     ASSERT_THAT_EXPECTED(MetadataOrError, Succeeded());
     EXPECT_EQ(*MetadataOrError, Dec);
   }
@@ -141,7 +141,7 @@ TEST(ELFTypesTest, PGOBBEntryMetadataInvalidEncodingTest) {
       "invalid encoding for BBEntry::Metadata: 0x80"};
   const std::array<uint32_t, 3> Values = {0xFFFF, 0x100001, 0x00c0};
   for (auto [Val, Err] : llvm::zip(Values, Errors)) {
-    EXPECT_THAT_ERROR(PGOBBAddrMap::BBEntry::decodeMD(Val).takeError(),
+    EXPECT_THAT_ERROR(PGOAnalysisMap::PGOBBEntry::decodeMD(Val).takeError(),
                       FailedWithMessage(Err));
   }
 }

>From 3411bb89ac5e4601fc69364eb509d32394279bed Mon Sep 17 00:00:00 2001
From: Micah Weston <micahsweston at gmail.com>
Date: Mon, 20 Nov 2023 19:49:53 -0500
Subject: [PATCH 3/8] Uses the new constructor from recent NFC commit.

---
 llvm/lib/Object/ELF.cpp                     |  2 +-
 llvm/unittests/Object/ELFObjectFileTest.cpp | 20 +++++++++-----------
 2 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Object/ELF.cpp b/llvm/lib/Object/ELF.cpp
index 6517ce3ead092ff..396bb3242ca46b5 100644
--- a/llvm/lib/Object/ELF.cpp
+++ b/llvm/lib/Object/ELF.cpp
@@ -818,7 +818,7 @@ decodeBBAddrMapImpl(const ELFFile<ELFT> &EF,
       PGOAnalyses->push_back({FuncEntryCount, std::move(PGOBBEntries),
                               FuncEntryCountEnabled, BBFreqEnabled,
                               BrProbEnabled});
-    FunctionEntries.push_back({Address, std::move(BBEntries)});
+    FunctionEntries.emplace_back(Address, std::move(BBEntries));
   }
   // Either Cur is in the error state, or we have an error in ULEBSizeErr or
   // MetadataDecodeErr (but not both), but we join all errors here to be safe.
diff --git a/llvm/unittests/Object/ELFObjectFileTest.cpp b/llvm/unittests/Object/ELFObjectFileTest.cpp
index da8af2f69195dea..ee8ffda11661a2e 100644
--- a/llvm/unittests/Object/ELFObjectFileTest.cpp
+++ b/llvm/unittests/Object/ELFObjectFileTest.cpp
@@ -979,14 +979,13 @@ TEST(ELFObjectFileTest, ReadPGOBBAddrMap) {
             Successors:    []
 )");
 
-  BBAddrMap E1 = {0x11111, {{1, 0x0, 0x1, {false, true, false, false, false}}}};
+  BBAddrMap E1(0x11111, {{1, 0x0, 0x1, {false, true, false, false, false}}});
   PGOAnalysisMap P1 = {892, {{}}, true, false, false};
-  BBAddrMap E2 = {0x22222, {{2, 0x0, 0x2, {false, false, true, false, false}}}};
+  BBAddrMap E2(0x22222, {{2, 0x0, 0x2, {false, false, true, false, false}}});
   PGOAnalysisMap P2 = {{}, {{BlockFrequency(343), {}}}, false, true, false};
-  BBAddrMap E3 = {0x33333,
-                  {{0, 0x0, 0x3, {false, true, true, false, false}},
-                   {1, 0x3, 0x3, {false, false, true, false, false}},
-                   {2, 0x6, 0x3, {false, false, false, false, false}}}};
+  BBAddrMap E3(0x33333, {{0, 0x0, 0x3, {false, true, true, false, false}},
+                         {1, 0x3, 0x3, {false, false, true, false, false}},
+                         {2, 0x6, 0x3, {false, false, false, false, false}}});
   PGOAnalysisMap P3 = {{},
                        {{{},
                          {{1, BranchProbability::getRaw(0x1111'1111)},
@@ -996,11 +995,10 @@ TEST(ELFObjectFileTest, ReadPGOBBAddrMap) {
                        false,
                        false,
                        true};
-  BBAddrMap E4 = {0x44444,
-                  {{0, 0x0, 0x4, {false, false, false, true, true}},
-                   {1, 0x4, 0x4, {false, false, false, false, false}},
-                   {2, 0x8, 0x4, {false, false, false, false, false}},
-                   {3, 0xc, 0x4, {false, false, false, false, false}}}};
+  BBAddrMap E4(0x44444, {{0, 0x0, 0x4, {false, false, false, true, true}},
+                         {1, 0x4, 0x4, {false, false, false, false, false}},
+                         {2, 0x8, 0x4, {false, false, false, false, false}},
+                         {3, 0xc, 0x4, {false, false, false, false, false}}});
   PGOAnalysisMap P4 = {
       1000,
       {{BlockFrequency(1000),

>From a8943fe2e7976623b9fdee164aba0f607f467cb2 Mon Sep 17 00:00:00 2001
From: Micah Weston <micahsweston at gmail.com>
Date: Tue, 21 Nov 2023 17:54:47 -0500
Subject: [PATCH 4/8] Address some feedback typos and unnecessary headers.

---
 llvm/include/llvm/Object/ELF.h           | 2 +-
 llvm/include/llvm/Object/ELFObjectFile.h | 6 +++---
 llvm/include/llvm/Object/ELFTypes.h      | 4 ++--
 llvm/lib/Object/ELF.cpp                  | 2 --
 4 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/llvm/include/llvm/Object/ELF.h b/llvm/include/llvm/Object/ELF.h
index cdc902559fb3bb3..0866614e4d49da6 100644
--- a/llvm/include/llvm/Object/ELF.h
+++ b/llvm/include/llvm/Object/ELF.h
@@ -413,7 +413,7 @@ 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
+  /// Optional out variable to collect all 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>>
diff --git a/llvm/include/llvm/Object/ELFObjectFile.h b/llvm/include/llvm/Object/ELFObjectFile.h
index 66171f1c355cc4b..91408724bcb2a02 100644
--- a/llvm/include/llvm/Object/ELFObjectFile.h
+++ b/llvm/include/llvm/Object/ELFObjectFile.h
@@ -111,9 +111,9 @@ 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. 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.
+  // specified, the vector is cleared then filled with extra PGO data.
+  // `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,
                 std::vector<PGOAnalysisMap> *PGOAnalyses = nullptr) const;
diff --git a/llvm/include/llvm/Object/ELFTypes.h b/llvm/include/llvm/Object/ELFTypes.h
index a89cdd58cf78afe..2e35fedae974f3c 100644
--- a/llvm/include/llvm/Object/ELFTypes.h
+++ b/llvm/include/llvm/Object/ELFTypes.h
@@ -896,7 +896,7 @@ struct PGOAnalysisMap {
   struct PGOBBEntry {
     using BaseMetadata = BBAddrMap::BBEntry::Metadata;
 
-    /// Enum indicating the how many successors a block has. This enum must fit
+    /// Enum indicating how many successors a block has. This enum must fit
     /// into two bits.
     enum class SuccessorsType {
       /// None should be present if BBAddrMap.feature has disabled branch
@@ -915,7 +915,7 @@ struct PGOAnalysisMap {
     struct SuccessorEntry {
       /// Unique ID of this successor basic block.
       uint32_t ID;
-      /// Branch Probability of the edge to this successor taken from MBPI
+      /// Branch Probability of the edge to this successor taken from MBPI.
       BranchProbability Prob;
 
       bool operator==(const SuccessorEntry &Other) const {
diff --git a/llvm/lib/Object/ELF.cpp b/llvm/lib/Object/ELF.cpp
index 396bb3242ca46b5..f6e2dcefb87f992 100644
--- a/llvm/lib/Object/ELF.cpp
+++ b/llvm/lib/Object/ELF.cpp
@@ -10,8 +10,6 @@
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/Support/DataExtractor.h"
-#include <type_traits>
-#include <vector>
 
 using namespace llvm;
 using namespace object;

>From 03ed2b2e8f3662a3a132b7fb168aa901890ddf70 Mon Sep 17 00:00:00 2001
From: Micah Weston <micahsweston at gmail.com>
Date: Tue, 21 Nov 2023 18:03:37 -0500
Subject: [PATCH 5/8] Removes another unnecessarily added header.

---
 llvm/lib/Object/ELFObjectFile.cpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/llvm/lib/Object/ELFObjectFile.cpp b/llvm/lib/Object/ELFObjectFile.cpp
index 630141b20a9ae0f..462cef1c6d4cca4 100644
--- a/llvm/lib/Object/ELFObjectFile.cpp
+++ b/llvm/lib/Object/ELFObjectFile.cpp
@@ -33,7 +33,6 @@
 #include <optional>
 #include <string>
 #include <utility>
-#include <vector>
 
 using namespace llvm;
 using namespace object;

>From 675b89b7a7de0c1777394d49250ab6ab019ba6da Mon Sep 17 00:00:00 2001
From: Micah Weston <micahsweston at gmail.com>
Date: Wed, 22 Nov 2023 12:24:42 -0500
Subject: [PATCH 6/8] Changes to SFINAE, removes radix, and removes old
 comment.

---
 llvm/include/llvm/Object/ELFTypes.h | 2 --
 llvm/lib/Object/ELF.cpp             | 6 ++----
 2 files changed, 2 insertions(+), 6 deletions(-)

diff --git a/llvm/include/llvm/Object/ELFTypes.h b/llvm/include/llvm/Object/ELFTypes.h
index 2e35fedae974f3c..5277f41085807a3 100644
--- a/llvm/include/llvm/Object/ELFTypes.h
+++ b/llvm/include/llvm/Object/ELFTypes.h
@@ -958,8 +958,6 @@ struct PGOAnalysisMap {
              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 FuncEntryCount;           // Prof count from IR function
   std::vector<PGOBBEntry> BBEntries; // Extended basic block entries
 
diff --git a/llvm/lib/Object/ELF.cpp b/llvm/lib/Object/ELF.cpp
index f6e2dcefb87f992..fa01d3ae55a99b7 100644
--- a/llvm/lib/Object/ELF.cpp
+++ b/llvm/lib/Object/ELF.cpp
@@ -10,6 +10,7 @@
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/Support/DataExtractor.h"
+#include <type_traits>
 
 using namespace llvm;
 using namespace object;
@@ -650,12 +651,9 @@ ELFFile<ELFT>::toMappedAddr(uint64_t VAddr, WarningHandler WarnHandler) const {
 // int limit.
 // Also returns zero if ULEBSizeErr is already in an error state.
 // ULEBSizeErr is an out variable if an error occurs.
-template <typename IntTy>
+template <typename IntTy, std::enable_if_t<std::is_unsigned_v<IntTy>, int> = 0>
 static IntTy readULEB128As(DataExtractor &Data, DataExtractor::Cursor &Cur,
                            Error &ULEBSizeErr) {
-  static_assert(std::is_unsigned_v<IntTy> &&
-                    (std::numeric_limits<IntTy>::radix == 2),
-                "only use unsigned radix 2");
   // Bail out and do not extract data if ULEBSizeErr is already set.
   if (ULEBSizeErr)
     return 0;

>From bad3043cdcf127f4ce6c72c84dc51c269ab9ae0d Mon Sep 17 00:00:00 2001
From: Micah Weston <micahsweston at gmail.com>
Date: Wed, 22 Nov 2023 13:49:18 -0500
Subject: [PATCH 7/8] Encodes the PGO analysis data after each function entry.

---
 llvm/include/llvm/Object/ELFTypes.h         | 51 +----------
 llvm/lib/Object/ELF.cpp                     | 97 ++++++++++-----------
 llvm/lib/Object/ELFObjectFile.cpp           |  1 +
 llvm/lib/ObjectYAML/ELFEmitter.cpp          | 62 +++++++------
 llvm/unittests/Object/ELFObjectFileTest.cpp | 14 +--
 llvm/unittests/Object/ELFTypesTest.cpp      | 47 +---------
 6 files changed, 88 insertions(+), 184 deletions(-)

diff --git a/llvm/include/llvm/Object/ELFTypes.h b/llvm/include/llvm/Object/ELFTypes.h
index 5277f41085807a3..64e87166f69b8c0 100644
--- a/llvm/include/llvm/Object/ELFTypes.h
+++ b/llvm/include/llvm/Object/ELFTypes.h
@@ -807,11 +807,6 @@ 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 &&
@@ -891,25 +886,9 @@ struct PGOAnalysisMap {
     BrProb = (1 << 2),
   };
 
-  /// Super-set of BBAddrMap::BBEntry with additional fields for block frequency
-  /// and branch probability.
+  /// Extra basic block data with fields for block frequency and branch
+  /// probability.
   struct PGOBBEntry {
-    using BaseMetadata = BBAddrMap::BBEntry::Metadata;
-
-    /// Enum indicating how many successors a block has. This enum must fit
-    /// into two bits.
-    enum class SuccessorsType {
-      /// None should be present if BBAddrMap.feature 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 {
@@ -928,36 +907,12 @@ struct PGOAnalysisMap {
     /// 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 PGOBBEntry &Other) const {
       return std::tie(BlockFreq, Successors) ==
              std::tie(Other.BlockFreq, Other.Successors);
     }
   };
+
   uint64_t FuncEntryCount;           // Prof count from IR function
   std::vector<PGOBBEntry> BBEntries; // Extended basic block entries
 
diff --git a/llvm/lib/Object/ELF.cpp b/llvm/lib/Object/ELF.cpp
index fa01d3ae55a99b7..2c9665b41ec2559 100644
--- a/llvm/lib/Object/ELF.cpp
+++ b/llvm/lib/Object/ELF.cpp
@@ -723,6 +723,12 @@ decodeBBAddrMapImpl(const ELFFile<ELFT> &EF,
           Feature & uint8_t(PGOAnalysisMap::Features::FuncEntryCnt);
       BBFreqEnabled = Feature & uint8_t(PGOAnalysisMap::Features::BBFreq);
       BrProbEnabled = Feature & uint8_t(PGOAnalysisMap::Features::BrProb);
+      if (Feature != 0 && Version < 2 && Cur)
+        return createError(
+            "version should be >= 2 for SHT_LLVM_BB_ADDR_MAP when "
+            "PGO features are enabled: version = " +
+            Twine(static_cast<int>(Version)) +
+            " feature = " + Twine(static_cast<int>(Feature)));
     }
     uint64_t SectionOffset = Cur.tell();
     auto Address =
@@ -741,18 +747,7 @@ decodeBBAddrMapImpl(const ELFFile<ELFT> &EF,
     }
     uint32_t NumBlocks = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
 
-    if (Feature != 0 && Version < 2 && Cur)
-      return createError("version should be >= 2 for SHT_LLVM_BB_ADDR_MAP when "
-                         "PGO features are enabled: version = " +
-                         Twine(static_cast<int>(Version)) +
-                         " feature = " + Twine(static_cast<int>(Feature)));
-
-    uint64_t FuncEntryCount =
-        FuncEntryCountEnabled ? readULEB128As<uint64_t>(Data, Cur, ULEBSizeErr)
-                              : 0;
-
     std::vector<BBAddrMap::BBEntry> BBEntries;
-    std::vector<PGOAnalysisMap::PGOBBEntry> PGOBBEntries;
     uint32_t PrevBBEndOffset = 0;
     for (uint32_t BlockIndex = 0;
          !MetadataDecodeErr && !ULEBSizeErr && Cur && (BlockIndex < NumBlocks);
@@ -768,53 +763,53 @@ decodeBBAddrMapImpl(const ELFFile<ELFT> &EF,
         Offset += PrevBBEndOffset;
         PrevBBEndOffset = Offset + Size;
       }
-
-      BBAddrMap::BBEntry::Metadata Metadata;
-      PGOAnalysisMap::PGOBBEntry::SuccessorsType ST{};
-      if (BrProbEnabled) {
-        auto MetadataOrErr = PGOAnalysisMap::PGOBBEntry::decodeMD(MD);
-        if (Error E = MetadataOrErr.takeError()) {
-          MetadataDecodeErr = std::move(E);
-          break;
-        }
-        std::tie(Metadata, ST) = *MetadataOrErr;
-      } else {
-        auto MetadataOrErr = BBAddrMap::BBEntry::Metadata::decode(MD);
-        if (Error E = MetadataOrErr.takeError()) {
-          MetadataDecodeErr = std::move(E);
-          break;
-        }
-        Metadata = *MetadataOrErr;
+      Expected<BBAddrMap::BBEntry::Metadata> MetadataOrErr =
+          BBAddrMap::BBEntry::Metadata::decode(MD);
+      if (!MetadataOrErr) {
+        MetadataDecodeErr = MetadataOrErr.takeError();
+        break;
       }
+      BBEntries.push_back({ID, Offset, Size, *MetadataOrErr});
+    }
+    FunctionEntries.emplace_back(Address, std::move(BBEntries));
 
-      uint64_t BBF =
-          BBFreqEnabled ? readULEB128As<uint64_t>(Data, Cur, ULEBSizeErr) : 0;
-
-      llvm::SmallVector<PGOAnalysisMap::PGOBBEntry::SuccessorEntry, 2>
-          Successors;
-      if (BrProbEnabled) {
-        auto SuccCount =
-            ST == PGOAnalysisMap::PGOBBEntry::SuccessorsType::Multiple
-                ? readULEB128As<uint64_t>(Data, Cur, ULEBSizeErr)
-                : uint64_t(ST);
-        for (uint64_t I = 0; I < SuccCount; ++I) {
-          uint32_t BBID = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-          uint32_t BrProb = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-          if (PGOAnalyses)
-            Successors.push_back({BBID, BranchProbability::getRaw(BrProb)});
+    if (FuncEntryCountEnabled || BBFreqEnabled || BrProbEnabled) {
+      // Function entry count
+      uint64_t FuncEntryCount =
+          FuncEntryCountEnabled
+              ? readULEB128As<uint64_t>(Data, Cur, ULEBSizeErr)
+              : 0;
+
+      std::vector<PGOAnalysisMap::PGOBBEntry> PGOBBEntries;
+      for (uint32_t BlockIndex = 0; !MetadataDecodeErr && !ULEBSizeErr && Cur &&
+                                    (BlockIndex < NumBlocks);
+           ++BlockIndex) {
+        // Block frequency
+        uint64_t BBF =
+            BBFreqEnabled ? readULEB128As<uint64_t>(Data, Cur, ULEBSizeErr) : 0;
+
+        // Branch probability
+        llvm::SmallVector<PGOAnalysisMap::PGOBBEntry::SuccessorEntry, 2>
+            Successors;
+        if (BrProbEnabled) {
+          auto SuccCount = readULEB128As<uint64_t>(Data, Cur, ULEBSizeErr);
+          for (uint64_t I = 0; I < SuccCount; ++I) {
+            uint32_t BBID = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
+            uint32_t BrProb = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
+            if (PGOAnalyses)
+              Successors.push_back({BBID, BranchProbability::getRaw(BrProb)});
+          }
         }
+
+        if (PGOAnalyses)
+          PGOBBEntries.push_back({BlockFrequency(BBF), std::move(Successors)});
       }
 
-      BBEntries.push_back({ID, Offset, Size, Metadata});
       if (PGOAnalyses)
-        PGOBBEntries.push_back({BlockFrequency(BBF), std::move(Successors)});
+        PGOAnalyses->push_back({FuncEntryCount, std::move(PGOBBEntries),
+                                FuncEntryCountEnabled, BBFreqEnabled,
+                                BrProbEnabled});
     }
-
-    if (PGOAnalyses)
-      PGOAnalyses->push_back({FuncEntryCount, std::move(PGOBBEntries),
-                              FuncEntryCountEnabled, BBFreqEnabled,
-                              BrProbEnabled});
-    FunctionEntries.emplace_back(Address, std::move(BBEntries));
   }
   // Either Cur is in the error state, or we have an error in ULEBSizeErr or
   // MetadataDecodeErr (but not both), but we join all errors here to be safe.
diff --git a/llvm/lib/Object/ELFObjectFile.cpp b/llvm/lib/Object/ELFObjectFile.cpp
index 462cef1c6d4cca4..9f035587837e505 100644
--- a/llvm/lib/Object/ELFObjectFile.cpp
+++ b/llvm/lib/Object/ELFObjectFile.cpp
@@ -839,3 +839,4 @@ Expected<std::vector<BBAddrMap>> ELFObjectFileBase::readBBAddrMap(
   return readBBAddrMapImpl(cast<ELF64BEObjectFile>(this)->getELFFile(),
                            TextSectionIndex, PGOAnalyses);
 }
+
diff --git a/llvm/lib/ObjectYAML/ELFEmitter.cpp b/llvm/lib/ObjectYAML/ELFEmitter.cpp
index 7df4ca786e21b1d..56f4995131bc581 100644
--- a/llvm/lib/ObjectYAML/ELFEmitter.cpp
+++ b/llvm/lib/ObjectYAML/ELFEmitter.cpp
@@ -176,9 +176,8 @@ struct Fragment {
 /// TODO: This class still has a ways to go before it is truly a "single
 /// point of truth".
 template <class ELFT> class ELFState {
-public:
   LLVM_ELF_IMPORT_TYPES_ELFT(ELFT)
-private:
+
   enum class SymtabType { Static, Dynamic };
 
   /// The future symbol table string section.
@@ -1410,8 +1409,6 @@ void ELFState<ELFT>::writeSectionContent(
   }
 
   for (const auto &[Idx, E] : llvm::enumerate(*Section.Entries)) {
-    const ELFYAML::PGOAnalysisMapEntry *PGOEntry =
-        PGOAnalyses ? &PGOAnalyses->at(Idx) : nullptr;
     // Write version and feature values.
     if (Section.Type == llvm::ELF::SHT_LLVM_BB_ADDR_MAP) {
       if (E.Version > 2)
@@ -1437,42 +1434,43 @@ void ELFState<ELFT>::writeSectionContent(
     uint64_t NumBlocks =
         E.NumBlocks.value_or(E.BBEntries ? E.BBEntries->size() : 0);
     SHeader.sh_size += sizeof(uintX_t) + CBA.writeULEB128(NumBlocks);
+    // Write all BBEntries.
+    if (E.BBEntries) {
+      for (const ELFYAML::BBAddrMapEntry::BBEntry &BBE : *E.BBEntries) {
+        if (Section.Type == llvm::ELF::SHT_LLVM_BB_ADDR_MAP && E.Version > 1)
+          SHeader.sh_size += CBA.writeULEB128(BBE.ID);
+        SHeader.sh_size += CBA.writeULEB128(BBE.AddressOffset) +
+                           CBA.writeULEB128(BBE.Size) +
+                           CBA.writeULEB128(BBE.Metadata);
+      }
+    }
+
+    if (!PGOAnalyses)
+      continue;
+    const ELFYAML::PGOAnalysisMapEntry &PGOEntry = PGOAnalyses->at(Idx);
 
-    if (PGOEntry && PGOEntry->FuncEntryCount)
-      SHeader.sh_size += CBA.writeULEB128(*PGOEntry->FuncEntryCount);
+    if (PGOEntry.FuncEntryCount)
+      SHeader.sh_size += CBA.writeULEB128(*PGOEntry.FuncEntryCount);
 
-    const auto *PGOBBEntries = PGOEntry && PGOEntry->PGOBBEntries
-                                   ? &PGOEntry->PGOBBEntries.value()
-                                   : nullptr;
-    if (PGOBBEntries && E.BBEntries &&
-        E.BBEntries->size() != PGOBBEntries->size()) {
-      PGOBBEntries = nullptr;
+    if (!PGOEntry.PGOBBEntries)
+      continue;
+
+    const auto &PGOBBEntries = PGOEntry.PGOBBEntries.value();
+    if (!E.BBEntries || E.BBEntries->size() != PGOBBEntries.size()) {
       WithColor::warning() << "PBOBBEntries must be the same length as "
                               "BBEntries in SHT_LLVM_BB_ADDR_MAP.\n"
                            << "Mismatch on function with address: "
                            << E.Address;
+      continue;
     }
 
-    // Write all BBEntries.
-    if (!E.BBEntries)
-      continue;
-    for (const auto &[BBIdx, BBE] : llvm::enumerate(*E.BBEntries)) {
-      const auto *PGOBBE = PGOBBEntries ? &PGOBBEntries->at(BBIdx) : nullptr;
-      if (Section.Type == llvm::ELF::SHT_LLVM_BB_ADDR_MAP && E.Version > 1)
-        SHeader.sh_size += CBA.writeULEB128(BBE.ID);
-      SHeader.sh_size += CBA.writeULEB128(BBE.AddressOffset) +
-                         CBA.writeULEB128(BBE.Size) +
-                         CBA.writeULEB128(BBE.Metadata);
-
-      if (PGOBBE) {
-        if (PGOBBE->BBFreq)
-          SHeader.sh_size += CBA.writeULEB128(*PGOBBE->BBFreq);
-        if (PGOBBE->Successors) {
-          if (PGOBBE->Successors->size() > 2)
-            SHeader.sh_size += CBA.writeULEB128(PGOBBE->Successors->size());
-          for (const auto &[ID, BrProb] : *PGOBBE->Successors)
-            SHeader.sh_size += CBA.writeULEB128(ID) + CBA.writeULEB128(BrProb);
-        }
+    for (const auto &PGOBBE : PGOBBEntries) {
+      if (PGOBBE.BBFreq)
+        SHeader.sh_size += CBA.writeULEB128(*PGOBBE.BBFreq);
+      if (PGOBBE.Successors) {
+        SHeader.sh_size += CBA.writeULEB128(PGOBBE.Successors->size());
+        for (const auto &[ID, BrProb] : *PGOBBE.Successors)
+          SHeader.sh_size += CBA.writeULEB128(ID) + CBA.writeULEB128(BrProb);
       }
     }
   }
diff --git a/llvm/unittests/Object/ELFObjectFileTest.cpp b/llvm/unittests/Object/ELFObjectFileTest.cpp
index ee8ffda11661a2e..15a32eeda6673f7 100644
--- a/llvm/unittests/Object/ELFObjectFileTest.cpp
+++ b/llvm/unittests/Object/ELFObjectFileTest.cpp
@@ -910,15 +910,15 @@ TEST(ELFObjectFileTest, ReadPGOBBAddrMap) {
           - ID:            0
             AddressOffset: 0x0
             Size:          0x3
-            Metadata:      0x46
+            Metadata:      0x6
           - ID:            1
             AddressOffset: 0x0
             Size:          0x3
-            Metadata:      0x24
+            Metadata:      0x4
           - ID:            2
             AddressOffset: 0x0
             Size:          0x3
-            Metadata:      0x00
+            Metadata:      0x0
     PGOAnalyses:
       - PGOBBEntries:
          - Successors:
@@ -941,19 +941,19 @@ TEST(ELFObjectFileTest, ReadPGOBBAddrMap) {
           - ID:            0
             AddressOffset: 0x0
             Size:          0x4
-            Metadata:      0x78
+            Metadata:      0x18
           - ID:            1
             AddressOffset: 0x0
             Size:          0x4
-            Metadata:      0x40
+            Metadata:      0x0
           - ID:            2
             AddressOffset: 0x0
             Size:          0x4
-            Metadata:      0x20
+            Metadata:      0x0
           - ID:            3
             AddressOffset: 0x0
             Size:          0x4
-            Metadata:      0x00
+            Metadata:      0x0
     PGOAnalyses:
       - FuncEntryCount: 1000
         PGOBBEntries:
diff --git a/llvm/unittests/Object/ELFTypesTest.cpp b/llvm/unittests/Object/ELFTypesTest.cpp
index d354b5e4570c352..1d36727d50ee683 100644
--- a/llvm/unittests/Object/ELFTypesTest.cpp
+++ b/llvm/unittests/Object/ELFTypesTest.cpp
@@ -73,12 +73,8 @@ TEST(ELFTypesTest, BBEntryMetadataEncodingTest) {
        {false, false, false, false, true},
        {true, true, true, true, true}}};
   const std::array<uint32_t, 7> Encoded = {{0, 1, 2, 4, 8, 16, 31}};
-  for (size_t i = 0; i < Decoded.size(); ++i) {
+  for (size_t i = 0; i < Decoded.size(); ++i)
     EXPECT_EQ(Decoded[i].encode(), Encoded[i]);
-    EXPECT_LT(Decoded[i].encode(),
-              uint32_t{1} << BBAddrMap::BBEntry::Metadata::NumberOfBits)
-        << "If a new bit was added, please update NumberOfBits.";
-  }
   for (size_t i = 0; i < Encoded.size(); ++i) {
     Expected<BBAddrMap::BBEntry::Metadata> MetadataOrError =
         BBAddrMap::BBEntry::Metadata::decode(Encoded[i]);
@@ -103,45 +99,4 @@ static_assert(
     std::is_same_v<decltype(PGOAnalysisMap::PGOBBEntry::SuccessorEntry::ID),
                    decltype(BBAddrMap::BBEntry::ID)>,
     "PGOAnalysisMap should use the same type for basic block ID as BBAddrMap");
-static_assert(BBAddrMap::BBEntry::Metadata::NumberOfBits <
-                  (sizeof(uint32_t) * 8) - 2,
-              "currently PGOAnalysisMap relies on having two bits of space to "
-              "encode number of successors, to add more we need increase the "
-              "encoded size of Metadata");
-
-TEST(ELFTypesTest, PGOBBEntryMetadataEncodingTest) {
-  using ST = PGOAnalysisMap::PGOBBEntry::SuccessorsType;
-  const std::array<std::pair<BBAddrMap::BBEntry::Metadata, ST>, 7> Decoded = {
-      {{{false, false, false, false, false}, ST::None},
-       {{true, false, false, false, false}, ST::One},
-       {{false, true, false, false, false}, ST::Two},
-       {{false, false, true, false, false}, ST::Multiple},
-       {{false, false, false, true, false}, ST::One},
-       {{false, false, false, false, true}, ST::Two},
-       {{true, true, true, true, true}, ST::Multiple}}};
-  const std::array<uint32_t, 7> Encoded = {{0b00'00000, 0b01'00001, 0b10'00010,
-                                            0b11'00100, 0b01'01000, 0b10'10000,
-                                            0b11'11111}};
-  for (auto [Enc, Dec] : llvm::zip(Encoded, Decoded)) {
-    auto [MD, SuccType] = Dec;
-    EXPECT_EQ(PGOAnalysisMap::PGOBBEntry::encodeMD(MD, SuccType), Enc);
-  }
-  for (auto [Enc, Dec] : llvm::zip(Encoded, Decoded)) {
-    Expected<std::pair<BBAddrMap::BBEntry::Metadata, ST>> MetadataOrError =
-        PGOAnalysisMap::PGOBBEntry::decodeMD(Enc);
-    ASSERT_THAT_EXPECTED(MetadataOrError, Succeeded());
-    EXPECT_EQ(*MetadataOrError, Dec);
-  }
-}
 
-TEST(ELFTypesTest, PGOBBEntryMetadataInvalidEncodingTest) {
-  const std::array<std::string, 3> Errors = {
-      "invalid encoding for BBEntry::Metadata: 0xff9f",
-      "invalid encoding for BBEntry::Metadata: 0x100001",
-      "invalid encoding for BBEntry::Metadata: 0x80"};
-  const std::array<uint32_t, 3> Values = {0xFFFF, 0x100001, 0x00c0};
-  for (auto [Val, Err] : llvm::zip(Values, Errors)) {
-    EXPECT_THAT_ERROR(PGOAnalysisMap::PGOBBEntry::decodeMD(Val).takeError(),
-                      FailedWithMessage(Err));
-  }
-}

>From 2d56621d67a14edc74d6748078ffe6b5ea2fe88e Mon Sep 17 00:00:00 2001
From: Micah Weston <micahsweston at gmail.com>
Date: Wed, 22 Nov 2023 13:54:58 -0500
Subject: [PATCH 8/8] Removes new line to fix formatter error.

---
 llvm/lib/Object/ELFObjectFile.cpp      | 1 -
 llvm/unittests/Object/ELFTypesTest.cpp | 1 -
 2 files changed, 2 deletions(-)

diff --git a/llvm/lib/Object/ELFObjectFile.cpp b/llvm/lib/Object/ELFObjectFile.cpp
index 9f035587837e505..462cef1c6d4cca4 100644
--- a/llvm/lib/Object/ELFObjectFile.cpp
+++ b/llvm/lib/Object/ELFObjectFile.cpp
@@ -839,4 +839,3 @@ Expected<std::vector<BBAddrMap>> ELFObjectFileBase::readBBAddrMap(
   return readBBAddrMapImpl(cast<ELF64BEObjectFile>(this)->getELFFile(),
                            TextSectionIndex, PGOAnalyses);
 }
-
diff --git a/llvm/unittests/Object/ELFTypesTest.cpp b/llvm/unittests/Object/ELFTypesTest.cpp
index 1d36727d50ee683..241d33149d92cf5 100644
--- a/llvm/unittests/Object/ELFTypesTest.cpp
+++ b/llvm/unittests/Object/ELFTypesTest.cpp
@@ -99,4 +99,3 @@ static_assert(
     std::is_same_v<decltype(PGOAnalysisMap::PGOBBEntry::SuccessorEntry::ID),
                    decltype(BBAddrMap::BBEntry::ID)>,
     "PGOAnalysisMap should use the same type for basic block ID as BBAddrMap");
-



More information about the llvm-commits mailing list