[ValueTracking] Sign bit of shl cannot be both known one and known zero. Fix PR28926

[lihuang]

Summary

[ValueTracking] Sign bit of shl cannot be both known one and known zero. Fix PR28926

This is to fix PR28296 caused by r278172 (D23296). When inferring the sign bit of a SHL from nsw flag, should update both KnownZero and KnownOne's sign bit. If the result of SHL is known to be non-negative, should set the sign bit of KnownZero to 1 and KnownOne to 0, vice versa. This cannot be implemented in KZF and KOF, because KZF only modifies knownZero and KOF only modifies knownOne.

Also made a little change to computeKnownBitsFromShiftOperator: extract the computeKnownBits(...) for operand 0 to its callsite, so that the "infer sign bit from SHL with nsw" logic can be implemented only at the SHL case.

[majnemer]

lib/Analysis/ValueTracking.cpp

842–850

Shouldn't we just add this logic to the ConstantInt branch above, right before the return?
Then we could teach InstSimplify to xform this into the non-undef arm of the select.

[sanjoy]

I've not reviewed the patch yet, but neither of the two PR numbers mentioned here seem related -- can you please check?

lib/Analysis/ValueTracking.cpp

1078

Can you do the same thing by a "normalization" step like:

// Bits that are known to be "both" zero and one are arbitrarily selected to be zero.
KnownZero |= (KnownZero & KnownOne);
KnownOne &= ~KnownZero;

?

If not, now that you have this logic here, can we clear out the logic
from KZF and KOF?

Also, were you able to pinpoint _why_ we ended up with a bit both proved
to be 0 and 1?

[lihuang]

Sorry, the related PR number is 28946
https://llvm.org/bugs/show_bug.cgi?id=28946

[lihuang]

lib/Analysis/ValueTracking.cpp

1078

We can't arbitrarily select 0 if they are both known to be zero, as this will hide some cases where the shift result could be actually inferred as having the same sign (negative) as its first operand. One example is the last test I added in known-signbit-shift.ll.

Sure, I should have done that, my mistake,

This is because we could have already shifted a KnownOne bit to the sign-bit by KOF when we are setting the inferred sign-bit of KnownZero in KZF, and also the other way around. For example:

  before shift:    KnownZero = 0x70000000
                         KnownOne =  0x40000000
after SHL by 1:  KnownZero = 0x00000000   inferred to 0x70000000
                          KnownOne = 0x70000000 (conflict)

We should keep the sign-bit of KnownZero and KnownOne to be the same as the first operand of shift. Actually this inspired me to fix the problem just in KZF and KOF. Anyway, I will make another patch after you revert r278172.

[lihuang]

lib/Analysis/ValueTracking.cpp

842–850

Thank you for the advice David. We could do this, but I think I still need to fix the logic in the SHL case of computeKnownBitsFromOperator as that's the root cause.

[lihuang]

Abandon this patch as r278172 has been reverted.

lib/Analysis/ValueTracking.cpp

// assume(v <_u c)
         KnownZero |=
           APInt::getHighBitsSet(BitWidth, RHSKnownZero.countLeadingOnes());
     }
   }
 }
 
 // Compute known bits from a shift operator, including those with a
 // non-constant shift amount. KnownZero and KnownOne are the outputs of this
-// function. KnownZero2 and KnownOne2 are pre-allocated temporaries with the
-// same bit width as KnownZero and KnownOne. KZF and KOF are operator-specific
+// function. Op0KnownZero and Op0KnownOne are pre-computed known-bit results
+// of the first operand of this shift. KZF and KOF are operator-specific
 // functors that, given the known-zero or known-one bits respectively, and a
 // shift amount, compute the implied known-zero or known-one bits of the shift
 // operator's result respectively for that shift amount. The results from calling
 // KZF and KOF are conservatively combined for all permitted shift amounts.
 template <typename KZFunctor, typename KOFunctor>
 static void computeKnownBitsFromShiftOperator(Operator *I,
               APInt &KnownZero, APInt &KnownOne,
-              APInt &KnownZero2, APInt &KnownOne2,
+              APInt &Op0KnownZero, APInt &Op0KnownOne,
               unsigned Depth, const Query &Q, KZFunctor KZF, KOFunctor KOF) {
   unsigned BitWidth = KnownZero.getBitWidth();
 
   if (auto *SA = dyn_cast<ConstantInt>(I->getOperand(1))) {
     unsigned ShiftAmt = SA->getLimitedValue(BitWidth-1);
 
-    computeKnownBits(I->getOperand(0), KnownZero, KnownOne, Depth + 1, Q);
-    KnownZero = KZF(KnownZero, ShiftAmt);
-    KnownOne  = KOF(KnownOne, ShiftAmt);
+    KnownZero = KZF(Op0KnownZero, ShiftAmt);
+    KnownOne  = KOF(Op0KnownOne, ShiftAmt);
     return;
   }
 
   computeKnownBits(I->getOperand(1), KnownZero, KnownOne, Depth + 1, Q);
 
   // Note: We cannot use KnownZero.getLimitedValue() here, because if
   // BitWidth > 64 and any upper bits are known, we'll end up returning the
   // limit value (which implies all bits are known).
static void computeKnownBitsFromShiftOperator(Operator *I,
   // Early exit if we can't constrain any well-defined shift amount.
   if (!(ShiftAmtKZ & (BitWidth - 1)) && !(ShiftAmtKO & (BitWidth - 1))) {
     ShifterOperandIsNonZero =
         isKnownNonZero(I->getOperand(1), Depth + 1, Q);
     if (!*ShifterOperandIsNonZero)
       return;
   }
 
-  computeKnownBits(I->getOperand(0), KnownZero2, KnownOne2, Depth + 1, Q);
-
   KnownZero = KnownOne = APInt::getAllOnesValue(BitWidth);
   for (unsigned ShiftAmt = 0; ShiftAmt < BitWidth; ++ShiftAmt) {
     // Combine the shifted known input bits only for those shift amounts
     // compatible with its known constraints.
     if ((ShiftAmt & ~ShiftAmtKZ) != ShiftAmt)
       continue;
     if ((ShiftAmt | ShiftAmtKO) != ShiftAmt)
       continue;
     // If we know the shifter is nonzero, we may be able to infer more known
     // bits. This check is sunk down as far as possible to avoid the expensive
     // call to isKnownNonZero if the cheaper checks above fail.
     if (ShiftAmt == 0) {
       if (!ShifterOperandIsNonZero.hasValue())
         ShifterOperandIsNonZero =
             isKnownNonZero(I->getOperand(1), Depth + 1, Q);
       if (*ShifterOperandIsNonZero)
         continue;
     }
 
-    KnownZero &= KZF(KnownZero2, ShiftAmt);
-    KnownOne  &= KOF(KnownOne2, ShiftAmt);
+    KnownZero &= KZF(Op0KnownZero, ShiftAmt);
+    KnownOne  &= KOF(Op0KnownOne, ShiftAmt);
   }
 
   // If there are no compatible shift amounts, then we've proven that the shift
   // amount must be >= the BitWidth, and the result is undefined. We could
   // return anything we'd like, but we need to make sure the sets of known bits
   // stay disjoint (it should be better for some other code to actually
   // propagate the undef than to pick a value here using known bits).
   if ((KnownZero & KnownOne) != 0) {
     KnownZero.clearAllBits();
     KnownOne.clearAllBits();
   }
 }
 
 static void computeKnownBitsFromOperator(Operator *I, APInt &KnownZero,
                                          APInt &KnownOne, unsigned Depth,
                                          const Query &Q) {
   unsigned BitWidth = KnownZero.getBitWidth();
 
   APInt KnownZero2(KnownZero), KnownOne2(KnownOne);
case Instruction::Shl: {
     };
 
     auto KOF = [BitWidth, NSW](const APInt &KnownOne, unsigned ShiftAmt) {
       APInt KOResult = KnownOne << ShiftAmt;
       if (NSW && KnownOne.isNegative())
         KOResult.setBit(BitWidth - 1);
       return KOResult;
     };
   
+    computeKnownBits(I->getOperand(0), KnownZero2, KnownOne2, Depth + 1, Q);
     computeKnownBitsFromShiftOperator(I, KnownZero, KnownOne,
                                       KnownZero2, KnownOne2, Depth, Q, KZF,
                                       KOF);
+    // If this shift has "nsw" keyword, then the result is either a poison 
+    // value or has the same sign bit as the first operand. 
+    if (cast<OverflowingBinaryOperator>(I)->hasNoSignedWrap()) {
+      if (KnownZero2.isNegative()) {
+          KnownZero.setBit(BitWidth - 1);
+          KnownOne.clearBit(BitWidth - 1);
+      }
+      else if (KnownOne2.isNegative()) {
+          KnownOne.setBit(BitWidth - 1);
+          KnownZero.clearBit(BitWidth - 1);
+      }
+    }
     break;
   }
   case Instruction::LShr: {
     // (ushr X, C1) & C2 == 0   iff  (-1 >> C1) & C2 == 0
     auto KZF = [BitWidth](const APInt &KnownZero, unsigned ShiftAmt) {
       return APIntOps::lshr(KnownZero, ShiftAmt) |
              // High bits known zero.
              APInt::getHighBitsSet(BitWidth, ShiftAmt);
     };
 
     auto KOF = [BitWidth](const APInt &KnownOne, unsigned ShiftAmt) {
       return APIntOps::lshr(KnownOne, ShiftAmt);
     };
 
+    computeKnownBits(I->getOperand(0), KnownZero2, KnownOne2, Depth + 1, Q);
     computeKnownBitsFromShiftOperator(I, KnownZero, KnownOne,
                                       KnownZero2, KnownOne2, Depth, Q, KZF,
                                       KOF);
     break;
   }
   case Instruction::AShr: {
     // (ashr X, C1) & C2 == 0   iff  (-1 >> C1) & C2 == 0
     auto KZF = [BitWidth](const APInt &KnownZero, unsigned ShiftAmt) {
       return APIntOps::ashr(KnownZero, ShiftAmt);
     };
 
     auto KOF = [BitWidth](const APInt &KnownOne, unsigned ShiftAmt) {
       return APIntOps::ashr(KnownOne, ShiftAmt);
     };
 
+    computeKnownBits(I->getOperand(0), KnownZero2, KnownOne2, Depth + 1, Q);
     computeKnownBitsFromShiftOperator(I, KnownZero, KnownOne,
                                       KnownZero2, KnownOne2, Depth, Q, KZF,
                                       KOF);
     break;
   }
   case Instruction::Sub: {
     bool NSW = cast<OverflowingBinaryOperator>(I)->hasNoSignedWrap();
     computeKnownBitsAddSub(false, I->getOperand(0), I->getOperand(1), NSW,

test/Analysis/ValueTracking/known-signbit-shift.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
-; Result of left shifting a non-negative integer 
+; Result of left shifting a non-negative integer
 ; with nsw flag should also be non-negative
 define i1 @test_shift_nonnegative(i32 %a) {
 ; CHECK-LABEL: @test_shift_nonnegative(
-; CHECK: ret i1 true
+; CHECK-NEXT:    ret i1 true
+;
   %b = lshr i32 %a, 2
   %shift = shl nsw i32 %b, 3
   %cmp = icmp sge i32 %shift, 0
   ret i1 %cmp
 }
 
 ; Result of left shifting a negative integer with
 ; nsw flag should also be negative
 define i1 @test_shift_negative(i32 %a, i32 %b) {
 ; CHECK-LABEL: @test_shift_negative(
-; CHECK: ret i1 true
+; CHECK-NEXT:    ret i1 true
+;
   %c = or i32 %a, -2147483648
   %d = and i32 %b, 7
   %shift = shl nsw i32 %c, %d
   %cmp = icmp slt i32 %shift, 0
   ret i1 %cmp
 }
+
+; If sign bit is a known zero, it cannot be a known one.
+; This test should not crash opt.
+define i32 @test_no_sign_bit_conflict1(i1 %b) {
+; CHECK-LABEL: @test_no_sign_bit_conflict1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 %b, i32 -2147221504, i32 -2147483648
+; CHECK-NEXT:    ret i32 [[SEL]]
+;
+entry:
+  %sel = select i1 %b, i32 8193, i32 8192
+  %mul = shl nsw i32 %sel, 18
+  ret i32 %mul
+}
+
+; If sign bit is a known one, it cannot be a known zero.
+; This test should not crash opt.
+define i32 @test_no_sign_bit_conflict2(i1 %b) {
+; CHECK-LABEL: @test_no_sign_bit_conflict2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 %b, i32 2147221504, i32 2146959360
+; CHECK-NEXT:    ret i32 [[SEL]]
+;
+entry:
+  %sel = select i1 %b, i32 -8193, i32 -8194
+  %mul = shl nsw i32 %sel, 18
+  ret i32 %mul
+}