Skip to content

Commit 8d205fc

Browse files
[LLVM][SVE] Add isel for bfloat based select operations.
Patch also adds missing tests for unpacked half and float types.
1 parent 01371d6 commit 8d205fc

File tree

3 files changed

+118
-0
lines changed

3 files changed

+118
-0
lines changed

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1701,6 +1701,8 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
17011701
setOperationAction(ISD::FP_ROUND, VT, Custom);
17021702
setOperationAction(ISD::MLOAD, VT, Custom);
17031703
setOperationAction(ISD::INSERT_SUBVECTOR, VT, Custom);
1704+
setOperationAction(ISD::SELECT, VT, Custom);
1705+
setOperationAction(ISD::SELECT_CC, VT, Expand);
17041706
setOperationAction(ISD::SPLAT_VECTOR, VT, Legal);
17051707
setOperationAction(ISD::VECTOR_SPLICE, VT, Custom);
17061708

llvm/lib/Target/AArch64/SVEInstrFormats.td

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1962,6 +1962,8 @@ multiclass sve_int_sel_vvv<string asm, SDPatternOperator op> {
19621962
def : SVE_3_Op_Pat<nxv2f64, op, nxv2i1, nxv2f64, nxv2f64, !cast<Instruction>(NAME # _D)>;
19631963

19641964
def : SVE_3_Op_Pat<nxv8bf16, op, nxv8i1, nxv8bf16, nxv8bf16, !cast<Instruction>(NAME # _H)>;
1965+
def : SVE_3_Op_Pat<nxv4bf16, op, nxv4i1, nxv4bf16, nxv4bf16, !cast<Instruction>(NAME # _S)>;
1966+
def : SVE_3_Op_Pat<nxv2bf16, op, nxv2i1, nxv2bf16, nxv2bf16, !cast<Instruction>(NAME # _D)>;
19651967

19661968
def : InstAlias<"mov $Zd, $Pg/m, $Zn",
19671969
(!cast<Instruction>(NAME # _B) ZPR8:$Zd, PPRAny:$Pg, ZPR8:$Zn, ZPR8:$Zd), 1>;

llvm/test/CodeGen/AArch64/sve-select.ll

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,30 @@ define <vscale x 8 x half> @select_nxv8f16(i1 %cond, <vscale x 8 x half> %a, <
109109
ret <vscale x 8 x half> %res
110110
}
111111

112+
define <vscale x 4 x half> @select_nxv4f16(i1 %cond, <vscale x 4 x half> %a, <vscale x 4 x half> %b) {
113+
; CHECK-LABEL: select_nxv4f16:
114+
; CHECK: // %bb.0:
115+
; CHECK-NEXT: // kill: def $w0 killed $w0 def $x0
116+
; CHECK-NEXT: sbfx x8, x0, #0, #1
117+
; CHECK-NEXT: whilelo p0.s, xzr, x8
118+
; CHECK-NEXT: sel z0.s, p0, z0.s, z1.s
119+
; CHECK-NEXT: ret
120+
%res = select i1 %cond, <vscale x 4 x half> %a, <vscale x 4 x half> %b
121+
ret <vscale x 4 x half> %res
122+
}
123+
124+
define <vscale x 2 x half> @select_nxv2f16(i1 %cond, <vscale x 2 x half> %a, <vscale x 2 x half> %b) {
125+
; CHECK-LABEL: select_nxv2f16:
126+
; CHECK: // %bb.0:
127+
; CHECK-NEXT: // kill: def $w0 killed $w0 def $x0
128+
; CHECK-NEXT: sbfx x8, x0, #0, #1
129+
; CHECK-NEXT: whilelo p0.d, xzr, x8
130+
; CHECK-NEXT: sel z0.d, p0, z0.d, z1.d
131+
; CHECK-NEXT: ret
132+
%res = select i1 %cond, <vscale x 2 x half> %a, <vscale x 2 x half> %b
133+
ret <vscale x 2 x half> %res
134+
}
135+
112136
define <vscale x 4 x float> @select_nxv4f32(i1 %cond, <vscale x 4 x float> %a, <vscale x 4 x float> %b) {
113137
; CHECK-LABEL: select_nxv4f32:
114138
; CHECK: // %bb.0:
@@ -121,6 +145,18 @@ define <vscale x 4 x float> @select_nxv4f32(i1 %cond, <vscale x 4 x float> %a,
121145
ret <vscale x 4 x float> %res
122146
}
123147

148+
define <vscale x 2 x float> @select_nxv2f32(i1 %cond, <vscale x 2 x float> %a, <vscale x 2 x float> %b) {
149+
; CHECK-LABEL: select_nxv2f32:
150+
; CHECK: // %bb.0:
151+
; CHECK-NEXT: // kill: def $w0 killed $w0 def $x0
152+
; CHECK-NEXT: sbfx x8, x0, #0, #1
153+
; CHECK-NEXT: whilelo p0.d, xzr, x8
154+
; CHECK-NEXT: sel z0.d, p0, z0.d, z1.d
155+
; CHECK-NEXT: ret
156+
%res = select i1 %cond, <vscale x 2 x float> %a, <vscale x 2 x float> %b
157+
ret <vscale x 2 x float> %res
158+
}
159+
124160
define <vscale x 2 x double> @select_nxv2f64(i1 %cond, <vscale x 2 x double> %a, <vscale x 2 x double> %b) {
125161
; CHECK-LABEL: select_nxv2f64:
126162
; CHECK: // %bb.0:
@@ -133,6 +169,42 @@ define <vscale x 2 x double> @select_nxv2f64(i1 %cond, <vscale x 2 x double> %
133169
ret <vscale x 2 x double> %res
134170
}
135171

172+
define <vscale x 8 x bfloat> @select_nxv8bf16(i1 %cond, <vscale x 8 x bfloat> %a, <vscale x 8 x bfloat> %b) {
173+
; CHECK-LABEL: select_nxv8bf16:
174+
; CHECK: // %bb.0:
175+
; CHECK-NEXT: // kill: def $w0 killed $w0 def $x0
176+
; CHECK-NEXT: sbfx x8, x0, #0, #1
177+
; CHECK-NEXT: whilelo p0.h, xzr, x8
178+
; CHECK-NEXT: sel z0.h, p0, z0.h, z1.h
179+
; CHECK-NEXT: ret
180+
%res = select i1 %cond, <vscale x 8 x bfloat> %a, <vscale x 8 x bfloat> %b
181+
ret <vscale x 8 x bfloat> %res
182+
}
183+
184+
define <vscale x 4 x bfloat> @select_nxv4bf16(i1 %cond, <vscale x 4 x bfloat> %a, <vscale x 4 x bfloat> %b) {
185+
; CHECK-LABEL: select_nxv4bf16:
186+
; CHECK: // %bb.0:
187+
; CHECK-NEXT: // kill: def $w0 killed $w0 def $x0
188+
; CHECK-NEXT: sbfx x8, x0, #0, #1
189+
; CHECK-NEXT: whilelo p0.s, xzr, x8
190+
; CHECK-NEXT: sel z0.s, p0, z0.s, z1.s
191+
; CHECK-NEXT: ret
192+
%res = select i1 %cond, <vscale x 4 x bfloat> %a, <vscale x 4 x bfloat> %b
193+
ret <vscale x 4 x bfloat> %res
194+
}
195+
196+
define <vscale x 2 x bfloat> @select_nxv2bf16(i1 %cond, <vscale x 2 x bfloat> %a, <vscale x 2 x bfloat> %b) {
197+
; CHECK-LABEL: select_nxv2bf16:
198+
; CHECK: // %bb.0:
199+
; CHECK-NEXT: // kill: def $w0 killed $w0 def $x0
200+
; CHECK-NEXT: sbfx x8, x0, #0, #1
201+
; CHECK-NEXT: whilelo p0.d, xzr, x8
202+
; CHECK-NEXT: sel z0.d, p0, z0.d, z1.d
203+
; CHECK-NEXT: ret
204+
%res = select i1 %cond, <vscale x 2 x bfloat> %a, <vscale x 2 x bfloat> %b
205+
ret <vscale x 2 x bfloat> %res
206+
}
207+
136208
define <vscale x 16 x i1> @select_nxv16i1(i1 %cond, <vscale x 16 x i1> %a, <vscale x 16 x i1> %b) {
137209
; CHECK-LABEL: select_nxv16i1:
138210
; CHECK: // %bb.0:
@@ -324,6 +396,20 @@ define <vscale x 2 x double> @icmp_select_nxv2f64(<vscale x 2 x double> %a, <vsc
324396
ret <vscale x 2 x double> %sel
325397
}
326398

399+
define <vscale x 2 x bfloat> @icmp_select_nxv2bf16(<vscale x 2 x bfloat> %a, <vscale x 2 x bfloat> %b, i64 %x0) {
400+
; CHECK-LABEL: icmp_select_nxv2bf16:
401+
; CHECK: // %bb.0:
402+
; CHECK-NEXT: cmp x0, #0
403+
; CHECK-NEXT: cset w8, eq
404+
; CHECK-NEXT: sbfx x8, x8, #0, #1
405+
; CHECK-NEXT: whilelo p0.d, xzr, x8
406+
; CHECK-NEXT: sel z0.d, p0, z0.d, z1.d
407+
; CHECK-NEXT: ret
408+
%mask = icmp eq i64 %x0, 0
409+
%sel = select i1 %mask, <vscale x 2 x bfloat> %a, <vscale x 2 x bfloat> %b
410+
ret <vscale x 2 x bfloat> %sel
411+
}
412+
327413
define <vscale x 4 x half> @icmp_select_nxv4f16(<vscale x 4 x half> %a, <vscale x 4 x half> %b, i64 %x0) {
328414
; CHECK-LABEL: icmp_select_nxv4f16:
329415
; CHECK: // %bb.0:
@@ -352,6 +438,20 @@ define <vscale x 4 x float> @icmp_select_nxv4f32(<vscale x 4 x float> %a, <vscal
352438
ret <vscale x 4 x float> %sel
353439
}
354440

441+
define <vscale x 4 x bfloat> @icmp_select_nxv4bf16(<vscale x 4 x bfloat> %a, <vscale x 4 x bfloat> %b, i64 %x0) {
442+
; CHECK-LABEL: icmp_select_nxv4bf16:
443+
; CHECK: // %bb.0:
444+
; CHECK-NEXT: cmp x0, #0
445+
; CHECK-NEXT: cset w8, eq
446+
; CHECK-NEXT: sbfx x8, x8, #0, #1
447+
; CHECK-NEXT: whilelo p0.s, xzr, x8
448+
; CHECK-NEXT: sel z0.s, p0, z0.s, z1.s
449+
; CHECK-NEXT: ret
450+
%mask = icmp eq i64 %x0, 0
451+
%sel = select i1 %mask, <vscale x 4 x bfloat> %a, <vscale x 4 x bfloat> %b
452+
ret <vscale x 4 x bfloat> %sel
453+
}
454+
355455
define <vscale x 8 x half> @icmp_select_nxv8f16(<vscale x 8 x half> %a, <vscale x 8 x half> %b, i64 %x0) {
356456
; CHECK-LABEL: icmp_select_nxv8f16:
357457
; CHECK: // %bb.0:
@@ -366,6 +466,20 @@ define <vscale x 8 x half> @icmp_select_nxv8f16(<vscale x 8 x half> %a, <vscale
366466
ret <vscale x 8 x half> %sel
367467
}
368468

469+
define <vscale x 8 x bfloat> @icmp_select_nxv8bf16(<vscale x 8 x bfloat> %a, <vscale x 8 x bfloat> %b, i64 %x0) {
470+
; CHECK-LABEL: icmp_select_nxv8bf16:
471+
; CHECK: // %bb.0:
472+
; CHECK-NEXT: cmp x0, #0
473+
; CHECK-NEXT: cset w8, eq
474+
; CHECK-NEXT: sbfx x8, x8, #0, #1
475+
; CHECK-NEXT: whilelo p0.h, xzr, x8
476+
; CHECK-NEXT: sel z0.h, p0, z0.h, z1.h
477+
; CHECK-NEXT: ret
478+
%mask = icmp eq i64 %x0, 0
479+
%sel = select i1 %mask, <vscale x 8 x bfloat> %a, <vscale x 8 x bfloat> %b
480+
ret <vscale x 8 x bfloat> %sel
481+
}
482+
369483
define <vscale x 1 x i64> @icmp_select_nxv1i64(<vscale x 1 x i64> %a, <vscale x 1 x i64> %b, i64 %x0) {
370484
; CHECK-LABEL: icmp_select_nxv1i64:
371485
; CHECK: // %bb.0:

0 commit comments

Comments
 (0)