[CSSPGO][llvm-profgen] Disassemble text sections

[wlei]

Summary

This stack of changes introduces llvm-profgen utility which generates a profile data file from given perf script data files for sample-based PGO. It’s part of(not only) the CSSPGO work. Specifically to support context-sensitive with/without pseudo probe profile, it implements a series of functionalities including perf trace parsing, instruction symbolization, LBR stack/call frame stack unwinding, pseudo probe decoding, etc. Also high throughput is achieved by multiple levels of sample aggregation and compatible format with one stop is generated at the end. Please refer to: https://groups.google.com/g/llvm-dev/c/1p1rdYbL93s for the CSSPGO RFC.

This change enables disassembling the text sections to build various address maps that are potentially used by the virtual unwinder. A switch --show-disassembly is being added to print the disassembly code.

Like the llvm-objdump tool, this change leverages existing LLVM components to parse and disassemble ELF binary files. So far X86 is supported.

Test Plan:

ninja check-llvm

[wlei]

rebase

[wmi]

llvm/test/tools/llvm-profgen/disassemble.s

24

Can we have a test with multiple .text sections and some of the sections have suffixes like ".text.hot" or ".text.unlikely"?

llvm/tools/llvm-profgen/ProfiledBinary.cpp

53

Can we use MCInstrDesc::isCall()?

73

Can we use MCInstrDesc::isReturn()?

127

Can we reverse the condition here?
auto Obj = dyn_cast<ELFObjectFileBase>(&Binary)
if (!Obj)

exitWithError("not a valid Elf image", Path);

...

261–292

Can we wrap it in a function like disassembleSymbol, to make the disassemble function smaller?

[wlei]

address reviewer's comments

[wlei]

llvm/test/tools/llvm-profgen/disassemble.s

24

@wmi Thanks for your helpful suggestions.
Modified this test with multiple .text sections and ".text.hot" and ".text.unlikely"

llvm/tools/llvm-profgen/ProfiledBinary.cpp

53

fixed, thanks!

73

fixed!

127

Good catch, fixed!

261–292

Yeah, It seems there are so many arguments after the refactoring, so I changed the scope of those MC ptrs, move into the class scope, also wrap them into setUpDisassembler function.

cc: @hoy

[hoy]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

115

No need to pass this in as a parameter. Can just use the member field PreferredBaseAddress instead.

117

StringRef type doesn't need to be passed by reference as a convention.

117

The function can return false upon a disassemble error and let the caller handle the error so that FileName is not needed here.

[wlei]

address reviewer's comments

[wlei]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

115

Thanks for your suggestions, all fixed

[wenlei]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

60

Can we sink the 4 getELFImageLMAForSec manually to avoid duplication and make it more readable?

[wenlei]

If it's not too much work, I'd suggest squash the refactoring of your later patch (RVA->Offset etc from https://reviews.llvm.org/D89715) into this one. Upstream patch does not have to map 1:1 to internal patch - squashing, minimizing code churn and grouping by functionality helps make it easier for other reviewers. Thanks.

[wlei]

move some renaming work into this from the child diff

[wlei]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

60

Seems we can't merge those IF together by "||", I tried and got error.
you can see in llvm/lib/DebugInfo/Symbolize/Symbolize.cpp, they also use the same style.

Optional<ArrayRef<uint8_t>> getBuildID(const ELFObjectFileBase *Obj) {
  Optional<ArrayRef<uint8_t>> BuildID;
  if (auto *O = dyn_cast<ELFObjectFile<ELF32LE>>(Obj))
    BuildID = getBuildID(O->getELFFile());
  else if (auto *O = dyn_cast<ELFObjectFile<ELF32BE>>(Obj))
    BuildID = getBuildID(O->getELFFile());
  else if (auto *O = dyn_cast<ELFObjectFile<ELF64LE>>(Obj))
    BuildID = getBuildID(O->getELFFile());
  else if (auto *O = dyn_cast<ELFObjectFile<ELF64BE>>(Obj))
    BuildID = getBuildID(O->getELFFile());
  else
    llvm_unreachable("unsupported file format");
  return BuildID;
}

(https://github.com/llvm/llvm-project/blob/master/llvm/lib/DebugInfo/Symbolize/Symbolize.cpp)

[wlei]

rebase

[wmi]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

233

Can we set PreferredBaseAddress here and don't have to iterate all the sections in setPreferredBaseAddress function?
PreferredBaseAddress = getELFImageLMAForSec(Section);

[wlei]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

233

Thanks for your suggestion. Considering the section number should not be very large, we intend to decouple them to make disassemble function focus on disassembling things and PreferredBaseAddress would also be used in other functions so we used an explicit set-up for better readability.

[wlei]

fix clang-tidy warning

[wmi]

LGTM.

llvm/tools/llvm-profgen/ProfiledBinary.cpp

233

Ok, thanks.

[wenlei]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

60

What I was thinking about was to hoist getELFImageLMAForSec only. But not a big deal, current one is fine too.

auto *ELFObj = dyn_cast<ELF32LEObjectFile>(Sec.getObject());
if (!ELFObj)
  ELFObj = dyn_cast<ELF32BEObjectFile>(Sec.getObject());
if (!ELFObj)
  ELFObj = dyn_cast<ELF64BEObjectFile>(Sec.getObject());
if (!ELFObj)
  ELFObj = dyn_cast<ELF64BEObjectFile>(Sec.getObject());

return getELFImageLMAForSec(ELFObj->getELFFile(), Sec, ELFObj->getFileName());

llvm/docs/CommandGuide/llvm-profgen.rst

 --------
 At least one of the following commands are required:
 
 .. option:: --perfscript=<string[,string,...]>
 
   Path of perf-script trace created by Linux perf tool with `script`
   command(the raw perf.data should be profiled with -b).
 
+.. option:: --binary=<string[,string,...]>
+
+  Path of the input profiled binary files.
+
 .. option:: --output=<string>
 
   Path of the output profile file.
 
 OPTIONS
 -------
 :program:`llvm-profgen` supports the following options:
 
-.. option:: --binary=<string[,string,...]>
-
-  Path of the input profiled binary files. If no file path is specified, the
-  path of the actual profiled binaries will be used instead.
-
 .. option:: --show-mmap-events
 
   Print mmap events.
+
+.. option:: --show-disassembly
+
+  Print disassembled code.
+
+.. option:: --x86-asm-syntax=[att|intel]
+
+  Specify whether to print assembly code in AT&T syntax (the default) or Intel
+  syntax.

llvm/test/tools/llvm-profgen/disassemble.s

+# REQUIRES: x86-registered-target
+# RUN: llvm-mc -filetype=obj -triple=x86_64 %s -o %t
+# RUN: llvm-profgen --binary=%t --perfscript=%s --output=%t1 -show-disassembly -x86-asm-syntax=intel | FileCheck %s --match-full-lines
+
+# CHECK: Disassembly of section .text [0x0, 0x66]:
+# CHECK: <foo1>:
+# CHECK:        0:  push  rbp
+# CHECK:        1:  mov rbp, rsp
+# CHECK:        4:  sub rsp, 16
+# CHECK:        8:  mov dword ptr [rbp - 4], 0
+# CHECK:        f:  mov edi, 1
+# CHECK:       14:  call  0x19
+# CHECK:       19:  mov edi, 2
+# CHECK:       1e:  mov dword ptr [rbp - 8], eax
+# CHECK:       21:  call  0x26
+# CHECK:       26:  mov ecx, dword ptr [rbp - 8]
+# CHECK:       29:  add ecx, eax
+# CHECK:       2b:  mov eax, ecx
+# CHECK:       2d:  add rsp, 16
+# CHECK:       31:  pop rbp
+# CHECK:       32:  ret
+
+# CHECK: <foo2>:
+# CHECK:       33:  push  rbp
+# CHECK:       34:  mov rbp, rsp
+# CHECK:       37:  sub rsp, 16
+# CHECK:       3b:  mov dword ptr [rbp - 4], 0
+# CHECK:       42:  mov edi, 1
+# CHECK:       47:  call  0x4c
+# CHECK:       4c:  mov edi, 2
+# CHECK:       51:  mov dword ptr [rbp - 8], eax
+# CHECK:       54:  call  0x59
+# CHECK:       59:  mov ecx, dword ptr [rbp - 8]
+# CHECK:       5c:  add ecx, eax
+# CHECK:       5e:  mov eax, ecx
+# CHECK:       60:  add rsp, 16
+# CHECK:       64:  pop rbp
+# CHECK:       65:  ret
+
+
+
+.section .text
+foo1:
+  pushq %rbp
+  movq  %rsp, %rbp
+  subq  $16, %rsp
+  movl  $0, -4(%rbp)
+  movl  $1, %edi
+  callq _Z5funcAi
+  movl  $2, %edi
+  movl  %eax, -8(%rbp)
+  callq _Z5funcBi
+  movl  -8(%rbp), %ecx
+  addl  %eax, %ecx
+  movl  %ecx, %eax
+  addq  $16, %rsp
+  popq  %rbp
+  retq
+
+.section .text
+foo2:
+  pushq %rbp
+  movq  %rsp, %rbp
+  subq  $16, %rsp
+  movl  $0, -4(%rbp)
+  movl  $1, %edi
+  callq _Z5funcBi
+  movl  $2, %edi
+  movl  %eax, -8(%rbp)
+  callq _Z5funcAi
+  movl  -8(%rbp), %ecx
+  addl  %eax, %ecx
+  movl  %ecx, %eax
+  addq  $16, %rsp
+  popq  %rbp
+  retq
+
+# CHECK: Disassembly of section .text.hot [0x0, 0x12]:
+# CHECK: <bar>:
+# CHECK:        0:  push  rbp
+# CHECK:        1:  mov rbp, rsp
+# CHECK:        4:  mov dword ptr [rbp - 4], edi
+# CHECK:        7:  mov dword ptr [rbp - 8], esi
+# CHECK:        a:  mov eax, dword ptr [rbp - 4]
+# CHECK:        d:  add eax, dword ptr [rbp - 8]
+# CHECK:       10:  pop rbp
+# CHECK:       11:  ret
+
+.section .text.hot
+bar:
+  pushq %rbp
+  movq  %rsp, %rbp
+  movl  %edi, -4(%rbp)
+  movl  %esi, -8(%rbp)
+  movl  -4(%rbp), %eax
+  addl  -8(%rbp), %eax
+  popq  %rbp
+  retq
+
+
+# CHECK: Disassembly of section .text.unlikely [0x0, 0x12]:
+# CHECK: <baz>:
+# CHECK:        0:  push  rbp
+# CHECK:        1:  mov rbp, rsp
+# CHECK:        4:  mov dword ptr [rbp - 4], edi
+# CHECK:        7:  mov dword ptr [rbp - 8], esi
+# CHECK:        a:  mov eax, dword ptr [rbp - 4]
+# CHECK:        d:  sub eax, dword ptr [rbp - 8]
+# CHECK:       10:  pop rbp
+# CHECK:       11:  ret
+
+.section .text.unlikely
+baz:
+  pushq %rbp
+  movq  %rsp, %rbp
+  movl  %edi, -4(%rbp)
+  movl  %esi, -8(%rbp)
+  movl  -4(%rbp), %eax
+  subl  -8(%rbp), %eax
+  popq  %rbp
+  retq

llvm/test/tools/llvm-profgen/lit.local.cfg

 import subprocess
 import lit.util
 
 config.suffixes = ['.test', '.ll', '.s', '.yaml']
 if not lit.util.which("llvm-profgen", config.llvm_tools_dir):
-  config.unsupported = True  
+  config.unsupported = True
+
+if not 'X86' in config.root.targets:
+  config.unsupported = True

llvm/test/tools/llvm-profgen/mmapEvent.test

-; RUN: llvm-profgen --perfscript=%s --output=%t --show-mmap-events | FileCheck %s
+; REQUIRES: x86-registered-target
+; RUN: llvm-mc -filetype=obj -triple=x86_64 %S/disassemble.s -o %t
+; RUN: llvm-profgen --perfscript=%s --binary=%t --output=%t --show-mmap-events | FileCheck %s
 
 PERF_RECORD_MMAP2 2580483/2580483: [0x400000(0x1000) @ 0 103:01 539973862 1972407324]: r-xp /home/a.out
 PERF_RECORD_MMAP2 2580483/2580483: [0x7f2505b40000(0x224000) @ 0 08:04 19532214 4169021329]: r-xp /usr/lib64/ld-2.17.so
 PERF_RECORD_MMAP2 2580483/2580483: [0x7ffe88097000(0x1000) @ 0 00:00 0 0]: r-xp [vdso]
 PERF_RECORD_MMAP2 2580483/2580483: [0x7f2505d56000(0xa000) @ 0 08:04 19530021 4190740662]: r-xp /usr/lib64/perf_fopen_hook.so
 PERF_RECORD_MMAP2 2580483/2580483: [0x7f250593c000(0x204000) @ 0 08:04 19532229 3585508847]: r-xp /usr/lib64/libdl-2.17.so
 PERF_RECORD_MMAP2 2580483/2580483: [0x7f250556e000(0x3ce000) @ 0 08:04 19532221 4003737677]: r-xp /usr/lib64/libc-2.17.so
 PERF_RECORD_MMAP2 2580483/2580483: [0x7f2505358000(0x216000) @ 0 08:04 19534595 2609212015]: r-xp /usr/lib64/libz.so.1.2.7

llvm/tools/llvm-profgen/CMakeLists.txt

+include_directories(
+  ${LLVM_MAIN_SRC_DIR}/lib/Target/X86
+  ${LLVM_BINARY_DIR}/lib/Target/X86
+  )
 set(LLVM_LINK_COMPONENTS
+  AllTargetsDescs
+  AllTargetsDisassemblers
   Core
-  ProfileData
+  MC
+  MCDisassembler
+  Object
   Support
-  Symbolize
   )
 
 add_llvm_tool(llvm-profgen
   llvm-profgen.cpp
   PerfReader.cpp
+  ProfiledBinary.cpp
   )

llvm/tools/llvm-profgen/ErrorHandling.h

 LLVM_ATTRIBUTE_NORETURN inline void
 exitWithError(std::error_code EC, StringRef Whence = StringRef()) {
   exitWithError(EC.message(), Whence);
 }
 
 LLVM_ATTRIBUTE_NORETURN inline void exitWithError(Error E, StringRef Whence) {
   exitWithError(errorToErrorCode(std::move(E)), Whence);
 }
+
+template <typename T, typename... Ts>
+T unwrapOrError(Expected<T> EO, Ts &&... Args) {
+  if (EO)
+    return std::move(*EO);
+  exitWithError(EO.takeError(), std::forward<Ts>(Args)...);
+}
 #endif

llvm/tools/llvm-profgen/LLVMBuild.txt

 ;   http://llvm.org/docs/LLVMBuild.html
 ;
 ;===------------------------------------------------------------------------===;
 
 [component_0]
 type = Tool
 name = llvm-profgen
 parent = Tools
-required_libraries = Support
+required_libraries = DebugInfoDWARF MC MCDisassembler MCParser Object all-targets Demangle Support

llvm/tools/llvm-profgen/PerfReader.cpp

 }
 
 void PerfReader::parsePerfTraces(cl::list<std::string> &PerfTraceFilenames) {
   // Parse perf traces.
   for (auto Filename : PerfTraceFilenames)
     parseTrace(Filename);
 }
 
-} // namespace sampleprof
-} // namespace llvm
+} // end namespace sampleprof
+} // end namespace llvm

llvm/tools/llvm-profgen/ProfiledBinary.h

 //===-- ProfiledBinary.h - Binary decoder -----------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TOOLS_LLVM_PROFGEN_PROFILEDBINARY_H
 #define LLVM_TOOLS_LLVM_PROFGEN_PROFILEDBINARY_H
 #include "llvm/ADT/StringRef.h"
+#include "llvm/MC/MCAsmInfo.h"
+#include "llvm/MC/MCContext.h"
+#include "llvm/MC/MCDisassembler/MCDisassembler.h"
+#include "llvm/MC/MCInst.h"
+#include "llvm/MC/MCInstPrinter.h"
+#include "llvm/MC/MCInstrAnalysis.h"
+#include "llvm/MC/MCInstrInfo.h"
+#include "llvm/MC/MCObjectFileInfo.h"
+#include "llvm/MC/MCRegisterInfo.h"
+#include "llvm/MC/MCSubtargetInfo.h"
+#include "llvm/MC/MCTargetOptions.h"
+#include "llvm/Object/ELFObjectFile.h"
 #include "llvm/Support/Path.h"
+#include <set>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+using namespace llvm::object;
 
 namespace llvm {
 namespace sampleprof {
 
 class ProfiledBinary {
+  // Absolute path of the binary.
   std::string Path;
+  // The target triple.
+  Triple TheTriple;
+  // The runtime base address that the executable sections are loaded at.
   mutable uint64_t BaseAddress = 0;
+  // The preferred base address that the executable sections are loaded at.
+  uint64_t PreferredBaseAddress = 0;
+  // Mutiple MC component info
+  std::unique_ptr<const MCRegisterInfo> MRI;
+  std::unique_ptr<const MCAsmInfo> AsmInfo;
+  std::unique_ptr<const MCSubtargetInfo> STI;
+  std::unique_ptr<const MCInstrInfo> MII;
+  std::unique_ptr<MCDisassembler> DisAsm;
+  std::unique_ptr<const MCInstrAnalysis> MIA;
+  std::unique_ptr<MCInstPrinter> IP;
+  // A list of text sections sorted by start RVA and size. Used to check
+  // if a given RVA is a valid code address.
+  std::set<std::pair<uint64_t, uint64_t>> TextSections;
+  // Function offset to name mapping.
+  std::unordered_map<uint64_t, std::string> FuncStartAddrMap;
+  // An array of offsets of all instructions sorted in increasing order. The
+  // sorting is needed to fast advance to the next forward/backward instruction.
+  std::vector<uint64_t> CodeAddrs;
+  // A set of call instruction offsets. Used by virtual unwinding.
+  std::unordered_set<uint64_t> CallAddrs;
+  // A set of return instruction offsets. Used by virtual unwinding.
+  std::unordered_set<uint64_t> RetAddrs;
+
+  void setPreferredBaseAddress(const ELFObjectFileBase *O);
+
+  // Set up disassembler and related components.
+  void setUpDisassembler(const ELFObjectFileBase *Obj);
+
+  /// Dissassemble the text section and build various address maps.
+  void disassemble(const ELFObjectFileBase *O);
+
+  /// Helper function to dissassemble the symbol and extract info for unwinding
+  bool dissassembleSymbol(std::size_t SI, ArrayRef<uint8_t> Bytes,
+                          SectionSymbolsTy &Symbols, const SectionRef &Section);
+
+  /// Decode the interesting parts of the binary and build internal data
+  /// structures. On high level, the parts of interest are:
+  ///   1. Text sections, including the main code section and the PLT
+  ///   entries that will be used to handle cross-module call transitions.
+  ///   2. The .debug_line section, used by Dwarf-based profile generation.
+  ///   3. Pseudo probe related sections, used by probe-based profile
+  ///   generation.
+  void load();
 
 public:
   ProfiledBinary(StringRef Path) : Path(Path) { load(); }
 
   const StringRef getPath() const { return Path; }
   const StringRef getName() const { return llvm::sys::path::filename(Path); }
   uint64_t getBaseAddress() const { return BaseAddress; }
   void setBaseAddress(uint64_t Address) { BaseAddress = Address; }
-
-private:
-  void load() {
-    // TODO:
-  }
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/ProfiledBinary.cpp

+//===-- ProfiledBinary.cpp - Binary decoder ---------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "ProfiledBinary.h"
+#include "ErrorHandling.h"
+#include "MCTargetDesc/X86MCTargetDesc.h"
+#include "llvm/ADT/Triple.h"
+#include "llvm/Demangle/Demangle.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/TargetRegistry.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Support/X86TargetParser.h"
+
+#define DEBUG_TYPE "load-binary"
+
+using namespace llvm;
+
+static cl::opt<bool> ShowDisassembly("show-disassembly", cl::ReallyHidden,
+                                     cl::init(false), cl::ZeroOrMore,
+                                     cl::desc("Print disassembled code."));
+
+namespace llvm {
+namespace sampleprof {
+
+static const Target *getTarget(const ObjectFile *Obj) {
+  Triple TheTriple = Obj->makeTriple();
+  std::string Error;
+  std::string ArchName;
+  const Target *TheTarget =
+      TargetRegistry::lookupTarget(ArchName, TheTriple, Error);
+  if (!TheTarget)
+    exitWithError(Error, Obj->getFileName());
+  return TheTarget;
+}
+
+template <class ELFT>
+static uint64_t getELFImageLMAForSec(const ELFFile<ELFT> *Obj,
+                                     const object::ELFSectionRef &Sec,
+                                     StringRef FileName) {
+  // Search for a PT_LOAD segment containing the requested section. Return this
+  // segment's p_addr as the image load address for the section.
+  const auto &PhdrRange = unwrapOrError(Obj->program_headers(), FileName);
+  for (const typename ELFT::Phdr &Phdr : PhdrRange)
+    if ((Phdr.p_type == ELF::PT_LOAD) && (Phdr.p_vaddr <= Sec.getAddress()) &&
+        (Phdr.p_vaddr + Phdr.p_memsz > Sec.getAddress()))
+      // Segments will always be loaded at a page boundary.
+      return Phdr.p_paddr & ~(Phdr.p_align - 1U);
+  return 0;
+}
+
+// Get the image load address for a specific section. Note that an image is
+// loaded by segments (a group of sections) and segments may not be consecutive
+// in memory.
+static uint64_t getELFImageLMAForSec(const object::ELFSectionRef &Sec) {
+  if (const auto *ELFObj = dyn_cast<ELF32LEObjectFile>(Sec.getObject()))
+    return getELFImageLMAForSec(ELFObj->getELFFile(), Sec,
+                                ELFObj->getFileName());
+  else if (const auto *ELFObj = dyn_cast<ELF32BEObjectFile>(Sec.getObject()))
+    return getELFImageLMAForSec(ELFObj->getELFFile(), Sec,
+                                ELFObj->getFileName());
+  else if (const auto *ELFObj = dyn_cast<ELF64LEObjectFile>(Sec.getObject()))
+    return getELFImageLMAForSec(ELFObj->getELFFile(), Sec,
+                                ELFObj->getFileName());
+  const auto *ELFObj = cast<ELF64BEObjectFile>(Sec.getObject());
+  return getELFImageLMAForSec(ELFObj->getELFFile(), Sec, ELFObj->getFileName());
+}
+
+void ProfiledBinary::load() {
+  // Attempt to open the binary.
+  OwningBinary<Binary> OBinary = unwrapOrError(createBinary(Path), Path);
+  Binary &Binary = *OBinary.getBinary();
+
+  auto *Obj = dyn_cast<ELFObjectFileBase>(&Binary);
+  if (!Obj)
+    exitWithError("not a valid Elf image", Path);
+
+  TheTriple = Obj->makeTriple();
+  // Current only support X86
+  if (!TheTriple.isX86())
+    exitWithError("unsupported target", TheTriple.getTriple());
+  LLVM_DEBUG(dbgs() << "Loading " << Path << "\n");
+
+  // Find the preferred base address for text sections.
+  setPreferredBaseAddress(Obj);
+
+  // Disassemble the text sections.
+  disassemble(Obj);
+
+  // TODO: decode other sections.
+
+  return;
+}
+
+void ProfiledBinary::setPreferredBaseAddress(const ELFObjectFileBase *Obj) {
+  for (section_iterator SI = Obj->section_begin(), SE = Obj->section_end();
+       SI != SE; ++SI) {
+    const SectionRef &Section = *SI;
+    if (Section.isText()) {
+      PreferredBaseAddress = getELFImageLMAForSec(Section);
+      return;
+    }
+  }
+  exitWithError("no text section found", Obj->getFileName());
+}
+
+bool ProfiledBinary::dissassembleSymbol(std::size_t SI, ArrayRef<uint8_t> Bytes,
+                                        SectionSymbolsTy &Symbols,
+                                        const SectionRef &Section) {
+
+  std::size_t SE = Symbols.size();
+  uint64_t SectionOffset = Section.getAddress() - PreferredBaseAddress;
+  uint64_t SectSize = Section.getSize();
+  uint64_t StartOffset = Symbols[SI].Addr - PreferredBaseAddress;
+  uint64_t EndOffset = (SI + 1 < SE)
+                           ? Symbols[SI + 1].Addr - PreferredBaseAddress
+                           : SectionOffset + SectSize;
+  if (StartOffset >= EndOffset)
+    return true;
+
+  std::string &&SymbolName = Symbols[SI].Name.str();
+  if (ShowDisassembly)
+    outs() << '<' << SymbolName << ">:\n";
+
+  uint64_t Offset = StartOffset;
+  while (Offset < EndOffset) {
+    MCInst Inst;
+    uint64_t Size;
+    // Disassemble an instruction.
+    if (!DisAsm->getInstruction(Inst, Size, Bytes.slice(Offset - SectionOffset),
+                                Offset + PreferredBaseAddress, nulls()))
+      return false;
+
+    if (ShowDisassembly) {
+      outs() << format("%8" PRIx64 ":", Offset);
+      IP->printInst(&Inst, Offset + Size, "", *STI.get(), outs());
+      outs() << "\n";
+    }
+
+    const MCInstrDesc &MCDesc = MII->get(Inst.getOpcode());
+
+    // Populate address maps.
+    CodeAddrs.push_back(Offset);
+    if (MCDesc.isCall())
+      CallAddrs.insert(Offset);
+    else if (MCDesc.isReturn())
+      RetAddrs.insert(Offset);
+
+    Offset += Size;
+  }
+
+  if (ShowDisassembly)
+    outs() << "\n";
+
+  FuncStartAddrMap[StartOffset] = Symbols[SI].Name.str();
+  return true;
+}
+
+void ProfiledBinary::setUpDisassembler(const ELFObjectFileBase *Obj) {
+  const Target *TheTarget = getTarget(Obj);
+  std::string TripleName = TheTriple.getTriple();
+  StringRef FileName = Obj->getFileName();
+
+  MRI.reset(TheTarget->createMCRegInfo(TripleName));
+  if (!MRI)
+    exitWithError("no register info for target " + TripleName, FileName);
+
+  MCTargetOptions MCOptions;
+  AsmInfo.reset(TheTarget->createMCAsmInfo(*MRI, TripleName, MCOptions));
+  if (!AsmInfo)
+    exitWithError("no assembly info for target " + TripleName, FileName);
+
+  SubtargetFeatures Features = Obj->getFeatures();
+  STI.reset(
+      TheTarget->createMCSubtargetInfo(TripleName, "", Features.getString()));
+  if (!STI)
+    exitWithError("no subtarget info for target " + TripleName, FileName);
+
+  MII.reset(TheTarget->createMCInstrInfo());
+  if (!MII)
+    exitWithError("no instruction info for target " + TripleName, FileName);
+
+  MCObjectFileInfo MOFI;
+  MCContext Ctx(AsmInfo.get(), MRI.get(), &MOFI);
+  MOFI.InitMCObjectFileInfo(Triple(TripleName), false, Ctx);
+  DisAsm.reset(TheTarget->createMCDisassembler(*STI, Ctx));
+  if (!DisAsm)
+    exitWithError("no disassembler for target " + TripleName, FileName);
+
+  MIA.reset(TheTarget->createMCInstrAnalysis(MII.get()));
+
+  int AsmPrinterVariant = AsmInfo->getAssemblerDialect();
+  IP.reset(TheTarget->createMCInstPrinter(Triple(TripleName), AsmPrinterVariant,
+                                          *AsmInfo, *MII, *MRI));
+  IP->setPrintBranchImmAsAddress(true);
+}
+
+void ProfiledBinary::disassemble(const ELFObjectFileBase *Obj) {
+  // Set up disassembler and related components.
+  setUpDisassembler(Obj);
+
+  // Create a mapping from virtual address to symbol name. The symbols in text
+  // sections are the candidates to dissassemble.
+  std::map<SectionRef, SectionSymbolsTy> AllSymbols;
+  StringRef FileName = Obj->getFileName();
+  for (const SymbolRef &Symbol : Obj->symbols()) {
+    const uint64_t Addr = unwrapOrError(Symbol.getAddress(), FileName);
+    const StringRef Name = unwrapOrError(Symbol.getName(), FileName);
+    section_iterator SecI = unwrapOrError(Symbol.getSection(), FileName);
+    if (SecI != Obj->section_end())
+      AllSymbols[*SecI].push_back(SymbolInfoTy(Addr, Name, ELF::STT_NOTYPE));
+  }
+
+  // Sort all the symbols. Use a stable sort to stabilize the output.
+  for (std::pair<const SectionRef, SectionSymbolsTy> &SecSyms : AllSymbols)
+    stable_sort(SecSyms.second);
+
+  if (ShowDisassembly)
+    outs() << "\nDisassembly of " << FileName << ":\n";
+
+  // Dissassemble a text section.
+  for (section_iterator SI = Obj->section_begin(), SE = Obj->section_end();
+       SI != SE; ++SI) {
+    const SectionRef &Section = *SI;
+    if (!Section.isText())
+      continue;
+
+    uint64_t ImageLoadAddr = PreferredBaseAddress;
+    uint64_t SectionOffset = Section.getAddress() - ImageLoadAddr;
+    uint64_t SectSize = Section.getSize();
+    if (!SectSize)
+      continue;
+
+    // Register the text section.
+    TextSections.insert({SectionOffset, SectSize});
+
+    if (ShowDisassembly) {
+      StringRef SectionName = unwrapOrError(Section.getName(), FileName);
+      outs() << "\nDisassembly of section " << SectionName;
+      outs() << " [" << format("0x%" PRIx64, SectionOffset) << ", "
+             << format("0x%" PRIx64, SectionOffset + SectSize) << "]:\n\n";
+    }
+
+    // Get the section data.
+    ArrayRef<uint8_t> Bytes =
+        arrayRefFromStringRef(unwrapOrError(Section.getContents(), FileName));
+
+    // Get the list of all the symbols in this section.
+    SectionSymbolsTy &Symbols = AllSymbols[Section];
+
+    // Disassemble symbol by symbol.
+    for (std::size_t SI = 0, SE = Symbols.size(); SI != SE; ++SI) {
+      if (!dissassembleSymbol(SI, Bytes, Symbols, Section))
+        exitWithError("disassembling error", FileName);
+    }
+  }
+}
+} // end namespace sampleprof
+} // end namespace llvm

llvm/tools/llvm-profgen/llvm-profgen.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "ErrorHandling.h"
 #include "PerfReader.h"
 #include "ProfiledBinary.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/TargetSelect.h"
 
 static cl::list<std::string> PerfTraceFilenames(
     "perfscript", cl::value_desc("perfscript"), cl::OneOrMore,
     llvm::cl::MiscFlags::CommaSeparated,
     cl::desc("Path of perf-script trace created by Linux perf tool with "
              "`script` command(the raw perf.data should be profiled with -b)"));
 
 static cl::list<std::string>
-    BinaryFilenames("binary", cl::value_desc("binary"), cl::ZeroOrMore,
+    BinaryFilenames("binary", cl::value_desc("binary"), cl::OneOrMore,
                     llvm::cl::MiscFlags::CommaSeparated,
                     cl::desc("Path of profiled binary files"));
 
 static cl::opt<std::string> OutputFilename("output", cl::value_desc("output"),
                                            cl::Required,
                                            cl::desc("Output profile file"));
 
 using namespace llvm;
 using namespace sampleprof;
 
 int main(int argc, const char *argv[]) {
   InitLLVM X(argc, argv);
 
   cl::ParseCommandLineOptions(argc, argv, "llvm SPGO profile generator\n");
 
+  // Initialize targets and assembly printers/parsers.
+  InitializeAllTargetInfos();
+  InitializeAllTargetMCs();
+  InitializeAllDisassemblers();
+
   // Load binaries and parse perf events and samples
   PerfReader Reader(BinaryFilenames);
   Reader.parsePerfTraces(PerfTraceFilenames);
 
   return EXIT_SUCCESS;
 }