Skip to content

[AArch64][PAC] Combine signing with address materialization #130809

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: users/atrosinenko/pauth-imm-modifier-auth-resign
Choose a base branch
from
Open

[AArch64][PAC] Combine signing with address materialization #130809

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
cf9b917
[AArch64][PAC] Precommit tests on merging MOVaddr/LOADgotAUTH with PAC*
atrosinenko Mar 10, 2025
3af1146
[AArch64][PAC] Combine signing with address materialization
atrosinenko Jun 30, 2025
3c4236f
Transform LOADgot in addition to LOADgotAUTH
atrosinenko Jul 23, 2025
62a98de
Handle non-register operands correctly
atrosinenko Sep 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 77 additions & 8 deletions llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand All @@ -3122,7 +3125,20 @@ 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;
// Handle cases like `ADDXri %stack.0.local_var, 0, 0`.
if (!DefMI->getOperand(1).isReg())
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;
}
Expand All @@ -3138,8 +3154,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".
Expand Down Expand Up @@ -3178,6 +3199,54 @@ 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::LOADgot ||
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 {

Expand Down Expand Up @@ -3288,7 +3357,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);
Expand Down
3 changes: 3 additions & 0 deletions llvm/lib/Target/AArch64/AArch64ISelLowering.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
3 changes: 3 additions & 0 deletions llvm/lib/Target/AArch64/AArch64InstrInfo.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
176 changes: 176 additions & 0 deletions llvm/test/CodeGen/AArch64/GlobalISel/ptrauth-constant-in-code.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,182 @@ 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: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)
%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 = LOADgot 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 x16, :got:const_table_got
; ISEL-ASM-ELF-NEXT: ldr x16, [x16, :got_lo12:const_table_got]
; 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: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
%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
}

;--- finalize-isel-elf-got.ll

; RUN: llc < finalize-isel-elf-got.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-ELF-GOT-MIR %s
; RUN: llc < finalize-isel-elf-got.ll -mtriple aarch64-elf -mattr=+pauth -global-isel=1 \
; RUN: -verify-machineinstrs -global-isel-abort=1 -asm-verbose=0 | \
; RUN: FileCheck --check-prefixes=ISEL-ELF-GOT-ASM %s

@const_table_got = constant [3 x ptr] [ptr null, ptr null, ptr null]

; Similar to finalize-isel.ll, but tests conversion of LOADgotAUTH to LOADgotPAC.
; This requires module-level attribute, thus a separate sub-file.

define void @store_signed_const_got(ptr %dest) {
; ISEL-ELF-GOT-MIR-LABEL: name: store_signed_const_got
; ISEL-ELF-GOT-MIR: body:
; ISEL-ELF-GOT-MIR: %0:gpr64common = COPY $x0
; ISEL-ELF-GOT-MIR-NEXT: %7:gpr64common = LOADgotAUTH target-flags(aarch64-got) @const_table_got
; ISEL-ELF-GOT-MIR-NEXT: %6:gpr64common = ADDXri %7, 8, 0
; ISEL-ELF-GOT-MIR-NEXT: %2:gpr64noip = MOVKXi %0, 1234, 48
; ISEL-ELF-GOT-MIR-NEXT: %12:gpr64noip = COPY %0
; ISEL-ELF-GOT-MIR-NEXT: LOADgotPAC target-flags(aarch64-got) @const_table_got + 8, 2, %12, 1234, implicit-def $x16, implicit-def $x17, implicit-def $nzcv
; ISEL-ELF-GOT-MIR-NEXT: %4:gpr64 = COPY $x16
; ISEL-ELF-GOT-MIR-NEXT: %10:gpr64 = COPY %4
; ISEL-ELF-GOT-MIR-NEXT: STRXui %10, %0, 0 :: (store (p0) into %ir.dest)
; ISEL-ELF-GOT-MIR-NEXT: RET_ReallyLR
;
; ISEL-ELF-GOT-ASM-LABEL: store_signed_const_got:
; ISEL-ELF-GOT-ASM-NEXT: .cfi_startproc
; ISEL-ELF-GOT-ASM-NEXT: adrp x17, :got_auth:const_table_got
; ISEL-ELF-GOT-ASM-NEXT: add x17, x17, :got_auth_lo12:const_table_got
; ISEL-ELF-GOT-ASM-NEXT: ldr x16, [x17]
; ISEL-ELF-GOT-ASM-NEXT: autda x16, x17
; ISEL-ELF-GOT-ASM-NEXT: mov x17, x16
; ISEL-ELF-GOT-ASM-NEXT: xpacd x17
; ISEL-ELF-GOT-ASM-NEXT: cmp x16, x17
; ISEL-ELF-GOT-ASM-NEXT: b.eq .Lauth_success_0
; ISEL-ELF-GOT-ASM-NEXT: brk #0xc472
; ISEL-ELF-GOT-ASM-NEXT: .Lauth_success_0:
; ISEL-ELF-GOT-ASM-NEXT: add x16, x16, #8
; ISEL-ELF-GOT-ASM-NEXT: mov x17, x0
; ISEL-ELF-GOT-ASM-NEXT: movk x17, #1234, lsl #48
; ISEL-ELF-GOT-ASM-NEXT: pacda x16, x17
; ISEL-ELF-GOT-ASM-NEXT: str x16, [x0]
; ISEL-ELF-GOT-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_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
}

!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 \
Expand Down
Loading