Fix frame pointer layout on AArch64 Linux.

When floating point callee-saved registers were used, the frame pointer would
incorrectly point to the bottom of the CSR space (containing saved floating-point
registers), rather than to the frame record.

While all frame offsets were calculated consistently, resulting in working code,
this prevented stack walkers from being about to traverse the frame list.

Author: resistor

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -1185,7 +1185,26 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
   // For funclets the FP belongs to the containing function.
   if (!IsFunclet && HasFP) {
     // Only set up FP if we actually need to.
-    int64_t FPOffset = isTargetDarwin(MF) ? (AFI->getCalleeSavedStackSize() - 16) : 0;
+    int64_t FPOffset;
+
+    // The frame pointer needs to point to the location of the frame record
+    // (x28 and x29) within the callee saved register space.
+    if (isTargetDarwin(MF)) {
+      // On Darwin, these are located at the top of the CSR space.
+      FPOffset = (AFI->getCalleeSavedStackSize() - 16);
+    } else {
+      // On other systems, these are located in the middle of the CSR space,
+      // after the other GPRs and before the FPRs.
+      assert(MFI.isCalleeSavedInfoValid() && "CalleeSavedInfo not calculated");
+      if (MFI.getCalleeSavedInfo().empty()) {
+        FPOffset = 0;
+      } else {
+        FPOffset = AFI->getCalleeSavedStackSize(MFI, [](unsigned Reg) {
+          return AArch64::FPR64RegClass.contains(Reg) ||
+                 AArch64::FPR128RegClass.contains(Reg);
+        });
+      }
+    }
 
     if (CombineSPBump)
       FPOffset += AFI->getLocalStackSize();
@@ -1842,8 +1861,16 @@ static StackOffset getFPOffset(const MachineFunction &MF, int64_t ObjectOffset)
 
   unsigned FixedObject =
       getFixedObjectSize(MF, AFI, IsWin64, /*IsFunclet=*/false);
-  unsigned FPAdjust = isTargetDarwin(MF)
-                        ? 16 : AFI->getCalleeSavedStackSize(MF.getFrameInfo());
+
+  // Compensate for the position of the frame record within the callee-saved
+  // register space.  On Darwin, this is a fixed offset.  On other systems,
+  // this is determined by the number of callee-saved GPRs, excluding FPRs.
+  unsigned FPAdjust =
+      isTargetDarwin(MF)
+          ? 16
+          : AFI->getCalleeSavedStackSize(MF.getFrameInfo(), [](unsigned Reg) {
+              return AArch64::GPR64RegClass.contains(Reg);
+            });
   return {ObjectOffset + FixedObject + FPAdjust, MVT::i8};
 }
 

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h

@@ -194,9 +194,15 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
   // When CalleeSavedStackSize has not been set (for example when
   // some MachineIR pass is run in isolation), then recalculate
   // the CalleeSavedStackSize directly from the CalleeSavedInfo.
+  // RegisterFilter is a predicate to calculate the stack size for
+  // subsets of the callee-saved registers.  It should return true
+  // for registers that should be included in the size calculation,
+  // and false otherwise.
   // Note: This information can only be recalculated after PEI
   // has assigned offsets to the callee save objects.
-  unsigned getCalleeSavedStackSize(const MachineFrameInfo &MFI) const {
+  unsigned getCalleeSavedStackSize(
+      const MachineFrameInfo &MFI,
+      llvm::function_ref<bool(unsigned)> RegisterFilter = nullptr) const {
     bool ValidateCalleeSavedStackSize = false;
 
 #ifndef NDEBUG
@@ -206,25 +212,40 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
     ValidateCalleeSavedStackSize = HasCalleeSavedStackSize;
 #endif
 
-    if (!HasCalleeSavedStackSize || ValidateCalleeSavedStackSize) {
+    if (RegisterFilter || !HasCalleeSavedStackSize ||
+        ValidateCalleeSavedStackSize) {
       assert(MFI.isCalleeSavedInfoValid() && "CalleeSavedInfo not calculated");
       if (MFI.getCalleeSavedInfo().empty())
         return 0;
 
       int64_t MinOffset = std::numeric_limits<int64_t>::max();
       int64_t MaxOffset = std::numeric_limits<int64_t>::min();
+
+      bool AnyRegistersCounted = false;
       for (const auto &Info : MFI.getCalleeSavedInfo()) {
+        if (RegisterFilter) {
+          unsigned Reg = Info.getReg();
+          if (!RegisterFilter(Reg))
+            continue;
+        }
+
+        AnyRegistersCounted = true;
         int FrameIdx = Info.getFrameIdx();
         if (MFI.getStackID(FrameIdx) != TargetStackID::Default)
           continue;
         int64_t Offset = MFI.getObjectOffset(FrameIdx);
         int64_t ObjSize = MFI.getObjectSize(FrameIdx);
         MinOffset = std::min<int64_t>(Offset, MinOffset);
         MaxOffset = std::max<int64_t>(Offset + ObjSize, MaxOffset);
+        AnyRegistersCounted = true;
       }
 
+      if (!AnyRegistersCounted)
+        return 0;
+
       unsigned Size = alignTo(MaxOffset - MinOffset, 16);
-      assert((!HasCalleeSavedStackSize || getCalleeSavedStackSize() == Size) &&
+      assert((RegisterFilter || !HasCalleeSavedStackSize ||
+              getCalleeSavedStackSize() == Size) &&
              "Invalid size calculated for callee saves");
       return Size;
     }