[AArch64][PAC] Combine signing with address materialization

In pauthtest ABI, it is common to store a pointer signed with address
diversity to a heap-allocated object (namely, storing a signed pointer
to VTable as part of new object construction). This patch tries to
prevent introducing a signing oracle by combining pointer
materialization and its (re)signing into a single pseudo instruction
which is not expanded until AsmPrinter, if possible.

One of the typical patterns is materializing an unsigned pointer with
`MOVaddr` pseudo and then signing it with `PAC[ID][AB]` instruction,
which can be moved far away from `MOVaddr` by one of the passes in the
machine pipeline. As the storage address is not a `Constant` value, one
cannot simply emit a `ptrauth` constant in the frontend, which would be
selected into `MOVaddrPAC` pseudo.

Another pattern is fetching a pointer to VTable from a signed GOT entry
using `LOADgotAUTH` pseudo, authenticating and checking it, and then
re-signing after adding an offset.

This commit adds an instruction insertion hook for `PAC[ID][AB]` which
detects the above patterns and replaces it either with `MOVaddrPAC` or
`LOADgotPAC` instruction.

Author: atrosinenko

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -3058,6 +3058,78 @@ AArch64TargetLowering::EmitGetSMESaveSize(MachineInstr &MI,
             MI.getOperand(0).getReg())
         .addReg(AArch64::XZR);
   BB->remove_instr(&MI);
+
+  return BB;
+}
+
+MachineBasicBlock *
+AArch64TargetLowering::tryRewritingPAC(MachineInstr &MI,
+                                       MachineBasicBlock *BB) const {
+  const TargetInstrInfo *TII = Subtarget->getInstrInfo();
+  MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
+  const DebugLoc &DL = MI.getDebugLoc();
+
+  MachineInstr *AddrInstr = nullptr;
+  int64_t Offset = 0;
+  // Try to find the address-setting instruction, accumulating the offset
+  // along the way. If any unknown pattern is found, keep everything as-is.
+  MachineOperand *CurOp = &MI.getOperand(1);
+  while (CurOp) {
+    MachineOperand *Def = MRI.getOneDef(CurOp->getReg());
+    if (!Def)
+      return BB;
+    MachineInstr *DefMI = Def->getParent();
+    assert(DefMI != nullptr);
+
+    switch (DefMI->getOpcode()) {
+    case AArch64::COPY:
+      CurOp = &DefMI->getOperand(1);
+      break;
+    case AArch64::ADDXri:
+      if (DefMI->getOperand(3).getImm() != 0) // shifts are not handled
+        return BB;
+      CurOp = &DefMI->getOperand(1);
+      Offset += DefMI->getOperand(2).getImm();
+      break;
+    case AArch64::MOVaddr:
+    case AArch64::LOADgotAUTH:
+      AddrInstr = DefMI;
+      CurOp = nullptr;
+      break;
+    default:
+      return BB;
+    }
+  }
+
+  unsigned NewOpcode = AddrInstr->getOpcode() == AArch64::LOADgotAUTH
+                           ? AArch64::LOADgotPAC
+                           : AArch64::MOVaddrPAC;
+  MachineOperand &AddrOp = AddrInstr->getOperand(1);
+  unsigned TargetFlags = AddrOp.getTargetFlags() & ~AArch64II::MO_PAGE;
+  Offset += AddrOp.getOffset();
+
+  // MOVaddrPAC and LOADgotPAC pseudos are expanded so that they use X16/X17
+  // internally, thus their restrictions on the register class of $AddrDisc
+  // operand are stricter than those of real PAC* instructions.
+  // If the original instruction accepts a discriminator operand, make sure
+  // it is moved out of X16/X17.
+  Register DiscReg = AArch64::XZR;
+  if (!isPACWithZeroDisc(MI.getOpcode())) {
+    DiscReg = MRI.createVirtualRegister(&AArch64::GPR64noipRegClass);
+    BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), DiscReg)
+        .addReg(MI.getOperand(2).getReg());
+  }
+
+  BuildMI(*BB, MI, DL, TII->get(NewOpcode))
+      .addGlobalAddress(AddrOp.getGlobal(), Offset, TargetFlags)
+      .addImm(getKeyForPACOpcode(MI.getOpcode()))
+      .addReg(DiscReg)
+      .addImm(0);
+
+  BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), MI.getOperand(0).getReg())
+      .addReg(AArch64::X16);
+
+  MI.removeFromParent();
   return BB;
 }
 
@@ -3159,6 +3231,16 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
     return EmitZTInstr(MI, BB, AArch64::ZERO_T, /*Op0IsDef=*/true);
   case AArch64::MOVT_TIZ_PSEUDO:
     return EmitZTInstr(MI, BB, AArch64::MOVT_TIZ, /*Op0IsDef=*/true);
+
+  case AArch64::PACDA:
+  case AArch64::PACDB:
+  case AArch64::PACIA:
+  case AArch64::PACIB:
+  case AArch64::PACDZA:
+  case AArch64::PACDZB:
+  case AArch64::PACIZA:
+  case AArch64::PACIZB:
+    return tryRewritingPAC(MI, BB);
   }
 }
 

llvm/lib/Target/AArch64/AArch64ISelLowering.h

@@ -181,6 +181,8 @@ class AArch64TargetLowering : public TargetLowering {
                                                MachineBasicBlock *BB) const;
   MachineBasicBlock *EmitGetSMESaveSize(MachineInstr &MI,
                                         MachineBasicBlock *BB) const;
+  MachineBasicBlock *tryRewritingPAC(MachineInstr &MI,
+                                     MachineBasicBlock *BB) const;
 
   MachineBasicBlock *
   EmitInstrWithCustomInserter(MachineInstr &MI,

llvm/lib/Target/AArch64/AArch64InstrInfo.h

@@ -790,6 +790,39 @@ static inline unsigned getPACOpcodeForKey(AArch64PACKey::ID K, bool Zero) {
   llvm_unreachable("Unhandled AArch64PACKey::ID enum");
 }
 
+static inline AArch64PACKey::ID getKeyForPACOpcode(unsigned Opcode) {
+  switch (Opcode) {
+  case AArch64::PACDA:
+  case AArch64::PACDZA:
+    return AArch64PACKey::DA;
+  case AArch64::PACDB:
+  case AArch64::PACDZB:
+    return AArch64PACKey::DB;
+  case AArch64::PACIA:
+  case AArch64::PACIZA:
+    return AArch64PACKey::IA;
+  case AArch64::PACIB:
+  case AArch64::PACIZB:
+    return AArch64PACKey::IB;
+  }
+  llvm_unreachable("Unhandled PAC opcode");
+}
+
+static inline bool isPACWithZeroDisc(unsigned Opcode) {
+  switch (Opcode) {
+  case AArch64::PACDA:
+  case AArch64::PACDB:
+  case AArch64::PACIA:
+  case AArch64::PACIB:
+    return false;
+  case AArch64::PACDZA:
+  case AArch64::PACDZB:
+  case AArch64::PACIZA:
+  case AArch64::PACIZB:
+    return true;
+  }
+  llvm_unreachable("Unhandled PAC opcode");
+}
 // struct TSFlags {
 #define TSFLAG_ELEMENT_SIZE_TYPE(X)      (X)        // 3-bits
 #define TSFLAG_DESTRUCTIVE_INST_TYPE(X) ((X) << 3)  // 4-bits

llvm/lib/Target/AArch64/AArch64InstrInfo.td

@@ -2018,7 +2018,9 @@ let Predicates = [HasPAuth] in {
     def DZB  : SignAuthZero<prefix_z,  0b11, !strconcat(asm, "dzb"), op>;
   }
 
+  let usesCustomInserter = true in
   defm PAC : SignAuth<0b000, 0b010, "pac", int_ptrauth_sign>;
+
   defm AUT : SignAuth<0b001, 0b011, "aut", null_frag>;
 
   def XPACI : ClearAuth<0, "xpaci">;

llvm/test/CodeGen/AArch64/GlobalISel/ptrauth-constant-in-code.ll

@@ -96,7 +96,9 @@ define void @store_signed_const_local(ptr %dest) {
 ; ISEL:         %0:gpr64common = COPY $x0
 ; ISEL-NEXT:    %10:gpr64common = MOVaddr target-flags(aarch64-page) @const_table_local + 8, target-flags(aarch64-pageoff, aarch64-nc) @const_table_local + 8
 ; ISEL-NEXT:    %2:gpr64common = PAUTH_BLEND %0, 1234
-; ISEL-NEXT:    %4:gpr64 = PACDA %10, %2
+; ISEL-NEXT:    %15:gpr64noip = COPY %2
+; ISEL-NEXT:    MOVaddrPAC @const_table_local + 8, 2, %15, 0, implicit-def $x16, implicit-def $x17
+; ISEL-NEXT:    %4:gpr64 = COPY $x16
 ; ISEL-NEXT:    %14:gpr64 = COPY %4
 ; ISEL-NEXT:    STRXui %14, %0, 0 :: (store (p0) into %ir.dest)
 ; ISEL-NEXT:    RET_ReallyLR
@@ -115,7 +117,9 @@ define void @store_signed_const_got(ptr %dest) {
 ; ISEL-ELF-NEXT:    %7:gpr64common = LOADgotAUTH target-flags(aarch64-got) @const_table_got
 ; ISEL-ELF-NEXT:    %6:gpr64common = ADDXri %7, 8, 0
 ; ISEL-ELF-NEXT:    %2:gpr64common = PAUTH_BLEND %0, 1234
-; ISEL-ELF-NEXT:    %4:gpr64 = PACDA %6, %2
+; ISEL-ELF-NEXT:    %12:gpr64noip = COPY %2
+; ISEL-ELF-NEXT:    LOADgotPAC target-flags(aarch64-got) @const_table_got + 8, 2, %12, 0, implicit-def $x16, implicit-def $x17, implicit-def $nzcv
+; ISEL-ELF-NEXT:    %4:gpr64 = COPY $x16
 ; ISEL-ELF-NEXT:    %10:gpr64 = COPY %4
 ; ISEL-ELF-NEXT:    STRXui %10, %0, 0 :: (store (p0) into %ir.dest)
 ; ISEL-ELF-NEXT:    RET_ReallyLR

llvm/test/CodeGen/AArch64/ptrauth-constant-in-code.ll

@@ -85,7 +85,9 @@ define void @store_signed_const_local(ptr %dest) {
 ; ISEL:         %0:gpr64common = COPY $x0
 ; ISEL-NEXT:    %1:gpr64common = PAUTH_BLEND %0, 1234
 ; ISEL-NEXT:    %2:gpr64common = MOVaddr target-flags(aarch64-page) @const_table_local + 8, target-flags(aarch64-pageoff, aarch64-nc) @const_table_local + 8
-; ISEL-NEXT:    %3:gpr64 = PACDA %2, killed %1
+; ISEL-NEXT:    %4:gpr64noip = COPY %1
+; ISEL-NEXT:    MOVaddrPAC @const_table_local + 8, 2, %4, 0, implicit-def $x16, implicit-def $x17
+; ISEL-NEXT:    %3:gpr64 = COPY $x16
 ; ISEL-NEXT:    STRXui killed %3, %0, 0 :: (store (s64) into %ir.dest)
 ; ISEL-NEXT:    RET_ReallyLR
   %dest.i = ptrtoint ptr %dest to i64
@@ -103,7 +105,9 @@ define void @store_signed_const_got(ptr %dest) {
 ; ISEL-ELF-NEXT:    %1:gpr64common = PAUTH_BLEND %0, 1234
 ; ISEL-ELF-NEXT:    %2:gpr64common = LOADgotAUTH target-flags(aarch64-got) @const_table_got, implicit-def dead $x16, implicit-def dead $x17, implicit-def dead $nzcv
 ; ISEL-ELF-NEXT:    %3:gpr64common = ADDXri killed %2, 8, 0
-; ISEL-ELF-NEXT:    %4:gpr64 = PACDA %3, killed %1
+; ISEL-ELF-NEXT:    %5:gpr64noip = COPY %1
+; ISEL-ELF-NEXT:    LOADgotPAC target-flags(aarch64-got) @const_table_got + 8, 2, %5, 0, implicit-def $x16, implicit-def $x17, implicit-def $nzcv
+; ISEL-ELF-NEXT:    %4:gpr64 = COPY $x16
 ; ISEL-ELF-NEXT:    STRXui killed %4, %0, 0 :: (store (s64) into %ir.dest)
 ; ISEL-ELF-NEXT:    RET_ReallyLR
   %dest.i = ptrtoint ptr %dest to i64