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

Ricky Zhou via llvm-commits llvm-commits at lists.llvm.org
Mon May 27 20:56:29 PDT 2024


https://github.com/rickyz updated https://github.com/llvm/llvm-project/pull/89360

>From abb580e748d9e31f794a993e9bb6323002d04cfd Mon Sep 17 00:00:00 2001
From: Ricky Zhou <ricky at rzhou.org>
Date: Fri, 19 Apr 2024 01:52:28 -0700
Subject: [PATCH 1/3] [xray] Fix stack alignment for custom event calls.

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.
---
 llvm/lib/Target/X86/X86ISelLowering.cpp       | 27 ++++++++++++++++++-
 llvm/lib/Target/X86/X86ISelLowering.h         |  3 +++
 .../CodeGen/X86/xray-leaf-stack-alignment.ll  | 18 +++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/X86/xray-leaf-stack-alignment.ll

diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index bedec0c8974a8..d45e7eb1d1a01 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -36254,6 +36254,31 @@ X86TargetLowering::EmitSjLjDispatchBlock(MachineInstr &MI,
   return BB;
 }
 
+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.
+  BB->getParent()->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 {
@@ -36484,7 +36509,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 e348ba6e8ac08..78bac6e077463 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.h
+++ b/llvm/lib/Target/X86/X86ISelLowering.h
@@ -1808,6 +1808,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-leaf-stack-alignment.ll b/llvm/test/CodeGen/X86/xray-leaf-stack-alignment.ll
new file mode 100644
index 0000000000000..f6dd8f527d39d
--- /dev/null
+++ b/llvm/test/CodeGen/X86/xray-leaf-stack-alignment.ll
@@ -0,0 +1,18 @@
+; RUN: llc -mtriple=x86_64 < %s | FileCheck %s
+
+ at leaf_func.event_id = internal constant i32 1, align 4
+
+define void @leaf_func() "xray-instruction-threshold"="999" "frame-pointer"="none" nounwind {
+  ; CHECK-LABEL: leaf_func:
+  ; CHECK-NEXT:  .Lfunc_begin0:
+  ; CHECK-NEXT:  # %bb.0:
+  ; CHECK-NEXT:    pushq %rax
+  ; CHECK-NEXT:    movl $leaf_func.event_id, %eax
+  ; CHECK-NEXT:    movl $8, %ecx
+  ; CHECK-NEXT:    .p2align 1, 0x90
+  ; CHECK-NEXT:  .Lxray_event_sled_0:
+  call void @llvm.xray.customevent(ptr @leaf_func.event_id, i64 8)
+  ret void
+}
+
+declare void @llvm.xray.customevent(ptr nocapture readonly, i64)

>From 5998490edce3ab5de0c7cf291e7bd329f5ae3508 Mon Sep 17 00:00:00 2001
From: Ricky Zhou <ricky at rzhou.org>
Date: Mon, 27 May 2024 20:53:26 -0700
Subject: [PATCH 2/3] fixup! [xray] Fix stack alignment for custom event calls.

Add file comment to test.
---
 llvm/test/CodeGen/X86/xray-leaf-stack-alignment.ll | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/llvm/test/CodeGen/X86/xray-leaf-stack-alignment.ll b/llvm/test/CodeGen/X86/xray-leaf-stack-alignment.ll
index f6dd8f527d39d..6061b2be129ed 100644
--- a/llvm/test/CodeGen/X86/xray-leaf-stack-alignment.ll
+++ b/llvm/test/CodeGen/X86/xray-leaf-stack-alignment.ll
@@ -1,5 +1,8 @@
 ; RUN: llc -mtriple=x86_64 < %s | FileCheck %s
 
+;; 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() "xray-instruction-threshold"="999" "frame-pointer"="none" nounwind {

>From d8b6765b8fc56ec5ae39b1ff28f6cfd279706418 Mon Sep 17 00:00:00 2001
From: Ricky Zhou <ricky at rzhou.org>
Date: Mon, 27 May 2024 20:56:11 -0700
Subject: [PATCH 3/3] fixup! fixup! [xray] Fix stack alignment for custom event
 calls.

Use MachineFunction::getFrameInfo
---
 llvm/lib/Target/X86/X86ISelLowering.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index d45e7eb1d1a01..0affe60026177 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -36264,7 +36264,7 @@ X86TargetLowering::emitPatchableEventCall(MachineInstr &MI,
   MachineFunction &MF = *BB->getParent();
 
   // Emit CALLSEQ_START right before the instruction.
-  BB->getParent()->getFrameInfo().setAdjustsStack(true);
+  MF.getFrameInfo().setAdjustsStack(true);
   unsigned AdjStackDown = TII.getCallFrameSetupOpcode();
   MachineInstrBuilder CallseqStart =
       BuildMI(MF, MIMD, TII.get(AdjStackDown)).addImm(0).addImm(0).addImm(0);



More information about the llvm-commits mailing list