faiss 0.6.2 → 0.6.3

data/vendor/faiss/faiss/impl/index_write.cpp

@@ -25,11 +25,13 @@
 #include <faiss/Index2Layer.h>
 #include <faiss/IndexAdditiveQuantizer.h>
 #include <faiss/IndexAdditiveQuantizerFastScan.h>
+#include <faiss/IndexEDEN.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexHNSW.h>
 #include <faiss/IndexIVF.h>
 #include <faiss/IndexIVFAdditiveQuantizer.h>
 #include <faiss/IndexIVFAdditiveQuantizerFastScan.h>
+#include <faiss/IndexIVFEDEN.h>
 #include <faiss/IndexIVFFlat.h>
 #include <faiss/IndexIVFFlatPanorama.h>
 #include <faiss/IndexIVFIndependentQuantizer.h>
@@ -63,6 +65,7 @@
 #include <faiss/IndexScalarQuantizer.h>
 #include <faiss/MetaIndexes.h>
 #include <faiss/VectorTransform.h>
+#include <faiss/impl/EDENQuantizer.h>
 
 #include <faiss/IndexBinaryFlat.h>
 #include <faiss/IndexBinaryFromFloat.h>
@@ -439,6 +442,20 @@ static void write_RaBitQuantizer(
     }
 }
 
+static void write_EDENScalarQuantizer(
+        const ScalarQuantizer* sq,
+        MetricType metric_type,
+        EDENScaleType scale_type,
+        size_t full_code_size,
+        IOWriter* f) {
+    WRITE1(sq->d);
+    WRITE1(full_code_size);
+    WRITE1(metric_type);
+    WRITE1(sq->bits);
+    int scale_type_int = static_cast<int>(scale_type);
+    WRITE1(scale_type_int);
+}
+
 static void write_direct_map(const DirectMap* dm, IOWriter* f) {
     char maintain_direct_map =
             (char)dm->type; // for backwards compatibility with bool
@@ -871,8 +888,12 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
                 : dynamic_cast<const IndexHNSWSQ*>(idx)     ? fourcc("IHNs")
                 : dynamic_cast<const IndexHNSW2Level*>(idx) ? fourcc("IHN2")
                 : dynamic_cast<const IndexHNSWCagra*>(idx)  ? fourcc("IHc2")
+                : typeid(*idx) == typeid(IndexHNSW)         ? fourcc("IH00")
                                                             : 0;
-        FAISS_THROW_IF_NOT(h != 0);
+        FAISS_THROW_IF_NOT_FMT(
+                h != 0,
+                "don't know how to serialize this IndexHNSW subtype: %s",
+                typeid(*idx).name());
         WRITE1(h);
         write_index_header(idxhnsw, f);
         if (h == fourcc("IHfP")) {
@@ -900,7 +921,8 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
                 : dynamic_cast<const IndexNSGPQ*>(idx)      ? fourcc("INSp")
                 : dynamic_cast<const IndexNSGSQ*>(idx)      ? fourcc("INSs")
                                                             : 0;
-        FAISS_THROW_IF_NOT(h != 0);
+        FAISS_THROW_IF_MSG(
+                h == 0, "don't know how to serialize this IndexNSG subtype");
         WRITE1(h);
         write_index_header(idxnsg, f);
         WRITE1(idxnsg->GK);
@@ -915,9 +937,11 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
             const IndexNNDescent* idxnnd =
                     dynamic_cast<const IndexNNDescent*>(idx)) {
         auto idxnndflat = dynamic_cast<const IndexNNDescentFlat*>(idx);
-        FAISS_THROW_IF_NOT(idxnndflat != nullptr);
+        FAISS_THROW_IF_NOT(idxnndflat);
         uint32_t h = fourcc("INNf");
-        FAISS_THROW_IF_NOT(h != 0);
+        FAISS_THROW_IF_MSG(
+                h == 0,
+                "don't know how to serialize this IndexNNDescent subtype");
         WRITE1(h);
         write_index_header(idxnnd, f);
         write_NNDescent(&idxnnd->nndescent, f);
@@ -965,6 +989,33 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
         WRITE1(h);
         write_index_header(imm_2, f);
         write_index(imm_2->index, f);
+    } else if (const IndexEDEN* idxe = dynamic_cast<const IndexEDEN*>(idx)) {
+        uint32_t h = fourcc("IxEe");
+        WRITE1(h);
+        write_index_header(idx, f);
+        write_EDENScalarQuantizer(
+                &idxe->sq,
+                idxe->metric_type,
+                idxe->scale_type,
+                idxe->code_size,
+                f);
+        WRITEVECTOR(idxe->codes);
+        WRITEVECTOR(idxe->center);
+    } else if (
+            const IndexIVFEDEN* iveden =
+                    dynamic_cast<const IndexIVFEDEN*>(idx)) {
+        uint32_t h = fourcc("IwEe");
+        WRITE1(h);
+        write_ivf_header(iveden, f);
+        write_EDENScalarQuantizer(
+                &iveden->sq,
+                iveden->metric_type,
+                iveden->scale_type,
+                iveden->code_size,
+                f);
+        WRITE1(iveden->code_size);
+        WRITE1(iveden->by_residual);
+        write_InvertedLists(iveden->invlists, f);
     } else if (
             const IndexRaBitQFastScan* idxqfs =
                     dynamic_cast<const IndexRaBitQFastScan*>(idx)) {

data/vendor/faiss/faiss/impl/lattice_Zn.cpp

@@ -557,12 +557,11 @@ ZnSphereCodecRec::ZnSphereCodecRec(int dim_in, int r2_in)
 
     decode_cache.resize((r2 + 1));
 
-    // The decode cache stores total_cache_entries * dimsub floats.
-    // Cap at 2^27 total floats (~512 MB), aligned with the nv_cum
-    // memory cap above which also uses 2^27 entries.  The entry count
-    // grows as the number of lattice points in dimension
-    // 2^cache_level, which is O(r2^(dim/2)) -- much faster than the
-    // O(r2^2) growth of nv_cum.
+    // The decode cache stores total_cache_entries * dimsub floats and is
+    // built with one decode() call per entry. The entry count grows as the
+    // number of lattice points in dimension 2^cache_level, which is
+    // O(r2^(dim/2)) -- much faster than the O(r2^2) growth of nv_cum, so it
+    // is bounded per-r2sub in the loop below.
     size_t total_cache_entries = 0;
     int dimsub = (1 << cache_level);
 
@@ -571,9 +570,9 @@ ZnSphereCodecRec::ZnSphereCodecRec(int dim_in, int r2_in)
         uint64_t nvi = get_nv(ld, r2sub);
         total_cache_entries += nvi;
         FAISS_THROW_IF_NOT_MSG(
-                total_cache_entries <= (size_t(1) << 27) / dimsub,
-                "ZnSphereCodecRec: r2 too large, decode cache "
-                "would require excessive memory");
+                total_cache_entries <= (size_t(1) << 27) / (size_t)dim,
+                "ZnSphereCodecRec: r2/dim too large, decode cache "
+                "would require excessive memory and computation");
         std::vector<float>& cache = decode_cache[r2sub];
         cache.resize(nvi * dimsub);
         std::vector<float> c(dim);

data/vendor/faiss/faiss/impl/platform_macros.h

@@ -25,7 +25,9 @@
 
 #define strtok_r strtok_s
 
-#ifdef _MSC_VER
+// clang-cl defines _MSC_VER but provides __PRETTY_FUNCTION__ itself, in
+// clang's format. Redefining it there breaks code that parses it.
+#if defined(_MSC_VER) && !defined(__clang__)
 #define __PRETTY_FUNCTION__ __FUNCSIG__
 #endif // _MSC_VER
 

data/vendor/faiss/faiss/impl/polysemous_training/avx512.cpp

@@ -0,0 +1,284 @@
+/*
+ * 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++ -*-
+
+#ifdef COMPILE_SIMD_AVX512
+
+#include <faiss/impl/PolysemousTraining.h>
+#include <faiss/impl/polysemous_training/dispatch.h>
+#include <faiss/utils/popcount.h>
+
+#include <immintrin.h>
+#include <cstdint>
+
+namespace faiss {
+namespace {
+
+inline int hamming_dis(uint64_t a, uint64_t b) {
+    return popcount64(a ^ b);
+}
+
+inline double sqr(double x) {
+    return x * x;
+}
+
+inline __m512i popcnt_512(__m512i v) {
+#ifdef __AVX512VPOPCNTDQ__
+    return _mm512_popcnt_epi64(v);
+#else
+    const __m128i nibble_popcount =
+            _mm_setr_epi8(0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4);
+    const __m512i lookup = _mm512_broadcast_i32x4(nibble_popcount);
+
+    const __m512i low_mask = _mm512_set1_epi8(0x0f);
+    const __m512i lo = _mm512_and_si512(v, low_mask);
+    const __m512i hi = _mm512_and_si512(_mm512_srli_epi16(v, 4), low_mask);
+
+    const __m512i popcnt_lo = _mm512_shuffle_epi8(lookup, lo);
+    const __m512i popcnt_hi = _mm512_shuffle_epi8(lookup, hi);
+    const __m512i popcnt_bytes = _mm512_add_epi8(popcnt_lo, popcnt_hi);
+
+    return _mm512_sad_epu8(popcnt_bytes, _mm512_setzero_si512());
+#endif
+}
+
+} // namespace
+
+namespace polysemous_training {
+
+template <>
+double hamming_compute_cost<SIMDLevel::AVX512>(
+        int n,
+        const int* perm,
+        const double* target_dis,
+        const double* weights) {
+    double total_cost = 0.0;
+    for (int i = 0; i < n; i++) {
+        __m512d cost_vec = _mm512_setzero_pd();
+        const __m512i perm_i_vec = _mm512_set1_epi64(perm[i]);
+        const int bro = i * n;
+        int j = 0;
+        for (; j <= n - 8; j += 8) {
+            const __m512d wanted_vec = _mm512_loadu_pd(&target_dis[bro + j]);
+            const __m512d w_vec = _mm512_loadu_pd(&weights[bro + j]);
+            const __m256i pj32 = _mm256_loadu_si256((__m256i const*)&perm[j]);
+            const __m512i pj64 = _mm512_cvtepi32_epi64(pj32);
+            const __m512i xor_res = _mm512_xor_si512(perm_i_vec, pj64);
+            const __m512d actual_vec = _mm512_cvtepi64_pd(popcnt_512(xor_res));
+            const __m512d diff = _mm512_sub_pd(wanted_vec, actual_vec);
+            cost_vec =
+                    _mm512_fmadd_pd(w_vec, _mm512_mul_pd(diff, diff), cost_vec);
+        }
+        total_cost += _mm512_reduce_add_pd(cost_vec);
+        for (; j < n; j++) {
+            double wanted = target_dis[bro + j];
+            double w = weights[bro + j];
+            double actual = hamming_dis(perm[i], perm[j]);
+            total_cost += w * sqr(wanted - actual);
+        }
+    }
+    return total_cost;
+}
+
+template <>
+double hamming_cost_update<SIMDLevel::AVX512>(
+        int n,
+        const int* perm,
+        int iw,
+        int jw,
+        const double* target_dis,
+        const double* weights) {
+    double delta_cost_scalar = 0;
+    const __m512i v_idx_base = _mm512_setr_epi64(0, 1, 2, 3, 4, 5, 6, 7);
+    // Loop-invariant across the j-loop and both process_row calls.
+    const __m512i v_iw = _mm512_set1_epi64(iw);
+    const __m512i v_jw = _mm512_set1_epi64(jw);
+    const __m512i v_perm_iw = _mm512_set1_epi64(perm[iw]);
+    const __m512i v_perm_jw = _mm512_set1_epi64(perm[jw]);
+    __m512d delta_cost_vec = _mm512_setzero_pd();
+
+    auto process_row = [&](int row, int old_pi, int new_pi) {
+        const int bro = row * n;
+        const __m512i v_old = _mm512_set1_epi64(old_pi);
+        const __m512i v_new = _mm512_set1_epi64(new_pi);
+        int j = 0;
+        for (; j <= n - 8; j += 8) {
+            __m512d wv = _mm512_loadu_pd(&target_dis[bro + j]);
+            __m512d ww = _mm512_loadu_pd(&weights[bro + j]);
+            __m256i pj32 = _mm256_loadu_si256((__m256i const*)&perm[j]);
+            __m512i pjv = _mm512_cvtepi32_epi64(pj32);
+            __m512d av = _mm512_cvtepi64_pd(
+                    popcnt_512(_mm512_xor_si512(v_old, pjv)));
+            __m512d to = _mm512_sub_pd(wv, av);
+            to = _mm512_mul_pd(to, to);
+            delta_cost_vec = _mm512_fnmadd_pd(ww, to, delta_cost_vec);
+
+            __m512i ji = _mm512_add_epi64(_mm512_set1_epi64(j), v_idx_base);
+            __mmask8 miw = _mm512_cmpeq_epi64_mask(ji, v_iw);
+            __mmask8 mjw = _mm512_cmpeq_epi64_mask(ji, v_jw);
+            __m512i pnj = _mm512_mask_blend_epi64(mjw, pjv, v_perm_iw);
+            pnj = _mm512_mask_blend_epi64(miw, pnj, v_perm_jw);
+            __m512d nav = _mm512_cvtepi64_pd(
+                    popcnt_512(_mm512_xor_si512(v_new, pnj)));
+            __m512d tn = _mm512_sub_pd(wv, nav);
+            tn = _mm512_mul_pd(tn, tn);
+            delta_cost_vec = _mm512_fmadd_pd(ww, tn, delta_cost_vec);
+        }
+        for (; j < n; j++) {
+            double wanted = target_dis[bro + j];
+            double w = weights[bro + j];
+            double actual = hamming_dis(old_pi, perm[j]);
+            delta_cost_scalar -= w * sqr(wanted - actual);
+            double new_actual = hamming_dis(
+                    new_pi,
+                    perm[j == iw           ? jw
+                                 : j == jw ? iw
+                                           : j]);
+            delta_cost_scalar += w * sqr(wanted - new_actual);
+        }
+    };
+    process_row(iw, perm[iw], perm[jw]);
+    process_row(jw, perm[jw], perm[iw]);
+
+    for (int i = 0; i < n; ++i) {
+        if (i == iw || i == jw)
+            continue;
+        int j = iw;
+        {
+            double wanted = target_dis[i * n + j];
+            double w = weights[i * n + j];
+            delta_cost_scalar -=
+                    w * sqr(wanted - hamming_dis(perm[i], perm[j]));
+            delta_cost_scalar +=
+                    w * sqr(wanted - hamming_dis(perm[i], perm[jw]));
+        }
+        j = jw;
+        {
+            double wanted = target_dis[i * n + j];
+            double w = weights[i * n + j];
+            delta_cost_scalar -=
+                    w * sqr(wanted - hamming_dis(perm[i], perm[j]));
+            delta_cost_scalar +=
+                    w * sqr(wanted - hamming_dis(perm[i], perm[iw]));
+        }
+    }
+    return _mm512_reduce_add_pd(delta_cost_vec) + delta_cost_scalar;
+}
+
+template <>
+double distances_compute_cost<SIMDLevel::AVX512>(
+        const ReproduceDistancesObjective& obj,
+        const int* perm) {
+    const int n = obj.n;
+    double total_cost = 0.0;
+    for (int i = 0; i < n; ++i) {
+        const int pi = perm[i];
+        const int bro_t = i * n;
+        const int bro_s = pi * n;
+        __m512d sum = _mm512_setzero_pd();
+        int j = 0;
+        for (; j <= n - 8; j += 8) {
+            __m512d wv = _mm512_loadu_pd(&obj.target_dis[bro_t + j]);
+            __m512d ww = _mm512_loadu_pd(&obj.weights[bro_t + j]);
+            __m256i pj = _mm256_loadu_si256(
+                    reinterpret_cast<const __m256i*>(&perm[j]));
+            __m256i idx = _mm256_add_epi32(_mm256_set1_epi32(bro_s), pj);
+            __m512d av = _mm512_i32gather_pd(idx, obj.source_dis.data(), 8);
+            __m512d d = _mm512_sub_pd(wv, av);
+            sum = _mm512_fmadd_pd(_mm512_mul_pd(d, d), ww, sum);
+        }
+        total_cost += _mm512_reduce_add_pd(sum);
+        for (; j < n; ++j) {
+            double wanted = obj.target_dis[bro_t + j];
+            double w = obj.weights[bro_t + j];
+            double actual = obj.get_source_dis(pi, perm[j]);
+            total_cost += w * sqr(wanted - actual);
+        }
+    }
+    return total_cost;
+}
+
+template <>
+double distances_cost_update<SIMDLevel::AVX512>(
+        const ReproduceDistancesObjective& obj,
+        const int* perm,
+        int iw,
+        int jw) {
+    const int n = obj.n;
+    double delta_cost = 0.0;
+    const int p_iw = perm[iw], p_jw = perm[jw];
+    const __m256i v_joff = _mm256_set_epi32(7, 6, 5, 4, 3, 2, 1, 0);
+    const __m256i vi = _mm256_set1_epi32(iw);
+    const __m256i vj = _mm256_set1_epi32(jw);
+    const __m256i vpi = _mm256_set1_epi32(p_iw);
+    const __m256i vpj = _mm256_set1_epi32(p_jw);
+    const __m256i vpin = _mm256_set1_epi32(p_iw * n);
+    const __m256i vpjn = _mm256_set1_epi32(p_jw * n);
+
+    auto process_row =
+            [&](int row, int old_p, int new_p, __m256i old_pn, __m256i new_pn) {
+                const int bro = row * n;
+                __m512d dv = _mm512_setzero_pd();
+                int j = 0;
+                for (; j <= n - 8; j += 8) {
+                    __m512d wv = _mm512_loadu_pd(&obj.target_dis[bro + j]);
+                    __m512d ww = _mm512_loadu_pd(&obj.weights[bro + j]);
+                    __m256i pjv = _mm256_loadu_si256(
+                            reinterpret_cast<const __m256i*>(&perm[j]));
+                    __m256i ia = _mm256_add_epi32(old_pn, pjv);
+                    __m512d av =
+                            _mm512_i32gather_pd(ia, obj.source_dis.data(), 8);
+                    __m512d da = _mm512_sub_pd(wv, av);
+                    dv = _mm512_fnmadd_pd(ww, _mm512_mul_pd(da, da), dv);
+                    __m256i vk = _mm256_add_epi32(_mm256_set1_epi32(j), v_joff);
+                    __mmask8 mi = _mm256_cmpeq_epi32_mask(vk, vi);
+                    __mmask8 mj = _mm256_cmpeq_epi32_mask(vk, vj);
+                    __m256i pnj = _mm256_mask_blend_epi32(mi, pjv, vpj);
+                    pnj = _mm256_mask_blend_epi32(mj, pnj, vpi);
+                    __m256i in2 = _mm256_add_epi32(new_pn, pnj);
+                    __m512d nav =
+                            _mm512_i32gather_pd(in2, obj.source_dis.data(), 8);
+                    __m512d dn = _mm512_sub_pd(wv, nav);
+                    dv = _mm512_fmadd_pd(ww, _mm512_mul_pd(dn, dn), dv);
+                }
+                delta_cost += _mm512_reduce_add_pd(dv);
+                for (; j < n; ++j) {
+                    double wanted = obj.target_dis[bro + j];
+                    double w = obj.weights[bro + j];
+                    double actual = obj.get_source_dis(old_p, perm[j]);
+                    delta_cost -= w * sqr(wanted - actual);
+                    int pnj = (j == iw) ? p_jw : ((j == jw) ? p_iw : perm[j]);
+                    double na = obj.get_source_dis(new_p, pnj);
+                    delta_cost += w * sqr(wanted - na);
+                }
+            };
+    process_row(iw, p_iw, p_jw, vpin, vpjn);
+    process_row(jw, p_jw, p_iw, vpjn, vpin);
+
+    for (int i = 0; i < n; ++i) {
+        if (i == iw || i == jw)
+            continue;
+        double wanted = obj.target_dis[i * n + iw], w = obj.weights[i * n + iw];
+        double actual = obj.get_source_dis(perm[i], p_iw);
+        delta_cost -= w * sqr(wanted - actual);
+        double na = obj.get_source_dis(perm[i], p_jw);
+        delta_cost += w * sqr(wanted - na);
+        wanted = obj.target_dis[i * n + jw];
+        w = obj.weights[i * n + jw];
+        actual = obj.get_source_dis(perm[i], p_jw);
+        delta_cost -= w * sqr(wanted - actual);
+        na = obj.get_source_dis(perm[i], p_iw);
+        delta_cost += w * sqr(wanted - na);
+    }
+    return delta_cost;
+}
+
+} // namespace polysemous_training
+} // namespace faiss
+
+#endif // COMPILE_SIMD_AVX512

data/vendor/faiss/faiss/impl/polysemous_training/dispatch.h

@@ -0,0 +1,115 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <faiss/utils/simd_levels.h>
+
+namespace faiss {
+
+struct ReproduceDistancesObjective;
+
+namespace polysemous_training {
+
+// Levels with a dedicated kernel: scalar (NONE) plus AVX-512. Any other runtime
+// level (AVX2, NEON, ...) falls back to NONE via with_selected_simd_levels /
+// get_simd_fallback, and AVX512_SPR reuses the AVX-512 kernel.
+constexpr int SIMD_LEVELS =
+        (1 << int(SIMDLevel::NONE)) | (1 << int(SIMDLevel::AVX512));
+
+/// compute_cost for ReproduceWithHammingObjective.
+/// Parameters mirror the objective's fields to avoid exposing the
+/// anonymous-namespace struct.
+template <SIMDLevel SL>
+double hamming_compute_cost(
+        int n,
+        const int* perm,
+        const double* target_dis,
+        const double* weights);
+
+/// cost_update for ReproduceWithHammingObjective.
+template <SIMDLevel SL>
+double hamming_cost_update(
+        int n,
+        const int* perm,
+        int iw,
+        int jw,
+        const double* target_dis,
+        const double* weights);
+
+/// compute_cost for ReproduceDistancesObjective.
+template <SIMDLevel SL>
+double distances_compute_cost(
+        const ReproduceDistancesObjective& obj,
+        const int* perm);
+
+/// cost_update for ReproduceDistancesObjective.
+template <SIMDLevel SL>
+double distances_cost_update(
+        const ReproduceDistancesObjective& obj,
+        const int* perm,
+        int iw,
+        int jw);
+
+// The scalar (NONE) specializations are defined in PolysemousTraining.cpp; the
+// AVX-512 specializations are defined in polysemous_training/avx512.cpp. Both
+// are declared here so callers instantiate the out-of-line definition rather
+// than implicitly instantiating the (undefined) primary template.
+template <>
+double hamming_compute_cost<SIMDLevel::NONE>(
+        int n,
+        const int* perm,
+        const double* target_dis,
+        const double* weights);
+template <>
+double hamming_compute_cost<SIMDLevel::AVX512>(
+        int n,
+        const int* perm,
+        const double* target_dis,
+        const double* weights);
+
+template <>
+double hamming_cost_update<SIMDLevel::NONE>(
+        int n,
+        const int* perm,
+        int iw,
+        int jw,
+        const double* target_dis,
+        const double* weights);
+template <>
+double hamming_cost_update<SIMDLevel::AVX512>(
+        int n,
+        const int* perm,
+        int iw,
+        int jw,
+        const double* target_dis,
+        const double* weights);
+
+template <>
+double distances_compute_cost<SIMDLevel::NONE>(
+        const ReproduceDistancesObjective& obj,
+        const int* perm);
+template <>
+double distances_compute_cost<SIMDLevel::AVX512>(
+        const ReproduceDistancesObjective& obj,
+        const int* perm);
+
+template <>
+double distances_cost_update<SIMDLevel::NONE>(
+        const ReproduceDistancesObjective& obj,
+        const int* perm,
+        int iw,
+        int jw);
+template <>
+double distances_cost_update<SIMDLevel::AVX512>(
+        const ReproduceDistancesObjective& obj,
+        const int* perm,
+        int iw,
+        int jw);
+
+} // namespace polysemous_training
+} // namespace faiss

data/vendor/faiss/faiss/impl/pq_code_distance/IVFPQ_QueryTables.cpp

@@ -7,7 +7,6 @@
 
 #include <faiss/impl/pq_code_distance/IVFPQ_QueryTables.h>
 
-#include <faiss/Clustering.h>
 #include <faiss/utils/distances_dispatch.h>
 
 namespace faiss {

data/vendor/faiss/faiss/impl/pq_code_distance/PQDistanceComputer_impl.h

@@ -20,23 +20,39 @@ namespace faiss {
 namespace pq_code_distance {
 
 template <class PQCodeDist>
-struct PQDistanceComputer : FlatCodesDistanceComputer {
+struct PQDistanceComputer final : FlatCodesDistanceComputer {
     using PQDecoder = typename PQCodeDist::PQDecoder;
-    size_t d;
     MetricType metric;
-    idx_t nb;
     const ProductQuantizer& pq;
     const float* sdc;
     std::vector<float> precomputed_table;
-    size_t ndis;
-    const float* q;
 
     float distance_to_code(const uint8_t* code) final {
-        ndis++;
-
-        float dis = PQCodeDist::distance_single_code(
+        return PQCodeDist::distance_single_code(
                 pq.M, pq.nbits, precomputed_table.data(), code);
-        return dis;
+    }
+
+    void distance_to_code_batch_4(
+            const uint8_t* c1,
+            const uint8_t* c2,
+            const uint8_t* c3,
+            const uint8_t* c4,
+            float& d1,
+            float& d2,
+            float& d3,
+            float& d4) override {
+        PQCodeDist::distance_four_codes(
+                pq.M,
+                pq.nbits,
+                precomputed_table.data(),
+                c1,
+                c2,
+                c3,
+                c4,
+                d1,
+                d2,
+                d3,
+                d4);
     }
 
     float symmetric_dis(idx_t i, idx_t j) override {
@@ -50,7 +66,6 @@ struct PQDistanceComputer : FlatCodesDistanceComputer {
             accu += sdci[codei.decode() + (codej.decode() << codei.nbits)];
             sdci += uint64_t(1) << (2 * codei.nbits);
         }
-        ndis++;
         return accu;
     }
 
@@ -58,18 +73,14 @@ struct PQDistanceComputer : FlatCodesDistanceComputer {
             : FlatCodesDistanceComputer(
                       storage.codes.data(),
                       storage.code_size),
-              pq(storage.pq),
-              q(nullptr) {
+              pq(storage.pq) {
         precomputed_table.resize(pq.M * pq.ksub);
-        nb = storage.ntotal;
-        d = storage.d;
         metric = storage.metric_type;
         if (pq.sdc_table.size() == pq.ksub * pq.ksub * pq.M) {
             sdc = pq.sdc_table.data();
         } else {
             sdc = nullptr;
         }
-        ndis = 0;
     }
 
     void set_query(const float* x) override {

data/vendor/faiss/faiss/impl/pq_code_distance/avx2.cpp

@@ -9,15 +9,15 @@
 
 #define THE_SIMD_LEVEL SIMDLevel::AVX2
 
-// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+// NOLINTNEXTLINE(facebook-hte-InlineHeader,facebook-unused-include-check)
 #include <faiss/impl/pq_code_distance/pq_code_distance-avx2.h>
-// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+// NOLINTNEXTLINE(facebook-hte-InlineHeader,facebook-unused-include-check)
 #include <faiss/utils/hamming_distance/hamming_computer-avx2.h>
 // NOLINTNEXTLINE(facebook-hte-InlineHeader,facebook-unused-include-check)
 #include <faiss/impl/pq_code_distance/pq_scan_impl.h>
-// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+// NOLINTNEXTLINE(facebook-hte-InlineHeader,facebook-unused-include-check)
 #include <faiss/impl/pq_code_distance/PQDistanceComputer_impl.h>
-// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+// NOLINTNEXTLINE(facebook-hte-InlineHeader,facebook-unused-include-check)
 #include <faiss/impl/pq_code_distance/IVFPQScanner_impl.h>
 
 #endif // COMPILE_SIMD_AVX2