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

[Fraser Cormack via llvm-commits]

Author: Fraser Cormack
Date: 2021-06-03T10:24:55+01:00
New Revision: 1de1887f5f18cdd2cffb756afcbfb5274b4fd92f

URL: https://github.com/llvm/llvm-project/commit/1de1887f5f18cdd2cffb756afcbfb5274b4fd92f
DIFF: https://github.com/llvm/llvm-project/commit/1de1887f5f18cdd2cffb756afcbfb5274b4fd92f.diff

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

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.

Reviewed By: craig.topper

Differential Revision: https://reviews.llvm.org/D103536

Added: 
    

Modified: 
    llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
    llvm/test/CodeGen/RISCV/rvv/vselect-fp-rv32.ll
    llvm/test/CodeGen/RISCV/rvv/vselect-fp-rv64.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
index 586df35777898..22adb9ae52f5c 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
@@ -4254,6 +4254,12 @@ SDValue DAGTypeLegalizer::WidenVSELECTMask(SDNode *N) {
     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();

diff  --git a/llvm/test/CodeGen/RISCV/rvv/vselect-fp-rv32.ll b/llvm/test/CodeGen/RISCV/rvv/vselect-fp-rv32.ll
index ead95d14eb0e1..3bd5ecb2406e0 100644
--- a/llvm/test/CodeGen/RISCV/rvv/vselect-fp-rv32.ll
+++ b/llvm/test/CodeGen/RISCV/rvv/vselect-fp-rv32.ll
@@ -420,3 +420,30 @@ define <vscale x 16 x double> @vselect_combine_regression(<vscale x 16 x i64> %v
   %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
+}

diff  --git a/llvm/test/CodeGen/RISCV/rvv/vselect-fp-rv64.ll b/llvm/test/CodeGen/RISCV/rvv/vselect-fp-rv64.ll
index fb61097129e8e..26de47d71a19a 100644
--- a/llvm/test/CodeGen/RISCV/rvv/vselect-fp-rv64.ll
+++ b/llvm/test/CodeGen/RISCV/rvv/vselect-fp-rv64.ll
@@ -420,3 +420,30 @@ define <vscale x 16 x double> @vselect_combine_regression(<vscale x 16 x i64> %v
   %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
+}