[ValueTracking] ComputeKnownBits - minimum leading/trailing zero bits in LSHR/SHL (PR44526)

[RKSimon]

Summary

Followup to D72573 - as detailed in https://blog.regehr.org/archives/1709 we don't make use of the known leading/trailing zeros for shifted values in cases where we don't know the shift amount value.

Stop ValueTracking returning zero for poison shift patterns and use the KnownBits shift helpers directly.

Extend KnownBits::shl to combine all possible shifted combinations if both min/max shift amount values are in range.

[RKSimon]

llvm/test/Transforms/InstCombine/and2.ll

155 ↗

(On Diff #301948)

These tests look like they might even need to be scrapped? And we just limit the transform to Shl instead of logical shifts in general.

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

191 ↗

(On Diff #301948)

Replace with -12345 so we don't have leading zeroes?

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

190 ↗

(On Diff #301948)

Replace with -2147483649 so we don't have trailing zeroes?

[efriedma]

llvm/lib/Analysis/ValueTracking.cpp

1245

We could also do AShr... not really useful for known zeros, but maybe for known ones.

llvm/test/Transforms/InstCombine/and2.ll

155 ↗

(On Diff #301948)

I suspect we might want to continue to favor icmp+zext over lshr.

[RKSimon]

Rebase now that computeKnownBitsFromShiftOperator's callbacks uses KnownBits shift handling.

[RKSimon]

Move the and(logical_shift(1, X), 1) --> zext(X == 0) earlier so SimplifyDemandedInstructionBits doesn't remove the pattern.

[spatel]

llvm/lib/Analysis/ValueTracking.cpp

996–1003

I'm not sure if it changes anything directly here, but we should revisit this conflict logic (similarly at the final block in the function) because we have a poison value in IR with D71126 (see also D92270).

[RKSimon]

llvm/lib/Analysis/ValueTracking.cpp

996–1003

TBH - I don't think we should be setting 'known bits' to arbitrary/poison numbers like this inside ValueTracking - poison handling should be inside the shift combines and we should just return 'unknown' here.

Is this something that needs to be addressed before this patch do you think? I've just kept to the current behaviour here so far.

[spatel]

llvm/lib/Analysis/ValueTracking.cpp

996–1003

I'm not sure if this patch exposes a new way for that clamp to zero to cause trouble, but I agree that we can return unknown state here and try to deal with the poison cases more directly in instsimplify or other passes.

[RKSimon]

Rebase + simplify after D95959

Poison nsw shl patterns no longer return zero and instead bails out - InstSimplify should handle this and return poison.

[RKSimon]

Add shl nsw poison handling to SimplifyShlInst

[spatel]

llvm/lib/Analysis/ValueTracking.cpp

1057–1058

Do you plan to change this independently (or is there some reason not to change it)?

[RKSimon]

llvm/lib/Analysis/ValueTracking.cpp

1057–1058

thanks - my rebase screwed up

[RKSimon]

llvm/lib/Analysis/ValueTracking.cpp

1057–1058

and fixing this unearthes another poison case - looking at this now

[RKSimon]

KnownBits::shl - combine all possible shifted combinations if both min/max shift amount values are in range - helps find some additional poision cases and prevents another regression due to ValueTracking clearing all KnownBits for poison shifts.

[RKSimon]

Add early-out from KnownBits::shl loop if we know no bits

[RKSimon]

rebase

[spatel]

cc @aqjune - another step away from undef. :)

LGTM - hoping that the 'nsw' check before the computeKnownBits means the compile-time impact is minimal, but that's something to watch out for.

[nikic]

llvm/lib/Support/KnownBits.cpp

207

Hm, this code basically repeats what ValueTracking does in a slightly weaker form (using all shifts between min and max, rather than only the possible ones). If we're going to do this inside the KnownBits function itself, then I think we should be doing it only here, not in both places.

[RKSimon]

llvm/lib/Support/KnownBits.cpp

207

I'm intending to clean this up as a followup - cheers

[MaggieYi]

It seems that this commit has caused several test failures.

You can see here: http://lab.llvm.org:8011/#/builders/52/builds/4550

Could you please fix them?
Thanks,

[RKSimon]

In D90479#2585284, @MaggieYi wrote:

It seems that this commit has caused several test failures.

You can see here: http://lab.llvm.org:8011/#/builders/52/builds/4550

Could you please fix them?
Thanks,

I've reverted for now - I took a quick look at the sanitizer test and couldn't grok whats going wrong, will try again when I'm awake :)

[MaggieYi]

I've reverted for now - I took a quick look at the sanitizer test and couldn't grok whats going wrong, will try again when I'm awake :)

Thanks for your quick response :).

llvm/lib/Analysis/InstructionSimplify.cpp

static bool isPoisonShift(Value *Amount, const SimplifyQuery &Q) {
   }
 
   return false;
 }
 
 /// Given operands for an Shl, LShr or AShr, see if we can fold the result.
 /// If not, this returns null.
 static Value *SimplifyShift(Instruction::BinaryOps Opcode, Value *Op0,
-                            Value *Op1, const SimplifyQuery &Q, unsigned MaxRecurse) {
+                            Value *Op1, bool isNSW, const SimplifyQuery &Q,
+                            unsigned MaxRecurse) {
   if (Constant *C = foldOrCommuteConstant(Opcode, Op0, Op1, Q))
     return C;
 
   // 0 shift by X -> 0
   if (match(Op0, m_Zero()))
     return Constant::getNullValue(Op0->getType());
 
   // X shift by 0 -> X
static Value *SimplifyShift(Instruction::BinaryOps Opcode, Value *Op0,
   // If the operation is with the result of a phi instruction, check whether
   // operating on all incoming values of the phi always yields the same value.
   if (isa<PHINode>(Op0) || isa<PHINode>(Op1))
     if (Value *V = ThreadBinOpOverPHI(Opcode, Op0, Op1, Q, MaxRecurse))
       return V;
 
   // If any bits in the shift amount make that value greater than or equal to
   // the number of bits in the type, the shift is undefined.
-  KnownBits Known = computeKnownBits(Op1, Q.DL, 0, Q.AC, Q.CxtI, Q.DT);
-  if (Known.getMinValue().uge(Known.getBitWidth()))
+  KnownBits KnownAmt = computeKnownBits(Op1, Q.DL, 0, Q.AC, Q.CxtI, Q.DT);
+  if (KnownAmt.getMinValue().uge(KnownAmt.getBitWidth()))
     return PoisonValue::get(Op0->getType());
 
+  // Check for nsw shl leading to a poison value.
+  if (isNSW) {
+    assert(Opcode == Instruction::Shl && "Expected shl for nsw instruction");
+    KnownBits KnownVal = computeKnownBits(Op0, Q.DL, 0, Q.AC, Q.CxtI, Q.DT);
+    KnownBits KnownShl = KnownBits::shl(KnownVal, KnownAmt);
+
+    if (KnownVal.Zero.isSignBitSet())
+      KnownShl.Zero.setSignBit();
+    if (KnownVal.One.isSignBitSet())
+      KnownShl.One.setSignBit();
+
+    if (KnownShl.hasConflict())
+      return PoisonValue::get(Op0->getType());
+  }
+
   // If all valid bits in the shift amount are known zero, the first operand is
   // unchanged.
-  unsigned NumValidShiftBits = Log2_32_Ceil(Known.getBitWidth());
-  if (Known.countMinTrailingZeros() >= NumValidShiftBits)
+  unsigned NumValidShiftBits = Log2_32_Ceil(KnownAmt.getBitWidth());
+  if (KnownAmt.countMinTrailingZeros() >= NumValidShiftBits)
     return Op0;
 
   return nullptr;
 }
 
 /// Given operands for an Shl, LShr or AShr, see if we can
 /// fold the result.  If not, this returns null.
 static Value *SimplifyRightShift(Instruction::BinaryOps Opcode, Value *Op0,
                                  Value *Op1, bool isExact, const SimplifyQuery &Q,
                                  unsigned MaxRecurse) {
-  if (Value *V = SimplifyShift(Opcode, Op0, Op1, Q, MaxRecurse))
+  if (Value *V = SimplifyShift(Opcode, Op0, Op1, false, Q, MaxRecurse))
     return V;
 
   // X >> X -> 0
   if (Op0 == Op1)
     return Constant::getNullValue(Op0->getType());
 
   // undef >> X -> 0
   // undef >> X -> undef (if it's exact)
static Value *SimplifyRightShift(Instruction::BinaryOps Opcode, Value *Op0,
 
   return nullptr;
 }
 
 /// Given operands for an Shl, see if we can fold the result.
 /// If not, this returns null.
 static Value *SimplifyShlInst(Value *Op0, Value *Op1, bool isNSW, bool isNUW,
                               const SimplifyQuery &Q, unsigned MaxRecurse) {
-  if (Value *V = SimplifyShift(Instruction::Shl, Op0, Op1, Q, MaxRecurse))
+  if (Value *V =
+          SimplifyShift(Instruction::Shl, Op0, Op1, isNSW, Q, MaxRecurse))
     return V;
 
   // undef << X -> 0
   // undef << X -> undef if (if it's NSW/NUW)
   if (Q.isUndefValue(Op0))
     return isNSW || isNUW ? Op0 : Constant::getNullValue(Op0->getType());
 
   // (X >> A) << A -> X

llvm/lib/Analysis/ValueTracking.cpp

static void computeKnownBitsFromShiftOperator(
 
   // Note: We cannot use Known.Zero.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).
   uint64_t ShiftAmtKZ = Known.Zero.zextOrTrunc(64).getZExtValue();
   uint64_t ShiftAmtKO = Known.One.zextOrTrunc(64).getZExtValue();
   bool ShiftAmtIsConstant = Known.isConstant();
   bool MaxShiftAmtIsOutOfRange = Known.getMaxValue().uge(BitWidth);
 
-  if (ShiftAmtIsConstant) {
-    Known = KF(Known2, Known);
+  // Use the KF callback to get an initial knownbits approximation.
 
-    // If the known bits conflict, this must be an overflowing left shift, so
-    // the shift result is poison. We can return anything we want. Choose 0 for
-    // the best folding opportunity.
-    if (Known.hasConflict())
-      Known.setAllZero();
-
+  // the shift result is poison.
+  if (Known.hasConflict()) {
+    Known.resetAll();
     return;
   }
 
+  // Constant shift amount - we're not going to improve on this.
+  if (ShiftAmtIsConstant)
+    return;
+
   // If the shift amount could be greater than or equal to the bit-width of the
   // LHS, the value could be poison, but bail out because the check below is
   // expensive.
   // TODO: Should we just carry on?
-  if (MaxShiftAmtIsOutOfRange) {
-    Known.resetAll();
+  if (MaxShiftAmtIsOutOfRange)
     return;
-  }
-
-  // It would be more-clearly correct to use the two temporaries for this
-  // calculation. Reusing the APInts here to prevent unnecessary allocations.
-  Known.resetAll();
 
   // If we know the shifter operand is nonzero, we can sometimes infer more
   // known bits. However this is expensive to compute, so be lazy about it and
   // only compute it when absolutely necessary.
   Optional<bool> ShifterOperandIsNonZero;
 
   // Early exit if we can't constrain any well-defined shift amount.
   if (!(ShiftAmtKZ & (PowerOf2Ceil(BitWidth) - 1)) &&
if (ShiftAmt == 0) {
       if (*ShifterOperandIsNonZero)
         continue;
     }
 
     Known = KnownBits::commonBits(
         Known, KF(Known2, KnownBits::makeConstant(APInt(32, ShiftAmt))));
   }
 
-  // If the known bits conflict, the result is poison. Return a 0 and hope the
-  // caller can further optimize that.
+  // If the known bits conflict, the result is poison.
   if (Known.hasConflict())
-    Known.setAllZero();
+    Known.resetAll();
 }
 
 static void computeKnownBitsFromOperator(const Operator *I,
                                          const APInt &DemandedElts,
                                          KnownBits &Known, unsigned Depth,
                                          const Query &Q) {
   unsigned BitWidth = Known.getBitWidth();
 
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);
     break;
   }
   case Instruction::Sub: {
     bool NSW = Q.IIQ.hasNoSignedWrap(cast<OverflowingBinaryOperator>(I));
     computeKnownBitsAddSub(false, I->getOperand(0), I->getOperand(1), NSW,
                            DemandedElts, Known, Known2, Depth, Q);
     break;
   }

llvm/lib/Support/KnownBits.cpp

KnownBits KnownBits::shl(const KnownBits &LHS, const KnownBits &RHS) {
 
   // Minimum shift amount low bits are known zero.
   APInt MinShiftAmount = RHS.getMinValue();
   if (MinShiftAmount.ult(BitWidth)) {
     MinTrailingZeros += MinShiftAmount.getZExtValue();
     MinTrailingZeros = std::min(MinTrailingZeros, BitWidth);
   }
 
+  // If the maximum shift is in range, then find the common bits from all
+  // possible shifts.
+  APInt MaxShiftAmount = RHS.getMaxValue();
+  if (MaxShiftAmount.ult(BitWidth)) {
+    assert(MinShiftAmount.ult(MaxShiftAmount) && "Illegal shift range");
+    Known.Zero.setAllBits();
+    Known.One.setAllBits();
+    for (uint64_t ShiftAmt = MinShiftAmount.getZExtValue(),
+                  MaxShiftAmt = MaxShiftAmount.getZExtValue();
+         ShiftAmt <= MaxShiftAmt; ++ShiftAmt) {
+      KnownBits SpecificShift;
+      SpecificShift.Zero = LHS.Zero << ShiftAmt;
+      SpecificShift.One = LHS.One << ShiftAmt;
+      Known = KnownBits::commonBits(Known, SpecificShift);
+      if (Known.isUnknown())
+        break;
+    }
+  }
+
   Known.Zero.setLowBits(MinTrailingZeros);
   return Known;
 }
 
 KnownBits KnownBits::lshr(const KnownBits &LHS, const KnownBits &RHS) {
   unsigned BitWidth = LHS.getBitWidth();
   KnownBits Known(BitWidth);
 

llvm/test/Transforms/InstCombine/known-signbit-shift.ll

;
   %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. The shift produces poison.
 define i32 @test_no_sign_bit_conflict1(i1 %b) {
 ; CHECK-LABEL: @test_no_sign_bit_conflict1(
-; CHECK-NEXT:    ret i32 undef
+; CHECK-NEXT:    ret i32 poison
 ;
   %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. The shift produces poison.
 define i32 @test_no_sign_bit_conflict2(i1 %b) {
 ; CHECK-LABEL: @test_no_sign_bit_conflict2(
-; CHECK-NEXT:    ret i32 undef
+; CHECK-NEXT:    ret i32 poison
 ;
   %sel = select i1 %b, i32 -8193, i32 -8194
   %mul = shl nsw i32 %sel, 18
   ret i32 %mul
 }

llvm/test/Transforms/InstSimplify/icmp-constant.ll

;
   ret <2 x i1> %cmp
 }
 
 ; PR34838 - https://bugs.llvm.org/show_bug.cgi?id=34838
 ; The shift is known to create poison, so we can simplify the cmp.
 
 define i1 @ne_shl_by_constant_produces_poison(i8 %x) {
 ; CHECK-LABEL: @ne_shl_by_constant_produces_poison(
-; CHECK-NEXT:    ret i1 true
+; CHECK-NEXT:    ret i1 poison
 ;
   %zx = zext i8 %x to i16      ; zx  = 0x00xx
   %xor = xor i16 %zx, 32767    ; xor = 0x7fyy
   %sub = sub nsw i16 %zx, %xor ; sub = 0x80zz  (the top bit is known one)
   %poison = shl nsw i16 %sub, 2    ; oops! this shl can't be nsw; that's POISON
   %cmp = icmp ne i16 %poison, 1
   ret i1 %cmp
 }
 
 define i1 @eq_shl_by_constant_produces_poison(i8 %x) {
 ; CHECK-LABEL: @eq_shl_by_constant_produces_poison(
-; CHECK-NEXT:    ret i1 false
+; CHECK-NEXT:    ret i1 poison
 ;
   %clear_high_bit = and i8 %x, 127                 ; 0x7f
   %set_next_high_bits = or i8 %clear_high_bit, 112 ; 0x70
   %poison = shl nsw i8 %set_next_high_bits, 3
   %cmp = icmp eq i8 %poison, 15
   ret i1 %cmp
 }
 
 ; Shift-by-variable that produces poison is more complicated but still possible.
 ; We guarantee that the shift will change the sign of the shifted value (and
 ; therefore produce poison) by limiting its range from 1 to 3.
 
 define i1 @eq_shl_by_variable_produces_poison(i8 %x) {
 ; CHECK-LABEL: @eq_shl_by_variable_produces_poison(
-; CHECK-NEXT:    ret i1 false
+; CHECK-NEXT:    ret i1 poison
 ;
   %clear_high_bit = and i8 %x, 127                 ; 0x7f
   %set_next_high_bits = or i8 %clear_high_bit, 112 ; 0x70
   %notundef_shiftamt = and i8 %x, 3
   %nonzero_shiftamt = or i8 %notundef_shiftamt, 1
   %poison = shl nsw i8 %set_next_high_bits, %nonzero_shiftamt
   %cmp = icmp eq i8 %poison, 15
   ret i1 %cmp