Skip to content

Commit eb02ee4

Browse files
atrosinenkoatrosinenko
committed
[AArch64] Move PAuth codegen down the machine pipeline
To simplify handling PAuth in the machine outliner, introduce a separate AArch64PointerAuth pass that is executed after both Prologue/Epilogue Inserter and Machine Outliner passes. After moving to AArch64PointerAuth, signLR and authenticateLR are not used outside of their class anymore, so make them private and simplify accordingly. The new pass is added via AArch64PassConfig::addPostBBSections(), so that it can change the code size before branch relaxation occurs. AArch64BranchTargets is placed there too, so it can take into account any PACI(A|B)SP instructions and not excessively add BTIs at the start of functions. Reviewed By: tmatheson Differential Revision: https://reviews.llvm.org/D159357
1 parent deac021 commit eb02ee4

14 files changed

+249
-135
lines changed

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ FunctionPass *createAArch64A57FPLoadBalancing();
5151
FunctionPass *createAArch64A53Fix835769();
5252
FunctionPass *createFalkorHWPFFixPass();
5353
FunctionPass *createFalkorMarkStridedAccessesPass();
54+
FunctionPass *createAArch64PointerAuthPass();
5455
FunctionPass *createAArch64BranchTargetsPass();
5556
FunctionPass *createAArch64MIPeepholeOptPass();
5657

@@ -74,6 +75,7 @@ ModulePass *createAArch64GlobalsTaggingPass();
7475
void initializeAArch64A53Fix835769Pass(PassRegistry&);
7576
void initializeAArch64A57FPLoadBalancingPass(PassRegistry&);
7677
void initializeAArch64AdvSIMDScalarPass(PassRegistry&);
78+
void initializeAArch64PointerAuthPass(PassRegistry&);
7779
void initializeAArch64BranchTargetsPass(PassRegistry&);
7880
void initializeAArch64CFIFixupPass(PassRegistry&);
7981
void initializeAArch64CollectLOHPass(PassRegistry &);

‎llvm/lib/Target/AArch64/AArch64FrameLowering.cpp‎

Lines changed: 13 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -1381,42 +1381,6 @@ static void emitDefineCFAWithFP(MachineFunction &MF, MachineBasicBlock &MBB,
13811381
.setMIFlags(MachineInstr::FrameSetup);
13821382
}
13831383

1384-
void AArch64FrameLowering::signLR(MachineFunction &MF, MachineBasicBlock &MBB,
1385-
MachineBasicBlock::iterator MBBI,
1386-
bool NeedsWinCFI, bool *HasWinCFI) {
1387-
const auto &MFnI = *MF.getInfo<AArch64FunctionInfo>();
1388-
const AArch64Subtarget &Subtarget = MF.getSubtarget<AArch64Subtarget>();
1389-
const TargetInstrInfo *TII = Subtarget.getInstrInfo();
1390-
bool EmitCFI = MFnI.needsDwarfUnwindInfo(MF);
1391-
1392-
// Debug location must be unknown, see emitPrologue().
1393-
DebugLoc DL;
1394-
1395-
if (MFnI.shouldSignWithBKey()) {
1396-
BuildMI(MBB, MBBI, DL, TII->get(AArch64::EMITBKEY))
1397-
.setMIFlag(MachineInstr::FrameSetup);
1398-
}
1399-
1400-
// No SEH opcode for this one; it doesn't materialize into an
1401-
// instruction on Windows.
1402-
BuildMI(
1403-
MBB, MBBI, DL,
1404-
TII->get(MFnI.shouldSignWithBKey() ? AArch64::PACIBSP : AArch64::PACIASP))
1405-
.setMIFlag(MachineInstr::FrameSetup);
1406-
1407-
if (EmitCFI) {
1408-
unsigned CFIIndex =
1409-
MF.addFrameInst(MCCFIInstruction::createNegateRAState(nullptr));
1410-
BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
1411-
.addCFIIndex(CFIIndex)
1412-
.setMIFlags(MachineInstr::FrameSetup);
1413-
} else if (NeedsWinCFI) {
1414-
*HasWinCFI = true;
1415-
BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_PACSignLR))
1416-
.setMIFlag(MachineInstr::FrameSetup);
1417-
}
1418-
}
1419-
14201384
void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
14211385
MachineBasicBlock &MBB) const {
14221386
MachineBasicBlock::iterator MBBI = MBB.begin();
@@ -1450,8 +1414,12 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
14501414
emitShadowCallStackPrologue(*TII, MF, MBB, MBBI, DL, NeedsWinCFI,
14511415
MFnI.needsDwarfUnwindInfo(MF));
14521416

1453-
if (MFnI.shouldSignReturnAddress(MF))
1454-
signLR(MF, MBB, MBBI, NeedsWinCFI, &HasWinCFI);
1417+
if (MFnI.shouldSignReturnAddress(MF)) {
1418+
BuildMI(MBB, MBBI, DL, TII->get(AArch64::PAUTH_PROLOGUE))
1419+
.setMIFlag(MachineInstr::FrameSetup);
1420+
if (NeedsWinCFI)
1421+
HasWinCFI = true; // AArch64PointerAuth pass will insert SEH_PACSignLR
1422+
}
14551423

14561424
if (EmitCFI && MFnI.isMTETagged()) {
14571425
BuildMI(MBB, MBBI, DL, TII->get(AArch64::EMITMTETAGGED))
@@ -1911,54 +1879,6 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
19111879
}
19121880
}
19131881

1914-
void AArch64FrameLowering::authenticateLR(MachineFunction &MF,
1915-
MachineBasicBlock &MBB,
1916-
bool NeedsWinCFI, bool *HasWinCFI) {
1917-
const auto &MFI = *MF.getInfo<AArch64FunctionInfo>();
1918-
const AArch64Subtarget &Subtarget = MF.getSubtarget<AArch64Subtarget>();
1919-
const TargetInstrInfo *TII = Subtarget.getInstrInfo();
1920-
bool EmitAsyncCFI = MFI.needsAsyncDwarfUnwindInfo(MF);
1921-
1922-
MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator();
1923-
DebugLoc DL;
1924-
if (MBBI != MBB.end())
1925-
DL = MBBI->getDebugLoc();
1926-
1927-
// The AUTIASP instruction assembles to a hint instruction before v8.3a so
1928-
// this instruction can safely used for any v8a architecture.
1929-
// From v8.3a onwards there are optimised authenticate LR and return
1930-
// instructions, namely RETA{A,B}, that can be used instead. In this case the
1931-
// DW_CFA_AARCH64_negate_ra_state can't be emitted.
1932-
bool TerminatorIsCombinable =
1933-
MBBI != MBB.end() && (MBBI->getOpcode() == AArch64::RET_ReallyLR ||
1934-
MBBI->getOpcode() == AArch64::RET);
1935-
if (Subtarget.hasPAuth() && TerminatorIsCombinable && !NeedsWinCFI &&
1936-
!MF.getFunction().hasFnAttribute(Attribute::ShadowCallStack)) {
1937-
BuildMI(MBB, MBBI, DL,
1938-
TII->get(MFI.shouldSignWithBKey() ? AArch64::RETAB : AArch64::RETAA))
1939-
.copyImplicitOps(*MBBI);
1940-
MBB.erase(MBBI);
1941-
} else {
1942-
BuildMI(
1943-
MBB, MBBI, DL,
1944-
TII->get(MFI.shouldSignWithBKey() ? AArch64::AUTIBSP : AArch64::AUTIASP))
1945-
.setMIFlag(MachineInstr::FrameDestroy);
1946-
1947-
if (EmitAsyncCFI) {
1948-
unsigned CFIIndex =
1949-
MF.addFrameInst(MCCFIInstruction::createNegateRAState(nullptr));
1950-
BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
1951-
.addCFIIndex(CFIIndex)
1952-
.setMIFlags(MachineInstr::FrameDestroy);
1953-
}
1954-
if (NeedsWinCFI) {
1955-
*HasWinCFI = true;
1956-
BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_PACSignLR))
1957-
.setMIFlag(MachineInstr::FrameDestroy);
1958-
}
1959-
}
1960-
}
1961-
19621882
static bool isFuncletReturnInstr(const MachineInstr &MI) {
19631883
switch (MI.getOpcode()) {
19641884
default:
@@ -1990,8 +1910,13 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
19901910
MachineBasicBlock::iterator EpilogStartI = MBB.end();
19911911

19921912
auto FinishingTouches = make_scope_exit([&]() {
1993-
if (AFI->shouldSignReturnAddress(MF))
1994-
authenticateLR(MF, MBB, NeedsWinCFI, &HasWinCFI);
1913+
if (AFI->shouldSignReturnAddress(MF)) {
1914+
BuildMI(MBB, MBB.getFirstTerminator(), DL,
1915+
TII->get(AArch64::PAUTH_EPILOGUE))
1916+
.setMIFlag(MachineInstr::FrameDestroy);
1917+
if (NeedsWinCFI)
1918+
HasWinCFI = true; // AArch64PointerAuth pass will insert SEH_PACSignLR
1919+
}
19951920
if (needsShadowCallStackPrologueEpilogue(MF))
19961921
emitShadowCallStackEpilogue(*TII, MF, MBB, MBB.getFirstTerminator(), DL);
19971922
if (EmitCFI)

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

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,13 +30,6 @@ class AArch64FrameLowering : public TargetFrameLowering {
3030
eliminateCallFramePseudoInstr(MachineFunction &MF, MachineBasicBlock &MBB,
3131
MachineBasicBlock::iterator I) const override;
3232

33-
static void signLR(MachineFunction &MF, MachineBasicBlock &MBB,
34-
MachineBasicBlock::iterator MBBI, bool NeedsWinCFI,
35-
bool *HasWinCFI);
36-
37-
static void authenticateLR(MachineFunction &MF, MachineBasicBlock &MBB,
38-
bool NeedsWinCFI, bool *HasWinCFI);
39-
4033
/// emitProlog/emitEpilog - These methods insert prolog and epilog code into
4134
/// the function.
4235
void emitPrologue(MachineFunction &MF, MachineBasicBlock &MBB) const override;

‎llvm/lib/Target/AArch64/AArch64InstrInfo.cpp‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7963,6 +7963,8 @@ AArch64InstrInfo::getOutliningTypeImpl(MachineBasicBlock::iterator &MIT,
79637963
case AArch64::RETAA:
79647964
case AArch64::RETAB:
79657965
case AArch64::EMITBKEY:
7966+
case AArch64::PAUTH_PROLOGUE:
7967+
case AArch64::PAUTH_EPILOGUE:
79667968
return outliner::InstrType::Illegal;
79677969
}
79687970

@@ -8114,15 +8116,16 @@ void AArch64InstrInfo::fixupPostOutline(MachineBasicBlock &MBB) const {
81148116
}
81158117

81168118
static void signOutlinedFunction(MachineFunction &MF, MachineBasicBlock &MBB,
8119+
const AArch64InstrInfo *TII,
81178120
bool ShouldSignReturnAddr) {
81188121
if (!ShouldSignReturnAddr)
81198122
return;
81208123

8121-
AArch64FrameLowering::signLR(MF, MBB, MBB.begin(),
8122-
/*NeedsWinCFI=*/false, /*HasWinCFI=*/nullptr);
8123-
8124-
AArch64FrameLowering::authenticateLR(MF, MBB, /*NeedsWinCFI=*/false,
8125-
/*HasWinCFI=*/nullptr);
8124+
BuildMI(MBB, MBB.begin(), DebugLoc(), TII->get(AArch64::PAUTH_PROLOGUE))
8125+
.setMIFlag(MachineInstr::FrameSetup);
8126+
BuildMI(MBB, MBB.getFirstInstrTerminator(), DebugLoc(),
8127+
TII->get(AArch64::PAUTH_EPILOGUE))
8128+
.setMIFlag(MachineInstr::FrameDestroy);
81268129
}
81278130

81288131
void AArch64InstrInfo::buildOutlinedFrame(
@@ -8227,7 +8230,7 @@ void AArch64InstrInfo::buildOutlinedFrame(
82278230
// If this is a tail call outlined function, then there's already a return.
82288231
if (OF.FrameConstructionID == MachineOutlinerTailCall ||
82298232
OF.FrameConstructionID == MachineOutlinerThunk) {
8230-
signOutlinedFunction(MF, MBB, ShouldSignReturnAddr);
8233+
signOutlinedFunction(MF, MBB, this, ShouldSignReturnAddr);
82318234
return;
82328235
}
82338236

@@ -8241,7 +8244,7 @@ void AArch64InstrInfo::buildOutlinedFrame(
82418244
.addReg(AArch64::LR);
82428245
MBB.insert(MBB.end(), ret);
82438246

8244-
signOutlinedFunction(MF, MBB, ShouldSignReturnAddr);
8247+
signOutlinedFunction(MF, MBB, this, ShouldSignReturnAddr);
82458248

82468249
FI->setOutliningStyle("Function");
82478250

‎llvm/lib/Target/AArch64/AArch64InstrInfo.td‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1477,6 +1477,15 @@ def : InstAlias<"autia1716", (AUTIA1716), 0>;
14771477
def : InstAlias<"autib1716", (AUTIB1716), 0>;
14781478
def : InstAlias<"xpaclri", (XPACLRI), 0>;
14791479

1480+
// Pseudos
1481+
1482+
// Insertion point of LR signing code.
1483+
def PAUTH_PROLOGUE : Pseudo<(outs), (ins), []>, Sched<[]>;
1484+
// Insertion point of LR authentication code.
1485+
// The RET terminator of the containing machine basic block may be replaced
1486+
// with a combined RETA(A|B) instruction when rewriting this Pseudo.
1487+
def PAUTH_EPILOGUE : Pseudo<(outs), (ins), []>, Sched<[]>;
1488+
14801489
// These pointer authentication instructions require armv8.3a
14811490
let Predicates = [HasPAuth] in {
14821491

‎llvm/lib/Target/AArch64/AArch64PointerAuth.cpp‎

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
1+
//===-- AArch64PointerAuth.cpp -- Harden code using PAuth ------------------==//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#include "AArch64.h"
10+
#include "AArch64MachineFunctionInfo.h"
11+
#include "AArch64Subtarget.h"
12+
#include "llvm/CodeGen/MachineBasicBlock.h"
13+
#include "llvm/CodeGen/MachineInstrBuilder.h"
14+
#include "llvm/CodeGen/MachineModuleInfo.h"
15+
16+
using namespace llvm;
17+
18+
#define AARCH64_POINTER_AUTH_NAME "AArch64 Pointer Authentication"
19+
20+
namespace {
21+
22+
class AArch64PointerAuth : public MachineFunctionPass {
23+
public:
24+
static char ID;
25+
26+
AArch64PointerAuth() : MachineFunctionPass(ID) {}
27+
28+
bool runOnMachineFunction(MachineFunction &MF) override;
29+
30+
StringRef getPassName() const override { return AARCH64_POINTER_AUTH_NAME; }
31+
32+
private:
33+
const AArch64Subtarget *Subtarget = nullptr;
34+
const AArch64InstrInfo *TII = nullptr;
35+
36+
void signLR(MachineFunction &MF, MachineBasicBlock::iterator MBBI) const;
37+
38+
void authenticateLR(MachineFunction &MF,
39+
MachineBasicBlock::iterator MBBI) const;
40+
};
41+
42+
} // end anonymous namespace
43+
44+
INITIALIZE_PASS(AArch64PointerAuth, "aarch64-ptrauth",
45+
AARCH64_POINTER_AUTH_NAME, false, false)
46+
47+
FunctionPass *llvm::createAArch64PointerAuthPass() {
48+
return new AArch64PointerAuth();
49+
}
50+
51+
char AArch64PointerAuth::ID = 0;
52+
53+
void AArch64PointerAuth::signLR(MachineFunction &MF,
54+
MachineBasicBlock::iterator MBBI) const {
55+
const AArch64FunctionInfo *MFnI = MF.getInfo<AArch64FunctionInfo>();
56+
bool UseBKey = MFnI->shouldSignWithBKey();
57+
bool EmitCFI = MFnI->needsDwarfUnwindInfo(MF);
58+
bool NeedsWinCFI = MF.hasWinCFI();
59+
60+
MachineBasicBlock &MBB = *MBBI->getParent();
61+
62+
// Debug location must be unknown, see AArch64FrameLowering::emitPrologue.
63+
DebugLoc DL;
64+
65+
if (UseBKey) {
66+
BuildMI(MBB, MBBI, DL, TII->get(AArch64::EMITBKEY))
67+
.setMIFlag(MachineInstr::FrameSetup);
68+
}
69+
70+
// No SEH opcode for this one; it doesn't materialize into an
71+
// instruction on Windows.
72+
BuildMI(MBB, MBBI, DL,
73+
TII->get(UseBKey ? AArch64::PACIBSP : AArch64::PACIASP))
74+
.setMIFlag(MachineInstr::FrameSetup);
75+
76+
if (EmitCFI) {
77+
unsigned CFIIndex =
78+
MF.addFrameInst(MCCFIInstruction::createNegateRAState(nullptr));
79+
BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
80+
.addCFIIndex(CFIIndex)
81+
.setMIFlags(MachineInstr::FrameSetup);
82+
} else if (NeedsWinCFI) {
83+
BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_PACSignLR))
84+
.setMIFlag(MachineInstr::FrameSetup);
85+
}
86+
}
87+
88+
void AArch64PointerAuth::authenticateLR(
89+
MachineFunction &MF, MachineBasicBlock::iterator MBBI) const {
90+
const AArch64FunctionInfo *MFnI = MF.getInfo<AArch64FunctionInfo>();
91+
bool UseBKey = MFnI->shouldSignWithBKey();
92+
bool EmitAsyncCFI = MFnI->needsAsyncDwarfUnwindInfo(MF);
93+
bool NeedsWinCFI = MF.hasWinCFI();
94+
95+
MachineBasicBlock &MBB = *MBBI->getParent();
96+
DebugLoc DL = MBBI->getDebugLoc();
97+
// MBBI points to a PAUTH_EPILOGUE instruction to be replaced and
98+
// TI points to a terminator instruction that may or may not be combined.
99+
// Note that inserting new instructions "before MBBI" and "before TI" is
100+
// not the same because if ShadowCallStack is enabled, its instructions
101+
// are placed between MBBI and TI.
102+
MachineBasicBlock::iterator TI = MBB.getFirstInstrTerminator();
103+
104+
// The AUTIASP instruction assembles to a hint instruction before v8.3a so
105+
// this instruction can safely used for any v8a architecture.
106+
// From v8.3a onwards there are optimised authenticate LR and return
107+
// instructions, namely RETA{A,B}, that can be used instead. In this case the
108+
// DW_CFA_AARCH64_negate_ra_state can't be emitted.
109+
bool TerminatorIsCombinable =
110+
TI != MBB.end() && TI->getOpcode() == AArch64::RET;
111+
if (Subtarget->hasPAuth() && TerminatorIsCombinable && !NeedsWinCFI &&
112+
!MF.getFunction().hasFnAttribute(Attribute::ShadowCallStack)) {
113+
unsigned CombinedRetOpcode = UseBKey ? AArch64::RETAB : AArch64::RETAA;
114+
BuildMI(MBB, TI, DL, TII->get(CombinedRetOpcode)).copyImplicitOps(*TI);
115+
MBB.erase(TI);
116+
} else {
117+
unsigned AutOpcode = UseBKey ? AArch64::AUTIBSP : AArch64::AUTIASP;
118+
BuildMI(MBB, MBBI, DL, TII->get(AutOpcode))
119+
.setMIFlag(MachineInstr::FrameDestroy);
120+
121+
if (EmitAsyncCFI) {
122+
unsigned CFIIndex =
123+
MF.addFrameInst(MCCFIInstruction::createNegateRAState(nullptr));
124+
BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
125+
.addCFIIndex(CFIIndex)
126+
.setMIFlags(MachineInstr::FrameDestroy);
127+
}
128+
if (NeedsWinCFI) {
129+
BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_PACSignLR))
130+
.setMIFlag(MachineInstr::FrameDestroy);
131+
}
132+
}
133+
}
134+
135+
bool AArch64PointerAuth::runOnMachineFunction(MachineFunction &MF) {
136+
if (!MF.getInfo<AArch64FunctionInfo>()->shouldSignReturnAddress(true))
137+
return false;
138+
139+
Subtarget = &MF.getSubtarget<AArch64Subtarget>();
140+
TII = Subtarget->getInstrInfo();
141+
142+
SmallVector<MachineBasicBlock::iterator> DeletedInstrs;
143+
bool Modified = false;
144+
145+
for (auto &MBB : MF) {
146+
for (auto &MI : MBB) {
147+
auto It = MI.getIterator();
148+
switch (MI.getOpcode()) {
149+
default:
150+
// do nothing
151+
break;
152+
case AArch64::PAUTH_PROLOGUE:
153+
signLR(MF, It);
154+
DeletedInstrs.push_back(It);
155+
Modified = true;
156+
break;
157+
case AArch64::PAUTH_EPILOGUE:
158+
authenticateLR(MF, It);
159+
DeletedInstrs.push_back(It);
160+
Modified = true;
161+
break;
162+
}
163+
}
164+
}
165+
166+
for (auto MI : DeletedInstrs)
167+
MI->eraseFromParent();
168+
169+
return Modified;
170+
}

0 commit comments

Comments
 (0)