[X86] Fix the base pointer save/restore bug

[LuoYuanke]

Summary

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.

[LuoYuanke]

llvm/test/CodeGen/X86/sjlj-baseptr.ll

27

The offset happen to be the same.

[LuoYuanke]

update test case.

[ArchDRobison]

It's been a decade since I've worked with LLVM. The change looks okay to me.

[LuoYuanke]

In D144625#4150401, @ArchDRobison wrote:

It's been a decade since I've worked with LLVM. The change looks okay to me.

Thanks for review. Would you accept the patch, so that I can land it?

[pengfei]

llvm/lib/Target/X86/X86FrameLowering.cpp

2737

Should add const here rather than remove from line 2703?

2740

Why need this->?

2844–2845

ditto.

2848

ditto.

llvm/lib/Target/X86/X86MachineFunctionInfo.cpp

38–41

This can be moved to X86MachineFunctionInfo.h

llvm/test/CodeGen/X86/sjlj-baseptr.ll

25

No quite understand why do we need push twice. Is it because the alignment of alloca is 16? How about align to 32?

I found previously only sjlj code calls setRestoreBasePointer, but the change here seems more general, should we add a general test for it?

[LuoYuanke]

llvm/lib/Target/X86/X86FrameLowering.cpp

2737

Yes, it looks better.

2740

There is "const TargetRegisterInfo *TRI" passed from argument. "this->TRI" is X86TRI.

llvm/test/CodeGen/X86/sjlj-baseptr.ll

25

It is not for alignment. It is used to allocate the spill stack slot for base pointer register.
There are 2 more test case in https://reviews.llvm.org/D144541, this patch focuses on the bug fix.

[LuoYuanke]

Address Phoebe's comments.

[pengfei]

LGTM.

llvm/lib/Target/X86/X86MachineFunctionInfo.cpp

37

Recover the change here.

llvm/lib/Target/X86/X86FrameLowering.cpp

if (IsWin64Prologue && IsFunclet && !IsClrFunclet) {
     MBB.addLiveIn(Establisher);
   }
 
   if (HasFP) {
     assert(MF.getRegInfo().isReserved(MachineFramePtr) && "FP reserved");
 
     // 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);
 
     // Callee-saved registers are pushed on stack before the stack is realigned.
     if (TRI->hasStackRealignment(MF) && !IsWin64Prologue)
       NumBytes = alignTo(NumBytes, MaxAlign);
 
     // Save EBP/RBP into the appropriate stack slot.
for (CalleeSavedInfo &I : llvm::reverse(CSI)) {
 
     SpillSlotOffset -= SlotSize;
     CalleeSavedFrameSize += SlotSize;
 
     int SlotIndex = MFI.CreateFixedSpillStackObject(SlotSize, SpillSlotOffset);
     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);
 
   // Assign slots for XMMs.
   for (CalleeSavedInfo &I : llvm::reverse(CSI)) {
     Register Reg = I.getReg();
     if (X86::GR64RegClass.contains(Reg) || X86::GR32RegClass.contains(Reg))
       continue;
bool X86FrameLowering::spillCalleeSavedRegisters(
   DebugLoc DL = MBB.findDebugLoc(MI);
 
   // Don't save CSRs in 32-bit EH funclets. The caller saves EBX, EBP, ESI, EDI
   // for us, and there are no XMM CSRs on Win32.
   if (MBB.isEHFuncletEntry() && STI.is32Bit() && STI.isOSWindows())
     return true;
 
   // Push GPRs. It increases frame size.
-  const MachineFunction &MF = *MBB.getParent();
+  MachineFunction &MF = *MBB.getParent();
   unsigned Opc = STI.is64Bit() ? X86::PUSH64r : X86::PUSH32r;
   for (const CalleeSavedInfo &I : llvm::reverse(CSI)) {
     Register Reg = I.getReg();
 
     if (!X86::GR64RegClass.contains(Reg) && !X86::GR32RegClass.contains(Reg))
       continue;
 
     const MachineRegisterInfo &MRI = MF.getRegInfo();
for (const CalleeSavedInfo &I : llvm::reverse(CSI)) {
     // happens with the @llvm-returnaddress intrinsic and with arguments
     // passed in callee saved registers.
     // Omitting the kill flags is conservatively correct even if the live-in
     // is not used after all.
     BuildMI(MBB, MI, DL, TII.get(Opc)).addReg(Reg, getKillRegState(CanKill))
       .setMIFlag(MachineInstr::FrameSetup);
   }
 
+  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)) {
     Register Reg = I.getReg();
     if (X86::GR64RegClass.contains(Reg) || X86::GR32RegClass.contains(Reg))
       continue;
 
     // If this is k-register make sure we lookup via the largest legal type.
for (const CalleeSavedInfo &I : CSI) {
     if (X86::VK16RegClass.contains(Reg))
       VT = STI.hasBWI() ? MVT::v64i1 : MVT::v16i1;
 
     const TargetRegisterClass *RC = TRI->getMinimalPhysRegClass(Reg, VT);
     TII.loadRegFromStackSlot(MBB, MI, Reg, I.getFrameIdx(), RC, TRI,
                              Register());
   }
 
+  // Clear the stack slot for spill base pointer register.
+  MachineFunction &MF = *MBB.getParent();
+  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) {
     Register Reg = I.getReg();
     if (!X86::GR64RegClass.contains(Reg) &&
         !X86::GR32RegClass.contains(Reg))
       continue;
 

llvm/lib/Target/X86/X86MachineFunctionInfo.h

public:
   bool getForceFramePointer() const { return ForceFramePointer;}
   void setForceFramePointer(bool forceFP) { ForceFramePointer = forceFP; }
 
   bool getHasPushSequences() const { return HasPushSequences; }
   void setHasPushSequences(bool HasPush) { HasPushSequences = HasPush; }
 
   bool getRestoreBasePointer() const { return RestoreBasePointerOffset!=0; }
   void setRestoreBasePointer(const MachineFunction *MF);
+  void setRestoreBasePointer(unsigned CalleeSavedFrameSize);
   int getRestoreBasePointerOffset() const {return RestoreBasePointerOffset; }
 
   DenseMap<int, unsigned>& getWinEHXMMSlotInfo() { return WinEHXMMSlotInfo; }
   const DenseMap<int, unsigned>& getWinEHXMMSlotInfo() const {
     return WinEHXMMSlotInfo; }
 
   unsigned getCalleeSavedFrameSize() const { return CalleeSavedFrameSize; }
   void setCalleeSavedFrameSize(unsigned bytes) { CalleeSavedFrameSize = bytes; }

llvm/lib/Target/X86/X86MachineFunctionInfo.cpp

if (!RestoreBasePointerOffset) {
     unsigned SlotSize = RegInfo->getSlotSize();
     for (const MCPhysReg *CSR = MF->getRegInfo().getCalleeSavedRegs();
          unsigned Reg = *CSR; ++CSR) {
       if (X86::GR64RegClass.contains(Reg) || X86::GR32RegClass.contains(Reg))
         RestoreBasePointerOffset -= SlotSize;
     }
   }
 }
 
+void X86MachineFunctionInfo::setRestoreBasePointer(
+    unsigned CalleeSavedFrameSize) {
+  RestoreBasePointerOffset = -CalleeSavedFrameSize;
+}

llvm/test/CodeGen/X86/sjlj-baseptr.ll

entry:
   %buf = alloca [5 x ptr], align 16
   %p = alloca ptr, align 8
   %q = alloca i8, align 64
   %s = alloca i8, i64 %n, align 1
   store ptr %s, ptr %p, align 8
   %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
 }