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

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

[spatel]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

14622–14624

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

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

Opened Bug 32719 to discuss the extra copy issue

[zvi]

ping

[spatel]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

14635–14642

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

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

14629

typo: mismatxh

14635–14641

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);

14636–14639

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

14629

Thanks

14635–14641

Ok

14636–14639

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.

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) {
+  EVT VT = Splat->getValueType(0);
+  unsigned NumElts = VT.getVectorNumElements();
+  SmallVector<int, 32> NewMask;
+  // Build the new mask by composing the user shuffle's mask with the
+  // splat-shuffle's mask. This is to ensure we don't introduce undef's
+  // incorrectly to the new mask.
+  bool IsSameMask = true;
+  for (unsigned i = 0; i < NumElts; ++i)
+    if (UserMask[i] < 0)
+      NewMask.push_back(Splat->getMaskElt(i));
+    else {
+      int NewIdx = Splat->getMaskElt(UserMask[i]);
+      IsSameMask &= NewIdx == Splat->getMaskElt(i);
+      NewMask.push_back(NewIdx);
+    }
+
+  // If the new composed mask is identical to the splat-shuffle's mask, simplify
+  // to the splat-shuffle itself. Otherwise create a new shuffle with the new
+  // mask.
+  if (IsSameMask)
+    return SDValue(Splat, 0);
+  return DAG.getVectorShuffle(VT, 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())
+      if (SDValue V = combineShuffleOfSplat(SVN->getMask(), N0Shuf, DAG))
+        return V;
+
   // 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.

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-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,2,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
 }