[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

2

Needs a standard license header.

bolt/include/bolt/Profile/BoltAddressTranslation.h

119 ↗

(On Diff #541549)

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

bolt/include/bolt/Rewrite/RewriteInstance.h

384 ↗

(On Diff #541549)

Maybe add a flag to BinarySection, instead? WDYT?

bolt/lib/Core/AddressMap.cpp

46

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

bolt/lib/Rewrite/RewriteInstance.cpp

3203

Can we make this call from inside emitBinaryContext()?

4040

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

46

TIL about DataExtractor. Thanks!

bolt/lib/Rewrite/RewriteInstance.cpp

3203

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().

4040

We don't :)

[maksfb]

LGTM

bolt/include/bolt/Core/AddressMap.h

+//===- bolt/Core/AddressMap.h - Input-output address map --------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Helper class to create a mapping from input to output addresses needed for
+// updating debugging symbols and BAT. We emit an MCSection containing
+// <Input address, Output MCSymbol> pairs to the object file and JITLink will
+// transform this in <Input address, Output address> pairs. The linker output
+// can then be parsed and used to establish the mapping.
+//
+//===----------------------------------------------------------------------===//
+//
+#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, const BinaryContext &BC);
+
+  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/Core/BinarySection.h

class BinarySection {
   StringRef OutputContents;        // Rewritten section contents.
   const uint64_t SectionNumber;    // Order in which the section was created.
   std::string SectionID;           // Unique ID used for address mapping.
                                    // Set by ExecutableFileMemoryManager.
   uint32_t Index{0};               // Section index in the output file.
   mutable bool IsReordered{false}; // Have the contents been reordered?
   bool IsAnonymous{false};         // True if the name should not be included
                                    // in the output file.
+  bool IsLinkOnly{false};          // True if the section should not be included
+                                   // in the output file.
 
   uint64_t hash(const BinaryData &BD,
                 std::map<const BinaryData *, uint64_t> &Cache) const;
 
   // non-copyable
   BinarySection(const BinarySection &) = delete;
   BinarySection(BinarySection &&) = delete;
   BinarySection &operator=(const BinarySection &) = delete;
public:
   void setOutputFileOffset(uint64_t Offset) { OutputFileOffset = Offset; }
   void setSectionID(StringRef ID) {
     assert(!hasValidSectionID() && "trying to set section id twice");
     SectionID = ID;
   }
   void setIndex(uint32_t I) { Index = I; }
   void setOutputName(const Twine &Name) { OutputName = Name.str(); }
   void setAnonymous(bool Flag) { IsAnonymous = Flag; }
+  bool isLinkOnly() const { return IsLinkOnly; }
+  void setLinkOnly() { IsLinkOnly = true; }
 
   /// Emit the section as data, possibly with relocations.
   /// Use name \p SectionName for the section during the emission.
   void emitAsData(MCStreamer &Streamer, const Twine &SectionName) const;
 
   using SymbolResolverFuncTy = llvm::function_ref<uint64_t(const MCSymbol *)>;
 
   /// Flush all pending relocations to patch original contents of sections

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"
+#include "llvm/Support/DataExtractor.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, const BinaryContext &BC) {
+  const auto EntrySize = 2 * BC.AsmInfo->getCodePointerSize();
+  assert(Buffer.size() % EntrySize == 0 && "Unexpected address map size");
+
+  DataExtractor DE(Buffer, BC.AsmInfo->isLittleEndian(),
+                   BC.AsmInfo->getCodePointerSize());
+  DataExtractor::Cursor Cursor(0);
+
+  AddressMap Parsed;
+  Parsed.Map.reserve(Buffer.size() / EntrySize);
+
+  while (Cursor && !DE.eof(Cursor)) {
+    const auto Input = DE.getAddress(Cursor);
+    const auto Output = DE.getAddress(Cursor);
+    Parsed.Map.insert({Input, Output});
+  }
+
+  assert(Cursor && "Error reading address map section");
+  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/BinaryEmitter.cpp

void BinaryEmitter::emitAll(StringRef OrgSecPrefix) {
   emitFunctions();
 
   if (opts::UpdateDebugSections) {
     emitDebugLineInfoForOriginalFunctions();
     DwarfLineTable::emit(BC, Streamer);
   }
 
   emitDataSections(OrgSecPrefix);
+
+  // TODO Enable for Mach-O once BinaryContext::getDataSection supports it.
+  if (BC.isELF())
+    AddressMap::emit(Streamer, BC);
 }
 
 void BinaryEmitter::emitFunctions() {
   auto emit = [&](const std::vector<BinaryFunction *> &Functions) {
     const bool HasProfile = BC.NumProfiledFuncs > 0;
     const bool OriginalAllowAutoPadding = Streamer.getAllowAutoPadding();
     for (BinaryFunction *Function : Functions) {
       if (!BC.shouldEmit(*Function))

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

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;
+  const auto &IOAddressMap =
+      BB.getFunction()->getBinaryContext().getIOAddressMap();
+
+  for (const auto &[InputOffset, Sym] : BB.getLocSyms()) {
+    const auto InputAddress = BB.getFunction()->getAddress() + InputOffset;
+    const auto OutputAddress = IOAddressMap.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(

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 (BinarySection *EHFrameSection = getSection(getEHFrameSectionName())) {
     BC->registerSection(".relocated.eh_frame", *EHFrameSection);
   }
   BC->registerOrUpdateSection(getNewSecPrefix() + ".gcc_except_table",
                               ELF::SHT_PROGBITS, ROFlags);
   BC->registerOrUpdateSection(getNewSecPrefix() + ".rodata", ELF::SHT_PROGBITS,
                               ROFlags);
   BC->registerOrUpdateSection(getNewSecPrefix() + ".rodata.cold",
                               ELF::SHT_PROGBITS, ROFlags);
+  BC->registerOrUpdateSection(AddressMap::SectionName, ELF::SHT_PROGBITS,
+                              ROFlags)
+      .setLinkOnly();
 }
 
 void RewriteInstance::emitAndLink() {
   NamedRegionTimer T("emitAndLink", "emit and link", TimerGroupName,
                      TimerGroupDesc, opts::TimeRewrite);
 
   SmallString<0> ObjectBuffer;
   raw_svector_ostream OS(ObjectBuffer);
if (opts::UseOldText || opts::StrictMode) {
       EHFrameSection->clearContents();
     } else {
       // Make .eh_frame relocatable.
       relocateEHFrameSection();
     }
   }
 
   emitBinaryContext(*Streamer, *BC, getOrgSecPrefix());
 
   Streamer->finish();
   if (Streamer->getContext().hadError()) {
     errs() << "BOLT-ERROR: Emission failed.\n";
     exit(1);
   }
 
   if (opts::KeepTmp) {
     SmallString<128> OutObjectPath;
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 (Section.isLinkOnly())
+        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);
+    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 (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();
     NewSection.sh_entsize = 0;
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 (Section.isLinkOnly())
+      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());