[llvm] [SmallPtrSet] Optimize Grow and inline FindBucketFor. NFC (PR #206304)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-support

Author: Fangrui Song (MaskRay)

<details>
<summary>Changes</summary>

During rehashing every key is known absent, so skip the equal-key comparison.
DenseMap got a similar optimization in #<!-- -->200595

FindBucketFor is then only used by insert_imp_big; inline its probe loop
there. The loop takes 7 inst instead of 9 in the clang-generated
assembly after removing the two `LLVM_LIKELY`.


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


2 Files Affected:

- (modified) llvm/include/llvm/ADT/SmallPtrSet.h (-1) 
- (modified) llvm/lib/Support/SmallPtrSet.cpp (+22-24) 


``````````diff
diff --git a/llvm/include/llvm/ADT/SmallPtrSet.h b/llvm/include/llvm/ADT/SmallPtrSet.h
index d7b4bbf6bd64f..d685279d3c31b 100644
--- a/llvm/include/llvm/ADT/SmallPtrSet.h
+++ b/llvm/include/llvm/ADT/SmallPtrSet.h
@@ -245,7 +245,6 @@ class SmallPtrSetImplBase : public DebugEpochBase {
   LLVM_ABI std::pair<const void *const *, bool> insert_imp_big(const void *Ptr);
 
   LLVM_ABI const void *const *doFind(const void *Ptr) const;
-  const void *const *FindBucketFor(const void *Ptr) const;
   LLVM_ABI void shrink_and_clear();
 
 protected:
diff --git a/llvm/lib/Support/SmallPtrSet.cpp b/llvm/lib/Support/SmallPtrSet.cpp
index d8e66b9c89c74..265b8eba4076d 100644
--- a/llvm/lib/Support/SmallPtrSet.cpp
+++ b/llvm/lib/Support/SmallPtrSet.cpp
@@ -44,16 +44,21 @@ SmallPtrSetImplBase::insert_imp_big(const void *Ptr) {
     Grow(CurArraySize < 64 ? 128 : CurArraySize * 2);
   }
 
-  // Okay, we know we have space.  Find a hash bucket.
-  const void **Bucket = const_cast<const void**>(FindBucketFor(Ptr));
-  if (*Bucket == Ptr)
-    return {Bucket, false}; // Already inserted, good.
+  // Find the first empty bucket or Ptr itself on the probe chain.
+  unsigned Mask = CurArraySize - 1;
+  unsigned I = DenseMapInfo<void *>::getHashValue(Ptr) & Mask;
+  const void **Array = CurArray;
+  while (Array[I] != getEmptyMarker()) {
+    if (Array[I] == Ptr)
+      return {Array + I, false};
+    I = (I + 1) & Mask;
+  }
 
-  // Otherwise, insert it.
+  // Insert into the empty bucket.
   ++NumEntries;
-  *Bucket = Ptr;
+  Array[I] = Ptr;
   incrementEpoch();
-  return {Bucket, true};
+  return {Array + I, true};
 }
 
 const void *const *SmallPtrSetImplBase::doFind(const void *Ptr) const {
@@ -69,19 +74,6 @@ const void *const *SmallPtrSetImplBase::doFind(const void *Ptr) const {
   }
 }
 
-const void *const *SmallPtrSetImplBase::FindBucketFor(const void *Ptr) const {
-  unsigned Mask = CurArraySize - 1;
-  unsigned Bucket = DenseMapInfo<void *>::getHashValue(Ptr) & Mask;
-  const void *const *Array = CurArray;
-  while (true) {
-    if (LLVM_LIKELY(Array[Bucket] == getEmptyMarker()))
-      return Array + Bucket;
-    if (LLVM_LIKELY(Array[Bucket] == Ptr))
-      return Array + Bucket;
-    Bucket = (Bucket + 1) & Mask;
-  }
-}
-
 void SmallPtrSetImplBase::eraseFromBucket(const void **Bucket) {
   // Knuth TAOCP 6.4 Algorithm R: walk forward sliding each following entry
   // whose probe path crosses the hole.
@@ -114,10 +106,16 @@ void SmallPtrSetImplBase::Grow(unsigned NewSize) {
   memset(CurArray, -1, NewSize*sizeof(void*));
 
   // Copy over all valid entries.
-  for (const void *&Bucket : OldBuckets) {
-    // Copy over the element if it is valid.
-    if (Bucket != getEmptyMarker())
-      *const_cast<void **>(FindBucketFor(Bucket)) = const_cast<void *>(Bucket);
+  unsigned Mask = CurArraySize - 1;
+  for (const void *Ptr : OldBuckets) {
+    if (Ptr == getEmptyMarker())
+      continue;
+    // Find the first empty bucket on this key's probe chain; there is no equal
+    // key, so nothing to compare against.
+    unsigned I = DenseMapInfo<void *>::getHashValue(Ptr) & Mask;
+    while (NewBuckets[I] != getEmptyMarker())
+      I = (I + 1) & Mask;
+    NewBuckets[I] = Ptr;
   }
 
   if (!WasSmall)

``````````

</details>


https://github.com/llvm/llvm-project/pull/206304