diff --git a/c/include/cuvs/neighbors/cagra.h b/c/include/cuvs/neighbors/cagra.h index 22809da37e..afae152d84 100644 --- a/c/include/cuvs/neighbors/cagra.h +++ b/c/include/cuvs/neighbors/cagra.h @@ -224,6 +224,19 @@ struct cuvsCagraIndexParams { * - Others: nullptr */ void* graph_build_params; + /** + * Whether to add the dataset content to the index after building the graph. + * + * - true (default): the index is filled with the dataset vectors and ready to search + * after build, but requires enough memory to hold an aligned copy of the dataset. + * - false: only the search graph is built. The user must call cuvsCagraUpdateDataset + * to attach the dataset before searching. This avoids duplicating the dataset in + * device memory during build, which is useful for memory-constrained scenarios. + * + * When compression is set, this parameter is ignored (compressed dataset is always + * added to the index). + */ + bool attach_dataset_on_build; }; typedef struct cuvsCagraIndexParams* cuvsCagraIndexParams_t; @@ -619,6 +632,26 @@ CUVS_EXPORT cuvsError_t cuvsCagraBuild(cuvsResources_t res, DLManagedTensor* dataset, cuvsCagraIndex_t index); +/** + * @brief Update (attach) a dataset to an existing CAGRA index. + * + * This is intended for use after building an index with attach_dataset_on_build = false. + * If the dataset rows are already aligned on 16 bytes and reside on the device, only a + * reference is stored (zero-copy). Otherwise, an aligned copy is made. + * + * It is the caller's responsibility to ensure that the same dataset used for building + * is supplied here. The dataset must remain valid for the lifetime of the index when + * zero-copy is used. + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] dataset DLManagedTensor* dataset to attach + * @param[in] index cuvsCagraIndex_t the index to update + * @return cuvsError_t + */ +cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res, + DLManagedTensor* dataset, + cuvsCagraIndex_t index); + /** * @} */ diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 081179ca46..172c1139e6 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -128,6 +128,28 @@ void* _build(cuvsResources_t res, cuvsCagraIndexParams params, DLManagedTensor* return index; } +template +void _update_dataset(cuvsResources_t res, + cuvsCagraIndex index, + DLManagedTensor* dataset_tensor) +{ + auto dataset = dataset_tensor->dl_tensor; + auto res_ptr = reinterpret_cast(res); + auto index_ptr = reinterpret_cast*>(index.addr); + + if (cuvs::core::is_dlpack_device_compatible(dataset)) { + using mdspan_type = raft::device_matrix_view; + auto mds = cuvs::core::from_dlpack(dataset_tensor); + index_ptr->update_dataset(*res_ptr, mds); + } else if (cuvs::core::is_dlpack_host_compatible(dataset)) { + using mdspan_type = raft::host_matrix_view; + auto mds = cuvs::core::from_dlpack(dataset_tensor); + index_ptr->update_dataset(*res_ptr, mds); + } else { + RAFT_FAIL("Unsupported dataset DLtensor device type: %d", dataset.device.device_type); + } +} + template void* _from_args(cuvsResources_t res, cuvsDistanceType _metric, @@ -442,6 +464,7 @@ void convert_c_index_params(cuvsCagraIndexParams params, out->metric = static_cast((int)params.metric); out->intermediate_graph_degree = params.intermediate_graph_degree; out->graph_degree = params.graph_degree; + out->attach_dataset_on_build = params.attach_dataset_on_build; _set_graph_build_params(out->graph_build_params, params, params.build_algo, n_rows, dim); if (auto* cparams = params.compression; cparams != nullptr) { @@ -588,6 +611,27 @@ extern "C" cuvsError_t cuvsCagraBuild(cuvsResources_t res, }); } +extern "C" cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res, + DLManagedTensor* dataset_tensor, + cuvsCagraIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + _update_dataset(res, *index, dataset_tensor); + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + _update_dataset(res, *index, dataset_tensor); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + _update_dataset(res, *index, dataset_tensor); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + _update_dataset(res, *index, dataset_tensor); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", + index->dtype.code, + index->dtype.bits); + } + }); +} + extern "C" cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res, cuvsDistanceType metric, DLManagedTensor* graph_tensor, @@ -736,7 +780,10 @@ extern "C" cuvsError_t cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t* params .intermediate_graph_degree = 128, .graph_degree = 64, .build_algo = IVF_PQ, - .nn_descent_niter = 20}; + .nn_descent_niter = 20, + .compression = nullptr, + .graph_build_params = nullptr, + .attach_dataset_on_build = true}; (*params)->graph_build_params = new cuvsIvfPqParams{nullptr, nullptr, 1}; }); } diff --git a/c/src/neighbors/tiered_index.cpp b/c/src/neighbors/tiered_index.cpp index 2a7d54b16d..1c21b0876a 100644 --- a/c/src/neighbors/tiered_index.cpp +++ b/c/src/neighbors/tiered_index.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -71,6 +71,9 @@ void* _build(cuvsResources_t res, cuvsTieredIndexParams params, DLManagedTensor* case CUVS_TIERED_INDEX_ALGO_CAGRA: { auto build_params = tiered_index::index_params(); convert_c_index_params(params, dataset.shape[0], dataset.shape[1], &build_params); + // The tiered index sub-CAGRA always needs the dataset attached for search. + // Force this in case the caller did not set the field (e.g. zero-initialized struct). + build_params.attach_dataset_on_build = true; return new tiered_index::index>( tiered_index::build(*res_ptr, build_params, mds)); } diff --git a/c/tests/neighbors/ann_cagra_c.cu b/c/tests/neighbors/ann_cagra_c.cu index 9c14bbea7d..5f9e95dc4a 100644 --- a/c/tests/neighbors/ann_cagra_c.cu +++ b/c/tests/neighbors/ann_cagra_c.cu @@ -337,6 +337,228 @@ TEST(CagraC, BuildExtendSearch) cuvsResourcesDestroy(res); } +TEST(CagraC, BuildNoDatasetThenUpdateAndSearch) +{ + // Test the attach_dataset_on_build = false workflow: + // 1. Build index without attaching dataset (saves a full dataset copy) + // 2. Attach dataset via cuvsCagraUpdateDataset + // 3. Search and verify correctness + + // create cuvsResources_t + cuvsResources_t res; + cuvsResourcesCreate(&res); + cudaStream_t stream; + cuvsStreamGet(res, &stream); + + // create dataset DLTensor + DLManagedTensor dataset_tensor; + dataset_tensor.dl_tensor.data = dataset; + dataset_tensor.dl_tensor.device.device_type = kDLCPU; + dataset_tensor.dl_tensor.ndim = 2; + dataset_tensor.dl_tensor.dtype.code = kDLFloat; + dataset_tensor.dl_tensor.dtype.bits = 32; + dataset_tensor.dl_tensor.dtype.lanes = 1; + int64_t dataset_shape[2] = {4, 2}; + dataset_tensor.dl_tensor.shape = dataset_shape; + dataset_tensor.dl_tensor.strides = nullptr; + + // create index + cuvsCagraIndex_t index; + cuvsCagraIndexCreate(&index); + + // build index with attach_dataset_on_build = false + cuvsCagraIndexParams_t build_params; + cuvsCagraIndexParamsCreate(&build_params); + build_params->attach_dataset_on_build = false; + ASSERT_EQ(cuvsCagraBuild(res, build_params, &dataset_tensor, index), CUVS_SUCCESS); + + // now attach the dataset + ASSERT_EQ(cuvsCagraUpdateDataset(res, &dataset_tensor, index), CUVS_SUCCESS); + + // create queries DLTensor + rmm::device_uvector queries_d(4 * 2, stream); + raft::copy(queries_d.data(), (float*)queries, 4 * 2, stream); + + DLManagedTensor queries_tensor; + queries_tensor.dl_tensor.data = queries_d.data(); + queries_tensor.dl_tensor.device.device_type = kDLCUDA; + queries_tensor.dl_tensor.ndim = 2; + queries_tensor.dl_tensor.dtype.code = kDLFloat; + queries_tensor.dl_tensor.dtype.bits = 32; + queries_tensor.dl_tensor.dtype.lanes = 1; + int64_t queries_shape[2] = {4, 2}; + queries_tensor.dl_tensor.shape = queries_shape; + queries_tensor.dl_tensor.strides = nullptr; + + // create neighbors DLTensor + rmm::device_uvector neighbors_d(4, stream); + + DLManagedTensor neighbors_tensor; + neighbors_tensor.dl_tensor.data = neighbors_d.data(); + neighbors_tensor.dl_tensor.device.device_type = kDLCUDA; + neighbors_tensor.dl_tensor.ndim = 2; + neighbors_tensor.dl_tensor.dtype.code = kDLUInt; + neighbors_tensor.dl_tensor.dtype.bits = 32; + neighbors_tensor.dl_tensor.dtype.lanes = 1; + int64_t neighbors_shape[2] = {4, 1}; + neighbors_tensor.dl_tensor.shape = neighbors_shape; + neighbors_tensor.dl_tensor.strides = nullptr; + + // create distances DLTensor + rmm::device_uvector distances_d(4, stream); + + DLManagedTensor distances_tensor; + distances_tensor.dl_tensor.data = distances_d.data(); + distances_tensor.dl_tensor.device.device_type = kDLCUDA; + distances_tensor.dl_tensor.ndim = 2; + distances_tensor.dl_tensor.dtype.code = kDLFloat; + distances_tensor.dl_tensor.dtype.bits = 32; + distances_tensor.dl_tensor.dtype.lanes = 1; + int64_t distances_shape[2] = {4, 1}; + distances_tensor.dl_tensor.shape = distances_shape; + distances_tensor.dl_tensor.strides = nullptr; + + cuvsFilter filter; + filter.type = NO_FILTER; + filter.addr = (uintptr_t)NULL; + + // search index + cuvsCagraSearchParams_t search_params; + cuvsCagraSearchParamsCreate(&search_params); + cuvsCagraSearch( + res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter); + + // verify output — should match the standard BuildSearch test results + ASSERT_TRUE( + cuvs::devArrMatchHost(neighbors_exp, neighbors_d.data(), 4, cuvs::Compare())); + ASSERT_TRUE(cuvs::devArrMatchHost( + distances_exp, distances_d.data(), 4, cuvs::CompareApprox(0.001f))); + + // de-allocate index and res + cuvsCagraSearchParamsDestroy(search_params); + cuvsCagraIndexParamsDestroy(build_params); + cuvsCagraIndexDestroy(index); + cuvsResourcesDestroy(res); +} + +TEST(CagraC, BuildNoDatasetThenUpdateDeviceAndSearch) +{ + // Test the motivating scenario: dataset already on device (kDLCUDA). + // Using attach_dataset_on_build = false avoids duplicating the device dataset, + // then cuvsCagraUpdateDataset attaches it (zero-copy when properly aligned). + + // create cuvsResources_t + cuvsResources_t res; + cuvsResourcesCreate(&res); + cudaStream_t stream; + cuvsStreamGet(res, &stream); + + // copy dataset to device memory (simulating a dataset that is already on GPU) + rmm::device_uvector dataset_d(4 * 2, stream); + raft::copy(dataset_d.data(), (float*)dataset, 4 * 2, stream); + + // create dataset DLTensor on CPU for building the graph + DLManagedTensor dataset_tensor; + dataset_tensor.dl_tensor.data = dataset; + dataset_tensor.dl_tensor.device.device_type = kDLCPU; + dataset_tensor.dl_tensor.ndim = 2; + dataset_tensor.dl_tensor.dtype.code = kDLFloat; + dataset_tensor.dl_tensor.dtype.bits = 32; + dataset_tensor.dl_tensor.dtype.lanes = 1; + int64_t dataset_shape[2] = {4, 2}; + dataset_tensor.dl_tensor.shape = dataset_shape; + dataset_tensor.dl_tensor.strides = nullptr; + + // create index + cuvsCagraIndex_t index; + cuvsCagraIndexCreate(&index); + + // build index with attach_dataset_on_build = false + cuvsCagraIndexParams_t build_params; + cuvsCagraIndexParamsCreate(&build_params); + build_params->attach_dataset_on_build = false; + ASSERT_EQ(cuvsCagraBuild(res, build_params, &dataset_tensor, index), CUVS_SUCCESS); + + // attach the device-resident dataset via cuvsCagraUpdateDataset (kDLCUDA path) + DLManagedTensor device_dataset_tensor; + device_dataset_tensor.dl_tensor.data = dataset_d.data(); + device_dataset_tensor.dl_tensor.device.device_type = kDLCUDA; + device_dataset_tensor.dl_tensor.device.device_id = 0; + device_dataset_tensor.dl_tensor.ndim = 2; + device_dataset_tensor.dl_tensor.dtype.code = kDLFloat; + device_dataset_tensor.dl_tensor.dtype.bits = 32; + device_dataset_tensor.dl_tensor.dtype.lanes = 1; + device_dataset_tensor.dl_tensor.shape = dataset_shape; + device_dataset_tensor.dl_tensor.strides = nullptr; + + ASSERT_EQ(cuvsCagraUpdateDataset(res, &device_dataset_tensor, index), CUVS_SUCCESS); + + // create queries DLTensor + rmm::device_uvector queries_d(4 * 2, stream); + raft::copy(queries_d.data(), (float*)queries, 4 * 2, stream); + + DLManagedTensor queries_tensor; + queries_tensor.dl_tensor.data = queries_d.data(); + queries_tensor.dl_tensor.device.device_type = kDLCUDA; + queries_tensor.dl_tensor.ndim = 2; + queries_tensor.dl_tensor.dtype.code = kDLFloat; + queries_tensor.dl_tensor.dtype.bits = 32; + queries_tensor.dl_tensor.dtype.lanes = 1; + int64_t queries_shape[2] = {4, 2}; + queries_tensor.dl_tensor.shape = queries_shape; + queries_tensor.dl_tensor.strides = nullptr; + + // create neighbors DLTensor + rmm::device_uvector neighbors_d(4, stream); + + DLManagedTensor neighbors_tensor; + neighbors_tensor.dl_tensor.data = neighbors_d.data(); + neighbors_tensor.dl_tensor.device.device_type = kDLCUDA; + neighbors_tensor.dl_tensor.ndim = 2; + neighbors_tensor.dl_tensor.dtype.code = kDLUInt; + neighbors_tensor.dl_tensor.dtype.bits = 32; + neighbors_tensor.dl_tensor.dtype.lanes = 1; + int64_t neighbors_shape[2] = {4, 1}; + neighbors_tensor.dl_tensor.shape = neighbors_shape; + neighbors_tensor.dl_tensor.strides = nullptr; + + // create distances DLTensor + rmm::device_uvector distances_d(4, stream); + + DLManagedTensor distances_tensor; + distances_tensor.dl_tensor.data = distances_d.data(); + distances_tensor.dl_tensor.device.device_type = kDLCUDA; + distances_tensor.dl_tensor.ndim = 2; + distances_tensor.dl_tensor.dtype.code = kDLFloat; + distances_tensor.dl_tensor.dtype.bits = 32; + distances_tensor.dl_tensor.dtype.lanes = 1; + int64_t distances_shape[2] = {4, 1}; + distances_tensor.dl_tensor.shape = distances_shape; + distances_tensor.dl_tensor.strides = nullptr; + + cuvsFilter filter; + filter.type = NO_FILTER; + filter.addr = (uintptr_t)NULL; + + // search index + cuvsCagraSearchParams_t search_params; + cuvsCagraSearchParamsCreate(&search_params); + cuvsCagraSearch( + res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter); + + // verify output — should match the standard BuildSearch test results + ASSERT_TRUE( + cuvs::devArrMatchHost(neighbors_exp, neighbors_d.data(), 4, cuvs::Compare())); + ASSERT_TRUE(cuvs::devArrMatchHost( + distances_exp, distances_d.data(), 4, cuvs::CompareApprox(0.001f))); + + // de-allocate index and res + cuvsCagraSearchParamsDestroy(search_params); + cuvsCagraIndexParamsDestroy(build_params); + cuvsCagraIndexDestroy(index); + cuvsResourcesDestroy(res); +} + TEST(CagraC, BuildSearchFiltered) { // create cuvsResources_t