diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 97cc9d8bac..ea023c82b6 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -179,13 +179,13 @@ static void merge_indices_for_layout( cuvs::neighbors::cagra::detail::merged_dataset_size( *res_ptr, index_ptrs, row_filter); auto const dim = static_cast(index_ptrs.front()->dim()); - auto const stride = static_cast(index_ptrs.front()->dataset().stride()); + auto const stride = static_cast(index_ptrs.front()->dataset().data_const_view().stride(0)); try { - auto matrix = raft::make_device_matrix(*res_ptr, final_row_count, stride); using owner_t = cuvs::neighbors::owning_dataset_for_view_t; - auto owner = std::make_unique(std::move(matrix), dim); - auto view = owner->as_dataset_view(); + auto owner = std::make_unique( + cuvs::neighbors::detail::make_dense_array(*res_ptr, final_row_count, dim, stride)); + auto view = std::as_const(*owner).view(); auto merged_idx = cuvs::neighbors::cagra::merge( *res_ptr, params_cpp, index_ptrs, view, merge_params, row_filter); @@ -218,8 +218,9 @@ static void merge_indices_for_layout( cuvs::neighbors::host_standard_dataset_view>; using host_owner_t = cuvs::neighbors::owning_dataset_for_view_t; - auto matrix = raft::make_host_matrix(final_row_count, stride); - std::fill_n(matrix.data_handle(), static_cast(matrix.size()), T{}); + auto matrix = cuvs::neighbors::detail::make_dense_array( + *res_ptr, final_row_count, dim, stride); + std::fill_n(matrix.data_handle(), matrix.mapping().required_span_size(), T{}); std::size_t row_offset = 0; auto stream = raft::resource::get_cuda_stream(*res_ptr); @@ -227,8 +228,8 @@ static void merge_indices_for_layout( auto const& input = index->dataset(); raft::copy_matrix(matrix.data_handle() + row_offset * static_cast(stride), static_cast(stride), - input.view().data_handle(), - static_cast(input.stride()), + input.data_const_view().data_handle(), + static_cast(input.data_const_view().stride(0)), static_cast(dim), static_cast(input.n_rows()), stream); @@ -236,8 +237,8 @@ static void merge_indices_for_layout( } raft::resource::sync_stream(*res_ptr); - auto owner = std::make_unique(std::move(matrix), dim); - auto view = owner->as_dataset_view(); + auto owner = std::make_unique(std::move(matrix)); + auto view = std::as_const(*owner).view(); auto merged_idx = cuvs::neighbors::cagra::build(*res_ptr, params_cpp, view); auto* holder = new cuvs_cagra_c_api_index_lifetime_holder{std::move(merged_idx)}; @@ -402,7 +403,7 @@ static void with_dataset_view(cuvsDataset_t dataset, Fn&& fn) RAFT_EXPECTS(dataset->addr != 0, "null dataset storage"); if (dataset->is_owning) { auto* owner = reinterpret_cast(dataset->addr); - auto view = owner->as_dataset_view(); + auto view = std::as_const(*owner).view(); fn(view); } else { fn(*reinterpret_cast(dataset->addr)); @@ -1383,7 +1384,7 @@ void get_dataset_view(cuvsCagraIndex_t index, DLManagedTensor* dataset) RAFT_FAIL( "cuvsCagraIndexGetDataset does not expose quantized datasets as dense DLPack tensors"); } else { - cuvs::core::to_dlpack(idx.dataset().view(), dataset); + cuvs::core::to_dlpack(cuvs::neighbors::dense_row_major_view(idx.dataset()), dataset); } }); } diff --git a/c/src/neighbors/mg_cagra.cpp b/c/src/neighbors/mg_cagra.cpp index 664883cd6f..5540b0fed2 100644 --- a/c/src/neighbors/mg_cagra.cpp +++ b/c/src/neighbors/mg_cagra.cpp @@ -79,7 +79,7 @@ static void with_device_padded_dataset_view(cuvsDataset_t dataset, Fn&& fn) using view_t = cuvs::neighbors::device_padded_dataset_view; if (dataset->is_owning) { auto* owner = reinterpret_cast(dataset->addr); - auto view = owner->as_dataset_view(); + auto view = std::as_const(*owner).view(); fn(view); } else { fn(*reinterpret_cast(dataset->addr)); diff --git a/c/src/preprocessing/quantize/pq.cpp b/c/src/preprocessing/quantize/pq.cpp index 1e3a48694a..169c24a001 100644 --- a/c/src/preprocessing/quantize/pq.cpp +++ b/c/src/preprocessing/quantize/pq.cpp @@ -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 */ @@ -244,7 +244,7 @@ extern "C" cuvsError_t cuvsProductQuantizerGetPqCodebook(cuvsProductQuantizer_t if (quantizer->dtype.code == kDLFloat && quantizer->dtype.bits == 32) { auto pq_mdspan = (reinterpret_cast*>(quant_addr)) - ->vpq_codebooks.pq_code_book.view(); + ->vpq_codebooks.dictionary_view().pq_code_book; cuvs::core::to_dlpack(pq_mdspan, pq_codebook); } else { RAFT_FAIL("Unsupported quantizer dtype: %d and bits: %d", @@ -266,7 +266,7 @@ extern "C" cuvsError_t cuvsProductQuantizerGetVqCodebook(cuvsProductQuantizer_t if (quantizer->dtype.code == kDLFloat && quantizer->dtype.bits == 32) { auto pq_mdspan = (reinterpret_cast*>(quant_addr)) - ->vpq_codebooks.vq_code_book.view(); + ->vpq_codebooks.dictionary_view().vq_code_book; cuvs::core::to_dlpack(pq_mdspan, vq_codebook); } else { RAFT_FAIL("Unsupported quantizer dtype: %d and bits: %d", diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h index da5e5baa9d..19ef1e4fd6 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h @@ -93,7 +93,12 @@ auto make_padded_view(const raft::resources& res, } } cuvs::neighbors::cagra::detail::copy_with_padding(res, buffer, src); - return {raft::make_const_mdspan(buffer.view()), static_cast(src.extent(1))}; + return cuvs::neighbors::detail::make_dense_view< + cuvs::neighbors::device_padded_dataset_view>( + buffer.data_handle(), + buffer.extent(0), + static_cast(src.extent(1)), + static_cast(buffer.extent(1))); } /** Grow `buffers` to `size` empty matrices (device matrices are not default-constructible). */ @@ -373,8 +378,9 @@ void cuvs_cagra::build(const T* dataset, size_t nrow) auto const stride = static_cast( cuvs::neighbors::cagra_required_row_width(static_cast(dim_))); *dataset_ = raft::make_device_matrix(handle_, merged_rows, stride); - auto merged_dataset_view = cuvs::neighbors::device_padded_dataset_view( - raft::make_const_mdspan(dataset_->view()), static_cast(dim_)); + auto merged_dataset_view = cuvs::neighbors::detail::make_dense_view< + cuvs::neighbors::device_padded_dataset_view>( + dataset_->data_handle(), merged_rows, static_cast(dim_), stride); index_ = std::make_shared(cuvs::neighbors::cagra::merge(handle_, params, @@ -413,11 +419,11 @@ void cuvs_cagra::compress_dataset(const T* dataset, size_t nrow) vpq_dataset_ = std::make_shared>( cuvs::preprocessing::quantize::pq::make_vpq_dataset(handle_, *index_params_.compression, src)); vpq_index_ = std::make_shared>( - handle_, parse_metric_type(metric_), vpq_dataset_->as_dataset_view(), index_->graph()); + handle_, parse_metric_type(metric_), vpq_dataset_->view(), index_->graph()); // Search runs on the compressed rows and the graph, so release the dense copy of the dataset. - cuvs::neighbors::device_padded_dataset_view empty_dv( - raft::make_device_matrix_view(static_cast(nullptr), 0, this->dim_), this->dim_); + auto empty_dv = cuvs::neighbors::detail::make_empty_dataset_view< + cuvs::neighbors::device_padded_dataset_view>(this->dim_); *index_ = cuvs::neighbors::cagra::update_dataset(handle_, std::move(*index_), empty_dv); *dataset_ = raft::make_device_matrix(handle_, 0, 0); need_dataset_update_ = false; @@ -482,9 +488,9 @@ void cuvs_cagra::set_search_param(const search_param_base& param, dataset_mem_ = sp.dataset_mem; // First free up existing memory - *dataset_ = raft::make_device_matrix(handle_, 0, 0); - cuvs::neighbors::device_padded_dataset_view empty_dv( - raft::make_device_matrix_view(static_cast(nullptr), 0, this->dim_), this->dim_); + *dataset_ = raft::make_device_matrix(handle_, 0, 0); + auto empty_dv = cuvs::neighbors::detail::make_empty_dataset_view< + cuvs::neighbors::device_padded_dataset_view>(this->dim_); *index_ = cuvs::neighbors::cagra::update_dataset(handle_, std::move(*index_), empty_dv); // Allocate space using the correct memory resource. @@ -494,10 +500,9 @@ void cuvs_cagra::set_search_param(const search_param_base& param, auto mr = get_mr(dataset_mem_); cuvs::neighbors::cagra::detail::copy_with_padding(handle_, *dataset_, *input_dataset_v_, mr); - cuvs::neighbors::device_padded_dataset_view dv( - raft::make_device_matrix_view( - dataset_->data_handle(), dataset_->extent(0), dataset_->extent(1)), - this->dim_); + auto dv = cuvs::neighbors::detail::make_dense_view< + cuvs::neighbors::device_padded_dataset_view>( + dataset_->data_handle(), dataset_->extent(0), this->dim_, dataset_->extent(1)); *index_ = cuvs::neighbors::cagra::update_dataset(handle_, std::move(*index_), dv); need_dataset_update_ = false; diff --git a/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h index 7a89b68394..89b5f6483f 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h @@ -93,13 +93,12 @@ void cuvs_mg_cagra::build(const T* dataset, size_t nrow) auto dataset_mds = raft::make_host_matrix_view(dataset, nrow, dim_); - // The row alignment of the host view is irrelevant: every per-rank device shard is padded - // individually during the multi-GPU build. - cuvs::neighbors::host_padded_dataset_view dataset_view(dataset_mds, - static_cast(dim_)); - auto idx = cuvs::neighbors::cagra::build(clique_, build_params, dataset_view); - index_ = std::make_shared< - cuvs::neighbors::mg_index, T, IdxT>>( + // A padded view cannot wrap rows that are not aligned, so the host rows get padded in a copy. + auto dataset_padded = cuvs::neighbors::make_host_padded_dataset(clique_, dataset_mds); + auto idx = + cuvs::neighbors::cagra::build(clique_, build_params, std::as_const(*dataset_padded).view()); + index_ = std::make_shared< + cuvs::neighbors::mg_index, T, IdxT>>( std::move(idx)); } diff --git a/cpp/cmake/patches/faiss-1.14-cuvs-26.08.diff b/cpp/cmake/patches/faiss-1.14-cuvs-26.08.diff index 738ae8fa3b..5e89668e3d 100644 --- a/cpp/cmake/patches/faiss-1.14-cuvs-26.08.diff +++ b/cpp/cmake/patches/faiss-1.14-cuvs-26.08.diff @@ -46,7 +46,7 @@ index 4be7cebb8..42315d027 100644 raft_handle, cuvs::distance::DistanceType::BitwiseHamming, - dataset_mds, -+ host_to_device_dataset_->as_dataset_view(), ++ std::as_const(*host_to_device_dataset_).view(), raft::make_const_mdspan(knn_graph_copy.view())); } else { FAISS_THROW_MSG( @@ -72,7 +72,7 @@ index 4be7cebb8..42315d027 100644 cuvs::neighbors::cagra::build( - raft_handle, index_params_, dataset)); + raft_handle, index_params_, -+ host_to_device_dataset_->as_dataset_view())); ++ std::as_const(*host_to_device_dataset_).view())); } } @@ -99,7 +99,7 @@ index 4be7cebb8..42315d027 100644 + *cuvs_index = cuvs::neighbors::cagra::update_dataset( + raft_handle, + std::move(*cuvs_index), -+ host_to_device_dataset_->as_dataset_view()); ++ std::as_const(*host_to_device_dataset_).view()); } store_dataset_ = true; } @@ -189,7 +189,7 @@ index 6ff3419b7..f03a5e5f0 100644 raft_handle, metricFaissToCuvs(metric_, false), - dataset_mds, -+ host_to_device_dataset_->as_dataset_view(), ++ std::as_const(*host_to_device_dataset_).view(), raft::make_const_mdspan(knn_graph_copy.view())); } else { FAISS_THROW_MSG( @@ -215,7 +215,7 @@ index 6ff3419b7..f03a5e5f0 100644 cuvs::neighbors::cagra::build( - raft_handle, index_params_, dataset)); + raft_handle, index_params_, -+ host_to_device_dataset_->as_dataset_view())); ++ std::as_const(*host_to_device_dataset_).view())); } } @@ -241,7 +241,7 @@ index 6ff3419b7..f03a5e5f0 100644 + *cuvs_index = cuvs::neighbors::cagra::update_dataset( + raft_handle, + std::move(*cuvs_index), -+ host_to_device_dataset_->as_dataset_view()); ++ std::as_const(*host_to_device_dataset_).view()); } store_dataset_ = true; } diff --git a/cpp/include/cuvs/core/dataset.hpp b/cpp/include/cuvs/core/dataset.hpp new file mode 100644 index 0000000000..c3b5f2cfe2 --- /dev/null +++ b/cpp/include/cuvs/core/dataset.hpp @@ -0,0 +1,17 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +/** + * The minimal viable dataset API: the owning dataset, its view, and the one spec that needs no + * storage backend at all. A dataset that does keep data around takes the spec that describes how - + * `cuvs/core/dataset/dense.hpp`, `sparse.hpp`, `vpq.hpp`, `sq.hpp` - so that including this header + * costs nothing beyond the types themselves. + */ + +#include +#include +#include diff --git a/cpp/include/cuvs/core/dataset/dense.hpp b/cpp/include/cuvs/core/dataset/dense.hpp new file mode 100644 index 0000000000..4238832c92 --- /dev/null +++ b/cpp/include/cuvs/core/dataset/dense.hpp @@ -0,0 +1,99 @@ +/* + * 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 + +namespace CUVS_EXPORT cuvs { + +/** + * Dense plain/padded dataset implemented via raft::mdarray. + */ +template +struct mdarray_spec { + template + struct apply { + // NOTE: index type != extents type + // The index type can vary depending on the use case; we often may want to store indices in + // 32-bit slots to save memory on large datasets. The extents are always fixed to 64-bit signed + // integers to simplify the conversion and avoid integer overflow issues. + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using data_type = + raft::mdarray, LayoutPolicy, ContainerPolicy>; + using data_view_type = typename data_type::view_type; + using const_data_view_type = typename data_type::const_view_type; + + [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> data_view_type + { + return data.view(); + } + + [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept + -> const_data_view_type + { + return data.view(); + } + // A mutable mdspan does not decay to a read-only one on its own; it has to be rebuilt over an + // accessor of `const T`. + [[nodiscard]] static auto get_data_const_view(data_view_type data) noexcept + -> const_data_view_type + { + return raft::make_const_mdspan(data); + } + // NOTE: the shape accessors take the data in either form - the storage an owning dataset holds, + // or the view a `dataset_view` holds. A plain spec defines no dictionary either, so it ignores + // whatever the dataset puts in that slot instead. + [[nodiscard]] static auto get_n_rows(const auto& data) noexcept -> index_type + { + return static_cast(data.extent(0)); + } + [[nodiscard]] static auto get_dim(const auto& data, const auto&) noexcept -> uint32_t + { + return static_cast(data.extent(1)); + } + }; +}; + +/** + * Row-major rows padded to 16 bytes, the alignment the CAGRA kernels require. A padded matrix keeps + * the logical dimension as its second extent; the row stride follows from the layout. + */ +template +using padded_layout = cuda::std::layout_right_padded<16 / std::gcd(16, sizeof(T))>; + +template +using contiguous_dataset = + dataset>; +template +using padded_dataset = dataset, ContainerPolicy>>; +template +using device_contiguous_dataset = + contiguous_dataset>>; +template +using device_padded_dataset = + padded_dataset>>; +template +using host_contiguous_dataset = + contiguous_dataset>>; +template +using host_padded_dataset = + padded_dataset>>; + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/core/dataset/empty.hpp b/cpp/include/cuvs/core/dataset/empty.hpp new file mode 100644 index 0000000000..00fde08e94 --- /dev/null +++ b/cpp/include/cuvs/core/dataset/empty.hpp @@ -0,0 +1,60 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include +#include + +namespace CUVS_EXPORT cuvs { + +/** + * Empty dataset specification: contains no data, but keeps the dimension of the dataset implement + * the common interface. + */ +struct empty_spec { + struct empty_dataset_rep { + uint32_t dim; + }; + + template + struct apply { + using data_type = empty_dataset_rep; + using data_view_type = empty_dataset_rep; + // NOTE: not `const empty_dataset_rep`: a read-only view holds this as a member, and a const + // member would delete its assignment operator. + using const_data_view_type = empty_dataset_rep; + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + + [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> data_view_type + { + return data; + } + // NOTE: the view of an empty dataset is the storage itself, so one overload serves both the + // owning dataset and the view. + [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept + -> const_data_view_type + { + return data; + } + // NOTE: the shape accessors take the data in either form - the storage an owning dataset holds, + // or the view a `dataset_view` holds. A plain spec defines no dictionary either, so it ignores + // whatever the dataset puts in that slot instead. + [[nodiscard]] static auto get_n_rows(const auto&) noexcept -> index_type { return 0; } + [[nodiscard]] static auto get_dim(const auto& data, const auto&) noexcept -> uint32_t + { + return data.dim; + } + }; +}; + +template +using empty_dataset = dataset; + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/core/dataset/sparse.hpp b/cpp/include/cuvs/core/dataset/sparse.hpp new file mode 100644 index 0000000000..d20d21e662 --- /dev/null +++ b/cpp/include/cuvs/core/dataset/sparse.hpp @@ -0,0 +1,129 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include +#include +#include +#include + +#include +#include + +namespace CUVS_EXPORT cuvs { + +/** + * Sparse dataset specification implemented via the raft sparse matrix hierarchy. `SparseLayoutT` + * selects the representation (`csr_layout`, `coo_layout`). + * + * NOTE: only the memory space of the container policy is used here; raft's sparse types bind the + * policy themselves, separately for the values, the offsets and the indices. + */ +template +struct sparse_spec { + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + + static constexpr raft::memory_type mem_type = ContainerPolicy::mem_type; + static_assert(mem_type == raft::memory_type::host || mem_type == raft::memory_type::device, + "raft's sparse hierarchy only provides host and device containers"); + + using data_type = std::conditional_t< + mem_type == raft::memory_type::device, + typename SparseLayoutT:: + template matrix_type, + typename SparseLayoutT:: + template matrix_type>; + using data_view_type = typename data_type::view_type; + // NOTE: raft's sparse hierarchy provides no read-only view yet (see the TODO in + // raft/core/sparse_types.hpp); until then, even reading the structure requires a mutable + // matrix. + // + // Consequently a read-only `dataset_view` over a sparse spec is read-only in name + // only: the two view types coincide, so writing through it still compiles. That resolves itself + // once raft grows a const sparse view. + using const_data_view_type = data_view_type; + + [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> data_view_type + { + return data.view(); + } + + [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept + -> const_data_view_type + { + return const_cast(data).view(); + } + // NOTE: there is no read-only view to convert to (see above), so the view passes through. + [[nodiscard]] static auto get_data_const_view(data_view_type data) noexcept + -> const_data_view_type + { + return data; + } + // NOTE: raft's sparse hierarchy exposes the structure through a mutable accessor only, so + // reading the shape of either the owning matrix or its view takes a cast. + [[nodiscard]] static auto get_structure_view(const auto& data) noexcept + { + return const_cast&>(data).structure_view(); + } + // NOTE: the shape accessors take the data in either form - the storage an owning dataset holds, + // or the view a `dataset_view` holds. A plain spec defines no dictionary either, so it ignores + // whatever the dataset puts in that slot instead. + [[nodiscard]] static auto get_n_rows(const auto& data) noexcept -> index_type + { + return static_cast(get_structure_view(data).get_n_rows()); + } + [[nodiscard]] static auto get_dim(const auto& data, const auto&) noexcept -> uint32_t + { + return static_cast(get_structure_view(data).get_n_cols()); + } + }; +}; + +/** + * CSR representation: one compressed row per vector of the dataset. + * + * NOTE: the row offsets are 64-bit for the same reason the dense extents are: the number of + * non-zeros is not bounded by the index type. The column indices are feature ids, so they are + * bounded by the dataset dimension rather than by its size. + */ +struct csr_layout { + template typename ContainerPolicy> + using matrix_type = raft::csr_matrix; +}; + +/** + * COO representation: one (row, column, value) triple per non-zero. + * + * NOTE: unlike CSR, the row component is an index into the dataset rather than an offset into the + * non-zeros, hence the dataset index type. + */ +struct coo_layout { + template typename ContainerPolicy> + using matrix_type = raft::coo_matrix; +}; + +template +using csr_dataset = dataset>; +template +using coo_dataset = dataset>; +template +using device_csr_dataset = + csr_dataset>>; +template +using device_coo_dataset = + coo_dataset>>; +template +using host_csr_dataset = csr_dataset>>; +template +using host_coo_dataset = coo_dataset>>; + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/core/dataset/sq.hpp b/cpp/include/cuvs/core/dataset/sq.hpp new file mode 100644 index 0000000000..b56fd1f4b6 --- /dev/null +++ b/cpp/include/cuvs/core/dataset/sq.hpp @@ -0,0 +1,78 @@ +/* + * 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 + +namespace CUVS_EXPORT cuvs { + +/** + * Scalar-quantized dataset specification: every component of a vector becomes one code, so the + * dictionary is only the range the codes are mapped back onto and the dimension is still a property + * of the data. + */ +template +struct sq_spec { + // One code per component, so the codes are laid out exactly like a dense dataset. + template + using storage_spec = + typename mdarray_spec::template apply; + + using math_type = MathT; + + // Nothing to allocate: the whole dictionary is the interval the codes are dequantized into, + // which is exactly what the scalar quantizer holds. + using dictionary_type = cuvs::preprocessing::quantize::scalar::quantizer; + using dictionary_view_type = dictionary_type; + + template + struct apply : storage_spec { + // The dataset stores codes, but it still represents vectors of `T`. + using value_type = std::remove_cv_t; + using code_type = int8_t; + // NOTE: re-exported rather than declared here, so that they do not belong to one instantiation + // of `apply`; qualified lookup into `apply` does not reach the enclosing spec on its own. + using math_type = sq_spec::math_type; + using dictionary_type = sq_spec::dictionary_type; + using dictionary_view_type = sq_spec::dictionary_view_type; + + // NOTE: every accessor is inherited; the dimension is still one code per component, so even + // `get_dim` is the dense one. + [[nodiscard]] static auto get_dictionary_view(const dictionary_type& dict) noexcept + -> dictionary_view_type + { + return dict; + } + // NOTE: the dictionary is its own view here, read-only or not, so one overload serves both the + // owning dataset and the view. + [[nodiscard]] static auto get_dictionary_const_view(const dictionary_type& dict) noexcept + -> dictionary_view_type + { + return dict; + } + }; +}; + +template +using sq_dataset = dataset>; +template +using device_sq_dataset = + sq_dataset>>; +template +using host_sq_dataset = + sq_dataset>>; + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/core/dataset/type.hpp b/cpp/include/cuvs/core/dataset/type.hpp new file mode 100644 index 0000000000..11baff1231 --- /dev/null +++ b/cpp/include/cuvs/core/dataset/type.hpp @@ -0,0 +1,150 @@ +/* + * 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 { + +/** + * An owning dataset: the data and, for a compressed dataset, the dictionary that gives it meaning. + * Ownership is unique, so the type is move-only whenever the underlying storage is. + */ +template +struct dataset { + // NOTE: element type for an owning dataset cannot be const, + // the constness of the dataset itself selects the flavor of the view it hands out. + static_assert(!std::is_const_v, "Element type for an owning dataset must not be const."); + + using view_type = dataset_view; + using const_view_type = dataset_view; + + using spec_type = typename view_type::spec_type; + using value_type = typename view_type::value_type; + using index_type = typename view_type::index_type; + using code_type = typename view_type::code_type; + + using data_type = typename view_type::data_type; + using data_view_type = typename view_type::data_view_type; + using const_data_view_type = typename view_type::const_data_view_type; + + /** + Optional slot that contains everything to give meaning to the encoded data; + `std::monostate` for an uncompressed dataset. + */ + using dictionary_type = typename view_type::dictionary_type; + using dictionary_view_type = typename view_type::dictionary_view_type; + using const_dictionary_view_type = typename view_type::const_dictionary_view_type; + + static constexpr bool is_compressed = view_type::is_compressed; + + explicit dataset(data_type&& data) + requires(!is_compressed) + : data_{std::move(data)} + { + } + + dataset(data_type&& data, dictionary_type&& dictionary) + requires(is_compressed) + : data_{std::move(data)}, dictionary_{std::move(dictionary)} + { + } + + /** + * @brief A non-owning view of the whole dataset, data and dictionary alike. + * + * @return A mutable view of the dataset. + */ + [[nodiscard]] auto view() noexcept -> view_type + { + if constexpr (is_compressed) { + return view_type{data_view(), dictionary_view()}; + } else { + return view_type{data_view()}; + } + } + /** + * @brief A read-only view of the whole dataset; the constness of the dataset selects this flavor. + * + * @return A read-only view of the dataset. + */ + [[nodiscard]] auto view() const noexcept -> const_view_type + { + if constexpr (is_compressed) { + return const_view_type{data_const_view(), const_dictionary_view()}; + } else { + return const_view_type{data_const_view()}; + } + } + + /** + * @brief The number of rows in the dataset. + * + * @return The number of rows. + */ + [[nodiscard]] auto n_rows() const noexcept -> index_type { return spec_type::get_n_rows(data_); } + /** + * @brief The dimension of the vectors in the dataset. + * + * @return The dimension. + */ + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + // NOTE: a compressed dataset may not be able to tell its dimension from the encoded data alone + return spec_type::get_dim(data_, dictionary_); + } + + /** + * @brief A non-owning view of the data, such as mdspan. + * + * @return A mutable view of the data. + */ + [[nodiscard]] auto data_view() noexcept -> data_view_type + { + return spec_type::get_data_view(data_); + } + /** + * @brief A read-only view of the data. + * + * @return A read-only view of the data. + */ + [[nodiscard]] auto data_const_view() const noexcept -> const_data_view_type + { + return spec_type::get_data_const_view(data_); + } + + /** + * @brief A non-owning view of the dictionary, such as the two codebooks of a VPQ dataset. + * + * @return A mutable view of the dictionary. + */ + [[nodiscard]] auto dictionary_view() noexcept -> dictionary_view_type + requires(is_compressed) + { + return spec_type::get_dictionary_view(dictionary_); + } + /** + * @brief A read-only view of the dictionary. + * + * @return A read-only view of the dictionary. + */ + [[nodiscard]] auto const_dictionary_view() const noexcept -> const_dictionary_view_type + requires(is_compressed) + { + return spec_type::get_dictionary_const_view(dictionary_); + } + + private: + data_type data_; + [[no_unique_address]] dictionary_type dictionary_; +}; + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/core/dataset/view.hpp b/cpp/include/cuvs/core/dataset/view.hpp new file mode 100644 index 0000000000..e34ac19261 --- /dev/null +++ b/cpp/include/cuvs/core/dataset/view.hpp @@ -0,0 +1,201 @@ +/* + * 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 { + +/** + * A spec is compressing if it defines a dictionary: the codebooks, the quantization range, or + * whatever else is needed to interpret the data. Plain specs define none and the dataset + * substitutes `std::monostate`, so "there is no dictionary" and "the dictionary slot is empty" are + * the same vocabulary type. + * + * NOTE: a spec declares its vocabulary types - the dictionary above, and any storage representation + * of its own - at spec scope and only re-exports them from `apply`. A type defined inside `apply` + * belongs to one instantiation of it, and the read-only and the mutable flavor of a view apply the + * spec to element types that differ in constness; the two would then disagree about a nested + * dictionary and could no longer be built from one another. + */ +template +concept compressed_dataset_spec = requires { typename SpecT::dictionary_type; }; + +/** + * A non-owning view of a dataset: a view of the data and, for a compressed dataset, a view of the + * dictionary that gives it meaning. Cheap to copy and to pass by value, like an mdspan. + */ +template +struct dataset_view { + using spec_type = typename SpecT::template apply; + // The read-only flavor, in which each view type collapses onto its const counterpart. + static constexpr bool is_const_view = std::is_const_v; + + using value_type = typename spec_type::value_type; + using index_type = typename spec_type::index_type; + // What one element of the data slot is + // - for compressed datasets it is the raw data element type; often is uint8_t but may be + // a larger type to enforce better alignment. + // - for uncompressed datasets it is the value type of the dataset. + using code_type = typename decltype([] { + if constexpr (requires { typename spec_type::code_type; }) { + return std::type_identity{}; + } else { + return std::type_identity{}; + } + }())::type; + + using data_type = typename spec_type::data_type; + using const_data_view_type = typename spec_type::const_data_view_type; + using data_view_type = + std::conditional_t; + + // Optional slot that contains everything to give meaning to the encoded data; + // `std::monostate` for an uncompressed dataset. + using dictionary_type = typename decltype([] { + if constexpr (requires { typename spec_type::dictionary_type; }) { + return std::type_identity{}; + } else { + return std::type_identity{}; + } + }())::type; + // Optional, so that the specs don't have to declare it; a spec that declares neither view uses + // the dictionary itself, and one that declares only the mutable view uses it for both. + using const_dictionary_view_type = typename decltype([] { + if constexpr (requires { typename spec_type::const_dictionary_view_type; }) { + return std::type_identity{}; + } else if constexpr (requires { typename spec_type::dictionary_view_type; }) { + return std::type_identity{}; + } else { + return std::type_identity{}; + } + }())::type; + // Optional, so that the specs don't have to declare it + using dictionary_view_type = typename decltype([] { + if constexpr (is_const_view) { + return std::type_identity{}; + } else if constexpr (requires { typename spec_type::dictionary_view_type; }) { + return std::type_identity{}; + } else { + return std::type_identity{}; + } + }())::type; + + static constexpr bool is_compressed = compressed_dataset_spec; + + /** An empty view, like a default-constructed mdspan: no rows, no data. */ + dataset_view() noexcept = default; + + explicit dataset_view(data_view_type data_view) + requires(!is_compressed) + : data_view_{data_view} + { + } + + dataset_view(data_view_type data_view, dictionary_view_type dictionary_view) + requires(is_compressed) + : data_view_{data_view}, dictionary_view_{dictionary_view} + { + } + + /** + * @brief A mutable view decays to a read-only one, the way an mdspan does. + * + * NOTE: a constructor template, because in the mutable flavor the parameter would spell the very + * same type; a non-template would be a copy constructor there, and declaring one - even a + * constrained-away one - suppresses the implicitly declared copy constructor. + * + * NOTE: the dictionary is assigned in the body rather than in the init-list, because the + * accessor it comes from is constrained on `is_compressed`. + */ + template + requires(is_const_view && std::same_as>) + dataset_view(const dataset_view& other) noexcept + : data_view_{other.data_const_view()} + { + if constexpr (is_compressed) { dictionary_view_ = other.const_dictionary_view(); } + } + + /** + * @brief The number of rows in the dataset. + * + * @return The number of rows. + */ + [[nodiscard]] auto n_rows() const noexcept -> index_type + { + return spec_type::get_n_rows(data_view_); + } + /** + * @brief The dimension of the vectors in the dataset. + * + * NOTE: a compressed dataset may not be able to tell its dimension from the encoded data alone. + * + * @return The dimension. + */ + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + return spec_type::get_dim(data_view_, dictionary_view_); + } + + /** + * @brief A non-owning view of the data, such as mdspan. + * + * @return A view of the data. + */ + [[nodiscard]] auto data_view() const noexcept -> data_view_type { return data_view_; } + /** + * @brief A read-only view of the data. + * + * NOTE: giving up write access is not a plain cast for most view types - a mutable mdspan, for + * one, has to be rebuilt over a const accessor - hence the spec does the conversion. + * + * @return A read-only view of the data. + */ + [[nodiscard]] auto data_const_view() const noexcept -> const_data_view_type + { + if constexpr (is_const_view) { + return data_view_; + } else { + return spec_type::get_data_const_view(data_view_); + } + } + + /** + * @brief A non-owning view of the dictionary, such as the two codebooks of a VPQ dataset. + * + * @return A view of the dictionary. + */ + [[nodiscard]] auto dictionary_view() const noexcept -> dictionary_view_type + requires(is_compressed) + { + return dictionary_view_; + } + /** + * @brief A read-only view of the dictionary. + * + * @return A read-only view of the dictionary. + */ + [[nodiscard]] auto const_dictionary_view() const noexcept -> const_dictionary_view_type + requires(is_compressed) + { + if constexpr (is_const_view) { + return dictionary_view_; + } else { + return spec_type::get_dictionary_const_view(dictionary_view_); + } + } + + private: + data_view_type data_view_; + [[no_unique_address]] dictionary_view_type dictionary_view_; +}; + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/core/dataset/vpq.hpp b/cpp/include/cuvs/core/dataset/vpq.hpp new file mode 100644 index 0000000000..cf8e7f49fa --- /dev/null +++ b/cpp/include/cuvs/core/dataset/vpq.hpp @@ -0,0 +1,136 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include + +namespace CUVS_EXPORT cuvs { + +/** + * VQ+PQ compressed dataset specification, mirroring the VPQ dataset of cuvs/neighbors/common.hpp: + * the two codebooks form the dictionary, while the data slot keeps the encoded rows - the VQ label + * in the row prefix, followed by the packed PQ codes. + * + * `BookPolicy` is the container policy of the codebooks (elements of `MathT`) and `CodePolicy` the + * one of the encoded rows (elements of `uint8_t`); a bound policy cannot be rebound to another + * element type, hence the two parameters. + */ +template +struct vpq_spec { + // The encoded rows are stored like the rows of any other dense dataset; only how a row is read + // and where the dimension comes from are specific to the compression. + template + using storage_spec = + typename mdarray_spec::template apply; + + // The type the codebooks are stored in; the rows themselves are always packed codes. + using math_type = MathT; + + // [vq_n_centers, dim] for the VQ codebook, [pq_n_centers, pq_len] for the PQ one. + using book_type = + raft::mdarray, raft::layout_c_contiguous, BookPolicy>; + + struct dictionary_type { + book_type vq_code_book; + book_type pq_code_book; + }; + + // The views of the codebooks, writable or read-only alike. + template + struct dictionary_view_of { + BookViewT vq_code_book; + BookViewT pq_code_book; + + // NOTE: the dimension of the dataset is a property of the VQ codebook: the encoded rows tell + // only how many bytes it took to compress a vector, which the padding makes ambiguous. + [[nodiscard]] auto dim() const noexcept -> uint32_t { return vq_code_book.extent(1); } + [[nodiscard]] auto vq_n_centers() const noexcept -> uint32_t { return vq_code_book.extent(0); } + [[nodiscard]] auto pq_n_centers() const noexcept -> uint32_t { return pq_code_book.extent(0); } + [[nodiscard]] auto pq_len() const noexcept -> uint32_t { return pq_code_book.extent(1); } + [[nodiscard]] auto pq_bits() const noexcept -> uint32_t + { + return static_cast(std::countr_zero(pq_n_centers())); + } + [[nodiscard]] auto pq_dim() const noexcept -> uint32_t + { + return raft::div_rounding_up_safe(dim(), pq_len()); + } + }; + using dictionary_view_type = dictionary_view_of; + using const_dictionary_view_type = dictionary_view_of; + + template + struct apply : storage_spec { + // The dataset stores codes, but it still represents vectors of `T`. + using value_type = std::remove_cv_t; + using code_type = uint8_t; + // NOTE: re-exported rather than declared here, so that they do not belong to one instantiation + // of `apply`; qualified lookup into `apply` does not reach the enclosing spec on its own. + using math_type = vpq_spec::math_type; + using book_type = vpq_spec::book_type; + using dictionary_type = vpq_spec::dictionary_type; + using dictionary_view_type = vpq_spec::dictionary_view_type; + using const_dictionary_view_type = vpq_spec::const_dictionary_view_type; + + // NOTE: the data accessors are inherited; only `get_dim` differs from a dense dataset. Like the + // inherited ones, it reads the dictionary in either form: the codebooks an owning dataset + // holds, or the views a `dataset_view` holds. + [[nodiscard]] static auto get_dim(const auto&, const auto& dict) noexcept -> uint32_t + { + return static_cast(dict.vq_code_book.extent(1)); + } + [[nodiscard]] static auto get_dictionary_view(dictionary_type& dict) noexcept + -> dictionary_view_type + { + return {dict.vq_code_book.view(), dict.pq_code_book.view()}; + } + [[nodiscard]] static auto get_dictionary_const_view(const dictionary_type& dict) noexcept + -> const_dictionary_view_type + { + return {dict.vq_code_book.view(), dict.pq_code_book.view()}; + } + [[nodiscard]] static auto get_dictionary_const_view(dictionary_view_type dict) noexcept + -> const_dictionary_view_type + { + return {raft::make_const_mdspan(dict.vq_code_book), + raft::make_const_mdspan(dict.pq_code_book)}; + } + // The length of an encoded row in bytes, including the inlined VQ label. + [[nodiscard]] static auto get_encoded_row_length(const auto& data) noexcept -> uint32_t + { + return static_cast(data.extent(1)); + } + }; +}; + +template +using vpq_dataset = dataset>; +template +using device_vpq_dataset = + vpq_dataset>, + raft::device_accessor>>; +template +using host_vpq_dataset = vpq_dataset>, + raft::host_accessor>>; + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 418db9d88f..532cbf56c2 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -348,10 +348,10 @@ struct index_params : cuvs::neighbors::index_params { * cagra::index_params index_params; * // Build graph only — caller attaches dataset later. * index_params.attach_dataset_on_build = false; - * auto index = cagra::build(res, index_params, dataset->as_dataset_view()); + * auto index = cagra::build(res, index_params, dataset->view()); * // ASSERT(index.size() == 0); // no dataset yet * // Attach with a view (storage owned by `dataset`). - * index = cagra::update_dataset(res, std::move(index), dataset->as_dataset_view()); + * index = cagra::update_dataset(res, std::move(index), dataset->view()); * cagra::search(res, search_params, index, queries, neighbors, distances); * @endcode */ @@ -592,7 +592,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { : cuvs::neighbors::index(), metric_(metric), graph_(raft::make_device_matrix(res, 0, 0)), - dataset_{}, + dataset_{cuvs::neighbors::detail::make_empty_dataset_view()}, dataset_norms_(std::nullopt) { } @@ -618,7 +618,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { * the index uses the dataset; the index does not take ownership of the buffer. * @code{.cpp} * auto padded_owner = cuvs::neighbors::make_device_padded_dataset(res, dataset_mdspan); - * auto view = padded_owner->as_dataset_view(); + * auto view = padded_owner->view(); * cuvs::neighbors::cagra::device_padded_index idx(res, metric, view, * raft::make_const_mdspan(graph)); * // `padded_owner` must outlive `idx` (do not let it go out of scope while `idx` is used). @@ -793,24 +793,9 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { // Re-open the file descriptor in read-only mode for subsequent operations dataset_fd_.emplace(std::move(fd)); - if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { - auto v = raft::make_device_matrix_view( - static_cast(nullptr), int64_t{0}, dim_); - dataset_ = DatasetViewT(v, dim_); - } else if constexpr (cuvs::neighbors::is_device_standard_dataset_view_v) { - auto v = raft::make_device_matrix_view( - static_cast(nullptr), int64_t{0}, dim_); - dataset_ = DatasetViewT(v); - } else if constexpr (cuvs::neighbors::is_host_padded_dataset_view_v) { - auto v = raft::make_host_matrix_view( - static_cast(nullptr), int64_t{0}, dim_); - dataset_ = DatasetViewT(v, dim_); - } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { - auto v = raft::make_host_matrix_view( - static_cast(nullptr), int64_t{0}, dim_); - dataset_ = DatasetViewT(v); - } else if constexpr (cuvs::neighbors::is_empty_dataset_view_v) { - dataset_ = DatasetViewT{dim_}; + if constexpr (cuvs::neighbors::is_dense_row_major_dataset_view_v || + cuvs::neighbors::is_empty_dataset_view_v) { + dataset_ = cuvs::neighbors::detail::make_empty_dataset_view(dim_); } else { RAFT_FAIL("update_dataset(fd): unsupported DatasetViewT for disk-backed dataset"); } @@ -1293,7 +1278,7 @@ auto build(raft::resources const& res, * using namespace cuvs::neighbors; * // Build `extended` = old || new on device, padded for CAGRA. * auto extended = make_device_padded_dataset(res, concatenated_view); - * auto extended_view = extended->as_dataset_view(); + * auto extended_view = extended->view(); * * cagra::extend_params params; * int64_t new_start_row = static_cast(index.size()); diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index baa5b44da2..36de166f40 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -24,6 +24,8 @@ #include #include +#include +#include #include #include @@ -192,27 +194,6 @@ using dataset_owning_accessor_for_view = std::conditional_t, host_owning_accessor>; -template -using dense_owning_matrix = std::conditional_t, - raft::host_matrix>; - -template -using dense_view_matrix = - std::conditional_t, - raft::host_matrix_view>; - -template -using vpq_vq_book_matrix = std::conditional_t, - raft::host_matrix>; - -template -using vpq_data_matrix = std::conditional_t, - raft::host_matrix>; - // ----------------------------------------------------------------------------- // empty // ----------------------------------------------------------------------------- @@ -232,186 +213,6 @@ using empty_dataset_owning_storage = empty_dataset_storage; template using empty_dataset_view_storage = empty_dataset_storage; -// ----------------------------------------------------------------------------- -// dense row-major (logical dim may differ from row pitch; shared by padded & standard) -// ----------------------------------------------------------------------------- - -/** - * Dense row-major owning storage shared by padded and standard dataset containers. - * - * Template parameters: - * - MatrixT: owning matrix type that stores the payload (host/device matrix). - * - ViewT: non-owning row-major view type returned by `view()`. - * - DataT: scalar element type of the dataset payload. - * - IdxT: index type used for row counts (`n_rows()` return type). - */ -template -struct dense_row_major_dataset_owning_storage { - MatrixT data_; - uint32_t logical_dim_; - - dense_row_major_dataset_owning_storage(MatrixT&& data, uint32_t logical_dim) noexcept - : data_{std::move(data)}, logical_dim_{logical_dim} - { - } - - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data_.extent(0); } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return logical_dim_; } - [[nodiscard]] auto stride() const noexcept -> uint32_t - { - return static_cast(data_.extent(1)); - } - [[nodiscard]] auto view() const noexcept -> ViewT { return data_.view(); } - [[nodiscard]] auto data_handle() noexcept -> DataT* { return data_.data_handle(); } - [[nodiscard]] auto data_handle() const noexcept -> const DataT* { return data_.data_handle(); } -}; - -template -struct dense_row_major_dataset_view_storage { - ViewT data_; - uint32_t logical_dim_; - - dense_row_major_dataset_view_storage() noexcept = default; - - explicit dense_row_major_dataset_view_storage(ViewT v) noexcept - : data_(v), logical_dim_(static_cast(v.extent(1))) - { - } - - dense_row_major_dataset_view_storage(ViewT v, uint32_t logical_dim) noexcept - : data_(v), logical_dim_(logical_dim) - { - } - - dense_row_major_dataset_view_storage(dense_row_major_dataset_view_storage const& other) noexcept - : data_(other.data_), logical_dim_(other.logical_dim_) - { - } - - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data_.extent(0); } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return logical_dim_; } - [[nodiscard]] auto stride() const noexcept -> uint32_t - { - return static_cast(data_.stride(0) > 0 ? data_.stride(0) : data_.extent(1)); - } - [[nodiscard]] auto view() const noexcept -> ViewT { return data_; } -}; - -template -using padded_dataset_owning_storage = - dense_row_major_dataset_owning_storage; - -template -using padded_dataset_view_storage = dense_row_major_dataset_view_storage; - -template -using standard_dataset_owning_storage = - dense_row_major_dataset_owning_storage; - -template -using standard_dataset_view_storage = dense_row_major_dataset_view_storage; - -// ----------------------------------------------------------------------------- -// VPQ compressed -// ----------------------------------------------------------------------------- - -/** - * Owning storage for VPQ-compressed datasets. - * - * Template parameters: - * - VqBookMatrixT: owning matrix type for the VQ codebook. - * - PqBookMatrixT: owning matrix type for the PQ codebook. - * - DataMatrixT: owning matrix type for encoded row data (uint8 codes). - * - MathT: floating-point type used by VQ/PQ codebooks. - * - IdxT: index type used for row counts (`n_rows()` return type). - */ -template -struct vpq_dataset_owning_storage { - /** Floating-point type used for VQ/PQ codebooks (rows are still uint8 codes). */ - using math_type = MathT; - - VqBookMatrixT vq_code_book; - PqBookMatrixT pq_code_book; - DataMatrixT data; - - vpq_dataset_owning_storage(VqBookMatrixT&& vq_code_book, - PqBookMatrixT&& pq_code_book, - DataMatrixT&& data) noexcept - : vq_code_book{std::move(vq_code_book)}, - pq_code_book{std::move(pq_code_book)}, - data{std::move(data)} - { - } - - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data.extent(0); } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return vq_code_book.extent(1); } - - [[nodiscard]] constexpr inline auto encoded_row_length() const noexcept -> uint32_t - { - return data.extent(1); - } - [[nodiscard]] constexpr inline auto vq_n_centers() const noexcept -> uint32_t - { - return vq_code_book.extent(0); - } - [[nodiscard]] constexpr inline auto pq_bits() const noexcept -> uint32_t - { - auto pq_width = pq_n_centers(); -#ifdef __cpp_lib_bitops - return std::countr_zero(pq_width); -#else - uint32_t pq_bits = 0; - while (pq_width > 1) { - pq_bits++; - pq_width >>= 1; - } - return pq_bits; -#endif - } - [[nodiscard]] constexpr inline auto pq_dim() const noexcept -> uint32_t - { - return raft::div_rounding_up_unsafe(dim(), pq_len()); - } - [[nodiscard]] constexpr inline auto pq_len() const noexcept -> uint32_t - { - return pq_code_book.extent(1); - } - [[nodiscard]] constexpr inline auto pq_n_centers() const noexcept -> uint32_t - { - return pq_code_book.extent(0); - } -}; - -template -struct vpq_dataset_view_storage { - using owning_dataset_type = - dataset>; - - owning_dataset_type const* dataset_{nullptr}; - - vpq_dataset_view_storage() = default; - - explicit vpq_dataset_view_storage(owning_dataset_type const* ptr) : dataset_(ptr) - { - RAFT_EXPECTS(ptr != nullptr, "vpq_dataset_view: null dataset pointer"); - } - - [[nodiscard]] auto n_rows() const noexcept - { - using idx_type = decltype(std::declval().n_rows()); - return dataset_ != nullptr ? dataset_->n_rows() : idx_type{0}; - } - [[nodiscard]] auto dim() const noexcept -> uint32_t - { - return dataset_ != nullptr ? dataset_->dim() : uint32_t{0}; - } - [[nodiscard]] owning_dataset_type const& dset() const noexcept { return *dataset_; } -}; - } // namespace detail // ----------------------------------------------------------------------------- @@ -425,55 +226,6 @@ struct empty_dataset_container { using view_storage = detail::empty_dataset_view_storage; }; -// ----------------------------------------------------------------------------- -// padded (row-major with logical dim vs stride) -// ----------------------------------------------------------------------------- - -struct padded_dataset_container { - template - using owning_storage = - detail::padded_dataset_owning_storage, - detail::dense_view_matrix, - DataT, - IdxT>; - template - using view_storage = detail:: - padded_dataset_view_storage, DataT, IdxT>; -}; - -// ----------------------------------------------------------------------------- -// standard (row-major with arbitrary stride; no CAGRA alignment requirement) -// ----------------------------------------------------------------------------- - -struct standard_dataset_container { - template - using owning_storage = - detail::standard_dataset_owning_storage, - detail::dense_view_matrix, - DataT, - IdxT>; - template - using view_storage = detail:: - standard_dataset_view_storage, DataT, IdxT>; -}; - -// ----------------------------------------------------------------------------- -// VPQ compressed -// ----------------------------------------------------------------------------- - -struct vpq_dataset_container { - template - using owning_storage = - detail::vpq_dataset_owning_storage, - detail::vpq_vq_book_matrix, - detail::vpq_data_matrix, - MathT, - IdxT>; - template - using view_storage = - detail::vpq_dataset_view_storage; -}; - template struct dataset { static_assert(!std::is_same_v, @@ -497,7 +249,7 @@ struct dataset using owning_storage_type = typename container_type::template owning_storage; using owning_storage_type::owning_storage_type; - [[nodiscard]] auto as_dataset_view() const noexcept + [[nodiscard]] auto view() const noexcept -> dataset_view using view_storage_type::view_storage_type; }; -// ----------------------------------------------------------------------------- -// standard (row-major with arbitrary stride) -// ----------------------------------------------------------------------------- - -template -struct dataset - : standard_dataset_container::template owning_storage { - using container_type = standard_dataset_container; - using owning_storage_type = - typename container_type::template owning_storage; - using owning_storage_type::owning_storage_type; - - [[nodiscard]] auto as_dataset_view() const noexcept - -> dataset_view> - { - return dataset_view>(this->view(), - this->dim()); - } -}; - -template -struct dataset_view - : standard_dataset_container::template view_storage { - using container_type = standard_dataset_container; - using view_storage_type = typename container_type::template view_storage; - using view_storage_type::view_storage_type; -}; - -// ----------------------------------------------------------------------------- -// padded (row-major with logical dim vs stride) -// ----------------------------------------------------------------------------- - -template -struct dataset - : padded_dataset_container::template owning_storage { - using container_type = padded_dataset_container; - using owning_storage_type = - typename container_type::template owning_storage; - using owning_storage_type::owning_storage_type; - - [[nodiscard]] auto as_dataset_view() const noexcept - -> dataset_view> - { - return dataset_view>(this->view(), - this->dim()); - } -}; - -template -struct dataset_view - : padded_dataset_container::template view_storage { - using container_type = padded_dataset_container; - using view_storage_type = typename container_type::template view_storage; - using view_storage_type::view_storage_type; -}; - -// ----------------------------------------------------------------------------- -// VPQ compressed (view holds non-owning pointer to owning dataset) -// ----------------------------------------------------------------------------- - -template -struct dataset - : vpq_dataset_container::template owning_storage { - using container_type = vpq_dataset_container; - using owning_storage_type = - typename container_type::template owning_storage; - using owning_storage_type::owning_storage_type; - - [[nodiscard]] auto as_dataset_view() const - -> dataset_view> - { - return dataset_view>{this}; - } -}; - -template -struct dataset_view - : vpq_dataset_container::template view_storage { - using container_type = vpq_dataset_container; - using view_storage_type = typename container_type::template view_storage; - using view_storage_type::view_storage_type; -}; - /** * @brief Aliases for concrete `dataset` / `dataset_view` layouts. */ @@ -638,53 +289,47 @@ template using host_empty_dataset_view = dataset_view>; -template -using device_padded_dataset = - dataset>; - -template -using device_padded_dataset_view = - dataset_view>; - -template -using host_padded_dataset = - dataset>; - -template -using host_padded_dataset_view = - dataset_view>; - -template -using device_standard_dataset = - dataset>; - -template -using device_standard_dataset_view = - dataset_view>; - -template -using host_standard_dataset = - dataset>; - -template -using host_standard_dataset_view = - dataset_view>; - -template -using device_vpq_dataset = - dataset>; - -template -using device_vpq_dataset_view = - dataset_view>; - -template -using host_vpq_dataset = - dataset>; +// ----------------------------------------------------------------------------- +// Compatibility spellings of the dataset types, so that the code written against the former +// `cuvs::neighbors` dataset API keeps compiling on top of `cuvs::dataset` and `cuvs::dataset_view` +// (the headers under cuvs/core/dataset). It is all names, with the `standard` dataset as the only +// new kind. Once the callers spell the types `cuvs::...` directly, this section can be deleted. +// ----------------------------------------------------------------------------- -template -using host_vpq_dataset_view = - dataset_view>; +using cuvs::device_padded_dataset; +using cuvs::host_padded_dataset; + +/* The `standard` dataset: row-major rows of any stride (at least the dimension), no alignment. */ +template +using device_standard_dataset = cuvs::dataset< + T, + IdxT, + cuvs::mdarray_spec>>>; +template +using host_standard_dataset = cuvs::dataset< + T, + IdxT, + cuvs::mdarray_spec>>>; + +/* The former spelling carried the math type as an optional third argument. */ +template +using device_vpq_dataset = cuvs::device_vpq_dataset; +template +using host_vpq_dataset = cuvs::host_vpq_dataset; + +/* The read-only views are the const views of the owning types. */ +template +using device_padded_dataset_view = typename device_padded_dataset::const_view_type; +template +using host_padded_dataset_view = typename host_padded_dataset::const_view_type; +template +using device_standard_dataset_view = typename device_standard_dataset::const_view_type; +template +using host_standard_dataset_view = typename host_standard_dataset::const_view_type; +template +using device_vpq_dataset_view = typename device_vpq_dataset::const_view_type; +template +using host_vpq_dataset_view = typename host_vpq_dataset::const_view_type; // Maps a dataset view type to its owning (allocating) dataset counterpart. // Used by serialize/deserialize to type the out_dataset output parameter; @@ -692,71 +337,14 @@ using host_vpq_dataset_view = template struct owning_dataset_for_view; -template -struct owning_dataset_for_view> { - using type = device_padded_dataset; -}; - -template -struct owning_dataset_for_view> { - using type = device_standard_dataset; -}; - -template -struct owning_dataset_for_view> { - using type = host_padded_dataset; -}; - -template -struct owning_dataset_for_view> { - using type = host_standard_dataset; -}; - -template -struct owning_dataset_for_view> { - using type = device_vpq_dataset; +template +struct owning_dataset_for_view> { + using type = cuvs::dataset; }; template using owning_dataset_for_view_t = typename owning_dataset_for_view::type; -template -struct is_padded_dataset : std::false_type {}; - -template -struct is_padded_dataset> - : std::true_type {}; - -template -struct is_padded_dataset> - : std::true_type {}; - -template -inline constexpr bool is_padded_dataset_v = is_padded_dataset::value; - -template -struct is_standard_dataset : std::false_type {}; - -template -struct is_standard_dataset> - : std::true_type {}; - -template -struct is_standard_dataset> - : std::true_type {}; - -template -inline constexpr bool is_standard_dataset_v = is_standard_dataset::value; - -template -struct is_vpq_dataset : std::false_type {}; - -template -struct is_vpq_dataset> : std::true_type {}; - -template -inline constexpr bool is_vpq_dataset_v = is_vpq_dataset::value; - // ----------------------------------------------------------------------------- // Dataset view compile-time classification (replaces runtime std::variant dispatch). // ----------------------------------------------------------------------------- @@ -794,18 +382,32 @@ struct dataset_view_kind_of -struct dataset_view_kind_of> { +// The migrated kinds are told apart by their spec, for a dataset as well as for its view. +template