[SystemZ] [z/OS] Use assignCalleeSavedSpillSlots() to mark handle special registers in CSR list instead of determineCalleeSaves

[Everybody0523]

Summary

This PR moves the handling of special registers that need to be saved/restored in the prolog/epilog respectively from determineCalleeSaves to assignCalleeSavedSpillSlots. The documentation of the parent function of assignCalleeSavedSpillSlots explicitly allows the modification of the CSI hence adding the special registers (the stack pointer register, the return address register, and the entry point register) to the CSI list at that stage should be permissible.

This cleans up the code a bit and makes it so that we do not have to place registers that are not actually considered CSRs by the spec in the CSR list, which is something of a hack.

[uweigand]

LGTM, thanks!

[Everybody0523]

rebase to latest main

llvm/lib/Target/SystemZ/SystemZCallingConv.td

 def CSR_SystemZ_AllRegs_Vector : CalleeSavedRegs<(add (sequence "R%dD", 2, 15),
                                                       (sequence "V%d", 0, 31))>;
 
 def CSR_SystemZ_NoRegs : CalleeSavedRegs<(add)>;
 
 //===----------------------------------------------------------------------===//
 // z/OS XPLINK64 callee-saved registers
 //===----------------------------------------------------------------------===//
-// %R7D is volatile by the spec, but it must be saved in the prologue by
-// any non-leaf function and restored in the epilogue for use by the
-// return instruction so it functions exactly like a callee-saved register.
-def CSR_SystemZ_XPLINK64 : CalleeSavedRegs<(add (sequence "R%dD", 7, 15),
-                                                (sequence "R%dD", 4, 4),
+def CSR_SystemZ_XPLINK64 : CalleeSavedRegs<(add (sequence "R%dD", 8, 15),
                                                 (sequence "F%dD", 15, 8))>;
 
 def CSR_SystemZ_XPLINK64_Vector : CalleeSavedRegs<(add CSR_SystemZ_XPLINK64,
                                                    (sequence "V%d", 23, 16))>;
 
 //===----------------------------------------------------------------------===//
 // z/OS XPLINK64 return value calling convention
 //===----------------------------------------------------------------------===//

llvm/lib/Target/SystemZ/SystemZFrameLowering.cpp

 
 bool SystemZXPLINKFrameLowering::assignCalleeSavedSpillSlots(
     MachineFunction &MF, const TargetRegisterInfo *TRI,
     std::vector<CalleeSavedInfo> &CSI) const {
   MachineFrameInfo &MFFrame = MF.getFrameInfo();
   SystemZMachineFunctionInfo *MFI = MF.getInfo<SystemZMachineFunctionInfo>();
   const SystemZSubtarget &Subtarget = MF.getSubtarget<SystemZSubtarget>();
   auto &Regs = Subtarget.getSpecialRegisters<SystemZXPLINK64Registers>();
+  auto &GRRegClass = SystemZ::GR64BitRegClass;
+
+  // For non-leaf functions:
+  // - the address of callee (entry point) register R6 must be saved
+  CSI.push_back(CalleeSavedInfo(Regs.getAddressOfCalleeRegister()));
+  CSI.back().setRestored(false);
+
+  // The return address register R7 must be saved and restored.
+  CSI.push_back(CalleeSavedInfo(Regs.getReturnFunctionAddressRegister()));
+
+  // If the function needs a frame pointer, or if the backchain pointer should
+  // be stored, then save the stack pointer register R4.
+  if (hasFP(MF) || MF.getFunction().hasFnAttribute("backchain"))
+    CSI.push_back(CalleeSavedInfo(Regs.getStackPointerRegister()));
 
   // Scan the call-saved GPRs and find the bounds of the register spill area.
-  unsigned LowGPR = 0;
-  int LowOffset = INT32_MAX;
-  unsigned HighGPR = LowGPR;
+  Register LowRestoreGPR = 0;
+  int LowRestoreOffset = INT32_MAX;
+  Register LowSpillGPR = 0;
+  int LowSpillOffset = INT32_MAX;
+  Register HighGPR = 0;
   int HighOffset = -1;
 
-  unsigned RegSP = Regs.getStackPointerRegister();
-  auto &GRRegClass = SystemZ::GR64BitRegClass;
-  const unsigned RegSize = 8;
-
-  auto ProcessCSI = [&](std::vector<CalleeSavedInfo> &CSIList) {
-    for (auto &CS : CSIList) {
-      Register Reg = CS.getReg();
-      int Offset = RegSpillOffsets[Reg];
-      if (Offset >= 0) {
-        if (GRRegClass.contains(Reg)) {
-          if (LowOffset > Offset) {
-            LowOffset = Offset;
-            LowGPR = Reg;
-          }
+  for (auto &CS : CSI) {
+        if (LowSpillOffset > Offset) {
+          LowSpillOffset = Offset;
+          LowSpillGPR = Reg;
+        }
+        if (CS.isRestored() && LowRestoreOffset > Offset) {
+          LowRestoreOffset = Offset;
+          LowRestoreGPR = Reg;
 
-          if (Offset > HighOffset) {
-            HighOffset = Offset;
-            HighGPR = Reg;
-          }
-        }
+        // Non-volatile GPRs are saved in the dedicated register save area at
+        // the bottom of the stack and are not truly part of the "normal" stack
+        // frame. Mark the frame index as NoAlloc to indicate it as such.
+        unsigned RegSize = 8;
         int FrameIdx = MFFrame.CreateFixedSpillStackObject(RegSize, Offset);
         CS.setFrameIdx(FrameIdx);
-      } else
-        CS.setFrameIdx(INT32_MAX);
+        MFFrame.setStackID(FrameIdx, TargetStackID::NoAlloc);
       }
-  };
-
-  std::vector<CalleeSavedInfo> Spills;
-
-  // For non-leaf functions:
-  // - the address of callee (entry point) register R6 must be saved
-  Spills.push_back(CalleeSavedInfo(Regs.getAddressOfCalleeRegister()));
-
-  // If the function needs a frame pointer, or if the backchain pointer should
-  // be stored, then save the stack pointer register R4.
-  if (hasFP(MF) || MF.getFunction().hasFnAttribute("backchain"))
-    Spills.push_back(CalleeSavedInfo(RegSP));
-
-  // Save the range of call-saved registers, for use by the
-  // prologue/epilogue inserters.
-  ProcessCSI(CSI);
-  MFI->setRestoreGPRRegs(LowGPR, HighGPR, LowOffset);
-
-  // Save the range of call-saved registers, for use by the epilogue inserter.
-  ProcessCSI(Spills);
-  MFI->setSpillGPRRegs(LowGPR, HighGPR, LowOffset);
-
-  // Create spill slots for the remaining registers.
-  for (auto &CS : CSI) {
-    if (CS.getFrameIdx() != INT32_MAX)
-      continue;
+    } else {
       Register Reg = CS.getReg();
       const TargetRegisterClass *RC = TRI->getMinimalPhysRegClass(Reg);
       Align Alignment = TRI->getSpillAlign(*RC);
       unsigned Size = TRI->getSpillSize(*RC);
       Alignment = std::min(Alignment, getStackAlign());
       int FrameIdx = MFFrame.CreateStackObject(Size, Alignment, true);
       CS.setFrameIdx(FrameIdx);
     }
+  }
+
+  // Save the range of call-saved registers, for use by the
+  // prologue/epilogue inserters.
+  if (LowRestoreGPR)
+    MFI->setRestoreGPRRegs(LowRestoreGPR, HighGPR, LowRestoreOffset);
+
+  // Save the range of call-saved registers, for use by the epilogue inserter.
+  assert(LowSpillGPR && "Expected registers to spill");
+  MFI->setSpillGPRRegs(LowSpillGPR, HighGPR, LowSpillOffset);
 
   return true;
 }
 
 void SystemZXPLINKFrameLowering::determineCalleeSaves(MachineFunction &MF,
                                                       BitVector &SavedRegs,
                                                       RegScavenger *RS) const {
   TargetFrameLowering::determineCalleeSaves(MF, SavedRegs, RS);
 
   bool HasFP = hasFP(MF);
   const SystemZSubtarget &Subtarget = MF.getSubtarget<SystemZSubtarget>();
   auto &Regs = Subtarget.getSpecialRegisters<SystemZXPLINK64Registers>();
 
   // If the function requires a frame pointer, record that the hard
   // frame pointer will be clobbered.
   if (HasFP)
     SavedRegs.set(Regs.getFramePointerRegister());
-
-  // If the function is not an XPLeaf function, we need to save the
-  // return address register. We also always use that register for
-  // the return instruction, so it needs to be restored in the
-  // epilogue even though that register is considered to be volatile.
-  // #TODO: Implement leaf detection.
-  SavedRegs.set(Regs.getReturnFunctionAddressRegister());
 }
 
 bool SystemZXPLINKFrameLowering::spillCalleeSavedRegisters(
     MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
     ArrayRef<CalleeSavedInfo> CSI, const TargetRegisterInfo *TRI) const {
   if (CSI.empty())
     return true;