[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 ↗

(On Diff #89656)

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 ↗

(On Diff #89656)

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 ↗

(On Diff #89656)

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 ↗

(On Diff #89656)

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 ↗

(On Diff #89656)

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.

llvm/trunk/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.

llvm/trunk/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, makeArrayRef(pVal + loWord, 1 + hiWord - 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();
 

llvm/trunk/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());

llvm/trunk/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();

llvm/trunk/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