From ea29bac75d0039ad9f3aa0351ac5cac6aabdb0bb Mon Sep 17 00:00:00 2001 From: ChenLiqing <23721160+CLiqing@users.noreply.github.com> Date: Fri, 2 Oct 2026 07:58:02 -0700 Subject: [PATCH] RaBitQ: add batch-four estimates and 9-bit SIMD scoring (#5647) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- benchs/CMakeLists.txt | 5 + benchs/bench_rabitq_batch_estimates.cpp | 96 ++++++ faiss/CMakeLists.txt | 8 +- faiss/impl/RaBitQUtils.h | 2 +- faiss/impl/RaBitQuantizer.cpp | 37 +++ faiss/impl/RaBitQuantizer.h | 13 + faiss/utils/rabitq_simd.h | 67 ++++- faiss/utils/simd_impl/rabitq_avx2.cpp | 52 +++- faiss/utils/simd_impl/rabitq_avx512.cpp | 55 ++-- .../utils/simd_impl/rabitq_avx512_vpopcnt.cpp | 61 +++- faiss/utils/simd_levels.cpp | 6 +- tests/CMakeLists.txt | 5 + tests/test_rabitq_batch_estimates.cpp | 278 ++++++++++++++++++ 13 files changed, 624 insertions(+), 61 deletions(-) create mode 100644 benchs/bench_rabitq_batch_estimates.cpp create mode 100644 tests/test_rabitq_batch_estimates.cpp diff --git a/benchs/CMakeLists.txt b/benchs/CMakeLists.txt index 478691d66..7c27804ae 100644 --- a/benchs/CMakeLists.txt +++ b/benchs/CMakeLists.txt @@ -7,6 +7,11 @@ find_package(BLAS REQUIRED) find_package(LAPACK REQUIRED) find_package(OpenMP REQUIRED) +add_executable(bench_rabitq_batch_estimates EXCLUDE_FROM_ALL + bench_rabitq_batch_estimates.cpp) +include(../cmake/link_to_faiss_lib.cmake) +link_to_faiss_lib(bench_rabitq_batch_estimates) + add_executable(bench_ivf_selector EXCLUDE_FROM_ALL bench_ivf_selector.cpp) target_link_libraries(bench_ivf_selector PRIVATE faiss_avx512 ${BLAS_LIBRARIES} ${LAPACK_LIBRARIES} OpenMP::OpenMP_CXX) target_compile_options(bench_ivf_selector PRIVATE $<$:-mavx2 -mfma -mf16c -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mpopcnt>) diff --git a/benchs/bench_rabitq_batch_estimates.cpp b/benchs/bench_rabitq_batch_estimates.cpp new file mode 100644 index 000000000..28cab0336 --- /dev/null +++ b/benchs/bench_rabitq_batch_estimates.cpp @@ -0,0 +1,96 @@ +/* + * 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. + */ + +// Microbenchmark: four single estimates against the batch API. +// Build bench_rabitq_batch_estimates. Run it with OMP_NUM_THREADS=1 +// and OPENBLAS_NUM_THREADS=1. +#include +#include +#include +#include +#include +#include + +#include + +int main() { + constexpr int n = 4096, repetitions = 100; + std::mt19937 rng(123); + std::normal_distribution normal; + std::printf("dim,round,single_ns_per_code,batch_ns_per_code,speedup\n"); + for (int d : {128, 768, 960, 1536}) { + std::vector data(size_t(n) * d), query(d); + for (auto& v : data) { + v = normal(rng); + } + for (auto& v : query) { + v = normal(rng); + } + faiss::IndexRaBitQ index(d, faiss::METRIC_L2, 8); + index.train(n, data.data()); + index.add(n, data.data()); + std::unique_ptr owner( + index.get_quantized_distance_computer(4, false)); + auto& dc = dynamic_cast(*owner); + dc.set_query(query.data()); + std::vector codes(n); + for (int i = 0; i < n; ++i) { + codes[i] = index.codes.data() + i * index.code_size; + } + std::shuffle(codes.begin(), codes.end(), rng); + auto measure = [&](bool batch) { + double checksum = 0; + const auto start = std::chrono::steady_clock::now(); + for (int r = 0; r < repetitions; ++r) { + for (int i = 0; i < n; i += 4) { + float distances[4]; + if (batch) { + dc.distance_to_code_1bit_batch_4( + codes.data() + i, distances); + } else { + for (int j = 0; j < 4; ++j) { + distances[j] = + dc.distance_to_code_1bit(codes[i + j]); + } + } + for (float distance : distances) { + checksum += distance; + } + } + } + const double ns = std::chrono::duration( + std::chrono::steady_clock::now() - start) + .count(); + std::fprintf( + stderr, + "d=%d batch=%d checksum=%.9g\n", + d, + batch, + checksum); + return ns / (n * repetitions); + }; + measure(false); + measure(true); + for (int round = 0; round < 5; ++round) { + double single, batch; + if (round % 2) { + batch = measure(true); + single = measure(false); + } else { + single = measure(false); + batch = measure(true); + } + std::printf( + "%d,%d,%.3f,%.3f,%.3f\n", + d, + round, + single, + batch, + single / batch); + } + } +} diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index b709bd0b1..f4f89f8d3 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -567,24 +567,24 @@ if(FAISS_OPT_LEVEL STREQUAL "dd") # Per-file SIMD flags (GCC/Clang) set_source_files_properties(${FAISS_SIMD_AVX2_SRC} TARGET_DIRECTORY faiss - PROPERTIES COMPILE_OPTIONS "-mavx2;-mfma;-mf16c;-mpopcnt" + PROPERTIES COMPILE_OPTIONS "-mavx2;-mfma;-mf16c;-mpopcnt;-mbmi2" ) set_source_files_properties(${FAISS_SIMD_AVX512_SRC} TARGET_DIRECTORY faiss PROPERTIES COMPILE_OPTIONS - "-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mfma;-mf16c;-mpopcnt" + "-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mfma;-mf16c;-mpopcnt;-mbmi2" ) # VPOPCNT sources require VPOPCNTDQ and BITALG on top of 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;-mavx512bitalg;-mpopcnt" + "-mavx2;-mfma;-mf16c;-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mavx512vpopcntdq;-mavx512bitalg;-mpopcnt;-mbmi2" ) # 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 - "-mavx2;-mfma;-mf16c;-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mavx512vpopcntdq;-mavx512bitalg;-mpopcnt;-mavx512vnni;-mavx512fp16;-mavx512bf16" + "-mavx2;-mfma;-mf16c;-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mavx512vpopcntdq;-mavx512bitalg;-mpopcnt;-mavx512vnni;-mavx512fp16;-mavx512bf16;-mbmi2" ) else() # Per-file SIMD flags (MSVC) diff --git a/faiss/impl/RaBitQUtils.h b/faiss/impl/RaBitQUtils.h index 4d46aff6c..2837e2e78 100644 --- a/faiss/impl/RaBitQUtils.h +++ b/faiss/impl/RaBitQUtils.h @@ -353,7 +353,7 @@ inline int extract_code_inline( * The multi-bit distance combines the sign bit (1-bit) with additional * magnitude bits (ex_bits) to compute a more accurate distance estimate. * Uses SIMD-optimized bit-plane decomposition (AVX2+BMI2) for ex_bits 1-7, - * with scalar fallback for non-x86 or non-BMI2 platforms. + * with scalar fallback for non-x86 platforms. * * @param sign_bits unpacked sign bits (1-bit codes in standard format) * @param ex_code packed ex-bit codes diff --git a/faiss/impl/RaBitQuantizer.cpp b/faiss/impl/RaBitQuantizer.cpp index 15fd8a20a..fc1581ed8 100644 --- a/faiss/impl/RaBitQuantizer.cpp +++ b/faiss/impl/RaBitQuantizer.cpp @@ -6,6 +6,7 @@ */ #include +#include #include #include @@ -578,6 +579,42 @@ struct RaBitQDistanceComputerQ final : RaBitQDistanceComputer { return distance_to_code_1bit_impl(code, base_fac, size); } + void distance_to_code_1bit_batch_4( + const uint8_t* const* codes_in, + float* distances) final { + if (qb != 4 || centered) { + RaBitQDistanceComputer::distance_to_code_1bit_batch_4( + codes_in, distances); + return; + } + const size_t size = (d + 7) / 8; + const size_t prefix = size + + (nb_bits == 1 ? sizeof(SignBitFactors) + : sizeof(SignBitFactorsWithError)); + for (int i = 0; i < 4; ++i) { + for (size_t offset = 0; offset < prefix; offset += 64) { + prefetch_L1(codes_in[i] + offset); + } + } + rabitq::BitwiseAndDotProductResult results[4]; + rabitq::bitwise_q4_batch_4( + rearranged_rotated_qq.data(), codes_in, size, results); + for (int i = 0; i < 4; ++i) { + const auto* factors = + reinterpret_cast(codes_in[i] + size); + float final_dot = 0; + final_dot += query_fac.c1 * results[i].dot_product; + final_dot += query_fac.c2 * results[i].popcount; + final_dot -= query_fac.c34; + const float pre_dist = factors->or_minus_c_l2sqr + + query_fac.qr_to_c_L2sqr - + 2 * factors->dp_multiplier * final_dot; + distances[i] = metric_type == METRIC_L2 + ? std::max(0.0f, pre_dist) + : -0.5f * (pre_dist - query_fac.qr_norm_L2sqr); + } + } + // Compute full distance using 1-bit + ex-bits (accurate) float distance_to_code_full(const uint8_t* code) final { FAISS_ASSERT(code != nullptr); diff --git a/faiss/impl/RaBitQuantizer.h b/faiss/impl/RaBitQuantizer.h index 01756a963..7ec23a428 100644 --- a/faiss/impl/RaBitQuantizer.h +++ b/faiss/impl/RaBitQuantizer.h @@ -159,6 +159,19 @@ struct RaBitQDistanceComputer : FlatCodesDistanceComputer { // Compute 1-bit distance estimate (fast) virtual float distance_to_code_1bit(const uint8_t* code) = 0; + /// Compute four one-bit estimates in input order. Caller must call + /// set_query() first. Each pointer addresses one complete encoded row; rows + /// need not be contiguous or distinct. The output holds four floats. + /// Does not refine candidates, update thresholds, or increment stats. The + /// default calls the single estimate four times. + virtual void distance_to_code_1bit_batch_4( + const uint8_t* const* codes_in, + float* distances) { + for (size_t i = 0; i < 4; ++i) { + distances[i] = distance_to_code_1bit(codes_in[i]); + } + } + // Compute full multi-bit distance (accurate) virtual float distance_to_code_full(const uint8_t* code) = 0; diff --git a/faiss/utils/rabitq_simd.h b/faiss/utils/rabitq_simd.h index 5dac1d60a..9f010e0cc 100644 --- a/faiss/utils/rabitq_simd.h +++ b/faiss/utils/rabitq_simd.h @@ -53,6 +53,28 @@ BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount( size_t size, size_t qb); +/// Four independent AND-dot-product/popcount pairs for a four-plane query. +/// Query contains 4 * size bytes; each data pointer contains `size` bytes. +/// Byte-aligned buffers and arbitrary byte counts are supported. +template +inline void bitwise_q4_batch_4( + const uint8_t* query, + const uint8_t* const* data, + size_t size, + BitwiseAndDotProductResult* results) { + for (size_t i = 0; i < 4; ++i) { + results[i] = bitwise_and_dot_product_with_popcount( + query, data[i], size, 4); + } +} + +template <> +void bitwise_q4_batch_4( + const uint8_t* query, + const uint8_t* const* data, + size_t size, + BitwiseAndDotProductResult* results); + /** * Compute dot product between query and binary data using popcount on XOR. * @@ -140,9 +162,11 @@ inline uint64_t bitwise_and_dot_product( uint64_t sum = 0; size_t offset = 0; for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); for (int j = 0; j < qb; j++) { - const auto qv = *(const uint64_t*)(query + j * size + offset); + uint64_t qv; + std::memcpy(&qv, query + j * size + offset, sizeof(qv)); sum += popcount64(qv & yv) << j; } } @@ -167,10 +191,12 @@ inline BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount< uint64_t popcount_sum = 0; size_t offset = 0; for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); popcount_sum += popcount64(yv); for (int j = 0; j < qb; j++) { - const auto qv = *(const uint64_t*)(query + j * size + offset); + uint64_t qv; + std::memcpy(&qv, query + j * size + offset, sizeof(qv)); dot_product += popcount64(qv & yv) << j; } } @@ -194,9 +220,11 @@ inline uint64_t bitwise_xor_dot_product( uint64_t sum = 0; size_t offset = 0; for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); for (int j = 0; j < qb; j++) { - const auto qv = *(const uint64_t*)(query + j * size + offset); + uint64_t qv; + std::memcpy(&qv, query + j * size + offset, sizeof(qv)); sum += popcount64(qv ^ yv) << j; } } @@ -215,7 +243,8 @@ inline uint64_t popcount(const uint8_t* data, size_t size) { uint64_t sum = 0; size_t offset = 0; for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); sum += popcount64(yv); } for (; offset < size; ++offset) { @@ -332,16 +361,12 @@ inline void quantize_query_values( /********************************************************* * Multi-bit RaBitQ inner product kernels. * - * Compute: sum_i rotated_q[i] * ((sign_bit_i << ex_bits) + ex_code_val_i + cb) + * Score = sum_i rotated_q[i] * ((sign_i << ex_bits) + ex_i + cb) * - * Strategy: - * ex_bits == 1: Specialized kernel — both sign_bits and ex_code are - * 1-bit-per-dim packed, enabling direct bit→mask→float - * conversion with zero per-element extraction. - * ex_bits >= 2: Bit-plane decomposition (BMI2 required) — PEXT extracts - * each bit plane in one instruction, then the same - * bit→mask→float kernel computes each plane's dot product. - * Fallback: Scalar extraction via 64-bit window read + shift + mask. + * ex_bits == 1: bits map straight to floats. No extraction. + * ex_bits 2..7: BMI2 PEXT extracts each bit plane. + * ex_bits == 8: whole bytes. No extraction and no BMI2. + * Fallback: scalar math reads one code at a time. *********************************************************/ namespace faiss::rabitq::multibit { @@ -357,6 +382,16 @@ inline float ip_scalar( size_t ex_bits, float cb) { float result = 0.0f; + if (ex_bits == 8) { + // RBQ9 is byte-aligned, including the last dimension. Do not require + // trailing factor bytes for the scalar reference or SIMD tail. + for (size_t i = start; i < d; ++i) { + const int sb = (sign_bits[i / 8] >> (i % 8)) & 1; + result += rotated_q[i] * + (static_cast((sb << 8) + ex_code[i]) + cb); + } + return result; + } const int sign_shift = static_cast(ex_bits); const uint64_t code_mask = (1ULL << ex_bits) - 1; for (size_t i = start; i < d; i++) { diff --git a/faiss/utils/simd_impl/rabitq_avx2.cpp b/faiss/utils/simd_impl/rabitq_avx2.cpp index 0d63504dc..c1c3ec053 100644 --- a/faiss/utils/simd_impl/rabitq_avx2.cpp +++ b/faiss/utils/simd_impl/rabitq_avx2.cpp @@ -275,9 +275,11 @@ uint64_t bitwise_and_dot_product( } sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const uint64_t yv = *(const uint64_t*)(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); for (int j = 0; j < qb; j++) { - const uint64_t qv = *(const uint64_t*)(query + j * size + offset); + uint64_t qv; + std::memcpy(&qv, query + j * size + offset, sizeof(qv)); sum += popcount64(qv & yv) << j; } } @@ -336,10 +338,12 @@ BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount< dot_product += reduce_add_128(dot_128); popcount_sum += reduce_add_128(pop_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const uint64_t yv = *(const uint64_t*)(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); popcount_sum += popcount64(yv); for (int j = 0; j < qb; j++) { - const uint64_t qv = *(const uint64_t*)(query + j * size + offset); + uint64_t qv; + std::memcpy(&qv, query + j * size + offset, sizeof(qv)); dot_product += popcount64(qv & yv) << j; } } @@ -391,9 +395,11 @@ uint64_t bitwise_xor_dot_product( } sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); for (int j = 0; j < qb; j++) { - const auto qv = *(const uint64_t*)(query + j * size + offset); + uint64_t qv; + std::memcpy(&qv, query + j * size + offset, sizeof(qv)); sum += popcount64(qv ^ yv) << j; } } @@ -427,7 +433,8 @@ uint64_t popcount(const uint8_t* data, size_t size) { } sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); sum += popcount64(yv); } for (; offset < size; ++offset) { @@ -517,7 +524,6 @@ inline float ip_1exbit_avx2( return result; } -#ifdef __BMI2__ inline float ip_bitplane_avx2( const uint8_t* __restrict sign_bits, const uint8_t* __restrict ex_code, @@ -571,7 +577,6 @@ inline float ip_bitplane_avx2( result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb); return result; } -#endif // __BMI2__ } // namespace @@ -583,15 +588,40 @@ float compute_inner_product( size_t d, size_t ex_bits, float cb) { + if (ex_bits == 8) { + // Eight byte-aligned extra codes, plus their independent sign bits. + __m256 acc = _mm256_setzero_ps(); + const __m256 weight = _mm256_set1_ps(256.f); + const __m256 offset = _mm256_set1_ps(cb); + const __m256i positions = + _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128); + size_t i = 0; + for (; i + 8 <= d; i += 8) { + const __m128i bytes = _mm_loadl_epi64( + reinterpret_cast(ex_code + i)); + const __m256 extra = + _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(bytes)); + const __m256i mask = _mm256_cmpgt_epi32( + _mm256_and_si256( + _mm256_set1_epi32(sign_bits[i / 8]), positions), + _mm256_setzero_si256()); + const __m256 recon = _mm256_add_ps( + extra, _mm256_and_ps(_mm256_castsi256_ps(mask), weight)); + acc = _mm256_fmadd_ps( + _mm256_loadu_ps(rotated_q + i), + _mm256_add_ps(recon, offset), + acc); + } + return hsum_avx2(acc) + + ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb); + } if (ex_bits == 1) { return ip_1exbit_avx2(sign_bits, ex_code, rotated_q, d, cb); } -#ifdef __BMI2__ if (ex_bits <= 7) { return ip_bitplane_avx2(sign_bits, ex_code, rotated_q, d, ex_bits, cb); } -#endif return ip_scalar(sign_bits, ex_code, rotated_q, 0, d, ex_bits, cb); } diff --git a/faiss/utils/simd_impl/rabitq_avx512.cpp b/faiss/utils/simd_impl/rabitq_avx512.cpp index d8db42e48..1e78f86ff 100644 --- a/faiss/utils/simd_impl/rabitq_avx512.cpp +++ b/faiss/utils/simd_impl/rabitq_avx512.cpp @@ -378,9 +378,11 @@ uint64_t bitwise_and_dot_product( } sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); for (int j = 0; j < qb; j++) { - const auto qv = *(const uint64_t*)(query + j * size + offset); + uint64_t qv; + std::memcpy(&qv, query + j * size + offset, sizeof(qv)); sum += popcount64(qv & yv) << j; } } @@ -457,10 +459,12 @@ BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount< dot_product += reduce_add_128(dot_128); popcount_sum += reduce_add_128(pop_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); popcount_sum += popcount64(yv); for (int j = 0; j < qb; j++) { - const auto qv = *(const uint64_t*)(query + j * size + offset); + uint64_t qv; + std::memcpy(&qv, query + j * size + offset, sizeof(qv)); dot_product += popcount64(qv & yv) << j; } } @@ -527,9 +531,11 @@ uint64_t bitwise_xor_dot_product( } sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); for (int j = 0; j < qb; j++) { - const auto qv = *(const uint64_t*)(query + j * size + offset); + uint64_t qv; + std::memcpy(&qv, query + j * size + offset, sizeof(qv)); sum += popcount64(qv ^ yv) << j; } } @@ -572,7 +578,8 @@ uint64_t popcount(const uint8_t* data, size_t size) { } sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); sum += popcount64(yv); } for (; offset < size; ++offset) { @@ -654,13 +661,9 @@ inline float ip_1exbit_avx512( return result; } -// Needs BMI2 for _pext_u64. Some AVX2 CPUs lack it, and FAISS_BMI2_FLAGS can -// be empty, so the dispatcher falls back to the scalar path without it. -#ifdef __BMI2__ -// Bitplane kernel for ex_bits >= 2, 16 dims per iteration. A bitplane is -// already a bitmask, so it goes into a mask register and one masked add -// applies its weight. Reads of ex_code run a few bytes past the ex-code -// section into the record's own trailing factors, so they stay in bounds. +// Bit-plane kernel for ex_bits >= 2. It scores 16 dims per pass. +// Each plane acts as a mask for one weighted add. Reads can run +// past the codes into the row tail, which stays in bounds. inline float ip_bitplane_avx512( const uint8_t* __restrict sign_bits, const uint8_t* __restrict ex_code, @@ -737,7 +740,6 @@ inline float ip_bitplane_avx512( result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb); return result; } -#endif // __BMI2__ } // namespace @@ -749,16 +751,35 @@ float compute_inner_product( size_t d, size_t ex_bits, float cb) { + if (ex_bits == 8) { + // RBQ9 has one byte per extra code: no bit-plane extraction or BMI2. + __m512 acc = _mm512_setzero_ps(); + const __m512 weight = _mm512_set1_ps(256.f); + const __m512 offset = _mm512_set1_ps(cb); + size_t i = 0; + for (; i + 16 <= d; i += 16) { + uint16_t signs; + memcpy(&signs, sign_bits + i / 8, sizeof(signs)); + const __m128i bytes = _mm_loadu_si128( + reinterpret_cast(ex_code + i)); + __m512 recon = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(bytes)); + recon = _mm512_mask_add_ps(recon, signs, recon, weight); + acc = _mm512_fmadd_ps( + _mm512_loadu_ps(rotated_q + i), + _mm512_add_ps(recon, offset), + acc); + } + return _mm512_reduce_add_ps(acc) + + ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb); + } if (ex_bits == 1) { return ip_1exbit_avx512(sign_bits, ex_code, rotated_q, d, cb); } -#ifdef __BMI2__ if (ex_bits <= 7) { return ip_bitplane_avx512( sign_bits, ex_code, rotated_q, d, ex_bits, cb); } -#endif return ip_scalar(sign_bits, ex_code, rotated_q, 0, d, ex_bits, cb); } diff --git a/faiss/utils/simd_impl/rabitq_avx512_vpopcnt.cpp b/faiss/utils/simd_impl/rabitq_avx512_vpopcnt.cpp index c52e2298e..499632b9a 100644 --- a/faiss/utils/simd_impl/rabitq_avx512_vpopcnt.cpp +++ b/faiss/utils/simd_impl/rabitq_avx512_vpopcnt.cpp @@ -63,6 +63,43 @@ uint64_t bitwise_xor_dot_product( template <> uint64_t popcount(const uint8_t* data, size_t size); +template <> +void bitwise_q4_batch_4( + const uint8_t* query, + const uint8_t* const* data, + size_t size, + BitwiseAndDotProductResult* results) { + __m512i dots[4], pops[4]; + for (int i = 0; i < 4; ++i) { + dots[i] = pops[i] = _mm512_setzero_si512(); + } + for (size_t off = 0; off < size; off += 64) { + const size_t count = std::min(size - off, size_t(64)); + const __mmask64 mask = + count == 64 ? ~__mmask64(0) : (__mmask64(1) << count) - 1; + __m512i x[4]; + for (int i = 0; i < 4; ++i) { + x[i] = _mm512_maskz_loadu_epi8(mask, data[i] + off); + pops[i] = _mm512_add_epi64(pops[i], _mm512_popcnt_epi64(x[i])); + } + for (int bit = 0; bit < 4; ++bit) { + // Load each query plane once for four independent database codes. + const __m512i q = + _mm512_maskz_loadu_epi8(mask, query + bit * size + off); + for (int i = 0; i < 4; ++i) { + const __m512i p = + _mm512_popcnt_epi64(_mm512_and_si512(q, x[i])); + dots[i] = _mm512_add_epi64(dots[i], _mm512_slli_epi64(p, bit)); + } + } + } + for (int i = 0; i < 4; ++i) { + results[i] = { + static_cast(_mm512_reduce_add_epi64(dots[i])), + static_cast(_mm512_reduce_add_epi64(pops[i]))}; + } +} + namespace { // 512-bit popcount using AVX-512 VPOPCNTDQ (vpopcntq). @@ -164,10 +201,11 @@ uint64_t bitwise_and_dot_product( // 64-bit scalar tail. for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *reinterpret_cast(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); for (size_t j = 0; j < qb; j++) { - const auto qv = *reinterpret_cast( - query + j * size + offset); + uint64_t qv; + std::memcpy(&qv, query + j * size + offset, sizeof(qv)); sum += static_cast(popcount64(qv & yv)) << j; } } @@ -255,11 +293,12 @@ BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount< popcount_sum += reduce_add_128(pop_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *reinterpret_cast(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); popcount_sum += popcount64(yv); for (size_t j = 0; j < qb; j++) { - const auto qv = *reinterpret_cast( - query + j * size + offset); + uint64_t qv; + std::memcpy(&qv, query + j * size + offset, sizeof(qv)); dot_product += static_cast(popcount64(qv & yv)) << j; } } @@ -336,10 +375,11 @@ uint64_t bitwise_xor_dot_product( sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *reinterpret_cast(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); for (size_t j = 0; j < qb; j++) { - const auto qv = *reinterpret_cast( - query + j * size + offset); + uint64_t qv; + std::memcpy(&qv, query + j * size + offset, sizeof(qv)); sum += static_cast(popcount64(qv ^ yv)) << j; } } @@ -389,7 +429,8 @@ uint64_t popcount(const uint8_t* data, size_t size) { sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *reinterpret_cast(data + offset); + uint64_t yv; + std::memcpy(&yv, data + offset, sizeof(yv)); sum += popcount64(yv); } for (; offset < size; ++offset) { diff --git a/faiss/utils/simd_levels.cpp b/faiss/utils/simd_levels.cpp index 5784275e6..5ef8973f9 100644 --- a/faiss/utils/simd_levels.cpp +++ b/faiss/utils/simd_levels.cpp @@ -274,7 +274,9 @@ SIMDLevel SIMDConfig::auto_detect_simd_level() { #if defined(COMPILE_SIMD_AVX2) || defined(COMPILE_SIMD_AVX512) bool has_avx2 = (ebx7 & (1 << 5)) != 0; - if (has_avx2) { + // The AVX2 and higher SIMD files compile with -mbmi2. + bool has_bmi2 = (ebx7 & (1 << 8)) != 0; + if (has_avx2 && has_bmi2) { supported_simd_levels |= (1 << static_cast(SIMDLevel::AVX2)); detected_level = SIMDLevel::AVX2; } @@ -282,7 +284,7 @@ SIMDLevel SIMDConfig::auto_detect_simd_level() { #if defined(COMPILE_SIMD_AVX512) bool cpu_has_avx512f = (ebx7 & (1 << 16)) != 0; bool os_supports_avx512 = (xcr0 & 0xE0) == 0xE0; - bool has_avx512f = cpu_has_avx512f && os_supports_avx512; + bool has_avx512f = cpu_has_avx512f && os_supports_avx512 && has_bmi2; if (has_avx512f) { bool has_avx512cd = (ebx7 & (1 << 28)) != 0; bool has_avx512vl = (ebx7 & (1 << 31)) != 0; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 103c8907e..64866262e 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -45,6 +45,7 @@ set(FAISS_TEST_SRC 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 ) @@ -69,6 +70,10 @@ if(FAISS_OPT_LEVEL STREQUAL "dd") 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 diff --git a/tests/test_rabitq_batch_estimates.cpp b/tests/test_rabitq_batch_estimates.cpp new file mode 100644 index 000000000..3c52d10cd --- /dev/null +++ b/tests/test_rabitq_batch_estimates.cpp @@ -0,0 +1,278 @@ +/* + * 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. + */ + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace { + +struct RestoreSIMD { + faiss::SIMDLevel saved = faiss::SIMDConfig::get_level(); + ~RestoreSIMD() { + faiss::SIMDConfig::set_level(saved); + } +}; + +std::vector available_levels() { + std::vector levels; + for (auto level : + {faiss::SIMDConfig::get_level(), + faiss::SIMDLevel::NONE, + faiss::SIMDLevel::AVX2, + faiss::SIMDLevel::AVX512, + faiss::SIMDLevel::AVX512_VPOPCNT}) { + if (faiss::SIMDConfig::is_simd_level_available(level) && + std::find(levels.begin(), levels.end(), level) == levels.end()) { + levels.push_back(level); + } + } + return levels; +} + +} // namespace + +TEST(RaBitQBatchEstimates, MatchesSingleAcrossPrecisionsMetricsAndSIMD) { + RestoreSIMD restore; + std::mt19937 rng(12345); + std::normal_distribution normal; + for (int d : {7, 64, 129, 769}) { + std::vector data(16 * d), queries(2 * d); + for (auto& v : data) { + v = normal(rng); + } + for (auto& v : queries) { + v = normal(rng); + } + for (auto metric : {faiss::METRIC_L2, faiss::METRIC_INNER_PRODUCT}) { + for (int bits = 1; bits <= 9; ++bits) { + faiss::IndexRaBitQ index(d, metric, bits); + index.train(16, data.data()); + index.add(16, data.data()); + const uint8_t* codes[] = { + index.codes.data() + 9 * index.code_size, + index.codes.data(), + index.codes.data() + 15 * index.code_size, + index.codes.data()}; + for (auto level : available_levels()) { + faiss::SIMDConfig::set_level(level); + for (int qb = 0; qb <= 8; ++qb) { + for (bool centered : {false, true}) { + SCOPED_TRACE( + ::testing::Message() + << d << '/' << bits << '/' << qb << '/' + << centered << '/' << int(level)); + std::unique_ptr + owner(index.get_quantized_distance_computer( + qb, centered)); + auto* dc = dynamic_cast< + faiss::RaBitQDistanceComputer*>( + owner.get()); + ASSERT_NE(dc, nullptr); + for (int q = 0; q < 3; ++q) { + // Include the degenerate zero-residual query. + dc->set_query( + q == 2 ? index.center.data() + : queries.data() + q * d); + float distances[4]; + dc->distance_to_code_1bit_batch_4( + codes, distances); + for (int i = 0; i < 4; ++i) { + ASSERT_TRUE(std::isfinite(distances[i])); + EXPECT_FLOAT_EQ( + distances[i], + dc->distance_to_code_1bit( + codes[i])); + } + EXPECT_EQ(dc->stats.n_1bit, 0); + EXPECT_EQ(dc->stats.n_refine, 0); + } + } + } + } + } + } + } +} + +TEST(RaBitQBatchEstimates, SerializationPreservesCodesAndEstimates) { + constexpr int d = 65; + std::vector data(8 * d); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(float(i)); + } + for (auto metric : {faiss::METRIC_L2, faiss::METRIC_INNER_PRODUCT}) { + for (int bits : {1, 4, 8, 9}) { + faiss::IndexRaBitQ index(d, metric, bits); + index.train(8, data.data()); + index.add(8, data.data()); + faiss::VectorIOWriter writer; + faiss::write_index(&index, &writer); + faiss::VectorIOReader reader; + reader.data = writer.data; + std::unique_ptr loaded(faiss::read_index(&reader)); + auto* copy = dynamic_cast(loaded.get()); + ASSERT_NE(copy, nullptr); + ASSERT_EQ(index.codes, copy->codes); + std::unique_ptr a( + index.get_FlatCodesDistanceComputer()); + std::unique_ptr b( + copy->get_FlatCodesDistanceComputer()); + a->set_query(data.data()); + b->set_query(data.data()); + for (int i = 0; i < 8; ++i) { + EXPECT_FLOAT_EQ((*a)(i), (*b)(i)); + } + } + } +} + +TEST(RaBitQBatchEstimates, RBQ9ByteCodesMatchIndependentReference) { + RestoreSIMD restore; + std::mt19937 rng(2718); + for (size_t d : {1, 7, 8, 9, 15, 16, 17, 63, 64, 65, 128, 257, 769}) { + // Exactly d extra bytes after an intentionally unaligned start. + std::vector signs(1 + (d + 7) / 8), extra(1 + d); + std::vector query(d); + for (auto& v : signs) { + v = rng(); + } + for (auto& v : extra) { + v = rng(); + } + for (auto& v : query) { + v = (int(rng() % 2001) - 1000) / 1000.f; + } + double expected = 0, magnitude = 0; + for (size_t i = 0; i < d; ++i) { + const int sign = (signs[1 + i / 8] >> (i % 8)) & 1; + const double term = query[i] * (extra[1 + i] + 256 * sign - 255.5); + expected += term; + magnitude += std::abs(term); + } + for (auto level : available_levels()) { + faiss::SIMDConfig::set_level(level); + const float actual = faiss::with_selected_simd_levels< + faiss::AVAILABLE_SIMD_LEVELS_BASE_WITH_VPOPCNT>( + [&]() { + return faiss::rabitq::multibit::compute_inner_product< + SL>( + signs.data() + 1, + extra.data() + 1, + query.data(), + d, + 8, + -255.5f); + }); + EXPECT_NEAR(actual, expected, 2e-6 * magnitude + 1e-5); + } + } +} + +TEST(RaBitQBatchEstimates, Q4KernelUnalignedTailsMatchByteReference) { + RestoreSIMD restore; + std::mt19937 rng(42); + for (size_t size : + {0, 1, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 65, 96, 97, 193}) { + std::vector query(1 + 4 * size); + std::vector storage[4]; + for (auto& v : query) { + v = rng(); + } + const uint8_t* codes[4]; + for (int i = 0; i < 4; ++i) { + storage[i].resize(i + 1 + size); + for (auto& v : storage[i]) { + v = rng(); + } + codes[i] = storage[i].data() + i + 1; + } + for (auto level : available_levels()) { + faiss::SIMDConfig::set_level(level); + faiss::rabitq::BitwiseAndDotProductResult results[4]; + faiss::with_selected_simd_levels< + faiss::AVAILABLE_SIMD_LEVELS_BASE_WITH_VPOPCNT>( + [&]() { + faiss::rabitq::bitwise_q4_batch_4( + query.data() + 1, codes, size, results); + }); + for (int i = 0; i < 4; ++i) { + uint64_t dot = 0, pop = 0; + for (size_t j = 0; j < size; ++j) { + for (int bit = 0; bit < 8; ++bit) { + const int on = (codes[i][j] >> bit) & 1; + pop += on; + for (int plane = 0; plane < 4; ++plane) { + dot += (on & (query[1 + plane * size + j] >> bit)) * + (1 << plane); + } + } + } + EXPECT_EQ(results[i].dot_product, dot); + EXPECT_EQ(results[i].popcount, pop); + } + } + } +} + +TEST(RaBitQBatchEstimates, MultiBitKernelsMatchIndependentReference) { + RestoreSIMD restore; + std::mt19937 rng(31415); + for (size_t d : {1, 7, 8, 9, 15, 16, 17, 65, 129}) { + for (size_t ex_bits = 1; ex_bits <= 8; ++ex_bits) { + // Existing sub-byte scorers may read a uint64 window into the + // trailing factor area; reserve those bytes as in an encoded row. + std::vector signs(1 + (d + 7) / 8); + std::vector extra(1 + (d * ex_bits + 7) / 8 + 8); + std::vector query(d); + for (auto& v : signs) { + v = rng(); + } + double expected = 0, magnitude = 0; + const float cb = -(float(1u << ex_bits) - 0.5f); + for (size_t i = 0; i < d; ++i) { + query[i] = (int(rng() % 2001) - 1000) / 1000.f; + const unsigned value = rng() % (1u << ex_bits); + for (size_t bit = 0; bit < ex_bits; ++bit) { + const size_t position = i * ex_bits + bit; + extra[1 + position / 8] |= ((value >> bit) & 1) + << (position % 8); + } + const int sign = (signs[1 + i / 8] >> (i % 8)) & 1; + const double term = + query[i] * (double(value) + (sign << ex_bits) + cb); + expected += term; + magnitude += std::abs(term); + } + for (auto level : available_levels()) { + faiss::SIMDConfig::set_level(level); + const float actual = faiss::with_selected_simd_levels< + faiss::AVAILABLE_SIMD_LEVELS_BASE_WITH_VPOPCNT>( + [&]() { + return faiss::rabitq::multibit:: + compute_inner_product( + signs.data() + 1, + extra.data() + 1, + query.data(), + d, + ex_bits, + cb); + }); + EXPECT_NEAR(actual, expected, 2e-6 * magnitude + 1e-5); + } + } + } +}