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

(On Diff #118239)

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 ↗

(On Diff #118239)

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 ↗

(On Diff #118239)

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 ↗

(On Diff #118239)

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!

llvm/trunk/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

 // 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
+// fold shuffles if the resulting (non-splat) 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 undef elements based on the shuffle mask.
+  bool IsSplat = false;
+  auto *BV0 = dyn_cast<BuildVectorSDNode>(N0);
+  auto *BV1 = dyn_cast<BuildVectorSDNode>(N1);
+  if (BV0 && BV1)
+    if (SDValue Splat0 = BV0->getSplatValue())
+      IsSplat = (Splat0 == BV1->getSplatValue());
+
   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) {
         Op = S.getOperand(Idx);
       } else if (S.getOpcode() == ISD::SCALAR_TO_VECTOR) {
         assert(Idx == 0 && "Unexpected SCALAR_TO_VECTOR operand index.");
         Op = S.getOperand(0);
       } else {
         // Operand can't be combined - bail out.
         return SDValue();
       }
     }
 
-    // Don't duplicate a non-constant BUILD_VECTOR operand; semantically, this is
-    // fine, but it's likely to generate low-quality code if the target can't
-    // reconstruct an appropriate shuffle.
+    // Don't duplicate a non-constant BUILD_VECTOR operand unless we're
+    // generating a splat; semantically, this is fine, but it's likely to
+    // generate low-quality code if the target can't reconstruct an appropriate
+    // shuffle.
     if (!Op.isUndef() && !isa<ConstantSDNode>(Op) && !isa<ConstantFPSDNode>(Op))
-      if (!DuplicateOps.insert(Op).second)
+      if (!IsSplat && !DuplicateOps.insert(Op).second)
         return SDValue();
 
     Ops.push_back(Op);
   }
 
   // BUILD_VECTOR requires all inputs to be of the same type, find the
   // maximum type and extend them all.
   EVT SVT = VT.getScalarType();

llvm/trunk/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
 }