[AArch64][PAC] Refine authenticated pointer check methods

Align the values of the immediate operand of BRK instruction with those
used by the existing arm64e implementation.

Make AuthCheckMethod::DummyLoad use the requested register instead of LR.

Author: atrosinenko

llvm/lib/Target/AArch64/AArch64PointerAuth.cpp

@@ -12,6 +12,7 @@
 #include "AArch64InstrInfo.h"
 #include "AArch64MachineFunctionInfo.h"
 #include "AArch64Subtarget.h"
+#include "Utils/AArch64BaseInfo.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
@@ -35,7 +36,10 @@ class AArch64PointerAuth : public MachineFunctionPass {
 
 private:
   /// An immediate operand passed to BRK instruction, if it is ever emitted.
-  const unsigned BrkOperand = 0xc471;
+  static unsigned BrkOperandForKey(AArch64PACKey::ID KeyId) {
+    const unsigned BrkOperandBase = 0xc470;
+    return BrkOperandBase + KeyId;
+  }
 
   const AArch64Subtarget *Subtarget = nullptr;
   const AArch64InstrInfo *TII = nullptr;
@@ -174,7 +178,7 @@ MachineBasicBlock &llvm::AArch64PAuth::checkAuthenticatedRegister(
     return MBB;
   case AuthCheckMethod::DummyLoad:
     BuildMI(MBB, MBBI, DL, TII->get(AArch64::LDRWui), getWRegFromXReg(TmpReg))
-        .addReg(AArch64::LR)
+        .addReg(AuthenticatedReg)
         .addImm(0)
         .addMemOperand(createCheckMemOperand(MF, Subtarget));
     return MBB;
@@ -250,6 +254,10 @@ unsigned llvm::AArch64PAuth::getCheckerSizeInBytes(AuthCheckMethod Method) {
 
 bool AArch64PointerAuth::checkAuthenticatedLR(
     MachineBasicBlock::iterator TI) const {
+  const AArch64FunctionInfo *MFnI = TI->getMF()->getInfo<AArch64FunctionInfo>();
+  AArch64PACKey::ID KeyId =
+      MFnI->shouldSignWithBKey() ? AArch64PACKey::IB : AArch64PACKey::IA;
+
   AuthCheckMethod Method = Subtarget->getAuthenticatedLRCheckMethod();
 
   if (Method == AuthCheckMethod::None)
@@ -290,7 +298,7 @@ bool AArch64PointerAuth::checkAuthenticatedLR(
          "More than a single register is used by TCRETURN");
 
   checkAuthenticatedRegister(TI, Method, AArch64::LR, TmpReg, /*UseIKey=*/true,
-                             BrkOperand);
+                             BrkOperandForKey(KeyId));
 
   return true;
 }

llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll

@@ -23,7 +23,7 @@ define i32 @tailcall_direct() "sign-return-address"="non-leaf" {
 ;
 ; COMMON-NEXT:    b callee
 ; BRK-NEXT:     .[[FAIL]]:
-; BRK-NEXT:       brk #0xc471
+; BRK-NEXT:       brk #0xc470
   tail call void asm sideeffect "", "~{lr}"()
   %call = tail call i32 @callee()
   ret i32 %call
@@ -48,7 +48,7 @@ define i32 @tailcall_indirect(ptr %fptr) "sign-return-address"="non-leaf" {
 ;
 ; COMMON-NEXT:    br x0
 ; BRK-NEXT:     .[[FAIL]]:
-; BRK-NEXT:       brk #0xc471
+; BRK-NEXT:       brk #0xc470
   tail call void asm sideeffect "", "~{lr}"()
   %call = tail call i32 %fptr()
   ret i32 %call
@@ -89,7 +89,7 @@ define i32 @tailcall_direct_noframe_sign_all() "sign-return-address"="all" {
 ;
 ; COMMON-NEXT:    b callee
 ; BRK-NEXT:     .[[FAIL]]:
-; BRK-NEXT:       brk #0xc471
+; BRK-NEXT:       brk #0xc470
   %call = tail call i32 @callee()
   ret i32 %call
 }
@@ -113,9 +113,20 @@ define i32 @tailcall_indirect_noframe_sign_all(ptr %fptr) "sign-return-address"=
 ;
 ; COMMON-NEXT:    br x0
 ; BRK-NEXT:     .[[FAIL]]:
-; BRK-NEXT:       brk #0xc471
+; BRK-NEXT:       brk #0xc470
   %call = tail call i32 %fptr()
   ret i32 %call
 }
 
+define i32 @tailcall_ib_key() "sign-return-address"="all" "sign-return-address-key"="b_key" {
+; COMMON-LABEL: tailcall_ib_key:
+;
+; COMMON:         b callee
+; BRK-NEXT:     .{{LBB.*}}:
+; BRK-NEXT:       brk #0xc471
+  tail call void asm sideeffect "", "~{lr}"()
+  %call = tail call i32 @callee()
+  ret i32 %call
+}
+
 declare i32 @callee()