Skip to content

[AArch64][PAC] Fix clobbering registers by BLRA and AUTH_TCRETURN #155373

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

Merged
merged 2 commits into from
Aug 27, 2025
Merged

[AArch64][PAC] Fix clobbering registers by BLRA and AUTH_TCRETURN #155373

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ce72add
[AArch64][PAC] Precommit test cases for BLRA/AUTH_TCRETURN clobbering…
atrosinenko Aug 25, 2025
b4fc557
[AArch64][PAC] Fix clobbering registers by BLRA and AUTH_TCRETURN
atrosinenko Aug 25, 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
28 changes: 23 additions & 5 deletions llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2229,13 +2229,24 @@ void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
if (BrTarget == AddrDisc)
report_fatal_error("Branch target is signed with its own value");

// If we are printing BLRA pseudo instruction, then x16 and x17 are
// implicit-def'ed by the MI and AddrDisc is not used as any other input, so
// try to save one MOV by setting MayUseAddrAsScratch.
// If we are printing BLRA pseudo, try to save one MOV by making use of the
// fact that x16 and x17 are described as clobbered by the MI instruction and
// AddrDisc is not used as any other input.
//
// Back in the day, emitPtrauthDiscriminator was restricted to only returning
// either x16 or x17, meaning the returned register is always among the
// implicit-def'ed registers of BLRA pseudo. Now this property can be violated
// if isX16X17Safer predicate is false, thus manually check if AddrDisc is
// among x16 and x17 to prevent clobbering unexpected registers.
//
// Unlike BLRA, BRA pseudo is used to perform computed goto, and thus not
// declared as clobbering x16/x17.
//
// FIXME: Make use of `killed` flags and register masks instead.
bool AddrDiscIsImplicitDef =
IsCall && (AddrDisc == AArch64::X16 || AddrDisc == AArch64::X17);
Register DiscReg = emitPtrauthDiscriminator(Disc, AddrDisc, AArch64::X17,
/*MayUseAddrAsScratch=*/IsCall);
AddrDiscIsImplicitDef);
bool IsZeroDisc = DiscReg == AArch64::XZR;

unsigned Opc;
Expand Down Expand Up @@ -2968,8 +2979,15 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
// See the comments in emitPtrauthBranch.
if (Callee == AddrDisc)
report_fatal_error("Call target is signed with its own value");

// After isX16X17Safer predicate was introduced, emitPtrauthDiscriminator is
// no longer restricted to only reusing AddrDisc when it is X16 or X17
// (which are implicit-def'ed by AUTH_TCRETURN pseudos), thus impose this
// restriction manually not to clobber an unexpected register.
bool AddrDiscIsImplicitDef =
AddrDisc == AArch64::X16 || AddrDisc == AArch64::X17;
Register DiscReg = emitPtrauthDiscriminator(Disc, AddrDisc, ScratchReg,
/*MayUseAddrAsScratch=*/true);
AddrDiscIsImplicitDef);

const bool IsZero = DiscReg == AArch64::XZR;
const unsigned Opcodes[2][2] = {{AArch64::BRAA, AArch64::BRAAZ},
Expand Down
101 changes: 97 additions & 4 deletions llvm/test/CodeGen/AArch64/ptrauth-call.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,8 +181,9 @@ define void @test_tailcall_omit_mov_x16_x16(ptr %objptr) #0 {
; ELF-NEXT: movk x8, #6503, lsl #48
; ELF-NEXT: autda x1, x8
; ELF-NEXT: ldr x2, [x1]
; ELF-NEXT: movk x1, #54167, lsl #48
; ELF-NEXT: braa x2, x1
; ELF-NEXT: mov x16, x1
; ELF-NEXT: movk x16, #54167, lsl #48
; ELF-NEXT: braa x2, x16
%vtable.signed = load ptr, ptr %objptr, align 8
%objptr.int = ptrtoint ptr %objptr to i64
%vtable.discr = tail call i64 @llvm.ptrauth.blend(i64 %objptr.int, i64 6503)
Expand Down Expand Up @@ -213,8 +214,9 @@ define i32 @test_call_omit_extra_moves(ptr %objptr) #0 {
; ELF-NEXT: movk x9, #6503, lsl #48
; ELF-NEXT: autda x8, x9
; ELF-NEXT: ldr x9, [x8]
; ELF-NEXT: movk x8, #34646, lsl #48
; ELF-NEXT: blraa x9, x8
; ELF-NEXT: mov x17, x8
; ELF-NEXT: movk x17, #34646, lsl #48
; ELF-NEXT: blraa x9, x17
; ELF-NEXT: mov w0, #42
; ELF-NEXT: ldr x30, [sp], #16
; CHECK-NEXT: ret
Expand All @@ -230,6 +232,97 @@ define i32 @test_call_omit_extra_moves(ptr %objptr) #0 {
ret i32 42
}

; The second BLRA instruction should not reuse its AddrDisc operand as a scratch register (returned later).
define i64 @test_call_discr_csr_live(ptr %fnptr, i64 %addr.discr) #0 {
; ELF-LABEL: test_call_discr_csr_live:
; ELF-NEXT: str x30, [sp, #-32]!
; ELF-NEXT: stp x20, x19, [sp, #16]
; ELF-DAG: mov x[[FNPTR:[0-9]+]], x0
; ELF-DAG: mov x[[ADDR_DISC:[0-9]+]], x1
; ELF-DAG: mov x17, x1
; ELF-NEXT: movk x17, #6503, lsl #48
; ELF-NEXT: blraa x0, x17
; ELF-NEXT: mov x17, x[[ADDR_DISC]]
; ELF-NEXT: movk x17, #6503, lsl #48
; ELF-NEXT: blraa x[[FNPTR]], x17
; ELF-NEXT: mov x0, x[[ADDR_DISC]]
; ELF-NEXT: ldp x20, x19, [sp, #16]
; ELF-NEXT: ldr x30, [sp], #32
; ELF-NEXT: ret
%discr = tail call i64 @llvm.ptrauth.blend(i64 %addr.discr, i64 6503)
tail call void %fnptr() [ "ptrauth"(i32 0, i64 %discr) ]
tail call void %fnptr() [ "ptrauth"(i32 0, i64 %discr) ]
ret i64 %addr.discr
}

; The second BLRA instruction may reuse its AddrDisc operand as a scratch register.
define i64 @test_call_discr_csr_killed(ptr %fnptr, i64 %addr.discr) #0 {
; ELF-LABEL: test_call_discr_csr_killed:
; ELF-NEXT: str x30, [sp, #-32]!
; ELF-NEXT: stp x20, x19, [sp, #16]
; ELF-DAG: mov x[[FNPTR:[0-9]+]], x0
; ELF-DAG: mov x[[ADDR_DISC:[0-9]+]], x1
; ELF-DAG: mov x17, x1
; ELF-NEXT: movk x17, #6503, lsl #48
; ELF-NEXT: blraa x0, x17
; ELF-DAG: mov x17, x[[ADDR_DISC]]
; ELF-NEXT: movk x17, #6503, lsl #48
; ELF-NEXT: blraa x[[FNPTR]], x17
; ELF-NEXT: ldp x20, x19, [sp, #16]
; ELF-NEXT: mov w0, #42
; ELF-NEXT: ldr x30, [sp], #32
; ELF-NEXT: ret
%discr = tail call i64 @llvm.ptrauth.blend(i64 %addr.discr, i64 6503)
tail call void %fnptr() [ "ptrauth"(i32 0, i64 %discr) ]
tail call void %fnptr() [ "ptrauth"(i32 0, i64 %discr) ]
ret i64 42
}

; BLRA instruction should not reuse its AddrDisc operand as a scratch register (function argument).
define i64 @test_call_discr_arg(ptr %fnptr, i64 %addr.discr) #0 {
; ELF-LABEL: test_call_discr_arg:
; ELF-NEXT: str x30, [sp, #-16]!
; ELF-NEXT: mov x8, x0
; ELF-NEXT: mov x0, xzr
; ELF-NEXT: mov x17, x1
; ELF-NEXT: movk x17, #6503, lsl #48
; ELF-NEXT: blraa x8, x17
; ELF-NEXT: mov w0, #42
; ELF-NEXT: ldr x30, [sp], #16
; ELF-NEXT: ret
%discr = tail call i64 @llvm.ptrauth.blend(i64 %addr.discr, i64 6503)
tail call void %fnptr(ptr null, i64 %addr.discr) [ "ptrauth"(i32 0, i64 %discr) ]
ret i64 42
}

; BLRA instruction may reuse its AddrDisc operand as a scratch register.
define i64 @test_call_discr_non_arg(ptr %fnptr, i64 %addr.discr) #0 {
; ELF-LABEL: test_call_discr_non_arg:
; ELF-NEXT: str x30, [sp, #-16]!
; ELF-NEXT: mov x17, x1
; ELF-NEXT: movk x17, #6503, lsl #48
; ELF-NEXT: blraa x0, x17
; ELF-NEXT: mov w0, #42
; ELF-NEXT: ldr x30, [sp], #16
; ELF-NEXT: ret
%discr = tail call i64 @llvm.ptrauth.blend(i64 %addr.discr, i64 6503)
tail call void %fnptr() [ "ptrauth"(i32 0, i64 %discr) ]
ret i64 42
}

; AUTH_TCRETURN instruction should not reuse its AddrDisc operand as a scratch register (function argument).
define i64 @test_tailcall_discr_arg(ptr %fnptr, i64 %addr.discr) #0 {
; ELF-LABEL: test_tailcall_discr_arg:
; ELF-NEXT: mov x2, x0
; ELF-NEXT: mov x0, xzr
; ELF-NEXT: mov x16, x1
; ELF-NEXT: movk x16, #6503, lsl #48
; ELF-NEXT: braa x2, x16
%discr = tail call i64 @llvm.ptrauth.blend(i64 %addr.discr, i64 6503)
%result = tail call i64 %fnptr(ptr null, i64 %addr.discr) [ "ptrauth"(i32 0, i64 %discr) ]
ret i64 %result
}

define i32 @test_call_ia_arg(ptr %arg0, i64 %arg1) #0 {
; DARWIN-LABEL: test_call_ia_arg:
; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
Expand Down