[ValueTracking] improve analysis for "C << X" and "C >> X"

[spatel]

Summary

This is based on the example/comments in:
https://llvm.org/PR48984

I tried just lifting the restriction in computeKnownBitsFromShiftOperator() as suggested, but that doesn't catch all of the cases shown here. I didn't step through to see exactly why that happened. But it seems like a reasonable compromise to cheaply check the special-case of shifting a constant.

There's a slight regression on a cmp transform as noted, but I think this is likely the more important/common pattern, so we can fix that icmp pattern later if needed.

[RKSimon]

This is very similar to D90479 which is trying to do this more generally by using the KnownBits shift helpers.

[spatel]

In D95959#2539683, @RKSimon wrote:

This is very similar to D90479 which is trying to do this more generally by using the KnownBits shift helpers.

Ah, I forgot about that patch. Was anything holding it up?
Yes, they're showing almost the same test diffs apart from the zext+icmp test. I don't have a preference for the fix, so we should push one of these.

[RKSimon]

In D95959#2539817, @spatel wrote:

In D95959#2539683, @RKSimon wrote:

This is very similar to D90479 which is trying to do this more generally by using the KnownBits shift helpers.

Ah, I forgot about that patch. Was anything holding it up?
Yes, they're showing almost the same test diffs apart from the zext+icmp test. I don't have a preference for the fix, so we should push one of these.

D90479 noticed that we were assuming a poison shift could be 'known' to be all zero - so I needed to go back and fix that first but I never did (got distracted by reductions, and other squirrels.....).

If you want to get this patch that's fine - D90479 will just remove this again when I get around to looking at it.

[spatel]

If there are no objections (and someone approves), I can push this change to make progress.

[RKSimon]

LGTM

llvm/lib/Analysis/ValueTracking.cpp

auto KF = [NSW](const KnownBits &KnownVal, const KnownBits &KnownAmt) {
           Result.Zero.setSignBit();
         if (KnownVal.One.isSignBitSet())
           Result.One.setSignBit();
       }
       return Result;
     };
     computeKnownBitsFromShiftOperator(I, DemandedElts, Known, Known2, Depth, Q,
                                       KF);
+    // Trailing zeros of a right-shifted constant never decrease.
+    const APInt *C;
+    if (match(I->getOperand(0), m_APInt(C)))
+      Known.Zero.setLowBits(C->countTrailingZeros());
     break;
   }
   case Instruction::LShr: {
     auto KF = [](const KnownBits &KnownVal, const KnownBits &KnownAmt) {
       return KnownBits::lshr(KnownVal, KnownAmt);
     };
     computeKnownBitsFromShiftOperator(I, DemandedElts, Known, Known2, Depth, Q,
                                       KF);
+    // Leading zeros of a left-shifted constant never decrease.
+    const APInt *C;
+    if (match(I->getOperand(0), m_APInt(C)))
+      Known.Zero.setHighBits(C->countLeadingZeros());
     break;
   }
   case Instruction::AShr: {
     auto KF = [](const KnownBits &KnownVal, const KnownBits &KnownAmt) {
       return KnownBits::ashr(KnownVal, KnownAmt);
     };
     computeKnownBitsFromShiftOperator(I, DemandedElts, Known, Known2, Depth, Q,
                                       KF);

llvm/test/Transforms/InstCombine/and2.ll

 ; CHECK-NEXT:    [[AND:%.*]] = zext <2 x i1> [[TMP1]] to <2 x i8>
 ; CHECK-NEXT:    ret <2 x i8> [[AND]]
 ;
   %sh = shl <2 x i8> <i8 1, i8 undef>, %x
   %and = and <2 x i8> %sh, <i8 1, i8 undef>
   ret <2 x i8> %and
 }
 
-; (1 >> x) & 1 --> zext(x == 0)
+; The mask is unnecessary.
 
 define i8 @and1_lshr1_is_cmp_eq_0(i8 %x) {
 ; CHECK-LABEL: @and1_lshr1_is_cmp_eq_0(
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i8 [[X:%.*]], 0
-; CHECK-NEXT:    [[AND:%.*]] = zext i1 [[TMP1]] to i8
-; CHECK-NEXT:    ret i8 [[AND]]
+; CHECK-NEXT:    [[SH:%.*]] = lshr i8 1, [[X:%.*]]
+; CHECK-NEXT:    ret i8 [[SH]]
 ;
   %sh = lshr i8 1, %x
   %and = and i8 %sh, 1
   ret i8 %and
 }
 
-; Don't do it if the shift has another use.
-
 define i8 @and1_lshr1_is_cmp_eq_0_multiuse(i8 %x) {
 ; CHECK-LABEL: @and1_lshr1_is_cmp_eq_0_multiuse(
 ; CHECK-NEXT:    [[SH:%.*]] = lshr i8 1, [[X:%.*]]
-; CHECK-NEXT:    [[AND:%.*]] = and i8 [[SH]], 1
-; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw i8 [[SH]], [[AND]]
+; CHECK-NEXT:    [[ADD:%.*]] = shl nuw nsw i8 [[SH]], 1
 ; CHECK-NEXT:    ret i8 [[ADD]]
 ;
   %sh = lshr i8 1, %x
   %and = and i8 %sh, 1
   %add = add i8 %sh, %and
   ret i8 %add
 }
 
-; (1 >> x) & 1 --> zext(x == 0)
+; The mask is unnecessary.
 
 define <2 x i8> @and1_lshr1_is_cmp_eq_0_vec(<2 x i8> %x) {
 ; CHECK-LABEL: @and1_lshr1_is_cmp_eq_0_vec(
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq <2 x i8> [[X:%.*]], zeroinitializer
-; CHECK-NEXT:    [[AND:%.*]] = zext <2 x i1> [[TMP1]] to <2 x i8>
-; CHECK-NEXT:    ret <2 x i8> [[AND]]
+; CHECK-NEXT:    [[SH:%.*]] = lshr <2 x i8> <i8 1, i8 1>, [[X:%.*]]
+; CHECK-NEXT:    ret <2 x i8> [[SH]]
 ;
   %sh = lshr <2 x i8> <i8 1, i8 1>, %x
   %and = and <2 x i8> %sh, <i8 1, i8 1>
   ret <2 x i8> %and
 }
 
 define <2 x i8> @and1_lshr1_is_cmp_eq_0_vec_undef(<2 x i8> %x) {
 ; CHECK-LABEL: @and1_lshr1_is_cmp_eq_0_vec_undef(

llvm/test/Transforms/InstCombine/lshr-and-negC-icmpeq-zero.ll

;
   %add = add i32 %lshr, %z  ; extra use of lshr
   store i32 %add, i32* %q
   %r = icmp eq i32 %and, 0
   ret i1 %r
 }
 
 ; Negative tests
 
-; X is constant
+; TODO: This could be reduced to lshr+icmp ult.
 
 define i1 @scalar_i32_lshr_and_negC_eq_X_is_constant1(i32 %y) {
 ; CHECK-LABEL: @scalar_i32_lshr_and_negC_eq_X_is_constant1(
 ; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 12345, [[Y:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = icmp ult i32 [[LSHR]], 8
+; CHECK-NEXT:    [[AND:%.*]] = and i32 [[LSHR]], 16376
+; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[AND]], 0
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %lshr = lshr i32 12345, %y
   %and = and i32 %lshr, 4294967288  ; ~7
   %r = icmp eq i32 %and, 0
   ret i1 %r
 }
 
+; TODO: This could be reduced to lshr+icmp ult.
+
 define i1 @scalar_i32_lshr_and_negC_eq_X_is_constant2(i32 %y) {
 ; CHECK-LABEL: @scalar_i32_lshr_and_negC_eq_X_is_constant2(
 ; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 268435456, [[Y:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = icmp ult i32 [[LSHR]], 8
+; CHECK-NEXT:    [[AND:%.*]] = and i32 [[LSHR]], 536870904
+; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[AND]], 0
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %lshr = lshr i32 268435456, %y
   %and = and i32 %lshr, 4294967288  ; ~7
   %r = icmp eq i32 %and, 0
   ret i1 %r
 }
 

llvm/test/Transforms/InstCombine/shift.ll

;
   %B28 = urem i177 %B24, %B6
   store i177 %B28, i177* %G62, align 4
   ret void
 }
 
 define i8 @lshr_mask_demand(i8 %x) {
 ; CHECK-LABEL: @lshr_mask_demand(
 ; CHECK-NEXT:    [[S:%.*]] = lshr i8 63, [[X:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = and i8 [[S]], -32
+; CHECK-NEXT:    [[R:%.*]] = and i8 [[S]], 32
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %s = lshr i8 63, %x ; 0b00111111
   %r = and i8 %s, 224 ; 0b11100000
   ret i8 %r
 }
 
 define i8 @shl_mask_demand(i8 %x) {
 ; CHECK-LABEL: @shl_mask_demand(
 ; CHECK-NEXT:    [[S:%.*]] = shl i8 12, [[X:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = and i8 [[S]], 7
+; CHECK-NEXT:    [[R:%.*]] = and i8 [[S]], 4
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %s = shl i8 12, %x ; 0b00001100
   %r = and i8 %s, 7  ; 0b00000111
   ret i8 %r
 }

llvm/test/Transforms/InstCombine/signbit-shl-and-icmpeq-zero.ll

;
   %r = icmp eq i32 %and, 0
   ret i1 %r
 }
 
 ; X is constant
 
 define i1 @scalar_i32_signbit_shl_and_eq_X_is_constant1(i32 %y) {
 ; CHECK-LABEL: @scalar_i32_signbit_shl_and_eq_X_is_constant1(
-; CHECK-NEXT:    [[SHL:%.*]] = shl i32 -2147483648, [[Y:%.*]]
-; CHECK-NEXT:    [[AND:%.*]] = and i32 [[SHL]], 12345
-; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[AND]], 0
-; CHECK-NEXT:    ret i1 [[R]]
+; CHECK-NEXT:    ret i1 true
 ;
   %shl = shl i32 2147483648, %y
   %and = and i32 %shl, 12345
   %r = icmp eq i32 %and, 0
   ret i1 %r
 }
 
 define i1 @scalar_i32_signbit_shl_and_eq_X_is_constant2(i32 %y) {
 ; CHECK-LABEL: @scalar_i32_signbit_shl_and_eq_X_is_constant2(
-; CHECK-NEXT:    [[SHL:%.*]] = shl i32 -2147483648, [[Y:%.*]]
-; CHECK-NEXT:    [[AND:%.*]] = and i32 [[SHL]], 1
-; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[AND]], 0
-; CHECK-NEXT:    ret i1 [[R]]
+; CHECK-NEXT:    ret i1 true
 ;
   %shl = shl i32 2147483648, %y
   %and = and i32 %shl, 1
   %r = icmp eq i32 %and, 0
   ret i1 %r
 }
 
 ; Negative tests
 
 ; Check 'slt' predicate
 
 define i1 @scalar_i32_signbit_shl_and_slt(i32 %x, i32 %y) {
 ; CHECK-LABEL: @scalar_i32_signbit_shl_and_slt(
 ; CHECK-NEXT:    [[SHL:%.*]] = shl i32 -2147483648, [[Y:%.*]]
 ; CHECK-NEXT:    [[AND:%.*]] = and i32 [[SHL]], [[X:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = icmp slt i32 [[AND]], 0
+; CHECK-NEXT:    [[R:%.*]] = icmp ne i32 [[AND]], 0
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %shl = shl i32 2147483648, %y
   %and = and i32 %shl, %x
   %r = icmp slt i32 %and, 0
   ret i1 %r
 }
 
 ; Compare with nonzero
 
 define i1 @scalar_i32_signbit_shl_and_eq_nonzero(i32 %x, i32 %y) {
 ; CHECK-LABEL: @scalar_i32_signbit_shl_and_eq_nonzero(
-; CHECK-NEXT:    [[SHL:%.*]] = shl i32 -2147483648, [[Y:%.*]]
-; CHECK-NEXT:    [[AND:%.*]] = and i32 [[SHL]], [[X:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[AND]], 1
-; CHECK-NEXT:    ret i1 [[R]]
+; CHECK-NEXT:    ret i1 false
 ;
   %shl = shl i32 2147483648, %y
   %and = and i32 %shl, %x
   %r = icmp eq i32 %and, 1  ; should be comparing with 0
   ret i1 %r
 }

llvm/test/Transforms/InstCombine/zext-or-icmp.ll

 define i32 @dont_widen_undef() {
 ; CHECK-LABEL: @dont_widen_undef(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[BLOCK2:%.*]]
 ; CHECK:       block1:
 ; CHECK-NEXT:    br label [[BLOCK2]]
 ; CHECK:       block2:
 ; CHECK-NEXT:    [[CMP_I:%.*]] = phi i1 [ false, [[BLOCK1:%.*]] ], [ true, [[ENTRY:%.*]] ]
-; CHECK-NEXT:    [[M_011:%.*]] = phi i32 [ 0, [[BLOCK1]] ], [ 33, [[ENTRY]] ]
-; CHECK-NEXT:    [[M_1_OP:%.*]] = lshr i32 1, [[M_011]]
-; CHECK-NEXT:    [[SEXT_MASK:%.*]] = and i32 [[M_1_OP]], 65535
-; CHECK-NEXT:    [[CMP115:%.*]] = icmp ne i32 [[SEXT_MASK]], 0
+; CHECK-NEXT:    [[CMP115:%.*]] = phi i1 [ true, [[BLOCK1]] ], [ false, [[ENTRY]] ]
 ; CHECK-NEXT:    [[CMP1:%.*]] = or i1 [[CMP_I]], [[CMP115]]
 ; CHECK-NEXT:    [[CONV2:%.*]] = zext i1 [[CMP1]] to i32
 ; CHECK-NEXT:    ret i32 [[CONV2]]
 ;
 entry:
   br label %block2
 
 block1:

 define i32 @dont_widen_undef_logical() {
 ; CHECK-LABEL: @dont_widen_undef_logical(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[BLOCK2:%.*]]
 ; CHECK:       block1:
 ; CHECK-NEXT:    br label [[BLOCK2]]
 ; CHECK:       block2:
 ; CHECK-NEXT:    [[CMP_I:%.*]] = phi i1 [ false, [[BLOCK1:%.*]] ], [ true, [[ENTRY:%.*]] ]
-; CHECK-NEXT:    [[M_011:%.*]] = phi i32 [ 0, [[BLOCK1]] ], [ 33, [[ENTRY]] ]
-; CHECK-NEXT:    [[M_1_OP:%.*]] = lshr i32 1, [[M_011]]
-; CHECK-NEXT:    [[SEXT_MASK:%.*]] = and i32 [[M_1_OP]], 65535
-; CHECK-NEXT:    [[CMP115:%.*]] = icmp ne i32 [[SEXT_MASK]], 0
+; CHECK-NEXT:    [[CMP115:%.*]] = phi i1 [ true, [[BLOCK1]] ], [ false, [[ENTRY]] ]
 ; CHECK-NEXT:    [[CMP1:%.*]] = or i1 [[CMP_I]], [[CMP115]]
 ; CHECK-NEXT:    [[CONV2:%.*]] = zext i1 [[CMP1]] to i32
 ; CHECK-NEXT:    ret i32 [[CONV2]]
 ;
 entry:
   br label %block2
 
 block1: