From 47bcf0d5a65ee7fc5d027b93c05c56ed4054abdc Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev@redking.me.uk>
Date: Tue, 8 Aug 2017 16:10:33 +0000
Subject: [PATCH] [DAGCombiner] simplifyShuffleMask - handle UNDEF inputs from
 shuffles as well as BUILD_VECTOR

Minor extension to D36393

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@310372 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 21 ++++++++++-----------
 test/CodeGen/X86/oddshuffles.ll          |  4 ++--
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 573002a6a4b..ee876b8eb8e 100644
--- a/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -15114,11 +15114,14 @@ static SDValue simplifyShuffleOperands(ShuffleVectorSDNode *SVN, SDValue N0,
 
 static SDValue simplifyShuffleMask(ShuffleVectorSDNode *SVN, SDValue N0,
                                    SDValue N1, SelectionDAG &DAG) {
-  // TODO - handle cases other than BUILD_VECTOR.
-  auto *BV0 = dyn_cast<BuildVectorSDNode>(N0);
-  auto *BV1 = dyn_cast<BuildVectorSDNode>(N1);
-  if (!BV0 && !BV1)
-    return SDValue();
+  auto isUndefElt = [](SDValue V, int Idx) {
+    // TODO - handle more cases as required.
+    if (V.getOpcode() == ISD::BUILD_VECTOR)
+      return V.getOperand(Idx).isUndef();
+    if (auto *SVN = dyn_cast<ShuffleVectorSDNode>(V))
+      return SVN->getMaskElt(Idx) < 0;
+    return false;
+  };
 
   EVT VT = SVN->getValueType(0);
   unsigned NumElts = VT.getVectorNumElements();
@@ -15127,12 +15130,8 @@ static SDValue simplifyShuffleMask(ShuffleVectorSDNode *SVN, SDValue N0,
   SmallVector<int, 8> NewMask;
   for (unsigned i = 0; i != NumElts; ++i) {
     int Idx = SVN->getMaskElt(i);
-    if (BV0 && 0 <= Idx && Idx < (int)NumElts &&
-        BV0->getOperand(Idx).isUndef()) {
-      Changed = true;
-      Idx = -1;
-    } else if (BV1 && Idx > (int)NumElts &&
-               BV1->getOperand(Idx - NumElts).isUndef()) {
+    if ((0 <= Idx && Idx < (int)NumElts && isUndefElt(N0, Idx)) ||
+        ((int)NumElts < Idx && isUndefElt(N1, Idx - NumElts))) {
       Changed = true;
       Idx = -1;
     }
diff --git a/test/CodeGen/X86/oddshuffles.ll b/test/CodeGen/X86/oddshuffles.ll
index 0bda41a30c6..5dba96ad56a 100644
--- a/test/CodeGen/X86/oddshuffles.ll
+++ b/test/CodeGen/X86/oddshuffles.ll
@@ -258,7 +258,7 @@ define void @v7i8(<4 x i8> %a, <4 x i8> %b, <7 x i8>* %p) nounwind {
 ; SSE42-NEXT:    pextrb $0, %xmm1, 6(%rdi)
 ; SSE42-NEXT:    pshufb {{.*#+}} xmm1 = xmm1[8,9,8,9,4,5,8,9,0,1,12,13,0,1,14,15]
 ; SSE42-NEXT:    pblendw {{.*#+}} xmm1 = xmm0[0],xmm1[1],xmm0[2],xmm1[3],xmm0[4],xmm1[5,6,7]
-; SSE42-NEXT:    pshufb {{.*#+}} xmm1 = xmm1[0,2,4,6,8,10,12,14,u,u,u,u,u,u,u,u]
+; SSE42-NEXT:    pshufb {{.*#+}} xmm1 = xmm1[0,2,4,6,8,10,12,u,u,u,u,u,u,u,u,u]
 ; SSE42-NEXT:    pextrw $2, %xmm1, 4(%rdi)
 ; SSE42-NEXT:    movd %xmm1, (%rdi)
 ; SSE42-NEXT:    retq
@@ -268,7 +268,7 @@ define void @v7i8(<4 x i8> %a, <4 x i8> %b, <7 x i8>* %p) nounwind {
 ; AVX-NEXT:    vpshufd {{.*#+}} xmm0 = xmm0[0,3,1,3]
 ; AVX-NEXT:    vpshufb {{.*#+}} xmm2 = xmm1[8,9,8,9,4,5,8,9,0,1,12,13,0,1,14,15]
 ; AVX-NEXT:    vpblendw {{.*#+}} xmm0 = xmm0[0],xmm2[1],xmm0[2],xmm2[3],xmm0[4],xmm2[5,6,7]
-; AVX-NEXT:    vpshufb {{.*#+}} xmm0 = xmm0[0,2,4,6,8,10,12,14,u,u,u,u,u,u,u,u]
+; AVX-NEXT:    vpshufb {{.*#+}} xmm0 = xmm0[0,2,4,6,8,10,12,u,u,u,u,u,u,u,u,u]
 ; AVX-NEXT:    vpextrb $0, %xmm1, 6(%rdi)
 ; AVX-NEXT:    vpextrw $2, %xmm0, 4(%rdi)
 ; AVX-NEXT:    vmovd %xmm0, (%rdi)
-- 
2.50.1