faiss 0.6.4 → 0.6.5

data/vendor/faiss/faiss/utils/simd_impl/exhaustive_L2sqr_blas_cmax.h

@@ -7,11 +7,93 @@
 
 #pragma once
 
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
 #include <faiss/impl/ResultHandler.h>
+#include <faiss/utils/distances.h>
 #include <faiss/utils/simd_levels.h>
+#include <faiss/utils/utils.h>
 
 namespace faiss {
 
+/** fp32 query tiles for the BLAS exhaustive-search kernels.
+ *
+ * A kernel calls prepare_norms() once, then tile(i0, i1, x_norms) for each
+ * query tile; the returned rows stay valid until the next tile() call.
+ */
+struct Fp32QueryTiles {
+    const float* x;
+    size_t d;
+
+    /// squared L2 norms of all nx queries
+    void prepare_norms(size_t nx, float* x_norms) const {
+        fvec_norms_L2sqr(x_norms, x, d, nx);
+    }
+
+    const float* tile(size_t i0, size_t /*i1*/, float* /*x_norms*/) const {
+        return x + i0 * d;
+    }
+};
+
+/** Packed IEEE binary16 query tiles for the BLAS exhaustive-search kernels.
+ *
+ * Queries are widened in chunks of whole tiles (at most chunk_bytes of fp32,
+ * at least one tile) into a buffer reused across chunks: the fp32 copy stays
+ * bounded and each widened tile serves all database tiles. Norms are computed
+ * from the widened rows, so results match the fp32 kernels on the widened
+ * queries.
+ */
+struct Fp16QueryTiles {
+    static constexpr size_t chunk_bytes = size_t(16) << 20;
+
+    const uint16_t* x;
+    size_t d;
+    size_t nx;
+    std::unique_ptr<float[]> buffer;
+    size_t chunk_rows = 0;
+    size_t chunk_begin = 0;
+    size_t chunk_end = 0;
+
+    Fp16QueryTiles(const uint16_t* x_in, size_t d_in, size_t nx_in)
+            : x(x_in), d(d_in), nx(nx_in) {}
+
+    void prepare_norms(size_t /*nx*/, float* /*x_norms*/) const {}
+
+    /** widened rows [i0, i1), and their norms if x_norms is not null
+     *
+     * Tiles are requested in order, all of the same size except the last.
+     */
+    const float* tile(size_t i0, size_t i1, float* x_norms) {
+        if (i1 > chunk_end) {
+            if (!buffer) {
+                const size_t tile_rows = i1 - i0;
+                chunk_rows =
+                        tile_rows *
+                        std::max<size_t>(
+                                1,
+                                chunk_bytes / (tile_rows * d * sizeof(float)));
+                chunk_rows = std::min(chunk_rows, nx - i0);
+                buffer.reset(new float[chunk_rows * d]);
+            }
+            chunk_begin = i0;
+            chunk_end = std::min(nx, i0 + chunk_rows);
+            const int64_t n = static_cast<int64_t>(chunk_end - chunk_begin);
+#pragma omp parallel for if (n * static_cast<int64_t>(d) > (1 << 16))
+            for (int64_t i = 0; i < n; i++) {
+                float* row = buffer.get() + i * d;
+                fp16_to_fp32(d, x + (chunk_begin + i) * d, row);
+                if (x_norms) {
+                    x_norms[chunk_begin + i] = fvec_norm_L2sqr(row, d);
+                }
+            }
+        }
+        return buffer.get() + (i0 - chunk_begin) * d;
+    }
+};
+
 /// BLAS-accelerated exhaustive L2 search for the k=1 (top-1) case.
 /// Specializations live in the per-SIMD translation units under simd_impl/.
 template <SIMDLevel>
@@ -24,4 +106,15 @@ void exhaustive_L2sqr_blas_cmax(
         Top1BlockResultHandler<CMax<float, int64_t>>& res,
         const float* y_norms);
 
+/// Same as exhaustive_L2sqr_blas_cmax for packed IEEE binary16 queries.
+template <SIMDLevel>
+void exhaustive_L2sqr_blas_cmax_fp16(
+        const uint16_t* x,
+        const float* y,
+        size_t d,
+        size_t nx,
+        size_t ny,
+        Top1BlockResultHandler<CMax<float, int64_t>>& res,
+        const float* y_norms);
+
 } // namespace faiss

data/vendor/faiss/faiss/utils/simd_impl/fp16_kernels.h

@@ -0,0 +1,76 @@
+/*
+ * 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 <cstddef>
+#include <cstdint>
+
+#include <faiss/utils/fp16.h>
+#include <faiss/utils/simd_levels.h>
+
+namespace faiss {
+namespace detail {
+
+// c[j] += fp32(x[j]) * w for j < d
+template <SIMDLevel Level>
+inline void fp16_madd(size_t d, const uint16_t* x, float w, float* c) {
+    for (size_t j = 0; j < d; ++j) {
+        c[j] += decode_fp16(x[j]) * w;
+    }
+}
+
+template <SIMDLevel Level>
+inline void fp16_to_fp32_kernel(size_t n, const uint16_t* x, float* out) {
+    for (size_t j = 0; j < n; ++j) {
+        out[j] = decode_fp16(x[j]);
+    }
+}
+
+#ifdef COMPILE_SIMD_AVX2
+template <>
+void fp16_madd<SIMDLevel::AVX2>(size_t d, const uint16_t* x, float w, float* c);
+
+template <>
+void fp16_to_fp32_kernel<SIMDLevel::AVX2>(
+        size_t n,
+        const uint16_t* x,
+        float* out);
+#endif
+
+#ifdef COMPILE_SIMD_AVX512
+template <>
+void fp16_madd<SIMDLevel::AVX512>(
+        size_t d,
+        const uint16_t* x,
+        float w,
+        float* c);
+
+template <>
+void fp16_to_fp32_kernel<SIMDLevel::AVX512>(
+        size_t n,
+        const uint16_t* x,
+        float* out);
+#endif
+
+#ifdef COMPILE_SIMD_ARM_NEON
+template <>
+void fp16_madd<SIMDLevel::ARM_NEON>(
+        size_t d,
+        const uint16_t* x,
+        float w,
+        float* c);
+
+template <>
+void fp16_to_fp32_kernel<SIMDLevel::ARM_NEON>(
+        size_t n,
+        const uint16_t* x,
+        float* out);
+#endif
+
+} // namespace detail
+} // namespace faiss

data/vendor/faiss/faiss/utils/simd_impl/fp16_kernels_avx2.cpp

@@ -0,0 +1,52 @@
+/*
+ * 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.
+ */
+
+#include <faiss/utils/simd_impl/fp16_kernels.h>
+
+#include <immintrin.h>
+
+namespace faiss {
+namespace detail {
+
+template <>
+void fp16_madd<SIMDLevel::AVX2>(
+        size_t d,
+        const uint16_t* x,
+        float w,
+        float* c) {
+    const __m256 wv = _mm256_set1_ps(w);
+    size_t j = 0;
+    for (; j + 8 <= d; j += 8) {
+        const __m256 xv = _mm256_cvtph_ps(
+                _mm_loadu_si128(reinterpret_cast<const __m128i*>(x + j)));
+        _mm256_storeu_ps(
+                c + j,
+                _mm256_add_ps(_mm256_loadu_ps(c + j), _mm256_mul_ps(xv, wv)));
+    }
+    for (; j < d; ++j) {
+        c[j] += decode_fp16(x[j]) * w;
+    }
+}
+
+template <>
+void fp16_to_fp32_kernel<SIMDLevel::AVX2>(
+        size_t n,
+        const uint16_t* x,
+        float* out) {
+    size_t j = 0;
+    for (; j + 8 <= n; j += 8) {
+        const __m128i h =
+                _mm_loadu_si128(reinterpret_cast<const __m128i*>(x + j));
+        _mm256_storeu_ps(out + j, _mm256_cvtph_ps(h));
+    }
+    for (; j < n; ++j) {
+        out[j] = decode_fp16(x[j]);
+    }
+}
+
+} // namespace detail
+} // namespace faiss

data/vendor/faiss/faiss/utils/simd_impl/fp16_kernels_avx512.cpp

@@ -0,0 +1,52 @@
+/*
+ * 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.
+ */
+
+#include <faiss/utils/simd_impl/fp16_kernels.h>
+
+#include <immintrin.h>
+
+namespace faiss {
+namespace detail {
+
+template <>
+void fp16_madd<SIMDLevel::AVX512>(
+        size_t d,
+        const uint16_t* x,
+        float w,
+        float* c) {
+    const __m512 wv = _mm512_set1_ps(w);
+    size_t j = 0;
+    for (; j + 16 <= d; j += 16) {
+        const __m512 xv = _mm512_cvtph_ps(
+                _mm256_loadu_si256(reinterpret_cast<const __m256i*>(x + j)));
+        _mm512_storeu_ps(
+                c + j,
+                _mm512_add_ps(_mm512_loadu_ps(c + j), _mm512_mul_ps(xv, wv)));
+    }
+    for (; j < d; ++j) {
+        c[j] += decode_fp16(x[j]) * w;
+    }
+}
+
+template <>
+void fp16_to_fp32_kernel<SIMDLevel::AVX512>(
+        size_t n,
+        const uint16_t* x,
+        float* out) {
+    size_t j = 0;
+    for (; j + 16 <= n; j += 16) {
+        const __m256i h =
+                _mm256_loadu_si256(reinterpret_cast<const __m256i*>(x + j));
+        _mm512_storeu_ps(out + j, _mm512_cvtph_ps(h));
+    }
+    for (; j < n; ++j) {
+        out[j] = decode_fp16(x[j]);
+    }
+}
+
+} // namespace detail
+} // namespace faiss

data/vendor/faiss/faiss/utils/simd_impl/fp16_kernels_neon.cpp

@@ -0,0 +1,48 @@
+/*
+ * 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.
+ */
+
+#include <faiss/utils/simd_impl/fp16_kernels.h>
+
+#include <arm_neon.h>
+
+namespace faiss {
+namespace detail {
+
+template <>
+void fp16_madd<SIMDLevel::ARM_NEON>(
+        size_t d,
+        const uint16_t* x,
+        float w,
+        float* c) {
+    const float32x4_t wv = vdupq_n_f32(w);
+    size_t j = 0;
+    for (; j + 4 <= d; j += 4) {
+        const float32x4_t xv =
+                vcvt_f32_f16(vreinterpret_f16_u16(vld1_u16(x + j)));
+        vst1q_f32(c + j, vaddq_f32(vld1q_f32(c + j), vmulq_f32(xv, wv)));
+    }
+    for (; j < d; ++j) {
+        c[j] += decode_fp16(x[j]) * w;
+    }
+}
+
+template <>
+void fp16_to_fp32_kernel<SIMDLevel::ARM_NEON>(
+        size_t n,
+        const uint16_t* x,
+        float* out) {
+    size_t j = 0;
+    for (; j + 4 <= n; j += 4) {
+        vst1q_f32(out + j, vcvt_f32_f16(vreinterpret_f16_u16(vld1_u16(x + j))));
+    }
+    for (; j < n; ++j) {
+        out[j] = decode_fp16(x[j]);
+    }
+}
+
+} // namespace detail
+} // namespace faiss

data/vendor/faiss/faiss/utils/simd_impl/masked_sum-avx2.h

@@ -0,0 +1,44 @@
+/*
+ * 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 <immintrin.h>
+
+#include <faiss/utils/simd_impl/masked_sum.h>
+
+namespace faiss {
+namespace {
+
+template <>
+inline float masked_sum<SIMDLevel::AVX2>(
+        const float* x,
+        const uint8_t* bits,
+        size_t d) {
+    const __m256i bit_masks = _mm256_set_epi32(128, 64, 32, 16, 8, 4, 2, 1);
+    __m256 acc = _mm256_setzero_ps();
+    const size_t full_bytes = d / 8;
+    for (size_t byte_idx = 0; byte_idx < full_bytes; byte_idx++) {
+        const __m256i byte_broadcast = _mm256_set1_epi32(bits[byte_idx]);
+        const __m256i masked = _mm256_and_si256(byte_broadcast, bit_masks);
+        const __m256 mask =
+                _mm256_castsi256_ps(_mm256_cmpeq_epi32(masked, bit_masks));
+        const __m256 vals = _mm256_loadu_ps(x + byte_idx * 8);
+        acc = _mm256_add_ps(acc, _mm256_and_ps(mask, vals));
+    }
+    __m128 sum128 = _mm_add_ps(
+            _mm256_castps256_ps128(acc), _mm256_extractf128_ps(acc, 1));
+    __m128 shuf = _mm_movehdup_ps(sum128);
+    __m128 sums = _mm_add_ps(sum128, shuf);
+    shuf = _mm_movehl_ps(shuf, sums);
+    sums = _mm_add_ss(sums, shuf);
+    return _mm_cvtss_f32(sums) +
+            masked_sum_scalar(x + full_bytes * 8, bits + full_bytes, d % 8);
+}
+
+} // namespace
+} // namespace faiss

data/vendor/faiss/faiss/utils/simd_impl/masked_sum-avx512.h

@@ -0,0 +1,43 @@
+/*
+ * 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 <cstring>
+
+#include <immintrin.h>
+
+#include <faiss/utils/simd_impl/masked_sum.h>
+
+namespace faiss {
+namespace {
+
+template <>
+inline float masked_sum<SIMDLevel::AVX512>(
+        const float* x,
+        const uint8_t* bits,
+        size_t d) {
+    __m512 acc = _mm512_setzero_ps();
+    size_t i = 0;
+    for (; i + 16 <= d; i += 16) {
+        uint16_t mask16;
+        memcpy(&mask16, bits + i / 8, sizeof(mask16));
+        acc = _mm512_mask_add_ps(acc, mask16, acc, _mm512_loadu_ps(x + i));
+    }
+    if (i < d) {
+        uint32_t mask = bits[i / 8];
+        if (d - i > 8) {
+            mask |= uint32_t(bits[i / 8 + 1]) << 8;
+        }
+        const __mmask16 k = mask & ((1u << (d - i)) - 1);
+        acc = _mm512_add_ps(acc, _mm512_maskz_loadu_ps(k, x + i));
+    }
+    return _mm512_reduce_add_ps(acc);
+}
+
+} // namespace
+} // namespace faiss

data/vendor/faiss/faiss/utils/simd_impl/masked_sum.h

@@ -0,0 +1,38 @@
+/*
+ * 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 <cstddef>
+#include <cstdint>
+
+#include <faiss/utils/simd_levels.h>
+
+namespace faiss {
+
+// Internal linkage: each file compiles these with its own flags.
+// masked_sum-avx2.h and masked_sum-avx512.h specialize masked_sum.
+namespace {
+
+inline float masked_sum_scalar(const float* x, const uint8_t* bits, size_t d) {
+    float sum = 0;
+    for (size_t j = 0; j < d; j++) {
+        if ((bits[j / 8] >> (j % 8)) & 1) {
+            sum += x[j];
+        }
+    }
+    return sum;
+}
+
+template <SIMDLevel SL>
+inline float masked_sum(const float* x, const uint8_t* bits, size_t d) {
+    return masked_sum_scalar(x, bits, d);
+}
+
+} // namespace
+
+} // namespace faiss

data/vendor/faiss/faiss/utils/simd_impl/partitioning_avx2.cpp

@@ -5,10 +5,6 @@
  * 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/simd_impl/partitioning_simdlib256.h>
-
-#endif // COMPILE_SIMD_AVX2

data/vendor/faiss/faiss/utils/simd_impl/partitioning_neon.cpp

@@ -5,10 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#ifdef COMPILE_SIMD_ARM_NEON
-
 #define THE_SIMD_LEVEL SIMDLevel::ARM_NEON
 // NOLINTNEXTLINE(facebook-hte-InlineHeader)
 #include <faiss/utils/simd_impl/partitioning_simdlib256.h>
-
-#endif // COMPILE_SIMD_ARM_NEON

data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp

@@ -5,8 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#ifdef COMPILE_SIMD_AVX2
-
 #include <faiss/utils/rabitq_simd.h>
 #include <immintrin.h>
 #include <limits>
@@ -275,9 +273,11 @@ uint64_t bitwise_and_dot_product<SIMDLevel::AVX2>(
     }
     sum += reduce_add_128(sum_128);
     for (size_t step = 64 / 8; offset + step <= size; offset += step) {
-        const uint64_t yv = *(const uint64_t*)(data + offset);
+        uint64_t yv;
+        std::memcpy(&yv, data + offset, sizeof(yv));
         for (int j = 0; j < qb; j++) {
-            const uint64_t qv = *(const uint64_t*)(query + j * size + offset);
+            uint64_t qv;
+            std::memcpy(&qv, query + j * size + offset, sizeof(qv));
             sum += popcount64(qv & yv) << j;
         }
     }
@@ -336,10 +336,12 @@ BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount<
     dot_product += reduce_add_128(dot_128);
     popcount_sum += reduce_add_128(pop_128);
     for (size_t step = 64 / 8; offset + step <= size; offset += step) {
-        const uint64_t yv = *(const uint64_t*)(data + offset);
+        uint64_t yv;
+        std::memcpy(&yv, data + offset, sizeof(yv));
         popcount_sum += popcount64(yv);
         for (int j = 0; j < qb; j++) {
-            const uint64_t qv = *(const uint64_t*)(query + j * size + offset);
+            uint64_t qv;
+            std::memcpy(&qv, query + j * size + offset, sizeof(qv));
             dot_product += popcount64(qv & yv) << j;
         }
     }
@@ -391,9 +393,11 @@ uint64_t bitwise_xor_dot_product<SIMDLevel::AVX2>(
     }
     sum += reduce_add_128(sum_128);
     for (size_t step = 64 / 8; offset + step <= size; offset += step) {
-        const auto yv = *(const uint64_t*)(data + offset);
+        uint64_t yv;
+        std::memcpy(&yv, data + offset, sizeof(yv));
         for (int j = 0; j < qb; j++) {
-            const auto qv = *(const uint64_t*)(query + j * size + offset);
+            uint64_t qv;
+            std::memcpy(&qv, query + j * size + offset, sizeof(qv));
             sum += popcount64(qv ^ yv) << j;
         }
     }
@@ -427,7 +431,8 @@ uint64_t popcount<SIMDLevel::AVX2>(const uint8_t* data, size_t size) {
     }
     sum += reduce_add_128(sum_128);
     for (size_t step = 64 / 8; offset + step <= size; offset += step) {
-        const auto yv = *(const uint64_t*)(data + offset);
+        uint64_t yv;
+        std::memcpy(&yv, data + offset, sizeof(yv));
         sum += popcount64(yv);
     }
     for (; offset < size; ++offset) {
@@ -517,7 +522,6 @@ inline float ip_1exbit_avx2(
     return result;
 }
 
-#ifdef __BMI2__
 inline float ip_bitplane_avx2(
         const uint8_t* __restrict sign_bits,
         const uint8_t* __restrict ex_code,
@@ -571,7 +575,6 @@ inline float ip_bitplane_avx2(
     result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb);
     return result;
 }
-#endif // __BMI2__
 
 } // namespace
 
@@ -583,18 +586,54 @@ float compute_inner_product<SIMDLevel::AVX2>(
         size_t d,
         size_t ex_bits,
         float cb) {
+    if (ex_bits == 8) {
+        // Eight byte-aligned extra codes, plus their independent sign bits.
+        // One sign-word broadcast serves 32 dims, to save shuffle-port uops.
+        __m256 acc0 = _mm256_setzero_ps();
+        __m256 acc1 = _mm256_setzero_ps();
+        const __m256 weight = _mm256_set1_ps(256.f);
+        const __m256 offset = _mm256_set1_ps(cb);
+        const __m256i positions =
+                _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128);
+        // cmpeq, not cmpgt: bit 31 of the shifted positions is a sign bit.
+        auto step = [&](__m256& acc, size_t i, __m256i signs, __m256i pos) {
+            const __m128i bytes = _mm_loadl_epi64(
+                    reinterpret_cast<const __m128i*>(ex_code + i));
+            const __m256 extra =
+                    _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(bytes));
+            const __m256i mask =
+                    _mm256_cmpeq_epi32(_mm256_and_si256(signs, pos), pos);
+            const __m256 recon = _mm256_add_ps(
+                    extra, _mm256_and_ps(_mm256_castsi256_ps(mask), weight));
+            acc = _mm256_fmadd_ps(
+                    _mm256_loadu_ps(rotated_q + i),
+                    _mm256_add_ps(recon, offset),
+                    acc);
+        };
+        size_t i = 0;
+        for (; i + 32 <= d; i += 32) {
+            uint32_t word;
+            memcpy(&word, sign_bits + i / 8, sizeof(word));
+            const __m256i signs = _mm256_set1_epi32(static_cast<int>(word));
+            step(acc0, i, signs, positions);
+            step(acc1, i + 8, signs, _mm256_slli_epi32(positions, 8));
+            step(acc0, i + 16, signs, _mm256_slli_epi32(positions, 16));
+            step(acc1, i + 24, signs, _mm256_slli_epi32(positions, 24));
+        }
+        for (; i + 8 <= d; i += 8) {
+            step(acc0, i, _mm256_set1_epi32(sign_bits[i / 8]), positions);
+        }
+        return hsum_avx2(_mm256_add_ps(acc0, acc1)) +
+                ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb);
+    }
     if (ex_bits == 1) {
         return ip_1exbit_avx2(sign_bits, ex_code, rotated_q, d, cb);
     }
 
-#ifdef __BMI2__
     if (ex_bits <= 7) {
         return ip_bitplane_avx2(sign_bits, ex_code, rotated_q, d, ex_bits, cb);
     }
-#endif
     return ip_scalar(sign_bits, ex_code, rotated_q, 0, d, ex_bits, cb);
 }
 
 } // namespace faiss::rabitq::multibit
-
-#endif // COMPILE_SIMD_AVX2