[BOLT] Fix incorrect basic block output addresses (llvm#70000)

Some optimization passes may duplicate basic blocks and assign the same
input offset to a number of different blocks in a function. This is done
e.g. to correctly map debugging ranges for duplicated code.

However, duplicate input offsets present a problem when we use
AddressMap to generate new addresses for basic blocks. The output
address is calculated based on the input offset and will be the same for
blocks with identical offsets. The result is potentially incorrect debug
info and BAT records.

To address the issue, we have to eliminate the dependency on input
offsets while generating output addresses for a basic block. Each block
has a unique label, hence we extend AddressMap to include address lookup
based on MCSymbol and use the new functionality to update block
addresses.

bolt/include/bolt/Core/AddressMap.h

@@ -6,18 +6,16 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// 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.
+// This file contains the declaration of the AddressMap class used for looking
+// up addresses in the output object.
 //
 //===----------------------------------------------------------------------===//
-//
+
 #ifndef BOLT_CORE_ADDRESS_MAP_H
 #define BOLT_CORE_ADDRESS_MAP_H
 
 #include "llvm/ADT/StringRef.h"
+#include "llvm/MC/MCSymbol.h"
 
 #include <optional>
 #include <unordered_map>
@@ -30,26 +28,48 @@ namespace bolt {
 
 class BinaryContext;
 
+/// Helper class to create a mapping from input entities to output addresses
+/// needed for updating debugging symbols and BAT. We emit a section containing
+/// <Input entity, Output MCSymbol> pairs to the object file and JITLink will
+/// transform this in <Input entity, Output address> pairs. The linker output
+/// can then be parsed and used to establish the mapping.
+///
+/// The entities that can be mapped to output address are input addresses and
+/// labels (MCSymbol). Input addresses support one-to-many mapping.
 class AddressMap {
-  using MapTy = std::unordered_multimap<uint64_t, uint64_t>;
-  MapTy Map;
+  static const char *const AddressSectionName;
+  static const char *const LabelSectionName;
 
-public:
-  static const char *const SectionName;
+  /// Map multiple <input address> to <output address>.
+  using Addr2AddrMapTy = std::unordered_multimap<uint64_t, uint64_t>;
+  Addr2AddrMapTy Address2AddressMap;
 
+  /// Map MCSymbol to its output address. Normally used for temp symbols that
+  /// are not updated by the linker.
+  using Label2AddrMapTy = DenseMap<const MCSymbol *, uint64_t>;
+  Label2AddrMapTy Label2AddrMap;
+
+public:
   static void emit(MCStreamer &Streamer, BinaryContext &BC);
-  static AddressMap parse(StringRef Buffer, const BinaryContext &BC);
+  static std::optional<AddressMap> parse(BinaryContext &BC);
 
   std::optional<uint64_t> lookup(uint64_t InputAddress) const {
-    auto It = Map.find(InputAddress);
-    if (It != Map.end())
+    auto It = Address2AddressMap.find(InputAddress);
+    if (It != Address2AddressMap.end())
+      return It->second;
+    return std::nullopt;
+  }
+
+  std::optional<uint64_t> lookup(const MCSymbol *Symbol) const {
+    auto It = Label2AddrMap.find(Symbol);
+    if (It != Label2AddrMap.end())
       return It->second;
     return std::nullopt;
   }
 
-  std::pair<MapTy::const_iterator, MapTy::const_iterator>
+  std::pair<Addr2AddrMapTy::const_iterator, Addr2AddrMapTy::const_iterator>
   lookupAll(uint64_t InputAddress) const {
-    return Map.equal_range(InputAddress);
+    return Address2AddressMap.equal_range(InputAddress);
   }
 };
 

bolt/lib/Core/AddressMap.cpp

@@ -1,22 +1,44 @@
+//===- bolt/Core/AddressMap.cpp - Input-output Address Map ----------------===//
+//
+// 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/Core/AddressMap.h"
 #include "bolt/Core/BinaryContext.h"
 #include "bolt/Core/BinaryFunction.h"
+#include "bolt/Core/BinarySection.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/Support/DataExtractor.h"
 
 namespace llvm {
 namespace bolt {
 
-const char *const AddressMap::SectionName = ".bolt.address_map";
+const char *const AddressMap::AddressSectionName = ".bolt.addr2addr_map";
+const char *const AddressMap::LabelSectionName = ".bolt.label2addr_map";
 
-static void emitLabel(MCStreamer &Streamer, uint64_t InputAddress,
-                      const MCSymbol *OutputLabel) {
+static void emitAddress(MCStreamer &Streamer, uint64_t InputAddress,
+                        const MCSymbol *OutputLabel) {
   Streamer.emitIntValue(InputAddress, 8);
   Streamer.emitSymbolValue(OutputLabel, 8);
 }
 
+static void emitLabel(MCStreamer &Streamer, const MCSymbol *OutputLabel) {
+  Streamer.emitIntValue(reinterpret_cast<uint64_t>(OutputLabel), 8);
+  Streamer.emitSymbolValue(OutputLabel, 8);
+}
+
 void AddressMap::emit(MCStreamer &Streamer, BinaryContext &BC) {
-  Streamer.switchSection(BC.getDataSection(SectionName));
+  // Mark map sections as link-only to avoid allocation in the output file.
+  const unsigned Flags = BinarySection::getFlags(/*IsReadOnly*/ true,
+                                                 /*IsText*/ false,
+                                                 /*IsAllocatable*/ true);
+  BC.registerOrUpdateSection(AddressSectionName, ELF::SHT_PROGBITS, Flags)
+      .setLinkOnly();
+  BC.registerOrUpdateSection(LabelSectionName, ELF::SHT_PROGBITS, Flags)
+      .setLinkOnly();
 
   for (const auto &[BFAddress, BF] : BC.getBinaryFunctions()) {
     if (!BF.requiresAddressMap())
@@ -26,37 +48,69 @@ void AddressMap::emit(MCStreamer &Streamer, BinaryContext &BC) {
       if (!BB.getLabel()->isDefined())
         continue;
 
-      emitLabel(Streamer, BFAddress + BB.getInputAddressRange().first,
-                BB.getLabel());
+      Streamer.switchSection(BC.getDataSection(LabelSectionName));
+      emitLabel(Streamer, BB.getLabel());
 
       if (!BB.hasLocSyms())
         continue;
 
+      Streamer.switchSection(BC.getDataSection(AddressSectionName));
       for (auto [Offset, Symbol] : BB.getLocSyms())
-        emitLabel(Streamer, BFAddress + Offset, Symbol);
+        emitAddress(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");
+std::optional<AddressMap> AddressMap::parse(BinaryContext &BC) {
+  auto AddressMapSection = BC.getUniqueSectionByName(AddressSectionName);
+  auto LabelMapSection = BC.getUniqueSectionByName(LabelSectionName);
 
-  DataExtractor DE(Buffer, BC.AsmInfo->isLittleEndian(),
-                   BC.AsmInfo->getCodePointerSize());
-  DataExtractor::Cursor Cursor(0);
+  if (!AddressMapSection && !LabelMapSection)
+    return std::nullopt;
 
   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);
-    if (!Parsed.Map.count(Input))
-      Parsed.Map.insert({Input, Output});
+  const size_t EntrySize = 2 * BC.AsmInfo->getCodePointerSize();
+  auto parseSection =
+      [&](BinarySection &Section,
+          function_ref<void(uint64_t, uint64_t)> InsertCallback) {
+        StringRef Buffer = Section.getOutputContents();
+        assert(Buffer.size() % EntrySize == 0 && "Unexpected address map size");
+
+        DataExtractor DE(Buffer, BC.AsmInfo->isLittleEndian(),
+                         BC.AsmInfo->getCodePointerSize());
+        DataExtractor::Cursor Cursor(0);
+
+        while (Cursor && !DE.eof(Cursor)) {
+          const uint64_t Input = DE.getAddress(Cursor);
+          const uint64_t Output = DE.getAddress(Cursor);
+          InsertCallback(Input, Output);
+        }
+
+        assert(Cursor && "Error reading address map section");
+        BC.deregisterSection(Section);
+      };
+
+  if (AddressMapSection) {
+    Parsed.Address2AddressMap.reserve(AddressMapSection->getOutputSize() /
+                                      EntrySize);
+    parseSection(*AddressMapSection, [&](uint64_t Input, uint64_t Output) {
+      if (!Parsed.Address2AddressMap.count(Input))
+        Parsed.Address2AddressMap.insert({Input, Output});
+    });
+  }
+
+  if (LabelMapSection) {
+    Parsed.Label2AddrMap.reserve(LabelMapSection->getOutputSize() / EntrySize);
+    parseSection(*LabelMapSection, [&](uint64_t Input, uint64_t Output) {
+      assert(!Parsed.Label2AddrMap.count(
+                 reinterpret_cast<const MCSymbol *>(Input)) &&
+             "Duplicate label entry detected.");
+      Parsed.Label2AddrMap.insert(
+          {reinterpret_cast<const MCSymbol *>(Input), Output});
+    });
   }
 
-  assert(Cursor && "Error reading address map section");
   return Parsed;
 }
 

bolt/lib/Core/BinaryFunction.cpp

@@ -4165,15 +4165,20 @@ void BinaryFunction::updateOutputValues(const BOLTLinker &Linker) {
       // Injected functions likely will fail lookup, as they have no
       // input range. Just assign the BB the output address of the
       // function.
-      auto MaybeBBAddress =
-          BC.getIOAddressMap().lookup(BB->getInputOffset() + getAddress());
+      auto MaybeBBAddress = BC.getIOAddressMap().lookup(BB->getLabel());
       const uint64_t BBAddress = MaybeBBAddress  ? *MaybeBBAddress
                                  : BB->isSplit() ? FF.getAddress()
                                                  : getOutputAddress();
       BB->setOutputStartAddress(BBAddress);
 
-      if (PrevBB)
+      if (PrevBB) {
+        assert(PrevBB->getOutputAddressRange().first <= BBAddress &&
+               "Bad output address for basic block.");
+        assert((PrevBB->getOutputAddressRange().first != BBAddress ||
+                !hasInstructions() || PrevBB->empty()) &&
+               "Bad output address for basic block.");
         PrevBB->setOutputEndAddress(BBAddress);
+      }
       PrevBB = BB;
     }
 

bolt/lib/Rewrite/RewriteInstance.cpp

@@ -3196,9 +3196,6 @@ void RewriteInstance::preregisterSections() {
                               ROFlags);
   BC->registerOrUpdateSection(getNewSecPrefix() + ".rodata.cold",
                               ELF::SHT_PROGBITS, ROFlags);
-  BC->registerOrUpdateSection(AddressMap::SectionName, ELF::SHT_PROGBITS,
-                              ROFlags)
-      .setLinkOnly();
 }
 
 void RewriteInstance::emitAndLink() {
@@ -3669,11 +3666,8 @@ void RewriteInstance::mapAllocatableSections(
 }
 
 void RewriteInstance::updateOutputValues(const BOLTLinker &Linker) {
-  if (auto MapSection = BC->getUniqueSectionByName(AddressMap::SectionName)) {
-    auto Map = AddressMap::parse(MapSection->getOutputContents(), *BC);
-    BC->setIOAddressMap(std::move(Map));
-    BC->deregisterSection(*MapSection);
-  }
+  if (std::optional<AddressMap> Map = AddressMap::parse(*BC))
+    BC->setIOAddressMap(std::move(*Map));
 
   for (BinaryFunction *Function : BC->getAllBinaryFunctions())
     Function->updateOutputValues(Linker);

bolt/test/X86/jump-table-icp.test

@@ -12,6 +12,7 @@ RUN: (llvm-bolt %t.exe --data %t.fdata -o %t --relocs \
 RUN:   --reorder-blocks=cache --split-functions --split-all-cold \
 RUN:   --use-gnu-stack --dyno-stats --indirect-call-promotion=jump-tables \
 RUN:   --print-icp -v=0 \
+RUN:   --enable-bat --print-cache-metrics \
 RUN:   --icp-jt-remaining-percent-threshold=10 \
 RUN:   --icp-jt-total-percent-threshold=2 \
 RUN:   --indirect-call-promotion-topn=1 \