diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index c4e7a0d66..8b1c2cf88 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -271,6 +271,7 @@ set(FAISS_HEADERS impl/fast_scan/decompose_qbs.h impl/fast_scan/kernels_simd256.h impl/fast_scan/kernels_simd512.h + impl/fast_scan/rabitq_dispatching.h impl/fast_scan/rabitq_result_handler.h impl/residual_quantizer_encode_steps.h impl/simd_dispatch.h diff --git a/faiss/IndexFastScan.cpp b/faiss/IndexFastScan.cpp index c23f139e9..045462d3a 100644 --- a/faiss/IndexFastScan.cpp +++ b/faiss/IndexFastScan.cpp @@ -211,6 +211,19 @@ void estimators_from_tables_generic( } // anonymous namespace +std::unique_ptr IndexFastScan::make_knn_scanner( + bool is_max, + idx_t n, + idx_t k, + size_t ntotal, + float* distances, + idx_t* labels, + const IDSelector* sel, + const FastScanDistancePostProcessing&) const { + return make_fast_scan_knn_scanner( + is_max, 0, n, ntotal, k, distances, labels, sel); +} + // Default implementation of make_knn_handler with centralized fallback logic SIMDResultHandlerToFloat* IndexFastScan::make_knn_handler( bool is_max, diff --git a/faiss/IndexFastScan.h b/faiss/IndexFastScan.h index 8b509ce81..2a4e81f65 100644 --- a/faiss/IndexFastScan.h +++ b/faiss/IndexFastScan.h @@ -7,8 +7,11 @@ #pragma once +#include + #include #include +#include #include namespace faiss { @@ -121,6 +124,22 @@ struct IndexFastScan : Index { const float* x, const FastScanDistancePostProcessing& context) const = 0; + /** Create a SIMD-dispatched scanner for knn search. + * + * Returns a FastScanCodeScanner that bundles handler + accumulation + * kernel behind the SIMD dispatch boundary. + * The scanner's accumulate methods dispatch to the optimal SIMD level. + */ + virtual std::unique_ptr make_knn_scanner( + bool is_max, + idx_t n, + idx_t k, + size_t ntotal, + float* distances, + idx_t* labels, + const IDSelector* sel, + const FastScanDistancePostProcessing& context = {}) const; + /** Create a KNN handler for this index type * * This method can be overridden by derived classes to provide diff --git a/faiss/IndexIVFFastScan.cpp b/faiss/IndexIVFFastScan.cpp index dcfef9649..b59265655 100644 --- a/faiss/IndexIVFFastScan.cpp +++ b/faiss/IndexIVFFastScan.cpp @@ -495,6 +495,26 @@ int compute_search_nslice( } // namespace +std::unique_ptr IndexIVFFastScan::make_knn_scanner( + bool is_max, + idx_t n, + idx_t k, + float* distances, + idx_t* labels, + const IDSelector* sel, + const FastScanDistancePostProcessing&) const { + return make_fast_scan_knn_scanner( + is_max, + 0, + n, + 0, + k, + distances, + labels, + sel, + /*with_id_map=*/true); +} + SIMDResultHandlerToFloat* IndexIVFFastScan::make_knn_handler( bool is_max, int impl, diff --git a/faiss/IndexIVFFastScan.h b/faiss/IndexIVFFastScan.h index 4dd9c9dd6..8ba299f63 100644 --- a/faiss/IndexIVFFastScan.h +++ b/faiss/IndexIVFFastScan.h @@ -9,6 +9,7 @@ #include #include +#include #include namespace faiss { @@ -221,6 +222,31 @@ struct IndexIVFFastScan : IndexIVF { RangeSearchResult* result, const SearchParameters* params = nullptr) const override; + /** Create a SIMD-dispatched scanner for knn search (IVF variant). + * + * Returns a FastScanCodeScanner that bundles handler + accumulation + * kernel behind the SIMD dispatch boundary. ntotal is not passed + * because IVF sets it per-list via handler->ntotal. + * Derived classes that need custom handlers (e.g. RaBitQ) override + * this to return nullptr, falling back to make_knn_handler. + * + * @param is_max whether to use CMax comparator (true) or CMin + * @param n number of queries + * @param k number of neighbors to find + * @param distances output distances array + * @param labels output labels array + * @param sel optional ID selector + * @return scanner, or nullptr if unsupported + */ + virtual std::unique_ptr make_knn_scanner( + bool is_max, + idx_t n, + idx_t k, + float* distances, + idx_t* labels, + const IDSelector* sel, + const FastScanDistancePostProcessing& context = {}) const; + /** Create a KNN handler for this index type * * This method can be overridden by derived classes to provide diff --git a/faiss/IndexIVFRaBitQFastScan.cpp b/faiss/IndexIVFRaBitQFastScan.cpp index afbba7c5a..5ef46d5b8 100644 --- a/faiss/IndexIVFRaBitQFastScan.cpp +++ b/faiss/IndexIVFRaBitQFastScan.cpp @@ -539,6 +539,19 @@ void IndexIVFRaBitQFastScan::decode_fastscan_to_residual( } } +std::unique_ptr IndexIVFRaBitQFastScan::make_knn_scanner( + bool is_max, + idx_t n, + idx_t k, + float* distances, + idx_t* labels, + const IDSelector* /*sel*/, + const FastScanDistancePostProcessing& context) const { + const bool is_multibit = (rabitq.nb_bits - 1) > 0; + return rabitq_ivf_make_knn_scanner( + is_max, this, n, k, distances, labels, &context, is_multibit); +} + // Implementation of virtual make_knn_handler method SIMDResultHandlerToFloat* IndexIVFRaBitQFastScan::make_knn_handler( bool is_max, @@ -562,266 +575,6 @@ SIMDResultHandlerToFloat* IndexIVFRaBitQFastScan::make_knn_handler( } } -/********************************************************* - * simd_result_handlers::IVFRaBitQHeapHandler implementation - *********************************************************/ - -template -simd_result_handlers::IVFRaBitQHeapHandler::IVFRaBitQHeapHandler( - const IndexIVFRaBitQFastScan* idx, - size_t nq_val, - size_t k_val, - float* distances, - int64_t* labels, - const FastScanDistancePostProcessing* ctx, - bool multibit) - : simd_result_handlers::ResultHandlerCompare( - nq_val, - 0, - nullptr), - index(idx), - heap_distances(distances), - heap_labels(labels), - nq(nq_val), - k(k_val), - context(ctx), - is_multibit(multibit), - storage_size(idx->compute_per_vector_storage_size()), - packed_block_size(((idx->M2 + 1) / 2) * idx->bbs), - full_block_size(idx->get_block_stride()), - packer(idx->get_CodePacker()) { - current_list_no = 0; - probe_indices.clear(); - - // Initialize heaps in constructor (standard pattern from HeapHandler) - for (int64_t q = 0; q < static_cast(nq); q++) { - float* heap_dis = heap_distances + q * k; - int64_t* heap_ids = heap_labels + q * k; - heap_heapify(k, heap_dis, heap_ids); - } -} - -template -void simd_result_handlers::IVFRaBitQHeapHandler::handle( - size_t q, - size_t b, - simd16uint16 d0, - simd16uint16 d1) { - // Store the original local query index before adjust_with_origin changes it - size_t local_q = q; - this->adjust_with_origin(q, d0, d1); - - ALIGNED(32) uint16_t d32tab[32]; - d0.store(d32tab); - d1.store(d32tab + 16); - - float* const heap_dis = heap_distances + q * k; - int64_t* const heap_ids = heap_labels + q * k; - - FAISS_THROW_IF_NOT_FMT( - !probe_indices.empty() && local_q < probe_indices.size(), - "set_list_context() must be called before handle() - probe_indices size: %zu, local_q: %zu, global_q: %zu", - probe_indices.size(), - local_q, - q); - - // Access query factors directly from array via ProcessingContext - if (!context || !context->query_factors) { - FAISS_THROW_MSG( - "Query factors not available: FastScanDistancePostProcessing with query_factors required"); - } - - // Use probe_rank from probe_indices for compact storage indexing - size_t probe_rank = probe_indices[local_q]; - size_t nprobe = context->nprobe > 0 ? context->nprobe : index->nprobe; - size_t storage_idx = q * nprobe + probe_rank; - - const auto& query_factors = context->query_factors[storage_idx]; - - const float one_a = - this->normalizers ? (1.0f / this->normalizers[2 * q]) : 1.0f; - const float bias = this->normalizers ? this->normalizers[2 * q + 1] : 0.0f; - - uint64_t idx_base = this->j0 + b * 32; - if (idx_base >= this->ntotal) { - return; - } - - size_t max_positions = std::min(32, this->ntotal - idx_base); - - // Stats tracking for two-stage search - // n_1bit_evaluations: candidates evaluated using 1-bit lower bound - // n_multibit_evaluations: candidates requiring full multi-bit distance - size_t local_1bit_evaluations = 0; - size_t local_multibit_evaluations = 0; - - // Process each candidate vector in the SIMD batch - for (size_t j = 0; j < max_positions; j++) { - const int64_t result_id = this->adjust_id(b, j); - - if (result_id < 0) { - continue; - } - - const float normalized_distance = d32tab[j] * one_a + bias; - - const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr( - list_codes_ptr, - idx_base + j, - index->bbs, - packed_block_size, - full_block_size, - storage_size); - - if (is_multibit) { - // Track candidates actually considered for two-stage filtering - local_1bit_evaluations++; - - // Multi-bit: use SignBitFactorsWithError and two-stage search - const SignBitFactorsWithError& full_factors = - *reinterpret_cast(base_ptr); - - // Compute 1-bit adjusted distance using shared helper - float dist_1bit = rabitq_utils::compute_1bit_adjusted_distance( - normalized_distance, - full_factors, - query_factors, - index->centered, - index->qb, - index->d); - - // Adaptive filtering: decide whether to compute full distance - const bool is_similarity = - index->metric_type == MetricType::METRIC_INNER_PRODUCT; - - float g_error = query_factors.g_error; - - bool should_refine = rabitq_utils::should_refine_candidate( - dist_1bit, - full_factors.f_error, - g_error, - heap_dis[0], - is_similarity); - if (should_refine) { - local_multibit_evaluations++; - - // Compute local_offset: position within current inverted list - size_t local_offset = this->j0 + b * 32 + j; - - // Compute full multi-bit distance - float dist_full = compute_full_multibit_distance( - result_id, local_q, q, local_offset); - - // Update heap if this distance is better - if (Cfloat::cmp(heap_dis[0], dist_full)) { - heap_replace_top( - k, heap_dis, heap_ids, dist_full, result_id); - nup++; - } - } - } else { - const auto& db_factors = - *reinterpret_cast(base_ptr); - - // Compute adjusted distance using shared helper - float adjusted_distance = - rabitq_utils::compute_1bit_adjusted_distance( - normalized_distance, - db_factors, - query_factors, - index->centered, - index->qb, - index->d); - - if (Cfloat::cmp(heap_dis[0], adjusted_distance)) { - heap_replace_top( - k, heap_dis, heap_ids, adjusted_distance, result_id); - nup++; - } - } - } - - // Update global stats atomically -#pragma omp atomic - rabitq_stats.n_1bit_evaluations += local_1bit_evaluations; -#pragma omp atomic - rabitq_stats.n_multibit_evaluations += local_multibit_evaluations; -} - -template -void simd_result_handlers::IVFRaBitQHeapHandler::set_list_context( - size_t list_no, - const std::vector& probe_map) { - current_list_no = list_no; - probe_indices = probe_map; - list_codes_ptr = index->invlists->get_codes(list_no); -} - -template -void simd_result_handlers::IVFRaBitQHeapHandler::begin( - const float* norms) { - this->normalizers = norms; -} - -template -void simd_result_handlers::IVFRaBitQHeapHandler::end() { -#pragma omp parallel for - for (int64_t q = 0; q < static_cast(nq); q++) { - float* heap_dis = heap_distances + q * k; - int64_t* heap_ids = heap_labels + q * k; - heap_reorder(k, heap_dis, heap_ids); - } -} - -template -float simd_result_handlers::IVFRaBitQHeapHandler:: - compute_full_multibit_distance( - size_t /*db_idx*/, - size_t local_q, - size_t global_q, - size_t local_offset) const { - const size_t ex_bits = index->rabitq.nb_bits - 1; - const size_t dim = index->d; - - const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr( - list_codes_ptr, - local_offset, - index->bbs, - packed_block_size, - full_block_size, - storage_size); - - const size_t ex_code_size = (dim * ex_bits + 7) / 8; - const uint8_t* ex_code = base_ptr + sizeof(SignBitFactorsWithError); - const ExtraBitsFactors& ex_fac = *reinterpret_cast( - base_ptr + sizeof(SignBitFactorsWithError) + ex_code_size); - - // Use local_q to access probe_indices (batch-local), global_q for storage - size_t probe_rank = probe_indices[local_q]; - size_t nprobe = context->nprobe > 0 ? context->nprobe : index->nprobe; - size_t storage_idx = global_q * nprobe + probe_rank; - const auto& query_factors = context->query_factors[storage_idx]; - - size_t list_no = current_list_no; - InvertedLists::ScopedCodes list_codes(index->invlists, list_no); - - std::vector unpacked_code(index->code_size); - packer->unpack_1(list_codes.get(), local_offset, unpacked_code.data()); - const uint8_t* sign_bits = unpacked_code.data(); - - return rabitq_utils::compute_full_multibit_distance( - sign_bits, - ex_code, - ex_fac, - query_factors.rotated_q.data(), - (index->metric_type == MetricType::METRIC_INNER_PRODUCT) - ? query_factors.q_dot_c - : query_factors.qr_to_c_L2sqr, - dim, - ex_bits, - index->metric_type); -} - /********************************************************* * IVFRaBitQFastScanScanner implementation *********************************************************/ @@ -905,17 +658,15 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner { std::vector curr_dists(k, distances[0]); std::vector curr_labels(k, labels[0]); - std::unique_ptr handler( - index.make_knn_handler( - !keep_max, - impl, - nq, - k, - curr_dists.data(), - curr_labels.data(), - sel, - context, - &normalizers[0])); + auto scanner = index.make_knn_scanner( + !keep_max, + nq, + k, + curr_dists.data(), + curr_labels.data(), + sel, + context); + auto* handler = scanner->handler(); int qmap1[1] = {0}; handler->q_map = qmap1; @@ -930,14 +681,13 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner { std::vector probe_map = {0}; handler->set_list_context(list_no, probe_map); - pq4_accumulate_loop( + scanner->accumulate_loop( 1, roundup(ntotal, index.bbs), index.bbs, static_cast(index.M2), codes, LUT, - *handler, 0, index.get_block_stride()); diff --git a/faiss/IndexIVFRaBitQFastScan.h b/faiss/IndexIVFRaBitQFastScan.h index fcbd5efa9..a326deb0b 100644 --- a/faiss/IndexIVFRaBitQFastScan.h +++ b/faiss/IndexIVFRaBitQFastScan.h @@ -148,6 +148,16 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan { const IVFSearchParameters* params = nullptr, IndexIVFStats* stats = nullptr) const override; + /// RaBitQ scanner via rabitq_ivf_make_knn_scanner + std::unique_ptr make_knn_scanner( + bool is_max, + idx_t n, + idx_t k, + float* distances, + idx_t* labels, + const IDSelector* sel, + const FastScanDistancePostProcessing& context = {}) const override; + /// Override to create RaBitQ-specific handlers SIMDResultHandlerToFloat* make_knn_handler( bool is_max, @@ -172,4 +182,226 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan { using IVFRaBitQHeapHandler = simd_result_handlers::IVFRaBitQHeapHandler; }; +/********************************************************* + * IVFRaBitQHeapHandler method bodies — defined here (after + * IndexIVFRaBitQFastScan is complete) to break the circular + * dependency between rabitq_result_handler.h and this header. + *********************************************************/ + +namespace simd_result_handlers { + +template +IVFRaBitQHeapHandler::IVFRaBitQHeapHandler( + const IndexIVFRaBitQFastScan* idx, + size_t nq_val, + size_t k_val, + float* distances, + int64_t* labels, + const FastScanDistancePostProcessing* ctx, + bool multibit) + : ResultHandlerCompare(nq_val, 0, nullptr), + index(idx), + heap_distances(distances), + heap_labels(labels), + nq(nq_val), + k(k_val), + context(ctx), + is_multibit(multibit), + storage_size(idx->compute_per_vector_storage_size()), + packed_block_size(((idx->M2 + 1) / 2) * idx->bbs), + full_block_size(idx->get_block_stride()), + packer(idx->get_CodePacker()) { + current_list_no = 0; + probe_indices.clear(); + for (int64_t q = 0; q < static_cast(nq); q++) { + heap_heapify(k, heap_distances + q * k, heap_labels + q * k); + } +} + +template +void IVFRaBitQHeapHandler::handle( + size_t q, + size_t b, + simd16uint16 d0, + simd16uint16 d1) { + size_t local_q = q; + this->adjust_with_origin(q, d0, d1); + + ALIGNED(32) uint16_t d32tab[32]; + d0.store(d32tab); + d1.store(d32tab + 16); + + float* const heap_dis = heap_distances + q * k; + int64_t* const heap_ids = heap_labels + q * k; + + FAISS_THROW_IF_NOT_FMT( + !probe_indices.empty() && local_q < probe_indices.size(), + "set_list_context() must be called before handle() - probe_indices size: %zu, local_q: %zu, global_q: %zu", + probe_indices.size(), + local_q, + q); + + if (!context || !context->query_factors) { + FAISS_THROW_MSG( + "Query factors not available: FastScanDistancePostProcessing with query_factors required"); + } + + size_t probe_rank = probe_indices[local_q]; + size_t nprobe = context->nprobe > 0 ? context->nprobe : index->nprobe; + size_t storage_idx = q * nprobe + probe_rank; + const auto& query_factors = context->query_factors[storage_idx]; + + const float one_a = + this->normalizers ? (1.0f / this->normalizers[2 * q]) : 1.0f; + const float bias = this->normalizers ? this->normalizers[2 * q + 1] : 0.0f; + + uint64_t idx_base = this->j0 + b * 32; + if (idx_base >= this->ntotal) { + return; + } + size_t max_positions = std::min(32, this->ntotal - idx_base); + + size_t local_1bit_evaluations = 0; + size_t local_multibit_evaluations = 0; + + for (size_t j = 0; j < max_positions; j++) { + const int64_t result_id = this->adjust_id(b, j); + if (result_id < 0) { + continue; + } + + const float normalized_distance = d32tab[j] * one_a + bias; + const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr( + list_codes_ptr, + idx_base + j, + index->bbs, + packed_block_size, + full_block_size, + storage_size); + + if (is_multibit) { + local_1bit_evaluations++; + const SignBitFactorsWithError& full_factors = + *reinterpret_cast(base_ptr); + + float dist_1bit = rabitq_utils::compute_1bit_adjusted_distance( + normalized_distance, + full_factors, + query_factors, + index->centered, + index->qb, + index->d); + + const bool is_similarity = + index->metric_type == MetricType::METRIC_INNER_PRODUCT; + bool should_refine = rabitq_utils::should_refine_candidate( + dist_1bit, + full_factors.f_error, + query_factors.g_error, + heap_dis[0], + is_similarity); + if (should_refine) { + local_multibit_evaluations++; + size_t local_offset = this->j0 + b * 32 + j; + float dist_full = compute_full_multibit_distance( + result_id, local_q, q, local_offset); + if (Cfloat::cmp(heap_dis[0], dist_full)) { + heap_replace_top( + k, heap_dis, heap_ids, dist_full, result_id); + nup++; + } + } + } else { + const auto& db_factors = + *reinterpret_cast(base_ptr); + float adjusted_distance = + rabitq_utils::compute_1bit_adjusted_distance( + normalized_distance, + db_factors, + query_factors, + index->centered, + index->qb, + index->d); + if (Cfloat::cmp(heap_dis[0], adjusted_distance)) { + heap_replace_top( + k, heap_dis, heap_ids, adjusted_distance, result_id); + nup++; + } + } + } + +#pragma omp atomic + rabitq_stats.n_1bit_evaluations += local_1bit_evaluations; +#pragma omp atomic + rabitq_stats.n_multibit_evaluations += local_multibit_evaluations; +} + +template +void IVFRaBitQHeapHandler::set_list_context( + size_t list_no, + const std::vector& probe_map) { + current_list_no = list_no; + probe_indices = probe_map; + list_codes_ptr = index->invlists->get_codes(list_no); +} + +template +void IVFRaBitQHeapHandler::begin(const float* norms) { + this->normalizers = norms; +} + +template +void IVFRaBitQHeapHandler::end() { +#pragma omp parallel for + for (int64_t q = 0; q < static_cast(nq); q++) { + heap_reorder(k, heap_distances + q * k, heap_labels + q * k); + } +} + +template +float IVFRaBitQHeapHandler::compute_full_multibit_distance( + size_t /*db_idx*/, + size_t local_q, + size_t global_q, + size_t local_offset) const { + const size_t ex_bits = index->rabitq.nb_bits - 1; + const size_t dim = index->d; + + const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr( + list_codes_ptr, + local_offset, + index->bbs, + packed_block_size, + full_block_size, + storage_size); + + const size_t ex_code_size = (dim * ex_bits + 7) / 8; + const uint8_t* ex_code = base_ptr + sizeof(SignBitFactorsWithError); + const ExtraBitsFactors& ex_fac = *reinterpret_cast( + base_ptr + sizeof(SignBitFactorsWithError) + ex_code_size); + + size_t probe_rank = probe_indices[local_q]; + size_t nprobe_val = context->nprobe > 0 ? context->nprobe : index->nprobe; + size_t storage_idx_val = global_q * nprobe_val + probe_rank; + const auto& query_factors = context->query_factors[storage_idx_val]; + + InvertedLists::ScopedCodes list_codes(index->invlists, current_list_no); + std::vector unpacked_code(index->code_size); + packer->unpack_1(list_codes.get(), local_offset, unpacked_code.data()); + + return rabitq_utils::compute_full_multibit_distance( + unpacked_code.data(), + ex_code, + ex_fac, + query_factors.rotated_q.data(), + (index->metric_type == MetricType::METRIC_INNER_PRODUCT) + ? query_factors.q_dot_c + : query_factors.qr_to_c_L2sqr, + dim, + ex_bits, + index->metric_type); +} + +} // namespace simd_result_handlers + } // namespace faiss diff --git a/faiss/IndexRaBitQFastScan.cpp b/faiss/IndexRaBitQFastScan.cpp index 0a613ffb8..0c54cfe64 100644 --- a/faiss/IndexRaBitQFastScan.cpp +++ b/faiss/IndexRaBitQFastScan.cpp @@ -522,241 +522,18 @@ void IndexRaBitQFastScan::search( } } -// Template implementations for RaBitQHeapHandler -template -RaBitQHeapHandler::RaBitQHeapHandler( - const IndexRaBitQFastScan* index, - size_t nq_val, - size_t k_val, +std::unique_ptr IndexRaBitQFastScan::make_knn_scanner( + bool is_max, + idx_t n, + idx_t k, + size_t /*ntotal*/, float* distances, - int64_t* labels, - const IDSelector* sel_in, - const FastScanDistancePostProcessing& ctx, - bool multi_bit) - : RHC(nq_val, index->ntotal, sel_in), - rabitq_index(index), - heap_distances(distances), - heap_labels(labels), - nq(nq_val), - k(k_val), - context(ctx), - is_multi_bit(multi_bit), - storage_size(index->compute_per_vector_storage_size()), - packed_block_size(((index->M2 + 1) / 2) * index->bbs), - full_block_size(index->get_block_stride()), - packer(index->get_CodePacker()) { - // Initialize heaps for all queries in constructor - // This allows us to support direct normalizer assignment -#pragma omp parallel for if (nq > 100) - for (int64_t q = 0; q < static_cast(nq); q++) { - float* heap_dis = heap_distances + q * k; - int64_t* heap_ids = heap_labels + q * k; - heap_heapify(k, heap_dis, heap_ids); - } -} - -template -void RaBitQHeapHandler::handle( - size_t q, - size_t b, - simd16uint16 d0, - simd16uint16 d1) { - ALIGNED(32) uint16_t d32tab[32]; - d0.store(d32tab); - d1.store(d32tab + 16); - - // Get heap pointers and query factors (computed once per batch) - float* const heap_dis = heap_distances + q * k; - int64_t* const heap_ids = heap_labels + q * k; - - // Access query factors from query_factors pointer - rabitq_utils::QueryFactorsData query_factors_data = {}; - if (context.query_factors != nullptr) { - query_factors_data = context.query_factors[q]; - } - - // Compute normalizers once per batch - const float one_a = normalizers ? (1.0f / normalizers[2 * q]) : 1.0f; - const float bias = normalizers ? normalizers[2 * q + 1] : 0.0f; - - // Compute loop bounds to avoid redundant bounds checking - const size_t base_db_idx = this->j0 + b * 32; - const size_t max_vectors = (base_db_idx < rabitq_index->ntotal) - ? std::min(32, rabitq_index->ntotal - base_db_idx) - : 0; - - // Compute block auxiliary region base pointer once per batch. - // Since bbs=32, each batch of 32 vectors aligns to one block. - const size_t block_idx = base_db_idx / rabitq_index->bbs; - const uint8_t* aux_base = rabitq_index->codes.get() + - block_idx * full_block_size + packed_block_size; - - // Stats tracking for multi-bit two-stage search only - // n_1bit_evaluations: candidates evaluated using 1-bit lower bound - // n_multibit_evaluations: candidates requiring full multi-bit distance - size_t local_1bit_evaluations = 0; - size_t local_multibit_evaluations = 0; - - // Process distances in batch - for (size_t i = 0; i < max_vectors; i++) { - const size_t db_idx = base_db_idx + i; - - // Normalize distance from LUT lookup - const float normalized_distance = d32tab[i] * one_a + bias; - - // Access factors from block auxiliary region - const uint8_t* base_ptr = aux_base + i * storage_size; - - if (is_multi_bit) { - // Track candidates actually considered for two-stage filtering - local_1bit_evaluations++; - - const SignBitFactorsWithError& full_factors = - *reinterpret_cast(base_ptr); - - float dist_1bit = rabitq_utils::compute_1bit_adjusted_distance( - normalized_distance, - full_factors, - query_factors_data, - rabitq_index->centered, - rabitq_index->qb, - rabitq_index->d); - - // Adaptive filtering: decide whether to compute full distance - const bool is_similarity = rabitq_index->metric_type == - MetricType::METRIC_INNER_PRODUCT; - bool should_refine = rabitq_utils::should_refine_candidate( - dist_1bit, - full_factors.f_error, - context.query_factors ? context.query_factors[q].g_error - : 0.0f, - heap_dis[0], - is_similarity); - - if (should_refine) { - local_multibit_evaluations++; - float dist_full = compute_full_multibit_distance(db_idx, q); - - if (Cfloat::cmp(heap_dis[0], dist_full)) { - heap_replace_top( - k, heap_dis, heap_ids, dist_full, db_idx); - } - } - } else { - const rabitq_utils::SignBitFactors& db_factors = - *reinterpret_cast( - base_ptr); - - float adjusted_distance = - rabitq_utils::compute_1bit_adjusted_distance( - normalized_distance, - db_factors, - query_factors_data, - rabitq_index->centered, - rabitq_index->qb, - rabitq_index->d); - - // Add to heap if better than current worst - if (Cfloat::cmp(heap_dis[0], adjusted_distance)) { - heap_replace_top( - k, heap_dis, heap_ids, adjusted_distance, db_idx); - } - } - } - - // Update global stats atomically -#pragma omp atomic - rabitq_stats.n_1bit_evaluations += local_1bit_evaluations; -#pragma omp atomic - rabitq_stats.n_multibit_evaluations += local_multibit_evaluations; -} - -template -void RaBitQHeapHandler::begin(const float* norms) { - normalizers = norms; - // Heap initialization is now done in constructor -} - -template -void RaBitQHeapHandler::end() { -// Reorder final results -#pragma omp parallel for if (nq > 100) - for (int64_t q = 0; q < static_cast(nq); q++) { - float* heap_dis = heap_distances + q * k; - int64_t* heap_ids = heap_labels + q * k; - heap_reorder(k, heap_dis, heap_ids); - } -} - -template -float RaBitQHeapHandler::compute_lower_bound( - float dist_1bit, - size_t db_idx, - size_t q) const { - // Access f_error from block auxiliary region - const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr( - rabitq_index->codes.get(), - db_idx, - rabitq_index->bbs, - packed_block_size, - full_block_size, - storage_size); - const SignBitFactorsWithError& db_factors = - *reinterpret_cast(base_ptr); - float f_error = db_factors.f_error; - - // Get g_error from query factors (query-dependent error term) - float g_error = 0.0f; - if (context.query_factors != nullptr) { - g_error = context.query_factors[q].g_error; - } - - // Compute error adjustment: f_error * g_error - float error_adjustment = f_error * g_error; - - return dist_1bit - error_adjustment; -} - -template -float RaBitQHeapHandler::compute_full_multibit_distance( - size_t db_idx, - size_t q) const { - const size_t ex_bits = rabitq_index->rabitq.nb_bits - 1; - const size_t dim = rabitq_index->d; - - const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr( - rabitq_index->codes.get(), - db_idx, - rabitq_index->bbs, - packed_block_size, - full_block_size, - storage_size); - - const size_t ex_code_size = (dim * ex_bits + 7) / 8; - const uint8_t* ex_code = base_ptr + sizeof(SignBitFactorsWithError); - const ExtraBitsFactors& ex_fac = *reinterpret_cast( - base_ptr + sizeof(SignBitFactorsWithError) + ex_code_size); - - // Get query factors reference (avoid copying) - const rabitq_utils::QueryFactorsData& query_factors = - context.query_factors[q]; - - // Get sign bits from FastScan packed format - std::vector unpacked_code(rabitq_index->code_size); - packer->unpack_1(rabitq_index->codes.get(), db_idx, unpacked_code.data()); - const uint8_t* sign_bits = unpacked_code.data(); - - return rabitq_utils::compute_full_multibit_distance( - sign_bits, - ex_code, - ex_fac, - query_factors.rotated_q.data(), - (rabitq_index->metric_type == MetricType::METRIC_INNER_PRODUCT) - ? query_factors.q_dot_c - : query_factors.qr_to_c_L2sqr, - dim, - ex_bits, - rabitq_index->metric_type); + idx_t* labels, + const IDSelector* sel, + const FastScanDistancePostProcessing& context) const { + const bool is_multi_bit = rabitq.nb_bits > 1; + return rabitq_make_knn_scanner( + this, is_max, n, k, distances, labels, sel, context, is_multi_bit); } // Implementation of virtual make_knn_handler method @@ -770,15 +547,14 @@ SIMDResultHandlerToFloat* IndexRaBitQFastScan::make_knn_handler( idx_t* labels, const IDSelector* sel, const FastScanDistancePostProcessing& context) const { - // Use runtime boolean for multi-bit mode const bool multi_bit = rabitq.nb_bits > 1; if (is_max) { return new RaBitQHeapHandler, false>( - this, n, k, distances, labels, sel, context, multi_bit); + this, n, k, distances, labels, sel, &context, multi_bit); } else { return new RaBitQHeapHandler, false>( - this, n, k, distances, labels, sel, context, multi_bit); + this, n, k, distances, labels, sel, &context, multi_bit); } } diff --git a/faiss/IndexRaBitQFastScan.h b/faiss/IndexRaBitQFastScan.h index a3e32a064..966f1e8c7 100644 --- a/faiss/IndexRaBitQFastScan.h +++ b/faiss/IndexRaBitQFastScan.h @@ -92,6 +92,17 @@ struct IndexRaBitQFastScan : IndexFastScan { idx_t* labels, const SearchParameters* params = nullptr) const override; + /// RaBitQ scanner wired through rabitq_make_knn_scanner + std::unique_ptr make_knn_scanner( + bool is_max, + idx_t n, + idx_t k, + size_t ntotal, + float* distances, + idx_t* labels, + const IDSelector* sel, + const FastScanDistancePostProcessing& context = {}) const override; + /// Override to create RaBitQ-specific handlers SIMDResultHandlerToFloat* make_knn_handler( bool is_max, @@ -122,18 +133,24 @@ struct IndexRaBitQFastScan : IndexFastScan { * * @tparam C Comparator type (CMin/CMax) for heap operations * @tparam with_id_map Whether to use id mapping (similar to HeapHandler) + * @tparam SL SIMD level for per-TU template instantiation */ -template +template < + class C, + bool with_id_map = false, + SIMDLevel SL = SINGLE_SIMD_LEVEL_256> struct RaBitQHeapHandler - : simd_result_handlers::ResultHandlerCompare { - using RHC = simd_result_handlers::ResultHandlerCompare; + : simd_result_handlers::ResultHandlerCompare { + using RHC = simd_result_handlers::ResultHandlerCompare; using RHC::normalizers; + static constexpr SIMDLevel SL256 = simd256_level_selector::value; + using simd16uint16 = simd16uint16_tpl; const IndexRaBitQFastScan* rabitq_index; float* heap_distances; // [nq * k] int64_t* heap_labels; // [nq * k] const size_t nq, k; - const FastScanDistancePostProcessing& + const FastScanDistancePostProcessing* context; // Processing context with query offset const bool is_multi_bit; // Runtime flag for multi-bit mode @@ -156,22 +173,176 @@ struct RaBitQHeapHandler float* distances, int64_t* labels, const IDSelector* sel_in, - const FastScanDistancePostProcessing& context, - bool multi_bit); + const FastScanDistancePostProcessing* ctx, + bool multi_bit) + : RHC(nq_val, index->ntotal, sel_in), + rabitq_index(index), + heap_distances(distances), + heap_labels(labels), + nq(nq_val), + k(k_val), + context(ctx), + is_multi_bit(multi_bit), + storage_size(index->compute_per_vector_storage_size()), + packed_block_size(((index->M2 + 1) / 2) * index->bbs), + full_block_size(index->get_block_stride()), + packer(index->get_CodePacker()) { +#pragma omp parallel for if (nq > 100) + for (int64_t q = 0; q < static_cast(nq); q++) { + float* heap_dis = heap_distances + q * k; + int64_t* heap_ids = heap_labels + q * k; + heap_heapify(k, heap_dis, heap_ids); + } + } - void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) override; + void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) override { + ALIGNED(32) uint16_t d32tab[32]; + d0.store(d32tab); + d1.store(d32tab + 16); - void begin(const float* norms); + float* const heap_dis = heap_distances + q * k; + int64_t* const heap_ids = heap_labels + q * k; - void end(); + rabitq_utils::QueryFactorsData query_factors_data = {}; + if (context && context->query_factors != nullptr) { + query_factors_data = context->query_factors[q]; + } + + const float one_a = normalizers ? (1.0f / normalizers[2 * q]) : 1.0f; + const float bias = normalizers ? normalizers[2 * q + 1] : 0.0f; + + const size_t base_db_idx = this->j0 + b * 32; + const size_t max_vectors = (base_db_idx < rabitq_index->ntotal) + ? std::min(32, rabitq_index->ntotal - base_db_idx) + : 0; + + const size_t block_idx = base_db_idx / rabitq_index->bbs; + const uint8_t* aux_base = rabitq_index->codes.get() + + block_idx * full_block_size + packed_block_size; + + size_t local_1bit_evaluations = 0; + size_t local_multibit_evaluations = 0; + + for (size_t i = 0; i < max_vectors; i++) { + const size_t db_idx = base_db_idx + i; + const float normalized_distance = d32tab[i] * one_a + bias; + const uint8_t* base_ptr = aux_base + i * storage_size; + + if (is_multi_bit) { + local_1bit_evaluations++; + + const SignBitFactorsWithError& full_factors = + *reinterpret_cast( + base_ptr); + + float dist_1bit = rabitq_utils::compute_1bit_adjusted_distance( + normalized_distance, + full_factors, + query_factors_data, + rabitq_index->centered, + rabitq_index->qb, + rabitq_index->d); + + const bool is_similarity = rabitq_index->metric_type == + MetricType::METRIC_INNER_PRODUCT; + bool should_refine = rabitq_utils::should_refine_candidate( + dist_1bit, + full_factors.f_error, + context && context->query_factors + ? context->query_factors[q].g_error + : 0.0f, + heap_dis[0], + is_similarity); + + if (should_refine) { + local_multibit_evaluations++; + float dist_full = compute_full_multibit_distance(db_idx, q); + + if (Cfloat::cmp(heap_dis[0], dist_full)) { + heap_replace_top( + k, heap_dis, heap_ids, dist_full, db_idx); + } + } + } else { + const rabitq_utils::SignBitFactors& db_factors = + *reinterpret_cast( + base_ptr); + + float adjusted_distance = + rabitq_utils::compute_1bit_adjusted_distance( + normalized_distance, + db_factors, + query_factors_data, + rabitq_index->centered, + rabitq_index->qb, + rabitq_index->d); + + if (Cfloat::cmp(heap_dis[0], adjusted_distance)) { + heap_replace_top( + k, heap_dis, heap_ids, adjusted_distance, db_idx); + } + } + } + +#pragma omp atomic + rabitq_stats.n_1bit_evaluations += local_1bit_evaluations; +#pragma omp atomic + rabitq_stats.n_multibit_evaluations += local_multibit_evaluations; + } + + void begin(const float* norms) { + normalizers = norms; + } + + void end() { +#pragma omp parallel for if (nq > 100) + for (int64_t q = 0; q < static_cast(nq); q++) { + float* heap_dis = heap_distances + q * k; + int64_t* heap_ids = heap_labels + q * k; + heap_reorder(k, heap_dis, heap_ids); + } + } private: - /// Compute full multi-bit distance for a candidate vector (multi-bit only) - float compute_full_multibit_distance(size_t db_idx, size_t q) const; + float compute_full_multibit_distance(size_t db_idx, size_t q) const { + const size_t ex_bits = rabitq_index->rabitq.nb_bits - 1; + const size_t dim = rabitq_index->d; - /// Compute lower bound using 1-bit distance and error bound (multi-bit - /// only) - float compute_lower_bound(float dist_1bit, size_t db_idx, size_t q) const; + const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr( + rabitq_index->codes.get(), + db_idx, + rabitq_index->bbs, + packed_block_size, + full_block_size, + storage_size); + + const size_t ex_code_size = (dim * ex_bits + 7) / 8; + const uint8_t* ex_code = base_ptr + sizeof(SignBitFactorsWithError); + const ExtraBitsFactors& ex_fac = + *reinterpret_cast( + base_ptr + sizeof(SignBitFactorsWithError) + + ex_code_size); + + const rabitq_utils::QueryFactorsData& query_factors = + context->query_factors[q]; + + std::vector unpacked_code(rabitq_index->code_size); + packer->unpack_1( + rabitq_index->codes.get(), db_idx, unpacked_code.data()); + const uint8_t* sign_bits = unpacked_code.data(); + + return rabitq_utils::compute_full_multibit_distance( + sign_bits, + ex_code, + ex_fac, + query_factors.rotated_q.data(), + (rabitq_index->metric_type == MetricType::METRIC_INNER_PRODUCT) + ? query_factors.q_dot_c + : query_factors.qr_to_c_L2sqr, + dim, + ex_bits, + rabitq_index->metric_type); + } }; } // namespace faiss diff --git a/faiss/impl/fast_scan/fast_scan.cpp b/faiss/impl/fast_scan/fast_scan.cpp index 342b8a2d6..d9eb1c84b 100644 --- a/faiss/impl/fast_scan/fast_scan.cpp +++ b/faiss/impl/fast_scan/fast_scan.cpp @@ -362,7 +362,8 @@ int pq4_pack_LUT_qbs_q_map( ***************************************************************/ #define THE_LEVEL_TO_DISPATCH SIMDLevel::NONE -#include // IWYU pragma: keep +#include // IWYU pragma: keep +#include // IWYU pragma: keep #undef THE_LEVEL_TO_DISPATCH namespace faiss { @@ -390,4 +391,48 @@ std::unique_ptr make_fast_scan_knn_scanner( with_id_map); } +std::unique_ptr rabitq_make_knn_scanner( + const IndexRaBitQFastScan* index, + bool is_max, + size_t nq, + int64_t k, + float* distances, + int64_t* ids, + const IDSelector* sel, + const FastScanDistancePostProcessing& context, + bool is_multi_bit) { + DISPATCH_SIMDLevel( + rabitq_make_knn_scanner_impl, + index, + is_max, + nq, + k, + distances, + ids, + sel, + context, + is_multi_bit); +} + +std::unique_ptr rabitq_ivf_make_knn_scanner( + bool is_max, + const IndexIVFRaBitQFastScan* index, + size_t nq, + size_t k, + float* distances, + int64_t* ids, + const FastScanDistancePostProcessing* context, + bool multi_bit) { + DISPATCH_SIMDLevel( + rabitq_ivf_make_knn_scanner_impl, + is_max, + index, + nq, + k, + distances, + ids, + context, + multi_bit); +} + } // namespace faiss diff --git a/faiss/impl/fast_scan/fast_scan.h b/faiss/impl/fast_scan/fast_scan.h index 972a1ac75..d20bcdd10 100644 --- a/faiss/impl/fast_scan/fast_scan.h +++ b/faiss/impl/fast_scan/fast_scan.h @@ -26,9 +26,9 @@ namespace faiss { +struct IDSelector; struct SIMDResultHandler; struct SIMDResultHandlerToFloat; -struct IDSelector; /** Pack codes for consumption by the SIMD kernels. * The unused bytes are set to 0. @@ -288,4 +288,61 @@ std::unique_ptr make_fast_scan_knn_scanner( const IDSelector* sel, bool with_id_map = false); +/*************************************************************** + * RaBitQ scanner factory: per-SIMD specializations live in + * rabitq_dispatching.h, included by each per-SIMD TU. + ***************************************************************/ + +struct IndexRaBitQFastScan; +struct IndexIVFRaBitQFastScan; +struct FastScanDistancePostProcessing; + +/// Per-SIMD factory (primary template; specializations in rabitq_dispatching.h) +template +std::unique_ptr rabitq_make_knn_scanner_impl( + const IndexRaBitQFastScan* index, + bool is_max, + size_t nq, + int64_t k, + float* distances, + int64_t* ids, + const IDSelector* sel, + const FastScanDistancePostProcessing& context, + bool is_multi_bit); + +/// Runtime dispatch wrapper for rabitq_make_knn_scanner_impl +std::unique_ptr rabitq_make_knn_scanner( + const IndexRaBitQFastScan* index, + bool is_max, + size_t nq, + int64_t k, + float* distances, + int64_t* ids, + const IDSelector* sel, + const FastScanDistancePostProcessing& context, + bool is_multi_bit); + +/// Per-SIMD IVF RaBitQ scanner factory. +template +std::unique_ptr rabitq_ivf_make_knn_scanner_impl( + bool is_max, + const IndexIVFRaBitQFastScan* index, + size_t nq, + size_t k, + float* distances, + int64_t* ids, + const FastScanDistancePostProcessing* context, + bool multi_bit); + +/// Runtime dispatch wrapper for IVF RaBitQ scanner. +std::unique_ptr rabitq_ivf_make_knn_scanner( + bool is_max, + const IndexIVFRaBitQFastScan* index, + size_t nq, + size_t k, + float* distances, + int64_t* ids, + const FastScanDistancePostProcessing* context, + bool multi_bit); + } // namespace faiss diff --git a/faiss/impl/fast_scan/impl-avx2.cpp b/faiss/impl/fast_scan/impl-avx2.cpp index b3f619d95..3dcac04f4 100644 --- a/faiss/impl/fast_scan/impl-avx2.cpp +++ b/faiss/impl/fast_scan/impl-avx2.cpp @@ -8,6 +8,7 @@ #ifdef COMPILE_SIMD_AVX2 #define THE_LEVEL_TO_DISPATCH SIMDLevel::AVX2 -#include // IWYU pragma: keep +#include // IWYU pragma: keep +#include // IWYU pragma: keep #endif // COMPILE_SIMD_AVX2 diff --git a/faiss/impl/fast_scan/impl-avx512.cpp b/faiss/impl/fast_scan/impl-avx512.cpp index 7f46c017d..2fac7920d 100644 --- a/faiss/impl/fast_scan/impl-avx512.cpp +++ b/faiss/impl/fast_scan/impl-avx512.cpp @@ -8,6 +8,7 @@ #ifdef COMPILE_SIMD_AVX512 #define THE_LEVEL_TO_DISPATCH SIMDLevel::AVX512 -#include // IWYU pragma: keep +#include // IWYU pragma: keep +#include // IWYU pragma: keep #endif // COMPILE_SIMD_AVX512 diff --git a/faiss/impl/fast_scan/impl-neon.cpp b/faiss/impl/fast_scan/impl-neon.cpp index 062034690..cf40e3967 100644 --- a/faiss/impl/fast_scan/impl-neon.cpp +++ b/faiss/impl/fast_scan/impl-neon.cpp @@ -8,7 +8,8 @@ #ifdef COMPILE_SIMD_ARM_NEON #define THE_LEVEL_TO_DISPATCH SIMDLevel::ARM_NEON -#include // IWYU pragma: keep +#include // IWYU pragma: keep +#include // IWYU pragma: keep // ARM_SVE: forward to ARM_NEON implementation until a dedicated SVE // specialization is written (same pattern as scalar_quantizer/sq-neon.cpp). @@ -32,6 +33,37 @@ std::unique_ptr make_fast_scan_scanner_impl< is_max, impl, nq, ntotal, k, distances, ids, sel, with_id_map); } +template <> +std::unique_ptr rabitq_make_knn_scanner_impl< + SIMDLevel::ARM_SVE>( + const IndexRaBitQFastScan* index, + bool is_max, + size_t nq, + int64_t k, + float* distances, + int64_t* ids, + const IDSelector* sel, + const FastScanDistancePostProcessing& context, + bool is_multi_bit) { + return rabitq_make_knn_scanner_impl( + index, is_max, nq, k, distances, ids, sel, context, is_multi_bit); +} + +template <> +std::unique_ptr rabitq_ivf_make_knn_scanner_impl< + SIMDLevel::ARM_SVE>( + bool is_max, + const IndexIVFRaBitQFastScan* index, + size_t nq, + size_t k, + float* distances, + int64_t* ids, + const FastScanDistancePostProcessing* context, + bool multi_bit) { + return rabitq_ivf_make_knn_scanner_impl( + is_max, index, nq, k, distances, ids, context, multi_bit); +} + } // namespace faiss #endif // COMPILE_SIMD_ARM_SVE diff --git a/faiss/impl/fast_scan/rabitq_dispatching.h b/faiss/impl/fast_scan/rabitq_dispatching.h new file mode 100644 index 000000000..4513c443f --- /dev/null +++ b/faiss/impl/fast_scan/rabitq_dispatching.h @@ -0,0 +1,81 @@ +/* + * 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. + */ + +#pragma once + +/** + * @file rabitq_dispatching.h + * @brief Per-SIMD TU dispatch for RaBitQ flat scanner. + * + * Included after dispatching.h in each per-SIMD TU, so that + * ScannerMixIn from dispatching.h is visible (same TU). + * + * Provides the rabitq_make_knn_scanner_impl specialization + * that wraps RaBitQHeapHandler in ScannerMixIn. + */ + +#ifndef THE_LEVEL_TO_DISPATCH +#error "Define THE_LEVEL_TO_DISPATCH before including this header" +#endif + +#include +#include +#include + +// ScannerMixIn is visible from dispatching.h (same TU) + +namespace faiss { + +template <> +std::unique_ptr rabitq_make_knn_scanner_impl< + THE_LEVEL_TO_DISPATCH>( + const IndexRaBitQFastScan* index, + bool is_max, + size_t nq, + int64_t k, + float* distances, + int64_t* ids, + const IDSelector* sel, + const FastScanDistancePostProcessing& context, + bool is_multi_bit) { + if (is_max) { + using H = RaBitQHeapHandler, false>; + return std::make_unique>( + index, nq, k, distances, ids, sel, &context, is_multi_bit); + } else { + using H = RaBitQHeapHandler, false>; + return std::make_unique>( + index, nq, k, distances, ids, sel, &context, is_multi_bit); + } +} + +// IVF RaBitQ scanner factory +template <> +std::unique_ptr rabitq_ivf_make_knn_scanner_impl< + THE_LEVEL_TO_DISPATCH>( + bool is_max, + const IndexIVFRaBitQFastScan* index, + size_t nq, + size_t k, + float* distances, + int64_t* ids, + const FastScanDistancePostProcessing* context, + bool multi_bit) { + if (is_max) { + using C = CMax; + using H = simd_result_handlers::IVFRaBitQHeapHandler; + return std::make_unique>( + index, nq, k, distances, ids, context, multi_bit); + } else { + using C = CMin; + using H = simd_result_handlers::IVFRaBitQHeapHandler; + return std::make_unique>( + index, nq, k, distances, ids, context, multi_bit); + } +} + +} // namespace faiss diff --git a/faiss/impl/fast_scan/rabitq_result_handler.h b/faiss/impl/fast_scan/rabitq_result_handler.h index 4fe4d44de..2427c8a5f 100644 --- a/faiss/impl/fast_scan/rabitq_result_handler.h +++ b/faiss/impl/fast_scan/rabitq_result_handler.h @@ -11,33 +11,29 @@ #include #include +#include +#include +#include #include #include +#include #include namespace faiss { -// Forward declaration — full definition needed only in implementation +// Forward declaration — full definition in IndexIVFRaBitQFastScan.h struct IndexIVFRaBitQFastScan; namespace simd_result_handlers { +// Import shared utilities from RaBitQUtils +using rabitq_utils::ExtraBitsFactors; +using rabitq_utils::SignBitFactors; +using rabitq_utils::SignBitFactorsWithError; + /** SIMD result handler for IndexIVFRaBitQFastScan that applies * RaBitQ-specific distance corrections during batch processing. * - * This handler processes batches of 32 distance computations from SIMD - * kernels, applies RaBitQ distance formula adjustments (factors and - * normalizers), and immediately updates result heaps. This eliminates the - * need for post-processing and provides significant performance benefits. - * - * Key optimizations: - * - Direct heap integration with no intermediate result storage - * - Batch-level computation of normalizers and query factors - * - Specialized handling for both centered and non-centered quantization - * modes - * - Efficient inner product metric corrections - * - Uses runtime boolean for multi-bit mode - * * @tparam C Comparator type (CMin/CMax) for heap operations * @tparam SL SIMD level for dynamic dispatch */ @@ -71,6 +67,9 @@ struct IVFRaBitQHeapHandler : ResultHandlerCompare { CMax, CMin>::type; + // Constructor and method bodies are defined at the bottom of + // IndexIVFRaBitQFastScan.h (after the full struct definition is + // available) to break the circular header dependency. IVFRaBitQHeapHandler( const IndexIVFRaBitQFastScan* idx, size_t nq_val, @@ -82,7 +81,6 @@ struct IVFRaBitQHeapHandler : ResultHandlerCompare { void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) override; - /// Override base class virtual method to receive context information void set_list_context(size_t list_no, const std::vector& probe_map) override; @@ -95,14 +93,8 @@ struct IVFRaBitQHeapHandler : ResultHandlerCompare { } private: - /// Compute full multi-bit distance for a candidate vector (multi-bit - /// only) - /// @param db_idx Global database vector index - /// @param local_q Batch-local query index (for probe_indices access) - /// @param global_q Global query index (for storage indexing) - /// @param local_offset Offset within the current inverted list float compute_full_multibit_distance( - size_t /*db_idx*/, + size_t db_idx, size_t local_q, size_t global_q, size_t local_offset) const; diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig index e58000e78..3547a6d19 100644 --- a/faiss/python/swigfaiss.swig +++ b/faiss/python/swigfaiss.swig @@ -599,11 +599,23 @@ void gpu_sync_all_devices() %include -// Ignore make_knn_handler methods — they return SIMDResultHandlerToFloat* -// which is an internal SIMD type not exposed to Python. +// Ignore FastScanCodeScanner and related factory functions — internal SIMD types. +%ignore faiss::FastScanCodeScanner; +%ignore faiss::make_fast_scan_scanner_impl; +%ignore faiss::make_fast_scan_knn_scanner; +%ignore faiss::rabitq_make_knn_scanner_impl; +%ignore faiss::rabitq_make_knn_scanner; +%ignore faiss::rabitq_ivf_make_knn_scanner_impl; +%ignore faiss::rabitq_ivf_make_knn_scanner; + +// Ignore make_knn_scanner / make_knn_handler — they return internal SIMD types. +%ignore faiss::IndexFastScan::make_knn_scanner; %ignore faiss::IndexFastScan::make_knn_handler; +%ignore faiss::IndexIVFFastScan::make_knn_scanner; %ignore faiss::IndexIVFFastScan::make_knn_handler; +%ignore faiss::IndexRaBitQFastScan::make_knn_scanner; %ignore faiss::IndexRaBitQFastScan::make_knn_handler; +%ignore faiss::IndexIVFRaBitQFastScan::make_knn_scanner; %ignore faiss::IndexIVFRaBitQFastScan::make_knn_handler; %include