faiss 0.6.4 → 0.6.5

data/vendor/faiss/faiss/IndexIVFFastScan.cpp

@@ -21,6 +21,7 @@
 #include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
 #include <faiss/impl/fast_scan/fast_scan.h>
 #include <faiss/impl/fast_scan/simd_result_handlers.h>
+#include <faiss/impl/platform_macros.h>
 #include <faiss/invlists/BlockInvertedLists.h>
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/quantize_lut.h>
@@ -646,8 +647,23 @@ void IndexIVFFastScan::search_dispatch_implem(
                 search_implem_14(
                         n, x, k, distances, labels, cq, impl, context, params);
             } else {
+                // Size each thread's LUT here, on the calling thread: a
+                // heap-profiler sample taken on an OpenMP worker unwinds
+                // through libomp, which is very slow on aarch64. The tables
+                // keep their capacity, so smaller slices do not reallocate.
+                std::vector<AlignedTable<uint8_t>> thread_dis_tables(
+                        omp_get_max_threads());
+                const size_t max_slice_n = (n + nslice - 1) / nslice;
+                const size_t lut_per_query =
+                        ksub * M2 * (lookup_table_is_3d() ? cur_nprobe : 1);
+                for (auto& t : thread_dis_tables) {
+                    t.keep_capacity = true;
+                    t.resize(max_slice_n * lut_per_query);
+                }
 #pragma omp parallel for reduction(+ : ndis, nlist_visited)
                 for (int slice = 0; slice < nslice; slice++) {
+                    AlignedTable<uint8_t>* dis_tables_buf =
+                            &thread_dis_tables.at(omp_get_thread_num());
                     idx_t i0 = n * slice / nslice;
                     idx_t i1 = n * (slice + 1) / nslice;
                     float* dis_i = distances + i0 * k;
@@ -663,6 +679,9 @@ void IndexIVFFastScan::search_dispatch_implem(
                     if (thread_context.query_factors != nullptr) {
                         thread_context.query_factors += i0 * cur_nprobe;
                     }
+                    if (thread_context.rotated_q != nullptr) {
+                        thread_context.rotated_q += i0 * cur_nprobe * d;
+                    }
 
                     auto scanner = make_knn_scanner(
                             is_max,
@@ -684,7 +703,8 @@ void IndexIVFFastScan::search_dispatch_implem(
                                 &nlist_visited,
                                 thread_context,
                                 params,
-                                *scanner);
+                                *scanner,
+                                dis_tables_buf);
                     } else {
                         search_implem_10(
                                 i1 - i0,
@@ -696,7 +716,8 @@ void IndexIVFFastScan::search_dispatch_implem(
                                 &nlist_visited,
                                 thread_context,
                                 params,
-                                *scanner);
+                                *scanner,
+                                dis_tables_buf);
                     }
                 }
             }
@@ -1020,9 +1041,12 @@ void IndexIVFFastScan::search_implem_10(
         size_t* nlist_out,
         const FastScanDistancePostProcessing& context,
         const IVFSearchParameters* params,
-        FastScanCodeScanner& scanner) const {
+        FastScanCodeScanner& scanner,
+        AlignedTable<uint8_t>* dis_tables_buf) const {
     size_t dim12 = ksub * M2;
-    AlignedTable<uint8_t> dis_tables;
+    AlignedTable<uint8_t> local_dis_tables;
+    AlignedTable<uint8_t>& dis_tables =
+            dis_tables_buf ? *dis_tables_buf : local_dis_tables;
     AlignedTable<uint16_t> biases;
     std::unique_ptr<float[]> normalizers(new float[2 * n]);
 
@@ -1060,6 +1084,7 @@ void IndexIVFFastScan::search_implem_10(
     // Allocate probe_map once and reuse it
     std::vector<int> probe_map;
     probe_map.reserve(1);
+    const size_t block_stride = get_block_stride();
 
     for (idx_t i = 0; i < n; i++) {
         const uint8_t* LUT = nullptr;
@@ -1134,7 +1159,7 @@ void IndexIVFFastScan::search_implem_10(
                     codes.get(),
                     LUT,
                     context.pq2x4_scale,
-                    get_block_stride());
+                    block_stride);
 
             ndis += ls;
             nlist_visited++;
@@ -1172,14 +1197,17 @@ void IndexIVFFastScan::search_implem_12(
         size_t* nlist_out,
         const FastScanDistancePostProcessing& context,
         const IVFSearchParameters* /* params */,
-        FastScanCodeScanner& scanner) const {
+        FastScanCodeScanner& scanner,
+        AlignedTable<uint8_t>* dis_tables_buf) const {
     if (n == 0) { // does not work well with reservoir
         return;
     }
     FAISS_THROW_IF_NOT(bbs == 32);
 
     size_t dim12 = ksub * M2;
-    AlignedTable<uint8_t> dis_tables;
+    AlignedTable<uint8_t> local_dis_tables;
+    AlignedTable<uint8_t>& dis_tables =
+            dis_tables_buf ? *dis_tables_buf : local_dis_tables;
     AlignedTable<uint16_t> biases;
     std::unique_ptr<float[]> normalizers(new float[2 * n]);
 
@@ -1196,6 +1224,7 @@ void IndexIVFFastScan::search_implem_12(
     size_t cur_nprobe = cq.nprobe;
 
     std::vector<QC> qcs;
+    qcs.reserve(n * cur_nprobe);
     {
         size_t ij = 0;
         for (idx_t i = 0; i < n; i++) {
@@ -1226,6 +1255,9 @@ void IndexIVFFastScan::search_implem_12(
     // Allocate vectors once and reuse them
     std::vector<int> probe_map;
     probe_map.reserve(actual_qbs2);
+    std::vector<int> q_map(actual_qbs2), lut_entries(actual_qbs2);
+    AlignedTable<uint8_t> LUT(actual_qbs2 * dim12);
+    const size_t block_stride = get_block_stride();
 
     size_t i0 = 0;
     uint64_t t_copy_pack = 0, t_scan = 0;
@@ -1252,8 +1284,6 @@ void IndexIVFFastScan::search_implem_12(
         // re-organize LUTs and biases into the right order
         int nc = static_cast<int>(i1 - i0);
 
-        std::vector<int> q_map(nc), lut_entries(nc);
-        AlignedTable<uint8_t> LUT(nc * dim12);
         memset(LUT.get(), -1, nc * dim12);
         int qbs_for_list = pq4_preferred_qbs(nc);
 
@@ -1303,7 +1333,7 @@ void IndexIVFFastScan::search_implem_12(
                 codes.get(),
                 LUT.get(),
                 context.pq2x4_scale,
-                get_block_stride());
+                block_stride);
         // prepare for next loop
         i0 = i1;
     }
@@ -1351,6 +1381,7 @@ void IndexIVFFastScan::search_implem_14(
     size_t cur_nprobe = cq.nprobe;
 
     std::vector<QC> qcs;
+    qcs.reserve(n * cur_nprobe);
     {
         size_t ij = 0;
         for (idx_t i = 0; i < n; i++) {
@@ -1461,8 +1492,11 @@ void IndexIVFFastScan::search_implem_14(
         // Allocate probe_map once per thread and reuse it
         std::vector<int> probe_map;
         probe_map.reserve(actual_qbs2);
+        std::vector<int> q_map(actual_qbs2), lut_entries(actual_qbs2);
+        AlignedTable<uint8_t> LUT(actual_qbs2 * dim12);
+        const size_t block_stride = get_block_stride();
 
-#pragma omp for schedule(dynamic)
+#pragma omp for schedule(FAISS_OMP_DYNAMIC)
         for (idx_t cluster = 0; cluster < static_cast<idx_t>(ses.size());
              cluster++) {
             size_t i0 = ses[cluster].start;
@@ -1474,8 +1508,6 @@ void IndexIVFFastScan::search_implem_14(
             // re-organize LUTs and biases into the right order
             int nc = static_cast<int>(i1 - i0);
 
-            std::vector<int> q_map(nc), lut_entries(nc);
-            AlignedTable<uint8_t> LUT(nc * dim12);
             memset(LUT.get(), -1, nc * dim12);
             int qbs_for_list = pq4_preferred_qbs(nc);
 
@@ -1526,7 +1558,7 @@ void IndexIVFFastScan::search_implem_14(
                     codes.get(),
                     LUT.get(),
                     context.pq2x4_scale,
-                    get_block_stride());
+                    block_stride);
         }
 
         // labels is in-place for HeapHC

data/vendor/faiss/faiss/IndexIVFFastScan.h

@@ -303,7 +303,8 @@ struct IndexIVFFastScan : IndexIVF {
             size_t* nlist_out,
             const FastScanDistancePostProcessing& context,
             const IVFSearchParameters* params,
-            FastScanCodeScanner& scanner) const;
+            FastScanCodeScanner& scanner,
+            AlignedTable<uint8_t>* dis_tables_buf = nullptr) const;
 
     void search_implem_12(
             idx_t n,
@@ -314,7 +315,8 @@ struct IndexIVFFastScan : IndexIVF {
             size_t* nlist_out,
             const FastScanDistancePostProcessing& context,
             const IVFSearchParameters* params,
-            FastScanCodeScanner& scanner) const;
+            FastScanCodeScanner& scanner,
+            AlignedTable<uint8_t>* dis_tables_buf = nullptr) const;
 
     // implem 14 is multithreaded internally across nprobes and queries
     void search_implem_14(

data/vendor/faiss/faiss/IndexIVFFlat.cpp

@@ -57,6 +57,25 @@ IndexIVFFlat::IndexIVFFlat() {
     by_residual = false;
 }
 
+void IndexIVFFlat::train_encoded(
+        idx_t n,
+        const uint8_t* x,
+        const Index* codec) {
+    FAISS_THROW_IF_NOT_MSG(codec, "encoded training requires a codec");
+    train_q1_encoded(n, x, codec, verbose, metric_type);
+    is_trained = true;
+}
+
+void IndexIVFFlat::train_ex(idx_t n, const void* x, NumericType numeric_type) {
+    if (numeric_type != NumericType::Float16) {
+        IndexIVF::train_ex(n, x, numeric_type);
+        return;
+    }
+    FAISS_THROW_IF_NOT_MSG(x, "training data must not be null");
+    train_q1_ex(n, x, numeric_type, verbose, metric_type);
+    is_trained = true;
+}
+
 void IndexIVFFlat::add_core(
         idx_t n,
         const float* x,
@@ -186,6 +205,20 @@ IndexIVFFlatDedup::IndexIVFFlatDedup(
                   metric_type_in,
                   own_invlists_in) {}
 
+void IndexIVFFlatDedup::train_encoded(idx_t, const uint8_t*, const Index*) {
+    FAISS_THROW_MSG("encoded training is not supported by IndexIVFFlatDedup");
+}
+
+void IndexIVFFlatDedup::train_ex(
+        idx_t n,
+        const void* x,
+        NumericType numeric_type) {
+    FAISS_THROW_IF_NOT_MSG(
+            numeric_type == NumericType::Float32,
+            "only fp32 training is supported by IndexIVFFlatDedup");
+    train(n, static_cast<const float*>(x));
+}
+
 void IndexIVFFlatDedup::train(idx_t n, const float* x) {
     std::unordered_map<uint64_t, idx_t> map;
     std::unique_ptr<float[]> x2(new float[n * d]);

data/vendor/faiss/faiss/IndexIVFFlat.h

@@ -29,6 +29,16 @@ struct IndexIVFFlat : IndexIVF {
             MetricType = METRIC_L2,
             bool own_invlists = true);
 
+    /** Train the coarse quantizer from vectors decoded on demand by `codec`.
+     *
+     * IndexIVFFlat has no trainable residual encoder, so this avoids a
+     * corpus-sized decoded buffer. The codec must remain alive for this call.
+     */
+    void train_encoded(idx_t n, const uint8_t* x, const Index* codec) override;
+
+    /// Float16 trains only the coarse quantizer, natively on the fp16 rows.
+    void train_ex(idx_t n, const void* x, NumericType numeric_type) override;
+
     void add_core(
             idx_t n,
             const float* x,
@@ -110,6 +120,12 @@ struct IndexIVFFlatDedup : IndexIVFFlat {
     /// also dedups the training set
     void train(idx_t n, const float* x) override;
 
+    /// Encoded training cannot preserve the byte-exact deduplication contract.
+    void train_encoded(idx_t n, const uint8_t* x, const Index* codec) override;
+
+    /// Non-fp32 input is rejected for the same reason as train_encoded.
+    void train_ex(idx_t n, const void* x, NumericType numeric_type) override;
+
     /// implemented for all IndexIVF* classes
     void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
 

data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp

@@ -334,6 +334,7 @@ struct IVFPQFastScanScanner : InvertedListScanner {
             const IDSelector* sel_in)
             : InvertedListScanner(store_pairs_in, sel_in), index(index_in) {
         this->keep_max = is_similarity_metric(index_in.metric_type);
+        this->code_size = index_in.code_size;
         residual.resize(index_in.d);
     }
 

data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp

@@ -162,8 +162,7 @@ struct RaBitInvertedListScanner : InvertedListScanner {
     std::vector<float> query_vector;
 
     std::unique_ptr<FlatCodesDistanceComputer> dc;
-    RaBitQDistanceComputer* rabitq_dc =
-            nullptr; // For multi-bit adaptive filtering
+    RaBitQDistanceComputer* rabitq_dc = nullptr;
 
     uint8_t qb = 0;
     bool centered = false;
@@ -205,30 +204,15 @@ struct RaBitInvertedListScanner : InvertedListScanner {
         return dc->distance_to_code(code);
     }
 
-    // redefiniing the scan_codes allows to inline the distance_to_code
-    // (this is unlikely to matter because it contains a virtual function call)
-    size_t scan_codes_1bit(
-            size_t list_size,
-            const uint8_t* codes,
-            const idx_t* ids,
-            ResultHandler& handler) const {
-        return run_scan_codes(*this, list_size, codes, ids, handler);
-    }
-
-    /// Override scan_codes to implement adaptive filtering for multi-bit codes
     size_t scan_codes(
             size_t list_size,
             const uint8_t* codes,
             const idx_t* ids,
             ResultHandler& handler) const override {
-        size_t ex_bits = ivf_rabitq.rabitq.nb_bits - 1;
-
-        // For 1-bit codes, use default implementation
-        if (ex_bits == 0 || rabitq_dc == nullptr) {
-            return scan_codes_1bit(list_size, codes, ids, handler);
+        if (qb == 0 && ivf_rabitq.rabitq.nb_bits == 1) {
+            return run_scan_codes(*this, list_size, codes, ids, handler);
         }
 
-        // Multi-bit: Two-stage search with adaptive filtering
         return rabitq_dc->scan_codes_multibit(
                 list_size,
                 codes,
@@ -247,7 +231,6 @@ struct RaBitInvertedListScanner : InvertedListScanner {
             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);
             }

data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.cpp

@@ -19,9 +19,11 @@
 #include <faiss/impl/ResultHandler.h>
 #include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
 #include <faiss/impl/fast_scan/fast_scan.h>
+#include <faiss/impl/platform_macros.h>
 #include <faiss/impl/simd_dispatch.h>
 #include <faiss/invlists/BlockInvertedLists.h>
 #include <faiss/utils/distances.h>
+#include <faiss/utils/quantize_lut.h>
 #include <faiss/utils/rabitq_simd.h>
 #include <faiss/utils/utils.h>
 
@@ -401,12 +403,14 @@ void IndexIVFRaBitQFastScan::compute_residual_LUT(
         uint8_t qb_param,
         bool centered_param,
         std::vector<float>& rotated_q,
-        std::vector<float>& centroid_buf) const {
+        std::vector<float>& centroid_buf,
+        std::vector<uint8_t>& rotated_qq,
+        float* rotated_q_out) const {
     const size_t d_val = static_cast<size_t>(d);
     FAISS_THROW_IF_NOT(d_val > 0);
     rotated_q.resize(d_val);
     centroid_buf.resize(d_val);
-    std::vector<uint8_t> rotated_qq(d_val);
+    rotated_qq.resize(d_val);
 
     // Compute residual
     quantizer->reconstruct(centroid_id, centroid_buf.data());
@@ -414,7 +418,9 @@ void IndexIVFRaBitQFastScan::compute_residual_LUT(
         rotated_q[i] = query[i] - centroid_buf[i];
     }
 
-    // Compute query factors using shared utility
+    // Keep the caller's rotated_q capacity across the assignment below, so a
+    // reused QueryFactorsData does not reallocate on every call.
+    std::vector<float> rotated_q_keep = std::move(query_factors.rotated_q);
     query_factors = rabitq_utils::compute_query_factors(
             rotated_q.data(),
             d_val,
@@ -432,7 +438,12 @@ void IndexIVFRaBitQFastScan::compute_residual_LUT(
     }
 
     if (rabitq.nb_bits > 1) {
-        query_factors.rotated_q = rotated_q;
+        if (rotated_q_out != nullptr) {
+            std::copy(rotated_q.begin(), rotated_q.end(), rotated_q_out);
+        } else {
+            rotated_q_keep.assign(rotated_q.begin(), rotated_q.end());
+            query_factors.rotated_q = std::move(rotated_q_keep);
+        }
     }
 
     // Build LUT using branchless subset-sum construction
@@ -528,8 +539,15 @@ void IndexIVFRaBitQFastScan::search_preassigned(
     }
 
     std::vector<QueryFactorsData> query_factors_storage(n * cur_nprobe);
+    // One block for the multi-bit rotated queries, allocated here rather than
+    // per (query, probe) on OpenMP workers: under jemalloc heap profiling on
+    // aarch64, worker-thread allocations are very expensive to sample.
+    // Left uninitialised; every slot read is written by compute_LUT first.
+    std::unique_ptr<float[]> rotated_q_storage(
+            rabitq.nb_bits > 1 ? new float[n * cur_nprobe * d] : nullptr);
     FastScanDistancePostProcessing context;
     context.query_factors = query_factors_storage.data();
+    context.rotated_q = rotated_q_storage.get();
     context.nprobe = cur_nprobe;
     context.qb = used_qb;
     context.centered = used_centered;
@@ -570,6 +588,7 @@ void IndexIVFRaBitQFastScan::compute_LUT(
     {
         std::vector<float> rotated_q(d);
         std::vector<float> centroid_buf(d);
+        std::vector<uint8_t> rotated_qq(d);
 
 #pragma omp for
         for (idx_t ij = 0; ij < static_cast<idx_t>(n * cq_nprobe); ij++) {
@@ -587,7 +606,10 @@ void IndexIVFRaBitQFastScan::compute_LUT(
                         used_qb,
                         used_centered,
                         rotated_q,
-                        centroid_buf);
+                        centroid_buf,
+                        rotated_qq,
+                        context.rotated_q ? context.rotated_q + ij * d
+                                          : nullptr);
 
                 if (context.query_factors != nullptr) {
                     context.query_factors[ij] = std::move(query_factors_data);
@@ -628,10 +650,12 @@ void IndexIVFRaBitQFastScan::compute_LUT_uint8(
         AlignedTable<float> lut_float(cur_nprobe * dim12);
         std::vector<float> rotated_q(d);
         std::vector<float> centroid_buf(d);
+        std::vector<uint8_t> rotated_qq(d);
         std::vector<float> all_mins(cur_nprobe * M);
         std::vector<float> probe_b(cur_nprobe);
+        std::vector<float> probe_span(cur_nprobe);
 
-#pragma omp for schedule(dynamic)
+#pragma omp for schedule(FAISS_OMP_DYNAMIC)
         for (int64_t i = 0; i < static_cast<int64_t>(n); i++) {
             const float* xi = x + i * d;
 
@@ -650,10 +674,13 @@ void IndexIVFRaBitQFastScan::compute_LUT_uint8(
                             used_qb,
                             used_centered,
                             rotated_q,
-                            centroid_buf);
+                            centroid_buf,
+                            rotated_qq,
+                            context.rotated_q ? context.rotated_q + ij * d
+                                              : nullptr);
 
                     if (context.query_factors != nullptr) {
-                        context.query_factors[ij] = qf;
+                        context.query_factors[ij] = std::move(qf);
                     }
                 } else {
                     memset(lut_float.get() + j * dim12,
@@ -686,13 +713,18 @@ void IndexIVFRaBitQFastScan::compute_LUT_uint8(
                                 span_j += span;
                             }
                             probe_b[j2] = b_j;
-                            glob_max_dis = std::max(glob_max_dis, span_j);
+                            probe_span[j2] = span_j;
                             glob_b = std::min(glob_b, b_j);
                         }
 
-                        a = std::min(
-                                255.0f / glob_max_span,
-                                65535.0f / glob_max_dis);
+                        for (size_t j2 = 0; j2 < cur_nprobe; j2++) {
+                            glob_max_dis = std::max(
+                                    glob_max_dis,
+                                    probe_span[j2] + probe_b[j2] - glob_b);
+                        }
+
+                        a = quantize_lut::fastscan_lut_scale(
+                                glob_max_span, glob_max_dis, M + 1);
 
                         // Second pass: quantize LUT and compute biasq.
                         uint8_t* out_base =
@@ -872,6 +904,7 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
     AlignedTable<float> lut_float;
     std::vector<float> rotated_q;
     std::vector<float> centroid_buf;
+    std::vector<uint8_t> rotated_qq;
     QueryFactorsData query_factors;
     FastScanDistancePostProcessing context;
     std::vector<int> probe_map;
@@ -893,6 +926,7 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
               lut_float(16 * index_in.M),
               rotated_q(index_in.d),
               centroid_buf(index_in.d),
+              rotated_qq(index_in.d),
               probe_map({0}),
               mins_buf(index_in.M) {
         this->keep_max = is_similarity_metric(index_in.metric_type);
@@ -924,7 +958,9 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
                 qb,
                 centered,
                 rotated_q,
-                centroid_buf);
+                centroid_buf,
+                rotated_qq,
+                nullptr);
 
         // Single-probe quantization (simplified inline, no OMP, no 3D)
         const size_t M = index.M;
@@ -952,7 +988,7 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
                         b += mn;
                     }
 
-                    a = std::min(255.0f / max_span, 65535.0f / max_dis);
+                    a = quantize_lut::fastscan_lut_scale(max_span, max_dis, M);
                     for (size_t m = 0; m < M; m++) {
                         const float* tab = lut_float.get() + m * ksub;
                         rabitq::lut_quantize_16_to_uint8<SL>(

data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.h

@@ -131,7 +131,9 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
             uint8_t qb_param,
             bool centered_param,
             std::vector<float>& rotated_q,
-            std::vector<float>& centroid_buf) const;
+            std::vector<float>& centroid_buf,
+            std::vector<uint8_t>& rotated_qq,
+            float* rotated_q_out) const;
 
    private:
     /// Decode FastScan code to RaBitQ residual vector with explicit
@@ -399,7 +401,8 @@ float IVFRaBitQHeapHandler<C, SL>::compute_full_multibit_distance(
             unpack_buf.data(),
             ex_code,
             ex_fac,
-            query_factors.rotated_q.data(),
+            context->rotated_q ? context->rotated_q + storage_idx_val * dim
+                               : query_factors.rotated_q.data(),
             is_similarity ? query_factors.q_dot_c : query_factors.qr_to_c_L2sqr,
             dim,
             ex_bits,