[X86] Use APInt instead of SmallBitVector for tracking undef elements in constant pool shuffle decoding

[craig.topper]

Summary

SmallBitVector uses a malloc for more than 58 bits on a 64-bit target and more than 27 bits on a 32-bit target. Some of the vector types we deal with here use more than those number of elements and therefore cause a malloc.

APInt on the other hand supports up to 64 bits without a malloc. That's the maximum number of bits we need here so we can avoid a malloc for all cases by using APInt. This will incur a minor increase in stack usage due to APInt storing the bit count separately from the data bits unlike SmallBitVector, but that should be ok.

[RKSimon]

LGTM

llvm/trunk/lib/Target/X86/X86ShuffleDecodeConstantPool.cpp

 //===-- X86ShuffleDecodeConstantPool.cpp - X86 shuffle decode -------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // Define several functions to decode x86 specific shuffle semantics using
 // constants from the constant pool.
 //
 //===----------------------------------------------------------------------===//
 
 #include "X86ShuffleDecodeConstantPool.h"
 #include "Utils/X86ShuffleDecode.h"
-#include "llvm/ADT/SmallBitVector.h"
+#include "llvm/ADT/APInt.h"
 #include "llvm/CodeGen/MachineValueType.h"
 #include "llvm/IR/Constants.h"
 
 //===----------------------------------------------------------------------===//
 //  Vector Mask Decoding
 //===----------------------------------------------------------------------===//
 
 namespace llvm {
 
 static bool extractConstantMask(const Constant *C, unsigned MaskEltSizeInBits,
-                                SmallBitVector &UndefElts,
+                                APInt &UndefElts,
                                 SmallVectorImpl<uint64_t> &RawMask) {
   // It is not an error for shuffle masks to not be a vector of
   // MaskEltSizeInBits because the constant pool uniques constants by their
   // bit representation.
   // e.g. the following take up the same space in the constant pool:
   //   i128 -170141183420855150465331762880109871104
   //
   //   <2 x i64> <i64 -9223372034707292160, i64 -9223372034707292160>
for (unsigned i = 0; i != NumCstElts; ++i) {
     MaskBits |= Elt->getValue().zextOrTrunc(CstSizeInBits).shl(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 = SmallBitVector(NumMaskElts, false);
+  UndefElts = APInt(NumMaskElts, 0);
   RawMask.resize(NumMaskElts, 0);
 
   for (unsigned i = 0; i != NumMaskElts; ++i) {
     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;
+      UndefElts.setBit(i);
       RawMask[i] = 0;
       continue;
     }
 
     APInt EltBits = MaskBits.extractBits(MaskEltSizeInBits, BitOffset);
     RawMask[i] = EltBits.getZExtValue();
   }
 
   return true;
 }
 
 void DecodePSHUFBMask(const Constant *C, SmallVectorImpl<int> &ShuffleMask) {
   Type *MaskTy = C->getType();
   unsigned MaskTySize = MaskTy->getPrimitiveSizeInBits();
   (void)MaskTySize;
   assert((MaskTySize == 128 || MaskTySize == 256 || MaskTySize == 512) &&
          "Unexpected vector size.");
 
   // The shuffle mask requires a byte vector.
-  SmallBitVector UndefElts;
+  APInt UndefElts;
   SmallVector<uint64_t, 32> RawMask;
   if (!extractConstantMask(C, 8, UndefElts, RawMask))
     return;
 
   unsigned NumElts = RawMask.size();
   assert((NumElts == 16 || NumElts == 32 || NumElts == 64) &&
          "Unexpected number of vector elements.");
 
void DecodeVPERMILPMask(const Constant *C, unsigned ElSize,
   Type *MaskTy = C->getType();
   unsigned MaskTySize = MaskTy->getPrimitiveSizeInBits();
   (void)MaskTySize;
   assert((MaskTySize == 128 || MaskTySize == 256 || MaskTySize == 512) &&
          "Unexpected vector size.");
   assert((ElSize == 32 || ElSize == 64) && "Unexpected vector element size.");
 
   // The shuffle mask requires elements the same size as the target.
-  SmallBitVector UndefElts;
+  APInt UndefElts;
   SmallVector<uint64_t, 8> RawMask;
   if (!extractConstantMask(C, ElSize, UndefElts, RawMask))
     return;
 
   unsigned NumElts = RawMask.size();
   unsigned NumEltsPerLane = 128 / ElSize;
   assert((NumElts == 2 || NumElts == 4 || NumElts == 8 || NumElts == 16) &&
          "Unexpected number of vector elements.");

 void DecodeVPERMIL2PMask(const Constant *C, unsigned M2Z, unsigned ElSize,
                          SmallVectorImpl<int> &ShuffleMask) {
   Type *MaskTy = C->getType();
   unsigned MaskTySize = MaskTy->getPrimitiveSizeInBits();
   (void)MaskTySize;
   assert((MaskTySize == 128 || MaskTySize == 256) && "Unexpected vector size.");
 
   // The shuffle mask requires elements the same size as the target.
-  SmallBitVector UndefElts;
+  APInt UndefElts;
   SmallVector<uint64_t, 8> RawMask;
   if (!extractConstantMask(C, ElSize, UndefElts, RawMask))
     return;
 
   unsigned NumElts = RawMask.size();
   unsigned NumEltsPerLane = 128 / ElSize;
   assert((NumElts == 2 || NumElts == 4 || NumElts == 8) &&
          "Unexpected number of vector elements.");
void DecodeVPERMIL2PMask(const Constant *C, unsigned M2Z, unsigned ElSize,
   }
 }
 
 void DecodeVPPERMMask(const Constant *C, SmallVectorImpl<int> &ShuffleMask) {
   assert(C->getType()->getPrimitiveSizeInBits() == 128 &&
          "Unexpected vector size.");
 
   // The shuffle mask requires a byte vector.
-  SmallBitVector UndefElts;
+  APInt UndefElts;
   SmallVector<uint64_t, 32> RawMask;
   if (!extractConstantMask(C, 8, UndefElts, RawMask))
     return;
 
   unsigned NumElts = RawMask.size();
   assert(NumElts == 16 && "Unexpected number of vector elements.");
 
   for (unsigned i = 0; i != NumElts; ++i) {
void DecodeVPERMVMask(const Constant *C, unsigned ElSize,
   unsigned MaskTySize = MaskTy->getPrimitiveSizeInBits();
   (void)MaskTySize;
   assert((MaskTySize == 128 || MaskTySize == 256 || MaskTySize == 512) &&
          "Unexpected vector size.");
   assert((ElSize == 8 || ElSize == 16 || ElSize == 32 || ElSize == 64) &&
          "Unexpected vector element size.");
 
   // The shuffle mask requires elements the same size as the target.
-  SmallBitVector UndefElts;
+  APInt UndefElts;
   SmallVector<uint64_t, 8> RawMask;
   if (!extractConstantMask(C, ElSize, UndefElts, RawMask))
     return;
 
   unsigned NumElts = RawMask.size();
 
   for (unsigned i = 0; i != NumElts; ++i) {
     if (UndefElts[i]) {
void DecodeVPERMV3Mask(const Constant *C, unsigned ElSize,
   unsigned MaskTySize = MaskTy->getPrimitiveSizeInBits();
   (void)MaskTySize;
   assert((MaskTySize == 128 || MaskTySize == 256 || MaskTySize == 512) &&
          "Unexpected vector size.");
   assert((ElSize == 8 || ElSize == 16 || ElSize == 32 || ElSize == 64) &&
          "Unexpected vector element size.");
 
   // The shuffle mask requires elements the same size as the target.
-  SmallBitVector UndefElts;
+  APInt UndefElts;
   SmallVector<uint64_t, 8> RawMask;
   if (!extractConstantMask(C, ElSize, UndefElts, RawMask))
     return;
 
   unsigned NumElts = RawMask.size();
 
   for (unsigned i = 0; i != NumElts; ++i) {
     if (UndefElts[i]) {
       ShuffleMask.push_back(SM_SentinelUndef);
       continue;
     }
     int Index = RawMask[i] & (NumElts*2 - 1);
     ShuffleMask.push_back(Index);
   }
 }
 } // llvm namespace