[llvm-branch-commits] [llvm] [AArch64][PAC] Eliminate excessive MOVs when computing blend (PR #115185)
Anatoly Trosinenko via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Mon Nov 11 08:38:25 PST 2024
https://github.com/atrosinenko updated https://github.com/llvm/llvm-project/pull/115185
>From c4fa68c7ad66911e006c175ec9acf9e3cca9c1d9 Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko <atrosinenko at accesssoftek.com>
Date: Mon, 21 Oct 2024 17:56:40 +0300
Subject: [PATCH 1/2] [AArch64][PAC] Eliminate excessive MOVs when computing
blend
As function calls do not generally preserve X16 and X17, it is beneficial
to allow AddrDisc operand of B(L)RA instruction to reside in these
registers and make use of this condition when computing the discriminator.
This can save up to two MOVs in cases such as loading a (signed) virtual
function pointer via a (signed) pointer to vtable, for example
ldr x9, [x16]
mov x8, x16
mov x17, x8
movk x17, #34646, lsl #48
blraa x9, x17
can be simplified to
ldr x8, [x16]
movk x16, #34646, lsl #48
blraa x8, x16
---
llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp | 94 +++++++++++--------
llvm/lib/Target/AArch64/AArch64InstrInfo.td | 18 ++--
llvm/test/CodeGen/AArch64/ptrauth-call.ll | 27 ++++++
3 files changed, 89 insertions(+), 50 deletions(-)
diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index 436bb332053f75..3263bb38ef1fcc 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -163,8 +163,15 @@ class AArch64AsmPrinter : public AsmPrinter {
// Emit the sequence for AUT or AUTPAC.
void emitPtrauthAuthResign(const MachineInstr *MI);
- // Emit the sequence to compute a discriminator into x17, or reuse AddrDisc.
- unsigned emitPtrauthDiscriminator(uint16_t Disc, unsigned AddrDisc);
+ // Emit the sequence to compute the discriminator.
+ // ScratchReg should be x16/x17.
+ // The returned register is either unmodified AddrDisc or x16/x17.
+ // 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.
+ Register emitPtrauthDiscriminator(uint16_t Disc, Register AddrDisc,
+ Register ScratchReg,
+ bool MayUseAddrAsScratch = false);
// Emit the sequence for LOADauthptrstatic
void LowerLOADauthptrstatic(const MachineInstr &MI);
@@ -1727,8 +1734,10 @@ void AArch64AsmPrinter::emitFMov0(const MachineInstr &MI) {
}
}
-unsigned AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
- unsigned AddrDisc) {
+Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
+ Register AddrDisc,
+ Register ScratchReg,
+ bool MayUseAddrAsScratch) {
// So far we've used NoRegister in pseudos. Now we need real encodings.
if (AddrDisc == AArch64::NoRegister)
AddrDisc = AArch64::XZR;
@@ -1738,16 +1747,24 @@ unsigned AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
if (!Disc)
return AddrDisc;
- // If there's only a constant discriminator, MOV it into x17.
+ // If there's only a constant discriminator, MOV it into the scratch register.
if (AddrDisc == AArch64::XZR) {
- emitMOVZ(AArch64::X17, Disc, 0);
- return AArch64::X17;
+ emitMOVZ(ScratchReg, Disc, 0);
+ return ScratchReg;
}
- // If there are both, emit a blend into x17.
- emitMovXReg(AArch64::X17, AddrDisc);
- emitMOVK(AArch64::X17, Disc, 48);
- return AArch64::X17;
+ // 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;
+ if (MayUseAddrAsScratch && AddrDiscIsSafe)
+ ScratchReg = AddrDisc;
+ else
+ emitMovXReg(ScratchReg, AddrDisc);
+
+ emitMOVK(ScratchReg, Disc, 48);
+ return ScratchReg;
}
/// Emits a code sequence to check an authenticated pointer value.
@@ -1964,7 +1981,8 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
// Compute aut discriminator into x17
assert(isUInt<16>(AUTDisc));
- unsigned AUTDiscReg = emitPtrauthDiscriminator(AUTDisc, AUTAddrDisc);
+ Register AUTDiscReg =
+ emitPtrauthDiscriminator(AUTDisc, AUTAddrDisc, AArch64::X17);
bool AUTZero = AUTDiscReg == AArch64::XZR;
unsigned AUTOpc = getAUTOpcodeForKey(AUTKey, AUTZero);
@@ -2005,7 +2023,8 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
// Compute pac discriminator into x17
assert(isUInt<16>(PACDisc));
- unsigned PACDiscReg = emitPtrauthDiscriminator(PACDisc, PACAddrDisc);
+ Register PACDiscReg =
+ emitPtrauthDiscriminator(PACDisc, PACAddrDisc, AArch64::X17);
bool PACZero = PACDiscReg == AArch64::XZR;
unsigned PACOpc = getPACOpcodeForKey(PACKey, PACZero);
@@ -2037,8 +2056,17 @@ void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
unsigned AddrDisc = MI->getOperand(3).getReg();
- // Compute discriminator into x17
- unsigned DiscReg = emitPtrauthDiscriminator(Disc, AddrDisc);
+ // Make sure AddrDisc is solely used to compute the discriminator.
+ // While hardly meaningful, it is still possible to describe an authentication
+ // of a pointer against its own value (instead of storage address) with
+ // intrinsics, so use report_fatal_error instead of assert.
+ if (BrTarget == AddrDisc)
+ report_fatal_error("Branch target is signed with its own value");
+
+ // x16 and x17 are implicit-def'ed by MI, and AddrDisc is not used as any
+ // other input, so try to save one MOV by setting MayUseAddrAsScratch.
+ Register DiscReg = emitPtrauthDiscriminator(Disc, AddrDisc, AArch64::X17,
+ /*MayUseAddrAsScratch=*/true);
bool IsZeroDisc = DiscReg == AArch64::XZR;
unsigned Opc;
@@ -2332,16 +2360,7 @@ void AArch64AsmPrinter::LowerMOVaddrPAC(const MachineInstr &MI) {
}
}
- unsigned DiscReg = AddrDisc;
- if (Disc != 0) {
- if (AddrDisc != AArch64::XZR) {
- emitMovXReg(AArch64::X17, AddrDisc);
- emitMOVK(AArch64::X17, Disc, 48);
- } else {
- emitMOVZ(AArch64::X17, Disc, 0);
- }
- DiscReg = AArch64::X17;
- }
+ Register DiscReg = emitPtrauthDiscriminator(Disc, AddrDisc, AArch64::X17);
auto MIB = MCInstBuilder(getPACOpcodeForKey(Key, DiscReg == AArch64::XZR))
.addReg(AArch64::X16)
@@ -2609,6 +2628,7 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
// instruction here.
case AArch64::AUTH_TCRETURN:
case AArch64::AUTH_TCRETURN_BTI: {
+ Register Callee = MI->getOperand(0).getReg();
const uint64_t Key = MI->getOperand(2).getImm();
assert((Key == AArch64PACKey::IA || Key == AArch64PACKey::IB) &&
"Invalid auth key for tail-call return");
@@ -2618,31 +2638,23 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
Register AddrDisc = MI->getOperand(4).getReg();
- Register ScratchReg = MI->getOperand(0).getReg() == AArch64::X16
- ? AArch64::X17
- : AArch64::X16;
+ Register ScratchReg = Callee == AArch64::X16 ? AArch64::X17 : AArch64::X16;
emitPtrauthTailCallHardening(MI);
- unsigned DiscReg = AddrDisc;
- if (Disc) {
- if (AddrDisc != AArch64::NoRegister) {
- if (ScratchReg != AddrDisc)
- emitMovXReg(ScratchReg, AddrDisc);
- emitMOVK(ScratchReg, Disc, 48);
- } else {
- emitMOVZ(ScratchReg, Disc, 0);
- }
- DiscReg = ScratchReg;
- }
+ // See the comments in emitPtrauthBranch.
+ if (Callee == AddrDisc)
+ report_fatal_error("Call target is signed with its own value");
+ Register DiscReg = emitPtrauthDiscriminator(Disc, AddrDisc, ScratchReg,
+ /*MayUseAddrAsScratch=*/true);
- const bool IsZero = DiscReg == AArch64::NoRegister;
+ const bool IsZero = DiscReg == AArch64::XZR;
const unsigned Opcodes[2][2] = {{AArch64::BRAA, AArch64::BRAAZ},
{AArch64::BRAB, AArch64::BRABZ}};
MCInst TmpInst;
TmpInst.setOpcode(Opcodes[Key][IsZero]);
- TmpInst.addOperand(MCOperand::createReg(MI->getOperand(0).getReg()));
+ TmpInst.addOperand(MCOperand::createReg(Callee));
if (!IsZero)
TmpInst.addOperand(MCOperand::createReg(DiscReg));
EmitToStreamer(*OutStreamer, TmpInst);
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
index dfa3209928068b..8c019aed3e9793 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.td
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
@@ -1840,28 +1840,28 @@ let Predicates = [HasPAuth] in {
// materialization here), in part because they're handled in a safer way by
// the kernel, notably on Darwin.
def BLRA : Pseudo<(outs), (ins GPR64noip:$Rn, i32imm:$Key, i64imm:$Disc,
- GPR64noip:$AddrDisc),
+ GPR64:$AddrDisc),
[(AArch64authcall GPR64noip:$Rn, timm:$Key, timm:$Disc,
- GPR64noip:$AddrDisc)]>, Sched<[]> {
+ GPR64:$AddrDisc)]>, Sched<[]> {
let isCodeGenOnly = 1;
let hasSideEffects = 1;
let mayStore = 0;
let mayLoad = 0;
let isCall = 1;
let Size = 12; // 4 fixed + 8 variable, to compute discriminator.
- let Defs = [X17,LR];
+ let Defs = [X16,X17,LR];
let Uses = [SP];
}
def BLRA_RVMARKER : Pseudo<
(outs), (ins i64imm:$rvfunc, GPR64noip:$Rn, i32imm:$Key, i64imm:$Disc,
- GPR64noip:$AddrDisc),
+ GPR64:$AddrDisc),
[(AArch64authcall_rvmarker tglobaladdr:$rvfunc,
GPR64noip:$Rn, timm:$Key, timm:$Disc,
- GPR64noip:$AddrDisc)]>, Sched<[]> {
+ GPR64:$AddrDisc)]>, Sched<[]> {
let isCodeGenOnly = 1;
let isCall = 1;
- let Defs = [X17,LR];
+ let Defs = [X16,X17,LR];
let Uses = [SP];
}
@@ -1869,7 +1869,7 @@ let Predicates = [HasPAuth] in {
// This directly manipulates x16/x17, which are the only registers the OS
// guarantees are safe to use for sensitive operations.
def BRA : Pseudo<(outs), (ins GPR64noip:$Rn, i32imm:$Key, i64imm:$Disc,
- GPR64noip:$AddrDisc), []>, Sched<[]> {
+ GPR64:$AddrDisc), []>, Sched<[]> {
let isCodeGenOnly = 1;
let hasNoSchedulingInfo = 1;
let hasSideEffects = 1;
@@ -1880,7 +1880,7 @@ let Predicates = [HasPAuth] in {
let isBarrier = 1;
let isIndirectBranch = 1;
let Size = 12; // 4 fixed + 8 variable, to compute discriminator.
- let Defs = [X17];
+ let Defs = [X16,X17];
}
let isReturn = 1, isTerminator = 1, isBarrier = 1 in {
@@ -1971,7 +1971,7 @@ let Predicates = [HasPAuth] in {
// make sure at least one register is usable as a scratch one - for that
// purpose, use tcGPRnotx16x17 register class for the second operand.
let isCall = 1, isTerminator = 1, isReturn = 1, isBarrier = 1, Size = 16,
- Uses = [SP] in {
+ Defs = [X16,X17], Uses = [SP] in {
def AUTH_TCRETURN
: Pseudo<(outs), (ins tcGPRnotx16x17:$dst, i32imm:$FPDiff, i32imm:$Key,
i64imm:$Disc, tcGPR64:$AddrDisc),
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-call.ll b/llvm/test/CodeGen/AArch64/ptrauth-call.ll
index c46fe79f754fef..bf35cf8fecbdb7 100644
--- a/llvm/test/CodeGen/AArch64/ptrauth-call.ll
+++ b/llvm/test/CodeGen/AArch64/ptrauth-call.ll
@@ -188,6 +188,33 @@ define void @test_tailcall_omit_mov_x16_x16(ptr %objptr) #0 {
ret void
}
+define i32 @test_call_omit_extra_moves(ptr %objptr) #0 {
+; CHECK-LABEL: test_call_omit_extra_moves:
+; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
+; ELF-NEXT: str x30, [sp, #-16]!
+; CHECK-NEXT: ldr x16, [x0]
+; CHECK-NEXT: mov x17, x0
+; CHECK-NEXT: movk x17, #6503, lsl #48
+; CHECK-NEXT: autda x16, x17
+; CHECK-NEXT: ldr x8, [x16]
+; CHECK-NEXT: movk x16, #34646, lsl #48
+; CHECK-NEXT: blraa x8, x16
+; CHECK-NEXT: mov w0, #42
+; DARWIN-NEXT: ldp x29, x30, [sp], #16
+; ELF-NEXT: ldr x30, [sp], #16
+; CHECK-NEXT: ret
+ %vtable.signed = load ptr, ptr %objptr
+ %objptr.int = ptrtoint ptr %objptr to i64
+ %vtable.discr = tail call i64 @llvm.ptrauth.blend(i64 %objptr.int, i64 6503)
+ %vtable.signed.int = ptrtoint ptr %vtable.signed to i64
+ %vtable.int = tail call i64 @llvm.ptrauth.auth(i64 %vtable.signed.int, i32 2, i64 %vtable.discr)
+ %vtable = inttoptr i64 %vtable.int to ptr
+ %callee.signed = load ptr, ptr %vtable
+ %callee.discr = tail call i64 @llvm.ptrauth.blend(i64 %vtable.int, i64 34646)
+ %call.result = tail call i32 %callee.signed(ptr %objptr) [ "ptrauth"(i32 0, i64 %callee.discr) ]
+ ret i32 42
+}
+
define i32 @test_call_ia_arg(ptr %arg0, i64 %arg1) #0 {
; DARWIN-LABEL: test_call_ia_arg:
; DARWIN-NEXT: stp x29, x30, [sp, #-16]!
>From 05c9852d7e27a79bdf9e291247ad9f731abdb500 Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko <atrosinenko at accesssoftek.com>
Date: Mon, 11 Nov 2024 17:48:40 +0300
Subject: [PATCH 2/2] Fixes
---
llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp | 24 +++++++++++++++----
llvm/lib/Target/AArch64/AArch64InstrInfo.td | 4 ++--
2 files changed, 22 insertions(+), 6 deletions(-)
diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index 3263bb38ef1fcc..9a502557c8daf6 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -164,11 +164,23 @@ class AArch64AsmPrinter : public AsmPrinter {
void emitPtrauthAuthResign(const MachineInstr *MI);
// Emit the sequence to compute the discriminator.
+ //
// ScratchReg should be x16/x17.
+ //
// The returned register is either unmodified AddrDisc or x16/x17.
+ //
// 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.
+ // is already in x16/x17 (i.e. return x16/x17 which is the *modified* AddrDisc
+ // register at the same time):
+ //
+ // mov x17, x16
+ // movk x17, #1234, lsl #48
+ // ; x16 is not used anymore
+ //
+ // can be replaced by
+ //
+ // movk x16, #1234, lsl #48
Register emitPtrauthDiscriminator(uint16_t Disc, Register AddrDisc,
Register ScratchReg,
bool MayUseAddrAsScratch = false);
@@ -1738,6 +1750,7 @@ Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
Register AddrDisc,
Register ScratchReg,
bool MayUseAddrAsScratch) {
+ assert(ScratchReg == AArch64::X16 || ScratchReg == AArch64::X17);
// So far we've used NoRegister in pseudos. Now we need real encodings.
if (AddrDisc == AArch64::NoRegister)
AddrDisc = AArch64::XZR;
@@ -2063,10 +2076,13 @@ void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
if (BrTarget == AddrDisc)
report_fatal_error("Branch target is signed with its own value");
- // x16 and x17 are implicit-def'ed by MI, and AddrDisc is not used as any
- // other input, so try to save one MOV by setting MayUseAddrAsScratch.
+ // If we are printing BLRA pseudo instruction, then x16 and x17 are
+ // implicit-def'ed by the MI and AddrDisc is not used as any other input, so
+ // try to save one MOV by setting MayUseAddrAsScratch.
+ // Unlike BLRA, BRA pseudo is used to perform computed goto, and thus not
+ // declared as clobbering x16/x17.
Register DiscReg = emitPtrauthDiscriminator(Disc, AddrDisc, AArch64::X17,
- /*MayUseAddrAsScratch=*/true);
+ /*MayUseAddrAsScratch=*/IsCall);
bool IsZeroDisc = DiscReg == AArch64::XZR;
unsigned Opc;
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
index 8c019aed3e9793..fb786f87e4b15f 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.td
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
@@ -1869,7 +1869,7 @@ let Predicates = [HasPAuth] in {
// This directly manipulates x16/x17, which are the only registers the OS
// guarantees are safe to use for sensitive operations.
def BRA : Pseudo<(outs), (ins GPR64noip:$Rn, i32imm:$Key, i64imm:$Disc,
- GPR64:$AddrDisc), []>, Sched<[]> {
+ GPR64noip:$AddrDisc), []>, Sched<[]> {
let isCodeGenOnly = 1;
let hasNoSchedulingInfo = 1;
let hasSideEffects = 1;
@@ -1880,7 +1880,7 @@ let Predicates = [HasPAuth] in {
let isBarrier = 1;
let isIndirectBranch = 1;
let Size = 12; // 4 fixed + 8 variable, to compute discriminator.
- let Defs = [X16,X17];
+ let Defs = [X17];
}
let isReturn = 1, isTerminator = 1, isBarrier = 1 in {
More information about the llvm-branch-commits
mailing list