diff --git a/faiss/IndexFlat.cpp b/faiss/IndexFlat.cpp index 9ccc4c8ed3..039eb3ca8b 100644 --- a/faiss/IndexFlat.cpp +++ b/faiss/IndexFlat.cpp @@ -269,10 +269,15 @@ struct FlatIPDis : FlatCodesDistanceComputer { FlatCodesDistanceComputer* IndexFlat::get_FlatCodesDistanceComputer() const { FlatCodesDistanceComputer* dc = nullptr; + // BASE_WITH_SVE rather than plain with_simd_level: the fvec_* kernels + // these computers call have dedicated ARM_SVE specializations, which the + // BASE mask would skip in favour of ARM_NEON. if (metric_type == METRIC_L2) { - with_simd_level([&]() { dc = new FlatL2Dis(*this); }); + with_simd_level_with_sve( + [&]() { dc = new FlatL2Dis(*this); }); } else if (metric_type == METRIC_INNER_PRODUCT) { - with_simd_level([&]() { dc = new FlatIPDis(*this); }); + with_simd_level_with_sve( + [&]() { dc = new FlatIPDis(*this); }); } else { dc = get_extra_distance_computer(d, metric_type, metric_arg, get_xb()); } @@ -404,7 +409,7 @@ FlatCodesDistanceComputer* IndexFlatL2::get_FlatCodesDistanceComputer() const { if (metric_type == METRIC_L2) { if (!cached_l2norms.empty()) { FlatCodesDistanceComputer* dc = nullptr; - with_simd_level([&]() { + with_simd_level_with_sve([&]() { dc = new FlatL2WithNormsDis(*this); }); return dc; @@ -784,7 +789,9 @@ void IndexFlatPanorama::search_subset( idx_t k, float* distances, idx_t* labels) const { - with_simd_level([&]() { + // BASE_WITH_SVE: the only SIMD-templated call in the loop body is + // fvec_inner_product, which has an ARM_SVE specialization. + with_simd_level_with_sve([&]() { with_metric_type(metric_type, [&]() { constexpr bool is_sim = is_similarity_metric(M); using C = std::conditional_t< diff --git a/faiss/SuperKMeans.cpp b/faiss/SuperKMeans.cpp index db9ad2b087..c457a2f2d4 100644 --- a/faiss/SuperKMeans.cpp +++ b/faiss/SuperKMeans.cpp @@ -620,8 +620,11 @@ void super_kmeans_assign_iteration( &ldc); } - // One SIMD dispatch per (xi, yj) tile. - with_simd_level([&]() { + // One SIMD dispatch per (xi, yj) tile. BASE_WITH_SVE rather than + // plain with_simd_level: block_l2 has a dedicated ARM_SVE kernel, + // which the BASE mask would skip in favour of ARM_NEON (no + // specialization, so the scalar primary template). + with_simd_level_with_sve([&]() { [[maybe_unused]] const int omp_chunk_local = cp.omp_chunk; int64_t tile_total = 0; int64_t tile_pruned = 0; diff --git a/faiss/docs/simd_dynamic_dispatch_migration.md b/faiss/docs/simd_dynamic_dispatch_migration.md index 3da8bcb9b3..0bf516e762 100644 --- a/faiss/docs/simd_dynamic_dispatch_migration.md +++ b/faiss/docs/simd_dynamic_dispatch_migration.md @@ -238,8 +238,9 @@ your own with `(1 << int(SIMDLevel::X)) | ...`): |------|--------|---------| | `AVAILABLE_SIMD_LEVELS_NONE` | NONE only | Scalar-only functions | | `AVAILABLE_SIMD_LEVELS_AVX2_NEON` | NONE, AVX2, ARM_NEON | 256-bit `simdlib` ops (`with_simd_level_256bit`) | -| `AVAILABLE_SIMD_LEVELS_A0` | NONE, AVX2, AVX512, ARM_NEON, RISCV_RVV | Default (`with_simd_level`) | -| `AVAILABLE_SIMD_LEVELS_A1` | A0 + ARM_SVE | Functions with dedicated SVE implementations | +| `AVAILABLE_SIMD_LEVELS_BASE` | NONE, AVX2, AVX512, ARM_NEON, RISCV_RVV | Default (`with_simd_level`). ARM_NEON is part of BASE: NEON is mandatory on aarch64, SVE is optional | +| `AVAILABLE_SIMD_LEVELS_BASE_WITH_SPR` | BASE + AVX512_SPR | Kernels needing the whole SPR feature set (`with_simd_level_with_spr`) | +| `AVAILABLE_SIMD_LEVELS_BASE_WITH_SVE` | BASE + ARM_SVE | Functions with dedicated SVE implementations (`with_simd_level_with_sve`) | | `AVAILABLE_SIMD_LEVELS_ALL` | All levels | Identity / diagnostic functions | ### Step 4: Register in build system diff --git a/faiss/impl/AdditiveQuantizer.cpp b/faiss/impl/AdditiveQuantizer.cpp index cab5c4c808..e55a3879f0 100644 --- a/faiss/impl/AdditiveQuantizer.cpp +++ b/faiss/impl/AdditiveQuantizer.cpp @@ -376,7 +376,8 @@ void AdditiveQuantizer::compute_centroid_norms(float* norms) const { size_t ntotal = (size_t)1 << tot_bits; int64_t ntotal_signed = ntotal; // TODO: make tree of partial sums - with_simd_level([&]() { + // BASE_WITH_SVE: fvec_norm_L2sqr has an ARM_SVE specialization. + with_simd_level_with_sve([&]() { #pragma omp parallel { std::vector tmp(d); diff --git a/faiss/impl/ProductQuantizer.cpp b/faiss/impl/ProductQuantizer.cpp index 1190b52574..eb3bac6093 100644 --- a/faiss/impl/ProductQuantizer.cpp +++ b/faiss/impl/ProductQuantizer.cpp @@ -280,8 +280,9 @@ void compute_1_code(const ProductQuantizer& pq, const float* x, uint8_t* code) { } // namespace void ProductQuantizer::compute_code(const float* x, uint8_t* code) const { - // a1: fvec_L2sqr_ny_nearest / _y_transposed have ARM_SVE specializations - with_simd_level_a1([&]() { + // with_sve: fvec_L2sqr_ny_nearest / _y_transposed have ARM_SVE + // specializations + with_simd_level_with_sve([&]() { switch (nbits) { case 8: compute_1_code(*this, x, code); @@ -295,7 +296,7 @@ void ProductQuantizer::compute_code(const float* x, uint8_t* code) const { compute_1_code(*this, x, code); break; } - }); // with_simd_level_a1 + }); // with_simd_level_with_sve } template @@ -443,8 +444,8 @@ void ProductQuantizer::compute_codes(const float* x, uint8_t* codes, size_t n) void ProductQuantizer::compute_distance_table(const float* x, float* dis_table) const { - // a1: fvec_L2sqr_ny / _transposed have ARM_SVE specializations - with_simd_level_a1([&]() { + // with_sve: fvec_L2sqr_ny / _transposed have ARM_SVE specializations + with_simd_level_with_sve([&]() { if (transposed_centroids.empty()) { // use regular version for (size_t m = 0; m < M; m++) { @@ -826,8 +827,8 @@ void ProductQuantizer::compute_sdc_table() { sdc_table.resize(M * ksub * ksub); if (dsub < 4) { - // a1: fvec_L2sqr_ny has an ARM_SVE specialization - with_simd_level_a1([&]() { + // with_sve: fvec_L2sqr_ny has an ARM_SVE specialization + with_simd_level_with_sve([&]() { #pragma omp parallel for for (int64_t mk = 0; mk < static_cast(M * ksub); mk++) { // allow omp to schedule in a more fine-grained way diff --git a/faiss/impl/RaBitQUtils.cpp b/faiss/impl/RaBitQUtils.cpp index 051874300e..f2eb77c7e3 100644 --- a/faiss/impl/RaBitQUtils.cpp +++ b/faiss/impl/RaBitQUtils.cpp @@ -329,7 +329,7 @@ float compute_full_multibit_distance( size_t d, size_t ex_bits, MetricType metric_type) { - return with_selected_simd_levels( + return with_selected_simd_levels( [&]() { return compute_full_multibit_distance( sign_bits, diff --git a/faiss/impl/RaBitQuantizer.cpp b/faiss/impl/RaBitQuantizer.cpp index 90c524f6f7..af5d563ad2 100644 --- a/faiss/impl/RaBitQuantizer.cpp +++ b/faiss/impl/RaBitQuantizer.cpp @@ -665,12 +665,12 @@ FlatCodesDistanceComputer* RaBitQuantizer::get_distance_computer( // call the SIMD-specialized rabitq functions directly (no per-call // with_simd_level overhead). // - // Use A0_SPR (which includes AVX512_SPR) so that on Sapphire Rapids + // Use BASE_WITH_SPR (which includes AVX512_SPR) so that on Sapphire Rapids // and later x86 microarchitectures the VPOPCNTDQ-based RaBitQ // specialization in rabitq_avx512_spr.cpp is selected. On AVX-512 // CPUs without VPOPCNTDQ, dispatch falls through to the AVX512 // specialization in rabitq_avx512.cpp. - return with_selected_simd_levels( + return with_selected_simd_levels( [&]() -> FlatCodesDistanceComputer* { if (qb == 0) { auto dc = diff --git a/faiss/impl/ScalarQuantizer.cpp b/faiss/impl/ScalarQuantizer.cpp index 11719e342e..0918d7ddb1 100644 --- a/faiss/impl/ScalarQuantizer.cpp +++ b/faiss/impl/ScalarQuantizer.cpp @@ -665,7 +665,7 @@ ScalarQuantizer::SQuantizer* ScalarQuantizer::select_quantizer() const { // A SIMD level's factory returns nullptr when the dimension is // incompatible (e.g. AVX-512 needs d % 16 == 0); the dispatcher then falls // back to the next-lower level (AVX-512 -> AVX2 -> scalar). - return with_simd_level_fallback( + return with_simd_level_fallback( [&]() -> SQuantizer* { return scalar_quantizer::sq_select_quantizer( qtype, d, trained); @@ -702,7 +702,7 @@ void ScalarQuantizer::decode(const uint8_t* codes, float* x, size_t n) const { ScalarQuantizer::SQDistanceComputer* ScalarQuantizer::get_distance_computer( MetricType metric) const { FAISS_THROW_IF_NOT(metric == METRIC_L2 || metric == METRIC_INNER_PRODUCT); - return with_simd_level_fallback( + return with_simd_level_fallback( [&]() -> SQDistanceComputer* { return scalar_quantizer::sq_select_distance_computer( metric, qtype, d, trained); @@ -715,7 +715,7 @@ InvertedListScanner* ScalarQuantizer::select_InvertedListScanner( bool store_pairs, const IDSelector* sel, bool by_residual) const { - return with_simd_level_fallback( + return with_simd_level_fallback( [&]() -> InvertedListScanner* { return scalar_quantizer::sq_select_InvertedListScanner( qtype, diff --git a/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp b/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp index d3a33dbea9..0b017bbc0e 100644 --- a/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp +++ b/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp @@ -7,13 +7,17 @@ // This TU provides non-templated PQ code distance dispatch wrappers // (pq_code_distance_8bit_single, pq_code_distance_8bit_four) declared -// in pq_code_distance-inl.h. These use with_simd_level to route to the +// in pq_code_distance-inl.h. These use with_simd_level_with_sve to route to the // best available SIMD implementation via pq_code_distance_8bit_*_impl // function template specializations. // // The NONE and ARM_NEON _impl specializations are defined inline in // pq_code_distance-generic.h (included transitively). The AVX2, AVX512, // and ARM_SVE specializations are in their respective per-SIMD files. +// +// BASE_WITH_SVE rather than plain with_simd_level: all three _impl +// functions have dedicated ARM_SVE specializations, which the BASE mask +// would skip in favour of ARM_NEON. #include @@ -34,7 +38,7 @@ void pq_scan_8bit( float* heap_dis, int64_t* heap_ids, bool max_heap) { - with_simd_level([&]() { + with_simd_level_with_sve([&]() { pq_scan_8bit_impl( M, dis_table, codes, ncodes, k, heap_dis, heap_ids, max_heap); }); @@ -44,7 +48,7 @@ float pq_code_distance_8bit_single( size_t M, const float* sim_table, const uint8_t* code) { - return with_simd_level([&]() { + return with_simd_level_with_sve([&]() { return pq_code_distance_8bit_single_impl(M, sim_table, code); }); } @@ -60,7 +64,7 @@ void pq_code_distance_8bit_four( float& result1, float& result2, float& result3) { - with_simd_level([&]() { + with_simd_level_with_sve([&]() { pq_code_distance_8bit_four_impl( M, sim_table, diff --git a/faiss/impl/scalar_quantizer/distance_computers.h b/faiss/impl/scalar_quantizer/distance_computers.h index e233e1cd72..4c9d61cb4f 100644 --- a/faiss/impl/scalar_quantizer/distance_computers.h +++ b/faiss/impl/scalar_quantizer/distance_computers.h @@ -78,8 +78,11 @@ template struct DistanceComputerByte : SQDistanceComputer {}; // Byte-domain distance computer for QT_8bit_direct_signed (storage is -// value+128). Only specialized for AVX512_SPR; other levels fall back to -// the float-domain DCTemplate path via the dispatch logic. +// value+128). Specialized for AVX2, AVX512, AVX512_SPR and ARM_NEON; levels +// without a specialization fall back to the float-domain DCTemplate path via +// the dispatch logic. Adding a level to the dispatch chain in sq-dispatch.h +// without a specialization here instantiates this empty primary template, +// whose pure virtuals are left unimplemented. template struct DistanceComputerByteSigned : SQDistanceComputer {}; diff --git a/faiss/impl/scalar_quantizer/sq-dispatch.h b/faiss/impl/scalar_quantizer/sq-dispatch.h index e42285c063..0dfdd830e2 100644 --- a/faiss/impl/scalar_quantizer/sq-dispatch.h +++ b/faiss/impl/scalar_quantizer/sq-dispatch.h @@ -531,7 +531,8 @@ SQDistanceComputer* select_distance_computer_body( return new DistanceComputerByte( static_cast(d), trained); } - } else if constexpr (SL2 == SIMDLevel::AVX2) { + } else if constexpr ( + SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) { if (d % 16 == 0) { return new DistanceComputerByte( static_cast(d), trained); @@ -551,7 +552,8 @@ SQDistanceComputer* select_distance_computer_body( return new DistanceComputerByteSigned( static_cast(d), trained); } - } else if constexpr (SL2 == SIMDLevel::AVX2) { + } else if constexpr ( + SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) { if (d % 16 == 0) { return new DistanceComputerByteSigned( static_cast(d), trained); @@ -744,7 +746,8 @@ InvertedListScanner* sq_select_InvertedListScanner( return scan.template operator()>(); } - } else if constexpr (SL2 == SIMDLevel::AVX2) { + } else if constexpr ( + SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) { if (d % 16 == 0) { return scan.template operator()>(); @@ -765,7 +768,8 @@ InvertedListScanner* sq_select_InvertedListScanner( return scan.template operator()< DistanceComputerByteSigned>(); } - } else if constexpr (SL2 == SIMDLevel::AVX2) { + } else if constexpr ( + SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) { if (d % 16 == 0) { return scan.template operator()< DistanceComputerByteSigned>(); diff --git a/faiss/impl/scalar_quantizer/sq-neon.cpp b/faiss/impl/scalar_quantizer/sq-neon.cpp index eb58673b58..802f285877 100644 --- a/faiss/impl/scalar_quantizer/sq-neon.cpp +++ b/faiss/impl/scalar_quantizer/sq-neon.cpp @@ -618,6 +618,73 @@ struct DCTemplate } }; +/********************************************************** + * Byte-domain (integer) kernels for QT_8bit_direct{,_signed} + * + * These stay in the integer domain end to end, unlike the DCTemplate path + * which widens each byte to float. They are only reached when d % 16 == 0 + * (enforced by the dispatch in sq-dispatch.h), so no loop needs a tail. + **********************************************************/ + +namespace { + +/// Squared L2 between two unsigned byte codes. |a-b| then squared is exact +/// and avoids needing a signed intermediate: vmull_u8 tops out at +/// 255*255 = 65025, which still fits a uint16 lane. +FAISS_ALWAYS_INLINE int neon_byte_l2sqr( + const uint8_t* code1, + const uint8_t* code2, + int d) { + uint32x4_t accu = vdupq_n_u32(0); + for (int i = 0; i < d; i += 16) { + const uint8x16_t diff = + vabdq_u8(vld1q_u8(code1 + i), vld1q_u8(code2 + i)); + accu = vpadalq_u16( + accu, vmull_u8(vget_low_u8(diff), vget_low_u8(diff))); + accu = vpadalq_u16( + accu, vmull_u8(vget_high_u8(diff), vget_high_u8(diff))); + } + return static_cast(vaddvq_u32(accu)); +} + +/// Inner product of two unsigned byte codes. +FAISS_ALWAYS_INLINE int neon_byte_ip( + const uint8_t* code1, + const uint8_t* code2, + int d) { + uint32x4_t accu = vdupq_n_u32(0); + for (int i = 0; i < d; i += 16) { + const uint8x16_t c1 = vld1q_u8(code1 + i); + const uint8x16_t c2 = vld1q_u8(code2 + i); + accu = vpadalq_u16(accu, vmull_u8(vget_low_u8(c1), vget_low_u8(c2))); + accu = vpadalq_u16(accu, vmull_u8(vget_high_u8(c1), vget_high_u8(c2))); + } + return static_cast(vaddvq_u32(accu)); +} + +/// Inner product of two bias-encoded byte codes, which store value + 128. +/// The bias has to come off before multiplying; for x in [0, 255], +/// x ^ 0x80 reinterpreted as int8 is exactly x - 128. vmull_s8 then tops +/// out at (-128)*(-128) = 16384, well inside an int16 lane. +FAISS_ALWAYS_INLINE int neon_byte_ip_unbias( + const uint8_t* code1, + const uint8_t* code2, + int d) { + const uint8x16_t bias = vdupq_n_u8(0x80); + int32x4_t accu = vdupq_n_s32(0); + for (int i = 0; i < d; i += 16) { + const int8x16_t c1 = + vreinterpretq_s8_u8(veorq_u8(vld1q_u8(code1 + i), bias)); + const int8x16_t c2 = + vreinterpretq_s8_u8(veorq_u8(vld1q_u8(code2 + i), bias)); + accu = vpadalq_s16(accu, vmull_s8(vget_low_s8(c1), vget_low_s8(c2))); + accu = vpadalq_s16(accu, vmull_s8(vget_high_s8(c1), vget_high_s8(c2))); + } + return static_cast(vaddvq_s32(accu)); +} + +} // namespace + template struct DistanceComputerByte : SQDistanceComputer { @@ -630,16 +697,11 @@ struct DistanceComputerByte int compute_code_distance(const uint8_t* code1, const uint8_t* code2) const { - int accu = 0; - for (int i = 0; i < d; i++) { - if (Sim::metric_type == METRIC_INNER_PRODUCT) { - accu += int(code1[i]) * code2[i]; - } else { - int diff = int(code1[i]) - code2[i]; - accu += diff * diff; - } + if constexpr (Sim::metric_type == METRIC_INNER_PRODUCT) { + return neon_byte_ip(code1, code2, d); + } else { + return neon_byte_l2sqr(code1, code2, d); } - return accu; } void set_query(const float* x) final { @@ -663,6 +725,49 @@ struct DistanceComputerByte } }; +template +struct DistanceComputerByteSigned + : SQDistanceComputer { + using Sim = Similarity; + + int d; + std::vector tmp; + + DistanceComputerByteSigned(int d, const std::vector&) + : d(d), tmp(d) {} + + int compute_code_distance(const uint8_t* code1, const uint8_t* code2) + const { + if constexpr (Sim::metric_type == METRIC_INNER_PRODUCT) { + return neon_byte_ip_unbias(code1, code2, d); + } else { + // The +128 bias is common to both codes, so it cancels in the + // difference and the unsigned kernel is already exact. + return neon_byte_l2sqr(code1, code2, d); + } + } + + void set_query(const float* x) final { + for (int i = 0; i < d; i++) { + tmp[i] = uint8_t(int(x[i]) + 128); + } + } + + int compute_distance(const float* x, const uint8_t* code) { + set_query(x); + return compute_code_distance(tmp.data(), code); + } + + float symmetric_dis(idx_t i, idx_t j) override { + return compute_code_distance( + codes + i * code_size, codes + j * code_size); + } + + float query_to_code(const uint8_t* code) const final { + return compute_code_distance(tmp.data(), code); + } +}; + /********************************************************** * TurboQuant masked_sum NEON specialization (scalar fallback) **********************************************************/ diff --git a/faiss/impl/simd_dispatch.h b/faiss/impl/simd_dispatch.h index cbfe9d72db..aa0b99c19d 100644 --- a/faiss/impl/simd_dispatch.h +++ b/faiss/impl/simd_dispatch.h @@ -32,24 +32,23 @@ constexpr int AVAILABLE_SIMD_LEVELS_NONE = (1 << int(SIMDLevel::NONE)); constexpr int AVAILABLE_SIMD_LEVELS_AVX2_NEON = AVAILABLE_SIMD_LEVELS_NONE | (1 << int(SIMDLevel::AVX2)) | (1 << int(SIMDLevel::ARM_NEON)); -// A0: same + AVX512 + RISCV_RVV -constexpr int AVAILABLE_SIMD_LEVELS_A0 = AVAILABLE_SIMD_LEVELS_AVX2_NEON | +// BASE: the default mask. Same as AVX2_NEON + AVX512 + RISCV_RVV, i.e. +// NONE + AVX2 + AVX512 + ARM_NEON + RISCV_RVV. ARM_NEON is part of BASE: +// NEON is mandatory on aarch64 (the ARM baseline, like AVX2 on x86) while +// ARM_SVE is optional, so the ARM fallback chain is SVE -> NEON -> NONE. +constexpr int AVAILABLE_SIMD_LEVELS_BASE = AVAILABLE_SIMD_LEVELS_AVX2_NEON | (1 << int(SIMDLevel::AVX512)) | (1 << int(SIMDLevel::RISCV_RVV)); -// A0_SPR: same as A0 + AVX512_SPR (for functions with a dedicated SPR +// BASE_WITH_SPR: BASE + AVX512_SPR (for functions with a dedicated SPR // specialization on top of an AVX512 fallback). Currently used by the // RaBitQ popcount kernels, which use VPOPCNTDQ on SPR+. -constexpr int AVAILABLE_SIMD_LEVELS_A0_SPR = - AVAILABLE_SIMD_LEVELS_A0 | (1 << int(SIMDLevel::AVX512_SPR)); +constexpr int AVAILABLE_SIMD_LEVELS_BASE_WITH_SPR = + AVAILABLE_SIMD_LEVELS_BASE | (1 << int(SIMDLevel::AVX512_SPR)); -// A1: same + ARM_SVE (for functions with dedicated SVE implementations) -constexpr int AVAILABLE_SIMD_LEVELS_A1 = - AVAILABLE_SIMD_LEVELS_A0 | (1 << int(SIMDLevel::ARM_SVE)); - -// A2: NONE + AVX2 + ARM_SVE only (for functions with only these +// BASE_WITH_SVE: BASE + ARM_SVE (for functions with dedicated SVE // implementations) -constexpr int AVAILABLE_SIMD_LEVELS_A2 = AVAILABLE_SIMD_LEVELS_NONE | - (1 << int(SIMDLevel::AVX2)) | (1 << int(SIMDLevel::ARM_SVE)); +constexpr int AVAILABLE_SIMD_LEVELS_BASE_WITH_SVE = + AVAILABLE_SIMD_LEVELS_BASE | (1 << int(SIMDLevel::ARM_SVE)); constexpr int AVAILABLE_SIMD_LEVELS_ALL = -1; @@ -224,7 +223,7 @@ inline auto with_simd_level_fallback(const LambdaType& action) { */ template inline auto with_simd_level(LambdaType&& action) { - return with_selected_simd_levels( + return with_selected_simd_levels( std::forward(action)); } @@ -239,25 +238,24 @@ inline auto with_simd_level_256bit(LambdaType&& action) { } /** - * Use for functions that have A0-level implementations plus a dedicated - * ARM_SVE specialization. Plain with_simd_level() uses A0, which omits the + * Use for functions that have BASE-level implementations plus a dedicated + * ARM_SVE specialization. Plain with_simd_level() uses BASE, which omits the * ARM_SVE bit, so on an SVE host the ARM_SVE case falls through to ARM_NEON * and the SVE specialization is never instantiated. */ template -inline auto with_simd_level_a1(LambdaType&& action) { - return with_selected_simd_levels( +inline auto with_simd_level_with_sve(LambdaType&& action) { + return with_selected_simd_levels( std::forward(action)); } /** - * Use for functions that have A0-level implementations plus an AVX512_SPR + * Use for functions that have BASE-level implementations plus an AVX512_SPR * specialization (e.g. using VPOPCNTDQ). */ template -inline auto with_simd_level_a0_spr(LambdaType&& action) { - return with_selected_simd_levels( +inline auto with_simd_level_with_spr(LambdaType&& action) { + return with_selected_simd_levels( std::forward(action)); } - } // namespace faiss diff --git a/faiss/utils/distances.cpp b/faiss/utils/distances.cpp index 60fc29d0c4..2702b8ffcc 100644 --- a/faiss/utils/distances.cpp +++ b/faiss/utils/distances.cpp @@ -541,17 +541,19 @@ void exhaustive_L2sqr_blas>>( return; } - with_selected_simd_levels([&]() { - if constexpr ( - SL == SIMDLevel::AVX2 || SL == SIMDLevel::AVX512 || - SL == SIMDLevel::ARM_SVE) { - exhaustive_L2sqr_blas_cmax(x, y, d, nx, ny, res, y_norms); - } else { - exhaustive_L2sqr_blas_default_impl< - Top1BlockResultHandler>>( - x, y, d, nx, ny, res, y_norms); - } - }); + with_selected_simd_levels( + [&]() { + if constexpr ( + SL == SIMDLevel::AVX2 || SL == SIMDLevel::AVX512 || + SL == SIMDLevel::ARM_SVE) { + exhaustive_L2sqr_blas_cmax( + x, y, d, nx, ny, res, y_norms); + } else { + exhaustive_L2sqr_blas_default_impl< + Top1BlockResultHandler>>( + x, y, d, nx, ny, res, y_norms); + } + }); } struct Run_search_inner_product { diff --git a/faiss/utils/distances_dispatch.h b/faiss/utils/distances_dispatch.h index 3fb916b87f..1cb8e338c0 100644 --- a/faiss/utils/distances_dispatch.h +++ b/faiss/utils/distances_dispatch.h @@ -29,22 +29,22 @@ namespace faiss { inline float fvec_L1_dispatch(const float* x, const float* y, size_t d) { - return with_selected_simd_levels( + return with_selected_simd_levels( [&]() { return fvec_L1(x, y, d); }); } inline float fvec_Linf_dispatch(const float* x, const float* y, size_t d) { - return with_selected_simd_levels( + return with_selected_simd_levels( [&]() { return fvec_Linf(x, y, d); }); } inline float fvec_norm_L2sqr_dispatch(const float* x, size_t d) { - return with_selected_simd_levels( + return with_selected_simd_levels( [&]() { return fvec_norm_L2sqr(x, d); }); } inline float fvec_L2sqr_dispatch(const float* x, const float* y, size_t d) { - return with_selected_simd_levels( + return with_selected_simd_levels( [&]() { return fvec_L2sqr(x, y, d); }); } @@ -52,7 +52,7 @@ inline float fvec_inner_product_dispatch( const float* x, const float* y, size_t d) { - return with_selected_simd_levels( + return with_selected_simd_levels( [&]() { return fvec_inner_product(x, y, d); }); } @@ -67,10 +67,11 @@ inline void fvec_inner_product_batch_4_dispatch( float& dis1, float& dis2, float& dis3) { - with_selected_simd_levels([&]() { - fvec_inner_product_batch_4( - x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3); - }); + with_selected_simd_levels( + [&]() { + fvec_inner_product_batch_4( + x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3); + }); } inline void fvec_L2sqr_batch_4_dispatch( @@ -84,9 +85,11 @@ inline void fvec_L2sqr_batch_4_dispatch( float& dis1, float& dis2, float& dis3) { - with_selected_simd_levels([&]() { - fvec_L2sqr_batch_4(x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3); - }); + with_selected_simd_levels( + [&]() { + fvec_L2sqr_batch_4( + x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3); + }); } inline void fvec_L2sqr_ny_transposed_dispatch( @@ -97,9 +100,11 @@ inline void fvec_L2sqr_ny_transposed_dispatch( size_t d, size_t d_offset, size_t ny) { - with_selected_simd_levels([&]() { - fvec_L2sqr_ny_transposed(dis, x, y, y_sqlen, d, d_offset, ny); - }); + with_selected_simd_levels( + [&]() { + fvec_L2sqr_ny_transposed( + dis, x, y, y_sqlen, d, d_offset, ny); + }); } inline void fvec_inner_products_ny_dispatch( @@ -108,9 +113,10 @@ inline void fvec_inner_products_ny_dispatch( const float* y, size_t d, size_t ny) { - with_selected_simd_levels([&]() { - fvec_inner_products_ny(ip, x, y, d, ny); - }); + with_selected_simd_levels( + [&]() { + fvec_inner_products_ny(ip, x, y, d, ny); + }); } inline void fvec_L2sqr_ny_dispatch( @@ -119,7 +125,7 @@ inline void fvec_L2sqr_ny_dispatch( const float* y, size_t d, size_t ny) { - with_selected_simd_levels( + with_selected_simd_levels( [&]() { fvec_L2sqr_ny(dis, x, y, d, ny); }); } @@ -129,7 +135,7 @@ inline size_t fvec_L2sqr_ny_nearest_dispatch( const float* y, size_t d, size_t ny) { - return with_selected_simd_levels( + return with_selected_simd_levels( [&]() { return fvec_L2sqr_ny_nearest( distances_tmp_buffer, x, y, d, ny); @@ -144,7 +150,7 @@ inline size_t fvec_L2sqr_ny_nearest_y_transposed_dispatch( size_t d, size_t d_offset, size_t ny) { - return with_selected_simd_levels( + return with_selected_simd_levels( [&]() { return fvec_L2sqr_ny_nearest_y_transposed( distances_tmp_buffer, x, y, y_sqlen, d, d_offset, ny); @@ -157,7 +163,7 @@ inline void fvec_madd_dispatch( float bf, const float* b, float* c) { - with_selected_simd_levels( + with_selected_simd_levels( [&]() { fvec_madd(n, a, bf, b, c); }); } @@ -167,7 +173,7 @@ inline int fvec_madd_and_argmin_dispatch( float bf, const float* b, float* c) { - return with_selected_simd_levels( + return with_selected_simd_levels( [&]() { return fvec_madd_and_argmin(n, a, bf, b, c); }); @@ -241,7 +247,11 @@ auto with_VectorDistance( if constexpr (!has_simd) { return call.template operator()(); } else { - return with_simd_level(call); + // BASE_WITH_SVE rather than plain with_simd_level: VectorDistance + // and IVFFlatScanner are both instantiated at ARM_SVE in + // distances_arm_sve.cpp, which the BASE mask would skip in favour + // of ARM_NEON. + return with_simd_level_with_sve(call); } }; return with_metric_type(metric, dispatch_metric); diff --git a/faiss/utils/distances_fused/distances_fused.cpp b/faiss/utils/distances_fused/distances_fused.cpp index fd39c11be3..fcbcd2d742 100644 --- a/faiss/utils/distances_fused/distances_fused.cpp +++ b/faiss/utils/distances_fused/distances_fused.cpp @@ -51,7 +51,7 @@ bool exhaustive_L2sqr_fused_cmax( return true; } - return with_selected_simd_levels( + return with_selected_simd_levels( [&]() { return exhaustive_L2sqr_fused_cmax( x, y, d, nx, ny, res, y_norms); diff --git a/faiss/utils/hamming.cpp b/faiss/utils/hamming.cpp index e5c3bdb7a5..870055486a 100644 --- a/faiss/utils/hamming.cpp +++ b/faiss/utils/hamming.cpp @@ -146,7 +146,7 @@ void hammings( size_t nb, size_t ncodes, hamdis_t* __restrict dis) { - with_simd_level_a0_spr([&]() { + with_simd_level_with_spr([&]() { // Ragged sizes have their own kernel; keeping it out of // hammings_fixSL() leaves the word-level paths untouched. if (ncodes % 8 != 0) { @@ -176,7 +176,7 @@ void hammings_knn_hc( int order, ApproxTopK_mode_t approx_topk_mode, const faiss::IDSelector* sel) { - with_simd_level_a0_spr([&]() { + with_simd_level_with_spr([&]() { hammings_knn_hc_fixSL( ha, a, b, nb, ncodes, order, approx_topk_mode, sel); }); @@ -192,7 +192,7 @@ void hammings_knn_mc( int32_t* __restrict distances, int64_t* __restrict labels, const faiss::IDSelector* sel) { - with_simd_level_a0_spr([&]() { + with_simd_level_with_spr([&]() { hammings_knn_mc_fixSL( a, b, na, nb, k, ncodes, distances, labels, sel); }); @@ -207,7 +207,7 @@ void hamming_range_search( size_t code_size, RangeSearchResult* result, const faiss::IDSelector* sel) { - with_simd_level_a0_spr([&]() { + with_simd_level_with_spr([&]() { hamming_range_search_fixSL( a, b, na, nb, radius, code_size, result, sel); }); @@ -221,7 +221,7 @@ void hamming_count_thres( hamdis_t ht, size_t ncodes, size_t* nptr) { - with_simd_level_a0_spr([&]() { + with_simd_level_with_spr([&]() { hamming_count_thres_fixSL(bs1, bs2, n1, n2, ht, ncodes, nptr); }); } @@ -232,7 +232,7 @@ void crosshamming_count_thres( hamdis_t ht, size_t ncodes, size_t* nptr) { - with_simd_level_a0_spr([&]() { + with_simd_level_with_spr([&]() { crosshamming_count_thres_fixSL(dbs, n, ht, ncodes, nptr); }); } @@ -246,7 +246,7 @@ size_t match_hamming_thres( size_t ncodes, int64_t* idx, hamdis_t* dis) { - return with_simd_level_a0_spr([&]() -> size_t { + return with_simd_level_with_spr([&]() -> size_t { return match_hamming_thres_fixSL( bs1, bs2, n1, n2, ht, ncodes, idx, dis); }); @@ -259,7 +259,7 @@ void generalized_hammings_knn_hc( size_t nb, size_t code_size, int ordered) { - with_simd_level_a0_spr([&]() { + with_simd_level_with_spr([&]() { generalized_hammings_knn_hc_fixSL(ha, a, b, nb, code_size, ordered); }); } diff --git a/faiss/utils/simd_impl/super_kmeans_dispatch.h b/faiss/utils/simd_impl/super_kmeans_dispatch.h index 33e7ac00d6..0be2337614 100644 --- a/faiss/utils/simd_impl/super_kmeans_dispatch.h +++ b/faiss/utils/simd_impl/super_kmeans_dispatch.h @@ -11,10 +11,10 @@ // highest available SIMD specialization at runtime (DD mode) or the // compiled-in level (static mode). // -// The A1 level mask is required: plain with_simd_level() omits the ARM_SVE bit, -// so on an SVE host the SVE specialization would never be instantiated. -// ARM_NEON has no specialization and falls through to the scalar primary -// template. +// The BASE_WITH_SVE level mask is required: plain with_simd_level() omits the +// ARM_SVE bit, so on an SVE host the SVE specialization would never be +// instantiated. ARM_NEON has no specialization and falls through to the scalar +// primary template. #include #include @@ -23,7 +23,7 @@ namespace faiss { namespace detail { inline float block_l2_dispatch(const float* x, const float* y, int n) { - return with_simd_level_a1( + return with_simd_level_with_sve( [&]() { return block_l2(x, y, n); }); } diff --git a/faiss/utils/simd_levels.cpp b/faiss/utils/simd_levels.cpp index 1450d5e191..e77ac9e667 100644 --- a/faiss/utils/simd_levels.cpp +++ b/faiss/utils/simd_levels.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #if defined(_MSC_VER) @@ -156,6 +157,35 @@ void detect_x86_uarch_flags() {} // NOLINTNEXTLINE(facebook-avoid-non-const-global-variables) static SIMDConfig simd_config_initializer; +namespace { + +/// Levels this binary actually contains code for. Mirrors the case labels in +/// with_selected_simd_levels; a level with no case label there is unreachable. +uint64_t compiled_simd_levels() { + uint64_t mask = uint64_t(1) << static_cast(SIMDLevel::NONE); +#ifdef COMPILE_SIMD_AVX2 + mask |= uint64_t(1) << static_cast(SIMDLevel::AVX2); +#endif +#ifdef COMPILE_SIMD_AVX512 + mask |= uint64_t(1) << static_cast(SIMDLevel::AVX512); +#endif +#ifdef COMPILE_SIMD_AVX512_SPR + mask |= uint64_t(1) << static_cast(SIMDLevel::AVX512_SPR); +#endif +#ifdef COMPILE_SIMD_ARM_NEON + mask |= uint64_t(1) << static_cast(SIMDLevel::ARM_NEON); +#endif +#ifdef COMPILE_SIMD_ARM_SVE + mask |= uint64_t(1) << static_cast(SIMDLevel::ARM_SVE); +#endif +#ifdef COMPILE_SIMD_RISCV_RVV + mask |= uint64_t(1) << static_cast(SIMDLevel::RISCV_RVV); +#endif + return mask; +} + +} // namespace + SIMDConfig::SIMDConfig(const char** faiss_simd_level_env) { // Support dependency injection for testing const char* env_var = faiss_simd_level_env ? *faiss_simd_level_env @@ -164,7 +194,27 @@ SIMDConfig::SIMDConfig(const char** faiss_simd_level_env) { if (!env_var) { level = auto_detect_simd_level(); } else { - level = to_simd_level(env_var); + // The override is honoured even when the running CPU lacks the + // requested level -- forcing a level is the point of it -- but not + // when the binary holds no code for that level. + // with_selected_simd_levels has no case label for an uncompiled level, + // so every dispatch would land on the default (NONE) and skip the + // levels in between. Walk down to the nearest compiled level instead, + // which on aarch64 turns ARM_SVE -> NONE into ARM_SVE -> ARM_NEON. NONE + // is always compiled, so this terminates. + const uint64_t compiled = compiled_simd_levels(); + const SIMDLevel requested = to_simd_level(env_var); + level = requested; + while (((compiled >> static_cast(level)) & 1) == 0) { + level = get_simd_fallback(level); + } + if (level != requested) { + fprintf(stderr, + "faiss: FAISS_SIMD_LEVEL=%s is not compiled into this " + "build, using %s instead\n", + to_string(requested).c_str(), + to_string(level).c_str()); + } supported_simd_levels = (1 << static_cast(level)); } supported_simd_levels |= (1 << static_cast(SIMDLevel::NONE)); diff --git a/tests/test_index_accuracy.py b/tests/test_index_accuracy.py index 8049310b0b..5b494aa1a7 100644 --- a/tests/test_index_accuracy.py +++ b/tests/test_index_accuracy.py @@ -386,7 +386,8 @@ def test_8bit_direct(self): faiss.ScalarQuantizer.QT_8bit_direct, faiss.ScalarQuantizer.QT_8bit_direct_signed, ): - # d % 16 / 32 / 64 exercise the AVX2 / AVX512 / SPR byte kernels + # d % 16 / 32 / 64 exercise the NEON / AVX2 / AVX512 / SPR byte + # kernels for d in 13, 16, 24, 32, 64, 128: for metric_type in faiss.METRIC_L2, faiss.METRIC_INNER_PRODUCT: self.subtest_8bit_direct(metric_type, d, quantizer) diff --git a/tests/test_scalar_quantizer_correctness.py b/tests/test_scalar_quantizer_correctness.py index be57d85ddb..259b145daf 100644 --- a/tests/test_scalar_quantizer_correctness.py +++ b/tests/test_scalar_quantizer_correctness.py @@ -603,3 +603,59 @@ def test_search_distances_sorted(self): for q in range(len(xq)): for k in range(1, 10): self.assertLessEqual(D[q, k - 1], D[q, k]) + + +@for_all_simd_levels +class TestSQByteDirectAcrossSIMDLevels(unittest.TestCase): + """QT_8bit_direct{,_signed} distances are exact integers, so every SIMD + level must agree with NONE bit for bit. + + d = 16 and 32 hit the d % 16 == 0 gate in sq-dispatch.h that routes to the + byte-domain DistanceComputerByte / DistanceComputerByteSigned kernels. The + two tests cover the two independent dispatch chains those kernels are + reached from: search() goes through sq_select_InvertedListScanner, and + get_distance_computer() through sq_select_distance_computer (the chain + Refine(SQ8) uses). + """ + + def do_test(self, measure): + if not faiss.SIMDConfig.is_simd_level_available(faiss.SIMDLevel_NONE): + self.skipTest("SIMDLevel.NONE not available") + rng = np.random.RandomState(1234) + for qtype in ( + faiss.ScalarQuantizer.QT_8bit_direct, + faiss.ScalarQuantizer.QT_8bit_direct_signed, + ): + lo, hi = ( + (-128, 128) + if qtype == faiss.ScalarQuantizer.QT_8bit_direct_signed + else (0, 256) + ) + for metric in (faiss.METRIC_L2, faiss.METRIC_INNER_PRODUCT): + for d in (16, 32): + with self.subTest(qtype=qtype, metric=metric, d=d): + xb = rng.randint(lo, hi, (500, d)).astype("float32") + xq = rng.randint(lo, hi, (10, d)).astype("float32") + index = faiss.IndexScalarQuantizer(d, qtype, metric) + index.add(xb) + got = measure(index, xq) + with NoneSIMDLevel(): + expect = measure(index, xq) + np.testing.assert_array_equal(got, expect) + + def test_scanner(self): + # One sequential scan per query at both levels, so identical distances + # imply identical tie-breaking and the labels are comparable too. + self.do_test(lambda index, xq: np.hstack(index.search(xq, 10))) + + def test_distance_computer(self): + def measure(index, xq): + dc = index.get_distance_computer() + out = np.empty((len(xq), index.ntotal), dtype="float32") + for q in range(len(xq)): + dc.set_query(faiss.swig_ptr(xq[q])) + for i in range(index.ntotal): + out[q, i] = dc(int(i)) + return out + + self.do_test(measure)