[llvm] [AArch64][PAC] Refine authenticated pointer check methods (PR #74074)

[Anatoly Trosinenko via llvm-commits]

https://github.com/atrosinenko updated https://github.com/llvm/llvm-project/pull/74074

>From e07476e1bef713a75f8a9be4a327976c9a79eb89 Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko <atrosinenko at accesssoftek.com>
Date: Wed, 29 Nov 2023 16:21:30 +0300
Subject: [PATCH] [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.
---
 .../lib/Target/AArch64/AArch64PointerAuth.cpp | 14 +++++++++++---
 .../AArch64/sign-return-address-tailcall.ll   | 19 +++++++++++++++----
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
index 7509afaeb5fef..91d97cb376303 100644
--- a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
+++ b/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;
@@ -238,7 +242,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;
@@ -314,6 +318,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)
@@ -354,7 +362,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;
 }
diff --git a/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll b/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll
index ec04e553cac6e..cf033cb8208cc 100644
--- a/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll
+++ b/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()