AArch64: Relax x16/x17 constraint on AUT in certain cases.

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -168,7 +168,12 @@ class AArch64AsmPrinter : public AsmPrinter {
   void emitPtrauthTailCallHardening(const MachineInstr *TC);
 
   // Emit the sequence for AUT or AUTPAC.
-  void emitPtrauthAuthResign(const MachineInstr *MI);
+  void emitPtrauthAuthResign(Register AUTVal, AArch64PACKey::ID AUTKey,
+                             uint64_t AUTDisc,
+                             const MachineOperand *AUTAddrDisc,
+                             Register Scratch,
+                             std::optional<AArch64PACKey::ID> PACKey,
+                             uint64_t PACDisc, unsigned PACAddrDisc);
 
   // Emit the sequence to compute the discriminator.
   //
@@ -1866,7 +1871,6 @@ Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
                                                      Register AddrDisc,
                                                      Register ScratchReg,
                                                      bool MayUseAddrAsScratch) {
-  assert(ScratchReg == AArch64::X16 || ScratchReg == AArch64::X17);
   // So far we've used NoRegister in pseudos.  Now we need real encodings.
   if (AddrDisc == AArch64::NoRegister)
     AddrDisc = AArch64::XZR;
@@ -1886,7 +1890,8 @@ Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
 
   // Check if we can save one MOV instruction.
   assert(MayUseAddrAsScratch || ScratchReg != AddrDisc);
-  bool AddrDiscIsSafe = AddrDisc == AArch64::X16 || AddrDisc == AArch64::X17;
+  bool AddrDiscIsSafe = AddrDisc == AArch64::X16 || AddrDisc == AArch64::X17 ||
+                        !STI->isX16X17Safer();
   if (MayUseAddrAsScratch && AddrDiscIsSafe)
     ScratchReg = AddrDisc;
   else
@@ -2064,8 +2069,12 @@ void AArch64AsmPrinter::emitPtrauthTailCallHardening(const MachineInstr *TC) {
       /*ShouldTrap=*/true, /*OnFailure=*/nullptr);
 }
 
-void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
-  const bool IsAUTPAC = MI->getOpcode() == AArch64::AUTPAC;
+void AArch64AsmPrinter::emitPtrauthAuthResign(
+    Register AUTVal, AArch64PACKey::ID AUTKey, uint64_t AUTDisc,
+    const MachineOperand *AUTAddrDisc, Register Scratch,
+    std::optional<AArch64PACKey::ID> PACKey, uint64_t PACDisc,
+    unsigned PACAddrDisc) {
+  const bool IsAUTPAC = PACKey.has_value();
 
   // We expand AUT/AUTPAC into a sequence of the form
   //
@@ -2104,23 +2113,19 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
     break;
   }
 
-  auto AUTKey = (AArch64PACKey::ID)MI->getOperand(0).getImm();
-  uint64_t AUTDisc = MI->getOperand(1).getImm();
-  unsigned AUTAddrDisc = MI->getOperand(2).getReg();
-
   // Compute aut discriminator into x17
   assert(isUInt<16>(AUTDisc));
-  Register AUTDiscReg =
-      emitPtrauthDiscriminator(AUTDisc, AUTAddrDisc, AArch64::X17);
+  Register AUTDiscReg = emitPtrauthDiscriminator(
+      AUTDisc, AUTAddrDisc->getReg(), Scratch, AUTAddrDisc->isKill());
   bool AUTZero = AUTDiscReg == AArch64::XZR;
   unsigned AUTOpc = getAUTOpcodeForKey(AUTKey, AUTZero);
 
   //  autiza x16      ; if  AUTZero
   //  autia x16, x17  ; if !AUTZero
   MCInst AUTInst;
   AUTInst.setOpcode(AUTOpc);
-  AUTInst.addOperand(MCOperand::createReg(AArch64::X16));
-  AUTInst.addOperand(MCOperand::createReg(AArch64::X16));
+  AUTInst.addOperand(MCOperand::createReg(AUTVal));
+  AUTInst.addOperand(MCOperand::createReg(AUTVal));
   if (!AUTZero)
     AUTInst.addOperand(MCOperand::createReg(AUTDiscReg));
   EmitToStreamer(*OutStreamer, AUTInst);
@@ -2135,7 +2140,7 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
     if (IsAUTPAC && !ShouldTrap)
       EndSym = createTempSymbol("resign_end_");
 
-    emitPtrauthCheckAuthenticatedValue(AArch64::X16, AArch64::X17, AUTKey,
+    emitPtrauthCheckAuthenticatedValue(AUTVal, Scratch, AUTKey,
                                        AArch64PAuth::AuthCheckMethod::XPAC,
                                        ShouldTrap, EndSym);
   }
@@ -2146,23 +2151,19 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
   if (!IsAUTPAC)
     return;
 
-  auto PACKey = (AArch64PACKey::ID)MI->getOperand(3).getImm();
-  uint64_t PACDisc = MI->getOperand(4).getImm();
-  unsigned PACAddrDisc = MI->getOperand(5).getReg();
-
   // Compute pac discriminator into x17
   assert(isUInt<16>(PACDisc));
   Register PACDiscReg =
-      emitPtrauthDiscriminator(PACDisc, PACAddrDisc, AArch64::X17);
+      emitPtrauthDiscriminator(PACDisc, PACAddrDisc, Scratch);
   bool PACZero = PACDiscReg == AArch64::XZR;
-  unsigned PACOpc = getPACOpcodeForKey(PACKey, PACZero);
+  unsigned PACOpc = getPACOpcodeForKey(*PACKey, PACZero);
 
   //  pacizb x16      ; if  PACZero
   //  pacib x16, x17  ; if !PACZero
   MCInst PACInst;
   PACInst.setOpcode(PACOpc);
-  PACInst.addOperand(MCOperand::createReg(AArch64::X16));
-  PACInst.addOperand(MCOperand::createReg(AArch64::X16));
+  PACInst.addOperand(MCOperand::createReg(AUTVal));
+  PACInst.addOperand(MCOperand::createReg(AUTVal));
   if (!PACZero)
     PACInst.addOperand(MCOperand::createReg(PACDiscReg));
   EmitToStreamer(*OutStreamer, PACInst);
@@ -2861,9 +2862,26 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
     return;
   }
 
-  case AArch64::AUT:
+  case AArch64::AUTx16x17:
+    emitPtrauthAuthResign(AArch64::X16,
+                          (AArch64PACKey::ID)MI->getOperand(0).getImm(),
+                          MI->getOperand(1).getImm(), &MI->getOperand(2),
+                          AArch64::X17, std::nullopt, 0, 0);
+    return;
+
+  case AArch64::AUTxMxN:
+    emitPtrauthAuthResign(MI->getOperand(0).getReg(),
+                          (AArch64PACKey::ID)MI->getOperand(3).getImm(),
+                          MI->getOperand(4).getImm(), &MI->getOperand(5),
+                          MI->getOperand(1).getReg(), std::nullopt, 0, 0);
+    return;
+
   case AArch64::AUTPAC:
-    emitPtrauthAuthResign(MI);
+    emitPtrauthAuthResign(
+        AArch64::X16, (AArch64PACKey::ID)MI->getOperand(0).getImm(),
+        MI->getOperand(1).getImm(), &MI->getOperand(2), AArch64::X17,
+        (AArch64PACKey::ID)MI->getOperand(3).getImm(),
+        MI->getOperand(4).getImm(), MI->getOperand(5).getReg());
     return;
 
   case AArch64::LOADauthptrstatic:

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

@@ -1534,12 +1534,20 @@ void AArch64DAGToDAGISel::SelectPtrauthAuth(SDNode *N) {
   std::tie(AUTConstDisc, AUTAddrDisc) =
       extractPtrauthBlendDiscriminators(AUTDisc, CurDAG);
 
-  SDValue X16Copy = CurDAG->getCopyToReg(CurDAG->getEntryNode(), DL,
-                                         AArch64::X16, Val, SDValue());
-  SDValue Ops[] = {AUTKey, AUTConstDisc, AUTAddrDisc, X16Copy.getValue(1)};
+  if (!Subtarget->isX16X17Safer()) {
+    SDValue Ops[] = {Val, AUTKey, AUTConstDisc, AUTAddrDisc};
 
-  SDNode *AUT = CurDAG->getMachineNode(AArch64::AUT, DL, MVT::i64, Ops);
-  ReplaceNode(N, AUT);
+    SDNode *AUT =
+        CurDAG->getMachineNode(AArch64::AUTxMxN, DL, MVT::i64, MVT::i64, Ops);
+    ReplaceNode(N, AUT);
+  } else {
+    SDValue X16Copy = CurDAG->getCopyToReg(CurDAG->getEntryNode(), DL,
+                                           AArch64::X16, Val, SDValue());
+    SDValue Ops[] = {AUTKey, AUTConstDisc, AUTAddrDisc, X16Copy.getValue(1)};
+
+    SDNode *AUT = CurDAG->getMachineNode(AArch64::AUTx16x17, DL, MVT::i64, Ops);
+    ReplaceNode(N, AUT);
+  }
 }
 
 void AArch64DAGToDAGISel::SelectPtrauthResign(SDNode *N) {

llvm/lib/Target/AArch64/AArch64InstrInfo.td

@@ -1938,10 +1938,11 @@ let Predicates = [HasPAuth] in {
   defm LDRAB  : AuthLoad<1, "ldrab", simm10Scaled>;
 
   // AUT pseudo.
-  // This directly manipulates x16/x17, which are the only registers the OS
-  // guarantees are safe to use for sensitive operations.
-  def AUT : Pseudo<(outs), (ins i32imm:$Key, i64imm:$Disc, GPR64noip:$AddrDisc),
-                   []>, Sched<[WriteI, ReadI]> {
+  // This directly manipulates x16/x17, which are the only registers that
+  // certain OSs guarantee are safe to use for sensitive operations.
+  def AUTx16x17 : Pseudo<(outs), (ins i32imm:$Key, i64imm:$Disc,
+                                      GPR64noip:$AddrDisc),
+                         []>, Sched<[WriteI, ReadI]> {
     let isCodeGenOnly = 1;
     let hasSideEffects = 1;
     let mayStore = 0;
@@ -1951,9 +1952,23 @@ let Predicates = [HasPAuth] in {
     let Uses = [X16];
   }
 
+  def AUTxMxN : Pseudo<(outs GPR64:$AuthVal, GPR64common:$Scratch),
+                       (ins GPR64:$Val, i32imm:$Key,
+                            i64imm:$Disc, GPR64:$AddrDisc),
+                       []>, Sched<[WriteI, ReadI]> {
+    let Constraints = "$AuthVal = $Val";
+    let isCodeGenOnly = 1;
+    let hasSideEffects = 0;
+    let mayStore = 0;
+    let mayLoad = 0;
+    let Size = 32;
+    let Defs = [NZCV];
+    let Uses = [];
+  }
+
   // AUT and re-PAC a value, using different keys/data.
-  // This directly manipulates x16/x17, which are the only registers the OS
-  // guarantees are safe to use for sensitive operations.
+  // This directly manipulates x16/x17, which are the only registers that
+  // certain OSs guarantee are safe to use for sensitive operations.
   def AUTPAC
       : Pseudo<(outs),
                (ins i32imm:$AUTKey, i64imm:$AUTDisc, GPR64noip:$AUTAddrDisc,

llvm/lib/Target/AArch64/AArch64Subtarget.cpp

@@ -664,6 +664,12 @@ AArch64Subtarget::getPtrAuthBlockAddressDiscriminatorIfEnabled(
       (Twine(ParentFn.getName()) + " blockaddress").str());
 }
 
+bool AArch64Subtarget::isX16X17Safer() const {
+  // The Darwin kernel implements special protections for x16 and x17 so we
+  // should prefer to use those registers on that platform.
+  return isTargetDarwin();
+}
+
 bool AArch64Subtarget::enableMachinePipeliner() const {
   return getSchedModel().hasInstrSchedModel();
 }

llvm/lib/Target/AArch64/AArch64Subtarget.h

@@ -318,6 +318,10 @@ class AArch64Subtarget final : public AArch64GenSubtargetInfo {
     }
   }
 
+  /// Returns whether the operating system makes it safer to store sensitive
+  /// values in x16 and x17 as opposed to other registers.
+  bool isX16X17Safer() const;
+
   /// ParseSubtargetFeatures - Parses features string setting specified
   /// subtarget options.  Definition of function is auto generated by tblgen.
   void ParseSubtargetFeatures(StringRef CPU, StringRef TuneCPU, StringRef FS);

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

@@ -6737,14 +6737,27 @@ bool AArch64InstructionSelector::selectIntrinsic(MachineInstr &I,
     std::tie(AUTConstDiscC, AUTAddrDisc) =
         extractPtrauthBlendDiscriminators(AUTDisc, MRI);
 
-    MIB.buildCopy({AArch64::X16}, {ValReg});
-    MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X17}, {});
-    MIB.buildInstr(AArch64::AUT)
-        .addImm(AUTKey)
-        .addImm(AUTConstDiscC)
-        .addUse(AUTAddrDisc)
-        .constrainAllUses(TII, TRI, RBI);
-    MIB.buildCopy({DstReg}, Register(AArch64::X16));
+    if (STI.isX16X17Safer()) {
+      MIB.buildCopy({AArch64::X16}, {ValReg});
+      MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X17}, {});
+      MIB.buildInstr(AArch64::AUTx16x17)
+          .addImm(AUTKey)
+          .addImm(AUTConstDiscC)
+          .addUse(AUTAddrDisc)
+          .constrainAllUses(TII, TRI, RBI);
+      MIB.buildCopy({DstReg}, Register(AArch64::X16));
+    } else {
+      Register ScratchReg =
+          MRI.createVirtualRegister(&AArch64::GPR64commonRegClass);
+      MIB.buildInstr(AArch64::AUTxMxN)
+          .addDef(DstReg)
+          .addDef(ScratchReg)
+          .addUse(ValReg)
+          .addImm(AUTKey)
+          .addImm(AUTConstDiscC)
+          .addUse(AUTAddrDisc)
+          .constrainAllUses(TII, TRI, RBI);
+    }
 
     RBI.constrainGenericRegister(DstReg, AArch64::GPR64RegClass, MRI);
     I.eraseFromParent();

llvm/test/CodeGen/AArch64/ptrauth-call.ll

@@ -169,13 +169,20 @@ define i32 @test_tailcall_ib_var(ptr %arg0, ptr %arg1) #0 {
 
 define void @test_tailcall_omit_mov_x16_x16(ptr %objptr) #0 {
 ; CHECK-LABEL: test_tailcall_omit_mov_x16_x16:
-; CHECK-NEXT:    ldr     x16, [x0]
-; CHECK-NEXT:    mov     x17, x0
-; CHECK-NEXT:    movk    x17, #6503, lsl #48
-; CHECK-NEXT:    autda   x16, x17
-; CHECK-NEXT:    ldr     x1, [x16]
-; CHECK-NEXT:    movk    x16, #54167, lsl #48
-; CHECK-NEXT:    braa    x1, x16
+; DARWIN-NEXT:    ldr     x16, [x0]
+; DARWIN-NEXT:    mov     x17, x0
+; DARWIN-NEXT:    movk    x17, #6503, lsl #48
+; DARWIN-NEXT:    autda   x16, x17
+; DARWIN-NEXT:    ldr     x1, [x16]
+; DARWIN-NEXT:    movk    x16, #54167, lsl #48
+; DARWIN-NEXT:    braa    x1, x16
+; ELF-NEXT:       ldr     x1, [x0]
+; ELF-NEXT:       mov     x8, x0
+; ELF-NEXT:       movk    x8, #6503, lsl #48
+; ELF-NEXT:       autda   x1, x8
+; ELF-NEXT:       ldr     x2, [x1]
+; ELF-NEXT:       movk    x1, #54167, lsl #48
+; ELF-NEXT:       braa    x2, x1
   %vtable.signed = load ptr, ptr %objptr, align 8
   %objptr.int = ptrtoint ptr %objptr to i64
   %vtable.discr = tail call i64 @llvm.ptrauth.blend(i64 %objptr.int, i64 6503)
@@ -191,16 +198,24 @@ define void @test_tailcall_omit_mov_x16_x16(ptr %objptr) #0 {
 define i32 @test_call_omit_extra_moves(ptr %objptr) #0 {
 ; CHECK-LABEL: test_call_omit_extra_moves:
 ; DARWIN-NEXT:   stp     x29, x30, [sp, #-16]!
-; ELF-NEXT:      str     x30, [sp, #-16]!
-; CHECK-NEXT:    ldr     x16, [x0]
-; CHECK-NEXT:    mov     x17, x0
-; CHECK-NEXT:    movk    x17, #6503, lsl #48
-; CHECK-NEXT:    autda   x16, x17
-; CHECK-NEXT:    ldr     x8, [x16]
-; CHECK-NEXT:    movk    x16, #34646, lsl #48
-; CHECK-NEXT:    blraa   x8, x16
-; CHECK-NEXT:    mov     w0, #42
+; DARWIN-NEXT:   ldr     x16, [x0]
+; DARWIN-NEXT:   mov     x17, x0
+; DARWIN-NEXT:   movk    x17, #6503, lsl #48
+; DARWIN-NEXT:   autda   x16, x17
+; DARWIN-NEXT:   ldr     x8, [x16]
+; DARWIN-NEXT:   movk    x16, #34646, lsl #48
+; DARWIN-NEXT:   blraa   x8, x16
+; DARWIN-NEXT:   mov     w0, #42
 ; DARWIN-NEXT:   ldp     x29, x30, [sp], #16
+; ELF-NEXT:      str     x30, [sp, #-16]!
+; ELF-NEXT:      ldr     x8, [x0]
+; ELF-NEXT:      mov     x9, x0
+; ELF-NEXT:      movk    x9, #6503, lsl #48
+; ELF-NEXT:      autda   x8, x9
+; ELF-NEXT:      ldr     x9, [x8]
+; ELF-NEXT:      movk    x8, #34646, lsl #48
+; ELF-NEXT:      blraa   x9, x8
+; ELF-NEXT:      mov     w0, #42
 ; ELF-NEXT:      ldr     x30, [sp], #16
 ; CHECK-NEXT:    ret
   %vtable.signed = load ptr, ptr %objptr