[AArch64][GlobalISel] Implement __builtin_return_address for PAC-RET

[chill]

Summary

This patch implements stripping of the PAC in the return address for GlobalISel.

Implementation for when not using GLobalISel is in https://reviews.llvm.org/D75044
The analogous GCC patch is https://gcc.gnu.org/git/?p=gcc.git;a=commitdiff;h=a70d5d81c41048556fd86eaa1036018a6bfba115

[arsenm]

Merge conflict warning with D84199

[sammccall]

(Sorry, I don't know anything about this code)

[ab]

This looks sensible, modulo a couple inline comments

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

4891–4892

I don't know what model you have for the attribute and the command line options, and I guess you're probably aware, but this seems unsafe for depths > 0. I'm not sure there's a better way to deal with this, short of just forcing the strips in any aarch64 code, which is not an option. Maybe emitting strips when +pa is available? I imagine that's not at all reliable either.

4892

-> hasPA() ?

llvm/test/CodeGen/AArch64/GlobalISel/builtin-return-address-pacret.ll

2

can you explicitly pass -global-isel ?

[chill]

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

4891–4892

... I guess you're probably aware, but this seems unsafe for depths > 0.

In what sense is it unsafe? The clear instructions do not require a valid PAC
to be present, so even if a caller is not using PAC-RET we still get the return
address.

4892

I'll make a note to prepare a follow-up patch to use hasPA across the board.
There are a few more pending patches that likely use hasV8_3aOps, I want to fix them in a single patch.

llvm/test/CodeGen/AArch64/GlobalISel/builtin-return-address-pacret.ll

2

Kept -O0 so scheduling does not introduced uninteresting changes in the output.

[ab]

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

4891–4892

I'm thinking of the opposite scenario, where a caller in an unrelated library does have a signed return address, but this function wasn't compiled with that support: it wouldn't know to strip.

4892

If you have more patches please CC me, we have very extensive codegen support downstream, I can help with reviews until we upstream our changes ;) In the meantime you can take a look at our partial upstreaming in github apple/llvm-project

llvm/test/CodeGen/AArch64/GlobalISel/builtin-return-address-pacret.ll

2

Makes sense, that's fine. Can you also add -global-isel? ;) I'd rather not tie something as general as the optimization level with the specific path being tested: that can change over time, across cherry-picks, in different release branches, etc..

[chill]

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

4891–4892

I see. Looks like checking for +pa, while not perfect, will work in more cases.

llvm/test/CodeGen/AArch64/GlobalISel/builtin-return-address-pacret.ll

2

Right, I meant I added -global-isel, but also kept the -O0.

[danielkiss]

This comment was removed by danielkiss.

[danielkiss]

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

4891–4892

XPACLRI can always be emitted because it is in the NOP space. Besides it is not nice to emit always one extra presumably unnecessary instruction but I don't see why it won't work.

[chill]

I'm in favour of changing these patches to unconditionally emit XPACLRI (i.e. HINT #whatever).

[danielkiss]

XPACI could be emit when it is available, only the if (MF.getFunction().hasFnAttribute("sign-return-address")) { need do be dropped from the patches IMHO.

[chill]

Updated the patch to unconditionally include PAC stripping instructions, NOP-space ones to non-NON-space ones, depending on
whether v8.3-a (soo to be '+pa') feature.

[arsenm]

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

4894

Why do you need this temporary copy?

[chill]

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

4894

Nevermind, the patch is completely broken, I'll work on a new version.

[chill]

Well, it wasn't completely broken, just a little bit :D

We have virtual register that keeps LR, I would like for that vreg to keep the stripped version of LR, but
couldn't figure out how to get the proper place to insert the clearing insn, so instead there's a one clearing insn
per __builtin_fucntion_address "invocation".

[chill]

Ping?

[ab]

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

4891–4892

Yeah, it should work to always emit it. On darwin we know we don't need to, but I imagine it's the exception rather than the rule. I think we'll address that when merging this into our ptrauth changes: either emit XPACI (as a strip intrinsic) or nothing, based on the target architecture. Either way this LGTM.

[chill]

Thanks!

[aemerson]

Did you commit this?

[chill]

No, not yet.

[chill]

Rebased.

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

case Intrinsic::returnaddress: {
     MachineFunction &MF = *I.getParent()->getParent();
     MachineFrameInfo &MFI = MF.getFrameInfo();
 
     unsigned Depth = I.getOperand(2).getImm();
     Register DstReg = I.getOperand(0).getReg();
     RBI.constrainGenericRegister(DstReg, AArch64::GPR64RegClass, MRI);
 
     if (Depth == 0 && IntrinID == Intrinsic::returnaddress) {
-      if (MFReturnAddr) {
-        MIRBuilder.buildCopy({DstReg}, MFReturnAddr);
-        I.eraseFromParent();
-        return true;
-      }
-
-      MFI.setReturnAddressIsTaken(true);
-
+      if (!MFReturnAddr) {
         // Insert the copy from LR/X30 into the entry block, before it can be
         // clobbered by anything.
-      Register LiveInLR = getFunctionLiveInPhysReg(MF, TII, AArch64::LR,
-                                                   AArch64::GPR64spRegClass);
-      MIRBuilder.buildCopy(DstReg, LiveInLR);
+        MFI.setReturnAddressIsTaken(true);
+        MFReturnAddr = getFunctionLiveInPhysReg(MF, TII, AArch64::LR,
+                                                AArch64::GPR64RegClass);
+      }
+
+      if (STI.hasV8_3aOps()) {
+        MIRBuilder.buildInstr(AArch64::XPACI, {DstReg}, {MFReturnAddr});
+      } else {
+        MIRBuilder.buildCopy({Register(AArch64::LR)}, {MFReturnAddr});
+        MIRBuilder.buildInstr(AArch64::XPACLRI);
+        MIRBuilder.buildCopy({DstReg}, {Register(AArch64::LR)});
+      }
 
-      MFReturnAddr = LiveInLR;
       I.eraseFromParent();
       return true;
     }
 
     MFI.setFrameAddressIsTaken(true);
     Register FrameAddr(AArch64::FP);
     while (Depth--) {
       Register NextFrame = MRI.createVirtualRegister(&AArch64::GPR64spRegClass);
       auto Ldr =
           MIRBuilder.buildInstr(AArch64::LDRXui, {NextFrame}, {FrameAddr})
               .addImm(0);
       constrainSelectedInstRegOperands(*Ldr, TII, TRI, RBI);
       FrameAddr = NextFrame;
     }
 
     if (IntrinID == Intrinsic::frameaddress)
       MIRBuilder.buildCopy({DstReg}, {FrameAddr});
     else {
       MFI.setReturnAddressIsTaken(true);
-      MIRBuilder.buildInstr(AArch64::LDRXui, {DstReg}, {FrameAddr}).addImm(1);
+
+      if (STI.hasV8_3aOps()) {
+        Register TmpReg = MRI.createVirtualRegister(&AArch64::GPR64RegClass);
+        MIRBuilder.buildInstr(AArch64::LDRXui, {TmpReg}, {FrameAddr}).addImm(1);
+        MIRBuilder.buildInstr(AArch64::XPACI, {DstReg}, {TmpReg});
+      } else {
+        MIRBuilder.buildInstr(AArch64::LDRXui, {Register(AArch64::LR)}, {FrameAddr}).addImm(1);
+        MIRBuilder.buildInstr(AArch64::XPACLRI);
+        MIRBuilder.buildCopy({DstReg}, {Register(AArch64::LR)});
+      }
     }
 
     I.eraseFromParent();
     return true;
   }
   }
   return false;
 }

llvm/test/CodeGen/AArch64/GlobalISel/builtin-return-address-pacret.ll

+;; RUN: llc -mtriple aarch64               -global-isel -O0 %s -o - | FileCheck %s --check-prefixes=CHECK,CHECK-NOP
+;; RUN: llc -mtriple aarch64 -mattr=+v8.3a -global-isel -O0 %s -o - | FileCheck %s --check-prefixes=CHECK,CHECK-V83
+declare void @g0() #1
+declare void @g1(i8*) #1
+declare void @g2(i32, i8*) #1
+
+declare i8* @llvm.returnaddress(i32 immarg) #2
+
+define i8* @f0() #0 {
+entry:
+  %0 = call i8* @llvm.returnaddress(i32 0)
+  call void @g1(i8* %0)
+  %1 = call i8* @llvm.returnaddress(i32 1)
+  call void @g2(i32 1, i8* %1)
+  %2 = call i8* @llvm.returnaddress(i32 2)
+  ret i8* %2
+}
+;; CHECK-LABEL:    f0:
+;; CHECK-NOT:      {{(mov|ldr)}} x30
+;; CHECK-NOP:      hint #7
+;; CHECK-V83:      xpaci x30
+;; CHECK:          bl g1
+;; CHECK:          ldr x[[T0:[0-9]+]], [x29]
+;; CHECK-NOP-NEXT: ldr x30, [x[[T0]], #8]
+;; CHECK-NOP-NEXT: hint #7
+;; CHECK-V83-NEXT: ldr x[[T0]], [x[[T0]], #8]
+;; CHECK-V83-NEXT: xpaci x[[T0]]
+;; CHECK:          bl g2
+;; CHECK:          ldr x[[T1:[0-9]+]], [x29]
+;; CHECK-NEXT:     ldr x[[T1]], [x[[T1]]]
+;; CHECK-NOP-NEXT: ldr x30, [x[[T1]], #8]
+;; CHECK-NOP-NEXT: hint #7
+;; CHECK-NOP-NEXT: mov x0, x30
+;; CHECK-V83-NEXT: ldr x[[T1]], [x[[T1]], #8]
+;; CHECK-V83-NEXT: xpaci x[[T1]]
+;; CHECK-V83-NEXT: mov x0, x[[T1]]
+
+define i8* @f1() #0 {
+entry:
+  %0 = call i8* @llvm.returnaddress(i32 1)
+  call void @g1(i8* %0)
+  %1 = call i8* @llvm.returnaddress(i32 2)
+  call void @g2(i32 1, i8* %1)
+  %2 = call i8* @llvm.returnaddress(i32 0)
+  ret i8* %2
+}
+;; CHECK-LABEL:    f1:
+;; CHECK-DAG:      ldr x[[T0:[0-9]+]], [x29]
+;; CHECK-NOP-DAG:  str x30, [sp, #[[OFF:[0-9]+]]
+;; CHECK-NOP:      ldr x30, [x[[T0]], #8]
+;; CHECK-NOP-NEXT: hint #7
+;; CHECK-V83:      ldr x[[T0]], [x[[T0]], #8]
+;; CHECK-V83-NEXT: xpaci x[[T0]]
+;; CHECK-V83:      str x30, [sp, #[[OFF:[0-9]+]]
+;; CHECK:          bl g1
+;; CHECK:          ldr x[[T1:[0-9]+]], [x29]
+;; CHECK-NEXT:     ldr x[[T1]], [x[[T1]]]
+;; CHECK-NOP-NEXT: ldr x30, [x[[T1]], #8]
+;; CHECK-NOP-NEXT: hint #7
+;; CHECK-V83-NEXT: ldr x[[T1]], [x[[T1]], #8]
+;; CHECK-V83-NEXT: xpaci x[[T1]]
+;; CHECK:          bl g2
+;; CHECK:          ldr x[[T2:[0-9]+]], [sp, #[[OFF]]]
+;; CHECK-NOP-NEXT: mov x30, x[[T2]]
+;; CHECK-NOP-NEXT: hint #7
+;; CHECK-NOP-NEXT: mov x0, x30
+;; CHECK-V83-NEXT: xpaci x[[T2]]
+;; CHECK-V83-NEXT: mov x0, x[[T2]]
+;; CHECK-NOT:      x0
+;; CHECK:          ret
+
+define i8* @f2() #0 {
+entry:
+  call void bitcast (void ()* @g0 to void ()*)()
+  %0 = call i8* @llvm.returnaddress(i32 0)
+  ret i8* %0
+}
+;; CHECK-LABEL:    f2
+;; CHECK:          bl g0
+;; CHECK:          ldr x[[T0:[0-9]+]], [sp,
+;; CHECK-NOP-NEXT: mov x30, x[[T2]]
+;; CHECK-NOP-NEXT: hint #7
+;; CHECK-NOP-NEXT: mov x0, x30
+;; CHECK-V83-NEXT: xpaci x[[T2]]
+;; CHECK-V83-NEXT: mov x0, x[[T2]]
+;; CHECK-NOT:      x0
+;; CHECK:          ret
+
+define i8* @f3() #0 {
+entry:
+  %0 = call i8* @llvm.returnaddress(i32 0)
+  ret i8* %0
+}
+;; CHECK-LABEL:    f3:
+;; CHECK:          str x30, [sp,
+;; CHECK-NOP-NEXT: hint #7
+;; CHECK-V83-NEXT: xpaci x30
+;; CHECK-NEXT:     mov x0, x30
+;; CHECK-NOT:      x0
+;; CHECK:          ret
+attributes #0 = { nounwind }
+attributes #1 = { nounwind }
+attributes #2 = { nounwind readnone }

llvm/test/CodeGen/AArch64/GlobalISel/select-returnaddr.ll

 ; RUN: llc -mtriple=arm64-apple-ios -global-isel -o - %s | FileCheck %s
 
 define i8* @rt0(i32 %x) nounwind readnone {
 entry:
 ; CHECK-LABEL: rt0:
-; CHECK-NOT: stp
-; CHECK: mov x0, x30
+; CHECK:       hint #7
+; CHECK-NEXT:  mov x0, x30
   %0 = tail call i8* @llvm.returnaddress(i32 0)
   ret i8* %0
 }
 
 define i8* @rt0_call_clobber(i32 %x) nounwind readnone {
 entry:
 ; CHECK-LABEL: rt0_call_clobber:
-; CHECK: stp     x20, x19, [sp, #-32]!
-; CHECK: stp     x29, x30, [sp, #16]
-; CHECK: mov     x19, x30
-; CHECK: bl      _foo
-; CHECK: ldp     x29, x30, [sp, #16]
-; CHECK: mov     x0, x19
+; CHECK:       stp x20, x19, [sp, #-32]!
+; CHECK:       stp x29, x30, [sp, #16]
+; CHECK:       mov x19, x30
+; CHECK:       bl _foo
+; CHECK:       mov x30, x19
+; CHECK-NEXT:  hint #7
+; CHECK-NEXT:  mov x0, x30
+; CHECK-NOT:   x0
+; CHECK:       ret
   %ret = call i32 @foo()
   %0 = tail call i8* @llvm.returnaddress(i32 0)
   ret i8* %0
 }
 
 define i8* @rt2() nounwind readnone {
 entry:
 ; CHECK-LABEL: rt2:
-; CHECK: ldr x[[reg:[0-9]+]], [x29]
-; CHECK: ldr x[[reg]], [x[[reg]]]
-; CHECK: ldr x0, [x[[reg]], #8]
+; CHECK:       ldr x[[reg:[0-9]+]], [x29]
+; CHECK:       ldr x[[reg]], [x[[reg]]]
+; CHECK:       ldr x30, [x[[reg]], #8]
+; CHECK:       hint #7
+; CHECK:       mov x0, x30
+; CHECK-NOT:   x0
+; CHECK:       ret
   %0 = tail call i8* @llvm.returnaddress(i32 2)
   ret i8* %0
 }
 
 
 declare i32 @foo()
 declare i8* @llvm.returnaddress(i32) nounwind readnone

llvm/test/CodeGen/AArch64/GlobalISel/select-returnaddress-liveins.mir

 legalized:       true
 regBankSelected: true
 tracksRegLiveness: true
 body:             |
   ; CHECK-LABEL: name: lr_other_block
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.1(0x80000000)
   ; CHECK:   liveins: $w0, $x0, $lr
-  ; CHECK:   [[COPY:%[0-9]+]]:gpr64sp = COPY $lr
+  ; CHECK:   [[COPY:%[0-9]+]]:gpr64 = COPY $lr
   ; CHECK: bb.1:
-  ; CHECK:   [[COPY1:%[0-9]+]]:gpr64 = COPY [[COPY]]
+  ; CHECK:   $lr = COPY [[COPY]]
+  ; CHECK:   XPACLRI implicit-def $lr, implicit $lr
+  ; CHECK:   [[COPY1:%[0-9]+]]:gpr64 = COPY $lr
   ; CHECK:   $x0 = COPY [[COPY1]]
   ; CHECK:   RET_ReallyLR implicit $x0
   ; LR should be added as a livein to the entry block.
 
   bb.0:
     ; We should have lr as a livein to bb.0, and a copy from LR.
     liveins: $w0, $x0
     G_BR %bb.1
   bb.1:
     %4:gpr(p0) = G_INTRINSIC intrinsic(@llvm.returnaddress), 0
     $x0 = COPY %4
     RET_ReallyLR implicit $x0
 ...
 ---
 name:            already_live_in
 alignment:       4
 legalized:       true
 regBankSelected: true
 tracksRegLiveness: true
 body:             |
   ; CHECK-LABEL: name: already_live_in
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.1(0x80000000)
   ; CHECK:   liveins: $w0, $x0, $lr
-  ; CHECK:   [[COPY:%[0-9]+]]:gpr64sp = COPY $lr
+  ; CHECK:   [[COPY:%[0-9]+]]:gpr64 = COPY $lr
   ; CHECK: bb.1:
-  ; CHECK:   [[COPY1:%[0-9]+]]:gpr64 = COPY [[COPY]]
+  ; CHECK:   $lr = COPY [[COPY]]
+  ; CHECK:   XPACLRI implicit-def $lr, implicit $lr
+  ; CHECK:   [[COPY1:%[0-9]+]]:gpr64 = COPY $lr
   ; CHECK:   $x0 = COPY [[COPY1]]
   ; CHECK:   RET_ReallyLR implicit $x0
   ; We should not have LR listed as a livein twice.
 
   bb.0:
     liveins: $w0, $x0, $lr
     G_BR %bb.1
   bb.1:

 legalized:       true
 regBankSelected: true
 tracksRegLiveness: true
 body:             |
   ; CHECK-LABEL: name: multi_use
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.1(0x80000000)
   ; CHECK:   liveins: $w0, $x0, $lr
-  ; CHECK:   [[COPY:%[0-9]+]]:gpr64sp = COPY $lr
-  ; CHECK:   [[COPY1:%[0-9]+]]:gpr64 = COPY [[COPY]]
+  ; CHECK:   [[COPY:%[0-9]+]]:gpr64 = COPY $lr
+  ; CHECK:   $lr = COPY [[COPY]]
+  ; CHECK:   XPACLRI implicit-def $lr, implicit $lr
+  ; CHECK:   [[COPY1:%[0-9]+]]:gpr64 = COPY $lr
   ; CHECK: bb.1:
   ; CHECK:   $x0 = COPY [[COPY1]]
-  ; CHECK:   [[COPY2:%[0-9]+]]:gpr64 = COPY [[COPY]]
+  ; CHECK:   $lr = COPY [[COPY]]
+  ; CHECK:   XPACLRI implicit-def $lr, implicit $lr
+  ; CHECK:   [[COPY2:%[0-9]+]]:gpr64 = COPY $lr
   ; CHECK:   RET_ReallyLR implicit [[COPY2]]
   bb.0:
     liveins: $w0, $x0, $lr
     %0:gpr(p0) = G_INTRINSIC intrinsic(@llvm.returnaddress), 0
     G_BR %bb.1
   bb.1:
     $x0 = COPY %0
     %1:gpr(p0) = G_INTRINSIC intrinsic(@llvm.returnaddress), 0
     RET_ReallyLR implicit %1
 ...