[AArch64][SVE] Don't align the last SVE callee save.

[sdesmalen]

Summary

Instead of aligning the last callee-saved-register slot to the stack
alignment (16 bytes), just align the SVE callee-saved block. This also
simplifies the code that allocates space for the callee-saves.

This change is needed to make sure the offset to which the callee-saved
register is spilled, corresponds to the offset used for e.g. unwind call
frame instructions.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

2045

Since ZPRs are architecturally defined in multiplies of 16-bytes I don't see why anything needs fixing up.

At the same time I think PPRs are more problematic because they're defined in multiples of 2-bytes and thus we'd need a group of 8 to fulfil alignment.

[efriedma]

llvm/test/CodeGen/AArch64/framelayout-sve.mir

368

It seems weird to introduce a gap in the middle of the callee-saves; I mean, the CPU doesn't care where the gap is, but it seems less readable for people. I'd prefer to keep this as-is, and change whatever code is assuming the other layout.

[efriedma]

LGTM

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

if (windowsRequiresStackProbe(MF, NumBytes)) {
     NumBytes = 0;
   }
 
   StackOffset AllocateBefore = SVEStackSize, AllocateAfter = {};
   MachineBasicBlock::iterator CalleeSavesBegin = MBBI, CalleeSavesEnd = MBBI;
 
   // Process the SVE callee-saves to determine what space needs to be
   // allocated.
-  if (AFI->getSVECalleeSavedStackSize()) {
+  if (int64_t CalleeSavedSize = AFI->getSVECalleeSavedStackSize()) {
     // Find callee save instructions in frame.
     CalleeSavesBegin = MBBI;
     assert(IsSVECalleeSave(CalleeSavesBegin) && "Unexpected instruction");
     while (IsSVECalleeSave(MBBI) && MBBI != MBB.getFirstTerminator())
       ++MBBI;
     CalleeSavesEnd = MBBI;
 
-    int64_t OffsetToFirstCalleeSaveFromSP =
-        MFI.getObjectOffset(AFI->getMaxSVECSFrameIndex());
-    StackOffset OffsetToCalleeSavesFromSP =
-        StackOffset(OffsetToFirstCalleeSaveFromSP, MVT::nxv1i8) + SVEStackSize;
-    AllocateBefore -= OffsetToCalleeSavesFromSP;
+    AllocateBefore = {CalleeSavedSize, MVT::nxv1i8};
     AllocateAfter = SVEStackSize - AllocateBefore;
   }
 
   // Allocate space for the callee saves (if any).
   emitFrameOffset(MBB, CalleeSavesBegin, DL, AArch64::SP, AArch64::SP,
                   -AllocateBefore, TII,
                   MachineInstr::FrameSetup);
 
void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
 
   NumBytes -= PrologueSaveSize;
   assert(NumBytes >= 0 && "Negative stack allocation size!?");
 
   // Process the SVE callee-saves to determine what space needs to be
   // deallocated.
   StackOffset DeallocateBefore = {}, DeallocateAfter = SVEStackSize;
   MachineBasicBlock::iterator RestoreBegin = LastPopI, RestoreEnd = LastPopI;
-  if (AFI->getSVECalleeSavedStackSize()) {
+  if (int64_t CalleeSavedSize = AFI->getSVECalleeSavedStackSize()) {
     RestoreBegin = std::prev(RestoreEnd);;
     while (IsSVECalleeSave(RestoreBegin) &&
            RestoreBegin != MBB.begin())
       --RestoreBegin;
     ++RestoreBegin;
 
     assert(IsSVECalleeSave(RestoreBegin) &&
            IsSVECalleeSave(std::prev(RestoreEnd)) && "Unexpected instruction");
 
-    int64_t OffsetToFirstCalleeSaveFromSP =
-        MFI.getObjectOffset(AFI->getMaxSVECSFrameIndex());
-    StackOffset OffsetToCalleeSavesFromSP =
-        StackOffset(OffsetToFirstCalleeSaveFromSP, MVT::nxv1i8) + SVEStackSize;
-    DeallocateBefore = OffsetToCalleeSavesFromSP;
-    DeallocateAfter = SVEStackSize - DeallocateBefore;
+    StackOffset CalleeSavedSizeAsOffset = {CalleeSavedSize, MVT::nxv1i8};
+    DeallocateBefore = SVEStackSize - CalleeSavedSizeAsOffset;
+    DeallocateAfter = CalleeSavedSizeAsOffset;
   }
 
   // Deallocate the SVE area.
   if (SVEStackSize) {
     if (AFI->isStackRealigned()) {
       if (AFI->getSVECalleeSavedStackSize())
         // Set SP to start of SVE area, from which the callee-save reloads
         // can be done. The code below will deallocate the stack space
if (i + 1 < Count) {
             !invalidateWindowsRegisterPairing(RPI.Reg1, NextReg, NeedsWinCFI))
           RPI.Reg2 = NextReg;
         break;
       case RegPairInfo::FPR128:
         if (AArch64::FPR128RegClass.contains(NextReg))
           RPI.Reg2 = NextReg;
         break;
       case RegPairInfo::PPR:
       case RegPairInfo::ZPR:
         break;
       }
     }
 
     // If either of the registers to be saved is the lr register, it means that
     // we also need to save lr in the shadow call stack.
     if ((RPI.Reg1 == AArch64::LR || RPI.Reg2 == AArch64::LR) &&
         MF.getFunction().hasFnAttribute(Attribute::ShadowCallStack)) {
auto Assign = [&MFI](int FI, int64_t Offset) {
     LLVM_DEBUG(dbgs() << "alloc FI(" << FI << ") at SP[" << Offset << "]\n");
     MFI.setObjectOffset(FI, Offset);
   };
 
   int64_t Offset = 0;
 
   // Then process all callee saved slots.
   if (getSVECalleeSaveSlotRange(MFI, MinCSFrameIndex, MaxCSFrameIndex)) {
-    // Make sure to align the last callee save slot.
-    MFI.setObjectAlignment(MaxCSFrameIndex, Align(16));
-
     // Assign offsets to the callee save slots.
     for (int I = MinCSFrameIndex; I <= MaxCSFrameIndex; ++I) {
       Offset += MFI.getObjectSize(I);
       Offset = alignTo(Offset, MFI.getObjectAlign(I));
       if (AssignOffsets)
         Assign(I, -Offset);
     }
   }
 
+  // Ensure that the Callee-save area is aligned to 16bytes.
+  Offset = alignTo(Offset, Align(16U));
+
   // Create a buffer of SVE objects to allocate and sort it.
   SmallVector<int, 8> ObjectsToAllocate;
   for (int I = 0, E = MFI.getObjectIndexEnd(); I != E; ++I) {
     unsigned StackID = MFI.getStackID(I);
     if (StackID != TargetStackID::SVEVector)
       continue;
     if (MaxCSFrameIndex >= I && I >= MinCSFrameIndex)
       continue;

llvm/test/CodeGen/AArch64/framelayout-sve.mir

bb.0.entry:
     RET_ReallyLR
 ---
 ...
 # CHECK-LABEL: name: save_restore_pregs_sve
 # CHECK: $sp = frame-setup ADDVL_XXI $sp, -1
 # CHECK: frame-setup STR_PXI killed $p6, $sp, 5
 # CHECK: frame-setup STR_PXI killed $p5, $sp, 6
 # CHECK: frame-setup STR_PXI killed $p4, $sp, 7
 # CHECK: $sp = frame-setup SUBXri $sp, 32, 0
 
 # CHECK: $sp = frame-destroy ADDXri $sp, 32, 0
 # CHECK: $p6 = frame-destroy LDR_PXI $sp, 5
 # CHECK: $p5 = frame-destroy LDR_PXI $sp, 6
 # CHECK: $p4 = frame-destroy LDR_PXI $sp, 7
 # CHECK: $sp = frame-destroy ADDVL_XXI $sp, 1
 # CHECK: RET_ReallyLR
 name: save_restore_pregs_sve

 # CHECK:        - { id: 5, name: '', type: spill-slot, offset: -146, size: 2, alignment: 2,
 # CHECK-NEXT:       stack-id: sve-vec,
 # CHECK:        - { id: 6, name: '', type: spill-slot, offset: -16, size: 16, alignment: 16,
 # CHECK-NEXT:       stack-id: sve-vec, callee-saved-register: '$z8',
 # CHECK:        - { id: 7, name: '', type: spill-slot, offset: -32, size: 16, alignment: 16,
 # CHECK-NEXT:       stack-id: sve-vec, callee-saved-register: '$z23',
 # CHECK:        - { id: 8, name: '', type: spill-slot, offset: -34, size: 2, alignment: 2,
 # CHECK-NEXT:       stack-id: sve-vec, callee-saved-register: '$p4',
-# CHECK:        - { id: 9, name: '', type: spill-slot, offset: -48, size: 2, alignment: 16,
+# CHECK:        - { id: 9, name: '', type: spill-slot, offset: -36, size: 2, alignment: 2,
 # CHECK-NEXT:       stack-id: sve-vec, callee-saved-register: '$p15',
 # CHECK:        - { id: 10, name: '', type: spill-slot, offset: -16, size: 8, alignment: 16,
 # CHECK-NEXT:       stack-id: default, callee-saved-register: '$fp',
 #
 # CHECK:      bb.0.entry:
 # CHECK-NEXT: $sp = frame-setup STRXpre killed $[[SCRATCH:[a-z0-9]+]], $sp, -16
 # CHECK-NEXT: $sp = frame-setup ADDVL_XXI $sp, -3
 # CHECK-NEXT: STR_PXI killed $p15, $sp, 6

llvm/test/CodeGen/AArch64/sve-trunc.ll

 
 define <vscale x 16 x i1> @trunc_i64toi1_split3(<vscale x 16 x i64> %in) {
 ; CHECK-LABEL: trunc_i64toi1_split3:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    addvl sp, sp, #-1
 ; CHECK-NEXT:    str p4, [sp, #7, mul vl] // 2-byte Folded Spill
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
-; CHECK-NEXT:    .cfi_offset p4, -16
+; CHECK-NEXT:    .cfi_offset p4, -2
 ; CHECK-NEXT:    .cfi_offset w29, -16
 ; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    and z7.d, z7.d, #0x1
 ; CHECK-NEXT:    and z6.d, z6.d, #0x1
 ; CHECK-NEXT:    and z5.d, z5.d, #0x1
 ; CHECK-NEXT:    and z4.d, z4.d, #0x1
 ; CHECK-NEXT:    and z3.d, z3.d, #0x1
 ; CHECK-NEXT:    and z2.d, z2.d, #0x1