Skip to content

Commit 72b5ad1

Browse files
atrosinenkoatrosinenko
committed
[AArch64][PAC] Factor out printing real AUT/PAC/BLRA encodings (NFC)
Separate the low-level emission of the appropriate variants of `AUT*`, `PAC*` and `B(L)RA*` instructions from the high-level logic of pseudo instruction expansion. Introduce `getBranchOpcodeForKey` helper function by analogy to `get(AUT|PAC)OpcodeForKey`.
1 parent d896965 commit 72b5ad1

File tree

2 files changed

+88
-105
lines changed

2 files changed

+88
-105
lines changed

‎llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp‎

Lines changed: 70 additions & 105 deletions
Original file line numberDiff line numberDiff line change
@@ -313,6 +313,11 @@ class AArch64AsmPrinter : public AsmPrinter {
313313
void emitMOVZ(Register Dest, uint64_t Imm, unsigned Shift);
314314
void emitMOVK(Register Dest, uint64_t Imm, unsigned Shift);
315315

316+
void emitAUT(AArch64PACKey::ID Key, Register Pointer, Register Disc);
317+
void emitPAC(AArch64PACKey::ID Key, Register Pointer, Register Disc);
318+
void emitBLRA(bool IsCall, AArch64PACKey::ID Key, Register Target,
319+
Register Disc);
320+
316321
/// Emit instruction to set float register to zero.
317322
void emitFMov0(const MachineInstr &MI);
318323
void emitFMov0AsFMov(const MachineInstr &MI, Register DestReg);
@@ -1834,6 +1839,55 @@ void AArch64AsmPrinter::emitMOVK(Register Dest, uint64_t Imm, unsigned Shift) {
18341839
.addImm(Shift));
18351840
}
18361841

1842+
void AArch64AsmPrinter::emitAUT(AArch64PACKey::ID Key, Register Pointer,
1843+
Register Disc) {
1844+
bool IsZeroDisc = Disc == AArch64::XZR;
1845+
unsigned Opcode = getAUTOpcodeForKey(Key, IsZeroDisc);
1846+
1847+
// autiza x16 ; if IsZeroDisc
1848+
// autia x16, x17 ; if !IsZeroDisc
1849+
MCInst AUTInst;
1850+
AUTInst.setOpcode(Opcode);
1851+
AUTInst.addOperand(MCOperand::createReg(Pointer));
1852+
AUTInst.addOperand(MCOperand::createReg(Pointer));
1853+
if (!IsZeroDisc)
1854+
AUTInst.addOperand(MCOperand::createReg(Disc));
1855+
1856+
EmitToStreamer(AUTInst);
1857+
}
1858+
1859+
void AArch64AsmPrinter::emitPAC(AArch64PACKey::ID Key, Register Pointer,
1860+
Register Disc) {
1861+
bool IsZeroDisc = Disc == AArch64::XZR;
1862+
unsigned Opcode = getPACOpcodeForKey(Key, IsZeroDisc);
1863+
1864+
// paciza x16 ; if IsZeroDisc
1865+
// pacia x16, x17 ; if !IsZeroDisc
1866+
MCInst PACInst;
1867+
PACInst.setOpcode(Opcode);
1868+
PACInst.addOperand(MCOperand::createReg(Pointer));
1869+
PACInst.addOperand(MCOperand::createReg(Pointer));
1870+
if (!IsZeroDisc)
1871+
PACInst.addOperand(MCOperand::createReg(Disc));
1872+
1873+
EmitToStreamer(PACInst);
1874+
}
1875+
1876+
void AArch64AsmPrinter::emitBLRA(bool IsCall, AArch64PACKey::ID Key,
1877+
Register Target, Register Disc) {
1878+
bool IsZeroDisc = Disc == AArch64::XZR;
1879+
unsigned Opcode = getBranchOpcodeForKey(IsCall, Key, IsZeroDisc);
1880+
1881+
// blraaz x16 ; if IsZeroDisc
1882+
// blraa x16, x17 ; if !IsZeroDisc
1883+
MCInst Inst;
1884+
Inst.setOpcode(Opcode);
1885+
Inst.addOperand(MCOperand::createReg(Target));
1886+
if (!IsZeroDisc)
1887+
Inst.addOperand(MCOperand::createReg(Disc));
1888+
EmitToStreamer(Inst);
1889+
}
1890+
18371891
void AArch64AsmPrinter::emitFMov0(const MachineInstr &MI) {
18381892
Register DestReg = MI.getOperand(0).getReg();
18391893
if (!STI->hasZeroCycleZeroingFPWorkaround() && STI->isNeonAvailable()) {
@@ -2162,18 +2216,7 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(
21622216
// Compute aut discriminator
21632217
Register AUTDiscReg = emitPtrauthDiscriminator(
21642218
AUTDisc, AUTAddrDisc->getReg(), Scratch, AUTAddrDisc->isKill());
2165-
bool AUTZero = AUTDiscReg == AArch64::XZR;
2166-
unsigned AUTOpc = getAUTOpcodeForKey(AUTKey, AUTZero);
2167-
2168-
// autiza x16 ; if AUTZero
2169-
// autia x16, x17 ; if !AUTZero
2170-
MCInst AUTInst;
2171-
AUTInst.setOpcode(AUTOpc);
2172-
AUTInst.addOperand(MCOperand::createReg(AUTVal));
2173-
AUTInst.addOperand(MCOperand::createReg(AUTVal));
2174-
if (!AUTZero)
2175-
AUTInst.addOperand(MCOperand::createReg(AUTDiscReg));
2176-
EmitToStreamer(*OutStreamer, AUTInst);
2219+
emitAUT(AUTKey, AUTVal, AUTDiscReg);
21772220

21782221
// Unchecked or checked-but-non-trapping AUT is just an "AUT": we're done.
21792222
if (!IsAUTPAC && (!ShouldCheck || !ShouldTrap))
@@ -2196,20 +2239,8 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(
21962239
return;
21972240

21982241
// Compute pac discriminator
2199-
Register PACDiscReg =
2200-
emitPtrauthDiscriminator(PACDisc, PACAddrDisc, Scratch);
2201-
bool PACZero = PACDiscReg == AArch64::XZR;
2202-
unsigned PACOpc = getPACOpcodeForKey(*PACKey, PACZero);
2203-
2204-
// pacizb x16 ; if PACZero
2205-
// pacib x16, x17 ; if !PACZero
2206-
MCInst PACInst;
2207-
PACInst.setOpcode(PACOpc);
2208-
PACInst.addOperand(MCOperand::createReg(AUTVal));
2209-
PACInst.addOperand(MCOperand::createReg(AUTVal));
2210-
if (!PACZero)
2211-
PACInst.addOperand(MCOperand::createReg(PACDiscReg));
2212-
EmitToStreamer(*OutStreamer, PACInst);
2242+
Register PACDiscReg = emitPtrauthDiscriminator(PACDisc, PACAddrDisc, Scratch);
2243+
emitPAC(*PACKey, AUTVal, PACDiscReg);
22132244

22142245
// Lend:
22152246
if (EndSym)
@@ -2232,28 +2263,14 @@ void AArch64AsmPrinter::emitPtrauthSign(const MachineInstr *MI) {
22322263
// Compute pac discriminator
22332264
Register DiscReg = emitPtrauthDiscriminator(
22342265
Disc, AddrDisc, ScratchReg, /*MayUseAddrAsScratch=*/AddrDiscKilled);
2235-
bool IsZeroDisc = DiscReg == AArch64::XZR;
2236-
unsigned Opc = getPACOpcodeForKey(Key, IsZeroDisc);
2237-
2238-
// paciza x16 ; if IsZeroDisc
2239-
// pacia x16, x17 ; if !IsZeroDisc
2240-
MCInst PACInst;
2241-
PACInst.setOpcode(Opc);
2242-
PACInst.addOperand(MCOperand::createReg(Val));
2243-
PACInst.addOperand(MCOperand::createReg(Val));
2244-
if (!IsZeroDisc)
2245-
PACInst.addOperand(MCOperand::createReg(DiscReg));
2246-
EmitToStreamer(*OutStreamer, PACInst);
2266+
emitPAC(Key, Val, DiscReg);
22472267
}
22482268

22492269
void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
22502270
bool IsCall = MI->getOpcode() == AArch64::BLRA;
22512271
unsigned BrTarget = MI->getOperand(0).getReg();
22522272

22532273
auto Key = (AArch64PACKey::ID)MI->getOperand(1).getImm();
2254-
assert((Key == AArch64PACKey::IA || Key == AArch64PACKey::IB) &&
2255-
"Invalid auth call key");
2256-
22572274
uint64_t Disc = MI->getOperand(2).getImm();
22582275

22592276
unsigned AddrDisc = MI->getOperand(3).getReg();
@@ -2283,27 +2300,7 @@ void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
22832300
IsCall && (AddrDisc == AArch64::X16 || AddrDisc == AArch64::X17);
22842301
Register DiscReg = emitPtrauthDiscriminator(Disc, AddrDisc, AArch64::X17,
22852302
AddrDiscIsImplicitDef);
2286-
bool IsZeroDisc = DiscReg == AArch64::XZR;
2287-
2288-
unsigned Opc;
2289-
if (IsCall) {
2290-
if (Key == AArch64PACKey::IA)
2291-
Opc = IsZeroDisc ? AArch64::BLRAAZ : AArch64::BLRAA;
2292-
else
2293-
Opc = IsZeroDisc ? AArch64::BLRABZ : AArch64::BLRAB;
2294-
} else {
2295-
if (Key == AArch64PACKey::IA)
2296-
Opc = IsZeroDisc ? AArch64::BRAAZ : AArch64::BRAA;
2297-
else
2298-
Opc = IsZeroDisc ? AArch64::BRABZ : AArch64::BRAB;
2299-
}
2300-
2301-
MCInst BRInst;
2302-
BRInst.setOpcode(Opc);
2303-
BRInst.addOperand(MCOperand::createReg(BrTarget));
2304-
if (!IsZeroDisc)
2305-
BRInst.addOperand(MCOperand::createReg(DiscReg));
2306-
EmitToStreamer(*OutStreamer, BRInst);
2303+
emitBLRA(IsCall, Key, BrTarget, DiscReg);
23072304
}
23082305

23092306
const MCExpr *
@@ -2506,22 +2503,14 @@ void AArch64AsmPrinter::LowerMOVaddrPAC(const MachineInstr &MI) {
25062503

25072504
assert(GAOp.isGlobal());
25082505
assert(GAOp.getGlobal()->getValueType() != nullptr);
2509-
unsigned AuthOpcode = GAOp.getGlobal()->getValueType()->isFunctionTy()
2510-
? AArch64::AUTIA
2511-
: AArch64::AUTDA;
2512-
2513-
EmitToStreamer(MCInstBuilder(AuthOpcode)
2514-
.addReg(AArch64::X16)
2515-
.addReg(AArch64::X16)
2516-
.addReg(AArch64::X17));
25172506

2518-
if (!STI->hasFPAC()) {
2519-
auto AuthKey = (AuthOpcode == AArch64::AUTIA ? AArch64PACKey::IA
2520-
: AArch64PACKey::DA);
2507+
bool IsFunctionTy = GAOp.getGlobal()->getValueType()->isFunctionTy();
2508+
auto AuthKey = IsFunctionTy ? AArch64PACKey::IA : AArch64PACKey::DA;
2509+
emitAUT(AuthKey, AArch64::X16, AArch64::X17);
25212510

2511+
if (!STI->hasFPAC())
25222512
emitPtrauthCheckAuthenticatedValue(AArch64::X16, AArch64::X17, AuthKey,
25232513
AArch64PAuth::AuthCheckMethod::XPAC);
2524-
}
25252514
} else {
25262515
EmitToStreamer(MCInstBuilder(AArch64::LDRXui)
25272516
.addReg(AArch64::X16)
@@ -2578,12 +2567,7 @@ void AArch64AsmPrinter::LowerMOVaddrPAC(const MachineInstr &MI) {
25782567

25792568
Register DiscReg = emitPtrauthDiscriminator(Disc, AddrDisc, AArch64::X17);
25802569

2581-
auto MIB = MCInstBuilder(getPACOpcodeForKey(Key, DiscReg == AArch64::XZR))
2582-
.addReg(AArch64::X16)
2583-
.addReg(AArch64::X16);
2584-
if (DiscReg != AArch64::XZR)
2585-
MIB.addReg(DiscReg);
2586-
EmitToStreamer(MIB);
2570+
emitPAC(Key, AArch64::X16, DiscReg);
25872571
}
25882572

25892573
void AArch64AsmPrinter::LowerLOADgotAUTH(const MachineInstr &MI) {
@@ -2637,21 +2621,15 @@ void AArch64AsmPrinter::LowerLOADgotAUTH(const MachineInstr &MI) {
26372621
}
26382622

26392623
assert(GAMO.getGlobal()->getValueType() != nullptr);
2640-
unsigned AuthOpcode = GAMO.getGlobal()->getValueType()->isFunctionTy()
2641-
? AArch64::AUTIA
2642-
: AArch64::AUTDA;
2643-
EmitToStreamer(MCInstBuilder(AuthOpcode)
2644-
.addReg(AuthResultReg)
2645-
.addReg(AuthResultReg)
2646-
.addReg(AArch64::X17));
2624+
2625+
bool IsFunctionTy = GAMO.getGlobal()->getValueType()->isFunctionTy();
2626+
auto AuthKey = IsFunctionTy ? AArch64PACKey::IA : AArch64PACKey::DA;
2627+
emitAUT(AuthKey, AuthResultReg, AArch64::X17);
26472628

26482629
if (GAMO.getGlobal()->hasExternalWeakLinkage())
26492630
OutStreamer->emitLabel(UndefWeakSym);
26502631

26512632
if (!STI->hasFPAC()) {
2652-
auto AuthKey =
2653-
(AuthOpcode == AArch64::AUTIA ? AArch64PACKey::IA : AArch64PACKey::DA);
2654-
26552633
emitPtrauthCheckAuthenticatedValue(AuthResultReg, AArch64::X17, AuthKey,
26562634
AArch64PAuth::AuthCheckMethod::XPAC);
26572635

@@ -3002,10 +2980,7 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
30022980
case AArch64::AUTH_TCRETURN:
30032981
case AArch64::AUTH_TCRETURN_BTI: {
30042982
Register Callee = MI->getOperand(0).getReg();
3005-
const uint64_t Key = MI->getOperand(2).getImm();
3006-
assert((Key == AArch64PACKey::IA || Key == AArch64PACKey::IB) &&
3007-
"Invalid auth key for tail-call return");
3008-
2983+
const auto Key = (AArch64PACKey::ID)MI->getOperand(2).getImm();
30092984
const uint64_t Disc = MI->getOperand(3).getImm();
30102985

30112986
Register AddrDisc = MI->getOperand(4).getReg();
@@ -3026,17 +3001,7 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
30263001
AddrDisc == AArch64::X16 || AddrDisc == AArch64::X17;
30273002
Register DiscReg = emitPtrauthDiscriminator(Disc, AddrDisc, ScratchReg,
30283003
AddrDiscIsImplicitDef);
3029-
3030-
const bool IsZero = DiscReg == AArch64::XZR;
3031-
const unsigned Opcodes[2][2] = {{AArch64::BRAA, AArch64::BRAAZ},
3032-
{AArch64::BRAB, AArch64::BRABZ}};
3033-
3034-
MCInst TmpInst;
3035-
TmpInst.setOpcode(Opcodes[Key][IsZero]);
3036-
TmpInst.addOperand(MCOperand::createReg(Callee));
3037-
if (!IsZero)
3038-
TmpInst.addOperand(MCOperand::createReg(DiscReg));
3039-
EmitToStreamer(*OutStreamer, TmpInst);
3004+
emitBLRA(/*IsCall*/ false, Key, Callee, DiscReg);
30403005
return;
30413006
}
30423007

‎llvm/lib/Target/AArch64/AArch64InstrInfo.h‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -806,6 +806,24 @@ static inline unsigned getPACOpcodeForKey(AArch64PACKey::ID K, bool Zero) {
806806
llvm_unreachable("Unhandled AArch64PACKey::ID enum");
807807
}
808808

809+
static inline unsigned getBranchOpcodeForKey(bool IsCall, AArch64PACKey::ID K,
810+
bool Zero) {
811+
using namespace AArch64PACKey;
812+
static const unsigned BranchOpcode[2][2] = {
813+
{AArch64::BRAA, AArch64::BRAAZ},
814+
{AArch64::BRAB, AArch64::BRABZ},
815+
};
816+
static const unsigned CallOpcode[2][2] = {
817+
{AArch64::BLRAA, AArch64::BLRAAZ},
818+
{AArch64::BLRAB, AArch64::BLRABZ},
819+
};
820+
821+
assert((K == IA || K == IB) && "I-key expected");
822+
if (IsCall)
823+
return CallOpcode[K == IB][Zero];
824+
return BranchOpcode[K == IB][Zero];
825+
}
826+
809827
// struct TSFlags {
810828
#define TSFLAG_ELEMENT_SIZE_TYPE(X) (X) // 3-bits
811829
#define TSFLAG_DESTRUCTIVE_INST_TYPE(X) ((X) << 3) // 4-bits

0 commit comments

Comments
 (0)