[APInt] Add APInt::insertBits() method to insert an APInt into a larger APInt

[RKSimon]

Summary

We currently have to insert bits via a temporary variable of the same size as the target with various shift/mask stages, resulting in further temporary variables, all of which require the allocation of memory for large APInts (MaskSizeInBits > 64).

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

This patch adds the APInt::insertBits() helper method which avoids the temporary memory allocation and masks/inserts the raw bits directly into the target.

[filcab]

LGTM, but maybe wait a bit to let other people chime in.
Thank you!

unittests/ADT/APIntTest.cpp

1674

Why not this?

EXPECT_EQ((i127&UINT64_MAX).getZExtValue(), 0x3456ffffffffffff);
EXPECT_EQ((i127>>64).getZExtValue(), 0x7ffffffffff8012);

(Please double check the values)

[hans]

lib/Support/APInt.cpp

592

I suppose subBitWidth can't be zero, but maybe assert that? Otherwise we'd run into trouble with the shifts and so on.

597

I would assert bitPosition == 0 here, just to make it easier for the reader

617

This code is the same as for the isSingleWord() case except for pVal[loWord] instead of VAL.

If you really wanted to unify the two, we could use getRawData()[loWord] & = instead.

Not sure if it's worth it though.

638

It seems we could do better. Can't we do something similar to setBitsSlowCase, so we process word-by-word even when bitPosition isn't aligned?

[RKSimon]

Thanks for the feedback guys.

lib/Support/APInt.cpp

592

No problem - but tbh many APInt methods would fail with zero BitWidths.

617

I can merge these no problem, it won't particularly affect readability.

638

I'll add a TODO - there is scope for improvement here, but I've been struggling to find use cases. All the places I've converted to insertBits have been on who byte boundaries or better, and none actually cross multiple words - all other cases are current only in the unit tests. If more appear I'll gladly revisit this.

unittests/ADT/APIntTest.cpp

1674

OK, although I might actually use the i127.extracteBits().getZExtValue() pattern instead for clarity.

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);
 
+  /// Insert the bits from a smaller APInt starting at bitPosition.
+  void insertBits(const APInt &SubBits, 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.

lib/CodeGen/SelectionDAG/SelectionDAG.cpp

bool BuildVectorSDNode::isConstantSplat(APInt &SplatValue,
   for (unsigned j = 0; j < nOps; ++j) {
     unsigned i = isBigEndian ? nOps-1-j : j;
     SDValue OpVal = getOperand(i);
     unsigned BitPos = j * EltBitSize;
 
     if (OpVal.isUndef())
       SplatUndef.setBits(BitPos, BitPos + EltBitSize);
     else if (ConstantSDNode *CN = dyn_cast<ConstantSDNode>(OpVal))
-      SplatValue |= CN->getAPIntValue().zextOrTrunc(EltBitSize).
-                    zextOrTrunc(sz) << BitPos;
+      SplatValue.insertBits(CN->getAPIntValue().zextOrTrunc(EltBitSize),
+                            BitPos);
     else if (ConstantFPSDNode *CN = dyn_cast<ConstantFPSDNode>(OpVal))
-      SplatValue |= CN->getValueAPF().bitcastToAPInt().zextOrTrunc(sz) <<BitPos;
-     else
+      SplatValue.insertBits(CN->getValueAPF().bitcastToAPInt(), BitPos);
       return false;
   }
 
   // The build_vector is all constants or undefs.  Find the smallest element
   // size that splats the vector.
 
   HasAnyUndefs = (SplatUndef != 0);
   while (sz > 8) {

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);
 }
 
+void APInt::insertBits(const APInt &subBits, unsigned bitPosition) {
+  unsigned subBitWidth = subBits.getBitWidth();
+  assert((subBitWidth + bitPosition) <= BitWidth && "Illegal bit insertion");
+
+  // Insertion is a direct copy.
+  if (subBitWidth == BitWidth) {
+    *this = subBits;
+    return;
+  }
+
+  // Single word result can be done as a direct bitmask.
+  if (isSingleWord()) {
+    uint64_t mask = UINT64_MAX >> (APINT_BITS_PER_WORD - subBitWidth);
+    VAL &= ~(mask << bitPosition);
+    VAL |= (subBits.VAL << bitPosition);
+    return;
+  }
+
+  unsigned loBit = whichBit(bitPosition);
+  unsigned loWord = whichWord(bitPosition);
+  unsigned hi1Word = whichWord(bitPosition + subBitWidth - 1);
+
+  // Insertion within a single word can be done as a direct bitmask.
+  if (loWord == hi1Word) {
+    uint64_t mask = UINT64_MAX >> (APINT_BITS_PER_WORD - subBitWidth);
+    pVal[loWord] &= ~(mask << loBit);
+    pVal[loWord] |= (subBits.VAL << loBit);
+    return;
+  }
+
+  // Insert on word boundaries.
+  if (loBit == 0) {
+    // Direct copy whole words.
+    unsigned numWholeSubWords = subBitWidth / APINT_BITS_PER_WORD;
+    memcpy(pVal + loWord, subBits.getRawData(),
+           numWholeSubWords * APINT_WORD_SIZE);
+
+    // Mask+insert remaining bits.
+    unsigned remainingBits = subBitWidth % APINT_BITS_PER_WORD;
+    if (remainingBits != 0) {
+      uint64_t mask = UINT64_MAX >> (APINT_BITS_PER_WORD - remainingBits);
+      pVal[hi1Word] &= ~mask;
+      pVal[hi1Word] |= subBits.getWord(subBitWidth - 1);
+    }
+    return;
+  }
+
+  // General case - set/clear individual bits in dst based on src.
+  for (unsigned i = 0; i != subBitWidth; ++i) {
+    if (subBits[i])
+      setBit(bitPosition + i);
+    else
+      clearBit(bitPosition + i);
+  }
+}
+
 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);
 

lib/Target/X86/X86ISelLowering.cpp

auto CollectConstantBits = [SizeInBits](const Constant *Cst, APInt &Mask,
     if (!Cst)
       return false;
     unsigned CstSizeInBits = Cst->getType()->getPrimitiveSizeInBits();
     if (isa<UndefValue>(Cst)) {
       Undefs.setBits(BitOffset, BitOffset + CstSizeInBits);
       return true;
     }
     if (auto *CInt = dyn_cast<ConstantInt>(Cst)) {
-      Mask |= CInt->getValue().zextOrTrunc(SizeInBits).shl(BitOffset);
+      Mask.insertBits(CInt->getValue(), BitOffset);
       return true;
     }
     if (auto *CFP = dyn_cast<ConstantFP>(Cst)) {
-      APInt CstBits = CFP->getValueAPF().bitcastToAPInt();
-      Mask |= CstBits.zextOrTrunc(SizeInBits).shl(BitOffset);
+      Mask.insertBits(CFP->getValueAPF().bitcastToAPInt(), BitOffset);
       return true;
     }
     return false;
   };
 
   // Extract constant bits from build vector.
   if (ISD::isBuildVectorOfConstantSDNodes(Op.getNode())) {
     for (unsigned i = 0, e = Op.getNumOperands(); i != e; ++i) {
       const SDValue &Src = Op.getOperand(i);
       unsigned BitOffset = i * SrcEltSizeInBits;
       if (Src.isUndef()) {
         UndefBits.setBits(BitOffset, BitOffset + SrcEltSizeInBits);
         continue;
       }
       auto *Cst = cast<ConstantSDNode>(Src);
       APInt Bits = Cst->getAPIntValue().zextOrTrunc(SrcEltSizeInBits);
-      MaskBits |= Bits.zext(SizeInBits).shl(BitOffset);
+      MaskBits.insertBits(Bits, BitOffset);
     }
     return SplitBitData();
   }
 
   // Extract constant bits from constant pool vector.
   if (auto *Cst = getTargetConstantFromNode(Op)) {
     Type *CstTy = Cst->getType();
     if (!CstTy->isVectorTy() || (SizeInBits != CstTy->getPrimitiveSizeInBits()))

lib/Target/X86/X86ShuffleDecodeConstantPool.cpp

for (unsigned i = 0; i != NumCstElts; ++i) {
 
     unsigned BitOffset = i * CstEltSizeInBits;
 
     if (isa<UndefValue>(COp)) {
       UndefBits.setBits(BitOffset, BitOffset + CstEltSizeInBits);
       continue;
     }
 
-    auto *Elt = cast<ConstantInt>(COp);
-    MaskBits |= Elt->getValue().zextOrTrunc(CstSizeInBits).shl(BitOffset);
+    MaskBits.insertBits(cast<ConstantInt>(COp)->getValue(), BitOffset);
   }
 
   // Now extract the undef/constant bit data into the raw shuffle masks.
   assert((CstSizeInBits % MaskEltSizeInBits) == 0 &&
          "Unaligned shuffle mask size");
 
   unsigned NumMaskElts = CstSizeInBits / MaskEltSizeInBits;
   UndefElts = APInt(NumMaskElts, 0);

unittests/ADT/APIntTest.cpp

for (unsigned I = 0; I < N; ++I) {
       APInt X = APInt::getOneBitSet(N, I);
       APInt Y = APInt::getOneBitSet(N, N - (I + 1));
       EXPECT_EQ(Y, X.reverseBits());
       EXPECT_EQ(X, Y.reverseBits());
     }
   }
 }
 
+TEST(APIntTest, insertBits) {
+  APInt iSrc(31, 0x00123456);
+
+  // Direct copy.
+  APInt i31(31, 0x76543210ull);
+  i31.insertBits(iSrc, 0);
+  EXPECT_EQ(static_cast<int64_t>(0x00123456ull), i31.getSExtValue());
+
+  // Single word src/dst insertion.
+  APInt i63(63, 0x01234567FFFFFFFFull);
+  i63.insertBits(iSrc, 4);
+  EXPECT_EQ(static_cast<int64_t>(0x012345600123456Full), i63.getSExtValue());
+
+  // Insert single word src into one word of dst.
+  APInt i120(120, UINT64_MAX, true);
+  i120.insertBits(iSrc, 8);
+  EXPECT_EQ(static_cast<int64_t>(0xFFFFFF80123456FFull), i120.getSExtValue());
+
+  // Insert single word src into two words of dst.
+  APInt i127(127, UINT64_MAX, true);
+  i127.insertBits(iSrc, 48);
+  EXPECT_EQ(48u, i127.countLeadingOnes());
+  EXPECT_EQ(127u, i127.getActiveBits());
+  EXPECT_EQ(48u, i127.countTrailingOnes());
+  EXPECT_EQ(105u, i127.countPopulation());
+
+  // Insert on word boundaries.
+  APInt i128(128, 0);
+  i128.insertBits(APInt(64, UINT64_MAX, true), 0);
+  i128.insertBits(APInt(64, UINT64_MAX, true), 64);
+  EXPECT_EQ(-1, i128.getSExtValue());
+
+  APInt i256(256, UINT64_MAX, true);
+  i256.insertBits(APInt(65, 0), 0);
+  i256.insertBits(APInt(69, 0), 64);
+  i256.insertBits(APInt(128, 0), 128);
+  EXPECT_EQ(0u, i256.getSExtValue());
+
+  APInt i257(257, 0);
+  i257.insertBits(APInt(96, UINT64_MAX, true), 64);
+  EXPECT_EQ(97u, i257.countLeadingZeros());
+  EXPECT_EQ(64u, i257.countTrailingZeros());
+  EXPECT_EQ(96u, i257.countPopulation());
+
+  // General insertion.
+  APInt i260(260, UINT64_MAX, true);
+  i260.insertBits(APInt(129, 1ull << 48), 15);
+  EXPECT_EQ(116u, i260.countLeadingOnes());
+  EXPECT_EQ(260u, i260.getActiveBits());
+  EXPECT_EQ(145u, i260.getMinSignedBits());
+  EXPECT_EQ(15u, i260.countTrailingOnes());
+  EXPECT_EQ(132u, i260.countPopulation());
+}
+
 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());