[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/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/InstSimplify/icmp-constant.ll

 }
 
 ; 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