[APInt] Add APInt::extractBits() method to extract APInt subrange

[RKSimon]

Summary

The current pattern for extract bits in range is typically:

Mask.lshr(BitOffset).trunc(SubSizeInBits);

Which can be particularly slow for large APInts (MaskSizeInBits > 64) as they require the allocation of memory for the temporary variable.

This is another of the compile time issues identified in PR32037 (see also D30265).

This patch adds the APInt::extractBits() helper method which avoids the temporary memory allocation.

[filcab]

LGTM

lib/Support/APInt.cpp

640

Looks deprecated

/// Equivalent to APInt(numBits, ArrayRef<uint64_t>(bigVal, numWords)), but
/// deprecated because this constructor is prone to ambiguity with the
/// APInt(unsigned, uint64_t, bool) constructor.
///
/// If this overload is ever deleted, care should be taken to prevent calls
/// from being incorrectly captured by the APInt(unsigned, uint64_t, bool)
/// constructor.
APInt(unsigned numBits, unsigned numWords, const uint64_t bigVal[]);

Maybe return APInt(numBits, makeArrayRef(pVal + loWord, 1 + hiWord - loWord));?

[davide]

LGTM

[hans]

lib/Support/APInt.cpp

624

This is very nit-picky, but I find numBits > 0 much more natural than 0 < numBits.
And I'd suggest splitting up the second assert into two, to make it clearer what went wrong if it fires.

628

For the single-word case we know bitPosition is in the first word, so there's no need for the remainder operation in whichBit. I'd suggest moving the if-statement before computing loBit, and just using bitPosition in the shift.

645

I'm not sure this loop is correct. It will never set Result.pVal[hiWord]?

The only test case I see that hits this loop is EXPECT_EQ(-8388481, i256.extractBits(128, 1).getSExtValue());

In that one, the result is computed from three words, but indeed fits into two words.

But what if it didn't -- what about i256.extractBits(129, 1)?

unittests/ADT/APIntTest.cpp

1441

Would it be possible to use hex constants here and in the expectations somehow? Especially for the last one, it's hard to see what the expected result should be..

[RKSimon]

Thanks for catching that Hans - I've reverted the commit and will recommit with your recommendations later on.

include/llvm/ADT/APInt.h

   }
 
   /// \brief Toggles a given bit to its opposite value.
   ///
   /// Toggle a given bit to its opposite value whose position is given
   /// as "bitPosition".
   void flipBit(unsigned bitPosition);
 
+  /// Return an APInt with the extracted bits [bitPosition,bitPosition+numBits).
+  APInt extractBits(unsigned numBits, unsigned bitPosition) const;
+
   /// @}
   /// \name Value Characterization Functions
   /// @{
 
   /// \brief Return the number of bits in the APInt.
   unsigned getBitWidth() const { return BitWidth; }
 
   /// \brief Get the number of words.

lib/Support/APInt.cpp

 /// as "bitPosition".
 /// @brief Toggles a given bit to its opposite value.
 void APInt::flipBit(unsigned bitPosition) {
   assert(bitPosition < BitWidth && "Out of the bit-width range!");
   if ((*this)[bitPosition]) clearBit(bitPosition);
   else setBit(bitPosition);
 }
 
+APInt APInt::extractBits(unsigned numBits, unsigned bitPosition) const {
+  assert(0 < numBits && "Can't extract zero bits");
+  assert(bitPosition < BitWidth && (numBits + bitPosition) <= BitWidth &&
+    "Illegal bit extraction");
+
+  unsigned loBit = whichBit(bitPosition);
+  if (isSingleWord())
+    return APInt(numBits, VAL >> loBit);
+
+  unsigned loWord = whichWord(bitPosition);
+  unsigned hiWord = whichWord(bitPosition + numBits - 1);
+
+  // Single word result extracting bits from a single word source.
+  if (loWord == hiWord)
+    return APInt(numBits, pVal[loWord] >> loBit);
+
+  // Extracting bits that start on a source word boundary can be done
+  // as a fast memory copy.
+  if (loBit == 0)
+    return APInt(numBits, 1 + hiWord - loWord, pVal + loWord);
+
+  // General case - shift + copy source words into place.
+  APInt Result(numBits, 0);
+  uint64_t *pDst = Result.pVal;
+  for (unsigned word = loWord; word < hiWord; ++word, ++pDst) {
+    uint64_t w0 = pVal[word + 0];
+    uint64_t w1 = pVal[word + 1];
+    *pDst = (w0 >> loBit) | (w1 << (APINT_BITS_PER_WORD - loBit));
+  }
+  return Result.clearUnusedBits();
+}
+
 unsigned APInt::getBitsNeeded(StringRef str, uint8_t radix) {
   assert(!str.empty() && "Invalid string length");
   assert((radix == 10 || radix == 8 || radix == 16 || radix == 2 ||
           radix == 36) &&
          "Radix should be 2, 8, 10, 16, or 36!");
 
   size_t slen = str.size();
 

lib/Target/X86/X86ISelLowering.cpp

auto SplitBitData = [&]() {
     bool AllowUndefs = AllowWholeUndefs || AllowPartialUndefs;
     if (UndefBits.getBoolValue() && !AllowUndefs)
       return false;
 
     UndefElts = SmallBitVector(NumElts, false);
     EltBits.resize(NumElts, APInt(EltSizeInBits, 0));
 
     for (unsigned i = 0; i != NumElts; ++i) {
-      APInt UndefEltBits = UndefBits.lshr(i * EltSizeInBits);
-      UndefEltBits = UndefEltBits.zextOrTrunc(EltSizeInBits);
+      unsigned BitOffset = i * EltSizeInBits;
+      APInt UndefEltBits = UndefBits.extractBits(EltSizeInBits, BitOffset);
 
       // Only treat an element as UNDEF if all bits are UNDEF.
       if (UndefEltBits.isAllOnesValue()) {
         if (!AllowWholeUndefs)
           return false;
         UndefElts[i] = true;
         continue;
       }
 
       // If only some bits are UNDEF then treat them as zero (or bail if not
       // supported).
       if (UndefEltBits.getBoolValue() && !AllowPartialUndefs)
         return false;
 
-      APInt Bits = MaskBits.lshr(i * EltSizeInBits).zextOrTrunc(EltSizeInBits);
+      APInt Bits = MaskBits.extractBits(EltSizeInBits, BitOffset);
       EltBits[i] = Bits.getZExtValue();
     }
     return true;
   };
 
   // Collect constant bits and insert into mask/undef bit masks.
   auto CollectConstantBits = [SizeInBits](const Constant *Cst, APInt &Mask,
                                           APInt &Undefs, unsigned BitOffset) {

 
 static Constant *getConstantVector(MVT VT, const APInt &SplatValue,
                                    unsigned SplatBitSize, LLVMContext &C) {
   unsigned ScalarSize = VT.getScalarSizeInBits();
   unsigned NumElm = SplatBitSize / ScalarSize;
 
   SmallVector<Constant *, 32> ConstantVec;
   for (unsigned i = 0; i < NumElm; i++) {
-    APInt Val = SplatValue.lshr(ScalarSize * i).trunc(ScalarSize);
+    APInt Val = SplatValue.extractBits(ScalarSize, ScalarSize * i);
     Constant *Const;
     if (VT.isFloatingPoint()) {
       assert((ScalarSize == 32 || ScalarSize == 64) &&
              "Unsupported floating point scalar size");
       if (ScalarSize == 32)
         Const = ConstantFP::get(Type::getFloatTy(C), Val.bitsToFloat());
       else
         Const = ConstantFP::get(Type::getDoubleTy(C), Val.bitsToDouble());

lib/Target/X86/X86ShuffleDecodeConstantPool.cpp

static bool extractConstantMask(const Constant *C, unsigned MaskEltSizeInBits,
   assert((CstSizeInBits % MaskEltSizeInBits) == 0 &&
          "Unaligned shuffle mask size");
 
   unsigned NumMaskElts = CstSizeInBits / MaskEltSizeInBits;
   UndefElts = SmallBitVector(NumMaskElts, false);
   RawMask.resize(NumMaskElts, 0);
 
   for (unsigned i = 0; i != NumMaskElts; ++i) {
-    APInt EltUndef = UndefBits.lshr(i * MaskEltSizeInBits);
-    EltUndef = EltUndef.zextOrTrunc(MaskEltSizeInBits);
+    unsigned BitOffset = i * MaskEltSizeInBits;
+    APInt EltUndef = UndefBits.extractBits(MaskEltSizeInBits, BitOffset);
 
     // Only treat the element as UNDEF if all bits are UNDEF, otherwise
     // treat it as zero.
     if (EltUndef.isAllOnesValue()) {
       UndefElts[i] = true;
       RawMask[i] = 0;
       continue;
     }
 
-    APInt EltBits = MaskBits.lshr(i * MaskEltSizeInBits);
-    EltBits = EltBits.zextOrTrunc(MaskEltSizeInBits);
+    APInt EltBits = MaskBits.extractBits(MaskEltSizeInBits, BitOffset);
     RawMask[i] = EltBits.getZExtValue();
   }
 
   return true;
 }
 
 void DecodePSHUFBMask(const Constant *C, SmallVectorImpl<int> &ShuffleMask) {
   Type *MaskTy = C->getType();

unittests/ADT/APIntTest.cpp

for (int N : { 1, 2, 3, 4, 7, 8, 16, 32, 64, 127, 128, 129, 256 }) {
     APInt One(N, 1);
     for (int I = 1; I <= N; ++I) {
       APInt MaskVal = One.shl(I) - 1;
       EXPECT_TRUE(APIntOps::isMask(MaskVal));
     }
   }
 }
 
+TEST(APIntTest, extractBits) {
+  APInt i32(32, 0x1234567);
+  EXPECT_EQ(0x3456, i32.extractBits(16, 4));
+
+  APInt i256(256, -16776961 /* 0xFFFFFFFFFFFFFFFFFFFFFFFFFF0000FF */, true);
+  EXPECT_EQ(255, i256.extractBits(16, 0));
+  EXPECT_EQ(127, i256.extractBits(16, 1));
+  EXPECT_EQ(-1, i256.extractBits(32, 64).getSExtValue());
+  EXPECT_EQ(-1, i256.extractBits(128, 128).getSExtValue());
+  EXPECT_EQ(-8388481, i256.extractBits(128, 1).getSExtValue());
+}
+
 #if defined(__clang__)
 // Disable the pragma warning from versions of Clang without -Wself-move
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wunknown-pragmas"
 // Disable the warning that triggers on exactly what is being tested.
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wself-move"
 #endif