[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

@@ -3101,10 +3101,13 @@ AArch64TargetLowering::EmitGetSMESaveSize(MachineInstr &MI,
   return BB;
 }
 
-// Helper function to find the instruction that defined a virtual register.
-// If unable to find such instruction, returns nullptr.
-static const MachineInstr *stripVRegCopies(const MachineRegisterInfo &MRI,
-                                           Register Reg) {
+// Helper function to find the instruction that defined a virtual register,
+// stripping and accumulating optional offset.
+// If unable to find such instruction, returns nullptr (Offset is unspecified).
+static const MachineInstr *
+stripAndAccumulateOffset(const MachineRegisterInfo &MRI, Register Reg,
+                         int64_t &Offset) {
+  Offset = 0;
   while (Reg.isVirtual()) {
     MachineInstr *DefMI = MRI.getVRegDef(Reg);
     assert(DefMI && "Virtual register definition not found");
@@ -3122,7 +3125,17 @@ static const MachineInstr *stripVRegCopies(const MachineRegisterInfo &MRI,
       continue;
     }
 
-    return DefMI;
+    // If this is neither a copy, nor inc/dec instruction, we are done.
+    if (Opcode != AArch64::ADDXri && Opcode != AArch64::SUBXri)
+      return DefMI;
+    // Inc/dec with shifted immediates are not handled.
+    if (DefMI->getOperand(3).getImm() != 0)
+      return DefMI;
+
+    int64_t Imm = DefMI->getOperand(2).getImm();
+    Offset += (Opcode == AArch64::ADDXri) ? Imm : -Imm;
+
+    Reg = DefMI->getOperand(1).getReg();
   }
   return nullptr;
 }
@@ -3138,8 +3151,13 @@ void AArch64TargetLowering::fixupPtrauthDiscriminator(
   int64_t IntDisc = IntDiscOp.getImm();
   assert(IntDisc == 0 && "Blend components are already expanded");
 
-  const MachineInstr *DiscMI = stripVRegCopies(MRI, AddrDisc);
-  if (DiscMI) {
+  int64_t Offset = 0;
+  const MachineInstr *DiscMI = stripAndAccumulateOffset(MRI, AddrDisc, Offset);
+
+  // The result of any recognized discriminator computation may be copied, but
+  // without adding any offset. Nevertheless, perform the remaining fix-ups
+  // even on an opaque, pre-computed discriminator.
+  if (DiscMI && Offset == 0) {
     switch (DiscMI->getOpcode()) {
     case AArch64::MOVKXi:
       // blend(addr, imm) which is lowered as "MOVK addr, #imm, #48".
@@ -3178,6 +3196,53 @@ void AArch64TargetLowering::fixupPtrauthDiscriminator(
   IntDiscOp.setImm(IntDisc);
 }
 
+MachineBasicBlock *
+AArch64TargetLowering::tryRewritingPAC(MachineInstr &MI,
+                                       MachineBasicBlock *BB) const {
+  const TargetInstrInfo *TII = Subtarget->getInstrInfo();
+  MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
+  const DebugLoc &DL = MI.getDebugLoc();
+
+  Register Val = MI.getOperand(1).getReg();
+  unsigned Key = MI.getOperand(2).getImm();
+  int64_t IntDisc = MI.getOperand(3).getImm();
+  Register AddrDisc = MI.getOperand(4).getReg();
+
+  // Try to find a known address-setting instruction, accumulating the offset
+  // along the way. If no known pattern is found, keep everything as-is.
+
+  int64_t AddrOffset = 0;
+  const MachineInstr *AddrDefInstr =
+      stripAndAccumulateOffset(MRI, Val, AddrOffset);
+  if (!AddrDefInstr)
+    return BB;
+
+  unsigned NewOpcode;
+  if (AddrDefInstr->getOpcode() == AArch64::LOADgotAUTH)
+    NewOpcode = AArch64::LOADgotPAC;
+  else if (AddrDefInstr->getOpcode() == AArch64::MOVaddr)
+    NewOpcode = AArch64::MOVaddrPAC;
+  else
+    return BB; // Unknown opcode.
+
+  const MachineOperand &AddrOp = AddrDefInstr->getOperand(1);
+  unsigned TargetFlags = AddrOp.getTargetFlags() & ~AArch64II::MO_PAGE;
+  const GlobalValue *GV = AddrOp.getGlobal();
+  AddrOffset += AddrOp.getOffset();
+
+  BuildMI(*BB, MI, DL, TII->get(NewOpcode))
+      .addGlobalAddress(GV, AddrOffset, TargetFlags)
+      .addImm(Key)
+      .addReg(AddrDisc)
+      .addImm(IntDisc);
+
+  BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), MI.getOperand(0).getReg())
+      .addReg(AArch64::X16);
+
+  MI.removeFromParent();
+  return BB;
+}
+
 MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
     MachineInstr &MI, MachineBasicBlock *BB) const {
 
@@ -3288,7 +3353,7 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
   case AArch64::PAC:
     fixupPtrauthDiscriminator(MI, BB, MI.getOperand(3), MI.getOperand(4),
                               &AArch64::GPR64noipRegClass);
-    return BB;
+    return tryRewritingPAC(MI, BB);
   case AArch64::AUTPAC:
     fixupPtrauthDiscriminator(MI, BB, MI.getOperand(1), MI.getOperand(2),
                               &AArch64::GPR64noipRegClass);

llvm/lib/Target/AArch64/AArch64ISelLowering.h

@@ -189,6 +189,9 @@ class AArch64TargetLowering : public TargetLowering {
                                  MachineOperand &AddrDiscOp,
                                  const TargetRegisterClass *AddrDiscRC) const;
 
+  MachineBasicBlock *tryRewritingPAC(MachineInstr &MI,
+                                     MachineBasicBlock *BB) const;
+
   MachineBasicBlock *
   EmitInstrWithCustomInserter(MachineInstr &MI,
                               MachineBasicBlock *MBB) const override;

llvm/lib/Target/AArch64/AArch64InstrInfo.td

@@ -2218,6 +2218,9 @@ let Predicates = [HasPAuth] in {
 
   def LOADgotAUTH : Pseudo<(outs GPR64common:$dst), (ins i64imm:$addr), []>,
                Sched<[WriteI, ReadI]> {
+    // Make it possible to eliminate dead instruction after folding it
+    // into LOADgotPAC.
+    let hasSideEffects = 0;
     let Defs = [X16,X17,NZCV];
     let Size = 44;
   }

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

@@ -82,25 +82,50 @@ define ptr @foo() {
 
 ; RUN: llc < finalize-isel.ll -mtriple aarch64-elf -mattr=+pauth -global-isel=1 \
 ; RUN:   -verify-machineinstrs -global-isel-abort=1 -stop-after=finalize-isel | \
-; RUN:   FileCheck --check-prefixes=ISEL,ISEL-ELF %s
+; RUN:   FileCheck --check-prefixes=ISEL-MIR,ISEL-MIR-ELF %s
 ; RUN: llc < finalize-isel.ll -mtriple arm64-apple-ios -mattr=+pauth -global-isel=1 \
 ; RUN:   -verify-machineinstrs -global-isel-abort=1 -stop-after=finalize-isel | \
-; RUN:   FileCheck --check-prefixes=ISEL %s
+; RUN:   FileCheck --check-prefixes=ISEL-MIR %s
+; RUN: llc < finalize-isel.ll -mtriple aarch64-elf -mattr=+pauth -global-isel=1 \
+; RUN:   -verify-machineinstrs -global-isel-abort=1 -asm-verbose=0 | \
+; RUN:   FileCheck --check-prefixes=ISEL-ASM,ISEL-ASM-ELF %s
+; RUN: llc < finalize-isel.ll -mtriple arm64-apple-ios -mattr=+pauth -global-isel=1 \
+; RUN:   -verify-machineinstrs -global-isel-abort=1 -asm-verbose=0 | \
+; RUN:   FileCheck --check-prefixes=ISEL-ASM,ISEL-ASM-MACHO %s
 
 @const_table_local = dso_local constant [3 x ptr] [ptr null, ptr null, ptr null]
 @const_table_got = constant [3 x ptr] [ptr null, ptr null, ptr null]
 
+; Test that after post-processing in finalize-isel, MOVaddrPAC (or LOADgotPAC,
+; respectively) has both $AddrDisc and $Disc operands set. MOVaddr (or LOADgotAUTH,
+; respectively) and MOVKXi are not used anymore and are dead-code-eliminated
+; by the later passes.
+
 define void @store_signed_const_local(ptr %dest) {
-; ISEL-LABEL: name: store_signed_const_local
-; ISEL:       body:
-; 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:gpr64noip = MOVKXi %0, 1234
-; ISEL-NEXT:    %15:gpr64noip = COPY %0
-; ISEL-NEXT:    %4:gpr64 = PAC %10, 2, 1234, %15, implicit-def dead $x16, implicit-def dead $x17
-; ISEL-NEXT:    %14:gpr64 = COPY %4
-; ISEL-NEXT:    STRXui %14, %0, 0 :: (store (p0) into %ir.dest)
-; ISEL-NEXT:    RET_ReallyLR
+; ISEL-MIR-LABEL: name: store_signed_const_local
+; ISEL-MIR:       body:
+; ISEL-MIR:         %0:gpr64common = COPY $x0
+; ISEL-MIR-NEXT:    %10:gpr64common = MOVaddr target-flags(aarch64-page) @const_table_local + 8, target-flags(aarch64-pageoff, aarch64-nc) @const_table_local + 8
+; ISEL-MIR-NEXT:    %2:gpr64noip = MOVKXi %0, 1234, 48
+; ISEL-MIR-NEXT:    %15:gpr64noip = COPY %0
+; ISEL-MIR-NEXT:    MOVaddrPAC @const_table_local + 8, 2, %15, 1234, implicit-def $x16, implicit-def $x17
+; ISEL-MIR-NEXT:    %4:gpr64 = COPY $x16
+; ISEL-MIR-NEXT:    %14:gpr64 = COPY %4
+; ISEL-MIR-NEXT:    STRXui %14, %0, 0 :: (store (p0) into %ir.dest)
+; ISEL-MIR-NEXT:    RET_ReallyLR
+;
+; ISEL-ASM-LABEL: store_signed_const_local:
+; ISEL-ASM-NEXT:    .cfi_startproc
+; ISEL-ASM-ELF-NEXT:    adrp    x16, const_table_local
+; ISEL-ASM-ELF-NEXT:    add     x16, x16, :lo12:const_table_local
+; ISEL-ASM-MACHO-NEXT:  adrp    x16, _const_table_local@PAGE
+; ISEL-ASM-MACHO-NEXT:  add     x16, x16, _const_table_local@PAGEOFF
+; ISEL-ASM-NEXT:    add     x16, x16, #8
+; ISEL-ASM-NEXT:    mov     x17, x0
+; ISEL-ASM-NEXT:    movk    x17, #1234, lsl #48
+; ISEL-ASM-NEXT:    pacda   x16, x17
+; ISEL-ASM-NEXT:    str     x16, [x0]
+; ISEL-ASM-NEXT:    ret
   %dest.i = ptrtoint ptr %dest to i64
   %discr = call i64 @llvm.ptrauth.blend(i64 %dest.i, i64 1234)
   %signed.i = call i64 @llvm.ptrauth.sign(i64 ptrtoint (ptr getelementptr ([2 x ptr], ptr @const_table_local, i32 0, i32 1) to i64), i32 2, i64 %discr)
@@ -110,17 +135,37 @@ define void @store_signed_const_local(ptr %dest) {
 }
 
 define void @store_signed_const_got(ptr %dest) {
-; ISEL-ELF-LABEL: name: store_signed_const_got
-; ISEL-ELF:       body:
-; ISEL-ELF:         %0:gpr64common = COPY $x0
-; 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:gpr64noip = MOVKXi %0, 1234
-; ISEL-ELF-NEXT:    %12:gpr64noip = COPY %0
-; ISEL-ELF-NEXT:    %4:gpr64 = PAC %6, 2, 1234, %12, implicit-def dead $x16, implicit-def dead $x17
-; ISEL-ELF-NEXT:    %10:gpr64 = COPY %4
-; ISEL-ELF-NEXT:    STRXui %10, %0, 0 :: (store (p0) into %ir.dest)
-; ISEL-ELF-NEXT:    RET_ReallyLR
+; ISEL-MIR-ELF-LABEL: name: store_signed_const_got
+; ISEL-MIR-ELF:       body:
+; ISEL-MIR-ELF:         %0:gpr64common = COPY $x0
+; ISEL-MIR-ELF-NEXT:    %7:gpr64common = LOADgotAUTH target-flags(aarch64-got) @const_table_got
+; ISEL-MIR-ELF-NEXT:    %6:gpr64common = ADDXri %7, 8, 0
+; ISEL-MIR-ELF-NEXT:    %2:gpr64noip = MOVKXi %0, 1234, 48
+; ISEL-MIR-ELF-NEXT:    %12:gpr64noip = COPY %0
+; ISEL-MIR-ELF-NEXT:    LOADgotPAC target-flags(aarch64-got) @const_table_got + 8, 2, %12, 1234, implicit-def $x16, implicit-def $x17, implicit-def $nzcv
+; ISEL-MIR-ELF-NEXT:    %4:gpr64 = COPY $x16
+; ISEL-MIR-ELF-NEXT:    %10:gpr64 = COPY %4
+; ISEL-MIR-ELF-NEXT:    STRXui %10, %0, 0 :: (store (p0) into %ir.dest)
+; ISEL-MIR-ELF-NEXT:    RET_ReallyLR
+;
+; ISEL-ASM-ELF-LABEL: store_signed_const_got:
+; ISEL-ASM-ELF-NEXT:    .cfi_startproc
+; ISEL-ASM-ELF-NEXT:    adrp    x17, :got_auth:const_table_got
+; ISEL-ASM-ELF-NEXT:    add     x17, x17, :got_auth_lo12:const_table_got
+; ISEL-ASM-ELF-NEXT:    ldr     x16, [x17]
+; ISEL-ASM-ELF-NEXT:    autda   x16, x17
+; ISEL-ASM-ELF-NEXT:    mov     x17, x16
+; ISEL-ASM-ELF-NEXT:    xpacd   x17
+; ISEL-ASM-ELF-NEXT:    cmp     x16, x17
+; ISEL-ASM-ELF-NEXT:    b.eq    .Lauth_success_0
+; ISEL-ASM-ELF-NEXT:    brk     #0xc472
+; ISEL-ASM-ELF-NEXT:  .Lauth_success_0:
+; ISEL-ASM-ELF-NEXT:    add     x16, x16, #8
+; ISEL-ASM-ELF-NEXT:    mov     x17, x0
+; ISEL-ASM-ELF-NEXT:    movk    x17, #1234, lsl #48
+; ISEL-ASM-ELF-NEXT:    pacda   x16, x17
+; ISEL-ASM-ELF-NEXT:    str     x16, [x0]
+; ISEL-ASM-ELF-NEXT:    ret
   %dest.i = ptrtoint ptr %dest to i64
   %discr = call i64 @llvm.ptrauth.blend(i64 %dest.i, i64 1234)
   %signed.i = call i64 @llvm.ptrauth.sign(i64 ptrtoint (ptr getelementptr ([2 x ptr], ptr @const_table_got, i32 0, i32 1) to i64), i32 2, i64 %discr)
@@ -130,17 +175,26 @@ define void @store_signed_const_got(ptr %dest) {
 }
 
 define void @store_signed_arg(ptr %dest, ptr %p) {
-; ISEL-LABEL: name: store_signed_arg
-; ISEL:       body:
-; ISEL:         %0:gpr64common = COPY $x0
-; ISEL-NEXT:    %1:gpr64common = COPY $x1
-; ISEL-NEXT:    %3:gpr64noip = MOVKXi %0, 1234
-; ISEL-NEXT:    %6:gpr64common = ADDXri %1, 8, 0
-; ISEL-NEXT:    %12:gpr64noip = COPY %0
-; ISEL-NEXT:    %8:gpr64 = PAC %6, 2, 1234, %12, implicit-def dead $x16, implicit-def dead $x17
-; ISEL-NEXT:    %10:gpr64 = COPY %8
-; ISEL-NEXT:    STRXui %10, %0, 0 :: (store (p0) into %ir.dest)
-; ISEL-NEXT:    RET_ReallyLR
+; ISEL-MIR-LABEL: name: store_signed_arg
+; ISEL-MIR:       body:
+; ISEL-MIR:         %0:gpr64common = COPY $x0
+; ISEL-MIR-NEXT:    %1:gpr64common = COPY $x1
+; ISEL-MIR-NEXT:    %3:gpr64noip = MOVKXi %0, 1234, 48
+; ISEL-MIR-NEXT:    %6:gpr64common = ADDXri %1, 8, 0
+; ISEL-MIR-NEXT:    %12:gpr64noip = COPY %0
+; ISEL-MIR-NEXT:    %8:gpr64 = PAC %6, 2, 1234, %12, implicit-def dead $x16, implicit-def dead $x17
+; ISEL-MIR-NEXT:    %10:gpr64 = COPY %8
+; ISEL-MIR-NEXT:    STRXui %10, %0, 0 :: (store (p0) into %ir.dest)
+; ISEL-MIR-NEXT:    RET_ReallyLR
+;
+; ISEL-ASM-LABEL: store_signed_arg:
+; ISEL-ASM-NEXT:    .cfi_startproc
+; ISEL-ASM-NEXT:    add     x8, x1, #8
+; ISEL-ASM-NEXT:    mov     x16, x0
+; ISEL-ASM-NEXT:    movk    x16, #1234, lsl #48
+; ISEL-ASM-NEXT:    pacda   x8, x16
+; ISEL-ASM-NEXT:    str     x8, [x0]
+; ISEL-ASM-NEXT:    ret
   %dest.i = ptrtoint ptr %dest to i64
   %discr = call i64 @llvm.ptrauth.blend(i64 %dest.i, i64 1234)
   %p.offset = getelementptr [2 x ptr], ptr %p, i32 0, i32 1