Skip to content

Commit 9157e20

Browse files
pccpcc
committed
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. Pull Request: llvm#132857
1 parent e9dbf31 commit 9157e20

10 files changed

+547
-352
lines changed

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

Lines changed: 42 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -168,7 +168,12 @@ class AArch64AsmPrinter : public AsmPrinter {
168168
void emitPtrauthTailCallHardening(const MachineInstr *TC);
169169

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

173178
// Emit the sequence to compute the discriminator.
174179
//
@@ -1866,7 +1871,6 @@ Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
18661871
Register AddrDisc,
18671872
Register ScratchReg,
18681873
bool MayUseAddrAsScratch) {
1869-
assert(ScratchReg == AArch64::X16 || ScratchReg == AArch64::X17);
18701874
// So far we've used NoRegister in pseudos. Now we need real encodings.
18711875
if (AddrDisc == AArch64::NoRegister)
18721876
AddrDisc = AArch64::XZR;
@@ -1886,7 +1890,8 @@ Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
18861890

18871891
// Check if we can save one MOV instruction.
18881892
assert(MayUseAddrAsScratch || ScratchReg != AddrDisc);
1889-
bool AddrDiscIsSafe = AddrDisc == AArch64::X16 || AddrDisc == AArch64::X17;
1893+
bool AddrDiscIsSafe = AddrDisc == AArch64::X16 || AddrDisc == AArch64::X17 ||
1894+
!STI->isX16X17Safer();
18901895
if (MayUseAddrAsScratch && AddrDiscIsSafe)
18911896
ScratchReg = AddrDisc;
18921897
else
@@ -2064,8 +2069,12 @@ void AArch64AsmPrinter::emitPtrauthTailCallHardening(const MachineInstr *TC) {
20642069
/*ShouldTrap=*/true, /*OnFailure=*/nullptr);
20652070
}
20662071

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

20702079
// We expand AUT/AUTPAC into a sequence of the form
20712080
//
@@ -2104,23 +2113,19 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
21042113
break;
21052114
}
21062115

2107-
auto AUTKey = (AArch64PACKey::ID)MI->getOperand(0).getImm();
2108-
uint64_t AUTDisc = MI->getOperand(1).getImm();
2109-
unsigned AUTAddrDisc = MI->getOperand(2).getReg();
2110-
21112116
// Compute aut discriminator into x17
21122117
assert(isUInt<16>(AUTDisc));
2113-
Register AUTDiscReg =
2114-
emitPtrauthDiscriminator(AUTDisc, AUTAddrDisc, AArch64::X17);
2118+
Register AUTDiscReg = emitPtrauthDiscriminator(
2119+
AUTDisc, AUTAddrDisc->getReg(), Scratch, AUTAddrDisc->isKill());
21152120
bool AUTZero = AUTDiscReg == AArch64::XZR;
21162121
unsigned AUTOpc = getAUTOpcodeForKey(AUTKey, AUTZero);
21172122

21182123
// autiza x16 ; if AUTZero
21192124
// autia x16, x17 ; if !AUTZero
21202125
MCInst AUTInst;
21212126
AUTInst.setOpcode(AUTOpc);
2122-
AUTInst.addOperand(MCOperand::createReg(AArch64::X16));
2123-
AUTInst.addOperand(MCOperand::createReg(AArch64::X16));
2127+
AUTInst.addOperand(MCOperand::createReg(AUTVal));
2128+
AUTInst.addOperand(MCOperand::createReg(AUTVal));
21242129
if (!AUTZero)
21252130
AUTInst.addOperand(MCOperand::createReg(AUTDiscReg));
21262131
EmitToStreamer(*OutStreamer, AUTInst);
@@ -2135,7 +2140,7 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
21352140
if (IsAUTPAC && !ShouldTrap)
21362141
EndSym = createTempSymbol("resign_end_");
21372142

2138-
emitPtrauthCheckAuthenticatedValue(AArch64::X16, AArch64::X17, AUTKey,
2143+
emitPtrauthCheckAuthenticatedValue(AUTVal, Scratch, AUTKey,
21392144
AArch64PAuth::AuthCheckMethod::XPAC,
21402145
ShouldTrap, EndSym);
21412146
}
@@ -2146,23 +2151,19 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
21462151
if (!IsAUTPAC)
21472152
return;
21482153

2149-
auto PACKey = (AArch64PACKey::ID)MI->getOperand(3).getImm();
2150-
uint64_t PACDisc = MI->getOperand(4).getImm();
2151-
unsigned PACAddrDisc = MI->getOperand(5).getReg();
2152-
21532154
// Compute pac discriminator into x17
21542155
assert(isUInt<16>(PACDisc));
21552156
Register PACDiscReg =
2156-
emitPtrauthDiscriminator(PACDisc, PACAddrDisc, AArch64::X17);
2157+
emitPtrauthDiscriminator(PACDisc, PACAddrDisc, Scratch);
21572158
bool PACZero = PACDiscReg == AArch64::XZR;
2158-
unsigned PACOpc = getPACOpcodeForKey(PACKey, PACZero);
2159+
unsigned PACOpc = getPACOpcodeForKey(*PACKey, PACZero);
21592160

21602161
// pacizb x16 ; if PACZero
21612162
// pacib x16, x17 ; if !PACZero
21622163
MCInst PACInst;
21632164
PACInst.setOpcode(PACOpc);
2164-
PACInst.addOperand(MCOperand::createReg(AArch64::X16));
2165-
PACInst.addOperand(MCOperand::createReg(AArch64::X16));
2165+
PACInst.addOperand(MCOperand::createReg(AUTVal));
2166+
PACInst.addOperand(MCOperand::createReg(AUTVal));
21662167
if (!PACZero)
21672168
PACInst.addOperand(MCOperand::createReg(PACDiscReg));
21682169
EmitToStreamer(*OutStreamer, PACInst);
@@ -2861,9 +2862,26 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
28612862
return;
28622863
}
28632864

2864-
case AArch64::AUT:
2865+
case AArch64::AUTx16x17:
2866+
emitPtrauthAuthResign(AArch64::X16,
2867+
(AArch64PACKey::ID)MI->getOperand(0).getImm(),
2868+
MI->getOperand(1).getImm(), &MI->getOperand(2),
2869+
AArch64::X17, std::nullopt, 0, 0);
2870+
return;
2871+
2872+
case AArch64::AUTxMxN:
2873+
emitPtrauthAuthResign(MI->getOperand(0).getReg(),
2874+
(AArch64PACKey::ID)MI->getOperand(3).getImm(),
2875+
MI->getOperand(4).getImm(), &MI->getOperand(5),
2876+
MI->getOperand(1).getReg(), std::nullopt, 0, 0);
2877+
return;
2878+
28652879
case AArch64::AUTPAC:
2866-
emitPtrauthAuthResign(MI);
2880+
emitPtrauthAuthResign(
2881+
AArch64::X16, (AArch64PACKey::ID)MI->getOperand(0).getImm(),
2882+
MI->getOperand(1).getImm(), &MI->getOperand(2), AArch64::X17,
2883+
(AArch64PACKey::ID)MI->getOperand(3).getImm(),
2884+
MI->getOperand(4).getImm(), MI->getOperand(5).getReg());
28672885
return;
28682886

28692887
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: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1938,10 +1938,11 @@ let Predicates = [HasPAuth] in {
19381938
defm LDRAB : AuthLoad<1, "ldrab", simm10Scaled>;
19391939

19401940
// AUT pseudo.
1941-
// This directly manipulates x16/x17, which are the only registers the OS
1942-
// guarantees are safe to use for sensitive operations.
1943-
def AUT : Pseudo<(outs), (ins i32imm:$Key, i64imm:$Disc, GPR64noip:$AddrDisc),
1944-
[]>, Sched<[WriteI, ReadI]> {
1941+
// This directly manipulates x16/x17, which are the only registers that
1942+
// certain OSs guarantee are safe to use for sensitive operations.
1943+
def AUTx16x17 : Pseudo<(outs), (ins i32imm:$Key, i64imm:$Disc,
1944+
GPR64noip:$AddrDisc),
1945+
[]>, Sched<[WriteI, ReadI]> {
19451946
let isCodeGenOnly = 1;
19461947
let hasSideEffects = 1;
19471948
let mayStore = 0;
@@ -1951,9 +1952,23 @@ let Predicates = [HasPAuth] in {
19511952
let Uses = [X16];
19521953
}
19531954

1955+
def AUTxMxN : Pseudo<(outs GPR64:$AuthVal, GPR64common:$Scratch),
1956+
(ins GPR64:$Val, i32imm:$Key,
1957+
i64imm:$Disc, GPR64:$AddrDisc),
1958+
[]>, Sched<[WriteI, ReadI]> {
1959+
let Constraints = "$AuthVal = $Val";
1960+
let isCodeGenOnly = 1;
1961+
let hasSideEffects = 0;
1962+
let mayStore = 0;
1963+
let mayLoad = 0;
1964+
let Size = 32;
1965+
let Defs = [NZCV];
1966+
let Uses = [];
1967+
}
1968+
19541969
// AUT and re-PAC a value, using different keys/data.
1955-
// This directly manipulates x16/x17, which are the only registers the OS
1956-
// guarantees are safe to use for sensitive operations.
1970+
// This directly manipulates x16/x17, which are the only registers that
1971+
// certain OSs guarantee are safe to use for sensitive operations.
19571972
def AUTPAC
19581973
: Pseudo<(outs),
19591974
(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
@@ -6737,14 +6737,27 @@ bool AArch64InstructionSelector::selectIntrinsic(MachineInstr &I,
67376737
std::tie(AUTConstDiscC, AUTAddrDisc) =
67386738
extractPtrauthBlendDiscriminators(AUTDisc, MRI);
67396739

6740-
MIB.buildCopy({AArch64::X16}, {ValReg});
6741-
MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X17}, {});
6742-
MIB.buildInstr(AArch64::AUT)
6743-
.addImm(AUTKey)
6744-
.addImm(AUTConstDiscC)
6745-
.addUse(AUTAddrDisc)
6746-
.constrainAllUses(TII, TRI, RBI);
6747-
MIB.buildCopy({DstReg}, Register(AArch64::X16));
6740+
if (STI.isX16X17Safer()) {
6741+
MIB.buildCopy({AArch64::X16}, {ValReg});
6742+
MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X17}, {});
6743+
MIB.buildInstr(AArch64::AUTx16x17)
6744+
.addImm(AUTKey)
6745+
.addImm(AUTConstDiscC)
6746+
.addUse(AUTAddrDisc)
6747+
.constrainAllUses(TII, TRI, RBI);
6748+
MIB.buildCopy({DstReg}, Register(AArch64::X16));
6749+
} else {
6750+
Register ScratchReg =
6751+
MRI.createVirtualRegister(&AArch64::GPR64commonRegClass);
6752+
MIB.buildInstr(AArch64::AUTxMxN)
6753+
.addDef(DstReg)
6754+
.addDef(ScratchReg)
6755+
.addUse(ValReg)
6756+
.addImm(AUTKey)
6757+
.addImm(AUTConstDiscC)
6758+
.addUse(AUTAddrDisc)
6759+
.constrainAllUses(TII, TRI, RBI);
6760+
}
67486761

67496762
RBI.constrainGenericRegister(DstReg, AArch64::GPR64RegClass, MRI);
67506763
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)