Skip to content

Commit 6f8a36e

Browse files
jroelofsjroelofs
committed
[aarch64][PAC] Move MachO (pro|epi)logue PAC to generic AArch64PointerAuth infra
This is _almost_ NFCI, though there are some small changes I expect re: adding cfi's that we were missing before, as well as some small improvements to epilogue sequences in a couple of cases. rdar://163365479
1 parent 8f019f1 commit 6f8a36e

File tree

6 files changed

+61
-94
lines changed

6 files changed

+61
-94
lines changed

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -86,8 +86,7 @@ static std::pair<bool, bool> GetSignReturnAddress(const Function &F,
8686
}
8787

8888
static bool ShouldSignWithBKey(const Function &F, const AArch64Subtarget &STI) {
89-
if (!STI.getTargetTriple().isOSBinFormatMachO() &&
90-
F.hasFnAttribute("ptrauth-returns"))
89+
if (F.hasFnAttribute("ptrauth-returns"))
9190
return true;
9291
if (!F.hasFnAttribute("sign-return-address-key")) {
9392
if (STI.getTargetTriple().isOSWindows())

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,17 @@ void AArch64PointerAuth::authenticateLR(
181181
TI != MBB.end() && TI->getOpcode() == AArch64::RET;
182182
MCSymbol *PACSym = MFnI->getSigningInstrLabel();
183183

184+
const MachineFrameInfo &MFI = MF.getFrameInfo();
185+
bool IsLRSpilled =
186+
llvm::any_of(MFI.getCalleeSavedInfo(), [](const CalleeSavedInfo &Info) {
187+
return Info.getReg() == AArch64::LR;
188+
});
189+
190+
MachineBasicBlock::iterator EpilogueEndI = MBB.getLastNonDebugInstr();
191+
bool IsSwiftCoroPartialReturn =
192+
MBB.end() != EpilogueEndI &&
193+
EpilogueEndI->getOpcode() == AArch64::RET_POPLESS;
194+
184195
if (Subtarget->hasPAuth() && TerminatorIsCombinable && !NeedsWinCFI &&
185196
!MF.getFunction().hasFnAttribute(Attribute::ShadowCallStack)) {
186197
if (MFnI->branchProtectionPAuthLR() && Subtarget->hasPAuthLR()) {
@@ -198,6 +209,29 @@ void AArch64PointerAuth::authenticateLR(
198209
.setMIFlag(MachineInstr::FrameDestroy);
199210
}
200211
MBB.erase(TI);
212+
} else if (IsLRSpilled && IsSwiftCoroPartialReturn) {
213+
const auto *TRI = Subtarget->getRegisterInfo();
214+
215+
MachineBasicBlock::iterator EpilogStartI = MBB.getFirstTerminator();
216+
MachineBasicBlock::iterator Begin = MBB.begin();
217+
while (EpilogStartI != Begin) {
218+
--EpilogStartI;
219+
if (!EpilogStartI->getFlag(MachineInstr::FrameDestroy)) {
220+
++EpilogStartI;
221+
break;
222+
}
223+
if (EpilogStartI->readsRegister(AArch64::X16, TRI) ||
224+
EpilogStartI->modifiesRegister(AArch64::X16, TRI))
225+
report_fatal_error("unable to use x16 for popless ret LR auth");
226+
}
227+
228+
emitFrameOffset(MBB, EpilogStartI, DL, AArch64::X16, AArch64::FP,
229+
StackOffset::getFixed(16), TII, MachineInstr::FrameDestroy);
230+
emitPACCFI(MBB, MBBI, MachineInstr::FrameDestroy, EmitAsyncCFI);
231+
BuildMI(MBB, TI, DL, TII->get(AArch64::AUTIB), AArch64::LR)
232+
.addUse(AArch64::LR)
233+
.addUse(AArch64::X16)
234+
.setMIFlag(MachineInstr::FrameDestroy);
201235
} else {
202236
if (MFnI->branchProtectionPAuthLR() && Subtarget->hasPAuthLR()) {
203237
assert(PACSym && "No PAC instruction to refer to");

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

Lines changed: 0 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -536,14 +536,6 @@ void AArch64PrologueEmitter::verifyPrologueClobbers() const {
536536
}
537537
#endif
538538

539-
static bool shouldAuthenticateLR(const MachineFunction &MF) {
540-
// Return address authentication can be enabled at the function level, using
541-
// the "ptrauth-returns" attribute.
542-
const AArch64Subtarget &Subtarget = MF.getSubtarget<AArch64Subtarget>();
543-
return Subtarget.isTargetMachO() &&
544-
MF.getFunction().hasFnAttribute("ptrauth-returns");
545-
}
546-
547539
void AArch64PrologueEmitter::determineLocalsStackSize(
548540
uint64_t StackSize, uint64_t PrologueSaveSize) {
549541
AFI->setLocalStackSize(StackSize - PrologueSaveSize);
@@ -725,18 +717,6 @@ void AArch64PrologueEmitter::emitPrologue() {
725717
BuildMI(MBB, PrologueBeginI, DL, TII->get(AArch64::EMITMTETAGGED))
726718
.setMIFlag(MachineInstr::FrameSetup);
727719

728-
// If we're saving LR, sign it first.
729-
if (shouldAuthenticateLR(MF)) {
730-
if (LLVM_UNLIKELY(!Subtarget.hasPAuth()))
731-
report_fatal_error("arm64e LR authentication requires ptrauth");
732-
for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo()) {
733-
if (Info.getReg() != AArch64::LR)
734-
continue;
735-
BuildMI(MBB, PrologueBeginI, DL, TII->get(AArch64::PACIBSP))
736-
.setMIFlags(MachineInstr::FrameSetup);
737-
}
738-
}
739-
740720
// We signal the presence of a Swift extended frame to external tools by
741721
// storing FP with 0b0001 in bits 63:60. In normal userland operation a simple
742722
// ORR is sufficient, it is assumed a Swift kernel would initialize the TBI
@@ -1436,66 +1416,6 @@ void AArch64EpilogueEmitter::emitEpilogue() {
14361416
if (MF.getFunction().getCallingConv() == CallingConv::GHC)
14371417
return;
14381418

1439-
// If we're restoring LR, authenticate it before returning.
1440-
// Use scope_exit to ensure we do that last on all return paths.
1441-
auto InsertAuthLROnExit = make_scope_exit([&]() {
1442-
if (shouldAuthenticateLR(MF)) {
1443-
if (LLVM_UNLIKELY(!Subtarget.hasPAuth()))
1444-
report_fatal_error("arm64e LR authentication requires ptrauth");
1445-
for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo()) {
1446-
if (Info.getReg() != AArch64::LR)
1447-
continue;
1448-
MachineBasicBlock::iterator TI = MBB.getFirstTerminator();
1449-
1450-
// When we're doing a popless ret (i.e., that doesn't restore SP), we
1451-
// can't rely on the exit SP being the same as the entry, but they need
1452-
// to match for the LR auth to succeed. Instead, derive the entry SP
1453-
// from our FP (using a -16 static offset for the size of the frame
1454-
// record itself), save that into X16, and use that as the discriminator
1455-
// in an AUTIB.
1456-
if (IsSwiftCoroPartialReturn) {
1457-
const auto *TRI = Subtarget.getRegisterInfo();
1458-
1459-
MachineBasicBlock::iterator EpilogStartI = MBB.getFirstTerminator();
1460-
MachineBasicBlock::iterator Begin = MBB.begin();
1461-
while (EpilogStartI != Begin) {
1462-
--EpilogStartI;
1463-
if (!EpilogStartI->getFlag(MachineInstr::FrameDestroy)) {
1464-
++EpilogStartI;
1465-
break;
1466-
}
1467-
if (EpilogStartI->readsRegister(AArch64::X16, TRI) ||
1468-
EpilogStartI->modifiesRegister(AArch64::X16, TRI))
1469-
report_fatal_error("unable to use x16 for popless ret LR auth");
1470-
}
1471-
1472-
emitFrameOffset(MBB, EpilogStartI, DL, AArch64::X16, AArch64::FP,
1473-
StackOffset::getFixed(16), TII,
1474-
MachineInstr::FrameDestroy);
1475-
BuildMI(MBB, TI, DL, TII->get(AArch64::AUTIB), AArch64::LR)
1476-
.addUse(AArch64::LR)
1477-
.addUse(AArch64::X16)
1478-
.setMIFlag(MachineInstr::FrameDestroy);
1479-
return;
1480-
}
1481-
1482-
if (TI != MBB.end() && TI->getOpcode() == AArch64::RET_ReallyLR) {
1483-
// If there is a terminator and it's a RET, we can fold AUTH into it.
1484-
// Be careful to keep the implicitly returned registers.
1485-
// By now, we don't need the ReallyLR pseudo, since it's only there
1486-
// to make it possible for LR to be used for non-RET purposes, and
1487-
// that happens in RA and PEI.
1488-
BuildMI(MBB, TI, DL, TII->get(AArch64::RETAB)).copyImplicitOps(*TI);
1489-
MBB.erase(TI);
1490-
} else {
1491-
// Otherwise, we could be in a shrink-wrapped or tail-calling block.
1492-
BuildMI(MBB, TI, DL, TII->get(AArch64::AUTIBSP));
1493-
}
1494-
}
1495-
}
1496-
});
1497-
1498-
14991419
// How much of the stack used by incoming arguments this function is expected
15001420
// to restore in this particular epilogue.
15011421
int64_t ArgumentStackToRestore = AFL.getArgumentStackToRestore(MF, MBB);

‎llvm/test/CodeGen/AArch64/ptrauth-invoke-wrapper-globals.ll‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,9 @@
77
; CHECK-NEXT: .cfi_personality 155, ___gxx_personality_v0
88
; CHECK-NEXT: .cfi_lsda 16, [[EXCEPT:Lexception[0-9]+]]
99
; CHECK-NEXT: ; %bb.0:
10+
; CHECK-NEXT: .cfi_b_key_frame
1011
; CHECK-NEXT: pacibsp
12+
; CHECK-NEXT: .cfi_negate_ra_state
1113
; CHECK-NEXT: stp x20, x19, [sp, #-32]!
1214
; CHECK-NEXT: stp x29, x30, [sp, #16]
1315
; CHECK-NEXT: .cfi_def_cfa_offset 32
@@ -57,7 +59,9 @@ continuebb:
5759
; CHECK-NEXT: .cfi_personality 155, ___gxx_personality_v0
5860
; CHECK-NEXT: .cfi_lsda 16, [[EXCEPT:Lexception[0-9]+]]
5961
; CHECK-NEXT: ; %bb.0:
62+
; CHECK-NEXT: .cfi_b_key_frame
6063
; CHECK-NEXT: pacibsp
64+
; CHECK-NEXT: .cfi_negate_ra_state
6165
; CHECK-NEXT: stp x20, x19, [sp, #-32]!
6266
; CHECK-NEXT: stp x29, x30, [sp, #16]
6367
; CHECK-NEXT: .cfi_def_cfa_offset 32
@@ -112,7 +116,9 @@ continuebb:
112116
; CHECK-NEXT: .cfi_personality 155, ___gxx_personality_v0
113117
; CHECK-NEXT: .cfi_lsda 16, [[EXCEPT:Lexception[0-9]+]]
114118
; CHECK-NEXT: ; %bb.0:
119+
; CHECK-NEXT: .cfi_b_key_frame
115120
; CHECK-NEXT: pacibsp
121+
; CHECK-NEXT: .cfi_negate_ra_state
116122
; CHECK-NEXT: stp x20, x19, [sp, #-32]!
117123
; CHECK-NEXT: stp x29, x30, [sp, #16]
118124
; CHECK-NEXT: .cfi_def_cfa_offset 32

‎llvm/test/CodeGen/AArch64/swiftcorocc-call.ll‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,9 @@ declare i64 @g(ptr, ptr)
1010
define i64 @test_call_to_swiftcoro() #0 {
1111
; CHECK-LABEL: test_call_to_swiftcoro:
1212
; CHECK: ; %bb.0:
13+
; CHECK-NEXT: .cfi_b_key_frame
1314
; CHECK-NEXT: pacibsp
15+
; CHECK-NEXT: .cfi_negate_ra_state
1416
; CHECK-NEXT: stp x26, x25, [sp, #-32]! ; 16-byte Folded Spill
1517
; CHECK-NEXT: stp x29, x30, [sp, #16] ; 16-byte Folded Spill
1618
; CHECK-NEXT: add x29, sp, #16
@@ -41,7 +43,9 @@ define i64 @test_call_to_swiftcoro() #0 {
4143
define i64 @test_call_to_normal() #0 {
4244
; CHECK-LABEL: test_call_to_normal:
4345
; CHECK: ; %bb.0:
46+
; CHECK-NEXT: .cfi_b_key_frame
4447
; CHECK-NEXT: pacibsp
48+
; CHECK-NEXT: .cfi_negate_ra_state
4549
; CHECK-NEXT: sub sp, sp, #48
4650
; CHECK-NEXT: stp x26, x25, [sp, #16] ; 16-byte Folded Spill
4751
; CHECK-NEXT: stp x29, x30, [sp, #32] ; 16-byte Folded Spill
@@ -71,7 +75,9 @@ define i64 @test_call_to_normal() #0 {
7175
define swiftcorocc i64 @test_call() #0 {
7276
; CHECK-LABEL: test_call:
7377
; CHECK: ; %bb.0:
78+
; CHECK-NEXT: .cfi_b_key_frame
7479
; CHECK-NEXT: pacibsp
80+
; CHECK-NEXT: .cfi_negate_ra_state
7581
; CHECK-NEXT: sub sp, sp, #48
7682
; CHECK-NEXT: stp x26, x25, [sp, #16] ; 16-byte Folded Spill
7783
; CHECK-NEXT: stp x29, x30, [sp, #32] ; 16-byte Folded Spill
@@ -99,7 +105,9 @@ define swiftcorocc i64 @test_call() #0 {
99105
define i64 @test_call_normal() #0 {
100106
; CHECK-LABEL: test_call_normal:
101107
; CHECK: ; %bb.0:
108+
; CHECK-NEXT: .cfi_b_key_frame
102109
; CHECK-NEXT: pacibsp
110+
; CHECK-NEXT: .cfi_negate_ra_state
103111
; CHECK-NEXT: sub sp, sp, #48
104112
; CHECK-NEXT: stp x26, x25, [sp, #16] ; 16-byte Folded Spill
105113
; CHECK-NEXT: stp x29, x30, [sp, #32] ; 16-byte Folded Spill

‎llvm/test/CodeGen/AArch64/swiftcorocc-ret-popless.ll‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -7,24 +7,26 @@ declare i64 @g(ptr, ptr)
77
define swiftcorocc i64 @test_intrin_basic() #0 {
88
; CHECK-LABEL: test_intrin_basic:
99
; CHECK: ; %bb.0:
10+
; CHECK-NEXT: .cfi_b_key_frame
1011
; CHECK-NEXT: pacibsp
12+
; CHECK-NEXT: .cfi_negate_ra_state
1113
; CHECK-NEXT: stp x29, x30, [sp, #-16]! ; 16-byte Folded Spill
1214
; CHECK-NEXT: mov x29, sp
1315
; CHECK-NEXT: .cfi_def_cfa w29, 16
1416
; CHECK-NEXT: .cfi_offset w30, -8
1517
; CHECK-NEXT: .cfi_offset w29, -16
16-
; CHECK-NEXT: add x16, x29, #16
1718
; CHECK-NEXT: ldp x29, x30, [x29] ; 16-byte Folded Reload
18-
; CHECK-NEXT: autib x30, x16
19-
; CHECK-NEXT: ret
19+
; CHECK-NEXT: retab
2020
musttail call void @llvm.ret.popless()
2121
ret i64 0
2222
}
2323

2424
define swiftcorocc i64 @test_intrin() #0 {
2525
; CHECK-LABEL: test_intrin:
2626
; CHECK: ; %bb.0:
27+
; CHECK-NEXT: .cfi_b_key_frame
2728
; CHECK-NEXT: pacibsp
29+
; CHECK-NEXT: .cfi_negate_ra_state
2830
; CHECK-NEXT: sub sp, sp, #48
2931
; CHECK-NEXT: stp x26, x25, [sp, #16] ; 16-byte Folded Spill
3032
; CHECK-NEXT: stp x29, x30, [sp, #32] ; 16-byte Folded Spill
@@ -41,11 +43,9 @@ define swiftcorocc i64 @test_intrin() #0 {
4143
; CHECK-NEXT: bl _g
4244
; CHECK-NEXT: cbz x0, LBB1_2
4345
; CHECK-NEXT: ; %bb.1: ; %else
44-
; CHECK-NEXT: add x16, x29, #16
4546
; CHECK-NEXT: ldp x26, x25, [x29, #-16] ; 16-byte Folded Reload
4647
; CHECK-NEXT: ldp x29, x30, [x29] ; 16-byte Folded Reload
47-
; CHECK-NEXT: autib x30, x16
48-
; CHECK-NEXT: ret
48+
; CHECK-NEXT: retab
4949
; CHECK-NEXT: LBB1_2: ; %then
5050
; CHECK-NEXT: add sp, sp, #16
5151
; CHECK-NEXT: ldp x29, x30, [sp, #16] ; 16-byte Folded Reload
@@ -67,7 +67,9 @@ else:
6767
define swiftcorocc i64 @test_vla(i32 %n) #0 {
6868
; SDISEL-LABEL: test_vla:
6969
; SDISEL: ; %bb.0:
70+
; SDISEL-NEXT: .cfi_b_key_frame
7071
; SDISEL-NEXT: pacibsp
72+
; SDISEL-NEXT: .cfi_negate_ra_state
7173
; SDISEL-NEXT: stp x26, x25, [sp, #-32]! ; 16-byte Folded Spill
7274
; SDISEL-NEXT: stp x29, x30, [sp, #16] ; 16-byte Folded Spill
7375
; SDISEL-NEXT: add x29, sp, #16
@@ -91,11 +93,9 @@ define swiftcorocc i64 @test_vla(i32 %n) #0 {
9193
; SDISEL-NEXT: bl _g
9294
; SDISEL-NEXT: cbz x0, LBB2_2
9395
; SDISEL-NEXT: ; %bb.1: ; %else
94-
; SDISEL-NEXT: add x16, x29, #16
9596
; SDISEL-NEXT: ldp x26, x25, [x29, #-16] ; 16-byte Folded Reload
9697
; SDISEL-NEXT: ldp x29, x30, [x29] ; 16-byte Folded Reload
97-
; SDISEL-NEXT: autib x30, x16
98-
; SDISEL-NEXT: ret
98+
; SDISEL-NEXT: retab
9999
; SDISEL-NEXT: LBB2_2: ; %then
100100
; SDISEL-NEXT: sub sp, x29, #16
101101
; SDISEL-NEXT: ldp x29, x30, [sp, #16] ; 16-byte Folded Reload
@@ -104,7 +104,9 @@ define swiftcorocc i64 @test_vla(i32 %n) #0 {
104104
;
105105
; GISEL-LABEL: test_vla:
106106
; GISEL: ; %bb.0:
107+
; GISEL-NEXT: .cfi_b_key_frame
107108
; GISEL-NEXT: pacibsp
109+
; GISEL-NEXT: .cfi_negate_ra_state
108110
; GISEL-NEXT: stp x26, x25, [sp, #-32]! ; 16-byte Folded Spill
109111
; GISEL-NEXT: stp x29, x30, [sp, #16] ; 16-byte Folded Spill
110112
; GISEL-NEXT: add x29, sp, #16
@@ -128,11 +130,9 @@ define swiftcorocc i64 @test_vla(i32 %n) #0 {
128130
; GISEL-NEXT: bl _g
129131
; GISEL-NEXT: cbz x0, LBB2_2
130132
; GISEL-NEXT: ; %bb.1: ; %else
131-
; GISEL-NEXT: add x16, x29, #16
132133
; GISEL-NEXT: ldp x26, x25, [x29, #-16] ; 16-byte Folded Reload
133134
; GISEL-NEXT: ldp x29, x30, [x29] ; 16-byte Folded Reload
134-
; GISEL-NEXT: autib x30, x16
135-
; GISEL-NEXT: ret
135+
; GISEL-NEXT: retab
136136
; GISEL-NEXT: LBB2_2: ; %then
137137
; GISEL-NEXT: sub sp, x29, #16
138138
; GISEL-NEXT: ldp x29, x30, [sp, #16] ; 16-byte Folded Reload

0 commit comments

Comments
 (0)