Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

atrosinenko
Copy link
Contributor

@atrosinenko atrosinenko commented Nov 16, 2023

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 @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).

@llvmbot
Copy link
Collaborator

llvmbot commented Nov 16, 2023

@llvm/pr-subscribers-backend-aarch64

Author: Anatoly Trosinenko (atrosinenko)

Changes

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 22 pseudos for expressing @llvm.ptrauth.resign the same way (or even 33 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 --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]

Copy link

github-actions bot commented Nov 16, 2023

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:
git-clang-format --diff 702664e7870c27f197dfb744a4db54aa259ce452 be81ce3e6e01fd7825ee9c4bbebe787ee8e8d5ae -- llvm/lib/Target/AArch64/AArch64ISelLowering.cpp llvm/lib/Target/AArch64/AArch64ISelLowering.h llvm/lib/Target/AArch64/AArch64PointerAuth.cpp llvm/lib/Target/AArch64/AArch64PointerAuth.h llvm/lib/Target/AArch64/AArch64Subtarget.cpp llvm/lib/Target/AArch64/AArch64Subtarget.h
View the diff from clang-format here.
diff --git a/llvm/lib/Target/AArch64/AArch64PointerAuth.h b/llvm/lib/Target/AArch64/AArch64PointerAuth.h
index e9f1379c66..10fa12bc85 100644
--- a/llvm/lib/Target/AArch64/AArch64PointerAuth.h
+++ b/llvm/lib/Target/AArch64/AArch64PointerAuth.h
@@ -102,7 +102,7 @@ enum class AuthCheckMethod {
 
 /// The methods applicable to any register authenticated using any key.
 #define AUTH_CHECK_METHOD_CL_VALUES_GENERIC                                    \
-      AUTH_CHECK_METHOD_CL_VALUES_COMMON,                                      \
+  AUTH_CHECK_METHOD_CL_VALUES_COMMON,                                          \
       clEnumValN(AArch64PAuth::AuthCheckMethod::XPAC, "xpac",                  \
                  "Compare with the result of XPAC (needs FEAT_PAuth)")
 

@atrosinenko
Copy link
Contributor Author

A few notes:

  • DummyLoad checker was fixed to use the register passed in argument of checkAuthenticatedRegister instead of LR, the operands of BRK instruction are updated to closely follow the numbers already used in @ahmedbougacha's branch
  • this patch has similar purpose to https://reviews.llvm.org/D132386, but pseudo instructions are generated by TableGen-erated code plus custom inserter and expanded before AsmPrinter
  • for uniformity, this patch computes a blend with one of the arguments being zero differently from the non-blended case:
    • blend(value, 0) clears top 16 bits of value (as it is done for any other integer modifier)
    • blend(NULL, n) shifts n left by 48 bits (as it is done for any other pointer)

@atrosinenko
Copy link
Contributor Author

To simplify this patch, two smaller PRs are opened:

@atrosinenko atrosinenko marked this pull request as draft December 19, 2023 10:19
@atrosinenko
Copy link
Contributor Author

Explicitly marked this PR as draft as it should be updated after processing #74074 and #74729 mentioned above.

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 @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).
@atrosinenko atrosinenko marked this pull request as ready for review February 5, 2024 16:45
@atrosinenko
Copy link
Contributor Author

Updated and un-drafted this PR because #74074 and #74729 were merged.

@ahmedbougacha Is the original implementation of @llvm.ptrauth.auth planned to be merged instead (#79024)?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
Status: In Progress
Development

Successfully merging this pull request may close these issues.

None yet

2 participants