[ARM] Fix llvm.returnaddress for Thumb1 with R11 frame-pointer (#117735)

When the llvm.returnaddress intrinsic is used, the LR is marked as
live-in to the function, so it must be preserved through the prologue.
This is normally fine, but there is one case for Thumb1 where we use LR
as a temporary in the prologue to set up a frame chain using r11 as the
frame pointer. There are no other registers guaranteed to be free to do
this, so we have to re-load LR from the stack after pushing the callee
saved registers.

Author: ostannard

llvm/lib/Target/ARM/ARMAsmPrinter.cpp

@@ -1309,6 +1309,10 @@ void ARMAsmPrinter::EmitUnwindingInstruction(const MachineInstr *MI) {
       default:
         MI->print(errs());
         llvm_unreachable("Unsupported opcode for unwinding information");
+      case ARM::tLDRspi:
+        // Used to restore LR in a prologue which uses it as a temporary, has
+        // no effect on unwind tables.
+        return;
       case ARM::MOVr:
       case ARM::tMOVr:
         Offset = 0;

llvm/lib/Target/ARM/Thumb1FrameLowering.cpp

@@ -43,6 +43,8 @@
 #include <iterator>
 #include <vector>
 
+#define DEBUG_TYPE "arm-frame-lowering"
+
 using namespace llvm;
 
 Thumb1FrameLowering::Thumb1FrameLowering(const ARMSubtarget &sti)
@@ -277,6 +279,20 @@ void Thumb1FrameLowering::emitPrologue(MachineFunction &MF,
     }
   }
 
+  // Skip past this code sequence, which is emitted to restore the LR if it is
+  // live-in and clobbered by the frame record setup code:
+  //   ldr rX, [sp, #Y]
+  //   mov lr, rX
+  if (MBBI != MBB.end() && MBBI->getOpcode() == ARM::tLDRspi &&
+      MBBI->getFlag(MachineInstr::FrameSetup)) {
+    ++MBBI;
+    if (MBBI != MBB.end() && MBBI->getOpcode() == ARM::tMOVr &&
+        MBBI->getOperand(0).getReg() == ARM::LR &&
+        MBBI->getFlag(MachineInstr::FrameSetup)) {
+      ++MBBI;
+    }
+  }
+
   // Determine starting offsets of spill areas.
   unsigned DPRCSOffset = NumBytes - ArgRegsSaveSize -
                          (FRSize + GPRCS1Size + GPRCS2Size + DPRCSSize);
@@ -857,7 +873,8 @@ static void pushRegsToStack(MachineBasicBlock &MBB,
                             MachineBasicBlock::iterator MI,
                             const TargetInstrInfo &TII,
                             const std::set<Register> &RegsToSave,
-                            const std::set<Register> &CopyRegs) {
+                            const std::set<Register> &CopyRegs,
+                            bool &UsedLRAsTemp) {
   MachineFunction &MF = *MBB.getParent();
   const MachineRegisterInfo &MRI = MF.getRegInfo();
   DebugLoc DL;
@@ -914,6 +931,8 @@ static void pushRegsToStack(MachineBasicBlock &MBB,
         bool isKill = !MRI.isLiveIn(*HiRegToSave);
         if (isKill && !MRI.isReserved(*HiRegToSave))
           MBB.addLiveIn(*HiRegToSave);
+        if (*CopyRegIt == ARM::LR)
+          UsedLRAsTemp = true;
 
         // Emit a MOV from the high reg to the low reg.
         BuildMI(MBB, MI, DL, TII.get(ARM::tMOVr))
@@ -1093,6 +1112,8 @@ bool Thumb1FrameLowering::spillCalleeSavedRegisters(
   // In case FP is a high reg, we need a separate push sequence to generate
   // a correct Frame Record
   bool NeedsFrameRecordPush = hasFP(MF) && ARM::hGPRRegClass.contains(FPReg);
+  bool LRLiveIn = MF.getRegInfo().isLiveIn(ARM::LR);
+  bool UsedLRAsTemp = false;
 
   std::set<Register> FrameRecord;
   std::set<Register> SpilledGPRs;
@@ -1104,7 +1125,22 @@ bool Thumb1FrameLowering::spillCalleeSavedRegisters(
       SpilledGPRs.insert(Reg);
   }
 
-  pushRegsToStack(MBB, MI, TII, FrameRecord, {ARM::LR});
+  // Determine intermediate registers which can be used for pushing the frame
+  // record:
+  // - Unused argument registers
+  // - LR: This is possible because the first PUSH will save it on the stack,
+  //       so it is free to be used as a temporary for the second. However, it
+  //       is possible for LR to be live-in to the function, in which case we
+  //       will need to restore it later in the prologue, so we only use this
+  //       if there are no free argument registers.
+  std::set<Register> FrameRecordCopyRegs;
+  for (unsigned ArgReg : {ARM::R0, ARM::R1, ARM::R2, ARM::R3})
+    if (!MF.getRegInfo().isLiveIn(ArgReg))
+      FrameRecordCopyRegs.insert(ArgReg);
+  if (FrameRecordCopyRegs.empty())
+    FrameRecordCopyRegs.insert(ARM::LR);
+
+  pushRegsToStack(MBB, MI, TII, FrameRecord, FrameRecordCopyRegs, UsedLRAsTemp);
 
   // Determine intermediate registers which can be used for pushing high regs:
   // - Spilled low regs
@@ -1118,7 +1154,33 @@ bool Thumb1FrameLowering::spillCalleeSavedRegisters(
     if (!MF.getRegInfo().isLiveIn(ArgReg))
       CopyRegs.insert(ArgReg);
 
-  pushRegsToStack(MBB, MI, TII, SpilledGPRs, CopyRegs);
+  pushRegsToStack(MBB, MI, TII, SpilledGPRs, CopyRegs, UsedLRAsTemp);
+
+  // If the push sequence used LR as a temporary, and LR is live-in (for
+  // example because it is used by the llvm.returnaddress intrinsic), then we
+  // need to reload it from the stack. Thumb1 does not have a load instruction
+  // which can use LR, so we need to load into a temporary low register and
+  // copy to LR.
+  if (LRLiveIn && UsedLRAsTemp) {
+    auto CopyRegIt = getNextOrderedReg(OrderedCopyRegs.rbegin(),
+                                       OrderedCopyRegs.rend(), CopyRegs);
+    assert(CopyRegIt != OrderedCopyRegs.rend());
+    unsigned NumRegsPushed = FrameRecord.size() + SpilledGPRs.size();
+    LLVM_DEBUG(
+        dbgs() << "LR is live-in but clobbered in prologue, restoring via "
+               << RegInfo->getName(*CopyRegIt) << "\n");
+
+    BuildMI(MBB, MI, DebugLoc(), TII.get(ARM::tLDRspi), *CopyRegIt)
+        .addReg(ARM::SP)
+        .addImm(NumRegsPushed - 1)
+        .add(predOps(ARMCC::AL))
+        .setMIFlags(MachineInstr::FrameSetup);
+
+    BuildMI(MBB, MI, DebugLoc(), TII.get(ARM::tMOVr), ARM::LR)
+        .addReg(*CopyRegIt)
+        .add(predOps(ARMCC::AL))
+        .setMIFlags(MachineInstr::FrameSetup);
+  }
 
   return true;
 }

llvm/test/CodeGen/Thumb/frame-chain.ll

@@ -23,9 +23,9 @@ define dso_local noundef i32 @leaf(i32 noundef %0) {
 ; LEAF-FP-AAPCS:       @ %bb.0:
 ; LEAF-FP-AAPCS-NEXT:    .save {lr}
 ; LEAF-FP-AAPCS-NEXT:    push {lr}
-; LEAF-FP-AAPCS-NEXT:    mov lr, r11
+; LEAF-FP-AAPCS-NEXT:    mov r3, r11
 ; LEAF-FP-AAPCS-NEXT:    .save {r11}
-; LEAF-FP-AAPCS-NEXT:    push {lr}
+; LEAF-FP-AAPCS-NEXT:    push {r3}
 ; LEAF-FP-AAPCS-NEXT:    .setfp r11, sp
 ; LEAF-FP-AAPCS-NEXT:    mov r11, sp
 ; LEAF-FP-AAPCS-NEXT:    .pad #4
@@ -80,9 +80,9 @@ define dso_local noundef i32 @non_leaf(i32 noundef %0) {
 ; FP-AAPCS:       @ %bb.0:
 ; FP-AAPCS-NEXT:    .save {lr}
 ; FP-AAPCS-NEXT:    push {lr}
-; FP-AAPCS-NEXT:    mov lr, r11
+; FP-AAPCS-NEXT:    mov r3, r11
 ; FP-AAPCS-NEXT:    .save {r11}
-; FP-AAPCS-NEXT:    push {lr}
+; FP-AAPCS-NEXT:    push {r3}
 ; FP-AAPCS-NEXT:    .setfp r11, sp
 ; FP-AAPCS-NEXT:    mov r11, sp
 ; FP-AAPCS-NEXT:    .pad #8
@@ -161,9 +161,9 @@ define dso_local void @required_fp(i32 %0, i32 %1) {
 ; FP-AAPCS:       @ %bb.0:
 ; FP-AAPCS-NEXT:    .save {lr}
 ; FP-AAPCS-NEXT:    push {lr}
-; FP-AAPCS-NEXT:    mov lr, r11
+; FP-AAPCS-NEXT:    mov r3, r11
 ; FP-AAPCS-NEXT:    .save {r11}
-; FP-AAPCS-NEXT:    push {lr}
+; FP-AAPCS-NEXT:    push {r3}
 ; FP-AAPCS-NEXT:    .setfp r11, sp
 ; FP-AAPCS-NEXT:    mov r11, sp
 ; FP-AAPCS-NEXT:    .save {r4, r6}
@@ -227,9 +227,9 @@ define dso_local void @required_fp(i32 %0, i32 %1) {
 ; NOFP-AAPCS:       @ %bb.0:
 ; NOFP-AAPCS-NEXT:    .save {lr}
 ; NOFP-AAPCS-NEXT:    push {lr}
-; NOFP-AAPCS-NEXT:    mov lr, r11
+; NOFP-AAPCS-NEXT:    mov r3, r11
 ; NOFP-AAPCS-NEXT:    .save {r11}
-; NOFP-AAPCS-NEXT:    push {lr}
+; NOFP-AAPCS-NEXT:    push {r3}
 ; NOFP-AAPCS-NEXT:    .setfp r11, sp
 ; NOFP-AAPCS-NEXT:    mov r11, sp
 ; NOFP-AAPCS-NEXT:    .save {r4, r6}