[llvm] [X86] Don't save/restore fp around longjmp instructions (PR #102556)

[via llvm-commits]

https://github.com/weiguozhi created https://github.com/llvm/llvm-project/pull/102556

Longjmp instructions can also modify fp register, it is expected behavior. We should not save/restore fp around these instructions.

>From 9cae8fbc5792b59026a48cffdc182eae32982206 Mon Sep 17 00:00:00 2001
From: Guozhi Wei <carrot at google.com>
Date: Thu, 8 Aug 2024 17:47:33 -0700
Subject: [PATCH] [X86] Don't save/restore fp around longjmp instructions

Longjmp instructions can also modify fp register, it is expected
behavior. We should not save/restore fp around these instructions.
---
 llvm/lib/Target/X86/X86FrameLowering.cpp     |  6 +++
 llvm/lib/Target/X86/X86ISelLowering.cpp      |  2 +
 llvm/lib/Target/X86/X86MachineFunctionInfo.h | 11 +++++
 llvm/test/CodeGen/X86/clobber_frame_ptr.ll   | 43 ++++++++++++++++++++
 4 files changed, 62 insertions(+)

diff --git a/llvm/lib/Target/X86/X86FrameLowering.cpp b/llvm/lib/Target/X86/X86FrameLowering.cpp
index 8404f2231680d..8760e97e684e3 100644
--- a/llvm/lib/Target/X86/X86FrameLowering.cpp
+++ b/llvm/lib/Target/X86/X86FrameLowering.cpp
@@ -4389,6 +4389,12 @@ bool X86FrameLowering::skipSpillFPBP(
       ++MI;
     return true;
   }
+
+  // FP clobbered by longjmp is expected.
+  X86MachineFunctionInfo *X86FI = MF.getInfo<X86MachineFunctionInfo>();
+  if (X86FI->containLongJmpMIClobberFP(&*MI))
+    return true;
+
   return false;
 }
 
diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index f69606783f25c..6974d18814eb7 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -36205,6 +36205,8 @@ X86TargetLowering::emitEHSjLjLongJmp(MachineInstr &MI,
       MIB.add(MO);
   }
   MIB.setMemRefs(MMOs);
+  auto *MFI = MF->getInfo<X86MachineFunctionInfo>();
+  MFI->insertLongJmpMIClobberFP(MIB.getInstr());
 
   // Reload IP
   MIB = BuildMI(*thisMBB, MI, MIMD, TII->get(PtrLoadOpc), Tmp);
diff --git a/llvm/lib/Target/X86/X86MachineFunctionInfo.h b/llvm/lib/Target/X86/X86MachineFunctionInfo.h
index 13d57c2fa9dfb..bf3ffdc8677bf 100644
--- a/llvm/lib/Target/X86/X86MachineFunctionInfo.h
+++ b/llvm/lib/Target/X86/X86MachineFunctionInfo.h
@@ -14,6 +14,7 @@
 #define LLVM_LIB_TARGET_X86_X86MACHINEFUNCTIONINFO_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/CodeGen/CallingConvLower.h"
 #include "llvm/CodeGen/MIRYamlMapping.h"
@@ -174,6 +175,9 @@ class X86MachineFunctionInfo : public MachineFunctionInfo {
   bool FPClobberedByCall = false;
   bool BPClobberedByCall = false;
 
+  // A set of instructions clobber FP generated from longjmp.
+  SmallSet<MachineInstr *, 2> LongJmpMIClobberFP;
+
 private:
   /// ForwardedMustTailRegParms - A list of virtual and physical registers
   /// that must be forwarded to every musttail call.
@@ -338,6 +342,13 @@ class X86MachineFunctionInfo : public MachineFunctionInfo {
 
   bool getBPClobberedByCall() const { return BPClobberedByCall; }
   void setBPClobberedByCall(bool C) { BPClobberedByCall = C; }
+
+  bool containLongJmpMIClobberFP(MachineInstr *MI) const {
+    return LongJmpMIClobberFP.contains(MI);
+  }
+  void insertLongJmpMIClobberFP(MachineInstr *MI) {
+    LongJmpMIClobberFP.insert(MI);
+  }
 };
 
 } // End llvm namespace
diff --git a/llvm/test/CodeGen/X86/clobber_frame_ptr.ll b/llvm/test/CodeGen/X86/clobber_frame_ptr.ll
index 6209e1a85e9e1..fd8ba7feb9f48 100644
--- a/llvm/test/CodeGen/X86/clobber_frame_ptr.ll
+++ b/llvm/test/CodeGen/X86/clobber_frame_ptr.ll
@@ -157,3 +157,46 @@ define i64 @test3(i64 %a0, i64 %a1, i64 %a2, i64 %a3, i64 %a4, i64 %a5, i64 %a6,
   %x = call cc 11 i64 @hipe2(i64 %a0, i64 %a1, i64 %a2, i64 %a3, i64 %a4, i64 %a5, i64 %a6, i64 %a7)
   ret i64 %x
 }
+
+ at buf = dso_local global [20 x ptr] zeroinitializer, align 16
+
+; longjmp modifies fp, it is expected behavior, wo should not save/restore fp
+; around it.
+define void @test4() {
+; CHECK-LABEL: test4:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    pushq %rbp
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    .cfi_offset %rbp, -16
+; CHECK-NEXT:    movq %rsp, %rbp
+; CHECK-NEXT:    .cfi_def_cfa_register %rbp
+; CHECK-NEXT:    pushq %r15
+; CHECK-NEXT:    pushq %r14
+; CHECK-NEXT:    pushq %r13
+; CHECK-NEXT:    pushq %r12
+; CHECK-NEXT:    pushq %rbx
+; CHECK-NEXT:    andq $-16, %rsp
+; CHECK-NEXT:    subq $16, %rsp
+; CHECK-NEXT:    .cfi_offset %rbx, -56
+; CHECK-NEXT:    .cfi_offset %r12, -48
+; CHECK-NEXT:    .cfi_offset %r13, -40
+; CHECK-NEXT:    .cfi_offset %r14, -32
+; CHECK-NEXT:    .cfi_offset %r15, -24
+; CHECK-NEXT:    pushq %rbp
+; CHECK-NEXT:    pushq %rax
+; CHECK-NEXT:    .cfi_remember_state
+; CHECK-NEXT:    .cfi_escape 0x0f, 0x06, 0x77, 0x08, 0x06, 0x11, 0x10, 0x22 #
+; CHECK-NEXT:    xorl %r13d, %r13d
+; CHECK-NEXT:    callq external at PLT
+; CHECK-NEXT:    addq $8, %rsp
+; CHECK-NEXT:    popq %rbp
+; CHECK-NEXT:    .cfi_restore_state
+; CHECK-NEXT:    movq buf(%rip), %rbp
+; CHECK-NEXT:    movq buf+8(%rip), %rax
+; CHECK-NEXT:    movq buf+16(%rip), %rsp
+; CHECK-NEXT:    jmpq *%rax
+entry:
+  %x = call ghccc i32 @external(i32 0)
+  call void @llvm.eh.sjlj.longjmp(ptr @buf)
+  unreachable
+}