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

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.

Differential Revision: https://reviews.llvm.org/D30336

llvm-svn: 296272

include/llvm/ADT/APInt.h

@@ -1272,6 +1272,9 @@ public:
   /// 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
   /// @{

lib/Support/APInt.cpp

@@ -618,6 +618,42 @@ void APInt::flipBit(unsigned bitPosition) {
   else setBit(bitPosition);
 }
 
+APInt APInt::extractBits(unsigned numBits, unsigned bitPosition) const {
+  assert(numBits > 0 && "Can't extract zero bits");
+  assert(bitPosition < BitWidth && (numBits + bitPosition) <= BitWidth &&
+         "Illegal bit extraction");
+
+  if (isSingleWord())
+    return APInt(numBits, VAL >> bitPosition);
+
+  unsigned loBit = whichBit(bitPosition);
+  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 directly into place.
+  APInt Result(numBits, 0);
+  unsigned NumSrcWords = getNumWords();
+  unsigned NumDstWords = Result.getNumWords();
+
+  for (unsigned word = 0; word < NumDstWords; ++word) {
+    uint64_t w0 = pVal[loWord + word];
+    uint64_t w1 =
+        (loWord + word + 1) < NumSrcWords ? pVal[loWord + word + 1] : 0;
+    Result.pVal[word] = (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 ||

lib/Target/X86/X86ISelLowering.cpp

@@ -5207,8 +5207,8 @@ static bool getTargetConstantBitsFromNode(SDValue Op, unsigned EltSizeInBits,
     EltBits.resize(NumElts, APInt(EltSizeInBits, 0));
 
     for (unsigned i = 0; i != NumElts; ++i) {
-      APInt UndefEltBits = UndefBits.lshr(i * EltSizeInBits);
+      unsigned BitOffset = i * EltSizeInBits;
-      UndefEltBits = UndefEltBits.zextOrTrunc(EltSizeInBits);
+      APInt UndefEltBits = UndefBits.extractBits(EltSizeInBits, BitOffset);
 
       // Only treat an element as UNDEF if all bits are UNDEF.
       if (UndefEltBits.isAllOnesValue()) {
@@ -5223,7 +5223,7 @@ static bool getTargetConstantBitsFromNode(SDValue Op, unsigned EltSizeInBits,
       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;
@@ -6421,7 +6421,7 @@ static Constant *getConstantVector(MVT VT, const APInt &SplatValue,
 
   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) &&

lib/Target/X86/X86ShuffleDecodeConstantPool.cpp

@@ -77,8 +77,8 @@ static bool extractConstantMask(const Constant *C, unsigned MaskEltSizeInBits,
   RawMask.resize(NumMaskElts, 0);
 
   for (unsigned i = 0; i != NumMaskElts; ++i) {
-    APInt EltUndef = UndefBits.lshr(i * MaskEltSizeInBits);
+    unsigned BitOffset = i * MaskEltSizeInBits;
-    EltUndef = EltUndef.zextOrTrunc(MaskEltSizeInBits);
+    APInt EltUndef = UndefBits.extractBits(MaskEltSizeInBits, BitOffset);
 
     // Only treat the element as UNDEF if all bits are UNDEF, otherwise
     // treat it as zero.
@@ -88,8 +88,7 @@ static bool extractConstantMask(const Constant *C, unsigned MaskEltSizeInBits,
       continue;
     }
 
-    APInt EltBits = MaskBits.lshr(i * MaskEltSizeInBits);
+    APInt EltBits = MaskBits.extractBits(MaskEltSizeInBits, BitOffset);
-    EltBits = EltBits.zextOrTrunc(MaskEltSizeInBits);
     RawMask[i] = EltBits.getZExtValue();
   }
 

unittests/ADT/APIntTest.cpp

@@ -1505,3 +1505,19 @@ TEST(APIntTest, reverseBits) {
     }
   }
 }
+
+TEST(APIntTest, extractBits) {
+  APInt i32(32, 0x1234567);
+  EXPECT_EQ(0x3456, i32.extractBits(16, 4));
+
+  APInt i257(257, 0xFFFFFFFFFF0000FFull, true);
+  EXPECT_EQ(0xFFu, i257.extractBits(16, 0));
+  EXPECT_EQ((0xFFu >> 1), i257.extractBits(16, 1));
+  EXPECT_EQ(-1, i257.extractBits(32, 64).getSExtValue());
+  EXPECT_EQ(-1, i257.extractBits(128, 128).getSExtValue());
+  EXPECT_EQ(-1, i257.extractBits(66, 191).getSExtValue());
+  EXPECT_EQ(static_cast<int64_t>(0xFFFFFFFFFF80007Full),
+            i257.extractBits(128, 1).getSExtValue());
+  EXPECT_EQ(static_cast<int64_t>(0xFFFFFFFFFF80007Full),
+            i257.extractBits(129, 1).getSExtValue());
+}