faiss 0.3.0 → 0.3.1

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

@@ -5,8 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 #include <faiss/impl/HNSW.h>
 
 #include <string>
@@ -14,6 +12,17 @@
 #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 {
 
@@ -212,7 +221,7 @@ int HNSW::prepare_level_tab(size_t n, bool preset_levels) {
     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.
  */
@@ -503,17 +512,15 @@ void HNSW::add_with_locks(
  **************************************************************/
 
 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,
@@ -529,15 +536,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 +571,86 @@ int search_from_candidates(
         size_t begin, end;
         hnsw.neighbor_range(v0, level, &begin, &end);
 
+        // // baseline version
+        // for (size_t j = begin; j < end; j++) {
+        //     int v1 = hnsw.neighbors[j];
+        //     if (v1 < 0)
+        //         break;
+        //     if (vt.get(v1)) {
+        //         continue;
+        //     }
+        //     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);
+        //         }
+        //     }
+        //     candidates.push(v1, d);
+        // }
+
+        // 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];
+
+        ndis += jmax - begin;
+        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;
+                    }
+                }
             }
+            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]);
                 }
+
+                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);
         }
 
         nstep++;
@@ -630,29 +700,92 @@ 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) {
+        // // baseline version
+        // for (size_t j = begin; j < end; ++j) {
+        //     int v1 = hnsw.neighbors[j];
+        //
+        //     if (v1 < 0) {
+        //         break;
+        //     }
+        //     if (vt->get(v1)) {
+        //         continue;
+        //     }
+        //
+        //     vt->set(v1);
+        //
+        //     float d1 = qdis(v1);
+        //     ++ndis;
+        //
+        //     if (top_candidates.top().first > d1 ||
+        //         top_candidates.size() < ef) {
+        //         candidates.emplace(d1, v1);
+        //         top_candidates.emplace(d1, v1);
+        //
+        //         if (top_candidates.size() > ef) {
+        //             top_candidates.pop();
+        //         }
+        //     }
+        // }
+
+        // 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;
+        }
+
+        int counter = 0;
+        size_t saved_j[4];
 
-            float d1 = qdis(v1);
-            ++ndis;
+        ndis += jmax - begin;
 
-            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]);
+                }
+
+                counter = 0;
+            }
+        }
+
+        for (size_t icnt = 0; icnt < counter; icnt++) {
+            float dis = qdis(saved_j[icnt]);
+            add_to_heap(saved_j[icnt], dis);
         }
     }
 
@@ -665,19 +798,28 @@ std::priority_queue<HNSW::Node> search_from_candidate_unbounded(
     return top_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;
+}
+
 } // anonymous namespace
 
 HNSWStats HNSW::search(
         DistanceComputer& qdis,
-        int k,
-        idx_t* I,
-        float* D,
+        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 (upper_beam == 1) {
         //  greedy search on upper levels
         storage_idx_t nearest = entry_point;
@@ -687,14 +829,14 @@ HNSWStats HNSW::search(
             greedy_update_nearest(*this, qdis, level, nearest, d_nearest);
         }
 
-        int ef = std::max(efSearch, k);
+        int ef = std::max(params ? params->efSearch : efSearch, k);
         if (search_bounded_queue) { // this is the most common branch
             MinimaxHeap candidates(ef);
 
             candidates.push(nearest, d_nearest);
 
             search_from_candidates(
-                    *this, qdis, k, I, D, candidates, vt, stats, 0, 0, params);
+                    *this, qdis, res, candidates, vt, stats, 0, 0, params);
         } else {
             std::priority_queue<Node> top_candidates =
                     search_from_candidate_unbounded(
@@ -709,12 +851,11 @@ HNSWStats HNSW::search(
                 top_candidates.pop();
             }
 
-            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);
+                res.add_result(d, label);
                 top_candidates.pop();
             }
         }
@@ -728,6 +869,10 @@ HNSWStats HNSW::search(
         std::vector<idx_t> I_to_next(candidates_size);
         std::vector<float> D_to_next(candidates_size);
 
+        HeapBlockResultHandler<C> block_resh(
+                1, D_to_next.data(), I_to_next.data(), candidates_size);
+        HeapBlockResultHandler<C>::SingleResultHandler resh(block_resh);
+
         int nres = 1;
         I_to_next[0] = entry_point;
         D_to_next[0] = qdis(entry_point);
@@ -743,18 +888,12 @@ HNSWStats HNSW::search(
 
             if (level == 0) {
                 nres = search_from_candidates(
-                        *this, qdis, k, I, D, candidates, vt, stats, 0);
+                        *this, qdis, res, candidates, vt, stats, 0);
             } else {
+                resh.begin(0);
                 nres = search_from_candidates(
-                        *this,
-                        qdis,
-                        candidates_size,
-                        I_to_next.data(),
-                        D_to_next.data(),
-                        candidates,
-                        vt,
-                        stats,
-                        level);
+                        *this, qdis, resh, candidates, vt, stats, level);
+                resh.end();
             }
             vt.advance();
         }
@@ -765,9 +904,7 @@ HNSWStats HNSW::search(
 
 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,
@@ -775,7 +912,7 @@ void HNSW::search_level_0(
         HNSWStats& search_stats,
         VisitedTable& vt) const {
     const HNSW& hnsw = *this;
-
+    int k = extract_k_from_ResultHandler(res);
     if (search_type == 1) {
         int nres = 0;
 
@@ -788,22 +925,13 @@ void HNSW::search_level_0(
             if (vt.get(cj))
                 continue;
 
-            int candidates_size = std::max(hnsw.efSearch, int(k));
+            int candidates_size = std::max(hnsw.efSearch, k);
             MinimaxHeap candidates(candidates_size);
 
             candidates.push(cj, nearest_d[j]);
 
             nres = search_from_candidates(
-                    hnsw,
-                    qdis,
-                    k,
-                    idxi,
-                    simi,
-                    candidates,
-                    vt,
-                    search_stats,
-                    0,
-                    nres);
+                    hnsw, qdis, res, candidates, vt, search_stats, 0, nres);
         }
     } else if (search_type == 2) {
         int candidates_size = std::max(hnsw.efSearch, int(k));
@@ -819,10 +947,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);
     }
 }
 
+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 +1013,105 @@ void HNSW::MinimaxHeap::clear() {
     nvalid = k = 0;
 }
 
+#ifdef __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 +1122,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;

data/vendor/faiss/faiss/impl/HNSW.h

@@ -5,8 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 #pragma once
 
 #include <queue>
@@ -42,6 +40,8 @@ namespace faiss {
 struct VisitedTable;
 struct DistanceComputer; // from AuxIndexStructures
 struct HNSWStats;
+template <class C>
+struct ResultHandler;
 
 struct SearchParametersHNSW : SearchParameters {
     int efSearch = 16;
@@ -54,6 +54,9 @@ struct HNSW {
     /// internal storage of vectors (32 bits: this is expensive)
     using storage_idx_t = int32_t;
 
+    // for now we do only these distances
+    using C = CMax<float, int64_t>;
+
     typedef std::pair<float, storage_idx_t> Node;
 
     /** Heap structure that allows fast
@@ -195,18 +198,14 @@ struct HNSW {
     /// search interface for 1 point, single thread
     HNSWStats search(
             DistanceComputer& qdis,
-            int k,
-            idx_t* I,
-            float* D,
+            ResultHandler<C>& res,
             VisitedTable& vt,
             const SearchParametersHNSW* params = nullptr) const;
 
     /// search only in level 0 from a given vertex
     void 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,
@@ -226,6 +225,8 @@ struct HNSW {
             std::priority_queue<NodeDistFarther>& input,
             std::vector<NodeDistFarther>& output,
             int max_size);
+
+    void permute_entries(const idx_t* map);
 };
 
 struct HNSWStats {

data/vendor/faiss/faiss/impl/IDSelector.h

@@ -10,7 +10,7 @@
 #include <unordered_set>
 #include <vector>
 
-#include <faiss/Index.h>
+#include <faiss/MetricType.h>
 
 /** IDSelector is intended to define a subset of vectors to handle (for removal
  * or as subset to search) */
@@ -140,7 +140,7 @@ struct IDSelectorAnd : IDSelector {
             : lhs(lhs), rhs(rhs) {}
     bool is_member(idx_t id) const final {
         return lhs->is_member(id) && rhs->is_member(id);
-    };
+    }
     virtual ~IDSelectorAnd() {}
 };
 
@@ -153,7 +153,7 @@ struct IDSelectorOr : IDSelector {
             : lhs(lhs), rhs(rhs) {}
     bool is_member(idx_t id) const final {
         return lhs->is_member(id) || rhs->is_member(id);
-    };
+    }
     virtual ~IDSelectorOr() {}
 };
 
@@ -166,7 +166,7 @@ struct IDSelectorXOr : IDSelector {
             : lhs(lhs), rhs(rhs) {}
     bool is_member(idx_t id) const final {
         return lhs->is_member(id) ^ rhs->is_member(id);
-    };
+    }
     virtual ~IDSelectorXOr() {}
 };
 

data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp

@@ -628,7 +628,9 @@ void LocalSearchQuantizer::icm_encode_step(
                         {
                             size_t binary_idx = (other_m + 1) * M * K * K +
                                     m * K * K + code2 * K + code;
-                            _mm_prefetch(binaries + binary_idx, _MM_HINT_T0);
+                            _mm_prefetch(
+                                    (const char*)(binaries + binary_idx),
+                                    _MM_HINT_T0);
                         }
                     }
 #endif