[ValueTracking] peek through 2-input shuffles in ComputeNumSignBits

[spatel]

Summary

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.

[RKSimon]

lib/Analysis/ValueTracking.cpp

2548

Don't you still have to check if a valid shuffle mask element references an UNDEF operand?

%1 = shufflevector <4 x float> %a0, <4 x float> undef, <4 x i32> <i32 4, i32 5, i32 0, i32 1>

Or is that guaranteed to have already been handled, the mask element set to undef and caught by containsUndefElement() ?

[spatel]

lib/Analysis/ValueTracking.cpp

2548

Yes, good catch. This is an analysis usable by any pass, so we can't assume canonical code. That means even the existing assert for 2 undefs is not allowed; it has to be a real check.

[RKSimon]

lib/Analysis/ValueTracking.cpp

2548

You can probably go further and replicate the SelectionDAG version's checking of the entire mask - bailing on undef mask elts and determining if it demands/references LHS/RHS operands - and then calling the LHS/RHS num sign bits accordingly.

Might be good prep work for PR36319.

[spatel]

Patch updated:
Copy the DAG logic completely to make sure undefs are handled properly and allow further enhancements.
I added a unit-test to show that the existing code had a bug.
Also, I fixed the 2-undef bug and added a unit-test with:
rL345988

[RKSimon]

LGTM - thanks

lib/Analysis/ValueTracking.cpp

static unsigned ComputeNumSignBitsImpl(const Value *V, unsigned Depth,
   case Instruction::ExtractElement:
     // Look through extract element. At the moment we keep this simple and skip
     // tracking the specific element. But at least we might find information
     // valid for all elements of the vector (for example if vector is sign
     // extended, shifted, etc).
     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))
-      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).
+    // 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;
+    assert(Tmp <= V->getType()->getScalarSizeInBits() &&
+           "Failed to determine minimum sign bits");
+    return Tmp;
   }
   }
 
   // Finally, if we can prove that the top bits of the result are 0's or 1's,
   // use this information.
 
   // If we can examine all elements of a vector constant successfully, we're
   // done (we can't do any better than that). If not, keep trying.

test/Transforms/InstCombine/logical-select.ll

;
   ret <4 x i32> %sel
 }
 
 define <4 x i32> @computesignbits_through_two_input_shuffle(<4 x i32> %x, <4 x i32> %y, <4 x i1> %cond1, <4 x i1> %cond2) {
 ; CHECK-LABEL: @computesignbits_through_two_input_shuffle(
 ; 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>
   %cond = shufflevector <4 x i32> %sext1, <4 x i32> %sext2, <4 x i32> <i32 0, i32 2, i32 4, i32 6>
   %notcond = xor <4 x i32> %cond, <i32 -1, i32 -1, i32 -1, i32 -1>
   %and1 = and <4 x i32> %notcond, %x
   %and2 = and <4 x i32> %cond, %y
   %sel = or <4 x i32> %and1, %and2
   ret <4 x i32> %sel
 }
 

unittests/Analysis/ValueTrackingTest.cpp

TEST(ValueTracking, ComputeNumSignBits_Shuffle) {
   auto *F = M->getFunction("f");
   assert(F && "Bad assembly?");
 
   auto *RVal =
       cast<ReturnInst>(F->getEntryBlock().getTerminator())->getOperand(0);
   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) {
   StringRef Assembly = "define <2 x i32> @f(<2 x i1> %x) { "
                        "  %sext = sext <2 x i1> %x to <2 x i32> "
                        "  %val = shufflevector <2 x i32> %sext, <2 x i32> undef, <2 x i32> <i32 0, i32 2> "
                        "  ret <2 x i32> %val "
                        "} ";
 
   LLVMContext Context;
   SMDiagnostic Error;
   auto M = parseAssemblyString(Assembly, Error, Context);
   assert(M && "Bad assembly?");
 
   auto *F = M->getFunction("f");
   assert(F && "Bad assembly?");
 
   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) {
   StringRef Assembly = "define i32 @f(i32 %a, i32 %b) { "
                        "  %ash = mul i32 %a, 8 "
                        "  %aad = add i32 %ash, 7 "
                        "  %aan = and i32 %aad, 4095 "
                        "  %bsh = shl i32 %b, 4 "