[AArch64][PAC] Cleanup AArch64AsmPrinter::emitPtrauthDiscriminator (NFC)

[atrosinenko]

Refactor emitPtrauthDiscriminator function: introduce isPtrauthRegSafe
function, update the comments and assertions for readability.

[atrosinenko]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• [AArch64][PAC] Factor out printing real AUT/PAC/BLRA encodings (NFC) #160901
• [AArch64][PAC] Cleanup AArch64AsmPrinter::emitPtrauthDiscriminator (NFC) #160900 👈 (View in Graphite)
• [AArch64][PAC] Simplify emission of authenticated pointer check (NFC) #160899
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Anatoly Trosinenko (atrosinenko)

Changes

Refactor emitPtrauthDiscriminator function: introduce isPtrauthRegSafe
function, update the comments and assertions for readability.

---

Full diff: https://github.com/llvm/llvm-project/pull/160900.diff

1 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp (+25-21)

diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index 1a2808f4d56d8..e7135da17a8d5 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -155,6 +155,13 @@ class AArch64AsmPrinter : public AsmPrinter {
 
   void emitSled(const MachineInstr &MI, SledKind Kind);
 
+  bool isPtrauthRegSafe(Register Reg) const {
+    if (STI->isX16X17Safer())
+      return Reg == AArch64::X16 || Reg == AArch64::X17;
+
+    return true;
+  }
+
   // Emit the sequence for BRA/BLRA (authenticate + branch/call).
   void emitPtrauthBranch(const MachineInstr *MI);
 
@@ -180,13 +187,15 @@ class AArch64AsmPrinter : public AsmPrinter {
 
   // Emit the sequence to compute the discriminator.
   //
-  // The returned register is either unmodified AddrDisc or ScratchReg.
+  // The Scratch register passed to this function must be safe, see
+  // isPtrauthRegSafe(Reg) function.
+  //
+  // The returned register is either ScratchReg or AddrDisc. Furthermore, it
+  // is safe, unless unsafe AddrDisc was passed-through unmodified.
   //
   // If the expanded pseudo is allowed to clobber AddrDisc register, setting
-  // MayUseAddrAsScratch may save one MOV instruction, provided the address
-  // is already in x16/x17 (i.e. return x16/x17 which is the *modified* AddrDisc
-  // register at the same time) or the OS doesn't make it safer to use x16/x17
-  // (see AArch64Subtarget::isX16X17Safer()):
+  // MayUseAddrAsScratch may save one MOV instruction, provided
+  // isPtrauthRegSafe(AddrDisc) is true:
   //
   //   mov   x17, x16
   //   movk  x17, #1234, lsl #48
@@ -195,7 +204,7 @@ class AArch64AsmPrinter : public AsmPrinter {
   // can be replaced by
   //
   //   movk  x16, #1234, lsl #48
-  Register emitPtrauthDiscriminator(uint16_t Disc, Register AddrDisc,
+  Register emitPtrauthDiscriminator(uint64_t Disc, Register AddrDisc,
                                     Register ScratchReg,
                                     bool MayUseAddrAsScratch = false);
 
@@ -1902,12 +1911,14 @@ void AArch64AsmPrinter::emitFMov0AsFMov(const MachineInstr &MI,
   EmitToStreamer(*OutStreamer, FMov);
 }
 
-Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
+Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint64_t Disc,
                                                      Register AddrDisc,
                                                      Register ScratchReg,
                                                      bool MayUseAddrAsScratch) {
-  assert(ScratchReg == AArch64::X16 || ScratchReg == AArch64::X17 ||
-         !STI->isX16X17Safer());
+  assert(isPtrauthRegSafe(ScratchReg) &&
+         "Safe scratch register must be provided by the caller");
+  assert(isUInt<16>(Disc) && "Constant discriminator is too wide");
+
   // So far we've used NoRegister in pseudos.  Now we need real encodings.
   if (AddrDisc == AArch64::NoRegister)
     AddrDisc = AArch64::XZR;
@@ -1926,13 +1937,13 @@ Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
   // If there are both, emit a blend into the scratch register.
 
   // Check if we can save one MOV instruction.
-  assert(MayUseAddrAsScratch || ScratchReg != AddrDisc);
-  bool AddrDiscIsSafe = AddrDisc == AArch64::X16 || AddrDisc == AArch64::X17 ||
-                        !STI->isX16X17Safer();
-  if (MayUseAddrAsScratch && AddrDiscIsSafe)
+  if (MayUseAddrAsScratch && isPtrauthRegSafe(AddrDisc)) {
     ScratchReg = AddrDisc;
-  else
+  } else {
     emitMovXReg(ScratchReg, AddrDisc);
+    assert(ScratchReg != AddrDisc &&
+           "Forbidden to clobber AddrDisc, but have to");
+  }
 
   emitMOVK(ScratchReg, Disc, 48);
   return ScratchReg;
@@ -2151,7 +2162,6 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(
   }
 
   // Compute aut discriminator
-  assert(isUInt<16>(AUTDisc));
   Register AUTDiscReg = emitPtrauthDiscriminator(
       AUTDisc, AUTAddrDisc->getReg(), Scratch, AUTAddrDisc->isKill());
   bool AUTZero = AUTDiscReg == AArch64::XZR;
@@ -2188,7 +2198,6 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(
     return;
 
   // Compute pac discriminator
-  assert(isUInt<16>(PACDisc));
   Register PACDiscReg =
       emitPtrauthDiscriminator(PACDisc, PACAddrDisc, Scratch);
   bool PACZero = PACDiscReg == AArch64::XZR;
@@ -2223,7 +2232,6 @@ void AArch64AsmPrinter::emitPtrauthSign(const MachineInstr *MI) {
          "Neither X16 nor X17 is available as a scratch register");
 
   // Compute pac discriminator
-  assert(isUInt<16>(Disc));
   Register DiscReg = emitPtrauthDiscriminator(
       Disc, AddrDisc, ScratchReg, /*MayUseAddrAsScratch=*/AddrDiscKilled);
   bool IsZeroDisc = DiscReg == AArch64::XZR;
@@ -2249,7 +2257,6 @@ void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
          "Invalid auth call key");
 
   uint64_t Disc = MI->getOperand(2).getImm();
-  assert(isUInt<16>(Disc));
 
   unsigned AddrDisc = MI->getOperand(3).getReg();
 
@@ -2422,8 +2429,6 @@ void AArch64AsmPrinter::LowerMOVaddrPAC(const MachineInstr &MI) {
   const auto Key = (AArch64PACKey::ID)KeyC;
   const unsigned AddrDisc = MI.getOperand(2).getReg();
   const uint64_t Disc = MI.getOperand(3).getImm();
-  assert(isUInt<16>(Disc) &&
-         "constant discriminator is out of range [0, 0xffff]");
 
   const int64_t Offset = GAOp.getOffset();
   GAOp.setOffset(0);
@@ -2995,7 +3000,6 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
            "Invalid auth key for tail-call return");
 
     const uint64_t Disc = MI->getOperand(3).getImm();
-    assert(isUInt<16>(Disc) && "Integer discriminator is too wide");
 
     Register AddrDisc = MI->getOperand(4).getReg();
 

[atrosinenko]

Ping.