[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

2526

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

2526

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

2526

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::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;
 
     assert((!isa<UndefValue>(U->getOperand(0)) ||
-            !isa<UndefValue>(U->getOperand(1)))
-           && "Should have simplified shuffle with 2 undef inputs");
+            !isa<UndefValue>(U->getOperand(1))) &&
+           "Should have simplified shuffle with 2 undef inputs");
 
-    // Look through shuffle of 1 source vector.
-    if (isa<UndefValue>(U->getOperand(0)))
-      return ComputeNumSignBits(U->getOperand(1), Depth + 1, Q);
-    if (isa<UndefValue>(U->getOperand(1)))
-      return ComputeNumSignBits(U->getOperand(0), 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).
+    Tmp = std::numeric_limits<unsigned>::max();
+    if (!isa<UndefValue>(U->getOperand(0))) {
+      Tmp = ComputeNumSignBits(U->getOperand(0), Depth + 1, Q);
+      if (Tmp == 1)
         break;
     }
+    if (!isa<UndefValue>(U->getOperand(1))) {
+      Tmp = std::min(Tmp, ComputeNumSignBits(U->getOperand(1), Depth + 1, Q));
+      if (Tmp == 1)
+        break;
+    }
+    // The input vector(s) provided extra signbit info.
+    assert(Tmp > 1 && Tmp < std::numeric_limits<unsigned>::max() &&
+           "Unexpected signbits computation for shufflevector");
+    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.
   if (unsigned VecSignBits = computeNumSignBitsVectorConstant(V, TyBits))
     return VecSignBits;

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
 }