faiss 0.6.4 → 0.6.5

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

@@ -5,12 +5,8 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#ifdef COMPILE_SIMD_AVX2
-
 #define THE_SIMD_LEVEL SIMDLevel::AVX2
 // NOLINTBEGIN(facebook-hte-InlineHeader,facebook-unused-include-check)
 #include <faiss/utils/hamming_distance/hamming_computer-avx2.h>
 #include <faiss/utils/hamming_distance/hamming_impl.h>
 // NOLINTEND(facebook-hte-InlineHeader,facebook-unused-include-check)
-
-#endif // COMPILE_SIMD_AVX2

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

@@ -5,11 +5,7 @@
  * 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_avx512_vpopcnt.cpp

@@ -5,8 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#ifdef COMPILE_SIMD_AVX512_VPOPCNT
-
 #define THE_SIMD_LEVEL SIMDLevel::AVX512_VPOPCNT
 #include <faiss/utils/hamming_distance/hamming_computer-avx512_vpopcnt.h>
 #include <faiss/utils/hamming_distance/hamming_impl.h>
@@ -20,5 +18,3 @@
 // Fails the build if the batch path stops being selected for this computer.
 static_assert(faiss::has_hamming_batch<
               faiss::HammingComputer20_tpl<faiss::SIMDLevel::AVX512_VPOPCNT>>);
-
-#endif // COMPILE_SIMD_AVX512_VPOPCNT

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

@@ -103,12 +103,40 @@ struct HammingComputerDefault_tpl<SIMDLevel::AVX512> {
         remainder8 = code_size % 8;
     }
 
+    /// No VPOPCNTDQ at this level: count the bits of each nibble with a
+    /// 16-entry table (vpshufb), then sum the bytes per 64-bit lane (vpsadbw).
+    /// The partial last block is a masked load, so no byte past the code is
+    /// read.
     int hamming(const uint8_t* b8) const {
-        const uint64_t* a64 = reinterpret_cast<const uint64_t*>(a8);
-        const uint64_t* b64 = reinterpret_cast<const uint64_t*>(b8);
-
-        return hamming_popcount_tail(
-                a64, b64, 0, quotient8, a8, b8, remainder8);
+        const int code_size = quotient8 * 8 + remainder8;
+        const __m512i lut = _mm512_broadcast_i32x4(
+                _mm_setr_epi8(0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4));
+        const __m512i low4 = _mm512_set1_epi8(0x0f);
+        const auto count = [&](__m512i x) {
+            const __m512i lo = _mm512_and_si512(x, low4);
+            const __m512i hi = _mm512_and_si512(_mm512_srli_epi16(x, 4), low4);
+            const __m512i c = _mm512_add_epi8(
+                    _mm512_shuffle_epi8(lut, lo), _mm512_shuffle_epi8(lut, hi));
+            return _mm512_sad_epu8(c, _mm512_setzero_si512());
+        };
+        __m512i accu = _mm512_setzero_si512();
+        int i = 0;
+        for (; i + 64 <= code_size; i += 64) {
+            accu = _mm512_add_epi64(
+                    accu,
+                    count(_mm512_xor_si512(
+                            _mm512_loadu_si512(a8 + i),
+                            _mm512_loadu_si512(b8 + i))));
+        }
+        if (i < code_size) {
+            const __mmask64 mask = _bzhi_u64(~uint64_t(0), code_size - i);
+            accu = _mm512_add_epi64(
+                    accu,
+                    count(_mm512_xor_si512(
+                            _mm512_maskz_loadu_epi8(mask, a8 + i),
+                            _mm512_maskz_loadu_epi8(mask, b8 + i))));
+        }
+        return static_cast<int>(_mm512_reduce_add_epi64(accu));
     }
 
     inline int get_code_size() const {

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

@@ -186,40 +186,25 @@ struct HammingComputerDefault_tpl<SIMDLevel::AVX512_VPOPCNT> {
         remainder8 = code_size % 8;
     }
 
+    /// One vector accumulator, reduced once. The partial last block is a
+    /// masked load, so no byte past the code is read.
     int hamming(const uint8_t* b8) const {
-        int accu = 0;
-
-        const uint64_t* a64 = reinterpret_cast<const uint64_t*>(a8);
-        const uint64_t* b64 = reinterpret_cast<const uint64_t*>(b8);
-
+        const int code_size = quotient8 * 8 + remainder8;
+        __m512i accu = _mm512_setzero_si512();
         int i = 0;
-        int quotient64 = quotient8 / 8;
-        for (; i < quotient64; ++i) {
-            __m512i vxor = _mm512_xor_si512(
-                    _mm512_loadu_si512(&a64[i * 8]),
-                    _mm512_loadu_si512(&b64[i * 8]));
-            __m512i vpcnt = _mm512_popcnt_epi64(vxor);
-            accu += _mm512_reduce_add_epi32(vpcnt);
+        for (; i + 64 <= code_size; i += 64) {
+            const __m512i vxor = _mm512_xor_si512(
+                    _mm512_loadu_si512(a8 + i), _mm512_loadu_si512(b8 + i));
+            accu = _mm512_add_epi64(accu, _mm512_popcnt_epi64(vxor));
         }
-        i *= 8;
-
-        // Handle 4-word (256-bit) remainder with VPOPCNTDQ VL
-        if (i + 4 <= quotient8) {
-            __m256i vxor = _mm256_xor_si256(
-                    _mm256_loadu_si256((const __m256i*)&a64[i]),
-                    _mm256_loadu_si256((const __m256i*)&b64[i]));
-            __m256i vpcnt = _mm256_popcnt_epi64(vxor);
-            __m128i lo = _mm256_castsi256_si128(vpcnt);
-            __m128i hi = _mm256_extracti128_si256(vpcnt, 1);
-            __m128i sum = _mm_add_epi64(lo, hi);
-            accu += static_cast<int>(
-                    _mm_extract_epi64(sum, 0) + _mm_extract_epi64(sum, 1));
-            i += 4;
+        if (i < code_size) {
+            const __mmask64 mask = _bzhi_u64(~uint64_t(0), code_size - i);
+            const __m512i vxor = _mm512_xor_si512(
+                    _mm512_maskz_loadu_epi8(mask, a8 + i),
+                    _mm512_maskz_loadu_epi8(mask, b8 + i));
+            accu = _mm512_add_epi64(accu, _mm512_popcnt_epi64(vxor));
         }
-
-        accu += hamming_popcount_tail(
-                a64, b64, i, quotient8, a8, b8, remainder8);
-        return accu;
+        return static_cast<int>(_mm512_reduce_add_epi64(accu));
     }
 
     inline int get_code_size() const {

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

@@ -16,8 +16,6 @@
 // hamming_computer-generic.h. SIMDLevel::ARM_NEON specializations for the
 // ISA-varying HammingComputer and GenHammingComputer structs live in this file.
 
-#ifdef __aarch64__
-
 #include <arm_neon.h>
 
 #include <cassert>
@@ -318,5 +316,3 @@ struct GenHammingComputerM8_tpl<SIMDLevel::ARM_NEON> {
 } // namespace faiss
 
 #endif
-
-#endif

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

@@ -102,7 +102,8 @@ void hammings_knn_hc_impl(
                 CMax<hamdis_t, int64_t>,                                     \
                 NB,                                                          \
                 BD,                                                          \
-                HammingComputer>::                                           \
+                HammingComputer,                                             \
+                THE_SIMD_LEVEL>::                                            \
                 addn(j1 - j0, hc, bs2_, k, bh_val_, bh_ids_, sel);           \
         break;
 

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

@@ -5,11 +5,7 @@
  * 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/hamming_distance/hamming_computer-neon.h>
 #include <faiss/utils/hamming_distance/hamming_impl.h>
-
-#endif // COMPILE_SIMD_ARM_NEON

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

@@ -5,11 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#ifdef COMPILE_SIMD_RISCV_RVV
-
 #define THE_SIMD_LEVEL SIMDLevel::RISCV_RVV
 // NOLINTNEXTLINE(facebook-hte-InlineHeader)
 #include <faiss/utils/hamming_distance/hamming_computer-rvv.h>
 #include <faiss/utils/hamming_distance/hamming_impl.h>
-
-#endif // COMPILE_SIMD_RISCV_RVV

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

@@ -53,10 +53,14 @@ namespace pq_code_distance {
  *********************************************************************/
 
 /// 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>
+// SL keeps each SIMD-level TU's instantiation a distinct symbol. Without it
+// the linker merges them and may keep one compiled with -mavx512*, which
+// then runs on the AVX2 path and raises SIGILL on CPUs without AVX512.
+// The default keeps callers outside faiss that name only the decoder (e.g.
+// unicorn/features/encoders/PQFSTable.h) compiling; faiss passes SL itself.
+template <typename PQDecoderT, SIMDLevel SL = SIMDLevel::NONE>
 struct PQCodeDistanceScalar {
     using PQDecoder = PQDecoderT;
 
@@ -160,7 +164,7 @@ struct PQCodeDistance {
         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(
+            return PQCodeDistanceScalar<PQDecoderT, SL>::distance_single_code(
                     M, nbits, sim_table, code);
         }
     }
@@ -191,7 +195,7 @@ struct PQCodeDistance {
                     result2,
                     result3);
         } else {
-            PQCodeDistanceScalar<PQDecoderT>::distance_four_codes(
+            PQCodeDistanceScalar<PQDecoderT, SL>::distance_four_codes(
                     M,
                     nbits,
                     sim_table,

data/vendor/faiss/faiss/utils/prefetch.h

@@ -25,34 +25,9 @@ inline void prefetch_L3(const void* address) {
     _mm_prefetch((const char*)address, _MM_HINT_T2);
 }
 
-#elif defined(__aarch64__)
-
-// ARM64
-
-#ifdef _MSC_VER
-
-// todo: arm on MSVC
-inline void prefetch_L1(const void* address) {}
-inline void prefetch_L2(const void* address) {}
-inline void prefetch_L3(const void* address) {}
-
-#else
-// arm on non-MSVC
-
-inline void prefetch_L1(const void* address) {
-    __builtin_prefetch(address, 0, 3);
-}
-inline void prefetch_L2(const void* address) {
-    __builtin_prefetch(address, 0, 2);
-}
-inline void prefetch_L3(const void* address) {
-    __builtin_prefetch(address, 0, 1);
-}
-#endif
-
 #else
 
-// a generic platform
+// other platforms, including ARM64
 
 #ifdef _MSC_VER
 

data/vendor/faiss/faiss/utils/quantize_lut.cpp

@@ -60,6 +60,19 @@ void round_tab(const float* tab, size_t n, float a, float bi, T* tab_out) {
 
 } // anonymous namespace
 
+float fastscan_lut_scale(
+        float max_column_span,
+        float max_sum_span,
+        size_t n_rounded_terms) {
+    const float headroom = 0.5f * float(n_rounded_terms);
+    const float sum_budget = 65535.0f - headroom;
+    const float a = std::min(
+            max_column_span > 0 ? 255.0f / max_column_span : HUGE_VALF,
+            (max_sum_span > 0 && sum_budget > 0) ? sum_budget / max_sum_span
+                                                 : HUGE_VALF);
+    return std::isfinite(a) ? a : 0.0f;
+}
+
 void round_uint8_per_column(
         float* tab,
         size_t n,
@@ -67,6 +80,7 @@ void round_uint8_per_column(
         float* a_out,
         float* b_out) {
     float max_span = 0;
+    float total_span = 0;
     std::vector<float> mins(n);
     for (size_t i = 0; i < n; i++) {
         mins[i] = tab_min(tab + i * d, d);
@@ -74,8 +88,9 @@ void round_uint8_per_column(
         if (span > max_span) {
             max_span = span;
         }
+        total_span += span;
     }
-    float a = max_span > 0 ? 255.0f / max_span : 0.0f;
+    float a = fastscan_lut_scale(max_span, total_span, n);
     float b = 0;
     for (size_t i = 0; i < n; i++) {
         b += mins[i];
@@ -97,6 +112,7 @@ void round_uint8_per_column_multi(
         float* a_out,
         float* b_out) {
     float max_span = 0;
+    float total_span = 0;
     std::vector<float> mins(n);
     for (size_t i = 0; i < n; i++) {
         float min_i = HUGE_VAL;
@@ -110,8 +126,9 @@ void round_uint8_per_column_multi(
         if (span > max_span) {
             max_span = span;
         }
+        total_span += span;
     }
-    float a = max_span > 0 ? 255.0f / max_span : 0.0f;
+    float a = fastscan_lut_scale(max_span, total_span, n);
     float b = 0;
     for (size_t i = 0; i < n; i++) {
         b += mins[i];
@@ -155,12 +172,7 @@ void quantize_LUT_and_bias(
             max_span_dis += span;
             b += mins[i];
         }
-        a = std::min(
-                max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF,
-                max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF);
-        if (!std::isfinite(a)) {
-            a = 0.0f;
-        }
+        a = fastscan_lut_scale(max_span_LUT, max_span_dis, M);
 
         for (size_t i = 0; i < M; i++) {
             round_tab(LUT + i * ksub, ksub, a, mins[i], LUTq + i * ksub);
@@ -180,12 +192,7 @@ void quantize_LUT_and_bias(
             max_span_dis += span;
             b += mins[i];
         }
-        a = std::min(
-                max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF,
-                max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF);
-        if (!std::isfinite(a)) {
-            a = 0.0f;
-        }
+        a = fastscan_lut_scale(max_span_LUT, max_span_dis, M);
         b += bias_min;
 
         for (size_t i = 0; i < M; i++) {
@@ -219,12 +226,7 @@ void quantize_LUT_and_bias(
             b = std::min(b, b2j);
         }
 
-        a = std::min(
-                max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF,
-                max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF);
-        if (!std::isfinite(a)) {
-            a = 0.0f;
-        }
+        a = fastscan_lut_scale(max_span_LUT, max_span_dis, M);
 
         ij = 0;
         size_t ij_2 = 0;
@@ -266,13 +268,15 @@ void quantize_LUT_and_bias(
         }
 
         float max_span = -HUGE_VAL;
+        float total_span = 0;
         b = 0;
         for (size_t i = 0; i < M; i++) {
             float span = maxs[i] - mins[i];
             max_span = std::max(max_span, span);
+            total_span += span;
             b += mins[i];
         }
-        a = max_span > 0 ? 255.0f / max_span : 0.0f;
+        a = fastscan_lut_scale(max_span, total_span, M);
         ij = 0;
         size_t ij_2 = 0;
         for (size_t j = 0; j < nprobe; j++) {
@@ -321,12 +325,7 @@ void aq_quantize_LUT_and_bias(
         max_span_dis += (i >= M - M_norm ? span * norm_scale : span);
         b += mins[i];
     }
-    a = std::min(
-            max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF,
-            max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF);
-    if (!std::isfinite(a)) {
-        a = 0.0f;
-    }
+    a = fastscan_lut_scale(max_span_LUT, max_span_dis, M);
     b += bias_min;
 
     for (size_t i = 0; i < M; i++) {

data/vendor/faiss/faiss/utils/quantize_lut.h

@@ -7,6 +7,8 @@
 
 #pragma once
 
+#include <cstddef>
+
 #include <cstdint>
 #include <cstdio>
 
@@ -25,6 +27,15 @@ namespace faiss {
 
 namespace quantize_lut {
 
+/** Scale for a fast scan LUT, bounding both the uint8 entry and the uint16
+ *  sum the kernel accumulates. `n_rounded_terms` is how many rounded values
+ *  land in one accumulator: the columns, plus one for a rounded bias term.
+ */
+float fastscan_lut_scale(
+        float max_column_span,
+        float max_sum_span,
+        size_t n_rounded_terms);
+
 /* affine quantizer, a and b are the affine coefficients, marginalize over d
  *
  * @param tab input/output, size (n, d)

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

@@ -53,6 +53,28 @@ BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount(
         size_t size,
         size_t qb);
 
+/// Four independent AND-dot-product/popcount pairs for a four-plane query.
+/// Query contains 4 * size bytes; each data pointer contains `size` bytes.
+/// Byte-aligned buffers and arbitrary byte counts are supported.
+template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
+inline void bitwise_q4_batch_4(
+        const uint8_t* query,
+        const uint8_t* const* data,
+        size_t size,
+        BitwiseAndDotProductResult* results) {
+    for (size_t i = 0; i < 4; ++i) {
+        results[i] = bitwise_and_dot_product_with_popcount<SL>(
+                query, data[i], size, 4);
+    }
+}
+
+template <>
+void bitwise_q4_batch_4<SIMDLevel::AVX512_VPOPCNT>(
+        const uint8_t* query,
+        const uint8_t* const* data,
+        size_t size,
+        BitwiseAndDotProductResult* results);
+
 /**
  * Compute dot product between query and binary data using popcount on XOR.
  *
@@ -140,9 +162,11 @@ inline uint64_t bitwise_and_dot_product<SIMDLevel::NONE>(
     uint64_t sum = 0;
     size_t offset = 0;
     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;
         }
     }
@@ -167,10 +191,12 @@ inline BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount<
     uint64_t popcount_sum = 0;
     size_t offset = 0;
     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));
         popcount_sum += popcount64(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));
             dot_product += popcount64(qv & yv) << j;
         }
     }
@@ -194,9 +220,11 @@ inline uint64_t bitwise_xor_dot_product<SIMDLevel::NONE>(
     uint64_t sum = 0;
     size_t offset = 0;
     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;
         }
     }
@@ -215,7 +243,8 @@ inline uint64_t popcount<SIMDLevel::NONE>(const uint8_t* data, size_t size) {
     uint64_t sum = 0;
     size_t offset = 0;
     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) {
@@ -332,16 +361,12 @@ inline void quantize_query_values<SIMDLevel::NONE>(
 /*********************************************************
  * Multi-bit RaBitQ inner product kernels.
  *
- * Compute: sum_i rotated_q[i] * ((sign_bit_i << ex_bits) + ex_code_val_i + cb)
+ * Score = sum_i rotated_q[i] * ((sign_i << ex_bits) + ex_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.
+ * ex_bits == 1: bits map straight to floats. No extraction.
+ * ex_bits 2..7: BMI2 PEXT extracts each bit plane.
+ * ex_bits == 8: whole bytes. No extraction and no BMI2.
+ * Fallback:     scalar math reads one code at a time.
  *********************************************************/
 namespace faiss::rabitq::multibit {
 
@@ -357,6 +382,16 @@ inline float ip_scalar(
         size_t ex_bits,
         float cb) {
     float result = 0.0f;
+    if (ex_bits == 8) {
+        // RBQ9 is byte-aligned, including the last dimension. Do not require
+        // trailing factor bytes for the scalar reference or SIMD tail.
+        for (size_t i = start; i < d; ++i) {
+            const int sb = (sign_bits[i / 8] >> (i % 8)) & 1;
+            result += rotated_q[i] *
+                    (static_cast<float>((sb << 8) + ex_code[i]) + cb);
+        }
+        return result;
+    }
     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++) {

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

@@ -6,8 +6,7 @@
  */
 
 #include <faiss/utils/distances.h>
-
-#ifdef __aarch64__
+#include <faiss/utils/simd_impl/masked_sum.h>
 
 #include <arm_neon.h>
 #include <limits>
@@ -41,6 +40,14 @@ void fvec_madd<SIMDLevel::ARM_NEON>(
         c[i] = a[i] + bf * b[i];
 }
 
+template <>
+float fvec_masked_sum<SIMDLevel::ARM_NEON>(
+        const float* x,
+        const uint8_t* bits,
+        size_t d) {
+    return masked_sum<SIMDLevel::ARM_NEON>(x, bits, d);
+}
+
 template <>
 void fvec_L2sqr_ny_transposed<SIMDLevel::ARM_NEON>(
         float* dis,
@@ -150,5 +157,3 @@ int fvec_madd_and_argmin<SIMDLevel::ARM_NEON>(
 }
 
 } // namespace faiss
-
-#endif // __aarch64__

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

@@ -758,9 +758,11 @@ size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::ARM_SVE>(
     return current_min_idx;
 }
 
-template <>
-void exhaustive_L2sqr_blas_cmax<SIMDLevel::ARM_SVE>(
-        const float* x,
+namespace {
+
+template <class QueryTiles>
+void exhaustive_L2sqr_blas_cmax_sve(
+        QueryTiles& queries,
         const float* y,
         size_t d,
         size_t nx,
@@ -779,7 +781,7 @@ void exhaustive_L2sqr_blas_cmax<SIMDLevel::ARM_SVE>(
     std::unique_ptr<float[]> x_norms(new float[nx]);
     std::unique_ptr<float[]> del2;
 
-    fvec_norms_L2sqr(x_norms.get(), x, d, nx);
+    queries.prepare_norms(nx, x_norms.get());
 
     const size_t lanes = svcntw();
 
@@ -795,6 +797,8 @@ void exhaustive_L2sqr_blas_cmax<SIMDLevel::ARM_SVE>(
         if (i1 > nx)
             i1 = nx;
 
+        const float* xt = queries.tile(i0, i1, x_norms.get());
+
         res.begin_multiple(i0, i1);
 
         for (size_t j0 = 0; j0 < ny; j0 += bs_y) {
@@ -813,7 +817,7 @@ void exhaustive_L2sqr_blas_cmax<SIMDLevel::ARM_SVE>(
                        &one,
                        y + j0 * d,
                        &di,
-                       x + i0 * d,
+                       xt,
                        &di,
                        &zero,
                        ip_block.get(),
@@ -938,4 +942,32 @@ void exhaustive_L2sqr_blas_cmax<SIMDLevel::ARM_SVE>(
     }
 }
 
+} // namespace
+
+template <>
+void exhaustive_L2sqr_blas_cmax<SIMDLevel::ARM_SVE>(
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t nx,
+        size_t ny,
+        Top1BlockResultHandler<CMax<float, int64_t>>& res,
+        const float* y_norms) {
+    Fp32QueryTiles queries{x, d};
+    exhaustive_L2sqr_blas_cmax_sve(queries, y, d, nx, ny, res, y_norms);
+}
+
+template <>
+void exhaustive_L2sqr_blas_cmax_fp16<SIMDLevel::ARM_SVE>(
+        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) {
+    Fp16QueryTiles queries(x, d, nx);
+    exhaustive_L2sqr_blas_cmax_sve(queries, y, d, nx, ny, res, y_norms);
+}
+
 } // namespace faiss