Spill/restore FP/BP around instructions in which they are clobbered

weiguozhi · weiguozhi · commit 619b9e24202d · 2024-02-07T22:23:45.000Z
This patch fixes llvm#17204. If a base pointer is used in a function, and it is clobbered by an instruction (typically an inline asm), current register allocator can't handle this situation, so BP becomes garbage after those instructions. It can also occur to FP in theory. We can spill and reload FP/BP registers around those instructions. But normal spill/reload instructions also use FP/BP, so we can't spill them into normal spill slots, instead we spill them into the top of stack by using SP register.
diff --git a/llvm/include/llvm/CodeGen/MachineFrameInfo.h b/llvm/include/llvm/CodeGen/MachineFrameInfo.h
@@ -599,6 +599,7 @@ class MachineFrameInfo {
   /// Return the alignment in bytes that this function must be aligned to,
   /// which is greater than the default stack alignment provided by the target.
   Align getMaxAlign() const { return MaxAlignment; }
+  Align getStackAlignment() const { return StackAlignment; }
 
   /// Make sure the function is at least Align bytes aligned.
   void ensureMaxAlignment(Align Alignment);
diff --git a/llvm/include/llvm/CodeGen/TargetFrameLowering.h b/llvm/include/llvm/CodeGen/TargetFrameLowering.h
@@ -466,6 +466,31 @@ class TargetFrameLowering {
   /// Return the frame base information to be encoded in the DWARF subprogram
   /// debug info.
   virtual DwarfFrameBase getDwarfFrameBase(const MachineFunction &MF) const;
+
+  /// Issue instructions to allocate stack space and spill FP and/or BP to stack
+  /// using SP register.
+  virtual void spillFPBPUsingSP(MachineFunction &MF,
+                                const MachineBasicBlock::iterator BeforeMI,
+                                bool SpillFP, bool SpillBP) const {
+    report_fatal_error("spillFPBPUsingSP is not implemented for this target!");
+  }
+
+  /// Issue instructions to restore FP and/or BP from stack using SP register,
+  /// and free stack space.
+  virtual void restoreFPBPUsingSP(MachineFunction &MF,
+                                  const MachineBasicBlock::iterator AfterMI,
+                                  bool SpillFP, bool SpillBP) const {
+    report_fatal_error("restoreFPBPUsingSP is not implemented for this "
+                       "target!");
+  }
+
+  // If MI uses fp/bp, but target can handle it, and doesn't want PEI to spill
+  // it, this function should return true, and update MI so PEI will not check
+  // any instructions from it and later.
+  virtual bool skipSpillFPBP(MachineFunction &MF,
+                             MachineBasicBlock::reverse_iterator &MI) const {
+    return false;
+  }
 };
 
 } // End llvm namespace
diff --git a/llvm/include/llvm/CodeGen/TargetRegisterInfo.h b/llvm/include/llvm/CodeGen/TargetRegisterInfo.h
@@ -963,6 +963,16 @@ class TargetRegisterInfo : public MCRegisterInfo {
     return false;
   }
 
+  /// Returns true if the target needs a base register.
+  virtual bool hasBaseRegister(const MachineFunction &MF) const {
+    return false;
+  }
+
+  /// Returns the physical base register used to access stack object.
+  virtual Register getBaseRegister(const MachineFunction &MF) const {
+    return 0;
+  }
+
   /// Return true if target has reserved a spill slot in the stack frame of
   /// the given function for the specified register. e.g. On x86, if the frame
   /// register is required, the first fixed stack object is reserved as its
diff --git a/llvm/lib/CodeGen/PrologEpilogInserter.cpp b/llvm/lib/CodeGen/PrologEpilogInserter.cpp
@@ -77,6 +77,15 @@ using MBBVector = SmallVector<MachineBasicBlock *, 4>;
 STATISTIC(NumLeafFuncWithSpills, "Number of leaf functions with CSRs");
 STATISTIC(NumFuncSeen, "Number of functions seen in PEI");
 
+static cl::opt<bool> SpillClobberedFP(
+    "spill-clobbered-fp",
+    cl::desc("Spill clobbered fp register to stack."),
+    cl::init(false), cl::Hidden);
+
+static cl::opt<bool> SpillClobberedBP(
+    "spill-clobbered-bp",
+    cl::desc("Spill clobbered bp register to stack."),
+    cl::init(true), cl::Hidden);
 
 namespace {
 
@@ -122,6 +131,7 @@ class PEI : public MachineFunctionPass {
   void calculateCallFrameInfo(MachineFunction &MF);
   void calculateSaveRestoreBlocks(MachineFunction &MF);
   void spillCalleeSavedRegs(MachineFunction &MF);
+  void spillFrameBasePointer(MachineFunction &MF);
 
   void calculateFrameObjectOffsets(MachineFunction &MF);
   void replaceFrameIndices(MachineFunction &MF);
@@ -228,6 +238,9 @@ bool PEI::runOnMachineFunction(MachineFunction &MF) {
   FrameIndexVirtualScavenging = TRI->requiresFrameIndexScavenging(MF);
   ORE = &getAnalysis<MachineOptimizationRemarkEmitterPass>().getORE();
 
+  // Handle FP and BP spilling and restoring, for instructions that need it.
+  spillFrameBasePointer(MF);
+
   // Calculate the MaxCallFrameSize and AdjustsStack variables for the
   // function's frame information. Also eliminates call frame pseudo
   // instructions.
@@ -1587,3 +1600,114 @@ void PEI::replaceFrameIndices(MachineBasicBlock *BB, MachineFunction &MF,
       ++I;
   }
 }
+
+static bool accessFrameBasePointer(const MachineInstr &MI, Register FP,
+                                   Register BP, bool &AccessFP, bool &AccessBP,
+                                   const TargetRegisterInfo *TRI) {
+  AccessFP = AccessBP = false;
+  if (FP) {
+    if (MI.findRegisterUseOperandIdx(FP, false, TRI) != -1 ||
+        MI.findRegisterDefOperandIdx(FP, false, true, TRI) != -1)
+      AccessFP = true;
+  }
+  if (BP) {
+    if (MI.findRegisterUseOperandIdx(BP, false, TRI) != -1 ||
+        MI.findRegisterDefOperandIdx(BP, false, true, TRI) != -1)
+      AccessBP = true;
+  }
+  return AccessFP || AccessBP;
+}
+
+/// If a function uses base pointer and the base pointer is clobbered by inline
+/// asm, RA doesn't detect this case, and after the inline asm, the base pointer
+/// cotains garbage value.
+/// For example if a 32b x86 function uses base pointer esi, and esi is
+/// clobbered by following inline asm
+///     asm("rep movsb" : "+D"(ptr), "+S"(x), "+c"(c)::"memory");
+/// We need to save esi before the asm and restore it after the asm.
+///
+/// The problem can also occur to frame pointer if there is a function call, and
+/// the callee uses a different calling convention and clobbers the fp.
+///
+/// Because normal frame objects (spill slots) are accessed through fp/bp
+/// register, so we can't spill fp/bp to normal spill slots.
+///
+/// FIXME: There are 2 possible enhancements:
+/// 1. In many cases there are different physical registers not clobbered by
+/// inline asm, we can use one of them as base pointer. Or use a virtual
+/// register as base pointer and let RA allocate a physical register to it.
+/// 2. If there is no other instructions access stack with fp/bp from the
+/// inline asm to the epilog, we can skip the save and restore operations.
+void PEI::spillFrameBasePointer(MachineFunction &MF) {
+  Register FP, BP;
+  const TargetFrameLowering &TFI = *MF.getSubtarget().getFrameLowering();
+  const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
+  if (TFI.hasFP(MF) && SpillClobberedFP)
+    FP = TRI->getFrameRegister(MF);
+  if (TRI->hasBaseRegister(MF) && SpillClobberedBP)
+    BP = TRI->getBaseRegister(MF);
+
+  for (MachineBasicBlock &MBB : MF) {
+    auto MI = MBB.rbegin(), ME = MBB.rend();
+    while (MI != ME) {
+      // Skip frame setup/destroy instructions.
+      // Skip instructions handled by target.
+      if (MI->getFlag(MachineInstr::MIFlag::FrameSetup) ||
+          MI->getFlag(MachineInstr::MIFlag::FrameDestroy) ||
+          TFI.skipSpillFPBP(MF, MI)) {
+        ++MI;
+        continue;
+      }
+
+      bool AccessFP, AccessBP;
+      // Check if fp or bp is used in MI.
+      if (!accessFrameBasePointer(*MI, FP, BP, AccessFP, AccessBP, TRI)) {
+        ++MI;
+        continue;
+      }
+
+      // Look for the range [Start, Stop] in which fp or bp is defined and used.
+      bool FPLive = false, BPLive = false;
+      bool SpillFP = false, SpillBP = false;
+      auto Start = MI, Stop = MI;
+      do {
+        SpillFP |= AccessFP;
+        SpillBP |= AccessBP;
+
+        // Maintain FPLive and BPLive.
+        if (FPLive && MI->findRegisterDefOperandIdx(FP, false, true, TRI) != -1)
+          FPLive = false;
+        if (FP && MI->findRegisterUseOperandIdx(FP, false, TRI) != -1)
+          FPLive = true;
+        if (BPLive && MI->findRegisterDefOperandIdx(BP, false, true, TRI) != -1)
+          BPLive = false;
+        if (BP && MI->findRegisterUseOperandIdx(BP, false, TRI) != -1)
+          BPLive = true;
+
+        Start = MI++;
+      } while ((MI != ME) && (FPLive || BPLive ||
+                accessFrameBasePointer(*MI, FP, BP, AccessFP, AccessBP, TRI)));
+
+      // If the bp is clobbered by a call, we should save and restore outside of
+      // the frame setup instructions.
+      if (Stop->isCall()) {
+        const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
+        auto FrameSetup = std::next(Start);
+        while (FrameSetup != ME && !TII.isFrameSetup(*FrameSetup) &&
+               !FrameSetup->isCall())
+          ++FrameSetup;
+        if (FrameSetup != ME && TII.isFrameSetup(*FrameSetup)) {
+          while (!TII.isFrameInstr(*Stop))
+            --Stop;
+          Start = FrameSetup;
+          MI = Start;
+          ++MI;
+        }
+      }
+
+      // Call target functions to spill and restore FP and BP registers.
+      TFI.spillFPBPUsingSP(MF, &(*Start), SpillFP, SpillBP);
+      TFI.restoreFPBPUsingSP(MF, &(*Stop), SpillFP, SpillBP);
+    }
+  }
+}
diff --git a/llvm/lib/Target/X86/X86FrameLowering.cpp b/llvm/lib/Target/X86/X86FrameLowering.cpp
@@ -4178,3 +4178,114 @@ void X86FrameLowering::restoreWinEHStackPointersInParent(
                                   /*RestoreSP=*/IsSEH);
   }
 }
+
+static int computeSPAdjust4SpillFPBP(MachineFunction &MF,
+                                     const TargetRegisterClass *RC,
+                                     unsigned SpillRegNum) {
+  const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
+  unsigned AllocSize = TRI->getSpillSize(*RC) * SpillRegNum;
+  Align StackAlign = MF.getFrameInfo().getStackAlignment();
+  unsigned AlignedSize = alignTo(AllocSize, StackAlign);
+  return AlignedSize - AllocSize;
+}
+
+void X86FrameLowering::spillFPBPUsingSP(
+    MachineFunction &MF, MachineBasicBlock::iterator BeforeMI,
+    bool SpillFP, bool SpillBP) const {
+  const TargetRegisterClass *RC;
+  unsigned RegNum = 0;
+  MachineBasicBlock *MBB = BeforeMI->getParent();
+  DebugLoc DL = BeforeMI->getDebugLoc();
+
+  // Spill FP.
+  if (SpillFP) {
+    Register FP = TRI->getFrameRegister(MF);
+    if (STI.isTarget64BitILP32())
+      FP = Register(getX86SubSuperRegister(FP, 64));
+    RC = TRI->getMinimalPhysRegClass(FP);
+    ++RegNum;
+
+    BuildMI(*MBB, BeforeMI, DL,
+            TII.get(getPUSHOpcode(MF.getSubtarget<X86Subtarget>())))
+        .addReg(FP);
+  }
+
+  // Spill BP.
+  if (SpillBP) {
+    Register BP = TRI->getBaseRegister(MF);
+    if (STI.isTarget64BitILP32())
+      BP = Register(getX86SubSuperRegister(BP, 64));
+    RC = TRI->getMinimalPhysRegClass(BP);
+    ++RegNum;
+
+    BuildMI(*MBB, BeforeMI, DL,
+            TII.get(getPUSHOpcode(MF.getSubtarget<X86Subtarget>())))
+        .addReg(BP);
+  }
+
+  // Make sure SP is aligned.
+  int SPAdjust = computeSPAdjust4SpillFPBP(MF, RC, RegNum);
+  if (SPAdjust)
+    emitSPUpdate(*MBB, BeforeMI, DL, -SPAdjust, false);
+}
+
+void X86FrameLowering::restoreFPBPUsingSP(
+    MachineFunction &MF, MachineBasicBlock::iterator AfterMI,
+    bool SpillFP, bool SpillBP) const {
+  Register FP, BP;
+  const TargetRegisterClass *RC;
+  unsigned RegNum = 0;
+  if (SpillFP) {
+    FP = TRI->getFrameRegister(MF);
+    if (STI.isTarget64BitILP32())
+      FP = Register(getX86SubSuperRegister(FP, 64));
+    RC = TRI->getMinimalPhysRegClass(FP);
+    ++RegNum;
+  }
+  if (SpillBP) {
+    BP = TRI->getBaseRegister(MF);
+    if (STI.isTarget64BitILP32())
+      BP = Register(getX86SubSuperRegister(BP, 64));
+    RC = TRI->getMinimalPhysRegClass(BP);
+    ++RegNum;
+  }
+
+  // Adjust SP so it points to spilled FP or BP.
+  MachineBasicBlock *MBB = AfterMI->getParent();
+  MachineBasicBlock::iterator Pos = std::next(AfterMI);
+  DebugLoc DL = AfterMI->getDebugLoc();
+  int SPAdjust = computeSPAdjust4SpillFPBP(MF, RC, RegNum);
+  if (SPAdjust)
+    emitSPUpdate(*MBB, Pos, DL, SPAdjust, false);
+
+  // Restore BP.
+  if (SpillBP) {
+    BuildMI(*MBB, Pos, DL,
+            TII.get(getPOPOpcode(MF.getSubtarget<X86Subtarget>())), BP);
+  }
+
+  // Restore FP.
+  if (SpillFP) {
+    BuildMI(*MBB, Pos, DL,
+            TII.get(getPOPOpcode(MF.getSubtarget<X86Subtarget>())), FP);
+  }
+}
+
+bool X86FrameLowering::skipSpillFPBP(
+    MachineFunction &MF, MachineBasicBlock::reverse_iterator &MI) const {
+  if (MI->getOpcode() == X86::LCMPXCHG16B_SAVE_RBX) {
+    // The pseudo instruction LCMPXCHG16B_SAVE_RBX is generated in the form
+    //     SaveRbx = COPY RBX
+    //     SaveRbx = LCMPXCHG16B_SAVE_RBX ..., SaveRbx, implicit-def rbx
+    // And later LCMPXCHG16B_SAVE_RBX is expanded to restore RBX from SaveRbx.
+    // We should skip this instruction sequence.
+    int FI;
+    unsigned Reg;
+    while (!(MI->getOpcode() == TargetOpcode::COPY &&
+             MI->getOperand(1).getReg() == X86::RBX) &&
+           !((Reg = TII.isStoreToStackSlot(*MI, FI)) && Reg == X86::RBX))
+      ++MI;
+    return true;
+  }
+  return false;
+}
diff --git a/llvm/lib/Target/X86/X86FrameLowering.h b/llvm/lib/Target/X86/X86FrameLowering.h
@@ -103,6 +103,20 @@ class X86FrameLowering : public TargetFrameLowering {
                               MutableArrayRef<CalleeSavedInfo> CSI,
                               const TargetRegisterInfo *TRI) const override;
 
+  void
+  spillFPBPUsingSP(MachineFunction &MF,
+                   const MachineBasicBlock::iterator BeforeMI,
+                   bool SpillFP, bool SpillBP) const override;
+
+  void
+  restoreFPBPUsingSP(MachineFunction &MF,
+                     const MachineBasicBlock::iterator AfterMI,
+                     bool SpillFP, bool SpillBP) const override;
+
+  bool
+  skipSpillFPBP(MachineFunction &MF,
+                MachineBasicBlock::reverse_iterator &MI) const override;
+
   bool hasFP(const MachineFunction &MF) const override;
   bool hasReservedCallFrame(const MachineFunction &MF) const override;
   bool canSimplifyCallFramePseudos(const MachineFunction &MF) const override;
diff --git a/llvm/lib/Target/X86/X86RegisterInfo.h b/llvm/lib/Target/X86/X86RegisterInfo.h
@@ -133,6 +133,9 @@ class X86RegisterInfo final : public X86GenRegisterInfo {
   void adjustStackMapLiveOutMask(uint32_t *Mask) const override;
 
   bool hasBasePointer(const MachineFunction &MF) const;
+  bool hasBaseRegister(const MachineFunction &MF) const override {
+    return hasBasePointer(MF);
+  }
 
   bool canRealignStack(const MachineFunction &MF) const override;
 
@@ -164,6 +167,9 @@ class X86RegisterInfo final : public X86GenRegisterInfo {
   unsigned getPtrSizedStackRegister(const MachineFunction &MF) const;
   Register getStackRegister() const { return StackPtr; }
   Register getBaseRegister() const { return BasePtr; }
+  Register getBaseRegister(const MachineFunction &MF) const override {
+    return BasePtr;
+  }
   /// Returns physical register used as frame pointer.
   /// This will always returns the frame pointer register, contrary to
   /// getFrameRegister() which returns the "base pointer" in situations
diff --git a/llvm/test/CodeGen/X86/clobber_base_ptr.ll b/llvm/test/CodeGen/X86/clobber_base_ptr.ll
diff --git a/llvm/test/CodeGen/X86/i386-baseptr.ll b/llvm/test/CodeGen/X86/i386-baseptr.ll
diff --git a/llvm/test/CodeGen/X86/x86-64-baseptr.ll b/llvm/test/CodeGen/X86/x86-64-baseptr.ll
