[llvm] [AArch64][PAC] Implement code generation for @llvm.ptrauth.auth (PR #72502)

via llvm-commits llvm-commits at lists.llvm.org
Thu Nov 16 03:33:25 PST 2023


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-aarch64

Author: Anatoly Trosinenko (atrosinenko)

<details>
<summary>Changes</summary>

This patch introduces PAUTH_AUTH pseudo instruction that can encode well-known discriminator computations in its operands:
- small immediate integer discriminator
- blend of a GPR64 register and an immediate integer
- arbitrary GPR64 register as a fallback

For convenience, a PAUTH_BLEND pseudo instruction is introduced as well that is selected for @<!-- -->llvm.ptrauth.blend intrinsic.

For @<!-- -->llvm.ptrauth.auth, the TableGen-erated code selects a PAUTH_AUTH instruction in its "fallback" form. After that, custom inserter tries to detect a well-known signing schema and refines the operands of PAUTH_AUTH instruction, if possible.

It may be necessary to use fixed X17 and X16 for pointer and scratch registers, either for security or compatibility with Armv8.2. For that purpose, implicit defs of X16 and X17 are added by TableGen-erated code, to make sure that custom inserter can safely use these registers as pointer and scratch operands. These temporary implicit-def operands are removed by custom inserter.

As it is worth asking register allocator to place $reg_disc right in the $scratch register, these operands are tied. Thus, as a special case it is permitted to assign XZR to $scratch and provide the real scratch register as an implicit-def operand. While it would be possible to use 2 separate pseudo instructions: one for immediate integer discriminator and one for everything else, it would require 2*2 pseudos for expressing @<!-- -->llvm.ptrauth.resign the same way (or even 3*3 if clearly separating register/immediate/blended discriminator cases).

---

Patch is 47.47 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/72502.diff


9 Files Affected:

- (modified) llvm/lib/Target/AArch64/AArch64ISelLowering.cpp (+127) 
- (modified) llvm/lib/Target/AArch64/AArch64ISelLowering.h (+3) 
- (modified) llvm/lib/Target/AArch64/AArch64InstrInfo.td (+26-3) 
- (modified) llvm/lib/Target/AArch64/AArch64PointerAuth.cpp (+224-3) 
- (modified) llvm/lib/Target/AArch64/AArch64PointerAuth.h (+19) 
- (modified) llvm/lib/Target/AArch64/AArch64Subtarget.cpp (+42) 
- (modified) llvm/lib/Target/AArch64/AArch64Subtarget.h (+6) 
- (added) llvm/test/CodeGen/AArch64/ptrauth-intrinsic-auth.ll (+501) 
- (modified) llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll (+4-4) 


``````````diff
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 9ff6d6f0f565edb..0e3ef3ffb3410e1 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -2801,6 +2801,130 @@ AArch64TargetLowering::EmitZero(MachineInstr &MI, MachineBasicBlock *BB) const {
   return BB;
 }
 
+MachineBasicBlock *
+AArch64TargetLowering::EmitPAuthInstr(MachineInstr &MI,
+                                      MachineBasicBlock *BB) const {
+  const AArch64InstrInfo *TII = Subtarget->getInstrInfo();
+  MachineRegisterInfo &MRI = BB->getParent()->getRegInfo();
+  DebugLoc DL = MI.getDebugLoc();
+
+  // The discriminator should be expressed by consecutive operands
+  // (raw_register, immediate_integer, is_blend). This function accepts
+  // (reg, 0, 0) operands generated by the instruction selector and tries
+  // to detect either small immediate discriminator expressed as
+  // (XZR, small_int, 0), blend(something, small_int) expressed as
+  // (something, small_int, 1) or keeps the operands as-is otherwise.
+  auto DetectDiscriminator = [&](unsigned RegDiscOpIndex) {
+    MachineOperand &RegOp = MI.getOperand(RegDiscOpIndex);
+    MachineOperand &ImmOp = MI.getOperand(RegDiscOpIndex + 1);
+    MachineOperand &IsBlendOp = MI.getOperand(RegDiscOpIndex + 2);
+    assert(ImmOp.getImm() == 0 && "Operand was already initialized");
+    assert(IsBlendOp.getImm() == 0 && "Operand was already initialized");
+
+    // Walk through the chain of copy-like instructions until we find
+    // a known signing schema, if any.
+    Register AddrDisc;
+    uint64_t ImmDisc;
+    for (Register DiscReg = RegOp.getReg(); DiscReg.isVirtual();) {
+      MachineInstr *DefiningMI = MRI.getVRegDef(DiscReg);
+      switch (DefiningMI->getOpcode()) {
+      case AArch64::COPY:
+        DiscReg = DefiningMI->getOperand(1).getReg();
+        if (DiscReg == AArch64::XZR) {
+          // Zero discriminator: (XZR, 0, 0).
+          RegOp.setReg(AArch64::XZR);
+          return;
+        }
+        break;
+      case AArch64::SUBREG_TO_REG:
+        DiscReg = DefiningMI->getOperand(2).getReg();
+        break;
+      case AArch64::MOVi32imm:
+        ImmDisc = DefiningMI->getOperand(1).getImm();
+        // If ImmDisc does not fit in 16 bits,
+        // consider it as custom computation.
+        if ((ImmDisc & 0xffff) == ImmDisc) {
+          // Small immediate integer: (XZR, imm, 0).
+          RegOp.setReg(AArch64::XZR);
+          ImmOp.setImm(ImmDisc);
+        }
+        return;
+      case AArch64::PAUTH_BLEND:
+        AddrDisc = DefiningMI->getOperand(1).getReg();
+        ImmDisc = DefiningMI->getOperand(2).getImm();
+        assert((ImmDisc & 0xffff) == ImmDisc &&
+               "Expected 16-bit integer operand in PAUTH_BLEND");
+        RegOp.setReg(AddrDisc);
+        ImmOp.setImm(ImmDisc);
+        IsBlendOp.setImm(1);
+        return;
+      default:
+        // Custom computation, leave it as-is.
+        return;
+      }
+    }
+  };
+
+  auto PopImplicitDef = [&](Register ExpectedReg) {
+    (void)ExpectedReg;
+    unsigned Index = MI.getNumOperands() - 1;
+    assert(MI.getOperand(Index).isImplicit());
+    assert(MI.getOperand(Index).isDef());
+    assert(MI.getOperand(Index).getReg() == ExpectedReg);
+    MI.removeOperand(Index);
+  };
+
+  auto AdjustDefinedRegisters = [&](unsigned TiedRegDiscOpIndex) {
+    Register RegDisc = MI.getOperand(TiedRegDiscOpIndex).getReg();
+
+    // The instruction, as selected by TableGen-erated code, has X16 and X17
+    // registers implicitly defined, to make sure they are safe to clobber.
+    //
+    // Remove these generic implicit defs here and re-add them as needed and
+    // if needed. If assertions are enabled, additionally check that the two
+    // implicit operands are the expected ones.
+    PopImplicitDef(AArch64::X17);
+    PopImplicitDef(AArch64::X16);
+
+    // $scratch operand is tied to $reg_disc, thus if an immediate integer
+    // discriminator is used, $scratch ends up being XZR. In that case, add
+    // an implicit-def scratch register - this is a special case known by
+    // aarch64-ptrauth pass.
+    MachineOperand *RealScratchOp = &MI.getOperand(1);
+    if (RegDisc == AArch64::XZR) {
+      MI.getOperand(1).setReg(AArch64::XZR);
+      Register ScratchReg = MRI.createVirtualRegister(&AArch64::GPR64RegClass);
+      MI.addOperand(MachineOperand::CreateReg(ScratchReg, /*isDef=*/true,
+                                              /*isImp=*/true));
+      RealScratchOp = &MI.getOperand(MI.getNumOperands() - 1);
+    }
+
+    assert((RegDisc == AArch64::XZR || RegDisc.isVirtual()) &&
+           "Accidentally clobbering register?");
+
+    // If target CPU does not support FEAT_PAuth, IA and IB keys are still
+    // usable via HINT-encoded instructions.
+    if (!Subtarget->hasPAuth()) {
+      Register AutedReg = MI.getOperand(0).getReg();
+
+      MI.getOperand(0).setReg(AArch64::X17);
+      RealScratchOp->setReg(AArch64::X16);
+      BuildMI(*BB, MI.getNextNode(), DL, TII->get(AArch64::COPY), AutedReg)
+          .addReg(AArch64::X17);
+    }
+  };
+
+  switch (MI.getOpcode()) {
+  default:
+    llvm_unreachable("Unhandled opcode");
+  case AArch64::PAUTH_AUTH:
+    DetectDiscriminator(/*RegDiscOpIndex=*/3);
+    AdjustDefinedRegisters(/*TiedRegDiscOpIndex=*/3);
+    break;
+  }
+  return BB;
+}
+
 MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
     MachineInstr &MI, MachineBasicBlock *BB) const {
 
@@ -2854,6 +2978,9 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
   case TargetOpcode::PATCHABLE_TYPED_EVENT_CALL:
     return BB;
 
+  case AArch64::PAUTH_AUTH:
+    return EmitPAuthInstr(MI, BB);
+
   case AArch64::CATCHRET:
     return EmitLoweredCatchRet(MI, BB);
   case AArch64::LD1_MXIPXX_H_PSEUDO_B:
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.h b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
index b638084f98dadb8..0ab8b113d3bbede 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.h
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
@@ -612,6 +612,9 @@ class AArch64TargetLowering : public TargetLowering {
                                      bool IsSpill) const;
   MachineBasicBlock *EmitZero(MachineInstr &MI, MachineBasicBlock *BB) const;
 
+  MachineBasicBlock *EmitPAuthInstr(MachineInstr &MI,
+                                    MachineBasicBlock *BB) const;
+
   MachineBasicBlock *
   EmitInstrWithCustomInserter(MachineInstr &MI,
                               MachineBasicBlock *MBB) const override;
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
index 290c79f7bacdb8f..f69d27c1269eba7 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.td
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
@@ -1533,6 +1533,27 @@ def PAUTH_PROLOGUE : Pseudo<(outs), (ins), []>, Sched<[]>;
 def PAUTH_EPILOGUE : Pseudo<(outs), (ins), []>, Sched<[]>;
 }
 
+def PAUTH_BLEND : Pseudo<(outs GPR64:$disc), (ins GPR64:$addr_disc, i32imm:$int_disc), []>, Sched<[]>;
+
+// Two tasks are handled by custom inserter:
+// 1. It tries to detect known signing schemas: either small immediate integer
+//    discriminator or an arbitrary register blended with a small integer -
+//    if such schema is detected, it is saved into the instruction's operands.
+// 2. It is worth to reuse $reg_disc as a scratch register unless we use
+//    immediate integer as a discriminator (in that case $reg_disc is XZR).
+//    In the latter case $scratch is technically XZR, but another def-ed
+//    register is added as an implicit operand by the inserter.
+//
+// See the comments in custom inserter code for explanation of the reason
+// to specify "Defs = [X16, X17]" here.
+let usesCustomInserter = 1, Defs = [X16, X17] in {
+def PAUTH_AUTH : Pseudo<(outs GPR64common:$auted, GPR64:$scratch),
+                        (ins GPR64common:$signed,
+                             GPR64:$reg_disc, i32imm:$int_disc,
+                             i32imm:$is_blended, i32imm:$key_id), [],
+                        "$auted = $signed, $scratch = $reg_disc">, Sched<[]>;
+}
+
 // These pointer authentication instructions require armv8.3a
 let Predicates = [HasPAuth] in {
 
@@ -9146,12 +9167,14 @@ let Predicates = [HasMOPS, HasMTE], Defs = [NZCV], Size = 12, mayLoad = 0, maySt
 //-----------------------------------------------------------------------------
 // v8.3 Pointer Authentication late patterns
 
-let Predicates = [HasPAuth] in {
 def : Pat<(int_ptrauth_blend GPR64:$Rd, imm64_0_65535:$imm),
-          (MOVKXi GPR64:$Rd, (trunc_imm imm64_0_65535:$imm), 48)>;
+          (PAUTH_BLEND GPR64:$Rd, (trunc_imm imm64_0_65535:$imm))>;
 def : Pat<(int_ptrauth_blend GPR64:$Rd, GPR64:$Rn),
           (BFMXri GPR64:$Rd, GPR64:$Rn, 16, 15)>;
-}
+
+def : Pat<(int_ptrauth_auth GPR64:$signed,
+                            timm:$key_id, GPR64:$reg_disc),
+          (PAUTH_AUTH GPR64:$signed, GPR64:$reg_disc, 0, 0, timm:$key_id)>;
 
 //-----------------------------------------------------------------------------
 
diff --git a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
index 7576d2a899d1afb..8ae38af4a3b43d2 100644
--- a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
+++ b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
@@ -12,9 +12,11 @@
 #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"
+#include "llvm/IR/Intrinsics.h"
 
 using namespace llvm;
 using namespace llvm::AArch64PAuth;
@@ -34,8 +36,8 @@ class AArch64PointerAuth : public MachineFunctionPass {
   StringRef getPassName() const override { return AARCH64_POINTER_AUTH_NAME; }
 
 private:
-  /// An immediate operand passed to BRK instruction, if it is ever emitted.
-  const unsigned BrkOperand = 0xc471;
+  /// An immediate operand passed to BRK instruction is (0xc470 + KeyId).
+  const unsigned BrkOperandBase = 0xc470;
 
   const AArch64Subtarget *Subtarget = nullptr;
   const AArch64InstrInfo *TII = nullptr;
@@ -46,9 +48,77 @@ class AArch64PointerAuth : public MachineFunctionPass {
   void authenticateLR(MachineFunction &MF,
                       MachineBasicBlock::iterator MBBI) const;
 
+  /// Stores blend(AddrDisc, IntDisc) to the Result register.
+  void emitBlend(MachineBasicBlock::iterator MBBI, Register Result,
+                 Register AddrDisc, unsigned IntDisc) const;
+
+  /// Stores the discriminator value to the Result register.
+  ///
+  /// This is an utility function used when expanding several PAUTH_*
+  /// pseudo instructions that follow the same conventions on expressing
+  /// the requested signing schema.
+  void storeDiscriminator(MachineBasicBlock::iterator MBBI, Register Result,
+                          bool IsBlended, Register RegDisc, unsigned IntDisc,
+                          bool ShouldStoreZero) const;
+
+  /// Emits discriminator computation followed by AUT* or PAC* instruction.
+  ///
+  /// By encoding the common cases of discriminator computation right in
+  /// the instruction's operands, this pass strives to emit blend operation
+  /// or move-immediate as close to AUT* or PAC* instruction as possible.
+  ///
+  /// The instruction is expected to have operand list starting with
+  /// - def $dst_ptr
+  /// - def $scratch
+  /// - use $src_ptr (tied to $dst_ptr)
+  /// and a contiguous sequence of operands describing the signing schema:
+  /// - use $reg_disc
+  /// - use $int_disc
+  /// - use $is_blended
+  /// - use $key_id
+  ///
+  /// As tying $scratch operand to (one of) $reg_disc is beneficial for
+  /// eliminating excessive register moves, it is allowed as an exception to
+  /// express immediate-only discriminator as $reg_disc (and thus $scratch)
+  /// being XZR. In that case, custom inserter should provide the actual
+  /// scratch register as an implicit def.
+  void expandAutOrSignWithDiscriminator(
+      MachineBasicBlock::iterator MBBI, unsigned RegDiscOpIndex,
+      unsigned (*GetHintOpcode)(Register, Register, AArch64PACKey::ID),
+      unsigned (*GetOpcode)(AArch64PACKey::ID, bool)) const;
+
+  /// Checks pointer that was authenticated.
+  ///
+  /// The instruction is expected to have operand list starting with
+  /// - def $dst_ptr
+  /// - def $scratch
+  /// as well as $key_id and optional implicit def
+  /// (see expandAutOrSignWithDiscriminator for explanation).
+  void expandAddressCheck(MachineBasicBlock::iterator MBBI,
+                          unsigned KeyIdOpIndex,
+                          Intrinsic::ID IntrinsicId) const;
+
+  /// Expands PAUTH_AUTH pseudo instruction.
+  void expandPAuthAuth(MachineBasicBlock::iterator MBBI) const;
+
   bool checkAuthenticatedLR(MachineBasicBlock::iterator TI) const;
 };
 
+unsigned getHintAUTOpcode(Register Pointer, Register Discriminator,
+                          AArch64PACKey::ID KeyId) {
+  if (Pointer != AArch64::X17 || Discriminator != AArch64::X16)
+    return 0;
+
+  switch (KeyId) {
+  case AArch64PACKey::IA:
+    return AArch64::AUTIA1716;
+  case AArch64PACKey::IB:
+    return AArch64::AUTIB1716;
+  default:
+    return 0;
+  }
+}
+
 } // end anonymous namespace
 
 INITIALIZE_PASS(AArch64PointerAuth, "aarch64-ptrauth",
@@ -174,7 +244,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;
@@ -230,6 +300,22 @@ MachineBasicBlock &llvm::AArch64PAuth::checkAuthenticatedRegister(
         .addImm(AArch64CC::NE)
         .addMBB(BreakBlock);
     return *SuccessBlock;
+  case AuthCheckMethod::XPAC:
+    BuildMI(CheckBlock, DL, TII->get(AArch64::ORRXrs), TmpReg)
+        .addReg(AArch64::XZR)
+        .addReg(AuthenticatedReg)
+        .addImm(0);
+    BuildMI(CheckBlock, DL, TII->get(UseIKey ? AArch64::XPACI : AArch64::XPACD),
+            TmpReg)
+        .addReg(TmpReg);
+    BuildMI(CheckBlock, DL, TII->get(AArch64::SUBSXrs), AArch64::XZR)
+        .addReg(TmpReg)
+        .addReg(AuthenticatedReg)
+        .addImm(0);
+    BuildMI(CheckBlock, DL, TII->get(AArch64::Bcc))
+        .addImm(AArch64CC::NE)
+        .addMBB(BreakBlock);
+    return *SuccessBlock;
   }
   llvm_unreachable("Unknown AuthCheckMethod enum");
 }
@@ -244,12 +330,18 @@ unsigned llvm::AArch64PAuth::getCheckerSizeInBytes(AuthCheckMethod Method) {
     return 12;
   case AuthCheckMethod::XPACHint:
     return 20;
+  case AuthCheckMethod::XPAC:
+    return 20;
   }
   llvm_unreachable("Unknown AuthCheckMethod enum");
 }
 
 bool AArch64PointerAuth::checkAuthenticatedLR(
     MachineBasicBlock::iterator TI) const {
+  const AArch64FunctionInfo *MFnI = TI->getMF()->getInfo<AArch64FunctionInfo>();
+  unsigned BrkOperand =
+      MFnI->shouldSignWithBKey() ? (BrkOperandBase + 1) : BrkOperandBase;
+
   AuthCheckMethod Method = Subtarget->getAuthenticatedLRCheckMethod();
 
   if (Method == AuthCheckMethod::None)
@@ -295,6 +387,123 @@ bool AArch64PointerAuth::checkAuthenticatedLR(
   return true;
 }
 
+void AArch64PointerAuth::emitBlend(MachineBasicBlock::iterator MBBI,
+                                   Register Result, Register AddrDisc,
+                                   unsigned IntDisc) const {
+  MachineBasicBlock &MBB = *MBBI->getParent();
+  DebugLoc DL = MBBI->getDebugLoc();
+
+  if (Result != AddrDisc)
+    BuildMI(MBB, MBBI, DL, TII->get(AArch64::ORRXrs), Result)
+        .addReg(AArch64::XZR)
+        .addReg(AddrDisc)
+        .addImm(0);
+
+  BuildMI(MBB, MBBI, DL, TII->get(AArch64::MOVKXi), Result)
+      .addReg(Result)
+      .addImm(IntDisc)
+      .addImm(48);
+}
+
+void AArch64PointerAuth::storeDiscriminator(MachineBasicBlock::iterator MBBI,
+                                            Register Result, bool IsBlended,
+                                            Register RegDisc, unsigned IntDisc,
+                                            bool ShouldStoreZero) const {
+  MachineBasicBlock &MBB = *MBBI->getParent();
+  DebugLoc DL = MBBI->getDebugLoc();
+
+  if (IsBlended) {
+    // Discriminator is blend(RegDisc, IntDisc).
+    emitBlend(MBBI, Result, RegDisc, IntDisc);
+  } else if (RegDisc != AArch64::XZR) {
+    // Discriminator is RegDisc.
+    assert(IntDisc == 0 &&
+           "Cannot use both reg and int discriminators without blending");
+    if (Result != RegDisc)
+      BuildMI(MBB, MBBI, DL, TII->get(AArch64::ORRXrs), Result)
+          .addReg(AArch64::XZR)
+          .addReg(RegDisc)
+          .addImm(0);
+  } else {
+    // Discriminator is IntDisc (possibly, zero at all).
+    assert(RegDisc == AArch64::XZR &&
+           "Cannot use both reg and int discriminators without blending");
+    if (IntDisc != 0 || ShouldStoreZero)
+      BuildMI(MBB, MBBI, DL, TII->get(AArch64::MOVZXi), Result)
+          .addImm(IntDisc)
+          .addImm(0);
+  }
+}
+
+void AArch64PointerAuth::expandAutOrSignWithDiscriminator(
+    MachineBasicBlock::iterator MBBI, unsigned RegDiscOpIndex,
+    unsigned (*GetHintOpcode)(Register, Register, AArch64PACKey::ID),
+    unsigned (*GetOpcode)(AArch64PACKey::ID, bool)) const {
+  MachineBasicBlock &MBB = *MBBI->getParent();
+  DebugLoc DL = MBBI->getDebugLoc();
+
+  Register PointerReg = MBBI->getOperand(0).getReg();
+  Register ScratchReg = MBBI->getOperand(1).getReg();
+  if (ScratchReg == AArch64::XZR)
+    ScratchReg = MBBI->getOperand(MBBI->getNumOperands() - 1).getReg();
+
+  Register RegDisc = MBBI->getOperand(RegDiscOpIndex).getReg();
+  unsigned IntDisc = MBBI->getOperand(RegDiscOpIndex + 1).getImm();
+  bool IsBlended = MBBI->getOperand(RegDiscOpIndex + 2).getImm();
+  auto KeyId = (AArch64PACKey::ID)MBBI->getOperand(RegDiscOpIndex + 3).getImm();
+
+  // Try using an instruction encoded as HINT, if possible.
+  unsigned HintOpcode = GetHintOpcode(PointerReg, ScratchReg, KeyId);
+  assert(HintOpcode || Subtarget->hasPAuth());
+
+  // Compute discriminator value.
+  storeDiscriminator(MBBI, ScratchReg, IsBlended, RegDisc, IntDisc,
+                     /*ShouldStoreZero=*/HintOpcode != 0);
+
+  // Now, the discriminator value is in the ScratchReg register,
+  // if not using zero-discriminator opcode variant.
+
+  if (HintOpcode) {
+    // All other opcodes require FEAT_PAuth, so check this first.
+    BuildMI(MBB, MBBI, DL, TII->get(HintOpcode));
+  } else {
+    bool IsZeroDisc = RegDisc == AArch64::XZR && IntDisc == 0;
+    unsigned RegularOpcode = GetOpcode(KeyId, IsZeroDisc);
+
+    if (IsZeroDisc) {
+      BuildMI(MBB, MBBI, DL, TII->get(RegularOpcode), PointerReg)
+          .addReg(PointerReg);
+    } else {
+      BuildMI(MBB, MBBI, DL, TII->get(RegularOpcode), PointerReg)
+          .addReg(PointerReg)
+          .addReg(ScratchReg);
+    }
+  }
+}
+
+void AArch64PointerAuth::expandAddressCheck(MachineBasicBlock::iterator MBBI,
+                                            unsigned KeyIdOpIndex,
+                                            Intrinsic::ID IntrinsicId) const {
+  Register PointerReg = MBBI->getOperand(0).getReg();
+  Register ScratchReg = MBBI->getOperand(1).getReg();
+  if (ScratchReg == AArch64::XZR)
+    ScratchReg = MBBI->getOperand(MBBI->getNumOperands() - 1).getReg();
+  auto KeyId = (AArch64PACKey::ID)MBBI->getOperand(KeyIdOpIndex).getImm();
+
+  bool UseIKey = KeyId == AArch64PACKey::IA || KeyId == AArch64PACKey::IB;
+  AuthCheckMethod Method =
+      Subtarget->getPAuthIntrinsicCheckMethod(IntrinsicId, KeyId);
+  checkAuthenticatedRegister(MBBI, Method, PointerReg, ScratchReg, UseIKey,
+                             BrkOperandBase + KeyId);
+}
+
+void AArch64PointerAuth::expandPAuthAuth(
+    MachineBasicBlock::iterator MBBI) const {
+  expandAutOrSignWithDiscriminator(MBBI, /*RegDiscOpIndex=*/3, getHintAUTOpcode,
+                                   getAUTOpcodeForKey);
+  expandAddressCheck(MBBI, /*KeyIdOpIndex=*/6, Intrinsic::ptrauth_auth);
+}
+
 bool AArch64PointerAuth::runOnMachineFunction(MachineFunction &MF) {
   const auto *MFnI = MF.getInfo<AArch64FunctionInfo>();
 
@@ -326,6...
[truncated]

``````````

</details>


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


More information about the llvm-commits mailing list