[llvm] a3fea06 - [SmallPtrSet] Optimize find/erase

[via llvm-commits]

Author: Fangrui Song
Date: 2024-08-19T09:42:14-07:00
New Revision: a3fea0643da544b2e5d5c48f2ea942dad22d01f2

URL: https://github.com/llvm/llvm-project/commit/a3fea0643da544b2e5d5c48f2ea942dad22d01f2
DIFF: https://github.com/llvm/llvm-project/commit/a3fea0643da544b2e5d5c48f2ea942dad22d01f2.diff

LOG: [SmallPtrSet] Optimize find/erase

Port #100517 for DenseMap.

Pull Request: https://github.com/llvm/llvm-project/pull/104740

Added: 
    

Modified: 
    llvm/include/llvm/ADT/SmallPtrSet.h
    llvm/lib/Support/SmallPtrSet.cpp

Removed: 
    


################################################################################
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;