From 1de1887f5f18cdd2cffb756afcbfb5274b4fd92f Mon Sep 17 00:00:00 2001
From: Fraser Cormack <fraser@codeplay.com>
Date: Wed, 2 Jun 2021 17:58:29 +0100
Subject: [PATCH] [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
---
 .../CodeGen/SelectionDAG/LegalizeVectorTypes.cpp   |  6 +++++
 llvm/test/CodeGen/RISCV/rvv/vselect-fp-rv32.ll     | 27 ++++++++++++++++++++++
 llvm/test/CodeGen/RISCV/rvv/vselect-fp-rv64.ll     | 27 ++++++++++++++++++++++
 3 files changed, 60 insertions(+)

diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
index 586df35..22adb9a 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 ead95d1..3bd5ecb 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 fb61097..26de47d 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
+}
-- 
2.7.4