[BOLT] Calculate input to output address map using BOLTLinker

[jobnoorman]

Summary

BOLT uses MCAsmLayout to calculate the output values of basic blocks.
This means output values are calculated based on a pre-linking state and
any changes to symbol values during linking will cause incorrect values
to be used.

This issue was first addressed in D154604 by adding all basic block
symbols to the symbol table for the linker to resolve them. However, the
runtime overhead of handling this huge symbol table turned out to be
prohibitively large.

This patch solves the issue in a different way. First, a temporary
section containing [input address, output symbol] pairs is emitted to the
intermediary object file. The linker will resolve all these references
so we end up with a section of [input address, output address] pairs.
This section is then parsed and used to:

• Replace BinaryBasicBlock::OffsetTranslationTable
• Replace BinaryFunction::InputOffsetToAddressMap
• Update BinaryBasicBlock::OutputAddressRange

Note that the reason this is more performant than the previous attempt
is that these symbol references do not cause entries to be added to the
symbol table. Instead, section-relative references are used for the
relocations.

[maksfb]

Apologies for the delay. Thanks for adding the map. I think we can eventually get rid of LocSyms, but we can do it in follow-up diffs.

bolt/include/bolt/Core/AddressMap.h

1

Needs a standard license header.

bolt/include/bolt/Profile/BoltAddressTranslation.h

119

Unneeded. Should be able to get BC via BB->BF->BC.

bolt/include/bolt/Rewrite/RewriteInstance.h

384

Maybe add a flag to BinarySection, instead? WDYT?

bolt/lib/Core/AddressMap.cpp

45

You can also use DataExtractor class for parsing and derive endianness from BC.AsmInfo->isLittleEndian().

bolt/lib/Rewrite/RewriteInstance.cpp

3194

Can we make this call from inside emitBinaryContext()?

4032

Do we need this check and the one below despite calling deregisterSection()?

[jobnoorman]

Address comments from @maksfb. Major changes:

• Use DataExtractor to parse address map
• Add IsLinkOnly flag to BinarySection

[jobnoorman]

bolt/lib/Core/AddressMap.cpp

45

TIL about DataExtractor. Thanks!

bolt/lib/Rewrite/RewriteInstance.cpp

3194

Yes, the reason I didn't do this is because it currently doesn't work for Mach-O (because BinaryContext::getDataSection doesn't support it). I moved the call but guarded it by isELF().

4032

We don't :)

[maksfb]

LGTM

bolt/include/bolt/Core/AddressMap.h

+#ifndef BOLT_CORE_ADDRESS_MAP_H
+#define BOLT_CORE_ADDRESS_MAP_H
+
+#include "llvm/ADT/StringRef.h"
+
+#include <optional>
+#include <unordered_map>
+
+namespace llvm {
+
+class MCStreamer;
+
+namespace bolt {
+
+class BinaryContext;
+
+class AddressMap {
+  using MapTy = std::unordered_multimap<uint64_t, uint64_t>;
+  MapTy Map;
+
+public:
+  static const char *const SectionName;
+
+  static void emit(MCStreamer &Streamer, BinaryContext &BC);
+  static AddressMap parse(StringRef Buffer);
+
+  std::optional<uint64_t> lookup(uint64_t InputAddress) const {
+    auto It = Map.find(InputAddress);
+    if (It != Map.end())
+      return It->second;
+    return std::nullopt;
+  }
+
+  std::pair<MapTy::const_iterator, MapTy::const_iterator>
+  lookupAll(uint64_t InputAddress) const {
+    return Map.equal_range(InputAddress);
+  }
+};
+
+} // namespace bolt
+} // namespace llvm
+
+#endif

bolt/include/bolt/Core/BinaryBasicBlock.h

private:
   std::pair<uint32_t, uint32_t> InputRange = {INVALID_OFFSET, INVALID_OFFSET};
 
   /// Map input offset (from function start) of an instruction to an output
   /// symbol. Enables writing BOLT address translation tables used for mapping
   /// control transfer in the output binary back to the original binary.
   using LocSymsTy = std::vector<std::pair<uint32_t, const MCSymbol *>>;
   std::unique_ptr<LocSymsTy> LocSyms;
 
-  /// After output/codegen, map output offsets of instructions in this basic
-  /// block to instruction offsets in the original function. Note that the
-  /// output basic block could be different from the input basic block.
-  /// We only map instruction of interest, such as calls and markers.
-  ///
-  /// We store the offset array in a basic block to facilitate BAT tables
-  /// generation. Otherwise, the mapping could be done at function level.
-  using OffsetTranslationTableTy = std::vector<std::pair<uint32_t, uint32_t>>;
-  std::unique_ptr<OffsetTranslationTableTy> OffsetTranslationTable;
-
   /// Alignment requirements for the block.
   uint32_t Alignment{1};
 
   /// Maximum number of bytes to use for alignment of the block.
   uint32_t AlignmentMaxBytes{0};
 
   /// Number of times this basic block was executed.
   uint64_t ExecutionCount{COUNT_NO_PROFILE};
std::pair<uint64_t, uint64_t> getInputAddressRange() const {
     return InputRange;
   }
 
   /// Gets the memory address range of this BB in the output binary.
   std::pair<uint64_t, uint64_t> getOutputAddressRange() const {
     return OutputAddressRange;
   }
 
-  /// Update addresses of special instructions inside this basic block.
-  void updateOutputValues(const MCAsmLayout &Layout);
+  bool hasLocSyms() const { return LocSyms != nullptr; }
 
   /// Return mapping of input offsets to symbols in the output.
   LocSymsTy &getLocSyms() {
     return LocSyms ? *LocSyms : *(LocSyms = std::make_unique<LocSymsTy>());
   }
 
   /// Return mapping of input offsets to symbols in the output.
   const LocSymsTy &getLocSyms() const {
     return const_cast<BinaryBasicBlock *>(this)->getLocSyms();
   }
 
-  /// Return offset translation table for the basic block.
-  OffsetTranslationTableTy &getOffsetTranslationTable() {
-    return OffsetTranslationTable
-               ? *OffsetTranslationTable
-               : *(OffsetTranslationTable =
-                       std::make_unique<OffsetTranslationTableTy>());
-  }
-
-  /// Return offset translation table for the basic block.
-  const OffsetTranslationTableTy &getOffsetTranslationTable() const {
-    return const_cast<BinaryBasicBlock *>(this)->getOffsetTranslationTable();
-  }
-
   /// Return size of the basic block in the output binary.
   uint64_t getOutputSize() const {
     return OutputAddressRange.second - OutputAddressRange.first;
   }
 
   BinaryFunction *getFunction() const { return Function; }
 
   /// Analyze and interpret the terminators of this basic block. TBB must be

bolt/include/bolt/Core/BinaryContext.h

 //===- bolt/Core/BinaryContext.h - Low-level context ------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Context for processing binary executable/library files.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef BOLT_CORE_BINARY_CONTEXT_H
 #define BOLT_CORE_BINARY_CONTEXT_H
 
+#include "bolt/Core/AddressMap.h"
 #include "bolt/Core/BinaryData.h"
 #include "bolt/Core/BinarySection.h"
 #include "bolt/Core/DebugData.h"
 #include "bolt/Core/JumpTable.h"
 #include "bolt/Core/MCPlusBuilder.h"
 #include "bolt/RuntimeLibs/RuntimeLibrary.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringMap.h"
class BinaryContext {
 
   /// A map of DWO Ids to CUs.
   using DWOIdToCUMapType = std::unordered_map<uint64_t, DWARFUnit *>;
   DWOIdToCUMapType DWOCUs;
 
   bool ContainsDwarf5{false};
   bool ContainsDwarfLegacy{false};
 
+  /// Mapping from input to output addresses.
+  std::optional<AddressMap> IOAddressMap;
+
   /// Preprocess DWO debug information.
   void preprocessDWODebugInfo();
 
   /// DWARF line info for CUs.
   std::map<unsigned, DwarfLineTable> DwarfLineTablesCUMap;
 
   /// Internal helper for removing section name from a lookup table.
   void deregisterSectionName(const BinarySection &Section);
createStreamer(llvm::raw_pwrite_stream &OS) const {
     std::unique_ptr<MCStreamer> Streamer(TheTarget->createMCObjectStreamer(
         *TheTriple, *Ctx, std::unique_ptr<MCAsmBackend>(MAB), std::move(OW),
         std::unique_ptr<MCCodeEmitter>(MCE), *STI,
         /* RelaxAll */ false,
         /* IncrementalLinkerCompatible */ false,
         /* DWARFMustBeAtTheEnd */ false));
     return Streamer;
   }
+
+  void setIOAddressMap(AddressMap Map) { IOAddressMap = std::move(Map); }
+  const AddressMap &getIOAddressMap() const {
+    assert(IOAddressMap && "Address map not set yet");
+    return *IOAddressMap;
+  }
 };
 
 template <typename T, typename = std::enable_if_t<sizeof(T) == 1>>
 inline raw_ostream &operator<<(raw_ostream &OS, const ArrayRef<T> &ByteArray) {
   const char *Sep = "";
   for (const auto Byte : ByteArray) {
     OS << Sep << format("%.2x", Byte);
     Sep = " ";
   }
   return OS;
 }
 
 } // namespace bolt
 } // namespace llvm
 
 #endif

bolt/include/bolt/Core/BinaryFunction.h

private:
   mutable SmallVector<MCSymbol *, 0> FunctionEndLabels;
 
   /// Unique number associated with the function.
   uint64_t FunctionNumber;
 
   /// Count the number of functions created.
   static uint64_t Count;
 
-  /// Map offsets of special instructions to addresses in the output.
-  InputOffsetToAddressMapTy InputOffsetToAddressMap;
-
   /// Register alternative function name.
   void addAlternativeName(std::string NewName) {
     Aliases.push_back(std::move(NewName));
   }
 
   /// Return a label at a given \p Address in the function. If the label does
   /// not exist - create it. Assert if the \p Address does not belong to
   /// the function. If \p CreatePastEnd is true, then return the function
void setPLTSymbol(const MCSymbol *Symbol) {
     assert(Size == 0 && "function size should be 0 for PLT functions");
     PLTSymbol = Symbol;
     IsPseudo = true;
   }
 
   /// Update output values of the function based on the final \p Layout.
   void updateOutputValues(const MCAsmLayout &Layout);
 
-  /// Return mapping of input to output addresses. Most users should call
-  /// translateInputToOutputAddress() for address translation.
-  InputOffsetToAddressMapTy &getInputOffsetToAddressMap() {
-    assert(isEmitted() && "cannot use address mapping before code emission");
-    return InputOffsetToAddressMap;
-  }
-
   /// Register relocation type \p RelType at a given \p Address in the function
   /// against \p Symbol.
   /// Assert if the \p Address is not inside this function.
   void addRelocation(uint64_t Address, MCSymbol *Symbol, uint64_t RelType,
                      uint64_t Addend, uint64_t Value);
 
   /// Return the name of the section this function originated from.
   std::optional<StringRef> getOriginSectionName() const {

   /// After reordering, this function checks the state of CFI and fixes it if it
   /// is corrupted. If it is unable to fix it, it returns false.
   bool finalizeCFIState();
 
   /// Return true if this function needs an address-transaltion table after
   /// its code emission.
   bool requiresAddressTranslation() const;
 
+  /// Return true if the linker needs to generate an address map for this
+  /// function. Used for keeping track of the mapping from input to out
+  /// addresses of basic blocks.
+  bool requiresAddressMap() const;
+
   /// Adjust branch instructions to match the CFG.
   ///
   /// As it comes to internal branches, the CFG represents "the ultimate source
   /// of truth". Transformations on functions and blocks have to update the CFG
   /// and fixBranches() would make sure the correct branch instructions are
   /// inserted at the end of basic blocks.
   ///
   /// We do require a conditional branch at the end of the basic block if

bolt/include/bolt/Profile/BoltAddressTranslation.h

public:
   bool enabledFor(llvm::object::ELFObjectFileBase *InputFile) const;
 
 private:
   /// Helper to update \p Map by inserting one or more BAT entries reflecting
   /// \p BB for function located at \p FuncAddress. At least one entry will be
   /// emitted for the start of the BB. More entries may be emitted to cover
   /// the location of calls or any instruction that may change control flow.
   void writeEntriesForBB(MapTy &Map, const BinaryBasicBlock &BB,
-                         uint64_t FuncAddress);
+                         uint64_t FuncAddress, const BinaryContext &BC);
 
   std::map<uint64_t, MapTy> Maps;
 
   /// Links outlined cold bocks to their original function
   std::map<uint64_t, uint64_t> ColdPartSource;
 
   /// Identifies the address of a control-flow changing instructions in a
   /// translation map entry
   const static uint32_t BRANCHENTRY = 0x80000000;
 };
 } // namespace bolt
 
 } // namespace llvm
 
 #endif

bolt/include/bolt/Rewrite/RewriteInstance.h

public:
   static std::vector<std::string> DebugSectionsToOverwrite;
 
   /// Return true if the section holds debug information.
   static bool isDebugSection(StringRef SectionName);
 
   /// Return true if the section holds linux kernel symbol information.
   static bool isKSymtabSection(StringRef SectionName);
 
+  /// Return true if the section is only used while linking and will not appear
+  /// in the output file.
+  static bool isLinkOnlySection(StringRef SectionName);
+
   /// Adds Debug section to overwrite.
   static void addToDebugSectionsToOverwrite(const char *Section) {
     DebugSectionsToOverwrite.emplace_back(Section);
   }
 
 private:
   /// Manage a pipeline of metadata handlers.
   class MetadataManager MetadataManager;

bolt/lib/Core/AddressMap.cpp

+#include "bolt/Core/AddressMap.h"
+#include "bolt/Core/BinaryContext.h"
+#include "bolt/Core/BinaryFunction.h"
+#include "llvm/MC/MCStreamer.h"
+
+namespace llvm {
+namespace bolt {
+
+const char *const AddressMap::SectionName = ".bolt.address_map";
+
+static void emitLabel(MCStreamer &Streamer, uint64_t InputAddress,
+                      const MCSymbol *OutputLabel) {
+  Streamer.emitIntValue(InputAddress, 8);
+  Streamer.emitSymbolValue(OutputLabel, 8);
+}
+
+void AddressMap::emit(MCStreamer &Streamer, BinaryContext &BC) {
+  Streamer.switchSection(BC.getDataSection(SectionName));
+
+  for (const auto &[BFAddress, BF] : BC.getBinaryFunctions()) {
+    if (!BF.requiresAddressMap())
+      continue;
+
+    for (const auto &BB : BF) {
+      if (!BB.getLabel()->isDefined())
+        continue;
+
+      emitLabel(Streamer, BFAddress + BB.getInputAddressRange().first,
+                BB.getLabel());
+
+      if (!BB.hasLocSyms())
+        continue;
+
+      for (auto [Offset, Symbol] : BB.getLocSyms())
+        emitLabel(Streamer, BFAddress + Offset, Symbol);
+    }
+  }
+}
+
+AddressMap AddressMap::parse(StringRef Buffer) {
+  assert(Buffer.size() % 16 == 0 && "Unexpected address map size");
+
+  AddressMap Parsed;
+  Parsed.Map.reserve(Buffer.size() / 16);
+  const auto *const MapData = Buffer.data();
+
+  for (size_t I = 0; I < Buffer.size(); I += 16) {
+    const auto Input = support::endian::read64le(MapData + I);
+    const auto Output = support::endian::read64le(MapData + I + 8);
+    Parsed.Map.insert({Input, Output});
+  }
+
+  return Parsed;
+}
+
+} // namespace bolt
+} // namespace llvm

bolt/lib/Core/BinaryBasicBlock.cpp

BinaryBasicBlock *BinaryBasicBlock::splitAt(iterator II) {
   // Move instructions over.
   adjustNumPseudos(II, end(), -1);
   NewBlock->addInstructions(II, end());
   Instructions.erase(II, end());
 
   return NewBlock;
 }
 
-void BinaryBasicBlock::updateOutputValues(const MCAsmLayout &Layout) {
-  if (!LocSyms)
-    return;
-
-  const uint64_t BBAddress = getOutputAddressRange().first;
-  const uint64_t BBOffset = Layout.getSymbolOffset(*getLabel());
-  for (const auto &LocSymKV : *LocSyms) {
-    const uint32_t InputFunctionOffset = LocSymKV.first;
-    const uint32_t OutputOffset = static_cast<uint32_t>(
-        Layout.getSymbolOffset(*LocSymKV.second) - BBOffset);
-    getOffsetTranslationTable().emplace_back(
-        std::make_pair(OutputOffset, InputFunctionOffset));
-
-    // Update reverse (relative to BAT) address lookup table for function.
-    if (getFunction()->requiresAddressTranslation()) {
-      getFunction()->getInputOffsetToAddressMap().emplace(
-          std::make_pair(InputFunctionOffset, OutputOffset + BBAddress));
-    }
-  }
-  LocSyms.reset(nullptr);
-}
-
 } // namespace bolt
 } // namespace llvm

bolt/lib/Core/BinaryFunction.cpp

bool BinaryFunction::finalizeCFIState() {
 
   return true;
 }
 
 bool BinaryFunction::requiresAddressTranslation() const {
   return opts::EnableBAT || hasSDTMarker() || hasPseudoProbe();
 }
 
+bool BinaryFunction::requiresAddressMap() const {
+  if (isInjected())
+    return false;
+
+  return opts::UpdateDebugSections || isMultiEntry() ||
+         requiresAddressTranslation();
+}
+
 uint64_t BinaryFunction::getInstructionCount() const {
   uint64_t Count = 0;
   for (const BinaryBasicBlock &BB : blocks())
     Count += BB.getNumNonPseudos();
   return Count;
 }
 
 void BinaryFunction::clearDisasmState() {
for (BinaryBasicBlock *const BB : FF) {
       assert(BB->getLabel()->isDefined() && "symbol should be defined");
       if (!BC.HasRelocations) {
         if (BB->isSplit())
           assert(FragmentBaseAddress == FF.getAddress());
         else
           assert(FragmentBaseAddress == getOutputAddress());
       }
 
-      const uint64_t BBOffset = Layout.getSymbolOffset(*BB->getLabel());
-      const uint64_t BBAddress = FragmentBaseAddress + BBOffset;
+      const uint64_t BBAddress =
+          *BC.getIOAddressMap().lookup(BB->getInputOffset() + getAddress());
       BB->setOutputStartAddress(BBAddress);
 
       if (PrevBB)
         PrevBB->setOutputEndAddress(BBAddress);
       PrevBB = BB;
-
-      BB->updateOutputValues(Layout);
     }
 
     PrevBB->setOutputEndAddress(PrevBB->isSplit()
                                     ? FF.getAddress() + FF.getImageSize()
                                     : getOutputAddress() + getOutputSize());
   }
 }
 
uint64_t BinaryFunction::translateInputToOutputAddress(uint64_t Address) const {
   if (!isEmitted())
     return Address;
 
   if (Address < getAddress())
     return 0;
 
   // Check if the address is associated with an instruction that is tracked
   // by address translation.
-  auto KV = InputOffsetToAddressMap.find(Address - getAddress());
-  if (KV != InputOffsetToAddressMap.end())
-    return KV->second;
+  if (auto OutputAddress = BC.getIOAddressMap().lookup(Address))
+    return *OutputAddress;
 
   // FIXME: #18950828 - we rely on relative offsets inside basic blocks to stay
   //        intact. Instead we can use pseudo instructions and/or annotations.
   const uint64_t Offset = Address - getAddress();
   const BinaryBasicBlock *BB = getBasicBlockContainingOffset(Offset);
   if (!BB) {
     // Special case for address immediately past the end of the function.
     if (Offset == getSize())

bolt/lib/Core/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   DebugInfoDWARF
   Demangle
   MC
   MCDisassembler
   Object
   Support
   BinaryFormat
   AsmPrinter
   TargetParser
   )
 
 add_llvm_library(LLVMBOLTCore
+  AddressMap.cpp
   BinaryBasicBlock.cpp
   BinaryContext.cpp
   BinaryData.cpp
   BinaryEmitter.cpp
   BinaryFunction.cpp
   BinaryFunctionProfile.cpp
   BinarySection.cpp
   DebugData.cpp

bolt/lib/Profile/BoltAddressTranslation.cpp

 
 namespace llvm {
 namespace bolt {
 
 const char *BoltAddressTranslation::SECTION_NAME = ".note.bolt_bat";
 
 void BoltAddressTranslation::writeEntriesForBB(MapTy &Map,
                                                const BinaryBasicBlock &BB,
-                                               uint64_t FuncAddress) {
+                                               uint64_t FuncAddress,
+                                               const BinaryContext &BC) {
   const uint64_t BBOutputOffset =
       BB.getOutputAddressRange().first - FuncAddress;
   const uint32_t BBInputOffset = BB.getInputOffset();
 
   // Every output BB must track back to an input BB for profile collection
   // in bolted binaries. If we are missing an offset, it means this block was
   // created by a pass. We will skip writing any entries for it, and this means
   // any traffic happening in this block will map to the previous block in the
void BoltAddressTranslation::writeEntriesForBB(MapTy &Map,
   // update takes precedence. Of course it is not ideal to have conflicts and
   // those happen when we have an empty BB that either contained only
   // NOPs or a jump to the next block (successor). Either way, the successor
   // and this deleted block will both share the same output address (the same
   // key), and we need to map back. We choose here to privilege the successor by
   // allowing it to overwrite the previously inserted key in the map.
   Map[BBOutputOffset] = BBInputOffset;
 
-  for (const auto &IOPair : BB.getOffsetTranslationTable()) {
-    const uint64_t OutputOffset = IOPair.first + BBOutputOffset;
-    const uint32_t InputOffset = IOPair.second;
+  for (const auto &[InputOffset, Sym] : BB.getLocSyms()) {
+    const auto InputAddress = BB.getFunction()->getAddress() + InputOffset;
+    const auto OutputAddress = BC.getIOAddressMap().lookup(InputAddress);
+    assert(OutputAddress && "Unknown instruction address");
+    const auto OutputOffset = *OutputAddress - FuncAddress;
 
     // Is this the first instruction in the BB? No need to duplicate the entry.
     if (OutputOffset == BBOutputOffset)
       continue;
 
     LLVM_DEBUG(dbgs() << "  Key: " << Twine::utohexstr(OutputOffset) << " Val: "
                       << Twine::utohexstr(InputOffset) << " (branch)\n");
     Map.insert(
for (auto &BFI : BC.getBinaryFunctions()) {
 
     LLVM_DEBUG(dbgs() << "Function name: " << Function.getPrintName() << "\n");
     LLVM_DEBUG(dbgs() << " Address reference: 0x"
                       << Twine::utohexstr(Function.getOutputAddress()) << "\n");
 
     MapTy Map;
     for (const BinaryBasicBlock *const BB :
          Function.getLayout().getMainFragment())
-      writeEntriesForBB(Map, *BB, Function.getOutputAddress());
+      writeEntriesForBB(Map, *BB, Function.getOutputAddress(), BC);
     Maps.emplace(Function.getOutputAddress(), std::move(Map));
 
     if (!Function.isSplit())
       continue;
 
     // Split maps
     LLVM_DEBUG(dbgs() << " Cold part\n");
     for (const FunctionFragment &FF :
          Function.getLayout().getSplitFragments()) {
       Map.clear();
       for (const BinaryBasicBlock *const BB : FF)
-        writeEntriesForBB(Map, *BB, FF.getAddress());
+        writeEntriesForBB(Map, *BB, FF.getAddress(), BC);
 
       Maps.emplace(FF.getAddress(), std::move(Map));
       ColdPartSource.emplace(FF.getAddress(), Function.getOutputAddress());
     }
   }
 
   const uint32_t NumFuncs = Maps.size();
   OS.write(reinterpret_cast<const char *>(&NumFuncs), 4);

bolt/lib/Rewrite/PseudoProbeRewriter.cpp

for (auto &AP : Address2ProbesMap) {
     std::list<MCDecodedPseudoProbe>::iterator Probe = AP.second.begin();
     while (ProbeTrack != 0) {
       if (Probe->isBlock()) {
         Probe->setAddress(BlkOutputAddress);
       } else if (Probe->isCall()) {
         // A call probe may be duplicated due to ICP
         // Go through output of InputOffsetToAddressMap to collect all related
         // probes
-        const InputOffsetToAddressMapTy &Offset2Addr =
-            F->getInputOffsetToAddressMap();
-        auto CallOutputAddresses = Offset2Addr.equal_range(Offset);
+        auto CallOutputAddresses = BC.getIOAddressMap().lookupAll(AP.first);
         auto CallOutputAddress = CallOutputAddresses.first;
         if (CallOutputAddress == CallOutputAddresses.second) {
           Probe->setAddress(INT64_MAX);
         } else {
           Probe->setAddress(CallOutputAddress->second);
           CallOutputAddress = std::next(CallOutputAddress);
         }
 

bolt/lib/Rewrite/RewriteInstance.cpp

 //===- bolt/Rewrite/RewriteInstance.cpp - ELF rewriter --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Rewrite/RewriteInstance.h"
+#include "bolt/Core/AddressMap.h"
 #include "bolt/Core/BinaryContext.h"
 #include "bolt/Core/BinaryEmitter.h"
 #include "bolt/Core/BinaryFunction.h"
 #include "bolt/Core/DebugData.h"
 #include "bolt/Core/Exceptions.h"
 #include "bolt/Core/FunctionLayout.h"
 #include "bolt/Core/MCPlusBuilder.h"
 #include "bolt/Core/ParallelUtilities.h"
if (opts::UseOldText || opts::StrictMode) {
       EHFrameSection->clearContents();
     } else {
       // Make .eh_frame relocatable.
       relocateEHFrameSection();
     }
   }
 
   emitBinaryContext(*Streamer, *BC, getOrgSecPrefix());
+  AddressMap::emit(*Streamer, *BC);
 
   Streamer->finish();
   if (Streamer->getContext().hadError()) {
     errs() << "BOLT-ERROR: Emission failed.\n";
     exit(1);
   }
 
   if (opts::KeepTmp) {
for (uint8_t SType = ST_READONLY; SType <= ST_READWRITE; ++SType) {
     const uint64_t LastNextAvailableAddress = NextAvailableAddress;
     if (SType == ST_READWRITE) {
       // Align R+W segment to regular page size
       NextAvailableAddress = alignTo(NextAvailableAddress, BC->RegularPageSize);
       NewWritableSegmentAddress = NextAvailableAddress;
     }
 
     for (BinarySection &Section : BC->allocatableSections()) {
+      if (isLinkOnlySection(Section.getName()))
+        continue;
+
       if (!Section.hasValidSectionID())
         continue;
 
       if (Section.isWritable() == (SType == ST_READONLY))
         continue;
 
       if (Section.getOutputAddress()) {
         LLVM_DEBUG({
if (SType == ST_READONLY) {
       // Restore NextAvailableAddress if no new writable sections
       if (!NewWritableSegmentSize)
         NextAvailableAddress = LastNextAvailableAddress;
     }
   }
 }
 
 void RewriteInstance::updateOutputValues(const MCAsmLayout &Layout) {
+  if (auto MapSection = BC->getUniqueSectionByName(AddressMap::SectionName)) {
+    auto Map = AddressMap::parse(MapSection->getOutputContents());
+    BC->setIOAddressMap(std::move(Map));
+    BC->deregisterSection(*MapSection);
+  }
+
   for (BinaryFunction *Function : BC->getAllBinaryFunctions())
     Function->updateOutputValues(Layout);
 }
 
 void RewriteInstance::patchELFPHDRTable() {
   auto ELF64LEFile = cast<ELF64LEObjectFile>(InputFile);
   const ELFFile<ELF64LE> &Obj = ELF64LEFile->getELFFile();
   raw_fd_ostream &OS = Out->os();
for (BinarySection &Section : BC->allocatableSections()) {
 
     if (Section.hasSectionRef() || Section.isAnonymous()) {
       if (opts::Verbosity)
         outs() << "BOLT-INFO: not writing section header for section "
                << Section.getOutputName() << '\n';
       continue;
     }
 
+    if (isLinkOnlySection(Section.getName()))
+      continue;
+
     if (opts::Verbosity >= 1)
       outs() << "BOLT-INFO: writing section header for "
              << Section.getOutputName() << '\n';
     ELFShdrTy NewSection;
     NewSection.sh_type = ELF::SHT_PROGBITS;
     NewSection.sh_addr = Section.getOutputAddress();
     NewSection.sh_offset = Section.getOutputFileOffset();
     NewSection.sh_size = Section.getOutputSize();
if (BC->HasRelocations && opts::TrapOldCode) {
     }
     OS.seek(SavedPos);
   }
 
   // Write all allocatable sections - reloc-mode text is written here as well
   for (BinarySection &Section : BC->allocatableSections()) {
     if (!Section.isFinalized() || !Section.getOutputData())
       continue;
+    if (isLinkOnlySection(Section.getName()))
+      continue;
 
     if (opts::Verbosity >= 1)
       outs() << "BOLT: writing new section " << Section.getName()
              << "\n data at 0x" << Twine::utohexstr(Section.getAllocAddress())
              << "\n of size " << Section.getOutputSize() << "\n at offset "
              << Section.getOutputFileOffset() << '\n';
     OS.pwrite(reinterpret_cast<const char *>(Section.getOutputData()),
               Section.getOutputSize(), Section.getOutputFileOffset());

 }
 
 bool RewriteInstance::isKSymtabSection(StringRef SectionName) {
   if (SectionName.startswith("__ksymtab"))
     return true;
 
   return false;
 }
+
+bool RewriteInstance::isLinkOnlySection(StringRef SectionName) {
+  if (SectionName == AddressMap::SectionName)
+    return true;
+
+  return false;
+}