[ipra] Fix missing save/restore of function return addresses

[cme]

Summary

This patch adds the architecture-neutral components of a fix for issue #57482 ( https://github.com/llvm/llvm-project/issues/57482 ), and the artitecture-specific parts of the fix for MIPS processors.

IPRA allows the save/restore sequences to be omitted from callees whose clobber sets have been fully accounted for by the calling function due to IPRA. However, since the return address to return from the callee is used in the *return* instruction of the calling instruction rather than in the caller, having the caller preserve a value is inadequate for return addresses.

No existing register property available in the framework encapsulates this "needed for return instruction" property, so this changed adds this as a TargerRegisterInfo query method.

[mehdi_amini]

I haven't touched this part of the project for multiple years, I'm not sure I'll be the best person to review.

[cme]

In D133301#3770280, @mehdi_amini wrote:

I haven't touched this part of the project for multiple years, I'm not sure I'll be the best person to review.

Okay, thanks. Hopefully @arsenm is better placed since they're using IPRA by default in AMDGPU

[cme]

In D133301#3770280, @mehdi_amini wrote:

I haven't touched this part of the project for multiple years, I'm not sure I'll be the best person to review.

Or, would you be able to suggest anyone else that might be a suitable reviewer? Thanks.

[arsenm]

For AMDGPU we special cased the RA handling in determineCalleeSaves

llvm/lib/CodeGen/TargetFrameLoweringImpl.cpp

130

It would be better to handle the list of special case registers rather than iterating through every CSR

[arsenm]

Should at least avoid iterating every register for this

llvm/include/llvm/CodeGen/TargetRegisterInfo.h

public:
   }
 
   /// Returns true if PhysReg is a general purpose register.
   virtual bool isGeneralPurposeRegister(const MachineFunction &MF,
                                         MCRegister PhysReg) const {
     return false;
   }
 
+  // Return true if the register is needed for returning from the
+  // function and so must be preserved in the callee even if preserved
+  // by the caller
+  virtual bool isNeededForReturn(MCRegister PhysReg,
+                                 const MachineFunction &MF) const {
+    return false;
+  }
+
   /// Prior to adding the live-out mask to a stackmap or patchpoint
   /// instruction, provide the target the opportunity to adjust it (mainly to
   /// remove pseudo-registers that should be ignored).
   virtual void adjustStackMapLiveOutMask(uint32_t *Mask) const {}
 
   /// Return a super-register of the specified register
   /// Reg so its sub-register of index SubIdx is Reg.
   MCRegister getMatchingSuperReg(MCRegister Reg, unsigned SubIdx,

llvm/lib/CodeGen/TargetFrameLoweringImpl.cpp

void TargetFrameLowering::determineCalleeSaves(MachineFunction &MF,
 
   // Resize before the early returns. Some backends expect that
   // SavedRegs.size() == TRI.getNumRegs() after this call even if there are no
   // saved registers.
   SavedRegs.resize(TRI.getNumRegs());
 
   // When interprocedural register allocation is enabled caller saved registers
   // are preferred over callee saved registers.
-  if (MF.getTarget().Options.EnableIPRA &&
+  bool NoCSR = (MF.getTarget().Options.EnableIPRA &&
                 isSafeForNoCSROpt(MF.getFunction()) &&
-      isProfitableForNoCSROpt(MF.getFunction()))
-    return;
+                isProfitableForNoCSROpt(MF.getFunction()));
 
   // Get the callee saved register list...
   const MCPhysReg *CSRegs = MF.getRegInfo().getCalleeSavedRegs();
 
   // Early exit if there are no callee saved registers.
   if (!CSRegs || CSRegs[0] == 0)
     return;
 
if (MF.getFunction().hasFnAttribute(Attribute::NoReturn) &&
         MF.getFunction().hasFnAttribute(Attribute::NoUnwind) &&
         !MF.getFunction().hasFnAttribute(Attribute::UWTable) &&
         enableCalleeSaveSkip(MF))
     return;
 
   // Functions which call __builtin_unwind_init get all their registers saved.
   bool CallsUnwindInit = MF.callsUnwindInit();
   const MachineRegisterInfo &MRI = MF.getRegInfo();
+  const TargetRegisterInfo *RI = MF.getSubtarget().getRegisterInfo();
   for (unsigned i = 0; CSRegs[i]; ++i) {
     unsigned Reg = CSRegs[i];
     if (CallsUnwindInit || MRI.isPhysRegModified(Reg))
+      if (!NoCSR || !MRI.isAllocatable(Reg) || RI->isNeededForReturn(Reg, MF))
         SavedRegs.set(Reg);
   }
 }
 
 unsigned TargetFrameLowering::getStackAlignmentSkew(
     const MachineFunction &MF) const {
   // When HHVM function is called, the stack is skewed as the return address
   // is removed from the stack before we enter the function.
   if (LLVM_UNLIKELY(MF.getFunction().getCallingConv() == CallingConv::HHVM))

llvm/lib/Target/Mips/MipsRegisterInfo.h

public:
   bool canRealignStack(const MachineFunction &MF) const override;
 
   /// Debug information queries.
   Register getFrameRegister(const MachineFunction &MF) const override;
 
   /// Return GPR register class.
   virtual const TargetRegisterClass *intRegClass(unsigned Size) const = 0;
 
+  virtual bool isNeededForReturn(MCRegister PhysReg,
+                                 const MachineFunction &MF) const override;
+
 private:
   virtual void eliminateFI(MachineBasicBlock::iterator II, unsigned OpNo,
                            int FrameIndex, uint64_t StackSize,
                            int64_t SPOffset) const = 0;
 };
 
 } // end namespace llvm
 
 #endif // LLVM_LIB_TARGET_MIPS_MIPSREGISTERINFO_H

llvm/lib/Target/Mips/MipsRegisterInfo.cpp

getReservedRegs(const MachineFunction &MF) const {
   return Reserved;
 }
 
 bool
 MipsRegisterInfo::requiresRegisterScavenging(const MachineFunction &MF) const {
   return true;
 }
 
+bool MipsRegisterInfo::isNeededForReturn(MCRegister PhysReg,
+                                         const MachineFunction &MF) const {
+  return PhysReg == Mips::RA || PhysReg == Mips::RA_64;
+}
+
 // FrameIndex represent objects inside a abstract stack.
 // We must replace FrameIndex with an stack/frame pointer
 // direct reference.
 void MipsRegisterInfo::
 eliminateFrameIndex(MachineBasicBlock::iterator II, int SPAdj,
                     unsigned FIOperandNum, RegScavenger *RS) const {
   MachineInstr &MI = *II;
   MachineFunction &MF = *MI.getParent()->getParent();