From d1c432fb5fd7d1bf96a73299e95fa6476b2825e3 Mon Sep 17 00:00:00 2001 From: Gergely Szilvasy Date: Tue, 31 Mar 2026 14:41:16 -0700 Subject: [PATCH] Convert rabitq_simd.h to runtime SIMD dispatch (#4912) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4912 Convert RaBitQ SIMD primitives from compile-time #ifdef cascades to the DD template specialization pattern. - Replace inline SIMD functions in rabitq_simd.h with template declarations and NONE (scalar) specializations - Create rabitq_avx2.cpp with AVX2+SSE4.1 specializations for bitwise_and_dot_product, bitwise_xor_dot_product, popcount, and compute_inner_product (including BMI2 bitplane kernel) - Create rabitq_avx512.cpp with AVX512 specializations (falls back to AVX2 bitplane kernel for compute_inner_product with ex_bits >= 2) - Add AVX512_SPR forwarding to AVX512 for all four functions - Create rabitq_neon.cpp forwarding ARM_NEON/ARM_SVE to NONE - Wrap calls in RaBitQuantizer.cpp and RaBitQUtils.cpp with with_simd_level dispatch - Add -mbmi2 to AVX2 and AVX512 SIMD flags (required by _pext_u64 in bitplane kernel; BMI2 is present on all Haswell+ CPUs) - Register new SIMD files in xplat.bzl and CMakeLists.txt Reviewed By: mdouze Differential Revision: D96298743 fbshipit-source-id: 1941bf099664ba2d3d1c1245d5c06d1775614955 --- faiss/CMakeLists.txt | 3 + faiss/impl/RaBitQUtils.cpp | 28 +- faiss/impl/RaBitQUtils.h | 30 ++ faiss/impl/RaBitQuantizer.cpp | 600 +++++++++++------------ faiss/utils/rabitq_simd.h | 606 +++--------------------- faiss/utils/simd_impl/rabitq_avx2.cpp | 355 ++++++++++++++ faiss/utils/simd_impl/rabitq_avx512.cpp | 477 +++++++++++++++++++ faiss/utils/simd_impl/rabitq_neon.cpp | 55 +++ 8 files changed, 1304 insertions(+), 850 deletions(-) create mode 100644 faiss/utils/simd_impl/rabitq_avx2.cpp create mode 100644 faiss/utils/simd_impl/rabitq_avx512.cpp create mode 100644 faiss/utils/simd_impl/rabitq_neon.cpp diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index 5ed365e6b..0469df3ba 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -15,6 +15,7 @@ set(FAISS_SIMD_AVX2_SRC impl/approx_topk/avx2.cpp utils/simd_impl/distances_avx2.cpp utils/distances_fused/simdlib_based.cpp + utils/simd_impl/rabitq_avx2.cpp ) set(FAISS_SIMD_AVX512_SRC impl/fast_scan/impl-avx512.cpp @@ -22,6 +23,7 @@ set(FAISS_SIMD_AVX512_SRC impl/scalar_quantizer/sq-avx512.cpp utils/simd_impl/distances_avx512.cpp utils/distances_fused/avx512.cpp + utils/simd_impl/rabitq_avx512.cpp ) set(FAISS_SIMD_NEON_SRC impl/fast_scan/impl-neon.cpp @@ -29,6 +31,7 @@ set(FAISS_SIMD_NEON_SRC impl/approx_topk/neon.cpp utils/simd_impl/distances_aarch64.cpp utils/distances_fused/simdlib_based_neon.cpp + utils/simd_impl/rabitq_neon.cpp ) set(FAISS_SIMD_SVE_SRC impl/pq_code_distance/pq_code_distance-sve.cpp diff --git a/faiss/impl/RaBitQUtils.cpp b/faiss/impl/RaBitQUtils.cpp index 5cb3ec21c..cde9fc84f 100644 --- a/faiss/impl/RaBitQUtils.cpp +++ b/faiss/impl/RaBitQUtils.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include #include @@ -306,6 +307,9 @@ size_t compute_per_vector_storage_size(size_t nb_bits, size_t d) { } } +// Non-template wrapper with dynamic dispatch (one dispatch per call). +// The hot path in RaBitQuantizer dispatches once at distance computer +// construction, so per-vector dispatch only affects this utility path. float compute_full_multibit_distance( const uint8_t* sign_bits, const uint8_t* ex_code, @@ -315,18 +319,18 @@ float compute_full_multibit_distance( size_t d, size_t ex_bits, MetricType metric_type) { - const float cb = -(static_cast(1 << ex_bits) - 0.5f); - - float ex_ip = rabitq::multibit::compute_inner_product( - sign_bits, ex_code, rotated_q, d, ex_bits, cb); - - float dist = qr_base + ex_fac.f_add_ex + ex_fac.f_rescale_ex * ex_ip; - - if (metric_type == MetricType::METRIC_L2) { - dist = std::max(0.0f, dist); - } - - return dist; + return with_selected_simd_levels( + [&]() { + return compute_full_multibit_distance( + sign_bits, + ex_code, + ex_fac, + rotated_q, + qr_base, + d, + ex_bits, + metric_type); + }); } void populate_block_aux_from_flat_storage( diff --git a/faiss/impl/RaBitQUtils.h b/faiss/impl/RaBitQUtils.h index 9030fba11..1eb2ebf04 100644 --- a/faiss/impl/RaBitQUtils.h +++ b/faiss/impl/RaBitQUtils.h @@ -10,6 +10,9 @@ #include #include #include +#include +#include +#include #include #include #include @@ -337,6 +340,33 @@ float compute_full_multibit_distance( size_t ex_bits, MetricType metric_type); +// SIMDLevel-templatized version — avoids per-call dynamic dispatch. +// Inline so it can be used from templatized distance computers without +// needing explicit instantiations in per-SIMD TUs. +template +inline float compute_full_multibit_distance( + const uint8_t* sign_bits, + const uint8_t* ex_code, + const ExtraBitsFactors& ex_fac, + const float* rotated_q, + float qr_base, + size_t d, + size_t ex_bits, + MetricType metric_type) { + const float cb = -(static_cast(1 << ex_bits) - 0.5f); + + float ex_ip = rabitq::multibit::compute_inner_product( + sign_bits, ex_code, rotated_q, d, ex_bits, cb); + + float dist = qr_base + ex_fac.f_add_ex + ex_fac.f_rescale_ex * ex_ip; + + if (metric_type == MetricType::METRIC_L2) { + dist = std::max(0.0f, dist); + } + + return dist; +} + /** Compute pointer to a vector's auxiliary data within block layout. */ template inline T* get_block_aux_ptr( diff --git a/faiss/impl/RaBitQuantizer.cpp b/faiss/impl/RaBitQuantizer.cpp index f169ebe01..e694ca9b6 100644 --- a/faiss/impl/RaBitQuantizer.cpp +++ b/faiss/impl/RaBitQuantizer.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include @@ -221,162 +222,159 @@ void RaBitQuantizer::decode_core( namespace { +// Distance computers templatized on SIMDLevel to avoid per-call dynamic +// dispatch. The SIMDLevel is baked in at construction time via +// get_distance_computer, so virtual calls through the base class go +// directly to the SIMD-specialized code. + +template struct RaBitQDistanceComputerNotQ : RaBitQDistanceComputer { // the rotated query (qr - c) std::vector rotated_q; // some additional numbers for the query QueryFactorsData query_fac; - RaBitQDistanceComputerNotQ(); + RaBitQDistanceComputerNotQ() = default; // Compute distance using only 1-bit codes (fast) - float distance_to_code_1bit(const uint8_t* code) override; + float distance_to_code_1bit(const uint8_t* code) override { + FAISS_ASSERT(code != nullptr); + FAISS_ASSERT( + (metric_type == MetricType::METRIC_L2 || + metric_type == MetricType::METRIC_INNER_PRODUCT)); + FAISS_ASSERT(rotated_q.size() == d); + + // split the code into parts + const uint8_t* binary_data = code; + + // Cast to appropriate type based on nb_bits + // For 1-bit: use SignBitFactors (8 bytes) + // For multi-bit: use SignBitFactorsWithError (12 bytes) which includes + // f_error + size_t ex_bits = nb_bits - 1; + const SignBitFactors* base_fac = (ex_bits == 0) + ? reinterpret_cast(code + (d + 7) / 8) + : reinterpret_cast( + code + (d + 7) / 8); + + // this is the baseline code + // + // compute using floats + float dot_qo = 0; + // It was a willful decision (after the discussion) to not to pre-cache + // the sum of all bits, just in order to reduce the overhead per + // vector. + uint64_t sum_q = 0; + + for (size_t i = 0; i < d; i++) { + // Extract i-th bit + bool bit = rabitq_utils::extract_bit_standard(binary_data, i); + // accumulate dp + dot_qo += bit ? rotated_q[i] : 0; + // accumulate sum-of-bits + sum_q += bit ? 1 : 0; + } + + // Apply query factors + float final_dot = + query_fac.c1 * dot_qo + query_fac.c2 * sum_q - query_fac.c34; + + // pre_dist = ||or - c||^2 + ||qr - c||^2 - + // 2 * ||or - c|| * ||qr - c|| * - (IP ? ||or||^2 : 0) + float pre_dist = base_fac->or_minus_c_l2sqr + query_fac.qr_to_c_L2sqr - + 2 * base_fac->dp_multiplier * final_dot; + + if (metric_type == MetricType::METRIC_L2) { + // ||or - q||^ 2 + return pre_dist; + } else { + // metric == MetricType::METRIC_INNER_PRODUCT + // 2 * (or, q) = (||or - q||^2 - ||q||^2 - ||or||^2) + return -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) override; + float distance_to_code_full(const uint8_t* code) override { + FAISS_ASSERT(code != nullptr); + FAISS_ASSERT( + (metric_type == MetricType::METRIC_L2 || + metric_type == MetricType::METRIC_INNER_PRODUCT)); + FAISS_ASSERT(rotated_q.size() == d); - void set_query(const float* x) override; + size_t ex_bits = nb_bits - 1; + + if (ex_bits == 0) { + // No ex-bits, just return 1-bit distance + return distance_to_code_1bit(code); + } + + // Extract pointers to code sections + const uint8_t* binary_data = code; + size_t offset = (d + 7) / 8 + sizeof(SignBitFactorsWithError); + const uint8_t* ex_code = code + offset; + const ExtraBitsFactors* ex_fac = + reinterpret_cast( + ex_code + (d * ex_bits + 7) / 8); + + float qr_base = (metric_type == MetricType::METRIC_INNER_PRODUCT) + ? query_fac.q_dot_c + : query_fac.qr_to_c_L2sqr; + return rabitq_utils::compute_full_multibit_distance( + binary_data, + ex_code, + *ex_fac, + rotated_q.data(), + qr_base, + d, + ex_bits, + metric_type); + } + + void set_query(const float* x) override { + q = x; + FAISS_ASSERT(x != nullptr); + FAISS_ASSERT( + (metric_type == MetricType::METRIC_L2 || + metric_type == MetricType::METRIC_INNER_PRODUCT)); + + // compute the distance from the query to the centroid + if (centroid != nullptr) { + query_fac.qr_to_c_L2sqr = fvec_L2sqr(x, centroid, d); + } else { + query_fac.qr_to_c_L2sqr = fvec_norm_L2sqr(x, d); + } + + // subtract c, obtain P^(-1)(qr - c) + rotated_q.resize(d); + for (size_t i = 0; i < d; i++) { + rotated_q[i] = x[i] - ((centroid == nullptr) ? 0 : centroid[i]); + } + + // Compute g_error = ||qr - c|| (L2 norm of rotated query) + g_error = std::sqrt(query_fac.qr_to_c_L2sqr); + + // compute some numbers — do not quantize the query + const float inv_d = (d == 0) ? 1.0f : (1.0f / std::sqrt((float)d)); + + float sum_q = 0; + for (size_t i = 0; i < d; i++) { + sum_q += rotated_q[i]; + } + + query_fac.c1 = 2 * inv_d; + query_fac.c2 = 0; + query_fac.c34 = sum_q * inv_d; + + if (metric_type == MetricType::METRIC_INNER_PRODUCT) { + query_fac.qr_norm_L2sqr = fvec_norm_L2sqr(x, d); + query_fac.q_dot_c = + centroid ? fvec_inner_product(x, centroid, d) : 0.0f; + } + } }; -RaBitQDistanceComputerNotQ::RaBitQDistanceComputerNotQ() = default; - -float RaBitQDistanceComputerNotQ::distance_to_code_1bit(const uint8_t* code) { - FAISS_ASSERT(code != nullptr); - FAISS_ASSERT( - (metric_type == MetricType::METRIC_L2 || - metric_type == MetricType::METRIC_INNER_PRODUCT)); - FAISS_ASSERT(rotated_q.size() == d); - - // split the code into parts - const uint8_t* binary_data = code; - - // Cast to appropriate type based on nb_bits - // For 1-bit: use SignBitFactors (8 bytes) - // For multi-bit: use SignBitFactorsWithError (12 bytes) which includes - // f_error - size_t ex_bits = nb_bits - 1; - const SignBitFactors* base_fac = (ex_bits == 0) - ? reinterpret_cast(code + (d + 7) / 8) - : reinterpret_cast( - code + (d + 7) / 8); - - // this is the baseline code - // - // compute using floats - float dot_qo = 0; - // It was a willful decision (after the discussion) to not to pre-cache - // the sum of all bits, just in order to reduce the overhead per vector. - uint64_t sum_q = 0; - - for (size_t i = 0; i < d; i++) { - // Extract i-th bit - bool bit = rabitq_utils::extract_bit_standard(binary_data, i); - // accumulate dp - dot_qo += bit ? rotated_q[i] : 0; - // accumulate sum-of-bits - sum_q += bit ? 1 : 0; - } - - // Apply query factors - float final_dot = - query_fac.c1 * dot_qo + query_fac.c2 * sum_q - query_fac.c34; - - // pre_dist = ||or - c||^2 + ||qr - c||^2 - - // 2 * ||or - c|| * ||qr - c|| * - (IP ? ||or||^2 : 0) - float pre_dist = base_fac->or_minus_c_l2sqr + query_fac.qr_to_c_L2sqr - - 2 * base_fac->dp_multiplier * final_dot; - - if (metric_type == MetricType::METRIC_L2) { - // ||or - q||^ 2 - return pre_dist; - } else { - // metric == MetricType::METRIC_INNER_PRODUCT - return -0.5f * (pre_dist - query_fac.qr_norm_L2sqr); - } -} - -float RaBitQDistanceComputerNotQ::distance_to_code_full(const uint8_t* code) { - FAISS_ASSERT(code != nullptr); - FAISS_ASSERT( - (metric_type == MetricType::METRIC_L2 || - metric_type == MetricType::METRIC_INNER_PRODUCT)); - FAISS_ASSERT(rotated_q.size() == d); - - size_t ex_bits = nb_bits - 1; - - if (ex_bits == 0) { - // No ex-bits, just return 1-bit distance - return distance_to_code_1bit(code); - } - - // Extract pointers to code sections - const uint8_t* binary_data = code; - size_t offset = (d + 7) / 8 + sizeof(SignBitFactorsWithError); - const uint8_t* ex_code = code + offset; - const ExtraBitsFactors* ex_fac = reinterpret_cast( - ex_code + (d * ex_bits + 7) / 8); - - // Call shared utility directly with rotated_q pointer - float qr_base = (metric_type == MetricType::METRIC_INNER_PRODUCT) - ? query_fac.q_dot_c - : query_fac.qr_to_c_L2sqr; - return rabitq_utils::compute_full_multibit_distance( - binary_data, - ex_code, - *ex_fac, - rotated_q.data(), - qr_base, - d, - ex_bits, - metric_type); -} - -void RaBitQDistanceComputerNotQ::set_query(const float* x) { - q = x; - FAISS_ASSERT(x != nullptr); - FAISS_ASSERT( - (metric_type == MetricType::METRIC_L2 || - metric_type == MetricType::METRIC_INNER_PRODUCT)); - - // compute the distance from the query to the centroid - if (centroid != nullptr) { - query_fac.qr_to_c_L2sqr = fvec_L2sqr(x, centroid, d); - } else { - query_fac.qr_to_c_L2sqr = fvec_norm_L2sqr(x, d); - } - - // subtract c, obtain P^(-1)(qr - c) - rotated_q.resize(d); - for (size_t i = 0; i < d; i++) { - rotated_q[i] = x[i] - ((centroid == nullptr) ? 0 : centroid[i]); - } - - // Compute g_error (query norm for lower bound computation) - // g_error = ||qr - c|| (L2 norm of rotated query) - g_error = std::sqrt(query_fac.qr_to_c_L2sqr); - - // compute some numbers - const float inv_d = (d == 0) ? 1.0f : (1.0f / std::sqrt((float)d)); - - // do not quantize the query - float sum_q = 0; - for (size_t i = 0; i < d; i++) { - sum_q += rotated_q[i]; - } - - query_fac.c1 = 2 * inv_d; - query_fac.c2 = 0; - query_fac.c34 = sum_q * inv_d; - - if (metric_type == MetricType::METRIC_INNER_PRODUCT) { - // precompute if needed - query_fac.qr_norm_L2sqr = fvec_norm_L2sqr(x, d); - query_fac.q_dot_c = - centroid ? fvec_inner_product(x, centroid, d) : 0.0f; - } -} - -// +template struct RaBitQDistanceComputerQ : RaBitQDistanceComputer { // the rotated and quantized query (qr - c) std::vector rotated_q; @@ -394,176 +392,188 @@ struct RaBitQDistanceComputerQ : RaBitQDistanceComputer { // the smallest value divisible by 8 that is not smaller than dim size_t popcount_aligned_dim = 0; - RaBitQDistanceComputerQ(); + RaBitQDistanceComputerQ() = default; // Compute distance using only 1-bit codes (fast) - float distance_to_code_1bit(const uint8_t* code) override; + float distance_to_code_1bit(const uint8_t* code) override { + FAISS_ASSERT(code != nullptr); + FAISS_ASSERT( + (metric_type == MetricType::METRIC_L2 || + metric_type == MetricType::METRIC_INNER_PRODUCT)); + + // split the code into parts + size_t size = (d + 7) / 8; + const uint8_t* binary_data = code; + + // Cast to appropriate type based on nb_bits + // For 1-bit: use SignBitFactors (8 bytes) + // For multi-bit: use SignBitFactorsWithError (12 bytes) which + // includes f_error + size_t ex_bits = nb_bits - 1; + const SignBitFactors* base_fac = (ex_bits == 0) + ? reinterpret_cast(code + size) + : reinterpret_cast(code + size); + + // this is ||or - c||^2 - (IP ? ||or||^2 : 0) + float final_dot = 0; + if (centered) { + int64_t int_dot = ((1 << qb) - 1) * d; + // See RaBitDistanceComputerNotQ::distance_to_code() for + // baseline code. + int_dot -= 2 * + rabitq::bitwise_xor_dot_product( + rearranged_rotated_qq.data(), + binary_data, + size, + qb); + final_dot += int_dot * query_fac.int_dot_scale; + } else { + auto dot_qo = rabitq::bitwise_and_dot_product( + rearranged_rotated_qq.data(), binary_data, size, qb); + // It was a willful decision (after the discussion) to not to + // pre-cache the sum of all bits, just in order to reduce the + // overhead per vector. + // process 64-bit popcounts + auto sum_q = rabitq::popcount(binary_data, size); + // dot-product itself + final_dot += query_fac.c1 * dot_qo; + // normalizer coefficients + final_dot += query_fac.c2 * sum_q; + // normalizer coefficients + final_dot -= query_fac.c34; + } + + const float pre_dist = base_fac->or_minus_c_l2sqr + + query_fac.qr_to_c_L2sqr - + 2 * base_fac->dp_multiplier * final_dot; + + if (metric_type == MetricType::METRIC_L2) { + // ||or - q||^ 2 + return pre_dist; + } else { + // metric == MetricType::METRIC_INNER_PRODUCT + // 2 * (or, q) = (||or - q||^2 - ||q||^2 - ||or||^2) + return -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) override; + float distance_to_code_full(const uint8_t* code) override { + FAISS_ASSERT(code != nullptr); + FAISS_ASSERT( + (metric_type == MetricType::METRIC_L2 || + metric_type == MetricType::METRIC_INNER_PRODUCT)); + FAISS_ASSERT(rotated_q.size() == d); - void set_query(const float* x) override; + size_t ex_bits = nb_bits - 1; + + if (ex_bits == 0) { + // No ex-bits, just return 1-bit distance + return distance_to_code_1bit(code); + } + + // Extract pointers to code sections + const uint8_t* binary_data = code; + size_t offset = (d + 7) / 8 + sizeof(SignBitFactorsWithError); + const uint8_t* ex_code = code + offset; + const ExtraBitsFactors* ex_fac = + reinterpret_cast( + ex_code + (d * ex_bits + 7) / 8); + + float qr_base = (metric_type == MetricType::METRIC_INNER_PRODUCT) + ? query_fac.q_dot_c + : query_fac.qr_to_c_L2sqr; + return rabitq_utils::compute_full_multibit_distance( + binary_data, + ex_code, + *ex_fac, + rotated_q.data(), + qr_base, + d, + ex_bits, + metric_type); + } + + void set_query(const float* x) override { + q = x; + FAISS_ASSERT(x != nullptr); + FAISS_ASSERT( + (metric_type == MetricType::METRIC_L2 || + metric_type == MetricType::METRIC_INNER_PRODUCT)); + FAISS_THROW_IF_NOT(qb <= 8); + FAISS_THROW_IF_NOT(qb > 0); + + // Use shared utilities for core query factor computation + // rotated_q is populated directly by compute_query_factors as an + // output parameter + query_fac = rabitq_utils::compute_query_factors( + x, + d, + centroid, + qb, + centered, + metric_type, + rotated_q, + rotated_qq); + + // Compute g_error (query norm for lower bound computation) + // g_error = ||qr - c|| (L2 norm of rotated query) + g_error = std::sqrt(query_fac.qr_to_c_L2sqr); + + // Rearrange the query vector for SIMD operations + // (RaBitQuantizer-specific) + popcount_aligned_dim = ((d + 7) / 8) * 8; + size_t offset = (d + 7) / 8; + + rearranged_rotated_qq.resize(offset * qb); + std::fill( + rearranged_rotated_qq.begin(), rearranged_rotated_qq.end(), 0); + + for (size_t idim = 0; idim < d; idim++) { + for (size_t iv = 0; iv < qb; iv++) { + const bool bit = ((rotated_qq[idim] & (1 << iv)) != 0); + rearranged_rotated_qq[iv * offset + idim / 8] |= + bit ? (1 << (idim % 8)) : 0; + } + } + } }; -RaBitQDistanceComputerQ::RaBitQDistanceComputerQ() = default; - -float RaBitQDistanceComputerQ::distance_to_code_1bit(const uint8_t* code) { - FAISS_ASSERT(code != nullptr); - FAISS_ASSERT( - (metric_type == MetricType::METRIC_L2 || - metric_type == MetricType::METRIC_INNER_PRODUCT)); - - // split the code into parts - size_t size = (d + 7) / 8; - const uint8_t* binary_data = code; - - // Cast to appropriate type based on nb_bits - // For 1-bit: use SignBitFactors (8 bytes) - // For multi-bit: use SignBitFactorsWithError (12 bytes) which includes - // f_error - size_t ex_bits = nb_bits - 1; - const SignBitFactors* base_fac = (ex_bits == 0) - ? reinterpret_cast(code + size) - : reinterpret_cast(code + size); - - // this is ||or - c||^2 - (IP ? ||or||^2 : 0) - float final_dot = 0; - if (centered) { - int64_t int_dot = ((1 << qb) - 1) * d; - // See RaBitDistanceComputerNotQ::distance_to_code() for baseline code. - int_dot -= 2 * - rabitq::bitwise_xor_dot_product( - rearranged_rotated_qq.data(), binary_data, size, qb); - final_dot += int_dot * query_fac.int_dot_scale; - } else { - auto dot_qo = rabitq::bitwise_and_dot_product( - rearranged_rotated_qq.data(), binary_data, size, qb); - // It was a willful decision (after the discussion) to not to pre-cache - // the sum of all bits, just in order to reduce the overhead per vector. - // process 64-bit popcounts - auto sum_q = rabitq::popcount(binary_data, size); - // dot-product itself - final_dot += query_fac.c1 * dot_qo; - // normalizer coefficients - final_dot += query_fac.c2 * sum_q; - // normalizer coefficients - final_dot -= query_fac.c34; - } - - // pre_dist = ||or - c||^2 + ||qr - c||^2 - - // 2 * ||or - c|| * ||qr - c|| * - (IP ? ||or||^2 : 0) - const float pre_dist = base_fac->or_minus_c_l2sqr + - query_fac.qr_to_c_L2sqr - 2 * base_fac->dp_multiplier * final_dot; - - if (metric_type == MetricType::METRIC_L2) { - // ||or - q||^ 2 - return pre_dist; - } else { - // metric == MetricType::METRIC_INNER_PRODUCT - // 2 * (or, q) = (||or - q||^2 - ||q||^2 - ||or||^2) - return -0.5f * (pre_dist - query_fac.qr_norm_L2sqr); - } -} - -float RaBitQDistanceComputerQ::distance_to_code_full(const uint8_t* code) { - FAISS_ASSERT(code != nullptr); - FAISS_ASSERT( - (metric_type == MetricType::METRIC_L2 || - metric_type == MetricType::METRIC_INNER_PRODUCT)); - FAISS_ASSERT(rotated_q.size() == d); - - size_t ex_bits = nb_bits - 1; - - if (ex_bits == 0) { - // No ex-bits, just return 1-bit distance - return distance_to_code_1bit(code); - } - - // Extract pointers to code sections - const uint8_t* binary_data = code; - size_t offset = (d + 7) / 8 + sizeof(SignBitFactorsWithError); - const uint8_t* ex_code = code + offset; - const ExtraBitsFactors* ex_fac = reinterpret_cast( - ex_code + (d * ex_bits + 7) / 8); - - // Call shared utility directly with rotated_q pointer - float qr_base = (metric_type == MetricType::METRIC_INNER_PRODUCT) - ? query_fac.q_dot_c - : query_fac.qr_to_c_L2sqr; - return rabitq_utils::compute_full_multibit_distance( - binary_data, - ex_code, - *ex_fac, - rotated_q.data(), - qr_base, - d, - ex_bits, - metric_type); -} - // Use shared constant from RaBitQUtils using rabitq_utils::Z_MAX_BY_QB; -void RaBitQDistanceComputerQ::set_query(const float* x) { - q = x; - FAISS_ASSERT(x != nullptr); - FAISS_ASSERT( - (metric_type == MetricType::METRIC_L2 || - metric_type == MetricType::METRIC_INNER_PRODUCT)); - FAISS_THROW_IF_NOT(qb <= 8); - FAISS_THROW_IF_NOT(qb > 0); - - // Use shared utilities for core query factor computation - // rotated_q is populated directly by compute_query_factors as an output - // parameter - query_fac = rabitq_utils::compute_query_factors( - x, d, centroid, qb, centered, metric_type, rotated_q, rotated_qq); - - // Compute g_error (query norm for lower bound computation) - // g_error = ||qr - c|| (L2 norm of rotated query) - g_error = std::sqrt(query_fac.qr_to_c_L2sqr); - - // Rearrange the query vector for SIMD operations (RaBitQuantizer-specific) - popcount_aligned_dim = ((d + 7) / 8) * 8; - size_t offset = (d + 7) / 8; - - rearranged_rotated_qq.resize(offset * qb); - std::fill(rearranged_rotated_qq.begin(), rearranged_rotated_qq.end(), 0); - - for (size_t idim = 0; idim < d; idim++) { - for (size_t iv = 0; iv < qb; iv++) { - const bool bit = ((rotated_qq[idim] & (1 << iv)) != 0); - rearranged_rotated_qq[iv * offset + idim / 8] |= - bit ? (1 << (idim % 8)) : 0; - } - } -} - } // anonymous namespace FlatCodesDistanceComputer* RaBitQuantizer::get_distance_computer( uint8_t qb, const float* centroid_in, bool centered) const { - if (qb == 0) { - auto dc = std::make_unique(); - dc->metric_type = metric_type; - dc->d = d; - dc->centroid = centroid_in; - dc->nb_bits = nb_bits; + // Dispatch on SIMDLevel once here so the distance computer methods + // call the SIMD-specialized rabitq functions directly (no per-call + // with_simd_level overhead). + return with_selected_simd_levels( + [&]() -> FlatCodesDistanceComputer* { + if (qb == 0) { + auto dc = + std::make_unique>(); + dc->metric_type = metric_type; + dc->d = d; + dc->centroid = centroid_in; + dc->nb_bits = nb_bits; - return dc.release(); - } else { - auto dc = std::make_unique(); - dc->metric_type = metric_type; - dc->d = d; - dc->centroid = centroid_in; - dc->qb = qb; - dc->centered = centered; - dc->nb_bits = nb_bits; + return dc.release(); + } else { + auto dc = std::make_unique>(); + dc->metric_type = metric_type; + dc->d = d; + dc->centroid = centroid_in; + dc->qb = qb; + dc->centered = centered; + dc->nb_bits = nb_bits; - return dc.release(); - } + return dc.release(); + } + }); } } // namespace faiss diff --git a/faiss/utils/rabitq_simd.h b/faiss/utils/rabitq_simd.h index 6178083a4..40e76c53d 100644 --- a/faiss/utils/rabitq_simd.h +++ b/faiss/utils/rabitq_simd.h @@ -11,282 +11,72 @@ #include #include -// Only include x86 SIMD intrinsics on x86/x86_64 architectures -#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || \ - defined(_M_IX86) -#include -#endif // defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || +#include namespace faiss::rabitq { -#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || \ - defined(_M_IX86) /** - * Returns the lookup table for AVX512 popcount operations. - * This table is used for lookup-based popcount implementation. + * Compute dot product between query and binary data using popcount on AND. * - * Source: https://github.com/WojciechMula/sse-popcount. - * - * @return Lookup table as __m512i register - */ -#if defined(__AVX512F__) -inline __m512i get_lookup_512() { - return _mm512_set_epi8( - /* f */ 4, - /* e */ 3, - /* d */ 3, - /* c */ 2, - /* b */ 3, - /* a */ 2, - /* 9 */ 2, - /* 8 */ 1, - /* 7 */ 3, - /* 6 */ 2, - /* 5 */ 2, - /* 4 */ 1, - /* 3 */ 2, - /* 2 */ 1, - /* 1 */ 1, - /* 0 */ 0, - /* f */ 4, - /* e */ 3, - /* d */ 3, - /* c */ 2, - /* b */ 3, - /* a */ 2, - /* 9 */ 2, - /* 8 */ 1, - /* 7 */ 3, - /* 6 */ 2, - /* 5 */ 2, - /* 4 */ 1, - /* 3 */ 2, - /* 2 */ 1, - /* 1 */ 1, - /* 0 */ 0, - /* f */ 4, - /* e */ 3, - /* d */ 3, - /* c */ 2, - /* b */ 3, - /* a */ 2, - /* 9 */ 2, - /* 8 */ 1, - /* 7 */ 3, - /* 6 */ 2, - /* 5 */ 2, - /* 4 */ 1, - /* 3 */ 2, - /* 2 */ 1, - /* 1 */ 1, - /* 0 */ 0, - /* f */ 4, - /* e */ 3, - /* d */ 3, - /* c */ 2, - /* b */ 3, - /* a */ 2, - /* 9 */ 2, - /* 8 */ 1, - /* 7 */ 3, - /* 6 */ 2, - /* 5 */ 2, - /* 4 */ 1, - /* 3 */ 2, - /* 2 */ 1, - /* 1 */ 1, - /* 0 */ 0); -} -#endif // defined(__AVX512F__) -#if defined(__AVX2__) -/** - * Returns the lookup table for AVX2 popcount operations. - * This table is used for lookup-based popcount implementation. - * - * @return Lookup table as __m256i register - */ -inline __m256i get_lookup_256() { - return _mm256_setr_epi8( - /* 0 */ 0, - /* 1 */ 1, - /* 2 */ 1, - /* 3 */ 2, - /* 4 */ 1, - /* 5 */ 2, - /* 6 */ 2, - /* 7 */ 3, - /* 8 */ 1, - /* 9 */ 2, - /* a */ 2, - /* b */ 3, - /* c */ 2, - /* d */ 3, - /* e */ 3, - /* f */ 4, - /* 0 */ 0, - /* 1 */ 1, - /* 2 */ 1, - /* 3 */ 2, - /* 4 */ 1, - /* 5 */ 2, - /* 6 */ 2, - /* 7 */ 3, - /* 8 */ 1, - /* 9 */ 2, - /* a */ 2, - /* b */ 3, - /* c */ 2, - /* d */ 3, - /* e */ 3, - /* f */ 4); -} -#endif // defined(__AVX2__) - -#if defined(__AVX512F__) -/** - * Popcount for a 512-bit register, using lookup tables if necessary. - * - * @param v Input vector to count bits in - * @return Vector int32_t[16] with popcount results. - */ -inline __m512i popcount_512(__m512i v) { -#if defined(__AVX512VPOPCNTDQ__) - return _mm512_popcnt_epi64(v); -#else - const __m512i lookup = get_lookup_512(); - const __m512i low_mask = _mm512_set1_epi8(0x0f); - - const __m512i lo = _mm512_and_si512(v, low_mask); - const __m512i hi = _mm512_and_si512(_mm512_srli_epi16(v, 4), low_mask); - const __m512i popcnt_lo = _mm512_shuffle_epi8(lookup, lo); - const __m512i popcnt_hi = _mm512_shuffle_epi8(lookup, hi); - const __m512i popcnt = _mm512_add_epi8(popcnt_lo, popcnt_hi); - return _mm512_sad_epu8(_mm512_setzero_si512(), popcnt); -#endif // defined(__AVX512VPOPCNTDQ__) -} -#endif // defined(__AVX512F__) - -#if defined(__AVX2__) -/** - * Popcount for a 256-bit register, using lookup tables if necessary. - * - * @param v Input vector to count bits in - * @return uint64_t[4] of popcounts for each portion of the input vector. - */ -inline __m256i popcount_256(__m256i v) { - const __m256i lookup = get_lookup_256(); - const __m256i low_mask = _mm256_set1_epi8(0x0f); - - const __m256i lo = _mm256_and_si256(v, low_mask); - const __m256i hi = _mm256_and_si256(_mm256_srli_epi16(v, 4), low_mask); - const __m256i popcnt_lo = _mm256_shuffle_epi8(lookup, lo); - const __m256i popcnt_hi = _mm256_shuffle_epi8(lookup, hi); - const __m256i popcnt = _mm256_add_epi8(popcnt_lo, popcnt_hi); - // Reduce uint8_t[32] into uint64_t[4] by addition. - return _mm256_sad_epu8(_mm256_setzero_si256(), popcnt); -} - -inline uint64_t reduce_add_256(__m256i v) { - alignas(32) uint64_t lanes[4]; - _mm256_store_si256((__m256i*)lanes, v); - return lanes[0] + lanes[1] + lanes[2] + lanes[3]; -} -#endif // defined(__AVX2__) - -#if defined(__SSE4_1__) -inline __m128i popcount_128(__m128i v) { - // Scalar popcount for each 64-bit lane - uint64_t lane0 = _mm_extract_epi64(v, 0); - uint64_t lane1 = _mm_extract_epi64(v, 1); - uint64_t pop0 = __builtin_popcountll(lane0); - uint64_t pop1 = __builtin_popcountll(lane1); - return _mm_set_epi64x(pop1, pop0); -} - -inline uint64_t reduce_add_128(__m128i v) { - alignas(16) uint64_t lanes[2]; - _mm_store_si128((__m128i*)lanes, v); - return lanes[0] + lanes[1]; -} -#endif // defined(__SSE4_1__) -#endif // defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || - -/** - * Compute dot product between query and binary data using popcount operations. - * - * @param query Pointer to rearranged rotated query data + * @param query Pointer to rearranged rotated query data * @param data Pointer to binary data - * @param d Dimension - * @param qb Number of quantization bits - * @return Unsigned integer dot product + * @param size Size in bytes + * @param qb Number of quantization bits + * @return Unsigned integer dot product */ -inline uint64_t bitwise_and_dot_product( +template +uint64_t bitwise_and_dot_product( + const uint8_t* query, + const uint8_t* data, + size_t size, + size_t qb); + +/** + * Compute dot product between query and binary data using popcount on XOR. + * + * @param query Pointer to rearranged rotated query data + * @param data Pointer to binary data + * @param size Size in bytes + * @param qb Number of quantization bits + * @return Unsigned integer dot product + */ +template +uint64_t bitwise_xor_dot_product( + const uint8_t* query, + const uint8_t* data, + size_t size, + size_t qb); + +/** + * Count total set bits in data. + * + * @param data Pointer to binary data + * @param size Size in bytes + * @return Total popcount + */ +template +uint64_t popcount(const uint8_t* data, size_t size); + +// NONE specializations — scalar fallbacks + +template <> +inline uint64_t bitwise_and_dot_product( const uint8_t* query, const uint8_t* data, size_t size, size_t qb) { uint64_t sum = 0; size_t offset = 0; -#if defined(__AVX512F__) - // Handle 512-bit chunks. - if (size_t step = 512 / 8; offset + step <= size) { - __m512i sum_512 = _mm512_setzero_si512(); - for (; offset + step <= size; offset += step) { - __m512i v_x = _mm512_loadu_si512((const __m512i*)(data + offset)); - for (size_t j = 0; j < qb; j++) { - __m512i v_q = _mm512_loadu_si512( - (const __m512i*)(query + j * size + offset)); - __m512i v_and = _mm512_and_si512(v_q, v_x); - __m512i v_popcnt = popcount_512(v_and); - __m512i v_shifted = _mm512_slli_epi64(v_popcnt, j); - sum_512 = _mm512_add_epi64(sum_512, v_shifted); - } - } - sum += _mm512_reduce_add_epi64(sum_512); - } -#endif // defined(__AVX512F__) -#if defined(__AVX2__) - if (size_t step = 256 / 8; offset + step <= size) { - __m256i sum_256 = _mm256_setzero_si256(); - for (; offset + step <= size; offset += step) { - __m256i v_x = _mm256_loadu_si256((const __m256i*)(data + offset)); - for (size_t j = 0; j < qb; j++) { - __m256i v_q = _mm256_loadu_si256( - (const __m256i*)(query + j * size + offset)); - __m256i v_and = _mm256_and_si256(v_q, v_x); - __m256i v_popcnt = popcount_256(v_and); - __m256i v_shifted = _mm256_slli_epi64(v_popcnt, j); - sum_256 = _mm256_add_epi64(sum_256, v_shifted); - } - } - sum += reduce_add_256(sum_256); - } -#endif // defined(__AVX2__) -#if defined(__SSE4_1__) - __m128i sum_128 = _mm_setzero_si128(); - for (size_t step = 128 / 8; offset + step <= size; offset += step) { - __m128i v_x = _mm_loadu_si128((const __m128i*)(data + offset)); - for (size_t j = 0; j < qb; j++) { - __m128i v_q = _mm_loadu_si128( - (const __m128i*)(query + j * size + offset)); - __m128i v_and = _mm_and_si128(v_q, v_x); - __m128i v_popcnt = popcount_128(v_and); - __m128i v_shifted = _mm_slli_epi64(v_popcnt, j); - sum_128 = _mm_add_epi64(sum_128, v_shifted); - } - } - sum += reduce_add_128(sum_128); -#endif // defined(__SSE4_1__) for (size_t step = 64 / 8; offset + step <= size; offset += step) { const auto yv = *(const uint64_t*)(data + offset); - for (size_t j = 0; j < qb; j++) { + for (int j = 0; j < qb; j++) { const auto qv = *(const uint64_t*)(query + j * size + offset); sum += __builtin_popcountll(qv & yv) << j; } } for (; offset < size; ++offset) { const auto yv = *(data + offset); - for (size_t j = 0; j < qb; j++) { + for (int j = 0; j < qb; j++) { const auto qv = *(query + j * size + offset); sum += __builtin_popcount(qv & yv) << j; } @@ -294,82 +84,24 @@ inline uint64_t bitwise_and_dot_product( return sum; } -/** - * Compute dot product between query and binary data using popcount operations. - * - * @param query Pointer to rearranged rotated query data - * @param data Pointer to binary data - * @param d Dimension - * @param qb Number of quantization bits - * @return Unsigned integer dot product - */ -inline uint64_t bitwise_xor_dot_product( +template <> +inline uint64_t bitwise_xor_dot_product( const uint8_t* query, const uint8_t* data, size_t size, size_t qb) { uint64_t sum = 0; size_t offset = 0; -#if defined(__AVX512F__) - // Handle 512-bit chunks. - if (size_t step = 512 / 8; offset + step <= size) { - __m512i sum_512 = _mm512_setzero_si512(); - for (; offset + step <= size; offset += step) { - __m512i v_x = _mm512_loadu_si512((const __m512i*)(data + offset)); - for (size_t j = 0; j < qb; j++) { - __m512i v_q = _mm512_loadu_si512( - (const __m512i*)(query + j * size + offset)); - __m512i v_xor = _mm512_xor_si512(v_q, v_x); - __m512i v_popcnt = popcount_512(v_xor); - __m512i v_shifted = _mm512_slli_epi64(v_popcnt, j); - sum_512 = _mm512_add_epi64(sum_512, v_shifted); - } - } - sum += _mm512_reduce_add_epi64(sum_512); - } -#endif -#if defined(__AVX2__) - if (size_t step = 256 / 8; offset + step <= size) { - __m256i sum_256 = _mm256_setzero_si256(); - for (; offset + step <= size; offset += step) { - __m256i v_x = _mm256_loadu_si256((const __m256i*)(data + offset)); - for (size_t j = 0; j < qb; j++) { - __m256i v_q = _mm256_loadu_si256( - (const __m256i*)(query + j * size + offset)); - __m256i v_xor = _mm256_xor_si256(v_q, v_x); - __m256i v_popcnt = popcount_256(v_xor); - __m256i v_shifted = _mm256_slli_epi64(v_popcnt, j); - sum_256 = _mm256_add_epi64(sum_256, v_shifted); - } - } - sum += reduce_add_256(sum_256); - } -#endif -#if defined(__SSE4_1__) - __m128i sum_128 = _mm_setzero_si128(); - for (size_t step = 128 / 8; offset + step <= size; offset += step) { - __m128i v_x = _mm_loadu_si128((const __m128i*)(data + offset)); - for (size_t j = 0; j < qb; j++) { - __m128i v_q = _mm_loadu_si128( - (const __m128i*)(query + j * size + offset)); - __m128i v_xor = _mm_xor_si128(v_q, v_x); - __m128i v_popcnt = popcount_128(v_xor); - __m128i v_shifted = _mm_slli_epi64(v_popcnt, j); - sum_128 = _mm_add_epi64(sum_128, v_shifted); - } - } - sum += reduce_add_128(sum_128); -#endif for (size_t step = 64 / 8; offset + step <= size; offset += step) { const auto yv = *(const uint64_t*)(data + offset); - for (size_t j = 0; j < qb; j++) { + for (int j = 0; j < qb; j++) { const auto qv = *(const uint64_t*)(query + j * size + offset); sum += __builtin_popcountll(qv ^ yv) << j; } } for (; offset < size; ++offset) { const auto yv = *(data + offset); - for (size_t j = 0; j < qb; j++) { + for (int j = 0; j < qb; j++) { const auto qv = *(query + j * size + offset); sum += __builtin_popcount(qv ^ yv) << j; } @@ -377,41 +109,10 @@ inline uint64_t bitwise_xor_dot_product( return sum; } -inline uint64_t popcount(const uint8_t* data, size_t size) { +template <> +inline uint64_t popcount(const uint8_t* data, size_t size) { uint64_t sum = 0; size_t offset = 0; -#if defined(__AVX512F__) - // Handle 512-bit chunks. - if (offset + 512 / 8 <= size) { - __m512i sum_512 = _mm512_setzero_si512(); - for (size_t end; (end = offset + 512 / 8) <= size; offset = end) { - __m512i v_x = _mm512_loadu_si512((const __m512i*)(data + offset)); - __m512i v_popcnt = popcount_512(v_x); - sum_512 = _mm512_add_epi64(sum_512, v_popcnt); - } - sum += _mm512_reduce_add_epi64(sum_512); - } -#endif // defined(__AVX512F__) -#if defined(__AVX2__) - if (offset + 256 / 8 <= size) { - __m256i sum_256 = _mm256_setzero_si256(); - for (size_t end; (end = offset + 256 / 8) <= size; offset = end) { - __m256i v_x = _mm256_loadu_si256((const __m256i*)(data + offset)); - __m256i v_popcnt = popcount_256(v_x); - sum_256 = _mm256_add_epi64(sum_256, v_popcnt); - } - sum += reduce_add_256(sum_256); - } -#endif // defined(__AVX2__) -#if defined(__SSE4_1__) - __m128i sum_128 = _mm_setzero_si128(); - for (size_t step = 128 / 8; offset + step <= size; offset += step) { - __m128i v_x = _mm_loadu_si128((const __m128i*)(data + offset)); - sum_128 = _mm_add_epi64(sum_128, popcount_128(v_x)); - } - sum += reduce_add_128(sum_128); -#endif // defined(__SSE4_1__) - for (size_t step = 64 / 8; offset + step <= size; offset += step) { const auto yv = *(const uint64_t*)(data + offset); sum += __builtin_popcountll(yv); @@ -469,186 +170,9 @@ inline float ip_scalar( return result; } -#if defined(__x86_64__) || defined(_M_X64) - -#if defined(__AVX2__) -/// Horizontal sum of 8 floats in a __m256 register. -inline float hsum_avx2(__m256 v) { - __m128 hi = _mm256_extractf128_ps(v, 1); - __m128 lo = _mm256_castps256_ps128(v); - lo = _mm_add_ps(lo, hi); - __m128 shuf = _mm_movehdup_ps(lo); - lo = _mm_add_ps(lo, shuf); - shuf = _mm_movehl_ps(shuf, lo); - return _mm_cvtss_f32(_mm_add_ss(lo, shuf)); -} -#endif // __AVX2__ - -/********************************************************* - * Specialized 1-bit kernels (ex_bits == 1). - * - * For 1 extra bit, both sign_bits and ex_code are 1-bit-per-dim packed, - * so we convert bits to floats directly — no extraction loops needed. - *********************************************************/ - -#if defined(__AVX512F__) -/// AVX-512: 16 dims/iter, ex_bits == 1. -inline float ip_1exbit_avx512( - const uint8_t* __restrict sign_bits, - const uint8_t* __restrict ex_code, - const float* __restrict rotated_q, - size_t d, - float cb) { - __m512 acc = _mm512_setzero_ps(); - const __m512 v_cb = _mm512_set1_ps(cb); - const __m512 v_two = _mm512_set1_ps(2.0f); - const __m512 v_one = _mm512_set1_ps(1.0f); - - size_t i = 0; - for (; i + 16 <= d; i += 16) { - uint16_t sb16; - memcpy(&sb16, sign_bits + i / 8, sizeof(uint16_t)); - uint16_t eb16; - memcpy(&eb16, ex_code + i / 8, sizeof(uint16_t)); - - __m512 sb_f = _mm512_maskz_mov_ps(_cvtu32_mask16(sb16), v_one); - __m512 eb_f = _mm512_maskz_mov_ps(_cvtu32_mask16(eb16), v_one); - - __m512 recon = _mm512_add_ps(_mm512_fmadd_ps(sb_f, v_two, eb_f), v_cb); - __m512 rq = _mm512_loadu_ps(rotated_q + i); - acc = _mm512_fmadd_ps(rq, recon, acc); - } - - float result = _mm512_reduce_add_ps(acc); - result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, 1, cb); - return result; -} -#endif // __AVX512F__ - -#if defined(__AVX2__) -/// AVX2: 8 dims/iter, ex_bits == 1. -inline float ip_1exbit_avx2( - const uint8_t* __restrict sign_bits, - const uint8_t* __restrict ex_code, - const float* __restrict rotated_q, - size_t d, - float cb) { - __m256 acc = _mm256_setzero_ps(); - const __m256 v_cb = _mm256_set1_ps(cb); - const __m256 v_two = _mm256_set1_ps(2.0f); - const __m256 v_one = _mm256_set1_ps(1.0f); - const __m256i bit_pos = _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128); - const __m256i zero = _mm256_setzero_si256(); - - size_t i = 0; - for (; i + 8 <= d; i += 8) { - uint8_t sb = sign_bits[i / 8]; - uint8_t eb = ex_code[i / 8]; - - __m256i sb_cmp = _mm256_cmpgt_epi32( - _mm256_and_si256(_mm256_set1_epi32(sb), bit_pos), zero); - __m256 sb_f = _mm256_and_ps(_mm256_castsi256_ps(sb_cmp), v_one); - - __m256i eb_cmp = _mm256_cmpgt_epi32( - _mm256_and_si256(_mm256_set1_epi32(eb), bit_pos), zero); - __m256 eb_f = _mm256_and_ps(_mm256_castsi256_ps(eb_cmp), v_one); - - __m256 recon = _mm256_add_ps(_mm256_fmadd_ps(sb_f, v_two, eb_f), v_cb); - __m256 rq = _mm256_loadu_ps(rotated_q + i); - acc = _mm256_fmadd_ps(rq, recon, acc); - } - - float result = hsum_avx2(acc); - result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, 1, cb); - return result; -} -#endif // __AVX2__ - -/********************************************************* - * Bit-plane decomposition kernels (ex_bits >= 2, BMI2 required). - * - * Decomposes the inner product as: - * ex_ip = (1 << ex_bits) * sign_dot - * + Σ_{b=0}^{ex_bits-1} (1 << b) * plane_dot_b - * + cb * total_q - * - * Each plane_dot_b is a float × bit-vector dot product, computed using - * the same bit→mask→float conversion as the 1-bit kernel. PEXT - * extracts each bit plane from the packed ex_code in one instruction - * per 8 dimensions. - *********************************************************/ - -#if defined(__AVX2__) && defined(__BMI2__) -/// AVX2 + BMI2 bit-plane decomposition: 8 dims/iter, ex_bits in [2, 7]. -/// Caller must ensure ex_bits <= 7 (pext_masks[7] / v_weights[8]). -inline float ip_bitplane_avx2( - const uint8_t* __restrict sign_bits, - const uint8_t* __restrict ex_code, - const float* __restrict rotated_q, - size_t d, - size_t ex_bits, - float cb) { - __m256 acc = _mm256_setzero_ps(); - const __m256 v_one = _mm256_set1_ps(1.0f); - const __m256i bit_pos = _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128); - const __m256i zero = _mm256_setzero_si256(); - const __m256 v_cb = _mm256_set1_ps(cb); - - // Precompute PEXT masks and plane weights - uint64_t pext_masks[7]; - __m256 v_weights[8]; - for (size_t b = 0; b < ex_bits; b++) { - uint64_t m = 0; - for (int j = 0; j < 8; j++) { - m |= (1ULL << (b + j * ex_bits)); - } - pext_masks[b] = m; - v_weights[b] = _mm256_set1_ps(static_cast(1u << b)); - } - v_weights[ex_bits] = _mm256_set1_ps(static_cast(1u << ex_bits)); - - size_t i = 0; - for (; i + 8 <= d; i += 8) { - // Sign bit → float via bit mask comparison - __m256i sb_cmp = _mm256_cmpgt_epi32( - _mm256_and_si256(_mm256_set1_epi32(sign_bits[i / 8]), bit_pos), - zero); - __m256 recon = _mm256_mul_ps( - _mm256_and_ps(_mm256_castsi256_ps(sb_cmp), v_one), - v_weights[ex_bits]); - - // Load packed ex_code for 8 dims (8 × ex_bits bits = ex_bits bytes) - uint64_t ex64 = 0; - memcpy(&ex64, ex_code + (i / 8) * ex_bits, sizeof(uint64_t)); - - // Extract each bit plane via PEXT → bit mask → float - for (size_t b = 0; b < ex_bits; b++) { - auto plane = static_cast(_pext_u64(ex64, pext_masks[b])); - __m256i p_cmp = _mm256_cmpgt_epi32( - _mm256_and_si256(_mm256_set1_epi32(plane), bit_pos), zero); - __m256 p_f = _mm256_and_ps(_mm256_castsi256_ps(p_cmp), v_one); - recon = _mm256_fmadd_ps(p_f, v_weights[b], recon); - } - - __m256 rq = _mm256_loadu_ps(rotated_q + i); - acc = _mm256_fmadd_ps(rq, _mm256_add_ps(recon, v_cb), acc); - } - - float result = hsum_avx2(acc); - result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb); - return result; -} -#endif // __AVX2__ && __BMI2__ - -#endif // x86_64 - /** * Dispatch to the best available kernel for the given ex_bits. * - * Routing (compile-time): - * ex_bits == 1: specialized 1-bit kernel (AVX-512 > AVX2 > scalar) - * ex_bits >= 2: bit-plane decomposition (AVX2+BMI2 > scalar) - * * @param sign_bits packed sign bits (1 bit/dim, standard byte packing) * @param ex_code packed extra-bit codes (ex_bits bits/dim) * @param rotated_q rotated query vector (float[d]) @@ -657,28 +181,24 @@ inline float ip_bitplane_avx2( * @param cb constant bias: -(2^ex_bits - 0.5) * @return inner product value */ -inline float compute_inner_product( +template +float compute_inner_product( + const uint8_t* __restrict sign_bits, + const uint8_t* __restrict ex_code, + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float cb); + +// NONE specialization — pure scalar +template <> +inline float compute_inner_product( const uint8_t* __restrict sign_bits, const uint8_t* __restrict ex_code, const float* __restrict rotated_q, size_t d, size_t ex_bits, float cb) { - if (ex_bits == 1) { -#if defined(__AVX512F__) - return ip_1exbit_avx512(sign_bits, ex_code, rotated_q, d, cb); -#elif defined(__AVX2__) - return ip_1exbit_avx2(sign_bits, ex_code, rotated_q, d, cb); -#else - return ip_scalar(sign_bits, ex_code, rotated_q, 0, d, 1, cb); -#endif - } - -#if defined(__AVX2__) && defined(__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_avx2.cpp b/faiss/utils/simd_impl/rabitq_avx2.cpp new file mode 100644 index 000000000..5feb2049f --- /dev/null +++ b/faiss/utils/simd_impl/rabitq_avx2.cpp @@ -0,0 +1,355 @@ +/* + * 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. + */ + +#ifdef COMPILE_SIMD_AVX2 + +#include +#include + +namespace faiss::rabitq { + +namespace { + +inline __m256i get_lookup_256() { + return _mm256_setr_epi8( + /* 0 */ 0, + /* 1 */ 1, + /* 2 */ 1, + /* 3 */ 2, + /* 4 */ 1, + /* 5 */ 2, + /* 6 */ 2, + /* 7 */ 3, + /* 8 */ 1, + /* 9 */ 2, + /* a */ 2, + /* b */ 3, + /* c */ 2, + /* d */ 3, + /* e */ 3, + /* f */ 4, + /* 0 */ 0, + /* 1 */ 1, + /* 2 */ 1, + /* 3 */ 2, + /* 4 */ 1, + /* 5 */ 2, + /* 6 */ 2, + /* 7 */ 3, + /* 8 */ 1, + /* 9 */ 2, + /* a */ 2, + /* b */ 3, + /* c */ 2, + /* d */ 3, + /* e */ 3, + /* f */ 4); +} + +inline __m256i popcount_256(__m256i v) { + const __m256i lookup = get_lookup_256(); + const __m256i low_mask = _mm256_set1_epi8(0x0f); + + const __m256i lo = _mm256_and_si256(v, low_mask); + const __m256i hi = _mm256_and_si256(_mm256_srli_epi16(v, 4), low_mask); + const __m256i popcnt_lo = _mm256_shuffle_epi8(lookup, lo); + const __m256i popcnt_hi = _mm256_shuffle_epi8(lookup, hi); + const __m256i popcnt = _mm256_add_epi8(popcnt_lo, popcnt_hi); + return _mm256_sad_epu8(_mm256_setzero_si256(), popcnt); +} + +inline uint64_t reduce_add_256(__m256i v) { + alignas(32) uint64_t lanes[4]; + _mm256_store_si256((__m256i*)lanes, v); + return lanes[0] + lanes[1] + lanes[2] + lanes[3]; +} + +inline __m128i popcount_128(__m128i v) { + uint64_t lane0 = _mm_extract_epi64(v, 0); + uint64_t lane1 = _mm_extract_epi64(v, 1); + uint64_t pop0 = __builtin_popcountll(lane0); + uint64_t pop1 = __builtin_popcountll(lane1); + return _mm_set_epi64x(pop1, pop0); +} + +inline uint64_t reduce_add_128(__m128i v) { + alignas(16) uint64_t lanes[2]; + _mm_store_si128((__m128i*)lanes, v); + return lanes[0] + lanes[1]; +} + +} // namespace + +template <> +uint64_t bitwise_and_dot_product( + const uint8_t* query, + const uint8_t* data, + size_t size, + size_t qb) { + uint64_t sum = 0; + size_t offset = 0; + if (size_t step = 256 / 8; offset + step <= size) { + __m256i sum_256 = _mm256_setzero_si256(); + for (; offset + step <= size; offset += step) { + __m256i v_x = _mm256_loadu_si256((const __m256i*)(data + offset)); + for (int j = 0; j < qb; j++) { + __m256i v_q = _mm256_loadu_si256( + (const __m256i*)(query + j * size + offset)); + __m256i v_and = _mm256_and_si256(v_q, v_x); + __m256i v_popcnt = popcount_256(v_and); + __m256i v_shifted = _mm256_slli_epi64(v_popcnt, j); + sum_256 = _mm256_add_epi64(sum_256, v_shifted); + } + } + sum += reduce_add_256(sum_256); + } + __m128i sum_128 = _mm_setzero_si128(); + for (size_t step = 128 / 8; offset + step <= size; offset += step) { + __m128i v_x = _mm_loadu_si128((const __m128i*)(data + offset)); + for (int j = 0; j < qb; j++) { + __m128i v_q = _mm_loadu_si128( + (const __m128i*)(query + j * size + offset)); + __m128i v_and = _mm_and_si128(v_q, v_x); + __m128i v_popcnt = popcount_128(v_and); + __m128i v_shifted = _mm_slli_epi64(v_popcnt, j); + sum_128 = _mm_add_epi64(sum_128, v_shifted); + } + } + 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); + for (int j = 0; j < qb; j++) { + const auto qv = *(const uint64_t*)(query + j * size + offset); + sum += __builtin_popcountll(qv & yv) << j; + } + } + for (; offset < size; ++offset) { + const auto yv = *(data + offset); + for (int j = 0; j < qb; j++) { + const auto qv = *(query + j * size + offset); + sum += __builtin_popcount(qv & yv) << j; + } + } + return sum; +} + +template <> +uint64_t bitwise_xor_dot_product( + const uint8_t* query, + const uint8_t* data, + size_t size, + size_t qb) { + uint64_t sum = 0; + size_t offset = 0; + if (size_t step = 256 / 8; offset + step <= size) { + __m256i sum_256 = _mm256_setzero_si256(); + for (; offset + step <= size; offset += step) { + __m256i v_x = _mm256_loadu_si256((const __m256i*)(data + offset)); + for (int j = 0; j < qb; j++) { + __m256i v_q = _mm256_loadu_si256( + (const __m256i*)(query + j * size + offset)); + __m256i v_xor = _mm256_xor_si256(v_q, v_x); + __m256i v_popcnt = popcount_256(v_xor); + __m256i v_shifted = _mm256_slli_epi64(v_popcnt, j); + sum_256 = _mm256_add_epi64(sum_256, v_shifted); + } + } + sum += reduce_add_256(sum_256); + } + __m128i sum_128 = _mm_setzero_si128(); + for (size_t step = 128 / 8; offset + step <= size; offset += step) { + __m128i v_x = _mm_loadu_si128((const __m128i*)(data + offset)); + for (int j = 0; j < qb; j++) { + __m128i v_q = _mm_loadu_si128( + (const __m128i*)(query + j * size + offset)); + __m128i v_xor = _mm_xor_si128(v_q, v_x); + __m128i v_popcnt = popcount_128(v_xor); + __m128i v_shifted = _mm_slli_epi64(v_popcnt, j); + sum_128 = _mm_add_epi64(sum_128, v_shifted); + } + } + 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); + for (int j = 0; j < qb; j++) { + const auto qv = *(const uint64_t*)(query + j * size + offset); + sum += __builtin_popcountll(qv ^ yv) << j; + } + } + for (; offset < size; ++offset) { + const auto yv = *(data + offset); + for (int j = 0; j < qb; j++) { + const auto qv = *(query + j * size + offset); + sum += __builtin_popcount(qv ^ yv) << j; + } + } + return sum; +} + +template <> +uint64_t popcount(const uint8_t* data, size_t size) { + uint64_t sum = 0; + size_t offset = 0; + if (offset + 256 / 8 <= size) { + __m256i sum_256 = _mm256_setzero_si256(); + for (size_t end; (end = offset + 256 / 8) <= size; offset = end) { + __m256i v_x = _mm256_loadu_si256((const __m256i*)(data + offset)); + __m256i v_popcnt = popcount_256(v_x); + sum_256 = _mm256_add_epi64(sum_256, v_popcnt); + } + sum += reduce_add_256(sum_256); + } + __m128i sum_128 = _mm_setzero_si128(); + for (size_t step = 128 / 8; offset + step <= size; offset += step) { + __m128i v_x = _mm_loadu_si128((const __m128i*)(data + offset)); + sum_128 = _mm_add_epi64(sum_128, popcount_128(v_x)); + } + 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); + sum += __builtin_popcountll(yv); + } + for (; offset < size; ++offset) { + const auto yv = *(data + offset); + sum += __builtin_popcount(yv); + } + return sum; +} + +} // namespace faiss::rabitq + +namespace faiss::rabitq::multibit { + +namespace { + +inline float hsum_avx2(__m256 v) { + __m128 hi = _mm256_extractf128_ps(v, 1); + __m128 lo = _mm256_castps256_ps128(v); + lo = _mm_add_ps(lo, hi); + __m128 shuf = _mm_movehdup_ps(lo); + lo = _mm_add_ps(lo, shuf); + shuf = _mm_movehl_ps(shuf, lo); + return _mm_cvtss_f32(_mm_add_ss(lo, shuf)); +} + +inline float ip_1exbit_avx2( + const uint8_t* __restrict sign_bits, + const uint8_t* __restrict ex_code, + const float* __restrict rotated_q, + size_t d, + float cb) { + __m256 acc = _mm256_setzero_ps(); + const __m256 v_cb = _mm256_set1_ps(cb); + const __m256 v_two = _mm256_set1_ps(2.0f); + const __m256 v_one = _mm256_set1_ps(1.0f); + const __m256i bit_pos = _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128); + const __m256i zero = _mm256_setzero_si256(); + + size_t i = 0; + for (; i + 8 <= d; i += 8) { + uint8_t sb = sign_bits[i / 8]; + uint8_t eb = ex_code[i / 8]; + + __m256i sb_cmp = _mm256_cmpgt_epi32( + _mm256_and_si256(_mm256_set1_epi32(sb), bit_pos), zero); + __m256 sb_f = _mm256_and_ps(_mm256_castsi256_ps(sb_cmp), v_one); + + __m256i eb_cmp = _mm256_cmpgt_epi32( + _mm256_and_si256(_mm256_set1_epi32(eb), bit_pos), zero); + __m256 eb_f = _mm256_and_ps(_mm256_castsi256_ps(eb_cmp), v_one); + + __m256 recon = _mm256_add_ps(_mm256_fmadd_ps(sb_f, v_two, eb_f), v_cb); + __m256 rq = _mm256_loadu_ps(rotated_q + i); + acc = _mm256_fmadd_ps(rq, recon, acc); + } + + float result = hsum_avx2(acc); + result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, 1, cb); + return result; +} + +#ifdef __BMI2__ +inline float ip_bitplane_avx2( + const uint8_t* __restrict sign_bits, + const uint8_t* __restrict ex_code, + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float cb) { + __m256 acc = _mm256_setzero_ps(); + const __m256 v_one = _mm256_set1_ps(1.0f); + const __m256i bit_pos = _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128); + const __m256i zero = _mm256_setzero_si256(); + const __m256 v_cb = _mm256_set1_ps(cb); + + uint64_t pext_masks[7]; + __m256 v_weights[8]; + for (size_t b = 0; b < ex_bits; b++) { + uint64_t m = 0; + for (int j = 0; j < 8; j++) { + m |= (1ULL << (b + j * ex_bits)); + } + pext_masks[b] = m; + v_weights[b] = _mm256_set1_ps(static_cast(1u << b)); + } + v_weights[ex_bits] = _mm256_set1_ps(static_cast(1u << ex_bits)); + + size_t i = 0; + for (; i + 8 <= d; i += 8) { + __m256i sb_cmp = _mm256_cmpgt_epi32( + _mm256_and_si256(_mm256_set1_epi32(sign_bits[i / 8]), bit_pos), + zero); + __m256 recon = _mm256_mul_ps( + _mm256_and_ps(_mm256_castsi256_ps(sb_cmp), v_one), + v_weights[ex_bits]); + + uint64_t ex64 = 0; + memcpy(&ex64, ex_code + (i / 8) * ex_bits, sizeof(uint64_t)); + + for (size_t b = 0; b < ex_bits; b++) { + auto plane = static_cast(_pext_u64(ex64, pext_masks[b])); + __m256i p_cmp = _mm256_cmpgt_epi32( + _mm256_and_si256(_mm256_set1_epi32(plane), bit_pos), zero); + __m256 p_f = _mm256_and_ps(_mm256_castsi256_ps(p_cmp), v_one); + recon = _mm256_fmadd_ps(p_f, v_weights[b], recon); + } + + __m256 rq = _mm256_loadu_ps(rotated_q + i); + acc = _mm256_fmadd_ps(rq, _mm256_add_ps(recon, v_cb), acc); + } + + float result = hsum_avx2(acc); + result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb); + return result; +} +#endif // __BMI2__ + +} // namespace + +template <> +float compute_inner_product( + const uint8_t* __restrict sign_bits, + const uint8_t* __restrict ex_code, + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float 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); +} + +} // namespace faiss::rabitq::multibit + +#endif // COMPILE_SIMD_AVX2 diff --git a/faiss/utils/simd_impl/rabitq_avx512.cpp b/faiss/utils/simd_impl/rabitq_avx512.cpp new file mode 100644 index 000000000..06413fe1f --- /dev/null +++ b/faiss/utils/simd_impl/rabitq_avx512.cpp @@ -0,0 +1,477 @@ +/* + * 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. + */ + +#ifdef COMPILE_SIMD_AVX512 + +#include +#include + +namespace faiss::rabitq { + +namespace { + +inline __m512i get_lookup_512() { + return _mm512_set_epi8( + /* f */ 4, + /* e */ 3, + /* d */ 3, + /* c */ 2, + /* b */ 3, + /* a */ 2, + /* 9 */ 2, + /* 8 */ 1, + /* 7 */ 3, + /* 6 */ 2, + /* 5 */ 2, + /* 4 */ 1, + /* 3 */ 2, + /* 2 */ 1, + /* 1 */ 1, + /* 0 */ 0, + /* f */ 4, + /* e */ 3, + /* d */ 3, + /* c */ 2, + /* b */ 3, + /* a */ 2, + /* 9 */ 2, + /* 8 */ 1, + /* 7 */ 3, + /* 6 */ 2, + /* 5 */ 2, + /* 4 */ 1, + /* 3 */ 2, + /* 2 */ 1, + /* 1 */ 1, + /* 0 */ 0, + /* f */ 4, + /* e */ 3, + /* d */ 3, + /* c */ 2, + /* b */ 3, + /* a */ 2, + /* 9 */ 2, + /* 8 */ 1, + /* 7 */ 3, + /* 6 */ 2, + /* 5 */ 2, + /* 4 */ 1, + /* 3 */ 2, + /* 2 */ 1, + /* 1 */ 1, + /* 0 */ 0, + /* f */ 4, + /* e */ 3, + /* d */ 3, + /* c */ 2, + /* b */ 3, + /* a */ 2, + /* 9 */ 2, + /* 8 */ 1, + /* 7 */ 3, + /* 6 */ 2, + /* 5 */ 2, + /* 4 */ 1, + /* 3 */ 2, + /* 2 */ 1, + /* 1 */ 1, + /* 0 */ 0); +} + +inline __m512i popcount_512(__m512i v) { +#if defined(__AVX512VPOPCNTDQ__) + return _mm512_popcnt_epi64(v); +#else + const __m512i lookup = get_lookup_512(); + const __m512i low_mask = _mm512_set1_epi8(0x0f); + + const __m512i lo = _mm512_and_si512(v, low_mask); + const __m512i hi = _mm512_and_si512(_mm512_srli_epi16(v, 4), low_mask); + const __m512i popcnt_lo = _mm512_shuffle_epi8(lookup, lo); + const __m512i popcnt_hi = _mm512_shuffle_epi8(lookup, hi); + const __m512i popcnt = _mm512_add_epi8(popcnt_lo, popcnt_hi); + return _mm512_sad_epu8(_mm512_setzero_si512(), popcnt); +#endif +} + +// AVX2 helpers needed for AVX512 fallback paths (compute_inner_product) +inline __m256i get_lookup_256() { + return _mm256_setr_epi8( + /* 0 */ 0, + /* 1 */ 1, + /* 2 */ 1, + /* 3 */ 2, + /* 4 */ 1, + /* 5 */ 2, + /* 6 */ 2, + /* 7 */ 3, + /* 8 */ 1, + /* 9 */ 2, + /* a */ 2, + /* b */ 3, + /* c */ 2, + /* d */ 3, + /* e */ 3, + /* f */ 4, + /* 0 */ 0, + /* 1 */ 1, + /* 2 */ 1, + /* 3 */ 2, + /* 4 */ 1, + /* 5 */ 2, + /* 6 */ 2, + /* 7 */ 3, + /* 8 */ 1, + /* 9 */ 2, + /* a */ 2, + /* b */ 3, + /* c */ 2, + /* d */ 3, + /* e */ 3, + /* f */ 4); +} + +inline __m256i popcount_256(__m256i v) { + const __m256i lookup = get_lookup_256(); + const __m256i low_mask = _mm256_set1_epi8(0x0f); + + const __m256i lo = _mm256_and_si256(v, low_mask); + const __m256i hi = _mm256_and_si256(_mm256_srli_epi16(v, 4), low_mask); + const __m256i popcnt_lo = _mm256_shuffle_epi8(lookup, lo); + const __m256i popcnt_hi = _mm256_shuffle_epi8(lookup, hi); + const __m256i popcnt = _mm256_add_epi8(popcnt_lo, popcnt_hi); + return _mm256_sad_epu8(_mm256_setzero_si256(), popcnt); +} + +inline uint64_t reduce_add_256(__m256i v) { + alignas(32) uint64_t lanes[4]; + _mm256_store_si256((__m256i*)lanes, v); + return lanes[0] + lanes[1] + lanes[2] + lanes[3]; +} + +inline __m128i popcount_128(__m128i v) { + uint64_t lane0 = _mm_extract_epi64(v, 0); + uint64_t lane1 = _mm_extract_epi64(v, 1); + uint64_t pop0 = __builtin_popcountll(lane0); + uint64_t pop1 = __builtin_popcountll(lane1); + return _mm_set_epi64x(pop1, pop0); +} + +inline uint64_t reduce_add_128(__m128i v) { + alignas(16) uint64_t lanes[2]; + _mm_store_si128((__m128i*)lanes, v); + return lanes[0] + lanes[1]; +} + +} // namespace + +template <> +uint64_t bitwise_and_dot_product( + const uint8_t* query, + const uint8_t* data, + size_t size, + size_t qb) { + uint64_t sum = 0; + size_t offset = 0; + if (size_t step = 512 / 8; offset + step <= size) { + __m512i sum_512 = _mm512_setzero_si512(); + for (; offset + step <= size; offset += step) { + __m512i v_x = _mm512_loadu_si512((const __m512i*)(data + offset)); + for (int j = 0; j < qb; j++) { + __m512i v_q = _mm512_loadu_si512( + (const __m512i*)(query + j * size + offset)); + __m512i v_and = _mm512_and_si512(v_q, v_x); + __m512i v_popcnt = popcount_512(v_and); + __m512i v_shifted = _mm512_slli_epi64(v_popcnt, j); + sum_512 = _mm512_add_epi64(sum_512, v_shifted); + } + } + sum += _mm512_reduce_add_epi64(sum_512); + } + if (size_t step = 256 / 8; offset + step <= size) { + __m256i sum_256 = _mm256_setzero_si256(); + for (; offset + step <= size; offset += step) { + __m256i v_x = _mm256_loadu_si256((const __m256i*)(data + offset)); + for (int j = 0; j < qb; j++) { + __m256i v_q = _mm256_loadu_si256( + (const __m256i*)(query + j * size + offset)); + __m256i v_and = _mm256_and_si256(v_q, v_x); + __m256i v_popcnt = popcount_256(v_and); + __m256i v_shifted = _mm256_slli_epi64(v_popcnt, j); + sum_256 = _mm256_add_epi64(sum_256, v_shifted); + } + } + sum += reduce_add_256(sum_256); + } + __m128i sum_128 = _mm_setzero_si128(); + for (size_t step = 128 / 8; offset + step <= size; offset += step) { + __m128i v_x = _mm_loadu_si128((const __m128i*)(data + offset)); + for (int j = 0; j < qb; j++) { + __m128i v_q = _mm_loadu_si128( + (const __m128i*)(query + j * size + offset)); + __m128i v_and = _mm_and_si128(v_q, v_x); + __m128i v_popcnt = popcount_128(v_and); + __m128i v_shifted = _mm_slli_epi64(v_popcnt, j); + sum_128 = _mm_add_epi64(sum_128, v_shifted); + } + } + 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); + for (int j = 0; j < qb; j++) { + const auto qv = *(const uint64_t*)(query + j * size + offset); + sum += __builtin_popcountll(qv & yv) << j; + } + } + for (; offset < size; ++offset) { + const auto yv = *(data + offset); + for (int j = 0; j < qb; j++) { + const auto qv = *(query + j * size + offset); + sum += __builtin_popcount(qv & yv) << j; + } + } + return sum; +} + +template <> +uint64_t bitwise_xor_dot_product( + const uint8_t* query, + const uint8_t* data, + size_t size, + size_t qb) { + uint64_t sum = 0; + size_t offset = 0; + if (size_t step = 512 / 8; offset + step <= size) { + __m512i sum_512 = _mm512_setzero_si512(); + for (; offset + step <= size; offset += step) { + __m512i v_x = _mm512_loadu_si512((const __m512i*)(data + offset)); + for (int j = 0; j < qb; j++) { + __m512i v_q = _mm512_loadu_si512( + (const __m512i*)(query + j * size + offset)); + __m512i v_xor = _mm512_xor_si512(v_q, v_x); + __m512i v_popcnt = popcount_512(v_xor); + __m512i v_shifted = _mm512_slli_epi64(v_popcnt, j); + sum_512 = _mm512_add_epi64(sum_512, v_shifted); + } + } + sum += _mm512_reduce_add_epi64(sum_512); + } + if (size_t step = 256 / 8; offset + step <= size) { + __m256i sum_256 = _mm256_setzero_si256(); + for (; offset + step <= size; offset += step) { + __m256i v_x = _mm256_loadu_si256((const __m256i*)(data + offset)); + for (int j = 0; j < qb; j++) { + __m256i v_q = _mm256_loadu_si256( + (const __m256i*)(query + j * size + offset)); + __m256i v_xor = _mm256_xor_si256(v_q, v_x); + __m256i v_popcnt = popcount_256(v_xor); + __m256i v_shifted = _mm256_slli_epi64(v_popcnt, j); + sum_256 = _mm256_add_epi64(sum_256, v_shifted); + } + } + sum += reduce_add_256(sum_256); + } + __m128i sum_128 = _mm_setzero_si128(); + for (size_t step = 128 / 8; offset + step <= size; offset += step) { + __m128i v_x = _mm_loadu_si128((const __m128i*)(data + offset)); + for (int j = 0; j < qb; j++) { + __m128i v_q = _mm_loadu_si128( + (const __m128i*)(query + j * size + offset)); + __m128i v_xor = _mm_xor_si128(v_q, v_x); + __m128i v_popcnt = popcount_128(v_xor); + __m128i v_shifted = _mm_slli_epi64(v_popcnt, j); + sum_128 = _mm_add_epi64(sum_128, v_shifted); + } + } + 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); + for (int j = 0; j < qb; j++) { + const auto qv = *(const uint64_t*)(query + j * size + offset); + sum += __builtin_popcountll(qv ^ yv) << j; + } + } + for (; offset < size; ++offset) { + const auto yv = *(data + offset); + for (int j = 0; j < qb; j++) { + const auto qv = *(query + j * size + offset); + sum += __builtin_popcount(qv ^ yv) << j; + } + } + return sum; +} + +template <> +uint64_t popcount(const uint8_t* data, size_t size) { + uint64_t sum = 0; + size_t offset = 0; + if (offset + 512 / 8 <= size) { + __m512i sum_512 = _mm512_setzero_si512(); + for (size_t end; (end = offset + 512 / 8) <= size; offset = end) { + __m512i v_x = _mm512_loadu_si512((const __m512i*)(data + offset)); + __m512i v_popcnt = popcount_512(v_x); + sum_512 = _mm512_add_epi64(sum_512, v_popcnt); + } + sum += _mm512_reduce_add_epi64(sum_512); + } + if (offset + 256 / 8 <= size) { + __m256i sum_256 = _mm256_setzero_si256(); + for (size_t end; (end = offset + 256 / 8) <= size; offset = end) { + __m256i v_x = _mm256_loadu_si256((const __m256i*)(data + offset)); + __m256i v_popcnt = popcount_256(v_x); + sum_256 = _mm256_add_epi64(sum_256, v_popcnt); + } + sum += reduce_add_256(sum_256); + } + __m128i sum_128 = _mm_setzero_si128(); + for (size_t step = 128 / 8; offset + step <= size; offset += step) { + __m128i v_x = _mm_loadu_si128((const __m128i*)(data + offset)); + sum_128 = _mm_add_epi64(sum_128, popcount_128(v_x)); + } + 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); + sum += __builtin_popcountll(yv); + } + for (; offset < size; ++offset) { + const auto yv = *(data + offset); + sum += __builtin_popcount(yv); + } + return sum; +} + +} // namespace faiss::rabitq + +namespace faiss::rabitq::multibit { + +namespace { + +inline float hsum_avx2(__m256 v) { + __m128 hi = _mm256_extractf128_ps(v, 1); + __m128 lo = _mm256_castps256_ps128(v); + lo = _mm_add_ps(lo, hi); + __m128 shuf = _mm_movehdup_ps(lo); + lo = _mm_add_ps(lo, shuf); + shuf = _mm_movehl_ps(shuf, lo); + return _mm_cvtss_f32(_mm_add_ss(lo, shuf)); +} + +inline float ip_1exbit_avx512( + const uint8_t* __restrict sign_bits, + const uint8_t* __restrict ex_code, + const float* __restrict rotated_q, + size_t d, + float cb) { + __m512 acc = _mm512_setzero_ps(); + const __m512 v_cb = _mm512_set1_ps(cb); + const __m512 v_two = _mm512_set1_ps(2.0f); + const __m512 v_one = _mm512_set1_ps(1.0f); + + size_t i = 0; + for (; i + 16 <= d; i += 16) { + uint16_t sb16; + memcpy(&sb16, sign_bits + i / 8, sizeof(uint16_t)); + uint16_t eb16; + memcpy(&eb16, ex_code + i / 8, sizeof(uint16_t)); + + __m512 sb_f = _mm512_maskz_mov_ps(_cvtu32_mask16(sb16), v_one); + __m512 eb_f = _mm512_maskz_mov_ps(_cvtu32_mask16(eb16), v_one); + + __m512 recon = _mm512_add_ps(_mm512_fmadd_ps(sb_f, v_two, eb_f), v_cb); + __m512 rq = _mm512_loadu_ps(rotated_q + i); + acc = _mm512_fmadd_ps(rq, recon, acc); + } + + float result = _mm512_reduce_add_ps(acc); + result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, 1, cb); + return result; +} + +// AVX2+BMI2 bitplane kernel used as fallback for ex_bits >= 2. +// AVX512 TU has AVX2 available. BMI2 guarded separately since +// VIA Eden X4 has AVX2 without BMI2. +#ifdef __BMI2__ +inline float ip_bitplane_avx2( + const uint8_t* __restrict sign_bits, + const uint8_t* __restrict ex_code, + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float cb) { + __m256 acc = _mm256_setzero_ps(); + const __m256 v_one = _mm256_set1_ps(1.0f); + const __m256i bit_pos = _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128); + const __m256i zero = _mm256_setzero_si256(); + const __m256 v_cb = _mm256_set1_ps(cb); + + uint64_t pext_masks[7]; + __m256 v_weights[8]; + for (size_t b = 0; b < ex_bits; b++) { + uint64_t m = 0; + for (int j = 0; j < 8; j++) { + m |= (1ULL << (b + j * ex_bits)); + } + pext_masks[b] = m; + v_weights[b] = _mm256_set1_ps(static_cast(1u << b)); + } + v_weights[ex_bits] = _mm256_set1_ps(static_cast(1u << ex_bits)); + + size_t i = 0; + for (; i + 8 <= d; i += 8) { + __m256i sb_cmp = _mm256_cmpgt_epi32( + _mm256_and_si256(_mm256_set1_epi32(sign_bits[i / 8]), bit_pos), + zero); + __m256 recon = _mm256_mul_ps( + _mm256_and_ps(_mm256_castsi256_ps(sb_cmp), v_one), + v_weights[ex_bits]); + + uint64_t ex64 = 0; + memcpy(&ex64, ex_code + (i / 8) * ex_bits, sizeof(uint64_t)); + + for (size_t b = 0; b < ex_bits; b++) { + auto plane = static_cast(_pext_u64(ex64, pext_masks[b])); + __m256i p_cmp = _mm256_cmpgt_epi32( + _mm256_and_si256(_mm256_set1_epi32(plane), bit_pos), zero); + __m256 p_f = _mm256_and_ps(_mm256_castsi256_ps(p_cmp), v_one); + recon = _mm256_fmadd_ps(p_f, v_weights[b], recon); + } + + __m256 rq = _mm256_loadu_ps(rotated_q + i); + acc = _mm256_fmadd_ps(rq, _mm256_add_ps(recon, v_cb), acc); + } + + float result = hsum_avx2(acc); + result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb); + return result; +} +#endif // __BMI2__ + +} // namespace + +template <> +float compute_inner_product( + const uint8_t* __restrict sign_bits, + const uint8_t* __restrict ex_code, + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float 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_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); +} + +} // namespace faiss::rabitq::multibit + +#endif // COMPILE_SIMD_AVX512 diff --git a/faiss/utils/simd_impl/rabitq_neon.cpp b/faiss/utils/simd_impl/rabitq_neon.cpp new file mode 100644 index 000000000..05be093c8 --- /dev/null +++ b/faiss/utils/simd_impl/rabitq_neon.cpp @@ -0,0 +1,55 @@ +/* + * 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 + +#ifdef COMPILE_SIMD_ARM_NEON + +namespace faiss::rabitq { + +template <> +uint64_t bitwise_and_dot_product( + const uint8_t* query, + const uint8_t* data, + size_t size, + size_t qb) { + return bitwise_and_dot_product(query, data, size, qb); +} + +template <> +uint64_t bitwise_xor_dot_product( + const uint8_t* query, + const uint8_t* data, + size_t size, + size_t qb) { + return bitwise_xor_dot_product(query, data, size, qb); +} + +template <> +uint64_t popcount(const uint8_t* data, size_t size) { + return popcount(data, size); +} + +} // namespace faiss::rabitq + +namespace faiss::rabitq::multibit { + +template <> +float compute_inner_product( + const uint8_t* __restrict sign_bits, + const uint8_t* __restrict ex_code, + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float cb) { + return compute_inner_product( + sign_bits, ex_code, rotated_q, d, ex_bits, cb); +} + +} // namespace faiss::rabitq::multibit + +#endif // COMPILE_SIMD_ARM_NEON