[AArch64][PAC] Rework the expansion of AUT/AUTPAC pseudos #169699
+192
−75
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This was referenced Nov 26, 2025
Contributor
Author
|
Warning This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
This stack of pull requests is managed by Graphite. Learn more about stacking. |
atrosinenko
requested review from
ahmedbougacha,
asl,
kbeyls,
kovdan01,
pcc and
stuij
Member
|
@llvm/pr-subscribers-backend-aarch64 Author: Anatoly Trosinenko (atrosinenko) ChangesRefactor
Full diff: https://github.com/llvm/llvm-project/pull/169699.diff 3 Files Affected:
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)
|
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
atrosinenko
force-pushed
the
users/atrosinenko/pauth-refactor-autpac
branch
from
0371fa4 to
19d2b2c
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Refactor
AArch64AsmPrinter::emitPtrauthAuthResignto improvereadability and fix the conditions of
emitPtrauthDiscriminatorbeingallowed to clobber AddrDisc:
AUTAddrDiscwhen computingAUTDiscRegon resigningif
AUTAddrDisc == PACAddrDisc, as it would prevent passing raw,64-bit value as the new discriminator
$Scratchoperand ofAUTxMxNas early-clobber (fixesassertions when emitting code at
-O0)ShouldCheckandShouldTrapconditions to aseparate function
struct PtrAuthSchemato pass arguments toemitPtrauthAuthResignin a better structured way