XRay: Add entry and exit sleds

[dberris]

Summary

In this patch we implement the following parts of XRay:

• Supporting a function attribute named 'function-instrument' which currently only supports 'xray-always'. We should be able to use this attribute for other instrumentation approaches.
• Supporting a function attribute named 'xray-instruction-threshold' used to determine whether a function is instrumented with a minimum number of instructions (IR instruction counts).
• X86-specific nop sleds as described in the white paper.
• A machine function pass that adds the different instrumentation marker instructions at a very late stage.
• A way of identifying which return opcode is considered "normal" for each architecture.

There are some caveats here:

1. We don't handle PATCHABLE_RET in platforms other than x86_64 yet -- this means if IR used PATCHABLE_RET directly instead of a normal ret, instruction lowering for that platform might do the wrong thing. We think this should be handled at instruction selection time to by default be unpacked for platforms where XRay is not availble yet.

2. The generated section for X86 is different from what is described from the white paper for the sole reason that LLVM allows us to do this neatly. We're taking the opportunity to deviate from the white paper from this perspective to allow us to get richer information from the runtime library.

[dberris]

BTW, It's worth mentioning that this is a prototype implementation and we're fine with making large changes for this patch.

[sanjoy]

Some minor comments inline. I'm not familiar with all of the intricacies of the backend, so I'll CC some other people who are likely to have more perspective on this.

include/llvm/Target/Target.td

957

Looks like this isn't used yet? If so, I'd suggest dropping it completely for now.

lib/CodeGen/XRayInstrumentation.cpp

80

Why do you care only about register operands here?

lib/Target/X86/X86MCInstLower.cpp

1068

Why not

AlwaysInstrument = !Attr.hasAttribute(Attribute::None) && Attr.isStringAttribute() &&
      Attr.getValueAsString() == "xray-always";

?

1069

I don't quite grok the Attributes API, but I think the !Attr.hasAttribute(Attribute::None) bit is redundant -- if isStringAttribute return true, then Attr.hasAttribute(Attribute::None) will be false.

1072

emplace_back ?

1100

Indent is off.

[majnemer]

lib/CodeGen/XRayInstrumentation.cpp

72–77

I don't think you want to replace whacky constructs like EH_RETURN and CLEANUPRET with PATCHABLE_RET, do you? Just "normal" returns, right?

[dberris]

Address comments from sanjoy and majnemer

lib/CodeGen/XRayInstrumentation.cpp

72–77

Yes, that's right -- does this version exclude those appropriately?

80

Oops, I had assumed that the RET instructions only ever had register operands. That's a faulty assumption, good catch.

[majnemer]

lib/CodeGen/XRayInstrumentation.cpp

72–77

I would recommend white-listing the opcodes you know are safe to replace instead of blacklisting the ones you know which are unsafe.
This would make the pass conservatively correct in the face of future changes to LLVM.

[dberris]

lib/CodeGen/XRayInstrumentation.cpp

72–77

That makes sense, thanks majnemer -- however it's not clear to me how I'd refer to the target-specific opcodes from here. Pardon the newbie question, but how do I say "I just want return instructions here"? The only options I can see are:

• Look at the original LLVM IR to see if it's a ReturnInst.
• Figure out a way for including the X86 opcodes here.

Maybe I'm missing something simpler here?

[dberris]

Update operand unpacking properly.

[dberris]

• clang-format

[majnemer]

lib/CodeGen/XRayInstrumentation.cpp

72–76

The typical way of doing this sort of thing, AFAICT, is to teach TargetInstrInfo about it.

For example, it implements getCatchReturnOpcode and getCallFrameSetupOpcode.

[dberris]

Add a check in TargetInstrInfo for whether a return is a "normal" return.

We also implement this somewhat correctly in X86, but use the default
implementation for other platforms.

[dberris]

lib/CodeGen/XRayInstrumentation.cpp

72–74

Thanks!

I've updated the patch to do this. PTAL?

[dberris]

Fix the uploaded diff.

[majnemer]

include/llvm/Target/TargetInstrInfo.h

155–166 ↗

(On Diff #56531)

I was thinking more of a whilelist-oriented solution, something like:

`unsigned getNormalReturnOpcode() const { return NormalRetOpcode; }`

[bmakam]

lib/Target/X86/X86MCInstLower.cpp

1082

Could you please expand on why you need 9 bytes of noops here? I am not quite familiar with x86_64 but was under the impression that on x86_64 the jmp instruction is 1 byte for opcode and 4 bytes for signed relative displacement, so shouldn't 5 bytes worth of nops be sufficient?

[dberris]

include/llvm/Target/TargetInstrInfo.h

155–166 ↗

(On Diff #56531)

But how does this work on platforms that can spell 'return' many different normal ways (like in X86)? There's RETL, RETQ, and all other versions of RET which are considered "normal"?

lib/Target/X86/X86MCInstLower.cpp

1082

Good question, thanks.

I have to check whether we're using the right version of JMP, but I'm specifically looking for the version that's one byte for the JMP instruction, and 8 bits (1 byte) for the relative offset. So far I haven't been able to spell jmp +0x09 correctly and have it work, without having an additional symbol as a target for the jump instruction. If we get that right, we can then add the 9 byte nops we need to get exactly 11 bytes for the function entry.

Is there a fool-proof way of spelling "JMP +0x09" with the builder interface? Or should I add another JMP instruction in X86 that supports the 8-bit displacement immediate operand?

The reason why I can't use a JMP that isn't 2 bytes, is because it's really hard to write just 5 bytes atomically. I can probably do something with an 8-byte atomic write, but enforcing that 8-byte write doesn't span cache lines is also very tricky to make safe.

[dberris]

lib/Target/X86/X86MCInstLower.cpp

1082

Actually now that I've had a look at the generated object file, I can confirm that we're using the two-byte version of JMP with this construct. Quick test:

test.cc:

#include <cstdio>

[[clang::xray_always_instrument]] void foo() { std::printf("Hello, XRay!\n"); }

int main(int argc, char* argv[]) { foo(); }

Compiled with (modified clang to emit IR that has annotated functions for XRay):

./bin/clang -fxray-instrument -fxray-instruction-threshold=1 -std=c++11 -x c++ -S test.cc -emit-llvm

Creates this IR listing:

; ModuleID = 'test.cc'
source_filename = "test.cc"
target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
target triple = "x86_64-unknown-linux-gnu"

@.str = private unnamed_addr constant [14 x i8] c"Hello, XRay!\0A\00", align 1

; Function Attrs: uwtable
define void @_Z3foov() #0 {
entry:
  %call = call i32 (i8*, ...) @printf(i8* getelementptr inbounds ([14 x i8], [14 x i8]* @.str, i32 0, i32 0))
  ret void
}

declare i32 @printf(i8*, ...) #1

; Function Attrs: norecurse uwtable
define i32 @main(i32 %argc, i8** %argv) #2 {
entry:
  %argc.addr = alloca i32, align 4
  %argv.addr = alloca i8**, align 8
  store i32 %argc, i32* %argc.addr, align 4
  store i8** %argv, i8*** %argv.addr, align 8
  call void @_Z3foov()
  ret i32 0
}

attributes #0 = { uwtable "disable-tail-calls"="false" "function-instrument"="xray-always" "less-precise-fpmad"="false" "no-frame-pointer-elim"="true" "no-frame-pointer-elim-non-leaf" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+fxsr,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" }
attributes #1 = { "disable-tail-calls"="false" "less-precise-fpmad"="false" "no-frame-pointer-elim"="true" "no-frame-pointer-elim-non-leaf" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+fxsr,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" }
attributes #2 = { norecurse uwtable "disable-tail-calls"="false" "less-precise-fpmad"="false" "no-frame-pointer-elim"="true" "no-frame-pointer-elim-non-leaf" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+fxsr,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" "xray-instruction-threshold"="1" }

!llvm.ident = !{!0}

!0 = !{!"clang version 3.9.0 (http://llvm.org/git/clang.git 2b9ed9227330789f24e56fe9e800e7be2111073b) (http://llvm.org/git/llvm.git 1f3e353113ff9e2f835955bf005a8a5e25f16ad1)"}

Then compiled+disassembled this way:

./bin/llc -filetype=obj -o - < test.ll | ./bin/llvm-objdump -disassemble -

Produces the following output:

<stdin>:        file format ELF64-x86-64

Disassembly of section .text:
_Z3foov:
       0:       eb 09   jmp     9 <_Z3foov+0xB>
       2:       66 0f 1f 84 00 00 02 00 00      nopw    512(%rax,%rax)
       b:       55      pushq   %rbp
       c:       48 89 e5        movq    %rsp, %rbp
       f:       bf 00 00 00 00  movl    $0, %edi
      14:       31 c0   xorl    %eax, %eax
      16:       e8 00 00 00 00  callq   0 <_Z3foov+0x1B>
      1b:       5d      popq    %rbp
      1c:       c3      retq
      1d:       2e 66 0f 1f 84 00 00 02 00 00   nopw    %cs:512(%rax,%rax)
      27:       66 0f 1f 84 00 00 00 00 00      nopw    (%rax,%rax)

main:
      30:       eb 09   jmp     9 <main+0xB>
      32:       66 0f 1f 84 00 00 02 00 00      nopw    512(%rax,%rax)
      3b:       55      pushq   %rbp
      3c:       48 89 e5        movq    %rsp, %rbp
      3f:       48 83 ec 10     subq    $16, %rsp
      43:       89 7d fc        movl    %edi, -4(%rbp)
      46:       48 89 75 f0     movq    %rsi, -16(%rbp)
      4a:       e8 b1 ff ff ff  callq   -79 <_Z3foov>
      4f:       31 c0   xorl    %eax, %eax
      51:       48 83 c4 10     addq    $16, %rsp
      55:       5d      popq    %rbp
      56:       c3      retq
      57:       2e 66 0f 1f 84 00 00 02 00 00   nopw    %cs:512(%rax,%rax)

[dberris]

• Add documentation on why we need 9 bytes of nops

[dberris]

lib/Target/X86/X86MCInstLower.cpp

1082

Updated the comments now to make it clear and specific.

Thanks!

[majnemer]

include/llvm/Target/TargetInstrInfo.h

155–166 ↗

(On Diff #56850)

The same way it is handled for X86::ADJCALLSTACKDOWN64 vs X86::ADJCALLSTACKDOWN32:
https://github.com/llvm-mirror/llvm/blob/master/lib/Target/X86/X86InstrInfo.cpp#L105

[dberris]

include/llvm/Target/TargetInstrInfo.h

155–166 ↗

(On Diff #56850)

That makes a lot of sense, thanks! I'll rework this a bit to make it more of a targeted specific OpCode.

I suppose it's fine to have multiples of these functions, each for relevant flavours of fairly common instruction classes, yes? So I suppose, for tail call exits/returns, we would have something similar for that?

[dberris]

• Make ReturnOpcode an argument to TargetInstrInfo constructor and whitelist one ReturnOpcode instead

[dberris]

include/llvm/Target/TargetInstrInfo.h

157–168 ↗

(On Diff #56851)

Updated now to white-list a specific return OpCode.

[dberris]

Ping -- is there something else missing for this change?

[dberris]

Marked more inline comments as "Done".

[echristo]

Drive by comment.

include/llvm/Target/TargetOpcodes.def

145

Remove the XRay specific stuff here and rewrite in a more generic fashion (and below).

[dberris]

• Update documentation to remove XRay references

[dberris]

include/llvm/Target/TargetOpcodes.def

145

Done

[dberris]

Ping?

[dberris]

Sync to master again.

[dberris]

• Fixes after merge
• Emit a common symbol for the instrumentation map
• Improve table creation and allow access to the table via xray_instr_map and xray_instr_map_end
• Fix the sled type encoding

[dberris]

• Use an explicit byte sequence for the jmp

[dberris]

• Make XRay write out tables per function

[dberris]

• Clean up local history, squashing to a single revision (rebased too)

[dberris]

• Rebase
• Add ELF::SHF_MERGE to flags for XRay section
• Update tests to reflect adjustments in implementation

[dberris]

• Un-break XRay section creation.

[rnk]

lgtm

lib/Target/X86/X86MCInstLower.cpp

1089

Does XRay support multiple DSOs in the process? If so, the XRay runtime may be more than 2GB away from the code being patched, and this offset will overflow.

1098–1099

I guess we can deal with this separately. IMO we should have a 'jmpb' instruction or something that forces a short jump or assembler error if the displacement is too large.

[dberris]

Thanks Reid -- do you or echristo@ mind landing this for me? I don't think I have commit powers yet.

lib/Target/X86/X86MCInstLower.cpp

1089

That's true -- the runtime patching code will fail to patch certain sleds that are farther than 32-bits away. In case we need more, in the future we can tweak this so that we can try to force the XRay runtime trampoline in the first 2GB of the process' memory and use an "absolute" address (which may need more bytes in the sled), or have a mode to support larger binaries (a flag that changes the emitted sleds, and have the runtime support both kinds of sleds).

I suspect if it comes up in practice that we can deal with this more effectively.

1098–1099

I'm happy to do this as a refactoring or a follow-up change. A new instruction sounds like the right action to take anyway.

[dberris]

Rebase

include/llvm/CodeGen/Passes.h

/// MachineDominanaceFrontier - This pass is a machine dominators analysis pass.
   extern char &OptimizePHIsID;
 
   /// StackSlotColoring - This pass performs stack slot coloring.
   extern char &StackSlotColoringID;
 
   /// \brief This pass lays out funclets contiguously.
   extern char &FuncletLayoutID;
 
+  /// This pass inserts the XRay instrumentation sleds if they are supported by
+  /// the target platform.
+  extern char &XRayInstrumentationID;
+
   /// \brief This pass implements the "patchable-function" attribute.
   extern char &PatchableFunctionID;
 
   /// createStackProtectorPass - This pass adds stack protectors to functions.
   ///
   FunctionPass *createStackProtectorPass(const TargetMachine *TM);
 
   /// createMachineVerifierPass - This pass verifies cenerated machine code

include/llvm/InitializePasses.h

 void initializeSjLjEHPreparePass(PassRegistry&);
 void initializeDemandedBitsWrapperPassPass(PassRegistry&);
 void initializeFuncletLayoutPass(PassRegistry &);
 void initializeLoopLoadEliminationPass(PassRegistry&);
 void initializeFunctionImportPassPass(PassRegistry &);
 void initializeLoopVersioningPassPass(PassRegistry &);
 void initializeWholeProgramDevirtPass(PassRegistry &);
 void initializePatchableFunctionPass(PassRegistry &);
+void initializeXRayInstrumentationPass(PassRegistry &);
 }
 
 #endif

include/llvm/MC/MCObjectFileInfo.h

protected:
   MCSection *ThreadLocalPointerSection;
 
   /// COFF specific sections.
   MCSection *DrectveSection;
   MCSection *PDataSection;
   MCSection *XDataSection;
   MCSection *SXDataSection;
 
+  /// XRay specific section.
+  MCSection *XRaySection;
+
 public:
   void InitMCObjectFileInfo(const Triple &TT, Reloc::Model RM,
                             CodeModel::Model CM, MCContext &ctx);
   LLVM_ATTRIBUTE_DEPRECATED(
       void InitMCObjectFileInfo(StringRef TT, Reloc::Model RM,
                                 CodeModel::Model CM, MCContext &ctx),
       "StringRef GNU Triple argument replaced by a llvm::Triple object");
 
public:
   }
 
   // COFF specific sections.
   MCSection *getDrectveSection() const { return DrectveSection; }
   MCSection *getPDataSection() const { return PDataSection; }
   MCSection *getXDataSection() const { return XDataSection; }
   MCSection *getSXDataSection() const { return SXDataSection; }
 
+  // XRay specific sections.
+  MCSection *getXRaySection() const { return XRaySection; }
+
   MCSection *getEHFrameSection() {
     return EHFrameSection;
   }
 
   enum Environment { IsMachO, IsELF, IsCOFF };
   Environment getObjectFileType() const { return Env; }
 
   Reloc::Model getRelocM() const { return RelocM; }

include/llvm/Target/Target.td

 def PATCHABLE_OP : Instruction {
   let OutOperandList = (outs unknown:$dst);
   let InOperandList = (ins variable_ops);
   let usesCustomInserter = 1;
   let mayLoad = 1;
   let mayStore = 1;
   let hasSideEffects = 1;
 }
+def PATCHABLE_FUNCTION_ENTER : Instruction {
+  let OutOperandList = (outs);
+  let InOperandList = (ins);
+  let AsmString = "# XRay Function Enter.";
+  let usesCustomInserter = 1;
+  let hasSideEffects = 0;
+}
+def PATCHABLE_RET : Instruction {
+  let OutOperandList = (outs unknown:$dst);
+  let InOperandList = (ins variable_ops);
+  let AsmString = "# XRay Function Exit.";
+  let usesCustomInserter = 1;
+  let hasSideEffects = 1;
+  let isReturn = 1;
+}
+def PATCHABLE_TAIL_CALL : Instruction {
+  let OutOperandList = (outs);
+  let InOperandList = (ins);
+  let AsmString = "# XRay Tail Call Exit.";
+  let usesCustomInserter = 1;
+  let hasSideEffects = 0;
+  let isTerminator = 1;
+  let isReturn = 1;
+}
 
 // Generic opcodes used in GlobalISel.
 include "llvm/Target/GenericOpcodes.td"
 
 }
 
 //===----------------------------------------------------------------------===//
 // AsmParser - This class can be implemented by targets that wish to implement

include/llvm/Target/TargetOpcodes.def

 /// instruction wrapped in PATCHABLE_OP has to either have a minimum
 /// size or be preceded with a nop of that size.  The first operand is
 /// an immediate denoting the minimum size of the instruction, the
 /// second operand is an immediate denoting the opcode of the original
 /// instruction.  The rest of the operands are the operands of the
 /// original instruction.
 HANDLE_TARGET_OPCODE(PATCHABLE_OP, 23)
 
+/// PATCHABLE_FUNCTION_ENTER - a marker for XRay function entry blocks.
+HANDLE_TARGET_OPCODE(PATCHABLE_FUNCTION_ENTER, 24)
+
+/// PATCHABLE_RET - a marker for XRay function exit blocks.
+HANDLE_TARGET_OPCODE(PATCHABLE_RET, 25)
+
+/// PATCHABLE_TAIL_CALL - a marker for XRay tail call exits.
+HANDLE_TARGET_OPCODE(PATCHABLE_TAIL_CALL, 26)
+
 /// The following generic opcodes are not supposed to appear after ISel.
 /// This is something we might want to relax, but for now, this is convenient
 /// to produce diagnostics.
 
 /// Generic ADD instruction. This is an integer add.
-HANDLE_TARGET_OPCODE(G_ADD, 24)
+HANDLE_TARGET_OPCODE(G_ADD, 27)
 HANDLE_TARGET_OPCODE_MARKER(PRE_ISEL_GENERIC_OPCODE_START, G_ADD)
 
 /// Generic BRANCH instruction. This is an unconditional branch.
-HANDLE_TARGET_OPCODE(G_BR, 25)
+HANDLE_TARGET_OPCODE(G_BR, 28)
 
 // TODO: Add more generic opcodes as we move along.
 
 /// Marker for the end of the generic opcode.
 /// This is used to check if an opcode is in the range of the
 /// generic opcodes.
 HANDLE_TARGET_OPCODE_MARKER(PRE_ISEL_GENERIC_OPCODE_END, G_BR)
 
 /// BUILTIN_OP_END - This must be the last enum value in this list.
 /// The target-specific post-isel opcode values start here.
 HANDLE_TARGET_OPCODE_MARKER(GENERIC_OP_END, PRE_ISEL_GENERIC_OPCODE_END)

lib/CodeGen/CMakeLists.txt

add_llvm_library(LLVMCodeGen
   TargetLoweringObjectFileImpl.cpp
   TargetOptionsImpl.cpp
   TargetRegisterInfo.cpp
   TargetSchedule.cpp
   TwoAddressInstructionPass.cpp
   UnreachableBlockElim.cpp
   VirtRegMap.cpp
   WinEHPrepare.cpp
+  XRayInstrumentation.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/CodeGen
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/CodeGen/PBQP
 
   LINK_LIBS ${PTHREAD_LIB}
   )
 
 add_dependencies(LLVMCodeGen intrinsics_gen)
 
 add_subdirectory(SelectionDAG)
 add_subdirectory(AsmPrinter)
 add_subdirectory(MIRParser)
 add_subdirectory(GlobalISel)

lib/CodeGen/CodeGen.cpp

void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeMachineFunctionPrinterPassPass(Registry);
   initializeMachineLICMPass(Registry);
   initializeMachineLoopInfoPass(Registry);
   initializeMachineModuleInfoPass(Registry);
   initializeMachinePostDominatorTreePass(Registry);
   initializeMachineSchedulerPass(Registry);
   initializeMachineSinkingPass(Registry);
   initializeMachineVerifierPassPass(Registry);
+  initializeXRayInstrumentationPass(Registry);
   initializePatchableFunctionPass(Registry);
   initializeOptimizePHIsPass(Registry);
   initializePEIPass(Registry);
   initializePHIEliminationPass(Registry);
   initializePeepholeOptimizerPass(Registry);
   initializePostMachineSchedulerPass(Registry);
   initializePostRAHazardRecognizerPass(Registry);
   initializePostRASchedulerPass(Registry);

lib/CodeGen/Passes.cpp

void TargetPassConfig::addMachinePasses() {
 
   addPreEmitPass();
 
   addPass(&FuncletLayoutID, false);
 
   addPass(&StackMapLivenessID, false);
   addPass(&LiveDebugValuesID, false);
 
+  addPass(&XRayInstrumentationID, false);
   addPass(&PatchableFunctionID, false);
 
   AddingMachinePasses = false;
 }
 
 /// Add passes that optimize machine instructions in SSA form.
 void TargetPassConfig::addMachineSSAOptimization() {
   // Pre-ra tail duplication.

lib/CodeGen/XRayInstrumentation.cpp

+//===-- XRayInstrumentation.cpp - Adds XRay instrumentation to functions. -===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a MachineFunctionPass that inserts the appropriate
+// XRay instrumentation instructions. We look for XRay-specific attributes
+// on the function to determine whether we should insert the replacement
+// operations.
+//
+//===---------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/Analysis.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/Target/TargetInstrInfo.h"
+#include "llvm/Target/TargetSubtargetInfo.h"
+
+using namespace llvm;
+
+namespace {
+struct XRayInstrumentation : public MachineFunctionPass {
+  static char ID;
+
+  XRayInstrumentation() : MachineFunctionPass(ID) {
+    initializeXRayInstrumentationPass(*PassRegistry::getPassRegistry());
+  }
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+};
+}
+
+bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
+  auto &F = *MF.getFunction();
+  auto InstrAttr = F.getFnAttribute("function-instrument");
+  bool AlwaysInstrument = !InstrAttr.hasAttribute(Attribute::None) &&
+                          InstrAttr.isStringAttribute() &&
+                          InstrAttr.getValueAsString() == "xray-always";
+  Attribute Attr = F.getFnAttribute("xray-instruction-threshold");
+  unsigned XRayThreshold = 0;
+  if (!AlwaysInstrument) {
+    if (Attr.hasAttribute(Attribute::None) || !Attr.isStringAttribute())
+      return false; // XRay threshold attribute not found.
+    if (Attr.getValueAsString().getAsInteger(10, XRayThreshold))
+      return false; // Invalid value for threshold.
+    if (F.size() < XRayThreshold)
+      return false; // Function is too small.
+  }
+
+  // FIXME: Do the loop triviality analysis here or in an earlier pass.
+
+  // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
+  // MachineFunction.
+  auto &FirstMBB = *MF.begin();
+  auto &FirstMI = *FirstMBB.begin();
+  auto *TII = MF.getSubtarget().getInstrInfo();
+  BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
+          TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
+
+  // Then we look for *all* terminators and returns, then replace those with
+  // PATCHABLE_RET instructions.
+  SmallVector<MachineInstr *, 4> Terminators;
+  for (auto &MBB : MF) {
+    for (auto &Terminator : MBB.terminators()) {
+      // FIXME: Handle tail calls here too?
+      if (Terminator.isReturn()) {
+        // Replace return instructions with:
+        //   PATCHABLE_RET <Opcode>, <Operand>...
+        auto MIB = BuildMI(MBB, Terminator, Terminator.getDebugLoc(),
+                           TII->get(TargetOpcode::PATCHABLE_RET))
+                       .addImm(Terminator.getOpcode());
+        for (auto &MO : Terminator.operands()) {
+          if (MO.isReg())
+            MIB.addReg(MO.getReg(), RegState::Implicit);
+        }
+        Terminators.push_back(&Terminator);
+      }
+    }
+  }
+
+  for (auto &I : Terminators)
+    I->eraseFromParent();
+
+  return true;
+}
+
+char XRayInstrumentation::ID = 0;
+char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
+INITIALIZE_PASS(XRayInstrumentation, "xray-instrumentation", "Insert XRay ops",
+                false, false);

lib/MC/MCObjectFileInfo.cpp

DwarfTUIndexSection =
                            SectionKind::getMetadata());
   StackMapSection = Ctx->getMachOSection("__LLVM_STACKMAPS", "__llvm_stackmaps",
                                          0, SectionKind::getMetadata());
 
   FaultMapSection = Ctx->getMachOSection("__LLVM_FAULTMAPS", "__llvm_faultmaps",
                                          0, SectionKind::getMetadata());
 
   TLSExtraDataSection = TLSTLVSection;
+
+  XRaySection = Ctx->getMachOSection("__XRAY_INSTR_MAP", "__xray_instr_map", 0,
+                                     SectionKind::getMetadata());
 }
 
 void MCObjectFileInfo::initELFMCObjectFileInfo(Triple T) {
   switch (T.getArch()) {
   case Triple::mips:
   case Triple::mipsel:
     FDECFIEncoding = dwarf::DW_EH_PE_sdata4;
     break;
void MCObjectFileInfo::initELFMCObjectFileInfo(Triple T) {
   StackMapSection =
       Ctx->getELFSection(".llvm_stackmaps", ELF::SHT_PROGBITS, ELF::SHF_ALLOC);
 
   FaultMapSection =
       Ctx->getELFSection(".llvm_faultmaps", ELF::SHT_PROGBITS, ELF::SHF_ALLOC);
 
   EHFrameSection =
       Ctx->getELFSection(".eh_frame", EHSectionType, EHSectionFlags);
+
+  XRaySection =
+      Ctx->getELFSection(".xray_instr_map", ELF::SHT_PROGBITS, ELF::SHF_ALLOC);
 }
 
 void MCObjectFileInfo::initCOFFMCObjectFileInfo(Triple T) {
   EHFrameSection = Ctx->getCOFFSection(
       ".eh_frame", COFF::IMAGE_SCN_CNT_INITIALIZED_DATA |
                        COFF::IMAGE_SCN_MEM_READ | COFF::IMAGE_SCN_MEM_WRITE,
       SectionKind::getData());
 

lib/Target/X86/X86AsmPrinter.h

private:
     // CurrentShadowSize counts the number of bytes encoded since the most
     // recently encountered STACKMAP, stopping when that number is greater than
     // or equal to RequiredShadowSize.
     unsigned RequiredShadowSize = 0, CurrentShadowSize = 0;
   };
 
   StackMapShadowTracker SMShadowTracker;
 
+  // This describes the kind of sled we're storing in the XRay table.
+  enum class SledKind : uint8_t {
+    FUNCTION_ENTER,
+    FUNCTION_EXIT,
+    TAIL_CALL,
+  };
+
+  // The table will contain these structs that point to the sled, the function
+  // containing the sled, and what kind of sled (and whether they should always
+  // be instrumented).
+  struct XRayFunctionEntry {
+    const MCSymbol *Sled;
+    const MCSymbol *Function;
+    SledKind Kind;
+    bool AlwaysInstrument;
+  };
+
+  // All the sleds to be emitted.
+  std::vector<XRayFunctionEntry> Sleds;
+
   // All instructions emitted by the X86AsmPrinter should use this helper
   // method.
   //
   // This helper function invokes the SMShadowTracker on each instruction before
   // outputting it to the OutStream. This allows the shadow tracker to minimise
   // the number of NOPs used for stackmap padding.
   void EmitAndCountInstruction(MCInst &Inst);
   void LowerSTACKMAP(const MachineInstr &MI);
   void LowerPATCHPOINT(const MachineInstr &MI, X86MCInstLower &MCIL);
   void LowerSTATEPOINT(const MachineInstr &MI, X86MCInstLower &MCIL);
   void LowerFAULTING_LOAD_OP(const MachineInstr &MI, X86MCInstLower &MCIL);
   void LowerPATCHABLE_OP(const MachineInstr &MI, X86MCInstLower &MCIL);
 
   void LowerTlsAddr(X86MCInstLower &MCInstLowering, const MachineInstr &MI);
 
- public:
-   explicit X86AsmPrinter(TargetMachine &TM,
-                          std::unique_ptr<MCStreamer> Streamer)
-       : AsmPrinter(TM, std::move(Streamer)), SM(*this), FM(*this) {}
+  // XRay-specific lowering for X86.
+  void LowerPATCHABLE_FUNCTION_ENTER(const MachineInstr &MI,
+                                     X86MCInstLower &MCIL);
+  void LowerPATCHABLE_RET(const MachineInstr &MI, X86MCInstLower &MCIL);
+  void LowerPATCHABLE_TAIL_CALL(const MachineInstr &MI, X86MCInstLower &MCIL);
+
+  // Helper function that emits all XRay sleds we've collected.
+  void EmitXRayTable();
+
+  // Helper function to record a given XRay sled.
+  void recordSled(MCSymbol *Sled, const MachineInstr &MI);
 
   const char *getPassName() const override {
     return "X86 Assembly / Object Emitter";
   }
 
   const X86Subtarget &getSubtarget() const { return *Subtarget; }
 
   void EmitStartOfAsmFile(Module &M) override;

lib/Target/X86/X86AsmPrinter.cpp

if (!Stubs.empty()) {
         emitNonLazySymbolPointer(*OutStreamer, Stub.first, Stub.second);
 
       Stubs.clear();
       OutStreamer->AddBlankLine();
     }
 
     SM.serializeToStackMapSection();
     FM.serializeToFaultMapSection();
+    EmitXRayTable();
 
     // Funny Darwin hack: This flag tells the linker that no global symbols
     // contain code that falls through to other global symbols (e.g. the obvious
     // implementation of multiple entry points).  If this doesn't occur, the
     // linker can safely perform dead code stripping.  Since LLVM never
     // generates code that does this, it is always safe to set.
     OutStreamer->EmitAssemblerFlag(MCAF_SubsectionsViaSymbols);
   }
if (TT.isOSBinFormatCOFF()) {
     }
 
     SM.serializeToStackMapSection();
   }
 
   if (TT.isOSBinFormatELF()) {
     SM.serializeToStackMapSection();
     FM.serializeToFaultMapSection();
+    EmitXRayTable();
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Target Registry Stuff
 //===----------------------------------------------------------------------===//
 
 // Force static initialization.
 extern "C" void LLVMInitializeX86AsmPrinter() {
   RegisterAsmPrinter<X86AsmPrinter> X(TheX86_32Target);
   RegisterAsmPrinter<X86AsmPrinter> Y(TheX86_64Target);
 }

lib/Target/X86/X86FrameLowering.cpp

static unsigned findDeadCallerSavedReg(MachineBasicBlock &MBB,
   if (!F || MF->getMMI().callsEHReturn())
     return 0;
 
   const TargetRegisterClass &AvailableRegs = *TRI->getGPRsForTailCall(*MF);
 
   unsigned Opc = MBBI->getOpcode();
   switch (Opc) {
   default: return 0;
+  case TargetOpcode::PATCHABLE_RET:
   case X86::RET:
   case X86::RETL:
   case X86::RETQ:
   case X86::RETIL:
   case X86::RETIQ:
   case X86::TCRETURNdi:
   case X86::TCRETURNri:
   case X86::TCRETURNmi:

lib/Target/X86/X86MCInstLower.cpp

 #include "llvm/IR/Mangler.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCCodeEmitter.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCFixup.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCInstBuilder.h"
+#include "llvm/MC/MCSection.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSymbol.h"
 #include "llvm/Support/TargetRegistry.h"
+#include "llvm/Target/TargetLoweringObjectFile.h"
+
 using namespace llvm;
 
 namespace {
 
 /// X86MCInstLower - This class is used to lower an MachineInstr into an MCInst.
 class X86MCInstLower {
   MCContext &Ctx;
   const MachineFunction &MF;
void X86AsmPrinter::LowerPATCHPOINT(const MachineInstr &MI,
   unsigned NumBytes = opers.getMetaOper(PatchPointOpers::NBytesPos).getImm();
   assert(NumBytes >= EncodedBytes &&
          "Patchpoint can't request size less than the length of a call.");
 
   EmitNops(*OutStreamer, NumBytes - EncodedBytes, Subtarget->is64Bit(),
            getSubtargetInfo());
 }
 
+void X86AsmPrinter::recordSled(MCSymbol *Sled, const MachineInstr &MI) {
+  auto Fn = MI.getParent()->getParent()->getFunction();
+  auto Attr = Fn->getFnAttribute("function-instrument");
+  bool AlwaysInstrument = false;
+  if (!Attr.hasAttribute(Attribute::None) && Attr.isStringAttribute() &&
+      Attr.getValueAsString() == "xray-always")
+    AlwaysInstrument = true;
+  Sleds.push_back(
+      {Sled, CurrentFnSym, SledKind::FUNCTION_ENTER, AlwaysInstrument});
+}
+
+void X86AsmPrinter::LowerPATCHABLE_FUNCTION_ENTER(const MachineInstr &MI,
+                                                  X86MCInstLower &MCIL) {
+  // We want to emit the following pattern:
+  //
+  // .Lxray_sled_N:
+  //   .palign 2, ...
+  //   jmp .tmpN
+  //   # 9 bytes worth of noops
+  // .tmpN
+  //
+  auto CurSled = OutContext.createTempSymbol("xray_sled_", true);
+  OutStreamer->EmitLabel(CurSled);
+  OutStreamer->EmitCodeAlignment(4);
+  auto Target = OutContext.createLinkerPrivateTempSymbol();
+  OutStreamer->EmitInstruction(
+      MCInstBuilder(X86::JMP_1)
+          .addExpr(MCSymbolRefExpr::create(Target, OutContext)),
+      getSubtargetInfo());
+  EmitNops(*OutStreamer, 9, Subtarget->is64Bit(), getSubtargetInfo());
+  OutStreamer->EmitLabel(Target);
+  recordSled(CurSled, MI);
+}
+
+void X86AsmPrinter::LowerPATCHABLE_RET(const MachineInstr &MI,
+                                            X86MCInstLower &MCIL) {
+  // Since PATCHABLE_RET takes the opcode of the return statement as an
+  // argument, we use that to emit the correct form of the RET that we want.
+  // i.e. when we see this:
+  //
+  //   PATCHABLE_RET X86::RET ...
+  //
+  // We should emit the RET followed by sleds.
+  //
+  // .Lxray_sled_N:
+  //   ret  # or equivalent instruction
+  //   # 10 bytes worth of noops
+  //
+  // This just makes sure that the alignment for the next instruction is 2.
+  auto CurSled = OutContext.createTempSymbol("xray_sled_", true);
+  OutStreamer->EmitLabel(CurSled);
+  unsigned OpCode = MI.getOperand(0).getImm();
+  MCInst Ret;
+  Ret.setOpcode(OpCode);
+  for (auto &MO : make_range(MI.operands_begin() + 1, MI.operands_end()))
+    if (auto MaybeOperand = MCIL.LowerMachineOperand(&MI, MO))
+      Ret.addOperand(MaybeOperand.getValue());
+  OutStreamer->EmitInstruction(Ret, getSubtargetInfo());
+  EmitNops(*OutStreamer, 10, Subtarget->is64Bit(), getSubtargetInfo());
+  recordSled(CurSled, MI);
+}
+
+void X86AsmPrinter::LowerPATCHABLE_TAIL_CALL(const MachineInstr &MI,
+                                        X86MCInstLower &MCIL) {
+  // FIXME: Do something later.
+}
+
+void X86AsmPrinter::EmitXRayTable() {
+  // Do not emit the XRay section if there were no sleds encountered.
+  if (Sleds.empty()) return;
+
+  MCSection *XRaySection = getObjFileLowering().getXRaySection();
+  assert(XRaySection && "There should be an XRay Section!");
+  auto PrevSection = OutStreamer->getCurrentSectionOnly();
+  OutStreamer->SwitchSection(XRaySection);
+
+  // We lay it out so that we have fixed-sized records, where we emit a label
+  // reference (8 bytes), then a byte for the kind of entry, another byte for
+  // whether to always instrument, and a 64-bit (8 byte) value.
+  for (const auto &Sled : Sleds) {
+    OutStreamer->EmitSymbolValue(Sled.Sled, 8);
+    OutStreamer->EmitSymbolValue(Sled.Function, 8);
+    uint8_t Kind = static_cast<uint8_t>(Sled.Kind);
+    OutStreamer->EmitBytes(StringRef(reinterpret_cast<const char *>(&Kind), 1));
+    uint8_t AlwaysInstrument = Sled.AlwaysInstrument ? 1 : 0;
+    OutStreamer->EmitBytes(
+        StringRef(reinterpret_cast<const char *>(&AlwaysInstrument), 1));
+  }
+
+  // Switch back to wherever we were.
+  OutStreamer->SwitchSection(PrevSection);
+}
+
 // Returns instruction preceding MBBI in MachineFunction.
 // If MBBI is the first instruction of the first basic block, returns null.
 static MachineBasicBlock::const_iterator
 PrevCrossBBInst(MachineBasicBlock::const_iterator MBBI) {
   const MachineBasicBlock *MBB = MBBI->getParent();
   while (MBBI == MBB->begin()) {
     if (MBB == &MBB->getParent()->front())
       return nullptr;
case TargetOpcode::PATCHABLE_OP:
     return LowerPATCHABLE_OP(*MI, MCInstLowering);
 
   case TargetOpcode::STACKMAP:
     return LowerSTACKMAP(*MI);
 
   case TargetOpcode::PATCHPOINT:
     return LowerPATCHPOINT(*MI, MCInstLowering);
 
+  case TargetOpcode::PATCHABLE_FUNCTION_ENTER:
+    return LowerPATCHABLE_FUNCTION_ENTER(*MI, MCInstLowering);
+
+  case TargetOpcode::PATCHABLE_RET:
+    return LowerPATCHABLE_RET(*MI, MCInstLowering);
+
+  case TargetOpcode::PATCHABLE_TAIL_CALL:
+    return LowerPATCHABLE_TAIL_CALL(*MI, MCInstLowering);
+
   case X86::MORESTACK_RET:
     EmitAndCountInstruction(MCInstBuilder(getRetOpcode(*Subtarget)));
     return;
 
   case X86::MORESTACK_RET_RESTORE_R10:
     // Return, then restore R10.
     EmitAndCountInstruction(MCInstBuilder(getRetOpcode(*Subtarget)));
     EmitAndCountInstruction(MCInstBuilder(X86::MOV64rr)

test/CodeGen/X86/xray-attribute-instrumentation.ll

+; RUN: llc -filetype=asm -o - -mtriple=x86_64-apple-macosx < %s | FileCheck %s
+
+define i32 @foo() nounwind noinline uwtable "function-instrument"="xray-always" {
+; CHECK-LABEL: Lxray_sled_0:
+; CHECK-NEXT:  .p2align 2, 0x90
+; CHECK-NEXT:  jmp ltmp0
+; CHECK-NEXT:  nopw 512(%rax,%rax)
+; CHECK-LABEL: ltmp0:
+  ret i32 0
+; CHECK-LABEL: Lxray_sled_1:
+; CHECK-NEXT:  retq
+; CHECK-NEXT:  nopw %cs:512(%rax,%rax)
+}

test/CodeGen/X86/xray-selective-instrumentation-miss.ll

+; RUN: llc -mcpu=nehalem < %s | not grep xray_sled_
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128"
+target triple = "x86_64-apple-darwin8"
+
+define i32 @foo() nounwind uwtable "xray-instruction-threshold"="3" {
+entry:
+  ret i32 0
+}

test/CodeGen/X86/xray-selective-instrumentation.ll

+; RUN: llc -mcpu=nehalem < %s | grep xray_sled_
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128"
+target triple = "x86_64-apple-darwin8"
+
+define i32 @foo() nounwind uwtable "xray-instruction-threshold"="1" {
+entry:
+  ret i32 0
+}

test/TableGen/trydecode-emission.td

def InstB : TestInstruction {
   let AsmString = "InstB";
   let DecoderMethod = "DecodeInstB";
   let hasCompleteDecoder = 0;
 }
 
 // CHECK:      /* 0 */       MCD::OPC_ExtractField, 4, 4,  // Inst{7-4} ...
 // CHECK-NEXT: /* 3 */       MCD::OPC_FilterValue, 0, 14, 0, // Skip to: 21
 // CHECK-NEXT: /* 7 */       MCD::OPC_CheckField, 2, 2, 0, 5, 0, // Skip to: 18
-// CHECK-NEXT: /* 13 */      MCD::OPC_TryDecode, 27, 0, 0, 0, // Opcode: InstB, skip to: 18
-// CHECK-NEXT: /* 18 */      MCD::OPC_Decode, 26, 1, // Opcode: InstA
+// CHECK-NEXT: /* 13 */      MCD::OPC_TryDecode, 30, 0, 0, 0, // Opcode: InstB, skip to: 18
+// CHECK-NEXT: /* 18 */      MCD::OPC_Decode, 29, 1, // Opcode: InstA
 // CHECK-NEXT: /* 21 */      MCD::OPC_Fail,
 
 // CHECK: if (DecodeInstB(MI, insn, Address, Decoder) == MCDisassembler::Fail) { DecodeComplete = false; return MCDisassembler::Fail; }

test/TableGen/trydecode-emission2.td

def InstB : TestInstruction {
   let hasCompleteDecoder = 0;
 }
 
 // CHECK:      /* 0 */       MCD::OPC_ExtractField, 2, 1,  // Inst{2} ...
 // CHECK-NEXT: /* 3 */       MCD::OPC_FilterValue, 0, 29, 0, // Skip to: 36
 // CHECK-NEXT: /* 7 */       MCD::OPC_ExtractField, 5, 3,  // Inst{7-5} ...
 // CHECK-NEXT: /* 10 */      MCD::OPC_FilterValue, 0, 22, 0, // Skip to: 36
 // CHECK-NEXT: /* 14 */      MCD::OPC_CheckField, 0, 2, 3, 5, 0, // Skip to: 25
-// CHECK-NEXT: /* 20 */      MCD::OPC_TryDecode, 27, 0, 0, 0, // Opcode: InstB, skip to: 25
+// CHECK-NEXT: /* 20 */      MCD::OPC_TryDecode, 30, 0, 0, 0, // Opcode: InstB, skip to: 25
 // CHECK-NEXT: /* 25 */      MCD::OPC_CheckField, 3, 2, 0, 5, 0, // Skip to: 36
-// CHECK-NEXT: /* 31 */      MCD::OPC_TryDecode, 26, 1, 0, 0, // Opcode: InstA, skip to: 36
+// CHECK-NEXT: /* 31 */      MCD::OPC_TryDecode, 29, 1, 0, 0, // Opcode: InstA, skip to: 36
 // CHECK-NEXT: /* 36 */      MCD::OPC_Fail,
 
 // CHECK: if (DecodeInstB(MI, insn, Address, Decoder) == MCDisassembler::Fail) { DecodeComplete = false; return MCDisassembler::Fail; }
 // CHECK: if (DecodeInstA(MI, insn, Address, Decoder) == MCDisassembler::Fail) { DecodeComplete = false; return MCDisassembler::Fail; }

test/TableGen/trydecode-emission3.td

def InstB : TestInstruction {
   let Inst{1-0} = op;
   let OutOperandList = (outs InstBOp:$op);
   let AsmString = "InstB";
 }
 
 // CHECK:      /* 0 */       MCD::OPC_ExtractField, 4, 4,  // Inst{7-4} ...
 // CHECK-NEXT: /* 3 */       MCD::OPC_FilterValue, 0, 14, 0, // Skip to: 21
 // CHECK-NEXT: /* 7 */       MCD::OPC_CheckField, 2, 2, 0, 5, 0, // Skip to: 18
-// CHECK-NEXT: /* 13 */      MCD::OPC_TryDecode, 27, 0, 0, 0, // Opcode: InstB, skip to: 18
-// CHECK-NEXT: /* 18 */      MCD::OPC_Decode, 26, 1, // Opcode: InstA
+// CHECK-NEXT: /* 13 */      MCD::OPC_TryDecode, 30, 0, 0, 0, // Opcode: InstB, skip to: 18
+// CHECK-NEXT: /* 18 */      MCD::OPC_Decode, 29, 1, // Opcode: InstA
 // CHECK-NEXT: /* 21 */      MCD::OPC_Fail,
 
 // CHECK: if (DecodeInstBOp(MI, tmp, Address, Decoder) == MCDisassembler::Fail) { DecodeComplete = false; return MCDisassembler::Fail; }