diff --git a/c/CMakeLists.txt b/c/CMakeLists.txt index eab7560dfc..59932c32ee 100644 --- a/c/CMakeLists.txt +++ b/c/CMakeLists.txt @@ -103,6 +103,7 @@ add_library( src/neighbors/all_neighbors.cpp src/preprocessing/pca.cpp src/preprocessing/quantize/binary.cpp + src/preprocessing/quantize/bbq.cpp src/preprocessing/quantize/pq.cpp src/preprocessing/quantize/scalar.cpp src/distance/pairwise_distance.cpp diff --git a/c/include/cuvs/core/all.h b/c/include/cuvs/core/all.h index 1b693a355e..80b95992eb 100644 --- a/c/include/cuvs/core/all.h +++ b/c/include/cuvs/core/all.h @@ -44,6 +44,7 @@ #include #include +#include #include #include diff --git a/c/include/cuvs/core/dataset.h b/c/include/cuvs/core/dataset.h index 0fbc168759..1a1ab2565b 100644 --- a/c/include/cuvs/core/dataset.h +++ b/c/include/cuvs/core/dataset.h @@ -21,7 +21,8 @@ extern "C" { typedef enum { CUVS_DATASET_LAYOUT_STANDARD = 0, CUVS_DATASET_LAYOUT_PADDED = 1, - CUVS_DATASET_LAYOUT_PQ = 2 + CUVS_DATASET_LAYOUT_PQ = 2, + CUVS_DATASET_LAYOUT_BBQ = 3 } cuvsDatasetLayout_t; /** diff --git a/c/include/cuvs/neighbors/cagra.h b/c/include/cuvs/neighbors/cagra.h index 6a8b854fc2..602f52dbfc 100644 --- a/c/include/cuvs/neighbors/cagra.h +++ b/c/include/cuvs/neighbors/cagra.h @@ -654,7 +654,8 @@ CUVS_EXPORT cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res, * * The memory space and layout \p dataset was constructed with select the C++ build overload. * Build the handle with an owning factory or the matching dataset view factory - * (`cuvsDatasetMakePaddedView` / `cuvsDatasetMakeStandardView`). + * (`cuvsDatasetMakePaddedView`, `cuvsDatasetMakeStandardView`, or + * `cuvsDatasetMakeBbqView`). * * Note that a dataset residing in host memory produces a host-backed index, which * must be made search-ready with `cuvsCagraUpdateDataset` (using a device-padded @@ -696,6 +697,8 @@ CUVS_EXPORT cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res, * A `CUVS_DATASET_LAYOUT_PQ` dataset created by `cuvsDatasetMakePQ` builds an iterative CAGRA-Q * index. VPQ input requires `L2Expanded` and `ITERATIVE_CAGRA_SEARCH` (or `AUTO_SELECT`), and the * VPQ dataset must outlive the index because the index stores a non-owning view. + * A `CUVS_DATASET_LAYOUT_BBQ` dataset builds a graph-only index; attach a searchable dataset with + * `cuvsCagraUpdateDataset` before search. * * @param[in] res cuvsResources_t opaque C handle * @param[in] params cuvsCagraIndexParams_t used to build CAGRA index @@ -860,8 +863,9 @@ CUVS_EXPORT cuvsError_t cuvsCagraSearchMultiPartition(cuvsResources_t res, /** * Save the CAGRA graph to file without its dataset. * - * This supports dense and PQ-backed indexes. The dataset must be attached separately after loading - * the graph. + * This supports dense, PQ-backed, and BBQ-built indexes. The serialized file does not contain + * vector data. After deserialization the index cannot be searched until a compatible dataset is + * attached with `cuvsCagraUpdateDataset`. * * Experimental, both the API and the serialization format are subject to change. * @@ -876,9 +880,10 @@ CUVS_EXPORT cuvsError_t cuvsCagraSerializeGraph(cuvsResources_t res, /** * Save the CAGRA graph and its attached dataset to file. * - * The index stores a non-owning dataset view. The caller must keep the dataset backing that view - * alive while this function runs. Returns CUVS_ERROR without modifying the destination file if - * the index has no attached dataset. PQ datasets are not serialized by this function. + * The index stores a non-owning dataset view. The caller must keep the memory of the dataset + * backing that view alive while this function runs. Returns CUVS_ERROR without modifying the + * destination file if the index has no attached dataset. PQ and BBQ datasets are not serialized + * by this function. * * Experimental, both the API and the serialization format are subject to change. * diff --git a/c/include/cuvs/preprocessing/quantize/bbq.h b/c/include/cuvs/preprocessing/quantize/bbq.h new file mode 100644 index 0000000000..b08a989a18 --- /dev/null +++ b/c/include/cuvs/preprocessing/quantize/bbq.h @@ -0,0 +1,142 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @defgroup preprocessing_c_bbq C API for Better Binary Quantization datasets + * @{ + */ + +/** + * Storage layout of BBQ/OSQ quantized component codes in each dataset row. + * CUVS_BBQ_CODE_LAYOUT_PACKED_1B: Each dimension is quantized to a single bit and packed into bytes. Reflects + * Lucene's OptimizedScalarQuantizer.packAsBinary. + * CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_2B: Each dimension is quantized to 2 bits, stored as 2 bitplanes. + * Reflects Lucene's OptimizedScalarQuantizer.transposeDibit. SIMT popc path only + * (paired with a transposed_4b or packed_1b operand); + * CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_4B: Each dimension is quantized to 4 bits, optimized for bitwise operations. + * Reflects Lucene's OptimizedScalarQuantizer.transposeHalfByte. the first bit of + * every dimension is in the first set dimensions bits, or (dimensions/8) + * bytes. The second, third, and fourth bits are in the second, third, and + * fourth set of dimensions bits, respectively. Format used for queries. + * CUVS_BBQ_CODE_LAYOUT_PACKED_4B: Each dimension is quantized to 4 bits, two values are packed into each output + * byte. + * CUVS_BBQ_CODE_LAYOUT_PACKED_7B: Each dimension is quantized to 7 bits and treated as a signed value. + * CUVS_BBQ_CODE_LAYOUT_PACKED_8B: Each dimension is quantized to 8 bits and treated as an unsigned value. + */ +typedef enum { + CUVS_BBQ_CODE_LAYOUT_PACKED_1B = 0, + CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_2B, + CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_4B, + CUVS_BBQ_CODE_LAYOUT_PACKED_4B, + CUVS_BBQ_CODE_LAYOUT_PACKED_7B, + CUVS_BBQ_CODE_LAYOUT_PACKED_8B +} cuvsBbqCodeLayout_t; + + +/** + * @brief Better Binary Quantization + * ([BBQ](https://www.elastic.co/search-labs/blog/better-binary-quantization-lucene-elasticsearch)) + * is a vector-quantization approach used in Elasticsearch and Apache Lucene. It builds on ideas + * introduced in RaBitQ([Gao and Long](https://arxiv.org/pdf/2405.12497, [Gao et + * al.](https://arxiv.org/pdf/2409.09913)): residual binary codes around a centroid, corrective + * factors, and efficient bitwise comparison of codes at different bit widths. Lucene implements + * this as optimized scalar quantization (OSQ) with packed and bit-plane layouts; Elasticsearch + * exposes it as BBQ. + * + * BBQ in cuVS designed to be compatible with the Lucene/Elasticsearch dataset: a single shared + * centroid, no random rotation, and OSQ codes. + * + * RaBitQ and BBQ in cuVS both compress centroid-relative vectors to low-bit codes and retain + * additional per-vector information so search is better than naïve sign-bit comparison. They differ + * in transformation and scale representation. RaBitQ commonly separates residual magnitude from + * direction, then applies a random orthogonal rotation before binary coding; BBQ uses per-vector + * scalar intervals to interpret the compressed residual codes. + */ +typedef struct cuvsBbqQuantizer { + uintptr_t addr; + void (*destroy_addr)(void*); + DLDataType dtype; + bool is_owning; +} cuvsBbqQuantizer; +typedef cuvsBbqQuantizer* cuvsBbqQuantizer_t; + +/** + * @brief Create a BBQ quantizer view from caller-owned device tensors. + * + * Tensors are not copied and must remain valid while a derived dataset is in use. + * + * @param[in] codes uint8 device matrix containing encoded rows + * @param[in] lower_intervals float32 device vector with one lower interval per row + * @param[in] upper_intervals float32 device vector with one upper interval per row + * @param[in] additional_corrections float32 device vector with one correction per row + * @param[in] quantized_component_sums int32 device vector with one component sum per row + * @param[in] centroid device vector containing the dataset centroid + * @param[in] dequant_delta float32 device vector with one dequantization delta per row + * @param[in] dequant_sum_delta float32 device vector with one delta-times-sum value per row + * @param[in] row_norm float32 device vector with one original-space squared norm per row + * @param[in] layout encoded code layout + * @param[in] metric distance metric associated with the encoded dataset + * @param[in] centroid_norm_sq squared norm of the centroid + * @param[out] quantizer newly allocated non-owning quantizer handle + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsBbqQuantizerCreateView( + DLManagedTensor* codes, + DLManagedTensor* lower_intervals, + DLManagedTensor* upper_intervals, + DLManagedTensor* additional_corrections, + DLManagedTensor* quantized_component_sums, + DLManagedTensor* centroid, + DLManagedTensor* dequant_delta, + DLManagedTensor* dequant_sum_delta, + DLManagedTensor* row_norm, + cuvsBbqCodeLayout_t layout, + cuvsDistanceType metric, + float centroid_norm_sq, + cuvsBbqQuantizer_t* quantizer); + +/** + * @brief Destroy a BBQ quantizer without destroying its caller-owned tensors. + * + * @param[in] quantizer quantizer handle to destroy + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsBbqQuantizerDestroy(cuvsBbqQuantizer_t quantizer); + +/** + * @brief Create a non-owning device BBQ dataset view. + * + * Accepts one symmetric quantizer or two compatible asymmetric quantizers. + * + * @param[in] res cuVS resources + * @param[in] quantizers array containing one or two BBQ quantizer handles + * @param[in] num_quantizers number of elements in `quantizers` + * @param[out] dataset newly allocated non-owning BBQ dataset handle + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsDatasetMakeBbqView(cuvsResources_t res, + cuvsBbqQuantizer_t* quantizers, + size_t num_quantizers, + cuvsDataset_t* dataset); + +/** @} */ + +#ifdef __cplusplus +} +#endif diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 57e8168c1e..97cc9d8bac 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -24,12 +24,14 @@ #include "../core/exceptions.hpp" #include "../core/interop.hpp" #include +#include #include #include #include #include #include #include +#include #include #include "../core/exceptions.hpp" #include "../core/interop.hpp" @@ -57,6 +59,7 @@ struct sg_cagra_c_api_index_box { device_padded, device_standard, device_vpq, + device_bbq, host_padded, host_standard } layout; @@ -72,6 +75,8 @@ constexpr auto sg_cagra_index_layout_from_view() return sg_cagra_c_api_index_box::dataset_layout::device_padded; } else if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { return sg_cagra_c_api_index_box::dataset_layout::device_vpq; + } else if constexpr (cuvs::neighbors::is_device_bbq_dataset_view_v) { + return sg_cagra_c_api_index_box::dataset_layout::device_bbq; } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { return sg_cagra_c_api_index_box::dataset_layout::host_standard; } else { @@ -79,7 +84,11 @@ constexpr auto sg_cagra_index_layout_from_view() } } -template +template static void with_index_by_layout(sg_cagra_c_api_index_box* box, const char* null_handle_err, const char* host_not_allowed_err, @@ -116,6 +125,16 @@ static void with_index_by_layout(sg_cagra_c_api_index_box* box, fn(*idx); break; } + case sg_cagra_c_api_index_box::dataset_layout::device_bbq: { + if constexpr (AllowBbq) { + using index_t = cuvs::neighbors::cagra::device_bbq_index; + auto* idx = reinterpret_cast(box->index_ptr); + fn(*idx); + } else { + RAFT_FAIL("This operation does not support a BBQ-backed index"); + } + break; + } case sg_cagra_c_api_index_box::dataset_layout::host_standard: { if constexpr (AllowHost) { auto* idx = @@ -390,6 +409,92 @@ static void with_dataset_view(cuvsDataset_t dataset, Fn&& fn) } } +template +void validate_bbq_layouts(cuvs::neighbors::device_bbq_dataset_view const& dataset) +{ + using layout_t = cuvs::preprocessing::quantize::bbq::bbq_code_layout; + if (dataset.quantizers.size() == 1) { + auto layout = dataset.quantizers.front().layout; + RAFT_EXPECTS(layout == layout_t::packed_1b || layout == layout_t::transposed_2b || + layout == layout_t::packed_4b || layout == layout_t::packed_7b || + layout == layout_t::packed_8b, + "cuvsDatasetMakeBbqView: transposed_4b requires an asymmetric layout pair"); + return; + } + + const bool has_1b = dataset.has_layout(layout_t::packed_1b); + const bool has_2bt = dataset.has_layout(layout_t::transposed_2b); + const bool has_4b = dataset.has_layout(layout_t::packed_4b); + const bool has_4bt = dataset.has_layout(layout_t::transposed_4b); + RAFT_EXPECTS((has_1b && (has_4b || has_2bt || has_4bt)) || (has_2bt && has_4bt), + "cuvsDatasetMakeBbqView: unsupported asymmetric BBQ layout pair"); +} + +template +auto get_cpp_bbq_quantizer_view(cuvsBbqQuantizer_t quantizer) + -> cuvs::preprocessing::quantize::bbq::quantizer_view +{ + using owner_t = cuvs::preprocessing::quantize::bbq::quantizer; + using view_t = cuvs::preprocessing::quantize::bbq::quantizer_view; + RAFT_EXPECTS(quantizer != nullptr && quantizer->addr != 0, + "cuvsDatasetMakeBbqView: null quantizer storage"); + if (quantizer->is_owning) { + return reinterpret_cast(quantizer->addr)->view(); + } + return *reinterpret_cast(quantizer->addr); +} + +template +auto make_bbq_dataset_view(cuvsBbqQuantizer_t* quantizers, std::size_t num_quantizers) + -> std::unique_ptr> +{ + using dataset_view_t = cuvs::neighbors::device_bbq_dataset_view; + auto dataset = std::make_unique(); + int64_t expected_rows{-1}; + uint32_t expected_dim{}; + cuvs::distance::DistanceType expected_metric{}; + + for (std::size_t i = 0; i < num_quantizers; ++i) { + RAFT_EXPECTS(quantizers[i] != nullptr, + "cuvsDatasetMakeBbqView: quantizer %zu is null", + i); + auto quantizer = get_cpp_bbq_quantizer_view(quantizers[i]); + if (i == 0) { + expected_rows = quantizer.n_rows(); + expected_dim = quantizer.dim(); + expected_metric = quantizer.metric; + } else { + RAFT_EXPECTS(quantizer.n_rows() == expected_rows && quantizer.dim() == expected_dim, + "cuvsDatasetMakeBbqView: all quantizers must have the same shape"); + RAFT_EXPECTS(quantizer.metric == expected_metric, + "cuvsDatasetMakeBbqView: all quantizers must use the same metric"); + } + dataset->add_quantizer(quantizer); + } + + return dataset; +} + +template +void make_and_bind_bbq_dataset(cuvsBbqQuantizer_t* quantizers, + std::size_t num_quantizers, + DLDataType dtype, + cuvsDataset_t* output) +{ + using dataset_view_t = cuvs::neighbors::device_bbq_dataset_view; + auto view = make_bbq_dataset_view(quantizers, num_quantizers); + validate_bbq_layouts(*view); + + auto handle = std::make_unique(); + handle->addr = reinterpret_cast(view.release()); + handle->destroy_addr = &destroy_typed_addr; + handle->dtype = dtype; + handle->mem_type = CUVS_DATASET_MEM_TYPE_DEVICE; + handle->layout = CUVS_DATASET_LAYOUT_BBQ; + handle->is_owning = false; + *output = handle.release(); +} + using device_vpq_owner_t = cuvs::neighbors::device_vpq_dataset; using device_vpq_view_t = cuvs::neighbors::device_vpq_dataset_view; @@ -845,7 +950,7 @@ void _search(cuvsResources_t res, { auto res_ptr = reinterpret_cast(res); auto* box = reinterpret_cast(index.addr); - with_index_by_layout( + with_index_by_layout( box, "cuvsCagraSearch: null index handle", "cuvsCagraSearch: host index must be converted to device first via " @@ -1014,7 +1119,11 @@ void _serialize(cuvsResources_t res, const char *filename, with_index_by_layout(box, null_handle_err, "", [&](auto &idx) { using index_dataset_view_t = std::remove_cvref_t; - if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { + if constexpr (cuvs::neighbors::is_bbq_dataset_view_v) { + RAFT_EXPECTS(!include_dataset, + "cuvsCagraSerializeGraphAndDataset is not supported for BBQ indices"); + cuvs::neighbors::cagra::serialize(*res_ptr, std::string(filename), idx); + } else if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { RAFT_EXPECTS( !include_dataset, "cuvsCagraSerializeGraphAndDataset is not supported for PQ indices; serialize the PQ " @@ -1186,8 +1295,9 @@ void _serialize_to_hnswlib(cuvsResources_t res, const char *filename, "cuvsCagraSerializeToHnswlib: host indices are allowed", [&](auto &idx) { using index_dataset_view_t = std::remove_cvref_t; - if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { - RAFT_FAIL("cuvsCagraSerializeToHnswlib is not supported for VPQ indices"); + if constexpr (cuvs::neighbors::is_vpq_dataset_view_v || + cuvs::neighbors::is_bbq_dataset_view_v) { + RAFT_FAIL("cuvsCagraSerializeToHnswlib is not supported for quantized dataset layouts"); } else { cuvs::neighbors::cagra::serialize_to_hnswlib( *res_ptr, std::string(filename), idx); @@ -1268,8 +1378,10 @@ void get_dataset_view(cuvsCagraIndex_t index, DLManagedTensor* dataset) "cuvsCagraIndexGetDataset: host indices are allowed", [&](auto& idx) { using index_dataset_view_t = std::remove_cvref_t; - if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { - RAFT_FAIL("cuvsCagraIndexGetDataset does not expose VPQ datasets as dense DLPack tensors"); + if constexpr (cuvs::neighbors::is_vpq_dataset_view_v || + cuvs::neighbors::is_bbq_dataset_view_v) { + RAFT_FAIL( + "cuvsCagraIndexGetDataset does not expose quantized datasets as dense DLPack tensors"); } else { cuvs::core::to_dlpack(idx.dataset().view(), dataset); } @@ -1654,6 +1766,46 @@ extern "C" cuvsError_t cuvsDatasetMakeStandardView(cuvsResources_t res, }); } +extern "C" cuvsError_t cuvsDatasetMakeBbqView(cuvsResources_t, + cuvsBbqQuantizer_t* quantizers, + size_t num_quantizers, + cuvsDataset_t* dataset) +{ + return cuvs::core::translate_exceptions([=] { + RAFT_EXPECTS(dataset != nullptr, "cuvsDatasetMakeBbqView: null output dataset"); + *dataset = nullptr; + RAFT_EXPECTS(quantizers != nullptr, "cuvsDatasetMakeBbqView: null quantizer array"); + RAFT_EXPECTS(num_quantizers == 1 || num_quantizers == 2, + "cuvsDatasetMakeBbqView: expected one or two quantizers"); + RAFT_EXPECTS(quantizers[0] != nullptr && quantizers[0]->addr != 0, + "cuvsDatasetMakeBbqView: first quantizer has null storage"); + + auto dtype = quantizers[0]->dtype; + for (size_t i = 1; i < num_quantizers; ++i) { + RAFT_EXPECTS(quantizers[i] != nullptr && quantizers[i]->addr != 0, + "cuvsDatasetMakeBbqView: quantizer %zu has null storage", + i); + RAFT_EXPECTS(quantizers[i]->dtype.code == dtype.code && + quantizers[i]->dtype.bits == dtype.bits && + quantizers[i]->dtype.lanes == dtype.lanes, + "cuvsDatasetMakeBbqView: all quantizers must use the same dtype"); + } + if (dtype.code == kDLFloat && dtype.bits == 32) { + make_and_bind_bbq_dataset(quantizers, num_quantizers, dtype, dataset); + } else if (dtype.code == kDLFloat && dtype.bits == 16) { + make_and_bind_bbq_dataset(quantizers, num_quantizers, dtype, dataset); + } else if (dtype.code == kDLInt && dtype.bits == 8) { + make_and_bind_bbq_dataset(quantizers, num_quantizers, dtype, dataset); + } else if (dtype.code == kDLUInt && dtype.bits == 8) { + make_and_bind_bbq_dataset(quantizers, num_quantizers, dtype, dataset); + } else { + RAFT_FAIL("cuvsDatasetMakeBbqView: unsupported centroid dtype: code=%d, bits=%d", + dtype.code, + dtype.bits); + } + }); +} + extern "C" cuvsError_t cuvsDatasetMakePQ(cuvsResources_t res, cuvsPqParams_t params, cuvsDataset_t dataset, @@ -1766,6 +1918,17 @@ static void build_dispatch_on_mem_type_and_layout(raft::resources* res_ptr, cuvsDataset_t dataset, cuvsCagraIndex_t index) { + if (dataset->layout == CUVS_DATASET_LAYOUT_BBQ) { + RAFT_EXPECTS(dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE, + "cuvsCagraBuild: BBQ dataset must be device-resident"); + using owner_t = cuvs::neighbors::device_bbq_dataset; + using view_t = cuvs::neighbors::device_bbq_dataset_view; + with_dataset_view(dataset, [&](auto const& view) { + build_index_from_dataset_view(res_ptr, params, view, index); + }); + return; + } + auto const is_padded = dataset->layout == CUVS_DATASET_LAYOUT_PADDED; if (dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE) { @@ -1925,6 +2088,9 @@ extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res, auto index = *index_c_ptr; auto* box = reinterpret_cast(index.addr); RAFT_EXPECTS(box != nullptr, "cuvsCagraSearch: null index handle"); + RAFT_EXPECTS(box->layout != sg_cagra_c_api_index_box::dataset_layout::device_bbq, + "cuvsCagraSearch: attach a device padded or PQ dataset to the BBQ-built index " + "with cuvsCagraUpdateDataset before searching"); RAFT_EXPECTS(box->layout == sg_cagra_c_api_index_box::dataset_layout::device_padded || box->layout == sg_cagra_c_api_index_box::dataset_layout::device_vpq, "cuvsCagraSearch: index must be device-padded or VPQ. For standard indices, call " diff --git a/c/src/preprocessing/quantize/bbq.cpp b/c/src/preprocessing/quantize/bbq.cpp new file mode 100644 index 0000000000..a2415c661d --- /dev/null +++ b/c/src/preprocessing/quantize/bbq.cpp @@ -0,0 +1,243 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include "../../core/exceptions.hpp" +#include "../../core/interop.hpp" + +#include +#include + +#include +#include + +namespace { + +template +void destroy_typed_addr(void* ptr) +{ + delete reinterpret_cast(ptr); +} + +auto to_cpp_layout(cuvsBbqCodeLayout_t layout) + -> cuvs::preprocessing::quantize::bbq::bbq_code_layout +{ + using layout_t = cuvs::preprocessing::quantize::bbq::bbq_code_layout; + switch (layout) { + case CUVS_BBQ_CODE_LAYOUT_PACKED_1B: return layout_t::packed_1b; + case CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_2B: return layout_t::transposed_2b; + case CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_4B: return layout_t::transposed_4b; + case CUVS_BBQ_CODE_LAYOUT_PACKED_4B: return layout_t::packed_4b; + case CUVS_BBQ_CODE_LAYOUT_PACKED_7B: return layout_t::packed_7b; + case CUVS_BBQ_CODE_LAYOUT_PACKED_8B: return layout_t::packed_8b; + } + RAFT_FAIL("cuvsBbqQuantizerCreateView: invalid BBQ code layout"); +} + +template +auto make_cpp_quantizer_view( + DLManagedTensor* codes_tensor, + DLManagedTensor* lower_intervals_tensor, + DLManagedTensor* upper_intervals_tensor, + DLManagedTensor* additional_corrections_tensor, + DLManagedTensor* quantized_component_sums_tensor, + DLManagedTensor* centroid_tensor, + DLManagedTensor* dequant_delta_tensor, + DLManagedTensor* dequant_sum_delta_tensor, + DLManagedTensor* row_norm_tensor, + cuvsBbqCodeLayout_t c_layout, + cuvsDistanceType metric, + float centroid_norm_sq) + -> cuvs::preprocessing::quantize::bbq::quantizer_view +{ + using quantizer_view_t = cuvs::preprocessing::quantize::bbq::quantizer_view; + using codes_view_t = raft::device_matrix_view; + using float_view_t = raft::device_vector_view; + using int_view_t = raft::device_vector_view; + using centroid_view_t = raft::device_vector_view; + + auto codes = cuvs::core::from_dlpack(codes_tensor); + auto lower = cuvs::core::from_dlpack(lower_intervals_tensor); + auto upper = cuvs::core::from_dlpack(upper_intervals_tensor); + auto corrections = cuvs::core::from_dlpack(additional_corrections_tensor); + auto sums = cuvs::core::from_dlpack(quantized_component_sums_tensor); + auto centroid = cuvs::core::from_dlpack(centroid_tensor); + auto delta = cuvs::core::from_dlpack(dequant_delta_tensor); + auto sum_delta = cuvs::core::from_dlpack(dequant_sum_delta_tensor); + auto row_norm = cuvs::core::from_dlpack(row_norm_tensor); + auto layout = to_cpp_layout(c_layout); + + const auto n_rows = codes.extent(0); + const auto dim = static_cast(centroid.extent(0)); + RAFT_EXPECTS(n_rows > 0 && dim > 0, "cuvsBbqQuantizerCreateView: quantizer must not be empty"); + RAFT_EXPECTS( + codes.extent(1) == + static_cast( + cuvs::preprocessing::quantize::bbq::get_encoded_row_length(dim, layout)), + "cuvsBbqQuantizerCreateView: code row length does not match the dimension and layout"); + RAFT_EXPECTS(lower.extent(0) == n_rows && upper.extent(0) == n_rows && + corrections.extent(0) == n_rows && sums.extent(0) == n_rows && + delta.extent(0) == n_rows && sum_delta.extent(0) == n_rows && + row_norm.extent(0) == n_rows, + "cuvsBbqQuantizerCreateView: every correction vector must contain one value per row"); + + return quantizer_view_t{codes, + lower, + upper, + corrections, + sums, + centroid, + delta, + sum_delta, + row_norm, + layout, + static_cast(metric), + centroid_norm_sq}; +} + +template +void make_and_bind_quantizer_view( + cuvsBbqQuantizer_t* output, + DLDataType dtype, + DLManagedTensor* codes, + DLManagedTensor* lower_intervals, + DLManagedTensor* upper_intervals, + DLManagedTensor* additional_corrections, + DLManagedTensor* quantized_component_sums, + DLManagedTensor* centroid, + DLManagedTensor* dequant_delta, + DLManagedTensor* dequant_sum_delta, + DLManagedTensor* row_norm, + cuvsBbqCodeLayout_t layout, + cuvsDistanceType metric, + float centroid_norm_sq) +{ + using view_t = cuvs::preprocessing::quantize::bbq::quantizer_view; + auto view = std::make_unique(make_cpp_quantizer_view(codes, + lower_intervals, + upper_intervals, + additional_corrections, + quantized_component_sums, + centroid, + dequant_delta, + dequant_sum_delta, + row_norm, + layout, + metric, + centroid_norm_sq)); + auto handle = std::make_unique(); + handle->addr = reinterpret_cast(view.release()); + handle->destroy_addr = &destroy_typed_addr; + handle->dtype = dtype; + handle->is_owning = false; + *output = handle.release(); +} + +} // namespace + +extern "C" cuvsError_t cuvsBbqQuantizerCreateView( + DLManagedTensor* codes, + DLManagedTensor* lower_intervals, + DLManagedTensor* upper_intervals, + DLManagedTensor* additional_corrections, + DLManagedTensor* quantized_component_sums, + DLManagedTensor* centroid, + DLManagedTensor* dequant_delta, + DLManagedTensor* dequant_sum_delta, + DLManagedTensor* row_norm, + cuvsBbqCodeLayout_t layout, + cuvsDistanceType metric, + float centroid_norm_sq, + cuvsBbqQuantizer_t* quantizer) +{ + return cuvs::core::translate_exceptions([=] { + RAFT_EXPECTS(quantizer != nullptr, "cuvsBbqQuantizerCreateView: null output"); + *quantizer = nullptr; + RAFT_EXPECTS(codes != nullptr && lower_intervals != nullptr && upper_intervals != nullptr && + additional_corrections != nullptr && quantized_component_sums != nullptr && + centroid != nullptr && dequant_delta != nullptr && + dequant_sum_delta != nullptr && row_norm != nullptr, + "cuvsBbqQuantizerCreateView: null tensor"); + auto dtype = centroid->dl_tensor.dtype; + if (dtype.code == kDLFloat && dtype.bits == 32) { + make_and_bind_quantizer_view(quantizer, + dtype, + codes, + lower_intervals, + upper_intervals, + additional_corrections, + quantized_component_sums, + centroid, + dequant_delta, + dequant_sum_delta, + row_norm, + layout, + metric, + centroid_norm_sq); + } else if (dtype.code == kDLFloat && dtype.bits == 16) { + make_and_bind_quantizer_view(quantizer, + dtype, + codes, + lower_intervals, + upper_intervals, + additional_corrections, + quantized_component_sums, + centroid, + dequant_delta, + dequant_sum_delta, + row_norm, + layout, + metric, + centroid_norm_sq); + } else if (dtype.code == kDLInt && dtype.bits == 8) { + make_and_bind_quantizer_view(quantizer, + dtype, + codes, + lower_intervals, + upper_intervals, + additional_corrections, + quantized_component_sums, + centroid, + dequant_delta, + dequant_sum_delta, + row_norm, + layout, + metric, + centroid_norm_sq); + } else if (dtype.code == kDLUInt && dtype.bits == 8) { + make_and_bind_quantizer_view(quantizer, + dtype, + codes, + lower_intervals, + upper_intervals, + additional_corrections, + quantized_component_sums, + centroid, + dequant_delta, + dequant_sum_delta, + row_norm, + layout, + metric, + centroid_norm_sq); + } else { + RAFT_FAIL("cuvsBbqQuantizerCreateView: unsupported centroid dtype: code=%d, bits=%d", + dtype.code, + dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsBbqQuantizerDestroy(cuvsBbqQuantizer_t quantizer) +{ + return cuvs::core::translate_exceptions([=] { + if (quantizer == nullptr) { return; } + if (quantizer->destroy_addr != nullptr && quantizer->addr != 0) { + quantizer->destroy_addr(reinterpret_cast(quantizer->addr)); + } + delete quantizer; + }); +} diff --git a/c/tests/core/dataset_c.cu b/c/tests/core/dataset_c.cu index 0d05c497ff..bebf3f3196 100644 --- a/c/tests/core/dataset_c.cu +++ b/c/tests/core/dataset_c.cu @@ -3,13 +3,17 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include "../../src/core/interop.hpp" + #include #include #include +#include #include #include #include +#include #include #include @@ -37,6 +41,27 @@ struct MatrixTensor { } }; +template +auto make_device_matrix_tensor(T* data, int64_t rows, int64_t columns) -> DLManagedTensor +{ + DLManagedTensor tensor{}; + cuvs::core::to_dlpack(raft::make_device_matrix_view(data, rows, columns), &tensor); + return tensor; +} + +template +auto make_device_vector_tensor(T* data, int64_t size) -> DLManagedTensor +{ + DLManagedTensor tensor{}; + cuvs::core::to_dlpack(raft::make_device_vector_view(data, size), &tensor); + return tensor; +} + +void free_tensor(DLManagedTensor& tensor) +{ + if (tensor.deleter != nullptr) { tensor.deleter(&tensor); } +} + } // namespace TEST(DatasetC, CreateDestroy) @@ -213,3 +238,109 @@ TEST(DatasetC, MakeHostPaddedFromDevice) ASSERT_EQ(cuvsDatasetDestroy(padded), CUVS_SUCCESS); ASSERT_EQ(cuvsResourcesDestroy(res), CUVS_SUCCESS); } + +TEST(DatasetC, MakeBbqView) +{ + constexpr int64_t n_rows = 4; + constexpr int64_t dim = 8; + + cuvsResources_t res; + ASSERT_EQ(cuvsResourcesCreate(&res), CUVS_SUCCESS); + cudaStream_t stream; + ASSERT_EQ(cuvsStreamGet(res, &stream), CUVS_SUCCESS); + + rmm::device_uvector codes_1b(n_rows, stream); + rmm::device_uvector lower(n_rows, stream); + rmm::device_uvector upper(n_rows, stream); + rmm::device_uvector corrections(n_rows, stream); + rmm::device_uvector sums(n_rows, stream); + rmm::device_uvector centroid(dim, stream); + rmm::device_uvector delta(n_rows, stream); + rmm::device_uvector sum_delta(n_rows, stream); + rmm::device_uvector row_norm(n_rows, stream); + + auto codes_tensor = make_device_matrix_tensor(codes_1b.data(), n_rows, 1); + auto lower_tensor = make_device_vector_tensor(lower.data(), n_rows); + auto upper_tensor = make_device_vector_tensor(upper.data(), n_rows); + auto corrections_tensor = make_device_vector_tensor(corrections.data(), n_rows); + auto sums_tensor = make_device_vector_tensor(sums.data(), n_rows); + auto centroid_tensor = make_device_vector_tensor(centroid.data(), dim); + auto delta_tensor = make_device_vector_tensor(delta.data(), n_rows); + auto sum_delta_tensor = make_device_vector_tensor(sum_delta.data(), n_rows); + auto row_norm_tensor = make_device_vector_tensor(row_norm.data(), n_rows); + + cuvsBbqQuantizer_t quantizer; + ASSERT_EQ(cuvsBbqQuantizerCreateView(&codes_tensor, + &lower_tensor, + &upper_tensor, + &corrections_tensor, + &sums_tensor, + ¢roid_tensor, + &delta_tensor, + &sum_delta_tensor, + &row_norm_tensor, + CUVS_BBQ_CODE_LAYOUT_PACKED_1B, + L2Expanded, + 0.0f, + &quantizer), + CUVS_SUCCESS); + + cuvsDataset_t dataset; + ASSERT_EQ(cuvsDatasetMakeBbqView(res, &quantizer, 1, &dataset), CUVS_SUCCESS); + + cuvsDatasetLayout_t layout; + cuvsDatasetMemType_t mem_type; + bool is_owning; + DLDataType dtype; + ASSERT_EQ(cuvsDatasetGetLayout(dataset, &layout), CUVS_SUCCESS); + ASSERT_EQ(cuvsDatasetGetMemType(dataset, &mem_type), CUVS_SUCCESS); + ASSERT_EQ(cuvsDatasetGetIsOwning(dataset, &is_owning), CUVS_SUCCESS); + ASSERT_EQ(cuvsDatasetGetDtype(dataset, &dtype), CUVS_SUCCESS); + EXPECT_EQ(layout, CUVS_DATASET_LAYOUT_BBQ); + EXPECT_EQ(mem_type, CUVS_DATASET_MEM_TYPE_DEVICE); + EXPECT_FALSE(is_owning); + EXPECT_EQ(dtype.code, kDLFloat); + EXPECT_EQ(dtype.bits, 32); + + ASSERT_EQ(cuvsDatasetDestroy(dataset), CUVS_SUCCESS); + + rmm::device_uvector codes_4t(n_rows * 4, stream); + auto codes_4t_tensor = make_device_matrix_tensor(codes_4t.data(), n_rows, 4); + cuvsBbqQuantizer_t quantizer_4t; + ASSERT_EQ(cuvsBbqQuantizerCreateView(&codes_4t_tensor, + &lower_tensor, + &upper_tensor, + &corrections_tensor, + &sums_tensor, + ¢roid_tensor, + &delta_tensor, + &sum_delta_tensor, + &row_norm_tensor, + CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_4B, + L2Expanded, + 0.0f, + &quantizer_4t), + CUVS_SUCCESS); + + cuvsBbqQuantizer_t quantizers[] = {quantizer, quantizer_4t}; + cuvsDataset_t asymmetric_dataset; + ASSERT_EQ(cuvsDatasetMakeBbqView(res, quantizers, 2, &asymmetric_dataset), CUVS_SUCCESS); + ASSERT_EQ(cuvsDatasetGetLayout(asymmetric_dataset, &layout), CUVS_SUCCESS); + EXPECT_EQ(layout, CUVS_DATASET_LAYOUT_BBQ); + + ASSERT_EQ(cuvsDatasetDestroy(asymmetric_dataset), CUVS_SUCCESS); + ASSERT_EQ(cuvsBbqQuantizerDestroy(quantizer_4t), CUVS_SUCCESS); + ASSERT_EQ(cuvsBbqQuantizerDestroy(quantizer), CUVS_SUCCESS); + ASSERT_EQ(cuvsResourcesDestroy(res), CUVS_SUCCESS); + + free_tensor(codes_tensor); + free_tensor(codes_4t_tensor); + free_tensor(lower_tensor); + free_tensor(upper_tensor); + free_tensor(corrections_tensor); + free_tensor(sums_tensor); + free_tensor(centroid_tensor); + free_tensor(delta_tensor); + free_tensor(sum_delta_tensor); + free_tensor(row_norm_tensor); +} diff --git a/c/tests/neighbors/ann_cagra_c.cu b/c/tests/neighbors/ann_cagra_c.cu index ba0e3ee310..e560962a42 100644 --- a/c/tests/neighbors/ann_cagra_c.cu +++ b/c/tests/neighbors/ann_cagra_c.cu @@ -4,10 +4,13 @@ */ #include "test_utils.cuh" +#include "../../src/core/interop.hpp" #include #include #include +#include #include +#include #include #include @@ -36,6 +39,31 @@ #include +namespace { + +template +auto make_device_matrix_tensor(T* data, int64_t rows, int64_t columns) -> DLManagedTensor +{ + DLManagedTensor tensor{}; + cuvs::core::to_dlpack(raft::make_device_matrix_view(data, rows, columns), &tensor); + return tensor; +} + +template +auto make_device_vector_tensor(T* data, int64_t size) -> DLManagedTensor +{ + DLManagedTensor tensor{}; + cuvs::core::to_dlpack(raft::make_device_vector_view(data, size), &tensor); + return tensor; +} + +void free_tensor(DLManagedTensor& tensor) +{ + if (tensor.deleter != nullptr) { tensor.deleter(&tensor); } +} + +} // namespace + float dataset[4][2] = {{0.74021935, 0.9209938}, {0.03902049, 0.9689629}, {0.92514056, 0.4463501}, @@ -179,6 +207,153 @@ TEST(CagraC, BuildSearch) cuvsResourcesDestroy(res); } +TEST(CagraC, BuildBbqGraphAndAttachDataset) +{ + constexpr int64_t n_rows = 64; + constexpr int64_t dim = 32; + + cuvsResources_t res; + ASSERT_EQ(cuvsResourcesCreate(&res), CUVS_SUCCESS); + cudaStream_t stream; + ASSERT_EQ(cuvsStreamGet(res, &stream), CUVS_SUCCESS); + + rmm::device_uvector codes(n_rows * dim, stream); + rmm::device_uvector lower(n_rows, stream); + rmm::device_uvector upper(n_rows, stream); + rmm::device_uvector corrections(n_rows, stream); + rmm::device_uvector sums(n_rows, stream); + rmm::device_uvector centroid(dim, stream); + rmm::device_uvector delta(n_rows, stream); + rmm::device_uvector sum_delta(n_rows, stream); + rmm::device_uvector row_norm(n_rows, stream); + rmm::device_uvector dense(n_rows * dim, stream); + ASSERT_EQ(cudaMemsetAsync(codes.data(), 0, codes.size(), stream), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(lower.data(), 0, lower.size() * sizeof(float), stream), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(upper.data(), 0, upper.size() * sizeof(float), stream), cudaSuccess); + ASSERT_EQ( + cudaMemsetAsync(corrections.data(), 0, corrections.size() * sizeof(float), stream), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(sums.data(), 0, sums.size() * sizeof(int32_t), stream), cudaSuccess); + ASSERT_EQ( + cudaMemsetAsync(centroid.data(), 0, centroid.size() * sizeof(float), stream), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(delta.data(), 0, delta.size() * sizeof(float), stream), cudaSuccess); + ASSERT_EQ( + cudaMemsetAsync(sum_delta.data(), 0, sum_delta.size() * sizeof(float), stream), cudaSuccess); + ASSERT_EQ( + cudaMemsetAsync(row_norm.data(), 0, row_norm.size() * sizeof(float), stream), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(dense.data(), 0, dense.size() * sizeof(float), stream), cudaSuccess); + + auto codes_tensor = make_device_matrix_tensor(codes.data(), n_rows, dim); + auto lower_tensor = make_device_vector_tensor(lower.data(), n_rows); + auto upper_tensor = make_device_vector_tensor(upper.data(), n_rows); + auto corrections_tensor = make_device_vector_tensor(corrections.data(), n_rows); + auto sums_tensor = make_device_vector_tensor(sums.data(), n_rows); + auto centroid_tensor = make_device_vector_tensor(centroid.data(), dim); + auto delta_tensor = make_device_vector_tensor(delta.data(), n_rows); + auto sum_delta_tensor = make_device_vector_tensor(sum_delta.data(), n_rows); + auto row_norm_tensor = make_device_vector_tensor(row_norm.data(), n_rows); + auto dense_tensor = make_device_matrix_tensor(dense.data(), n_rows, dim); + + cuvsBbqQuantizer_t quantizer; + ASSERT_EQ(cuvsBbqQuantizerCreateView(&codes_tensor, + &lower_tensor, + &upper_tensor, + &corrections_tensor, + &sums_tensor, + ¢roid_tensor, + &delta_tensor, + &sum_delta_tensor, + &row_norm_tensor, + CUVS_BBQ_CODE_LAYOUT_PACKED_8B, + L2Expanded, + 0.0f, + &quantizer), + CUVS_SUCCESS); + + cuvsDataset_t bbq_dataset; + ASSERT_EQ(cuvsDatasetMakeBbqView(res, &quantizer, 1, &bbq_dataset), CUVS_SUCCESS); + + cuvsCagraIndexParams_t params; + ASSERT_EQ(cuvsCagraIndexParamsCreate(¶ms), CUVS_SUCCESS); + params->metric = L2Expanded; + params->build_algo = NN_DESCENT; + params->graph_degree = 8; + params->intermediate_graph_degree = 16; + params->nn_descent_niter = 5; + + cuvsCagraIndex_t index; + ASSERT_EQ(cuvsCagraIndexCreate(&index), CUVS_SUCCESS); + ASSERT_EQ(cuvsCagraBuild(res, params, bbq_dataset, index), CUVS_SUCCESS); + + DLManagedTensor graph_tensor{}; + ASSERT_EQ(cuvsCagraIndexGetGraph(index, &graph_tensor), CUVS_SUCCESS); + ASSERT_EQ(graph_tensor.dl_tensor.shape[0], n_rows); + ASSERT_EQ(graph_tensor.dl_tensor.shape[1], params->graph_degree); + std::vector graph(n_rows * params->graph_degree); + raft::copy(graph.data(), + static_cast(graph_tensor.dl_tensor.data), + graph.size(), + stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + for (auto neighbor : graph) { + EXPECT_LT(neighbor, static_cast(n_rows)); + } + + rmm::device_uvector query(dim, stream); + rmm::device_uvector neighbors(1, stream); + rmm::device_uvector distances(1, stream); + ASSERT_EQ(cudaMemsetAsync(query.data(), 0, query.size() * sizeof(float), stream), cudaSuccess); + auto query_tensor = make_device_matrix_tensor(query.data(), 1, dim); + auto neighbors_tensor = make_device_matrix_tensor(neighbors.data(), 1, 1); + auto distances_tensor = make_device_matrix_tensor(distances.data(), 1, 1); + cuvsCagraSearchParams_t search_params; + ASSERT_EQ(cuvsCagraSearchParamsCreate(&search_params), CUVS_SUCCESS); + cuvsFilter filter{.addr = 0, .type = NO_FILTER}; + EXPECT_EQ(cuvsCagraSearch(res, + search_params, + index, + &query_tensor, + &neighbors_tensor, + &distances_tensor, + filter), + CUVS_ERROR); + EXPECT_NE(std::string(cuvsGetLastErrorText()).find("BBQ-built"), std::string::npos); + + cuvsDataset_t padded_dataset; + ASSERT_EQ(cuvsDatasetMakePaddedView(res, &dense_tensor, &padded_dataset), CUVS_SUCCESS); + ASSERT_EQ(cuvsCagraUpdateDataset(res, padded_dataset, index), CUVS_SUCCESS); + ASSERT_EQ(cuvsCagraSearch(res, + search_params, + index, + &query_tensor, + &neighbors_tensor, + &distances_tensor, + filter), + CUVS_SUCCESS); + + ASSERT_EQ(cuvsCagraSearchParamsDestroy(search_params), CUVS_SUCCESS); + if (graph_tensor.deleter != nullptr) { graph_tensor.deleter(&graph_tensor); } + ASSERT_EQ(cuvsCagraIndexDestroy(index), CUVS_SUCCESS); + ASSERT_EQ(cuvsDatasetDestroy(padded_dataset), CUVS_SUCCESS); + ASSERT_EQ(cuvsCagraIndexParamsDestroy(params), CUVS_SUCCESS); + ASSERT_EQ(cuvsDatasetDestroy(bbq_dataset), CUVS_SUCCESS); + ASSERT_EQ(cuvsBbqQuantizerDestroy(quantizer), CUVS_SUCCESS); + ASSERT_EQ(cuvsResourcesDestroy(res), CUVS_SUCCESS); + + free_tensor(codes_tensor); + free_tensor(lower_tensor); + free_tensor(upper_tensor); + free_tensor(corrections_tensor); + free_tensor(sums_tensor); + free_tensor(centroid_tensor); + free_tensor(delta_tensor); + free_tensor(sum_delta_tensor); + free_tensor(row_norm_tensor); + free_tensor(dense_tensor); + free_tensor(query_tensor); + free_tensor(neighbors_tensor); + free_tensor(distances_tensor); +} + // CAGRA operations that need a search-ready index must reject host / non-device-padded // datasets rather than succeeding and producing undefined behavior. TEST(CagraC, DatasetContractFailures) diff --git a/ci/validate_wheel.sh b/ci/validate_wheel.sh index 8ed1b8cd8b..30de27cfc8 100755 --- a/ci/validate_wheel.sh +++ b/ci/validate_wheel.sh @@ -21,11 +21,11 @@ PYDISTCHECK_ARGS=( if [[ "${package_dir}" == "python/libcuvs" ]]; then if [[ "${RAPIDS_CUDA_MAJOR}" == "12" ]]; then PYDISTCHECK_ARGS+=( - --max-allowed-size-compressed '365Mi' + --max-allowed-size-compressed '375Mi' ) else PYDISTCHECK_ARGS+=( - --max-allowed-size-compressed '225Mi' + --max-allowed-size-compressed '235Mi' ) fi elif [[ "${package_dir}" != "python/cuvs" ]]; then diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 1beaf2484f..ef9977f72c 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1533,6 +1533,7 @@ if(NOT BUILD_CPU_ONLY) src/preprocessing/quantize/scalar.cu src/preprocessing/quantize/binary.cu src/preprocessing/quantize/pq.cu + src/preprocessing/quantize/bbq.cu src/preprocessing/spectral/spectral_embedding.cu src/preprocessing/pca/pca.cu ${select_k_inst_files} @@ -1838,9 +1839,10 @@ endif() # ################################################################################################## # * build test executable ---------------------------------------------------- +add_subdirectory(internal) + if(BUILD_TESTS) enable_testing() - add_subdirectory(internal) add_subdirectory(tests) endif() diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 90d49bfe16..b17fb8165a 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -956,6 +957,10 @@ template using device_pq_index = index>; +/** CAGRA index with a device-resident BBQ-quantized dataset. */ +template +using device_bbq_index = index>; + /** Index type returned by `cagra::build(res, params, dataset_view)`. */ template using cagra_index_t = @@ -1214,6 +1219,49 @@ auto build(raft::resources const& res, cuvs::neighbors::host_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_standard_index; +/** + * @brief Build from a device BBQ-quantized dataset view. + * + * The kNN graph is built from the quantized codes alone, so the uncompressed vectors are never + * needed and peak memory is driven by the code size. Only nn-descent graph construction is + * available (IVF-PQ, iterative CAGRA search, and ACE all read uncompressed vectors), and the + * metric must be one of L2Expanded, L2SqrtExpanded, CosineExpanded, or InnerProduct and must match + * the metric the quantizer corrections were generated for. + * + * The returned index cannot be searched: CAGRA has no BBQ search kernels. Call the type-changing + * `update_dataset` with an uncompressed device-padded dataset to search the resulting graph. + * + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset device BBQ dataset view [n_rows, dim] + * @return built `device_bbq_index` + */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::device_bbq_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_bbq_index; + +/** @copydoc build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, + * cuvs::neighbors::device_bbq_dataset_view const& dataset) */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::device_bbq_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_bbq_index; + +/** @copydoc build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, + * cuvs::neighbors::device_bbq_dataset_view const& dataset) */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::device_bbq_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_bbq_index; + +/** @copydoc build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, + * cuvs::neighbors::device_bbq_dataset_view const& dataset) */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::device_bbq_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_bbq_index; + /** * @} */ @@ -2965,6 +3013,75 @@ void deserialize(raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::device_pq_index* index); +/* device_bbq_index graph-only overloads (CAGRA + BBQ). + * + * These overloads persist the graph and index metadata, but not the attached BBQ dataset + * Attach a compatible dataset with `update_dataset` after deserialization before searching. + */ +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_bbq_index* index); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_bbq_index* index); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_bbq_index* index); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_bbq_index* index); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_bbq_index* index); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_bbq_index* index); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_bbq_index* index); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_bbq_index* index); + /** @copydoc serialize */ void serialize(raft::resources const& handle, const std::string& filename, @@ -4948,6 +5065,48 @@ auto update_dataset(raft::resources const& res, device_vpq_dataset_view dataset) -> index>; +auto update_dataset(raft::resources const& res, + index>&& cagra_index, + device_padded_dataset_view dataset) + -> index>; + +auto update_dataset(raft::resources const& res, + index>&& cagra_index, + device_padded_dataset_view dataset) + -> index>; + +auto update_dataset(raft::resources const& res, + index>&& cagra_index, + device_padded_dataset_view dataset) + -> index>; + +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + device_padded_dataset_view dataset) + -> index>; + +auto update_dataset(raft::resources const& res, + index>&& cagra_index, + device_vpq_dataset_view dataset) + -> index>; + +auto update_dataset(raft::resources const& res, + index>&& cagra_index, + device_vpq_dataset_view dataset) + -> index>; + +auto update_dataset(raft::resources const& res, + index>&& cagra_index, + device_vpq_dataset_view dataset) + -> index>; + +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + device_vpq_dataset_view dataset) + -> index>; + } // namespace cagra } // namespace neighbors } // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 0db1485a2f..baa5b44da2 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -779,6 +779,7 @@ enum class dataset_view_kind { standard, vpq_f16, vpq_f32, + bbq, }; /** Primary template returns `unknown` so traits safely return `false` for non-dataset-view types. diff --git a/cpp/include/cuvs/neighbors/nn_descent.hpp b/cpp/include/cuvs/neighbors/nn_descent.hpp index 4c031049e2..1e1546c651 100644 --- a/cpp/include/cuvs/neighbors/nn_descent.hpp +++ b/cpp/include/cuvs/neighbors/nn_descent.hpp @@ -1,11 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include +#include #include #include @@ -523,6 +524,58 @@ auto build(raft::resources const& res, std::optional> graph = std::nullopt) -> cuvs::neighbors::nn_descent::index; +/** + * @brief Build an NN-Descent index directly from a device-resident BBQ dataset. + * + * Symmetric compressed-code distances are used during graph construction. Supported metrics are + * L2Expanded, L2SqrtExpanded, CosineExpanded, and InnerProduct. The dataset's correction terms + * must have been generated for the selected metric. + * + * @param res raft resources + * @param params NN-Descent build parameters + * @param dataset BBQ codes and correction terms in device memory + * @param graph optional caller-owned host graph + * @return index containing the all-neighbors graph + */ +auto build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::device_bbq_dataset_view dataset, + std::optional> graph = + std::nullopt) -> cuvs::neighbors::nn_descent::index; +/** + * @copydoc build(raft::resources const& res, index_params const& params, + * cuvs::neighbors::device_bbq_dataset_view dataset, + * std::optional> graph = std::nullopt) + * -> cuvs::neighbors::nn_descent::index + */ +auto build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::device_bbq_dataset_view dataset, + std::optional> graph = + std::nullopt) -> cuvs::neighbors::nn_descent::index; +/** + * @copydoc build(raft::resources const& res, index_params const& params, + * cuvs::neighbors::device_bbq_dataset_view dataset, + * std::optional> graph = std::nullopt) + * -> cuvs::neighbors::nn_descent::index + */ +auto build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::device_bbq_dataset_view dataset, + std::optional> graph = + std::nullopt) -> cuvs::neighbors::nn_descent::index; +/** + * @copydoc build(raft::resources const& res, index_params const& params, + * cuvs::neighbors::device_bbq_dataset_view dataset, + * std::optional> graph = std::nullopt) + * -> cuvs::neighbors::nn_descent::index + */ +auto build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::device_bbq_dataset_view dataset, + std::optional> graph = + std::nullopt) -> cuvs::neighbors::nn_descent::index; + /** @} */ /** * @brief Test if we have enough GPU memory to run NN descent algorithm. diff --git a/cpp/include/cuvs/preprocessing/quantize/bbq.hpp b/cpp/include/cuvs/preprocessing/quantize/bbq.hpp new file mode 100644 index 0000000000..21366d47d0 --- /dev/null +++ b/cpp/include/cuvs/preprocessing/quantize/bbq.hpp @@ -0,0 +1,388 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace CUVS_EXPORT cuvs { + +namespace preprocessing::quantize::bbq { + +/** + * @defgroup bbq Better Binary Quantization utilities + * @{ + */ + +/** + * Storage layout of BBQ/OSQ quantized component codes in each dataset row. + * packed_1b: Each dimension is quantized to a single bit and packed into bytes. Reflects + * Lucene's OptimizedScalarQuantizer.packAsBinary. + * transposed_2b: Each dimension is quantized to 2 bits, stored as 2 bitplanes. + * Reflects Lucene's OptimizedScalarQuantizer.transposeDibit. SIMT popc path only + * (paired with a transposed_4b or packed_1b operand); + * transposed_4b: Each dimension is quantized to 4 bits, optimized for bitwise operations. + * Reflects Lucene's OptimizedScalarQuantizer.transposeHalfByte. the first bit of + * every dimension is in the first set dimensions bits, or (dimensions/8) + * bytes. The second, third, and fourth bits are in the second, third, and + * fourth set of dimensions bits, respectively. Format used for queries. + * packed_4b: Each dimension is quantized to 4 bits, two values are packed into each output + * byte. + * packed_7b: Each dimension is quantized to 7 bits and treated as a signed value. + * packed_8b: Each dimension is quantized to 8 bits and treated as an unsigned value. + */ +enum class bbq_code_layout { + packed_1b, + transposed_2b, + transposed_4b, + packed_4b, + packed_7b, + packed_8b, +}; + +/** + * Bit width of a layout. + */ +constexpr auto get_bit_width(bbq_code_layout layout) noexcept -> uint32_t +{ + switch (layout) { + case bbq_code_layout::packed_1b: return 1; + case bbq_code_layout::transposed_2b: return 2; + case bbq_code_layout::transposed_4b: + case bbq_code_layout::packed_4b: return 4; + case bbq_code_layout::packed_7b: return 7; + case bbq_code_layout::packed_8b: return 8; + } + return 0; +} + +/** Bytes one row of @p dim components occupies once encoded in @p layout. */ +constexpr auto get_encoded_row_length(uint32_t dim, bbq_code_layout layout) noexcept -> uint32_t +{ + switch (layout) { + case bbq_code_layout::packed_1b: return raft::div_rounding_up_safe(dim, 8u); + case bbq_code_layout::transposed_2b: return 2 * raft::div_rounding_up_safe(dim, 8u); + case bbq_code_layout::transposed_4b: return 4 * raft::div_rounding_up_safe(dim, 8u); + case bbq_code_layout::packed_4b: return raft::div_rounding_up_safe(dim, 2u); + case bbq_code_layout::packed_7b: return dim; + case bbq_code_layout::packed_8b: return dim; + } + return 0; +} + +template +struct quantizer_view; + +/** + * @brief Better Binary Quantization + * ([BBQ](https://www.elastic.co/search-labs/blog/better-binary-quantization-lucene-elasticsearch)) + * is a vector-quantization approach used in Elasticsearch and Apache Lucene. It builds on ideas + * introduced in RaBitQ([Gao and Long](https://arxiv.org/pdf/2405.12497, [Gao et + * al.](https://arxiv.org/pdf/2409.09913)): residual binary codes around a centroid, corrective + * factors, and efficient bitwise comparison of codes at different bit widths. Lucene implements + * this as optimized scalar quantization (OSQ) with packed and bit-plane layouts; Elasticsearch + * exposes it as BBQ. + * + * BBQ in cuVS designed to be compatible with the Lucene/Elasticsearch dataset: a single shared + * centroid, no random rotation, and OSQ codes. + * + * RaBitQ and BBQ in cuVS both compress centroid-relative vectors to low-bit codes and retain + * additional per-vector information so search is better than naïve sign-bit comparison. They differ + * in transformation and scale representation. RaBitQ commonly separates residual magnitude from + * direction, then applies a random orthogonal rotation before binary coding; BBQ uses per-vector + * scalar intervals to interpret the compressed residual codes. + */ +template +struct quantizer { + raft::device_matrix codes; + raft::device_vector lower_intervals; + raft::device_vector upper_intervals; + raft::device_vector additional_corrections; + raft::device_vector quantized_component_sums; + raft::device_vector centroid; + /** Precomputed per-row dequantization factors, derived once (offline) from + * lower/upper_intervals and quantized_component_sums: dequant_delta = (upper-lower)/(2^bits-1) */ + raft::device_vector dequant_delta; + /** Precomputed per-row dequantization factors, derived once (offline) from dequant_delta and + * quantized_component_sums: dequant_sum_delta = dequant_delta * quantized_component_sums. */ + raft::device_vector dequant_sum_delta; + /** Squared norm of the row in original (un-centered) vector space, ||x||^2 */ + raft::device_vector row_norm; + + bbq_code_layout layout{bbq_code_layout::packed_1b}; + cuvs::distance::DistanceType metric{cuvs::distance::DistanceType::L2Expanded}; + float centroid_norm_sq{}; + + quantizer(raft::resources const& res, + IdxT n_rows, + uint32_t dim, + bbq_code_layout layout, + cuvs::distance::DistanceType metric) + : codes{raft::make_device_matrix( + res, n_rows, static_cast(get_encoded_row_length(dim, layout)))}, + lower_intervals{raft::make_device_vector(res, n_rows)}, + upper_intervals{raft::make_device_vector(res, n_rows)}, + additional_corrections{raft::make_device_vector(res, n_rows)}, + quantized_component_sums{raft::make_device_vector(res, n_rows)}, + centroid{raft::make_device_vector(res, static_cast(dim))}, + dequant_delta{raft::make_device_vector(res, n_rows)}, + dequant_sum_delta{raft::make_device_vector(res, n_rows)}, + row_norm{raft::make_device_vector(res, n_rows)}, + layout{layout}, + metric{metric} + { + } + + [[nodiscard]] auto n_rows() const noexcept -> IdxT { return codes.extent(0); } + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + return static_cast(centroid.extent(0)); + } + [[nodiscard]] constexpr auto encoded_row_length() const noexcept -> uint32_t + { + return get_encoded_row_length(dim(), layout); + } + [[nodiscard]] auto view() const noexcept -> quantizer_view + { + return quantizer_view{codes.view(), + lower_intervals.view(), + upper_intervals.view(), + additional_corrections.view(), + quantized_component_sums.view(), + centroid.view(), + dequant_delta.view(), + dequant_sum_delta.view(), + row_norm.view(), + layout, + metric, + centroid_norm_sq}; + } +}; + +template +struct quantizer_view { + raft::device_matrix_view codes; + raft::device_vector_view lower_intervals; + raft::device_vector_view upper_intervals; + raft::device_vector_view additional_corrections; + raft::device_vector_view quantized_component_sums; + raft::device_vector_view centroid; + raft::device_vector_view dequant_delta; + raft::device_vector_view dequant_sum_delta; + raft::device_vector_view row_norm; + + bbq_code_layout layout{bbq_code_layout::packed_1b}; + cuvs::distance::DistanceType metric{cuvs::distance::DistanceType::L2Expanded}; + float centroid_norm_sq{}; + + quantizer_view(raft::device_matrix_view codes_, + raft::device_vector_view lower_intervals_, + raft::device_vector_view upper_intervals_, + raft::device_vector_view additional_corrections_, + raft::device_vector_view quantized_component_sums_, + raft::device_vector_view centroid_, + raft::device_vector_view dequant_delta_, + raft::device_vector_view dequant_sum_delta_, + raft::device_vector_view row_norm_, + bbq_code_layout layout_, + cuvs::distance::DistanceType metric_, + float centroid_norm_sq_) noexcept + : codes{codes_}, + lower_intervals{lower_intervals_}, + upper_intervals{upper_intervals_}, + additional_corrections{additional_corrections_}, + quantized_component_sums{quantized_component_sums_}, + centroid{centroid_}, + dequant_delta{dequant_delta_}, + dequant_sum_delta{dequant_sum_delta_}, + row_norm{row_norm_}, + layout{layout_}, + metric{metric_}, + centroid_norm_sq{centroid_norm_sq_} + { + } + + [[nodiscard]] constexpr auto n_rows() const noexcept -> IdxT { return codes.extent(0); } + [[nodiscard]] constexpr auto dim() const noexcept -> uint32_t + { + return static_cast(centroid.extent(0)); + } +}; + +namespace helpers { +/** + * Derives dequant_delta from lower/upper_intervals and the layout's code width, and + * dequant_sum_delta from that delta and quantized_component_sums. + */ +void resolve_dequant_factors( + raft::resources const& res, + raft::device_vector_view dequant_delta, + raft::device_vector_view dequant_sum_delta, + raft::device_vector_view lower_intervals, + raft::device_vector_view upper_intervals, + raft::device_vector_view quantized_component_sums, + bbq_code_layout layout); +} // namespace helpers +/** @} */ // end of bbq group + +} // namespace preprocessing::quantize::bbq + +namespace neighbors { +struct bbq_dataset_container { + template + using owning_storage = cuvs::preprocessing::quantize::bbq::quantizer; + template + using view_storage = cuvs::preprocessing::quantize::bbq::quantizer_view; +}; + +template +struct dataset { + using owning_storage_type = bbq_dataset_container::owning_storage; + std::vector quantizers; + + dataset(owning_storage_type&& quantizer) noexcept { add_quantizer(std::move(quantizer)); } + [[nodiscard]] auto as_dataset_view() const noexcept + -> dataset_view> + { + return dataset_view>{quantizers}; + } + [[nodiscard]] constexpr auto n_rows() const noexcept -> IdxT + { + return quantizers.size() > 0 ? quantizers[0].n_rows() : 0; + } + [[nodiscard]] constexpr auto dim() const noexcept -> uint32_t + { + return quantizers.size() > 0 ? quantizers[0].dim() : 0; + } + + void add_quantizer(owning_storage_type&& quantizer) + { + RAFT_EXPECTS(!has_layout(quantizer.layout), "Quantizer already exists with layout."); + quantizers.push_back(std::move(quantizer)); + } + bool has_layout(cuvs::preprocessing::quantize::bbq::bbq_code_layout layout) const noexcept + { + for (uint32_t i = 0; i < quantizers.size(); i++) { + if (quantizers[i].layout == layout) { return true; } + } + return false; + } +}; + +template +struct dataset_view { + using owning_storage_type = bbq_dataset_container::owning_storage; + using view_storage_type = bbq_dataset_container::view_storage; + std::vector quantizers; + + dataset_view() noexcept = default; + + dataset_view(const std::vector& quantizers) noexcept + { + for (const auto& quantizer : quantizers) { + add_quantizer(quantizer); + } + } + [[nodiscard]] constexpr auto n_rows() const noexcept -> IdxT + { + return quantizers.size() > 0 ? quantizers[0].n_rows() : 0; + } + [[nodiscard]] constexpr auto dim() const noexcept -> uint32_t + { + return quantizers.size() > 0 ? quantizers[0].dim() : 0; + } + + void add_quantizer(view_storage_type quantizer) + { + RAFT_EXPECTS(!has_layout(quantizer.layout), "Quantizer already exists with layout."); + quantizers.push_back(quantizer); + } + void add_quantizer(const owning_storage_type& quantizer) + { + RAFT_EXPECTS(!has_layout(quantizer.layout), "Quantizer already exists with layout."); + quantizers.push_back(quantizer.view()); + } + bool has_layout(cuvs::preprocessing::quantize::bbq::bbq_code_layout layout) const noexcept + { + for (uint32_t i = 0; i < quantizers.size(); i++) { + if (quantizers[i].layout == layout) { return true; } + } + return false; + } + view_storage_type get_quantizer(cuvs::preprocessing::quantize::bbq::bbq_code_layout layout) const + { + for (uint32_t i = 0; i < quantizers.size(); i++) { + if (quantizers[i].layout == layout) { return quantizers[i]; } + } + throw std::runtime_error("No quantizer found with layout."); + } +}; + +template +using device_bbq_dataset = + dataset>; + +template +using device_bbq_dataset_view = + dataset_view>; + +template +struct owning_dataset_for_view> { + using type = device_bbq_dataset; +}; + +template +struct is_bbq_dataset : std::false_type {}; + +template +struct is_bbq_dataset> : std::true_type {}; + +template +inline constexpr bool is_bbq_dataset_v = is_bbq_dataset::value; + +template +struct dataset_view_kind_of> { + static constexpr dataset_view_kind value = dataset_view_kind::bbq; +}; + +template +struct cagra_view_element_type> { + using type = DataT; +}; + +template +inline constexpr bool is_device_bbq_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::bbq && dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_bbq_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::bbq && !dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_bbq_dataset_view_v = + is_device_bbq_dataset_view_v || is_host_bbq_dataset_view_v; + +} // namespace neighbors + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/internal/CMakeLists.txt b/cpp/internal/CMakeLists.txt index 4a50bfef58..403d06310d 100644 --- a/cpp/internal/CMakeLists.txt +++ b/cpp/internal/CMakeLists.txt @@ -1,14 +1,17 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= -if(BUILD_TESTS OR BUILD_MICRO_BENCH) +if(BUILD_TESTS + OR BUILD_MICRO_BENCH + OR BUILD_CUVS_BENCH +) add_library(cuvs_internal INTERFACE) target_include_directories( - cuvs_internal INTERFACE "$" + cuvs_internal INTERFACE "$" ) target_compile_features(cuvs_internal INTERFACE cxx_std_17 $) endif() diff --git a/cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp b/cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp new file mode 100644 index 0000000000..603517e518 --- /dev/null +++ b/cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp @@ -0,0 +1,629 @@ +/* + * SPDX-FileCopyrightText: Copyright the Apache Software Foundation (ASF) + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +/** + * Host-side reference implementation of Lucene's OptimizedScalarQuantizer, used to feed the + * BBQ NN-Descent build path. This is prototype scaffolding: it is a slow CPU quantizer that + * exists so the compressed build can be exercised end to end, and it is expected to be replaced + * by a GPU quantizer. It is shared between the unit tests and the ann-bench CAGRA wrapper so the + * two can never disagree about the code format; libcuvs itself only ever sees the uploaded codes. + */ +namespace cuvs_internal::bbq { + +using cuvs::preprocessing::quantize::bbq::bbq_code_layout; +using cuvs::preprocessing::quantize::bbq::get_bit_width; +using cuvs::preprocessing::quantize::bbq::get_encoded_row_length; + +/** + * Host-resident mirror of the arrays in a BBQ quantizer. The library quantizer is device-only, + * so the CPU reference implementation and its on-disk cache need their own staging type; it is + * uploaded by copy_bbq_owning_storage_host_to_device below. + */ +struct host_quantizer_storage { + raft::host_matrix codes; + raft::host_vector lower_intervals; + raft::host_vector upper_intervals; + raft::host_vector additional_corrections; + raft::host_vector quantized_component_sums; + raft::host_vector centroid; + raft::host_vector dequant_delta; + raft::host_vector dequant_sum_delta; + raft::host_vector row_norm; + bbq_code_layout layout{bbq_code_layout::packed_1b}; + cuvs::distance::DistanceType metric{cuvs::distance::DistanceType::L2Expanded}; + float centroid_norm_sq{}; +}; + +/** + * Fills the derived per-row dequantization factors. The device quantizer requires them up front + * (deriving them there would need a stream to allocate with), so they are computed here, where + * the intervals are already host-resident. + */ +inline void derive_dequant_factors(host_quantizer_storage& q) +{ + const float scale = 1.0f / static_cast((uint32_t{1} << get_bit_width(q.layout)) - 1); + for (int64_t i = 0; i < q.lower_intervals.extent(0); ++i) { + const float delta = (q.upper_intervals(i) - q.lower_intervals(i)) * scale; + q.dequant_delta(i) = delta; + q.dequant_sum_delta(i) = delta * static_cast(q.quantized_component_sums(i)); + } +} + +constexpr float kMinimumMseGrid[8][2] = {{-0.798f, 0.798f}, + {-1.493f, 1.493f}, + {-2.051f, 2.051f}, + {-2.514f, 2.514f}, + {-2.916f, 2.916f}, + {-3.278f, 3.278f}, + {-3.611f, 3.611f}, + {-3.922f, 3.922f}}; + +constexpr float kDefaultLambda = 0.1f; +constexpr int kDefaultIters = 5; + +inline long round(double x) { return static_cast(std::floor(x + 0.5)); } + +inline double clamp(double x, double a, double b) { return std::min(std::max(x, a), b); } + +inline double loss( + const std::vector& vector, const float interval[2], int points, float norm2, float lambda) +{ + const double a = interval[0]; + const double b = interval[1]; + const double step = (b - a) / (points - 1.0); + const double step_inv = 1.0 / step; + double xe = 0.0; + double e = 0.0; + for (double xi : vector) { + const double xiq = a + step * static_cast(round((clamp(xi, a, b) - a) * step_inv)); + xe += xi * (xi - xiq); + e += (xi - xiq) * (xi - xiq); + } + return (1.0 - lambda) * xe * xe / norm2 + lambda * e; +} + +inline void optimize_intervals(float interval[2], + const std::vector& vector, + float norm2, + int points, + float lambda = kDefaultLambda, + int iters = kDefaultIters) +{ + double initial_loss = loss(vector, interval, points, norm2, lambda); + const float scale = (1.0f - lambda) / norm2; + if (!std::isfinite(scale)) { return; } + for (int i = 0; i < iters; ++i) { + const float a = interval[0]; + const float b = interval[1]; + const float step_inv = (points - 1.0f) / (b - a); + double daa = 0.0, dab = 0.0, dbb = 0.0, dax = 0.0, dbx = 0.0; + for (float xi : vector) { + const float k = + static_cast(round(static_cast((clamp(xi, a, b) - a) * step_inv))); + const float s = k / (points - 1); + daa += (1.0 - s) * (1.0 - s); + dab += (1.0 - s) * s; + dbb += s * s; + dax += xi * (1.0 - s); + dbx += xi * s; + } + const double m0 = scale * dax * dax + lambda * daa; + const double m1 = scale * dax * dbx + lambda * dab; + const double m2 = scale * dbx * dbx + lambda * dbb; + const double det = m0 * m2 - m1 * m1; + if (det == 0) { return; } + const float a_opt = static_cast((m2 * dax - m1 * dbx) / det); + const float b_opt = static_cast((m0 * dbx - m1 * dax) / det); + if (std::abs(interval[0] - a_opt) < 1e-8 && std::abs(interval[1] - b_opt) < 1e-8) { return; } + float new_interval[2] = {a_opt, b_opt}; + const double new_loss = loss(vector, new_interval, points, norm2, lambda); + if (new_loss > initial_loss) { return; } + interval[0] = a_opt; + interval[1] = b_opt; + initial_loss = new_loss; + } +} + +struct row_result { + float lower_interval; + float upper_interval; + float additional_correction; + int32_t quantized_component_sum; +}; + +inline row_result scalar_quantize(std::vector& vector, + std::vector& destination, + uint8_t bits, + const float* centroid, + bool euclidean) +{ + const int n = static_cast(vector.size()); + const int points = 1 << bits; + double vec_mean = 0.0; + double vec_var = 0.0; + float norm2 = 0.0f; + float centroid_dot = 0.0f; + float min = FLT_MAX; + float max = -FLT_MAX; + for (int i = 0; i < n; ++i) { + if (!euclidean) { centroid_dot += vector[i] * centroid[i]; } + vector[i] = vector[i] - centroid[i]; + min = std::min(min, vector[i]); + max = std::max(max, vector[i]); + norm2 += vector[i] * vector[i]; + const double delta = vector[i] - vec_mean; + vec_mean += delta / (i + 1); + vec_var += delta * (vector[i] - vec_mean); + } + vec_var /= n; + const double vec_std = std::sqrt(vec_var); + + float interval[2]; + interval[0] = + static_cast(clamp(kMinimumMseGrid[bits - 1][0] * vec_std + vec_mean, min, max)); + interval[1] = + static_cast(clamp(kMinimumMseGrid[bits - 1][1] * vec_std + vec_mean, min, max)); + optimize_intervals(interval, vector, norm2, points); + + const float n_steps = static_cast((1 << bits) - 1); + const float a = interval[0]; + const float b = interval[1]; + const float step = (b - a) / n_steps; + int sum_query = 0; + for (int h = 0; h < n; ++h) { + const float xi = static_cast(clamp(vector[h], a, b)); + const int assignment = static_cast(round((xi - a) / step)); + sum_query += assignment; + destination[h] = static_cast(assignment); + } + return row_result{interval[0], interval[1], euclidean ? norm2 : centroid_dot, sum_query}; +} + +// Packs one-byte-per-component codes into packed_1b / transposed_2b / +// packed_4b / transposed_4b (or leaves unpacked). Matches Lucene packAsBinary, +// packNibbles, transposeDibit, transposeHalfByte. +inline std::vector pack_codes(const std::vector& unpacked, + size_t n_rows, + size_t dim, + bbq_code_layout layout) +{ + const uint32_t bits = get_bit_width(layout); + const size_t row_length = get_encoded_row_length(static_cast(dim), layout); + if (layout == bbq_code_layout::packed_8b || layout == bbq_code_layout::packed_7b) { + return unpacked; + } + + std::vector packed(n_rows * row_length, 0); +#pragma omp parallel for + for (int64_t row = 0; row < static_cast(n_rows); ++row) { + auto* output = packed.data() + static_cast(row) * row_length; + const auto* input = unpacked.data() + static_cast(row) * dim; + if (layout == bbq_code_layout::packed_4b) { + // Contiguous: dims 2k / 2k+1 share byte k. NOT Lucene packNibbles, which pairs dim i with + // dim dim/2 + i. A self-join is position-agnostic so either works there, but an asymmetric + // pair (packed_1b or transposed_2b document promoted to 4-bit width against this query) + // needs dimension k of both operands in the same slot -- halves-pairing silently + // multiplies mismatched dimensions and costs recall. + const size_t pairs = dim / 2; + for (size_t i = 0; i < pairs; ++i) { + output[i] = static_cast((input[2 * i] << 4) | (input[2 * i + 1] & 0x0f)); + } + continue; + } + for (size_t d = 0; d < dim; ++d) { + const uint8_t code = input[d]; + if (layout == bbq_code_layout::packed_1b) { + for (uint32_t bit = 0; bit < bits; ++bit) { + const size_t position = d * bits + bit; + output[position / 8] |= + static_cast(((code >> (bits - 1 - bit)) & 1u) << (7 - position % 8)); + } + } else { + // transposed_2b / transposed_4b bit-planes (LSB plane first) + const size_t stripe = (dim + 7) / 8; + for (uint32_t bit = 0; bit < bits; ++bit) { + output[bit * stripe + d / 8] |= static_cast(((code >> bit) & 1u) << (7 - d % 8)); + } + } + } + } + return packed; +} + +inline host_quantizer_storage quantize(const float* data, + int64_t n_rows, + int64_t dim, + cuvs::distance::DistanceType metric, + bbq_code_layout layout = bbq_code_layout::packed_8b) +{ + const auto bits = static_cast(get_bit_width(layout)); + const bool euclidean = metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded; + + auto centroid = raft::make_host_vector(dim); + float centroid_norm_sq = 0.0f; + std::fill_n(centroid.data_handle(), static_cast(dim), 0.0f); + for (int64_t i = 0; i < n_rows; ++i) { + for (int64_t d = 0; d < dim; ++d) { + centroid(d) += data[i * dim + d]; + } + } + for (int64_t d = 0; d < dim; ++d) { + centroid(d) /= static_cast(n_rows); + centroid_norm_sq += centroid(d) * centroid(d); + } + + std::vector unpacked(static_cast(n_rows * dim)); + auto lower_intervals = raft::make_host_vector(n_rows); + auto upper_intervals = raft::make_host_vector(n_rows); + auto additional_corrections = raft::make_host_vector(n_rows); + auto quantized_component_sums = raft::make_host_vector(n_rows); + auto row_norm = raft::make_host_vector(n_rows); + +#pragma omp parallel for + for (int64_t i = 0; i < n_rows; ++i) { + std::vector row(data + i * dim, data + (i + 1) * dim); + std::vector codes(dim); + float orig_norm2 = 0.0f; + for (int64_t d = 0; d < dim; ++d) { + orig_norm2 += row[d] * row[d]; + } + row_norm(i) = orig_norm2; + const auto result = scalar_quantize(row, codes, bits, centroid.data_handle(), euclidean); + std::copy(codes.begin(), codes.end(), unpacked.begin() + i * dim); + lower_intervals(i) = result.lower_interval; + upper_intervals(i) = result.upper_interval; + additional_corrections(i) = result.additional_correction; + quantized_component_sums(i) = result.quantized_component_sum; + } + + auto packed = pack_codes(unpacked, static_cast(n_rows), static_cast(dim), layout); + auto codes = raft::make_host_matrix( + n_rows, get_encoded_row_length(static_cast(dim), layout)); + std::copy(packed.begin(), packed.end(), codes.data_handle()); + + host_quantizer_storage out{std::move(codes), + std::move(lower_intervals), + std::move(upper_intervals), + std::move(additional_corrections), + std::move(quantized_component_sums), + std::move(centroid), + raft::make_host_vector(n_rows), + raft::make_host_vector(n_rows), + std::move(row_norm), + layout, + metric, + centroid_norm_sq}; + derive_dequant_factors(out); + return out; +} + +/** The CPU reference quantizer works in float, so it only ever feeds a float-valued dataset. */ +template +auto copy_bbq_owning_storage_host_to_device(raft::resources const& res, + host_quantizer_storage const& host_storage) -> + typename cuvs::neighbors::device_bbq_dataset::owning_storage_type +{ + using device_storage = + typename cuvs::neighbors::device_bbq_dataset::owning_storage_type; + auto stream = raft::resource::get_cuda_stream(res); + device_storage device{res, + static_cast(host_storage.codes.extent(0)), + static_cast(host_storage.centroid.extent(0)), + host_storage.layout, + host_storage.metric}; + + raft::copy( + device.codes.data_handle(), host_storage.codes.data_handle(), device.codes.size(), stream); + raft::copy(device.lower_intervals.data_handle(), + host_storage.lower_intervals.data_handle(), + device.lower_intervals.size(), + stream); + raft::copy(device.upper_intervals.data_handle(), + host_storage.upper_intervals.data_handle(), + device.upper_intervals.size(), + stream); + raft::copy(device.additional_corrections.data_handle(), + host_storage.additional_corrections.data_handle(), + device.additional_corrections.size(), + stream); + raft::copy(device.quantized_component_sums.data_handle(), + host_storage.quantized_component_sums.data_handle(), + device.quantized_component_sums.size(), + stream); + raft::copy(device.centroid.data_handle(), + host_storage.centroid.data_handle(), + device.centroid.size(), + stream); + raft::copy(device.dequant_delta.data_handle(), + host_storage.dequant_delta.data_handle(), + device.dequant_delta.size(), + stream); + raft::copy(device.dequant_sum_delta.data_handle(), + host_storage.dequant_sum_delta.data_handle(), + device.dequant_sum_delta.size(), + stream); + raft::copy(device.row_norm.data_handle(), + host_storage.row_norm.data_handle(), + device.row_norm.size(), + stream); + device.centroid_norm_sq = host_storage.centroid_norm_sq; + return device; +} + +template +auto make_device_bbq_dataset(raft::resources const& res, + std::vector const& host) + -> cuvs::neighbors::device_bbq_dataset +{ + RAFT_EXPECTS(host.size() != 0, "host BBQ dataset has no storage"); + cuvs::neighbors::device_bbq_dataset device{ + copy_bbq_owning_storage_host_to_device(res, host[0])}; + for (std::size_t i = 1; i < host.size(); ++i) { + device.add_quantizer(copy_bbq_owning_storage_host_to_device(res, host[i])); + } + return device; +} + +struct bbq_layout_token { + std::string_view token; + bbq_code_layout layout; +}; + +// Same convention as my_tests/bbq's CLI tokens: bare N = densely packed (tensor-core-eligible), +// N + "t" = transposed/bitplane (SIMT). At 1 bit the two coincide, so there is no "1t". There is +// no densely-packed 2-bit layout at all (packed_2b was retired): transposed_2b is SIMT-only, so +// "2t" is the only 2-bit token and it is not tensor-core-eligible. +constexpr bbq_layout_token kBbqLayoutTokens[] = { + {"1", bbq_code_layout::packed_1b}, + {"2t", bbq_code_layout::transposed_2b}, + {"4", bbq_code_layout::packed_4b}, + {"4t", bbq_code_layout::transposed_4b}, + {"7", bbq_code_layout::packed_7b}, + {"8", bbq_code_layout::packed_8b}, +}; + +/** Parses a layout token. Width follows from the layout (bits=4 is packed_4b vs transposed_4b). */ +inline auto parse_bbq_layout_token(std::string_view token) -> bbq_code_layout +{ + for (const auto& t : kBbqLayoutTokens) { + if (t.token == token) { return t.layout; } + } + RAFT_FAIL("Unknown BBQ layout token '%s'; expected one of 1, 2t, 4, 4t, 7, 8.", + std::string(token).c_str()); +} + +/** + * Reject layout pairs the local-join kernels cannot serve. Without this the library either + * silently falls back to a symmetric join on the first quantizer (which looks like a working + * asymmetric run), or reaches an unsupported-layout failure deep inside GNND::build. Must be kept + * in sync with the dispatch in GNND::local_join (nn_descent.cuh). + */ +inline void validate_layout_pair(bbq_code_layout query_layout, bbq_code_layout doc_layout) +{ + if (query_layout == doc_layout) { + const bool supported = query_layout == bbq_code_layout::packed_1b || + query_layout == bbq_code_layout::transposed_2b || + query_layout == bbq_code_layout::packed_4b || + query_layout == bbq_code_layout::packed_7b || + query_layout == bbq_code_layout::packed_8b; + RAFT_EXPECTS(supported, + "Symmetric BBQ NN-Descent has no local-join kernel for layout %d -- " + "transposed_4b is only supported as one side of an asymmetric pair.", + static_cast(query_layout)); + return; + } + const bool supported = + (doc_layout == bbq_code_layout::packed_1b && query_layout == bbq_code_layout::packed_4b) || + (doc_layout == bbq_code_layout::packed_1b && query_layout == bbq_code_layout::transposed_2b) || + (doc_layout == bbq_code_layout::packed_1b && query_layout == bbq_code_layout::transposed_4b) || + (doc_layout == bbq_code_layout::transposed_2b && + query_layout == bbq_code_layout::transposed_4b); + RAFT_EXPECTS(supported, + "Asymmetric BBQ NN-Descent supports only (doc, query) layouts of (packed_1b, " + "packed_4b), (packed_1b, transposed_2b), (packed_1b, transposed_4b), or " + "(transposed_2b, transposed_4b); got (%d, %d).", + static_cast(doc_layout), + static_cast(query_layout)); +} + +/** Every parameter that affects the codes is in the name, so the cache self-invalidates. */ +inline auto cache_path(int64_t n_rows, + int64_t dim, + bbq_code_layout layout, + cuvs::distance::DistanceType metric) -> std::string +{ + const char* dir = std::getenv("CUVS_BBQ_CACHE_DIR"); + return std::string{dir != nullptr && dir[0] != '\0' ? dir : "/tmp"} + "/bbq-n" + + std::to_string(n_rows) + "-d" + std::to_string(dim) + "-b" + + std::to_string(get_bit_width(layout)) + "-l" + std::to_string(static_cast(layout)) + + "-m" + std::to_string(static_cast(metric)) + ".bin"; +} + +/** Visits every raw buffer of @p q in a fixed order; this is the on-disk layout. */ +template +void for_each_buffer(host_quantizer_storage& q, OpT op) +{ + op(q.codes.data_handle(), q.codes.size() * sizeof(uint8_t)); + op(q.lower_intervals.data_handle(), q.lower_intervals.size() * sizeof(float)); + op(q.upper_intervals.data_handle(), q.upper_intervals.size() * sizeof(float)); + op(q.additional_corrections.data_handle(), q.additional_corrections.size() * sizeof(float)); + op(q.quantized_component_sums.data_handle(), q.quantized_component_sums.size() * sizeof(int32_t)); + op(q.centroid.data_handle(), q.centroid.size() * sizeof(float)); + // Unlike dequant_delta/dequant_sum_delta, row_norm (original-space ||x||^2) isn't derivable + // from the other cached fields, so it has to round-trip through the cache. + op(q.row_norm.data_handle(), q.row_norm.size() * sizeof(float)); +} + +/** Allocates the arrays of the given shape, leaving their contents undefined. */ +inline auto make_host_quantizer_storage(int64_t n_rows, + int64_t dim, + bbq_code_layout layout, + cuvs::distance::DistanceType metric) + -> host_quantizer_storage +{ + return host_quantizer_storage{ + raft::make_host_matrix( + n_rows, get_encoded_row_length(static_cast(dim), layout)), + raft::make_host_vector(n_rows), + raft::make_host_vector(n_rows), + raft::make_host_vector(n_rows), + raft::make_host_vector(n_rows), + raft::make_host_vector(dim), + raft::make_host_vector(n_rows), + raft::make_host_vector(n_rows), + raft::make_host_vector(n_rows), + layout, + metric, + 0.0f}; +} + +/** Reads the codes back if the file is present and has exactly the expected length. */ +inline auto cache_load(const std::string& path, + int64_t n_rows, + int64_t dim, + bbq_code_layout layout, + cuvs::distance::DistanceType metric) -> std::optional +{ + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) { return std::nullopt; } + auto q = make_host_quantizer_storage(n_rows, dim, layout, metric); + std::streamoff bytes = 0; + for_each_buffer(q, [&bytes](void*, size_t n) { bytes += static_cast(n); }); + if (f.tellg() != bytes) { + RAFT_LOG_WARN("Ignoring BBQ cache of unexpected size: %s", path.c_str()); + return std::nullopt; + } + f.seekg(0); + for_each_buffer( + q, [&f](void* p, size_t n) { f.read(static_cast(p), static_cast(n)); }); + if (!f) { + RAFT_LOG_WARN("Failed to read BBQ cache, re-quantizing: %s", path.c_str()); + return std::nullopt; + } + // Cheaper to recompute than to store and validate. dequant_delta/dequant_sum_delta are likewise + // a deterministic function of the cached lower/upper_intervals, quantized_component_sums, and + // layout, so they are recomputed here rather than added to the on-disk layout. + for (int64_t d = 0; d < dim; ++d) { + q.centroid_norm_sq += q.centroid(d) * q.centroid(d); + } + derive_dequant_factors(q); + RAFT_LOG_INFO("Loaded BBQ codes from %s", path.c_str()); + return q; +} + +/** Writes via a temporary so an interrupted run cannot leave a truncated cache behind. */ +inline void cache_store(const std::string& path, host_quantizer_storage& q) +{ + const std::string tmp = path + ".tmp"; + { + std::ofstream f(tmp, std::ios::binary | std::ios::trunc); + if (!f) { + RAFT_LOG_WARN("Cannot write BBQ cache: %s", tmp.c_str()); + return; + } + for_each_buffer(q, [&f](void* p, size_t n) { + f.write(static_cast(p), static_cast(n)); + }); + if (!f) { + RAFT_LOG_WARN("Failed to write BBQ cache: %s", tmp.c_str()); + f.close(); + std::remove(tmp.c_str()); + return; + } + } + if (std::rename(tmp.c_str(), path.c_str()) != 0) { + RAFT_LOG_WARN("Cannot move BBQ cache into place: %s", path.c_str()); + std::remove(tmp.c_str()); + return; + } + RAFT_LOG_INFO("Wrote BBQ codes to %s", path.c_str()); +} + +inline auto quantize_cached(const float* rows, + int64_t n_rows, + int64_t dim, + bbq_code_layout layout, + cuvs::distance::DistanceType metric) -> host_quantizer_storage +{ + const auto path = cache_path(n_rows, dim, layout, metric); + if (auto cached = cache_load(path, n_rows, dim, layout, metric)) { return std::move(*cached); } + auto quantized = quantize(rows, n_rows, dim, metric, layout); + cache_store(path, quantized); + return quantized; +} + +/** + * Quantize host-resident @p rows and upload the codes, ready for the BBQ NN-Descent build. + * + * `query_token == doc_token` produces a single quantizer and a symmetric join; differing tokens + * produce two quantizers and an asymmetric join, where the wider codes serve the query side. + */ +inline auto quantize_to_device(raft::resources const& res, + const float* rows, + int64_t n_rows, + int64_t dim, + cuvs::distance::DistanceType metric, + bbq_code_layout query_layout, + bbq_code_layout doc_layout) + -> cuvs::neighbors::device_bbq_dataset +{ + validate_layout_pair(query_layout, doc_layout); + std::vector host; + host.push_back(quantize_cached(rows, n_rows, dim, query_layout, metric)); + if (doc_layout != query_layout) { + host.push_back(quantize_cached(rows, n_rows, dim, doc_layout, metric)); + } + auto device = make_device_bbq_dataset(res, host); + // The uploads are stream-ordered against `host`, which dies with this frame. + raft::resource::sync_stream(res); + return device; +} + +inline auto quantize_to_device(raft::resources const& res, + const float* rows, + int64_t n_rows, + int64_t dim, + cuvs::distance::DistanceType metric, + std::string_view query_token, + std::string_view doc_token) + -> cuvs::neighbors::device_bbq_dataset +{ + return quantize_to_device(res, + rows, + n_rows, + dim, + metric, + parse_bbq_layout_token(query_token), + parse_bbq_layout_token(doc_token)); +} +} // namespace cuvs_internal::bbq diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 1c03cb8731..311211ff05 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -45,7 +45,8 @@ CUVS_EXPORT void index::compute_dataset_norms_(raft::reso if constexpr (nb::is_padded_dataset_view_v || nb::is_standard_dataset_view_v) { rm_dataset = dataset_.view(); - } else if constexpr (nb::is_vpq_dataset_view_v) { + } else if constexpr (nb::is_vpq_dataset_view_v || + nb::is_bbq_dataset_view_v) { skip_norms = true; } @@ -301,7 +302,10 @@ auto build(raft::resources const& res, const index_params& params, DatasetViewT // Dense paths build the graph and optionally attach the input dataset view. Host indexes remain // non-searchable until the type-changing update_dataset(...) supplies a device-padded dataset. - if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { + if constexpr (cuvs::neighbors::is_device_bbq_dataset_view_v) { + return cuvs::neighbors::cagra::detail::build_from_bbq_dataset( + res, params, dataset); + } else if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { auto effective_params = params; if (std::holds_alternative(effective_params.graph_build_params)) { effective_params.graph_build_params = graph_build_params::iterative_search_params{}; diff --git a/cpp/src/neighbors/cagra_build_inst.cu.in b/cpp/src/neighbors/cagra_build_inst.cu.in index b4c83c6191..2a3e2b43c1 100644 --- a/cpp/src/neighbors/cagra_build_inst.cu.in +++ b/cpp/src/neighbors/cagra_build_inst.cu.in @@ -20,6 +20,7 @@ using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_vie using inst_host_padded_view_t = cuvs::neighbors::host_padded_dataset_view; using inst_host_standard_view_t = cuvs::neighbors::host_standard_dataset_view; using inst_vpq_view_t = cuvs::neighbors::device_vpq_dataset_view; +using inst_bbq_view_t = cuvs::neighbors::device_bbq_dataset_view; } // namespace namespace cuvs::neighbors::cagra { @@ -32,6 +33,8 @@ extern template void index::compute_d raft::resources const&); extern template void index::compute_dataset_norms_( raft::resources const&); +extern template void index::compute_dataset_norms_( + raft::resources const&); void build_knn_graph(raft::resources const& handle, raft::host_matrix_view dataset, raft::host_matrix_view knn_graph, @@ -62,6 +65,9 @@ CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_vpq_view_t, cuvs::neighbors::cagra::device_pq_index); #endif +CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_bbq_view_t, + cuvs::neighbors::cagra::device_bbq_index); + #undef CUVS_DEFINE_CAGRA_BUILD_OVERLOAD #define CUVS_INST_CAGRA_UPDATE_DATASET(T, IdxT, SrcDatasetViewT, DstDatasetViewT) \ @@ -100,6 +106,8 @@ CUVS_INST_CAGRA_UPDATE_DATASET(data_t, index_t, inst_device_standard_view_t, ins CUVS_INST_CAGRA_UPDATE_DATASET(data_t, index_t, inst_device_padded_view_t, inst_vpq_view_t); CUVS_INST_CAGRA_UPDATE_DATASET(data_t, index_t, inst_vpq_view_t, inst_device_padded_view_t); CUVS_INST_CAGRA_UPDATE_DATASET(data_t, index_t, inst_vpq_view_t, inst_vpq_view_t); +CUVS_INST_CAGRA_UPDATE_DATASET(data_t, index_t, inst_bbq_view_t, inst_device_padded_view_t); +CUVS_INST_CAGRA_UPDATE_DATASET(data_t, index_t, inst_bbq_view_t, inst_vpq_view_t); #undef CUVS_INST_CAGRA_UPDATE_DATASET diff --git a/cpp/src/neighbors/cagra_search_inst.cu.in b/cpp/src/neighbors/cagra_search_inst.cu.in index c5e35819f7..9fff49abd6 100644 --- a/cpp/src/neighbors/cagra_search_inst.cu.in +++ b/cpp/src/neighbors/cagra_search_inst.cu.in @@ -14,6 +14,7 @@ using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view< using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; using inst_vpq_f16_view_t = cuvs::neighbors::device_vpq_dataset_view; using inst_vpq_f32_view_t = cuvs::neighbors::device_vpq_dataset_view; +using inst_bbq_view_t = cuvs::neighbors::device_bbq_dataset_view; using inst_empty_view_t = cuvs::neighbors::device_empty_dataset_view; } // namespace @@ -57,6 +58,8 @@ index::compute_dataset_norms_(raft::resou template CUVS_EXPORT __attribute__((externally_visible)) void index::compute_dataset_norms_(raft::resources const& res); template CUVS_EXPORT __attribute__((externally_visible)) void +index::compute_dataset_norms_(raft::resources const& res); +template CUVS_EXPORT __attribute__((externally_visible)) void index::compute_dataset_norms_(raft::resources const& res); #define CUVS_INST_CAGRA_SEARCH_MULTI_PARTITION(T, IdxT, OutputIdxT) \ diff --git a/cpp/src/neighbors/cagra_serialize.cuh b/cpp/src/neighbors/cagra_serialize.cuh index c05e4c0deb..18e807b90f 100644 --- a/cpp/src/neighbors/cagra_serialize.cuh +++ b/cpp/src/neighbors/cagra_serialize.cuh @@ -184,6 +184,35 @@ namespace cuvs::neighbors::cagra { cuvs::neighbors::cagra::detail::deserialize(handle, is, index, nullptr); \ } \ \ + void serialize(raft::resources const& handle, \ + const std::string& filename, \ + const cuvs::neighbors::cagra::device_bbq_index& index) \ + { \ + cuvs::neighbors::cagra::detail::serialize(handle, filename, index, false); \ + } \ + \ + void deserialize(raft::resources const& handle, \ + const std::string& filename, \ + cuvs::neighbors::cagra::device_bbq_index* index) \ + { \ + cuvs::neighbors::cagra::detail::deserialize( \ + handle, filename, index, nullptr); \ + } \ + \ + void serialize(raft::resources const& handle, \ + std::ostream& os, \ + const cuvs::neighbors::cagra::device_bbq_index& index) \ + { \ + cuvs::neighbors::cagra::detail::serialize(handle, os, index, false); \ + } \ + \ + void deserialize(raft::resources const& handle, \ + std::istream& is, \ + cuvs::neighbors::cagra::device_bbq_index* index) \ + { \ + cuvs::neighbors::cagra::detail::deserialize(handle, is, index, nullptr); \ + } \ + \ void serialize_to_hnswlib( \ raft::resources const& handle, \ std::ostream& os, \ diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 5fb60581e2..d90eca17a4 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -2212,6 +2212,36 @@ void build_knn_graph( res, build_params.metric, dataset, knn_graph_internal); } +template +void build_knn_graph(raft::resources const& res, + cuvs::neighbors::device_bbq_dataset_view dataset, + raft::host_matrix_view knn_graph, + cuvs::neighbors::nn_descent::index_params build_params) +{ + raft::common::nvtx::range fun_scope( + "cagra::build_knn_graph(%zu, %zu, %u)", + size_t(dataset.n_rows()), + size_t(dataset.dim()), + size_t(knn_graph.extent(1))); + + std::optional> graph_view = knn_graph; + auto nn_descent_idx = cuvs::neighbors::nn_descent::build(res, build_params, dataset, graph_view); + + using internal_IdxT = typename std::make_unsigned::type; + using g_accessor = typename decltype(nn_descent_idx.graph())::accessor_type; + using g_accessor_internal = + raft::host_device_accessor, g_accessor::mem_type>; + + auto knn_graph_internal = + raft::mdspan, raft::row_major, g_accessor_internal>( + reinterpret_cast(nn_descent_idx.graph().data_handle()), + nn_descent_idx.graph().extent(0), + nn_descent_idx.graph().extent(1)); + + cuvs::neighbors::cagra::detail::graph::sort_knn_graph_bbq( + res, build_params.metric, dataset, knn_graph_internal); +} + template , raft::memory_type::host>> @@ -2865,4 +2895,94 @@ auto build_from_device_matrix(raft::resources const& res, RAFT_LOG_TRACE("Graph optimized, creating index"); return idx; } + +[[nodiscard]] inline auto resolve_bbq_knn_graph_build_params(index_params const& params, + size_t intermediate_degree) + -> cuvs::neighbors::nn_descent::index_params +{ + if (std::holds_alternative(params.graph_build_params)) { + return cagra::graph_build_params::nn_descent_params(intermediate_degree, params.metric); + } + + auto nn_descent_params = + std::get(params.graph_build_params); + if (nn_descent_params.metric != params.metric) { + RAFT_LOG_WARN( + "Metric (%lu) for nn-descent needs to match cagra metric (%lu), " + "aligning nn-descent metric.", + nn_descent_params.metric, + params.metric); + nn_descent_params.metric = params.metric; + } + if (nn_descent_params.graph_degree != intermediate_degree) { + RAFT_LOG_WARN( + "Graph degree (%lu) for nn-descent needs to match cagra intermediate graph degree (%lu), " + "aligning nn-descent graph_degree.", + nn_descent_params.graph_degree, + intermediate_degree); + nn_descent_params = + cagra::graph_build_params::nn_descent_params(intermediate_degree, params.metric); + } + return nn_descent_params; +} + +/** + * Build from a device-resident BBQ-quantized dataset: the whole graph construction runs on the + * compressed codes. + * + * The returned index cannot be searched, because CAGRA has no BBQ search kernels. Pass an + * uncompressed device-padded dataset to `cagra::update_dataset` to obtain a searchable index over + * the same graph. + */ +template + requires cuvs::neighbors::is_device_bbq_dataset_view_v +auto build_from_bbq_dataset(raft::resources const& res, + const index_params& params, + DatasetViewT const& dataset) + -> cuvs::neighbors::cagra::index +{ + RAFT_EXPECTS(!dataset.quantizers.empty(), "cagra::build: the BBQ dataset is empty."); + RAFT_EXPECTS(params.metric == cuvs::distance::DistanceType::L2Expanded || + params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || + params.metric == cuvs::distance::DistanceType::CosineExpanded || + params.metric == cuvs::distance::DistanceType::InnerProduct, + "cagra::build: a BBQ-quantized dataset supports L2Expanded, L2SqrtExpanded, " + "CosineExpanded, and InnerProduct."); + RAFT_EXPECTS(std::holds_alternative(params.graph_build_params) || + std::holds_alternative( + params.graph_build_params), + "cagra::build: a BBQ-quantized dataset requires nn-descent graph construction."); + + size_t intermediate_degree = params.intermediate_graph_degree; + size_t graph_degree = params.graph_degree; + common::nvtx::range function_scope( + "cagra::detail::build_from_bbq_dataset(%zu, %zu)", intermediate_degree, graph_degree); + auto const n_rows = static_cast(dataset.n_rows()); + check_graph_degree(intermediate_degree, graph_degree, static_cast(n_rows)); + + auto nn_descent_params = resolve_bbq_knn_graph_build_params(params, intermediate_degree); + nn_descent_params.return_distances = false; + + auto cagra_graph = [&]() -> raft::host_matrix { + std::optional> knn_graph( + raft::make_host_matrix(n_rows, intermediate_degree)); + build_knn_graph(res, dataset, knn_graph->view(), nn_descent_params); + + auto optimized = raft::make_host_matrix(n_rows, graph_degree); + RAFT_LOG_TRACE("optimizing graph"); + optimize(res, knn_graph->view(), optimized.view(), params.guarantee_connectivity); + knn_graph.reset(); + return optimized; + }(); + + RAFT_LOG_TRACE("Graph optimized, creating index"); + + if (params.attach_dataset_on_build) { + return cuvs::neighbors::cagra::index( + res, params.metric, dataset, raft::make_const_mdspan(cagra_graph.view())); + } + cuvs::neighbors::cagra::index idx(res, params.metric); + idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view())); + return idx; +} } // namespace cuvs::neighbors::cagra::detail diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index 438303a607..94ea20c1ad 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -77,6 +77,19 @@ constexpr bool is_valid_serialized_dataset_kind(std::uint32_t raw) return raw <= static_cast(kind::host_standard); } +/** + * Quantized datasets (PQ, BBQ) are owned outside the index and carry codebooks the index file has + * no representation for, so such indexes serialize the graph alone. + */ +template +inline constexpr bool is_graph_only_dataset_view_v = + cuvs::neighbors::is_vpq_dataset_view_v || + cuvs::neighbors::is_bbq_dataset_view_v; + +inline constexpr char const* kGraphOnlyDatasetMessage = + "CAGRA indexes with a quantized dataset store only the graph; serialize the quantized dataset " + "separately and reattach it with update_dataset()"; + template void serialize_index_mdspan(raft::resources const& res, std::ostream& os, MdspanT const& mdspan) { @@ -111,10 +124,8 @@ void serialize(raft::resources const& res, include_dataset &= (index_.dataset().n_rows() > 0); auto dataset_kind = cuvs::neighbors::cagra::serialized_dataset_kind::none; - if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { - RAFT_EXPECTS(!include_dataset, - "CAGRA PQ index serialization stores the graph only; serialize the PQ dataset " - "separately"); + if constexpr (is_graph_only_dataset_view_v) { + RAFT_EXPECTS(!include_dataset, kGraphOnlyDatasetMessage); } else { if (include_dataset) { dataset_kind = serialized_dataset_kind_for_view(); } } @@ -140,8 +151,8 @@ void serialize(raft::resources const& res, RAFT_LOG_DEBUG("Saving CAGRA index with dataset"); if constexpr (cuvs::neighbors::is_dense_row_major_dataset_view_v) { neighbors::detail::serialize_cagra_dense_dataset(res, os, index_.dataset()); - } else if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { - RAFT_FAIL("CAGRA PQ index serialization stores the graph only"); + } else if constexpr (is_graph_only_dataset_view_v) { + RAFT_FAIL(kGraphOnlyDatasetMessage); } else { // A further dataset type requires a new branch here and a corresponding deserialize branch. // Use static_assert to catch unsupported types at compile time. @@ -588,14 +599,14 @@ void deserialize_impl( std::unique_ptr dataset_owner{}; if (has_dataset) { if (out_dataset == nullptr) { - if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { - RAFT_FAIL("cagra::deserialize: PQ index files must contain only the graph"); + if constexpr (is_graph_only_dataset_view_v) { + RAFT_FAIL("cagra::deserialize: quantized index files must contain only the graph"); } else { cuvs::neighbors::detail::skip_dense_dataset(res, is); } } else { - if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { - RAFT_FAIL("cagra::deserialize: PQ index files must contain only the graph"); + if constexpr (is_graph_only_dataset_view_v) { + RAFT_FAIL("cagra::deserialize: quantized index files must contain only the graph"); } else { auto const expected_kind = serialized_dataset_kind_for_view(); RAFT_EXPECTS( diff --git a/cpp/src/neighbors/detail/cagra/factory.cuh b/cpp/src/neighbors/detail/cagra/factory.cuh index cdcc18867b..ce522533a8 100644 --- a/cpp/src/neighbors/detail/cagra/factory.cuh +++ b/cpp/src/neighbors/detail/cagra/factory.cuh @@ -82,6 +82,7 @@ namespace descriptor_cache { */ struct key { uint64_t data_ptr; + uint64_t norms_ptr; uint64_t n_rows; uint32_t dim; uint32_t extra_val; // this one has different meanings for different descriptor types @@ -93,10 +94,11 @@ struct key { template auto make_key(const cagra::search_params& params, const DatasetT& dataset, - cuvs::distance::DistanceType metric) - -> std::enable_if_t, key> + cuvs::distance::DistanceType metric, + const void* dataset_norms) -> std::enable_if_t, key> { return key{reinterpret_cast(dataset.view().data_handle()), + reinterpret_cast(dataset_norms), uint64_t(dataset.n_rows()), dataset.dim(), dataset.stride(), @@ -108,10 +110,11 @@ auto make_key(const cagra::search_params& params, template auto make_key(const cagra::search_params& params, const DatasetT& dataset, - cuvs::distance::DistanceType metric) - -> std::enable_if_t, key> + cuvs::distance::DistanceType metric, + const void* dataset_norms) -> std::enable_if_t, key> { return key{reinterpret_cast(dataset.data.data_handle()), + reinterpret_cast(dataset_norms), uint64_t(dataset.n_rows()), dataset.dim(), uint32_t(reinterpret_cast(dataset.pq_code_book.data_handle()) >> 6), @@ -122,15 +125,16 @@ auto make_key(const cagra::search_params& params, inline auto operator==(const key& a, const key& b) -> bool { - return a.data_ptr == b.data_ptr && a.n_rows == b.n_rows && a.dim == b.dim && - a.extra_val == b.extra_val && a.team_size == b.team_size && a.metric == b.metric && - a.smem_dtype == b.smem_dtype; + return a.data_ptr == b.data_ptr && a.norms_ptr == b.norms_ptr && a.n_rows == b.n_rows && + a.dim == b.dim && a.extra_val == b.extra_val && a.team_size == b.team_size && + a.metric == b.metric && a.smem_dtype == b.smem_dtype; } struct key_hash { inline auto operator()(const key& x) const noexcept -> std::size_t { - return size_t{x.data_ptr} + size_t{x.n_rows} * size_t{x.dim} * size_t{x.extra_val} + + return size_t{x.data_ptr} + size_t{x.norms_ptr} + + size_t{x.n_rows} * size_t{x.dim} * size_t{x.extra_val} + (size_t{x.team_size} ^ size_t{x.metric}) + size_t{x.smem_dtype}; } }; @@ -164,7 +168,7 @@ auto dataset_descriptor_init_with_cache(const raft::resources& res, const DistanceT* dataset_norms = nullptr) -> dataset_descriptor_host { - auto key = descriptor_cache::make_key(params, dataset, metric); + auto key = descriptor_cache::make_key(params, dataset, metric, dataset_norms); auto& cache = raft::resource::get_custom_resource>(res) ->value; diff --git a/cpp/src/neighbors/detail/cagra/graph_shared.cu b/cpp/src/neighbors/detail/cagra/graph_shared.cu index c088b8c477..8dcb4d8b6f 100644 --- a/cpp/src/neighbors/detail/cagra/graph_shared.cu +++ b/cpp/src/neighbors/detail/cagra/graph_shared.cu @@ -8,13 +8,18 @@ #include "utils.hpp" // TODO: This shouldn't be invoking anything from spatial/knn +#include "../../../preprocessing/quantize/detail/bbq_distance.cuh" #include "../ann_utils.cuh" +#include +#include #include +#include #include #include #include +#include namespace cuvs::neighbors::cagra::detail::graph { namespace { @@ -156,6 +161,190 @@ void launch_sort_knn_graph_impl(raft::resources const& res, RAFT_CUDA_TRY(cudaGetLastError()); } +template +using bbq_quantizer_view_t = cuvs::preprocessing::quantize::bbq::quantizer_view; + +/** + * Distance between two dataset rows in quantized space + * + */ +template +__device__ __forceinline__ float bbq_row_distance( + const bbq_quantizer_view_t& quantizer_document, + const bbq_quantizer_view_t& quantizer_query, + cuvs::distance::DistanceType metric, + int64_t row_document, + int64_t row_query) +{ + namespace bbq = cuvs::preprocessing::quantize::bbq; + const uint32_t raw = + bbq::code_inner_product(quantizer_document, quantizer_query, row_document, row_query); + return bbq::bbq_calculate_metric(raw, + bbq::get_dequant_factors(quantizer_document, row_document), + bbq::get_dequant_factors(quantizer_query, row_query), + quantizer_document, + quantizer_query, + metric, + raft::identity_op{}, + row_document, + row_query); +} + +template +__global__ void kern_sort_bbq(const bbq_quantizer_view_t quantizer_document, + const bbq_quantizer_view_t quantizer_query, + uint32_t* const knn_graph, // [graph_size, graph_degree] + const uint32_t graph_size, + const uint32_t graph_degree, + const cuvs::distance::DistanceType metric) +{ + const uint32_t src_node = (blockDim.x * blockIdx.x + threadIdx.x) / raft::WarpSize; + if (src_node >= graph_size) { return; } + + const uint32_t lane_id = threadIdx.x % raft::WarpSize; + + float my_keys[numElementsPerThread]; + uint32_t my_vals[numElementsPerThread]; + + for (int i = 0; i < numElementsPerThread; i++) { + const uint32_t k = i * raft::WarpSize + lane_id; + if (k >= graph_degree) { + my_keys[i] = utils::get_max_value(); + my_vals[i] = utils::get_max_value(); + continue; + } + const uint32_t dst_node = knn_graph[k + static_cast(graph_degree) * src_node]; + my_keys[i] = bbq_row_distance(quantizer_document, quantizer_query, metric, src_node, dst_node); + my_vals[i] = dst_node; + } + + raft::util::bitonic(true).sort(my_keys, my_vals); + + for (int i = 0; i < numElementsPerThread; i++) { + const uint32_t k = i * raft::WarpSize + lane_id; + if (k < graph_degree) { + knn_graph[k + (static_cast(graph_degree) * src_node)] = my_vals[i]; + } + } +} + +template +using sort_bbq_kernel_type = void (*)(bbq_quantizer_view_t, + bbq_quantizer_view_t, + uint32_t*, + uint32_t, + uint32_t, + cuvs::distance::DistanceType); + +template +auto select_sort_bbq_kernel(uint32_t degree) -> sort_bbq_kernel_type +{ + if (degree <= raft::WarpSize * 1) { return kern_sort_bbq; } + if (degree <= raft::WarpSize * 2) { return kern_sort_bbq; } + if (degree <= raft::WarpSize * 4) { return kern_sort_bbq; } + if (degree <= raft::WarpSize * 8) { return kern_sort_bbq; } + if (degree <= raft::WarpSize * 16) { return kern_sort_bbq; } + if (degree <= kMaxSortDegree) { return kern_sort_bbq; } + RAFT_FAIL( + "The degree of input knn graph is too large (%u). It must be equal to or smaller than %lu.", + degree, + kMaxSortDegree); +} + +template +auto select_sort_quantizers(cuvs::neighbors::device_bbq_dataset_view const& dataset) + -> std::pair, bbq_quantizer_view_t> +{ + using bbq_code_layout = cuvs::preprocessing::quantize::bbq::bbq_code_layout; + + if (dataset.quantizers.size() == 1) { return {dataset.quantizers[0], dataset.quantizers[0]}; } + + const bool has_1b = dataset.has_layout(bbq_code_layout::packed_1b); + const bool has_4b = dataset.has_layout(bbq_code_layout::packed_4b); + const bool has_2bt = dataset.has_layout(bbq_code_layout::transposed_2b); + const bool has_4bt = dataset.has_layout(bbq_code_layout::transposed_4b); + + const bool tc_pair = has_1b && has_4b; + const bool simt_pair = (has_4bt && (has_1b || has_2bt)) || (has_2bt && has_1b); + RAFT_EXPECTS(tc_pair || simt_pair, + "Unsupported BBQ layout pair for an asymmetric dataset. Supported (document, " + "query) pairs: (packed_1b, packed_4b), (packed_1b, transposed_2b), " + "(packed_1b, transposed_4b), (transposed_2b, transposed_4b)."); + return {has_1b ? dataset.get_quantizer(bbq_code_layout::packed_1b) + : dataset.get_quantizer(bbq_code_layout::transposed_2b), + tc_pair ? dataset.get_quantizer(bbq_code_layout::packed_4b) + : has_4bt ? dataset.get_quantizer(bbq_code_layout::transposed_4b) + : dataset.get_quantizer(bbq_code_layout::transposed_2b)}; +} + +template +void sort_knn_graph_bbq_impl(raft::resources const& res, + cuvs::distance::DistanceType metric, + cuvs::neighbors::device_bbq_dataset_view dataset, + raft::host_matrix_view knn_graph) +{ + namespace bbq = cuvs::preprocessing::quantize::bbq; + + RAFT_EXPECTS(!dataset.quantizers.empty(), "the BBQ dataset holds no quantizer"); + RAFT_EXPECTS(dataset.n_rows() == knn_graph.extent(0), + "dataset size is expected to have the same number of graph index size"); + RAFT_EXPECTS(metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::InnerProduct || + metric == cuvs::distance::DistanceType::CosineExpanded, + "Unsupported metric. Only L2Expanded, L2SqrtExpanded, InnerProduct and " + "CosineExpanded are supported for a BBQ-quantized dataset"); + + auto const graph_size = static_cast(knn_graph.extent(0)); + auto const graph_degree = static_cast(knn_graph.extent(1)); + auto kernel = select_sort_bbq_kernel(graph_degree); + auto const [quantizer_document, quantizer_query] = select_sort_quantizers(dataset); + // The code inner products read both rows as uint32_t words, and a bit-sliced layout starts + // every plane at a multiple of the plane stride, so each plane must be 4-byte aligned. + for (const auto& quantizer : {quantizer_document, quantizer_query}) { + const auto row_length = bbq::get_encoded_row_length(quantizer); + const auto planes = static_cast(bbq::get_code_planes(quantizer.layout)); + RAFT_EXPECTS(row_length % (4u * planes) == 0, + "Sorting a BBQ-quantized kNN graph requires the encoded row length to be a " + "multiple of 4*n_planes for 32-bit aligned plane loads, got %u with n_planes = %u", + row_length, + planes); + } + // A packed_1b document is promoted to 4-bit width one 32-dimension word at a time, which only + // covers the packed_4b query row exactly when the dimensionality is a multiple of 32. + RAFT_EXPECTS(quantizer_query.layout != bbq::bbq_code_layout::packed_4b || + quantizer_document.layout == quantizer_query.layout || + quantizer_document.dim() % 32 == 0, + "Sorting a BBQ-quantized kNN graph with packed_1b codes against packed_4b ones " + "requires the dataset dim to be a multiple of 32, got %u", + quantizer_document.dim()); + + const double time_sort_start = cur_time(); + RAFT_LOG_DEBUG("# Sorting kNN Graph on GPUs\n"); + + auto large_tmp_mr = raft::resource::get_large_workspace_resource_ref(res); + auto d_input_graph = raft::make_device_mdarray( + res, large_tmp_mr, raft::make_extents(graph_size, graph_degree)); + raft::copy(res, d_input_graph.view(), knn_graph); + + constexpr uint32_t block_size = 256; + auto const warps = block_size / raft::WarpSize; + auto const blocks = (graph_size + warps - 1) / warps; + kernel<<>>( + quantizer_document, + quantizer_query, + d_input_graph.data_handle(), + graph_size, + graph_degree, + metric); + RAFT_CUDA_TRY(cudaGetLastError()); + raft::resource::sync_stream(res); + raft::copy(res, knn_graph, raft::make_const_mdspan(d_input_graph.view())); + + const double time_sort_end = cur_time(); + RAFT_LOG_DEBUG("# Sorting kNN graph time: %.1lf sec\n", time_sort_end - time_sort_start); +} + } // namespace #define CUVS_DEFINE_CAGRA_GRAPH_SORT(DataT) \ @@ -178,6 +367,22 @@ CUVS_DEFINE_CAGRA_GRAPH_SORT(uint8_t) #undef CUVS_DEFINE_CAGRA_GRAPH_SORT +#define CUVS_DEFINE_CAGRA_GRAPH_SORT_BBQ(DataT) \ + void sort_knn_graph_bbq(raft::resources const& res, \ + cuvs::distance::DistanceType metric, \ + cuvs::neighbors::device_bbq_dataset_view dataset, \ + raft::host_matrix_view knn_graph) \ + { \ + sort_knn_graph_bbq_impl(res, metric, dataset, knn_graph); \ + } + +CUVS_DEFINE_CAGRA_GRAPH_SORT_BBQ(float) +CUVS_DEFINE_CAGRA_GRAPH_SORT_BBQ(half) +CUVS_DEFINE_CAGRA_GRAPH_SORT_BBQ(int8_t) +CUVS_DEFINE_CAGRA_GRAPH_SORT_BBQ(uint8_t) + +#undef CUVS_DEFINE_CAGRA_GRAPH_SORT_BBQ + void optimize_device_graph( raft::resources const& res, raft::device_matrix_view knn_graph, diff --git a/cpp/src/neighbors/detail/cagra/graph_shared.cuh b/cpp/src/neighbors/detail/cagra/graph_shared.cuh index 0f2bc0f7dc..33ec7f92ed 100644 --- a/cpp/src/neighbors/detail/cagra/graph_shared.cuh +++ b/cpp/src/neighbors/detail/cagra/graph_shared.cuh @@ -6,8 +6,10 @@ #include #include +#include #include +#include #include #include @@ -35,6 +37,20 @@ CUVS_DECL_CAGRA_GRAPH_SORT(uint8_t); #undef CUVS_DECL_CAGRA_GRAPH_SORT +#define CUVS_DECL_CAGRA_GRAPH_SORT_BBQ(DataT) \ + CUVS_EXPORT void sort_knn_graph_bbq( \ + raft::resources const& res, \ + cuvs::distance::DistanceType metric, \ + cuvs::neighbors::device_bbq_dataset_view dataset, \ + raft::host_matrix_view knn_graph) + +CUVS_DECL_CAGRA_GRAPH_SORT_BBQ(float); +CUVS_DECL_CAGRA_GRAPH_SORT_BBQ(half); +CUVS_DECL_CAGRA_GRAPH_SORT_BBQ(int8_t); +CUVS_DECL_CAGRA_GRAPH_SORT_BBQ(uint8_t); + +#undef CUVS_DECL_CAGRA_GRAPH_SORT_BBQ + /** Run the existing CAGRA optimizer through one compiled instantiation instead of rematerializing * its reverse-graph, prune, merge, and MST kernels in every Fastener dtype TU. */ CUVS_EXPORT void optimize_device_graph( diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index ec92e1a36f..b1260784b1 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -9,9 +9,12 @@ #include "neighbors_device_intrinsics.cuh" #include "nn_descent_gnnd.hpp" +#include "../../core/nvtx.hpp" #include "../../core/omp_wrapper.hpp" +#include "../../preprocessing/quantize/detail/bbq_distance.cuh" #include #include +#include #include #include @@ -38,6 +41,7 @@ #include +#include #include #include #include @@ -47,6 +51,9 @@ namespace cuvs::neighbors::nn_descent::detail { +using cuvs::preprocessing::quantize::bbq::bbq_code_layout; +using cuvs::preprocessing::quantize::bbq::quantizer_view; + template struct ResultItem; @@ -423,7 +430,9 @@ __device__ ResultItem get_min_item(const Index_t id, const int idx_in_list, const Index_t* neighbs, const DistData_t* distances, - const bool find_in_row = true) + const bool find_in_row = true, + const int stride = SKEWED_MAX_NUM_BI_SAMPLES, + const int neighbs_size = MAX_NUM_BI_SAMPLES) { int lane_id = threadIdx.x % raft::warp_size(); @@ -434,16 +443,18 @@ __device__ ResultItem get_min_item(const Index_t id, idx[0] = lane_id; idx[1] = raft::warp_size() + lane_id; - if (neighbs[idx[0]] != id) { - dist[0] = find_in_row ? distances[idx_in_list * SKEWED_MAX_NUM_BI_SAMPLES + lane_id] - : distances[idx_in_list + lane_id * SKEWED_MAX_NUM_BI_SAMPLES]; + // neighbs_size defaults to the full width, so existing (dense/calculate_metric-based) callers + // that pre-fill `distances` out to MAX_NUM_BI_SAMPLES are unaffected. Callers that don't pre-fill + // (e.g. the BBQ SIMT kernel, which only ever writes real cells) pass the real list size instead, + // so out-of-range entries just never get read -- cheaper than a separate fill pass. + if (idx[0] < neighbs_size && neighbs[idx[0]] != id) { + dist[0] = find_in_row ? distances[idx_in_list * stride + lane_id] + : distances[idx_in_list + lane_id * stride]; } - if (neighbs[idx[1]] != id) { - dist[1] = - find_in_row - ? distances[idx_in_list * SKEWED_MAX_NUM_BI_SAMPLES + raft::warp_size() + lane_id] - : distances[idx_in_list + (raft::warp_size() + lane_id) * SKEWED_MAX_NUM_BI_SAMPLES]; + if (idx[1] < neighbs_size && neighbs[idx[1]] != id) { + dist[1] = find_in_row ? distances[idx_in_list * stride + raft::warp_size() + lane_id] + : distances[idx_in_list + (raft::warp_size() + lane_id) * stride]; } if (dist[1] < dist[0]) { @@ -523,7 +534,7 @@ __device__ __forceinline__ void calculate_metric(float* s_distances, { // if we have a distance epilogue, distances need to be fully calculated instead of postprocessing // them. - bool can_postprocess_dist = std::is_same_v; + constexpr bool can_postprocess_dist = std::is_same_v; for (int i = threadIdx.x; i < MAX_NUM_BI_SAMPLES * SKEWED_MAX_NUM_BI_SAMPLES; i += blockDim.x) { int row_id = i / SKEWED_MAX_NUM_BI_SAMPLES; @@ -583,6 +594,11 @@ struct DistAccumulator { // SIMT kernel: scalar element-wise distance computation. // Used for fp32 data (all metrics) and L1 distance computation for all dtypes (which cannot use // tensor cores). +// -------------------------------------------------------------------------- +// Dense local-join kernels +// Full-precision joins over the original vectors: SIMT and fp16 tensor core. +// -------------------------------------------------------------------------- + template , @@ -848,13 +864,6 @@ __launch_bounds__(BLOCK_SIZE) #endif } -// launch_bounds here denote BLOCK_SIZE = 512 and MIN_BLOCKS_PER_SM = 4 -// Per -// https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#features-and-technical-specifications, -// MAX_RESIDENT_THREAD_PER_SM = BLOCK_SIZE * BLOCKS_PER_SM = 2048 -// For architectures 750 and 860 (890), the values for MAX_RESIDENT_THREAD_PER_SM -// is 1024 and 1536 respectively, which means the bounds don't work anymore -// Used for fp32 data downcast to fp16, and all types using non-L1 distance metric. template , @@ -1115,6 +1124,905 @@ __launch_bounds__(BLOCK_SIZE) } #endif } +// -------------------------------------------------------------------------- +// BBQ local-join kernels +// Quantized-code joins: SIMT (popc / dp4a) and int4 tensor core (u4 wmma). +// -------------------------------------------------------------------------- + +// The code promotion, the dequant factors and the raw-dot-product-to-distance conversion live +// next to the code inner products in bbq.cuh, so the CAGRA graph sort ranks by the very same +// distance these kernels build the lists with. +using cuvs::preprocessing::quantize::bbq::bbq_calculate_metric; +using cuvs::preprocessing::quantize::bbq::bbq_dequant_factors; +using cuvs::preprocessing::quantize::bbq::get_dequant_factors; +using cuvs::preprocessing::quantize::bbq::packed_1b_to_4b; + +// Stages one K-tile of `count` neighbor rows into SMEM. Shared skeleton: one warp per row, +// lanes strided along native uint32 words, last-tile zero-pad so compute reads a full tile. +// +// stage_tile_simt copies native codes in storage layout (Planes > 1 gathers bit-sliced planes). +// stage_promoted_tile expands packed_1b to u4 for int4 MMA (packed_4b is a plain copy). +template +__device__ __forceinline__ void stage_tile_simt(uint8_t (*dst)[RowStride], + const quantizer_view& quantizer, + const Index_t* neighbors, + const int count, + const size_t base, + const int plane_extent, + const int num_load_u32, + const int plane_tile_u32, + const bool last_tile, + const int warp_id, + const int lane_id) +{ + constexpr int num_warps = BLOCK_SIZE / raft::warp_size(); + for (int i = 0; i < MAX_NUM_BI_SAMPLES / num_warps; ++i) { + const int idx = i * num_warps + warp_id; + if (idx >= count) continue; + auto* s = reinterpret_cast(dst[idx]); + const auto* src = reinterpret_cast(&quantizer.codes(neighbors[idx], base)); + for (int w = lane_id; w < num_load_u32; w += raft::warp_size()) { +#pragma unroll + for (int p = 0; p < Planes; ++p) { + s[p * plane_tile_u32 + w] = src[p * plane_extent + w]; + } + } + if (last_tile) { + for (int w = num_load_u32 + lane_id; w < plane_tile_u32; w += raft::warp_size()) { +#pragma unroll + for (int p = 0; p < Planes; ++p) { + s[p * plane_tile_u32 + w] = 0; + } + } + } + } +} + +// `native_row_bytes` is the layout's own encoded row length. One native tile is +// RowStride-independent: BBQ_ROW_BYTES / expansion native bytes promote to exactly BBQ_ROW_BYTES +// promoted bytes, so a single n_tiles drives every operand regardless of how compact each one's +// on-disk format is. +template +__device__ __forceinline__ void stage_promoted_tile(uint8_t (*dst)[RowStride], + const quantizer_view& quantizer, + const Index_t* neighbors, + const int count, + const int step, + const int native_row_bytes, + const int warp_id, + const int lane_id) +{ + static_assert(Layout == bbq_code_layout::packed_1b || Layout == bbq_code_layout::packed_4b, + "int4 MMA path supports packed_1b (1b), packed_4b (4b)"); + // A u4 MMA fragment needs 4 bits per value, so a layout storing `bits` bits per value expands + // one native word into 4/bits promoted words. packed_4b is the identity case (a plain word + // copy), which is why the symmetric kernel needs no separate "no promotion" path -- and why a + // promoted query works exactly like a promoted document. + constexpr int expansion = Layout == bbq_code_layout::packed_1b ? 4 : 1; + constexpr int native_tile = TileBytes / expansion; + constexpr int native_tile_u32 = native_tile / 4; + constexpr int num_warps = BLOCK_SIZE / raft::warp_size(); + // packed_4b's native tile is 32 u32s (one per lane). packed_1b's is 8, so a warp-per-row + // map would leave 24 lanes idle; instead pack warp_size/8 = 4 rows into one pass. + static_assert(raft::warp_size() % native_tile_u32 == 0); + constexpr int rows_per_pass = raft::warp_size() / native_tile_u32; + static_assert(MAX_NUM_BI_SAMPLES % (num_warps * rows_per_pass) == 0); + + const int base = step * native_tile; + const int remaining = native_row_bytes - base; + const int num_load_u32 = + (remaining < native_tile ? (remaining > 0 ? remaining : 0) : native_tile) / 4; + + for (int i = 0; i < MAX_NUM_BI_SAMPLES / (num_warps * rows_per_pass); ++i) { + const int idx = (i * num_warps + warp_id) * rows_per_pass + lane_id / native_tile_u32; + const int w = lane_id % native_tile_u32; + if (idx >= count) continue; + auto* s = reinterpret_cast(dst[idx]); + const auto* src = reinterpret_cast(&quantizer.codes(neighbors[idx], base)); + if (w < num_load_u32) { + if constexpr (expansion == 4) { + reinterpret_cast(s)[w] = packed_1b_to_4b(src[w]); + } else { + s[w] = src[w]; + } + } else { + if constexpr (expansion == 4) { + reinterpret_cast(s)[w] = uint4{0, 0, 0, 0}; + } else { + s[w] = 0; + } + } + } +} + +// Stages this block's new/old neighbour lists into shared memory and drops duplicates. Identical +// in the SIMT and wmma BBQ kernels, so it lives here rather than twice. The caller keeps the +// `if (!new_size) return;` early-out, which has to stay at kernel scope. +template +__device__ __forceinline__ void stage_neighbor_lists(Index_t* new_neighbors, + Index_t* old_neighbors, + int* s_unique_counter, + const Index_t* graph_new, + const Index_t* rev_graph_new, + const Index_t* graph_old, + const Index_t* rev_graph_old, + const size_t list_id, + const int width, + const int2 new_size2, + const int2 old_size2, + int& new_size, + int& old_size) +{ + const int tx = threadIdx.x; + if (tx < new_size2.x) { + new_neighbors[tx] = graph_new[list_id * width + tx]; + } else if (tx < new_size) { + new_neighbors[tx] = rev_graph_new[list_id * width + tx - new_size2.x]; + } + if (tx < old_size2.x) { + old_neighbors[tx] = graph_old[list_id * width + tx]; + } else if (tx < old_size) { + old_neighbors[tx] = rev_graph_old[list_id * width + tx - old_size2.x]; + } + __syncthreads(); + + remove_duplicates( + new_neighbors, new_size2.x, new_neighbors + new_size2.x, new_size2.y, s_unique_counter[0], 0); + remove_duplicates( + old_neighbors, old_size2.x, old_neighbors + old_size2.x, old_size2.y, s_unique_counter[1], 1); + __syncthreads(); + new_size = new_size2.x + s_unique_counter[0]; + old_size = old_size2.x + s_unique_counter[1]; +} + +template , + typename DistEpilogue_t> +RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) + local_join_kernel_bbq_simt(const Index_t* graph_new, + const Index_t* rev_graph_new, + const int2* sizes_new, + const Index_t* graph_old, + const Index_t* rev_graph_old, + const int2* sizes_old, + const int width, + quantizer_view dataset_document, + quantizer_view dataset_query, + ID_t* graph, + DistData_t* dists, + int graph_width, + int* locks, + cuvs::distance::DistanceType metric, + DistEpilogue_t dist_epilogue) +{ + constexpr int document_planes = + cuvs::preprocessing::quantize::bbq::get_code_planes(DocumentLayout); + constexpr int query_planes = cuvs::preprocessing::quantize::bbq::get_code_planes(QueryLayout); + static_assert(!SelfJoin || DocumentLayout == QueryLayout, + "a self-join must use the same layout on both operands"); + + // Both operands are tiled at the same per-plane tile so each step covers the same dimension + // range on both sides. QUERY_ROW_BYTES fixes the query row width; the document row width then + // follows as QUERY_ROW_BYTES / (query_planes / document_planes), i.e. query_plane_tile scaled + // by the document's own plane count. Worked out per supported pair: + // + // pair query_plane_tile doc_row_bytes doc_stride query_stride + // ------ ---------------- -------------- ---------- ------------ + // 1 x 1 128 128 * 1 = 128 128 128 + // 2t x 2t 64 64 * 2 = 128 64 64 + // 1 + 2t 64 64 * 1 = 64 64 64 + // 1 + 4t 32 32 * 1 = 32 32 32 + // 2t + 4t 32 32 * 2 = 64 32 32 + // + constexpr int QUERY_ROW_BYTES = 128; + constexpr int BBQ_PAD = alignof(uint32_t); + // The document buffer is normally only the A operand (two rows broadcast across a warp), so it + // needs no bank-conflict pad. Under SelfJoin it doubles as the B operand in phase 1 (32 + // consecutive columns at one byte offset), so it needs the same skew the query buffer gets. + constexpr int DOC_PAD = SelfJoin ? BBQ_PAD : 0; + static_assert((QUERY_ROW_BYTES + BBQ_PAD) % alignof(uint32_t) == 0); + + __shared__ int s_list[MAX_NUM_BI_SAMPLES * 2]; + __shared__ __align__(alignof(uint32_t)) uint8_t + s_doc_vec[MAX_NUM_BI_SAMPLES][QUERY_ROW_BYTES / query_planes * document_planes + DOC_PAD]; + __shared__ __align__(alignof(uint32_t)) + uint8_t s_query_vec[MAX_NUM_BI_SAMPLES][QUERY_ROW_BYTES + BBQ_PAD]; + // Holds the final (post-metric, post-epilogue) float distance per cell + __shared__ DistData_t s_distances[MAX_NUM_BI_SAMPLES * SKEWED_MAX_NUM_BI_SAMPLES]; + __shared__ int s_unique_counter[2]; + // Document-side (row axis) dequant factors, indexed by list position -- shared by both phases + // below since rows are always the `new_neighbors`/document side in both, so this is staged once + // and never restaged. At a given accumulation step, up to MAX_NUM_BI_SAMPLES threads share the + // same row0 (see the pair_idx assignment below), so caching avoids up to 64x redundant scattered + // global reads + __shared__ bbq_dequant_factors s_document_factors[MAX_NUM_BI_SAMPLES]; + + if (threadIdx.x == 0) { + s_unique_counter[0] = 0; + s_unique_counter[1] = 0; + } + + Index_t* new_neighbors = s_list; + Index_t* old_neighbors = s_list + MAX_NUM_BI_SAMPLES; + const size_t list_id = blockIdx.x; + const int2 new_size2 = sizes_new[list_id]; + const int2 old_size2 = sizes_old[list_id]; + int new_size = new_size2.x + new_size2.y; + int old_size = old_size2.x + old_size2.y; + const int tx = threadIdx.x; + + if (!new_size) return; + stage_neighbor_lists(new_neighbors, + old_neighbors, + s_unique_counter, + graph_new, + rev_graph_new, + graph_old, + rev_graph_old, + list_id, + width, + new_size2, + old_size2, + new_size, + old_size); + + const int warp_id = threadIdx.x / raft::warp_size(); + const int lane_id = threadIdx.x % raft::warp_size(); + constexpr int num_warps = BLOCK_SIZE / raft::warp_size(); + + // Rows are always new_neighbors/document in both phases below, so this runs exactly once. + for (int i = tx; i < new_size; i += BLOCK_SIZE) { + s_document_factors[i] = get_dequant_factors(dataset_document, new_neighbors[i]); + } + __syncthreads(); + + // Each plane gets an equal slice of the row in shared memory, so the cached bytes always form a + // valid encoded chunk. + // Bytes per plane = encoded row length / plane count. Do NOT assume ceildiv(dim, 8): that is + // bytes-per-plane only for the bit-plane layouts (packed_1b, transposed_2b, transposed_4b), + // where it happens to equal encoded/planes for all three. The dense byte layouts + // (packed_7b/packed_8b) are `dim` bytes in one plane, and would read 1/8 of each row. + const int plane_bytes = + static_cast(cuvs::preprocessing::quantize::bbq::get_encoded_row_length(dataset_document)) / + document_planes; + assert(plane_bytes == static_cast(cuvs::preprocessing::quantize::bbq::get_encoded_row_length( + dataset_query)) / + query_planes); + constexpr int query_plane_tile = QUERY_ROW_BYTES / query_planes; + constexpr int plane_tile = query_plane_tile; + constexpr int doc_row_bytes = query_plane_tile * document_planes; + static_assert(plane_tile % 4 == 0, "plane_tile must be 4-byte aligned for uint32 loads"); + // Row strides too: rows are indexed as base + idx * stride and then read as uint32_t, so a + // stride that is not a multiple of 4 misaligns every odd row. Derived from QUERY_ROW_BYTES, so + // this is what catches an ill-chosen QUERY_ROW_BYTES rather than letting it fault at runtime. + static_assert((doc_row_bytes + DOC_PAD) % alignof(uint32_t) == 0, + "document row stride must be 4-byte aligned for uint32 loads"); + static_assert( + doc_row_bytes % document_planes == 0 && doc_row_bytes / document_planes == query_plane_tile, + "document plane stride must match query plane stride"); + // plane_bytes is the per-plane stride in bytes; plane_extent is the same in uint32 elements, + // computed once so call sites don't re-derive it. Alignment (plane_bytes % 4 == 0, i.e. + // dataset dim % 32 == 0) is enforced by the launcher. + const int plane_extent = plane_bytes / 4; + constexpr int plane_tile_u32 = plane_tile / 4; + constexpr int query_plane_tile_u32 = query_plane_tile / 4; + const int n_tiles = raft::ceildiv(plane_bytes, plane_tile); + + // with NUM_SAMPLES=32, BLOCK_SIZE=256, pairs_per_thread = 32 * 64 / 512 = 4 + constexpr int num_row_pairs = MAX_NUM_BI_SAMPLES / 2; + constexpr int num_pairs = num_row_pairs * MAX_NUM_BI_SAMPLES; + constexpr int pairs_per_thread = (num_pairs + BLOCK_SIZE - 1) / BLOCK_SIZE; + + // Every thread's column is invariant across its pairs_per_thread iterations, so fetch it once + // here and reuse the register copy in both phases' store loops below instead of a shared array. + const int my_col = tx % MAX_NUM_BI_SAMPLES; + bbq_dequant_factors my_col_factors{}; + if (my_col < new_size) { + my_col_factors = get_dequant_factors(dataset_query, new_neighbors[my_col]); + } + + uint32_t acc0[pairs_per_thread] = {}; + uint32_t acc1[pairs_per_thread] = {}; + for (int step = 0; step < n_tiles; ++step) { + const bool last_tile = (step == n_tiles - 1); + const int num_load = last_tile ? plane_bytes - step * plane_tile : plane_tile; + const int num_load_u32 = num_load / 4; + const size_t base = static_cast(step) * plane_tile; + stage_tile_simt(s_doc_vec, + dataset_document, + new_neighbors, + new_size, + base, + plane_extent, + num_load_u32, + plane_tile_u32, + last_tile, + warp_id, + lane_id); + + // Query and document tiles cover the same dimension range per step (both tile at + // query_plane_tile), so load the query tile once and run the dot product directly -- no + // per-step query sub-tile loop. Under SelfJoin phase 1 is new x new on a single quantizer, + // so the document buffer already holds exactly what the query buffer would: skip the load + // and point the B operand at s_doc_vec. + if constexpr (!SelfJoin) { + stage_tile_simt(s_query_vec, + dataset_query, + new_neighbors, + new_size, + base, + plane_extent, + num_load_u32, + query_plane_tile_u32, + last_tile, + warp_id, + lane_id); + } + __syncthreads(); + +#pragma unroll + for (int k = 0; k < pairs_per_thread; ++k) { + const int pair_idx = tx + k * BLOCK_SIZE; + if (pair_idx >= num_pairs) continue; + const int row0 = (pair_idx / MAX_NUM_BI_SAMPLES) * 2; + const int col = pair_idx % MAX_NUM_BI_SAMPLES; + if (col >= new_size) continue; + // Phase 1 B operand: s_doc_vec under SelfJoin (see the staging note above). + const uint8_t* row_b; + if constexpr (SelfJoin) { + row_b = s_doc_vec[col]; + } else { + row_b = s_query_vec[col]; + } + cuvs::preprocessing::quantize::bbq::bbq_code_inner_product_2x1( + s_doc_vec[row0], s_doc_vec[row0 + 1], row_b, acc0[k], acc1[k]); + } + __syncthreads(); + } +#pragma unroll + for (int k = 0; k < pairs_per_thread; ++k) { + const int pair_idx = tx + k * BLOCK_SIZE; + if (pair_idx >= num_pairs) continue; + const int row0 = (pair_idx / MAX_NUM_BI_SAMPLES) * 2; + const int col = pair_idx % MAX_NUM_BI_SAMPLES; + if (col >= new_size) continue; + const int distance0 = row0 * SKEWED_MAX_NUM_BI_SAMPLES + col; + const Index_t query_id = new_neighbors[col]; + // row0 (unlike row0 + 1 below) isn't otherwise bounded by new_size -- gate it explicitly so a + // stale/garbage new_neighbors[row0] past the real list never reaches a dataset_document lookup + // (row_norm, for CosineExpanded); the cell is provably never read downstream either way, since + // the min-search loops below only visit idx_in_list < new_size. + if (row0 < new_size) { + const Index_t doc_id0 = new_neighbors[row0]; + s_distances[distance0] = bbq_calculate_metric(acc0[k], + s_document_factors[row0], + my_col_factors, + dataset_document, + dataset_query, + metric, + dist_epilogue, + doc_id0, + query_id); + } + if (row0 + 1 < new_size) { + const Index_t doc_id1 = new_neighbors[row0 + 1]; + s_distances[distance0 + SKEWED_MAX_NUM_BI_SAMPLES] = + bbq_calculate_metric(acc1[k], + s_document_factors[row0 + 1], + my_col_factors, + dataset_document, + dataset_query, + metric, + dist_epilogue, + doc_id1, + query_id); + } + } + __syncthreads(); + + for (int step = 0; step < raft::ceildiv(new_size, num_warps); ++step) { + const int idx_in_list = step * num_warps + tx / raft::warp_size(); + if (idx_in_list >= new_size) continue; + auto min_elem = get_min_item(s_list[idx_in_list], + idx_in_list, + new_neighbors, + s_distances, + true, + SKEWED_MAX_NUM_BI_SAMPLES, + new_size); + if (min_elem.id() < gridDim.x) { + insert_to_global_graph(min_elem, s_list[idx_in_list], graph, dists, graph_width, locks); + } + } + + if (!old_size) return; + __syncthreads(); + + // Phase 2's column list is old_neighbors, not new_neighbors -- refetch this thread's column + // factors (rows/s_document_factors stay valid unchanged, since rows are new_neighbors in both + // phases). + const int my_old_col = tx % MAX_NUM_BI_SAMPLES; + bbq_dequant_factors my_old_col_factors{}; + if (my_old_col < old_size) { + my_old_col_factors = get_dequant_factors(dataset_query, old_neighbors[my_old_col]); + } + + uint32_t acc0_old[pairs_per_thread] = {}; + uint32_t acc1_old[pairs_per_thread] = {}; + for (int step = 0; step < n_tiles; ++step) { + const bool last_tile = (step == n_tiles - 1); + const int num_load = last_tile ? plane_bytes - step * plane_tile : plane_tile; + const int num_load_u32 = num_load / 4; + const size_t base = static_cast(step) * plane_tile; + if (n_tiles > 1) { + stage_tile_simt(s_doc_vec, + dataset_document, + new_neighbors, + new_size, + base, + plane_extent, + num_load_u32, + plane_tile_u32, + last_tile, + warp_id, + lane_id); + } + stage_tile_simt(s_query_vec, + dataset_query, + old_neighbors, + old_size, + base, + plane_extent, + num_load_u32, + query_plane_tile_u32, + last_tile, + warp_id, + lane_id); + __syncthreads(); + +#pragma unroll + for (int k = 0; k < pairs_per_thread; ++k) { + const int pair_idx = tx + k * BLOCK_SIZE; + if (pair_idx >= num_pairs) continue; + const int row0 = (pair_idx / MAX_NUM_BI_SAMPLES) * 2; + const int col = pair_idx % MAX_NUM_BI_SAMPLES; + if (col >= old_size) continue; + cuvs::preprocessing::quantize::bbq::bbq_code_inner_product_2x1( + s_doc_vec[row0], s_doc_vec[row0 + 1], s_query_vec[col], acc0_old[k], acc1_old[k]); + } + __syncthreads(); + } +#pragma unroll + for (int k = 0; k < pairs_per_thread; ++k) { + const int pair_idx = tx + k * BLOCK_SIZE; + if (pair_idx >= num_pairs) continue; + const int row0 = (pair_idx / MAX_NUM_BI_SAMPLES) * 2; + const int col = pair_idx % MAX_NUM_BI_SAMPLES; + if (col >= old_size) continue; + const int distance0 = row0 * SKEWED_MAX_NUM_BI_SAMPLES + col; + const Index_t query_id = old_neighbors[col]; + // See the identical row0 < new_size gate in phase 1 above -- same reasoning (row0 isn't + // otherwise bounded, and this cell is never read downstream either way). + if (row0 < new_size) { + const Index_t doc_id0 = new_neighbors[row0]; + s_distances[distance0] = bbq_calculate_metric(acc0_old[k], + s_document_factors[row0], + my_old_col_factors, + dataset_document, + dataset_query, + metric, + dist_epilogue, + doc_id0, + query_id); + } + if (row0 + 1 < new_size) { + const Index_t doc_id1 = new_neighbors[row0 + 1]; + s_distances[distance0 + SKEWED_MAX_NUM_BI_SAMPLES] = + bbq_calculate_metric(acc1_old[k], + s_document_factors[row0 + 1], + my_old_col_factors, + dataset_document, + dataset_query, + metric, + dist_epilogue, + doc_id1, + query_id); + } + } + __syncthreads(); + + for (int step = 0; step < raft::ceildiv(MAX_NUM_BI_SAMPLES, num_warps); ++step) { + const int idx_in_list = step * num_warps + tx / raft::warp_size(); + if (idx_in_list >= new_size) continue; + auto min_elem = get_min_item(s_list[idx_in_list], + idx_in_list, + old_neighbors, + s_distances, + true, + SKEWED_MAX_NUM_BI_SAMPLES, + old_size); + if (min_elem.id() < gridDim.x) { + insert_to_global_graph(min_elem, s_list[idx_in_list], graph, dists, graph_width, locks); + } + } + + for (int step = 0; step < raft::ceildiv(MAX_NUM_BI_SAMPLES, num_warps); ++step) { + const int idx_in_list = step * num_warps + tx / raft::warp_size(); + if (idx_in_list >= old_size) continue; + const int list_idx = idx_in_list + MAX_NUM_BI_SAMPLES; + auto min_elem = get_min_item(s_list[list_idx], + idx_in_list, + new_neighbors, + s_distances, + false, + SKEWED_MAX_NUM_BI_SAMPLES, + new_size); + if (min_elem.id() < gridDim.x) { + insert_to_global_graph(min_elem, s_list[list_idx], graph, dists, graph_width, locks); + } + } +} + +// int4 tensor-core BBQ local join, covering both the symmetric (one quantizer, self-join) and +// asymmetric (two quantizers) cases. Modeled on local_join_kernel_wmma: nvcuda::wmma fragments +// (u4 x u4 -> s32, shape m8n8k32) replace the popc/dp4a inner product; the accumulator lives in +// registers across the whole K reduction and is stored to s_distances once per (row-tile, +// col-tile), not accumulated into shared memory every step like the scalar kernel. +// +// Warp tiling: num_warps = BLOCK_SIZE/32 warps arranged as a WARPS_PER_DIM x WARPS_PER_DIM square +// grid (WARPS_PER_DIM=4 so 4x4=16=num_warps), each warp owning a (MAX_NUM_BI_SAMPLES/WARPS_PER_DIM) +// region of the MAX_NUM_BI_SAMPLES x MAX_NUM_BI_SAMPLES output matrix, same as +// local_join_kernel_wmma's WMMA_M=N=16 warp assignment. Since int4 MMA tiles are MMA_M x MMA_N +// (8x8, the only shape nvcuda::wmma exposes for u4), each warp covers its region via a +// SUB_PER_DIM x SUB_PER_DIM grid of native tiles instead of a single call -- SUB_PER_DIM is a +// forced consequence of (MAX_NUM_BI_SAMPLES/MMA_M) / WARPS_PER_DIM, not an arbitrary choice. +template , + typename DistEpilogue_t> +RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) + local_join_kernel_bbq_wmma(const Index_t* graph_new, + const Index_t* rev_graph_new, + const int2* sizes_new, + const Index_t* graph_old, + const Index_t* rev_graph_old, + const int2* sizes_old, + const int width, + const quantizer_view dataset_document, + const quantizer_view dataset_query, + ID_t* graph, + DistData_t* dists, + int graph_width, + int* locks, + cuvs::distance::DistanceType metric, + DistEpilogue_t dist_epilogue) +{ +// int4 sub-byte MMA (nvcuda::wmma experimental::precision::u4) still compiles on every Blackwell +// variant (sm_100/103/110/120/121, verified by disassembly): ptxas lowers it to the same software +// path it's always used since Turing -- unpack each u4 nibble pair into two u8 operands and run +// two native u8 IMMA instructions, summing the partial products. +#if (__CUDA_ARCH__ >= 750) + using namespace nvcuda; + constexpr int MMA_M = 8; + constexpr int MMA_N = 8; + constexpr int MMA_K = 32; + // num_warps = BLOCK_SIZE/32 = 16, arranged as a square WARPS_PER_DIM x WARPS_PER_DIM grid since + // 4*4=16 matches exactly; the static_assert is what actually enforces this holds for the + // current BLOCK_SIZE, WARPS_PER_DIM itself isn't derived (no trivial constexpr integer sqrt). + constexpr int WARPS_PER_DIM = 4; + static_assert(WARPS_PER_DIM * WARPS_PER_DIM == BLOCK_SIZE / raft::warp_size(), + "warp grid must be square and match num_warps = BLOCK_SIZE/32"); + // Each warp owns a WARP_TILE x WARP_TILE region of the MAX_NUM_BI_SAMPLES x MAX_NUM_BI_SAMPLES + // output matrix. TILES_PER_DIM is how many native MMA_M x MMA_N tiles span one output dimension; + // SUB_PER_DIM (native tiles per warp per dim) is a forced consequence of TILES_PER_DIM / + // WARPS_PER_DIM, not an arbitrary choice -- it's 2 here only because 8/4=2 for these particular + // MAX_NUM_BI_SAMPLES/MMA_M/WARPS_PER_DIM values. + static_assert(MAX_NUM_BI_SAMPLES % MMA_M == 0 && MMA_M == MMA_N, + "MAX_NUM_BI_SAMPLES must divide evenly into square MMA_MxMMA_N tiles"); + constexpr int TILES_PER_DIM = MAX_NUM_BI_SAMPLES / MMA_M; + static_assert(TILES_PER_DIM % WARPS_PER_DIM == 0, + "warps must evenly tile the native MMA tiles in each output dimension"); + constexpr int SUB_PER_DIM = TILES_PER_DIM / WARPS_PER_DIM; + constexpr int WARP_TILE = SUB_PER_DIM * MMA_M; + + // Promoted (4-bit-width) staging tile, 128 B/row. Row stride is BBQ_ROW_BYTES + MMA_PAD, not + // just BBQ_ROW_BYTES: sub-byte IMMA loads need at least 16-byte row alignment, and MMA_PAD must + // be a multiple of 16 to preserve that -- but BBQ_ROW_BYTES=128 alone is *also* exactly 32 + // shared-memory banks (4 B/bank), so every row would land on the same bank offset and any + // multi-row access load_matrix_sync does internally would conflict. MMA_PAD=16 breaks that + // exact-32-bank alignment (144 B/row is not a multiple of 128 B) while staying a multiple of 16 + // for the IMMA alignment requirement. + constexpr int BBQ_ROW_BYTES = 128; + constexpr int MMA_PAD = 16; + static_assert(MMA_PAD % 16 == 0, "row padding must preserve 16-byte IMMA row alignment"); + constexpr int ELEMS_PER_TILE = BBQ_ROW_BYTES * 2; // 2 u4 elements/byte + constexpr int K_STEPS_PER_TILE = ELEMS_PER_TILE / MMA_K; + constexpr int ROW_STRIDE_U4 = (BBQ_ROW_BYTES + MMA_PAD) * 2; // row-to-row stride, u4 elements + + constexpr int MMA_STORE_STRIDE = SKEWED_MAX_NUM_BI_SAMPLES; + + // s_row_vec is the A operand: always the `new` list, document quantizer. s_col_vec is the B + // operand: the `new` list in phase 1 and the `old` list in phase 2, query quantizer. Under + // SelfJoin, phase 1 leaves s_col_vec untouched and both fragments read s_row_vec. + __shared__ int s_list[MAX_NUM_BI_SAMPLES * 2]; + __shared__ __align__(16) uint8_t s_row_vec[MAX_NUM_BI_SAMPLES][BBQ_ROW_BYTES + MMA_PAD]; + // s_col_vec aliases s_distances's memory instead of its own array: their live ranges never + // overlap (col_vec fully consumed by the kk-loop before distances are stored; distances fully + // consumed by the min-search loop before the next phase re-stages col_vec). + __shared__ __align__(16) DistData_t s_distances[MAX_NUM_BI_SAMPLES * MMA_STORE_STRIDE]; + static_assert(sizeof(s_distances) >= MAX_NUM_BI_SAMPLES * (BBQ_ROW_BYTES + MMA_PAD), + "s_col_vec aliases s_distances's memory and must fit inside it"); + auto(*s_col_vec)[BBQ_ROW_BYTES + MMA_PAD] = + reinterpret_cast(s_distances); + __shared__ int s_unique_counter[2]; + // Document-side (row axis) dequant factors -- see the identical buffer and full rationale in + // local_join_kernel_bbq_simt. Staged once below, reused unchanged by both phases (rows are + // always new_neighbors/document in both). + __shared__ bbq_dequant_factors s_document_factors[MAX_NUM_BI_SAMPLES]; + + if (threadIdx.x == 0) { + s_unique_counter[0] = 0; + s_unique_counter[1] = 0; + } + + Index_t* new_neighbors = s_list; + Index_t* old_neighbors = s_list + MAX_NUM_BI_SAMPLES; + const size_t list_id = blockIdx.x; + const int2 new_size2 = sizes_new[list_id]; + const int2 old_size2 = sizes_old[list_id]; + int new_size = new_size2.x + new_size2.y; + int old_size = old_size2.x + old_size2.y; + const int tx = threadIdx.x; + + if (!new_size) return; + stage_neighbor_lists(new_neighbors, + old_neighbors, + s_unique_counter, + graph_new, + rev_graph_new, + graph_old, + rev_graph_old, + list_id, + width, + new_size2, + old_size2, + new_size, + old_size); + + const int warp_id = threadIdx.x / raft::warp_size(); + const int lane_id = threadIdx.x % raft::warp_size(); + constexpr int num_warps = BLOCK_SIZE / raft::warp_size(); + // n_tiles is driven by the *promoted* (4-bit-width) row length, which both operands share at a + // given dim; each side's own native row length only sets how many native bytes + // stage_promoted_tile reads per tile (dim % 32 == 0 guarantees no rounding either way). + const int doc_row_bytes = + static_cast(cuvs::preprocessing::quantize::bbq::get_encoded_row_length(dataset_document)); + const int query_row_bytes = + static_cast(cuvs::preprocessing::quantize::bbq::get_encoded_row_length(dataset_query)); + const int promoted_row_bytes = static_cast((dataset_query.dim() + 1) / 2); + const int n_tiles = raft::ceildiv(promoted_row_bytes, BBQ_ROW_BYTES); + + const int warp_id_y = warp_id / WARPS_PER_DIM; + const int warp_id_x = warp_id % WARPS_PER_DIM; + + // One phase: accumulate s_row_vec (rows, `new` list) against col_buf over the whole K range, + // then store the accumulators to s_distances (converted to final float distances in place right + // after). col_buf is s_row_vec itself when phase 1 is a + // self-join, otherwise s_col_vec. col_neighbors/col_size select which list the B operand stages. + // alias_tag: staging skip and b_frag source + // row_resident: the A operand is already staged from a previous phase, which is only true + // when the whole row fits one tile (n_tiles == 1) so nothing overwrote it. + auto run_phase = + [&](const Index_t* col_neighbors, int col_size, auto alias_tag, bool row_resident) { + constexpr bool alias_col = decltype(alias_tag)::value; + wmma::fragment c_frag[SUB_PER_DIM][SUB_PER_DIM]; +#pragma unroll + for (int msub = 0; msub < SUB_PER_DIM; ++msub) { +#pragma unroll + for (int nsub = 0; nsub < SUB_PER_DIM; ++nsub) { + wmma::fill_fragment(c_frag[msub][nsub], 0); + } + } + + for (int step = 0; step < n_tiles; ++step) { + if (!row_resident) { + stage_promoted_tile(s_row_vec, + dataset_document, + new_neighbors, + new_size, + step, + doc_row_bytes, + warp_id, + lane_id); + } + if constexpr (!alias_col) { + stage_promoted_tile(s_col_vec, + dataset_query, + col_neighbors, + col_size, + step, + query_row_bytes, + warp_id, + lane_id); + } + __syncthreads(); + + // a_frag depends only on (msub, kk); b_frag depends only on (nsub, kk) -- load each once + // per kk and reuse across the other sub-tile index, instead of reloading redundantly inside + // a full msub x nsub x kk cross product. + // Deliberately not #pragma unroll'd: full unrolling here keeps more fragment live ranges + // simultaneous, driving register pressure up (64/thread, tied with SMEM for the occupancy + // cap) -- letting the compiler pick reduces that at the cost of some intra-warp ILP. + for (int kk = 0; kk < K_STEPS_PER_TILE; ++kk) { + wmma::fragment + a_frag[SUB_PER_DIM]; + wmma::fragment + b_frag[SUB_PER_DIM]; + const auto* col_buf = alias_col ? s_row_vec : s_col_vec; // compile-time select +#pragma unroll + for (int msub = 0; msub < SUB_PER_DIM; ++msub) { + const int row0 = warp_id_y * WARP_TILE + msub * MMA_M; + wmma::load_matrix_sync(a_frag[msub], s_row_vec[row0] + kk * (MMA_K / 2), ROW_STRIDE_U4); + } +#pragma unroll + for (int nsub = 0; nsub < SUB_PER_DIM; ++nsub) { + const int col0 = warp_id_x * WARP_TILE + nsub * MMA_N; + wmma::load_matrix_sync(b_frag[nsub], col_buf[col0] + kk * (MMA_K / 2), ROW_STRIDE_U4); + } +#pragma unroll + for (int msub = 0; msub < SUB_PER_DIM; ++msub) { +#pragma unroll + for (int nsub = 0; nsub < SUB_PER_DIM; ++nsub) { + wmma::mma_sync(c_frag[msub][nsub], a_frag[msub], b_frag[nsub], c_frag[msub][nsub]); + } + } + } + __syncthreads(); + } + +#pragma unroll + for (int msub = 0; msub < SUB_PER_DIM; ++msub) { + const int row0 = warp_id_y * WARP_TILE + msub * MMA_M; +#pragma unroll + for (int nsub = 0; nsub < SUB_PER_DIM; ++nsub) { + const int col0 = warp_id_x * WARP_TILE + nsub * MMA_N; + wmma::store_matrix_sync( + reinterpret_cast(s_distances) + row0 * MMA_STORE_STRIDE + col0, + c_frag[msub][nsub], + MMA_STORE_STRIDE, + wmma::mem_row_major); + } + } + __syncthreads(); + + // Converts store_matrix_sync's raw int32 dot products into final float distances, in place + // col = i % MAX_NUM_BI_SAMPLES is invariant across a thread's own iterations + // (BLOCK_SIZE is a multiple of MAX_NUM_BI_SAMPLES), so this thread's column factors are + // fetched once and reused below + const int my_col = tx % MAX_NUM_BI_SAMPLES; + bbq_dequant_factors my_col_factors{}; + if (my_col < col_size) { + my_col_factors = get_dequant_factors(dataset_query, col_neighbors[my_col]); + } + + // Cell (row, col) is read (as raw int, via this alias) and written (as the final float) at + // the same address by the same thread + auto* raw_view = reinterpret_cast(s_distances); + constexpr int total_cells = MAX_NUM_BI_SAMPLES * MAX_NUM_BI_SAMPLES; + for (int i = tx; i < total_cells; i += BLOCK_SIZE) { + const int row = i / MAX_NUM_BI_SAMPLES; + const int col = i % MAX_NUM_BI_SAMPLES; + if (row >= new_size || col >= col_size) continue; + const int distance0 = row * MMA_STORE_STRIDE + col; + const Index_t doc_id = new_neighbors[row]; + const Index_t query_id = col_neighbors[col]; + const uint32_t raw = static_cast(raw_view[distance0]); + s_distances[distance0] = bbq_calculate_metric(raw, + s_document_factors[row], + my_col_factors, + dataset_document, + dataset_query, + metric, + dist_epilogue, + doc_id, + query_id); + } + __syncthreads(); + }; + + // Rows are always new_neighbors/document in both phases below, so this runs exactly once -- see + // the identical staging and rationale in local_join_kernel_bbq_simt. + for (int i = tx; i < new_size; i += BLOCK_SIZE) { + s_document_factors[i] = get_dequant_factors(dataset_document, new_neighbors[i]); + } + __syncthreads(); + + // ---- Phase 1: new x new ---- + run_phase(new_neighbors, new_size, std::integral_constant{}, false); + + for (int step = 0; step < raft::ceildiv(new_size, num_warps); ++step) { + const int idx_in_list = step * num_warps + tx / raft::warp_size(); + if (idx_in_list >= new_size) continue; + auto min_elem = get_min_item(s_list[idx_in_list], + idx_in_list, + new_neighbors, + s_distances, + true, + MMA_STORE_STRIDE, + new_size); + if (min_elem.id() < gridDim.x) { + insert_to_global_graph(min_elem, s_list[idx_in_list], graph, dists, graph_width, locks); + } + } + + if (!old_size) return; + __syncthreads(); + + // ---- Phase 2: new x old ---- + run_phase(old_neighbors, old_size, std::false_type{}, n_tiles == 1); + + for (int step = 0; step < raft::ceildiv(MAX_NUM_BI_SAMPLES, num_warps); ++step) { + const int idx_in_list = step * num_warps + tx / raft::warp_size(); + if (idx_in_list >= new_size) continue; + auto min_elem = get_min_item(s_list[idx_in_list], + idx_in_list, + old_neighbors, + s_distances, + true, + MMA_STORE_STRIDE, + old_size); + if (min_elem.id() < gridDim.x) { + insert_to_global_graph(min_elem, s_list[idx_in_list], graph, dists, graph_width, locks); + } + } + + for (int step = 0; step < raft::ceildiv(MAX_NUM_BI_SAMPLES, num_warps); ++step) { + const int idx_in_list = step * num_warps + tx / raft::warp_size(); + if (idx_in_list >= old_size) continue; + const int list_idx = idx_in_list + MAX_NUM_BI_SAMPLES; + auto min_elem = get_min_item( + s_list[list_idx], idx_in_list, new_neighbors, s_distances, false, MMA_STORE_STRIDE, new_size); + if (min_elem.id() < gridDim.x) { + insert_to_global_graph(min_elem, s_list[list_idx], graph, dists, graph_width, locks); + } + } +#endif // (__CUDA_ARCH__ >= 750) +} + +// launch_bounds here denote BLOCK_SIZE = 512 and MIN_BLOCKS_PER_SM = 4 +// Per +// https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#features-and-technical-specifications, +// MAX_RESIDENT_THREAD_PER_SM = BLOCK_SIZE * BLOCKS_PER_SM = 2048 +// For architectures 750 and 860 (890), the values for MAX_RESIDENT_THREAD_PER_SM +// is 1024 and 1536 respectively, which means the bounds don't work anymore +// Used for fp32 data downcast to fp16, and all types using non-L1 distance metric. namespace { template @@ -1147,6 +2055,75 @@ int insert_to_ordered_list(InternalID_t* list, return idx_insert; }; +// Copies `h_graph` (stride `in_degree`) into `out` (stride `out_degree`), dropping self-edges and +// repeated ids, then pads each short list with distinct random ids. +// +// The device-side lists can hold the same id twice: insert_to_global_graph only rejects an +// incoming element when it compares equal to its immediate sorted neighbours, which misses a +// second copy that arrived with a different distance. +template +void shrink_graph_removing_duplicates(Index_t* out, + const InternalID_t* h_graph, + const size_t nrow, + const size_t in_degree, + const size_t out_degree) +{ + // Copy the output graph while removing duplicates. Each thread keeps a bit-packed "seen" + // array, one bit per dataset row, to test and mark ids in O(1) as it scans a row's + // candidates. Only the ids actually placed for a row are ever set, and they're cleared again + // immediately after that row is done, so beyond the one-time zero-initialization when each + // thread starts, no reset across the full array is ever needed. + const size_t num_dedup_words = (nrow + 63) / 64; +#pragma omp parallel + { + std::vector seen_bits(num_dedup_words, 0); + + auto test_and_set = [&](size_t idx) -> bool { + uint64_t mask = uint64_t{1} << (idx & 63); + if (seen_bits[idx >> 6] & mask) { return false; } + seen_bits[idx >> 6] |= mask; + return true; + }; + +#pragma omp for + for (size_t i = 0; i < nrow; i++) { + auto* output_neighbor_list_ptr = out + i * out_degree; + + size_t out_j = 0; + // Copy neighbor list while removing duplicates. + for (size_t in_j = 0; in_j < out_degree; in_j++) { + size_t idx = h_graph[i * in_degree + in_j].id(); + // Unfilled slots carry a sentinel id; leave them to the random fill below. + if (idx >= nrow || idx == i || !test_and_set(idx)) { continue; } + output_neighbor_list_ptr[out_j] = static_cast(idx); + out_j++; + } + + // Fill with random nodes if the length of the filled neighbor list is less than the degree. + for (size_t j = out_j; j < out_degree; j++) { + uint64_t rnd = static_cast(i * out_degree + j + 1); + uint64_t idx = 0; + bool placed = false; + for (size_t attempts = 0; !placed && attempts < out_degree; attempts++) { + rnd = cuvs::neighbors::detail::device::xorshift64(rnd); + idx = rnd % nrow; + placed = (idx != i) && test_and_set(idx); + } + output_neighbor_list_ptr[j] = static_cast(idx); + } + + // Unset every bit this row touched so the array is back to all-zero for the next row this + // thread processes. Since seen_bits is all-zero on entry to this row and thread-local, the + // only bits set anywhere are ones this row's own entries set, so the whole word covering + // idx can be zeroed outright rather than masking off just its one bit. + for (size_t k = 0; k < out_degree; k++) { + size_t idx = static_cast(output_neighbor_list_ptr[k]); + seen_bits[idx >> 6] = 0; + } + } + } +} + } // namespace template @@ -1496,6 +2473,145 @@ void GNND::local_join(cudaStream_t stream, DistEpilogue_t dist_ } } +template +template +void GNND::local_join( + cudaStream_t stream, + cuvs::neighbors::device_bbq_dataset_view, int64_t> dataset, + DistEpilogue_t dist_epilogue) +{ + raft::matrix::fill(res, dists_buffer_.view(), std::numeric_limits::max()); + + // Both kernels take the same (document, query) pair, so there is no symmetric/asymmetric split + // here: a single quantizer just means the same one on both operands, which is exactly what + // SelfJoin encodes. Picking the two quantizers is all that differs. + const bool self_join = dataset.quantizers.size() == 1; + const bool has_1b = dataset.has_layout(bbq_code_layout::packed_1b); + const bool has_4b = dataset.has_layout(bbq_code_layout::packed_4b); + const bool has_2bt = dataset.has_layout(bbq_code_layout::transposed_2b); + const bool has_4bt = dataset.has_layout(bbq_code_layout::transposed_4b); + + // Asymmetric: a packed_4b query selects the tensor-core path, a transposed query the SIMT one. + // Only packed_1b promotes to the tensor-core path; transposed_2b is SIMT-only (it would need + // promoting from two bitplanes at once, which the int4 MMA staging doesn't support). + const bool tc_pair = has_4b && has_1b; + const bool simt_pair = (has_4bt && (has_1b || has_2bt)) || (has_2bt && has_1b); + RAFT_EXPECTS(self_join || tc_pair || simt_pair, + "Unsupported BBQ layout pair for asymmetric local join. Supported: " + "packed_1b x packed_4b (tensor core); packed_1b x transposed_2b, " + "packed_1b x transposed_4b, transposed_2b x transposed_4b (SIMT)."); + auto quantizer_query = + self_join ? dataset.quantizers[0] + : (tc_pair ? dataset.get_quantizer(bbq_code_layout::packed_4b) + : (has_4bt ? dataset.get_quantizer(bbq_code_layout::transposed_4b) + : dataset.get_quantizer(bbq_code_layout::transposed_2b))); + auto quantizer_document = self_join ? dataset.quantizers[0] + : has_1b ? dataset.get_quantizer(bbq_code_layout::packed_1b) + : dataset.get_quantizer(bbq_code_layout::transposed_2b); + + // stage_tile_simt / stage_promoted_tile cast code buffers to uint32_t*, so every plane stride + // must be 4-byte aligned. + { + const auto len = cuvs::preprocessing::quantize::bbq::get_encoded_row_length(quantizer_query); + const int n_planes = + cuvs::preprocessing::quantize::bbq::get_code_planes(quantizer_query.layout); + RAFT_EXPECTS(len % (4u * static_cast(n_planes)) == 0, + "BBQ local join requires the encoded row length to be a multiple of 4*n_planes " + "for 32-bit aligned plane loads, got %u with n_planes = %d", + len, + n_planes); + RAFT_EXPECTS(quantizer_document.dim() % 32 == 0, + "BBQ local join requires dataset dim to be a multiple of 32, got %lld", + static_cast(quantizer_document.dim())); + } + + // One launch site for both kernels: they take identical arguments, and the query's layout picks + // the path -- packed_4b is the only layout the int4 tensor-core kernel is dispatched for. + auto launch = [&](auto document_layout, auto query_layout, auto self_join_tag) { + constexpr auto D = decltype(document_layout)::value; + constexpr auto Q = decltype(query_layout)::value; + constexpr bool S = decltype(self_join_tag)::value; + auto launch_local_join = [&](auto kernel) { + kernel<<>>(graph_.h_graph_new.data_handle(), + h_rev_graph_new_.data_handle(), + d_list_sizes_new_.data_handle(), + h_graph_old_.data_handle(), + h_rev_graph_old_.data_handle(), + d_list_sizes_old_.data_handle(), + NUM_SAMPLES, + quantizer_document, + quantizer_query, + graph_buffer_.data_handle(), + dists_buffer_.data_handle(), + DEGREE_ON_DEVICE, + d_locks_.data_handle(), + build_config_.metric, + dist_epilogue); + }; + // Naming a kernel rather than launching it directly leaves the trailing template parameters + // nothing to deduce from, so they are spelled out here. + using kernel_data_t = std::remove_const_t; + using kernel_id_t = InternalID_t; + if constexpr (Q == bbq_code_layout::packed_4b) { + launch_local_join( + local_join_kernel_bbq_wmma); + } else { + launch_local_join( + local_join_kernel_bbq_simt); + } + }; + const bbq_code_layout d = quantizer_document.layout; + const bbq_code_layout q = quantizer_query.layout; + if (self_join) { + switch (d) { + case bbq_code_layout::packed_1b: + launch(std::integral_constant{}, + std::integral_constant{}, + std::true_type{}); + break; + case bbq_code_layout::transposed_2b: + launch(std::integral_constant{}, + std::integral_constant{}, + std::true_type{}); + break; + case bbq_code_layout::packed_4b: + launch(std::integral_constant{}, + std::integral_constant{}, + std::true_type{}); + break; + case bbq_code_layout::packed_7b: + launch(std::integral_constant{}, + std::integral_constant{}, + std::true_type{}); + break; + case bbq_code_layout::packed_8b: + launch(std::integral_constant{}, + std::integral_constant{}, + std::true_type{}); + break; + default: RAFT_FAIL("Unsupported BBQ layout for symmetric local join on this branch."); + } + } else if (d == bbq_code_layout::packed_1b && q == bbq_code_layout::packed_4b) { + launch(std::integral_constant{}, + std::integral_constant{}, + std::false_type{}); + } else if (d == bbq_code_layout::packed_1b && q == bbq_code_layout::transposed_2b) { + launch(std::integral_constant{}, + std::integral_constant{}, + std::false_type{}); + } else if (d == bbq_code_layout::packed_1b && q == bbq_code_layout::transposed_4b) { + launch(std::integral_constant{}, + std::integral_constant{}, + std::false_type{}); + } else if (d == bbq_code_layout::transposed_2b && q == bbq_code_layout::transposed_4b) { + launch(std::integral_constant{}, + std::integral_constant{}, + std::false_type{}); + } else { + RAFT_FAIL("Unsupported BBQ layout pair for asymmetric local join."); + } +} + template template void GNND::build(Data_t* data, @@ -1727,71 +2843,216 @@ void GNND::build(Data_t* data, Index_t* graph_shrink_buffer = (Index_t*)graph_.h_dists.data_handle(); - // Copy the output graph while removing duplicates. Each thread keeps a bit-packed "seen" - // array, one bit per dataset row, to test and mark ids in O(1) as it scans a row's - // candidates. Only the ids actually placed for a row are ever set, and they're cleared again - // immediately after that row is done, so beyond the one-time zero-initialization when each - // thread starts, no reset across the full array is ever needed. - const size_t num_dedup_words = (static_cast(nrow_) + 63) / 64; -#pragma omp parallel - { - std::vector seen_bits(num_dedup_words, 0); + shrink_graph_removing_duplicates( + graph_shrink_buffer, graph_.h_graph, nrow_, graph_.node_degree, build_config_.node_degree); + graph_.h_graph = nullptr; - auto test_and_set = [&](size_t idx) -> bool { - uint64_t mask = uint64_t{1} << (idx & 63); - if (seen_bits[idx >> 6] & mask) { return false; } - seen_bits[idx >> 6] |= mask; - return true; - }; +#pragma omp parallel for + for (size_t i = 0; i < (size_t)nrow_; i++) { + for (size_t j = 0; j < build_config_.node_degree; j++) { + output_graph[i * build_config_.node_degree + j] = + graph_shrink_buffer[i * build_config_.node_degree + j]; + } + } +} -#pragma omp for - for (size_t i = 0; i < (size_t)nrow_; i++) { - auto output_neighbor_list_ptr = graph_shrink_buffer + i * build_config_.node_degree; +template +template +void GNND::build( + cuvs::neighbors::device_bbq_dataset_view, int64_t> dataset, + Index_t* output_graph, + bool return_distances, + DistData_t* output_distances, + DistEpilogue_t dist_epilogue) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); + nrow_ = static_cast(dataset.n_rows()); + graph_.nrow = nrow_; + graph_.bloom_filter.set_nrow(nrow_); + update_counter_ = 0; + graph_.h_graph = reinterpret_cast*>(output_graph); - size_t out_j = 0; + graph_.clear(); + graph_.init_random_graph(); + graph_.sample_graph(true); - // Copy neighbor list while removing duplicates. - for (size_t in_j = 0; in_j < build_config_.node_degree; in_j++) { - size_t idx = graph_.h_graph[i * graph_.node_degree + in_j].id(); - if (idx >= (size_t)nrow_ || idx == i || !test_and_set(idx)) continue; - output_neighbor_list_ptr[out_j] = idx; - out_j++; + auto update_and_sample = [&](bool update_graph) { + if (update_graph) { + update_counter_ = 0; + graph_.update_graph(graph_host_buffer_.data_handle(), + dists_host_buffer_.data_handle(), + DEGREE_ON_DEVICE, + update_counter_); + if (update_counter_ < build_config_.termination_threshold * nrow_ * + build_config_.dataset_dim / counter_interval) { + update_counter_ = -1; } + } + graph_.sample_graph(false); + }; - // Fill with random nodes if the length of the filled neighbor list is less than the degree. - for (size_t j = out_j; j < build_config_.node_degree; j++) { - uint64_t rnd = static_cast(i * build_config_.node_degree + j + 1); - uint64_t idx = 0; - bool placed = false; - for (size_t attempts = 0; !placed && attempts < build_config_.node_degree; attempts++) { - rnd = cuvs::neighbors::detail::device::xorshift64(rnd); - idx = rnd % nrow_; - placed = (idx != i) && test_and_set(idx); - } - output_neighbor_list_ptr[j] = static_cast(idx); - } + for (size_t it = 0; it < build_config_.max_iterations; ++it) { + raft::copy(res, d_list_sizes_new_.view(), graph_.h_list_sizes_new.view()); + raft::copy(res, h_graph_old_.view(), graph_.h_graph_old.view()); + raft::copy(res, d_list_sizes_old_.view(), graph_.h_list_sizes_old.view()); + raft::resource::sync_stream(res); - // Unset every bit this row touched so the array is back to all-zero for the next row this - // thread processes. Since seen_bits is all-zero on entry to this row and thread-local, the - // only bits set anywhere are ones this row's own entries set, so the whole word covering - // idx can be zeroed outright rather than masking off just its one bit. - for (size_t k = 0; k < build_config_.node_degree; k++) { - size_t idx = static_cast(output_neighbor_list_ptr[k]); - seen_bits[idx >> 6] = 0; + std::thread update_and_sample_thread(update_and_sample, it); + RAFT_LOG_DEBUG("# GNND iteration: %lu / %lu", it + 1, build_config_.max_iterations); + + static_assert(DEGREE_ON_DEVICE * sizeof(*(dists_buffer_.data_handle())) >= + NUM_SAMPLES * sizeof(*(graph_buffer_.data_handle()))); + add_reverse_edges(graph_.h_graph_new.data_handle(), + h_rev_graph_new_.data_handle(), + reinterpret_cast(dists_buffer_.data_handle()), + d_list_sizes_new_.data_handle(), + stream); + add_reverse_edges(h_graph_old_.data_handle(), + h_rev_graph_old_.data_handle(), + reinterpret_cast(dists_buffer_.data_handle()), + d_list_sizes_old_.data_handle(), + stream); + + local_join(stream, dataset, dist_epilogue); + update_and_sample_thread.join(); + if (update_counter_ == -1) { break; } + raft::copy(res, graph_host_buffer_.view(), graph_buffer_.view()); + raft::copy(res, dists_host_buffer_.view(), dists_buffer_.view()); + raft::resource::sync_stream(res); + graph_.sample_graph_new(graph_host_buffer_.data_handle(), DEGREE_ON_DEVICE); + } + + graph_.update_graph(graph_host_buffer_.data_handle(), + dists_host_buffer_.data_handle(), + DEGREE_ON_DEVICE, + update_counter_); + raft::resource::sync_stream(res); + graph_.sort_lists(); + + static_assert(sizeof(decltype(*(graph_.h_dists.data_handle()))) >= sizeof(Index_t)); + if (return_distances) { + auto graph_h_dists = raft::make_host_matrix( + nrow_, build_config_.output_graph_degree); +#pragma omp parallel for + for (size_t i = 0; i < nrow_; ++i) { + for (size_t j = 0; j < build_config_.output_graph_degree; ++j) { + graph_h_dists(i, j) = graph_.h_dists(i, j); } } + raft::copy( + res, + raft::make_device_vector_view(output_distances, nrow_ * build_config_.output_graph_degree), + raft::make_host_vector_view(graph_h_dists.data_handle(), + nrow_ * build_config_.output_graph_degree)); + + auto output_dist_view = raft::make_device_matrix_view( + output_distances, nrow_, build_config_.output_graph_degree); + const bool can_postprocess_dist = std::is_same_v; + if (build_config_.metric == cuvs::distance::DistanceType::L2SqrtExpanded && + can_postprocess_dist) { + raft::linalg::map( + res, output_dist_view, raft::sqrt_op{}, raft::make_const_mdspan(output_dist_view)); + } else if (!cuvs::distance::is_min_close(build_config_.metric) && can_postprocess_dist) { + raft::linalg::map(res, + output_dist_view, + raft::mul_const_op(-1), + raft::make_const_mdspan(output_dist_view)); + } + raft::resource::sync_stream(res); } + + auto* graph_shrink_buffer = reinterpret_cast(graph_.h_dists.data_handle()); + + shrink_graph_removing_duplicates( + graph_shrink_buffer, graph_.h_graph, nrow_, graph_.node_degree, build_config_.node_degree); graph_.h_graph = nullptr; #pragma omp parallel for - for (size_t i = 0; i < (size_t)nrow_; i++) { - for (size_t j = 0; j < build_config_.node_degree; j++) { + for (size_t i = 0; i < nrow_; ++i) { + for (size_t j = 0; j < build_config_.node_degree; ++j) { output_graph[i * build_config_.node_degree + j] = graph_shrink_buffer[i * build_config_.node_degree + j]; } } } +template +void build(raft::resources const& res, + const index_params& params, + cuvs::neighbors::device_bbq_dataset_view dataset, + index& idx) +{ + RAFT_EXPECTS(dataset.quantizers.size() > 0, "BBQ dataset must not be empty."); + auto front_quantizer = dataset.quantizers[0]; + cuvs::common::nvtx::range fun_scope( + "neighbors::nn_descent::detail::build-bbq(%zu, %zu, %zu, %zu, %zu)", + size_t(dataset.n_rows()), + size_t(dataset.dim()), + size_t(idx.graph().extent(1)), + size_t(idx.metric()), + size_t(cuvs::preprocessing::quantize::bbq::get_bit_width(front_quantizer.layout))); + RAFT_EXPECTS(idx.metric() == cuvs::distance::DistanceType::L2Expanded || + idx.metric() == cuvs::distance::DistanceType::L2SqrtExpanded || + idx.metric() == cuvs::distance::DistanceType::CosineExpanded || + idx.metric() == cuvs::distance::DistanceType::InnerProduct, + "BBQ NN-Descent supports L2Expanded, L2SqrtExpanded, CosineExpanded, and " + "InnerProduct."); + RAFT_EXPECTS(idx.metric() == front_quantizer.metric, + "BBQ dataset metric does not match the NN-Descent metric."); + // packed_4b is the only layout dispatched to local_join_kernel_bbq_wmma, and the int4 MMA that + // kernel is built around only exists from sm_75 on. + if (dataset.has_layout(bbq_code_layout::packed_4b)) { + auto kernel = local_join_kernel_bbq_wmma, + raft::identity_op>; + auto runtime_arch = raft::util::arch::kernel_virtual_arch(reinterpret_cast(kernel)); + RAFT_EXPECTS( + raft::util::arch::SM_range(raft::util::arch::SM_75(), raft::util::arch::SM_future()) + .contains(runtime_arch), + "The BBQ packed_4b layout requires int4 tensor cores (compute capability 7.5 or newer), but " + "the local join kernel resolves to %d.%d here. Use transposed_4b for 4-bit codes instead.", + runtime_arch.value() / 100, + (runtime_arch.value() / 10) % 10); + } + + size_t extended_graph_degree; + size_t graph_degree; + auto build_config = get_build_config(res, + params, + dataset.n_rows(), + dataset.dim(), + idx.metric(), + extended_graph_degree, + graph_degree); + auto int_graph = + raft::make_host_matrix(dataset.n_rows(), extended_graph_degree); + GNND nnd(res, build_config); + + if (idx.distances().has_value() || !params.return_distances) { + nnd.build(dataset, + int_graph.data_handle(), + params.return_distances, + idx.distances() + .value_or(raft::make_device_matrix(res, 0, 0).view()) + .data_handle()); + } else { + RAFT_FAIL( + "Distance view not allocated. Using return_distances set to true requires " + "distance view to be allocated."); + } + +#pragma omp parallel for + for (size_t i = 0; i < static_cast(dataset.n_rows()); ++i) { + for (size_t j = 0; j < graph_degree; ++j) { + idx.graph()(i, j) = int_graph(i, j); + } + } +} + template +index build(raft::resources const& res, + const index_params& params, + cuvs::neighbors::device_bbq_dataset_view dataset) +{ + size_t graph_degree = params.graph_degree; + if (params.intermediate_graph_degree < graph_degree) { + RAFT_LOG_WARN( + "Graph degree (%lu) cannot be larger than intermediate graph degree (%lu), reducing " + "graph_degree.", + graph_degree, + params.intermediate_graph_degree); + graph_degree = params.intermediate_graph_degree; + } + + index idx{res, + static_cast(dataset.n_rows()), + static_cast(graph_degree), + params.return_distances, + params.metric}; + detail::build(res, params, dataset, idx); + return idx; +} + template + void build(cuvs::neighbors::device_bbq_dataset_view, int64_t> dataset, + Index_t* output_graph, + bool return_distances, + DistData_t* output_distances, + DistEpilogue_t dist_epilogue = DistEpilogue_t{}); ~GNND() = default; using ID_t = InternalID_t; void reset(raft::resources const& res); @@ -219,6 +226,12 @@ class CUVS_EXPORT GNND { template void local_join(cudaStream_t stream = 0, DistEpilogue_t dist_epilogue = DistEpilogue_t{}); + template + void local_join( + cudaStream_t stream, + cuvs::neighbors::device_bbq_dataset_view, int64_t> dataset, + DistEpilogue_t dist_epilogue = DistEpilogue_t{}); + raft::resources const& res; BuildConfig build_config_; diff --git a/cpp/src/neighbors/nn_descent.cuh b/cpp/src/neighbors/nn_descent.cuh index abe45f0a6f..9c9b3c5540 100644 --- a/cpp/src/neighbors/nn_descent.cuh +++ b/cpp/src/neighbors/nn_descent.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -184,6 +184,23 @@ void build(raft::resources const& res, detail::build(res, params, dataset, idx); } +template +auto build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::device_bbq_dataset_view dataset) -> index +{ + return detail::build(res, params, dataset); +} + +template +void build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::device_bbq_dataset_view dataset, + index& idx) +{ + detail::build(res, params, dataset, idx); +} + /** @} */ // end group nn-descent } // namespace cuvs::neighbors::nn_descent diff --git a/cpp/src/neighbors/nn_descent_gnnd_inst.cu b/cpp/src/neighbors/nn_descent_gnnd_inst.cu index 0e4bb3a534..a5b2517b9c 100644 --- a/cpp/src/neighbors/nn_descent_gnnd_inst.cu +++ b/cpp/src/neighbors/nn_descent_gnnd_inst.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -39,6 +39,12 @@ template void detail::GNND::build( bool return_distances, float* output_distances, raft::identity_op dist_epilogue); +template void detail::GNND::build( + cuvs::neighbors::device_bbq_dataset_view, int64_t> dataset, + int* output_graph, + bool return_distances, + float* output_distances, + raft::identity_op dist_epilogue); template void detail::GNND::local_join( cudaStream_t stream, raft::identity_op dist_epilogue); diff --git a/cpp/src/neighbors/nn_descent_inst.cu.in b/cpp/src/neighbors/nn_descent_inst.cu.in index fbbeb58f7a..62d749e5b6 100644 --- a/cpp/src/neighbors/nn_descent_inst.cu.in +++ b/cpp/src/neighbors/nn_descent_inst.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -52,4 +52,21 @@ auto build(raft::resources const& handle, } }; +auto build(raft::resources const& handle, + const cuvs::neighbors::nn_descent::index_params& params, + cuvs::neighbors::device_bbq_dataset_view dataset, + std::optional> graph) + -> cuvs::neighbors::nn_descent::index +{ + if (!graph.has_value()) { + return cuvs::neighbors::nn_descent::build(handle, params, dataset); + } else { + std::optional> distances = + std::nullopt; + cuvs::neighbors::nn_descent::index idx{ + handle, graph.value(), distances, params.metric}; + cuvs::neighbors::nn_descent::build(handle, params, dataset, idx); + return idx; + }; +} } // namespace cuvs::neighbors::nn_descent diff --git a/cpp/src/preprocessing/quantize/bbq.cu b/cpp/src/preprocessing/quantize/bbq.cu new file mode 100644 index 0000000000..991f47b77e --- /dev/null +++ b/cpp/src/preprocessing/quantize/bbq.cu @@ -0,0 +1,48 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include +#include + +namespace CUVS_EXPORT cuvs { + +namespace preprocessing::quantize::bbq { + +namespace helpers { +void resolve_dequant_factors( + raft::resources const& res, + raft::device_vector_view dequant_delta, + raft::device_vector_view dequant_sum_delta, + raft::device_vector_view lower_intervals, + raft::device_vector_view upper_intervals, + raft::device_vector_view quantized_component_sums, + bbq_code_layout layout) +{ + const auto n_rows = dequant_delta.extent(0); + RAFT_EXPECTS(dequant_sum_delta.extent(0) == n_rows && lower_intervals.extent(0) == n_rows && + upper_intervals.extent(0) == n_rows && + quantized_component_sums.extent(0) == n_rows, + "resolve_dequant_factors: all vectors must have the same length"); + const float scale = 1.0f / static_cast((uint32_t{1} << get_bit_width(layout)) - 1); + auto* sum_delta = dequant_sum_delta.data_handle(); + raft::linalg::map_offset(res, + dequant_delta, + [scale, + sum_delta, + lower = lower_intervals.data_handle(), + upper = upper_intervals.data_handle(), + sums = quantized_component_sums.data_handle()] __device__(int64_t i) { + const float delta = (upper[i] - lower[i]) * scale; + sum_delta[i] = delta * static_cast(sums[i]); + return delta; + }); +} +} // namespace helpers +} // namespace preprocessing::quantize::bbq +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh b/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh new file mode 100644 index 0000000000..4d426d5454 --- /dev/null +++ b/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh @@ -0,0 +1,515 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include + +namespace CUVS_EXPORT cuvs { +namespace preprocessing::quantize::bbq { + +#ifdef __CUDACC__ + +// -------------------------------------------------------------------------- +// Layout geometry +// Encoded row length of a quantizer, and the bit planes its layout slices a row into. +// -------------------------------------------------------------------------- + +template +_RAFT_HOST_DEVICE constexpr uint32_t get_encoded_row_length( + const quantizer_view& dataset) +{ + return get_encoded_row_length(dataset.dim(), dataset.layout); +} + +/** + * Bit planes a row is sliced into: the transposed layouts hold one plane per code bit, the packed + * ones are a single dense plane. This is all the tiling and the cross-plane inner products need + * to know about a layout. + */ +_RAFT_HOST_DEVICE constexpr int get_code_planes(const bbq_code_layout layout) +{ + switch (layout) { + case bbq_code_layout::transposed_2b: return 2; + case bbq_code_layout::transposed_4b: return 4; + default: return 1; + } +} + +// -------------------------------------------------------------------------- +// Code inner products (1x1) +// Raw byte-level dot products over one row pair, per code layout. +// -------------------------------------------------------------------------- + +/** + * Cross-plane binary inner product over `Planes` bit planes, shifting each (i, j) plane pair by + * i + j. `Planes == 1` is a plain binary product, so this covers packed_1b as well as the + * transposed layouts. The 4-byte body needs both rows 4-byte aligned; the byte tail handles a + * stripe whose length is not a multiple of 4. + */ +template +__device__ __forceinline__ uint32_t code_inner_product_transposed(const uint8_t* row_a, + const uint8_t* row_b, + size_t n_bytes, + uint32_t result = 0) +{ + const size_t stripe = n_bytes / Planes; + // Plane p starts at row + p * stripe and is read as uint32_t words, so a stripe that is not a + // multiple of 4 misaligns every plane past the first -- the byte tail below only covers a ragged + // stripe *length*, not a ragged stripe *offset*. Planes == 1 has no offset and so is exempt. + assert(Planes == 1 || stripe % sizeof(uint32_t) == 0); +#pragma unroll + for (int i = 0; i < Planes; ++i) { +#pragma unroll + for (int j = 0; j < Planes; ++j) { + const uint8_t* a = row_a + i * stripe; + const uint8_t* b = row_b + j * stripe; + uint32_t partial = 0; + size_t k = 0; +#pragma unroll 4 + for (; k + sizeof(uint32_t) <= stripe; k += sizeof(uint32_t)) { + partial += __popc(*reinterpret_cast(a + k) & + *reinterpret_cast(b + k)); + } + for (; k < stripe; ++k) { + partial += __popc(static_cast(a[k] & b[k])); + } + result += partial << (i + j); + } + } + return result; +} + +/** One word of two packed_4b rows: the two nibble halves are two masked dp4a products. */ +__device__ __forceinline__ uint32_t dp4a_packed_4b_word(uint32_t a, uint32_t b, uint32_t total) +{ + constexpr uint32_t nibble_mask = 0x0F0F0F0Fu; + total = __dp4a(a & nibble_mask, b & nibble_mask, total); + return __dp4a((a >> 4) & nibble_mask, (b >> 4) & nibble_mask, total); +} + +/** Symmetric for packNibbles (Lucene int4DotProductBothPacked). */ +__device__ __forceinline__ uint32_t code_inner_product_packed_4b(const uint8_t* row_a, + const uint8_t* row_b, + size_t n_bytes, + uint32_t total = 0) +{ + size_t i = 0; +#pragma unroll 4 + for (; i + 4 <= n_bytes; i += 4) { + total = dp4a_packed_4b_word(*reinterpret_cast(row_a + i), + *reinterpret_cast(row_b + i), + total); + } + for (; i < n_bytes; ++i) { + const unsigned a = row_a[i]; + const unsigned b = row_b[i]; + total += (a & 0x0Fu) * (b & 0x0Fu); + total += ((a >> 4) & 0x0Fu) * ((b >> 4) & 0x0Fu); + } + return total; +} + +/** One-byte-per-code dot product, optionally masking unused high bits. */ +__device__ __forceinline__ uint32_t code_inner_product_packed_8b(const uint8_t* row_a, + const uint8_t* row_b, + size_t n_bytes, + uint32_t result = 0, + uint8_t code_mask = 0xFFu) +{ + const uint32_t word_mask = uint32_t{code_mask} * 0x01010101u; + size_t i = 0; +#pragma unroll 4 + for (; i + 4 <= n_bytes; i += 4) { + const auto a = *reinterpret_cast(row_a + i) & word_mask; + const auto b = *reinterpret_cast(row_b + i) & word_mask; + result = __dp4a(a, b, result); + } + for (; i < n_bytes; ++i) { + result += + static_cast(row_a[i] & code_mask) * static_cast(row_b[i] & code_mask); + } + return result; +} + +/** + * Integer inner product between two encoded rows. + * + * The uint32_t result bounds every BBQ layout to 66,050 dimensions: the worst case is + * `packed_8b`, where `dim * 255 * 255` must not exceed UINT32_MAX. + */ +__device__ __forceinline__ uint32_t code_inner_product(const uint8_t* row_a, + const uint8_t* row_b, + const bbq_code_layout layout, + const size_t n_bytes, + uint32_t result = 0) +{ + switch (layout) { + case bbq_code_layout::packed_1b: + return code_inner_product_transposed<1>(row_a, row_b, n_bytes, result); + case bbq_code_layout::transposed_2b: + return code_inner_product_transposed<2>(row_a, row_b, n_bytes, result); + case bbq_code_layout::packed_4b: + return code_inner_product_packed_4b(row_a, row_b, n_bytes, result); + case bbq_code_layout::transposed_4b: + return code_inner_product_transposed<4>(row_a, row_b, n_bytes, result); + case bbq_code_layout::packed_8b: + return code_inner_product_packed_8b(row_a, row_b, n_bytes, result); + case bbq_code_layout::packed_7b: + default: + return code_inner_product_packed_8b( + row_a, + row_b, + n_bytes, + result, + static_cast((uint32_t{1} << get_bit_width(layout)) - 1)); + } +} + +// -------------------------------------------------------------------------- +// Code promotion +// Widens a narrower layout to 4-bit width, so an asymmetric pair can meet in one format. +// -------------------------------------------------------------------------- + +// Promotes one native word of dense packed_1b codes (1 bit/value, 8 values/byte, MSB-first: +// the value at position 8*byte+i sits at bit (7-i)) into four 4-bit-width, packed_4b-style +// output words, packed into one uint4 for a single vectorized store. +// +// Branch-free SWAR: extract the 4 (2-bit) fields of each byte lane-wise across all 4 native +// bytes at once, spread each field 0-3 into a nibble value lane-wise, then transpose the 4 +// resulting field-words into the 4 per-native-byte output words with chained __byte_perm pairs. +__device__ __forceinline__ uint4 packed_1b_to_4b(uint32_t native_word) +{ + const uint32_t spread0 = ((native_word >> 3) & 0x10101010u) | ((native_word >> 6) & 0x01010101u); + const uint32_t spread1 = ((native_word >> 1) & 0x10101010u) | ((native_word >> 4) & 0x01010101u); + const uint32_t spread2 = ((native_word << 1) & 0x10101010u) | ((native_word >> 2) & 0x01010101u); + const uint32_t spread3 = ((native_word << 3) & 0x10101010u) | (native_word & 0x01010101u); + const uint32_t t0 = __byte_perm(spread0, spread1, 0x5140); + const uint32_t t1 = __byte_perm(spread0, spread1, 0x7362); + const uint32_t t2 = __byte_perm(spread2, spread3, 0x5140); + const uint32_t t3 = __byte_perm(spread2, spread3, 0x7362); + return uint4{__byte_perm(t0, t2, 0x5410), + __byte_perm(t0, t2, 0x7632), + __byte_perm(t1, t3, 0x5410), + __byte_perm(t1, t3, 0x7632)}; +} + +// -------------------------------------------------------------------------- +// Cross-layout inner products (1x1) +// One document row against one query row encoded in a different layout. +// -------------------------------------------------------------------------- + +/** + * Cross-plane binary inner product over two bit-sliced rows whose plane counts need not match, + * shifting each (i, j) plane pair by i + j. The runtime-width counterpart of + * code_inner_product_transposed, which the local-join kernels reach with both counts known at + * compile time; a single plane is a plain binary product, so this covers packed_1b too. + * + * All three bit-sliced layouts store ceildiv(dim, 8) bytes per plane, which is what lets one + * @p plane_bytes describe both operands. + */ +__device__ __forceinline__ uint32_t code_inner_product_planes(const uint8_t* row_document, + int document_planes, + const uint8_t* row_query, + int query_planes, + size_t plane_bytes, + uint32_t result = 0) +{ + // Planes past the first start at a multiple of plane_bytes and are read as uint32_t words. + assert((document_planes == 1 && query_planes == 1) || plane_bytes % sizeof(uint32_t) == 0); + for (int i = 0; i < document_planes; ++i) { + for (int j = 0; j < query_planes; ++j) { + const uint8_t* a = row_document + i * plane_bytes; + const uint8_t* b = row_query + j * plane_bytes; + uint32_t partial = 0; + size_t k = 0; +#pragma unroll 4 + for (; k + sizeof(uint32_t) <= plane_bytes; k += sizeof(uint32_t)) { + partial += __popc(*reinterpret_cast(a + k) & + *reinterpret_cast(b + k)); + } + for (; k < plane_bytes; ++k) { + partial += __popc(static_cast(a[k] & b[k])); + } + result += partial << (i + j); + } + } + return result; +} + +/** + * packed_1b document against a packed_4b query. The document is promoted to 4-bit width one + * native word at a time -- 32 dimensions, which is exactly the four query words covering the same + * range -- and the pair is multiplied as two packed_4b rows. The SIMT equivalent of what + * stage_promoted_tile plus the u4 MMA do for this layout pair in the tensor-core local join, and + * it pairs codes the same way, since packed_1b_to_4b emits packed_4b-style words. + * + * Requires dim % 32 == 0, so that the promoted document covers the query row exactly. + */ +__device__ __forceinline__ uint32_t code_inner_product_1b_x_packed_4b(const uint8_t* row_document, + const uint8_t* row_query, + size_t document_bytes, + uint32_t result = 0) +{ + assert(document_bytes % sizeof(uint32_t) == 0); + for (size_t i = 0; i + sizeof(uint32_t) <= document_bytes; i += sizeof(uint32_t)) { + uint32_t promoted[4]; + reinterpret_cast(promoted) = + packed_1b_to_4b(*reinterpret_cast(row_document + i)); + const auto* query_words = reinterpret_cast(row_query + 4 * i); +#pragma unroll + for (int e = 0; e < 4; ++e) { + result = dp4a_packed_4b_word(promoted[e], query_words[e], result); + } + } + return result; +} + +/** + * Integer inner product between a document row and a query row, each read through its own + * quantizer. The two layouts may differ: the supported (document, query) pairs are the ones + * nn-descent's local join dispatches on, and passing the same view twice is the symmetric case. + */ +template +__device__ __forceinline__ uint32_t +code_inner_product(const quantizer_view& quantizer_document, + const quantizer_view& quantizer_query, + int64_t row_document, + int64_t row_query) +{ + const uint8_t* document = &quantizer_document.codes(row_document, 0); + const uint8_t* query = &quantizer_query.codes(row_query, 0); + + if (quantizer_document.layout == quantizer_query.layout) { + return code_inner_product( + document, query, quantizer_document.layout, get_encoded_row_length(quantizer_document)); + } + if (quantizer_document.layout == bbq_code_layout::packed_1b && + quantizer_query.layout == bbq_code_layout::packed_4b) { + return code_inner_product_1b_x_packed_4b( + document, query, get_encoded_row_length(quantizer_document)); + } + // Every remaining supported pair is bit-sliced on both sides: (packed_1b, transposed_2b), + // (packed_1b, transposed_4b) and (transposed_2b, transposed_4b). + assert(quantizer_document.layout == bbq_code_layout::packed_1b || + quantizer_document.layout == bbq_code_layout::transposed_2b); + assert(quantizer_query.layout == bbq_code_layout::transposed_2b || + quantizer_query.layout == bbq_code_layout::transposed_4b); + return code_inner_product_planes(document, + get_code_planes(quantizer_document.layout), + query, + get_code_planes(quantizer_query.layout), + (quantizer_document.dim() + 7) / 8); +} + +// -------------------------------------------------------------------------- +// Fused inner products (2x1) +// Two left rows against a shared right operand, for the local-join inner loop. +// -------------------------------------------------------------------------- + +/** + * Two cross-plane inner products over statically known document and query plane counts. + * Not asymmetric-specific: a symmetric pair is just document_planes == query_planes, which is how + * the 1x1 and 2t x 2t self-joins are computed. + */ +template +__device__ inline void code_inner_product_planes_2x1(const uint8_t* row_a0, + const uint8_t* row_a1, + const uint8_t* row_b, + uint32_t& total0, + uint32_t& total1) +{ + constexpr size_t document_plane_stride = document_row_bytes / document_planes; + constexpr size_t query_plane_stride = query_row_bytes / query_planes; + // Both operands are stepped by their own plane stride and then read as uint32_t, so both + // strides -- not just the query's -- must be 4-byte aligned. + static_assert(query_plane_stride % sizeof(uint32_t) == 0); + static_assert(document_plane_stride % sizeof(uint32_t) == 0); +#pragma unroll + for (int p_query = 0; p_query < query_planes; ++p_query) { +#pragma unroll + for (int p_document = 0; p_document < document_planes; ++p_document) { + const uint8_t* a0 = row_a0 + p_document * document_plane_stride; + const uint8_t* a1 = row_a1 + p_document * document_plane_stride; + const uint8_t* b = row_b + p_query * query_plane_stride; + uint32_t partial0 = 0; + uint32_t partial1 = 0; +#pragma unroll 4 + for (size_t i = 0; i < query_plane_stride; i += sizeof(uint32_t)) { + const auto wa0 = *reinterpret_cast(a0 + i); + const auto wa1 = *reinterpret_cast(a1 + i); + const auto wb = *reinterpret_cast(b + i); + partial0 += __popc(wa0 & wb); + partial1 += __popc(wa1 & wb); + } + total0 += partial0 << (p_document + p_query); + total1 += partial1 << (p_document + p_query); + } + } +} + +template +__device__ inline void code_inner_product_packed_4b_symmetric_2x1(const uint8_t* row_a0, + const uint8_t* row_a1, + const uint8_t* row_b, + uint32_t& total0, + uint32_t& total1) +{ + static_assert(n_bytes % sizeof(uint32_t) == 0); + constexpr uint32_t nibble_mask = 0x0F0F0F0Fu; +#pragma unroll 4 + for (size_t i = 0; i < n_bytes; i += sizeof(uint32_t)) { + const auto a0 = *reinterpret_cast(row_a0 + i); + const auto a1 = *reinterpret_cast(row_a1 + i); + const auto b = *reinterpret_cast(row_b + i); + const auto b_low = b & nibble_mask; + const auto b_high = (b >> 4) & nibble_mask; + total0 = __dp4a(a0 & nibble_mask, b_low, total0); + total0 = __dp4a((a0 >> 4) & nibble_mask, b_high, total0); + total1 = __dp4a(a1 & nibble_mask, b_low, total1); + total1 = __dp4a((a1 >> 4) & nibble_mask, b_high, total1); + } +} + +template +__device__ inline void code_inner_product_packed_8b_2x1(const uint8_t* row_a0, + const uint8_t* row_a1, + const uint8_t* row_b, + uint32_t& total0, + uint32_t& total1, + uint8_t code_mask = 0xFFu) +{ + static_assert(n_bytes % sizeof(uint32_t) == 0); + const uint32_t word_mask = uint32_t{code_mask} * 0x01010101u; +#pragma unroll 4 + for (size_t i = 0; i < n_bytes; i += sizeof(uint32_t)) { + const auto a0 = *reinterpret_cast(row_a0 + i) & word_mask; + const auto a1 = *reinterpret_cast(row_a1 + i) & word_mask; + const auto b = *reinterpret_cast(row_b + i) & word_mask; + total0 = __dp4a(a0, b, total0); + total1 = __dp4a(a1, b, total1); + } +} + +// Selects the SIMT inner product for a (document, query) layout pair: bit-sliced layouts go to the +// cross-plane popc, densely-packed ones to dp4a. dp4a needs both operands in the same packing, so +// those forms apply whenever the two layouts match. packed_7b and packed_8b reach this from +// GNND::local_join; packed_4b does not (symmetric packed_4b goes to the wmma kernel), but is kept +// as a SIMT reference point. +template +__device__ __forceinline__ void bbq_code_inner_product_2x1(const uint8_t* row_a0, + const uint8_t* row_a1, + const uint8_t* row_b, + uint32_t& total0, + uint32_t& total1) +{ + namespace bbq = cuvs::preprocessing::quantize::bbq; + if constexpr (DocumentLayout == QueryLayout && DocumentLayout == bbq_code_layout::packed_4b) { + bbq::code_inner_product_packed_4b_symmetric_2x1( + row_a0, row_a1, row_b, total0, total1); + } else if constexpr (DocumentLayout == QueryLayout && + (DocumentLayout == bbq_code_layout::packed_8b || + DocumentLayout == bbq_code_layout::packed_7b)) { + // packed_7b is packed_8b with the top bit masked off, matching code_inner_product's + // (1 << bits) - 1 mask for the same two layouts. + constexpr uint8_t code_mask = DocumentLayout == bbq_code_layout::packed_7b ? 0x7Fu : 0xFFu; + bbq::code_inner_product_packed_8b_2x1( + row_a0, row_a1, row_b, total0, total1, code_mask); + } else { + bbq:: + code_inner_product_planes_2x1( + row_a0, row_a1, row_b, total0, total1); + } +} + +// -------------------------------------------------------------------------- +// Dequantization +// Turns a raw code inner product into a final float distance. +// -------------------------------------------------------------------------- + +// Per-row dequantization terms needed by bbq_calculate_metric. +// row_norm is read directly from the quantizer view instead, +// since register pressure in the BBQ local-join kernels is already tight. +struct bbq_dequant_factors { + float lower; + float delta; + float sum_delta; + float corrections; +}; + +template +__device__ __forceinline__ bbq_dequant_factors +get_dequant_factors(const quantizer_view& quantizer, int64_t row) +{ + return bbq_dequant_factors{quantizer.lower_intervals(row), + quantizer.dequant_delta(row), + quantizer.dequant_sum_delta(row), + quantizer.additional_corrections(row)}; +} + +// Converts one raw BBQ dot product into a final (post-epilogue) float distance, given both +// operands' precomputed dequant factors. Evaluated exactly once per matrix cell +// dim/centroid_norm_sq/row_norm both come directly from quantizer_document/quantizer_query rather +// than being passed separately, since every caller reads them the same id-indexed way -- row_norm +// is only read for CosineExpanded (skipped entirely otherwise). +template +__device__ __forceinline__ float bbq_calculate_metric( + uint32_t raw, + const bbq_dequant_factors& doc_factors, + const bbq_dequant_factors& query_factors, + const quantizer_view& quantizer_document, + const quantizer_view& quantizer_query, + cuvs::distance::DistanceType metric, + DistEpilogue_t dist_epilogue, + Index_t document_id, + Index_t query_id) +{ + constexpr bool can_postprocess_dist = std::is_same_v; + const float dim = static_cast(quantizer_document.dim()); + + const float centered = dim * doc_factors.lower * query_factors.lower + + query_factors.lower * doc_factors.sum_delta + + doc_factors.lower * query_factors.sum_delta + + doc_factors.delta * query_factors.delta * static_cast(raw); + const float corrections = doc_factors.corrections + query_factors.corrections; + float d; + if (metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded) { + const float raw_distance = corrections - 2.0f * centered; + d = raw_distance < 0.0f ? 0.0f : raw_distance; + if (!can_postprocess_dist && metric == cuvs::distance::DistanceType::L2SqrtExpanded) { + d = sqrtf(d); + } + } else if (metric == cuvs::distance::DistanceType::InnerProduct) { + d = -(centered + corrections - quantizer_document.centroid_norm_sq); + } else { // CosineExpanded + const float norm_product = + quantizer_document.row_norm(document_id) * quantizer_query.row_norm(query_id); + const float dot = centered + corrections - quantizer_document.centroid_norm_sq; + d = norm_product > 0.0f ? 1.0f - dot / sqrtf(norm_product) : 0.0f; + } + return dist_epilogue(d, document_id, query_id); +} + +#endif // __CUDACC__ + +} // namespace preprocessing::quantize::bbq +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index c6a434e4a9..2eed6a3853 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -41,6 +41,7 @@ function(ConfigureTest) ${TEST_NAME} PRIVATE cuvs cuvs::cuvs + cuvs_internal raft::raft rtcx::rtcx GTest::gtest @@ -274,11 +275,19 @@ ConfigureTest( PERCENT 100 ) +ConfigureTest( + NAME NEIGHBORS_ANN_CAGRA_BBQ_UINT32_TEST + PATH neighbors/ann_cagra/test_bbq_uint32_t.cu + GPUS 1 + PERCENT 100 +) + ConfigureTest( NAME NEIGHBORS_ANN_NN_DESCENT_TEST PATH neighbors/ann_nn_descent/test_float_uint32_t.cu neighbors/ann_nn_descent/test_int8_t_uint32_t.cu neighbors/ann_nn_descent/test_uint8_t_uint32_t.cu + neighbors/ann_nn_descent/test_bbq_uint32_t.cu GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/neighbors/ann_cagra/test_bbq_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_bbq_uint32_t.cu new file mode 100644 index 0000000000..f3c0b2a1ff --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/test_bbq_uint32_t.cu @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../ann_cagra_bbq.cuh" + +#include + +namespace cuvs::neighbors::cagra { + +TEST_P(AnnCagraBbqTest, AnnCagraBbqSearchRecall) { this->testSearchRecall(); } +TEST_P(AnnCagraBbqTest, AnnCagraBbqGraphShape) { this->testGraphShape(); } +TEST_P(AnnCagraBbqTest, AnnCagraBbqGraphOnlyBuild) { this->testGraphOnlyBuild(); } +TEST_P(AnnCagraBbqTest, AnnCagraBbqSerializeRoundTrip) { this->testSerializeRoundTrip(); } +TEST_P(AnnCagraBbqTest, AnnCagraBbqUnsupportedParams) { this->testUnsupportedParams(); } + +INSTANTIATE_TEST_CASE_P(AnnCagraBbqTest, AnnCagraBbqTest, ::testing::ValuesIn(bbq_inputs)); + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/neighbors/ann_cagra_bbq.cuh b/cpp/tests/neighbors/ann_cagra_bbq.cuh new file mode 100644 index 0000000000..9ecc4bd1d4 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra_bbq.cuh @@ -0,0 +1,379 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include "ann_cagra.cuh" + +#include +#include +#include + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra { + +struct AnnCagraBbqInputs { + int n_queries; + int n_rows; + int dim; + int k; + int graph_degree; + cuvs::distance::DistanceType metric; + /** Query-side layout; the code width follows from it. */ + cuvs::preprocessing::quantize::bbq::bbq_code_layout layout; + /** + * Document-side layout of a second, coarser quantizer. When set, the dataset carries two + * quantizers and the asymmetric NN-descent path is used. + */ + std::optional second_layout; + /** Fraction of the full-precision build's recall that the BBQ build must retain. */ + double min_recall_ratio; +}; + +inline ::std::ostream& operator<<(::std::ostream& os, const AnnCagraBbqInputs& p) +{ + os << "{n_queries=" << p.n_queries << ", n_rows=" << p.n_rows << ", dim=" << p.dim + << ", k=" << p.k << ", graph_degree=" << p.graph_degree << ", metric=" + << cuvs::neighbors::print_metric{static_cast((int)p.metric)} + << ", layout=" << static_cast(p.layout) << ", second_layout=" + << (p.second_layout.has_value() ? static_cast(*p.second_layout) : -1) << "}"; + return os; +} + +class AnnCagraBbqTest : public ::testing::TestWithParam { + public: + AnnCagraBbqTest() + : stream_(raft::resource::get_cuda_stream(handle_)), + ps(::testing::TestWithParam::GetParam()), + database(0, stream_), + search_queries(0, stream_) + { + } + + protected: + /** Quantize the float database on the host and upload the codes. */ + auto quantize_database() -> cuvs::neighbors::device_bbq_dataset + { + std::vector host_data(static_cast(ps.n_rows) * ps.dim); + raft::update_host(host_data.data(), database.data(), host_data.size(), stream_); + raft::resource::sync_stream(handle_); + + return cuvs_internal::bbq::quantize_to_device(handle_, + host_data.data(), + ps.n_rows, + ps.dim, + ps.metric, + ps.layout, + ps.second_layout.value_or(ps.layout)); + } + + [[nodiscard]] auto default_index_params() const -> cagra::index_params + { + cagra::index_params params; + params.metric = ps.metric; + params.graph_degree = ps.graph_degree; + params.intermediate_graph_degree = 2 * ps.graph_degree; + // The BBQ path always uses nn-descent; pin the dense baseline to it as well so the two builds + // differ only in the precision of the distances driving graph construction. + params.graph_build_params = + cagra::graph_build_params::nn_descent_params(params.intermediate_graph_degree, ps.metric); + return params; + } + + /** Brute-force top-k over the current database, as the ground truth for recall. */ + auto naive_neighbours() -> std::vector + { + size_t queries_size = static_cast(ps.n_queries) * ps.k; + std::vector indices_naive(queries_size); + 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.n_queries, + ps.n_rows, + ps.dim, + ps.k, + ps.metric); + raft::update_host(indices_naive.data(), indices_naive_dev.data(), queries_size, stream_); + raft::resource::sync_stream(handle_); + return indices_naive; + } + + /** Top-k neighbors @p index returns for the test queries. */ + template + auto search_neighbors(IndexT const& index, cagra::search_params search_params = {}) + -> std::vector + { + size_t queries_size = static_cast(ps.n_queries) * ps.k; + rmm::device_uvector distances_dev(queries_size, stream_); + rmm::device_uvector indices_dev(queries_size, stream_); + + auto search_queries_view = raft::make_device_matrix_view( + search_queries.data(), ps.n_queries, ps.dim); + auto indices_out_view = + raft::make_device_matrix_view(indices_dev.data(), ps.n_queries, ps.k); + auto dists_out_view = + raft::make_device_matrix_view(distances_dev.data(), ps.n_queries, ps.k); + + cagra::search( + handle_, search_params, index, search_queries_view, indices_out_view, dists_out_view); + + std::vector indices_cagra(queries_size); + raft::update_host(indices_cagra.data(), indices_dev.data(), queries_size, stream_); + raft::resource::sync_stream(handle_); + return indices_cagra; + } + + template + auto search_recall(IndexT const& index, std::vector const& ground_truth) -> double + { + auto [recall, match_count, total_count] = + calc_recall(ground_truth, search_neighbors(index), ps.n_queries, ps.k); + return recall; + } + + /** + * The end-to-end contract: a graph built purely from quantized codes navigates the uncompressed + * vectors nearly as well as one built from the full-precision vectors. Comparing against a dense + * build on the same data and parameters keeps the bar independent of how much recall the search + * configuration itself can reach. + */ + void testSearchRecall() + { + // CAGRA search has no L2SqrtExpanded kernels, so this metric is only exercised by the + // build-side tests below. + if (ps.metric == cuvs::distance::DistanceType::L2SqrtExpanded) { + GTEST_SKIP() << "CAGRA search does not support L2SqrtExpanded"; + } + + auto ground_truth = naive_neighbours(); + auto database_view = + raft::make_device_matrix_view(database.data(), ps.n_rows, ps.dim); + cuvs::neighbors::test::padded_device_matrix_for_cagra device_padded(handle_, + database_view); + + double reference_recall = 0.0; + { + auto dense_index = cagra::build(handle_, default_index_params(), device_padded.view); + dense_index = cagra::update_dataset(handle_, std::move(dense_index), device_padded.view); + reference_recall = search_recall(dense_index, ground_truth); + } + + double bbq_recall = 0.0; + { + // `owning_codes` backs the view held by the built index, so it must outlive it. + auto owning_codes = quantize_database(); + auto graph_index = + cagra::build(handle_, default_index_params(), owning_codes.as_dataset_view()); + ASSERT_EQ(graph_index.graph_size(), static_cast(ps.n_rows)); + + // Rebinding the uncompressed vectors is what makes the BBQ-built graph searchable. + auto index = cagra::update_dataset(handle_, std::move(graph_index), device_padded.view); + bbq_recall = search_recall(index, ground_truth); + } + + const double min_recall = ps.min_recall_ratio * reference_recall; + RAFT_LOG_INFO( + "CAGRA BBQ build (layout=%d, second_layout=%d): recall=%f, dense reference=%f, " + "retained=%.1f%% (required >= %.1f%%)", + static_cast(ps.layout), + ps.second_layout.has_value() ? static_cast(*ps.second_layout) : -1, + bbq_recall, + reference_recall, + 100.0 * bbq_recall / reference_recall, + 100.0 * ps.min_recall_ratio); + + // A broken baseline would make the ratio check meaningless. + ASSERT_GT(reference_recall, 0.8) << "the full-precision CAGRA baseline is unexpectedly poor"; + EXPECT_GE(bbq_recall, min_recall) + << "recall " << bbq_recall << " retains only " << 100.0 * bbq_recall / reference_recall + << "% of the full-precision build's " << reference_recall; + } + + /** The optimized graph has the requested shape and refers only to existing rows. */ + void testGraphShape() + { + auto owning_codes = quantize_database(); + auto index = cagra::build(handle_, default_index_params(), owning_codes.as_dataset_view()); + + ASSERT_EQ(index.graph_size(), static_cast(ps.n_rows)); + ASSERT_EQ(index.graph_degree(), static_cast(ps.graph_degree)); + // The BBQ dataset view is attached by default, so the index reports the quantized shape. + ASSERT_EQ(index.size(), static_cast(ps.n_rows)); + ASSERT_EQ(index.dim(), static_cast(ps.dim)); + + auto graph_host = raft::make_host_matrix(ps.n_rows, ps.graph_degree); + raft::copy(graph_host.data_handle(), index.graph().data_handle(), graph_host.size(), stream_); + raft::resource::sync_stream(handle_); + + for (int64_t i = 0; i < ps.n_rows; i++) { + for (int64_t j = 0; j < ps.graph_degree; j++) { + ASSERT_LT(graph_host(i, j), static_cast(ps.n_rows)) + << "graph node " << i << " has an out-of-range neighbor at position " << j; + } + } + } + + /** + * A BBQ index file holds the graph alone: the codes and their quantizers live outside the + * index, so the restored graph only becomes searchable once a dataset is reattached, and then + * it answers exactly as the original did. + */ + void testSerializeRoundTrip() + { + if (ps.metric == cuvs::distance::DistanceType::L2SqrtExpanded) { + GTEST_SKIP() << "CAGRA search does not support L2SqrtExpanded"; + } + + auto database_view = + raft::make_device_matrix_view(database.data(), ps.n_rows, ps.dim); + cuvs::neighbors::test::padded_device_matrix_for_cagra device_padded(handle_, + database_view); + + auto owning_codes = quantize_database(); + auto graph_index = + cagra::build(handle_, default_index_params(), owning_codes.as_dataset_view()); + + std::stringstream stored; + cagra::serialize(handle_, stored, graph_index); + + device_bbq_index restored{handle_}; + cagra::deserialize(handle_, stored, &restored); + + ASSERT_EQ(restored.size(), graph_index.size()); + ASSERT_EQ(restored.graph_size(), graph_index.graph_size()); + ASSERT_EQ(restored.graph_degree(), graph_index.graph_degree()); + EXPECT_EQ(restored.metric(), graph_index.metric()); + EXPECT_EQ(restored.dataset().n_rows(), 0); + EXPECT_TRUE(restored.dataset().quantizers.empty()); + + auto original = cagra::update_dataset(handle_, std::move(graph_index), device_padded.view); + auto reattached = cagra::update_dataset(handle_, std::move(restored), device_padded.view); + // Use single-CTA kernel to check for exact search results. + cagra::search_params search_params; + search_params.algo = cagra::search_algo::SINGLE_CTA; + EXPECT_EQ(search_neighbors(reattached, search_params), + search_neighbors(original, search_params)); + } + + /** `attach_dataset_on_build = false` yields a graph without any dataset binding. */ + void testGraphOnlyBuild() + { + auto params = default_index_params(); + params.attach_dataset_on_build = false; + + auto owning_codes = quantize_database(); + auto index = cagra::build(handle_, params, owning_codes.as_dataset_view()); + + ASSERT_EQ(index.graph_size(), static_cast(ps.n_rows)); + ASSERT_EQ(index.graph_degree(), static_cast(ps.graph_degree)); + EXPECT_EQ(index.dataset().n_rows(), 0); + EXPECT_TRUE(index.dataset().quantizers.empty()); + } + + /** Only NN-descent graph construction and the four BBQ metrics are accepted. */ + void testUnsupportedParams() + { + auto owning_codes = quantize_database(); + auto dataset = owning_codes.as_dataset_view(); + + auto ivf_pq_params = default_index_params(); + ivf_pq_params.graph_build_params = cagra::graph_build_params::ivf_pq_params( + raft::matrix_extent(ps.n_rows, ps.dim), ps.metric); + EXPECT_THROW(cagra::build(handle_, ivf_pq_params, dataset), raft::exception); + + auto iterative_params = default_index_params(); + iterative_params.graph_build_params = cagra::graph_build_params::iterative_search_params(); + EXPECT_THROW(cagra::build(handle_, iterative_params, dataset), raft::exception); + + auto l1_params = default_index_params(); + l1_params.metric = cuvs::distance::DistanceType::L1; + EXPECT_THROW(cagra::build(handle_, l1_params, dataset), raft::exception); + } + + void SetUp() override + { + // nn-descent rejects packed_4b below sm_75. + if (ps.layout == cuvs::preprocessing::quantize::bbq::bbq_code_layout::packed_4b && + cuvs::neighbors::device_compute_capability() < 75) { + GTEST_SKIP() << "packed_4b requires int4 tensor cores (compute capability 7.5 or newer)"; + } + database.resize(static_cast(ps.n_rows) * ps.dim, stream_); + search_queries.resize(static_cast(ps.n_queries) * ps.dim, stream_); + raft::random::RngState r(1234ULL); + InitDataset(handle_, database.data(), ps.n_rows, ps.dim, ps.metric, r); + InitDataset(handle_, search_queries.data(), ps.n_queries, ps.dim, ps.metric, r); + raft::resource::sync_stream(handle_); + } + + void TearDown() override + { + raft::resource::sync_stream(handle_); + database.resize(0, stream_); + search_queries.resize(0, stream_); + } + + private: + raft::resources handle_; + rmm::cuda_stream_view stream_; + AnnCagraBbqInputs ps; + rmm::device_uvector database; + rmm::device_uvector search_queries; +}; + +/** + * Search runs on exact distances, so recall mostly reflects how well a graph built from `bits`-wide + * codes navigates. Quality degrades gracefully as the codes get coarser, so the fraction of the + * full-precision build's recall that must be retained is looser for the narrow codes. + */ +inline const std::vector bbq_inputs = [] { + using cuvs::preprocessing::quantize::bbq::bbq_code_layout; + using opt_layout = std::optional; + const std::vector> code_specs{ + {bbq_code_layout::packed_1b, opt_layout{}, 0.85}, + {bbq_code_layout::transposed_2b, opt_layout{}, 0.92}, + {bbq_code_layout::transposed_2b, opt_layout{bbq_code_layout::packed_1b}, 0.92}, + {bbq_code_layout::packed_4b, opt_layout{}, 0.95}, + {bbq_code_layout::packed_4b, opt_layout{bbq_code_layout::packed_1b}, 0.90}, + {bbq_code_layout::transposed_4b, opt_layout{bbq_code_layout::packed_1b}, 0.90}, + {bbq_code_layout::transposed_4b, opt_layout{bbq_code_layout::transposed_2b}, 0.88}, + {bbq_code_layout::packed_7b, opt_layout{}, 0.95}, + {bbq_code_layout::packed_8b, opt_layout{}, 0.95}}; + + std::vector out; + for (const auto& [layout, second_layout, min_recall_ratio] : code_specs) { + const auto batch = raft::util::itertools::product( + {200}, // n_queries + {4000}, // n_rows + {128}, // dim + {10}, // k + {32}, // graph_degree + {cuvs::distance::DistanceType::L2Expanded, + // cuvs::distance::DistanceType::L2SqrtExpanded, + cuvs::distance::DistanceType::InnerProduct, + cuvs::distance::DistanceType::CosineExpanded}, + {layout}, + {second_layout}, + {min_recall_ratio}); + out.insert(out.end(), batch.begin(), batch.end()); + } + return out; +}(); + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/neighbors/ann_nn_descent/test_bbq_uint32_t.cu b/cpp/tests/neighbors/ann_nn_descent/test_bbq_uint32_t.cu new file mode 100644 index 0000000000..bea903ef49 --- /dev/null +++ b/cpp/tests/neighbors/ann_nn_descent/test_bbq_uint32_t.cu @@ -0,0 +1,14 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../ann_nn_descent_bbq.cuh" + +#include + +namespace cuvs::neighbors::nn_descent { +TEST_P(AnnNNDescentBbqTest, AnnNNDescentBbq) { this->testNNDescent(); } + +INSTANTIATE_TEST_CASE_P(AnnNNDescentBbqTest, AnnNNDescentBbqTest, ::testing::ValuesIn(bbq_inputs)); +} // namespace cuvs::neighbors::nn_descent diff --git a/cpp/tests/neighbors/ann_nn_descent_bbq.cuh b/cpp/tests/neighbors/ann_nn_descent_bbq.cuh new file mode 100644 index 0000000000..295cae935f --- /dev/null +++ b/cpp/tests/neighbors/ann_nn_descent_bbq.cuh @@ -0,0 +1,208 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "ann_nn_descent.cuh" + +#include +#include + +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace cuvs::neighbors::nn_descent { +// The host reference quantizer is shared with the ann-bench CAGRA wrapper, so these tests and +// the benchmark can never disagree about the code format. +namespace cpu_bbq = cuvs_internal::bbq; +using cuvs::preprocessing::quantize::bbq::get_bit_width; +using cuvs_internal::bbq::make_device_bbq_dataset; + +struct AnnNNDescentBbqInputs : AnnNNDescentInputs { + cuvs::preprocessing::quantize::bbq::bbq_code_layout layout; + std::optional second_dataset_layout; +}; + +inline ::std::ostream& operator<<(::std::ostream& os, const AnnNNDescentBbqInputs& p) +{ + os << "dataset shape=" << p.n_rows << "x" << p.dim << ", graph_degree=" << p.graph_degree + << ", metric=" + << cuvs::neighbors::print_metric{static_cast((int)p.metric)} + << (p.host_dataset ? ", host" : ", device") << ", layout=" << static_cast(p.layout) + << ", second_dataset_layout=" + << (p.second_dataset_layout.has_value() ? static_cast(p.second_dataset_layout.value()) + : -1) + << std::endl; + return os; +} + +class AnnNNDescentBbqTest : public ::testing::TestWithParam { + public: + AnnNNDescentBbqTest() + : stream_(raft::resource::get_cuda_stream(handle_)), + ps(::testing::TestWithParam::GetParam()), + database(raft::make_device_matrix(handle_, ps.n_rows, ps.dim)) + { + } + + protected: + void testNNDescent() + { + // nn-descent rejects packed_4b below sm_75. + if (ps.layout == cuvs::preprocessing::quantize::bbq::bbq_code_layout::packed_4b && + cuvs::neighbors::device_compute_capability() < 75) { + GTEST_SKIP() << "packed_4b requires int4 tensor cores (compute capability 7.5 or newer)"; + } + if (ps.second_dataset_layout.has_value()) { + // The document must be strictly coarser than the query; the pair's layouts are stated in + // the spec, so validity of the layout combination is the spec's business, not inferred here. + const uint32_t query_bits = get_bit_width(ps.layout); + const uint32_t second_bits = get_bit_width(ps.second_dataset_layout.value()); + if (query_bits > 4 || query_bits == second_bits || query_bits == 1) { + GTEST_SKIP() << "Second dataset is N/A: bits=" << query_bits + << ", layout=" << static_cast(ps.layout) + << " and second bits=" << second_bits; + } + } + size_t queries_size = ps.n_rows * ps.graph_degree; + std::vector indices_NNDescent(queries_size); + std::vector distances_NNDescent(queries_size); + std::vector indices_naive(queries_size); + std::vector distances_naive(queries_size); + + { + rmm::device_uvector distances_naive_dev(queries_size, stream_); + rmm::device_uvector indices_naive_dev(queries_size, stream_); + naive_knn(handle_, + distances_naive_dev.data(), + indices_naive_dev.data(), + database.data_handle(), + database.data_handle(), + ps.n_rows, + ps.n_rows, + ps.dim, + ps.graph_degree, + ps.metric); + raft::update_host(indices_naive.data(), indices_naive_dev.data(), queries_size, stream_); + raft::update_host(distances_naive.data(), distances_naive_dev.data(), queries_size, stream_); + raft::resource::sync_stream(handle_); + } + + { + std::vector host_data(static_cast(ps.n_rows) * ps.dim); + raft::update_host(host_data.data(), database.data_handle(), host_data.size(), stream_); + raft::resource::sync_stream(handle_); + + auto owning_dataset = + cuvs_internal::bbq::quantize_to_device(handle_, + host_data.data(), + ps.n_rows, + ps.dim, + ps.metric, + ps.layout, + ps.second_dataset_layout.value_or(ps.layout)); + auto dataset = owning_dataset.as_dataset_view(); + nn_descent::index_params index_params; + index_params.metric = ps.metric; + index_params.graph_degree = ps.graph_degree; + index_params.intermediate_graph_degree = 2 * ps.graph_degree; + index_params.max_iterations = 100; + index_params.return_distances = true; + + auto index = nn_descent::build(handle_, index_params, dataset); + + raft::copy(indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_); + ASSERT_TRUE(index.distances().has_value()); + raft::copy( + distances_NNDescent.data(), index.distances().value().data_handle(), queries_size, stream_); + raft::resource::sync_stream(handle_); + } + + EXPECT_TRUE(eval_neighbours(indices_naive, + indices_NNDescent, + distances_naive, + distances_NNDescent, + ps.n_rows, + ps.graph_degree, + 0.001, + ps.min_recall)); + } + + void SetUp() override + { + raft::random::RngState r(1234ULL); + raft::random::normal(handle_, r, database.data_handle(), ps.n_rows * ps.dim, 0.1f, 2.0f); + raft::resource::sync_stream(handle_); + } + + void TearDown() override { raft::resource::sync_stream(handle_); } + + private: + raft::resources handle_; + rmm::cuda_stream_view stream_; + AnnNNDescentBbqInputs ps; + raft::device_matrix database; +}; + +// Estimated recall based on bruteforce (InnerProduct): 1: 0.23, 2: 0.52, 4: 0.85, 7: 0.98, 8: 0.99. +const std::vector bbq_inputs = [] { + using cuvs::preprocessing::quantize::bbq::bbq_code_layout; + const std::vector>> + code_specifications{ + // min_recall, query layout, document layout (both widths follow from the layouts) + {0.15, bbq_code_layout::packed_1b, std::optional{}}, + {0.50, bbq_code_layout::transposed_2b, std::optional{}}, + {0.27, + bbq_code_layout::transposed_2b, + std::optional{bbq_code_layout::packed_1b}}, + {0.80, bbq_code_layout::packed_4b, std::optional{}}, + // Asymmetric packed_4b queries take the int4 wmma path (SelfJoin = false). packed_1b is the + // only document layout that promotes to it -- transposed_2b never reaches this kernel. At + // dim=256 this is also the only coverage of the phase-2 staging skip (n_tiles == 1) outside + // a self-join. + {0.35, + bbq_code_layout::packed_4b, + std::optional{bbq_code_layout::packed_1b}}, + // Asymmetric transposed_4b queries (1 + 4t, 2t + 4t) take the SIMT path. + {0.35, + bbq_code_layout::transposed_4b, + std::optional{bbq_code_layout::packed_1b}}, + {0.65, + bbq_code_layout::transposed_4b, + std::optional{bbq_code_layout::transposed_2b}}, + {0.80, bbq_code_layout::packed_7b, std::optional{}}, + {0.80, bbq_code_layout::packed_8b, std::optional{}}}; + std::vector out; + for (const auto& [min_recall, layout, second_layout] : code_specifications) { + const auto batch = raft::util::itertools::product( + {2000}, + {256}, // dim + {64}, // graph_degree + {cuvs::distance::DistanceType::L2Expanded, + cuvs::distance::DistanceType::L2SqrtExpanded, + cuvs::distance::DistanceType::InnerProduct, + cuvs::distance::DistanceType::CosineExpanded}, + {false}, // host_dataset + {min_recall}, + {layout}, + {second_layout}); + out.insert(out.end(), batch.begin(), batch.end()); + } + return out; +}(); + +} // namespace cuvs::neighbors::nn_descent diff --git a/cpp/tests/neighbors/ann_utils.cuh b/cpp/tests/neighbors/ann_utils.cuh index a45628ed62..c1022eba17 100644 --- a/cpp/tests/neighbors/ann_utils.cuh +++ b/cpp/tests/neighbors/ann_utils.cuh @@ -28,6 +28,13 @@ namespace cuvs::neighbors { +/** Compute capability of the current device as major * 10 + minor. */ +inline auto device_compute_capability() -> int +{ + auto [major, minor] = raft::getComputeCapability(); + return major * 10 + minor; +} + struct print_dtype { cudaDataType_t value; }; diff --git a/fern/docs.yml b/fern/docs.yml index c6c334fe54..a2ab50408e 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -314,6 +314,8 @@ navigation: path: "./pages/c_api/c-api-neighbors-tiered-index.md" - page: "Neighbors Vamana" path: "./pages/c_api/c-api-neighbors-vamana.md" + - page: "Preprocessing Quantize Bbq" + path: "./pages/c_api/c-api-preprocessing-quantize-bbq.md" - page: "Preprocessing Quantize Binary" path: "./pages/c_api/c-api-preprocessing-quantize-binary.md" - page: "Preprocessing PCA" @@ -398,6 +400,8 @@ navigation: path: "./pages/cpp_api/cpp-api-neighbors-tiered-index.md" - page: "Neighbors Vamana" path: "./pages/cpp_api/cpp-api-neighbors-vamana.md" + - page: "Preprocessing Quantize Bbq" + path: "./pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md" - page: "Preprocessing Quantize Binary" path: "./pages/cpp_api/cpp-api-preprocessing-quantize-binary.md" - page: "Preprocessing PCA" @@ -472,6 +476,8 @@ navigation: - section: "Java API Documentation" path: "./pages/java_api/index.md" contents: + - page: "BbqQuantizer" + path: "./pages/java_api/java-api-com-nvidia-cuvs-bbqquantizer.md" - page: "CuVSDeviceMatrix" path: "./pages/java_api/java-api-com-nvidia-cuvs-cuvsdevicematrix.md" - page: "CuVSHostMatrix" diff --git a/fern/pages/c_api/c-api-core-dataset.md b/fern/pages/c_api/c-api-core-dataset.md index 915c16cbbc..9d989183d6 100644 --- a/fern/pages/c_api/c-api-core-dataset.md +++ b/fern/pages/c_api/c-api-core-dataset.md @@ -17,7 +17,8 @@ Generic dataset layout kind for C API dataset handles. typedef enum { CUVS_DATASET_LAYOUT_STANDARD = 0, CUVS_DATASET_LAYOUT_PADDED = 1, - CUVS_DATASET_LAYOUT_PQ = 2 + CUVS_DATASET_LAYOUT_PQ = 2, + CUVS_DATASET_LAYOUT_BBQ = 3 } cuvsDatasetLayout_t; ``` @@ -28,6 +29,7 @@ typedef enum { | `CUVS_DATASET_LAYOUT_STANDARD` | `0` | | `CUVS_DATASET_LAYOUT_PADDED` | `1` | | `CUVS_DATASET_LAYOUT_PQ` | `2` | +| `CUVS_DATASET_LAYOUT_BBQ` | `3` | ### cuvsDatasetMemType_t diff --git a/fern/pages/c_api/c-api-neighbors-cagra.md b/fern/pages/c_api/c-api-neighbors-cagra.md index 7ac45afa5a..13c663a7c8 100644 --- a/fern/pages/c_api/c-api-neighbors-cagra.md +++ b/fern/pages/c_api/c-api-neighbors-cagra.md @@ -827,11 +827,11 @@ cuvsDataset_t dataset, cuvsCagraIndex_t index); ``` -The memory space and layout `dataset` was constructed with select the C++ build overload. Build the handle with an owning factory or the matching dataset view factory (`cuvsDatasetMakePaddedView` / `cuvsDatasetMakeStandardView`). +The memory space and layout `dataset` was constructed with select the C++ build overload. Build the handle with an owning factory or the matching dataset view factory (`cuvsDatasetMakePaddedView`, `cuvsDatasetMakeStandardView`, or `cuvsDatasetMakeBbqView`). Note that a dataset residing in host memory produces a host-backed index, which must be made search-ready with `cuvsCagraUpdateDataset` (using a device-padded dataset) before calling `cuvsCagraSearch`. -A `CUVS_DATASET_LAYOUT_PQ` dataset created by `cuvsDatasetMakePQ` builds an iterative CAGRA-Q index. VPQ input requires `L2Expanded` and `ITERATIVE_CAGRA_SEARCH` (or `AUTO_SELECT`), and the VPQ dataset must outlive the index because the index stores a non-owning view. +A `CUVS_DATASET_LAYOUT_PQ` dataset created by `cuvsDatasetMakePQ` builds an iterative CAGRA-Q index. VPQ input requires `L2Expanded` and `ITERATIVE_CAGRA_SEARCH` (or `AUTO_SELECT`), and the VPQ dataset must outlive the index because the index stores a non-owning view. A `CUVS_DATASET_LAYOUT_BBQ` dataset builds a graph-only index; attach a searchable dataset with `cuvsCagraUpdateDataset` before search. **Parameters** @@ -935,7 +935,7 @@ const char* filename, cuvsCagraIndex_t index); ``` -This supports dense and PQ-backed indexes. The dataset must be attached separately after loading the graph. +This supports dense, PQ-backed, and BBQ-built indexes. The serialized file does not contain vector data. After deserialization the index cannot be searched until a compatible dataset is attached with `cuvsCagraUpdateDataset`. Experimental, both the API and the serialization format are subject to change. @@ -962,7 +962,7 @@ const char* filename, cuvsCagraIndex_t index); ``` -The index stores a non-owning dataset view. The caller must keep the dataset backing that view alive while this function runs. Returns CUVS_ERROR without modifying the destination file if the index has no attached dataset. PQ datasets are not serialized by this function. +The index stores a non-owning dataset view. The caller must keep the memory of the dataset backing that view alive while this function runs. Returns CUVS_ERROR without modifying the destination file if the index has no attached dataset. PQ and BBQ datasets are not serialized by this function. Experimental, both the API and the serialization format are subject to change. diff --git a/fern/pages/c_api/c-api-preprocessing-quantize-bbq.md b/fern/pages/c_api/c-api-preprocessing-quantize-bbq.md new file mode 100644 index 0000000000..8f57b980c7 --- /dev/null +++ b/fern/pages/c_api/c-api-preprocessing-quantize-bbq.md @@ -0,0 +1,153 @@ +--- +slug: api-reference/c-api-preprocessing-quantize-bbq +--- + +# Bbq + +_Source header: `cuvs/preprocessing/quantize/bbq.h`_ + +## C API for Better Binary Quantization datasets + + +### cuvsBbqCodeLayout_t + +Storage layout of BBQ/OSQ quantized component codes in each dataset row. + +```c +typedef enum { + CUVS_BBQ_CODE_LAYOUT_PACKED_1B = 0, + CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_2B, + CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_4B, + CUVS_BBQ_CODE_LAYOUT_PACKED_4B, + CUVS_BBQ_CODE_LAYOUT_PACKED_7B, + CUVS_BBQ_CODE_LAYOUT_PACKED_8B +} cuvsBbqCodeLayout_t; +``` + +**Values** + +| Name | Value | Description | +| --- | --- | --- | +| `CUVS_BBQ_CODE_LAYOUT_PACKED_1B` | `0` | Each dimension is quantized to a single bit and packed into bytes. Reflects Lucene's OptimizedScalarQuantizer.packAsBinary. | +| `CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_2B` | `` | Each dimension is quantized to 2 bits, stored as 2 bitplanes. Reflects Lucene's OptimizedScalarQuantizer.transposeDibit. SIMT popc path only (paired with a transposed_4b or packed_1b operand); | +| `CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_4B` | `` | Each dimension is quantized to 4 bits, optimized for bitwise operations. Reflects Lucene's OptimizedScalarQuantizer.transposeHalfByte. the first bit of every dimension is in the first set dimensions bits, or (dimensions/8) bytes. The second, third, and fourth bits are in the second, third, and fourth set of dimensions bits, respectively. Format used for queries. | +| `CUVS_BBQ_CODE_LAYOUT_PACKED_4B` | `` | Each dimension is quantized to 4 bits, two values are packed into each output byte. | +| `CUVS_BBQ_CODE_LAYOUT_PACKED_7B` | `` | Each dimension is quantized to 7 bits and treated as a signed value. | +| `CUVS_BBQ_CODE_LAYOUT_PACKED_8B` | `` | Each dimension is quantized to 8 bits and treated as an unsigned value. | + + +### cuvsBbqQuantizer + +Better Binary Quantization ([BBQ](https://www.elastic.co/search-labs/blog/better-binary-quantization-lucene-elasticsearch)) is a vector-quantization approach used in Elasticsearch and Apache Lucene. It builds on ideas introduced in RaBitQ([Gao and Long](https://arxiv.org/pdf/2405.12497, [Gao et al.](https://arxiv.org/pdf/2409.09913)): residual binary codes around a centroid, corrective factors, and efficient bitwise comparison of codes at different bit widths. Lucene implements this as optimized scalar quantization (OSQ) with packed and bit-plane layouts; Elasticsearch exposes it as BBQ. + +BBQ in cuVS designed to be compatible with the Lucene/Elasticsearch dataset: a single shared centroid, no random rotation, and OSQ codes. + +RaBitQ and BBQ in cuVS both compress centroid-relative vectors to low-bit codes and retain additional per-vector information so search is better than naïve sign-bit comparison. They differ in transformation and scale representation. RaBitQ commonly separates residual magnitude from direction, then applies a random orthogonal rotation before binary coding; BBQ uses per-vector scalar intervals to interpret the compressed residual codes. + +```c +typedef struct cuvsBbqQuantizer { + uintptr_t addr; + DLDataType dtype; + bool is_owning; +} cuvsBbqQuantizer; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `addr` | `uintptr_t` | | +| `dtype` | `DLDataType` | | +| `is_owning` | `bool` | | + + +### cuvsBbqQuantizerCreateView + +Create a BBQ quantizer view from caller-owned device tensors. + +```c +cuvsError_t cuvsBbqQuantizerCreateView( +DLManagedTensor* codes, +DLManagedTensor* lower_intervals, +DLManagedTensor* upper_intervals, +DLManagedTensor* additional_corrections, +DLManagedTensor* quantized_component_sums, +DLManagedTensor* centroid, +DLManagedTensor* dequant_delta, +DLManagedTensor* dequant_sum_delta, +DLManagedTensor* row_norm, +cuvsBbqCodeLayout_t layout, +cuvsDistanceType metric, +float centroid_norm_sq, +cuvsBbqQuantizer_t* quantizer); +``` + +Tensors are not copied and must remain valid while a derived dataset is in use. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `codes` | in | `DLManagedTensor*` | uint8 device matrix containing encoded rows | +| `lower_intervals` | in | `DLManagedTensor*` | float32 device vector with one lower interval per row | +| `upper_intervals` | in | `DLManagedTensor*` | float32 device vector with one upper interval per row | +| `additional_corrections` | in | `DLManagedTensor*` | float32 device vector with one correction per row | +| `quantized_component_sums` | in | `DLManagedTensor*` | int32 device vector with one component sum per row | +| `centroid` | in | `DLManagedTensor*` | device vector containing the dataset centroid | +| `dequant_delta` | in | `DLManagedTensor*` | float32 device vector with one dequantization delta per row | +| `dequant_sum_delta` | in | `DLManagedTensor*` | float32 device vector with one delta-times-sum value per row | +| `row_norm` | in | `DLManagedTensor*` | float32 device vector with one original-space squared norm per row | +| `layout` | in | [`cuvsBbqCodeLayout_t`](/api-reference/c-api-preprocessing-quantize-bbq#cuvsbbqcodelayout-t) | encoded code layout | +| `metric` | in | [`cuvsDistanceType`](/api-reference/c-api-distance-distance#cuvsdistancetype) | distance metric associated with the encoded dataset | +| `centroid_norm_sq` | in | `float` | squared norm of the centroid | +| `quantizer` | out | [`cuvsBbqQuantizer_t*`](/api-reference/c-api-preprocessing-quantize-bbq#cuvsbbqquantizer) | newly allocated non-owning quantizer handle | + +**Returns** + +[`cuvsError_t`](/api-reference/c-api-core-c-api#cuvserror-t) + + +### cuvsBbqQuantizerDestroy + +Destroy a BBQ quantizer without destroying its caller-owned tensors. + +```c +cuvsError_t cuvsBbqQuantizerDestroy(cuvsBbqQuantizer_t quantizer); +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `quantizer` | in | [`cuvsBbqQuantizer_t`](/api-reference/c-api-preprocessing-quantize-bbq#cuvsbbqquantizer) | quantizer handle to destroy | + +**Returns** + +[`cuvsError_t`](/api-reference/c-api-core-c-api#cuvserror-t) + + +### cuvsDatasetMakeBbqView + +Create a non-owning device BBQ dataset view. + +```c +cuvsError_t cuvsDatasetMakeBbqView(cuvsResources_t res, +cuvsBbqQuantizer_t* quantizers, +size_t num_quantizers, +cuvsDataset_t* dataset); +``` + +Accepts one symmetric quantizer or two compatible asymmetric quantizers. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | in | [`cuvsResources_t`](/api-reference/c-api-core-c-api#cuvsresources-t) | cuVS resources | +| `quantizers` | in | [`cuvsBbqQuantizer_t*`](/api-reference/c-api-preprocessing-quantize-bbq#cuvsbbqquantizer) | array containing one or two BBQ quantizer handles | +| `num_quantizers` | in | `size_t` | number of elements in `quantizers` | +| `dataset` | out | `cuvsDataset_t*` | newly allocated non-owning BBQ dataset handle | + +**Returns** + +[`cuvsError_t`](/api-reference/c-api-core-c-api#cuvserror-t) diff --git a/fern/pages/c_api/index.md b/fern/pages/c_api/index.md index 968ef01b0e..cd70102764 100644 --- a/fern/pages/c_api/index.md +++ b/fern/pages/c_api/index.md @@ -41,6 +41,7 @@ These pages are generated from the documented public headers in the cuVS source ## Preprocessing +- [Bbq](/api-reference/c-api-preprocessing-quantize-bbq) - [Binary](/api-reference/c-api-preprocessing-quantize-binary) - [PCA](/api-reference/c-api-preprocessing-pca) - [PQ](/api-reference/c-api-preprocessing-quantize-pq) diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index 9097f4cb47..2560f59ef1 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -1098,6 +1098,104 @@ cuvs::neighbors::host_standard_dataset_view const& dataset) built `host_standard_index<uint8_t, uint32_t>` +**Additional overload:** `neighbors::cagra::build` + +Build from a device BBQ-quantized dataset view. + +```cpp +auto build(raft::resources const& res, +const cuvs::neighbors::cagra::index_params& params, +cuvs::neighbors::device_bbq_dataset_view const& dataset) +-> cuvs::neighbors::cagra::device_bbq_index; +``` + +The kNN graph is built from the quantized codes alone, so the uncompressed vectors are never needed and peak memory is driven by the code size. Only nn-descent graph construction is available (IVF-PQ, iterative CAGRA search, and ACE all read uncompressed vectors), and the metric must be one of L2Expanded, L2SqrtExpanded, CosineExpanded, or InnerProduct and must match the metric the quantizer corrections were generated for. + +The returned index cannot be searched: CAGRA has no BBQ search kernels. Call the type-changing `update_dataset` with an uncompressed device-padded dataset to search the resulting graph. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | in | `raft::resources const&` | raft resources | +| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `dataset` | in | `cuvs::neighbors::device_bbq_dataset_view const&` | device BBQ dataset view [n_rows, dim] | + +**Returns** + +`cuvs::neighbors::cagra::device_bbq_index` + +built `device_bbq_index<float, uint32_t>` + +**Additional overload:** `neighbors::cagra::build` + +cuvs::neighbors::device_bbq_dataset_view<float, int64_t> const& dataset) + +```cpp +auto build(raft::resources const& res, +const cuvs::neighbors::cagra::index_params& params, +cuvs::neighbors::device_bbq_dataset_view const& dataset) +-> cuvs::neighbors::cagra::device_bbq_index; +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | +| `params` | | `const cuvs::neighbors::cagra::index_params&` | | +| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view const&` | | + +**Returns** + +`cuvs::neighbors::cagra::device_bbq_index` + +**Additional overload:** `neighbors::cagra::build` + +cuvs::neighbors::device_bbq_dataset_view<float, int64_t> const& dataset) + +```cpp +auto build(raft::resources const& res, +const cuvs::neighbors::cagra::index_params& params, +cuvs::neighbors::device_bbq_dataset_view const& dataset) +-> cuvs::neighbors::cagra::device_bbq_index; +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | +| `params` | | `const cuvs::neighbors::cagra::index_params&` | | +| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view const&` | | + +**Returns** + +`cuvs::neighbors::cagra::device_bbq_index` + +**Additional overload:** `neighbors::cagra::build` + +cuvs::neighbors::device_bbq_dataset_view<float, int64_t> const& dataset) + +```cpp +auto build(raft::resources const& res, +const cuvs::neighbors::cagra::index_params& params, +cuvs::neighbors::device_bbq_dataset_view const& dataset) +-> cuvs::neighbors::cagra::device_bbq_index; +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | +| `params` | | `const cuvs::neighbors::cagra::index_params&` | | +| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view const&` | | + +**Returns** + +`cuvs::neighbors::cagra::device_bbq_index` + ## CAGRA extend functions diff --git a/fern/pages/cpp_api/cpp-api-neighbors-nn-descent.md b/fern/pages/cpp_api/cpp-api-neighbors-nn-descent.md index abbb9cbe58..396966c202 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-nn-descent.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-nn-descent.md @@ -528,3 +528,111 @@ Usage example: **Returns** [`cuvs::neighbors::nn_descent::index`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index) + +**Additional overload:** `neighbors::nn_descent::build` + +Build an NN-Descent index directly from a device-resident BBQ dataset. + +```cpp +auto build(raft::resources const& res, +index_params const& params, +cuvs::neighbors::device_bbq_dataset_view dataset, +std::optional> graph = +std::nullopt) -> cuvs::neighbors::nn_descent::index; +``` + +Symmetric compressed-code distances are used during graph construction. Supported metrics are L2Expanded, L2SqrtExpanded, CosineExpanded, and InnerProduct. The dataset's correction terms must have been generated for the selected metric. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | raft resources | +| `params` | | [`index_params const&`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index-params) | NN-Descent build parameters | +| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view` | BBQ codes and correction terms in device memory | +| `graph` | | `std::optional>` | optional caller-owned host graph
Default: `std::nullopt`. | + +**Returns** + +[`cuvs::neighbors::nn_descent::index`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index) + +**Additional overload:** `neighbors::nn_descent::build` + +cuvs::neighbors::device_bbq_dataset_view<float, int64_t> dataset, + +```cpp +auto build(raft::resources const& res, +index_params const& params, +cuvs::neighbors::device_bbq_dataset_view dataset, +std::optional> graph = +std::nullopt) -> cuvs::neighbors::nn_descent::index; +``` + +std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph = std::nullopt) -> cuvs::neighbors::nn_descent::index<uint32_t> + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | +| `params` | | [`index_params const&`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index-params) | | +| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view` | | +| `graph` | | `std::optional>` | Default: `std::nullopt`. | + +**Returns** + +[`cuvs::neighbors::nn_descent::index`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index) + +**Additional overload:** `neighbors::nn_descent::build` + +cuvs::neighbors::device_bbq_dataset_view<float, int64_t> dataset, + +```cpp +auto build(raft::resources const& res, +index_params const& params, +cuvs::neighbors::device_bbq_dataset_view dataset, +std::optional> graph = +std::nullopt) -> cuvs::neighbors::nn_descent::index; +``` + +std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph = std::nullopt) -> cuvs::neighbors::nn_descent::index<uint32_t> + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | +| `params` | | [`index_params const&`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index-params) | | +| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view` | | +| `graph` | | `std::optional>` | Default: `std::nullopt`. | + +**Returns** + +[`cuvs::neighbors::nn_descent::index`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index) + +**Additional overload:** `neighbors::nn_descent::build` + +cuvs::neighbors::device_bbq_dataset_view<float, int64_t> dataset, + +```cpp +auto build(raft::resources const& res, +index_params const& params, +cuvs::neighbors::device_bbq_dataset_view dataset, +std::optional> graph = +std::nullopt) -> cuvs::neighbors::nn_descent::index; +``` + +std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph = std::nullopt) -> cuvs::neighbors::nn_descent::index<uint32_t> + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | +| `params` | | [`index_params const&`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index-params) | | +| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view` | | +| `graph` | | `std::optional>` | Default: `std::nullopt`. | + +**Returns** + +[`cuvs::neighbors::nn_descent::index`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index) diff --git a/fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md b/fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md new file mode 100644 index 0000000000..67fa8291c1 --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md @@ -0,0 +1,153 @@ +--- +slug: api-reference/cpp-api-preprocessing-quantize-bbq +--- + +# Bbq + +_Source header: `cuvs/preprocessing/quantize/bbq.hpp`_ + +## Better Binary Quantization utilities + + +### preprocessing::quantize::bbq::bbq_code_layout + +Storage layout of BBQ/OSQ quantized component codes in each dataset row. + +```cpp +enum class bbq_code_layout { + packed_1b, + transposed_2b, + transposed_4b, + packed_4b, + packed_7b, + packed_8b +}; +``` + +**Values** + +| Name | Value | Description | +| --- | --- | --- | +| `packed_1b` | `` | Each dimension is quantized to a single bit and packed into bytes. Reflects Lucene's OptimizedScalarQuantizer.packAsBinary. | +| `transposed_2b` | `` | Each dimension is quantized to 2 bits, stored as 2 bitplanes. Reflects Lucene's OptimizedScalarQuantizer.transposeDibit. SIMT popc path only (paired with a transposed_4b or packed_1b operand); | +| `transposed_4b` | `` | Each dimension is quantized to 4 bits, optimized for bitwise operations. Reflects Lucene's OptimizedScalarQuantizer.transposeHalfByte. the first bit of every dimension is in the first set dimensions bits, or (dimensions/8) bytes. The second, third, and fourth bits are in the second, third, and fourth set of dimensions bits, respectively. Format used for queries. | +| `packed_4b` | `` | Each dimension is quantized to 4 bits, two values are packed into each output byte. | +| `packed_7b` | `` | Each dimension is quantized to 7 bits and treated as a signed value. | +| `packed_8b` | `` | Each dimension is quantized to 8 bits and treated as an unsigned value. | + + +### preprocessing::quantize::bbq::get_bit_width + +Bit width of a layout. + +```cpp +constexpr auto get_bit_width(bbq_code_layout layout) noexcept -> uint32_t; +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `layout` | | [`bbq_code_layout`](/api-reference/cpp-api-preprocessing-quantize-bbq#preprocessing-quantize-bbq-bbq-code-layout) | | + +**Returns** + +`uint32_t` + + +### preprocessing::quantize::bbq::get_encoded_row_length + +Bytes one row of `dim` components occupies once encoded in `layout`. + +```cpp +constexpr auto get_encoded_row_length(uint32_t dim, bbq_code_layout layout) noexcept -> uint32_t; +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `dim` | | `uint32_t` | | +| `layout` | | [`bbq_code_layout`](/api-reference/cpp-api-preprocessing-quantize-bbq#preprocessing-quantize-bbq-bbq-code-layout) | | + +**Returns** + +`uint32_t` + + +### preprocessing::quantize::bbq::quantizer + +Better Binary Quantization ([BBQ](https://www.elastic.co/search-labs/blog/better-binary-quantization-lucene-elasticsearch)) is a vector-quantization approach used in Elasticsearch and Apache Lucene. It builds on ideas introduced in RaBitQ([Gao and Long](https://arxiv.org/pdf/2405.12497, [Gao et al.](https://arxiv.org/pdf/2409.09913)): residual binary codes around a centroid, corrective factors, and efficient bitwise comparison of codes at different bit widths. Lucene implements this as optimized scalar quantization (OSQ) with packed and bit-plane layouts; Elasticsearch exposes it as BBQ. + +BBQ in cuVS designed to be compatible with the Lucene/Elasticsearch dataset: a single shared centroid, no random rotation, and OSQ codes. + +RaBitQ and BBQ in cuVS both compress centroid-relative vectors to low-bit codes and retain additional per-vector information so search is better than naïve sign-bit comparison. They differ in transformation and scale representation. RaBitQ commonly separates residual magnitude from direction, then applies a random orthogonal rotation before binary coding; BBQ uses per-vector scalar intervals to interpret the compressed residual codes. + +```cpp +template +struct quantizer { + raft::device_vector dequant_delta; + raft::device_vector dequant_sum_delta; + raft::device_vector row_norm; + raft::device_matrix codes; + raft::device_vector lower_intervals; + raft::device_vector upper_intervals; + raft::device_vector additional_corrections; + raft::device_vector quantized_component_sums; + raft::device_vector centroid; + bbq_code_layout layout; + cuvs::distance::DistanceType metric; + float centroid_norm_sq; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `dequant_delta` | `raft::device_vector` | Precomputed per-row dequantization factors, derived once (offline) from lower/upper_intervals and quantized_component_sums: dequant_delta = (upper-lower)/(2^bits-1) | +| `dequant_sum_delta` | `raft::device_vector` | Precomputed per-row dequantization factors, derived once (offline) from dequant_delta and quantized_component_sums: dequant_sum_delta = dequant_delta * quantized_component_sums. | +| `row_norm` | `raft::device_vector` | Squared norm of the row in original (un-centered) vector space, \|\|x\|\|^2 | +| `codes` | `raft::device_matrix` | | +| `lower_intervals` | `raft::device_vector` | | +| `upper_intervals` | `raft::device_vector` | | +| `additional_corrections` | `raft::device_vector` | | +| `quantized_component_sums` | `raft::device_vector` | | +| `centroid` | `raft::device_vector` | | +| `layout` | [`bbq_code_layout`](/api-reference/cpp-api-preprocessing-quantize-bbq#preprocessing-quantize-bbq-bbq-code-layout) | | +| `metric` | [`cuvs::distance::DistanceType`](/api-reference/cpp-api-distance-distance#distance-distancetype) | | +| `centroid_norm_sq` | `float` | | + + +### preprocessing::quantize::bbq::helpers::resolve_dequant_factors + +Derives dequant_delta from lower/upper_intervals and the layout's code width, and + +```cpp +void resolve_dequant_factors( +raft::resources const& res, +raft::device_vector_view dequant_delta, +raft::device_vector_view dequant_sum_delta, +raft::device_vector_view lower_intervals, +raft::device_vector_view upper_intervals, +raft::device_vector_view quantized_component_sums, +bbq_code_layout layout); +``` + +dequant_sum_delta from that delta and quantized_component_sums. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | +| `dequant_delta` | | `raft::device_vector_view` | | +| `dequant_sum_delta` | | `raft::device_vector_view` | | +| `lower_intervals` | | `raft::device_vector_view` | | +| `upper_intervals` | | `raft::device_vector_view` | | +| `quantized_component_sums` | | `raft::device_vector_view` | | +| `layout` | | [`bbq_code_layout`](/api-reference/cpp-api-preprocessing-quantize-bbq#preprocessing-quantize-bbq-bbq-code-layout) | | + +**Returns** + +`void` diff --git a/fern/pages/cpp_api/index.md b/fern/pages/cpp_api/index.md index 3c489df78b..07115a5eca 100644 --- a/fern/pages/cpp_api/index.md +++ b/fern/pages/cpp_api/index.md @@ -43,6 +43,7 @@ These pages are generated from the documented public headers in the cuVS source ## Preprocessing +- [Bbq](/api-reference/cpp-api-preprocessing-quantize-bbq) - [Binary](/api-reference/cpp-api-preprocessing-quantize-binary) - [PCA](/api-reference/cpp-api-preprocessing-pca) - [PQ](/api-reference/cpp-api-preprocessing-quantize-pq) diff --git a/fern/pages/java_api/index.md b/fern/pages/java_api/index.md index 1adeb32dbc..f89a93bc93 100644 --- a/fern/pages/java_api/index.md +++ b/fern/pages/java_api/index.md @@ -6,6 +6,7 @@ For the Apache Lucene codecs built on this API, see the [Lucene API Documentatio ## Common +- [BbqQuantizer](/api-reference/java-api-com-nvidia-cuvs-bbqquantizer) - [CuVSDeviceMatrix](/api-reference/java-api-com-nvidia-cuvs-cuvsdevicematrix) - [CuVSHostMatrix](/api-reference/java-api-com-nvidia-cuvs-cuvshostmatrix) - [CuVSMatrix](/api-reference/java-api-com-nvidia-cuvs-cuvsmatrix) diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-bbqquantizer.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-bbqquantizer.md new file mode 100644 index 0000000000..5f47befedf --- /dev/null +++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-bbqquantizer.md @@ -0,0 +1,31 @@ +--- +slug: api-reference/java-api-com-nvidia-cuvs-bbqquantizer +--- + +# BbqQuantizer + +_Java package: `com.nvidia.cuvs`_ + +```java +public final class BbqQuantizer +``` + +Caller-owned tensors describing one encoded BBQ dataset representation. + +The index stores views over these matrices rather than copying them, so they must stay open +for as long as any index built from them is in use. + +## Public Members + +### withCentroidNormSq + +```java +public Builder withCentroidNormSq(float value) +``` + +Sets the squared L2 norm of the centroid. Required: there is no meaningful default, and +leaving it at zero skews inner-product and cosine distances without reporting an error. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/BbqQuantizer.java:185`_ + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/BbqQuantizer.java:15`_ diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindex.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindex.md index 0484513500..b26f36e2d1 100644 --- a/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindex.md +++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindex.md @@ -613,6 +613,24 @@ an instance of this Builder _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:433`_ +### withBbqDataset + +```java +Builder withBbqDataset(BbqQuantizer... quantizers) +``` + +Builds the graph from one or two encoded BBQ representations. An optional dense dataset +supplied with `#withDataset(CuVSMatrix)` is attached before search; otherwise call +`CagraIndex#updateDataset(PaddedDatasetView)` or +`CagraIndex#updateDataset(PaddedDataset)` before searching. + +The index stores views over the quantizer tensors rather than copying them, so they must +stay open for as long as the index is in use. A dense dataset passed to +`#withDataset(CuVSMatrix)` is owned by the index, as it is for a non-BBQ build, and is +closed with it. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:446`_ + ### withIndexParams ```java @@ -632,7 +650,7 @@ Builder. An instance of this Builder. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:442`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:455`_ ### build @@ -646,6 +664,6 @@ Builds and returns an instance of CagraIndex. an instance of CagraIndex -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:449`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:462`_ _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:26`_ diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/BbqQuantizer.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/BbqQuantizer.java new file mode 100644 index 0000000000..04c193cd47 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/BbqQuantizer.java @@ -0,0 +1,195 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs; + +import java.util.Objects; + +/** + * Caller-owned tensors describing one encoded BBQ dataset representation. + * + *

The index stores views over these matrices rather than copying them, so they must stay open + * for as long as any index built from them is in use. + */ +public final class BbqQuantizer { + public enum CodeLayout { + PACKED_1B(0), + TRANSPOSED_2B(1), + TRANSPOSED_4B(2), + PACKED_4B(3), + PACKED_7B(4), + PACKED_8B(5); + + public final int value; + + CodeLayout(int value) { + this.value = value; + } + } + + private final CuVSMatrix codes; + private final CuVSMatrix lowerIntervals; + private final CuVSMatrix upperIntervals; + private final CuVSMatrix additionalCorrections; + private final CuVSMatrix quantizedComponentSums; + private final CuVSMatrix centroid; + private final CuVSMatrix dequantDelta; + private final CuVSMatrix dequantSumDelta; + private final CuVSMatrix rowNorm; + private final CodeLayout layout; + private final CagraIndexParams.CuvsDistanceType metric; + private final float centroidNormSq; + + private BbqQuantizer(Builder builder) { + codes = Objects.requireNonNull(builder.codes, "codes"); + lowerIntervals = Objects.requireNonNull(builder.lowerIntervals, "lowerIntervals"); + upperIntervals = Objects.requireNonNull(builder.upperIntervals, "upperIntervals"); + additionalCorrections = + Objects.requireNonNull(builder.additionalCorrections, "additionalCorrections"); + quantizedComponentSums = + Objects.requireNonNull(builder.quantizedComponentSums, "quantizedComponentSums"); + centroid = Objects.requireNonNull(builder.centroid, "centroid"); + dequantDelta = Objects.requireNonNull(builder.dequantDelta, "dequantDelta"); + dequantSumDelta = Objects.requireNonNull(builder.dequantSumDelta, "dequantSumDelta"); + rowNorm = Objects.requireNonNull(builder.rowNorm, "rowNorm"); + layout = Objects.requireNonNull(builder.layout, "layout"); + metric = Objects.requireNonNull(builder.metric, "metric"); + centroidNormSq = Objects.requireNonNull(builder.centroidNormSq, "centroidNormSq"); + if (!Float.isFinite(centroidNormSq) || centroidNormSq < 0.0f) { + throw new IllegalArgumentException( + "centroidNormSq must be a finite, non-negative squared norm, but was " + centroidNormSq); + } + } + + public CuVSMatrix getCodes() { + return codes; + } + + public CuVSMatrix getLowerIntervals() { + return lowerIntervals; + } + + public CuVSMatrix getUpperIntervals() { + return upperIntervals; + } + + public CuVSMatrix getAdditionalCorrections() { + return additionalCorrections; + } + + public CuVSMatrix getQuantizedComponentSums() { + return quantizedComponentSums; + } + + public CuVSMatrix getCentroid() { + return centroid; + } + + public CuVSMatrix getDequantDelta() { + return dequantDelta; + } + + public CuVSMatrix getDequantSumDelta() { + return dequantSumDelta; + } + + public CuVSMatrix getRowNorm() { + return rowNorm; + } + + public CodeLayout getLayout() { + return layout; + } + + public CagraIndexParams.CuvsDistanceType getMetric() { + return metric; + } + + public float getCentroidNormSq() { + return centroidNormSq; + } + + public static final class Builder { + private CuVSMatrix codes; + private CuVSMatrix lowerIntervals; + private CuVSMatrix upperIntervals; + private CuVSMatrix additionalCorrections; + private CuVSMatrix quantizedComponentSums; + private CuVSMatrix centroid; + private CuVSMatrix dequantDelta; + private CuVSMatrix dequantSumDelta; + private CuVSMatrix rowNorm; + private CodeLayout layout; + private CagraIndexParams.CuvsDistanceType metric; + private Float centroidNormSq; + + public Builder withCodes(CuVSMatrix value) { + codes = value; + return this; + } + + public Builder withLowerIntervals(CuVSMatrix value) { + lowerIntervals = value; + return this; + } + + public Builder withUpperIntervals(CuVSMatrix value) { + upperIntervals = value; + return this; + } + + public Builder withAdditionalCorrections(CuVSMatrix value) { + additionalCorrections = value; + return this; + } + + public Builder withQuantizedComponentSums(CuVSMatrix value) { + quantizedComponentSums = value; + return this; + } + + public Builder withCentroid(CuVSMatrix value) { + centroid = value; + return this; + } + + public Builder withDequantDelta(CuVSMatrix value) { + dequantDelta = value; + return this; + } + + public Builder withDequantSumDelta(CuVSMatrix value) { + dequantSumDelta = value; + return this; + } + + public Builder withRowNorm(CuVSMatrix value) { + rowNorm = value; + return this; + } + + public Builder withLayout(CodeLayout value) { + layout = value; + return this; + } + + public Builder withMetric(CagraIndexParams.CuvsDistanceType value) { + metric = value; + return this; + } + + /** + * Sets the squared L2 norm of the centroid. Required: there is no meaningful default, and + * leaving it at zero skews inner-product and cosine distances without reporting an error. + */ + public Builder withCentroidNormSq(float value) { + centroidNormSq = value; + return this; + } + + public BbqQuantizer build() { + return new BbqQuantizer(this); + } + } +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index d9472c001d..5320afe4af 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java @@ -433,6 +433,19 @@ interface Builder { */ Builder withDataset(CuVSMatrix dataset); + /** + * Builds the graph from one or two encoded BBQ representations. An optional dense dataset + * supplied with {@link #withDataset(CuVSMatrix)} is attached before search; otherwise call + * {@link CagraIndex#updateDataset(PaddedDatasetView)} or + * {@link CagraIndex#updateDataset(PaddedDataset)} before searching. + * + *

The index stores views over the quantizer tensors rather than copying them, so they must + * stay open for as long as the index is in use. A dense dataset passed to + * {@link #withDataset(CuVSMatrix)} is owned by the index, as it is for a non-BBQ build, and is + * closed with it. + */ + Builder withBbqDataset(BbqQuantizer... quantizers); + /** * Registers an instance of configured {@link CagraIndexParams} with this * Builder. diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index a205931034..491c5213fc 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -50,6 +50,8 @@ * @since 25.02 */ public class CagraIndexImpl implements CagraIndex { + + private static final System.Logger LOG = System.getLogger(CagraIndexImpl.class.getName()); private final CuVSResources resources; private final IndexReference cagraIndexReference; private boolean destroyed; @@ -70,6 +72,19 @@ private CagraIndexImpl( this.cagraIndexReference = build(indexParameters, (CuVSMatrixInternal) dataset); } + private CagraIndexImpl( + CagraIndexParams indexParameters, + CuVSMatrix dataset, + BbqQuantizer[] bbqQuantizers, + CuVSResources resources) { + this.resources = resources; + if (dataset != null && !(dataset instanceof CuVSMatrixInternal)) { + throw new IllegalArgumentException("dataset must be a native CuVS matrix"); + } + this.cagraIndexReference = + buildBbq(indexParameters, (CuVSMatrixInternal) dataset, bbqQuantizers); + } + /** * Constructor for loading the index from an {@link InputStream} * @@ -138,14 +153,14 @@ private void checkNotDestroyed() { public void close() { checkNotDestroyed(); try { - int returnValue = cuvsCagraIndexDestroy(cagraIndexReference.getMemorySegment()); - checkCuVSError(returnValue, "cuvsCagraIndexDestroy"); + // Both are released whatever either reports: `destroyed` below makes a retry impossible. + quietly( + () -> + checkCuVSError( + cuvsCagraIndexDestroy(cagraIndexReference.getMemorySegment()), + "cuvsCagraIndexDestroy")); if (cagraIndexReference.datasetOwner != null) { - try { - cagraIndexReference.datasetOwner.close(); - } catch (Exception e) { - throw new RuntimeException("Failed to destroy CAGRA dataset", e); - } + quietly(cagraIndexReference.datasetOwner::close); } } finally { destroyed = true; @@ -163,59 +178,224 @@ private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixInterna long rows = dataset.size(); try (var indexParams = segmentFromIndexParams(indexParameters); - var localArena = Arena.ofConfined()) { + var localArena = Arena.ofConfined(); + var resourcesAccessor = resources.access()) { MemorySegment indexParamsMemorySegment = indexParams.handle(); int numWriterThreads = indexParameters != null ? indexParameters.getNumWriterThreads() : 1; - omp_set_num_threads(numWriterThreads); - - var datasetTensor = dataset.toTensor(localArena); + MemorySegment datasetView = MemorySegment.NULL; + boolean success = false; var index = createCagraIndex(); - if (cuvsCagraIndexParams.build_algo(indexParamsMemorySegment) - == 1) { // when build algo is IVF_PQ - MemorySegment cuvsIvfPqIndexParamsMS = - cuvsIvfPqParams.ivf_pq_build_params( - cuvsCagraIndexParams.graph_build_params(indexParamsMemorySegment)); - int n_lists = cuvsIvfPqIndexParams.n_lists(cuvsIvfPqIndexParamsMS); - // As rows cannot be less than n_lists value so trim down. - cuvsIvfPqIndexParams.n_lists( - cuvsIvfPqIndexParamsMS, (int) (rows < n_lists ? rows : n_lists)); - } - try (var resourcesAccessor = resources.access()) { + try { + omp_set_num_threads(numWriterThreads); + + if (cuvsCagraIndexParams.build_algo(indexParamsMemorySegment) + == 1) { // when build algo is IVF_PQ + MemorySegment cuvsIvfPqIndexParamsMS = + cuvsIvfPqParams.ivf_pq_build_params( + cuvsCagraIndexParams.graph_build_params(indexParamsMemorySegment)); + int n_lists = cuvsIvfPqIndexParams.n_lists(cuvsIvfPqIndexParamsMS); + // As rows cannot be less than n_lists value so trim down. + cuvsIvfPqIndexParams.n_lists( + cuvsIvfPqIndexParamsMS, (int) (rows < n_lists ? rows : n_lists)); + } + var cuvsRes = resourcesAccessor.handle(); // TODO: do we need a stream sync here? - var returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); - MemorySegment datasetView = MemorySegment.NULL; - try { - MemorySegment datasetViewPtr = localArena.allocate(cuvsDataset_t); - if (isCagraPaddedLayout(dataset)) { - returnValue = cuvsDatasetMakePaddedView(cuvsRes, datasetTensor, datasetViewPtr); - checkCuVSError(returnValue, "cuvsDatasetMakePaddedView"); - } else { - returnValue = cuvsDatasetMakeStandardView(cuvsRes, datasetTensor, datasetViewPtr); - checkCuVSError(returnValue, "cuvsDatasetMakeStandardView"); - } - datasetView = datasetViewPtr.get(cuvsDataset_t, 0); + datasetView = + isCagraPaddedLayout(dataset) + ? makePaddedViewHandle(cuvsRes, dataset, localArena) + : makeStandardViewHandle(cuvsRes, dataset, localArena); - returnValue = cuvsCagraBuild(cuvsRes, indexParamsMemorySegment, datasetView, index); - checkCuVSError(returnValue, "cuvsCagraBuild"); - } finally { - if (datasetView.address() != 0) { - checkCuVSError(cuvsDatasetDestroy(datasetView), "cuvsDatasetDestroy"); - } + checkCuVSError( + cuvsCagraBuild(cuvsRes, indexParamsMemorySegment, datasetView, index), + "cuvsCagraBuild"); + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); + + success = true; + return new IndexReference(index, dataset); + } finally { + MemorySegment view = datasetView; + if (view.address() != 0) { + quietly(() -> checkCuVSError(cuvsDatasetDestroy(view), "cuvsDatasetDestroy")); + } + // Ownership transfers only on success, so on any other exit the index is ours and the + // caller's matrix is not. + if (!success) { + quietly(() -> checkCuVSError(cuvsCagraIndexDestroy(index), "cuvsCagraIndexDestroy")); } + quietly(() -> omp_set_num_threads(1)); + } + } + } - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); + private IndexReference buildBbq( + CagraIndexParams indexParameters, CuVSMatrixInternal dataset, BbqQuantizer[] bbqQuantizers) { + if (bbqQuantizers == null || bbqQuantizers.length < 1 || bbqQuantizers.length > 2) { + throw new IllegalArgumentException("BBQ build requires one or two quantizers"); + } + for (BbqQuantizer quantizer : bbqQuantizers) { + Objects.requireNonNull(quantizer); + } + + try (var indexParams = segmentFromIndexParams(indexParameters); + var localArena = Arena.ofConfined(); + var resourcesAccessor = resources.access()) { + var cuvsRes = resourcesAccessor.handle(); + var quantizerHandles = new ArrayList(bbqQuantizers.length); + MemorySegment bbqDataset = MemorySegment.NULL; + MemorySegment paddedDataset = MemorySegment.NULL; + int numWriterThreads = indexParameters != null ? indexParameters.getNumWriterThreads() : 1; + boolean success = false; + var index = createCagraIndex(); + + try { + // nn-descent, which a BBQ build always uses, does a substantial amount of its work on + // the host under OpenMP. + omp_set_num_threads(numWriterThreads); + for (BbqQuantizer quantizer : bbqQuantizers) { + var quantizerPtr = localArena.allocate(cuvsBbqQuantizer_t); + checkCuVSError( + cuvsBbqQuantizerCreateView( + tensor(quantizer.getCodes(), "codes", localArena), + vectorTensor(quantizer.getLowerIntervals(), "lowerIntervals", localArena), + vectorTensor(quantizer.getUpperIntervals(), "upperIntervals", localArena), + vectorTensor( + quantizer.getAdditionalCorrections(), "additionalCorrections", localArena), + vectorTensor( + quantizer.getQuantizedComponentSums(), "quantizedComponentSums", localArena), + vectorTensor(quantizer.getCentroid(), "centroid", localArena), + vectorTensor(quantizer.getDequantDelta(), "dequantDelta", localArena), + vectorTensor(quantizer.getDequantSumDelta(), "dequantSumDelta", localArena), + vectorTensor(quantizer.getRowNorm(), "rowNorm", localArena), + quantizer.getLayout().value, + quantizer.getMetric().value, + quantizer.getCentroidNormSq(), + quantizerPtr), + "cuvsBbqQuantizerCreateView"); + quantizerHandles.add(quantizerPtr.get(cuvsBbqQuantizer_t, 0)); + } + + var quantizerArray = localArena.allocate(ValueLayout.ADDRESS, quantizerHandles.size()); + for (int i = 0; i < quantizerHandles.size(); ++i) { + quantizerArray.setAtIndex(ValueLayout.ADDRESS, i, quantizerHandles.get(i)); + } + var bbqDatasetPtr = localArena.allocate(cuvsDataset_t); + checkCuVSError( + cuvsDatasetMakeBbqView(cuvsRes, quantizerArray, quantizerHandles.size(), bbqDatasetPtr), + "cuvsDatasetMakeBbqView"); + bbqDataset = bbqDatasetPtr.get(cuvsDataset_t, 0); + + checkCuVSError( + cuvsCagraBuild(cuvsRes, indexParams.handle(), bbqDataset, index), "cuvsCagraBuild"); + + IndexReference result; + if (dataset == null) { + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); + result = new IndexReference(index, null, null); + } else { + paddedDataset = makeDevicePaddedDataset(cuvsRes, dataset, localArena); + checkCuVSError( + cuvsCagraUpdateDataset(cuvsRes, paddedDataset, index), "cuvsCagraUpdateDataset"); + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); + + // The index owns both the padded dataset and the caller's matrix from here on. + result = + new IndexReference( + index, + dataset, + new PaddedDatasetOwner(new DatasetCloseDelegate(paddedDataset), dataset)); + } + success = true; + return result; + } finally { + // Ownership transfers only on success, so on any other exit the index and the padded + // dataset are ours and the caller's matrix is not. Destroy the index first: it holds a + // view into the padded dataset. + if (!success) { + quietly(() -> checkCuVSError(cuvsCagraIndexDestroy(index), "cuvsCagraIndexDestroy")); + MemorySegment padded = paddedDataset; + if (padded.address() != 0) { + quietly(() -> checkCuVSError(cuvsDatasetDestroy(padded), "cuvsDatasetDestroy")); + } + } + MemorySegment bbqView = bbqDataset; + if (bbqView.address() != 0) { + quietly(() -> checkCuVSError(cuvsDatasetDestroy(bbqView), "cuvsDatasetDestroy")); + } + for (MemorySegment quantizer : quantizerHandles) { + quietly( + () -> checkCuVSError(cuvsBbqQuantizerDestroy(quantizer), "cuvsBbqQuantizerDestroy")); + } + quietly(() -> omp_set_num_threads(1)); } - omp_set_num_threads(1); + } + } - return new IndexReference(index, dataset); + /** + * Validates a caller-supplied BBQ component and returns it in its internal form. + * + *

The native side reads every one of these through {@code from_dlpack}, which accepts only + * compact row-major memory, so a padded matrix has to be rejected here. Left to the C layer it + * surfaces as an error naming the quantizer entry point, with no indication of which of the nine + * tensors was at fault. + */ + private static CuVSMatrixInternal bbqComponent(CuVSMatrix matrix, String name) { + if (!(matrix instanceof CuVSDeviceMatrix) || !(matrix instanceof CuVSMatrixInternal internal)) { + throw new IllegalArgumentException("BBQ tensor '" + name + "' must be a device CuVS matrix"); } + long rowStride = internal.rowStride(); + if (rowStride > 0 && rowStride != matrix.columns()) { + throw new IllegalArgumentException( + "BBQ tensor '" + + name + + "' must be contiguous, but its row stride (" + + rowStride + + ") does not match its column count (" + + matrix.columns() + + ")"); + } + return internal; + } + + /** + * The rank-2 codes tensor. + * + *

Built without strides rather than through {@code toTensor}, which emits + * {@code {rowStride, columnStride}} whenever a row stride is set and so would pass a column + * stride of -1 for a matrix that only declared a row stride. Omitting strides is correct here + * because {@link #bbqComponent} has already rejected anything non-contiguous. + */ + private static MemorySegment tensor(CuVSMatrix matrix, String name, Arena arena) { + var internal = bbqComponent(matrix, name); + return prepareTensor( + arena, + internal.memorySegment(), + new long[] {matrix.size(), matrix.columns()}, + internal.code(), + internal.bits(), + kDLCUDA()); + } + + /** + * One of the rank-1 correction vectors, flattened from the caller's matrix. + * + *

Flattening to {@code size * columns} is only valid because {@link #bbqComponent} has + * rejected anything non-contiguous, so there is no row padding to skip. + */ + private static MemorySegment vectorTensor(CuVSMatrix matrix, String name, Arena arena) { + var internal = bbqComponent(matrix, name); + return prepareTensor( + arena, + internal.memorySegment(), + new long[] {matrix.size() * matrix.columns()}, + internal.code(), + internal.bits(), + kDLCUDA()); } private static MemorySegment createCagraIndex() { @@ -422,27 +602,73 @@ public MemorySegment getIndexHandle() { return cagraIndexReference.getMemorySegment(); } - @Override - public CagraIndex.PaddedDataset makePaddedDataset(CuVSMatrix dataset) throws Throwable { - checkNotDestroyed(); + /** Narrows a caller-supplied dataset to the internal form the native layer can read. */ + private static CuVSMatrixInternal requireInternalMatrix(CuVSMatrix dataset) { Objects.requireNonNull(dataset); - if (!(dataset instanceof CuVSMatrixInternal datasetInternal)) { + if (!(dataset instanceof CuVSMatrixInternal internal)) { throw new IllegalArgumentException("dataset must be a CuVSMatrixInternal matrix"); } + return internal; + } + + /** Creates a native padded dataset that views the caller's storage rather than copying it. */ + private static MemorySegment makePaddedViewHandle( + long cuvsRes, CuVSMatrixInternal dataset, Arena arena) { + MemorySegment out = arena.allocate(cuvsDataset_t); + checkCuVSError( + cuvsDatasetMakePaddedView(cuvsRes, dataset.toTensor(arena), out), + "cuvsDatasetMakePaddedView"); + return out.get(cuvsDataset_t, 0); + } + + /** Creates a native padded dataset owning its own copy of the caller's storage. */ + private static MemorySegment makePaddedCopyHandle( + long cuvsRes, CuVSMatrixInternal dataset, int targetMemType, Arena arena) { + MemorySegment out = arena.allocate(cuvsDataset_t); + checkCuVSError( + cuvsDatasetMakePadded(cuvsRes, dataset.toTensor(arena), targetMemType, out), + "cuvsDatasetMakePadded"); + return out.get(cuvsDataset_t, 0); + } + + /** Creates a native standard (unpadded) dataset that views the caller's storage. */ + private static MemorySegment makeStandardViewHandle( + long cuvsRes, CuVSMatrixInternal dataset, Arena arena) { + MemorySegment out = arena.allocate(cuvsDataset_t); + checkCuVSError( + cuvsDatasetMakeStandardView(cuvsRes, dataset.toTensor(arena), out), + "cuvsDatasetMakeStandardView"); + return out.get(cuvsDataset_t, 0); + } + + /** + * Produces a device padded dataset for {@code dataset}: a view when it already has CAGRA's padded + * row width, a copy otherwise. Always device-resident, because {@code cuvsCagraUpdateDataset} + * accepts nothing else. + */ + private static MemorySegment makeDevicePaddedDataset( + long cuvsRes, CuVSMatrixInternal dataset, Arena arena) { + return dataset instanceof CuVSDeviceMatrix && isCagraPaddedLayout(dataset) + ? makePaddedViewHandle(cuvsRes, dataset, arena) + : makePaddedCopyHandle(cuvsRes, dataset, CUVS_DATASET_MEM_TYPE_DEVICE(), arena); + } + + @Override + public CagraIndex.PaddedDataset makePaddedDataset(CuVSMatrix dataset) throws Throwable { + checkNotDestroyed(); + var datasetInternal = requireInternalMatrix(dataset); try (var localArena = Arena.ofConfined(); var resourcesAccessor = resources.access()) { var cuvsRes = resourcesAccessor.handle(); - var datasetTensor = datasetInternal.toTensor(localArena); + // A padded copy stays wherever the source lives; only attaching one to an index forces it + // onto the device. int targetMemType = (datasetInternal instanceof CuVSHostMatrixImpl) ? CUVS_DATASET_MEM_TYPE_HOST() : CUVS_DATASET_MEM_TYPE_DEVICE(); - MemorySegment paddedDatasetPtr = localArena.allocate(cuvsDataset_t); - var returnValue = - cuvsDatasetMakePadded(cuvsRes, datasetTensor, targetMemType, paddedDatasetPtr); - checkCuVSError(returnValue, "cuvsDatasetMakePadded"); - MemorySegment paddedDataset = paddedDatasetPtr.get(cuvsDataset_t, 0); + MemorySegment paddedDataset = + makePaddedCopyHandle(cuvsRes, datasetInternal, targetMemType, localArena); var out = new CagraIndex.PaddedDataset(); out.setDelegate(new DatasetCloseDelegate(paddedDataset), paddedDataset.address()); @@ -453,19 +679,12 @@ public CagraIndex.PaddedDataset makePaddedDataset(CuVSMatrix dataset) throws Thr @Override public CagraIndex.PaddedDatasetView makePaddedDatasetView(CuVSMatrix dataset) throws Throwable { checkNotDestroyed(); - Objects.requireNonNull(dataset); - if (!(dataset instanceof CuVSMatrixInternal datasetInternal)) { - throw new IllegalArgumentException("dataset must be a CuVSMatrixInternal matrix"); - } + var datasetInternal = requireInternalMatrix(dataset); try (var localArena = Arena.ofConfined(); var resourcesAccessor = resources.access()) { var cuvsRes = resourcesAccessor.handle(); - var datasetTensor = datasetInternal.toTensor(localArena); - MemorySegment paddedViewPtr = localArena.allocate(cuvsDataset_t); - var returnValue = cuvsDatasetMakePaddedView(cuvsRes, datasetTensor, paddedViewPtr); - checkCuVSError(returnValue, "cuvsDatasetMakePaddedView"); - MemorySegment paddedView = paddedViewPtr.get(cuvsDataset_t, 0); + MemorySegment paddedView = makePaddedViewHandle(cuvsRes, datasetInternal, localArena); var out = new CagraIndex.PaddedDatasetView(); out.setDelegate(new DatasetCloseDelegate(paddedView), paddedView.address()); @@ -477,19 +696,12 @@ public CagraIndex.PaddedDatasetView makePaddedDatasetView(CuVSMatrix dataset) th public CagraIndex.StandardDatasetView makeStandardDatasetView(CuVSMatrix dataset) throws Throwable { checkNotDestroyed(); - Objects.requireNonNull(dataset); - if (!(dataset instanceof CuVSMatrixInternal datasetInternal)) { - throw new IllegalArgumentException("dataset must be a CuVSMatrixInternal matrix"); - } + var datasetInternal = requireInternalMatrix(dataset); try (var localArena = Arena.ofConfined(); var resourcesAccessor = resources.access()) { var cuvsRes = resourcesAccessor.handle(); - var datasetTensor = datasetInternal.toTensor(localArena); - MemorySegment standardViewPtr = localArena.allocate(cuvsDataset_t); - var returnValue = cuvsDatasetMakeStandardView(cuvsRes, datasetTensor, standardViewPtr); - checkCuVSError(returnValue, "cuvsDatasetMakeStandardView"); - MemorySegment standardView = standardViewPtr.get(cuvsDataset_t, 0); + MemorySegment standardView = makeStandardViewHandle(cuvsRes, datasetInternal, localArena); var out = new CagraIndex.StandardDatasetView(); out.setDelegate(new DatasetCloseDelegate(standardView), standardView.address()); @@ -760,6 +972,44 @@ private IndexReference deserialize( } } + /** A cleanup action that is allowed to fail, including with a checked exception. */ + @FunctionalInterface + private interface CleanupStep { + void run() throws Exception; + } + + /** + * Runs a cleanup step, logging any failure rather than propagating it. + * + * TODO: replace this with a more comprehensive mechanism as part of https://github.com/NVIDIA/cuvs/issues/2670 + */ + private static void quietly(CleanupStep cleanup) { + try { + cleanup.run(); + } catch (Exception thrown) { + // Error is left to propagate: an OutOfMemoryError or a linkage failure says the JVM is in + // no state to carry on, and swallowing it would hide that. + LOG.log( + System.Logger.Level.WARNING, + "Failed to release a native cuVS resource; it has probably been leaked", + thrown); + } + } + + /** + * Owns both halves of an attached dataset: the native {@code cuvsDataset_t} and the matrix whose + * storage it may be viewing. The native dataset is released first, since closing the matrix + * underneath a live view would leave it dangling. + */ + private record PaddedDatasetOwner(DatasetCloseDelegate nativeDataset, CuVSMatrix matrix) + implements AutoCloseable { + @Override + public void close() { + quietly(nativeDataset::close); + quietly(matrix::close); + } + } + private static final class DatasetCloseDelegate implements AutoCloseable { private MemorySegment handle; @@ -1096,6 +1346,7 @@ public static class Builder implements CagraIndex.Builder { private CagraIndexParams cagraIndexParams; private final CuVSResources cuvsResources; private CuVSMatrix graph; + private BbqQuantizer[] bbqQuantizers; public Builder(CuVSResources cuvsResources) { this.cuvsResources = cuvsResources; @@ -1133,6 +1384,12 @@ public Builder withDataset(CuVSMatrix dataset) { return this; } + @Override + public Builder withBbqDataset(BbqQuantizer... quantizers) { + this.bbqQuantizers = quantizers == null ? null : quantizers.clone(); + return this; + } + @Override public Builder withIndexParams(CagraIndexParams cagraIndexParameters) { this.cagraIndexParams = cagraIndexParameters; @@ -1153,6 +1410,8 @@ public CagraIndexImpl build() throws Throwable { } return new CagraIndexImpl( cagraIndexParams.getCuvsDistanceType(), graph, dataset, cuvsResources); + } else if (bbqQuantizers != null) { + return new CagraIndexImpl(cagraIndexParams, dataset, bbqQuantizers, cuvsResources); } else if (dataset != null) { return new CagraIndexImpl(cagraIndexParams, dataset, cuvsResources); } else { diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBbqIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBbqIT.java new file mode 100644 index 0000000000..31e659cb96 --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBbqIT.java @@ -0,0 +1,673 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs; + +import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import java.lang.invoke.MethodHandles; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import org.junit.Before; +import org.junit.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * CAGRA indices built from BBQ quantizers via {@link CagraIndex.Builder#withBbqDataset}. + * + *

{@code gpuSearchBaselineWithoutBbq} is a control: the same search over a graph plain CAGRA + * built from the same vectors. If it regresses alongside the BBQ case, look at the search + * parameters or the data before suspecting quantization. + */ +public class CagraBbqIT extends CuVSTestCase { + + private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); + + private static final int ROWS = 20000; + private static final int DIM = 64; + private static final int QUERIES = 32; + private static final int CLUSTERS = 40; + private static final double CLUSTER_SPREAD = 3.0; + + /** A neighbor within this much of the true k-th distance is as good an answer. */ + private static final float TIE_TOLERANCE = 0.01f; + + private static final int TOP_K = 5; + + /** + * Floor for recall@{@value #TOP_K}. Both searches measure 1.0 on this data with the pinned search + * parameters, so this leaves room for nn-descent's run-to-run variation without being loose + * enough to pass a real regression. + */ + private static final double RECALL_FLOOR = 0.95; + + /** PACKED_1B encodes one bit per dimension. */ + private static final int CODE_BYTES = (DIM + 7) / 8; + + private float[][] dataset; + private float[][] queries; + + @Before + public void setup() { + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); + // Fixed seed so the dataset and queries are stable. Note this does not make recall stable: + // nn-descent is stochastic, so the measured values still move by a few points between runs. + random = new java.util.Random(42); + float[][] centers = randomVectors(CLUSTERS); + dataset = clusteredVectors(centers, ROWS); + queries = clusteredVectors(centers, QUERIES); + } + + /** Builds the graph from quantizers with a dense dataset attached, then searches on the GPU. */ + @Test + public void gpuBuildGpuSearch() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + // Owned by the index once built, so it is not closed here. + CuVSMatrix dense = newDenseDataset(resources); + try (CagraIndex index = bbqIndex(resources, quantizer, dense); + CuVSMatrix queryVectors = CuVSMatrix.ofArray(queries)) { + CagraQuery query = + new CagraQuery.Builder(resources) + .withTopK(TOP_K) + .withSearchParams(searchParams()) + .withQueryVectors(queryVectors) + .withMapping(SearchResults.IDENTITY_MAPPING) + .build(); + assertRecall("gpu-search", index.search(query)); + } + } + } + + /** + * Plain CAGRA over the same vectors, as a reference point for the BBQ recall above. Both sides + * run the same GPU search, so the difference between the two numbers is the cost of building the + * graph from quantized codes. + */ + @Test + public void gpuSearchBaselineWithoutBbq() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create()) { + // Owned by the index once built, so it is not closed here. + CuVSMatrix dense = newDenseDataset(resources); + try (CagraIndex index = + CagraIndex.newBuilder(resources) + .withIndexParams( + new CagraIndexParams.Builder() + .withMetric(CuvsDistanceType.L2Expanded) + .withGraphDegree(32) + .withIntermediateGraphDegree(64) + .build()) + .withDataset(dense) + .build(); + CuVSMatrix queryVectors = CuVSMatrix.ofArray(queries)) { + CagraQuery query = + new CagraQuery.Builder(resources) + .withTopK(TOP_K) + .withSearchParams(searchParams()) + .withQueryVectors(queryVectors) + .withMapping(SearchResults.IDENTITY_MAPPING) + .build(); + assertRecall("baseline-gpu-search", index.search(query)); + } + } + } + + /** + * Extracting the graph from a graph-only BBQ index, which is how cuvs-lucene consumes one: no + * dense dataset is attached, so this is the memory-efficient shape BBQ exists for, and nothing + * native dereferences the BBQ dataset view on this path. + */ + @Test + public void graphOnlyBbqBuildProducesUsableGraph() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + try (CagraIndex index = bbqIndex(resources, quantizer, null); + CuVSMatrix graph = index.getGraph()) { + assertEquals("one adjacency row per vector", ROWS, graph.size()); + int degree = (int) graph.columns(); + assertTrue("graph degree should be positive", degree > 0); + + int[][] adjacency = readAdjacency(graph); + int selfLoops = 0; + for (int node = 0; node < ROWS; node++) { + for (int neighbor : adjacency[node]) { + assertTrue( + "neighbor id " + neighbor + " of node " + node + " is out of range", + neighbor >= 0 && neighbor < ROWS); + if (neighbor == node) { + selfLoops++; + } + } + } + // A handful of self-references is normal padding; a graph made mostly of them is not. + assertTrue( + "graph is mostly self-references (" + selfLoops + " of " + ROWS * degree + ")", + selfLoops < ROWS * degree / 10); + } + } + } + + /** + * A dense dataset handed to the builder is owned by the resulting index, exactly as it is for a + * non-BBQ build, so closing the index releases it. Host matrices are arena-backed, which makes + * the release observable: reads fail once the arena is gone. + */ + @Test + public void indexOwnsTheDenseDatasetItWasGiven() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + // hostBuilder yields an arena-backed matrix; ofArray's close() is a no-op and would make + // this test pass whether or not the index released anything. + var denseBuilder = CuVSMatrix.hostBuilder(ROWS, DIM, CuVSMatrix.DataType.FLOAT); + for (float[] row : dataset) { + denseBuilder.addVector(row); + } + CuVSMatrix dense = denseBuilder.build(); + try (CagraIndex index = bbqIndex(resources, quantizer, dense)) { + assertNotNull("the dataset should be readable while the index is open", dense.getRow(0)); + } + // The read, not getRow itself, is what validates the arena's scope. + assertThrows( + "closing the index should have released the dataset it was given", + IllegalStateException.class, + () -> dense.getRow(0).getAsFloat(0)); + } + } + + /** + * Every quantizer component is required. {@code centroidNormSq} is the one that cannot rely on a + * null check, since it is a primitive whose unset state would otherwise read as a legitimate + * zero and quietly skew inner-product and cosine distances. + */ + @Test + public void quantizerMustRejectMissingCentroidNormSq() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + var q = quantizer.quantizer(); + assertThrows( + "omitting withCentroidNormSq should fail fast, like every other component does", + RuntimeException.class, + () -> + new BbqQuantizer.Builder() + .withCodes(q.getCodes()) + .withLowerIntervals(q.getLowerIntervals()) + .withUpperIntervals(q.getUpperIntervals()) + .withAdditionalCorrections(q.getAdditionalCorrections()) + .withQuantizedComponentSums(q.getQuantizedComponentSums()) + .withCentroid(q.getCentroid()) + .withDequantDelta(q.getDequantDelta()) + .withDequantSumDelta(q.getDequantSumDelta()) + .withRowNorm(q.getRowNorm()) + .withLayout(BbqQuantizer.CodeLayout.PACKED_1B) + .withMetric(CuvsDistanceType.L2Expanded) + // withCentroidNormSq deliberately omitted + .build()); + } + } + + /** A build takes one or two encoded representations; anything else is a usage error. */ + @Test + public void buildRejectsAnUnusableQuantizerCount() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + var q = quantizer.quantizer(); + assertThrows( + "a build needs at least one representation", + IllegalArgumentException.class, + () -> bbqIndex(resources, new BbqQuantizer[] {}, null)); + assertThrows( + "three representations are more than the native side accepts", + IllegalArgumentException.class, + () -> bbqIndex(resources, new BbqQuantizer[] {q, q, q}, null)); + assertThrows( + "a null representation should be caught before it reaches native code", + NullPointerException.class, + () -> bbqIndex(resources, new BbqQuantizer[] {q, null}, null)); + } + } + + /** Quantizer components live on the device; a host matrix cannot be read by the native side. */ + @Test + public void hostQuantizerComponentIsRejected() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + var q = quantizer.quantizer(); + var hostBuilder = CuVSMatrix.hostBuilder(ROWS, 1, CuVSMatrix.DataType.FLOAT); + for (int i = 0; i < ROWS; i++) { + hostBuilder.addVector(new float[] {1.0f}); + } + try (CuVSMatrix hostRowNorm = hostBuilder.build()) { + var withHostComponent = + new BbqQuantizer.Builder() + .withCodes(q.getCodes()) + .withLowerIntervals(q.getLowerIntervals()) + .withUpperIntervals(q.getUpperIntervals()) + .withAdditionalCorrections(q.getAdditionalCorrections()) + .withQuantizedComponentSums(q.getQuantizedComponentSums()) + .withCentroid(q.getCentroid()) + .withDequantDelta(q.getDequantDelta()) + .withDequantSumDelta(q.getDequantSumDelta()) + .withRowNorm(hostRowNorm) + .withLayout(BbqQuantizer.CodeLayout.PACKED_1B) + .withMetric(CuvsDistanceType.L2Expanded) + .withCentroidNormSq(q.getCentroidNormSq()) + .build(); + var failure = + assertThrows( + IllegalArgumentException.class, + () -> bbqIndex(resources, new BbqQuantizer[] {withHostComponent}, null)); + assertTrue( + "the message should name the offending component, but was: " + failure.getMessage(), + failure.getMessage().contains("rowNorm")); + } + } + } + + /** A graph-only index carries no vectors, so searching it before attaching a dataset fails. */ + @Test + public void graphOnlyIndexCannotBeSearchedUntilADatasetIsAttached() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + try (CagraIndex index = bbqIndex(resources, quantizer, null); + CuVSMatrix queryVectors = CuVSMatrix.ofArray(queries)) { + CagraQuery query = + new CagraQuery.Builder(resources) + .withTopK(TOP_K) + .withSearchParams(searchParams()) + .withQueryVectors(queryVectors) + .withMapping(SearchResults.IDENTITY_MAPPING) + .build(); + var failure = assertThrows(RuntimeException.class, () -> index.search(query)); + assertTrue( + "the message should point at updateDataset, but was: " + failure.getMessage(), + failure.getMessage().contains("cuvsCagraUpdateDataset")); + } + } + } + + /** The documented follow-up: attach a dataset to a graph-only index and it becomes searchable. */ + @Test + public void attachingADatasetMakesAGraphOnlyIndexSearchable() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources); + CuVSMatrix dense = newDenseDataset(resources)) { + try (CagraIndex index = bbqIndex(resources, quantizer, null); + CuVSMatrix queryVectors = CuVSMatrix.ofArray(queries)) { + try (CagraIndex.PaddedDatasetView padded = index.makePaddedDatasetView(dense)) { + index.updateDataset(padded); + CagraQuery query = + new CagraQuery.Builder(resources) + .withTopK(TOP_K) + .withSearchParams(searchParams()) + .withQueryVectors(queryVectors) + .withMapping(SearchResults.IDENTITY_MAPPING) + .build(); + var results = index.search(query).getResults(); + assertEquals("one result set per query", QUERIES, results.size()); + for (var perQuery : results) { + assertEquals("each query should return topK neighbors", TOP_K, perQuery.size()); + } + } + } + } + } + + /** + * A component may declare a row stride explicitly as long as it equals the column count. That is + * still contiguous, so it has to be accepted and passed to the native side without strides. + */ + @Test + public void explicitlyContiguousCodesMatrixIsAccepted() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + var q = quantizer.quantizer(); + var codesBuilder = + CuVSMatrix.deviceBuilder( + resources, ROWS, CODE_BYTES, CODE_BYTES, -1, CuVSMatrix.DataType.BYTE); + for (int i = 0; i < ROWS; i++) { + codesBuilder.addVector(new byte[CODE_BYTES]); + } + try (CuVSMatrix stridedCodes = codesBuilder.build()) { + var withStridedCodes = + new BbqQuantizer.Builder() + .withCodes(stridedCodes) + .withLowerIntervals(q.getLowerIntervals()) + .withUpperIntervals(q.getUpperIntervals()) + .withAdditionalCorrections(q.getAdditionalCorrections()) + .withQuantizedComponentSums(q.getQuantizedComponentSums()) + .withCentroid(q.getCentroid()) + .withDequantDelta(q.getDequantDelta()) + .withDequantSumDelta(q.getDequantSumDelta()) + .withRowNorm(q.getRowNorm()) + .withLayout(BbqQuantizer.CodeLayout.PACKED_1B) + .withMetric(CuvsDistanceType.L2Expanded) + .withCentroidNormSq(q.getCentroidNormSq()) + .build(); + try (CagraIndex index = bbqIndex(resources, new BbqQuantizer[] {withStridedCodes}, null); + CuVSMatrix graph = index.getGraph()) { + assertEquals("one adjacency row per vector", ROWS, graph.size()); + } + } + } + } + + /** + * Every quantizer component has to be contiguous, because the native side reads them through + * {@code from_dlpack}, which accepts only compact row-major memory. The check belongs on the Java + * side so the message can name the offending tensor rather than just the entry point. + */ + @Test + public void paddedCodesMatrixMustBeRejectedClearly() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + var q = quantizer.quantizer(); + var paddedCodesBuilder = + CuVSMatrix.deviceBuilder( + resources, ROWS, CODE_BYTES, CODE_BYTES * 2, -1, CuVSMatrix.DataType.BYTE); + for (int i = 0; i < ROWS; i++) { + paddedCodesBuilder.addVector(new byte[CODE_BYTES]); + } + try (CuVSMatrix paddedCodes = paddedCodesBuilder.build()) { + var padded = + new BbqQuantizer.Builder() + .withCodes(paddedCodes) + .withLowerIntervals(q.getLowerIntervals()) + .withUpperIntervals(q.getUpperIntervals()) + .withAdditionalCorrections(q.getAdditionalCorrections()) + .withQuantizedComponentSums(q.getQuantizedComponentSums()) + .withCentroid(q.getCentroid()) + .withDequantDelta(q.getDequantDelta()) + .withDequantSumDelta(q.getDequantSumDelta()) + .withRowNorm(q.getRowNorm()) + .withLayout(BbqQuantizer.CodeLayout.PACKED_1B) + .withMetric(CuvsDistanceType.L2Expanded) + .withCentroidNormSq(q.getCentroidNormSq()) + .build(); + var failure = + assertThrows( + "a padded codes matrix should be rejected by the Java guard, as vector tensors are", + IllegalArgumentException.class, + () -> + CagraIndex.newBuilder(resources) + .withIndexParams( + new CagraIndexParams.Builder() + .withMetric(CuvsDistanceType.L2Expanded) + .build()) + .withBbqDataset(padded) + .build()); + assertTrue( + "the message should name the offending component, but was: " + failure.getMessage(), + failure.getMessage().contains("codes")); + } + } + } + + private static int[][] readAdjacency(CuVSMatrix graph) { + int nodes = (int) graph.size(); + int degree = (int) graph.columns(); + int[][] adjacency = new int[nodes][degree]; + for (int i = 0; i < nodes; i++) { + RowView row = graph.getRow(i); + for (int j = 0; j < degree; j++) { + adjacency[i][j] = row.getAsInt(j); + } + } + return adjacency; + } + + /** + * Recall@{@value #TOP_K} against brute-force ground truth, scored by distance rather than by id. + * + *

A returned neighbor counts when it is at least as close as the true {@value #TOP_K}-th, so + * an answer that is equally good but differently tied is not penalised. Scoring by id instead + * measures how the data happens to break ties, which on synthetic vectors is mostly noise. + */ + private void assertRecall(String label, SearchResults results) { + int hits = 0; + int total = 0; + for (int q = 0; q < QUERIES; q++) { + float[] distances = new float[ROWS]; + for (int i = 0; i < ROWS; i++) { + distances[i] = squaredDistance(queries[q], dataset[i]); + } + float[] sorted = distances.clone(); + Arrays.sort(sorted); + float acceptable = sorted[TOP_K - 1] * (1.0f + TIE_TOLERANCE); + for (int id : results.getResults().get(q).keySet()) { + total++; + // CAGRA reports a missing neighbor as an out-of-range id; count it as a miss rather than + // indexing with it. + if (id >= 0 && id < ROWS && distances[id] <= acceptable) { + hits++; + } + } + } + double recall = (double) hits / total; + log.info("{} recall@{} = {}", label, TOP_K, recall); + assertTrue(label + " recall@" + TOP_K + " was only " + recall, recall >= RECALL_FLOOR); + } + + /** + * Search effort is pinned rather than left to the defaults: recall is being asserted here, and a + * default that changes would silently move the number the assertion is checking. + */ + private static CagraSearchParams searchParams() { + return new CagraSearchParams.Builder() + .withAlgo(CagraSearchParams.SearchAlgo.SINGLE_CTA) + .withItopkSize(256) + .build(); + } + + private static float squaredDistance(float[] a, float[] b) { + float sum = 0.0f; + for (int i = 0; i < a.length; i++) { + float diff = a[i] - b[i]; + sum += diff * diff; + } + return sum; + } + + private static CagraIndex bbqIndex( + CuVSResources resources, Quantizer quantizer, CuVSMatrix dataset) throws Throwable { + return bbqIndex(resources, new BbqQuantizer[] {quantizer.quantizer()}, dataset); + } + + private static CagraIndex bbqIndex( + CuVSResources resources, BbqQuantizer[] quantizers, CuVSMatrix dataset) throws Throwable { + var params = + new CagraIndexParams.Builder() + .withMetric(CuvsDistanceType.L2Expanded) + .withGraphDegree(32) + .withIntermediateGraphDegree(64) + .build(); + var builder = + CagraIndex.newBuilder(resources).withIndexParams(params).withBbqDataset(quantizers); + if (dataset != null) { + builder = builder.withDataset(dataset); + } + return builder.build(); + } + + private float[][] randomVectors(int count) { + float[][] vectors = new float[count][DIM]; + for (float[] vector : vectors) { + for (int d = 0; d < DIM; d++) { + vector[d] = random.nextFloat() * 100.0f; + } + } + return vectors; + } + + /** + * Vectors drawn around a set of cluster centers. + * + *

Uniformly random high-dimensional vectors are close to equidistant from one another, so a + * graph index has nothing to exploit and recall says more about the data than the + * implementation. Clustered data is both more representative and a meaningful thing to assert on. + */ + private float[][] clusteredVectors(float[][] centers, int count) { + float[][] vectors = new float[count][DIM]; + for (float[] vector : vectors) { + float[] center = centers[random.nextInt(centers.length)]; + for (int d = 0; d < DIM; d++) { + vector[d] = center[d] + (float) (random.nextGaussian() * CLUSTER_SPREAD); + } + } + return vectors; + } + + private CuVSMatrix newDenseDataset(CuVSResources resources) { + var builder = CuVSMatrix.deviceBuilder(resources, ROWS, DIM, CuVSMatrix.DataType.FLOAT); + for (float[] row : dataset) { + builder.addVector(row); + } + return builder.build(); + } + + /** A quantizer plus the component matrices it was built from, closed together. */ + private record Quantizer(BbqQuantizer quantizer, List matrices) + implements AutoCloseable { + @Override + public void close() { + matrices.forEach(CuVSMatrix::close); + } + } + + /** + * Encodes {@link #dataset} as PACKED_1B BBQ codes and uploads every component to device memory. + * + *

This mirrors {@code cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp} closely + * enough for the codes to describe the real data, which is what makes the end-to-end searches + * below meaningful. It deliberately skips that file's {@code optimize_intervals} MSE refinement + * and keeps the initial grid estimate: the graph only has to be good, not optimal, because both + * searches score candidates against the dense dataset rather than against the codes. + */ + private Quantizer newQuantizer(CuVSResources resources) { + float[] centroid = new float[DIM]; + for (float[] row : dataset) { + for (int d = 0; d < DIM; d++) { + centroid[d] += row[d]; + } + } + float centroidNormSq = 0.0f; + for (int d = 0; d < DIM; d++) { + centroid[d] /= ROWS; + centroidNormSq += centroid[d] * centroid[d]; + } + + byte[][] codes = new byte[ROWS][CODE_BYTES]; + float[] lower = new float[ROWS]; + float[] upper = new float[ROWS]; + float[] corrections = new float[ROWS]; + int[] sums = new int[ROWS]; + float[] rowNorms = new float[ROWS]; + float[] delta = new float[ROWS]; + float[] sumDelta = new float[ROWS]; + + for (int i = 0; i < ROWS; i++) { + float[] centred = new float[DIM]; + float origNormSq = 0.0f; + float min = Float.MAX_VALUE; + float max = -Float.MAX_VALUE; + float centredNormSq = 0.0f; + double mean = 0.0; + double var = 0.0; + for (int d = 0; d < DIM; d++) { + origNormSq += dataset[i][d] * dataset[i][d]; + centred[d] = dataset[i][d] - centroid[d]; + min = Math.min(min, centred[d]); + max = Math.max(max, centred[d]); + centredNormSq += centred[d] * centred[d]; + double diff = centred[d] - mean; + mean += diff / (d + 1); + var += diff * (centred[d] - mean); + } + double stddev = Math.sqrt(var / DIM); + + // kMinimumMseGrid[0] from the C++ reference, i.e. the 1-bit row. + float a = (float) clamp(-0.798 * stddev + mean, min, max); + float b = (float) clamp(0.798 * stddev + mean, min, max); + // One step for 1-bit codes, so each component lands on 0 or 1. + float step = b - a; + int sum = 0; + for (int d = 0; d < DIM; d++) { + int code = step == 0.0f ? 0 : Math.round((float) (clamp(centred[d], a, b) - a) / step); + sum += code; + codes[i][d / 8] |= (byte) ((code & 1) << (7 - (d % 8))); + } + + lower[i] = a; + upper[i] = b; + // L2Expanded is euclidean, so the correction carries the centred norm. + corrections[i] = centredNormSq; + sums[i] = sum; + rowNorms[i] = origNormSq; + delta[i] = b - a; + sumDelta[i] = delta[i] * sum; + } + + List owned = new ArrayList<>(); + var codesBuilder = + CuVSMatrix.deviceBuilder(resources, ROWS, CODE_BYTES, CuVSMatrix.DataType.BYTE); + for (byte[] row : codes) { + codesBuilder.addVector(row); + } + + BbqQuantizer quantizer = + new BbqQuantizer.Builder() + .withCodes(track(owned, codesBuilder.build())) + .withLowerIntervals(track(owned, deviceVector(resources, lower))) + .withUpperIntervals(track(owned, deviceVector(resources, upper))) + .withAdditionalCorrections(track(owned, deviceVector(resources, corrections))) + .withQuantizedComponentSums(track(owned, deviceVector(resources, sums))) + .withCentroid(track(owned, deviceRow(resources, centroid))) + .withDequantDelta(track(owned, deviceVector(resources, delta))) + .withDequantSumDelta(track(owned, deviceVector(resources, sumDelta))) + .withRowNorm(track(owned, deviceVector(resources, rowNorms))) + .withLayout(BbqQuantizer.CodeLayout.PACKED_1B) + .withMetric(CuvsDistanceType.L2Expanded) + .withCentroidNormSq(centroidNormSq) + .build(); + return new Quantizer(quantizer, owned); + } + + private static double clamp(double value, double min, double max) { + return Math.min(Math.max(value, min), max); + } + + private static CuVSMatrix track(List owned, CuVSMatrix matrix) { + owned.add(matrix); + return matrix; + } + + /** One value per row, which the BBQ view reads as a length-ROWS vector. */ + private static CuVSMatrix deviceVector(CuVSResources resources, float[] values) { + var builder = CuVSMatrix.deviceBuilder(resources, values.length, 1, CuVSMatrix.DataType.FLOAT); + for (float value : values) { + builder.addVector(new float[] {value}); + } + return builder.build(); + } + + private static CuVSMatrix deviceVector(CuVSResources resources, int[] values) { + var builder = CuVSMatrix.deviceBuilder(resources, values.length, 1, CuVSMatrix.DataType.INT); + for (int value : values) { + builder.addVector(new int[] {value}); + } + return builder.build(); + } + + private static CuVSMatrix deviceRow(CuVSResources resources, float[] row) { + var builder = CuVSMatrix.deviceBuilder(resources, 1, row.length, CuVSMatrix.DataType.FLOAT); + builder.addVector(row); + return builder.build(); + } +}