faiss 0.5.2 → 0.6.0

data/vendor/faiss/faiss/utils/pq_code_distance.h

@@ -0,0 +1,251 @@
+/*
+ * 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 <type_traits>
+
+#include <faiss/impl/ProductQuantizer.h>
+#include <faiss/impl/platform_macros.h>
+#include <faiss/utils/simd_levels.h>
+
+namespace faiss {
+namespace pq_code_distance {
+
+/*********************************************************************
+ * PQCodeDistance — SIMD-dispatched PQ code distance
+ *
+ * Computes the distance from a PQ-encoded vector to a query vector,
+ * given a precomputed table of sub-distances (one per subquantizer
+ * per centroid). Originally extracted from IndexIVFPQ.cpp.
+ *
+ * DESIGN:
+ *
+ * PQCodeDistance<PQDecoderT, SL> computes PQ code distances at a given
+ * SIMD level. The dispatch site (IndexIVFPQ.cpp, IndexPQ.cpp) uses
+ * DISPATCH_SIMDLevel to select SL at runtime, which instantiates
+ * PQCodeDistance for ALL decoder types (PQDecoder8, PQDecoder16,
+ * PQDecoderGeneric) at the chosen level.
+ *
+ * Only PQDecoder8 has SIMD-optimized implementations (AVX2, AVX512,
+ * ARM_SVE). The other decoders always use scalar code — their decode()
+ * method is inherently sequential, so SIMD doesn't help.
+ *
+ * The primary template is always complete (no forward declarations
+ * needed). For PQDecoder8, it delegates to _impl dispatch bridge
+ * functions whose specializations are defined in per-SIMD .cpp files
+ * and resolved at link time. For other decoders, it uses scalar.
+ *
+ * ADDING A NEW SIMD LEVEL:
+ *
+ *   1. Add the level to SIMDLevel enum (simd_levels.h)
+ *   2. Add dispatch_config entry (simd_dispatch.bzl)
+ *   3. Define pq_code_distance_single_impl<NEW_LEVEL> and
+ *      pq_code_distance_four_impl<NEW_LEVEL> specializations in a
+ *      new .cpp file compiled with appropriate SIMD flags
+ *   4. Add the .cpp to the build (CMakeLists.txt, xplat.bzl)
+ *********************************************************************/
+
+/// Scalar PQ code distance implementation.
+/// Templated only on decoder type, independent of SIMD level.
+/// Used directly by non-PQDecoder8 decoders (PQDecoder16,
+/// PQDecoderGeneric) and as fallback for PQDecoder8 at NONE/NEON.
+template <typename PQDecoderT>
+struct PQCodeDistanceScalar {
+    using PQDecoder = PQDecoderT;
+
+    static float distance_single_code(
+            // number of subquantizers
+            size_t M,
+            size_t nbits,
+            // precomputed distances, layout (M, ksub)
+            const float* sim_table,
+            const uint8_t* code) {
+        PQDecoderT decoder(code, nbits);
+        const size_t ksub = 1 << nbits;
+
+        const float* tab = sim_table;
+        float result = 0;
+
+        for (size_t m = 0; m < M; m++) {
+            result += tab[decoder.decode()];
+            tab += ksub;
+        }
+
+        return result;
+    }
+
+    static void distance_four_codes(
+            size_t M,
+            size_t nbits,
+            const float* sim_table,
+            const uint8_t* __restrict code0,
+            const uint8_t* __restrict code1,
+            const uint8_t* __restrict code2,
+            const uint8_t* __restrict code3,
+            float& result0,
+            float& result1,
+            float& result2,
+            float& result3) {
+        PQDecoderT decoder0(code0, nbits);
+        PQDecoderT decoder1(code1, nbits);
+        PQDecoderT decoder2(code2, nbits);
+        PQDecoderT decoder3(code3, nbits);
+        const size_t ksub = 1 << nbits;
+
+        const float* tab = sim_table;
+        result0 = 0;
+        result1 = 0;
+        result2 = 0;
+        result3 = 0;
+
+        for (size_t m = 0; m < M; m++) {
+            result0 += tab[decoder0.decode()];
+            result1 += tab[decoder1.decode()];
+            result2 += tab[decoder2.decode()];
+            result3 += tab[decoder3.decode()];
+            tab += ksub;
+        }
+    }
+};
+
+/*********************************************************************
+ * Dispatch bridge — function templates for PQDecoder8 SIMD dispatch.
+ *
+ * Primary declarations only; specializations are defined in per-SIMD
+ * .cpp files (AVX2, AVX512, ARM_SVE) and pq_code_distance-generic.cpp
+ * (NONE, ARM_NEON). Same pattern as fvec_L2sqr et al. in distances.h.
+ *********************************************************************/
+
+template <SIMDLevel SL>
+float pq_code_distance_single_impl(
+        size_t M,
+        size_t nbits,
+        const float* sim_table,
+        const uint8_t* code);
+
+template <SIMDLevel SL>
+void pq_code_distance_four_impl(
+        size_t M,
+        size_t nbits,
+        const float* sim_table,
+        const uint8_t* __restrict code0,
+        const uint8_t* __restrict code1,
+        const uint8_t* __restrict code2,
+        const uint8_t* __restrict code3,
+        float& result0,
+        float& result1,
+        float& result2,
+        float& result3);
+
+/// Primary template — always complete.
+/// For PQDecoder8, delegates to _impl dispatch bridges (resolved at
+/// link time to per-SIMD implementations). For other decoders, uses
+/// scalar — their sequential decode() methods don't benefit from SIMD.
+template <typename PQDecoderT, SIMDLevel SL>
+struct PQCodeDistance {
+    using PQDecoder = PQDecoderT;
+
+    static float distance_single_code(
+            size_t M,
+            size_t nbits,
+            const float* sim_table,
+            const uint8_t* code) {
+        if constexpr (std::is_same_v<PQDecoderT, PQDecoder8>) {
+            return pq_code_distance_single_impl<SL>(M, nbits, sim_table, code);
+        } else {
+            return PQCodeDistanceScalar<PQDecoderT>::distance_single_code(
+                    M, nbits, sim_table, code);
+        }
+    }
+
+    static void distance_four_codes(
+            size_t M,
+            size_t nbits,
+            const float* sim_table,
+            const uint8_t* __restrict code0,
+            const uint8_t* __restrict code1,
+            const uint8_t* __restrict code2,
+            const uint8_t* __restrict code3,
+            float& result0,
+            float& result1,
+            float& result2,
+            float& result3) {
+        if constexpr (std::is_same_v<PQDecoderT, PQDecoder8>) {
+            pq_code_distance_four_impl<SL>(
+                    M,
+                    nbits,
+                    sim_table,
+                    code0,
+                    code1,
+                    code2,
+                    code3,
+                    result0,
+                    result1,
+                    result2,
+                    result3);
+        } else {
+            PQCodeDistanceScalar<PQDecoderT>::distance_four_codes(
+                    M,
+                    nbits,
+                    sim_table,
+                    code0,
+                    code1,
+                    code2,
+                    code3,
+                    result0,
+                    result1,
+                    result2,
+                    result3);
+        }
+    }
+};
+
+/*********************************************************************
+ * Non-templated PQ code distance dispatch (PQDecoder8 only).
+ *
+ * These follow the same pattern as distances.h: the caller does not
+ * name a SIMDLevel. Internally they dispatch via DISPATCH_SIMDLevel
+ * to the best available SIMD implementation (DD: runtime detection,
+ * static: compile-time selection). Definitions are in
+ * pq_code_distance-generic.cpp.
+ *********************************************************************/
+
+/// Compute PQ distance for a single code, dispatching to the best
+/// available SIMD level.
+FAISS_API float pq_code_distance_single(
+        size_t M,
+        size_t nbits,
+        const float* sim_table,
+        const uint8_t* code);
+
+/// Compute PQ distances for four codes simultaneously, dispatching
+/// to the best available SIMD level.
+FAISS_API void pq_code_distance_four(
+        size_t M,
+        size_t nbits,
+        const float* sim_table,
+        const uint8_t* __restrict code0,
+        const uint8_t* __restrict code1,
+        const uint8_t* __restrict code2,
+        const uint8_t* __restrict code3,
+        float& result0,
+        float& result1,
+        float& result2,
+        float& result3);
+
+} // namespace pq_code_distance
+
+// Re-export public API into namespace faiss for convenience
+using pq_code_distance::pq_code_distance_four;
+using pq_code_distance::pq_code_distance_single;
+using pq_code_distance::PQCodeDistance;
+using pq_code_distance::PQCodeDistanceScalar;
+
+} // namespace faiss

data/vendor/faiss/faiss/utils/rabitq_simd.h

@@ -9,6 +9,7 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <cstring>
 
 // Only include x86 SIMD intrinsics on x86/x86_64 architectures
 #if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || \
@@ -423,3 +424,262 @@ inline uint64_t popcount(const uint8_t* data, size_t size) {
 }
 
 } // namespace faiss::rabitq
+
+/*********************************************************
+ * Multi-bit RaBitQ inner product kernels.
+ *
+ * Compute: sum_i rotated_q[i] * ((sign_bit_i << ex_bits) + ex_code_val_i + cb)
+ *
+ * Strategy:
+ *   ex_bits == 1: Specialized kernel — both sign_bits and ex_code are
+ *                 1-bit-per-dim packed, enabling direct bit→mask→float
+ *                 conversion with zero per-element extraction.
+ *   ex_bits >= 2: Bit-plane decomposition (BMI2 required) — PEXT extracts
+ *                 each bit plane in one instruction, then the same
+ *                 bit→mask→float kernel computes each plane's dot product.
+ *   Fallback:     Scalar extraction via 64-bit window read + shift + mask.
+ *********************************************************/
+namespace faiss::rabitq::multibit {
+
+/// Scalar inner product for multi-bit RaBitQ.
+/// Extracts each code value in O(1) via 64-bit window read + shift + mask.
+/// Also serves as the tail handler for SIMD kernels via the @p start parameter.
+inline float ip_scalar(
+        const uint8_t* __restrict sign_bits,
+        const uint8_t* __restrict ex_code,
+        const float* __restrict rotated_q,
+        size_t start,
+        size_t d,
+        size_t ex_bits,
+        float cb) {
+    float result = 0.0f;
+    const int sign_shift = static_cast<int>(ex_bits);
+    const uint64_t code_mask = (1ULL << ex_bits) - 1;
+    for (size_t i = start; i < d; i++) {
+        int sb = (sign_bits[i / 8] >> (i % 8)) & 1;
+        size_t bit_pos = i * ex_bits;
+        size_t byte_idx = bit_pos / 8;
+        size_t bit_offset = bit_pos % 8;
+        uint64_t raw = 0;
+        memcpy(&raw, ex_code + byte_idx, sizeof(uint64_t));
+        int ex_val = static_cast<int>((raw >> bit_offset) & code_mask);
+        result += rotated_q[i] *
+                (static_cast<float>((sb << sign_shift) + ex_val) + cb);
+    }
+    return result;
+}
+
+#if defined(__x86_64__) || defined(_M_X64)
+
+#if defined(__AVX2__)
+/// Horizontal sum of 8 floats in a __m256 register.
+inline float hsum_avx2(__m256 v) {
+    __m128 hi = _mm256_extractf128_ps(v, 1);
+    __m128 lo = _mm256_castps256_ps128(v);
+    lo = _mm_add_ps(lo, hi);
+    __m128 shuf = _mm_movehdup_ps(lo);
+    lo = _mm_add_ps(lo, shuf);
+    shuf = _mm_movehl_ps(shuf, lo);
+    return _mm_cvtss_f32(_mm_add_ss(lo, shuf));
+}
+#endif // __AVX2__
+
+/*********************************************************
+ * Specialized 1-bit kernels (ex_bits == 1).
+ *
+ * For 1 extra bit, both sign_bits and ex_code are 1-bit-per-dim packed,
+ * so we convert bits to floats directly — no extraction loops needed.
+ *********************************************************/
+
+#if defined(__AVX512F__)
+/// AVX-512: 16 dims/iter, ex_bits == 1.
+inline float ip_1exbit_avx512(
+        const uint8_t* __restrict sign_bits,
+        const uint8_t* __restrict ex_code,
+        const float* __restrict rotated_q,
+        size_t d,
+        float cb) {
+    __m512 acc = _mm512_setzero_ps();
+    const __m512 v_cb = _mm512_set1_ps(cb);
+    const __m512 v_two = _mm512_set1_ps(2.0f);
+    const __m512 v_one = _mm512_set1_ps(1.0f);
+
+    size_t i = 0;
+    for (; i + 16 <= d; i += 16) {
+        uint16_t sb16;
+        memcpy(&sb16, sign_bits + i / 8, sizeof(uint16_t));
+        uint16_t eb16;
+        memcpy(&eb16, ex_code + i / 8, sizeof(uint16_t));
+
+        __m512 sb_f = _mm512_maskz_mov_ps(_cvtu32_mask16(sb16), v_one);
+        __m512 eb_f = _mm512_maskz_mov_ps(_cvtu32_mask16(eb16), v_one);
+
+        __m512 recon = _mm512_add_ps(_mm512_fmadd_ps(sb_f, v_two, eb_f), v_cb);
+        __m512 rq = _mm512_loadu_ps(rotated_q + i);
+        acc = _mm512_fmadd_ps(rq, recon, acc);
+    }
+
+    float result = _mm512_reduce_add_ps(acc);
+    result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, 1, cb);
+    return result;
+}
+#endif // __AVX512F__
+
+#if defined(__AVX2__)
+/// AVX2: 8 dims/iter, ex_bits == 1.
+inline float ip_1exbit_avx2(
+        const uint8_t* __restrict sign_bits,
+        const uint8_t* __restrict ex_code,
+        const float* __restrict rotated_q,
+        size_t d,
+        float cb) {
+    __m256 acc = _mm256_setzero_ps();
+    const __m256 v_cb = _mm256_set1_ps(cb);
+    const __m256 v_two = _mm256_set1_ps(2.0f);
+    const __m256 v_one = _mm256_set1_ps(1.0f);
+    const __m256i bit_pos = _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128);
+    const __m256i zero = _mm256_setzero_si256();
+
+    size_t i = 0;
+    for (; i + 8 <= d; i += 8) {
+        uint8_t sb = sign_bits[i / 8];
+        uint8_t eb = ex_code[i / 8];
+
+        __m256i sb_cmp = _mm256_cmpgt_epi32(
+                _mm256_and_si256(_mm256_set1_epi32(sb), bit_pos), zero);
+        __m256 sb_f = _mm256_and_ps(_mm256_castsi256_ps(sb_cmp), v_one);
+
+        __m256i eb_cmp = _mm256_cmpgt_epi32(
+                _mm256_and_si256(_mm256_set1_epi32(eb), bit_pos), zero);
+        __m256 eb_f = _mm256_and_ps(_mm256_castsi256_ps(eb_cmp), v_one);
+
+        __m256 recon = _mm256_add_ps(_mm256_fmadd_ps(sb_f, v_two, eb_f), v_cb);
+        __m256 rq = _mm256_loadu_ps(rotated_q + i);
+        acc = _mm256_fmadd_ps(rq, recon, acc);
+    }
+
+    float result = hsum_avx2(acc);
+    result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, 1, cb);
+    return result;
+}
+#endif // __AVX2__
+
+/*********************************************************
+ * Bit-plane decomposition kernels (ex_bits >= 2, BMI2 required).
+ *
+ * Decomposes the inner product as:
+ *   ex_ip = (1 << ex_bits) * sign_dot
+ *         + Σ_{b=0}^{ex_bits-1} (1 << b) * plane_dot_b
+ *         + cb * total_q
+ *
+ * Each plane_dot_b is a float × bit-vector dot product, computed using
+ * the same bit→mask→float conversion as the 1-bit kernel. PEXT
+ * extracts each bit plane from the packed ex_code in one instruction
+ * per 8 dimensions.
+ *********************************************************/
+
+#if defined(__AVX2__) && defined(__BMI2__)
+/// AVX2 + BMI2 bit-plane decomposition: 8 dims/iter, ex_bits in [2, 7].
+/// Caller must ensure ex_bits <= 7 (pext_masks[7] / v_weights[8]).
+inline float ip_bitplane_avx2(
+        const uint8_t* __restrict sign_bits,
+        const uint8_t* __restrict ex_code,
+        const float* __restrict rotated_q,
+        size_t d,
+        size_t ex_bits,
+        float cb) {
+    __m256 acc = _mm256_setzero_ps();
+    const __m256 v_one = _mm256_set1_ps(1.0f);
+    const __m256i bit_pos = _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128);
+    const __m256i zero = _mm256_setzero_si256();
+    const __m256 v_cb = _mm256_set1_ps(cb);
+
+    // Precompute PEXT masks and plane weights
+    uint64_t pext_masks[7];
+    __m256 v_weights[8];
+    for (size_t b = 0; b < ex_bits; b++) {
+        uint64_t m = 0;
+        for (int j = 0; j < 8; j++) {
+            m |= (1ULL << (b + j * ex_bits));
+        }
+        pext_masks[b] = m;
+        v_weights[b] = _mm256_set1_ps(static_cast<float>(1u << b));
+    }
+    v_weights[ex_bits] = _mm256_set1_ps(static_cast<float>(1u << ex_bits));
+
+    size_t i = 0;
+    for (; i + 8 <= d; i += 8) {
+        // Sign bit → float via bit mask comparison
+        __m256i sb_cmp = _mm256_cmpgt_epi32(
+                _mm256_and_si256(_mm256_set1_epi32(sign_bits[i / 8]), bit_pos),
+                zero);
+        __m256 recon = _mm256_mul_ps(
+                _mm256_and_ps(_mm256_castsi256_ps(sb_cmp), v_one),
+                v_weights[ex_bits]);
+
+        // Load packed ex_code for 8 dims (8 × ex_bits bits = ex_bits bytes)
+        uint64_t ex64 = 0;
+        memcpy(&ex64, ex_code + (i / 8) * ex_bits, sizeof(uint64_t));
+
+        // Extract each bit plane via PEXT → bit mask → float
+        for (size_t b = 0; b < ex_bits; b++) {
+            auto plane = static_cast<uint8_t>(_pext_u64(ex64, pext_masks[b]));
+            __m256i p_cmp = _mm256_cmpgt_epi32(
+                    _mm256_and_si256(_mm256_set1_epi32(plane), bit_pos), zero);
+            __m256 p_f = _mm256_and_ps(_mm256_castsi256_ps(p_cmp), v_one);
+            recon = _mm256_fmadd_ps(p_f, v_weights[b], recon);
+        }
+
+        __m256 rq = _mm256_loadu_ps(rotated_q + i);
+        acc = _mm256_fmadd_ps(rq, _mm256_add_ps(recon, v_cb), acc);
+    }
+
+    float result = hsum_avx2(acc);
+    result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb);
+    return result;
+}
+#endif // __AVX2__ && __BMI2__
+
+#endif // x86_64
+
+/**
+ * Dispatch to the best available kernel for the given ex_bits.
+ *
+ * Routing (compile-time):
+ *   ex_bits == 1:  specialized 1-bit kernel (AVX-512 > AVX2 > scalar)
+ *   ex_bits >= 2:  bit-plane decomposition (AVX2+BMI2 > scalar)
+ *
+ * @param sign_bits  packed sign bits (1 bit/dim, standard byte packing)
+ * @param ex_code    packed extra-bit codes (ex_bits bits/dim)
+ * @param rotated_q  rotated query vector (float[d])
+ * @param d          dimensionality
+ * @param ex_bits    number of extra bits per dimension (nb_bits - 1)
+ * @param cb         constant bias: -(2^ex_bits - 0.5)
+ * @return           inner product value
+ */
+inline float compute_inner_product(
+        const uint8_t* __restrict sign_bits,
+        const uint8_t* __restrict ex_code,
+        const float* __restrict rotated_q,
+        size_t d,
+        size_t ex_bits,
+        float cb) {
+    if (ex_bits == 1) {
+#if defined(__AVX512F__)
+        return ip_1exbit_avx512(sign_bits, ex_code, rotated_q, d, cb);
+#elif defined(__AVX2__)
+        return ip_1exbit_avx2(sign_bits, ex_code, rotated_q, d, cb);
+#else
+        return ip_scalar(sign_bits, ex_code, rotated_q, 0, d, 1, cb);
+#endif
+    }
+
+#if defined(__AVX2__) && defined(__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

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

@@ -0,0 +1,150 @@
+/*
+ * 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/distances.h>
+
+#ifdef __aarch64__
+
+#include <arm_neon.h>
+#include <limits>
+
+#define AUTOVEC_LEVEL SIMDLevel::ARM_NEON
+#include <faiss/utils/simd_impl/distances_autovec-inl.h>
+
+namespace faiss {
+
+template <>
+void fvec_madd<SIMDLevel::ARM_NEON>(
+        size_t n,
+        const float* a,
+        float bf,
+        const float* b,
+        float* c) {
+    const size_t n_simd = n - (n & 3);
+    const float32x4_t bfv = vdupq_n_f32(bf);
+    size_t i;
+    for (i = 0; i < n_simd; i += 4) {
+        const float32x4_t ai = vld1q_f32(a + i);
+        const float32x4_t bi = vld1q_f32(b + i);
+        const float32x4_t ci = vfmaq_f32(ai, bfv, bi);
+        vst1q_f32(c + i, ci);
+    }
+    for (; i < n; ++i)
+        c[i] = a[i] + bf * b[i];
+}
+
+template <>
+void fvec_L2sqr_ny_transposed<SIMDLevel::ARM_NEON>(
+        float* dis,
+        const float* x,
+        const float* y,
+        const float* y_sqlen,
+        size_t d,
+        size_t d_offset,
+        size_t ny) {
+    // Use autovectorized implementation
+    fvec_L2sqr_ny_transposed<SIMDLevel::NONE>(
+            dis, x, y, y_sqlen, d, d_offset, ny);
+}
+
+template <>
+void fvec_inner_products_ny<SIMDLevel::ARM_NEON>(
+        float* dis,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t ny) {
+    fvec_inner_products_ny<SIMDLevel::NONE>(dis, x, y, d, ny);
+}
+
+template <>
+void fvec_L2sqr_ny<SIMDLevel::ARM_NEON>(
+        float* dis,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t ny) {
+    fvec_L2sqr_ny<SIMDLevel::NONE>(dis, x, y, d, ny);
+}
+
+template <>
+size_t fvec_L2sqr_ny_nearest<SIMDLevel::ARM_NEON>(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t ny) {
+    return fvec_L2sqr_ny_nearest<SIMDLevel::NONE>(
+            distances_tmp_buffer, x, y, d, ny);
+}
+
+template <>
+size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::ARM_NEON>(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        const float* y_sqlen,
+        size_t d,
+        size_t d_offset,
+        size_t ny) {
+    return fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::NONE>(
+            distances_tmp_buffer, x, y, y_sqlen, d, d_offset, ny);
+}
+
+template <>
+int fvec_madd_and_argmin<SIMDLevel::ARM_NEON>(
+        size_t n,
+        const float* a,
+        float bf,
+        const float* b,
+        float* c) {
+    float32x4_t vminv = vdupq_n_f32(1e20);
+    uint32x4_t iminv = vdupq_n_u32(static_cast<uint32_t>(-1));
+    size_t i;
+    {
+        const size_t n_simd = n - (n & 3);
+        const uint32_t iota[] = {0, 1, 2, 3};
+        uint32x4_t iv = vld1q_u32(iota);
+        const uint32x4_t incv = vdupq_n_u32(4);
+        const float32x4_t bfv = vdupq_n_f32(bf);
+        for (i = 0; i < n_simd; i += 4) {
+            const float32x4_t ai = vld1q_f32(a + i);
+            const float32x4_t bi = vld1q_f32(b + i);
+            const float32x4_t ci = vfmaq_f32(ai, bfv, bi);
+            vst1q_f32(c + i, ci);
+            const uint32x4_t less_than = vcltq_f32(ci, vminv);
+            vminv = vminq_f32(ci, vminv);
+            iminv = vorrq_u32(
+                    vandq_u32(less_than, iv),
+                    vandq_u32(vmvnq_u32(less_than), iminv));
+            iv = vaddq_u32(iv, incv);
+        }
+    }
+    float vmin = vminvq_f32(vminv);
+    uint32_t imin;
+    {
+        const float32x4_t vminy = vdupq_n_f32(vmin);
+        const uint32x4_t equals = vceqq_f32(vminv, vminy);
+        imin = vminvq_u32(vorrq_u32(
+                vandq_u32(equals, iminv),
+                vandq_u32(
+                        vmvnq_u32(equals),
+                        vdupq_n_u32(std::numeric_limits<uint32_t>::max()))));
+    }
+    for (; i < n; ++i) {
+        c[i] = a[i] + bf * b[i];
+        if (c[i] < vmin) {
+            vmin = c[i];
+            imin = static_cast<uint32_t>(i);
+        }
+    }
+    return static_cast<int>(imin);
+}
+
+} // namespace faiss
+
+#endif // __aarch64__