From e5da12363a550c10f0e3663f17810585162a9e9e Mon Sep 17 00:00:00 2001
From: Deyu Kong <deyukong@microsoft.com>
Date: Mon, 8 Jun 2026 13:20:01 +0800
Subject: [PATCH 1/2] perf: ANN optimizations - SwissHashSet, AVX2 distance,
 insert_fast

Port hot-path optimizations from parent repo PR 6610806 (commit 844664001):

- swiss_hash_set.h: New SSE2/SIMD-optimized Swiss table hash set for
  uint32_t keys with O(occupied_slots) clear and ~87.5% load factor.
  Replaces tsl::robin_set for inserted_into_pool_rs.
- neighbor.h: Add NeighborMinHeap::insert_fast() with AVX2-accelerated
  linear scan over a contiguous float distance array (_dist_arr) for
  faster insertion-point finding.
- index.cpp: Refactor iterate_to_fixed_point() to interleave prefetch,
  distance computation, and immediate insert; remove separate
  dist_scratch allocation; mark nodes visited and process neighbors
  under the same shared lock (eliminates redundant lock/unlock).
- pq_l2_distance.cpp: AVX2 inner loop for chunk-wise PQ centroid
  distance computation (8-wide _mm256_fmadd_ps).
- scratch.h/.cpp: Switch _inserted_into_pool_rs type to SwissHashSet
  and skip clear() on empty containers.
---
 include/neighbor.h       |  72 ++++++++-
 include/scratch.h        |   6 +-
 include/swiss_hash_set.h | 318 +++++++++++++++++++++++++++++++++++++++
 src/index.cpp            |  79 +++++-----
 src/pq_l2_distance.cpp   |  13 +-
 src/scratch.cpp          |  21 ++-
 6 files changed, 458 insertions(+), 51 deletions(-)
 create mode 100644 include/swiss_hash_set.h

diff --git a/include/neighbor.h b/include/neighbor.h
index 0f7d96a5d..b3092cb6a 100644
--- a/include/neighbor.h
+++ b/include/neighbor.h
@@ -7,6 +7,7 @@
 #include <mutex>
 #include <vector>
 #include <assert.h>
+#include <immintrin.h>
 #include <tsl/robin_map.h>
 //#include "utils.h"
 #include "color_info.h"
@@ -90,6 +91,7 @@ class NeighborVector : public NeighborVectorBase
     NeighborVector(size_t capacity)
         : _capacity(capacity)
         , _data(capacity + 1, Neighbor(std::numeric_limits<uint32_t>::max(), std::numeric_limits<float>::max()))
+        , _dist_arr(capacity + 1, std::numeric_limits<float>::max())
     {
     }
 
@@ -108,12 +110,71 @@ class NeighborVector : public NeighborVectorBase
     virtual void insert(const Neighbor& nbr, int lo, int /*kick_loc*/, int size) override
     {
         std::memmove(&_data[lo + 1], &_data[lo], (size - lo) * sizeof(Neighbor));
+        std::memmove(&_dist_arr[lo + 1], &_dist_arr[lo], (size - lo) * sizeof(float));
         _data[lo] = { nbr.id, nbr.distance };
+        _dist_arr[lo] = nbr.distance;
+    }
+
+    // SIMD lower-bound insertion (Rust-style): finds first position where
+    // distance >= target, inserts there. Caller's visited set prevents duplicates.
+    // Returns insertion index, or capacity+1 if rejected.
+    size_t insert_fast(const Neighbor& nbr, size_t size, size_t capacity)
+    {
+        float dist = nbr.distance;
+
+        if (size >= capacity && dist >= _dist_arr[size - 1])
+        {
+            if (dist > _dist_arr[size - 1])
+                return capacity + 1;
+            if (_data[size - 1].id <= nbr.id)
+                return capacity + 1;
+        }
+
+        // Lower-bound scan: find first position where _dist_arr[i] > dist,
+        // or _dist_arr[i] == dist && _data[i].id >= nbr.id
+        size_t lo = 0;
+#ifdef USE_AVX2
+        __m256 target = _mm256_set1_ps(dist);
+        for (; lo + 8 <= size; lo += 8)
+        {
+            __m256 d = _mm256_loadu_ps(&_dist_arr[lo]);
+            int ge_mask = _mm256_movemask_ps(_mm256_cmp_ps(d, target, _CMP_GE_OQ));
+            if (ge_mask)
+            {
+                lo += _tzcnt_u32(ge_mask);
+                while (lo < size && _dist_arr[lo] == dist && _data[lo].id < nbr.id)
+                    lo++;
+                goto found;
+            }
+        }
+#endif
+        for (; lo < size; lo++)
+        {
+            if (_dist_arr[lo] > dist)
+                goto found;
+            if (_dist_arr[lo] == dist && _data[lo].id >= nbr.id)
+                goto found;
+        }
+
+    found:
+        if (lo < capacity)
+        {
+            size_t move_count = (size < capacity ? size : capacity - 1) - lo;
+            if (move_count > 0)
+            {
+                std::memmove(&_data[lo + 1], &_data[lo], move_count * sizeof(Neighbor));
+                std::memmove(&_dist_arr[lo + 1], &_dist_arr[lo], move_count * sizeof(float));
+            }
+            _data[lo] = { nbr.id, nbr.distance };
+            _dist_arr[lo] = nbr.distance;
+        }
+        return lo;
     }
 
     void resize(size_t capacity)
     {
         _data.resize(capacity + 1, Neighbor(std::numeric_limits<uint32_t>::max(), std::numeric_limits<float>::max()));
+        _dist_arr.resize(capacity + 1, std::numeric_limits<float>::max());
         _capacity = capacity;
     }
 
@@ -127,8 +188,11 @@ class NeighborVector : public NeighborVectorBase
         // no need additional clear up
     }
 
+    float* distances() { return _dist_arr.data(); }
+
 private:
     std::vector<Neighbor> _data;
+    std::vector<float> _dist_arr;  // SoA distance array for SIMD scan
     size_t _capacity = 0;
 };
 
@@ -499,7 +563,13 @@ class NeighborPriorityQueue : public NeighborPriorityQueueBase
 
     virtual void insert(const Neighbor& nbr) override
     {
-        NeighborPriorityQueueBase::insert(nbr, _data);
+        size_t lo = _data.insert_fast(nbr, _size, _capacity);
+        if (lo > _capacity)
+            return;
+        if (_size < _capacity)
+            _size++;
+        if (lo < _cur)
+            _cur = lo;
     }
 
     virtual Neighbor closest_unexpanded() override
diff --git a/include/scratch.h b/include/scratch.h
index ce30f1249..2884c862a 100644
--- a/include/scratch.h
+++ b/include/scratch.h
@@ -13,6 +13,8 @@
 
 #include "aligned_file_reader.h"
 #include "abstract_scratch.h"
+
+#include "swiss_hash_set.h"
 #include "neighbor.h"
 #include "defaults.h"
 #include "concurrent_queue.h"
@@ -73,7 +75,7 @@ template <typename T> class InMemQueryScratch : public AbstractScratch<T>
     {
         return _candidate_pick_flags;
     }
-    inline tsl::robin_set<uint32_t> &inserted_into_pool_rs()
+    inline SwissHashSet &inserted_into_pool_rs()
     {
         return _inserted_into_pool_rs;
     }
@@ -133,7 +135,7 @@ template <typename T> class InMemQueryScratch : public AbstractScratch<T>
     std::vector<bool> _candidate_pick_flags;
 
     // Capacity initialized to 20L
-    tsl::robin_set<uint32_t> _inserted_into_pool_rs;
+    SwissHashSet _inserted_into_pool_rs;
 
     // Use a pointer here to allow for forward declaration of dynamic_bitset
     // in public headers to avoid making boost a dependency for clients
diff --git a/include/swiss_hash_set.h b/include/swiss_hash_set.h
new file mode 100644
index 000000000..fb079f4f9
--- /dev/null
+++ b/include/swiss_hash_set.h
@@ -0,0 +1,318 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+//
+// Swiss table-style hash set for uint32_t keys.
+// Inspired by Google's Abseil flat_hash_set / Rust's hashbrown.
+// Uses SSE2 SIMD control byte matching for fast probing.
+// clear() is O(occupied_slots) not O(capacity).
+
+#pragma once
+
+#include <cstdint>
+#include <cstring>
+#include <cassert>
+#include <algorithm>
+#include <immintrin.h>
+
+namespace diskann
+{
+
+class SwissHashSet
+{
+  public:
+    // Control byte values
+    static constexpr int8_t kEmpty = -128;   // 0b10000000
+    static constexpr int8_t kDeleted = -2;   // 0b11111110 (unused — we don't do tombstones)
+
+    SwissHashSet() : _ctrl(nullptr), _slots(nullptr), _capacity(0), _size(0), _growth_left(0)
+    {
+    }
+
+    ~SwissHashSet()
+    {
+        _free();
+    }
+
+    SwissHashSet(const SwissHashSet &) = delete;
+    SwissHashSet &operator=(const SwissHashSet &) = delete;
+
+    SwissHashSet(SwissHashSet &&other) noexcept
+        : _ctrl(other._ctrl), _slots(other._slots), _capacity(other._capacity), _size(other._size),
+          _growth_left(other._growth_left)
+    {
+        other._ctrl = nullptr;
+        other._slots = nullptr;
+        other._capacity = 0;
+        other._size = 0;
+        other._growth_left = 0;
+    }
+
+    void reserve(size_t n)
+    {
+        if (n == 0)
+            return;
+        // Target load factor ~7/8 (87.5%), so we need capacity >= n * 8/7
+        size_t needed = n + n / 7;
+        needed = _round_up_power_of_2(std::max(needed, (size_t)16));
+        if (needed > _capacity)
+        {
+            _resize(needed);
+        }
+    }
+
+    // Returns true if inserted, false if already present
+    bool insert(uint32_t key)
+    {
+        if (_capacity == 0)
+        {
+            _resize(16);
+        }
+
+        size_t hash = _hash(key);
+        size_t pos = _find_or_prepare_insert(hash, key);
+        if (_slots[pos] == key && _ctrl[pos] >= 0)
+        {
+            return false; // already exists
+        }
+
+        // Insert
+        _set_ctrl(pos, _h2(hash));
+        _slots[pos] = key;
+        _size++;
+        _growth_left--;
+
+        if (_growth_left == 0)
+        {
+            _resize(_capacity * 2);
+        }
+        return true;
+    }
+
+    bool find(uint32_t key) const
+    {
+        if (_capacity == 0)
+            return false;
+
+        size_t hash = _hash(key);
+        size_t idx = hash & (_capacity - 1);
+        size_t group_start = idx & ~(size_t)15; // align to 16-byte group
+
+        // Probe groups
+        for (size_t probe = 0; probe <= _capacity; probe += 16)
+        {
+            size_t g = (group_start + probe) & (_capacity - 1);
+            __m128i ctrl_group = _mm_loadu_si128(reinterpret_cast<const __m128i *>(_ctrl + g));
+            int8_t h2 = _h2(hash);
+
+            // Match h2
+            __m128i match = _mm_cmpeq_epi8(ctrl_group, _mm_set1_epi8(h2));
+            uint32_t mask = static_cast<uint32_t>(_mm_movemask_epi8(match));
+
+            while (mask != 0)
+            {
+                int bit = _bit_scan_forward(mask);
+                size_t candidate = (g + bit) & (_capacity - 1);
+                if (_slots[candidate] == key)
+                {
+                    return true;
+                }
+                mask &= mask - 1; // clear lowest bit
+            }
+
+            // Check for empty — if any slot in group is empty, key is absent
+            __m128i empty_match = _mm_cmpeq_epi8(ctrl_group, _mm_set1_epi8(kEmpty));
+            if (_mm_movemask_epi8(empty_match) != 0)
+            {
+                return false;
+            }
+        }
+        return false;
+    }
+
+    // O(size) clear — only walks occupied slots tracked via control bytes
+    void clear()
+    {
+        if (_size == 0)
+            return;
+        // Walk in 16-byte groups, only memset groups that have occupied slots
+        for (size_t g = 0; g < _capacity; g += 16)
+        {
+            __m128i ctrl_group = _mm_loadu_si128(reinterpret_cast<const __m128i *>(_ctrl + g));
+            __m128i empty_cmp = _mm_cmpeq_epi8(ctrl_group, _mm_set1_epi8(kEmpty));
+            int all_empty = _mm_movemask_epi8(empty_cmp);
+            if (all_empty != 0xFFFF)
+            {
+                // This group has at least one occupied slot — clear it
+                _mm_storeu_si128(reinterpret_cast<__m128i *>(_ctrl + g), _mm_set1_epi8(kEmpty));
+            }
+        }
+        // Also clear the mirror/clamp bytes at end
+        if (_capacity > 0)
+        {
+            std::memset(_ctrl + _capacity, kEmpty, 16);
+        }
+        _size = 0;
+        _growth_left = _capacity_to_growth(_capacity);
+    }
+
+    bool empty() const
+    {
+        return _size == 0;
+    }
+
+    size_t size() const
+    {
+        return _size;
+    }
+
+  private:
+    int8_t *_ctrl;      // control bytes: kEmpty or h2 (top 7 bits of hash)
+    uint32_t *_slots;   // key storage
+    size_t _capacity;   // always power of 2
+    size_t _size;
+    size_t _growth_left;
+
+    static size_t _round_up_power_of_2(size_t n)
+    {
+        n--;
+        n |= n >> 1;
+        n |= n >> 2;
+        n |= n >> 4;
+        n |= n >> 8;
+        n |= n >> 16;
+        n |= n >> 32;
+        return n + 1;
+    }
+
+    static size_t _capacity_to_growth(size_t cap)
+    {
+        // 7/8 load factor
+        return cap - cap / 8;
+    }
+
+    static size_t _hash(uint32_t key)
+    {
+        // Fast integer hash (murmurhash3 finalizer)
+        uint64_t h = key;
+        h ^= h >> 16;
+        h *= 0x85ebca6b;
+        h ^= h >> 13;
+        h *= 0xc2b2ae35;
+        h ^= h >> 16;
+        return static_cast<size_t>(h);
+    }
+
+    static int8_t _h2(size_t hash)
+    {
+        // Top 7 bits, ensuring positive (bit 7 = 0)
+        return static_cast<int8_t>(hash >> 57) & 0x7F;
+    }
+
+    static int _bit_scan_forward(uint32_t mask)
+    {
+#ifdef _MSC_VER
+        unsigned long idx;
+        _BitScanForward(&idx, mask);
+        return static_cast<int>(idx);
+#else
+        return __builtin_ctz(mask);
+#endif
+    }
+
+    void _set_ctrl(size_t pos, int8_t h2)
+    {
+        _ctrl[pos] = h2;
+        // Mirror the first 16 bytes at the end for wrap-around SIMD loads
+        if (pos < 16)
+        {
+            _ctrl[_capacity + pos] = h2;
+        }
+    }
+
+    size_t _find_or_prepare_insert(size_t hash, uint32_t key)
+    {
+        size_t idx = hash & (_capacity - 1);
+        size_t group_start = idx & ~(size_t)15;
+        int8_t h2 = _h2(hash);
+
+        for (size_t probe = 0; probe <= _capacity; probe += 16)
+        {
+            size_t g = (group_start + probe) & (_capacity - 1);
+            __m128i ctrl_group = _mm_loadu_si128(reinterpret_cast<const __m128i *>(_ctrl + g));
+
+            // Check for existing key
+            __m128i match = _mm_cmpeq_epi8(ctrl_group, _mm_set1_epi8(h2));
+            uint32_t mask = static_cast<uint32_t>(_mm_movemask_epi8(match));
+            while (mask != 0)
+            {
+                int bit = _bit_scan_forward(mask);
+                size_t candidate = (g + bit) & (_capacity - 1);
+                if (_slots[candidate] == key)
+                {
+                    return candidate; // found existing
+                }
+                mask &= mask - 1;
+            }
+
+            // Check for empty slot
+            __m128i empty_match = _mm_cmpeq_epi8(ctrl_group, _mm_set1_epi8(kEmpty));
+            uint32_t empty_mask = static_cast<uint32_t>(_mm_movemask_epi8(empty_match));
+            if (empty_mask != 0)
+            {
+                int bit = _bit_scan_forward(empty_mask);
+                return (g + bit) & (_capacity - 1);
+            }
+        }
+        // Should never reach here if load factor is maintained
+        assert(false && "SwissHashSet: table is full");
+        return 0;
+    }
+
+    void _resize(size_t new_capacity)
+    {
+        int8_t *old_ctrl = _ctrl;
+        uint32_t *old_slots = _slots;
+        size_t old_capacity = _capacity;
+
+        _capacity = new_capacity;
+        // Allocate ctrl bytes: capacity + 16 (mirror) bytes, 16-byte aligned
+        _ctrl = static_cast<int8_t *>(_aligned_malloc(_capacity + 16, 16));
+        _slots = static_cast<uint32_t *>(std::malloc(_capacity * sizeof(uint32_t)));
+        std::memset(_ctrl, kEmpty, _capacity + 16);
+        _size = 0;
+        _growth_left = _capacity_to_growth(_capacity);
+
+        // Re-insert old elements
+        if (old_ctrl != nullptr)
+        {
+            for (size_t i = 0; i < old_capacity; i++)
+            {
+                if (old_ctrl[i] >= 0) // occupied
+                {
+                    insert(old_slots[i]);
+                }
+            }
+            _aligned_free(old_ctrl);
+            std::free(old_slots);
+        }
+    }
+
+    void _free()
+    {
+        if (_ctrl != nullptr)
+        {
+            _aligned_free(_ctrl);
+            _ctrl = nullptr;
+        }
+        if (_slots != nullptr)
+        {
+            std::free(_slots);
+            _slots = nullptr;
+        }
+        _capacity = 0;
+        _size = 0;
+        _growth_left = 0;
+    }
+};
+
+} // namespace diskann
diff --git a/src/index.cpp b/src/index.cpp
index 5261adb5e..3d0e326a7 100644
--- a/src/index.cpp
+++ b/src/index.cpp
@@ -935,17 +935,14 @@ std::pair<uint32_t, uint32_t> Index<T, TagT, LabelT>::iterate_to_fixed_point(
         best_L_nodes = &(best_diverse_nodes_ref);
     }
     
-    tsl::robin_set<uint32_t> &inserted_into_pool_rs = scratch->inserted_into_pool_rs();
+    SwissHashSet &inserted_into_pool_rs = scratch->inserted_into_pool_rs();
     boost::dynamic_bitset<> &inserted_into_pool_bs = scratch->inserted_into_pool_bs();
     std::vector<uint32_t> &id_scratch = scratch->id_scratch();
-    std::vector<float> &dist_scratch = scratch->dist_scratch();
     assert(id_scratch.size() == 0);
 
     T *aligned_query = scratch->aligned_query();
     std::vector<std::uint64_t>& query_bitmask_buf = scratch->query_label_bitmask();
 
-    float *pq_dists = nullptr;
-
     _pq_data_store->preprocess_query(aligned_query, scratch);
 
     if (expanded_nodes.size() > 0 || id_scratch.size() > 0)
@@ -971,12 +968,7 @@ std::pair<uint32_t, uint32_t> Index<T, TagT, LabelT>::iterate_to_fixed_point(
     // Lambda to determine if a node has been visited
     auto is_not_visited = [this, fast_iterate, &inserted_into_pool_bs, &inserted_into_pool_rs](const uint32_t id) {
         return fast_iterate ? inserted_into_pool_bs[id] == 0
-                            : inserted_into_pool_rs.find(id) == inserted_into_pool_rs.end();
-    };
-
-    // Lambda to batch compute query<-> node distances in PQ space
-    auto compute_dists = [this, scratch, pq_dists](const std::vector<uint32_t> &ids, std::vector<float> &dists_out) {
-        _pq_data_store->get_distance(scratch->aligned_query(), ids, dists_out, scratch);
+                            : !inserted_into_pool_rs.find(id);
     };
 
     label_filter_match_holder match_proxy(
@@ -1034,7 +1026,6 @@ std::pair<uint32_t, uint32_t> Index<T, TagT, LabelT>::iterate_to_fixed_point(
     uint32_t hops = 0;
     uint32_t cmps = 0;
     cmps += static_cast<uint32_t>(init_ids.size());
-    std::vector<location_t> tmp_neighbor_list;
 
     while (best_L_nodes->has_unexpanded_node())
     {
@@ -1061,7 +1052,6 @@ std::pair<uint32_t, uint32_t> Index<T, TagT, LabelT>::iterate_to_fixed_point(
 
         // Find which of the nodes in des have not been visited before
         id_scratch.clear();
-        dist_scratch.clear();
         if (_dynamic_index)
         {
             LockGuard guard(_locks[n]);
@@ -1085,19 +1075,25 @@ std::pair<uint32_t, uint32_t> Index<T, TagT, LabelT>::iterate_to_fixed_point(
                     }
                 }
 
+                if (fast_iterate)
+                {
+                    inserted_into_pool_bs[id] = 1;
+                }
+                else
+                {
+                    inserted_into_pool_rs.insert(id);
+                }
                 id_scratch.push_back(id);
 
             }
         }
         else
         {
-            tmp_neighbor_list.clear();
+            // Static graph: iterate neighbors directly under shared lock,
+            // mark visited and collect unvisited into id_scratch
             _locks[n].lock_shared();
             auto nbrs = _graph_store->get_neighbours(n);
-            tmp_neighbor_list.resize(nbrs.size());
-            memcpy(tmp_neighbor_list.data(), nbrs.data(), nbrs.size() * sizeof(location_t));
-            _locks[n].unlock_shared();
-            for (auto id : tmp_neighbor_list)
+            for (auto id : nbrs)
             {
                 assert(id < _max_points);
 
@@ -1116,36 +1112,41 @@ std::pair<uint32_t, uint32_t> Index<T, TagT, LabelT>::iterate_to_fixed_point(
                     }
                 }
 
-
+                if (fast_iterate)
+                {
+                    inserted_into_pool_bs[id] = 1;
+                }
+                else
+                {
+                    inserted_into_pool_rs.insert(id);
+                }
                 id_scratch.push_back(id);
-
             }
+            _locks[n].unlock_shared();
         }
 
-        // Mark nodes visited
-        for (auto id : id_scratch)
+        // Sliding-window prefetch (lookahead=8) + compute + immediate insert.
+        // Pre-prefetch the first batch, then maintain the window during the loop.
+        constexpr size_t PREFETCH_LOOKAHEAD = 8;
+        cmps += (uint32_t)id_scratch.size();
+        const size_t id_count = id_scratch.size();
+        const size_t prefetch_init = std::min(PREFETCH_LOOKAHEAD, id_count);
+        for (size_t p = 0; p < prefetch_init; ++p)
         {
-            if (fast_iterate)
-            {
-                inserted_into_pool_bs[id] = 1;
-            }
-            else
-            {
-                inserted_into_pool_rs.insert(id);
-            }
+            _pq_data_store->prefetch_vector(id_scratch[p]);
         }
-
-        dist_scratch.resize(id_scratch.size());
-        //assert(dist_scratch.capacity() >= id_scratch.size());
-        compute_dists(id_scratch, dist_scratch);
-        cmps += (uint32_t)id_scratch.size();
-
-        // Insert <id, dist> pairs into the pool of candidates
-        for (size_t m = 0; m < id_scratch.size(); ++m)
+        for (size_t m = 0; m < id_count; ++m)
         {
-            best_L_nodes->insert(Neighbor(id_scratch[m], dist_scratch[m]));
+            if (m + PREFETCH_LOOKAHEAD < id_count)
+            {
+                _pq_data_store->prefetch_vector(id_scratch[m + PREFETCH_LOOKAHEAD]);
+            }
+            float distances[] = {std::numeric_limits<float>::max()};
+            uint32_t ids[] = {id_scratch[m]};
+            _pq_data_store->get_distance(aligned_query, ids, 1, distances, scratch);
+            best_L_nodes->insert(Neighbor(id_scratch[m], distances[0]));
             if (debug_info != nullptr)
-                debug_info->record_visited(id_scratch[m], dist_scratch[m]);
+                debug_info->record_visited(id_scratch[m], distances[0]);
         }
     }
     return std::make_pair(hops, cmps);
diff --git a/src/pq_l2_distance.cpp b/src/pq_l2_distance.cpp
index c08744c35..de9345f8f 100644
--- a/src/pq_l2_distance.cpp
+++ b/src/pq_l2_distance.cpp
@@ -2,6 +2,7 @@
 #include "pq.h"
 #include "pq_l2_distance.h"
 #include "pq_scratch.h"
+#include <immintrin.h>
 
 // block size for reading/processing large files and matrices in blocks
 #define BLOCK_SIZE 5000000
@@ -260,18 +261,20 @@ template <typename data_t>
 void PQL2Distance<data_t>::prepopulate_chunkwise_distances(const float *query_vec, float *dist_vec)
 {
     memset(dist_vec, 0, 256 * _num_chunks * sizeof(float));
-    // chunk wise distance computation
     for (size_t chunk = 0; chunk < _num_chunks; chunk++)
     {
-        // sum (q-c)^2 for the dimensions associated with this chunk
         float *chunk_dists = dist_vec + (256 * chunk);
         for (size_t j = _chunk_offsets[chunk]; j < _chunk_offsets[chunk + 1]; j++)
         {
             const float *centers_dim_vec = _tables_tr + (256 * j);
-            for (size_t idx = 0; idx < 256; idx++)
+            __m256 q_val = _mm256_set1_ps(query_vec[j]);
+            for (size_t idx = 0; idx < 256; idx += 8)
             {
-                double diff = centers_dim_vec[idx] - (query_vec[j]);
-                chunk_dists[idx] += (float)(diff * diff);
+                __m256 c = _mm256_loadu_ps(centers_dim_vec + idx);
+                __m256 d = _mm256_sub_ps(c, q_val);
+                __m256 acc = _mm256_loadu_ps(chunk_dists + idx);
+                acc = _mm256_fmadd_ps(d, d, acc);
+                _mm256_storeu_ps(chunk_dists + idx, acc);
             }
         }
     }
diff --git a/src/scratch.cpp b/src/scratch.cpp
index d1fcb7462..b93b7c4a1 100644
--- a/src/scratch.cpp
+++ b/src/scratch.cpp
@@ -50,18 +50,31 @@ InMemQueryScratch<T>::InMemQueryScratch(uint32_t search_l, uint32_t indexing_l,
 
 template <typename T> void InMemQueryScratch<T>::clear()
 {
+    // Clear vectors - std::vector::clear() is O(1) for POD types (just resets size)
     _pool.clear();
     _best_l_nodes.clear();
     _occlude_factor.clear();
     _candidate_pick_flags.clear();
 
-    _inserted_into_pool_rs.clear();
-    _inserted_into_pool_bs->reset();
+    // SwissHashSet::clear() is O(occupied_slots), only clear if non-empty
+    if (!_inserted_into_pool_rs.empty())
+    {
+        _inserted_into_pool_rs.clear();
+    }
+    // dynamic_bitset::reset() is O(n) where n is bitset size, only reset if has bits set
+    if (_inserted_into_pool_bs->size() > 0 && _inserted_into_pool_bs->any())
+    {
+        _inserted_into_pool_bs->reset();
+    }
 
     _id_scratch.clear();
     _dist_scratch.clear();
 
-    _expanded_nodes_set.clear();
+    // These are typically only used during indexing, not search - skip if empty
+    if (!_expanded_nodes_set.empty())
+    {
+        _expanded_nodes_set.clear();
+    }
     _expanded_nghrs_vec.clear();
     _occlude_list_output.clear();
     _query_label_bitmask.clear();
@@ -78,7 +91,7 @@ template <typename T> void InMemQueryScratch<T>::resize_for_new_L(uint32_t new_l
         _pool.reserve(3 * _L + _R);
         _best_l_nodes.reserve(_L);
 
-        _inserted_into_pool_rs.reserve(20 * _L);
+        _inserted_into_pool_rs.reserve((size_t)std::ceil(1.5 * _R * _L));
     }
 }
 

From 3cfa14b71e31e24cf847658bb4b3a23aa1cd7ddb Mon Sep 17 00:00:00 2001
From: Deyu Kong <deyukong@microsoft.com>
Date: Thu, 11 Jun 2026 11:04:26 +0800
Subject: [PATCH 2/2] ci: bump actions/upload|download-artifact from v3 to v4

v3 artifacts were retired by GitHub on 2024-04-16 and now fail builds.
Updates 10 references across 9 workflows and 1 composite action.

push-test.yml is intentionally left on @v3 because its upload step uses
a non-unique artifact name (`dependencies`) across a 3-entry os matrix,
which v4 rejects (each upload requires a unique name). That file needs
both a name suffix (e.g. `dependencies-${{ matrix.os }}`) and a version
bump together.
---
 .github/actions/python-wheel/action.yml    | 2 +-
 .github/workflows/disk-pq.yml              | 2 +-
 .github/workflows/dynamic-labels.yml       | 2 +-
 .github/workflows/dynamic.yml              | 2 +-
 .github/workflows/in-mem-no-pq.yml         | 2 +-
 .github/workflows/in-mem-pq.yml            | 2 +-
 .github/workflows/labels.yml               | 2 +-
 .github/workflows/multi-sector-disk-pq.yml | 2 +-
 .github/workflows/perf.yml                 | 2 +-
 .github/workflows/python-release.yml       | 2 +-
 10 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/.github/actions/python-wheel/action.yml b/.github/actions/python-wheel/action.yml
index 6a2880c6d..c58352ad8 100644
--- a/.github/actions/python-wheel/action.yml
+++ b/.github/actions/python-wheel/action.yml
@@ -16,7 +16,7 @@ runs:
       env:
         CIBW_BUILD: ${{inputs.cibw-identifier}}
       shell: bash
-    - uses: actions/upload-artifact@v3
+    - uses: actions/upload-artifact@v4
       with:
         name: wheels
         path: ./dist/*.whl
diff --git a/.github/workflows/disk-pq.yml b/.github/workflows/disk-pq.yml
index 6e71e7999..f30f1483b 100644
--- a/.github/workflows/disk-pq.yml
+++ b/.github/workflows/disk-pq.yml
@@ -109,7 +109,7 @@ jobs:
           dist/bin/search_disk_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_mips_rand_float_10D_10K_norm1.0_diskpq_sharded --result_path /tmp/res --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/mips_rand_float_10D_10K_norm1.0_10D_1K_norm1.0_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16
 
       - name: upload data and bin
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: disk-pq
           path: |
diff --git a/.github/workflows/dynamic-labels.yml b/.github/workflows/dynamic-labels.yml
index 0f3b56eb9..43cb6e6bd 100644
--- a/.github/workflows/dynamic-labels.yml
+++ b/.github/workflows/dynamic-labels.yml
@@ -94,7 +94,7 @@ jobs:
           dist/bin/search_memory_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_rand_ins_del.after-concurrent-delete-del2500-7500 --result_path res_stream --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/gt100_rand_random10D_1K -K 10 -L 20 40 60 80 100 -T 64
 
       - name: upload data and bin
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: dynamic
           path: |
diff --git a/.github/workflows/dynamic.yml b/.github/workflows/dynamic.yml
index 35eb6d42d..5be88ddf5 100644
--- a/.github/workflows/dynamic.yml
+++ b/.github/workflows/dynamic.yml
@@ -67,7 +67,7 @@ jobs:
           dist/bin/search_memory_index --data_type uint8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_ins_del.after-concurrent-delete-del2500-7500 --result_path data/res_ins_del --query_file data/rand_uint8_10D_1K_norm50.0.bin --gt_file data/gt100_random10D_10K-conc-2500-7500 -K 10 -L 20 40 60 80 100 -T 8 --dynamic true --tags 1
 
       - name: upload data and bin
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: dynamic
           path: |
diff --git a/.github/workflows/in-mem-no-pq.yml b/.github/workflows/in-mem-no-pq.yml
index 0039754d2..84263e1d5 100644
--- a/.github/workflows/in-mem-no-pq.yml
+++ b/.github/workflows/in-mem-no-pq.yml
@@ -73,7 +73,7 @@ jobs:
           dist/bin/search_memory_index --data_type uint8 --dist_fn cosine --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_uint8_10D_10K_norm50.0 --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/cosine_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 -L  16 32
 
       - name: upload data and bin
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: in-memory-no-pq
           path: |
diff --git a/.github/workflows/in-mem-pq.yml b/.github/workflows/in-mem-pq.yml
index f9276adfc..8726d2c65 100644
--- a/.github/workflows/in-mem-pq.yml
+++ b/.github/workflows/in-mem-pq.yml
@@ -48,7 +48,7 @@ jobs:
           dist/bin/search_memory_index --data_type uint8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_uint8_10D_10K_norm50.0_buildpq5 --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/l2_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 -L  16 32
 
       - name: upload data and bin
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: in-memory-pq
           path: |
diff --git a/.github/workflows/labels.yml b/.github/workflows/labels.yml
index 5555f7f84..e08a1571e 100644
--- a/.github/workflows/labels.yml
+++ b/.github/workflows/labels.yml
@@ -111,7 +111,7 @@ jobs:
           dist/bin/search_memory_index --num_threads 48 --data_type uint8 --dist_fn l2 --filter_label 5 --index_path_prefix data/stit_zipf_32_100_64_new --query_file data/rand_uint8_10D_1K_norm50.0.bin --result_path data/zipf_stit_96_10_90_new --gt_file data/l2_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel -K 10 -L 16 32 150
 
       - name: upload data and bin
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: labels
           path: |
diff --git a/.github/workflows/multi-sector-disk-pq.yml b/.github/workflows/multi-sector-disk-pq.yml
index 8ea55c88d..391d80d64 100644
--- a/.github/workflows/multi-sector-disk-pq.yml
+++ b/.github/workflows/multi-sector-disk-pq.yml
@@ -52,7 +52,7 @@ jobs:
           dist/bin/search_disk_index --data_type int8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_int8_4096D_5K_norm1.0_diskfull_oneshot --result_path /tmp/res --query_file data/rand_int8_4096D_1K_norm1.0.bin --gt_file data/l2_rand_int8_4096D_5K_norm1.0_4096D_1K_norm1.0_gt100 --recall_at 5 -L 250 -W 2 --num_nodes_to_cache 100 -T 16
 
       - name: upload data and bin
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: multi-sector-disk-pq
           path: |
diff --git a/.github/workflows/perf.yml b/.github/workflows/perf.yml
index 1595a4221..b399b70cb 100644
--- a/.github/workflows/perf.yml
+++ b/.github/workflows/perf.yml
@@ -19,7 +19,7 @@ jobs:
           mkdir metrics
           docker run -v ./metrics:/app/logs perf &> ./metrics/combined_stdouterr.log
       - name: Upload Metrics Logs
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: metrics
           path: |
diff --git a/.github/workflows/python-release.yml b/.github/workflows/python-release.yml
index a15d4d161..64398a9d5 100644
--- a/.github/workflows/python-release.yml
+++ b/.github/workflows/python-release.yml
@@ -17,7 +17,7 @@ jobs:
     runs-on: ubuntu-latest
     needs: python-release-wheels
     steps:
-      - uses: actions/download-artifact@v3
+      - uses: actions/download-artifact@v4
         with:
           name: wheels
           path: dist/