[DAGCombine] Permit combining of shuffle of equivalent splat BUILD_VECTORs

[RKSimon]

Summary

combineShuffleOfScalars is very conservative about shuffled BUILD_VECTOsR that can be combined together.

I'd like to add one additional case - if both BUILD_VECTORs represent splats of the same scalar value but with different UNDEF elements, then we should create a single splat BUILD_VECTOR, sharing only the common UNDEF elements.

[spatel]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15492–15493

Is this more conservative than necessary? If the shuffle mask is choosing an undef lane from either input, then the output would still be undef?

[RKSimon]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15492–15493

Agreed, I started trying to integrate this into the existing code below but I gave up as it was getting very convoluted. I'll update the code so that it just finds a single splat value.

[andreadb]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15467–15470

This comment can be updated.
Your change would allow duplicate non-constant values in the resulting build vector if shuffle operands are both splat of the same value.

15494–15500

Instead of conservatively merging common undef elements, have you considered the possibility of reusing the loop at line 15512?

You can use the code from your patch at lines [15494,15500] to match your particular shuffle(splat,splat) pattern.
Then you can slightly modify the check at line 15532 , so that the restriction on non-constant duplicate values is not applied to your particular case.

[RKSimon]

Updated to reuse more of the existing BUILD_VECTOR code.

[RKSimon]

ping?

[andreadb]

LGTM.
Thanks Simon!

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

 // a simplification in some sense, but it isn't appropriate in general: some
 // BUILD_VECTORs are substantially cheaper than others. The general case
 // of a BUILD_VECTOR requires inserting each element individually (or
 // performing the equivalent in a temporary stack variable). A BUILD_VECTOR of
 // all constants is a single constant pool load.  A BUILD_VECTOR where each
 // element is identical is a splat.  A BUILD_VECTOR where most of the operands
 // are undef lowers to a small number of element insertions.
 //
 // To deal with this, we currently use a bunch of mostly arbitrary heuristics.
 // We don't fold shuffles where one side is a non-zero constant, and we don't
 // fold shuffles if the resulting BUILD_VECTOR would have duplicate
 // non-constant operands. This seems to work out reasonably well in practice.
 static SDValue combineShuffleOfScalars(ShuffleVectorSDNode *SVN,
                                        SelectionDAG &DAG,
                                        const TargetLowering &TLI) {
   EVT VT = SVN->getValueType(0);
   unsigned NumElts = VT.getVectorNumElements();
   SDValue N0 = SVN->getOperand(0);
   SDValue N1 = SVN->getOperand(1);
 
   if (!N0->hasOneUse() || !N1->hasOneUse())
     return SDValue();
   // If only one of N1,N2 is constant, bail out if it is not ALL_ZEROS as
   // discussed above.
   if (!N1.isUndef()) {
     bool N0AnyConst = isAnyConstantBuildVector(N0.getNode());
     bool N1AnyConst = isAnyConstantBuildVector(N1.getNode());
     if (N0AnyConst && !N1AnyConst && !ISD::isBuildVectorAllZeros(N0.getNode()))
       return SDValue();
     if (!N0AnyConst && N1AnyConst && !ISD::isBuildVectorAllZeros(N1.getNode()))
       return SDValue();
   }
 
+  // If both inputs are splats of the same value then we can safely merge this
+  // to a single BUILD_VECTOR with only the common undef elements.
+  auto *BV0 = dyn_cast<BuildVectorSDNode>(N0);
+  auto *BV1 = dyn_cast<BuildVectorSDNode>(N1);
+  if (BV0 && BV1) {
+    BitVector Undefs0, Undefs1;
+    SDValue Splat0 = BV0->getSplatValue(&Undefs0);
+    SDValue Splat1 = BV1->getSplatValue(&Undefs1);
+    if (Splat0 && Splat0 == Splat1) {
+      Undefs0 &= Undefs1;
+      SmallVector<SDValue, 8> Ops(NumElts, Splat0);
+      for (unsigned i = 0; i != NumElts; ++i)
+        if (Undefs0[0])
+          Ops[i] = DAG.getUNDEF(Splat0.getValueType());
+      return DAG.getBuildVector(VT, SDLoc(SVN), Ops);
+    }
+  }
+
   SmallVector<SDValue, 8> Ops;
   SmallSet<SDValue, 16> DuplicateOps;
   for (int M : SVN->getMask()) {
     SDValue Op = DAG.getUNDEF(VT.getScalarType());
     if (M >= 0) {
       int Idx = M < (int)NumElts ? M : M - NumElts;
       SDValue &S = (M < (int)NumElts ? N0 : N1);
       if (S.getOpcode() == ISD::BUILD_VECTOR) {

test/CodeGen/X86/avx-vbroadcast.ll

; X64-NEXT:    retq
   %3 = shufflevector <2 x double> %2, <2 x double> undef, <4 x i32> <i32 1, i32 0, i32 0, i32 0>
   ret <4 x double> %3
 }
 
 define <4 x double> @broadcast_shuffle1032(double* %p) {
 ; X32-LABEL: broadcast_shuffle1032:
 ; X32:       ## BB#0:
 ; X32-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; X32-NEXT:    vmovddup {{.*#+}} xmm0 = mem[0,0]
-; X32-NEXT:    vinsertf128 $1, %xmm0, %ymm0, %ymm0
+; X32-NEXT:    vbroadcastsd (%eax), %ymm0
 ; X32-NEXT:    retl
 ;
 ; X64-LABEL: broadcast_shuffle1032:
 ; X64:       ## BB#0:
-; X64-NEXT:    vmovddup {{.*#+}} xmm0 = mem[0,0]
-; X64-NEXT:    vinsertf128 $1, %xmm0, %ymm0, %ymm0
+; X64-NEXT:    vbroadcastsd (%rdi), %ymm0
 ; X64-NEXT:    retq
   %1 = load double, double* %p
   %2 = insertelement <2 x double> undef, double %1, i32 1
   %3 = insertelement <2 x double> undef, double %1, i32 0
   %4 = shufflevector <2 x double> %2, <2 x double> %3, <4 x i32> <i32 1, i32 0, i32 3, i32 2>
   ret <4 x double> %4
 }