mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Summary: Extend RaBitQ's distance-computation API with four independent one-bit estimates and add optimized kernels for standard RaBitQ codes. - Add `RaBitQDistanceComputer::distance_to_code_1bit_batch_4`, with a default implementation that delegates to four single-code evaluations. - For non-centered, four-bit query quantization, reuse query bit planes across four database codes. The optimized kernel uses the existing `AVX512_VPOPCNT` dispatch level; other SIMD levels and query modes retain fallbacks. - Add AVX2 and AVX-512 full-distance kernels for 9-bit database codes, exploiting byte-aligned extra codes rather than scalar extraction. The scalar tail reads only the required extra bytes. - Add `-mbmi2` to the x86 SIMD compile flags in `xplat.bzl` and `CMakeLists.txt`. The `AVX2` and higher SIMD levels now require CPUID BMI2. The RaBitQ kernels call `_pext_u64` without `__BMI2__` guards. - Use local fixed-size `memcpy` for byte-aligned scalar bit-plane loads. ## API contract and scope Call `set_query()` before evaluating a batch. The four row pointers may be non-contiguous or repeated. Results are returned in input order. The method computes estimates only: it does not refine candidates, update thresholds, or increment search counters. Existing query modes and database precisions remain supported. No changes to HNSW traversal, filtering, graph construction, defaults, or serialized code layout are included. This is an opt-in batch API; existing search call sites are unchanged. ## Tests and benchmark New standalone tests cover batch/single agreement for database bits 1–9, query bits 0–8, both centering modes, L2/IP, available SIMD overrides, repeated/non-contiguous row pointers, query reuse, and unchanged statistics. Independent references check packed multi-bit values and unaligned/tail buffers, including exactly sized RBQ9 extra codes. A round-trip test checks stored codes and distances. `bench_rabitq_batch_estimates` compares four single-code estimates with the batch API at dimensions 128, 768, 960, and 1536. It prepares queries outside timing, shuffles code pointers, warms both paths, alternates measurement order, and prints per-code times and checksums. It is a scorer microbenchmark, not an end-to-end search benchmark. Local validation (GCC 12, Release dynamic-dispatch build, Intel Xeon Platinum 8375C): - Built `faiss_test` and `bench_rabitq_batch_estimates`. - 61 tests passed with `--gtest_filter='TestLowLevelIVF.*:*RaBitQ*:*HNSW*'`, including all five new tests. - 47 individual-process tests passed with `ctest -R 'RaBitQ|HNSW'`. - All five new tests passed with the changed RaBitQ translation units and test file instrumented by ASan/UBSan, including leak detection. Unchanged dependencies were linked from the Release build; this was not a fully instrumented library build. - Checked NONE, AVX2, AVX512 and AVX512_VPOPCNT on the local host. Clang 15 syntax checks passed for the three affected SIMD units, with baseline BMI2 disabled for the AVX2/AVX512 checks. ARM, non-BMI2 hardware, and native SPR execution were not performed locally. A CPU without BMI2 now uses the `NONE` level. - No on-disk format changes; code and distance round-trip checks passed. Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5647 Reviewed By: junjieqi Differential Revision: D122390129 Pulled By: mnorris11 fbshipit-source-id: f3f1b1817e3fbcbf7896f3b4cf86eebcbba6f6b2
192 lines
6.2 KiB
CMake
192 lines
6.2 KiB
CMake
# 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.
|
|
|
|
set(FAISS_TEST_SRC
|
|
test_binary_flat.cpp
|
|
test_dealloc_invlists.cpp
|
|
test_ivfpq_codec.cpp
|
|
test_ivfpq_indexing.cpp
|
|
test_lowlevel_ivf.cpp
|
|
test_ivf_index.cpp
|
|
test_omp_threads.cpp
|
|
test_pairs_decoding.cpp
|
|
test_params_override.cpp
|
|
test_pq_encoding.cpp
|
|
test_sliding_ivf.cpp
|
|
test_threaded_index.cpp
|
|
test_transfer_invlists.cpp
|
|
test_cppcontrib_sa_decode.cpp
|
|
test_cppcontrib_uintreader.cpp
|
|
test_simdlib.cpp
|
|
test_approx_topk.cpp
|
|
test_RCQ_cropping.cpp
|
|
test_heap.cpp
|
|
test_pq_code_distance.cpp
|
|
test_hnsw.cpp
|
|
test_NSG_compressed_graph.cpp
|
|
test_partitioning.cpp
|
|
test_fastscan_perf.cpp
|
|
test_pqfs_unaligned.cpp
|
|
test_common_ivf_empty_index.cpp
|
|
test_ivf_early_termination.cpp
|
|
test_callback.cpp
|
|
test_utils.cpp
|
|
test_hamming.cpp
|
|
test_zerocopy.cpp
|
|
test_scalar_quantizer.cpp
|
|
test_factory_tools.cpp
|
|
test_custom_result_handler.cpp
|
|
test_fastscan_filter.cpp
|
|
test_single_query_collect_handler.cpp
|
|
# These tests work in both static and DD modes (uniform SIMDConfig API)
|
|
test_distances_simd.cpp
|
|
test_simd_levels.cpp
|
|
test_fast_scan_distance_to_code.cpp
|
|
test_rabitq_fastscan.cpp
|
|
test_rabitq_batch_estimates.cpp
|
|
test_super_kmeans_foundations.cpp
|
|
)
|
|
|
|
# Tests that use POSIX APIs or have platform-specific issues on Windows
|
|
if(NOT WIN32)
|
|
list(APPEND FAISS_TEST_SRC
|
|
test_merge.cpp
|
|
test_ondisk_ivf.cpp
|
|
test_disable_pq_sdc_tables.cpp
|
|
test_mmap.cpp
|
|
)
|
|
endif()
|
|
|
|
if(FAISS_ENABLE_SVS)
|
|
list(APPEND FAISS_TEST_SRC test_svs.cpp)
|
|
endif()
|
|
|
|
# DD-only tests; x86-specific raw-intrinsic tests are added separately below.
|
|
if(FAISS_OPT_LEVEL STREQUAL "dd")
|
|
list(APPEND FAISS_TEST_SRC
|
|
test_distances_dispatch.cpp
|
|
test_simd_perf.cpp
|
|
)
|
|
if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86_64|amd64|AMD64)")
|
|
# Enable VPOPCNT dispatch for this test. The common
|
|
# target enables AVX2/AVX512 only.
|
|
set_source_files_properties(test_rabitq_batch_estimates.cpp
|
|
PROPERTIES COMPILE_DEFINITIONS COMPILE_SIMD_AVX512_VPOPCNT)
|
|
list(APPEND FAISS_TEST_SRC
|
|
test_rabitq_simd.cpp
|
|
test_rabitq_simd_dd.cpp
|
|
test_simd_levels_x86_avx2.cpp
|
|
test_simd_levels_x86_avx512.cpp
|
|
)
|
|
# Set SIMD compile flags per file
|
|
# test_simd_levels_x86_avx2.cpp and test_simd_levels_x86_avx512.cpp
|
|
# contain raw intrinsics and need their respective SIMD flags.
|
|
if(NOT WIN32)
|
|
set_source_files_properties(test_simd_levels_x86_avx2.cpp
|
|
PROPERTIES COMPILE_OPTIONS "-mavx2;-mfma;-mf16c;-mpopcnt"
|
|
)
|
|
set_source_files_properties(test_simd_levels_x86_avx512.cpp
|
|
PROPERTIES COMPILE_OPTIONS "-mavx2;-mfma;-mf16c;-mpopcnt;-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw"
|
|
)
|
|
else()
|
|
# MSVC
|
|
set_source_files_properties(test_simd_levels_x86_avx2.cpp
|
|
PROPERTIES COMPILE_OPTIONS "/arch:AVX2"
|
|
)
|
|
set_source_files_properties(test_simd_levels_x86_avx512.cpp
|
|
PROPERTIES COMPILE_OPTIONS "/arch:AVX512"
|
|
)
|
|
endif()
|
|
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(aarch64|arm64|ARM64)")
|
|
list(APPEND FAISS_TEST_SRC
|
|
test_rabitq_simd_neon.cpp
|
|
)
|
|
endif()
|
|
endif()
|
|
|
|
# test_sq_rvv_tail_and_semantics.cpp references
|
|
# sq_select_distance_computer<SIMDLevel::RISCV_RVV>, whose specialization is
|
|
# only instantiated by sq-rvv.cpp when the RVV implementation is compiled
|
|
# (riscv64 builds). Referencing it unconditionally breaks linking on other
|
|
# architectures — a runtime GTEST_SKIP() cannot remove the link-time symbol
|
|
# requirement — so add the source only inside the RISC-V condition. It uses
|
|
# raw RVV intrinsics (guarded in-source by __riscv && COMPILE_SIMD_RISCV_RVV).
|
|
# The faiss target's COMPILE_SIMD_RISCV_RVV is PRIVATE and does not propagate,
|
|
# so give this TU the same flags the library's RVV sources get — otherwise the
|
|
# compile-time-intrinsic cases (D5b) skip on riscv64.
|
|
if(CMAKE_SYSTEM_PROCESSOR MATCHES "(riscv64|riscv)")
|
|
list(APPEND FAISS_TEST_SRC test_sq_rvv_tail_and_semantics.cpp)
|
|
set_source_files_properties(test_sq_rvv_tail_and_semantics.cpp
|
|
PROPERTIES
|
|
COMPILE_OPTIONS "-march=rv64gcv_zvfhmin;-mabi=lp64d"
|
|
COMPILE_DEFINITIONS "COMPILE_SIMD_RISCV_RVV"
|
|
)
|
|
endif()
|
|
|
|
add_executable(faiss_test ${FAISS_TEST_SRC})
|
|
|
|
if(FAISS_ENABLE_SVS)
|
|
target_compile_definitions(faiss_test PRIVATE FAISS_ENABLE_SVS FAISS_SVS_RUNTIME_VERSION=${FAISS_SVS_RUNTIME_VERSION})
|
|
endif()
|
|
|
|
include(../cmake/link_to_faiss_lib.cmake)
|
|
|
|
link_to_faiss_lib(faiss_test)
|
|
|
|
if (FAISS_ENABLE_PYTHON)
|
|
add_dependencies(faiss_test faiss_example_external_module)
|
|
endif()
|
|
|
|
include(FetchContent)
|
|
FetchContent_Declare(
|
|
googletest
|
|
GIT_REPOSITORY https://github.com/google/googletest.git
|
|
GIT_TAG 58d77fa8070e8cec2dc1ed015d66b454c8d78850 # release-1.12.1
|
|
OVERRIDE_FIND_PACKAGE)
|
|
set(BUILD_GMOCK CACHE BOOL OFF)
|
|
set(INSTALL_GTEST CACHE BOOL OFF)
|
|
FetchContent_MakeAvailable(googletest)
|
|
|
|
if(NOT EXISTS ${CMAKE_FIND_PACKAGE_REDIRECTS_DIR}/gtest-config.cmake
|
|
AND NOT EXISTS ${CMAKE_FIND_PACKAGE_REDIRECTS_DIR}/GTestConfig.cmake)
|
|
file(
|
|
WRITE ${CMAKE_FIND_PACKAGE_REDIRECTS_DIR}/gtest-config.cmake
|
|
[=[
|
|
include(CMakeFindDependencyMacro)
|
|
find_dependency(googletest)
|
|
if(NOT TARGET GTest::GTest)
|
|
add_library(GTest::GTest INTERFACE IMPORTED)
|
|
target_link_libraries(GTest::GTest INTERFACE GTest::gtest)
|
|
endif()
|
|
if(NOT TARGET GTest::Main)
|
|
add_library(GTest::Main INTERFACE IMPORTED)
|
|
target_link_libraries(GTest::Main INTERFACE GTest::gtest_main)
|
|
endif()
|
|
]=])
|
|
endif()
|
|
|
|
find_package(OpenMP REQUIRED)
|
|
find_package(GTest CONFIG REQUIRED)
|
|
|
|
target_link_libraries(faiss_test PRIVATE
|
|
OpenMP::OpenMP_CXX
|
|
GTest::gtest_main
|
|
$<$<BOOL:${FAISS_ENABLE_ROCM}>:hip::host>
|
|
)
|
|
|
|
# Defines `gtest_discover_tests()`.
|
|
include(GoogleTest)
|
|
if(CMAKE_CROSSCOMPILING)
|
|
# Skip test discovery when cross-compiling: the binary cannot run on the
|
|
# host, so gtest_discover_tests() would fail. The build itself verifies
|
|
# that the binary compiles and links correctly.
|
|
elseif(WIN32)
|
|
# On Windows, defer test discovery to ctest time (PRE_TEST) instead of
|
|
# build time (POST_BUILD) because faiss.dll is not in PATH during build.
|
|
gtest_discover_tests(faiss_test DISCOVERY_MODE PRE_TEST)
|
|
else()
|
|
gtest_discover_tests(faiss_test)
|
|
endif()
|