diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 08009b731d..3108574963 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1476,6 +1476,8 @@ if(NOT BUILD_CPU_ONLY) src/neighbors/composite/index.cu $<$:src/neighbors/cagra.cpp> $<$:src/neighbors/hnsw.cpp> + src/neighbors/ivf_binary_index.cpp + src/neighbors/ivf_binary/ivf_binary.cu src/neighbors/ivf_common.cu src/neighbors/ivf_flat_index.cpp ${ivf_flat_build_extend_inst_files} diff --git a/cpp/include/cuvs/cluster/kmeans.hpp b/cpp/include/cuvs/cluster/kmeans.hpp index a805aff0a4..6ac1b881f6 100644 --- a/cpp/include/cuvs/cluster/kmeans.hpp +++ b/cpp/include/cuvs/cluster/kmeans.hpp @@ -216,6 +216,19 @@ struct balanced_params : base_params { * average cluster size; in that mode, `balance_upper_tolerance` does not control donor selection. */ balanced_donor_selection donor_selection = balanced_donor_selection::SizeSorted; + + /** + * If true, treats uint8_t input data as bit-packed binary data where each byte contains 8 bits. + * Bits are expanded on-the-fly, least-significant bit first, to {-1, +1} floats + * during training and prediction. Other input types are rejected when this flag is set. + * When enabled: + * - Input data dimension represents packed dimension (actual_dim / 8) + * - Output centroids dimension is expanded (packed_dim * 8) + * - The metric operates on the expanded floating-point vectors (for example L2Expanded), + * not on the packed bytes; BitwiseHamming is not a balanced k-means training metric. + * - CosineExpanded is not supported. + */ + bool is_packed_binary = false; }; /** @@ -717,6 +730,12 @@ void fit(const raft::resources& handle, /** * @brief Find balanced clusters with k-means algorithm. * + * @note When `params.is_packed_binary` is true, `X.extent(1)` counts packed bytes, + * and centroids must have `8 * X.extent(1)` floating-point coordinates. Bits are + * expanded least-significant bit first to {-1, +1}; the selected metric operates + * on those expanded vectors. CosineExpanded is not supported in packed binary mode. + * With the flag disabled, uint8_t values are numeric. + * * @code{.cpp} * #include * #include @@ -1215,6 +1234,12 @@ void predict(const raft::resources& handle, /** * @brief Predict the closest cluster each sample in X belongs to. * + * @note When `params.is_packed_binary` is true, `X.extent(1)` counts packed bytes, + * and centroids must have `8 * X.extent(1)` floating-point coordinates. Bits are + * expanded least-significant bit first to {-1, +1}; the selected metric operates + * on those expanded vectors. CosineExpanded is not supported in packed binary mode. + * With the flag disabled, uint8_t values are numeric. + * * @code{.cpp} * #include * #include diff --git a/cpp/include/cuvs/core/dataset/binary.hpp b/cpp/include/cuvs/core/dataset/binary.hpp new file mode 100644 index 0000000000..4a04072977 --- /dev/null +++ b/cpp/include/cuvs/core/dataset/binary.hpp @@ -0,0 +1,98 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include +#include + +#include + +namespace CUVS_EXPORT cuvs { + +/** + * @brief Non-owning view of the dictionary that binarizes a dataset: the per-dimension thresholds + * of a binary quantizer. An empty threshold vector stands for the zero threshold. + */ +template +struct binary_dictionary_view { + raft::device_vector_view threshold; +}; + +/** + * @brief Non-owning view of a binary-quantized dataset. + * + * The data are bit-packed codes, `ceil(dim / 8)` bytes per row; bit `i` of a row is stored in bit + * `i % 8` of byte `i / 8`, and the unused bits of the last byte are zero. The dictionary is the + * quantizer that produced the codes, so the logical dimension `dim` (the number of bits) never has + * to be guessed from the padded byte width. + * + * @tparam MathT the type of the source (unquantized) values + * @tparam IdxT the type of the row indices + */ +template +struct device_binary_dataset_view { + using value_type = MathT; + using index_type = IdxT; + using code_type = uint8_t; + using data_view_type = raft::device_matrix_view; + using dictionary_view_type = binary_dictionary_view; + + /** + * @param codes row-major codes [n_rows, ceil(dim / 8)] + * @param dictionary the thresholds used to produce the codes; empty or of size `dim` + * @param dim the number of bits in each vector + */ + device_binary_dataset_view(data_view_type codes, dictionary_view_type dictionary, uint32_t dim) + : codes_{codes}, dictionary_{dictionary}, dim_{dim} + { + RAFT_EXPECTS(dim > 0, "The dimension of a binary dataset must be positive"); + RAFT_EXPECTS(static_cast((uint64_t{dim} + 7) / 8) == codes.extent(1), + "The row length of the codes must be ceil(dim / 8) bytes"); + RAFT_EXPECTS(dictionary.threshold.extent(0) == 0 || + dictionary.threshold.extent(0) == static_cast(dim), + "The threshold vector must be empty or have one element per dimension"); + } + + /** The number of vectors. */ + [[nodiscard]] auto n_rows() const noexcept -> index_type + { + return static_cast(codes_.extent(0)); + } + /** The number of bits in each vector. */ + [[nodiscard]] auto dim() const noexcept -> uint32_t { return dim_; } + /** The number of bytes in each row of the codes. */ + [[nodiscard]] auto encoded_row_length() const noexcept -> int64_t { return codes_.extent(1); } + + [[nodiscard]] auto data_const_view() const noexcept -> data_view_type { return codes_; } + [[nodiscard]] auto const_dictionary_view() const noexcept -> dictionary_view_type + { + return dictionary_; + } + + private: + data_view_type codes_; + dictionary_view_type dictionary_; + uint32_t dim_; +}; + +/** + * @brief View the codes produced by `preprocessing::quantize::binary::transform` together with the + * quantizer that produced them. + */ +template +auto make_device_binary_dataset_view( + raft::device_matrix_view codes, + const preprocessing::quantize::binary::quantizer& quantizer, + uint32_t dim) -> device_binary_dataset_view +{ + return device_binary_dataset_view{ + codes, binary_dictionary_view{quantizer.threshold.view()}, dim}; +} + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp index 56180e3434..4f7ec74b15 100644 --- a/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp @@ -12,6 +12,7 @@ struct tag_h {}; struct tag_d {}; struct tag_i8 {}; struct tag_u8 {}; +struct tag_u32 {}; struct tag_filter_none {}; struct tag_filter_bitset {}; struct tag_filter_bloom_filter {}; diff --git a/cpp/include/cuvs/detail/jit_lto/ivf_flat/interleaved_scan_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/ivf_flat/interleaved_scan_fragments.hpp index 6f65837741..4c3230c920 100644 --- a/cpp/include/cuvs/detail/jit_lto/ivf_flat/interleaved_scan_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/ivf_flat/interleaved_scan_fragments.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -16,6 +16,7 @@ struct tag_acc_u32 {}; // Tag types for distance metrics struct tag_metric_euclidean {}; struct tag_metric_inner_product {}; +struct tag_metric_bitwise_hamming {}; struct tag_metric_custom_udf {}; // Tag types for post-processing diff --git a/cpp/include/cuvs/detail/jit_lto/pairwise_matrix/pairwise_matrix_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/pairwise_matrix/pairwise_matrix_fragments.hpp index d9b7f4ec8d..acd431e61b 100644 --- a/cpp/include/cuvs/detail/jit_lto/pairwise_matrix/pairwise_matrix_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/pairwise_matrix/pairwise_matrix_fragments.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -13,6 +13,7 @@ struct tag_layout_col {}; struct tag_fin_op_identity {}; struct tag_fin_op_rbf {}; +struct tag_distance_bitwise_hamming {}; struct tag_distance_canberra {}; struct tag_distance_correlation {}; struct tag_distance_cosine {}; diff --git a/cpp/include/cuvs/neighbors/ivf_binary.hpp b/cpp/include/cuvs/neighbors/ivf_binary.hpp new file mode 100644 index 0000000000..9f9857b115 --- /dev/null +++ b/cpp/include/cuvs/neighbors/ivf_binary.hpp @@ -0,0 +1,246 @@ +/* + * 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 neighbors { +namespace ivf_binary { + +namespace detail { +struct index_access; +} // namespace detail + +/** + * @defgroup ivf_binary_cpp_index_params IVF-Binary index build parameters + * @{ + */ +struct index_params { + /** The number of inverted lists (clusters). */ + uint32_t n_lists = 1024; + /** The number of iterations searching for k-means centers. */ + uint32_t kmeans_n_iters = 20; + /** The fraction of data to use during iterative k-means building. */ + double kmeans_trainset_fraction = 0.5; + /** + * By default (adaptive_centers = false), the cluster centers are trained in `build` and never + * modified in `extend`. + * + * With adaptive_centers = true, `extend` updates the centers to the majority bit of the vectors + * added to each cluster, with ties resolved to 0. The index keeps exact per-bit counts of ones to + * make the update independent of how the data are split between `extend` calls. + */ + bool adaptive_centers = false; + /** + * Allocate the minimum amount of memory required for the inverted lists, instead of reserving + * extra space to amortize repeated calls of `extend`. + */ + bool conservative_memory_allocation = false; + /** + * Whether to add the dataset content to the index in `build`. If false, `build` only trains the + * clustering and the index is populated later with `extend`. + */ + bool add_data_on_build = true; +}; +/** @} */ + +/** + * @defgroup ivf_binary_cpp_search_params IVF-Binary index search parameters + * @{ + */ +struct search_params : cuvs::neighbors::search_params { + /** The number of clusters to search. */ + uint32_t n_probes = 20; +}; +/** @} */ + +/** + * @defgroup ivf_binary_cpp_index IVF-Binary index + * @{ + */ +/** + * @brief IVF index of bit-packed binary vectors searched with the Hamming distance. + * + * The index stores the codes of a binary-quantized dataset (see `device_binary_dataset_view`), + * together with the quantizer that produced them. Queries must be supplied as bit-packed bytes. + * Neighbors are returned with their Hamming distances, i.e. the number of differing bits. + */ +struct index : cuvs::neighbors::index { + using index_params_type = ivf_binary::index_params; + using search_params_type = ivf_binary::search_params; + using index_type = int64_t; + + public: + index(const index&) = delete; + index(index&&) = default; + index& operator=(const index&) = delete; + index& operator=(index&&) = default; + ~index() = default; + + /** Construct an empty index. It needs to be trained with `build` or loaded with `deserialize`. */ + explicit index(raft::resources const& res); + + /** Construct an empty index of `dim`-bit vectors. It needs to be trained and then populated. */ + index(raft::resources const& res, const index_params& params, uint32_t dim); + + /** The number of bits in each vector. */ + uint32_t dim() const noexcept; + /** The number of bytes in each encoded vector. */ + uint32_t encoded_dim() const noexcept; + /** The number of clusters/inverted lists. */ + uint32_t n_lists() const noexcept; + /** The total number of vectors in the index. */ + int64_t size() const noexcept; + /** Whether `centers()` change upon extending the index. */ + bool adaptive_centers() const noexcept; + /** Whether the lists are allocated without extra space (see index_params). */ + bool conservative_memory_allocation() const noexcept; + + /** Sizes of the lists (clusters) [n_lists]. */ + raft::device_vector_view list_sizes() const noexcept; + /** Bit-packed cluster centers [n_lists, encoded_dim]. */ + raft::device_matrix_view centers() const noexcept; + /** + * Exact number of ones at each bit of the vectors in each cluster [n_lists, 8 * encoded_dim]; + * empty unless the centers are adaptive. + */ + raft::device_matrix_view center_counts() const noexcept; + /** The build dataset's quantizer, available for external encoding of queries or new vectors. */ + const preprocessing::quantize::binary::quantizer& quantizer() const noexcept; + + private: + friend struct detail::index_access; + + uint32_t dim_; + bool adaptive_centers_; + preprocessing::quantize::binary::quantizer quantizer_; + raft::device_matrix centers_; + raft::device_matrix center_counts_; + // The inverted lists of codes, scanned by the interleaved Hamming kernel of IVF-Flat. + ivf_flat::index lists_; +}; +/** @} */ + +/** + * @defgroup ivf_binary_cpp_index_build IVF-Binary index build + * @{ + */ +/** + * @brief Build the index from a binary-quantized dataset. + * + * The thresholds of the dataset dictionary are copied into the index. + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * namespace binary = cuvs::preprocessing::quantize::binary; + * auto quantizer = binary::train(handle, binary::params{}, raw_dataset); + * auto codes = raft::make_device_matrix(handle, n_rows, (dim + 7) / 8); + * binary::transform(handle, quantizer, raw_dataset, codes.view()); + * auto dataset = cuvs::make_device_binary_dataset_view( + * raft::make_const_mdspan(codes.view()), quantizer, dim); + * auto index = ivf_binary::build(handle, ivf_binary::index_params{}, dataset); + * @endcode + * + * @param[in] handle + * @param[in] index_params configure the index building + * @param[in] dataset the binary-quantized vectors + * + * @return the constructed index + */ +auto build(raft::resources const& handle, + const index_params& index_params, + device_binary_dataset_view dataset) -> index; + +/** + * @brief Add bit-packed vectors to the index. + * + * The supplied codes are used as-is. The dictionary of the new vectors is not compared with or + * copied into the index; choosing the quantizer is the caller's responsibility. + * + * @param[in] handle + * @param[in] new_vectors the binary-quantized vectors to add + * @param[in] new_indices ids of the new vectors; required if the index is not empty, otherwise the + * ids start from zero + * @param[inout] index the index to extend + */ +void extend(raft::resources const& handle, + device_binary_dataset_view new_vectors, + std::optional> new_indices, + index* index); +/** @} */ + +/** + * @defgroup ivf_binary_cpp_index_search IVF-Binary index search + * @{ + */ +/** + * @brief Search the index with bit-packed byte queries using Hamming distance. + * + * Each query has `index.encoded_dim()` bytes, with bit `i` stored in bit `i % 8` of byte `i / 8`. + * Unused bits in the final byte must be zero. The supplied codes are searched as-is, without + * quantization or threshold comparisons; choosing the quantizer is the caller's responsibility. + * + * @param[in] handle + * @param[in] params configure the search + * @param[in] index the index + * @param[in] queries a device matrix [n_queries, index.encoded_dim()] of bit-packed queries + * @param[out] neighbors ids of the k nearest neighbors [n_queries, k] + * @param[out] distances Hamming distances to the neighbors [n_queries, k] + * @param[in] sample_filter an optional device filter on the dataset ids + */ +void search(raft::resources const& handle, + const search_params& params, + const index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); +/** @} */ + +/** + * @defgroup ivf_binary_cpp_serialize IVF-Binary index serialize + * @{ + */ +/** + * Save the index to an output stream. + * + * Experimental, both the API and the serialization format are subject to change. + */ +void serialize(raft::resources const& handle, std::ostream& os, const index& index); + +/** Save the index to a file. */ +void serialize(raft::resources const& handle, const std::string& filename, const index& index); + +/** + * Load the index from an input stream. + * + * Experimental, both the API and the serialization format are subject to change. + */ +void deserialize(raft::resources const& handle, std::istream& is, index* index); + +/** Load the index from a file. */ +void deserialize(raft::resources const& handle, const std::string& filename, index* index); +/** @} */ + +} // namespace ivf_binary +} // namespace neighbors +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index bb96c9115b..321acc9895 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -10,6 +10,7 @@ #include "../../core/nvtx.hpp" #include "../../distance/distance.cuh" +#include "../../distance/fused_distance_nn.cuh" #include #include @@ -38,9 +39,11 @@ #include #include +#include #include #include +#include "../../neighbors/detail/ann_utils.cuh" #include #include #include @@ -51,6 +54,47 @@ namespace cuvs::cluster::kmeans::detail { +/** Validate the input type and calculate the width used by the floating-point centers. */ +template +IdxT centers_dim(IdxT dim, bool is_packed_binary) +{ + RAFT_EXPECTS(dim > 0, "The number of features must be strictly positive"); + if (!is_packed_binary) { return dim; } + if constexpr (std::is_same_v) { + RAFT_EXPECTS(dim <= std::numeric_limits::max() / 8, + "The chosen index type cannot represent the expanded binary dimension"); + return dim * 8; + } else { + RAFT_FAIL("Packed binary mode is only supported for uint8_t data type"); + } +} + +inline void validate_packed_binary_metric(const cuvs::cluster::kmeans::balanced_params& params) +{ + RAFT_EXPECTS( + !params.is_packed_binary || params.metric != cuvs::distance::DistanceType::CosineExpanded, + "CosineExpanded is not supported for packed binary input"); +} + +/** + * @brief Create a transform iterator for on-the-fly bit expansion + * + * This helper function creates a thrust transform iterator that expands packed + * uint8_t data into float values on-the-fly (bit 1 → +1.0f, bit 0 → -1.0f), + * + * @tparam IdxT index type + * + * @param packed_data Pointer to row-major packed uint8_t data + * @return A transform iterator that yields float values for each bit + */ +template +auto make_bitwise_expanded_iterator(const uint8_t* packed_data) +{ + auto counting_iter = thrust::make_counting_iterator(0); + auto decoder = cuvs::spatial::knn::detail::utils::bitwise_decode_op(packed_data); + return thrust::make_transform_iterator(counting_iter, decoder); +} + /** * @brief Predict labels for the dataset; floating-point types only. * @@ -150,6 +194,86 @@ inline std::enable_if_t> predict_core( } } +/** + * @brief Predict labels for the dataset; uint8_t only (specialization for BitwiseHamming). + */ +template +inline void predict_bitwise_hamming(const raft::resources& handle, + const cuvs::cluster::kmeans::balanced_params& params, + const uint8_t* centers, + IdxT n_clusters, + IdxT dim, + const uint8_t* dataset, + const uint8_t* dataset_norm, + IdxT n_rows, + LabelT* labels, + rmm::device_async_resource_ref mr) +{ + RAFT_EXPECTS(params.metric == cuvs::distance::DistanceType::BitwiseHamming, + "uint8_t data only supports BitwiseHamming distance"); + + RAFT_EXPECTS(n_clusters > 0 && dim > 0, "Centers must have nonzero dimensions"); + if (n_rows == 0) { return; } + + auto workspace = raft::make_device_mdarray( + handle, mr, raft::make_extents((sizeof(int)) * n_rows)); + + auto minClusterAndDistance = raft::make_device_mdarray, IdxT>( + handle, mr, raft::make_extents(n_rows)); + raft::KeyValuePair initial_value(0, std::numeric_limits::max()); + raft::matrix::fill(handle, minClusterAndDistance.view(), initial_value); + + cuvs::distance::fusedDistanceNNMinReduce, IdxT>( + handle, + minClusterAndDistance.data_handle(), + dataset, + centers, + nullptr, + nullptr, + n_rows, + n_clusters, + dim, + (void*)workspace.data_handle(), + false, + false, + true, + params.metric, + 0.0f); + + raft::linalg::map(handle, + raft::make_const_mdspan(minClusterAndDistance.view()), + raft::make_device_vector_view(labels, n_rows), + raft::compose_op, raft::key_op>()); +} + +template +inline void predict_bitwise_hamming(const raft::resources& handle, + raft::device_matrix_view dataset, + raft::device_matrix_view centers, + raft::device_vector_view labels) +{ + RAFT_EXPECTS(dataset.extent(1) == centers.extent(1), + "Number of features in dataset and centroids are different"); + RAFT_EXPECTS(dataset.extent(0) == labels.extent(0), + "Number of rows in dataset and labels are different"); + RAFT_EXPECTS(static_cast(centers.extent(0)) <= + static_cast(std::numeric_limits::max()), + "The chosen label type cannot represent all cluster labels"); + cuvs::cluster::kmeans::balanced_params params; + params.metric = cuvs::distance::DistanceType::BitwiseHamming; + + predict_bitwise_hamming(handle, + params, + centers.data_handle(), + centers.extent(0), + centers.extent(1), + dataset.data_handle(), + nullptr, + dataset.extent(0), + labels.data_handle(), + raft::resource::get_workspace_resource_ref(handle)); +} + /** * @brief Suggest a minibatch size for kmeans prediction. * @@ -195,25 +319,30 @@ auto calc_minibatch_size(const raft::resources& handle, } // If we need to convert to MathT, space required for the converted batch. - if (!needs_conversion) { mem_per_row += sizeof(MathT) * dim; } + if (needs_conversion) { mem_per_row += sizeof(MathT) * size_t(dim); } - // Heuristic: calculate the minibatch size in order to use at most 80% or 512MB workspace memory. - // We go below 1GB here as the allocation is mostly done in a single chunk which - // is problematic if e.g. a pool allocator manages its own chunks <= 1GB. + // Include row norms in the workspace estimate. Cap large contiguous allocations at + // 80% of free workspace or 512 MiB, preserving upstream's pool-friendly allocation policy. + mem_per_row += sizeof(MathT); const auto free_ws_size = raft::resource::get_workspace_free_bytes(handle); const auto available_ws_size = - std::min((free_ws_size * size_t{8}) / size_t{10}, size_t{1} << 29); - - IdxT minibatch_size = std::max(IdxT{1}, static_cast(available_ws_size / mem_per_row)); - - minibatch_size = raft::round_down_safe(minibatch_size, IdxT{64}); - minibatch_size = std::min(minibatch_size, n_rows); - return std::make_tuple(minibatch_size, mem_per_row); + std::min((free_ws_size / size_t{10}) * size_t{8}, size_t{1} << 29); + size_t minibatch_size = std::max(1, available_ws_size / mem_per_row); + // Always make progress, including when fewer than 64 rows fit the budget. + if (minibatch_size >= 64) { minibatch_size = (minibatch_size / 64) * 64; } + IdxT batch_size = static_cast(std::min(minibatch_size, size_t(n_rows))); + return std::make_tuple(batch_size, mem_per_row); } /** * @brief Given the data and labels, calculate cluster centers and sizes in one sweep. * + * This function supports two modes: + * 1. Regular mode: Works with any data type T with optional type conversion via mapping_op + * 2. Packed binary mode: When T=uint8_t and is_packed_binary=true, treats data as bit-packed + * and expands bits on-the-fly (bit 1 → +1, bit 0 → -1) into float centers. + * In this mode, dim represents the packed dimension (dim_expanded / 8). + * * @note all pointers must be accessible on the device. * * @tparam T element type @@ -224,10 +353,10 @@ auto calc_minibatch_size(const raft::resources& handle, * @tparam MappingOpT type of the mapping operation * * @param[in] handle The raft handle. - * @param[inout] centers Pointer to the output [n_clusters, dim] + * @param[inout] centers Pointer to the output [n_clusters, dim] or [n_clusters, dim*8] if packed * @param[inout] cluster_sizes Number of rows in each cluster [n_clusters] * @param[in] n_clusters Number of clusters/centers - * @param[in] dim Dimensionality of the data + * @param[in] dim Dimensionality of the data (or packed dim if is_packed_binary=true) * @param[in] dataset Pointer to the data [n_rows, dim] * @param[in] n_rows Number of samples in the `dataset` * @param[in] labels Output predictions [n_rows] @@ -236,6 +365,8 @@ auto calc_minibatch_size(const raft::resources& handle, * the weighted average principle. * @param[in] mapping_op Mapping operation from T to MathT * @param[inout] mr (optional) Memory resource to use for temporary allocations on the device + * @param[in] is_packed_binary If true and T=uint8_t, treats data as bit-packed and expands + * on-the-fly */ template (centers, n_clusters, dim); + // For packed binary, dim is packed dimension, centers are in expanded dimension (dim * 8) + IdxT centers_dim = detail::centers_dim(dim, is_packed_binary); + + auto centersView = raft::make_device_matrix_view(centers, n_clusters, centers_dim); auto clusterSizesView = raft::make_device_vector_view(cluster_sizes, n_clusters); if (!reset_counters) { @@ -276,8 +411,26 @@ void calc_centers_and_sizes(const raft::resources& handle, temp_sizes = temp_cluster_sizes.data(); } + // Handle packed binary data with on-the-fly bit expansion + if (is_packed_binary) { + if constexpr (std::is_same_v) { + auto decoded_dataset_iter = make_bitwise_expanded_iterator(dataset); + raft::linalg::reduce_rows_by_key(decoded_dataset_iter, + centers_dim, + labels, + nullptr, + n_rows, + centers_dim, + n_clusters, + centers, + stream.get(), + reset_counters); + } else { + RAFT_FAIL("Packed binary mode is only supported for uint8_t data type"); + } + } // Apply mapping only when the data and math types are different. - if constexpr (std::is_same_v) { + else if constexpr (std::is_same_v) { raft::linalg::reduce_rows_by_key(dataset, dim, labels, @@ -533,18 +686,27 @@ void predict(const raft::resources& handle, auto stream = raft::resource::get_cuda_stream(handle); raft::common::nvtx::range fun_scope( "predict(%zu, %u)", static_cast(n_rows), n_clusters); - auto mem_res = mr.value_or(raft::resource::get_workspace_resource_ref(handle)); + auto mem_res = mr.value_or(raft::resource::get_workspace_resource_ref(handle)); + IdxT transformed_dim = centers_dim(dim, params.is_packed_binary); + if (n_rows == 0) { return; } auto [max_minibatch_size, _mem_per_row] = calc_minibatch_size( - handle, n_clusters, n_rows, dim, params.metric, std::is_same_v); + handle, n_clusters, n_rows, transformed_dim, params.metric, !std::is_same_v); rmm::device_uvector cur_dataset( - std::is_same_v ? 0 : max_minibatch_size * dim, stream, mem_res); + std::is_same_v ? 0 : max_minibatch_size * transformed_dim, stream, mem_res); constexpr bool native_half = std::is_same_v && std::is_same_v; - bool need_compute_norm = + bool need_norm = dataset_norm == nullptr && (params.metric == cuvs::distance::DistanceType::L2Expanded || params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || params.metric == cuvs::distance::DistanceType::CosineExpanded); + bool need_compute_norm = need_norm && !params.is_packed_binary; rmm::device_uvector cur_dataset_norm( - need_compute_norm || native_half ? max_minibatch_size : 0, stream, mem_res); + need_norm || native_half ? max_minibatch_size : 0, stream, mem_res); + if (need_norm && params.is_packed_binary) { + raft::matrix::fill( + handle, + raft::make_device_matrix_view(cur_dataset_norm.data(), max_minibatch_size, 1), + static_cast(transformed_dim)); + } const auto native_centers_size = native_half ? static_cast(n_clusters) * static_cast(dim) : 0; std::optional native_half_scratch; @@ -552,8 +714,9 @@ void predict(const raft::resources& handle, native_half_scratch.emplace( native_centers_size, n_clusters, max_minibatch_size, stream, mem_res); } - const MathT* dataset_norm_ptr = nullptr; - auto cur_dataset_ptr = cur_dataset.data(); + const MathT* dataset_norm_ptr = + need_norm && params.is_packed_binary ? cur_dataset_norm.data() : nullptr; + auto cur_dataset_ptr = cur_dataset.data(); for (IdxT offset = 0; offset < n_rows; offset += max_minibatch_size) { IdxT minibatch_size = std::min(max_minibatch_size, n_rows - offset); @@ -575,6 +738,14 @@ void predict(const raft::resources& handle, } if constexpr (std::is_same_v) { cur_dataset_ptr = const_cast(dataset + offset * dim); + } else if (params.is_packed_binary) { + if constexpr (std::is_same_v) { + raft::linalg::map_offset(handle, + raft::make_device_matrix_view( + cur_dataset_ptr, minibatch_size, transformed_dim), + cuvs::spatial::knn::detail::utils::bitwise_decode_op( + dataset + offset * dim)); + } } else { raft::linalg::map( handle, @@ -612,7 +783,7 @@ void predict(const raft::resources& handle, params, centers, n_clusters, - dim, + transformed_dim, cur_dataset_ptr, dataset_norm_ptr, minibatch_size, @@ -622,7 +793,7 @@ void predict(const raft::resources& handle, } template average_size * balance_upper_tolerance * balance_upper_tolerance > 1 * @param[in] centroid_offset offset from the donor cluster centroid towards a donor point - * @param[in] mapping_op Mapping operation from T to MathT + * @param[in] mapping_op Mapping operation from dataset values to MathT * @param[inout] device_memory memory resource to use for temporary allocations * * @return whether any of the centers has been updated (and thus, `labels` need to be recalculated). */ -template (dim, params.is_packed_binary); for (uint32_t iter = 0; iter < n_iters; iter++) { // Balancing step - move the centers around to equalize cluster sizes // (but not on the first iteration) - if (iter > 0 && adjust_centers(handle, - cluster_centers, - n_clusters, - dim, - dataset, - n_rows, - cluster_labels, - cluster_sizes, - balance_lower_tolerance, - balance_upper_tolerance, - static_cast(params.centroid_offset), - params.donor_selection, - mapping_op, - device_memory)) { + bool did_adjust = false; + if (iter > 0) { + auto adjust = [&](auto data, auto data_mapping) { + return adjust_centers(handle, + cluster_centers, + n_clusters, + transformed_dim, + data, + n_rows, + cluster_labels, + cluster_sizes, + balance_lower_tolerance, + balance_upper_tolerance, + static_cast(params.centroid_offset), + params.donor_selection, + data_mapping, + device_memory); + }; + if (params.is_packed_binary) { + if constexpr (std::is_same_v) { + did_adjust = + adjust(make_bitwise_expanded_iterator(dataset), raft::identity_op{}); + } + } else { + did_adjust = adjust(dataset, mapping_op); + } + } + if (did_adjust) { if (balancing_counter++ >= balancing_pullback) { balancing_counter -= balancing_pullback; n_iters++; @@ -993,9 +1179,9 @@ void balancing_em_iters(const raft::resources& handle, case cuvs::distance::DistanceType::CosineExpanded: case cuvs::distance::DistanceType::CorrelationExpanded: { auto clusters_in_view = raft::make_device_matrix_view( - cluster_centers, n_clusters, dim); + cluster_centers, n_clusters, transformed_dim); auto clusters_out_view = raft::make_device_matrix_view( - cluster_centers, n_clusters, dim); + cluster_centers, n_clusters, transformed_dim); raft::linalg::row_normalize( handle, clusters_in_view, clusters_out_view); break; @@ -1024,6 +1210,7 @@ void balancing_em_iters(const raft::resources& handle, n_rows, cluster_labels, true, + params.is_packed_binary, mapping_op, device_memory); } @@ -1066,6 +1253,7 @@ void build_clusters(const raft::resources& handle, n_rows, cluster_labels, true, + params.is_packed_binary, mapping_op, device_memory); @@ -1190,7 +1378,8 @@ auto build_fine_clusters(const raft::resources& handle, rmm::device_async_resource_ref managed_memory, rmm::device_async_resource_ref device_memory) -> IdxT { - auto stream = raft::resource::get_cuda_stream(handle); + auto stream = raft::resource::get_cuda_stream(handle); + IdxT transformed_dim = centers_dim(dim, params.is_packed_binary); rmm::device_uvector mc_trainset_ids_buf(mesocluster_size_max, stream, managed_memory); // for small cluster counts the maximum mesocluster size is proportional to the number of rows, so // we use large workspace @@ -1205,7 +1394,7 @@ auto build_fine_clusters(const raft::resources& handle, rmm::device_uvector mc_trainset_labels(mesocluster_size_max, stream, device_memory); rmm::device_uvector mc_trainset_ccenters( - fine_clusters_nums_max * dim, stream, device_memory); + fine_clusters_nums_max * transformed_dim, stream, device_memory); // number of vectors in each cluster rmm::device_uvector mc_trainset_csizes_tmp( fine_clusters_nums_max, stream, device_memory); @@ -1237,11 +1426,18 @@ auto build_fine_clusters(const raft::resources& handle, if (params.metric == cuvs::distance::DistanceType::L2Expanded || params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || params.metric == cuvs::distance::DistanceType::CosineExpanded) { - thrust::gather(raft::resource::get_thrust_policy(handle), - mc_trainset_ids, - mc_trainset_ids + k, - dataset_norm_mptr, - mc_trainset_norm); + if (params.is_packed_binary) { + // Expanded bits have a constant squared row norm. + raft::matrix::fill(handle, + raft::make_device_matrix_view(mc_trainset_norm, k, 1), + static_cast(transformed_dim)); + } else { + thrust::gather(raft::resource::get_thrust_policy(handle), + mc_trainset_ids, + mc_trainset_ids + k, + dataset_norm_mptr, + mc_trainset_norm); + } } build_clusters(handle, @@ -1256,12 +1452,12 @@ auto build_fine_clusters(const raft::resources& handle, mapping_op, device_memory, mc_trainset_norm); - - raft::copy(handle, - raft::make_device_vector_view(cluster_centers + (dim * fine_clusters_csum[i]), - fine_clusters_nums[i] * dim), - raft::make_device_vector_view(mc_trainset_ccenters.data(), - fine_clusters_nums[i] * dim)); + raft::copy( + handle, + raft::make_device_vector_view(cluster_centers + (transformed_dim * fine_clusters_csum[i]), + fine_clusters_nums[i] * transformed_dim), + raft::make_device_vector_view(mc_trainset_ccenters.data(), + fine_clusters_nums[i] * transformed_dim)); raft::resource::sync_stream(handle, stream); n_clusters_done += fine_clusters_nums[i]; } @@ -1299,8 +1495,9 @@ void build_hierarchical(const raft::resources& handle, MappingOpT mapping_op, MathT* inertia = nullptr) { - auto stream = raft::resource::get_cuda_stream(handle); - using LabelT = uint32_t; + auto stream = raft::resource::get_cuda_stream(handle); + using LabelT = uint32_t; + IdxT transformed_dim = centers_dim(dim, params.is_packed_binary); raft::common::nvtx::range fun_scope( "build_hierarchical(%zu, %u)", static_cast(n_rows), n_clusters); @@ -1312,21 +1509,22 @@ void build_hierarchical(const raft::resources& handle, rmm::mr::managed_memory_resource managed_memory; rmm::device_async_resource_ref device_memory = raft::resource::get_workspace_resource_ref(handle); auto [max_minibatch_size, mem_per_row] = calc_minibatch_size( - handle, n_clusters, n_rows, dim, params.metric, std::is_same_v); + handle, n_clusters, n_rows, transformed_dim, params.metric, !std::is_same_v); // Precompute the L2 norm of the dataset if relevant and not yet computed. rmm::device_uvector dataset_norm_buf(0, stream, device_memory); const MathT* dataset_norm = nullptr; if ((params.metric == cuvs::distance::DistanceType::L2Expanded || params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || - params.metric == cuvs::distance::DistanceType::CosineExpanded)) { + params.metric == cuvs::distance::DistanceType::CosineExpanded) && + !params.is_packed_binary) { dataset_norm_buf.resize(n_rows, stream); for (IdxT offset = 0; offset < n_rows; offset += max_minibatch_size) { IdxT minibatch_size = std::min(max_minibatch_size, n_rows - offset); if (params.metric == cuvs::distance::DistanceType::CosineExpanded) compute_norm(handle, dataset_norm_buf.data() + offset, - dataset + dim * offset, + dataset + offset * dim, dim, minibatch_size, mapping_op, @@ -1335,14 +1533,23 @@ void build_hierarchical(const raft::resources& handle, else compute_norm(handle, dataset_norm_buf.data() + offset, - dataset + dim * offset, + dataset + offset * dim, dim, minibatch_size, mapping_op, raft::identity_op{}, device_memory); } - dataset_norm = (const MathT*)dataset_norm_buf.data(); + dataset_norm = dataset_norm_buf.data(); + } else if (params.is_packed_binary && + (params.metric == cuvs::distance::DistanceType::L2Expanded || + params.metric == cuvs::distance::DistanceType::L2SqrtExpanded)) { + dataset_norm_buf.resize(n_rows, stream); + raft::matrix::fill( + handle, + raft::make_device_matrix_view(dataset_norm_buf.data(), n_rows, 1), + static_cast(transformed_dim)); + dataset_norm = dataset_norm_buf.data(); } /* Temporary workaround to cub::DeviceHistogram not supporting any type that isn't natively @@ -1354,7 +1561,8 @@ void build_hierarchical(const raft::resources& handle, rmm::device_uvector mesocluster_labels_buf(n_rows, stream, managed_memory); rmm::device_uvector mesocluster_sizes_buf(n_mesoclusters, stream, managed_memory); { - rmm::device_uvector mesocluster_centers_buf(n_mesoclusters * dim, stream, device_memory); + rmm::device_uvector mesocluster_centers_buf( + n_mesoclusters * transformed_dim, stream, device_memory); build_clusters(handle, params, dim, diff --git a/cpp/src/cluster/kmeans_balanced.cuh b/cpp/src/cluster/kmeans_balanced.cuh index 68c4925252..9f2833f50b 100644 --- a/cpp/src/cluster/kmeans_balanced.cuh +++ b/cpp/src/cluster/kmeans_balanced.cuh @@ -73,10 +73,12 @@ void fit(const raft::resources& handle, MappingOpT mapping_op = raft::identity_op(), std::optional> inertia = std::nullopt) { - RAFT_EXPECTS(X.extent(1) == centroids.extent(1), + cuvs::cluster::kmeans::detail::validate_packed_binary_metric(params); + auto centers_dim = + cuvs::cluster::kmeans::detail::centers_dim(X.extent(1), params.is_packed_binary); + RAFT_EXPECTS(centers_dim == centroids.extent(1), "Number of features in dataset and centroids are different"); - RAFT_EXPECTS(static_cast(X.extent(0)) * static_cast(X.extent(1)) <= - static_cast(std::numeric_limits::max()), + RAFT_EXPECTS(X.extent(0) <= std::numeric_limits::max() / centers_dim, "The chosen index type cannot represent all indices for the given dataset"); RAFT_EXPECTS(centroids.extent(0) > IndexT{0} && centroids.extent(0) <= X.extent(0), "The number of centroids must be strictly positive and cannot exceed the number of " @@ -135,12 +137,16 @@ void predict(const raft::resources& handle, raft::device_vector_view labels, MappingOpT mapping_op = raft::identity_op()) { + cuvs::cluster::kmeans::detail::validate_packed_binary_metric(params); RAFT_EXPECTS(X.extent(0) == labels.extent(0), "Number of rows in dataset and labels are different"); - RAFT_EXPECTS(X.extent(1) == centroids.extent(1), + auto centers_dim = + cuvs::cluster::kmeans::detail::centers_dim(X.extent(1), params.is_packed_binary); + RAFT_EXPECTS(centers_dim == centroids.extent(1), "Number of features in dataset and centroids are different"); - RAFT_EXPECTS(static_cast(X.extent(0)) * static_cast(X.extent(1)) <= - static_cast(std::numeric_limits::max()), + RAFT_EXPECTS(centroids.extent(0) > 0, "The number of centroids must be strictly positive"); + RAFT_EXPECTS(X.extent(0) <= std::numeric_limits::max() / centers_dim && + centroids.extent(0) <= std::numeric_limits::max() / centers_dim, "The chosen index type cannot represent all indices for the given dataset"); RAFT_EXPECTS(static_cast(centroids.extent(0)) <= static_cast(std::numeric_limits::max()), @@ -266,13 +272,16 @@ EXTERN_TEMPLATE_BUILD_CLUSTERS( * @param[in] X Dataset for which to calculate cluster centers. The data must be in * row-major format. [dim = n_samples x n_features] * @param[in] labels The input labels [dim = n_samples] - * @param[out] centroids The output centroids [dim = n_clusters x n_features] + * @param[out] centroids The output centroids + * [dim = n_clusters x + * (is_packed_binary ? 8 * n_features : n_features)] * @param[out] cluster_sizes Size of each cluster [dim = n_clusters] * @param[in] reset_counters Whether to clear the output arrays before calculating. * When set to `false`, this function may be used to update existing * centers and sizes using the weighted average principle. * @param[in] mapping_op (optional) Functor to convert from the input datatype to the * arithmetic datatype. If DataT == MathT, this must be the identity. + * @param[in] is_packed_binary Treat uint8_t X as packed bits. Requires DataT == uint8_t. */ template centroids, raft::device_vector_view cluster_sizes, bool reset_counters = true, - MappingOpT mapping_op = raft::identity_op()) + MappingOpT mapping_op = raft::identity_op(), + bool is_packed_binary = false) { RAFT_EXPECTS(X.extent(0) == labels.extent(0), "Number of rows in dataset and labels are different"); - RAFT_EXPECTS(X.extent(1) == centroids.extent(1), + auto centers_dim = + cuvs::cluster::kmeans::detail::centers_dim(X.extent(1), is_packed_binary); + RAFT_EXPECTS(centers_dim == centroids.extent(1), "Number of features in dataset and centroids are different"); + RAFT_EXPECTS(X.extent(0) <= std::numeric_limits::max() / centers_dim && + centroids.extent(0) <= std::numeric_limits::max() / centers_dim, + "The chosen index type cannot represent all indices for the given dataset"); RAFT_EXPECTS(centroids.extent(0) == cluster_sizes.extent(0), - "Number of rows in centroids and clusyer_sizes are different"); + "Number of rows in centroids and cluster_sizes are different"); cuvs::cluster::kmeans::detail::calc_centers_and_sizes( handle, @@ -305,6 +320,7 @@ void calc_centers_and_sizes(const raft::resources& handle, X.extent(0), labels.data_handle(), reset_counters, + is_packed_binary, mapping_op, raft::resource::get_workspace_resource_ref(handle)); } diff --git a/cpp/src/cluster/kmeans_balanced_build_clusters_impl.cuh b/cpp/src/cluster/kmeans_balanced_build_clusters_impl.cuh index 2bce856c6c..e160cc710f 100644 --- a/cpp/src/cluster/kmeans_balanced_build_clusters_impl.cuh +++ b/cpp/src/cluster/kmeans_balanced_build_clusters_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -50,10 +50,19 @@ void build_clusters(const raft::resources& handle, MappingOpT mapping_op, std::optional> X_norm) { + cuvs::cluster::kmeans::detail::validate_packed_binary_metric(params); + RAFT_EXPECTS(centroids.extent(0) > IndexT{0}, + "The number of centroids must be strictly positive"); RAFT_EXPECTS(X.extent(0) == labels.extent(0), "Number of rows in dataset and labels are different"); - RAFT_EXPECTS(X.extent(1) == centroids.extent(1), + auto centers_dim = + cuvs::cluster::kmeans::detail::centers_dim(X.extent(1), params.is_packed_binary); + RAFT_EXPECTS(centers_dim == centroids.extent(1), "Number of features in dataset and centroids are different"); + RAFT_EXPECTS(X.extent(0) <= std::numeric_limits::max() / centers_dim, + "The chosen index type cannot represent all indices for the given dataset"); + RAFT_EXPECTS(!X_norm.has_value() || X_norm->extent(0) == X.extent(0), + "Number of rows in dataset and norms are different"); RAFT_EXPECTS(centroids.extent(0) == cluster_sizes.extent(0), "Number of rows in centroids and clusyer_sizes are different"); diff --git a/cpp/src/distance/detail/distance_ops/all_ops.cuh b/cpp/src/distance/detail/distance_ops/all_ops.cuh index f0a3984eb6..104ce00f01 100644 --- a/cpp/src/distance/detail/distance_ops/all_ops.cuh +++ b/cpp/src/distance/detail/distance_ops/all_ops.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,6 +9,7 @@ #include "cutlass.cuh" // The distance operations: +#include "../distance_ops/bitwise_hamming.cuh" #include "../distance_ops/canberra.cuh" #include "../distance_ops/correlation.cuh" #include "../distance_ops/cosine.cuh" diff --git a/cpp/src/distance/detail/distance_ops/bitwise_hamming.cuh b/cpp/src/distance/detail/distance_ops/bitwise_hamming.cuh new file mode 100644 index 0000000000..5f1808143e --- /dev/null +++ b/cpp/src/distance/detail/distance_ops/bitwise_hamming.cuh @@ -0,0 +1,70 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include +#include +#include +#include + +namespace cuvs::distance::detail::ops { + +/** + * @brief the Bitwise Hamming distance matrix calculation + * It computes the following equation: + * + * c_ij = sum_k popcount(x_ik XOR y_kj) + * + * where x and y are binary data packed as uint8_t + */ +template +struct bitwise_hamming_distance_op { + using DataT = DataType; + using AccT = AccType; + using IdxT = IdxType; + + IdxT k; + + bitwise_hamming_distance_op(IdxT k_) : k(k_) + { + static_assert(std::is_same_v, "BitwiseHamming only supports uint8_t"); + static_assert(std::is_same_v, "BitwiseHamming requires a uint32_t accumulator"); + RAFT_EXPECTS(k >= 0 && static_cast(k) <= std::numeric_limits::max() / 8, + "BitwiseHamming dimension exceeds the uint32_t accumulator range"); + } + + static constexpr bool use_norms = false; + static constexpr bool expensive_inner_loop = false; + + template + static constexpr size_t shared_mem_size() + { + return Policy::SmemSize; + } + + __device__ __forceinline__ void core(AccT& acc, DataT& x, DataT& y) const + { + static_assert(std::is_same_v, "BitwiseHamming only supports uint8_t"); + // Ensure proper masking and casting to avoid undefined behavior + uint32_t xor_val = static_cast(static_cast(x ^ y)); + uint32_t masked_val = xor_val & 0xffu; + int popcount = __popc(masked_val); + acc += static_cast(popcount); + } + + template + __device__ __forceinline__ void epilog(AccT acc[Policy::AccRowsPerTh][Policy::AccColsPerTh], + AccT* regxn, + AccT* regyn, + IdxT gridStrideX, + IdxT gridStrideY) const + { + } +}; + +} // namespace cuvs::distance::detail::ops diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index 632780fc33..e100ace651 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -10,6 +10,7 @@ #include "fused_distance_nn/cutile/fused_1nn_tile.hpp" #endif #include "fused_distance_nn/cutlass_base.cuh" +#include "fused_distance_nn/fused_bitwise_hamming_nn.cuh" #include "fused_distance_nn/fused_cosine_nn.cuh" #include "fused_distance_nn/fused_l2_nn.cuh" #include "fused_distance_nn/helper_structs.cuh" @@ -180,28 +181,46 @@ void fusedDistanceNNImpl(raft::resources const& handle, dim3 blk(P::Nthreads); auto nblks = raft::ceildiv(m, P::Nthreads); - constexpr auto maxVal = std::numeric_limits::max(); - typedef raft::KeyValuePair KVPair; + using AccT = std::conditional_t, uint32_t, DataT>; + constexpr auto maxVal = std::numeric_limits::max(); RAFT_CUDA_TRY(cudaMemsetAsync(workspace, 0, sizeof(int) * m, stream.get())); if (initOutBuffer) { - initKernel + initKernel <<>>(min, m, maxVal, redOp); RAFT_CUDA_TRY(cudaGetLastError()); } + // An empty candidate set leaves the initialized (or supplied) result unchanged. + if (n == 0) { return; } + switch (metric) { case cuvs::distance::DistanceType::CosineExpanded: - fusedCosineNN( - min, x, y, xn, yn, m, n, k, workspace, redOp, pairRedOp, sqrt, stream.get()); + if constexpr (std::is_same_v || std::is_same_v) { + RAFT_FAIL("Cosine distance is not supported for uint8_t/int8_t data types"); + } else { + fusedCosineNN( + min, x, y, xn, yn, m, n, k, workspace, redOp, pairRedOp, sqrt, stream.get()); + } break; case cuvs::distance::DistanceType::L2SqrtExpanded: case cuvs::distance::DistanceType::L2Expanded: - // initOutBuffer is take care by fusedDistanceNNImpl() so we set it false to fusedL2NNImpl. - fusedL2NNImpl( - min, x, y, xn, yn, m, n, k, workspace, redOp, pairRedOp, sqrt, false, stream.get()); + if constexpr (std::is_same_v || std::is_same_v) { + RAFT_FAIL("L2 distance is not supported for uint8_t/int8_t data types"); + } else { + fusedL2NNImpl( + min, x, y, xn, yn, m, n, k, workspace, redOp, pairRedOp, sqrt, false, stream.get()); + } + break; + case cuvs::distance::DistanceType::BitwiseHamming: + if constexpr (std::is_same_v) { + fusedBitwiseHammingNN( + min, x, y, xn, yn, m, n, k, workspace, redOp, pairRedOp, sqrt, stream.get()); + } else { + RAFT_FAIL("BitwiseHamming distance only supports uint8_t data type"); + } break; - default: RAFT_FAIL("Only cosine and L2 metrics are supported by fusedDistanceNN"); + default: RAFT_FAIL("only cosine/l2/bitwise hamming metric is supported with fusedDistanceNN"); } } diff --git a/cpp/src/distance/detail/fused_distance_nn/fused_bitwise_hamming_nn.cuh b/cpp/src/distance/detail/fused_distance_nn/fused_bitwise_hamming_nn.cuh new file mode 100644 index 0000000000..e0daa55a8d --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/fused_bitwise_hamming_nn.cuh @@ -0,0 +1,81 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include "../distance_ops/bitwise_hamming.cuh" // ops::bitwise_hamming_distance_op +#include "../pairwise_distance_base.cuh" // PairwiseDistances +#include "helper_structs.cuh" +#include "simt_kernel.cuh" + +namespace cuvs { +namespace distance { +namespace detail { + +/** + * @brief Fused BitwiseHamming distance and 1-nearest-neighbor + * + * This implementation is only meaningful for uint8_t data type. + * The if constexpr in fusedDistanceNNImpl ensures it's only called for uint8_t. + */ +template +void fusedBitwiseHammingNN(OutT* min, + const DataT* x, + const DataT* y, + const DataT* xn, + const DataT* yn, + IdxT m, + IdxT n, + IdxT k, + int* workspace, + ReduceOpT redOp, + KVPReduceOpT pairRedOp, + bool sqrt, + cudaStream_t stream) +{ + typedef Policy P; + + dim3 blk(P::Nthreads); + constexpr auto maxVal = std::numeric_limits::max(); + using distance_op_type = ops::bitwise_hamming_distance_op; + distance_op_type distance_op{k}; + auto kernel = fusedDistanceNNkernel; + + constexpr size_t shmemSize = P::SmemSize; + + dim3 grid = launchConfigGenerator

(m, n, shmemSize, kernel); + + kernel<<>>(min, + x, + y, + nullptr, + nullptr, + m, + n, + k, + maxVal, + workspace, + redOp, + pairRedOp, + distance_op, + raft::identity_op{}); + + RAFT_CUDA_TRY(cudaGetLastError()); +} + +} // namespace detail +} // namespace distance +} // namespace cuvs diff --git a/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh b/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh index 762c720568..eeccdd5e45 100644 --- a/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh +++ b/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -27,8 +27,15 @@ namespace detail { template struct KVPMinReduceImpl { typedef raft::KeyValuePair KVP; - DI KVP operator()(LabelT rit, const KVP& a, const KVP& b) { return b.value < a.value ? b : a; } - DI KVP operator()(const KVP& a, const KVP& b) { return b.value < a.value ? b : a; } + // Use index as tiebreaker for consistent behavior when distances are equal + DI KVP operator()(LabelT rit, const KVP& a, const KVP& b) + { + return (b.value < a.value || (b.value == a.value && b.key < a.key)) ? b : a; + } + DI KVP operator()(const KVP& a, const KVP& b) + { + return (b.value < a.value || (b.value == a.value && b.key < a.key)) ? b : a; + } }; // KVPMinReduce @@ -38,14 +45,16 @@ struct MinAndDistanceReduceOpImpl { DI void operator()(LabelT rid, KVP* out, const KVP& other) const { - if (other.value < out->value) { + // Use index as tiebreaker for consistent behavior when distances are equal + if (other.value < out->value || (other.value == out->value && other.key < out->key)) { out->key = other.key; out->value = other.value; } } DI void operator()(LabelT rid, volatile KVP* out, const KVP& other) const { - if (other.value < out->value) { + // Use index as tiebreaker for consistent behavior when distances are equal + if (other.value < out->value || (other.value == out->value && other.key < out->key)) { out->key = other.key; out->value = other.value; } @@ -75,7 +84,7 @@ struct MinAndDistanceReduceOpImpl { DI void init(KVP* out, DataT maxVal) const { out->value = maxVal; - out->key = 0xfffffff0; + out->key = std::numeric_limits::max(); } DI void init_key(DataT& out, LabelT idx) const { return; } diff --git a/cpp/src/distance/detail/fused_distance_nn/simt_kernel.cuh b/cpp/src/distance/detail/fused_distance_nn/simt_kernel.cuh index 4211ff653d..eccdc3de67 100644 --- a/cpp/src/distance/detail/fused_distance_nn/simt_kernel.cuh +++ b/cpp/src/distance/detail/fused_distance_nn/simt_kernel.cuh @@ -1,14 +1,15 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include "../distance_ops/l2_exp.cuh" // ops::l2_exp_distance_op -#include "../pairwise_distance_base.cuh" // PairwiseDistances -#include // raft::KeyValuePair -#include // Policy +#include "../distance_ops/bitwise_hamming.cuh" // ops::bitwise_hamming_distance_op +#include "../distance_ops/l2_exp.cuh" // ops::l2_exp_distance_op +#include "../pairwise_distance_base.cuh" // PairwiseDistances +#include // raft::KeyValuePair +#include // Policy #include // size_t #include // std::numeric_limits @@ -64,111 +65,121 @@ __launch_bounds__(P::Nthreads, 2) RAFT_KERNEL fusedDistanceNNkernel(OutT* min, IdxT m, IdxT n, IdxT k, - DataT maxVal, + typename OpT::AccT maxVal, int* mutex, ReduceOpT redOp, KVPReduceOpT pairRedOp, OpT distance_op, FinalLambda fin_op) { -// compile only if below non-ampere arch. -#if __CUDA_ARCH__ < 800 - extern __shared__ char smem[]; + // For hamming-like distances, we need this kernel on all architectures + // For other distances, only use for pre-ampere architectures + +#if __CUDA_ARCH__ >= 800 + static constexpr bool compile = + std::is_same_v>; +#else + static constexpr bool compile = true; +#endif - typedef raft::KeyValuePair KVPair; - KVPair val[P::AccRowsPerTh]; -#pragma unroll - for (int i = 0; i < P::AccRowsPerTh; ++i) { - val[i] = {0, maxVal}; - } + if constexpr (compile) { + extern __shared__ char smem[]; - // epilogue operation lambda for final value calculation - auto epilog_lambda = [n, pairRedOp, &val, maxVal] __device__( - DataT acc[P::AccRowsPerTh][P::AccColsPerTh], - DataT * regxn, - DataT * regyn, - IdxT gridStrideX, - IdxT gridStrideY) { - KVPReduceOpT pairRed_op(pairRedOp); - - // intra thread reduce - const auto acccolid = threadIdx.x % P::AccThCols; - const auto accrowid = threadIdx.x / P::AccThCols; + using AccT = typename OpT::AccT; + typedef raft::KeyValuePair KVPair; + KVPair val[P::AccRowsPerTh]; #pragma unroll for (int i = 0; i < P::AccRowsPerTh; ++i) { -#pragma unroll - for (int j = 0; j < P::AccColsPerTh; ++j) { - auto tmpkey = acccolid + j * P::AccThCols + gridStrideX; - KVPair tmp = {tmpkey, acc[i][j]}; - if (tmpkey < n) { - val[i] = pairRed_op(accrowid + i * P::AccThRows + gridStrideY, tmp, val[i]); - } - } + val[i] = {0, maxVal}; } - }; - auto rowEpilog_lambda = - [m, mutex, min, pairRedOp, redOp, &val, maxVal] __device__(IdxT gridStrideY) { + // epilogue operation lambda for final value calculation + auto epilog_lambda = [n, pairRedOp, &val, maxVal] __device__( + AccT acc[P::AccRowsPerTh][P::AccColsPerTh], + AccT * regxn, + AccT * regyn, + IdxT gridStrideX, + IdxT gridStrideY) { KVPReduceOpT pairRed_op(pairRedOp); - ReduceOpT red_op(redOp); + // intra thread reduce + const auto acccolid = threadIdx.x % P::AccThCols; const auto accrowid = threadIdx.x / P::AccThCols; - const auto lid = raft::laneId(); - - // reduce #pragma unroll for (int i = 0; i < P::AccRowsPerTh; ++i) { #pragma unroll - for (int j = P::AccThCols / 2; j > 0; j >>= 1) { - // Actually, the srcLane (lid +j) should be (lid +j) % P:AccThCols, - // but the shfl op applies the modulo internally. - auto tmpkey = raft::shfl(val[i].key, lid + j, P::AccThCols); - auto tmpvalue = raft::shfl(val[i].value, lid + j, P::AccThCols); - KVPair tmp = {tmpkey, tmpvalue}; - val[i] = pairRed_op(accrowid + i * P::AccThRows + gridStrideY, tmp, val[i]); + for (int j = 0; j < P::AccColsPerTh; ++j) { + auto tmpkey = acccolid + j * P::AccThCols + gridStrideX; + KVPair tmp = {tmpkey, acc[i][j]}; + if (tmpkey < n) { + val[i] = pairRed_op(accrowid + i * P::AccThRows + gridStrideY, tmp, val[i]); + } } } + }; - updateReducedVal(mutex, min, val, red_op, m, gridStrideY); + auto rowEpilog_lambda = + [m, mutex, min, pairRedOp, redOp, &val, maxVal] __device__(IdxT gridStrideY) { + KVPReduceOpT pairRed_op(pairRedOp); + ReduceOpT red_op(redOp); - // reset the val array. + const auto accrowid = threadIdx.x / P::AccThCols; + const auto lid = raft::laneId(); + + // reduce #pragma unroll - for (int i = 0; i < P::AccRowsPerTh; ++i) { - val[i] = {0, maxVal}; - } - }; + for (int i = 0; i < P::AccRowsPerTh; ++i) { +#pragma unroll + for (int j = P::AccThCols / 2; j > 0; j >>= 1) { + // Actually, the srcLane (lid +j) should be (lid +j) % P:AccThCols, + // but the shfl op applies the modulo internally. + auto tmpkey = raft::shfl(val[i].key, lid + j, P::AccThCols); + auto tmpvalue = raft::shfl(val[i].value, lid + j, P::AccThCols); + KVPair tmp = {tmpkey, tmpvalue}; + val[i] = pairRed_op(accrowid + i * P::AccThRows + gridStrideY, tmp, val[i]); + } + } - IdxT lda = k, ldb = k, ldd = n; - constexpr bool row_major = true; - constexpr bool write_out = false; - PairwiseDistances - obj(x, - y, - m, - n, - k, - lda, - ldb, - ldd, - xn, - yn, - nullptr, // Output pointer - smem, - distance_op, - epilog_lambda, - fin_op, - rowEpilog_lambda); - obj.run(); -#endif + updateReducedVal(mutex, min, val, red_op, m, gridStrideY); + + // reset the val array. +#pragma unroll + for (int i = 0; i < P::AccRowsPerTh; ++i) { + val[i] = {0, maxVal}; + } + }; + + IdxT lda = k, ldb = k, ldd = n; + constexpr bool row_major = true; + constexpr bool write_out = false; + PairwiseDistances + obj(x, + y, + m, + n, + k, + lda, + ldb, + ldd, + reinterpret_cast(xn), + reinterpret_cast(yn), + nullptr, // Output pointer + smem, + distance_op, + epilog_lambda, + fin_op, + rowEpilog_lambda); + obj.run(); + } } } // namespace detail diff --git a/cpp/src/distance/detail/pairwise_matrix/dispatch-ext.cuh b/cpp/src/distance/detail/pairwise_matrix/dispatch-ext.cuh index c93a2f3f2b..e6878c80b2 100644 --- a/cpp/src/distance/detail/pairwise_matrix/dispatch-ext.cuh +++ b/cpp/src/distance/detail/pairwise_matrix/dispatch-ext.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -67,6 +67,11 @@ void pairwise_matrix_dispatch(OpT distance_op, instantiate_cuvs_distance_detail_pairwise_matrix_dispatch( \ OpT, half, float, float, FinOpT, IdxT); +#define instantiate_cuvs_distance_detail_pairwise_matrix_dispatch_by_algo_bitwise_hamming(OpT, \ + IdxT) \ + instantiate_cuvs_distance_detail_pairwise_matrix_dispatch( \ + OpT, uint8_t, uint32_t, uint32_t, raft::identity_op, IdxT); + /* * Hierarchy of instantiations: * @@ -112,5 +117,9 @@ instantiate_cuvs_distance_detail_pairwise_matrix_dispatch_by_algo( instantiate_cuvs_distance_detail_pairwise_matrix_dispatch_by_algo_default( cuvs::distance::detail::ops::l2_exp_distance_op, int64_t); +instantiate_cuvs_distance_detail_pairwise_matrix_dispatch_by_algo_bitwise_hamming( + cuvs::distance::detail::ops::bitwise_hamming_distance_op, int64_t); + +#undef instantiate_cuvs_distance_detail_pairwise_matrix_dispatch_by_algo_bitwise_hamming #undef instantiate_cuvs_distance_detail_pairwise_matrix_dispatch_by_algo #undef instantiate_cuvs_distance_detail_pairwise_matrix_dispatch diff --git a/cpp/src/distance/detail/pairwise_matrix/dispatch-inl.cuh b/cpp/src/distance/detail/pairwise_matrix/dispatch-inl.cuh index aff78d87f9..691f319b72 100644 --- a/cpp/src/distance/detail/pairwise_matrix/dispatch-inl.cuh +++ b/cpp/src/distance/detail/pairwise_matrix/dispatch-inl.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 */ #pragma once @@ -79,6 +79,9 @@ void pairwise_matrix_dispatch(OpT distance_op, cudaStream_t stream, bool is_row_major) { + RAFT_EXPECTS(m >= 0 && n >= 0 && k >= 0, "Distance dimensions must be non-negative"); + if (m == 0 || n == 0) { return; } + // Create kernel parameter struct. Flip x and y if column major. IdxT ldx = is_row_major ? k : m; IdxT ldy = is_row_major ? k : n; diff --git a/cpp/src/distance/detail/pairwise_matrix/dispatch_matrix.json b/cpp/src/distance/detail/pairwise_matrix/dispatch_matrix.json index bf6be6bed6..d7cdb07668 100644 --- a/cpp/src/distance/detail/pairwise_matrix/dispatch_matrix.json +++ b/cpp/src/distance/detail/pairwise_matrix/dispatch_matrix.json @@ -124,5 +124,26 @@ ], "index_type": "int64_t", "index_abbrev": "i64" + }, + { + "_data": [ + { + "data_type": "uint8_t", + "data_abbrev": "u8", + "acc_type": "uint32_t", + "acc_abbrev": "u32", + "out_type": "uint32_t", + "out_abbrev": "u32" + } + ], + "_op": [ + { + "op_type": "cuvs::distance::detail::ops::bitwise_hamming_distance_op", + "op_abbrev": "bitwise_hamming", + "arch_includes": "" + } + ], + "index_type": "int64_t", + "index_abbrev": "i64" } ] diff --git a/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/compute_distance_epilog_matrix.json b/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/compute_distance_epilog_matrix.json index 3c03a381f7..91b6e15b6f 100644 --- a/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/compute_distance_epilog_matrix.json +++ b/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/compute_distance_epilog_matrix.json @@ -258,5 +258,64 @@ "index_abbrev": "i64" } ] + }, + { + "_distance": [ + { + "distance_name": "bitwise_hamming", + "distance_abbrev": "bitwise_hamming", + "op_type": "cuvs::distance::detail::ops::bitwise_hamming_distance_op", + "header_file": "distance/detail/distance_ops/bitwise_hamming.cuh" + } + ], + "_data": [ + { + "data_type": "uint8_t", + "data_abbrev": "u8", + "type_abbrev": "u8", + "acc_type": "uint32_t", + "acc_abbrev": "u32", + "out_type": "uint32_t", + "out_abbrev": "u32" + } + ], + "_index": [ + { + "index_type": "int64_t", + "index_abbrev": "i64" + } + ], + "_policy": [ + { + "policy_type": "Policy", + "layout_abbrev": "row", + "veclen": "1" + }, + { + "policy_type": "ColPolicy", + "layout_abbrev": "col", + "veclen": "1" + }, + { + "policy_type": "Policy", + "layout_abbrev": "row", + "veclen": "2" + }, + { + "policy_type": "ColPolicy", + "layout_abbrev": "col", + "veclen": "2" + }, + { + "policy_type": "Policy", + "layout_abbrev": "row", + "veclen": "4" + }, + { + "policy_type": "ColPolicy", + "layout_abbrev": "col", + "veclen": "4" + } + ] } ] diff --git a/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/compute_distance_matrix.json b/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/compute_distance_matrix.json index 0c3b0c4164..91c9a89678 100644 --- a/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/compute_distance_matrix.json +++ b/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/compute_distance_matrix.json @@ -151,5 +151,32 @@ "index_abbrev": "i64" } ] + }, + { + "_distance": [ + { + "distance_name": "bitwise_hamming", + "distance_abbrev": "bitwise_hamming", + "op_type": "cuvs::distance::detail::ops::bitwise_hamming_distance_op", + "header_file": "distance/detail/distance_ops/bitwise_hamming.cuh" + } + ], + "_data": [ + { + "data_type": "uint8_t", + "data_abbrev": "u8", + "type_abbrev": "u8", + "acc_type": "uint32_t", + "acc_abbrev": "u32", + "out_type": "uint32_t", + "out_abbrev": "u32" + } + ], + "_index": [ + { + "index_type": "int64_t", + "index_abbrev": "i64" + } + ] } ] diff --git a/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_jit.cuh b/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_jit.cuh index 17526e75b7..d8d665b112 100644 --- a/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_jit.cuh +++ b/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_jit.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -35,6 +35,10 @@ constexpr auto get_pairwise_scalar_type_tag() return cuvs::neighbors::detail::tag_f{}; } else if constexpr (std::is_same_v) { return cuvs::neighbors::detail::tag_d{}; + } else if constexpr (std::is_same_v) { + return cuvs::neighbors::detail::tag_u8{}; + } else if constexpr (std::is_same_v) { + return cuvs::neighbors::detail::tag_u32{}; } else if constexpr (std::is_same_v || std::is_same_v) { return cuvs::neighbors::detail::tag_h{}; } else { @@ -84,6 +88,11 @@ struct pairwise_distance_op_tag { "Pairwise matrix JIT LTO does not have a distance-op tag for this type"); }; +template +struct pairwise_distance_op_tag> { + using type = tag_distance_bitwise_hamming; +}; + template struct pairwise_distance_op_tag> { using type = tag_distance_canberra; diff --git a/cpp/src/distance/fused_distance_nn-inl.cuh b/cpp/src/distance/fused_distance_nn-inl.cuh index bbf3bd1e68..15e9a6bb3c 100644 --- a/cpp/src/distance/fused_distance_nn-inl.cuh +++ b/cpp/src/distance/fused_distance_nn-inl.cuh @@ -69,7 +69,8 @@ namespace distance { * @param[in] initOutBuffer whether to initialize the output buffer before the * main kernel launch * @param[in] isRowMajor whether the input/output is row or column major. - * @param[in] metric Distance metric to be used (supports L2, cosine) + * @param[in] metric Distance metric to be used (supports L2, cosine, and packed uint8_t + * Hamming) * @param[in] metric_arg power argument for distances like Minkowski (not supported for now) * @param[in] handle RAFT resources containing the caller-provided CUDA stream */ @@ -92,110 +93,123 @@ void fusedDistanceNN(raft::resources const& handle, cuvs::distance::DistanceType metric, float metric_arg) { - ASSERT(isRowMajor, "fusedDistanceNN only supports row major inputs"); + RAFT_EXPECTS(isRowMajor, "fusedDistanceNN only supports row major inputs"); + RAFT_EXPECTS(m >= 0 && n >= 0 && k >= 0, "Distance dimensions must be non-negative"); + if constexpr (std::is_same_v) { + RAFT_EXPECTS(metric == cuvs::distance::DistanceType::BitwiseHamming, + "uint8_t fused distance only supports BitwiseHamming"); + RAFT_EXPECTS(static_cast(k) <= std::numeric_limits::max() / 8, + "BitwiseHamming distance exceeds the uint32_t accumulator range"); + } else { + RAFT_EXPECTS(metric == cuvs::distance::DistanceType::CosineExpanded || + metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded, + "Floating-point fused distance only supports cosine and L2 metrics"); + } + if (m == 0) { return; } + // When k is smaller than 32, the Policy4x4 results in redundant calculations // as it uses tiles that have k=32. Therefore, use a "skinny" policy instead // that uses tiles with a smaller value of k. bool is_skinny = k < 32; + // Packed bytes use at most four elements per vectorized load. + constexpr int veclen16 = std::is_same_v ? 4 : 16 / sizeof(DataT); + constexpr int veclen8 = std::is_same_v ? 4 : 8 / sizeof(DataT); size_t bytes = sizeof(DataT) * k; auto px = reinterpret_cast(x); auto py = reinterpret_cast(y); if (16 % sizeof(DataT) == 0 && bytes % 16 == 0 && px % 16 == 0 && py % 16 == 0) { if (is_skinny) { - detail::fusedDistanceNNImpl< - DataT, - OutT, - IdxT, - typename raft::linalg::Policy4x4Skinny::Policy, - ReduceOpT>(handle, - min, - x, - y, - xn, - yn, - m, - n, - k, - (int*)workspace, - redOp, - pairRedOp, - sqrt, - initOutBuffer, - isRowMajor, - metric, - metric_arg); + detail::fusedDistanceNNImpl::Policy, + ReduceOpT>(handle, + min, + x, + y, + xn, + yn, + m, + n, + k, + (int*)workspace, + redOp, + pairRedOp, + sqrt, + initOutBuffer, + isRowMajor, + metric, + metric_arg); } else { - detail::fusedDistanceNNImpl< - DataT, - OutT, - IdxT, - typename raft::linalg::Policy4x4::Policy, - ReduceOpT>(handle, - min, - x, - y, - xn, - yn, - m, - n, - k, - (int*)workspace, - redOp, - pairRedOp, - sqrt, - initOutBuffer, - isRowMajor, - metric, - metric_arg); + detail::fusedDistanceNNImpl::Policy, + ReduceOpT>(handle, + min, + x, + y, + xn, + yn, + m, + n, + k, + (int*)workspace, + redOp, + pairRedOp, + sqrt, + initOutBuffer, + isRowMajor, + metric, + metric_arg); } } else if (8 % sizeof(DataT) == 0 && bytes % 8 == 0 && px % 8 == 0 && py % 8 == 0) { if (is_skinny) { - detail::fusedDistanceNNImpl< - DataT, - OutT, - IdxT, - typename raft::linalg::Policy4x4Skinny::Policy, - ReduceOpT>(handle, - min, - x, - y, - xn, - yn, - m, - n, - k, - (int*)workspace, - redOp, - pairRedOp, - sqrt, - initOutBuffer, - isRowMajor, - metric, - metric_arg); + detail::fusedDistanceNNImpl::Policy, + ReduceOpT>(handle, + min, + x, + y, + xn, + yn, + m, + n, + k, + (int*)workspace, + redOp, + pairRedOp, + sqrt, + initOutBuffer, + isRowMajor, + metric, + metric_arg); } else { - detail::fusedDistanceNNImpl< - DataT, - OutT, - IdxT, - typename raft::linalg::Policy4x4::Policy, - ReduceOpT>(handle, - min, - x, - y, - xn, - yn, - m, - n, - k, - (int*)workspace, - redOp, - pairRedOp, - sqrt, - initOutBuffer, - isRowMajor, - metric, - metric_arg); + detail::fusedDistanceNNImpl::Policy, + ReduceOpT>(handle, + min, + x, + y, + xn, + yn, + m, + n, + k, + (int*)workspace, + redOp, + pairRedOp, + sqrt, + initOutBuffer, + isRowMajor, + metric, + metric_arg); } } else { if (is_skinny) { @@ -273,7 +287,8 @@ void fusedDistanceNN(raft::resources const& handle, * @param[in] initOutBuffer whether to initialize the output buffer before the * main kernel launch * @param[in] isRowMajor whether the input/output is row or column major. - * @param[in] metric Distance metric to be used (supports L2, cosine) + * @param[in] metric Distance metric to be used (supports L2, cosine, and packed uint8_t + * Hamming) * @param[in] metric_arg power argument for distances like Minkowski (not supported for now) * @param[in] handle RAFT resources containing the caller-provided CUDA stream */ @@ -294,30 +309,56 @@ void fusedDistanceNNMinReduce(raft::resources const& handle, cuvs::distance::DistanceType metric, float metric_arg) { - static_assert( - std::is_same_v> || std::is_same_v, - "fusedDistanceNNMinReduce supports KVP or scalar distance output"); - detail::Top1nnTuning tuning{}; - const auto workspace_bytes = - top_1_nn_workspace_size(m, n, k, tuning, detail::Top1nnBackend::Cutlass); - top_1_nn(handle, - min, - x, - y, - xn, - yn, - m, - n, - k, - tuning, - workspace, - workspace_bytes, - sqrt, - initOutBuffer, - isRowMajor, - metric, - metric_arg, - detail::Top1nnBackend::Cutlass); + if constexpr (std::is_same_v) { + using AccT = uint32_t; + static_assert( + std::is_same_v> || std::is_same_v, + "BitwiseHamming supports uint32_t KVP or scalar distance output"); + MinAndDistanceReduceOp red_op; + KVPMinReduce pair_red_op; + fusedDistanceNN(handle, + min, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + red_op, + pair_red_op, + sqrt, + initOutBuffer, + isRowMajor, + metric, + metric_arg); + } else { + static_assert( + std::is_same_v> || std::is_same_v, + "fusedDistanceNNMinReduce supports KVP or scalar distance output"); + detail::Top1nnTuning tuning{}; + const auto workspace_bytes = + top_1_nn_workspace_size(m, n, k, tuning, detail::Top1nnBackend::Cutlass); + top_1_nn(handle, + min, + x, + y, + xn, + yn, + m, + n, + k, + tuning, + workspace, + workspace_bytes, + sqrt, + initOutBuffer, + isRowMajor, + metric, + metric_arg, + detail::Top1nnBackend::Cutlass); + } } namespace detail { diff --git a/cpp/src/neighbors/detail/ann_utils.cuh b/cpp/src/neighbors/detail/ann_utils.cuh index deddc47dcb..06ba30a1a3 100644 --- a/cpp/src/neighbors/detail/ann_utils.cuh +++ b/cpp/src/neighbors/detail/ann_utils.cuh @@ -216,6 +216,18 @@ HDI constexpr auto mapping::operator()(const float& x) const -> int8_t return static_cast(std::clamp(x * 128.0f, -128.0f, 127.0f)); } +template +struct bitwise_decode_op { + explicit bitwise_decode_op(const uint8_t* binary_vecs) : binary_vecs(binary_vecs) {} + const uint8_t* binary_vecs; + HDI constexpr auto operator()(const IdxT& i) const -> OutT + { + // Rows contain complete bytes, so flattened bit offsets directly address the packed input. + // Avoid multiplying the packed dimension (or row offset) in a potentially narrow index type. + return ((binary_vecs[i >> 3] >> (i & 7)) & 1) ? OutT{1} : OutT{-1}; + } +}; + /** * @brief Sets the first num bytes of the block of memory pointed by ptr to the specified value. * diff --git a/cpp/src/neighbors/ivf_binary/ivf_binary.cu b/cpp/src/neighbors/ivf_binary/ivf_binary.cu new file mode 100644 index 0000000000..7cfb7ec9a6 --- /dev/null +++ b/cpp/src/neighbors/ivf_binary/ivf_binary.cu @@ -0,0 +1,153 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../../util/kvikio_serialize.hpp" +#include "../../util/serialize_validation.hpp" +#include "../ivf_flat/ivf_flat_serialize.cuh" +#include "ivf_binary_build.cuh" +#include "ivf_binary_search.cuh" + +#include + +#include + +namespace cuvs::neighbors::ivf_binary { +namespace { + +// Serialization version of the index: the header, the quantizer and the centers, followed by +// the inverted lists in the IVF-Flat format. +constexpr int kSerializationVersion = 1; + +void serialize_impl(raft::resources const& handle, std::ostream& os, const index& idx) +{ + raft::serialize_scalar(handle, os, kSerializationVersion); + raft::serialize_scalar(handle, os, idx.dim()); + raft::serialize_scalar(handle, os, idx.n_lists()); + raft::serialize_scalar(handle, os, idx.adaptive_centers()); + raft::serialize_scalar(handle, os, idx.conservative_memory_allocation()); + raft::serialize_scalar(handle, os, idx.quantizer().threshold.extent(0)); + if (idx.quantizer().threshold.extent(0) > 0) { + cuvs::util::detail::serialize_mdspan(handle, os, idx.quantizer().threshold.view()); + } + cuvs::util::detail::serialize_mdspan(handle, os, idx.centers()); + if (idx.adaptive_centers()) { + cuvs::util::detail::serialize_mdspan(handle, os, idx.center_counts()); + } + ivf_flat::detail::serialize(handle, os, detail::index_access::lists(idx)); +} + +template +auto deserialize_impl(raft::resources const& handle, Input& input) -> index +{ + auto& is = cuvs::util::detail::input_stream(input); + + auto version = raft::deserialize_scalar(handle, is); + RAFT_EXPECTS(version == kSerializationVersion, + "ivf_binary::deserialize: serialization version mismatch, expected %d, got %d", + kSerializationVersion, + version); + auto dim = raft::deserialize_scalar(handle, is); + auto n_lists = raft::deserialize_scalar(handle, is); + bool adaptive = raft::deserialize_scalar(handle, is); + bool cma = raft::deserialize_scalar(handle, is); + auto n_threshold = raft::deserialize_scalar(handle, is); + + RAFT_EXPECTS(dim > 0 && n_lists > 0 && n_lists <= cuvs::util::kMaxIvfNLists, + "ivf_binary::deserialize: invalid dimension %u or number of lists %u", + dim, + n_lists); + RAFT_EXPECTS(n_threshold == 0 || n_threshold == static_cast(dim), + "ivf_binary::deserialize: invalid threshold size %ld", + static_cast(n_threshold)); + + index_params params; + params.n_lists = n_lists; + params.adaptive_centers = adaptive; + params.conservative_memory_allocation = cma; + index idx(handle, params, dim); + + auto& quantizer = detail::index_access::quantizer(idx); + quantizer.threshold = raft::make_device_vector(handle, n_threshold); + if (n_threshold > 0) { + cuvs::util::detail::deserialize_mdspan(handle, input, quantizer.threshold.view()); + } + cuvs::util::detail::deserialize_mdspan(handle, input, detail::index_access::centers(idx).view()); + if (adaptive) { + cuvs::util::detail::deserialize_mdspan( + handle, input, detail::index_access::center_counts(idx).view()); + } + + auto lists = ivf_flat::detail::deserialize_impl(handle, input); + RAFT_EXPECTS(lists.n_lists() == n_lists && lists.dim() == idx.encoded_dim() && + lists.metric() == cuvs::distance::DistanceType::BitwiseHamming, + "ivf_binary::deserialize: the lists do not match the index header"); + detail::index_access::lists(idx) = std::move(lists); + return idx; +} + +} // namespace + +auto build(raft::resources const& handle, + const index_params& index_params, + device_binary_dataset_view dataset) -> index +{ + return detail::build(handle, index_params, dataset); +} + +void extend(raft::resources const& handle, + device_binary_dataset_view new_vectors, + std::optional> new_indices, + index* index) +{ + RAFT_EXPECTS(index != nullptr, "index cannot be empty."); + RAFT_EXPECTS(new_vectors.dim() == index->dim(), + "The dimension of the vectors must equal the dimension of the index."); + if (new_indices.has_value()) { + RAFT_EXPECTS(new_vectors.n_rows() == new_indices->extent(0), + "new_vectors and new_indices have different number of rows"); + } + detail::extend(handle, + index, + new_vectors.data_const_view().data_handle(), + new_indices.has_value() ? new_indices->data_handle() : nullptr, + new_vectors.n_rows()); +} + +void search(raft::resources const& handle, + const search_params& params, + const index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter) +{ + detail::search(handle, params, index, queries, neighbors, distances, sample_filter); +} + +void serialize(raft::resources const& handle, std::ostream& os, const index& index) +{ + serialize_impl(handle, os, index); +} + +void serialize(raft::resources const& handle, const std::string& filename, const index& index) +{ + cuvs::util::kvikio_ofstream of(filename); + serialize_impl(handle, of, index); + of.close(); + if (!of) { RAFT_FAIL("Error writing output %s", filename.c_str()); } +} + +void deserialize(raft::resources const& handle, std::istream& is, index* index) +{ + *index = deserialize_impl(handle, is); +} + +void deserialize(raft::resources const& handle, const std::string& filename, index* index) +{ + cuvs::util::kvikio_file_reader reader(filename); + *index = deserialize_impl(handle, reader); +} + +} // namespace cuvs::neighbors::ivf_binary diff --git a/cpp/src/neighbors/ivf_binary/ivf_binary_build.cuh b/cpp/src/neighbors/ivf_binary/ivf_binary_build.cuh new file mode 100644 index 0000000000..fd63479ebe --- /dev/null +++ b/cpp/src/neighbors/ivf_binary/ivf_binary_build.cuh @@ -0,0 +1,235 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "../../cluster/kmeans_balanced.cuh" +#include "../../core/nvtx.hpp" +#include "../ivf_common.cuh" +#include "../ivf_flat/ivf_flat_build.cuh" +#include "../ivf_list.cuh" +#include "ivf_binary_index_access.hpp" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace cuvs::neighbors::ivf_binary::detail { + +using namespace cuvs::spatial::knn::detail; // NOLINT + +// One thread per input byte; list sizes bound every per-bit count by uint32_t. +RAFT_KERNEL accumulate_center_counts( + const uint8_t* data, const uint32_t* labels, uint32_t* counts, uint64_t n_elements, uint32_t dim) +{ + const uint64_t offset = uint64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (offset >= n_elements) { return; } + const auto row = offset / dim; + const auto col = offset % dim; + auto* output = counts + (uint64_t(labels[row]) * dim + col) * 8; + const uint8_t value = data[offset]; + for (int bit = 0; bit < 8; ++bit) { + if ((value >> bit) & 1) { atomicAdd(output + bit, 1u); } + } +} + +/** See cuvs::neighbors::ivf_binary::extend docs; the vectors are validated by the caller. */ +inline void extend(raft::resources const& handle, + index* idx, + const uint8_t* new_vectors, + const int64_t* new_indices, + int64_t n_rows) +{ + using LabelT = uint32_t; + RAFT_EXPECTS(new_indices != nullptr || idx->size() == 0, + "You must pass data indices when the index is non-empty."); + if (n_rows == 0) { return; } + + auto stream = raft::resource::get_cuda_stream(handle); + auto& lists = index_access::lists(*idx); + auto n_lists = idx->n_lists(); + auto dim = idx->encoded_dim(); + cuvs::common::nvtx::range fun_scope( + "ivf_binary::extend(%zu, %u)", size_t(n_rows), dim); + + auto& centers = index_access::centers(*idx); + auto labels = + raft::make_device_mdarray(handle, + raft::resource::get_large_workspace_resource_ref(handle), + raft::make_extents(n_rows)); + cuvs::cluster::kmeans::detail::predict_bitwise_hamming( + handle, + raft::make_device_matrix_view(new_vectors, n_rows, dim), + raft::make_device_matrix_view(centers.data_handle(), n_lists, dim), + raft::make_device_vector_view(labels.data_handle(), n_rows)); + + auto* list_sizes_ptr = lists.list_sizes().data_handle(); + auto old_list_sizes_dev = + raft::make_device_mdarray(handle, + raft::resource::get_workspace_resource_ref(handle), + raft::make_extents(n_lists)); + raft::copy(handle, + old_list_sizes_dev.view(), + raft::make_device_vector_view(list_sizes_ptr, n_lists)); + + raft::stats::histogram(raft::stats::HistTypeAuto, + reinterpret_cast(list_sizes_ptr), + int64_t(n_lists), + labels.data_handle(), + n_rows, + 1, + stream.get()); + raft::linalg::add(handle, + raft::make_device_vector_view(list_sizes_ptr, n_lists), + raft::make_device_vector_view( + old_list_sizes_dev.data_handle(), n_lists), + raft::make_device_vector_view(list_sizes_ptr, n_lists)); + + // Adaptive centers are the majority bits of everything added so far. The exact per-bit counts are + // kept, because rounding a previous mean back to a sum could flip an exact majority tie. + if (idx->adaptive_centers()) { + auto& counts = index_access::center_counts(*idx); + const auto n_elements = uint64_t(n_rows) * dim; + accumulate_center_counts<<>>( + new_vectors, labels.data_handle(), counts.data_handle(), n_elements, dim); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + const auto* counts_ptr = counts.data_handle(); + const auto* centers_ptr = centers.data_handle(); + raft::linalg::map_offset( + handle, + centers.view(), + [counts_ptr, centers_ptr, list_sizes_ptr, dim] __device__(int64_t offset) { + const auto list_size = list_sizes_ptr[offset / dim]; + if (list_size == 0) { return centers_ptr[offset]; } // keep the trained center + uint8_t packed = 0; + for (int bit = 0; bit < 8; ++bit) { + if (uint64_t(counts_ptr[offset * 8 + bit]) * 2 > list_size) { + packed |= uint8_t(1u << bit); + } + } + return packed; + }); + } + + // Calculate and allocate new list data + ivf_flat::list_spec list_device_spec{ + dim, idx->conservative_memory_allocation()}; + std::vector new_list_sizes(n_lists); + std::vector old_list_sizes(n_lists); + raft::copy(handle, + raft::make_host_vector_view(old_list_sizes.data(), n_lists), + raft::make_device_vector_view(old_list_sizes_dev.data_handle(), n_lists)); + raft::copy(handle, + raft::make_host_vector_view(new_list_sizes.data(), n_lists), + raft::make_device_vector_view(list_sizes_ptr, n_lists)); + raft::resource::sync_stream(handle); + for (uint32_t label = 0; label < n_lists; label++) { + ivf::resize_list( + handle, lists.lists()[label], list_device_spec, new_list_sizes[label], old_list_sizes[label]); + } + ivf::detail::recompute_internal_state(handle, lists); + // Restart from the old sizes: the kernel below uses them as the atomic counters of the lists. + raft::copy( + handle, + raft::make_device_vector_view(list_sizes_ptr, n_lists), + raft::make_device_vector_view(old_list_sizes_dev.data_handle(), n_lists)); + + const dim3 block_dim(256); + const dim3 grid_dim(raft::ceildiv(n_rows, block_dim.x)); + ivf_flat::detail::build_index_kernel + <<>>(labels.data_handle(), + new_vectors, + new_indices, + lists.data_ptrs().data_handle(), + lists.inds_ptrs().data_handle(), + list_sizes_ptr, + n_rows, + dim, + lists.veclen()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); +} + +/** See cuvs::neighbors::ivf_binary::build docs */ +inline auto build(raft::resources const& handle, + const index_params& params, + device_binary_dataset_view dataset) -> index +{ + auto stream = raft::resource::get_cuda_stream(handle); + const int64_t n_rows = dataset.n_rows(); + const auto codes = dataset.data_const_view(); + const auto threshold = dataset.const_dictionary_view().threshold; + cuvs::common::nvtx::range fun_scope( + "ivf_binary::build(%zu, %u)", size_t(n_rows), dataset.dim()); + + RAFT_EXPECTS(params.n_lists > 0, "n_lists must be positive"); + RAFT_EXPECTS(n_rows > 0, "empty dataset"); + RAFT_EXPECTS(n_rows >= params.n_lists, "number of rows can't be less than n_lists"); + + index idx(handle, params, dataset.dim()); + const auto dim = idx.encoded_dim(); + auto& lists = index_access::lists(idx); + + auto& quantizer = index_access::quantizer(idx); + quantizer.threshold = raft::make_device_vector(handle, threshold.extent(0)); + if (threshold.extent(0) > 0) { raft::copy(handle, quantizer.threshold.view(), threshold); } + + utils::memzero( + lists.accum_sorted_sizes().data_handle(), lists.accum_sorted_sizes().size(), stream); + utils::memzero(lists.list_sizes().data_handle(), lists.list_sizes().size(), stream); + utils::memzero(lists.data_ptrs().data_handle(), lists.data_ptrs().size(), stream); + utils::memzero(lists.inds_ptrs().data_handle(), lists.inds_ptrs().size(), stream); + + // Train the k-means clustering on the expanded bits and pack the centers again + { + auto trainset_ratio = std::max( + 1, n_rows / std::max(params.kmeans_trainset_fraction * n_rows, params.n_lists)); + auto n_rows_train = n_rows / trainset_ratio; + + rmm::device_uvector trainset( + n_rows_train * dim, stream, raft::resource::get_large_workspace_resource_ref(handle)); + // TODO: a proper sampling + raft::copy_matrix( + trainset.data(), dim, codes.data_handle(), dim * trainset_ratio, dim, n_rows_train, stream); + + cuvs::cluster::kmeans::balanced_params kmeans_params; + kmeans_params.n_iters = params.kmeans_n_iters; + kmeans_params.metric = cuvs::distance::DistanceType::L2Expanded; + kmeans_params.is_packed_binary = true; + + rmm::device_uvector decoded_centers( + size_t(params.n_lists) * dim * 8, stream, raft::resource::get_workspace_resource_ref(handle)); + auto decoded_centers_view = raft::make_device_matrix_view( + decoded_centers.data(), params.n_lists, int64_t(dim) * 8); + cuvs::cluster::kmeans_balanced::fit( + handle, + kmeans_params, + raft::make_device_matrix_view(trainset.data(), n_rows_train, dim), + decoded_centers_view, + raft::identity_op{}); + + cuvs::preprocessing::quantize::binary::quantizer zero_threshold(handle); + cuvs::preprocessing::quantize::binary::transform(handle, + zero_threshold, + raft::make_const_mdspan(decoded_centers_view), + index_access::centers(idx).view()); + } + + if (params.add_data_on_build) { extend(handle, &idx, codes.data_handle(), nullptr, n_rows); } + return idx; +} + +} // namespace cuvs::neighbors::ivf_binary::detail diff --git a/cpp/src/neighbors/ivf_binary/ivf_binary_index_access.hpp b/cpp/src/neighbors/ivf_binary/ivf_binary_index_access.hpp new file mode 100644 index 0000000000..3362d79b5f --- /dev/null +++ b/cpp/src/neighbors/ivf_binary/ivf_binary_index_access.hpp @@ -0,0 +1,33 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +namespace cuvs::neighbors::ivf_binary::detail { + +/** Internal access to the mutable state of the index. */ +struct index_access { + static auto lists(index& idx) -> ivf_flat::index& { return idx.lists_; } + static auto lists(const index& idx) -> const ivf_flat::index& + { + return idx.lists_; + } + static auto centers(index& idx) -> raft::device_matrix& + { + return idx.centers_; + } + static auto center_counts(index& idx) -> raft::device_matrix& + { + return idx.center_counts_; + } + static auto quantizer(index& idx) -> preprocessing::quantize::binary::quantizer& + { + return idx.quantizer_; + } +}; + +} // namespace cuvs::neighbors::ivf_binary::detail diff --git a/cpp/src/neighbors/ivf_binary/ivf_binary_search.cuh b/cpp/src/neighbors/ivf_binary/ivf_binary_search.cuh new file mode 100644 index 0000000000..ec0f8152c1 --- /dev/null +++ b/cpp/src/neighbors/ivf_binary/ivf_binary_search.cuh @@ -0,0 +1,333 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "../../core/nvtx.hpp" +#include "../../distance/detail/distance_ops/bitwise_hamming.cuh" +#include "../../distance/detail/pairwise_matrix/dispatch.cuh" +#include "../detail/ann_utils.cuh" +#include "../ivf_common.cuh" +#include "../ivf_flat/ivf_flat_interleaved_scan_ext.cuh" +#include "ivf_binary_index_access.hpp" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include + +#include +#include + +namespace cuvs::neighbors::ivf_binary::detail { + +using namespace cuvs::spatial::knn::detail; // NOLINT + +constexpr auto kMetric = cuvs::distance::DistanceType::BitwiseHamming; + +inline auto is_local_topk_feasible(uint32_t k) -> bool +{ + return k <= raft::matrix::detail::select::warpsort::kMaxCapacity; +} + +template +void search_impl(raft::resources const& handle, + const index& idx, + const uint8_t* queries, + uint32_t n_queries, + uint32_t queries_offset, + uint32_t k, + uint32_t n_probes, + uint32_t max_samples, + int64_t* neighbors, + float* distances, + rmm::device_async_resource_ref search_mr, + IvfSampleFilterT sample_filter) +{ + using IdxT = int64_t; + constexpr bool select_min = true; + auto stream = raft::resource::get_cuda_stream(handle); + const auto& lists = index_access::lists(idx); + const uint32_t n_lists = idx.n_lists(); + const uint32_t dim = idx.encoded_dim(); + std::size_t n_queries_probes = std::size_t(n_queries) * std::size_t(n_probes); + + // The distance value of cluster(list) and queries + rmm::device_uvector distance_buffer_dev(size_t(n_queries) * n_lists, stream, search_mr); + // The topk distance value of cluster(list) and queries + rmm::device_uvector coarse_distances_dev(n_queries_probes, stream, search_mr); + // The topk index of cluster(list) and queries + rmm::device_uvector coarse_indices_dev(n_queries_probes, stream, search_mr); + + // Optional structures if postprocessing is required + // The topk distance value of candidate vectors from each cluster(list) + rmm::device_uvector distances_tmp_dev(0, stream, search_mr); + // Number of samples for each query + rmm::device_uvector num_samples(0, stream, search_mr); + // Offsets per probe for each query + rmm::device_uvector chunk_index(0, stream, search_mr); + // The topk index of candidate vectors from each cluster(list), local index offset + // also we might need additional storage for select_k + rmm::device_uvector indices_tmp_dev(0, stream, search_mr); + rmm::device_uvector neighbors_uint32_buf(0, stream, search_mr); + auto distance_buffer_dev_view = + raft::make_device_matrix_view(distance_buffer_dev.data(), n_queries, n_lists); + + // Coarse search: Hamming distances to the centers + { + cuvs::distance::detail::ops::bitwise_hamming_distance_op distance_op{ + static_cast(dim)}; + rmm::device_uvector uint32_distances(size_t(n_queries) * n_lists, stream, search_mr); + cuvs::distance::detail::pairwise_matrix_dispatch(distance_op, + static_cast(n_queries), + static_cast(n_lists), + static_cast(dim), + queries, + idx.centers().data_handle(), + nullptr, + nullptr, + uint32_distances.data(), + raft::identity_op{}, + stream.get(), + true); + raft::linalg::map( + handle, + distance_buffer_dev_view, + raft::cast_op{}, + raft::make_const_mdspan(raft::make_device_matrix_view( + uint32_distances.data(), n_queries, n_lists))); + } + + cuvs::selection::select_k( + handle, + raft::make_const_mdspan(distance_buffer_dev_view), + std::nullopt, + raft::make_device_matrix_view(coarse_distances_dev.data(), n_queries, n_probes), + raft::make_device_matrix_view( + coarse_indices_dev.data(), n_queries, n_probes), + select_min); + + using ScanAccT = typename utils::config::value_t; + uint32_t grid_dim_x = 0; + if (n_probes > 1) { + // query the gridDimX size to store probes topK output + ivf_flat::detail::ivfflat_interleaved_scan( + lists, + nullptr, + nullptr, + n_queries, + queries_offset, + kMetric, + n_probes, + k, + 0, + nullptr, + select_min, + sample_filter, + nullptr, + nullptr, + grid_dim_x, + stream, + std::nullopt); + } else { + grid_dim_x = 1; + } + + num_samples.resize(n_queries, stream); + chunk_index.resize(n_queries_probes, stream); + + ivf::detail::calc_chunk_indices::configure(n_probes, n_queries)(lists.list_sizes().data_handle(), + coarse_indices_dev.data(), + chunk_index.data(), + num_samples.data(), + stream); + + auto distances_dev_ptr = distances; + + neighbors_uint32_buf.resize(std::size_t(n_queries) * std::size_t(k), stream); + uint32_t* neighbors_uint32 = neighbors_uint32_buf.data(); + + uint32_t* indices_dev_ptr = nullptr; + + bool manage_local_topk = is_local_topk_feasible(k); + if (!manage_local_topk || grid_dim_x > 1) { + auto target_size = std::size_t(n_queries) * (manage_local_topk ? grid_dim_x * k : max_samples); + + distances_tmp_dev.resize(target_size, stream); + if (manage_local_topk) indices_tmp_dev.resize(target_size, stream); + + distances_dev_ptr = distances_tmp_dev.data(); + indices_dev_ptr = indices_tmp_dev.data(); + } else { + indices_dev_ptr = neighbors_uint32; + } + + ivf_flat::detail::ivfflat_interleaved_scan( + lists, + queries, + coarse_indices_dev.data(), + n_queries, + queries_offset, + kMetric, + n_probes, + k, + max_samples, + chunk_index.data(), + select_min, + sample_filter, + indices_dev_ptr, + distances_dev_ptr, + grid_dim_x, + stream, + std::nullopt); + + // Merge topk values from different blocks + if (!manage_local_topk || grid_dim_x > 1) { + std::optional> num_samples_vector; + if (!manage_local_topk) { + num_samples_vector = + raft::make_device_vector_view(num_samples.data(), n_queries); + } + + auto cols = manage_local_topk ? (k * grid_dim_x) : max_samples; + + cuvs::selection::select_k( + handle, + raft::make_device_matrix_view( + distances_tmp_dev.data(), n_queries, cols), + raft::make_device_matrix_view( + indices_tmp_dev.data(), n_queries, cols), + raft::make_device_matrix_view(distances, n_queries, k), + raft::make_device_matrix_view(neighbors_uint32, n_queries, k), + select_min, + false, + cuvs::selection::SelectAlgo::kAuto, + num_samples_vector); + } + if (!manage_local_topk) { + // post process distances && neighbor IDs + ivf::detail::postprocess_distances( + handle, distances, distances, kMetric, n_queries, k, 1.0, false); + } + ivf::detail::postprocess_neighbors(neighbors, + neighbors_uint32, + lists.inds_ptrs().data_handle(), + coarse_indices_dev.data(), + chunk_index.data(), + n_queries, + n_probes, + k, + stream); +} + +template +void search_with_filtering( + raft::resources const& handle, + const search_params& params, + const index& idx, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + IvfSampleFilterT sample_filter) +{ + RAFT_EXPECTS( + queries.extent(0) == neighbors.extent(0) && queries.extent(0) == distances.extent(0), + "Number of rows in output neighbors and distances matrices must equal the number of queries."); + RAFT_EXPECTS(neighbors.extent(1) == distances.extent(1), + "Number of columns in output neighbors and distances matrices must be equal"); + RAFT_EXPECTS(queries.extent(1) == idx.encoded_dim(), + "Number of query bytes should equal the number of bytes in the index vectors."); + RAFT_EXPECTS(params.n_probes > 0, + "n_probes (number of clusters to probe in the search) must be positive."); + const uint32_t n_queries = queries.extent(0); + const uint32_t k = neighbors.extent(1); + if (n_queries == 0 || k == 0) { return; } + + common::nvtx::range fun_scope( + "ivf_binary::search(k = %u, n_queries = %u, dim = %u)", k, n_queries, idx.dim()); + + const auto& lists = index_access::lists(idx); + auto n_probes = std::min(params.n_probes, idx.n_lists()); + bool manage_local_topk = is_local_topk_feasible(k); + + uint32_t max_samples = 0; + if (!manage_local_topk) { + int64_t ms = raft::Pow2<128 / sizeof(float)>::roundUp( + std::max(lists.accum_sorted_sizes()(n_probes), k)); + RAFT_EXPECTS(ms <= int64_t(std::numeric_limits::max()), + "The maximum sample size is too big."); + max_samples = ms; + } + + // a batch size heuristic: try to keep the workspace within the specified size + constexpr uint64_t kExpectedWsSize = 1024 * 1024 * 1024; + uint64_t max_ws_size = + std::min(raft::resource::get_workspace_free_bytes(handle), kExpectedWsSize); + + uint64_t ws_size_per_query = 4ull * (2ull * n_probes + 2ull * idx.n_lists() + 1) + + (manage_local_topk ? ((sizeof(int64_t) + 4) * n_probes * k) + : (4ull * (max_samples + n_probes + 1))); + + const uint32_t max_queries = + std::min(n_queries, raft::div_rounding_up_safe(max_ws_size, ws_size_per_query)); + + for (uint32_t offset_q = 0; offset_q < n_queries; offset_q += max_queries) { + uint32_t queries_batch = raft::min(max_queries, n_queries - offset_q); + + search_impl(handle, + idx, + queries.data_handle() + size_t(offset_q) * idx.encoded_dim(), + queries_batch, + offset_q, + k, + n_probes, + max_samples, + neighbors.data_handle() + size_t(offset_q) * k, + distances.data_handle() + size_t(offset_q) * k, + raft::resource::get_workspace_resource_ref(handle), + sample_filter); + } +} + +/** Search with the binary-quantized queries; see cuvs::neighbors::ivf_binary::search. */ +inline void search(raft::resources const& handle, + const search_params& params, + const index& idx, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter_ref) +{ + try { + auto& sample_filter = + dynamic_cast(sample_filter_ref); + return search_with_filtering(handle, params, idx, queries, neighbors, distances, sample_filter); + } catch (const std::bad_cast&) { + } + + try { + auto& sample_filter = + dynamic_cast&>( + sample_filter_ref); + return search_with_filtering(handle, params, idx, queries, neighbors, distances, sample_filter); + } catch (const std::bad_cast&) { + RAFT_FAIL("Unsupported sample filter type"); + } +} + +} // namespace cuvs::neighbors::ivf_binary::detail diff --git a/cpp/src/neighbors/ivf_binary_index.cpp b/cpp/src/neighbors/ivf_binary_index.cpp new file mode 100644 index 0000000000..8cb98a06fa --- /dev/null +++ b/cpp/src/neighbors/ivf_binary_index.cpp @@ -0,0 +1,104 @@ +/* + * 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::neighbors::ivf_binary { +namespace { + +uint32_t encoded_length(uint32_t dim) +{ + RAFT_EXPECTS(dim > 0, "The dimension of a binary index must be positive"); + RAFT_EXPECTS(dim <= std::numeric_limits::max() / 8, + "The binary dimension is too large for the per-bit center statistics"); + return (dim + 7) / 8; +} + +template +void zero_fill(raft::resources const& res, T* ptr, size_t n) +{ + if (n != 0) { + RAFT_CUDA_TRY( + cudaMemsetAsync(ptr, 0, n * sizeof(T), raft::resource::get_cuda_stream(res).get())); + } +} + +} // namespace + +index::index(raft::resources const& res) + : dim_{0}, + adaptive_centers_{false}, + quantizer_{res}, + centers_{raft::make_device_matrix(res, 0, 0)}, + center_counts_{raft::make_device_matrix(res, 0, 0)}, + lists_{res} +{ +} + +index::index(raft::resources const& res, const index_params& params, uint32_t dim) + : dim_{dim}, + adaptive_centers_{params.adaptive_centers}, + quantizer_{res}, + centers_{raft::make_device_matrix(res, params.n_lists, encoded_length(dim))}, + center_counts_{raft::make_device_matrix( + res, + params.adaptive_centers ? params.n_lists : 0, + params.adaptive_centers ? int64_t{8} * encoded_length(dim) : 0)}, + lists_{res, + cuvs::distance::DistanceType::BitwiseHamming, + params.n_lists, + false, + params.conservative_memory_allocation, + encoded_length(dim)} +{ + zero_fill(res, center_counts_.data_handle(), center_counts_.size()); + // The lists reserve floating-point centers that are not used by the binary index. + zero_fill(res, lists_.centers().data_handle(), lists_.centers().size()); +} + +uint32_t index::dim() const noexcept { return dim_; } + +uint32_t index::encoded_dim() const noexcept { return static_cast(centers_.extent(1)); } + +uint32_t index::n_lists() const noexcept { return lists_.n_lists(); } + +int64_t index::size() const noexcept { return lists_.size(); } + +bool index::adaptive_centers() const noexcept { return adaptive_centers_; } + +bool index::conservative_memory_allocation() const noexcept +{ + return lists_.conservative_memory_allocation(); +} + +raft::device_vector_view index::list_sizes() const noexcept +{ + return lists_.list_sizes(); +} + +raft::device_matrix_view index::centers() const noexcept +{ + return centers_.view(); +} + +raft::device_matrix_view index::center_counts() + const noexcept +{ + return center_counts_.view(); +} + +const preprocessing::quantize::binary::quantizer& index::quantizer() const noexcept +{ + return quantizer_; +} + +} // namespace cuvs::neighbors::ivf_binary diff --git a/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/metric_impl.cuh b/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/metric_impl.cuh index 252daaf2aa..b9ea405f2e 100644 --- a/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/metric_impl.cuh +++ b/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/metric_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -48,4 +48,20 @@ __device__ void compute_dist_inner_product_impl(AccT& acc, AccT x, AccT y) } } +// Bitwise Hamming over byte-packed binary vectors (uint8_t data, uint32_t acc). +// `x` and `y` are already promoted to uint32_t by the load fragment: for Veclen >= 4 they pack +// 4 source bytes per accumulator word, for Veclen == 2 the upper 16 bits are zero from the +// uint16->uint32 zero-extension, and for Veclen == 1 only the low 8 bits carry a byte. +template +__device__ void compute_dist_bitwise_hamming_impl(AccT& acc, AccT x, AccT y) +{ + static_assert(std::is_same_v && std::is_same_v, + "compute_dist_bitwise_hamming_impl is only valid for uint8_t/uint32_t"); + if constexpr (Veclen > 1) { + acc += __popc(x ^ y); + } else { + acc += __popc(static_cast(x ^ y) & 0xffu); + } +} + } // namespace cuvs::neighbors::ivf_flat::detail diff --git a/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/metric_matrix.json b/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/metric_matrix.json index 5670ee0d84..5db309adb7 100644 --- a/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/metric_matrix.json +++ b/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/metric_matrix.json @@ -1,33 +1,47 @@ -{ - "metric_name": ["euclidean", "inner_product"], - "_data_type": [ - { - "data_type": "float", - "data_abbrev": "f", - "acc_type": "float", - "acc_abbrev": "f", - "veclen": ["1", "4"] - }, - { - "data_type": "__half", - "data_abbrev": "h", - "acc_type": "__half", - "acc_abbrev": "h", - "veclen": ["1", "8"] - }, - { - "data_type": "uint8_t", - "data_abbrev": "u8", - "acc_type": "uint32_t", - "acc_abbrev": "u32", - "veclen": ["1", "16"] - }, - { - "data_type": "int8_t", - "data_abbrev": "i8", - "acc_type": "int32_t", - "acc_abbrev": "i32", - "veclen": ["1", "16"] - } - ] -} +[ + { + "metric_name": ["euclidean", "inner_product"], + "_data_type": [ + { + "data_type": "float", + "data_abbrev": "f", + "acc_type": "float", + "acc_abbrev": "f", + "veclen": ["1", "4"] + }, + { + "data_type": "__half", + "data_abbrev": "h", + "acc_type": "__half", + "acc_abbrev": "h", + "veclen": ["1", "8"] + }, + { + "data_type": "uint8_t", + "data_abbrev": "u8", + "acc_type": "uint32_t", + "acc_abbrev": "u32", + "veclen": ["1", "16"] + }, + { + "data_type": "int8_t", + "data_abbrev": "i8", + "acc_type": "int32_t", + "acc_abbrev": "i32", + "veclen": ["1", "16"] + } + ] + }, + { + "metric_name": ["bitwise_hamming"], + "_data_type": [ + { + "data_type": "uint8_t", + "data_abbrev": "u8", + "acc_type": "uint32_t", + "acc_abbrev": "u32", + "veclen": ["1", "16"] + } + ] + } +] diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh index 294ae91f5b..6d6c1cab67 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh @@ -237,6 +237,8 @@ void extend(raft::resources const& handle, auto list_sizes_view = raft::make_device_vector_view, IdxT>( list_sizes_ptr, n_lists); + vec_batches.reset(); + vec_batches.prefetch_next_batch(); for (const auto& batch : vec_batches) { auto batch_data_view = raft::make_device_matrix_view(batch.data(), batch.size(), index->dim()); @@ -249,6 +251,9 @@ void extend(raft::resources const& handle, list_sizes_view, false, utils::mapping{}); + vec_batches.prefetch_next_batch(); + // The copy stream must not overwrite a batch while the main stream still reads it. + if (enable_prefetch) { raft::resource::sync_stream(handle); } } } else { raft::stats::histogram(raft::stats::HistTypeAuto, diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh index 48aa62501c..2b98316851 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh @@ -313,6 +313,21 @@ void launch_with_fixed_consts(cuvs::distance::DistanceType metric, Args&&... arg tag_post_process_compose>( std::forward(args)...); // NB: update the description of `knn::ivf_flat::build` when // adding here a new metric. + case cuvs::distance::DistanceType::BitwiseHamming: + if constexpr (std::is_same_v) { + return launch_kernel(std::forward(args)...); + } else { + RAFT_FAIL("BitwiseHamming distance is only supported with uint8_t data type"); + } case cuvs::distance::DistanceType::CustomUDF: return launch_kernel& index, cuda::stream_ref stream, const std::optional& metric_udf) { + if (metric == cuvs::distance::DistanceType::BitwiseHamming && !std::is_same_v) { + RAFT_FAIL("BitwiseHamming distance is only supported with uint8_t data type, got %s", + typeid(T).name()); + } + const uint32_t n_probes_clamped = std::min(n_probes, index.n_lists()); const int capacity = raft::bound_by_power_of_two(k); diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 38758dec90..f5b32007a7 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -178,6 +178,13 @@ ConfigureTest( PERCENT 100 ) +ConfigureTest( + NAME NEIGHBORS_ANN_IVF_BINARY_TEST + PATH neighbors/ann_ivf_binary.cu + GPUS 1 + PERCENT 100 +) + ConfigureTest( NAME NEIGHBORS_ANN_IVF_SQ_TEST PATH neighbors/ann_ivf_sq/test_float_int64_t.cu neighbors/ann_ivf_sq/test_half_int64_t.cu diff --git a/cpp/tests/cluster/kmeans_balanced.cu b/cpp/tests/cluster/kmeans_balanced.cu index 777cb82b41..3d44444475 100644 --- a/cpp/tests/cluster/kmeans_balanced.cu +++ b/cpp/tests/cluster/kmeans_balanced.cu @@ -3,6 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include "../../src/cluster/detail/kmeans_balanced.cuh" #include "../test_utils.h" #include @@ -297,4 +298,199 @@ KB_TEST((KmeansBalancedTest, KmeansBalancedTestFI8I32I32_SEP, inputsf_i32); +TEST(KmeansBalancedBinary, PredictMatchesExpandedInput) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int64_t n_rows = 97; + constexpr int64_t n_clusters = 7; + for (int64_t packed_dim : {int64_t{1}, int64_t{3}, int64_t{192}}) { + SCOPED_TRACE(packed_dim); + const int64_t dim = packed_dim * 8; + std::vector packed(n_rows * packed_dim); + std::vector expanded(n_rows * dim); + std::vector centers(n_clusters * dim); + for (size_t i = 0; i < packed.size(); ++i) { + packed[i] = uint8_t((i * 73 + i / 7 + 19) % 256); + } + for (size_t i = 0; i < expanded.size(); ++i) { + expanded[i] = ((packed[i / 8] >> (i % 8)) & 1) ? 1.0f : -1.0f; + } + for (size_t i = 0; i < centers.size(); ++i) { + centers[i] = (int((i * 19 + i / dim * 5) % 31) - 15) / 16.0f; + } + auto X = raft::make_device_matrix(handle, n_rows, packed_dim); + auto X_expanded = raft::make_device_matrix(handle, n_rows, dim); + auto C = raft::make_device_matrix(handle, n_clusters, dim); + auto labels = raft::make_device_vector(handle, n_rows); + auto expected = raft::make_device_vector(handle, n_rows); + raft::update_device(X.data_handle(), packed.data(), packed.size(), stream); + raft::update_device(X_expanded.data_handle(), expanded.data(), expanded.size(), stream); + raft::update_device(C.data_handle(), centers.data(), centers.size(), stream); + for (auto metric : {cuvs::distance::DistanceType::L2Expanded, + cuvs::distance::DistanceType::L2SqrtExpanded, + cuvs::distance::DistanceType::InnerProduct}) { + SCOPED_TRACE(int(metric)); + cuvs::cluster::kmeans::balanced_params params; + params.metric = metric; + params.is_packed_binary = true; + cuvs::cluster::kmeans::predict(handle, + params, + raft::make_const_mdspan(X.view()), + raft::make_const_mdspan(C.view()), + labels.view()); + params.is_packed_binary = false; + cuvs::cluster::kmeans::predict(handle, + params, + raft::make_const_mdspan(X_expanded.view()), + raft::make_const_mdspan(C.view()), + expected.view()); + std::vector actual_labels(n_rows), expected_labels(n_rows); + raft::update_host(actual_labels.data(), labels.data_handle(), n_rows, stream); + raft::update_host(expected_labels.data(), expected.data_handle(), n_rows, stream); + raft::resource::sync_stream(handle); + EXPECT_EQ(actual_labels, expected_labels); + } + } +} + +TEST(KmeansBalancedBinary, HierarchicalCentersPreserveEveryExpandedCoordinate) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int64_t n_rows = 4096, packed_dim = 32, dim = 256, n_clusters = 256; + std::vector packed(n_rows * packed_dim); + for (int64_t row = 0; row < n_rows; ++row) { + // Initial mesocluster labels select the low four bits of the prototype ID. + // Within each mesocluster, fine-cluster labels select its high four bits. + // Every true centroid therefore has an exact {-1,+1} CPU oracle, with no + // dependence on centroid ordering or on private library symbols. + for (int64_t byte = 0; byte < packed_dim; ++byte) { + packed[row * packed_dim + byte] = (row & (int64_t{1} << (byte / 4))) ? 0xff : 0; + } + } + auto X = raft::make_device_matrix(handle, n_rows, packed_dim); + auto C = raft::make_device_matrix(handle, n_clusters, dim); + auto labels = raft::make_device_vector(handle, n_rows); + raft::update_device(X.data_handle(), packed.data(), packed.size(), stream); + raft::matrix::fill(handle, C.view(), 123.0f); + cuvs::cluster::kmeans::balanced_params params; + params.is_packed_binary = true; + params.n_iters = 1; + cuvs::cluster::kmeans::fit(handle, params, raft::make_const_mdspan(X.view()), C.view()); + cuvs::cluster::kmeans::predict(handle, + params, + raft::make_const_mdspan(X.view()), + raft::make_const_mdspan(C.view()), + labels.view()); + std::vector centers(n_clusters * dim); + std::vector actual_labels(n_rows); + raft::update_host(centers.data(), C.data_handle(), centers.size(), stream); + raft::update_host(actual_labels.data(), labels.data_handle(), actual_labels.size(), stream); + raft::resource::sync_stream(handle); + std::vector seen(n_clusters, false); + for (int64_t row = 0; row < n_rows; ++row) { + ASSERT_LT(actual_labels[row], n_clusters); + seen[actual_labels[row]] = true; + for (int64_t bit = 0; bit < dim; ++bit) { + const auto byte = packed[row * packed_dim + bit / 8]; + ASSERT_EQ(centers[actual_labels[row] * dim + bit], ((byte >> (bit % 8)) & 1) ? 1.0f : -1.0f) + << "row=" << row << ", bit=" << bit; + } + } + EXPECT_EQ(std::count(seen.begin(), seen.end(), true), n_clusters); +} + +TEST(KmeansBalancedBinary, FitAndPredictRecoverPackedClusters) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int64_t n_rows = 128, packed_dim = 3, dim = 24, n_clusters = 4; + std::vector packed(n_rows * packed_dim); + for (int64_t row = 0; row < n_rows; ++row) { + packed[row * packed_dim] = (row & 1) ? 0xff : 0; + packed[row * packed_dim + 1] = (row & 2) ? 0xff : 0; + packed[row * packed_dim + 2] = (row & 2) ? 0xa5 : 0x5a; + } + auto X = raft::make_device_matrix(handle, n_rows, packed_dim); + auto C = raft::make_device_matrix(handle, n_clusters, dim); + auto labels = raft::make_device_vector(handle, n_rows); + raft::update_device(X.data_handle(), packed.data(), packed.size(), stream); + cuvs::cluster::kmeans::balanced_params params; + params.is_packed_binary = true; + params.n_iters = 2; + cuvs::cluster::kmeans::fit(handle, params, raft::make_const_mdspan(X.view()), C.view()); + cuvs::cluster::kmeans::predict(handle, + params, + raft::make_const_mdspan(X.view()), + raft::make_const_mdspan(C.view()), + labels.view()); + std::vector actual_labels(n_rows); + std::vector centers(n_clusters * dim); + raft::update_host(actual_labels.data(), labels.data_handle(), n_rows, stream); + raft::update_host(centers.data(), C.data_handle(), centers.size(), stream); + raft::resource::sync_stream(handle); + for (int64_t row = 0; row < n_rows; ++row) { + ASSERT_LT(actual_labels[row], n_clusters); + for (int64_t bit = 0; bit < dim; ++bit) { + const auto byte = packed[row * packed_dim + bit / 8]; + EXPECT_EQ(centers[actual_labels[row] * dim + bit], ((byte >> (bit % 8)) & 1) ? 1.0f : -1.0f); + } + } +} + +TEST(KmeansBalancedBinary, IncrementalCentersMatchSinglePass) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int64_t n_rows = 9, packed_dim = 3, dim = 24, n_clusters = 3; + std::vector packed(n_rows * packed_dim); + std::vector labels(n_rows); + for (size_t i = 0; i < packed.size(); ++i) { + packed[i] = uint8_t(i * 73 + 19); + } + for (int64_t row = 0; row < n_rows; ++row) { + labels[row] = row % n_clusters; + } + auto X = raft::make_device_matrix(handle, n_rows, packed_dim); + auto L = raft::make_device_vector(handle, n_rows); + auto C = raft::make_device_matrix(handle, n_clusters, dim); + auto sizes = raft::make_device_vector(handle, n_clusters); + raft::update_device(X.data_handle(), packed.data(), packed.size(), stream); + raft::update_device(L.data_handle(), labels.data(), labels.size(), stream); + auto update = [&](int64_t offset, int64_t rows, bool reset) { + cuvs::cluster::kmeans::detail::calc_centers_and_sizes( + handle, + C.data_handle(), + sizes.data_handle(), + n_clusters, + packed_dim, + X.data_handle() + offset * packed_dim, + rows, + L.data_handle() + offset, + reset, + true, + raft::identity_op{}, + raft::resource::get_workspace_resource_ref(handle)); + }; + update(0, 4, true); + update(4, 5, false); + std::vector actual(n_clusters * dim); + std::vector actual_sizes(n_clusters); + raft::update_host(actual.data(), C.data_handle(), actual.size(), stream); + raft::update_host(actual_sizes.data(), sizes.data_handle(), actual_sizes.size(), stream); + raft::resource::sync_stream(handle); + for (int64_t cluster = 0; cluster < n_clusters; ++cluster) { + EXPECT_EQ(actual_sizes[cluster], 3); + for (int64_t bit = 0; bit < dim; ++bit) { + float expected = 0; + for (int64_t row = cluster; row < n_rows; row += n_clusters) { + const auto byte = packed[row * packed_dim + bit / 8]; + expected += ((byte >> (bit % 8)) & 1) ? 1.0f : -1.0f; + } + EXPECT_NEAR(actual[cluster * dim + bit], expected / 3.0f, 1e-6f); + } + } +} + } // namespace cuvs diff --git a/cpp/tests/neighbors/ann_ivf_binary.cu b/cpp/tests/neighbors/ann_ivf_binary.cu new file mode 100644 index 0000000000..fcbb450f92 --- /dev/null +++ b/cpp/tests/neighbors/ann_ivf_binary.cu @@ -0,0 +1,441 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../test_utils.cuh" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::ivf_binary { +namespace { + +namespace binary = cuvs::preprocessing::quantize::binary; + +auto zero_threshold() -> binary_dictionary_view +{ + return {raft::make_device_vector_view(nullptr, 0)}; +} + +template +auto to_device(raft::resources const& handle, + const std::vector& host, + int64_t n_rows, + int64_t dim) +{ + auto out = raft::make_device_matrix(handle, n_rows, dim); + raft::update_device( + out.data_handle(), host.data(), host.size(), raft::resource::get_cuda_stream(handle)); + return out; +} + +template +auto to_host(raft::resources const& handle, const T* device, size_t n) +{ + std::vector out(n); + raft::update_host(out.data(), device, n, raft::resource::get_cuda_stream(handle)); + raft::resource::sync_stream(handle); + return out; +} + +// Codes with the zero threshold, 8 bits per byte +auto zero_threshold_dataset(const raft::device_matrix& codes) + -> device_binary_dataset_view +{ + return {raft::make_const_mdspan(codes.view()), zero_threshold(), uint32_t(codes.extent(1) * 8)}; +} + +auto hamming(const uint8_t* a, const uint8_t* b, int64_t n_bytes) -> float +{ + int d = 0; + for (int64_t i = 0; i < n_bytes; ++i) { + d += __builtin_popcount(unsigned(a[i] ^ b[i])); + } + return float(d); +} + +index_params small_params(uint32_t n_lists) +{ + index_params params; + params.n_lists = n_lists; + params.kmeans_n_iters = 3; + params.kmeans_trainset_fraction = 1.0; + return params; +} + +TEST(IvfBinary, ExhaustiveSearchMatchesBruteForce) +{ + raft::resources handle; + constexpr int64_t n_rows = 1536; + constexpr int64_t n_queries = 3; + std::mt19937 rng(42); + std::normal_distribution normal(0.5f, 1.0f); + for (uint32_t dim : {1u, 3u, 16u, 100u, 192u}) { + for (auto threshold : {binary::bit_threshold::zero, binary::bit_threshold::mean}) { + SCOPED_TRACE(dim); + const int64_t n_bytes = (dim + 7) / 8; + std::vector data(n_rows * dim); + std::vector queries(n_queries * dim); + for (auto& x : data) { + x = normal(rng); + } + for (auto& x : queries) { + x = normal(rng); + } + auto data_dev = to_device(handle, data, n_rows, dim); + auto queries_dev = to_device(handle, queries, n_queries, dim); + + binary::params quantizer_params; + quantizer_params.threshold = threshold; + auto quantizer = + binary::train(handle, quantizer_params, raft::make_const_mdspan(data_dev.view())); + auto codes = raft::make_device_matrix(handle, n_rows, n_bytes); + auto query_codes = raft::make_device_matrix(handle, n_queries, n_bytes); + binary::transform(handle, quantizer, raft::make_const_mdspan(data_dev.view()), codes.view()); + auto dataset = make_device_binary_dataset_view( + raft::make_const_mdspan(codes.view()), quantizer, dim); + + auto params = small_params(16); + auto idx = build(handle, params, dataset); + ASSERT_EQ(idx.dim(), dim); + ASSERT_EQ(idx.encoded_dim(), n_bytes); + ASSERT_EQ(idx.size(), n_rows); + + auto codes_host = to_host(handle, codes.data_handle(), codes.size()); + + search_params search_params; + search_params.n_probes = params.n_lists; + // Query codes may be produced using a different quantizer from the build dataset. + for (auto query_threshold : {binary::bit_threshold::zero, binary::bit_threshold::mean}) { + binary::params query_quantizer_params; + query_quantizer_params.threshold = query_threshold; + auto query_quantizer = + binary::train(handle, query_quantizer_params, raft::make_const_mdspan(data_dev.view())); + binary::transform( + handle, query_quantizer, raft::make_const_mdspan(queries_dev.view()), query_codes.view()); + auto query_codes_host = to_host(handle, query_codes.data_handle(), query_codes.size()); + // Include a k larger than the local top-k capacity. + for (int64_t k : {10, 100, 1025}) { + SCOPED_TRACE(k); + auto ids_dev = raft::make_device_matrix(handle, n_queries, k); + auto distances_dev = raft::make_device_matrix(handle, n_queries, k); + search(handle, + search_params, + idx, + raft::make_const_mdspan(query_codes.view()), + ids_dev.view(), + distances_dev.view()); + auto ids = to_host(handle, ids_dev.data_handle(), ids_dev.size()); + auto distances = to_host(handle, distances_dev.data_handle(), distances_dev.size()); + for (int64_t q = 0; q < n_queries; ++q) { + std::vector expected(n_rows); + for (int64_t row = 0; row < n_rows; ++row) { + expected[row] = + hamming(&codes_host[row * n_bytes], &query_codes_host[q * n_bytes], n_bytes); + } + auto sorted = expected; + std::sort(sorted.begin(), sorted.end()); + std::vector seen(n_rows); + std::vector actual(distances.begin() + q * k, distances.begin() + (q + 1) * k); + std::sort(actual.begin(), actual.end()); + for (int64_t j = 0; j < k; ++j) { + // Compare distances rather than ids: the k-th neighbor often has many ties + auto id = ids[q * k + j]; + ASSERT_GE(id, 0); + ASSERT_LT(id, n_rows); + ASSERT_FALSE(seen[id]); + seen[id] = true; + ASSERT_EQ(distances[q * k + j], expected[id]); + ASSERT_EQ(actual[j], sorted[j]); + } + } + } + } + } + } +} + +TEST(IvfBinary, SerializationRoundTrip) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int64_t n_rows = 512; + constexpr int64_t n_bytes = 13; + constexpr int64_t n_queries = 4; + std::mt19937 rng(7); + std::vector data(n_rows * n_bytes); + std::vector queries(n_queries * n_bytes); + for (auto& x : data) { + x = static_cast(rng()); + } + for (auto& x : queries) { + x = static_cast(rng()); + } + const uint32_t dim = 100; + // The unused high bits of the final byte are not part of the logical vector. + for (int64_t row = 0; row < n_rows; ++row) { + data[(row + 1) * n_bytes - 1] &= (1u << (dim % 8)) - 1; + } + for (int64_t row = 0; row < n_queries; ++row) { + queries[(row + 1) * n_bytes - 1] &= (1u << (dim % 8)) - 1; + } + auto codes = to_device(handle, data, n_rows, n_bytes); + auto query_codes = to_device(handle, queries, n_queries, n_bytes); + const device_binary_dataset_view dataset{ + raft::make_const_mdspan(codes.view()), zero_threshold(), dim}; + auto encoded_queries = raft::make_const_mdspan(query_codes.view()); + + auto params = small_params(8); + params.adaptive_centers = true; + auto idx = build(handle, params, dataset); + + std::stringstream bytes; + serialize(handle, bytes, idx); + index restored(handle); + deserialize(handle, bytes, &restored); + ASSERT_EQ(restored.dim(), dim); + ASSERT_EQ(restored.size(), n_rows); + ASSERT_EQ(restored.n_lists(), params.n_lists); + ASSERT_TRUE(restored.adaptive_centers()); + ASSERT_TRUE(cuvs::devArrMatch(idx.centers().data_handle(), + restored.centers().data_handle(), + idx.centers().size(), + cuvs::Compare(), + stream.get())); + ASSERT_TRUE(cuvs::devArrMatch(idx.center_counts().data_handle(), + restored.center_counts().data_handle(), + idx.center_counts().size(), + cuvs::Compare(), + stream.get())); + + constexpr int64_t k = 10; + search_params search_params; + search_params.n_probes = params.n_lists; + auto distances_a = raft::make_device_matrix(handle, n_queries, k); + auto distances_b = raft::make_device_matrix(handle, n_queries, k); + auto ids = raft::make_device_matrix(handle, n_queries, k); + search(handle, search_params, idx, encoded_queries, ids.view(), distances_a.view()); + search(handle, search_params, restored, encoded_queries, ids.view(), distances_b.view()); + ASSERT_TRUE(cuvs::devArrMatch(distances_a.data_handle(), + distances_b.data_handle(), + distances_a.size(), + cuvs::Compare(), + stream.get())); +} + +TEST(IvfBinary, AdaptiveCentersKeepExactCounts) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + auto data = to_device(handle, std::vector{1, 1, 1, 0}, 4, 1); + auto params = small_params(1); + params.adaptive_centers = true; + auto idx = build(handle, params, zero_threshold_dataset(data)); + ASSERT_EQ(idx.center_counts().extent(1), 8); + ASSERT_EQ(to_host(handle, idx.centers().data_handle(), 1)[0], 1); + ASSERT_EQ(to_host(handle, idx.center_counts().data_handle(), 1)[0], 3); + + std::stringstream bytes; + serialize(handle, bytes, idx); + index restored(handle); + deserialize(handle, bytes, &restored); + + auto more = to_device(handle, std::vector{0, 0, 0}, 3, 1); + auto ids = raft::make_device_vector(handle, 3); + const std::vector ids_host{4, 5, 6}; + raft::update_device(ids.data_handle(), ids_host.data(), ids_host.size(), stream); + extend(handle, + zero_threshold_dataset(more), + std::make_optional(raft::make_const_mdspan(ids.view())), + &restored); + EXPECT_EQ(restored.size(), 7); + EXPECT_EQ(to_host(handle, restored.centers().data_handle(), 1)[0], 0); // 3 ones out of 7 + EXPECT_EQ(to_host(handle, restored.center_counts().data_handle(), 1)[0], 3); + EXPECT_EQ(idx.size(), 4); +} + +TEST(IvfBinary, AdaptiveExactTieAfterTrainingOnly) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + // mean = 7 / 13 rounds upward in float; reconstructing its sum would make a tie positive. + std::vector initial(13); + std::fill(initial.begin(), initial.begin() + 10, 1); + auto data = to_device(handle, initial, 13, 1); + auto params = small_params(1); + params.adaptive_centers = true; + params.add_data_on_build = false; + auto idx = build(handle, params, zero_threshold_dataset(data)); + EXPECT_EQ(idx.size(), 0); + EXPECT_EQ(to_host(handle, idx.center_counts().data_handle(), 8), std::vector(8, 0)); + + extend(handle, zero_threshold_dataset(data), std::nullopt, &idx); + auto more = to_device(handle, std::vector(7, 0), 7, 1); + auto ids = raft::make_device_vector(handle, 7); + const std::vector ids_host{13, 14, 15, 16, 17, 18, 19}; + raft::update_device(ids.data_handle(), ids_host.data(), ids_host.size(), stream); + extend(handle, + zero_threshold_dataset(more), + std::make_optional(raft::make_const_mdspan(ids.view())), + &idx); + EXPECT_EQ(idx.size(), 20); + EXPECT_EQ(to_host(handle, idx.center_counts().data_handle(), 1)[0], 10); + EXPECT_EQ(to_host(handle, idx.centers().data_handle(), 1)[0], 0); // exact ties are zeros +} + +TEST(IvfBinary, RejectsMismatchedInputs) +{ + raft::resources handle; + auto data = to_device(handle, std::vector(32, 0x55), 32, 1); + auto idx = build(handle, small_params(1), zero_threshold_dataset(data)); + auto ids = raft::make_device_matrix(handle, 1, 1); + auto dist = raft::make_device_matrix(handle, 1, 1); + + // The row length of the codes must match the dimension + EXPECT_THROW((device_binary_dataset_view{ + raft::make_const_mdspan(data.view()), zero_threshold(), 17}), + raft::logic_error); + + // Reject a different logical dimension, even when the encoded byte width matches. + const device_binary_dataset_view other_dim{ + raft::make_const_mdspan(data.view()), zero_threshold(), 7}; + EXPECT_THROW(extend(handle, other_dim, std::nullopt, &idx), raft::logic_error); + + // Query width is measured in packed bytes. + auto queries = raft::make_device_matrix(handle, 1, 2); + EXPECT_THROW(search(handle, + search_params{}, + idx, + raft::make_const_mdspan(queries.view()), + ids.view(), + dist.view()), + raft::logic_error); +} + +TEST(IvfBinary, ExtendAcceptsDifferentThresholds) +{ + raft::resources handle; + auto data = to_device(handle, std::vector{1, 3}, 2, 1); + auto build_threshold = raft::make_device_vector(handle, 8); + raft::matrix::fill(handle, build_threshold.view(), 0.0f); + const device_binary_dataset_view dataset{ + raft::make_const_mdspan(data.view()), + binary_dictionary_view{raft::make_const_mdspan(build_threshold.view())}, + 8}; + auto params = small_params(1); + params.adaptive_centers = true; + auto idx = build(handle, params, dataset); + + auto more = to_device(handle, std::vector{0xf0}, 1, 1); + auto new_threshold = raft::make_device_vector(handle, 8); + raft::matrix::fill(handle, new_threshold.view(), 1.0f); + const device_binary_dataset_view new_vectors{ + raft::make_const_mdspan(more.view()), + binary_dictionary_view{raft::make_const_mdspan(new_threshold.view())}, + 8}; + auto new_ids = to_device(handle, std::vector{2}, 1, 1); + extend(handle, + new_vectors, + std::make_optional( + raft::make_device_vector_view(new_ids.data_handle(), 1)), + &idx); + ASSERT_EQ(idx.size(), 3); + EXPECT_EQ(to_host(handle, idx.center_counts().data_handle(), 8), + (std::vector{2, 1, 0, 0, 1, 1, 1, 1})); + // Extending does not replace the quantizer metadata captured at build time. + EXPECT_EQ(to_host(handle, idx.quantizer().threshold.data_handle(), 8), + std::vector(8, 0.0f)); + + auto ids = raft::make_device_matrix(handle, 1, 3); + auto distances = raft::make_device_matrix(handle, 1, 3); + search_params search_params; + search_params.n_probes = 1; + search( + handle, search_params, idx, raft::make_const_mdspan(more.view()), ids.view(), distances.view()); + auto result_ids = to_host(handle, ids.data_handle(), 3); + auto result_distances = to_host(handle, distances.data_handle(), 3); + const std::vector expected{5, 6, 0}; + for (int64_t i = 0; i < 3; ++i) { + ASSERT_GE(result_ids[i], 0); + ASSERT_LT(result_ids[i], 3); + EXPECT_EQ(result_distances[i], expected[result_ids[i]]); + } + std::sort(result_ids.begin(), result_ids.end()); + EXPECT_EQ(result_ids, (std::vector{0, 1, 2})); +} + +TEST(IvfBinary, SearchHonorsSampleFilter) +{ + raft::resources handle; + constexpr int64_t n_rows = 1536; + constexpr int64_t n_bytes = 8; + constexpr int64_t n_removed = 512; + constexpr int64_t k = 10; + std::mt19937 rng(3); + std::vector data(n_rows * n_bytes); + std::vector query(n_bytes); + for (auto& x : data) { + x = static_cast(rng()); + } + for (auto& x : query) { + x = static_cast(rng()); + } + auto codes = to_device(handle, data, n_rows, n_bytes); + auto query_codes = to_device(handle, query, 1, n_bytes); + auto params = small_params(16); + auto idx = build(handle, params, zero_threshold_dataset(codes)); + + auto removed = raft::make_device_vector(handle, n_removed); + raft::linalg::map_offset(handle, removed.view(), raft::identity_op{}); + raft::resource::sync_stream(handle); + cuvs::core::bitset removed_bitset(handle, removed.view(), n_rows); + auto filter = cuvs::neighbors::filtering::bitset_filter(removed_bitset.view()); + + search_params search_params; + search_params.n_probes = params.n_lists; + auto ids_dev = raft::make_device_matrix(handle, 1, k); + auto distances_dev = raft::make_device_matrix(handle, 1, k); + search(handle, + search_params, + idx, + raft::make_const_mdspan(query_codes.view()), + ids_dev.view(), + distances_dev.view(), + filter); + auto ids = to_host(handle, ids_dev.data_handle(), k); + auto distances = to_host(handle, distances_dev.data_handle(), k); + + std::vector expected; + for (int64_t row = n_removed; row < n_rows; ++row) { + expected.push_back(hamming(&data[row * n_bytes], query.data(), n_bytes)); + } + std::sort(expected.begin(), expected.end()); + for (int64_t j = 0; j < k; ++j) { + ASSERT_GE(ids[j], n_removed); + ASSERT_EQ(distances[j], hamming(&data[ids[j] * n_bytes], query.data(), n_bytes)); + } + std::sort(distances.begin(), distances.end()); + for (int64_t j = 0; j < k; ++j) { + ASSERT_EQ(distances[j], expected[j]); + } +} + +} // namespace +} // namespace cuvs::neighbors::ivf_binary diff --git a/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu b/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu index fee22db3c5..8812b242d5 100644 --- a/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu +++ b/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu @@ -7,13 +7,17 @@ #include "../ann_ivf_flat.cuh" +#include #include +#include #include #include +#include #include #include #include +#include namespace cuvs::neighbors::ivf_flat { @@ -27,6 +31,55 @@ TEST_P(AnnIVFFlatTestF_float, AnnIVFFlat) INSTANTIATE_TEST_CASE_P(AnnIVFFlatTest, AnnIVFFlatTestF_float, ::testing::ValuesIn(inputs)); +TEST(AnnIVFFlatTest, AdaptiveCentersWithPrefetchedHostBatches) +{ + raft::resources handle; + // Avoid implicit cudaFree synchronization hiding staging-buffer reuse races. + raft::resource::set_workspace_to_pool_resource(handle, 64 * 1024 * 1024); + raft::resource::set_cuda_stream_pool(handle, std::make_shared(1)); + auto stream = raft::resource::get_cuda_stream(handle); + + // Cross multiple staging-buffer boundaries, including a partial final batch. Distinct values + // in each batch expose overwrites by the copy stream during center accumulation. + constexpr int64_t batch_rows = 65536; + constexpr int64_t rows = 3 * batch_rows + 17; + constexpr int64_t dim = 32; + // Pinned host memory permits the next H2D copy to overlap the center-update kernel. + auto dataset = raft::make_pinned_matrix(handle, rows, dim); + std::vector sums(dim, 0.0); + for (int64_t row = 0; row < rows; ++row) { + for (int64_t col = 0; col < dim; ++col) { + auto value = float((row / batch_rows + 1) * (col + 1)); + dataset(row, col) = value; + sums[col] += value; + } + } + + index_params params; + params.n_lists = 1; + params.adaptive_centers = true; + params.add_data_on_build = false; + params.kmeans_trainset_fraction = 1.0; + auto training_view = + raft::make_host_matrix_view(dataset.data_handle(), 1, dim); + auto idx = build(handle, params, training_view); + auto data_view = + raft::make_host_matrix_view(dataset.data_handle(), rows, dim); + extend(handle, data_view, std::nullopt, &idx); + + std::vector center(dim); + uint32_t list_size = 0; + raft::update_host(center.data(), idx.centers().data_handle(), dim, stream); + raft::update_host(&list_size, idx.list_sizes().data_handle(), 1, stream); + raft::resource::sync_stream(handle); + + EXPECT_EQ(idx.size(), rows); + EXPECT_EQ(list_size, rows); + for (int64_t col = 0; col < dim; ++col) { + EXPECT_NEAR(center[col], sums[col] / rows, 1e-4) << "column " << col; + } +} + TEST(IVFFlatSerialization, StreamRoundTrip) { raft::resources handle; diff --git a/cpp/tests/neighbors/ann_utils.cuh b/cpp/tests/neighbors/ann_utils.cuh index 346972d79d..35e940dc6f 100644 --- a/cpp/tests/neighbors/ann_utils.cuh +++ b/cpp/tests/neighbors/ann_utils.cuh @@ -111,6 +111,7 @@ inline auto operator<<(std::ostream& os, const print_metric& p) -> std::ostream& break; case cuvs::distance::DistanceType::DiceExpanded: os << "distance::DiceExpanded"; break; case cuvs::distance::DistanceType::Precomputed: os << "distance::Precomputed"; break; + case cuvs::distance::DistanceType::BitwiseHamming: os << "distance::BitwiseHamming"; break; default: RAFT_FAIL("unreachable code"); } return os; diff --git a/fern/docs.yml b/fern/docs.yml index 77ee0dae2a..83f08cbd11 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -335,6 +335,8 @@ navigation: path: "./pages/cpp_api/cpp-api-cluster-kmeans.md" - page: "Cluster Spectral" path: "./pages/cpp_api/cpp-api-cluster-spectral.md" + - page: "Core Dataset Binary" + path: "./pages/cpp_api/cpp-api-core-dataset-binary.md" - page: "Core Bloom Filter" path: "./pages/cpp_api/cpp-api-core-bloom-filter.md" - section: "Common Types" @@ -382,6 +384,8 @@ navigation: path: "./pages/cpp_api/cpp-api-neighbors-hnsw.md" - page: "Neighbors Composite Index" path: "./pages/cpp_api/cpp-api-neighbors-composite-index.md" + - page: "Neighbors IVF Binary" + path: "./pages/cpp_api/cpp-api-neighbors-ivf-binary.md" - page: "Neighbors IVF Flat" path: "./pages/cpp_api/cpp-api-neighbors-ivf-flat.md" - page: "Neighbors IVF PQ" diff --git a/fern/pages/cpp_api/cpp-api-cluster-kmeans.md b/fern/pages/cpp_api/cpp-api-cluster-kmeans.md index 854db8c317..c05c2fa0dd 100644 --- a/fern/pages/cpp_api/cpp-api-cluster-kmeans.md +++ b/fern/pages/cpp_api/cpp-api-cluster-kmeans.md @@ -423,6 +423,8 @@ raft::device_matrix_view centroids, std::optional> inertia = std::nullopt); ``` +**Note:** When `params.is_packed_binary` is true, `X.extent(1)` counts packed bytes,
and centroids must have `8 * X.extent(1)` floating-point coordinates. Bits are
expanded least-significant bit first to \{-1, +1\}; the selected metric operates
on those expanded vectors. CosineExpanded is not supported in packed binary mode.
With the flag disabled, uint8_t values are numeric. + **Parameters** | Name | Direction | Type | Description | @@ -708,6 +710,8 @@ raft::device_matrix_view centroids, raft::device_vector_view labels); ``` +**Note:** When `params.is_packed_binary` is true, `X.extent(1)` counts packed bytes,
and centroids must have `8 * X.extent(1)` floating-point coordinates. Bits are
expanded least-significant bit first to \{-1, +1\}; the selected metric operates
on those expanded vectors. CosineExpanded is not supported in packed binary mode.
With the flag disabled, uint8_t values are numeric. + **Parameters** | Name | Direction | Type | Description | diff --git a/fern/pages/cpp_api/cpp-api-core-dataset-binary.md b/fern/pages/cpp_api/cpp-api-core-dataset-binary.md new file mode 100644 index 0000000000..04ac158aca --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-core-dataset-binary.md @@ -0,0 +1,39 @@ +--- +slug: api-reference/cpp-api-core-dataset-binary +--- + +# Binary + +_Source header: `cuvs/core/dataset/binary.hpp`_ + +## Types + + +### binary_dictionary_view + +Non-owning view of the dictionary that binarizes a dataset: the per-dimension thresholds of a binary quantizer. An empty threshold vector stands for the zero threshold. + +```cpp +template +struct binary_dictionary_view { + raft::device_vector_view threshold; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `threshold` | `raft::device_vector_view` | | + + +### device_binary_dataset_view + +Non-owning view of a binary-quantized dataset. + +The data are bit-packed codes, `ceil(dim / 8)` bytes per row; bit `i` of a row is stored in bit `i % 8` of byte `i / 8`, and the unused bits of the last byte are zero. The dictionary is the quantizer that produced the codes, so the logical dimension `dim` (the number of bits) never has to be guessed from the padded byte width. + +```cpp +template +struct device_binary_dataset_view; +``` diff --git a/fern/pages/cpp_api/cpp-api-neighbors-ivf-binary.md b/fern/pages/cpp_api/cpp-api-neighbors-ivf-binary.md new file mode 100644 index 0000000000..fcc5bc1e3f --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-neighbors-ivf-binary.md @@ -0,0 +1,383 @@ +--- +slug: api-reference/cpp-api-neighbors-ivf-binary +--- + +# IVF Binary + +_Source header: `cuvs/neighbors/ivf_binary.hpp`_ + +## IVF-Binary index build parameters + + +### neighbors::ivf_binary::index_params + +IVF-Binary index build parameters + +```cpp +struct index_params { + uint32_t n_lists; + uint32_t kmeans_n_iters; + double kmeans_trainset_fraction; + bool adaptive_centers; + bool conservative_memory_allocation; + bool add_data_on_build; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `n_lists` | `uint32_t` | The number of inverted lists (clusters). | +| `kmeans_n_iters` | `uint32_t` | The number of iterations searching for k-means centers. | +| `kmeans_trainset_fraction` | `double` | The fraction of data to use during iterative k-means building. | +| `adaptive_centers` | `bool` | By default (adaptive_centers = false), the cluster centers are trained in `build` and never modified in `extend`.

With adaptive_centers = true, `extend` updates the centers to the majority bit of the vectors added to each cluster, with ties resolved to 0. The index keeps exact per-bit counts of ones to make the update independent of how the data are split between `extend` calls. | +| `conservative_memory_allocation` | `bool` | Allocate the minimum amount of memory required for the inverted lists, instead of reserving extra space to amortize repeated calls of `extend`. | +| `add_data_on_build` | `bool` | Whether to add the dataset content to the index in `build`. If false, `build` only trains the clustering and the index is populated later with `extend`. | + +## IVF-Binary index search parameters + + +### neighbors::ivf_binary::search_params + +IVF-Binary index search parameters + +```cpp +struct search_params : cuvs::neighbors::search_params { + uint32_t n_probes; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `n_probes` | `uint32_t` | The number of clusters to search. | + +## IVF-Binary index + + +### neighbors::ivf_binary::index + +IVF index of bit-packed binary vectors searched with the Hamming distance. + +The index stores the codes of a binary-quantized dataset (see `device_binary_dataset_view`), together with the quantizer that produced them. Queries must be supplied as bit-packed bytes. Neighbors are returned with their Hamming distances, i.e. the number of differing bits. + +```cpp +struct index; +``` + + +### neighbors::ivf_binary::index::index + +Construct an empty index. It needs to be trained with `build` or loaded with `deserialize`. + +```cpp +explicit index(raft::resources const& res); +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | + +**Returns** + +`explicit` + +**Additional overload:** `neighbors::ivf_binary::index::index` + +Construct an empty index of `dim`-bit vectors. It needs to be trained and then populated. + +```cpp +index(raft::resources const& res, const index_params& params, uint32_t dim); +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | +| `params` | | [`const index_params&`](/api-reference/cpp-api-neighbors-ivf-binary#neighbors-ivf-binary-index-params) | | +| `dim` | | `uint32_t` | | + +**Returns** + +`void` + + +### neighbors::ivf_binary::index::dim + +The number of bits in each vector. + +```cpp +uint32_t dim() const noexcept; +``` + +**Returns** + +`uint32_t` + + +### neighbors::ivf_binary::index::encoded_dim + +The number of bytes in each encoded vector. + +```cpp +uint32_t encoded_dim() const noexcept; +``` + +**Returns** + +`uint32_t` + + +### neighbors::ivf_binary::index::n_lists + +The number of clusters/inverted lists. + +```cpp +uint32_t n_lists() const noexcept; +``` + +**Returns** + +`uint32_t` + + +### neighbors::ivf_binary::index::size + +The total number of vectors in the index. + +```cpp +int64_t size() const noexcept; +``` + +**Returns** + +`int64_t` + + +### neighbors::ivf_binary::index::adaptive_centers + +Whether `centers()` change upon extending the index. + +```cpp +bool adaptive_centers() const noexcept; +``` + +**Returns** + +`bool` + + +### neighbors::ivf_binary::index::conservative_memory_allocation + +Whether the lists are allocated without extra space (see index_params). + +```cpp +bool conservative_memory_allocation() const noexcept; +``` + +**Returns** + +`bool` + + +### neighbors::ivf_binary::index::list_sizes + +Sizes of the lists (clusters) [n_lists]. + +```cpp +raft::device_vector_view list_sizes() const noexcept; +``` + +**Returns** + +`raft::device_vector_view` + + +### neighbors::ivf_binary::index::centers + +Bit-packed cluster centers [n_lists, encoded_dim]. + +```cpp +raft::device_matrix_view centers() const noexcept; +``` + +**Returns** + +`raft::device_matrix_view` + + +### neighbors::ivf_binary::index::center_counts + +Exact number of ones at each bit of the vectors in each cluster [n_lists, 8 * encoded_dim]; + +```cpp +raft::device_matrix_view center_counts() const noexcept; +``` + +empty unless the centers are adaptive. + +**Returns** + +`raft::device_matrix_view` + + +### neighbors::ivf_binary::index::quantizer + +The build dataset's quantizer, available for external encoding of queries or new vectors. + +```cpp +const preprocessing::quantize::binary::quantizer& quantizer() const noexcept; +``` + +**Returns** + +[`const preprocessing::quantize::binary::quantizer&`](/api-reference/cpp-api-preprocessing-quantize-binary#preprocessing-quantize-binary-quantizer) + +## IVF-Binary index build + + +### neighbors::ivf_binary::build + +Build the index from a binary-quantized dataset. + +```cpp +auto build(raft::resources const& handle, +const index_params& index_params, +device_binary_dataset_view dataset) -> index; +``` + +The thresholds of the dataset dictionary are copied into the index. + +Usage example: + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `handle` | in | `raft::resources const&` | | +| `index_params` | in | [`const index_params&`](/api-reference/cpp-api-neighbors-ivf-binary#neighbors-ivf-binary-index-params) | configure the index building | +| `dataset` | in | [`device_binary_dataset_view`](/api-reference/cpp-api-core-dataset-binary#device-binary-dataset-view) | the binary-quantized vectors | + +**Returns** + +[`index`](/api-reference/cpp-api-neighbors-ivf-binary#neighbors-ivf-binary-index) + + +### neighbors::ivf_binary::extend + +Add bit-packed vectors to the index. + +```cpp +void extend(raft::resources const& handle, +device_binary_dataset_view new_vectors, +std::optional> new_indices, +index* index); +``` + +The supplied codes are used as-is. The dictionary of the new vectors is not compared with or copied into the index; choosing the quantizer is the caller's responsibility. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `handle` | in | `raft::resources const&` | | +| `new_vectors` | in | [`device_binary_dataset_view`](/api-reference/cpp-api-core-dataset-binary#device-binary-dataset-view) | the binary-quantized vectors to add | +| `new_indices` | in | `std::optional>` | ids of the new vectors; required if the index is not empty, otherwise the ids start from zero | +| `index` | inout | [`index*`](/api-reference/cpp-api-neighbors-ivf-binary#neighbors-ivf-binary-index) | the index to extend | + +**Returns** + +`void` + +## IVF-Binary index serialize + + +### neighbors::ivf_binary::serialize + +Save the index to an output stream. + +```cpp +void serialize(raft::resources const& handle, std::ostream& os, const index& index); +``` + +Experimental, both the API and the serialization format are subject to change. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `handle` | | `raft::resources const&` | | +| `os` | | `std::ostream&` | | +| `index` | | [`const index&`](/api-reference/cpp-api-neighbors-ivf-binary#neighbors-ivf-binary-index) | | + +**Returns** + +`void` + +**Additional overload:** `neighbors::ivf_binary::serialize` + +Save the index to a file. + +```cpp +void serialize(raft::resources const& handle, const std::string& filename, const index& index); +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `handle` | | `raft::resources const&` | | +| `filename` | | `const std::string&` | | +| `index` | | [`const index&`](/api-reference/cpp-api-neighbors-ivf-binary#neighbors-ivf-binary-index) | | + +**Returns** + +`void` + + +### neighbors::ivf_binary::deserialize + +Load the index from an input stream. + +```cpp +void deserialize(raft::resources const& handle, std::istream& is, index* index); +``` + +Experimental, both the API and the serialization format are subject to change. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `handle` | | `raft::resources const&` | | +| `is` | | `std::istream&` | | +| `index` | | [`index*`](/api-reference/cpp-api-neighbors-ivf-binary#neighbors-ivf-binary-index) | | + +**Returns** + +`void` + +**Additional overload:** `neighbors::ivf_binary::deserialize` + +Load the index from a file. + +```cpp +void deserialize(raft::resources const& handle, const std::string& filename, index* index); +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `handle` | | `raft::resources const&` | | +| `filename` | | `const std::string&` | | +| `index` | | [`index*`](/api-reference/cpp-api-neighbors-ivf-binary#neighbors-ivf-binary-index) | | + +**Returns** + +`void` diff --git a/fern/pages/cpp_api/index.md b/fern/pages/cpp_api/index.md index 07115a5eca..44f7a2f95b 100644 --- a/fern/pages/cpp_api/index.md +++ b/fern/pages/cpp_api/index.md @@ -11,6 +11,7 @@ These pages are generated from the documented public headers in the cuVS source ## Common +- [Binary](/api-reference/cpp-api-core-dataset-binary) - [Bloom Filter](/api-reference/cpp-api-core-bloom-filter) - [Common Types](/api-reference/cpp-api-common-types) - [Roaring Allowlist](/api-reference/cpp-api-core-roaring-allowlist) @@ -31,6 +32,7 @@ These pages are generated from the documented public headers in the cuVS source - [Epsilon Neighborhood](/api-reference/cpp-api-neighbors-epsilon-neighborhood) - [HNSW](/api-reference/cpp-api-neighbors-hnsw) - [Index](/api-reference/cpp-api-neighbors-composite-index) +- [IVF Binary](/api-reference/cpp-api-neighbors-ivf-binary) - [IVF Flat](/api-reference/cpp-api-neighbors-ivf-flat) - [IVF PQ](/api-reference/cpp-api-neighbors-ivf-pq) - [IVF Rabitq](/api-reference/cpp-api-neighbors-ivf-rabitq)