[DAGCombine] A shuffle of a splat is always the splat itself

[zvi]

Summary

Add a simplification:
shuffle (splat-shuffle), undef, M --> splat-shuffle

Fixes pr32449

Patch by Sanjay Patel

[spatel]

This seems unnecessarily complicated. Why can't we just match any shuffle of a splat before we get here and simplify it?

Also, would you be interested in fixing this in IR? :)

$ ./opt -instsimplify shufsplat.ll -S

define <4 x i32> @foo(<4 x i32> %x) {
  %splat = shufflevector <4 x i32> %x, <4 x i32> undef, <4 x i32> zeroinitializer
  %shuf = shufflevector <4 x i32> %splat, <4 x i32> undef, <4 x i32> <i32 0, i32 3, i32 2, i32 1>
  ret <4 x i32> %shuf
}

This example is handled by instcombine...but in a very inefficient way.

[zvi]

I initially tried handling all cases of shuffle of shuffle-splat but the improvements from D27793 regressed, so this is what i came up with.
Should i change this to an X86-specific combine to follow the spirit of D27793?

About doing this in instcombine, yes, i plan to look into that too, but the pattern optimized by this patch is born out of thin air in another patch I'm working on that improves the combine of BUILD_VECTOR into composed shuffles.

[spatel]

In D31426#712864, @zvi wrote:

I initially tried handling all cases of shuffle of shuffle-splat but the improvements from D27793 regressed, so this is what i came up with.
Should i change this to an X86-specific combine to follow the spirit of D27793?

I think D27793 overstepped what it was trying to protect against. In the case where N1 is undef, we probably still want to do more combining? But that doesn't solve the tests here. I haven't looked at what is happening below that check.

But the transform you want is a "simplify" not a "combine" (we're returning an existing node), so we shouldn't need to check uses, legality, or anything else. This diff solves the tests in this patch with no regressions on other tests:

Index: lib/CodeGen/SelectionDAG/DAGCombiner.cpp
===================================================================
--- lib/CodeGen/SelectionDAG/DAGCombiner.cpp	(revision 298954)
+++ lib/CodeGen/SelectionDAG/DAGCombiner.cpp	(working copy)
@@ -14646,6 +14646,12 @@
       return DAG.getVectorShuffle(VT, SDLoc(N), N0, N1, NewMask);
   }
 
+  // A shuffle of a splat is always the splat itself:
+  // shuffle (splat-shuffle), undef, M --> splat-shuffle
+  if (auto *N0Shuf = dyn_cast<ShuffleVectorSDNode>(N0))
+    if (N1.isUndef() && N0Shuf->isSplat())
+      return N0;
+
   // If it is a splat, check if the argument vector is another splat or a
   // build_vector.
   if (SVN->isSplat() && SVN->getSplatIndex() < (int)NumElts) {

[zvi]

Thanks, Sanjay, for suggesting this superior patch. I will use it for the next revision.

[spatel]

LGTM of course. :)
If anyone who knows the shuffle combiner better wants to confirm that I haven't misdiagnosed the problem, that would be good too.

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

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 splat is always the splat itself:
+  // shuffle (splat-shuffle), undef, M --> splat-shuffle
+  if (auto *N0Shuf = dyn_cast<ShuffleVectorSDNode>(N0))
+    if (N1.isUndef() && N0Shuf->isSplat())
+      return N0;
+
   // 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
 }