[ARM] fix bug causing shrinkwrapping not always being off using PAC

[stuij]

Summary

If you want to check for all uses of PAC, the SpillsLR argument to
shouldSignReturnAddress should be true instead of false, as that value will be
returned from the function if the other checks fall through.

[miyuki]

llvm/lib/Target/ARM/ARMFrameLowering.cpp

1814

Can we use MF.getInfo<ARMFunctionInfo>()->shouldSignReturnAddress() instead (i.e. base the decision on whether or not the function actually spills LR)? Or is there a pass later in the pipeline that can transform the function in such a way that it starts spilling LR?

[stuij]

Yea, good point. I amended the patch.

I don't think we should worry about a pass spilling the LR and we don't do that anywhere else in the code. Practically shrinkwrapping is adjacent to the prolog/epilog inserter where we make the decision to insert PACs.

[miyuki]

LGTM, but please update the commit message. The part saying

SpillsLR argument to shouldSignReturnAddress should be true instead of false

No longer describes the actual change in the code.

[stuij]

I did some more testing, and unfortunately we can't rely on LRSpilled being set correctly.

LRSpilled is being set by ARMFrameLowering::determineCalleeSaves, which is called a couple of times throughout the code. I experimented a bit, and it's a bit dirty to call it from within enableShrinkWrapping, as you need to cast away the const on MachineFunction, and it's a very big slab of code that clearly triggers some side effect (besides setting LRSPilled) as calling it here will make some tests fail. Looking at the failed tests (not just PACBTI tests) it doesn't look like anything wrong is going on, but it's all a bit too icky for my taste.

So I reverted to the previous state, as it's not worth the time-effort to possibly eek out a slight performance edge for a temporary stop-gap which we'll remove once we fix the issue proper. I feel I already spent a bit too much time on this.

[miyuki]

OK, in this case, it makes sense to be conservative and disable shrink-wrapping for all functions that use PAC. LGTM.

llvm/lib/Target/ARM/ARMFrameLowering.cpp

bool ARMFrameLowering::enableShrinkWrapping(const MachineFunction &MF) const {
   // upon function entry (resp. restore it immmediately before return)
   if (STI.hasV8_1MMainlineOps() &&
       MF.getInfo<ARMFunctionInfo>()->isCmseNSEntryFunction())
     return false;
 
   // We are disabling shrinkwrapping for now when PAC is enabled, as
   // shrinkwrapping can cause clobbering of r12 when the PAC code is
   // generated. A follow-up patch will fix this in a more performant manner.
   if (MF.getInfo<ARMFunctionInfo>()->shouldSignReturnAddress(
-          false /*SpillsLR */))
+          true /* SpillsLR */))
     return false;
 
   return true;
 }
 
 void ARMFrameLowering::determineCalleeSaves(MachineFunction &MF,
                                             BitVector &SavedRegs,
                                             RegScavenger *RS) const {

llvm/test/CodeGen/Thumb2/pacbti-m-outliner-3.ll

if.end:                                           ; preds = %entry
   br label %return
 
 return:                                           ; preds = %entry, %if.end
   %retval.0 = phi i32 [ %add3, %if.end ], [ -1, %entry ]
   ret i32 %retval.0
 }
 
 ; CHECK-LABEL: f:
-; CHECK:       bmi    .LBB
 ; ...
 ; CHECK:       pac    r12, lr, sp
 ; CHECK-NEXT:  .save  {r4, r5, r6, lr}
 ; CHECK-NEXT:  push   {r4, r5, r6, lr}
 ; CHECK-NEXT:  .cfi_def_cfa_offset 16
 ; CHECK-NEXT:  .cfi_offset lr, -4
 ; CHECK-NEXT:  .cfi_offset r6, -8
 ; CHECK-NEXT:  .cfi_offset r5, -12
 ; CHECK-NEXT:  .cfi_offset r4, -16
 ; CHECK-NEXT:  .save  {ra_auth_code}
 ; CHECK-NEXT:  str    r12, [sp, #-4]!
 ; CHECK-NEXT:  .cfi_def_cfa_offset 20
 ; CHECK-NEXT:  .cfi_offset ra_auth_code, -20
 ; CHECK-NEXT:  .pad    #4
 ; CHECK-NEXT:  sub     sp, #4
 ; CHECK-NEXT:  .cfi_def_cfa_offset 24
 ; ...
+; CHECK:        bl  OUTLINED_FUNCTION_0
+; ...
 ; CHECK:        add    sp, #4
 ; CHECK-NEXT:   ldr    r12, [sp], #4
 ; CHECK-NEXT:   pop.w  {r4, r5, r6, lr}
 ; CHECK-NEXT:   aut    r12, lr, sp
 ; CHECK-NEXT:   bx     lr
-; ...
-; CHECK: .LBB
-; CHECK:        bx    lr
 
 
 define hidden i32 @g(i32 %a, i32 %b, i32 %c, i32 %d) local_unnamed_addr #0 {
 entry:
   %cmp = icmp slt i32 %a, 0
   br i1 %cmp, label %return, label %if.end
 
 if.end:                                           ; preds = %entry
   %mul = mul nsw i32 %a, 11
   %sub = sub nsw i32 %mul, %b
   %call = tail call i32 @h(i32 %sub, i32 %b)
   %add = add nsw i32 %call, %b
   %mul1 = mul nsw i32 %add, %call
   %add2 = add nsw i32 %d, %c
   %div = sdiv i32 %mul1, %add2
   %add3 = add nsw i32 %div, 1
   br label %return
 
 return:                                           ; preds = %entry, %if.end
   %retval.0 = phi i32 [ %add3, %if.end ], [ -1, %entry ]
   ret i32 %retval.0
 }
 ; CHECK-LABEL: g:
-; CHECK:       bmi    .LBB
-; ...
 ; CHECK:       pac    r12, lr, sp
 ; CHECK-NEXT:  .save  {r4, r5, r6, lr}
 ; CHECK-NEXT:  push   {r4, r5, r6, lr}
 ; CHECK-NEXT:  .cfi_def_cfa_offset 16
 ; CHECK-NEXT:  .cfi_offset lr, -4
 ; CHECK-NEXT:  .cfi_offset r6, -8
 ; CHECK-NEXT:  .cfi_offset r5, -12
 ; CHECK-NEXT:  .cfi_offset r4, -16
 ; CHECK-NEXT:  .save  {ra_auth_code}
 ; CHECK-NEXT:  str    r12, [sp, #-4]!
 ; CHECK-NEXT:  .cfi_def_cfa_offset 20
 ; CHECK-NEXT:  .cfi_offset ra_auth_code, -20
 ; CHECK-NEXT:  .pad   #4
 ; CHECK-NEXT:  sub    sp, #4
 ; CHECK-NEXT:  .cfi_def_cfa_offset 24
 ; ...
+; CHECK:        bl  OUTLINED_FUNCTION_0
+; ...
 ; CHECK:       add    sp, #4
 ; CHECK-NEXT:  ldr    r12, [sp], #4
 ; CHECK-NEXT:  pop.w  {r4, r5, r6, lr}
 ; CHECK-NEXT:  aut    r12, lr, sp
 ; CHECK-NEXT:  bx     lr
-; ...
-; CHECK: .LBB
-; CHECK:       bx     lr
 
 ; CHECK-LABEL: OUTLINED_FUNCTION_0:
 ; CHECK:       pac    r12, lr, sp
 ; CHECK-NEXT:  .save  {ra_auth_code, lr}
 ; CHECK-NEXT:  strd    r12, lr, [sp, #-8]!
 ; CHECK-NEXT:  .cfi_def_cfa_offset 8
 ; CHECK-NEXT:  .cfi_offset lr, -4
 ; CHECK-NEXT:  .cfi_offset ra_auth_code, -8

llvm/test/CodeGen/Thumb2/pacbti-m-outliner-4.ll

 
 if.end:                                           ; preds = %entry
   %add = add nsw i32 %b, %a
   ret i32 %add
 }
 
 ; CHECK-LABEL: _Z1hii:
 ; ...
-; CHECK:    bxgt    lr
-; ...
 ; CHECK:    pac    r12, lr, sp
 ; CHECK-NEXT:    .save    {r7, lr}
 ; CHECK-NEXT:    push    {r7, lr}
 ; CHECK-NEXT:    .cfi_def_cfa_offset 8
 ; CHECK-NEXT:    .cfi_offset lr, -4
 ; CHECK-NEXT:    .cfi_offset r7, -8
 ; CHECK-NEXT:    .save    {ra_auth_code}
 ; CHECK-NEXT:    str    r12, [sp, #-4]!
 ; CHECK-NEXT:    .cfi_def_cfa_offset 12
 ; CHECK-NEXT:    .cfi_offset ra_auth_code, -12
 ; CHECK-NEXT:    .pad    #4
 ; CHECK-NEXT:    sub    sp, #4
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; ...
 ; CHECK-NOT: pac
-; CHECK-NOT: aut
+; CHECK: aut
 ; CHECK:     .cfi_endproc
 
 declare dso_local i8* @__cxa_allocate_exception(i32) local_unnamed_addr
 
 declare dso_local void @__cxa_throw(i8*, i8*, i8*) local_unnamed_addr
 
 define hidden i32 @_Z1fiiii(i32 %a, i32 %b, i32 %c, i32 %d) local_unnamed_addr #0 {
 entry:

 
 return:                                           ; preds = %entry, %if.end
   %retval.0 = phi i32 [ %add2, %if.end ], [ -1, %entry ]
   ret i32 %retval.0
 }
 
 ; CHECK-LABEL: _Z1fiiii:
 ; ...
-; CHECK:    bmi    .L[[B:[a-zA-Z0-9]*]]
-; ...
 ; CHECK:    pac    r12, lr, sp
 ; CHECK-NEXT:    .save    {r4, r5, r6, lr}
 ; CHECK-NEXT:    push    {r4, r5, r6, lr}
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    .cfi_offset lr, -4
 ; CHECK-NEXT:    .cfi_offset r6, -8
 ; CHECK-NEXT:    .cfi_offset r5, -12
 ; CHECK-NEXT:    .cfi_offset r4, -16
 ; CHECK-NEXT:    .save    {ra_auth_code}
 ; CHECK-NEXT:    str    r12, [sp, #-4]!
 ; CHECK-NEXT:    .cfi_def_cfa_offset 20
 ; CHECK-NEXT:    .cfi_offset ra_auth_code, -20
 ; CHECK-NEXT:    .pad    #4
 ; CHECK-NEXT:    sub    sp, #4
 ; CHECK-NEXT:    .cfi_def_cfa_offset 24
 ; ...
+; CHECK:    bl  OUTLINED_FUNCTION_0
+; ...
 ; CHECK:    add    sp, #4
 ; CHECK-NEXT:    ldr    r12, [sp], #4
 ; CHECK-NEXT:    pop.w    {r4, r5, r6, lr}
 ; CHECK-NEXT:    aut    r12, lr, sp
 ; CHECK-NEXT:    bx    lr
-; ...
-; CHECK: .L[[B]]
-; ...
-; CHECK:    bx    lr
 
 
 
 define hidden i32 @_Z1giiii(i32 %a, i32 %b, i32 %c, i32 %d) local_unnamed_addr #0 {
 entry:
   %cmp = icmp slt i32 %a, 0
   br i1 %cmp, label %return, label %if.end
 
 if.end:                                           ; preds = %entry
   %call = tail call i32 @_Z1hii(i32 %a, i32 %b)
   %add = add nsw i32 %call, %b
   %mul = mul nsw i32 %add, %call
   %add1 = add nsw i32 %d, %c
   %div = sdiv i32 %mul, %add1
   %add2 = add nsw i32 %div, 1
   br label %return
 
 return:                                           ; preds = %entry, %if.end
   %retval.0 = phi i32 [ %add2, %if.end ], [ -1, %entry ]
   ret i32 %retval.0
 }
 
 ; CHECK-LABEL: _Z1giiii:
 ; ...
-; CHECK:    bmi    .L[[B:[a-zA-Z0-9]*]]
-; ...
 ; CHECK:    pac    r12, lr, sp
 ; CHECK-NEXT:    .save    {r4, r5, r6, lr}
 ; CHECK-NEXT:    push    {r4, r5, r6, lr}
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    .cfi_offset lr, -4
 ; CHECK-NEXT:    .cfi_offset r6, -8
 ; CHECK-NEXT:    .cfi_offset r5, -12
 ; CHECK-NEXT:    .cfi_offset r4, -16
 ; CHECK-NEXT:    .save    {ra_auth_code}
 ; CHECK-NEXT:    str    r12, [sp, #-4]!
 ; CHECK-NEXT:    .cfi_def_cfa_offset 20
 ; CHECK-NEXT:    .cfi_offset ra_auth_code, -20
 ; CHECK-NEXT:    .pad    #4
 ; CHECK-NEXT:    sub    sp, #4
 ; CHECK-NEXT:    .cfi_def_cfa_offset 24
 ; ...
+; CHECK:    bl  OUTLINED_FUNCTION_0
+; ...
 ; CHECK:    add    sp, #4
 ; CHECK-NEXT:    ldr    r12, [sp], #4
 ; CHECK-NEXT:    pop.w    {r4, r5, r6, lr}
 ; CHECK-NEXT:    aut    r12, lr, sp
 ; CHECK-NEXT:    bx    lr
-; ...
-; CHECK: .L[[B]]
-; ...
-; CHECK:    bx    lr
 
 
 ; CHEK-LABEL: OUTLINED_FUNCTION_0:
 ; CHECK-NOT: pac
 ; CHECK-NOT: aut
 ; CHECK:        b    _Z1hii
 
 attributes #0 = { minsize noinline optsize "denormal-fp-math"="preserve-sign,preserve-sign" "denormal-fp-math-f32"="ieee,ieee" "disable-tail-calls"="false" "frame-pointer"="none" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="true" "no-jump-tables"="false" "no-nans-fp-math"="true" "no-signed-zeros-fp-math"="true" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="cortex-m3" "target-features"="+armv7-m,+hwdiv,+thumb-mode" "unsafe-fp-math"="false" "use-soft-float"="false" }

 ; UNWIND-LABEL: FunctionAddress: 0x0
 ; UNWIND:       Opcodes
 ; UNWIND-NEXT:  0x00      ; vsp = vsp + 4
 ; UNWIND-NEXT:  0xB4      ; pop ra_auth_code
 ; UNWIND-NEXT:  0x84 0x08 ; pop {r7, lr}
 ; UNWIND-NEXT:  0xB0      ; finish
 ; UNWIND-NEXT:  0xB0      ; finish
 
-; UNWIND-LABEL: FunctionAddress: 0x2C
+; UNWIND-LABEL: FunctionAddress: 0x3C
 ; UNWIND:       Opcodes
 ; UNWIND-NEXT:  0x00      ; vsp = vsp + 4
 ; UNWIND-NEXT:  0xB4      ; pop ra_auth_code
 ; UNWIND-NEXT:  0xAA      ; pop {r4, r5, r6, lr}
 
-; UNWIND-LABEL: FunctionAddress: 0x62
+; UNWIND-LABEL: FunctionAddress: 0x72
 ; UNWIND:       Opcodes
 ; UNWIND-NEXT:  0x00      ; vsp = vsp + 4
 ; UNWIND-NEXT:  0xB4      ; pop ra_auth_code
 ; UNWIND-NEXT:  0xAA      ; pop {r4, r5, r6, lr}
 
-; UNWIND-LABEL: FunctionAddress: 0x98
+; UNWIND-LABEL: FunctionAddress: 0xA8
 ; UNWIND:       Opcodes
 ; UNWIND-NEXT:  0xB0      ; finish
 ; UNWIND-NEXT:  0xB0      ; finish
 ; UNWIND-NEXT:  0xB0      ; finish
 
-; UNWIND: 00000099 {{.*}} OUTLINED_FUNCTION_0
-; UWNIND: 0000002d {{.*}} _Z1fiiii
-; UWNIND: 00000063 {{.*}} _Z1giiii
+; UNWIND: 000000a9 {{.*}} OUTLINED_FUNCTION_0
 ; UWNIND: 00000001 {{.*}} _Z1hii
+; UWNIND: 0000003d {{.*}} _Z1fiiii
+; UWNIND: 00000073 {{.*}} _Z1giiii