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..ee5cef4694 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,47 @@ 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 +330,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 +372,57 @@ 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."); + auto rows = static_cast(nrow); + // The codebooks are trained on the device; upload the rows first if the GPU cannot read them. + auto staging = raft::make_device_matrix(handle_, 0, 0); + auto src = raft::make_device_matrix_view(dataset, rows, dim_); + if (!detail::is_device_accessible(dataset)) { + staging = raft::make_device_matrix(handle_, rows, dim_); + raft::copy(staging.data_handle(), + dataset, + static_cast(rows) * dim_, + raft::resource::get_cuda_stream(handle_)); + src = raft::make_const_mdspan(staging.view()); + } + 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 +472,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 +506,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,10 +545,10 @@ 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; @@ -544,33 +557,21 @@ void cuvs_cagra::set_search_dataset(const T* dataset, size_t nrow) auto sub_dev = raft::make_device_matrix_view( sub_ptr, static_cast(rows), static_cast(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 || @@ -595,7 +596,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,23 +617,16 @@ 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()); } } @@ -683,6 +677,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()); }