[KnownBits] Simplify shl. NFCI.

[foad]

@nikic this was a result of me trying to understand your implementation. I tried to:

• remove special cases (but perhaps some of these were deliberate early-outs to speed up to the common cases?)
• work in unsigned instead of APInt as much as possible
• take advantage of APInt::*shl_ov

[nikic]

In D151421#4371619, @foad wrote:

@nikic this was a result of me trying to understand your implementation. I tried to:

• remove special cases (but perhaps some of these were deliberate early-outs to speed up to the common cases?)

Yes, these are intentional fast paths. In fact, the ones we currently have are insufficient, and we'll need another one for the case where the shift amount is unknown, because the current implementation is too slow. (This path is mostly not exposed right now due to special code in ValueTracking, but it if does get exposed it causes substantial compile-time regressions.)

It's possible that dropping the special handling for constants would be fine, but we need to at least keep the code for handling unknown LHS.

• work in unsigned instead of APInt as much as possible
• take advantage of APInt::*shl_ov

These changes look good to me and should make the implementation faster.

[nikic]

In D151421#4371774, @nikic wrote:

In D151421#4371619, @foad wrote:

@nikic this was a result of me trying to understand your implementation. I tried to:

• remove special cases (but perhaps some of these were deliberate early-outs to speed up to the common cases?)

Yes, these are intentional fast paths. In fact, the ones we currently have are insufficient, and we'll need another one for the case where the shift amount is unknown, because the current implementation is too slow. (This path is mostly not exposed right now due to special code in ValueTracking, but it if does get exposed it causes substantial compile-time regressions.)

For reference, that extra fast path is going to look something like this: https://github.com/llvm/llvm-project/commit/6c242bc3d81f87b819e6d705d73a43fb4bf1c5d0 (Which does fix the compile-time regressions.) But I'll return to this once you're done with your changes.

[foad]

Restore and simplify the fast path for unknown LHS.

[nikic]

LGTM

llvm/lib/Support/KnownBits.cpp

210–214

Seems like pulling this out of the isUnknown() branch would be cleaner, as it is not specific to unknown values.

[foad]

llvm/lib/Support/KnownBits.cpp

210–214

Yeah, maybe. This was a result of me trying to work out why this code is required, and it's only required in the LHS.isUnknown() case, otherwise it will naturally fall out in the wash. And since it will almost never happen in practice, I chose to keep it out of the main code path.

210–214

It is also only required in order to guarantee that we return zero for poison results. I've just put D151456 to check that property so we don't regress it.

llvm/lib/Support/KnownBits.cpp

KnownBits KnownBits::smin(const KnownBits &LHS, const KnownBits &RHS) {
   };
   return Flip(umax(Flip(LHS), Flip(RHS)));
 }
 
 KnownBits KnownBits::shl(const KnownBits &LHS, const KnownBits &RHS, bool NUW,
                          bool NSW) {
   unsigned BitWidth = LHS.getBitWidth();
   auto ShiftByConst = [&](const KnownBits &LHS,
-                          uint64_t ShiftAmt) -> std::optional<KnownBits> {
+                          unsigned ShiftAmt) -> std::optional<KnownBits> {
     KnownBits Known;
-    Known.Zero = LHS.Zero << ShiftAmt;
+    bool ShiftedOutZero, ShiftedOutOne;
+    Known.Zero = LHS.Zero.ushl_ov(ShiftAmt, ShiftedOutZero);
     Known.Zero.setLowBits(ShiftAmt);
-    Known.One = LHS.One << ShiftAmt;
-    if ((!NUW && !NSW) || ShiftAmt == 0)
-      return Known;
+    Known.One = LHS.One.ushl_ov(ShiftAmt, ShiftedOutOne);
 
-    KnownBits ShiftedOutBits = LHS.extractBits(ShiftAmt, BitWidth - ShiftAmt);
-    if (NUW && !ShiftedOutBits.One.isZero())
+    if (NUW) {
+      if (ShiftedOutOne)
         // One bit has been shifted out.
         return std::nullopt;
+      if (ShiftAmt != 0)
+        // NUW means we can assume anything shifted out was a zero.
+        ShiftedOutZero = true;
+    }
+
     if (NSW) {
-      if (!ShiftedOutBits.Zero.isZero() && !ShiftedOutBits.One.isZero())
+      if (ShiftedOutZero && ShiftedOutOne)
         // Both zeros and ones have been shifted out.
         return std::nullopt;
-      if (NUW || !ShiftedOutBits.Zero.isZero()) {
+      if (ShiftedOutZero) {
         if (Known.isNegative())
           // Zero bit has been shifted out, but result sign is negative.
           return std::nullopt;
         Known.makeNonNegative();
-      } else if (!ShiftedOutBits.One.isZero()) {
+      } else if (ShiftedOutOne) {
         if (Known.isNonNegative())
-          // One bit has been shifted out, but result sign is negative.
+          // One bit has been shifted out, but result sign is non-negative.
           return std::nullopt;
         Known.makeNegative();
       }
     }
     return Known;
   };
 
-  // If the shift amount is a valid constant then transform LHS directly.
-  if (RHS.isConstant() && RHS.getConstant().ult(BitWidth)) {
-    if (auto Res = ShiftByConst(LHS, RHS.getConstant().getZExtValue()))
-      return *Res;
-    KnownBits Known(BitWidth);
-    Known.setAllZero();
-    return Known;
-  }
-
+  // Fast path for a common case when LHS is completely unknown.
   KnownBits Known(BitWidth);
-  APInt MinShiftAmount = RHS.getMinValue();
-  if (MinShiftAmount.uge(BitWidth)) {
+  unsigned MinShiftAmount = RHS.getMinValue().getLimitedValue(BitWidth);
+  if (LHS.isUnknown()) {
+    if (MinShiftAmount == BitWidth) {
       // Always poison. Return zero because we don't like returning conflict.
       Known.setAllZero();
       return Known;
     }
-
-  if (LHS.isUnknown()) {
-    // No matter the shift amount, the trailing zeros will stay zero.
-    unsigned MinTrailingZeros = LHS.countMinTrailingZeros();
-    // Minimum shift amount low bits are known zero.
-    MinTrailingZeros += MinShiftAmount.getZExtValue();
-    MinTrailingZeros = std::min(MinTrailingZeros, BitWidth);
-    Known.Zero.setLowBits(MinTrailingZeros);
-    if (NUW && NSW && !MinShiftAmount.isZero())
+    Known.Zero.setLowBits(MinShiftAmount);
+    if (NUW && NSW && MinShiftAmount != 0)
       Known.makeNonNegative();
     return Known;
   }
 
   // Find the common bits from all possible shifts.
-  APInt MaxShiftAmount = RHS.getMaxValue();
-  uint64_t ShiftAmtZeroMask = (~RHS.Zero).zextOrTrunc(64).getZExtValue();
-  uint64_t ShiftAmtOneMask = RHS.One.zextOrTrunc(64).getZExtValue();
-  assert(MinShiftAmount.ult(MaxShiftAmount) && "Illegal shift range");
+  unsigned MaxShiftAmount = RHS.getMaxValue().getLimitedValue(BitWidth - 1);
+  unsigned ShiftAmtZeroMask = RHS.Zero.zextOrTrunc(32).getZExtValue();
+  unsigned ShiftAmtOneMask = RHS.One.zextOrTrunc(32).getZExtValue();
   Known.Zero.setAllBits();
   Known.One.setAllBits();
-  for (uint64_t ShiftAmt = MinShiftAmount.getZExtValue(),
-                MaxShiftAmt = MaxShiftAmount.getLimitedValue(BitWidth - 1);
-       ShiftAmt <= MaxShiftAmt; ++ShiftAmt) {
+  for (unsigned ShiftAmt = MinShiftAmount; ShiftAmt <= MaxShiftAmount;
+       ++ShiftAmt) {
     // Skip if the shift amount is impossible.
-    if ((ShiftAmtZeroMask & ShiftAmt) != ShiftAmt ||
+    if ((ShiftAmtZeroMask & ShiftAmt) != 0 ||
         (ShiftAmtOneMask | ShiftAmt) != ShiftAmt)
       continue;
     auto Res = ShiftByConst(LHS, ShiftAmt);
     if (!Res)
       // All larger shift amounts will overflow as well.
       break;
     Known = Known.intersectWith(*Res);
     if (Known.isUnknown())
       break;
   }
 
   // All shift amounts may result in poison.
-  if (Known.hasConflict()) {
-    assert((NUW || NSW) && "Can only happen with nowrap flags");
+  if (Known.hasConflict())
     Known.setAllZero();
-  }
   return Known;
 }
 
 KnownBits KnownBits::lshr(const KnownBits &LHS, const KnownBits &RHS) {
   unsigned BitWidth = LHS.getBitWidth();
   KnownBits Known(BitWidth);
 
   if (RHS.isConstant() && RHS.getConstant().ult(BitWidth)) {