[ValueTracking] peek through 2-input shuffles in ComputeNumSignBits

This patch gives the IR ComputeNumSignBits the same functionality as the 
DAG version (the code is derived from the existing code).

This an extension of the single input shuffle analysis added with D53659.

Differential Revision: https://reviews.llvm.org/D53987

llvm-svn: 346071

3 files changed, 38 insertions, 26 deletions

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 89a621576ec..6e08272c64e 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -2512,26 +2512,41 @@ static unsigned ComputeNumSignBitsImpl(const Value *V, unsigned Depth,
     return ComputeNumSignBits(U->getOperand(0), Depth + 1, Q);
 
   case Instruction::ShuffleVector: {
-    // If the shuffle mask contains any undefined elements, that element of the
-    // result is undefined. Propagating information from a source operand may
-    // not be correct in that case, so just bail out.
-    if (cast<ShuffleVectorInst>(U)->getMask()->containsUndefElement())
-      break;
-
-    // If everything is undef, we can't say anything. This should be simplified.
-    Value *Op0 = U->getOperand(0), *Op1 = U->getOperand(1);
-    if (isa<UndefValue>(Op0) && isa<UndefValue>(Op1))
+    // TODO: This is copied almost directly from the SelectionDAG version of
+    //       ComputeNumSignBits. It would be better if we could share common
+    //       code. If not, make sure that changes are translated to the DAG.
+
+    // Collect the minimum number of sign bits that are shared by every vector
+    // element referenced by the shuffle.
+    auto *Shuf = cast<ShuffleVectorInst>(U);
+    int NumElts = Shuf->getOperand(0)->getType()->getVectorNumElements();
+    int NumMaskElts = Shuf->getMask()->getType()->getVectorNumElements();
+    APInt DemandedLHS(NumElts, 0), DemandedRHS(NumElts, 0);
+    for (int i = 0; i != NumMaskElts; ++i) {
+      int M = Shuf->getMaskValue(i);
+      assert(M < NumElts * 2 && "Invalid shuffle mask constant");
+      // For undef elements, we don't know anything about the common state of
+      // the shuffle result.
+      if (M == -1)
+        return 1;
+      if (M < NumElts)
+        DemandedLHS.setBit(M % NumElts);
+      else
+        DemandedRHS.setBit(M % NumElts);
+    }
+    Tmp = std::numeric_limits<unsigned>::max();
+    if (!!DemandedLHS)
+      Tmp = ComputeNumSignBits(Shuf->getOperand(0), Depth + 1, Q);
+    if (!!DemandedRHS) {
+      Tmp2 = ComputeNumSignBits(Shuf->getOperand(1), Depth + 1, Q);
+      Tmp = std::min(Tmp, Tmp2);
+    }
+    // If we don't know anything, early out and try computeKnownBits fall-back.
+    if (Tmp == 1)
       break;
-
-    // Look through shuffle of 1 source vector.
-    if (isa<UndefValue>(Op0))
-      return ComputeNumSignBits(Op1, Depth + 1, Q);
-    if (isa<UndefValue>(Op1))
-      return ComputeNumSignBits(Op0, Depth + 1, Q);
-
-    // TODO: We can look through shuffles of 2 sources by computing the minimum
-    // sign bits for each operand (similar to what we do for binops).
-    break;
+    assert(Tmp <= V->getType()->getScalarSizeInBits() &&
+           "Failed to determine minimum sign bits");
+    return Tmp;
   }
   }
 
diff --git a/llvm/test/Transforms/InstCombine/logical-select.ll b/llvm/test/Transforms/InstCombine/logical-select.ll
index 888c6a544ea..999e4512723 100644
--- a/llvm/test/Transforms/InstCombine/logical-select.ll
+++ b/llvm/test/Transforms/InstCombine/logical-select.ll
@@ -621,11 +621,9 @@ define <4 x i32> @computesignbits_through_two_input_shuffle(<4 x i32> %x, <4 x i
 ; CHECK-NEXT:    [[SEXT1:%.*]] = sext <4 x i1> [[COND1:%.*]] to <4 x i32>
 ; CHECK-NEXT:    [[SEXT2:%.*]] = sext <4 x i1> [[COND2:%.*]] to <4 x i32>
 ; CHECK-NEXT:    [[COND:%.*]] = shufflevector <4 x i32> [[SEXT1]], <4 x i32> [[SEXT2]], <4 x i32> <i32 0, i32 2, i32 4, i32 6>
-; CHECK-NEXT:    [[NOTCOND:%.*]] = xor <4 x i32> [[COND]], <i32 -1, i32 -1, i32 -1, i32 -1>
-; CHECK-NEXT:    [[AND1:%.*]] = and <4 x i32> [[NOTCOND]], [[X:%.*]]
-; CHECK-NEXT:    [[AND2:%.*]] = and <4 x i32> [[COND]], [[Y:%.*]]
-; CHECK-NEXT:    [[SEL:%.*]] = or <4 x i32> [[AND1]], [[AND2]]
-; CHECK-NEXT:    ret <4 x i32> [[SEL]]
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc <4 x i32> [[COND]] to <4 x i1>
+; CHECK-NEXT:    [[TMP2:%.*]] = select <4 x i1> [[TMP1]], <4 x i32> [[Y:%.*]], <4 x i32> [[X:%.*]]
+; CHECK-NEXT:    ret <4 x i32> [[TMP2]]
 ;
   %sext1 = sext <4 x i1> %cond1 to <4 x i32>
   %sext2 = sext <4 x i1> %cond2 to <4 x i32>
diff --git a/llvm/unittests/Analysis/ValueTrackingTest.cpp b/llvm/unittests/Analysis/ValueTrackingTest.cpp
index e66d8f77fd4..b13948adad7 100644
--- a/llvm/unittests/Analysis/ValueTrackingTest.cpp
+++ b/llvm/unittests/Analysis/ValueTrackingTest.cpp
@@ -514,7 +514,6 @@ TEST(ValueTracking, ComputeNumSignBits_Shuffle) {
   EXPECT_EQ(ComputeNumSignBits(RVal, M->getDataLayout()), 1u);
 }
 
-// FIXME:
 // No guarantees for canonical IR in this analysis, so a shuffle element that
 // references an undef value means this can't return any extra information. 
 TEST(ValueTracking, ComputeNumSignBits_Shuffle2) {
@@ -534,7 +533,7 @@ TEST(ValueTracking, ComputeNumSignBits_Shuffle2) {
 
   auto *RVal =
       cast<ReturnInst>(F->getEntryBlock().getTerminator())->getOperand(0);
-  EXPECT_EQ(ComputeNumSignBits(RVal, M->getDataLayout()), 32u);
+  EXPECT_EQ(ComputeNumSignBits(RVal, M->getDataLayout()), 1u);
 }
 
 TEST(ValueTracking, ComputeKnownBits) {