diff --git a/benchs/bench_rabitq.py b/benchs/bench_rabitq.py index 3162cebde..fa5f91a5a 100644 --- a/benchs/bench_rabitq.py +++ b/benchs/bench_rabitq.py @@ -23,8 +23,8 @@ from faiss.contrib.datasets import SyntheticDataset # d=512 -> size=64 -> 512-bit zmm, 1 iteration per bit-plane # d=768 -> size=96 -> 512-bit zmm + 256-bit tail # d=1024 -> size=128 -> 512-bit zmm only, 2 iterations per bit-plane -# Sweeping these is useful for verifying the AVX512_SPR (vpopcntdq) -# specialization in faiss/utils/simd_impl/rabitq_avx512_spr.cpp and for +# Sweeping these is useful for verifying the AVX512_VPOPCNT specialization in +# faiss/utils/simd_impl/rabitq_avx512_vpopcnt.cpp and for # profiling perf-record annotations across SIMD-width tiers. DIMENSIONS = [256, 512, 768, 1024] nlist: int = 1000 diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index 83e5a8753..9d893d886 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -39,15 +39,16 @@ set(FAISS_SIMD_AVX512_SRC utils/simd_impl/rabitq_avx512.cpp utils/simd_impl/super_kmeans_kernels_avx512.cpp ) -# AVX-512 sources that additionally require AVX512_VPOPCNTDQ -# (Sapphire Rapids and later). Compiled into faiss_avx512_spr only, -# and into the DD faiss target with an extra per-file -mavx512vpopcntdq -# flag. Not compiled into faiss_avx512 — that target stops at the -# baseline AVX-512F/CD/VL/DQ/BW feature set. +# AVX-512 sources that only require AVX512_VPOPCNTDQ on top of the baseline +# AVX-512F/CD/VL/DQ/BW feature set. In DD builds these are available on CPUs +# such as Ice Lake and Zen 4 without requiring the full SPR feature set. +set(FAISS_SIMD_AVX512_VPOPCNT_SRC + utils/simd_impl/rabitq_avx512_vpopcnt.cpp + utils/hamming_distance/hamming_avx512_vpopcnt.cpp +) +# AVX-512 sources that require the full Sapphire Rapids feature set. set(FAISS_SIMD_AVX512_SPR_SRC impl/scalar_quantizer/sq-avx512-spr.cpp - utils/simd_impl/rabitq_avx512_spr.cpp - utils/hamming_distance/hamming_avx512_spr.cpp ) set(FAISS_SIMD_NEON_SRC impl/fast_scan/impl-neon.cpp @@ -80,7 +81,8 @@ if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86_64|amd64|AMD64)") set(FAISS_SIMD_SRC ${FAISS_SIMD_AVX2_SRC} ${FAISS_SIMD_AVX512_SRC} - ${FAISS_SIMD_AVX512_SPR_SRC} ${FAISS_SIMD_AVX512_SPR_SRC}) + ${FAISS_SIMD_AVX512_VPOPCNT_SRC} + ${FAISS_SIMD_AVX512_SPR_SRC}) elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(aarch64|arm64|ARM64)") set(FAISS_SIMD_SRC ${FAISS_SIMD_NEON_SRC} ${FAISS_SIMD_SVE_SRC}) elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(riscv64|riscv)") @@ -404,7 +406,7 @@ set(FAISS_HEADERS utils/hamming_distance/hamming_computer.h utils/hamming_distance/hamming_computer-avx2.h utils/hamming_distance/hamming_computer-avx512.h - utils/hamming_distance/hamming_computer-avx512_spr.h + utils/hamming_distance/hamming_computer-avx512_vpopcnt.h utils/hamming_distance/hamming_computer-generic.h utils/hamming_distance/hamming_computer-neon.h utils/hamming_distance/hamming_computer-rvv.h @@ -496,8 +498,16 @@ else() # we need bigobj for the swig wrapper add_compile_options(/bigobj) endif() -target_sources(faiss_avx512_spr PRIVATE ${FAISS_SIMD_AVX2_SRC} ${FAISS_SIMD_AVX512_SRC} ${FAISS_SIMD_AVX512_SPR_SRC}) -target_compile_definitions(faiss_avx512_spr PRIVATE COMPILE_SIMD_AVX2 COMPILE_SIMD_AVX512 COMPILE_SIMD_AVX512_SPR) +target_sources(faiss_avx512_spr PRIVATE + ${FAISS_SIMD_AVX2_SRC} + ${FAISS_SIMD_AVX512_SRC} + ${FAISS_SIMD_AVX512_VPOPCNT_SRC} + ${FAISS_SIMD_AVX512_SPR_SRC}) +target_compile_definitions(faiss_avx512_spr PRIVATE + COMPILE_SIMD_AVX2 + COMPILE_SIMD_AVX512 + COMPILE_SIMD_AVX512_VPOPCNT + COMPILE_SIMD_AVX512_SPR) add_library(faiss_sve ${FAISS_SRC}) if(NOT FAISS_OPT_LEVEL STREQUAL "sve") @@ -539,7 +549,10 @@ if(FAISS_OPT_LEVEL STREQUAL "dd") # Architecture-specific SIMD definitions for Dynamic Dispatch if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86_64|amd64|AMD64)") target_compile_definitions(faiss PRIVATE - COMPILE_SIMD_AVX2 COMPILE_SIMD_AVX512 COMPILE_SIMD_AVX512_SPR) + COMPILE_SIMD_AVX2 + COMPILE_SIMD_AVX512 + COMPILE_SIMD_AVX512_VPOPCNT + COMPILE_SIMD_AVX512_SPR) if(NOT WIN32) # Baseline flags for common files (prevents auto-vectorization) target_compile_options(faiss PRIVATE @@ -554,9 +567,13 @@ if(FAISS_OPT_LEVEL STREQUAL "dd") PROPERTIES COMPILE_OPTIONS "-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mfma;-mf16c;-mpopcnt" ) - # SPR-only sources additionally require AVX512_VPOPCNTDQ. - # vpopcntq is the whole point of these files, so the extra flag - # is mandatory; without it the intrinsics fail to compile. + # VPOPCNT sources only require VPOPCNTDQ in addition to baseline AVX-512. + set_source_files_properties(${FAISS_SIMD_AVX512_VPOPCNT_SRC} + TARGET_DIRECTORY faiss + PROPERTIES COMPILE_OPTIONS + "-mavx2;-mfma;-mf16c;-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mavx512vpopcntdq;-mpopcnt" + ) + # The remaining SPR sources require the full SPR feature set. set_source_files_properties(${FAISS_SIMD_AVX512_SPR_SRC} TARGET_DIRECTORY faiss PROPERTIES COMPILE_OPTIONS @@ -573,11 +590,11 @@ if(FAISS_OPT_LEVEL STREQUAL "dd") TARGET_DIRECTORY faiss PROPERTIES COMPILE_OPTIONS "/arch:AVX512" ) - # MSVC has no per-feature flag for VPOPCNTDQ; /arch:AVX512 - # enables it as part of the AVX-512 feature set used by recent - # toolchains. (Newer MSVC supports __isa_available checks at - # runtime; the SPR specialization is gated by COMPILE_SIMD_AVX512_SPR - # at the source level, so this is safe.) + # MSVC has no per-feature flag for VPOPCNTDQ. + set_source_files_properties(${FAISS_SIMD_AVX512_VPOPCNT_SRC} + TARGET_DIRECTORY faiss + PROPERTIES COMPILE_OPTIONS "/arch:AVX512" + ) set_source_files_properties(${FAISS_SIMD_AVX512_SPR_SRC} TARGET_DIRECTORY faiss PROPERTIES COMPILE_OPTIONS "/arch:AVX512" diff --git a/faiss/docs/simd_dynamic_dispatch_migration.md b/faiss/docs/simd_dynamic_dispatch_migration.md index 0bf516e76..66c45521a 100644 --- a/faiss/docs/simd_dynamic_dispatch_migration.md +++ b/faiss/docs/simd_dynamic_dispatch_migration.md @@ -215,8 +215,8 @@ ARM_NEON + RISCV_RVV implementations exist. If your function has another subset of available implementations, it can be passed with `with_selected_simd_levels` with a bitmask of available levels. Missing levels in the mask cause the dispatch to **fall through** to the next lower -level in the same architecture family (x86: AVX512_SPR → AVX512 → AVX2 → -NONE; ARM: ARM_SVE → ARM_NEON → NONE; RISC-V: RISCV_RVV → NONE — +level in the same architecture family (x86: AVX512_SPR → AVX512_VPOPCNT → +AVX512 → AVX2 → NONE; ARM: ARM_SVE → ARM_NEON → NONE; RISC-V: RISCV_RVV → NONE — architecture chains are independent): ```cpp @@ -239,7 +239,8 @@ 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_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_VPOPCNT` | BASE + AVX512_VPOPCNT | Kernels needing only VPOPCNTDQ (`with_simd_level_with_vpopcnt`) | +| `AVAILABLE_SIMD_LEVELS_BASE_WITH_SPR` | BASE + AVX512_SPR | Kernels needing the whole SPR feature set | | `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 | diff --git a/faiss/impl/RaBitQUtils.cpp b/faiss/impl/RaBitQUtils.cpp index f2eb77c7e..dae57cd09 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 04c3900cc..0bbf758d2 100644 --- a/faiss/impl/RaBitQuantizer.cpp +++ b/faiss/impl/RaBitQuantizer.cpp @@ -721,12 +721,10 @@ FlatCodesDistanceComputer* RaBitQuantizer::get_distance_computer( // call the SIMD-specialized rabitq functions directly (no per-call // with_simd_level overhead). // - // 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( + // VPOPCNT rather than SPR: Ice Lake and Zen 4 have VPOPCNTDQ without the + // rest of the SPR feature set. Below it, dispatch falls through to + // rabitq_avx512.cpp. + return with_selected_simd_levels( [&]() -> FlatCodesDistanceComputer* { if (qb == 0) { auto dc = diff --git a/faiss/impl/fast_scan/decompose_qbs.h b/faiss/impl/fast_scan/decompose_qbs.h index 48ca40210..dc44c3d60 100644 --- a/faiss/impl/fast_scan/decompose_qbs.h +++ b/faiss/impl/fast_scan/decompose_qbs.h @@ -40,6 +40,7 @@ void kernel_accumulate_block( #ifdef __AVX512F__ if constexpr ( KernelSL == SIMDLevel::AVX512 || + KernelSL == SIMDLevel::AVX512_VPOPCNT || KernelSL == SIMDLevel::AVX512_SPR) { pq4_kernel_qbs_512(nsq, codes, LUT, res, scaler); } else { diff --git a/faiss/impl/scalar_quantizer/sq-dispatch.h b/faiss/impl/scalar_quantizer/sq-dispatch.h index 0dfdd830e..a02f8fe82 100644 --- a/faiss/impl/scalar_quantizer/sq-dispatch.h +++ b/faiss/impl/scalar_quantizer/sq-dispatch.h @@ -28,6 +28,14 @@ namespace scalar_quantizer { // Define SL as alias for THE_LEVEL_TO_DISPATCH for use in this file constexpr SIMDLevel SL = THE_LEVEL_TO_DISPATCH; +// The SPR scalar-quantizer implementation uses VPOPCNTDQ for its RaBitQ +// popcount helpers, but those helpers belong to the narrower VPOPCNT +// capability. Keep the SQ implementation at AVX512_SPR while explicitly +// reusing the independently dispatched VPOPCNT kernels. +template +inline constexpr SIMDLevel rabitq_popcount_level = + SL0 == SIMDLevel::AVX512_SPR ? SIMDLevel::AVX512_VPOPCNT : SL0; + /******************************************************************* * TurboQuant SIMD kernel: masked_sum * Compute sum of arr[j] where bit j of the bitmask is set. @@ -259,9 +267,11 @@ struct DCTurboQuantFull : ScalarQuantizer::TurboQuantRefine::DistanceComputer { if (qb > 0) { // Integer popcount path for 1-bit MSE size_t byte_size = (d + 7) / 8; - uint64_t and_result = rabitq::bitwise_and_dot_product( + uint64_t and_result = rabitq::bitwise_and_dot_product< + rabitq_popcount_level>( rearranged_q.data(), code, byte_size, qb); - uint64_t pop = rabitq::popcount(code, byte_size); + uint64_t pop = rabitq::popcount>( + code, byte_size); mse_dot = mse_base + mse_int_scale * static_cast(and_result) + mse_popcnt_scale * static_cast(pop); @@ -316,9 +326,11 @@ struct DCTurboQuantFull : ScalarQuantizer::TurboQuantRefine::DistanceComputer { float qjl_dot; if (qb > 0 && int_qjl) { size_t byte_size = (d + 7) / 8; - uint64_t and_result = rabitq::bitwise_and_dot_product( - rearranged_qproj.data(), qjl_code, byte_size, qb); - uint64_t pop = rabitq::popcount(qjl_code, byte_size); + uint64_t and_result = + rabitq::bitwise_and_dot_product>( + rearranged_qproj.data(), qjl_code, byte_size, qb); + uint64_t pop = rabitq::popcount>( + qjl_code, byte_size); float pos_sum = qjl_popcnt_scale * static_cast(pop) + qjl_int_scale * static_cast(and_result); qjl_dot = qjl_coeff * gamma * (2.0f * pos_sum - total_qproj_sum); @@ -342,10 +354,17 @@ struct DCTurboQuantFull : ScalarQuantizer::TurboQuantRefine::DistanceComputer { } }; +// True for every level that runs the 512-bit kernels. AVX512_VPOPCNT is +// unreachable here until an SQ translation unit is compiled at that level; +// naming it keeps the 512-bit alignment rule correct when one is. +template +constexpr bool is_avx512_family = SL2 == SIMDLevel::AVX512 || + SL2 == SIMDLevel::AVX512_VPOPCNT || SL2 == SIMDLevel::AVX512_SPR; + // Returns true if dimension d is compatible with the given SIMD level template constexpr bool is_dimension_compatible(size_t d) { - if constexpr (SL2 == SIMDLevel::AVX512 || SL2 == SIMDLevel::AVX512_SPR) { + if constexpr (is_avx512_family) { return d % 16 == 0; } else if constexpr (SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) { return d % 8 == 0; @@ -525,8 +544,7 @@ SQDistanceComputer* select_distance_computer_body( return new DCTemplate, Sim, SL2>(d, trained); case ScalarQuantizer::QT_8bit_direct: - if constexpr ( - SL2 == SIMDLevel::AVX512 || SL2 == SIMDLevel::AVX512_SPR) { + if constexpr (is_avx512_family) { if (d % 32 == 0) { return new DistanceComputerByte( static_cast(d), trained); @@ -547,7 +565,7 @@ SQDistanceComputer* select_distance_computer_body( return new DistanceComputerByteSigned( static_cast(d), trained); } - } else if constexpr (SL2 == SIMDLevel::AVX512) { + } else if constexpr (is_avx512_family) { if (d % 32 == 0) { return new DistanceComputerByteSigned( static_cast(d), trained); @@ -739,9 +757,7 @@ InvertedListScanner* sq_select_InvertedListScanner( return scan.template operator(), Similarity, SL2>>(); case ScalarQuantizer::QT_8bit_direct: - if constexpr ( - SL2 == SIMDLevel::AVX512 || - SL2 == SIMDLevel::AVX512_SPR) { + if constexpr (is_avx512_family) { if (d % 32 == 0) { return scan.template operator()>(); @@ -763,7 +779,7 @@ InvertedListScanner* sq_select_InvertedListScanner( return scan.template operator()< DistanceComputerByteSigned>(); } - } else if constexpr (SL2 == SIMDLevel::AVX512) { + } else if constexpr (is_avx512_family) { if (d % 32 == 0) { return scan.template operator()< DistanceComputerByteSigned>(); diff --git a/faiss/impl/simd_dispatch.h b/faiss/impl/simd_dispatch.h index d59169e3e..8a275327d 100644 --- a/faiss/impl/simd_dispatch.h +++ b/faiss/impl/simd_dispatch.h @@ -36,8 +36,13 @@ constexpr int AVAILABLE_SIMD_LEVELS_AVX2_NEON = AVAILABLE_SIMD_LEVELS_NONE | constexpr int AVAILABLE_SIMD_LEVELS_BASE = AVAILABLE_SIMD_LEVELS_AVX2_NEON | (1 << int(SIMDLevel::AVX512)) | (1 << int(SIMDLevel::RISCV_RVV)); -// BASE_WITH_SPR: BASE + AVX512_SPR, for functions with a dedicated SPR -// specialization on top of an AVX512 fallback. +// BASE_WITH_VPOPCNT: BASE + AVX512_VPOPCNT, for kernels that need only +// VPOPCNTDQ on top of baseline AVX-512 (Ice Lake, Zen 4, Zen 5). +constexpr int AVAILABLE_SIMD_LEVELS_BASE_WITH_VPOPCNT = + AVAILABLE_SIMD_LEVELS_BASE | (1 << int(SIMDLevel::AVX512_VPOPCNT)); + +// BASE_WITH_SPR: BASE + AVX512_SPR, for kernels that need the whole SPR +// feature set rather than VPOPCNTDQ alone. constexpr int AVAILABLE_SIMD_LEVELS_BASE_WITH_SPR = AVAILABLE_SIMD_LEVELS_BASE | (1 << int(SIMDLevel::AVX512_SPR)); @@ -51,6 +56,8 @@ constexpr int AVAILABLE_SIMD_LEVELS_ALL = -1; constexpr SIMDLevel get_simd_fallback(SIMDLevel level) { switch (level) { case SIMDLevel::AVX512_SPR: + return SIMDLevel::AVX512_VPOPCNT; + case SIMDLevel::AVX512_VPOPCNT: return SIMDLevel::AVX512; case SIMDLevel::AVX512: return SIMDLevel::AVX2; @@ -119,6 +126,15 @@ inline auto with_selected_simd_levels(LambdaType&& action) { [[fallthrough]]; #endif +#ifdef COMPILE_SIMD_AVX512_VPOPCNT + case SIMDLevel::AVX512_VPOPCNT: + if constexpr ( + available_levels & (1 << int(SIMDLevel::AVX512_VPOPCNT))) { + return action.template operator()(); + } + [[fallthrough]]; +#endif + #ifdef COMPILE_SIMD_AVX512 case SIMDLevel::AVX512: if constexpr (available_levels & (1 << int(SIMDLevel::AVX512))) { @@ -166,7 +182,7 @@ inline auto with_selected_simd_levels(LambdaType&& action) { // In static mode, SINGLE_SIMD_LEVEL is a constexpr resolved at compile // time. We mirror the DD fallthrough behavior at compile time via // dispatch_with_fallback, which recursively walks get_simd_fallback: - // x86: AVX512_SPR -> AVX512 -> AVX2 -> NONE + // x86: AVX512_SPR -> AVX512_VPOPCNT -> AVX512 -> AVX2 -> NONE // ARM: ARM_SVE -> ARM_NEON -> NONE // RISCV: RISCV_RVV -> NONE // The first level in the chain that appears in available_levels is @@ -233,25 +249,21 @@ inline auto with_simd_level_256bit(LambdaType&& action) { std::forward(action)); } -/** - * 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. - */ +// Plain with_simd_level() uses BASE, which omits the optional levels below. +// A call site must opt in, or its specialization is never instantiated. + +/// BASE + ARM_SVE. template inline auto with_simd_level_with_sve(LambdaType&& action) { return with_selected_simd_levels( std::forward(action)); } -/** - * Use for functions that have BASE-level implementations plus an AVX512_SPR - * specialization (e.g. using VPOPCNTDQ). - */ +/// BASE + AVX512_VPOPCNT. template -inline auto with_simd_level_with_spr(LambdaType&& action) { - return with_selected_simd_levels( +inline auto with_simd_level_with_vpopcnt(LambdaType&& action) { + return with_selected_simd_levels( std::forward(action)); } + } // namespace faiss diff --git a/faiss/python/__init__.pyi b/faiss/python/__init__.pyi index f0c2c8658..bc1e24ed9 100644 --- a/faiss/python/__init__.pyi +++ b/faiss/python/__init__.pyi @@ -3087,8 +3087,11 @@ SIMDLevel_AVX512_SPR: int SIMDLevel_ARM_NEON: int SIMDLevel_ARM_SVE: int SIMDLevel_RISCV_RVV: int +SIMDLevel_AVX512_VPOPCNT: int SIMDLevel_COUNT: int +def compiled_simd_levels() -> int: ... + class SIMDConfig: level: int supported_simd_levels: int diff --git a/faiss/utils/hamming.cpp b/faiss/utils/hamming.cpp index 870055486..51f1ef727 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_with_spr([&]() { + with_simd_level_with_vpopcnt([&]() { // 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_with_spr([&]() { + with_simd_level_with_vpopcnt([&]() { 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_with_spr([&]() { + with_simd_level_with_vpopcnt([&]() { 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_with_spr([&]() { + with_simd_level_with_vpopcnt([&]() { 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_with_spr([&]() { + with_simd_level_with_vpopcnt([&]() { 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_with_spr([&]() { + with_simd_level_with_vpopcnt([&]() { 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_with_spr([&]() -> size_t { + return with_simd_level_with_vpopcnt([&]() -> 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_with_spr([&]() { + with_simd_level_with_vpopcnt([&]() { generalized_hammings_knn_hc_fixSL(ha, a, b, nb, code_size, ordered); }); } diff --git a/faiss/utils/hamming_distance/hamming_avx512_spr.cpp b/faiss/utils/hamming_distance/hamming_avx512_vpopcnt.cpp similarity index 59% rename from faiss/utils/hamming_distance/hamming_avx512_spr.cpp rename to faiss/utils/hamming_distance/hamming_avx512_vpopcnt.cpp index 1912bcc69..69b0e3d1b 100644 --- a/faiss/utils/hamming_distance/hamming_avx512_spr.cpp +++ b/faiss/utils/hamming_distance/hamming_avx512_vpopcnt.cpp @@ -5,11 +5,11 @@ * LICENSE file in the root directory of this source tree. */ -#ifdef COMPILE_SIMD_AVX512_SPR +#ifdef COMPILE_SIMD_AVX512_VPOPCNT -#define THE_SIMD_LEVEL SIMDLevel::AVX512_SPR +#define THE_SIMD_LEVEL SIMDLevel::AVX512_VPOPCNT // NOLINTNEXTLINE(facebook-hte-InlineHeader) -#include +#include #include -#endif // COMPILE_SIMD_AVX512_SPR +#endif // COMPILE_SIMD_AVX512_VPOPCNT diff --git a/faiss/utils/hamming_distance/hamming_computer-avx512.h b/faiss/utils/hamming_distance/hamming_computer-avx512.h index c44941bb4..a52b62e5a 100644 --- a/faiss/utils/hamming_distance/hamming_computer-avx512.h +++ b/faiss/utils/hamming_distance/hamming_computer-avx512.h @@ -12,7 +12,7 @@ // Types without custom AVX512 code inherit from the NONE specializations // in hamming_computer-generic.h. HammingComputer64 and // HammingComputerDefault use scalar popcount here; the VPOPCNTDQ fast -// path lives in hamming_computer-avx512_spr.h (AVX512_SPR level). +// path lives in hamming_computer-avx512_vpopcnt.h (AVX512_VPOPCNT level). // GenHammingComputer classes leverage SSE/AVX2 intrinsics. #include diff --git a/faiss/utils/hamming_distance/hamming_computer-avx512_spr.h b/faiss/utils/hamming_distance/hamming_computer-avx512_vpopcnt.h similarity index 78% rename from faiss/utils/hamming_distance/hamming_computer-avx512_spr.h rename to faiss/utils/hamming_distance/hamming_computer-avx512_vpopcnt.h index fa3e063ac..751eede05 100644 --- a/faiss/utils/hamming_distance/hamming_computer-avx512_spr.h +++ b/faiss/utils/hamming_distance/hamming_computer-avx512_vpopcnt.h @@ -5,14 +5,12 @@ * LICENSE file in the root directory of this source tree. */ -#ifndef HAMMING_COMPUTER_AVX512_SPR_H -#define HAMMING_COMPUTER_AVX512_SPR_H +#ifndef HAMMING_COMPUTER_AVX512_VPOPCNT_H +#define HAMMING_COMPUTER_AVX512_VPOPCNT_H -// AVX512_SPR HammingComputer specializations using VPOPCNTDQ. -// On Sapphire Rapids+, _mm512_popcnt_epi64 (and _mm256_popcnt_epi64 with VL) -// are unconditionally available. This gives a faster path than the scalar -// popcount fallback used in the base AVX512 specializations when compiled -// without -mavx512vpopcntdq. +// AVX512_VPOPCNT HammingComputer specializations using VPOPCNTDQ. This gives +// a faster path than the scalar popcount fallback used in the base AVX512 +// specializations when compiled without -mavx512vpopcntdq. #include #include @@ -25,31 +23,31 @@ namespace faiss { /*************************************************************************** - * AVX512_SPR inheriting specializations for types without custom SPR code. + * AVX512_VPOPCNT inheriting specializations without custom VPOPCNT code. ***************************************************************************/ -#define FAISS_INHERIT_HAMMING_SPR(Class) \ - template <> \ - struct Class## \ - _tpl : Class##_tpl { \ - using Class##_tpl::Class##_tpl; \ +#define FAISS_INHERIT_HAMMING_VPOPCNT(Class) \ + template <> \ + struct Class## \ + _tpl : Class##_tpl { \ + using Class##_tpl::Class##_tpl; \ } -FAISS_INHERIT_HAMMING_SPR(HammingComputer16); -FAISS_INHERIT_HAMMING_SPR(HammingComputer20); -FAISS_INHERIT_HAMMING_SPR(GenHammingComputer8); -FAISS_INHERIT_HAMMING_SPR(GenHammingComputer16); -FAISS_INHERIT_HAMMING_SPR(GenHammingComputer32); -FAISS_INHERIT_HAMMING_SPR(GenHammingComputerM8); +FAISS_INHERIT_HAMMING_VPOPCNT(HammingComputer16); +FAISS_INHERIT_HAMMING_VPOPCNT(HammingComputer20); +FAISS_INHERIT_HAMMING_VPOPCNT(GenHammingComputer8); +FAISS_INHERIT_HAMMING_VPOPCNT(GenHammingComputer16); +FAISS_INHERIT_HAMMING_VPOPCNT(GenHammingComputer32); +FAISS_INHERIT_HAMMING_VPOPCNT(GenHammingComputerM8); -#undef FAISS_INHERIT_HAMMING_SPR +#undef FAISS_INHERIT_HAMMING_VPOPCNT /*************************************************************************** - * Custom AVX512_SPR specializations using VPOPCNTDQ. + * Custom AVX512_VPOPCNT specializations using VPOPCNTDQ. ***************************************************************************/ template <> -struct HammingComputer32_tpl { +struct HammingComputer32_tpl { const uint8_t* a8; HammingComputer32_tpl() {} @@ -81,7 +79,7 @@ struct HammingComputer32_tpl { }; template <> -struct HammingComputer64_tpl { +struct HammingComputer64_tpl { const uint8_t* a8; HammingComputer64_tpl() {} @@ -108,7 +106,7 @@ struct HammingComputer64_tpl { }; template <> -struct HammingComputerDefault_tpl { +struct HammingComputerDefault_tpl { const uint8_t* a8; int quotient8; int remainder8; diff --git a/faiss/utils/simd_impl/rabitq_avx512_spr.cpp b/faiss/utils/simd_impl/rabitq_avx512_vpopcnt.cpp similarity index 91% rename from faiss/utils/simd_impl/rabitq_avx512_spr.cpp rename to faiss/utils/simd_impl/rabitq_avx512_vpopcnt.cpp index 0f8951f80..c52e2298e 100644 --- a/faiss/utils/simd_impl/rabitq_avx512_spr.cpp +++ b/faiss/utils/simd_impl/rabitq_avx512_vpopcnt.cpp @@ -6,34 +6,31 @@ */ /** - * @file rabitq_avx512_spr.cpp + * @file rabitq_avx512_vpopcnt.cpp * - * RaBitQ SIMD kernels specialized for SIMDLevel::AVX512_SPR. + * RaBitQ SIMD kernels specialized for SIMDLevel::AVX512_VPOPCNT. * - * Sapphire Rapids (SPR) and later Intel microarchitectures expose - * AVX-512 VPOPCNTDQ (vpopcntq), which performs a per-lane 64-bit - * popcount in a single instruction. This is used here to replace the - * multi-step shuffle/pshufb-based popcount used by the generic AVX-512 - * specialization in rabitq_avx512.cpp. The popcount-heavy kernels - * (bitwise_and_dot_product, bitwise_xor_dot_product, popcount) become - * substantially shorter and faster on SPR+ as a result. + * AVX-512 VPOPCNTDQ performs a per-lane 64-bit popcount in a single + * instruction. It is available on CPUs including Ice Lake, Zen 4, and + * Sapphire Rapids, independently of the other SPR-only extensions. This + * replaces the multi-step shuffle-based popcount used by the generic + * AVX-512 specialization in rabitq_avx512.cpp. * * Build / dispatch behavior: * - faiss_avx512 (AVX-512 only, no SPR features): NOT compiled. * The existing AVX512 specialization in rabitq_avx512.cpp is used. - * - faiss_avx512_spr (statically built for SPR+): compiled. The - * SINGLE_SIMD_LEVEL is AVX512_SPR, so this specialization is - * selected by static dispatch. + * - faiss_avx512_spr: compiled alongside the full SPR specialization and + * selected through the SPR -> VPOPCNT fallback. * - faiss with FAISS_OPT_LEVEL=dd (dynamic dispatch): compiled with * -mavx512vpopcntdq as a per-file flag. Selected at runtime when - * SIMDConfig::level == SIMDLevel::AVX512_SPR. + * the CPU exposes AVX512_VPOPCNTDQ. * * The floating-point multi-bit inner-product kernel does not benefit - * from VPOPCNTDQ, so this TU forwards compute_inner_product to + * from VPOPCNTDQ, so this TU forwards compute_inner_product to * the AVX512 implementation to avoid duplicating that code path. */ -#ifdef COMPILE_SIMD_AVX512_SPR +#ifdef COMPILE_SIMD_AVX512_VPOPCNT #include #include @@ -47,7 +44,7 @@ namespace faiss::rabitq { // Forward declarations for the AVX512 specializations defined in -// rabitq_avx512.cpp. They live in the same TU group on SPR builds, so +// rabitq_avx512.cpp. They live in the same TU group in supported builds, so // we can reuse them as a tail handler / fallback. Declaring rather // than redefining avoids ODR risk and keeps a single source of truth // for the floating-point kernel. @@ -75,8 +72,8 @@ inline __m512i popcount_512_vpopcntdq(__m512i v) { } // 256-bit popcount using AVX-512VL VPOPCNTDQ. -// AVX512VL is part of the SPR feature set, so vpopcntq is available -// on 256-bit registers via _mm256_popcnt_epi64. +// Baseline AVX-512 includes AVX512VL, so VPOPCNTDQ is also available on +// 256-bit registers via _mm256_popcnt_epi64. inline __m256i popcount_256_vpopcntdq(__m256i v) { return _mm256_popcnt_epi64(v); } @@ -101,7 +98,7 @@ inline uint64_t reduce_add_128(__m128i v) { } // namespace template <> -uint64_t bitwise_and_dot_product( +uint64_t bitwise_and_dot_product( const uint8_t* query, const uint8_t* data, size_t size, @@ -187,7 +184,7 @@ uint64_t bitwise_and_dot_product( template <> BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount< - SIMDLevel::AVX512_SPR>( + SIMDLevel::AVX512_VPOPCNT>( const uint8_t* query, const uint8_t* data, size_t size, @@ -278,7 +275,7 @@ BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount< } template <> -uint64_t bitwise_xor_dot_product( +uint64_t bitwise_xor_dot_product( const uint8_t* query, const uint8_t* data, size_t size, @@ -357,7 +354,7 @@ uint64_t bitwise_xor_dot_product( } template <> -uint64_t popcount(const uint8_t* data, size_t size) { +uint64_t popcount(const uint8_t* data, size_t size) { uint64_t sum = 0; size_t offset = 0; @@ -419,7 +416,7 @@ float compute_inner_product( float cb); template <> -float compute_inner_product( +float compute_inner_product( const uint8_t* __restrict sign_bits, const uint8_t* __restrict ex_code, const float* __restrict rotated_q, @@ -432,4 +429,4 @@ float compute_inner_product( } // namespace faiss::rabitq::multibit -#endif // COMPILE_SIMD_AVX512_SPR +#endif // COMPILE_SIMD_AVX512_VPOPCNT diff --git a/faiss/utils/simd_levels.cpp b/faiss/utils/simd_levels.cpp index f1c3fecdc..576d1d150 100644 --- a/faiss/utils/simd_levels.cpp +++ b/faiss/utils/simd_levels.cpp @@ -147,20 +147,10 @@ void detect_x86_uarch_flags() {} } // namespace -#ifdef FAISS_ENABLE_DD - -// ============================================================================= -// Dynamic Dispatch (DD) mode implementation -// ============================================================================= - -// Static initializer to run constructor at load time -// NOLINTNEXTLINE(facebook-avoid-non-const-global-variables) -static SIMDConfig simd_config_initializer; - -namespace { - -/// Levels this binary holds code for. Must mirror the case labels in -/// with_selected_simd_levels. +/// Must mirror the case labels in with_selected_simd_levels. A static build +/// defines a COMPILE_SIMD_* macro for every level whose sources it compiles, +/// not only for SINGLE_SIMD_LEVEL, so this reports a correct set in both +/// modes and the result is not a single level. uint64_t compiled_simd_levels() { uint64_t mask = uint64_t(1) << static_cast(SIMDLevel::NONE); #ifdef COMPILE_SIMD_AVX2 @@ -169,6 +159,9 @@ uint64_t compiled_simd_levels() { #ifdef COMPILE_SIMD_AVX512 mask |= uint64_t(1) << static_cast(SIMDLevel::AVX512); #endif +#ifdef COMPILE_SIMD_AVX512_VPOPCNT + mask |= uint64_t(1) << static_cast(SIMDLevel::AVX512_VPOPCNT); +#endif #ifdef COMPILE_SIMD_AVX512_SPR mask |= uint64_t(1) << static_cast(SIMDLevel::AVX512_SPR); #endif @@ -184,7 +177,11 @@ uint64_t compiled_simd_levels() { return mask; } -} // namespace +#ifdef FAISS_ENABLE_DD + +// Static initializer to run constructor at load time +// NOLINTNEXTLINE(facebook-avoid-non-const-global-variables) +static SIMDConfig simd_config_initializer; SIMDConfig::SIMDConfig(const char** faiss_simd_level_env) { // Support dependency injection for testing @@ -265,6 +262,11 @@ SIMDLevel SIMDConfig::auto_detect_simd_level() { // needed for the SPR detection below. Kept in a local so a later // xgetbv cannot clobber it. unsigned int cpuid7_edx = regs[3]; + // Leaf 7 subleaf 0, not leaf 1: leaf 1 ECX holds unrelated bits at + // these positions. + unsigned int ecx7 = regs[2]; + [[maybe_unused]] bool has_avx512_vnni = (ecx7 & (1 << 11)) != 0; + [[maybe_unused]] bool has_avx512_vpopcntdq = (ecx7 & (1 << 14)) != 0; uint64_t xcr0 = xgetbv0(); @@ -289,19 +291,28 @@ SIMDLevel SIMDConfig::auto_detect_simd_level() { supported_simd_levels |= (1 << static_cast(SIMDLevel::AVX512)); +#if defined(COMPILE_SIMD_AVX512_VPOPCNT) + if (has_avx512_vpopcntdq) { + detected_level = SIMDLevel::AVX512_VPOPCNT; + supported_simd_levels |= + (1 << static_cast(SIMDLevel::AVX512_VPOPCNT)); + } +#endif + #if defined(COMPILE_SIMD_AVX512_SPR) // Check for Sapphire Rapids features. - // The SPR code path is compiled with -mavx512fp16, so we - // must verify both AVX512_BF16 and AVX512_FP16 before - // dispatching to it. AMD Zen 4 (bergamo) has BF16 but - // not FP16 — using SPR code there causes SIGILL. + // The SPR code path is compiled with AVX512_VNNI, BF16, + // FP16, and VPOPCNTDQ, so all four features are required. + // AMD Zen 4 has VPOPCNTDQ and BF16 but not FP16, and must + // remain on the AVX512_VPOPCNT level. // CPUID EAX=7, ECX=1: EAX bit 5 = AVX512_BF16 // CPUID EAX=7, ECX=0: EDX bit 23 = AVX512_FP16 // (Linux: X86_FEATURE_AVX512_FP16 = 18*32+23) bool has_avx512_fp16 = (cpuid7_edx & (1 << 23)) != 0; cpuid_count(7, 1, regs); const bool has_avx512_bf16 = (regs[0] & (1 << 5)) != 0; - if (has_avx512_bf16 && has_avx512_fp16) { + if (has_avx512_vnni && has_avx512_vpopcntdq && + has_avx512_bf16 && has_avx512_fp16) { detected_level = SIMDLevel::AVX512_SPR; supported_simd_levels |= (1 << static_cast(SIMDLevel::AVX512_SPR)); @@ -395,6 +406,8 @@ SIMDLevel SIMDConfig::auto_detect_simd_level() { // In static mode, return the compiled-in level #if defined(COMPILE_SIMD_AVX512_SPR) return SIMDLevel::AVX512_SPR; +#elif defined(COMPILE_SIMD_AVX512_VPOPCNT) + return SIMDLevel::AVX512_VPOPCNT; #elif defined(COMPILE_SIMD_AVX512) return SIMDLevel::AVX512; #elif defined(COMPILE_SIMD_AVX2) @@ -429,6 +442,8 @@ std::string to_string(SIMDLevel level) { return "AVX2"; case SIMDLevel::AVX512: return "AVX512"; + case SIMDLevel::AVX512_VPOPCNT: + return "AVX512_VPOPCNT"; case SIMDLevel::AVX512_SPR: return "AVX512_SPR"; case SIMDLevel::ARM_NEON: @@ -453,6 +468,9 @@ SIMDLevel to_simd_level(const std::string& level_str) { if (level_str == "AVX512") { return SIMDLevel::AVX512; } + if (level_str == "AVX512_VPOPCNT") { + return SIMDLevel::AVX512_VPOPCNT; + } if (level_str == "AVX512_SPR") { return SIMDLevel::AVX512_SPR; } diff --git a/faiss/utils/simd_levels.h b/faiss/utils/simd_levels.h index a7eba7ede..781d0b800 100644 --- a/faiss/utils/simd_levels.h +++ b/faiss/utils/simd_levels.h @@ -21,13 +21,17 @@ enum class SIMDLevel { // x86 AVX2, AVX512, - AVX512_SPR, // Sapphire Rapids: AVX512 + BF16 + FP16 + VNNI + AVX512_SPR, // Sapphire Rapids: AVX512 + BF16 + FP16 + VNNI + VPOPCNTDQ // arm & aarch64 ARM_NEON, ARM_SVE, // Scalable Vector Extension (ARMv8.2+) // riscv RISCV_RVV, // RISC-V Vector Extension (rv64gcv) + // Appended to preserve the numeric values of the existing public enum. + // AVX-512 core features plus AVX512_VPOPCNTDQ (Ice Lake, Zen 4, etc.). + AVX512_VPOPCNT, + COUNT }; @@ -52,6 +56,8 @@ inline constexpr SIMDLevel SINGLE_SIMD_LEVEL = SIMDLevel::NONE; #else #if defined(COMPILE_SIMD_AVX512_SPR) inline constexpr SIMDLevel SINGLE_SIMD_LEVEL = SIMDLevel::AVX512_SPR; +#elif defined(COMPILE_SIMD_AVX512_VPOPCNT) +inline constexpr SIMDLevel SINGLE_SIMD_LEVEL = SIMDLevel::AVX512_VPOPCNT; #elif defined(COMPILE_SIMD_AVX512) inline constexpr SIMDLevel SINGLE_SIMD_LEVEL = SIMDLevel::AVX512; #elif defined(COMPILE_SIMD_AVX2) @@ -71,7 +77,7 @@ inline constexpr SIMDLevel SINGLE_SIMD_LEVEL = SIMDLevel::NONE; * Helper to select the appropriate 256-bit SIMD level. * * For 256-bit SIMD types (simd16uint16, simd32uint8, etc.), maps: - * AVX512/AVX512_SPR → AVX2 (256-bit ops use AVX2 instructions) + * AVX512/AVX512_VPOPCNT/AVX512_SPR → AVX2 * AVX2 → AVX2 * ARM_NEON/ARM_SVE → ARM_NEON * NONE → NONE @@ -79,7 +85,8 @@ inline constexpr SIMDLevel SINGLE_SIMD_LEVEL = SIMDLevel::NONE; template struct simd256_level_selector { static constexpr SIMDLevel value = - (SL == SIMDLevel::AVX512 || SL == SIMDLevel::AVX512_SPR) + (SL == SIMDLevel::AVX512 || SL == SIMDLevel::AVX512_VPOPCNT || + SL == SIMDLevel::AVX512_SPR) ? SIMDLevel::AVX2 : (SL == SIMDLevel::ARM_SVE ? SIMDLevel::ARM_NEON : SL == SIMDLevel::RISCV_RVV ? SIMDLevel::NONE @@ -96,13 +103,14 @@ inline constexpr SIMDLevel SINGLE_SIMD_LEVEL_256 = * Helper to select the appropriate 512-bit SIMD level. * * For 512-bit SIMD types (simd32uint16, simd64uint8, etc.), maps: - * AVX512_SPR → AVX512 (512-bit ops share the same instructions) + * AVX512_VPOPCNT/AVX512_SPR → AVX512 for generic 512-bit operations * AVX512 → AVX512 * NONE → NONE ***************************************************************/ template struct simd512_level_selector { - static constexpr SIMDLevel value = (SL == SIMDLevel::AVX512_SPR) + static constexpr SIMDLevel value = + (SL == SIMDLevel::AVX512_SPR || SL == SIMDLevel::AVX512_VPOPCNT) ? SIMDLevel::AVX512 : (SL == SIMDLevel::RISCV_RVV) ? SIMDLevel::NONE : SL; @@ -124,14 +132,21 @@ constexpr int simd_width() { static_assert( SL != SIMDLevel::RISCV_RVV, "simd_width is not supported: RVV is variable-width"); - if constexpr (SL == SIMDLevel::AVX512 || SL == SIMDLevel::AVX512_SPR) + if constexpr ( + SL == SIMDLevel::AVX512 || SL == SIMDLevel::AVX512_VPOPCNT || + SL == SIMDLevel::AVX512_SPR) { return 16; - else if constexpr (SL == SIMDLevel::AVX2 || SL == SIMDLevel::ARM_NEON) + } else if constexpr (SL == SIMDLevel::AVX2 || SL == SIMDLevel::ARM_NEON) { return 8; - else + } else { return 1; + } } +/// Bitmask (1 << SIMDLevel) of the levels this binary holds code for, +/// whether or not the CPU supports them. +uint64_t compiled_simd_levels(); + /// Convert SIMDLevel to string. Throws FaissException for invalid level. std::string to_string(SIMDLevel level); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 28535a58e..dc6e80e21 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -71,6 +71,7 @@ if(FAISS_OPT_LEVEL STREQUAL "dd") if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86_64|amd64|AMD64)") list(APPEND FAISS_TEST_SRC test_rabitq_simd.cpp + test_rabitq_simd_vpopcnt.cpp test_simd_levels_x86_avx2.cpp test_simd_levels_x86_avx512.cpp ) diff --git a/tests/common_faiss_tests.py b/tests/common_faiss_tests.py index d65caf89a..02b33ec82 100644 --- a/tests/common_faiss_tests.py +++ b/tests/common_faiss_tests.py @@ -148,6 +148,7 @@ _SIMD_LEVEL_NAMES = { faiss.SIMDLevel_ARM_NEON: "ARM_NEON", faiss.SIMDLevel_ARM_SVE: "ARM_SVE", faiss.SIMDLevel_RISCV_RVV: "RISCV_RVV", + faiss.SIMDLevel_AVX512_VPOPCNT: "AVX512_VPOPCNT", } diff --git a/tests/test_distances_dispatch.cpp b/tests/test_distances_dispatch.cpp index 1e2782e09..91e4911ab 100644 --- a/tests/test_distances_dispatch.cpp +++ b/tests/test_distances_dispatch.cpp @@ -61,6 +61,7 @@ std::vector available_levels() { SIMDLevel::AVX2, SIMDLevel::AVX512, SIMDLevel::AVX512_SPR, + SIMDLevel::AVX512_VPOPCNT, SIMDLevel::ARM_NEON, SIMDLevel::ARM_SVE, }; diff --git a/tests/test_rabitq_simd.cpp b/tests/test_rabitq_simd.cpp index 35c956994..5da71a872 100644 --- a/tests/test_rabitq_simd.cpp +++ b/tests/test_rabitq_simd.cpp @@ -39,8 +39,23 @@ static std::vector random_bytes(size_t n, uint32_t seed) { } // 32-d chunks and chunk boundaries. -static const std::vector kDims = - {1, 8, 16, 31, 32, 33, 255, 256, 257, 512, 1024, 2048}; +static const std::vector kDims = { + 1, + 8, + 16, + 31, + 32, + 33, + 100, + 128, + 255, + 256, + 257, + 384, + 512, + 768, + 1024, + 2048}; template static void check_quantization_matches_scalar() { diff --git a/tests/test_rabitq_simd_vpopcnt.cpp b/tests/test_rabitq_simd_vpopcnt.cpp new file mode 100644 index 000000000..81082404c --- /dev/null +++ b/tests/test_rabitq_simd_vpopcnt.cpp @@ -0,0 +1,102 @@ +/* + * 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. + */ + +/* + * AVX512_VPOPCNT RaBitQ kernels against the scalar reference. + * + * Separate from test_rabitq_simd.cpp because the AVX512_VPOPCNT + * specializations only exist in a dynamic dispatch build, and that file also + * builds in static mode. The target is dynamic dispatch only. + */ + +#include + +#include +#include +#include +#include + +#include +#include + +using faiss::SIMDLevel; + +namespace { + +std::vector random_bytes(size_t n, uint32_t seed) { + std::mt19937 rng(seed); + std::vector v(n); + for (size_t i = 0; i < n; i++) { + v[i] = static_cast(rng()); + } + return v; +} + +// Widths on either side of every register boundary the kernels step over. +const std::vector kDims = { + 1, + 8, + 16, + 31, + 32, + 33, + 100, + 128, + 255, + 256, + 257, + 384, + 512, + 768, + 1024, + 2048}; + +} // namespace + +TEST(RaBitQVpopcnt, BitwiseKernelsMatchScalarAcrossTails) { + if (!faiss::SIMDConfig::is_simd_level_available( + SIMDLevel::AVX512_VPOPCNT)) { + GTEST_SKIP() << "AVX512_VPOPCNT is not available on this CPU"; + } + + for (size_t d : kDims) { + for (size_t qb = 1; qb <= 8; qb++) { + const size_t size = (d + 7) / 8; + const auto data = random_bytes(size, 19717); + const auto q = random_bytes(size * qb, 51691); + + const uint64_t expected_and = + faiss::rabitq::bitwise_and_dot_product( + q.data(), data.data(), size, qb); + const uint64_t expected_xor = + faiss::rabitq::bitwise_xor_dot_product( + q.data(), data.data(), size, qb); + const uint64_t expected_pop = + faiss::rabitq::popcount(data.data(), size); + + const uint64_t actual_and = faiss::rabitq::bitwise_and_dot_product< + SIMDLevel::AVX512_VPOPCNT>(q.data(), data.data(), size, qb); + const uint64_t actual_xor = faiss::rabitq::bitwise_xor_dot_product< + SIMDLevel::AVX512_VPOPCNT>(q.data(), data.data(), size, qb); + const uint64_t actual_pop = + faiss::rabitq::popcount( + data.data(), size); + const auto actual_fused = + faiss::rabitq::bitwise_and_dot_product_with_popcount< + SIMDLevel::AVX512_VPOPCNT>( + q.data(), data.data(), size, qb); + + EXPECT_EQ(actual_and, expected_and) << "d=" << d << " qb=" << qb; + EXPECT_EQ(actual_xor, expected_xor) << "d=" << d << " qb=" << qb; + EXPECT_EQ(actual_pop, expected_pop) << "d=" << d << " qb=" << qb; + EXPECT_EQ(actual_fused.dot_product, expected_and) + << "d=" << d << " qb=" << qb; + EXPECT_EQ(actual_fused.popcount, expected_pop) + << "d=" << d << " qb=" << qb; + } + } +} diff --git a/tests/test_simd_dispatch.py b/tests/test_simd_dispatch.py index 12851777e..16e0e4e4f 100644 --- a/tests/test_simd_dispatch.py +++ b/tests/test_simd_dispatch.py @@ -38,26 +38,52 @@ def get_cpu_flags(): return set() +def is_simd_level_compiled(name): + """Whether this build holds code for the named level. + + is_simd_level_available() reports compiled and CPU-supported together, + so it cannot tell a level the build left out from one the CPU lacks. + """ + import faiss + + level = getattr(faiss, f"SIMDLevel_{name}", None) + if level is None: + return False + return ((faiss.compiled_simd_levels() >> int(level)) & 1) == 1 + + def get_available_simd_levels(): """ - Returns SIMD levels that are available on the current platform. - NONE is always available. Others depend on architecture. + Returns the SIMD levels this build can be forced to with FAISS_SIMD_LEVEL. + + Asks faiss rather than reading /proc/cpuinfo: a level is only honoured + when it is compiled into the build, and a CPU feature does not imply + that. is_simd_level_available() reports compiled and supported, which is + a subset of compiled, so every level it returns is one the override + accepts. + + Setting FAISS_SIMD_LEVEL collapses the supported set to the forced level, + so the caller must not already have it set. """ - import platform - - arch = platform.machine().lower() - - # SIMDLevel enum names - levels = ["NONE"] - - if arch in ("x86_64", "amd64"): - levels.extend(["AVX2", "AVX512"]) - # AVX512_SPR is typically not enabled in DD builds - elif arch in ("aarch64", "arm64"): - levels.append("ARM_NEON") - elif arch in ("riscv64", "riscv"): - levels.append("RISCV_RVV") + import faiss + names = [ + "NONE", + "AVX2", + "AVX512", + "AVX512_VPOPCNT", + "AVX512_SPR", + "ARM_NEON", + "ARM_SVE", + "RISCV_RVV", + ] + levels = [] + for name in names: + level = getattr(faiss, f"SIMDLevel_{name}", None) + if level is not None and faiss.SIMDConfig.is_simd_level_available( + level + ): + levels.append(name) return levels @@ -109,6 +135,12 @@ class TestSIMDDispatch(unittest.TestCase): except ImportError: self.skipTest("faiss not available") + if os.environ.get("FAISS_SIMD_LEVEL"): + self.skipTest( + "FAISS_SIMD_LEVEL is already set, which collapses the " + "supported set" + ) + levels = get_available_simd_levels() # Get the PYTHONPATH that includes faiss module @@ -224,7 +256,7 @@ for lvl in range(int(faiss.SIMDLevel_COUNT)): """SPR detection must agree with the CPU's real feature flags. The SPR code path is compiled with -mavx512fp16, so AVX512_SPR must be reported available if and only if the CPU actually has the full - AVX512 core feature set AND AVX512_BF16 AND AVX512_FP16. + AVX512 core feature set, VNNI, VPOPCNTDQ, BF16, and FP16. """ import platform @@ -241,6 +273,15 @@ for lvl in range(int(faiss.SIMDLevel_COUNT)): if platform.machine().lower() not in ("x86_64", "amd64"): self.skipTest("x86_64-only test") + if not is_simd_level_compiled("AVX512_SPR"): + self.skipTest("AVX512_SPR is not compiled into this build") + + if os.environ.get("FAISS_SIMD_LEVEL"): + self.skipTest( + "FAISS_SIMD_LEVEL is already set, which collapses the " + "supported set" + ) + flags = get_cpu_flags() if flags is None: self.skipTest("/proc/cpuinfo not available") @@ -255,6 +296,8 @@ for lvl in range(int(faiss.SIMDLevel_COUNT)): } spr_capable = ( avx512_core <= flags + and "avx512_vnni" in flags + and "avx512_vpopcntdq" in flags and "avx512_bf16" in flags and "avx512_fp16" in flags ) @@ -269,11 +312,65 @@ for lvl in range(int(faiss.SIMDLevel_COUNT)): "AVX512_SPR detection disagrees with /proc/cpuinfo: " f"detected={spr_detected}, cpu_spr_capable={spr_capable} " f"(avx512_core={avx512_core <= flags}, " - f"bf16={'avx512_bf16' in flags}, fp16={'avx512_fp16' in flags}). " + f"vnni={'avx512_vnni' in flags}, " + f"vpopcntdq={'avx512_vpopcntdq' in flags}, " + f"bf16={'avx512_bf16' in flags}, " + f"fp16={'avx512_fp16' in flags}). " "detected=True with fp16=False is the D107684495 regression " "(AMD Zen 4 mis-detected as SPR).", ) + def test_vpopcnt_detection_matches_cpu_features(self): + """The VPOPCNT level requires baseline AVX-512 and VPOPCNTDQ.""" + import platform + + try: + import faiss + + if "DD" not in faiss.get_compile_options(): + self.skipTest( + "Not a DD build - SIMD level is fixed at compile time" + ) + except ImportError: + self.skipTest("faiss not available") + + if platform.machine().lower() not in ("x86_64", "amd64"): + self.skipTest("x86_64-only test") + + if not is_simd_level_compiled("AVX512_VPOPCNT"): + self.skipTest("AVX512_VPOPCNT is not compiled into this build") + + if os.environ.get("FAISS_SIMD_LEVEL"): + self.skipTest( + "FAISS_SIMD_LEVEL is already set, which collapses the " + "supported set" + ) + + flags = get_cpu_flags() + if flags is None: + self.skipTest("/proc/cpuinfo not available") + + avx512_core = { + "avx512f", + "avx512cd", + "avx512vl", + "avx512dq", + "avx512bw", + } + capable = ( + avx512_core <= flags and "avx512_vpopcntdq" in flags + ) + detected = faiss.SIMDConfig.is_simd_level_available( + faiss.SIMDLevel_AVX512_VPOPCNT + ) + + self.assertEqual( + detected, + capable, + "AVX512_VPOPCNT detection disagrees with /proc/cpuinfo: " + f"detected={detected}, cpu_capable={capable}", + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_simd_levels.cpp b/tests/test_simd_levels.cpp index 881523bee..48a58a34d 100644 --- a/tests/test_simd_levels.cpp +++ b/tests/test_simd_levels.cpp @@ -118,6 +118,9 @@ TEST(SIMDLevel, to_string_all_levels) { EXPECT_EQ("NONE", faiss::to_string(faiss::SIMDLevel::NONE)); EXPECT_EQ("AVX2", faiss::to_string(faiss::SIMDLevel::AVX2)); EXPECT_EQ("AVX512", faiss::to_string(faiss::SIMDLevel::AVX512)); + EXPECT_EQ( + "AVX512_VPOPCNT", + faiss::to_string(faiss::SIMDLevel::AVX512_VPOPCNT)); EXPECT_EQ("AVX512_SPR", faiss::to_string(faiss::SIMDLevel::AVX512_SPR)); EXPECT_EQ("ARM_NEON", faiss::to_string(faiss::SIMDLevel::ARM_NEON)); EXPECT_EQ("ARM_SVE", faiss::to_string(faiss::SIMDLevel::ARM_SVE)); @@ -132,6 +135,9 @@ TEST(SIMDLevel, to_simd_level_all_strings) { EXPECT_EQ(faiss::SIMDLevel::NONE, faiss::to_simd_level("NONE")); EXPECT_EQ(faiss::SIMDLevel::AVX2, faiss::to_simd_level("AVX2")); EXPECT_EQ(faiss::SIMDLevel::AVX512, faiss::to_simd_level("AVX512")); + EXPECT_EQ( + faiss::SIMDLevel::AVX512_VPOPCNT, + faiss::to_simd_level("AVX512_VPOPCNT")); EXPECT_EQ(faiss::SIMDLevel::AVX512_SPR, faiss::to_simd_level("AVX512_SPR")); EXPECT_EQ(faiss::SIMDLevel::ARM_NEON, faiss::to_simd_level("ARM_NEON")); EXPECT_EQ(faiss::SIMDLevel::ARM_SVE, faiss::to_simd_level("ARM_SVE"));