Skip to content

Commit 0c9a20c

Browse files
committed
[AArch64][PAC] Prevent immediate modifier substitution of AUT/PAC
Expand `blend(addr, imm)` in AsmPrinter, so that the discriminator cannot be changed between blending `imm` in `addr` and using the discriminator for signing or authenticating a pointer. For that purpose, a new PAC pseudo instruction is introduced that has separate operands for register and immediate modifiers. Instruction selection is refactored along the way, so that AUT and PAC instructions are selected by tablegen-erated patterns (always with a zero immediate modifier at first) and a unified fixupBlendComponents function is used by a custom inserter to fetch operands of `blend()` when possible. Aside from unification, this helps inspecting operands of blend operation that is located in another basic block which can be tricky for DAGISel. This patch is somewhat related to llvm#132857.
1 parent b4c02ae commit 0c9a20c

10 files changed

+402
-110
lines changed

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,9 @@ class AArch64AsmPrinter : public AsmPrinter {
171171
// Emit the sequence for AUT or AUTPAC.
172172
void emitPtrauthAuthResign(const MachineInstr *MI);
173173

174+
// Emit the sequence for PAC.
175+
void emitPtrauthSign(const MachineInstr *MI);
176+
174177
// Emit the sequence to compute the discriminator.
175178
//
176179
// ScratchReg should be x16/x17.
@@ -2173,6 +2176,30 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
21732176
OutStreamer->emitLabel(EndSym);
21742177
}
21752178

2179+
void AArch64AsmPrinter::emitPtrauthSign(const MachineInstr *MI) {
2180+
Register Val = MI->getOperand(1).getReg();
2181+
auto Key = (AArch64PACKey::ID)MI->getOperand(2).getImm();
2182+
uint64_t Disc = MI->getOperand(3).getImm();
2183+
Register AddrDisc = MI->getOperand(4).getReg();
2184+
// Can isKill() be reliably used at this point?
2185+
2186+
// Compute aut discriminator into x17
2187+
assert(isUInt<16>(Disc));
2188+
Register DiscReg = emitPtrauthDiscriminator(Disc, AddrDisc, AArch64::X17);
2189+
bool IsZeroDisc = DiscReg == AArch64::XZR;
2190+
unsigned Opc = getPACOpcodeForKey(Key, IsZeroDisc);
2191+
2192+
// paciza x16 ; if IsZeroDisc
2193+
// pacia x16, x17 ; if !IsZeroDisc
2194+
MCInst PACInst;
2195+
PACInst.setOpcode(Opc);
2196+
PACInst.addOperand(MCOperand::createReg(Val));
2197+
PACInst.addOperand(MCOperand::createReg(Val));
2198+
if (!IsZeroDisc)
2199+
PACInst.addOperand(MCOperand::createReg(DiscReg));
2200+
EmitToStreamer(*OutStreamer, PACInst);
2201+
}
2202+
21762203
void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
21772204
bool IsCall = MI->getOpcode() == AArch64::BLRA;
21782205
unsigned BrTarget = MI->getOperand(0).getReg();
@@ -2867,6 +2894,10 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
28672894
emitPtrauthAuthResign(MI);
28682895
return;
28692896

2897+
case AArch64::PAC:
2898+
emitPtrauthSign(MI);
2899+
return;
2900+
28702901
case AArch64::LOADauthptrstatic:
28712902
LowerLOADauthptrstatic(*MI);
28722903
return;

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -361,7 +361,6 @@ class AArch64DAGToDAGISel : public SelectionDAGISel {
361361

362362
bool tryIndexedLoad(SDNode *N);
363363

364-
void SelectPtrauthAuth(SDNode *N);
365364
void SelectPtrauthResign(SDNode *N);
366365

367366
bool trySelectStackSlotTagP(SDNode *N);
@@ -1520,28 +1519,6 @@ extractPtrauthBlendDiscriminators(SDValue Disc, SelectionDAG *DAG) {
15201519
AddrDisc);
15211520
}
15221521

1523-
void AArch64DAGToDAGISel::SelectPtrauthAuth(SDNode *N) {
1524-
SDLoc DL(N);
1525-
// IntrinsicID is operand #0
1526-
SDValue Val = N->getOperand(1);
1527-
SDValue AUTKey = N->getOperand(2);
1528-
SDValue AUTDisc = N->getOperand(3);
1529-
1530-
unsigned AUTKeyC = cast<ConstantSDNode>(AUTKey)->getZExtValue();
1531-
AUTKey = CurDAG->getTargetConstant(AUTKeyC, DL, MVT::i64);
1532-
1533-
SDValue AUTAddrDisc, AUTConstDisc;
1534-
std::tie(AUTConstDisc, AUTAddrDisc) =
1535-
extractPtrauthBlendDiscriminators(AUTDisc, CurDAG);
1536-
1537-
SDValue X16Copy = CurDAG->getCopyToReg(CurDAG->getEntryNode(), DL,
1538-
AArch64::X16, Val, SDValue());
1539-
SDValue Ops[] = {AUTKey, AUTConstDisc, AUTAddrDisc, X16Copy.getValue(1)};
1540-
1541-
SDNode *AUT = CurDAG->getMachineNode(AArch64::AUT, DL, MVT::i64, Ops);
1542-
ReplaceNode(N, AUT);
1543-
}
1544-
15451522
void AArch64DAGToDAGISel::SelectPtrauthResign(SDNode *N) {
15461523
SDLoc DL(N);
15471524
// IntrinsicID is operand #0
@@ -5780,10 +5757,6 @@ void AArch64DAGToDAGISel::Select(SDNode *Node) {
57805757
SelectTagP(Node);
57815758
return;
57825759

5783-
case Intrinsic::ptrauth_auth:
5784-
SelectPtrauthAuth(Node);
5785-
return;
5786-
57875760
case Intrinsic::ptrauth_resign:
57885761
SelectPtrauthResign(Node);
57895762
return;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

Lines changed: 60 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -3103,7 +3103,14 @@ void AArch64TargetLowering::fixupBlendComponents(
31033103
Register AddrDisc = AddrDiscOp.getReg();
31043104
int64_t IntDisc = IntDiscOp.getImm();
31053105

3106-
assert(IntDisc == 0 && "Blend components are already expanded");
3106+
// This function tries to expand an opaque vreg discriminator into address
3107+
// modifier and immediate modifier components. If the immediate modifier
3108+
// operand is already non-zero, no fixup needed.
3109+
if (IntDisc != 0) {
3110+
assert(AddrDisc == AArch64::XZR || AddrDisc == AArch64::NoRegister ||
3111+
MRI.getRegClass(AddrDisc) == AddrDiscRC);
3112+
return;
3113+
}
31073114

31083115
int64_t Offset = 0;
31093116
MachineInstr *MaybeBlend = stripAndAccumulateOffset(MRI, AddrDisc, Offset);
@@ -3130,51 +3137,70 @@ void AArch64TargetLowering::fixupBlendComponents(
31303137
}
31313138

31323139
MachineBasicBlock *
3133-
AArch64TargetLowering::tryRewritingPAC(MachineInstr &MI,
3134-
MachineBasicBlock *BB) const {
3140+
AArch64TargetLowering::EmitAUTxMxN(MachineInstr &MI,
3141+
MachineBasicBlock *BB) const {
3142+
const TargetInstrInfo *TII = Subtarget->getInstrInfo();
3143+
const DebugLoc &DL = MI.getDebugLoc();
3144+
3145+
Register AuthVal = MI.getOperand(0).getReg();
3146+
Register Val = MI.getOperand(2).getReg();
3147+
unsigned Key = MI.getOperand(3).getImm();
3148+
int64_t IntDisc = MI.getOperand(4).getImm();
3149+
Register AddrDisc = MI.getOperand(5).getReg();
3150+
3151+
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), AArch64::X16).addReg(Val);
3152+
BuildMI(*BB, MI, DL, TII->get(AArch64::AUT))
3153+
.addImm(Key)
3154+
.addImm(IntDisc)
3155+
.addReg(AddrDisc);
3156+
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), AuthVal).addReg(AArch64::X16);
3157+
3158+
MI.eraseFromParent();
3159+
return BB;
3160+
}
3161+
3162+
void AArch64TargetLowering::tryRewritingPAC(MachineInstr &MI,
3163+
MachineBasicBlock *BB) const {
31353164
const TargetInstrInfo *TII = Subtarget->getInstrInfo();
31363165
MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
31373166
const DebugLoc &DL = MI.getDebugLoc();
31383167

3168+
Register Val = MI.getOperand(1).getReg();
3169+
unsigned Key = MI.getOperand(2).getImm();
3170+
int64_t IntDisc = MI.getOperand(3).getImm();
3171+
Register AddrDisc = MI.getOperand(4).getReg();
3172+
31393173
// Try to find a known address-setting instruction, accumulating the offset
31403174
// along the way. If no known pattern is found, keep everything as-is.
31413175

31423176
int64_t AddrOffset = 0;
3143-
MachineInstr *AddrDefInstr =
3144-
stripAndAccumulateOffset(MRI, MI.getOperand(1).getReg(), AddrOffset);
3177+
MachineInstr *AddrDefInstr = stripAndAccumulateOffset(MRI, Val, AddrOffset);
31453178
if (!AddrDefInstr)
3146-
return BB;
3179+
return;
31473180

31483181
unsigned NewOpcode;
31493182
if (AddrDefInstr->getOpcode() == AArch64::LOADgotAUTH)
31503183
NewOpcode = AArch64::LOADgotPAC;
31513184
else if (AddrDefInstr->getOpcode() == AArch64::MOVaddr)
31523185
NewOpcode = AArch64::MOVaddrPAC;
31533186
else
3154-
return BB; // Unknown opcode.
3187+
return; // Unknown opcode.
31553188

31563189
MachineOperand &AddrOp = AddrDefInstr->getOperand(1);
31573190
unsigned TargetFlags = AddrOp.getTargetFlags() & ~AArch64II::MO_PAGE;
31583191
const GlobalValue *GV = AddrOp.getGlobal();
31593192
AddrOffset += AddrOp.getOffset();
31603193

3161-
Register DiscReg = isPACWithZeroDisc(MI.getOpcode())
3162-
? AArch64::XZR
3163-
: MI.getOperand(2).getReg();
3164-
3165-
MachineInstr *NewMI = BuildMI(*BB, MI, DL, TII->get(NewOpcode))
3166-
.addGlobalAddress(GV, AddrOffset, TargetFlags)
3167-
.addImm(getKeyForPACOpcode(MI.getOpcode()))
3168-
.addReg(DiscReg)
3169-
.addImm(0);
3170-
fixupBlendComponents(*NewMI, BB, NewMI->getOperand(3), NewMI->getOperand(2),
3171-
&AArch64::GPR64noipRegClass);
3194+
BuildMI(*BB, MI, DL, TII->get(NewOpcode))
3195+
.addGlobalAddress(GV, AddrOffset, TargetFlags)
3196+
.addImm(Key)
3197+
.addReg(AddrDisc)
3198+
.addImm(IntDisc);
31723199

31733200
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), MI.getOperand(0).getReg())
31743201
.addReg(AArch64::X16);
31753202

31763203
MI.removeFromParent();
3177-
return BB;
31783204
}
31793205

31803206
MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
@@ -3276,15 +3302,21 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
32763302
case AArch64::MOVT_TIZ_PSEUDO:
32773303
return EmitZTInstr(MI, BB, AArch64::MOVT_TIZ, /*Op0IsDef=*/true);
32783304

3279-
case AArch64::PACDA:
3280-
case AArch64::PACDB:
3281-
case AArch64::PACIA:
3282-
case AArch64::PACIB:
3283-
case AArch64::PACDZA:
3284-
case AArch64::PACDZB:
3285-
case AArch64::PACIZA:
3286-
case AArch64::PACIZB:
3287-
return tryRewritingPAC(MI, BB);
3305+
case AArch64::PAC:
3306+
fixupBlendComponents(MI, BB, MI.getOperand(3), MI.getOperand(4),
3307+
&AArch64::GPR64noipRegClass);
3308+
tryRewritingPAC(MI, BB);
3309+
return BB;
3310+
case AArch64::AUTxMxN:
3311+
fixupBlendComponents(MI, BB, MI.getOperand(4), MI.getOperand(5),
3312+
&AArch64::GPR64noipRegClass);
3313+
return EmitAUTxMxN(MI, BB);
3314+
case AArch64::AUTPAC:
3315+
fixupBlendComponents(MI, BB, MI.getOperand(1), MI.getOperand(2),
3316+
&AArch64::GPR64noipRegClass);
3317+
fixupBlendComponents(MI, BB, MI.getOperand(4), MI.getOperand(5),
3318+
&AArch64::GPR64noipRegClass);
3319+
return BB;
32883320
}
32893321
}
32903322

llvm/lib/Target/AArch64/AArch64ISelLowering.h

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -187,8 +187,9 @@ class AArch64TargetLowering : public TargetLowering {
187187
MachineOperand &AddrDiscOp,
188188
const TargetRegisterClass *AddrDiscRC) const;
189189

190-
MachineBasicBlock *tryRewritingPAC(MachineInstr &MI,
191-
MachineBasicBlock *BB) const;
190+
MachineBasicBlock *EmitAUTxMxN(MachineInstr &MI, MachineBasicBlock *BB) const;
191+
192+
void tryRewritingPAC(MachineInstr &MI, MachineBasicBlock *BB) const;
192193

193194
MachineBasicBlock *
194195
EmitInstrWithCustomInserter(MachineInstr &MI,

llvm/lib/Target/AArch64/AArch64InstrInfo.td

Lines changed: 37 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1988,6 +1988,8 @@ def PAUTH_PROLOGUE : Pseudo<(outs), (ins), []>, Sched<[]> {
19881988
def PAUTH_EPILOGUE : Pseudo<(outs), (ins), []>, Sched<[]>;
19891989
}
19901990

1991+
def ScratchReg : OptionalDefOperand<OtherVT, (ops GPR64), (ops (i64 zero_reg))>;
1992+
19911993
// These pointer authentication instructions require armv8.3a
19921994
let Predicates = [HasPAuth] in {
19931995

@@ -2018,9 +2020,7 @@ let Predicates = [HasPAuth] in {
20182020
def DZB : SignAuthZero<prefix_z, 0b11, !strconcat(asm, "dzb"), op>;
20192021
}
20202022

2021-
let usesCustomInserter = true in
2022-
defm PAC : SignAuth<0b000, 0b010, "pac", int_ptrauth_sign>;
2023-
2023+
defm PAC : SignAuth<0b000, 0b010, "pac", null_frag>;
20242024
defm AUT : SignAuth<0b001, 0b011, "aut", null_frag>;
20252025

20262026
def XPACI : ClearAuth<0, "xpaci">;
@@ -2126,6 +2126,39 @@ let Predicates = [HasPAuth] in {
21262126
let Uses = [X16];
21272127
}
21282128

2129+
// ScratchReg could be used by
2130+
// https://github.com/llvm/llvm-project/pull/132857.
2131+
def AUTxMxN : Pseudo<(outs GPR64:$AuthVal, GPR64:$ScratchWb),
2132+
(ins GPR64:$Val, i32imm:$Key,
2133+
i64imm:$Disc, GPR64:$AddrDisc, ScratchReg:$Scratch),
2134+
[], "$AuthVal = $Val, $ScratchWb = $Scratch">,
2135+
Sched<[WriteI, ReadI]> {
2136+
let isCodeGenOnly = 1;
2137+
let hasSideEffects = 0;
2138+
let mayStore = 0;
2139+
let mayLoad = 0;
2140+
let Size = 32;
2141+
let Defs = [NZCV];
2142+
let usesCustomInserter = true;
2143+
}
2144+
2145+
def PAC : Pseudo<(outs GPR64:$SignedVal),
2146+
(ins GPR64:$Val, i32imm:$Key, i64imm:$Disc, GPR64noip:$AddrDisc),
2147+
[], "$SignedVal = $Val">, Sched<[WriteI, ReadI]> {
2148+
let isCodeGenOnly = 1;
2149+
let hasSideEffects = 0;
2150+
let mayStore = 0;
2151+
let mayLoad = 0;
2152+
let Size = 12;
2153+
let Defs = [X17];
2154+
let usesCustomInserter = true;
2155+
}
2156+
2157+
def : Pat<(int_ptrauth_auth GPR64:$Val, timm:$Key, GPR64:$AddrDisc),
2158+
(AUTxMxN GPR64:$Val, $Key, 0, GPR64:$AddrDisc)>;
2159+
def : Pat<(int_ptrauth_sign GPR64:$Val, timm:$Key, GPR64noip:$AddrDisc),
2160+
(PAC GPR64:$Val, $Key, 0, GPR64noip:$AddrDisc)>;
2161+
21292162
// AUT and re-PAC a value, using different keys/data.
21302163
// This directly manipulates x16/x17, which are the only registers the OS
21312164
// guarantees are safe to use for sensitive operations.
@@ -2141,6 +2174,7 @@ let Predicates = [HasPAuth] in {
21412174
let Size = 48;
21422175
let Defs = [X16,X17,NZCV];
21432176
let Uses = [X16];
2177+
let usesCustomInserter = true;
21442178
}
21452179

21462180
// Materialize a signed global address, with adrp+add and PAC.

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

Lines changed: 0 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -6727,30 +6727,6 @@ bool AArch64InstructionSelector::selectIntrinsic(MachineInstr &I,
67276727
I.eraseFromParent();
67286728
return true;
67296729
}
6730-
case Intrinsic::ptrauth_auth: {
6731-
Register DstReg = I.getOperand(0).getReg();
6732-
Register ValReg = I.getOperand(2).getReg();
6733-
uint64_t AUTKey = I.getOperand(3).getImm();
6734-
Register AUTDisc = I.getOperand(4).getReg();
6735-
6736-
Register AUTAddrDisc = AUTDisc;
6737-
uint16_t AUTConstDiscC = 0;
6738-
std::tie(AUTConstDiscC, AUTAddrDisc) =
6739-
extractPtrauthBlendDiscriminators(AUTDisc, MRI);
6740-
6741-
MIB.buildCopy({AArch64::X16}, {ValReg});
6742-
MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X17}, {});
6743-
MIB.buildInstr(AArch64::AUT)
6744-
.addImm(AUTKey)
6745-
.addImm(AUTConstDiscC)
6746-
.addUse(AUTAddrDisc)
6747-
.constrainAllUses(TII, TRI, RBI);
6748-
MIB.buildCopy({DstReg}, Register(AArch64::X16));
6749-
6750-
RBI.constrainGenericRegister(DstReg, AArch64::GPR64RegClass, MRI);
6751-
I.eraseFromParent();
6752-
return true;
6753-
}
67546730
case Intrinsic::frameaddress:
67556731
case Intrinsic::returnaddress: {
67566732
MachineFunction &MF = *I.getParent()->getParent();

llvm/test/CodeGen/AArch64/GlobalISel/ptrauth-constant-in-code.ll

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ define void @store_signed_const_local(ptr %dest) {
106106
; ISEL-MIR: body:
107107
; ISEL-MIR: %0:gpr64common = COPY $x0
108108
; ISEL-MIR-NEXT: %10:gpr64common = MOVaddr target-flags(aarch64-page) @const_table_local + 8, target-flags(aarch64-pageoff, aarch64-nc) @const_table_local + 8
109-
; ISEL-MIR-NEXT: %2:gpr64common = MOVKXi %0, 1234, 48
109+
; ISEL-MIR-NEXT: %2:gpr64noip = MOVKXi %0, 1234, 48
110110
; ISEL-MIR-NEXT: %15:gpr64noip = COPY %0
111111
; ISEL-MIR-NEXT: MOVaddrPAC @const_table_local + 8, 2, %15, 1234, implicit-def $x16, implicit-def $x17
112112
; ISEL-MIR-NEXT: %4:gpr64 = COPY $x16
@@ -140,7 +140,7 @@ define void @store_signed_const_got(ptr %dest) {
140140
; ISEL-MIR-ELF: %0:gpr64common = COPY $x0
141141
; ISEL-MIR-ELF-NEXT: %7:gpr64common = LOADgotAUTH target-flags(aarch64-got) @const_table_got
142142
; ISEL-MIR-ELF-NEXT: %6:gpr64common = ADDXri %7, 8, 0
143-
; ISEL-MIR-ELF-NEXT: %2:gpr64common = MOVKXi %0, 1234, 48
143+
; ISEL-MIR-ELF-NEXT: %2:gpr64noip = MOVKXi %0, 1234, 48
144144
; ISEL-MIR-ELF-NEXT: %12:gpr64noip = COPY %0
145145
; ISEL-MIR-ELF-NEXT: LOADgotPAC target-flags(aarch64-got) @const_table_got + 8, 2, %12, 1234, implicit-def $x16, implicit-def $x17, implicit-def $nzcv
146146
; ISEL-MIR-ELF-NEXT: %4:gpr64 = COPY $x16
@@ -179,21 +179,21 @@ define void @store_signed_arg(ptr %dest, ptr %p) {
179179
; ISEL-MIR: body:
180180
; ISEL-MIR: %0:gpr64common = COPY $x0
181181
; ISEL-MIR-NEXT: %1:gpr64common = COPY $x1
182-
; ISEL-MIR-NEXT: %3:gpr64common = MOVKXi %0, 1234, 48
182+
; ISEL-MIR-NEXT: %3:gpr64noip = MOVKXi %0, 1234, 48
183183
; ISEL-MIR-NEXT: %6:gpr64common = ADDXri %1, 8, 0
184-
; Check that no implicit defs are added to PACDA instruction.
185-
; ISEL-MIR-NEXT: %8:gpr64 = PACDA %6, %3{{$}}
184+
; ISEL-MIR-NEXT: %12:gpr64noip = COPY %0
185+
; ISEL-MIR-NEXT: %8:gpr64 = PAC %6, 2, 1234, %12, implicit-def dead $x17
186186
; ISEL-MIR-NEXT: %10:gpr64 = COPY %8
187187
; ISEL-MIR-NEXT: STRXui %10, %0, 0 :: (store (p0) into %ir.dest)
188188
; ISEL-MIR-NEXT: RET_ReallyLR
189189
;
190190
; ISEL-ASM-LABEL: store_signed_arg:
191191
; ISEL-ASM-NEXT: .cfi_startproc
192-
; ISEL-ASM-NEXT: mov x8, x0
193-
; ISEL-ASM-NEXT: add x9, x1, #8
194-
; ISEL-ASM-NEXT: movk x8, #1234, lsl #48
195-
; ISEL-ASM-NEXT: pacda x9, x8
196-
; ISEL-ASM-NEXT: str x9, [x0]
192+
; ISEL-ASM-NEXT: add x8, x1, #8
193+
; ISEL-ASM-NEXT: mov x17, x0
194+
; ISEL-ASM-NEXT: movk x17, #1234, lsl #48
195+
; ISEL-ASM-NEXT: pacda x8, x17
196+
; ISEL-ASM-NEXT: str x8, [x0]
197197
; ISEL-ASM-NEXT: ret
198198
%dest.i = ptrtoint ptr %dest to i64
199199
%discr = call i64 @llvm.ptrauth.blend(i64 %dest.i, i64 1234)

0 commit comments

Comments
 (0)