[RISCV] Check the legality of source vector types in matchSplatAsGather

[mshockwave]

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

[llvmbot]

@llvm/pr-subscribers-backend-risc-v

Author: Min-Yih Hsu (mshockwave)

Changes

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

---

Full diff: https://github.com/llvm/llvm-project/pull/133028.diff

2 Files Affected:

• (modified) llvm/lib/Target/RISCV/RISCVISelLowering.cpp (+2-1)
• (modified) llvm/test/CodeGen/RISCV/rvv/splat-vectors.ll (+15)

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index f09b8e42b9005..26738856cd3fe 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -3558,7 +3558,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:

[lukel97]

I thought matchSplatAsGather was only supposed to be run with all types legal. Maybe the best place to fix this would be in performDAGCombine?

llvm-project/llvm/lib/Target/RISCV/RISCVISelLowering.cpp

Lines 19133 to 19142 in 5d0c18e

	case ISD::SPLAT_VECTOR: {
	EVT VT = N->getValueType(0);
	// Only perform this combine on legal MVT types.
	if (!isTypeLegal(VT))
	break;
	if (auto Gather = matchSplatAsGather(N->getOperand(0), VT.getSimpleVT(), N,
	DAG, Subtarget))
	return Gather;
	break;
	}

Can we replace the !isTypeLegal(VT) with DCI.isBeforeLegalize()? I presume that would also cover the source type being illegal.

[topperc]

I thought matchSplatAsGather was only supposed to be run with all types legal. Maybe the best place to fix this would be in performDAGCombine?

llvm-project/llvm/lib/Target/RISCV/RISCVISelLowering.cpp

Lines 19133 to 19142 in 5d0c18e

	case ISD::SPLAT_VECTOR: {
	EVT VT = N->getValueType(0);
	// Only perform this combine on legal MVT types.
	if (!isTypeLegal(VT))
	break;
	if (auto Gather = matchSplatAsGather(N->getOperand(0), VT.getSimpleVT(), N,
	DAG, Subtarget))
	return Gather;
	break;
	}

Can we replace the !isTypeLegal(VT) with DCI.isBeforeLegalize()? I presume that would also cover the source type being illegal.

I'm not sure if that's a good idea. If all types are legal, there isn't a DAGCombine run after the type legalizer. We'll go straight to LegalizeVectorOps. This would mean the optimization won't run at all before LegalizeVectorOps if all types are legal.

[lukel97]

I thought matchSplatAsGather was only supposed to be run with all types legal. Maybe the best place to fix this would be in performDAGCombine?

llvm-project/llvm/lib/Target/RISCV/RISCVISelLowering.cpp

Lines 19133 to 19142 in 5d0c18e

	case ISD::SPLAT_VECTOR: {
	EVT VT = N->getValueType(0);
	// Only perform this combine on legal MVT types.
	if (!isTypeLegal(VT))
	break;
	if (auto Gather = matchSplatAsGather(N->getOperand(0), VT.getSimpleVT(), N,
	DAG, Subtarget))
	return Gather;
	break;
	}

Can we replace the !isTypeLegal(VT) with DCI.isBeforeLegalize()? I presume that would also cover the source type being illegal.

I'm not sure if that's a good idea. If all types are legal, there isn't a DAGCombine run after the type legalizer. We'll go straight to LegalizeVectorOps. This would mean the optimization won't run at all before LegalizeVectorOps if all types are legal.

Oh I see, I checked and that does seem to cause problems. I think adding this check is fine then

[mikhailramalho]

LGTM

[topperc]

LGTM