[ValueTracking] Do more thorough non-zero check in isKnownToBePowerOfTwo when OrZero is no set.

We can cover more cases by directly checking if the result is
known-nonzero for common patterns when they are missing `OrZero`.

This patch add `isKnownNonZero` checks for `shl`, `lshr`, `and`, and `mul`.

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

Author: goldsteinn

llvm/lib/Analysis/ValueTracking.cpp

@@ -2061,20 +2061,19 @@ bool isKnownToBeAPowerOfTwo(const Value *V, bool OrZero, unsigned Depth,
       return isKnownToBeAPowerOfTwo(I->getOperand(0), OrZero, Depth, Q);
     return false;
   case Instruction::Mul:
-    return OrZero &&
-           isKnownToBeAPowerOfTwo(I->getOperand(1), OrZero, Depth, Q) &&
-           isKnownToBeAPowerOfTwo(I->getOperand(0), OrZero, Depth, Q);
+    return isKnownToBeAPowerOfTwo(I->getOperand(1), OrZero, Depth, Q) &&
+           isKnownToBeAPowerOfTwo(I->getOperand(0), OrZero, Depth, Q) &&
+           (OrZero || isKnownNonZero(I, Depth, Q));
   case Instruction::And:
-    if (OrZero) {
-      // A power of two and'd with anything is a power of two or zero.
-      if (isKnownToBeAPowerOfTwo(I->getOperand(1), /*OrZero*/ true, Depth, Q) ||
-          isKnownToBeAPowerOfTwo(I->getOperand(0), /*OrZero*/ true, Depth, Q))
-        return true;
-      // X & (-X) is always a power of two or zero.
-      if (match(I->getOperand(0), m_Neg(m_Specific(I->getOperand(1)))) ||
-          match(I->getOperand(1), m_Neg(m_Specific(I->getOperand(0)))))
-        return true;
-    }
+    // A power of two and'd with anything is a power of two or zero.
+    if (OrZero &&
+        (isKnownToBeAPowerOfTwo(I->getOperand(1), /*OrZero*/ true, Depth, Q) ||
+         isKnownToBeAPowerOfTwo(I->getOperand(0), /*OrZero*/ true, Depth, Q)))
+      return true;
+    // X & (-X) is always a power of two or zero.
+    if (match(I->getOperand(0), m_Neg(m_Specific(I->getOperand(1)))) ||
+        match(I->getOperand(1), m_Neg(m_Specific(I->getOperand(0)))))
+      return OrZero || isKnownNonZero(I->getOperand(0), Depth, Q);
     return false;
   case Instruction::Add: {
     // Adding a power-of-two or zero to the same power-of-two or zero yields

llvm/test/Analysis/ValueTracking/known-power-of-two.ll

@@ -584,9 +584,9 @@ define i1 @and_is_pow2(i16 %x, i16 %y) {
 ; CHECK-SAME: (i16 [[X:%.*]], i16 [[Y:%.*]]) {
 ; CHECK-NEXT:    [[XNZ:%.*]] = or i16 [[X]], 4
 ; CHECK-NEXT:    [[X_NEG:%.*]] = sub nsw i16 0, [[XNZ]]
-; CHECK-NEXT:    [[XX:%.*]] = and i16 [[XNZ]], [[X_NEG]]
-; CHECK-NEXT:    [[AND:%.*]] = and i16 [[XX]], [[Y]]
-; CHECK-NEXT:    [[R:%.*]] = icmp eq i16 [[AND]], [[XX]]
+; CHECK-NEXT:    [[TMP1:%.*]] = and i16 [[X_NEG]], [[Y]]
+; CHECK-NEXT:    [[AND:%.*]] = and i16 [[TMP1]], [[XNZ]]
+; CHECK-NEXT:    [[R:%.*]] = icmp ne i16 [[AND]], 0
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %xnz = or i16 %x, 4

llvm/test/Transforms/InstSimplify/ctpop-pow2.ll

@@ -41,10 +41,7 @@ define i16 @ctpop_x_and_negx(i16 %x) {
 
 define i8 @ctpop_x_nz_and_negx(i8 %x) {
 ; CHECK-LABEL: @ctpop_x_nz_and_negx(
-; CHECK-NEXT:    [[X1:%.*]] = or i8 [[X:%.*]], 1
-; CHECK-NEXT:    [[V0:%.*]] = sub i8 0, [[X1]]
-; CHECK-NEXT:    [[V1:%.*]] = and i8 [[X1]], [[V0]]
-; CHECK-NEXT:    ret i8 [[V1]]
+; CHECK-NEXT:    ret i8 1
 ;
   %x1 = or i8 %x, 1
   %v0 = sub i8 0, %x1