Recommit "[ValueTracking] Apply the isKnownNonZero techniques in `ash…

…r`/`lshl` to `shl` and vice-versa" (2nd Try)

Wasn't related to the bug it was original thought to be causing.

llvm/lib/Analysis/ValueTracking.cpp

@@ -2512,6 +2512,57 @@ static bool isNonZeroRecurrence(const PHINode *PN) {
   }
 }
 
+static bool isNonZeroShift(const Operator *I, const APInt &DemandedElts,
+                           unsigned Depth, const Query &Q,
+                           const KnownBits &KnownVal) {
+  auto ShiftOp = [&](const APInt &Lhs, const APInt &Rhs) {
+    switch (I->getOpcode()) {
+    case Instruction::Shl:
+      return Lhs.shl(Rhs);
+    case Instruction::LShr:
+      return Lhs.lshr(Rhs);
+    case Instruction::AShr:
+      return Lhs.ashr(Rhs);
+    default:
+      llvm_unreachable("Unknown Shift Opcode");
+    }
+  };
+
+  auto InvShiftOp = [&](const APInt &Lhs, const APInt &Rhs) {
+    switch (I->getOpcode()) {
+    case Instruction::Shl:
+      return Lhs.lshr(Rhs);
+    case Instruction::LShr:
+    case Instruction::AShr:
+      return Lhs.shl(Rhs);
+    default:
+      llvm_unreachable("Unknown Shift Opcode");
+    }
+  };
+
+  if (KnownVal.isUnknown())
+    return false;
+
+  KnownBits KnownCnt =
+      computeKnownBits(I->getOperand(1), DemandedElts, Depth, Q);
+  APInt MaxShift = KnownCnt.getMaxValue();
+  unsigned NumBits = KnownVal.getBitWidth();
+  if (MaxShift.uge(NumBits))
+    return false;
+
+  if (!ShiftOp(KnownVal.One, MaxShift).isZero())
+    return true;
+
+  // If all of the bits shifted out are known to be zero, and Val is known
+  // non-zero then at least one non-zero bit must remain.
+  if (InvShiftOp(KnownVal.Zero, NumBits - MaxShift)
+          .eq(InvShiftOp(APInt::getAllOnes(NumBits), NumBits - MaxShift)) &&
+      isKnownNonZero(I->getOperand(0), DemandedElts, Depth, Q))
+    return true;
+
+  return false;
+}
+
 /// Return true if the given value is known to be non-zero when defined. For
 /// vectors, return true if every demanded element is known to be non-zero when
 /// defined. For pointers, if the context instruction and dominator tree are
@@ -2682,16 +2733,7 @@ bool isKnownNonZero(const Value *V, const APInt &DemandedElts, unsigned Depth,
     if (Known.One[0])
       return true;
 
-    if (!Known.isUnknown()) {
-      KnownBits KnownCnt =
-          computeKnownBits(I->getOperand(1), DemandedElts, Depth, Q);
-
-      if (KnownCnt.getMaxValue().ult(Known.getBitWidth()) &&
-          !Known.One.shl(KnownCnt.getMaxValue()).isZero())
-        return true;
-    }
-
-    break;
+    return isNonZeroShift(I, DemandedElts, Depth, Q, Known);
   }
   case Instruction::LShr:
   case Instruction::AShr: {
@@ -2707,19 +2749,7 @@ bool isKnownNonZero(const Value *V, const APInt &DemandedElts, unsigned Depth,
     if (Known.isNegative())
       return true;
 
-    // If the shifter operand is a constant, and all of the bits shifted
-    // out are known to be zero, and X is known non-zero then at least one
-    // non-zero bit must remain.
-    if (ConstantInt *Shift = dyn_cast<ConstantInt>(I->getOperand(1))) {
-      auto ShiftVal = Shift->getLimitedValue(BitWidth - 1);
-      // Is there a known one in the portion not shifted out?
-      if (Known.countMaxLeadingZeros() < BitWidth - ShiftVal)
-        return true;
-      // Are all the bits to be shifted out known zero?
-      if (Known.countMinTrailingZeros() >= ShiftVal)
-        return isKnownNonZero(I->getOperand(0), DemandedElts, Depth, Q);
-    }
-    break;
+    return isNonZeroShift(I, DemandedElts, Depth, Q, Known);
   }
   case Instruction::UDiv:
   case Instruction::SDiv:

llvm/test/Analysis/ValueTracking/known-non-zero.ll

@@ -261,10 +261,7 @@ define i1 @lshr_nz_bounded_cnt(i32 %cnt, i32 %y) {
 ; CHECK-LABEL: @lshr_nz_bounded_cnt(
 ; CHECK-NEXT:    [[CNT_ULT4:%.*]] = icmp ult i32 [[CNT:%.*]], 4
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CNT_ULT4]])
-; CHECK-NEXT:    [[VAL:%.*]] = or i32 [[Y:%.*]], 90
-; CHECK-NEXT:    [[SHL:%.*]] = lshr i32 [[VAL]], [[CNT]]
-; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[SHL]], 0
-; CHECK-NEXT:    ret i1 [[R]]
+; CHECK-NEXT:    ret i1 false
 ;
   %cnt_ult4 = icmp ult i32 %cnt, 4
   call void @llvm.assume(i1 %cnt_ult4)
@@ -276,11 +273,7 @@ define i1 @lshr_nz_bounded_cnt(i32 %cnt, i32 %y) {
 
 define <2 x i1> @ashr_nz_bounded_cnt_vec(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @ashr_nz_bounded_cnt_vec(
-; CHECK-NEXT:    [[CNT:%.*]] = and <2 x i32> [[X:%.*]], <i32 16, i32 24>
-; CHECK-NEXT:    [[VAL:%.*]] = or <2 x i32> [[Y:%.*]], <i32 402784272, i32 268697601>
-; CHECK-NEXT:    [[SHL:%.*]] = ashr <2 x i32> [[VAL]], [[CNT]]
-; CHECK-NEXT:    [[R:%.*]] = icmp eq <2 x i32> [[SHL]], zeroinitializer
-; CHECK-NEXT:    ret <2 x i1> [[R]]
+; CHECK-NEXT:    ret <2 x i1> zeroinitializer
 ;
   %cnt = and <2 x i32> %x, <i32 16, i32 24>
   %val = or <2 x i32> %y, <i32 402784272, i32 268697601>
@@ -332,9 +325,7 @@ define i1 @lshr_nonzero_and_shift_out_zeros(i32 %cnt, i32 %y) {
 ; CHECK-NEXT:    [[VAL:%.*]] = and i32 [[Y:%.*]], -131072
 ; CHECK-NEXT:    [[VAL_NZ:%.*]] = icmp ne i32 [[VAL]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[VAL_NZ]])
-; CHECK-NEXT:    [[SHL:%.*]] = lshr i32 [[VAL]], [[CNT]]
-; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[SHL]], 0
-; CHECK-NEXT:    ret i1 [[R]]
+; CHECK-NEXT:    ret i1 false
 ;
   %cnt_ult = icmp ult i32 %cnt, 4
   call void @llvm.assume(i1 %cnt_ult)
@@ -352,13 +343,10 @@ define i1 @lshr_nonzero_and_shift_out_zeros(i32 %cnt, i32 %y) {
 
 define i1 @ashr_nonzero_and_shift_out_zeros(i32 %ccnt, i32 %y) {
 ; CHECK-LABEL: @ashr_nonzero_and_shift_out_zeros(
-; CHECK-NEXT:    [[CNT:%.*]] = and i32 [[CCNT:%.*]], 7
 ; CHECK-NEXT:    [[VAL:%.*]] = and i32 [[Y:%.*]], -131072
 ; CHECK-NEXT:    [[VAL_NZ:%.*]] = icmp ne i32 [[VAL]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[VAL_NZ]])
-; CHECK-NEXT:    [[SHL:%.*]] = ashr i32 [[VAL]], [[CNT]]
-; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[SHL]], 0
-; CHECK-NEXT:    ret i1 [[R]]
+; CHECK-NEXT:    ret i1 false
 ;
   %cnt = and i32 %ccnt, 7
   %val = and i32 %y, -131072
@@ -372,13 +360,10 @@ define i1 @ashr_nonzero_and_shift_out_zeros(i32 %ccnt, i32 %y) {
 
 define i1 @shl_nonzero_and_shift_out_zeros(i32 %ccnt, i32 %y) {
 ; CHECK-LABEL: @shl_nonzero_and_shift_out_zeros(
-; CHECK-NEXT:    [[CNT:%.*]] = and i32 [[CCNT:%.*]], 6
 ; CHECK-NEXT:    [[VAL:%.*]] = and i32 [[Y:%.*]], 131071
 ; CHECK-NEXT:    [[VAL_NZ:%.*]] = icmp ne i32 [[VAL]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[VAL_NZ]])
-; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[VAL]], [[CNT]]
-; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[SHL]], 0
-; CHECK-NEXT:    ret i1 [[R]]
+; CHECK-NEXT:    ret i1 false
 ;
   %cnt = and i32 %ccnt, 6
   %val = and i32 %y, 131071

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

@@ -116,11 +116,7 @@ define <2 x i32> @ctpop_lshr_intmin_intmin_plus1_vec_nz(<2 x i32> %x) {
 
 define <2 x i32> @ctpop_shl2_1_vec_nz(<2 x i32> %x) {
 ; CHECK-LABEL: @ctpop_shl2_1_vec_nz(
-; CHECK-NEXT:    [[AND:%.*]] = and <2 x i32> [[X:%.*]], <i32 15, i32 15>
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i32> <i32 2, i32 1>, [[AND]]
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne <2 x i32> [[SHL]], zeroinitializer
-; CHECK-NEXT:    [[CNT:%.*]] = zext <2 x i1> [[TMP1]] to <2 x i32>
-; CHECK-NEXT:    ret <2 x i32> [[CNT]]
+; CHECK-NEXT:    ret <2 x i32> <i32 1, i32 1>
 ;
   %and = and <2 x i32> %x, <i32 15 ,i32 15>
   %shl = shl <2 x i32> <i32 2 ,i32 1>, %and