[RISCV] Ensure shuffle splat operands are type-legal

[frasercrmck]

Summary

The use of SelectionDAG::getSplatValue isn't guaranteed to return a
type-legal splat value as it may implicitly extract a vector element
from another shuffle. It is not permitted to introduce an illegal type
when lowering shuffles.

This patch addresses the crash by adding a boolean flag to
getSplatValue, defaulting to false, which when set will ensure a
type-legal return value. If it is unable to do that it will fail to
return a splat value.

I've been through the existing uses of getSplatValue in other targets
and was unable to find a need or test cases showing a need to update
their uses. In some cases, the call is made during LegalizeVectorOps
which may still produce illegal scalar types. In other situations, the
illegally-typed splat value may be quickly patched up to a legal type
(such as any-extending the returned extract_vector_elt up to a legal
type) before LegalizeDAG notices.

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

1853

This still creates the illegal type. The node just ends up unused with this patch and gets deleted before anyone notices?

Could we call getSplatSourceVector instead and create an XLenVT extract_vector_elt ourselves? Or teach getSplatValue to obey NewNodesMustHaveLegalTypes.

[frasercrmck]

• move the illegal-type fix into getSplatValue
• create a separate X86 version as it relies on illegal types

[frasercrmck]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

1853

Yeah true, I wasn't too happy with this solution to be honest.

I gave a go at fixing it in getSplatValue. Unfortunately I found that X86 relies on illegal types being returned. What's worse is the X86 backend gets into a loop with this test case when using the type-legal getSplatValue. So for now I added an X86-specific version (the original) to X86ISelLowering.

; RUN: llc < %s -mtriple=i686-unknown -mattr=+sse2

define void @shift1a(<2 x i64> %val, <2 x i64>* %dst, <2 x i64> %sh) nounwind {
entry:
  %shamt = shufflevector <2 x i64> %sh, <2 x i64> undef, <2 x i32> <i32 0, i32 0>
  %shl = shl <2 x i64> %val, %shamt
  store <2 x i64> %shl, <2 x i64>* %dst
  ret void
}

Perhaps this isn't the way forward but I thought I'd throw it up in case people had comments.

The alternative for RISC-V is, as you say, calling getSplatSourceVector ourselves. But it does feel like getSplatValue should obey NewNoesMustHaveLegalTypes.

[frasercrmck]

• restore RISCVISelLowering back to 'main'

[craig.topper]

I think the issue for X86 might be that NewNodesMustHaveLegalTypes is set after type legalization, but vector legalization is allowed to create illegal scalar types. A second type legalization runs after. Maybe we should add a bool parameter to the function that defaults to false, but can be set to true to request a legal type?

[frasercrmck]

• pass bool to getSplatValue for legal types

[frasercrmck]

In D102687#2769104, @craig.topper wrote:

I think the issue for X86 might be that NewNodesMustHaveLegalTypes is set after type legalization, but vector legalization is allowed to create illegal scalar types. A second type legalization runs after. Maybe we should add a bool parameter to the function that defaults to false, but can be set to true to request a legal type?

Yeah that might be going on but I'm not sure. It's pretty finicky. I think X86 is calling this during lowering, which should be after all legalization? I think it's just managing to get the extracted scalar into a legal form before returning from the lowering method.

From a brief look, I believe that WebAssembly's use of this method would crash, but it unconditionally "any exts or truncs" the extracted value to i32, which I think (for their set of illegal types) is always folded into the extract_vector_elt and so the illegal type is quickly removed.

AArch64's use also looks suspicious but since it's for scatter/gather I found it difficult to get an IR test case together which produces the right DAG.

[craig.topper]

In D102687#2770717, @frasercrmck wrote:

In D102687#2769104, @craig.topper wrote:

I think the issue for X86 might be that NewNodesMustHaveLegalTypes is set after type legalization, but vector legalization is allowed to create illegal scalar types. A second type legalization runs after. Maybe we should add a bool parameter to the function that defaults to false, but can be set to true to request a legal type?

Yeah that might be going on but I'm not sure. It's pretty finicky. I think X86 is calling this during lowering, which should be after all legalization? I think it's just managing to get the extracted scalar into a legal form before returning from the lowering method.

There are 2 different lowering stages. One in the LegalizeVectorOps and one in LegalizeDAG. There is a scalar type legalization step in between if needed. It looks like it was being called during shift lowering which would be in LegalizeVectorOps. While our SHUFFLE_VECTOR usage would be in LegalizeDAG.

From a brief look, I believe that WebAssembly's use of this method would crash, but it unconditionally "any exts or truncs" the extracted value to i32, which I think (for their set of illegal types) is always folded into the extract_vector_elt and so the illegal type is quickly removed.

AArch64's use also looks suspicious but since it's for scatter/gather I found it difficult to get an IR test case together which produces the right DAG.

[craig.topper]

LGTM

llvm/lib/Target/RISCV/RISCVISelLowering.h

private:
   SDValue lowerSELECT(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerBRCOND(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerVASTART(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerFRAMEADDR(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerRETURNADDR(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerShiftLeftParts(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerShiftRightParts(SDValue Op, SelectionDAG &DAG, bool IsSRA) const;
   SDValue lowerSPLAT_VECTOR_PARTS(SDValue Op, SelectionDAG &DAG) const;
+  SDValue lowerVECTOR_SHUFFLE(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerVectorMaskSplat(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerVectorMaskExt(SDValue Op, SelectionDAG &DAG,
                              int64_t ExtTrueVal) const;
   SDValue lowerVectorMaskTrunc(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerINSERT_VECTOR_ELT(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerEXTRACT_VECTOR_ELT(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerINTRINSIC_WO_CHAIN(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerINTRINSIC_W_CHAIN(SDValue Op, SelectionDAG &DAG) const;

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

static SDValue lowerScalarSplat(SDValue Scalar, SDValue VL, MVT VT, SDLoc DL,
 
   assert(XLenVT == MVT::i32 && Scalar.getValueType() == MVT::i64 &&
          "Unexpected scalar for splat lowering!");
 
   // Otherwise use the more complicated splatting algorithm.
   return splatSplitI64WithVL(DL, VT, Scalar, VL, DAG);
 }
 
-static SDValue lowerVECTOR_SHUFFLE(SDValue Op, SelectionDAG &DAG,
-                                   const RISCVSubtarget &Subtarget) {
+SDValue RISCVTargetLowering::lowerVECTOR_SHUFFLE(SDValue Op,
+                                                 SelectionDAG &DAG) const {
   SDValue V1 = Op.getOperand(0);
   SDValue V2 = Op.getOperand(1);
   SDLoc DL(Op);
   MVT XLenVT = Subtarget.getXLenVT();
   MVT VT = Op.getSimpleValueType();
   unsigned NumElts = VT.getVectorNumElements();
   ShuffleVectorSDNode *SVN = cast<ShuffleVectorSDNode>(Op.getNode());
 
-  MVT ContainerVT = getContainerForFixedLengthVector(DAG, VT, Subtarget);
+  MVT ContainerVT = getContainerForFixedLengthVector(VT);
 
   SDValue TrueMask, VL;
   std::tie(TrueMask, VL) = getDefaultVLOps(VT, ContainerVT, DL, DAG, Subtarget);
 
   if (SVN->isSplat()) {
     const int Lane = SVN->getSplatIndex();
     if (Lane >= 0) {
       MVT SVT = VT.getVectorElementType();
if (IndexVT.getScalarType().bitsGT(XLenVT)) {
     GatherOpc = RISCVISD::VRGATHEREI16_VV_VL;
     IndexVT = IndexVT.changeVectorElementType(MVT::i16);
   }
 
   MVT IndexContainerVT =
       ContainerVT.changeVectorElementType(IndexVT.getScalarType());
 
   SDValue Gather;
+  // We can't introduce illegal types, so ensure that a splatted element is
+  // either a constant integer (and can be freely extended to XLenVT) or
+  // type-legal.
   // TODO: This doesn't trigger for i64 vectors on RV32, since there we
   // encounter a bitcasted BUILD_VECTOR with low/high i32 values.
-  if (SDValue SplatValue = DAG.getSplatValue(V1)) {
+  SDValue SplatValue = DAG.getSplatValue(V1);
+  if (SplatValue && (isa<ConstantSDNode>(SplatValue) ||
+                     isTypeLegal(SplatValue.getValueType()))) {
     Gather = lowerScalarSplat(SplatValue, VL, ContainerVT, DL, DAG, Subtarget);
   } else {
     SDValue LHSIndices = DAG.getBuildVector(IndexVT, DL, GatherIndicesLHS);
     LHSIndices =
         convertToScalableVector(IndexContainerVT, LHSIndices, DAG, Subtarget);
 
     V1 = convertToScalableVector(ContainerVT, V1, DAG, Subtarget);
     Gather =
case ISD::VECTOR_REVERSE:
     return lowerVECTOR_REVERSE(Op, DAG);
   case ISD::BUILD_VECTOR:
     return lowerBUILD_VECTOR(Op, DAG, Subtarget);
   case ISD::SPLAT_VECTOR:
     if (Op.getValueType().getVectorElementType() == MVT::i1)
       return lowerVectorMaskSplat(Op, DAG);
     return lowerSPLAT_VECTOR(Op, DAG, Subtarget);
   case ISD::VECTOR_SHUFFLE:
-    return lowerVECTOR_SHUFFLE(Op, DAG, Subtarget);
+    return lowerVECTOR_SHUFFLE(Op, DAG);
   case ISD::CONCAT_VECTORS: {
     // Split CONCAT_VECTORS into a series of INSERT_SUBVECTOR nodes. This is
     // better than going through the stack, as the default expansion does.
     SDLoc DL(Op);
     MVT VT = Op.getSimpleValueType();
     unsigned NumOpElts =
         Op.getOperand(0).getSimpleValueType().getVectorMinNumElements();
     SDValue Vec = DAG.getUNDEF(VT);

llvm/test/CodeGen/RISCV/rvv/fixed-vectors-int-shuffles.ll

 ; RV64-NEXT:    vmv.v.i v28, 5
 ; RV64-NEXT:    vsetivli a0, 8, e64,m4,tu,mu
 ; RV64-NEXT:    vrgather.vv v28, v8, v12, v0.t
 ; RV64-NEXT:    vmv4r.v v8, v28
 ; RV64-NEXT:    ret
   %s = shufflevector <8 x i64> %x, <8 x i64> <i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5>, <8 x i32> <i32 0, i32 3, i32 10, i32 9, i32 4, i32 1, i32 7, i32 14>
   ret <8 x i64> %s
 }
+
+define <4 x i8> @interleave_shuffles(<4 x i8> %x) {
+; CHECK-LABEL: interleave_shuffles:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetivli a0, 4, e8,mf4,ta,mu
+; CHECK-NEXT:    vrgather.vi v25, v8, 0
+; CHECK-NEXT:    vrgather.vi v26, v8, 1
+; CHECK-NEXT:    addi a0, zero, 1
+; CHECK-NEXT:    vmv.s.x v27, a0
+; CHECK-NEXT:    vmv.v.i v28, 0
+; CHECK-NEXT:    vsetivli a0, 4, e8,mf4,tu,mu
+; CHECK-NEXT:    vmv1r.v v29, v28
+; CHECK-NEXT:    vslideup.vi v29, v27, 3
+; CHECK-NEXT:    vsetivli a0, 3, e8,mf4,tu,mu
+; CHECK-NEXT:    vslideup.vi v28, v27, 2
+; CHECK-NEXT:    vsetivli a0, 4, e8,mf4,ta,mu
+; CHECK-NEXT:    vrgather.vv v8, v25, v28
+; CHECK-NEXT:    addi a0, zero, 10
+; CHECK-NEXT:    vsetivli a1, 1, e8,mf8,ta,mu
+; CHECK-NEXT:    vmv.s.x v0, a0
+; CHECK-NEXT:    vsetivli a0, 4, e8,mf4,tu,mu
+; CHECK-NEXT:    vrgather.vv v8, v26, v29, v0.t
+; CHECK-NEXT:    ret
+  %y = shufflevector <4 x i8> %x, <4 x i8> undef, <4 x i32> <i32 0, i32 0, i32 0, i32 0>
+  %z = shufflevector <4 x i8> %x, <4 x i8> undef, <4 x i32> <i32 1, i32 1, i32 1, i32 1>
+  %w = shufflevector <4 x i8> %y, <4 x i8> %z, <4 x i32> <i32 0, i32 4, i32 1, i32 5>
+  ret <4 x i8> %w
+}