faiss 0.3.1 → 0.3.2

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

@@ -46,6 +46,7 @@ struct ResultHandler;
 struct SearchParametersHNSW : SearchParameters {
     int efSearch = 16;
     bool check_relative_distance = true;
+    bool bounded_queue = true;
 
     ~SearchParametersHNSW() {}
 };
@@ -141,9 +142,6 @@ struct HNSW {
     /// enough?
     bool check_relative_distance = true;
 
-    /// number of entry points in levels > 0.
-    int upper_beam = 1;
-
     /// use bounded queue during exploration
     bool search_bounded_queue = true;
 
@@ -184,7 +182,8 @@ struct HNSW {
             float d_nearest,
             int level,
             omp_lock_t* locks,
-            VisitedTable& vt);
+            VisitedTable& vt,
+            bool keep_max_size_level0 = false);
 
     /** add point pt_id on all levels <= pt_level and build the link
      * structure for them. */
@@ -193,7 +192,8 @@ struct HNSW {
             int pt_level,
             int pt_id,
             std::vector<omp_lock_t>& locks,
-            VisitedTable& vt);
+            VisitedTable& vt,
+            bool keep_max_size_level0 = false);
 
     /// search interface for 1 point, single thread
     HNSWStats search(
@@ -211,7 +211,8 @@ struct HNSW {
             const float* nearest_d,
             int search_type,
             HNSWStats& search_stats,
-            VisitedTable& vt) const;
+            VisitedTable& vt,
+            const SearchParametersHNSW* params = nullptr) const;
 
     void reset();
 
@@ -224,40 +225,60 @@ struct HNSW {
             DistanceComputer& qdis,
             std::priority_queue<NodeDistFarther>& input,
             std::vector<NodeDistFarther>& output,
-            int max_size);
+            int max_size,
+            bool keep_max_size_level0 = false);
 
     void permute_entries(const idx_t* map);
 };
 
 struct HNSWStats {
-    size_t n1, n2, n3;
-    size_t ndis;
-    size_t nreorder;
-
-    HNSWStats(
-            size_t n1 = 0,
-            size_t n2 = 0,
-            size_t n3 = 0,
-            size_t ndis = 0,
-            size_t nreorder = 0)
-            : n1(n1), n2(n2), n3(n3), ndis(ndis), nreorder(nreorder) {}
+    size_t n1 = 0; /// number of vectors searched
+    size_t n2 =
+            0; /// number of queries for which the candidate list is exhausted
+    size_t ndis = 0;  /// number of distances computed
+    size_t nhops = 0; /// number of hops aka number of edges traversed
 
     void reset() {
-        n1 = n2 = n3 = 0;
+        n1 = n2 = 0;
         ndis = 0;
-        nreorder = 0;
+        nhops = 0;
     }
 
     void combine(const HNSWStats& other) {
         n1 += other.n1;
         n2 += other.n2;
-        n3 += other.n3;
         ndis += other.ndis;
-        nreorder += other.nreorder;
+        nhops += other.nhops;
     }
 };
 
 // global var that collects them all
 FAISS_API extern HNSWStats hnsw_stats;
 
+int search_from_candidates(
+        const HNSW& hnsw,
+        DistanceComputer& qdis,
+        ResultHandler<HNSW::C>& res,
+        HNSW::MinimaxHeap& candidates,
+        VisitedTable& vt,
+        HNSWStats& stats,
+        int level,
+        int nres_in = 0,
+        const SearchParametersHNSW* params = nullptr);
+
+HNSWStats greedy_update_nearest(
+        const HNSW& hnsw,
+        DistanceComputer& qdis,
+        int level,
+        HNSW::storage_idx_t& nearest,
+        float& d_nearest);
+
+std::priority_queue<HNSW::Node> search_from_candidate_unbounded(
+        const HNSW& hnsw,
+        const HNSW::Node& node,
+        DistanceComputer& qdis,
+        int ef,
+        VisitedTable* vt,
+        HNSWStats& stats);
+
 } // namespace faiss

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

@@ -104,10 +104,10 @@ int dgemm_(
 
 namespace {
 
-void fmat_inverse(float* a, int n) {
-    int info;
-    int lwork = n * n;
-    std::vector<int> ipiv(n);
+void fmat_inverse(float* a, FINTEGER n) {
+    FINTEGER info;
+    FINTEGER lwork = n * n;
+    std::vector<FINTEGER> ipiv(n);
     std::vector<float> workspace(lwork);
 
     sgetrf_(&n, &n, a, &n, ipiv.data(), &info);
@@ -123,10 +123,10 @@ void dfvec_add(size_t d, const double* a, const float* b, double* c) {
     }
 }
 
-void dmat_inverse(double* a, int n) {
-    int info;
-    int lwork = n * n;
-    std::vector<int> ipiv(n);
+void dmat_inverse(double* a, FINTEGER n) {
+    FINTEGER info;
+    FINTEGER lwork = n * n;
+    std::vector<FINTEGER> ipiv(n);
     std::vector<double> workspace(lwork);
 
     dgetrf_(&n, &n, a, &n, ipiv.data(), &info);

data/vendor/faiss/faiss/impl/LookupTableScaler.h

@@ -38,6 +38,23 @@ struct DummyScaler {
         return simd16uint16(0);
     }
 
+#ifdef __AVX512F__
+    inline simd64uint8 lookup(const simd64uint8&, const simd64uint8&) const {
+        FAISS_THROW_MSG("DummyScaler::lookup should not be called.");
+        return simd64uint8(0);
+    }
+
+    inline simd32uint16 scale_lo(const simd64uint8&) const {
+        FAISS_THROW_MSG("DummyScaler::scale_lo should not be called.");
+        return simd32uint16(0);
+    }
+
+    inline simd32uint16 scale_hi(const simd64uint8&) const {
+        FAISS_THROW_MSG("DummyScaler::scale_hi should not be called.");
+        return simd32uint16(0);
+    }
+#endif
+
     template <class dist_t>
     inline dist_t scale_one(const dist_t&) const {
         FAISS_THROW_MSG("DummyScaler::scale_one should not be called.");
@@ -67,6 +84,23 @@ struct NormTableScaler {
         return (simd16uint16(res) >> 8) * scale_simd;
     }
 
+#ifdef __AVX512F__
+    inline simd64uint8 lookup(const simd64uint8& lut, const simd64uint8& c)
+            const {
+        return lut.lookup_4_lanes(c);
+    }
+
+    inline simd32uint16 scale_lo(const simd64uint8& res) const {
+        auto scale_simd_wide = simd32uint16(scale_simd, scale_simd);
+        return simd32uint16(res) * scale_simd_wide;
+    }
+
+    inline simd32uint16 scale_hi(const simd64uint8& res) const {
+        auto scale_simd_wide = simd32uint16(scale_simd, scale_simd);
+        return (simd32uint16(res) >> 8) * scale_simd_wide;
+    }
+#endif
+
     // for non-SIMD implem 2, 3, 4
     template <class dist_t>
     inline dist_t scale_one(const dist_t& x) const {

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

@@ -154,15 +154,20 @@ NNDescent::NNDescent(const int d, const int K) : K(K), d(d) {
 NNDescent::~NNDescent() {}
 
 void NNDescent::join(DistanceComputer& qdis) {
+    idx_t check_period = InterruptCallback::get_period_hint(d * search_L);
+    for (idx_t i0 = 0; i0 < (idx_t)ntotal; i0 += check_period) {
+        idx_t i1 = std::min(i0 + check_period, (idx_t)ntotal);
 #pragma omp parallel for default(shared) schedule(dynamic, 100)
-    for (int n = 0; n < ntotal; n++) {
-        graph[n].join([&](int i, int j) {
-            if (i != j) {
-                float dist = qdis.symmetric_dis(i, j);
-                graph[i].insert(j, dist);
-                graph[j].insert(i, dist);
-            }
-        });
+        for (idx_t n = i0; n < i1; n++) {
+            graph[n].join([&](int i, int j) {
+                if (i != j) {
+                    float dist = qdis.symmetric_dis(i, j);
+                    graph[i].insert(j, dist);
+                    graph[j].insert(i, dist);
+                }
+            });
+        }
+        InterruptCallback::check();
     }
 }
 

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

@@ -25,35 +25,6 @@ namespace {
 // It needs to be smaller than 0
 constexpr int EMPTY_ID = -1;
 
-/* Wrap the distance computer into one that negates the
-   distances. This makes supporting INNER_PRODUCE search easier */
-
-struct NegativeDistanceComputer : DistanceComputer {
-    /// owned by this
-    DistanceComputer* basedis;
-
-    explicit NegativeDistanceComputer(DistanceComputer* basedis)
-            : basedis(basedis) {}
-
-    void set_query(const float* x) override {
-        basedis->set_query(x);
-    }
-
-    /// compute distance of vector i to current query
-    float operator()(idx_t i) override {
-        return -(*basedis)(i);
-    }
-
-    /// compute distance between two stored vectors
-    float symmetric_dis(idx_t i, idx_t j) override {
-        return -basedis->symmetric_dis(i, j);
-    }
-
-    ~NegativeDistanceComputer() override {
-        delete basedis;
-    }
-};
-
 } // namespace
 
 DistanceComputer* storage_distance_computer(const Index* storage) {

data/vendor/faiss/faiss/impl/ProductQuantizer.cpp

@@ -61,6 +61,7 @@ void ProductQuantizer::set_derived_values() {
             "The dimension of the vector (d) should be a multiple of the number of subquantizers (M)");
     dsub = d / M;
     code_size = (nbits * M + 7) / 8;
+    FAISS_THROW_IF_MSG(nbits > 24, "nbits larger than 24 is not practical.");
     ksub = 1 << nbits;
     centroids.resize(d * ksub);
     verbose = false;

data/vendor/faiss/faiss/impl/ProductQuantizer.h

@@ -21,7 +21,11 @@
 
 namespace faiss {
 
-/** Product Quantizer. Implemented only for METRIC_L2 */
+/** Product Quantizer.
+ * PQ is trained using k-means, minimizing the L2 distance to centroids.
+ * PQ supports L2 and Inner Product search, however the quantization error is
+ * biased towards L2 distance.
+ */
 struct ProductQuantizer : Quantizer {
     size_t M;     ///< number of subquantizers
     size_t nbits; ///< number of bits per quantization index

data/vendor/faiss/faiss/impl/ResultHandler.h

@@ -12,9 +12,14 @@
 #pragma once
 
 #include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissException.h>
+#include <faiss/impl/IDSelector.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/partitioning.h>
 
+#include <algorithm>
+#include <iostream>
+
 namespace faiss {
 
 /*****************************************************************
@@ -24,16 +29,21 @@ namespace faiss {
  * - by instanciating a SingleResultHandler that tracks results for a single
  *   query
  * - with begin_multiple/add_results/end_multiple calls where a whole block of
- *   resutls is submitted
+ *   results is submitted
  * All classes are templated on C which to define wheter the min or the max of
- * results is to be kept.
+ * results is to be kept, and on sel, so that the codepaths for with / without
+ * selector can be separated at compile time.
  *****************************************************************/
 
-template <class C>
+template <class C, bool use_sel = false>
 struct BlockResultHandler {
     size_t nq; // number of queries for which we search
+    const IDSelector* sel;
 
-    explicit BlockResultHandler(size_t nq) : nq(nq) {}
+    explicit BlockResultHandler(size_t nq, const IDSelector* sel = nullptr)
+            : nq(nq), sel(sel) {
+        assert(!use_sel || sel);
+    }
 
     // currently handled query range
     size_t i0 = 0, i1 = 0;
@@ -51,13 +61,17 @@ struct BlockResultHandler {
     virtual void end_multiple() {}
 
     virtual ~BlockResultHandler() {}
+
+    bool is_in_selection(idx_t i) const {
+        return !use_sel || sel->is_member(i);
+    }
 };
 
 // handler for a single query
 template <class C>
 struct ResultHandler {
     // if not better than threshold, then not necessary to call add_result
-    typename C::T threshold = 0;
+    typename C::T threshold = C::neutral();
 
     // return whether threshold was updated
     virtual bool add_result(typename C::T dis, typename C::TI idx) = 0;
@@ -71,20 +85,26 @@ struct ResultHandler {
  * some temporary data in memory.
  *****************************************************************/
 
-template <class C>
-struct Top1BlockResultHandler : BlockResultHandler<C> {
+template <class C, bool use_sel = false>
+struct Top1BlockResultHandler : BlockResultHandler<C, use_sel> {
     using T = typename C::T;
     using TI = typename C::TI;
-    using BlockResultHandler<C>::i0;
-    using BlockResultHandler<C>::i1;
+    using BlockResultHandler<C, use_sel>::i0;
+    using BlockResultHandler<C, use_sel>::i1;
 
     // contains exactly nq elements
     T* dis_tab;
     // contains exactly nq elements
     TI* ids_tab;
 
-    Top1BlockResultHandler(size_t nq, T* dis_tab, TI* ids_tab)
-            : BlockResultHandler<C>(nq), dis_tab(dis_tab), ids_tab(ids_tab) {}
+    Top1BlockResultHandler(
+            size_t nq,
+            T* dis_tab,
+            TI* ids_tab,
+            const IDSelector* sel = nullptr)
+            : BlockResultHandler<C, use_sel>(nq, sel),
+              dis_tab(dis_tab),
+              ids_tab(ids_tab) {}
 
     struct SingleResultHandler : ResultHandler<C> {
         Top1BlockResultHandler& hr;
@@ -163,12 +183,12 @@ struct Top1BlockResultHandler : BlockResultHandler<C> {
  * Heap based result handler
  *****************************************************************/
 
-template <class C>
-struct HeapBlockResultHandler : BlockResultHandler<C> {
+template <class C, bool use_sel = false>
+struct HeapBlockResultHandler : BlockResultHandler<C, use_sel> {
     using T = typename C::T;
     using TI = typename C::TI;
-    using BlockResultHandler<C>::i0;
-    using BlockResultHandler<C>::i1;
+    using BlockResultHandler<C, use_sel>::i0;
+    using BlockResultHandler<C, use_sel>::i1;
 
     T* heap_dis_tab;
     TI* heap_ids_tab;
@@ -179,8 +199,9 @@ struct HeapBlockResultHandler : BlockResultHandler<C> {
             size_t nq,
             T* heap_dis_tab,
             TI* heap_ids_tab,
-            size_t k)
-            : BlockResultHandler<C>(nq),
+            size_t k,
+            const IDSelector* sel = nullptr)
+            : BlockResultHandler<C, use_sel>(nq, sel),
               heap_dis_tab(heap_dis_tab),
               heap_ids_tab(heap_ids_tab),
               k(k) {}
@@ -345,12 +366,12 @@ struct ReservoirTopN : ResultHandler<C> {
     }
 };
 
-template <class C>
-struct ReservoirBlockResultHandler : BlockResultHandler<C> {
+template <class C, bool use_sel = false>
+struct ReservoirBlockResultHandler : BlockResultHandler<C, use_sel> {
     using T = typename C::T;
     using TI = typename C::TI;
-    using BlockResultHandler<C>::i0;
-    using BlockResultHandler<C>::i1;
+    using BlockResultHandler<C, use_sel>::i0;
+    using BlockResultHandler<C, use_sel>::i1;
 
     T* heap_dis_tab;
     TI* heap_ids_tab;
@@ -362,8 +383,9 @@ struct ReservoirBlockResultHandler : BlockResultHandler<C> {
             size_t nq,
             T* heap_dis_tab,
             TI* heap_ids_tab,
-            size_t k)
-            : BlockResultHandler<C>(nq),
+            size_t k,
+            const IDSelector* sel = nullptr)
+            : BlockResultHandler<C, use_sel>(nq, sel),
               heap_dis_tab(heap_dis_tab),
               heap_ids_tab(heap_ids_tab),
               k(k) {
@@ -458,18 +480,23 @@ struct ReservoirBlockResultHandler : BlockResultHandler<C> {
  * Result handler for range searches
  *****************************************************************/
 
-template <class C>
-struct RangeSearchBlockResultHandler : BlockResultHandler<C> {
+template <class C, bool use_sel = false>
+struct RangeSearchBlockResultHandler : BlockResultHandler<C, use_sel> {
     using T = typename C::T;
     using TI = typename C::TI;
-    using BlockResultHandler<C>::i0;
-    using BlockResultHandler<C>::i1;
+    using BlockResultHandler<C, use_sel>::i0;
+    using BlockResultHandler<C, use_sel>::i1;
 
     RangeSearchResult* res;
     T radius;
 
-    RangeSearchBlockResultHandler(RangeSearchResult* res, float radius)
-            : BlockResultHandler<C>(res->nq), res(res), radius(radius) {}
+    RangeSearchBlockResultHandler(
+            RangeSearchResult* res,
+            float radius,
+            const IDSelector* sel = nullptr)
+            : BlockResultHandler<C, use_sel>(res->nq, sel),
+              res(res),
+              radius(radius) {}
 
     /******************************************************
      * API for 1 result at a time (each SingleResultHandler is
@@ -504,7 +531,15 @@ struct RangeSearchBlockResultHandler : BlockResultHandler<C> {
         void end() {}
 
         ~SingleResultHandler() {
-            pres.finalize();
+            try {
+                // finalize the partial result
+                pres.finalize();
+            } catch (const faiss::FaissException& e) {
+                // Do nothing if allocation fails in finalizing partial results.
+#ifndef NDEBUG
+                std::cerr << e.what() << std::endl;
+#endif
+            }
         }
     };
 
@@ -559,10 +594,94 @@ struct RangeSearchBlockResultHandler : BlockResultHandler<C> {
     }
 
     ~RangeSearchBlockResultHandler() {
-        if (partial_results.size() > 0) {
-            RangeSearchPartialResult::merge(partial_results);
+        try {
+            if (partial_results.size() > 0) {
+                RangeSearchPartialResult::merge(partial_results);
+            }
+        } catch (const faiss::FaissException& e) {
+            // Do nothing if allocation fails in merge.
+#ifndef NDEBUG
+            std::cerr << e.what() << std::endl;
+#endif
         }
     }
 };
 
+/*****************************************************************
+ * Dispatcher function to choose the right knn result handler depending on k
+ *****************************************************************/
+
+// declared in distances.cpp
+FAISS_API extern int distance_compute_min_k_reservoir;
+
+template <class Consumer, class... Types>
+typename Consumer::T dispatch_knn_ResultHandler(
+        size_t nx,
+        float* vals,
+        int64_t* ids,
+        size_t k,
+        MetricType metric,
+        const IDSelector* sel,
+        Consumer& consumer,
+        Types... args) {
+#define DISPATCH_C_SEL(C, use_sel)                                          \
+    if (k == 1) {                                                           \
+        Top1BlockResultHandler<C, use_sel> res(nx, vals, ids, sel);         \
+        return consumer.template f<>(res, args...);                         \
+    } else if (k < distance_compute_min_k_reservoir) {                      \
+        HeapBlockResultHandler<C, use_sel> res(nx, vals, ids, k, sel);      \
+        return consumer.template f<>(res, args...);                         \
+    } else {                                                                \
+        ReservoirBlockResultHandler<C, use_sel> res(nx, vals, ids, k, sel); \
+        return consumer.template f<>(res, args...);                         \
+    }
+
+    if (is_similarity_metric(metric)) {
+        using C = CMin<float, int64_t>;
+        if (sel) {
+            DISPATCH_C_SEL(C, true);
+        } else {
+            DISPATCH_C_SEL(C, false);
+        }
+    } else {
+        using C = CMax<float, int64_t>;
+        if (sel) {
+            DISPATCH_C_SEL(C, true);
+        } else {
+            DISPATCH_C_SEL(C, false);
+        }
+    }
+#undef DISPATCH_C_SEL
+}
+
+template <class Consumer, class... Types>
+typename Consumer::T dispatch_range_ResultHandler(
+        RangeSearchResult* res,
+        float radius,
+        MetricType metric,
+        const IDSelector* sel,
+        Consumer& consumer,
+        Types... args) {
+#define DISPATCH_C_SEL(C, use_sel)                                    \
+    RangeSearchBlockResultHandler<C, use_sel> resb(res, radius, sel); \
+    return consumer.template f<>(resb, args...);
+
+    if (is_similarity_metric(metric)) {
+        using C = CMin<float, int64_t>;
+        if (sel) {
+            DISPATCH_C_SEL(C, true);
+        } else {
+            DISPATCH_C_SEL(C, false);
+        }
+    } else {
+        using C = CMax<float, int64_t>;
+        if (sel) {
+            DISPATCH_C_SEL(C, true);
+        } else {
+            DISPATCH_C_SEL(C, false);
+        }
+    }
+#undef DISPATCH_C_SEL
+}
+
 } // namespace faiss