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

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 AUT
to only using them. Therefore, adjust the backend to only use the AUT
pseudo (renamed AUTx16x17 for clarity) on Darwin platforms, or if traps
are requested. All other platforms use the unconstrained variant of the
instruction for selection.

Pull Request: llvm#132857

Author: pcc

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