[BOLT] Fix incorrect basic block output 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 \