From e6bd1d252caa5b9315a04390269d382845984ae2 Mon Sep 17 00:00:00 2001 From: lyang24 Date: Thu, 10 Sep 2026 09:32:36 -0700 Subject: [PATCH] Add AVX512 VPOPCNT dynamic dispatch level (#5531) Summary: ## What Split the existing VPOPCNTDQ RaBitQ and Hamming kernels from the full Sapphire Rapids SIMD level. DD builds now expose an `AVX512_VPOPCNT` capability for CPUs such as Ice Lake and Zen 4. The fallback chain is: `AVX512_SPR -> AVX512_VPOPCNT -> AVX512 -> AVX2 -> NONE` The ordinary static AVX512 target remains unchanged and contains no VPOPCNT instructions. ## Why The VPOPCNT kernels only require baseline AVX-512 plus `AVX512_VPOPCNTDQ`. Tying them to the SPR level unnecessarily excluded CPUs that support VPOPCNTDQ but not AVX512-FP16, notably Zen 4. ## Changes made while importing Two changes were needed on top of the pull request. **Feature detection read the wrong CPUID leaf.** `AVX512_VPOPCNTDQ` and `AVX512_VNNI` were read from leaf 1 ECX. Both live in leaf 7 subleaf 0 ECX, at bits 14 and 11. Leaf 1 ECX holds unrelated bits at those positions, so the new level was never selected on any CPU. Measured on an AMD Genoa host before the fix: the level reported as unavailable and detection returned `AVX512`. After the fix the same host reports `AVX512_VPOPCNT`. **The build definitions outside CMake were not updated.** The pull request renames three files. The other build description still named the old paths and had no entry for the new level, so the library did not build. This adds the level constant, its compiler flags, its `COMPILE_SIMD_AVX512_VPOPCNT` define, and the renamed paths. **The new level was not registered in the compiled-level mask.** `compiled_simd_levels()` reports which levels a binary holds, and `FAISS_SIMD_LEVEL` walks down from the requested level until it finds one that is present. The new level had no entry, so the override always stepped past it to `AVX512` and printed that the level was not compiled in. This was reported by `test_dispatch_with_env_var` on the open-source dynamic-dispatch build. **The test helper inferred available levels from `/proc/cpuinfo`.** A CPU feature does not imply that a level is compiled into the build, so the helper listed a level the override then refused. It now asks `SIMDConfig::is_simd_level_available()`, which reports compiled and supported. That is a subset of compiled, so every level it returns is one the override accepts. The mask and helper names also follow the naming used in the parent diff, so `AVAILABLE_SIMD_LEVELS_BASE_WITH_VPOPCNT` and `with_simd_level_with_vpopcnt`. Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5531 Test Plan: ``` buck2 test fbcode//mode/opt -c faiss.dynamic_dispatch=true \ fbcode//faiss/tests:test_simd_levels fbcode//faiss/tests:test_rabitq_simd ``` Result: all pass. The build failures the imported version reported are cleared. All three configurations build: ``` buck2 build fbcode//mode/dev fbcode//faiss:faiss fbcode//faiss:faiss_no_multithreading fbcode//faiss:faiss_omp_mock buck2 build fbcode//mode/opt -c faiss.dynamic_dispatch=true buck2 build fbcode//mode/dev -c fbcode.arch=aarch64 ``` Detection on an AMD Genoa host, which has AVX512_VPOPCNTDQ and AVX512_BF16 but not AVX512_FP16: | | before the CPUID fix | after | | --- | --- | --- | | detected level | `AVX512` | `AVX512_VPOPCNT` | | `AVX512_VPOPCNT` available | 0 | 1 | | `AVX512_SPR` available | 0 | 0 | Reviewed By: alibeklfc Differential Revision: D118969754 Pulled By: mnorris11 fbshipit-source-id: d1ca52b614829a3d4b322c9fecf0c004f2174dab --- benchs/bench_rabitq.py | 4 +- faiss/CMakeLists.txt | 57 +++++--- faiss/docs/simd_dynamic_dispatch_migration.md | 7 +- faiss/impl/RaBitQUtils.cpp | 2 +- faiss/impl/RaBitQuantizer.cpp | 10 +- faiss/impl/fast_scan/decompose_qbs.h | 1 + faiss/impl/scalar_quantizer/sq-dispatch.h | 42 ++++-- faiss/impl/simd_dispatch.h | 42 ++++-- faiss/python/__init__.pyi | 3 + faiss/utils/hamming.cpp | 16 +-- ...512_spr.cpp => hamming_avx512_vpopcnt.cpp} | 8 +- .../hamming_computer-avx512.h | 2 +- ...pr.h => hamming_computer-avx512_vpopcnt.h} | 46 +++--- ...x512_spr.cpp => rabitq_avx512_vpopcnt.cpp} | 45 +++--- faiss/utils/simd_levels.cpp | 58 +++++--- faiss/utils/simd_levels.h | 31 ++-- tests/CMakeLists.txt | 1 + tests/common_faiss_tests.py | 1 + tests/test_distances_dispatch.cpp | 1 + tests/test_rabitq_simd.cpp | 19 ++- tests/test_rabitq_simd_vpopcnt.cpp | 102 ++++++++++++++ tests/test_simd_dispatch.py | 133 +++++++++++++++--- tests/test_simd_levels.cpp | 6 + 23 files changed, 468 insertions(+), 169 deletions(-) rename faiss/utils/hamming_distance/{hamming_avx512_spr.cpp => hamming_avx512_vpopcnt.cpp} (59%) rename faiss/utils/hamming_distance/{hamming_computer-avx512_spr.h => hamming_computer-avx512_vpopcnt.h} (78%) rename faiss/utils/simd_impl/{rabitq_avx512_spr.cpp => rabitq_avx512_vpopcnt.cpp} (91%) create mode 100644 tests/test_rabitq_simd_vpopcnt.cpp 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"));