faiss 0.3.0 → 0.3.2

data/vendor/faiss/faiss/impl/HNSW.cpp

@@ -5,15 +5,25 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 #include <faiss/impl/HNSW.h>
 
+#include <cstddef>
 #include <string>
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/IDSelector.h>
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/utils/prefetch.h>
+
+#include <faiss/impl/platform_macros.h>
+
+#ifdef __AVX2__
+#include <immintrin.h>
+
+#include <limits>
+#include <type_traits>
+#endif
 
 namespace faiss {
 
@@ -101,8 +111,8 @@ void HNSW::print_neighbor_stats(int level) const {
            level,
            nb_neighbors(level));
     size_t tot_neigh = 0, tot_common = 0, tot_reciprocal = 0, n_node = 0;
-#pragma omp parallel for reduction(+: tot_neigh) reduction(+: tot_common) \
-  reduction(+: tot_reciprocal) reduction(+: n_node)
+#pragma omp parallel for reduction(+ : tot_neigh) reduction(+ : tot_common) \
+        reduction(+ : tot_reciprocal) reduction(+ : n_node)
     for (int i = 0; i < levels.size(); i++) {
         if (levels[i] > level) {
             n_node++;
@@ -206,13 +216,13 @@ int HNSW::prepare_level_tab(size_t n, bool preset_levels) {
         if (pt_level > max_level)
             max_level = pt_level;
         offsets.push_back(offsets.back() + cum_nb_neighbors(pt_level + 1));
-        neighbors.resize(offsets.back(), -1);
     }
+    neighbors.resize(offsets.back(), -1);
 
     return max_level;
 }
 
-/** Enumerate vertices from farthest to nearest from query, keep a
+/** Enumerate vertices from nearest to farthest from query, keep a
  * neighbor only if there is no previous neighbor that is closer to
  * that vertex than the query.
  */
@@ -220,7 +230,14 @@ void HNSW::shrink_neighbor_list(
         DistanceComputer& qdis,
         std::priority_queue<NodeDistFarther>& input,
         std::vector<NodeDistFarther>& output,
-        int max_size) {
+        int max_size,
+        bool keep_max_size_level0) {
+    // This prevents number of neighbors at
+    // level 0 from being shrunk to less than 2 * M.
+    // This is essential in making sure
+    // `faiss::gpu::GpuIndexCagra::copyFrom(IndexHNSWCagra*)` is functional
+    std::vector<NodeDistFarther> outsiders;
+
     while (input.size() > 0) {
         NodeDistFarther v1 = input.top();
         input.pop();
@@ -241,8 +258,15 @@ void HNSW::shrink_neighbor_list(
             if (output.size() >= max_size) {
                 return;
             }
+        } else if (keep_max_size_level0) {
+            outsiders.push_back(v1);
         }
     }
+    size_t idx = 0;
+    while (keep_max_size_level0 && (output.size() < max_size) &&
+           (idx < outsiders.size())) {
+        output.push_back(outsiders[idx++]);
+    }
 }
 
 namespace {
@@ -259,7 +283,8 @@ using NodeDistFarther = HNSW::NodeDistFarther;
 void shrink_neighbor_list(
         DistanceComputer& qdis,
         std::priority_queue<NodeDistCloser>& resultSet1,
-        int max_size) {
+        int max_size,
+        bool keep_max_size_level0 = false) {
     if (resultSet1.size() < max_size) {
         return;
     }
@@ -271,7 +296,8 @@ void shrink_neighbor_list(
         resultSet1.pop();
     }
 
-    HNSW::shrink_neighbor_list(qdis, resultSet, returnlist, max_size);
+    HNSW::shrink_neighbor_list(
+            qdis, resultSet, returnlist, max_size, keep_max_size_level0);
 
     for (NodeDistFarther curen2 : returnlist) {
         resultSet1.emplace(curen2.d, curen2.id);
@@ -285,7 +311,8 @@ void add_link(
         DistanceComputer& qdis,
         storage_idx_t src,
         storage_idx_t dest,
-        int level) {
+        int level,
+        bool keep_max_size_level0 = false) {
     size_t begin, end;
     hnsw.neighbor_range(src, level, &begin, &end);
     if (hnsw.neighbors[end - 1] == -1) {
@@ -310,7 +337,7 @@ void add_link(
         resultSet.emplace(qdis.symmetric_dis(src, neigh), neigh);
     }
 
-    shrink_neighbor_list(qdis, resultSet, end - begin);
+    shrink_neighbor_list(qdis, resultSet, end - begin, keep_max_size_level0);
 
     // ...and back
     size_t i = begin;
@@ -333,6 +360,9 @@ void search_neighbors_to_add(
         float d_entry_point,
         int level,
         VisitedTable& vt) {
+    // selects a version
+    const bool reference_version = false;
+
     // top is nearest candidate
     std::priority_queue<NodeDistFarther> candidates;
 
@@ -354,59 +384,90 @@ void search_neighbors_to_add(
         // loop over neighbors
         size_t begin, end;
         hnsw.neighbor_range(currNode, level, &begin, &end);
-        for (size_t i = begin; i < end; i++) {
-            storage_idx_t nodeId = hnsw.neighbors[i];
-            if (nodeId < 0)
-                break;
-            if (vt.get(nodeId))
-                continue;
-            vt.set(nodeId);
-
-            float dis = qdis(nodeId);
-            NodeDistFarther evE1(dis, nodeId);
 
-            if (results.size() < hnsw.efConstruction || results.top().d > dis) {
-                results.emplace(dis, nodeId);
-                candidates.emplace(dis, nodeId);
-                if (results.size() > hnsw.efConstruction) {
-                    results.pop();
+        // select a version, based on a flag
+        if (reference_version) {
+            // a reference version
+            for (size_t i = begin; i < end; i++) {
+                storage_idx_t nodeId = hnsw.neighbors[i];
+                if (nodeId < 0)
+                    break;
+                if (vt.get(nodeId))
+                    continue;
+                vt.set(nodeId);
+
+                float dis = qdis(nodeId);
+                NodeDistFarther evE1(dis, nodeId);
+
+                if (results.size() < hnsw.efConstruction ||
+                    results.top().d > dis) {
+                    results.emplace(dis, nodeId);
+                    candidates.emplace(dis, nodeId);
+                    if (results.size() > hnsw.efConstruction) {
+                        results.pop();
+                    }
                 }
             }
-        }
-    }
-    vt.advance();
-}
+        } else {
+            // a faster version
+
+            // the following version processes 4 neighbors at a time
+            auto update_with_candidate = [&](const storage_idx_t idx,
+                                             const float dis) {
+                if (results.size() < hnsw.efConstruction ||
+                    results.top().d > dis) {
+                    results.emplace(dis, idx);
+                    candidates.emplace(dis, idx);
+                    if (results.size() > hnsw.efConstruction) {
+                        results.pop();
+                    }
+                }
+            };
 
-/**************************************************************
- * Searching subroutines
- **************************************************************/
+            int n_buffered = 0;
+            storage_idx_t buffered_ids[4];
 
-/// greedily update a nearest vector at a given level
-void greedy_update_nearest(
-        const HNSW& hnsw,
-        DistanceComputer& qdis,
-        int level,
-        storage_idx_t& nearest,
-        float& d_nearest) {
-    for (;;) {
-        storage_idx_t prev_nearest = nearest;
+            for (size_t j = begin; j < end; j++) {
+                storage_idx_t nodeId = hnsw.neighbors[j];
+                if (nodeId < 0)
+                    break;
+                if (vt.get(nodeId)) {
+                    continue;
+                }
+                vt.set(nodeId);
+
+                buffered_ids[n_buffered] = nodeId;
+                n_buffered += 1;
+
+                if (n_buffered == 4) {
+                    float dis[4];
+                    qdis.distances_batch_4(
+                            buffered_ids[0],
+                            buffered_ids[1],
+                            buffered_ids[2],
+                            buffered_ids[3],
+                            dis[0],
+                            dis[1],
+                            dis[2],
+                            dis[3]);
+
+                    for (size_t id4 = 0; id4 < 4; id4++) {
+                        update_with_candidate(buffered_ids[id4], dis[id4]);
+                    }
 
-        size_t begin, end;
-        hnsw.neighbor_range(nearest, level, &begin, &end);
-        for (size_t i = begin; i < end; i++) {
-            storage_idx_t v = hnsw.neighbors[i];
-            if (v < 0)
-                break;
-            float dis = qdis(v);
-            if (dis < d_nearest) {
-                nearest = v;
-                d_nearest = dis;
+                    n_buffered = 0;
+                }
+            }
+
+            // process leftovers
+            for (size_t icnt = 0; icnt < n_buffered; icnt++) {
+                float dis = qdis(buffered_ids[icnt]);
+                update_with_candidate(buffered_ids[icnt], dis);
             }
-        }
-        if (nearest == prev_nearest) {
-            return;
         }
     }
+
+    vt.advance();
 }
 
 } // namespace
@@ -420,7 +481,8 @@ void HNSW::add_links_starting_from(
         float d_nearest,
         int level,
         omp_lock_t* locks,
-        VisitedTable& vt) {
+        VisitedTable& vt,
+        bool keep_max_size_level0) {
     std::priority_queue<NodeDistCloser> link_targets;
 
     search_neighbors_to_add(
@@ -429,13 +491,13 @@ void HNSW::add_links_starting_from(
     // but we can afford only this many neighbors
     int M = nb_neighbors(level);
 
-    ::faiss::shrink_neighbor_list(ptdis, link_targets, M);
+    ::faiss::shrink_neighbor_list(ptdis, link_targets, M, keep_max_size_level0);
 
     std::vector<storage_idx_t> neighbors;
     neighbors.reserve(link_targets.size());
     while (!link_targets.empty()) {
         storage_idx_t other_id = link_targets.top().id;
-        add_link(*this, ptdis, pt_id, other_id, level);
+        add_link(*this, ptdis, pt_id, other_id, level, keep_max_size_level0);
         neighbors.push_back(other_id);
         link_targets.pop();
     }
@@ -443,7 +505,7 @@ void HNSW::add_links_starting_from(
     omp_unset_lock(&locks[pt_id]);
     for (storage_idx_t other_id : neighbors) {
         omp_set_lock(&locks[other_id]);
-        add_link(*this, ptdis, other_id, pt_id, level);
+        add_link(*this, ptdis, other_id, pt_id, level, keep_max_size_level0);
         omp_unset_lock(&locks[other_id]);
     }
     omp_set_lock(&locks[pt_id]);
@@ -458,7 +520,8 @@ void HNSW::add_with_locks(
         int pt_level,
         int pt_id,
         std::vector<omp_lock_t>& locks,
-        VisitedTable& vt) {
+        VisitedTable& vt,
+        bool keep_max_size_level0) {
     //  greedy search on upper levels
 
     storage_idx_t nearest;
@@ -487,7 +550,14 @@ void HNSW::add_with_locks(
 
     for (; level >= 0; level--) {
         add_links_starting_from(
-                ptdis, pt_id, nearest, d_nearest, level, locks.data(), vt);
+                ptdis,
+                pt_id,
+                nearest,
+                d_nearest,
+                level,
+                locks.data(),
+                vt,
+                keep_max_size_level0);
     }
 
     omp_unset_lock(&locks[pt_id]);
@@ -502,24 +572,20 @@ void HNSW::add_with_locks(
  * Searching
  **************************************************************/
 
-namespace {
-
 using MinimaxHeap = HNSW::MinimaxHeap;
 using Node = HNSW::Node;
+using C = HNSW::C;
 /** Do a BFS on the candidates list */
-
 int search_from_candidates(
         const HNSW& hnsw,
         DistanceComputer& qdis,
-        int k,
-        idx_t* I,
-        float* D,
+        ResultHandler<C>& res,
         MinimaxHeap& candidates,
         VisitedTable& vt,
         HNSWStats& stats,
         int level,
-        int nres_in = 0,
-        const SearchParametersHNSW* params = nullptr) {
+        int nres_in,
+        const SearchParametersHNSW* params) {
     int nres = nres_in;
     int ndis = 0;
 
@@ -529,15 +595,16 @@ int search_from_candidates(
     int efSearch = params ? params->efSearch : hnsw.efSearch;
     const IDSelector* sel = params ? params->sel : nullptr;
 
+    C::T threshold = res.threshold;
     for (int i = 0; i < candidates.size(); i++) {
         idx_t v1 = candidates.ids[i];
         float d = candidates.dis[i];
         FAISS_ASSERT(v1 >= 0);
         if (!sel || sel->is_member(v1)) {
-            if (nres < k) {
-                faiss::maxheap_push(++nres, D, I, d, v1);
-            } else if (d < D[0]) {
-                faiss::maxheap_replace_top(nres, D, I, d, v1);
+            if (d < threshold) {
+                if (res.add_result(d, v1)) {
+                    threshold = res.threshold;
+                }
             }
         }
         vt.set(v1);
@@ -563,24 +630,70 @@ int search_from_candidates(
         size_t begin, end;
         hnsw.neighbor_range(v0, level, &begin, &end);
 
+        // a faster version: reference version in unit test test_hnsw.cpp
+        // the following version processes 4 neighbors at a time
+        size_t jmax = begin;
         for (size_t j = begin; j < end; j++) {
             int v1 = hnsw.neighbors[j];
             if (v1 < 0)
                 break;
-            if (vt.get(v1)) {
-                continue;
+
+            prefetch_L2(vt.visited.data() + v1);
+            jmax += 1;
+        }
+
+        int counter = 0;
+        size_t saved_j[4];
+
+        threshold = res.threshold;
+
+        auto add_to_heap = [&](const size_t idx, const float dis) {
+            if (!sel || sel->is_member(idx)) {
+                if (dis < threshold) {
+                    if (res.add_result(dis, idx)) {
+                        threshold = res.threshold;
+                        nres += 1;
+                    }
+                }
             }
+            candidates.push(idx, dis);
+        };
+
+        for (size_t j = begin; j < jmax; j++) {
+            int v1 = hnsw.neighbors[j];
+
+            bool vget = vt.get(v1);
             vt.set(v1);
-            ndis++;
-            float d = qdis(v1);
-            if (!sel || sel->is_member(v1)) {
-                if (nres < k) {
-                    faiss::maxheap_push(++nres, D, I, d, v1);
-                } else if (d < D[0]) {
-                    faiss::maxheap_replace_top(nres, D, I, d, v1);
+            saved_j[counter] = v1;
+            counter += vget ? 0 : 1;
+
+            if (counter == 4) {
+                float dis[4];
+                qdis.distances_batch_4(
+                        saved_j[0],
+                        saved_j[1],
+                        saved_j[2],
+                        saved_j[3],
+                        dis[0],
+                        dis[1],
+                        dis[2],
+                        dis[3]);
+
+                for (size_t id4 = 0; id4 < 4; id4++) {
+                    add_to_heap(saved_j[id4], dis[id4]);
                 }
+
+                ndis += 4;
+
+                counter = 0;
             }
-            candidates.push(v1, d);
+        }
+
+        for (size_t icnt = 0; icnt < counter; icnt++) {
+            float dis = qdis(saved_j[icnt]);
+            add_to_heap(saved_j[icnt], dis);
+
+            ndis += 1;
         }
 
         nstep++;
@@ -594,7 +707,8 @@ int search_from_candidates(
         if (candidates.size() == 0) {
             stats.n2++;
         }
-        stats.n3 += ndis;
+        stats.ndis += ndis;
+        stats.nhops += nstep;
     }
 
     return nres;
@@ -630,151 +744,241 @@ std::priority_queue<HNSW::Node> search_from_candidate_unbounded(
         size_t begin, end;
         hnsw.neighbor_range(v0, 0, &begin, &end);
 
-        for (size_t j = begin; j < end; ++j) {
+        // a faster version: reference version in unit test test_hnsw.cpp
+        // the following version processes 4 neighbors at a time
+        size_t jmax = begin;
+        for (size_t j = begin; j < end; j++) {
             int v1 = hnsw.neighbors[j];
-
-            if (v1 < 0) {
+            if (v1 < 0)
                 break;
-            }
-            if (vt->get(v1)) {
-                continue;
-            }
 
-            vt->set(v1);
+            prefetch_L2(vt->visited.data() + v1);
+            jmax += 1;
+        }
 
-            float d1 = qdis(v1);
-            ++ndis;
+        int counter = 0;
+        size_t saved_j[4];
 
-            if (top_candidates.top().first > d1 || top_candidates.size() < ef) {
-                candidates.emplace(d1, v1);
-                top_candidates.emplace(d1, v1);
+        auto add_to_heap = [&](const size_t idx, const float dis) {
+            if (top_candidates.top().first > dis ||
+                top_candidates.size() < ef) {
+                candidates.emplace(dis, idx);
+                top_candidates.emplace(dis, idx);
 
                 if (top_candidates.size() > ef) {
                     top_candidates.pop();
                 }
             }
+        };
+
+        for (size_t j = begin; j < jmax; j++) {
+            int v1 = hnsw.neighbors[j];
+
+            bool vget = vt->get(v1);
+            vt->set(v1);
+            saved_j[counter] = v1;
+            counter += vget ? 0 : 1;
+
+            if (counter == 4) {
+                float dis[4];
+                qdis.distances_batch_4(
+                        saved_j[0],
+                        saved_j[1],
+                        saved_j[2],
+                        saved_j[3],
+                        dis[0],
+                        dis[1],
+                        dis[2],
+                        dis[3]);
+
+                for (size_t id4 = 0; id4 < 4; id4++) {
+                    add_to_heap(saved_j[id4], dis[id4]);
+                }
+
+                ndis += 4;
+
+                counter = 0;
+            }
+        }
+
+        for (size_t icnt = 0; icnt < counter; icnt++) {
+            float dis = qdis(saved_j[icnt]);
+            add_to_heap(saved_j[icnt], dis);
+
+            ndis += 1;
         }
+
+        stats.nhops += 1;
     }
 
     ++stats.n1;
     if (candidates.size() == 0) {
         ++stats.n2;
     }
-    stats.n3 += ndis;
+    stats.ndis += ndis;
 
     return top_candidates;
 }
 
-} // anonymous namespace
-
-HNSWStats HNSW::search(
+/// greedily update a nearest vector at a given level
+HNSWStats greedy_update_nearest(
+        const HNSW& hnsw,
         DistanceComputer& qdis,
-        int k,
-        idx_t* I,
-        float* D,
-        VisitedTable& vt,
-        const SearchParametersHNSW* params) const {
+        int level,
+        storage_idx_t& nearest,
+        float& d_nearest) {
     HNSWStats stats;
-    if (entry_point == -1) {
-        return stats;
-    }
-    if (upper_beam == 1) {
-        //  greedy search on upper levels
-        storage_idx_t nearest = entry_point;
-        float d_nearest = qdis(nearest);
 
-        for (int level = max_level; level >= 1; level--) {
-            greedy_update_nearest(*this, qdis, level, nearest, d_nearest);
-        }
+    for (;;) {
+        storage_idx_t prev_nearest = nearest;
 
-        int ef = std::max(efSearch, k);
-        if (search_bounded_queue) { // this is the most common branch
-            MinimaxHeap candidates(ef);
+        size_t begin, end;
+        hnsw.neighbor_range(nearest, level, &begin, &end);
 
-            candidates.push(nearest, d_nearest);
+        size_t ndis = 0;
 
-            search_from_candidates(
-                    *this, qdis, k, I, D, candidates, vt, stats, 0, 0, params);
-        } else {
-            std::priority_queue<Node> top_candidates =
-                    search_from_candidate_unbounded(
-                            *this,
-                            Node(d_nearest, nearest),
-                            qdis,
-                            ef,
-                            &vt,
-                            stats);
-
-            while (top_candidates.size() > k) {
-                top_candidates.pop();
+        // a faster version: reference version in unit test test_hnsw.cpp
+        // the following version processes 4 neighbors at a time
+        auto update_with_candidate = [&](const storage_idx_t idx,
+                                         const float dis) {
+            if (dis < d_nearest) {
+                nearest = idx;
+                d_nearest = dis;
             }
+        };
+
+        int n_buffered = 0;
+        storage_idx_t buffered_ids[4];
+
+        for (size_t j = begin; j < end; j++) {
+            storage_idx_t v = hnsw.neighbors[j];
+            if (v < 0)
+                break;
+            ndis += 1;
+
+            buffered_ids[n_buffered] = v;
+            n_buffered += 1;
+
+            if (n_buffered == 4) {
+                float dis[4];
+                qdis.distances_batch_4(
+                        buffered_ids[0],
+                        buffered_ids[1],
+                        buffered_ids[2],
+                        buffered_ids[3],
+                        dis[0],
+                        dis[1],
+                        dis[2],
+                        dis[3]);
+
+                for (size_t id4 = 0; id4 < 4; id4++) {
+                    update_with_candidate(buffered_ids[id4], dis[id4]);
+                }
 
-            int nres = 0;
-            while (!top_candidates.empty()) {
-                float d;
-                storage_idx_t label;
-                std::tie(d, label) = top_candidates.top();
-                faiss::maxheap_push(++nres, D, I, d, label);
-                top_candidates.pop();
+                n_buffered = 0;
             }
         }
 
-        vt.advance();
+        // process leftovers
+        for (size_t icnt = 0; icnt < n_buffered; icnt++) {
+            float dis = qdis(buffered_ids[icnt]);
+            update_with_candidate(buffered_ids[icnt], dis);
+        }
 
-    } else {
-        int candidates_size = upper_beam;
-        MinimaxHeap candidates(candidates_size);
+        // update stats
+        stats.ndis += ndis;
+        stats.nhops += 1;
 
-        std::vector<idx_t> I_to_next(candidates_size);
-        std::vector<float> D_to_next(candidates_size);
+        if (nearest == prev_nearest) {
+            return stats;
+        }
+    }
+}
 
-        int nres = 1;
-        I_to_next[0] = entry_point;
-        D_to_next[0] = qdis(entry_point);
+namespace {
+using MinimaxHeap = HNSW::MinimaxHeap;
+using Node = HNSW::Node;
+using C = HNSW::C;
 
-        for (int level = max_level; level >= 0; level--) {
-            // copy I, D -> candidates
+// just used as a lower bound for the minmaxheap, but it is set for heap search
+int extract_k_from_ResultHandler(ResultHandler<C>& res) {
+    using RH = HeapBlockResultHandler<C>;
+    if (auto hres = dynamic_cast<RH::SingleResultHandler*>(&res)) {
+        return hres->k;
+    }
+    return 1;
+}
 
-            candidates.clear();
+} // namespace
 
-            for (int i = 0; i < nres; i++) {
-                candidates.push(I_to_next[i], D_to_next[i]);
-            }
+HNSWStats HNSW::search(
+        DistanceComputer& qdis,
+        ResultHandler<C>& res,
+        VisitedTable& vt,
+        const SearchParametersHNSW* params) const {
+    HNSWStats stats;
+    if (entry_point == -1) {
+        return stats;
+    }
+    int k = extract_k_from_ResultHandler(res);
 
-            if (level == 0) {
-                nres = search_from_candidates(
-                        *this, qdis, k, I, D, candidates, vt, stats, 0);
-            } else {
-                nres = search_from_candidates(
-                        *this,
-                        qdis,
-                        candidates_size,
-                        I_to_next.data(),
-                        D_to_next.data(),
-                        candidates,
-                        vt,
-                        stats,
-                        level);
-            }
-            vt.advance();
+    bool bounded_queue =
+            params ? params->bounded_queue : this->search_bounded_queue;
+
+    //  greedy search on upper levels
+    storage_idx_t nearest = entry_point;
+    float d_nearest = qdis(nearest);
+
+    for (int level = max_level; level >= 1; level--) {
+        HNSWStats local_stats =
+                greedy_update_nearest(*this, qdis, level, nearest, d_nearest);
+        stats.combine(local_stats);
+    }
+
+    int ef = std::max(params ? params->efSearch : efSearch, k);
+    if (bounded_queue) { // this is the most common branch
+        MinimaxHeap candidates(ef);
+
+        candidates.push(nearest, d_nearest);
+
+        search_from_candidates(
+                *this, qdis, res, candidates, vt, stats, 0, 0, params);
+    } else {
+        std::priority_queue<Node> top_candidates =
+                search_from_candidate_unbounded(
+                        *this, Node(d_nearest, nearest), qdis, ef, &vt, stats);
+
+        while (top_candidates.size() > k) {
+            top_candidates.pop();
+        }
+
+        while (!top_candidates.empty()) {
+            float d;
+            storage_idx_t label;
+            std::tie(d, label) = top_candidates.top();
+            res.add_result(d, label);
+            top_candidates.pop();
         }
     }
 
+    vt.advance();
+
     return stats;
 }
 
 void HNSW::search_level_0(
         DistanceComputer& qdis,
-        int k,
-        idx_t* idxi,
-        float* simi,
+        ResultHandler<C>& res,
         idx_t nprobe,
         const storage_idx_t* nearest_i,
         const float* nearest_d,
         int search_type,
         HNSWStats& search_stats,
-        VisitedTable& vt) const {
+        VisitedTable& vt,
+        const SearchParametersHNSW* params) const {
     const HNSW& hnsw = *this;
+    auto efSearch = params ? params->efSearch : hnsw.efSearch;
+    int k = extract_k_from_ResultHandler(res);
 
     if (search_type == 1) {
         int nres = 0;
@@ -788,7 +992,7 @@ void HNSW::search_level_0(
             if (vt.get(cj))
                 continue;
 
-            int candidates_size = std::max(hnsw.efSearch, int(k));
+            int candidates_size = std::max(efSearch, k);
             MinimaxHeap candidates(candidates_size);
 
             candidates.push(cj, nearest_d[j]);
@@ -796,17 +1000,17 @@ void HNSW::search_level_0(
             nres = search_from_candidates(
                     hnsw,
                     qdis,
-                    k,
-                    idxi,
-                    simi,
+                    res,
                     candidates,
                     vt,
                     search_stats,
                     0,
-                    nres);
+                    nres,
+                    params);
+            nres = std::min(nres, candidates_size);
         }
     } else if (search_type == 2) {
-        int candidates_size = std::max(hnsw.efSearch, int(k));
+        int candidates_size = std::max(efSearch, int(k));
         candidates_size = std::max(candidates_size, int(nprobe));
 
         MinimaxHeap candidates(candidates_size);
@@ -819,10 +1023,43 @@ void HNSW::search_level_0(
         }
 
         search_from_candidates(
-                hnsw, qdis, k, idxi, simi, candidates, vt, search_stats, 0);
+                hnsw, qdis, res, candidates, vt, search_stats, 0, 0, params);
     }
 }
 
+void HNSW::permute_entries(const idx_t* map) {
+    // remap levels
+    storage_idx_t ntotal = levels.size();
+    std::vector<storage_idx_t> imap(ntotal); // inverse mapping
+    // map: new index -> old index
+    // imap: old index -> new index
+    for (int i = 0; i < ntotal; i++) {
+        assert(map[i] >= 0 && map[i] < ntotal);
+        imap[map[i]] = i;
+    }
+    if (entry_point != -1) {
+        entry_point = imap[entry_point];
+    }
+    std::vector<int> new_levels(ntotal);
+    std::vector<size_t> new_offsets(ntotal + 1);
+    std::vector<storage_idx_t> new_neighbors(neighbors.size());
+    size_t no = 0;
+    for (int i = 0; i < ntotal; i++) {
+        storage_idx_t o = map[i]; // corresponding "old" index
+        new_levels[i] = levels[o];
+        for (size_t j = offsets[o]; j < offsets[o + 1]; j++) {
+            storage_idx_t neigh = neighbors[j];
+            new_neighbors[no++] = neigh >= 0 ? imap[neigh] : neigh;
+        }
+        new_offsets[i + 1] = no;
+    }
+    assert(new_offsets[ntotal] == offsets[ntotal]);
+    // swap everyone
+    std::swap(levels, new_levels);
+    std::swap(offsets, new_offsets);
+    std::swap(neighbors, new_neighbors);
+}
+
 /**************************************************************
  * MinimaxHeap
  **************************************************************/
@@ -852,17 +1089,197 @@ void HNSW::MinimaxHeap::clear() {
     nvalid = k = 0;
 }
 
+#ifdef __AVX512F__
+
+int HNSW::MinimaxHeap::pop_min(float* vmin_out) {
+    assert(k > 0);
+    static_assert(
+            std::is_same<storage_idx_t, int32_t>::value,
+            "This code expects storage_idx_t to be int32_t");
+
+    int32_t min_idx = -1;
+    float min_dis = std::numeric_limits<float>::infinity();
+
+    __m512i min_indices = _mm512_set1_epi32(-1);
+    __m512 min_distances =
+            _mm512_set1_ps(std::numeric_limits<float>::infinity());
+    __m512i current_indices = _mm512_setr_epi32(
+            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
+    __m512i offset = _mm512_set1_epi32(16);
+
+    // The following loop tracks the rightmost index with the min distance.
+    // -1 index values are ignored.
+    const int k16 = (k / 16) * 16;
+    for (size_t iii = 0; iii < k16; iii += 16) {
+        __m512i indices =
+                _mm512_loadu_si512((const __m512i*)(ids.data() + iii));
+        __m512 distances = _mm512_loadu_ps(dis.data() + iii);
+
+        // This mask filters out -1 values among indices.
+        __mmask16 m1mask =
+                _mm512_cmpgt_epi32_mask(_mm512_setzero_si512(), indices);
+
+        __mmask16 dmask =
+                _mm512_cmp_ps_mask(min_distances, distances, _CMP_LT_OS);
+        __mmask16 finalmask = m1mask | dmask;
+
+        const __m512i min_indices_new = _mm512_mask_blend_epi32(
+                finalmask, current_indices, min_indices);
+        const __m512 min_distances_new =
+                _mm512_mask_blend_ps(finalmask, distances, min_distances);
+
+        min_indices = min_indices_new;
+        min_distances = min_distances_new;
+
+        current_indices = _mm512_add_epi32(current_indices, offset);
+    }
+
+    // leftovers
+    if (k16 != k) {
+        const __mmask16 kmask = (1 << (k - k16)) - 1;
+
+        __m512i indices = _mm512_mask_loadu_epi32(
+                _mm512_set1_epi32(-1), kmask, ids.data() + k16);
+        __m512 distances = _mm512_maskz_loadu_ps(kmask, dis.data() + k16);
+
+        // This mask filters out -1 values among indices.
+        __mmask16 m1mask =
+                _mm512_cmpgt_epi32_mask(_mm512_setzero_si512(), indices);
+
+        __mmask16 dmask =
+                _mm512_cmp_ps_mask(min_distances, distances, _CMP_LT_OS);
+        __mmask16 finalmask = m1mask | dmask;
+
+        const __m512i min_indices_new = _mm512_mask_blend_epi32(
+                finalmask, current_indices, min_indices);
+        const __m512 min_distances_new =
+                _mm512_mask_blend_ps(finalmask, distances, min_distances);
+
+        min_indices = min_indices_new;
+        min_distances = min_distances_new;
+    }
+
+    // grab min distance
+    min_dis = _mm512_reduce_min_ps(min_distances);
+    // blend
+    __mmask16 mindmask =
+            _mm512_cmpeq_ps_mask(min_distances, _mm512_set1_ps(min_dis));
+    // pick the max one
+    min_idx = _mm512_mask_reduce_max_epi32(mindmask, min_indices);
+
+    if (min_idx == -1) {
+        return -1;
+    }
+
+    if (vmin_out) {
+        *vmin_out = min_dis;
+    }
+    int ret = ids[min_idx];
+    ids[min_idx] = -1;
+    --nvalid;
+    return ret;
+}
+
+#elif __AVX2__
+
+int HNSW::MinimaxHeap::pop_min(float* vmin_out) {
+    assert(k > 0);
+    static_assert(
+            std::is_same<storage_idx_t, int32_t>::value,
+            "This code expects storage_idx_t to be int32_t");
+
+    int32_t min_idx = -1;
+    float min_dis = std::numeric_limits<float>::infinity();
+
+    size_t iii = 0;
+
+    __m256i min_indices = _mm256_setr_epi32(-1, -1, -1, -1, -1, -1, -1, -1);
+    __m256 min_distances =
+            _mm256_set1_ps(std::numeric_limits<float>::infinity());
+    __m256i current_indices = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
+    __m256i offset = _mm256_set1_epi32(8);
+
+    // The baseline version is available in non-AVX2 branch.
+
+    // The following loop tracks the rightmost index with the min distance.
+    // -1 index values are ignored.
+    const int k8 = (k / 8) * 8;
+    for (; iii < k8; iii += 8) {
+        __m256i indices =
+                _mm256_loadu_si256((const __m256i*)(ids.data() + iii));
+        __m256 distances = _mm256_loadu_ps(dis.data() + iii);
+
+        // This mask filters out -1 values among indices.
+        __m256i m1mask = _mm256_cmpgt_epi32(_mm256_setzero_si256(), indices);
+
+        __m256i dmask = _mm256_castps_si256(
+                _mm256_cmp_ps(min_distances, distances, _CMP_LT_OS));
+        __m256 finalmask = _mm256_castsi256_ps(_mm256_or_si256(m1mask, dmask));
+
+        const __m256i min_indices_new = _mm256_castps_si256(_mm256_blendv_ps(
+                _mm256_castsi256_ps(current_indices),
+                _mm256_castsi256_ps(min_indices),
+                finalmask));
+
+        const __m256 min_distances_new =
+                _mm256_blendv_ps(distances, min_distances, finalmask);
+
+        min_indices = min_indices_new;
+        min_distances = min_distances_new;
+
+        current_indices = _mm256_add_epi32(current_indices, offset);
+    }
+
+    // Vectorizing is doable, but is not practical
+    int32_t vidx8[8];
+    float vdis8[8];
+    _mm256_storeu_ps(vdis8, min_distances);
+    _mm256_storeu_si256((__m256i*)vidx8, min_indices);
+
+    for (size_t j = 0; j < 8; j++) {
+        if (min_dis > vdis8[j] || (min_dis == vdis8[j] && min_idx < vidx8[j])) {
+            min_idx = vidx8[j];
+            min_dis = vdis8[j];
+        }
+    }
+
+    // process last values. Vectorizing is doable, but is not practical
+    for (; iii < k; iii++) {
+        if (ids[iii] != -1 && dis[iii] <= min_dis) {
+            min_dis = dis[iii];
+            min_idx = iii;
+        }
+    }
+
+    if (min_idx == -1) {
+        return -1;
+    }
+
+    if (vmin_out) {
+        *vmin_out = min_dis;
+    }
+    int ret = ids[min_idx];
+    ids[min_idx] = -1;
+    --nvalid;
+    return ret;
+}
+
+#else
+
+// baseline non-vectorized version
 int HNSW::MinimaxHeap::pop_min(float* vmin_out) {
     assert(k > 0);
     // returns min. This is an O(n) operation
     int i = k - 1;
     while (i >= 0) {
-        if (ids[i] != -1)
+        if (ids[i] != -1) {
             break;
+        }
         i--;
     }
-    if (i == -1)
+    if (i == -1) {
         return -1;
+    }
     int imin = i;
     float vmin = dis[i];
     i--;
@@ -873,14 +1290,16 @@ int HNSW::MinimaxHeap::pop_min(float* vmin_out) {
         }
         i--;
     }
-    if (vmin_out)
+    if (vmin_out) {
         *vmin_out = vmin;
+    }
     int ret = ids[imin];
     ids[imin] = -1;
     --nvalid;
 
     return ret;
 }
+#endif
 
 int HNSW::MinimaxHeap::count_below(float thresh) {
     int n_below = 0;