Skip to content

Commit 919cbdf

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 ced57f0 commit 919cbdf

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
@@ -170,6 +170,9 @@ class AArch64AsmPrinter : public AsmPrinter {
170170
// Emit the sequence for AUT or AUTPAC.
171171
void emitPtrauthAuthResign(const MachineInstr *MI);
172172

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

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

2896+
case AArch64::PAC:
2897+
emitPtrauthSign(MI);
2898+
return;
2899+
28692900
case AArch64::LOADauthptrstatic:
28702901
LowerLOADauthptrstatic(*MI);
28712902
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
@@ -5645,10 +5622,6 @@ void AArch64DAGToDAGISel::Select(SDNode *Node) {
56455622
SelectTagP(Node);
56465623
return;
56475624

5648-
case Intrinsic::ptrauth_auth:
5649-
SelectPtrauthAuth(Node);
5650-
return;
5651-
56525625
case Intrinsic::ptrauth_resign:
56535626
SelectPtrauthResign(Node);
56545627
return;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

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

3397-
assert(IntDisc == 0 && "Blend components are already expanded");
3397+
// This function tries to expand an opaque vreg discriminator into address
3398+
// modifier and immediate modifier components. If the immediate modifier
3399+
// operand is already non-zero, no fixup needed.
3400+
if (IntDisc != 0) {
3401+
assert(AddrDisc == AArch64::XZR || AddrDisc == AArch64::NoRegister ||
3402+
MRI.getRegClass(AddrDisc) == AddrDiscRC);
3403+
return;
3404+
}
33983405

33993406
int64_t Offset = 0;
34003407
MachineInstr *MaybeBlend = stripAndAccumulateOffset(MRI, AddrDisc, Offset);
@@ -3421,51 +3428,70 @@ void AArch64TargetLowering::fixupBlendComponents(
34213428
}
34223429

34233430
MachineBasicBlock *
3424-
AArch64TargetLowering::tryRewritingPAC(MachineInstr &MI,
3425-
MachineBasicBlock *BB) const {
3431+
AArch64TargetLowering::EmitAUTxMxN(MachineInstr &MI,
3432+
MachineBasicBlock *BB) const {
3433+
const TargetInstrInfo *TII = Subtarget->getInstrInfo();
3434+
const DebugLoc &DL = MI.getDebugLoc();
3435+
3436+
Register AuthVal = MI.getOperand(0).getReg();
3437+
Register Val = MI.getOperand(2).getReg();
3438+
unsigned Key = MI.getOperand(3).getImm();
3439+
int64_t IntDisc = MI.getOperand(4).getImm();
3440+
Register AddrDisc = MI.getOperand(5).getReg();
3441+
3442+
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), AArch64::X16).addReg(Val);
3443+
BuildMI(*BB, MI, DL, TII->get(AArch64::AUT))
3444+
.addImm(Key)
3445+
.addImm(IntDisc)
3446+
.addReg(AddrDisc);
3447+
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), AuthVal).addReg(AArch64::X16);
3448+
3449+
MI.eraseFromParent();
3450+
return BB;
3451+
}
3452+
3453+
void AArch64TargetLowering::tryRewritingPAC(MachineInstr &MI,
3454+
MachineBasicBlock *BB) const {
34263455
const TargetInstrInfo *TII = Subtarget->getInstrInfo();
34273456
MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
34283457
const DebugLoc &DL = MI.getDebugLoc();
34293458

3459+
Register Val = MI.getOperand(1).getReg();
3460+
unsigned Key = MI.getOperand(2).getImm();
3461+
int64_t IntDisc = MI.getOperand(3).getImm();
3462+
Register AddrDisc = MI.getOperand(4).getReg();
3463+
34303464
// Try to find a known address-setting instruction, accumulating the offset
34313465
// along the way. If no known pattern is found, keep everything as-is.
34323466

34333467
int64_t AddrOffset = 0;
3434-
MachineInstr *AddrDefInstr =
3435-
stripAndAccumulateOffset(MRI, MI.getOperand(1).getReg(), AddrOffset);
3468+
MachineInstr *AddrDefInstr = stripAndAccumulateOffset(MRI, Val, AddrOffset);
34363469
if (!AddrDefInstr)
3437-
return BB;
3470+
return;
34383471

34393472
unsigned NewOpcode;
34403473
if (AddrDefInstr->getOpcode() == AArch64::LOADgotAUTH)
34413474
NewOpcode = AArch64::LOADgotPAC;
34423475
else if (AddrDefInstr->getOpcode() == AArch64::MOVaddr)
34433476
NewOpcode = AArch64::MOVaddrPAC;
34443477
else
3445-
return BB; // Unknown opcode.
3478+
return; // Unknown opcode.
34463479

34473480
MachineOperand &AddrOp = AddrDefInstr->getOperand(1);
34483481
unsigned TargetFlags = AddrOp.getTargetFlags() & ~AArch64II::MO_PAGE;
34493482
const GlobalValue *GV = AddrOp.getGlobal();
34503483
AddrOffset += AddrOp.getOffset();
34513484

3452-
Register DiscReg = isPACWithZeroDisc(MI.getOpcode())
3453-
? AArch64::XZR
3454-
: MI.getOperand(2).getReg();
3455-
3456-
MachineInstr *NewMI = BuildMI(*BB, MI, DL, TII->get(NewOpcode))
3457-
.addGlobalAddress(GV, AddrOffset, TargetFlags)
3458-
.addImm(getKeyForPACOpcode(MI.getOpcode()))
3459-
.addReg(DiscReg)
3460-
.addImm(0);
3461-
fixupBlendComponents(*NewMI, BB, NewMI->getOperand(3), NewMI->getOperand(2),
3462-
&AArch64::GPR64noipRegClass);
3485+
BuildMI(*BB, MI, DL, TII->get(NewOpcode))
3486+
.addGlobalAddress(GV, AddrOffset, TargetFlags)
3487+
.addImm(Key)
3488+
.addReg(AddrDisc)
3489+
.addImm(IntDisc);
34633490

34643491
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), MI.getOperand(0).getReg())
34653492
.addReg(AArch64::X16);
34663493

34673494
MI.removeFromParent();
3468-
return BB;
34693495
}
34703496

34713497
MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
@@ -3567,15 +3593,21 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
35673593
case AArch64::MOVT_TIZ_PSEUDO:
35683594
return EmitZTInstr(MI, BB, AArch64::MOVT_TIZ, /*Op0IsDef=*/true);
35693595

3570-
case AArch64::PACDA:
3571-
case AArch64::PACDB:
3572-
case AArch64::PACIA:
3573-
case AArch64::PACIB:
3574-
case AArch64::PACDZA:
3575-
case AArch64::PACDZB:
3576-
case AArch64::PACIZA:
3577-
case AArch64::PACIZB:
3578-
return tryRewritingPAC(MI, BB);
3596+
case AArch64::PAC:
3597+
fixupBlendComponents(MI, BB, MI.getOperand(3), MI.getOperand(4),
3598+
&AArch64::GPR64noipRegClass);
3599+
tryRewritingPAC(MI, BB);
3600+
return BB;
3601+
case AArch64::AUTxMxN:
3602+
fixupBlendComponents(MI, BB, MI.getOperand(4), MI.getOperand(5),
3603+
&AArch64::GPR64noipRegClass);
3604+
return EmitAUTxMxN(MI, BB);
3605+
case AArch64::AUTPAC:
3606+
fixupBlendComponents(MI, BB, MI.getOperand(1), MI.getOperand(2),
3607+
&AArch64::GPR64noipRegClass);
3608+
fixupBlendComponents(MI, BB, MI.getOperand(4), MI.getOperand(5),
3609+
&AArch64::GPR64noipRegClass);
3610+
return BB;
35793611
}
35803612
}
35813613

llvm/lib/Target/AArch64/AArch64ISelLowering.h

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

688-
MachineBasicBlock *tryRewritingPAC(MachineInstr &MI,
689-
MachineBasicBlock *BB) const;
688+
MachineBasicBlock *EmitAUTxMxN(MachineInstr &MI, MachineBasicBlock *BB) const;
689+
690+
void tryRewritingPAC(MachineInstr &MI, MachineBasicBlock *BB) const;
690691

691692
MachineBasicBlock *
692693
EmitInstrWithCustomInserter(MachineInstr &MI,

llvm/lib/Target/AArch64/AArch64InstrInfo.td

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

1818+
def ScratchReg : OptionalDefOperand<OtherVT, (ops GPR64), (ops (i64 zero_reg))>;
1819+
18181820
// These pointer authentication instructions require armv8.3a
18191821
let Predicates = [HasPAuth] in {
18201822

@@ -1845,9 +1847,7 @@ let Predicates = [HasPAuth] in {
18451847
def DZB : SignAuthZero<prefix_z, 0b11, !strconcat(asm, "dzb"), op>;
18461848
}
18471849

1848-
let usesCustomInserter = true in
1849-
defm PAC : SignAuth<0b000, 0b010, "pac", int_ptrauth_sign>;
1850-
1850+
defm PAC : SignAuth<0b000, 0b010, "pac", null_frag>;
18511851
defm AUT : SignAuth<0b001, 0b011, "aut", null_frag>;
18521852

18531853
def XPACI : ClearAuth<0, "xpaci">;
@@ -1953,6 +1953,39 @@ let Predicates = [HasPAuth] in {
19531953
let Uses = [X16];
19541954
}
19551955

1956+
// ScratchReg could be used by
1957+
// https://github.com/llvm/llvm-project/pull/132857.
1958+
def AUTxMxN : Pseudo<(outs GPR64:$AuthVal, GPR64:$ScratchWb),
1959+
(ins GPR64:$Val, i32imm:$Key,
1960+
i64imm:$Disc, GPR64:$AddrDisc, ScratchReg:$Scratch),
1961+
[], "$AuthVal = $Val, $ScratchWb = $Scratch">,
1962+
Sched<[WriteI, ReadI]> {
1963+
let isCodeGenOnly = 1;
1964+
let hasSideEffects = 0;
1965+
let mayStore = 0;
1966+
let mayLoad = 0;
1967+
let Size = 32;
1968+
let Defs = [NZCV];
1969+
let usesCustomInserter = true;
1970+
}
1971+
1972+
def PAC : Pseudo<(outs GPR64:$SignedVal),
1973+
(ins GPR64:$Val, i32imm:$Key, i64imm:$Disc, GPR64noip:$AddrDisc),
1974+
[], "$SignedVal = $Val">, Sched<[WriteI, ReadI]> {
1975+
let isCodeGenOnly = 1;
1976+
let hasSideEffects = 0;
1977+
let mayStore = 0;
1978+
let mayLoad = 0;
1979+
let Size = 12;
1980+
let Defs = [X17];
1981+
let usesCustomInserter = true;
1982+
}
1983+
1984+
def : Pat<(int_ptrauth_auth GPR64:$Val, timm:$Key, GPR64:$AddrDisc),
1985+
(AUTxMxN GPR64:$Val, $Key, 0, GPR64:$AddrDisc)>;
1986+
def : Pat<(int_ptrauth_sign GPR64:$Val, timm:$Key, GPR64noip:$AddrDisc),
1987+
(PAC GPR64:$Val, $Key, 0, GPR64noip:$AddrDisc)>;
1988+
19561989
// AUT and re-PAC a value, using different keys/data.
19571990
// This directly manipulates x16/x17, which are the only registers the OS
19581991
// guarantees are safe to use for sensitive operations.
@@ -1968,6 +2001,7 @@ let Predicates = [HasPAuth] in {
19682001
let Size = 48;
19692002
let Defs = [X16,X17,NZCV];
19702003
let Uses = [X16];
2004+
let usesCustomInserter = true;
19712005
}
19722006

19732007
// 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
@@ -6726,30 +6726,6 @@ bool AArch64InstructionSelector::selectIntrinsic(MachineInstr &I,
67266726
I.eraseFromParent();
67276727
return true;
67286728
}
6729-
case Intrinsic::ptrauth_auth: {
6730-
Register DstReg = I.getOperand(0).getReg();
6731-
Register ValReg = I.getOperand(2).getReg();
6732-
uint64_t AUTKey = I.getOperand(3).getImm();
6733-
Register AUTDisc = I.getOperand(4).getReg();
6734-
6735-
Register AUTAddrDisc = AUTDisc;
6736-
uint16_t AUTConstDiscC = 0;
6737-
std::tie(AUTConstDiscC, AUTAddrDisc) =
6738-
extractPtrauthBlendDiscriminators(AUTDisc, MRI);
6739-
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));
6748-
6749-
RBI.constrainGenericRegister(DstReg, AArch64::GPR64RegClass, MRI);
6750-
I.eraseFromParent();
6751-
return true;
6752-
}
67536729
case Intrinsic::frameaddress:
67546730
case Intrinsic::returnaddress: {
67556731
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)