faiss 0.4.2 → 0.5.0

data/vendor/faiss/faiss/IndexIVFFlatPanorama.h

@@ -0,0 +1,64 @@
+/*
+ * 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++ -*-
+
+#ifndef FAISS_INDEX_IVF_FLAT_PANORAMA_H
+#define FAISS_INDEX_IVF_FLAT_PANORAMA_H
+
+#include <stdint.h>
+#include "faiss/IndexIVFFlat.h"
+
+namespace faiss {
+
+/// Panorama adaptation of IndexIVFFlat following
+/// https://www.arxiv.org/pdf/2510.00566.
+///
+/// IDEA:
+/// Panorama adapts the storage layout within each cluster and uses
+/// pruning with bounds to improve the search performance.
+/// Combined with orthogonal transforms upstream that concentrate the energy
+/// in the early dimensions (like PCA, Cayley, etc.), Panorama can prune up
+/// to 95% of the vectors in the cluster.
+///
+/// OVERHEAD:
+/// To be more efficient, we compute the residual energies at insertion time
+/// and store them along the vectors, which comes with an additional storage
+/// overhead of exactly (nlevels + 1) floats per vector. Add time is also
+/// slightly higher due to the overhead of transposing the vectors.
+///
+/// NOTE:
+/// We inherit from IndexIVFFlat instead of IndexIVF so we can keep the same
+/// insertion logic. The code responsible for level-oriented storage is in
+/// `ArrayInvertedListsPanorama`, which is a struct member of `IndexIVF`.
+struct IndexIVFFlatPanorama : IndexIVFFlat {
+    size_t n_levels;
+
+    std::vector<MaybeOwnedVector<float>> cum_sums;
+
+    explicit IndexIVFFlatPanorama(
+            Index* quantizer,
+            size_t d,
+            size_t nlist_,
+            int n_levels,
+            MetricType = METRIC_L2,
+            bool own_invlists = true);
+
+    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;
+
+    IndexIVFFlatPanorama();
+};
+
+} // namespace faiss
+
+#endif

data/vendor/faiss/faiss/IndexIVFPQ.cpp

@@ -1163,7 +1163,9 @@ struct IVFPQScannerT : QueryTables {
         }
 
 #pragma omp critical
-        { indexIVFPQ_stats.n_hamming_pass += n_hamming_pass; }
+        {
+            indexIVFPQ_stats.n_hamming_pass += n_hamming_pass;
+        }
     }
 
     template <class SearchResultType>

data/vendor/faiss/faiss/IndexIVFPQ.h

@@ -113,7 +113,7 @@ struct IndexIVFPQ : IndexIVF {
      */
     size_t find_duplicates(idx_t* ids, size_t* lims) const;
 
-    // map a vector to a binary code knowning the index
+    // map a vector to a binary code knowing the index
     void encode(idx_t key, const float* x, uint8_t* code) const;
 
     /** Encode multiple vectors

data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp

@@ -7,6 +7,7 @@
 
 #include <faiss/IndexIVFPQFastScan.h>
 
+#include <array>
 #include <cassert>
 #include <cstdio>
 
@@ -14,6 +15,7 @@
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/simdlib.h>
 
@@ -210,10 +212,11 @@ void IndexIVFPQFastScan::compute_LUT(
         const float* x,
         const CoarseQuantized& cq,
         AlignedTable<float>& dis_tables,
-        AlignedTable<float>& biases) const {
+        AlignedTable<float>& biases,
+        const FastScanDistancePostProcessing&) const {
     size_t dim12 = pq.ksub * pq.M;
     size_t d = pq.d;
-    size_t nprobe = this->nprobe;
+    size_t nprobe = cq.nprobe;
 
     if (by_residual) {
         if (metric_type == METRIC_L2) {
@@ -292,4 +295,133 @@ void IndexIVFPQFastScan::compute_LUT(
     }
 }
 
+/*********************************************************
+ * InvertedListScanner for IVFPQFS
+ *********************************************************/
+
+namespace {
+
+struct IVFPQFastScanScanner : InvertedListScanner {
+    static constexpr int impl = 10; // based on search_implem_10
+    static constexpr size_t nq = 1; // 1 query at a time.
+    const IndexIVFPQFastScan& index;
+    AlignedTable<uint8_t> dis_tables;
+    AlignedTable<uint16_t> biases;
+    std::array<float, 2> normalizers{};
+    const float* xi = nullptr;
+
+    IVFPQFastScanScanner(
+            const IndexIVFPQFastScan& index,
+            bool store_pairs,
+            const IDSelector* sel)
+            : InvertedListScanner(store_pairs, sel), index(index) {
+        this->keep_max = is_similarity_metric(index.metric_type);
+    }
+
+    void set_query(const float* query) override {
+        this->xi = query;
+    }
+
+    void set_list(idx_t list_no, float coarse_dis) override {
+        this->list_no = list_no;
+        IndexIVFFastScan::CoarseQuantized cq{
+                .nprobe = 1,        // 1 due to explicitly passing in list_no
+                .dis = &coarse_dis, // dis from query to list_no centroid.
+                .ids = &list_no,    // id of the current list we are scanning
+        };
+        FastScanDistancePostProcessing empty_context{};
+        index.compute_LUT_uint8(
+                1, xi, cq, dis_tables, biases, &normalizers[0], empty_context);
+    }
+
+    float distance_to_code(const uint8_t* /* code */) const override {
+        // It's not really possible to implement a distance_to_code since codes
+        // for 32 database vectors are intermixed.
+        FAISS_THROW_MSG("not implemented");
+    }
+
+    // Based on IVFFastScan search_implem_10, since it also deals with 1 query
+    // at a time.
+    size_t scan_codes(
+            size_t ntotal,
+            const uint8_t* codes,
+            const idx_t* ids,
+            float* distances,
+            idx_t* labels,
+            size_t k) const override {
+        // initialize the current iteration heap to the worst possible value of
+        // the prior loop
+        std::vector<float> curr_dists(k, distances[0]);
+        std::vector<idx_t> curr_labels(k, labels[0]);
+        FastScanDistancePostProcessing empty_context{};
+        std::unique_ptr<SIMDResultHandlerToFloat> handler(
+                index.make_knn_handler(
+                        !keep_max,
+                        impl,
+                        nq,
+                        k,
+                        curr_dists.data(),
+                        curr_labels.data(),
+                        sel,
+                        empty_context,
+                        &normalizers[0]));
+
+        // This does not quite match search_implem_10, but it is fine because
+        // the scanner operates on a single query at a time, and this value is
+        // used as the query index. For a single query, the value is always 0.
+        int qmap1[1] = {0};
+
+        handler->q_map = qmap1;
+        handler->begin(&normalizers[0]);
+
+        const uint8_t* LUT = dis_tables.get();
+        handler->dbias = biases.get();
+
+        handler->ntotal = ntotal;
+        handler->id_map = ids;
+
+        pq4_accumulate_loop(
+                1,
+                roundup(ntotal, index.bbs),
+                index.bbs,
+                static_cast<int>(index.M2),
+                codes,
+                LUT,
+                *handler,
+                nullptr);
+
+        // The handler is for the results of this iteration.
+        // Then we need a second heap to combine across iterations.
+        handler->end();
+        if (keep_max) {
+            minheap_addn(
+                    k,
+                    distances,
+                    labels,
+                    curr_dists.data(),
+                    curr_labels.data(),
+                    k);
+        } else {
+            maxheap_addn(
+                    k,
+                    distances,
+                    labels,
+                    curr_dists.data(),
+                    curr_labels.data(),
+                    k);
+        }
+
+        return handler->num_updates();
+    }
+};
+
+} // anonymous namespace
+
+InvertedListScanner* IndexIVFPQFastScan::get_InvertedListScanner(
+        bool store_pairs,
+        const IDSelector* sel,
+        const IVFSearchParameters*) const {
+    return new IVFPQFastScanScanner(*this, store_pairs, sel);
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexIVFPQFastScan.h

@@ -80,7 +80,13 @@ struct IndexIVFPQFastScan : IndexIVFFastScan {
             const float* x,
             const CoarseQuantized& cq,
             AlignedTable<float>& dis_tables,
-            AlignedTable<float>& biases) const override;
+            AlignedTable<float>& biases,
+            const FastScanDistancePostProcessing& context) const override;
+
+    InvertedListScanner* get_InvertedListScanner(
+            bool store_pairs,
+            const IDSelector* sel,
+            const IVFSearchParameters*) const override;
 };
 
 } // namespace faiss

data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp

@@ -155,15 +155,18 @@ struct RaBitInvertedListScanner : InvertedListScanner {
     std::unique_ptr<FlatCodesDistanceComputer> dc;
 
     uint8_t qb = 0;
+    bool centered = false;
 
-    RaBitInvertedListScanner(
+    explicit RaBitInvertedListScanner(
             const IndexIVFRaBitQ& ivf_rabitq_in,
             bool store_pairs = false,
             const IDSelector* sel = nullptr,
-            uint8_t qb_in = 0)
+            uint8_t qb_in = 0,
+            bool centered = false)
             : InvertedListScanner(store_pairs, sel),
               ivf_rabitq{ivf_rabitq_in},
-              qb{qb_in} {
+              qb{qb_in},
+              centered(centered) {
         keep_max = is_similarity_metric(ivf_rabitq.metric_type);
         code_size = ivf_rabitq.code_size;
     }
@@ -196,7 +199,7 @@ struct RaBitInvertedListScanner : InvertedListScanner {
             // both query_vector and centroid are available!
             // set up DistanceComputer
             dc.reset(ivf_rabitq.rabitq.get_distance_computer(
-                    qb, reconstructed_centroid.data()));
+                    qb, reconstructed_centroid.data(), centered));
 
             dc->set_query(query_vector.data());
         }
@@ -208,12 +211,15 @@ InvertedListScanner* IndexIVFRaBitQ::get_InvertedListScanner(
         const IDSelector* sel,
         const IVFSearchParameters* search_params_in) const {
     uint8_t used_qb = qb;
+    bool centered = false;
     if (auto params = dynamic_cast<const IVFRaBitQSearchParameters*>(
                 search_params_in)) {
         used_qb = params->qb;
+        centered = params->centered;
     }
 
-    return new RaBitInvertedListScanner(*this, store_pairs, sel, used_qb);
+    return new RaBitInvertedListScanner(
+            *this, store_pairs, sel, used_qb, centered);
 }
 
 void IndexIVFRaBitQ::reconstruct_from_offset(
@@ -278,7 +284,8 @@ float IVFRaBitDistanceComputer::operator()(idx_t i) {
     float distance = 0;
 
     std::unique_ptr<FlatCodesDistanceComputer> dc(
-            parent->rabitq.get_distance_computer(parent->qb, centroid.data()));
+            parent->rabitq.get_distance_computer(
+                    parent->qb, centroid.data(), /*centered=*/false));
     dc->set_query(q);
     distance = dc->distance_to_code(code);
 

data/vendor/faiss/faiss/IndexIVFRaBitQ.h

@@ -19,6 +19,7 @@ namespace faiss {
 
 struct IVFRaBitQSearchParameters : IVFSearchParameters {
     uint8_t qb = 0;
+    bool centered = false;
 };
 
 // * by_residual is true, just by design