[ADT] Fix SmallDenseMap assertion with large InlineBuckets

[nikic]

Summary

Fixes issue encountered in D56362, where I tried to use an SmallSetVector<Instruction*, 128> with an excessively large number of inline elements. This triggers an "Must allocate more buckets than are inline" assertion inside allocateBuckets() under certain usage patterns.

The issue is as follows: The grow() method is used either to grow the map, or to rehash it and remove tombstones. The latter is done if the fraction of empty (non-used, non-tombstone) elements is below 1/8. In this case grow() is invoked with the current number of buckets.

This is currently incorrectly handled for dense maps using the small rep. The current implementation will switch them over to the large rep, which violates the invariant that the large rep is only used if there are more than InlineBuckets buckets.

This patch fixes the issue by staying in the small rep and only moving the buckets. Alternatively, if we do want to switch to the large rep in this case, would be to relax the assertion in allocateBuckets().

[nikic]

I should probably mention that this also changes the behavior for smaller inline buckets. If say we are in the small rep with InlineBuckets=16 and call grow(16), then before this change it would switch it to the large rep with 64 buckets, while after this change it will stay in the small rep with 16 buckets. If we want to keep it this way, then an alternative fix here would be to simply adjust the assertion and allow creating a large rep with InlineBuckets buckets.

[nikic]

Ping

[nikic]

ping

[atrick]

Can this problem be fixed without losing the invariant that memory is never allocated until NumBuckets > InlineBuckets, regardless of the load? I don't like the removal of the fast path if (AtLeast < InlineBuckets) and the load doesn't matter at all in small mode!
e.g. If I ask to reserve 100 out of 128 entries, it should not attempt to copy all the entries into temporary storage.

[nikic]

@atrick That invariant doesn't exist currently. If a SmallDenseMap with less than InlineBuckets elements is compacted, it will change to the large rep (with an allocation). After this change it will instead stay in the small rep (without an allocation).

The short-circuit condition AtLeast < InlineBuckets is dead code right now, I think, because NumBuckets is >= InlineBuckets and AtLeast is either NumBuckets or NumBuckets*2.

[atrick]

Hmm... I guess the intent of the original design is not to overload the table even if it's in small mode. I'm not sure that makes sense, but I'll accept it as status quo. Anyway, it does sort of make sense that we can compress the small table by removing tombstones.

Now that I fully understand what's happening, this looks like a good fix!

llvm/include/llvm/ADT/DenseMap.h

void init(unsigned InitBuckets) {
     if (InitBuckets > InlineBuckets) {
       Small = false;
       new (getLargeRep()) LargeRep(allocateBuckets(InitBuckets));
     }
     this->BaseT::initEmpty();
   }
 
   void grow(unsigned AtLeast) {
-    if (AtLeast >= InlineBuckets)
+    if (AtLeast > InlineBuckets)
       AtLeast = std::max<unsigned>(64, NextPowerOf2(AtLeast-1));
 
     if (Small) {
-      if (AtLeast < InlineBuckets)
-        return; // Nothing to do.
-
       // First move the inline buckets into a temporary storage.
       AlignedCharArrayUnion<BucketT[InlineBuckets]> TmpStorage;
       BucketT *TmpBegin = reinterpret_cast<BucketT *>(TmpStorage.buffer);
       BucketT *TmpEnd = TmpBegin;
 
       // Loop over the buckets, moving non-empty, non-tombstones into the
       // temporary storage. Have the loop move the TmpEnd forward as it goes.
       const KeyT EmptyKey = this->getEmptyKey();
       const KeyT TombstoneKey = this->getTombstoneKey();
       for (BucketT *P = getBuckets(), *E = P + InlineBuckets; P != E; ++P) {
         if (!KeyInfoT::isEqual(P->getFirst(), EmptyKey) &&
             !KeyInfoT::isEqual(P->getFirst(), TombstoneKey)) {
           assert(size_t(TmpEnd - TmpBegin) < InlineBuckets &&
                  "Too many inline buckets!");
           ::new (&TmpEnd->getFirst()) KeyT(std::move(P->getFirst()));
           ::new (&TmpEnd->getSecond()) ValueT(std::move(P->getSecond()));
           ++TmpEnd;
           P->getSecond().~ValueT();
         }
         P->getFirst().~KeyT();
       }
 
-      // Now make this map use the large rep, and move all the entries back
-      // into it.
+      // AtLeast == InlineBuckets can happen if there are many tombstones,
+      // and grow() is used to remove them. Usually we always switch to the
+      // large rep here.
+      if (AtLeast > InlineBuckets) {
         Small = false;
         new (getLargeRep()) LargeRep(allocateBuckets(AtLeast));
+      }
       this->moveFromOldBuckets(TmpBegin, TmpEnd);
       return;
     }
 
     LargeRep OldRep = std::move(*getLargeRep());
     getLargeRep()->~LargeRep();
     if (AtLeast <= InlineBuckets) {
       Small = true;

llvm/unittests/ADT/DenseMapTest.cpp

TEST(DenseMapCustomTest, SmallDenseMapGrowTest) {
 
   // Try to find an element which doesn't exist.  There was a bug in
   // SmallDenseMap which led to a map with num elements == small capacity not
   // having an empty bucket any more.  Finding an element not in the map would
   // therefore never terminate.
   EXPECT_TRUE(map.find(32) == map.end());
 }
 
+TEST(DenseMapCustomTest, LargeSmallDenseMapCompaction) {
+  SmallDenseMap<unsigned, unsigned, 128, ContiguousDenseMapInfo> map;
+  // Fill to < 3/4 load.
+  for (unsigned i = 0; i < 95; ++i)
+    map[i] = i;
+  // And erase, leaving behind tombstones.
+  for (unsigned i = 0; i < 95; ++i)
+    map.erase(i);
+  // Fill further, so that less than 1/8 are empty, but still below 3/4 load.
+  for (unsigned i = 95; i < 128; ++i)
+    map[i] = i;
+
+  EXPECT_EQ(33u, map.size());
+  // Similar to the previous test, check for a non-existing element, as an
+  // indirect check that tombstones have been removed.
+  EXPECT_TRUE(map.find(0) == map.end());
+}
+
 TEST(DenseMapCustomTest, TryEmplaceTest) {
   DenseMap<int, std::unique_ptr<int>> Map;
   std::unique_ptr<int> P(new int(2));
   auto Try1 = Map.try_emplace(0, new int(1));
   EXPECT_TRUE(Try1.second);
   auto Try2 = Map.try_emplace(0, std::move(P));
   EXPECT_FALSE(Try2.second);
   EXPECT_EQ(Try1.first, Try2.first);