diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 227c2906cc..2250e7c4e1 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -684,6 +684,192 @@ if(NOT BUILD_CPU_ONLY) OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_pq/increment_score" 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_compute_inner_products_with_lut" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_kernel.cu.in" + FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_lut" + FRAGMENT_TAG_HEADER_FILES "" + OUTPUT_DIRECTORY + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_inner_products_with_lut" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + ) + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT "ivf_rabitq_lut_emit_distances_@with_ex_descriptor@" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in" + FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_lut_emit_distances<@with_ex_value@>" + FRAGMENT_TAG_HEADER_FILES "" + OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/lut_emit_distances" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + ) + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT "ivf_rabitq_compute_inner_products_with_lut_block_sort" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_kernel.cu.in" + FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_lut_block_sort" + FRAGMENT_TAG_HEADER_FILES "" + OUTPUT_DIRECTORY + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_inner_products_with_lut_block_sort" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + ) + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT "ivf_rabitq_lut_block_sort_emit_topk_@with_ex_descriptor@" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in" + FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_lut_block_sort_emit_topk<@with_ex_value@>" + FRAGMENT_TAG_HEADER_FILES "" + OUTPUT_DIRECTORY + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/lut_block_sort_emit_topk" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + ) + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT "ivf_rabitq_extract_code_@ex_bits@bit" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_kernel.cu.in" + FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_extract_code<@ex_bits@>" + FRAGMENT_TAG_HEADER_FILES "" + OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/extract_code" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + ) + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT "ivf_rabitq_compute_lut_ip_for_vec_@lut_dtype_descriptor@" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_kernel.cu.in" + FRAGMENT_TAG_FORMAT + "${ivf_rabitq_ns}::fragment_tag_compute_lut_ip_for_vec<${ivf_rabitq_ns}::tag_lut_dtype_@lut_dtype_descriptor@>" + FRAGMENT_TAG_HEADER_FILES "" + OUTPUT_DIRECTORY + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_lut_ip_for_vec" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + ) + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT "ivf_rabitq_compute_bitwise_quantized_ip_for_vec_@num_bits@bit" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_kernel.cu.in" + FRAGMENT_TAG_FORMAT + "${ivf_rabitq_ns}::fragment_tag_compute_bitwise_quantized_ip_for_vec<@num_bits@>" + FRAGMENT_TAG_HEADER_FILES "" + OUTPUT_DIRECTORY + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_bitwise_quantized_ip_for_vec" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + ) + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT "ivf_rabitq_compute_inner_products_with_lut16_opt" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_kernel.cu.in" + FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_lut16_opt" + FRAGMENT_TAG_HEADER_FILES "" + OUTPUT_DIRECTORY + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_inner_products_with_lut16_opt" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + ) + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT "ivf_rabitq_lut16_opt_emit_distances_@with_ex_descriptor@" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in" + FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_lut16_opt_emit_distances<@with_ex_value@>" + FRAGMENT_TAG_HEADER_FILES "" + OUTPUT_DIRECTORY + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/lut16_opt_emit_distances" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + ) + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT "ivf_rabitq_compute_inner_products_with_lut16_opt_block_sort_@with_ex_descriptor@" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_kernel.cu.in" + FRAGMENT_TAG_FORMAT + "${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_lut16_opt_block_sort<@with_ex_value@>" + FRAGMENT_TAG_HEADER_FILES "" + OUTPUT_DIRECTORY + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_inner_products_with_lut16_opt_block_sort" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + ) + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT "ivf_rabitq_compute_inner_products_with_bitwise" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_kernel.cu.in" + FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_bitwise" + FRAGMENT_TAG_HEADER_FILES "" + OUTPUT_DIRECTORY + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_inner_products_with_bitwise" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + ) + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT "ivf_rabitq_bitwise_emit_distances_@with_ex_descriptor@" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in" + FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_bitwise_emit_distances<@with_ex_value@>" + FRAGMENT_TAG_HEADER_FILES "" + OUTPUT_DIRECTORY + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/bitwise_emit_distances" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + ) + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT "ivf_rabitq_compute_inner_products_with_bitwise_block_sort" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in" + FRAGMENT_TAG_FORMAT + "${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_bitwise_block_sort" + FRAGMENT_TAG_HEADER_FILES "" + OUTPUT_DIRECTORY + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_inner_products_with_bitwise_block_sort" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + ) + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT "ivf_rabitq_bitwise_block_sort_emit_topk_@with_ex_descriptor@" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in" + FRAGMENT_TAG_FORMAT + "${ivf_rabitq_ns}::fragment_tag_bitwise_block_sort_emit_topk<@with_ex_value@>" + FRAGMENT_TAG_HEADER_FILES "" + OUTPUT_DIRECTORY + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/bitwise_block_sort_emit_topk" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + ) set(cagra_ns "cuvs::neighbors::cagra::detail") generate_jit_lto_kernels( jit_lto_files @@ -1200,6 +1386,15 @@ if(NOT BUILD_CPU_ONLY) src/neighbors/ivf_pq/detail/ivf_pq_process_and_fill_codes.cu ${ivf_pq_search_inst_files} ${ivf_pq_transform_inst_files} + src/neighbors/ivf_rabitq.cu + src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu + src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu + src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu + src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu + src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu + src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu + src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu + src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu src/neighbors/ivf_sq_index.cpp src/neighbors/ivf_sq/ivf_sq_build_extend_float_uint8_t_int64_t.cu src/neighbors/ivf_sq/ivf_sq_build_extend_half_uint8_t_int64_t.cu diff --git a/cpp/bench/ann/CMakeLists.txt b/cpp/bench/ann/CMakeLists.txt index c7884286e6..90d23d9aef 100644 --- a/cpp/bench/ann/CMakeLists.txt +++ b/cpp/bench/ann/CMakeLists.txt @@ -23,6 +23,9 @@ option(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_FLAT "Include faiss' cpu ivf flat algori ON ) option(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_PQ "Include faiss' cpu ivf pq algorithm in benchmark" ON) +option(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_RABITQ + "Include faiss' cpu ivf rabitq algorithm in benchmark" ON +) option(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_SQ "Include faiss' cpu ivf sq algorithm in benchmark" ON) option(CUVS_ANN_BENCH_USE_FAISS_CPU_HNSW_FLAT "Include faiss' hnsw algorithm in benchmark" ON) 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 "Make a single executable with benchmark as shared library modules" OFF ) option(CUVS_KNN_BENCH_USE_CUVS_BRUTE_FORCE "Include cuVS brute force knn in benchmark" ON) +option(CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ "Include cuVS ivf RaBitQ algorithm in benchmark" ON) # ################################################################################################## # * Process options ---------------------------------------------------------- @@ -85,6 +89,7 @@ if(BUILD_CPU_ONLY) set(CUVS_ANN_BENCH_USE_CUVS_IVF_FLAT OFF) set(CUVS_ANN_BENCH_USE_CUVS_IVF_SQ OFF) set(CUVS_ANN_BENCH_USE_CUVS_IVF_PQ OFF) + set(CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ OFF) set(CUVS_ANN_BENCH_USE_CUVS_CAGRA OFF) set(CUVS_ANN_BENCH_USE_CUVS_BRUTE_FORCE OFF) set(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB OFF) @@ -102,6 +107,7 @@ if(CUVS_ANN_BENCH_USE_CUVS_IVF_PQ OR CUVS_ANN_BENCH_USE_CUVS_BRUTE_FORCE OR CUVS_ANN_BENCH_USE_CUVS_IVF_FLAT OR CUVS_ANN_BENCH_USE_CUVS_IVF_SQ + OR CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ OR CUVS_ANN_BENCH_USE_CUVS_CAGRA OR CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB OR CUVS_KNN_BENCH_USE_CUVS_BRUTE_FORCE @@ -249,6 +255,12 @@ if(CUVS_ANN_BENCH_USE_CUVS_IVF_FLAT) ) endif() +if(CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ) + ConfigureAnnBench( + NAME CUVS_IVF_RABITQ PATH src/cuvs/cuvs_benchmark.cu src/cuvs/cuvs_ivf_rabitq.cu LINKS cuvs + ) +endif() + if(CUVS_ANN_BENCH_USE_CUVS_IVF_SQ) ConfigureAnnBench( 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) ) endif() +if(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_RABITQ) + ConfigureAnnBench( + NAME FAISS_CPU_IVF_RABITQ PATH src/faiss/faiss_cpu_benchmark.cpp LINKS ${CUVS_FAISS_TARGETS} + ) +endif() + if(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_SQ) ConfigureAnnBench( NAME FAISS_CPU_IVF_SQ PATH src/faiss/faiss_cpu_benchmark.cpp LINKS ${CUVS_FAISS_TARGETS} diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h index ea26ceed78..57b47d97db 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h @@ -24,6 +24,12 @@ extern template class cuvs::bench::cuvs_ivf_pq; extern template class cuvs::bench::cuvs_ivf_pq; extern template class cuvs::bench::cuvs_ivf_pq; #endif +#if defined(CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ) +#include "cuvs_ivf_rabitq_wrapper.h" +#endif +#ifdef CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ +extern template class cuvs::bench::cuvs_ivf_rabitq; +#endif #if defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA) || defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB) || \ defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_DISKANN) #include "cuvs_cagra_wrapper.h" @@ -207,6 +213,47 @@ void parse_search_param(const nlohmann::json& conf, } #endif +#if defined(CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ) +template +void parse_build_param(const nlohmann::json& conf, + typename cuvs::bench::cuvs_ivf_rabitq::build_param& param) +{ + if (conf.contains("nlist")) { param.n_lists = conf.at("nlist"); } + if (conf.contains("niter")) { param.kmeans_n_iters = conf.at("niter"); } + if (conf.contains("max_points_per_cluster")) { + param.max_train_points_per_cluster = conf.at("max_points_per_cluster"); + } + if (conf.contains("bits_per_dim")) { param.bits_per_dim = conf.at("bits_per_dim"); } + if (conf.contains("fast_quantize_flag")) { + param.fast_quantize_flag = conf.at("fast_quantize_flag"); + } + if (conf.contains("force_streaming")) { param.force_streaming = conf.at("force_streaming"); } +} + +template +void parse_search_param(const nlohmann::json& conf, + typename cuvs::bench::cuvs_ivf_rabitq::search_param& param) +{ + if (conf.contains("nprobe")) { param.rabitq_param.n_probes = conf.at("nprobe"); } + + if (conf.contains("mode")) { + std::string mode = conf.at("mode"); + if (mode == "lut16") { + param.rabitq_param.mode = cuvs::neighbors::ivf_rabitq::search_mode::LUT16; + } else if (mode == "lut32") { + param.rabitq_param.mode = cuvs::neighbors::ivf_rabitq::search_mode::LUT32; + } else if (mode == "quant4") { + param.rabitq_param.mode = cuvs::neighbors::ivf_rabitq::search_mode::QUANT4; + } else if (mode == "quant8") { + param.rabitq_param.mode = cuvs::neighbors::ivf_rabitq::search_mode::QUANT8; + } else { + throw std::runtime_error("mode: '" + mode + + "', should be either 'lut16', 'lut32', 'quant4' or 'quant8'"); + } + } +} +#endif + #if defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA) || defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB) || \ defined(CUVS_ANN_BENCH_USE_CUVS_MG) || defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_DISKANN) template diff --git a/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu b/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu index 22aeb31c38..3056ddc365 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu @@ -100,6 +100,15 @@ auto create_algo(const std::string& algo_name, a = std::make_unique>(metric, dim, param); } #endif +#ifdef CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ + if constexpr (std::is_same_v) { + if (algo_name == "cuvs_ivf_rabitq") { + typename cuvs::bench::cuvs_ivf_rabitq::build_param param; + parse_build_param(conf, param); + a = std::make_unique>(metric, dim, param); + } + } +#endif #ifdef CUVS_ANN_BENCH_USE_CUVS_CAGRA if (algo_name == "raft_cagra" || algo_name == "cuvs_cagra") { typename cuvs::bench::cuvs_cagra::build_param param; @@ -176,6 +185,16 @@ auto create_search_param(const std::string& algo_name, const nlohmann::json& con return param; } #endif +#ifdef CUVS_ANN_BENCH_USE_CUVS_IVF_RABITQ + if constexpr (std::is_same_v) { + if (algo_name == "cuvs_ivf_rabitq") { + auto param = + std::make_unique::search_param>(); + parse_search_param(conf, *param); + return param; + } + } +#endif #ifdef CUVS_ANN_BENCH_USE_CUVS_CAGRA if (algo_name == "raft_cagra" || algo_name == "cuvs_cagra") { auto param = std::make_unique::search_param>(); diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq.cu b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq.cu new file mode 100644 index 0000000000..1c6772d994 --- /dev/null +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq.cu @@ -0,0 +1,9 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#include "cuvs_ivf_rabitq_wrapper.h" + +namespace cuvs::bench { +template class cuvs_ivf_rabitq; +} // namespace cuvs::bench diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h new file mode 100644 index 0000000000..ca8f77b808 --- /dev/null +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h @@ -0,0 +1,166 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include "../common/ann_types.hpp" +#include "cuvs_ann_bench_utils.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace cuvs::bench { + +template +class cuvs_ivf_rabitq : public algo, public algo_gpu { + public: + using search_param_base = typename algo::search_param; + using algo::dim_; + + struct search_param : public search_param_base { + cuvs::neighbors::ivf_rabitq::search_params rabitq_param; + float refine_ratio = 1.0f; + [[nodiscard]] auto needs_dataset() const -> bool override { return refine_ratio > 1.0f; } + }; + + using build_param = cuvs::neighbors::ivf_rabitq::index_params; + + cuvs_ivf_rabitq(Metric metric, int dim, const build_param& param) + : algo(metric, dim), index_params_(param), dimension_(dim) + { + } + + void build(const T* dataset, size_t nrow) final; + + void set_search_param(const search_param_base& param, const void* filter_bitset) override; + void set_search_dataset(const T* dataset, size_t nrow) override; + + void search(const T* queries, + int batch_size, + int k, + algo_base::index_type* neighbors, + float* distances) const override; + + [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override + { + return handle_.get_sync_stream(); + } + + // to enable dataset access from GPU memory + [[nodiscard]] auto get_preference() const -> algo_property override + { + algo_property property; + property.dataset_memory_type = MemoryType::kHost; + property.query_memory_type = MemoryType::kDevice; + return property; + } + void save(const std::string& file) const override; + void load(const std::string&) override; + std::unique_ptr> copy() override; + + private: + // handle_ must go first to make sure it dies last and all memory allocated in pool + configured_raft_resources handle_{}; + build_param index_params_; + cuvs::neighbors::ivf_rabitq::search_params search_params_; + std::shared_ptr> index_; + int dimension_; + float refine_ratio_ = 1.0; + raft::device_matrix_view dataset_; +}; + +template +void cuvs_ivf_rabitq::save(const std::string& file) const +{ + cuvs::neighbors::ivf_rabitq::serialize(handle_, file, *index_); +} + +template +void cuvs_ivf_rabitq::load(const std::string& file) +{ + index_ = std::make_shared>(handle_); + cuvs::neighbors::ivf_rabitq::deserialize(handle_, file, index_.get()); +} + +template +void cuvs_ivf_rabitq::build(const T* dataset, size_t nrow) +{ + // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping. + size_t n_streams = 1; + raft::resource::set_cuda_stream_pool(handle_, std::make_shared(n_streams)); + auto dataset_v = raft::make_device_matrix_view(dataset, IdxT(nrow), dim_); + std::make_shared>( + std::move(cuvs::neighbors::ivf_rabitq::build(handle_, index_params_, dataset_v))) + .swap(index_); + // Note: internally the IVF-RaBitQ build works with simple pointers, and accepts both host and + // device pointer. Therefore, although we provide here a device_mdspan, this works with host + // pointer too. +} + +template +std::unique_ptr> cuvs_ivf_rabitq::copy() +{ + return std::make_unique>(*this); // use copy constructor +} + +template +void cuvs_ivf_rabitq::set_search_param(const search_param_base& param, const void*) +{ + auto sp = dynamic_cast(param); + search_params_ = sp.rabitq_param; + refine_ratio_ = sp.refine_ratio; + assert(search_params_.n_probes <= index_params_.n_lists); +} + +template +void cuvs_ivf_rabitq::set_search_dataset(const T* dataset, size_t nrow) +{ + dataset_ = raft::make_device_matrix_view(dataset, nrow, index_->dim()); +} + +template +void cuvs_ivf_rabitq::search( + const T* queries, int batch_size, int k, algo_base::index_type* neighbors, float* distances) const +{ + static_assert(std::is_integral_v); + static_assert(std::is_integral_v); + + IdxT* neighbors_idx; + std::optional> neighbors_storage{std::nullopt}; + if constexpr (sizeof(IdxT) == sizeof(algo_base::index_type)) { + neighbors_idx = reinterpret_cast(neighbors); + } else { + neighbors_storage.emplace(batch_size * k, raft::resource::get_cuda_stream(handle_)); + neighbors_idx = neighbors_storage->data(); + } + + auto queries_view = + raft::make_device_matrix_view(queries, batch_size, dimension_); + auto neighbors_view = raft::make_device_matrix_view(neighbors_idx, batch_size, k); + 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); + + if constexpr (sizeof(IdxT) != sizeof(algo_base::index_type)) { + raft::linalg::unaryOp(neighbors, + neighbors_idx, + batch_size * k, + raft::cast_op(), + raft::resource::get_cuda_stream(handle_)); + } +} +} // namespace cuvs::bench diff --git a/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp b/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp index 4a78f7c0a3..1a99e56028 100644 --- a/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp +++ b/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -49,6 +49,13 @@ void parse_build_param(const nlohmann::json& conf, } } +template +void parse_build_param(const nlohmann::json& conf, + typename cuvs::bench::faiss_cpu_ivfrabitq::build_param& param) +{ + parse_base_build_param(conf, param); +} + template void parse_build_param(const nlohmann::json& conf, typename cuvs::bench::faiss_cpu_ivfsq::build_param& param) @@ -114,6 +121,8 @@ auto create_algo(const std::string& algo_name, a = make_algo(metric, dim, conf); } else if (algo_name == "faiss_cpu_ivf_pq") { a = make_algo(metric, dim, conf); + } else if (algo_name == "faiss_cpu_ivf_rabitq") { + a = make_algo(metric, dim, conf); } else if (algo_name == "faiss_cpu_ivf_sq") { a = make_algo(metric, dim, conf); } else if (algo_name == "faiss_cpu_flat") { @@ -135,7 +144,7 @@ auto create_search_param(const std::string& algo_name, const nlohmann::json& con -> std::unique_ptr::search_param> { if (algo_name == "faiss_cpu_ivf_flat" || algo_name == "faiss_cpu_ivf_pq" || - algo_name == "faiss_cpu_ivf_sq") { + algo_name == "faiss_cpu_ivf_rabitq" || algo_name == "faiss_cpu_ivf_sq") { auto param = std::make_unique::search_param>(); parse_search_param(conf, *param); return param; diff --git a/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h b/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h index ef2b7d2c25..282d57dc2e 100644 --- a/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h +++ b/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -241,6 +242,31 @@ class faiss_cpu_ivfpq : public faiss_cpu { } }; +template +class faiss_cpu_ivfrabitq : public faiss_cpu { + public: + struct build_param : public faiss_cpu::build_param {}; + + faiss_cpu_ivfrabitq(Metric metric, int dim, const build_param& param) + : faiss_cpu(metric, dim, param) + { + this->init_quantizer(dim); + this->index_ = std::make_shared( + this->quantizer_.get(), dim, param.nlist, this->metric_type_); + } + + void save(const std::string& file) const override + { + this->template save_(file); + } + void load(const std::string& file) override { this->template load_(file); } + + std::unique_ptr> copy() + { + return std::make_unique>(*this); // use copy constructor + } +}; + // TODO(snanditale): Enable this in cmake // ref: https://github.com/rapidsai/raft/issues/1876 template diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index 507f44d364..82a34f9242 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -39,6 +39,9 @@ endfunction() function(populate_matrix_variables matrix_json_entry) string(JSON len LENGTH "${matrix_json_entry}") + if(len EQUAL 0) + return() + endif() math(EXPR last "${len} - 1") # cmake-lint: disable=C0103,E1120 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 new file mode 100644 index 0000000000..30885947b5 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp @@ -0,0 +1,50 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +namespace cuvs::neighbors::ivf_rabitq::detail { + +struct fragment_tag_compute_inner_products_with_lut {}; + +template +struct fragment_tag_lut_emit_distances {}; + +struct fragment_tag_compute_inner_products_with_lut_block_sort {}; + +template +struct fragment_tag_lut_block_sort_emit_topk {}; + +struct fragment_tag_compute_inner_products_with_lut16_opt {}; + +template +struct fragment_tag_lut16_opt_emit_distances {}; + +template +struct fragment_tag_compute_inner_products_with_lut16_opt_block_sort {}; + +struct fragment_tag_compute_inner_products_with_bitwise {}; + +template +struct fragment_tag_bitwise_emit_distances {}; + +struct fragment_tag_compute_inner_products_with_bitwise_block_sort {}; + +template +struct fragment_tag_bitwise_block_sort_emit_topk {}; + +template +struct fragment_tag_extract_code {}; + +struct tag_lut_dtype_f32 {}; +struct tag_lut_dtype_f16 {}; + +template +struct fragment_tag_compute_lut_ip_for_vec {}; + +template +struct fragment_tag_compute_bitwise_quantized_ip_for_vec {}; + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/include/cuvs/neighbors/ivf_rabitq.hpp b/cpp/include/cuvs/neighbors/ivf_rabitq.hpp new file mode 100644 index 0000000000..d26bc04022 --- /dev/null +++ b/cpp/include/cuvs/neighbors/ivf_rabitq.hpp @@ -0,0 +1,308 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace CUVS_EXPORT cuvs { +namespace neighbors { +namespace ivf_rabitq { + +// forward declaration +namespace detail { +class IVFGPU; +} + +/** + * @defgroup ivf_rabitq_cpp_index_params IVF-RaBitQ index build parameters + * @{ + */ +struct index_params : cuvs::neighbors::index_params { + /** + * The number of inverted lists (clusters) + * + * Hint: Increasing this parameter may alleviate shared memory pressure. + */ + uint32_t n_lists = 1024; + /** + * The total number of bits per dimension (single bit required for the binary RaBitQ algorithm + + * additional bits for extended RaBitQ). + * + * Supported values: [1, 2, 3, 4, 5, 6, 7, 8, 9]. + * + * Hint: the smaller the 'bits_per_dim', the smaller the index size and the better the search + * performance, but the lower the recall. + */ + uint32_t bits_per_dim = 3; + /** The number of iterations searching for kmeans centers (index building). */ + uint32_t kmeans_n_iters = 20; + /** The number of data vectors (per cluster) to use during iterative kmeans building. */ + uint32_t max_train_points_per_cluster = 256; + /** Flag for using the fast quantize method */ + bool fast_quantize_flag = true; + /** + * Maximum number of vectors per batch when using streaming construction. + * + * This parameter controls the batch size during streaming construction from host memory. + * Batches contain complete clusters only (no partial clusters across batch boundaries). + * + * Note: Streaming construction is automatically used when the dataset doesn't fit + * comfortably in GPU memory (determined by available workspace and kTolerableRatio). + */ + size_t streaming_batch_size = 100000; + /** + * Force streaming construction regardless of dataset size. + * + * When set to true, streaming construction will be used even if the dataset would fit + * in GPU memory. This is useful for testing or when you want explicit control over + * the construction method. + * + * Note: This parameter only applies when the input dataset is in host memory. If the + * dataset is already in device memory, streaming construction is not applicable and + * this parameter has no effect. + * + * Default: false (auto-detect based on available memory) + */ + bool force_streaming = false; +}; +/** + * @} + */ + +/** + * @defgroup ivf_rabitq_cpp_search_params IVF-RaBitQ index search parameters + * @{ + */ +/** A type for specifying the mode for searching the RaBitQ index. */ +enum class search_mode { + LUT16 = 0, + LUT32 = 1, + QUANT4 = 2, + QUANT8 = 3, +}; + +struct search_params : cuvs::neighbors::search_params { + /** The number of clusters to search. */ + uint32_t n_probes = 20; + /** The search mode to be used. */ + search_mode mode = search_mode::QUANT4; +}; +/** + * @} + */ + +static_assert(std::is_aggregate_v); +static_assert(std::is_aggregate_v); + +/** + * @defgroup ivf_rabitq_cpp_index IVF-RaBitQ index + * @{ + */ +/** + * @brief IVF-RaBitQ index. + * @tparam IdxT type of the indices in the source dataset + * + */ +template +struct index : cuvs::neighbors::index { + using index_params_type = ivf_rabitq::index_params; + using search_params_type = ivf_rabitq::search_params; + using index_type = IdxT; + static_assert(!raft::is_narrowing_v, + "IdxT must be able to represent all values of uint32_t"); + + public: + index(const index&) = delete; + index( + index&&); // declaration only; definition in impl to allow member unique_ptr of incomplete type + auto operator=(const index&) -> index& = delete; + auto operator=(index&&) -> index&; // declaration only; definition in impl to allow member + // unique_ptr of incomplete type + ~index(); // declaration only; definition in impl to allow member unique_ptr of incomplete type + + /** + * @brief Construct an empty index yet to be populated. + * + */ + index(raft::resources const& handle); + + /** Construct an empty index yet to be populated. */ + index(raft::resources const& handle, + size_t n_rows, + uint32_t dim, + uint32_t n_lists, + uint32_t bits_per_dim); + + /** Dimensionality of the input data. */ + uint32_t dim() const noexcept; + + /** Total length of the index. */ + IdxT size() const noexcept; + + /** Accessor for underlying RaBitQ index */ + detail::IVFGPU& rabitq_index() noexcept; + + private: + std::unique_ptr rabitq_index_; +}; +/** + * @} + */ + +/** + * @defgroup ivf_rabitq_cpp_index_build IVF-RaBitQ index build + * @{ + */ +/** + * @brief Build the index from the dataset for efficient search. + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * // use default index parameters + * ivf_rabitq::index_params index_params; + * // create and fill the index from a [N, D] dataset + * auto index = ivf_rabitq::build(handle, index_params, dataset); + * @endcode + * + * @param[in] handle + * @param[in] index_params configure the index building + * @param[in] dataset a device_matrix_view to a row-major matrix [n_rows, dim] + * @return the constructed ivf-rabitq index + * + */ +auto build(raft::resources const& handle, + const cuvs::neighbors::ivf_rabitq::index_params& index_params, + raft::device_matrix_view dataset) + -> cuvs::neighbors::ivf_rabitq::index; + +/** + * @brief Build the index from the dataset for efficient search. + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * // use default index parameters + * ivf_rabitq::index_params index_params; + * // create and fill the index from a [N, D] dataset + * auto index = ivf_rabitq::build(handle, index_params, dataset); + * @endcode + * + * @param[in] handle + * @param[in] index_params configure the index building + * @param[in] dataset a host_matrix_view to a row-major matrix [n_rows, dim] + * @return the constructed ivf-rabitq index + * + */ +auto build(raft::resources const& handle, + const cuvs::neighbors::ivf_rabitq::index_params& index_params, + raft::host_matrix_view dataset) + -> cuvs::neighbors::ivf_rabitq::index; +/** + * @} + */ + +/** + * @defgroup ivf_rabitq_cpp_index_search IVF-RaBitQ index search + * @{ + */ +/** + * @brief Search ANN using the constructed index. + * + * @param[in] handle + * @param[in] search_params configure the search + * @param[in] index ivf-rabitq constructed index + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + */ +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); + +/** + * @} + */ + +/** + * @defgroup ivf_rabitq_cpp_serialize IVF-RaBitQ index serialize + * @{ + */ +/** + * Save the index to file. + * + * @code{.cpp} + * #include + * + * raft::resources handle; + * + * // create a string with a filepath + * std::string filename("/path/to/index"); + * // create an index with `auto index = ivf_rabitq::build(...);` + * cuvs::neighbors::ivf_rabitq::serialize(handle, filename, index); + * @endcode + * + * @param[in] handle the raft handle + * @param[in] filename the file name for saving the index + * @param[in] index IVF-RaBitQ index + * + */ +void serialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::ivf_rabitq::index& index); + +/** + * Load index from file. + * + * @code{.cpp} + * #include + * + * raft::resources handle; + * + * // create a string with a filepath + * std::string filename("/path/to/index"); + * using IdxT = int64_t; // type of the index + * // create an empty index + * ivf_rabitq::index index(handle); + * + * cuvs::neighbors::ivf_rabitq::deserialize(handle, filename, &index); + * @endcode + * + * @param[in] handle the raft handle + * @param[in] filename the name of the file that stores the index + * @param[out] index IVF-PQ index + * + */ +void deserialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::ivf_rabitq::index* index); +/** + * @} + */ + +} // namespace ivf_rabitq +} // namespace neighbors +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/src/neighbors/ivf_rabitq.cu b/cpp/src/neighbors/ivf_rabitq.cu new file mode 100644 index 0000000000..d5572d4039 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq.cu @@ -0,0 +1,394 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "ivf_rabitq/gpu_index/ivf_gpu.cuh" +#include "ivf_rabitq/gpu_index/searcher_gpu.cuh" +#include +#include + +#include "../core/nvtx.hpp" +#include "detail/ann_utils.cuh" + +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq { + +namespace detail { + +using namespace cuvs::spatial::knn::detail; // NOLINT + +template +auto build(raft::resources const& handle, + const index_params& params, + raft::mdspan, raft::row_major, accessor> dataset) + -> cuvs::neighbors::ivf_rabitq::index +{ + IdxT n_rows = dataset.extent(0); + IdxT dim = dataset.extent(1); + raft::common::nvtx::range fun_scope( + "ivf_rabitq::build(%zu, %u)", size_t(n_rows), dim); + static_assert(std::is_same_v, "Unsupported data type"); + + RAFT_EXPECTS(n_rows > 0 && dim > 0, "empty dataset"); + RAFT_EXPECTS(n_rows >= params.n_lists, "number of rows can't be less than n_lists"); + + // Calculate dataset size and available workspace once + size_t dataset_bytes = sizeof(T) * n_rows * dim; + size_t available_workspace = raft::resource::get_workspace_free_bytes(handle); + constexpr size_t kTolerableRatio = 4; + + rmm::device_async_resource_ref device_memory = raft::resource::get_workspace_resource_ref(handle); + // If the dataset is small enough to comfortably fit into device memory, put it there. + // Otherwise, use the managed memory. + rmm::device_async_resource_ref big_memory_resource = + raft::resource::get_large_workspace_resource_ref(handle); + if (dataset_bytes * kTolerableRatio < available_workspace) { + big_memory_resource = device_memory; + } + + auto stream = raft::resource::get_cuda_stream(handle); + + // Determine if we should use streaming construction + bool use_streaming = false; + const float* host_dataset_ptr = nullptr; + + // Check if dataset is already on device + auto dataset_residency = utils::check_pointer_residency(dataset.data_handle()); + bool dataset_on_device = (dataset_residency == utils::pointer_residency::device_only); + + // If dataset is on host, check if we should use streaming construction + if (!dataset_on_device) { + // Use streaming if explicitly requested or if dataset doesn't fit comfortably + if (params.force_streaming || dataset_bytes * kTolerableRatio >= available_workspace) { + use_streaming = true; + host_dataset_ptr = dataset.data_handle(); + if (params.force_streaming) { + RAFT_LOG_INFO( + "Using streaming construction: explicitly requested via force_streaming parameter"); + } else { + RAFT_LOG_INFO( + "Using streaming construction: dataset size (%.2f GB) exceeds comfortable GPU memory " + "limit", + dataset_bytes / (1024.0 * 1024.0 * 1024.0)); + } + } + } + + // create device view of dataset (only if not using streaming) + auto d_dataset_array = + raft::make_device_mdarray(handle, big_memory_resource, raft::make_extents(0, 0)); + auto d_dataset_view = + raft::make_mdspan(dataset.data_handle(), raft::make_extents(n_rows, dim)); + + if (!use_streaming && !dataset_on_device) { + try { + d_dataset_array = raft::make_device_mdarray( + handle, big_memory_resource, raft::make_extents(n_rows, dim)); + } catch (raft::logic_error& e) { + RAFT_LOG_ERROR( + "Insufficient memory for full GPU construction. Please decrease " + "dataset size, or set large_workspace_resource appropriately."); + throw; + } + raft::copy(d_dataset_array.data_handle(), dataset.data_handle(), n_rows * dim, stream); + d_dataset_view = d_dataset_array.view(); + } + + auto dataset_const_view = raft::make_const_mdspan(d_dataset_view); + rmm::device_uvector cluster_centers(params.n_lists * dim, stream, big_memory_resource); + rmm::device_uvector labels(n_rows, stream, big_memory_resource); + + // Scope for kmeans training set allocation. + { + raft::random::RngState random_state{137}; + size_t n_rows_train = + std::min(static_cast(n_rows), + static_cast(params.max_train_points_per_cluster) * params.n_lists); + + // Besides just sampling, we transform the input dataset into floats to make it easier + // to use gemm operations from cublas. + auto trainset = + raft::make_device_mdarray(handle, big_memory_resource, raft::make_extents(0, 0)); + try { + trainset = raft::make_device_mdarray( + 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."); + throw; + } + raft::matrix::sample_rows(handle, random_state, dataset, trainset.view()); + + // perform k-means clustering + // NB: here cluster_centers is used as if it is [n_clusters, data_dim] not [n_clusters, + // dim_ext]! + auto centers_view = + raft::make_device_matrix_view(cluster_centers.data(), params.n_lists, dim); + cuvs::cluster::kmeans::balanced_params kmeans_params; + kmeans_params.n_iters = params.kmeans_n_iters; + kmeans_params.metric = cuvs::distance::DistanceType::L2Expanded; + // find cluster labels for dataset vectors + auto labels_view = raft::make_device_vector_view(labels.data(), n_rows); + auto centers_const_view = raft::make_device_matrix_view( + cluster_centers.data(), params.n_lists, dim); + cuvs::cluster::kmeans::fit( + handle, kmeans_params, raft::make_const_mdspan(trainset.view()), centers_view); + if (use_streaming) { + // The full dataset lives on host (it doesn't fit comfortably in device memory, which is + // why we stream). `predict` launches device kernels, so it can't read the host pointer + // directly. Compute cluster labels in device-sized batches: copy each chunk to the GPU and + // predict into the matching slice of `labels`. + size_t batch_rows = + std::min(static_cast(n_rows), std::max(1, params.streaming_batch_size)); + auto d_batch = raft::make_device_mdarray( + handle, big_memory_resource, raft::make_extents(batch_rows, dim)); + for (size_t row_start = 0; row_start < static_cast(n_rows); row_start += batch_rows) { + size_t rows = std::min(batch_rows, static_cast(n_rows) - row_start); + raft::copy(d_batch.data_handle(), host_dataset_ptr + row_start * dim, rows * dim, stream); + auto batch_view = + raft::make_device_matrix_view(d_batch.data_handle(), rows, dim); + auto batch_labels = + raft::make_device_vector_view(labels.data() + row_start, rows); + cuvs::cluster::kmeans::predict( + handle, kmeans_params, batch_view, centers_const_view, batch_labels); + } + } else { + cuvs::cluster::kmeans::predict( + handle, kmeans_params, dataset_const_view, centers_const_view, labels_view); + } + } + + index index(handle, n_rows, dim, params.n_lists, params.bits_per_dim); + + // Call RaBitQ index construct - use streaming if dataset doesn't fit in GPU memory + if (use_streaming) { + index.rabitq_index().construct_on_gpu_streaming(host_dataset_ptr, + cluster_centers.data(), + labels.data(), + params.fast_quantize_flag, + params.streaming_batch_size); + } else { + index.rabitq_index().construct_on_gpu(d_dataset_view.data_handle(), + cluster_centers.data(), + labels.data(), + params.fast_quantize_flag); + } + + return index; +} + +template +void search(raft::resources const& handle, + const search_params& params, + index& idx, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances) +{ + auto stream = raft::resource::get_cuda_stream(handle).value(); + + size_t NQ = queries.extent(0); + size_t k = neighbors.extent(1); + if (NQ == 0 || k == 0 || params.n_probes == 0) return; + size_t dim = queries.extent(1); + + RAFT_EXPECTS(neighbors.extent(0) == static_cast(NQ), + "neighbors rows must equal queries rows"); + RAFT_EXPECTS(distances.extent(0) == static_cast(NQ), + "distances rows must equal queries rows"); + RAFT_EXPECTS(distances.extent(1) == static_cast(k), + "distances cols must equal neighbors cols"); + RAFT_EXPECTS(dim == static_cast(idx.dim()), + "query dimensionality (%zu) must match index dimensionality (%u)", + dim, + idx.dim()); + RAFT_EXPECTS(params.n_probes <= idx.rabitq_index().get_num_centroids(), + "n_probes (%u) cannot exceed number of IVF lists (%zu)", + params.n_probes, + idx.rabitq_index().get_num_centroids()); + + auto padded_dim = idx.rabitq_index().get_num_padded_dim(); + auto rotated_queries = raft::make_device_matrix(handle, NQ, padded_dim); + if (padded_dim == dim) { + // TODO: replace RotatorGPU::rotate with cuVS/RAFT primitives + idx.rabitq_index().rotator().rotate( + handle, queries.data_handle(), rotated_queries.data_handle(), NQ); + } else { + auto padded_queries = raft::make_device_matrix(handle, NQ, padded_dim); + RAFT_CUDA_TRY( + cudaMemsetAsync(padded_queries.data_handle(), 0, sizeof(T) * NQ * padded_dim, stream)); + RAFT_CUDA_TRY(cudaMemcpy2DAsync(padded_queries.data_handle(), + sizeof(T) * padded_dim, + queries.data_handle(), + sizeof(T) * dim, + sizeof(T) * dim, + NQ, + cudaMemcpyDefault, + stream)); + // TODO: replace RotatorGPU::rotate with cuVS/RAFT primitives + idx.rabitq_index().rotator().rotate( + handle, padded_queries.data_handle(), rotated_queries.data_handle(), NQ); + } + + auto search_mode_to_string = [](search_mode mode) -> std::string { + switch (mode) { + case search_mode::LUT16: return std::string("lut16"); + case search_mode::LUT32: return std::string("lut32"); + case search_mode::QUANT4: return std::string("quant4"); + case search_mode::QUANT8: return std::string("quant8"); + default: RAFT_FAIL("Invalid search mode"); + } + }; + detail::SearcherGPU searcher(handle, + rotated_queries.data_handle(), + padded_dim, + idx.rabitq_index().get_ex_bits(), + search_mode_to_string(params.mode), + idx.rabitq_index().quantizer().get_query_scaling_factor(), + /* rabitq_quantize_flag = */ true); + searcher.AllocateSearcherSpace(idx.rabitq_index().get_num_centroids(), NQ); + + auto final_ids = raft::make_device_matrix(handle, NQ, k); + auto queries_view = + raft::make_device_matrix_view(rotated_queries.data_handle(), NQ, padded_dim); + + if (params.mode == search_mode::LUT32) { + idx.rabitq_index().BatchClusterSearch( + queries_view, k, params.n_probes, searcher, NQ, distances, final_ids.view()); + } 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()); + } else if (params.mode == search_mode::QUANT8) { + idx.rabitq_index().BatchClusterSearchQuantizeQuery( + queries_view, k, params.n_probes, searcher, NQ, distances, final_ids.view(), 8); + } else if (params.mode == search_mode::QUANT4) { + idx.rabitq_index().BatchClusterSearchQuantizeQuery( + queries_view, k, params.n_probes, searcher, NQ, distances, final_ids.view(), 4); + } + + // cast data in final_ids to array of IdxT in neighbors + raft::linalg::map( + handle, + neighbors, + raft::cast_op{}, + raft::make_device_vector_view(final_ids.data_handle(), NQ * k)); +} + +template +void serialize(raft::resources const& handle, const std::string& filename, index& index) +{ + // Save the index to a file. + index.rabitq_index().save(filename.c_str()); +} + +template +void deserialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::ivf_rabitq::index* index) +{ + index->rabitq_index().load_transposed(filename.c_str()); +} + +} // namespace detail + +template +index::index(raft::resources const& handle) + // this constructor is just for a temporary index, for use in the deserialization + // api. all the parameters here will get replaced with loaded values - that aren't + // necessarily known ahead of time before deserialization. + // TODO: do we even need a handle here - could just construct one? + : rabitq_index_(std::make_unique(handle)) +{ +} + +template +index::index(raft::resources const& handle, + size_t n_rows, + uint32_t dim, + uint32_t n_lists, + uint32_t bits_per_dim) +{ + RAFT_EXPECTS(bits_per_dim >= 1 && bits_per_dim <= 9, "Unsupported bits_per_dim"); + rabitq_index_ = std::make_unique(handle, n_rows, dim, n_lists, bits_per_dim); +} + +template +index::index(index&&) = default; + +template +auto index::operator=(index&&) -> index& = default; + +template +index::~index() = default; + +template struct index; + +template +detail::IVFGPU& index::rabitq_index() noexcept +{ + return *rabitq_index_; +} + +template +uint32_t index::dim() const noexcept +{ + return rabitq_index_->get_num_dimensions(); +} + +template +IdxT index::size() const noexcept +{ + return rabitq_index_->get_num_vectors(); +} + +auto build(raft::resources const& handle, + const cuvs::neighbors::ivf_rabitq::index_params& index_params, + raft::device_matrix_view dataset) + -> cuvs::neighbors::ivf_rabitq::index +{ + return cuvs::neighbors::ivf_rabitq::detail::build(handle, index_params, dataset); +} + +auto build(raft::resources const& handle, + const cuvs::neighbors::ivf_rabitq::index_params& index_params, + raft::host_matrix_view dataset) + -> cuvs::neighbors::ivf_rabitq::index +{ + return cuvs::neighbors::ivf_rabitq::detail::build(handle, index_params, dataset); +} + +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) +{ + cuvs::neighbors::ivf_rabitq::detail::search( + handle, search_params, index, queries, neighbors, distances); +} + +void serialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::ivf_rabitq::index& index) +{ + cuvs::neighbors::ivf_rabitq::detail::serialize(handle, filename, index); +} + +void deserialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::ivf_rabitq::index* index) +{ + if (!index) { RAFT_FAIL("Invalid index pointer"); } + cuvs::neighbors::ivf_rabitq::detail::deserialize(handle, filename, index); +} + +} // namespace cuvs::neighbors::ivf_rabitq diff --git a/cpp/src/neighbors/ivf_rabitq/defines.hpp b/cpp/src/neighbors/ivf_rabitq/defines.hpp new file mode 100644 index 0000000000..f91c06a1ff --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/defines.hpp @@ -0,0 +1,30 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +using PID = uint32_t; + +struct Candidate { + PID id; + float distance; + + Candidate() = default; + Candidate(PID id, float distance) : id(id), distance(distance) {} + + bool operator<(const Candidate& other) const { return distance < other.distance; } + + bool operator>(const Candidate& other) const { return !(*this < other); } +}; + +struct ExFactor { + float xipnorm; +}; + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu new file mode 100644 index 0000000000..6755bfc3de --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +// +// Created by Stardust on 3/3/25. +// + +#include "initializer_gpu.cuh" + +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +FlatInitializerGPU::FlatInitializerGPU(raft::resources const& handle, size_t d, size_t k) + : InitializerGPU(handle, d, k), + centroids_(raft::make_device_matrix(handle_, K, D)) +{ +} + +__host__ __device__ float* FlatInitializerGPU::GetCentroid(PID id) const +{ + return const_cast(centroids_.data_handle()) + id * D; +} + +void FlatInitializerGPU::AddVectors(const float* cent) +{ + RAFT_EXPECTS(cent != nullptr, "FlatInitializerGPU::AddVectors: cent is null"); + raft::copy(centroids_.data_handle(), cent, data_elements(), stream_); +} + +// Loads the centroids from a file. +// The centroids are read into a temporary host buffer and then copied to the +// device memory. +void FlatInitializerGPU::LoadCentroids(std::ifstream& input, const char* filename) +{ + // Allocate temporary host buffer for centroids. + auto host_buf = raft::make_host_vector(K * D); + // Read the raw data from the file. + input.read(reinterpret_cast(host_buf.data_handle()), data_bytes()); + RAFT_EXPECTS(input && input.gcount() == static_cast(data_bytes()), + "Failed to read centroids from file: %s", + filename); + // Copy the data from host to device. + raft::copy(centroids_.data_handle(), host_buf.data_handle(), data_elements(), stream_); + raft::resource::sync_stream(handle_); +} + +// Saves the centroids to a file. +// The centroids are first copied from device memory to a temporary host buffer, +// and then the entire block of data is written to the output stream. +void FlatInitializerGPU::SaveCentroids(std::ofstream& output, const char* filename) const +{ + // Allocate temporary host buffer for centroids. + auto host_buf = raft::make_host_vector(K * D); + // Copy centroids from device to host. + raft::copy(host_buf.data_handle(), centroids_.data_handle(), data_elements(), stream_); + raft::resource::sync_stream(handle_); + // Write the raw data to the file. + output.write(reinterpret_cast(host_buf.data_handle()), data_bytes()); + RAFT_EXPECTS(static_cast(output), "Failed to write centroids to file: %s", filename); +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh new file mode 100644 index 0000000000..0640148814 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh @@ -0,0 +1,108 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +// +// Created by Stardust on 3/3/25. +// + +#pragma once + +#include "../defines.hpp" + +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +class InitializerGPU { + public: + /** + * @brief Constructor. + * + * @param d dimension of vectors + * @param k number of centroids + */ + explicit InitializerGPU(raft::resources const& handle, size_t d, size_t k) + : D(d), K(k), handle_(handle) + { + RAFT_EXPECTS(d > 0, "InitializerGPU: d must be > 0"); + RAFT_EXPECTS(k > 0, "InitializerGPU: k must be > 0"); + } + + virtual ~InitializerGPU() = default; + + /** + * @brief Get a getCentroidbyId's vector data. Use ID as index. + * + * @param id Index. + * @return + */ + virtual __host__ __device__ float* GetCentroid(PID id) const = 0; + + /** + * @brief Copies centroids from cent into device memory (K * D floats). + * + * The copy is issued asynchronously on the internal CUDA stream. The caller must keep cent + * valid until that stream has been synchronized (or until any subsequent work on the stream + * consuming the centroids has been submitted). + * + * @param cent pointer to centroids (host or device) + */ + virtual void AddVectors(const float* cent) = 0; + + /** + * @brief LoadCentroids centroids' information from files. + * + * @param input + * @param filename + */ + virtual void LoadCentroids(std::ifstream& input, const char* filename) = 0; + + /** + * @brief SaveCentroids centroids' information from files. + * + * @param save + * @param filename + */ + virtual void SaveCentroids(std::ofstream& output, const char* filename) const = 0; + + protected: + size_t D; // Dimension + size_t K; // Num of Centroids + raft::resources const& handle_; // reusable resource handle + rmm::cuda_stream_view stream_ = + raft::resource::get_cuda_stream(handle_); // CUDA stream obtained from handle_ +}; + +class FlatInitializerGPU : public InitializerGPU { + public: + explicit FlatInitializerGPU(raft::resources const& handle, size_t d, size_t k); + + [[nodiscard]] __host__ __device__ float* GetCentroid(PID id) const override; + + void AddVectors(const float* cent) override; + + void LoadCentroids(std::ifstream& input, const char* filename) override; + + void SaveCentroids(std::ofstream& output, const char* filename) const override; + + private: + // D, K are inherited from parent + + raft::device_matrix + centroids_; // Stored in GPU device memory. Points to the parent centroids' array + + [[nodiscard]] size_t data_bytes() const noexcept { return sizeof(float) * K * D; } + [[nodiscard]] size_t data_elements() const noexcept { return K * D; } +}; + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu new file mode 100644 index 0000000000..49070acbcf --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu @@ -0,0 +1,1206 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +// +// Created by Stardust on 2/23/25. +// + +#include "../../../util/serialize_validation.hpp" +#include "../../detail/ann_utils.cuh" +#include "../utils/reductions.cuh" +#include "ivf_gpu.cuh" +#include "searcher_gpu.cuh" +#include + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include + +#include + +#include +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +IVFGPU::IVFGPU(raft::resources const& handle, size_t n, size_t dim, size_t k, size_t bits_per_dim) + : handle_(handle), + num_vectors(n), + num_dimensions(dim), + num_padded_dim(raft::round_up_safe(dim, 64)), + num_centroids(k), + ex_bits(bits_per_dim - 1), + initializer(nullptr), + DQ(std::make_unique(handle_, dim, bits_per_dim - 1)), + Rota(std::make_unique(handle_, dim)) +{ +} + +void IVFGPU::AllocateDeviceMemory() +{ + this->initializer = std::make_unique(handle_, num_padded_dim, num_centroids); + + // Compute the sizes for each data array. + size_t short_data_size = GetShortDataBytesSimple(); + size_t long_code_size = GetLongCodeBytes(); + size_t ex_factor_size = GetExFactorBytes(); + size_t pids_size = GetPIDsBytes(); + + // Allocate memory for the quantized arrays (SoA layout). + short_data_ = + raft::make_device_vector(handle_, short_data_size / sizeof(uint32_t)); + short_factors_batch_ = raft::make_device_vector( + handle_, GetShortDataFactorBytesBatch() / sizeof(float)); + long_code_ = + raft::make_device_vector(handle_, long_code_size / sizeof(uint8_t)); + ex_factor_ = + raft::make_device_vector(handle_, ex_factor_size / sizeof(ExFactor)); + ids_ = raft::make_device_vector(handle_, pids_size / sizeof(PID)); + + // Allocate memory for the per-cluster metadata and centroids. + cluster_meta_ = raft::make_device_vector(handle_, num_centroids); + raft::resource::sync_stream(handle_); +} + +void IVFGPU::AllocateHostMemory() +{ + // Compute the sizes for each data array. + size_t short_data_size = GetShortDataBytesSimple(); + size_t long_code_size = GetLongCodeBytes(); + size_t ex_factor_size = GetExFactorBytes(); + size_t pids_size = GetPIDsBytes(); + + this->short_data_host_ = + raft::make_host_vector(short_data_size / sizeof(uint32_t)); + this->long_code_host_ = + raft::make_host_vector(long_code_size / sizeof(uint8_t)); + this->ex_factor_host_ = + raft::make_host_vector(ex_factor_size / sizeof(ExFactor)); + this->ids_host_ = raft::make_host_vector(pids_size / sizeof(PID)); +} + +// load transposed data for short codes +void IVFGPU::load_transposed(const char* filename) +{ + std::ifstream input(filename, std::ios::binary); + RAFT_EXPECTS(input.is_open(), "failed to open file: %s", filename); + + auto read_exact = [&](void* ptr, size_t n_bytes) { + input.read(reinterpret_cast(ptr), n_bytes); + RAFT_EXPECTS(input.gcount() == static_cast(n_bytes), + "unexpected EOF reading header from: %s", + filename); + }; + + // Load metadata. + read_exact(&this->num_vectors, sizeof(size_t)); + read_exact(&this->num_dimensions, sizeof(size_t)); + // Compute padded dimension. + this->num_padded_dim = raft::round_up_safe(this->num_dimensions, 64); + read_exact(&this->num_centroids, sizeof(size_t)); + read_exact(&this->ex_bits, sizeof(size_t)); + + // Skip legacy batch_flag field for backward compatibility + bool legacy_batch_flag; + read_exact(&legacy_batch_flag, sizeof(bool)); + + // Validate header fields to guard against malicious or corrupted files. With these bounds, + // downstream `num_vectors * per_vector_bytes` products in AllocateDeviceMemory/the Get*Bytes + // helpers cannot overflow size_t (per_vector_bytes is bounded by num_padded_dim). + RAFT_EXPECTS(num_dimensions > 0, "ivf_rabitq::deserialize: num_dimensions=0 in: %s", filename); + RAFT_EXPECTS(ex_bits < 9, + "ivf_rabitq::deserialize: ex_bits=%zu out of valid range [0, 9) in: %s", + 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", + 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 " + "in: %s", + num_vectors, + num_padded_dim, + filename); + + // Initialize quantizer and rotator (host objects that drive GPU routines). + this->DQ = std::make_unique(handle_, num_dimensions, ex_bits); + read_exact(this->DQ->get_query_scaling_factor(), sizeof(DataQuantizerGPU::FastQuantizeFactors)); + this->Rota = std::make_unique(handle_, num_dimensions); + // Load cluster sizes. + std::vector cluster_sizes(num_centroids, 0); + read_exact(cluster_sizes.data(), sizeof(size_t) * num_centroids); + RAFT_EXPECTS( + std::accumulate(cluster_sizes.begin(), cluster_sizes.end(), size_t(0)) == num_vectors, + "cluster sizes do not sum to num_vectors in: %s", + filename); + + // Load rotator from file. + this->rotator().load(handle_, input); + // Allocate device memory based on the cluster sizes. + AllocateDeviceMemory(); + // Load initializer data (e.g., centroids) from file. + this->initializer->LoadCentroids(input, filename); + compute_centroid_norms(); + // Read raw arrays from file into device memory. + auto read_into_device = [&](void* d_ptr, size_t n_bytes) { + std::vector h_buf(n_bytes); // host staging buffer + input.read(reinterpret_cast(h_buf.data()), n_bytes); + if (input.gcount() != static_cast(n_bytes)) + throw std::runtime_error("unexpected EOF"); + + raft::copy(static_cast(d_ptr), h_buf.data(), n_bytes, stream_); + raft::resource::sync_stream(handle_); + }; + + auto read_into_device_host = [&](void* d_ptr, void* h_ptr, size_t n_bytes) { + std::vector h_buf(n_bytes); // host staging buffer + auto before = input.tellg(); + input.read(reinterpret_cast(h_buf.data()), n_bytes); + auto got = static_cast(input.gcount()); + RAFT_EXPECTS(got == n_bytes, + "unexpected EOF: wanted %zu bytes at offset %ld, got %zu%s%s", + n_bytes, + static_cast(before), + got, + input.eof() ? " (hit EOF)" : "", + input.bad() ? " (I/O error)" : ""); + + raft::copy(static_cast(d_ptr), h_buf.data(), n_bytes, stream_); + raft::resource::sync_stream(handle_); + memcpy(h_ptr, h_buf.data(), n_bytes); + }; + + auto read_into_device_host_transposed_short = + [&](void* d_ptr, void* h_ptr, size_t n_bytes, size_t& max_cluster_size) { + // Read all clusters' data into staging buffer (still in sequential format) + std::vector h_buf(n_bytes); + input.read(reinterpret_cast(h_buf.data()), n_bytes); + if (input.gcount() != static_cast(n_bytes)) + throw std::runtime_error("unexpected EOF"); + + // Create transposed buffer + std::vector h_transposed(n_bytes); + + // Process each cluster and transpose its vectors + size_t src_offset = 0; // offset in original sequential buffer + size_t dst_offset = 0; // offset in transposed buffer + + // Also evaluate maximum cluster size during this loop + max_cluster_size = 0; + for (size_t cluster_id = 0; cluster_id < num_centroids; cluster_id++) { + size_t cluster_size = cluster_sizes[cluster_id]; + max_cluster_size = max(max_cluster_size, cluster_size); + if (cluster_size == 0) continue; + + // Calculate dimensions per vector + size_t bytes_per_vector = DQ->block_bytes(); + size_t uint32s_per_vector = bytes_per_vector / sizeof(uint32_t); + + // Get pointers to source (sequential) and destination (transposed) data + uint32_t* src_cluster = reinterpret_cast(h_buf.data() + src_offset); + uint32_t* dst_cluster = reinterpret_cast(h_transposed.data() + dst_offset); + + // Transpose the cluster: + // From: vec1[all_dims], vec2[all_dims], ..., vecn[all_dims] + // To: vec1[dim0-31], vec2[dim0-31], ..., vecn[dim0-31], + // vec1[dim32-63], vec2[dim32-63], ..., vecn[dim32-63], ... + + for (size_t dim_chunk = 0; dim_chunk < uint32s_per_vector; dim_chunk++) { + for (size_t vec_id = 0; vec_id < cluster_size; vec_id++) { + // Source: vector vec_id, dimension chunk dim_chunk + size_t src_idx = vec_id * uint32s_per_vector + dim_chunk; + // Destination: dimension chunk dim_chunk, vector vec_id + size_t dst_idx = dim_chunk * cluster_size + vec_id; + + dst_cluster[dst_idx] = src_cluster[src_idx]; + } + } + + // Update offsets for next cluster + size_t cluster_bytes = cluster_size * bytes_per_vector; + src_offset += cluster_bytes; + dst_offset += cluster_bytes; + } + + // Copy transposed data to device and host + raft::copy(static_cast(d_ptr), h_transposed.data(), n_bytes, stream_); + raft::resource::sync_stream(handle_); + memcpy(h_ptr, h_transposed.data(), n_bytes); + }; + + // New change: host copy of ivf. + AllocateHostMemory(); + read_into_device_host_transposed_short(short_data_.data_handle(), + short_data_host_.data_handle(), + GetShortDataBytesSimple(), + this->max_cluster_length); + read_into_device(short_factors_batch_.data_handle(), + GetShortDataFactorBytesBatch()); // SoA layout - no copy on CPU + read_into_device_host( + long_code_.data_handle(), long_code_host_.data_handle(), GetLongCodeBytes()); + read_into_device_host( + ex_factor_.data_handle(), ex_factor_host_.data_handle(), GetExFactorBytes()); + read_into_device_host(ids_.data_handle(), ids_host_.data_handle(), GetPIDsBytes()); + + // Initialize cluster metadata (host side) based on the loaded cluster sizes. + init_clusters(cluster_sizes); + + input.close(); +} + +void IVFGPU::compute_centroid_norms() +{ + centroid_norms_ = raft::make_device_vector(handle_, num_centroids); + auto centroids_view = raft::make_device_matrix_view( + initializer->GetCentroid(0), num_centroids, num_padded_dim); + raft::linalg::norm( + handle_, centroids_view, centroid_norms_.view()); +} + +void IVFGPU::init_clusters(const std::vector& cluster_sizes) +{ + // Allocate a host vector to hold cluster metadata. + cluster_meta_host_ = raft::make_host_vector(num_centroids); + + size_t added_vectors = 0; + for (size_t i = 0; i < num_centroids; ++i) { + size_t num = cluster_sizes[i]; + cluster_meta_host_(i) = {num, added_vectors}; + added_vectors += num; + } + + // Copy the host cluster metadata to device memory. + // cluster_meta_ must have been allocated with size: num_centroids * sizeof(GPUClusterMeta) + raft::copy(cluster_meta_.data_handle(), cluster_meta_host_.data_handle(), num_centroids, stream_); + raft::resource::sync_stream(handle_); +} + +// for debug use + +// Extract i-th code (MSB-first packing) from a bitstream. +// ex_bits in [1..7]. bytes_len is the byte-length of the code stream. +static inline uint32_t extract_code_msb(const uint8_t* bytes, + size_t bytes_len, + size_t i, + int ex_bits) +{ + const size_t bit_index = i * ex_bits; // absolute bit index in stream + const size_t byte_idx = bit_index / 8; // byte holding the MSB of this code + const int bit_off = bit_index % 8; // 0 => starts at MSB of bytes[byte_idx] + + // Build a 16-bit big-endian window: [bytes[byte_idx] | bytes[byte_idx+1]] + const uint32_t b0 = (byte_idx < bytes_len) ? bytes[byte_idx] : 0u; + const uint32_t b1 = (byte_idx + 1 < bytes_len) ? bytes[byte_idx + 1] : 0u; + const uint32_t val = (b0 << 8) | b1; + + // Align MSB-first window so that the ex_bits land in the LSBs. + // Example: ex_bits=3, bit_off=0 -> take bits 15..13 -> shift by 13. + const int shift = 16 - ex_bits - bit_off; + return (val >> shift) & ((1u << ex_bits) - 1u); +} + +void print_first_vector(uint8_t* ex_data_, float* ex_factors_, size_t dim, int ex_bits) +{ + // 1) Number of bytes for the bit-packed codes (ceil). + const size_t code_bytes = (dim * ex_bits + 7) / 8; + + // If your format pads to 4-byte alignment before the floats, use this instead: + // const size_t code_bytes_padded = (code_bytes + 3) & ~size_t(3); + // and then read floats from ex_data_ + code_bytes_padded. + const uint8_t* ex_code = reinterpret_cast(ex_data_); + + // 2) Safely read the floats (avoid potential unaligned pointer UB). + const float* f_exadd = reinterpret_cast(ex_factors_); + const float* f_exrescale = reinterpret_cast(ex_factors_ + 1); + + // 3) Print first 20 codes (MSB-first, left->right within each byte). + std::cout << "ExCode (first 20 dims): "; + const int to_print = static_cast(std::min(20, dim)); + for (int i = 0; i < to_print; ++i) { + uint32_t code = extract_code_msb(ex_code, code_bytes, static_cast(i), ex_bits); + std::cout << code << (i + 1 == to_print ? '\n' : ' '); + } + + std::cout << "F_exadd = " << *f_exadd << "\n"; + std::cout << "F_exrescale = " << *f_exrescale << "\n"; +} + +void IVFGPU::save(const char* filename) const +{ + RAFT_EXPECTS(num_centroids > 0, "IVF index has not been constructed"); + + std::ofstream output(filename, std::ios::binary); + RAFT_EXPECTS(output.is_open(), "failed to open file for saving: %s", filename); + + auto write_exact = [&](const void* ptr, size_t n_bytes) { + output.write(reinterpret_cast(ptr), n_bytes); + RAFT_EXPECTS(static_cast(output), "write failed to: %s", filename); + }; + + // Save meta data. + write_exact(&num_vectors, sizeof(size_t)); + write_exact(&num_dimensions, sizeof(size_t)); + write_exact(&num_centroids, sizeof(size_t)); + write_exact(&ex_bits, sizeof(size_t)); + // Write legacy batch_flag=true for backward compatibility + bool legacy_batch_flag = true; + write_exact(&legacy_batch_flag, sizeof(bool)); + write_exact(DQ->get_query_scaling_factor(), sizeof(DataQuantizerGPU::FastQuantizeFactors)); + + // Save number of vectors of each cluster. + std::vector h_cluster_meta(num_centroids); + raft::copy(h_cluster_meta.data(), cluster_meta_.data_handle(), num_centroids, stream_); + std::vector cluster_sizes(num_centroids); + raft::resource::sync_stream(handle_); + for (size_t i = 0; i < num_centroids; i++) { + cluster_sizes[i] = h_cluster_meta[i].num; + } + write_exact(cluster_sizes.data(), sizeof(size_t) * num_centroids); + + // Save rotator. + this->rotator().save(handle_, output); + + // Save initializer data. + this->initializer->SaveCentroids(output, filename); + + // Compute sizes for device arrays. + size_t short_data_size = GetShortDataBytesSimple(); + size_t long_code_size = GetLongCodeBytes(); + size_t ex_factor_size = GetExFactorBytes(); + size_t ids_size = GetPIDsBytes(); + // for batch data + size_t short_factors_size = GetShortDataFactorBytesBatch(); + + // Allocate temporary host buffers. + auto h_short_data_buf = raft::make_host_vector(short_data_size); + auto h_long_code_buf = raft::make_host_vector(long_code_size); + auto h_ex_factor_buf = + raft::make_host_vector(ex_factor_size / sizeof(ExFactor)); + auto h_ids_buf = raft::make_host_vector(ids_size / sizeof(PID)); + auto h_short_factors_batch_buf = raft::make_host_vector(short_factors_size); + + // Copy device data to host (SoA layout). + raft::copy(h_short_data_buf.data_handle(), + reinterpret_cast(short_data_.data_handle()), + short_data_size, + stream_); + raft::copy(h_short_factors_batch_buf.data_handle(), + reinterpret_cast(short_factors_batch_.data_handle()), + short_factors_size, + stream_); + raft::copy(h_long_code_buf.data_handle(), long_code_.data_handle(), long_code_size, stream_); + raft::copy(h_ex_factor_buf.data_handle(), + ex_factor_.data_handle(), + ex_factor_size / sizeof(ExFactor), + stream_); + raft::copy(h_ids_buf.data_handle(), ids_.data_handle(), ids_size / sizeof(PID), stream_); + raft::resource::sync_stream(handle_); + + // Write raw arrays to file. + write_exact(h_short_data_buf.data_handle(), short_data_size); + write_exact(h_short_factors_batch_buf.data_handle(), short_factors_size); + write_exact(h_long_code_buf.data_handle(), long_code_size); + write_exact(h_ex_factor_buf.data_handle(), ex_factor_size); + write_exact(h_ids_buf.data_handle(), ids_size); + + output.close(); +} + +/** + * @brief Build GPUClusterMeta from counts and offsets on GPU + */ +__global__ void build_cluster_meta_kernel(IVFGPU::GPUClusterMeta* d_cluster_meta, + const unsigned long long* d_counts, + const size_t* d_offsets, + size_t num_centroids) +{ + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_centroids) { + d_cluster_meta[idx] = IVFGPU::GPUClusterMeta(d_counts[idx], d_offsets[idx]); + } +} + +/** + * @brief Scatter PIDs into flat array based on cluster assignment + * Uses atomics to determine write position within each cluster + */ +__global__ void scatter_pids_kernel(PID* d_flat_pids, + const PID* d_cluster_ids, + const size_t* d_offsets, + size_t* d_atomic_counters, // Initialized to d_offsets values + size_t num_vectors) +{ + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_vectors) { + PID cid = d_cluster_ids[idx]; + // Atomically get position and increment + size_t pos = atomicAdd((unsigned long long*)&d_atomic_counters[cid], 1ULL); + d_flat_pids[pos] = static_cast(idx); + } +} + +// ============================================ +// Modified construct function +// ============================================ + +void IVFGPU::construct_on_gpu(const float* device_data, + const float* device_centroids, + const PID* device_cluster_ids, + bool fast_quantize) +{ + DQ->fast_quantize_flag = fast_quantize; + + // pre-compute rescaling factors for search + DQ->compute_query_scaling_factors(this->num_padded_dim); + + // compute rescaling factors for query if needed + if (DQ->fast_quantize_flag) { DQ->compute_quantize_scaling_factors(); } + + if (DQ->fast_quantize_flag) { + if (ex_bits == 3) { + DQ->set_quantize_scaling_factors(DQ->get_query_scaling_factor()->const_scaling_factor_4bit); + } else if (ex_bits == 7) { + DQ->set_quantize_scaling_factors(DQ->get_query_scaling_factor()->const_scaling_factor_8bit); + } else { + DQ->compute_quantize_scaling_factors(); + } + } + + // ------------------------- + // 2. Validate cluster IDs on GPU + // ------------------------- + int block_size = 256; + int num_blocks = (num_vectors + block_size - 1) / block_size; + + // ------------------------- + // 9. Allocate device memory for IVF arrays + // ------------------------- + std::vector counts(num_centroids); // counts is not used here + AllocateDeviceMemory(); + + // ------------------------- + // 3. Compute histogram (cluster sizes) on GPU using CUB + // ------------------------- + + // Use CUB's DeviceHistogram + using CounterT = unsigned long long; + + int num_levels = num_centroids + 1; + int lower_level = 0; + int upper_level = num_centroids; + + rmm::device_uvector d_histogram(num_centroids, stream_); + + // Determine temporary device storage requirements + size_t temp_storage_bytes = 0; + RAFT_CUDA_TRY(cub::DeviceHistogram::HistogramEven(nullptr, + temp_storage_bytes, + device_cluster_ids, + d_histogram.data(), + num_levels, + lower_level, + upper_level, + num_vectors, + stream_)); + + { + rmm::device_buffer d_temp_storage(temp_storage_bytes, stream_); + RAFT_CUDA_TRY(cub::DeviceHistogram::HistogramEven(d_temp_storage.data(), + temp_storage_bytes, + device_cluster_ids, + d_histogram.data(), + num_levels, + lower_level, + upper_level, + num_vectors, + stream_)); + } + + // ------------------------- + // 4. Compute prefix sum (offsets) on GPU using CUB + // ------------------------- + rmm::device_uvector d_offsets(num_centroids + 1, stream_); + + temp_storage_bytes = 0; + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum( + nullptr, temp_storage_bytes, d_histogram.data(), d_offsets.data(), num_centroids, stream_)); + + { + rmm::device_buffer d_temp_storage(temp_storage_bytes, stream_); + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(d_temp_storage.data(), + temp_storage_bytes, + d_histogram.data(), + d_offsets.data(), + num_centroids, + stream_)); + } + + // Set the last offset element + raft::copy(d_offsets.data() + num_centroids, &num_vectors, 1, stream_); + + // ------------------------- + // 5. Build cluster metadata on GPU + // ------------------------- + GPUClusterMeta* d_cluster_meta_temp = cluster_meta_.data_handle(); + + num_blocks = (num_centroids + block_size - 1) / block_size; + build_cluster_meta_kernel<<>>( + d_cluster_meta_temp, d_histogram.data(), d_offsets.data(), num_centroids); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // ------------------------- + // 6. Scatter PIDs to flat array on GPU + // ------------------------- + PID* d_flat_pids = ids_.data_handle(); + + // Create atomic counters initialized with offsets + rmm::device_uvector d_atomic_counters(num_centroids, stream_); + raft::copy(d_atomic_counters.data(), d_offsets.data(), num_centroids, stream_); + + num_blocks = (num_vectors + block_size - 1) / block_size; + scatter_pids_kernel<<>>( + d_flat_pids, device_cluster_ids, d_offsets.data(), d_atomic_counters.data(), num_vectors); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // ------------------------- + // 7. Copy cluster metadata back to host + // ------------------------- + std::vector h_cluster_meta(num_centroids); + raft::copy(h_cluster_meta.data(), d_cluster_meta_temp, num_centroids, stream_); + raft::resource::sync_stream(handle_); + + // ------------------------- + // 10. Allocate and process rotated centroids + // ------------------------- + rmm::device_uvector d_rotated_centroids(num_centroids * num_padded_dim, stream_); + + // ------------------------- + // 10.5 get max cluster length and allocate temporary buffer for quantization + // Do note that this update will disable the ability of quantization with multiple streams + // ------------------------- + max_cluster_length = 0; + for (const auto& meta : h_cluster_meta) { + max_cluster_length = std::max(max_cluster_length, meta.num); + } + DQ->alloc_buffers(max_cluster_length); + + // Process clusters sequentially + for (size_t i = 0; i < num_centroids; ++i) { + const float* cur_centroid = device_centroids + i * num_dimensions; + float* cur_rotated_c = d_rotated_centroids.data() + i * num_padded_dim; + GPUClusterMeta& cp = h_cluster_meta[i]; + quantize_cluster(cp, device_data, cur_centroid, cur_rotated_c); + } + + // Add rotated centroids + initializer->AddVectors(d_rotated_centroids.data()); + + raft::resource::sync_stream(handle_); +} + +void IVFGPU::construct_on_gpu_streaming(const float* host_data, + const float* device_centroids, + const PID* device_cluster_ids, + bool fast_quantize, + size_t batch_size_vectors) +{ + DQ->fast_quantize_flag = fast_quantize; + + // pre-compute rescaling factors for search + DQ->compute_query_scaling_factors(this->num_padded_dim); + + // compute rescaling factors for query if needed + if (DQ->fast_quantize_flag) { DQ->compute_quantize_scaling_factors(); } + + if (DQ->fast_quantize_flag) { + if (ex_bits == 3) { + DQ->set_quantize_scaling_factors(DQ->get_query_scaling_factor()->const_scaling_factor_4bit); + } else if (ex_bits == 7) { + DQ->set_quantize_scaling_factors(DQ->get_query_scaling_factor()->const_scaling_factor_8bit); + } else { + DQ->compute_quantize_scaling_factors(); + } + } + + // ------------------------- + // 1. Use cluster IDs already on device + // ------------------------- + + // ------------------------- + // 2. Allocate device memory for IVF arrays + // ------------------------- + AllocateDeviceMemory(); + + // ------------------------- + // 3. Compute histogram (cluster sizes) on GPU using CUB + // ------------------------- + + using CounterT = unsigned long long; + + int num_levels = num_centroids + 1; + int lower_level = 0; + int upper_level = num_centroids; + + auto d_histogram = raft::make_device_vector(handle_, num_centroids); + + void* d_temp_storage = nullptr; + size_t temp_storage_bytes = 0; + RAFT_CUDA_TRY(cub::DeviceHistogram::HistogramEven(d_temp_storage, + temp_storage_bytes, + device_cluster_ids, + d_histogram.data_handle(), + num_levels, + lower_level, + upper_level, + num_vectors, + stream_)); + + auto d_temp_storage_vec = raft::make_device_vector(handle_, temp_storage_bytes); + d_temp_storage = d_temp_storage_vec.data_handle(); + + RAFT_CUDA_TRY(cub::DeviceHistogram::HistogramEven(d_temp_storage, + temp_storage_bytes, + device_cluster_ids, + d_histogram.data_handle(), + num_levels, + lower_level, + upper_level, + num_vectors, + stream_)); + + // ------------------------- + // 4. Compute prefix sum (offsets) on GPU using CUB + // ------------------------- + auto d_offsets = raft::make_device_vector(handle_, num_centroids + 1); + + d_temp_storage = nullptr; + temp_storage_bytes = 0; + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(d_temp_storage, + temp_storage_bytes, + d_histogram.data_handle(), + d_offsets.data_handle(), + num_centroids, + stream_)); + + d_temp_storage_vec = raft::make_device_vector(handle_, temp_storage_bytes); + d_temp_storage = d_temp_storage_vec.data_handle(); + + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(d_temp_storage, + temp_storage_bytes, + d_histogram.data_handle(), + d_offsets.data_handle(), + num_centroids, + stream_)); + + // Set the last offset element + raft::copy(d_offsets.data_handle() + num_centroids, &num_vectors, 1, stream_); + + // ------------------------- + // 5. Build cluster metadata on GPU + // ------------------------- + GPUClusterMeta* d_cluster_meta_temp = cluster_meta_.data_handle(); + + int block_size = 256; + int num_blocks = (num_centroids + block_size - 1) / block_size; + build_cluster_meta_kernel<<>>( + d_cluster_meta_temp, d_histogram.data_handle(), d_offsets.data_handle(), num_centroids); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // ------------------------- + // 6. Scatter PIDs to flat array on GPU + // ------------------------- + PID* d_flat_pids = ids_.data_handle(); + + // Create atomic counters initialized with offsets + auto d_atomic_counters = raft::make_device_vector(handle_, num_centroids); + raft::copy(d_atomic_counters.data_handle(), d_offsets.data_handle(), num_centroids, stream_); + + num_blocks = (num_vectors + block_size - 1) / block_size; + scatter_pids_kernel<<>>(d_flat_pids, + device_cluster_ids, + d_offsets.data_handle(), + d_atomic_counters.data_handle(), + num_vectors); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // ------------------------- + // 7. Copy cluster metadata back to host for batching + // ------------------------- + std::vector h_cluster_meta(num_centroids); + raft::copy(h_cluster_meta.data(), d_cluster_meta_temp, num_centroids, stream_); + raft::resource::sync_stream(handle_); + + // ------------------------- + // 8. Compute max cluster length and allocate buffers + // ------------------------- + max_cluster_length = 0; + for (size_t i = 0; i < num_centroids; ++i) { + max_cluster_length = std::max(max_cluster_length, h_cluster_meta[i].num); + } + RAFT_EXPECTS(max_cluster_length <= batch_size_vectors, + "max cluster size (%zu) exceeds batch_size_vectors (%zu); " + "increase batch_size_vectors so every cluster fits in one batch", + max_cluster_length, + batch_size_vectors); + DQ->alloc_buffers(max_cluster_length); + + // ------------------------- + // 9. Allocate rotated centroids buffer + // ------------------------- + auto d_rotated_centroids = + raft::make_device_vector(handle_, num_centroids * num_padded_dim); + + // ------------------------- + // 10. Batch clusters and stream to GPU + // ------------------------- + // Preallocate batch buffers (reused across all batch iterations) + auto batch_pids = raft::make_host_vector(batch_size_vectors); + auto batch_data = raft::make_host_vector(batch_size_vectors * num_dimensions); + auto d_batch_data = + raft::make_device_vector(handle_, batch_size_vectors * num_dimensions); + + size_t batch_count = 0; + + // Calculate number of threads for host data gathering + const size_t num_threads = omp_get_max_threads(); + + size_t cluster_idx = 0; + + while (cluster_idx < num_centroids) { + batch_count++; + // Determine batch: complete clusters only + size_t batch_start_cluster = cluster_idx; + size_t batch_vectors = 0; + size_t batch_start_offset = h_cluster_meta[cluster_idx].start_index; + + while (cluster_idx < num_centroids && + batch_vectors + h_cluster_meta[cluster_idx].num <= batch_size_vectors) { + batch_vectors += h_cluster_meta[cluster_idx].num; + cluster_idx++; + } + + if (batch_vectors == 0) break; // No more clusters to process + + // ------------------------- + // 11. Gather and transfer batch data to GPU using reordered PIDs + // ------------------------- + + // First, copy the reordered PIDs for this batch to host (reuse preallocated buffer) + raft::copy(batch_pids.data_handle(), d_flat_pids + batch_start_offset, batch_vectors, stream_); + raft::resource::sync_stream(handle_); + + // Gather data on host using the reordered PIDs (reuse preallocated buffer) + // Use OpenMP parallel region to leverage persistent thread pool across batches +#pragma omp parallel num_threads(num_threads) + { + size_t tid = omp_get_thread_num(); + size_t chunk_size = batch_vectors / num_threads; + size_t start = tid * chunk_size; + size_t end = (tid == num_threads - 1) ? batch_vectors : start + chunk_size; + + for (size_t i = start; i < end; ++i) { + std::memcpy(&batch_data.data_handle()[i * num_dimensions], + &host_data[batch_pids.data_handle()[i] * num_dimensions], + num_dimensions * sizeof(float)); + } + } + // OpenMP implicit barrier ensures all threads complete before continuing + + // Transfer batch to GPU (reuse preallocated buffer) + raft::copy(d_batch_data.data_handle(), + batch_data.data_handle(), + batch_vectors * num_dimensions, + stream_); + + // ------------------------- + // 12. Quantize each cluster in the batch + // ------------------------- + + size_t batch_offset = 0; + for (size_t c = batch_start_cluster; c < cluster_idx; ++c) { + size_t cluster_size = h_cluster_meta[c].num; + if (cluster_size == 0) continue; + + const float* cur_centroid = device_centroids + c * num_dimensions; + float* cur_rotated_c = d_rotated_centroids.data_handle() + c * num_padded_dim; + GPUClusterMeta& cp = h_cluster_meta[c]; + + // Quantize this cluster's data (contiguous in batch, SoA layout) + const float* cluster_data = d_batch_data.data_handle() + batch_offset * num_dimensions; + + // Quantize using batch-optimized path + DQ->quantize_batch_opt_contiguous(cluster_data, + cur_centroid, + cluster_size, + this->rotator(), + cp.first_block(*this), + cp.short_factor_batch(*this, 0), + cp.long_code(*this, 0, DQ->long_code_length()), + reinterpret_cast(cp.ex_factor(*this, 0)), + cur_rotated_c); + + batch_offset += cluster_size; + } + } + + // ------------------------- + // 13. Add rotated centroids + // ------------------------- + initializer->AddVectors(d_rotated_centroids.data_handle()); + + raft::resource::sync_stream(handle_); +} + +void IVFGPU::quantize_cluster(GPUClusterMeta& cp, + const float* d_data, // device pointer + const float* d_centroid, // device pointer + float* d_rotated_c) const +{ // device pointer for rotated centroid + size_t num = cp.num; + // Copy the IDs for this cluster into device memory. + // Note: cp.ids(this) returns ids_ + cp.start_index. + const PID* idp = cp.ids(*this); + + // Call the GPU quantization function (SoA layout). + DQ->quantize_batch_opt(d_data, + d_centroid, + idp, + num, + this->rotator(), + cp.first_block(*this), + cp.short_factor_batch(*this, 0), + cp.long_code(*this, 0, DQ->long_code_length()), + reinterpret_cast(cp.ex_factor(*this, 0)), + d_rotated_c); +} + +// Optimized kernel to prepare keys and values from d_raft_idx +__global__ void prepare_keys_values( + const int* d_raft_idx, int* d_cluster_keys, int* d_query_values, int batch_size, int nprobe) +{ + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int total_pairs = batch_size * nprobe; + + if (tid < total_pairs) { + // Coalesced memory read + d_cluster_keys[tid] = d_raft_idx[tid]; + d_query_values[tid] = tid / nprobe; // query index + } +} + +// Optimized kernel to combine sorted keys/values into pairs +__global__ void combine_to_pairs(const int* d_sorted_clusters, + const int* d_sorted_queries, + ClusterQueryPair* d_sorted_pairs, + int total_pairs) +{ + int tid = blockIdx.x * blockDim.x + threadIdx.x; + + if (tid < total_pairs) { + // Coalesced writes to pair structure + d_sorted_pairs[tid].cluster_idx = d_sorted_clusters[tid]; + d_sorted_pairs[tid].query_idx = d_sorted_queries[tid]; + } +} + +// Sorting cluster_query pairs using cub sort +void sort_cluster_query_pairs(raft::resources const& handle, + int* d_raft_idx, // Input: cluster indices + ClusterQueryPair* d_sorted_pairs, // Output: sorted pairs + int batch_size, + int nprobe) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + int total_pairs = batch_size * nprobe; + + // Allocate temporary arrays for sorting + auto d_cluster_keys = raft::make_device_vector(handle, total_pairs); + auto d_query_values = raft::make_device_vector(handle, total_pairs); + auto d_sorted_clusters = raft::make_device_vector(handle, total_pairs); + auto d_sorted_queries = raft::make_device_vector(handle, total_pairs); + + // Prepare data for sorting + int threads_per_block = 256; + int num_blocks = (total_pairs + threads_per_block - 1) / threads_per_block; + + prepare_keys_values<<>>( + d_raft_idx, d_cluster_keys.data_handle(), d_query_values.data_handle(), batch_size, nprobe); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // Determine temporary storage requirements for CUB + // ---- sort by cluster id + size_t temp_storage_bytes = 0; + + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs(nullptr, + temp_storage_bytes, + d_cluster_keys.data_handle(), + d_sorted_clusters.data_handle(), + d_query_values.data_handle(), + d_sorted_queries.data_handle(), + total_pairs, + 0, + sizeof(int) * 8, + stream)); + + rmm::device_buffer d_temp_storage(temp_storage_bytes, stream); + + // Perform radix sort + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs(d_temp_storage.data(), + temp_storage_bytes, + d_cluster_keys.data_handle(), + d_sorted_clusters.data_handle(), + d_query_values.data_handle(), + d_sorted_queries.data_handle(), + total_pairs, + 0, + sizeof(int) * 8, + stream)); + + // Combine sorted results into pairs + combine_to_pairs<<>>( + d_sorted_clusters.data_handle(), d_sorted_queries.data_handle(), d_sorted_pairs, total_pairs); + RAFT_CUDA_TRY(cudaPeekAtLastError()); +} + +// Generic warp-level kernel: Multiple queries per block, one query per warp +template +__global__ void computeQueryFactorsWarpKernel(const T* d_query, // num_queries x D matrix + float* d_G_k1xSumq, // output: num_queries array + float* d_G_kbxSumq, // output: num_queries array + size_t num_queries, + size_t D, // dimension of each query + size_t ex_bits) // extra bits parameter +{ + // Calculate which query this warp is processing + const int warp_id = threadIdx.x / 32; + const int lane_id = threadIdx.x % 32; + const int warps_per_block = blockDim.x / 32; + const int query_idx = blockIdx.x * warps_per_block + warp_id; + + // Early exit if beyond valid queries + if (query_idx >= num_queries) return; + + // Compute constants + float c_1 = -static_cast((1 << 1) - 1) / 2.0f; // -0.5f + float c_b = -static_cast((1 << (ex_bits + 1)) - 1) / 2.0f; + + // Pointer to the current query + const T* query = d_query + query_idx * D; + + // Each thread in the warp accumulates multiple elements + T sum = 0; + + // Coalesced memory access: warp reads consecutive elements + for (int i = lane_id; i < D; i += 32) { + sum += query[i]; + } + + // Warp-level reduction using shuffle instructions + sum = warpReduceSum(sum); + + // Lane 0 writes the result + if (lane_id == 0) { + d_G_k1xSumq[query_idx] = sum * c_1; + d_G_kbxSumq[query_idx] = sum * c_b; + } +} + +// Wrapper function +template +void computeQueryFactors(raft::device_matrix_view queries, + raft::device_vector_view d_G_k1xSumq, + raft::device_vector_view d_G_kbxSumq, + size_t ex_bits, + cudaStream_t stream) +{ + const size_t num_queries = queries.extent(0); + const size_t D = queries.extent(1); + + // Use 256 threads per block (8 warps per block) + const int threads_per_block = 256; + const int warps_per_block = threads_per_block / 32; + const int blocks = (num_queries + warps_per_block - 1) / warps_per_block; + + computeQueryFactorsWarpKernel + <<>>(queries.data_handle(), + d_G_k1xSumq.data_handle(), + d_G_kbxSumq.data_handle(), + num_queries, + D, + ex_bits); + RAFT_CUDA_TRY(cudaPeekAtLastError()); +} + +// normal way to first sort (cluster, query) pairs, then use a CTA to do the search +void IVFGPU::PrepareClusterSearchInputs( + raft::device_matrix_view queries, + size_t nprobe, + SearcherGPU& searcher, + raft::device_vector& d_sorted_pairs, + raft::device_vector& d_G_k1xSumq, + raft::device_vector& d_G_kbxSumq) +{ + raft::resources const& searcher_handle = searcher.get_handle(); + rmm::cuda_stream_view searcher_stream = searcher.get_stream(); + const size_t batch_size = queries.extent(0); + + // Compute ||q - c||^2 = -2 * q . c + ||q||^2 + ||c||^2 into centroid_distances: + // (1) query norms (squared L2, no sqrt - default fin_op on L2Norm returns sum of squares; + // centroid norms were precomputed in compute_centroid_norms() and cached in centroid_norms_) + // (2) outer_add writes (q_norm + c_norm) into centroid_distances, overwriting the prior + // uninitialized contents of the freshly-allocated SearcherGPU scratch buffer + // (3) matmul with beta=1 accumulates -2 * Q * C^T on top + raft::linalg::norm( + searcher_handle, + queries, + raft::make_device_vector_view(searcher.get_q_norms(), batch_size)); + cuvs::spatial::knn::detail::utils::outer_add(searcher.get_q_norms(), + batch_size, + centroid_norms_.data_handle(), + num_centroids, + searcher.get_centroid_distances(), + searcher_stream); + const float alpha = -2.f, beta = 1.f; + raft::linalg::detail::matmul(searcher_handle, + /* trans_a = */ true, + /* trans_b = */ false, + num_centroids, + batch_size, + num_padded_dim, + &alpha, + initializer->GetCentroid(0), + num_padded_dim, + queries.data_handle(), + num_padded_dim, + &beta, + searcher.get_centroid_distances(), + num_centroids); + + // Step 4: select top-nprobe clusters per query + auto d_raft_vals = raft::make_device_matrix(searcher_handle, batch_size, nprobe); + auto d_raft_idx = raft::make_device_matrix(searcher_handle, batch_size, nprobe); + auto in_view = raft::make_device_matrix_view( + searcher.get_centroid_distances(), batch_size, num_centroids); + cuvs::selection::select_k(searcher_handle, + in_view, + std::nullopt, + d_raft_vals.view(), + d_raft_idx.view(), + /*select_min=*/true, + /*sorted=*/true); + + // Step 5: sort (cluster, query) pairs + d_sorted_pairs = + raft::make_device_vector(searcher_handle, batch_size * nprobe); + sort_cluster_query_pairs( + searcher_handle, d_raft_idx.data_handle(), d_sorted_pairs.data_handle(), batch_size, nprobe); + + // Step 6: compute query factors + d_G_k1xSumq = raft::make_device_vector(searcher_handle, batch_size); + d_G_kbxSumq = raft::make_device_vector(searcher_handle, batch_size); + computeQueryFactors( + queries, d_G_k1xSumq.view(), d_G_kbxSumq.view(), ex_bits, searcher_stream); + + // Sync here to ensure all outputs are visible before SearcherGPU reads them. + raft::resource::sync_stream(searcher_handle); +} + +void IVFGPU::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) +{ + auto d_sorted_pairs = + raft::make_device_vector(searcher.get_handle(), 0); + auto d_G_k1xSumq = raft::make_device_vector(searcher.get_handle(), 0); + auto d_G_kbxSumq = raft::make_device_vector(searcher.get_handle(), 0); + PrepareClusterSearchInputs(queries, nprobe, searcher, d_sorted_pairs, d_G_k1xSumq, d_G_kbxSumq); + searcher.SearchClusterQueryPairs(*this, + cluster_meta_.data_handle(), + d_sorted_pairs.data_handle(), + batch_size, + queries, + d_G_k1xSumq.data_handle(), + d_G_kbxSumq.data_handle(), + nprobe, + k, + d_final_dists, + d_final_pids); +} + +void IVFGPU::BatchClusterSearchLUT16( + 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) +{ + auto d_sorted_pairs = + raft::make_device_vector(searcher.get_handle(), 0); + auto d_G_k1xSumq = raft::make_device_vector(searcher.get_handle(), 0); + auto d_G_kbxSumq = raft::make_device_vector(searcher.get_handle(), 0); + PrepareClusterSearchInputs(queries, nprobe, searcher, d_sorted_pairs, d_G_k1xSumq, d_G_kbxSumq); + searcher.SearchClusterQueryPairsSharedMemOpt(*this, + cluster_meta_.data_handle(), + d_sorted_pairs.data_handle(), + batch_size, + queries, + d_G_k1xSumq.data_handle(), + d_G_kbxSumq.data_handle(), + nprobe, + k, + d_final_dists, + d_final_pids); +} + +void IVFGPU::BatchClusterSearchQuantizeQuery( + 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, + int query_bits) +{ + auto d_sorted_pairs = + raft::make_device_vector(searcher.get_handle(), 0); + auto d_G_k1xSumq = raft::make_device_vector(searcher.get_handle(), 0); + auto d_G_kbxSumq = raft::make_device_vector(searcher.get_handle(), 0); + PrepareClusterSearchInputs(queries, nprobe, searcher, d_sorted_pairs, d_G_k1xSumq, d_G_kbxSumq); + searcher.SearchClusterQueryPairsQuantizeQuery(*this, + cluster_meta_.data_handle(), + d_sorted_pairs.data_handle(), + batch_size, + queries, + d_G_k1xSumq.data_handle(), + d_G_kbxSumq.data_handle(), + nprobe, + k, + d_final_dists, + d_final_pids, + query_bits == 4); +} + +} // 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 new file mode 100644 index 0000000000..75ddaee865 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh @@ -0,0 +1,394 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +// +// Created by Stardust on 2/23/25. +// + +#pragma once + +#include "../defines.hpp" +#include "initializer_gpu.cuh" +#include "quantizer_gpu.cuh" +#include "rotator_gpu.cuh" + +#include +#include +#include +#include +#include + +#include + +#include + +#include +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +class SearcherGPU; + +// Structure for cluster-query pairs +struct ClusterQueryPair { + int cluster_idx; + int query_idx; +}; + +class IVFGPU { + public: + class GPUClusterMeta { + public: + size_t num; // Number of vectors in this cluster. + size_t start_index; // Combined offset: index of first vector in the flattened arrays. + + __host__ __device__ GPUClusterMeta(size_t num, size_t start_idx) + : num(num), start_index(start_idx) + { + } + + // default constructor (useful for vector initialize) + GPUClusterMeta() : num(0), start_index(0) {} + + // Copy constructor. + GPUClusterMeta(const GPUClusterMeta& other) = default; + + // Copy assignment. + GPUClusterMeta& operator=(const GPUClusterMeta& other) = default; + + // Destructor. + ~GPUClusterMeta() = default; + + /** + * @brief Compute pointer to the first block of quantized short data. + * Factors are stored separately in SoA layout. + * + * @param parent Pointer to the IVFGPU instance that holds the base pointer. + * @return Pointer to the first block of this cluster's short data. + */ + __host__ uint32_t* first_block(const IVFGPU& parent) const + { + return parent.get_short_data_device() + + start_index * (parent.quantizer().short_code_length()); + } + + __device__ uint32_t* first_block_gpu(uint32_t* d_short_data, size_t short_code_length) const + { + return d_short_data + start_index * short_code_length; + } + + __host__ uint32_t* first_block_host(const IVFGPU& parent) const + { + return parent.get_short_data_host() + start_index * (parent.quantizer().short_code_length()); + } + + float* short_factor_batch(const IVFGPU& parent, size_t i) const + { + return parent.get_short_factors_batch_device() + (start_index + i) * 3; + } + + /** + * @brief Get the pointer to the long code for the i-th vector in this cluster. + * + * @param parent Pointer to the IVFGPU instance that holds the base pointer. + * @param i Index within the cluster (0 ≤ i < Num). + * @param long_code_length Fixed length (in bytes) for each vector’s long code. + * @return Pointer to the long code for the i-th vector. + */ + __host__ __device__ uint8_t* long_code(const IVFGPU& parent, + size_t i, + size_t long_code_length) const + { + return parent.get_long_code_device() + (start_index + i) * long_code_length; + } + + __host__ __device__ uint8_t* long_code_host(const IVFGPU& parent, + size_t i, + size_t long_code_length) const + { + return parent.get_long_code_host() + (start_index + i) * long_code_length; + } + + /** + * @brief Get the pointer to the extra factor for the i-th vector in this cluster. + * Two ex_factors are stored (f_add_ex and f_rescale_ex) in SoA layout. + * + * @param parent Pointer to the IVFGPU instance that holds the base pointer. + * @param i Index within the cluster. + * @return Pointer to the extra factor for the i-th vector. + */ + ExFactor* ex_factor(const IVFGPU& parent, size_t i) const + { + return parent.get_ex_factor_device() + (start_index + i) * 2; + } + + ExFactor* ex_factor_host(const IVFGPU& parent, size_t i) const + { + return parent.get_ex_factor_host() + (start_index + i) * 2; + } + + /** + * @brief Get the pointer to the vector IDs for this cluster. + * + * @param parent Pointer to the IVFGPU instance that holds the base pointer. + * @return Pointer to the first ID in this cluster. + */ + __host__ __device__ PID* ids(const IVFGPU& parent) const + { + return parent.get_ids_device() + start_index; + } + + __host__ __device__ PID* ids_host(const IVFGPU& parent) const + { + return parent.get_ids_host() + start_index; + } + }; + + /** + * @brief initialize function (no memory allocated yet) + * + * @param n Num of data points. + * @param dim Dimension of data points. + * @param k Num of centroids. + * @param bits_per_dim totalbits = EX_BITS+1 + */ + IVFGPU(raft::resources const& handle, size_t n, size_t dim, size_t k, size_t bits_per_dim); + IVFGPU(raft::resources const& handle) + : handle_(handle), initializer(nullptr), Rota(std::make_unique(handle_, 128)) + { + } + + /** + * @brief Build function + * + * @param device_data pointer to device data. + * @param device_centroids pointer to centroids. + * @param device_cluster_ids PIDs of vectors. + */ + void construct_on_gpu(const float* device_data, + const float* device_centroids, + const PID* device_cluster_ids, + bool fast_quantize); + + /** + * @brief Build function with streaming from host memory + * + * @param host_data pointer to host-resident data. + * @param device_centroids pointer to centroids on device. + * @param device_cluster_ids cluster assignments for each vector (on device). + * @param fast_quantize whether to use fast quantization. + * @param batch_size_vectors maximum vectors per batch (upper bound). + */ + void construct_on_gpu_streaming(const float* host_data, + const float* device_centroids, + const PID* device_cluster_ids, + bool fast_quantize, + size_t batch_size_vectors = 100000); + + void save(const char* filename) const; + + // load_transposed only applies for new batch index + void load_transposed(const char* filename); + + // device data getters + __host__ __device__ uint32_t* get_short_data_device() const noexcept + { + return const_cast(this->short_data_.data_handle()); + } + __host__ __device__ uint8_t* get_long_code_device() const noexcept + { + return const_cast(this->long_code_.data_handle()); + } + __host__ __device__ ExFactor* get_ex_factor_device() const noexcept + { + return const_cast(this->ex_factor_.data_handle()); + } + __host__ __device__ PID* get_ids_device() const noexcept + { + return const_cast(this->ids_.data_handle()); + } + __host__ __device__ float* get_short_factors_batch_device() const noexcept + { + return const_cast(this->short_factors_batch_.data_handle()); + } + raft::device_vector const& get_cluster_meta() const + { + return cluster_meta_; + } + + // host data getters + raft::host_vector const& get_cluster_meta_host() const + { + return cluster_meta_host_; + } + __host__ __device__ uint32_t* get_short_data_host() const noexcept + { + return const_cast(short_data_host_.data_handle()); + } + __host__ __device__ uint8_t* get_long_code_host() const noexcept + { + return const_cast(long_code_host_.data_handle()); + } + __host__ __device__ ExFactor* get_ex_factor_host() const noexcept + { + return const_cast(ex_factor_host_.data_handle()); + } + __host__ __device__ PID* get_ids_host() const noexcept + { + return const_cast(ids_host_.data_handle()); + } + + // metadata getters + size_t get_num_vectors() const noexcept { return this->num_vectors; } + size_t get_num_dimensions() const noexcept { return this->num_dimensions; } + size_t get_num_padded_dim() const noexcept { return this->num_padded_dim; } + size_t get_num_centroids() const { return num_centroids; } + size_t get_max_cluster_length() const noexcept { return max_cluster_length; } + size_t get_ex_bits() const noexcept { return ex_bits; } + + // member object getters + 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 BatchClusterSearchLUT16( + 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 BatchClusterSearchQuantizeQuery( + 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, + int query_bits); + + private: + void PrepareClusterSearchInputs( + raft::device_matrix_view queries, + size_t nprobe, + SearcherGPU& searcher, + raft::device_vector& d_sorted_pairs, + raft::device_vector& d_G_k1xSumq, + raft::device_vector& d_G_kbxSumq); + + /** + * @brief function to allocate memory based on the cluster + * + */ + void AllocateDeviceMemory(); + + // Following are inline functions to compute spaces for memory allocation + + size_t GetShortDataBytesSimple() const + { + // assert(num_clusters == num_centroids); // num of clusters + // size_t total_blocks = 0; + // for (auto s = 0; s < num_clusters; s++) { + // total_blocks += cluster_sizes[s]; + // } + return num_vectors * this->quantizer().block_bytes(); + } + + size_t GetShortDataFactorBytesBatch() const + { + // assert(num_clusters == num_centroids); // num of clusters + // size_t total_blocks = 0; + // for (auto s = 0; s < num_clusters; s++) { + // total_blocks += cluster_sizes[s]; + // } + return num_vectors * 3 * sizeof(float); + } + + size_t GetExFactorBytes() const + { + return 2 * sizeof(float) * num_vectors; // only f_add_ex and f_rescale_ex + } + + size_t GetPIDsBytes() const { return sizeof(PID) * num_vectors; } + + size_t GetLongCodeBytes() const + { + return sizeof(uint8_t) * quantizer().long_code_length() * num_vectors; + } + void init_clusters(const std::vector& cluster_sizes); + + void quantize_cluster(GPUClusterMeta& cp, + const float* data, + const float* cur_centroid, + float* rotated_c) const; + + void compute_centroid_norms(); + + void AllocateHostMemory(); + + raft::resources const& handle_; // reusable resource handle + rmm::cuda_stream_view stream_ = + raft::resource::get_cuda_stream(handle_); // CUDA stream obtained from handle_ + + // Device pointers for each data array. + raft::device_vector short_data_ = + raft::make_device_vector(handle_, 0); // RaBitQ code and factors. + raft::device_vector long_code_ = + raft::make_device_vector(handle_, 0); // ExRaBitQ code. + raft::device_vector ex_factor_ = + raft::make_device_vector(handle_, 0); // ExRaBitQ factor. + raft::device_vector ids_ = + raft::make_device_vector(handle_, 0); // PID of vectors. + raft::device_vector cluster_meta_ = + raft::make_device_vector(handle_, + 0); // Device-side array of clusters. + raft::device_vector centroid_norms_ = + raft::make_device_vector(handle_, 0); + + // batch-data (SoA layout - factors stored separately) + raft::device_vector short_factors_batch_ = + raft::make_device_vector(handle_, 0); // N * 3 float rabitq factors + + // host-side copies + raft::host_vector short_data_host_ = raft::make_host_vector( + 0); // TODO: CPU side, we need on factors from short_data_host_, so no need to + // store all these codes + raft::host_vector long_code_host_ = + raft::make_host_vector(0); // ExRaBitQ code. + raft::host_vector ex_factor_host_ = + raft::make_host_vector(0); // ExRaBitQ factor. + raft::host_vector ids_host_ = + raft::make_host_vector(0); // PID of vectors. + raft::host_vector cluster_meta_host_ = + raft::make_host_vector(0); // Host-side copy of clusters + + // Index meta-data. + size_t num_vectors; // Total number of vectors. + size_t num_dimensions; // Dimensionality of the input vectors. + size_t num_padded_dim; // padded dimensions + size_t num_centroids; // Centroids && Clusters + size_t max_cluster_length; // Maximum length of clusters + size_t ex_bits; // Extra bits parameter for quantization. + + std::unique_ptr + initializer; // Initializer, indicates which initializer to use (currently only FlatIVF) + std::unique_ptr DQ; // Data quantizer. + std::unique_ptr Rota; // Data Rotator. +}; + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu new file mode 100644 index 0000000000..af0876acca --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu @@ -0,0 +1,1391 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +// +// Created by Stardust on 3/10/25. +// + +#include "../../detail/smem_utils.cuh" +#include "../utils/reductions.cuh" +#include "quantizer_gpu.cuh" + +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +namespace { + +// Forward declarations for free functions +void data_transformation_batch_opt(const float* d_data, + const float* d_centroid, + const PID* d_IDs, + size_t num_points, + const RotatorGPU& rotator, + float* d_rotated_c, + float* d_XP_norm, + int* d_bin_XP, + float* d_XP_output, + float* d_X_and_C_pad, + size_t DIM, + size_t D, + raft::resources const& handle); + +void data_transformation_batch_opt_contiguous(const float* d_contiguous_data, + const float* d_centroid, + size_t num_points, + const RotatorGPU& rotator, + float* d_rotated_c, + float* d_XP_norm, + int* d_bin_XP, + float* d_XP_output, + float* d_X_and_C_pad, + size_t DIM, + size_t D, + raft::resources const& handle); + +void rabitq_codes_and_factors_fused(const float* d_rotated_c, + const int* d_bin_XP, + const float* d_XP, + uint32_t* d_short_data, + float* d_short_data_factors, + size_t num_points, + size_t D, + raft::resources const& handle); + +void exrabitq_codes_and_factors_fused(const int* d_bin_XP, + const float* d_XP_norm, + float* d_XP, + uint8_t* d_long_code, + float* d_ex_factor, + const float* d_centroid, + size_t num_points, + size_t D, + size_t EX_BITS, + float const_scaling_factor, + raft::resources const& handle); + +void exrabitq_codes_and_factors_fused_ori(const int* d_bin_XP, + const float* d_XP_norm, + float* d_XP, + uint8_t* d_long_code, + float* d_ex_factor, + const float* d_centroid, + size_t num_points, + size_t D, + size_t EX_BITS, + raft::resources const& handle); + +//--------------------------------------------------------------------------- +// Kernel: subtract_normalize_binarize_Kernel (Fused) +// This single kernel performs three sequential operations: +// 1. Subtracts the rotated centroid (d_CP) from each row of rotated data (d_XP). +// 2. Normalizes each resulting row (residual). +// 3. Binarizes the normalized row. +// It uses a block-per-row strategy with shared memory for the reduction. +// +// Template parameters: +// - BlockSize: The number of threads per block, should be a power of 2. +template +__global__ void subtract_normalize_binarize_Kernel( + const float* __restrict__ d_XP, + const float* __restrict__ d_CP, + float* d_XP_residuals, // Output for final residuals + float* d_XP_norm, // Output for normalized data + int* d_bin_XP, // Output for binarized data + size_t num_points, + size_t D) +{ + // Use one block to process one row. + int row = blockIdx.x; + if (row >= num_points) { return; } + + // Shared memory for the reduction and for broadcasting the final norm. + extern __shared__ float s_mem[]; + float* s_partials = s_mem; // Used for the reduction sum + + int tid_in_block = threadIdx.x; + float thread_sum_sq = 0.0f; + + // Each thread processes multiple elements if D > BlockSize + for (int j = tid_in_block; j < D; j += BlockSize) { + // Step 1: Calculate the residual (XP - CP) + float residual = d_XP[row * D + j] - d_CP[j]; + + // Step 2: Store the residual for later use and accumulate sum of squares + thread_sum_sq += residual * residual; + + // Step 3 (part 1): Write the residual to its final output location + d_XP_residuals[row * D + j] = residual; + } + + s_partials[tid_in_block] = thread_sum_sq; + __syncthreads(); + + // Step 4: Parallel reduction within the block to find the total sum of squares + for (unsigned int s = BlockSize / 2; s > 0; s >>= 1) { + if (tid_in_block < s) { s_partials[tid_in_block] += s_partials[tid_in_block + s]; } + __syncthreads(); + } + + // The final L2 norm is calculated by thread 0 and broadcast via shared memory + if (tid_in_block == 0) { s_partials[0] = sqrtf(s_partials[0]); } + __syncthreads(); + + float norm = s_partials[0]; + + // Step 5: Normalize and Binarize + // Each thread calculates and writes the final normalized and binary values + if (norm > 0 /*1e-6f*/) { // Avoid division by zero + for (int j = tid_in_block; j < D; j += BlockSize) { + float residual = d_XP_residuals[row * D + j]; // Read the just-written residual + float normalized_val = residual / norm; + d_XP_norm[row * D + j] = normalized_val; + d_bin_XP[row * D + j] = (normalized_val > 0.0f) ? 1 : 0; + } + } else { // Handle zero-norm case + for (int j = tid_in_block; j < D; j += BlockSize) { + d_XP_norm[row * D + j] = 0.0f; + d_bin_XP[row * D + j] = 0; + } + } +} + +//--------------------------------------------------------------------------- +// Kernel: gatherAndPadKernel (Combined) +// Populates a (num_points + 1) x D matrix, d_X_and_C_pad. +// - The first num_points rows are filled with data points padded to length D. +// If d_IDs is non-null, rows are gathered from d_data via d_IDs; +// if d_IDs is null, d_data is assumed contiguous and copied directly. +// - The last row (index num_points) is filled with the padded centroid. +__global__ void gatherAndPadKernel(const float* __restrict__ d_data, + const PID* __restrict__ d_IDs, + const float* __restrict__ d_centroid, + float* d_X_and_C_pad, + size_t num_points, + size_t DIM, + size_t D) +{ + // Calculate the global thread ID for a 1D grid over all elements. + int global_idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t total_elements = (num_points + 1) * D; + + if (global_idx < total_elements) { + // Determine which row and column this thread is responsible for. + int row_idx = global_idx / D; + int col_idx = global_idx % D; + + if (row_idx < num_points) { + // This thread is processing a data point. + if (col_idx < DIM) { + // Gather via d_IDs when provided; otherwise treat d_data as contiguous. + size_t src_row = (d_IDs != nullptr) ? static_cast(d_IDs[row_idx]) : row_idx; + d_X_and_C_pad[global_idx] = d_data[src_row * DIM + col_idx]; + } else { + // Pad with zero. + d_X_and_C_pad[global_idx] = 0.0f; + } + } else { + // This thread is processing the centroid (row_idx == num_points). + if (col_idx < DIM) { + // Copy from the centroid vector. + d_X_and_C_pad[global_idx] = d_centroid[col_idx]; + } else { + // Pad with zero. + d_X_and_C_pad[global_idx] = 0.0f; + } + } + } +} + +//--------------------------------------------------------------------------- +// Kernel: pack_and_compute_factors_kernel (Fused) +// +// This single kernel performs two independent operations for each data point (row): +// 1. Computes the RaBitQ factors (logic from RowwisePackedKernel). +// 2. Packs the binary representation into uint32_t codes (logic from pack_binary_kernel). +// +// It is launched with one block per data point (row). +__global__ void pack_and_compute_factors_kernel( + // Inputs for factor computation + const float* __restrict__ d_centroid, // [D] + const int* __restrict__ d_bin_XP, // [N*D], values in {0,1} + const float* __restrict__ d_XP, // [N*D], residuals + float kConstEpsilon, // e.g. 1.9f + // Inputs for packing (d_bin_XP is shared) + // Common parameters + size_t N, + size_t D, + // Outputs + float* __restrict__ d_factors, // [3*N], packed factors + uint32_t* __restrict__ d_packed_code // [N * (D/32)], packed codes +) +{ + // Each block is responsible for one data point (row). + size_t row = blockIdx.x; + if (row >= N) return; + + //========================================================================= + // Part 1: Compute RaBitQ factors (from RowwisePackedKernel) + //========================================================================= + float l2_sqr = 0.f; + float ip_resi_xucb = 0.f; + float ip_cent_xucb = 0.f; + float xu_sq = 0.f; + + // Each thread in the block processes a subset of the columns + for (size_t d = threadIdx.x; d < D; d += blockDim.x) { + float res = d_XP[row * D + d]; + float xu = float(d_bin_XP[row * D + d]) - 0.5f; + float c = d_centroid[d]; + + l2_sqr += res * res; + ip_resi_xucb += res * xu; + ip_cent_xucb += c * xu; + xu_sq += xu * xu; + } + + // Perform parallel reduction within the block + l2_sqr = blockReduceSum(l2_sqr); + ip_resi_xucb = blockReduceSum(ip_resi_xucb); + ip_cent_xucb = blockReduceSum(ip_cent_xucb); + xu_sq = blockReduceSum(xu_sq); + + // Thread 0 performs the final calculations and writes the factors + if (threadIdx.x == 0) { + if (ip_resi_xucb == 0.0f) ip_resi_xucb = INFINITY; + + float l2_norm = sqrtf(fmaxf(l2_sqr, 0.f)); + float denom = ip_resi_xucb; + + float fadd = l2_sqr + 2.f * l2_sqr * (ip_cent_xucb / denom); + float frescale = -2.f * l2_sqr / denom; + + float ratio = (l2_sqr * xu_sq) / (denom * denom); + float inner = (ratio - 1.f) / fmaxf(float(D - 1), 1.f); + inner = fmaxf(inner, 0.f); + float tmp_error = l2_norm * kConstEpsilon * sqrtf(inner); + float ferr = 2.f * tmp_error; + + size_t base = 3 * row; + d_factors[base + 0] = fadd; + d_factors[base + 1] = frescale; + d_factors[base + 2] = ferr; + } + + //========================================================================= + // Part 2: Pack binary codes (from pack_binary_kernel) + //========================================================================= + // Ensure all threads have finished with d_bin_XP for factor calculation + // before potentially reusing registers for packing. A sync is good practice here. + __syncthreads(); + + // Each thread processes one or more blocks of 32 bits for the current row. + size_t blocks_per_point = D / 32; + for (size_t block_id = threadIdx.x; block_id < blocks_per_point; block_id += blockDim.x) { + uint32_t cur = 0; + // Process 32 bits for this block_id + for (int i = 0; i < 32; i++) { + int bit = d_bin_XP[row * D + block_id * 32 + i]; + cur |= ((uint32_t)bit << (31 - i)); + } + d_packed_code[row * blocks_per_point + block_id] = cur; + } +} + +//--------------------------------------------------------------------------- +// Kernel: exrabitq_fused_kernel_batch (Definitive Optimized Version) +// +// Fuses ExRaBitQ code generation and factor computation into a single kernel. +// Uses a "one block per row" strategy with shared memory, coalesced access, +// efficient reductions, and a fully parallel, minimal-read packing algorithm. +// +template +__global__ void exrabitq_fused_kernel_batch( + // Inputs + const int* __restrict__ d_bin_XP, + const float* __restrict__ d_XP_norm, + const float* __restrict__ d_XP, + const float* __restrict__ d_centroid, + size_t num_points, + size_t D, + size_t EX_BITS, + float const_scaling_factor, + float kConstEpsilon, + // Outputs + uint8_t* d_long_code, + float* d_ex_factor) +{ + //========================================================================= + // Setup: One block per row + //========================================================================= + int row = blockIdx.x; + if (row >= num_points) return; + + // Dynamically allocated shared memory for one row's data. + extern __shared__ float s_mem[]; + float* s_xp_norm = s_mem; + uint8_t* s_tmp_code = (uint8_t*)(s_xp_norm + D); + + int tid = threadIdx.x; + + //========================================================================= + // Step 1: Coalesced load of all necessary data into shared memory + //========================================================================= + for (int j = tid; j < D; j += BlockSize) { + s_xp_norm[j] = d_XP_norm[row * D + j]; + } + + //========================================================================= + // Part A: ExRaBitQ Code Generation + //========================================================================= + const int mask = (1 << EX_BITS) - 1; + float thread_ipnorm_sum = 0.0f; + + // Parallel quantization and start of ip_norm reduction + for (int j = tid; j < D; j += BlockSize) { + float val = fabsf(s_xp_norm[j]); + int code_val = static_cast((const_scaling_factor * val) + 1e-5f); + if (code_val >= (1 << EX_BITS)) code_val = (1 << EX_BITS) - 1; + s_tmp_code[j] = static_cast(code_val); + thread_ipnorm_sum += (code_val + 0.5f) * val; + } + + // Parallel bit-flipping + for (int j = tid; j < D; j += BlockSize) { + if (d_bin_XP[row * D + j] == 0) { s_tmp_code[j] = (~s_tmp_code[j]) & mask; } + } + + // Finish ip_norm reduction + float total_ipnorm = blockReduceSum(thread_ipnorm_sum); + float ip_norm_inv = 1.0f; + if (tid == 0) { + float inv = 1.0f / total_ipnorm; + ip_norm_inv = isfinite(inv) ? inv : 1.0f; + } + // Broadcast ip_norm_inv to all threads in the block + ip_norm_inv = __shfl_sync(0xffffffff, ip_norm_inv, 0); + + //========================================================================= + // Part B: Factor Computation + //========================================================================= + float l2_sqr = 0.f, ip_resi_xucb = 0.f, ip_cent_xucb = 0.f, xu_sq = 0.f; + + for (size_t j = tid; j < D; j += BlockSize) { + float res = d_XP[row * D + j]; + int xu_pre = s_tmp_code[j]; + xu_pre += static_cast(res >= 0) << EX_BITS; + float xu = float(xu_pre) - (static_cast(1 << EX_BITS) - 0.5f); + float c = d_centroid[j]; + l2_sqr += res * res; + ip_resi_xucb += res * xu; + ip_cent_xucb += c * xu; + xu_sq += xu * xu; + } + + // Perform parallel reductions for all factor components + l2_sqr = blockReduceSum(l2_sqr); + ip_resi_xucb = blockReduceSum(ip_resi_xucb); + ip_cent_xucb = blockReduceSum(ip_cent_xucb); + xu_sq = blockReduceSum(xu_sq); + + // Thread 0 computes and writes the final factors + if (tid == 0) { + float denom = ip_resi_xucb; + if (denom == 0.0f) denom = INFINITY; + float l2_norm = sqrtf(fmaxf(l2_sqr, 0.f)); + + float fadd = l2_sqr + 2.f * l2_sqr * ip_cent_xucb / denom; + float frescale = -2.f * l2_norm * ip_norm_inv; + + float ratio = (l2_sqr * xu_sq) / (ip_resi_xucb * ip_resi_xucb); + float inner = (ratio - 1.f) / fmaxf(float(D - 1), 1.f); + float tmp_error = l2_norm * kConstEpsilon * sqrtf(fmaxf(inner, 0.f)); + + size_t base = 2 * row; + d_ex_factor[base + 0] = fadd; + d_ex_factor[base + 1] = frescale; + } + + //========================================================================= + // Part C: Pack and Write Long Code (MINIMAL READS, PARALLEL, COALESCED) + //========================================================================= + int long_code_length = (D * EX_BITS + 7) / 8; + uint8_t* out_ptr = d_long_code + row * long_code_length; + + if (EX_BITS == 8) { + for (int j = tid; j < D; j += BlockSize) + out_ptr[j] = s_tmp_code[j]; + } else { + const int num_codes_to_read = (8 + EX_BITS - 2) / EX_BITS + 1; + for (int out_byte_idx = tid; out_byte_idx < long_code_length; out_byte_idx += BlockSize) { + int start_bit = out_byte_idx * 8; + int start_code_idx = start_bit / EX_BITS; + int bit_offset = start_bit % EX_BITS; + uint64_t bit_buffer = 0; + + switch (num_codes_to_read) { + case 2: + if (start_code_idx < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx]; + bit_buffer <<= EX_BITS; + if (start_code_idx + 1 < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx + 1]; + break; + case 3: + if (start_code_idx < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx]; + bit_buffer <<= EX_BITS; + if (start_code_idx + 1 < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx + 1]; + bit_buffer <<= EX_BITS; + if (start_code_idx + 2 < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx + 2]; + break; + case 4: + if (start_code_idx < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx]; + bit_buffer <<= EX_BITS; + if (start_code_idx + 1 < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx + 1]; + bit_buffer <<= EX_BITS; + if (start_code_idx + 2 < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx + 2]; + bit_buffer <<= EX_BITS; + if (start_code_idx + 3 < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx + 3]; + break; + default: + for (int k = 0; k < num_codes_to_read; ++k) { + bit_buffer <<= EX_BITS; + if (start_code_idx + k < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx + k]; + } + break; + } + int total_bits_in_buffer = num_codes_to_read * EX_BITS; + int shift = total_bits_in_buffer - bit_offset - 8; + uint8_t out_byte = (shift >= 0) ? (uint8_t)(bit_buffer >> shift) : 0; + out_ptr[out_byte_idx] = out_byte; + } + } +} + +void data_transformation_batch_opt(const float* d_data, + const float* d_centroid, + const PID* d_IDs, + size_t num_points, + const RotatorGPU& rotator, + float* d_rotated_c, + float* d_XP_norm, + int* d_bin_XP, + float* d_XP_output, + float* d_X_and_C_pad, + size_t DIM, + size_t D, + raft::resources const& handle) +{ + auto stream = raft::resource::get_cuda_stream(handle); + + // 1. Allocate a single temporary buffer for both padded data and the padded centroid. + + // 2. Launch a single kernel to gather and pad both data and centroid. + int blockSize = D < 256 ? 128 : 256; + size_t totalPadElements = (num_points + 1) * D; + int gridPadSize = (totalPadElements + blockSize - 1) / blockSize; + gatherAndPadKernel<<>>( + d_data, d_IDs, d_centroid, d_X_and_C_pad, num_points, DIM, D); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // 3. Allocate a single output buffer for both rotated data (XP) and rotated centroid (CP). + float* d_XP_and_CP = d_XP_output; + + // 4. Perform a single, combined rotation. + // The input is d_X_and_C_pad, output is d_XP_and_CP. The number of "points" is num_points + 1. + rotator.rotate(handle, d_X_and_C_pad, d_XP_and_CP, num_points + 1); + + // Create pointers to the specific results within the combined buffer. + float* d_XP = d_XP_and_CP; + float* d_CP = d_XP_and_CP + num_points * D; + + // 5. Save the rotated centroid: copy CP into d_rotated_c. + raft::copy(d_rotated_c, d_CP, D, stream); + + // 6. Launch the single FUSED kernel for subtract, normalize, and binarize. + const unsigned int FusedBlockSize = 256; // A good default, can be tuned. + dim3 gridDim(num_points); + dim3 blockDim(FusedBlockSize); + size_t sharedMemSize = FusedBlockSize * sizeof(float); + + subtract_normalize_binarize_Kernel + <<>>(d_XP, // Input: Rotated data + d_CP, // Input: Rotated centroid + d_XP_output, // Output 1: Final residuals + d_XP_norm, // Output 2: Normalized residuals + d_bin_XP, // Output 3: Binarized data + num_points, + D); + RAFT_CUDA_TRY(cudaPeekAtLastError()); +} + +// Fused function to compute RaBitQ codes and factors in a single pass. +void rabitq_codes_and_factors_fused(const float* d_rotated_c, + const int* d_bin_XP, + const float* d_XP, + uint32_t* d_short_data, + float* d_short_data_factors, + size_t num_points, + size_t D, + raft::resources const& handle) +{ + auto stream = raft::resource::get_cuda_stream(handle); + + int threads_per_block = 256; // A good default, can be tuned + dim3 grid(num_points); + dim3 block(threads_per_block); + + pack_and_compute_factors_kernel<<>>( + d_rotated_c, + d_bin_XP, + d_XP, + 1.9f, // kConstEpsilon, hardcoded from original call + num_points, + D, + d_short_data_factors, + d_short_data); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(handle); +} + +//--------------------------------------------------------------------------- +// Host function: exrabitq_codes_batch (Optimized) +// Launches a single fused kernel to compute ExRaBitQ codes and factors. +void exrabitq_codes_and_factors_fused(const int* d_bin_XP, + const float* d_XP_norm, + float* d_XP, + uint8_t* d_long_code, + float* d_ex_factor, + const float* d_centroid, + size_t num_points, + size_t D, + size_t EX_BITS, + float const_scaling_factor, + raft::resources const& handle) +{ + auto stream = raft::resource::get_cuda_stream(handle); + + const unsigned int BlockSize = 256; // A good default, can be tuned + dim3 gridDim(num_points); + dim3 blockDim(BlockSize); + + // Calculate required shared memory size + size_t shared_mem_size = D * sizeof(float) + // s_xp_norm + D * sizeof(uint8_t) + // s_tmp_code + BlockSize * sizeof(float); // s_partials for reduction + + exrabitq_fused_kernel_batch + <<>>(d_bin_XP, + d_XP_norm, + d_XP, + d_centroid, + num_points, + D, + EX_BITS, + const_scaling_factor, + 1.9f, // kConstEpsilon + d_long_code, + d_ex_factor); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(handle); +} + +} // anonymous namespace + +void DataQuantizerGPU::quantize_batch_opt(const float* d_data, + const float* d_centroid, + const PID* d_IDs, + size_t num_points, + const RotatorGPU& rotator, + uint32_t* d_short_data, + float* d_short_data_factors, + uint8_t* d_long_code, + float* d_ex_factor, + float* d_rotated_c) +{ + // 1. Data Transformation: + data_transformation_batch_opt(d_data, + d_centroid, + d_IDs, + num_points, + rotator, + d_rotated_c, + d_XP_norm.data_handle(), + d_bin_XP.data_handle(), + d_XP.data_handle(), + d_X_and_C_pad.data_handle(), + DIM, + D, + handle_); + + rabitq_codes_and_factors_fused(d_rotated_c, + d_bin_XP.data_handle(), + d_XP.data_handle(), + d_short_data, + d_short_data_factors, + num_points, + D, + handle_); + + // 5. Compute ExRaBitQ quantization codes. + if (fast_quantize_flag) { + exrabitq_codes_and_factors_fused(d_bin_XP.data_handle(), + d_XP_norm.data_handle(), + d_XP.data_handle(), + d_long_code, + d_ex_factor, + d_rotated_c, + num_points, + D, + EX_BITS, + const_scaling_factor, + handle_); + } else { + exrabitq_codes_and_factors_fused_ori(d_bin_XP.data_handle(), + d_XP_norm.data_handle(), + d_XP.data_handle(), + d_long_code, + d_ex_factor, + d_rotated_c, + num_points, + D, + EX_BITS, + handle_); + } +} + +namespace { + +void data_transformation_batch_opt_contiguous(const float* d_contiguous_data, + const float* d_centroid, + size_t num_points, + const RotatorGPU& rotator, + float* d_rotated_c, + float* d_XP_norm, + int* d_bin_XP, + float* d_XP_output, + float* d_X_and_C_pad, + size_t DIM, + size_t D, + raft::resources const& handle) +{ + auto stream = raft::resource::get_cuda_stream(handle); + + // 1. Allocate a single temporary buffer for both padded data and the padded centroid. + + // 2. Launch a single kernel to copy, pad, and add centroid. + int blockSize = D < 256 ? 128 : 256; + size_t totalPadElements = (num_points + 1) * D; + int gridPadSize = (totalPadElements + blockSize - 1) / blockSize; + gatherAndPadKernel<<>>( + d_contiguous_data, nullptr, d_centroid, d_X_and_C_pad, num_points, DIM, D); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // 3. Allocate a single output buffer for both rotated data (XP) and rotated centroid (CP). + float* d_XP_and_CP = d_XP_output; + + // 4. Perform a single, combined rotation. + // The input is d_X_and_C_pad, output is d_XP_and_CP. The number of "points" is num_points + 1. + rotator.rotate(handle, d_X_and_C_pad, d_XP_and_CP, num_points + 1); + + // Create pointers to the specific results within the combined buffer. + float* d_XP = d_XP_and_CP; + float* d_CP = d_XP_and_CP + num_points * D; + + // 5. Save the rotated centroid: copy CP into d_rotated_c. + raft::copy(d_rotated_c, d_CP, D, stream); + + // 6. Launch the single FUSED kernel for subtract, normalize, and binarize. + const unsigned int FusedBlockSize = 256; // A good default, can be tuned. + dim3 gridDim(num_points); + dim3 blockDim(FusedBlockSize); + size_t sharedMemSize = FusedBlockSize * sizeof(float); + + subtract_normalize_binarize_Kernel + <<>>(d_XP, // Input: Rotated data + d_CP, // Input: Rotated centroid + d_XP_output, // Output 1: Final residuals + d_XP_norm, // Output 2: Normalized residuals + d_bin_XP, // Output 3: Binarized data + num_points, + D); + RAFT_CUDA_TRY(cudaPeekAtLastError()); +} + +} // anonymous namespace + +void DataQuantizerGPU::quantize_batch_opt_contiguous(const float* d_contiguous_data, + const float* d_centroid, + size_t num_points, + const RotatorGPU& rotator, + uint32_t* d_short_data, + float* d_short_data_factors, + uint8_t* d_long_code, + float* d_ex_factor, + float* d_rotated_c) +{ + // 1. Data Transformation: + data_transformation_batch_opt_contiguous(d_contiguous_data, + d_centroid, + num_points, + rotator, + d_rotated_c, + d_XP_norm.data_handle(), + d_bin_XP.data_handle(), + d_XP.data_handle(), + d_X_and_C_pad.data_handle(), + DIM, + D, + handle_); + + rabitq_codes_and_factors_fused(d_rotated_c, + d_bin_XP.data_handle(), + d_XP.data_handle(), + d_short_data, + d_short_data_factors, + num_points, + D, + handle_); + + // 5. Compute ExRaBitQ quantization codes. + if (fast_quantize_flag) { + exrabitq_codes_and_factors_fused(d_bin_XP.data_handle(), + d_XP_norm.data_handle(), + d_XP.data_handle(), + d_long_code, + d_ex_factor, + d_rotated_c, + num_points, + D, + EX_BITS, + const_scaling_factor, + handle_); + } else { + exrabitq_codes_and_factors_fused_ori(d_bin_XP.data_handle(), + d_XP_norm.data_handle(), + d_XP.data_handle(), + d_long_code, + d_ex_factor, + d_rotated_c, + num_points, + D, + EX_BITS, + handle_); + } +} + +namespace { + +constexpr std::array kTightStart = { + 0, + 0.15, + 0.20, + 0.52, + 0.59, + 0.71, + 0.75, + 0.77, + 0.81, +}; + +template +inline double best_rescale_factor(const T* o_abs, size_t dim, size_t ex_bits) +{ + constexpr double kEps = 1e-5; + constexpr int kNEnum = 10; + double max_o = *std::max_element(o_abs, o_abs + dim); + + double t_end = static_cast(((1 << ex_bits) - 1) + kNEnum) / max_o; + double t_start = t_end * kTightStart[ex_bits]; + + std::vector cur_o_bar(dim); + double sqr_denominator = static_cast(dim) * 0.25; + double numerator = 0; + + for (size_t i = 0; i < dim; ++i) { + int cur = static_cast((t_start * o_abs[i]) + kEps); + cur_o_bar[i] = cur; + sqr_denominator += cur * cur + cur; + numerator += (cur + 0.5) * o_abs[i]; + } + + std::priority_queue, + std::vector>, + std::greater<>> + next_t; + + for (size_t i = 0; i < dim; ++i) { + next_t.emplace(static_cast(cur_o_bar[i] + 1) / o_abs[i], i); + } + + double max_ip = 0; + double t = 0; + + while (!next_t.empty()) { + double cur_t = next_t.top().first; + size_t update_id = next_t.top().second; + next_t.pop(); + + cur_o_bar[update_id]++; + int update_o_bar = cur_o_bar[update_id]; + sqr_denominator += 2 * update_o_bar; + numerator += o_abs[update_id]; + + double cur_ip = numerator / std::sqrt(sqr_denominator); + if (cur_ip > max_ip) { + max_ip = cur_ip; + t = cur_t; + } + + if (update_o_bar < (1 << ex_bits) - 1) { + double t_next = static_cast(update_o_bar + 1) / o_abs[update_id]; + if (t_next < t_end) { next_t.emplace(t_next, update_id); } + } + } + + return t; +} + +} // anonymous namespace + +float DataQuantizerGPU::get_const_scaling_factors(raft::resources const& handle, + size_t dim, + size_t ex_bits) +{ + constexpr long kConstNum = 100; + + // random matrix with normal distribution + auto rand = raft::make_device_matrix(handle, kConstNum, dim); + raft::random::RngState rng(7ULL); + raft::random::normal(handle, rng, rand.data_handle(), kConstNum * dim, 0., 1.); + // row-wise normalization + auto row_normalized = raft::make_device_matrix(handle, kConstNum, dim); + raft::linalg::row_normalize( + handle, rand.view(), row_normalized.view()); + // take abs values (reusing memory allocation for `rand`) + raft::linalg::map(handle, + rand.view(), + raft::abs_op{}, + raft::make_device_vector_view( + row_normalized.data_handle(), kConstNum * dim)); + auto h_rand_row_normalized_abs = raft::make_host_matrix(kConstNum, dim); + raft::copy(h_rand_row_normalized_abs.data_handle(), + rand.data_handle(), + kConstNum * dim, + raft::resource::get_cuda_stream(handle)); + raft::resource::sync_stream(handle); + + double sum = 0; + for (long j = 0; j < kConstNum; ++j) { + sum += best_rescale_factor(&h_rand_row_normalized_abs(j, 0), dim, ex_bits); + } + + double t_const = sum / kConstNum; + + return (float)t_const; +} + +namespace { + +extern __constant__ float d_kTightStart_opt[9] = { + 0.0f, + 0.15f, + 0.20f, + 0.52f, + 0.59f, + 0.71f, + 0.75f, + 0.77f, + 0.81f, +}; + +__device__ float compute_best_rescale_parallel( + float* s_xp_norm, int D, int EX_BITS, float* reuse_space, int BlockSize) +{ + int tid = threadIdx.x; + constexpr float kEps = 1e-5f; + constexpr int kNEnum = 10; + + //========================================================================= + // Step 1: Find maximum value using parallel reduction + //========================================================================= + float local_max = 0.0f; + for (int i = tid; i < D; i += BlockSize) { + local_max = fmaxf(local_max, fabsf(s_xp_norm[i])); + } + + // Block-level reduction for max + float* s_reduce; + s_reduce = reuse_space; + s_reduce[tid] = local_max; + __syncthreads(); + + for (int stride = BlockSize / 2; stride > 0; stride >>= 1) { + if (tid < stride) { s_reduce[tid] = fmaxf(s_reduce[tid], s_reduce[tid + stride]); } + __syncthreads(); + } + + float max_o = s_reduce[0]; // full max, already visible to all threads via the reduction's sync + if (max_o < kEps) return 1.0f; + + //========================================================================= + // Step 2: Each thread evaluates multiple critical points + //========================================================================= + float t_end = static_cast((1 << EX_BITS) - 1 + kNEnum) / max_o; + float t_start = t_end * d_kTightStart_opt[EX_BITS]; + + //========================================================================= + // Phase 1: Coarse grid search + //========================================================================= + const int COARSE_SAMPLES = 64; // Though not fully utilize each thread, but very fast + float best_coarse_ip = 0.0f; + float best_coarse_t = t_start; + + for (int i = tid; i < COARSE_SAMPLES; i += BlockSize) { + float t = t_start + (t_end - t_start) * i / (COARSE_SAMPLES - 1); + + // Quick approximate evaluation (sample subset of dimensions) + float numerator = 0.0f; + float sqr_denominator = static_cast(D) * 0.25f; + + // Full evaluation for coarse grid + for (int j = 0; j < D; j++) { + float val = fabsf(s_xp_norm[j]); + int quantized = min(static_cast((t * val) + kEps), (1 << EX_BITS) - 1); + numerator += (quantized + 0.5f) * val; + sqr_denominator += quantized * quantized + quantized; + } + + float ip = numerator / sqrtf(sqr_denominator); + if (ip > best_coarse_ip) { + best_coarse_ip = ip; + best_coarse_t = t; + } + } + + // Parallel reduction to find best coarse point + float* s_coarse_ip; + s_coarse_ip = reuse_space + BlockSize; + float* s_coarse_t; + s_coarse_t = s_coarse_ip + BlockSize; + s_coarse_ip[tid] = best_coarse_ip; + s_coarse_t[tid] = best_coarse_t; + __syncthreads(); + + for (int stride = BlockSize / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + if (s_coarse_ip[tid + stride] > s_coarse_ip[tid]) { + s_coarse_ip[tid] = s_coarse_ip[tid + stride]; + s_coarse_t[tid] = s_coarse_t[tid + stride]; + } + } + __syncthreads(); + } + + //========================================================================= + // Phase 2: Fine search around best coarse point + //========================================================================= + float center_t = s_coarse_t[0]; + float range = (t_end - t_start) / COARSE_SAMPLES; + float fine_start = fmaxf(t_start, center_t - range); + float fine_end = fminf(t_end, center_t + range); + + const int FINE_SAMPLES = 32; + float best_fine_ip = 0.0f; + float best_fine_t = center_t; + + for (int i = tid; i < FINE_SAMPLES; i += BlockSize) { + float t = fine_start + (fine_end - fine_start) * i / (FINE_SAMPLES - 1); + + // Full evaluation + float numerator = 0.0f; + float sqr_denominator = static_cast(D) * 0.25f; + + for (int j = 0; j < D; j++) { + float val = fabsf(s_xp_norm[j]); + int quantized = min(static_cast((t * val) + kEps), (1 << EX_BITS) - 1); + numerator += (quantized + 0.5f) * val; + sqr_denominator += quantized * quantized + quantized; + } + + float ip = numerator / sqrtf(sqr_denominator); + if (ip > best_fine_ip) { + best_fine_ip = ip; + best_fine_t = t; + } + } + + // Final reduction + float* s_fine_ip; + s_fine_ip = s_coarse_ip; + float* s_fine_t; + s_fine_t = s_coarse_t; + s_fine_ip[tid] = best_fine_ip; + s_fine_t[tid] = best_fine_t; + __syncthreads(); + + for (int stride = BlockSize / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + if (s_fine_ip[tid + stride] > s_fine_ip[tid]) { + s_fine_ip[tid] = s_fine_ip[tid + stride]; + s_fine_t[tid] = s_fine_t[tid + stride]; + } + } + __syncthreads(); + } + + return s_fine_t[0]; +} + +template +__global__ void exrabitq_fused_kernel_batch_ori( + // Inputs + const int* __restrict__ d_bin_XP, + const float* __restrict__ d_XP_norm, + const float* __restrict__ d_XP, + const float* __restrict__ d_centroid, + size_t num_points, + size_t D, + size_t EX_BITS, + float kConstEpsilon, + // Outputs + uint8_t* d_long_code, + float* d_ex_factor) +{ + //========================================================================= + // Setup: One block per row + //========================================================================= + int row = blockIdx.x; + if (row >= num_points) return; + + // Dynamically allocated shared memory for one row's data. + extern __shared__ float s_mem[]; + float* s_xp_norm = s_mem; + uint8_t* s_tmp_code = (uint8_t*)(s_xp_norm + D); + + int tid = threadIdx.x; + + //========================================================================= + // Step 1: Coalesced load of all necessary data into shared memory + //========================================================================= + // Load data + for (int j = tid; j < D; j += BlockSize) { + s_xp_norm[j] = d_XP_norm[row * D + j]; + // Note: s_bin_xp, s_xp not loaded yet - we'll reuse as workspace + } + // Compute scaling factor (reusing unused buffers) + float const_scaling_factor = + compute_best_rescale_parallel(s_xp_norm, + D, + EX_BITS, + (s_xp_norm + D), // Reused as workspace + BlockSize); + + //========================================================================= + // Part A: ExRaBitQ Code Generation + //========================================================================= + const int mask = (1 << EX_BITS) - 1; + float thread_ipnorm_sum = 0.0f; + + // Parallel quantization and start of ip_norm reduction + for (int j = tid; j < D; j += BlockSize) { + float val = fabsf(s_xp_norm[j]); + int code_val = static_cast((const_scaling_factor * val) + 1e-5f); + if (code_val >= (1 << EX_BITS)) code_val = (1 << EX_BITS) - 1; + s_tmp_code[j] = static_cast(code_val); + thread_ipnorm_sum += (code_val + 0.5f) * val; + } + + // Parallel bit-flipping + for (int j = tid; j < D; j += BlockSize) { + if (d_bin_XP[row * D + j] == 0) { s_tmp_code[j] = (~s_tmp_code[j]) & mask; } + } + + // Finish ip_norm reduction + float total_ipnorm = blockReduceSum(thread_ipnorm_sum); + float ip_norm_inv = 1.0f; + if (tid == 0) { + float inv = 1.0f / total_ipnorm; + ip_norm_inv = isfinite(inv) ? inv : 1.0f; + } + // Broadcast ip_norm_inv to all threads in the block + ip_norm_inv = __shfl_sync(0xffffffff, ip_norm_inv, 0); + + //========================================================================= + // Part B: Factor Computation + //========================================================================= + float l2_sqr = 0.f, ip_resi_xucb = 0.f, ip_cent_xucb = 0.f, xu_sq = 0.f; + + for (size_t j = tid; j < D; j += BlockSize) { + float res = d_XP[row * D + j]; + int xu_pre = s_tmp_code[j]; + xu_pre += static_cast(res >= 0) << EX_BITS; + float xu = float(xu_pre) - (static_cast(1 << EX_BITS) - 0.5f); + float c = d_centroid[j]; + l2_sqr += res * res; + ip_resi_xucb += res * xu; + ip_cent_xucb += c * xu; + xu_sq += xu * xu; + } + + // Perform parallel reductions for all factor components + l2_sqr = blockReduceSum(l2_sqr); + ip_resi_xucb = blockReduceSum(ip_resi_xucb); + ip_cent_xucb = blockReduceSum(ip_cent_xucb); + xu_sq = blockReduceSum(xu_sq); + + // Thread 0 computes and writes the final factors + if (tid == 0) { + float denom = ip_resi_xucb; + if (denom == 0.0f) denom = INFINITY; + float l2_norm = sqrtf(fmaxf(l2_sqr, 0.f)); + + float fadd = l2_sqr + 2.f * l2_sqr * ip_cent_xucb / denom; + float frescale = -2.f * l2_norm * ip_norm_inv; + + float ratio = (l2_sqr * xu_sq) / (ip_resi_xucb * ip_resi_xucb); + float inner = (ratio - 1.f) / fmaxf(float(D - 1), 1.f); + float tmp_error = l2_norm * kConstEpsilon * sqrtf(fmaxf(inner, 0.f)); + + size_t base = 2 * row; + d_ex_factor[base + 0] = fadd; + d_ex_factor[base + 1] = frescale; + } + + //========================================================================= + // Part C: Pack and Write Long Code (MINIMAL READS, PARALLEL, COALESCED) + //========================================================================= + int long_code_length = (D * EX_BITS + 7) / 8; + uint8_t* out_ptr = d_long_code + row * long_code_length; + + if (EX_BITS == 8) { + for (int j = tid; j < D; j += BlockSize) + out_ptr[j] = s_tmp_code[j]; + } else { + const int num_codes_to_read = (8 + EX_BITS - 2) / EX_BITS + 1; + for (int out_byte_idx = tid; out_byte_idx < long_code_length; out_byte_idx += BlockSize) { + int start_bit = out_byte_idx * 8; + int start_code_idx = start_bit / EX_BITS; + int bit_offset = start_bit % EX_BITS; + uint64_t bit_buffer = 0; + + switch (num_codes_to_read) { + case 2: + if (start_code_idx < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx]; + bit_buffer <<= EX_BITS; + if (start_code_idx + 1 < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx + 1]; + break; + case 3: + if (start_code_idx < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx]; + bit_buffer <<= EX_BITS; + if (start_code_idx + 1 < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx + 1]; + bit_buffer <<= EX_BITS; + if (start_code_idx + 2 < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx + 2]; + break; + case 4: + if (start_code_idx < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx]; + bit_buffer <<= EX_BITS; + if (start_code_idx + 1 < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx + 1]; + bit_buffer <<= EX_BITS; + if (start_code_idx + 2 < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx + 2]; + bit_buffer <<= EX_BITS; + if (start_code_idx + 3 < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx + 3]; + break; + default: + for (int k = 0; k < num_codes_to_read; ++k) { + bit_buffer <<= EX_BITS; + if (start_code_idx + k < D) bit_buffer |= (uint64_t)s_tmp_code[start_code_idx + k]; + } + break; + } + int total_bits_in_buffer = num_codes_to_read * EX_BITS; + int shift = total_bits_in_buffer - bit_offset - 8; + uint8_t out_byte = (shift >= 0) ? (uint8_t)(bit_buffer >> shift) : 0; + out_ptr[out_byte_idx] = out_byte; + } + } +} + +//--------------------------------------------------------------------------- +// Host function: exrabitq_codes_batch (Optimized) +// No fast quantization version +void exrabitq_codes_and_factors_fused_ori(const int* d_bin_XP, + const float* d_XP_norm, + float* d_XP, + uint8_t* d_long_code, + float* d_ex_factor, + const float* d_centroid, + size_t num_points, + size_t D, + size_t EX_BITS, + raft::resources const& handle) +{ + auto stream = raft::resource::get_cuda_stream(handle); + + const unsigned int BlockSize = 256; // A good default, can be tuned + dim3 gridDim(num_points); + dim3 blockDim(BlockSize); + + // Calculate required shared memory size + size_t shared_mem_size = D * sizeof(float) + // s_xp_norm + D * sizeof(uint8_t) + // s_tmp_code + BlockSize * sizeof(float); // s_partials for reduction + + if (shared_mem_size <= D * sizeof(float) + 3 * BlockSize * sizeof(float)) { + shared_mem_size = + D * sizeof(float) + + 3 * BlockSize * + sizeof(float); // Make sure there are space to reuse for finding scaling factors + } + + exrabitq_fused_kernel_batch_ori + <<>>(d_bin_XP, + d_XP_norm, + d_XP, + d_centroid, + num_points, + D, + EX_BITS, + 1.9f, // kConstEpsilon + d_long_code, + d_ex_factor); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(handle); +} + +__global__ void fully_fused_kernel(float* __restrict__ output_factors, + const int rows, + const int cols, + const int ex_bits, + unsigned long long seed) +{ + const int row_id = blockIdx.x; + if (row_id >= rows) return; + + const int tid = threadIdx.x; + const int block_size = blockDim.x; + + // Calculate shared memory layout + // row_data: cols floats + // reuse_space: 3 * block_size floats (for best_rescale computation) + extern __shared__ float shared_mem[]; + float* row_data = shared_mem; + float* reuse_space = &row_data[cols]; // No reduction_buffer needed! + + // Initialize RNG state per thread + curandState rng_state; + curand_init(seed, row_id * block_size + tid, 0, &rng_state); + + // Generate random Gaussian values and calculate L2 norm + float local_sum = 0.0f; + for (int i = tid; i < cols; i += block_size) { + float val = curand_normal(&rng_state); + row_data[i] = val; + local_sum += val * val; + } + + float norm_squared = blockReduceSum(local_sum); + + // full block sum only available on lane 0 of warp 0 (see definition of blockReduceSum); need to + // make inv_norm a shared variable + __shared__ float inv_norm; + if (tid == 0) { inv_norm = rsqrtf(norm_squared); } + __syncthreads(); + + for (int i = tid; i < cols; i += block_size) { + row_data[i] = fabsf(row_data[i] * inv_norm); + } + + float rescale_factor = + compute_best_rescale_parallel(row_data, cols, ex_bits, reuse_space, block_size); + + if (tid == 0) { output_factors[row_id] = rescale_factor; } +} + +} // anonymous namespace + +float DataQuantizerGPU::get_const_scaling_factors_fully_gpu(size_t dim, size_t ex_bits) +{ + constexpr long kConstNum = 100; + + rmm::device_uvector d_factors(kConstNum, stream_); + rmm::device_uvector d_sum(1, stream_); + + // Calculate block size (must be power of 2 for reductions) + int block_size = 256; + if (dim <= 512) block_size = 128; + if (dim >= 1536) block_size = 512; + + // Calculate shared memory size + size_t shared_mem_size = (dim + // row_data + 3 * block_size // reuse_space for best_rescale + ) * + sizeof(float); + + unsigned long long seed = time(nullptr); + auto kernel_fn = fully_fused_kernel; // decay to function pointer for std::atomic + auto const& kernel_launcher = [&]() { + kernel_fn<<>>( + d_factors.data(), kConstNum, dim, ex_bits, seed); + }; + cudaKernel_t cuda_kernel; + RAFT_CUDA_TRY(cudaGetKernel(&cuda_kernel, reinterpret_cast(kernel_fn))); + cuvs::neighbors::detail::safely_launch_kernel_with_smem_size>( + static_cast(shared_mem_size), kernel_launcher, cuda_kernel); + RAFT_CUDA_TRY(cudaGetLastError()); + + // Use CUB for reduction - handles any size optimally + size_t temp_storage_bytes = 0; + cub::DeviceReduce::Sum( + nullptr, temp_storage_bytes, d_factors.data(), d_sum.data(), kConstNum, stream_); + + { + rmm::device_buffer d_temp_storage(temp_storage_bytes, stream_); + cub::DeviceReduce::Sum(d_temp_storage.data(), + temp_storage_bytes, + d_factors.data(), + d_sum.data(), + kConstNum, + stream_); + } + RAFT_CUDA_TRY(cudaGetLastError()); + + // Copy single value back and sync before reading host variable + float sum; + raft::copy(&sum, d_sum.data(), 1, stream_); + raft::resource::sync_stream(handle_); + + return sum / kConstNum; +} + +void DataQuantizerGPU::alloc_buffers(size_t num_points) +{ + const int64_t size_norm = static_cast(num_points) * D; + const int64_t size_bin = static_cast(num_points) * D; + const int64_t size_xp = static_cast(num_points + 1) * D; + + // Overwrite RAFT device vectors with new allocations + d_XP_norm = raft::make_device_vector(handle_, size_norm); + d_bin_XP = raft::make_device_vector(handle_, size_bin); + d_XP = raft::make_device_vector(handle_, size_xp); + d_X_and_C_pad = raft::make_device_vector(handle_, size_xp); +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh new file mode 100644 index 0000000000..90ef71ed12 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh @@ -0,0 +1,170 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +// +// Created by Stardust on 3/10/25. +// + +#pragma once + +#include "../defines.hpp" +#include "rotator_gpu.cuh" +#include + +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +typedef uint32_t PID; + +/*! + * @brief GPU version of DataQuantizer for ExRaBitQ. + * + * This class quantizes data into two sets of codes. + * Its interface closely follows the CPU version. Later, you would replace + * stub functions with CUDA kernels and device‐side implementations. + */ +class DataQuantizerGPU { + public: + // Nested struct for fast quantization factors + struct FastQuantizeFactors { + float const_scaling_factor_4bit; + float const_scaling_factor_8bit; + }; + + // Static methods + static float get_const_scaling_factors(raft::resources const& handle, size_t dim, size_t ex_bits); + + float get_const_scaling_factors_fully_gpu(size_t dim, size_t ex_bits); + // Constructor: initialize from dimension and bit count. + explicit DataQuantizerGPU(raft::resources const& handle, size_t dim, size_t b) + : DIM(dim), + D(raft::round_up_safe(dim, 64)), + EX_BITS(b), + SHORT_CODE_LENGTH((D + 31) / 32), + LONG_CODE_LENGTH(D * EX_BITS / 8) // Simplified for now. + , + FAC_NORM(1 / std::sqrt((double)D)), + FAC_ERR(2.0 / std::sqrt((double)(D - 1))), + fast_quantize_flag(false), + const_scaling_factor(0.0f), + handle_(handle) + { + } + + // Disable copy assignment + DataQuantizerGPU& operator=(const DataQuantizerGPU& other) = delete; + + // Accessor functions. + size_t short_code_length() const { return SHORT_CODE_LENGTH; } + size_t long_code_length() const { return LONG_CODE_LENGTH; } + // Block size for SoA layout (factors stored separately) + size_t block_bytes() const + { + // 3 factors for batch: f_add, f_rescale and f_error are stored separately + return SHORT_CODE_LENGTH * sizeof(uint32_t); + } + static constexpr size_t num_short_factors() { return NUM_SHORT_FACTORS; } + const FastQuantizeFactors* get_query_scaling_factor() const { return &fast_quantize_factors; } + FastQuantizeFactors* get_query_scaling_factor() { return &fast_quantize_factors; } + void compute_query_scaling_factors(size_t dim) + { + fast_quantize_factors.const_scaling_factor_4bit = get_const_scaling_factors(handle_, dim, 3); + fast_quantize_factors.const_scaling_factor_8bit = get_const_scaling_factors_fully_gpu(dim, 7); + } + void compute_quantize_scaling_factors() + { + const_scaling_factor = get_const_scaling_factors_fully_gpu(D, EX_BITS); + } + void set_quantize_scaling_factors(float value) { const_scaling_factor = value; } + + // Public configuration flag + bool fast_quantize_flag; + + // functions to malloc temp buffers for gpu + void alloc_buffers(size_t num_points); + + /*! + * @brief Quantize the input data for batch layout. + * + * @param data Pointer to the input data (host pointer or device pointer as needed). + * @param centroid Pointer to the corresponding centroid. + * @param pids Vector of vector IDs. + * @param rotator A RotatorGPU instance. + * @param outShort Output buffer for short codes. + * @param outLong Output buffer for long codes. + * @param outExFactor Output buffer for extra factors. + * @param outTemp Temporary buffer. + */ + void quantize_batch_opt(const float* d_data, + const float* d_centroid, + const PID* d_IDs, + size_t num_points, + const RotatorGPU& rotator, + uint32_t* d_short_data, + float* short_data_factors, + uint8_t* d_long_code, + float* d_ex_factor, + float* d_rotated_c); + + /*! + * @brief Quantize contiguous cluster data for batch layout (without PID gathering). + * + * @param d_contiguous_data Pointer to contiguous cluster data on device. + * @param d_centroid Pointer to the corresponding centroid on device. + * @param num_points Number of points in the cluster. + * @param rotator A RotatorGPU instance. + * @param d_short_data Output buffer for short codes. + * @param short_data_factors Output buffer for short code factors. + * @param d_long_code Output buffer for long codes. + * @param d_ex_factor Output buffer for extra factors. + * @param d_rotated_c Output buffer for rotated centroid. + */ + void quantize_batch_opt_contiguous(const float* d_contiguous_data, + const float* d_centroid, + size_t num_points, + const RotatorGPU& rotator, + uint32_t* d_short_data, + float* short_data_factors, + uint8_t* d_long_code, + float* d_ex_factor, + float* d_rotated_c); + + private: + // Dimension and quantization parameters + size_t DIM; // Original data dimension. + size_t D; // Padded dimension (multiple of 64). + size_t EX_BITS; // Number of bits for ExRaBitQ. + size_t SHORT_CODE_LENGTH; // Number of uint32_t to store 1-bit code for a vector. + size_t LONG_CODE_LENGTH; // Number of uint8_t to store EX_BITS code for a vector. + double FAC_NORM; + double FAC_ERR; + float const_scaling_factor; + FastQuantizeFactors fast_quantize_factors; + static constexpr size_t NUM_SHORT_FACTORS = 1; + + // RAFT resources + raft::resources const& handle_; // reusable resource handle + rmm::cuda_stream_view stream_ = + raft::resource::get_cuda_stream(handle_); // CUDA stream obtained from handle_ + + // Device temporary buffers for quantization + raft::device_vector d_XP_norm = + raft::make_device_vector(handle_, 0); + raft::device_vector d_bin_XP = raft::make_device_vector(handle_, 0); + raft::device_vector d_XP = raft::make_device_vector(handle_, 0); + raft::device_vector d_X_and_C_pad = + raft::make_device_vector(handle_, 0); +}; + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu new file mode 100644 index 0000000000..59a3e575d9 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu @@ -0,0 +1,81 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +// +// Created by Stardust on 3/24/25. +// + +#include "rotator_gpu.cuh" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +RotatorGPU::RotatorGPU(raft::resources const& handle, uint32_t dim) + : D(raft::round_up_safe(dim, 64u)), + rotation_matrix_(raft::make_device_matrix(handle, D, D)) +{ + raft::random::RngState rng(7ULL); + raft::random::normal(handle, rng, rotation_matrix_.data_handle(), D * D, 0.0f, 1.0f); + raft::linalg::detail::qrGetQ_inplace( + handle, rotation_matrix_.data_handle(), D, D, raft::resource::get_cuda_stream(handle)); +} + +size_t RotatorGPU::size() const { return D; } + +void RotatorGPU::load(raft::resources const& handle, std::ifstream& input) +{ + auto stream = raft::resource::get_cuda_stream(handle); + auto host_buf = raft::make_host_vector(D * D); + const auto matrix_size = static_cast(sizeof(float) * D * D); + input.read(reinterpret_cast(host_buf.data_handle()), matrix_size); + RAFT_EXPECTS(input.gcount() == matrix_size, "unexpected EOF reading rotator matrix"); + raft::copy(rotation_matrix_.data_handle(), host_buf.data_handle(), D * D, stream); + raft::resource::sync_stream(handle); +} + +void RotatorGPU::save(raft::resources const& handle, std::ofstream& output) const +{ + auto stream = raft::resource::get_cuda_stream(handle); + auto host_buf = raft::make_host_vector(D * D); + raft::copy(host_buf.data_handle(), rotation_matrix_.data_handle(), D * D, stream); + raft::resource::sync_stream(handle); + output.write(reinterpret_cast(host_buf.data_handle()), + static_cast(sizeof(float) * D * D)); + RAFT_EXPECTS(static_cast(output), "failed to write rotator matrix"); +} + +// Rotate the matrix A and store the result in RAND_A on the GPU. +// A and RAND_A are assumed to be stored in row-major order. +// A is of size N x D and P is of size D x D, so the result is N x D. +// This function uses cuBLAS to perform the matrix multiplication. +void RotatorGPU::rotate(raft::resources const& handle, + const float* d_A, + float* d_RAND_A, + size_t N) const +{ + // cuBLAS assumes column-major storage by default. Since our matrices are in row-major order, + // we can perform the multiplication as: + // RAND_A^T = P^T * A^T + // which is equivalent to RAND_A = A * P, if we interpret the data as row-major. + // Note that in Cublas it is RAND_A^T in column major, which is what we want in row-major + // Here, we use the RAFT wrapper for gemm. + raft::linalg::gemm( + handle, + raft::make_device_matrix_view( + const_cast(rotation_matrix_.data_handle()), D, D), + raft::make_device_matrix_view(const_cast(d_A), D, N), + raft::make_device_matrix_view(d_RAND_A, D, N)); +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh new file mode 100644 index 0000000000..8db4bb464c --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh @@ -0,0 +1,74 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +// +// Created by Stardust on 3/24/25. +// + +#pragma once + +#include "../defines.hpp" + +#include +#include +#include + +#include + +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// The RotatorGPU class holds a rotation matrix (P) on the GPU. The matrix is computed +// on the CPU (using Eigen, similar to your CPU code) and then copied to device memory. +// The rotate() function uses cuBLAS to compute the product: RAND_A = A * P. +// It is assumed that A and RAND_A reside in GPU memory. +class RotatorGPU { + public: /** + * @brief Construct a new RotatorGPU object. + * @param dim The original dimension; the padded dimension D is computed as + * rd_up_to_multiple_of(dim, 64). + * + * The constructor generates a random rotation matrix on the CPU (using Eigen) and then + * copies it into device memory in column-major order. + */ + explicit RotatorGPU(raft::resources const& handle, uint32_t dim); + + // Disable copy assignment + RotatorGPU& operator=(const RotatorGPU& other) = delete; + + size_t size() const; + + /** + * @brief Load the rotation matrix from a file. + * @param input Input stream (the file stores the matrix in row-major order). + * + * The function reads the D×D matrix from the file, transposes it into column-major order, + * and copies it into device memory. + */ + void load(raft::resources const& handle, std::ifstream& input); + + /** + * @brief Save the rotation matrix to a file. + * @param handle Resource handle + * @param output Output stream. + * + * The function copies the rotation matrix from device memory, transposes it from column-major to + * row-major, and writes it to the file. + */ + void save(raft::resources const& handle, std::ofstream& output) const; + + // Rotate matrix A and store the result in RAND_A. + // A and RAND_A are device pointers representing matrices of size N x D. + // This function computes: RAND_A = A * P using cuBLAS. + void rotate(raft::resources const& handle, const float* d_A, float* d_RAND_A, size_t N) const; + + private: + size_t D; // Padded dimension + raft::device_matrix rotation_matrix_; +}; + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu new file mode 100644 index 0000000000..22bb74c682 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu @@ -0,0 +1,290 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +// +// Created by Stardust on 4/14/25. +// + +// This file implements `SearcherGPU::SearchClusterQueryPairs`. +#include "../../detail/smem_utils.cuh" +#include "../jit_lto_kernels/kernel_def.hpp" +#include "../jit_lto_kernels/launcher_factory.hpp" +#include "../utils/memory.hpp" +#include "../utils/searcher_gpu_utils.hpp" +#include "searcher_gpu.cuh" +#include "searcher_gpu_common.cuh" + +#include +#include + +#include + +#include + +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +SearcherGPU::SearcherGPU(raft::resources const& handle, + const float* q, + size_t d, + size_t ex_bits, + std::string mode, + DataQuantizerGPU::FastQuantizeFactors* fast_quantize_factors, + bool rabitq_quantize_flag) + : D(d), query_(q), rabitq_quantize_flag_(rabitq_quantize_flag), mode_(mode), handle_(handle) +{ + RAFT_EXPECTS(D % 32 == 0, "D (%zu) must be divisible by 32", D); + RAFT_EXPECTS( + D % BITS_PER_CHUNK == 0, "D (%zu) must be divisible by BITS_PER_CHUNK (%d)", D, BITS_PER_CHUNK); + set_unit_q(memory::align_mm<64, float>(D * sizeof(float))); + set_quant_query(memory::align_mm<64, int16_t>(D * sizeof(int16_t))); + if (mode_ == "quant4" || mode_ == "quant8") { + RAFT_EXPECTS(fast_quantize_factors != nullptr, + "fast_quantize_factors must be set for quant4/quant8 mode"); + if (mode_ == "quant4") { + best_rescaling_factor = fast_quantize_factors->const_scaling_factor_4bit; + } else if (mode_ == "quant8") { + best_rescaling_factor = fast_quantize_factors->const_scaling_factor_8bit; + } + } +} + +void SearcherGPU::AllocateSearcherSpace(size_t num_centroids, size_t num_queries) +{ + centroid_distances_ = + raft::make_device_vector(handle_, num_queries * num_centroids); + q_norms_ = raft::make_device_vector(handle_, num_queries); +}; + +__global__ void precomputeAllLUTs(const float* d_query, // Query vectors + float* d_lut_for_queries, // Output LUTs for all queries + size_t num_queries, // Number of queries + size_t D // Dimension +) +{ + // Each block handles one query + const int query_idx = blockIdx.x; + if (query_idx >= num_queries) return; + + const int tid = threadIdx.x; + const int num_threads = blockDim.x; + + // Calculate LUT parameters + const size_t num_chunks = D / BITS_PER_CHUNK; + const size_t lut_per_query_size = num_chunks * LUT_SIZE; + + // Pointer to this query's LUT in global memory + float* query_lut = d_lut_for_queries + query_idx * lut_per_query_size; + + // Pointer to this query's vector + const float* query_vec = d_query + query_idx * D; + + // Each thread computes part of the LUT + for (size_t chunk_idx = tid; chunk_idx < num_chunks; chunk_idx += num_threads) { + size_t dim_start = chunk_idx * BITS_PER_CHUNK; + + // Compute LUT entries for this chunk + for (int lut_entry = 0; lut_entry < LUT_SIZE; lut_entry++) { + float sum = 0.0f; + + // For each bit in the 4-bit pattern + for (int bit_idx = 0; bit_idx < BITS_PER_CHUNK; bit_idx++) { + size_t dim = dim_start + bit_idx; + if (dim < D) { // Check if within actual dimension + // Check if bit is set in the pattern + if (lut_entry & (1 << (BITS_PER_CHUNK - 1 - bit_idx))) { sum += query_vec[dim]; } + } + } + + // Store in global LUT + size_t lut_offset = chunk_idx * LUT_SIZE + lut_entry; + query_lut[lut_offset] = sum; + } + } +} + +// Launch function for precomputing LUTs +void launchPrecomputeLUTs(const float* d_query, + float* d_lut_for_queries, + size_t num_queries, + size_t D, + rmm::cuda_stream_view stream) +{ + // Launch precompute kernel + dim3 gridDim(num_queries, 1, 1); + dim3 blockDim(256, 1, 1); // Can tune this + + precomputeAllLUTs<<>>(d_query, d_lut_for_queries, num_queries, D); + RAFT_CUDA_TRY(cudaPeekAtLastError()); +} + +void SearcherGPU::SearchClusterQueryPairs( + const IVFGPU& cur_ivf, + IVFGPU::GPUClusterMeta* d_cluster_meta, + ClusterQueryPair* d_sorted_pairs, + size_t num_queries, + raft::device_matrix_view queries, + const float* d_G_k1xSumq, + const float* d_G_kbxSumq, + size_t nprobe, + size_t topk, + raft::device_matrix_view d_final_dists, + raft::device_matrix_view d_final_pids) +{ + // First allocate space for LUT + size_t lut_size = num_queries * (D / BITS_PER_CHUNK) * LUT_SIZE * sizeof(float); + rmm::device_uvector d_lut_for_queries(lut_size / sizeof(float), stream_); + thrust::fill(thrust::cuda::par.on(stream_), + d_lut_for_queries.data(), + d_lut_for_queries.data() + d_lut_for_queries.size(), + -std::numeric_limits::infinity()); + // precompute LUTS + launchPrecomputeLUTs(queries.data_handle(), d_lut_for_queries.data(), num_queries, D, stream_); + + // 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 + const bool use_block_sort{topk <= kMaxTopKBlockSort}; + + // We minimize max_cluster_size to reduce shared memory usage when the probe clusters do not + // include the largest cluster. This optimization is expected to be more effective when + // num_queries and/or nprobe are low. + uint32_t max_cluster_size; + + // For the intermediate distances (and associated IDs), we want to minimize the allocation both to + // reduce memory footprint and to avoid unnecessary passes in the final RAFT select_k call. For + // `use_block_sort = true`, the required allocation is simply num_queries * nprobe * topk. For + // `use_block_sort = false`, the strategy here is to compute the sum of cluster sizes over all + // probed clusters for each query, and use the maximum of these sums as the allocation size needed + // per query. This avoids negative performance impact from any abnormally large cluster when using + // the global maximum cluster size as the allocation size per query per probe. + std::optional max_probed_vectors_count = + use_block_sort ? std::nullopt : std::optional{0}; + + // call utility function to evaluate max_cluster_size and max_probed_vectors_count + get_max_probed_cluster_size_and_vectors_count(handle_, + d_sorted_pairs, + num_queries * nprobe, + cur_ivf.get_cluster_meta().data_handle(), + num_queries, + max_cluster_size, + max_probed_vectors_count); + + // allocate memory for intermediate output + size_t n_cols = use_block_sort ? nprobe * topk : max_probed_vectors_count.value(); + auto d_topk_dists = raft::make_device_matrix(handle_, num_queries, n_cols); + auto d_topk_pids = raft::make_device_matrix(handle_, num_queries, n_cols); + + // initialize distances + thrust::fill(thrust::cuda::par.on(stream_), + d_topk_dists.data_handle(), + d_topk_dists.data_handle() + d_topk_dists.size(), + std::numeric_limits::infinity()); + + rmm::device_uvector d_query_write_counters(num_queries, stream_); + thrust::fill(thrust::cuda::par.on(stream_), + d_query_write_counters.data(), + d_query_write_counters.data() + num_queries, + 0); + + rmm::device_uvector d_topk_threshold_batch(use_block_sort ? num_queries : 0, stream_); + if (use_block_sort) { + thrust::fill(thrust::cuda::par.on(stream_), + d_topk_threshold_batch.data(), + d_topk_threshold_batch.data() + num_queries, + std::numeric_limits::infinity()); + } + // Then launch kernel for computation + size_t num_pairs = num_queries * nprobe; + uint32_t gridDim{static_cast(num_pairs)}; + uint32_t blockDim{256}; + size_t num_chunks = D / BITS_PER_CHUNK; + size_t candidate_storage = + max_cluster_size * (use_block_sort ? (2 * sizeof(float) /* ip, ip2 */ + sizeof(int) /* idx */) + : sizeof(float) /* ip2 */); + const int queue_buffer_smem_bytes = + use_block_sort ? raft::matrix::detail::select::warpsort::calc_smem_size_for_block_wide( + blockDim / raft::WarpSize, kMaxTopKBlockSort) + : 0; + + ComputeInnerProductsKernelParams kernelParams; + kernelParams.d_sorted_pairs = d_sorted_pairs; + kernelParams.d_query = queries.data_handle(); + kernelParams.d_short_data = cur_ivf.get_short_data_device(); + kernelParams.d_cluster_meta = d_cluster_meta; + kernelParams.d_lut_for_queries_float = d_lut_for_queries.data(); + kernelParams.d_short_factors = cur_ivf.get_short_factors_batch_device(); + kernelParams.d_G_k1xSumq = d_G_k1xSumq; + kernelParams.d_G_kbxSumq = d_G_kbxSumq; + kernelParams.d_centroid_distances = get_centroid_distances(); + kernelParams.topk = topk; + kernelParams.num_queries = num_queries; + kernelParams.nprobe = nprobe; + kernelParams.num_pairs = num_pairs; + kernelParams.num_centroids = cur_ivf.get_num_centroids(); + kernelParams.D = D; + kernelParams.d_threshold = d_topk_threshold_batch.data(); + 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.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) + : make_compute_inner_products_with_lut_launcher( + cur_ivf.get_ex_bits(), /*with_ex=*/true); + auto const& kernel_launcher = [&]() -> void { + jit_launcher->dispatch( + stream_, gridDim, blockDim, shared_mem_size, kernelParams); + }; + cuvs::neighbors::detail::safely_launch_kernel_with_smem_size< + compute_inner_products_with_lut_func_t>( + static_cast(shared_mem_size), kernel_launcher, jit_launcher->get_kernel()); + } else { + size_t shared_mem_size = + 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) + : make_compute_inner_products_with_lut_launcher( + /*ex_bits=*/0, /*with_ex=*/false); + auto const& kernel_launcher = [&]() -> void { + jit_launcher->dispatch( + stream_, gridDim, blockDim, shared_mem_size, kernelParams); + }; + cuvs::neighbors::detail::safely_launch_kernel_with_smem_size< + compute_inner_products_with_lut_func_t>( + static_cast(shared_mem_size), kernel_launcher, jit_launcher->get_kernel()); + } + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // merge results from different blocks + cuvs::selection::select_k( + handle_, + raft::make_device_matrix_view( + d_topk_dists.data_handle(), num_queries, n_cols), + std::make_optional(raft::make_device_matrix_view( + d_topk_pids.data_handle(), num_queries, n_cols)), + d_final_dists, + d_final_pids, + /*select_min=*/true, + /*sorted=*/false); +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh new file mode 100644 index 0000000000..73557f57ea --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh @@ -0,0 +1,134 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +// +// Created by Stardust on 4/14/25. +// + +#pragma once + +#include "ivf_gpu.cuh" +#include "quantizer_gpu.cuh" + +#include +#include +#include +#include + +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// block-level sorting used if topk <= kMaxTopKBlockSort; must be power of 2; increases shared mem +// usage +static constexpr int kMaxTopKBlockSort = 64; + +class SearcherGPU { + public: + /** + * @param d vector dimension; must be divisible by 32 and by BITS_PER_CHUNK. Pass the padded + * dimension (e.g. IVFGPU::get_num_padded_dim()), not the raw dataset dimension. + */ + explicit SearcherGPU(raft::resources const& handle, + const float* q, + size_t d, + size_t ex_bits, + std::string mode = "", + DataQuantizerGPU::FastQuantizeFactors* fast_quantize_factors = nullptr, + bool rabitq_quantize_flag = true); + + SearcherGPU(SearcherGPU&& other) = default; + + /** + * @brief Allocate temp search space in GPU memory + * + * @param num_centroids number of cluster centroids + * @param num_queries number of queries + */ + void AllocateSearcherSpace(size_t num_centroids, size_t num_queries); + + // Getter methods + std::string const& get_mode() { return mode_; } + raft::resources const& get_handle() const { return handle_; } + rmm::cuda_stream_view get_stream() const { return stream_; } + float* get_centroid_distances() { return centroid_distances_.data_handle(); } + float* get_q_norms() { return q_norms_.data_handle(); } + + // Setter methods + void set_query(const float* query) { query_ = query; } + void set_quant_query(int16_t* quant_query) + { + quant_query_ = std::unique_ptr(quant_query, std::free); + } // quant_query must be allocated with `malloc` or similar, as opposed to `new` + void set_unit_q(float* unit_q) + { + unit_q_ = std::unique_ptr(unit_q, std::free); + } // unit_q must be allocated with `malloc` or similar, as opposed to `new` + + /** @brief Search pairs */ + void SearchClusterQueryPairs( + const IVFGPU& cur_ivf, + IVFGPU::GPUClusterMeta* d_cluster_meta, + ClusterQueryPair* d_sorted_pairs, + size_t num_queries, + raft::device_matrix_view queries, + const float* d_G_k1xSumq, + const float* d_G_kbxSumq, + size_t nprobe, + size_t topk, + raft::device_matrix_view d_final_dists, + raft::device_matrix_view d_final_pids); + + void SearchClusterQueryPairsSharedMemOpt( + const IVFGPU& cur_ivf, + IVFGPU::GPUClusterMeta* d_cluster_meta, + ClusterQueryPair* d_sorted_pairs, + size_t num_queries, + raft::device_matrix_view queries, + const float* d_G_k1xSumq, + const float* d_G_kbxSumq, + size_t nprobe, + size_t topk, + raft::device_matrix_view d_final_dists, + raft::device_matrix_view d_final_pids); + + void SearchClusterQueryPairsQuantizeQuery( + const IVFGPU& cur_ivf, + IVFGPU::GPUClusterMeta* d_cluster_meta, + ClusterQueryPair* d_sorted_pairs, + size_t num_queries, + raft::device_matrix_view queries, + const float* d_G_k1xSumq, + const float* d_G_kbxSumq, + size_t nprobe, + size_t topk, + raft::device_matrix_view d_final_dists, + raft::device_matrix_view d_final_pids, + bool use_4bit = false); + + private: + raft::resources const& handle_; // reusable resource handle + rmm::cuda_stream_view stream_ = + raft::resource::get_cuda_stream(handle_); // CUDA stream obtained from handle_ + size_t D; // number of dimension + const float* query_ = nullptr; // rotated query (non-owning) + std::unique_ptr quant_query_ = { + nullptr, std::free}; // quantized query (to 2 bytes) + std::unique_ptr unit_q_ = {nullptr, std::free}; + bool rabitq_quantize_flag_; + std::string mode_; + + // batch + raft::device_vector centroid_distances_ = + raft::make_device_vector( + handle_, 0); // stores distances between centroids and queries (l2 square) + raft::device_vector q_norms_ = + raft::make_device_vector(handle_, 0); // query norms + + // quantization for queries + float best_rescaling_factor = 0.0f; +}; + +} // namespace cuvs::neighbors::ivf_rabitq::detail 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 new file mode 100644 index 0000000000..1e004d9164 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh @@ -0,0 +1,65 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +// +// Created by Stardust on 4/14/25. +// + +#include "searcher_gpu.cuh" + +#include + +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +static constexpr int BITS_PER_CHUNK = 4; +static constexpr int LUT_SIZE = (1 << BITS_PER_CHUNK); // 16 + +// --- Tunables --- +using T = float; +using IdxT = uint32_t; + +using lut_dtype = __half; // FP16 + +// POD struct consolidating parameters for all computeInnerProducts* kernels +struct ComputeInnerProductsKernelParams { + const ClusterQueryPair* d_sorted_pairs = nullptr; + const float* d_query = nullptr; + const uint32_t* d_short_data = nullptr; + const IVFGPU::GPUClusterMeta* d_cluster_meta = nullptr; + float* d_lut_for_queries_float = nullptr; + lut_dtype* d_lut_for_queries_half = nullptr; + const uint32_t* d_packed_queries = nullptr; // Packed query bit planes + const float* d_widths = nullptr; // Query scaling factors + const float* d_short_factors = nullptr; + const float* d_G_k1xSumq = nullptr; + const float* d_G_kbxSumq = nullptr; + const float* d_centroid_distances = nullptr; + uint32_t topk = 0; + uint32_t num_queries = 0; + uint32_t nprobe = 0; + uint32_t num_pairs = 0; + uint32_t num_centroids = 0; + uint32_t D = 0; + const float* d_threshold = nullptr; // threshold for each query + uint32_t max_candidates_per_pair = 0; // max storage per pair, 1000 suggested + uint32_t max_candidates_per_query = + 0; // max number of vectors in probed clusters for any particular query + uint32_t ex_bits = 0; // bits per dimension in ex codes + 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 + float* d_topk_dists = nullptr; // output top-k distances + PID* d_topk_pids = nullptr; // output top-k PIDs + int* d_query_write_counters = nullptr; + uint32_t num_bits = 0; // number of bits (8 for int8) + uint32_t num_words = 0; // approx. D/32 +}; + +} // namespace cuvs::neighbors::ivf_rabitq::detail 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 new file mode 100644 index 0000000000..92906fb0e6 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu @@ -0,0 +1,549 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +// +// Created by Stardust on 4/14/25. +// + +// This file implements `SearcherGPU::SearchClusterQueryPairsQuantizeQuery`. +#include "../../detail/smem_utils.cuh" +#include "../jit_lto_kernels/kernel_def.hpp" +#include "../jit_lto_kernels/launcher_factory.hpp" +#include "../utils/reductions.cuh" +#include "../utils/searcher_gpu_utils.hpp" +#include "searcher_gpu.cuh" +#include "searcher_gpu_common.cuh" + +#include +#include + +#include + +#include + +#include +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +//--------------------------------------------------------------------------- +// Kernel: exrabitq_quantize_query +// +// Quantize queries using exrabitq implementation, the output are always int8_t array +// +template +__global__ void exrabitq_quantize_query( + // Inputs + const float* __restrict__ d_XP, + size_t num_points, + size_t D, + size_t EX_BITS, + float const_scaling_factor, + float kConstEpsilon, + // Outputs + int8_t* d_long_code, + float* d_delta) +{ + //========================================================================= + // Setup: One block per row + //========================================================================= + int row = blockIdx.x; + if (row >= num_points) return; + + // Dynamically allocated shared memory for one row's data. + extern __shared__ float s_mem[]; + float* s_xp = s_mem; + float* s_reduce = s_xp + D; // For reduction + + int tid = threadIdx.x; + + //========================================================================= + // Step 0: Load XP and compute L2 nrom && normalize + //========================================================================= + float thread_sum_sq = 0.0f; + + // local L2 norm + for (int j = tid; j < D; j += BlockSize) { + float xp_val = d_XP[row * D + j]; + s_xp[j] = xp_val; + thread_sum_sq += xp_val * xp_val; // Direct L2 norm of XP + } + + s_reduce[tid] = thread_sum_sq; + __syncthreads(); + + // global reduction + for (unsigned int stride = BlockSize / 2; stride > 0; stride >>= 1) { + if (tid < stride) { s_reduce[tid] += s_reduce[tid + stride]; } + __syncthreads(); + } + + float norm = sqrtf(s_reduce[0]); + float norm_inv = (norm > 0) ? (1.0f / norm) : 0.0f; + + //========================================================================= + // Part A+B+C: Quantize, accumulate factors, and write codes + //========================================================================= + // Each thread fully owns the j values it visits, so code_val stays in a + // register: there is no need to materialize an int8_t scratch in shared + // memory or to read it back in separate passes. + const int max_code = (1 << (EX_BITS - 1)) - 1; + const int min_code = -(1 << (EX_BITS - 1)); + int8_t* out_ptr = d_long_code + row * D; // long_code_length == D + + float ip_resi_xucb = 0.f, xu_sq = 0.f; + for (int j = tid; j < D; j += BlockSize) { + float val = s_xp[j] * norm_inv; + int code_val = __float2int_rn(const_scaling_factor * val); // round-to-nearest-even + code_val = std::clamp(code_val, min_code, max_code); + + float xu = float(code_val); + // just ignore the 0.5 since we are not going to store extra shift + ip_resi_xucb += s_xp[j] * xu; // for cos_similarity + xu_sq += xu * xu; // norm_quan^2 + out_ptr[j] = static_cast(code_val); + } + + // only thread 0 in the block needs the results, so simply use blockReduceSum + ip_resi_xucb = blockReduceSum(ip_resi_xucb); + xu_sq = blockReduceSum(xu_sq); + + // Thread 0 computes and writes the final factors + if (tid == 0) { + float norm_quan = sqrtf(fmaxf(xu_sq, 0.f)); + float delta; + if (norm > 1e-6f && norm_quan > 1e-6f) { + float cos_similarity = ip_resi_xucb / (norm * norm_quan); + delta = norm / norm_quan * cos_similarity; + } else { + delta = 0.f; + } + + size_t base = row; + d_delta[base] = delta; + } +} + +__global__ void findQueryRanges(const float* __restrict__ queries, + float* __restrict__ query_ranges, + int num_queries, + int num_dimensions) +{ + const int query_idx = blockIdx.x; + if (query_idx >= num_queries) return; + + const float* query = queries + query_idx * num_dimensions; + + using BlockReduceFloat = cub::BlockReduce; + __shared__ typename BlockReduceFloat::TempStorage temp_storage_min; + __shared__ typename BlockReduceFloat::TempStorage temp_storage_max; + + float local_min = FLT_MAX; + float local_max = -FLT_MAX; + + for (int i = threadIdx.x; i < num_dimensions; i += blockDim.x) { + float val = query[i]; + local_min = fminf(local_min, val); + local_max = fmaxf(local_max, val); + } + + float block_min = BlockReduceFloat(temp_storage_min).Reduce(local_min, cuda::minimum<>{}); + float block_max = BlockReduceFloat(temp_storage_max).Reduce(local_max, cuda::maximum<>{}); + + if (threadIdx.x == 0) { + query_ranges[query_idx * 2] = block_min; + query_ranges[query_idx * 2 + 1] = block_max; + } +} + +__global__ void quantizeQueriesToInt8(const float* __restrict__ queries, + const float* __restrict__ query_ranges, + int8_t* __restrict__ quantized_queries, + float* __restrict__ widths, + int num_queries, + int num_dimensions) +{ + const int BQ = 8; // Use full 8-bit range + const float max_int_val = (1 << (BQ - 1)) - 1; // 127 + + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total_elements = num_queries * num_dimensions; + + for (int i = idx; i < total_elements; i += gridDim.x * blockDim.x) { + int query_idx = i / num_dimensions; + int dim_idx = i % num_dimensions; + + float vmin = query_ranges[query_idx * 2]; + float vmax = query_ranges[query_idx * 2 + 1]; + float vmax_abs = fmaxf(fabsf(vmin), fabsf(vmax)); + + float width = vmax_abs / max_int_val; + float one_over_width = (width > 0) ? 1.0f / width : 0.0f; + + if (dim_idx == 0) { widths[query_idx] = width; } + + float val = queries[query_idx * num_dimensions + dim_idx]; + float scaled = val * one_over_width; + scaled = fmaxf(-128.0f, fminf(127.0f, scaled)); + int8_t quantized = (int8_t)__float2int_rn(scaled); + + quantized_queries[query_idx * num_dimensions + dim_idx] = quantized; + } +} + +__global__ void packInt8QueryBitPlanes(const int8_t* __restrict__ queries, + uint32_t* __restrict__ packed_queries, + int num_queries, + int num_dimensions) +{ + const int dims_per_word = 32; + const int num_words = (num_dimensions + dims_per_word - 1) / dims_per_word; + const int num_bits = 8; + + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total_elements = num_queries * num_bits * num_words; + + for (int i = idx; i < total_elements; i += gridDim.x * blockDim.x) { + int query_idx = i / (num_bits * num_words); + int remainder = i % (num_bits * num_words); + int bit_idx = remainder / num_words; + int word_idx = remainder % num_words; + + uint32_t packed_word = 0; + +#pragma unroll 8 + for (int d = 0; d < dims_per_word; ++d) { + int dim_idx = word_idx * dims_per_word + d; + if (dim_idx < num_dimensions) { + uint8_t val = (uint8_t)queries[query_idx * num_dimensions + dim_idx]; + uint32_t bit_val = (val >> bit_idx) & 1; + + // FIXED: Match data bit ordering - dim 0 goes to bit 31, dim 31 to bit 0 + int bit_position = 31 - d; // Reverse bit ordering! + packed_word |= (bit_val << bit_position); + } + } + + packed_queries[i] = packed_word; + } +} + +__global__ void quantizeQueriesToInt4( + const float* __restrict__ queries, + const float* __restrict__ query_ranges, + int8_t* __restrict__ quantized_queries, // Still use int8_t for storage + float* __restrict__ widths, + int num_queries, + int num_dimensions) +{ + const int BQ = 4; // Use 4-bit range + const float max_int_val = (1 << (BQ - 1)) - 1; // 2^3 - 1 = 7 + + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total_elements = num_queries * num_dimensions; + + for (int i = idx; i < total_elements; i += gridDim.x * blockDim.x) { + int query_idx = i / num_dimensions; + int dim_idx = i % num_dimensions; + + float vmin = query_ranges[query_idx * 2]; + float vmax = query_ranges[query_idx * 2 + 1]; + float vmax_abs = fmaxf(fabsf(vmin), fabsf(vmax)); + + float width = vmax_abs / max_int_val; + float one_over_width = (width > 0) ? 1.0f / width : 0.0f; + + if (dim_idx == 0) { widths[query_idx] = width; } + + float val = queries[query_idx * num_dimensions + dim_idx]; + float scaled = val * one_over_width; + + // Clamp to 4-bit range [-8, 7] + scaled = fmaxf(-8.0f, fminf(7.0f, scaled)); + int8_t quantized = (int8_t)__float2int_rn(scaled); + + quantized_queries[query_idx * num_dimensions + dim_idx] = quantized; + } +} + +__global__ void packInt4QueryBitPlanes(const int8_t* __restrict__ queries, + uint32_t* __restrict__ packed_queries, + int num_queries, + int num_dimensions) +{ + const int dims_per_word = 32; + const int num_words = (num_dimensions + dims_per_word - 1) / dims_per_word; + const int num_bits = 4; // Only 4 bit planes! + + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total_elements = num_queries * num_bits * num_words; + + for (int i = idx; i < total_elements; i += gridDim.x * blockDim.x) { + int query_idx = i / (num_bits * num_words); + int remainder = i % (num_bits * num_words); + int bit_idx = remainder / num_words; + int word_idx = remainder % num_words; + + uint32_t packed_word = 0; + +#pragma unroll 8 + for (int d = 0; d < dims_per_word; ++d) { + int dim_idx = word_idx * dims_per_word + d; + if (dim_idx < num_dimensions) { + // For 4-bit values, we only care about the lower 4 bits + // But need to handle sign extension properly + uint8_t val = (uint8_t)(queries[query_idx * num_dimensions + dim_idx] & 0xF); + uint32_t bit_val = (val >> bit_idx) & 1; + + // Match data bit ordering - dim 0 goes to bit 31 + int bit_position = 31 - d; + packed_word |= (bit_val << bit_position); + } + } + + packed_queries[i] = packed_word; + } +} + +// Search with qunatized query vectors +void SearcherGPU::SearchClusterQueryPairsQuantizeQuery( + const IVFGPU& cur_ivf, + IVFGPU::GPUClusterMeta* d_cluster_meta, + ClusterQueryPair* d_sorted_pairs, + size_t num_queries, + raft::device_matrix_view queries, + const float* d_G_k1xSumq, + const float* d_G_kbxSumq, + size_t nprobe, + 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 +) +{ + // 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 + const bool use_block_sort{topk <= kMaxTopKBlockSort}; + + // query quantize + const int num_bits = use_4bit ? 4 : 8; // Choose bit width + const int num_words = (cur_ivf.get_num_padded_dim() + 31) / 32; + + // Allocate memory for quantization + auto d_query_write_counters = raft::make_device_vector(handle_, num_queries); + auto d_query_ranges = raft::make_device_vector(handle_, 0); + auto d_widths = raft::make_device_vector(handle_, 0); + auto d_quantized_queries = raft::make_device_vector(handle_, 0); + auto d_packed_queries = raft::make_device_vector(handle_, 0); + auto d_topk_threshold_batch = raft::make_device_vector(handle_, 0); + if (use_block_sort) { + d_query_ranges = raft::make_device_vector(handle_, num_queries * 2); + d_widths = raft::make_device_vector(handle_, num_queries); + d_quantized_queries = raft::make_device_vector( + handle_, num_queries * cur_ivf.get_num_padded_dim()); + d_packed_queries = + raft::make_device_vector(handle_, num_queries * num_bits * num_words); + d_topk_threshold_batch = raft::make_device_vector(handle_, num_queries); + } + + if (use_block_sort) { + if (rabitq_quantize_flag_) { + const int block_size = 256; + const int grid_size = num_queries; + size_t shared_mem = D * sizeof(float) + block_size * sizeof(float); + exrabitq_quantize_query + <<>>(queries.data_handle(), + num_queries, + D, + num_bits, + best_rescaling_factor, + 1.9f, + d_quantized_queries.data_handle(), + d_widths.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } else { // scalar quantize + // Step 1: Find min/max for each query + const int block_size = 256; + const int grid_size = num_queries; + findQueryRanges<<>>(queries.data_handle(), + d_query_ranges.data_handle(), + num_queries, + cur_ivf.get_num_padded_dim()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // Step 2: Quantize queries to int8_t with BQ=8 + if (use_4bit) { + quantizeQueriesToInt4<<>>( + queries.data_handle(), + d_query_ranges.data_handle(), + d_quantized_queries.data_handle(), + d_widths.data_handle(), + num_queries, + cur_ivf.get_num_padded_dim()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } else { + quantizeQueriesToInt8<<>>( + queries.data_handle(), + d_query_ranges.data_handle(), + d_quantized_queries.data_handle(), + d_widths.data_handle(), + num_queries, + cur_ivf.get_num_padded_dim()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + } + } + + // Step 3: Pack quantized queries into bit planes + if (use_block_sort) { + const int block_size = 256; + const int grid_size = (num_queries * num_bits * num_words + block_size - 1) / block_size; + + if (use_4bit) { + packInt4QueryBitPlanes<<>>( + d_quantized_queries.data_handle(), + d_packed_queries.data_handle(), + num_queries, + cur_ivf.get_num_padded_dim()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } else { + packInt8QueryBitPlanes<<>>( + d_quantized_queries.data_handle(), + d_packed_queries.data_handle(), + num_queries, + cur_ivf.get_num_padded_dim()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + } + + // We minimize max_cluster_size to reduce shared memory usage when the probe clusters do not + // include the largest cluster. This optimization is expected to be more effective when + // num_queries and/or nprobe are low. + uint32_t max_cluster_size; + + // For the intermediate distances (and associated IDs), we want to minimize the allocation both to + // reduce memory footprint and to avoid unnecessary passes in the final RAFT select_k call. For + // `use_block_sort = true`, the required allocation is simply num_queries * nprobe * topk. For + // `use_block_sort = false`, the strategy here is to compute the sum of cluster sizes over all + // probed clusters for each query, and use the maximum of these sums as the allocation size needed + // per query. This avoids negative performance impact from any abnormally large cluster when using + // the global maximum cluster size as the allocation size per query per probe. + std::optional max_probed_vectors_count = + use_block_sort ? std::nullopt : std::optional{0}; + + // call utility function to evaluate max_cluster_size and max_probed_vectors_count + get_max_probed_cluster_size_and_vectors_count(handle_, + d_sorted_pairs, + num_queries * nprobe, + cur_ivf.get_cluster_meta().data_handle(), + num_queries, + max_cluster_size, + max_probed_vectors_count); + + // allocate memory for intermediate output + size_t n_cols = use_block_sort ? nprobe * topk : max_probed_vectors_count.value(); + auto d_topk_dists = raft::make_device_matrix(handle_, num_queries, n_cols); + auto d_topk_pids = raft::make_device_matrix(handle_, num_queries, n_cols); + + // initialize distances + thrust::fill(thrust::cuda::par.on(stream_), + d_topk_dists.data_handle(), + d_topk_dists.data_handle() + d_topk_dists.size(), + std::numeric_limits::infinity()); + + thrust::fill(thrust::cuda::par.on(stream_), + d_query_write_counters.data_handle(), + d_query_write_counters.data_handle() + num_queries, + 0); + + if (use_block_sort) { + thrust::fill(thrust::cuda::par.on(stream_), + d_topk_threshold_batch.data_handle(), + d_topk_threshold_batch.data_handle() + num_queries, + std::numeric_limits::infinity()); + } + + // Launch modified kernel with packed queries instead of LUT + size_t num_pairs = num_queries * nprobe; + uint32_t gridDim{static_cast(num_pairs)}; + uint32_t blockDim{256}; + + // Recalculate shared memory for new approach + size_t query_storage = D * sizeof(float); // For shared query vector + const int queue_buffer_smem_bytes = + use_block_sort ? raft::matrix::detail::select::warpsort::calc_smem_size_for_block_wide( + blockDim / raft::WarpSize, kMaxTopKBlockSort) + : 0; + + // Now we need: packed query bits, candidate storage, and query vector + // this part is also used to store ip2 results + size_t packed_query_size = max((use_block_sort ? (num_bits * num_words * sizeof(uint32_t)) : 0), + max_cluster_size * sizeof(float)); + size_t candidate_storage = use_block_sort ? max_cluster_size * (sizeof(float) + sizeof(int)) : 0; + size_t shared_mem_size = + max(packed_query_size + candidate_storage + query_storage, (size_t)queue_buffer_smem_bytes); + + ComputeInnerProductsKernelParams kernelParams; + kernelParams.d_sorted_pairs = d_sorted_pairs; + kernelParams.d_query = queries.data_handle(); + kernelParams.d_short_data = cur_ivf.get_short_data_device(); + kernelParams.d_cluster_meta = d_cluster_meta; + kernelParams.d_packed_queries = d_packed_queries.data_handle(); + kernelParams.d_widths = d_widths.data_handle(); + kernelParams.d_short_factors = cur_ivf.get_short_factors_batch_device(); + kernelParams.d_G_k1xSumq = d_G_k1xSumq; + kernelParams.d_G_kbxSumq = d_G_kbxSumq; + kernelParams.d_centroid_distances = get_centroid_distances(); + kernelParams.topk = topk; + kernelParams.num_queries = num_queries; + kernelParams.nprobe = nprobe; + kernelParams.num_pairs = num_pairs; + kernelParams.num_centroids = cur_ivf.get_num_centroids(); + kernelParams.D = D; + kernelParams.d_threshold = d_topk_threshold_batch.data_handle(); + 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.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) + : make_compute_inner_products_with_bitwise_launcher( + cur_ivf.get_ex_bits(), with_ex); + auto const& kernel_launcher = [&]() -> void { + jit_launcher->dispatch( + stream_, gridDim, blockDim, shared_mem_size, kernelParams); + }; + cuvs::neighbors::detail::safely_launch_kernel_with_smem_size< + compute_inner_products_with_lut_func_t>( + static_cast(shared_mem_size), kernel_launcher, jit_launcher->get_kernel()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // Merge results + cuvs::selection::select_k( + handle_, + raft::make_device_matrix_view( + d_topk_dists.data_handle(), num_queries, n_cols), + std::make_optional(raft::make_device_matrix_view( + d_topk_pids.data_handle(), num_queries, n_cols)), + d_final_dists, + d_final_pids, + /*select_min=*/true, + /*sorted=*/false); +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail 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 new file mode 100644 index 0000000000..8357af5fbb --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu @@ -0,0 +1,264 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +// +// Created by Stardust on 4/14/25. +// + +// This file implements `SearcherGPU::SearchClusterQueryPairsSharedMemOpt`. +#include "../../detail/smem_utils.cuh" +#include "../jit_lto_kernels/kernel_def.hpp" +#include "../jit_lto_kernels/launcher_factory.hpp" +#include "../utils/searcher_gpu_utils.hpp" +#include "searcher_gpu.cuh" +#include "searcher_gpu_common.cuh" + +#include +#include + +#include + +#include + +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +__global__ void precomputeAllLUTs_fp16(const float* d_query, + lut_dtype* d_lut_for_queries, + size_t num_queries, + size_t D) +{ + const int query_idx = blockIdx.x; + if (query_idx >= num_queries) return; + + const int tid = threadIdx.x; + const int num_threads = blockDim.x; + + const size_t num_chunks = D / BITS_PER_CHUNK; + const size_t lut_per_query_size = num_chunks * LUT_SIZE; + + lut_dtype* query_lut = d_lut_for_queries + query_idx * lut_per_query_size; + const float* query_vec = d_query + query_idx * D; + + for (size_t chunk_idx = tid; chunk_idx < num_chunks; chunk_idx += num_threads) { + size_t dim_start = chunk_idx * BITS_PER_CHUNK; + + for (int lut_entry = 0; lut_entry < LUT_SIZE; lut_entry++) { + float sum = 0.0f; // Compute in FP32 + + for (int bit_idx = 0; bit_idx < BITS_PER_CHUNK; bit_idx++) { + size_t dim = dim_start + bit_idx; + if (dim < D) { + if (lut_entry & (1 << (BITS_PER_CHUNK - 1 - bit_idx))) { sum += query_vec[dim]; } + } + } + + size_t lut_offset = chunk_idx * LUT_SIZE + lut_entry; + query_lut[lut_offset] = __float2half(sum); + } + } +} + +void launchPrecomputeLUTs_fp16(const float* d_query, + lut_dtype* d_lut_for_queries, + size_t num_queries, + size_t D, + cudaStream_t stream) +{ + dim3 gridDim(num_queries, 1, 1); + dim3 blockDim(256, 1, 1); + + precomputeAllLUTs_fp16<<>>( + d_query, d_lut_for_queries, num_queries, D); + RAFT_CUDA_TRY(cudaPeekAtLastError()); +} + +// optimize shared memory usage using (1) fp16 (2) shared memory reuse in the kernel +void SearcherGPU::SearchClusterQueryPairsSharedMemOpt( + const IVFGPU& cur_ivf, + IVFGPU::GPUClusterMeta* d_cluster_meta, + ClusterQueryPair* d_sorted_pairs, + size_t num_queries, + raft::device_matrix_view queries, + const float* d_G_k1xSumq, + const float* d_G_kbxSumq, + size_t nprobe, + size_t topk, + raft::device_matrix_view d_final_dists, + raft::device_matrix_view d_final_pids) +{ + // Using FP16 for storage + + // Allocate space for LUT with reduced precision + size_t lut_elements = num_queries * (cur_ivf.get_num_padded_dim() / BITS_PER_CHUNK) * LUT_SIZE; + size_t lut_size = lut_elements * sizeof(lut_dtype); + + rmm::device_uvector d_lut_for_queries(lut_elements, stream_); + + // Initialize with -infinity (convert to FP16) + lut_dtype neg_inf_fp16 = __float2half(-std::numeric_limits::infinity()); + thrust::fill(thrust::cuda::par.on(stream_), + d_lut_for_queries.data(), + d_lut_for_queries.data() + lut_elements, + neg_inf_fp16); + + // Precompute LUTs + launchPrecomputeLUTs_fp16(queries.data_handle(), + d_lut_for_queries.data(), + num_queries, + cur_ivf.get_num_padded_dim(), + stream_); + + // 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 + const bool use_block_sort{topk <= kMaxTopKBlockSort}; + + // We minimize max_cluster_size to reduce shared memory usage when the probe clusters do not + // include the largest cluster. This optimization is expected to be more effective when + // num_queries and/or nprobe are low. + uint32_t max_cluster_size; + + // For the intermediate distances (and associated IDs), we want to minimize the allocation both to + // reduce memory footprint and to avoid unnecessary passes in the final RAFT select_k call. For + // `use_block_sort = true`, the required allocation is simply num_queries * nprobe * topk. For + // `use_block_sort = false`, the strategy here is to compute the sum of cluster sizes over all + // probed clusters for each query, and use the maximum of these sums as the allocation size needed + // per query. This avoids negative performance impact from any abnormally large cluster when using + // the global maximum cluster size as the allocation size per query per probe. + std::optional max_probed_vectors_count = + use_block_sort ? std::nullopt : std::optional{0}; + + // call utility function to evaluate max_cluster_size and max_probed_vectors_count + get_max_probed_cluster_size_and_vectors_count(handle_, + d_sorted_pairs, + num_queries * nprobe, + cur_ivf.get_cluster_meta().data_handle(), + num_queries, + max_cluster_size, + max_probed_vectors_count); + + // allocate memory for intermediate output + size_t n_cols = use_block_sort ? nprobe * topk : max_probed_vectors_count.value(); + auto d_topk_dists = raft::make_device_matrix(handle_, num_queries, n_cols); + auto d_topk_pids = raft::make_device_matrix(handle_, num_queries, n_cols); + + // initialize distances + thrust::fill(thrust::cuda::par.on(stream_), + d_topk_dists.data_handle(), + d_topk_dists.data_handle() + d_topk_dists.size(), + std::numeric_limits::infinity()); + + rmm::device_uvector d_query_write_counters(num_queries, stream_); + thrust::fill(thrust::cuda::par.on(stream_), + d_query_write_counters.data(), + d_query_write_counters.data() + num_queries, + 0); + + rmm::device_uvector d_topk_threshold_batch(use_block_sort ? num_queries : 0, stream_); + if (use_block_sort) { + thrust::fill(thrust::cuda::par.on(stream_), + d_topk_threshold_batch.data(), + d_topk_threshold_batch.data() + num_queries, + std::numeric_limits::infinity()); + } + // Then launch kernel for computation + size_t num_pairs = num_queries * nprobe; + uint32_t gridDim{static_cast(num_pairs)}; + uint32_t blockDim{256}; + const int queue_buffer_smem_bytes = + use_block_sort ? raft::matrix::detail::select::warpsort::calc_smem_size_for_block_wide( + blockDim / raft::WarpSize, kMaxTopKBlockSort) + : 0; + ComputeInnerProductsKernelParams kernelParams; + kernelParams.d_sorted_pairs = d_sorted_pairs; + kernelParams.d_query = queries.data_handle(); + kernelParams.d_short_data = cur_ivf.get_short_data_device(); + kernelParams.d_cluster_meta = d_cluster_meta; + kernelParams.d_lut_for_queries_half = d_lut_for_queries.data(); + kernelParams.d_short_factors = cur_ivf.get_short_factors_batch_device(); + kernelParams.d_G_k1xSumq = d_G_k1xSumq; + kernelParams.d_G_kbxSumq = d_G_kbxSumq; + kernelParams.d_centroid_distances = get_centroid_distances(); + kernelParams.topk = topk; + kernelParams.num_queries = num_queries; + kernelParams.nprobe = nprobe; + kernelParams.num_pairs = num_pairs; + kernelParams.num_centroids = cur_ivf.get_num_centroids(); + kernelParams.D = D; + kernelParams.d_threshold = d_topk_threshold_batch.data(); + 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.d_query_write_counters = d_query_write_counters.data(); + + if (cur_ivf.get_ex_bits() != 0) { + size_t query_storage = D * sizeof(float); // For shared query vector + size_t first_part_shared_mem = + use_block_sort ? max(lut_size / num_queries, max_cluster_size * (sizeof(float))) + : (lut_size / num_queries); + size_t second_part_shared_mem = + max_cluster_size * (sizeof(float) + (use_block_sort ? sizeof(int) : 0)); + size_t third_part_shared_mem = use_block_sort ? query_storage : 0; + // queue buffer reuses first 3 parts + 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) + : make_compute_inner_products_with_lut16_opt_launcher( + cur_ivf.get_ex_bits(), /*with_ex=*/true); + auto const& kernel_launcher = [&]() -> void { + jit_launcher->dispatch( + stream_, gridDim, blockDim, shared_mem_size, kernelParams); + }; + cuvs::neighbors::detail::safely_launch_kernel_with_smem_size< + compute_inner_products_with_lut_func_t>( + static_cast(shared_mem_size), kernel_launcher, jit_launcher->get_kernel()); + } else { + size_t first_part_shared_mem = lut_size / num_queries; + size_t second_part_shared_mem = + use_block_sort ? max_cluster_size * (sizeof(float) + sizeof(int)) : 0; + // 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) + : make_compute_inner_products_with_lut16_opt_launcher( + /*ex_bits=*/0, /*with_ex=*/false); + auto const& kernel_launcher = [&]() -> void { + jit_launcher->dispatch( + stream_, gridDim, blockDim, shared_mem_size, kernelParams); + }; + cuvs::neighbors::detail::safely_launch_kernel_with_smem_size< + compute_inner_products_with_lut_func_t>( + static_cast(shared_mem_size), kernel_launcher, jit_launcher->get_kernel()); + } + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // merge results from different blocks + cuvs::selection::select_k( + handle_, + raft::make_device_matrix_view( + d_topk_dists.data_handle(), num_queries, n_cols), + std::make_optional(raft::make_device_matrix_view( + d_topk_pids.data_handle(), num_queries, n_cols)), + d_final_dists, + d_final_pids, + /*select_min=*/true, + /*sorted=*/false); +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in new file mode 100644 index 0000000000..d8dcff324d --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in @@ -0,0 +1,185 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include + +namespace { + +constexpr bool with_ex = @with_ex_value@; + +} // namespace + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// Builds and stores the per-block top-k for the bitwise BlockSort kernel. +// WithEx refines candidate distances with IP2 from long codes; NoEx uses the +// 1-bit IPs directly. Lifted into its own JIT-LTO fragment so the kernel emits +// a single fragment regardless of WithEx, and nvJitLink resolves the matching +// body at link time. Must be called with all threads of the block and only +// when `num_candidates > 0`. Re-derives the shared-memory pointers using the +// same layout as the caller. +__device__ void bitwise_block_sort_emit_topk(const ComputeInnerProductsKernelParams params, + int num_candidates, + int query_idx, + int cluster_idx, + size_t num_vectors_in_cluster, + size_t cluster_start_index, + float q_g_add, + float q_k1xsumq, + float threshold) +{ + const int tid = threadIdx.x; + const int num_threads = blockDim.x; + + // Shared-memory layout matches the caller's setup + extern __shared__ __align__(256) char shared_mem_raw_2[]; + size_t packed_query_bytes = max(params.num_bits * params.num_words * sizeof(uint32_t), + params.max_candidates_per_pair * sizeof(float)); + float* shared_candidate_ips = reinterpret_cast(shared_mem_raw_2 + packed_query_bytes); + int* shared_candidate_indices = + reinterpret_cast(shared_candidate_ips + params.max_candidates_per_pair); + float* shared_query = (float*)(shared_candidate_indices + params.max_candidates_per_pair); + const size_t short_code_length = params.D / 32; + + const int candidates_per_thread = (num_candidates + num_threads - 1) / num_threads; + __shared__ int probe_slot; + + if constexpr (with_ex) { + // Phase 2 (WithEx): Compute exact 1-bit IPs and store for IP2 refinement + for (int c = 0; c < candidates_per_thread; ++c) { + int cand_idx = tid + c * num_threads; + if (cand_idx < num_candidates) { + int vec_idx = shared_candidate_indices[cand_idx]; + float exact_ip = compute_bitwise_1bit_ip_for_vec(params.d_short_data, + shared_query, + cluster_start_index, + num_vectors_in_cluster, + vec_idx, + short_code_length, + params.D); + shared_candidate_ips[cand_idx] = exact_ip; + } + } + __syncthreads(); + + // Phase 3 (WithEx): Warp-level IP2 computation + block-sort queue + using block_sort_t = typename rabitq_block_sort::type; + block_sort_t queue(params.topk); + + float q_kbxsumq = params.d_G_kbxSumq[query_idx]; + const uint32_t long_code_size = (params.D * params.ex_bits + 7) / 8; + float* shared_ip2_results = reinterpret_cast(shared_mem_raw_2); + + const int warp_id = tid / raft::WarpSize; + const int lane_id = tid % raft::WarpSize; + const int num_warps = num_threads / raft::WarpSize; + + for (int cand_idx = warp_id; cand_idx < num_candidates; cand_idx += num_warps) { + size_t global_vec_idx = cluster_start_index + shared_candidate_indices[cand_idx]; + const uint8_t* vec_long_code = params.d_long_code + global_vec_idx * long_code_size; + + float ip2 = compute_ip2_from_long_codes_warp(vec_long_code, shared_query, params.D, lane_id); + if (lane_id == 0) { shared_ip2_results[cand_idx] = ip2; } + } + __syncthreads(); + + for (int round = 0; round < candidates_per_thread; round++) { + int cand_idx = tid + round * num_threads; + + float ex_dist; + uint32_t pid; + if (cand_idx < num_candidates) { + float ip = shared_candidate_ips[cand_idx]; + float ip2 = shared_ip2_results[cand_idx]; + int local_vec_idx = shared_candidate_indices[cand_idx]; + size_t global_vec_idx = cluster_start_index + local_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; + + ex_dist = f_ex_add + q_g_add + + f_ex_rescale * (static_cast(1 << params.ex_bits) * ip + ip2 + q_kbxsumq); + pid = (uint32_t)params.d_pids[global_vec_idx]; + } else { + ex_dist = INFINITY; + pid = 0; + } + queue.add(ex_dist, pid); + } + __syncthreads(); + + queue.done((uint8_t*)shared_mem_raw_2); + if (tid == 0) { probe_slot = atomicAdd(¶ms.d_query_write_counters[query_idx], 1); } + __syncthreads(); + + // Preserve the original WithEx behavior: skip both the queue store and the + // threshold update when this query has already filled all of its probe slots. + if (probe_slot >= params.nprobe) { return; } + + uint32_t output_offset = query_idx * (params.topk * params.nprobe) + probe_slot * params.topk; + queue.store(params.d_topk_dists + output_offset, + (uint32_t*)(params.d_topk_pids + output_offset)); + } else { + // Phase 2+3 (NoEx): Compute exact 1-bit IPs and add directly to queue + using block_sort_t = typename rabitq_block_sort::type; + block_sort_t queue(params.topk); + + float final_dist; + PID final_pid; + for (int c = 0; c < candidates_per_thread; ++c) { + int cand_idx = tid + c * num_threads; + if (cand_idx < num_candidates) { + int vec_idx = shared_candidate_indices[cand_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; + float f_rescale = factors.y; + size_t global_vec_idx = cluster_start_index + vec_idx; + + float exact_ip = compute_bitwise_1bit_ip_for_vec(params.d_short_data, + shared_query, + cluster_start_index, + num_vectors_in_cluster, + vec_idx, + short_code_length, + params.D); + final_dist = f_add + q_g_add + f_rescale * (exact_ip + q_k1xsumq); + final_pid = (uint32_t)params.d_pids[global_vec_idx]; + } else { + final_dist = INFINITY; + final_pid = 0; + } + queue.add(final_dist, final_pid); + } + __syncthreads(); + + queue.done((uint8_t*)shared_mem_raw_2); + if (tid == 0) { probe_slot = atomicAdd(¶ms.d_query_write_counters[query_idx], 1); } + __syncthreads(); + + uint32_t output_offset = query_idx * (params.topk * params.nprobe) + probe_slot * params.topk; + queue.store(params.d_topk_dists + output_offset, + (uint32_t*)(params.d_topk_pids + output_offset)); + } + + if (num_candidates >= params.topk) { + update_threshold_atomicmin(params.d_topk_dists, + params.d_threshold, + query_idx, + params.topk, + params.nprobe, + probe_slot, + threshold, + tid); + } +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_matrix.json new file mode 100644 index 0000000000..7e0f0f3603 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_matrix.json @@ -0,0 +1,12 @@ +{ + "_with_ex": [ + { + "with_ex_descriptor": "with_ex", + "with_ex_value": "true" + }, + { + "with_ex_descriptor": "no_ex", + "with_ex_value": "false" + } + ] +} 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 new file mode 100644 index 0000000000..20c5aec6a1 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in @@ -0,0 +1,114 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include + +namespace { + +constexpr bool with_ex = @with_ex_value@; + +} // namespace + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// Computes and writes per-vector distances for the bitwise (non-BlockSort) +// kernel. WithEx precomputes warp-level IP2 from long codes and refines with +// ex_factors; NoEx uses 1-bit short codes only. Lifted into its own JIT-LTO +// fragment so the kernel emits a single fragment regardless of WithEx, and +// nvJitLink resolves the matching body at link time. Must be called by all +// threads of the block, with `shared_query` already populated. +__device__ void bitwise_emit_distances(const ComputeInnerProductsKernelParams params, + int query_idx, + int cluster_idx, + size_t num_vectors_in_cluster, + size_t cluster_start_index, + float q_g_add) +{ + const int tid = threadIdx.x; + const int num_threads = blockDim.x; + + extern __shared__ __align__(256) char shared_mem_raw[]; + float* shared_query = reinterpret_cast(shared_mem_raw); + + if constexpr (with_ex) { + // Precompute warp-level IP2 for every vector, stored in shared memory + float* shared_ip2_results = shared_query + params.D; + const uint32_t long_code_size = (params.D * params.ex_bits + 7) / 8; + const int warp_id = tid / raft::WarpSize; + const int lane_id = tid % raft::WarpSize; + const int num_warps = num_threads / raft::WarpSize; + + for (int cand_idx = warp_id; cand_idx < num_vectors_in_cluster; cand_idx += num_warps) { + size_t global_vec_idx = cluster_start_index + cand_idx; + const uint8_t* vec_long_code = params.d_long_code + global_vec_idx * long_code_size; + float ip2 = compute_ip2_from_long_codes_warp(vec_long_code, shared_query, params.D, lane_id); + if (lane_id == 0) { shared_ip2_results[cand_idx] = ip2; } + } + __syncthreads(); + } + + const size_t short_code_length = params.D / 32; + + __shared__ int probe_slot; + if (tid == 0) { + probe_slot = atomicAdd(¶ms.d_query_write_counters[query_idx], num_vectors_in_cluster); + } + __syncthreads(); + + uint32_t output_offset = query_idx * params.max_candidates_per_query + probe_slot; + + if constexpr (with_ex) { + float* shared_ip2_results = shared_query + params.D; + 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) { + float exact_ip = compute_bitwise_1bit_ip_for_vec(params.d_short_data, + shared_query, + cluster_start_index, + num_vectors_in_cluster, + vec_idx, + 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; + + params.d_topk_dists[output_offset + vec_idx] = + f_ex_add + q_g_add + + f_ex_rescale * (static_cast(1 << params.ex_bits) * exact_ip + ip2 + q_kbxsumq); + params.d_topk_pids[output_offset + vec_idx] = params.d_pids[global_vec_idx]; + } + } else { + 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; + + float exact_ip = compute_bitwise_1bit_ip_for_vec(params.d_short_data, + shared_query, + cluster_start_index, + num_vectors_in_cluster, + vec_idx, + short_code_length, + params.D); + + params.d_topk_dists[output_offset + vec_idx] = + f_add + q_g_add + f_rescale * (exact_ip + q_k1xsumq); + params.d_topk_pids[output_offset + vec_idx] = params.d_pids[global_vec_idx]; + } + } +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_matrix.json new file mode 100644 index 0000000000..7e0f0f3603 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_matrix.json @@ -0,0 +1,12 @@ +{ + "_with_ex": [ + { + "with_ex_descriptor": "with_ex", + "with_ex_value": "true" + }, + { + "with_ex_descriptor": "no_ex", + "with_ex_value": "false" + } + ] +} diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/block_sort.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/block_sort.cuh new file mode 100644 index 0000000000..a16058c2af --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/block_sort.cuh @@ -0,0 +1,38 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "../../ivf_common.cuh" + +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// IVF-RaBitQ block-level top-k sort wrapper. Mirrors the pattern used by +// ivf_pq's pq_block_sort and ivf_sq's sq_block_sort: hide the warp_sort +// variant and Ascending policy that this algorithm doesn't vary, while still +// letting consumers pick the value/index types. Capacity == 0 falls back to +// the shared dummy_block_sort_t. +template +struct rabitq_block_sort { + using type = raft::matrix::detail::select::warpsort::block_sort< + raft::matrix::detail::select::warpsort::warp_sort_filtered, + Capacity, + /*Ascending=*/true, + T, + IdxT>; +}; + +template +struct rabitq_block_sort<0, T, IdxT> + : ivf::detail::dummy_block_sort_t { + using type = ivf::detail::dummy_block_sort_t; +}; + +template +using block_sort_t = typename rabitq_block_sort::type; + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_kernel.cu.in new file mode 100644 index 0000000000..e4e2aa6d79 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_kernel.cu.in @@ -0,0 +1,46 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +namespace { + +constexpr int num_bits = @num_bits@; + +} // namespace + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// Non-inline __device__ definition. num_bits is captured from the matrix.json +// axis into a constexpr so the inner #pragma-unrolled (b) loop has a +// compile-time bound and the (num_bits - 1) shift becomes an immediate. Each +// matrix.json value produces a separate fragment with the same mangled symbol +// name; the planner adds exactly one via +// add_compute_bitwise_quantized_ip_for_vec_device_function(), so +// nvJitLink resolves consumer calls to the matching specialization with no +// runtime dispatch. +__device__ int32_t compute_bitwise_quantized_ip_for_vec(const uint32_t* d_short_data, + const uint32_t* shared_packed_query, + size_t cluster_start_index, + size_t num_vectors_in_cluster, + size_t vec_idx, + uint32_t num_words) +{ + int32_t accumulator = 0; + for (uint32_t word = 0; word < num_words; ++word) { + size_t data_offset = cluster_start_index * num_words + word * num_vectors_in_cluster + vec_idx; + uint32_t data_word = d_short_data[data_offset]; + +#pragma unroll + for (int b = 0; b < num_bits - 1; b++) { + accumulator += __popc(shared_packed_query[b * num_words + word] & data_word) << b; + } + accumulator -= __popc(shared_packed_query[(num_bits - 1) * num_words + word] & data_word) + << (num_bits - 1); + } + return accumulator; +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_matrix.json new file mode 100644 index 0000000000..30b12de467 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_matrix.json @@ -0,0 +1,3 @@ +{ + "num_bits": ["4", "8"] +} 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 new file mode 100644 index 0000000000..a577ec44ce --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh @@ -0,0 +1,133 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "../../detail/smem_utils.cuh" +#include "../gpu_index/searcher_gpu_common.cuh" +#include "block_sort.cuh" +#include "device_functions.cuh" + +#include + +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// Unified kernel using BlockSort. The WithEx-axis-specific tail (refine + sort +// for WithEx=true, direct sort for WithEx=false) lives in the +// bitwise_block_sort_emit_topk JIT-LTO fragment, so this kernel is emitted as +// a single fragment regardless of WithEx. The num_bits-dependent popc +// inner-product loop is dispatched at runtime through the +// compute_bitwise_quantized_ip_for_vec JIT-LTO fragment, so this kernel is +// also not templated on num_bits. +__device__ void compute_inner_products_with_bitwise_block_sort_impl( + const ComputeInnerProductsKernelParams params) +{ + const int block_id = blockIdx.x; + if (block_id >= params.num_pairs) return; + + ClusterQueryPair pair = params.d_sorted_pairs[block_id]; + int cluster_idx = pair.cluster_idx; + int query_idx = pair.query_idx; + + if (cluster_idx >= params.num_centroids || query_idx >= params.num_queries) return; + + size_t num_vectors_in_cluster = params.d_cluster_meta[cluster_idx].num; + size_t cluster_start_index = params.d_cluster_meta[cluster_idx].start_index; + + extern __shared__ __align__(256) char shared_mem_raw_2[]; + uint32_t* shared_packed_query = reinterpret_cast(shared_mem_raw_2); + + const int tid = threadIdx.x; + const int num_threads = blockDim.x; + + const uint32_t* query_packed_ptr = + params.d_packed_queries + query_idx * params.num_bits * params.num_words; + for (uint32_t i = tid; i < params.num_bits * params.num_words; i += num_threads) { + shared_packed_query[i] = query_packed_ptr[i]; + } + + float query_width = params.d_widths[query_idx]; + + __shared__ int num_candidates; + float q_g_add = params.d_centroid_distances[query_idx * params.num_centroids + cluster_idx]; + float q_k1xsumq = params.d_G_k1xSumq[query_idx]; + float q_g_error = sqrtf(q_g_add); + float threshold = params.d_threshold[query_idx]; + + if (tid == 0) { num_candidates = 0; } + __syncthreads(); + + size_t packed_query_bytes = max(params.num_bits * params.num_words * sizeof(uint32_t), + params.max_candidates_per_pair * sizeof(float)); + float* shared_candidate_ips = reinterpret_cast(shared_mem_raw_2 + packed_query_bytes); + int* shared_candidate_indices = + reinterpret_cast(shared_candidate_ips + params.max_candidates_per_pair); + float* shared_query = (float*)(shared_candidate_indices + params.max_candidates_per_pair); + + // Phase 1: Bitwise inner product filter + for (size_t vec_base = 0; vec_base < num_vectors_in_cluster; vec_base += num_threads) { + size_t vec_idx = vec_base + tid; + + bool is_candidate = false; + float local_ip_quantized = 0; + + if (vec_idx < num_vectors_in_cluster) { + 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; + float f_error = factors.z; + + int32_t accumulator = compute_bitwise_quantized_ip_for_vec(params.d_short_data, + shared_packed_query, + cluster_start_index, + num_vectors_in_cluster, + vec_idx, + params.num_words); + + float ip = (float)accumulator * query_width; + float est_dist = f_add + q_g_add + f_rescale * (ip + q_k1xsumq); + float low_dist = est_dist - f_error * q_g_error; + + if (low_dist < threshold) { + is_candidate = true; + local_ip_quantized = ip; + } + } + + __syncwarp(); + + if (is_candidate) { + int candidate_slot = atomicAdd(&num_candidates, 1); + if (candidate_slot < params.max_candidates_per_pair) { + shared_candidate_ips[candidate_slot] = local_ip_quantized; + shared_candidate_indices[candidate_slot] = vec_idx; + } + } + } + + __syncthreads(); + + if (num_candidates > 0) { + for (size_t i = tid; i < params.D; i += num_threads) { + shared_query[i] = params.d_query[query_idx * params.D + i]; + } + __syncthreads(); + + bitwise_block_sort_emit_topk(params, + num_candidates, + query_idx, + cluster_idx, + num_vectors_in_cluster, + cluster_start_index, + q_g_add, + q_k1xsumq, + threshold); + } +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in new file mode 100644 index 0000000000..7bbdc24139 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +extern "C" __global__ void compute_inner_products_with_bitwise_block_sort( + const ComputeInnerProductsKernelParams params) +{ + compute_inner_products_with_bitwise_block_sort_impl(params); +} + +static_assert(std::is_same_v); + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_matrix.json new file mode 100644 index 0000000000..0967ef424b --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_matrix.json @@ -0,0 +1 @@ +{} 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 new file mode 100644 index 0000000000..dbe6ec5fd4 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp @@ -0,0 +1,46 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +struct ComputeInnerProductsWithBitwiseBlockSortPlanner : AlgorithmPlanner { + inline static LauncherJitCache launcher_jit_cache{}; + + ComputeInnerProductsWithBitwiseBlockSortPlanner() + : AlgorithmPlanner("compute_inner_products_with_bitwise_block_sort", launcher_jit_cache) + { + } + + void add_entrypoint() + { + this->add_static_fragment(); + } + + template + void add_bitwise_block_sort_emit_topk_device_function() + { + this->add_static_fragment>(); + } + + template + void add_extract_code_device_function() + { + this->add_static_fragment>(); + } + + template + void add_compute_bitwise_quantized_ip_for_vec_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_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_impl.cuh new file mode 100644 index 0000000000..4ff53b067c --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_impl.cuh @@ -0,0 +1,53 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "../gpu_index/searcher_gpu_common.cuh" +#include "device_functions.cuh" + +#include + +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// Unified kernel without BlockSort. The WithEx-axis-specific tail (IP2 +// precompute + per-vector distance write for WithEx=true, simple distance +// write for WithEx=false) lives in the bitwise_emit_distances JIT-LTO +// fragment, so this kernel is emitted as a single fragment regardless of +// WithEx. +__device__ void compute_inner_products_with_bitwise_impl( + const ComputeInnerProductsKernelParams params) +{ + const int block_id = blockIdx.x; + if (block_id >= params.num_pairs) return; + + ClusterQueryPair pair = params.d_sorted_pairs[block_id]; + int cluster_idx = pair.cluster_idx; + int query_idx = pair.query_idx; + + if (cluster_idx >= params.num_centroids || query_idx >= params.num_queries) return; + + size_t num_vectors_in_cluster = params.d_cluster_meta[cluster_idx].num; + size_t cluster_start_index = params.d_cluster_meta[cluster_idx].start_index; + + extern __shared__ __align__(256) char shared_mem_raw[]; + const int tid = threadIdx.x; + const int num_threads = blockDim.x; + + float* shared_query = reinterpret_cast(shared_mem_raw); + for (size_t i = tid; i < params.D; i += num_threads) { + shared_query[i] = params.d_query[query_idx * params.D + i]; + } + __syncthreads(); + + float q_g_add = params.d_centroid_distances[query_idx * params.num_centroids + cluster_idx]; + + bitwise_emit_distances( + params, query_idx, cluster_idx, num_vectors_in_cluster, cluster_start_index, q_g_add); +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_kernel.cu.in new file mode 100644 index 0000000000..1f7355c7fe --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_kernel.cu.in @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +extern "C" __global__ void compute_inner_products_with_bitwise( + const ComputeInnerProductsKernelParams params) +{ + compute_inner_products_with_bitwise_impl(params); +} + +static_assert(std::is_same_v); + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_matrix.json new file mode 100644 index 0000000000..0967ef424b --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_matrix.json @@ -0,0 +1 @@ +{} 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 new file mode 100644 index 0000000000..3ec4809395 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp @@ -0,0 +1,40 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +struct ComputeInnerProductsWithBitwisePlanner : AlgorithmPlanner { + inline static LauncherJitCache launcher_jit_cache{}; + + ComputeInnerProductsWithBitwisePlanner() + : AlgorithmPlanner("compute_inner_products_with_bitwise", launcher_jit_cache) + { + } + + void add_entrypoint() + { + this->add_static_fragment(); + } + + template + void add_bitwise_emit_distances_device_function() + { + this->add_static_fragment>(); + } + + template + void add_extract_code_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 new file mode 100644 index 0000000000..45c8793b59 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh @@ -0,0 +1,229 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "../../detail/smem_utils.cuh" +#include "../gpu_index/searcher_gpu_common.cuh" +#include "block_sort.cuh" +#include "device_functions.cuh" + +#include + +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +template +__device__ void compute_inner_products_with_lut16_opt_block_sort_impl( + const ComputeInnerProductsKernelParams params) +{ + const int block_id = blockIdx.x; + if (block_id >= params.num_pairs) return; + + ClusterQueryPair pair = params.d_sorted_pairs[block_id]; + int cluster_idx = pair.cluster_idx; + int query_idx = pair.query_idx; + + if (cluster_idx >= params.num_centroids || query_idx >= params.num_queries) return; + + size_t num_vectors_in_cluster = params.d_cluster_meta[cluster_idx].num; + size_t cluster_start_index = params.d_cluster_meta[cluster_idx].start_index; + + const uint32_t num_chunks = params.D / BITS_PER_CHUNK; + const uint32_t lut_per_query_size = num_chunks * LUT_SIZE; + + extern __shared__ __align__(256) char shared_mem_raw[]; + lut_dtype* shared_lut_fp16 = reinterpret_cast(shared_mem_raw); + uint32_t lut_bytes = num_chunks * LUT_SIZE * sizeof(lut_dtype); + + const int tid = threadIdx.x; + const int num_threads = blockDim.x; + + lut_dtype* query_lut = params.d_lut_for_queries_half + query_idx * lut_per_query_size; + for (uint32_t i = tid; i < lut_per_query_size; i += num_threads) { + shared_lut_fp16[i] = query_lut[i]; + } + __syncthreads(); + + __shared__ int num_candidates; + float q_g_add = params.d_centroid_distances[query_idx * params.num_centroids + cluster_idx]; + float q_k1xsumq = params.d_G_k1xSumq[query_idx]; + float threshold = params.d_threshold[query_idx]; + + float q_g_error; + if constexpr (WithEx) { q_g_error = sqrtf(q_g_add); } + + if (tid == 0) { num_candidates = 0; } + __syncthreads(); + + // WithEx: candidates placed at max(lut_bytes, max_candidates_per_pair * sizeof(float)) + // NoEX: candidates placed right after LUT + float* shared_candidate_ips; + if constexpr (WithEx) { + if (lut_bytes < params.max_candidates_per_pair * sizeof(float)) { + shared_candidate_ips = + reinterpret_cast(shared_mem_raw + params.max_candidates_per_pair * sizeof(float)); + } else { + shared_candidate_ips = reinterpret_cast(shared_mem_raw + lut_bytes); + } + } else { + shared_candidate_ips = reinterpret_cast(shared_mem_raw + lut_bytes); + } + // WithEx: stores local vec_idx; NoEX: stores PID + int* shared_candidate_indices = (int*)(shared_candidate_ips + params.max_candidates_per_pair); + + const uint32_t short_code_length = params.D / 32; + + for (size_t vec_base = 0; vec_base < num_vectors_in_cluster; vec_base += num_threads) { + size_t vec_idx = vec_base + tid; + + float local_ip = 0.0f; + bool is_candidate = false; + + if (vec_idx < num_vectors_in_cluster) { + 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; + + float ip = compute_lut_ip_for_vec<__half>(params.d_short_data, + shared_lut_fp16, + cluster_start_index, + num_vectors_in_cluster, + vec_idx, + short_code_length); + + float est_dist = f_add + q_g_add + f_rescale * (ip + q_k1xsumq); + + if constexpr (WithEx) { + float f_error = factors.z; + float low_dist = est_dist - f_error * q_g_error; + if (low_dist < threshold) { + is_candidate = true; + local_ip = ip; + } + } else { + if (est_dist < threshold) { + is_candidate = true; + local_ip = est_dist; + } + } + } + + __syncwarp(); + if (is_candidate) { + int candidate_slot = atomicAdd(&num_candidates, 1); + if (candidate_slot < params.max_candidates_per_pair) { + shared_candidate_ips[candidate_slot] = local_ip; + if constexpr (WithEx) { + shared_candidate_indices[candidate_slot] = (int)vec_idx; + } else { + shared_candidate_indices[candidate_slot] = + (int)params.d_pids[cluster_start_index + vec_idx]; + } + } + } + } + __syncthreads(); + if (num_candidates > 0) { + __shared__ int probe_slot; + uint32_t output_offset; + { + using block_sort_t = typename rabitq_block_sort::type; + block_sort_t queue(params.topk); + + if constexpr (WithEx) { + float q_kbxsumq = params.d_G_kbxSumq[query_idx]; + const uint32_t long_code_size = (params.D * params.ex_bits + 7) / 8; + + float* shared_query = (float*)(shared_candidate_indices + params.max_candidates_per_pair); + for (uint32_t i = tid; i < params.D; i += num_threads) { + shared_query[i] = params.d_query[query_idx * params.D + i]; + } + __syncthreads(); + + float* shared_ip2_results = reinterpret_cast(shared_lut_fp16); + const int warp_id = tid / raft::WarpSize; + const int lane_id = tid % raft::WarpSize; + const int num_warps = num_threads / raft::WarpSize; + + for (int cand_idx = warp_id; cand_idx < num_candidates; cand_idx += num_warps) { + size_t global_vec_idx = cluster_start_index + shared_candidate_indices[cand_idx]; + const uint8_t* vec_long_code = params.d_long_code + global_vec_idx * long_code_size; + float ip2 = + compute_ip2_from_long_codes_warp(vec_long_code, shared_query, params.D, lane_id); + if (lane_id == 0) { shared_ip2_results[cand_idx] = ip2; } + } + __syncthreads(); + + const int adds_per_thread = (num_candidates + num_threads - 1) / num_threads; + for (int round = 0; round < adds_per_thread; round++) { + int cand_idx = tid + round * num_threads; + float ex_dist; + uint32_t pid; + if (cand_idx < num_candidates) { + float ip = shared_candidate_ips[cand_idx]; + float ip2 = shared_ip2_results[cand_idx]; + int local_vec_idx = shared_candidate_indices[cand_idx]; + size_t global_vec_idx = cluster_start_index + local_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; + ex_dist = + f_ex_add + q_g_add + + f_ex_rescale * (static_cast(1 << params.ex_bits) * ip + ip2 + q_kbxsumq); + pid = (uint32_t)params.d_pids[global_vec_idx]; + } else { + ex_dist = INFINITY; + pid = 0; + } + queue.add(ex_dist, pid); + } + } else { + const int candidates_per_thread = (num_candidates + num_threads - 1) / num_threads; + for (int c = 0; c < candidates_per_thread; ++c) { + int cand_idx = tid + c * num_threads; + float dist; + uint32_t pid; + if (cand_idx < num_candidates) { + dist = shared_candidate_ips[cand_idx]; + pid = (uint32_t)shared_candidate_indices[cand_idx]; + } else { + dist = INFINITY; + pid = 0; + } + queue.add(dist, pid); + } + } + + __syncthreads(); + + queue.done((uint8_t*)shared_lut_fp16); + + if (tid == 0) { probe_slot = atomicAdd(¶ms.d_query_write_counters[query_idx], 1); } + __syncthreads(); + + if (probe_slot >= params.nprobe) { return; } + + output_offset = query_idx * (params.topk * params.nprobe) + probe_slot * params.topk; + queue.store(params.d_topk_dists + output_offset, + (uint32_t*)(params.d_topk_pids + output_offset)); + } + + if (num_candidates >= params.topk) { + update_threshold_atomicmin(params.d_topk_dists, + params.d_threshold, + query_idx, + params.topk, + params.nprobe, + probe_slot, + threshold, + tid); + } + } +} + +} // 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_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_kernel.cu.in new file mode 100644 index 0000000000..8e45f99c0f --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_kernel.cu.in @@ -0,0 +1,26 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +namespace { + +constexpr bool with_ex = @with_ex_value@; + +} // namespace + +namespace cuvs::neighbors::ivf_rabitq::detail { + +extern "C" __global__ void compute_inner_products_with_lut16_opt_block_sort( + const ComputeInnerProductsKernelParams params) +{ + compute_inner_products_with_lut16_opt_block_sort_impl(params); +} + +static_assert(std::is_same_v); + +} // 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_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_matrix.json new file mode 100644 index 0000000000..7e0f0f3603 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_matrix.json @@ -0,0 +1,12 @@ +{ + "_with_ex": [ + { + "with_ex_descriptor": "with_ex", + "with_ex_value": "true" + }, + { + "with_ex_descriptor": "no_ex", + "with_ex_value": "false" + } + ] +} 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 new file mode 100644 index 0000000000..acf4e27e3a --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp @@ -0,0 +1,41 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +struct ComputeInnerProductsWithLut16OptBlockSortPlanner : AlgorithmPlanner { + inline static LauncherJitCache launcher_jit_cache{}; + + ComputeInnerProductsWithLut16OptBlockSortPlanner() + : AlgorithmPlanner("compute_inner_products_with_lut16_opt_block_sort", launcher_jit_cache) + { + } + + template + void add_entrypoint() + { + this->add_static_fragment< + fragment_tag_compute_inner_products_with_lut16_opt_block_sort>(); + } + + template + void add_extract_code_device_function() + { + this->add_static_fragment>(); + } + + void add_compute_lut_ip_for_vec_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_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_impl.cuh new file mode 100644 index 0000000000..fc88d79fc4 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_impl.cuh @@ -0,0 +1,42 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "../gpu_index/searcher_gpu_common.cuh" +#include "device_functions.cuh" + +#include + +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// Unified non-BlockSort half-precision LUT kernel. The WithEx-axis-specific +// body (IP2 precompute + per-vector distance write for WithEx=true, FP16-LUT +// distance write for WithEx=false) lives in the lut16_opt_emit_distances +// JIT-LTO fragment, so this kernel is emitted as a single fragment regardless +// of WithEx. +__device__ void compute_inner_products_with_lut16_opt_impl( + const ComputeInnerProductsKernelParams params) +{ + const int block_id = blockIdx.x; + if (block_id >= params.num_pairs) return; + + ClusterQueryPair pair = params.d_sorted_pairs[block_id]; + int cluster_idx = pair.cluster_idx; + int query_idx = pair.query_idx; + + if (cluster_idx >= params.num_centroids || query_idx >= params.num_queries) return; + + size_t num_vectors_in_cluster = params.d_cluster_meta[cluster_idx].num; + size_t cluster_start_index = params.d_cluster_meta[cluster_idx].start_index; + float q_g_add = params.d_centroid_distances[query_idx * params.num_centroids + cluster_idx]; + + lut16_opt_emit_distances( + params, query_idx, cluster_idx, num_vectors_in_cluster, cluster_start_index, q_g_add); +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_kernel.cu.in new file mode 100644 index 0000000000..4160aa7a33 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_kernel.cu.in @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +extern "C" __global__ void compute_inner_products_with_lut16_opt( + const ComputeInnerProductsKernelParams params) +{ + compute_inner_products_with_lut16_opt_impl(params); +} + +static_assert(std::is_same_v); + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_matrix.json new file mode 100644 index 0000000000..0967ef424b --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_matrix.json @@ -0,0 +1 @@ +{} 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 new file mode 100644 index 0000000000..95e0694c9c --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp @@ -0,0 +1,45 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +struct ComputeInnerProductsWithLut16OptPlanner : AlgorithmPlanner { + inline static LauncherJitCache launcher_jit_cache{}; + + ComputeInnerProductsWithLut16OptPlanner() + : AlgorithmPlanner("compute_inner_products_with_lut16_opt", launcher_jit_cache) + { + } + + void add_entrypoint() + { + this->add_static_fragment(); + } + + template + void add_lut16_opt_emit_distances_device_function() + { + this->add_static_fragment>(); + } + + template + void add_extract_code_device_function() + { + this->add_static_fragment>(); + } + + void add_compute_lut_ip_for_vec_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_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_impl.cuh new file mode 100644 index 0000000000..615824621f --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_impl.cuh @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "../../detail/smem_utils.cuh" +#include "../gpu_index/searcher_gpu_common.cuh" +#include "block_sort.cuh" +#include "device_functions.cuh" + +#include + +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// Unified BlockSort LUT kernel. The WithEx-axis-specific body (LUT filter + +// IP2-refined block sort for WithEx=true, LUT filter + direct block sort for +// WithEx=false) lives in the lut_block_sort_emit_topk JIT-LTO fragment, so +// this kernel is emitted as a single fragment regardless of WithEx. +__device__ void compute_inner_products_with_lut_block_sort_impl( + const ComputeInnerProductsKernelParams params) +{ + const int block_id = blockIdx.x; + if (block_id >= params.num_pairs) return; + + ClusterQueryPair pair = params.d_sorted_pairs[block_id]; + int cluster_idx = pair.cluster_idx; + int query_idx = pair.query_idx; + + if (cluster_idx >= params.num_centroids || query_idx >= params.num_queries) return; + + const size_t num_chunks = params.D / BITS_PER_CHUNK; + const size_t lut_per_query_size = num_chunks * LUT_SIZE; + + extern __shared__ __align__(256) char shared_mem_raw[]; + float* shared_lut = reinterpret_cast(shared_mem_raw); + + const int tid = threadIdx.x; + const int num_threads = blockDim.x; + + float* query_lut = params.d_lut_for_queries_float + query_idx * lut_per_query_size; + for (size_t i = tid; i < lut_per_query_size; i += num_threads) { + shared_lut[i] = query_lut[i]; + } + __syncthreads(); + + size_t num_vectors_in_cluster = params.d_cluster_meta[cluster_idx].num; + size_t cluster_start_index = params.d_cluster_meta[cluster_idx].start_index; + float q_g_add = params.d_centroid_distances[query_idx * params.num_centroids + cluster_idx]; + float q_k1xsumq = params.d_G_k1xSumq[query_idx]; + float threshold = params.d_threshold[query_idx]; + + lut_block_sort_emit_topk(params, + query_idx, + cluster_idx, + num_vectors_in_cluster, + cluster_start_index, + q_g_add, + q_k1xsumq, + threshold); +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_kernel.cu.in new file mode 100644 index 0000000000..014cd12ac2 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_kernel.cu.in @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +extern "C" __global__ void compute_inner_products_with_lut_block_sort( + const ComputeInnerProductsKernelParams params) +{ + compute_inner_products_with_lut_block_sort_impl(params); +} + +static_assert(std::is_same_v); + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_matrix.json new file mode 100644 index 0000000000..0967ef424b --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_matrix.json @@ -0,0 +1 @@ +{} 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 new file mode 100644 index 0000000000..9db3839880 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp @@ -0,0 +1,45 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +struct ComputeInnerProductsWithLutBlockSortPlanner : AlgorithmPlanner { + inline static LauncherJitCache launcher_jit_cache{}; + + ComputeInnerProductsWithLutBlockSortPlanner() + : AlgorithmPlanner("compute_inner_products_with_lut_block_sort", launcher_jit_cache) + { + } + + void add_entrypoint() + { + this->add_static_fragment(); + } + + template + void add_lut_block_sort_emit_topk_device_function() + { + this->add_static_fragment>(); + } + + template + void add_extract_code_device_function() + { + this->add_static_fragment>(); + } + + void add_compute_lut_ip_for_vec_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_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_impl.cuh new file mode 100644 index 0000000000..1a44453887 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_impl.cuh @@ -0,0 +1,40 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "../gpu_index/searcher_gpu_common.cuh" +#include "device_functions.cuh" + +#include + +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// Unified non-BlockSort LUT kernel. The WithEx-axis-specific body (IP2 +// precompute + per-vector distance write for WithEx=true, LUT-only distance +// write for WithEx=false) lives in the lut_emit_distances JIT-LTO fragment, so +// this kernel is emitted as a single fragment regardless of WithEx. +__device__ void compute_inner_products_with_lut_impl(const ComputeInnerProductsKernelParams params) +{ + const int block_id = blockIdx.x; + if (block_id >= params.num_pairs) return; + + ClusterQueryPair pair = params.d_sorted_pairs[block_id]; + int cluster_idx = pair.cluster_idx; + int query_idx = pair.query_idx; + + if (cluster_idx >= params.num_centroids || query_idx >= params.num_queries) return; + + size_t num_vectors_in_cluster = params.d_cluster_meta[cluster_idx].num; + size_t cluster_start_index = params.d_cluster_meta[cluster_idx].start_index; + float q_g_add = params.d_centroid_distances[query_idx * params.num_centroids + cluster_idx]; + + lut_emit_distances( + params, query_idx, cluster_idx, num_vectors_in_cluster, cluster_start_index, q_g_add); +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_kernel.cu.in new file mode 100644 index 0000000000..7761879ce0 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_kernel.cu.in @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +extern "C" __global__ void compute_inner_products_with_lut( + const ComputeInnerProductsKernelParams params) +{ + compute_inner_products_with_lut_impl(params); +} + +static_assert(std::is_same_v); + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_matrix.json new file mode 100644 index 0000000000..0967ef424b --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_matrix.json @@ -0,0 +1 @@ +{} 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 new file mode 100644 index 0000000000..3559a9bee1 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp @@ -0,0 +1,45 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +struct ComputeInnerProductsWithLutPlanner : AlgorithmPlanner { + inline static LauncherJitCache launcher_jit_cache{}; + + ComputeInnerProductsWithLutPlanner() + : AlgorithmPlanner("compute_inner_products_with_lut", launcher_jit_cache) + { + } + + void add_entrypoint() + { + this->add_static_fragment(); + } + + template + void add_lut_emit_distances_device_function() + { + this->add_static_fragment>(); + } + + template + void add_extract_code_device_function() + { + this->add_static_fragment>(); + } + + void add_compute_lut_ip_for_vec_device_function() + { + this->add_static_fragment>(); + } +}; + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_kernel.cu.in new file mode 100644 index 0000000000..0828b0581a --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_kernel.cu.in @@ -0,0 +1,58 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include + +#include + +namespace { + +using lut_t = @lut_dtype@; + +} // namespace + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// Primary template definition. The fragment is built once per @lut_dtype@ +// value declared in compute_lut_ip_for_vec_matrix.json; the explicit +// instantiation at the bottom emits one specialization (via the lut_t alias) +// per build, and nvJitLink resolves consumer calls to the matching +// specialization. +template +__device__ float compute_lut_ip_for_vec(const uint32_t* d_short_data, + const LutT* shared_lut, + size_t cluster_start_index, + size_t num_vectors_in_cluster, + size_t vec_idx, + size_t short_code_length) +{ + constexpr int chunks_per_uint32 = 32 / BITS_PER_CHUNK; + float ip = 0.0f; + for (size_t uint32_idx = 0; uint32_idx < short_code_length; uint32_idx++) { + size_t short_code_offset = + cluster_start_index * short_code_length + uint32_idx * num_vectors_in_cluster + vec_idx; + uint32_t short_code_chunk = d_short_data[short_code_offset]; + for (int chunk_in_uint32 = 0; chunk_in_uint32 < chunks_per_uint32; chunk_in_uint32++) { + int shift = 28 - (chunk_in_uint32 * BITS_PER_CHUNK); + int pattern = (short_code_chunk >> shift) & 0xF; + size_t lut_chunk_idx = uint32_idx * chunks_per_uint32 + chunk_in_uint32; + size_t lut_offset = lut_chunk_idx * LUT_SIZE + pattern; + if constexpr (std::is_same_v) { + ip += __half2float(shared_lut[lut_offset]); + } else { + ip += shared_lut[lut_offset]; + } + } + } + return ip; +} + +template __device__ float compute_lut_ip_for_vec( + const uint32_t*, const lut_t*, size_t, size_t, size_t, size_t); + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_matrix.json new file mode 100644 index 0000000000..d53a17ca38 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_matrix.json @@ -0,0 +1,6 @@ +{ + "_lut_dtype": [ + {"lut_dtype": "float", "lut_dtype_descriptor": "f32"}, + {"lut_dtype": "__half", "lut_dtype_descriptor": "f16"} + ] +} 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 new file mode 100644 index 0000000000..7939b8ec13 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh @@ -0,0 +1,169 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "../gpu_index/searcher_gpu_common.cuh" +#include "../utils/reductions.cuh" + +#include + +#include +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// Cross-fragment device function declarations. Definitions live in their own +// JIT-LTO fragments and are resolved at nvJitLink time when a consumer planner +// adds them via add_*_device_function(). +__device__ uint32_t extract_code(const uint8_t* codes, size_t d); + +// 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. +__inline__ __device__ float compute_ip2_from_long_codes_warp(const uint8_t* vec_long_code, + const float* shared_query, + size_t D, + int lane_id) +{ + float ip2 = 0.0f; + for (size_t d = lane_id; d < D; d += raft::WarpSize) { + uint32_t code_val = extract_code(vec_long_code, d); + ip2 += shared_query[d] * static_cast(code_val); + } + return warpReduceSum(ip2); +} + +// Primary template; explicit instantiations live in the +// compute_lut_ip_for_vec_kernel.cu.in fragment, one per @lut_dtype@ substitution +// (float and __half today). Consumers call the specialization matching the +// element type of their `shared_lut`. +template +__device__ float compute_lut_ip_for_vec(const uint32_t* d_short_data, + const LutT* shared_lut, + size_t cluster_start_index, + size_t num_vectors_in_cluster, + size_t vec_idx, + size_t short_code_length); + +__inline__ __device__ float compute_bitwise_1bit_ip_for_vec(const uint32_t* d_short_data, + const float* shared_query, + size_t cluster_start_index, + size_t num_vectors_in_cluster, + size_t vec_idx, + size_t short_code_length, + size_t D) +{ + float exact_ip = 0.0f; + for (size_t uint32_idx = 0; uint32_idx < short_code_length; uint32_idx++) { + size_t short_code_offset = + cluster_start_index * short_code_length + uint32_idx * num_vectors_in_cluster + vec_idx; + uint32_t short_code_chunk = d_short_data[short_code_offset]; +#pragma unroll 8 + for (int bit_idx = 0; bit_idx < 32; bit_idx++) { + size_t dim = uint32_idx * 32 + bit_idx; + if (dim < D) { + if ((short_code_chunk >> (31 - bit_idx)) & 0x1) { exact_ip += shared_query[dim]; } + } + } + } + return exact_ip; +} + +// Note: contains __syncthreads(); must be called by all threads of the block. +__inline__ __device__ void update_threshold_atomicmin(const float* d_topk_dists, + const float* d_threshold, + uint32_t query_idx, + uint32_t topk, + uint32_t nprobe, + int probe_slot, + float threshold, + int tid) +{ + float max_topk_dist; + if (tid == 0) { + max_topk_dist = -INFINITY; + uint32_t output_offset = query_idx * (topk * nprobe) + probe_slot * topk; + for (uint32_t i = 0; i < topk; i++) { + float dist = d_topk_dists[output_offset + i]; + if (dist > 0 && dist > max_topk_dist && dist < INFINITY) { max_topk_dist = dist; } + } + } + __syncthreads(); + if (tid == 0 && max_topk_dist > 0 && max_topk_dist < threshold) { + int* threshold_ptr = (int*)(d_threshold + query_idx); + atomicMin(threshold_ptr, __float_as_int(max_topk_dist)); + } +} + +__device__ int32_t compute_bitwise_quantized_ip_for_vec(const uint32_t* d_short_data, + const uint32_t* shared_packed_query, + size_t cluster_start_index, + size_t num_vectors_in_cluster, + size_t vec_idx, + uint32_t num_words); + +// WithEx-axis-specific final step of the bitwise (non-BlockSort) kernel: +// optionally precomputes IP2 in shared memory and writes per-vector distances +// to the output. Two fragments (with_ex / no_ex) define this symbol; nvJitLink +// resolves the call. Pre-computed state passed in to avoid redundant global +// loads and give the JIT-LTO inliner direct visibility into the helper inputs. +__device__ void bitwise_emit_distances(const ComputeInnerProductsKernelParams params, + int query_idx, + int cluster_idx, + size_t num_vectors_in_cluster, + size_t cluster_start_index, + float q_g_add); + +// WithEx-axis-specific final step of the LUT (non-BlockSort) kernel: computes +// and writes per-vector distances, with optional IP2 refinement when WithEx. +__device__ void lut_emit_distances(const ComputeInnerProductsKernelParams params, + int query_idx, + int cluster_idx, + size_t num_vectors_in_cluster, + size_t cluster_start_index, + float q_g_add); + +// WithEx-axis-specific final step of the half-precision LUT (non-BlockSort) +// kernel: computes and writes per-vector distances, with optional IP2 +// refinement when WithEx. +__device__ void lut16_opt_emit_distances(const ComputeInnerProductsKernelParams params, + int query_idx, + int cluster_idx, + size_t num_vectors_in_cluster, + size_t cluster_start_index, + float q_g_add); + +// WithEx-axis-specific final step of the LUT BlockSort kernel: filters +// candidates via the LUT and builds the per-block top-k (with IP2 refinement +// when WithEx). Two fragments (with_ex / no_ex) define this symbol. +__device__ void lut_block_sort_emit_topk(const ComputeInnerProductsKernelParams params, + int query_idx, + int cluster_idx, + size_t num_vectors_in_cluster, + size_t cluster_start_index, + float q_g_add, + float q_k1xsumq, + float threshold); + +// WithEx-axis-specific final step of the bitwise BlockSort kernel: builds and +// stores the per-block top-k for the current cluster/query pair. Two fragments +// (with_ex / no_ex) define this symbol; nvJitLink resolves the call. State the +// kernel already computed for Phase 1 is passed in to avoid redundant global +// loads and to give the JIT-LTO inliner the same view as the previous +// templated kernel. +__device__ void bitwise_block_sort_emit_topk(const ComputeInnerProductsKernelParams params, + int num_candidates, + int query_idx, + int cluster_idx, + size_t num_vectors_in_cluster, + size_t cluster_start_index, + float q_g_add, + float q_k1xsumq, + float threshold); + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_kernel.cu.in new file mode 100644 index 0000000000..c2b7c21726 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_kernel.cu.in @@ -0,0 +1,43 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +namespace { + +constexpr int ex_bits = @ex_bits@; + +} // namespace + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// Non-inline __device__ definition. ex_bits is captured from the matrix.json +// axis into a constexpr so the compiler can constant-fold the bit-extraction +// arithmetic, eliminate the spanning-byte branch when impossible (ex_bits in +// {1, 2, 4, 8}), and collapse the function to a single byte load for +// ex_bits=8. Each matrix.json value produces a separate fragment with the +// same mangled C++ symbol; the planner adds exactly one via +// add_extract_code_device_function(), so nvJitLink resolves consumer +// calls to the matching specialization with no runtime dispatch. +__device__ uint32_t extract_code(const uint8_t* codes, size_t d) +{ + if constexpr (ex_bits == 8) { + return codes[d]; + } else if constexpr (ex_bits == 4) { + return (codes[d >> 1] >> (4 * (1 - (d & 1)))) & 0xF; + } else { + size_t bitPos = d * ex_bits; + size_t byteIdx = bitPos >> 3; + size_t bitOffset = bitPos & 7; + uint32_t v = codes[byteIdx] << 8; + if constexpr (ex_bits != 1 && ex_bits != 2) { + if (bitOffset + ex_bits > 8) { v |= codes[byteIdx + 1]; } + } + int shift = 16 - (bitOffset + ex_bits); + return (v >> shift) & ((1u << ex_bits) - 1); + } +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_matrix.json new file mode 100644 index 0000000000..f377fa041c --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_matrix.json @@ -0,0 +1,3 @@ +{ + "ex_bits": ["1", "2", "3", "4", "5", "6", "7", "8"] +} diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/kernel_def.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/kernel_def.hpp new file mode 100644 index 0000000000..2314d20b62 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/kernel_def.hpp @@ -0,0 +1,14 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "../gpu_index/searcher_gpu_common.cuh" + +namespace cuvs::neighbors::ivf_rabitq::detail { + +using compute_inner_products_with_lut_func_t = void(const ComputeInnerProductsKernelParams); + +} // namespace cuvs::neighbors::ivf_rabitq::detail 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 new file mode 100644 index 0000000000..d9c94821f1 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp @@ -0,0 +1,136 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#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" +#include "compute_inner_products_with_lut16_opt_planner.hpp" +#include "compute_inner_products_with_lut_block_sort_planner.hpp" +#include "compute_inner_products_with_lut_planner.hpp" + +#include + +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +namespace { + +template +inline void add_ex_bits_device_functions(Planner& planner, int ex_bits) +{ + switch (ex_bits) { + case 1: planner.template add_extract_code_device_function<1>(); break; + case 2: planner.template add_extract_code_device_function<2>(); break; + case 3: planner.template add_extract_code_device_function<3>(); break; + case 4: planner.template add_extract_code_device_function<4>(); break; + case 5: planner.template add_extract_code_device_function<5>(); break; + case 6: planner.template add_extract_code_device_function<6>(); break; + case 7: planner.template add_extract_code_device_function<7>(); break; + case 8: planner.template add_extract_code_device_function<8>(); break; + default: assert(false); + } +} + +} // namespace + +inline std::shared_ptr make_compute_inner_products_with_lut_launcher( + int ex_bits, bool with_ex) +{ + ComputeInnerProductsWithLutPlanner planner; + planner.add_entrypoint(); + planner.add_compute_lut_ip_for_vec_device_function(); + if (with_ex) { + planner.add_lut_emit_distances_device_function(); + add_ex_bits_device_functions(planner, ex_bits); + } else { + planner.add_lut_emit_distances_device_function(); + } + return planner.get_launcher(); +} + +inline std::shared_ptr make_compute_inner_products_with_lut_block_sort_launcher( + int ex_bits, bool with_ex) +{ + ComputeInnerProductsWithLutBlockSortPlanner planner; + planner.add_entrypoint(); + planner.add_compute_lut_ip_for_vec_device_function(); + if (with_ex) { + planner.add_lut_block_sort_emit_topk_device_function(); + add_ex_bits_device_functions(planner, ex_bits); + } else { + planner.add_lut_block_sort_emit_topk_device_function(); + } + return planner.get_launcher(); +} + +inline std::shared_ptr make_compute_inner_products_with_lut16_opt_launcher( + int ex_bits, bool with_ex) +{ + ComputeInnerProductsWithLut16OptPlanner planner; + planner.add_entrypoint(); + planner.add_compute_lut_ip_for_vec_device_function(); + if (with_ex) { + planner.add_lut16_opt_emit_distances_device_function(); + add_ex_bits_device_functions(planner, ex_bits); + } else { + planner.add_lut16_opt_emit_distances_device_function(); + } + return planner.get_launcher(); +} + +inline std::shared_ptr +make_compute_inner_products_with_lut16_opt_block_sort_launcher(int ex_bits, bool with_ex) +{ + ComputeInnerProductsWithLut16OptBlockSortPlanner planner; + planner.add_compute_lut_ip_for_vec_device_function(); + if (with_ex) { + planner.add_entrypoint(); + add_ex_bits_device_functions(planner, ex_bits); + } else { + planner.add_entrypoint(); + } + return planner.get_launcher(); +} + +inline std::shared_ptr make_compute_inner_products_with_bitwise_launcher( + int ex_bits, bool with_ex) +{ + ComputeInnerProductsWithBitwisePlanner planner; + planner.add_entrypoint(); + if (with_ex) { + planner.add_bitwise_emit_distances_device_function(); + add_ex_bits_device_functions(planner, ex_bits); + } else { + planner.add_bitwise_emit_distances_device_function(); + } + return planner.get_launcher(); +} + +inline std::shared_ptr +make_compute_inner_products_with_bitwise_block_sort_launcher(int num_bits, + int ex_bits, + bool with_ex) +{ + ComputeInnerProductsWithBitwiseBlockSortPlanner planner; + planner.add_entrypoint(); + if (num_bits == 4) { + planner.add_compute_bitwise_quantized_ip_for_vec_device_function<4>(); + } else { + planner.add_compute_bitwise_quantized_ip_for_vec_device_function<8>(); + } + if (with_ex) { + planner.add_bitwise_block_sort_emit_topk_device_function(); + add_ex_bits_device_functions(planner, ex_bits); + } else { + planner.add_bitwise_block_sort_emit_topk_device_function(); + } + return planner.get_launcher(); +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail 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 new file mode 100644 index 0000000000..24aee45e03 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in @@ -0,0 +1,145 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include + +namespace { + +constexpr bool with_ex = @with_ex_value@; + +} // namespace + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// Computes and writes per-vector distances for the half-precision LUT +// (non-BlockSort) kernel. WithEx precomputes warp-level IP2 from long codes; +// NoEx uses the FP16 LUT only. Lifted into its own JIT-LTO fragment so the +// kernel emits a single fragment regardless of WithEx, and nvJitLink resolves +// the matching body at link time. +__device__ void lut16_opt_emit_distances(const ComputeInnerProductsKernelParams params, + int query_idx, + int cluster_idx, + size_t num_vectors_in_cluster, + size_t cluster_start_index, + float q_g_add) +{ + const uint32_t num_chunks = params.D / BITS_PER_CHUNK; + const uint32_t lut_per_query_size = num_chunks * LUT_SIZE; + + extern __shared__ __align__(256) char shared_mem_raw[]; + + const int tid = threadIdx.x; + const int num_threads = blockDim.x; + + if constexpr (with_ex) { + // Shared memory layout: [ip2_results (max_candidates_per_pair floats)][lut_fp16 / query] + float* shared_ip2_results = reinterpret_cast(shared_mem_raw); + lut_dtype* shared_lut_fp16 = + reinterpret_cast(shared_ip2_results + params.max_candidates_per_pair); + + const uint32_t long_code_size = (params.D * params.ex_bits + 7) / 8; + + // Load query into LUT region (reused; D floats always fit in lut_per_query_size lut_dtype + // slots) + float* shared_query = reinterpret_cast(shared_lut_fp16); + for (uint32_t i = tid; i < params.D; i += num_threads) { + shared_query[i] = params.d_query[query_idx * params.D + i]; + } + __syncthreads(); + + const int warp_id = tid / raft::WarpSize; + const int lane_id = tid % raft::WarpSize; + const int num_warps = num_threads / raft::WarpSize; + + for (int cand_idx = warp_id; cand_idx < num_vectors_in_cluster; cand_idx += num_warps) { + size_t global_vec_idx = cluster_start_index + cand_idx; + const uint8_t* vec_long_code = params.d_long_code + global_vec_idx * long_code_size; + + float ip2 = compute_ip2_from_long_codes_warp(vec_long_code, shared_query, params.D, lane_id); + if (lane_id == 0) { shared_ip2_results[cand_idx] = ip2; } + } + __syncthreads(); + + // Load LUT (overwrites query) + lut_dtype* query_lut = params.d_lut_for_queries_half + query_idx * lut_per_query_size; + for (uint32_t i = tid; i < lut_per_query_size; i += num_threads) { + shared_lut_fp16[i] = query_lut[i]; + } + + const uint32_t short_code_length = params.D / 32; + + __shared__ int probe_slot; + if (tid == 0) { + probe_slot = atomicAdd(¶ms.d_query_write_counters[query_idx], num_vectors_in_cluster); + } + __syncthreads(); + uint32_t output_offset = query_idx * params.max_candidates_per_query + probe_slot; + + 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) { + float ip = compute_lut_ip_for_vec<__half>(params.d_short_data, + shared_lut_fp16, + cluster_start_index, + num_vectors_in_cluster, + vec_idx, + short_code_length); + + 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; + + params.d_topk_dists[output_offset + vec_idx] = + f_ex_add + q_g_add + + f_ex_rescale * (static_cast(1 << params.ex_bits) * ip + ip2 + q_kbxsumq); + params.d_topk_pids[output_offset + vec_idx] = params.d_pids[global_vec_idx]; + } + } else { + // Shared memory layout: [lut_fp16 (lut_per_query_size lut_dtype elements)] + lut_dtype* shared_lut_fp16 = reinterpret_cast(shared_mem_raw); + + lut_dtype* query_lut = params.d_lut_for_queries_half + query_idx * lut_per_query_size; + for (uint32_t i = tid; i < lut_per_query_size; i += num_threads) { + shared_lut_fp16[i] = query_lut[i]; + } + __syncthreads(); + + float q_k1xsumq = params.d_G_k1xSumq[query_idx]; + + const uint32_t short_code_length = params.D / 32; + + __shared__ int probe_slot; + if (tid == 0) { + probe_slot = atomicAdd(¶ms.d_query_write_counters[query_idx], num_vectors_in_cluster); + } + __syncthreads(); + 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; + + float ip = compute_lut_ip_for_vec<__half>(params.d_short_data, + shared_lut_fp16, + cluster_start_index, + num_vectors_in_cluster, + vec_idx, + short_code_length); + + params.d_topk_dists[output_offset + vec_idx] = f_add + q_g_add + f_rescale * (ip + q_k1xsumq); + params.d_topk_pids[output_offset + vec_idx] = params.d_pids[cluster_start_index + vec_idx]; + } + } +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_matrix.json new file mode 100644 index 0000000000..7e0f0f3603 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_matrix.json @@ -0,0 +1,12 @@ +{ + "_with_ex": [ + { + "with_ex_descriptor": "with_ex", + "with_ex_value": "true" + }, + { + "with_ex_descriptor": "no_ex", + "with_ex_value": "false" + } + ] +} 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 new file mode 100644 index 0000000000..5828027b6d --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in @@ -0,0 +1,286 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include + +namespace { + +constexpr bool with_ex = @with_ex_value@; + +} // namespace + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// Builds and stores the per-block top-k for the LUT BlockSort kernel. +// WithEx filters via the LUT with an error bound and refines via IP2 from long +// codes; NoEx filters via the LUT and adds 1-bit distances directly. Lifted +// into its own JIT-LTO fragment so the kernel emits a single fragment +// regardless of WithEx, and nvJitLink resolves the matching body at link time. +// Must be called by all threads of the block, with `shared_lut` already +// populated. +__device__ void lut_block_sort_emit_topk(const ComputeInnerProductsKernelParams params, + int query_idx, + int cluster_idx, + size_t num_vectors_in_cluster, + size_t cluster_start_index, + float q_g_add, + float q_k1xsumq, + float threshold) +{ + const size_t num_chunks = params.D / BITS_PER_CHUNK; + const size_t lut_per_query_size = num_chunks * LUT_SIZE; + const size_t lut_bytes = lut_per_query_size * sizeof(float); + + extern __shared__ __align__(256) char shared_mem_raw[]; + float* shared_lut = reinterpret_cast(shared_mem_raw); + + const int tid = threadIdx.x; + const int num_threads = blockDim.x; + + __shared__ int num_candidates; + if (tid == 0) { num_candidates = 0; } + __syncthreads(); + + const size_t short_code_length = params.D / 32; + const int vectors_per_iteration = num_threads; + + if constexpr (with_ex) { + float q_g_error = sqrtf(q_g_add); + + // Shared memory layout after LUT: [cand_dists][cand_ips][cand_indices][queue_buffer] + float* shared_candidate_dists = reinterpret_cast(shared_mem_raw + lut_bytes); + float* shared_candidate_ips = shared_candidate_dists + params.max_candidates_per_pair; + int* shared_candidate_indices = + reinterpret_cast(shared_candidate_ips + params.max_candidates_per_pair); + int* shared_buffer = shared_candidate_indices + params.max_candidates_per_pair; + + for (size_t vec_base = 0; vec_base < num_vectors_in_cluster; + vec_base += vectors_per_iteration) { + size_t vec_idx = vec_base + tid; + + float local_low_dist = INFINITY; + float local_ip = 0.0f; + bool is_candidate = false; + + if (vec_idx < num_vectors_in_cluster) { + 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; + float f_error = factors.z; + + float ip = compute_lut_ip_for_vec(params.d_short_data, + shared_lut, + cluster_start_index, + num_vectors_in_cluster, + vec_idx, + short_code_length); + + float est_dist = f_add + q_g_add + f_rescale * (ip + q_k1xsumq); + float low_dist = est_dist - f_error * q_g_error; + + if (low_dist < threshold) { + is_candidate = true; + local_low_dist = est_dist; + local_ip = ip; + } + } + __syncwarp(); + + if (is_candidate) { + int candidate_slot = atomicAdd(&num_candidates, 1); + if (candidate_slot < params.max_candidates_per_pair) { + shared_candidate_dists[candidate_slot] = local_low_dist; + shared_candidate_ips[candidate_slot] = local_ip; + shared_candidate_indices[candidate_slot] = vec_idx; + } + } + } + __syncthreads(); + + int final_num_candidates = min(num_candidates, (int)params.max_candidates_per_pair); + + if (final_num_candidates > 0) { + using block_sort_t = typename rabitq_block_sort::type; + block_sort_t queue(params.topk); + + float q_kbxsumq = params.d_G_kbxSumq[query_idx]; + const size_t long_code_size = (params.D * params.ex_bits + 7) / 8; + + float* shared_query = reinterpret_cast(shared_lut); + for (size_t i = tid; i < params.D; i += num_threads) { + shared_query[i] = params.d_query[query_idx * params.D + i]; + } + __syncthreads(); + + float* shared_ip2_results = shared_candidate_dists; + + const int warp_id = tid / raft::WarpSize; + const int lane_id = tid % raft::WarpSize; + const int num_warps = num_threads / raft::WarpSize; + + for (int cand_idx = warp_id; cand_idx < final_num_candidates; cand_idx += num_warps) { + int local_vec_idx = shared_candidate_indices[cand_idx]; + size_t global_vec_idx = cluster_start_index + local_vec_idx; + const uint8_t* vec_long_code = params.d_long_code + global_vec_idx * long_code_size; + + float ip2 = + compute_ip2_from_long_codes_warp(vec_long_code, shared_query, params.D, lane_id); + if (lane_id == 0) { shared_ip2_results[cand_idx] = ip2; } + } + __syncthreads(); + + const int adds_per_thread = (final_num_candidates + num_threads - 1) / num_threads; + for (int round = 0; round < adds_per_thread; round++) { + int cand_idx = tid + round * num_threads; + + float ex_dist; + uint32_t pid; + + if (cand_idx < final_num_candidates) { + float ip = shared_candidate_ips[cand_idx]; + float ip2 = shared_ip2_results[cand_idx]; + int local_vec_idx = shared_candidate_indices[cand_idx]; + size_t global_vec_idx = cluster_start_index + local_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; + + ex_dist = f_ex_add + q_g_add + + f_ex_rescale * (static_cast(1 << params.ex_bits) * ip + ip2 + q_kbxsumq); + pid = (uint32_t)params.d_pids[global_vec_idx]; + } else { + ex_dist = INFINITY; + pid = 0; + } + queue.add(ex_dist, pid); + } + __syncthreads(); + + queue.done((uint8_t*)shared_buffer); + + __shared__ int probe_slot; + if (tid == 0) { probe_slot = atomicAdd(¶ms.d_query_write_counters[query_idx], 1); } + __syncthreads(); + + if (probe_slot >= params.nprobe) { return; } + + size_t output_offset = query_idx * (params.topk * params.nprobe) + probe_slot * params.topk; + queue.store(params.d_topk_dists + output_offset, + (uint32_t*)(params.d_topk_pids + output_offset)); + + if (final_num_candidates >= params.topk) { + update_threshold_atomicmin(params.d_topk_dists, + params.d_threshold, + query_idx, + params.topk, + params.nprobe, + probe_slot, + threshold, + tid); + } + } + } else { + // Shared memory layout after LUT: [cand_ips][cand_pids] + float* shared_candidate_ips = reinterpret_cast(shared_mem_raw + lut_bytes); + int* shared_candidate_pids = + reinterpret_cast(shared_candidate_ips + params.max_candidates_per_pair); + + float final_1bit_dist; + PID final_1bit_pid; + + for (size_t vec_base = 0; vec_base < num_vectors_in_cluster; + vec_base += vectors_per_iteration) { + size_t vec_idx = vec_base + tid; + + bool is_candidate = false; + + if (vec_idx < num_vectors_in_cluster) { + 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; + + float ip = compute_lut_ip_for_vec(params.d_short_data, + shared_lut, + cluster_start_index, + num_vectors_in_cluster, + vec_idx, + short_code_length); + + final_1bit_dist = f_add + q_g_add + f_rescale * (ip + q_k1xsumq); + + if (final_1bit_dist < threshold) { + is_candidate = true; + final_1bit_pid = params.d_pids[cluster_start_index + vec_idx]; + } + } + __syncwarp(); + + if (is_candidate) { + int candidate_slot = atomicAdd(&num_candidates, 1); + if (candidate_slot < params.max_candidates_per_pair) { + shared_candidate_ips[candidate_slot] = final_1bit_dist; + shared_candidate_pids[candidate_slot] = final_1bit_pid; + } + } + } + __syncthreads(); + if (num_candidates > 0) { + __shared__ int probe_slot; + { + using block_sort_t = typename rabitq_block_sort::type; + block_sort_t queue(params.topk); + + const int candidates_per_thread = (num_candidates + num_threads - 1) / num_threads; + + for (int c = 0; c < candidates_per_thread; ++c) { + int cand_idx = tid + c * num_threads; + + if (cand_idx < num_candidates) { + final_1bit_dist = shared_candidate_ips[cand_idx]; + final_1bit_pid = shared_candidate_pids[cand_idx]; + } else { + final_1bit_dist = INFINITY; + final_1bit_pid = 0; + } + queue.add(final_1bit_dist, final_1bit_pid); + } + __syncthreads(); + + queue.done((uint8_t*)shared_lut); + + if (tid == 0) { probe_slot = atomicAdd(¶ms.d_query_write_counters[query_idx], 1); } + __syncthreads(); + + if (probe_slot >= params.nprobe) { return; } + + uint32_t output_offset = + query_idx * (params.topk * params.nprobe) + probe_slot * params.topk; + queue.store(params.d_topk_dists + output_offset, + (uint32_t*)(params.d_topk_pids + output_offset)); + } + + if (num_candidates >= params.topk) { + update_threshold_atomicmin(params.d_topk_dists, + params.d_threshold, + query_idx, + params.topk, + params.nprobe, + probe_slot, + threshold, + tid); + } + } + } +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_matrix.json new file mode 100644 index 0000000000..7e0f0f3603 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_matrix.json @@ -0,0 +1,12 @@ +{ + "_with_ex": [ + { + "with_ex_descriptor": "with_ex", + "with_ex_value": "true" + }, + { + "with_ex_descriptor": "no_ex", + "with_ex_value": "false" + } + ] +} 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 new file mode 100644 index 0000000000..28d508e977 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in @@ -0,0 +1,142 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include + +namespace { + +constexpr bool with_ex = @with_ex_value@; + +} // namespace + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// Computes and writes per-vector distances for the LUT (non-BlockSort) kernel. +// WithEx precomputes warp-level IP2 from long codes; NoEx uses the LUT only. +// Lifted into its own JIT-LTO fragment so the kernel emits a single fragment +// regardless of WithEx, and nvJitLink resolves the matching body at link time. +__device__ void lut_emit_distances(const ComputeInnerProductsKernelParams params, + int query_idx, + int cluster_idx, + size_t num_vectors_in_cluster, + size_t cluster_start_index, + float q_g_add) +{ + const size_t num_chunks = params.D / BITS_PER_CHUNK; + const size_t lut_per_query_size = num_chunks * LUT_SIZE; + + extern __shared__ __align__(256) char shared_mem_raw[]; + + const int tid = threadIdx.x; + const int num_threads = blockDim.x; + + if constexpr (with_ex) { + // Shared memory layout: [ip2_results (max_candidates_per_pair floats)][lut / query (floats)] + float* shared_ip2_results = reinterpret_cast(shared_mem_raw); + float* shared_lut = shared_ip2_results + params.max_candidates_per_pair; + + const size_t long_code_size = (params.D * params.ex_bits + 7) / 8; + + // Load query into the LUT region (reused; query fits since D ≤ lut_per_query_size) + float* shared_query = shared_lut; + for (size_t i = tid; i < params.D; i += num_threads) { + shared_query[i] = params.d_query[query_idx * params.D + i]; + } + __syncthreads(); + + const int warp_id = tid / raft::WarpSize; + const int lane_id = tid % raft::WarpSize; + const int num_warps = num_threads / raft::WarpSize; + + for (int cand_idx = warp_id; cand_idx < num_vectors_in_cluster; cand_idx += num_warps) { + size_t global_vec_idx = cluster_start_index + cand_idx; + const uint8_t* vec_long_code = params.d_long_code + global_vec_idx * long_code_size; + + float ip2 = compute_ip2_from_long_codes_warp(vec_long_code, shared_query, params.D, lane_id); + if (lane_id == 0) { shared_ip2_results[cand_idx] = ip2; } + } + __syncthreads(); + + // Load LUT into the same region (overwrites query) + float* query_lut = params.d_lut_for_queries_float + query_idx * lut_per_query_size; + for (size_t i = tid; i < lut_per_query_size; i += num_threads) { + shared_lut[i] = query_lut[i]; + } + + const size_t short_code_length = params.D / 32; + + __shared__ int probe_slot; + if (tid == 0) { + probe_slot = atomicAdd(¶ms.d_query_write_counters[query_idx], num_vectors_in_cluster); + } + __syncthreads(); + + float q_kbxsumq = params.d_G_kbxSumq[query_idx]; + 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) { + float ip = compute_lut_ip_for_vec(params.d_short_data, + shared_lut, + cluster_start_index, + num_vectors_in_cluster, + vec_idx, + short_code_length); + + 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; + + params.d_topk_dists[output_offset + vec_idx] = + f_ex_add + q_g_add + + f_ex_rescale * (static_cast(1 << params.ex_bits) * ip + ip2 + q_kbxsumq); + params.d_topk_pids[output_offset + vec_idx] = params.d_pids[global_vec_idx]; + } + } else { + // Shared memory layout: [lut (lut_per_query_size floats)] + float* shared_lut = reinterpret_cast(shared_mem_raw); + + float* query_lut = params.d_lut_for_queries_float + query_idx * lut_per_query_size; + for (size_t i = tid; i < lut_per_query_size; i += num_threads) { + shared_lut[i] = query_lut[i]; + } + __syncthreads(); + + float q_k1xsumq = params.d_G_k1xSumq[query_idx]; + + const size_t short_code_length = params.D / 32; + + __shared__ int probe_slot; + if (tid == 0) { + probe_slot = atomicAdd(¶ms.d_query_write_counters[query_idx], num_vectors_in_cluster); + } + __syncthreads(); + 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; + + float ip = compute_lut_ip_for_vec(params.d_short_data, + shared_lut, + cluster_start_index, + num_vectors_in_cluster, + vec_idx, + short_code_length); + + params.d_topk_dists[output_offset + vec_idx] = f_add + q_g_add + f_rescale * (ip + q_k1xsumq); + params.d_topk_pids[output_offset + vec_idx] = params.d_pids[cluster_start_index + vec_idx]; + } + } +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_matrix.json new file mode 100644 index 0000000000..7e0f0f3603 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_matrix.json @@ -0,0 +1,12 @@ +{ + "_with_ex": [ + { + "with_ex_descriptor": "with_ex", + "with_ex_value": "true" + }, + { + "with_ex_descriptor": "no_ex", + "with_ex_value": "false" + } + ] +} diff --git a/cpp/src/neighbors/ivf_rabitq/utils/IO.hpp b/cpp/src/neighbors/ivf_rabitq/utils/IO.hpp new file mode 100644 index 0000000000..0a6b00ba14 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/utils/IO.hpp @@ -0,0 +1,109 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include +#include + +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +size_t get_filesize(const char* filename) +{ + struct stat64 stat_buf; + int rc = stat64(filename, &stat_buf); + return rc == 0 ? stat_buf.st_size : -1; +} + +bool file_exits(const char* filename) +{ + std::ifstream f(filename); + if (!f.good()) { + f.close(); + return false; + } + f.close(); + return true; +} + +template +void load_vecs(const char* filename, M& Mat) +{ + RAFT_EXPECTS(file_exits(filename), "File %s not exists", filename); + + static_assert(std::is_same_v, "T must match M::element_type"); + + uint32_t tmp; + size_t file_size = get_filesize(filename); + std::ifstream input(filename, std::ios::binary); + + input.read((char*)&tmp, sizeof(uint32_t)); + + size_t cols = tmp; + size_t rows = file_size / (cols * sizeof(T) + sizeof(uint32_t)); + Mat = raft::make_host_matrix(rows, cols); + + input.seekg(0, input.beg); + + for (size_t i = 0; i < rows; i++) { + input.read((char*)&tmp, sizeof(uint32_t)); + input.read((char*)&Mat(i, 0), sizeof(T) * cols); + } + + RAFT_LOG_DEBUG("File %s loaded", filename); + RAFT_LOG_DEBUG("Rows %zu Cols %zu", rows, cols); + input.close(); +} + +// load_vecs, but duplicate k times +template +void load_vecs_k(const char* filename, M& Mat, size_t k) +{ + static_assert(std::is_same_v, "T must match M::element_type"); + RAFT_EXPECTS(file_exits(filename), "File %s not exists", filename); + + uint32_t tmp; + size_t file_size = get_filesize(filename); + std::ifstream input(filename, std::ios::binary); + + // read `cols` under assumption of file format + input.read(reinterpret_cast(&tmp), sizeof(uint32_t)); + const size_t cols = tmp; + + // calculate rows under assumption of file format + const size_t rows = file_size / (cols * sizeof(T) + sizeof(uint32_t)); + + // allocate final size: rows * k + Mat = raft::make_host_matrix(rows * k, cols); + + // return to start of file and read line by line + input.seekg(0, input.beg); + + for (size_t i = 0; i < rows; ++i) { + // read header for current line + input.read(reinterpret_cast(&tmp), sizeof(uint32_t)); + // read rest of line into i-th row of output + T* base = &Mat(i, 0); + input.read(reinterpret_cast(base), sizeof(T) * cols); + + // create (k-1) copies + for (size_t t = 1; t < k; ++t) { + std::memcpy(&Mat(i + t * rows, 0), base, sizeof(T) * cols); + } + } + + RAFT_LOG_DEBUG("File %s loaded", filename); + RAFT_LOG_DEBUG("Rows %zu Cols %zu", rows * k, cols); + input.close(); +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/utils/StopW.hpp b/cpp/src/neighbors/ivf_rabitq/utils/StopW.hpp new file mode 100644 index 0000000000..8b2934b771 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/utils/StopW.hpp @@ -0,0 +1,45 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +class StopW { + std::chrono::steady_clock::time_point time_begin; + + public: + StopW() { time_begin = std::chrono::steady_clock::now(); } + + float getElapsedTimeSec() + { + std::chrono::steady_clock::time_point time_end = std::chrono::steady_clock::now(); + return (std::chrono::duration_cast(time_end - time_begin).count()); + } + + float getElapsedTimeMili() + { + std::chrono::steady_clock::time_point time_end = std::chrono::steady_clock::now(); + return (std::chrono::duration_cast(time_end - time_begin).count()); + } + + float getElapsedTimeMicro() + { + std::chrono::steady_clock::time_point time_end = std::chrono::steady_clock::now(); + return (std::chrono::duration_cast(time_end - time_begin).count()); + } + + float getElapsedTimeNano() + { + std::chrono::steady_clock::time_point time_end = std::chrono::steady_clock::now(); + return (std::chrono::duration_cast(time_end - time_begin).count()); + } + + void reset() { time_begin = std::chrono::steady_clock::now(); } +}; + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/utils/memory.hpp b/cpp/src/neighbors/ivf_rabitq/utils/memory.hpp new file mode 100644 index 0000000000..0050ddee72 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/utils/memory.hpp @@ -0,0 +1,39 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include +#include +#include + +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +#define PORTABLE_ALIGN32 __attribute__((aligned(32))) +#define PORTABLE_ALIGN64 __attribute__((aligned(64))) + +namespace memory { +template +inline T* align_mm(size_t nbytes) +{ + static_assert(alignment != 0 && (alignment & (alignment - 1)) == 0); + static_assert(alignment % alignof(void*) == 0); + + size_t size = raft::round_up_safe(nbytes, alignment); + if (size == 0) { size = alignment; } + void* p = std::aligned_alloc(alignment, size); + if (p == nullptr) { throw std::bad_alloc{}; } + if constexpr (HUGE_PAGE) { madvise(p, size, MADV_HUGEPAGE); } + std::memset(p, 0, size); + return static_cast(p); +} + +} // namespace memory +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh b/cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh new file mode 100644 index 0000000000..51d3651b45 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh @@ -0,0 +1,43 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// Warp-level sum reduction using shuffle instructions. +template +__inline__ __device__ T warpReduceSum(T val) +{ +#pragma unroll + for (int offset = raft::WarpSize / 2; offset > 0; offset >>= 1) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + return val; +} + +// Block-level sum reduction via per-warp shfl + shared-memory exchange. +// Returns the full block sum in lane 0 of warp 0; other lanes hold partial values. +template +__inline__ __device__ T blockReduceSum(T val) +{ + __shared__ T shared[32]; // up to 1024 threads -> 32 warps + int lane = threadIdx.x & 31; + int wid = threadIdx.x >> 5; + + val = warpReduceSum(val); + if (lane == 0) shared[wid] = val; + __syncthreads(); + + T out = (threadIdx.x < blockDim.x / 32) ? shared[lane] : T(0); + if (wid == 0) out = warpReduceSum(out); + return out; +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu b/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu new file mode 100644 index 0000000000..2bb9fb2174 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu @@ -0,0 +1,67 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "searcher_gpu_utils.hpp" + +#include +#include +#include + +#include +#include +#include + +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// finds the maximum cluster size among clusters to be probed and, optionally, the maximum number of +// vectors probed for any single query +void get_max_probed_cluster_size_and_vectors_count( + raft::resources const& handle, + const ClusterQueryPair* d_cluster_query_pairs, + const size_t num_pairs, + const IVFGPU::GPUClusterMeta* d_cluster_meta, // actually just need the cluster sizes + const size_t num_queries, + uint32_t& max_probed_cluster_size, + std::optional& max_probed_vectors_count) +{ + auto stream = raft::resource::get_cuda_stream(handle); + + // check if max_probed_vectors_count is requested for evaluation + const bool get_max_probed_vectors_count = max_probed_vectors_count.has_value(); + + auto d_max_probed_cluster_size = raft::make_device_scalar(handle, 0); + auto d_probed_vectors_count = raft::make_device_vector( + handle, get_max_probed_vectors_count ? num_queries : 0); + // raw pointers for passing by value to device lambda + auto d_max_probed_cluster_size_ptr = d_max_probed_cluster_size.data_handle(); + auto d_probed_vectors_count_ptr = d_probed_vectors_count.data_handle(); + if (get_max_probed_vectors_count) { + RAFT_CUDA_TRY(cudaMemsetAsync( + d_probed_vectors_count_ptr, 0, num_queries * sizeof(size_t), stream)); // Initialize to 0 + } + + auto count = thrust::make_counting_iterator(0); + thrust::for_each( + raft::resource::get_thrust_policy(handle), count, count + num_pairs, [=] __device__(int64_t i) { + auto [cluster_idx, query_idx] = d_cluster_query_pairs[i]; + auto cluster_size = d_cluster_meta[cluster_idx].num; + atomicMax(d_max_probed_cluster_size_ptr, cluster_size); + if (get_max_probed_vectors_count) + atomicAdd(&d_probed_vectors_count_ptr[query_idx], + static_cast(cluster_size)); + }); + raft::copy(&max_probed_cluster_size, d_max_probed_cluster_size_ptr, 1, stream); + if (get_max_probed_vectors_count) { + max_probed_vectors_count = thrust::reduce(raft::resource::get_thrust_policy(handle), + d_probed_vectors_count_ptr, + d_probed_vectors_count_ptr + num_queries, + 0, + thrust::maximum()); + } +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.hpp b/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.hpp new file mode 100644 index 0000000000..5126cfbdff --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.hpp @@ -0,0 +1,28 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "../gpu_index/ivf_gpu.cuh" + +#include + +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +// finds the maximum cluster size among clusters to be probed and, optionally, the maximum number of +// vectors probed for any single query +void get_max_probed_cluster_size_and_vectors_count( + raft::resources const& handle, + const ClusterQueryPair* d_cluster_query_pairs, + const size_t num_pairs, + const IVFGPU::GPUClusterMeta* d_cluster_meta, // actually just need the cluster sizes + const size_t num_queries, + uint32_t& max_probed_cluster_size, + std::optional& max_probed_vectors_count); + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/utils/space.hpp b/cpp/src/neighbors/ivf_rabitq/utils/space.hpp new file mode 100644 index 0000000000..df756de7e0 --- /dev/null +++ b/cpp/src/neighbors/ivf_rabitq/utils/space.hpp @@ -0,0 +1,55 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq::detail { + +inline float normalize_query16_scalar(float* unit_q, // out: normalised vector (length D) + const float* __restrict__ q, + const float* __restrict__ c, + float norm, + size_t D) +{ + constexpr float eps = 1e-5f; + + if (norm > eps) { + const float inv_norm = 1.0f / norm; + float sum = 0.0f; + + for (size_t i = 0; i < D; ++i) { + float u = (q[i] - c[i]) * inv_norm; // (q - c) / norm + unit_q[i] = u; + sum += u; // running sum + } + return sum; // same as _mm512_reduce_add_ps + } else // q‑c is (almost) the zero vector + { + float value = 1.0f / std::sqrt(static_cast(D)); // 1 / √D + std::fill(unit_q, unit_q + D, value); + return static_cast(D) * value; // == √D + } +} + +inline float L2SqrCPU_STL(const float* h_x, const float* h_y, size_t N) +{ + // Using STL algorithms + return std::inner_product( + h_x, + h_x + N, // First range + h_y, // Second range begin + 0.0f, // Initial value + std::plus(), // Sum operation + [](float a, float b) { // Product operation (replaced with squared difference) + float diff = a - b; + return diff * diff; + }); +} + +} // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 9b96f94bf0..13a07b10b5 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -155,6 +155,13 @@ ConfigureTest( PERCENT 100 ) +ConfigureTest( + NAME NEIGHBORS_ANN_IVF_RABITQ_TEST + PATH neighbors/ann_ivf_rabitq/test_float_int64_t.cu + GPUS 1 + PERCENT 100 +) + ConfigureTest( NAME NEIGHBORS_ANN_SCANN_TEST PATH neighbors/ann_scann/test_float_int64_t.cu diff --git a/cpp/tests/neighbors/ann_ivf_rabitq.cuh b/cpp/tests/neighbors/ann_ivf_rabitq.cuh new file mode 100644 index 0000000000..938f41f846 --- /dev/null +++ b/cpp/tests/neighbors/ann_ivf_rabitq.cuh @@ -0,0 +1,396 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include "../test_utils.cuh" +#include "ann_utils.cuh" +#include "naive_knn.cuh" +#include +#include + +namespace cuvs::neighbors::ivf_rabitq { + +struct ivf_rabitq_inputs { + uint32_t num_db_vecs = 4096; + uint32_t num_queries = 1024; + uint32_t dim = 64; + uint32_t k = 10; + std::optional min_recall = std::nullopt; + + cuvs::neighbors::ivf_rabitq::index_params index_params; + cuvs::neighbors::ivf_rabitq::search_params search_params; + + // Set some default parameters for tests + ivf_rabitq_inputs() { index_params.n_lists = max(32u, min(1024u, num_db_vecs / 128u)); } +}; + +inline auto operator<<(std::ostream& os, const ivf_rabitq::search_mode& p) -> std::ostream& +{ + switch (p) { + case ivf_rabitq::search_mode::LUT16: os << "search_mode::LUT16"; break; + case ivf_rabitq::search_mode::LUT32: os << "search_mode::LUT32"; break; + case ivf_rabitq::search_mode::QUANT4: os << "search_mode::QUANT4"; break; + case ivf_rabitq::search_mode::QUANT8: os << "search_mode::QUANT8"; break; + default: RAFT_FAIL("unreachable code"); + } + return os; +} + +inline auto operator<<(std::ostream& os, const ivf_rabitq_inputs& p) -> std::ostream& +{ + ivf_rabitq_inputs dflt; + bool need_comma = false; +#define PRINT_DIFF_V(spec, val) \ + do { \ + if (dflt spec != p spec) { \ + if (need_comma) { os << ", "; } \ + os << #spec << " = " << val; \ + need_comma = true; \ + } \ + } while (0) +#define PRINT_DIFF(spec) PRINT_DIFF_V(spec, p spec) + + os << "ivf_rabitq_inputs {"; + PRINT_DIFF(.num_db_vecs); + PRINT_DIFF(.num_queries); + PRINT_DIFF(.dim); + PRINT_DIFF(.k); + PRINT_DIFF_V(.min_recall, p.min_recall.value_or(0)); + PRINT_DIFF(.index_params.n_lists); + PRINT_DIFF(.index_params.bits_per_dim); + PRINT_DIFF(.index_params.kmeans_n_iters); + PRINT_DIFF(.index_params.fast_quantize_flag); + PRINT_DIFF(.search_params.n_probes); + PRINT_DIFF(.search_params.mode); + os << "}"; + return os; +} + +template +class ivf_rabitq_test : public ::testing::TestWithParam { + public: + ivf_rabitq_test() + : stream_(raft::resource::get_cuda_stream(handle_)), + ps(::testing::TestWithParam::GetParam()), + database(0, stream_), + search_queries(0, stream_) + { + } + + void gen_data() + { + database.resize(size_t{ps.num_db_vecs} * size_t{ps.dim}, stream_); + search_queries.resize(size_t{ps.num_queries} * size_t{ps.dim}, stream_); + + raft::random::RngState r(1234ULL); + if constexpr (std::is_same{}) { + raft::random::uniform( + handle_, r, database.data(), ps.num_db_vecs * ps.dim, DataT(0.1), DataT(2.0)); + raft::random::uniform( + handle_, r, search_queries.data(), ps.num_queries * ps.dim, DataT(0.1), DataT(2.0)); + } else { + raft::random::uniformInt( + handle_, r, database.data(), ps.num_db_vecs * ps.dim, DataT(1), DataT(20)); + raft::random::uniformInt( + handle_, r, search_queries.data(), ps.num_queries * ps.dim, DataT(1), DataT(20)); + } + raft::resource::sync_stream(handle_); + } + + void calc_ref() + { + 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_); + cuvs::neighbors::naive_knn( + handle_, + distances_naive_dev.data(), + indices_naive_dev.data(), + search_queries.data(), + database.data(), + ps.num_queries, + ps.num_db_vecs, + ps.dim, + ps.k, + static_cast((int)ps.index_params.metric)); + distances_ref.resize(queries_size); + raft::update_host(distances_ref.data(), distances_naive_dev.data(), queries_size, stream_); + indices_ref.resize(queries_size); + raft::update_host(indices_ref.data(), indices_naive_dev.data(), queries_size, stream_); + raft::resource::sync_stream(handle_); + } + + auto build_only() + { + auto ipams = ps.index_params; + + auto database_view = + raft::make_device_matrix_view(database.data(), ps.num_db_vecs, ps.dim); + return cuvs::neighbors::ivf_rabitq::build(handle_, ipams, database_view); + } + + auto build_only_host_input() + { + auto ipams = ps.index_params; + + auto host_database = raft::make_host_matrix(ps.num_db_vecs, ps.dim); + raft::copy(host_database.data_handle(), database.data(), ps.num_db_vecs * ps.dim, stream_); + auto database_view = raft::make_host_matrix_view( + host_database.data_handle(), ps.num_db_vecs, ps.dim); + return cuvs::neighbors::ivf_rabitq::build(handle_, ipams, database_view); + } + + auto build_serialize() + { + tmp_index_file index_file; + auto idx_to_serialize = build_only(); + cuvs::neighbors::ivf_rabitq::serialize(handle_, index_file.filename, idx_to_serialize); + cuvs::neighbors::ivf_rabitq::index deserialized_index(handle_); + cuvs::neighbors::ivf_rabitq::deserialize(handle_, index_file.filename, &deserialized_index); + return deserialized_index; + } + + auto build_host_input_serialize() + { + tmp_index_file index_file; + auto idx_to_serialize = build_only_host_input(); + cuvs::neighbors::ivf_rabitq::serialize(handle_, index_file.filename, idx_to_serialize); + cuvs::neighbors::ivf_rabitq::index deserialized_index(handle_); + cuvs::neighbors::ivf_rabitq::deserialize(handle_, index_file.filename, &deserialized_index); + return deserialized_index; + } + + auto build_with_forced_streaming() + { + tmp_index_file index_file; + auto ipams = ps.index_params; + // Force streaming construction even if dataset fits in GPU memory + ipams.force_streaming = true; + // Use batch size that ensures at least 2-3 batches for typical test sizes + // but scales reasonably for larger datasets + ipams.streaming_batch_size = std::max(size_t{1000}, size_t{ps.num_db_vecs / 3}); + + auto host_database = raft::make_host_matrix(ps.num_db_vecs, ps.dim); + raft::copy(host_database.data_handle(), database.data(), ps.num_db_vecs * ps.dim, stream_); + auto database_view = raft::make_host_matrix_view( + host_database.data_handle(), ps.num_db_vecs, ps.dim); + auto idx_to_serialize = cuvs::neighbors::ivf_rabitq::build(handle_, ipams, database_view); + + // Serialize and deserialize to reorganize data for efficient search + cuvs::neighbors::ivf_rabitq::serialize(handle_, index_file.filename, idx_to_serialize); + cuvs::neighbors::ivf_rabitq::index deserialized_index(handle_); + cuvs::neighbors::ivf_rabitq::deserialize(handle_, index_file.filename, &deserialized_index); + return deserialized_index; + } + + template + void run(BuildIndex build_index) + { + index index = build_index(); + + double compression_ratio = sizeof(DataT) * 8 / ps.index_params.bits_per_dim; + + size_t queries_size = ps.num_queries * ps.k; + std::vector indices_ivf_rabitq(queries_size); + std::vector distances_ivf_rabitq(queries_size); + + rmm::device_uvector distances_ivf_rabitq_dev(queries_size, stream_); + rmm::device_uvector indices_ivf_rabitq_dev(queries_size, stream_); + + auto query_view = + raft::make_device_matrix_view(search_queries.data(), ps.num_queries, ps.dim); + auto inds_view = raft::make_device_matrix_view( + indices_ivf_rabitq_dev.data(), ps.num_queries, ps.k); + auto dists_view = raft::make_device_matrix_view( + distances_ivf_rabitq_dev.data(), ps.num_queries, ps.k); + + cuvs::neighbors::ivf_rabitq::search( + handle_, ps.search_params, index, query_view, inds_view, dists_view); + + raft::update_host( + distances_ivf_rabitq.data(), distances_ivf_rabitq_dev.data(), queries_size, stream_); + raft::update_host( + indices_ivf_rabitq.data(), indices_ivf_rabitq_dev.data(), queries_size, stream_); + raft::resource::sync_stream(handle_); + + // A very conservative lower bound on recall + double min_recall = 0.5; + // Use explicit per-test min recall value if provided. + min_recall = ps.min_recall.value_or(min_recall); + + ASSERT_TRUE(cuvs::neighbors::eval_neighbours(indices_ref, + indices_ivf_rabitq, + distances_ref, + distances_ivf_rabitq, + ps.num_queries, + ps.k, + 0.0001 * compression_ratio, + min_recall)) + << ps; + } + + void SetUp() override // NOLINT + { + gen_data(); + calc_ref(); + } + + void TearDown() override // NOLINT + { + cudaGetLastError(); + raft::resource::sync_stream(handle_); + database.resize(0, stream_); + search_queries.resize(0, stream_); + } + + private: + raft::resources handle_; + rmm::cuda_stream_view stream_; + ivf_rabitq_inputs ps; // NOLINT + rmm::device_uvector database; // NOLINT + rmm::device_uvector search_queries; // NOLINT + std::vector indices_ref; // NOLINT + std::vector distances_ref; // NOLINT +}; + +/* Test cases */ +using test_cases_t = std::vector; + +// concatenate parameter sets for different type +template +auto operator+(const std::vector& a, const std::vector& b) -> std::vector +{ + std::vector res = a; + res.insert(res.end(), b.begin(), b.end()); + return res; +} + +inline auto defaults() -> test_cases_t { return {ivf_rabitq_inputs{}}; } + +template +auto map(const std::vector& xs, F f) -> std::vector +{ + std::vector ys(xs.size()); + std::transform(xs.begin(), xs.end(), ys.begin(), f); + return ys; +} + +inline auto with_dims(const std::vector& dims) -> test_cases_t +{ + return map(dims, [](uint32_t d) { + ivf_rabitq_inputs x; + x.dim = d; + return x; + }); +} + +inline auto small_dims() -> test_cases_t { return with_dims({1, 2, 3, 4, 5, 6, 7, 8}); } + +inline auto big_dims() -> test_cases_t +{ + return with_dims({512, 513, 1023, 1024, 1025, 2048, 2049, 2050}); +} + +inline auto var_n_probes() -> test_cases_t +{ + ivf_rabitq_inputs dflt; + std::vector xs; + for (auto x = dflt.index_params.n_lists; x >= 1; x /= 2) { + xs.push_back(x); + } + return map(xs, [](uint32_t n_probes) { + ivf_rabitq_inputs x; + x.search_params.n_probes = n_probes; + // reduce `min_recall` for low `n_probes` + if (n_probes <= 5) { x.min_recall = 0.08 * n_probes; } + return x; + }); +} + +inline auto var_k() -> test_cases_t +{ + return map({1, 5, 10, 32, 64, 16384}, [](uint32_t k) { + ivf_rabitq_inputs x; + x.k = k; + x.num_db_vecs = max(x.num_db_vecs, k * 2); + if (k > 64) x.num_queries = 64; // reduce runtime of large-k tests + return x; + }); +} + +inline auto var_bits_per_dim() -> test_cases_t +{ + ivf_rabitq_inputs dflt; + std::vector xs; + for (auto x = 1; x <= 9; ++x) { + xs.push_back(x); + } + return map(xs, [](uint32_t bits_per_dim) { + ivf_rabitq_inputs x; + x.index_params.bits_per_dim = bits_per_dim; + if (bits_per_dim == 1) { x.min_recall = 0.3; } + return x; + }); +} + +inline auto var_search_mode() -> test_cases_t +{ + ivf_rabitq_inputs dflt; + std::vector xs{ivf_rabitq::search_mode::LUT16, + ivf_rabitq::search_mode::LUT32, + ivf_rabitq::search_mode::QUANT4, + ivf_rabitq::search_mode::QUANT8}; + + return map(xs, [](cuvs::neighbors::ivf_rabitq::search_mode mode) { + ivf_rabitq_inputs x; + x.search_params.mode = mode; + return x; + }); +} + +inline auto var_search_mode_1_bit() -> test_cases_t +{ + ivf_rabitq_inputs dflt; + std::vector xs{ivf_rabitq::search_mode::LUT16, + ivf_rabitq::search_mode::LUT32, + ivf_rabitq::search_mode::QUANT4, + ivf_rabitq::search_mode::QUANT8}; + + return map(xs, [](cuvs::neighbors::ivf_rabitq::search_mode mode) { + ivf_rabitq_inputs x; + x.search_params.mode = mode; + x.index_params.bits_per_dim = 1; + x.min_recall = 0.3; + return x; + }); +} + +/* Test instantiations */ + +// Currently IVF-RaBitQ deserialization reorganizes data for efficient search and is required for +// producing correct results. + +#define TEST_BUILD_SERIALIZE_SEARCH(type) \ + TEST_P(type, build_serialize_search) /* NOLINT */ \ + { \ + this->run([this]() { return this->build_serialize(); }); \ + } + +#define TEST_BUILD_HOST_INPUT_SERIALIZE_SEARCH(type) \ + TEST_P(type, build_host_input_serialize_search) /* NOLINT */ \ + { \ + this->run([this]() { return this->build_host_input_serialize(); }); \ + } + +#define TEST_BUILD_FORCED_STREAMING(type) \ + TEST_P(type, build_forced_streaming) /* NOLINT */ \ + { \ + this->run([this]() { return this->build_with_forced_streaming(); }); \ + } + +#define INSTANTIATE(type, vals) \ + INSTANTIATE_TEST_SUITE_P(IvfRabitq, type, ::testing::ValuesIn(vals)); /* NOLINT */ + +} // namespace cuvs::neighbors::ivf_rabitq 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 new file mode 100644 index 0000000000..2b412f3401 --- /dev/null +++ b/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu @@ -0,0 +1,19 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../ann_ivf_rabitq.cuh" + +namespace cuvs::neighbors::ivf_rabitq { + +using f32_f32_i64 = ivf_rabitq_test; + +TEST_BUILD_SERIALIZE_SEARCH(f32_f32_i64) +TEST_BUILD_HOST_INPUT_SERIALIZE_SEARCH(f32_f32_i64) +TEST_BUILD_FORCED_STREAMING(f32_f32_i64) +INSTANTIATE(f32_f32_i64, + defaults() + small_dims() + big_dims() + var_n_probes() + var_k() + var_bits_per_dim() + + var_search_mode() + var_search_mode_1_bit()); + +} // namespace cuvs::neighbors::ivf_rabitq diff --git a/fern/docs.yml b/fern/docs.yml index 5e6eab3d42..882dd665aa 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -367,6 +367,8 @@ navigation: path: "./pages/cpp_api/cpp-api-neighbors-ivf-flat.md" - page: "Neighbors IVF PQ" path: "./pages/cpp_api/cpp-api-neighbors-ivf-pq.md" + - page: "Neighbors IVF RaBitQ" + path: "./pages/cpp_api/cpp-api-neighbors-ivf-rabitq.md" - page: "Neighbors IVF SQ" path: "./pages/cpp_api/cpp-api-neighbors-ivf-sq.md" - page: "Neighbors NN Descent" diff --git a/fern/pages/cpp_api/cpp-api-neighbors-ivf-rabitq.md b/fern/pages/cpp_api/cpp-api-neighbors-ivf-rabitq.md new file mode 100644 index 0000000000..fa39e69e98 --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-neighbors-ivf-rabitq.md @@ -0,0 +1,291 @@ +--- +slug: api-reference/cpp-api-neighbors-ivf-rabitq +--- + +# IVF RaBitQ + +_Source header: `cuvs/neighbors/ivf_rabitq.hpp`_ + +## IVF-RaBitQ index build parameters + + +### neighbors::ivf_rabitq::index_params + +IVF-RaBitQ index build parameters + +```cpp +struct index_params : cuvs::neighbors::index_params { + uint32_t n_lists; + uint32_t bits_per_dim; + uint32_t kmeans_n_iters; + uint32_t max_train_points_per_cluster; + bool fast_quantize_flag; + size_t streaming_batch_size; + bool force_streaming; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `n_lists` | `uint32_t` | The number of inverted lists (clusters)

Hint: Increasing this parameter may alleviate shared memory pressure. | +| `bits_per_dim` | `uint32_t` | The total number of bits per dimension (single bit required for the binary RaBitQ algorithm + additional bits for extended RaBitQ).

Supported values: [1, 2, 3, 4, 5, 6, 7, 8, 9].

Hint: the smaller the 'bits_per_dim', the smaller the index size and the better the search performance, but the lower the recall. | +| `kmeans_n_iters` | `uint32_t` | The number of iterations searching for kmeans centers (index building). | +| `max_train_points_per_cluster` | `uint32_t` | The number of data vectors (per cluster) to use during iterative kmeans building. | +| `fast_quantize_flag` | `bool` | Flag for using the fast quantize method | +| `streaming_batch_size` | `size_t` | Maximum number of vectors per batch when using streaming construction.

This parameter controls the batch size during streaming construction from host memory. Batches contain complete clusters only (no partial clusters across batch boundaries).

Note: Streaming construction is automatically used when the dataset doesn't fit comfortably in GPU memory (determined by available workspace and kTolerableRatio). | +| `force_streaming` | `bool` | Force streaming construction regardless of dataset size.

When set to true, streaming construction will be used even if the dataset would fit in GPU memory. This is useful for testing or when you want explicit control over the construction method.

Note: This parameter only applies when the input dataset is in host memory. If the dataset is already in device memory, streaming construction is not applicable and this parameter has no effect.

Default: false (auto-detect based on available memory) | + +## IVF-RaBitQ index search parameters + + +### neighbors::ivf_rabitq::search_mode + +A type for specifying the mode for searching the RaBitQ index. + +```cpp +enum class search_mode { + LUT16 = 0, + LUT32 = 1, + QUANT4 = 2, + QUANT8 = 3 +}; +``` + +**Values** + +| Name | Value | +| --- | --- | +| `LUT16` | `0` | +| `LUT32` | `1` | +| `QUANT4` | `2` | +| `QUANT8` | `3` | + +## IVF-RaBitQ index + + +### neighbors::ivf_rabitq::index + +IVF-RaBitQ index. + +```cpp +template +struct index; +``` + + +### neighbors::ivf_rabitq::index::index + +Construct an empty index yet to be populated. + +```cpp +index(raft::resources const& handle); +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `handle` | | `raft::resources const&` | | + +**Returns** + +`void` + +**Additional overload:** `neighbors::ivf_rabitq::index::index` + +Construct an empty index yet to be populated. + +```cpp +index(raft::resources const& handle, +size_t n_rows, +uint32_t dim, +uint32_t n_lists, +uint32_t bits_per_dim); +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `handle` | | `raft::resources const&` | | +| `n_rows` | | `size_t` | | +| `dim` | | `uint32_t` | | +| `n_lists` | | `uint32_t` | | +| `bits_per_dim` | | `uint32_t` | | + +**Returns** + +`void` + + +### neighbors::ivf_rabitq::index::dim + +Dimensionality of the input data. + +```cpp +uint32_t dim() const noexcept; +``` + +**Returns** + +`uint32_t` + + +### neighbors::ivf_rabitq::index::size + +Total length of the index. + +```cpp +IdxT size() const noexcept; +``` + +**Returns** + +`IdxT` + + +### neighbors::ivf_rabitq::index::rabitq_index + +Accessor for underlying RaBitQ index + +```cpp +detail::IVFGPU& rabitq_index() noexcept; +``` + +**Returns** + +`detail::IVFGPU&` + +## IVF-RaBitQ index build + + +### neighbors::ivf_rabitq::build + +Build the index from the dataset for efficient search. + +```cpp +auto build(raft::resources const& handle, +const cuvs::neighbors::ivf_rabitq::index_params& index_params, +raft::device_matrix_view dataset) +-> cuvs::neighbors::ivf_rabitq::index; +``` + +Usage example: + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `handle` | in | `raft::resources const&` | | +| `index_params` | in | [`const cuvs::neighbors::ivf_rabitq::index_params&`](/api-reference/cpp-api-neighbors-ivf-rabitq#neighbors-ivf-rabitq-index-params) | configure the index building | +| `dataset` | in | `raft::device_matrix_view` | a device_matrix_view to a row-major matrix [n_rows, dim] | + +**Returns** + +[`cuvs::neighbors::ivf_rabitq::index`](/api-reference/cpp-api-neighbors-ivf-rabitq#neighbors-ivf-rabitq-index) + +**Additional overload:** `neighbors::ivf_rabitq::build` + +Build the index from the dataset for efficient search. + +```cpp +auto build(raft::resources const& handle, +const cuvs::neighbors::ivf_rabitq::index_params& index_params, +raft::host_matrix_view dataset) +-> cuvs::neighbors::ivf_rabitq::index; +``` + +Usage example: + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `handle` | in | `raft::resources const&` | | +| `index_params` | in | [`const cuvs::neighbors::ivf_rabitq::index_params&`](/api-reference/cpp-api-neighbors-ivf-rabitq#neighbors-ivf-rabitq-index-params) | configure the index building | +| `dataset` | in | `raft::host_matrix_view` | a host_matrix_view to a row-major matrix [n_rows, dim] | + +**Returns** + +[`cuvs::neighbors::ivf_rabitq::index`](/api-reference/cpp-api-neighbors-ivf-rabitq#neighbors-ivf-rabitq-index) + +## IVF-RaBitQ index search + + +### neighbors::ivf_rabitq::search + +Search ANN using the constructed index. + +```cpp +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); +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `handle` | in | `raft::resources const&` | | +| `search_params` | in | `const cuvs::neighbors::ivf_rabitq::search_params&` | configure the search | +| `index` | in | [`cuvs::neighbors::ivf_rabitq::index&`](/api-reference/cpp-api-neighbors-ivf-rabitq#neighbors-ivf-rabitq-index) | ivf-rabitq constructed index | +| `queries` | in | `raft::device_matrix_view` | a device matrix view to a row-major matrix [n_queries, index->dim()] | +| `neighbors` | out | `raft::device_matrix_view` | a device matrix view to the indices of the neighbors in the source dataset [n_queries, k] | +| `distances` | out | `raft::device_matrix_view` | a device matrix view to the distances to the selected neighbors [n_queries, k] | + +**Returns** + +`void` + +## IVF-RaBitQ index serialize + + +### neighbors::ivf_rabitq::serialize + +Save the index to file. + +```cpp +void serialize(raft::resources const& handle, +const std::string& filename, +cuvs::neighbors::ivf_rabitq::index& index); +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `handle` | in | `raft::resources const&` | the raft handle | +| `filename` | in | `const std::string&` | the file name for saving the index | +| `index` | in | [`cuvs::neighbors::ivf_rabitq::index&`](/api-reference/cpp-api-neighbors-ivf-rabitq#neighbors-ivf-rabitq-index) | IVF-RaBitQ index | + +**Returns** + +`void` + + +### neighbors::ivf_rabitq::deserialize + +Load index from file. + +```cpp +void deserialize(raft::resources const& handle, +const std::string& filename, +cuvs::neighbors::ivf_rabitq::index* index); +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `handle` | in | `raft::resources const&` | the raft handle | +| `filename` | in | `const std::string&` | the name of the file that stores the index | +| `index` | out | [`cuvs::neighbors::ivf_rabitq::index*`](/api-reference/cpp-api-neighbors-ivf-rabitq#neighbors-ivf-rabitq-index) | IVF-PQ index | + +**Returns** + +`void` diff --git a/fern/pages/cpp_api/index.md b/fern/pages/cpp_api/index.md index d0bad37b16..e61927491a 100644 --- a/fern/pages/cpp_api/index.md +++ b/fern/pages/cpp_api/index.md @@ -30,6 +30,7 @@ These pages are generated from the documented public headers in the cuVS source - [Index](/api-reference/cpp-api-neighbors-composite-index) - [IVF Flat](/api-reference/cpp-api-neighbors-ivf-flat) - [IVF PQ](/api-reference/cpp-api-neighbors-ivf-pq) +- [IVF RaBitQ](/api-reference/cpp-api-neighbors-ivf-rabitq) - [IVF SQ](/api-reference/cpp-api-neighbors-ivf-sq) - [NN Descent](/api-reference/cpp-api-neighbors-nn-descent) - [Refine](/api-reference/cpp-api-neighbors-refine) diff --git a/python/cuvs_bench/cuvs_bench/config/algorithms.yaml b/python/cuvs_bench/cuvs_bench/config/algorithms.yaml index f181e549eb..5f77c43623 100644 --- a/python/cuvs_bench/cuvs_bench/config/algorithms.yaml +++ b/python/cuvs_bench/cuvs_bench/config/algorithms.yaml @@ -25,6 +25,9 @@ faiss_cpu_ivf_flat: faiss_cpu_ivf_pq: executable: FAISS_CPU_IVF_PQ_ANN_BENCH requires_gpu: false +faiss_cpu_ivf_rabitq: + executable: FAISS_CPU_IVF_RABITQ_ANN_BENCH + requires_gpu: false faiss_cpu_ivf_sq: executable: FAISS_CPU_IVF_SQ_ANN_BENCH requires_gpu: false @@ -37,6 +40,9 @@ cuvs_ivf_flat: cuvs_ivf_pq: executable: CUVS_IVF_PQ_ANN_BENCH requires_gpu: true +cuvs_ivf_rabitq: + executable: CUVS_IVF_RABITQ_ANN_BENCH + requires_gpu: true cuvs_ivf_sq: executable: CUVS_IVF_SQ_ANN_BENCH requires_gpu: true