[APInt] Add ashrInPlace method and rewrite in ashr using it.

[craig.topper]

Summary

This patch adds an in place version of ashr to match lshr and shl which were recently added.

I've tried to make this similar to the lshr code with additions to handle the sign extension. I've also tried to do this with less if checks than the current ashr code by sign extending the original result to a word boundary before doing any of the shifting. This removes a lot of the complexity of determining where to fill in sign bits after the shifting.

[hans]

lgtm with some comments

include/llvm/ADT/APInt.h

911 ↗

(On Diff #96345)

What does the "// undefined" comment mean? That we don't really define what happens when ShiftAmt == BitWidth but we try to do something reasonable? Maybe spell that out, I hadn't seen that before.

lib/Support/APInt.cpp

1029 ↗

(On Diff #96345)

Maybe this should be inline since it's so simple?

1039 ↗

(On Diff #96345)

Did you lose the "Invalid shift amount" assert?

1043 ↗

(On Diff #96345)

nit: "is is intra-part shift"

1051 ↗

(On Diff #96345)

Maybe run clang-format on this line? There should usually be spaces around the '-' operator for example.

1060 ↗

(On Diff #96345)

nit: spaces around the - here too

[craig.topper]

lib/Support/APInt.cpp

1029 ↗

(On Diff #96345)

Possibly. Though both ashrInPlace and getLimitedValue will expand to include slow case calls.

All 3 shifts types have this out of line so I'll look at them together.

1039 ↗

(On Diff #96345)

It was moved to the inline single word case.

1043 ↗

(On Diff #96345)

Seems that tcShiftRight and tcShiftLeft have the same mistake. Thanks!

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

private:
   void initSlowCase(const APInt &that);
 
   /// out-of-line slow case for shl
   void shlSlowCase(unsigned ShiftAmt);
 
   /// out-of-line slow case for lshr.
   void lshrSlowCase(unsigned ShiftAmt);
 
+  /// out-of-line slow case for ashr.
+  void ashrSlowCase(unsigned ShiftAmt);
+
   /// out-of-line slow case for operator=
   void AssignSlowCase(const APInt &RHS);
 
   /// out-of-line slow case for operator==
   bool EqualSlowCase(const APInt &RHS) const LLVM_READONLY;
 
   /// out-of-line slow case for countLeadingZeros
   unsigned countLeadingZerosSlowCase() const LLVM_READONLY;
public:
   /// \brief Left logical shift operator.
   ///
   /// Shifts this APInt left by \p Bits and returns the result.
   APInt operator<<(const APInt &Bits) const { return shl(Bits); }
 
   /// \brief Arithmetic right-shift function.
   ///
   /// Arithmetic right-shift this APInt by shiftAmt.
-  APInt ashr(unsigned shiftAmt) const;
+  APInt ashr(unsigned ShiftAmt) const {
+    APInt R(*this);
+    R.ashrInPlace(ShiftAmt);
+    return R;
+  }
+
+  /// Arithmetic right-shift this APInt by ShiftAmt in place.
+  void ashrInPlace(unsigned ShiftAmt) {
+    assert(ShiftAmt <= BitWidth && "Invalid shift amount");
+    if (isSingleWord()) {
+      int64_t SExtVAL = SignExtend64(VAL, BitWidth);
+      if (ShiftAmt == BitWidth)
+        VAL = SExtVAL >> (APINT_BITS_PER_WORD - 1); // Fill with sign bit.
+      else
+        VAL = SExtVAL >> ShiftAmt;
+      clearUnusedBits();
+      return;
+    }
+    ashrSlowCase(ShiftAmt);
+  }
 
   /// \brief Logical right-shift function.
   ///
   /// Logical right-shift this APInt by shiftAmt.
   APInt lshr(unsigned shiftAmt) const {
     APInt R(*this);
     R.lshrInPlace(shiftAmt);
     return R;
public:
   APInt rotl(unsigned rotateAmt) const;
 
   /// \brief Rotate right by rotateAmt.
   APInt rotr(unsigned rotateAmt) const;
 
   /// \brief Arithmetic right-shift function.
   ///
   /// Arithmetic right-shift this APInt by shiftAmt.
-  APInt ashr(const APInt &shiftAmt) const;
+  APInt ashr(const APInt &ShiftAmt) const {
+    APInt R(*this);
+    R.ashrInPlace(ShiftAmt);
+    return R;
+  }
+
+  /// Arithmetic right-shift this APInt by shiftAmt in place.
+  void ashrInPlace(const APInt &shiftAmt);
 
   /// \brief Logical right-shift function.
   ///
   /// Logical right-shift this APInt by shiftAmt.
   APInt lshr(const APInt &ShiftAmt) const {
     APInt R(*this);
     R.lshrInPlace(ShiftAmt);
     return R;

llvm/trunk/include/llvm/ADT/APSInt.h

public:
 
   APSInt operator>>(unsigned Amt) const {
     return IsUnsigned ? APSInt(lshr(Amt), true) : APSInt(ashr(Amt), false);
   }
   APSInt& operator>>=(unsigned Amt) {
     if (IsUnsigned)
       lshrInPlace(Amt);
     else
-      *this = *this >> Amt;
+      ashrInPlace(Amt);
     return *this;
   }
 
   inline bool operator<(const APSInt& RHS) const {
     assert(IsUnsigned == RHS.IsUnsigned && "Signedness mismatch!");
     return IsUnsigned ? ult(RHS) : slt(RHS);
   }
   inline bool operator>(const APSInt& RHS) const {

llvm/trunk/lib/Support/APInt.cpp

 APInt APInt::sextOrSelf(unsigned width) const {
   if (BitWidth < width)
     return sext(width);
   return *this;
 }
 
 /// Arithmetic right-shift this APInt by shiftAmt.
 /// @brief Arithmetic right-shift function.
-APInt APInt::ashr(const APInt &shiftAmt) const {
-  return ashr((unsigned)shiftAmt.getLimitedValue(BitWidth));
+void APInt::ashrInPlace(const APInt &shiftAmt) {
+  ashrInPlace((unsigned)shiftAmt.getLimitedValue(BitWidth));
 }
 
 /// Arithmetic right-shift this APInt by shiftAmt.
 /// @brief Arithmetic right-shift function.
-APInt APInt::ashr(unsigned shiftAmt) const {
-  assert(shiftAmt <= BitWidth && "Invalid shift amount");
-  // Handle a degenerate case
-  if (shiftAmt == 0)
-    return *this;
+void APInt::ashrSlowCase(unsigned ShiftAmt) {
+  // Don't bother performing a no-op shift.
+  if (!ShiftAmt)
+    return;
 
-  // Handle single word shifts with built-in ashr
-  if (isSingleWord()) {
-    if (shiftAmt == BitWidth)
-      // Undefined
-      return APInt(BitWidth,
-                   SignExtend64(VAL, BitWidth) >> (APINT_BITS_PER_WORD - 1));
-    return APInt(BitWidth, SignExtend64(VAL, BitWidth) >> shiftAmt);
-  }
+  // Save the original sign bit for later.
+  bool Negative = isNegative();
 
-  // If all the bits were shifted out, the result is, technically, undefined.
-  // We return -1 if it was negative, 0 otherwise. We check this early to avoid
-  // issues in the algorithm below.
-  if (shiftAmt == BitWidth) {
-    if (isNegative())
-      return APInt(BitWidth, WORD_MAX, true);
-    else
-      return APInt(BitWidth, 0);
-  }
+  // WordShift is the inter-part shift; BitShift is is intra-part shift.
+  unsigned WordShift = ShiftAmt / APINT_BITS_PER_WORD;
+  unsigned BitShift = ShiftAmt % APINT_BITS_PER_WORD;
+
+  unsigned WordsToMove = getNumWords() - WordShift;
+  if (WordsToMove != 0) {
+    // Sign extend the last word to fill in the unused bits.
+    pVal[getNumWords() - 1] = SignExtend64(
+        pVal[getNumWords() - 1], ((BitWidth - 1) % APINT_BITS_PER_WORD) + 1);
 
-  // Create some space for the result.
-  uint64_t * val = new uint64_t[getNumWords()];
-
-  // Compute some values needed by the following shift algorithms
-  unsigned wordShift = shiftAmt % APINT_BITS_PER_WORD; // bits to shift per word
-  unsigned offset = shiftAmt / APINT_BITS_PER_WORD; // word offset for shift
-  unsigned breakWord = getNumWords() - 1 - offset; // last word affected
-  unsigned bitsInWord = whichBit(BitWidth); // how many bits in last word?
-  if (bitsInWord == 0)
-    bitsInWord = APINT_BITS_PER_WORD;
-
-  // If we are shifting whole words, just move whole words
-  if (wordShift == 0) {
-    // Move the words containing significant bits
-    for (unsigned i = 0; i <= breakWord; ++i)
-      val[i] = pVal[i+offset]; // move whole word
-
-    // Adjust the top significant word for sign bit fill, if negative
-    if (isNegative())
-      if (bitsInWord < APINT_BITS_PER_WORD)
-        val[breakWord] |= WORD_MAX << bitsInWord; // set high bits
-  } else {
-    // Shift the low order words
-    for (unsigned i = 0; i < breakWord; ++i) {
-      // This combines the shifted corresponding word with the low bits from
-      // the next word (shifted into this word's high bits).
-      val[i] = (pVal[i+offset] >> wordShift) |
-               (pVal[i+offset+1] << (APINT_BITS_PER_WORD - wordShift));
-    }
-
-    // Shift the break word. In this case there are no bits from the next word
-    // to include in this word.
-    val[breakWord] = pVal[breakWord+offset] >> wordShift;
+    // Fastpath for moving by whole words.
+    if (BitShift == 0) {
+      std::memmove(pVal, pVal + WordShift, WordsToMove * APINT_WORD_SIZE);
+    } else {
+      // Move the words containing significant bits.
+      for (unsigned i = 0; i != WordsToMove - 1; ++i)
+        pVal[i] = (pVal[i + WordShift] >> BitShift) |
+                  (pVal[i + WordShift + 1] << (APINT_BITS_PER_WORD - BitShift));
 
-    // Deal with sign extension in the break word, and possibly the word before
-    // it.
-    if (isNegative()) {
-      if (wordShift > bitsInWord) {
-        if (breakWord > 0)
-          val[breakWord-1] |=
-            WORD_MAX << (APINT_BITS_PER_WORD - (wordShift - bitsInWord));
-        val[breakWord] |= WORD_MAX;
-      } else
-        val[breakWord] |= WORD_MAX << (bitsInWord - wordShift);
+      // Handle the last word which has no high bits to copy.
+      pVal[WordsToMove - 1] = pVal[WordShift + WordsToMove - 1] >> BitShift;
+      // Sign extend one more time.
+      pVal[WordsToMove - 1] =
+          SignExtend64(pVal[WordsToMove - 1], APINT_BITS_PER_WORD - BitShift);
     }
   }
 
-  // Remaining words are 0 or -1, just assign them.
-  uint64_t fillValue = (isNegative() ? WORD_MAX : 0);
-  for (unsigned i = breakWord+1; i < getNumWords(); ++i)
-    val[i] = fillValue;
-  APInt Result(val, BitWidth);
-  Result.clearUnusedBits();
-  return Result;
+  // Fill in the remainder based on the original sign.
+  std::memset(pVal + WordsToMove, Negative ? -1 : 0,
+              WordShift * APINT_WORD_SIZE);
+  clearUnusedBits();
 }
 
 /// Logical right-shift this APInt by shiftAmt.
 /// @brief Logical right-shift function.
 void APInt::lshrInPlace(const APInt &shiftAmt) {
   lshrInPlace((unsigned)shiftAmt.getLimitedValue(BitWidth));
 }
 

llvm/trunk/unittests/ADT/APIntTest.cpp

TEST(APIntTest, LogicalRightShift) {
   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(128));
 }
 
+TEST(APIntTest, ArithmeticRightShift) {
+  APInt i72(APInt::getHighBitsSet(72, 1));
+  i72.ashrInPlace(46);
+  EXPECT_EQ(47U, i72.countLeadingOnes());
+  EXPECT_EQ(25U, i72.countTrailingZeros());
+  EXPECT_EQ(47U, i72.countPopulation());
+
+  i72 = APInt::getHighBitsSet(72, 1);
+  i72.ashrInPlace(64);
+  EXPECT_EQ(65U, i72.countLeadingOnes());
+  EXPECT_EQ(7U, i72.countTrailingZeros());
+  EXPECT_EQ(65U, i72.countPopulation());
+
+  APInt i128(APInt::getHighBitsSet(128, 1));
+  i128.ashrInPlace(64);
+  EXPECT_EQ(65U, i128.countLeadingOnes());
+  EXPECT_EQ(63U, i128.countTrailingZeros());
+  EXPECT_EQ(65U, i128.countPopulation());
+
+  // Ensure we handle large shifts of multi-word.
+  const APInt signmin32(APInt::getSignedMinValue(32));
+  EXPECT_TRUE(signmin32.ashr(32).isAllOnesValue());
+
+  // Ensure we handle large shifts of multi-word.
+  const APInt umax32(APInt::getSignedMaxValue(32));
+  EXPECT_EQ(0, umax32.ashr(32));
+
+  // Ensure we handle large shifts of multi-word.
+  const APInt signmin128(APInt::getSignedMinValue(128));
+  EXPECT_TRUE(signmin128.ashr(128).isAllOnesValue());
+
+  // Ensure we handle large shifts of multi-word.
+  const APInt umax128(APInt::getSignedMaxValue(128));
+  EXPECT_EQ(0, umax128.ashr(128));
+}
+
 TEST(APIntTest, LeftShift) {
   APInt i256(APInt::getLowBitsSet(256, 2));
 
   i256 <<= 1;
   EXPECT_EQ(253U, i256.countLeadingZeros());
   EXPECT_EQ(1U, i256.countTrailingZeros());
   EXPECT_EQ(2U, i256.countPopulation());