faiss 0.5.2 → 0.6.0

data/vendor/faiss/faiss/impl/Panorama.cpp

@@ -10,10 +10,46 @@
 #include <algorithm>
 #include <cmath>
 #include <cstring>
-#include <vector>
+
+#include <faiss/impl/FaissAssert.h>
 
 namespace faiss {
 
+namespace {
+
+/// Helper function to compute cumulative sums by iterating backwards through
+/// levels. This is the core logic shared by compute_cumulative_sums and
+/// compute_query_cum_sums.
+template <typename OffsetFunc>
+inline void compute_cum_sums_impl(
+        const float* vector,
+        float* output,
+        size_t d,
+        size_t n_levels,
+        size_t level_width_floats,
+        OffsetFunc&& get_offset) {
+    // Iterate backwards through levels, accumulating sum as we go.
+    // This avoids computing the suffix sum for each vector, which takes
+    // extra memory.
+    float sum = 0.0f;
+
+    for (int level = n_levels - 1; level >= 0; level--) {
+        size_t start_idx = level * level_width_floats;
+        size_t end_idx = std::min(
+                (level + 1) * level_width_floats, static_cast<size_t>(d));
+
+        for (size_t j = start_idx; j < end_idx; j++) {
+            sum += vector[j] * vector[j];
+        }
+
+        output[get_offset(level)] = std::sqrt(sum);
+    }
+
+    output[get_offset(n_levels)] = 0.0f;
+}
+
+} // namespace
+
 /**************************************************************
  * Panorama structure implementation
  **************************************************************/
@@ -24,6 +60,7 @@ Panorama::Panorama(size_t code_size, size_t n_levels, size_t batch_size)
 }
 
 void Panorama::set_derived_values() {
+    FAISS_THROW_IF_NOT_MSG(n_levels > 0, "Panorama: n_levels must be > 0");
     this->d = code_size / sizeof(float);
     this->level_width_floats = ((d + n_levels - 1) / n_levels);
     this->level_width = this->level_width_floats * sizeof(float);
@@ -69,64 +106,34 @@ void Panorama::compute_cumulative_sums(
         float* cumsum_base,
         size_t offset,
         size_t n_entry,
-        const float* vectors) {
-    std::vector<float> suffix_sums(d + 1);
-
+        const float* vectors) const {
     for (size_t entry_idx = 0; entry_idx < n_entry; entry_idx++) {
         size_t current_pos = offset + entry_idx;
         size_t batch_no = current_pos / batch_size;
         size_t pos_in_batch = current_pos % batch_size;
 
         const float* vector = vectors + entry_idx * d;
-
-        // Compute suffix sums of squared values.
-        suffix_sums[d] = 0.0f;
-        for (int j = d - 1; j >= 0; j--) {
-            float squared_val = vector[j] * vector[j];
-            suffix_sums[j] = suffix_sums[j + 1] + squared_val;
-        }
-
-        // Store cumulative sums in batch-oriented layout.
         size_t cumsum_batch_offset = batch_no * batch_size * (n_levels + 1);
 
-        for (size_t level = 0; level < n_levels; level++) {
-            size_t start_idx = level * level_width_floats;
-            size_t cumsum_offset =
-                    cumsum_batch_offset + level * batch_size + pos_in_batch;
-            if (start_idx < d) {
-                cumsum_base[cumsum_offset] = std::sqrt(suffix_sums[start_idx]);
-            } else {
-                cumsum_base[cumsum_offset] = 0.0f;
-            }
-        }
-
-        // Last level sum is always 0.
-        size_t cumsum_offset =
-                cumsum_batch_offset + n_levels * batch_size + pos_in_batch;
-        cumsum_base[cumsum_offset] = 0.0f;
+        auto get_offset = [&](size_t level) {
+            return cumsum_batch_offset + level * batch_size + pos_in_batch;
+        };
+
+        compute_cum_sums_impl(
+                vector,
+                cumsum_base,
+                d,
+                n_levels,
+                level_width_floats,
+                get_offset);
     }
 }
 
 void Panorama::compute_query_cum_sums(const float* query, float* query_cum_sums)
         const {
-    std::vector<float> suffix_sums(d + 1);
-    suffix_sums[d] = 0.0f;
-
-    for (int j = d - 1; j >= 0; j--) {
-        float squared_val = query[j] * query[j];
-        suffix_sums[j] = suffix_sums[j + 1] + squared_val;
-    }
-
-    for (size_t level = 0; level < n_levels; level++) {
-        size_t start_idx = level * level_width_floats;
-        if (start_idx < d) {
-            query_cum_sums[level] = std::sqrt(suffix_sums[start_idx]);
-        } else {
-            query_cum_sums[level] = 0.0f;
-        }
-    }
-
-    query_cum_sums[n_levels] = 0.0f;
+    auto get_offset = [](size_t level) { return level; };
+    compute_cum_sums_impl(
+            query, query_cum_sums, d, n_levels, level_width_floats, get_offset);
 }
 
 void Panorama::reconstruct(idx_t key, float* recons, const uint8_t* codes_base)

data/vendor/faiss/faiss/impl/Panorama.h

@@ -10,6 +10,7 @@
 #ifndef FAISS_PANORAMA_H
 #define FAISS_PANORAMA_H
 
+#include <faiss/MetricType.h>
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/PanoramaStats.h>
 #include <faiss/utils/distances.h>
@@ -67,7 +68,7 @@ struct Panorama {
             float* cumsum_base,
             size_t offset,
             size_t n_entry,
-            const float* vectors);
+            const float* vectors) const;
 
     /// Compute the cumulative sums of the query vector.
     void compute_query_cum_sums(const float* query, float* query_cum_sums)
@@ -97,7 +98,7 @@ struct Panorama {
     /// 4. After all levels, survivors are exact distances; update heap.
     /// This achieves early termination while maintaining SIMD-friendly
     /// sequential access patterns in the level-oriented storage layout.
-    template <typename C>
+    template <typename C, MetricType M>
     size_t progressive_filter_batch(
             const uint8_t* codes_base,
             const float* cum_sums,
@@ -116,7 +117,7 @@ struct Panorama {
     void reconstruct(idx_t key, float* recons, const uint8_t* codes_base) const;
 };
 
-template <typename C>
+template <typename C, MetricType M>
 size_t Panorama::progressive_filter_batch(
         const uint8_t* codes_base,
         const float* cum_sums,
@@ -151,7 +152,12 @@ size_t Panorama::progressive_filter_batch(
 
         active_indices[num_active] = i;
         float cum_sum = batch_cum_sums[i];
-        exact_distances[i] = cum_sum * cum_sum + q_norm;
+
+        if constexpr (M == METRIC_INNER_PRODUCT) {
+            exact_distances[i] = 0.0f;
+        } else {
+            exact_distances[i] = cum_sum * cum_sum + q_norm;
+        }
 
         num_active += include;
     }
@@ -183,10 +189,20 @@ size_t Panorama::progressive_filter_batch(
             float dot_product =
                     fvec_inner_product(query_level, yj, actual_level_width);
 
-            exact_distances[idx] -= 2.0f * dot_product;
+            if constexpr (M == METRIC_INNER_PRODUCT) {
+                exact_distances[idx] += dot_product;
+            } else {
+                exact_distances[idx] -= 2.0f * dot_product;
+            }
 
             float cum_sum = level_cum_sums[idx];
-            float cauchy_schwarz_bound = 2.0f * cum_sum * query_cum_norm;
+            float cauchy_schwarz_bound;
+            if constexpr (M == METRIC_INNER_PRODUCT) {
+                cauchy_schwarz_bound = -cum_sum * query_cum_norm;
+            } else {
+                cauchy_schwarz_bound = 2.0f * cum_sum * query_cum_norm;
+            }
+
             float lower_bound = exact_distances[idx] - cauchy_schwarz_bound;
 
             active_indices[next_active] = idx;

data/vendor/faiss/faiss/impl/PolysemousTraining.cpp

@@ -18,6 +18,7 @@
 #include <cstring>
 #include <memory>
 
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/random.h>
@@ -431,6 +432,8 @@ void ReproduceDistancesObjective::set_affine_target_dis(
  * Cost functions: RankingScore
  ****************************************************/
 
+namespace {
+
 /// Maintains a 3D table of elementary costs.
 /// Accumulates elements based on Hamming distance comparisons
 template <typename Ttab, typename Taccu>
@@ -756,6 +759,8 @@ struct RankingScore2 : Score3Computer<float, double> {
     }
 };
 
+} // namespace
+
 /*****************************************
  * PolysemousTraining
  ******************************************/
@@ -798,12 +803,18 @@ void PolysemousTraining::optimize_reproduce_distances(
 
         float* centroids = pq.get_centroids(m, 0);
 
-        for (int i = 0; i < n; i++) {
-            for (int j = 0; j < n; j++) {
-                dis_table.push_back(fvec_L2sqr(
-                        centroids + i * dsub, centroids + j * dsub, dsub));
+        auto compute_dis_table = [&]<SIMDLevel SL>() {
+            for (int i = 0; i < n; i++) {
+                for (int j = 0; j < n; j++) {
+                    dis_table.push_back(
+                            fvec_L2sqr<SL>(
+                                    centroids + i * dsub,
+                                    centroids + j * dsub,
+                                    dsub));
+                }
             }
-        }
+        };
+        with_simd_level(compute_dis_table);
 
         std::vector<int> perm(n);
         ReproduceWithHammingObjective obj(nbits, dis_table, dis_weight_factor);

data/vendor/faiss/faiss/impl/ProductQuantizer.cpp

@@ -19,6 +19,7 @@
 #include <faiss/IndexFlat.h>
 #include <faiss/VectorTransform.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
 
 extern "C" {
@@ -56,14 +57,15 @@ ProductQuantizer::ProductQuantizer() : ProductQuantizer(0, 1, 0) {}
 
 void ProductQuantizer::set_derived_values() {
     // quite a few derived values
+    FAISS_THROW_IF_NOT_MSG(M > 0, "M must be > 0");
     FAISS_THROW_IF_NOT_MSG(
             d % M == 0,
             "The dimension of the vector (d) should be a multiple of the number of subquantizers (M)");
     dsub = d / M;
-    code_size = (nbits * M + 7) / 8;
     FAISS_THROW_IF_MSG(nbits > 24, "nbits larger than 24 is not practical.");
+    code_size = (nbits * M + 7) / 8;
     ksub = 1 << nbits;
-    centroids.resize(d * ksub);
+    centroids.resize(mul_no_overflow(d, (size_t)ksub, "PQ centroids"));
     verbose = false;
     train_type = Train_default;
 }
@@ -201,8 +203,10 @@ void ProductQuantizer::train(size_t n, const float* x) {
     }
 }
 
-template <class PQEncoder>
-void compute_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
+namespace {
+
+template <class PQEncoder, SIMDLevel SL>
+void compute_1_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
     std::vector<float> distances(pq.ksub);
 
     // It seems to be meaningless to allocate std::vector<float> distances.
@@ -248,7 +252,7 @@ void compute_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
         uint64_t idxm = 0;
         if (pq.transposed_centroids.empty()) {
             // the regular version
-            idxm = fvec_L2sqr_ny_nearest(
+            idxm = fvec_L2sqr_ny_nearest<SL>(
                     distances.data(),
                     xsub,
                     pq.get_centroids(m, 0),
@@ -256,7 +260,7 @@ void compute_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
                     pq.ksub);
         } else {
             // transposed centroids are available, use'em
-            idxm = fvec_L2sqr_ny_nearest_y_transposed(
+            idxm = fvec_L2sqr_ny_nearest_y_transposed<SL>(
                     distances.data(),
                     xsub,
                     pq.transposed_centroids.data() + m * pq.ksub,
@@ -270,20 +274,24 @@ void compute_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
     }
 }
 
+} // namespace
+
 void ProductQuantizer::compute_code(const float* x, uint8_t* code) const {
-    switch (nbits) {
-        case 8:
-            faiss::compute_code<PQEncoder8>(*this, x, code);
-            break;
+    with_simd_level([&]<SIMDLevel SL>() {
+        switch (nbits) {
+            case 8:
+                compute_1_code<PQEncoder8, SL>(*this, x, code);
+                break;
 
-        case 16:
-            faiss::compute_code<PQEncoder16>(*this, x, code);
-            break;
+            case 16:
+                compute_1_code<PQEncoder16, SL>(*this, x, code);
+                break;
 
-        default:
-            faiss::compute_code<PQEncoderGeneric>(*this, x, code);
-            break;
-    }
+            default:
+                compute_1_code<PQEncoderGeneric, SL>(*this, x, code);
+                break;
+        }
+    }); // with_simd_level
 }
 
 template <class PQDecoder>
@@ -428,44 +436,46 @@ void ProductQuantizer::compute_codes(const float* x, uint8_t* codes, size_t n)
 
 void ProductQuantizer::compute_distance_table(const float* x, float* dis_table)
         const {
-    if (transposed_centroids.empty()) {
-        // use regular version
-        for (size_t m = 0; m < M; m++) {
-            fvec_L2sqr_ny(
-                    dis_table + m * ksub,
-                    x + m * dsub,
-                    get_centroids(m, 0),
-                    dsub,
-                    ksub);
+    with_simd_level([&]<SIMDLevel SL>() {
+        if (transposed_centroids.empty()) {
+            // use regular version
+            for (size_t m = 0; m < M; m++) {
+                fvec_L2sqr_ny<SL>(
+                        dis_table + m * ksub,
+                        x + m * dsub,
+                        get_centroids(m, 0),
+                        dsub,
+                        ksub);
+            }
+        } else {
+            // transposed centroids are available, use'em
+            for (size_t m = 0; m < M; m++) {
+                fvec_L2sqr_ny_transposed<SL>(
+                        dis_table + m * ksub,
+                        x + m * dsub,
+                        transposed_centroids.data() + m * ksub,
+                        centroids_sq_lengths.data() + m * ksub,
+                        dsub,
+                        M * ksub,
+                        ksub);
+            }
         }
-    } else {
-        // transposed centroids are available, use'em
+    });
+}
+
+void ProductQuantizer::compute_inner_prod_table(
+        const float* x,
+        float* dis_table) const {
+    with_simd_level([&]<SIMDLevel SL>() {
         for (size_t m = 0; m < M; m++) {
-            fvec_L2sqr_ny_transposed(
+            fvec_inner_products_ny<SL>(
                     dis_table + m * ksub,
                     x + m * dsub,
-                    transposed_centroids.data() + m * ksub,
-                    centroids_sq_lengths.data() + m * ksub,
+                    get_centroids(m, 0),
                     dsub,
-                    M * ksub,
                     ksub);
         }
-    }
-}
-
-void ProductQuantizer::compute_inner_prod_table(
-        const float* x,
-        float* dis_table) const {
-    size_t m;
-
-    for (m = 0; m < M; m++) {
-        fvec_inner_products_ny(
-                dis_table + m * ksub,
-                x + m * dsub,
-                get_centroids(m, 0),
-                dsub,
-                ksub);
-    }
+    });
 }
 
 void ProductQuantizer::compute_distance_tables(
@@ -785,17 +795,19 @@ void ProductQuantizer::compute_sdc_table() {
     sdc_table.resize(M * ksub * ksub);
 
     if (dsub < 4) {
+        with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for
-        for (int mk = 0; mk < M * ksub; mk++) {
-            // allow omp to schedule in a more fine-grained way
-            // `collapse` is not supported in OpenMP 2.x
-            int m = mk / ksub;
-            int k = mk % ksub;
-            const float* cents = centroids.data() + m * ksub * dsub;
-            const float* centi = cents + k * dsub;
-            float* dis_tab = sdc_table.data() + m * ksub * ksub;
-            fvec_L2sqr_ny(dis_tab + k * ksub, centi, cents, dsub, ksub);
-        }
+            for (int mk = 0; mk < M * ksub; mk++) {
+                // allow omp to schedule in a more fine-grained way
+                // `collapse` is not supported in OpenMP 2.x
+                int m = mk / ksub;
+                int k = mk % ksub;
+                const float* cents = centroids.data() + m * ksub * dsub;
+                const float* centi = cents + k * dsub;
+                float* dis_tab = sdc_table.data() + m * ksub * ksub;
+                fvec_L2sqr_ny<SL>(dis_tab + k * ksub, centi, cents, dsub, ksub);
+            }
+        });
     } else {
         // NOTE: it would disable the omp loop in pairwise_L2sqr
         // but still accelerate especially when M >= 4

data/vendor/faiss/faiss/impl/RaBitQUtils.cpp

@@ -9,8 +9,10 @@
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/distances.h>
+#include <faiss/utils/rabitq_simd.h>
 #include <algorithm>
 #include <cmath>
+#include <cstring>
 #include <limits>
 
 namespace faiss {
@@ -242,8 +244,12 @@ QueryFactorsData compute_query_factors(
 
     // Compute query norm for inner product metric
     query_factors.qr_norm_L2sqr = 0.0f;
+    query_factors.q_dot_c = 0.0f;
     if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
         query_factors.qr_norm_L2sqr = fvec_norm_L2sqr(query, d);
+        if (centroid != nullptr) {
+            query_factors.q_dot_c = fvec_inner_product(query, centroid, d);
+        }
     }
 
     return query_factors;
@@ -290,5 +296,91 @@ void set_bit_fastscan(uint8_t* code, size_t bit_index) {
     }
 }
 
+size_t compute_per_vector_storage_size(size_t nb_bits, size_t d) {
+    const size_t ex_bits = nb_bits - 1;
+    if (ex_bits == 0) {
+        return sizeof(SignBitFactors);
+    } else {
+        return sizeof(SignBitFactorsWithError) + sizeof(ExtraBitsFactors) +
+                (d * ex_bits + 7) / 8;
+    }
+}
+
+float compute_full_multibit_distance(
+        const uint8_t* sign_bits,
+        const uint8_t* ex_code,
+        const ExtraBitsFactors& ex_fac,
+        const float* rotated_q,
+        float qr_base,
+        size_t d,
+        size_t ex_bits,
+        MetricType metric_type) {
+    const float cb = -(static_cast<float>(1 << ex_bits) - 0.5f);
+
+    float ex_ip = rabitq::multibit::compute_inner_product(
+            sign_bits, ex_code, rotated_q, d, ex_bits, cb);
+
+    float dist = qr_base + ex_fac.f_add_ex + ex_fac.f_rescale_ex * ex_ip;
+
+    if (metric_type == MetricType::METRIC_L2) {
+        dist = std::max(0.0f, dist);
+    }
+
+    return dist;
+}
+
+void populate_block_aux_from_flat_storage(
+        const std::vector<uint8_t>& flat_storage,
+        AlignedTable<uint8_t>& codes,
+        size_t num_vectors,
+        size_t bbs,
+        size_t M2,
+        size_t old_block_stride,
+        size_t new_block_stride,
+        size_t storage_size,
+        const int64_t* id_map) {
+    if (flat_storage.empty() || num_vectors == 0) {
+        return;
+    }
+
+    const size_t packed_block_size = ((M2 + 1) / 2) * bbs;
+    const size_t n_blocks = (num_vectors + bbs - 1) / bbs;
+
+    if (old_block_stride < new_block_stride) {
+        AlignedTable<uint8_t> old_data;
+        old_data.resize(codes.size());
+        memcpy(old_data.data(), codes.data(), codes.size());
+
+        codes.resize(n_blocks * new_block_stride);
+        memset(codes.data(), 0, n_blocks * new_block_stride);
+        for (size_t b = 0; b < n_blocks; b++) {
+            memcpy(codes.data() + b * new_block_stride,
+                   old_data.data() + b * old_block_stride,
+                   packed_block_size);
+        }
+    }
+
+    for (size_t offset = 0; offset < num_vectors; offset++) {
+        const int64_t global_id =
+                id_map ? id_map[offset] : static_cast<int64_t>(offset);
+        FAISS_THROW_IF_NOT_MSG(
+                global_id >= 0 &&
+                        static_cast<size_t>(global_id) * storage_size +
+                                        storage_size <=
+                                flat_storage.size(),
+                "global_id out of bounds for flat_storage during migration");
+
+        const uint8_t* src = flat_storage.data() + global_id * storage_size;
+        uint8_t* dst = get_block_aux_ptr(
+                codes.data(),
+                offset,
+                bbs,
+                packed_block_size,
+                new_block_stride,
+                storage_size);
+        memcpy(dst, src, storage_size);
+    }
+}
+
 } // namespace rabitq_utils
 } // namespace faiss