Skip to content

update_dataset CAGRA #2427

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aamijar wants to merge 5 commits into
base: main
Choose a base branch
from
Open

update_dataset CAGRA #2427

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
491ebc0
update_dataset
aamijar Aug 8, 2026
596dd19
Merge branch 'main' into update_dataset
aamijar Aug 8, 2026
eb34e5d
remove attach_dataset
aamijar Aug 12, 2026
75c3248
C layer remove attach_dataset
aamijar Aug 12, 2026
3261373
Merge branch 'main' into update_dataset
aamijar Aug 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
119 changes: 24 additions & 95 deletions c/src/neighbors/cagra.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -524,7 +524,7 @@ static void make_host_standard_dataset_view(raft::resources*,
}

template <typename T>
static void attach_dataset(raft::resources* res_ptr,
static void update_dataset(raft::resources* res_ptr,
cuvsDataset_t device_padded_dataset,
cuvsCagraIndex_t index)
{
Expand All @@ -537,7 +537,7 @@ static void attach_dataset(raft::resources* res_ptr,
auto* box = reinterpret_cast<sg_cagra_c_api_index_box*>(index->addr);
RAFT_EXPECTS(device_padded_dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE &&
device_padded_dataset->layout == CUVS_DATASET_LAYOUT_PADDED,
"cuvsCagraAttachDataset: dataset must be device padded");
"cuvsCagraUpdateDataset: dataset must be device padded");

using owner_t = cuvs::neighbors::device_padded_dataset<T, int64_t>;
using view_t = cuvs::neighbors::device_padded_dataset_view<T, int64_t>;
Expand All @@ -547,63 +547,24 @@ static void attach_dataset(raft::resources* res_ptr,
"cuvsCagraUpdateDataset: null index handle",
"cuvsCagraUpdateDataset: host index layout is allowed for this operation",
[&](auto& idx) {
auto padded_idx = cuvs::neighbors::cagra::attach_dataset(*res_ptr, idx, padded_view);
auto* holder =
new cuvs_cagra_c_api_index_lifetime_holder<T, view_t>{std::move(padded_idx)};
destroy_sg_cagra_c_api_box(index->addr);
index->addr = 0;
bind_index_lifetime_holder_to_C_index<T, view_t>(index, index->dtype, holder);
using index_t = std::remove_cvref_t<decltype(idx)>;
if constexpr (std::is_same_v<
index_t,
cuvs::neighbors::cagra::device_padded_index<T, uint32_t>>) {
idx.update_dataset(*res_ptr, padded_view);
} else {
auto padded_idx =
cuvs::neighbors::cagra::update_dataset(*res_ptr, std::move(idx), padded_view);
auto* holder =
new cuvs_cagra_c_api_index_lifetime_holder<T, view_t>{std::move(padded_idx)};
destroy_sg_cagra_c_api_box(index->addr);
index->addr = 0;
bind_index_lifetime_holder_to_C_index<T, view_t>(index, index->dtype, holder);
}
});
});
}

template <typename T>
static void update_device_dataset_same_layout(raft::resources* res_ptr,
cuvsDataset_t device_dataset,
cuvsCagraIndex_t index)
{
RAFT_EXPECTS(device_dataset != nullptr, "cuvsCagraUpdateDataset: null dataset");
RAFT_EXPECTS(index != nullptr, "cuvsCagraUpdateDataset: null index handle");
RAFT_EXPECTS(index->addr != 0, "cuvsCagraUpdateDataset: null index storage");
RAFT_EXPECTS(device_dataset->addr != 0, "cuvsCagraUpdateDataset: null dataset storage");

auto* box = reinterpret_cast<sg_cagra_c_api_index_box*>(index->addr);
if (box->layout == sg_cagra_c_api_index_box::dataset_layout::device_padded) {
RAFT_EXPECTS(device_dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE &&
device_dataset->layout == CUVS_DATASET_LAYOUT_PADDED,
"cuvsCagraUpdateDeviceDatasetSameLayout: device-padded index "
"requires a "
"device-padded dataset");
using owner_t = cuvs::neighbors::device_padded_dataset<T, int64_t>;
using view_t = cuvs::neighbors::device_padded_dataset_view<T, int64_t>;
with_dataset_view<owner_t, view_t>(device_dataset, [&](auto const& dataset_view) {
auto* idx =
reinterpret_cast<cuvs::neighbors::cagra::device_padded_index<T, uint32_t>*>(box->index_ptr);
RAFT_EXPECTS(idx != nullptr, "cuvsCagraUpdateDataset: null index handle");
idx->update_device_dataset_same_layout(*res_ptr, dataset_view);
});
} else if (box->layout == sg_cagra_c_api_index_box::dataset_layout::device_standard) {
RAFT_EXPECTS(device_dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE &&
device_dataset->layout == CUVS_DATASET_LAYOUT_STANDARD,
"cuvsCagraUpdateDeviceDatasetSameLayout: device-standard "
"index requires a "
"device-standard dataset");
using owner_t = cuvs::neighbors::device_standard_dataset<T, int64_t>;
using view_t = cuvs::neighbors::device_standard_dataset_view<T, int64_t>;
with_dataset_view<owner_t, view_t>(device_dataset, [&](auto const& dataset_view) {
auto* idx =
reinterpret_cast<cuvs::neighbors::cagra::device_standard_index<T, uint32_t>*>(box->index_ptr);
RAFT_EXPECTS(idx != nullptr, "cuvsCagraUpdateDataset: null index handle");
idx->update_device_dataset_same_layout(*res_ptr, dataset_view);
});
} else {
RAFT_FAIL(
"cuvsCagraUpdateDataset: C++ "
"update_device_dataset_same_layout "
"requires a device index and dataset");
}
}

static void _set_graph_build_params(
std::variant<std::monostate,
cuvs::neighbors::cagra::graph_build_params::ivf_pq_params,
Expand Down Expand Up @@ -711,7 +672,7 @@ void _from_args(cuvsResources_t res,
auto dataset_view = cuvs::neighbors::make_device_padded_dataset_view(*res_ptr, mds);
auto* raw = new cuvs::neighbors::cagra::device_padded_index<T, uint32_t>(
*res_ptr, metric);
raw->update_device_dataset_same_layout(*res_ptr, dataset_view);
raw->update_dataset(*res_ptr, dataset_view);
update_graph_from_dlpack(raw);
wrap_CPP_index_in_lifetime_holder_and_bind_to_C_index<
T,
Expand All @@ -721,7 +682,7 @@ void _from_args(cuvsResources_t res,
auto dataset_view = cuvs::neighbors::make_device_standard_dataset_view(mds);
auto* raw = new cuvs::neighbors::cagra::device_standard_index<T, uint32_t>(
*res_ptr, metric);
raw->update_device_dataset_same_layout(*res_ptr, dataset_view);
raw->update_dataset(*res_ptr, dataset_view);
update_graph_from_dlpack(raw);
wrap_CPP_index_in_lifetime_holder_and_bind_to_C_index<
T,
Expand Down Expand Up @@ -1578,7 +1539,7 @@ extern "C" cuvsError_t cuvsDatasetMakeStandardView(cuvsResources_t res,
});
}

static cuvsError_t dispatch_attach_dataset(cuvsResources_t res,
static cuvsError_t dispatch_update_dataset(cuvsResources_t res,
cuvsDataset_t device_padded_dataset,
cuvsCagraIndex_t index)
{
Expand All @@ -1592,40 +1553,13 @@ static cuvsError_t dispatch_attach_dataset(cuvsResources_t res,
index->dtype.bits == device_padded_dataset->dtype.bits,
"cuvsCagraUpdateDataset: dtype mismatch between index and dataset");
if (index->dtype.code == kDLFloat && index->dtype.bits == 32) {
attach_dataset<float>(res_ptr, device_padded_dataset, index);
update_dataset<float>(res_ptr, device_padded_dataset, index);
} else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) {
attach_dataset<half>(res_ptr, device_padded_dataset, index);
update_dataset<half>(res_ptr, device_padded_dataset, index);
} else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
attach_dataset<int8_t>(res_ptr, device_padded_dataset, index);
update_dataset<int8_t>(res_ptr, device_padded_dataset, index);
} else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
attach_dataset<uint8_t>(res_ptr, device_padded_dataset, index);
} else {
RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits);
}
});
}

static cuvsError_t dispatch_update_device_dataset_same_layout(cuvsResources_t res,
cuvsDataset_t device_dataset,
cuvsCagraIndex_t index)
{
return cuvs::core::translate_exceptions([=] {
auto* res_ptr = reinterpret_cast<raft::resources*>(res);
RAFT_EXPECTS(index != nullptr, "cuvsCagraUpdateDataset: null index handle");
RAFT_EXPECTS(device_dataset != nullptr,
"cuvsCagraUpdateDataset: null dataset view");
RAFT_EXPECTS(index->dtype.code == device_dataset->dtype.code &&
index->dtype.bits == device_dataset->dtype.bits,
"cuvsCagraUpdateDataset: dtype mismatch "
"between index and dataset");
if (index->dtype.code == kDLFloat && index->dtype.bits == 32) {
update_device_dataset_same_layout<float>(res_ptr, device_dataset, index);
} else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) {
update_device_dataset_same_layout<half>(res_ptr, device_dataset, index);
} else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
update_device_dataset_same_layout<int8_t>(res_ptr, device_dataset, index);
} else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
update_device_dataset_same_layout<uint8_t>(res_ptr, device_dataset, index);
update_dataset<uint8_t>(res_ptr, device_padded_dataset, index);
} else {
RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits);
}
Expand All @@ -1650,12 +1584,7 @@ extern "C" cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res,
"cuvsCagraUpdateDataset: dtype mismatch between index and dataset");
});
if (status != CUVS_SUCCESS) { return status; }

auto* box = reinterpret_cast<sg_cagra_c_api_index_box*>(index->addr);
if (box->layout == sg_cagra_c_api_index_box::dataset_layout::device_padded) {
return dispatch_update_device_dataset_same_layout(res, device_padded_dataset, index);
}
return dispatch_attach_dataset(res, device_padded_dataset, index);
return dispatch_update_dataset(res, device_padded_dataset, index);
}

/**
Expand Down
6 changes: 3 additions & 3 deletions c/src/neighbors/mg_cagra.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,16 +229,16 @@ void _mg_update_dataset(cuvsResources_t res,
using padded_ann_t = cuvs::neighbors::cagra::device_padded_index<T, uint32_t>;
auto* standard_index = reinterpret_cast<mg_cagra_index_t<T, standard_ann_t>*>(box->index_ptr);
auto* padded_index = new mg_cagra_index_t<T, padded_ann_t>(
cuvs::neighbors::cagra::attach_dataset(*res_ptr, *standard_index, padded_view));
cuvs::neighbors::cagra::update_dataset(
*res_ptr, std::move(*standard_index), padded_view));
auto* padded_box =
make_mg_cagra_box<T, padded_ann_t>(padded_index, mg_cagra_dataset_layout::device_padded);
destroy_mg_cagra_c_api_box(index->addr);
index->addr = reinterpret_cast<uintptr_t>(padded_box);
} else if (box->layout == mg_cagra_dataset_layout::device_padded) {
using padded_ann_t = cuvs::neighbors::cagra::device_padded_index<T, uint32_t>;
auto* padded_index = reinterpret_cast<mg_cagra_index_t<T, padded_ann_t>*>(box->index_ptr);
cuvs::neighbors::cagra::update_device_dataset_same_layout(
*res_ptr, *padded_index, padded_view);
cuvs::neighbors::cagra::update_dataset(*res_ptr, *padded_index, padded_view);
} else {
RAFT_FAIL("cuvsMultiGpuCagraUpdateDataset: unsupported index dataset layout");
}
Expand Down
36 changes: 18 additions & 18 deletions cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -244,8 +244,8 @@ void cuvs_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
}
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());
auto ace_index = cuvs::neighbors::cagra::update_dataset(
handle_, std::move(ace_host_index), padded->as_dataset_view());
*dataset_ = std::move(padded->data_);
index_ = std::make_shared<index_type>(std::move(ace_index));
} else {
Expand Down Expand Up @@ -283,13 +283,13 @@ void cuvs_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
*input_dataset_v_ = raft::make_device_matrix_view<const T, int64_t, raft::row_major>(
mds.data_handle(), static_cast<int64_t>(nrow), static_cast<int64_t>(dim_));
auto index = cuvs::neighbors::cagra::build(handle_, params, pdv);
index.update_device_dataset_same_layout(handle_, pdv);
index.update_dataset(handle_, pdv);
index_ = std::make_shared<index_type>(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);
index.update_dataset(handle_, view);
*dataset_ = std::move(padded->data_);
index_ = std::make_shared<index_type>(std::move(index));
}
Expand Down Expand Up @@ -318,17 +318,17 @@ void cuvs_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
raft::resource::get_cuda_stream(handle_));
cuvs::neighbors::device_padded_dataset_view<T, int64_t> dv(
raft::make_const_mdspan(sub_dataset_buffers_->back().view()), dim_);
sub_index.update_device_dataset_same_layout(handle_, dv);
sub_index.update_dataset(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);
sub_index.update_dataset(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<T, int64_t> pdv(
raft::make_const_mdspan(sub_dataset_buffers_->back().view()), dim_);
sub_index.update_device_dataset_same_layout(handle_, pdv);
sub_index.update_dataset(handle_, pdv);
}
}
}
Expand All @@ -348,8 +348,8 @@ void cuvs_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
}
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_index = cuvs::neighbors::cagra::update_dataset(
handle_, std::move(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<T, int64_t>(
Expand All @@ -370,12 +370,12 @@ void cuvs_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
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);
sub_index.update_dataset(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);
index.update_dataset(handle_, view);
sub_dataset_buffers_->push_back(std::move(padded_sub->data_));
sub_index = std::move(index);
}
Expand All @@ -391,12 +391,12 @@ void cuvs_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
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);
sub_index.update_dataset(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);
index.update_dataset(handle_, view);
sub_dataset_buffers_->push_back(std::move(padded_sub->data_));
sub_index = std::move(index);
}
Expand Down Expand Up @@ -484,7 +484,7 @@ void cuvs_cagra<T, IdxT>::set_search_param(const search_param_base& param,
*dataset_ = raft::make_device_matrix<T, int64_t>(handle_, 0, 0);
cuvs::neighbors::device_padded_dataset_view<T, int64_t> empty_dv(
raft::make_device_matrix_view(static_cast<T const*>(nullptr), 0, this->dim_), this->dim_);
index_->update_device_dataset_same_layout(handle_, empty_dv);
index_->update_dataset(handle_, empty_dv);

// Allocate space using the correct memory resource.
RAFT_LOG_DEBUG("moving dataset to new memory space: %s",
Expand All @@ -497,7 +497,7 @@ void cuvs_cagra<T, IdxT>::set_search_param(const search_param_base& param,
raft::make_device_matrix_view(
dataset_->data_handle(), dataset_->extent(0), dataset_->extent(1)),
this->dim_);
index_->update_device_dataset_same_layout(handle_, dv);
index_->update_dataset(handle_, dv);

need_dataset_update_ = false;
needs_dynamic_batcher_update = true;
Expand Down Expand Up @@ -554,17 +554,17 @@ void cuvs_cagra<T, IdxT>::set_search_dataset(const T* dataset, size_t nrow)
raft::resource::get_cuda_stream(handle_));
cuvs::neighbors::device_padded_dataset_view<T, int64_t> dv(
raft::make_const_mdspan(sub_dataset_buffers_->back().view()), dim_);
sub_index->update_device_dataset_same_layout(handle_, dv);
sub_index->update_dataset(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);
sub_index->update_dataset(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<T, int64_t> pdv(
raft::make_const_mdspan(sub_dataset_buffers_->back().view()), dim_);
sub_index->update_device_dataset_same_layout(handle_, pdv);
sub_index->update_dataset(handle_, pdv);
}
}
}
Expand Down
Loading
Loading