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

[dberris]

Summary

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.

Depends on D21982 for the in-memory logging in compiler-rt.
Depends on D21983 for the changes to flags in clang.

[dberris]

Thanks @dblaikie -- I'll make some changes and ping when ready. :)

tools/llvm-xray/xray-extract.cc

52–58

Could we detect which kind of file we're reading with file magic instead of passing an argument to specify the input format?

Not with how we've defined it in compiler-rt, no. :(

Is it useful to specify the input format as YAML if the only task to be performed is to generate YAML? Perhaps the only input format should be ELF/object file.

Good point. For extraction it doesn't make sense. But the class that does extraction can be used (and actually is used in later patches) to load an instrumentation map in YAML. I'll remove the flag and assume that extraction will only ever make sense on ELF/object files.

63–68

LLVM isn't too fussy about using declarations/directives - you're welcome to "using llvm::yaml" if that seems suitable, I think? (maybe the yaml namespace has particularly vague names that are problematic?)

Right -- I was merely following the examples. I'll try with just using namespace llvm::yaml here.

(also we don't tend to bother with the '::' global scope qualifier)

Will fix this -- force of habit. :)

105

If you can get away without needing to be convertible to error_code, that's probably a good thing (it's just there for backwards compatibility in APIs where we haven't migrated the whole stack to llvm::Error yet)

Unfortunately that seems to be a pure virtual function in ErrorInfo<...>. :(

I like StringError though, I'll change this to use StringError instead (less types, better).

280

Yeah, we're constrained by the structure that's in disk. I thought about only ever using the YAML struct for everything, and it boils down to allowing it to be ported/used by other classes/functions that need to load the instrumentation map and don't need any of the YAML functionality.

286–287

This is vestigial -- it used to be that there was a function defined in the namespace, and it's changed to this global registration approach. I'll remove the namespaces. :)

289

Good idea, yeah I'll change this to make the function return an llvm::Error and consolidate the logging and exit in main.

[dberris]

• Address some review comments

[dberris]

I haven't added the tests, will do so when the style questions have been resolved.

tools/llvm-xray/xray-extract.cc

131–133

I'll add more test cases in the next revision.

[dblaikie]

tools/llvm-xray/xray-extract.cc

238–240

Previous code only printed this "Cannot extract" message, new code will print that as well as whatever text is in Err, right? Is that a desired change?

Is there some nice/easy way to append/prepend the "Cannot extract" text to the exsiting Error to pass up to main to print there? (splitting the diagnostic printing between the two places seems a bit awkward)

246–247

Ideally pass up the string and error code here, rather than printing then passing?

280

Still, think it might be nice to move away from it being splatted back and forth to memory like that - and having the expected C++ types (like enums and bools).

[dberris]

• Use errors exclusively

[dberris]

Please have another look @dblaikie -- before I add more tests?

tools/llvm-xray/xray-extract.cc

280

I agree -- on the ELF section (and the stuff embedded in the binaries), I think we're a bit constrained by what we can effectively write on the compiler-side (and read on the tooling side). On the trace files though I'm a bit more open to working with a unified encoding format, that was a bit more clever than it currently is.

[dblaikie]

tools/llvm-xray/xray-extract.cc

280

I agree -- on the ELF section (and the stuff embedded in the binaries), I think we're a bit constrained by what we can effectively write on the compiler-side (and read on the tooling side). On the trace files though I'm a bit more open to working with a unified encoding format, that was a bit more clever than it currently is.

Sorry I think I'm still confused/we're talking past each other. All I mean is that, rather than this tool splatting into a carefully laid out struct, it would use something like DataExtractor to extract the carefully laid out bytes into a more generic struct.

The file formats and writer would remain the same - just the reader would be more robust and more a usable type for the rest of the APIs.

(I suppose put another way: Using the YAML struct with its YAML types isn't really ideal for general tool code that's trying to process these records. Equally, using a carefully laid out struct with surprising types (char instead of bool/enum, explicit padding, etc) has the same problem. The main interchange structure should probably be the generic semantic thing, not YAML or binary file related - and once that happens, we could skip the packed struct entirely and just read in the few records with DataExtractor or a similar API)

[dberris]

• Use a common sled representation, hide ELF64-specific implementation detail

[dberris]

Please have another look @dblaikie.

tools/llvm-xray/xray-extract.cc

280

Ah, yes, this makes sense. I've hidden the ELF64-specific sled layout and defined a higher level SledEntry type that could be exposed later.

[dblaikie]

tools/llvm-xray/xray-extract.cc

142–145

I'm still a bit concerned about doing this via memory mapping (what if we are on a different platform that happens to add some extra padding between fields, etc?)

So I'd suggest using DataExtractor or similar (then you don't have to worry about alignment either) techniques/devices/tools, probably?

Maybe it just seems like this function is harder to read than it is in the code review, but might consider breaking it up - perhaps taking all the error handling at the beginning and turning it into one utility function, so this function (LoadBinaryInstrELF) can focus on the parsing, etc (or break out this parsing into another function, etc - or even both).

263–267

Would it make sense to sink this code into InstrumentationMapExtractor's ctor? (it has all the information - it knows it's trying to extract an instrumentation map, and the name of the input - so it seems it should be responsible for creating this message, maybe?)

[dberris]

• Use DataExtractor to load data from the memory mapped object file.

[dberris]

PTAL @dblaikie

tools/llvm-xray/xray-extract.cc

142–145

Ah, I get it now. I'm using llvm::DataExtractor now.

Refactoring this further seems pre-mature, given that this is already an implementation detail. When supporting other formats, it would be something to think about, I agree.

263–267

Actually, nope -- because users of this class could choose to ignore the errors (i.e. treat it as if there was no available instrumentation map). It's just that for this sub-command, it wouldn't work if extraction actually failed. :)

[dblaikie]

Generally looking pretty good - a few little nits & more test coverage (error cases, for example) is probably a good idea.

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

6–7

Convenient how?
I suppose what i'm trying to understand is: why /both/? (which is to say, what does each offer that the other doesn't)

13

Is it worth checking the specific addresses? If this test is simple enough to provide fairly stable values here - to make sure the algorithms, etc, are working correctly & not just printing garbage values?

tools/llvm-xray/xray-extract.cc

86

I'd sink this down into the function/near the use (the general principle of keeping variables in the smallest scope needed)

164

Drop the "? false : true" here, and use !=, perhaps:

"Entry.AlwaysInstrument = AlwaysInstrument != 0" (potentially even drop the "!= 0", but I can see how that helps readability.

182

Move assignment rather than swap? (but equally I wouldn't mind the old code that left the OutputSleds in an unspecified state on failure)

234–238

This code (& LoadYAMLInstrFile) is dead/untested - perhaps should be moved into another/separate change.

263–267

I'm not sure I understand this - could you rephrase?

What I mean is: The joining of this "Cannot extract instrumentation map from" seems like it could go inside the InstrumentationMapExtractor ctor (since it has all the context to make that message and the message seems appropriate at that level). Then at this level (in the CommandRegistration) we just propagate the error up ("if (Err) return Err; or whatever)

tools/llvm-xray/xray-extract.h

39

Doesn't look like this needs to be a member? (it's only used during construction, if I'm reading it correctly)

51

Unused/dead code

tools/llvm-xray/xray-registry.h

32–33

Could just make this a struct - since it's only member's public anyway.

37

It doesn't seem like it actually requires that SC is not null (& pedantically the terminology would be "not null" (or "not a null pointer" - 'nullptr' is just some specific null pointer literal)) - but, sure - seems OK to say that's not acceptable even if it'd be fine for the current implementation.

Is it worth just defining this function to only work if the SubCommand is registered? (assert when not found, instead of returning the empty std::function)

[dberris]

• Add failure tests
• Address more review comments

[dberris]

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

6–7

JSON was convenient for making it easily parse-able and view-able by anything that will show it on a browser. :)

13

Yep, added a few more functions to be a bit more complete.

tools/llvm-xray/xray-extract.cc

263–267

Ah, right -- I've moved the concatenation of the Twine into the error generation parts.

tools/llvm-xray/xray-extract.h

51

Unused/dead code

Yes, but this is actually used in a later (stacked) change...

[dberris]

This is ready for another look @dblaikie -- added some failure tests that were easy to catch (not found, no instrumentation map) and provided some sample binaries to use instead of creating one on the fly (as it seemed harder to generate a final linked executable from a .ll).

If you have ideas about how to make the sequence:

llc ...; <link the .o into a final binary> ;  llvm-xray extract <final binary>

work in a .ll test, I'd be very interested to try.

[dblaikie]

In D21987#576231, @dberris wrote:

This is ready for another look @dblaikie -- added some failure tests that were easy to catch (not found, no instrumentation map) and provided some sample binaries to use instead of creating one on the fly (as it seemed harder to generate a final linked executable from a .ll).

If you have ideas about how to make the sequence:

llc ...; <link the .o into a final binary> ;  llvm-xray extract <final binary>

work in a .ll test, I'd be very interested to try.

Could you explain why you need a linked binary rather than just an object file? What changes? Relocations, I guess? All the addresses look like zero?

Yeah, doubt there's much to do about that - no way to produce linked binaries with the LLVM tools available to the test suite, and no point implementing support in the tool to handle unapplied relocations (unlike dwarfdump which does handle them, and there's more reason to run dwarfdump on unlinked objects, so it's worth having that support for more than just testing).

tools/llvm-xray/xray-extract.cc

89–92

Untested (do we need this test? Presumably we just won't be able to open the file)

97–98

I'm guessing this is the path we take for "no such file or directory"? (ie: covered by the test for that)

100–103

Untested

118–121

Untested (though, granted - not sure quite how to test this, but could look further into libObject to see how getContents can fail)

127–131

Untested

206–208

Guess this should just be an assert for now? (represents a programmer error/isn't reachable/testable/etc?)

tools/llvm-xray/xray-registry.cc

34–35

Replace branch-to-unreachable with assert.

[dberris]

In D21987#576626, @dblaikie wrote:

In D21987#576231, @dberris wrote:

This is ready for another look @dblaikie -- added some failure tests that were easy to catch (not found, no instrumentation map) and provided some sample binaries to use instead of creating one on the fly (as it seemed harder to generate a final linked executable from a .ll).

If you have ideas about how to make the sequence:

llc ...; <link the .o into a final binary> ;  llvm-xray extract <final binary>

work in a .ll test, I'd be very interested to try.

Could you explain why you need a linked binary rather than just an object file? What changes? Relocations, I guess? All the addresses look like zero?

The addresses are basically zero on a .o :/

Yeah, doubt there's much to do about that - no way to produce linked binaries with the LLVM tools available to the test suite, and no point implementing support in the tool to handle unapplied relocations (unlike dwarfdump which does handle them, and there's more reason to run dwarfdump on unlinked objects, so it's worth having that support for more than just testing).

Yep, I think the binary should be fine here.

However, I'm having some trouble generating malformed binaries without resorting to hex-editing some existing files. Any recommendations for faking "bad" inputs to test the error conditions?

tools/llvm-xray/xray-extract.cc

89–92

Good point, removed.

97–98

Yes. that's correct.

[dberris]

• Add test for empty, convert some if-unreachable to assert

[dberris]

• Add elf32 sample and test for unsupportedness
• Add test for badly sized instrumentation map entries.

[dberris]

Added a couple more tests, ptal -- I haven't quite resorted to hex-editing files, but close enough with deliberately generating a bad XRay instrumentation map with a modified clang+llvm local build.

tools/llvm-xray/xray-extract.cc

118–121

I looked, and it had something to do with the ELF encoding of a file (say, if the supposed size of the section defined in the header is different in reality (through some checks)). Not sure how to properly test this yet.

[dblaikie]

tools/llvm-xray/llvm-xray.cc

36

Perhaps it'd be better to catch this case earlier - disallow registering with a null/empty function (just assert) rather than quietly not executing anything?

[dberris]

Squash local commits and rebase before landing.

[dberris]

Thanks for the review @dblaikie!

tools/llvm-xray/llvm-xray.cc

36

Good idea. Added an assert on the registration path to make sure the function is not "empty".

[dberris]

Update description.

[dberris]

Landed in rL285155.

[dberris]

Rolled back in rL285155, broke the tests.

[dberris]

Landed as rL285165.

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

+; 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

+config.suffixes = ['.yaml', '.ll', '.txt']

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

+; 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

+; 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

tools/llvm-xray/CMakeLists.txt

+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

+//===- 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.h

+//===- 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-extract.cc

+//===- 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) {
+  if (Filename.empty())
+    return make_error<StringError>(
+        "Provided filename is empty.",
+        std::make_error_code(std::errc::no_such_file_or_directory));
+
+  auto ObjectFile = object::ObjectFile::createObjectFile(Filename);
+
+  // FIXME: Maybe support other ELF formats. For now, 64-bit Little Endian only.
+  if (!ObjectFile)
+    return ObjectFile.takeError();
+
+  if (!ObjectFile->getBinary()->isELF())
+    return make_error<StringError>(
+        "File format not supported (only does ELF).",
+        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:
+    EC =
+        make_error<StringError>("Input format type not supported yet.",
+                                std::make_error_code(std::errc::not_supported));
+    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-registry.h

+//===- 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-registry.cc

+//===- 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");
+  (*Commands)[SC] = Command;
+}
+
+HandlerType dispatch(cl::SubCommand *SC) {
+  auto It = Commands->find(SC);
+  if (It == Commands->end())
+    llvm_unreachable("Attempting to dispatch on un-registered SubCommand.");
+  return It->second;
+}
+
+} // namespace xray
+} // namespace llvm

tools/llvm-xray/xray-sleds.h

+//===- 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