faiss 0.5.2 → 0.6.0

data/vendor/faiss/faiss/impl/FaissAssert.h

@@ -12,6 +12,7 @@
 
 #include <faiss/impl/FaissException.h>
 #include <faiss/impl/platform_macros.h>
+#include <cinttypes>
 #include <cstdio>
 #include <cstdlib>
 #include <string>
@@ -69,7 +70,7 @@
 
 #define FAISS_THROW_MSG(MSG)                                   \
     do {                                                       \
-        throw faiss::FaissException(                           \
+        throw ::faiss::FaissException(                         \
                 MSG, __PRETTY_FUNCTION__, __FILE__, __LINE__); \
     } while (false)
 
@@ -79,7 +80,7 @@
         int __size = snprintf(nullptr, 0, FMT, __VA_ARGS__);   \
         __s.resize(__size + 1);                                \
         snprintf(&__s[0], __s.size(), FMT, __VA_ARGS__);       \
-        throw faiss::FaissException(                           \
+        throw ::faiss::FaissException(                         \
                 __s, __PRETTY_FUNCTION__, __FILE__, __LINE__); \
     } while (false)
 
@@ -110,4 +111,61 @@
         }                                                                 \
     } while (false)
 
+///
+/// Safe arithmetic
+///
+
+#include <limits>
+
+namespace faiss {
+
+/// Multiplication that throws on overflow instead of wrapping
+inline size_t mul_no_overflow(size_t a, size_t b, const char* context) {
+    if (a != 0 && b > (std::numeric_limits<size_t>::max)() / a) {
+        FAISS_THROW_FMT("integer overflow in %s: %zu * %zu", context, a, b);
+    }
+    return a * b;
+}
+
+/// Addition that throws on overflow instead of wrapping
+inline size_t add_no_overflow(size_t a, size_t b, const char* context) {
+    if (a > (std::numeric_limits<size_t>::max)() - b) {
+        FAISS_THROW_FMT("integer overflow in %s: %zu + %zu", context, a, b);
+    }
+    return a + b;
+}
+
+} // namespace faiss
+
+///
+/// Bounds checking
+///
+
+/// Check that val is in half-open range [lo, hi). Throws with the
+/// stringified expression, its value, and the bounds on failure.
+#define FAISS_CHECK_RANGE(val, lo, hi)                                         \
+    FAISS_THROW_IF_NOT_FMT(                                                    \
+            (int64_t)(val) >= (int64_t)(lo) && (int64_t)(val) < (int64_t)(hi), \
+            "%s (= %" PRId64 ") out of range [%" PRId64 ", %" PRId64 ")",      \
+            #val,                                                              \
+            (int64_t)(val),                                                    \
+            (int64_t)(lo),                                                     \
+            (int64_t)(hi))
+
+/// Debug-only variant of FAISS_CHECK_RANGE. Aborts in debug builds
+/// (when NDEBUG is not defined); compiled out in release builds to
+/// avoid overhead in hot paths.
+#ifndef NDEBUG
+#define FAISS_CHECK_RANGE_DEBUG(val, lo, hi)                                   \
+    FAISS_ASSERT_FMT(                                                          \
+            (int64_t)(val) >= (int64_t)(lo) && (int64_t)(val) < (int64_t)(hi), \
+            "%s (= %" PRId64 ") out of range [%" PRId64 ", %" PRId64 ")",      \
+            #val,                                                              \
+            (int64_t)(val),                                                    \
+            (int64_t)(lo),                                                     \
+            (int64_t)(hi))
+#else
+#define FAISS_CHECK_RANGE_DEBUG(val, lo, hi) ((void)0)
+#endif
+
 #endif

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

@@ -7,15 +7,15 @@
 
 #include <faiss/impl/HNSW.h>
 
+#include <cinttypes>
 #include <cstddef>
 
 #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>
+#include <faiss/impl/VisitedTable.h>
 
 #ifdef __AVX2__
 #include <immintrin.h>
@@ -45,11 +45,15 @@ void HNSW::set_nb_neighbors(int level_no, int 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);
@@ -407,10 +411,9 @@ 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);
@@ -448,10 +451,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;
@@ -616,7 +618,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 +677,7 @@ int search_from_candidates(
                 break;
             }
 
-            prefetch_L2(vt.visited.data() + v1);
+            vt.prefetch(v1);
             jmax += 1;
         }
 
@@ -699,10 +701,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];
@@ -755,7 +755,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,13 +802,8 @@ 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;
@@ -847,21 +842,19 @@ 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;
         }
 
         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;
@@ -1037,7 +1030,7 @@ std::priority_queue<HNSW::Node> search_from_candidate_unbounded(
                 break;
             }
 
-            prefetch_L2(vt->visited.data() + v1);
+            vt->prefetch(v1);
             jmax += 1;
         }
 
@@ -1059,10 +1052,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];
@@ -1187,7 +1178,7 @@ 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 +1191,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;
@@ -1277,7 +1268,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,

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

@@ -7,6 +7,7 @@
 
 #pragma once
 
+#include <optional>
 #include <queue>
 #include <vector>
 
@@ -45,8 +46,6 @@ struct IndexHNSWFlatPanorama;
 struct VisitedTable;
 struct DistanceComputer; // from AuxIndexStructures
 struct HNSWStats;
-template <class C>
-struct ResultHandler;
 
 struct SearchParametersHNSW : SearchParameters {
     int efSearch = 16;
@@ -153,6 +152,9 @@ struct HNSW {
     /// use Panorama progressive pruning in search
     bool is_panorama = false;
 
+    // See impl/VisitedTable.h.
+    std::optional<bool> use_visited_hashset;
+
     // methods that initialize the tree sizes
 
     /// initialize the assign_probas and cum_nneighbor_per_level to
@@ -212,14 +214,14 @@ struct HNSW {
     HNSWStats search(
             DistanceComputer& qdis,
             const IndexHNSW* index,
-            ResultHandler<C>& res,
+            ResultHandler& res,
             VisitedTable& vt,
             const SearchParameters* params = nullptr) const;
 
     /// search only in level 0 from a given vertex
     void search_level_0(
             DistanceComputer& qdis,
-            ResultHandler<C>& res,
+            ResultHandler& res,
             idx_t nprobe,
             const storage_idx_t* nearest_i,
             const float* nearest_d,
@@ -272,7 +274,7 @@ FAISS_API extern HNSWStats hnsw_stats;
 int search_from_candidates(
         const HNSW& hnsw,
         DistanceComputer& qdis,
-        ResultHandler<HNSW::C>& res,
+        ResultHandler& res,
         HNSW::MinimaxHeap& candidates,
         VisitedTable& vt,
         HNSWStats& stats,
@@ -288,7 +290,7 @@ int search_from_candidates_panorama(
         const HNSW& hnsw,
         const IndexHNSW* index,
         DistanceComputer& qdis,
-        ResultHandler<HNSW::C>& res,
+        ResultHandler& res,
         HNSW::MinimaxHeap& candidates,
         VisitedTable& vt,
         HNSWStats& stats,

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

@@ -18,6 +18,7 @@
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/utils.h>
 
@@ -684,22 +685,24 @@ void LocalSearchQuantizer::perturb_codes(
 void LocalSearchQuantizer::compute_binary_terms(float* binaries) const {
     LSQTimerScope scope(&lsq_timer, "compute_binary_terms");
 
+    with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for
-    for (int64_t m12 = 0; m12 < M * M; m12++) {
-        size_t m1 = m12 / M;
-        size_t m2 = m12 % M;
-
-        for (size_t code1 = 0; code1 < K; code1++) {
-            for (size_t code2 = 0; code2 < K; code2++) {
-                const float* c1 = codebooks.data() + m1 * K * d + code1 * d;
-                const float* c2 = codebooks.data() + m2 * K * d + code2 * d;
-                float ip = fvec_inner_product(c1, c2, d);
-                // binaries[m1, m2, code1, code2] = ip * 2
-                binaries[m1 * M * K * K + m2 * K * K + code1 * K + code2] =
-                        ip * 2;
+        for (int64_t m12 = 0; m12 < M * M; m12++) {
+            size_t m1 = m12 / M;
+            size_t m2 = m12 % M;
+
+            for (size_t code1 = 0; code1 < K; code1++) {
+                for (size_t code2 = 0; code2 < K; code2++) {
+                    const float* c1 = codebooks.data() + m1 * K * d + code1 * d;
+                    const float* c2 = codebooks.data() + m2 * K * d + code2 * d;
+                    float ip = fvec_inner_product<SL>(c1, c2, d);
+                    // binaries[m1, m2, code1, code2] = ip * 2
+                    binaries[m1 * M * K * K + m2 * K * K + code1 * K + code2] =
+                            ip * 2;
+                }
             }
         }
-    }
+    });
 }
 
 void LocalSearchQuantizer::compute_unary_terms(
@@ -760,23 +763,27 @@ float LocalSearchQuantizer::evaluate(
     std::vector<float> decoded_x(n * d, 0.0f);
     float obj = 0.0f;
 
-#pragma omp parallel for reduction(+ : obj)
-    for (int64_t i = 0; i < n; i++) {
-        const auto code = codes + i * M;
-        const auto decoded_i = decoded_x.data() + i * d;
-        for (size_t m = 0; m < M; m++) {
-            // c = codebooks[m, code[m]]
-            const auto c = codebooks.data() + m * K * d + code[m] * d;
-            fvec_add(d, decoded_i, c, decoded_i);
-        }
+    with_simd_level([&]<SIMDLevel SL>() {
+        float local_obj = 0.0f;
+#pragma omp parallel for reduction(+ : local_obj)
+        for (int64_t i = 0; i < n; i++) {
+            const auto code = codes + i * M;
+            const auto decoded_i = decoded_x.data() + i * d;
+            for (size_t m = 0; m < M; m++) {
+                // c = codebooks[m, code[m]]
+                const auto c = codebooks.data() + m * K * d + code[m] * d;
+                fvec_add(d, decoded_i, c, decoded_i);
+            }
 
-        float err = faiss::fvec_L2sqr(x + i * d, decoded_i, d);
-        obj += err;
+            float err = fvec_L2sqr<SL>(x + i * d, decoded_i, d);
+            local_obj += err;
 
-        if (objs) {
-            objs[i] = err;
+            if (objs) {
+                objs[i] = err;
+            }
         }
-    }
+        obj = local_obj;
+    });
 
     obj = obj / n;
     return obj;

data/vendor/faiss/faiss/impl/NNDescent.cpp

@@ -10,10 +10,10 @@
 #include <faiss/impl/NNDescent.h>
 
 #include <mutex>
-#include <string>
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/DistanceComputer.h>
+#include <faiss/impl/VisitedTable.h>
 
 namespace faiss {
 

data/vendor/faiss/faiss/impl/NSG.cpp

@@ -13,6 +13,7 @@
 #include <stack>
 
 #include <faiss/impl/DistanceComputer.h>
+#include <faiss/impl/VisitedTable.h>
 
 namespace faiss {
 
@@ -233,7 +234,7 @@ void NSG::init_graph(Index* storage, const nsg::Graph<idx_t>& knn_graph) {
     std::unique_ptr<DistanceComputer> dis(storage_distance_computer(storage));
 
     dis->set_query(center.get());
-    VisitedTable vt(ntotal);
+    VisitedTable vt(ntotal, use_visited_hashset);
 
     // Do not collect the visited nodes
     search_on_graph<false>(knn_graph, *dis, vt, ep, L, retset, tmpset);
@@ -304,7 +305,7 @@ void NSG::search_on_graph(
             size_t nneigh_for_n = graph.get_neighbors(n, neighbors.data());
             for (int m = 0; m < nneigh_for_n; m++) {
                 int id = neighbors[m];
-                if (id > ntotal || vt.get(id)) {
+                if (id >= ntotal || vt.get(id)) {
                     continue;
                 }
                 vt.set(id);
@@ -341,7 +342,7 @@ void NSG::link(
         std::vector<Node> pool;
         std::vector<Neighbor> tmp;
 
-        VisitedTable vt(ntotal);
+        VisitedTable vt(ntotal, use_visited_hashset);
         std::unique_ptr<DistanceComputer> dis(
                 storage_distance_computer(storage));
 
@@ -513,8 +514,8 @@ void NSG::add_reverse_links(
 
 int NSG::tree_grow(Index* storage, std::vector<int>& degrees) {
     int root = enterpoint;
-    VisitedTable vt(ntotal);
-    VisitedTable vt2(ntotal);
+    VisitedTable vt(ntotal, use_visited_hashset);
+    VisitedTable vt2(ntotal, use_visited_hashset);
 
     int num_attached = 0;
     int cnt = 0;

data/vendor/faiss/faiss/impl/NSG.h

@@ -9,12 +9,12 @@
 
 #include <memory>
 #include <mutex>
+#include <optional>
 #include <vector>
 
 #include <omp.h>
 
 #include <faiss/Index.h>
-#include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/random.h>
@@ -38,6 +38,7 @@ namespace faiss {
  */
 
 struct DistanceComputer; // from AuxIndexStructures
+struct VisitedTable;
 
 namespace nsg {
 
@@ -65,12 +66,14 @@ struct Graph {
     // construct an empty graph
     // NOTE: the newly allocated data needs to be destroyed at destruction time
     Graph(int N, int K) : K(K), N(N), own_fields(true) {
-        data = new node_t[N * K];
+        size_t total = faiss::mul_no_overflow(
+                (size_t)N, (size_t)K, "Graph allocation");
+        data = new node_t[total];
     }
 
     // copy constructor
     Graph(const Graph& g) : Graph(g.N, g.K) {
-        memcpy(data, g.data, N * K * sizeof(node_t));
+        memcpy(data, g.data, (size_t)N * (size_t)K * sizeof(node_t));
     }
 
     // release the allocated memory if needed
@@ -82,21 +85,25 @@ struct Graph {
 
     // access the j-th neighbor of node i
     inline node_t at(int i, int j) const {
-        return data[i * K + j];
+        FAISS_CHECK_RANGE_DEBUG(i, 0, N);
+        FAISS_CHECK_RANGE_DEBUG(j, 0, K);
+        return data[(size_t)i * K + j];
     }
 
     // access the j-th neighbor of node i by reference
     inline node_t& at(int i, int j) {
-        return data[i * K + j];
+        FAISS_CHECK_RANGE_DEBUG(i, 0, N);
+        FAISS_CHECK_RANGE_DEBUG(j, 0, K);
+        return data[(size_t)i * K + j];
     }
 
     // get all neighbors of node i (used during search only)
     virtual size_t get_neighbors(int i, node_t* neighbors) const {
         for (int j = 0; j < K; j++) {
-            if (data[i * K + j] < 0) {
+            if (data[(size_t)i * K + j] < 0) {
                 return j;
             }
-            neighbors[j] = data[i * K + j];
+            neighbors[j] = data[(size_t)i * K + j];
         }
         return K;
     }
@@ -122,6 +129,9 @@ struct NSG {
     // search-time parameters
     int search_L = 16; ///< length of the search path
 
+    // See impl/VisitedTable.h.
+    std::optional<bool> use_visited_hashset;
+
     int enterpoint; ///< enterpoint
 
     std::shared_ptr<nsg::Graph<int32_t>> final_graph; ///< NSG graph structure