[AArch64][PAC] Select XPAC for ptrauth.strip intrinsic.

[ab]

Summary

This is mostly straightforward, with an additional optimization: XPACIuntied allows strips of LR (which we emit in __builtin_return_address) to avoid clobbering LR. (This can be useful, e.g., for a certain kind of wrapper functions that just forward the caller address, for logging and such)

[tschuett]

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

5873

Can the Key be 2? Is an if-statement easier to comprehend than a ternary operator?

[apazos]

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

1227 ↗

(On Diff #454524)

there is no unit test case for this scenario

[kristof.beyls]

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

5873

I guess having an enum for the valid key descriptors for AArch64 would make this and other code working with keys more readable. I'm not sure if that's typically done or where to define that enum?
Than in this location, maybe a switch on that enum would be clearer?

I see these constants 0,1,2,3 are also needed in tablegen files. Not sure if something similar to an enum exists in Tablegen?

[ab]

• Extract out XPACIuntied, leaving it for ptrauth-returns.
• Pull in key/opcode helper functions from later patch in the series.

[ab]

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

1227 ↗

(On Diff #454524)

Yeah, XPACIuntied is really intended to be used in @llvm.returnaddress when return address signing (/PAC-RET) is enabled. I pulled it out of this patch into that one, where it actually gets exercised.

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

5873

A later patch in the series actually does introduce helper functions in AArch64BaseInfo.h/AArch64InstrInfo.h, so it's a simple matter of using them here! I combined the two in this review, to be committed separately.

[kristof.beyls]

LGTM!
I just added a few minor nitpicks - I'll leave it to you to decide if following those suggestions would result in an improvement or not.

llvm/lib/Target/AArch64/AArch64InstrInfo.h

514 ↗

(On Diff #456771)

nitpick, maybe do s/ptrauth auth/pointer authentication/?

526 ↗

(On Diff #456771)

nitpick, maybe do s/ptrauth auth/pointer signing/?

llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h

798 ↗

(On Diff #456771)

I guess that if we added "LAST = 3" here, the somewhat cryptic if statement in AArch64InstructionSelector.cpp that says

if (Key > 3)
      return false;

could instead say something like

if (AArch64PACKey::LAST > 3)
      return false;

It seems the word "LAST" is already used in a few other places in LLVM to indicate the highest defined value in an enum. Maybe doing this change would be a minor improvement in readability/maintainability?

llvm/lib/Target/AArch64/AArch64InstrInfo.td

multiclass SignAuth<bits<3> prefix, bits<3> prefix_z, string asm,
     def IZB  : SignAuthZero<prefix_z,  0b01, !strconcat(asm, "izb"), op>;
     def DZB  : SignAuthZero<prefix_z,  0b11, !strconcat(asm, "dzb"), op>;
   }
 
   defm PAC : SignAuth<0b000, 0b010, "pac", int_ptrauth_sign>;
   defm AUT : SignAuth<0b001, 0b011, "aut", null_frag>;
 
   def XPACI : ClearAuth<0, "xpaci">;
+  def : Pat<(int_ptrauth_strip GPR64:$Rd, 0), (XPACI GPR64:$Rd)>;
+  def : Pat<(int_ptrauth_strip GPR64:$Rd, 1), (XPACI GPR64:$Rd)>;
+
   def XPACD : ClearAuth<1, "xpacd">;
+  def : Pat<(int_ptrauth_strip GPR64:$Rd, 2), (XPACD GPR64:$Rd)>;
+  def : Pat<(int_ptrauth_strip GPR64:$Rd, 3), (XPACD GPR64:$Rd)>;
 
   def PACGA : SignAuthTwoOperand<0b1100, "pacga", int_ptrauth_sign_generic>;
 
   // Combined Instructions
   let isBranch = 1, isTerminator = 1, isBarrier = 1, isIndirectBranch = 1  in {
     def BRAA    : AuthBranchTwoOperands<0, 0, "braa">;
     def BRAB    : AuthBranchTwoOperands<0, 1, "brab">;
   }

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

if (!IsDiscZero) {
       PAC.addUse(DiscReg);
       RBI.constrainGenericRegister(DiscReg, AArch64::GPR64spRegClass, MRI);
     }
 
     RBI.constrainGenericRegister(DstReg, AArch64::GPR64RegClass, MRI);
     I.eraseFromParent();
     return true;
   }
+  case Intrinsic::ptrauth_strip: {
+    Register DstReg = I.getOperand(0).getReg();
+    Register ValReg = I.getOperand(2).getReg();
+    uint64_t Key = I.getOperand(3).getImm();
+
+    if (Key > AArch64PACKey::LAST)
+      return false;
+    unsigned Opcode = getXPACOpcodeForKey((AArch64PACKey::ID)Key);
+
+    MIB.buildInstr(Opcode, {DstReg}, {ValReg});
+
+    RBI.constrainGenericRegister(DstReg, AArch64::GPR64RegClass, MRI);
+    RBI.constrainGenericRegister(ValReg, AArch64::GPR64RegClass, MRI);
+    I.eraseFromParent();
+    return true;
+  }
   case Intrinsic::frameaddress:
   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);

llvm/test/CodeGen/AArch64/ptrauth-intrinsic-strip.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple arm64e-apple-darwin -verify-machineinstrs -global-isel=0 | FileCheck %s
+; RUN: llc < %s -mtriple arm64e-apple-darwin -verify-machineinstrs -global-isel=1 -global-isel-abort=1 | FileCheck %s
+
+target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+
+define i64 @test_strip_ia(i64 %arg) {
+; CHECK-LABEL: test_strip_ia:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    xpaci x0
+; CHECK-NEXT:    ret
+  %tmp = call i64 @llvm.ptrauth.strip(i64 %arg, i32 0)
+  ret i64 %tmp
+}
+
+define i64 @test_strip_ib(i64 %arg) {
+; CHECK-LABEL: test_strip_ib:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    xpaci x0
+; CHECK-NEXT:    ret
+  %tmp = call i64 @llvm.ptrauth.strip(i64 %arg, i32 1)
+  ret i64 %tmp
+}
+
+define i64 @test_strip_da(i64 %arg) {
+; CHECK-LABEL: test_strip_da:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    xpacd x0
+; CHECK-NEXT:    ret
+  %tmp = call i64 @llvm.ptrauth.strip(i64 %arg, i32 2)
+  ret i64 %tmp
+}
+
+define i64 @test_strip_db(i64 %arg) {
+; CHECK-LABEL: test_strip_db:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    xpacd x0
+; CHECK-NEXT:    ret
+  %tmp = call i64 @llvm.ptrauth.strip(i64 %arg, i32 3)
+  ret i64 %tmp
+}
+
+declare i64 @llvm.ptrauth.strip(i64, i32)