[AArch64][PAC] Combine signing with address materialization

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -3058,6 +3058,122 @@ AArch64TargetLowering::EmitGetSMESaveSize(MachineInstr &MI,
             MI.getOperand(0).getReg())
         .addReg(AArch64::XZR);
   BB->remove_instr(&MI);
+
+  return BB;
+}
+
+// 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 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");
+    unsigned Opcode = DefMI->getOpcode();
+
+    if (Opcode == AArch64::COPY) {
+      Reg = DefMI->getOperand(1).getReg();
+      continue;
+    }
+
+    // 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;
+}
+
+void AArch64TargetLowering::fixupBlendComponents(
+    MachineInstr &MI, MachineBasicBlock *BB, MachineOperand &IntDiscOp,
+    MachineOperand &AddrDiscOp, const TargetRegisterClass *AddrDiscRC) const {
+  const TargetInstrInfo *TII = Subtarget->getInstrInfo();
+  MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
+  const DebugLoc &DL = MI.getDebugLoc();
+
+  Register AddrDisc = AddrDiscOp.getReg();
+  int64_t IntDisc = IntDiscOp.getImm();
+
+  assert(IntDisc == 0 && "Blend components are already expanded");
+
+  int64_t Offset = 0;
+  MachineInstr *MaybeBlend = stripAndAccumulateOffset(MRI, AddrDisc, Offset);
+
+  // Detect blend(addr, imm) which is lowered as MOVK addr, #imm, #48.
+  // The result of MOVK may be copied, but without adding any offset.
+  if (MaybeBlend && Offset == 0 && MaybeBlend->getOpcode() == AArch64::MOVKXi &&
+      MaybeBlend->getOperand(3).getImm() == 48) {
+    AddrDisc = MaybeBlend->getOperand(1).getReg();
+    IntDisc = MaybeBlend->getOperand(2).getImm();
+  }
+
+  if (AddrDisc == AArch64::NoRegister)
+    AddrDisc = AArch64::XZR;
+
+  if (AddrDisc != AArch64::XZR && MRI.getRegClass(AddrDisc) != AddrDiscRC) {
+    Register TmpReg = MRI.createVirtualRegister(AddrDiscRC);
+    BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), TmpReg).addReg(AddrDisc);
+    AddrDisc = TmpReg;
+  }
+
+  AddrDiscOp.setReg(AddrDisc);
+  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();
+
+  // 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;
+  MachineInstr *AddrDefInstr =
+      stripAndAccumulateOffset(MRI, MI.getOperand(1).getReg(), 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.
+
+  MachineOperand &AddrOp = AddrDefInstr->getOperand(1);
+  unsigned TargetFlags = AddrOp.getTargetFlags() & ~AArch64II::MO_PAGE;
+  const GlobalValue *GV = AddrOp.getGlobal();
+  AddrOffset += AddrOp.getOffset();
+
+  Register DiscReg = isPACWithZeroDisc(MI.getOpcode())
+                         ? AArch64::XZR
+                         : MI.getOperand(2).getReg();
+
+  MachineInstr *NewMI = BuildMI(*BB, MI, DL, TII->get(NewOpcode))
+                            .addGlobalAddress(GV, AddrOffset, TargetFlags)
+                            .addImm(getKeyForPACOpcode(MI.getOpcode()))
+                            .addReg(DiscReg)
+                            .addImm(0);
+  fixupBlendComponents(*NewMI, BB, NewMI->getOperand(3), NewMI->getOperand(2),
+                       &AArch64::GPR64noipRegClass);
+
+  BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), MI.getOperand(0).getReg())
+      .addReg(AArch64::X16);
+
+  MI.removeFromParent();
   return BB;
 }
 
@@ -3159,6 +3275,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

@@ -182,6 +182,14 @@ class AArch64TargetLowering : public TargetLowering {
   MachineBasicBlock *EmitGetSMESaveSize(MachineInstr &MI,
                                         MachineBasicBlock *BB) const;
 
+  void fixupBlendComponents(MachineInstr &MI, MachineBasicBlock *BB,
+                            MachineOperand &IntDiscOp,
+                            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.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">;
@@ -2166,6 +2168,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

@@ -78,6 +78,136 @@ define ptr @foo() {
   ret ptr ptrauth (ptr @g, i32 0)
 }
 
+;--- finalize-isel.ll
+
+; 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-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-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-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:gpr64common = 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)
+  %signed.ptr = inttoptr i64 %signed.i to ptr
+  store ptr %signed.ptr, ptr %dest
+  ret void
+}
+
+define void @store_signed_const_got(ptr %dest) {
+; 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:gpr64common = 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)
+  %signed.ptr = inttoptr i64 %signed.i to ptr
+  store ptr %signed.ptr, ptr %dest
+  ret void
+}
+
+define void @store_signed_arg(ptr %dest, ptr %p) {
+; 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:gpr64common = MOVKXi %0, 1234, 48
+; ISEL-MIR-NEXT:    %6:gpr64common = ADDXri %1, 8, 0
+; Check that no implicit defs are added to PACDA instruction.
+; ISEL-MIR-NEXT:    %8:gpr64 = PACDA %6, %3{{$}}
+; 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:    mov     x8, x0
+; ISEL-ASM-NEXT:    add     x9, x1, #8
+; ISEL-ASM-NEXT:    movk    x8, #1234, lsl #48
+; ISEL-ASM-NEXT:    pacda   x9, x8
+; ISEL-ASM-NEXT:    str     x9, [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
+  %p.offset.i = ptrtoint ptr %p.offset to i64
+  %signed.i = call i64 @llvm.ptrauth.sign(i64 %p.offset.i, i32 2, i64 %discr)
+  %signed.ptr = inttoptr i64 %signed.i to ptr
+  store ptr %signed.ptr, ptr %dest
+  ret void
+}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 8, !"ptrauth-elf-got", i32 1}
+
 ;--- ok.ll
 
 ; RUN: llc < ok.ll -mtriple aarch64-elf -mattr=+pauth -global-isel=1 \