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

via llvm-commits llvm-commits at lists.llvm.org
Mon Oct 23 19:54:12 PDT 2023


https://github.com/maksfb created https://github.com/llvm/llvm-project/pull/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.

>From 2ebde064bbb2021dba66c9e88887e4bbfafba483 Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks at fb.com>
Date: Mon, 23 Oct 2023 19:45:09 -0700
Subject: [PATCH] [BOLT] Fix incorrect basic block output addresses

Summary:
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  | 50 ++++++++++-----
 bolt/lib/Core/AddressMap.cpp         | 92 ++++++++++++++++++++++------
 bolt/lib/Core/BinaryFunction.cpp     | 11 +++-
 bolt/lib/Rewrite/RewriteInstance.cpp | 10 +--
 bolt/test/X86/jump-table-icp.test    |  1 +
 5 files changed, 118 insertions(+), 46 deletions(-)

diff --git a/bolt/include/bolt/Core/AddressMap.h b/bolt/include/bolt/Core/AddressMap.h
index 16c2727b6943fc6..85a9ab4473aafed 100644
--- a/bolt/include/bolt/Core/AddressMap.h
+++ b/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);
   }
 };
 
diff --git a/bolt/lib/Core/AddressMap.cpp b/bolt/lib/Core/AddressMap.cpp
index c5f628d87864060..1ef4550884bbefb 100644
--- a/bolt/lib/Core/AddressMap.cpp
+++ b/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,67 @@ 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<MCSymbol *>(Input)));
+      Parsed.Label2AddrMap.insert(
+          {reinterpret_cast<MCSymbol *>(Input), Output});
+    });
   }
 
-  assert(Cursor && "Error reading address map section");
   return Parsed;
 }
 
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 65ebb8599ff75c4..2e1fa739e53f20b 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/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;
     }
 
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index b3de3b96b3ab8e6..0d78c9b75e03d32 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/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);
diff --git a/bolt/test/X86/jump-table-icp.test b/bolt/test/X86/jump-table-icp.test
index 708f1273af3f196..34339dc327fae52 100644
--- a/bolt/test/X86/jump-table-icp.test
+++ b/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 \



More information about the llvm-commits mailing list