[llvm] 25c5bad - [RISCV] Check the legality of source vector types in matchSplatAsGather (#133028)

[via llvm-commits]

Author: Min-Yih Hsu
Date: 2025-03-26T12:07:02-07:00
New Revision: 25c5bad2f26e96a44b84765bedd82ee9bcffd22c

URL: https://github.com/llvm/llvm-project/commit/25c5bad2f26e96a44b84765bedd82ee9bcffd22c
DIFF: https://github.com/llvm/llvm-project/commit/25c5bad2f26e96a44b84765bedd82ee9bcffd22c.diff

LOG: [RISCV] Check the legality of source vector types in matchSplatAsGather (#133028)

When we're trying to lower `extractelement + splat` with
vrgather.vi/.vx, we should also check the legality of source vector type
from `extractelement`, as the entire transformation assumes legal types.

Fixes #133020

Added: 
    

Modified: 
    llvm/lib/Target/RISCV/RISCVISelLowering.cpp
    llvm/test/CodeGen/RISCV/rvv/splat-vectors.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index 752fed3d479a2..5b5dca4b541df 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -3566,7 +3566,8 @@ static SDValue matchSplatAsGather(SDValue SplatVal, MVT VT, const SDLoc &DL,
   // FIXME: Support i1 vectors, maybe by promoting to i8?
   MVT EltTy = VT.getVectorElementType();
   MVT SrcVT = Src.getSimpleValueType();
-  if (EltTy == MVT::i1 || EltTy != SrcVT.getVectorElementType())
+  if (EltTy == MVT::i1 || EltTy != SrcVT.getVectorElementType() ||
+      !DAG.getTargetLoweringInfo().isTypeLegal(SrcVT))
     return SDValue();
   SDValue Idx = SplatVal.getOperand(1);
   // The index must be a legal type.

diff  --git a/llvm/test/CodeGen/RISCV/rvv/splat-vectors.ll b/llvm/test/CodeGen/RISCV/rvv/splat-vectors.ll
index 47db3da3fbe7a..e1d3bbea18b73 100644
--- a/llvm/test/CodeGen/RISCV/rvv/splat-vectors.ll
+++ b/llvm/test/CodeGen/RISCV/rvv/splat-vectors.ll
@@ -230,6 +230,21 @@ define <8 x float> @splat_idx_nxv4f32_v8f32_constant_7(<vscale x 4 x float> %v)
   ret <8 x float> %splat
 }
 
+; This test shouldn't crash.
+define <vscale x 2 x float> @splat_idx_illegal_type(<3 x float> %v) {
+; CHECK-LABEL: splat_idx_illegal_type:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vsetvli a0, zero, e32, m1, ta, ma
+; CHECK-NEXT:    vrgather.vi v9, v8, 0
+; CHECK-NEXT:    vmv.v.v v8, v9
+; CHECK-NEXT:    ret
+entry:
+  %x = extractelement <3 x float> %v, i64 0
+  %ins = insertelement <vscale x 2 x float> poison, float %x, i64 0
+  %splat = shufflevector <vscale x 2 x float> %ins, <vscale x 2 x float> poison, <vscale x 2 x i32> zeroinitializer
+  ret <vscale x 2 x float> %splat
+}
+
 ; Negative test, vscale might be 4
 define <8 x float> @splat_idx_nxv4f32_v8f32_constant_8(<vscale x 4 x float> %v) {
 ; CHECK-LABEL: splat_idx_nxv4f32_v8f32_constant_8: