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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-aarch64

Author: Anatoly Trosinenko (atrosinenko)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/74074.diff


2 Files Affected:

- (modified) llvm/lib/Target/AArch64/AArch64PointerAuth.cpp (+11-3) 
- (modified) llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll (+15-4) 


``````````diff
diff --git a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
index 7576d2a899d1afb..8597700d5e6f080 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;
@@ -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;
 }
diff --git a/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll b/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll
index ec04e553cac6e37..cf033cb8208cc90 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()

``````````

</details>


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