faiss 0.6.0 → 0.6.1

data/vendor/faiss/faiss/impl/approx_topk/rq_beam_search_tab-inl.h

@@ -0,0 +1,169 @@
+/*
+ * 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.
+ */
+
+// Definitions of the SIMDLevel-templatized accum_and_*_tab functions.
+// Only included by per-ISA .cpp files (avx2.cpp, neon.cpp).
+// Do NOT include this from common translation units.
+//
+// Common TUs include rq_beam_search_tab.h (declarations only).
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+#include <faiss/impl/approx_topk/rq_beam_search_tab.h>
+#include <faiss/impl/simdlib/simdlib.h>
+
+namespace faiss {
+
+template <size_t M, size_t NK, SIMDLevel SL>
+void accum_and_store_tab(
+        const size_t m_offset,
+        const float* const __restrict codebook_cross_norms,
+        const uint64_t* const __restrict codebook_offsets,
+        const int32_t* const __restrict codes_i,
+        const size_t b,
+        const size_t ldc,
+        const size_t K,
+        float* const __restrict output) {
+    using simd_float = simd8float32_tpl<SL>;
+
+    const float* cbs[M];
+    for (size_t ij = 0; ij < M; ij++) {
+        const size_t code = static_cast<size_t>(codes_i[b * m_offset + ij]);
+        cbs[ij] = &codebook_cross_norms[(codebook_offsets[ij] + code) * ldc];
+    }
+
+    const size_t K8 = (K / (8 * NK)) * (8 * NK);
+
+    for (size_t kk = 0; kk < K8; kk += 8 * NK) {
+        simd_float regs[NK];
+        for (size_t ik = 0; ik < NK; ik++) {
+            regs[ik] = simd_float(cbs[0] + kk + ik * 8);
+        }
+
+        for (size_t ij = 1; ij < M; ij++) {
+            for (size_t ik = 0; ik < NK; ik++) {
+                regs[ik] += simd_float(cbs[ij] + kk + ik * 8);
+            }
+        }
+
+        for (size_t ik = 0; ik < NK; ik++) {
+            regs[ik].storeu(output + kk + ik * 8);
+        }
+    }
+
+    for (size_t kk = K8; kk < K; kk++) {
+        float reg = cbs[0][kk];
+        for (size_t ij = 1; ij < M; ij++) {
+            reg += cbs[ij][kk];
+        }
+        output[kk] = reg;
+    }
+}
+
+template <size_t M, size_t NK, SIMDLevel SL>
+void accum_and_add_tab(
+        const size_t m_offset,
+        const float* const __restrict codebook_cross_norms,
+        const uint64_t* const __restrict codebook_offsets,
+        const int32_t* const __restrict codes_i,
+        const size_t b,
+        const size_t ldc,
+        const size_t K,
+        float* const __restrict output) {
+    using simd_float = simd8float32_tpl<SL>;
+
+    const float* cbs[M];
+    for (size_t ij = 0; ij < M; ij++) {
+        const size_t code = static_cast<size_t>(codes_i[b * m_offset + ij]);
+        cbs[ij] = &codebook_cross_norms[(codebook_offsets[ij] + code) * ldc];
+    }
+
+    const size_t K8 = (K / (8 * NK)) * (8 * NK);
+
+    for (size_t kk = 0; kk < K8; kk += 8 * NK) {
+        simd_float regs[NK];
+        for (size_t ik = 0; ik < NK; ik++) {
+            regs[ik] = simd_float(cbs[0] + kk + ik * 8);
+        }
+
+        for (size_t ij = 1; ij < M; ij++) {
+            for (size_t ik = 0; ik < NK; ik++) {
+                regs[ik] += simd_float(cbs[ij] + kk + ik * 8);
+            }
+        }
+
+        for (size_t ik = 0; ik < NK; ik++) {
+            simd_float existing(output + kk + ik * 8);
+            existing += regs[ik];
+            existing.storeu(output + kk + ik * 8);
+        }
+    }
+
+    for (size_t kk = K8; kk < K; kk++) {
+        float reg = cbs[0][kk];
+        for (size_t ij = 1; ij < M; ij++) {
+            reg += cbs[ij][kk];
+        }
+        output[kk] += reg;
+    }
+}
+
+template <size_t M, size_t NK, SIMDLevel SL>
+void accum_and_finalize_tab(
+        const float* const __restrict codebook_cross_norms,
+        const uint64_t* const __restrict codebook_offsets,
+        const int32_t* const __restrict codes_i,
+        const size_t b,
+        const size_t ldc,
+        const size_t K,
+        const float* const __restrict distances_i,
+        const float* const __restrict cd_common,
+        float* const __restrict output) {
+    using simd_float = simd8float32_tpl<SL>;
+
+    const float* cbs[M];
+    for (size_t ij = 0; ij < M; ij++) {
+        const size_t code = static_cast<size_t>(codes_i[b * M + ij]);
+        cbs[ij] = &codebook_cross_norms[(codebook_offsets[ij] + code) * ldc];
+    }
+
+    const size_t K8 = (K / (8 * NK)) * (8 * NK);
+
+    for (size_t kk = 0; kk < K8; kk += 8 * NK) {
+        simd_float regs[NK];
+        for (size_t ik = 0; ik < NK; ik++) {
+            regs[ik] = simd_float(cbs[0] + kk + ik * 8);
+        }
+
+        for (size_t ij = 1; ij < M; ij++) {
+            for (size_t ik = 0; ik < NK; ik++) {
+                regs[ik] += simd_float(cbs[ij] + kk + ik * 8);
+            }
+        }
+
+        simd_float two(2.0f);
+        for (size_t ik = 0; ik < NK; ik++) {
+            simd_float common_v(cd_common + kk + ik * 8);
+            common_v = fmadd(two, regs[ik], common_v);
+            common_v += simd_float(distances_i[b]);
+            common_v.storeu(output + b * K + kk + ik * 8);
+        }
+    }
+
+    for (size_t kk = K8; kk < K; kk++) {
+        float reg = cbs[0][kk];
+        for (size_t ij = 1; ij < M; ij++) {
+            reg += cbs[ij][kk];
+        }
+        output[b * K + kk] = distances_i[b] + cd_common[kk] + 2 * reg;
+    }
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/approx_topk/rq_beam_search_tab.h

@@ -0,0 +1,117 @@
+/*
+ * 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.
+ */
+
+/// @file rq_beam_search_tab.h
+/// @brief Declarations for SIMDLevel-templatized codebook accumulation
+/// functions.
+///
+/// These functions accumulate codebook cross-norm tables for beam search
+/// encoding in the Residual Quantizer. They compute the distance
+/// contributions from previously encoded codebooks using SIMD-accelerated
+/// register accumulation.
+///
+/// Definitions are in rq_beam_search_tab-inl.h (only included by per-ISA
+/// .cpp files). The common TU only sees these declarations, so no extern
+/// template suppression is needed — the linker resolves to the explicit
+/// instantiations in avx2.cpp / neon.cpp.
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+#include <faiss/utils/simd_levels.h>
+
+namespace faiss {
+
+/// Accumulate cross-norms for M codebooks and store the result.
+///
+/// Loads M codebook rows (selected by codes_i) and sums them using
+/// NK×8-wide SIMD chunks, writing the result to output. Used to
+/// initialize the temporary buffer in the m≥8 path.
+///
+/// @tparam M      number of codebook rows to accumulate
+/// @tparam NK     number of 8-float SIMD chunks per loop iteration
+/// @tparam SL     SIMD level (AVX2, ARM_NEON, etc.)
+/// @param m_offset  stride between beam entries in codes_i
+/// @param codebook_cross_norms  cross-norm table, shape (total_codes, ldc)
+/// @param codebook_offsets      per-codebook offset into cross-norm table
+/// @param codes_i               code indices for the current query
+/// @param b                     beam index
+/// @param ldc                   leading dimension of cross-norm table (≥ K)
+/// @param K                     number of centroids in the current codebook
+/// @param output                output buffer, size K (overwritten)
+template <size_t M, size_t NK, SIMDLevel SL>
+void accum_and_store_tab(
+        size_t m_offset,
+        const float* __restrict codebook_cross_norms,
+        const uint64_t* __restrict codebook_offsets,
+        const int32_t* __restrict codes_i,
+        size_t b,
+        size_t ldc,
+        size_t K,
+        float* __restrict output);
+
+/// Accumulate cross-norms for M codebooks and add to existing output.
+///
+/// Like accum_and_store_tab, but adds the accumulated result to the
+/// existing values in output (output[k] += sum). Used for subsequent
+/// chunks of 8 codebooks in the m≥8 path.
+///
+/// @tparam M      number of codebook rows to accumulate
+/// @tparam NK     number of 8-float SIMD chunks per loop iteration
+/// @tparam SL     SIMD level (AVX2, ARM_NEON, etc.)
+/// @param m_offset  stride between beam entries in codes_i
+/// @param codebook_cross_norms  cross-norm table
+/// @param codebook_offsets      per-codebook offset
+/// @param codes_i               code indices
+/// @param b                     beam index
+/// @param ldc                   leading dimension of cross-norm table
+/// @param K                     number of centroids
+/// @param output                output buffer, size K (accumulated into)
+template <size_t M, size_t NK, SIMDLevel SL>
+void accum_and_add_tab(
+        size_t m_offset,
+        const float* __restrict codebook_cross_norms,
+        const uint64_t* __restrict codebook_offsets,
+        const int32_t* __restrict codes_i,
+        size_t b,
+        size_t ldc,
+        size_t K,
+        float* __restrict output);
+
+/// Accumulate cross-norms for M codebooks and finalize distances.
+///
+/// Accumulates M codebook rows, then computes the final centroid distance:
+///   output[b*K + k] = distances_i[b] + cd_common[k] + 2 * sum[k]
+/// Used for m=1..7 where the entire accumulation fits in registers.
+///
+/// @tparam M      number of codebook rows to accumulate (equals m)
+/// @tparam NK     number of 8-float SIMD chunks per loop iteration
+/// @tparam SL     SIMD level (AVX2, ARM_NEON, etc.)
+/// @param codebook_cross_norms  cross-norm table
+/// @param codebook_offsets      per-codebook offset
+/// @param codes_i               code indices (stride is M)
+/// @param b                     beam index
+/// @param ldc                   leading dimension of cross-norm table
+/// @param K                     number of centroids
+/// @param distances_i           per-beam input distances, size beam_size
+/// @param cd_common             common distance term, size K
+/// @param output                output centroid distances (b*K offset)
+template <size_t M, size_t NK, SIMDLevel SL>
+void accum_and_finalize_tab(
+        const float* __restrict codebook_cross_norms,
+        const uint64_t* __restrict codebook_offsets,
+        const int32_t* __restrict codes_i,
+        size_t b,
+        size_t ldc,
+        size_t K,
+        const float* __restrict distances_i,
+        const float* __restrict cd_common,
+        float* __restrict output);
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/approx_topk/simdlib256-inl.h

@@ -0,0 +1,146 @@
+/*
+ * 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.
+ */
+
+// Out-of-line definition of HeapWithBucketsCMaxFloat::bs_addn using
+// simdlib types. Only included by per-ISA .cpp files (avx2.cpp, neon.cpp).
+// Do NOT include this from common translation units.
+
+#pragma once
+
+#include <cstdint>
+#include <limits>
+
+#include <faiss/impl/approx_topk/approx_topk.h>
+#include <faiss/impl/simdlib/simdlib.h>
+#include <faiss/utils/Heap.h>
+#include <faiss/utils/simd_levels.h>
+
+namespace faiss {
+
+// Element-wise max of two simd8float32 vectors, implemented via
+// cmplt_min_max_fast (which computes both min and max).
+template <SIMDLevel SL>
+inline simd8float32_tpl<SL> simd8float32_max(
+        simd8float32_tpl<SL> a,
+        simd8float32_tpl<SL> b) {
+    simd8float32_tpl<SL> min_val, max_val;
+    simd8uint32_tpl<SL> dummy(0u), dmin, dmax;
+    cmplt_min_max_fast(a, dummy, b, dummy, min_val, dmin, max_val, dmax);
+    return max_val;
+}
+
+template <uint32_t NBUCKETS, uint32_t N, SIMDLevel SL>
+void HeapWithBucketsCMaxFloat<NBUCKETS, N, SL>::bs_addn(
+        const uint32_t beam_size,
+        const uint32_t n_per_beam,
+        const float* const __restrict distances,
+        const uint32_t k,
+        float* const __restrict bh_val,
+        int32_t* const __restrict bh_ids) {
+    using C = CMax<float, int>;
+    using simd_float = simd8float32_tpl<SL>;
+    using simd_uint = simd8uint32_tpl<SL>;
+
+    for (uint32_t beam_index = 0; beam_index < beam_size; beam_index++) {
+        simd_float min_distances_i[NBUCKETS / 8][N];
+        simd_uint min_indices_i[NBUCKETS / 8][N];
+
+        for (uint32_t j = 0; j < NBUCKETS / 8; j++) {
+            for (uint32_t p = 0; p < N; p++) {
+                min_distances_i[j][p] =
+                        simd_float(std::numeric_limits<float>::max());
+                min_indices_i[j][p] = simd_uint(0u, 1u, 2u, 3u, 4u, 5u, 6u, 7u);
+            }
+        }
+
+        simd_uint current_indices(0u, 1u, 2u, 3u, 4u, 5u, 6u, 7u);
+        simd_uint indices_delta(NBUCKETS);
+
+        const uint32_t nb = (n_per_beam / NBUCKETS) * NBUCKETS;
+
+        // put the data into buckets
+        for (uint32_t ip = 0; ip < nb; ip += NBUCKETS) {
+            for (uint32_t j = 0; j < NBUCKETS / 8; j++) {
+                const simd_float distances_reg(
+                        distances + j * 8 + ip + n_per_beam * beam_index);
+
+                simd_float distance_candidate = distances_reg;
+                simd_uint indices_candidate = current_indices;
+
+                for (uint32_t p = 0; p < N; p++) {
+                    // Use cmplt_min_max_fast for comparison, min values,
+                    // min indices, and max indices.
+                    simd_float min_d_new, max_d_unused;
+                    simd_uint min_idx_new, max_idx_new;
+                    cmplt_min_max_fast(
+                            distance_candidate,
+                            indices_candidate,
+                            min_distances_i[j][p],
+                            min_indices_i[j][p],
+                            min_d_new,
+                            min_idx_new,
+                            max_d_unused,
+                            max_idx_new);
+
+                    // The max distance uses distances_reg (the original
+                    // input), NOT distance_candidate. This is a deliberate
+                    // approximation that breaks the data dependency chain.
+                    simd_float max_d_new = simd8float32_max<SL>(
+                            min_distances_i[j][p], distances_reg);
+
+                    distance_candidate = max_d_new;
+                    indices_candidate = max_idx_new;
+
+                    min_distances_i[j][p] = min_d_new;
+                    min_indices_i[j][p] = min_idx_new;
+                }
+            }
+
+            current_indices = current_indices + indices_delta;
+        }
+
+        // fix the indices
+        for (uint32_t j = 0; j < NBUCKETS / 8; j++) {
+            const simd_uint offset(n_per_beam * beam_index + j * 8);
+            for (uint32_t p = 0; p < N; p++) {
+                min_indices_i[j][p] = min_indices_i[j][p] + offset;
+            }
+        }
+
+        // merge every bucket into the regular heap
+        for (uint32_t p = 0; p < N; p++) {
+            for (uint32_t j = 0; j < NBUCKETS / 8; j++) {
+                uint32_t min_indices_scalar[8];
+                float min_distances_scalar[8];
+
+                min_indices_i[j][p].storeu(min_indices_scalar);
+                min_distances_i[j][p].storeu(min_distances_scalar);
+
+                for (size_t j8 = 0; j8 < 8; j8++) {
+                    const auto value = min_distances_scalar[j8];
+                    const auto index =
+                            static_cast<int32_t>(min_indices_scalar[j8]);
+                    if (C::cmp2(bh_val[0], value, bh_ids[0], index)) {
+                        heap_replace_top<C>(k, bh_val, bh_ids, value, index);
+                    }
+                }
+            }
+        }
+
+        // process leftovers
+        for (uint32_t ip = nb; ip < n_per_beam; ip++) {
+            const int32_t index = ip + n_per_beam * beam_index;
+            const float value = distances[index];
+
+            if (C::cmp(bh_val[0], value)) {
+                heap_replace_top<C>(k, bh_val, bh_ids, value, index);
+            }
+        }
+    }
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/binary_hamming/IndexBinaryHNSW_impl.h

@@ -0,0 +1,73 @@
+/*
+ * 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.
+ */
+
+/*
+ * Per-ISA implementation of Hamming distance computation for
+ * IndexBinaryHNSW. Included once per SIMD TU with THE_SIMD_LEVEL
+ * set to the desired SIMDLevel.
+ */
+
+#pragma once
+
+#ifndef THE_SIMD_LEVEL
+#error "THE_SIMD_LEVEL must be defined before including this file"
+#endif
+
+// The including TU (or the per-ISA hamming_computer-*.h it pulls in first)
+// is responsible for providing the HammingComputer*_tpl<SL> specializations;
+// this header only needs the forward declarations and with_HammingComputer<SL>
+// dispatcher from hamming_computer.h.
+#include <faiss/utils/hamming_distance/hamming_computer.h>
+
+#include <faiss/IndexBinaryFlat.h>
+#include <faiss/impl/DistanceComputer.h>
+#include <faiss/impl/binary_hamming/dispatch.h>
+#include <faiss/utils/hamming.h>
+
+namespace faiss {
+
+namespace {
+
+template <class HammingComputer>
+struct FlatHammingDis : DistanceComputer {
+    const int code_size;
+    const uint8_t* b;
+    HammingComputer hc;
+
+    float operator()(idx_t i) override {
+        return hc.hamming(b + i * code_size);
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) override {
+        return HammingComputerDefault_tpl<THE_SIMD_LEVEL>(
+                       b + j * code_size, code_size)
+                .hamming(b + i * code_size);
+    }
+
+    explicit FlatHammingDis(const IndexBinaryFlat& storage)
+            : code_size(storage.code_size), b(storage.xb.data()), hc() {}
+
+    // NOTE: Pointers are cast from float in order to reuse the floating-point
+    //   DistanceComputer.
+    void set_query(const float* x) override {
+        hc.set((uint8_t*)x, code_size);
+    }
+};
+
+} // anonymous namespace
+
+template <>
+DistanceComputer* make_binary_hnsw_distance_computer_fixSL<THE_SIMD_LEVEL>(
+        int code_size,
+        IndexBinaryFlat* flat_storage) {
+    return with_HammingComputer<THE_SIMD_LEVEL>(
+            code_size, [&]<class HammingComputer>() -> DistanceComputer* {
+                return new FlatHammingDis<HammingComputer>(*flat_storage);
+            });
+}
+
+} // namespace faiss