perf: ANN optimizations - SwissHashSet, AVX2 distance, insert_fast

.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

.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: |

.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: |

.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: |

.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: |

.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: |

.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: |

.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: |

.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: |

.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/

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

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