diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h index c5228f8580..45eb945fbc 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h @@ -416,6 +416,14 @@ void parse_build_param(const nlohmann::json& conf, throw std::runtime_error("invalid value for merge_type"); } } + + nlohmann::json comp_search_conf = collect_conf_with_prefix(conf, "compression_"); + if (!comp_search_conf.empty()) { + auto vpq_pams = param.compression.value_or(cuvs::neighbors::vpq_params{}); + parse_build_param(comp_search_conf, vpq_pams); + param.compression.emplace(vpq_pams); + } + param.cagra_params = [conf](raft::matrix_extent extents, cuvs::distance::DistanceType dist_type) { // Delayed parsing/initialization of cagra_params - it's called once the dataset shape is known diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_diskann_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_diskann_wrapper.h index a65e32a8a3..6a0628b913 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_diskann_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_diskann_wrapper.h @@ -8,12 +8,18 @@ #include #include +#include #include +#include +#include +#include #include #include #include #include "../common/ann_types.hpp" +#include "../common/blob.hpp" +#include "../common/conf.hpp" #include "../diskann/diskann_wrapper.h" #include "cuvs_ann_bench_utils.h" #include @@ -165,51 +171,35 @@ void cuvs_cagra_diskann::save(const std::string& file) const index_of.close(); if (!index_of) { RAFT_FAIL("Error writing output %s", file.c_str()); } - // try allocating a buffer for the dataset on host - try { - auto const* idx_ptr = cagra_build_.get_index(); - std::optional> h_dataset = std::nullopt; - auto const& data_view = idx_ptr->dataset(); - if constexpr (cuvs::neighbors::is_padded_dataset_view_v>) { - auto const& v = data_view; - auto n_rows = v.n_rows(); - auto dim = v.dim(); - auto stride = v.stride(); - h_dataset.emplace(raft::make_host_matrix(n_rows, dim)); - raft::copy_matrix(h_dataset->data_handle(), - dim, - v.view().data_handle(), - stride, - dim, - n_rows, - raft::resource::get_cuda_stream(handle_)); - } else { - RAFT_LOG_DEBUG("dataset serialization: index dataset is not device_padded_dataset_view"); - } - - if (h_dataset.has_value()) { - raft::resource::sync_stream(handle_); - std::string dataset_base_file = file + ".data"; - std::ofstream dataset_of(dataset_base_file, std::ios::out | std::ios::binary); - if (!dataset_of) { RAFT_FAIL("Cannot open file %s", dataset_base_file.c_str()); } - size_t dataset_file_offset = 0; - int size = static_cast(cagra_build_.get_index()->size()); - int dim = static_cast(cagra_build_.get_index()->dim()); - dataset_of.seekp(dataset_file_offset, dataset_of.beg); - dataset_of.write((char*)&size, sizeof(int)); - dataset_of.write((char*)&dim, sizeof(int)); - for (int i = 0; i < size; i++) { - dataset_of.write((char*)(h_dataset->data_handle() + i * h_dataset->extent(1)), - dim * sizeof(T)); - } - dataset_of.close(); - if (!dataset_of) { RAFT_FAIL("Error writing output %s", dataset_base_file.c_str()); } - } - } catch (std::bad_alloc& e) { - RAFT_LOG_INFO("Failed to serialize dataset"); - } catch (raft::logic_error& e) { - RAFT_LOG_INFO("Failed to serialize dataset"); - } + // Write the rows next to the graph; diskann::Index::load() reads them from `.data`. + // The benchmark base file is already in the same bin format, so copy it rather than pull the + // rows out of memory - this way `save()` does not care where the dataset was allocated. + const auto& ds_conf = configuration::singleton().get_dataset_conf(); + blob_file base{ds_conf.base_file, ds_conf.subset_first_row, ds_conf.subset_size}; + int size = static_cast(base.rows_limit()); + int dim = static_cast(base.n_cols()); + RAFT_EXPECTS(dim == this->dim_, "base_file dimensionality does not match the index"); + + size_t header_bytes = 2 * sizeof(uint32_t); + size_t skip_bytes = sizeof(T) * static_cast(base.rows_offset()) * dim; + size_t copy_bytes = sizeof(T) * static_cast(size) * dim; + RAFT_EXPECTS(std::filesystem::file_size(base.path()) >= header_bytes + skip_bytes + copy_bytes, + "base_file is shorter than its header claims"); + + std::ifstream base_in(base.path(), std::ios::in | std::ios::binary); + if (!base_in) { RAFT_FAIL("Cannot open file %s", base.path().c_str()); } + base_in.seekg(header_bytes + skip_bytes); + + std::string dataset_base_file = file + ".data"; + std::ofstream dataset_of(dataset_base_file, std::ios::out | std::ios::binary); + if (!dataset_of) { RAFT_FAIL("Cannot open file %s", dataset_base_file.c_str()); } + dataset_of.write((char*)&size, sizeof(int)); + dataset_of.write((char*)&dim, sizeof(int)); + std::copy_n(std::istreambuf_iterator(base_in), + copy_bytes, + std::ostreambuf_iterator(dataset_of)); + dataset_of.close(); + if (!base_in || !dataset_of) { RAFT_FAIL("Error writing output %s", dataset_base_file.c_str()); } } template diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h index 6b275068b0..23eeabb00f 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -65,6 +66,69 @@ inline void maybe_log_cagra_persistent_concurrency_hint(bool persistent_search) threads_rec); } +/** Whether the GPU can dereference this pointer: device, managed or mapped host memory. */ +inline auto is_device_accessible(const void* ptr) -> bool +{ + cudaPointerAttributes attrs{}; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&attrs, ptr)); + return attrs.type == cudaMemoryTypeDevice || attrs.type == cudaMemoryTypeManaged || + attrs.devicePointer != nullptr; +} + +/** + * A CAGRA-padded device view of `src`. + * + * The view points at `src` itself if that is device-accessible and its rows already have the width + * CAGRA requires. Otherwise `buffer` receives a single padded copy and the view points at that. + */ +template +auto make_padded_view(const raft::resources& res, + SrcT src, + raft::device_matrix& buffer) + -> cuvs::neighbors::device_padded_dataset_view +{ + if constexpr (SrcT::accessor_type::is_device_accessible) { + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(src)) { + return cuvs::neighbors::make_device_padded_dataset_view(res, src); + } + } + cuvs::neighbors::cagra::detail::copy_with_padding(res, buffer, src); + return {raft::make_const_mdspan(buffer.view()), static_cast(src.extent(1))}; +} + +/** Grow `buffers` to `size` empty matrices (device matrices are not default-constructible). */ +template +void grow_buffers(const raft::resources& res, + std::vector>& buffers, + size_t size) +{ + while (buffers.size() < size) { + buffers.emplace_back(raft::make_device_matrix(res, 0, 0)); + } +} + +/** + * Turn a host-resident index into a searchable device index carrying only the graph. + * + * The graph is already in device memory, owned by `host_index`; the returned index only views it, + * so the caller must keep `host_index` alive. The dataset is attached later, by `set_search_param` + * or `set_search_dataset`. + */ +template +auto to_graph_only_index(const raft::resources& res, HostIndexT& host_index) + -> cuvs::neighbors::cagra::device_padded_index +{ + cuvs::neighbors::cagra::device_padded_index index(res, host_index.metric()); + if (host_index.graph_fd().has_value()) { + // ACE in disk mode keeps the graph and the rows in files; hand the descriptors over instead + // of copying anything. + cuvs::neighbors::cagra::detail::fd_transfer::steal_disk_fds_to(res, host_index, index); + } else { + index.update_graph(res, host_index.graph()); + } + return index; +} + } // namespace detail enum class AllocatorType { kHostPinned, kHostHugePage, kDevice }; @@ -75,6 +139,7 @@ template class cuvs_cagra : public algo, public algo_gpu { public: using index_type = cuvs::neighbors::cagra::device_padded_index; + using host_index_type = cuvs::neighbors::cagra::host_standard_index; using search_param_base = typename algo::search_param; using algo::dim_; using algo::metric_; @@ -100,8 +165,9 @@ class cuvs_cagra : public algo, public algo_gpu { using dataset_dependent_params = std::function, cuvs::distance::DistanceType)>; dataset_dependent_params cagra_params; - size_t num_dataset_splits = 1; - CagraMergeType merge_type = CagraMergeType::kPhysical; + std::optional compression = std::nullopt; + size_t num_dataset_splits = 1; + CagraMergeType merge_type = CagraMergeType::kPhysical; }; cuvs_cagra(Metric metric, int dim, const build_param& param, int concurrent_searches = 1) @@ -166,6 +232,9 @@ class cuvs_cagra : public algo, public algo_gpu { auto get_index() const -> const index_type* { return index_.get(); } private: + /** Train the VPQ codebooks and create the CAGRA-Q index sharing the graph of `index_`. */ + void compress_dataset(const T* dataset, size_t nrow); + // handle_ must go first to make sure it dies last and all memory allocated in pool configured_raft_resources handle_{}; rmm::mr::pinned_host_memory_resource mr_pinned_; @@ -177,6 +246,9 @@ class cuvs_cagra : public algo, public algo_gpu { bool need_dataset_update_{true}; cuvs::neighbors::cagra::search_params search_params_; std::shared_ptr index_; + // Owns the graph viewed by index_ when it was constructed from host memory; dropped as soon as + // graph_ takes over in set_search_param. + std::shared_ptr host_index_; std::shared_ptr> graph_; std::shared_ptr> dataset_; std::shared_ptr> input_dataset_v_; @@ -190,12 +262,12 @@ class cuvs_cagra : public algo, public algo_gpu { std::shared_ptr filter_; std::vector> sub_indices_; + std::vector> sub_host_indices_; std::shared_ptr>> sub_dataset_buffers_ = std::make_shared>>(); - std::shared_ptr> deserialized_dataset_; - std::vector>> - sub_deserialized_datasets_; + std::shared_ptr> vpq_dataset_; + std::shared_ptr> vpq_index_; inline rmm::device_async_resource_ref get_mr(AllocatorType mem_type) { @@ -210,93 +282,43 @@ class cuvs_cagra : public algo, public algo_gpu { template void cuvs_cagra::build(const T* dataset, size_t nrow) { - auto dataset_extents = raft::make_extents(nrow, dim_); + auto dataset_extents = raft::make_extents(nrow, dim_); auto params = index_params_.cagra_params(dataset_extents, parse_metric_type(metric_)); + // The host paths keep the graph only, so the index need not hold a view of the caller's rows. + auto host_params = params; + host_params.attach_dataset_on_build = false; - // Use int64_t throughout so that device copies are compatible with dataset_ (device_matrix) and so that host padded dataset views carry the correct index type. - auto dataset_extents_i64 = - raft::make_extents(static_cast(nrow), static_cast(dim_)); auto dataset_view_host = - raft::make_mdspan(dataset, dataset_extents_i64); - bool dataset_is_on_host = raft::get_device_for_address(dataset) == -1; - // Host mdspan + ace_params: `cagra::build` dispatches to ACE. Non-ACE from host uses padded - // uses `cagra::build(res, params, dataset_view)` with a padded device dataset (or upload - // host data first). Used for both single-split and logical multi-split build paths. - bool const use_ace_host = - dataset_is_on_host && std::holds_alternative( - params.graph_build_params); + raft::make_mdspan(dataset, dataset_extents); + auto dataset_view_device = + raft::make_mdspan(dataset, dataset_extents); + // Pinned and managed allocations are readable by the GPU and hence take the device path, where + // CAGRA can use them without a copy. + bool dataset_is_on_host = !detail::is_device_accessible(dataset); if (index_params_.num_dataset_splits <= 1) { - if (use_ace_host) { - // ACE build is always graph-only; build the graph from a host_padded_dataset_view (required - // by the new build() API), then upload and attach a device padded copy for search. - // The input data may not satisfy CAGRA's per-row alignment; create an owning host-padded - // copy when needed, or a zero-copy view when the stride already matches. - const uint32_t req_stride = - cuvs::neighbors::cagra_required_row_width(static_cast(dim_), 16); - std::unique_ptr> host_padded_own; - std::optional> host_pdv; - if (static_cast(dim_) == req_stride) { - host_pdv = cuvs::neighbors::make_host_padded_dataset_view(dataset_view_host); - } else { - host_padded_own = cuvs::neighbors::make_host_padded_dataset(handle_, dataset_view_host); - host_pdv = host_padded_own->as_dataset_view(); - } - auto ace_host_index = cuvs::neighbors::cagra::build(handle_, params, *host_pdv); - auto padded = cuvs::neighbors::make_device_padded_dataset(handle_, dataset_view_host); - auto ace_index = - cuvs::neighbors::cagra::attach_dataset(handle_, ace_host_index, padded->as_dataset_view()); - *dataset_ = std::move(padded->data_); - index_ = std::make_shared(std::move(ace_index)); + if (dataset_is_on_host) { + // Construct the graph straight from host memory: cagra::build streams the rows in batches + // (and dispatches to ACE when it is configured), so the dataset is not uploaded here at all. + // The single device copy needed for search is made later, by set_search_param. + host_index_ = std::make_shared(cuvs::neighbors::cagra::build( + handle_, host_params, cuvs::neighbors::make_host_standard_dataset_view(dataset_view_host))); + index_ = + std::make_shared(detail::to_graph_only_index(handle_, *host_index_)); + // The graph moved into the index along with the file descriptors; nothing views the host + // index any more. + if (index_->graph_fd().has_value()) { host_index_.reset(); } } else { - // Non-ACE CAGRA build must use cagra::build(res, params, dataset_view) from - // make_device_padded_dataset / make_device_padded_dataset_view; the host mdspan and raw - // device mdspan entry points are not valid for these graph types. - // Host + non-ACE: copy to a device buffer first, then use the same path - // as a native device pointer. - raft::device_matrix_view mds; - if (dataset_is_on_host) { - *dataset_ = std::move(raft::make_device_matrix( - handle_, static_cast(nrow), static_cast(dim_))); - raft::copy(dataset_->data_handle(), - dataset, - static_cast(nrow) * dim_, - raft::resource::get_cuda_stream(handle_)); - mds = raft::make_device_matrix_view( - dataset_->data_handle(), static_cast(nrow), static_cast(dim_)); - } else { - mds = raft::make_device_matrix_view( - dataset, static_cast(nrow), static_cast(dim_)); - } - const uint32_t required_stride = - cuvs::neighbors::cagra_required_row_width(static_cast(mds.extent(1)), 16); - const uint32_t src_stride = mds.stride(0) > 0 ? static_cast(mds.stride(0)) - : static_cast(mds.extent(1)); - cudaPointerAttributes ptr_attrs{}; - RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, mds.data_handle())); - const bool device_src = (reinterpret_cast(ptr_attrs.devicePointer) != nullptr); - // `cagra::index` is move-only; use a non-const `index` per branch so - // `std::move(index)` moves (a const `index` would try to copy the deleted - // cagra::index copy ctor). - if (device_src && src_stride == required_stride) { - auto const pdv = cuvs::neighbors::make_device_padded_dataset_view(handle_, mds); - *input_dataset_v_ = raft::make_device_matrix_view( - mds.data_handle(), static_cast(nrow), static_cast(dim_)); - auto index = cuvs::neighbors::cagra::build(handle_, params, pdv); - index.update_device_dataset_same_layout(handle_, pdv); - index_ = std::make_shared(std::move(index)); - } else { - auto padded = cuvs::neighbors::make_device_padded_dataset(handle_, mds); - auto view = padded->as_dataset_view(); - auto index = cuvs::neighbors::cagra::build(handle_, params, view); - index.update_device_dataset_same_layout(handle_, view); - *dataset_ = std::move(padded->data_); - index_ = std::make_shared(std::move(index)); - } + index_ = std::make_shared(cuvs::neighbors::cagra::build( + handle_, params, detail::make_padded_view(handle_, dataset_view_device, *dataset_))); + // The index views either the caller's memory or the padded copy in dataset_; either way + // there is nothing left for set_search_param to upload. + *input_dataset_v_ = dataset_view_device; + need_dataset_update_ = false; } } else { IdxT rows_per_split = raft::ceildiv(nrow, static_cast(index_params_.num_dataset_splits)); + detail::grow_buffers(handle_, *sub_dataset_buffers_, index_params_.num_dataset_splits); for (size_t i = 0; i < index_params_.num_dataset_splits; ++i) { IdxT start = static_cast(i * rows_per_split); if (start >= nrow) break; @@ -304,106 +326,35 @@ void cuvs_cagra::build(const T* dataset, size_t nrow) const T* sub_ptr = dataset + static_cast(start) * dim_; auto sub_host = raft::make_host_matrix_view(sub_ptr, rows, dim_); - auto sub_dev = raft::make_device_matrix_view( - sub_ptr, static_cast(rows), static_cast(dim_)); + auto sub_dev = + raft::make_device_matrix_view(sub_ptr, rows, dim_); + auto& sub_dataset_buffer = (*sub_dataset_buffers_)[i]; auto sub_index = index_type(handle_, params.metric); if (index_params_.merge_type == CagraMergeType::kPhysical) { + // Physical merge only needs the rows of every split; cagra::merge builds the graph. if (dataset_is_on_host) { - sub_dataset_buffers_->emplace_back( - raft::make_device_matrix(handle_, rows, dim_)); - raft::copy(sub_dataset_buffers_->back().data_handle(), - sub_ptr, - static_cast(rows) * dim_, - raft::resource::get_cuda_stream(handle_)); - cuvs::neighbors::device_padded_dataset_view dv( - raft::make_const_mdspan(sub_dataset_buffers_->back().view()), dim_); - sub_index.update_device_dataset_same_layout(handle_, dv); + sub_index.update_device_dataset_same_layout( + handle_, detail::make_padded_view(handle_, sub_host, sub_dataset_buffer)); } else { - if (cuvs::neighbors::matrix_row_width_matches_cagra_required(sub_dev)) { - auto pdv = cuvs::neighbors::make_device_padded_dataset_view(handle_, sub_dev); - sub_index.update_device_dataset_same_layout(handle_, pdv); - } else { - auto padded = cuvs::neighbors::make_device_padded_dataset(handle_, sub_dev); - sub_dataset_buffers_->push_back(std::move(padded->data_)); - cuvs::neighbors::device_padded_dataset_view pdv( - raft::make_const_mdspan(sub_dataset_buffers_->back().view()), dim_); - sub_index.update_device_dataset_same_layout(handle_, pdv); - } + sub_index.update_device_dataset_same_layout( + handle_, detail::make_padded_view(handle_, sub_dev, sub_dataset_buffer)); } } if (index_params_.merge_type == CagraMergeType::kLogical) { - if (use_ace_host) { - // ACE build is always graph-only; build the graph from a host_padded_dataset_view - // (required by the new build() API), then upload and attach a device padded copy. - const uint32_t req_stride_sub = - cuvs::neighbors::cagra_required_row_width(static_cast(dim_), 16); - std::unique_ptr> host_padded_sub_own; - std::optional> host_pdv_sub; - if (static_cast(dim_) == req_stride_sub) { - host_pdv_sub = cuvs::neighbors::make_host_padded_dataset_view(sub_host); - } else { - host_padded_sub_own = cuvs::neighbors::make_host_padded_dataset(handle_, sub_host); - host_pdv_sub = host_padded_sub_own->as_dataset_view(); - } - auto ace_host_index = cuvs::neighbors::cagra::build(handle_, params, *host_pdv_sub); - auto padded_sub = cuvs::neighbors::make_device_padded_dataset(handle_, sub_host); - sub_index = cuvs::neighbors::cagra::attach_dataset( - handle_, ace_host_index, padded_sub->as_dataset_view()); - sub_dataset_buffers_->push_back(std::move(padded_sub->data_)); - } else if (dataset_is_on_host) { - sub_dataset_buffers_->emplace_back(raft::make_device_matrix( - handle_, static_cast(rows), static_cast(dim_))); - raft::copy(sub_dataset_buffers_->back().data_handle(), - sub_ptr, - static_cast(rows) * dim_, - raft::resource::get_cuda_stream(handle_)); - auto mds_sub = raft::make_device_matrix_view( - sub_dataset_buffers_->back().data_handle(), static_cast(rows), dim_); - const uint32_t req_sub = cuvs::neighbors::cagra_required_row_width( - static_cast(mds_sub.extent(1)), 16); - const uint32_t src_sub = mds_sub.stride(0) > 0 ? static_cast(mds_sub.stride(0)) - : static_cast(mds_sub.extent(1)); - cudaPointerAttributes sub_attrs{}; - RAFT_CUDA_TRY(cudaPointerGetAttributes(&sub_attrs, mds_sub.data_handle())); - const bool sub_device = (reinterpret_cast(sub_attrs.devicePointer) != nullptr); - if (sub_device && src_sub == req_sub) { - auto pdv_sub = cuvs::neighbors::make_device_padded_dataset_view(handle_, mds_sub); - sub_index = cuvs::neighbors::cagra::build(handle_, params, pdv_sub); - sub_index.update_device_dataset_same_layout(handle_, pdv_sub); - } else { - auto padded_sub = cuvs::neighbors::make_device_padded_dataset(handle_, mds_sub); - auto view = padded_sub->as_dataset_view(); - auto index = cuvs::neighbors::cagra::build(handle_, params, view); - index.update_device_dataset_same_layout(handle_, view); - sub_dataset_buffers_->push_back(std::move(padded_sub->data_)); - sub_index = std::move(index); - } + if (dataset_is_on_host) { + // As in the single-split case: graph only, the rows are uploaded by set_search_dataset. + sub_host_indices_.push_back( + std::make_shared(cuvs::neighbors::cagra::build( + handle_, host_params, cuvs::neighbors::make_host_standard_dataset_view(sub_host)))); + sub_index = detail::to_graph_only_index(handle_, *sub_host_indices_.back()); + if (sub_index.graph_fd().has_value()) { sub_host_indices_.pop_back(); } } else { - auto mds_sub = sub_dev; - const uint32_t req_sub = cuvs::neighbors::cagra_required_row_width( - static_cast(mds_sub.extent(1)), 16); - const uint32_t src_sub = mds_sub.stride(0) > 0 ? static_cast(mds_sub.stride(0)) - : static_cast(mds_sub.extent(1)); - cudaPointerAttributes sub_attrs{}; - RAFT_CUDA_TRY(cudaPointerGetAttributes(&sub_attrs, mds_sub.data_handle())); - const bool sub_device = (reinterpret_cast(sub_attrs.devicePointer) != nullptr); - if (sub_device && src_sub == req_sub) { - auto pdv_sub = cuvs::neighbors::make_device_padded_dataset_view(handle_, mds_sub); - sub_index = cuvs::neighbors::cagra::build(handle_, params, pdv_sub); - sub_index.update_device_dataset_same_layout(handle_, pdv_sub); - } else { - auto padded_sub = cuvs::neighbors::make_device_padded_dataset(handle_, mds_sub); - auto view = padded_sub->as_dataset_view(); - auto index = cuvs::neighbors::cagra::build(handle_, params, view); - index.update_device_dataset_same_layout(handle_, view); - sub_dataset_buffers_->push_back(std::move(padded_sub->data_)); - sub_index = std::move(index); - } + sub_index = cuvs::neighbors::cagra::build( + handle_, params, detail::make_padded_view(handle_, sub_dev, sub_dataset_buffer)); } } - auto sub_index_shared = std::make_shared(std::move(sub_index)); - sub_indices_.push_back(std::move(sub_index_shared)); + sub_indices_.push_back(std::make_shared(std::move(sub_index))); } if (index_params_.merge_type == CagraMergeType::kPhysical) { std::vector indices; @@ -417,14 +368,52 @@ void cuvs_cagra::build(const T* dataset, size_t nrow) for (auto* index : indices) { merged_rows += static_cast(index->size()); } - auto const stride = static_cast(indices.front()->dataset().stride()); + 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_)); index_ = std::make_shared(cuvs::neighbors::cagra::merge( handle_, params, indices, merged_dataset_view, merge_row_filter)); + // The merged index holds all the rows now; drop the splits rather than keep a second copy + // of the dataset on the device for the rest of the run. + sub_indices_.clear(); + sub_host_indices_.clear(); + sub_dataset_buffers_->clear(); + *input_dataset_v_ = dataset_view_device; + need_dataset_update_ = false; } } + if (index_params_.compression.has_value()) { + RAFT_EXPECTS(index_params_.num_dataset_splits <= 1, + "cagra: compression_* (CAGRA-Q) cannot be combined with num_dataset_splits > 1."); + compress_dataset(dataset, nrow); + } +} + +template +void cuvs_cagra::compress_dataset(const T* dataset, size_t nrow) +{ + RAFT_EXPECTS(parse_metric_type(metric_) == cuvs::distance::DistanceType::L2Expanded, + "cagra: compression_* (CAGRA-Q) requires the L2Expanded metric."); + RAFT_EXPECTS(!index_->graph_fd().has_value(), + "cagra: compression_* (CAGRA-Q) requires the graph in memory; it cannot be combined " + "with a disk-resident (ACE) graph."); + auto rows = static_cast(nrow); + // make_vpq_dataset() reads the rows wherever they are: host-resident ones are subsampled and + // encoded in bounded batches instead of being staged on the device. + auto src = raft::make_device_matrix_view(dataset, rows, dim_); + 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()); + + // 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_); + index_->update_device_dataset_same_layout(handle_, empty_dv); + *dataset_ = raft::make_device_matrix(handle_, 0, 0); + need_dataset_update_ = false; } inline auto allocator_to_string(AllocatorType mem_type) -> std::string @@ -474,10 +463,15 @@ void cuvs_cagra::set_search_param(const search_param_base& param, // NB: update_graph() only stores a view in the index. We need to keep the graph object alive. index_->update_graph(handle_, make_const_mdspan(graph_->view())); + if (vpq_index_) { vpq_index_->update_graph(handle_, make_const_mdspan(graph_->view())); } + // graph_ owns the graph now, so release the host index that used to own it. + host_index_.reset(); needs_dynamic_batcher_update = true; } - if (sp.dataset_mem != dataset_mem_ || need_dataset_update_) { + // CAGRA-Q searches the compressed rows in vpq_index_, so the dense dataset is never needed. + if (!index_params_.compression.has_value() && + (sp.dataset_mem != dataset_mem_ || need_dataset_update_)) { dataset_mem_ = sp.dataset_mem; // First free up existing memory @@ -503,8 +497,18 @@ void cuvs_cagra::set_search_param(const search_param_base& param, needs_dynamic_batcher_update = true; } + if (index_params_.compression.has_value() && !vpq_index_) { + // The codebooks are not part of the serialized index, so they have to be trained again after + // load(). Unlike before, the reported build time therefore excludes the compression. + compress_dataset(input_dataset_v_->data_handle(), + static_cast(input_dataset_v_->extent(0))); + needs_dynamic_batcher_update = true; + } + // dynamic batching if (sp.dynamic_batching) { + RAFT_EXPECTS(!index_params_.compression.has_value(), + "cagra: dynamic batching is not supported together with compression_* (CAGRA-Q)."); if (!dynamic_batcher_ || needs_dynamic_batcher_update) { dynamic_batcher_ = std::make_shared>( @@ -532,52 +536,42 @@ void cuvs_cagra::set_search_dataset(const T* dataset, size_t nrow) { if (index_params_.num_dataset_splits > 1 && index_params_.merge_type == CagraMergeType::kLogical) { - bool dataset_is_on_host = raft::get_device_for_address(dataset) == -1; - if (dataset_is_on_host) { sub_dataset_buffers_->clear(); } + bool dataset_is_on_host = !detail::is_device_accessible(dataset); IdxT rows_per_split = raft::ceildiv(nrow, static_cast(index_params_.num_dataset_splits)); + detail::grow_buffers(handle_, *sub_dataset_buffers_, sub_indices_.size()); for (size_t i = 0; i < sub_indices_.size(); ++i) { IdxT start = static_cast(i * rows_per_split); if (start >= nrow) break; IdxT rows = std::min(rows_per_split, static_cast(nrow) - start); const T* sub_ptr = dataset + static_cast(start) * dim_; - auto sub_dev = raft::make_device_matrix_view( - sub_ptr, static_cast(rows), static_cast(dim_)); + auto sub_host = + raft::make_host_matrix_view(sub_ptr, rows, dim_); + auto sub_dev = + raft::make_device_matrix_view(sub_ptr, rows, dim_); auto sub_index = sub_indices_[i].get(); - if (index_params_.merge_type == CagraMergeType::kLogical) { - if (dataset_is_on_host) { - sub_dataset_buffers_->emplace_back( - raft::make_device_matrix(handle_, rows, dim_)); - raft::copy(sub_dataset_buffers_->back().data_handle(), - sub_ptr, - static_cast(rows) * dim_, - raft::resource::get_cuda_stream(handle_)); - cuvs::neighbors::device_padded_dataset_view dv( - raft::make_const_mdspan(sub_dataset_buffers_->back().view()), dim_); - sub_index->update_device_dataset_same_layout(handle_, dv); - } else { - if (cuvs::neighbors::matrix_row_width_matches_cagra_required(sub_dev)) { - auto pdv = cuvs::neighbors::make_device_padded_dataset_view(handle_, sub_dev); - sub_index->update_device_dataset_same_layout(handle_, pdv); - } else { - auto padded = cuvs::neighbors::make_device_padded_dataset(handle_, sub_dev); - sub_dataset_buffers_->push_back(std::move(padded->data_)); - cuvs::neighbors::device_padded_dataset_view pdv( - raft::make_const_mdspan(sub_dataset_buffers_->back().view()), dim_); - sub_index->update_device_dataset_same_layout(handle_, pdv); - } - } + // Release the storage of this split before allocating its replacement, so that the device + // never holds two copies of a split at once. + auto& sub_dataset_buffer = (*sub_dataset_buffers_)[i]; + sub_dataset_buffer = raft::make_device_matrix(handle_, 0, 0); + if (dataset_is_on_host) { + sub_index->update_device_dataset_same_layout( + handle_, detail::make_padded_view(handle_, sub_host, sub_dataset_buffer)); + } else { + sub_index->update_device_dataset_same_layout( + handle_, detail::make_padded_view(handle_, sub_dev, sub_dataset_buffer)); } } need_dataset_update_ = false; } else { + bool is_vpq = index_params_.compression.has_value(); // It can happen that we are re-using a previous algo object which already has // the dataset set. Check if we need update. if (static_cast(input_dataset_v_->extent(0)) != nrow || input_dataset_v_->data_handle() != dataset) { *input_dataset_v_ = raft::make_device_matrix_view(dataset, nrow, this->dim_); - need_dataset_update_ = true; + need_dataset_update_ = !is_vpq; // ignore update if this is a VPQ dataset. } } } @@ -595,7 +589,7 @@ void cuvs_cagra::save(const std::string& file) const f << sub_indices_.size(); f.close(); } else { - cuvs::neighbors::cagra::serialize(handle_, file, *index_, true); + cuvs::neighbors::cagra::serialize(handle_, file, *index_, false); } } @@ -616,52 +610,23 @@ void cuvs_cagra::load(const std::string& file) meta >> count; meta.close(); sub_indices_.clear(); - sub_deserialized_datasets_.resize(count); + sub_host_indices_.clear(); for (size_t i = 0; i < count; ++i) { std::string subfile = file + (i == 0 ? "" : ".subidx." + std::to_string(i)); auto sub_index = std::make_shared(handle_); - std::unique_ptr> tmp_ds; - cuvs::neighbors::cagra::deserialize(handle_, subfile, sub_index.get(), &tmp_ds); - sub_deserialized_datasets_[i] = - std::shared_ptr>(std::move(tmp_ds)); + cuvs::neighbors::cagra::deserialize(handle_, subfile, sub_index.get()); sub_indices_.push_back(std::move(sub_index)); } } else { index_ = std::make_shared(handle_); - deserialized_dataset_.reset(); - std::unique_ptr> tmp_ds; - cuvs::neighbors::cagra::deserialize(handle_, file, index_.get(), &tmp_ds); - deserialized_dataset_ = - std::shared_ptr>(std::move(tmp_ds)); + cuvs::neighbors::cagra::deserialize(handle_, file, index_.get()); } } template std::unique_ptr> cuvs_cagra::copy() { - auto out = std::make_unique>(metric_, dim_, index_params_); - out->refine_ratio_ = refine_ratio_; - out->graph_mem_ = graph_mem_; - out->dataset_mem_ = dataset_mem_; - out->need_dataset_update_ = need_dataset_update_; - out->search_params_ = search_params_; - out->index_ = index_; - out->graph_ = graph_; - out->dataset_ = dataset_; - out->input_dataset_v_ = - std::make_shared>( - *input_dataset_v_); - out->dynamic_batcher_ = dynamic_batcher_; - out->dynamic_batcher_sp_ = dynamic_batcher_sp_; - out->dynamic_batching_max_batch_size_ = dynamic_batching_max_batch_size_; - out->dynamic_batching_n_queues_ = dynamic_batching_n_queues_; - out->dynamic_batching_conservative_dispatch_ = dynamic_batching_conservative_dispatch_; - out->filter_ = filter_; - out->sub_indices_ = sub_indices_; - out->sub_dataset_buffers_ = sub_dataset_buffers_; - out->deserialized_dataset_ = deserialized_dataset_; - out->sub_deserialized_datasets_ = sub_deserialized_datasets_; - return out; + return std::make_unique>(std::cref(*this)); // use copy constructor } template @@ -683,6 +648,9 @@ void cuvs_cagra::search_base( queries_view, neighbors_view, distances_view); + } else if (vpq_index_) { + cuvs::neighbors::cagra::search( + handle_, search_params_, *vpq_index_, queries_view, neighbors_view, distances_view, *filter_); } else { if (index_params_.num_dataset_splits <= 1 || index_params_.merge_type == CagraMergeType::kPhysical) { 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 8df12b9f2a..6d94e9495f 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h @@ -78,7 +78,7 @@ class cuvs_mg_cagra : public algo, public algo_gpu { build_param index_params_; cuvs::neighbors::mg_search_params search_params_; std::shared_ptr< - cuvs::neighbors::mg_index, T, IdxT>> + cuvs::neighbors::mg_index, T, IdxT>> index_; }; @@ -93,10 +93,13 @@ void cuvs_mg_cagra::build(const T* dataset, size_t nrow) auto dataset_mds = raft::make_host_matrix_view(dataset, nrow, dim_); - auto dataset_view = cuvs::neighbors::make_host_standard_dataset_view(dataset_mds); - auto idx = cuvs::neighbors::cagra::build(clique_, build_params, dataset_view); - index_ = std::make_shared< - cuvs::neighbors::mg_index, T, IdxT>>( + // 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>>( std::move(idx)); } @@ -129,7 +132,7 @@ template void cuvs_mg_cagra::load(const std::string& file) { index_ = std::make_shared< - cuvs::neighbors::mg_index, T, IdxT>>( + cuvs::neighbors::mg_index, T, IdxT>>( clique_, index_params_.mode); cuvs::neighbors::cagra::deserialize(clique_, file, index_.get()); } diff --git a/cpp/include/cuvs/preprocessing/quantize/pq.hpp b/cpp/include/cuvs/preprocessing/quantize/pq.hpp index 29d0bbb1e3..112341f2ad 100644 --- a/cpp/include/cuvs/preprocessing/quantize/pq.hpp +++ b/cpp/include/cuvs/preprocessing/quantize/pq.hpp @@ -10,9 +10,11 @@ #include #include #include +#include #include #include +#include #include namespace CUVS_EXPORT cuvs { @@ -246,11 +248,18 @@ void inverse_transform( namespace detail { -template -[[nodiscard]] cuvs::neighbors::device_vpq_dataset vpq_train_from_device_rows( +// Trains from `n_rows` rows of `stride` elements each, whether they are device-accessible or +// host-resident; the residency is detected from the pointer. +// +// NB: the element type is erased into `dtype` so that this stays a plain function: under hidden +// default visibility, an instantiation cannot be exported from the shared library when one of its +// template arguments (`half`, or any mdspan type) is itself hidden, because the visibility of an +// instantiation is capped by that of its template arguments. +[[nodiscard]] CUVS_EXPORT cuvs::neighbors::device_vpq_dataset vpq_train_from_rows( raft::resources const& res, cuvs::neighbors::vpq_params const& params, - T const* src_ptr, + void const* src_ptr, + cudaDataType_t dtype, int64_t n_rows, int64_t dim, int64_t stride); @@ -258,12 +267,17 @@ template } // namespace detail /** - * @brief Train VPQ storage (codebooks + encoded rows) from a device row-major mdspan/matrix. + * @brief Train VPQ storage (codebooks + encoded rows) from a row-major mdspan/mdarray/dataset. * - * Accepts any device-accessible mdspan with `value_type`, `extent`, `stride`, and `data_handle` - * (same pattern as `cuvs::neighbors::make_device_padded_dataset`). Row-major tight storage (logical - * stride equals dimension) is passed through to training without an extra pack copy; wider row - * pitch triggers a contiguous dense copy first. Empty sources are rejected. + * Accepts either a row-major mdspan with `value_type`, `extent`, `stride`, and `data_handle` (same + * pattern as `cuvs::neighbors::make_device_padded_dataset`), or any cuVS dense dataset / dataset + * view exposing `view`, `dim` and `stride`, in which case the logical `dim()` is quantized and the + * row padding is skipped. The rows may be device-accessible or host-resident. Device-accessible + * rows (device, managed or pinned) with tight row-major storage (logical stride equals dimension) + * are passed through to training as they are; a wider row pitch triggers a contiguous dense copy + * first. Host-resident rows are subsampled for training and encoded in bounded batches, so the + * dense dataset is never staged on the device in full; they must be tightly packed. Empty sources + * are rejected. The element type must be `float`, `half`, `int8_t` or `uint8_t`. * * Typical **CAGRA** usage: build the graph on dense vectors, then attach VPQ for search (metric * must remain `L2Expanded` for this path). Train VPQ from the same CAGRA-padded device layout you @@ -277,7 +291,7 @@ template * // `idx` is a `cagra::index` with graph built on dense rows. * // `padded` is a `device_padded_dataset_view` view of those same rows. * cuvs::neighbors::vpq_params vpq_params{}; - * auto vpq = cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, vpq_params, padded.view()); + * auto vpq = cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, vpq_params, padded); * idx.update_device_dataset_same_layout(res, vpq.as_dataset_view()); * @endcode */ @@ -287,16 +301,34 @@ template SrcT const& src) -> cuvs::neighbors::device_vpq_dataset { - using T = typename SrcT::value_type; - RAFT_EXPECTS(src.extent(0) > 0, "make_vpq_dataset: dataset is empty"); - cudaPointerAttributes ptr_attrs; - RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src.data_handle())); - auto const* device_ptr = reinterpret_cast(ptr_attrs.devicePointer); - RAFT_EXPECTS(device_ptr != nullptr, "make_vpq_dataset: source must be device-accessible."); - const int64_t n_rows = src.extent(0); - const int64_t dim = src.extent(1); - const int64_t stride = src.stride(0) > 0 ? src.stride(0) : dim; - return detail::vpq_train_from_device_rows(res, params, device_ptr, n_rows, dim, stride); + // A cuVS dataset keeps its logical width in `dim()` while `view()` spans the full row pitch. + if constexpr (requires { + src.view(); + src.dim(); + src.stride(); + }) { + auto const rows = src.view(); + using value_type = typename decltype(rows)::value_type; + using extents_type = raft::matrix_extent; + return make_vpq_dataset( + res, + params, + raft::mdspan{ + rows.data_handle(), + raft::make_strided_layout(extents_type{rows.extent(0), int64_t{src.dim()}}, + cuda::std::array{int64_t{src.stride()}, 1})}); + } else { + using value_type = typename SrcT::value_type; + static_assert(std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v, + "make_vpq_dataset: element type must be float, half, int8_t or uint8_t"); + const int64_t n_rows = src.extent(0); + const int64_t dim = src.extent(1); + const int64_t stride = src.stride(0) > 0 ? src.stride(0) : dim; + RAFT_EXPECTS(n_rows > 0, "make_vpq_dataset: dataset is empty"); + return detail::vpq_train_from_rows( + res, params, src.data_handle(), raft::get_cuda_data_type(), n_rows, dim, stride); + } } /** @} */ // end of group product diff --git a/cpp/src/preprocessing/quantize/pq.cu b/cpp/src/preprocessing/quantize/pq.cu index 68068aa35f..20b8f21d36 100644 --- a/cpp/src/preprocessing/quantize/pq.cu +++ b/cpp/src/preprocessing/quantize/pq.cu @@ -79,55 +79,63 @@ CUVS_INST_VPQ_BUILD(uint8_t); namespace detail { template -auto vpq_train_from_device_rows(raft::resources const& res, - cuvs::neighbors::vpq_params const& params, - T const* src_ptr, - int64_t n_rows, - int64_t dim, - int64_t stride) - -> cuvs::neighbors::device_vpq_dataset +auto train_from_rows(raft::resources const& res, + cuvs::neighbors::vpq_params const& params, + T const* src_ptr, + int64_t n_rows, + int64_t dim, + int64_t stride) -> cuvs::neighbors::device_vpq_dataset { - auto stream = raft::resource::get_cuda_stream(res); + cudaPointerAttributes ptr_attrs; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src_ptr)); + auto const* device_ptr = reinterpret_cast(ptr_attrs.devicePointer); + if (device_ptr == nullptr) { + // A host mdspan makes training subsample the rows and encoding stream them in bounded batches, + // so the dense dataset is never staged on the device. + RAFT_EXPECTS(stride == dim, "make_vpq_dataset: host input must be tightly packed"); + auto row_view = raft::make_host_matrix_view(src_ptr, n_rows, dim); + return detail::vpq_build_half(res, params, row_view); + } if (stride != dim) { auto dense = raft::make_device_matrix(res, n_rows, dim); - raft::copy_matrix(dense.data_handle(), dim, src_ptr, stride, dim, n_rows, stream); + raft::copy_matrix(dense.data_handle(), + dim, + device_ptr, + stride, + dim, + n_rows, + raft::resource::get_cuda_stream(res)); auto dense_view = raft::make_device_matrix_view(dense.data_handle(), n_rows, dim); return detail::vpq_build_half(res, params, dense_view); } - auto row_view = raft::make_device_matrix_view(src_ptr, n_rows, dim); + auto row_view = raft::make_device_matrix_view(device_ptr, n_rows, dim); return detail::vpq_build_half(res, params, row_view); } -} // namespace detail +auto vpq_train_from_rows(raft::resources const& res, + cuvs::neighbors::vpq_params const& params, + void const* src_ptr, + cudaDataType_t dtype, + int64_t n_rows, + int64_t dim, + int64_t stride) -> cuvs::neighbors::device_vpq_dataset +{ + switch (dtype) { + case CUDA_R_32F: + return train_from_rows(res, params, static_cast(src_ptr), n_rows, dim, stride); + case CUDA_R_16F: + return train_from_rows(res, params, static_cast(src_ptr), n_rows, dim, stride); + case CUDA_R_8I: + return train_from_rows(res, params, static_cast(src_ptr), n_rows, dim, stride); + case CUDA_R_8U: + return train_from_rows( + res, params, static_cast(src_ptr), n_rows, dim, stride); + default: + RAFT_FAIL("make_vpq_dataset: unsupported dataset element type %d", static_cast(dtype)); + } +} -template cuvs::neighbors::device_vpq_dataset -detail::vpq_train_from_device_rows(raft::resources const&, - cuvs::neighbors::vpq_params const&, - float const*, - int64_t, - int64_t, - int64_t); -template cuvs::neighbors::device_vpq_dataset -detail::vpq_train_from_device_rows(raft::resources const&, - cuvs::neighbors::vpq_params const&, - half const*, - int64_t, - int64_t, - int64_t); -template cuvs::neighbors::device_vpq_dataset -detail::vpq_train_from_device_rows(raft::resources const&, - cuvs::neighbors::vpq_params const&, - int8_t const*, - int64_t, - int64_t, - int64_t); -template cuvs::neighbors::device_vpq_dataset -detail::vpq_train_from_device_rows(raft::resources const&, - cuvs::neighbors::vpq_params const&, - uint8_t const*, - int64_t, - int64_t, - int64_t); +} // namespace detail } // namespace cuvs::preprocessing::quantize::pq diff --git a/cpp/tests/preprocessing/product_quantization.cu b/cpp/tests/preprocessing/product_quantization.cu index 6ea881cd0b..a392e7e1db 100644 --- a/cpp/tests/preprocessing/product_quantization.cu +++ b/cpp/tests/preprocessing/product_quantization.cu @@ -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 */ @@ -312,6 +312,59 @@ TEST(ProductQuantizationTestF, Parameters) raft::logic_error); } +TEST(ProductQuantizationTestF, MakeVpqDatasetFromHost) +{ + raft::resources handle; + constexpr int64_t n_rows = 64; + constexpr int64_t dim = 16; + auto dataset = raft::make_host_matrix(n_rows, dim); + for (std::size_t i = 0; i < dataset.size(); ++i) { + dataset.data_handle()[i] = static_cast(i % 31) / 31.0f; + } + + cuvs::neighbors::vpq_params params{ + .pq_bits = 4, .pq_dim = 4, .vq_n_centers = 1, .kmeans_n_iters = 2}; + auto vpq = make_vpq_dataset(handle, params, raft::make_const_mdspan(dataset.view())); + raft::resource::sync_stream(handle); + + EXPECT_EQ(vpq.n_rows(), n_rows); + EXPECT_EQ(vpq.dim(), dim); + EXPECT_NE(vpq.data.data_handle(), nullptr); +} + +TEST(ProductQuantizationTestF, MakeVpqDatasetFromPaddedView) +{ + raft::resources handle; + constexpr int64_t n_rows = 64; + constexpr int64_t dim = 16; + constexpr int64_t stride = 24; // row pitch wider than the logical width + + auto host_rows = raft::make_host_matrix(n_rows, stride); + for (int64_t i = 0; i < n_rows; i++) { + for (int64_t j = 0; j < stride; j++) { + // The padding is far away from the payload: quantizing it would be plainly visible. + host_rows(i, j) = j < dim ? static_cast((i * dim + j) % 31) / 31.0f : 1e3f; + } + } + auto device_rows = raft::make_device_matrix(handle, n_rows, stride); + raft::copy(device_rows.data_handle(), + host_rows.data_handle(), + host_rows.size(), + raft::resource::get_cuda_stream(handle)); + cuvs::neighbors::device_padded_dataset_view padded( + raft::make_device_matrix_view(device_rows.data_handle(), n_rows, stride), + dim); + + cuvs::neighbors::vpq_params params{ + .pq_bits = 4, .pq_dim = 4, .vq_n_centers = 1, .kmeans_n_iters = 2}; + auto vpq = make_vpq_dataset(handle, params, padded); + raft::resource::sync_stream(handle); + + EXPECT_EQ(vpq.n_rows(), n_rows); + EXPECT_EQ(vpq.dim(), dim); + EXPECT_NE(vpq.data.data_handle(), nullptr); +} + // Define test cases with different parameters template const std::vector> inputs = {