https://github.com/atrosinenko created https://github.com/llvm/llvm-project/pull/169699
Refactor `AArch64AsmPrinter::emitPtrauthAuthResign` to improve readability and fix the conditions of `emitPtrauthDiscriminator` being allowed to clobber AddrDisc: * do not clobber `AUTAddrDisc` when computing `AUTDiscReg` on resigning if `AUTAddrDisc == PACAddrDisc`, as it would prevent passing raw, 64-bit value as the new discriminator * mark the `$Scratch` operand of `AUTxMxN` as early-clobber (fixes assertions when emitting code at `-O0`) * move the code computing `ShouldCheck` and `ShouldTrap` conditions to a separate function * define helper `struct PtrAuthSchema` to pass arguments to `emitPtrauthAuthResign` in a better structured way >From 0371fa4afcee42f5d1d8affd32126dadea173bd2 Mon Sep 17 00:00:00 2001 From: Anatoly Trosinenko <[email protected]> Date: Thu, 25 Sep 2025 22:28:14 +0300 Subject: [PATCH] [AArch64][PAC] Rework the expansion of AUT/AUTPAC pseudos Refactor `AArch64AsmPrinter::emitPtrauthAuthResign` to improve readability and fix the conditions of `emitPtrauthDiscriminator` being allowed to clobber AddrDisc: * do not clobber `AUTAddrDisc` when computing `AUTDiscReg` on resigning if `AUTAddrDisc == PACAddrDisc`, as it would prevent passing raw, 64-bit value as the new discriminator * mark the `$Scratch` operand of `AUTxMxN` as early-clobber (fixes assertions when emitting code at `-O0`) * move the code computing `ShouldCheck` and `ShouldTrap` conditions to a separate function * define helper `struct PtrAuthSchema` to pass arguments to `emitPtrauthAuthResign` in a better structured way --- llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp | 177 +++++++++++------- llvm/lib/Target/AArch64/AArch64InstrInfo.td | 13 +- ...trauth-intrinsic-auth-resign-with-blend.ll | 77 +++++++- 3 files changed, 192 insertions(+), 75 deletions(-) diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp index 73c91c97bc1db..594a0bbff0028 100644 --- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp +++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp @@ -179,13 +179,21 @@ class AArch64AsmPrinter : public AsmPrinter { // Check authenticated LR before tail calling. void emitPtrauthTailCallHardening(const MachineInstr *TC); + struct PtrAuthSchema { + PtrAuthSchema(AArch64PACKey::ID Key, uint64_t Disc, + const MachineOperand &AddrDiscOp); + + AArch64PACKey::ID Key; + uint64_t Disc; + Register AddrDisc; + bool AddrDiscIsKilled; + }; + // Emit the sequence for AUT or AUTPAC. - void emitPtrauthAuthResign(Register AUTVal, AArch64PACKey::ID AUTKey, - uint64_t AUTDisc, - const MachineOperand *AUTAddrDisc, - Register Scratch, - std::optional<AArch64PACKey::ID> PACKey, - uint64_t PACDisc, Register PACAddrDisc); + void + emitPtrauthAuthResign(Register Pointer, Register Scratch, + PtrAuthSchema AuthSchema, + std::optional<PtrAuthSchema> SignSchema = std::nullopt); // Emit the sequence for PAC. void emitPtrauthSign(const MachineInstr *MI); @@ -2175,23 +2183,9 @@ void AArch64AsmPrinter::emitPtrauthTailCallHardening(const MachineInstr *TC) { LRCheckMethod); } -void AArch64AsmPrinter::emitPtrauthAuthResign( - Register AUTVal, AArch64PACKey::ID AUTKey, uint64_t AUTDisc, - const MachineOperand *AUTAddrDisc, Register Scratch, - std::optional<AArch64PACKey::ID> PACKey, uint64_t PACDisc, - Register PACAddrDisc) { - const bool IsAUTPAC = PACKey.has_value(); - - // We expand AUT/AUTPAC into a sequence of the form - // - // ; authenticate x16 - // ; check pointer in x16 - // Lsuccess: - // ; sign x16 (if AUTPAC) - // Lend: ; if not trapping on failure - // - // with the checking sequence chosen depending on whether/how we should check - // the pointer and whether we should trap on failure. +static std::pair<bool, bool> getCheckAndTrapMode(const MachineFunction *MF, + bool IsResign) { + const AArch64Subtarget &STI = MF->getSubtarget<AArch64Subtarget>(); // By default, auth/resign sequences check for auth failures. bool ShouldCheck = true; @@ -2200,7 +2194,7 @@ void AArch64AsmPrinter::emitPtrauthAuthResign( // On an FPAC CPU, you get traps whether you want them or not: there's // no point in emitting checks or traps. - if (STI->hasFPAC()) + if (STI.hasFPAC()) ShouldCheck = ShouldTrap = false; // However, command-line flags can override this, for experimentation. @@ -2219,38 +2213,79 @@ void AArch64AsmPrinter::emitPtrauthAuthResign( break; } - // Compute aut discriminator - Register AUTDiscReg = emitPtrauthDiscriminator( - AUTDisc, AUTAddrDisc->getReg(), Scratch, AUTAddrDisc->isKill()); - emitAUT(AUTKey, AUTVal, AUTDiscReg); + // Checked-but-not-trapping mode ("poison") only applies to resigning, + // replace with "unchecked" for standalone AUT. + if (!IsResign && ShouldCheck && !ShouldTrap) + ShouldCheck = ShouldTrap = false; - // Unchecked or checked-but-non-trapping AUT is just an "AUT": we're done. - if (!IsAUTPAC && (!ShouldCheck || !ShouldTrap)) - return; + return std::make_pair(ShouldCheck, ShouldTrap); +} - MCSymbol *EndSym = nullptr; +AArch64AsmPrinter::PtrAuthSchema::PtrAuthSchema( + AArch64PACKey::ID Key, uint64_t Disc, const MachineOperand &AddrDiscOp) + : Key(Key), Disc(Disc), AddrDisc(AddrDiscOp.getReg()), + AddrDiscIsKilled(AddrDiscOp.isKill()) {} - if (ShouldCheck) { - if (IsAUTPAC && !ShouldTrap) - EndSym = createTempSymbol("resign_end_"); +// We expand AUTx16x17/AUTxMxN into a sequence of the form +// +// ; authenticate Pointer +// ; check that Pointer is valid (optional, traps on failure) +// +// We expand AUTPAC into a sequence of the form +// +// ; authenticate Pointer +// ; check that Pointer is valid (optional, traps on failure) +// ; sign Pointer +// +// or +// +// ; authenticate Pointer +// ; check that Pointer is valid (skips re-sign on failure) +// ; sign Pointer +// Lon_failure: +// +void AArch64AsmPrinter::emitPtrauthAuthResign( + Register Pointer, Register Scratch, PtrAuthSchema AuthSchema, + std::optional<PtrAuthSchema> SignSchema) { + const bool IsResign = SignSchema.has_value(); + + const auto [ShouldCheck, ShouldTrap] = getCheckAndTrapMode(MF, IsResign); + const bool ShouldSkipSignOnAuthFailure = ShouldCheck && !ShouldTrap; + assert((ShouldCheck || !ShouldTrap) && "ShouldTrap implies ShouldCheck"); + + // It is hardly meaningful to authenticate or sign a pointer using its own + // value, thus we only have to take care not to early-clobber + // AuthSchema.AddrDisc that is aliased with SignSchema->AddrDisc. + assert(Pointer != AuthSchema.AddrDisc); + assert(!SignSchema || Pointer != SignSchema->AddrDisc); + bool IsResignWithAliasedAddrDiscs = + IsResign && AuthSchema.AddrDisc == SignSchema->AddrDisc; + bool MayReuseAUTAddrDisc = + !IsResignWithAliasedAddrDiscs && AuthSchema.AddrDiscIsKilled; + Register AUTDiscReg = emitPtrauthDiscriminator( + AuthSchema.Disc, AuthSchema.AddrDisc, Scratch, MayReuseAUTAddrDisc); + emitAUT(AuthSchema.Key, Pointer, AUTDiscReg); - emitPtrauthCheckAuthenticatedValue( - AUTVal, Scratch, AUTKey, AArch64PAuth::AuthCheckMethod::XPAC, EndSym); - } + MCSymbol *OnFailure = + ShouldSkipSignOnAuthFailure ? createTempSymbol("resign_end_") : nullptr; + + if (ShouldCheck) + emitPtrauthCheckAuthenticatedValue(Pointer, Scratch, AuthSchema.Key, + AArch64PAuth::AuthCheckMethod::XPAC, + OnFailure); - // We already emitted unchecked and checked-but-non-trapping AUTs. - // That left us with trapping AUTs, and AUTPACs. - // Trapping AUTs don't need PAC: we're done. - if (!IsAUTPAC) + if (!IsResign) { + assert(!OnFailure && "Poison mode only applies to resigning"); return; + } - // Compute pac discriminator - Register PACDiscReg = emitPtrauthDiscriminator(PACDisc, PACAddrDisc, Scratch); - emitPAC(*PACKey, AUTVal, PACDiscReg); + Register PACDiscReg = + emitPtrauthDiscriminator(SignSchema->Disc, SignSchema->AddrDisc, Scratch, + SignSchema->AddrDiscIsKilled); + emitPAC(SignSchema->Key, Pointer, PACDiscReg); - // Lend: - if (EndSym) - OutStreamer->emitLabel(EndSym); + if (OnFailure) + OutStreamer->emitLabel(OnFailure); } void AArch64AsmPrinter::emitPtrauthSign(const MachineInstr *MI) { @@ -2936,27 +2971,41 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) { return; } - case AArch64::AUTx16x17: - emitPtrauthAuthResign(AArch64::X16, - (AArch64PACKey::ID)MI->getOperand(0).getImm(), - MI->getOperand(1).getImm(), &MI->getOperand(2), - AArch64::X17, std::nullopt, 0, 0); + case AArch64::AUTx16x17: { + const Register Pointer = AArch64::X16; + const Register Scratch = AArch64::X17; + + PtrAuthSchema AuthSchema((AArch64PACKey::ID)MI->getOperand(0).getImm(), + MI->getOperand(1).getImm(), MI->getOperand(2)); + + emitPtrauthAuthResign(Pointer, Scratch, AuthSchema); return; + } + + case AArch64::AUTxMxN: { + const Register Pointer = MI->getOperand(0).getReg(); + const Register Scratch = MI->getOperand(1).getReg(); - case AArch64::AUTxMxN: - emitPtrauthAuthResign(MI->getOperand(0).getReg(), - (AArch64PACKey::ID)MI->getOperand(3).getImm(), - MI->getOperand(4).getImm(), &MI->getOperand(5), - MI->getOperand(1).getReg(), std::nullopt, 0, 0); + PtrAuthSchema AuthSchema((AArch64PACKey::ID)MI->getOperand(3).getImm(), + MI->getOperand(4).getImm(), MI->getOperand(5)); + + emitPtrauthAuthResign(Pointer, Scratch, AuthSchema); return; + } - case AArch64::AUTPAC: - emitPtrauthAuthResign( - AArch64::X16, (AArch64PACKey::ID)MI->getOperand(0).getImm(), - MI->getOperand(1).getImm(), &MI->getOperand(2), AArch64::X17, - (AArch64PACKey::ID)MI->getOperand(3).getImm(), - MI->getOperand(4).getImm(), MI->getOperand(5).getReg()); + case AArch64::AUTPAC: { + const Register Pointer = AArch64::X16; + const Register Scratch = AArch64::X17; + + PtrAuthSchema AuthSchema((AArch64PACKey::ID)MI->getOperand(0).getImm(), + MI->getOperand(1).getImm(), MI->getOperand(2)); + + PtrAuthSchema SignSchema((AArch64PACKey::ID)MI->getOperand(3).getImm(), + MI->getOperand(4).getImm(), MI->getOperand(5)); + + emitPtrauthAuthResign(Pointer, Scratch, AuthSchema, SignSchema); return; + } case AArch64::PAC: emitPtrauthSign(MI); diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td index 03bad8ff8ac8a..f0e3f9289a770 100644 --- a/llvm/lib/Target/AArch64/AArch64InstrInfo.td +++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.td @@ -2188,10 +2188,21 @@ let Predicates = [HasPAuth] in { let Uses = [X16]; } + // AsmPrinter can clobber the $AddrDisc register as long as it is marked + // "killed", otherwise an assertion checks that $Scratch != $AddrDisc. + // The problem is that it is always correct to *omit* the "killed" flag, + // thus an instruction like this is can be passed to AsmPrinter: + // + // $x8, $x9 = AUTxMxN $x8, 0, 51756, $x9, implicit-def $nzcv + // + // While it is possible to check if "$Scratch == $AddrDisc OR $AddrDisc + // is killed" instead, it is easier and more straightforward to mark + // $Scratch as @earlyclobber. Furthermore, this would be mandatory for a + // (hypotetical at the time of writing this comment) AUTPACxMxN pseudo. def AUTxMxN : Pseudo<(outs GPR64:$AuthVal, GPR64common:$Scratch), (ins GPR64:$Val, i32imm:$Key, i64imm:$Disc, GPR64:$AddrDisc), - [], "$AuthVal = $Val">, Sched<[WriteI, ReadI]> { + [], "$AuthVal = $Val,@earlyclobber $Scratch">, Sched<[WriteI, ReadI]> { let isCodeGenOnly = 1; let hasSideEffects = 1; let mayStore = 0; diff --git a/llvm/test/CodeGen/AArch64/ptrauth-intrinsic-auth-resign-with-blend.ll b/llvm/test/CodeGen/AArch64/ptrauth-intrinsic-auth-resign-with-blend.ll index e2aea6df78250..ae5fa509d439b 100644 --- a/llvm/test/CodeGen/AArch64/ptrauth-intrinsic-auth-resign-with-blend.ll +++ b/llvm/test/CodeGen/AArch64/ptrauth-intrinsic-auth-resign-with-blend.ll @@ -1,4 +1,3 @@ -; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py ; RUN: llc < %s -mtriple arm64e-apple-darwin -global-isel=0 -verify-machineinstrs \ ; RUN: -aarch64-ptrauth-auth-checks=none | FileCheck %s -DL="L" --check-prefixes=UNCHECKED,UNCHECKED-DARWIN ; RUN: llc < %s -mtriple arm64e-apple-darwin -global-isel -global-isel-abort=1 -verify-machineinstrs \ @@ -84,12 +83,14 @@ define i64 @test_resign_blend(i64 %arg, i64 %arg1, i64 %arg2) { ; UNCHECKED-NEXT: mov x16, x0 ; UNCHECKED-ELF-NEXT: movk x1, #12345, lsl #48 ; UNCHECKED-ELF-NEXT: autda x16, x1 +; UNCHECKED-ELF-NEXT: movk x2, #56789, lsl #48 +; UNCHECKED-ELF-NEXT: pacdb x16, x2 ; UNCHECKED-DARWIN-NEXT: mov x17, x1 ; UNCHECKED-DARWIN-NEXT: movk x17, #12345, lsl #48 ; UNCHECKED-DARWIN-NEXT: autda x16, x17 -; UNCHECKED-NEXT: mov x17, x2 -; UNCHECKED-NEXT: movk x17, #56789, lsl #48 -; UNCHECKED-NEXT: pacdb x16, x17 +; UNCHECKED-DARWIN-NEXT: mov x17, x2 +; UNCHECKED-DARWIN-NEXT: movk x17, #56789, lsl #48 +; UNCHECKED-DARWIN-NEXT: pacdb x16, x17 ; UNCHECKED-NEXT: mov x0, x16 ; UNCHECKED-NEXT: ret ; @@ -108,9 +109,11 @@ define i64 @test_resign_blend(i64 %arg, i64 %arg1, i64 %arg2) { ; CHECKED-NEXT: mov x16, x17 ; CHECKED-NEXT: b [[L]]resign_end_0 ; CHECKED-NEXT: Lauth_success_0: -; CHECKED-NEXT: mov x17, x2 -; CHECKED-NEXT: movk x17, #56789, lsl #48 -; CHECKED-NEXT: pacdb x16, x17 +; CHECKED-ELF-NEXT: movk x2, #56789, lsl #48 +; CHECKED-ELF-NEXT: pacdb x16, x2 +; CHECKED-DARWIN-NEXT: mov x17, x2 +; CHECKED-DARWIN-NEXT: movk x17, #56789, lsl #48 +; CHECKED-DARWIN-NEXT: pacdb x16, x17 ; CHECKED-NEXT: Lresign_end_0: ; CHECKED-NEXT: mov x0, x16 ; CHECKED-NEXT: ret @@ -129,9 +132,11 @@ define i64 @test_resign_blend(i64 %arg, i64 %arg1, i64 %arg2) { ; TRAP-NEXT: b.eq [[L]]auth_success_1 ; TRAP-NEXT: brk #0xc472 ; TRAP-NEXT: Lauth_success_1: -; TRAP-NEXT: mov x17, x2 -; TRAP-NEXT: movk x17, #56789, lsl #48 -; TRAP-NEXT: pacdb x16, x17 +; TRAP-ELF-NEXT: movk x2, #56789, lsl #48 +; TRAP-ELF-NEXT: pacdb x16, x2 +; TRAP-DARWIN-NEXT: mov x17, x2 +; TRAP-DARWIN-NEXT: movk x17, #56789, lsl #48 +; TRAP-DARWIN-NEXT: pacdb x16, x17 ; TRAP-NEXT: mov x0, x16 ; TRAP-NEXT: ret %tmp0 = call i64 @llvm.ptrauth.blend(i64 %arg1, i64 12345) @@ -299,6 +304,58 @@ define i64 @test_auth_too_large_discriminator(i64 %arg, i64 %arg1) { ret i64 %tmp1 } +; As long as we support raw, non-blended 64-bit discriminators (which might be +; useful for low-level code such as dynamic loaders), the "auth" part of resign +; must not clobber %arg, if its upper bits are later used by the "sign" part. +define i64 @test_resign_aliased_discs_raw_sign_disc(i64 %p, i64 %arg) { +; UNCHECKED-LABEL: test_resign_aliased_discs_raw_sign_disc: +; UNCHECKED: %bb.0: +; UNCHECKED-NEXT: mov x16, x0 +; UNCHECKED-NEXT: mov x17, x1 +; UNCHECKED-NEXT: movk x17, #12345, lsl #48 +; UNCHECKED-NEXT: autda x16, x17 +; UNCHECKED-NEXT: pacdb x16, x1 +; UNCHECKED-NEXT: mov x0, x16 +; UNCHECKED-NEXT: ret +; +; CHECKED-LABEL: test_resign_aliased_discs_raw_sign_disc: +; CHECKED: %bb.0: +; CHECKED-NEXT: mov x16, x0 +; CHECKED-NEXT: mov x17, x1 +; CHECKED-NEXT: movk x17, #12345, lsl #48 +; CHECKED-NEXT: autda x16, x17 +; CHECKED-NEXT: mov x17, x16 +; CHECKED-NEXT: xpacd x17 +; CHECKED-NEXT: cmp x16, x17 +; CHECKED-NEXT: b.eq [[L]]auth_success_3 +; CHECKED-NEXT: mov x16, x17 +; CHECKED-NEXT: b [[L]]resign_end_3 +; CHECKED-NEXT: Lauth_success_3: +; CHECKED-NEXT: pacdb x16, x1 +; CHECKED-NEXT: Lresign_end_3: +; CHECKED-NEXT: mov x0, x16 +; CHECKED-NEXT: ret +; +; TRAP-LABEL: test_resign_aliased_discs_raw_sign_disc: +; TRAP: %bb.0: +; TRAP-NEXT: mov x16, x0 +; TRAP-NEXT: mov x17, x1 +; TRAP-NEXT: movk x17, #12345, lsl #48 +; TRAP-NEXT: autda x16, x17 +; TRAP-NEXT: mov x17, x16 +; TRAP-NEXT: xpacd x17 +; TRAP-NEXT: cmp x16, x17 +; TRAP-NEXT: b.eq [[L]]auth_success_5 +; TRAP-NEXT: brk #0xc472 +; TRAP-NEXT: Lauth_success_5: +; TRAP-NEXT: pacdb x16, x1 +; TRAP-NEXT: mov x0, x16 +; TRAP-NEXT: ret + %auth.disc = call i64 @llvm.ptrauth.blend(i64 %arg, i64 12345) + %res = call i64 @llvm.ptrauth.resign(i64 %p, i32 2, i64 %auth.disc, i32 3, i64 %arg) + ret i64 %res +} + declare i64 @llvm.ptrauth.auth(i64, i32, i64) declare i64 @llvm.ptrauth.resign(i64, i32, i64, i32, i64) declare i64 @llvm.ptrauth.blend(i64, i64) _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
