[XRay] Detect and emit sleds for sibling/tail calls

Summary:
This change promotes the 'isTailCall(...)' member function to
TargetInstrInfo as a query interface for determining on a per-target
basis whether a given MachineInstr is a tail call instruction. We build
upon this in the XRay instrumentation pass to emit special sleds for
tail call optimisations, where we emit the correct kind of sled.

The tail call sleds look like a mix between the function entry and
function exit sleds. Form-wise, the sled comes before the "jmp"
instruction that implements the tail call similar to how we do it for
the function entry sled. Functionally, because we know this is a tail
call, it behaves much like an exit sled -- i.e. at runtime we may use
the exit trampolines instead of a different kind of trampoline.

A follow-up change to recognise these sleds will be done in compiler-rt,
so that we can start intercepting these initially as exits, but also
have the option to have different log entries to more accurately reflect
that this is actually a tail call.

Reviewers: echristo, rSerge, majnemer

Subscribers: mehdi_amini, dberris, llvm-commits

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

llvm-svn: 280334

include/llvm/Target/Target.td

@@ -961,6 +961,14 @@ def PATCHABLE_RET : Instruction {
   let hasSideEffects = 1;
   let isReturn = 1;
 }
+def PATCHABLE_TAIL_CALL : Instruction {
+  let OutOperandList = (outs unknown:$dst);
+  let InOperandList = (ins variable_ops);
+  let AsmString = "# XRay Tail Call Exit.";
+  let usesCustomInserter = 1;
+  let hasSideEffects = 1;
+  let isReturn = 1;
+}
 
 // Generic opcodes used in GlobalISel.
 include "llvm/Target/GenericOpcodes.td"

include/llvm/Target/TargetOpcodes.def

@@ -155,6 +155,11 @@ HANDLE_TARGET_OPCODE(PATCHABLE_FUNCTION_ENTER)
 /// instrumentation instructions at runtime.
 HANDLE_TARGET_OPCODE(PATCHABLE_RET)
 
+/// Wraps a tail call instruction and its operands to enable adding nop sleds
+/// either before or after the tail exit. We use this as a disambiguation from
+/// PATCHABLE_RET which specifically only works for return instructions.
+HANDLE_TARGET_OPCODE(PATCHABLE_TAIL_CALL)
+
 /// 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.

lib/CodeGen/XRayInstrumentation.cpp

@@ -69,12 +69,19 @@ bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
   SmallVector<MachineInstr *, 4> Terminators;
   for (auto &MBB : MF) {
     for (auto &T : MBB.terminators()) {
-      // FIXME: Handle tail calls here too?
+      unsigned Opc = 0;
       if (T.isReturn() && T.getOpcode() == TII->getReturnOpcode()) {
         // Replace return instructions with:
         //   PATCHABLE_RET <Opcode>, <Operand>...
-        auto MIB = BuildMI(MBB, T, T.getDebugLoc(),
-                           TII->get(TargetOpcode::PATCHABLE_RET))
+        Opc = TargetOpcode::PATCHABLE_RET;
+      }
+      if (TII->isTailCall(T)) {
+        // Treat the tail call as a return instruction, which has a
+        // different-looking sled than the normal return case.
+        Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
+      }
+      if (Opc != 0) {
+        auto MIB = BuildMI(MBB, T, T.getDebugLoc(), TII->get(Opc))
                        .addImm(T.getOpcode());
         for (auto &MO : T.operands())
           MIB.addOperand(MO);

lib/Target/X86/X86MCInstLower.cpp

@@ -1093,6 +1093,39 @@ void X86AsmPrinter::LowerPATCHABLE_RET(const MachineInstr &MI,
   recordSled(CurSled, MI, SledKind::FUNCTION_EXIT);
 }
 
+void X86AsmPrinter::LowerPATCHABLE_TAIL_CALL(const MachineInstr &MI, X86MCInstLower &MCIL) {
+  // Like PATCHABLE_RET, we have the actual instruction in the operands to this
+  // instruction so we lower that particular instruction and its operands.
+  // Unlike PATCHABLE_RET though, we put the sled before the JMP, much like how
+  // we do it for PATCHABLE_FUNCTION_ENTER. The sled should be very similar to
+  // the PATCHABLE_FUNCTION_ENTER case, followed by the lowering of the actual
+  // tail call much like how we have it in PATCHABLE_RET.
+  auto CurSled = OutContext.createTempSymbol("xray_sled_", true);
+  OutStreamer->EmitCodeAlignment(2);
+  OutStreamer->EmitLabel(CurSled);
+  auto Target = OutContext.createTempSymbol();
+
+  // Use a two-byte `jmp`. This version of JMP takes an 8-bit relative offset as
+  // an operand (computed as an offset from the jmp instruction).
+  // FIXME: Find another less hacky way do force the relative jump.
+  OutStreamer->EmitBytes("\xeb\x09");
+  EmitNops(*OutStreamer, 9, Subtarget->is64Bit(), getSubtargetInfo());
+  OutStreamer->EmitLabel(Target);
+  recordSled(CurSled, MI, SledKind::TAIL_CALL);
+
+  unsigned OpCode = MI.getOperand(0).getImm();
+  MCInst TC;
+  TC.setOpcode(OpCode);
+
+  // Before emitting the instruction, add a comment to indicate that this is
+  // indeed a tail call.
+  OutStreamer->AddComment("TAILCALL");
+  for (auto &MO : make_range(MI.operands_begin() + 1, MI.operands_end()))
+    if (auto MaybeOperand = MCIL.LowerMachineOperand(&MI, MO))
+      TC.addOperand(MaybeOperand.getValue());
+  OutStreamer->EmitInstruction(TC, getSubtargetInfo());
+}
+
 void X86AsmPrinter::EmitXRayTable() {
   if (Sleds.empty())
     return;
@@ -1383,6 +1416,9 @@ void X86AsmPrinter::EmitInstruction(const MachineInstr *MI) {
   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;

test/CodeGen/X86/xray-tail-call-sled.ll

@@ -0,0 +1,41 @@
+; RUN: llc -filetype=asm -o - -mtriple=x86_64-unknown-linux-gnu < %s | FileCheck %s
+
+define i32 @callee() nounwind noinline uwtable "function-instrument"="xray-always" {
+; CHECK:       .p2align 1, 0x90
+; CHECK-LABEL: Lxray_sled_0:
+; CHECK-NEXT:  .ascii "\353\t"
+; CHECK-NEXT:  nopw 512(%rax,%rax)
+; CHECK-LABEL: Ltmp0:
+  ret i32 0
+; CHECK:       .p2align 1, 0x90
+; CHECK-LABEL: Lxray_sled_1:
+; CHECK-NEXT:  retq
+; CHECK-NEXT:  nopw %cs:512(%rax,%rax)
+}
+; CHECK:       .p2align 4, 0x90
+; CHECK-NEXT:  .quad .Lxray_synthetic_0
+; CHECK-NEXT:  .section xray_instr_map,{{.*}}
+; CHECK-LABEL: Lxray_synthetic_0:
+; CHECK:       .quad .Lxray_sled_0
+; CHECK:       .quad .Lxray_sled_1
+
+define i32 @caller() nounwind noinline uwtable "function-instrument"="xray-always" {
+; CHECK:       .p2align 1, 0x90
+; CHECK-LABEL: Lxray_sled_2:
+; CHECK-NEXT:  .ascii "\353\t"
+; CHECK-NEXT:  nopw 512(%rax,%rax)
+; CHECK-LABEL: Ltmp1:
+; CHECK:       .p2align 1, 0x90
+; CHECK-LABEL: Lxray_sled_3:
+; CHECK-NEXT:  .ascii "\353\t"
+; CHECK-NEXT:  nopw 512(%rax,%rax)
+; CHECK-LABEL: Ltmp2:
+  %retval = tail call i32 @callee()
+; CHECK:       jmp callee  # TAILCALL
+  ret i32 %retval
+}
+; CHECK:       .p2align 4, 0x90
+; CHECK-NEXT:  .quad .Lxray_synthetic_1
+; CHECK-LABEL: Lxray_synthetic_1:
+; CHECK:       .quad .Lxray_sled_2
+; CHECK:       .quad .Lxray_sled_3