[DAGCombiner][WebAssembly] Combine shuffles of more splat vals

[tlively]

Summary

combineShuffleOfSplatVal previously only combined shuffles of splat
values that were themselves shuffles. This patch generalizes the
combine to also combine away shuffles of arbitrary splat values
recognized by SelectionDAG::isSplatValue, as long as doing so does not
create any new undefined lanes.

On the WebAssembly side, this patch also introduces a new custom
combine to remove undefined lanes from splatting
build_vectors. Without this extra combine, the new generic shuffle
combine would be inhibited on interesting cases such as the
shl_abs_add function in simd-shift-complex-splats.ll because it would
expose the undefined lanes.

Depends on D83605.

[tlively]

• Remove TODO

[tlively]

• Update another x86 test

[aheejin]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

19841

We should fix this comment too now that we allow non-shuffles in side

19852

The previous version seems to allow undef elements in splats, but this forces all elements to be demanded. Is this OK?

19902

It looks like this method deals with a simpler form in which both inner and outer masks are splats, but it seems this has the same limitation that only shuffles are considered. Can this be improved too? (Not necessarily in this patch though)

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

1749

The same question as above. The previous version of combineShuffleOfSplatVal seems to allow undef lanes from splats. The previous version calls ShuffleVectorSDNode::isSplat, which calls ShuffleVectorSDNode::isSplatMask, which seems to allow undefs. Also the comments there seem to take undef elements in the splat mask in consideration. Do we need to, or, it safe to change it to disallow it? If we allow undef splats in combineShuffleOfSplatVal, we don't need this separate target combine function, right?

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

static SDValue combineTruncationShuffle(ShuffleVectorSDNode *SVN,
   // the same scale as the extension.
   if (isTruncate(ExtScale))
     return DAG.getBitcast(VT, N00);
 
   return SDValue();
 }
 
 // Combine shuffles of splat-shuffles of the form:
 // shuffle (shuffle V, undef, splat-mask), undef, M
 // If splat-mask contains undef elements, we need to be careful about
 // introducing undef's in the folded mask which are not the result of composing
 // the masks of the shuffles.
 static SDValue combineShuffleOfSplatVal(ShuffleVectorSDNode *Shuf,
                                         SelectionDAG &DAG) {
   if (!Shuf->getOperand(1).isUndef())
     return SDValue();
-  auto *Splat = dyn_cast<ShuffleVectorSDNode>(Shuf->getOperand(0));
-  if (!Splat || !Splat->isSplat())
+
+  APInt UndefElts;
+  APInt DemandedElts =
+      APInt::getAllOnesValue(Shuf->getValueType(0).getVectorNumElements());
+  if (!DAG.isSplatValue(Shuf->getOperand(0), DemandedElts, UndefElts))
     return SDValue();
 
-  ArrayRef<int> ShufMask = Shuf->getMask();
-  ArrayRef<int> SplatMask = Splat->getMask();
-  assert(ShufMask.size() == SplatMask.size() && "Mask length mismatch");
+  ArrayRef<int> UserMask = Shuf->getMask();
 
   // Prefer simplifying to the splat-shuffle, if possible. This is legal if
   // every undef mask element in the splat-shuffle has a corresponding undef
   // element in the user-shuffle's mask or if the composition of mask elements
   // would result in undef.
   // Examples for (shuffle (shuffle v, undef, SplatMask), undef, UserMask):
   // * UserMask=[0,2,u,u], SplatMask=[2,u,2,u] -> [2,2,u,u]
   //   In this case it is not legal to simplify to the splat-shuffle because we
   //   may be exposing the users of the shuffle an undef element at index 1
   //   which was not there before the combine.
   // * UserMask=[0,u,2,u], SplatMask=[2,u,2,u] -> [2,u,2,u]
   //   In this case the composition of masks yields SplatMask, so it's ok to
   //   simplify to the splat-shuffle.
   // * UserMask=[3,u,2,u], SplatMask=[2,u,2,u] -> [u,u,2,u]
   //   In this case the composed mask includes all undef elements of SplatMask
   //   and in addition sets element zero to undef. It is safe to simplify to
   //   the splat-shuffle.
-  auto CanSimplifyToExistingSplat = [](ArrayRef<int> UserMask,
-                                       ArrayRef<int> SplatMask) {
-    for (unsigned i = 0, e = UserMask.size(); i != e; ++i)
-      if (UserMask[i] != -1 && SplatMask[i] == -1 &&
-          SplatMask[UserMask[i]] != -1)
-        return false;
-    return true;
-  };
-  if (CanSimplifyToExistingSplat(ShufMask, SplatMask))
+  bool CanSimplifyToExistingSplat = true;
+    if (UserMask[i] != -1 && UndefElts[i] && !UndefElts[UserMask[i]])
+      CanSimplifyToExistingSplat = false;
+
+  if (CanSimplifyToExistingSplat)
     return Shuf->getOperand(0);
 
+  // Unless the splat value is a shuffle as well, there is nothing more to do
+  auto *Splat = dyn_cast<ShuffleVectorSDNode>(Shuf->getOperand(0));
+  if (!Splat)
+    return SDValue();
+
   // Create a new shuffle with a mask that is composed of the two shuffles'
   // masks.
+  ArrayRef<int> SplatMask = Splat->getMask();
   SmallVector<int, 32> NewMask;
-  for (int Idx : ShufMask)
+  for (int Idx : UserMask)
     NewMask.push_back(Idx == -1 ? -1 : SplatMask[Idx]);
 
   return DAG.getVectorShuffle(Splat->getValueType(0), SDLoc(Splat),
                               Splat->getOperand(0), Splat->getOperand(1),
                               NewMask);
 }
 
 /// Combine shuffle of shuffle of the form:
 /// shuf (shuf X, undef, InnerMask), undef, OuterMask --> splat X
 static SDValue formSplatFromShuffles(ShuffleVectorSDNode *OuterShuf,
                                      SelectionDAG &DAG) {
   if (!OuterShuf->getOperand(1).isUndef())
     return SDValue();
   auto *InnerShuf = dyn_cast<ShuffleVectorSDNode>(OuterShuf->getOperand(0));
   if (!InnerShuf || !InnerShuf->getOperand(1).isUndef())
     return SDValue();
 
   ArrayRef<int> OuterMask = OuterShuf->getMask();
   ArrayRef<int> InnerMask = InnerShuf->getMask();

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

if (Subtarget->hasSIMD128())
         setOperationAction(Op, T, Expand);
   }
 
   // SIMD-specific configuration
   if (Subtarget->hasSIMD128()) {
     // Hoist bitcasts out of shuffles
     setTargetDAGCombine(ISD::VECTOR_SHUFFLE);
 
+    // Remove undef lanes from splats
+    setTargetDAGCombine(ISD::BUILD_VECTOR);
+
     // Scalarize extracts of abs splats
     setTargetDAGCombine(ISD::EXTRACT_VECTOR_ELT);
 
     // Support saturating add for i8x16 and i16x8
     for (auto Op : {ISD::SADDSAT, ISD::UADDSAT})
       for (auto T : {MVT::v16i8, MVT::v8i16})
         setOperationAction(Op, T, Legal);
 
if (!SrcType.is128BitVector() ||
       SrcType.getVectorNumElements() != DstType.getVectorNumElements())
     return SDValue();
   SDValue NewShuffle = DAG.getVectorShuffle(
       SrcType, SDLoc(N), CastOp, DAG.getUNDEF(SrcType), Shuffle->getMask());
   return DAG.getBitcast(DstType, NewShuffle);
 }
 
 static SDValue
+performBUILD_VECTORCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI) {
+  auto &DAG = DCI.DAG;
+  auto *Build = cast<BuildVectorSDNode>(N);
+
+  // Remove undef lanes from splats. They don't allow us to make any extra
+  // optimizations and they can inhibit splat scalarization combines.
+  BitVector UndefElts;
+  if (SDValue SplatVal = Build->getSplatValue(&UndefElts))
+    if (UndefElts.any())
+      return DAG.getSplatBuildVector(N->getValueType(0), SDLoc(N), SplatVal);
+
+  return SDValue();
+}
+
+static SDValue
 performEXTRACT_VECTOR_ELTCombine(SDNode *N,
                                  TargetLowering::DAGCombinerInfo &DCI) {
   auto &DAG = DCI.DAG;
   SDLoc DL(N);
   EVT ScalarVT = N->getValueType(0);
   SDValue VecOp = N->getOperand(0);
   SDValue IndexOp = N->getOperand(1);
 

 SDValue
 WebAssemblyTargetLowering::PerformDAGCombine(SDNode *N,
                                              DAGCombinerInfo &DCI) const {
   switch (N->getOpcode()) {
   default:
     return SDValue();
   case ISD::VECTOR_SHUFFLE:
     return performVECTOR_SHUFFLECombine(N, DCI);
+  case ISD::BUILD_VECTOR:
+    return performBUILD_VECTORCombine(N, DCI);
   case ISD::EXTRACT_VECTOR_ELT:
     return performEXTRACT_VECTOR_ELTCombine(N, DCI);
   }
 }

llvm/test/CodeGen/WebAssembly/simd-build-vector.ll

define <8 x i16> @splat_common_arg_i16x8(i16 %a, i16 %b, i16 %c) {
   %v3 = insertelement <8 x i16> %v2, i16 %c, i32 3
   %v4 = insertelement <8 x i16> %v3, i16 %b, i32 4
   %v5 = insertelement <8 x i16> %v4, i16 %c, i32 5
   %v6 = insertelement <8 x i16> %v5, i16 %c, i32 6
   %v7 = insertelement <8 x i16> %v6, i16 %b, i32 7
   ret <8 x i16> %v7
 }
 
+;; TODO: This should be a swizzle, but will need a custom combine to
+;; preempt the combine that removes undef lanes from splat
+;; build_vectors, since swizzle lowering depends on those lanes being
+;; undef.
+
 ; CHECK-LABEL: swizzle_one_i8x16:
 ; CHECK-NEXT:  .functype       swizzle_one_i8x16 (v128, v128) -> (v128)
-; CHECK-NEXT:  v8x16.swizzle   $push[[L0:[0-9]+]]=, $0, $1
-; CHECK-NEXT:  return          $pop[[L0]]
+; CHECK-NOT:  v8x16.swizzle
 define <16 x i8> @swizzle_one_i8x16(<16 x i8> %src, <16 x i8> %mask) {
   %m0 = extractelement <16 x i8> %mask, i32 0
   %s0 = extractelement <16 x i8> %src, i8 %m0
   %v0 = insertelement <16 x i8> undef, i8 %s0, i32 0
   ret <16 x i8> %v0
 }
 
+; CHECK-LABEL: swizzle_two_i8x16:
+; CHECK-NEXT:  .functype       swizzle_two_i8x16 (v128, v128) -> (v128)
+; CHECK-NEXT:  v8x16.swizzle
+; CHECK-NEXT:  return          $pop[[L0]]
+define <16 x i8> @swizzle_two_i8x16(<16 x i8> %src, <16 x i8> %mask) {
+  %m0 = extractelement <16 x i8> %mask, i32 0
+  %s0 = extractelement <16 x i8> %src, i8 %m0
+  %v0 = insertelement <16 x i8> undef, i8 %s0, i32 0
+  %m1 = extractelement <16 x i8> %mask, i32 1
+  %s1 = extractelement <16 x i8> %src, i8 %m1
+  %v1 = insertelement <16 x i8> %v0, i8 %s1, i32 1
+  ret <16 x i8> %v1
+}
+
 ; CHECK-LABEL: swizzle_all_i8x16:
 ; CHECK-NEXT:  .functype       swizzle_all_i8x16 (v128, v128) -> (v128)
 ; CHECK-NEXT:  v8x16.swizzle   $push[[L0:[0-9]+]]=, $0, $1
 ; CHECK-NEXT:  return          $pop[[L0]]
 define <16 x i8> @swizzle_all_i8x16(<16 x i8> %src, <16 x i8> %mask) {
   %m0 = extractelement <16 x i8> %mask, i32 0
   %s0 = extractelement <16 x i8> %src, i8 %m0
   %v0 = insertelement <16 x i8> undef, i8 %s0, i32 0
define <16 x i8> @mashup_splat_i8x16(<16 x i8> %src, <16 x i8> %mask, i8 %splatted) {
   %v2 = insertelement <16 x i8> %v1, i8 %splatted, i32 12
   ; const 4
   %v3 = insertelement <16 x i8> %v2, i8 42, i32 4
   ret <16 x i8> %v3
 }
 
 ; CHECK-LABEL: undef_const_insert_f32x4:
 ; CHECK-NEXT:  .functype       undef_const_insert_f32x4 () -> (v128)
-; UNIMP-NEXT:  v128.const      $push[[L0:[0-9]+]]=, 0x0p0, 0x1.5p5, 0x0p0, 0x0p0
+; UNIMP-NEXT:  v128.const      $push[[L0:[0-9]+]]=, 0x1.5p5, 0x1.5p5, 0x1.5p5, 0x1.5p5
 ; UNIMP-NEXT:  return          $pop[[L0]]
 ; SIMD-VM: f32x4.splat
 define <4 x float> @undef_const_insert_f32x4() {
   %v = insertelement <4 x float> undef, float 42., i32 1
   ret <4 x float> %v
 }
 
 ; CHECK-LABEL: undef_arg_insert_i32x4:

llvm/test/CodeGen/WebAssembly/simd-shift-complex-splats.ll

define <16 x i8> @shl_abs(<16 x i8> %v, i8 %a) {
   %r = shl <16 x i8> %v, %shift
   ret <16 x i8> %r
 }
 
 ;; TODO: Complete scalarization by removing the unnecessary shuffle
 
 ; CHECK-LABEL: shl_abs_add:
 ; CHECK-NEXT: .functype shl_abs_add (v128, i32, i32) -> (v128)
-; CHECK-NEXT: i8x16.splat $push1=, $1
-; CHECK-NEXT: i8x16.splat $push0=, $2
-; CHECK-NEXT: i8x16.add $push2=, $pop1, $pop0
-; CHECK-NEXT: v8x16.shuffle $push3=, $pop2, $0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
-; CHECK-NEXT: i8x16.extract_lane_u $push11=, $pop3, 0
-; CHECK-NEXT: local.tee $push10=, $1=, $pop11
-; CHECK-NEXT: i32.const $push4=, 31
-; CHECK-NEXT: i32.shr_s $push9=, $1, $pop4
-; CHECK-NEXT: local.tee $push8=, $1=, $pop9
-; CHECK-NEXT: i32.add $push5=, $pop10, $pop8
-; CHECK-NEXT: i32.xor $push6=, $pop5, $1
-; CHECK-NEXT: i8x16.shl $push7=, $0, $pop6
-; CHECK-NEXT: return $pop7
+; CHECK-NEXT: i32.add $push7=, $1, $2
+; CHECK-NEXT: local.tee $push6=, $1=, $pop7
+; CHECK-NEXT: i32.const $push0=, 31
+; CHECK-NEXT: i32.shr_s $push5=, $1, $pop0
+; CHECK-NEXT: local.tee $push4=, $1=, $pop5
+; CHECK-NEXT: i32.add $push1=, $pop6, $pop4
+; CHECK-NEXT: i32.xor $push2=, $pop1, $1
+; CHECK-NEXT: i8x16.shl $push3=, $0, $pop2
+; CHECK-NEXT: return  $pop3
 define <16 x i8> @shl_abs_add(<16 x i8> %v, i8 %a, i8 %b) {
   %t1 = insertelement <16 x i8> undef, i8 %a, i32 0
   %va = shufflevector <16 x i8> %t1, <16 x i8> undef, <16 x i32> zeroinitializer
   %t2 = insertelement <16 x i8> undef, i8 %b, i32 0
   %vb = shufflevector <16 x i8> %t2, <16 x i8> undef, <16 x i32> zeroinitializer
   %vadd = add <16 x i8> %va, %vb
   %nvadd = sub <16 x i8> zeroinitializer, %vadd
   %c = icmp sgt <16 x i8> %vadd, zeroinitializer
   %shift = select <16 x i1> %c, <16 x i8> %vadd, <16 x i8> %nvadd
   %r = shl <16 x i8> %v, %shift
   ret <16 x i8> %r
 }

llvm/test/CodeGen/X86/vector-fshr-rot-256.ll

 ;
 ; Uniform Variable Shifts
 ;
 
 define <4 x i64> @splatvar_funnnel_v4i64(<4 x i64> %x, <4 x i64> %amt) nounwind {
 ; AVX1-LABEL: splatvar_funnnel_v4i64:
 ; AVX1:       # %bb.0:
 ; AVX1-NEXT:    vpshufd {{.*#+}} xmm2 = xmm1[0,1,0,1]
-; AVX1-NEXT:    vpxor %xmm3, %xmm3, %xmm3
-; AVX1-NEXT:    vpsubq %xmm2, %xmm3, %xmm2
 ; AVX1-NEXT:    vmovdqa {{.*#+}} xmm3 = [63,63]
+; AVX1-NEXT:    vpand %xmm3, %xmm1, %xmm4
+; AVX1-NEXT:    vextractf128 $1, %ymm0, %xmm5
+; AVX1-NEXT:    vpsrlq %xmm4, %xmm5, %xmm6
+; AVX1-NEXT:    vpsrlq %xmm4, %xmm0, %xmm4
+; AVX1-NEXT:    vinsertf128 $1, %xmm6, %ymm4, %ymm4
+; AVX1-NEXT:    vpxor %xmm6, %xmm6, %xmm6
+; AVX1-NEXT:    vpsubq %xmm2, %xmm6, %xmm2
 ; AVX1-NEXT:    vpand %xmm3, %xmm2, %xmm2
-; AVX1-NEXT:    vextractf128 $1, %ymm0, %xmm4
-; AVX1-NEXT:    vpsllq %xmm2, %xmm4, %xmm5
-; AVX1-NEXT:    vpshufd {{.*#+}} xmm6 = xmm2[2,3,0,1]
-; AVX1-NEXT:    vpsllq %xmm6, %xmm4, %xmm7
-; AVX1-NEXT:    vpblendw {{.*#+}} xmm5 = xmm5[0,1,2,3],xmm7[4,5,6,7]
-; AVX1-NEXT:    vpsllq %xmm2, %xmm0, %xmm2
-; AVX1-NEXT:    vpsllq %xmm6, %xmm0, %xmm6
-; AVX1-NEXT:    vpblendw {{.*#+}} xmm2 = xmm2[0,1,2,3],xmm6[4,5,6,7]
-; AVX1-NEXT:    vinsertf128 $1, %xmm5, %ymm2, %ymm2
+; AVX1-NEXT:    vpsllq %xmm2, %xmm5, %xmm2
+; AVX1-NEXT:    vpsubq %xmm1, %xmm6, %xmm1
 ; AVX1-NEXT:    vpand %xmm3, %xmm1, %xmm1
-; AVX1-NEXT:    vpsrlq %xmm1, %xmm4, %xmm3
-; AVX1-NEXT:    vpsrlq %xmm1, %xmm0, %xmm0
-; AVX1-NEXT:    vinsertf128 $1, %xmm3, %ymm0, %ymm0
-; AVX1-NEXT:    vorps %ymm0, %ymm2, %ymm0
+; AVX1-NEXT:    vpsllq %xmm1, %xmm0, %xmm0
+; AVX1-NEXT:    vinsertf128 $1, %xmm2, %ymm0, %ymm0
+; AVX1-NEXT:    vorps %ymm4, %ymm0, %ymm0
 ; AVX1-NEXT:    retq
 ;
 ; AVX2-LABEL: splatvar_funnnel_v4i64:
 ; AVX2:       # %bb.0:
 ; AVX2-NEXT:    vmovdqa {{.*#+}} xmm2 = [63,63]
 ; AVX2-NEXT:    vpand %xmm2, %xmm1, %xmm3
 ; AVX2-NEXT:    vpsrlq %xmm3, %ymm0, %ymm3
 ; AVX2-NEXT:    vpxor %xmm4, %xmm4, %xmm4