[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

4720

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.

4721

-> hasPA() ?

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

1

can you explicitly pass -global-isel ?

[chill]

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

4720

... 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.

4721

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

1

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

[ab]

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

4720

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.

4721

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

1

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

4720

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

1

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

4720

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

4716

Why do you need this temporary copy?

[chill]

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

4716

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

4720

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

if (Depth == 0 && IntrinID == Intrinsic::returnaddress) {
         I.eraseFromParent();
         return true;
       }
       MFI.setReturnAddressIsTaken(true);
       MF.addLiveIn(AArch64::LR, &AArch64::GPR64spRegClass);
       // Insert the copy from LR/X30 into the entry block, before it can be
       // clobbered by anything.
       MachineBasicBlock &EntryBlock = *MF.begin();
       if (!EntryBlock.isLiveIn(AArch64::LR))
         EntryBlock.addLiveIn(AArch64::LR);
       MachineIRBuilder EntryBuilder(MF);
       EntryBuilder.setInstr(*EntryBlock.begin());
+      if (MF.getFunction().hasFnAttribute("sign-return-address")) {
+        if (STI.hasV8_3aOps()) {
+          Register TmpReg = MRI.createVirtualRegister(&AArch64::GPR64RegClass);
+          EntryBuilder.buildCopy({TmpReg}, {Register(AArch64::LR)});
+          EntryBuilder.buildInstr(AArch64::XPACI, {DstReg}, {TmpReg});
+        } else {
+          EntryBuilder.buildInstr(AArch64::XPACLRI);
+          EntryBuilder.buildCopy({DstReg}, {Register(AArch64::LR)});
+        }
+      } else {
         EntryBuilder.buildCopy({DstReg}, {Register(AArch64::LR)});
+      }
+
       MFReturnAddr = DstReg;
       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);
+
+      if (MF.getFunction().hasFnAttribute("sign-return-address")) {
+        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)});
+        }
+      } else {
         MIRBuilder.buildInstr(AArch64::LDRXui, {DstReg}, {FrameAddr}).addImm(1);
       }
+    }
 
     I.eraseFromParent();
     return true;
   }
   }
   return false;
 }
 

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

+;; RUN: llc -mtriple aarch64               -O0 %s -o - | FileCheck %s --check-prefixes=CHECK,CHECK-NOP
+;; RUN: llc -mtriple aarch64 -mattr=+v8.3a -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-V83-NEXT: ldr x[[T1]], [x[[T1]], #8]
+;; CHECK-V83-NEXT: xpaci 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-NOT:      {{(mov|ldr)}} x30
+;; CHECK-NOP:      hint #7
+;; CHECK-V83:      xpaci x30
+;; CHECK-NOP-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:[0-9]+]], [x29]
+;; CHECK-V83-NEXT: 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-NOP-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]]]
+;; CHECK-V83-NEXT: ldr x[[T1]], [x[[T1]], #8]
+;; CHECK-V83-NEXT: xpaci x[[T1]]
+;; CHECK:          ldr x0, [sp, #[[OFF]]]
+;; 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:          str x30, [sp,
+;; CHECK-NOP-NEXT: hint #7
+;; CHECK-V83-NEXT: xpaci x30
+;; CHECK-NEXT:     str x30, [sp, #[[OFF:[0-9]+]]]
+;; CHECK:          bl g0
+;; CHECK:          ldr x0, [sp, #[[OFF]]]
+;; 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,
+;; CHECK-NOP-NEXT: hint #7
+;; CHECK-V83-NEXT: xpaci x30
+;; CHECK-NEXT:     mov x0, x30
+;; CHECK-NOT:      x0
+;; CHECK:          ret
+
+attributes #0 = { nounwind "sign-return-address"="all" "sign-return-address-key"="a_key" }
+attributes #1 = { nounwind }
+attributes #2 = { nounwind readnone }