[APInt] Rounding right-shifts

[lebedev.ri]

Summary

There's already rounding division, which is used in ConstantRange to
implement e.g. makeExactMulNUWRegion()/makeExactMulNUWRegion()
but there are no versions for right-shifts.
I'd like to try to extend ConstantRange::makeGuaranteedNoWrapRegion()
to deal with Instruction::Shl so i believe i need rounding right shifts.

The test coverage is confusing.
The existing RoundingSDiv() didn't actually pass - it never runs.
I've fixed that, and adjusted it to pass - i think we should be checking signed predicates?

Then i've added APIntOps::RoundingLShr(), APIntOps::RoundingAShr(),
standalone test coverage for them, and a cross-test to verify that
they produce identical output as compared to their div friends.
They do.

My understanding:

• LShr is "dividing" by positive 2^shamt, so it rounds down. Since the operation is unsigned it also rounds towards zero.
  • (33 l>> 1) << 1 == 32
  • (-31 l>> 1) << 1 == -32
    
    Which means that if we want to round up, if there's remainder then we just need to always increment the quotent.
  • ((33 l>> 1) + 1) << 1 == 34
  • ((-31 l>> 1) + 1) << 1 == -30
• AShr is "dividing" by positive 2^shamt, so it rounds down: (identical to LShr)
  • (33 a>> 1) << 1 == 32
  • (-31 a>> 1) << 1 == -32
    
    Again, if we want to round up, if there's remainder then we just need to always increment the quotent. (identical to LShr)
  • ((33 a>> 1) + 1) << 1 == 34
  • ((-31 a>> 1) + 1) << 1 == -30
    
    However, if we want to round towards zero, the behavior is different.
  • Since it's default behavior is to round down, then for positive numbers rounds towards zero is also the default:
    • (33 a>> 1) << 1 == 32
  • But for negatives by default we'd go away from zero:
    • (-31 l>> 1) << 1 == -32
      
      while the expected answer is -30; which is the answer for rounding up:
    • ((-31 a>> 1) + 1) << 1 == -30
      
      So if the value is negative then if quotent is non-zero - increment it, else do nothing.

[lebedev.ri]

llvm/unittests/ADT/APIntTest.cpp

2528–2541

@sanjoy @timshen FIXME: that's the intended check, right?

[lebedev.ri]

Upload correct patch - i've meant to disable the RoundingSDiv test - it doesn't pass after making it actually run.

[lebedev.ri]

Ok, last update - detect existence of quotent by counting trailing zeros, improve testss.

[nikic]

I'd like to try to extend ConstantRange::makeGuaranteedNoWrapRegion()
to deal with Instruction::Shl so i believe i need rounding right shifts.

I don't think rounding shifts are strictly necessary for this purpose, the correct behavior should fall out when applying the normal lshr/ashr operations to -1 / signed_min / signed_max.

[lebedev.ri]

In D68672#1707312, @nikic wrote:

I'd like to try to extend ConstantRange::makeGuaranteedNoWrapRegion()
to deal with Instruction::Shl so i believe i need rounding right shifts.

I don't think rounding shifts are strictly necessary for this purpose,

the correct behavior should fall out when applying the normal lshr/ashr operations to -1 / signed_min / signed_max.

I'm not sure i can parse that, fall out?

[lebedev.ri]

Bump

[nikic]

In D68672#1707404, @lebedev.ri wrote:

I'm not sure i can parse that, fall out?

Sorry, what I meant is that the normal lshr/ashr will behave as desired for the purposes of determining a shl GNWR range. Something along the lines of...

case Instruction::Shl: {
  APInt UMax = Other.getUnsignedMax();
  unsigned MaxShift = Unsigned ? BitWidth - 1 : BitWidth - 2;
  if (UMax.uge(MaxShift)) // (An intersect would be more precise)
    return ConstantRange(APInt::getNullValue(BitWidth));
  if (Unsigned)
    return getNonEmpty(APInt::getNullValue(BitWidth),
                       APInt::getMaxValue(BitWidth).lshr(UMax) + 1);
  return getNonEmpty(APInt::getSignedMinValue(BitWidth).ashr(UMax)
                     APInt::getSignedMaxValue(BitWidth).ashr(UMax) + 1);
}

We can still add these methods, they just don't seem necessary for the use-case you mentioned, unless I'm missing something.

[lebedev.ri]

In D68672#1714990, @nikic wrote:

In D68672#1707404, @lebedev.ri wrote:

I'm not sure i can parse that, fall out?

Sorry, what I meant is that the normal lshr/ashr will behave as desired for the purposes of determining a shl GNWR range. Something along the lines of...

case Instruction::Shl: {
  APInt UMax = Other.getUnsignedMax();
  unsigned MaxShift = Unsigned ? BitWidth - 1 : BitWidth - 2;
  if (UMax.uge(MaxShift)) // (An intersect would be more precise)
    return ConstantRange(APInt::getNullValue(BitWidth));
  if (Unsigned)
    return getNonEmpty(APInt::getNullValue(BitWidth),
                       APInt::getMaxValue(BitWidth).lshr(UMax) + 1);
  return getNonEmpty(APInt::getSignedMinValue(BitWidth).ashr(UMax)
                     APInt::getSignedMaxValue(BitWidth).ashr(UMax) + 1);
}

We can still add these methods, they just don't seem necessary for the use-case you mentioned, unless I'm missing something.

As i have discovered since then, it may make sense to first handle mul in CVP before looking into shifts.
So in this case it would be mainly good to have these for consistency.

[lebedev.ri]

Indeed, these turned out not to be needed for ConstanRange.
Do we want to have them for consistency regardless?

llvm/include/llvm/ADT/APInt.h

 }
 
 /// Return A unsign-divided by B, rounded by the given rounding mode.
 APInt RoundingUDiv(const APInt &A, const APInt &B, APInt::Rounding RM);
 
 /// Return A sign-divided by B, rounded by the given rounding mode.
 APInt RoundingSDiv(const APInt &A, const APInt &B, APInt::Rounding RM);
 
+/// Return A logically shifted right by B, rounded by the given rounding mode.
+APInt RoundingLShr(const APInt &A, const APInt &B, APInt::Rounding RM);
+
+/// Return A arithmetically shifted right by B, rounded by the given rounding
+/// mode.
+APInt RoundingAShr(const APInt &A, const APInt &B, APInt::Rounding RM);
+
 /// Let q(n) = An^2 + Bn + C, and BW = bit width of the value range
 /// (e.g. 32 for i32).
 /// This function finds the smallest number n, such that
 /// (a) n >= 0 and q(n) = 0, or
 /// (b) n >= 1 and q(n-1) and q(n), when evaluated in the set of all
 ///     integers, belong to two different intervals [Rk, Rk+R),
 ///     where R = 2^BW, and k is an integer.
 /// The idea here is to find when q(n) "overflows" 2^BW, while at the

llvm/lib/Support/APInt.cpp

APInt llvm::APIntOps::RoundingSDiv(const APInt &A, const APInt &B,
   }
   // Currently sdiv rounds twards zero.
   case APInt::Rounding::TOWARD_ZERO:
     return A.sdiv(B);
   }
   llvm_unreachable("Unknown APInt::Rounding enum");
 }
 
+APInt llvm::APIntOps::RoundingLShr(const APInt &A, const APInt &B,
+                                   APInt::Rounding RM) {
+  // Currently udivrem always rounds down.
+  switch (RM) {
+  case APInt::Rounding::DOWN:
+  case APInt::Rounding::TOWARD_ZERO:
+    return A.lshr(B);
+  case APInt::Rounding::UP: {
+    APInt Quo = A.lshr(B);
+    bool HasRem = A != Quo.shl(B);
+    if (!HasRem)
+      return Quo;
+    return Quo + 1;
+  }
+  }
+  llvm_unreachable("Unknown APInt::Rounding enum");
+}
+
+APInt llvm::APIntOps::RoundingAShr(const APInt &A, const APInt &B,
+                                   APInt::Rounding RM) {
+  switch (RM) {
+  case APInt::Rounding::TOWARD_ZERO:
+  case APInt::Rounding::UP: {
+    APInt Quo = A.ashr(B);
+    APInt RoundedA = Quo.shl(B);
+    bool HasRem = A != RoundedA;
+    if (!HasRem)
+      return Quo;
+    if (RM == APInt::Rounding::UP)
+      return Quo + 1;
+    if (A.isNegative())
+      return Quo + 1;
+    return Quo;
+  }
+  // Currently ashr rounds down.
+  case APInt::Rounding::DOWN:
+    return A.ashr(B);
+  }
+  llvm_unreachable("Unknown APInt::Rounding enum");
+}
+
 Optional<APInt>
 llvm::APIntOps::SolveQuadraticEquationWrap(APInt A, APInt B, APInt C,
                                            unsigned RangeWidth) {
   unsigned CoeffWidth = A.getBitWidth();
   assert(CoeffWidth == B.getBitWidth() && CoeffWidth == C.getBitWidth());
   assert(RangeWidth <= CoeffWidth &&
          "Value range width should be less than coefficient width");
   assert(RangeWidth > 1 && "Value range bit width should be > 1");

llvm/unittests/ADT/APIntTest.cpp

for (uint64_t Bi = 1; Bi <= 255; Bi++) {
         APInt Quo = A.udiv(B);
         EXPECT_EQ(Quo, APIntOps::RoundingUDiv(A, B, APInt::Rounding::TOWARD_ZERO));
         EXPECT_EQ(Quo, APIntOps::RoundingUDiv(A, B, APInt::Rounding::DOWN));
       }
     }
   }
 }
 
-TEST(APIntTest, RoundingSDiv) {
+TEST(APIntTest, DISABLED_RoundingSDiv) {
   for (int64_t Ai = -128; Ai <= 127; Ai++) {
     APInt A(8, Ai);
 
     if (Ai != 0) {
       APInt Zero(8, 0);
       EXPECT_EQ(0, APIntOps::RoundingSDiv(Zero, A, APInt::Rounding::UP));
       EXPECT_EQ(0, APIntOps::RoundingSDiv(Zero, A, APInt::Rounding::DOWN));
       EXPECT_EQ(0, APIntOps::RoundingSDiv(Zero, A, APInt::Rounding::TOWARD_ZERO));
     }
 
-    for (uint64_t Bi = -128; Bi <= 127; Bi++) {
+    for (int64_t Bi = -128; Bi <= 127; Bi++) {
       if (Bi == 0)
         continue;
 
       APInt B(8, Bi);
       {
         APInt Quo = APIntOps::RoundingSDiv(A, B, APInt::Rounding::UP);
         auto Prod = Quo.sext(16) * B.sext(16);
-        EXPECT_TRUE(Prod.uge(A));
-        if (Prod.ugt(A)) {
-          EXPECT_TRUE(((Quo - 1).sext(16) * B.sext(16)).ult(A));
+        // FIXME: these checks are incorrect.
+        EXPECT_TRUE(Prod.uge(Ai));
+        if (Prod.ugt(Ai)) {
+          EXPECT_TRUE(((Quo - 1).sext(16) * B.sext(16)).ult(Ai));
         }
       }
       {
         APInt Quo = APIntOps::RoundingSDiv(A, B, APInt::Rounding::DOWN);
         auto Prod = Quo.sext(16) * B.sext(16);
-        EXPECT_TRUE(Prod.ule(A));
-        if (Prod.ult(A)) {
-          EXPECT_TRUE(((Quo + 1).sext(16) * B.sext(16)).ugt(A));
+        // FIXME: these checks are incorrect.
+        EXPECT_TRUE(Prod.ule(Ai));
+        if (Prod.ult(Ai)) {
+          EXPECT_TRUE(((Quo + 1).sext(16) * B.sext(16)).ugt(Ai));
         }
       }
       {
         APInt Quo = A.sdiv(B);
         EXPECT_EQ(Quo, APIntOps::RoundingSDiv(A, B, APInt::Rounding::TOWARD_ZERO));
       }
     }
   }
 }
 
+TEST(APIntTest, RoundingLShr) {
+  for (uint64_t Ai = 0; Ai <= 255; Ai++) {
+    APInt A(8, Ai);
+
+    for (uint64_t Bi = 0; Bi <= 7; Bi++) {
+      APInt B(8, Bi);
+
+      if (Bi == 0) {
+        EXPECT_EQ(A, APIntOps::RoundingLShr(A, B, APInt::Rounding::UP));
+        EXPECT_EQ(A, APIntOps::RoundingLShr(A, B, APInt::Rounding::DOWN));
+        EXPECT_EQ(A,
+                  APIntOps::RoundingLShr(A, B, APInt::Rounding::TOWARD_ZERO));
+      }
+
+      {
+        APInt Quo = APIntOps::RoundingLShr(A, B, APInt::Rounding::UP);
+        auto Prod = Quo.zext(16) << B.zext(16);
+        EXPECT_TRUE(Prod.uge(Ai));
+        if (Prod.ugt(Ai)) {
+          EXPECT_TRUE(((Quo - 1).zext(16) << B.zext(16)).ult(Ai));
+        }
+      }
+      {
+        APInt Quo = A.lshr(B);
+        EXPECT_EQ(Quo,
+                  APIntOps::RoundingLShr(A, B, APInt::Rounding::TOWARD_ZERO));
+        EXPECT_EQ(Quo, APIntOps::RoundingLShr(A, B, APInt::Rounding::DOWN));
+      }
+    }
+  }
+}
+
+TEST(APIntTest, RoundingAShr) {
+  for (int64_t Ai = -128; Ai <= 127; Ai++) {
+    APInt A(8, Ai);
+
+    for (uint64_t Bi = 0; Bi <= 7; Bi++) {
+      APInt B(8, Bi);
+
+      if (Bi == 0) {
+        EXPECT_EQ(A, APIntOps::RoundingAShr(A, B, APInt::Rounding::UP));
+        EXPECT_EQ(A, APIntOps::RoundingAShr(A, B, APInt::Rounding::DOWN));
+        EXPECT_EQ(A,
+                  APIntOps::RoundingAShr(A, B, APInt::Rounding::TOWARD_ZERO));
+      }
+      {
+        APInt Quo = APIntOps::RoundingAShr(A, B, APInt::Rounding::UP);
+        auto Prod = Quo.sext(16) << B.sext(16);
+        EXPECT_TRUE(Prod.sge(Ai));
+        if (Prod.sgt(Ai)) {
+          EXPECT_TRUE(((Quo - 1).sext(16) << B.sext(16)).slt(Ai));
+        }
+      }
+      {
+        APInt Quo = APIntOps::RoundingAShr(A, B, APInt::Rounding::TOWARD_ZERO);
+        auto Prod = Quo.sext(16) << B.sext(16);
+        EXPECT_TRUE(Prod.ule(Ai));
+        if (Prod.ult(Ai)) {
+          EXPECT_TRUE(((Quo + 1).sext(16) << B.sext(16)).sgt(Ai));
+        }
+      }
+      {
+        APInt Quo = A.ashr(B);
+        EXPECT_EQ(Quo, APIntOps::RoundingAShr(A, B, APInt::Rounding::DOWN));
+      }
+    }
+  }
+}
+
+TEST(APIntTest, RoundingLShrVsUDiv) {
+  for (uint64_t Ai = 0; Ai <= 255; Ai++) {
+    APInt A(8, Ai);
+
+    for (uint64_t Bi = 0; Bi <= 7; Bi++) {
+      APInt B(8, Bi);
+
+      APInt Divisor(8, 1 << Bi);
+
+      EXPECT_EQ(APIntOps::RoundingUDiv(A, Divisor, APInt::Rounding::UP),
+                APIntOps::RoundingLShr(A, B, APInt::Rounding::UP));
+      EXPECT_EQ(APIntOps::RoundingUDiv(A, Divisor, APInt::Rounding::DOWN),
+                APIntOps::RoundingLShr(A, B, APInt::Rounding::DOWN));
+      EXPECT_EQ(
+          APIntOps::RoundingUDiv(A, Divisor, APInt::Rounding::TOWARD_ZERO),
+          APIntOps::RoundingLShr(A, B, APInt::Rounding::TOWARD_ZERO));
+    }
+  }
+}
+
+TEST(APIntTest, RoundingAShrVsSDiv) {
+  for (uint64_t Ai = 0; Ai <= 255; Ai++) {
+    APInt A(8, Ai);
+
+    for (uint64_t Bi = 0; Bi <= 6; Bi++) {
+      APInt B(8, Bi);
+
+      APInt Divisor(8, 1 << Bi);
+
+      EXPECT_EQ(APIntOps::RoundingSDiv(A, Divisor, APInt::Rounding::UP),
+                APIntOps::RoundingAShr(A, B, APInt::Rounding::UP));
+      EXPECT_EQ(APIntOps::RoundingSDiv(A, Divisor, APInt::Rounding::DOWN),
+                APIntOps::RoundingAShr(A, B, APInt::Rounding::DOWN));
+      EXPECT_EQ(
+          APIntOps::RoundingSDiv(A, Divisor, APInt::Rounding::TOWARD_ZERO),
+          APIntOps::RoundingAShr(A, B, APInt::Rounding::TOWARD_ZERO));
+    }
+  }
+}
+
 TEST(APIntTest, umul_ov) {
   const std::pair<uint64_t, uint64_t> Overflows[] = {
       {0x8000000000000000, 2},
       {0x5555555555555556, 3},
       {4294967296, 4294967296},
       {4294967295, 4294967298},
   };
   const std::pair<uint64_t, uint64_t> NonOverflows[] = {