From 9e72cbf4cd711b520270ce033b5bda782fc86d3e Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 28 Jul 2026 14:49:22 -0700 Subject: [PATCH 01/10] Add filtered search support to IVF-RaBitQ Support bitset filters across all IVF-RaBitQ search modes and kernel paths. Use infinite distances as the sole invalid-result sentinel, preserving all finite vector IDs, including UINT32_MAX, and map invalid results to the public out-of-bounds record. Add L2 and inner-product coverage across search modes, bit widths, moderate and 1% pass rates, and under-filled result sets. Validate bitset coverage and skip distance computation for filtered samples. --- cpp/CMakeLists.txt | 14 ++ .../ivf_rabitq/ivf_rabitq_fragments.hpp | 3 + cpp/include/cuvs/neighbors/ivf_rabitq.hpp | 14 +- cpp/src/neighbors/ivf_rabitq.cu | 77 ++++++-- .../neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu | 27 ++- .../ivf_rabitq/gpu_index/ivf_gpu.cuh | 23 ++- .../ivf_rabitq/gpu_index/searcher_gpu.cu | 38 ++-- .../ivf_rabitq/gpu_index/searcher_gpu.cuh | 20 +- .../gpu_index/searcher_gpu_common.cuh | 3 + .../gpu_index/searcher_gpu_quantize_query.cu | 35 ++-- .../gpu_index/searcher_gpu_shared_mem_opt.cu | 39 ++-- .../bitwise_emit_distances_kernel.cu.in | 27 ++- ..._products_with_bitwise_block_sort_impl.cuh | 3 +- ...oducts_with_bitwise_block_sort_planner.hpp | 5 + ...te_inner_products_with_bitwise_planner.hpp | 5 + ...roducts_with_lut16_opt_block_sort_impl.cuh | 3 +- ...ucts_with_lut16_opt_block_sort_planner.hpp | 5 + ..._inner_products_with_lut16_opt_planner.hpp | 5 + ...r_products_with_lut_block_sort_planner.hpp | 5 + ...ompute_inner_products_with_lut_planner.hpp | 5 + .../jit_lto_kernels/device_functions.cuh | 12 ++ .../jit_lto_kernels/launcher_factory.hpp | 41 ++++- .../lut16_opt_emit_distances_kernel.cu.in | 26 ++- .../lut_block_sort_emit_topk_kernel.cu.in | 6 +- .../lut_emit_distances_kernel.cu.in | 26 ++- .../sample_filter_kernel.cu.in | 45 +++++ .../jit_lto_kernels/sample_filter_matrix.json | 4 + cpp/tests/neighbors/ann_ivf_rabitq.cuh | 174 +++++++++++++++++- .../ann_ivf_rabitq/test_float_int64_t.cu | 7 +- 29 files changed, 569 insertions(+), 128 deletions(-) create mode 100644 cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/sample_filter_kernel.cu.in create mode 100644 cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/sample_filter_matrix.json diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index abbca74fe7..f423ee8132 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -683,6 +683,20 @@ if(NOT BUILD_CPU_ONLY) KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements ) set(ivf_rabitq_ns "cuvs::neighbors::ivf_rabitq::detail") + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT "ivf_rabitq_sample_filter_@filter_name@" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/sample_filter_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/sample_filter_kernel.cu.in" + FRAGMENT_TAG_FORMAT + "${ivf_rabitq_ns}::fragment_tag_sample_filter<${neighbors_ns}::tag_filter_@filter_name@>" + FRAGMENT_TAG_HEADER_FILES "" + "" + OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/sample_filter" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + ) generate_jit_lto_kernels( jit_lto_files NAME_FORMAT "ivf_rabitq_compute_inner_products_with_lut" diff --git a/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp index 5efb2b5051..24547eacf5 100644 --- a/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp @@ -48,4 +48,7 @@ struct fragment_tag_compute_lut_ip_for_vec {}; template struct fragment_tag_compute_bitwise_quantized_ip_for_vec {}; +template +struct fragment_tag_sample_filter {}; + } // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/include/cuvs/neighbors/ivf_rabitq.hpp b/cpp/include/cuvs/neighbors/ivf_rabitq.hpp index 659b95fcc3..c207548365 100644 --- a/cpp/include/cuvs/neighbors/ivf_rabitq.hpp +++ b/cpp/include/cuvs/neighbors/ivf_rabitq.hpp @@ -18,6 +18,7 @@ #include #include +#include #include #include #include @@ -100,6 +101,13 @@ enum class search_mode { QUANT8 = 3, }; +/** + * Default value returned by `search` when fewer than k samples pass the filter or the combined + * size of the probed clusters is smaller than k. + */ +template +constexpr static IdxT kOutOfBoundsRecord = std::numeric_limits::max(); + struct search_params : cuvs::neighbors::search_params { /** The number of clusters to search. */ uint32_t n_probes = 20; @@ -240,13 +248,17 @@ auto build(raft::resources const& handle, * [n_queries, k] * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. Supports `none_sample_filter` and `bitset_filter`. */ void search(raft::resources const& handle, const cuvs::neighbors::ivf_rabitq::search_params& search_params, cuvs::neighbors::ivf_rabitq::index& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, - raft::device_matrix_view distances); + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); /** * @} diff --git a/cpp/src/neighbors/ivf_rabitq.cu b/cpp/src/neighbors/ivf_rabitq.cu index 964349b280..174d4600da 100644 --- a/cpp/src/neighbors/ivf_rabitq.cu +++ b/cpp/src/neighbors/ivf_rabitq.cu @@ -73,10 +73,12 @@ auto build(raft::resources const& handle, host_dataset_ptr = dataset.data_handle(); if (params.force_streaming) { RAFT_LOG_INFO( - "Using streaming construction: explicitly requested via force_streaming parameter"); + "Using streaming construction: explicitly requested " + "via force_streaming parameter"); } else { RAFT_LOG_INFO( - "Using streaming construction: dataset size (%.2f GB) exceeds comfortable GPU memory " + "Using streaming construction: dataset size (%.2f GB) " + "exceeds comfortable GPU memory " "limit", dataset_bytes / (1024.0 * 1024.0 * 1024.0)); } @@ -123,8 +125,10 @@ auto build(raft::resources const& handle, handle, big_memory_resource, raft::make_extents(n_rows_train, dim)); } catch (raft::logic_error& e) { RAFT_LOG_ERROR( - "Insufficient memory for kmeans training set allocation. Please decrease " - "max_train_points_per_cluster, or set large_workspace_resource appropriately."); + "Insufficient memory for kmeans training set " + "allocation. Please decrease " + "max_train_points_per_cluster, or set " + "large_workspace_resource appropriately."); throw; } raft::matrix::sample_rows(handle, random_state, dataset, trainset.view()); @@ -193,7 +197,8 @@ void search(raft::resources const& handle, index& idx, raft::device_matrix_view queries, raft::device_matrix_view neighbors, - raft::device_matrix_view distances) + raft::device_matrix_view distances, + SampleFilterParams sample_filter) { auto stream = raft::resource::get_cuda_stream(handle).value(); @@ -264,25 +269,46 @@ void search(raft::resources const& handle, if (params.mode == search_mode::LUT32) { idx.rabitq_index().BatchClusterSearch( - queries_view, k, params.n_probes, searcher, NQ, distances, final_ids.view()); + queries_view, k, params.n_probes, searcher, NQ, distances, final_ids.view(), sample_filter); } else if (params.mode == search_mode::LUT16) { // test v3 lut using fp16 idx.rabitq_index().BatchClusterSearchLUT16( - queries_view, k, params.n_probes, searcher, NQ, distances, final_ids.view()); + queries_view, k, params.n_probes, searcher, NQ, distances, final_ids.view(), sample_filter); } else if (params.mode == search_mode::QUANT8) { - idx.rabitq_index().BatchClusterSearchQuantizeQuery( - queries_view, k, params.n_probes, searcher, NQ, distances, final_ids.view(), 8); + idx.rabitq_index().BatchClusterSearchQuantizeQuery(queries_view, + k, + params.n_probes, + searcher, + NQ, + distances, + final_ids.view(), + 8, + sample_filter); } else if (params.mode == search_mode::QUANT4) { - idx.rabitq_index().BatchClusterSearchQuantizeQuery( - queries_view, k, params.n_probes, searcher, NQ, distances, final_ids.view(), 4); + idx.rabitq_index().BatchClusterSearchQuantizeQuery(queries_view, + k, + params.n_probes, + searcher, + NQ, + distances, + final_ids.view(), + 4, + sample_filter); } - // cast data in final_ids to array of IdxT in neighbors + // Dummy candidates have an infinite distance. InnerProduct search negates its pseudo-distances + // before returning, so accept either sign here. Preserve every finite-result PID, including + // UINT32_MAX, and represent missing results with the public out-of-bounds sentinel. raft::linalg::map( handle, neighbors, - raft::cast_op{}, - raft::make_device_vector_view(final_ids.data_handle(), NQ * k)); + [] __device__(uint32_t id, float distance) { + return distance == raft::upper_bound() || distance == raft::lower_bound() + ? kOutOfBoundsRecord + : static_cast(id); + }, + raft::make_device_vector_view(final_ids.data_handle(), NQ * k), + raft::make_device_vector_view(distances.data_handle(), NQ * k)); } template @@ -381,10 +407,29 @@ void search(raft::resources const& handle, cuvs::neighbors::ivf_rabitq::index& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, - raft::device_matrix_view distances) + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter_ref) { + detail::SampleFilterParams sample_filter; + try { + dynamic_cast(sample_filter_ref); + } catch (const std::bad_cast&) { + try { + auto& bitset_filter = + dynamic_cast&>( + sample_filter_ref); + RAFT_EXPECTS(bitset_filter.view().size() >= index.size(), + "Bitset filter must contain at least index.size() bits"); + sample_filter.type = detail::SampleFilterType::Bitset; + sample_filter.bitset_ptr = bitset_filter.view().data(); + sample_filter.bitset_len = bitset_filter.view().size(); + sample_filter.original_nbits = bitset_filter.view().get_original_nbits(); + } catch (const std::bad_cast&) { + RAFT_FAIL("Unsupported sample filter type"); + } + } cuvs::neighbors::ivf_rabitq::detail::search( - handle, search_params, index, queries, neighbors, distances); + handle, search_params, index, queries, neighbors, distances, sample_filter); } void serialize(raft::resources const& handle, diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu index d0b1fedd95..ad4905134a 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu @@ -125,7 +125,8 @@ void IVFGPU::load_transposed(const char* filename) std::uint32_t version = 0; read_exact(&version, sizeof(std::uint32_t)); RAFT_EXPECTS(version == kSerializationVersion, - "ivf_rabitq::deserialize: serialization version mismatch (got %u, expected %u) " + "ivf_rabitq::deserialize: serialization version mismatch (got " + "%u, expected %u) " "in: %s", version, kSerializationVersion, @@ -162,12 +163,14 @@ void IVFGPU::load_transposed(const char* filename) ex_bits, filename); RAFT_EXPECTS(num_centroids > 0 && num_centroids <= cuvs::util::kMaxIvfNLists, - "ivf_rabitq::deserialize: num_centroids=%zu out of valid range (0, %u] in: %s", + "ivf_rabitq::deserialize: num_centroids=%zu out of valid " + "range (0, %u] in: %s", num_centroids, cuvs::util::kMaxIvfNLists, filename); RAFT_EXPECTS(cuvs::util::is_mul_no_overflow(num_vectors, num_padded_dim), - "ivf_rabitq::deserialize: num_vectors=%zu * num_padded_dim=%zu overflows size_t " + "ivf_rabitq::deserialize: num_vectors=%zu * " + "num_padded_dim=%zu overflows size_t " "in: %s", num_vectors, num_padded_dim, @@ -1211,7 +1214,8 @@ void IVFGPU::BatchClusterSearch( SearcherGPU& searcher, size_t batch_size, raft::device_matrix_view d_final_dists, - raft::device_matrix_view d_final_pids) + raft::device_matrix_view d_final_pids, + SampleFilterParams sample_filter) { auto d_sorted_pairs = raft::make_device_vector(searcher.get_handle(), 0); @@ -1228,7 +1232,8 @@ void IVFGPU::BatchClusterSearch( nprobe, k, d_final_dists, - d_final_pids); + d_final_pids, + sample_filter); } void IVFGPU::BatchClusterSearchLUT16( @@ -1238,7 +1243,8 @@ void IVFGPU::BatchClusterSearchLUT16( SearcherGPU& searcher, size_t batch_size, raft::device_matrix_view d_final_dists, - raft::device_matrix_view d_final_pids) + raft::device_matrix_view d_final_pids, + SampleFilterParams sample_filter) { auto d_sorted_pairs = raft::make_device_vector(searcher.get_handle(), 0); @@ -1255,7 +1261,8 @@ void IVFGPU::BatchClusterSearchLUT16( nprobe, k, d_final_dists, - d_final_pids); + d_final_pids, + sample_filter); } void IVFGPU::BatchClusterSearchQuantizeQuery( @@ -1266,7 +1273,8 @@ void IVFGPU::BatchClusterSearchQuantizeQuery( size_t batch_size, raft::device_matrix_view d_final_dists, raft::device_matrix_view d_final_pids, - int query_bits) + int query_bits, + SampleFilterParams sample_filter) { auto d_sorted_pairs = raft::make_device_vector(searcher.get_handle(), 0); @@ -1284,7 +1292,8 @@ void IVFGPU::BatchClusterSearchQuantizeQuery( k, d_final_dists, d_final_pids, - query_bits == 4); + query_bits == 4, + sample_filter); } } // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh index f05ca871d1..98cb242d52 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh @@ -34,6 +34,7 @@ namespace cuvs::neighbors::ivf_rabitq::detail { class SearcherGPU; +struct SampleFilterParams; // Structure for cluster-query pairs struct ClusterQueryPair { @@ -276,14 +277,14 @@ class IVFGPU { DataQuantizerGPU& quantizer() const { return *(this->DQ); } RotatorGPU& rotator() const { return *(this->Rota); } - void BatchClusterSearch( - raft::device_matrix_view queries, - size_t k, - size_t nprobe, - SearcherGPU& searcher, - size_t batch_size, - raft::device_matrix_view d_final_dists, - raft::device_matrix_view d_final_pids); + void BatchClusterSearch(raft::device_matrix_view queries, + size_t k, + size_t nprobe, + SearcherGPU& searcher, + size_t batch_size, + raft::device_matrix_view d_final_dists, + raft::device_matrix_view d_final_pids, + SampleFilterParams sample_filter); void BatchClusterSearchLUT16( raft::device_matrix_view queries, @@ -292,7 +293,8 @@ class IVFGPU { SearcherGPU& searcher, size_t batch_size, raft::device_matrix_view d_final_dists, - raft::device_matrix_view d_final_pids); + raft::device_matrix_view d_final_pids, + SampleFilterParams sample_filter); void BatchClusterSearchQuantizeQuery( raft::device_matrix_view queries, @@ -302,7 +304,8 @@ class IVFGPU { size_t batch_size, raft::device_matrix_view d_final_dists, raft::device_matrix_view d_final_pids, - int query_bits); + int query_bits, + SampleFilterParams sample_filter); private: void PrepareClusterSearchInputs( diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu index 213c5d7a19..66d30d4b46 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu @@ -138,7 +138,8 @@ void SearcherGPU::SearchClusterQueryPairs( size_t nprobe, size_t topk, raft::device_matrix_view d_final_dists, - raft::device_matrix_view d_final_pids) + raft::device_matrix_view d_final_pids, + SampleFilterParams sample_filter) { // First allocate space for LUT size_t lut_size = num_queries * (D / BITS_PER_CHUNK) * LUT_SIZE * sizeof(float); @@ -240,22 +241,26 @@ void SearcherGPU::SearchClusterQueryPairs( kernelParams.max_candidates_per_pair = max_cluster_size; kernelParams.max_candidates_per_query = use_block_sort ? 0 /* unused */ : max_probed_vectors_count.value(); - kernelParams.ex_bits = cur_ivf.get_ex_bits(); - kernelParams.d_long_code = cur_ivf.get_long_code_device(); - kernelParams.d_ex_factor = reinterpret_cast(cur_ivf.get_ex_factor_device()); - kernelParams.d_pids = cur_ivf.get_ids_device(); - kernelParams.d_topk_dists = d_topk_dists.data_handle(); - kernelParams.d_topk_pids = d_topk_pids.data_handle(); + kernelParams.ex_bits = cur_ivf.get_ex_bits(); + kernelParams.d_long_code = cur_ivf.get_long_code_device(); + kernelParams.d_ex_factor = reinterpret_cast(cur_ivf.get_ex_factor_device()); + kernelParams.d_pids = cur_ivf.get_ids_device(); + kernelParams.bitset_ptr = sample_filter.bitset_ptr; + kernelParams.bitset_len = sample_filter.bitset_len; + kernelParams.original_nbits = sample_filter.original_nbits; + kernelParams.d_topk_dists = d_topk_dists.data_handle(); + kernelParams.d_topk_pids = d_topk_pids.data_handle(); kernelParams.d_query_write_counters = d_query_write_counters.data(); if (cur_ivf.get_ex_bits() != 0) { size_t shared_mem_size = num_chunks * LUT_SIZE * sizeof(float) + candidate_storage + queue_buffer_smem_bytes; - auto jit_launcher = use_block_sort - ? make_compute_inner_products_with_lut_block_sort_launcher( - cur_ivf.get_ex_bits(), /*with_ex=*/true, is_inner_product) - : make_compute_inner_products_with_lut_launcher( - cur_ivf.get_ex_bits(), /*with_ex=*/true, is_inner_product); + auto jit_launcher = + use_block_sort + ? make_compute_inner_products_with_lut_block_sort_launcher( + cur_ivf.get_ex_bits(), /*with_ex=*/true, is_inner_product, sample_filter.type) + : make_compute_inner_products_with_lut_launcher( + cur_ivf.get_ex_bits(), /*with_ex=*/true, is_inner_product, sample_filter.type); auto const& kernel_launcher = [&]() -> void { jit_launcher->dispatch( stream_, gridDim, blockDim, shared_mem_size, kernelParams); @@ -268,10 +273,11 @@ void SearcherGPU::SearchClusterQueryPairs( max(num_chunks * LUT_SIZE * sizeof(float) + (use_block_sort ? max_cluster_size * (sizeof(float) + sizeof(int)) : 0), (size_t)queue_buffer_smem_bytes); - auto jit_launcher = use_block_sort ? make_compute_inner_products_with_lut_block_sort_launcher( - /*ex_bits=*/0, /*with_ex=*/false, is_inner_product) - : make_compute_inner_products_with_lut_launcher( - /*ex_bits=*/0, /*with_ex=*/false, is_inner_product); + auto jit_launcher = + use_block_sort ? make_compute_inner_products_with_lut_block_sort_launcher( + /*ex_bits=*/0, /*with_ex=*/false, is_inner_product, sample_filter.type) + : make_compute_inner_products_with_lut_launcher( + /*ex_bits=*/0, /*with_ex=*/false, is_inner_product, sample_filter.type); auto const& kernel_launcher = [&]() -> void { jit_launcher->dispatch( stream_, gridDim, blockDim, shared_mem_size, kernelParams); diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh index 73557f57ea..2bc2955456 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -21,6 +21,15 @@ namespace cuvs::neighbors::ivf_rabitq::detail { +enum class SampleFilterType { None, Bitset }; + +struct SampleFilterParams { + SampleFilterType type = SampleFilterType::None; + uint32_t* bitset_ptr = nullptr; + int64_t bitset_len = 0; + int64_t original_nbits = 0; +}; + // block-level sorting used if topk <= kMaxTopKBlockSort; must be power of 2; increases shared mem // usage static constexpr int kMaxTopKBlockSort = 64; @@ -79,7 +88,8 @@ class SearcherGPU { size_t nprobe, size_t topk, raft::device_matrix_view d_final_dists, - raft::device_matrix_view d_final_pids); + raft::device_matrix_view d_final_pids, + SampleFilterParams sample_filter); void SearchClusterQueryPairsSharedMemOpt( const IVFGPU& cur_ivf, @@ -92,7 +102,8 @@ class SearcherGPU { size_t nprobe, size_t topk, raft::device_matrix_view d_final_dists, - raft::device_matrix_view d_final_pids); + raft::device_matrix_view d_final_pids, + SampleFilterParams sample_filter); void SearchClusterQueryPairsQuantizeQuery( const IVFGPU& cur_ivf, @@ -106,7 +117,8 @@ class SearcherGPU { size_t topk, raft::device_matrix_view d_final_dists, raft::device_matrix_view d_final_pids, - bool use_4bit = false); + bool use_4bit, + SampleFilterParams sample_filter); private: raft::resources const& handle_; // reusable resource handle diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh index fc9aaf0189..8a218027ac 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh @@ -60,6 +60,9 @@ struct ComputeInnerProductsKernelParams { const uint8_t* d_long_code = nullptr; // long codes for all vectors const float* d_ex_factor = nullptr; // ex factors for distance computation const PID* d_pids = nullptr; // PIDs for all vectors + uint32_t* bitset_ptr = nullptr; + int64_t bitset_len = 0; + int64_t original_nbits = 0; float* d_topk_dists = nullptr; // output top-k distances PID* d_topk_pids = nullptr; // output top-k PIDs int* d_query_write_counters = nullptr; diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu index 747ec85f32..a62c8f05ca 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu @@ -324,8 +324,8 @@ void SearcherGPU::SearchClusterQueryPairsQuantizeQuery( size_t topk, raft::device_matrix_view d_final_dists, raft::device_matrix_view d_final_pids, - bool use_4bit // Add parameter to choose 4-bit or 8-bit -) + bool use_4bit, + SampleFilterParams sample_filter) { // check if the inner products kernel should use block sort to keep a top-k priority queue vs. // outputting distances from all vectors in probed clusters @@ -516,24 +516,31 @@ void SearcherGPU::SearchClusterQueryPairsQuantizeQuery( kernelParams.max_candidates_per_pair = max_cluster_size; kernelParams.max_candidates_per_query = use_block_sort ? 0 /* unused */ : max_probed_vectors_count.value(); - kernelParams.ex_bits = cur_ivf.get_ex_bits(); - kernelParams.d_long_code = cur_ivf.get_long_code_device(); - kernelParams.d_ex_factor = reinterpret_cast(cur_ivf.get_ex_factor_device()); - kernelParams.d_pids = cur_ivf.get_ids_device(); - kernelParams.d_topk_dists = d_topk_dists.data_handle(); - kernelParams.d_topk_pids = d_topk_pids.data_handle(); + kernelParams.ex_bits = cur_ivf.get_ex_bits(); + kernelParams.d_long_code = cur_ivf.get_long_code_device(); + kernelParams.d_ex_factor = reinterpret_cast(cur_ivf.get_ex_factor_device()); + kernelParams.d_pids = cur_ivf.get_ids_device(); + kernelParams.bitset_ptr = sample_filter.bitset_ptr; + kernelParams.bitset_len = sample_filter.bitset_len; + kernelParams.original_nbits = sample_filter.original_nbits; + kernelParams.d_topk_dists = d_topk_dists.data_handle(); + kernelParams.d_topk_pids = d_topk_pids.data_handle(); kernelParams.d_query_write_counters = d_query_write_counters.data_handle(); kernelParams.num_bits = num_bits; kernelParams.num_words = num_words; const int num_bits_for_dispatch = use_4bit ? 4 : 8; const bool with_ex = (cur_ivf.get_ex_bits() != 0); - auto jit_launcher = use_block_sort - ? make_compute_inner_products_with_bitwise_block_sort_launcher( - num_bits_for_dispatch, cur_ivf.get_ex_bits(), with_ex, is_inner_product) - : make_compute_inner_products_with_bitwise_launcher( - cur_ivf.get_ex_bits(), with_ex, is_inner_product); - auto const& kernel_launcher = [&]() -> void { + auto jit_launcher = + use_block_sort + ? make_compute_inner_products_with_bitwise_block_sort_launcher(num_bits_for_dispatch, + cur_ivf.get_ex_bits(), + with_ex, + is_inner_product, + sample_filter.type) + : make_compute_inner_products_with_bitwise_launcher( + cur_ivf.get_ex_bits(), with_ex, is_inner_product, sample_filter.type); + auto const& kernel_launcher = [&]() -> void { jit_launcher->dispatch( stream_, gridDim, blockDim, shared_mem_size, kernelParams); }; diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu index 85cda66513..4153c86943 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu @@ -92,7 +92,8 @@ void SearcherGPU::SearchClusterQueryPairsSharedMemOpt( size_t nprobe, size_t topk, raft::device_matrix_view d_final_dists, - raft::device_matrix_view d_final_pids) + raft::device_matrix_view d_final_pids, + SampleFilterParams sample_filter) { // Using FP16 for storage @@ -201,12 +202,15 @@ void SearcherGPU::SearchClusterQueryPairsSharedMemOpt( kernelParams.max_candidates_per_pair = max_cluster_size; kernelParams.max_candidates_per_query = use_block_sort ? 0 /* unused */ : max_probed_vectors_count.value(); - kernelParams.ex_bits = cur_ivf.get_ex_bits(); - kernelParams.d_long_code = cur_ivf.get_long_code_device(); - kernelParams.d_ex_factor = reinterpret_cast(cur_ivf.get_ex_factor_device()); - kernelParams.d_pids = cur_ivf.get_ids_device(); - kernelParams.d_topk_dists = d_topk_dists.data_handle(); - kernelParams.d_topk_pids = d_topk_pids.data_handle(); + kernelParams.ex_bits = cur_ivf.get_ex_bits(); + kernelParams.d_long_code = cur_ivf.get_long_code_device(); + kernelParams.d_ex_factor = reinterpret_cast(cur_ivf.get_ex_factor_device()); + kernelParams.d_pids = cur_ivf.get_ids_device(); + kernelParams.bitset_ptr = sample_filter.bitset_ptr; + kernelParams.bitset_len = sample_filter.bitset_len; + kernelParams.original_nbits = sample_filter.original_nbits; + kernelParams.d_topk_dists = d_topk_dists.data_handle(); + kernelParams.d_topk_pids = d_topk_pids.data_handle(); kernelParams.d_query_write_counters = d_query_write_counters.data(); if (cur_ivf.get_ex_bits() != 0) { @@ -221,11 +225,12 @@ void SearcherGPU::SearchClusterQueryPairsSharedMemOpt( size_t shared_mem_size = max(first_part_shared_mem + second_part_shared_mem + third_part_shared_mem, (size_t)queue_buffer_smem_bytes); - auto jit_launcher = use_block_sort - ? make_compute_inner_products_with_lut16_opt_block_sort_launcher( - cur_ivf.get_ex_bits(), /*with_ex=*/true, is_inner_product) - : make_compute_inner_products_with_lut16_opt_launcher( - cur_ivf.get_ex_bits(), /*with_ex=*/true, is_inner_product); + auto jit_launcher = + use_block_sort + ? make_compute_inner_products_with_lut16_opt_block_sort_launcher( + cur_ivf.get_ex_bits(), /*with_ex=*/true, is_inner_product, sample_filter.type) + : make_compute_inner_products_with_lut16_opt_launcher( + cur_ivf.get_ex_bits(), /*with_ex=*/true, is_inner_product, sample_filter.type); auto const& kernel_launcher = [&]() -> void { jit_launcher->dispatch( stream_, gridDim, blockDim, shared_mem_size, kernelParams); @@ -240,11 +245,11 @@ void SearcherGPU::SearchClusterQueryPairsSharedMemOpt( // queue buffer reuses first 2 parts size_t shared_mem_size = max(first_part_shared_mem + second_part_shared_mem, (size_t)queue_buffer_smem_bytes); - auto jit_launcher = use_block_sort - ? make_compute_inner_products_with_lut16_opt_block_sort_launcher( - /*ex_bits=*/0, /*with_ex=*/false, is_inner_product) - : make_compute_inner_products_with_lut16_opt_launcher( - /*ex_bits=*/0, /*with_ex=*/false, is_inner_product); + auto jit_launcher = + use_block_sort ? make_compute_inner_products_with_lut16_opt_block_sort_launcher( + /*ex_bits=*/0, /*with_ex=*/false, is_inner_product, sample_filter.type) + : make_compute_inner_products_with_lut16_opt_launcher( + /*ex_bits=*/0, /*with_ex=*/false, is_inner_product, sample_filter.type); auto const& kernel_launcher = [&]() -> void { jit_launcher->dispatch( stream_, gridDim, blockDim, shared_mem_size, kernelParams); diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in index 906c07cd57..37eb425883 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in @@ -71,6 +71,12 @@ __device__ void bitwise_emit_distances(const ComputeInnerProductsKernelParams pa float q_kbxsumq = params.d_G_kbxSumq[query_idx]; for (size_t vec_idx = tid; vec_idx < num_vectors_in_cluster; vec_idx += num_threads) { + size_t global_vec_idx = cluster_start_index + vec_idx; + if (!is_sample_allowed(params, global_vec_idx)) { + params.d_topk_dists[output_offset + vec_idx] = INFINITY; + continue; + } + float exact_ip = compute_bitwise_1bit_ip_for_vec(params.d_short_data, shared_query, cluster_start_index, @@ -79,11 +85,10 @@ __device__ void bitwise_emit_distances(const ComputeInnerProductsKernelParams pa short_code_length, params.D); - float ip2 = shared_ip2_results[vec_idx]; - size_t global_vec_idx = cluster_start_index + vec_idx; - float2 ex_factors = reinterpret_cast(params.d_ex_factor)[global_vec_idx]; - float f_ex_add = ex_factors.x; - float f_ex_rescale = ex_factors.y; + float ip2 = shared_ip2_results[vec_idx]; + float2 ex_factors = reinterpret_cast(params.d_ex_factor)[global_vec_idx]; + float f_ex_add = ex_factors.x; + float f_ex_rescale = ex_factors.y; float dist = f_ex_add + q_g_add + @@ -98,11 +103,15 @@ __device__ void bitwise_emit_distances(const ComputeInnerProductsKernelParams pa float q_k1xsumq = params.d_G_k1xSumq[query_idx]; for (size_t vec_idx = tid; vec_idx < num_vectors_in_cluster; vec_idx += num_threads) { - size_t factor_offset = cluster_start_index + vec_idx; - float3 factors = reinterpret_cast(params.d_short_factors)[factor_offset]; - float f_add = factors.x; - float f_rescale = factors.y; size_t global_vec_idx = cluster_start_index + vec_idx; + if (!is_sample_allowed(params, global_vec_idx)) { + params.d_topk_dists[output_offset + vec_idx] = INFINITY; + continue; + } + + float3 factors = reinterpret_cast(params.d_short_factors)[global_vec_idx]; + float f_add = factors.x; + float f_rescale = factors.y; float exact_ip = compute_bitwise_1bit_ip_for_vec(params.d_short_data, shared_query, diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh index 26411f4db7..18d9e3f56d 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh @@ -78,7 +78,8 @@ __device__ void compute_inner_products_with_bitwise_block_sort_impl( bool is_candidate = false; float local_ip_quantized = 0; - if (vec_idx < num_vectors_in_cluster) { + if (vec_idx < num_vectors_in_cluster && + is_sample_allowed(params, cluster_start_index + vec_idx)) { size_t factor_offset = cluster_start_index + vec_idx; float3 factors = reinterpret_cast(params.d_short_factors)[factor_offset]; float f_add = factors.x; diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp index 5c57224f96..2b56879ac5 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp @@ -43,6 +43,11 @@ struct ComputeInnerProductsWithBitwiseBlockSortPlanner : AlgorithmPlanner { { this->add_static_fragment>(); } + template + void add_sample_filter_device_function() + { + this->add_static_fragment>(); + } }; } // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp index 67bfd636d0..f08713f315 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp @@ -35,6 +35,11 @@ struct ComputeInnerProductsWithBitwisePlanner : AlgorithmPlanner { { this->add_static_fragment>(); } + template + void add_sample_filter_device_function() + { + this->add_static_fragment>(); + } }; } // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh index 2126c1c581..fdf5e6e386 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh @@ -85,7 +85,8 @@ __device__ void compute_inner_products_with_lut16_opt_block_sort_impl( float local_ip = 0.0f; bool is_candidate = false; - if (vec_idx < num_vectors_in_cluster) { + if (vec_idx < num_vectors_in_cluster && + is_sample_allowed(params, cluster_start_index + vec_idx)) { size_t factor_offset = cluster_start_index + vec_idx; float3 factors = reinterpret_cast(params.d_short_factors)[factor_offset]; float f_add = factors.x; diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp index 7cb0d4b387..d9786ab4dd 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp @@ -36,6 +36,11 @@ struct ComputeInnerProductsWithLut16OptBlockSortPlanner : AlgorithmPlanner { { this->add_static_fragment>(); } + template + void add_sample_filter_device_function() + { + this->add_static_fragment>(); + } }; } // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp index 37b7bde28e..9e2ece23e4 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp @@ -40,6 +40,11 @@ struct ComputeInnerProductsWithLut16OptPlanner : AlgorithmPlanner { { this->add_static_fragment>(); } + template + void add_sample_filter_device_function() + { + this->add_static_fragment>(); + } }; } // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp index 8299305fd6..06f51182ba 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp @@ -40,6 +40,11 @@ struct ComputeInnerProductsWithLutBlockSortPlanner : AlgorithmPlanner { { this->add_static_fragment>(); } + template + void add_sample_filter_device_function() + { + this->add_static_fragment>(); + } }; } // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp index 0d1efc1cf1..cd5648a96a 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp @@ -40,6 +40,11 @@ struct ComputeInnerProductsWithLutPlanner : AlgorithmPlanner { { this->add_static_fragment>(); } + template + void add_sample_filter_device_function() + { + this->add_static_fragment>(); + } }; } // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh index 296060ffd4..0093565d9e 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh @@ -22,6 +22,18 @@ namespace cuvs::neighbors::ivf_rabitq::detail { // adds them via add_*_device_function(). __device__ uint32_t extract_code(const uint8_t* codes, size_t d); +__device__ bool sample_filter(PID source_id, + uint32_t* bitset_ptr, + int64_t bitset_len, + int64_t original_nbits); + +__forceinline__ __device__ bool is_sample_allowed(const ComputeInnerProductsKernelParams& params, + size_t global_vec_idx) +{ + return sample_filter( + params.d_pids[global_vec_idx], params.bitset_ptr, params.bitset_len, params.original_nbits); +} + // The body does not depend on ex_bits; the inner extract_code call is the // ex_bits-specialized fragment, resolved at nvJitLink time based on which // extract_code variant the planner adds. diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp index 2891979e17..6670a8c754 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp @@ -5,6 +5,7 @@ #pragma once +#include "../gpu_index/searcher_gpu.cuh" #include "compute_inner_products_with_bitwise_block_sort_planner.hpp" #include "compute_inner_products_with_bitwise_planner.hpp" #include "compute_inner_products_with_lut16_opt_block_sort_planner.hpp" @@ -13,6 +14,7 @@ #include "compute_inner_products_with_lut_planner.hpp" #include +#include #include #include @@ -21,6 +23,22 @@ namespace cuvs::neighbors::ivf_rabitq::detail { namespace { +template +inline void add_sample_filter_device_function(Planner& planner, SampleFilterType filter_type) +{ + switch (filter_type) { + case SampleFilterType::None: + planner + .template add_sample_filter_device_function(); + break; + case SampleFilterType::Bitset: + planner + .template add_sample_filter_device_function(); + break; + default: assert(false); + } +} + template inline void add_ex_bits_device_functions(Planner& planner, int ex_bits) { @@ -40,9 +58,10 @@ inline void add_ex_bits_device_functions(Planner& planner, int ex_bits) } // namespace inline std::shared_ptr make_compute_inner_products_with_lut_launcher( - int ex_bits, bool with_ex, bool is_inner_product) + int ex_bits, bool with_ex, bool is_inner_product, SampleFilterType filter_type) { ComputeInnerProductsWithLutPlanner planner; + add_sample_filter_device_function(planner, filter_type); planner.add_entrypoint(); planner.add_compute_lut_ip_for_vec_device_function(); if (with_ex) { @@ -63,9 +82,10 @@ inline std::shared_ptr make_compute_inner_products_with_lut_l } inline std::shared_ptr make_compute_inner_products_with_lut_block_sort_launcher( - int ex_bits, bool with_ex, bool is_inner_product) + int ex_bits, bool with_ex, bool is_inner_product, SampleFilterType filter_type) { ComputeInnerProductsWithLutBlockSortPlanner planner; + add_sample_filter_device_function(planner, filter_type); planner.add_entrypoint(); planner.add_compute_lut_ip_for_vec_device_function(); if (with_ex) { @@ -86,9 +106,10 @@ inline std::shared_ptr make_compute_inner_products_with_lut_b } inline std::shared_ptr make_compute_inner_products_with_lut16_opt_launcher( - int ex_bits, bool with_ex, bool is_inner_product) + int ex_bits, bool with_ex, bool is_inner_product, SampleFilterType filter_type) { ComputeInnerProductsWithLut16OptPlanner planner; + add_sample_filter_device_function(planner, filter_type); planner.add_entrypoint(); planner.add_compute_lut_ip_for_vec_device_function(); if (with_ex) { @@ -111,9 +132,11 @@ inline std::shared_ptr make_compute_inner_products_with_lut16 inline std::shared_ptr make_compute_inner_products_with_lut16_opt_block_sort_launcher(int ex_bits, bool with_ex, - bool is_inner_product) + bool is_inner_product, + SampleFilterType filter_type) { ComputeInnerProductsWithLut16OptBlockSortPlanner planner; + add_sample_filter_device_function(planner, filter_type); planner.add_compute_lut_ip_for_vec_device_function(); if (with_ex) { if (is_inner_product) { @@ -133,9 +156,10 @@ make_compute_inner_products_with_lut16_opt_block_sort_launcher(int ex_bits, } inline std::shared_ptr make_compute_inner_products_with_bitwise_launcher( - int ex_bits, bool with_ex, bool is_inner_product) + int ex_bits, bool with_ex, bool is_inner_product, SampleFilterType filter_type) { ComputeInnerProductsWithBitwisePlanner planner; + add_sample_filter_device_function(planner, filter_type); planner.add_entrypoint(); if (with_ex) { if (is_inner_product) { @@ -155,12 +179,11 @@ inline std::shared_ptr make_compute_inner_products_with_bitwi } inline std::shared_ptr -make_compute_inner_products_with_bitwise_block_sort_launcher(int num_bits, - int ex_bits, - bool with_ex, - bool is_inner_product) +make_compute_inner_products_with_bitwise_block_sort_launcher( + int num_bits, int ex_bits, bool with_ex, bool is_inner_product, SampleFilterType filter_type) { ComputeInnerProductsWithBitwiseBlockSortPlanner planner; + add_sample_filter_device_function(planner, filter_type); if (is_inner_product) { planner.add_entrypoint(); } else { diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in index db87e7b734..0163c8d986 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in @@ -87,6 +87,12 @@ __device__ void lut16_opt_emit_distances(const ComputeInnerProductsKernelParams float q_kbxsumq = params.d_G_kbxSumq[query_idx]; for (size_t vec_idx = tid; vec_idx < num_vectors_in_cluster; vec_idx += num_threads) { + size_t global_vec_idx = cluster_start_index + vec_idx; + if (!is_sample_allowed(params, global_vec_idx)) { + params.d_topk_dists[output_offset + vec_idx] = INFINITY; + continue; + } + float ip = compute_lut_ip_for_vec<__half>(params.d_short_data, shared_lut_fp16, cluster_start_index, @@ -94,8 +100,7 @@ __device__ void lut16_opt_emit_distances(const ComputeInnerProductsKernelParams vec_idx, short_code_length); - float ip2 = shared_ip2_results[vec_idx]; - size_t global_vec_idx = cluster_start_index + vec_idx; + float ip2 = shared_ip2_results[vec_idx]; float2 ex_factors = reinterpret_cast(params.d_ex_factor)[global_vec_idx]; float f_ex_add = ex_factors.x; @@ -131,10 +136,15 @@ __device__ void lut16_opt_emit_distances(const ComputeInnerProductsKernelParams uint32_t output_offset = query_idx * params.max_candidates_per_query + probe_slot; for (size_t vec_idx = tid; vec_idx < num_vectors_in_cluster; vec_idx += num_threads) { - size_t factor_offset = cluster_start_index + vec_idx; - float3 factors = reinterpret_cast(params.d_short_factors)[factor_offset]; - float f_add = factors.x; - float f_rescale = factors.y; + size_t global_vec_idx = cluster_start_index + vec_idx; + if (!is_sample_allowed(params, global_vec_idx)) { + params.d_topk_dists[output_offset + vec_idx] = INFINITY; + continue; + } + + float3 factors = reinterpret_cast(params.d_short_factors)[global_vec_idx]; + float f_add = factors.x; + float f_rescale = factors.y; float ip = compute_lut_ip_for_vec<__half>(params.d_short_data, shared_lut_fp16, @@ -145,10 +155,10 @@ __device__ void lut16_opt_emit_distances(const ComputeInnerProductsKernelParams float dist = f_add + q_g_add + f_rescale * (ip + q_k1xsumq); if constexpr (is_signed) { - dist = (dist - q_sqr_norm - params.d_vec_sqr_norms[cluster_start_index + vec_idx]) * 0.5f; + dist = (dist - q_sqr_norm - params.d_vec_sqr_norms[global_vec_idx]) * 0.5f; } params.d_topk_dists[output_offset + vec_idx] = dist; - params.d_topk_pids[output_offset + vec_idx] = params.d_pids[cluster_start_index + vec_idx]; + params.d_topk_pids[output_offset + vec_idx] = params.d_pids[global_vec_idx]; } } } diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in index 33ae494212..f99b66599e 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in @@ -73,7 +73,8 @@ __device__ void lut_block_sort_emit_topk(const ComputeInnerProductsKernelParams float local_ip = 0.0f; bool is_candidate = false; - if (vec_idx < num_vectors_in_cluster) { + if (vec_idx < num_vectors_in_cluster && + is_sample_allowed(params, cluster_start_index + vec_idx)) { size_t factor_offset = cluster_start_index + vec_idx; float3 factors = reinterpret_cast(params.d_short_factors)[factor_offset]; float f_add = factors.x; @@ -214,7 +215,8 @@ __device__ void lut_block_sort_emit_topk(const ComputeInnerProductsKernelParams bool is_candidate = false; - if (vec_idx < num_vectors_in_cluster) { + if (vec_idx < num_vectors_in_cluster && + is_sample_allowed(params, cluster_start_index + vec_idx)) { size_t factor_offset = cluster_start_index + vec_idx; float3 factors = reinterpret_cast(params.d_short_factors)[factor_offset]; float f_add = factors.x; diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in index 63739ce168..44ac9366f0 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in @@ -84,6 +84,12 @@ __device__ void lut_emit_distances(const ComputeInnerProductsKernelParams params size_t output_offset = query_idx * params.max_candidates_per_query + probe_slot; for (size_t vec_idx = tid; vec_idx < num_vectors_in_cluster; vec_idx += num_threads) { + size_t global_vec_idx = cluster_start_index + vec_idx; + if (!is_sample_allowed(params, global_vec_idx)) { + params.d_topk_dists[output_offset + vec_idx] = INFINITY; + continue; + } + float ip = compute_lut_ip_for_vec(params.d_short_data, shared_lut, cluster_start_index, @@ -91,8 +97,7 @@ __device__ void lut_emit_distances(const ComputeInnerProductsKernelParams params vec_idx, short_code_length); - float ip2 = shared_ip2_results[vec_idx]; - size_t global_vec_idx = cluster_start_index + vec_idx; + float ip2 = shared_ip2_results[vec_idx]; float2 ex_factors = reinterpret_cast(params.d_ex_factor)[global_vec_idx]; float f_ex_add = ex_factors.x; @@ -128,10 +133,15 @@ __device__ void lut_emit_distances(const ComputeInnerProductsKernelParams params size_t output_offset = query_idx * params.max_candidates_per_query + probe_slot; for (size_t vec_idx = tid; vec_idx < num_vectors_in_cluster; vec_idx += num_threads) { - size_t factor_offset = cluster_start_index + vec_idx; - float3 factors = reinterpret_cast(params.d_short_factors)[factor_offset]; - float f_add = factors.x; - float f_rescale = factors.y; + size_t global_vec_idx = cluster_start_index + vec_idx; + if (!is_sample_allowed(params, global_vec_idx)) { + params.d_topk_dists[output_offset + vec_idx] = INFINITY; + continue; + } + + float3 factors = reinterpret_cast(params.d_short_factors)[global_vec_idx]; + float f_add = factors.x; + float f_rescale = factors.y; float ip = compute_lut_ip_for_vec(params.d_short_data, shared_lut, @@ -142,10 +152,10 @@ __device__ void lut_emit_distances(const ComputeInnerProductsKernelParams params float dist = f_add + q_g_add + f_rescale * (ip + q_k1xsumq); if constexpr (is_signed) { - dist = (dist - q_sqr_norm - params.d_vec_sqr_norms[cluster_start_index + vec_idx]) * 0.5f; + dist = (dist - q_sqr_norm - params.d_vec_sqr_norms[global_vec_idx]) * 0.5f; } params.d_topk_dists[output_offset + vec_idx] = dist; - params.d_topk_pids[output_offset + vec_idx] = params.d_pids[cluster_start_index + vec_idx]; + params.d_topk_pids[output_offset + vec_idx] = params.d_pids[global_vec_idx]; } } } diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/sample_filter_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/sample_filter_kernel.cu.in new file mode 100644 index 0000000000..d6aff236cf --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/sample_filter_kernel.cu.in @@ -0,0 +1,45 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +namespace { + +template +__device__ bool sample_filter_none(PID, uint32_t*, int64_t, int64_t) +{ + return true; +} + +template +__device__ bool sample_filter_bitset(PID source_id, + uint32_t* bitset_ptr, + int64_t bitset_len, + int64_t original_nbits) +{ + auto bitset_view = + raft::core::bitset_view{bitset_ptr, bitset_len, original_nbits}; + return bitset_view.test(source_id); +} + +constexpr auto sample_filter_impl = sample_filter_@filter_name@<>; + +} // namespace + +__device__ bool sample_filter(PID source_id, + uint32_t* bitset_ptr, + int64_t bitset_len, + int64_t original_nbits) +{ + return sample_filter_impl(source_id, bitset_ptr, bitset_len, original_nbits); +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/sample_filter_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/sample_filter_matrix.json new file mode 100644 index 0000000000..09c929ec45 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/sample_filter_matrix.json @@ -0,0 +1,4 @@ +[ + {"filter_name": "none"}, + {"filter_name": "bitset"} +] diff --git a/cpp/tests/neighbors/ann_ivf_rabitq.cuh b/cpp/tests/neighbors/ann_ivf_rabitq.cuh index 3d3ba38ff8..8d897abb4d 100644 --- a/cpp/tests/neighbors/ann_ivf_rabitq.cuh +++ b/cpp/tests/neighbors/ann_ivf_rabitq.cuh @@ -7,8 +7,15 @@ #include "../test_utils.cuh" #include "ann_utils.cuh" #include "naive_knn.cuh" +#include #include #include +#include + +#include + +#include +#include namespace cuvs::neighbors::ivf_rabitq { @@ -17,6 +24,7 @@ struct ivf_rabitq_inputs { uint32_t num_queries = 1024; uint32_t dim = 64; uint32_t k = 10; + uint32_t filter_pass_count = 0; std::optional min_recall = std::nullopt; // Generate signed, mean-zero data instead of the default positive-only data. Used by the // InnerProduct cases to produce mixed-sign inner products (which exercise both branches of the @@ -61,8 +69,10 @@ inline auto operator<<(std::ostream& os, const ivf_rabitq_inputs& p) -> std::ost PRINT_DIFF(.num_queries); PRINT_DIFF(.dim); PRINT_DIFF(.k); + PRINT_DIFF(.filter_pass_count); PRINT_DIFF_V(.min_recall, p.min_recall.value_or(0)); PRINT_DIFF(.signed_data); + PRINT_DIFF_V(.index_params.metric, static_cast(p.index_params.metric)); PRINT_DIFF(.index_params.n_lists); PRINT_DIFF(.index_params.bits_per_dim); PRINT_DIFF(.index_params.kmeans_n_iters); @@ -109,6 +119,16 @@ class ivf_rabitq_test : public ::testing::TestWithParam { void calc_ref() { + const IdxT filter_offset = static_cast(ps.num_db_vecs - ps.filter_pass_count); + const bool filtered = ps.filter_pass_count > 0; + const IdxT ref_offset = filtered ? filter_offset : IdxT{0}; + const uint32_t ref_size = filtered ? ps.filter_pass_count : ps.num_db_vecs; + if (ref_size < ps.k) { + indices_ref.clear(); + distances_ref.clear(); + return; + } + size_t queries_size = size_t{ps.num_queries} * size_t{ps.k}; rmm::device_uvector distances_naive_dev(queries_size, stream_); rmm::device_uvector indices_naive_dev(queries_size, stream_); @@ -117,12 +137,16 @@ class ivf_rabitq_test : public ::testing::TestWithParam { distances_naive_dev.data(), indices_naive_dev.data(), search_queries.data(), - database.data(), + database.data() + static_cast(ref_offset) * ps.dim, ps.num_queries, - ps.num_db_vecs, + ref_size, ps.dim, ps.k, static_cast((int)ps.index_params.metric)); + if (filtered) { + raft::linalg::addScalar( + indices_naive_dev.data(), indices_naive_dev.data(), ref_offset, queries_size, stream_); + } distances_ref.resize(queries_size); raft::update_host(distances_ref.data(), distances_naive_dev.data(), queries_size, stream_); indices_ref.resize(queries_size); @@ -239,6 +263,79 @@ class ivf_rabitq_test : public ::testing::TestWithParam { << ps; } + void run_filter() + { + ASSERT_GT(ps.filter_pass_count, 0); + ASSERT_LE(ps.filter_pass_count, ps.num_db_vecs); + const IdxT filter_offset = static_cast(ps.num_db_vecs - ps.filter_pass_count); + + index index = build_serialize(); + const size_t queries_size = size_t{ps.num_queries} * ps.k; + auto distances = raft::make_device_matrix(handle_, ps.num_queries, ps.k); + auto indices = raft::make_device_matrix(handle_, ps.num_queries, ps.k); + + auto removed_indices = raft::make_device_vector(handle_, filter_offset); + thrust::sequence(raft::resource::get_thrust_policy(handle_), + thrust::device_pointer_cast(removed_indices.data_handle()), + thrust::device_pointer_cast(removed_indices.data_handle() + filter_offset)); + cuvs::core::bitset bitset(handle_, removed_indices.view(), ps.num_db_vecs); + auto filter = cuvs::neighbors::filtering::bitset_filter(bitset.view()); + + auto queries = raft::make_device_matrix_view( + search_queries.data(), ps.num_queries, ps.dim); + cuvs::neighbors::ivf_rabitq::search( + handle_, ps.search_params, index, queries, indices.view(), distances.view(), filter); + + std::vector host_indices(queries_size); + std::vector host_distances(queries_size); + raft::update_host(host_indices.data(), indices.data_handle(), queries_size, stream_); + raft::update_host(host_distances.data(), distances.data_handle(), queries_size, stream_); + raft::resource::sync_stream(handle_); + + const uint32_t expected_valid = std::min(ps.filter_pass_count, ps.k); + for (uint32_t query_ix = 0; query_ix < ps.num_queries; ++query_ix) { + uint32_t valid_count = 0; + uint32_t oob_count = 0; + std::vector seen(ps.filter_pass_count, false); + for (uint32_t neighbor_ix = 0; neighbor_ix < ps.k; ++neighbor_ix) { + const size_t flat_ix = size_t{query_ix} * ps.k + neighbor_ix; + const IdxT id = host_indices[flat_ix]; + if (id == kOutOfBoundsRecord) { + ++oob_count; + ASSERT_TRUE(std::isinf(host_distances[flat_ix])) << ps; + continue; + } + ASSERT_GE(id, filter_offset) << ps; + ASSERT_LT(id, static_cast(ps.num_db_vecs)) << ps; + const size_t allowed_ix = static_cast(id - filter_offset); + ASSERT_FALSE(seen[allowed_ix]) << ps; + seen[allowed_ix] = true; + ++valid_count; + } + ASSERT_EQ(valid_count, expected_valid) << ps; + ASSERT_EQ(oob_count, ps.k - expected_valid) << ps; + } + + if (ps.filter_pass_count < ps.k) { + RAFT_LOG_INFO("Filter validation = %u valid, %u out-of-bounds results per query.", + expected_valid, + ps.k - expected_valid); + return; + } + + double compression_ratio = sizeof(DataT) * 8 / ps.index_params.bits_per_dim; + const double min_recall = ps.min_recall.value_or(0.5); + ASSERT_TRUE(cuvs::neighbors::eval_neighbours(indices_ref, + host_indices, + distances_ref, + host_distances, + ps.num_queries, + ps.k, + 0.0001 * compression_ratio, + min_recall)) + << ps; + } + void SetUp() override // NOLINT { gen_data(); @@ -263,6 +360,9 @@ class ivf_rabitq_test : public ::testing::TestWithParam { std::vector distances_ref; // NOLINT }; +template +class ivf_rabitq_filter_test : public ivf_rabitq_test {}; + /* Test cases */ using test_cases_t = std::vector; @@ -343,6 +443,76 @@ inline auto var_bits_per_dim() -> test_cases_t }); } +inline auto filter_cases() -> test_cases_t +{ + test_cases_t cases; + const std::vector modes{ + search_mode::LUT16, search_mode::LUT32, search_mode::QUANT4, search_mode::QUANT8}; + const std::vector metrics{ + cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct}; + // Keep the existing 70.7%-pass coverage and add an exact 1%-pass dataset. + // For the latter, k=65 exercises the non-block path without the largely + // forced overlap produced by selecting 96 of only 100 eligible vectors. + for (auto [num_db_vecs, filter_pass_count] : {std::pair{1024u, 724u}, std::pair{10000u, 100u}}) { + const uint32_t non_block_k = filter_pass_count == 100 ? 65u : 96u; + for (auto mode : modes) { + for (uint32_t bits_per_dim : {1u, 3u}) { + for (uint32_t k : {10u, non_block_k}) { + for (auto metric : metrics) { + ivf_rabitq_inputs x; + x.num_db_vecs = num_db_vecs; + x.num_queries = 4; + x.k = k; + x.filter_pass_count = filter_pass_count; + x.index_params.n_lists = 16; + x.index_params.bits_per_dim = bits_per_dim; + x.index_params.metric = metric; + x.signed_data = metric == cuvs::distance::DistanceType::InnerProduct; + // Large-k L2 searches should recover nearly all eligible neighbors. At small k, + // account for coarser 1-bit quantization while requiring higher recall when only 1% + // of the database can pass the filter. InnerProduct uses the conservative bounds from + // the existing metric sweep until filtered recall measurements are available. + if (metric == cuvs::distance::DistanceType::InnerProduct) { + x.min_recall = bits_per_dim == 1 ? 0.30 : 0.50; + } else if (k > 64) { + x.min_recall = 0.90; + } else if (filter_pass_count == 100) { + x.min_recall = bits_per_dim == 1 ? 0.50 : 0.85; + } else { + x.min_recall = bits_per_dim == 1 ? 0.45 : 0.70; + } + x.search_params.n_probes = 16; + x.search_params.mode = mode; + cases.push_back(x); + } + } + } + } + } + // Exercise under-filled output for both block-sort and non-block-sort search. + for (auto mode : modes) { + for (uint32_t bits_per_dim : {1u, 3u}) { + for (uint32_t k : {10u, 65u}) { + for (auto metric : metrics) { + ivf_rabitq_inputs x; + x.num_db_vecs = 1024; + x.num_queries = 4; + x.k = k; + x.filter_pass_count = 5; + x.index_params.n_lists = 16; + x.index_params.bits_per_dim = bits_per_dim; + x.index_params.metric = metric; + x.signed_data = metric == cuvs::distance::DistanceType::InnerProduct; + x.search_params.n_probes = 16; + x.search_params.mode = mode; + cases.push_back(x); + } + } + } + } + return cases; +} + inline auto var_search_mode() -> test_cases_t { ivf_rabitq_inputs dflt; diff --git a/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu b/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu index 3106be9f85..c53570758c 100644 --- a/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu +++ b/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu @@ -7,7 +7,12 @@ namespace cuvs::neighbors::ivf_rabitq { -using f32_f32_i64 = ivf_rabitq_test; +using f32_f32_i64 = ivf_rabitq_test; +using f32_f32_i64_filter = ivf_rabitq_filter_test; + +TEST_P(f32_f32_i64_filter, build_serialize_search) { this->run_filter(); } + +INSTANTIATE(f32_f32_i64_filter, filter_cases()); TEST_BUILD_SERIALIZE_SEARCH(f32_f32_i64) TEST_BUILD_HOST_INPUT_SERIALIZE_SEARCH(f32_f32_i64) From f909356aad9676a7abc99f1a881d19bb7024eaa5 Mon Sep 17 00:00:00 2001 From: James Xia Date: Mon, 18 May 2026 15:46:59 -0700 Subject: [PATCH 02/10] Add prefiltered ground truth generation to cuvs_bench Extends generate_groundtruth with optional bitset prefiltering so that exact-NN ground truth can be computed over a filtered subset of the dataset. New function create_bitset_filter(n_samples, filter_reject_rate) returns a packed uint32 numpy array (LSB-first, matching cuVS bitset layout) where bit i is set iff vector i passes the filter. Uses a modulo-1000 bucket scheme: vector i passes when i % 1000 >= round(filter_reject_rate * 1000), giving a reject rate within 0.1% of the requested value. calc_truth gains an optional bitset parameter: - GPU path: slices the packed words for each batch, transfers to device, and passes as a cuVS bitset prefilter to brute_force.search - CPU path: unpacks the same words to a boolean accept_mask and masks rejected distances to inf before argpartition Two new mutually exclusive CLI arguments: - --bitset : load a pre-saved .npy uint32 bitset from file - --filter_reject_rate : generate a bitset in memory via create_bitset_filter (cherry picked from commit 09281c15effe310739914e52acf8f9bdf43a373f) --- .../generate_groundtruth/__main__.py | 134 +++++++++++++++++- 1 file changed, 129 insertions(+), 5 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py b/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py index d527a24fc6..9908399c96 100644 --- a/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py +++ b/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py @@ -84,6 +84,7 @@ def force_fallback_to_numpy(): from rmm.allocators.cupy import rmm_cupy_allocator from cuvs.common import Resources + from cuvs.neighbors import filters from cuvs.neighbors.brute_force import build, search except ImportError: # RMM is available, cupy is available, but cuVS is not @@ -128,7 +129,37 @@ def choose_random_queries_with_jitter(dataset, n_queries, seed=12345): return add_jitter(sampled, rng, normalize=False) -def cpu_search(dataset, queries, k, metric="squeclidean"): +def create_bitset_filter(n_samples, filter_reject_rate): + """ + Creates a packed uint32 bitset where bit i is set iff vector i passes the + filter. Uses a modulo-1000 bucket scheme: vector i passes when + ``i % 1000 >= round(filter_reject_rate * 1000)``, giving a reject rate + within 0.1% of the requested value. + + Parameters + ---------- + n_samples : int + Number of vectors in the dataset. + filter_reject_rate : float + Fraction of vectors to reject, in [0.0, 1.0). + + Returns + ------- + numpy.ndarray + Packed uint32 array of shape ``(ceil(n_samples / 32),)``. + """ + import numpy as np + + fail_buckets = round(filter_reject_rate * 1000) + n_padded = ((n_samples + 31) // 32) * 32 + bool_mask = np.zeros(n_padded, dtype=bool) + bool_mask[:n_samples] = (np.arange(n_samples) % 1000) >= fail_buckets + # Pack with little-endian bit order: bit j maps to bit (j%32) of uint32 + # word (j//32), LSB first — matching cuVS bitset layout. + return np.packbits(bool_mask, bitorder="little").view(np.uint32) + + +def cpu_search(dataset, queries, k, metric="squeclidean", accept_mask=None): """ Find the k nearest neighbors for each query point in the dataset using the specified metric. @@ -145,6 +176,9 @@ def cpu_search(dataset, queries, k, metric="squeclidean"): metric : str, optional The distance metric to use. Can be 'squeclidean' or 'inner_product'. Default is 'squeclidean'. + accept_mask : numpy.ndarray, optional + Boolean array of shape (n_samples,). Where False, the corresponding + dataset vector is excluded from results. Returns ------- @@ -161,6 +195,9 @@ def cpu_search(dataset, queries, k, metric="squeclidean"): diff = queries[:, xp.newaxis, :] - dataset[xp.newaxis, :, :] dist_sq = xp.sum(diff**2, axis=2) # Shape: (n_queries, n_samples) + if accept_mask is not None: + dist_sq[:, ~accept_mask] = xp.inf + indices = xp.argpartition(dist_sq, kth=k - 1, axis=1)[:, :k] distances = xp.take_along_axis(dist_sq, indices, axis=1) @@ -173,6 +210,9 @@ def cpu_search(dataset, queries, k, metric="squeclidean"): queries, dataset.T ) # Shape: (n_queries, n_samples) + if accept_mask is not None: + similarities[:, ~accept_mask] = -xp.inf + neg_similarities = -similarities indices = xp.argpartition(neg_similarities, kth=k - 1, axis=1)[:, :k] distances = xp.take_along_axis(similarities, indices, axis=1) @@ -192,7 +232,27 @@ def cpu_search(dataset, queries, k, metric="squeclidean"): return distances, indices -def calc_truth(dataset, queries, k, metric="sqeuclidean"): +def calc_truth(dataset, queries, k, metric="sqeuclidean", bitset=None): + """ + Calculate exact nearest neighbors, optionally with a prefilter. + + Parameters + ---------- + dataset : array-like + Dataset of shape (n_samples, n_features). + queries : array-like + Queries of shape (n_queries, n_features). + k : int + Number of neighbors. + metric : str + Distance metric. + bitset : numpy.ndarray, optional + Packed uint32 array of shape (ceil(n_samples / 32),) as returned by + :func:`create_bitset_filter`. Bit i set means vector i passes the + filter. When None, all vectors are considered. + """ + import numpy as np + n_samples = dataset.shape[0] n = 500000 # batch size for processing neighbors i = 0 @@ -211,10 +271,28 @@ def calc_truth(dataset, queries, k, metric="sqeuclidean"): if gpu_system: index = build(X, metric=metric, resources=resources) - D, Ind = search(index, queries, k, resources=resources) + prefilter = None + if bitset is not None: + word_start = i // 32 + word_end = (i + n_batch + 31) // 32 + batch_words = xp.asarray(bitset[word_start:word_end]) + prefilter = filters.from_bitset(batch_words) + D, Ind = search( + index, queries, k, resources=resources, prefilter=prefilter + ) resources.sync() else: - D, Ind = cpu_search(X, queries, k, metric=metric) + accept_mask = None + if bitset is not None: + word_start = i // 32 + word_end = (i + n_batch + 31) // 32 + batch_bytes = bitset[word_start:word_end].view(np.uint8) + accept_mask = np.unpackbits(batch_bytes, bitorder="little")[ + :n_batch + ].astype(bool) + D, Ind = cpu_search( + X, queries, k, metric=metric, accept_mask=accept_mask + ) D, Ind = xp.asarray(D), xp.asarray(Ind) Ind = offset_neighbor_indices(Ind, i, n_samples) @@ -268,6 +346,16 @@ def main(): # Jittered queries (following the logic of cuvs_bench.synthesize_dataset) python -m cuvs_bench.generate_groundtruth /dataset/base.fbin \ --output=groundtruth_dir --queries=random-jitter --n_queries=10000 + + # Prefiltered ground truth using a saved bitset file + python -m cuvs_bench.generate_groundtruth --dataset /dataset/base.\ +fbin --output=groundtruth_dir --queries=/dataset/query.fbin \ +--bitset=/dataset/filter.npy + + # Prefiltered ground truth generated on-the-fly from a reject rate + python -m cuvs_bench.generate_groundtruth --dataset /dataset/base.\ +fbin --output=groundtruth_dir --queries=/dataset/query.fbin \ +--filter_reject_rate=0.1 """, formatter_class=argparse.RawDescriptionHelpFormatter, ) @@ -336,6 +424,25 @@ def main(): " commonly used with cuVS are 'sqeuclidean' and 'inner_product'", ) + filter_group = parser.add_mutually_exclusive_group() + filter_group.add_argument( + "--bitset", + type=str, + default=None, + help="Path to a .npy file containing a packed uint32 prefilter " + "bitset of shape (ceil(n_samples / 32),). Bit i set means vector i " + "passes the filter. Mutually exclusive with --filter_reject_rate.", + ) + filter_group.add_argument( + "--filter_reject_rate", + type=float, + default=None, + help="Fraction of vectors to reject, in [0.0, 1.0). Generates a " + "bitset using a modulo-1000 bucket scheme (vector i passes when " + "i %% 1000 >= round(filter_reject_rate * 1000)). Mutually exclusive " + "with --bitset.", + ) + if len(sys.argv) == 1: parser.print_help() sys.exit(1) @@ -351,6 +458,7 @@ def main(): args.dataset, args.dtype, shape=(args.rows, args.cols) ) n_features = dataset.shape[1] + n_samples = dataset.shape[0] dtype = dataset.dtype print( @@ -389,8 +497,24 @@ def main(): print("Reading queries from file", args.queries) queries = memmap_bin_file(args.queries, dtype) + # Resolve prefilter bitset. + bitset = None + if args.bitset is not None: + import numpy as np + + print("Loading prefilter bitset from", args.bitset) + bitset = np.load(args.bitset) + elif args.filter_reject_rate is not None: + print( + f"Generating prefilter bitset for filter_reject_rate=" + f"{args.filter_reject_rate}" + ) + bitset = create_bitset_filter(n_samples, args.filter_reject_rate) + print("Calculating true nearest neighbors") - distances, indices = calc_truth(dataset, queries, args.k, args.metric) + distances, indices = calc_truth( + dataset, queries, args.k, args.metric, bitset=bitset + ) n_base = dataset.shape[0] write_groundtruth_neighbors( From 373008a5e97ce473c4b266e93bcad180586b4cd4 Mon Sep 17 00:00:00 2001 From: James Xia Date: Wed, 27 May 2026 06:53:45 -0700 Subject: [PATCH 03/10] Persist prefilter bitset and switch bitset I/O to cuVS .bin format The --filter_reject_rate path previously generated a bitset in memory, used it to compute ground truth, then discarded it. The benchmark side (cuvs-bench C++) needs to run searches against the *exact same* bitset that was used to filter the GT, otherwise recall is computed against a stale GT. Persisting the generated bitset is the simplest way to keep producer and consumer in sync. The on-disk format is also switched from .npy to the cuVS .bin format used by the rest of the script (base / query / GT files). .npy is a Python-only format and cannot be read by the C++ blob loader. A single .bin format gives one round-trip across the Python producer, the Python --bitset loader, and the C++ benchmark consumer. With the bitset now persisted, create_bitset_filter no longer needs to be deterministic for reproducibility. Switch it to true Bernoulli sampling via numpy's default_rng to match the C++ generate_bernoulli helper in cpp/bench/ann/src/common/dataset.hpp -- one less subtle divergence between the two sides, and avoids the regular-pattern artifact of the old modulo-1000 scheme. - --filter_reject_rate now writes the generated bitset to /groundtruth.filter.bin via write_bin, with shape (ceil(N/32), 1) -- matching the C++ blob layout used by cuvs-bench's dataset loader. - --bitset loader switched from np.load to memmap_bin_file with dtype=np.uint32; the (ceil(N/32), 1) memmap is raveled to 1-D for use in calc_truth. - create_bitset_filter switched to numpy Bernoulli sampling (rng.random(n) >= reject_rate) for parity with the C++ side. (cherry picked from commit deabe0839cb052c9f76b5fbcfe2b39d4640e1e47) --- .../generate_groundtruth/__main__.py | 36 +++++++++++-------- 1 file changed, 21 insertions(+), 15 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py b/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py index 9908399c96..f489b75dc6 100644 --- a/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py +++ b/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py @@ -132,16 +132,16 @@ def choose_random_queries_with_jitter(dataset, n_queries, seed=12345): def create_bitset_filter(n_samples, filter_reject_rate): """ Creates a packed uint32 bitset where bit i is set iff vector i passes the - filter. Uses a modulo-1000 bucket scheme: vector i passes when - ``i % 1000 >= round(filter_reject_rate * 1000)``, giving a reject rate - within 0.1% of the requested value. + filter. Each vector independently passes with probability + ``1.0 - filter_reject_rate`` (Bernoulli), matching the C++ benchmark's + ``generate_bernoulli`` helper in ``cpp/bench/ann/src/common/dataset.hpp``. Parameters ---------- n_samples : int Number of vectors in the dataset. filter_reject_rate : float - Fraction of vectors to reject, in [0.0, 1.0). + Per-vector rejection probability, in [0.0, 1.0). Returns ------- @@ -150,10 +150,10 @@ def create_bitset_filter(n_samples, filter_reject_rate): """ import numpy as np - fail_buckets = round(filter_reject_rate * 1000) n_padded = ((n_samples + 31) // 32) * 32 bool_mask = np.zeros(n_padded, dtype=bool) - bool_mask[:n_samples] = (np.arange(n_samples) % 1000) >= fail_buckets + rng = np.random.default_rng() + bool_mask[:n_samples] = rng.random(n_samples) >= filter_reject_rate # Pack with little-endian bit order: bit j maps to bit (j%32) of uint32 # word (j//32), LSB first — matching cuVS bitset layout. return np.packbits(bool_mask, bitorder="little").view(np.uint32) @@ -350,9 +350,12 @@ def main(): # Prefiltered ground truth using a saved bitset file python -m cuvs_bench.generate_groundtruth --dataset /dataset/base.\ fbin --output=groundtruth_dir --queries=/dataset/query.fbin \ ---bitset=/dataset/filter.npy +--bitset=/dataset/groundtruth.filter.bin - # Prefiltered ground truth generated on-the-fly from a reject rate + # Generate a prefilter bitset on the fly from a reject rate, use it to + # compute the ground truth, and save the bitset to disk so the benchmark + # can later run searches against the exact same filter. The bitset is + # written to /groundtruth.filter.bin alongside the GT files. python -m cuvs_bench.generate_groundtruth --dataset /dataset/base.\ fbin --output=groundtruth_dir --queries=/dataset/query.fbin \ --filter_reject_rate=0.1 @@ -429,18 +432,19 @@ def main(): "--bitset", type=str, default=None, - help="Path to a .npy file containing a packed uint32 prefilter " - "bitset of shape (ceil(n_samples / 32),). Bit i set means vector i " + help="Path to a cuVS .bin file containing a packed uint32 prefilter " + "bitset of shape (ceil(n_samples / 32), 1). Bit i set means vector i " "passes the filter. Mutually exclusive with --filter_reject_rate.", ) filter_group.add_argument( "--filter_reject_rate", type=float, default=None, - help="Fraction of vectors to reject, in [0.0, 1.0). Generates a " - "bitset using a modulo-1000 bucket scheme (vector i passes when " - "i %% 1000 >= round(filter_reject_rate * 1000)). Mutually exclusive " - "with --bitset.", + help="Per-vector rejection probability, in [0.0, 1.0). Generates a " + "Bernoulli bitset where each vector independently passes with " + "probability (1 - filter_reject_rate). The generated bitset is " + "saved to /groundtruth.filter.bin. Mutually exclusive with " + "--bitset.", ) if len(sys.argv) == 1: @@ -503,13 +507,15 @@ def main(): import numpy as np print("Loading prefilter bitset from", args.bitset) - bitset = np.load(args.bitset) + bitset = np.asarray(memmap_bin_file(args.bitset, np.uint32)).ravel() elif args.filter_reject_rate is not None: print( f"Generating prefilter bitset for filter_reject_rate=" f"{args.filter_reject_rate}" ) bitset = create_bitset_filter(n_samples, args.filter_reject_rate) + bitset_filename = os.path.join(args.output, "groundtruth.filter.bin") + write_bin(bitset_filename, bitset.reshape(-1, 1)) print("Calculating true nearest neighbors") distances, indices = calc_truth( From b3e1e8e3d1aca3ec24b67f18de3661e263ee36f1 Mon Sep 17 00:00:00 2001 From: James Xia Date: Wed, 27 May 2026 12:39:07 -0700 Subject: [PATCH 04/10] Add filter_bitset_file dataset option to cuvs-bench Loads a precomputed bitset from disk (cuVS .bin format) and feeds it through the existing filter plumbing in the IVF-Flat / CAGRA / brute- force wrappers. Pairs with the Python generate_groundtruth output at /groundtruth.filter.bin so the benchmark searches against the exact bitset that was used to compute the prefiltered ground truth. - conf.hpp: new optional 'filter_bitset_file' field in dataset_conf; rejected at parse time if 'filtering_rate' is also set. - dataset.hpp: dataset<> ctor accepts the path and loads via the same blob(file) machinery used for base/query sets. When this path is used, the ground_truth_map is told the GT is pre-filtered; it then counts bitset rejections during GT-map construction and logs a warning if any occur (signals a bitset/GT mismatch). - benchmark.hpp: forwards dataset_conf.filter_bitset_file at the dataset construction site. (cherry picked from commit e2d014acd8e61b91f587611d6696160e24a0deaa) --- cpp/bench/ann/src/common/benchmark.hpp | 19 +++++----- cpp/bench/ann/src/common/conf.hpp | 10 ++++++ cpp/bench/ann/src/common/dataset.hpp | 50 +++++++++++++++++++++----- 3 files changed, 62 insertions(+), 17 deletions(-) diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp index a588b1e2a6..f29859e3ff 100644 --- a/cpp/bench/ann/src/common/benchmark.hpp +++ b/cpp/bench/ann/src/common/benchmark.hpp @@ -537,15 +537,16 @@ void dispatch_benchmark(std::string cmdline, auto base_file = dataset_conf.base_file; auto query_file = dataset_conf.query_file; auto gt_file = dataset_conf.groundtruth_neighbors_file; - auto dataset = - std::make_shared>(dataset_conf.name, - base_file, - dataset_conf.subset_first_row, - dataset_conf.subset_size, - query_file, - dataset_conf.distance, - gt_file, - search_mode ? dataset_conf.filtering_rate : std::nullopt); + auto dataset = std::make_shared>( + dataset_conf.name, + base_file, + dataset_conf.subset_first_row, + dataset_conf.subset_size, + query_file, + dataset_conf.distance, + gt_file, + search_mode ? dataset_conf.filtering_rate : std::nullopt, + search_mode ? dataset_conf.filter_bitset_file : std::nullopt); ::benchmark::AddCustomContext("dataset", dataset_conf.name); ::benchmark::AddCustomContext("distance", dataset_conf.distance); std::vector& indices = conf.get_indices(); diff --git a/cpp/bench/ann/src/common/conf.hpp b/cpp/bench/ann/src/common/conf.hpp index afc7bc0a1f..66fbc52fc0 100644 --- a/cpp/bench/ann/src/common/conf.hpp +++ b/cpp/bench/ann/src/common/conf.hpp @@ -8,6 +8,7 @@ #include #include +#include #include #include #include @@ -44,6 +45,7 @@ class configuration { std::string dtype; std::optional filtering_rate{std::nullopt}; + std::optional filter_bitset_file{std::nullopt}; }; [[nodiscard]] inline auto get_dataset_conf() const -> const dataset_conf& @@ -89,6 +91,14 @@ class configuration { if (conf.contains("filtering_rate")) { dataset_conf_.filtering_rate.emplace(conf.at("filtering_rate")); } + if (conf.contains("filter_bitset_file")) { + dataset_conf_.filter_bitset_file = combine_path(data_prefix, conf.at("filter_bitset_file")); + } + if (dataset_conf_.filtering_rate.has_value() && dataset_conf_.filter_bitset_file.has_value()) { + throw std::invalid_argument( + "dataset config must set at most one of 'filtering_rate' or " + "'filter_bitset_file'; got both"); + } if (conf.contains("groundtruth_neighbors_file")) { dataset_conf_.groundtruth_neighbors_file = diff --git a/cpp/bench/ann/src/common/dataset.hpp b/cpp/bench/ann/src/common/dataset.hpp index 4dc43c343c..e84b8d3b5c 100644 --- a/cpp/bench/ann/src/common/dataset.hpp +++ b/cpp/bench/ann/src/common/dataset.hpp @@ -41,7 +41,8 @@ struct ground_truth_map { explicit ground_truth_map(std::string file_name, uint32_t n_queries, - std::optional>& filter_bitset) + std::optional>& filter_bitset, + bool gt_is_prefiltered = false) : gt_maps_(n_queries) { // Eagerly iterate over and optionally filter the ground truth set to build gt_maps_ for up to @@ -59,12 +60,19 @@ struct ground_truth_map { */ auto ground_truth_set = blob(file_name); max_k_ = ground_truth_set.n_cols(); - auto filter = [&](T i) -> bool { + // Counts ground-truth ids that the bitset rejects. When the GT was + // produced with a matching pre-filter (gt_is_prefiltered), every id + // is expected to pass, so any nonzero count signals a bitset/GT + // mismatch and is reported as a warning after the workers join. + std::atomic filtered_out_count{0}; + auto filter = [&](T i) -> bool { if (!filter_bitset.has_value()) { return true; } // bitset is `32 = bitset_carrier_type * 8` times more dense than the data // use bitwise arithmetic to get the `row_id` and correct bit pos in the `word` - auto word = filter_bitset->data(MemoryType::kHostMmap)[i >> 5]; - return word & (1 << (i & 31)); + auto word = filter_bitset->data(MemoryType::kHostMmap)[i >> 5]; + bool passes = word & (1 << (i & 31)); + if (!passes) { filtered_out_count.fetch_add(1, std::memory_order_relaxed); } + return passes; }; // Avoid CPU oversubscription when parallelizing recall calculation loop int num_map_building_worker_threads = @@ -113,6 +121,18 @@ struct ground_truth_map { for (auto& worker : gt_map_building_workers) { worker.join(); } + if (gt_is_prefiltered) { + auto drops = filtered_out_count.load(std::memory_order_relaxed); + if (drops > 0) { + log_warn( + "ground truth file '%s' was declared pre-filtered, but %zu of its " + "neighbor ids are rejected by the supplied filter bitset. This " + "indicates the bitset and ground truth files are mismatched; " + "recall numbers will be unreliable.", + file_name.c_str(), + drops); + } + } } [[nodiscard]] auto max_k() const -> uint32_t { return max_k_; } @@ -179,13 +199,24 @@ struct dataset { std::string query_file, std::string distance, std::optional groundtruth_neighbors_file, - std::optional filtering_rate = std::nullopt) + std::optional filtering_rate = std::nullopt, + std::optional filter_bitset_file = std::nullopt) : name_{std::move(name)}, distance_{std::move(distance)}, base_set_{base_file, subset_first_row, subset_size}, query_set_{query_file} { - if (filtering_rate.has_value()) { + if (filtering_rate.has_value() && filter_bitset_file.has_value()) { + throw std::invalid_argument( + "dataset: at most one of 'filtering_rate' or 'filter_bitset_file' may be set"); + } + + if (filter_bitset_file.has_value()) { + // Load a pre-generated bitset from disk. The producer + // (cuvs_bench.generate_groundtruth) writes it in cuVS .bin format with + // shape (ceil(n_rows / 32), 1), matching the in-memory layout below. + filter_bitset_.emplace(blob{filter_bitset_file.value()}); + } else if (filtering_rate.has_value()) { // Generate a random bitset for filtering auto n_rows = static_cast(subset_size) + static_cast(subset_first_row); if (subset_size == 0) { @@ -203,8 +234,11 @@ struct dataset { } if (groundtruth_neighbors_file.has_value()) { - ground_truth_map_.emplace(ground_truth_map{ - groundtruth_neighbors_file.value(), query_set_.n_rows(), filter_bitset_}); + ground_truth_map_.emplace( + ground_truth_map{groundtruth_neighbors_file.value(), + query_set_.n_rows(), + filter_bitset_, + /*gt_is_prefiltered=*/filter_bitset_file.has_value()}); } } From 27b50bc61fc6bb49177807609a2a1cddfc2274b6 Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 28 Jul 2026 17:34:38 -0700 Subject: [PATCH 05/10] Fix dataset argument in prefiltered ground truth usage examples The two prefilter examples in the --help epilog were written against an older CLI that took the dataset path as --dataset. It has been a positional argument for some time, and the surrounding examples were corrected earlier in this series, leaving these two as the only stale ones. Copying either of them would fail with an unrecognized-argument error. Also drop the mid-word "base.\fbin" line continuation so all six examples wrap the same way. --- .../cuvs_bench/generate_groundtruth/__main__.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py b/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py index f489b75dc6..9c4a3ee63e 100644 --- a/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py +++ b/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py @@ -348,16 +348,16 @@ def main(): --output=groundtruth_dir --queries=random-jitter --n_queries=10000 # Prefiltered ground truth using a saved bitset file - python -m cuvs_bench.generate_groundtruth --dataset /dataset/base.\ -fbin --output=groundtruth_dir --queries=/dataset/query.fbin \ + python -m cuvs_bench.generate_groundtruth /dataset/base.fbin \ +--output=groundtruth_dir --queries=/dataset/query.fbin \ --bitset=/dataset/groundtruth.filter.bin # Generate a prefilter bitset on the fly from a reject rate, use it to # compute the ground truth, and save the bitset to disk so the benchmark # can later run searches against the exact same filter. The bitset is # written to /groundtruth.filter.bin alongside the GT files. - python -m cuvs_bench.generate_groundtruth --dataset /dataset/base.\ -fbin --output=groundtruth_dir --queries=/dataset/query.fbin \ + python -m cuvs_bench.generate_groundtruth /dataset/base.fbin \ +--output=groundtruth_dir --queries=/dataset/query.fbin \ --filter_reject_rate=0.1 """, formatter_class=argparse.RawDescriptionHelpFormatter, From 3503226f3ed443ce7dc03c30bd94c1fe586c7085 Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 28 Jul 2026 17:39:48 -0700 Subject: [PATCH 06/10] Fix metric name typo in generate_groundtruth CPU search --- .../cuvs_bench/generate_groundtruth/__main__.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py b/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py index 9c4a3ee63e..7a61f72114 100644 --- a/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py +++ b/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py @@ -159,7 +159,7 @@ def create_bitset_filter(n_samples, filter_reject_rate): return np.packbits(bool_mask, bitorder="little").view(np.uint32) -def cpu_search(dataset, queries, k, metric="squeclidean", accept_mask=None): +def cpu_search(dataset, queries, k, metric="sqeuclidean", accept_mask=None): """ Find the k nearest neighbors for each query point in the dataset using the specified metric. @@ -174,8 +174,8 @@ def cpu_search(dataset, queries, k, metric="squeclidean", accept_mask=None): k : int The number of nearest neighbors to find. metric : str, optional - The distance metric to use. Can be 'squeclidean' or 'inner_product'. - Default is 'squeclidean'. + The distance metric to use. Can be 'sqeuclidean' or 'inner_product'. + Default is 'sqeuclidean'. accept_mask : numpy.ndarray, optional Boolean array of shape (n_samples,). Where False, the corresponding dataset vector is excluded from results. @@ -184,14 +184,14 @@ def cpu_search(dataset, queries, k, metric="squeclidean", accept_mask=None): ------- distances : numpy.ndarray An array of shape (n_queries, k) containing the distances - (for 'squeclidean') or similarities + (for 'sqeuclidean') or similarities (for 'inner_product') to the k nearest neighbors for each query. indices : numpy.ndarray An array of shape (n_queries, k) containing the indices of the k nearest neighbors in the dataset for each query. """ - if metric == "squeclidean": + if metric == "sqeuclidean": diff = queries[:, xp.newaxis, :] - dataset[xp.newaxis, :, :] dist_sq = xp.sum(diff**2, axis=2) # Shape: (n_queries, n_samples) @@ -222,7 +222,7 @@ def cpu_search(dataset, queries, k, metric="squeclidean", accept_mask=None): else: raise ValueError( "Unsupported metric in cuvs-bench-cpu. " - "Use 'squeclidean' or 'inner_product' or use the GPU package" + "Use 'sqeuclidean' or 'inner_product', or use the GPU package " "to use any distance supported by cuVS." ) From 916545e825eb25bdefbc272220cd0bdf99461e5b Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 28 Jul 2026 19:04:18 -0700 Subject: [PATCH 07/10] Validate the prefilter bitset and re-base it correctly per batch Two related defects in the prefiltered ground truth path. Batch re-basing was only correct for 32-aligned offsets. calc_truth builds an index per batch holding rows [i, i + n_batch), so the prefilter handed to it must be re-based to that batch. Slicing the bitset by word starts at global bit (i / 32) * 32, which equals i only when i is a multiple of 32, making correctness depend on the batch size constant. Replace the word slicing with slice_bitset(), which unpacks, realigns by i % 32, and repacks with a zero-padded tail, so any offset works. The GPU and CPU branches now share one implementation instead of duplicating the arithmetic. Neither side checked that the bitset covers the dataset. A short file was silently accepted, since NumPy clamps the out-of-range slice, and the uncovered rows were read out of bounds. The C++ consumer has the same hole, where the overrun is past the end of an mmap. Both sides now check the covered row count up front and fail with the actual numbers. The check is word-granular, since the file records only a word count. A bitset short by fewer than 32 rows needs the same number of words and is indistinguishable from a correct one; those rows read padding zeros, so they are rejected rather than read out of bounds. The C++ half is restructured so both filter branches derive the row count from one place. It stays behind the has_value() guard, so the base_set_size() header read is still skipped when no filter is configured. --- cpp/bench/ann/src/common/dataset.hpp | 45 ++++++---- .../generate_groundtruth/__main__.py | 87 +++++++++++++++---- 2 files changed, 101 insertions(+), 31 deletions(-) diff --git a/cpp/bench/ann/src/common/dataset.hpp b/cpp/bench/ann/src/common/dataset.hpp index e84b8d3b5c..a439c5a899 100644 --- a/cpp/bench/ann/src/common/dataset.hpp +++ b/cpp/bench/ann/src/common/dataset.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -211,26 +211,41 @@ struct dataset { "dataset: at most one of 'filtering_rate' or 'filter_bitset_file' may be set"); } - if (filter_bitset_file.has_value()) { - // Load a pre-generated bitset from disk. The producer - // (cuvs_bench.generate_groundtruth) writes it in cuVS .bin format with - // shape (ceil(n_rows / 32), 1), matching the in-memory layout below. - filter_bitset_.emplace(blob{filter_bitset_file.value()}); - } else if (filtering_rate.has_value()) { - // Generate a random bitset for filtering + if (filtering_rate.has_value() || filter_bitset_file.has_value()) { auto n_rows = static_cast(subset_size) + static_cast(subset_first_row); if (subset_size == 0) { // Read the base set size as a last resort only - for better laziness n_rows = base_set_size(); } auto bitset_size = (n_rows - 1) / kBitsPerCarrierValue + 1; - blob_file bitset_blob_file{static_cast(bitset_size), 1}; - blob_mmap bitset_blob{ - std::move(bitset_blob_file), false, HugePages::kDisable}; - generate_bernoulli(const_cast(bitset_blob.data()), - bitset_size, - 1.0 - filtering_rate.value()); - filter_bitset_.emplace(std::move(bitset_blob)); + + if (filter_bitset_file.has_value()) { + // Load a pre-generated bitset from disk. The producer + // (cuvs_bench.generate_groundtruth) writes it in cuVS .bin format with + // shape (ceil(n_rows / 32), 1), matching the in-memory layout below. + filter_bitset_.emplace(blob{filter_bitset_file.value()}); + // Reject a bitset that does not cover the whole dataset up front: the + // filter is indexed by absolute row id, so a short file would be read + // out of bounds (past the end of the mmap) rather than fail cleanly. + auto loaded_words = static_cast(filter_bitset_->n_rows()) * + static_cast(filter_bitset_->n_cols()); + if (loaded_words < bitset_size) { + throw std::invalid_argument("dataset: filter bitset file '" + filter_bitset_file.value() + + "' holds " + std::to_string(loaded_words) + " words, but " + + std::to_string(bitset_size) + " are needed to cover " + + std::to_string(n_rows) + + " rows; the bitset and dataset do not match"); + } + } else { + // Generate a random bitset for filtering + blob_file bitset_blob_file{static_cast(bitset_size), 1}; + blob_mmap bitset_blob{ + std::move(bitset_blob_file), false, HugePages::kDisable}; + generate_bernoulli(const_cast(bitset_blob.data()), + bitset_size, + 1.0 - filtering_rate.value()); + filter_bitset_.emplace(std::move(bitset_blob)); + } } if (groundtruth_neighbors_file.has_value()) { diff --git a/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py b/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py index 7a61f72114..7d052030db 100644 --- a/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py +++ b/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py @@ -159,6 +159,64 @@ def create_bitset_filter(n_samples, filter_reject_rate): return np.packbits(bool_mask, bitorder="little").view(np.uint32) +def slice_bitset(bitset, start, count): + """ + Extract rows ``[start, start + count)`` of a packed bitset. + + Returns ``(mask, words)`` where ``mask`` is a boolean array of length + ``count`` and ``words`` is a packed uint32 bitset whose *bit 0* is row + ``start``. Ground truth is computed in batches against an index holding + only the batch's rows, so the prefilter handed to that index has to be + re-based to the batch rather than to the whole dataset. + + Slicing by word (``bitset[start // 32:]``) would only be correct when + ``start`` is a multiple of 32. This unpacks, realigns by ``start % 32``, + and repacks, so any ``start`` works. The repacked tail is zero-padded out + to a whole word; those bits sit past the end of the batch index and are + never consulted, and zero can only ever reject, never admit a row that + does not exist. + + Parameters + ---------- + bitset : numpy.ndarray + Packed uint32 bitset covering the whole dataset. + start : int + First dataset row of the batch. + count : int + Number of rows in the batch. + + Returns + ------- + mask : numpy.ndarray + Boolean array of shape ``(count,)``; True means the row passes. + words : numpy.ndarray + Packed uint32 array of shape ``(ceil(count / 32),)``. + """ + import numpy as np + + n_words = int(np.asarray(bitset).size) + if n_words * 32 < start + count: + raise ValueError( + f"prefilter bitset covers {n_words * 32} vectors, but rows " + f"[{start}, {start + count}) were requested; the bitset does not " + f"match the dataset" + ) + + word_start = start // 32 + word_end = (start + count + 31) // 32 + bits = np.unpackbits( + np.asarray(bitset[word_start:word_end]).view(np.uint8), + bitorder="little", + ) + offset = start - word_start * 32 + mask = bits[offset : offset + count].astype(bool) + + padded = np.zeros(((count + 31) // 32) * 32, dtype=np.uint8) + padded[:count] = mask + words = np.packbits(padded, bitorder="little").view(np.uint32) + return mask, words + + def cpu_search(dataset, queries, k, metric="sqeuclidean", accept_mask=None): """ Find the k nearest neighbors for each query point in the dataset using the @@ -251,8 +309,6 @@ def calc_truth(dataset, queries, k, metric="sqeuclidean", bitset=None): :func:`create_bitset_filter`. Bit i set means vector i passes the filter. When None, all vectors are considered. """ - import numpy as np - n_samples = dataset.shape[0] n = 500000 # batch size for processing neighbors i = 0 @@ -269,27 +325,26 @@ def calc_truth(dataset, queries, k, metric="sqeuclidean", bitset=None): X = xp.asarray(dataset[i : i + n_batch, :], xp.float32) + # Re-base the prefilter onto this batch: the index holds only rows + # [i, i + n_batch), so bit 0 of the filter must be row i. + accept_mask, batch_words = ( + (None, None) + if bitset is None + else slice_bitset(bitset, i, n_batch) + ) + if gpu_system: index = build(X, metric=metric, resources=resources) - prefilter = None - if bitset is not None: - word_start = i // 32 - word_end = (i + n_batch + 31) // 32 - batch_words = xp.asarray(bitset[word_start:word_end]) - prefilter = filters.from_bitset(batch_words) + prefilter = ( + None + if batch_words is None + else filters.from_bitset(xp.asarray(batch_words)) + ) D, Ind = search( index, queries, k, resources=resources, prefilter=prefilter ) resources.sync() else: - accept_mask = None - if bitset is not None: - word_start = i // 32 - word_end = (i + n_batch + 31) // 32 - batch_bytes = bitset[word_start:word_end].view(np.uint8) - accept_mask = np.unpackbits(batch_bytes, bitorder="little")[ - :n_batch - ].astype(bool) D, Ind = cpu_search( X, queries, k, metric=metric, accept_mask=accept_mask ) From 7914bf4deca8c2fdcd5b91c8a85621dc92b7298d Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 28 Jul 2026 19:18:09 -0700 Subject: [PATCH 08/10] Reject prefilters that accept fewer than k vectors, and allow seeding If fewer vectors pass the filter than the requested k, no query can have k valid neighbors and the ground truth is padded with filtered-out ids at infinite distance. That was written out silently. Count the accepted vectors up front and fail before doing any search work. Only the global count matters, since results are merged across batches. Padding bits in the final word are excluded so a bitset with stray padding cannot inflate the count. The count uses a byte-wise popcount table rather than unpackbits, which would allocate eight bytes per bit. Also add --filter_seed, so a generated bitset can be reproduced without the saved file. It is deliberately not named --seed: query generation has its own seed and is unaffected. --- .../generate_groundtruth/__main__.py | 73 ++++++++++++++++++- 1 file changed, 70 insertions(+), 3 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py b/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py index 7d052030db..fd54dd7bc1 100644 --- a/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py +++ b/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py @@ -129,7 +129,7 @@ def choose_random_queries_with_jitter(dataset, n_queries, seed=12345): return add_jitter(sampled, rng, normalize=False) -def create_bitset_filter(n_samples, filter_reject_rate): +def create_bitset_filter(n_samples, filter_reject_rate, seed=None): """ Creates a packed uint32 bitset where bit i is set iff vector i passes the filter. Each vector independently passes with probability @@ -142,6 +142,10 @@ def create_bitset_filter(n_samples, filter_reject_rate): Number of vectors in the dataset. filter_reject_rate : float Per-vector rejection probability, in [0.0, 1.0). + seed : int, optional + Seed for the sampling. ``None`` draws fresh entropy, so the bitset + differs between runs; the generated bitset is written to disk either + way, so a seed is only needed to reproduce one without that file. Returns ------- @@ -152,13 +156,52 @@ def create_bitset_filter(n_samples, filter_reject_rate): n_padded = ((n_samples + 31) // 32) * 32 bool_mask = np.zeros(n_padded, dtype=bool) - rng = np.random.default_rng() + rng = np.random.default_rng(seed) bool_mask[:n_samples] = rng.random(n_samples) >= filter_reject_rate # Pack with little-endian bit order: bit j maps to bit (j%32) of uint32 # word (j//32), LSB first — matching cuVS bitset layout. return np.packbits(bool_mask, bitorder="little").view(np.uint32) +def count_accepted(bitset, n_samples): + """ + Count how many of the first ``n_samples`` rows the bitset accepts. + + Padding bits in the final word are excluded, so a bitset whose padding + happens to be set does not inflate the count. + + Parameters + ---------- + bitset : numpy.ndarray + Packed uint32 bitset covering the dataset. + n_samples : int + Number of dataset rows to consider. + + Returns + ------- + int + Number of accepted rows. + """ + import numpy as np + + words = np.asarray(bitset) + # Byte-wise popcount table; avoids unpackbits, which would allocate eight + # bytes per bit and is a real cost on billion-row datasets. + table = np.array([bin(i).count("1") for i in range(256)], dtype=np.uint8) + + n_full_words = n_samples // 32 + total = int(table[words[:n_full_words].view(np.uint8)].sum(dtype=np.int64)) + + remainder = n_samples - n_full_words * 32 + if remainder: + tail = np.unpackbits( + words[n_full_words : n_full_words + 1].view(np.uint8), + bitorder="little", + ) + total += int(tail[:remainder].sum()) + return total + + def slice_bitset(bitset, start, count): """ Extract rows ``[start, start + count)`` of a packed bitset. @@ -501,6 +544,14 @@ def main(): "saved to /groundtruth.filter.bin. Mutually exclusive with " "--bitset.", ) + parser.add_argument( + "--filter_seed", + type=int, + default=None, + help="Seed for --filter_reject_rate sampling. Unseeded by default; " + "the generated bitset is saved either way, so this is only needed to " + "reproduce one without that file. Does not affect query generation.", + ) if len(sys.argv) == 1: parser.print_help() @@ -568,10 +619,26 @@ def main(): f"Generating prefilter bitset for filter_reject_rate=" f"{args.filter_reject_rate}" ) - bitset = create_bitset_filter(n_samples, args.filter_reject_rate) + bitset = create_bitset_filter( + n_samples, args.filter_reject_rate, args.filter_seed + ) bitset_filename = os.path.join(args.output, "groundtruth.filter.bin") write_bin(bitset_filename, bitset.reshape(-1, 1)) + if bitset is not None: + # Fewer than k accepted vectors cannot produce k valid neighbors for + # any query: the top-k would be padded with rejected ids at infinite + # distance, which is silently wrong ground truth. Only the global + # count matters, since results are merged across batches. + accepted = count_accepted(bitset, n_samples) + if accepted < args.k: + raise ValueError( + f"prefilter accepts only {accepted} of {n_samples} vectors, " + f"which is fewer than k={args.k}; the ground truth would be " + f"padded with filtered-out ids. Lower k or the reject rate." + ) + print(f"Prefilter accepts {accepted}/{n_samples} vectors") + print("Calculating true nearest neighbors") distances, indices = calc_truth( dataset, queries, args.k, args.metric, bitset=bitset From a90521bb972e7b1f98eeaeb23b31805400b93ca1 Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 28 Jul 2026 19:19:16 -0700 Subject: [PATCH 09/10] Update documentation --- fern/pages/cuvs_bench/datasets.md | 42 ++++++++++++++++++++++++++++--- 1 file changed, 39 insertions(+), 3 deletions(-) diff --git a/fern/pages/cuvs_bench/datasets.md b/fern/pages/cuvs_bench/datasets.md index 3b9a1a9bc3..83fab11947 100644 --- a/fern/pages/cuvs_bench/datasets.md +++ b/fern/pages/cuvs_bench/datasets.md @@ -83,17 +83,35 @@ If a dataset does not include ground truth, generate it with `cuvs_bench.generat ```bash # With an existing query file -python -m cuvs_bench.generate_groundtruth --dataset /dataset/base.fbin --output=groundtruth_dir --queries=/dataset/query.public.10K.fbin +python -m cuvs_bench.generate_groundtruth /dataset/base.fbin --output=groundtruth_dir --queries=/dataset/query.public.10K.fbin # With randomly generated queries -python -m cuvs_bench.generate_groundtruth --dataset /dataset/base.fbin --output=groundtruth_dir --queries=random --n_queries=10000 +python -m cuvs_bench.generate_groundtruth /dataset/base.fbin --output=groundtruth_dir --queries=random --n_queries=10000 # With random queries selected from a subset of the dataset -python -m cuvs_bench.generate_groundtruth --dataset /dataset/base.fbin --nrows=2000000 --output=groundtruth_dir --queries=random-choice --n_queries=10000 +python -m cuvs_bench.generate_groundtruth /dataset/base.fbin --rows=2000000 --output=groundtruth_dir --queries=random-choice --n_queries=10000 ``` For billion-scale sources that provide ground truth for only the first 10M or 100M base vectors, use `subset_size` in the dataset configuration so the benchmark uses the matching prefix of the base file. +### Prefiltered ground truth + +To benchmark filtered search, generate ground truth against the same filter the benchmark will apply. `--filter_reject_rate` builds a random bitset, uses it while computing the neighbors, and saves it next to the ground truth as `groundtruth.filter.bin`: + +```bash +python -m cuvs_bench.generate_groundtruth /dataset/base.fbin --output=groundtruth_dir --queries=/dataset/query.fbin --filter_reject_rate=0.9 +``` + +`--bitset` reuses an existing one instead, so several ground-truth sets can share a filter. The two options are mutually exclusive: + +```bash +python -m cuvs_bench.generate_groundtruth /dataset/base.fbin --output=groundtruth_dir --queries=/dataset/query.fbin --bitset=groundtruth_dir/groundtruth.filter.bin +``` + +The alternative is `filtering_rate` (below), which filters unfiltered ground truth on the fly. That is simpler to configure, but at high reject rates it discards most ground-truth neighbors and leaves fewer than `k` valid entries per query, so recall becomes unreliable. Prefiltered ground truth keeps `k` valid neighbors regardless of the reject rate. + +The bitset is indexed by absolute row id, so it must cover the whole dataset and match the ground truth it was generated with. Both the generator and the benchmark reject a bitset that is too short; a mismatched one of the right size is reported as a warning during the run. + ## Dataset configurations Each benchmark dataset needs a YAML descriptor with file names and basic properties. Common descriptors are available in [datasets.yaml](https://github.com/rapidsai/cuvs/blob/branch-25.04/python/cuvs_bench/cuvs_bench/config/datasets/datasets.yaml). @@ -123,6 +141,24 @@ For a new dataset, create a descriptor such as `mydataset.yaml`: Choose any `name` and pass it as `--dataset`. File paths are relative to `--dataset-path`. The optional `subset_size` uses the first `subset_size` vectors, which lets you benchmark subsets without duplicating base files. Generate separate ground truth for each subset. +Two optional keys enable filtered search, and at most one may be set: + +| Key | Purpose | +| --- | --- | +| `filtering_rate` | Fraction of vectors to reject. Generates a random bitset at startup and filters the ground truth on the fly. Needs no extra files, but see the caveat above at high reject rates. | +| `filter_bitset_file` | Path to a bitset written by `generate_groundtruth`, applied during search. Pair it with the prefiltered `groundtruth_neighbors_file` produced by the same run. | + +```yaml +- name: mydata-1M-filtered + base_file: mydata-1M/base.100M.u8bin + subset_size: 1000000 + dims: 128 + query_file: mydata-10M/queries.u8bin + groundtruth_neighbors_file: mydata-1M/groundtruth.neighbors.ibin + filter_bitset_file: mydata-1M/groundtruth.filter.bin + distance: euclidean +``` + Run the custom dataset with: ```bash From d9d78e6bb98041b2ba9a3b62dccb38a4776e66f8 Mon Sep 17 00:00:00 2001 From: James Xia Date: Wed, 29 Jul 2026 13:00:49 -0700 Subject: [PATCH 10/10] Add filtered IVF-RaBitQ support to cuVS Bench Forward filtering configuration through the Python orchestration and C++ benchmark layers, and pass generated or file-backed bitsets to IVF-RaBitQ search. Make filter residency independently configurable, defaulting to device memory. Add IVF-RaBitQ tuning configuration and constraints, document the new parameters, and test filter configuration propagation and validation. --- cpp/bench/ann/src/common/ann_types.hpp | 4 +- cpp/bench/ann/src/common/benchmark.hpp | 7 ++- .../ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h | 7 ++- fern/pages/cuvs_bench/param_tuning.md | 18 +++++++ python/cuvs_bench/cuvs_bench/backends/base.py | 10 +++- .../cuvs_bench/backends/cpp_gbench.py | 14 +++++- .../config/algos/constraints/__init__.py | 8 ++- .../config/algos/cuvs_ivf_rabitq.yaml | 31 ++++++++++++ .../cuvs_bench/orchestrator/config_loaders.py | 14 +++++- .../cuvs_bench/orchestrator/orchestrator.py | 4 +- .../cuvs_bench/tests/test_cpp_gbench.py | 50 ++++++++++++++++++- .../cuvs_bench/cuvs_bench/tests/test_utils.py | 35 ++++++++++++- 12 files changed, 190 insertions(+), 12 deletions(-) create mode 100644 python/cuvs_bench/cuvs_bench/config/algos/cuvs_ivf_rabitq.yaml diff --git a/cpp/bench/ann/src/common/ann_types.hpp b/cpp/bench/ann/src/common/ann_types.hpp index bd669dff78..c4b5e36120 100644 --- a/cpp/bench/ann/src/common/ann_types.hpp +++ b/cpp/bench/ann/src/common/ann_types.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -78,6 +78,8 @@ struct algo_property { MemoryType dataset_memory_type; // neighbors/distances should have same memory type as queries MemoryType query_memory_type; + // GPU filters should be device-resident by default, independently of the dataset. + MemoryType filter_memory_type = MemoryType::kDevice; }; class algo_base { diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp index f29859e3ff..86b254b8d7 100644 --- a/cpp/bench/ann/src/common/benchmark.hpp +++ b/cpp/bench/ann/src/common/benchmark.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -102,6 +102,9 @@ inline auto parse_algo_property(algo_property prop, const nlohmann::json& conf) if (conf.contains("query_memory_type")) { prop.query_memory_type = parse_memory_type(conf.at("query_memory_type")); } + if (conf.contains("filter_memory_type")) { + prop.filter_memory_type = parse_memory_type(conf.at("filter_memory_type")); + } return prop; }; @@ -263,7 +266,7 @@ void bench_search(::benchmark::State& state, } try { a->set_search_param(*search_param, - dataset->filter_bitset(current_algo_props->dataset_memory_type)); + dataset->filter_bitset(current_algo_props->filter_memory_type)); } catch (const std::exception& ex) { state.SkipWithError("An error occurred setting search parameters: " + std::string(ex.what())); return; diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h index fcd6b2aa70..ec89ab1b71 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h @@ -81,6 +81,7 @@ class cuvs_ivf_rabitq : public algo, public algo_gpu { int dimension_; float refine_ratio_ = 1.0; raft::device_matrix_view dataset_; + std::shared_ptr filter_; }; template @@ -118,8 +119,10 @@ std::unique_ptr> cuvs_ivf_rabitq::copy() } template -void cuvs_ivf_rabitq::set_search_param(const search_param_base& param, const void*) +void cuvs_ivf_rabitq::set_search_param(const search_param_base& param, + const void* filter_bitset) { + filter_ = make_cuvs_filter(filter_bitset, index_->size()); auto sp = dynamic_cast(param); search_params_ = sp.rabitq_param; refine_ratio_ = sp.refine_ratio; @@ -154,7 +157,7 @@ void cuvs_ivf_rabitq::search( auto distances_view = raft::make_device_matrix_view(distances, batch_size, k); cuvs::neighbors::ivf_rabitq::search( - handle_, search_params_, *index_, queries_view, neighbors_view, distances_view); + handle_, search_params_, *index_, queries_view, neighbors_view, distances_view, *filter_); if constexpr (sizeof(IdxT) != sizeof(algo_base::index_type)) { raft::linalg::unaryOp(neighbors, diff --git a/fern/pages/cuvs_bench/param_tuning.md b/fern/pages/cuvs_bench/param_tuning.md index 2805516d30..6cf986d9a4 100644 --- a/fern/pages/cuvs_bench/param_tuning.md +++ b/fern/pages/cuvs_bench/param_tuning.md @@ -35,6 +35,8 @@ The parameter tables below describe the build and search knobs that sweep mode v ## NVIDIA cuVS Indexes +Filter-capable GPU algorithms store the filter in device memory by default, independently of the base dataset. Set `filter_memory_type` in a search configuration to override this. Accepted values are `device`, `managed`, `pinned`, `host`, and `mmap`; `host` and `mmap` require a platform where the GPU can directly access pageable host memory. + ### cuvs_brute_force Use NVIDIA cuVS brute-force index for exact search. Brute-force has no further build or search parameters. @@ -73,6 +75,22 @@ IVF-pq is an inverted-file index, which partitions the vectors into a series of | `smemLutDtype` | `search` | N | [`float`, `half`, `fp8`] | `half` | The precision to use for the lookup table in shared memory. Lower precision can increase performance at the cost of accuracy. | | `refine_ratio` | `search` | N | Positive integer >0 | 1 | `refine_ratio * k` nearest neighbors are queried from the index initially and an additional refinement step improves recall by selecting only the best `k` neighbors. | +### cuvs_ivf_rabitq + +IVF-RaBitQ partitions vectors into inverted lists and compresses each dimension to a small number of bits. Its search modes trade lookup-table precision and query quantization cost for throughput. + +| Parameter | Type | Required | Data Type | Default | Description | +| --- | --- | --- | --- | --- | --- | +| `nlist` | `build` | N | Positive integer >0 | 1024 | Number of inverted lists. | +| `niter` | `build` | N | Positive integer >0 | 20 | Number of k-means iterations. | +| `max_points_per_cluster` | `build` | N | Positive integer >0 | 256 | Maximum training points per cluster. | +| `bits_per_dim` | `build` | N | Integer [1-9] | 3 | Number of bits used to encode each dimension. | +| `fast_quantize_flag` | `build` | N | Boolean | `true` | Enable the faster index quantization path. | +| `force_streaming` | `build` | N | Boolean | `false` | Force host-data streaming during construction. | +| `nprobe` | `search` | N | Positive integer <= `nlist` | 20 | Number of inverted lists searched per query. | +| `mode` | `search` | N | [`lut16`, `lut32`, `quant4`, `quant8`] | `quant4` | Distance evaluation mode. | +| `filter_memory_type` | `search` | N | [`device`, `managed`, `pinned`, `host`, `mmap`] | `device` | Residency used for a configured filter bitset. | + ### cuvs_cagra CAGRA uses a graph-based index, which creates an intermediate, approximate kNN graph using IVF-PQ and then further refining and optimizing to create a final kNN graph. This kNN graph is used by CAGRA as an index for search. diff --git a/python/cuvs_bench/cuvs_bench/backends/base.py b/python/cuvs_bench/cuvs_bench/backends/base.py index bf802d2128..3c24afb25b 100644 --- a/python/cuvs_bench/cuvs_bench/backends/base.py +++ b/python/cuvs_bench/cuvs_bench/backends/base.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -57,6 +57,10 @@ class Dataset: Path to ground truth distances file metadata : Optional[Dict[str, Any]] Additional dataset metadata like {"subset_size": 10000} + filtering_rate : Optional[float] + Fraction of dataset vectors rejected by a generated random filter + filter_bitset_file : Optional[str] + Path to a persisted cuVS bitset file used during search """ def __init__( @@ -72,6 +76,8 @@ def __init__( groundtruth_neighbors_file: Optional[str] = None, groundtruth_distances_file: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, + filtering_rate: Optional[float] = None, + filter_bitset_file: Optional[str] = None, ): self.name = name # Vectors are stored privately to support lazy loading. @@ -96,6 +102,8 @@ def __init__( self.query_file = query_file self.groundtruth_neighbors_file = groundtruth_neighbors_file self.groundtruth_distances_file = groundtruth_distances_file + self.filtering_rate = filtering_rate + self.filter_bitset_file = filter_bitset_file self.metadata = metadata or {} @property diff --git a/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py b/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py index a888e957e5..6e508e0c77 100644 --- a/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py +++ b/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -166,6 +166,12 @@ def build( dataset_config["groundtruth_neighbors_file"] = ( dataset.groundtruth_neighbors_file ) + if dataset.filtering_rate is not None: + dataset_config["filtering_rate"] = dataset.filtering_rate + if dataset.filter_bitset_file is not None: + dataset_config["filter_bitset_file"] = ( + dataset.filter_bitset_file + ) # Build index list from ALL IndexConfig objects (matches runners.py) index_list = [ @@ -409,6 +415,12 @@ def search( dataset_config["groundtruth_neighbors_file"] = ( dataset.groundtruth_neighbors_file ) + if dataset.filtering_rate is not None: + dataset_config["filtering_rate"] = dataset.filtering_rate + if dataset.filter_bitset_file is not None: + dataset_config["filter_bitset_file"] = ( + dataset.filter_bitset_file + ) # Build index list from ALL IndexConfig objects (matches runners.py) index_list = [ diff --git a/python/cuvs_bench/cuvs_bench/config/algos/constraints/__init__.py b/python/cuvs_bench/cuvs_bench/config/algos/constraints/__init__.py index f22852f0ae..fed41d20ef 100644 --- a/python/cuvs_bench/cuvs_bench/config/algos/constraints/__init__.py +++ b/python/cuvs_bench/cuvs_bench/config/algos/constraints/__init__.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 @@ -56,6 +56,12 @@ def cuvs_ivf_sq_search(params, build_params, k, batch_size): return True +def cuvs_ivf_rabitq_search(params, build_params, k, batch_size): + if "nlist" in build_params and "nprobe" in params: + return build_params["nlist"] >= params["nprobe"] + return True + + ############################################################################### # FAISS constraints # ############################################################################### diff --git a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_ivf_rabitq.yaml b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_ivf_rabitq.yaml new file mode 100644 index 0000000000..2c157f6a8c --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_ivf_rabitq.yaml @@ -0,0 +1,31 @@ +name: cuvs_ivf_rabitq +constraints: + search: cuvs_bench.config.algos.constraints.cuvs_ivf_rabitq_search +groups: + base: + build: + nlist: [1024, 2048, 4096, 8192] + bits_per_dim: [1, 2, 4] + max_points_per_cluster: [256] + niter: [20] + search: + nprobe: [1, 5, 10, 20, 50, 100, 200] + mode: ["lut16", "lut32", "quant4", "quant8"] + large: + build: + nlist: [8192, 16384, 32768] + bits_per_dim: [1, 2, 4] + max_points_per_cluster: [256] + niter: [20] + search: + nprobe: [10, 20, 50, 100, 200, 500, 1000] + mode: ["lut16", "lut32", "quant4", "quant8"] + test: + build: + nlist: [1024] + bits_per_dim: [1, 4] + max_points_per_cluster: [256] + niter: [20] + search: + nprobe: [1, 5] + mode: ["lut16", "lut32", "quant4", "quant8"] diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py b/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py index e3af2fe58e..a64aca9ee0 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -123,6 +123,8 @@ class DatasetConfig: distance: str = "euclidean" dims: Optional[int] = None subset_size: Optional[int] = None + filtering_rate: Optional[float] = None + filter_bitset_file: Optional[str] = None class ConfigLoader(ABC): @@ -282,6 +284,14 @@ def _resolve(rel): return os.path.join(dataset_path, rel) return rel + filtering_rate = dataset_conf.get("filtering_rate") + filter_bitset_file = _resolve(dataset_conf.get("filter_bitset_file")) + if filtering_rate is not None and filter_bitset_file is not None: + raise ValueError( + "dataset config must set at most one of filtering_rate or " + "filter_bitset_file" + ) + return DatasetConfig( name=dataset_conf["name"], base_file=_resolve(dataset_conf.get("base_file")), @@ -292,6 +302,8 @@ def _resolve(rel): groundtruth_distances_file=_resolve( dataset_conf.get("groundtruth_distances_file") ), + filtering_rate=filtering_rate, + filter_bitset_file=filter_bitset_file, distance=dataset_conf.get("distance", "euclidean"), dims=dataset_conf.get("dims"), subset_size=subset_size, diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py index 580291c2f0..694c6836e3 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -644,6 +644,8 @@ def _create_dataset(self, dataset_config: DatasetConfig) -> Dataset: query_file=dataset_config.query_file, groundtruth_neighbors_file=dataset_config.groundtruth_neighbors_file, groundtruth_distances_file=dataset_config.groundtruth_distances_file, + filtering_rate=dataset_config.filtering_rate, + filter_bitset_file=dataset_config.filter_bitset_file, distance_metric=dataset_config.distance, metadata={"subset_size": dataset_config.subset_size}, ) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_cpp_gbench.py b/python/cuvs_bench/cuvs_bench/tests/test_cpp_gbench.py index a25b4fbe04..0ccf4a310a 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_cpp_gbench.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_cpp_gbench.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -305,6 +305,54 @@ def test_build_dry_run(self): finally: temp_exec.unlink(missing_ok=True) + @pytest.mark.parametrize( + "filter_config", + [{"filtering_rate": 0.9}, {"filter_bitset_file": "filter.bin"}], + ) + def test_search_serializes_filter_configuration( + self, monkeypatch, tmp_path, filter_config + ): + executable = tmp_path / "benchmark" + executable.touch() + executable.chmod(0o755) + backend = CppGoogleBenchmarkBackend( + { + "name": "test_backend", + "executable_path": str(executable), + "data_prefix": str(tmp_path), + "dataset": "filtered", + "output_filename": ("build", "search"), + } + ) + dataset = Dataset( + name="filtered", + base_file="base.fbin", + query_file="query.fbin", + groundtruth_neighbors_file="groundtruth.ibin", + **filter_config, + ) + indexes = [ + IndexConfig( + name="rabitq", + algo="cuvs_ivf_rabitq", + build_param={"nlist": 16}, + search_params=[{"nprobe": 4}], + file="index", + ) + ] + captured = {} + + def capture_config(config, *_args, **_kwargs): + captured.update(config) + + monkeypatch.setattr( + "cuvs_bench.backends.cpp_gbench.json.dump", capture_config + ) + backend.search(dataset, indexes, k=10, batch_size=4, dry_run=True) + + for key, value in filter_config.items(): + assert captured["dataset"][key] == value + def test_search_dry_run(self): """Test search dry run mode.""" with tempfile.NamedTemporaryFile(delete=False) as f: diff --git a/python/cuvs_bench/cuvs_bench/tests/test_utils.py b/python/cuvs_bench/cuvs_bench/tests/test_utils.py index 01dd803a07..dda09ae8f6 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_utils.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_utils.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -669,6 +669,39 @@ def test_get_dataset_configuration(self): assert result["name"] == "glove-100" assert result["dims"] == 100 + def test_build_dataset_config_preserves_filter_configuration(self): + loader = CppGBenchConfigLoader() + result = loader.build_dataset_config( + { + "name": "filtered", + "base_file": "base.fbin", + "query_file": "query.fbin", + "distance": "euclidean", + "filter_bitset_file": "filter.bin", + }, + "/data", + None, + ) + + assert result.filtering_rate is None + assert result.filter_bitset_file == "/data/filter.bin" + + def test_build_dataset_config_rejects_multiple_filter_sources(self): + loader = CppGBenchConfigLoader() + with pytest.raises(ValueError, match="at most one"): + loader.build_dataset_config( + { + "name": "filtered", + "base_file": "base.fbin", + "query_file": "query.fbin", + "distance": "euclidean", + "filtering_rate": 0.5, + "filter_bitset_file": "filter.bin", + }, + "/data", + None, + ) + def test_get_dataset_configuration_not_found(self): """Test that missing datasets raise ValueError.""" loader = CppGBenchConfigLoader()