diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 3be6cf1fa1..5d2cece12a 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -308,6 +308,15 @@ using standard_dataset_owning_storage = template using standard_dataset_view_storage = dense_row_major_dataset_view_storage; +template +using dense_owning_vector = std::conditional_t, + raft::host_vector>; + +template +using dense_view_vector = std::conditional_t, + raft::host_vector_view>; // ----------------------------------------------------------------------------- // VPQ compressed // ----------------------------------------------------------------------------- @@ -776,6 +785,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/preprocessing/quantize/bbq.hpp b/cpp/include/cuvs/preprocessing/quantize/bbq.hpp new file mode 100644 index 0000000000..3c06ed8b8e --- /dev/null +++ b/cpp/include/cuvs/preprocessing/quantize/bbq.hpp @@ -0,0 +1,317 @@ +/* + * 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 + +namespace CUVS_EXPORT cuvs { + +namespace preprocessing::quantize::bbq { + +/** + * @defgroup bbq Better Binary Quantization utilities + * @{ + */ + +/** + * Layout of BBQ quantized codes in each dataset row. + */ +enum class bbq_code_layout { + single_bit, /** Each dimension is quantized to a single bit and packed into bytes. Reflects + * OptimizedScalarQuantizer.packAsBinary. */ + dibit, /** Each dimension is quantized to 2 bits (dibit) and transposed for bitwise operations. + * Same principle as transpose_half_byte, but for 2 bits. + * Reflects OptimizedScalarQuantizer.transposeDibit. */ + transpose_half_byte, /** Each dimension is quantized to 4 bits, optimized for bitwise operations. + * Reflects 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_nibble, /** Each dimension is quantized to 4 bits, two values are packed into each output + * byte. Reflects OffHeapScalarQuantizedVectorValues.packNibbles. */ + seven_bit, /** Each dimension is quantized to 7 bits and treated as a signed value. */ + unsigned_byte, /** Each dimension is quantized to 8 bits and treated as an unsigned value. */ + +}; + +/** Owning structure for BBQ quantizer data. */ +template +struct bbq_quantizer { + template + using dense_owning_matrix = cuvs::neighbors::detail::dense_owning_matrix; + template + using dense_owning_vector = cuvs::neighbors::detail::dense_owning_vector; + dense_owning_matrix codes; + dense_owning_vector lower_intervals; + dense_owning_vector upper_intervals; + dense_owning_vector additional_corrections; + dense_owning_vector quantized_component_sums; + dense_owning_vector centroid; + + uint32_t bits{}; + bbq_code_layout layout{bbq_code_layout::single_bit}; + cuvs::distance::DistanceType metric{cuvs::distance::DistanceType::L2Expanded}; + float centroid_norm_sq{}; + + bbq_quantizer(dense_owning_matrix&& codes, + dense_owning_vector&& lower_intervals, + dense_owning_vector&& upper_intervals, + dense_owning_vector&& additional_corrections, + dense_owning_vector&& quantized_component_sums, + dense_owning_vector&& centroid, + uint32_t bits, + bbq_code_layout layout, + cuvs::distance::DistanceType metric, + float centroid_norm_sq) noexcept + : codes{std::move(codes)}, + lower_intervals{std::move(lower_intervals)}, + upper_intervals{std::move(upper_intervals)}, + additional_corrections{std::move(additional_corrections)}, + quantized_component_sums{std::move(quantized_component_sums)}, + centroid{std::move(centroid)}, + bits{bits}, + layout{layout}, + metric{metric}, + centroid_norm_sq{centroid_norm_sq} + { + const auto n_rows = static_cast(codes.extent(0)); + RAFT_EXPECTS(bits >= 1 && bits <= 8, "BBQ bits must be in [1, 8]."); + RAFT_EXPECTS(codes.extent(1) == static_cast(encoded_row_length()), + "BBQ code row length does not match dim, bits, and layout."); + RAFT_EXPECTS(lower_intervals.extent(0) == n_rows && upper_intervals.extent(0) == n_rows && + additional_corrections.extent(0) == n_rows && + quantized_component_sums.extent(0) == n_rows, + "Every BBQ correction array must contain one value per row."); + } + + [[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 + { + auto const d = dim(); + switch (layout) { + case bbq_code_layout::single_bit: return (d * bits + 7) / 8; + case bbq_code_layout::dibit: return bits * ((d + 7) / 8); + case bbq_code_layout::packed_nibble: return (d + 1) / 2; + case bbq_code_layout::seven_bit: return d; + case bbq_code_layout::unsigned_byte: return d; + case bbq_code_layout::transpose_half_byte: return 4 * ((d + 7) / 8); + } + return 0; + } +}; + +/** View structure for BBQ quantizer data. */ +template +struct bbq_quantizer_view { + using owning_accessor = + cuvs::neighbors::detail::dataset_owning_accessor_for_view; + using owning_storage = bbq_quantizer; + template + using dense_view_matrix = cuvs::neighbors::detail::dense_view_matrix; + template + using dense_view_vector = cuvs::neighbors::detail::dense_view_vector; + dense_view_matrix codes; + dense_view_vector lower_intervals; + dense_view_vector upper_intervals; + dense_view_vector additional_corrections; + dense_view_vector quantized_component_sums; + dense_view_vector centroid; + + uint32_t bits{}; + bbq_code_layout layout{bbq_code_layout::single_bit}; + cuvs::distance::DistanceType metric{cuvs::distance::DistanceType::L2Expanded}; + float centroid_norm_sq{}; + + bbq_quantizer_view(const owning_storage& quantizer) noexcept + : codes{quantizer.codes.view()}, + lower_intervals{quantizer.lower_intervals.view()}, + upper_intervals{quantizer.upper_intervals.view()}, + additional_corrections{quantizer.additional_corrections.view()}, + quantized_component_sums{quantizer.quantized_component_sums.view()}, + centroid{quantizer.centroid.view()}, + bits{quantizer.bits}, + layout{quantizer.layout}, + metric{quantizer.metric}, + centroid_norm_sq{quantizer.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)); + } +}; + +/** @} */ // end of bbq group + +} // namespace preprocessing::quantize::bbq + +namespace neighbors { +struct bbq_dataset_container { + template + using owning_storage = cuvs::preprocessing::quantize::bbq::bbq_quantizer; + template + using view_storage = + cuvs::preprocessing::quantize::bbq::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_bit_and_layout(quantizer.bits, quantizer.layout), + "Quantizer already exists with bits and layout."); + this->quantizers.push_back(std::move(quantizer)); + } + bool has_bit_and_layout(uint32_t bits, + cuvs::preprocessing::quantize::bbq::bbq_code_layout layout) const noexcept + { + for (uint32_t i = 0; i < quantizers.size(); i++) { + if (quantizers[i].bits == bits && 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(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_bit_and_layout(quantizer.bits, quantizer.layout), + "Quantizer already exists with bits and layout."); + this->quantizers.push_back(quantizer); + } + void add_quantizer(const owning_storage_type& quantizer) + { + RAFT_EXPECTS(!has_bit_and_layout(quantizer.bits, quantizer.layout), + "Quantizer already exists with bits and layout."); + this->quantizers.push_back(view_storage_type(quantizer)); + } + bool has_bit_and_layout(uint32_t bits, + cuvs::preprocessing::quantize::bbq::bbq_code_layout layout) const noexcept + { + for (uint32_t i = 0; i < quantizers.size(); i++) { + if (quantizers[i].bits == bits && quantizers[i].layout == layout) { return true; } + } + return false; + } + view_storage_type get_quantizer(uint32_t bits, + cuvs::preprocessing::quantize::bbq::bbq_code_layout layout) const + { + for (uint32_t i = 0; i < quantizers.size(); i++) { + if (quantizers[i].bits == bits && quantizers[i].layout == layout) { return quantizers[i]; } + } + throw std::runtime_error("No quantizer found with bits and layout."); + } +}; + +template +using device_bbq_dataset = + dataset>; + +template +using device_bbq_dataset_view = + dataset_view>; + +template +using host_bbq_dataset = + dataset>; + +template +using host_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 +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/fern/docs.yml b/fern/docs.yml index 2aafaea69e..0256a9db87 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -390,6 +390,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" 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..ad26a9422d --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md @@ -0,0 +1,108 @@ +--- +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 + +Layout of BBQ quantized codes in each dataset row. + +```cpp +enum class bbq_code_layout { + single_bit, + dibit, + transpose_half_byte, + packed_nibble, + seven_bit, + unsigned_byte +}; +``` + +**Values** + +| Name | Value | +| --- | --- | +| `single_bit` | `` | +| `dibit` | `` | +| `transpose_half_byte` | `` | +| `packed_nibble` | `` | +| `seven_bit` | `` | +| `unsigned_byte` | `` | + + +### preprocessing::quantize::bbq::bbq_quantizer + +Owning structure for BBQ quantizer data. + +```cpp +template +struct bbq_quantizer { + dense_owning_matrix codes; + dense_owning_vector lower_intervals; + dense_owning_vector upper_intervals; + dense_owning_vector additional_corrections; + dense_owning_vector quantized_component_sums; + dense_owning_vector centroid; + uint32_t bits; + bbq_code_layout layout; + cuvs::distance::DistanceType metric; + float centroid_norm_sq; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `codes` | `dense_owning_matrix` | | +| `lower_intervals` | `dense_owning_vector` | | +| `upper_intervals` | `dense_owning_vector` | | +| `additional_corrections` | `dense_owning_vector` | | +| `quantized_component_sums` | `dense_owning_vector` | | +| `centroid` | `dense_owning_vector` | | +| `bits` | `uint32_t` | | +| `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::bbq_quantizer_view + +View structure for BBQ quantizer data. + +```cpp +template +struct bbq_quantizer_view { + dense_view_matrix codes; + dense_view_vector lower_intervals; + dense_view_vector upper_intervals; + dense_view_vector additional_corrections; + dense_view_vector quantized_component_sums; + dense_view_vector centroid; + uint32_t bits; + bbq_code_layout layout; + cuvs::distance::DistanceType metric; + float centroid_norm_sq; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `codes` | `dense_view_matrix` | | +| `lower_intervals` | `dense_view_vector` | | +| `upper_intervals` | `dense_view_vector` | | +| `additional_corrections` | `dense_view_vector` | | +| `quantized_component_sums` | `dense_view_vector` | | +| `centroid` | `dense_view_vector` | | +| `bits` | `uint32_t` | | +| `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` | | diff --git a/fern/pages/cpp_api/index.md b/fern/pages/cpp_api/index.md index d5904b5db4..8a348ce0f1 100644 --- a/fern/pages/cpp_api/index.md +++ b/fern/pages/cpp_api/index.md @@ -41,6 +41,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)