Skip to content

Commit 71306ec

Browse files
Stardust-SJFStardust-SJF
authored
Support IVF-RaBitQ in cuVS Library (#1866)
This PR introduces IVF-RaBitQ, a GPU-native ANNS solution that integrates the cluster-based method IVF with RaBitQ quantization into an efficient GPU index build/search pipeline. It can achieve a strong recall–throughput trade-off while having fast index build speed and a small storage footprint. Authors: - https://github.com/Stardust-SJF - James Xia (https://github.com/jamxia155) - Tamas Bela Feher (https://github.com/tfeher) - Kyle Edwards (https://github.com/KyleFromNVIDIA) Approvers: - Divye Gala (https://github.com/divyegala) - Micka (https://github.com/lowener) - Tamas Bela Feher (https://github.com/tfeher) URL: #1866
1 parent 0b090ba commit 71306ec

84 files changed

Lines changed: 9593 additions & 3 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎cpp/CMakeLists.txt‎

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -684,6 +684,192 @@ if(NOT BUILD_CPU_ONLY)
684684
OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_pq/increment_score"
685685
KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
686686
)
687+
set(ivf_rabitq_ns "cuvs::neighbors::ivf_rabitq::detail")
688+
generate_jit_lto_kernels(
689+
jit_lto_files
690+
NAME_FORMAT "ivf_rabitq_compute_inner_products_with_lut"
691+
MATRIX_JSON_FILE
692+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_matrix.json"
693+
KERNEL_INPUT_FILE
694+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_kernel.cu.in"
695+
FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_lut"
696+
FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
697+
OUTPUT_DIRECTORY
698+
"${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_inner_products_with_lut"
699+
KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
700+
)
701+
generate_jit_lto_kernels(
702+
jit_lto_files
703+
NAME_FORMAT "ivf_rabitq_lut_emit_distances_@with_ex_descriptor@"
704+
MATRIX_JSON_FILE
705+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_matrix.json"
706+
KERNEL_INPUT_FILE
707+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in"
708+
FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_lut_emit_distances<@with_ex_value@>"
709+
FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
710+
OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/lut_emit_distances"
711+
KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
712+
)
713+
generate_jit_lto_kernels(
714+
jit_lto_files
715+
NAME_FORMAT "ivf_rabitq_compute_inner_products_with_lut_block_sort"
716+
MATRIX_JSON_FILE
717+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_matrix.json"
718+
KERNEL_INPUT_FILE
719+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_kernel.cu.in"
720+
FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_lut_block_sort"
721+
FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
722+
OUTPUT_DIRECTORY
723+
"${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_inner_products_with_lut_block_sort"
724+
KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
725+
)
726+
generate_jit_lto_kernels(
727+
jit_lto_files
728+
NAME_FORMAT "ivf_rabitq_lut_block_sort_emit_topk_@with_ex_descriptor@"
729+
MATRIX_JSON_FILE
730+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_matrix.json"
731+
KERNEL_INPUT_FILE
732+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in"
733+
FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_lut_block_sort_emit_topk<@with_ex_value@>"
734+
FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
735+
OUTPUT_DIRECTORY
736+
"${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/lut_block_sort_emit_topk"
737+
KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
738+
)
739+
generate_jit_lto_kernels(
740+
jit_lto_files
741+
NAME_FORMAT "ivf_rabitq_extract_code_@ex_bits@bit"
742+
MATRIX_JSON_FILE
743+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_matrix.json"
744+
KERNEL_INPUT_FILE
745+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_kernel.cu.in"
746+
FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_extract_code<@ex_bits@>"
747+
FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
748+
OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/extract_code"
749+
KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
750+
)
751+
generate_jit_lto_kernels(
752+
jit_lto_files
753+
NAME_FORMAT "ivf_rabitq_compute_lut_ip_for_vec_@lut_dtype_descriptor@"
754+
MATRIX_JSON_FILE
755+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_matrix.json"
756+
KERNEL_INPUT_FILE
757+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_kernel.cu.in"
758+
FRAGMENT_TAG_FORMAT
759+
"${ivf_rabitq_ns}::fragment_tag_compute_lut_ip_for_vec<${ivf_rabitq_ns}::tag_lut_dtype_@lut_dtype_descriptor@>"
760+
FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
761+
OUTPUT_DIRECTORY
762+
"${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_lut_ip_for_vec"
763+
KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
764+
)
765+
generate_jit_lto_kernels(
766+
jit_lto_files
767+
NAME_FORMAT "ivf_rabitq_compute_bitwise_quantized_ip_for_vec_@num_bits@bit"
768+
MATRIX_JSON_FILE
769+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_matrix.json"
770+
KERNEL_INPUT_FILE
771+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_kernel.cu.in"
772+
FRAGMENT_TAG_FORMAT
773+
"${ivf_rabitq_ns}::fragment_tag_compute_bitwise_quantized_ip_for_vec<@num_bits@>"
774+
FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
775+
OUTPUT_DIRECTORY
776+
"${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_bitwise_quantized_ip_for_vec"
777+
KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
778+
)
779+
generate_jit_lto_kernels(
780+
jit_lto_files
781+
NAME_FORMAT "ivf_rabitq_compute_inner_products_with_lut16_opt"
782+
MATRIX_JSON_FILE
783+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_matrix.json"
784+
KERNEL_INPUT_FILE
785+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_kernel.cu.in"
786+
FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_lut16_opt"
787+
FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
788+
OUTPUT_DIRECTORY
789+
"${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_inner_products_with_lut16_opt"
790+
KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
791+
)
792+
generate_jit_lto_kernels(
793+
jit_lto_files
794+
NAME_FORMAT "ivf_rabitq_lut16_opt_emit_distances_@with_ex_descriptor@"
795+
MATRIX_JSON_FILE
796+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_matrix.json"
797+
KERNEL_INPUT_FILE
798+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in"
799+
FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_lut16_opt_emit_distances<@with_ex_value@>"
800+
FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
801+
OUTPUT_DIRECTORY
802+
"${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/lut16_opt_emit_distances"
803+
KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
804+
)
805+
generate_jit_lto_kernels(
806+
jit_lto_files
807+
NAME_FORMAT "ivf_rabitq_compute_inner_products_with_lut16_opt_block_sort_@with_ex_descriptor@"
808+
MATRIX_JSON_FILE
809+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_matrix.json"
810+
KERNEL_INPUT_FILE
811+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_kernel.cu.in"
812+
FRAGMENT_TAG_FORMAT
813+
"${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_lut16_opt_block_sort<@with_ex_value@>"
814+
FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
815+
OUTPUT_DIRECTORY
816+
"${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_inner_products_with_lut16_opt_block_sort"
817+
KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
818+
)
819+
generate_jit_lto_kernels(
820+
jit_lto_files
821+
NAME_FORMAT "ivf_rabitq_compute_inner_products_with_bitwise"
822+
MATRIX_JSON_FILE
823+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_matrix.json"
824+
KERNEL_INPUT_FILE
825+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_kernel.cu.in"
826+
FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_bitwise"
827+
FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
828+
OUTPUT_DIRECTORY
829+
"${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_inner_products_with_bitwise"
830+
KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
831+
)
832+
generate_jit_lto_kernels(
833+
jit_lto_files
834+
NAME_FORMAT "ivf_rabitq_bitwise_emit_distances_@with_ex_descriptor@"
835+
MATRIX_JSON_FILE
836+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_matrix.json"
837+
KERNEL_INPUT_FILE
838+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in"
839+
FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_bitwise_emit_distances<@with_ex_value@>"
840+
FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
841+
OUTPUT_DIRECTORY
842+
"${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/bitwise_emit_distances"
843+
KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
844+
)
845+
generate_jit_lto_kernels(
846+
jit_lto_files
847+
NAME_FORMAT "ivf_rabitq_compute_inner_products_with_bitwise_block_sort"
848+
MATRIX_JSON_FILE
849+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_matrix.json"
850+
KERNEL_INPUT_FILE
851+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in"
852+
FRAGMENT_TAG_FORMAT
853+
"${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_bitwise_block_sort"
854+
FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
855+
OUTPUT_DIRECTORY
856+
"${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_inner_products_with_bitwise_block_sort"
857+
KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
858+
)
859+
generate_jit_lto_kernels(
860+
jit_lto_files
861+
NAME_FORMAT "ivf_rabitq_bitwise_block_sort_emit_topk_@with_ex_descriptor@"
862+
MATRIX_JSON_FILE
863+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_matrix.json"
864+
KERNEL_INPUT_FILE
865+
"${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in"
866+
FRAGMENT_TAG_FORMAT
867+
"${ivf_rabitq_ns}::fragment_tag_bitwise_block_sort_emit_topk<@with_ex_value@>"
868+
FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
869+
OUTPUT_DIRECTORY
870+
"${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/bitwise_block_sort_emit_topk"
871+
KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
872+
)
687873
set(cagra_ns "cuvs::neighbors::cagra::detail")
688874
generate_jit_lto_kernels(
689875
jit_lto_files
@@ -1201,6 +1387,15 @@ if(NOT BUILD_CPU_ONLY)
12011387
src/neighbors/ivf_pq/detail/ivf_pq_process_and_fill_codes.cu
12021388
${ivf_pq_search_inst_files}
12031389
${ivf_pq_transform_inst_files}
1390+
src/neighbors/ivf_rabitq.cu
1391+
src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu
1392+
src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu
1393+
src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu
1394+
src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu
1395+
src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu
1396+
src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu
1397+
src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu
1398+
src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu
12041399
src/neighbors/ivf_sq_index.cpp
12051400
src/neighbors/ivf_sq/ivf_sq_build_extend_float_uint8_t_int64_t.cu
12061401
src/neighbors/ivf_sq/ivf_sq_build_extend_half_uint8_t_int64_t.cu

‎cpp/bench/ann/CMakeLists.txt‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,9 @@ option(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_FLAT "Include faiss' cpu ivf flat algori
2323
ON
2424
)
2525
option(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_PQ "Include faiss' cpu ivf pq algorithm in benchmark" ON)
26+
option(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_RABITQ
27+
"Include faiss' cpu ivf rabitq algorithm in benchmark" ON
28+
)
2629
option(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_SQ "Include faiss' cpu ivf sq algorithm in benchmark" ON)
2730
option(CUVS_ANN_BENCH_USE_FAISS_CPU_HNSW_FLAT "Include faiss' hnsw algorithm in benchmark" ON)
2831
option(CUVS_ANN_BENCH_USE_CUVS_IVF_FLAT "Include cuVS ivf flat algorithm in benchmark" ON)
@@ -48,6 +51,7 @@ option(CUVS_ANN_BENCH_SINGLE_EXE
4851
"Make a single executable with benchmark as shared library modules" OFF
4952
)
5053
option(CUVS_KNN_BENCH_USE_CUVS_BRUTE_FORCE "Include cuVS brute force knn in benchmark" ON)
54+
option(CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ "Include cuVS ivf RaBitQ algorithm in benchmark" ON)
5155

5256
# ##################################################################################################
5357
# * Process options ----------------------------------------------------------
@@ -85,6 +89,7 @@ if(BUILD_CPU_ONLY)
8589
set(CUVS_ANN_BENCH_USE_CUVS_IVF_FLAT OFF)
8690
set(CUVS_ANN_BENCH_USE_CUVS_IVF_SQ OFF)
8791
set(CUVS_ANN_BENCH_USE_CUVS_IVF_PQ OFF)
92+
set(CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ OFF)
8893
set(CUVS_ANN_BENCH_USE_CUVS_CAGRA OFF)
8994
set(CUVS_ANN_BENCH_USE_CUVS_BRUTE_FORCE OFF)
9095
set(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB OFF)
@@ -102,6 +107,7 @@ if(CUVS_ANN_BENCH_USE_CUVS_IVF_PQ
102107
OR CUVS_ANN_BENCH_USE_CUVS_BRUTE_FORCE
103108
OR CUVS_ANN_BENCH_USE_CUVS_IVF_FLAT
104109
OR CUVS_ANN_BENCH_USE_CUVS_IVF_SQ
110+
OR CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ
105111
OR CUVS_ANN_BENCH_USE_CUVS_CAGRA
106112
OR CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB
107113
OR CUVS_KNN_BENCH_USE_CUVS_BRUTE_FORCE
@@ -249,6 +255,12 @@ if(CUVS_ANN_BENCH_USE_CUVS_IVF_FLAT)
249255
)
250256
endif()
251257

258+
if(CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ)
259+
ConfigureAnnBench(
260+
NAME CUVS_IVF_RABITQ PATH src/cuvs/cuvs_benchmark.cu src/cuvs/cuvs_ivf_rabitq.cu LINKS cuvs
261+
)
262+
endif()
263+
252264
if(CUVS_ANN_BENCH_USE_CUVS_IVF_SQ)
253265
ConfigureAnnBench(
254266
NAME CUVS_IVF_SQ PATH src/cuvs/cuvs_benchmark.cu src/cuvs/cuvs_ivf_sq.cu LINKS cuvs
@@ -320,6 +332,12 @@ if(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_PQ)
320332
)
321333
endif()
322334

335+
if(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_RABITQ)
336+
ConfigureAnnBench(
337+
NAME FAISS_CPU_IVF_RABITQ PATH src/faiss/faiss_cpu_benchmark.cpp LINKS ${CUVS_FAISS_TARGETS}
338+
)
339+
endif()
340+
323341
if(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_SQ)
324342
ConfigureAnnBench(
325343
NAME FAISS_CPU_IVF_SQ PATH src/faiss/faiss_cpu_benchmark.cpp LINKS ${CUVS_FAISS_TARGETS}

‎cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,12 @@ extern template class cuvs::bench::cuvs_ivf_pq<float, int64_t>;
2424
extern template class cuvs::bench::cuvs_ivf_pq<uint8_t, int64_t>;
2525
extern template class cuvs::bench::cuvs_ivf_pq<int8_t, int64_t>;
2626
#endif
27+
#if defined(CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ)
28+
#include "cuvs_ivf_rabitq_wrapper.h"
29+
#endif
30+
#ifdef CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ
31+
extern template class cuvs::bench::cuvs_ivf_rabitq<float, int64_t>;
32+
#endif
2733
#if defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA) || defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB) || \
2834
defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_DISKANN)
2935
#include "cuvs_cagra_wrapper.h"
@@ -207,6 +213,47 @@ void parse_search_param(const nlohmann::json& conf,
207213
}
208214
#endif
209215

216+
#if defined(CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ)
217+
template <typename T, typename IdxT>
218+
void parse_build_param(const nlohmann::json& conf,
219+
typename cuvs::bench::cuvs_ivf_rabitq<T, IdxT>::build_param& param)
220+
{
221+
if (conf.contains("nlist")) { param.n_lists = conf.at("nlist"); }
222+
if (conf.contains("niter")) { param.kmeans_n_iters = conf.at("niter"); }
223+
if (conf.contains("max_points_per_cluster")) {
224+
param.max_train_points_per_cluster = conf.at("max_points_per_cluster");
225+
}
226+
if (conf.contains("bits_per_dim")) { param.bits_per_dim = conf.at("bits_per_dim"); }
227+
if (conf.contains("fast_quantize_flag")) {
228+
param.fast_quantize_flag = conf.at("fast_quantize_flag");
229+
}
230+
if (conf.contains("force_streaming")) { param.force_streaming = conf.at("force_streaming"); }
231+
}
232+
233+
template <typename T, typename IdxT>
234+
void parse_search_param(const nlohmann::json& conf,
235+
typename cuvs::bench::cuvs_ivf_rabitq<T, IdxT>::search_param& param)
236+
{
237+
if (conf.contains("nprobe")) { param.rabitq_param.n_probes = conf.at("nprobe"); }
238+
239+
if (conf.contains("mode")) {
240+
std::string mode = conf.at("mode");
241+
if (mode == "lut16") {
242+
param.rabitq_param.mode = cuvs::neighbors::ivf_rabitq::search_mode::LUT16;
243+
} else if (mode == "lut32") {
244+
param.rabitq_param.mode = cuvs::neighbors::ivf_rabitq::search_mode::LUT32;
245+
} else if (mode == "quant4") {
246+
param.rabitq_param.mode = cuvs::neighbors::ivf_rabitq::search_mode::QUANT4;
247+
} else if (mode == "quant8") {
248+
param.rabitq_param.mode = cuvs::neighbors::ivf_rabitq::search_mode::QUANT8;
249+
} else {
250+
throw std::runtime_error("mode: '" + mode +
251+
"', should be either 'lut16', 'lut32', 'quant4' or 'quant8'");
252+
}
253+
}
254+
}
255+
#endif
256+
210257
#if defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA) || defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB) || \
211258
defined(CUVS_ANN_BENCH_USE_CUVS_MG) || defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_DISKANN)
212259
template <typename T, typename IdxT>

‎cpp/bench/ann/src/cuvs/cuvs_benchmark.cu‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,15 @@ auto create_algo(const std::string& algo_name,
100100
a = std::make_unique<cuvs::bench::cuvs_ivf_pq<T, int64_t>>(metric, dim, param);
101101
}
102102
#endif
103+
#ifdef CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ
104+
if constexpr (std::is_same_v<T, float>) {
105+
if (algo_name == "cuvs_ivf_rabitq") {
106+
typename cuvs::bench::cuvs_ivf_rabitq<T, int64_t>::build_param param;
107+
parse_build_param<T, int64_t>(conf, param);
108+
a = std::make_unique<cuvs::bench::cuvs_ivf_rabitq<T, int64_t>>(metric, dim, param);
109+
}
110+
}
111+
#endif
103112
#ifdef CUVS_ANN_BENCH_USE_CUVS_CAGRA
104113
if (algo_name == "raft_cagra" || algo_name == "cuvs_cagra") {
105114
typename cuvs::bench::cuvs_cagra<T, uint32_t>::build_param param;
@@ -176,6 +185,16 @@ auto create_search_param(const std::string& algo_name, const nlohmann::json& con
176185
return param;
177186
}
178187
#endif
188+
#ifdef CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ
189+
if constexpr (std::is_same_v<T, float>) {
190+
if (algo_name == "cuvs_ivf_rabitq") {
191+
auto param =
192+
std::make_unique<typename cuvs::bench::cuvs_ivf_rabitq<T, int64_t>::search_param>();
193+
parse_search_param<T, int64_t>(conf, *param);
194+
return param;
195+
}
196+
}
197+
#endif
179198
#ifdef CUVS_ANN_BENCH_USE_CUVS_CAGRA
180199
if (algo_name == "raft_cagra" || algo_name == "cuvs_cagra") {
181200
auto param = std::make_unique<typename cuvs::bench::cuvs_cagra<T, uint32_t>::search_param>();

‎cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq.cu‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
/*
2+
* SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
#include "cuvs_ivf_rabitq_wrapper.h"
6+
7+
namespace cuvs::bench {
8+
template class cuvs_ivf_rabitq<float, int64_t>;
9+
} // namespace cuvs::bench

0 commit comments

Comments
 (0)