[llvm] ba8a520 - [X86] Fix the base pointer save/restore bug

[via llvm-commits]

Author: Luo, Yuanke
Date: 2023-02-25T20:52:13+08:00
New Revision: ba8a520512b7c36a9958ddbf9be96a33ed0279b8

URL: https://github.com/llvm/llvm-project/commit/ba8a520512b7c36a9958ddbf9be96a33ed0279b8
DIFF: https://github.com/llvm/llvm-project/commit/ba8a520512b7c36a9958ddbf9be96a33ed0279b8.diff

LOG: [X86] Fix the base pointer save/restore bug

Previous the stack slot for spilling base pointer register is allocated
after the stack realignment. When the stack is naturally aligned the
stack slot is share with other data that allocated from stack and cause
data corrupt. Another issue is the stack slot for save/restore the
callee saved register is not fixed for each function. It depends on the
register usage of them in each function.

This patch is to recalculate the offset the stack slot for base pointer
register during the prolog/epilog insert pass, and allocate the stack
slot when spilling callee saved registers.

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

Added: 
    

Modified: 
    llvm/lib/Target/X86/X86FrameLowering.cpp
    llvm/lib/Target/X86/X86MachineFunctionInfo.h
    llvm/test/CodeGen/X86/sjlj-baseptr.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/X86/X86FrameLowering.cpp b/llvm/lib/Target/X86/X86FrameLowering.cpp
index de33699810a05..1f827587d1151 100644
--- a/llvm/lib/Target/X86/X86FrameLowering.cpp
+++ b/llvm/lib/Target/X86/X86FrameLowering.cpp
@@ -1628,10 +1628,6 @@ void X86FrameLowering::emitPrologue(MachineFunction &MF,
 
     // Calculate required stack adjustment.
     uint64_t FrameSize = StackSize - SlotSize;
-    // If required, include space for extra hidden slot for stashing base pointer.
-    if (X86FI->getRestoreBasePointer())
-      FrameSize += SlotSize;
-
     NumBytes = FrameSize -
                (X86FI->getCalleeSavedFrameSize() + TailCallArgReserveSize);
 
@@ -2646,6 +2642,16 @@ bool X86FrameLowering::assignCalleeSavedSpillSlots(
     I.setFrameIdx(SlotIndex);
   }
 
+  // Adjust the offset of spill slot as we know the accurate callee saved frame
+  // size.
+  if (X86FI->getRestoreBasePointer()) {
+    SpillSlotOffset -= SlotSize;
+    CalleeSavedFrameSize += SlotSize;
+
+    MFI.CreateFixedSpillStackObject(SlotSize, SpillSlotOffset);
+    // TODO: saving the slot index is better?
+    X86FI->setRestoreBasePointer(CalleeSavedFrameSize);
+  }
   X86FI->setCalleeSavedFrameSize(CalleeSavedFrameSize);
   MFI.setCVBytesOfCalleeSavedRegisters(CalleeSavedFrameSize);
 
@@ -2728,6 +2734,15 @@ bool X86FrameLowering::spillCalleeSavedRegisters(
       .setMIFlag(MachineInstr::FrameSetup);
   }
 
+  const X86MachineFunctionInfo *X86FI = MF.getInfo<X86MachineFunctionInfo>();
+  if (X86FI->getRestoreBasePointer()) {
+    unsigned Opc = STI.is64Bit() ? X86::PUSH64r : X86::PUSH32r;
+    Register BaseReg = this->TRI->getBaseRegister();
+    BuildMI(MBB, MI, DL, TII.get(Opc))
+        .addReg(BaseReg, getKillRegState(true))
+        .setMIFlag(MachineInstr::FrameSetup);
+  }
+
   // Make XMM regs spilled. X86 does not have ability of push/pop XMM.
   // It can be done by spilling XMMs to stack frame.
   for (const CalleeSavedInfo &I : llvm::reverse(CSI)) {
@@ -2825,6 +2840,16 @@ bool X86FrameLowering::restoreCalleeSavedRegisters(
                              Register());
   }
 
+  // Clear the stack slot for spill base pointer register.
+  MachineFunction &MF = *MBB.getParent();
+  const X86MachineFunctionInfo *X86FI = MF.getInfo<X86MachineFunctionInfo>();
+  if (X86FI->getRestoreBasePointer()) {
+    unsigned Opc = STI.is64Bit() ? X86::POP64r : X86::POP32r;
+    Register BaseReg = this->TRI->getBaseRegister();
+    BuildMI(MBB, MI, DL, TII.get(Opc), BaseReg)
+        .setMIFlag(MachineInstr::FrameDestroy);
+  }
+
   // POP GPRs.
   unsigned Opc = STI.is64Bit() ? X86::POP64r : X86::POP32r;
   for (const CalleeSavedInfo &I : CSI) {

diff  --git a/llvm/lib/Target/X86/X86MachineFunctionInfo.h b/llvm/lib/Target/X86/X86MachineFunctionInfo.h
index 6d7cf165699b4..372838a212cc8 100644
--- a/llvm/lib/Target/X86/X86MachineFunctionInfo.h
+++ b/llvm/lib/Target/X86/X86MachineFunctionInfo.h
@@ -149,6 +149,9 @@ class X86MachineFunctionInfo : public MachineFunctionInfo {
 
   bool getRestoreBasePointer() const { return RestoreBasePointerOffset!=0; }
   void setRestoreBasePointer(const MachineFunction *MF);
+  void setRestoreBasePointer(unsigned CalleeSavedFrameSize) {
+    RestoreBasePointerOffset = -CalleeSavedFrameSize;
+  }
   int getRestoreBasePointerOffset() const {return RestoreBasePointerOffset; }
 
   DenseMap<int, unsigned>& getWinEHXMMSlotInfo() { return WinEHXMMSlotInfo; }

diff  --git a/llvm/test/CodeGen/X86/sjlj-baseptr.ll b/llvm/test/CodeGen/X86/sjlj-baseptr.ll
index f7daa13d47752..e414b08d3e232 100644
--- a/llvm/test/CodeGen/X86/sjlj-baseptr.ll
+++ b/llvm/test/CodeGen/X86/sjlj-baseptr.ll
@@ -21,16 +21,24 @@ entry:
   %t = call i32 @llvm.eh.sjlj.setjmp(ptr %s)
   call void @whatever(i64 %n, ptr %f, ptr %p, ptr %q, ptr %s, i32 %t) #1
   ret i32 0
+; X86: pushl   %esi
+; X86-NEXT: pushl   %esi
 ; X86: movl    %esp, %esi
 ; X86: movl    %esp, -16(%ebp)
 ; X86: {{.LBB.*:}}
 ; X86: movl    -16(%ebp), %esi
 ; X86: {{.LBB.*:}}
+; X86: popl    %esi
+; X86-NEXT: popl    %esi
+; X64: pushq   %rbx
+; X64-NEXT: pushq   %rbx
 ; X64: movq    %rsp, %rbx
 ; X64: movq    %rsp, -48(%rbp)
 ; X64: {{.LBB.*:}}
 ; X64: movq    -48(%rbp), %rbx
 ; X64: {{.LBB.*:}}
+; X64: popq    %rbx
+; X64-NEXT: popq    %rbx
 }