[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 ↗

(On Diff #346158)

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 ↗

(On Diff #346158)

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/CodeGen/SelectionDAG/SelectionDAG.cpp

SDValue SelectionDAG::getSplatSourceVector(SDValue V, int &SplatIdx) {
   }
   }
 
   return SDValue();
 }
 
 SDValue SelectionDAG::getSplatValue(SDValue V) {
   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 (NewNodesMustHaveLegalTypes && !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/X86/X86ISelLowering.cpp

if (Op.getOpcode() == ISD::SRA) {
       return Res;
     }
     llvm_unreachable("Unknown shift opcode.");
   }
 
   return SDValue();
 }
 
+// A version of SelectionDAG::getSplatValue which may return illegal types,
+// even after type legalization.
+static SDValue getSplatValue(SDValue V, SelectionDAG &DAG) {
+  int SplatIdx;
+  if (SDValue SrcVector = DAG.getSplatSourceVector(V, SplatIdx))
+    return DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SDLoc(V),
+                       SrcVector.getValueType().getScalarType(), SrcVector,
+                       DAG.getVectorIdxConstant(SplatIdx, SDLoc(V)));
+  return SDValue();
+}
+
 static SDValue LowerScalarVariableShift(SDValue Op, SelectionDAG &DAG,
                                         const X86Subtarget &Subtarget) {
   MVT VT = Op.getSimpleValueType();
   SDLoc dl(Op);
   SDValue R = Op.getOperand(0);
   SDValue Amt = Op.getOperand(1);
   unsigned Opcode = Op.getOpcode();
   unsigned X86OpcI = getTargetVShiftUniformOpcode(Opcode, false);
   unsigned X86OpcV = getTargetVShiftUniformOpcode(Opcode, true);
 
-  if (SDValue BaseShAmt = DAG.getSplatValue(Amt)) {
+  if (SDValue BaseShAmt = getSplatValue(Amt, DAG)) {
     if (SupportedVectorShiftWithBaseAmnt(VT, Subtarget, Opcode)) {
       MVT EltVT = VT.getVectorElementType();
       assert(EltVT.bitsLE(MVT::i64) && "Unexpected element type!");
       if (EltVT != MVT::i64 && EltVT.bitsGT(MVT::i32))
         BaseShAmt = DAG.getNode(ISD::ZERO_EXTEND, dl, MVT::i64, BaseShAmt);
       else if (EltVT.bitsLT(MVT::i32))
         BaseShAmt = DAG.getNode(ISD::ZERO_EXTEND, dl, MVT::i32, BaseShAmt);
 

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
+}