faiss 0.6.0 → 0.6.2

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

@@ -13,12 +13,15 @@
 #include <cstdio>
 #include <cstring>
 #include <memory>
+#include <type_traits>
 
 #include <algorithm>
 
 #include <faiss/IndexFlat.h>
 #include <faiss/VectorTransform.h>
 #include <faiss/impl/FaissAssert.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
 #include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
 
@@ -48,8 +51,8 @@ namespace faiss {
  * PQ implementation
  *********************************************/
 
-ProductQuantizer::ProductQuantizer(size_t d, size_t M, size_t nbits)
-        : Quantizer(d, 0), M(M), nbits(nbits), assign_index(nullptr) {
+ProductQuantizer::ProductQuantizer(size_t d_in, size_t M_in, size_t nbits_in)
+        : Quantizer(d_in, 0), M(M_in), nbits(nbits_in), assign_index(nullptr) {
     set_derived_values();
 }
 
@@ -139,8 +142,8 @@ void ProductQuantizer::train(size_t n, const float* x) {
         }
 
         std::unique_ptr<float[]> xslice(new float[n * dsub]);
-        for (int m = 0; m < M; m++) {
-            for (int j = 0; j < n; j++)
+        for (size_t m = 0; m < M; m++) {
+            for (size_t j = 0; j < n; j++)
                 memcpy(xslice.get() + j * dsub,
                        x + j * d + m * dsub,
                        dsub * sizeof(float));
@@ -179,7 +182,7 @@ void ProductQuantizer::train(size_t n, const float* x) {
 
             if (verbose) {
                 clus.verbose = true;
-                printf("Training PQ slice %d/%zd\n", m, M);
+                printf("Training PQ slice %zd/%zd\n", m, M);
             }
             IndexFlatL2 index(dsub);
             clus.train(n, xslice.get(), assign_index ? *assign_index : index);
@@ -197,7 +200,7 @@ void ProductQuantizer::train(size_t n, const float* x) {
         IndexFlatL2 index(dsub);
 
         clus.train(n * M, x, assign_index ? *assign_index : index);
-        for (int m = 0; m < M; m++) {
+        for (size_t m = 0; m < M; m++) {
             set_params(clus.centroids.data(), m);
         }
     }
@@ -322,8 +325,9 @@ void ProductQuantizer::decode(const uint8_t* code, float* x) const {
 }
 
 void ProductQuantizer::decode(const uint8_t* code, float* x, size_t n) const {
+    int64_t n_signed = n;
 #pragma omp parallel for if (n > 100)
-    for (int64_t i = 0; i < n; i++) {
+    for (int64_t i = 0; i < n_signed; i++) {
         this->decode(code + code_size * i, x + d * i);
     }
 }
@@ -353,7 +357,8 @@ void ProductQuantizer::compute_codes_with_assign_index(
         const float* x,
         uint8_t* codes,
         size_t n) {
-    FAISS_THROW_IF_NOT(assign_index && assign_index->d == dsub);
+    FAISS_THROW_IF_NOT(
+            assign_index && static_cast<size_t>(assign_index->d) == dsub);
 
     for (size_t m = 0; m < M; m++) {
         assign_index->reset();
@@ -415,10 +420,11 @@ void ProductQuantizer::compute_codes(const float* x, uint8_t* codes, size_t n)
         return;
     }
 
+    int64_t n_signed = n;
     if (dsub < 16) { // simple direct computation
 
 #pragma omp parallel for
-        for (int64_t i = 0; i < n; i++)
+        for (int64_t i = 0; i < n_signed; i++)
             compute_code(x + i * d, codes + i * code_size);
 
     } else { // worthwhile to use BLAS
@@ -426,7 +432,7 @@ void ProductQuantizer::compute_codes(const float* x, uint8_t* codes, size_t n)
         compute_distance_tables(n, x, dis_tables.get());
 
 #pragma omp parallel for
-        for (int64_t i = 0; i < n; i++) {
+        for (int64_t i = 0; i < n_signed; i++) {
             uint8_t* code = codes + i * code_size;
             const float* tab = dis_tables.get() + i * ksub * M;
             compute_code_from_distance_table(tab, code);
@@ -482,7 +488,8 @@ void ProductQuantizer::compute_distance_tables(
         size_t nx,
         const float* x,
         float* dis_tables) const {
-#if defined(__AVX2__) || defined(__aarch64__)
+    int64_t nx_signed = nx;
+#if defined(COMPILE_SIMD_AVX2) || defined(COMPILE_SIMD_ARM_NEON)
     if (dsub == 2 && nbits < 8) { // interesting for a narrow range of settings
         compute_PQ_dis_tables_dsub2(
                 d, ksub, centroids.data(), nx, x, false, dis_tables);
@@ -491,13 +498,13 @@ void ProductQuantizer::compute_distance_tables(
             if (dsub < 16) {
 
 #pragma omp parallel for if (nx > 1)
-        for (int64_t i = 0; i < nx; i++) {
+        for (int64_t i = 0; i < nx_signed; i++) {
             compute_distance_table(x + i * d, dis_tables + i * ksub * M);
         }
 
     } else { // use BLAS
 
-        for (int m = 0; m < M; m++) {
+        for (size_t m = 0; m < M; m++) {
             pairwise_L2sqr(
                     dsub,
                     nx,
@@ -516,7 +523,8 @@ void ProductQuantizer::compute_inner_prod_tables(
         size_t nx,
         const float* x,
         float* dis_tables) const {
-#if defined(__AVX2__) || defined(__aarch64__)
+    int64_t nx_signed = nx;
+#if defined(COMPILE_SIMD_AVX2) || defined(COMPILE_SIMD_ARM_NEON)
     if (dsub == 2 && nbits < 8) {
         compute_PQ_dis_tables_dsub2(
                 d, ksub, centroids.data(), nx, x, true, dis_tables);
@@ -525,14 +533,14 @@ void ProductQuantizer::compute_inner_prod_tables(
             if (dsub < 16) {
 
 #pragma omp parallel for if (nx > 1)
-        for (int64_t i = 0; i < nx; i++) {
+        for (int64_t i = 0; i < nx_signed; i++) {
             compute_inner_prod_table(x + i * d, dis_tables + i * ksub * M);
         }
 
     } else { // use BLAS
 
         // compute distance tables
-        for (int m = 0; m < M; m++) {
+        for (size_t m = 0; m < M; m++) {
             FINTEGER ldc = ksub * M, nxi = nx, ksubi = ksub, dsubi = dsub,
                      di = d;
             float one = 1.0, zero = 0;
@@ -576,7 +584,7 @@ void pq_estimators_from_tables_Mmul4(
         float dis = 0;
         const float* dt = dis_table;
 
-        for (size_t m = 0; m < M; m += 4) {
+        for (int m = 0; m < M; m += 4) {
             float dism = 0;
             dism = dt[*codes++];
             dt += ksub;
@@ -648,7 +656,7 @@ void pq_estimators_from_tables(
     for (size_t j = 0; j < ncodes; j++) {
         float dis = 0;
         const float* __restrict dt = dis_table;
-        for (int m = 0; m < M; m++) {
+        for (size_t m = 0; m < M; m++) {
             dis += dt[*codes++];
             dt += ksub;
         }
@@ -696,10 +704,11 @@ void pq_knn_search_with_tables(
         HeapArray<C>* res,
         bool init_finalize_heap) {
     size_t k = res->k, nx = res->nh;
+    int64_t nx_signed = nx;
     size_t ksub = pq.ksub, M = pq.M;
 
 #pragma omp parallel for if (nx > 1)
-    for (int64_t i = 0; i < nx; i++) {
+    for (int64_t i = 0; i < nx_signed; i++) {
         /* query preparation for asymmetric search: compute look-up tables */
         const float* dis_table = dis_tables + i * ksub * M;
 
@@ -713,8 +722,28 @@ void pq_knn_search_with_tables(
 
         switch (nbits) {
             case 8:
-                pq_estimators_from_tables<uint8_t, C>(
-                        pq, codes, ncodes, dis_table, k, heap_dis, heap_ids);
+                if (ksub == 256) {
+                    constexpr bool max_heap =
+                            std::is_same_v<C, CMax<float, int64_t>>;
+                    pq_code_distance::pq_scan_8bit(
+                            M,
+                            dis_table,
+                            codes,
+                            ncodes,
+                            k,
+                            heap_dis,
+                            heap_ids,
+                            max_heap);
+                } else {
+                    pq_estimators_from_tables<uint8_t, C>(
+                            pq,
+                            codes,
+                            ncodes,
+                            dis_table,
+                            k,
+                            heap_dis,
+                            heap_ids);
+                }
                 break;
 
             case 16:
@@ -797,7 +826,7 @@ void ProductQuantizer::compute_sdc_table() {
     if (dsub < 4) {
         with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for
-            for (int mk = 0; mk < M * ksub; mk++) {
+            for (int64_t mk = 0; mk < static_cast<int64_t>(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;
@@ -812,7 +841,7 @@ void ProductQuantizer::compute_sdc_table() {
         // NOTE: it would disable the omp loop in pairwise_L2sqr
         // but still accelerate especially when M >= 4
 #pragma omp parallel for
-        for (int m = 0; m < M; m++) {
+        for (int64_t m = 0; m < static_cast<int64_t>(M); m++) {
             const float* cents = centroids.data() + m * ksub * dsub;
             float* dis_tab = sdc_table.data() + m * ksub * ksub;
             pairwise_L2sqr(
@@ -831,9 +860,10 @@ void ProductQuantizer::search_sdc(
     FAISS_THROW_IF_NOT(sdc_table.size() == M * ksub * ksub);
     FAISS_THROW_IF_NOT(nbits == 8);
     size_t k = res->k;
+    int64_t nq_signed = nq;
 
 #pragma omp parallel for
-    for (int64_t i = 0; i < nq; i++) {
+    for (int64_t i = 0; i < nq_signed; i++) {
         /* Compute distances and keep smallest values */
         idx_t* heap_ids = res->ids + i * k;
         float* heap_dis = res->val + i * k;
@@ -846,7 +876,7 @@ void ProductQuantizer::search_sdc(
         for (size_t j = 0; j < nb; j++) {
             float dis = 0;
             const float* tab = sdc_table.data();
-            for (int m = 0; m < M; m++) {
+            for (size_t m = 0; m < M; m++) {
                 dis += tab[bcode[m] + qcode[m] * ksub];
                 tab += ksub * ksub;
             }

data/vendor/faiss/faiss/impl/Quantizer.h

@@ -16,8 +16,8 @@ struct Quantizer {
     size_t d;         ///< size of the input vectors
     size_t code_size; ///< bytes per indexed vector
 
-    explicit Quantizer(size_t d = 0, size_t code_size = 0)
-            : d(d), code_size(code_size) {}
+    explicit Quantizer(size_t d_in = 0, size_t code_size_in = 0)
+            : d(d_in), code_size(code_size_in) {}
 
     /** Train the quantizer
      *

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

@@ -8,6 +8,7 @@
 #include <faiss/impl/RaBitQUtils.h>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/rabitq_simd.h>
 #include <algorithm>
@@ -153,6 +154,7 @@ QueryFactorsData compute_query_factors(
         std::vector<uint8_t>& rotated_qq) {
     FAISS_THROW_IF_NOT(qb <= 8);
     FAISS_THROW_IF_NOT(qb > 0);
+    FAISS_THROW_IF_NOT(d > 0);
 
     QueryFactorsData query_factors;
 
@@ -165,38 +167,42 @@ QueryFactorsData compute_query_factors(
     query_factors.g_error = std::sqrt(query_factors.qr_to_c_L2sqr);
 
     // Rotate the query (subtract centroid)
+    // Save aliasing state before resize(), which may reallocate the buffer.
+    const bool query_aliased = (query == rotated_q.data());
+    FAISS_THROW_IF_NOT_MSG(
+            !query_aliased || centroid == nullptr,
+            "query aliasing is only supported in the IVF residual path "
+            "(centroid == nullptr)");
     rotated_q.resize(d);
-    for (size_t i = 0; i < d; i++) {
-        if (i < rotated_q.size()) {
-            rotated_q[i] =
-                    query[i] - ((centroid == nullptr) ? 0.0f : centroid[i]);
+    if (centroid == nullptr) {
+        // Caller may pass query == rotated_q.data() (IVF residual path);
+        // memcpy with overlapping src/dst is UB, so skip the copy in that case.
+        if (!query_aliased) {
+            memcpy(rotated_q.data(), query, d * sizeof(float));
+        }
+    } else {
+        for (size_t i = 0; i < d; i++) {
+            rotated_q[i] = query[i] - centroid[i];
         }
     }
 
-    const float inv_d_sqrt =
-            (d == 0) ? 1.0f : (1.0f / std::sqrt(static_cast<float>(d)));
+    const float inv_d_sqrt = 1.0f / std::sqrt(static_cast<float>(d));
 
     // Compute quantization range
     float v_min = std::numeric_limits<float>::max();
     float v_max = std::numeric_limits<float>::lowest();
 
+    const float* rq = rotated_q.data();
     if (centered) {
         float z_max = Z_MAX_BY_QB[qb - 1];
         float v_radius = z_max * std::sqrt(query_factors.qr_to_c_L2sqr / d);
         v_min = -v_radius;
         v_max = v_radius;
     } else {
-        // Only compute min/max if we have dimensions to process
-        if (d > 0 && !rotated_q.empty()) {
-            for (size_t i = 0; i < d; i++) {
-                const float v_q = rotated_q[i];
-                v_min = std::min(v_min, v_q);
-                v_max = std::max(v_max, v_q);
-            }
-        } else {
-            // For empty dimensions, use default range
-            v_min = 0.0f;
-            v_max = 1.0f;
+        for (size_t i = 0; i < d; i++) {
+            const float v_q = rq[i];
+            v_min = std::min(v_min, v_q);
+            v_max = std::max(v_max, v_q);
         }
     }
 
@@ -209,25 +215,18 @@ QueryFactorsData compute_query_factors(
     size_t sum_qq = 0;
     int64_t sum2_signed_odd_int = 0;
 
-    // Process arrays - throw error if they are unexpectedly empty
-    if (d > 0 && !rotated_q.empty() && !rotated_qq.empty()) {
-        for (size_t i = 0; i < d; i++) {
-            const float v_q = rotated_q[i];
-            // Non-randomized scalar quantization
-            const uint8_t v_qq = std::clamp<float>(
-                    std::round((v_q - v_min) * inv_delta), 0, max_code);
-            rotated_qq[i] = v_qq;
-            sum_qq += v_qq;
-
-            if (centered) {
-                int64_t signed_odd_int = int64_t(v_qq) * 2 - max_code;
-                sum2_signed_odd_int += signed_odd_int * signed_odd_int;
-            }
+    uint8_t* rqq = rotated_qq.data();
+    for (size_t i = 0; i < d; i++) {
+        const float v_q = rq[i];
+        const uint8_t v_qq = std::clamp<float>(
+                std::round((v_q - v_min) * inv_delta), 0, max_code);
+        rqq[i] = v_qq;
+        sum_qq += v_qq;
+
+        if (centered) {
+            int64_t signed_odd_int = int64_t(v_qq) * 2 - max_code;
+            sum2_signed_odd_int += signed_odd_int * signed_odd_int;
         }
-    } else {
-        FAISS_THROW_MSG(
-                "Arrays unexpectedly empty when d=" + std::to_string(d) +
-                "or d is incorrectly set");
     }
 
     // Compute query factors
@@ -242,11 +241,15 @@ QueryFactorsData compute_query_factors(
         query_factors.int_dot_scale = 1.0f;
     }
 
-    // Compute query norm for inner product metric
+    // Compute query norm for inner product metric.
+    // When centroid is nullptr (IVF residual path), qr_to_c_L2sqr already
+    // holds fvec_norm_L2sqr(query, d) from line 164, so reuse it.
     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);
+        query_factors.qr_norm_L2sqr = (centroid == nullptr)
+                ? query_factors.qr_to_c_L2sqr
+                : fvec_norm_L2sqr(query, d);
         if (centroid != nullptr) {
             query_factors.q_dot_c = fvec_inner_product(query, centroid, d);
         }
@@ -306,6 +309,9 @@ size_t compute_per_vector_storage_size(size_t nb_bits, size_t d) {
     }
 }
 
+// Non-template wrapper with dynamic dispatch (one dispatch per call).
+// The hot path in RaBitQuantizer dispatches once at distance computer
+// construction, so per-vector dispatch only affects this utility path.
 float compute_full_multibit_distance(
         const uint8_t* sign_bits,
         const uint8_t* ex_code,
@@ -315,18 +321,18 @@ float compute_full_multibit_distance(
         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;
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0_SPR>(
+            [&]<SIMDLevel SL>() {
+                return compute_full_multibit_distance<SL>(
+                        sign_bits,
+                        ex_code,
+                        ex_fac,
+                        rotated_q,
+                        qr_base,
+                        d,
+                        ex_bits,
+                        metric_type);
+            });
 }
 
 void populate_block_aux_from_flat_storage(

data/vendor/faiss/faiss/impl/RaBitQUtils.h

@@ -10,6 +10,9 @@
 #include <faiss/MetricType.h>
 #include <faiss/impl/platform_macros.h>
 #include <faiss/utils/AlignedTable.h>
+#include <faiss/utils/rabitq_simd.h>
+#include <faiss/utils/simd_levels.h>
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
@@ -337,6 +340,33 @@ float compute_full_multibit_distance(
         size_t ex_bits,
         MetricType metric_type);
 
+// SIMDLevel-templatized version — avoids per-call dynamic dispatch.
+// Inline so it can be used from templatized distance computers without
+// needing explicit instantiations in per-SIMD TUs.
+template <SIMDLevel SL>
+inline 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<SL>(
+            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;
+}
+
 /** Compute pointer to a vector's auxiliary data within block layout. */
 template <typename T>
 inline T* get_block_aux_ptr(
@@ -350,6 +380,41 @@ inline T* get_block_aux_ptr(
             (vec_pos % bbs) * storage_size;
 }
 
+/// Extract sign bits from PQ4-interleaved block into flat byte packing.
+/// Like CodePackerRaBitQ::unpack_1 but sign-bits-only and with the
+/// vector's in-block address hoisted out of the per-SQ loop.
+inline void unpack_sign_bits_from_packed(
+        const uint8_t* block,
+        size_t bbs,
+        size_t nsq,
+        size_t offset,
+        size_t block_stride,
+        uint8_t* sign_bits_out) {
+    block += (offset / bbs) * block_stride;
+    offset = offset % bbs;
+
+    const bool nibble_high = offset > 15;
+    const size_t vid = offset & 15;
+    const size_t in_group_addr =
+            (vid < 8) ? (vid << 1) : (((vid - 8) << 1) + 1);
+
+    const size_t num_pairs = nsq / 2;
+    for (size_t k = 0; k < num_pairs; k++) {
+        const size_t base = k * bbs;
+        const uint8_t raw_even = block[base + in_group_addr];
+        const uint8_t raw_odd = block[base + in_group_addr + 16];
+
+        const uint8_t nib0 = nibble_high ? (raw_even >> 4) : (raw_even & 0xF);
+        const uint8_t nib1 = nibble_high ? (raw_odd >> 4) : (raw_odd & 0xF);
+        sign_bits_out[k] = nib0 | (nib1 << 4);
+    }
+
+    if (nsq & 1) {
+        const uint8_t raw = block[num_pairs * bbs + in_group_addr];
+        sign_bits_out[num_pairs] = nibble_high ? (raw >> 4) : (raw & 0xF);
+    }
+}
+
 /** Compute per-vector auxiliary storage size.
  *
  * @param nb_bits  number of quantization bits (1 = sign-bit only)