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

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -69,15 +69,6 @@
 
 using namespace llvm;
 
-enum PtrauthCheckMode { Default, Unchecked, Poison, Trap };
-static cl::opt<PtrauthCheckMode> PtrauthAuthChecks(
-    "aarch64-ptrauth-auth-checks", cl::Hidden,
-    cl::values(clEnumValN(Unchecked, "none", "don't test for failure"),
-               clEnumValN(Poison, "poison", "poison on failure"),
-               clEnumValN(Trap, "trap", "trap on failure")),
-    cl::desc("Check pointer authentication auth/resign failures"),
-    cl::init(Default));
-
 #define DEBUG_TYPE "asm-printer"
 
 namespace {
@@ -2866,7 +2857,7 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
     return;
   }
 
-  case AArch64::AUT:
+  case AArch64::AUTx16x17:
   case AArch64::AUTPAC:
     emitPtrauthAuthResign(MI);
     return;

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

@@ -361,7 +361,7 @@ class AArch64DAGToDAGISel : public SelectionDAGISel {
 
   bool tryIndexedLoad(SDNode *N);
 
-  void SelectPtrauthAuth(SDNode *N);
+  void SelectPtrauthAuthX16X17(SDNode *N);
   void SelectPtrauthResign(SDNode *N);
 
   bool trySelectStackSlotTagP(SDNode *N);
@@ -1521,7 +1521,7 @@ extractPtrauthBlendDiscriminators(SDValue Disc, SelectionDAG *DAG) {
       AddrDisc);
 }
 
-void AArch64DAGToDAGISel::SelectPtrauthAuth(SDNode *N) {
+void AArch64DAGToDAGISel::SelectPtrauthAuthX16X17(SDNode *N) {
   SDLoc DL(N);
   // IntrinsicID is operand #0
   SDValue Val = N->getOperand(1);
@@ -1539,7 +1539,7 @@ void AArch64DAGToDAGISel::SelectPtrauthAuth(SDNode *N) {
                                          AArch64::X16, Val, SDValue());
   SDValue Ops[] = {AUTKey, AUTConstDisc, AUTAddrDisc, X16Copy.getValue(1)};
 
-  SDNode *AUT = CurDAG->getMachineNode(AArch64::AUT, DL, MVT::i64, Ops);
+  SDNode *AUT = CurDAG->getMachineNode(AArch64::AUTx16x17, DL, MVT::i64, Ops);
   ReplaceNode(N, AUT);
 }
 
@@ -5613,7 +5613,9 @@ void AArch64DAGToDAGISel::Select(SDNode *Node) {
       return;
 
     case Intrinsic::ptrauth_auth:
-      SelectPtrauthAuth(Node);
+      if (!Subtarget->isX16X17Safer(CurDAG->getMachineFunction()))
+        break;
+      SelectPtrauthAuthX16X17(Node);
       return;
 
     case Intrinsic::ptrauth_resign:

llvm/lib/Target/AArch64/AArch64InstrInfo.td

@@ -66,6 +66,14 @@ def HasLOR           : Predicate<"Subtarget->hasLOR()">,
 def HasPAuth         : Predicate<"Subtarget->hasPAuth()">,
                        AssemblerPredicateWithAll<(all_of FeaturePAuth), "pauth">;
 
+// On most operating systems, the x16 and x17 registers are not special, so
+// there is no benefit, and only a code size cost, to constraining PAC
+// instructions to only using them. This predicate may be used to guard patterns
+// that allow PAC instructions to be used with any register.
+let RecomputePerFunction = 1 in {
+  def X16X17NotSafer   : Predicate<"!Subtarget->isX16X17Safer(*MF)">;
+}
+
 def HasPAuthLR       : Predicate<"Subtarget->hasPAuthLR()">,
                        AssemblerPredicateWithAll<(all_of FeaturePAuthLR), "pauth-lr">;
 
@@ -1820,7 +1828,9 @@ let Predicates = [HasPAuth] in {
   }
 
   defm PAC : SignAuth<0b000, 0b010, "pac", int_ptrauth_sign>;
-  defm AUT : SignAuth<0b001, 0b011, "aut", null_frag>;
+  let Predicates = [HasPAuth, X16X17NotSafer] in {
+    defm AUT : SignAuth<0b001, 0b011, "aut", int_ptrauth_auth>;
+  }
 
   def XPACI : ClearAuth<0, "xpaci">;
   def : Pat<(int_ptrauth_strip GPR64:$Rd, 0), (XPACI GPR64:$Rd)>;
@@ -1912,10 +1922,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;
@@ -1926,8 +1937,8 @@ let Predicates = [HasPAuth] in {
   }
 
   // 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

@@ -102,6 +102,16 @@ static cl::opt<bool>
     UseScalarIncVL("sve-use-scalar-inc-vl", cl::init(false), cl::Hidden,
                    cl::desc("Prefer add+cnt over addvl/inc/dec"));
 
+cl::opt<PtrauthCheckMode> llvm::PtrauthAuthChecks(
+    "aarch64-ptrauth-auth-checks", cl::Hidden,
+    cl::values(clEnumValN(PtrauthCheckMode::Unchecked, "none",
+                          "don't test for failure"),
+               clEnumValN(PtrauthCheckMode::Poison, "poison",
+                          "poison on failure"),
+               clEnumValN(PtrauthCheckMode::Trap, "trap", "trap on failure")),
+    cl::desc("Check pointer authentication auth/resign failures"),
+    cl::init(PtrauthCheckMode::Default));
+
 unsigned AArch64Subtarget::getVectorInsertExtractBaseCost() const {
   if (OverrideVectorInsertExtractBaseCost.getNumOccurrences() > 0)
     return OverrideVectorInsertExtractBaseCost;
@@ -663,6 +673,19 @@ AArch64Subtarget::getPtrAuthBlockAddressDiscriminatorIfEnabled(
       (Twine(ParentFn.getName()) + " blockaddress").str());
 }
 
+bool AArch64Subtarget::isX16X17Safer(const MachineFunction &MF) const {
+  // The Darwin kernel implements special protections for x16 and x17 so we
+  // should prefer to use those registers on that platform.
+  if (isTargetDarwin())
+    return true;
+  // Traps are only implemented for the pseudo instructions, but are only
+  // necessary if FEAT_FPAC is not implemented.
+  if (hasFPAC())
+    return false;
+  return MF.getFunction().hasFnAttribute("ptrauth-auth-traps") ||
+         PtrauthAuthChecks == PtrauthCheckMode::Trap;
+}
+
 bool AArch64Subtarget::enableMachinePipeliner() const {
   return getSchedModel().hasInstrSchedModel();
 }

llvm/lib/Target/AArch64/AArch64Subtarget.h

@@ -26,6 +26,7 @@
 #include "llvm/CodeGen/RegisterBankInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/DataLayout.h"
+#include "llvm/Support/CommandLine.h"
 
 #define GET_SUBTARGETINFO_HEADER
 #include "AArch64GenSubtargetInfo.inc"
@@ -35,6 +36,9 @@ class GlobalValue;
 class StringRef;
 class Triple;
 
+enum class PtrauthCheckMode { Default, Unchecked, Poison, Trap };
+extern cl::opt<PtrauthCheckMode> PtrauthAuthChecks;
+
 class AArch64Subtarget final : public AArch64GenSubtargetInfo {
 public:
   enum ARMProcFamilyEnum : uint8_t {
@@ -318,6 +322,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 MachineFunction &MF) 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

@@ -6735,7 +6735,7 @@ bool AArch64InstructionSelector::selectIntrinsic(MachineInstr &I,
 
     MIB.buildCopy({AArch64::X16}, {ValReg});
     MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X17}, {});
-    MIB.buildInstr(AArch64::AUT)
+    MIB.buildInstr(AArch64::AUTx16x17)
         .addImm(AUTKey)
         .addImm(AUTConstDiscC)
         .addUse(AUTAddrDisc)

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

@@ -169,13 +169,23 @@ 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]
+; FIXME: Get rid of the x16/x17 constraint on non-Darwin so we can eliminate
+; this mov.
+; ELF-NEXT:       mov     x16, x1
+; ELF-NEXT:       movk    x16, #54167, lsl #48
+; ELF-NEXT:       braa    x2, x16
   %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 +201,27 @@ 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]
+; FIXME: Get rid of the x16/x17 constraint on non-Darwin so we can eliminate
+; this mov.
+; ELF-NEXT:      mov     x17, x8
+; ELF-NEXT:      movk    x17, #34646, lsl #48
+; ELF-NEXT:      blraa   x9, x17
+; ELF-NEXT:      mov     w0, #42
 ; ELF-NEXT:      ldr     x30, [sp], #16
 ; CHECK-NEXT:    ret
   %vtable.signed = load ptr, ptr %objptr