Skip to content

Commit 1620290

Browse files
atrosinenkoatrosinenko
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 ec918ef commit 1620290

10 files changed

+421
-125
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: 75 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -3384,78 +3384,108 @@ static MachineInstr *stripAndAccumulateOffset(const MachineRegisterInfo &MRI,
33843384
return nullptr;
33853385
}
33863386

3387-
static std::pair<Register, unsigned>
3388-
detectBlendComponents(const MachineRegisterInfo &MRI, Register Reg) {
3387+
void AArch64TargetLowering::fixupBlendComponents(
3388+
MachineInstr &It, MachineBasicBlock *BB, MachineOperand &IntDiscOp,
3389+
MachineOperand &AddrDiscOp) const {
3390+
const TargetInstrInfo *TII = Subtarget->getInstrInfo();
3391+
MachineRegisterInfo &MRI = It.getMF()->getRegInfo();
3392+
const DebugLoc &DL = It.getDebugLoc();
3393+
3394+
Register AddrDisc = AddrDiscOp.getReg();
3395+
int64_t IntDisc = IntDiscOp.getImm();
3396+
// If the immediate modifier is non-zero, no fixup needed.
3397+
if (IntDisc != 0)
3398+
return;
3399+
33893400
int64_t Offset = 0;
3390-
MachineInstr *MaybeBlend = stripAndAccumulateOffset(MRI, Reg, Offset);
3391-
// This should be a plain copy, without adding any offset.
3392-
if (!MaybeBlend || Offset != 0)
3393-
return std::make_pair(Reg, 0);
3401+
MachineInstr *MaybeBlend = stripAndAccumulateOffset(MRI, AddrDisc, Offset);
33943402

3403+
// This should be a plain copy, without adding any offset.
33953404
// Detect blend(addr, imm) which is lowered as MOVK addr, #imm, 48.
3396-
if (MaybeBlend->getOpcode() != AArch64::MOVKXi ||
3397-
MaybeBlend->getOperand(3).getImm() != 48)
3398-
return std::make_pair(Reg, 0);
3405+
if (MaybeBlend && Offset == 0 && MaybeBlend->getOpcode() == AArch64::MOVKXi &&
3406+
MaybeBlend->getOperand(3).getImm() == 48) {
3407+
AddrDisc = MaybeBlend->getOperand(1).getReg();
3408+
IntDisc = MaybeBlend->getOperand(2).getImm();
3409+
}
3410+
3411+
if (AddrDisc == AArch64::NoRegister)
3412+
AddrDisc = AArch64::XZR;
3413+
3414+
if (AddrDisc != AArch64::XZR) {
3415+
Register TmpReg = MRI.createVirtualRegister(&AArch64::GPR64noipRegClass);
3416+
BuildMI(*BB, It, DL, TII->get(AArch64::COPY), TmpReg).addReg(AddrDisc);
3417+
AddrDisc = TmpReg;
3418+
}
33993419

3400-
return std::make_pair(MaybeBlend->getOperand(1).getReg(),
3401-
MaybeBlend->getOperand(2).getImm());
3420+
AddrDiscOp.setReg(AddrDisc);
3421+
IntDiscOp.setImm(IntDisc);
34023422
}
34033423

34043424
MachineBasicBlock *
3405-
AArch64TargetLowering::tryRewritingPAC(MachineInstr &MI,
3406-
MachineBasicBlock *BB) const {
3425+
AArch64TargetLowering::EmitAUTxMxN(MachineInstr &MI,
3426+
MachineBasicBlock *BB) const {
3427+
const TargetInstrInfo *TII = Subtarget->getInstrInfo();
3428+
const DebugLoc &DL = MI.getDebugLoc();
3429+
3430+
Register AuthVal = MI.getOperand(0).getReg();
3431+
Register Val = MI.getOperand(2).getReg();
3432+
unsigned Key = MI.getOperand(3).getImm();
3433+
int64_t IntDisc = MI.getOperand(4).getImm();
3434+
Register AddrDisc = MI.getOperand(5).getReg();
3435+
3436+
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), AArch64::X16).addReg(Val);
3437+
BuildMI(*BB, MI, DL, TII->get(AArch64::AUT))
3438+
.addImm(Key)
3439+
.addImm(IntDisc)
3440+
.addReg(AddrDisc);
3441+
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), AuthVal).addReg(AArch64::X16);
3442+
3443+
MI.eraseFromParent();
3444+
return BB;
3445+
}
3446+
3447+
void AArch64TargetLowering::tryRewritingPAC(MachineInstr &MI,
3448+
MachineBasicBlock *BB) const {
34073449
const TargetInstrInfo *TII = Subtarget->getInstrInfo();
34083450
MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
34093451
const DebugLoc &DL = MI.getDebugLoc();
34103452

3453+
Register Val = MI.getOperand(1).getReg();
3454+
unsigned Key = MI.getOperand(2).getImm();
3455+
int64_t IntDisc = MI.getOperand(3).getImm();
3456+
Register AddrDisc = MI.getOperand(4).getReg();
3457+
34113458
// Try to find a known address-setting instruction, accumulating the offset
34123459
// along the way. If no known pattern is found, keep everything as-is.
34133460

34143461
int64_t AddrOffset = 0;
3415-
MachineInstr *AddrDefInstr =
3416-
stripAndAccumulateOffset(MRI, MI.getOperand(1).getReg(), AddrOffset);
3462+
MachineInstr *AddrDefInstr = stripAndAccumulateOffset(MRI, Val, AddrOffset);
34173463
if (!AddrDefInstr)
3418-
return BB;
3464+
return;
34193465

34203466
unsigned NewOpcode;
34213467
if (AddrDefInstr->getOpcode() == AArch64::LOADgotAUTH)
34223468
NewOpcode = AArch64::LOADgotPAC;
34233469
else if (AddrDefInstr->getOpcode() == AArch64::MOVaddr)
34243470
NewOpcode = AArch64::MOVaddrPAC;
34253471
else
3426-
return BB; // Unknown opcode.
3472+
return; // Unknown opcode.
34273473

34283474
MachineOperand &AddrOp = AddrDefInstr->getOperand(1);
34293475
unsigned TargetFlags = AddrOp.getTargetFlags() & ~AArch64II::MO_PAGE;
34303476
const GlobalValue *GV = AddrOp.getGlobal();
34313477
AddrOffset += AddrOp.getOffset();
34323478

3433-
// Analyze the discriminator operand.
3434-
Register OriginalDisc = isPACWithZeroDisc(MI.getOpcode())
3435-
? AArch64::XZR
3436-
: MI.getOperand(2).getReg();
3437-
auto [AddrDisc, IntDisc] = detectBlendComponents(MRI, OriginalDisc);
3438-
3439-
// MOVaddrPAC and LOADgotPAC pseudos are expanded so that they use X16/X17
3440-
// internally, thus their restrictions on the register class of $AddrDisc
3441-
// operand are stricter than those of MOVKXi and PAC* instructions.
3442-
if (AddrDisc != AArch64::XZR) {
3443-
Register TmpReg = MRI.createVirtualRegister(&AArch64::GPR64noipRegClass);
3444-
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), TmpReg).addReg(AddrDisc);
3445-
AddrDisc = TmpReg;
3446-
}
3447-
34483479
BuildMI(*BB, MI, DL, TII->get(NewOpcode))
34493480
.addGlobalAddress(GV, AddrOffset, TargetFlags)
3450-
.addImm(getKeyForPACOpcode(MI.getOpcode()))
3481+
.addImm(Key)
34513482
.addReg(AddrDisc)
34523483
.addImm(IntDisc);
34533484

34543485
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), MI.getOperand(0).getReg())
34553486
.addReg(AArch64::X16);
34563487

34573488
MI.removeFromParent();
3458-
return BB;
34593489
}
34603490

34613491
MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
@@ -3557,15 +3587,17 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
35573587
case AArch64::MOVT_TIZ_PSEUDO:
35583588
return EmitZTInstr(MI, BB, AArch64::MOVT_TIZ, /*Op0IsDef=*/true);
35593589

3560-
case AArch64::PACDA:
3561-
case AArch64::PACDB:
3562-
case AArch64::PACIA:
3563-
case AArch64::PACIB:
3564-
case AArch64::PACDZA:
3565-
case AArch64::PACDZB:
3566-
case AArch64::PACIZA:
3567-
case AArch64::PACIZB:
3568-
return tryRewritingPAC(MI, BB);
3590+
case AArch64::PAC:
3591+
fixupBlendComponents(MI, BB, MI.getOperand(3), MI.getOperand(4));
3592+
tryRewritingPAC(MI, BB);
3593+
return BB;
3594+
case AArch64::AUTxMxN:
3595+
fixupBlendComponents(MI, BB, MI.getOperand(4), MI.getOperand(5));
3596+
return EmitAUTxMxN(MI, BB);
3597+
case AArch64::AUTPAC:
3598+
fixupBlendComponents(MI, BB, MI.getOperand(1), MI.getOperand(2));
3599+
fixupBlendComponents(MI, BB, MI.getOperand(4), MI.getOperand(5));
3600+
return BB;
35693601
}
35703602
}
35713603

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -680,8 +680,13 @@ class AArch64TargetLowering : public TargetLowering {
680680
MachineBasicBlock *EmitGetSMESaveSize(MachineInstr &MI,
681681
MachineBasicBlock *BB) const;
682682

683-
MachineBasicBlock *tryRewritingPAC(MachineInstr &MI,
684-
MachineBasicBlock *BB) const;
683+
void fixupBlendComponents(MachineInstr &It, MachineBasicBlock *BB,
684+
MachineOperand &IntDiscOp,
685+
MachineOperand &AddrDiscOp) const;
686+
687+
MachineBasicBlock *EmitAUTxMxN(MachineInstr &MI, MachineBasicBlock *BB) const;
688+
689+
void tryRewritingPAC(MachineInstr &MI, MachineBasicBlock *BB) const;
685690

686691
MachineBasicBlock *
687692
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();

0 commit comments

Comments
 (0)