[llvm] 105adf2 - [SHT_LLVM_BB_ADDR_MAP] Implements PGOAnalysisMap in Object and ObjectYAML with tests.

via llvm-commits llvm-commits at lists.llvm.org
Tue Dec 12 07:23:21 PST 2023


Author: Micah Weston
Date: 2023-12-12T10:23:16-05:00
New Revision: 105adf2cd9588b4839fbdc7287bb76a962fdb8ca

URL: https://github.com/llvm/llvm-project/commit/105adf2cd9588b4839fbdc7287bb76a962fdb8ca
DIFF: https://github.com/llvm/llvm-project/commit/105adf2cd9588b4839fbdc7287bb76a962fdb8ca.diff

LOG: [SHT_LLVM_BB_ADDR_MAP] Implements PGOAnalysisMap in Object and ObjectYAML with tests.

Reviewed in PR (#71750). A part of [RFC - PGO Accuracy Metrics: Emitting and Evaluating Branch
and Block
Analysis](https://discourse.llvm.org/t/rfc-pgo-accuracy-metrics-emitting-and-evaluating-branch-and-block-analysis/73902).

This PR adds the PGOAnalysisMap data structure and implements encoding and
decoding through Object and ObjectYAML along with associated tests. When
emitted into the bb-addr-map section, each function is followed by the associated
pgo-analysis-map for that function. The emitting of each analysis in the map is
controlled by a bit in the bb-addr-map feature byte. All existing bb-addr-map
code can ignore the pgo-analysis-map if the caller does not request the data.

Added: 
    

Modified: 
    llvm/include/llvm/Object/ELF.h
    llvm/include/llvm/Object/ELFObjectFile.h
    llvm/include/llvm/Object/ELFTypes.h
    llvm/include/llvm/ObjectYAML/ELFYAML.h
    llvm/lib/Object/ELF.cpp
    llvm/lib/Object/ELFObjectFile.cpp
    llvm/lib/ObjectYAML/ELFEmitter.cpp
    llvm/lib/ObjectYAML/ELFYAML.cpp
    llvm/unittests/Object/ELFObjectFileTest.cpp
    llvm/unittests/Object/ELFTypesTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Object/ELF.h b/llvm/include/llvm/Object/ELF.h
index 3fca00592f73b4..78b6b40cbddf67 100644
--- a/llvm/include/llvm/Object/ELF.h
+++ b/llvm/include/llvm/Object/ELF.h
@@ -457,8 +457,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 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>>
-  decodeBBAddrMap(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 761c532b9bf1f4..de418a1782acd2 100644
--- a/llvm/include/llvm/Object/ELFObjectFile.h
+++ b/llvm/include/llvm/Object/ELFObjectFile.h
@@ -110,9 +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.
+  // `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;
+  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 4670abc867de63..d3351a2d1650ed 100644
--- a/llvm/include/llvm/Object/ELFTypes.h
+++ b/llvm/include/llvm/Object/ELFTypes.h
@@ -13,6 +13,8 @@
 #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"
@@ -875,6 +877,79 @@ struct BBAddrMap {
   std::vector<BBEntry> BBEntries; // Basic block entries for this function.
 };
 
+/// A feature extension of BBAddrMap that holds information relevant to PGO.
+struct PGOAnalysisMap {
+  /// Bitfield of optional features to include in the PGO extended map.
+  struct Features {
+    bool FuncEntryCount : 1;
+    bool BBFreq : 1;
+    bool BrProb : 1;
+
+    // Encodes to minimum bit width representation.
+    uint8_t encode() const {
+      return (static_cast<uint8_t>(FuncEntryCount) << 0) |
+             (static_cast<uint8_t>(BBFreq) << 1) |
+             (static_cast<uint8_t>(BrProb) << 2);
+    }
+
+    // Decodes from minimum bit width representation and validates no
+    // unnecessary bits are used.
+    static Expected<Features> decode(uint8_t Val) {
+      Features Feat{static_cast<bool>(Val & (1 << 0)),
+                    static_cast<bool>(Val & (1 << 1)),
+                    static_cast<bool>(Val & (1 << 2))};
+      if (Feat.encode() != Val)
+        return createStringError(
+            std::error_code(),
+            "invalid encoding for PGOAnalysisMap::Features: 0x%x", Val);
+      return Feat;
+    }
+
+    bool operator==(const Features &Other) const {
+      return std::tie(FuncEntryCount, BBFreq, BrProb) ==
+             std::tie(Other.FuncEntryCount, Other.BBFreq, Other.BrProb);
+    }
+  };
+
+  /// Extra basic block data with fields for block frequency and branch
+  /// probability.
+  struct PGOBBEntry {
+    /// 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);
+      }
+    };
+
+    /// Block frequency taken from MBFI
+    BlockFrequency BlockFreq;
+    /// List of successors of the current block
+    llvm::SmallVector<SuccessorEntry, 2> Successors;
+
+    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
+
+  // Flags to indicate if each PGO related info was enabled in this function
+  Features FeatEnable;
+
+  bool operator==(const PGOAnalysisMap &Other) const {
+    return std::tie(FuncEntryCount, BBEntries, FeatEnable) ==
+           std::tie(Other.FuncEntryCount, Other.BBEntries, Other.FeatEnable);
+  }
+};
+
 } // end namespace object.
 } // end namespace llvm.
 

diff  --git a/llvm/include/llvm/ObjectYAML/ELFYAML.h b/llvm/include/llvm/ObjectYAML/ELFYAML.h
index 1ba41232f552e3..12b47c271da2cd 100644
--- a/llvm/include/llvm/ObjectYAML/ELFYAML.h
+++ b/llvm/include/llvm/ObjectYAML/ELFYAML.h
@@ -170,6 +170,19 @@ struct BBAddrMapEntry {
   std::optional<std::vector<BBEntry>> BBEntries;
 };
 
+struct PGOAnalysisMapEntry {
+  struct PGOBBEntry {
+    struct SuccessorEntry {
+      uint32_t ID;
+      llvm::yaml::Hex32 BrProb;
+    };
+    std::optional<uint64_t> BBFreq;
+    std::optional<std::vector<SuccessorEntry>> Successors;
+  };
+  std::optional<uint64_t> FuncEntryCount;
+  std::optional<std::vector<PGOBBEntry>> PGOBBEntries;
+};
+
 struct StackSizeEntry {
   llvm::yaml::Hex64 Address;
   llvm::yaml::Hex64 Size;
@@ -317,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) {}
 
@@ -737,6 +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::PGOAnalysisMapEntry)
+LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::ELFYAML::PGOAnalysisMapEntry::PGOBBEntry)
+LLVM_YAML_IS_SEQUENCE_VECTOR(
+    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)
@@ -905,6 +923,21 @@ template <> struct MappingTraits<ELFYAML::BBAddrMapEntry::BBEntry> {
   static void mapping(IO &IO, ELFYAML::BBAddrMapEntry::BBEntry &Rel);
 };
 
+template <> struct MappingTraits<ELFYAML::PGOAnalysisMapEntry> {
+  static void mapping(IO &IO, ELFYAML::PGOAnalysisMapEntry &Rel);
+};
+
+template <> struct MappingTraits<ELFYAML::PGOAnalysisMapEntry::PGOBBEntry> {
+  static void mapping(IO &IO, ELFYAML::PGOAnalysisMapEntry::PGOBBEntry &Rel);
+};
+
+template <>
+struct MappingTraits<ELFYAML::PGOAnalysisMapEntry::PGOBBEntry::SuccessorEntry> {
+  static void
+  mapping(IO &IO,
+          ELFYAML::PGOAnalysisMapEntry::PGOBBEntry::SuccessorEntry &Rel);
+};
+
 template <> struct MappingTraits<ELFYAML::GnuHashHeader> {
   static void mapping(IO &IO, ELFYAML::GnuHashHeader &Rel);
 };

diff  --git a/llvm/lib/Object/ELF.cpp b/llvm/lib/Object/ELF.cpp
index 36847d1a2a4223..300639f2bfa0af 100644
--- a/llvm/lib/Object/ELF.cpp
+++ b/llvm/lib/Object/ELF.cpp
@@ -646,11 +646,36 @@ 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, std::enable_if_t<std::is_unsigned_v<IntTy>, int> = 0>
+static IntTy readULEB128As(DataExtractor &Data, DataExtractor::Cursor &Cur,
+                           Error &ULEBSizeErr) {
+  // 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);
+}
+
+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
   // reference to the function in the SHT_LLVM_BB_ADDR_MAP section) to the
@@ -660,44 +685,28 @@ 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);
+  DataExtractor Data(Content, EF.isLE(), ELFT::Is64Bits ? 8 : 4);
   std::vector<BBAddrMap> 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);
-  };
 
   uint8_t Version = 0;
+  uint8_t Feature = 0;
+  PGOAnalysisMap::Features FeatEnable{};
   while (!ULEBSizeErr && !MetadataDecodeErr && Cur &&
          Cur.tell() < Content.size()) {
     if (Sec.sh_type == ELF::SHT_LLVM_BB_ADDR_MAP) {
@@ -707,10 +716,24 @@ ELFFile<ELFT>::decodeBBAddrMap(const Elf_Shdr &Sec,
       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
+      if (!Cur)
+        break;
+      auto FeatEnableOrErr = PGOAnalysisMap::Features::decode(Feature);
+      if (!FeatEnableOrErr)
+        return FeatEnableOrErr.takeError();
+      FeatEnable =
+          FeatEnableOrErr ? *FeatEnableOrErr : PGOAnalysisMap::Features{};
+      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();
-    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) {
@@ -719,20 +742,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();
+    uint32_t NumBlocks = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
+
     std::vector<BBAddrMap::BBEntry> BBEntries;
     uint32_t PrevBBEndOffset = 0;
     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();
+      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;
@@ -747,6 +773,44 @@ ELFFile<ELFT>::decodeBBAddrMap(const Elf_Shdr &Sec,
       BBEntries.push_back({ID, Offset, Size, *MetadataOrErr});
     }
     FunctionEntries.emplace_back(Address, std::move(BBEntries));
+
+    if (FeatEnable.FuncEntryCount || FeatEnable.BBFreq || FeatEnable.BrProb) {
+      // Function entry count
+      uint64_t FuncEntryCount =
+          FeatEnable.FuncEntryCount
+              ? 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 = FeatEnable.BBFreq
+                           ? readULEB128As<uint64_t>(Data, Cur, ULEBSizeErr)
+                           : 0;
+
+        // Branch probability
+        llvm::SmallVector<PGOAnalysisMap::PGOBBEntry::SuccessorEntry, 2>
+            Successors;
+        if (FeatEnable.BrProb) {
+          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)});
+      }
+
+      if (PGOAnalyses)
+        PGOAnalyses->push_back(
+            {FuncEntryCount, std::move(PGOBBEntries), FeatEnable});
+    }
   }
   // 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.
@@ -756,6 +820,18 @@ 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,
+                               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>
 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 25dbcbd68431f5..3c86b0f25ddac1 100644
--- a/llvm/lib/Object/ELFObjectFile.cpp
+++ b/llvm/lib/Object/ELFObjectFile.cpp
@@ -716,10 +716,13 @@ std::vector<ELFPltEntry> ELFObjectFileBase::getPltEntries() const {
 
 template <class ELFT>
 Expected<std::vector<BBAddrMap>> static readBBAddrMapImpl(
-    const ELFFile<ELFT> &EF, std::optional<unsigned> TextSectionIndex) {
+    const ELFFile<ELFT> &EF, std::optional<unsigned> TextSectionIndex,
+    std::vector<PGOAnalysisMap> *PGOAnalyses) {
   using Elf_Shdr = typename ELFT::Shdr;
   bool IsRelocatable = EF.getHeader().e_type == ELF::ET_REL;
   std::vector<BBAddrMap> BBAddrMaps;
+  if (PGOAnalyses)
+    PGOAnalyses->clear();
 
   const auto &Sections = cantFail(EF.sections());
   auto IsMatch = [&](const Elf_Shdr &Sec) -> Expected<bool> {
@@ -748,10 +751,13 @@ Expected<std::vector<BBAddrMap>> static readBBAddrMapImpl(
       return createError("unable to get relocation section for " +
                          describe(EF, *Sec));
     Expected<std::vector<BBAddrMap>> BBAddrMapOrErr =
-        EF.decodeBBAddrMap(*Sec, RelocSec);
-    if (!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));
   }
@@ -828,13 +834,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);
+                           TextSectionIndex, PGOAnalyses);
 }

diff  --git a/llvm/lib/ObjectYAML/ELFEmitter.cpp b/llvm/lib/ObjectYAML/ELFEmitter.cpp
index d54faec2086676..94b0529f761052 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;
 
@@ -1390,10 +1391,24 @@ template <class ELFT>
 void ELFState<ELFT>::writeSectionContent(
     Elf_Shdr &SHeader, const ELFYAML::BBAddrMapSection &Section,
     ContiguousBlobAccumulator &CBA) {
-  if (!Section.Entries)
+  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 ELFYAML::BBAddrMapEntry &E : *Section.Entries) {
+  for (const auto &[Idx, E] : llvm::enumerate(*Section.Entries)) {
     // Write version and feature values.
     if (Section.Type == llvm::ELF::SHT_LLVM_BB_ADDR_MAP) {
       if (E.Version > 2)
@@ -1404,6 +1419,14 @@ void ELFState<ELFT>::writeSectionContent(
       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.Address, ELFT::TargetEndianness);
     // Write number of BBEntries (number of basic blocks in the function). This
@@ -1412,14 +1435,43 @@ void ELFState<ELFT>::writeSectionContent(
         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)
+    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;
-    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);
+    const ELFYAML::PGOAnalysisMapEntry &PGOEntry = PGOAnalyses->at(Idx);
+
+    if (PGOEntry.FuncEntryCount)
+      SHeader.sh_size += CBA.writeULEB128(*PGOEntry.FuncEntryCount);
+
+    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;
+    }
+
+    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/lib/ObjectYAML/ELFYAML.cpp b/llvm/lib/ObjectYAML/ELFYAML.cpp
index 9d845a0c7c4870..6ad4a067415ac8 100644
--- a/llvm/lib/ObjectYAML/ELFYAML.cpp
+++ b/llvm/lib/ObjectYAML/ELFYAML.cpp
@@ -1390,6 +1390,7 @@ static void sectionMapping(IO &IO, ELFYAML::BBAddrMapSection &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) {
@@ -1825,6 +1826,28 @@ void MappingTraits<ELFYAML::BBAddrMapEntry::BBEntry>::mapping(
   IO.mapRequired("Metadata", E.Metadata);
 }
 
+void MappingTraits<ELFYAML::PGOAnalysisMapEntry>::mapping(
+    IO &IO, ELFYAML::PGOAnalysisMapEntry &E) {
+  assert(IO.getContext() && "The IO context is not initialized");
+  IO.mapOptional("FuncEntryCount", E.FuncEntryCount);
+  IO.mapOptional("PGOBBEntries", E.PGOBBEntries);
+}
+
+void MappingTraits<ELFYAML::PGOAnalysisMapEntry::PGOBBEntry>::mapping(
+    IO &IO, ELFYAML::PGOAnalysisMapEntry::PGOBBEntry &E) {
+  assert(IO.getContext() && "The IO context is not initialized");
+  IO.mapOptional("BBFreq", E.BBFreq);
+  IO.mapOptional("Successors", E.Successors);
+}
+
+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);
+}
+
 void MappingTraits<ELFYAML::GnuHashHeader>::mapping(IO &IO,
                                                     ELFYAML::GnuHashHeader &E) {
   assert(IO.getContext() && "The IO context is not initialized");

diff  --git a/llvm/unittests/Object/ELFObjectFileTest.cpp b/llvm/unittests/Object/ELFObjectFileTest.cpp
index 17402f39a5df83..2878ca088cd79d 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"
@@ -683,7 +684,7 @@ TEST(ELFObjectFileTest, ReadBBAddrMap) {
         ElfOrErr->getELFFile().getSection(1);
     ASSERT_THAT_EXPECTED(BBAddrMapSecOrErr, Succeeded());
     auto BBAddrMaps = ElfOrErr->readBBAddrMap(TextSectionIndex);
-    EXPECT_THAT_EXPECTED(BBAddrMaps, Succeeded());
+    ASSERT_THAT_EXPECTED(BBAddrMaps, Succeeded());
     EXPECT_EQ(*BBAddrMaps, ExpectedResult);
   };
 
@@ -744,6 +745,389 @@ TEST(ELFObjectFileTest, ReadBBAddrMap) {
                   Section1BBAddrMaps);
 }
 
+// Tests for error paths of the ELFFile::decodeBBAddrMap with PGOAnalysisMap
+// API.
+TEST(ELFObjectFileTest, InvalidDecodePGOAnalysisMap) {
+  if (IsHostWindows())
+    GTEST_SKIP();
+  StringRef CommonYamlString(R"(
+--- !ELF
+FileHeader:
+  Class: ELFCLASS64
+  Data:  ELFDATA2LSB
+  Type:  ET_EXEC
+Sections:
+  - Type: SHT_LLVM_BB_ADDR_MAP
+    Name: .llvm_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 *> BBAddrMapSecOrErr =
+        Elf.getSection(1);
+    ASSERT_THAT_EXPECTED(BBAddrMapSecOrErr, Succeeded());
+
+    std::vector<PGOAnalysisMap> PGOAnalyses;
+    EXPECT_THAT_ERROR(
+        Elf.decodeBBAddrMap(**BBAddrMapSecOrErr, nullptr, &PGOAnalyses)
+            .takeError(),
+        FailedWithMessage(ErrMsg));
+  };
+
+  // Check that we can detect unsupported versions that are too old.
+  SmallString<128> UnsupportedLowVersionYamlString(CommonYamlString);
+  UnsupportedLowVersionYamlString += R"(
+        Version: 1
+        Feature: 0x4
+        BBEntries:
+          - AddressOffset: 0x0
+            Size:          0x1
+            Metadata:      0x2
+)";
+
+  DoCheck(UnsupportedLowVersionYamlString,
+          "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"(
+        Version: 2
+        BBEntries:
+          - ID:            1
+            AddressOffset: 0x0
+            Size:          0x1
+            Metadata:      0x2
+)";
+
+  // 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: 0x04
+        BBEntries:
+          - ID:            1
+            AddressOffset: 0x0
+            Size:          0x1
+            Metadata:      0x6
+          - ID:            2
+            AddressOffset: 0x1
+            Size:          0x1
+            Metadata:      0x2
+          - ID:            3
+            AddressOffset: 0x2
+            Size:          0x1
+            Metadata:      0x2
+    PGOAnalyses:
+      - PGOBBEntries:
+         - Successors:
+            - ID:          2
+              BrProb:      0x80000000
+            - ID:          3
+              BrProb:      0x80000000
+         - Successors:
+            - ID:          3
+              BrProb:      0xF0000000
+)";
+
+  DoCheck(MissingBrProb, "unable to decode LEB128 at offset 0x00000017: "
+                         "malformed uleb128, extends past end");
+}
+
+// Test for the ELFObjectFile::readBBAddrMap API with PGOAnalysisMap.
+TEST(ELFObjectFileTest, ReadPGOAnalysisMap) {
+  if (IsHostWindows())
+    GTEST_SKIP();
+  StringRef CommonYamlString(R"(
+--- !ELF
+FileHeader:
+  Class: ELFCLASS64
+  Data:  ELFDATA2LSB
+  Type:  ET_EXEC
+Sections:
+  - Name: .llvm_bb_addr_map_1
+    Type: SHT_LLVM_BB_ADDR_MAP
+    Link: 1
+    Entries:
+      - Version: 2
+        Address: 0x11111
+        Feature: 0x1
+        BBEntries:
+          - ID:            1
+            AddressOffset: 0x0
+            Size:          0x1
+            Metadata:      0x2
+    PGOAnalyses:
+      - FuncEntryCount: 892
+  - Name: .llvm_bb_addr_map_2
+    Type: SHT_LLVM_BB_ADDR_MAP
+    Link: 1
+    Entries:
+      - Version: 2
+        Address: 0x22222
+        Feature: 0x2
+        BBEntries:
+          - ID:            2
+            AddressOffset: 0x0
+            Size:          0x2
+            Metadata:      0x4
+    PGOAnalyses:
+      - PGOBBEntries:
+         - BBFreq:         343
+  - Name: .llvm_bb_addr_map_3
+    Type: SHT_LLVM_BB_ADDR_MAP
+    Link: 2
+    Entries:
+      - Version: 2
+        Address: 0x33333
+        Feature: 0x4
+        BBEntries:
+          - ID:            0
+            AddressOffset: 0x0
+            Size:          0x3
+            Metadata:      0x6
+          - ID:            1
+            AddressOffset: 0x0
+            Size:          0x3
+            Metadata:      0x4
+          - ID:            2
+            AddressOffset: 0x0
+            Size:          0x3
+            Metadata:      0x0
+    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
+        BBEntries:
+          - ID:            0
+            AddressOffset: 0x0
+            Size:          0x4
+            Metadata:      0x18
+          - ID:            1
+            AddressOffset: 0x0
+            Size:          0x4
+            Metadata:      0x0
+          - ID:            2
+            AddressOffset: 0x0
+            Size:          0x4
+            Metadata:      0x0
+          - ID:            3
+            AddressOffset: 0x0
+            Size:          0x4
+            Metadata:      0x0
+    PGOAnalyses:
+      - FuncEntryCount: 1000
+        PGOBBEntries:
+          - BBFreq:         1000
+            Successors:
+            - ID:          1
+              BrProb:      0x22222222
+            - ID:          2
+              BrProb:      0x33333333
+            - ID:          3
+              BrProb:      0xaaaaaaaa
+          - BBFreq:         133
+            Successors:
+            - ID:          2
+              BrProb:      0x11111111
+            - ID:          3
+              BrProb:      0xeeeeeeee
+          - BBFreq:         18
+            Successors:
+            - ID:          3
+              BrProb:      0xffffffff
+          - BBFreq:         1000
+            Successors:    []
+)");
+
+  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}};
+
+  std::vector<BBAddrMap> Section0BBAddrMaps = {E4};
+  std::vector<BBAddrMap> Section1BBAddrMaps = {E3};
+  std::vector<BBAddrMap> Section2BBAddrMaps = {E1, E2};
+  std::vector<BBAddrMap> AllBBAddrMaps = {E1, E2, E3, E4};
+
+  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);
+        ASSERT_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,
+                          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());
+    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,
+                  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, {}, 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.
+  SmallString<128> InvalidLinkedYamlString(CommonYamlString);
+  InvalidLinkedYamlString += R"(
+    Link: 10
+)";
+
+  DoCheckFails(InvalidLinkedYamlString, /*TextSectionIndex=*/4,
+               "unable to get the linked-to section for "
+               "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,
+                  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);
+  TruncatedYamlString += R"(
+    ShSize: 0xa
+)";
+
+  DoCheckFails(TruncatedYamlString, /*TextSectionIndex=*/std::nullopt,
+               "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,
+                  Section1BBAddrMaps, std::nullopt);
+  DoCheckSucceeds(TruncatedYamlString, /*TextSectionIndex=*/2,
+                  Section1BBAddrMaps, Section1PGOAnalysisMaps);
+}
+
 // 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 2c9e8b7aebac29..f09ab5e1243846 100644
--- a/llvm/unittests/Object/ELFTypesTest.cpp
+++ b/llvm/unittests/Object/ELFTypesTest.cpp
@@ -94,3 +94,40 @@ TEST(ELFTypesTest, BBEntryMetadataInvalidEncodingTest) {
         FailedWithMessage(Errors[i]));
   }
 }
+
+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");
+
+TEST(ELFTypesTest, PGOAnalysisMapFeaturesEncodingTest) {
+  const std::array<PGOAnalysisMap::Features, 7> Decoded = {
+      {{false, false, false},
+       {true, false, false},
+       {false, true, false},
+       {false, false, true},
+       {true, true, false},
+       {false, true, true},
+       {true, true, true}}};
+  const std::array<uint8_t, 7> Encoded = {
+      {0b000, 0b001, 0b010, 0b100, 0b011, 0b110, 0b111}};
+  for (const auto &[Feat, EncodedVal] : llvm::zip(Decoded, Encoded))
+    EXPECT_EQ(Feat.encode(), EncodedVal);
+  for (const auto &[Feat, EncodedVal] : llvm::zip(Decoded, Encoded)) {
+    Expected<PGOAnalysisMap::Features> FeatEnableOrError =
+        PGOAnalysisMap::Features::decode(EncodedVal);
+    ASSERT_THAT_EXPECTED(FeatEnableOrError, Succeeded());
+    EXPECT_EQ(*FeatEnableOrError, Feat);
+  }
+}
+
+TEST(ELFTypesTest, PGOAnalysisMapFeaturesInvalidEncodingTest) {
+  const std::array<std::string, 2> Errors = {
+      "invalid encoding for PGOAnalysisMap::Features: 0x8",
+      "invalid encoding for PGOAnalysisMap::Features: 0xff"};
+  const std::array<uint8_t, 2> Values = {{0b1000, 0b1111'1111}};
+  for (const auto &[Val, Error] : llvm::zip(Values, Errors)) {
+    EXPECT_THAT_ERROR(PGOAnalysisMap::Features::decode(Val).takeError(),
+                      FailedWithMessage(Error));
+  }
+}


        


More information about the llvm-commits mailing list