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

[dberris]

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.

[dberris]

Rebase

[echristo]

Probably want to split out just the NFC of making isTailCall a function on TII.

[dberris]

In D23986#531036, @echristo wrote:

Probably want to split out just the NFC of making isTailCall a function on TII.

Good idea! Let me do that first, then rebase again. :)

[dberris]

Rebase

[echristo]

LGTM.

-eric

[rSerge]

llvm/trunk/lib/Target/X86/X86MCInstLower.cpp

1126 ↗

(On Diff #69939)

Why all the hassle with passing the original tail call instruction and its operands in the arguments of the patchable tail call instruction, instead of just prepending the original tail call instruction with the sled code?

[rSerge]

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

12 ↗

(On Diff #69939)

Isn't it missing setting EAX to 0? The function in the disassembly seems to just return;, rather than return 0; as, I guess, ret i32 0 requires.

[rSerge]

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

18–20 ↗

(On Diff #69939)

Perhaps I missed some patches. Where are these generated? And also .quad .Lxray_synthetic_0 2 lines above.

[dberris]

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

12 ↗

(On Diff #69939)

Nope, this is in a test -- the set to eax 0 happens before the return. We're only looking for the return instruction sled here.

18–20 ↗

(On Diff #69939)

This was introduced in D23398. The discussion of why is there.

include/llvm/Target/Target.td

 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 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"
 
 }
 
 //===----------------------------------------------------------------------===//
 // AsmParser - This class can be implemented by targets that wish to implement

include/llvm/Target/TargetOpcodes.def

 /// for inserting instrumentation instructions at runtime.
 HANDLE_TARGET_OPCODE(PATCHABLE_FUNCTION_ENTER)
 
 /// Wraps a return instruction and its operands to enable adding nop sleds
 /// either before or after the return. The nop sleds are useful for inserting
 /// 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.
 
 /// Generic ADD instruction. This is an integer add.
 HANDLE_TARGET_OPCODE(G_ADD)
 HANDLE_TARGET_OPCODE_MARKER(PRE_ISEL_GENERIC_OPCODE_START, G_ADD)
 

lib/CodeGen/XRayInstrumentation.cpp

bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
   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 &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);
         Terminators.push_back(&T);
       }
     }
   }
 

lib/Target/X86/X86MCInstLower.cpp

void X86AsmPrinter::LowerPATCHABLE_RET(const MachineInstr &MI,
   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, 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;
   if (Subtarget->isTargetELF()) {
     auto PrevSection = OutStreamer->getCurrentSectionOnly();
     auto Fn = MF->getFunction();
     MCSection *Section = nullptr;
     if (Fn->hasComdat()) {
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-tail-call-sled.ll

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