[llvm] [AArch64][PAC] Refactor discriminator emission in AsmPrinter (NFC) (PR #160900)
Anatoly Trosinenko via llvm-commits
llvm-commits at lists.llvm.org
Thu Dec 4 09:52:36 PST 2025
https://github.com/atrosinenko updated https://github.com/llvm/llvm-project/pull/160900
>From 1a731412a3bdc9ab3ab06cc921a16d35cc6169d9 Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko <atrosinenko at accesssoftek.com>
Date: Wed, 24 Sep 2025 19:46:13 +0300
Subject: [PATCH 1/2] [AArch64][PAC] Cleanup
AArch64AsmPrinter::emitPtrauthDiscriminator (NFC)
Refactor emitPtrauthDiscriminator function: introduce `isPtrauthRegSafe`
function, update the comments and assertions for readability.
---
llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp | 46 ++++++++++---------
1 file changed, 25 insertions(+), 21 deletions(-)
diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index 8267414e78955..7fe66592babf8 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -158,6 +158,13 @@ class AArch64AsmPrinter : public AsmPrinter {
void emitSled(const MachineInstr &MI, SledKind Kind);
+ bool isPtrauthRegSafe(Register Reg) const {
+ if (STI->isX16X17Safer())
+ return Reg == AArch64::X16 || Reg == AArch64::X17;
+
+ return true;
+ }
+
// Emit the sequence for BRA/BLRA (authenticate + branch/call).
void emitPtrauthBranch(const MachineInstr *MI);
@@ -188,13 +195,15 @@ class AArch64AsmPrinter : public AsmPrinter {
// Emit the sequence to compute the discriminator.
//
- // The returned register is either unmodified AddrDisc or ScratchReg.
+ // The Scratch register passed to this function must be safe, see
+ // isPtrauthRegSafe(Reg) function.
+ //
+ // The returned register is either ScratchReg or AddrDisc. Furthermore, it
+ // is safe, unless unsafe AddrDisc was passed-through unmodified.
//
// If the expanded pseudo is allowed to clobber AddrDisc register, setting
- // MayUseAddrAsScratch may save one MOV instruction, provided the address
- // is already in x16/x17 (i.e. return x16/x17 which is the *modified* AddrDisc
- // register at the same time) or the OS doesn't make it safer to use x16/x17
- // (see AArch64Subtarget::isX16X17Safer()):
+ // MayUseAddrAsScratch may save one MOV instruction, provided
+ // isPtrauthRegSafe(AddrDisc) is true:
//
// mov x17, x16
// movk x17, #1234, lsl #48
@@ -203,7 +212,7 @@ class AArch64AsmPrinter : public AsmPrinter {
// can be replaced by
//
// movk x16, #1234, lsl #48
- Register emitPtrauthDiscriminator(uint16_t Disc, Register AddrDisc,
+ Register emitPtrauthDiscriminator(uint64_t Disc, Register AddrDisc,
Register ScratchReg,
bool MayUseAddrAsScratch = false);
@@ -1912,12 +1921,14 @@ void AArch64AsmPrinter::emitFMov0AsFMov(const MachineInstr &MI,
EmitToStreamer(*OutStreamer, FMov);
}
-Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
+Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint64_t Disc,
Register AddrDisc,
Register ScratchReg,
bool MayUseAddrAsScratch) {
- assert(ScratchReg == AArch64::X16 || ScratchReg == AArch64::X17 ||
- !STI->isX16X17Safer());
+ assert(isPtrauthRegSafe(ScratchReg) &&
+ "Safe scratch register must be provided by the caller");
+ assert(isUInt<16>(Disc) && "Constant discriminator is too wide");
+
// So far we've used NoRegister in pseudos. Now we need real encodings.
if (AddrDisc == AArch64::NoRegister)
AddrDisc = AArch64::XZR;
@@ -1936,13 +1947,13 @@ Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
// If there are both, emit a blend into the scratch register.
// Check if we can save one MOV instruction.
- assert(MayUseAddrAsScratch || ScratchReg != AddrDisc);
- bool AddrDiscIsSafe = AddrDisc == AArch64::X16 || AddrDisc == AArch64::X17 ||
- !STI->isX16X17Safer();
- if (MayUseAddrAsScratch && AddrDiscIsSafe)
+ if (MayUseAddrAsScratch && isPtrauthRegSafe(AddrDisc)) {
ScratchReg = AddrDisc;
- else
+ } else {
emitMovXReg(ScratchReg, AddrDisc);
+ assert(ScratchReg != AddrDisc &&
+ "Forbidden to clobber AddrDisc, but have to");
+ }
emitMOVK(ScratchReg, Disc, 48);
return ScratchReg;
@@ -2181,7 +2192,6 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(
}
// Compute aut discriminator
- assert(isUInt<16>(AUTDisc));
Register AUTDiscReg = emitPtrauthDiscriminator(
AUTDisc, AUTAddrDisc->getReg(), Scratch, AUTAddrDisc->isKill());
bool AUTZero = AUTDiscReg == AArch64::XZR;
@@ -2220,7 +2230,6 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(
return;
// Compute pac discriminator
- assert(isUInt<16>(PACDisc));
Register PACDiscReg =
emitPtrauthDiscriminator(PACDisc, PACAddrDisc, Scratch);
bool PACZero = PACDiscReg == AArch64::XZR;
@@ -2255,7 +2264,6 @@ void AArch64AsmPrinter::emitPtrauthSign(const MachineInstr *MI) {
"Neither X16 nor X17 is available as a scratch register");
// Compute pac discriminator
- assert(isUInt<16>(Disc));
Register DiscReg = emitPtrauthDiscriminator(
Disc, AddrDisc, ScratchReg, /*MayUseAddrAsScratch=*/AddrDiscKilled);
bool IsZeroDisc = DiscReg == AArch64::XZR;
@@ -2284,7 +2292,6 @@ void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
"Invalid auth call key");
uint64_t Disc = MI->getOperand(2).getImm();
- assert(isUInt<16>(Disc));
unsigned AddrDisc = MI->getOperand(3).getReg();
@@ -2679,8 +2686,6 @@ void AArch64AsmPrinter::LowerMOVaddrPAC(const MachineInstr &MI) {
const auto Key = (AArch64PACKey::ID)KeyC;
const unsigned AddrDisc = MI.getOperand(2).getReg();
const uint64_t Disc = MI.getOperand(3).getImm();
- assert(isUInt<16>(Disc) &&
- "constant discriminator is out of range [0, 0xffff]");
const int64_t Offset = GAOp.getOffset();
GAOp.setOffset(0);
@@ -3263,7 +3268,6 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
"Invalid auth key for tail-call return");
const uint64_t Disc = MI->getOperand(3).getImm();
- assert(isUInt<16>(Disc) && "Integer discriminator is too wide");
Register AddrDisc = MI->getOperand(4).getReg();
>From b40807a3a8e92270db7eed3b8feeaf76cc83e022 Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko <atrosinenko at accesssoftek.com>
Date: Tue, 25 Nov 2025 21:41:20 +0300
Subject: [PATCH 2/2] Update the comments, rename MayUseAddrAsScratch
---
llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp | 24 ++++++++++++-------
1 file changed, 15 insertions(+), 9 deletions(-)
diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index 7fe66592babf8..7f1d331982bc2 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -158,6 +158,11 @@ class AArch64AsmPrinter : public AsmPrinter {
void emitSled(const MachineInstr &MI, SledKind Kind);
+ // Returns whether Reg may be used to store sensitive temporary values when
+ // expanding PtrAuth pseudos. Some OSes may take extra care to protect a
+ // small subset of GPRs on context switches - use these registers then.
+ //
+ // If there are no preferred registers, returns true for any Reg.
bool isPtrauthRegSafe(Register Reg) const {
if (STI->isX16X17Safer())
return Reg == AArch64::X16 || Reg == AArch64::X17;
@@ -195,14 +200,15 @@ class AArch64AsmPrinter : public AsmPrinter {
// Emit the sequence to compute the discriminator.
//
- // The Scratch register passed to this function must be safe, see
- // isPtrauthRegSafe(Reg) function.
+ // The Scratch register passed to this function must be safe, as returned by
+ // isPtrauthRegSafe(ScratchReg).
//
- // The returned register is either ScratchReg or AddrDisc. Furthermore, it
- // is safe, unless unsafe AddrDisc was passed-through unmodified.
+ // The returned register is either ScratchReg, AddrDisc, or XZR. Furthermore,
+ // it is guaranteed to be safe (or XZR), with the only exception of
+ // passing-through an *unmodified* unsafe AddrDisc register.
//
// If the expanded pseudo is allowed to clobber AddrDisc register, setting
- // MayUseAddrAsScratch may save one MOV instruction, provided
+ // MayClobberAddrDisc may save one MOV instruction, provided
// isPtrauthRegSafe(AddrDisc) is true:
//
// mov x17, x16
@@ -214,7 +220,7 @@ class AArch64AsmPrinter : public AsmPrinter {
// movk x16, #1234, lsl #48
Register emitPtrauthDiscriminator(uint64_t Disc, Register AddrDisc,
Register ScratchReg,
- bool MayUseAddrAsScratch = false);
+ bool MayClobberAddrDisc = false);
// Emit the sequence for LOADauthptrstatic
void LowerLOADauthptrstatic(const MachineInstr &MI);
@@ -1924,7 +1930,7 @@ void AArch64AsmPrinter::emitFMov0AsFMov(const MachineInstr &MI,
Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint64_t Disc,
Register AddrDisc,
Register ScratchReg,
- bool MayUseAddrAsScratch) {
+ bool MayClobberAddrDisc) {
assert(isPtrauthRegSafe(ScratchReg) &&
"Safe scratch register must be provided by the caller");
assert(isUInt<16>(Disc) && "Constant discriminator is too wide");
@@ -1947,7 +1953,7 @@ Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint64_t Disc,
// If there are both, emit a blend into the scratch register.
// Check if we can save one MOV instruction.
- if (MayUseAddrAsScratch && isPtrauthRegSafe(AddrDisc)) {
+ if (MayClobberAddrDisc && isPtrauthRegSafe(AddrDisc)) {
ScratchReg = AddrDisc;
} else {
emitMovXReg(ScratchReg, AddrDisc);
@@ -2265,7 +2271,7 @@ void AArch64AsmPrinter::emitPtrauthSign(const MachineInstr *MI) {
// Compute pac discriminator
Register DiscReg = emitPtrauthDiscriminator(
- Disc, AddrDisc, ScratchReg, /*MayUseAddrAsScratch=*/AddrDiscKilled);
+ Disc, AddrDisc, ScratchReg, /*MayClobberAddrDisc=*/AddrDiscKilled);
bool IsZeroDisc = DiscReg == AArch64::XZR;
unsigned Opc = getPACOpcodeForKey(Key, IsZeroDisc);
More information about the llvm-commits
mailing list