[AArch64][SVE] Fix crash with icmp+select

[CarolineConcatto]

Summary

This patch changes the lowering of SELECT_CC from Legal to Expand for scalable
vector and adds support for scalable vectors in performSelectCombine.

When selecting the nodes to lower in visitSELECT it checks if it is possible to
use SELECT_CC in cases where SETCC is followed by SELECT. visistSELECT checks
if SELECT_CC is legal or custom to replace SELECT by SELECT_CC.
SELECT_CC used to be legal for scalable vector, so the node changes to
SELECT_CC. This used to crash the compiler as there is no support for SELECT_CC
with scalable vectors. So now the compiler lowers to VSELECT instead of
SELECT_CC.

[sdesmalen]

Can you also add support for predicates?

This currently also fails:

define <vscale x 16 x i1> @icmp_select_nxv16i1(<vscale x 16 x i1> %a, <vscale x 16 x i1> %b, i64 %x0) {
  %mask = icmp eq i64 %x0, 0
  %sel = select i1 %mask, <vscale x 16 x i1> %a, <vscale x 16 x i1> %b
  ret <vscale x 16 x i1> %sel
}

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

550 ↗

(On Diff #337458)

I don't think testing these wider types is necessary, as they will just rely on splitting the select, for which we already have tests.

[CarolineConcatto]

Addressing reviewer's comment

• add support for missing types: nxv2i1, nxv4i1, nxv8i1, nxv16i1
• remove redundant tests

[CarolineConcatto]

Thank you @sdesmalen for pointing out that I missed some types.
Now nxv[2|4|8|16]i1 works. I also removed the redundant tests that you've asked for.

[sdesmalen]

Thanks for the changes @CarolineConcatto

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

222 ↗

(On Diff #337682)

You can also remove the unpacked integer tests (e.g. <vscale x [2|4|8] x i8>, <vscale x [2|4] x i16>, <vscale x 2 x i32>`) as those types are promoted to a legal type.

[c-rhodes]

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

135 ↗

(On Diff #337682)

Just an observation, I noticed we have 2 tests for SVE selects:

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

Can these be combined?

[CarolineConcatto]

• Rebase with main
• Remove Expand for MVT::nxv8i8, MVT::nxv4i16, MVT::nxv2i32, as it is done by default
• Remove tests for MVT::nxv8i8, MVT::nxv4i16, MVT::nxv2i32

[CarolineConcatto]

Thank you @sdesmalen and @c-rhodes for the review.
I'll comment you reviews in here, because the file select-sve.ll is deleted and the tests moved to sve-select.ll
@sdesmalen I have removed the tests you have asked.
I also remove the Expand in the Lowering for these types:
{MVT::nxv8i8, MVT::nxv4i16, MVT::nxv2i32}
as I believe this happens by default.
@c-rhodes the tests from select-sve.ll are now in sve-select.ll and select-sve.ll deleted.
That is already on main.

[sdesmalen]

LGTM, thanks @CarolineConcatto

[efriedma]

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

255

Not really related to this patch, but maybe looking into why we aren't combining the cset/sbfx into csinv.

[CarolineConcatto]

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

255

Thank you @efriedma.
We will have a look and create a patch for this in the future.
But like you wrote it is not related to this patch.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

for (auto VT : {MVT::nxv16i8, MVT::nxv8i16, MVT::nxv4i32, MVT::nxv2i64}) {
       setOperationAction(ISD::VECREDUCE_UMIN, VT, Custom);
       setOperationAction(ISD::VECREDUCE_UMAX, VT, Custom);
       setOperationAction(ISD::VECREDUCE_SMIN, VT, Custom);
       setOperationAction(ISD::VECREDUCE_SMAX, VT, Custom);
       setOperationAction(ISD::STEP_VECTOR, VT, Custom);
 
       setOperationAction(ISD::UMUL_LOHI, VT, Expand);
       setOperationAction(ISD::SMUL_LOHI, VT, Expand);
+      setOperationAction(ISD::SELECT_CC, VT, Expand);
     }
 
     // Illegal unpacked integer vector types.
     for (auto VT : {MVT::nxv8i8, MVT::nxv4i16, MVT::nxv2i32}) {
       setOperationAction(ISD::EXTRACT_SUBVECTOR, VT, Custom);
       setOperationAction(ISD::INSERT_SUBVECTOR, VT, Custom);
     }
 
     for (auto VT : {MVT::nxv16i1, MVT::nxv8i1, MVT::nxv4i1, MVT::nxv2i1}) {
       setOperationAction(ISD::CONCAT_VECTORS, VT, Custom);
       setOperationAction(ISD::SELECT, VT, Custom);
       setOperationAction(ISD::SETCC, VT, Custom);
       setOperationAction(ISD::SPLAT_VECTOR, VT, Custom);
       setOperationAction(ISD::TRUNCATE, VT, Custom);
       setOperationAction(ISD::VECREDUCE_AND, VT, Custom);
       setOperationAction(ISD::VECREDUCE_OR, VT, Custom);
       setOperationAction(ISD::VECREDUCE_XOR, VT, Custom);
 
+      setOperationAction(ISD::SELECT_CC, VT, Expand);
+
       // There are no legal MVT::nxv16f## based types.
       if (VT != MVT::nxv16i1) {
         setOperationAction(ISD::SINT_TO_FP, VT, Custom);
         setOperationAction(ISD::UINT_TO_FP, VT, Custom);
       }
     }
 
     for (auto VT : {MVT::nxv2f16, MVT::nxv4f16, MVT::nxv8f16, MVT::nxv2f32,
for (auto VT : {MVT::nxv2f16, MVT::nxv4f16, MVT::nxv8f16, MVT::nxv2f32,
       setOperationAction(ISD::FSQRT, VT, Custom);
       setOperationAction(ISD::FABS, VT, Custom);
       setOperationAction(ISD::FP_EXTEND, VT, Custom);
       setOperationAction(ISD::FP_ROUND, VT, Custom);
       setOperationAction(ISD::VECREDUCE_FADD, VT, Custom);
       setOperationAction(ISD::VECREDUCE_FMAX, VT, Custom);
       setOperationAction(ISD::VECREDUCE_FMIN, VT, Custom);
       setOperationAction(ISD::VECREDUCE_SEQ_FADD, VT, Custom);
+
+      setOperationAction(ISD::SELECT_CC, VT, Expand);
     }
 
     for (auto VT : {MVT::nxv2bf16, MVT::nxv4bf16, MVT::nxv8bf16}) {
       setOperationAction(ISD::CONCAT_VECTORS, VT, Custom);
       setOperationAction(ISD::MGATHER, VT, Custom);
       setOperationAction(ISD::MSCATTER, VT, Custom);
     }
 
static SDValue performSelectCombine(SDNode *N,
                                     TargetLowering::DAGCombinerInfo &DCI) {
   SelectionDAG &DAG = DCI.DAG;
   SDValue N0 = N->getOperand(0);
   EVT ResVT = N->getValueType(0);
 
   if (N0.getOpcode() != ISD::SETCC)
     return SDValue();
 
+  if (ResVT.isScalableVector())
+    return SDValue();
+
   // Make sure the SETCC result is either i1 (initial DAG), or i32, the lowered
   // scalar SetCCResultType. We also don't expect vectors, because we assume
   // that selects fed by vector SETCCs are canonicalized to VSELECT.
   assert((N0.getValueType() == MVT::i1 || N0.getValueType() == MVT::i32) &&
          "Scalar-SETCC feeding SELECT has unexpected result type!");
 
   // If NumMaskElts == 0, the comparison is larger than select result. The
   // largest real NEON comparison is 64-bits per lane, which means the result is

llvm/lib/Target/AArch64/SVEInstrFormats.td

multiclass sve_int_sel_vvv<string asm, SDPatternOperator op> {
   def _D : sve_int_sel_vvv<0b11, asm, ZPR64>;
 
   def : SVE_3_Op_Pat<nxv16i8, op, nxv16i1, nxv16i8, nxv16i8, !cast<Instruction>(NAME # _B)>;
   def : SVE_3_Op_Pat<nxv8i16, op, nxv8i1,  nxv8i16, nxv8i16, !cast<Instruction>(NAME # _H)>;
   def : SVE_3_Op_Pat<nxv4i32, op, nxv4i1,  nxv4i32, nxv4i32, !cast<Instruction>(NAME # _S)>;
   def : SVE_3_Op_Pat<nxv2i64, op, nxv2i1,  nxv2i64, nxv2i64, !cast<Instruction>(NAME # _D)>;
 
   def : SVE_3_Op_Pat<nxv8f16, op, nxv8i1,  nxv8f16, nxv8f16, !cast<Instruction>(NAME # _H)>;
+  def : SVE_3_Op_Pat<nxv4f16, op, nxv4i1,  nxv4f16, nxv4f16, !cast<Instruction>(NAME # _S)>;
   def : SVE_3_Op_Pat<nxv4f32, op, nxv4i1,  nxv4f32, nxv4f32, !cast<Instruction>(NAME # _S)>;
+  def : SVE_3_Op_Pat<nxv2f16, op, nxv2i1,  nxv2f16, nxv2f16, !cast<Instruction>(NAME # _D)>;
   def : SVE_3_Op_Pat<nxv2f32, op, nxv2i1,  nxv2f32, nxv2f32, !cast<Instruction>(NAME # _D)>;
   def : SVE_3_Op_Pat<nxv2f64, op, nxv2i1,  nxv2f64, nxv2f64, !cast<Instruction>(NAME # _D)>;
 
   def : SVE_3_Op_Pat<nxv8bf16, op, nxv8i1,  nxv8bf16, nxv8bf16, !cast<Instruction>(NAME # _H)>;
 
   def : InstAlias<"mov $Zd, $Pg/m, $Zn",
                   (!cast<Instruction>(NAME # _B) ZPR8:$Zd, PPRAny:$Pg, ZPR8:$Zn, ZPR8:$Zd), 1>;
   def : InstAlias<"mov $Zd, $Pg/m, $Zn",

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

define <vscale x 2 x double> @sel_nxv8f64(<vscale x 2 x i1> %p,
                                           <vscale x 2 x double> %dst,
                                           <vscale x 2 x double> %a) {
 ; CHECK-LABEL: sel_nxv8f64:
 ; CHECK:         mov z0.d, p0/m, z1.d
 ; CHECK-NEXT:    ret
   %sel = select <vscale x 2 x i1> %p, <vscale x 2 x double> %a, <vscale x 2 x double> %dst
   ret <vscale x 2 x double> %sel
 }
+
+; Check icmp+select
+
+define <vscale x 2 x half> @icmp_select_nxv2f16(<vscale x 2 x half> %a, <vscale x 2 x half> %b, i64 %x0) {
+; CHECK-LABEL: icmp_select_nxv2f16
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmp  x0, #0
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo  p0.d, xzr, x8
+; CHECK-NEXT:    sel  z0.d, p0, z0.d, z1.d
+; CHECK-NEXT:    ret
+  %mask = icmp eq i64 %x0, 0
+  %sel = select i1 %mask, <vscale x 2 x half> %a, <vscale x 2 x half> %b
+  ret <vscale x 2 x half> %sel
+}
+
+define <vscale x 2 x float> @icmp_select_nxv2f32(<vscale x 2 x float> %a, <vscale x 2 x float> %b, i64 %x0) {
+; CHECK-LABEL: icmp_select_nxv2f32
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmp  x0, #0
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo  p0.d, xzr, x8
+; CHECK-NEXT:    sel  z0.d, p0, z0.d, z1.d
+; CHECK-NEXT:    ret
+  %mask = icmp eq i64 %x0, 0
+  %sel = select i1 %mask, <vscale x 2 x float> %a, <vscale x 2 x float> %b
+  ret <vscale x 2 x float> %sel
+}
+
+define <vscale x 2 x double> @icmp_select_nxv2f64(<vscale x 2 x double> %a, <vscale x 2 x double> %b, i64 %x0) {
+; CHECK-LABEL: icmp_select_nxv2f64
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmp  x0, #0
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo  p0.d, xzr, x8
+; CHECK-NEXT:    sel  z0.d, p0, z0.d, z1.d
+; CHECK-NEXT:    ret
+  %mask = icmp eq i64 %x0, 0
+  %sel = select i1 %mask, <vscale x 2 x double> %a, <vscale x 2 x double> %b
+  ret <vscale x 2 x double> %sel
+}
+
+define <vscale x 4 x half> @icmp_select_nxv4f16(<vscale x 4 x half> %a, <vscale x 4 x half> %b, i64 %x0) {
+; CHECK-LABEL: icmp_select_nxv4f16
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmp  x0, #0
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo  p0.s, xzr, x8
+; CHECK-NEXT:    sel  z0.s, p0, z0.s, z1.s
+; CHECK-NEXT:    ret
+  %mask = icmp eq i64 %x0, 0
+  %sel = select i1 %mask, <vscale x 4 x half> %a, <vscale x 4 x half> %b
+  ret <vscale x 4 x half> %sel
+}
+
+define <vscale x 4 x float> @icmp_select_nxv4f32(<vscale x 4 x float> %a, <vscale x 4 x float> %b, i64 %x0) {
+; CHECK-LABEL: icmp_select_nxv4f32
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmp  x0, #0
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo  p0.s, xzr, x8
+; CHECK-NEXT:    sel  z0.s, p0, z0.s, z1.s
+; CHECK-NEXT:    ret
+  %mask = icmp eq i64 %x0, 0
+  %sel = select i1 %mask, <vscale x 4 x float> %a, <vscale x 4 x float> %b
+  ret <vscale x 4 x float> %sel
+}
+
+define <vscale x 8 x half> @icmp_select_nxv8f16(<vscale x 8 x half> %a, <vscale x 8 x half> %b, i64 %x0) {
+; CHECK-LABEL: icmp_select_nxv8f16
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmp  x0, #0
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo  p0.h, xzr, x8
+; CHECK-NEXT:    sel  z0.h, p0, z0.h, z1.h
+; CHECK-NEXT:    ret
+  %mask = icmp eq i64 %x0, 0
+  %sel = select i1 %mask, <vscale x 8 x half> %a, <vscale x 8 x half> %b
+  ret <vscale x 8 x half> %sel
+}
+
+define <vscale x 2 x i64> @icmp_select_nxv2i64(<vscale x 2 x i64> %a, <vscale x 2 x i64> %b, i64 %x0) {
+; CHECK-LABEL: icmp_select_nxv2i64
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmp  x0, #0
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo  p0.d, xzr, x8
+; CHECK-NEXT:    sel  z0.d, p0, z0.d, z1.d
+; CHECK-NEXT:    ret
+  %mask = icmp eq i64 %x0, 0
+  %sel = select i1 %mask, <vscale x 2 x i64> %a, <vscale x 2 x i64> %b
+  ret <vscale x 2 x i64> %sel
+}
+
+define <vscale x 4 x i32> @icmp_select_nxv4i32(<vscale x 4 x i32> %a, <vscale x 4 x i32> %b, i64 %x0) {
+; CHECK-LABEL: icmp_select_nxv4i32
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmp  x0, #0
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo  p0.s, xzr, x8
+; CHECK-NEXT:    sel  z0.s, p0, z0.s, z1.s
+; CHECK-NEXT:    ret
+  %mask = icmp eq i64 %x0, 0
+  %sel = select i1 %mask, <vscale x 4 x i32> %a, <vscale x 4 x i32> %b
+  ret <vscale x 4 x i32> %sel
+}
+
+define <vscale x 8 x i16> @icmp_select_nxv8i16(<vscale x 8 x i16> %a, <vscale x 8 x i16> %b, i64 %x0) {
+; CHECK-LABEL: icmp_select_nxv8i16
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmp  x0, #0
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo  p0.h, xzr, x8
+; CHECK-NEXT:    sel  z0.h, p0, z0.h, z1.h
+; CHECK-NEXT:    ret
+  %mask = icmp eq i64 %x0, 0
+  %sel = select i1 %mask, <vscale x 8 x i16> %a, <vscale x 8 x i16> %b
+  ret <vscale x 8 x i16> %sel
+}
+
+define  <vscale x 16 x i8> @icmp_select_nxv16i8(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b, i64 %x0) {
+; CHECK-LABEL: icmp_select_nxv16i8
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmp  x0, #0
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo  p0.b, xzr, x8
+; CHECK-NEXT:    sel  z0.b, p0, z0.b, z1.b
+; CHECK-NEXT:    ret
+  %mask = icmp eq i64 %x0, 0
+  %sel = select i1 %mask, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b
+  ret <vscale x 16 x i8> %sel
+}
+
+define <vscale x 2 x i1> @icmp_select_nxv2i1(<vscale x 2 x i1> %a, <vscale x 2 x i1> %b, i64 %x0) {
+; CHECK-LABEL: icmp_select_nxv2i1
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmp  x0, #0
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo  p2.d, xzr, x8
+; CHECK-NEXT:    sel  p0.b, p2, p0.b, p1.b
+; CHECK-NEXT:    ret
+    %mask = icmp eq i64 %x0, 0
+    %sel = select i1 %mask, <vscale x 2 x i1> %a, <vscale x 2 x i1> %b
+    ret <vscale x 2 x i1> %sel
+}
+define <vscale x 4 x i1> @icmp_select_nxv4i1(<vscale x 4 x i1> %a, <vscale x 4 x i1> %b, i64 %x0) {
+; CHECK-LABEL: icmp_select_nxv4i1
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmp  x0, #0
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo  p2.s, xzr, x8
+; CHECK-NEXT:    sel  p0.b, p2, p0.b, p1.b
+; CHECK-NEXT:    ret
+    %mask = icmp eq i64 %x0, 0
+    %sel = select i1 %mask, <vscale x 4 x i1> %a, <vscale x 4 x i1> %b
+    ret <vscale x 4 x i1> %sel
+}
+define <vscale x 8 x i1> @icmp_select_nxv8i1(<vscale x 8 x i1> %a, <vscale x 8 x i1> %b, i64 %x0) {
+; CHECK-LABEL: icmp_select_nxv8i1
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmp  x0, #0
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo  p2.h, xzr, x8
+; CHECK-NEXT:    sel  p0.b, p2, p0.b, p1.b
+; CHECK-NEXT:    ret
+    %mask = icmp eq i64 %x0, 0
+    %sel = select i1 %mask, <vscale x 8 x i1> %a, <vscale x 8 x i1> %b
+    ret <vscale x 8 x i1> %sel
+}
+define <vscale x 16 x i1> @icmp_select_nxv16i1(<vscale x 16 x i1> %a, <vscale x 16 x i1> %b, i64 %x0) {
+; CHECK-LABEL: icmp_select_nxv16i1
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmp  x0, #0
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo  p2.b, xzr, x8
+; CHECK-NEXT:    sel  p0.b, p2, p0.b, p1.b
+; CHECK-NEXT:    ret
+    %mask = icmp eq i64 %x0, 0
+    %sel = select i1 %mask, <vscale x 16 x i1> %a, <vscale x 16 x i1> %b
+    ret <vscale x 16 x i1> %sel
+}