faiss 0.6.2 → 0.6.3

data/vendor/faiss/faiss/IndexIVFEDEN.cpp

@@ -0,0 +1,302 @@
+/*
+ * 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.
+ */
+
+#include <faiss/IndexIVFEDEN.h>
+
+#include <omp.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <memory>
+#include <vector>
+
+#include <faiss/impl/EDENQuantizer.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/expanded_scanners.h>
+
+namespace faiss {
+
+IndexIVFEDEN::IndexIVFEDEN(
+        Index* quantizer_in,
+        const size_t d_in,
+        const size_t nlist_in,
+        MetricType metric,
+        bool own_invlists_in,
+        uint8_t nb_bits_in,
+        EDENScaleType scale_type_in)
+        : IndexIVF(quantizer_in, d_in, nlist_in, 0, metric, own_invlists_in),
+          sq(d_in, eden_utils::quantizer_type_for_bits(nb_bits_in)),
+          scale_type(scale_type_in) {
+    FAISS_THROW_IF_NOT_MSG(
+            scale_type == EDENScaleType_UNBIASED ||
+                    scale_type == EDENScaleType_BIASED,
+            "invalid EDEN scale type");
+    sq.train(0, nullptr);
+    code_size = eden_utils::code_size(d, sq.bits);
+    if (own_invlists_in) {
+        invlists->code_size = code_size;
+    }
+    is_trained = false;
+    by_residual = true;
+}
+
+IndexIVFEDEN::IndexIVFEDEN() {
+    by_residual = true;
+}
+
+void IndexIVFEDEN::train_encoder(
+        idx_t n,
+        const float* x,
+        const idx_t* /*assign*/) {
+    sq.train(n, x);
+}
+
+void IndexIVFEDEN::encode_vectors(
+        idx_t n,
+        const float* x,
+        const idx_t* list_nos,
+        uint8_t* codes,
+        bool include_listnos) const {
+    const size_t coarse_size = include_listnos ? coarse_code_size() : 0;
+    memset(codes, 0, (code_size + coarse_size) * n);
+
+#pragma omp parallel if (n > 1000)
+    {
+        std::vector<float> centroid(d);
+
+#pragma omp for
+        for (idx_t i = 0; i < n; i++) {
+            const int64_t list_no = list_nos[i];
+            if (list_no >= 0) {
+                const float* xi = x + i * d;
+                uint8_t* code = codes + i * (code_size + coarse_size);
+
+                quantizer->reconstruct(list_no, centroid.data());
+                eden_utils::compute_codes(
+                        sq,
+                        metric_type,
+                        scale_type,
+                        xi,
+                        code + coarse_size,
+                        1,
+                        centroid.data());
+
+                if (coarse_size) {
+                    encode_listno(list_no, code);
+                }
+            }
+        }
+    }
+}
+
+void IndexIVFEDEN::decode_vectors(
+        idx_t n,
+        const uint8_t* codes,
+        const idx_t* listnos,
+        float* x) const {
+#pragma omp parallel
+    {
+        std::vector<float> centroid(d);
+
+#pragma omp for
+        for (idx_t i = 0; i < n; i++) {
+            const uint8_t* code = codes + i * code_size;
+            const int64_t list_no = listnos[i];
+            float* xi = x + i * d;
+
+            quantizer->reconstruct(list_no, centroid.data());
+            eden_utils::decode(sq, code, xi, 1, centroid.data());
+        }
+    }
+}
+
+void IndexIVFEDEN::add_core(
+        idx_t n,
+        const float* x,
+        const idx_t* xids,
+        const idx_t* precomputed_idx,
+        void* inverted_list_context) {
+    FAISS_THROW_IF_NOT(is_trained);
+
+    DirectMapAdd dm_add(direct_map, n, xids);
+
+#pragma omp parallel
+    {
+        std::vector<uint8_t> one_code(code_size);
+        std::vector<float> centroid(d);
+
+        const int nt = omp_get_num_threads();
+        const int rank = omp_get_thread_num();
+
+        for (idx_t i = 0; i < n; i++) {
+            const int64_t list_no = precomputed_idx[i];
+            if (list_no >= 0 && list_no % nt == rank) {
+                const int64_t id = xids ? xids[i] : ntotal + i;
+                const float* xi = x + i * d;
+
+                quantizer->reconstruct(list_no, centroid.data());
+                eden_utils::compute_codes(
+                        sq,
+                        metric_type,
+                        scale_type,
+                        xi,
+                        one_code.data(),
+                        1,
+                        centroid.data());
+
+                const size_t ofs = invlists->add_entry(
+                        list_no, id, one_code.data(), inverted_list_context);
+                dm_add.add(i, list_no, ofs);
+            } else if (rank == 0 && list_no == -1) {
+                dm_add.add(i, -1, 0);
+            }
+        }
+    }
+
+    ntotal += n;
+}
+
+namespace {
+
+struct EDENInvertedListScanner : InvertedListScanner {
+    using InvertedListScanner::scan_codes;
+
+    const IndexIVFEDEN& ivf_eden;
+    std::vector<float> reconstructed_centroid;
+    std::vector<float> query_vector;
+    std::unique_ptr<FlatCodesDistanceComputer> dc;
+
+    explicit EDENInvertedListScanner(
+            const IndexIVFEDEN& ivf_eden_in,
+            bool store_pairs_in = false,
+            const IDSelector* sel_in = nullptr)
+            : InvertedListScanner(store_pairs_in, sel_in),
+              ivf_eden(ivf_eden_in) {
+        keep_max = is_similarity_metric(ivf_eden.metric_type);
+        code_size = ivf_eden.code_size;
+    }
+
+    void set_query(const float* query_vector_in) override {
+        query_vector.assign(query_vector_in, query_vector_in + ivf_eden.d);
+        internal_try_setup_dc();
+    }
+
+    void set_list(idx_t list_no_in, float /*coarse_dis*/) override {
+        list_no = list_no_in;
+
+        reconstructed_centroid.resize(ivf_eden.d);
+        ivf_eden.quantizer->reconstruct(
+                list_no_in, reconstructed_centroid.data());
+        internal_try_setup_dc();
+    }
+
+    float distance_to_code(const uint8_t* code) const final {
+        return dc->distance_to_code(code);
+    }
+
+    size_t scan_codes(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            ResultHandler& handler) const override {
+        return run_scan_codes(*this, list_size, codes, ids, handler);
+    }
+
+    void internal_try_setup_dc() {
+        if (!query_vector.empty() && !reconstructed_centroid.empty()) {
+            dc.reset(
+                    eden_utils::get_distance_computer(
+                            ivf_eden.sq,
+                            ivf_eden.metric_type,
+                            reconstructed_centroid.data()));
+            dc->set_query(query_vector.data());
+        }
+    }
+};
+
+} // namespace
+
+InvertedListScanner* IndexIVFEDEN::get_InvertedListScanner(
+        bool store_pairs,
+        const IDSelector* sel,
+        const IVFSearchParameters* /*params*/) const {
+    return new EDENInvertedListScanner(*this, store_pairs, sel);
+}
+
+void IndexIVFEDEN::reconstruct_from_offset(
+        int64_t list_no,
+        int64_t offset,
+        float* recons) const {
+    const uint8_t* code = invlists->get_single_code(list_no, offset);
+
+    std::vector<float> centroid(d);
+    quantizer->reconstruct(list_no, centroid.data());
+    eden_utils::decode(sq, code, recons, 1, centroid.data());
+
+    invlists->release_codes(list_no, code);
+}
+
+void IndexIVFEDEN::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
+    const size_t coarse_size = coarse_code_size();
+
+#pragma omp parallel
+    {
+        std::vector<float> centroid(d);
+
+#pragma omp for
+        for (idx_t i = 0; i < n; i++) {
+            const uint8_t* code = bytes + i * (code_size + coarse_size);
+            const int64_t list_no = decode_listno(code);
+            float* xi = x + i * d;
+
+            quantizer->reconstruct(list_no, centroid.data());
+            eden_utils::decode(sq, code + coarse_size, xi, 1, centroid.data());
+        }
+    }
+}
+
+struct IVFEDENDistanceComputer : DistanceComputer {
+    const float* q = nullptr;
+    const IndexIVFEDEN* parent = nullptr;
+
+    void set_query(const float* x) override {
+        q = x;
+    }
+
+    float operator()(idx_t i) override {
+        const idx_t lo = parent->direct_map.get(i);
+        const uint64_t list_no = lo_listno(lo);
+        const uint64_t offset = lo_offset(lo);
+        const uint8_t* code =
+                parent->invlists->get_single_code(list_no, offset);
+
+        std::vector<float> centroid(parent->d);
+        parent->quantizer->reconstruct(list_no, centroid.data());
+
+        std::unique_ptr<FlatCodesDistanceComputer> dc(
+                eden_utils::get_distance_computer(
+                        parent->sq, parent->metric_type, centroid.data()));
+        dc->set_query(q);
+        const float distance = dc->distance_to_code(code);
+
+        parent->invlists->release_codes(list_no, code);
+        return distance;
+    }
+
+    float symmetric_dis(idx_t /*i*/, idx_t /*j*/) override {
+        FAISS_THROW_MSG("Not implemented");
+    }
+};
+
+DistanceComputer* IndexIVFEDEN::get_distance_computer() const {
+    IVFEDENDistanceComputer* dc = new IVFEDENDistanceComputer;
+    dc->parent = this;
+    return dc;
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexIVFEDEN.h

@@ -0,0 +1,70 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+#include <faiss/IndexIVF.h>
+#include <faiss/impl/EDENQuantizer.h>
+
+namespace faiss {
+
+struct IndexIVFEDEN : IndexIVF {
+    ScalarQuantizer sq;
+    EDENScaleType scale_type = EDENScaleType_UNBIASED;
+
+    // Factory strings: IVF<nlist>,EDEN, IVF<nlist>,EDEN<n>, and the same
+    // forms with the BIASED suffix for EDEN's MSE-minimizing scale.
+    IndexIVFEDEN(
+            Index* quantizer,
+            const size_t d,
+            const size_t nlist,
+            MetricType metric = METRIC_L2,
+            bool own_invlists = true,
+            uint8_t nb_bits = 1,
+            EDENScaleType scale_type = EDENScaleType_UNBIASED);
+
+    IndexIVFEDEN();
+
+    void train_encoder(idx_t n, const float* x, const idx_t* assign) override;
+
+    void encode_vectors(
+            idx_t n,
+            const float* x,
+            const idx_t* list_nos,
+            uint8_t* codes,
+            bool include_listnos = false) const override;
+
+    void decode_vectors(
+            idx_t n,
+            const uint8_t* codes,
+            const idx_t* list_nos,
+            float* x) const override;
+
+    void add_core(
+            idx_t n,
+            const float* x,
+            const idx_t* xids,
+            const idx_t* precomputed_idx,
+            void* inverted_list_context = nullptr) override;
+
+    InvertedListScanner* get_InvertedListScanner(
+            bool store_pairs,
+            const IDSelector* sel,
+            const IVFSearchParameters* params) const override;
+
+    void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
+            const override;
+
+    void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
+
+    DistanceComputer* get_distance_computer() const override;
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexIVFFastScan.cpp

@@ -376,9 +376,8 @@ void IndexIVFFastScan::search_preassigned(
         cur_nprobe = params->nprobe;
     }
 
-    FAISS_THROW_IF_NOT_MSG(
-            !store_pairs, "store_pairs not supported for this index");
-    FAISS_THROW_IF_NOT_MSG(!stats, "stats not supported for this index");
+    FAISS_THROW_IF_MSG(store_pairs, "store_pairs not supported for this index");
+    FAISS_THROW_IF_MSG(stats, "stats not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
     FastScanDistancePostProcessing empty_context{};
 
@@ -404,8 +403,8 @@ void IndexIVFFastScan::range_search(
                 params->max_lists_num == 0,
                 "max_lists_num is a knn knob and is not honored by "
                 "fastscan range search");
-        FAISS_THROW_IF_NOT_MSG(
-                !params->ensure_topk_full,
+        FAISS_THROW_IF_MSG(
+                params->ensure_topk_full,
                 "ensure_topk_full is a knn knob and is not honored by "
                 "fastscan range search");
         FAISS_THROW_IF_NOT_MSG(
@@ -1592,7 +1591,7 @@ void IndexIVFFastScan::reconstruct_from_offset(
 }
 
 void IndexIVFFastScan::reconstruct_orig_invlists() {
-    FAISS_THROW_IF_NOT(orig_invlists != nullptr);
+    FAISS_THROW_IF_NOT(orig_invlists);
     FAISS_THROW_IF_NOT(orig_invlists->list_size(0) == 0);
 
 #pragma omp parallel for if (nlist > 100)

data/vendor/faiss/faiss/IndexIVFFlat.cpp

@@ -65,7 +65,7 @@ void IndexIVFFlat::add_core(
         void* inverted_list_context) {
     FAISS_THROW_IF_NOT(is_trained);
     FAISS_THROW_IF_NOT(coarse_idx);
-    FAISS_THROW_IF_NOT(!by_residual);
+    FAISS_THROW_IF_MSG(by_residual, "by_residual not supported for this index");
     FAISS_THROW_IF_NOT_MSG(invlists, "invlists not initialized");
     direct_map.check_can_add(xids);
 
@@ -110,7 +110,7 @@ void IndexIVFFlat::encode_vectors(
         const idx_t* list_nos,
         uint8_t* codes,
         bool include_listnos) const {
-    FAISS_THROW_IF_NOT(!by_residual);
+    FAISS_THROW_IF_MSG(by_residual, "by_residual not supported for this index");
     if (!include_listnos) {
         memcpy(codes, x, code_size * n);
     } else {
@@ -291,8 +291,7 @@ void IndexIVFFlatDedup::search_preassigned(
         bool store_pairs,
         const IVFSearchParameters* params,
         IndexIVFStats* /*stats*/) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !store_pairs, "store_pairs not supported in IVFDedup");
+    FAISS_THROW_IF_MSG(store_pairs, "store_pairs not supported in IVFDedup");
 
     IndexIVFFlat::search_preassigned(
             n, x, k, assign, centroid_dis, distances, labels, false, params);

data/vendor/faiss/faiss/IndexIVFIndependentQuantizer.cpp

@@ -90,7 +90,7 @@ void IndexIVFIndependentQuantizer::search(
         float* distances,
         idx_t* labels,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(!params, "search parameters not supported");
+    FAISS_THROW_IF_MSG(params, "search parameters not supported");
     size_t nprobe = index_ivf->nprobe;
     std::vector<float> D(n * nprobe);
     std::vector<idx_t> I(n * nprobe);

data/vendor/faiss/faiss/IndexIVFPQ.cpp

@@ -14,6 +14,8 @@
 #include <cstdint>
 #include <cstdio>
 
+#include <omp.h>
+
 #include <algorithm>
 
 #include <faiss/utils/distances_dispatch.h>
@@ -28,6 +30,7 @@
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/ProductQuantizer.h>
 #include <faiss/impl/ResultHandler.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader,facebook-unused-include-check)
 #include <faiss/impl/pq_code_distance/pq_code_distance-generic.h>
 #include <faiss/impl/simd_dispatch.h>
 
@@ -292,34 +295,49 @@ void IndexIVFPQ::add_core_o(
     pq.compute_codes(to_encode, xcodes.get(), n);
 
     double t2 = getmillisecs();
-    // TODO: parallelize?
     size_t n_ignore = 0;
-    for (idx_t i = 0; i < n; i++) {
-        idx_t key = idx[i];
-        idx_t id = xids ? xids[i] : ntotal + i;
-        if (key < 0) {
-            direct_map.add_single_id(id, -1, 0);
-            n_ignore++;
-            if (residuals_2) {
-                memset(residuals_2, 0, sizeof(*residuals_2) * d);
+    DirectMapAdd dm_adder(direct_map, n, xids);
+
+#pragma omp parallel reduction(+ : n_ignore)
+    {
+        int nt = omp_get_num_threads();
+        int rank = omp_get_thread_num();
+
+        // each thread takes care of a subset of lists
+        for (idx_t i = 0; i < n; i++) {
+            idx_t key = idx[i];
+            if (key < 0) {
+                if (rank == 0) {
+                    dm_adder.add(i, -1, 0);
+                    n_ignore++;
+                    if (residuals_2) {
+                        memset(residuals_2 + i * d,
+                               0,
+                               sizeof(*residuals_2) * d);
+                    }
+                }
+                continue;
+            }
+            if (key % nt != rank) {
+                continue;
             }
-            continue;
-        }
 
-        uint8_t* code = xcodes.get() + i * code_size;
-        size_t offset =
-                invlists->add_entry(key, id, code, inverted_list_context);
+            idx_t id = xids ? xids[i] : ntotal + i;
+            uint8_t* code = xcodes.get() + i * code_size;
+            size_t offset =
+                    invlists->add_entry(key, id, code, inverted_list_context);
 
-        if (residuals_2) {
-            float* res2 = residuals_2 + i * d;
-            const float* xi = to_encode + i * d;
-            pq.decode(code, res2);
-            for (int j = 0; j < d; j++) {
-                res2[j] = xi[j] - res2[j];
+            if (residuals_2) {
+                float* res2 = residuals_2 + i * d;
+                const float* xi = to_encode + i * d;
+                pq.decode(code, res2);
+                for (int j = 0; j < d; j++) {
+                    res2[j] = xi[j] - res2[j];
+                }
             }
-        }
 
-        direct_map.add_single_id(id, key, offset);
+            dm_adder.add(i, key, offset);
+        }
     }
 
     double t3 = getmillisecs();

data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp

@@ -16,7 +16,6 @@
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/ResultHandler.h>
 #include <faiss/impl/simdlib/simdlib_dispatch.h>
-#include <faiss/utils/distances.h>
 #include <faiss/utils/distances_dispatch.h>
 #include <faiss/utils/extra_distances.h>
 

data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp

@@ -229,60 +229,30 @@ struct RaBitInvertedListScanner : InvertedListScanner {
         }
 
         // Multi-bit: Two-stage search with adaptive filtering
-        size_t nup = 0;
-
-        for (size_t j = 0; j < list_size; j++) {
-            if (sel != nullptr) {
-                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-                if (!sel->is_member(id)) {
-                    codes += code_size;
-                    continue;
-                }
-            }
-
-            float est_distance = rabitq_dc->distance_to_code_1bit(codes);
-
-            size_t code_size_base = (ivf_rabitq.d + 7) / 8;
-            const rabitq_utils::SignBitFactorsWithError* base_fac =
-                    reinterpret_cast<
-                            const rabitq_utils::SignBitFactorsWithError*>(
-                            codes + code_size_base);
-
-            bool should_refine = rabitq_utils::should_refine_candidate(
-                    est_distance,
-                    base_fac->f_error,
-                    rabitq_dc->g_error,
-                    handler.threshold,
-                    keep_max);
-            if (should_refine) {
-                // Refining computes the full distance — counts as a
-                // post-filter "distance computed" for stats purposes.
-                handler.stats.scan_cnt++;
-                float dis = distance_to_code(codes);
-                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-
-                if (handler.add_result(dis, id)) {
-                    handler.stats.nheap_updates++;
-                    nup++;
-                }
-            }
-            codes += code_size;
-        }
-
-        return nup;
+        return rabitq_dc->scan_codes_multibit(
+                list_size,
+                codes,
+                ids,
+                code_size,
+                list_no,
+                store_pairs,
+                sel,
+                keep_max,
+                handler);
     }
 
     void internal_try_setup_dc() {
         if (!query_vector.empty() && !reconstructed_centroid.empty()) {
-            // both query_vector and centroid are available!
-            // set up DistanceComputer
-            dc.reset(ivf_rabitq.rabitq.get_distance_computer(
-                    qb, reconstructed_centroid.data(), centered));
-
+            // both query_vector and centroid are available
+            if (!dc) {
+                dc.reset(ivf_rabitq.rabitq.get_distance_computer(
+                        qb, nullptr, centered));
+                // Try to cast to RaBitQDistanceComputer for multi-bit support
+                rabitq_dc = dynamic_cast<RaBitQDistanceComputer*>(dc.get());
+                FAISS_THROW_IF_NOT(rabitq_dc);
+            }
+            rabitq_dc->set_centroid(reconstructed_centroid.data());
             dc->set_query(query_vector.data());
-
-            // Try to cast to RaBitQDistanceComputer for multi-bit support
-            rabitq_dc = dynamic_cast<RaBitQDistanceComputer*>(dc.get());
         }
     }
 };