[APInt] Merge the multiword code from lshrInPlace and tcShiftRight into a single implementation

[craig.topper]

Summary

This merges the two different multiword shift right implementations into a single version located in tcShiftRight. lshrInPlace now calls tcShiftRight for the multiword case.

I retained the memmove fast path from lshrInPlace and used a memset for the zeroing. The for loop is basically tcShiftRight's implementation with the zeroing and the intra-shift of 0 removed.

[craig.topper]

Fix diff context

[hans]

lgtm with a naming suggestion

Are there tests covering this? I couldn't find any direct tests, but are there at least indirect ones that would break if for example you change the function to not shift at all?

Do we have the same situation with tcShiftLeft and shlSlowCase? I suppose it's not quite the same since shlSlowCase isn't in-place, but maybe there's some code to be shared anyway?

lib/Support/APInt.cpp

2690 ↗

(On Diff #95399)

Hmm, seems it's spelled COUNT like this in many other parts of the file, but I don't think we do that generally in LLVM; I'd prefer Count. Maybe not worth changing now :-/

2697 ↗

(On Diff #95399)

Use // instead of /*

2698 ↗

(On Diff #95399)

Actually, Jump doesn't seem like a great name. Maybe we could have WordShift and BitShift?

[craig.topper]

I want to look at left shift at some point. Also making an in place version of ashr as well. lshrInPlace was just added last week which is why I focused on it.

[craig.topper]

Update comments and variable names. Add some basic tests.

[hans]

lgtm2

llvm/trunk/include/llvm/ADT/APInt.h

public:
   ///
   /// Logical right-shift this APInt by shiftAmt.
   APInt lshr(unsigned shiftAmt) const {
     APInt R(*this);
     R.lshrInPlace(shiftAmt);
     return R;
   }
 
-  /// Logical right-shift this APInt by shiftAmt in place.
-  void lshrInPlace(unsigned shiftAmt);
+  /// Logical right-shift this APInt by ShiftAmt in place.
+  void lshrInPlace(unsigned ShiftAmt);
 
   /// \brief Left-shift function.
   ///
   /// Left-shift this APInt by shiftAmt.
   APInt shl(unsigned shiftAmt) const {
     assert(shiftAmt <= BitWidth && "Invalid shift amount");
     if (isSingleWord()) {
       if (shiftAmt >= BitWidth)

   static int tcDivide(WordType *lhs, const WordType *rhs,
                       WordType *remainder, WordType *scratch,
                       unsigned parts);
 
   /// Shift a bignum left COUNT bits.  Shifted in bits are zero.  There are no
   /// restrictions on COUNT.
   static void tcShiftLeft(WordType *, unsigned parts, unsigned count);
 
-  /// Shift a bignum right COUNT bits.  Shifted in bits are zero.  There are no
-  /// restrictions on COUNT.
-  static void tcShiftRight(WordType *, unsigned parts, unsigned count);
+  /// Shift a bignum right Count bits.  Shifted in bits are zero.  There are no
+  /// restrictions on Count.
+  static void tcShiftRight(WordType *, unsigned Words, unsigned Count);
 
   /// The obvious AND, OR and XOR and complement operations.
   static void tcAnd(WordType *, const WordType *, unsigned);
   static void tcOr(WordType *, const WordType *, unsigned);
   static void tcXor(WordType *, const WordType *, unsigned);
   static void tcComplement(WordType *, unsigned);
 
   /// Comparison (unsigned) of two bignums.

llvm/trunk/lib/Support/APInt.cpp

 }
 
 /// Logical right-shift this APInt by shiftAmt.
 /// @brief Logical right-shift function.
 APInt APInt::lshr(const APInt &shiftAmt) const {
   return lshr((unsigned)shiftAmt.getLimitedValue(BitWidth));
 }
 
-/// Perform a logical right-shift from Src to Dst of Words words, by Shift,
-/// which must be less than 64. If the source and destination ranges overlap,
-/// we require that Src >= Dst (put another way, we require that the overall
-/// operation is a right shift on the combined range).
-static void lshrWords(APInt::WordType *Dst, APInt::WordType *Src,
-                      unsigned Words, unsigned Shift) {
-  assert(Shift < APInt::APINT_BITS_PER_WORD);
-
-  if (!Words)
-    return;
-
-  if (Shift == 0) {
-    std::memmove(Dst, Src, Words * APInt::APINT_WORD_SIZE);
-    return;
-  }
-
-  uint64_t Low = Src[0];
-  for (unsigned I = 1; I != Words; ++I) {
-    uint64_t High = Src[I];
-    Dst[I - 1] =
-        (Low >> Shift) | (High << (APInt::APINT_BITS_PER_WORD - Shift));
-    Low = High;
-  }
-  Dst[Words - 1] = Low >> Shift;
-}
-
 /// Logical right-shift this APInt by shiftAmt.
 /// @brief Logical right-shift function.
-void APInt::lshrInPlace(unsigned shiftAmt) {
+void APInt::lshrInPlace(unsigned ShiftAmt) {
   if (isSingleWord()) {
-    if (shiftAmt >= BitWidth)
+    if (ShiftAmt >= BitWidth)
       VAL = 0;
     else
-      VAL >>= shiftAmt;
+      VAL >>= ShiftAmt;
     return;
   }
 
-  // Don't bother performing a no-op shift.
-  if (!shiftAmt)
-    return;
-
-  // Find number of complete words being shifted out and zeroed.
-  const unsigned Words = getNumWords();
-  const unsigned ShiftFullWords =
-      std::min(shiftAmt / APINT_BITS_PER_WORD, Words);
-
-  // Fill in first Words - ShiftFullWords by shifting.
-  lshrWords(pVal, pVal + ShiftFullWords, Words - ShiftFullWords,
-            shiftAmt % APINT_BITS_PER_WORD);
-
-  // The remaining high words are all zero.
-  for (unsigned I = Words - ShiftFullWords; I != Words; ++I)
-    pVal[I] = 0;
+  return tcShiftRight(pVal, getNumWords(), ShiftAmt);
 }
 
 /// Left-shift this APInt by shiftAmt.
 /// @brief Left-shift function.
 APInt APInt::shl(const APInt &shiftAmt) const {
   // It's undefined behavior in C to shift by BitWidth or greater.
   return shl((unsigned)shiftAmt.getLimitedValue(BitWidth));
 }
while (parts > jump) {
       dst[parts] = part;
     }
 
     while (parts > 0)
       dst[--parts] = 0;
   }
 }
 
-/* Shift a bignum right COUNT bits in-place.  Shifted in bits are
-   zero.  There are no restrictions on COUNT.  */
-void APInt::tcShiftRight(WordType *dst, unsigned parts, unsigned count) {
-  if (count) {
-    /* Jump is the inter-part jump; shift is is intra-part shift.  */
-    unsigned jump = count / APINT_BITS_PER_WORD;
-    unsigned shift = count % APINT_BITS_PER_WORD;
+/// Shift a bignum right Count bits in-place. Shifted in bits are zero. There
+/// are no restrictions on Count.
+void APInt::tcShiftRight(WordType *Dst, unsigned Words, unsigned Count) {
+  // Don't bother performing a no-op shift.
+  if (!Count)
+    return;
 
-    /* Perform the shift.  This leaves the most significant COUNT bits
-       of the result at zero.  */
-    for (unsigned i = 0; i < parts; i++) {
-      WordType part;
+  // WordShift is the inter-part shift; BitShift is is intra-part shift.
+  unsigned WordShift = std::min(Count / APINT_BITS_PER_WORD, Words);
+  unsigned BitShift = Count % APINT_BITS_PER_WORD;
 
-      if (i + jump >= parts) {
-        part = 0;
-      } else {
-        part = dst[i + jump];
-        if (shift) {
-          part >>= shift;
-          if (i + jump + 1 < parts)
-            part |= dst[i + jump + 1] << (APINT_BITS_PER_WORD - shift);
-        }
-      }
+  unsigned WordsToMove = Words - WordShift;
+  // Fastpath for moving by whole words.
+  if (BitShift == 0) {
+    std::memmove(Dst, Dst + WordShift, WordsToMove * APINT_WORD_SIZE);
+    for (unsigned i = 0; i != WordsToMove; ++i) {
+      Dst[i] = Dst[i + WordShift] >> BitShift;
+      if (i + 1 != WordsToMove)
+        Dst[i] |= Dst[i + WordShift + 1] << (APINT_BITS_PER_WORD - BitShift);
 
-      dst[i] = part;
-    }
-  }
+  // Fill in the remainder with 0s.
+  std::memset(Dst + WordsToMove, 0, WordShift * APINT_WORD_SIZE);
 }
 
 /* Bitwise and of two bignums.  */
 void APInt::tcAnd(WordType *dst, const WordType *rhs, unsigned parts) {
   for (unsigned i = 0; i < parts; i++)
     dst[i] &= rhs[i];
 }
 

llvm/trunk/unittests/ADT/APIntTest.cpp

TEST(APIntTest, ShiftLeftByZero) {
   EXPECT_FALSE(Shl[1]);
 }
 
 TEST(APIntTest, i64_ArithmeticRightShiftNegative) {
   const APInt neg_one(64, static_cast<uint64_t>(-1), true);
   EXPECT_EQ(neg_one, neg_one.ashr(7));
 }
 
-TEST(APIntTest, i64_LogicalRightShiftNegative) {
-  const APInt neg_one(128, static_cast<uint64_t>(-1), true);
-  EXPECT_EQ(0, neg_one.lshr(257));
-}
-
 TEST(APIntTest, i128_NegativeCount) {
   APInt Minus3(128, static_cast<uint64_t>(-3), true);
   EXPECT_EQ(126u, Minus3.countLeadingOnes());
   EXPECT_EQ(-3, Minus3.getSExtValue());
 
   APInt Minus1(128, static_cast<uint64_t>(-1), true);
   EXPECT_EQ(0u, Minus1.countLeadingZeros());
   EXPECT_EQ(128u, Minus1.countLeadingOnes());
TEST(APIntTest, GCD) {
 
   // 9931 and 123456 are coprime.
   APInt A = HugePrime * APInt(BitWidth, 9931);
   APInt B = HugePrime * APInt(BitWidth, 123456);
   APInt C = GreatestCommonDivisor(A, B);
   EXPECT_EQ(C, HugePrime);
 }
 
+TEST(APIntTest, LogicalRightShift) {
+  APInt i256(APInt::getHighBitsSet(256, 2));
+
+  i256.lshrInPlace(1);
+  EXPECT_EQ(1U, i256.countLeadingZeros());
+  EXPECT_EQ(253U, i256.countTrailingZeros());
+  EXPECT_EQ(2U, i256.countPopulation());
+
+  i256.lshrInPlace(62);
+  EXPECT_EQ(63U, i256.countLeadingZeros());
+  EXPECT_EQ(191U, i256.countTrailingZeros());
+  EXPECT_EQ(2U, i256.countPopulation());
+
+  i256.lshrInPlace(65);
+  EXPECT_EQ(128U, i256.countLeadingZeros());
+  EXPECT_EQ(126U, i256.countTrailingZeros());
+  EXPECT_EQ(2U, i256.countPopulation());
+
+  i256.lshrInPlace(64);
+  EXPECT_EQ(192U, i256.countLeadingZeros());
+  EXPECT_EQ(62U, i256.countTrailingZeros());
+  EXPECT_EQ(2U, i256.countPopulation());
+
+  i256.lshrInPlace(63);
+  EXPECT_EQ(255U, i256.countLeadingZeros());
+  EXPECT_EQ(0U, i256.countTrailingZeros());
+  EXPECT_EQ(1U, i256.countPopulation());
+
+  // Ensure we handle large shifts of multi-word.
+  const APInt neg_one(128, static_cast<uint64_t>(-1), true);
+  EXPECT_EQ(0, neg_one.lshr(257));
+}
+
 } // end anonymous namespace