DAGCombine: Combine shuffles of splat-shuffles

[zvi]

Summary

Reapply r299047, but this time handle correctly splat-masks with undef elements.

[zvi]

test/CodeGen/X86/shuffle-of-splat-multiuses.ll

77 ↗

(On Diff #94895)

These two moves seem like a result of suboptimal scheduling + register allocation choices.

[spatel]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

14622–14624 ↗

(On Diff #94895)

Do you think we'll need to distinguish between the "strict" splat and splat-with-undef cases anywhere else? I'm wondering if it would be better to add a bool "HasAnyUndefs" param to ShuffleVectorSDNode::isSplat(). Even if we don't have another use case at the moment, that would at least make the API more like BuildVectorSDNode::isConstantSplat().

[zvi]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

14622–14624 ↗

(On Diff #94895)

Not sure about that, i can try to survey the code for such case.
I could find only one real use of "HasAnyUndefs" for the build_vector API in PPCISelLowering.cpp.

test/CodeGen/X86/shuffle-of-splat-multiuses.ll

77 ↗

(On Diff #94895)

Opened Bug 32719 to discuss the extra copy issue

[zvi]

ping

[spatel]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

14635–14642 ↗

(On Diff #95946)

I'm scared that I'm not seeing all of the undef possibilities here (since I definitely missed them the first time!).

How about writing this to make the options more visible (not sure if it's more or less efficient), but I'd find this easier to read (warning: untested code):

// If the splat mask has no undefs, then return it as-is. 
// It doesn't matter if the user mask has undefs. We choose the splat mask element because it's efficient to return an existing splat node.
ArrayRef<int> SplatMask = Splat->getMask();
if (find(SplatMask, -1) == SplatMask.end())
  return SDValue(Splat, 0);

// The splat mask has undefs. If the user shuffle does not choose undef elements where the splat does not have undef elements, return the existing splat node.
if (all_of(UserMask, [](int Idx) { return Idx == -1 || Splat->getMaskElt(Idx) == Splat->getSplatIndex(); }))
  return SDValue(Splat, 0); 
  
// The splat mask has undefs, and the user shuffle chooses undef elements where the splat has defined elements. Create a more liberal splat mask (one with more undefs).
 for (unsigned i = 0; i < NumElts; ++i)
  // No need for "IsSameMask" here because we already handled any case where we wanted to return the existing node.

[zvi]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

14635–14642 ↗

(On Diff #95946)

if (all_of(UserMask, [](int Idx) { return Idx == -1 || Splat->getMaskElt(Idx) == Splat->getSplatIndex(); }))
  return SDValue(Splat, 0);

This is not correct. For example UserMask = [0,2,u,u], SplatMask=[2,u,2,u]. It would be incorrect to return 'Splat' because it would result in an undef at element index=1 which the composition of masks does not yield.
But i will try to rewrite this code to follow the spirit of your suggestion: first try to identify cases where we can return the splat-shuffle itself, if not, build a new shuffle. Thanks!

[zvi]

Restructure in spirit of Sanjay's proposal: first try to identify cases where we can return the splat-shuffle itself, if not, build a new shuffle mask and create a new shuffle node.

[spatel]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

14880 ↗

(On Diff #98157)

typo: mismatxh

14886–14892 ↗

(On Diff #98157)

I'm a coding caveman, so I had not seen this syntax before. I know it takes a line or 2 more, but I think it's easier to read when you give the lambda a name, so something like:

auto canSimplifyToExistingSplat = [UserMask, SplatMask]() { ... }
if (canSimplifyToExistingSplat())
  return SDValue(Splat, 0);

14887–14890 ↗

(On Diff #98157)

Given that I got this wrong twice, I'd like to see some mask examples for various undef cases in the code comments. I think the tests cover these, so it should just be a matter of listing those cases here similar to the explanatory note in this review:

UserMask = [0,2,u,u], SplatMask=[2,u,2,u] --> ?

Bonus points for adding similar comments in the test file to describe the undef handling there too.

[zvi]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

14880 ↗

(On Diff #98157)

Thanks

14886–14892 ↗

(On Diff #98157)

Ok

14887–14890 ↗

(On Diff #98157)

I agree this will improve the readability of the code and point-out this tricky subtlety.

[zvi]

Update according to Sanjay's comments. I updated the comments with some representative examples, but gave up on annotating the tests since after Isel it's hard to see from the text which elements were undef as we choose arbitrary values for undef elements.

[spatel]

LGTM. Thanks for adding the examples in the code comments.

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

SDValue combineTruncationShuffle(ShuffleVectorSDNode *SVN, SelectionDAG &DAG) {
   // power-of-2 truncations as they are the most likely.
   for (unsigned Scale = 2; Scale < NumElts; Scale *= 2)
     if (isTruncate(Scale))
       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 combineShuffleOfSplat(ArrayRef<int> UserMask,
+                                     ShuffleVectorSDNode *Splat,
+                                     SelectionDAG &DAG) {
+  ArrayRef<int> SplatMask = Splat->getMask();
+  assert(UserMask.size() == SplatMask.size() && "Mask length mismatch");
+
+  // 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(UserMask, SplatMask))
+    return SDValue(Splat, 0);
+
+  // Create a new shuffle with a mask that is composed of the two shuffles'
+  // masks.
+  SmallVector<int, 32> NewMask;
+  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);
+}
+
 SDValue DAGCombiner::visitVECTOR_SHUFFLE(SDNode *N) {
   EVT VT = N->getValueType(0);
   unsigned NumElts = VT.getVectorNumElements();
 
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
 
   assert(N0.getValueType() == VT && "Vector shuffle must be normalized in DAG");
for (unsigned i = 0; i != NumElts; ++i) {
         Changed = true;
       }
       NewMask.push_back(Idx);
     }
     if (Changed)
       return DAG.getVectorShuffle(VT, SDLoc(N), N0, N1, NewMask);
   }
 
+  // A shuffle of a single vector that is a splat can always be folded.
+  if (auto *N0Shuf = dyn_cast<ShuffleVectorSDNode>(N0))
+    if (N1->isUndef() && N0Shuf->isSplat())
+      return combineShuffleOfSplat(SVN->getMask(), N0Shuf, DAG);
+
   // If it is a splat, check if the argument vector is another splat or a
   // build_vector.
   if (SVN->isSplat() && SVN->getSplatIndex() < (int)NumElts) {
     SDNode *V = N0.getNode();
 
     // If this is a bit convert that changes the element type of the vector but
     // not the number of vector elements, look through it.  Be careful not to
     // look though conversions that change things like v4f32 to v2f64.

llvm/trunk/test/CodeGen/X86/shuffle-of-splat-multiuses.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=x86_64-unknown -mattr=+avx2 | FileCheck %s --check-prefix=AVX2
 ; PR32449
 
 define <2 x double> @foo2(<2 x double> %v, <2 x double> *%p) nounwind {
 ; AVX2-LABEL: foo2:
 ; AVX2:       # BB#0:
-; AVX2-NEXT:    vpermilpd {{.*#+}} xmm1 = xmm0[1,1]
-; AVX2-NEXT:    vpermilpd {{.*#+}} xmm0 = xmm1[1,0]
-; AVX2-NEXT:    vmovapd %xmm1, (%rdi)
+; AVX2-NEXT:    vpermilpd {{.*#+}} xmm0 = xmm0[1,1]
+; AVX2-NEXT:    vmovapd %xmm0, (%rdi)
 ; AVX2-NEXT:    retq
   %res = shufflevector <2 x double> %v, <2 x double> undef, <2 x i32> <i32 1, i32 1>
   %res1 = shufflevector<2 x double> %res, <2 x double> undef, <2 x i32> <i32 1, i32 undef>
   store <2 x double> %res, <2 x double>* %p
   ret <2 x double> %res1
 }
 
 define <4 x double> @foo4(<4 x double> %v, <4 x double> *%p) nounwind {
 ; AVX2-LABEL: foo4:
 ; AVX2:       # BB#0:
-; AVX2-NEXT:    vpermpd {{.*#+}} ymm1 = ymm0[2,2,2,2]
-; AVX2-NEXT:    vpermpd {{.*#+}} ymm0 = ymm1[2,0,2,3]
-; AVX2-NEXT:    vmovapd %ymm1, (%rdi)
+; AVX2-NEXT:    vpermpd {{.*#+}} ymm0 = ymm0[2,2,2,2]
+; AVX2-NEXT:    vmovapd %ymm0, (%rdi)
 ; AVX2-NEXT:    retq
   %res = shufflevector <4 x double> %v, <4 x double> undef, <4 x i32> <i32 2, i32 2, i32 2, i32 2>
   %res1 = shufflevector<4 x double> %res, <4 x double> undef, <4 x i32> <i32 2, i32 0, i32 undef, i32 undef>
   store <4 x double> %res, <4 x double>* %p
   ret <4 x double> %res1
 }
 
 define <8 x float> @foo8(<8 x float> %v, <8 x float> *%p) nounwind {
 ; AVX2-LABEL: foo8:
 ; AVX2:       # BB#0:
 ; AVX2-NEXT:    vmovshdup {{.*#+}} ymm0 = ymm0[1,1,3,3,5,5,7,7]
-; AVX2-NEXT:    vpermpd {{.*#+}} ymm1 = ymm0[2,2,2,2]
-; AVX2-NEXT:    vmovaps {{.*#+}} ymm0 = <2,0,u,u,5,1,3,7>
-; AVX2-NEXT:    vpermps %ymm1, %ymm0, %ymm0
-; AVX2-NEXT:    vmovapd %ymm1, (%rdi)
+; AVX2-NEXT:    vpermpd {{.*#+}} ymm0 = ymm0[2,2,2,2]
+; AVX2-NEXT:    vmovapd %ymm0, (%rdi)
 ; AVX2-NEXT:    retq
   %res = shufflevector <8 x float> %v, <8 x float> undef, <8 x i32> <i32 5, i32 5, i32 5, i32 5, i32 5, i32 5, i32 5, i32 5>
   %res1 = shufflevector<8 x float> %res, <8 x float> undef, <8 x i32> <i32 2, i32 0, i32 undef, i32 undef, i32 5, i32 1, i32 3, i32 7>
   store <8 x float> %res, <8 x float>* %p
   ret <8 x float> %res1
 }
 
 define <4 x i32> @undef_splatmask(<4 x i32> %v) nounwind {
 ; AVX2-LABEL: undef_splatmask:
 ; AVX2:       # BB#0:
-; AVX2-NEXT:    vpshufd {{.*#+}} xmm0 = xmm0[2,2,2,3]
+; AVX2-NEXT:    vpshufd {{.*#+}} xmm0 = xmm0[2,2,3,3]
 ; AVX2-NEXT:    retq
   %res = shufflevector <4 x i32> %v, <4 x i32> undef, <4 x i32> <i32 2, i32 undef, i32 2, i32 undef>
   %res1 = shufflevector <4 x i32> %res, <4 x i32> undef, <4 x i32> <i32 0, i32 2, i32 undef, i32 undef>
   ret <4 x i32> %res1
 }
 
 define <4 x i32> @undef_splatmask2(<4 x i32> %v) nounwind {
 ; AVX2-LABEL: undef_splatmask2:
 ; AVX2:       # BB#0:
 ; AVX2-NEXT:    vpshufd {{.*#+}} xmm0 = xmm0[2,2,3,3]
 ; AVX2-NEXT:    retq
   %res = shufflevector <4 x i32> %v, <4 x i32> undef, <4 x i32> <i32 2, i32 1, i32 2, i32 undef>
   %res1 = shufflevector <4 x i32> %res, <4 x i32> undef, <4 x i32> <i32 0, i32 2, i32 undef, i32 undef>
   ret <4 x i32> %res1
 }
 
 define <4 x i32> @undef_splatmask3(<4 x i32> %v) nounwind {
 ; AVX2-LABEL: undef_splatmask3:
 ; AVX2:       # BB#0:
-; AVX2-NEXT:    vpshufd {{.*#+}} xmm0 = xmm0[2,2,2,3]
+; AVX2-NEXT:    vpshufd {{.*#+}} xmm0 = xmm0[2,2,3,3]
 ; AVX2-NEXT:    retq
   %res = shufflevector <4 x i32> %v, <4 x i32> undef, <4 x i32> <i32 2, i32 undef, i32 2, i32 undef>
   %res1 = shufflevector <4 x i32> %res, <4 x i32> undef, <4 x i32> <i32 0, i32 2, i32 undef, i32 3>
   ret <4 x i32> %res1
 }
 
 define <4 x i32> @undef_splatmask4(<4 x i32> %v, <4 x i32>* %p) nounwind {
 ; AVX2-LABEL: undef_splatmask4:
 ; AVX2:       # BB#0:
-; AVX2-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[2,3,2,3]
-; AVX2-NEXT:    vpshufd {{.*#+}} xmm0 = xmm1[0,2,2,3]
-; AVX2-NEXT:    vmovdqa %xmm1, (%rdi)
+; AVX2-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[2,2,3,3]
+; AVX2-NEXT:    vpshufd {{.*#+}} xmm0 = xmm0[2,3,2,3]
+; AVX2-NEXT:    vmovdqa %xmm0, (%rdi)
+; AVX2-NEXT:    vmovdqa %xmm1, %xmm0
 ; AVX2-NEXT:    retq
   %res = shufflevector <4 x i32> %v, <4 x i32> undef, <4 x i32> <i32 2, i32 undef, i32 2, i32 undef>
   %res1 = shufflevector <4 x i32> %res, <4 x i32> undef, <4 x i32> <i32 0, i32 2, i32 undef, i32 undef>
   store <4 x i32> %res, <4 x i32>* %p
   ret <4 x i32> %res1
 }
 
 define <4 x i32> @undef_splatmask5(<4 x i32> %v, <4 x i32>* %p) nounwind {
 ; AVX2-LABEL: undef_splatmask5:
 ; AVX2:       # BB#0:
-; AVX2-NEXT:    vpbroadcastq %xmm0, %xmm1
-; AVX2-NEXT:    vpshufd {{.*#+}} xmm0 = xmm1[0,2,2,3]
-; AVX2-NEXT:    vmovdqa %xmm1, (%rdi)
+; AVX2-NEXT:    vpbroadcastd %xmm0, %xmm1
+; AVX2-NEXT:    vpbroadcastq %xmm0, %xmm0
+; AVX2-NEXT:    vmovdqa %xmm0, (%rdi)
+; AVX2-NEXT:    vmovdqa %xmm1, %xmm0
 ; AVX2-NEXT:    retq
   %res = shufflevector <4 x i32> %v, <4 x i32> undef, <4 x i32> <i32 0, i32 undef, i32 0, i32 undef>
   %res1 = shufflevector <4 x i32> %res, <4 x i32> undef, <4 x i32> <i32 0, i32 2, i32 undef, i32 3>
   store <4 x i32> %res, <4 x i32>* %p
   ret <4 x i32> %res1
 }