faiss 0.6.2 → 0.6.3

data/vendor/faiss/faiss/impl/result_handler/ResultHandler.cpp

@@ -0,0 +1,195 @@
+/*
+ * 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.
+ */
+
+// Generic (NONE) implementations of Top1 and Reservoir add_results, plus the
+// runtime-dispatch method bodies.  SIMD specialisations live in
+// result_handler_avx2.cpp and result_handler_avx512.cpp.
+
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/simd_dispatch.h>
+
+namespace faiss {
+
+// ----------------------------------------------------------------
+// SIMD-level masks
+// ----------------------------------------------------------------
+
+// Top-1: scalar fallback + AVX2 (8-wide) + AVX512 (16-wide).
+constexpr int TOP1_SIMD_LEVELS = (1 << int(SIMDLevel::NONE)) |
+        (1 << int(SIMDLevel::AVX2)) | (1 << int(SIMDLevel::AVX512));
+
+// Reservoir: scalar fallback + AVX512 compress path.
+// VPCOMPRESSPS/VPCOMPRESSD require AVX512F so there is no AVX2 path.
+// On non-AVX512 hosts the dispatch falls back to NONE automatically.
+constexpr int RESERVOIR_SIMD_LEVELS =
+        (1 << int(SIMDLevel::NONE)) | (1 << int(SIMDLevel::AVX512));
+
+// ----------------------------------------------------------------
+// Scalar (NONE) helper implementations
+// ----------------------------------------------------------------
+
+namespace {
+
+template <class C, bool use_sel>
+void top1_add_results_none(
+        Top1BlockResultHandler<C, use_sel>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_tab_in) {
+    using T = typename C::T;
+    using TI = typename C::TI;
+
+    for (size_t qi = self->i0; qi < self->i1; qi++) {
+        const T* dis_tab_i = dis_tab_in + (j1 - j0) * (qi - self->i0) - j0;
+
+        // Hoist best_dis / best_idx into locals so the compiler keeps them in
+        // registers across the inner loop (no aliasing with dis_tab reads).
+        T best_dis = self->dis_tab[qi];
+        TI best_idx = self->ids_tab[qi];
+
+        for (size_t j = j0; j < j1; j++) {
+            if (C::cmp(best_dis, dis_tab_i[j])) {
+                best_dis = dis_tab_i[j];
+                best_idx = (TI)j;
+            }
+        }
+
+        self->dis_tab[qi] = best_dis;
+        self->ids_tab[qi] = best_idx;
+    }
+}
+
+template <class C, bool use_sel>
+void reservoir_add_results_none(
+        ReservoirBlockResultHandler<C, use_sel>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_in) {
+    using T = typename C::T;
+    using TI = typename C::TI;
+
+#pragma omp parallel for
+    for (int64_t qi = (int64_t)self->i0; qi < (int64_t)self->i1; qi++) {
+        ReservoirTopN<C>& res = self->reservoirs[qi - (int64_t)self->i0];
+        const T* dis_tab_i = dis_in + (j1 - j0) * (qi - (int64_t)self->i0) - j0;
+
+        // Hoist res.i and res.threshold into locals so the compiler keeps
+        // them in registers.
+        size_t ri = res.i;
+        T thresh = res.threshold;
+
+        for (size_t j = j0; j < j1; j++) {
+            T dis = dis_tab_i[j];
+            if (C::cmp(thresh, dis)) {
+                res.vals[ri] = dis;
+                res.ids[ri] = (TI)j;
+                ri++;
+                if (ri >= res.capacity) {
+                    res.i = ri;
+                    res.shrink_fuzzy();
+                    ri = res.i;
+                    thresh = res.threshold;
+                }
+            }
+        }
+        res.i = ri;
+    }
+}
+
+} // namespace
+
+// ----------------------------------------------------------------
+// SIMDLevel::NONE explicit specialisations
+// ----------------------------------------------------------------
+
+// Instantiate top1_add_results_tpl<C, use_sel, SIMDLevel::NONE> and
+// reservoir_add_results_tpl<C, use_sel, SIMDLevel::NONE> for all
+// (C, use_sel) combinations that the rest of FAISS uses.
+#define INSTANTIATE_NONE(C, use_sel)                                  \
+    template <>                                                       \
+    void top1_add_results_tpl<C, use_sel, SIMDLevel::NONE>(           \
+            Top1BlockResultHandler<C, use_sel> * self,                \
+            size_t j0,                                                \
+            size_t j1,                                                \
+            const float* dis_tab) {                                   \
+        top1_add_results_none<C, use_sel>(self, j0, j1, dis_tab);     \
+    }                                                                 \
+    template <>                                                       \
+    void reservoir_add_results_tpl<C, use_sel, SIMDLevel::NONE>(      \
+            ReservoirBlockResultHandler<C, use_sel> * self,           \
+            size_t j0,                                                \
+            size_t j1,                                                \
+            const float* dis_in) {                                    \
+        reservoir_add_results_none<C, use_sel>(self, j0, j1, dis_in); \
+    }
+
+// Type aliases so the comma in CMax<float, int64_t> doesn't split macro args.
+using CMaxFI = CMax<float, int64_t>;
+using CMinFI = CMin<float, int64_t>;
+
+INSTANTIATE_NONE(CMaxFI, false)
+INSTANTIATE_NONE(CMaxFI, true)
+INSTANTIATE_NONE(CMinFI, false)
+INSTANTIATE_NONE(CMinFI, true)
+
+#undef INSTANTIATE_NONE
+
+// ----------------------------------------------------------------
+// add_results method definitions — dispatch to the right SL kernel
+// ----------------------------------------------------------------
+
+template <class C, bool use_sel>
+void Top1BlockResultHandler<C, use_sel>::add_results(
+        size_t j0,
+        size_t j1,
+        const T* dis_tab_2) {
+    with_selected_simd_levels<TOP1_SIMD_LEVELS>([&]<SIMDLevel SL>() {
+        top1_add_results_tpl<C, use_sel, SL>(this, j0, j1, dis_tab_2);
+    });
+}
+
+template <class C, bool use_sel>
+void ReservoirBlockResultHandler<C, use_sel>::add_results(
+        size_t j0,
+        size_t j1,
+        const T* dis_in) {
+    with_selected_simd_levels<RESERVOIR_SIMD_LEVELS>([&]<SIMDLevel SL>() {
+        reservoir_add_results_tpl<C, use_sel, SL>(this, j0, j1, dis_in);
+    });
+}
+
+// ----------------------------------------------------------------
+// Explicit class-template instantiations (force linkage)
+// ----------------------------------------------------------------
+
+template void Top1BlockResultHandler<CMax<float, int64_t>, false>::add_results(
+        size_t,
+        size_t,
+        const float*);
+template void Top1BlockResultHandler<CMax<float, int64_t>, true>::add_results(
+        size_t,
+        size_t,
+        const float*);
+template void Top1BlockResultHandler<CMin<float, int64_t>, false>::add_results(
+        size_t,
+        size_t,
+        const float*);
+template void Top1BlockResultHandler<CMin<float, int64_t>, true>::add_results(
+        size_t,
+        size_t,
+        const float*);
+
+template void ReservoirBlockResultHandler<CMax<float, int64_t>, false>::
+        add_results(size_t, size_t, const float*);
+template void ReservoirBlockResultHandler<CMax<float, int64_t>, true>::
+        add_results(size_t, size_t, const float*);
+template void ReservoirBlockResultHandler<CMin<float, int64_t>, false>::
+        add_results(size_t, size_t, const float*);
+template void ReservoirBlockResultHandler<CMin<float, int64_t>, true>::
+        add_results(size_t, size_t, const float*);
+
+} // namespace faiss

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

@@ -0,0 +1,133 @@
+/*
+ * 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.
+ */
+
+// AVX2 specialisation of Top1 add_results (8-wide branchless argmin/argmax).
+// Reservoir stays on the NONE path — VPCOMPRESSPS requires AVX512F.
+
+#ifdef COMPILE_SIMD_AVX2
+
+#include <faiss/impl/ResultHandler.h>
+
+#include <immintrin.h>
+#include <type_traits>
+
+namespace faiss {
+
+namespace {
+
+/// Templated AVX2 implementation of Top1 add_results for both CMax (keeps the
+/// smallest distance) and CMin (keeps the largest similarity).
+template <class C, bool use_sel>
+void top1_add_results_avx2(
+        Top1BlockResultHandler<C, use_sel>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_tab_in) {
+    static_assert(
+            std::is_same<typename C::T, float>::value,
+            "This code expects float distances");
+    using TI = typename C::TI;
+    const __m256i vstep = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
+
+    for (size_t qi = self->i0; qi < self->i1; qi++) {
+        const float* dis_tab_i = dis_tab_in + (j1 - j0) * (qi - self->i0) - j0;
+
+        // Hoist best_dis / best_idx into locals so the compiler keeps them in
+        // registers across the inner loop (no aliasing with dis_tab reads).
+        float best_dis = self->dis_tab[qi];
+        TI best_idx = self->ids_tab[qi];
+        size_t j = j0;
+
+        __m256 vbest = _mm256_set1_ps(best_dis);
+        __m256i vbest_idx = _mm256_set1_epi32((int32_t)best_idx);
+
+        for (; j + 8 <= j1; j += 8) {
+            __m256 vdis = _mm256_loadu_ps(dis_tab_i + j);
+            __m256i vidx =
+                    _mm256_add_epi32(_mm256_set1_epi32((int32_t)j), vstep);
+
+            // CMax (L2 nearest neighbour): keep lane if dis < best.
+            // CMin (inner product):        keep lane if dis > best.
+            __m256 mask;
+            if constexpr (C::is_max) {
+                mask = _mm256_cmp_ps(vdis, vbest, _CMP_LT_OS);
+            } else {
+                mask = _mm256_cmp_ps(vdis, vbest, _CMP_GT_OS);
+            }
+            vbest = _mm256_blendv_ps(vbest, vdis, mask);
+            vbest_idx = _mm256_blendv_epi8(
+                    vbest_idx, vidx, _mm256_castps_si256(mask));
+        }
+
+        // Horizontal reduction across 8 lanes.
+        alignas(32) float best_arr[8];
+        alignas(32) int32_t idx_arr[8];
+        _mm256_store_ps(best_arr, vbest);
+        _mm256_store_si256((__m256i*)idx_arr, vbest_idx);
+        for (int k = 0; k < 8; k++) {
+            if (C::cmp(best_dis, best_arr[k])) {
+                best_dis = best_arr[k];
+                best_idx = (TI)idx_arr[k];
+            }
+        }
+
+        // Scalar tail.
+        for (; j < j1; j++) {
+            if (C::cmp(best_dis, dis_tab_i[j])) {
+                best_dis = dis_tab_i[j];
+                best_idx = (TI)j;
+            }
+        }
+
+        self->dis_tab[qi] = best_dis;
+        self->ids_tab[qi] = best_idx;
+    }
+}
+
+} // namespace
+
+// Explicit specialisations for AVX2
+
+template <>
+void top1_add_results_tpl<CMax<float, int64_t>, false, SIMDLevel::AVX2>(
+        Top1BlockResultHandler<CMax<float, int64_t>, false>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_tab) {
+    top1_add_results_avx2<CMax<float, int64_t>, false>(self, j0, j1, dis_tab);
+}
+
+template <>
+void top1_add_results_tpl<CMax<float, int64_t>, true, SIMDLevel::AVX2>(
+        Top1BlockResultHandler<CMax<float, int64_t>, true>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_tab) {
+    top1_add_results_avx2<CMax<float, int64_t>, true>(self, j0, j1, dis_tab);
+}
+
+template <>
+void top1_add_results_tpl<CMin<float, int64_t>, false, SIMDLevel::AVX2>(
+        Top1BlockResultHandler<CMin<float, int64_t>, false>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_tab) {
+    top1_add_results_avx2<CMin<float, int64_t>, false>(self, j0, j1, dis_tab);
+}
+
+template <>
+void top1_add_results_tpl<CMin<float, int64_t>, true, SIMDLevel::AVX2>(
+        Top1BlockResultHandler<CMin<float, int64_t>, true>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_tab) {
+    top1_add_results_avx2<CMin<float, int64_t>, true>(self, j0, j1, dis_tab);
+}
+
+} // namespace faiss
+
+#endif // COMPILE_SIMD_AVX2

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

@@ -0,0 +1,281 @@
+/*
+ * 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.
+ */
+
+// AVX-512 specialisations of Top1 and Reservoir add_results.
+//
+// Top-1:    16-wide branchless argmin/argmax via mask_blend.
+// Reservoir: VPCOMPRESSPS / VPCOMPRESSD bulk-insert of passing elements,
+//            eliminating the per-element threshold branch entirely.
+
+#ifdef COMPILE_SIMD_AVX512
+
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/utils/popcount.h>
+
+#include <immintrin.h>
+#include <type_traits>
+
+namespace faiss {
+
+namespace {
+
+/// Templated AVX-512 implementation of Top1 add_results for both CMax (keeps
+/// the smallest distance) and CMin (keeps the largest similarity).
+template <class C, bool use_sel>
+void top1_add_results_avx512(
+        Top1BlockResultHandler<C, use_sel>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_tab_in) {
+    static_assert(
+            std::is_same<typename C::T, float>::value,
+            "This code expects float distances");
+    using TI = typename C::TI;
+    const __m512i vstep = _mm512_set_epi32(
+            15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0);
+
+    for (size_t qi = self->i0; qi < self->i1; qi++) {
+        const float* dis_tab_i = dis_tab_in + (j1 - j0) * (qi - self->i0) - j0;
+
+        // Hoist best_dis / best_idx into locals so the compiler keeps them in
+        // registers across the inner loop (no aliasing with dis_tab reads).
+        float best_dis = self->dis_tab[qi];
+        TI best_idx = self->ids_tab[qi];
+        size_t j = j0;
+
+        __m512 vbest = _mm512_set1_ps(best_dis);
+        __m512i vbest_idx = _mm512_set1_epi32((int32_t)best_idx);
+
+        for (; j + 16 <= j1; j += 16) {
+            __m512 vdis = _mm512_loadu_ps(dis_tab_i + j);
+            __m512i vidx =
+                    _mm512_add_epi32(_mm512_set1_epi32((int32_t)j), vstep);
+
+            // CMax (L2 nearest neighbour): keep lane if dis < best.
+            // CMin (inner product):        keep lane if dis > best.
+            __mmask16 mask;
+            if constexpr (C::is_max) {
+                mask = _mm512_cmp_ps_mask(vdis, vbest, _CMP_LT_OS);
+            } else {
+                mask = _mm512_cmp_ps_mask(vdis, vbest, _CMP_GT_OS);
+            }
+            vbest = _mm512_mask_blend_ps(mask, vbest, vdis);
+            vbest_idx = _mm512_mask_blend_epi32(mask, vbest_idx, vidx);
+        }
+
+        // Horizontal reduction across 16 lanes.
+        alignas(64) float best_arr[16];
+        alignas(64) int32_t idx_arr[16];
+        _mm512_store_ps(best_arr, vbest);
+        _mm512_store_si512((__m512i*)idx_arr, vbest_idx);
+        for (int k = 0; k < 16; k++) {
+            if (C::cmp(best_dis, best_arr[k])) {
+                best_dis = best_arr[k];
+                best_idx = (TI)idx_arr[k];
+            }
+        }
+
+        // Scalar tail.
+        for (; j < j1; j++) {
+            if (C::cmp(best_dis, dis_tab_i[j])) {
+                best_dis = dis_tab_i[j];
+                best_idx = (TI)j;
+            }
+        }
+
+        self->dis_tab[qi] = best_dis;
+        self->ids_tab[qi] = best_idx;
+    }
+}
+
+/// Templated AVX-512 implementation of Reservoir add_results for both CMax
+/// and CMin.  Uses VPCOMPRESSPS / VPCOMPRESSD to bulk-insert all elements that
+/// beat the current threshold in a single pass, avoiding the per-element branch
+/// that dominates the scalar path.
+///
+/// Falls back to the scalar NONE path for small reservoirs (capacity < 32)
+/// where the compress-path setup cost outweighs its throughput benefit.
+template <class C, bool use_sel>
+void reservoir_add_results_avx512(
+        ReservoirBlockResultHandler<C, use_sel>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_in) {
+    static_assert(
+            std::is_same<typename C::T, float>::value,
+            "This code expects float distances");
+    static_assert(
+            std::is_same<typename C::TI, int64_t>::value,
+            "This code expects int64_t indices");
+
+    // AVX-512 compress amortizes its setup cost only for large reservoirs.
+    // Benchmarks show a ~9% regression at k=10 (capacity≈20) and a ~25%
+    // gain at k=100 (capacity≈200).  32 = 2 × lane-width is the crossover.
+    // All reservoirs in a handler share the same capacity, so check once.
+    constexpr size_t AVX512_RESERVOIR_MIN_CAPACITY = 32;
+    if (self->i0 < self->i1 &&
+        self->reservoirs[0].capacity < AVX512_RESERVOIR_MIN_CAPACITY) {
+        reservoir_add_results_tpl<C, use_sel, SIMDLevel::NONE>(
+                self, j0, j1, dis_in);
+        return;
+    }
+
+    const __m512i vstep = _mm512_set_epi32(
+            15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0);
+
+#pragma omp parallel for
+    for (int64_t qi = (int64_t)self->i0; qi < (int64_t)self->i1; qi++) {
+        ReservoirTopN<C>& res = self->reservoirs[qi - (int64_t)self->i0];
+        const float* dis_tab_i =
+                dis_in + (j1 - j0) * (qi - (int64_t)self->i0) - j0;
+        size_t j = j0;
+
+        for (; j + 16 <= j1; j += 16) {
+            // Near-capacity: fewer than 16 free slots remain.  Fall back to
+            // the scalar add_result which handles overflow correctly.
+            if (res.i + 16 > res.capacity) {
+                for (size_t jj = j; jj < j + 16; jj++)
+                    res.add_result(dis_tab_i[jj], jj);
+                continue;
+            }
+
+            __m512 vthresh = _mm512_set1_ps(res.threshold);
+            __m512 vdis = _mm512_loadu_ps(dis_tab_i + j);
+
+            // CMax (L2): keep elements with dis < threshold.
+            // CMin (IP): keep elements with dis > threshold.
+            __mmask16 mask;
+            if constexpr (C::is_max) {
+                mask = _mm512_cmp_ps_mask(vdis, vthresh, _CMP_LT_OS);
+            } else {
+                mask = _mm512_cmp_ps_mask(vdis, vthresh, _CMP_GT_OS);
+            }
+            if (mask == 0)
+                continue;
+
+            int count = popcount32(mask);
+
+            // Compress passing distances into a contiguous run (VPCOMPRESSPS).
+            __m512 passing_dis = _mm512_maskz_compress_ps(mask, vdis);
+
+            // Compress the sequential indices j..j+15 (VPCOMPRESSD).
+            __m512i vidx32 =
+                    _mm512_add_epi32(_mm512_set1_epi32((int32_t)j), vstep);
+            __m512i passing_idx32 = _mm512_maskz_compress_epi32(mask, vidx32);
+
+            // Unconditional 16-element stores are safe: the res.i + 16 <=
+            // res.capacity check above guarantees slots res.i .. res.i+15
+            // are all unused.  The (count .. 15) tail positions get garbage
+            // that res.i never reaches.
+            _mm512_storeu_ps(res.vals + res.i, passing_dis);
+
+            // Widen int32 indices to int64 (TI = int64_t) in two 8-element
+            // halves and store them.
+            _mm512_storeu_si512(
+                    (void*)(res.ids + res.i),
+                    _mm512_cvtepi32_epi64(
+                            _mm512_castsi512_si256(passing_idx32)));
+            _mm512_storeu_si512(
+                    (void*)(res.ids + res.i + 8),
+                    _mm512_cvtepi32_epi64(
+                            _mm512_extracti64x4_epi64(passing_idx32, 1)));
+
+            res.i += count;
+            if (res.i >= res.capacity) {
+                res.shrink_fuzzy();
+            }
+        }
+
+        // Scalar tail.
+        for (; j < j1; j++)
+            res.add_result(dis_tab_i[j], j);
+    }
+}
+
+} // namespace
+
+// Explicit specialisations for AVX-512
+
+template <>
+void top1_add_results_tpl<CMax<float, int64_t>, false, SIMDLevel::AVX512>(
+        Top1BlockResultHandler<CMax<float, int64_t>, false>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_tab) {
+    top1_add_results_avx512<CMax<float, int64_t>, false>(self, j0, j1, dis_tab);
+}
+
+template <>
+void top1_add_results_tpl<CMax<float, int64_t>, true, SIMDLevel::AVX512>(
+        Top1BlockResultHandler<CMax<float, int64_t>, true>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_tab) {
+    top1_add_results_avx512<CMax<float, int64_t>, true>(self, j0, j1, dis_tab);
+}
+
+template <>
+void top1_add_results_tpl<CMin<float, int64_t>, false, SIMDLevel::AVX512>(
+        Top1BlockResultHandler<CMin<float, int64_t>, false>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_tab) {
+    top1_add_results_avx512<CMin<float, int64_t>, false>(self, j0, j1, dis_tab);
+}
+
+template <>
+void top1_add_results_tpl<CMin<float, int64_t>, true, SIMDLevel::AVX512>(
+        Top1BlockResultHandler<CMin<float, int64_t>, true>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_tab) {
+    top1_add_results_avx512<CMin<float, int64_t>, true>(self, j0, j1, dis_tab);
+}
+
+template <>
+void reservoir_add_results_tpl<CMax<float, int64_t>, false, SIMDLevel::AVX512>(
+        ReservoirBlockResultHandler<CMax<float, int64_t>, false>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_in) {
+    reservoir_add_results_avx512<CMax<float, int64_t>, false>(
+            self, j0, j1, dis_in);
+}
+
+template <>
+void reservoir_add_results_tpl<CMax<float, int64_t>, true, SIMDLevel::AVX512>(
+        ReservoirBlockResultHandler<CMax<float, int64_t>, true>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_in) {
+    reservoir_add_results_avx512<CMax<float, int64_t>, true>(
+            self, j0, j1, dis_in);
+}
+
+template <>
+void reservoir_add_results_tpl<CMin<float, int64_t>, false, SIMDLevel::AVX512>(
+        ReservoirBlockResultHandler<CMin<float, int64_t>, false>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_in) {
+    reservoir_add_results_avx512<CMin<float, int64_t>, false>(
+            self, j0, j1, dis_in);
+}
+
+template <>
+void reservoir_add_results_tpl<CMin<float, int64_t>, true, SIMDLevel::AVX512>(
+        ReservoirBlockResultHandler<CMin<float, int64_t>, true>* self,
+        size_t j0,
+        size_t j1,
+        const float* dis_in) {
+    reservoir_add_results_avx512<CMin<float, int64_t>, true>(
+            self, j0, j1, dis_in);
+}
+
+} // namespace faiss
+
+#endif // COMPILE_SIMD_AVX512

data/vendor/faiss/faiss/impl/scalar_quantizer/EDENQuantizer-avx2.cpp

@@ -0,0 +1,72 @@
+/*
+ * 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.
+ */
+
+#ifdef COMPILE_SIMD_AVX2
+
+#include <faiss/impl/EDENQuantizer.h>
+#include <faiss/impl/simdlib/simdlib_avx2.h>
+
+#include <immintrin.h>
+
+namespace faiss {
+
+namespace eden_distance {
+
+namespace {
+
+constexpr int kCodeDotLUTKindHalfByte = 2;
+
+inline __m256i load_code_bytes_8(const uint8_t* const code[8], size_t i) {
+    return _mm256_setr_epi32(
+            code[0][i],
+            code[1][i],
+            code[2][i],
+            code[3][i],
+            code[4][i],
+            code[5][i],
+            code[6][i],
+            code[7][i]);
+}
+
+} // namespace
+
+void compute_code_dot_lut_batch_8_avx2(
+        const uint8_t* const code[8],
+        const float* lut,
+        int lut_kind,
+        size_t packed_size,
+        float dots[8]) {
+    __m256 acc = _mm256_setzero_ps();
+    const __m256i low_mask = _mm256_set1_epi32(0x0f);
+    const __m256i high_offset = _mm256_set1_epi32(16);
+
+    if (lut_kind == kCodeDotLUTKindHalfByte) {
+        for (size_t i = 0; i < packed_size; i++) {
+            const __m256i bytes = load_code_bytes_8(code, i);
+            const __m256i low = _mm256_and_si256(bytes, low_mask);
+            const __m256i high =
+                    _mm256_add_epi32(_mm256_srli_epi32(bytes, 4), high_offset);
+            acc = _mm256_add_ps(acc, _mm256_i32gather_ps(lut, low, 4));
+            acc = _mm256_add_ps(acc, _mm256_i32gather_ps(lut, high, 4));
+            lut += 32;
+        }
+    } else {
+        for (size_t i = 0; i < packed_size; i++) {
+            const __m256i bytes = load_code_bytes_8(code, i);
+            acc = _mm256_add_ps(acc, _mm256_i32gather_ps(lut, bytes, 4));
+            lut += 256;
+        }
+    }
+
+    _mm256_storeu_ps(dots, acc);
+}
+
+} // namespace eden_distance
+
+} // namespace faiss
+
+#endif