Skip to content

Commit f1b2dd2

Browse files
atrosinenkoatrosinenko
authored
[AArch64][BTI] Prevent Machine Scheduler from moving branch targets (#68313)
Moving instructions that are recognized as branch targets by BTI can result in runtime crash. In outliner tests, replaced "BRK 1" with "HINT 0" (a.k.a. NOP) as a generic outlinable instruction.
1 parent 07f7f1c commit f1b2dd2

File tree

5 files changed

+237
-11
lines changed

5 files changed

+237
-11
lines changed

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

Lines changed: 33 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1133,6 +1133,11 @@ bool AArch64InstrInfo::isSchedulingBoundary(const MachineInstr &MI,
11331133
const MachineFunction &MF) const {
11341134
if (TargetInstrInfo::isSchedulingBoundary(MI, MBB, MF))
11351135
return true;
1136+
1137+
// Do not move an instruction that can be recognized as a branch target.
1138+
if (hasBTISemantics(MI))
1139+
return true;
1140+
11361141
switch (MI.getOpcode()) {
11371142
case AArch64::HINT:
11381143
// CSDB hints are scheduling barriers.
@@ -4081,6 +4086,32 @@ bool AArch64InstrInfo::isQForm(const MachineInstr &MI) {
40814086
return llvm::any_of(MI.operands(), IsQFPR);
40824087
}
40834088

4089+
bool AArch64InstrInfo::hasBTISemantics(const MachineInstr &MI) {
4090+
switch (MI.getOpcode()) {
4091+
case AArch64::BRK:
4092+
case AArch64::HLT:
4093+
case AArch64::PACIASP:
4094+
case AArch64::PACIBSP:
4095+
// Implicit BTI behavior.
4096+
return true;
4097+
case AArch64::PAUTH_PROLOGUE:
4098+
// PAUTH_PROLOGUE expands to PACI(A|B)SP.
4099+
return true;
4100+
case AArch64::HINT: {
4101+
unsigned Imm = MI.getOperand(0).getImm();
4102+
// Explicit BTI instruction.
4103+
if (Imm == 32 || Imm == 34 || Imm == 36 || Imm == 38)
4104+
return true;
4105+
// PACI(A|B)SP instructions.
4106+
if (Imm == 25 || Imm == 27)
4107+
return true;
4108+
return false;
4109+
}
4110+
default:
4111+
return false;
4112+
}
4113+
}
4114+
40844115
bool AArch64InstrInfo::isFpOrNEON(const MachineInstr &MI) {
40854116
auto IsFPR = [&](const MachineOperand &Op) {
40864117
if (!Op.isReg())
@@ -8829,11 +8860,8 @@ AArch64InstrInfo::getOutliningTypeImpl(MachineBasicBlock::iterator &MIT,
88298860

88308861
// Don't outline BTI instructions, because that will prevent the outlining
88318862
// site from being indirectly callable.
8832-
if (MI.getOpcode() == AArch64::HINT) {
8833-
int64_t Imm = MI.getOperand(0).getImm();
8834-
if (Imm == 32 || Imm == 34 || Imm == 36 || Imm == 38)
8835-
return outliner::InstrType::Illegal;
8836-
}
8863+
if (hasBTISemantics(MI))
8864+
return outliner::InstrType::Illegal;
88378865

88388866
return outliner::InstrType::Legal;
88398867
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,9 @@ class AArch64InstrInfo final : public AArch64GenInstrInfo {
120120
/// Returns whether the instruction is in Q form (128 bit operands)
121121
static bool isQForm(const MachineInstr &MI);
122122

123+
/// Returns whether the instruction can be compatible with non-zero BTYPE.
124+
static bool hasBTISemantics(const MachineInstr &MI);
125+
123126
/// Returns the index for the immediate for a given instruction.
124127
static unsigned getLoadStoreImmIdx(unsigned Opc);
125128

‎llvm/test/CodeGen/AArch64/machine-outliner-noreturn-no-stack.mir

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ body: |
3434
$x0 = ORRXrs $xzr, $x1, 0
3535
$x1 = ORRXrs $xzr, $x2, 0
3636
BL @baz, implicit-def dead $lr, implicit $sp, implicit $x8, implicit $x0, implicit $x1, implicit $x3, implicit $x4, implicit-def $sp, implicit-def $x5, implicit-def $x6, implicit-def $x7, implicit-def $x8, implicit-def $x9, implicit-def $x10, implicit-def $x11, implicit-def $x12, implicit-def $x13, implicit-def $x14, implicit-def $x15, implicit-def $x18
37-
BRK 1
37+
HINT 0
3838
...
3939
---
4040
name: stack_2
@@ -56,7 +56,7 @@ body: |
5656
$x0 = ORRXrs $xzr, $x1, 0
5757
$x1 = ORRXrs $xzr, $x2, 0
5858
BL @baz, implicit-def dead $lr, implicit $sp, implicit $x8, implicit $x0, implicit $x1, implicit $x3, implicit $x4, implicit-def $sp, implicit-def $x5, implicit-def $x6, implicit-def $x7, implicit-def $x8, implicit-def $x9, implicit-def $x10, implicit-def $x11, implicit-def $x12, implicit-def $x13, implicit-def $x14, implicit-def $x15, implicit-def $x18
59-
BRK 1
59+
HINT 0
6060
...
6161
---
6262
name: baz

‎llvm/test/CodeGen/AArch64/machine-outliner-noreturn-save-lr.mir

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ body: |
3131
; CHECK: $lr = ORRXrs $xzr, $x0, 0
3232
$w3 = ORRWri $wzr, 1
3333
$w4 = ORRWri $wzr, 1
34-
BRK 1
34+
HINT 0
3535
$w5 = ORRWri $wzr, 1
3636
$w6 = ORRWri $wzr, 1
3737
...
@@ -53,7 +53,7 @@ body: |
5353
; CHECK: $lr = ORRXrs $xzr, $x0, 0
5454
$w3 = ORRWri $wzr, 1
5555
$w4 = ORRWri $wzr, 1
56-
BRK 1
56+
HINT 0
5757
$w5 = ORRWri $wzr, 1
5858
$w6 = ORRWri $wzr, 1
5959
...
@@ -75,7 +75,7 @@ body: |
7575
; CHECK: $lr = ORRXrs $xzr, $x0, 0
7676
$w3 = ORRWri $wzr, 1
7777
$w4 = ORRWri $wzr, 1
78-
BRK 1
78+
HINT 0
7979
$w5 = ORRWri $wzr, 1
8080
$w6 = ORRWri $wzr, 1
8181
...
@@ -97,7 +97,7 @@ body: |
9797
; CHECK: $lr = ORRXrs $xzr, $x0, 0
9898
$w3 = ORRWri $wzr, 1
9999
$w4 = ORRWri $wzr, 1
100-
BRK 1
100+
HINT 0
101101
$w5 = ORRWri $wzr, 1
102102
$w6 = ORRWri $wzr, 1
103103
...

‎llvm/test/CodeGen/AArch64/misched-branch-targets.mir

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,195 @@
1+
# RUN: llc -o - -run-pass=machine-scheduler -misched=shuffle %s | FileCheck %s
2+
# RUN: llc -o - -run-pass=postmisched %s | FileCheck %s
3+
4+
# Check that instructions that are recognized as branch targets by BTI
5+
# are not reordered by machine instruction schedulers.
6+
7+
--- |
8+
target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
9+
target triple = "aarch64-unknown-linux-gnu"
10+
11+
define i32 @f_pac_pseudo(i32 %a, i32 %b, i32 %c) #0 "sign-return-address"="all" {
12+
entry:
13+
ret i32 0
14+
}
15+
16+
define i32 @f_pac(i32 %a, i32 %b, i32 %c) #0 "sign-return-address"="all" {
17+
entry:
18+
ret i32 0
19+
}
20+
21+
define i32 @f_bti(i32 %a, i32 %b, i32 %c) #0 {
22+
entry:
23+
ret i32 0
24+
}
25+
26+
define i32 @f_brk(i32 %a, i32 %b, i32 %c) #0 {
27+
entry:
28+
ret i32 0
29+
}
30+
31+
define i32 @f_hlt(i32 %a, i32 %b, i32 %c) #0 {
32+
entry:
33+
ret i32 0
34+
}
35+
36+
define i32 @f_nop(i32 %a, i32 %b, i32 %c) #0 {
37+
entry:
38+
ret i32 0
39+
}
40+
41+
attributes #0 = { nounwind memory(none) "target-features"="+v8.2a" }
42+
43+
...
44+
---
45+
name: f_pac_pseudo
46+
alignment: 4
47+
tracksRegLiveness: true
48+
body: |
49+
bb.0.entry:
50+
liveins: $w0, $w1, $w2, $lr
51+
52+
frame-setup PAUTH_PROLOGUE implicit-def $lr, implicit $lr, implicit $sp
53+
$w8 = ADDWrs $w0, $w1, 0
54+
$w0 = MADDWrrr $w8, $w2, $wzr
55+
RET undef $lr, implicit $w0
56+
57+
# PAUTH_EPILOGUE instruction is omitted for simplicity as it is technically possible
58+
# to move it, so it may end up at a less obvious position in a basic block.
59+
60+
# CHECK-LABEL: name: f_pac_pseudo
61+
# CHECK: body: |
62+
# CHECK-NEXT: bb.0.entry:
63+
# CHECK-NEXT: liveins: $w0, $w1, $w2, $lr
64+
#
65+
# CHECK: frame-setup PAUTH_PROLOGUE implicit-def $lr, implicit {{.*}}$lr, implicit $sp
66+
# CHECK-NEXT: $w8 = ADDWrs {{.*}}$w0, {{.*}}$w1, 0
67+
# CHECK-NEXT: $w0 = MADDWrrr {{.*}}$w8, {{.*}}$w2, $wzr
68+
# CHECK-NEXT: RET undef $lr, implicit {{.*}}$w0
69+
70+
...
71+
---
72+
name: f_pac
73+
alignment: 4
74+
tracksRegLiveness: true
75+
body: |
76+
bb.0.entry:
77+
liveins: $w0, $w1, $w2, $lr
78+
79+
frame-setup PACIASP implicit-def $lr, implicit $lr, implicit $sp
80+
$w8 = ADDWrs $w0, $w1, 0
81+
$w0 = MADDWrrr $w8, $w2, $wzr
82+
RET undef $lr, implicit $w0
83+
84+
# AUTIASP is omitted, see above.
85+
86+
# CHECK-LABEL: name: f_pac
87+
# CHECK: body: |
88+
# CHECK-NEXT: bb.0.entry:
89+
# CHECK-NEXT: liveins: $w0, $w1, $w2, $lr
90+
#
91+
# CHECK: frame-setup PACIASP implicit-def $lr, implicit {{.*}}$lr, implicit $sp
92+
# CHECK-NEXT: $w8 = ADDWrs {{.*}}$w0, {{.*}}$w1, 0
93+
# CHECK-NEXT: $w0 = MADDWrrr {{.*}}$w8, {{.*}}$w2, $wzr
94+
# CHECK-NEXT: RET undef $lr, implicit {{.*}}$w0
95+
96+
...
97+
---
98+
name: f_bti
99+
alignment: 4
100+
tracksRegLiveness: true
101+
body: |
102+
bb.0.entry:
103+
liveins: $w0, $w1, $w2, $lr
104+
105+
HINT 34
106+
$w8 = ADDWrs $w0, $w1, 0
107+
$w0 = MADDWrrr $w8, $w2, $wzr
108+
RET undef $lr, implicit $w0
109+
110+
# CHECK-LABEL: name: f_bti
111+
# CHECK: body: |
112+
# CHECK-NEXT: bb.0.entry:
113+
# CHECK-NEXT: liveins: $w0, $w1, $w2, $lr
114+
#
115+
# CHECK: HINT 34
116+
# CHECK-NEXT: $w8 = ADDWrs {{.*}}$w0, {{.*}}$w1, 0
117+
# CHECK-NEXT: $w0 = MADDWrrr {{.*}}$w8, {{.*}}$w2, $wzr
118+
# CHECK-NEXT: RET undef $lr, implicit {{.*}}$w0
119+
120+
...
121+
---
122+
name: f_brk
123+
alignment: 4
124+
tracksRegLiveness: true
125+
body: |
126+
bb.0.entry:
127+
liveins: $w0, $w1, $w2, $lr
128+
129+
BRK 1
130+
$w8 = ADDWrs $w0, $w1, 0
131+
$w0 = MADDWrrr $w8, $w2, $wzr
132+
RET undef $lr, implicit $w0
133+
134+
# CHECK-LABEL: name: f_brk
135+
# CHECK: body: |
136+
# CHECK-NEXT: bb.0.entry:
137+
# CHECK-NEXT: liveins: $w0, $w1, $w2, $lr
138+
#
139+
# CHECK: BRK 1
140+
# CHECK-NEXT: $w8 = ADDWrs {{.*}}$w0, {{.*}}$w1, 0
141+
# CHECK-NEXT: $w0 = MADDWrrr {{.*}}$w8, {{.*}}$w2, $wzr
142+
# CHECK-NEXT: RET undef $lr, implicit {{.*}}$w0
143+
144+
...
145+
---
146+
name: f_hlt
147+
alignment: 4
148+
tracksRegLiveness: true
149+
body: |
150+
bb.0.entry:
151+
liveins: $w0, $w1, $w2, $lr
152+
153+
HLT 1
154+
$w8 = ADDWrs $w0, $w1, 0
155+
$w0 = MADDWrrr $w8, $w2, $wzr
156+
RET undef $lr, implicit $w0
157+
158+
# CHECK-LABEL: name: f_hlt
159+
# CHECK: body: |
160+
# CHECK-NEXT: bb.0.entry:
161+
# CHECK-NEXT: liveins: $w0, $w1, $w2, $lr
162+
#
163+
# CHECK: HLT 1
164+
# CHECK-NEXT: $w8 = ADDWrs {{.*}}$w0, {{.*}}$w1, 0
165+
# CHECK-NEXT: $w0 = MADDWrrr {{.*}}$w8, {{.*}}$w2, $wzr
166+
# CHECK-NEXT: RET undef $lr, implicit {{.*}}$w0
167+
168+
...
169+
---
170+
name: f_nop
171+
alignment: 4
172+
tracksRegLiveness: true
173+
body: |
174+
bb.0.entry:
175+
liveins: $w0, $w1, $w2, $lr
176+
177+
HINT 0
178+
$w8 = ADDWrs $w0, $w1, 0
179+
$w0 = MADDWrrr $w8, $w2, $wzr
180+
RET undef $lr, implicit $w0
181+
182+
# Check that BTI-related instructions are left intact not because *anything*
183+
# is left intact.
184+
185+
# CHECK-LABEL: name: f_nop
186+
# CHECK: body: |
187+
# CHECK-NEXT: bb.0.entry:
188+
# CHECK-NEXT: liveins: $w0, $w1, $w2, $lr
189+
#
190+
# CHECK: $w8 = ADDWrs {{.*}}$w0, {{.*}}$w1, 0
191+
# CHECK-DAG: $w0 = MADDWrrr {{.*}}$w8, {{.*}}$w2, $wzr
192+
# CHECK-DAG: HINT 0
193+
# CHECK-NEXT: RET undef $lr, implicit {{.*}}$w0
194+
195+
...

0 commit comments

Comments
 (0)