[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

1852

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

1852

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/include/llvm/CodeGen/SelectionDAG.h

public:
 
   /// Test whether \p V has a splatted value.
   bool isSplatValue(SDValue V, bool AllowUndefs = false);
 
   /// If V is a splatted value, return the source vector and its splat index.
   SDValue getSplatSourceVector(SDValue V, int &SplatIndex);
 
   /// If V is a splat vector, return its scalar source operand by extracting
-  /// that element from the source vector.
-  SDValue getSplatValue(SDValue V);
+  /// that element from the source vector. If LegalTypes is true, this method
+  /// may only return a legally-typed splat value. If it cannot legalize the
+  /// splatted value it will return SDValue().
+  SDValue getSplatValue(SDValue V, bool LegalTypes = false);
 
   /// If a SHL/SRA/SRL node \p V has a constant or splat constant shift amount
   /// that is less than the element bit-width of the shift node, return it.
   const APInt *getValidShiftAmountConstant(SDValue V,
                                            const APInt &DemandedElts) const;
 
   /// If a SHL/SRA/SRL node \p V has constant shift amounts that are all less
   /// than the element bit-width of the shift node, return the minimum value.

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

case ISD::VECTOR_SHUFFLE: {
     SplatIdx = Idx % NumElts;
     return V.getOperand(Idx / NumElts);
   }
   }
 
   return SDValue();
 }
 
-SDValue SelectionDAG::getSplatValue(SDValue V) {
+SDValue SelectionDAG::getSplatValue(SDValue V, bool LegalTypes) {
   int SplatIdx;
-  if (SDValue SrcVector = getSplatSourceVector(V, SplatIdx))
-    return getNode(ISD::EXTRACT_VECTOR_ELT, SDLoc(V),
-                   SrcVector.getValueType().getScalarType(), SrcVector,
+  if (SDValue SrcVector = getSplatSourceVector(V, SplatIdx)) {
+    EVT SVT = SrcVector.getValueType().getScalarType();
+    EVT LegalSVT = SVT;
+    if (LegalTypes && !TLI->isTypeLegal(SVT)) {
+      if (!SVT.isInteger())
+        return SDValue();
+      LegalSVT = TLI->getTypeToTransformTo(*getContext(), LegalSVT);
+      if (LegalSVT.bitsLT(SVT))
+        return SDValue();
+    }
+    return getNode(ISD::EXTRACT_VECTOR_ELT, SDLoc(V), LegalSVT, SrcVector,
                    getVectorIdxConstant(SplatIdx, SDLoc(V)));
+  }
   return SDValue();
 }
 
 const APInt *
 SelectionDAG::getValidShiftAmountConstant(SDValue V,
                                           const APInt &DemandedElts) const {
   assert((V.getOpcode() == ISD::SHL || V.getOpcode() == ISD::SRL ||
           V.getOpcode() == ISD::SRA) &&

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

static SDValue lowerVECTOR_SHUFFLE(SDValue Op, SelectionDAG &DAG,
   }
 
   MVT IndexContainerVT =
       ContainerVT.changeVectorElementType(IndexVT.getScalarType());
 
   SDValue Gather;
   // 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)) {
+  if (SDValue SplatValue = DAG.getSplatValue(V1, /*LegalTypes*/ true)) {
     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 =

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:    vsetvli zero, zero, e8,mf4,ta,mu
+; CHECK-NEXT:    vmv.x.s a0, v8
+; CHECK-NEXT:    vsetivli a1, 4, e8,mf4,ta,mu
+; CHECK-NEXT:    vrgather.vi v25, v8, 1
+; CHECK-NEXT:    addi a1, zero, 1
+; CHECK-NEXT:    vmv.s.x v26, a1
+; CHECK-NEXT:    vmv.v.i v27, 0
+; CHECK-NEXT:    vsetivli a1, 4, e8,mf4,tu,mu
+; CHECK-NEXT:    vslideup.vi v27, v26, 3
+; CHECK-NEXT:    addi a1, zero, 10
+; CHECK-NEXT:    vsetivli a2, 1, e8,mf8,ta,mu
+; CHECK-NEXT:    vmv.s.x v0, a1
+; CHECK-NEXT:    vsetivli a1, 4, e8,mf4,ta,mu
+; CHECK-NEXT:    vmv.v.x v8, a0
+; CHECK-NEXT:    vsetivli a0, 4, e8,mf4,tu,mu
+; CHECK-NEXT:    vrgather.vv v8, v25, v27, 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
+}