[PPC][AIX] Save/restore r31 when using base pointer (#100182)

When the base pointer r30 is used to hold the stack pointer, r30 is
spilled in the prologue. On AIX registers are saved from highest to
lowest, so r31 also needs to be saved.

Fixes #96411

Author: syzaara

llvm/lib/Target/PowerPC/PPCFrameLowering.cpp

@@ -1007,7 +1007,7 @@ void PPCFrameLowering::emitPrologue(MachineFunction &MF,
         // R0 cannot be used as a base register, but it can be used as an
         // index in a store-indexed.
         int LastOffset = 0;
-        if (HasFP)  {
+        if (HasFP) {
           // R0 += (FPOffset-LastOffset).
           // Need addic, since addi treats R0 as 0.
           BuildMI(MBB, MBBI, dl, TII.get(PPC::ADDIC), ScratchReg)
@@ -2025,8 +2025,18 @@ void PPCFrameLowering::determineCalleeSaves(MachineFunction &MF,
   // code. Same goes for the base pointer and the PIC base register.
   if (needsFP(MF))
     SavedRegs.reset(isPPC64 ? PPC::X31 : PPC::R31);
-  if (RegInfo->hasBasePointer(MF))
+  if (RegInfo->hasBasePointer(MF)) {
     SavedRegs.reset(RegInfo->getBaseRegister(MF));
+    // On AIX, when BaseRegister(R30) is used, need to spill r31 too to match
+    // AIX trackback table requirement.
+    if (!needsFP(MF) && !SavedRegs.test(isPPC64 ? PPC::X31 : PPC::R31) &&
+        Subtarget.isAIXABI()) {
+      assert(
+          (RegInfo->getBaseRegister(MF) == (isPPC64 ? PPC::X30 : PPC::R30)) &&
+          "Invalid base register on AIX!");
+      SavedRegs.set(isPPC64 ? PPC::X31 : PPC::R31);
+    }
+  }
   if (FI->usesPICBase())
     SavedRegs.reset(PPC::R30);
 

llvm/test/CodeGen/PowerPC/aix-base-pointer.ll

@@ -6,6 +6,7 @@
 
 ; Use an overaligned buffer to force base-pointer usage. Test verifies:
 ; - base pointer register (r30) is saved/defined/restored.
+; - frame pointer register (r31) is saved/defined/restored.
 ; - stack frame is allocated with correct alignment.
 ; - Address of %AlignedBuffer is calculated based off offset from the stack
 ;   pointer.
@@ -25,7 +26,9 @@ declare void @callee(ptr)
 ; 32BIT:         subfic 0, 0, -224
 ; 32BIT:         stwux 1, 1, 0
 ; 32BIT:         addi 3, 1, 64
+; 32BIT:         stw 31, -12(30)
 ; 32BIT:         bl .callee
+; 32BIT:         lwz 31, -12(30)
 ; 32BIT:         mr 1, 30
 ; 32BIT:         lwz 30, -16(1)
 
@@ -36,6 +39,8 @@ declare void @callee(ptr)
 ; 64BIT:         subfic 0, 0, -288
 ; 64BIT:         stdux 1, 1, 0
 ; 64BIT:         addi 3, 1, 128
+; 64BIT:         std 31, -16(30)
 ; 64BIT:         bl .callee
+; 64BIT:         ld 31, -16(30)
 ; 64BIT:         mr 1, 30
 ; 64BIT:         ld 30, -24(1)