faiss 0.3.1 → 0.3.3

data/vendor/faiss/faiss/IndexFastScan.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -17,11 +17,7 @@
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/LookupTableScaler.h>
 #include <faiss/impl/ResultHandler.h>
-#include <faiss/utils/distances.h>
-#include <faiss/utils/extra_distances.h>
 #include <faiss/utils/hamming.h>
-#include <faiss/utils/random.h>
-#include <faiss/utils/utils.h>
 
 #include <faiss/impl/pq4_fast_scan.h>
 #include <faiss/impl/simd_result_handlers.h>
@@ -37,22 +33,22 @@ inline size_t roundup(size_t a, size_t b) {
 
 void IndexFastScan::init_fastscan(
         int d,
-        size_t M_2,
-        size_t nbits_2,
+        size_t M_init,
+        size_t nbits_init,
         MetricType metric,
         int bbs) {
-    FAISS_THROW_IF_NOT(nbits_2 == 4);
+    FAISS_THROW_IF_NOT(nbits_init == 4);
     FAISS_THROW_IF_NOT(bbs % 32 == 0);
     this->d = d;
-    this->M = M_2;
-    this->nbits = nbits_2;
+    this->M = M_init;
+    this->nbits = nbits_init;
     this->metric_type = metric;
     this->bbs = bbs;
-    ksub = (1 << nbits_2);
+    ksub = (1 << nbits_init);
 
-    code_size = (M_2 * nbits_2 + 7) / 8;
+    code_size = (M_init * nbits_init + 7) / 8;
     ntotal = ntotal2 = 0;
-    M2 = roundup(M_2, 2);
+    M2 = roundup(M_init, 2);
     is_trained = false;
 }
 
@@ -189,6 +185,7 @@ void estimators_from_tables_generic(
                 dt += index.ksub;
             }
         }
+
         if (C::cmp(heap_dis[0], dis)) {
             heap_pop<C>(k, heap_dis, heap_ids);
             heap_push<C>(k, heap_dis, heap_ids, dis, j);
@@ -203,17 +200,18 @@ ResultHandlerCompare<C, false>* make_knn_handler(
         idx_t k,
         size_t ntotal,
         float* distances,
-        idx_t* labels) {
+        idx_t* labels,
+        const IDSelector* sel = nullptr) {
     using HeapHC = HeapHandler<C, false>;
     using ReservoirHC = ReservoirHandler<C, false>;
     using SingleResultHC = SingleResultHandler<C, false>;
 
     if (k == 1) {
-        return new SingleResultHC(n, ntotal, distances, labels);
+        return new SingleResultHC(n, ntotal, distances, labels, sel);
     } else if (impl % 2 == 0) {
-        return new HeapHC(n, ntotal, k, distances, labels);
+        return new HeapHC(n, ntotal, k, distances, labels, sel);
     } else /* if (impl % 2 == 1) */ {
-        return new ReservoirHC(n, ntotal, k, 2 * k, distances, labels);
+        return new ReservoirHC(n, ntotal, k, 2 * k, distances, labels, sel);
     }
 }
 
@@ -547,6 +545,22 @@ void IndexFastScan::search_implem_14(
     }
 }
 
+template void IndexFastScan::search_dispatch_implem<true>(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const NormTableScaler* scaler) const;
+
+template void IndexFastScan::search_dispatch_implem<false>(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const NormTableScaler* scaler) const;
+
 void IndexFastScan::reconstruct(idx_t key, float* recons) const {
     std::vector<uint8_t> code(code_size, 0);
     BitstringWriter bsw(code.data(), code_size);

data/vendor/faiss/faiss/IndexFastScan.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -133,6 +133,15 @@ struct IndexFastScan : Index {
 
     void merge_from(Index& otherIndex, idx_t add_id = 0) override;
     void check_compatible_for_merge(const Index& otherIndex) const override;
+
+    /// standalone codes interface (but the codes are flattened)
+    size_t sa_code_size() const override {
+        return code_size;
+    }
+
+    void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override {
+        compute_codes(bytes, n, x);
+    }
 };
 
 struct FastScanStats {

data/vendor/faiss/faiss/IndexFlat.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -16,7 +16,6 @@
 #include <faiss/utils/extra_distances.h>
 #include <faiss/utils/prefetch.h>
 #include <faiss/utils/sorting.h>
-#include <faiss/utils/utils.h>
 #include <cstring>
 
 namespace faiss {
@@ -41,15 +40,19 @@ void IndexFlat::search(
     } else if (metric_type == METRIC_L2) {
         float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances};
         knn_L2sqr(x, get_xb(), d, n, ntotal, &res, nullptr, sel);
-    } else if (is_similarity_metric(metric_type)) {
-        float_minheap_array_t res = {size_t(n), size_t(k), labels, distances};
-        knn_extra_metrics(
-                x, get_xb(), d, n, ntotal, metric_type, metric_arg, &res);
     } else {
-        FAISS_THROW_IF_NOT(!sel);
-        float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances};
+        FAISS_THROW_IF_NOT(!sel); // TODO implement with selector
         knn_extra_metrics(
-                x, get_xb(), d, n, ntotal, metric_type, metric_arg, &res);
+                x,
+                get_xb(),
+                d,
+                n,
+                ntotal,
+                metric_type,
+                metric_arg,
+                k,
+                distances,
+                labels);
     }
 }
 

data/vendor/faiss/faiss/IndexFlat.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/IndexFlatCodes.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -12,6 +12,8 @@
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/IDSelector.h>
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/utils/extra_distances.h>
 
 namespace faiss {
 
@@ -30,6 +32,15 @@ void IndexFlatCodes::add(idx_t n, const float* x) {
     ntotal += n;
 }
 
+void IndexFlatCodes::add_sa_codes(
+        idx_t n,
+        const uint8_t* codes_in,
+        const idx_t* /* xids */) {
+    codes.resize((ntotal + n) * code_size);
+    memcpy(codes.data() + (ntotal * code_size), codes_in, n * code_size);
+    ntotal += n;
+}
+
 void IndexFlatCodes::reset() {
     codes.clear();
     ntotal = 0;
@@ -70,11 +81,6 @@ void IndexFlatCodes::reconstruct(idx_t key, float* recons) const {
     reconstruct_n(key, 1, recons);
 }
 
-FlatCodesDistanceComputer* IndexFlatCodes::get_FlatCodesDistanceComputer()
-        const {
-    FAISS_THROW_MSG("not implemented");
-}
-
 void IndexFlatCodes::check_compatible_for_merge(const Index& otherIndex) const {
     // minimal sanity checks
     const IndexFlatCodes* other =
@@ -114,4 +120,161 @@ void IndexFlatCodes::permute_entries(const idx_t* perm) {
     std::swap(codes, new_codes);
 }
 
+namespace {
+
+template <class VD>
+struct GenericFlatCodesDistanceComputer : FlatCodesDistanceComputer {
+    const IndexFlatCodes& codec;
+    const VD vd;
+    // temp buffers
+    std::vector<uint8_t> code_buffer;
+    std::vector<float> vec_buffer;
+    const float* query = nullptr;
+
+    GenericFlatCodesDistanceComputer(const IndexFlatCodes* codec, const VD& vd)
+            : FlatCodesDistanceComputer(codec->codes.data(), codec->code_size),
+              codec(*codec),
+              vd(vd),
+              code_buffer(codec->code_size * 4),
+              vec_buffer(codec->d * 4) {}
+
+    void set_query(const float* x) override {
+        query = x;
+    }
+
+    float operator()(idx_t i) override {
+        codec.sa_decode(1, codes + i * code_size, vec_buffer.data());
+        return vd(query, vec_buffer.data());
+    }
+
+    float distance_to_code(const uint8_t* code) override {
+        codec.sa_decode(1, code, vec_buffer.data());
+        return vd(query, vec_buffer.data());
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) override {
+        codec.sa_decode(1, codes + i * code_size, vec_buffer.data());
+        codec.sa_decode(1, codes + j * code_size, vec_buffer.data() + vd.d);
+        return vd(vec_buffer.data(), vec_buffer.data() + vd.d);
+    }
+
+    void distances_batch_4(
+            const idx_t idx0,
+            const idx_t idx1,
+            const idx_t idx2,
+            const idx_t idx3,
+            float& dis0,
+            float& dis1,
+            float& dis2,
+            float& dis3) override {
+        uint8_t* cp = code_buffer.data();
+        for (idx_t i : {idx0, idx1, idx2, idx3}) {
+            memcpy(cp, codes + i * code_size, code_size);
+            cp += code_size;
+        }
+        // potential benefit is if batch decoding is more efficient than 1 by 1
+        // decoding
+        codec.sa_decode(4, code_buffer.data(), vec_buffer.data());
+        dis0 = vd(query, vec_buffer.data());
+        dis1 = vd(query, vec_buffer.data() + vd.d);
+        dis2 = vd(query, vec_buffer.data() + 2 * vd.d);
+        dis3 = vd(query, vec_buffer.data() + 3 * vd.d);
+    }
+};
+
+struct Run_get_distance_computer {
+    using T = FlatCodesDistanceComputer*;
+
+    template <class VD>
+    FlatCodesDistanceComputer* f(const VD& vd, const IndexFlatCodes* codec) {
+        return new GenericFlatCodesDistanceComputer<VD>(codec, vd);
+    }
+};
+
+template <class BlockResultHandler>
+struct Run_search_with_decompress {
+    using T = void;
+
+    template <class VectorDistance>
+    void f(VectorDistance& vd,
+           const IndexFlatCodes* index_ptr,
+           const float* xq,
+           BlockResultHandler& res) {
+        // Note that there seems to be a clang (?) bug that "sometimes" passes
+        // the const Index & parameters by value, so to be on the safe side,
+        // it's better to use pointers.
+        const IndexFlatCodes& index = *index_ptr;
+        size_t ntotal = index.ntotal;
+        using SingleResultHandler =
+                typename BlockResultHandler::SingleResultHandler;
+        using DC = GenericFlatCodesDistanceComputer<VectorDistance>;
+#pragma omp parallel // if (res.nq > 100)
+        {
+            std::unique_ptr<DC> dc(new DC(&index, vd));
+            SingleResultHandler resi(res);
+#pragma omp for
+            for (int64_t q = 0; q < res.nq; q++) {
+                resi.begin(q);
+                dc->set_query(xq + vd.d * q);
+                for (size_t i = 0; i < ntotal; i++) {
+                    if (res.is_in_selection(i)) {
+                        float dis = (*dc)(i);
+                        resi.add_result(dis, i);
+                    }
+                }
+                resi.end();
+            }
+        }
+    }
+};
+
+struct Run_search_with_decompress_res {
+    using T = void;
+
+    template <class ResultHandler>
+    void f(ResultHandler& res, const IndexFlatCodes* index, const float* xq) {
+        Run_search_with_decompress<ResultHandler> r;
+        dispatch_VectorDistance(
+                index->d,
+                index->metric_type,
+                index->metric_arg,
+                r,
+                index,
+                xq,
+                res);
+    }
+};
+
+} // anonymous namespace
+
+FlatCodesDistanceComputer* IndexFlatCodes::get_FlatCodesDistanceComputer()
+        const {
+    Run_get_distance_computer r;
+    return dispatch_VectorDistance(d, metric_type, metric_arg, r, this);
+}
+
+void IndexFlatCodes::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const SearchParameters* params) const {
+    Run_search_with_decompress_res r;
+    const IDSelector* sel = params ? params->sel : nullptr;
+    dispatch_knn_ResultHandler(
+            n, distances, labels, k, metric_type, sel, r, this, x);
+}
+
+void IndexFlatCodes::range_search(
+        idx_t n,
+        const float* x,
+        float radius,
+        RangeSearchResult* result,
+        const SearchParameters* params) const {
+    const IDSelector* sel = params ? params->sel : nullptr;
+    Run_search_with_decompress_res r;
+    dispatch_range_ResultHandler(result, radius, metric_type, sel, r, this, x);
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexFlatCodes.h

@@ -1,12 +1,10 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 #pragma once
 
 #include <faiss/Index.h>
@@ -45,13 +43,32 @@ struct IndexFlatCodes : Index {
      * different from the usual ones: the new ids are shifted */
     size_t remove_ids(const IDSelector& sel) override;
 
-    /** a FlatCodesDistanceComputer offers a distance_to_code method */
+    /** a FlatCodesDistanceComputer offers a distance_to_code method
+     *
+     * The default implementation explicitly decodes the vector with sa_decode.
+     */
     virtual FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const;
 
     DistanceComputer* get_distance_computer() const override {
         return get_FlatCodesDistanceComputer();
     }
 
+    /** Search implemented by decoding */
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    void range_search(
+            idx_t n,
+            const float* x,
+            float radius,
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
+
     // returns a new instance of a CodePacker
     CodePacker* get_CodePacker() const;
 
@@ -59,6 +76,9 @@ struct IndexFlatCodes : Index {
 
     virtual void merge_from(Index& otherIndex, idx_t add_id = 0) override;
 
+    virtual void add_sa_codes(idx_t n, const uint8_t* x, const idx_t* xids)
+            override;
+
     // permute_entries. perm of size ntotal maps new to old positions
     void permute_entries(const idx_t* perm);
 };