[APInt] Add APInt::setBits() method to set all bits in range

[RKSimon]

Summary

The current pattern for setting bits in range is typically:

Mask |= APInt::getBitsSet(MaskSizeInBits, LoPos, HiPos);

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

This is one of the key compile time issues identified in PR32037.

This patch adds the APInt::setBits() helper method which avoids the temporary memory allocation completely, this first implementation uses setBit() internally instead but already significantly reduces the regression in PR32037 (~10% drop). Additional optimization may be possible.

I investigated whether there is need for APInt::clearBits() and APInt::flipBits() equivalents but haven't seen these patterns to be particularly common.

[hans]

lib/Support/APInt.cpp

574

Is the hiBit < loBit case common?

If not (or maybe even then) perhaps just setBits(0, hiBit); setBits(loBit, BitWidth); return and focus on the loBit < hiBit case below?

580

Would it be worth trying to do this word-by-word instead?

[RKSimon]

lib/Support/APInt.cpp

574

No, I was just matching the behaviour of the '|= APInt::getBitsSet()' pattern but it doesn't seem common compared to using APInt::getHighBitsSet and APInt::getLowBitsSet instead and in those cases its more important to be able to support the 'all bits case'.

I'll change it.

580

Yes I can do that.

[filcab]

include/llvm/ADT/APInt.h

1244

I don't think you need to repeat that sentence.
I'd also suggest removing \brief, but I'm ok with keeping it consistent.

unittests/ADT/APIntTest.cpp

152

Missing a test for !isSingleWord() && loBit <= HiBit
(I'd actually prefer one for loBit < hiBit and one for loBit == HiBit, to make it explicit what you expect to happen on the edge case)

[RKSimon]

Dropped the wrap around functionality as it doesn't appear to have many use cases - and makes it easier to drop in as a replacement for APInt::getLowBitsSet/getHighBitsSet.

Removed use of APInt::setBit and do everything at the word level instead.

Added extra test cases that should cover all mask permutations.

[hans]

Nice!

lib/Support/APInt.cpp

582

Other parts of the file uses UINT64_MAX instead, which I think is nicer.

591

I'd suggest computing the final mask here, instead of starting in this declaration and then doing some more computation when or-ing below.

592

This is elegant, but maybe UINT64_MAX << maskBit(loBit) is more efficient?

593

Would this be more efficient? UINT64_MAX >> (APINT_BITS_PER_WORD - maskBit(hiBit))

594

I would have gone for word < hiWord here, not that it should matter.

595

UINT64_MAX here too.

unittests/ADT/APIntTest.cpp

204

I think there's still no test for loBit == hiBit?

[RKSimon]

Updated based on Hans' feedback.

I didn't do the UINT64_MAX >> (64 - whichBit(hiBit - 1)) change to avoid having to handle the shift by 64 case

[hans]

In D30265#684965, @RKSimon wrote:

I didn't do the UINT64_MAX >> (64 - whichBit(hiBit - 1)) change to avoid having to handle the shift by 64 case

Oh, I used whichBit(hiBit) instead of whichBit(hiBit - 1), which isn't right because hiWord refers to the highest word where we want set at least one bit.

So how about UINT64_MAX >> (64 - whichBit(hiBit - 1) - 1)?

Your solution is fine too, but it fells like this should be possible.

[RKSimon]

Updated hiMask shifts

[hans]

Looks good to me.

include/llvm/ADT/APInt.h

void setAllBits() {
     clearUnusedBits();
   }
 
   /// \brief Set a given bit to 1.
   ///
   /// Set the given bit to 1 whose position is given as "bitPosition".
   void setBit(unsigned bitPosition);
 
+  /// Set the bits from loBit (inclusive) to hiBit (exclusive) to 1.
+  void setBits(unsigned loBit, unsigned hiBit);
+
   /// \brief Set every bit to 0.
   void clearAllBits() {
     if (isSingleWord())
       VAL = 0;
     else
       memset(pVal, 0, getNumWords() * APINT_WORD_SIZE);
   }
 

lib/Support/APInt.cpp

 
 void APInt::setBit(unsigned bitPosition) {
   if (isSingleWord())
     VAL |= maskBit(bitPosition);
   else
     pVal[whichWord(bitPosition)] |= maskBit(bitPosition);
 }
 
+void APInt::setBits(unsigned loBit, unsigned hiBit) {
+  assert(hiBit <= BitWidth && "hiBit out of range");
+  assert(loBit <= hiBit && loBit <= BitWidth && "loBit out of range");
+
+  if (loBit == hiBit)
+    return;
+
+  if (isSingleWord())
+    *this |= APInt::getBitsSet(BitWidth, loBit, hiBit);
+  else {
+    unsigned loWord = whichWord(loBit);
+    unsigned hiWord = whichWord(hiBit - 1);
+    if (loWord == hiWord) {
+      // Set bits are all within the same word, create a [loBit,hiBit) mask.
+      uint64_t mask = ~uint64_t(0ULL);
+      mask >>= (APINT_BITS_PER_WORD - (hiBit - loBit));
+      mask <<= whichBit(loBit);
+      pVal[loWord] |= mask;
+    } else {
+      // Set bits span multiple words, create a lo mask with set bits starting
+      // at loBit, a hi mask with set bits below hiBit and set all bits of the
+      // words in between.
+      uint64_t loBitMask = maskBit(loBit);
+      uint64_t hiBitMask = maskBit(hiBit - 1);
+      pVal[loWord] |= ~(loBitMask - 1);
+      pVal[hiWord] |= hiBitMask | (hiBitMask - 1);
+      for (unsigned word = loWord + 1; word != hiWord; ++word)
+        pVal[word] = ~uint64_t(0ULL);
+    }
+  }
+}
+
 /// Set the given bit to 0 whose position is given as "bitPosition".
 /// @brief Set a given bit to 0.
 void APInt::clearBit(unsigned bitPosition) {
   if (isSingleWord())
     VAL &= ~maskBit(bitPosition);
   else
     pVal[whichWord(bitPosition)] &= ~maskBit(bitPosition);
 }

lib/Target/X86/X86ISelLowering.cpp

static bool getTargetConstantBitsFromNode(SDValue Op, unsigned EltSizeInBits,
 
   // Collect constant bits and insert into mask/undef bit masks.
   auto CollectConstantBits = [SizeInBits](const Constant *Cst, APInt &Mask,
                                           APInt &Undefs, unsigned BitOffset) {
     if (!Cst)
       return false;
     unsigned CstSizeInBits = Cst->getType()->getPrimitiveSizeInBits();
     if (isa<UndefValue>(Cst)) {
-      unsigned HiBits = BitOffset + CstSizeInBits;
-      Undefs |= APInt::getBitsSet(SizeInBits, BitOffset, HiBits);
+      Undefs.setBits(BitOffset, BitOffset + CstSizeInBits);
       return true;
     }
     if (auto *CInt = dyn_cast<ConstantInt>(Cst)) {
       Mask |= CInt->getValue().zextOrTrunc(SizeInBits).shl(BitOffset);
       return true;
     }
     if (auto *CFP = dyn_cast<ConstantFP>(Cst)) {
       APInt CstBits = CFP->getValueAPF().bitcastToAPInt();
       Mask |= CstBits.zextOrTrunc(SizeInBits).shl(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()) {
-        unsigned HiBits = BitOffset + SrcEltSizeInBits;
-        UndefBits |= APInt::getBitsSet(SizeInBits, BitOffset, HiBits);
+        UndefBits.setBits(BitOffset, BitOffset + SrcEltSizeInBits);
         continue;
       }
       auto *Cst = cast<ConstantSDNode>(Src);
       APInt Bits = Cst->getAPIntValue().zextOrTrunc(SrcEltSizeInBits);
       MaskBits |= Bits.zext(SizeInBits).shl(BitOffset);
     }
     return SplitBitData();
   }
void X86TargetLowering::computeKnownBitsForTargetNode(const SDValue Op,
   case X86ISD::OR:
   case X86ISD::XOR:
   case X86ISD::AND:
     // These nodes' second result is a boolean.
     if (Op.getResNo() == 0)
       break;
     LLVM_FALLTHROUGH;
   case X86ISD::SETCC:
-    KnownZero |= APInt::getHighBitsSet(BitWidth, BitWidth - 1);
+    KnownZero.setBits(1, BitWidth);
     break;
   case X86ISD::MOVMSK: {
     unsigned NumLoBits = Op.getOperand(0).getValueType().getVectorNumElements();
-    KnownZero = APInt::getHighBitsSet(BitWidth, BitWidth - NumLoBits);
+    KnownZero.setBits(NumLoBits, BitWidth);
     break;
   }
   case X86ISD::VZEXT: {
     SDValue N0 = Op.getOperand(0);
     unsigned NumElts = Op.getValueType().getVectorNumElements();
 
     EVT SrcVT = N0.getValueType();
     unsigned InNumElts = SrcVT.getVectorNumElements();
     unsigned InBitWidth = SrcVT.getScalarSizeInBits();
     assert(InNumElts >= NumElts && "Illegal VZEXT input");
 
     KnownZero = KnownOne = APInt(InBitWidth, 0);
     APInt DemandedSrcElts = APInt::getLowBitsSet(InNumElts, NumElts);
     DAG.computeKnownBits(N0, KnownZero, KnownOne, DemandedSrcElts, Depth + 1);
     KnownOne = KnownOne.zext(BitWidth);
     KnownZero = KnownZero.zext(BitWidth);
-    KnownZero |= APInt::getHighBitsSet(BitWidth, BitWidth - InBitWidth);
+    KnownZero.setBits(InBitWidth, BitWidth);
     break;
   }
   }
 }
 
 unsigned X86TargetLowering::ComputeNumSignBitsForTargetNode(
     SDValue Op, const SelectionDAG &DAG, unsigned Depth) const {
   // SETCC_CARRY sets the dest to ~0 for true or 0 for false.

lib/Target/X86/X86ShuffleDecodeConstantPool.cpp

static bool extractConstantMask(const Constant *C, unsigned MaskEltSizeInBits,
   for (unsigned i = 0; i != NumCstElts; ++i) {
     Constant *COp = C->getAggregateElement(i);
     if (!COp || (!isa<UndefValue>(COp) && !isa<ConstantInt>(COp)))
       return false;
 
     unsigned BitOffset = i * CstEltSizeInBits;
 
     if (isa<UndefValue>(COp)) {
-      unsigned HiBits = BitOffset + CstEltSizeInBits;
-      UndefBits |= APInt::getBitsSet(CstSizeInBits, BitOffset, HiBits);
+      UndefBits.setBits(BitOffset, BitOffset + CstEltSizeInBits);
       continue;
     }
 
     auto *Elt = cast<ConstantInt>(COp);
     MaskBits |= Elt->getValue().zextOrTrunc(CstSizeInBits).shl(BitOffset);
   }
 
   // Now extract the undef/constant bit data into the raw shuffle masks.

unittests/ADT/APIntTest.cpp

TEST(APIntTest, i33_Count) {
   EXPECT_EQ(32u, i33minus2.countLeadingOnes());
   EXPECT_EQ(33u, i33minus2.getActiveBits());
   EXPECT_EQ(1u, i33minus2.countTrailingZeros());
   EXPECT_EQ(32u, i33minus2.countPopulation());
   EXPECT_EQ(-2, i33minus2.getSExtValue());
   EXPECT_EQ(((uint64_t)-2)&((1ull<<33) -1), i33minus2.getZExtValue());
 }
 
+TEST(APIntTest, i61_Count) {
+  APInt i61(61, 1 << 15);
+  EXPECT_EQ(45u, i61.countLeadingZeros());
+  EXPECT_EQ(0u, i61.countLeadingOnes());
+  EXPECT_EQ(16u, i61.getActiveBits());
+  EXPECT_EQ(15u, i61.countTrailingZeros());
+  EXPECT_EQ(1u, i61.countPopulation());
+  EXPECT_EQ((1 << 15), i61.getSExtValue());
+  EXPECT_EQ((1 << 15), i61.getZExtValue());
+
+  i61.setBits(8, 19);
+  EXPECT_EQ(42u, i61.countLeadingZeros());
+  EXPECT_EQ(0u, i61.countLeadingOnes());
+  EXPECT_EQ(19u, i61.getActiveBits());
+  EXPECT_EQ(8u, i61.countTrailingZeros());
+  EXPECT_EQ(11u, i61.countPopulation());
+  EXPECT_EQ((1 << 19) - (1 << 8), i61.getSExtValue());
+  EXPECT_EQ((1 << 19) - (1 << 8), i61.getZExtValue());
+}
+
 TEST(APIntTest, i65_Count) {
   APInt i65(65, 0, true);
   EXPECT_EQ(65u, i65.countLeadingZeros());
   EXPECT_EQ(0u, i65.countLeadingOnes());
   EXPECT_EQ(0u, i65.getActiveBits());
   EXPECT_EQ(1u, i65.getActiveWords());
   EXPECT_EQ(65u, i65.countTrailingZeros());
   EXPECT_EQ(0u, i65.countPopulation());
TEST(APIntTest, i128_PositiveCount) {
   EXPECT_EQ(127u, one.countLeadingZeros());
   EXPECT_EQ(0u, one.countLeadingOnes());
   EXPECT_EQ(1u, one.getActiveBits());
   EXPECT_EQ(0u, one.countTrailingZeros());
   EXPECT_EQ(1u, one.countTrailingOnes());
   EXPECT_EQ(1u, one.countPopulation());
   EXPECT_EQ(1, one.getSExtValue());
   EXPECT_EQ(1u, one.getZExtValue());
+
+  APInt s128(128, 2, true);
+  EXPECT_EQ(126u, s128.countLeadingZeros());
+  EXPECT_EQ(0u, s128.countLeadingOnes());
+  EXPECT_EQ(2u, s128.getActiveBits());
+  EXPECT_EQ(1u, s128.countTrailingZeros());
+  EXPECT_EQ(0u, s128.countTrailingOnes());
+  EXPECT_EQ(1u, s128.countPopulation());
+  EXPECT_EQ(2, s128.getSExtValue());
+  EXPECT_EQ(2u, s128.getZExtValue());
+
+  s128.setBits(3, 32);
+  EXPECT_EQ(96u, s128.countLeadingZeros());
+  EXPECT_EQ(0u, s128.countLeadingOnes());
+  EXPECT_EQ(32u, s128.getActiveBits());
+  EXPECT_EQ(33u, s128.getMinSignedBits());
+  EXPECT_EQ(1u, s128.countTrailingZeros());
+  EXPECT_EQ(0u, s128.countTrailingOnes());
+  EXPECT_EQ(30u, s128.countPopulation());
+  EXPECT_EQ(static_cast<uint32_t>((~0u << 3) | 2), s128.getZExtValue());
+
+  s128.setBits(62, 128);
+  EXPECT_EQ(0u, s128.countLeadingZeros());
+  EXPECT_EQ(66u, s128.countLeadingOnes());
+  EXPECT_EQ(128u, s128.getActiveBits());
+  EXPECT_EQ(63u, s128.getMinSignedBits());
+  EXPECT_EQ(1u, s128.countTrailingZeros());
+  EXPECT_EQ(0u, s128.countTrailingOnes());
+  EXPECT_EQ(96u, s128.countPopulation());
+  EXPECT_EQ(static_cast<int64_t>((3ull << 62) |
+                                 static_cast<uint32_t>((~0u << 3) | 2)),
+            s128.getSExtValue());
+}
+
+TEST(APIntTest, i256) {
+  APInt s256(256, 15, true);
+  EXPECT_EQ(252u, s256.countLeadingZeros());
+  EXPECT_EQ(0u, s256.countLeadingOnes());
+  EXPECT_EQ(4u, s256.getActiveBits());
+  EXPECT_EQ(0u, s256.countTrailingZeros());
+  EXPECT_EQ(4u, s256.countTrailingOnes());
+  EXPECT_EQ(4u, s256.countPopulation());
+  EXPECT_EQ(15, s256.getSExtValue());
+  EXPECT_EQ(15u, s256.getZExtValue());
+
+  s256.setBits(62, 66);
+  EXPECT_EQ(190u, s256.countLeadingZeros());
+  EXPECT_EQ(0u, s256.countLeadingOnes());
+  EXPECT_EQ(66u, s256.getActiveBits());
+  EXPECT_EQ(67u, s256.getMinSignedBits());
+  EXPECT_EQ(0u, s256.countTrailingZeros());
+  EXPECT_EQ(4u, s256.countTrailingOnes());
+  EXPECT_EQ(8u, s256.countPopulation());
+
+  s256.setBits(60, 256);
+  EXPECT_EQ(0u, s256.countLeadingZeros());
+  EXPECT_EQ(196u, s256.countLeadingOnes());
+  EXPECT_EQ(256u, s256.getActiveBits());
+  EXPECT_EQ(61u, s256.getMinSignedBits());
+  EXPECT_EQ(0u, s256.countTrailingZeros());
+  EXPECT_EQ(4u, s256.countTrailingOnes());
+  EXPECT_EQ(200u, s256.countPopulation());
+  EXPECT_EQ(static_cast<int64_t>((~0ull << 60) | 15), s256.getSExtValue());
 }
 
 TEST(APIntTest, i1) {
   const APInt neg_two(1, static_cast<uint64_t>(-2), true);
   const APInt neg_one(1, static_cast<uint64_t>(-1), true);
   const APInt zero(1, 0);
   const APInt one(1, 1);
   const APInt two(1, 2);