[AArch64] Simplify prolog/epilog callee save/restore. NFC.

[gberry]

Summary

Factor out common code for callee-save register pair calculation. This
is intended to simplify follow-on changes that reduce the number of
registers saved/restored.

Depends on D16732

[george.burgess.iv]

Thanks for doing this!

I just have two tiny nits for you. :)

lib/Target/AArch64/AArch64FrameLowering.cpp

682

Can this take a SmallVectorImpl<RegPairInfo> & instead?

687

It seems that you always add Count/2 elements to RegPairs. Is there any harm in reserving this amount of space in RegPairs ahead of time?

[gberry]

lib/Target/AArch64/AArch64FrameLowering.cpp

682

Thanks, I hadn't learned that trick yet.

687

The number of RegPairs added is going to change in my next patch (and be less predictable up front), so I'm not sure this makes sense. It will still be possible to reserve an upper bound in the next change (with the possibility of not using all of the reserved space). Do you still think this is worth doing in light of that?

[george.burgess.iv]

lib/Target/AArch64/AArch64FrameLowering.cpp

687

Ah, if the result size isn't going to be basically handed to us for free, then I agree with you that it's probably better to not resize upfront. :)

lib/Target/AArch64/AArch64FrameLowering.cpp

if (Reg != AArch64::LR)
     return getKillRegState(true);
 
   // LR maybe referred to later by an @llvm.returnaddress intrinsic.
   bool LRLiveIn = MF.getRegInfo().isLiveIn(AArch64::LR);
   bool LRKill = !(LRLiveIn && MF.getFrameInfo()->isReturnAddressTaken());
   return getKillRegState(LRKill);
 }
 
-bool AArch64FrameLowering::spillCalleeSavedRegisters(
-    MachineBasicBlock &MBB, MachineBasicBlock::iterator MI,
-    const std::vector<CalleeSavedInfo> &CSI,
-    const TargetRegisterInfo *TRI) const {
-  MachineFunction &MF = *MBB.getParent();
-  const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
+struct RegPairInfo {
+  RegPairInfo() : Reg1(AArch64::NoRegister), Reg2(AArch64::NoRegister) {}
+  unsigned Reg1;
+  unsigned Reg2;
+  int FrameIdx;
+  int Offset;
+  bool IsGPR;
+};
+
+static void computeCalleeSaveRegisterPairs(
+    const std::vector<CalleeSavedInfo> &CSI, const TargetRegisterInfo *TRI,
+    SmallVector<RegPairInfo, 8> &RegPairs) {
+
   unsigned Count = CSI.size();
-  DebugLoc DL;
   assert((Count & 1) == 0 && "Odd number of callee-saved regs to spill!");
 
   for (unsigned i = 0; i < Count; i += 2) {
     unsigned idx = Count - i - 2;
-    unsigned Reg1 = CSI[idx].getReg();
-    unsigned Reg2 = CSI[idx + 1].getReg();
+    RegPairInfo RPI;
+    RPI.Reg1 = CSI[idx].getReg();
+    RPI.Reg2 = CSI[idx + 1].getReg();
+
     // GPRs and FPRs are saved in pairs of 64-bit regs. We expect the CSI
     // list to come in sorted by frame index so that we can issue the store
     // pair instructions directly. Assert if we see anything otherwise.
     //
     // The order of the registers in the list is controlled by
     // getCalleeSavedRegs(), so they will always be in-order, as well.
     assert(CSI[idx].getFrameIdx() + 1 == CSI[idx + 1].getFrameIdx() &&
            "Out of order callee saved regs!");
-    unsigned StrOpc;
     assert((Count & 1) == 0 && "Odd number of callee-saved regs to spill!");
     assert((i & 1) == 0 && "Odd index for callee-saved reg spill!");
+    RPI.FrameIdx = CSI[idx + 1].getFrameIdx();
+
+    if (AArch64::GPR64RegClass.contains(RPI.Reg1))
+      RPI.IsGPR = true;
+    else if (AArch64::FPR64RegClass.contains(RPI.Reg1))
+      RPI.IsGPR = false;
+    else
+      llvm_unreachable("Unexpected callee saved register!");
+    // Compute offset: i = 0 => offset = Count;
+    //                 i = 2 => offset = -(Count - 2) + Count = 2 = i; etc.
+    RPI.Offset = (i == 0) ? Count : i;
+    assert((RPI.Offset >= -64 && RPI.Offset <= 63) &&
+           "Offset out of bounds for LDP/STP immediate");
+
+    RegPairs.push_back(RPI);
+  }
+
+}
+
+bool AArch64FrameLowering::spillCalleeSavedRegisters(
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator MI,
+    const std::vector<CalleeSavedInfo> &CSI,
+    const TargetRegisterInfo *TRI) const {
+  MachineFunction &MF = *MBB.getParent();
+  const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
+  DebugLoc DL;
+  SmallVector<RegPairInfo, 8> RegPairs;
+
+  computeCalleeSaveRegisterPairs(CSI, TRI, RegPairs);
+
+  for (auto RPII = RegPairs.begin(), RPIE = RegPairs.end(); RPII != RPIE;
+       ++RPII) {
+    RegPairInfo RPI = *RPII;
+    unsigned Reg1 = RPI.Reg1;
+    unsigned Reg2 = RPI.Reg2;
+    unsigned StrOpc;
+
     // Issue sequence of non-sp increment and pi sp spills for cs regs. The
     // first spill is a pre-increment that allocates the stack.
     // For example:
     //    stp     x22, x21, [sp, #-48]!   // addImm(-6)
     //    stp     x20, x19, [sp, #16]    // addImm(+2)
     //    stp     fp, lr, [sp, #32]      // addImm(+4)
     // Rationale: This sequence saves uop updates compared to a sequence of
     // pre-increment spills like stp xi,xj,[sp,#-16]!
-    // Note: Similar rational and sequence for restores in epilog.
-    if (AArch64::GPR64RegClass.contains(Reg1)) {
-      assert(AArch64::GPR64RegClass.contains(Reg2) &&
-             "Expected GPR64 callee-saved register pair!");
+    // Note: Similar rationale and sequence for restores in epilog.
+    bool BumpSP = RPII == RegPairs.begin();
+    if (RPI.IsGPR) {
       // For first spill use pre-increment store.
-      if (i == 0)
+      if (BumpSP)
         StrOpc = AArch64::STPXpre;
       else
         StrOpc = AArch64::STPXi;
-    } else if (AArch64::FPR64RegClass.contains(Reg1)) {
-      assert(AArch64::FPR64RegClass.contains(Reg2) &&
-             "Expected FPR64 callee-saved register pair!");
+    } else {
       // For first spill use pre-increment store.
-      if (i == 0)
+      if (BumpSP)
         StrOpc = AArch64::STPDpre;
       else
         StrOpc = AArch64::STPDi;
-    } else
-      llvm_unreachable("Unexpected callee saved register!");
+    }
     DEBUG(dbgs() << "CSR spill: (" << TRI->getName(Reg1) << ", "
-                 << TRI->getName(Reg2) << ") -> fi#(" << CSI[idx].getFrameIdx()
-                 << ", " << CSI[idx + 1].getFrameIdx() << ")\n");
-    // Compute offset: i = 0 => offset = -Count;
-    //                 i = 2 => offset = -(Count - 2) + Count = 2 = i; etc.
-    const int Offset = (i == 0) ? -Count : i;
-    assert((Offset >= -64 && Offset <= 63) &&
-           "Offset out of bounds for STP immediate");
+                 << TRI->getName(Reg2) << ") -> fi#(" << RPI.FrameIdx
+                 << ", " << RPI.FrameIdx+1 << ")\n");
+
+    const int Offset = BumpSP ? -RPI.Offset : RPI.Offset;
     MachineInstrBuilder MIB = BuildMI(MBB, MI, DL, TII.get(StrOpc));
-    if (StrOpc == AArch64::STPDpre || StrOpc == AArch64::STPXpre)
+    if (BumpSP)
       MIB.addReg(AArch64::SP, RegState::Define);
 
     MBB.addLiveIn(Reg1);
     MBB.addLiveIn(Reg2);
     MIB.addReg(Reg2, getPrologueDeath(MF, Reg2))
         .addReg(Reg1, getPrologueDeath(MF, Reg1))
         .addReg(AArch64::SP)
         .addImm(Offset) // [sp, #offset * 8], where factor * 8 is implicit
         .setMIFlag(MachineInstr::FrameSetup);
   }
   return true;
 }
 
 bool AArch64FrameLowering::restoreCalleeSavedRegisters(
     MachineBasicBlock &MBB, MachineBasicBlock::iterator MI,
     const std::vector<CalleeSavedInfo> &CSI,
     const TargetRegisterInfo *TRI) const {
   MachineFunction &MF = *MBB.getParent();
   const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
-  unsigned Count = CSI.size();
   DebugLoc DL;
-  assert((Count & 1) == 0 && "Odd number of callee-saved regs to spill!");
+  SmallVector<RegPairInfo, 8> RegPairs;
 
   if (MI != MBB.end())
     DL = MI->getDebugLoc();
 
-  for (unsigned i = 0; i < Count; i += 2) {
-    unsigned Reg1 = CSI[i].getReg();
-    unsigned Reg2 = CSI[i + 1].getReg();
-    // GPRs and FPRs are saved in pairs of 64-bit regs. We expect the CSI
-    // list to come in sorted by frame index so that we can issue the store
-    // pair instructions directly. Assert if we see anything otherwise.
-    assert(CSI[i].getFrameIdx() + 1 == CSI[i + 1].getFrameIdx() &&
-           "Out of order callee saved regs!");
+  computeCalleeSaveRegisterPairs(CSI, TRI, RegPairs);
+
+  for (auto RPII = RegPairs.rbegin(), RPIE = RegPairs.rend(); RPII != RPIE;
+       ++RPII) {
+    RegPairInfo RPI = *RPII;
+    unsigned Reg1 = RPI.Reg1;
+    unsigned Reg2 = RPI.Reg2;
+
     // Issue sequence of non-sp increment and sp-pi restores for cs regs. Only
     // the last load is sp-pi post-increment and de-allocates the stack:
     // For example:
     //    ldp     fp, lr, [sp, #32]       // addImm(+4)
     //    ldp     x20, x19, [sp, #16]     // addImm(+2)
     //    ldp     x22, x21, [sp], #48     // addImm(+6)
     // Note: see comment in spillCalleeSavedRegisters()
     unsigned LdrOpc;
-
-    assert((Count & 1) == 0 && "Odd number of callee-saved regs to spill!");
-    assert((i & 1) == 0 && "Odd index for callee-saved reg spill!");
-    if (AArch64::GPR64RegClass.contains(Reg1)) {
-      assert(AArch64::GPR64RegClass.contains(Reg2) &&
-             "Expected GPR64 callee-saved register pair!");
-      if (i == Count - 2)
+    bool BumpSP = RPII == std::prev(RegPairs.rend());
+    if (RPI.IsGPR) {
+      if (BumpSP)
         LdrOpc = AArch64::LDPXpost;
       else
         LdrOpc = AArch64::LDPXi;
-    } else if (AArch64::FPR64RegClass.contains(Reg1)) {
-      assert(AArch64::FPR64RegClass.contains(Reg2) &&
-             "Expected FPR64 callee-saved register pair!");
-      if (i == Count - 2)
+    } else {
+      if (BumpSP)
         LdrOpc = AArch64::LDPDpost;
       else
         LdrOpc = AArch64::LDPDi;
-    } else
-      llvm_unreachable("Unexpected callee saved register!");
+    }
     DEBUG(dbgs() << "CSR restore: (" << TRI->getName(Reg1) << ", "
-                 << TRI->getName(Reg2) << ") -> fi#(" << CSI[i].getFrameIdx()
-                 << ", " << CSI[i + 1].getFrameIdx() << ")\n");
+                 << TRI->getName(Reg2) << ") -> fi#(" << RPI.FrameIdx
+                 << ", " << RPI.FrameIdx+1 << ")\n");
 
-    // Compute offset: i = 0 => offset = Count - 2; i = 2 => offset = Count - 4;
-    // etc.
-    const int Offset = (i == Count - 2) ? Count : Count - i - 2;
-    assert((Offset >= -64 && Offset <= 63) &&
-           "Offset out of bounds for LDP immediate");
+    const int Offset = RPI.Offset;
     MachineInstrBuilder MIB = BuildMI(MBB, MI, DL, TII.get(LdrOpc));
-    if (LdrOpc == AArch64::LDPXpost || LdrOpc == AArch64::LDPDpost)
+    if (BumpSP)
       MIB.addReg(AArch64::SP, RegState::Define);
 
     MIB.addReg(Reg2, getDefRegState(true))
         .addReg(Reg1, getDefRegState(true))
         .addReg(AArch64::SP)
         .addImm(Offset) // [sp], #offset * 8  or [sp, #offset * 8]
                         // where the factor * 8 is implicit
         .setMIFlag(MachineInstr::FrameDestroy);
if (AArch64::GPR64RegClass.contains(OddReg)) {
       // unspilled register that isn't reserved, if there is one.
       if (Reg != AArch64::NoRegister && !RegInfo->isReservedReg(MF, Reg))
         ExtraCSSpill = true;
     } else
       NumFPRSpilled += 2;
 
     CanEliminateFrame = false;
   }
+  DEBUG(dbgs() << "\n");
 
   // FIXME: Set BigStack if any stack slot references may be out of range.
   // For now, just conservatively guestimate based on unscaled indexing
   // range. We'll end up allocating an unnecessary spill slot a lot, but
   // realistically that's not a big deal at this stage of the game.
   // The CSR spill slots have not been allocated yet, so estimateStackSize
   // won't include them.
   MachineFrameInfo *MFI = MF.getFrameInfo();