faiss 0.5.3 → 0.6.1

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

@@ -7,22 +7,17 @@
 
 #include <faiss/impl/HNSW.h>
 
+#include <cinttypes>
 #include <cstddef>
+#include <cstdlib>
 
 #include <faiss/IndexHNSW.h>
 
-#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>
-
-#ifdef __AVX2__
-#include <immintrin.h>
-
-#include <limits>
-#include <type_traits>
-#endif
+#include <faiss/impl/VisitedTable.h>
+#include <faiss/impl/hnsw/MinimaxHeap.h>
 
 namespace faiss {
 
@@ -31,7 +26,8 @@ namespace faiss {
  **************************************************************/
 
 int HNSW::nb_neighbors(int layer_no) const {
-    FAISS_THROW_IF_NOT(layer_no + 1 < cum_nneighbor_per_level.size());
+    FAISS_THROW_IF_NOT(
+            static_cast<size_t>(layer_no + 1) < cum_nneighbor_per_level.size());
     return cum_nneighbor_per_level[layer_no + 1] -
             cum_nneighbor_per_level[layer_no];
 }
@@ -39,17 +35,21 @@ int HNSW::nb_neighbors(int layer_no) const {
 void HNSW::set_nb_neighbors(int level_no, int n) {
     FAISS_THROW_IF_NOT(levels.size() == 0);
     int cur_n = nb_neighbors(level_no);
-    for (int i = level_no + 1; i < cum_nneighbor_per_level.size(); i++) {
+    for (size_t i = level_no + 1; i < cum_nneighbor_per_level.size(); i++) {
         cum_nneighbor_per_level[i] += n - cur_n;
     }
 }
 
 int HNSW::cum_nb_neighbors(int layer_no) const {
+    FAISS_CHECK_RANGE_DEBUG(layer_no, 0, (int)cum_nneighbor_per_level.size());
     return cum_nneighbor_per_level[layer_no];
 }
 
 void HNSW::neighbor_range(idx_t no, int layer_no, size_t* begin, size_t* end)
         const {
+    FAISS_CHECK_RANGE_DEBUG(no, 0, (idx_t)offsets.size());
+    FAISS_CHECK_RANGE_DEBUG(
+            layer_no, 0, (int)cum_nneighbor_per_level.size() - 1);
     size_t o = offsets[no];
     *begin = o + cum_nb_neighbors(layer_no);
     *end = o + cum_nb_neighbors(layer_no + 1);
@@ -63,7 +63,7 @@ HNSW::HNSW(int M) : rng(12345) {
 int HNSW::random_level() {
     double f = rng.rand_float();
     // could be a bit faster with bisection
-    for (int level = 0; level < assign_probas.size(); level++) {
+    for (size_t level = 0; level < assign_probas.size(); level++) {
         if (f < assign_probas[level]) {
             return level;
         }
@@ -88,7 +88,7 @@ void HNSW::set_default_probas(int M, float levelMult) {
 }
 
 void HNSW::clear_neighbor_tables(int level) {
-    for (int i = 0; i < levels.size(); i++) {
+    for (size_t i = 0; i < levels.size(); i++) {
         size_t begin, end;
         neighbor_range(i, level, &begin, &end);
         for (size_t j = begin; j < end; j++) {
@@ -107,14 +107,15 @@ void HNSW::reset() {
 }
 
 void HNSW::print_neighbor_stats(int level) const {
-    FAISS_THROW_IF_NOT(level < cum_nneighbor_per_level.size());
+    FAISS_THROW_IF_NOT(
+            static_cast<size_t>(level) < cum_nneighbor_per_level.size());
     printf("stats on level %d, max %d neighbors per vertex:\n",
            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)
-    for (int i = 0; i < levels.size(); i++) {
+    for (idx_t i = 0; i < static_cast<idx_t>(levels.size()); i++) {
         if (levels[i] > level) {
             n_node++;
             size_t begin, end;
@@ -126,7 +127,7 @@ void HNSW::print_neighbor_stats(int level) const {
                 }
                 neighset.insert(neighbors[j]);
             }
-            int n_neigh = neighset.size();
+            size_t n_neigh = neighset.size();
             int n_common = 0;
             int n_reciprocal = 0;
             for (size_t j = begin; j < end; j++) {
@@ -175,7 +176,7 @@ void HNSW::fill_with_random_links(size_t n) {
 
     for (int level = max_level_2 - 1; level >= 0; --level) {
         std::vector<int> elts;
-        for (int i = 0; i < n; i++) {
+        for (size_t i = 0; i < n; i++) {
             if (levels[i] > level) {
                 elts.push_back(i);
             }
@@ -186,10 +187,10 @@ void HNSW::fill_with_random_links(size_t n) {
             continue;
         }
 
-        for (int ii = 0; ii < elts.size(); ii++) {
+        for (size_t ii = 0; ii < elts.size(); ii++) {
             int i = elts[ii];
             size_t begin, end;
-            neighbor_range(i, 0, &begin, &end);
+            neighbor_range(i, level, &begin, &end);
             for (size_t j = begin; j < end; j++) {
                 int other = 0;
                 do {
@@ -209,14 +210,14 @@ int HNSW::prepare_level_tab(size_t n, bool preset_levels) {
         FAISS_ASSERT(n0 + n == levels.size());
     } else {
         FAISS_ASSERT(n0 == levels.size());
-        for (int i = 0; i < n; i++) {
+        for (size_t i = 0; i < n; i++) {
             int pt_level = random_level();
             levels.push_back(pt_level + 1);
         }
     }
 
     int max_level_2 = 0;
-    for (int i = 0; i < n; i++) {
+    for (size_t i = 0; i < n; i++) {
         int pt_level = levels[i + n0] - 1;
         if (pt_level > max_level_2) {
             max_level_2 = pt_level;
@@ -236,7 +237,7 @@ void HNSW::shrink_neighbor_list(
         DistanceComputer& qdis,
         std::priority_queue<NodeDistFarther>& input,
         std::vector<NodeDistFarther>& output,
-        int max_size,
+        size_t max_size,
         bool keep_max_size_level0) {
     // This prevents number of neighbors at
     // level 0 from being shrunk to less than 2 * M.
@@ -261,7 +262,7 @@ void HNSW::shrink_neighbor_list(
 
         if (good) {
             output.push_back(v1);
-            if (output.size() >= max_size) {
+            if (output.size() >= static_cast<size_t>(max_size)) {
                 return;
             }
         } else if (keep_max_size_level0) {
@@ -269,7 +270,8 @@ void HNSW::shrink_neighbor_list(
         }
     }
     size_t idx = 0;
-    while (keep_max_size_level0 && (output.size() < max_size) &&
+    while (keep_max_size_level0 &&
+           (output.size() < static_cast<size_t>(max_size)) &&
            (idx < outsiders.size())) {
         output.push_back(outsiders[idx++]);
     }
@@ -289,9 +291,9 @@ using NodeDistFarther = HNSW::NodeDistFarther;
 void shrink_neighbor_list(
         DistanceComputer& qdis,
         std::priority_queue<NodeDistCloser>& resultSet1,
-        int max_size,
+        size_t max_size,
         bool keep_max_size_level0 = false) {
-    if (resultSet1.size() < max_size) {
+    if (resultSet1.size() < static_cast<size_t>(max_size)) {
         return;
     }
     std::priority_queue<NodeDistFarther> resultSet;
@@ -344,7 +346,9 @@ void add_link(
         resultSet.emplace(qdis.symmetric_dis(src, neigh), neigh);
     }
 
-    shrink_neighbor_list(qdis, resultSet, end - begin, keep_max_size_level0);
+    size_t max_size = end - begin;
+    max_size -= max_size * std::clamp(hnsw.prune_headroom, 0.0f, 0.5f);
+    shrink_neighbor_list(qdis, resultSet, max_size, keep_max_size_level0);
 
     // ...and back
     size_t i = begin;
@@ -407,19 +411,19 @@ void search_neighbors_to_add(
                 if (nodeId < 0) {
                     break;
                 }
-                if (vt.get(nodeId)) {
+                if (!vt.set(nodeId)) {
                     continue;
                 }
-                vt.set(nodeId);
 
                 float dis = qdis(nodeId);
                 NodeDistFarther evE1(dis, nodeId);
 
-                if (results.size() < hnsw.efConstruction ||
+                if (results.size() < static_cast<size_t>(hnsw.efConstruction) ||
                     results.top().d > dis) {
                     results.emplace(dis, nodeId);
                     candidates.emplace(dis, nodeId);
-                    if (results.size() > hnsw.efConstruction) {
+                    if (results.size() >
+                        static_cast<size_t>(hnsw.efConstruction)) {
                         results.pop();
                     }
                 }
@@ -430,11 +434,12 @@ void search_neighbors_to_add(
             // 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 ||
+                if (results.size() < static_cast<size_t>(hnsw.efConstruction) ||
                     results.top().d > dis) {
                     results.emplace(dis, idx);
                     candidates.emplace(dis, idx);
-                    if (results.size() > hnsw.efConstruction) {
+                    if (results.size() >
+                        static_cast<size_t>(hnsw.efConstruction)) {
                         results.pop();
                     }
                 }
@@ -448,10 +453,9 @@ void search_neighbors_to_add(
                 if (nodeId < 0) {
                     break;
                 }
-                if (vt.get(nodeId)) {
+                if (!vt.set(nodeId)) {
                     continue;
                 }
-                vt.set(nodeId);
 
                 buffered_ids[n_buffered] = nodeId;
                 n_buffered += 1;
@@ -477,7 +481,7 @@ void search_neighbors_to_add(
             }
 
             // process leftovers
-            for (size_t icnt = 0; icnt < n_buffered; icnt++) {
+            for (int icnt = 0; icnt < n_buffered; icnt++) {
                 float dis = qdis(buffered_ids[icnt]);
                 update_with_candidate(buffered_ids[icnt], dis);
             }
@@ -495,7 +499,7 @@ void HNSW::add_links_starting_from(
         storage_idx_t nearest,
         float d_nearest,
         int level,
-        omp_lock_t* locks,
+        LockVector& locks,
         VisitedTable& vt,
         bool keep_max_size_level0) {
     std::priority_queue<NodeDistCloser> link_targets;
@@ -517,13 +521,13 @@ void HNSW::add_links_starting_from(
         link_targets.pop();
     }
 
-    omp_unset_lock(&locks[pt_id]);
+    locks.unlock(pt_id);
     for (storage_idx_t other_id : neighbors_to_add) {
-        omp_set_lock(&locks[other_id]);
+        locks.lock(other_id);
         add_link(*this, ptdis, other_id, pt_id, level, keep_max_size_level0);
-        omp_unset_lock(&locks[other_id]);
+        locks.unlock(other_id);
     }
-    omp_set_lock(&locks[pt_id]);
+    locks.lock(pt_id);
 }
 
 /**************************************************************
@@ -534,19 +538,19 @@ void HNSW::add_with_locks(
         DistanceComputer& ptdis,
         int pt_level,
         int pt_id,
-        std::vector<omp_lock_t>& locks,
+        LockVector& locks,
         VisitedTable& vt,
         bool keep_max_size_level0) {
-    //  greedy search on upper levels
-
-    storage_idx_t nearest;
+    storage_idx_t nearest = entry_point;
+    if (nearest == -1) { // avoid locking after the first point.
 #pragma omp critical
-    {
-        nearest = entry_point;
-
-        if (nearest == -1) {
+        if (entry_point == -1) { // double-check under lock.
             max_level = pt_level;
             entry_point = pt_id;
+            // leave nearest = -1 to trigger early exit after critical block.
+        } else {
+            // else: Another thread set the entry point.
+            nearest = entry_point;
         }
     }
 
@@ -554,11 +558,12 @@ void HNSW::add_with_locks(
         return;
     }
 
-    omp_set_lock(&locks[pt_id]);
+    locks.lock(pt_id);
 
     int level = max_level; // level at which we start adding neighbors
     float d_nearest = ptdis(nearest);
 
+    //  greedy search on upper levels
     for (; level > pt_level; level--) {
         greedy_update_nearest(*this, ptdis, level, nearest, d_nearest);
     }
@@ -570,16 +575,19 @@ void HNSW::add_with_locks(
                 nearest,
                 d_nearest,
                 level,
-                locks.data(),
+                locks,
                 vt,
                 keep_max_size_level0);
     }
 
-    omp_unset_lock(&locks[pt_id]);
+    locks.unlock(pt_id);
 
-    if (pt_level > max_level) {
-        max_level = pt_level;
-        entry_point = pt_id;
+#pragma omp critical
+    {
+        if (pt_level > max_level) {
+            max_level = pt_level;
+            entry_point = pt_id;
+        }
     }
 }
 
@@ -587,7 +595,6 @@ void HNSW::add_with_locks(
  * Searching
  **************************************************************/
 
-using MinimaxHeap = HNSW::MinimaxHeap;
 using Node = HNSW::Node;
 using C = HNSW::C;
 
@@ -616,7 +623,7 @@ static inline void extract_search_params(
 int search_from_candidates(
         const HNSW& hnsw,
         DistanceComputer& qdis,
-        ResultHandler<C>& res,
+        ResultHandler& res,
         MinimaxHeap& candidates,
         VisitedTable& vt,
         HNSWStats& stats,
@@ -675,7 +682,7 @@ int search_from_candidates(
                 break;
             }
 
-            prefetch_L2(vt.visited.data() + v1);
+            vt.prefetch(v1);
             jmax += 1;
         }
 
@@ -699,10 +706,8 @@ int search_from_candidates(
         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;
+            counter += vt.set(v1) ? 1 : 0;
 
             if (counter == 4) {
                 float dis[4];
@@ -726,7 +731,7 @@ int search_from_candidates(
             }
         }
 
-        for (size_t icnt = 0; icnt < counter; icnt++) {
+        for (int icnt = 0; icnt < counter; icnt++) {
             float dis = qdis(saved_j[icnt]);
             add_to_heap(saved_j[icnt], dis);
 
@@ -755,7 +760,7 @@ int search_from_candidates_panorama(
         const HNSW& hnsw,
         const IndexHNSW* index,
         DistanceComputer& qdis,
-        ResultHandler<C>& res,
+        ResultHandler& res,
         MinimaxHeap& candidates,
         VisitedTable& vt,
         HNSWStats& stats,
@@ -802,16 +807,15 @@ int search_from_candidates_panorama(
     std::vector<float> exact_distances(M);
 
     const float* query = flat_codes_qdis->q;
-    std::vector<float> query_cum_sums(panorama_index->num_panorama_levels + 1);
-    IndexHNSWFlatPanorama::compute_cum_sums(
-            query,
-            query_cum_sums.data(),
-            panorama_index->d,
-            panorama_index->num_panorama_levels,
-            panorama_index->panorama_level_width);
+    std::vector<float> query_cum_sums(panorama_index->pano.n_levels + 1);
+    panorama_index->pano.compute_query_cum_sums(query, query_cum_sums.data());
     float query_norm_sq = query_cum_sums[0] * query_cum_sums[0];
 
     int nstep = 0;
+    const size_t d = static_cast<size_t>(panorama_index->d);
+
+    PanoramaStats local_pano_stats;
+    local_pano_stats.reset();
 
     while (candidates.size() > 0) {
         float d0 = 0;
@@ -847,21 +851,20 @@ int search_from_candidates_panorama(
                     query_norm_sq + cum_sums_v1[0] * cum_sums_v1[0];
 
             bool is_selected = !sel || sel->is_member(v1);
-            initial_size += is_selected && !vt.get(v1) ? 1 : 0;
-
-            vt.set(v1);
+            initial_size += is_selected && vt.set(v1) ? 1 : 0;
         }
 
+        local_pano_stats.total_dims += initial_size * d;
         size_t batch_size = initial_size;
         size_t curr_panorama_level = 0;
-        const size_t num_panorama_levels = panorama_index->num_panorama_levels;
+        const size_t num_panorama_levels = panorama_index->pano.n_levels;
         while (curr_panorama_level < num_panorama_levels && batch_size > 0) {
             float query_cum_norm = query_cum_sums[curr_panorama_level + 1];
 
-            const size_t panorama_level_width =
-                    panorama_index->panorama_level_width;
-            size_t start_dim = curr_panorama_level * panorama_level_width;
-            size_t end_dim = (curr_panorama_level + 1) * panorama_level_width;
+            size_t start_dim = curr_panorama_level *
+                    panorama_index->pano.level_width_floats;
+            size_t end_dim = (curr_panorama_level + 1) *
+                    panorama_index->pano.level_width_floats;
             end_dim = std::min(end_dim, static_cast<size_t>(panorama_index->d));
 
             size_t i = 0;
@@ -967,6 +970,8 @@ int search_from_candidates_panorama(
                 }
             }
 
+            local_pano_stats.total_dims_scanned +=
+                    batch_size * (end_dim - start_dim);
             batch_size = next_batch_size;
             curr_panorama_level++;
         }
@@ -975,6 +980,7 @@ int search_from_candidates_panorama(
         for (size_t i = 0; i < batch_size; i++) {
             idx_t idx = index_array[i];
             if (res.add_result(exact_distances[i], idx)) {
+                threshold = res.threshold;
                 nres += 1;
             }
             candidates.push(idx, exact_distances[i]);
@@ -995,9 +1001,22 @@ int search_from_candidates_panorama(
         stats.nhops += nstep;
     }
 
+    indexPanorama_stats.add(local_pano_stats);
     return nres;
 }
 
+template <typename T, typename Container, typename Compare>
+void reservePriorityQueue(
+        std::priority_queue<T, Container, Compare>& q,
+        std::size_t size) {
+    struct Access : std::priority_queue<T, Container, Compare> {
+        using std::priority_queue<T, Container, Compare>::c;
+    };
+    Access access{std::move(q)};
+    access.c.reserve(size);
+    q = std::move(access);
+}
+
 std::priority_queue<HNSW::Node> search_from_candidate_unbounded(
         const HNSW& hnsw,
         const Node& node,
@@ -1007,7 +1026,10 @@ std::priority_queue<HNSW::Node> search_from_candidate_unbounded(
         HNSWStats& stats) {
     int ndis = 0;
     std::priority_queue<Node> top_candidates;
+    reservePriorityQueue(top_candidates, ef);
+
     std::priority_queue<Node, std::vector<Node>, std::greater<Node>> candidates;
+    reservePriorityQueue(candidates, ef);
 
     top_candidates.push(node);
     candidates.push(node);
@@ -1037,7 +1059,7 @@ std::priority_queue<HNSW::Node> search_from_candidate_unbounded(
                 break;
             }
 
-            prefetch_L2(vt->visited.data() + v1);
+            vt->prefetch(v1);
             jmax += 1;
         }
 
@@ -1059,10 +1081,8 @@ std::priority_queue<HNSW::Node> search_from_candidate_unbounded(
         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;
+            counter += vt->set(v1) ? 1 : 0;
 
             if (counter == 4) {
                 float dis[4];
@@ -1086,7 +1106,7 @@ std::priority_queue<HNSW::Node> search_from_candidate_unbounded(
             }
         }
 
-        for (size_t icnt = 0; icnt < counter; icnt++) {
+        for (int icnt = 0; icnt < counter; icnt++) {
             float dis = qdis(saved_j[icnt]);
             add_to_heap(saved_j[icnt], dis);
 
@@ -1166,7 +1186,7 @@ HNSWStats greedy_update_nearest(
         }
 
         // process leftovers
-        for (size_t icnt = 0; icnt < n_buffered; icnt++) {
+        for (int icnt = 0; icnt < n_buffered; icnt++) {
             float dis = qdis(buffered_ids[icnt]);
             update_with_candidate(buffered_ids[icnt], dis);
         }
@@ -1182,12 +1202,11 @@ HNSWStats greedy_update_nearest(
 }
 
 namespace {
-using MinimaxHeap = HNSW::MinimaxHeap;
 using Node = HNSW::Node;
 using C = HNSW::C;
 
 // just used as a lower bound for the minmaxheap, but it is set for heap search
-int extract_k_from_ResultHandler(ResultHandler<C>& res) {
+int extract_k_from_ResultHandler(ResultHandler& res) {
     using RH = HeapBlockResultHandler<C>;
     if (auto hres = dynamic_cast<RH::SingleResultHandler*>(&res)) {
         return hres->k;
@@ -1200,7 +1219,7 @@ int extract_k_from_ResultHandler(ResultHandler<C>& res) {
 HNSWStats HNSW::search(
         DistanceComputer& qdis,
         const IndexHNSW* index,
-        ResultHandler<C>& res,
+        ResultHandler& res,
         VisitedTable& vt,
         const SearchParameters* params) const {
     HNSWStats stats;
@@ -1210,12 +1229,12 @@ HNSWStats HNSW::search(
     int k = extract_k_from_ResultHandler(res);
 
     bool bounded_queue = this->search_bounded_queue;
-    int efSearch = this->efSearch;
+    int cur_efSearch = this->efSearch;
     if (params) {
         if (const SearchParametersHNSW* hnsw_params =
                     dynamic_cast<const SearchParametersHNSW*>(params)) {
             bounded_queue = hnsw_params->bounded_queue;
-            efSearch = hnsw_params->efSearch;
+            cur_efSearch = hnsw_params->efSearch;
         }
     }
 
@@ -1229,7 +1248,7 @@ HNSWStats HNSW::search(
         stats.combine(local_stats);
     }
 
-    int ef = std::max(efSearch, k);
+    int ef = std::max(cur_efSearch, k);
     if (bounded_queue) { // this is the most common branch, for now we only
                          // support Panorama search in this branch
         MinimaxHeap candidates(ef);
@@ -1257,7 +1276,7 @@ HNSWStats HNSW::search(
                 search_from_candidate_unbounded(
                         *this, Node(d_nearest, nearest), qdis, ef, &vt, stats);
 
-        while (top_candidates.size() > k) {
+        while (top_candidates.size() > static_cast<size_t>(k)) {
             top_candidates.pop();
         }
 
@@ -1277,7 +1296,7 @@ HNSWStats HNSW::search(
 
 void HNSW::search_level_0(
         DistanceComputer& qdis,
-        ResultHandler<C>& res,
+        ResultHandler& res,
         idx_t nprobe,
         const storage_idx_t* nearest_i,
         const float* nearest_d,
@@ -1287,11 +1306,11 @@ void HNSW::search_level_0(
         const SearchParameters* params) const {
     const HNSW& hnsw = *this;
 
-    auto efSearch = hnsw.efSearch;
+    auto cur_efSearch = hnsw.efSearch;
     if (params) {
         if (const SearchParametersHNSW* hnsw_params =
                     dynamic_cast<const SearchParametersHNSW*>(params)) {
-            efSearch = hnsw_params->efSearch;
+            cur_efSearch = hnsw_params->efSearch;
         }
     }
 
@@ -1300,7 +1319,7 @@ void HNSW::search_level_0(
     if (search_type == 1) {
         int nres = 0;
 
-        for (int j = 0; j < nprobe; j++) {
+        for (idx_t j = 0; j < nprobe; j++) {
             storage_idx_t cj = nearest_i[j];
 
             if (cj < 0) {
@@ -1311,7 +1330,7 @@ void HNSW::search_level_0(
                 continue;
             }
 
-            int candidates_size = std::max(efSearch, k);
+            int candidates_size = std::max(cur_efSearch, k);
             MinimaxHeap candidates(candidates_size);
 
             candidates.push(cj, nearest_d[j]);
@@ -1329,11 +1348,11 @@ void HNSW::search_level_0(
             nres = std::min(nres, candidates_size);
         }
     } else if (search_type == 2) {
-        int candidates_size = std::max(efSearch, int(k));
+        int candidates_size = std::max(cur_efSearch, int(k));
         candidates_size = std::max(candidates_size, int(nprobe));
 
         MinimaxHeap candidates(candidates_size);
-        for (int j = 0; j < nprobe; j++) {
+        for (idx_t j = 0; j < nprobe; j++) {
             storage_idx_t cj = nearest_i[j];
 
             if (cj < 0) {
@@ -1380,257 +1399,4 @@ void HNSW::permute_entries(const idx_t* map) {
     neighbors = std::move(new_neighbors);
 }
 
-/**************************************************************
- * MinimaxHeap
- **************************************************************/
-
-void HNSW::MinimaxHeap::push(storage_idx_t i, float v) {
-    if (k == n) {
-        if (v >= dis[0]) {
-            return;
-        }
-        if (ids[0] != -1) {
-            --nvalid;
-        }
-        faiss::heap_pop<HC>(k--, dis.data(), ids.data());
-    }
-    faiss::heap_push<HC>(++k, dis.data(), ids.data(), v, i);
-    ++nvalid;
-}
-
-float HNSW::MinimaxHeap::max() const {
-    return dis[0];
-}
-
-int HNSW::MinimaxHeap::size() const {
-    return nvalid;
-}
-
-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) {
-            break;
-        }
-        i--;
-    }
-    if (i == -1) {
-        return -1;
-    }
-    int imin = i;
-    float vmin = dis[i];
-    i--;
-    while (i >= 0) {
-        if (ids[i] != -1 && dis[i] < vmin) {
-            vmin = dis[i];
-            imin = i;
-        }
-        i--;
-    }
-    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;
-    for (int i = 0; i < k; i++) {
-        if (dis[i] < thresh) {
-            n_below++;
-        }
-    }
-
-    return n_below;
-}
-
 } // namespace faiss