[XRay] Implement llvm-xray extract, start of the llvm-xray tool

Usage:

  llvm-xray extract <object file> [-o <filename or '-'>]

The tool gets the XRay instrumentation map from an object file and turns
it into YAML.  We first support ELF64 sleds on x86_64 binaries, with
provision for supporting other supported platforms and formats later.

This is the first of a many-part change to fully implement the
`llvm-xray` tool.

We also define a subcommand registration and dispatch mechanism to be
used by other further subcommand implementations for llvm-xray.

Diffusion Revision: https://reviews.llvm.org/D21987

llvm-svn: 285165

test/lit.cfg

@@ -310,6 +310,7 @@ for pattern in [r"\bbugpoint\b(?!-)",
                 r"\bllvm-tblgen\b",
                 r"\bllvm-c-test\b",
                 r"\bllvm-cxxfilt\b",
+                r"\bllvm-xray\b",
                 NOJUNK + r"\bllvm-symbolizer\b",
                 NOJUNK + r"\bopt\b",
                 r"\bFileCheck\b",

test/tools/llvm-xray/X86/bad-instrmap-sizes.bin

@@ -0,0 +1,3 @@
+; RUN: not llvm-xray extract %S/Inputs/elf64-badentrysizes.bin 2>&1 | FileCheck %s
+; CHECK: llvm-xray: Cannot extract instrumentation map from '{{.*}}elf64-badentrysizes.bin'.
+; CHECK-NEXT: Instrumentation map entries not evenly divisible by size of an XRay sled entry in ELF64.

test/tools/llvm-xray/X86/empty.txt

@@ -0,0 +1,4 @@
+; RUN: not llvm-xray extract %S/Inputs/empty-file.bin 2>&1 | FileCheck %s
+
+; CHECK: llvm-xray: Cannot extract instrumentation map from '{{.*}}empty-file.bin'.
+; CHECK-NEXT: The file was not recognized as a valid object file

test/tools/llvm-xray/X86/extract-instrmap.ll

@@ -0,0 +1,15 @@
+; This test makes sure we can extract the instrumentation map from an
+; XRay-instrumented object file.
+;
+; RUN: llvm-xray extract %S/Inputs/elf64-example.bin | FileCheck %s
+
+; CHECK:      ---
+; CHECK-NEXT: - { id: 1, address: 0x000000000041C900, function: 0x000000000041C900, kind: function-enter,
+; CHECK-NEXT:     always-instrument: true }
+; CHECK-NEXT: - { id: 1, address: 0x000000000041C912, function: 0x000000000041C900, kind: function-exit,
+; CHECK-NEXT:     always-instrument: true }
+; CHECK-NEXT: - { id: 2, address: 0x000000000041C930, function: 0x000000000041C930, kind: function-enter,
+; CHECK-NEXT:    always-instrument: true }
+; CHECK-NEXT: - { id: 2, address: 0x000000000041C946, function: 0x000000000041C930, kind: function-exit,
+; CHECK-NEXT:     always-instrument: true }
+; CHECK-NEXT: ...

test/tools/llvm-xray/X86/lit.local.cfg

@@ -0,0 +1 @@
+config.suffixes = ['.yaml', '.ll', '.txt']

test/tools/llvm-xray/X86/no-instr-map.txt

@@ -0,0 +1,4 @@
+; RUN: not llvm-xray extract %S/Inputs/elf64-noinstr-map.bin 2>&1 | FileCheck %s
+
+; CHECK: llvm-xray: Cannot extract instrumentation map from '{{.*}}elf64-noinstr-map.bin'.
+; CHECK-NEXT: Failed to find XRay instrumentation map.

test/tools/llvm-xray/X86/no-such-file.txt

@@ -0,0 +1,4 @@
+; RUN: not llvm-xray extract no-such-file 2>&1 | FileCheck %s
+
+; CHECK: llvm-xray: Cannot extract instrumentation map from 'no-such-file'.
+; CHECK-NEXT: No such file or directory

test/tools/llvm-xray/X86/unsupported-elf32.txt

@@ -0,0 +1,3 @@
+; RUN: not llvm-xray extract %S/Inputs/elf32-noxray.bin 2>&1 | FileCheck %s
+; CHECK: llvm-xray: Cannot extract instrumentation map from '{{.*}}elf32-noxray.bin'.
+; CHECK-NEXT: File format not supported (only does ELF little endian 64-bit).

tools/llvm-xray/CMakeLists.txt

@@ -0,0 +1,10 @@
+set(LLVM_LINK_COMPONENTS
+  ${LLVM_TARGETS_TO_BUILD}
+  Support
+  Object)
+
+set(LLVM_XRAY_TOOLS
+  xray-extract.cc
+  xray-registry.cc)
+
+add_llvm_tool(llvm-xray llvm-xray.cc ${LLVM_XRAY_TOOLS})

tools/llvm-xray/llvm-xray.cc

@@ -0,0 +1,42 @@
+//===- llvm-xray.cc - XRay Tool Main Program ------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the main entry point for the suite of XRay tools. All
+// additional functionality are implemented as subcommands.
+//
+//===----------------------------------------------------------------------===//
+//
+// Basic usage:
+//
+//   llvm-xray [options] <subcommand> [subcommand-specific options]
+//
+#include "xray-registry.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/raw_ostream.h"
+#include <unistd.h>
+
+using namespace llvm;
+using namespace llvm::xray;
+
+int main(int argc, char *argv[]) {
+  cl::ParseCommandLineOptions(argc, argv,
+                              "XRay Tools\n\n"
+                              "  This program consolidates multiple XRay trace "
+                              "processing tools for convenient access.\n");
+  for (auto *SC : cl::getRegisteredSubcommands()) {
+    if (*SC)
+      if (auto C = dispatch(SC)) {
+        ExitOnError("llvm-xray: ")(C());
+        return 0;
+      }
+  }
+
+  cl::PrintHelpMessage(false, true);
+}

tools/llvm-xray/xray-extract.cc

@@ -0,0 +1,236 @@
+//===- xray-extract.cc - XRay Instrumentation Map Extraction --------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Implementation of the xray-extract.h interface.
+//
+// FIXME: Support other XRay-instrumented binary formats other than ELF.
+//
+//===----------------------------------------------------------------------===//
+
+#include <type_traits>
+#include <unistd.h>
+#include <utility>
+
+#include "xray-extract.h"
+
+#include "xray-registry.h"
+#include "xray-sleds.h"
+#include "llvm/Object/ELF.h"
+#include "llvm/Object/ObjectFile.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/DataExtractor.h"
+#include "llvm/Support/ELF.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Format.h"
+#include "llvm/Support/YAMLTraits.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace llvm;
+using namespace llvm::xray;
+using namespace llvm::yaml;
+
+// llvm-xray extract
+// ----------------------------------------------------------------------------
+static cl::SubCommand Extract("extract", "Extract instrumentation maps");
+static cl::opt<std::string> ExtractInput(cl::Positional,
+                                         cl::desc("<input file>"), cl::Required,
+                                         cl::sub(Extract));
+static cl::opt<std::string>
+    ExtractOutput("output", cl::value_desc("output file"), cl::init("-"),
+                  cl::desc("output file; use '-' for stdout"),
+                  cl::sub(Extract));
+static cl::alias ExtractOutput2("o", cl::aliasopt(ExtractOutput),
+                                cl::desc("Alias for -output"),
+                                cl::sub(Extract));
+
+struct YAMLXRaySledEntry {
+  int32_t FuncId;
+  Hex64 Address;
+  Hex64 Function;
+  SledEntry::FunctionKinds Kind;
+  bool AlwaysInstrument;
+};
+
+template <> struct ScalarEnumerationTraits<SledEntry::FunctionKinds> {
+  static void enumeration(IO &IO, SledEntry::FunctionKinds &Kind) {
+    IO.enumCase(Kind, "function-enter", SledEntry::FunctionKinds::ENTRY);
+    IO.enumCase(Kind, "function-exit", SledEntry::FunctionKinds::EXIT);
+    IO.enumCase(Kind, "tail-exit", SledEntry::FunctionKinds::TAIL);
+  }
+};
+
+template <> struct MappingTraits<YAMLXRaySledEntry> {
+  static void mapping(IO &IO, YAMLXRaySledEntry &Entry) {
+    IO.mapRequired("id", Entry.FuncId);
+    IO.mapRequired("address", Entry.Address);
+    IO.mapRequired("function", Entry.Function);
+    IO.mapRequired("kind", Entry.Kind);
+    IO.mapRequired("always-instrument", Entry.AlwaysInstrument);
+  }
+
+  static constexpr bool flow = true;
+};
+
+LLVM_YAML_IS_SEQUENCE_VECTOR(YAMLXRaySledEntry);
+
+namespace {
+
+llvm::Error LoadBinaryInstrELF(
+    StringRef Filename, std::deque<SledEntry> &OutputSleds,
+    InstrumentationMapExtractor::FunctionAddressMap &InstrMap,
+    InstrumentationMapExtractor::FunctionAddressReverseMap &FunctionIds) {
+  auto ObjectFile = object::ObjectFile::createObjectFile(Filename);
+
+  if (!ObjectFile)
+    return ObjectFile.takeError();
+
+  // FIXME: Maybe support other ELF formats. For now, 64-bit Little Endian only.
+  if (!ObjectFile->getBinary()->isELF())
+    return make_error<StringError>(
+        "File format not supported (only does ELF).",
+        std::make_error_code(std::errc::not_supported));
+  if (ObjectFile->getBinary()->getArch() != Triple::x86_64)
+    return make_error<StringError>(
+        "File format not supported (only does ELF little endian 64-bit).",
+        std::make_error_code(std::errc::not_supported));
+
+  // Find the section named "xray_instr_map".
+  StringRef Contents = "";
+  const auto &Sections = ObjectFile->getBinary()->sections();
+  auto I = find_if(Sections, [&](object::SectionRef Section) {
+    StringRef Name = "";
+    if (Section.getName(Name))
+      return false;
+    return Name == "xray_instr_map";
+  });
+  if (I == Sections.end())
+    return make_error<StringError>(
+        "Failed to find XRay instrumentation map.",
+        std::make_error_code(std::errc::not_supported));
+  if (I->getContents(Contents))
+    return make_error<StringError>(
+        "Failed to get contents of 'xray_instr_map' section.",
+        std::make_error_code(std::errc::executable_format_error));
+
+  // Copy the instrumentation map data into the Sleds data structure.
+  auto C = Contents.bytes_begin();
+  static constexpr size_t ELF64SledEntrySize = 32;
+
+  if ((C - Contents.bytes_end()) % ELF64SledEntrySize != 0)
+    return make_error<StringError>(
+        "Instrumentation map entries not evenly divisible by size of an XRay "
+        "sled entry in ELF64.",
+        std::make_error_code(std::errc::executable_format_error));
+
+  int32_t FuncId = 1;
+  uint64_t CurFn = 0;
+  std::deque<SledEntry> Sleds;
+  for (; C != Contents.bytes_end(); C += ELF64SledEntrySize) {
+    DataExtractor Extractor(
+        StringRef(reinterpret_cast<const char *>(C), ELF64SledEntrySize), true,
+        8);
+    Sleds.push_back({});
+    auto &Entry = Sleds.back();
+    uint32_t OffsetPtr = 0;
+    Entry.Address = Extractor.getU64(&OffsetPtr);
+    Entry.Function = Extractor.getU64(&OffsetPtr);
+    auto Kind = Extractor.getU8(&OffsetPtr);
+    switch (Kind) {
+    case 0: // ENTRY
+      Entry.Kind = SledEntry::FunctionKinds::ENTRY;
+      break;
+    case 1: // EXIT
+      Entry.Kind = SledEntry::FunctionKinds::EXIT;
+      break;
+    case 2: // TAIL
+      Entry.Kind = SledEntry::FunctionKinds::TAIL;
+      break;
+    default:
+      return make_error<StringError>(
+          Twine("Encountered unknown sled type ") + "'" + Twine(int32_t{Kind}) +
+              "'.",
+          std::make_error_code(std::errc::protocol_error));
+    }
+    auto AlwaysInstrument = Extractor.getU8(&OffsetPtr);
+    Entry.AlwaysInstrument = AlwaysInstrument != 0;
+
+    // We replicate the function id generation scheme implemented in the runtime
+    // here. Ideally we should be able to break it out, or output this map from
+    // the runtime, but that's a design point we can discuss later on. For now,
+    // we replicate the logic and move on.
+    if (CurFn == 0) {
+      CurFn = Entry.Function;
+      InstrMap[FuncId] = Entry.Function;
+      FunctionIds[Entry.Function] = FuncId;
+    }
+    if (Entry.Function != CurFn) {
+      ++FuncId;
+      CurFn = Entry.Function;
+      InstrMap[FuncId] = Entry.Function;
+      FunctionIds[Entry.Function] = FuncId;
+    }
+  }
+  OutputSleds = std::move(Sleds);
+  return llvm::Error::success();
+}
+
+} // namespace
+
+InstrumentationMapExtractor::InstrumentationMapExtractor(std::string Filename,
+                                                         InputFormats Format,
+                                                         Error &EC) {
+  ErrorAsOutParameter ErrAsOutputParam(&EC);
+  switch (Format) {
+  case InputFormats::ELF: {
+    EC = handleErrors(
+        LoadBinaryInstrELF(Filename, Sleds, FunctionAddresses, FunctionIds),
+        [](std::unique_ptr<ErrorInfoBase> E) {
+          return joinErrors(
+              make_error<StringError>(
+                  Twine("Cannot extract instrumentation map from '") +
+                      ExtractInput + "'.",
+                  std::make_error_code(std::errc::protocol_error)),
+              std::move(E));
+        });
+    break;
+  }
+  default:
+    llvm_unreachable("Input format type not supported yet.");
+    break;
+  }
+}
+
+void InstrumentationMapExtractor::exportAsYAML(raw_ostream &OS) {
+  // First we translate the sleds into the YAMLXRaySledEntry objects in a deque.
+  std::vector<YAMLXRaySledEntry> YAMLSleds;
+  YAMLSleds.reserve(Sleds.size());
+  for (const auto &Sled : Sleds) {
+    YAMLSleds.push_back({FunctionIds[Sled.Function], Sled.Address,
+                         Sled.Function, Sled.Kind, Sled.AlwaysInstrument});
+  }
+  Output Out(OS);
+  Out << YAMLSleds;
+}
+
+static CommandRegistration Unused(&Extract, [] {
+  Error Err;
+  xray::InstrumentationMapExtractor Extractor(
+      ExtractInput, InstrumentationMapExtractor::InputFormats::ELF, Err);
+  if (Err)
+    return Err;
+
+  std::error_code EC;
+  raw_fd_ostream OS(ExtractOutput, EC, sys::fs::OpenFlags::F_Text);
+  if (EC)
+    return make_error<StringError>(
+        Twine("Cannot open file '") + ExtractOutput + "' for writing.", EC);
+  Extractor.exportAsYAML(OS);
+  return Error::success();
+});

tools/llvm-xray/xray-extract.h

@@ -0,0 +1,58 @@
+//===- xray-extract.h - XRay Instrumentation Map Extraction ---------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines the interface for extracting the instrumentation map from an
+// XRay-instrumented binary.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TOOLS_XRAY_EXTRACT_H
+#define LLVM_TOOLS_XRAY_EXTRACT_H
+
+#include <deque>
+#include <map>
+#include <string>
+#include <unordered_map>
+
+#include "xray-sleds.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace llvm {
+namespace xray {
+
+class InstrumentationMapExtractor {
+public:
+  typedef std::unordered_map<int32_t, uint64_t> FunctionAddressMap;
+  typedef std::unordered_map<uint64_t, int32_t> FunctionAddressReverseMap;
+
+  enum class InputFormats { ELF, YAML };
+
+private:
+  std::deque<SledEntry> Sleds;
+  FunctionAddressMap FunctionAddresses;
+  FunctionAddressReverseMap FunctionIds;
+
+public:
+  /// Loads the instrumentation map from |Filename|. Updates |EC| in case there
+  /// were errors encountered opening the file. |Format| defines what the input
+  /// instrumentation map is in.
+  InstrumentationMapExtractor(std::string Filename, InputFormats Format,
+                              Error &EC);
+
+  const FunctionAddressMap &getFunctionAddresses() { return FunctionAddresses; }
+
+  /// Exports the loaded function address map as YAML through |OS|.
+  void exportAsYAML(raw_ostream &OS);
+};
+
+} // namespace xray
+} // namespace llvm
+
+#endif // LLVM_TOOLS_XRAY_EXTRACT_H

tools/llvm-xray/xray-registry.cc

@@ -0,0 +1,41 @@
+//===- xray-registry.cc - Implement a command registry. -------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Implement a simple subcommand registry.
+//
+//===----------------------------------------------------------------------===//
+#include "xray-registry.h"
+
+#include "llvm/Support/ManagedStatic.h"
+#include <unordered_map>
+
+namespace llvm {
+namespace xray {
+
+using HandlerType = std::function<Error()>;
+
+ManagedStatic<std::unordered_map<cl::SubCommand *, HandlerType>> Commands;
+
+CommandRegistration::CommandRegistration(cl::SubCommand *SC,
+                                         HandlerType Command) {
+  assert(Commands->count(SC) == 0 &&
+         "Attempting to overwrite a command handler");
+  assert(Command && "Attempting to register an empty std::function<Error()>");
+  (*Commands)[SC] = Command;
+}
+
+HandlerType dispatch(cl::SubCommand *SC) {
+  auto It = Commands->find(SC);
+  assert(It != Commands->end() &&
+         "Attempting to dispatch on un-registered SubCommand.");
+  return It->second;
+}
+
+} // namespace xray
+} // namespace llvm

tools/llvm-xray/xray-registry.h

@@ -0,0 +1,41 @@
+//===- xray-registry.h - Define registry mechanism for commands. ----------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Implement a simple subcommand registry.
+//
+//===----------------------------------------------------------------------===//
+#ifndef TOOLS_LLVM_XRAY_XRAY_REGISTRY_H
+#define TOOLS_LLVM_XRAY_XRAY_REGISTRY_H
+
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Error.h"
+
+namespace llvm {
+namespace xray {
+
+// Use |CommandRegistration| as a global initialiser that registers a function
+// and associates it with |SC|. This requires that a command has not been
+// registered to a given |SC|.
+//
+// Usage:
+//
+//   // At namespace scope.
+//   static CommandRegistration Unused(&MySubCommand, [] { ... });
+//
+struct CommandRegistration {
+  CommandRegistration(cl::SubCommand *SC, std::function<Error()> Command);
+};
+
+// Requires that |SC| is not null and has an associated function to it.
+std::function<Error()> dispatch(cl::SubCommand *SC);
+
+} // namespace xray
+} // namespace llvm
+
+#endif // TOOLS_LLVM_XRAY_XRAY_REGISTRY_H

tools/llvm-xray/xray-sleds.h

@@ -0,0 +1,32 @@
+//===- xray-sleds.h - XRay Sleds Data Structure ---------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines the structure used to represent XRay instrumentation map entries.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TOOLS_LLVM_XRAY_XRAY_SLEDS_H
+#define LLVM_TOOLS_LLVM_XRAY_XRAY_SLEDS_H
+
+namespace llvm {
+namespace xray {
+
+struct SledEntry {
+  enum class FunctionKinds { ENTRY, EXIT, TAIL };
+
+  uint64_t Address;
+  uint64_t Function;
+  FunctionKinds Kind;
+  bool AlwaysInstrument;
+};
+
+} // namespace xray
+} // namespace llvm
+
+#endif // LLVM_TOOLS_LLVM_XRAY_XRAY_SLEDS_H