Add BAT testing framework
authorRafael Auler <rafaelauler@fb.com>
Wed, 6 Jul 2022 22:53:27 +0000 (15:53 -0700)
committerRafael Auler <rafaelauler@fb.com>
Fri, 29 Jul 2022 21:55:04 +0000 (14:55 -0700)
This patch refactors BAT to be testable as a library, so we
can have open-source tests on it. This further fixes an issue with
basic blocks that lack a valid input offset, making BAT omit those
when writing translation tables.

Test Plan: new testcases added, new testing tool added (llvm-bat-dump)

Differential Revision: https://reviews.llvm.org/D129382

12 files changed:
bolt/include/bolt/Profile/BoltAddressTranslation.h
bolt/lib/Passes/ValidateInternalCalls.cpp
bolt/lib/Profile/BoltAddressTranslation.cpp
bolt/lib/Profile/DataAggregator.cpp
bolt/lib/Rewrite/RewriteInstance.cpp
bolt/test/CMakeLists.txt
bolt/test/X86/bolt-address-translation-internal-call.test [new file with mode: 0644]
bolt/test/X86/bolt-address-translation.test [new file with mode: 0644]
bolt/test/lit.cfg.py
bolt/tools/CMakeLists.txt
bolt/tools/bat-dump/CMakeLists.txt [new file with mode: 0644]
bolt/tools/bat-dump/bat-dump.cpp [new file with mode: 0644]

index 2b1d471..43988ad 100644 (file)
@@ -76,26 +76,29 @@ public:
   /// output binary
   static const char *SECTION_NAME;
 
-  BoltAddressTranslation(BinaryContext &BC) : BC(BC) {}
+  BoltAddressTranslation() {}
 
   /// Write the serialized address translation tables for each reordered
   /// function
-  void write(raw_ostream &OS);
+  void write(const BinaryContext &BC, raw_ostream &OS);
 
   /// Read the serialized address translation tables and load them internally
   /// in memory. Return a parse error if failed.
   std::error_code parse(StringRef Buf);
 
+  /// Dump the parsed address translation tables
+  void dump(raw_ostream &OS);
+
   /// If the maps are loaded in memory, perform the lookup to translate LBR
-  /// addresses in \p Func.
-  uint64_t translate(const BinaryFunction &Func, uint64_t Offset,
+  /// addresses in function located at \p FuncAddress.
+  uint64_t translate(uint64_t FuncAddress, uint64_t Offset,
                      bool IsBranchSrc) const;
 
   /// Use the map keys containing basic block addresses to infer fall-throughs
   /// taken in the path started at FirstLBR.To and ending at SecondLBR.From.
   /// Return NoneType if trace is invalid or the list of fall-throughs
   /// otherwise.
-  Optional<FallthroughListTy> getFallthroughsInTrace(const BinaryFunction &Func,
+  Optional<FallthroughListTy> getFallthroughsInTrace(uint64_t FuncAddress,
                                                      uint64_t From,
                                                      uint64_t To) const;
 
@@ -115,8 +118,6 @@ private:
   void writeEntriesForBB(MapTy &Map, const BinaryBasicBlock &BB,
                          uint64_t FuncAddress);
 
-  BinaryContext &BC;
-
   std::map<uint64_t, MapTy> Maps;
 
   /// Links outlined cold bocks to their original function
index 4e2e527..ba7e791 100644 (file)
@@ -117,7 +117,6 @@ void ValidateInternalCalls::fixCFGForPIC(BinaryFunction &Function) const {
     if (!MovedInsts.empty()) {
       // Split this block at the call instruction.
       std::unique_ptr<BinaryBasicBlock> NewBB = Function.createBasicBlock();
-      NewBB->setOffset(0);
       NewBB->addInstructions(MovedInsts.begin(), MovedInsts.end());
       BB.moveAllSuccessorsTo(NewBB.get());
 
index b25329b..3059623 100644 (file)
@@ -25,9 +25,14 @@ void BoltAddressTranslation::writeEntriesForBB(MapTy &Map,
       BB.getOutputAddressRange().first - FuncAddress;
   const uint32_t BBInputOffset = BB.getInputOffset();
 
-  assert(BBInputOffset != BinaryBasicBlock::INVALID_OFFSET &&
-         "Every output BB must track back to an input BB for profile "
-         "collection in bolted binaries");
+  // 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
+  // layout. This covers the case where an input basic block is split into two,
+  // and the second one lacks any offset.
+  if (BBInputOffset == BinaryBasicBlock::INVALID_OFFSET)
+    return;
 
   LLVM_DEBUG(dbgs() << "BB " << BB.getName() << "\n");
   LLVM_DEBUG(dbgs() << "  Key: " << Twine::utohexstr(BBOutputOffset)
@@ -56,13 +61,13 @@ void BoltAddressTranslation::writeEntriesForBB(MapTy &Map,
   }
 }
 
-void BoltAddressTranslation::write(raw_ostream &OS) {
+void BoltAddressTranslation::write(const BinaryContext &BC, raw_ostream &OS) {
   LLVM_DEBUG(dbgs() << "BOLT-DEBUG: Writing BOLT Address Translation Tables\n");
   for (auto &BFI : BC.getBinaryFunctions()) {
-    BinaryFunction &Function = BFI.second;
+    const BinaryFunction &Function = BFI.second;
     // We don't need a translation table if the body of the function hasn't
     // changed
-    if (!BC.HasRelocations && !Function.isSimple())
+    if (Function.isIgnored() || (!BC.HasRelocations && !Function.isSimple()))
       continue;
 
     LLVM_DEBUG(dbgs() << "Function name: " << Function.getPrintName() << "\n");
@@ -70,7 +75,7 @@ void BoltAddressTranslation::write(raw_ostream &OS) {
                       << Twine::utohexstr(Function.getOutputAddress()) << "\n");
     MapTy Map;
     const bool IsSplit = Function.isSplit();
-    for (const BinaryBasicBlock *BB : Function.getLayout().blocks()) {
+    for (const BinaryBasicBlock *const BB : Function.getLayout().blocks()) {
       if (IsSplit && BB->isCold())
         break;
       writeEntriesForBB(Map, *BB, Function.getOutputAddress());
@@ -83,7 +88,7 @@ void BoltAddressTranslation::write(raw_ostream &OS) {
     // Cold map
     Map.clear();
     LLVM_DEBUG(dbgs() << " Cold part\n");
-    for (const BinaryBasicBlock *BB : Function.getLayout().blocks()) {
+    for (const BinaryBasicBlock *const BB : Function.getLayout().blocks()) {
       if (!BB->isCold())
         continue;
       writeEntriesForBB(Map, *BB, Function.cold().getAddress());
@@ -193,10 +198,39 @@ std::error_code BoltAddressTranslation::parse(StringRef Buf) {
   return std::error_code();
 }
 
-uint64_t BoltAddressTranslation::translate(const BinaryFunction &Func,
+void BoltAddressTranslation::dump(raw_ostream &OS) {
+  const size_t NumTables = Maps.size();
+  OS << "BAT tables for " << NumTables << " functions:\n";
+  for (const auto &MapEntry : Maps) {
+    OS << "Function Address: 0x" << Twine::utohexstr(MapEntry.first) << "\n";
+    OS << "BB mappings:\n";
+    for (const auto &Entry : MapEntry.second) {
+      const bool IsBranch = Entry.second & BRANCHENTRY;
+      const uint32_t Val = Entry.second & ~BRANCHENTRY;
+      OS << "0x" << Twine::utohexstr(Entry.first) << " -> "
+         << "0x" << Twine::utohexstr(Val);
+      if (IsBranch)
+        OS << " (branch)";
+      OS << "\n";
+    }
+    OS << "\n";
+  }
+  const size_t NumColdParts = ColdPartSource.size();
+  if (!NumColdParts)
+    return;
+
+  OS << NumColdParts << " cold mappings:\n";
+  for (const auto &Entry : ColdPartSource) {
+    OS << "0x" << Twine::utohexstr(Entry.first) << " -> "
+       << Twine::utohexstr(Entry.second) << "\n";
+  }
+  OS << "\n";
+}
+
+uint64_t BoltAddressTranslation::translate(uint64_t FuncAddress,
                                            uint64_t Offset,
                                            bool IsBranchSrc) const {
-  auto Iter = Maps.find(Func.getAddress());
+  auto Iter = Maps.find(FuncAddress);
   if (Iter == Maps.end())
     return Offset;
 
@@ -217,7 +251,7 @@ uint64_t BoltAddressTranslation::translate(const BinaryFunction &Func,
 }
 
 Optional<BoltAddressTranslation::FallthroughListTy>
-BoltAddressTranslation::getFallthroughsInTrace(const BinaryFunction &Func,
+BoltAddressTranslation::getFallthroughsInTrace(uint64_t FuncAddress,
                                                uint64_t From,
                                                uint64_t To) const {
   SmallVector<std::pair<uint64_t, uint64_t>, 16> Res;
@@ -226,10 +260,10 @@ BoltAddressTranslation::getFallthroughsInTrace(const BinaryFunction &Func,
   if (From >= To)
     return Res;
 
-  From -= Func.getAddress();
-  To -= Func.getAddress();
+  From -= FuncAddress;
+  To -= FuncAddress;
 
-  auto Iter = Maps.find(Func.getAddress());
+  auto Iter = Maps.find(FuncAddress);
   if (Iter == Maps.end())
     return NoneType();
 
index 39f17c4..885a54f 100644 (file)
@@ -699,7 +699,7 @@ bool DataAggregator::doSample(BinaryFunction &Func, uint64_t Address,
 
   Address -= Func.getAddress();
   if (BAT)
-    Address = BAT->translate(Func, Address, /*IsBranchSrc=*/false);
+    Address = BAT->translate(Func.getAddress(), Address, /*IsBranchSrc=*/false);
 
   I->second.bumpCount(Address, Count);
   return true;
@@ -722,8 +722,8 @@ bool DataAggregator::doIntraBranch(BinaryFunction &Func, uint64_t From,
                     << Func.getPrintName() << " @ " << Twine::utohexstr(To)
                     << '\n');
   if (BAT) {
-    From = BAT->translate(Func, From, /*IsBranchSrc=*/true);
-    To = BAT->translate(Func, To, /*IsBranchSrc=*/false);
+    From = BAT->translate(Func.getAddress(), From, /*IsBranchSrc=*/true);
+    To = BAT->translate(Func.getAddress(), To, /*IsBranchSrc=*/false);
     LLVM_DEBUG(dbgs() << "BOLT-DEBUG: BAT translation on bumpBranchCount: "
                       << Func.getPrintName() << " @ " << Twine::utohexstr(From)
                       << " -> " << Func.getPrintName() << " @ "
@@ -752,7 +752,7 @@ bool DataAggregator::doInterBranch(BinaryFunction *FromFunc,
     }
     From -= FromFunc->getAddress();
     if (BAT)
-      From = BAT->translate(*FromFunc, From, /*IsBranchSrc=*/true);
+      From = BAT->translate(FromFunc->getAddress(), From, /*IsBranchSrc=*/true);
 
     recordExit(*FromFunc, From, Mispreds, Count);
   }
@@ -766,7 +766,7 @@ bool DataAggregator::doInterBranch(BinaryFunction *FromFunc,
     }
     To -= ToFunc->getAddress();
     if (BAT)
-      To = BAT->translate(*ToFunc, To, /*IsBranchSrc=*/false);
+      To = BAT->translate(ToFunc->getAddress(), To, /*IsBranchSrc=*/false);
 
     recordEntry(*ToFunc, To, Mispreds, Count);
   }
@@ -822,7 +822,8 @@ bool DataAggregator::doTrace(const LBREntry &First, const LBREntry &Second,
   }
 
   Optional<BoltAddressTranslation::FallthroughListTy> FTs =
-      BAT ? BAT->getFallthroughsInTrace(*FromFunc, First.To, Second.From)
+      BAT ? BAT->getFallthroughsInTrace(FromFunc->getAddress(), First.To,
+                                        Second.From)
           : getFallthroughsInTrace(*FromFunc, First, Second, Count);
   if (!FTs) {
     LLVM_DEBUG(
index 5bc6374..8c171a8 100644 (file)
@@ -356,7 +356,7 @@ RewriteInstance::RewriteInstance(ELFObjectFileBase *File, const int Argc,
   BC->initializeTarget(std::unique_ptr<MCPlusBuilder>(createMCPlusBuilder(
       BC->TheTriple->getArch(), BC->MIA.get(), BC->MII.get(), BC->MRI.get())));
 
-  BAT = std::make_unique<BoltAddressTranslation>(*BC);
+  BAT = std::make_unique<BoltAddressTranslation>();
 
   if (opts::UpdateDebugSections)
     DebugInfoRewriter = std::make_unique<DWARFRewriter>(*BC);
@@ -4163,7 +4163,7 @@ void RewriteInstance::encodeBATSection() {
   std::string DescStr;
   raw_string_ostream DescOS(DescStr);
 
-  BAT->write(DescOS);
+  BAT->write(*BC, DescOS);
   DescOS.flush();
 
   const std::string BoltInfo =
index 2ac5b47..898ac3e 100644 (file)
@@ -39,6 +39,7 @@ list(APPEND BOLT_TEST_DEPS
   llvm-bolt
   llvm-boltdiff
   llvm-bolt-heatmap
+  llvm-bat-dump
   llvm-dwarfdump
   llvm-dwp
   llvm-mc
diff --git a/bolt/test/X86/bolt-address-translation-internal-call.test b/bolt/test/X86/bolt-address-translation-internal-call.test
new file mode 100644 (file)
index 0000000..edc32d9
--- /dev/null
@@ -0,0 +1,52 @@
+# This checks for an issue with internal calls and BAT (BOLT address
+# translation). BAT needs to map every output block back to an input
+# block, but passes that introduce new blocks (such as validate
+# internal calls) might create new blocks without a mapping to an
+# input block.
+
+# REQUIRES: system-linux,bolt-runtime
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
+# Delete our BB symbols so BOLT doesn't mark them as entry points
+# RUN: llvm-strip --strip-unneeded %t.o
+# RUN: %clang %t.o -o %t.exe -Wl,-q
+
+# RUN: llvm-bolt --enable-bat %t.exe --relocs -o %t.out | FileCheck %s
+# CHECK: BOLT-INFO: Wrote {{.*}} BAT maps
+
+# RUN: llvm-bat-dump %t.out --dump-all | \
+# RUN:   FileCheck %s --check-prefix=CHECK-BAT-DUMP
+# CHECK-BAT-DUMP:  BAT tables for {{.*}} functions
+
+  .text
+  .globl  main
+  .type main, %function
+  .p2align  4
+main:
+  push   %rbp
+  mov    %rsp,%rbp
+  push   %r12
+  push   %rbx
+  sub    $0x120,%rsp
+  mov    $0x3,%rbx
+.J1:
+  cmp    $0x0,%rbx
+  je     .J2
+  callq  .J3
+  nopl   (%rax,%rax,1)
+  movabs $0xdeadbeef,%rax
+  retq
+.J2:
+  add    $0x120,%rsp
+  pop    %rbx
+  pop    %r12
+  jmp    .J4
+.J3:
+  pop    %rax
+  add    $0x4,%rax
+  dec    %rbx
+  jmp    .J1
+.J4:
+  pop    %rbp
+  retq
+  .size main, .-main
diff --git a/bolt/test/X86/bolt-address-translation.test b/bolt/test/X86/bolt-address-translation.test
new file mode 100644 (file)
index 0000000..f68a8f7
--- /dev/null
@@ -0,0 +1,146 @@
+# Check a common case for BOLT address translation tables. These tables are used
+# to translate profile activity happening in a bolted binary back to the
+# original binary, so you can run BOLT again, with updated profile collected
+# in a production environment that only runs bolted binaries. As BOLT only
+# takes no-bolt binaries as inputs, this translation is necessary to cover
+# this scenario.
+#
+# RUN: yaml2obj %p/Inputs/blarge.yaml &> %t.exe
+# RUN: llvm-bolt %t.exe -o %t.out --data %p/Inputs/blarge.fdata \
+# RUN:   --reorder-blocks=normal --split-functions --enable-bat 2>&1 | FileCheck %s
+# RUN: llvm-bat-dump %t.out --dump-all \
+# RUN:   --translate=0x401180 | FileCheck %s --check-prefix=CHECK-BAT-DUMP
+#
+# In this test we focus on function usqrt at address 0x401170. This is a
+# non-reloc binary case, so we don't expect this address to change, that's
+# why we hardcode its address here. This address also comes hardcoded in the
+# blarge.yaml input file.
+#
+# This is the layout of the function before BOLT reorder blocks:
+#
+#    BB Layout   : .LBB02, .Ltmp39, .LFT1, .Ltmp38, .LFT2
+#
+# This is the layout of the function after BOLT reorder blocks:
+#
+#    BB Layout   : .LBB02, .Ltmp38, .Ltmp39, .LFT2, .LFT3
+#
+# .Ltmp38 is originally at offset 0x39 but gets moved to 0xc (see full dump
+# below).
+#
+# We check that BAT is able to translate references happening in .Ltmp38 to
+# its original offset.
+#
+
+# This binary has 3 functions with profile, all of them are split, so 6 maps.
+# BAT creates one map per function fragment.
+#
+# CHECK:      BOLT: 3 out of 7 functions were overwritten.
+# CHECK:      BOLT-INFO: Wrote 6 BAT maps
+# CHECK:      BOLT-INFO: Wrote 3 BAT cold-to-hot entries
+#
+# usqrt mappings (hot part). We match against any key (left side containing
+# the bolted binary offsets) because BOLT may change where it puts instructions
+# depending on whether it is relaxing a branch or not. But the original input
+# binary offsets (right side) should be the same because these addresses are
+# hardcoded in the blarge.yaml file.
+#
+# CHECK-BAT-DUMP:      Function Address: 0x401170
+# CHECK-BAT-DUMP-NEXT: BB mappings:
+# CHECK-BAT-DUMP-NEXT: 0x0 -> 0x0
+# CHECK-BAT-DUMP-NEXT: 0x8 -> 0x8 (branch)
+# CHECK-BAT-DUMP-NEXT: 0x{{.*}} -> 0x39
+# CHECK-BAT-DUMP-NEXT: 0x{{.*}} -> 0x3d (branch)
+# CHECK-BAT-DUMP-NEXT: 0x{{.*}} -> 0x10
+# CHECK-BAT-DUMP-NEXT: 0x{{.*}} -> 0x30 (branch)
+#
+# CHECK-BAT-DUMP: 3 cold mappings
+#
+# Now check that the translation 0x401180 maps back to its correct
+# input offset (offset 3d in the usqrt input function).
+#
+# COM: CHECK-BAT-DUMP: Translating addresses according to parsed BAT tables:
+#      CHECK-BAT-DUMP: 0x401180 -> usqrt + 0x3d
+
+# -------------------------
+# Full dump for reference (this is not checked):
+# -------------------------
+
+Binary Function "usqrt" after finalize-functions
+  Number      : 7
+  State       : CFG finalized
+  Address     : 0x401170
+  Size        : 0x43
+  MaxSize     : 0x43
+  Offset      : 0xcb0
+  Section     : .text
+  Orc Section : .local.text.usqrt
+  LSDA        : 0x0
+  IsSimple    : 1
+  IsMultiEntry: 0
+  IsSplit     : 1
+  BB Count    : 5
+  Hash        : a6468f132ec176ca
+  BB Layout   : .LBB02, .Ltmp38, .Ltmp39, .LFT2, .LFT3
+  Exec Count  : 199
+  Profile Acc : 100.0%
+
+.LBB02 (4 instructions, align : 1)
+  Entry Point
+  Exec Count : 199
+  CFI State : 0
+  Input offset: 0
+    00000000:   movl    $0x20, %r8d
+    00000006:   xorl    %eax, %eax
+    00000008:   xorl    %edx, %edx # Offset: 8
+    0000000a:   jmp     .Ltmp39
+  Successors: .Ltmp39 (mispreds: 0, count: 0)
+
+.Ltmp38 (2 instructions, align : 1)
+  Exec Count : 4711
+  CFI State : 0
+  Input offset: 39
+  Predecessors: .Ltmp39, .LFT2
+    0000000c:   subl    $0x1, %r8d
+    00000010:   je      .LFT3 # Offset: 61
+  Successors: .LFT3 (mispreds: 0, count: 0), .Ltmp39 (mispreds: 33, count: 4711)
+
+.Ltmp39 (10 instructions, align : 1)
+  Exec Count : 4711
+  CFI State : 0
+  Input offset: 10
+  Predecessors: .Ltmp38, .LBB02
+    00000012:   movq    %rdi, %rcx
+    00000015:   addq    %rax, %rax
+    00000018:   shlq    $0x2, %rdi
+    0000001c:   andl    $0xc0000000, %ecx
+    00000022:   shrq    $0x1e, %rcx
+    00000026:   leaq    (%rcx,%rdx,4), %rdx
+    0000002a:   leaq    0x1(%rax,%rax), %rcx
+    0000002f:   cmpq    %rcx, %rdx
+    00000032:   jb      .Ltmp38 # Offset: 48
+    00000034:   jmp     .LFT2
+  Successors: .Ltmp38 (mispreds: 171, count: 2886), .LFT2 (mispreds: 0, count: 0)
+
+-------   HOT-COLD SPLIT POINT   -------
+
+.LFT2 (3 instructions, align : 1)
+  Exec Count : 0
+  CFI State : 0
+  Input offset: 32
+  Predecessors: .Ltmp39
+    00000036:   subq    %rcx, %rdx
+    00000039:   addq    $0x1, %rax # Offset: 53
+    0000003d:   jmp     .Ltmp38
+  Successors: .Ltmp38 (mispreds: 0, count: 0)
+
+.LFT3 (2 instructions, align : 1)
+  Exec Count : 0
+  CFI State : 0
+  Input offset: 3f
+  Predecessors: .Ltmp38
+    0000003f:   movq    %rax, (%rsi)
+    00000042:   retq # Offset: 66
+
+DWARF CFI Instructions:
+    <empty>
+End of Function "usqrt"
index 526458e..87d437c 100644 (file)
@@ -80,6 +80,7 @@ tools = [
     ToolSubst('llvm-bolt', unresolved='fatal'),
     ToolSubst('llvm-boltdiff', unresolved='fatal'),
     ToolSubst('llvm-bolt-heatmap', unresolved='fatal'),
+    ToolSubst('llvm-bat-dump', unresolved='fatal'),
     ToolSubst('perf2bolt', unresolved='fatal'),
     ToolSubst('yaml2obj', unresolved='fatal'),
     ToolSubst('llvm-mc', unresolved='fatal'),
index 91b00a5..52050fb 100644 (file)
@@ -14,5 +14,6 @@ endmacro()
 
 add_subdirectory(driver)
 add_subdirectory(llvm-bolt-fuzzer)
+add_subdirectory(bat-dump)
 add_subdirectory(merge-fdata)
 add_subdirectory(heatmap)
diff --git a/bolt/tools/bat-dump/CMakeLists.txt b/bolt/tools/bat-dump/CMakeLists.txt
new file mode 100644 (file)
index 0000000..b7a7c45
--- /dev/null
@@ -0,0 +1,15 @@
+set(LLVM_LINK_COMPONENTS
+  Object
+  Support
+  )
+
+add_llvm_tool(llvm-bat-dump
+  bat-dump.cpp
+  )
+
+target_link_libraries(llvm-bat-dump
+  PRIVATE
+  LLVMBOLTProfile
+  )
+
+set_target_properties(llvm-bat-dump PROPERTIES FOLDER "BOLT")
diff --git a/bolt/tools/bat-dump/bat-dump.cpp b/bolt/tools/bat-dump/bat-dump.cpp
new file mode 100644 (file)
index 0000000..87a396e
--- /dev/null
@@ -0,0 +1,175 @@
+#include "bolt/Profile/BoltAddressTranslation.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/ADT/iterator_range.h"
+#include "llvm/Object/Binary.h"
+#include "llvm/Object/ELFObjectFile.h"
+#include "llvm/Object/Error.h"
+#include "llvm/Object/ObjectFile.h"
+#include "llvm/Object/SymbolicFile.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Errc.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/ErrorOr.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Program.h"
+#include "llvm/Support/raw_ostream.h"
+#include <algorithm>
+#include <assert.h>
+#include <cstdint>
+#include <map>
+#include <stdlib.h>
+#include <string>
+#include <system_error>
+#include <type_traits>
+#include <utility>
+#include <vector>
+
+using namespace llvm;
+using namespace bolt;
+
+namespace opts {
+
+cl::OptionCategory BatDumpCategory("BAT dump options");
+
+static cl::OptionCategory *BatDumpCategories[] = {&BatDumpCategory};
+
+static cl::opt<std::string> InputFilename(cl::Positional,
+                                          cl::desc("<executable>"),
+                                          cl::Required,
+                                          cl::cat(BatDumpCategory));
+
+static cl::list<uint64_t> Translate("translate",
+                                    cl::desc("translate addresses using BAT"),
+                                    cl::value_desc("addr"),
+                                    cl::cat(BatDumpCategory));
+
+static cl::opt<bool> DumpAll("dump-all", cl::desc("dump all BAT tables"),
+                             cl::cat(BatDumpCategory));
+
+} // namespace opts
+
+static StringRef ToolName;
+
+static void report_error(StringRef Message, std::error_code EC) {
+  assert(EC);
+  errs() << ToolName << ": '" << Message << "': " << EC.message() << ".\n";
+  exit(1);
+}
+
+static void report_error(StringRef Message, Error E) {
+  assert(E);
+  errs() << ToolName << ": '" << Message << "': " << toString(std::move(E))
+         << ".\n";
+  exit(1);
+}
+
+static std::string GetExecutablePath(const char *Argv0) {
+  SmallString<256> ExecutablePath(Argv0);
+  // Do a PATH lookup if Argv0 isn't a valid path.
+  if (!llvm::sys::fs::exists(ExecutablePath))
+    if (llvm::ErrorOr<std::string> P =
+            llvm::sys::findProgramByName(ExecutablePath))
+      ExecutablePath = *P;
+  return std::string(ExecutablePath.str());
+}
+
+void dumpBATFor(llvm::object::ELFObjectFileBase *InputFile) {
+  BoltAddressTranslation BAT;
+  if (!BAT.enabledFor(InputFile)) {
+    errs() << "error: no BAT table found.\n";
+    exit(1);
+  }
+
+  // Look for BAT section
+  bool Found = false;
+  StringRef SectionContents;
+  for (const llvm::object::SectionRef &Section : InputFile->sections()) {
+    Expected<StringRef> SectionNameOrErr = Section.getName();
+    if (Error E = SectionNameOrErr.takeError())
+      continue;
+
+    if (SectionNameOrErr.get() != BoltAddressTranslation::SECTION_NAME)
+      continue;
+
+    Found = true;
+    Expected<StringRef> ContentsOrErr = Section.getContents();
+    if (Error E = ContentsOrErr.takeError())
+      continue;
+    SectionContents = ContentsOrErr.get();
+  }
+
+  if (!Found) {
+    errs() << "BOLT-ERROR: failed to parse BOLT address translation "
+              "table. No BAT section found\n";
+    exit(1);
+  }
+
+  if (std::error_code EC = BAT.parse(SectionContents)) {
+    errs() << "BOLT-ERROR: failed to parse BOLT address translation "
+              "table. Malformed BAT section\n";
+    exit(1);
+  }
+
+  if (opts::DumpAll)
+    BAT.dump(outs());
+
+  if (!opts::Translate.empty()) {
+    // Build map of <Address, SymbolName> for InputFile
+    std::map<uint64_t, StringRef> FunctionsMap;
+    for (const llvm::object::ELFSymbolRef &Symbol : InputFile->symbols()) {
+      Expected<StringRef> NameOrError = Symbol.getName();
+      if (NameOrError.takeError())
+        continue;
+      if (cantFail(Symbol.getType()) != llvm::object::SymbolRef::ST_Function)
+        continue;
+      const StringRef Name = *NameOrError;
+      const uint64_t Address = cantFail(Symbol.getAddress());
+      FunctionsMap[Address] = Name;
+    }
+
+    outs() << "Translating addresses according to parsed BAT tables:\n";
+    for (uint64_t Address : opts::Translate) {
+      auto FI = FunctionsMap.upper_bound(Address);
+      if (FI == FunctionsMap.begin()) {
+        outs() << "No function symbol found for 0x" << Twine::utohexstr(Address)
+               << "\n";
+        continue;
+      }
+      --FI;
+      outs() << "0x" << Twine::utohexstr(Address) << " -> " << FI->second
+             << " + 0x"
+             << Twine::utohexstr(
+                    BAT.translate(FI->first, Address - FI->first, false))
+             << "\n";
+    }
+  }
+}
+
+int main(int argc, char **argv) {
+  cl::HideUnrelatedOptions(makeArrayRef(opts::BatDumpCategories));
+  cl::ParseCommandLineOptions(argc, argv, "");
+
+  if (!sys::fs::exists(opts::InputFilename))
+    report_error(opts::InputFilename, errc::no_such_file_or_directory);
+
+  ToolName = argv[0];
+  std::string ToolPath = GetExecutablePath(argv[0]);
+  Expected<llvm::object::OwningBinary<llvm::object::Binary>> BinaryOrErr =
+      llvm::object::createBinary(opts::InputFilename);
+  if (Error E = BinaryOrErr.takeError())
+    report_error(opts::InputFilename, std::move(E));
+  llvm::object::Binary &Binary = *BinaryOrErr.get().getBinary();
+
+  if (auto *InputFile = dyn_cast<llvm::object::ELFObjectFileBase>(&Binary))
+    dumpBATFor(InputFile);
+  else
+    report_error(opts::InputFilename,
+                 llvm::object::object_error::invalid_file_type);
+
+  return EXIT_SUCCESS;
+}