faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp

@@ -8,44 +8,51 @@
 #include <faiss/IndexIVFPQFastScan.h>
 
 #include <array>
-#include <cassert>
 #include <cstdio>
 
 #include <memory>
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
-#include <faiss/utils/Heap.h>
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/simdlib/simdlib_dispatch.h>
 #include <faiss/utils/distances.h>
-#include <faiss/utils/simdlib.h>
+#include <faiss/utils/distances_dispatch.h>
+#include <faiss/utils/extra_distances.h>
 
 #include <faiss/invlists/BlockInvertedLists.h>
 
-#include <faiss/impl/pq4_fast_scan.h>
-#include <faiss/impl/simd_result_handlers.h>
+#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
+#include <faiss/impl/fast_scan/fast_scan.h>
+#include <faiss/impl/fast_scan/simd_result_handlers.h>
 
 namespace faiss {
 
-using namespace simd_result_handlers;
-
 inline size_t roundup(size_t a, size_t b) {
     return (a + b - 1) / b * b;
 }
 
 IndexIVFPQFastScan::IndexIVFPQFastScan(
-        Index* quantizer,
-        size_t d,
-        size_t nlist,
-        size_t M,
-        size_t nbits,
+        Index* quantizer_in,
+        size_t d_in,
+        size_t nlist_in,
+        size_t M_in,
+        size_t nbits_in,
         MetricType metric,
-        int bbs,
-        bool own_invlists)
-        : IndexIVFFastScan(quantizer, d, nlist, 0, metric, own_invlists),
-          pq(d, M, nbits) {
+        int bbs_in,
+        bool own_invlists_in)
+        : IndexIVFFastScan(
+                  quantizer_in,
+                  d_in,
+                  nlist_in,
+                  0,
+                  metric,
+                  own_invlists_in),
+          pq(d_in, M_in, nbits_in) {
     by_residual = false; // set to false by default because it's faster
 
-    init_fastscan(&pq, M, nbits, nlist, metric, bbs, own_invlists);
+    init_fastscan(
+            &pq, M_in, nbits_in, nlist_in, metric, bbs_in, own_invlists_in);
 }
 
 IndexIVFPQFastScan::IndexIVFPQFastScan() {
@@ -54,7 +61,7 @@ IndexIVFPQFastScan::IndexIVFPQFastScan() {
     M2 = 0;
 }
 
-IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
+IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs_in)
         : IndexIVFFastScan(
                   orig.quantizer,
                   orig.d,
@@ -71,7 +78,7 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
             orig.pq.nbits,
             orig.nlist,
             orig.metric_type,
-            bbs,
+            bbs_in,
             orig.own_invlists);
 
     by_residual = orig.by_residual;
@@ -88,7 +95,7 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
     }
 
 #pragma omp parallel for if (nlist > 100)
-    for (idx_t i = 0; i < nlist; i++) {
+    for (idx_t i = 0; i < static_cast<idx_t>(nlist); i++) {
         size_t nb = orig.invlists->list_size(i);
         size_t nb2 = roundup(nb, bbs);
         AlignedTable<uint8_t> tmp(nb2 * M2 / 2);
@@ -110,6 +117,10 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
     orig_invlists = orig.invlists;
 }
 
+size_t IndexIVFPQFastScan::fast_scan_code_size() const {
+    return M2 / 2;
+}
+
 /*********************************************************
  * Training
  *********************************************************/
@@ -117,7 +128,7 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
 void IndexIVFPQFastScan::train_encoder(
         idx_t n,
         const float* x,
-        const idx_t* assign) {
+        const idx_t* /*assign*/) {
     pq.verbose = verbose;
     pq.train(n, x);
 
@@ -152,7 +163,7 @@ void IndexIVFPQFastScan::encode_vectors(
         bool include_listnos) const {
     if (by_residual) {
         AlignedTable<float> residuals(n * d);
-        for (size_t i = 0; i < n; i++) {
+        for (idx_t i = 0; i < n; i++) {
             if (list_nos[i] < 0) {
                 memset(residuals.data() + i * d, 0, sizeof(residuals[0]) * d);
             } else {
@@ -179,16 +190,19 @@ void IndexIVFPQFastScan::encode_vectors(
  * Look-Up Table functions
  *********************************************************/
 
+// Explicit SIMD-level alias (no global bare aliases).
+using simd8float32 = simd8float32_tpl<SINGLE_SIMD_LEVEL_256>;
+
 void fvec_madd_simd(
         size_t n,
         const float* a,
         float bf,
         const float* b,
         float* c) {
-    assert(is_aligned_pointer(a));
-    assert(is_aligned_pointer(b));
-    assert(is_aligned_pointer(c));
-    assert(n % 8 == 0);
+    FAISS_THROW_IF_NOT_MSG(is_aligned_pointer(a), "pointer a is not aligned");
+    FAISS_THROW_IF_NOT_MSG(is_aligned_pointer(b), "pointer b is not aligned");
+    FAISS_THROW_IF_NOT_MSG(is_aligned_pointer(c), "pointer c is not aligned");
+    FAISS_THROW_IF_NOT_MSG(n % 8 == 0, "n must be a multiple of 8");
     simd8float32 bf8(bf);
     n /= 8;
     for (size_t i = 0; i < n; i++) {
@@ -215,23 +229,24 @@ void IndexIVFPQFastScan::compute_LUT(
         AlignedTable<float>& biases,
         const FastScanDistancePostProcessing&) const {
     size_t dim12 = pq.ksub * pq.M;
-    size_t d = pq.d;
-    size_t nprobe = cq.nprobe;
+    size_t pq_d = pq.d;
+    size_t cq_nprobe = cq.nprobe;
 
     if (by_residual) {
         if (metric_type == METRIC_L2) {
-            dis_tables.resize(n * nprobe * dim12);
+            dis_tables.resize(n * cq_nprobe * dim12);
 
             if (use_precomputed_table == 1) {
-                biases.resize(n * nprobe);
-                memcpy(biases.get(), cq.dis, sizeof(float) * n * nprobe);
+                biases.resize(n * cq_nprobe);
+                memcpy(biases.get(), cq.dis, sizeof(float) * n * cq_nprobe);
 
                 AlignedTable<float> ip_table(n * dim12);
                 pq.compute_inner_prod_tables(n, x, ip_table.get());
 
-#pragma omp parallel for if (n * nprobe > 8000)
-                for (idx_t ij = 0; ij < n * nprobe; ij++) {
-                    idx_t i = ij / nprobe;
+#pragma omp parallel for if (n * cq_nprobe > 8000)
+                for (idx_t ij = 0; ij < static_cast<idx_t>(n * cq_nprobe);
+                     ij++) {
+                    idx_t i = ij / cq_nprobe;
                     float* tab = dis_tables.get() + ij * dim12;
                     idx_t cij = cq.ids[ij];
 
@@ -250,26 +265,27 @@ void IndexIVFPQFastScan::compute_LUT(
                 }
 
             } else {
-                std::unique_ptr<float[]> xrel(new float[n * nprobe * d]);
-                biases.resize(n * nprobe);
-                memset(biases.get(), 0, sizeof(float) * n * nprobe);
-
-#pragma omp parallel for if (n * nprobe > 8000)
-                for (idx_t ij = 0; ij < n * nprobe; ij++) {
-                    idx_t i = ij / nprobe;
-                    float* xij = &xrel[ij * d];
+                std::unique_ptr<float[]> xrel(new float[n * cq_nprobe * pq_d]);
+                biases.resize(n * cq_nprobe);
+                memset(biases.get(), 0, sizeof(float) * n * cq_nprobe);
+
+#pragma omp parallel for if (n * cq_nprobe > 8000)
+                for (idx_t ij = 0; ij < static_cast<idx_t>(n * cq_nprobe);
+                     ij++) {
+                    idx_t i = ij / cq_nprobe;
+                    float* xij = &xrel[ij * pq_d];
                     idx_t cij = cq.ids[ij];
 
                     if (cij >= 0) {
-                        quantizer->compute_residual(x + i * d, xij, cij);
+                        quantizer->compute_residual(x + i * pq_d, xij, cij);
                     } else {
                         // will fill with NaNs
-                        memset(xij, -1, sizeof(float) * d);
+                        memset(xij, -1, sizeof(float) * pq_d);
                     }
                 }
 
                 pq.compute_distance_tables(
-                        n * nprobe, xrel.get(), dis_tables.get());
+                        n * cq_nprobe, xrel.get(), dis_tables.get());
             }
 
         } else if (metric_type == METRIC_INNER_PRODUCT) {
@@ -277,8 +293,8 @@ void IndexIVFPQFastScan::compute_LUT(
             pq.compute_inner_prod_tables(n, x, dis_tables.get());
             // compute_inner_prod_tables(pq, n, x, dis_tables.get());
 
-            biases.resize(n * nprobe);
-            memcpy(biases.get(), cq.dis, sizeof(float) * n * nprobe);
+            biases.resize(n * cq_nprobe);
+            memcpy(biases.get(), cq.dis, sizeof(float) * n * cq_nprobe);
         } else {
             FAISS_THROW_FMT("metric %d not supported", metric_type);
         }
@@ -302,42 +318,74 @@ void IndexIVFPQFastScan::compute_LUT(
 namespace {
 
 struct IVFPQFastScanScanner : InvertedListScanner {
-    static constexpr int impl = 10; // based on search_implem_10
+    using InvertedListScanner::scan_codes;
+    [[maybe_unused]] 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::vector<float> residual;
     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);
+            const IndexIVFPQFastScan& index_in,
+            bool store_pairs_in,
+            const IDSelector* sel_in)
+            : InvertedListScanner(store_pairs_in, sel_in), index(index_in) {
+        this->keep_max = is_similarity_metric(index_in.metric_type);
+        residual.resize(index_in.d);
     }
 
     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;
+    void set_list(idx_t list_no_in, float coarse_dis_in) override {
+        this->list_no = list_no_in;
         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
+                .nprobe = 1,           // 1 due to explicitly passing in list_no
+                .dis = &coarse_dis_in, // dis from query to list_no centroid.
+                .ids = &list_no_in,    // 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);
+        // used in distance_to_code
+        index.quantizer->compute_residual(
+                this->xi, residual.data(), this->list_no);
     }
 
-    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");
+    float distance_to_code(const uint8_t* code) const override {
+        // directly use the PQ tables to compute the distance
+        const ProductQuantizer& pq = index.pq;
+        // when by_residual, codes are residuals so compare against query
+        // residual; otherwise codes are raw vectors so compare against raw
+        // query
+        const float* x = index.by_residual ? residual.data() : this->xi;
+        float accu = 0;
+        // implemented for all vector distances, although only L2 and IP are
+        // supported by FastScan
+        with_VectorDistance(pq.dsub, index.metric_type, 0.0, [&](auto vd) {
+            int m;
+            for (m = 0; m + 1 < pq.M; m += 2) {
+                const float* cent;
+                uint8_t c = *code++;
+                cent = pq.get_centroids(m, c & 15);
+                accu += vd(cent, x);
+                x += pq.dsub;
+                cent = pq.get_centroids(m + 1, c >> 4);
+                accu += vd(cent, x);
+                x += pq.dsub;
+            }
+            if (m < pq.M) { // leftover
+                uint8_t c = *code++;
+                const float* cent = pq.get_centroids(m, c & 15);
+                accu += vd(cent, x);
+            }
+        });
+        return accu;
     }
 
     // Based on IVFFastScan search_implem_10, since it also deals with 1 query
@@ -346,72 +394,86 @@ struct IVFPQFastScanScanner : InvertedListScanner {
             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(
+            ResultHandler& handler) const override {
+        auto scan_with_heap = [&](auto* heap_handler) -> size_t {
+            const size_t k = heap_handler->k;
+            if (k == 0) {
+                return 0;
+            }
+
+            // initialize the current iteration heap to the worst possible value
+            // of the caller-owned result handler.
+            std::vector<float> curr_dists(k, handler.threshold);
+            std::vector<idx_t> curr_labels(k, -1);
+
+            auto scanner = index.make_knn_scanner(
+                    !keep_max,
+                    nq,
                     k,
-                    distances,
-                    labels,
                     curr_dists.data(),
                     curr_labels.data(),
-                    k);
+                    sel);
+
+            SIMDResultHandlerToFloat* rh = scanner->handler();
+
+            // 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};
+
+            rh->q_map = qmap1;
+            rh->begin(&normalizers[0]);
+
+            rh->dbias = biases.get();
+            rh->ntotal = ntotal;
+            rh->id_map = ids;
+
+            scanner->accumulate_loop(
+                    1,
+                    roundup(ntotal, index.bbs),
+                    index.bbs,
+                    static_cast<int>(index.M2),
+                    codes,
+                    dis_tables.get(),
+                    0,
+                    index.get_block_stride());
+
+            const size_t scan_cnt = rh->count_scanned_rows();
+            rh->end();
+
+            handler.stats.scan_cnt += scan_cnt;
+            size_t nup = 0;
+            for (size_t j = 0; j < k; j++) {
+                if (curr_labels[j] < 0) {
+                    continue;
+                }
+                if (handler.add_result(curr_dists[j], curr_labels[j])) {
+                    handler.stats.nheap_updates++;
+                    nup++;
+                }
+            }
+            return nup;
+        };
+
+        if (!keep_max) {
+            using C = CMax<float, idx_t>;
+            if (auto* heap_handler =
+                        dynamic_cast<HeapResultHandler<C, false>*>(&handler)) {
+                return scan_with_heap(heap_handler);
+            }
         } else {
-            maxheap_addn(
-                    k,
-                    distances,
-                    labels,
-                    curr_dists.data(),
-                    curr_labels.data(),
-                    k);
+            using C = CMin<float, idx_t>;
+            if (auto* heap_handler =
+                        dynamic_cast<HeapResultHandler<C, false>*>(&handler)) {
+                return scan_with_heap(heap_handler);
+            }
         }
 
-        return handler->num_updates();
+        FAISS_THROW_MSG(
+                "IVFPQFastScanScanner::scan_codes requires HeapResultHandler; "
+                "custom ResultHandler scan is not supported by this optimized "
+                "scanner");
     }
 };
 

data/vendor/faiss/faiss/IndexIVFPQFastScan.h

@@ -55,6 +55,9 @@ struct IndexIVFPQFastScan : IndexIVFFastScan {
     // built from an IndexIVFPQ
     explicit IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs = 32);
 
+    /// Packed code size: M2 / 2 bytes (4-bit PQ sub-quantizer nibbles)
+    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;

data/vendor/faiss/faiss/IndexIVFPQR.cpp

@@ -11,6 +11,7 @@
 
 #include <cinttypes>
 
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/utils.h>
@@ -24,23 +25,23 @@ namespace faiss {
  ******************************************/
 
 IndexIVFPQR::IndexIVFPQR(
-        Index* quantizer,
-        size_t d,
-        size_t nlist,
+        Index* quantizer_in,
+        size_t d_in,
+        size_t nlist_in,
         size_t M,
         size_t nbits_per_idx,
         size_t M_refine,
         size_t nbits_per_idx_refine,
-        bool own_invlists)
+        bool own_invlists_in)
         : IndexIVFPQ(
-                  quantizer,
-                  d,
-                  nlist,
+                  quantizer_in,
+                  d_in,
+                  nlist_in,
                   M,
                   nbits_per_idx,
                   METRIC_L2,
-                  own_invlists),
-          refine_pq(d, M_refine, nbits_per_idx_refine),
+                  own_invlists_in),
+          refine_pq(d_in, M_refine, nbits_per_idx_refine),
           k_factor(4) {
     by_residual = true;
     refine_pq.cp.max_points_per_centroid = 1000;
@@ -125,10 +126,10 @@ void IndexIVFPQR::search_preassigned(
         idx_t* labels,
         bool store_pairs,
         const IVFSearchParameters* params,
-        IndexIVFStats* stats) const {
+        IndexIVFStats* /*stats*/) const {
     uint64_t t0;
     TIC;
-    size_t k_coarse = long(k * k_factor);
+    size_t k_coarse = long((size_t)k * k_factor);
     std::unique_ptr<idx_t[]> coarse_labels(new idx_t[k_coarse * n]);
     {
         // query with quantizer levels 1 and 2.
@@ -165,18 +166,22 @@ void IndexIVFPQR::search_preassigned(
             idx_t* heap_ids = labels + k * i;
             maxheap_heapify(k, heap_sim, heap_ids);
 
-            for (int j = 0; j < k_coarse; j++) {
+            for (size_t j = 0; j < k_coarse; j++) {
                 idx_t sl = shortlist[j];
 
                 if (sl == -1) {
                     continue;
                 }
 
-                int list_no = lo_listno(sl);
-                int ofs = lo_offset(sl);
+                int list_no = static_cast<int>(lo_listno(sl));
+                int ofs = static_cast<int>(lo_offset(sl));
 
-                assert(list_no >= 0 && list_no < nlist);
-                assert(ofs >= 0 && ofs < invlists->list_size(list_no));
+                FAISS_THROW_IF_NOT(
+                        list_no >= 0 && static_cast<size_t>(list_no) < nlist);
+                FAISS_THROW_IF_NOT(
+                        ofs >= 0 &&
+                        static_cast<size_t>(ofs) <
+                                invlists->list_size(list_no));
 
                 // 1st level residual
                 quantizer->compute_residual(xq, residual_1.get(), list_no);
@@ -191,7 +196,7 @@ void IndexIVFPQR::search_preassigned(
 
                 // 3rd level residual's approximation
                 idx_t id = invlists->get_single_id(list_no, ofs);
-                assert(0 <= id && id < ntotal);
+                FAISS_THROW_IF_NOT(0 <= id && id < ntotal);
                 refine_pq.decode(
                         &refine_codes[id * refine_pq.code_size],
                         residual_1.get());
@@ -218,7 +223,7 @@ void IndexIVFPQR::reconstruct_from_offset(
     IndexIVFPQ::reconstruct_from_offset(list_no, offset, recons);
 
     idx_t id = invlists->get_single_id(list_no, offset);
-    assert(0 <= id && id < ntotal);
+    FAISS_THROW_IF_NOT(0 <= id && id < ntotal);
 
     std::vector<float> r3(d);
     refine_pq.decode(&refine_codes[id * refine_pq.code_size], r3.data());