Skip to content

Commit 7297214

Browse files
atrosinenkoatrosinenko
committed
Rewrite discriminator matching and add more test cases
1 parent 36bd45b commit 7297214

File tree

3 files changed

+216
-22
lines changed

3 files changed

+216
-22
lines changed

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

Lines changed: 23 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -3100,8 +3100,8 @@ AArch64TargetLowering::EmitGetSMESaveSize(MachineInstr &MI,
31003100

31013101
// Helper function to find the instruction that defined a virtual register.
31023102
// If unable to find such instruction, returns nullptr.
3103-
static MachineInstr *stripVRegCopies(const MachineRegisterInfo &MRI,
3104-
Register Reg) {
3103+
static const MachineInstr *stripVRegCopies(const MachineRegisterInfo &MRI,
3104+
Register Reg) {
31053105
while (Reg.isVirtual()) {
31063106
MachineInstr *DefMI = MRI.getVRegDef(Reg);
31073107
assert(DefMI && "Virtual register definition not found");
@@ -3133,22 +3133,29 @@ void AArch64TargetLowering::fixupPtrauthDiscriminator(
31333133

31343134
Register AddrDisc = AddrDiscOp.getReg();
31353135
int64_t IntDisc = IntDiscOp.getImm();
3136-
31373136
assert(IntDisc == 0 && "Blend components are already expanded");
31383137

3139-
MachineInstr *MaybeBlend = stripVRegCopies(MRI, AddrDisc);
3140-
3141-
if (MaybeBlend) {
3142-
// Detect blend(addr, imm) which is lowered as "MOVK addr, #imm, #48".
3143-
// Alternatively, recognize small immediate modifier passed via VReg.
3144-
if (MaybeBlend->getOpcode() == AArch64::MOVKXi &&
3145-
MaybeBlend->getOperand(3).getImm() == 48) {
3146-
AddrDisc = MaybeBlend->getOperand(1).getReg();
3147-
IntDisc = MaybeBlend->getOperand(2).getImm();
3148-
} else if (MaybeBlend->getOpcode() == AArch64::MOVi32imm &&
3149-
isUInt<16>(MaybeBlend->getOperand(1).getImm())) {
3150-
AddrDisc = AArch64::NoRegister;
3151-
IntDisc = MaybeBlend->getOperand(1).getImm();
3138+
const MachineInstr *DiscMI = stripVRegCopies(MRI, AddrDisc);
3139+
if (DiscMI) {
3140+
switch (DiscMI->getOpcode()) {
3141+
case AArch64::MOVKXi:
3142+
// blend(addr, imm) which is lowered as "MOVK addr, #imm, #48".
3143+
// #imm should be an immediate and not a global symbol, for example.
3144+
if (DiscMI->getOperand(2).isImm() &&
3145+
DiscMI->getOperand(3).getImm() == 48) {
3146+
AddrDisc = DiscMI->getOperand(1).getReg();
3147+
IntDisc = DiscMI->getOperand(2).getImm();
3148+
}
3149+
break;
3150+
case AArch64::MOVi32imm:
3151+
case AArch64::MOVi64imm:
3152+
// Small immediate integer constant passed via VReg.
3153+
if (DiscMI->getOperand(1).isImm() &&
3154+
isUInt<16>(DiscMI->getOperand(1).getImm())) {
3155+
AddrDisc = AArch64::NoRegister;
3156+
IntDisc = DiscMI->getOperand(1).getImm();
3157+
}
3158+
break;
31523159
}
31533160
}
31543161

‎llvm/test/CodeGen/AArch64/ptrauth-isel.ll‎

Lines changed: 70 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,17 +13,17 @@
1313

1414
@discvar = dso_local global i64 0
1515

16-
; Make sure the components of blend(addr, imm) are recognized and passed to
17-
; PAC pseudo via separate operands to prevent substitution of the immediate
18-
; modifier.
16+
; Make sure the components of blend(addr, imm) and integer constants are
17+
; recognized and passed to PAC pseudo via separate operands to prevent
18+
; substitution of the immediate modifier.
1919
;
2020
; MIR output of the instruction selector is inspected, as it is hard to reliably
2121
; distinguish MOVKXi immediately followed by a pseudo from a standalone pseudo
2222
; instruction carrying address and immediate modifiers in its separate operands
2323
; by only observing the final asm output.
2424

25-
define i64 @small_imm_disc(i64 %addr) {
26-
; DAGISEL-LABEL: name: small_imm_disc
25+
define i64 @small_imm_disc_optimized(i64 %addr) {
26+
; DAGISEL-LABEL: name: small_imm_disc_optimized
2727
; DAGISEL: bb.0.entry:
2828
; DAGISEL-NEXT: liveins: $x0
2929
; DAGISEL-NEXT: {{ $}}
@@ -34,7 +34,7 @@ define i64 @small_imm_disc(i64 %addr) {
3434
; DAGISEL-NEXT: $x0 = COPY [[PAC]]
3535
; DAGISEL-NEXT: RET_ReallyLR implicit $x0
3636
;
37-
; GISEL-LABEL: name: small_imm_disc
37+
; GISEL-LABEL: name: small_imm_disc_optimized
3838
; GISEL: bb.1.entry:
3939
; GISEL-NEXT: liveins: $x0
4040
; GISEL-NEXT: {{ $}}
@@ -49,6 +49,35 @@ entry:
4949
ret i64 %signed
5050
}
5151

52+
; Without optimization, MOVi64imm may be used for small i64 constants as well.
53+
define i64 @small_imm_disc_non_optimized(i64 %addr) noinline optnone {
54+
; DAGISEL-LABEL: name: small_imm_disc_non_optimized
55+
; DAGISEL: bb.0.entry:
56+
; DAGISEL-NEXT: liveins: $x0
57+
; DAGISEL-NEXT: {{ $}}
58+
; DAGISEL-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
59+
; DAGISEL-NEXT: [[COPY1:%[0-9]+]]:gpr64 = COPY killed [[COPY]]
60+
; DAGISEL-NEXT: [[MOVi32imm:%[0-9]+]]:gpr32 = MOVi32imm 42
61+
; DAGISEL-NEXT: [[SUBREG_TO_REG:%[0-9]+]]:gpr64noip = SUBREG_TO_REG 0, killed [[MOVi32imm]], %subreg.sub_32
62+
; DAGISEL-NEXT: [[PAC:%[0-9]+]]:gpr64 = PAC [[COPY1]], 2, 42, killed $noreg, implicit-def dead $x17
63+
; DAGISEL-NEXT: [[COPY2:%[0-9]+]]:gpr64all = COPY [[PAC]]
64+
; DAGISEL-NEXT: $x0 = COPY [[COPY2]]
65+
; DAGISEL-NEXT: RET_ReallyLR implicit $x0
66+
;
67+
; GISEL-LABEL: name: small_imm_disc_non_optimized
68+
; GISEL: bb.1.entry:
69+
; GISEL-NEXT: liveins: $x0
70+
; GISEL-NEXT: {{ $}}
71+
; GISEL-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
72+
; GISEL-NEXT: [[MOVi64imm:%[0-9]+]]:gpr64noip = MOVi64imm 42
73+
; GISEL-NEXT: [[PAC:%[0-9]+]]:gpr64 = PAC [[COPY]], 2, 42, $noreg, implicit-def dead $x17
74+
; GISEL-NEXT: $x0 = COPY [[PAC]]
75+
; GISEL-NEXT: RET_ReallyLR implicit $x0
76+
entry:
77+
%signed = call i64 @llvm.ptrauth.sign(i64 %addr, i32 2, i64 42)
78+
ret i64 %signed
79+
}
80+
5281
define i64 @large_imm_disc_wreg(i64 %addr) {
5382
; DAGISEL-LABEL: name: large_imm_disc_wreg
5483
; DAGISEL: bb.0.entry:
@@ -101,6 +130,41 @@ entry:
101130
ret i64 %signed
102131
}
103132

133+
; Make sure blend() is lowered as expected when optimization is disabled.
134+
define i64 @blended_disc_non_optimized(i64 %addr, i64 %addrdisc) noinline optnone {
135+
; DAGISEL-LABEL: name: blended_disc_non_optimized
136+
; DAGISEL: bb.0.entry:
137+
; DAGISEL-NEXT: liveins: $x0, $x1
138+
; DAGISEL-NEXT: {{ $}}
139+
; DAGISEL-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x1
140+
; DAGISEL-NEXT: [[COPY1:%[0-9]+]]:gpr64 = COPY $x0
141+
; DAGISEL-NEXT: [[COPY2:%[0-9]+]]:gpr64 = COPY killed [[COPY1]]
142+
; DAGISEL-NEXT: [[COPY3:%[0-9]+]]:gpr64 = COPY killed [[COPY]]
143+
; DAGISEL-NEXT: [[MOVKXi:%[0-9]+]]:gpr64 = MOVKXi [[COPY3]], 42, 48
144+
; DAGISEL-NEXT: [[COPY4:%[0-9]+]]:gpr64noip = COPY [[MOVKXi]]
145+
; DAGISEL-NEXT: [[COPY5:%[0-9]+]]:gpr64noip = COPY [[COPY3]]
146+
; DAGISEL-NEXT: [[PAC:%[0-9]+]]:gpr64 = PAC [[COPY2]], 2, 42, [[COPY5]], implicit-def dead $x17
147+
; DAGISEL-NEXT: [[COPY6:%[0-9]+]]:gpr64all = COPY [[PAC]]
148+
; DAGISEL-NEXT: $x0 = COPY [[COPY6]]
149+
; DAGISEL-NEXT: RET_ReallyLR implicit $x0
150+
;
151+
; GISEL-LABEL: name: blended_disc_non_optimized
152+
; GISEL: bb.1.entry:
153+
; GISEL-NEXT: liveins: $x0, $x1
154+
; GISEL-NEXT: {{ $}}
155+
; GISEL-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
156+
; GISEL-NEXT: [[COPY1:%[0-9]+]]:gpr64 = COPY $x1
157+
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:gpr64noip = MOVKXi [[COPY1]], 42, 48
158+
; GISEL-NEXT: [[COPY2:%[0-9]+]]:gpr64noip = COPY [[COPY1]]
159+
; GISEL-NEXT: [[PAC:%[0-9]+]]:gpr64 = PAC [[COPY]], 2, 42, [[COPY2]], implicit-def dead $x17
160+
; GISEL-NEXT: $x0 = COPY [[PAC]]
161+
; GISEL-NEXT: RET_ReallyLR implicit $x0
162+
entry:
163+
%disc = call i64 @llvm.ptrauth.blend(i64 %addrdisc, i64 42)
164+
%signed = call i64 @llvm.ptrauth.sign(i64 %addr, i32 2, i64 %disc)
165+
ret i64 %signed
166+
}
167+
104168
define i64 @blend_and_sign_same_bb(i64 %addr) {
105169
; DAGISEL-LABEL: name: blend_and_sign_same_bb
106170
; DAGISEL: bb.0.entry:

‎llvm/test/CodeGen/AArch64/ptrauth-isel.mir‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5
2+
# RUN: llc -o - %s -mtriple arm64e-apple-darwin -verify-machineinstrs \
3+
# RUN: -stop-after=finalize-isel -start-before=finalize-isel | FileCheck %s
4+
# RUN: llc -o - %s -mtriple aarch64-linux-gnu -mattr=+pauth -verify-machineinstrs \
5+
# RUN: -stop-after=finalize-isel -start-before=finalize-isel | FileCheck %s
6+
7+
# This MIR-based test contains several test cases that are hard to implement
8+
# via an LLVM IR input. Most other test cases are in ptrauth-isel.ll file.
9+
10+
--- |
11+
@globalvar = dso_local global i64 0
12+
13+
define i64 @movk_correct_blend(i64 %a, i64 %b) {
14+
entry:
15+
ret i64 0
16+
}
17+
18+
define i64 @movk_wrong_shift_amount(i64 %a, i64 %b) {
19+
entry:
20+
ret i64 0
21+
}
22+
23+
define i64 @movk_non_immediate_operand(i64 %a, i64 %b) {
24+
entry:
25+
ret i64 0
26+
}
27+
28+
define i64 @movi64imm_non_immediate_operand(i64 %a) {
29+
entry:
30+
ret i64 0
31+
}
32+
...
33+
---
34+
name: movk_correct_blend
35+
tracksRegLiveness: true
36+
body: |
37+
bb.0.entry:
38+
liveins: $x0, $x1
39+
40+
; CHECK-LABEL: name: movk_correct_blend
41+
; CHECK: liveins: $x0, $x1
42+
; CHECK-NEXT: {{ $}}
43+
; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
44+
; CHECK-NEXT: [[COPY1:%[0-9]+]]:gpr64 = COPY $x1
45+
; CHECK-NEXT: [[MOVKXi:%[0-9]+]]:gpr64noip = MOVKXi [[COPY1]], 42, 48
46+
; CHECK-NEXT: [[COPY2:%[0-9]+]]:gpr64noip = COPY [[COPY1]]
47+
; CHECK-NEXT: [[PAC:%[0-9]+]]:gpr64 = PAC [[COPY]], 2, 42, killed [[COPY2]], implicit-def dead $x17
48+
; CHECK-NEXT: $x0 = COPY [[PAC]]
49+
; CHECK-NEXT: RET_ReallyLR implicit $x0
50+
%0:gpr64 = COPY $x0
51+
%1:gpr64 = COPY $x1
52+
%2:gpr64noip = MOVKXi %1, 42, 48
53+
%3:gpr64 = PAC %0, 2, 0, killed %2, implicit-def dead $x17
54+
$x0 = COPY %3
55+
RET_ReallyLR implicit $x0
56+
...
57+
---
58+
name: movk_wrong_shift_amount
59+
tracksRegLiveness: true
60+
body: |
61+
bb.0.entry:
62+
liveins: $x0, $x1
63+
64+
; CHECK-LABEL: name: movk_wrong_shift_amount
65+
; CHECK: liveins: $x0, $x1
66+
; CHECK-NEXT: {{ $}}
67+
; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
68+
; CHECK-NEXT: [[COPY1:%[0-9]+]]:gpr64 = COPY $x1
69+
; CHECK-NEXT: [[MOVKXi:%[0-9]+]]:gpr64noip = MOVKXi [[COPY1]], 42, 0
70+
; CHECK-NEXT: [[PAC:%[0-9]+]]:gpr64 = PAC [[COPY]], 2, 0, killed [[MOVKXi]], implicit-def dead $x17
71+
; CHECK-NEXT: $x0 = COPY [[PAC]]
72+
; CHECK-NEXT: RET_ReallyLR implicit $x0
73+
%0:gpr64 = COPY $x0
74+
%1:gpr64 = COPY $x1
75+
%2:gpr64noip = MOVKXi %1, 42, 0
76+
%3:gpr64 = PAC %0, 2, 0, killed %2, implicit-def dead $x17
77+
$x0 = COPY %3
78+
RET_ReallyLR implicit $x0
79+
...
80+
---
81+
name: movk_non_immediate_operand
82+
tracksRegLiveness: true
83+
body: |
84+
bb.0.entry:
85+
liveins: $x0, $x1
86+
87+
; CHECK-LABEL: name: movk_non_immediate_operand
88+
; CHECK: liveins: $x0, $x1
89+
; CHECK-NEXT: {{ $}}
90+
; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
91+
; CHECK-NEXT: [[COPY1:%[0-9]+]]:gpr64 = COPY $x1
92+
; CHECK-NEXT: [[MOVKXi:%[0-9]+]]:gpr64noip = MOVKXi [[COPY1]], target-flags(aarch64-pageoff, aarch64-nc) @globalvar, 48
93+
; CHECK-NEXT: [[PAC:%[0-9]+]]:gpr64 = PAC [[COPY]], 2, 0, killed [[MOVKXi]], implicit-def dead $x17
94+
; CHECK-NEXT: $x0 = COPY [[PAC]]
95+
; CHECK-NEXT: RET_ReallyLR implicit $x0
96+
%0:gpr64 = COPY $x0
97+
%1:gpr64 = COPY $x1
98+
%2:gpr64noip = MOVKXi %1, target-flags(aarch64-pageoff, aarch64-nc) @globalvar, 48
99+
%3:gpr64 = PAC %0, 2, 0, killed %2, implicit-def dead $x17
100+
$x0 = COPY %3
101+
RET_ReallyLR implicit $x0
102+
...
103+
---
104+
name: movi64imm_non_immediate_operand
105+
tracksRegLiveness: true
106+
body: |
107+
bb.0.entry:
108+
liveins: $x0, $x1
109+
110+
; CHECK-LABEL: name: movi64imm_non_immediate_operand
111+
; CHECK: liveins: $x0, $x1
112+
; CHECK-NEXT: {{ $}}
113+
; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
114+
; CHECK-NEXT: [[MOVi64imm:%[0-9]+]]:gpr64noip = MOVi64imm target-flags(aarch64-pageoff, aarch64-nc) @globalvar
115+
; CHECK-NEXT: [[PAC:%[0-9]+]]:gpr64 = PAC [[COPY]], 2, 0, killed [[MOVi64imm]], implicit-def dead $x17
116+
; CHECK-NEXT: $x0 = COPY [[PAC]]
117+
; CHECK-NEXT: RET_ReallyLR implicit $x0
118+
%0:gpr64 = COPY $x0
119+
%1:gpr64noip = MOVi64imm target-flags(aarch64-pageoff, aarch64-nc) @globalvar
120+
%2:gpr64 = PAC %0, 2, 0, killed %1, implicit-def dead $x17
121+
$x0 = COPY %2
122+
RET_ReallyLR implicit $x0
123+
...

0 commit comments

Comments
 (0)