[SmallPtrSet] Optimize find/erase

[MaskRay]

Port #100517 for DenseMap.

[llvmbot]

@llvm/pr-subscribers-llvm-support

@llvm/pr-subscribers-llvm-adt

Author: Fangrui Song (MaskRay)

Changes

Port #100517 for DenseMap.

---

Full diff: https://github.com/llvm/llvm-project/pull/104740.diff

2 Files Affected:

• (modified) llvm/include/llvm/ADT/SmallPtrSet.h (+4-4)
• (modified) llvm/lib/Support/SmallPtrSet.cpp (+19-1)

diff --git a/llvm/include/llvm/ADT/SmallPtrSet.h b/llvm/include/llvm/ADT/SmallPtrSet.h
index e668df08c3e636..1fc2318342ae78 100644
--- a/llvm/include/llvm/ADT/SmallPtrSet.h
+++ b/llvm/include/llvm/ADT/SmallPtrSet.h
@@ -185,8 +185,8 @@ class SmallPtrSetImplBase : public DebugEpochBase {
       return false;
     }
 
-    auto *Bucket = FindBucketFor(Ptr);
-    if (*Bucket != Ptr)
+    auto *Bucket = doFind(Ptr);
+    if (!Bucket)
       return false;
 
     *const_cast<const void **>(Bucket) = getTombstoneMarker();
@@ -211,8 +211,7 @@ class SmallPtrSetImplBase : public DebugEpochBase {
     }
 
     // Big set case.
-    auto *Bucket = FindBucketFor(Ptr);
-    if (*Bucket == Ptr)
+    if (auto *Bucket = doFind(Ptr))
       return Bucket;
     return EndPointer();
   }
@@ -222,6 +221,7 @@ class SmallPtrSetImplBase : public DebugEpochBase {
 private:
   std::pair<const void *const *, bool> insert_imp_big(const void *Ptr);
 
+  const void *const *doFind(const void *Ptr) const;
   const void * const *FindBucketFor(const void *Ptr) const;
   void shrink_and_clear();
 
diff --git a/llvm/lib/Support/SmallPtrSet.cpp b/llvm/lib/Support/SmallPtrSet.cpp
index cbb87ea8717cfc..1bec911f39b13e 100644
--- a/llvm/lib/Support/SmallPtrSet.cpp
+++ b/llvm/lib/Support/SmallPtrSet.cpp
@@ -62,7 +62,25 @@ SmallPtrSetImplBase::insert_imp_big(const void *Ptr) {
   return std::make_pair(Bucket, true);
 }
 
-const void * const *SmallPtrSetImplBase::FindBucketFor(const void *Ptr) const {
+const void *const *SmallPtrSetImplBase::doFind(const void *Ptr) const {
+  unsigned BucketNo =
+      DenseMapInfo<void *>::getHashValue(Ptr) & (CurArraySize - 1);
+  unsigned ProbeAmt = 1;
+  while (true) {
+    const void *const *Bucket = CurArray + BucketNo;
+    if (LLVM_LIKELY(*Bucket == Ptr))
+      return Bucket;
+    if (LLVM_LIKELY(*Bucket == getEmptyMarker()))
+      return nullptr;
+
+    // Otherwise, it's a hash collision or a tombstone, continue quadratic
+    // probing.
+    BucketNo += ProbeAmt++;
+    BucketNo &= CurArraySize - 1;
+  }
+}
+
+const void *const *SmallPtrSetImplBase::FindBucketFor(const void *Ptr) const {
   unsigned Bucket = DenseMapInfo<void *>::getHashValue(Ptr) & (CurArraySize-1);
   unsigned ArraySize = CurArraySize;
   unsigned ProbeAmt = 1;

[MaskRay]

https://llvm-compile-time-tracker.com/compare.php?from=c64ce8bf283120fd145a57d0e61f9697f719139d&to=0de5f8109403426ad82a0a48c2a43eaec15ab476&stat=instructions%3Au

stage1-O3:

Benchmark	Old	New
kimwitu++	42462M	42452M (-0.02%)
sqlite3	38313M	38264M (-0.13%)
consumer-typeset	34798M	34768M (-0.09%)
Bullet	103581M	103529M (-0.05%)
tramp3d-v4	86502M	86370M (-0.15%)
mafft	36182M	36158M (-0.07%)
ClamAV	55277M	55214M (-0.11%)
lencod	66504M	66467M (-0.05%)
SPASS	46805M	46733M (-0.15%)
7zip	212496M	212337M (-0.07%)
geomean	60628M	60573M (-0.09%)

stage2-O3:

Benchmark	Old	New
kimwitu++	37790M	37852M (+0.16%)
sqlite3	33990M	33946M (-0.13%)
consumer-typeset	31188M	31149M (-0.12%)
Bullet	91393M	91366M (-0.03%)
tramp3d-v4	76268M	76091M (-0.23%)
mafft	32014M	31974M (-0.12%)
ClamAV	49141M	49066M (-0.15%)
lencod	59224M	59167M (-0.10%)
SPASS	41660M	41613M (-0.11%)
7zip	188382M	188302M (-0.04%)
geomean	53827M	53780M (-0.09%)

[aengelke]

LGTM

[kazutakahirata]

LGTM. Thanks!