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

atrosinenko · 2025-08-26T08:45:43Z

After isX16X17Safer predicate was introduced, it became possible for emitPtrauthDiscriminator function to return its AddrDisc argument when it is neither X16 nor X17 (which are declared as implicit-def'ed by BLRA and AUTH_TCRETURN[_BTI] pseudo instructions). This resulted in the above pseudos being able to accidentally clobber unexpected register.

As a quick fix for miscompilation possibility, this patch virtually restores the old behavior for the affected pseudo instructions. It should be possible to improve the efficiency via subsequent patches by accounting for killed flags and register masks.

… AddrDisc

After isX16X17Safer predicate was introduced, it became possible for emitPtrauthDiscriminator function to return its AddrDisc argument when it is neither X16 nor X17 (which are declared as implicit-def'ed by BLRA and AUTH_TCRETURN[_BTI] pseudo instructions). This resulted in the above pseudos being able to accidentally clobber unexpected register. As a quick fix for miscompilation possibility, this patch virtually restores the old behavior for the affected pseudo instructions. It should be possible to improve the efficiency via subsequent patches by accounting for `killed` flags and register masks.

atrosinenko · 2025-08-26T10:01:28Z

The issue seems to be caused by combination of #115185 and #132857.

llvmbot · 2025-08-26T10:04:11Z

@llvm/pr-subscribers-backend-aarch64

Author: Anatoly Trosinenko (atrosinenko)

Changes

After isX16X17Safer predicate was introduced, it became possible for emitPtrauthDiscriminator function to return its AddrDisc argument when it is neither X16 nor X17 (which are declared as implicit-def'ed by BLRA and AUTH_TCRETURN[_BTI] pseudo instructions). This resulted in the above pseudos being able to accidentally clobber unexpected register.

As a quick fix for miscompilation possibility, this patch virtually restores the old behavior for the affected pseudo instructions. It should be possible to improve the efficiency via subsequent patches by accounting for killed flags and register masks.

Full diff: https://github.com/llvm/llvm-project/pull/155373.diff

2 Files Affected:

(modified) llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp (+23-5)
(modified) llvm/test/CodeGen/AArch64/ptrauth-call.ll (+97-4)

diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index c52487ab8a79a..74eab99c2e2db 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -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;
@@ -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},
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-call.ll b/llvm/test/CodeGen/AArch64/ptrauth-call.ll
index 700686b9f194b..fc555a882be2c 100644
--- a/llvm/test/CodeGen/AArch64/ptrauth-call.ll
+++ b/llvm/test/CodeGen/AArch64/ptrauth-call.ll
@@ -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)
@@ -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
@@ -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]!

atrosinenko · 2025-08-28T16:06:08Z

/cherry-pick c7f3bdb

llvmbot · 2025-08-28T16:12:33Z

/pull-request #155877

…vm#155373) After `isX16X17Safer` predicate was introduced, it became possible for `emitPtrauthDiscriminator` function to return its `AddrDisc` argument when it is neither X16 nor X17 (which are declared as implicit-def'ed by `BLRA` and `AUTH_TCRETURN[_BTI]` pseudo instructions). This resulted in the above pseudos being able to accidentally clobber unexpected register. As a quick fix for miscompilation possibility, this patch virtually restores the old behavior for the affected pseudo instructions. It should be possible to improve the efficiency via subsequent patches by accounting for `killed` flags and register masks. (cherry picked from commit c7f3bdb)

atrosinenko added 2 commits August 25, 2025 23:58

[AArch64][PAC] Precommit test cases for BLRA/AUTH_TCRETURN clobbering…

ce72add

… AddrDisc

atrosinenko requested review from ahmedbougacha, asl, kbeyls, kovdan01, ojhunt and pcc August 26, 2025 10:02

atrosinenko marked this pull request as ready for review August 26, 2025 10:03

llvmbot added the backend:AArch64 label Aug 26, 2025

pcc approved these changes Aug 26, 2025

View reviewed changes

atrosinenko merged commit c7f3bdb into main Aug 27, 2025
13 checks passed

atrosinenko deleted the users/atrosinenko/auth-call-register-clobbering branch August 27, 2025 16:36

atrosinenko added this to the LLVM 21.x Release milestone Aug 28, 2025

github-project-automation bot added this to LLVM Release Status Aug 28, 2025

github-project-automation bot moved this to Needs Triage in LLVM Release Status Aug 28, 2025

llvmbot moved this from Needs Triage to Done in LLVM Release Status Aug 28, 2025

dwblaikie mentioned this pull request Sep 16, 2025

port 5b4819e to release #159209

Merged

jhuber6 mentioned this pull request Sep 17, 2025

[LLVM] Update CUDA ELF flags for their new ABI #159451

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

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

Uh oh!

atrosinenko commented Aug 26, 2025

Uh oh!

atrosinenko commented Aug 26, 2025

Uh oh!

llvmbot commented Aug 26, 2025

Uh oh!

Uh oh!

atrosinenko commented Aug 28, 2025

Uh oh!

llvmbot commented Aug 28, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

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

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

Uh oh!

Conversation

atrosinenko commented Aug 26, 2025

Uh oh!

atrosinenko commented Aug 26, 2025

Uh oh!

llvmbot commented Aug 26, 2025

Uh oh!

Uh oh!

atrosinenko commented Aug 28, 2025

Uh oh!

llvmbot commented Aug 28, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants