Skip to content

Commit 8ed3a4e

Browse files
Ludu-nuvaiclaude
Ludu-nuvai
and
claude
committed
feat(c-api): expose attach_dataset_on_build and cuvsCagraUpdateDataset
Cherry-picked from upstream PR NVIDIA#1842. Adds C API functions to control dataset attachment during CAGRA build, halving VRAM usage for large datasets. Also adds cuvsCagraUpdateDataset for replacing the dataset after build without rebuilding the graph. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent ee4ef8a commit 8ed3a4e

4 files changed

Lines changed: 307 additions & 2 deletions

File tree

‎c/include/cuvs/neighbors/cagra.h‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,19 @@ struct cuvsCagraIndexParams {
222222
* - Others: nullptr
223223
*/
224224
void* graph_build_params;
225+
/**
226+
* Whether to add the dataset content to the index after building the graph.
227+
*
228+
* - true (default): the index is filled with the dataset vectors and ready to search
229+
* after build, but requires enough memory to hold an aligned copy of the dataset.
230+
* - false: only the search graph is built. The user must call cuvsCagraUpdateDataset
231+
* to attach the dataset before searching. This avoids duplicating the dataset in
232+
* device memory during build, which is useful for memory-constrained scenarios.
233+
*
234+
* When compression is set, this parameter is ignored (compressed dataset is always
235+
* added to the index).
236+
*/
237+
bool attach_dataset_on_build;
225238
};
226239

227240
typedef struct cuvsCagraIndexParams* cuvsCagraIndexParams_t;
@@ -617,6 +630,26 @@ cuvsError_t cuvsCagraBuild(cuvsResources_t res,
617630
DLManagedTensor* dataset,
618631
cuvsCagraIndex_t index);
619632

633+
/**
634+
* @brief Update (attach) a dataset to an existing CAGRA index.
635+
*
636+
* This is intended for use after building an index with attach_dataset_on_build = false.
637+
* If the dataset rows are already aligned on 16 bytes and reside on the device, only a
638+
* reference is stored (zero-copy). Otherwise, an aligned copy is made.
639+
*
640+
* It is the caller's responsibility to ensure that the same dataset used for building
641+
* is supplied here. The dataset must remain valid for the lifetime of the index when
642+
* zero-copy is used.
643+
*
644+
* @param[in] res cuvsResources_t opaque C handle
645+
* @param[in] dataset DLManagedTensor* dataset to attach
646+
* @param[in] index cuvsCagraIndex_t the index to update
647+
* @return cuvsError_t
648+
*/
649+
cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res,
650+
DLManagedTensor* dataset,
651+
cuvsCagraIndex_t index);
652+
620653
/**
621654
* @}
622655
*/

‎c/src/neighbors/cagra.cpp‎

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,28 @@ void* _build(cuvsResources_t res, cuvsCagraIndexParams params, DLManagedTensor*
128128
return index;
129129
}
130130

131+
template <typename T>
132+
void _update_dataset(cuvsResources_t res,
133+
cuvsCagraIndex index,
134+
DLManagedTensor* dataset_tensor)
135+
{
136+
auto dataset = dataset_tensor->dl_tensor;
137+
auto res_ptr = reinterpret_cast<raft::resources*>(res);
138+
auto index_ptr = reinterpret_cast<cuvs::neighbors::cagra::index<T, uint32_t>*>(index.addr);
139+
140+
if (cuvs::core::is_dlpack_device_compatible(dataset)) {
141+
using mdspan_type = raft::device_matrix_view<T const, int64_t, raft::row_major>;
142+
auto mds = cuvs::core::from_dlpack<mdspan_type>(dataset_tensor);
143+
index_ptr->update_dataset(*res_ptr, mds);
144+
} else if (cuvs::core::is_dlpack_host_compatible(dataset)) {
145+
using mdspan_type = raft::host_matrix_view<T const, int64_t, raft::row_major>;
146+
auto mds = cuvs::core::from_dlpack<mdspan_type>(dataset_tensor);
147+
index_ptr->update_dataset(*res_ptr, mds);
148+
} else {
149+
RAFT_FAIL("Unsupported dataset DLtensor device type: %d", dataset.device.device_type);
150+
}
151+
}
152+
131153
template <typename T>
132154
void* _from_args(cuvsResources_t res,
133155
cuvsDistanceType _metric,
@@ -442,6 +464,7 @@ void convert_c_index_params(cuvsCagraIndexParams params,
442464
out->metric = static_cast<cuvs::distance::DistanceType>((int)params.metric);
443465
out->intermediate_graph_degree = params.intermediate_graph_degree;
444466
out->graph_degree = params.graph_degree;
467+
out->attach_dataset_on_build = params.attach_dataset_on_build;
445468
_set_graph_build_params(out->graph_build_params, params, params.build_algo, n_rows, dim);
446469

447470
if (auto* cparams = params.compression; cparams != nullptr) {
@@ -588,6 +611,27 @@ extern "C" cuvsError_t cuvsCagraBuild(cuvsResources_t res,
588611
});
589612
}
590613

614+
extern "C" cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res,
615+
DLManagedTensor* dataset_tensor,
616+
cuvsCagraIndex_t index)
617+
{
618+
return cuvs::core::translate_exceptions([=] {
619+
if (index->dtype.code == kDLFloat && index->dtype.bits == 32) {
620+
_update_dataset<float>(res, *index, dataset_tensor);
621+
} else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) {
622+
_update_dataset<half>(res, *index, dataset_tensor);
623+
} else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
624+
_update_dataset<int8_t>(res, *index, dataset_tensor);
625+
} else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
626+
_update_dataset<uint8_t>(res, *index, dataset_tensor);
627+
} else {
628+
RAFT_FAIL("Unsupported index dtype: %d and bits: %d",
629+
index->dtype.code,
630+
index->dtype.bits);
631+
}
632+
});
633+
}
634+
591635
extern "C" cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res,
592636
cuvsDistanceType metric,
593637
DLManagedTensor* graph_tensor,
@@ -736,7 +780,10 @@ extern "C" cuvsError_t cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t* params
736780
.intermediate_graph_degree = 128,
737781
.graph_degree = 64,
738782
.build_algo = IVF_PQ,
739-
.nn_descent_niter = 20};
783+
.nn_descent_niter = 20,
784+
.compression = nullptr,
785+
.graph_build_params = nullptr,
786+
.attach_dataset_on_build = true};
740787
(*params)->graph_build_params = new cuvsIvfPqParams{nullptr, nullptr, 1};
741788
});
742789
}

‎c/src/neighbors/tiered_index.cpp‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

@@ -71,6 +71,9 @@ void* _build(cuvsResources_t res, cuvsTieredIndexParams params, DLManagedTensor*
7171
case CUVS_TIERED_INDEX_ALGO_CAGRA: {
7272
auto build_params = tiered_index::index_params<cagra::index_params>();
7373
convert_c_index_params(params, dataset.shape[0], dataset.shape[1], &build_params);
74+
// The tiered index sub-CAGRA always needs the dataset attached for search.
75+
// Force this in case the caller did not set the field (e.g. zero-initialized struct).
76+
build_params.attach_dataset_on_build = true;
7477
return new tiered_index::index<cagra::index<T, uint32_t>>(
7578
tiered_index::build(*res_ptr, build_params, mds));
7679
}

‎c/tests/neighbors/ann_cagra_c.cu‎

Lines changed: 222 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -337,6 +337,228 @@ TEST(CagraC, BuildExtendSearch)
337337
cuvsResourcesDestroy(res);
338338
}
339339

340+
TEST(CagraC, BuildNoDatasetThenUpdateAndSearch)
341+
{
342+
// Test the attach_dataset_on_build = false workflow:
343+
// 1. Build index without attaching dataset (saves a full dataset copy)
344+
// 2. Attach dataset via cuvsCagraUpdateDataset
345+
// 3. Search and verify correctness
346+
347+
// create cuvsResources_t
348+
cuvsResources_t res;
349+
cuvsResourcesCreate(&res);
350+
cudaStream_t stream;
351+
cuvsStreamGet(res, &stream);
352+
353+
// create dataset DLTensor
354+
DLManagedTensor dataset_tensor;
355+
dataset_tensor.dl_tensor.data = dataset;
356+
dataset_tensor.dl_tensor.device.device_type = kDLCPU;
357+
dataset_tensor.dl_tensor.ndim = 2;
358+
dataset_tensor.dl_tensor.dtype.code = kDLFloat;
359+
dataset_tensor.dl_tensor.dtype.bits = 32;
360+
dataset_tensor.dl_tensor.dtype.lanes = 1;
361+
int64_t dataset_shape[2] = {4, 2};
362+
dataset_tensor.dl_tensor.shape = dataset_shape;
363+
dataset_tensor.dl_tensor.strides = nullptr;
364+
365+
// create index
366+
cuvsCagraIndex_t index;
367+
cuvsCagraIndexCreate(&index);
368+
369+
// build index with attach_dataset_on_build = false
370+
cuvsCagraIndexParams_t build_params;
371+
cuvsCagraIndexParamsCreate(&build_params);
372+
build_params->attach_dataset_on_build = false;
373+
ASSERT_EQ(cuvsCagraBuild(res, build_params, &dataset_tensor, index), CUVS_SUCCESS);
374+
375+
// now attach the dataset
376+
ASSERT_EQ(cuvsCagraUpdateDataset(res, &dataset_tensor, index), CUVS_SUCCESS);
377+
378+
// create queries DLTensor
379+
rmm::device_uvector<float> queries_d(4 * 2, stream);
380+
raft::copy(queries_d.data(), (float*)queries, 4 * 2, stream);
381+
382+
DLManagedTensor queries_tensor;
383+
queries_tensor.dl_tensor.data = queries_d.data();
384+
queries_tensor.dl_tensor.device.device_type = kDLCUDA;
385+
queries_tensor.dl_tensor.ndim = 2;
386+
queries_tensor.dl_tensor.dtype.code = kDLFloat;
387+
queries_tensor.dl_tensor.dtype.bits = 32;
388+
queries_tensor.dl_tensor.dtype.lanes = 1;
389+
int64_t queries_shape[2] = {4, 2};
390+
queries_tensor.dl_tensor.shape = queries_shape;
391+
queries_tensor.dl_tensor.strides = nullptr;
392+
393+
// create neighbors DLTensor
394+
rmm::device_uvector<uint32_t> neighbors_d(4, stream);
395+
396+
DLManagedTensor neighbors_tensor;
397+
neighbors_tensor.dl_tensor.data = neighbors_d.data();
398+
neighbors_tensor.dl_tensor.device.device_type = kDLCUDA;
399+
neighbors_tensor.dl_tensor.ndim = 2;
400+
neighbors_tensor.dl_tensor.dtype.code = kDLUInt;
401+
neighbors_tensor.dl_tensor.dtype.bits = 32;
402+
neighbors_tensor.dl_tensor.dtype.lanes = 1;
403+
int64_t neighbors_shape[2] = {4, 1};
404+
neighbors_tensor.dl_tensor.shape = neighbors_shape;
405+
neighbors_tensor.dl_tensor.strides = nullptr;
406+
407+
// create distances DLTensor
408+
rmm::device_uvector<float> distances_d(4, stream);
409+
410+
DLManagedTensor distances_tensor;
411+
distances_tensor.dl_tensor.data = distances_d.data();
412+
distances_tensor.dl_tensor.device.device_type = kDLCUDA;
413+
distances_tensor.dl_tensor.ndim = 2;
414+
distances_tensor.dl_tensor.dtype.code = kDLFloat;
415+
distances_tensor.dl_tensor.dtype.bits = 32;
416+
distances_tensor.dl_tensor.dtype.lanes = 1;
417+
int64_t distances_shape[2] = {4, 1};
418+
distances_tensor.dl_tensor.shape = distances_shape;
419+
distances_tensor.dl_tensor.strides = nullptr;
420+
421+
cuvsFilter filter;
422+
filter.type = NO_FILTER;
423+
filter.addr = (uintptr_t)NULL;
424+
425+
// search index
426+
cuvsCagraSearchParams_t search_params;
427+
cuvsCagraSearchParamsCreate(&search_params);
428+
cuvsCagraSearch(
429+
res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter);
430+
431+
// verify output — should match the standard BuildSearch test results
432+
ASSERT_TRUE(
433+
cuvs::devArrMatchHost(neighbors_exp, neighbors_d.data(), 4, cuvs::Compare<uint32_t>()));
434+
ASSERT_TRUE(cuvs::devArrMatchHost(
435+
distances_exp, distances_d.data(), 4, cuvs::CompareApprox<float>(0.001f)));
436+
437+
// de-allocate index and res
438+
cuvsCagraSearchParamsDestroy(search_params);
439+
cuvsCagraIndexParamsDestroy(build_params);
440+
cuvsCagraIndexDestroy(index);
441+
cuvsResourcesDestroy(res);
442+
}
443+
444+
TEST(CagraC, BuildNoDatasetThenUpdateDeviceAndSearch)
445+
{
446+
// Test the motivating scenario: dataset already on device (kDLCUDA).
447+
// Using attach_dataset_on_build = false avoids duplicating the device dataset,
448+
// then cuvsCagraUpdateDataset attaches it (zero-copy when properly aligned).
449+
450+
// create cuvsResources_t
451+
cuvsResources_t res;
452+
cuvsResourcesCreate(&res);
453+
cudaStream_t stream;
454+
cuvsStreamGet(res, &stream);
455+
456+
// copy dataset to device memory (simulating a dataset that is already on GPU)
457+
rmm::device_uvector<float> dataset_d(4 * 2, stream);
458+
raft::copy(dataset_d.data(), (float*)dataset, 4 * 2, stream);
459+
460+
// create dataset DLTensor on CPU for building the graph
461+
DLManagedTensor dataset_tensor;
462+
dataset_tensor.dl_tensor.data = dataset;
463+
dataset_tensor.dl_tensor.device.device_type = kDLCPU;
464+
dataset_tensor.dl_tensor.ndim = 2;
465+
dataset_tensor.dl_tensor.dtype.code = kDLFloat;
466+
dataset_tensor.dl_tensor.dtype.bits = 32;
467+
dataset_tensor.dl_tensor.dtype.lanes = 1;
468+
int64_t dataset_shape[2] = {4, 2};
469+
dataset_tensor.dl_tensor.shape = dataset_shape;
470+
dataset_tensor.dl_tensor.strides = nullptr;
471+
472+
// create index
473+
cuvsCagraIndex_t index;
474+
cuvsCagraIndexCreate(&index);
475+
476+
// build index with attach_dataset_on_build = false
477+
cuvsCagraIndexParams_t build_params;
478+
cuvsCagraIndexParamsCreate(&build_params);
479+
build_params->attach_dataset_on_build = false;
480+
ASSERT_EQ(cuvsCagraBuild(res, build_params, &dataset_tensor, index), CUVS_SUCCESS);
481+
482+
// attach the device-resident dataset via cuvsCagraUpdateDataset (kDLCUDA path)
483+
DLManagedTensor device_dataset_tensor;
484+
device_dataset_tensor.dl_tensor.data = dataset_d.data();
485+
device_dataset_tensor.dl_tensor.device.device_type = kDLCUDA;
486+
device_dataset_tensor.dl_tensor.device.device_id = 0;
487+
device_dataset_tensor.dl_tensor.ndim = 2;
488+
device_dataset_tensor.dl_tensor.dtype.code = kDLFloat;
489+
device_dataset_tensor.dl_tensor.dtype.bits = 32;
490+
device_dataset_tensor.dl_tensor.dtype.lanes = 1;
491+
device_dataset_tensor.dl_tensor.shape = dataset_shape;
492+
device_dataset_tensor.dl_tensor.strides = nullptr;
493+
494+
ASSERT_EQ(cuvsCagraUpdateDataset(res, &device_dataset_tensor, index), CUVS_SUCCESS);
495+
496+
// create queries DLTensor
497+
rmm::device_uvector<float> queries_d(4 * 2, stream);
498+
raft::copy(queries_d.data(), (float*)queries, 4 * 2, stream);
499+
500+
DLManagedTensor queries_tensor;
501+
queries_tensor.dl_tensor.data = queries_d.data();
502+
queries_tensor.dl_tensor.device.device_type = kDLCUDA;
503+
queries_tensor.dl_tensor.ndim = 2;
504+
queries_tensor.dl_tensor.dtype.code = kDLFloat;
505+
queries_tensor.dl_tensor.dtype.bits = 32;
506+
queries_tensor.dl_tensor.dtype.lanes = 1;
507+
int64_t queries_shape[2] = {4, 2};
508+
queries_tensor.dl_tensor.shape = queries_shape;
509+
queries_tensor.dl_tensor.strides = nullptr;
510+
511+
// create neighbors DLTensor
512+
rmm::device_uvector<uint32_t> neighbors_d(4, stream);
513+
514+
DLManagedTensor neighbors_tensor;
515+
neighbors_tensor.dl_tensor.data = neighbors_d.data();
516+
neighbors_tensor.dl_tensor.device.device_type = kDLCUDA;
517+
neighbors_tensor.dl_tensor.ndim = 2;
518+
neighbors_tensor.dl_tensor.dtype.code = kDLUInt;
519+
neighbors_tensor.dl_tensor.dtype.bits = 32;
520+
neighbors_tensor.dl_tensor.dtype.lanes = 1;
521+
int64_t neighbors_shape[2] = {4, 1};
522+
neighbors_tensor.dl_tensor.shape = neighbors_shape;
523+
neighbors_tensor.dl_tensor.strides = nullptr;
524+
525+
// create distances DLTensor
526+
rmm::device_uvector<float> distances_d(4, stream);
527+
528+
DLManagedTensor distances_tensor;
529+
distances_tensor.dl_tensor.data = distances_d.data();
530+
distances_tensor.dl_tensor.device.device_type = kDLCUDA;
531+
distances_tensor.dl_tensor.ndim = 2;
532+
distances_tensor.dl_tensor.dtype.code = kDLFloat;
533+
distances_tensor.dl_tensor.dtype.bits = 32;
534+
distances_tensor.dl_tensor.dtype.lanes = 1;
535+
int64_t distances_shape[2] = {4, 1};
536+
distances_tensor.dl_tensor.shape = distances_shape;
537+
distances_tensor.dl_tensor.strides = nullptr;
538+
539+
cuvsFilter filter;
540+
filter.type = NO_FILTER;
541+
filter.addr = (uintptr_t)NULL;
542+
543+
// search index
544+
cuvsCagraSearchParams_t search_params;
545+
cuvsCagraSearchParamsCreate(&search_params);
546+
cuvsCagraSearch(
547+
res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter);
548+
549+
// verify output — should match the standard BuildSearch test results
550+
ASSERT_TRUE(
551+
cuvs::devArrMatchHost(neighbors_exp, neighbors_d.data(), 4, cuvs::Compare<uint32_t>()));
552+
ASSERT_TRUE(cuvs::devArrMatchHost(
553+
distances_exp, distances_d.data(), 4, cuvs::CompareApprox<float>(0.001f)));
554+
555+
// de-allocate index and res
556+
cuvsCagraSearchParamsDestroy(search_params);
557+
cuvsCagraIndexParamsDestroy(build_params);
558+
cuvsCagraIndexDestroy(index);
559+
cuvsResourcesDestroy(res);
560+
}
561+
340562
TEST(CagraC, BuildSearchFiltered)
341563
{
342564
// create cuvsResources_t

0 commit comments

Comments
 (0)