faiss 0.6.0 → 0.6.1

data/vendor/faiss/faiss/utils/hamming.h

@@ -30,11 +30,15 @@
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/platform_macros.h>
 #include <faiss/utils/Heap.h>
+#include <faiss/utils/popcount.h>
+#include <faiss/utils/simd_levels.h>
 
 // Low-level Hamming distance computations and hamdis_t.
-#include <faiss/utils/hamming_distance/hamdis-inl.h>
+#include <faiss/utils/hamming_distance/hamming_computer.h>
+// Scalar (NONE) HammingComputer struct specializations.
+#include <faiss/utils/hamming_distance/hamming_computer-generic.h>
 
-#include <faiss/utils/approx_topk/mode.h>
+#include <faiss/impl/approx_topk/approx_topk.h>
 
 namespace faiss {
 
@@ -284,6 +288,92 @@ void unpack_bitstrings(
         size_t code_size,
         int32_t* unpacked);
 
+/** Per-SIMD-level Hamming function implementations.
+ *  Specializations live in per-ISA TUs (hamming_avx2.cpp, etc.).
+ *  The `_fixSL` suffix marks these as pinned-to-a-SIMDLevel (no dispatch),
+ *  distinguishing them from the non-templated DD entry points above. */
+
+template <SIMDLevel SL>
+void hammings_knn_hc_fixSL(
+        int_maxheap_array_t* ha,
+        const uint8_t* a,
+        const uint8_t* b,
+        size_t nb,
+        size_t ncodes,
+        int ordered,
+        ApproxTopK_mode_t approx_topk_mode,
+        const IDSelector* sel);
+
+template <SIMDLevel SL>
+void hammings_knn_mc_fixSL(
+        const uint8_t* a,
+        const uint8_t* b,
+        size_t na,
+        size_t nb,
+        size_t k,
+        size_t ncodes,
+        int32_t* distances,
+        int64_t* labels,
+        const IDSelector* sel);
+
+template <SIMDLevel SL>
+void hamming_range_search_fixSL(
+        const uint8_t* a,
+        const uint8_t* b,
+        size_t na,
+        size_t nb,
+        int radius,
+        size_t code_size,
+        RangeSearchResult* result,
+        const IDSelector* sel);
+
+template <SIMDLevel SL>
+void hammings_fixSL(
+        const uint8_t* a,
+        const uint8_t* b,
+        size_t na,
+        size_t nb,
+        size_t ncodes,
+        hamdis_t* dis);
+
+template <SIMDLevel SL>
+void generalized_hammings_knn_hc_fixSL(
+        int_maxheap_array_t* ha,
+        const uint8_t* a,
+        const uint8_t* b,
+        size_t nb,
+        size_t code_size,
+        int ordered);
+
+template <SIMDLevel SL>
+void hamming_count_thres_fixSL(
+        const uint8_t* bs1,
+        const uint8_t* bs2,
+        size_t n1,
+        size_t n2,
+        hamdis_t ht,
+        size_t ncodes,
+        size_t* nptr);
+
+template <SIMDLevel SL>
+void crosshamming_count_thres_fixSL(
+        const uint8_t* dbs,
+        size_t n,
+        hamdis_t ht,
+        size_t ncodes,
+        size_t* nptr);
+
+template <SIMDLevel SL>
+size_t match_hamming_thres_fixSL(
+        const uint8_t* bs1,
+        const uint8_t* bs2,
+        size_t n1,
+        size_t n2,
+        hamdis_t ht,
+        size_t ncodes,
+        int64_t* idx,
+        hamdis_t* dis);
+
 } // namespace faiss
 
 #include <faiss/utils/hamming-inl.h>

data/vendor/faiss/faiss/utils/hamming_distance/common.h

@@ -8,25 +8,25 @@
 #ifndef FAISS_hamming_common_h
 #define FAISS_hamming_common_h
 
+#include <cassert>
+#include <cstddef>
 #include <cstdint>
 
 #include <faiss/impl/platform_macros.h>
+#include <faiss/utils/popcount.h>
+#include <faiss/utils/simd_levels.h>
+
+#ifdef __aarch64__
+// Include <arm_neon.h> at global scope so the NEON types it declares
+// (uint8x16_t, uint16x8_t, ...) end up in `::` and not inside `faiss::`.
+#include <arm_neon.h>
+#endif
 
 /* The Hamming distance type */
 using hamdis_t = int32_t;
 
 namespace faiss {
 
-// trust the compiler to provide efficient popcount implementations
-inline int popcount32(uint32_t x) {
-    return __builtin_popcount(x);
-}
-
-// popcount
-inline int popcount64(uint64_t x) {
-    return __builtin_popcountl(x);
-}
-
 // This table was moved from .cpp to .h file, because
 // otherwise it was causing compilation errors while trying to
 // compile swig modules on Windows.
@@ -43,6 +43,283 @@ inline constexpr uint8_t hamdis_tab_ham_bytes[256] = {
         4, 5, 5, 6, 5, 6, 6, 7, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
         4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8};
 
+/***************************************************************************
+ * generalized Hamming = number of bytes that are different between
+ * two codes.
+ ***************************************************************************/
+
+inline int generalized_hamming_64(uint64_t a) {
+    a |= a >> 1;
+    a |= a >> 2;
+    a |= a >> 4;
+    a &= 0x0101010101010101UL;
+    return popcount64(a);
+}
+
+/***************************************************************************
+ * Free-function hamming<nbits>() templates.
+ * hamming<64> is architecture-independent.
+ * hamming<128>, hamming<256>, the primary template, and the runtime-nwords
+ * overload have NEON-optimized versions on aarch64.
+ ***************************************************************************/
+
+#ifndef SWIG
+
+/* hamming<64> — identical on all architectures */
+template <size_t nbits>
+inline hamdis_t hamming(const uint64_t* bs1, const uint64_t* bs2);
+
+template <>
+inline hamdis_t hamming<64>(const uint64_t* pa, const uint64_t* pb) {
+    return popcount64(pa[0] ^ pb[0]);
+}
+
+#ifdef __aarch64__
+
+/* Hamming distances for multiples of 64 bits — NEON version */
+template <size_t nbits>
+inline hamdis_t hamming(const uint64_t* pa, const uint64_t* pb) {
+    constexpr size_t nwords256 = nbits / 256;
+    constexpr size_t nwords128 = (nbits - nwords256 * 256) / 128;
+    constexpr size_t nwords64 =
+            (nbits - nwords256 * 256 - nwords128 * 128) / 64;
+
+    hamdis_t h = 0;
+    if (nwords256 > 0) {
+        for (size_t i = 0; i < nwords256; i++) {
+            h += hamming<256>(pa, pb);
+            pa += 4;
+            pb += 4;
+        }
+    }
+
+    if (nwords128 > 0) {
+        h += hamming<128>(pa, pb);
+        pa += 2;
+        pb += 2;
+    }
+
+    if (nwords64 > 0) {
+        h += hamming<64>(pa, pb);
+    }
+
+    return h;
+}
+
+template <>
+inline hamdis_t hamming<128>(const uint64_t* pa, const uint64_t* pb) {
+    const uint8_t* pa8 = reinterpret_cast<const uint8_t*>(pa);
+    const uint8_t* pb8 = reinterpret_cast<const uint8_t*>(pb);
+    uint8x16_t or0 = veorq_u8(vld1q_u8(pa8), vld1q_u8(pb8));
+    uint8x16_t c0 = vcntq_u8(or0);
+    auto dis = vaddvq_u8(c0);
+    return dis;
+}
+
+template <>
+inline hamdis_t hamming<256>(const uint64_t* pa, const uint64_t* pb) {
+    const uint8_t* pa8 = reinterpret_cast<const uint8_t*>(pa);
+    const uint8_t* pb8 = reinterpret_cast<const uint8_t*>(pb);
+    uint8x16_t or0 = veorq_u8(vld1q_u8(pa8), vld1q_u8(pb8));
+    uint8x16_t or1 = veorq_u8(vld1q_u8(pa8 + 16), vld1q_u8(pb8 + 16));
+    uint8x16_t c0 = vcntq_u8(or0);
+    uint8x16_t c1 = vcntq_u8(or1);
+    uint8x16_t ca = vpaddq_u8(c0, c1);
+    auto dis = vaddvq_u8(ca);
+    return dis;
+}
+
+/* Hamming distances for multiple of 64 bits */
+inline hamdis_t hamming(const uint64_t* pa, const uint64_t* pb, size_t nwords) {
+    const size_t nwords256 = nwords / 4;
+    const size_t nwords128 = (nwords % 4) / 2;
+    const size_t nwords64 = nwords % 2;
+
+    hamdis_t h = 0;
+    if (nwords256 > 0) {
+        for (size_t i = 0; i < nwords256; i++) {
+            h += hamming<256>(pa, pb);
+            pa += 4;
+            pb += 4;
+        }
+    }
+
+    if (nwords128 > 0) {
+        h += hamming<128>(pa, pb);
+        pa += 2;
+        pb += 2;
+    }
+
+    if (nwords64 > 0) {
+        h += hamming<64>(pa, pb);
+    }
+
+    return h;
+}
+
+#else // !__aarch64__
+
+/* Hamming distances for multiples of 64 bits — scalar version */
+template <size_t nbits>
+inline hamdis_t hamming(const uint64_t* bs1, const uint64_t* bs2) {
+    const size_t nwords = nbits / 64;
+    size_t i;
+    hamdis_t h = 0;
+    for (i = 0; i < nwords; i++) {
+        h += popcount64(bs1[i] ^ bs2[i]);
+    }
+    return h;
+}
+
+/* specialized (optimized) functions */
+template <>
+inline hamdis_t hamming<128>(const uint64_t* pa, const uint64_t* pb) {
+    return popcount64(pa[0] ^ pb[0]) + popcount64(pa[1] ^ pb[1]);
+}
+
+template <>
+inline hamdis_t hamming<256>(const uint64_t* pa, const uint64_t* pb) {
+    return popcount64(pa[0] ^ pb[0]) + popcount64(pa[1] ^ pb[1]) +
+            popcount64(pa[2] ^ pb[2]) + popcount64(pa[3] ^ pb[3]);
+}
+
+/* Hamming distances for multiple of 64 bits */
+inline hamdis_t hamming(
+        const uint64_t* bs1,
+        const uint64_t* bs2,
+        size_t nwords) {
+    hamdis_t h = 0;
+    for (size_t i = 0; i < nwords; i++) {
+        h += popcount64(bs1[i] ^ bs2[i]);
+    }
+    return h;
+}
+
+#endif // __aarch64__
+
+/***************************************************************************
+ * Bit-level Hamming distance implementation functions.
+ * These depend only on the hamming<nbits>() free functions above.
+ ***************************************************************************/
+
+template <size_t nbits>
+inline void hammings_impl(
+        const uint64_t* __restrict bs1,
+        const uint64_t* __restrict bs2,
+        size_t n1,
+        size_t n2,
+        hamdis_t* __restrict dis) {
+    size_t i, j;
+    const size_t nwords = nbits / 64;
+    for (i = 0; i < n1; i++) {
+        const uint64_t* __restrict bs1_ = bs1 + i * nwords;
+        hamdis_t* __restrict dis_ = dis + i * n2;
+        for (j = 0; j < n2; j++) {
+            dis_[j] = hamming<nbits>(bs1_, bs2 + j * nwords);
+        }
+    }
+}
+
+inline void hammings_impl_runtime(
+        const uint64_t* __restrict bs1,
+        const uint64_t* __restrict bs2,
+        size_t n1,
+        size_t n2,
+        size_t nbits,
+        hamdis_t* __restrict dis) {
+    size_t i, j;
+    const size_t nwords = nbits / 64;
+    for (i = 0; i < n1; i++) {
+        const uint64_t* __restrict bs1_ = bs1 + i * nwords;
+        hamdis_t* __restrict dis_ = dis + i * n2;
+        for (j = 0; j < n2; j++) {
+            dis_[j] = hamming(bs1_, bs2 + j * nwords, nwords);
+        }
+    }
+}
+
+template <size_t nbits>
+inline void hamming_count_thres_impl(
+        const uint64_t* __restrict bs1,
+        const uint64_t* __restrict bs2,
+        size_t n1,
+        size_t n2,
+        hamdis_t ht,
+        size_t* __restrict nptr) {
+    const size_t nwords = nbits / 64;
+    size_t i, j, posm = 0;
+    const uint64_t* bs2_ = bs2;
+
+    for (i = 0; i < n1; i++) {
+        bs2 = bs2_;
+        for (j = 0; j < n2; j++) {
+            if (hamming<nbits>(bs1, bs2) <= ht) {
+                posm++;
+            }
+            bs2 += nwords;
+        }
+        bs1 += nwords;
+    }
+    *nptr = posm;
+}
+
+template <size_t nbits>
+inline void crosshamming_count_thres_impl(
+        const uint64_t* __restrict dbs,
+        size_t n,
+        int ht,
+        size_t* __restrict nptr) {
+    const size_t nwords = nbits / 64;
+    size_t i, j, posm = 0;
+    const uint64_t* bs1 = dbs;
+    for (i = 0; i < n; i++) {
+        const uint64_t* bs2 = bs1 + nwords;
+        for (j = i + 1; j < n; j++) {
+            if (hamming<nbits>(bs1, bs2) <= ht) {
+                posm++;
+            }
+            bs2 += nwords;
+        }
+        bs1 += nwords;
+    }
+    *nptr = posm;
+}
+
+template <size_t nbits>
+inline size_t match_hamming_thres_impl(
+        const uint64_t* __restrict bs1,
+        const uint64_t* __restrict bs2,
+        size_t n1,
+        size_t n2,
+        int ht,
+        int64_t* __restrict idx,
+        hamdis_t* __restrict hams) {
+    const size_t nwords = nbits / 64;
+    size_t i, j, posm = 0;
+    hamdis_t h;
+    const uint64_t* bs2_ = bs2;
+    for (i = 0; i < n1; i++) {
+        bs2 = bs2_;
+        for (j = 0; j < n2; j++) {
+            h = hamming<nbits>(bs1, bs2);
+            if (h <= ht) {
+                *idx = i;
+                idx++;
+                *idx = j;
+                idx++;
+                *hams = h;
+                hams++;
+                posm++;
+            }
+            bs2 += nwords;
+        }
+        bs1 += nwords;
+    }
+    return posm;
+}
+
+#endif // SWIG
+
 } // namespace faiss
 
 #endif

data/vendor/faiss/faiss/utils/hamming_distance/hamming_avx2.cpp

@@ -0,0 +1,15 @@
+/*
+ * 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
+
+#define THE_SIMD_LEVEL SIMDLevel::AVX2
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/utils/hamming_distance/hamming_computer-avx2.h>
+#include <faiss/utils/hamming_distance/hamming_impl.h>
+
+#endif // COMPILE_SIMD_AVX2

data/vendor/faiss/faiss/utils/hamming_distance/hamming_avx512.cpp

@@ -0,0 +1,15 @@
+/*
+ * 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_AVX512
+
+#define THE_SIMD_LEVEL SIMDLevel::AVX512
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/utils/hamming_distance/hamming_computer-avx512.h>
+#include <faiss/utils/hamming_distance/hamming_impl.h>
+
+#endif // COMPILE_SIMD_AVX512

data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-avx2.h

@@ -0,0 +1,142 @@
+/*
+ * 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.
+ */
+
+#ifndef HAMMING_COMPUTER_AVX2_H
+#define HAMMING_COMPUTER_AVX2_H
+
+// AVX2 HammingComputer and GenHammingComputer specializations.
+// Types without custom AVX2 code inherit from the NONE specializations
+// in hamming_computer-generic.h. Custom AVX2 specializations for
+// GenHammingComputer classes leverage SSE/AVX2 intrinsics.
+
+#include <cassert>
+#include <cstdint>
+
+#include <faiss/impl/platform_macros.h>
+#include <faiss/utils/hamming_distance/hamming_computer-generic.h>
+
+#include <immintrin.h>
+
+namespace faiss {
+
+/***************************************************************************
+ * AVX2 inheriting specializations for types without custom AVX2 code.
+ * These explicitly inherit the scalar (NONE) implementation so that
+ * every SIMDLevel has a concrete specialization.
+ ***************************************************************************/
+
+#define FAISS_INHERIT_HAMMING(Class)                                     \
+    template <>                                                          \
+    struct Class##_tpl<SIMDLevel::AVX2> : Class##_tpl<SIMDLevel::NONE> { \
+        using Class##_tpl<SIMDLevel::NONE>::Class##_tpl;                 \
+    }
+
+FAISS_INHERIT_HAMMING(HammingComputer16);
+FAISS_INHERIT_HAMMING(HammingComputer20);
+FAISS_INHERIT_HAMMING(HammingComputer32);
+FAISS_INHERIT_HAMMING(HammingComputer64);
+FAISS_INHERIT_HAMMING(HammingComputerDefault);
+FAISS_INHERIT_HAMMING(GenHammingComputer8);
+
+#undef FAISS_INHERIT_HAMMING
+
+/***************************************************************************
+ * Custom AVX2 specializations for GenHammingComputer classes.
+ ***************************************************************************/
+
+// I'm not sure whether this version is faster of slower, tbh
+// todo: test on different CPUs
+template <>
+struct GenHammingComputer16_tpl<SIMDLevel::AVX2> {
+    __m128i a;
+
+    GenHammingComputer16_tpl(
+            const uint8_t* a8,
+            FAISS_MAYBE_UNUSED int code_size) {
+        assert(code_size == 16);
+        a = _mm_loadu_si128((const __m128i_u*)a8);
+    }
+
+    inline int hamming(const uint8_t* b8) const {
+        const __m128i b = _mm_loadu_si128((const __m128i_u*)b8);
+        const __m128i cmp = _mm_cmpeq_epi8(a, b);
+        const auto movemask = _mm_movemask_epi8(cmp);
+        return 16 - popcount32(movemask);
+    }
+
+    inline static constexpr int get_code_size() {
+        return 16;
+    }
+};
+
+template <>
+struct GenHammingComputer32_tpl<SIMDLevel::AVX2> {
+    __m256i a;
+
+    GenHammingComputer32_tpl(
+            const uint8_t* a8,
+            FAISS_MAYBE_UNUSED int code_size) {
+        assert(code_size == 32);
+        a = _mm256_loadu_si256((const __m256i_u*)a8);
+    }
+
+    inline int hamming(const uint8_t* b8) const {
+        const __m256i b = _mm256_loadu_si256((const __m256i_u*)b8);
+        const __m256i cmp = _mm256_cmpeq_epi8(a, b);
+        const uint32_t movemask = _mm256_movemask_epi8(cmp);
+        return 32 - popcount32(movemask);
+    }
+
+    inline static constexpr int get_code_size() {
+        return 32;
+    }
+};
+
+// A specialized version might be needed for the very long
+// GenHamming code_size. In such a case, one may accumulate
+// counts using _mm256_sub_epi8 and then compute a horizontal
+// sum (using _mm256_sad_epu8, maybe, in blocks of no larger
+// than 256 * 32 bytes).
+
+template <>
+struct GenHammingComputerM8_tpl<SIMDLevel::AVX2> {
+    const uint64_t* a;
+    int n;
+
+    GenHammingComputerM8_tpl(const uint8_t* a8, int code_size) {
+        assert(code_size % 8 == 0);
+        a = (uint64_t*)a8;
+        n = code_size / 8;
+    }
+
+    int hamming(const uint8_t* b8) const {
+        const uint64_t* b = (uint64_t*)b8;
+        int accu = 0;
+
+        int i = 0;
+        int n4 = (n / 4) * 4;
+        for (; i < n4; i += 4) {
+            const __m256i av = _mm256_loadu_si256((const __m256i_u*)(a + i));
+            const __m256i bv = _mm256_loadu_si256((const __m256i_u*)(b + i));
+            const __m256i cmp = _mm256_cmpeq_epi8(av, bv);
+            const uint32_t movemask = _mm256_movemask_epi8(cmp);
+            accu += 32 - popcount32(movemask);
+        }
+
+        for (; i < n; i++)
+            accu += generalized_hamming_64(a[i] ^ b[i]);
+        return accu;
+    }
+
+    inline int get_code_size() const {
+        return n * 8;
+    }
+};
+
+} // namespace faiss
+
+#endif