[ADT] APIntOps::ScaleBitMask(): introduce 'BitMergingApproach' switch

[lebedev.ri]

Summary

Not always do we want to greedily merge the bits

[RKSimon]

TBH I think you'd be better off splitting ScaleBitMask into 2 methods rather than using an enum - unless you have a strong use case for needing to select between the 2 at runtime

[RKSimon]

TBH I think you'd be better off splitting ScaleBitMask into 2 methods rather than using an enum - unless you have a strong use case for needing to select between the 2 at runtime

llvm/include/llvm/ADT/APInt.h

 /// Compare two values, and if they are different, return the position of the
 /// most significant bit that is different in the values.
 Optional<unsigned> GetMostSignificantDifferentBit(const APInt &A,
                                                   const APInt &B);
 
 /// Splat/Merge neighboring bits to widen/narrow the bitmask represented
 /// by \param A to \param NewBitWidth bits.
 ///
-/// e.g. ScaleBitMask(0b0101, 8) -> 0b00110011
-/// e.g. ScaleBitMask(0b00011011, 4) -> 0b0111
+/// e.g. ScaleBitMask(0b0101, 8, <any>) -> 0b00110011
+/// e.g. ScaleBitMask(0b00011011, 4, BitMergingApproach::Greedy) -> 0b0111
+/// e.g. ScaleBitMask(0b00011011, 4, BitMergingApproach::Lossy)  -> 0b0001
 /// A.getBitwidth() or NewBitWidth must be a whole multiples of the other.
 ///
-/// TODO: Do we need a mode where all bits must be set when merging down?
-APInt ScaleBitMask(const APInt &A, unsigned NewBitWidth);
+/// How should we merge neighboring bits?
+enum struct BitMergingApproach {
+  // If *any* source bitisset, then the destination bit is set.
+  Greedy,
+  // If *all* source bits are set, only then the destination bit is set.
+  Lossy,
+};
+APInt ScaleBitMask(const APInt &A, unsigned NewBitWidth,
+                   BitMergingApproach Approach);
 } // namespace APIntOps
 
 // See friend declaration above. This additional declaration is required in
 // order to compile LLVM with IBM xlC compiler.
 hash_code hash_value(const APInt &Arg);
 
 /// StoreIntToMemory - Fills the StoreBytes bytes of memory starting from Dst
 /// with the integer held in IntVal.

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

case ISD::BITCAST: {
     if ((SubBitWidth % BitWidth) == 0) {
       assert(Op.getValueType().isVector() && "Expected bitcast to vector");
 
       // Collect known bits for the (smaller) output by collecting the known
       // bits from the overlapping larger input elements and extracting the
       // sub sections we actually care about.
       unsigned SubScale = SubBitWidth / BitWidth;
       APInt SubDemandedElts =
-          APIntOps::ScaleBitMask(DemandedElts, NumElts / SubScale);
+          APIntOps::ScaleBitMask(DemandedElts, NumElts / SubScale,
+                                 APIntOps::BitMergingApproach::Greedy);
       Known2 = computeKnownBits(N0, SubDemandedElts, Depth + 1);
 
       Known.Zero.setAllBits(); Known.One.setAllBits();
       for (unsigned i = 0; i != NumElts; ++i)
         if (DemandedElts[i]) {
           unsigned Shifts = IsLE ? i : NumElts - 1 - i;
           unsigned Offset = (Shifts % SubScale) * BitWidth;
           Known = KnownBits::commonBits(Known,
case ISD::BITCAST: {
 
     bool IsLE = getDataLayout().isLittleEndian();
 
     // Bitcast 'large element' scalar/vector to 'small element' vector.
     if ((SrcBits % VTBits) == 0) {
       assert(VT.isVector() && "Expected bitcast to vector");
 
       unsigned Scale = SrcBits / VTBits;
-      APInt SrcDemandedElts =
-          APIntOps::ScaleBitMask(DemandedElts, NumElts / Scale);
+      APInt SrcDemandedElts = APIntOps::ScaleBitMask(
+          DemandedElts, NumElts / Scale, APIntOps::BitMergingApproach::Greedy);
 
       // Fast case - sign splat can be simply split across the small elements.
       Tmp = ComputeNumSignBits(N0, SrcDemandedElts, Depth + 1);
       if (Tmp == SrcBits)
         return VTBits;
 
       // Slow case - determine how far the sign extends into each sub-element.
       Tmp2 = VTBits;

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

if (NumSrcElts == NumElts)
                                         KnownZero, TLO, Depth + 1);
 
     APInt SrcDemandedElts, SrcZero, SrcUndef;
 
     // Bitcast from 'large element' src vector to 'small element' vector, we
     // must demand a source element if any DemandedElt maps to it.
     if ((NumElts % NumSrcElts) == 0) {
       unsigned Scale = NumElts / NumSrcElts;
-      SrcDemandedElts = APIntOps::ScaleBitMask(DemandedElts, NumSrcElts);
+      SrcDemandedElts = APIntOps::ScaleBitMask(
+          DemandedElts, NumSrcElts, APIntOps::BitMergingApproach::Greedy);
       if (SimplifyDemandedVectorElts(Src, SrcDemandedElts, SrcUndef, SrcZero,
                                      TLO, Depth + 1))
         return true;
 
       // Try calling SimplifyDemandedBits, converting demanded elts to the bits
       // of the large element.
       // TODO - bigendian once we have test coverage.
       if (TLO.DAG.getDataLayout().isLittleEndian()) {
if ((NumElts % NumSrcElts) == 0) {
       }
     }
 
     // Bitcast from 'small element' src vector to 'large element' vector, we
     // demand all smaller source elements covered by the larger demanded element
     // of this vector.
     if ((NumSrcElts % NumElts) == 0) {
       unsigned Scale = NumSrcElts / NumElts;
-      SrcDemandedElts = APIntOps::ScaleBitMask(DemandedElts, NumSrcElts);
+      SrcDemandedElts = APIntOps::ScaleBitMask(
+          DemandedElts, NumSrcElts, APIntOps::BitMergingApproach::Greedy);
       if (SimplifyDemandedVectorElts(Src, SrcDemandedElts, SrcUndef, SrcZero,
                                      TLO, Depth + 1))
         return true;
 
       // If all the src elements covering an output element are zero/undef, then
       // the output element will be as well, assuming it was demanded.
       for (unsigned i = 0; i != NumElts; ++i) {
         if (DemandedElts[i]) {

llvm/lib/Support/APInt.cpp

 Optional<unsigned>
 llvm::APIntOps::GetMostSignificantDifferentBit(const APInt &A, const APInt &B) {
   assert(A.getBitWidth() == B.getBitWidth() && "Must have the same bitwidth");
   if (A == B)
     return llvm::None;
   return A.getBitWidth() - ((A ^ B).countLeadingZeros() + 1);
 }
 
-APInt llvm::APIntOps::ScaleBitMask(const APInt &A, unsigned NewBitWidth) {
+APInt llvm::APIntOps::ScaleBitMask(const APInt &A, unsigned NewBitWidth,
+                                   BitMergingApproach Approach) {
   unsigned OldBitWidth = A.getBitWidth();
   assert((((OldBitWidth % NewBitWidth) == 0) ||
           ((NewBitWidth % OldBitWidth) == 0)) &&
          "One size should be a multiple of the other one. "
          "Can't do fractional scaling.");
 
   // Check for matching bitwidths.
   if (OldBitWidth == NewBitWidth)
if (NewBitWidth > OldBitWidth) {
     // Repeat bits.
     unsigned Scale = NewBitWidth / OldBitWidth;
     for (unsigned i = 0; i != OldBitWidth; ++i)
       if (A[i])
         NewA.setBits(i * Scale, (i + 1) * Scale);
   } else {
     // Merge bits - if any old bit is set, then set scale equivalent new bit.
     unsigned Scale = OldBitWidth / NewBitWidth;
-    for (unsigned i = 0; i != NewBitWidth; ++i)
-      if (!A.extractBits(Scale, i * Scale).isNullValue())
+    for (unsigned i = 0; i != NewBitWidth; ++i) {
+      APInt Part = A.extractBits(Scale, i * Scale);
+      if (Approach == BitMergingApproach::Greedy ? !Part.isNullValue()
+                                                 : Part.isAllOnes())
         NewA.setBit(i);
     }
+  }
 
   return NewA;
 }
 
 /// StoreIntToMemory - Fills the StoreBytes bytes of memory starting from Dst
 /// with the integer held in IntVal.
 void llvm::StoreIntToMemory(const APInt &IntVal, uint8_t *Dst,
                             unsigned StoreBytes) {

llvm/lib/Target/X86/X86ISelLowering.cpp

case X86ISD::PSADBW: {
     assert(VT.getScalarType() == MVT::i64 &&
            LHS.getValueType() == RHS.getValueType() &&
            LHS.getValueType().getScalarType() == MVT::i8 &&
            "Unexpected PSADBW types");
 
     // Aggressively peek through ops to get at the demanded elts.
     if (!DemandedElts.isAllOnesValue()) {
       unsigned NumSrcElts = LHS.getValueType().getVectorNumElements();
-      APInt DemandedSrcElts = APIntOps::ScaleBitMask(DemandedElts, NumSrcElts);
+      APInt DemandedSrcElts = APIntOps::ScaleBitMask(
+          DemandedElts, NumSrcElts, APIntOps::BitMergingApproach::Greedy);
       SDValue NewLHS = SimplifyMultipleUseDemandedVectorElts(
           LHS, DemandedSrcElts, TLO.DAG, Depth + 1);
       SDValue NewRHS = SimplifyMultipleUseDemandedVectorElts(
           RHS, DemandedSrcElts, TLO.DAG, Depth + 1);
       if (NewLHS || NewRHS) {
         NewLHS = NewLHS ? NewLHS : LHS;
         NewRHS = NewRHS ? NewRHS : RHS;
         return TLO.CombineTo(

llvm/unittests/ADT/APIntTest.cpp

TEST(APIntTest, ZeroWidth) {
   // Move ctor (keep at the end of the method since moves are destructive).
   APInt MZW1(std::move(ZW));
   EXPECT_EQ(0U, MZW1.getBitWidth());
   // Move Assignment
   MZW1 = std::move(ZW2);
   EXPECT_EQ(0U, MZW1.getBitWidth());
 }
 
-TEST(APIntTest, ScaleBitMask) {
-  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(2, 0x00), 8), APInt(8, 0x00));
-  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(2, 0x01), 8), APInt(8, 0x0F));
-  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(2, 0x02), 8), APInt(8, 0xF0));
-  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(2, 0x03), 8), APInt(8, 0xFF));
-
-  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(8, 0x00), 4), APInt(4, 0x00));
-  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(8, 0xFF), 4), APInt(4, 0x0F));
-  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(8, 0xE4), 4), APInt(4, 0x0E));
+TEST(APIntTest, ScaleBitMaskGreedily) {
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(2, 0x00), 8,
+                                   APIntOps::BitMergingApproach::Greedy),
+            APInt(8, 0x00));
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(2, 0x01), 8,
+                                   APIntOps::BitMergingApproach::Greedy),
+            APInt(8, 0x0F));
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(2, 0x02), 8,
+                                   APIntOps::BitMergingApproach::Greedy),
+            APInt(8, 0xF0));
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(2, 0x03), 8,
+                                   APIntOps::BitMergingApproach::Greedy),
+            APInt(8, 0xFF));
+
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(8, 0x00), 4,
+                                   APIntOps::BitMergingApproach::Greedy),
+            APInt(4, 0x00));
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(8, 0xFF), 4,
+                                   APIntOps::BitMergingApproach::Greedy),
+            APInt(4, 0x0F));
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(8, 0xE4), 4,
+                                   APIntOps::BitMergingApproach::Greedy),
+            APInt(4, 0x0E));
+
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(8, 0x00), 8,
+                                   APIntOps::BitMergingApproach::Greedy),
+            APInt(8, 0x00));
 
-  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(8, 0x00), 8), APInt(8, 0x00));
-
-  EXPECT_EQ(APIntOps::ScaleBitMask(APInt::getNullValue(1024), 4096),
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt::getNullValue(1024), 4096,
+                                   APIntOps::BitMergingApproach::Greedy),
             APInt::getNullValue(4096));
-  EXPECT_EQ(APIntOps::ScaleBitMask(APInt::getAllOnes(4096), 256),
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt::getAllOnes(4096), 256,
+                                   APIntOps::BitMergingApproach::Greedy),
             APInt::getAllOnes(256));
-  EXPECT_EQ(APIntOps::ScaleBitMask(APInt::getOneBitSet(4096, 32), 256),
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt::getOneBitSet(4096, 32), 256,
+                                   APIntOps::BitMergingApproach::Greedy),
             APInt::getOneBitSet(256, 2));
 }
 
+TEST(APIntTest, ScaleBitMaskLossily) {
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(2, 0x00), 8,
+                                   APIntOps::BitMergingApproach::Lossy),
+            APInt(8, 0x00));
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(2, 0x01), 8,
+                                   APIntOps::BitMergingApproach::Lossy),
+            APInt(8, 0x0F));
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(2, 0x02), 8,
+                                   APIntOps::BitMergingApproach::Lossy),
+            APInt(8, 0xF0));
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(2, 0x03), 8,
+                                   APIntOps::BitMergingApproach::Lossy),
+            APInt(8, 0xFF));
+
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(8, 0x00), 4,
+                                   APIntOps::BitMergingApproach::Lossy),
+            APInt(4, 0x00));
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(8, 0xFF), 4,
+                                   APIntOps::BitMergingApproach::Lossy),
+            APInt(4, 0x0F));
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(8, 0xE4), 4,
+                                   APIntOps::BitMergingApproach::Lossy),
+            APInt(4, 0x08));
+
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt(8, 0x00), 8,
+                                   APIntOps::BitMergingApproach::Lossy),
+            APInt(8, 0x00));
+
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt::getNullValue(1024), 4096,
+                                   APIntOps::BitMergingApproach::Lossy),
+            APInt::getNullValue(4096));
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt::getAllOnes(4096), 256,
+                                   APIntOps::BitMergingApproach::Lossy),
+            APInt::getAllOnes(256));
+  EXPECT_EQ(APIntOps::ScaleBitMask(APInt::getOneBitSet(4096, 32), 256,
+                                   APIntOps::BitMergingApproach::Lossy),
+            APInt::getNullValue(256));
+}
+
 } // end anonymous namespace