Skip to content

Commit 2197671

Browse files
pccpcc
authored
AArch64: Relax x16/x17 constraint on AUT in certain cases.
On most operating systems, the x16 and x17 registers are not special, so there is no benefit, and only a code size cost, to constraining AUT to only using them. Therefore, adjust the backend to only use the AUT pseudo (renamed AUTx16x17 for clarity) on Darwin platforms. All other platforms use an unconstrained variant of the pseudo, AUTxMxN, for selection. Reviewers: ahmedbougacha, kovdan01, atrosinenko Reviewed By: atrosinenko Pull Request: llvm/llvm-project#132857
1 parent 4ea0ef2 commit 2197671

10 files changed

+589
-354
lines changed

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

Lines changed: 49 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -169,18 +169,22 @@ class AArch64AsmPrinter : public AsmPrinter {
169169
void emitPtrauthTailCallHardening(const MachineInstr *TC);
170170

171171
// Emit the sequence for AUT or AUTPAC.
172-
void emitPtrauthAuthResign(const MachineInstr *MI);
172+
void emitPtrauthAuthResign(Register AUTVal, AArch64PACKey::ID AUTKey,
173+
uint64_t AUTDisc,
174+
const MachineOperand *AUTAddrDisc,
175+
Register Scratch,
176+
std::optional<AArch64PACKey::ID> PACKey,
177+
uint64_t PACDisc, Register PACAddrDisc);
173178

174179
// Emit the sequence to compute the discriminator.
175180
//
176-
// ScratchReg should be x16/x17.
177-
//
178-
// The returned register is either unmodified AddrDisc or x16/x17.
181+
// The returned register is either unmodified AddrDisc or ScratchReg.
179182
//
180183
// If the expanded pseudo is allowed to clobber AddrDisc register, setting
181184
// MayUseAddrAsScratch may save one MOV instruction, provided the address
182185
// is already in x16/x17 (i.e. return x16/x17 which is the *modified* AddrDisc
183-
// register at the same time):
186+
// register at the same time) or the OS doesn't make it safer to use x16/x17
187+
// (see AArch64Subtarget::isX16X17Safer()):
184188
//
185189
// mov x17, x16
186190
// movk x17, #1234, lsl #48
@@ -1867,7 +1871,8 @@ Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
18671871
Register AddrDisc,
18681872
Register ScratchReg,
18691873
bool MayUseAddrAsScratch) {
1870-
assert(ScratchReg == AArch64::X16 || ScratchReg == AArch64::X17);
1874+
assert(ScratchReg == AArch64::X16 || ScratchReg == AArch64::X17 ||
1875+
!STI->isX16X17Safer());
18711876
// So far we've used NoRegister in pseudos. Now we need real encodings.
18721877
if (AddrDisc == AArch64::NoRegister)
18731878
AddrDisc = AArch64::XZR;
@@ -1887,7 +1892,8 @@ Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
18871892

18881893
// Check if we can save one MOV instruction.
18891894
assert(MayUseAddrAsScratch || ScratchReg != AddrDisc);
1890-
bool AddrDiscIsSafe = AddrDisc == AArch64::X16 || AddrDisc == AArch64::X17;
1895+
bool AddrDiscIsSafe = AddrDisc == AArch64::X16 || AddrDisc == AArch64::X17 ||
1896+
!STI->isX16X17Safer();
18911897
if (MayUseAddrAsScratch && AddrDiscIsSafe)
18921898
ScratchReg = AddrDisc;
18931899
else
@@ -2065,8 +2071,12 @@ void AArch64AsmPrinter::emitPtrauthTailCallHardening(const MachineInstr *TC) {
20652071
/*ShouldTrap=*/true, /*OnFailure=*/nullptr);
20662072
}
20672073

2068-
void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
2069-
const bool IsAUTPAC = MI->getOpcode() == AArch64::AUTPAC;
2074+
void AArch64AsmPrinter::emitPtrauthAuthResign(
2075+
Register AUTVal, AArch64PACKey::ID AUTKey, uint64_t AUTDisc,
2076+
const MachineOperand *AUTAddrDisc, Register Scratch,
2077+
std::optional<AArch64PACKey::ID> PACKey, uint64_t PACDisc,
2078+
Register PACAddrDisc) {
2079+
const bool IsAUTPAC = PACKey.has_value();
20702080

20712081
// We expand AUT/AUTPAC into a sequence of the form
20722082
//
@@ -2105,23 +2115,19 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
21052115
break;
21062116
}
21072117

2108-
auto AUTKey = (AArch64PACKey::ID)MI->getOperand(0).getImm();
2109-
uint64_t AUTDisc = MI->getOperand(1).getImm();
2110-
unsigned AUTAddrDisc = MI->getOperand(2).getReg();
2111-
2112-
// Compute aut discriminator into x17
2118+
// Compute aut discriminator
21132119
assert(isUInt<16>(AUTDisc));
2114-
Register AUTDiscReg =
2115-
emitPtrauthDiscriminator(AUTDisc, AUTAddrDisc, AArch64::X17);
2120+
Register AUTDiscReg = emitPtrauthDiscriminator(
2121+
AUTDisc, AUTAddrDisc->getReg(), Scratch, AUTAddrDisc->isKill());
21162122
bool AUTZero = AUTDiscReg == AArch64::XZR;
21172123
unsigned AUTOpc = getAUTOpcodeForKey(AUTKey, AUTZero);
21182124

21192125
// autiza x16 ; if AUTZero
21202126
// autia x16, x17 ; if !AUTZero
21212127
MCInst AUTInst;
21222128
AUTInst.setOpcode(AUTOpc);
2123-
AUTInst.addOperand(MCOperand::createReg(AArch64::X16));
2124-
AUTInst.addOperand(MCOperand::createReg(AArch64::X16));
2129+
AUTInst.addOperand(MCOperand::createReg(AUTVal));
2130+
AUTInst.addOperand(MCOperand::createReg(AUTVal));
21252131
if (!AUTZero)
21262132
AUTInst.addOperand(MCOperand::createReg(AUTDiscReg));
21272133
EmitToStreamer(*OutStreamer, AUTInst);
@@ -2136,7 +2142,7 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
21362142
if (IsAUTPAC && !ShouldTrap)
21372143
EndSym = createTempSymbol("resign_end_");
21382144

2139-
emitPtrauthCheckAuthenticatedValue(AArch64::X16, AArch64::X17, AUTKey,
2145+
emitPtrauthCheckAuthenticatedValue(AUTVal, Scratch, AUTKey,
21402146
AArch64PAuth::AuthCheckMethod::XPAC,
21412147
ShouldTrap, EndSym);
21422148
}
@@ -2147,23 +2153,19 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
21472153
if (!IsAUTPAC)
21482154
return;
21492155

2150-
auto PACKey = (AArch64PACKey::ID)MI->getOperand(3).getImm();
2151-
uint64_t PACDisc = MI->getOperand(4).getImm();
2152-
unsigned PACAddrDisc = MI->getOperand(5).getReg();
2153-
2154-
// Compute pac discriminator into x17
2156+
// Compute pac discriminator
21552157
assert(isUInt<16>(PACDisc));
21562158
Register PACDiscReg =
2157-
emitPtrauthDiscriminator(PACDisc, PACAddrDisc, AArch64::X17);
2159+
emitPtrauthDiscriminator(PACDisc, PACAddrDisc, Scratch);
21582160
bool PACZero = PACDiscReg == AArch64::XZR;
2159-
unsigned PACOpc = getPACOpcodeForKey(PACKey, PACZero);
2161+
unsigned PACOpc = getPACOpcodeForKey(*PACKey, PACZero);
21602162

21612163
// pacizb x16 ; if PACZero
21622164
// pacib x16, x17 ; if !PACZero
21632165
MCInst PACInst;
21642166
PACInst.setOpcode(PACOpc);
2165-
PACInst.addOperand(MCOperand::createReg(AArch64::X16));
2166-
PACInst.addOperand(MCOperand::createReg(AArch64::X16));
2167+
PACInst.addOperand(MCOperand::createReg(AUTVal));
2168+
PACInst.addOperand(MCOperand::createReg(AUTVal));
21672169
if (!PACZero)
21682170
PACInst.addOperand(MCOperand::createReg(PACDiscReg));
21692171
EmitToStreamer(*OutStreamer, PACInst);
@@ -2866,9 +2868,26 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
28662868
return;
28672869
}
28682870

2869-
case AArch64::AUT:
2871+
case AArch64::AUTx16x17:
2872+
emitPtrauthAuthResign(AArch64::X16,
2873+
(AArch64PACKey::ID)MI->getOperand(0).getImm(),
2874+
MI->getOperand(1).getImm(), &MI->getOperand(2),
2875+
AArch64::X17, std::nullopt, 0, 0);
2876+
return;
2877+
2878+
case AArch64::AUTxMxN:
2879+
emitPtrauthAuthResign(MI->getOperand(0).getReg(),
2880+
(AArch64PACKey::ID)MI->getOperand(3).getImm(),
2881+
MI->getOperand(4).getImm(), &MI->getOperand(5),
2882+
MI->getOperand(1).getReg(), std::nullopt, 0, 0);
2883+
return;
2884+
28702885
case AArch64::AUTPAC:
2871-
emitPtrauthAuthResign(MI);
2886+
emitPtrauthAuthResign(
2887+
AArch64::X16, (AArch64PACKey::ID)MI->getOperand(0).getImm(),
2888+
MI->getOperand(1).getImm(), &MI->getOperand(2), AArch64::X17,
2889+
(AArch64PACKey::ID)MI->getOperand(3).getImm(),
2890+
MI->getOperand(4).getImm(), MI->getOperand(5).getReg());
28722891
return;
28732892

28742893
case AArch64::LOADauthptrstatic:

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

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1534,12 +1534,20 @@ void AArch64DAGToDAGISel::SelectPtrauthAuth(SDNode *N) {
15341534
std::tie(AUTConstDisc, AUTAddrDisc) =
15351535
extractPtrauthBlendDiscriminators(AUTDisc, CurDAG);
15361536

1537-
SDValue X16Copy = CurDAG->getCopyToReg(CurDAG->getEntryNode(), DL,
1538-
AArch64::X16, Val, SDValue());
1539-
SDValue Ops[] = {AUTKey, AUTConstDisc, AUTAddrDisc, X16Copy.getValue(1)};
1537+
if (!Subtarget->isX16X17Safer()) {
1538+
SDValue Ops[] = {Val, AUTKey, AUTConstDisc, AUTAddrDisc};
15401539

1541-
SDNode *AUT = CurDAG->getMachineNode(AArch64::AUT, DL, MVT::i64, Ops);
1542-
ReplaceNode(N, AUT);
1540+
SDNode *AUT =
1541+
CurDAG->getMachineNode(AArch64::AUTxMxN, DL, MVT::i64, MVT::i64, Ops);
1542+
ReplaceNode(N, AUT);
1543+
} else {
1544+
SDValue X16Copy = CurDAG->getCopyToReg(CurDAG->getEntryNode(), DL,
1545+
AArch64::X16, Val, SDValue());
1546+
SDValue Ops[] = {AUTKey, AUTConstDisc, AUTAddrDisc, X16Copy.getValue(1)};
1547+
1548+
SDNode *AUT = CurDAG->getMachineNode(AArch64::AUTx16x17, DL, MVT::i64, Ops);
1549+
ReplaceNode(N, AUT);
1550+
}
15431551
}
15441552

15451553
void AArch64DAGToDAGISel::SelectPtrauthResign(SDNode *N) {

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

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2129,10 +2129,11 @@ let Predicates = [HasPAuth] in {
21292129
defm LDRAB : AuthLoad<1, "ldrab", simm10Scaled>;
21302130

21312131
// AUT pseudo.
2132-
// This directly manipulates x16/x17, which are the only registers the OS
2133-
// guarantees are safe to use for sensitive operations.
2134-
def AUT : Pseudo<(outs), (ins i32imm:$Key, i64imm:$Disc, GPR64noip:$AddrDisc),
2135-
[]>, Sched<[WriteI, ReadI]> {
2132+
// This directly manipulates x16/x17, which are the only registers that
2133+
// certain OSs guarantee are safe to use for sensitive operations.
2134+
def AUTx16x17 : Pseudo<(outs), (ins i32imm:$Key, i64imm:$Disc,
2135+
GPR64noip:$AddrDisc),
2136+
[]>, Sched<[WriteI, ReadI]> {
21362137
let isCodeGenOnly = 1;
21372138
let hasSideEffects = 1;
21382139
let mayStore = 0;
@@ -2142,9 +2143,22 @@ let Predicates = [HasPAuth] in {
21422143
let Uses = [X16];
21432144
}
21442145

2146+
def AUTxMxN : Pseudo<(outs GPR64:$AuthVal, GPR64common:$Scratch),
2147+
(ins GPR64:$Val, i32imm:$Key,
2148+
i64imm:$Disc, GPR64:$AddrDisc),
2149+
[], "$AuthVal = $Val">, Sched<[WriteI, ReadI]> {
2150+
let isCodeGenOnly = 1;
2151+
let hasSideEffects = 0;
2152+
let mayStore = 0;
2153+
let mayLoad = 0;
2154+
let Size = 32;
2155+
let Defs = [NZCV];
2156+
let Uses = [];
2157+
}
2158+
21452159
// AUT and re-PAC a value, using different keys/data.
2146-
// This directly manipulates x16/x17, which are the only registers the OS
2147-
// guarantees are safe to use for sensitive operations.
2160+
// This directly manipulates x16/x17, which are the only registers that
2161+
// certain OSs guarantee are safe to use for sensitive operations.
21482162
def AUTPAC
21492163
: Pseudo<(outs),
21502164
(ins i32imm:$AUTKey, i64imm:$AUTDisc, GPR64noip:$AUTAddrDisc,

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -664,6 +664,12 @@ AArch64Subtarget::getPtrAuthBlockAddressDiscriminatorIfEnabled(
664664
(Twine(ParentFn.getName()) + " blockaddress").str());
665665
}
666666

667+
bool AArch64Subtarget::isX16X17Safer() const {
668+
// The Darwin kernel implements special protections for x16 and x17 so we
669+
// should prefer to use those registers on that platform.
670+
return isTargetDarwin();
671+
}
672+
667673
bool AArch64Subtarget::enableMachinePipeliner() const {
668674
return getSchedModel().hasInstrSchedModel();
669675
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,10 @@ class AArch64Subtarget final : public AArch64GenSubtargetInfo {
318318
}
319319
}
320320

321+
/// Returns whether the operating system makes it safer to store sensitive
322+
/// values in x16 and x17 as opposed to other registers.
323+
bool isX16X17Safer() const;
324+
321325
/// ParseSubtargetFeatures - Parses features string setting specified
322326
/// subtarget options. Definition of function is auto generated by tblgen.
323327
void ParseSubtargetFeatures(StringRef CPU, StringRef TuneCPU, StringRef FS);

‎llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6762,14 +6762,27 @@ bool AArch64InstructionSelector::selectIntrinsic(MachineInstr &I,
67626762
std::tie(AUTConstDiscC, AUTAddrDisc) =
67636763
extractPtrauthBlendDiscriminators(AUTDisc, MRI);
67646764

6765-
MIB.buildCopy({AArch64::X16}, {ValReg});
6766-
MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X17}, {});
6767-
MIB.buildInstr(AArch64::AUT)
6768-
.addImm(AUTKey)
6769-
.addImm(AUTConstDiscC)
6770-
.addUse(AUTAddrDisc)
6771-
.constrainAllUses(TII, TRI, RBI);
6772-
MIB.buildCopy({DstReg}, Register(AArch64::X16));
6765+
if (STI.isX16X17Safer()) {
6766+
MIB.buildCopy({AArch64::X16}, {ValReg});
6767+
MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X17}, {});
6768+
MIB.buildInstr(AArch64::AUTx16x17)
6769+
.addImm(AUTKey)
6770+
.addImm(AUTConstDiscC)
6771+
.addUse(AUTAddrDisc)
6772+
.constrainAllUses(TII, TRI, RBI);
6773+
MIB.buildCopy({DstReg}, Register(AArch64::X16));
6774+
} else {
6775+
Register ScratchReg =
6776+
MRI.createVirtualRegister(&AArch64::GPR64commonRegClass);
6777+
MIB.buildInstr(AArch64::AUTxMxN)
6778+
.addDef(DstReg)
6779+
.addDef(ScratchReg)
6780+
.addUse(ValReg)
6781+
.addImm(AUTKey)
6782+
.addImm(AUTConstDiscC)
6783+
.addUse(AUTAddrDisc)
6784+
.constrainAllUses(TII, TRI, RBI);
6785+
}
67736786

67746787
RBI.constrainGenericRegister(DstReg, AArch64::GPR64RegClass, MRI);
67756788
I.eraseFromParent();

‎llvm/test/CodeGen/AArch64/ptrauth-call.ll

Lines changed: 31 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -169,13 +169,20 @@ define i32 @test_tailcall_ib_var(ptr %arg0, ptr %arg1) #0 {
169169

170170
define void @test_tailcall_omit_mov_x16_x16(ptr %objptr) #0 {
171171
; CHECK-LABEL: test_tailcall_omit_mov_x16_x16:
172-
; CHECK-NEXT: ldr x16, [x0]
173-
; CHECK-NEXT: mov x17, x0
174-
; CHECK-NEXT: movk x17, #6503, lsl #48
175-
; CHECK-NEXT: autda x16, x17
176-
; CHECK-NEXT: ldr x1, [x16]
177-
; CHECK-NEXT: movk x16, #54167, lsl #48
178-
; CHECK-NEXT: braa x1, x16
172+
; DARWIN-NEXT: ldr x16, [x0]
173+
; DARWIN-NEXT: mov x17, x0
174+
; DARWIN-NEXT: movk x17, #6503, lsl #48
175+
; DARWIN-NEXT: autda x16, x17
176+
; DARWIN-NEXT: ldr x1, [x16]
177+
; DARWIN-NEXT: movk x16, #54167, lsl #48
178+
; DARWIN-NEXT: braa x1, x16
179+
; ELF-NEXT: ldr x1, [x0]
180+
; ELF-NEXT: mov x8, x0
181+
; ELF-NEXT: movk x8, #6503, lsl #48
182+
; ELF-NEXT: autda x1, x8
183+
; ELF-NEXT: ldr x2, [x1]
184+
; ELF-NEXT: movk x1, #54167, lsl #48
185+
; ELF-NEXT: braa x2, x1
179186
%vtable.signed = load ptr, ptr %objptr, align 8
180187
%objptr.int = ptrtoint ptr %objptr to i64
181188
%vtable.discr = tail call i64 @llvm.ptrauth.blend(i64 %objptr.int, i64 6503)
@@ -191,16 +198,24 @@ define void @test_tailcall_omit_mov_x16_x16(ptr %objptr) #0 {
191198
define i32 @test_call_omit_extra_moves(ptr %objptr) #0 {
192199
; CHECK-LABEL: test_call_omit_extra_moves:
193200
; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
194-
; ELF-NEXT: str x30, [sp, #-16]!
195-
; CHECK-NEXT: ldr x16, [x0]
196-
; CHECK-NEXT: mov x17, x0
197-
; CHECK-NEXT: movk x17, #6503, lsl #48
198-
; CHECK-NEXT: autda x16, x17
199-
; CHECK-NEXT: ldr x8, [x16]
200-
; CHECK-NEXT: movk x16, #34646, lsl #48
201-
; CHECK-NEXT: blraa x8, x16
202-
; CHECK-NEXT: mov w0, #42
201+
; DARWIN-NEXT: ldr x16, [x0]
202+
; DARWIN-NEXT: mov x17, x0
203+
; DARWIN-NEXT: movk x17, #6503, lsl #48
204+
; DARWIN-NEXT: autda x16, x17
205+
; DARWIN-NEXT: ldr x8, [x16]
206+
; DARWIN-NEXT: movk x16, #34646, lsl #48
207+
; DARWIN-NEXT: blraa x8, x16
208+
; DARWIN-NEXT: mov w0, #42
203209
; DARWIN-NEXT: ldp x29, x30, [sp], #16
210+
; ELF-NEXT: str x30, [sp, #-16]!
211+
; ELF-NEXT: ldr x8, [x0]
212+
; ELF-NEXT: mov x9, x0
213+
; ELF-NEXT: movk x9, #6503, lsl #48
214+
; ELF-NEXT: autda x8, x9
215+
; ELF-NEXT: ldr x9, [x8]
216+
; ELF-NEXT: movk x8, #34646, lsl #48
217+
; ELF-NEXT: blraa x9, x8
218+
; ELF-NEXT: mov w0, #42
204219
; ELF-NEXT: ldr x30, [sp], #16
205220
; CHECK-NEXT: ret
206221
%vtable.signed = load ptr, ptr %objptr

0 commit comments

Comments
 (0)