faiss 0.6.4 → 0.6.5

data/vendor/faiss/faiss/IndexIVFSQFastScan.h

@@ -0,0 +1,145 @@
+/*
+ * 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 <faiss/IndexIVFFastScan.h>
+#include <faiss/IndexScalarQuantizer.h>
+#include <faiss/impl/ScalarQuantizer.h>
+#include <faiss/impl/sq_fastscan_utils.h>
+
+namespace faiss {
+
+/** Fast scan version of IndexIVFScalarQuantizer.
+ *
+ * Uses the PQ4 FastScan SIMD infrastructure (vpshufb) for accelerated
+ * scanning within inverted lists.  Supported quantizer types:
+ *
+ *   - QT_4bit, QT_4bit_uniform: native 4-bit codes are packed directly
+ *     into the SIMD block layout.  No precision loss.
+ *
+ *   - QT_6bit, QT_8bit, QT_8bit_uniform, QT_8bit_direct,
+ *     QT_8bit_direct_signed: re-quantised to 4-bit for the SIMD scan,
+ *     then reranked with exact original-precision distances.
+ *
+ *   - All other types (QT_fp16, QT_bf16, TurboQuant): fall back to the
+ *     ScalarQuantizer's own SIMD-optimised InvertedListScanner (same
+ *     behavior as IndexIVFScalarQuantizer, no fast-scan acceleration).
+ *
+ * For reranked types the block inverted lists hold the packed high nibbles
+ * and a parallel ArrayInvertedLists (lo_codes_invlists) holds the low ones.
+ */
+struct IndexIVFSQFastScan : IndexIVFFastScan {
+    ScalarQuantizer sq;
+
+    /// Overselection ratio for reranking (default 2).
+    float rerank_factor = 2;
+
+    /// Parallel inverted lists holding the leftover low bits of each code.
+    /// The top 4 bits of every code live in the packed scan blocks.
+    /// Owned by this index. nullptr for native 4-bit and fallback types.
+    InvertedLists* lo_codes_invlists = nullptr;
+
+    /// Bytes per vector in lo_codes_invlists: the bits of each code left
+    /// over once the scan has taken the top 4.
+    size_t lo_code_size() const {
+        return sq_fastscan::sq_rerank_size(d, sq.qtype);
+    }
+
+    /// Rebuild the code at (list_no, offset) from its two halves. `scratch`
+    /// needs M2 / 2 bytes, `code_out` needs sq.code_size.
+    void recombine_code(
+            int64_t list_no,
+            int64_t offset,
+            uint8_t* scratch,
+            int* values,
+            uint8_t* code_out) const;
+
+    IndexIVFSQFastScan(
+            Index* quantizer,
+            size_t d,
+            size_t nlist,
+            ScalarQuantizer::QuantizerType qtype,
+            MetricType metric = METRIC_L2,
+            int bbs = 32,
+            bool by_residual = true);
+
+    IndexIVFSQFastScan();
+
+    /// Build from an existing IndexIVFScalarQuantizer.
+    explicit IndexIVFSQFastScan(
+            const IndexIVFScalarQuantizer& orig,
+            int bbs = 32);
+
+    ~IndexIVFSQFastScan() override;
+
+    size_t fast_scan_code_size() const override;
+
+    void train_encoder(idx_t n, const float* x, const idx_t* assign) override;
+
+    idx_t train_encoder_num_vectors() const override;
+
+    void encode_vectors(
+            idx_t n,
+            const float* x,
+            const idx_t* list_nos,
+            uint8_t* codes,
+            bool include_listno = false) const override;
+
+    bool lookup_table_is_3d() const override;
+
+    void compute_LUT(
+            size_t n,
+            const float* x,
+            const CoarseQuantized& cq,
+            AlignedTable<float>& dis_tables,
+            AlignedTable<float>& biases,
+            const FastScanDistancePostProcessing& context) const override;
+
+    void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
+
+    void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
+            const override;
+
+    InvertedListScanner* get_InvertedListScanner(
+            bool store_pairs,
+            const IDSelector* sel,
+            const IVFSearchParameters*) const override;
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    void search_preassigned(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            const idx_t* assign,
+            const float* centroid_dis,
+            float* distances,
+            idx_t* labels,
+            bool store_pairs,
+            const IVFSearchParameters* params = nullptr,
+            IndexIVFStats* stats = nullptr) const override;
+
+    void range_search(
+            idx_t n,
+            const float* x,
+            float radius,
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
+
+    void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
+
+    void reset() override;
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexNNDescent.cpp

@@ -14,9 +14,6 @@
 #include <cstdio>
 #include <cstdlib>
 
-#ifdef __SSE__
-#endif
-
 #include <faiss/IndexFlat.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>

data/vendor/faiss/faiss/IndexPiPNN.cpp

@@ -0,0 +1,201 @@
+/*
+ * 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/IndexPiPNN.h>
+
+#include <omp.h>
+#include <algorithm>
+#include <atomic>
+#include <exception>
+#include <memory>
+#include <string>
+#include <typeinfo>
+
+#include <faiss/IndexFlat.h>
+#include <faiss/IndexScalarQuantizer.h>
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/DistanceComputer.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/FaissException.h>
+#include <faiss/impl/VisitedTable.h>
+
+namespace faiss {
+
+IndexPiPNN::IndexPiPNN() = default;
+
+IndexPiPNN::IndexPiPNN(Index* storage_in, int R)
+        : pipnn(R), storage(storage_in) {
+    FAISS_THROW_IF_NOT_MSG(storage, "IndexPiPNN has no storage");
+    check_storage_supported(*storage);
+    d = storage->d;
+    metric_type = storage->metric_type;
+    is_trained = storage->is_trained;
+    ntotal = storage->ntotal;
+}
+
+IndexPiPNN::~IndexPiPNN() {
+    if (own_fields) {
+        delete storage;
+    }
+}
+
+void IndexPiPNN::train(idx_t n, const float* x) {
+    FAISS_THROW_IF_NOT_MSG(storage, "IndexPiPNN has no storage");
+    storage->train(n, x);
+    is_trained = storage->is_trained;
+}
+
+void IndexPiPNN::add(idx_t n, const float* x) {
+    FAISS_THROW_IF_NOT_MSG(storage, "IndexPiPNN has no storage");
+    FAISS_THROW_IF_MSG(
+            is_built(), "IndexPiPNN does not support incremental addition");
+    storage->add(n, x);
+    ntotal = storage->ntotal;
+    if (!defer_build && ntotal > 0) {
+        build();
+    }
+}
+
+void IndexPiPNN::build() {
+    FAISS_THROW_IF_NOT_MSG(storage, "IndexPiPNN has no storage");
+    check_storage_supported(*storage);
+    FAISS_THROW_IF_NOT_MSG(
+            storage->ntotal > 0, "IndexPiPNN: cannot build over 0 vectors");
+    pipnn.build(*storage, verbose);
+    ntotal = storage->ntotal;
+}
+
+void IndexPiPNN::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT(k > 0);
+    FAISS_THROW_IF_NOT_MSG(storage, "IndexPiPNN has no storage");
+    FAISS_THROW_IF_NOT_MSG(
+            storage->is_trained,
+            "IndexPiPNN: storage must be trained before search()");
+    FAISS_THROW_IF_NOT_MSG(
+            pipnn.ntotal == ntotal && storage->ntotal == ntotal,
+            "IndexPiPNN: the index is not built over its storage; call "
+            "build()");
+    const auto* pipnn_params =
+            dynamic_cast<const SearchParametersPiPNN*>(params);
+    const int search_L = pipnn_params && pipnn_params->search_L != 0
+            ? pipnn_params->search_L
+            : pipnn.search_L;
+    const IDSelector* sel = params ? params->sel : nullptr;
+    const bool negate = metric_type == METRIC_INNER_PRODUCT;
+
+    const idx_t check_period =
+            InterruptCallback::get_period_hint(
+                    size_t(d) * size_t(std::max(idx_t(search_L), k))) *
+            omp_get_max_threads();
+
+    for (idx_t i0 = 0; i0 < n; i0 += check_period) {
+        const idx_t i1 = std::min(i0 + check_period, n);
+        std::exception_ptr ex;
+        std::atomic<bool> interrupt{false};
+
+#pragma omp parallel if (i1 - i0 > 1)
+        {
+            VisitedTable* vt = nullptr;
+            std::unique_ptr<DistanceComputer> dis;
+            try {
+                vt = &VisitedTable::get_reusable(ntotal, use_visited_hashset);
+                dis.reset(
+                        negate ? new NegativeDistanceComputer(
+                                         storage->get_distance_computer())
+                               : storage->get_distance_computer());
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
+
+#pragma omp for schedule(guided)
+            for (idx_t i = i0; i < i1; i++) {
+                if (interrupt.load(std::memory_order_relaxed)) {
+                    continue;
+                }
+                try {
+                    dis->set_query(x + i * d);
+                    pipnn.search(
+                            *dis,
+                            search_L,
+                            k,
+                            labels + i * k,
+                            distances + i * k,
+                            *vt,
+                            sel);
+                } catch (...) {
+                    omp_capture_exception(ex, [&] { interrupt = true; });
+                }
+            }
+        }
+        omp_rethrow_if_exception(ex);
+        InterruptCallback::check();
+    }
+
+    if (negate) {
+        for (idx_t i = 0; i < n * k; i++) {
+            distances[i] = -distances[i];
+        }
+    }
+}
+
+void IndexPiPNN::reconstruct(idx_t key, float* recons) const {
+    FAISS_THROW_IF_NOT_MSG(storage, "IndexPiPNN has no storage");
+    storage->reconstruct(key, recons);
+}
+
+void IndexPiPNN::reset() {
+    FAISS_THROW_IF_NOT_MSG(storage, "IndexPiPNN has no storage");
+    storage->reset();
+    pipnn.reset();
+    ntotal = 0;
+}
+
+void IndexPiPNN::check_storage_supported(const Index& storage) {
+    const std::type_info& t = typeid(storage);
+    bool ok = t == typeid(IndexFlat) || t == typeid(IndexFlatL2) ||
+            t == typeid(IndexFlatIP);
+    if (t == typeid(IndexScalarQuantizer)) {
+        const auto qt =
+                static_cast<const IndexScalarQuantizer&>(storage).sq.qtype;
+        ok = qt != ScalarQuantizer::QT_0bit &&
+                !ScalarQuantizer::TurboQuantRefine::is_turboq_full(qt);
+    }
+    ok = ok &&
+            (storage.metric_type == METRIC_L2 ||
+             storage.metric_type == METRIC_INNER_PRODUCT);
+    FAISS_THROW_IF_NOT_FMT(
+            ok,
+            "IndexPiPNN: unsupported storage %s; use IndexFlat or "
+            "IndexScalarQuantizer (not QT_0bit or QT_*bit_tq), L2 or IP",
+            demangle_cpp_symbol(t.name()).c_str());
+}
+
+IndexPiPNNFlat::IndexPiPNNFlat() = default;
+
+IndexPiPNNFlat::IndexPiPNNFlat(int d_in, int R, MetricType metric)
+        : IndexPiPNN(new IndexFlat(d_in, metric), R) {
+    own_fields = true;
+}
+
+IndexPiPNNSQ::IndexPiPNNSQ() = default;
+
+IndexPiPNNSQ::IndexPiPNNSQ(
+        int d_in,
+        ScalarQuantizer::QuantizerType qtype,
+        int R,
+        MetricType metric)
+        : IndexPiPNN(new IndexScalarQuantizer(d_in, qtype, metric), R) {
+    own_fields = true;
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexPiPNN.h

@@ -0,0 +1,83 @@
+/*
+ * 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 <optional>
+
+#include <faiss/Index.h>
+#include <faiss/impl/PiPNN.h>
+#include <faiss/impl/ScalarQuantizer.h>
+
+namespace faiss {
+
+struct SearchParametersPiPNN : SearchParameters {
+    /// beam width L; if 0, pipnn.search_L is used
+    int search_L = 0;
+};
+
+/** The PiPNN index is a normal random-access index with a PiPNN graph built
+ * on top of a flat storage (IndexFlat or IndexScalarQuantizer). */
+struct IndexPiPNN : Index {
+    PiPNN pipnn;
+    Index* storage = nullptr;
+    bool own_fields = false;
+    /// if true, add() only appends to storage; call build() afterwards
+    bool defer_build = false;
+    /// See impl/VisitedTable.h.
+    std::optional<bool> use_visited_hashset;
+
+    IndexPiPNN();
+    IndexPiPNN(Index* storage, int R = 64);
+
+    ~IndexPiPNN() override;
+
+    void train(idx_t n, const float* x) override;
+
+    void add(idx_t n, const float* x) override;
+
+    /// Builds the graph over all vectors in storage.
+    void build();
+
+    bool is_built() const {
+        return pipnn.ntotal > 0;
+    }
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    void reconstruct(idx_t key, float* recons) const override;
+
+    /// Empties storage (whatever own_fields says) and drops the graph.
+    void reset() override;
+
+    /// Throws unless IndexPiPNN supports the type and metric of `storage`.
+    static void check_storage_supported(const Index& storage);
+};
+
+/// IndexPiPNN over an owned IndexFlat.
+struct IndexPiPNNFlat : IndexPiPNN {
+    IndexPiPNNFlat();
+    IndexPiPNNFlat(int d, int R, MetricType metric = METRIC_L2);
+};
+
+/// IndexPiPNN over an owned IndexScalarQuantizer.
+struct IndexPiPNNSQ : IndexPiPNN {
+    IndexPiPNNSQ();
+    IndexPiPNNSQ(
+            int d,
+            ScalarQuantizer::QuantizerType qtype,
+            int R,
+            MetricType metric = METRIC_L2);
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexRaBitQFastScan.cpp

@@ -512,7 +512,8 @@ void IndexRaBitQFastScan::search(
         float* distances,
         idx_t* labels,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
+    // The kernel applies the selector after each SIMD block.
+    const IDSelector* rabitq_sel = params ? params->sel : nullptr;
 
     // Create query factors array on stack - memory managed by caller
     std::vector<rabitq_utils::QueryFactorsData> query_factors_storage(n);
@@ -523,9 +524,11 @@ void IndexRaBitQFastScan::search(
     FastScanDistancePostProcessing context;
     context.query_factors = query_factors_storage.data();
     if (metric_type == METRIC_L2) {
-        search_dispatch_implem<true>(n, x, k, distances, labels, context);
+        search_dispatch_implem<true>(
+                n, x, k, distances, labels, context, rabitq_sel);
     } else {
-        search_dispatch_implem<false>(n, x, k, distances, labels, context);
+        search_dispatch_implem<false>(
+                n, x, k, distances, labels, context, rabitq_sel);
     }
 }
 

data/vendor/faiss/faiss/IndexRaBitQFastScan.h

@@ -195,10 +195,12 @@ struct RaBitQHeapHandler
         float* const heap_dis = heap_distances + q * k;
         int64_t* const heap_ids = heap_labels + q * k;
 
-        rabitq_utils::QueryFactorsData query_factors_data = {};
-        if (context && context->query_factors != nullptr) {
-            query_factors_data = context->query_factors[q];
-        }
+        // Not a copy: QueryFactorsData owns a std::vector (malloc per block).
+        static const rabitq_utils::QueryFactorsData kNoQueryFactors{};
+        const rabitq_utils::QueryFactorsData& query_factors_data =
+                (context && context->query_factors != nullptr)
+                ? context->query_factors[q]
+                : kNoQueryFactors;
 
         const float one_a = normalizers ? (1.0f / normalizers[2 * q]) : 1.0f;
         const float bias = normalizers ? normalizers[2 * q + 1] : 0.0f;
@@ -218,6 +220,10 @@ struct RaBitQHeapHandler
 
         for (size_t i = 0; i < max_vectors; i++) {
             const size_t db_idx = base_db_idx + i;
+            if (this->sel != nullptr &&
+                !this->sel->is_member(static_cast<int64_t>(db_idx))) {
+                continue;
+            }
             const float normalized_distance = d32tab[i] * one_a + bias;
             const uint8_t* base_ptr = aux_base + i * storage_size;
 

data/vendor/faiss/faiss/IndexReplicas.cpp

@@ -111,6 +111,26 @@ void IndexReplicasTemplate<IndexT>::add(idx_t n, const component_t* x) {
     syncWithSubIndexes();
 }
 
+template <typename IndexT>
+void IndexReplicasTemplate<IndexT>::add_with_ids(
+        idx_t n,
+        const component_t* x,
+        const idx_t* xids) {
+    this->runOnIndex([n, x, xids](int, IndexT* index) {
+        index->add_with_ids(n, x, xids);
+    });
+    syncWithSubIndexes();
+}
+
+template <typename IndexT>
+size_t IndexReplicasTemplate<IndexT>::remove_ids(const IDSelector& sel) {
+    const idx_t before = this->ntotal;
+    this->runOnIndex([&sel](int, IndexT* index) { index->remove_ids(sel); });
+    // Throws when the replicas removed different counts.
+    syncWithSubIndexes();
+    return before - this->ntotal;
+}
+
 template <typename IndexT>
 void IndexReplicasTemplate<IndexT>::reconstruct(idx_t n, component_t* x) const {
     FAISS_THROW_IF_NOT_MSG(this->count() > 0, "no replicas in index");

data/vendor/faiss/faiss/IndexReplicas.h

@@ -56,6 +56,15 @@ class IndexReplicasTemplate : public ThreadedIndex<IndexT> {
     /// All indices receive the same call
     void add(idx_t n, const component_t* x) override;
 
+    /// faiss::Index API
+    /// All indices receive the same call
+    void add_with_ids(idx_t n, const component_t* x, const idx_t* xids)
+            override;
+
+    /// faiss::Index API
+    /// All indices receive the same call
+    size_t remove_ids(const IDSelector& sel) override;
+
     /// faiss::Index API
     /// Query is partitioned into a slice for each sub-index
     /// split by ceil(n / #indices) for our sub-indices