[XRay][X86] Fix stack alignment for custom event calls (#89360)

Calls to @llvm.xray.{custom,typed}event are lowered to an nop sled that
can be
patched with a call to an instrumentation function at runtime. Prior to
this
change, x86 codegen did not guarantee that these patched calls run with
the
appropriate stack alignment (particularly in leaf functions, where
normal stack
alignment assumptions may not hold).

This leads to crashes on x86, as the custom event hook can end up using
SSE
instructions that assume a 16-byte aligned stack.

Fix this by wrapping custom event hooks in CALLSEQ_START/CALLSEQ_END as
done for
regular function calls. One downside of this approach is that on
functions whose
stacks aren't already aligned, we may end up running stack alignment
fixup
instructions even when instrumentation is disabled. An alternative could
be to
make the custom event assembly trampolines stack-alignment-agnostic.
This was
the case in the past, but

https://github.com/llvm/llvm-project/commit/b46c89892fe25bec197fd30f09b3a312da126422
removed this due to complexity in maintaining CFI directives for these
stack
adjustments. Since we are already willing to pay the call argument setup
cost
for custom hooks when instrumentation is disabled, I am hoping that an
extra
push/pop in this hopefully uncommon unaligned stack case is tolerable.

3 files changed, 44 insertions, 1 deletions

diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index ce6adaa..7df8ffb 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -36378,6 +36378,31 @@ X86TargetLowering::EmitSjLjDispatchBlock(MachineInstr &MI,
 }
 
 MachineBasicBlock *
+X86TargetLowering::emitPatchableEventCall(MachineInstr &MI,
+                                          MachineBasicBlock *BB) const {
+  // Wrap patchable event calls in CALLSEQ_START/CALLSEQ_END, as tracing
+  // calls may require proper stack alignment.
+  const TargetInstrInfo &TII = *Subtarget.getInstrInfo();
+  const MIMetadata MIMD(MI);
+  MachineFunction &MF = *BB->getParent();
+
+  // Emit CALLSEQ_START right before the instruction.
+  MF.getFrameInfo().setAdjustsStack(true);
+  unsigned AdjStackDown = TII.getCallFrameSetupOpcode();
+  MachineInstrBuilder CallseqStart =
+      BuildMI(MF, MIMD, TII.get(AdjStackDown)).addImm(0).addImm(0).addImm(0);
+  BB->insert(MachineBasicBlock::iterator(MI), CallseqStart);
+
+  // Emit CALLSEQ_END right after the instruction.
+  unsigned AdjStackUp = TII.getCallFrameDestroyOpcode();
+  MachineInstrBuilder CallseqEnd =
+      BuildMI(MF, MIMD, TII.get(AdjStackUp)).addImm(0).addImm(0);
+  BB->insertAfter(MachineBasicBlock::iterator(MI), CallseqEnd);
+
+  return BB;
+}
+
+MachineBasicBlock *
 X86TargetLowering::EmitInstrWithCustomInserter(MachineInstr &MI,
                                                MachineBasicBlock *BB) const {
   MachineFunction *MF = BB->getParent();
@@ -36607,7 +36632,7 @@ X86TargetLowering::EmitInstrWithCustomInserter(MachineInstr &MI,
 
   case TargetOpcode::PATCHABLE_EVENT_CALL:
   case TargetOpcode::PATCHABLE_TYPED_EVENT_CALL:
-    return BB;
+    return emitPatchableEventCall(MI, BB);
 
   case X86::LCMPXCHG8B: {
     const X86RegisterInfo *TRI = Subtarget.getRegisterInfo();
diff --git a/llvm/lib/Target/X86/X86ISelLowering.h b/llvm/lib/Target/X86/X86ISelLowering.h
index 1facd1d..b0efa07 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.h
+++ b/llvm/lib/Target/X86/X86ISelLowering.h
@@ -1798,6 +1798,9 @@ namespace llvm {
     MachineBasicBlock *EmitSjLjDispatchBlock(MachineInstr &MI,
                                              MachineBasicBlock *MBB) const;
 
+    MachineBasicBlock *emitPatchableEventCall(MachineInstr &MI,
+                                              MachineBasicBlock *MBB) const;
+
     /// Emit flags for the given setcc condition and operands. Also returns the
     /// corresponding X86 condition code constant in X86CC.
     SDValue emitFlagsForSetcc(SDValue Op0, SDValue Op1, ISD::CondCode CC,
diff --git a/llvm/test/CodeGen/X86/xray-custom-log.ll b/llvm/test/CodeGen/X86/xray-custom-log.ll
index c41aa1e..9ccda8e 100644
--- a/llvm/test/CodeGen/X86/xray-custom-log.ll
+++ b/llvm/test/CodeGen/X86/xray-custom-log.ll
@@ -75,6 +75,21 @@ define i32 @typedevent() nounwind "function-instrument"="xray-always" !dbg !2 {
 ; CHECK-LABEL: Lxray_sleds_start1:
 ; CHECK:       .quad {{.*}}xray_typed_event_sled_0
 
+; Verify that custom event calls are done with proper stack alignment,
+; even in leaf functions.
+@leaf_func.event_id = internal constant i32 1, align 4
+define void @leaf_func() "function-instrument"="xray-always" "frame-pointer"="none" nounwind {
+  ; CHECK-LABEL: leaf_func:
+  ; CHECK-NEXT:  .Lfunc_begin2:
+  ; CHECK:         pushq %rax
+  ; CHECK:         movl $leaf_func.event_id, %eax
+  ; CHECK-NEXT:    movl $4, %ecx
+  ; CHECK-NEXT:    .p2align 1, 0x90
+  ; CHECK-NEXT:  .Lxray_event_sled_1:
+  call void @llvm.xray.customevent(ptr @leaf_func.event_id, i64 4)
+  ret void
+}
+
 declare void @llvm.xray.customevent(ptr, i64)
 declare void @llvm.xray.typedevent(i64, ptr, i64)