[KnownBits] Add support for nuw/nsw on shifts

[nikic]

Summary

Implement precise nuw/nsw support in the KnownBits implementation, replacing the rather crude handling in ValueTracking.

[goldstein.w.n]

llvm/lib/Support/KnownBits.cpp

176

nit style: Instead of nesting if statements, can you invert them and just return immediately? Think it makes code easier to read.

213

Imo, this check should just be at the top. Then you can drop the RHS.getConstant().ult(BitWidth) check before ShiftByConst

232

in the NSW/NUW we should be able to further bound MaxShiftAmount with LHS.countLeadingZeros no? (and poison if MinShiftAmunt > LHS.countLeadingZeros)

253–258

can you assert no conflict? Its not super obvious just reading the code that the loop will be enough to clear the all zeros/ones initial state in Known.

[nikic]

Rebase, check for conflicts, use early return.

[nikic]

llvm/lib/Support/KnownBits.cpp

176

I've done this for the first check here and also for the isUnknown checks in https://github.com/llvm/llvm-project/commit/2d48a771fc00681414c54ea3936bff91f3b253c4.

232

Yes, but this is already implicitly handled by the loop below. It stops once it gets the first shift amount that results in poison due to nuw/nsw.

253–258

You are right, it's necessary to check for conflicts here. I've added a general assertion for this in the exhaustive tests, so we don't have to remember checking it for every function: https://github.com/llvm/llvm-project/commit/d4bfc144eaa6e34eafaab1e8574421ebd2bb78ed

[goldstein.w.n]

LGTM.

llvm/lib/Support/KnownBits.cpp

232

Ah, fair enough.

255

This should only ever occur if NUW or NSW right? Maybe assert that?

llvm/include/llvm/Support/KnownBits.h

public:
   /// Compute known bits for smax(LHS, RHS).
   static KnownBits smax(const KnownBits &LHS, const KnownBits &RHS);
 
   /// Compute known bits for smin(LHS, RHS).
   static KnownBits smin(const KnownBits &LHS, const KnownBits &RHS);
 
   /// Compute known bits for shl(LHS, RHS).
   /// NOTE: RHS (shift amount) bitwidth doesn't need to be the same as LHS.
-  static KnownBits shl(const KnownBits &LHS, const KnownBits &RHS);
+  static KnownBits shl(const KnownBits &LHS, const KnownBits &RHS,
+                       bool NUW = false, bool NSW = false);
 
   /// Compute known bits for lshr(LHS, RHS).
   /// NOTE: RHS (shift amount) bitwidth doesn't need to be the same as LHS.
   static KnownBits lshr(const KnownBits &LHS, const KnownBits &RHS);
 
   /// Compute known bits for ashr(LHS, RHS).
   /// NOTE: RHS (shift amount) bitwidth doesn't need to be the same as LHS.
   static KnownBits ashr(const KnownBits &LHS, const KnownBits &RHS);

llvm/lib/Analysis/ValueTracking.cpp

case Instruction::SExt: {
     Known = Known.trunc(SrcBitWidth);
     computeKnownBits(I->getOperand(0), Known, Depth + 1, Q);
     // If the sign bit of the input is known set or clear, then we know the
     // top bits of the result.
     Known = Known.sext(BitWidth);
     break;
   }
   case Instruction::Shl: {
+    bool NUW = Q.IIQ.hasNoUnsignedWrap(cast<OverflowingBinaryOperator>(I));
     bool NSW = Q.IIQ.hasNoSignedWrap(cast<OverflowingBinaryOperator>(I));
-    auto KF = [NSW](const KnownBits &KnownVal, const KnownBits &KnownAmt) {
-      KnownBits Result = KnownBits::shl(KnownVal, KnownAmt);
-      // 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 (NSW) {
-        if (KnownVal.Zero.isSignBitSet())
-          Result.Zero.setSignBit();
-        if (KnownVal.One.isSignBitSet())
-          Result.One.setSignBit();
-        if (Result.hasConflict())
-          Result.setAllZero();
-      }
-      return Result;
+    auto KF = [NUW, NSW](const KnownBits &KnownVal, const KnownBits &KnownAmt) {
+      return KnownBits::shl(KnownVal, KnownAmt, NUW, NSW);
     };
     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->countr_zero());
     break;

llvm/lib/Support/KnownBits.cpp

auto Flip = [](const KnownBits &Val) {
     APInt One = Val.Zero;
     Zero.setBitVal(SignBitPosition, Val.Zero[SignBitPosition]);
     One.setBitVal(SignBitPosition, Val.One[SignBitPosition]);
     return KnownBits(Zero, One);
   };
   return Flip(umax(Flip(LHS), Flip(RHS)));
 }
 
-KnownBits KnownBits::shl(const KnownBits &LHS, const KnownBits &RHS) {
+KnownBits KnownBits::shl(const KnownBits &LHS, const KnownBits &RHS, bool NUW,
+                         bool NSW) {
   unsigned BitWidth = LHS.getBitWidth();
-  KnownBits Known(BitWidth);
+  auto ShiftByConst = [&](const KnownBits &LHS,
+                          uint64_t ShiftAmt) -> std::optional<KnownBits> {
+    KnownBits Known;
+    Known.Zero = LHS.Zero << ShiftAmt;
+    Known.Zero.setLowBits(ShiftAmt);
+    Known.One = LHS.One << ShiftAmt;
+    if ((!NUW && !NSW) || ShiftAmt == 0)
+      return Known;
+
+    KnownBits ShiftedOutBits = LHS.extractBits(ShiftAmt, BitWidth - ShiftAmt);
+    if (NUW && !ShiftedOutBits.One.isZero())
+      // One bit has been shifted out.
+      return std::nullopt;
+    if (NSW) {
+      if (!ShiftedOutBits.Zero.isZero() && !ShiftedOutBits.One.isZero())
+        // Both zeros and ones have been shifted out.
+        return std::nullopt;
+      if (NUW || !ShiftedOutBits.Zero.isZero()) {
+        if (Known.isNegative())
+          // Zero bit has been shifted out, but result sign is negative.
+          return std::nullopt;
+        Known.makeNonNegative();
+      } else if (!ShiftedOutBits.One.isZero()) {
+        if (Known.isNonNegative())
+          // One bit has been shifted out, but result sign is 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)) {
-    unsigned Shift = RHS.getConstant().getZExtValue();
-    Known = LHS;
-    Known.Zero <<= Shift;
-    Known.One <<= Shift;
-    // Low bits are known zero.
-    Known.Zero.setLowBits(Shift);
+    if (auto Res = ShiftByConst(LHS, RHS.getConstant().getZExtValue()))
+      return *Res;
+    KnownBits Known(BitWidth);
+    Known.setAllZero();
     return Known;
   }
 
+  KnownBits Known(BitWidth);
   APInt MinShiftAmount = RHS.getMinValue();
   if (MinShiftAmount.uge(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.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");
   Known.Zero.setAllBits();
   Known.One.setAllBits();
   for (uint64_t ShiftAmt = MinShiftAmount.getZExtValue(),
                 MaxShiftAmt = MaxShiftAmount.getLimitedValue(BitWidth - 1);
        ShiftAmt <= MaxShiftAmt; ++ShiftAmt) {
     // Skip if the shift amount is impossible.
     if ((ShiftAmtZeroMask & ShiftAmt) != ShiftAmt ||
         (ShiftAmtOneMask | ShiftAmt) != ShiftAmt)
       continue;
-    KnownBits SpecificShift;
-    SpecificShift.Zero = LHS.Zero << ShiftAmt;
-    SpecificShift.Zero.setLowBits(ShiftAmt);
-    SpecificShift.One = LHS.One << ShiftAmt;
-    Known = Known.intersectWith(SpecificShift);
+    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");
+    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)) {

llvm/test/Transforms/ConstraintElimination/geps-unsigned-predicates.ll

 ; CHECK-NEXT:    ret void
 ; CHECK:       check.idx:
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[IDX:%.*]], 5
 ; CHECK-NEXT:    br i1 [[CMP]], label [[CHECK_MAX:%.*]], label [[TRAP]]
 ; CHECK:       check.max:
 ; CHECK-NEXT:    [[IDX_SHL_1:%.*]] = shl nuw nsw i8 [[IDX]], 1
 ; CHECK-NEXT:    [[ADD_PTR_SHL_1:%.*]] = getelementptr inbounds i32, ptr [[SRC]], i8 [[IDX_SHL_1]]
 ; CHECK-NEXT:    [[C_MAX_0:%.*]] = icmp ult ptr [[ADD_PTR_SHL_1]], [[MAX]]
-; CHECK-NEXT:    call void @use(i1 [[C_MAX_0]])
+; CHECK-NEXT:    call void @use(i1 true)
 ; CHECK-NEXT:    [[IDX_SHL_2:%.*]] = shl nuw i8 [[IDX]], 2
 ; CHECK-NEXT:    [[ADD_PTR_SHL_2:%.*]] = getelementptr inbounds i32, ptr [[SRC]], i8 [[IDX_SHL_2]]
 ; CHECK-NEXT:    [[C_MAX_1:%.*]] = icmp ult ptr [[ADD_PTR_SHL_2]], [[MAX]]
 ; CHECK-NEXT:    call void @use(i1 [[C_MAX_1]])
 ; CHECK-NEXT:    [[IDX_SHL_NOT_NUW:%.*]] = shl i8 [[IDX]], 1
 ; CHECK-NEXT:    [[ADD_PTR_SHL_NOT_NUW:%.*]] = getelementptr inbounds i32, ptr [[SRC]], i8 [[IDX_SHL_NOT_NUW]]
 ; CHECK-NEXT:    [[C_MAX_2:%.*]] = icmp ult ptr [[ADD_PTR_SHL_NOT_NUW]], [[MAX]]
 ; CHECK-NEXT:    call void @use(i1 [[C_MAX_2]])

llvm/test/Transforms/LoopVectorize/AArch64/sve-interleaved-accesses.ll

 ; CHECK-NEXT:    [[LOAD2:%.*]] = load i32, ptr [[ARRAYIDX1]], align 4
 ; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[LOAD1]], [[C]]
 ; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i32 [[LOAD2]], [[D]]
 ; CHECK-NEXT:    [[ARRAYIDX2:%.*]] = getelementptr inbounds [1024 x i32], ptr @CD, i64 0, i64 [[INDVARS_IV]]
 ; CHECK-NEXT:    store i32 [[ADD]], ptr [[ARRAYIDX2]], align 4
 ; CHECK-NEXT:    [[ARRAYIDX3:%.*]] = getelementptr inbounds [1024 x i32], ptr @CD, i64 0, i64 [[OR]]
 ; CHECK-NEXT:    store i32 [[MUL]], ptr [[ARRAYIDX3]], align 4
 ; CHECK-NEXT:    [[INDVARS_IV_NEXT]] = add nuw nsw i64 [[INDVARS_IV]], 2
-; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i64 [[INDVARS_IV]], 1022
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[INDVARS_IV]], 1022
 ; CHECK-NEXT:    br i1 [[CMP]], label [[FOR_BODY]], label [[FOR_END]], !llvm.loop [[LOOP3:![0-9]+]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %for.body
 
 for.body:                                         ; preds = %for.body, %entry

llvm/unittests/Support/KnownBitsTest.cpp

testBinaryOpExhaustive(
       },
       [](const APInt &N1, const APInt &N2) -> std::optional<APInt> {
         if (N2.uge(N2.getBitWidth()))
           return std::nullopt;
         return N1.shl(N2);
       });
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
+        return KnownBits::shl(Known1, Known2, /* NUW */ true);
+      },
+      [](const APInt &N1, const APInt &N2) -> std::optional<APInt> {
+        bool Overflow;
+        APInt Res = N1.ushl_ov(N2, Overflow);
+        if (Overflow)
+          return std::nullopt;
+        return Res;
+      });
+  testBinaryOpExhaustive(
+      [](const KnownBits &Known1, const KnownBits &Known2) {
+        return KnownBits::shl(Known1, Known2, /* NUW */ false, /* NSW */ true);
+      },
+      [](const APInt &N1, const APInt &N2) -> std::optional<APInt> {
+        bool Overflow;
+        APInt Res = N1.sshl_ov(N2, Overflow);
+        if (Overflow)
+          return std::nullopt;
+        return Res;
+      });
+  testBinaryOpExhaustive(
+      [](const KnownBits &Known1, const KnownBits &Known2) {
+        return KnownBits::shl(Known1, Known2, /* NUW */ true, /* NSW */ true);
+      },
+      [](const APInt &N1, const APInt &N2) -> std::optional<APInt> {
+        bool OverflowUnsigned, OverflowSigned;
+        APInt Res = N1.ushl_ov(N2, OverflowUnsigned);
+        (void)N1.sshl_ov(N2, OverflowSigned);
+        if (OverflowUnsigned || OverflowSigned)
+          return std::nullopt;
+        return Res;
+      });
+
+  testBinaryOpExhaustive(
+      [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::lshr(Known1, Known2);
       },
       [](const APInt &N1, const APInt &N2) -> std::optional<APInt> {
         if (N2.uge(N2.getBitWidth()))
           return std::nullopt;
         return N1.lshr(N2);
       });
   testBinaryOpExhaustive(