[llvm-objdump] Add support for symbolizing PGOBBAddrMap Info (#76386)

This patch adds in support for symbolizing PGO information contained
within the SHT_LLVM_BB_ADDR_MAP section in llvm-objdump. The outputs are
simply the raw values contained within the section.

2 files changed, 263 insertions, 17 deletions

diff --git a/llvm/test/tools/llvm-objdump/X86/elf-pgoanalysismap.yaml b/llvm/test/tools/llvm-objdump/X86/elf-pgoanalysismap.yaml
new file mode 100644
index 0000000..c4bf443
--- /dev/null
+++ b/llvm/test/tools/llvm-objdump/X86/elf-pgoanalysismap.yaml
@@ -0,0 +1,180 @@
+## Test that in the presence of SHT_LLVM_BB_ADDR_MAP sections which also
+## contain PGO data, --symbolize-operands is able to label the basic blocks
+## correctly.
+
+## Check the case where we only have entry counts.
+
+# RUN: yaml2obj --docnum=1 %s -o %t1
+# RUN: llvm-objdump %t1 -d --symbolize-operands --no-show-raw-insn --no-leading-addr | \
+# RUN:   FileCheck %s --check-prefix=ENTRYCOUNT
+
+--- !ELF
+FileHeader:
+  Class:   ELFCLASS64
+  Data:    ELFDATA2LSB
+  Type:    ET_EXEC
+  Machine: EM_X86_64
+Sections:
+  - Name:    .text.foo
+    Type:    SHT_PROGBITS
+    Address: 0x0
+    Flags:   [SHF_ALLOC, SHF_EXECINSTR]
+    Content: '50'
+  - Name:   .llvm_bb_addr_map.foo
+    Type:   SHT_LLVM_BB_ADDR_MAP
+    Link:   .text.foo
+    Entries:
+      - Version: 2
+        Address: 0x0
+        Feature: 0x1
+        BBEntries:
+          - ID:            3
+            AddressOffset: 0x0
+            Size:          0x1
+            Metadata:      0x1
+    PGOAnalyses:
+      - FuncEntryCount: 1000
+Symbols:
+  - Name:    foo
+    Section: .text.foo
+    Value:   0x0
+
+# ENTRYCOUNT: <foo>:
+# ENTRYCOUNT: <BB3> (Entry count: 1000):
+
+## Check the case where we have entry points and block frequency information
+
+# RUN: yaml2obj %s --docnum=2 -o %t2
+# RUN: llvm-objdump %t2 -d --symbolize-operands --no-show-raw-insn --no-leading-addr | \
+# RUN:   FileCheck %s --check-prefix=ENTRYCOUNT-BLOCKFREQ
+
+--- !ELF
+FileHeader:
+  Class:   ELFCLASS64
+  Data:    ELFDATA2LSB
+  Type:    ET_EXEC
+  Machine: EM_X86_64
+Sections:
+  - Name:    .text.foo
+    Type:    SHT_PROGBITS
+    Address: 0x0
+    Flags:   [SHF_ALLOC, SHF_EXECINSTR]
+    Content: '503b0505200000907d02ebf5c3'
+  - Name:   .llvm_bb_addr_map.foo
+    Type:   SHT_LLVM_BB_ADDR_MAP
+    Link:   .text.foo
+    Entries:
+      - Version: 2
+        Address: 0x0
+        Feature: 0x3
+        BBEntries:
+          - ID:            3
+            AddressOffset: 0x0
+            Size:          0x1
+            Metadata:      0x1
+          - ID:            1
+            AddressOffset: 0x0
+            Size:          0x6
+            Metadata:      0x0
+          - ID:            2
+            AddressOffset: 0x1
+            Size:          0x4
+            Metadata:      0x0
+          - ID:            5
+            AddressOffset: 0x0
+            Size:          0x1
+            Metadata:      0x2
+    PGOAnalyses:
+      - FuncEntryCount: 1000
+        PGOBBEntries:
+          - BBFreq: 1000
+          - BBFreq: 133
+          - BBFreq: 18
+          - BBFreq: 1000
+Symbols:
+  - Name:    foo
+    Section: .text.foo
+    Value:   0x0
+
+# ENTRYCOUNT-BLOCKFREQ: <foo>:
+# ENTRYCOUNT-BLOCKFREQ: <BB3> (Entry count: 1000, Frequency: 1000):
+# ENTRYCOUNT-BLOCKFREQ: <BB1> (Frequency: 133):
+# ENTRYCOUNT-BLOCKFREQ: <BB2> (Frequency: 18):
+# ENTRYCOUNT-BLOCKFREQ: <BB5> (Frequency: 1000):
+
+## Check the case where we have entry points, block frequency, and branch
+## proabability information.
+
+# RUN: yaml2obj %s --docnum=3 -o %t3
+# RUN: llvm-objdump %t3 -d --symbolize-operands --no-show-raw-insn --no-leading-addr | \
+# RUN:   FileCheck %s --check-prefix=ENTRY-FREQ-PROB
+
+--- !ELF
+FileHeader:
+  Class:   ELFCLASS64
+  Data:    ELFDATA2LSB
+  Type:    ET_EXEC
+  Machine: EM_X86_64
+Sections:
+  - Name:    .text.foo
+    Type:    SHT_PROGBITS
+    Address: 0x0
+    Flags:   [SHF_ALLOC, SHF_EXECINSTR]
+    Content: '503b0505200000907d02ebf5c3'
+  - Name:   .llvm_bb_addr_map.foo
+    Type:   SHT_LLVM_BB_ADDR_MAP
+    Link:   .text.foo
+    Entries:
+      - Version: 2
+        Address: 0x0
+        Feature: 0x7
+        BBEntries:
+          - ID:            3
+            AddressOffset: 0x0
+            Size:          0x1
+            Metadata:      0x1
+          - ID:            1
+            AddressOffset: 0x0
+            Size:          0x6
+            Metadata:      0x0
+          - ID:            2
+            AddressOffset: 0x1
+            Size:          0x4
+            Metadata:      0x0
+          - ID:            5
+            AddressOffset: 0x0
+            Size:          0x1
+            Metadata:      0x2
+    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:    []
+Symbols:
+  - Name:    foo
+    Section: .text.foo
+    Value:   0x0
+
+# ENTRY-FREQ-PROB: <foo>:
+# ENTRY-FREQ-PROB: <BB3> (Entry count: 1000, Frequency: 1000, Successors: BB1:22222222, BB2:33333333, BB3:aaaaaaaa):
+# ENTRY-FREQ-PROB: <BB1> (Frequency: 133, Successors: BB2:11111111, BB3:eeeeeeee):
+# ENTRY-FREQ-PROB: <BB2> (Frequency: 18, Successors: BB3:ffffffff):
+# ENTRY-FREQ-PROB: <BB5> (Frequency: 1000):
diff --git a/llvm/tools/llvm-objdump/llvm-objdump.cpp b/llvm/tools/llvm-objdump/llvm-objdump.cpp
index 8944965..22b427f 100644
--- a/llvm/tools/llvm-objdump/llvm-objdump.cpp
+++ b/llvm/tools/llvm-objdump/llvm-objdump.cpp
@@ -1263,10 +1263,57 @@ static SymbolInfoTy createDummySymbolInfo(const ObjectFile &Obj,
     return SymbolInfoTy(Addr, Name, Type);
 }
 
-static void
-collectBBAddrMapLabels(const std::unordered_map<uint64_t, BBAddrMap> &AddrToBBAddrMap,
-                       uint64_t SectionAddr, uint64_t Start, uint64_t End,
-                       std::unordered_map<uint64_t, std::vector<std::string>> &Labels) {
+struct BBAddrMapLabel {
+  std::string BlockLabel;
+  std::string PGOAnalysis;
+};
+
+static std::string constructPGOLabelString(const PGOAnalysisMap &PGOMap,
+                                           size_t BBEntryIndex) {
+  std::string PGOString;
+  raw_string_ostream PGOSS(PGOString);
+
+  PGOSS << " (";
+  if (PGOMap.FeatEnable.FuncEntryCount && BBEntryIndex == 0) {
+    PGOSS << "Entry count: " << Twine(PGOMap.FuncEntryCount);
+    if (PGOMap.FeatEnable.BBFreq || PGOMap.FeatEnable.BrProb) {
+      PGOSS << ", ";
+    }
+  }
+
+  if (PGOMap.FeatEnable.BBFreq || PGOMap.FeatEnable.BrProb) {
+    assert(BBEntryIndex < PGOMap.BBEntries.size() &&
+           "Expected PGOAnalysisMap and BBAddrMap to have the same entires");
+    const PGOAnalysisMap::PGOBBEntry &PGOBBEntry =
+        PGOMap.BBEntries[BBEntryIndex];
+
+    if (PGOMap.FeatEnable.BBFreq) {
+      PGOSS << "Frequency: " << Twine(PGOBBEntry.BlockFreq.getFrequency());
+      if (PGOMap.FeatEnable.BrProb && PGOBBEntry.Successors.size() > 0) {
+        PGOSS << ", ";
+      }
+    }
+    if (PGOMap.FeatEnable.BrProb && PGOBBEntry.Successors.size() > 0) {
+      PGOSS << "Successors: ";
+      interleaveComma(
+          PGOBBEntry.Successors, PGOSS,
+          [&PGOSS](const PGOAnalysisMap::PGOBBEntry::SuccessorEntry &SE) {
+            PGOSS << "BB" << SE.ID << ":";
+            PGOSS.write_hex(SE.Prob.getNumerator());
+          });
+    }
+  }
+  PGOSS << ")";
+
+  return PGOString;
+}
+
+static void collectBBAddrMapLabels(
+    const std::unordered_map<uint64_t, BBAddrMap> &AddrToBBAddrMap,
+    const std::unordered_map<uint64_t, PGOAnalysisMap> &AddrToPGOAnalysisMap,
+    uint64_t SectionAddr, uint64_t Start, uint64_t End,
+    std::unordered_map<uint64_t, std::vector<BBAddrMapLabel>> &Labels,
+    const StringRef FileName) {
   if (AddrToBBAddrMap.empty())
     return;
   Labels.clear();
@@ -1275,11 +1322,21 @@ collectBBAddrMapLabels(const std::unordered_map<uint64_t, BBAddrMap> &AddrToBBAd
   auto Iter = AddrToBBAddrMap.find(StartAddress);
   if (Iter == AddrToBBAddrMap.end())
     return;
-  for (const BBAddrMap::BBEntry &BBEntry : Iter->second.getBBEntries()) {
+  auto PGOIter = AddrToPGOAnalysisMap.find(StartAddress);
+
+  for (size_t I = 0; I < Iter->second.getBBEntries().size(); ++I) {
+    const BBAddrMap::BBEntry &BBEntry = Iter->second.getBBEntries()[I];
     uint64_t BBAddress = BBEntry.Offset + Iter->second.getFunctionAddress();
     if (BBAddress >= EndAddress)
       continue;
-    Labels[BBAddress].push_back(("BB" + Twine(BBEntry.ID)).str());
+
+    std::string LabelString = ("BB" + Twine(BBEntry.ID)).str();
+    std::string PGOString;
+
+    if (PGOIter != AddrToPGOAnalysisMap.end())
+      PGOString = constructPGOLabelString(PGOIter->second, I);
+
+    Labels[BBAddress].push_back({LabelString, PGOString});
   }
 }
 
@@ -1637,18 +1694,24 @@ disassembleObject(ObjectFile &Obj, const ObjectFile &DbgObj,
   LLVM_DEBUG(LVP.dump());
 
   std::unordered_map<uint64_t, BBAddrMap> AddrToBBAddrMap;
+  std::unordered_map<uint64_t, PGOAnalysisMap> AddrToPGOAnalysisMap;
   auto ReadBBAddrMap = [&](std::optional<unsigned> SectionIndex =
                                std::nullopt) {
     AddrToBBAddrMap.clear();
     if (const auto *Elf = dyn_cast<ELFObjectFileBase>(&Obj)) {
-      auto BBAddrMapsOrErr = Elf->readBBAddrMap(SectionIndex);
+      std::vector<PGOAnalysisMap> PGOAnalyses;
+      auto BBAddrMapsOrErr = Elf->readBBAddrMap(SectionIndex, &PGOAnalyses);
       if (!BBAddrMapsOrErr) {
         reportWarning(toString(BBAddrMapsOrErr.takeError()), Obj.getFileName());
         return;
       }
-      for (auto &FunctionBBAddrMap : *BBAddrMapsOrErr)
-        AddrToBBAddrMap.emplace(FunctionBBAddrMap.Addr,
-                                std::move(FunctionBBAddrMap));
+      for (const auto &[FunctionBBAddrMap, FunctionPGOAnalysis] :
+           zip_equal(*std::move(BBAddrMapsOrErr), std::move(PGOAnalyses))) {
+        uint64_t Addr = FunctionBBAddrMap.Addr;
+        AddrToBBAddrMap.emplace(Addr, std::move(FunctionBBAddrMap));
+        if (FunctionPGOAnalysis.FeatEnable.anyEnabled())
+          AddrToPGOAnalysisMap.emplace(Addr, std::move(FunctionPGOAnalysis));
+      }
     }
   };
 
@@ -1977,14 +2040,15 @@ disassembleObject(ObjectFile &Obj, const ObjectFile &DbgObj,
         FOS.SetUnbuffered();
 
       std::unordered_map<uint64_t, std::string> AllLabels;
-      std::unordered_map<uint64_t, std::vector<std::string>> BBAddrMapLabels;
+      std::unordered_map<uint64_t, std::vector<BBAddrMapLabel>> BBAddrMapLabels;
       if (SymbolizeOperands) {
         collectLocalBranchTargets(Bytes, DT->InstrAnalysis.get(),
                                   DT->DisAsm.get(), DT->InstPrinter.get(),
                                   PrimaryTarget.SubtargetInfo.get(),
                                   SectionAddr, Index, End, AllLabels);
-        collectBBAddrMapLabels(AddrToBBAddrMap, SectionAddr, Index, End,
-                               BBAddrMapLabels);
+        collectBBAddrMapLabels(AddrToBBAddrMap, AddrToPGOAnalysisMap,
+                               SectionAddr, Index, End, BBAddrMapLabels,
+                               FileName);
       }
 
       if (DT->InstrAnalysis)
@@ -2082,8 +2146,9 @@ disassembleObject(ObjectFile &Obj, const ObjectFile &DbgObj,
           // Print local label if there's any.
           auto Iter1 = BBAddrMapLabels.find(SectionAddr + Index);
           if (Iter1 != BBAddrMapLabels.end()) {
-            for (StringRef Label : Iter1->second)
-              FOS << "<" << Label << ">:\n";
+            for (const auto &BBLabel : Iter1->second)
+              FOS << "<" << BBLabel.BlockLabel << ">" << BBLabel.PGOAnalysis
+                  << ":\n";
           } else {
             auto Iter2 = AllLabels.find(SectionAddr + Index);
             if (Iter2 != AllLabels.end())
@@ -2260,7 +2325,7 @@ disassembleObject(ObjectFile &Obj, const ObjectFile &DbgObj,
                 } else if (!Disp) {
                   *TargetOS << TargetName;
                 } else if (BBAddrMapLabelAvailable) {
-                  *TargetOS << BBAddrMapLabels[Target].front();
+                  *TargetOS << BBAddrMapLabels[Target].front().BlockLabel;
                 } else if (LabelAvailable) {
                   *TargetOS << AllLabels[Target];
                 } else {
@@ -2276,7 +2341,8 @@ disassembleObject(ObjectFile &Obj, const ObjectFile &DbgObj,
                 }
 
               } else if (BBAddrMapLabelAvailable) {
-                *TargetOS << " <" << BBAddrMapLabels[Target].front() << ">";
+                *TargetOS << " <" << BBAddrMapLabels[Target].front().BlockLabel
+                          << ">";
               } else if (LabelAvailable) {
                 *TargetOS << " <" << AllLabels[Target] << ">";
               }