Skip to content

[AArch64][PAC] Rework discriminator analysis in AUT and AUTPAC #146489

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: users/atrosinenko/pauth-imm-modifier-sign
Choose a base branch
from
+256 −69
Open

[AArch64][PAC] Rework discriminator analysis in AUT and AUTPAC #146489

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 5 additions & 46 deletions llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1487,39 +1487,6 @@ void AArch64DAGToDAGISel::SelectTable(SDNode *N, unsigned NumVecs, unsigned Opc,
ReplaceNode(N, CurDAG->getMachineNode(Opc, dl, VT, Ops));
}

static std::tuple<SDValue, SDValue>
extractPtrauthBlendDiscriminators(SDValue Disc, SelectionDAG *DAG) {
SDLoc DL(Disc);
SDValue AddrDisc;
SDValue ConstDisc;

// If this is a blend, remember the constant and address discriminators.
// Otherwise, it's either a constant discriminator, or a non-blended
// address discriminator.
if (Disc->getOpcode() == ISD::INTRINSIC_WO_CHAIN &&
Disc->getConstantOperandVal(0) == Intrinsic::ptrauth_blend) {
AddrDisc = Disc->getOperand(1);
ConstDisc = Disc->getOperand(2);
} else {
ConstDisc = Disc;
}

// If the constant discriminator (either the blend RHS, or the entire
// discriminator value) isn't a 16-bit constant, bail out, and let the
// discriminator be computed separately.
auto *ConstDiscN = dyn_cast<ConstantSDNode>(ConstDisc);
if (!ConstDiscN || !isUInt<16>(ConstDiscN->getZExtValue()))
return std::make_tuple(DAG->getTargetConstant(0, DL, MVT::i64), Disc);

// If there's no address discriminator, use XZR directly.
if (!AddrDisc)
AddrDisc = DAG->getRegister(AArch64::XZR, MVT::i64);

return std::make_tuple(
DAG->getTargetConstant(ConstDiscN->getZExtValue(), DL, MVT::i64),
AddrDisc);
}

void AArch64DAGToDAGISel::SelectPtrauthAuth(SDNode *N) {
SDLoc DL(N);
// IntrinsicID is operand #0
Expand All @@ -1530,13 +1497,11 @@ void AArch64DAGToDAGISel::SelectPtrauthAuth(SDNode *N) {
unsigned AUTKeyC = cast<ConstantSDNode>(AUTKey)->getZExtValue();
AUTKey = CurDAG->getTargetConstant(AUTKeyC, DL, MVT::i64);

SDValue AUTAddrDisc, AUTConstDisc;
std::tie(AUTConstDisc, AUTAddrDisc) =
extractPtrauthBlendDiscriminators(AUTDisc, CurDAG);
SDValue Zero = CurDAG->getTargetConstant(0, DL, MVT::i64);

SDValue X16Copy = CurDAG->getCopyToReg(CurDAG->getEntryNode(), DL,
AArch64::X16, Val, SDValue());
SDValue Ops[] = {AUTKey, AUTConstDisc, AUTAddrDisc, X16Copy.getValue(1)};
SDValue Ops[] = {AUTKey, Zero, AUTDisc, X16Copy.getValue(1)};

SDNode *AUT = CurDAG->getMachineNode(AArch64::AUT, DL, MVT::i64, Ops);
ReplaceNode(N, AUT);
Expand All @@ -1557,19 +1522,13 @@ void AArch64DAGToDAGISel::SelectPtrauthResign(SDNode *N) {
AUTKey = CurDAG->getTargetConstant(AUTKeyC, DL, MVT::i64);
PACKey = CurDAG->getTargetConstant(PACKeyC, DL, MVT::i64);

SDValue AUTAddrDisc, AUTConstDisc;
std::tie(AUTConstDisc, AUTAddrDisc) =
extractPtrauthBlendDiscriminators(AUTDisc, CurDAG);

SDValue PACAddrDisc, PACConstDisc;
std::tie(PACConstDisc, PACAddrDisc) =
extractPtrauthBlendDiscriminators(PACDisc, CurDAG);
SDValue Zero = CurDAG->getTargetConstant(0, DL, MVT::i64);

SDValue X16Copy = CurDAG->getCopyToReg(CurDAG->getEntryNode(), DL,
AArch64::X16, Val, SDValue());

SDValue Ops[] = {AUTKey, AUTConstDisc, AUTAddrDisc, PACKey,
PACConstDisc, PACAddrDisc, X16Copy.getValue(1)};
SDValue Ops[] = {
AUTKey, Zero, AUTDisc, PACKey, Zero, PACDisc, X16Copy.getValue(1)};

SDNode *AUTPAC = CurDAG->getMachineNode(AArch64::AUTPAC, DL, MVT::i64, Ops);
ReplaceNode(N, AUTPAC);
Expand Down
10 changes: 10 additions & 0 deletions llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3242,10 +3242,20 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
case AArch64::MOVT_TIZ_PSEUDO:
return EmitZTInstr(MI, BB, AArch64::MOVT_TIZ, /*Op0IsDef=*/true);

case AArch64::AUT:
fixupBlendComponents(MI, BB, MI.getOperand(1), MI.getOperand(2),
&AArch64::GPR64noipRegClass);
return BB;
case AArch64::PAC:
fixupBlendComponents(MI, BB, MI.getOperand(3), MI.getOperand(4),
&AArch64::GPR64noipRegClass);
return BB;
case AArch64::AUTPAC:
fixupBlendComponents(MI, BB, MI.getOperand(1), MI.getOperand(2),
&AArch64::GPR64noipRegClass);
fixupBlendComponents(MI, BB, MI.getOperand(4), MI.getOperand(5),
&AArch64::GPR64noipRegClass);
return BB;
}
}

Expand Down
2 changes: 2 additions & 0 deletions llvm/lib/Target/AArch64/AArch64InstrInfo.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2134,6 +2134,7 @@ let Predicates = [HasPAuth] in {
let Size = 32;
let Defs = [X16,X17,NZCV];
let Uses = [X16];
let usesCustomInserter = 1;
}

// PAC pseudo instruction. Is AsmPrinter, it is expanded into an actual PAC*
Expand Down Expand Up @@ -2170,6 +2171,7 @@ let Predicates = [HasPAuth] in {
let Size = 48;
let Defs = [X16,X17,NZCV];
let Uses = [X16];
let usesCustomInserter = 1;
}

// Materialize a signed global address, with adrp+add and PAC.
Expand Down
27 changes: 6 additions & 21 deletions llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6725,25 +6725,15 @@ bool AArch64InstructionSelector::selectIntrinsic(MachineInstr &I,
uint64_t PACKey = I.getOperand(5).getImm();
Register PACDisc = I.getOperand(6).getReg();

Register AUTAddrDisc = AUTDisc;
uint16_t AUTConstDiscC = 0;
std::tie(AUTConstDiscC, AUTAddrDisc) =
extractPtrauthBlendDiscriminators(AUTDisc, MRI);

Register PACAddrDisc = PACDisc;
uint16_t PACConstDiscC = 0;
std::tie(PACConstDiscC, PACAddrDisc) =
extractPtrauthBlendDiscriminators(PACDisc, MRI);

MIB.buildCopy({AArch64::X16}, {ValReg});
MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X17}, {});
MIB.buildInstr(AArch64::AUTPAC)
.addImm(AUTKey)
.addImm(AUTConstDiscC)
.addUse(AUTAddrDisc)
.addImm(0)
.addUse(AUTDisc)
.addImm(PACKey)
.addImm(PACConstDiscC)
.addUse(PACAddrDisc)
.addImm(0)
.addUse(PACDisc)
.constrainAllUses(TII, TRI, RBI);
MIB.buildCopy({DstReg}, Register(AArch64::X16));

Expand All @@ -6757,17 +6747,12 @@ bool AArch64InstructionSelector::selectIntrinsic(MachineInstr &I,
uint64_t AUTKey = I.getOperand(3).getImm();
Register AUTDisc = I.getOperand(4).getReg();

Register AUTAddrDisc = AUTDisc;
uint16_t AUTConstDiscC = 0;
std::tie(AUTConstDiscC, AUTAddrDisc) =
extractPtrauthBlendDiscriminators(AUTDisc, MRI);

MIB.buildCopy({AArch64::X16}, {ValReg});
MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X17}, {});
MIB.buildInstr(AArch64::AUT)
.addImm(AUTKey)
.addImm(AUTConstDiscC)
.addUse(AUTAddrDisc)
.addImm(0)
.addUse(AUTDisc)
.constrainAllUses(TII, TRI, RBI);
MIB.buildCopy({DstReg}, Register(AArch64::X16));

Expand Down
Loading
Loading