[ValueTracking] improve ComputeNumSignBits for vectors

[spatel]

Summary

This is similar to the computeKnownBits improvement in rL268479. There's probably more to do for vector logic instructions, but this should let us see non-splat constant masking ops that can become vector selects instead of and/andn/or sequences.

Side note - I don't know how this line under the computeKnownBits() call in this patch could ever shift by any amount other than zero:

Mask <<= Mask.getBitWidth()-TyBits;

We can delete that unless I'm misunderstanding.

[eli.friedman]

LGTM.

I did a bit of digging about the "Mask <<= Mask.getBitWidth()-TyBits;" thing; it looks like r47654 is at fault. (The original version of the code used uint64_t instead of APInt.)

[spatel]

In D21610#464745, @eli.friedman wrote:

LGTM.

I did a bit of digging about the "Mask <<= Mask.getBitWidth()-TyBits;" thing; it looks like r47654 is at fault. (The original version of the code used uint64_t instead of APInt.)

Thanks for finding that. I'll delete that line as a separate commit then.

llvm/trunk/include/llvm/Analysis/ValueTracking.h

bool MaskedValueIsZero(Value *V, const APInt &Mask, const DataLayout &DL,
                          unsigned Depth = 0, AssumptionCache *AC = nullptr,
                          const Instruction *CxtI = nullptr,
                          const DominatorTree *DT = nullptr);
 
   /// Return the number of times the sign bit of the register is replicated into
   /// the other bits. We know that at least 1 bit is always equal to the sign
   /// bit (itself), but other cases can give us information. For example,
   /// immediately after an "ashr X, 2", we know that the top 3 bits are all
-  /// equal to each other, so we return 3.
-  ///
-  /// 'Op' must have a scalar integer type.
+  /// equal to each other, so we return 3. For vectors, return the number of
+  /// sign bits for the vector element with the mininum number of known sign
+  /// bits.
   unsigned ComputeNumSignBits(Value *Op, const DataLayout &DL,
                               unsigned Depth = 0, AssumptionCache *AC = nullptr,
                               const Instruction *CxtI = nullptr,
                               const DominatorTree *DT = nullptr);
 
   /// This function computes the integer multiple of Base that equals V. If
   /// successful, it returns true and returns the multiple in Multiple. If
   /// unsuccessful, it returns false. Also, if V can be simplified to an

llvm/trunk/lib/Analysis/ValueTracking.cpp

 /// for all of the elements in the vector.
 bool MaskedValueIsZero(Value *V, const APInt &Mask, unsigned Depth,
                        const Query &Q) {
   APInt KnownZero(Mask.getBitWidth(), 0), KnownOne(Mask.getBitWidth(), 0);
   computeKnownBits(V, KnownZero, KnownOne, Depth, Q);
   return (KnownZero & Mask) == Mask;
 }
 
+/// For vector constants, loop over the elements and find the constant with the
+/// minimum number of sign bits. Return 0 if the value is not a vector constant
+/// or if any element was not analyzed; otherwise, return the count for the
+/// element with the minimum number of sign bits.
+static unsigned computeNumSignBitsVectorConstant(Value *V, unsigned TyBits) {
+  auto *CV = dyn_cast<Constant>(V);
+  if (!CV || !CV->getType()->isVectorTy())
+    return 0;
+
+  unsigned MinSignBits = TyBits;
+  unsigned NumElts = CV->getType()->getVectorNumElements();
+  for (unsigned i = 0; i != NumElts; ++i) {
+    // If we find a non-ConstantInt, bail out.
+    auto *Elt = dyn_cast_or_null<ConstantInt>(CV->getAggregateElement(i));
+    if (!Elt)
+      return 0;
+
+    // If the sign bit is 1, flip the bits, so we always count leading zeros.
+    APInt EltVal = Elt->getValue();
+    if (EltVal.isNegative())
+      EltVal = ~EltVal;
+    MinSignBits = std::min(MinSignBits, EltVal.countLeadingZeros());
+  }
 
+  return MinSignBits;
+}
 
 /// Return the number of times the sign bit of the register is replicated into
 /// the other bits. We know that at least 1 bit is always equal to the sign bit
 /// (itself), but other cases can give us information. For example, immediately
 /// after an "ashr X, 2", we know that the top 3 bits are all equal to each
-/// other, so we return 3.
-///
-/// 'Op' must have a scalar integer type.
-///
+/// other, so we return 3. For vectors, return the number of sign bits for the
+/// vector element with the mininum number of known sign bits.
 unsigned ComputeNumSignBits(Value *V, unsigned Depth, const Query &Q) {
   unsigned TyBits = Q.DL.getTypeSizeInBits(V->getType()->getScalarType());
   unsigned Tmp, Tmp2;
   unsigned FirstAnswer = 1;
 
   // Note that ConstantInt is handled by the general computeKnownBits case
   // below.
 
unsigned ComputeNumSignBits(Value *V, unsigned Depth, const Query &Q) {
   case Instruction::Trunc:
     // FIXME: it's tricky to do anything useful for this, but it is an important
     // case for targets like X86.
     break;
   }
 
   // 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;
+
   APInt KnownZero(TyBits, 0), KnownOne(TyBits, 0);
   APInt Mask;
   computeKnownBits(V, KnownZero, KnownOne, Depth, Q);
 
   if (KnownZero.isNegative()) {        // sign bit is 0
     Mask = KnownZero;
   } else if (KnownOne.isNegative()) {  // sign bit is 1;
     Mask = KnownOne;

llvm/trunk/test/Transforms/InstSimplify/shr-nop.ll

 ; CHECK-NEXT:    ret <2 x i64> <i64 -1, i64 -1>
 ;
   %shr = ashr <2 x i64> <i64 -1, i64 -1>, %shiftval
   ret <2 x i64> %shr
 }
 
 define <2 x i4> @ashr_zero_minus1_vec(<2 x i4> %shiftval) {
 ; CHECK-LABEL: @ashr_zero_minus1_vec(
-; CHECK-NEXT:    [[SHR:%.*]] = ashr <2 x i4> <i4 0, i4 -1>, %shiftval
-; CHECK-NEXT:    ret <2 x i4> [[SHR]]
+; CHECK-NEXT:    ret <2 x i4> <i4 0, i4 -1>
 ;
   %shr = ashr <2 x i4> <i4 0, i4 -1>, %shiftval
   ret <2 x i4> %shr
 }