[ValueTracking] Apply the isKnownNonZero techniques in `ashr`/`lshl` to `shl` and vice-versa

[goldstein.w.n]

Summary

For all shifts we can apply the same two optimizations.

1. ShiftOp(KnownVal.One, Max(KnownCnt)) != 0 -> result is non-zero
2. If already known Val != 0 and we only shift out zeros (based on Max(KnownCnt)) -> result is non-zero

The former exists for shl and the latter (for constant Cnt) exists
for ashr/lshr.

This patch improves the latter to use Max(KnownCnt) instead of
relying on a constant shift Cnt and applies both techniques for all
shift ops.

[goldstein.w.n]

Remove extraneous changes

[nikic]

llvm/lib/Analysis/ValueTracking.cpp

2518

const APInt &

2543

clang-format

2735–2736

This and below should directly return, because we don't want to fall through to the generic known bits logic anymore.

[goldstein.w.n]

Make args const ref. Make new helper unconditional return

[nikic]

LGTM, though I suspect you'd need less code if you didn't try to factor this into one function and just handled left and right shifts separately.

[goldstein.w.n]

This doesn't appear to have been related to PR62175 so repushing.

llvm/lib/Analysis/ValueTracking.cpp

static bool isNonZeroRecurrence(const PHINode *PN) {
   case Instruction::AShr:
   case Instruction::LShr:
     return BO->isExact();
   default:
     return false;
   }
 }
 
+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
 /// specified, perform context-sensitive analysis and return true if the
 /// pointer couldn't possibly be null at the specified instruction.
 /// Supports values with integer or pointer type and vectors of integers.
 bool isKnownNonZero(const Value *V, const APInt &DemandedElts, unsigned Depth,
                     const Query &Q) {
if (Q.IIQ.hasNoUnsignedWrap(BO) || Q.IIQ.hasNoSignedWrap(BO))
       return isKnownNonZero(I->getOperand(0), Depth, Q);
 
     // shl X, Y != 0 if X is odd.  Note that the value of the shift is undefined
     // if the lowest bit is shifted off the end.
     KnownBits Known(BitWidth);
     computeKnownBits(I->getOperand(0), DemandedElts, Known, Depth, Q);
     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: {
     // shr exact can only shift out zero bits.
     const PossiblyExactOperator *BO = cast<PossiblyExactOperator>(V);
     if (BO->isExact())
       return isKnownNonZero(I->getOperand(0), Depth, Q);
 
     // shr X, Y != 0 if X is negative.  Note that the value of the shift is not
     // defined if the sign bit is shifted off the end.
     KnownBits Known =
         computeKnownBits(I->getOperand(0), DemandedElts, Depth, Q);
     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:
     // div exact can only produce a zero if the dividend is zero.
     if (cast<PossiblyExactOperator>(I)->isExact())
       return isKnownNonZero(I->getOperand(0), DemandedElts, Depth, Q);
     break;
   case Instruction::Add: {

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

;
   %z = icmp eq i32 %shl, 0
   ret i1 %z
 }
 
 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)
   %val = or i32 %y, 90
   %shl = lshr i32 %val, %cnt
   %r = icmp eq i32 %shl, 0
   ret i1 %r
 }
 
 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>
   %shl = ashr <2 x i32> %val, %cnt
   %r = icmp eq <2 x i32> %shl, zeroinitializer
   ret <2 x i1> %r
 }
 

 
 define i1 @lshr_nonzero_and_shift_out_zeros(i32 %cnt, i32 %y) {
 ; CHECK-LABEL: @lshr_nonzero_and_shift_out_zeros(
 ; CHECK-NEXT:    [[CNT_ULT:%.*]] = icmp ult i32 [[CNT:%.*]], 4
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CNT_ULT]])
 ; 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)
 
   %val = and i32 %y, -131072
 
   %val_nz =  icmp ne i32 %val, 0
   call void @llvm.assume(i1 %val_nz)
 
   %shl = lshr i32 %val, %cnt
   %r = icmp eq i32 %shl, 0
   ret i1 %r
 }
 
 
 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
   %val_nz =  icmp ne i32 %val, 0
   call void @llvm.assume(i1 %val_nz)
   %shl = ashr i32 %val, %cnt
   %r = icmp eq i32 %shl, 0
   ret i1 %r
 }
 
 
 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
   %val_nz =  icmp ne i32 %val, 0
   call void @llvm.assume(i1 %val_nz)
   %shl = shl i32 %val, %cnt
   %r = icmp eq i32 %shl, 0
   ret i1 %r

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

;
   %shr = lshr <2 x i32> <i32 2147483648 ,i32 2147483649>, %x1
   %cnt = call <2 x i32> @llvm.ctpop.v2i32(<2 x i32> %shr)
   ret <2 x i32> %cnt
 }
 
 
 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
   %cnt = call <2 x i32> @llvm.ctpop.v2i32(<2 x i32> %shl)
   ret <2 x i32> %cnt
 }
 
 define <2 x i64> @ctpop_x_and_negx_vec(<2 x i64> %x) {