[ARM] Fix prologue/epilogue for pacbti-m leaf functions

[john.brawn]

Summary

R12 is callee-saved in functions with pacbti-m enabled, but this is done in assignCalleeSavedSpillSlots, meaning that in determineCalleeSaves we have to manually set CanEliminateFrame.

This fixes a bug where in leaf functions with no other callee-saved registers the aut instruction wouldn't be emitted and stack offsets of arguments passed on the stack would be incorrect.

[olista01]

LGTM

llvm/lib/Target/ARM/ARMFrameLowering.cpp

void ARMFrameLowering::determineCalleeSaves(MachineFunction &MF,
   // Spill the BasePtr if it's used.
   if (RegInfo->hasBasePointer(MF))
     SavedRegs.set(RegInfo->getBaseRegister());
 
   // On v8.1-M.Main CMSE entry functions save/restore FPCXT.
   if (STI.hasV8_1MMainlineOps() && AFI->isCmseNSEntryFunction())
     CanEliminateFrame = false;
 
+  // When return address signing is enabled R12 is treated as callee-saved.
+  if (AFI->shouldSignReturnAddress())
+    CanEliminateFrame = false;
+
   // Don't spill FP if the frame can be eliminated. This is determined
   // by scanning the callee-save registers to see if any is modified.
   const MCPhysReg *CSRegs = RegInfo->getCalleeSavedRegs(&MF);
   for (unsigned i = 0; CSRegs[i]; ++i) {
     unsigned Reg = CSRegs[i];
     bool Spilled = false;
     if (SavedRegs.test(Reg)) {
       Spilled = true;

llvm/test/CodeGen/Thumb2/pacbti-m-stack-arg.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
+; RUN: llc -mtriple=thumbv8m.main-none-eabi %s -o - | FileCheck %s
+
+define i32 @test_leaf(i32 %arg0, i32 %arg1, i32 %arg2, i32 %arg3, i32 %x) {
+; CHECK-LABEL: test_leaf:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    pac r12, lr, sp
+; CHECK-NEXT:    .save {ra_auth_code}
+; CHECK-NEXT:    str r12, [sp, #-4]!
+; CHECK-NEXT:    ldr r0, [sp, #4]
+; CHECK-NEXT:    ldr r12, [sp], #4
+; CHECK-NEXT:    aut r12, lr, sp
+; CHECK-NEXT:    bx lr
+entry:
+  ret i32 %x
+}
+
+define i32 @test_non_leaf(i32 %arg0, i32 %arg1, i32 %arg2, i32 %arg3, i32 %x) {
+; CHECK-LABEL: test_non_leaf:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    pac r12, lr, sp
+; CHECK-NEXT:    .save {r7, lr}
+; CHECK-NEXT:    push {r7, lr}
+; CHECK-NEXT:    .save {ra_auth_code}
+; CHECK-NEXT:    str r12, [sp, #-4]!
+; CHECK-NEXT:    .pad #4
+; CHECK-NEXT:    sub sp, #4
+; CHECK-NEXT:    bl otherfn
+; CHECK-NEXT:    ldr r0, [sp, #16]
+; CHECK-NEXT:    add sp, #4
+; CHECK-NEXT:    ldr r12, [sp], #4
+; CHECK-NEXT:    pop.w {r7, lr}
+; CHECK-NEXT:    aut r12, lr, sp
+; CHECK-NEXT:    bx lr
+entry:
+  call void @otherfn()
+  ret i32 %x
+}
+
+declare void @otherfn(...)
+
+!llvm.module.flags = !{!0, !1, !2}
+
+!0 = !{i32 8, !"branch-target-enforcement", i32 0}
+!1 = !{i32 8, !"sign-return-address", i32 1}
+!2 = !{i32 8, !"sign-return-address-all", i32 1}