Files
faiss/tests/CMakeLists.txt
T
ChenLiqing ea29bac75d RaBitQ: add batch-four estimates and 9-bit SIMD scoring (#5647)
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
2026-10-02 07:58:02 -07:00

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()