LLVM笔记(19) - ADT介绍(二) DenseMap

这周定位一个与DenseMap有关的问题, 正好趁机过一下它的实现.
DenseMap(稠密映射)是LLVM自定义的关系类的容器, 我们可以像std::map一样使用它, 但要注意两者之间稍稍有些区别.

概述

类似于unorderd_map, DenseMap也是通过哈希实现, 区别在于DenseMap使用单一分配方式一次性分配所有bucket的内存(用来存储key-value pair), 因此具有较好的局部性.
由于其预分配内存的原因, 扩增哈希表会导致重新分配内存与拷贝(否则无法维护连续内存的特性), 因此增删元素会导致迭代器失效(这点与std::map不同). 另外预分配也导致额外内存开销(实际插入元素个数少于分配个数).
最后对于特殊的Key类型需要特化对应类型的DenseMapInfo结构, 提供包括获取空键/获取tombstone/获取哈希结果/比较Key值四个接口.
官方文档里有对它的介绍.

实现

LLVM提供了两类稠密映射, DenseMap与SmallDenseMap, 两者区别在于使用场景不同导致的内存分配方式不同, 因此我们重点关注DenseMap(defined in include/llvm/ADT/DenseMap.h).

template <typename DerivedT, typename KeyT, typename ValueT, typename KeyInfoT, typename BucketT>
class DenseMapBase : public DebugEpochBase {
protected:
  DenseMapBase() = default;
};

template <typename KeyT, typename ValueT,
          typename KeyInfoT = DenseMapInfo<KeyT>,
          typename BucketT = llvm::detail::DenseMapPair<KeyT, ValueT>>
class DenseMap : public DenseMapBase<DenseMap<KeyT, ValueT, KeyInfoT, BucketT>,
                                     KeyT, ValueT, KeyInfoT, BucketT> {
  friend class DenseMapBase<DenseMap, KeyT, ValueT, KeyInfoT, BucketT>;

  using BaseT = DenseMapBase<DenseMap, KeyT, ValueT, KeyInfoT, BucketT>;

  BucketT *Buckets;
  unsigned NumEntries;
  unsigned NumTombstones;
  unsigned NumBuckets;

public:
  explicit DenseMap(unsigned InitialReserve = 0) { init(InitialReserve); }

  ~DenseMap() {
    this->destroyAll();
    deallocate_buffer(Buckets, sizeof(BucketT) * NumBuckets, alignof(BucketT));
  }

  void init(unsigned InitNumEntries) {
    auto InitBuckets = BaseT::getMinBucketToReserveForEntries(InitNumEntries);
    if (allocateBuckets(InitBuckets)) {
      this->BaseT::initEmpty();
    } else {
      NumEntries = 0;
      NumTombstones = 0;
    }
  }

  bool allocateBuckets(unsigned Num) {
    NumBuckets = Num;
    if (NumBuckets == 0) {
      Buckets = nullptr;
      return false;
    }

    Buckets = static_cast<BucketT *>(allocate_buffer(sizeof(BucketT) * NumBuckets, alignof(BucketT)));
    return true;
  }
};

DenseMapBase是DenseMap与SmallDenseMap的基类, 其仅定义了通用的Map操作, 本身并不存储数据, 继承类实现了容器功能并向父类声明友元以供父类访问, 借此实现算法与容器实现的分离.
DenseMap包含四个成员, 其中指针Buckets指向一块存放类型为BucketT的数组地址, NumBuckets即当前数组个数, NumEntries即已使用的条目个数, NumTombstones表示被废弃且未重新映射的条目个数.
这里模板参数BucketT允许用户自定义Bucket类型, 只要提供了对应获取Key与Value的接口即可(i.e. 可以用PointerIntPair). 默认使用DenseMapPair, 它是对std::pair的封装.

template <typename KeyT, typename ValueT>
struct DenseMapPair : public std::pair<KeyT, ValueT> {
  using std::pair<KeyT, ValueT>::pair;

  KeyT &getFirst() { return std::pair<KeyT, ValueT>::first; }
  const KeyT &getFirst() const { return std::pair<KeyT, ValueT>::first; }
  ValueT &getSecond() { return std::pair<KeyT, ValueT>::second; }
  const ValueT &getSecond() const { return std::pair<KeyT, ValueT>::second; }
};

回到DenseMap的构造函数, 其调用DenseMap::init()接受一个(默认为0的)参数InitialReserve作为DenseMap起始元素个数并申请对应元素个数的空间, 申请成功再调用BaseT::initEmpty()将所有的Key初始化为空.

void DenseMapBase<>::initEmpty() {
  setNumEntries(0);
  setNumTombstones(0);

  const KeyT EmptyKey = getEmptyKey();
  for (BucketT *B = getBuckets(), *E = getBucketsEnd(); B != E; ++B)
    ::new (&B->getFirst()) KeyT(EmptyKey);
}

static const KeyT DenseMapBase<>::getEmptyKey() {
  static_assert(std::is_base_of<DenseMapBase, DerivedT>::value, "Must pass the derived type to this template!");
  return KeyInfoT::getEmptyKey();
}

注意initEmpty()里获取空键方式是调用KeyInfoT::getEmptyKey(), 而KeyInfoT作为模板参数指定了哈希的方式, LLVM提供了默认的哈希方式是DenseMapInfo(defined in include/llvm/ADT/DenseMapInfo.h).

template<typename T> struct DenseMapInfo {
  //static inline T getEmptyKey();
  //static inline T getTombstoneKey();
  //static unsigned getHashValue(const T &Val);
  //static bool isEqual(const T &LHS, const T &RHS);
};

template<typename T> struct DenseMapInfo<T*> {
  static constexpr uintptr_t Log2MaxAlign = 12;
  static inline T* getEmptyKey() {
    uintptr_t Val = static_cast<uintptr_t>(-1);
    Val <<= Log2MaxAlign;
    return reinterpret_cast<T*>(Val);
  }
  static inline T* getTombstoneKey() {
    uintptr_t Val = static_cast<uintptr_t>(-2);
    Val <<= Log2MaxAlign;
    return reinterpret_cast<T*>(Val);
  }
  static unsigned getHashValue(const T *PtrVal) {
    return (unsigned((uintptr_t)PtrVal) >> 4) ^
           (unsigned((uintptr_t)PtrVal) >> 9);
  }
  static bool isEqual(const T *LHS, const T *RHS) { return LHS == RHS; }
};

模板类DenseMapInfo需要实现四个静态方法:

  1. getEmptyKey(): 获取空键, 默认初始化时会将所有的元素设置为空键以表明该bucket可用.
  2. getTombstoneKey(): 获取tombstone Key. 由于存在可能的哈希冲突, 一个哈希值对应多个Key的情况, 此时删除中间的元素时不能将键值设为空键, 否则查找无法继续, 因此使用特殊标记(tombstone).
  3. getHashValue(): 对输入Key做哈希运算.
  4. isEqual(): 由于存在可能的哈希冲突, 需要一个方法判断给定的两个Key是否为同一Key.

注意DenseMapInfo.h中已经特化了包含指针与基础类型在内的许多数据结构的DenseMapInfo, 上面列举了指针的特化模板作为参考. 如果你的键值是自定义类型, 那么需要实现对应的DenseMapInfo.

查询与增删

DenseMap使用与std::map一致的接口, 但实现上稍有不同. 先来看下find().

template<typename LookupKeyT>
bool DenseMapBase<>::LookupBucketFor(const LookupKeyT &Val, const BucketT *&FoundBucket) const {
  const BucketT *BucketsPtr = getBuckets();
  const unsigned NumBuckets = getNumBuckets();

  if (NumBuckets == 0) {
    FoundBucket = nullptr;
    return false;
  }

  const BucketT *FoundTombstone = nullptr;
  const KeyT EmptyKey = getEmptyKey();
  const KeyT TombstoneKey = getTombstoneKey();
  assert(!KeyInfoT::isEqual(Val, EmptyKey) && !KeyInfoT::isEqual(Val, TombstoneKey) &&
         "Empty/Tombstone value shouldn't be inserted into map!");

  unsigned BucketNo = getHashValue(Val) & (NumBuckets-1);
  unsigned ProbeAmt = 1;
  while (true) {
    const BucketT *ThisBucket = BucketsPtr + BucketNo;
    if (LLVM_LIKELY(KeyInfoT::isEqual(Val, ThisBucket->getFirst()))) {
      FoundBucket = ThisBucket;
      return true;
    }

    if (LLVM_LIKELY(KeyInfoT::isEqual(ThisBucket->getFirst(), EmptyKey))) {
      FoundBucket = FoundTombstone ? FoundTombstone : ThisBucket;
      return false;
    }

    if (KeyInfoT::isEqual(ThisBucket->getFirst(), TombstoneKey) && !FoundTombstone)
      FoundTombstone = ThisBucket;

    BucketNo += ProbeAmt++;
    BucketNo &= (NumBuckets-1);
  }
}

iterator DenseMapBase<>::find(const_arg_type_t<KeyT> Val) {
  BucketT *TheBucket;
  if (LookupBucketFor(Val, TheBucket))
    return makeIterator(TheBucket,
                        shouldReverseIterate<KeyT>() ? getBuckets() : getBucketsEnd(),
                        *this, true);
  return end();
}

LookupBucketFor()是DenseMap底层的查找接口实现, 其返回给定键值的bucket指针(通过FoundBucket), 如果该键值存在则返回true, 否则返回false.
当键值不存在时LookupBucketFor()会记录下第一个找到的tombstone或空键(方便插入时选择LRU的索引), 但是注意只有仅当找到空键时才会返回(意味着没有更多的冲突元素), 如果整个表里已完全被tombstone包含会导致此处死循环(需要在插入时保证避免出现这类情况).
DenseMapBase额外提供了find_as()同样提供查询功能, 其与find()的区别在于前者支持模板参数而不使用键值的类型, 这在构造键值元素比较困难(但比较操作比较简单)的场景下可以提供一个高效的查询方式.
举个例子: 一个Value中通常只支持一个PoisoningVH, 那对于以PoisoningVH为键值的DenseMap无法构造两个相同的键值, 此时可以使用其指向的Value本身查询. 注意在使用find_as()时DenseMapInfo需要额外实现对应模板类型的getHashValue()与isEqual()接口.

template<class LookupKeyT>
iterator DenseMapBase<>::find_as(const LookupKeyT &Val) {
  BucketT *TheBucket;
  if (LookupBucketFor(Val, TheBucket))
    return makeIterator(TheBucket,
                        shouldReverseIterate<KeyT>() ? getBuckets() : getBucketsEnd(),
                        *this, true);
  return end();
}

再来看下插入接口, 若该键值对不存在则就地构造该值, insert()返回true, 否则返回false且不更新该值. InsertIntoBucket()尝试寻找一个空bucket, 然后将KeyT与ValueT移动赋值给bucket. 这里查找bucket的步骤如下:

  1. 如果当前元素个数超过容量的3/4则将容器扩大一倍, 然后重新从扩容后的容器中查找一个bucket.
  2. 否则如果空键值的元素个数(即不包含使用中的元素与tombstone元素)小于总容量的1/8则将tombstone元素恢复为空键(防止上文查找中死循环).
  3. 自增使用计数, 如果插入位置是tombstone还要自减tombstone的使用计数.
template <typename LookupKeyT>
BucketT *DenseMapBase<>::InsertIntoBucketImpl(const KeyT &Key, const LookupKeyT &Lookup, BucketT *TheBucket) {
  incrementEpoch();

  unsigned NewNumEntries = getNumEntries() + 1;
  unsigned NumBuckets = getNumBuckets();
  if (LLVM_UNLIKELY(NewNumEntries * 4 >= NumBuckets * 3)) {
    this->grow(NumBuckets * 2);
    LookupBucketFor(Lookup, TheBucket);
    NumBuckets = getNumBuckets();
  } else if (LLVM_UNLIKELY(NumBuckets-(NewNumEntries+getNumTombstones()) <= NumBuckets/8)) {
    this->grow(NumBuckets);
    LookupBucketFor(Lookup, TheBucket);
  }
  assert(TheBucket);

  incrementNumEntries();

  const KeyT EmptyKey = getEmptyKey();
  if (!KeyInfoT::isEqual(TheBucket->getFirst(), EmptyKey))
    decrementNumTombstones();

  return TheBucket;
}

template <typename KeyArg, typename... ValueArgs>
BucketT *DenseMapBase<>::InsertIntoBucket(BucketT *TheBucket, KeyArg &&Key, ValueArgs &&... Values) {
  TheBucket = InsertIntoBucketImpl(Key, Key, TheBucket);

  TheBucket->getFirst() = std::forward<KeyArg>(Key);
  ::new (&TheBucket->getSecond()) ValueT(std::forward<ValueArgs>(Values)...);
  return TheBucket;
}

template <typename... Ts>
std::pair<iterator, bool> DenseMapBase<>::try_emplace(KeyT &&Key, Ts &&... Args) {
  BucketT *TheBucket;
  if (LookupBucketFor(Key, TheBucket))
    return std::make_pair(makeIterator(TheBucket,
                                       shouldReverseIterate<KeyT>() ? getBuckets() : getBucketsEnd(),
                                       *this, true),
                          false);

  TheBucket = InsertIntoBucket(TheBucket, std::move(Key), std::forward<Ts>(Args)...);
  return std::make_pair(makeIterator(TheBucket,
                                     shouldReverseIterate<KeyT>() ? getBuckets() : getBucketsEnd(),
                                     *this, true),
                        true);
}

std::pair<iterator, bool> DenseMapBase<>::insert(const std::pair<KeyT, ValueT> &KV) {
  return try_emplace(KV.first, KV.second);
}

在空余元素个数不足时DenseMap会调用继承类的grow()接口为容器扩容. 不同与std::map, DenseMap采用重新分配内存并拷贝旧元素的方式保持数据的连续性. 注意:

  1. grow()要求容器最少个数保持64个.
  2. 拷贝旧元素时仅仅对使用中的键值元素做拷贝, 因此即使grow()不增长实际大小, 也会释放tombstone增加可用空间.
void DenseMapBase<>::moveFromOldBuckets(BucketT *OldBucketsBegin, BucketT *OldBucketsEnd) {
  initEmpty();

  const KeyT EmptyKey = getEmptyKey();
  const KeyT TombstoneKey = getTombstoneKey();
  for (BucketT *B = OldBucketsBegin, *E = OldBucketsEnd; B != E; ++B) {
    if (!KeyInfoT::isEqual(B->getFirst(), EmptyKey) &&
        !KeyInfoT::isEqual(B->getFirst(), TombstoneKey)) {
      BucketT *DestBucket;
      bool FoundVal = LookupBucketFor(B->getFirst(), DestBucket);
      assert(!FoundVal && "Key already in new map?");
      DestBucket->getFirst() = std::move(B->getFirst());
      ::new (&DestBucket->getSecond()) ValueT(std::move(B->getSecond()));
      incrementNumEntries();

      B->getSecond().~ValueT();
    }
    B->getFirst().~KeyT();
  }
}

void DenseMap::grow(unsigned AtLeast) {
  unsigned OldNumBuckets = NumBuckets;
  BucketT *OldBuckets = Buckets;

  allocateBuckets(std::max<unsigned>(64, static_cast<unsigned>(NextPowerOf2(AtLeast-1))));
  assert(Buckets);
  if (!OldBuckets) {
    this->BaseT::initEmpty();
    return;
  }

  this->moveFromOldBuckets(OldBuckets, OldBuckets+OldNumBuckets);
  deallocate_buffer(OldBuckets, sizeof(BucketT) * OldNumBuckets, alignof(BucketT));
}

迭代器

与其它LLVM自定义容器一样, DenseMap也不支持标准的迭代器, 因此需要实现自定义的迭代器.
DenseMapIterator(defined in include/llvm/ADT/DenseMap.h)包含两个成员, Ptr指向当前bucket, End指向容器结束地址.
迭代器通过调用AdvancePastEmptyBuckets()/RetreatPastEmptyBuckets()遍历所有元素.

template <typename KeyT, typename ValueT, typename KeyInfoT, typename Bucket, bool IsConst>
class DenseMapIterator : DebugEpochBase::HandleBase {
  friend class DenseMapIterator<KeyT, ValueT, KeyInfoT, Bucket, true>;
  friend class DenseMapIterator<KeyT, ValueT, KeyInfoT, Bucket, false>;

public:
  using value_type = typename std::conditional<IsConst, const Bucket, Bucket>::type;
  using pointer = value_type *;

private:
  pointer Ptr = nullptr;
  pointer End = nullptr;

public:
  DenseMapIterator(pointer Pos, pointer E, const DebugEpochBase &Epoch, bool NoAdvance = false)
      : DebugEpochBase::HandleBase(&Epoch), Ptr(Pos), End(E) {
    assert(isHandleInSync() && "invalid construction!");

    if (NoAdvance) return;
    if (shouldReverseIterate<KeyT>()) {
      RetreatPastEmptyBuckets();
      return;
    }
    AdvancePastEmptyBuckets();
  }

private:
  void AdvancePastEmptyBuckets() {
    assert(Ptr <= End);
    const KeyT Empty = KeyInfoT::getEmptyKey();
    const KeyT Tombstone = KeyInfoT::getTombstoneKey();

    while (Ptr != End && (KeyInfoT::isEqual(Ptr->getFirst(), Empty) ||
                          KeyInfoT::isEqual(Ptr->getFirst(), Tombstone)))
      ++Ptr;
  }

  void RetreatPastEmptyBuckets() {
    assert(Ptr >= End);
    const KeyT Empty = KeyInfoT::getEmptyKey();
    const KeyT Tombstone = KeyInfoT::getTombstoneKey();

    while (Ptr != End && (KeyInfoT::isEqual(Ptr[-1].getFirst(), Empty) ||
                          KeyInfoT::isEqual(Ptr[-1].getFirst(), Tombstone)))
      --Ptr;
  }
};

DenseMapIterator向DenseMapBase声明了友元, 因此DenseMapBase可以通过makeIterator()构造迭代器.

iterator DenseMapBase<>::makeIterator(BucketT *P, BucketT *E,
                                      DebugEpochBase &Epoch, bool NoAdvance=false) {
  if (shouldReverseIterate<KeyT>()) {
    BucketT *B = P == getBucketsEnd() ? getBuckets() : P + 1;
    return iterator(B, E, Epoch, NoAdvance);
  }
  return iterator(P, E, Epoch, NoAdvance);
}

inline DenseMapBase<>::iterator begin() {
  if (empty())
    return end();
  if (shouldReverseIterate<KeyT>())
    return makeIterator(getBucketsEnd() - 1, getBuckets(), *this);
  return makeIterator(getBuckets(), getBucketsEnd(), *this);
}

空间优化的DenseMap

注意到DenseMap最小要求分配64个元素, 如果map中存储元素较少造成很大的浪费, 所以LLVM又定义了一个针对少量元素的SmallDenseMap.
SmallDenseMap与DenseMap的唯一区别是前者假定元素个数通常小于一个给定值, 因此默认初始化时会静态初始化一个bucket数组. 在分配元素超过了限制后会退化为DenseMap.

template <typename KeyT, typename ValueT, unsigned InlineBuckets = 4,
          typename KeyInfoT = DenseMapInfo<KeyT>,
          typename BucketT = llvm::detail::DenseMapPair<KeyT, ValueT>>
class SmallDenseMap : public DenseMapBase<
                          SmallDenseMap<KeyT, ValueT, InlineBuckets, KeyInfoT, BucketT>, KeyT, ValueT, KeyInfoT, BucketT> {
  friend class DenseMapBase<SmallDenseMap, KeyT, ValueT, KeyInfoT, BucketT>;

  using BaseT = DenseMapBase<SmallDenseMap, KeyT, ValueT, KeyInfoT, BucketT>;

  unsigned Small : 1;
  unsigned NumEntries : 31;
  unsigned NumTombstones;

  struct LargeRep {
    BucketT *Buckets;
    unsigned NumBuckets;
  };

  AlignedCharArrayUnion<BucketT[InlineBuckets], LargeRep> storage;
};

SmallDenseMap与DenseMap基本类似, 区别在于:

  1. SmallDenseMap定义一个union(AlignedCharArrayUnion), 保存了InlineBuckets个BucketT元素的数组, 或是一个LargeRep结构. 前者是静态构造的bucket数组, 后者等同于DenseMap中动态分配的bucket数组.
  2. SmallDenseMap定义一个Small标记指示如何理解storage的类型, 为true代表此时使用静态数组, 否则为动态分配数组.

小结

  1. DenseMap通过预分配方式申请内存, 在插入元素时在已分配的地址上构造元素. 预分配的方式是申请一块连续内存, 当需要扩容时会重新申请内存并发生内容拷贝.
  2. 增删元素会导致DenseMap扩容/缩减, 进而引起迭代器失效. 另外扩容时会删除旧缓存中的元素(析构Key与Value), 因此需要注意Key与Value成员的所有权问题.
  3. 使用自定义数据类型做Key时需要注意是否实现了对应的DenseMapInfo类, 如果不想关心这些琐事, 保险的做法是指针用对象指针做Key.
  4. SmallDenseMap与DenseMap的主要区别是默认预分配内存方式不同, 前者假定容器通常不会插入超过给定个数的元素, 因此使用静态数组. 当超过给定大小的元素后会退化为DenseMap一样的实现, 因此在元素个数可估算时通常使用SmallDenseMap更为高效.
posted @ 2021-01-03 20:16  Five100Miles  阅读(3692)  评论(0编辑  收藏  举报