KnownBits: refine high-bits of mul in signed case

KnownBits::mul suffers from the deficiency that it doesn't account for
signed inputs. Fix it by refining known leading zeros when both inputs
are signed, and setting known leading ones when one of the inputs is
signed. The strategy we've used is to still use umul_ov, after adjusting
for signed inputs, and setting known leading ones from the negation of
the result, when it is known to be negative, noting that a possibly-zero
result is a special case.

Author: artagnon

llvm/lib/Support/KnownBits.cpp

@@ -796,19 +796,26 @@ KnownBits KnownBits::mul(const KnownBits &LHS, const KnownBits &RHS,
   assert((!NoUndefSelfMultiply || LHS == RHS) &&
          "Self multiplication knownbits mismatch");
 
-  // Compute the high known-0 bits by multiplying the unsigned max of each side.
-  // Conservatively, M active bits * N active bits results in M + N bits in the
-  // result. But if we know a value is a power-of-2 for example, then this
-  // computes one more leading zero.
-  // TODO: This could be generalized to number of sign bits (negative numbers).
-  APInt UMaxLHS = LHS.getMaxValue();
-  APInt UMaxRHS = RHS.getMaxValue();
-
-  // For leading zeros in the result to be valid, the unsigned max product must
+  // Compute the high known-0 or known-1 bits by multiplying the max of each
+  // side. Conservatively, M active bits * N active bits results in M + N bits
+  // in the result. But if we know a value is a power-of-2 for example, then
+  // this computes one more leading zero or one.
+  APInt MaxLHS = LHS.isNegative() ? LHS.getMinValue().abs() : LHS.getMaxValue(),
+        MaxRHS = RHS.isNegative() ? RHS.getMinValue().abs() : RHS.getMaxValue();
+
+  // For leading zeros or ones in the result to be valid, the max product must
   // fit in the bitwidth (it must not overflow).
   bool HasOverflow;
-  APInt UMaxResult = UMaxLHS.umul_ov(UMaxRHS, HasOverflow);
-  unsigned LeadZ = HasOverflow ? 0 : UMaxResult.countl_zero();
+  APInt Result = MaxLHS.umul_ov(MaxRHS, HasOverflow);
+  bool NegResult = LHS.isNegative() ^ RHS.isNegative();
+  unsigned LeadZ = 0, LeadO = 0;
+  if (!HasOverflow) {
+    // Do not set leading ones unless the result is known to be non-zero.
+    if (NegResult && LHS.isNonZero() && RHS.isNonZero())
+      LeadO = (-Result).countLeadingOnes();
+    else if (!NegResult)
+      LeadZ = Result.countLeadingZeros();
+  }
 
   // The result of the bottom bits of an integer multiply can be
   // inferred by looking at the bottom bits of both operands and
@@ -873,8 +880,9 @@ KnownBits KnownBits::mul(const KnownBits &LHS, const KnownBits &RHS,
 
   KnownBits Res(BitWidth);
   Res.Zero.setHighBits(LeadZ);
+  Res.One.setHighBits(LeadO);
   Res.Zero |= (~BottomKnown).getLoBits(ResultBitsKnown);
-  Res.One = BottomKnown.getLoBits(ResultBitsKnown);
+  Res.One |= BottomKnown.getLoBits(ResultBitsKnown);
 
   // If we're self-multiplying then bit[1] is guaranteed to be zero.
   if (NoUndefSelfMultiply && BitWidth > 1) {

llvm/test/Analysis/ValueTracking/knownbits-mul.ll

@@ -72,12 +72,7 @@ define i8 @mul_high_bits_know(i8 %xx, i8 %yy) {
 define i8 @mul_high_bits_know2(i8 %xx, i8 %yy) {
 ; CHECK-LABEL: define i8 @mul_high_bits_know2(
 ; CHECK-SAME: i8 [[XX:%.*]], i8 [[YY:%.*]]) {
-; CHECK-NEXT:    [[X:%.*]] = or i8 [[XX]], -2
-; CHECK-NEXT:    [[Y:%.*]] = and i8 [[YY]], 4
-; CHECK-NEXT:    [[Y_NONZERO:%.*]] = or disjoint i8 [[Y]], 1
-; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i8 [[X]], [[Y_NONZERO]]
-; CHECK-NEXT:    [[R:%.*]] = and i8 [[MUL]], -16
-; CHECK-NEXT:    ret i8 [[R]]
+; CHECK-NEXT:    ret i8 -16
 ;
   %x = or i8 %xx, -2
   %y = and i8 %yy, 4
@@ -90,11 +85,7 @@ define i8 @mul_high_bits_know2(i8 %xx, i8 %yy) {
 define i8 @mul_high_bits_know3(i8 %xx, i8 %yy) {
 ; CHECK-LABEL: define i8 @mul_high_bits_know3(
 ; CHECK-SAME: i8 [[XX:%.*]], i8 [[YY:%.*]]) {
-; CHECK-NEXT:    [[X:%.*]] = or i8 [[XX]], 124
-; CHECK-NEXT:    [[Y:%.*]] = or i8 [[YY]], 126
-; CHECK-NEXT:    [[MUL:%.*]] = mul i8 [[X]], [[Y]]
-; CHECK-NEXT:    [[R:%.*]] = and i8 [[MUL]], 112
-; CHECK-NEXT:    ret i8 [[R]]
+; CHECK-NEXT:    ret i8 0
 ;
   %x = or i8 %xx, -4
   %y = or i8 %yy, -2