[SHT_LLVM_BB_ADDR_MAP] Encode and decode callsite offsets in a newly-introduced SHT_LLVM_BB_ADDR_MAP version.

llvm/include/llvm/Object/ELFTypes.h

@@ -831,6 +831,7 @@ struct BBAddrMap {
     bool BrProb : 1;
     bool MultiBBRange : 1;
     bool OmitBBEntries : 1;
+    bool CallsiteOffsets : 1;
 
     bool hasPGOAnalysis() const { return FuncEntryCount || BBFreq || BrProb; }
 
@@ -842,7 +843,8 @@ struct BBAddrMap {
              (static_cast<uint8_t>(BBFreq) << 1) |
              (static_cast<uint8_t>(BrProb) << 2) |
              (static_cast<uint8_t>(MultiBBRange) << 3) |
-             (static_cast<uint8_t>(OmitBBEntries) << 4);
+             (static_cast<uint8_t>(OmitBBEntries) << 4) |
+             (static_cast<uint8_t>(CallsiteOffsets) << 5);
     }
 
     // Decodes from minimum bit width representation and validates no
@@ -851,7 +853,7 @@ struct BBAddrMap {
       Features Feat{
           static_cast<bool>(Val & (1 << 0)), static_cast<bool>(Val & (1 << 1)),
           static_cast<bool>(Val & (1 << 2)), static_cast<bool>(Val & (1 << 3)),
-          static_cast<bool>(Val & (1 << 4))};
+          static_cast<bool>(Val & (1 << 4)), static_cast<bool>(Val & (1 << 5))};
       if (Feat.encode() != Val)
         return createStringError(
             std::error_code(), "invalid encoding for BBAddrMap::Features: 0x%x",
@@ -861,9 +863,10 @@ struct BBAddrMap {
 
     bool operator==(const Features &Other) const {
       return std::tie(FuncEntryCount, BBFreq, BrProb, MultiBBRange,
-                      OmitBBEntries) ==
+                      OmitBBEntries, CallsiteOffsets) ==
              std::tie(Other.FuncEntryCount, Other.BBFreq, Other.BrProb,
-                      Other.MultiBBRange, Other.OmitBBEntries);
+                      Other.MultiBBRange, Other.OmitBBEntries,
+                      Other.CallsiteOffsets);
     }
   };
 
@@ -914,13 +917,18 @@ struct BBAddrMap {
     uint32_t Size = 0;   // Size of the basic block.
     Metadata MD = {false, false, false, false,
                    false}; // Metdata for this basic block.
+    // Offsets of callsites (end of call instructions), relative to the basic
+    // block start.
+    SmallVector<uint32_t, 1> CallsiteOffsets;
 
-    BBEntry(uint32_t ID, uint32_t Offset, uint32_t Size, Metadata MD)
-        : ID(ID), Offset(Offset), Size(Size), MD(MD){};
+    BBEntry(uint32_t ID, uint32_t Offset, uint32_t Size, Metadata MD,
+            SmallVector<uint32_t, 1> CallsiteOffsets)
+        : ID(ID), Offset(Offset), Size(Size), MD(MD),
+          CallsiteOffsets(std::move(CallsiteOffsets)) {}
 
     bool operator==(const BBEntry &Other) const {
       return ID == Other.ID && Offset == Other.Offset && Size == Other.Size &&
-             MD == Other.MD;
+             MD == Other.MD && CallsiteOffsets == Other.CallsiteOffsets;
     }
 
     bool hasReturn() const { return MD.HasReturn; }

llvm/include/llvm/ObjectYAML/ELFYAML.h

@@ -162,6 +162,7 @@ struct BBAddrMapEntry {
     llvm::yaml::Hex64 AddressOffset;
     llvm::yaml::Hex64 Size;
     llvm::yaml::Hex64 Metadata;
+    std::optional<std::vector<llvm::yaml::Hex64>> CallsiteOffsets;
   };
   uint8_t Version;
   llvm::yaml::Hex8 Feature;
@@ -180,6 +181,20 @@ struct BBAddrMapEntry {
       return 0;
     return BBRanges->front().BaseAddress;
   }
+
+  // Returns if any BB entries have non-empty callsite offsets.
+  bool hasAnyCallsiteOffsets() const {
+    if (!BBRanges)
+      return false;
+    for (const ELFYAML::BBAddrMapEntry::BBRangeEntry &BBR : *BBRanges) {
+      if (!BBR.BBEntries)
+        continue;
+      for (const ELFYAML::BBAddrMapEntry::BBEntry &BBE : *BBR.BBEntries)
+        if (BBE.CallsiteOffsets && !BBE.CallsiteOffsets->empty())
+          return true;
+    }
+    return false;
+  }
 };
 
 struct PGOAnalysisMapEntry {

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -1426,9 +1426,12 @@ getBBAddrMapFeature(const MachineFunction &MF, int NumMBBSectionRanges) {
         "BB entries info is required for BBFreq and BrProb "
         "features");
   }
-  return {FuncEntryCountEnabled, BBFreqEnabled, BrProbEnabled,
+  return {FuncEntryCountEnabled,
+          BBFreqEnabled,
+          BrProbEnabled,
           MF.hasBBSections() && NumMBBSectionRanges > 1,
-          static_cast<bool>(BBAddrMapSkipEmitBBEntries)};
+          static_cast<bool>(BBAddrMapSkipEmitBBEntries),
+          false};
 }
 
 void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {

llvm/lib/Object/ELF.cpp

@@ -837,7 +837,7 @@ decodeBBAddrMapImpl(const ELFFile<ELFT> &EF,
       Version = Data.getU8(Cur);
       if (!Cur)
         break;
-      if (Version > 2)
+      if (Version > 3)
         return createError("unsupported SHT_LLVM_BB_ADDR_MAP version: " +
                            Twine(static_cast<int>(Version)));
       Feature = Data.getU8(Cur); // Feature byte
@@ -847,12 +847,18 @@ decodeBBAddrMapImpl(const ELFFile<ELFT> &EF,
       if (!FeatEnableOrErr)
         return FeatEnableOrErr.takeError();
       FeatEnable = *FeatEnableOrErr;
-      if (Feature != 0 && Version < 2 && Cur)
+      if (FeatEnable.hasPGOAnalysis() && Version < 2)
         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)));
+      if (FeatEnable.CallsiteOffsets && Version < 3)
+        return createError(
+            "version should be >= 3 for SHT_LLVM_BB_ADDR_MAP when "
+            "callsite offsets feature is enabled: version = " +
+            Twine(static_cast<int>(Version)) +
+            " feature = " + Twine(static_cast<int>(Feature)));
     }
     uint32_t NumBlocksInBBRange = 0;
     uint32_t NumBBRanges = 1;
@@ -893,7 +899,23 @@ decodeBBAddrMapImpl(const ELFFile<ELFT> &EF,
                             ? 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);
+          // Read the callsite offsets.
+          uint32_t LastCallsiteOffset = 0;
+          SmallVector<uint32_t, 1> CallsiteOffsets;
+          if (FeatEnable.CallsiteOffsets) {
+            uint32_t NumCallsites =
+                readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
+            CallsiteOffsets.reserve(NumCallsites);
+            for (uint32_t CallsiteIndex = 0;
+                 !ULEBSizeErr && Cur && (CallsiteIndex < NumCallsites);
+                 ++CallsiteIndex) {
+              LastCallsiteOffset +=
+                  readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
+              CallsiteOffsets.push_back(LastCallsiteOffset);
+            }
+          }
+          uint32_t Size = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr) +
+                          LastCallsiteOffset;
           uint32_t MD = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
           if (Version >= 1) {
             // Offset is calculated relative to the end of the previous BB.
@@ -906,7 +928,8 @@ decodeBBAddrMapImpl(const ELFFile<ELFT> &EF,
             MetadataDecodeErr = MetadataOrErr.takeError();
             break;
           }
-          BBEntries.push_back({ID, Offset, Size, *MetadataOrErr});
+          BBEntries.push_back(
+              {ID, Offset, Size, *MetadataOrErr, CallsiteOffsets});
         }
         TotalNumBlocks += BBEntries.size();
       }

llvm/lib/ObjectYAML/ELFEmitter.cpp

@@ -1452,7 +1452,7 @@ void ELFState<ELFT>::writeSectionContent(
   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)
+      if (E.Version > 3)
         WithColor::warning() << "unsupported SHT_LLVM_BB_ADDR_MAP version: "
                              << static_cast<int>(E.Version)
                              << "; encoding using the most recent version";
@@ -1483,6 +1483,8 @@ void ELFState<ELFT>::writeSectionContent(
     if (!E.BBRanges)
       continue;
     uint64_t TotalNumBlocks = 0;
+    bool EmitCallsiteOffsets =
+        FeatureOrErr->CallsiteOffsets || E.hasAnyCallsiteOffsets();
     for (const ELFYAML::BBAddrMapEntry::BBRangeEntry &BBR : *E.BBRanges) {
       // Write the base address of the range.
       CBA.write<uintX_t>(BBR.BaseAddress, ELFT::Endianness);
@@ -1500,6 +1502,15 @@ void ELFState<ELFT>::writeSectionContent(
         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);
+        if (EmitCallsiteOffsets) {
+          size_t NumCallsiteOffsets =
+              BBE.CallsiteOffsets ? BBE.CallsiteOffsets->size() : 0;
+          SHeader.sh_size += CBA.writeULEB128(NumCallsiteOffsets);
+          if (BBE.CallsiteOffsets) {
+            for (uint32_t Offset : *BBE.CallsiteOffsets)
+              SHeader.sh_size += CBA.writeULEB128(Offset);
+          }
+        }
         SHeader.sh_size += CBA.writeULEB128(BBE.Size);
         SHeader.sh_size += CBA.writeULEB128(BBE.Metadata);
       }

llvm/lib/ObjectYAML/ELFYAML.cpp

@@ -1882,6 +1882,7 @@ void MappingTraits<ELFYAML::BBAddrMapEntry::BBEntry>::mapping(
   IO.mapRequired("AddressOffset", E.AddressOffset);
   IO.mapRequired("Size", E.Size);
   IO.mapRequired("Metadata", E.Metadata);
+  IO.mapOptional("CallsiteOffsets", E.CallsiteOffsets);
 }
 
 void MappingTraits<ELFYAML::PGOAnalysisMapEntry>::mapping(

llvm/test/tools/llvm-readobj/ELF/bb-addr-map.test

@@ -49,7 +49,8 @@
 # CHECK-NEXT:           {
 # CHECK-NEXT:             ID: 2
 # CHECK-NEXT:             Offset: 0x3
-# CHECK-NEXT:             Size: 0x4
+# CHECK-NEXT:             Callsite Offsets: [1, 3]
+# CHECK-NEXT:             Size: 0x7
 # CHECK-NEXT:             HasReturn: Yes
 # CHECK-NEXT:             HasTailCall: No
 # CHECK-NEXT:             IsEHPad: Yes
@@ -75,7 +76,7 @@
 # CHECK-NEXT:             HasTailCall: No
 # CHECK-NEXT:             IsEHPad: No
 # CHECK-NEXT:             CanFallThrough: Yes
-# CHECK-NEXT:            HasIndirectBranch: No
+# CHECK-NEXT:             HasIndirectBranch: No
 # CHECK-NEXT:           }
 # CHECK-NEXT:         ]
 # CHECK-NEXT:       }
@@ -143,8 +144,8 @@ Sections:
     ShSize: [[SIZE=<none>]]
     Link:   .text
     Entries:
-      - Version: 2
-        Feature: 0x8
+      - Version: 3
+        Feature: 0x28
         BBRanges:
           - BaseAddress: [[ADDR=0x11111]]
             BBEntries:
@@ -158,6 +159,7 @@ Sections:
                 AddressOffset: 0x3
                 Size:          0x4
                 Metadata:      0x15
+                CallsiteOffsets: [ 0x1 , 0x2 ]
       - Version: 2
         BBRanges:
           - BaseAddress: 0x22222

llvm/test/tools/obj2yaml/ELF/bb-addr-map.yaml

@@ -14,7 +14,7 @@
 # VALID-NEXT:   - Name: .llvm_bb_addr_map
 # VALID-NEXT:     Type: SHT_LLVM_BB_ADDR_MAP
 # VALID-NEXT:     Entries:
-# VALID-NEXT:       - Version: 2
+# VALID-NEXT:       - Version: 3
 # VALID-NEXT:         BBRanges:
 ## The 'BaseAddress' field is omitted when it's zero.
 # VALID-NEXT:           - BBEntries:
@@ -30,15 +30,16 @@
 # VALID-NEXT:               AddressOffset: 0xFFFFFFFFFFFFFFF7
 # VALID-NEXT:               Size:          0xFFFFFFFFFFFFFFF8
 # VALID-NEXT:               Metadata:      0xFFFFFFFFFFFFFFF9
-# VALID-NEXT:       - Version: 2
-# VALID-NEXT:         Feature: 0x8
+# VALID-NEXT:       - Version: 3
+# VALID-NEXT:         Feature: 0x28
 # VALID-NEXT:         BBRanges:
 # VALID-NEXT:           - BaseAddress: 0xFFFFFFFFFFFFFF20
 # VALID-NEXT:             BBEntries:
-# VALID-NEXT:               - ID:            6
-# VALID-NEXT:                 AddressOffset: 0xA
-# VALID-NEXT:                 Size:          0xB
-# VALID-NEXT:                 Metadata:      0xC
+# VALID-NEXT:               - ID:              6
+# VALID-NEXT:                 AddressOffset:   0xA
+# VALID-NEXT:                 Size:            0xB
+# VALID-NEXT:                 Metadata:        0xC
+# VALID-NEXT:                 CallsiteOffsets: [ 0x1, 0x2 ]
 
 --- !ELF
 FileHeader:
@@ -50,7 +51,7 @@ Sections:
     Type:   SHT_LLVM_BB_ADDR_MAP
     ShSize: [[SIZE=<none>]]
     Entries:
-      - Version: 2
+      - Version: 3
         Feature: 0x0
         BBRanges:
           - BaseAddress: 0x0
@@ -67,17 +68,18 @@ Sections:
                 AddressOffset: 0xFFFFFFFFFFFFFFF7
                 Size:          0xFFFFFFFFFFFFFFF8
                 Metadata:      0xFFFFFFFFFFFFFFF9
-      - Version:   2
-        Feature:   0x8
+      - Version:   3
+        Feature:   0x28
         NumBBRanges: [[NUMBBRANGES=<none>]]
         BBRanges:
           - BaseAddress:   0xFFFFFFFFFFFFFF20
             NumBlocks: [[NUMBLOCKS=<none>]]
             BBEntries:
-             - ID:            6
-               AddressOffset: 0xA
-               Size:          0xB
-               Metadata:      0xC
+             - ID:              6
+               AddressOffset:   0xA
+               Size:            0xB
+               Metadata:        0xC
+               CallsiteOffsets: [ 0x1, 0x2 ]
 
 ## Check obj2yaml can dump empty .llvm_bb_addr_map sections.