[CodeGen] Fix a scalable-vector crash in VSELECT legalization

[frasercrmck]

Summary

The DAGTypeLegalizer::WidenVSELECTMask function is not (yet) ready for
scalable vector types, and has numerous places in which it tries to grab
either the fixed size or number of elements of its types.

I believe that it should be possible to update this method to properly
account for scalable-vector types, but we don't have test cases for
that; RISC-V bails out early on as it has legal i1 vector masks. As
such, this patch just prevents it from crashing.

[craig.topper]

LGTM

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

4258

nit: Can you add a FIXME to say that this support could be added in the future?
(or append "yet" after "vector types" in your current comment)

[frasercrmck]

• rebase & add a fixme

[frasercrmck]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

4258

Aye, I've just done so. I originally had a "yet" but didn't want it to just become some stale comment if there were no plans to change it. Do you think AArch64 may have a use for this function?

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

SDValue DAGTypeLegalizer::WidenVSELECTMask(SDNode *N) {
 
   // If this is a splitted VSELECT that was previously already handled, do
   // nothing.
   EVT CondVT = Cond->getValueType(0);
   if (CondVT.getScalarSizeInBits() != 1)
     return SDValue();
 
   EVT VSelVT = N->getValueType(0);
+
+  // This method can't handle scalable vector types.
+  // FIXME: This support could be added in the future.
+  if (VSelVT.isScalableVector())
+    return SDValue();
+
   // Only handle vector types which are a power of 2.
   if (!isPowerOf2_64(VSelVT.getSizeInBits()))
     return SDValue();
 
   // Don't touch if this will be scalarized.
   EVT FinalVT = VSelVT;
   while (getTypeAction(FinalVT) == TargetLowering::TypeSplitVector)
     FinalVT = FinalVT.getHalfNumVectorElementsVT(Ctx);

llvm/test/CodeGen/RISCV/rvv/vselect-fp-rv32.ll

 ; CHECK-NEXT:    slli a0, a0, 3
 ; CHECK-NEXT:    add sp, sp, a0
 ; CHECK-NEXT:    addi sp, sp, 16
 ; CHECK-NEXT:    ret
   %cond = icmp eq <vscale x 16 x i64> %va, zeroinitializer
   %sel = select <vscale x 16 x i1> %cond, <vscale x 16 x double> %vb, <vscale x 16 x double> zeroinitializer
   ret <vscale x 16 x double> %sel
 }
+
+define void @vselect_legalize_regression(<vscale x 16 x double> %a, <vscale x 16 x i1> %ma, <vscale x 16 x i1> %mb, <vscale x 16 x double>* %out) {
+; CHECK-LABEL: vselect_legalize_regression:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a2, zero, e8,m2,ta,mu
+; CHECK-NEXT:    vle1.v v25, (a0)
+; CHECK-NEXT:    vmand.mm v25, v0, v25
+; CHECK-NEXT:    csrr a0, vlenb
+; CHECK-NEXT:    srli a0, a0, 3
+; CHECK-NEXT:    vsetvli a2, zero, e8,mf4,ta,mu
+; CHECK-NEXT:    vslidedown.vx v0, v25, a0
+; CHECK-NEXT:    vmv1r.v v2, v25
+; CHECK-NEXT:    vsetvli a2, zero, e64,m8,ta,mu
+; CHECK-NEXT:    vmv.v.i v24, 0
+; CHECK-NEXT:    vmerge.vvm v16, v24, v16, v0
+; CHECK-NEXT:    vmv1r.v v0, v2
+; CHECK-NEXT:    vmerge.vvm v8, v24, v8, v0
+; CHECK-NEXT:    vs8r.v v8, (a1)
+; CHECK-NEXT:    slli a0, a0, 6
+; CHECK-NEXT:    add a0, a1, a0
+; CHECK-NEXT:    vs8r.v v16, (a0)
+; CHECK-NEXT:    ret
+  %cond = and <vscale x 16 x i1> %ma, %mb
+  %sel = select <vscale x 16 x i1> %cond, <vscale x 16 x double> %a, <vscale x 16 x double> zeroinitializer
+  store <vscale x 16 x double> %sel, <vscale x 16 x double>* %out
+  ret void
+}

llvm/test/CodeGen/RISCV/rvv/vselect-fp-rv64.ll

 ; CHECK-NEXT:    slli a0, a0, 3
 ; CHECK-NEXT:    add sp, sp, a0
 ; CHECK-NEXT:    addi sp, sp, 16
 ; CHECK-NEXT:    ret
   %cond = icmp eq <vscale x 16 x i64> %va, zeroinitializer
   %sel = select <vscale x 16 x i1> %cond, <vscale x 16 x double> %vb, <vscale x 16 x double> zeroinitializer
   ret <vscale x 16 x double> %sel
 }
+
+define void @vselect_legalize_regression(<vscale x 16 x double> %a, <vscale x 16 x i1> %ma, <vscale x 16 x i1> %mb, <vscale x 16 x double>* %out) {
+; CHECK-LABEL: vselect_legalize_regression:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a2, zero, e8,m2,ta,mu
+; CHECK-NEXT:    vle1.v v25, (a0)
+; CHECK-NEXT:    vmand.mm v25, v0, v25
+; CHECK-NEXT:    csrr a0, vlenb
+; CHECK-NEXT:    srli a0, a0, 3
+; CHECK-NEXT:    vsetvli a2, zero, e8,mf4,ta,mu
+; CHECK-NEXT:    vslidedown.vx v0, v25, a0
+; CHECK-NEXT:    vmv1r.v v2, v25
+; CHECK-NEXT:    vsetvli a2, zero, e64,m8,ta,mu
+; CHECK-NEXT:    vmv.v.i v24, 0
+; CHECK-NEXT:    vmerge.vvm v16, v24, v16, v0
+; CHECK-NEXT:    vmv1r.v v0, v2
+; CHECK-NEXT:    vmerge.vvm v8, v24, v8, v0
+; CHECK-NEXT:    vs8r.v v8, (a1)
+; CHECK-NEXT:    slli a0, a0, 6
+; CHECK-NEXT:    add a0, a1, a0
+; CHECK-NEXT:    vs8r.v v16, (a0)
+; CHECK-NEXT:    ret
+  %cond = and <vscale x 16 x i1> %ma, %mb
+  %sel = select <vscale x 16 x i1> %cond, <vscale x 16 x double> %a, <vscale x 16 x double> zeroinitializer
+  store <vscale x 16 x double> %sel, <vscale x 16 x double>* %out
+  ret void
+}