diff --git a/c/include/cuvs/neighbors/cagra.h b/c/include/cuvs/neighbors/cagra.h index 25555867b9..3b85ba6542 100644 --- a/c/include/cuvs/neighbors/cagra.h +++ b/c/include/cuvs/neighbors/cagra.h @@ -992,9 +992,39 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res, * @{ */ +/** + * @brief Compute per-index write offsets for a merged dataset buffer. + * + * `cuvsCagraMerge`/`cuvsCagraMergeWithParams` require the caller to have already concatenated + * every input index's dataset (in `indices` order, applying `filter` if any) into a single buffer + * and to know each index's starting row within it. For `filter.type == NO_FILTER`, those offsets + * are just the cumulative sizes of `indices` and this function is not needed. For `BITSET`, the + * number of surviving rows per index cannot be derived any other way, so call this first. + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] indices Array of input cuvsCagraIndex_t handles that will be passed to merge + * @param[in] num_indices Number of input indices + * @param[in] filter Filter that will be passed to merge. Only `NO_FILTER` and `BITSET` supported. + * @param[out] offsets Caller-allocated array of `num_indices + 1` int64_t. Entry `i` is the row at + * which `indices[i]`'s surviving rows must start in the merged buffer; the + * last entry is the total row count of the merged buffer. + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsCagraMergedDatasetOffsets(cuvsResources_t res, + cuvsCagraIndex_t* indices, + size_t num_indices, + cuvsFilter filter, + int64_t* offsets); + /** * @brief Merge multiple CAGRA indices into a single CAGRA index. * + * The caller is responsible for concatenating every input index's dataset (applying `filter` if + * any) into a single `merged_dataset` buffer before calling this, and for computing `offsets` + * (see `cuvsCagraMergedDatasetOffsets`). This function only builds/merges the graph and rebinds + * the output index to `merged_dataset` -- it never allocates or copies dataset rows itself. This + * mirrors the `cuvsCagraExtend` contract. + * * All input indices must have been built with the same data type (`index.dtype`) and * have the same dimensionality (`index.dims`). The merged index uses the output * parameters specified in `cuvsCagraIndexParams`. The merge algorithm is selected automatically. @@ -1027,12 +1057,14 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res, * cuvsCagraIndexParams_t merge_params; * cuvsError_t params_create_status = cuvsCagraIndexParamsCreate(&merge_params); * - * cuvsDataset_t merged_dataset; - * cuvsDatasetCreate(&merged_dataset); + * // Build `merged_dataset` as the caller-owned concatenation of index1 || index2 (e.g. via + * // cuvsDatasetMakePadded over a device buffer you populated yourself). + * cuvsDataset_t merged_dataset = ...; + * int64_t offsets[3] = {0, index1_size, index1_size + index2_size}; * cuvsFilter filter = {.type = NO_FILTER, .addr = 0}; * * cuvsError_t merge_status = cuvsCagraMerge(res, merge_params, (cuvsCagraIndex_t[]){index1, - * index2}, 2, filter, merged_dataset, merged_index); + * index2}, 2, filter, merged_dataset, offsets, merged_index); * * // Use merged_index for search operations * @@ -1046,13 +1078,15 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res, * @param[in] params cuvsCagraIndexParams_t parameters for the output index * @param[in] indices Array of input cuvsCagraIndex_t handles to merge * @param[in] num_indices Number of input indices - * @param[in] filter Filter that can be used to filter out vectors from the merged index - * @param[out] merged_dataset Empty owning dataset handle. Merge first attempts to allocate and - * populate device storage with the same layout as the input indices. For - * an unfiltered merge, if device allocation fails, it falls back to host - * storage and returns a host-backed output index. Keep this dataset alive - * while using \p output_index. A host-backed output index must be updated - * with `cuvsCagraUpdateDataset` before device search. + * @param[in] filter Filter, already applied by the caller while building `merged_dataset` + * @param[in] merged_dataset Caller-owned dataset handle already containing the concatenated (and, + * if `filter` is set, already-filtered) dataset, with the same layout as + * the input indices. Keep this dataset alive while using + * \p output_index. A host-backed dataset must be updated with + * `cuvsCagraUpdateDataset` before device search. + * @param[in] offsets Per-index starting row within `merged_dataset`, as returned by + * `cuvsCagraMergedDatasetOffsets`. Array of `num_indices + 1` int64_t; the last + * entry must equal `merged_dataset`'s row count. * @param[out] output_index Output handle that will store the merged index. * Must be initialized using `cuvsCagraIndexCreate` before use. */ @@ -1062,25 +1096,27 @@ CUVS_EXPORT cuvsError_t cuvsCagraMerge(cuvsResources_t res, size_t num_indices, cuvsFilter filter, cuvsDataset_t merged_dataset, + const int64_t* offsets, cuvsCagraIndex_t output_index); /** * @brief Merge multiple CAGRA indices with explicit merge parameters. * + * See `cuvsCagraMerge` for the full `merged_dataset`/`offsets` contract. + * * @param[in] res cuvsResources_t opaque C handle * @param[in] params cuvsCagraIndexParams_t parameters for the output index * @param[in] merge_params cuvsCagraMergeParams_t parameters controlling the merge algorithm, or * NULL to use AUTO defaults * @param[in] indices Array of input cuvsCagraIndex_t handles to merge * @param[in] num_indices Number of input indices - * @param[in] filter Filter that can be used to filter out vectors from the merged index - * @param[out] merged_dataset Empty owning dataset handle. Merge first attempts to allocate and - * populate device storage with the same layout as the input indices. For - * an unfiltered merge, AUTO and REBUILD can fall back to host storage if - * device allocation fails; explicit FASTENER reports the allocation - * failure instead. Keep this dataset alive while using `output_index`. - * A host-backed output index must be updated with - * `cuvsCagraUpdateDataset` before device search. + * @param[in] filter Filter, already applied by the caller while building `merged_dataset` + * @param[in] merged_dataset Caller-owned dataset handle already containing the concatenated (and, + * if `filter` is set, already-filtered) dataset. Keep this dataset alive + * while using `output_index`. A host-backed dataset must be updated with + * `cuvsCagraUpdateDataset` before device search. + * @param[in] offsets Per-index starting row within `merged_dataset`, as returned by + * `cuvsCagraMergedDatasetOffsets`. Array of `num_indices + 1` int64_t. * @param[out] output_index Output handle initialized with `cuvsCagraIndexCreate` */ CUVS_EXPORT cuvsError_t cuvsCagraMergeWithParams(cuvsResources_t res, @@ -1090,6 +1126,7 @@ CUVS_EXPORT cuvsError_t cuvsCagraMergeWithParams(cuvsResources_t res, size_t num_indices, cuvsFilter filter, cuvsDataset_t merged_dataset, + const int64_t* offsets, cuvsCagraIndex_t output_index); /** diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 82dba5424d..a135f1df13 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -116,6 +116,28 @@ static void with_index_by_layout(sg_cagra_c_api_index_box* box, template static void destroy_typed_addr(void* ptr); +template +static void with_dataset_view(cuvsDataset_t dataset, Fn&& fn); + +/** Build a `cuvsFilter` into the matching `cuvs::neighbors::filtering` object and invoke `fn` with + * it. `merged_row_count` is only used to size the bitset view for `BITSET`. */ +template +static void with_row_filter(cuvsFilter filter, int64_t merged_row_count, Fn&& fn) +{ + if (filter.type == NO_FILTER) { + fn(cuvs::neighbors::filtering::none_sample_filter{}); + } else if (filter.type == BITSET) { + using filter_mdspan_type = raft::device_vector_view; + auto removed_indices_tensor = reinterpret_cast(filter.addr); + auto removed_indices = cuvs::core::from_dlpack(removed_indices_tensor); + cuvs::core::bitset_view removed_indices_bitset( + removed_indices, merged_row_count); + fn(cuvs::neighbors::filtering::bitset_filter(removed_indices_bitset)); + } else { + RAFT_FAIL("Unsupported filter type: BITMAP"); + } +} + template static void merge_indices_for_layout( raft::resources* res_ptr, @@ -124,113 +146,38 @@ static void merge_indices_for_layout( cuvsFilter filter, cuvs::neighbors::cagra::merge_params const& merge_params, cuvsDataset_t merged_dataset, + std::vector const& offsets, cuvsCagraIndex_t output_index) { RAFT_EXPECTS(merged_dataset != nullptr, "cuvsCagraMerge: null merged dataset handle"); - RAFT_EXPECTS(merged_dataset->addr == 0, - "cuvsCagraMerge: merged dataset handle must be empty"); + RAFT_EXPECTS(merged_dataset->addr != 0, + "cuvsCagraMerge: merged_dataset must already contain the caller-concatenated " + "(and, if filtered, already-filtered) dataset"); - constexpr auto output_layout = + constexpr auto expected_layout = cuvs::neighbors::is_padded_dataset_view_v ? CUVS_DATASET_LAYOUT_PADDED : CUVS_DATASET_LAYOUT_STANDARD; + RAFT_EXPECTS(merged_dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE && + merged_dataset->layout == expected_layout, + "cuvsCagraMerge: merged_dataset must be a device dataset matching the input " + "indices' layout"); int64_t merged_row_count = 0; for (auto* idx_ptr : index_ptrs) { merged_row_count += static_cast(idx_ptr->size()); } - auto merge_into_dataset = [&](auto const& row_filter) { - auto const final_row_count = - cuvs::neighbors::cagra::detail::merged_dataset_size( - *res_ptr, index_ptrs, row_filter); - auto const dim = static_cast(index_ptrs.front()->dim()); - auto const stride = static_cast(index_ptrs.front()->dataset().stride()); - - try { - auto matrix = raft::make_device_matrix(*res_ptr, final_row_count, stride); - using owner_t = cuvs::neighbors::owning_dataset_for_view_t; - auto owner = std::make_unique(std::move(matrix), dim); - auto view = owner->as_dataset_view(); - auto merged_idx = - cuvs::neighbors::cagra::merge( - *res_ptr, params_cpp, index_ptrs, view, merge_params, row_filter); + using owner_t = cuvs::neighbors::owning_dataset_for_view_t; + with_dataset_view(merged_dataset, [&](auto const& view) { + with_row_filter(filter, merged_row_count, [&](auto const& row_filter) { + auto merged_idx = cuvs::neighbors::cagra::merge( + *res_ptr, params_cpp, index_ptrs, view, offsets, merge_params, row_filter); auto* holder = new cuvs_cagra_c_api_index_lifetime_holder{std::move(merged_idx)}; bind_index_lifetime_holder_to_C_index( output_index, output_index->dtype, holder); - - merged_dataset->addr = reinterpret_cast(owner.release()); - merged_dataset->destroy_addr = &destroy_typed_addr; - merged_dataset->dtype = output_index->dtype; - merged_dataset->mem_type = CUVS_DATASET_MEM_TYPE_DEVICE; - merged_dataset->layout = output_layout; - merged_dataset->is_owning = true; - return; - } catch (std::bad_alloc const& failure) { - if (merge_params.algo == cuvs::neighbors::cagra::merge_algo::FASTENER) { - RAFT_FAIL("FASTENER cagra::merge could not allocate device memory: %s", failure.what()); - } - // Filtered merge gathers rows with device-only primitives, matching the restriction on the - // legacy host fallback. - RAFT_EXPECTS(filter.type == NO_FILTER, - "Filtered merge isn't available with the host-memory OOM fallback"); - RAFT_LOG_DEBUG("cagra::merge: device allocation failed; using host memory for merged dataset"); - } - - using host_view_t = std::conditional_t< - cuvs::neighbors::is_padded_dataset_view_v, - cuvs::neighbors::host_padded_dataset_view, - cuvs::neighbors::host_standard_dataset_view>; - using host_owner_t = cuvs::neighbors::owning_dataset_for_view_t; - - auto matrix = raft::make_host_matrix(final_row_count, stride); - std::fill_n(matrix.data_handle(), static_cast(matrix.size()), T{}); - - std::size_t row_offset = 0; - auto stream = raft::resource::get_cuda_stream(*res_ptr); - for (auto* index : index_ptrs) { - auto const& input = index->dataset(); - raft::copy_matrix(matrix.data_handle() + row_offset * static_cast(stride), - static_cast(stride), - input.view().data_handle(), - static_cast(input.stride()), - static_cast(dim), - static_cast(input.n_rows()), - stream); - row_offset += static_cast(input.n_rows()); - } - raft::resource::sync_stream(*res_ptr); - - auto owner = std::make_unique(std::move(matrix), dim); - auto view = owner->as_dataset_view(); - auto merged_idx = cuvs::neighbors::cagra::build(*res_ptr, params_cpp, view); - auto* holder = - new cuvs_cagra_c_api_index_lifetime_holder{std::move(merged_idx)}; - bind_index_lifetime_holder_to_C_index( - output_index, output_index->dtype, holder); - - merged_dataset->addr = reinterpret_cast(owner.release()); - merged_dataset->destroy_addr = &destroy_typed_addr; - merged_dataset->dtype = output_index->dtype; - merged_dataset->mem_type = CUVS_DATASET_MEM_TYPE_HOST; - merged_dataset->layout = output_layout; - merged_dataset->is_owning = true; - }; - - if (filter.type == NO_FILTER) { - merge_into_dataset(cuvs::neighbors::filtering::none_sample_filter{}); - } else if (filter.type == BITSET) { - using filter_mdspan_type = raft::device_vector_view; - auto removed_indices_tensor = reinterpret_cast(filter.addr); - auto removed_indices = cuvs::core::from_dlpack(removed_indices_tensor); - cuvs::core::bitset_view removed_indices_bitset( - removed_indices, merged_row_count); - auto bitset_filter_obj = - cuvs::neighbors::filtering::bitset_filter(removed_indices_bitset); - merge_into_dataset(bitset_filter_obj); - } else { - RAFT_FAIL("Unsupported filter type: BITMAP"); - } + }); + }); } template @@ -1095,6 +1042,7 @@ void _merge(cuvsResources_t res, cuvsFilter filter, const cuvs::neighbors::cagra::merge_params& merge_params, cuvsDataset_t merged_dataset, + std::vector const& offsets, cuvsCagraIndex_t output_index) { auto res_ptr = reinterpret_cast(res); @@ -1141,13 +1089,58 @@ void _merge(cuvsResources_t res, convert_opaque_indices_to_concrete_types>( indices, num_indices); merge_indices_for_layout>( - res_ptr, params_cpp, index_ptrs, filter, merge_params, merged_dataset, output_index); + res_ptr, params_cpp, index_ptrs, filter, merge_params, merged_dataset, offsets, output_index); } else { auto index_ptrs = convert_opaque_indices_to_concrete_types>( indices, num_indices); merge_indices_for_layout>( - res_ptr, params_cpp, index_ptrs, filter, merge_params, merged_dataset, output_index); + res_ptr, params_cpp, index_ptrs, filter, merge_params, merged_dataset, offsets, output_index); + } +} + +template +void _merged_dataset_offsets(cuvsResources_t res, + cuvsCagraIndex_t* indices, + size_t num_indices, + cuvsFilter filter, + int64_t* offsets) +{ + auto res_ptr = reinterpret_cast(res); + auto* first_box = reinterpret_cast(indices[0]->addr); + RAFT_EXPECTS(first_box != nullptr, "cuvsCagraMergedDatasetOffsets: null index handle"); + auto layout = first_box->layout; + RAFT_EXPECTS(layout == sg_cagra_c_api_index_box::dataset_layout::device_padded || + layout == sg_cagra_c_api_index_box::dataset_layout::device_standard, + "cuvsCagraMergedDatasetOffsets: host indices are not supported; attach a device " + "dataset to each host index first."); + + int64_t merged_row_count = 0; + for (size_t i = 0; i < num_indices; ++i) { + auto* box = reinterpret_cast(indices[i]->addr); + RAFT_EXPECTS(box != nullptr, "cuvsCagraMergedDatasetOffsets: null index handle"); + RAFT_EXPECTS(box->layout == layout, + "cuvsCagraMergedDatasetOffsets: all input indices must share the same dataset " + "layout"); + with_index_by_layout( + box, + "cuvsCagraMergedDatasetOffsets: null index handle", + "cuvsCagraMergedDatasetOffsets: host indices are not supported", + [&](auto& idx) { merged_row_count += static_cast(idx.size()); }); + } + + auto compute = [&](auto const& index_ptrs) { + with_row_filter(filter, merged_row_count, [&](auto const& row_filter) { + auto result = cuvs::neighbors::cagra::merged_dataset_offsets(*res_ptr, index_ptrs, row_filter); + std::copy(result.begin(), result.end(), offsets); + }); + }; + if (layout == sg_cagra_c_api_index_box::dataset_layout::device_padded) { + compute(convert_opaque_indices_to_concrete_types>( + indices, num_indices)); + } else { + compute(convert_opaque_indices_to_concrete_types>( + indices, num_indices)); } } @@ -1848,10 +1841,11 @@ extern "C" cuvsError_t cuvsCagraMerge(cuvsResources_t res, size_t num_indices, cuvsFilter filter, cuvsDataset_t merged_dataset, + const int64_t* offsets, cuvsCagraIndex_t output_index) { return cuvsCagraMergeWithParams( - res, params, nullptr, indices, num_indices, filter, merged_dataset, output_index); + res, params, nullptr, indices, num_indices, filter, merged_dataset, offsets, output_index); } extern "C" cuvsError_t cuvsCagraMergeWithParams(cuvsResources_t res, @@ -1861,6 +1855,7 @@ extern "C" cuvsError_t cuvsCagraMergeWithParams(cuvsResources_t res, size_t num_indices, cuvsFilter filter, cuvsDataset_t merged_dataset, + const int64_t* offsets, cuvsCagraIndex_t output_index) { return cuvs::core::translate_exceptions([=] { @@ -1893,23 +1888,87 @@ extern "C" cuvsError_t cuvsCagraMergeWithParams(cuvsResources_t res, } RAFT_EXPECTS(output_index != nullptr, "Output index pointer must not be null"); RAFT_EXPECTS(merged_dataset != nullptr, "Merged dataset handle must not be null"); - RAFT_EXPECTS(merged_dataset->addr == 0, "Merged dataset handle must be empty"); + RAFT_EXPECTS(merged_dataset->addr != 0, + "merged_dataset must already contain the caller-concatenated dataset"); + RAFT_EXPECTS(offsets != nullptr, "offsets must not be null"); + auto offsets_vec = std::vector(offsets, offsets + num_indices + 1); output_index->dtype = dtype; // output index type matches inputs destroy_sg_cagra_c_api_box(output_index->addr); output_index->addr = 0; // Dispatch based on data type if (dtype.code == kDLFloat && dtype.bits == 32) { - _merge( - res, *params, indices, num_indices, filter, merge_params_cpp, merged_dataset, output_index); + _merge(res, + *params, + indices, + num_indices, + filter, + merge_params_cpp, + merged_dataset, + offsets_vec, + output_index); + } else if (dtype.code == kDLFloat && dtype.bits == 16) { + _merge(res, + *params, + indices, + num_indices, + filter, + merge_params_cpp, + merged_dataset, + offsets_vec, + output_index); + } else if (dtype.code == kDLInt && dtype.bits == 8) { + _merge(res, + *params, + indices, + num_indices, + filter, + merge_params_cpp, + merged_dataset, + offsets_vec, + output_index); + } else if (dtype.code == kDLUInt && dtype.bits == 8) { + _merge(res, + *params, + indices, + num_indices, + filter, + merge_params_cpp, + merged_dataset, + offsets_vec, + output_index); + } else { + RAFT_FAIL("Unsupported index data type: code=%d, bits=%d", dtype.code, dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsCagraMergedDatasetOffsets(cuvsResources_t res, + cuvsCagraIndex_t* indices, + size_t num_indices, + cuvsFilter filter, + int64_t* offsets) +{ + return cuvs::core::translate_exceptions([=] { + RAFT_EXPECTS(indices != nullptr && num_indices > 0, "indices array cannot be null or empty"); + RAFT_EXPECTS(indices[0] != nullptr && indices[0]->addr != 0, + "All input indices must be built (non-empty)"); + RAFT_EXPECTS(offsets != nullptr, "offsets must not be null"); + + auto dtype = (*indices[0]).dtype; + for (size_t i = 1; i < num_indices; ++i) { + RAFT_EXPECTS(indices[i] != nullptr && indices[i]->addr != 0, + "All input indices must be built (non-empty)"); + RAFT_EXPECTS((*indices[i]).dtype.code == dtype.code && (*indices[i]).dtype.bits == dtype.bits, + "All input indices must have the same data type"); + } + if (dtype.code == kDLFloat && dtype.bits == 32) { + _merged_dataset_offsets(res, indices, num_indices, filter, offsets); } else if (dtype.code == kDLFloat && dtype.bits == 16) { - _merge( - res, *params, indices, num_indices, filter, merge_params_cpp, merged_dataset, output_index); + _merged_dataset_offsets(res, indices, num_indices, filter, offsets); } else if (dtype.code == kDLInt && dtype.bits == 8) { - _merge( - res, *params, indices, num_indices, filter, merge_params_cpp, merged_dataset, output_index); + _merged_dataset_offsets(res, indices, num_indices, filter, offsets); } else if (dtype.code == kDLUInt && dtype.bits == 8) { - _merge( - res, *params, indices, num_indices, filter, merge_params_cpp, merged_dataset, output_index); + _merged_dataset_offsets(res, indices, num_indices, filter, offsets); } else { RAFT_FAIL("Unsupported index data type: code=%d, bits=%d", dtype.code, dtype.bits); } diff --git a/c/src/neighbors/cagra.hpp b/c/src/neighbors/cagra.hpp index 562f78f124..ab1cae9f05 100644 --- a/c/src/neighbors/cagra.hpp +++ b/c/src/neighbors/cagra.hpp @@ -18,12 +18,4 @@ void convert_c_search_params(cuvsCagraSearchParams params, /** Resolves `cuvsCagraIndex::addr` to `cagra::index*`; nullptr if the handle is empty. */ void* cagra_c_api_index_ptr(cuvsCagraIndex const* idx); - -namespace detail { -template -int64_t merged_dataset_size( - raft::resources const& res, - std::vector*> const& indices, - cuvs::neighbors::filtering::base_filter const& row_filter); -} // namespace detail } // namespace cuvs::neighbors::cagra diff --git a/c/tests/neighbors/ann_cagra_c.cu b/c/tests/neighbors/ann_cagra_c.cu index ba0e3ee310..d295608a37 100644 --- a/c/tests/neighbors/ann_cagra_c.cu +++ b/c/tests/neighbors/ann_cagra_c.cu @@ -291,12 +291,27 @@ TEST(CagraC, DatasetContractFailures) ASSERT_EQ(cuvsCagraBuild(res, build_params, host_standard_view, host_index_2), CUVS_SUCCESS); cuvsCagraIndex_t merge_out; ASSERT_EQ(cuvsCagraIndexCreate(&merge_out), CUVS_SUCCESS); + // Non-empty so the rejection below is actually the host-index-layout check, not a bounce off an + // empty merged_dataset handle. + rmm::device_uvector host_merge_dummy_d(16, stream); + DLManagedTensor host_merge_dummy_tensor = device_tensor; + host_merge_dummy_tensor.dl_tensor.data = host_merge_dummy_d.data(); + int64_t host_merge_dummy_shape[2] = {8, 2}; + host_merge_dummy_tensor.dl_tensor.shape = host_merge_dummy_shape; cuvsDataset_t merged_dataset; - ASSERT_EQ(cuvsDatasetCreate(&merged_dataset), CUVS_SUCCESS); + ASSERT_EQ(cuvsDatasetMakeStandardView(res, &host_merge_dummy_tensor, &merged_dataset), + CUVS_SUCCESS); cuvsCagraIndex_t host_indices[2] = {host_index, host_index_2}; - EXPECT_EQ( - cuvsCagraMerge(res, build_params, host_indices, 2, filter, merged_dataset, merge_out), - CUVS_ERROR); + int64_t host_merge_offsets[3] = {0, 4, 8}; + EXPECT_EQ(cuvsCagraMerge(res, + build_params, + host_indices, + 2, + filter, + merged_dataset, + host_merge_offsets, + merge_out), + CUVS_ERROR); ASSERT_EQ(cuvsCagraExtendParamsDestroy(extend_params), CUVS_SUCCESS); ASSERT_EQ(cuvsCagraSearchParamsDestroy(search_params), CUVS_SUCCESS); @@ -864,26 +879,9 @@ TEST(CagraC, BuildMergeSearch) filter.addr = 0; cuvsCagraIndex_t index_array[2] = {index_main, index_add}; - cuvsDataset_t merged_dataset; - ASSERT_EQ(cuvsDatasetCreate(&merged_dataset), CUVS_SUCCESS); - cuvsCagraMergeParams_t merge_params; - ASSERT_EQ(cuvsCagraMergeParamsCreate(&merge_params), CUVS_SUCCESS); - EXPECT_EQ(merge_params->algo, CUVS_CAGRA_MERGE_AUTO); - merge_params->algo = CUVS_CAGRA_MERGE_REBUILD; - ASSERT_EQ(cuvsCagraMergeWithParams( - res, build_params, merge_params, index_array, 2, filter, merged_dataset, index_merged), - CUVS_SUCCESS); - { - cuvsDatasetMemType_t mem_type{}; - cuvsDatasetLayout_t layout{}; - ASSERT_EQ(cuvsDatasetGetMemType(merged_dataset, &mem_type), CUVS_SUCCESS); - ASSERT_EQ(cuvsDatasetGetLayout(merged_dataset, &layout), CUVS_SUCCESS); - EXPECT_EQ(layout, CUVS_DATASET_LAYOUT_STANDARD); - EXPECT_EQ(mem_type, CUVS_DATASET_MEM_TYPE_DEVICE); - } - // Merge of standard-layout device inputs yields a standard index. Under the explicit C API - // contract, attach a padded dataset before calling search. + // The caller concatenates every input index's dataset itself (main || additional) and passes + // that pre-populated buffer plus per-index offsets -- merge() only merges the graph. rmm::device_uvector merged_d(14, stream); raft::copy(merged_d.data(), main_d.data(), main_d.size(), stream); raft::copy(merged_d.data() + main_d.size(), additional_d.data(), additional_d.size(), stream); @@ -900,6 +898,35 @@ TEST(CagraC, BuildMergeSearch) merged_dataset_tensor.dl_tensor.shape = merged_shape; merged_dataset_tensor.dl_tensor.strides = nullptr; + cuvsDataset_t merged_dataset; + ASSERT_EQ(cuvsDatasetMakeStandardView(res, &merged_dataset_tensor, &merged_dataset), + CUVS_SUCCESS); + int64_t merge_offsets[3] = {0, 4, 7}; + cuvsCagraMergeParams_t merge_params; + ASSERT_EQ(cuvsCagraMergeParamsCreate(&merge_params), CUVS_SUCCESS); + EXPECT_EQ(merge_params->algo, CUVS_CAGRA_MERGE_AUTO); + merge_params->algo = CUVS_CAGRA_MERGE_REBUILD; + ASSERT_EQ(cuvsCagraMergeWithParams(res, + build_params, + merge_params, + index_array, + 2, + filter, + merged_dataset, + merge_offsets, + index_merged), + CUVS_SUCCESS); + { + cuvsDatasetMemType_t mem_type{}; + cuvsDatasetLayout_t layout{}; + ASSERT_EQ(cuvsDatasetGetMemType(merged_dataset, &mem_type), CUVS_SUCCESS); + ASSERT_EQ(cuvsDatasetGetLayout(merged_dataset, &layout), CUVS_SUCCESS); + EXPECT_EQ(layout, CUVS_DATASET_LAYOUT_STANDARD); + EXPECT_EQ(mem_type, CUVS_DATASET_MEM_TYPE_DEVICE); + } + + // Merge of standard-layout device inputs yields a standard index. Under the explicit C API + // contract, attach a padded dataset before calling search. cuvsDataset_t padded_dataset_owner; ASSERT_EQ(cuvsDatasetMakePadded( res, &merged_dataset_tensor, CUVS_DATASET_MEM_TYPE_DEVICE, &padded_dataset_owner), diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h index ed067fac39..01109283ac 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h @@ -366,20 +366,28 @@ void cuvs_cagra::build(const T* dataset, size_t nrow) } cuvs::neighbors::filtering::none_sample_filter merge_row_filter; - int64_t merged_rows = 0; + std::vector offsets; + offsets.reserve(indices.size() + 1); + offsets.push_back(0); for (auto* index : indices) { - merged_rows += static_cast(index->size()); + offsets.push_back(offsets.back() + static_cast(index->size())); } - 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_)); + + // Each sub_index was built from a contiguous row-range slice of the single, original + // `dataset` buffer (see the split loop above), so the concatenation of the splits in + // `indices` order is exactly the original, unsplit `dataset`. Reuse the same padded-view + // helper the no-split branch above uses (branching on `dataset_is_on_host` the same way) to + // build `merged_dataset_view` directly from `dataset`, instead of letting `merge()` + // allocate an empty buffer and copy into it internally. + auto merged_dataset_view = + dataset_is_on_host ? detail::make_padded_view(handle_, dataset_view_host, *dataset_) + : detail::make_padded_view(handle_, dataset_view_device, *dataset_); index_ = std::make_shared(cuvs::neighbors::cagra::merge(handle_, params, indices, merged_dataset_view, + offsets, index_params_.merge_params, merge_row_filter)); // The merged index holds all the rows now; drop the splits rather than keep a second copy diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 84eaa7ac69..f5949c930c 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -3468,11 +3468,42 @@ struct merge_params { uint32_t leaf_degree = 4; }; +/** @brief Compute per-index write offsets for a merged dataset buffer. + * + * `merge()` requires the caller to have already concatenated every input index's dataset (in + * `indices` order, applying `row_filter` if any) into a single buffer and to know each index's + * starting row within it. For `row_filter = none_sample_filter{}`, those offsets are just the + * cumulative sizes of `indices` and the caller does not need this function. For a bitset + * `row_filter`, the number of surviving rows per index cannot be derived from any other public + * API, so use this to compute them before allocating and populating the merged buffer. + * + * @param[in] res RAFT resources. + * @param[in] indices CAGRA indices that will be passed to `merge()`. + * @param[in] row_filter Row filter that will be passed to `merge()`. Only `none_sample_filter` + * and `bitset_filter` are supported. + * @return A vector of `indices.size() + 1` offsets: entry `i` is the row at which + * `indices[i]`'s surviving rows must start in the merged buffer, and the last entry is the total + * row count of the merged buffer. + */ +template +auto merged_dataset_offsets( + raft::resources const& res, + std::vector*> const& indices, + const cuvs::neighbors::filtering::base_filter& row_filter) -> std::vector; + /** @brief Merge multiple physical CAGRA indices into one. + * + * The caller is responsible for concatenating every input index's dataset (applying `row_filter` + * if any) into a single `merged_dataset` buffer before calling this, and for computing `offsets` + * (see `merged_dataset_offsets()`). `merge()` only builds/merges the graph and rebinds the + * returned index to `merged_dataset` -- it never allocates or copies dataset rows itself. This + * mirrors the `extend()` contract. * * The overload without `merge_params` uses `merge_algo::AUTO`. AUTO runs Fastener only after a * non-mutating preflight validates every input and option; otherwise it calls the existing rebuild - * implementation. REBUILD always preserves every input dataset. + * implementation. REBUILD always preserves every input dataset. Fastener never applies a row + * filter -- `offsets` must equal the cumulative unfiltered sizes of `indices` for Fastener to be + * eligible; a filtered merge always uses rebuild. * * Fastener supports unfiltered, uncompressed `float`, `half`, `int8_t`, and `uint8_t` * indices using L2Expanded and `uint32_t` graph IDs. Fastener uses `root_fanout` at the first @@ -3481,22 +3512,38 @@ struct merge_params { * through 8 are supported. The configured spill width times `leaf_degree` must not exceed 255. * `index_params::graph_degree` is the final output degree. * - * Fastener copies the input datasets into `merged_dataset` and never mutates the input indices. - * The caller owns `merged_dataset` and must keep it alive for the lifetime of the returned index, - * which holds only a view of it. + * `merge()` never mutates the input indices. The caller owns `merged_dataset` and must keep it + * alive for the lifetime of the returned index, which holds only a view of it. * * @note This API only supports physical merge (`merge_strategy = MERGE_STRATEGY_PHYSICAL`). * All input indices must use the same `DatasetViewT` (dense padded or standard device views), * and must share one row stride. * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * // Compute this index's slice, e.g. via merged_dataset_offsets() for a bitset row_filter, or + * // cumulative index sizes for an unfiltered merge. + * std::vector offsets = {0, index0.size(), index0.size() + index1.size()}; + * // Build `merged` = concatenated(index0, index1) on device, padded for CAGRA. + * auto merged = make_device_padded_dataset(res, concatenated_view); + * auto merged_view = merged->as_dataset_view(); + * + * std::vector*> indices{&index0, &index1}; + * auto merged_index = cagra::merge(res, index_params, indices, merged_view, offsets); + * @endcode + * * @param[in] res RAFT resources used for the merge. * @param[in] params Parameters for the returned CAGRA index. * @param[in] indices CAGRA indices to merge. - * @param[in] merged_dataset Caller-owned storage for the consolidated dataset. Must have one row - * per retained vector after applying `row_filter`, and the same dimension and stride as the input - * datasets. - * @param[in] row_filter Optional row filter. Any filter selects rebuild in AUTO and is rejected by - * explicit FASTENER. + * @param[in] merged_dataset Caller-owned storage already containing the concatenated (and, if + * `row_filter` is set, already-filtered) dataset. Must have one row per retained vector, and the + * same dimension and stride as the input datasets. + * @param[in] offsets Per-index starting row within `merged_dataset`, as returned by + * `merged_dataset_offsets()`. Length `indices.size() + 1`; the last entry must equal + * `merged_dataset`'s row count. + * @param[in] row_filter Optional row filter, applied by the caller before this call. Any filter + * selects rebuild in AUTO and is rejected by explicit FASTENER. * @return The merged physical CAGRA index. */ template @@ -3504,6 +3551,7 @@ auto merge(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, std::vector*>& indices, DatasetViewT merged_dataset, + std::vector const& offsets, const cuvs::neighbors::filtering::base_filter& row_filter = cuvs::neighbors::filtering::none_sample_filter{}) -> cuvs::neighbors::cagra::index; @@ -3516,6 +3564,7 @@ auto merge(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, std::vector*>& indices, DatasetViewT merged_dataset, + std::vector const& offsets, const cuvs::neighbors::cagra::merge_params& merge_params, const cuvs::neighbors::filtering::base_filter& row_filter = cuvs::neighbors::filtering::none_sample_filter{}) diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 7af1251ba0..497a37097b 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -482,10 +482,11 @@ cuvs::neighbors::cagra::index merge( const cagra::index_params& params, std::vector*>& indices, DatasetViewT merged_dataset, + std::vector const& offsets, const cuvs::neighbors::filtering::base_filter& row_filter) { return cagra::detail::merge( - handle, params, indices, merged_dataset, row_filter); + handle, params, indices, merged_dataset, offsets, row_filter); } template @@ -494,11 +495,25 @@ cuvs::neighbors::cagra::index merge( const cagra::index_params& params, std::vector*>& indices, DatasetViewT merged_dataset, + std::vector const& offsets, const cagra::merge_params& merge_params, const cuvs::neighbors::filtering::base_filter& row_filter) { return cagra::detail::merge( - handle, params, indices, merged_dataset, merge_params, row_filter); + handle, params, indices, merged_dataset, offsets, merge_params, row_filter); +} + +/** @brief Compute per-index write offsets for a merged dataset buffer, needed to lay out a + * caller-populated `merged_dataset` for `merge()`. Only required for a bitset `row_filter`: for an + * unfiltered merge, offsets are just the cumulative sizes of `indices` and can be computed + * directly without calling this. See `merge()` for the full contract. */ +template +std::vector merged_dataset_offsets( + raft::resources const& handle, + std::vector*> const& indices, + const cuvs::neighbors::filtering::base_filter& row_filter) +{ + return cagra::detail::merged_dataset_offsets(handle, indices, row_filter); } template @@ -594,8 +609,8 @@ auto update_dataset(raft::resources const& res, } // namespace cuvs::neighbors::cagra #define CUVS_INST_CAGRA_MERGE(T, IdxT, DatasetViewT) \ - template CUVS_EXPORT int64_t \ - cuvs::neighbors::cagra::detail::merged_dataset_size( \ + template CUVS_EXPORT std::vector \ + cuvs::neighbors::cagra::merged_dataset_offsets( \ raft::resources const& handle, \ std::vector*> const& indices, \ cuvs::neighbors::filtering::base_filter const& row_filter); \ @@ -605,6 +620,7 @@ auto update_dataset(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, \ std::vector*>& indices, \ DatasetViewT merged_dataset, \ + std::vector const& offsets, \ cuvs::neighbors::filtering::base_filter const& row_filter); \ template CUVS_EXPORT cuvs::neighbors::cagra::index \ cuvs::neighbors::cagra::merge( \ @@ -612,5 +628,6 @@ auto update_dataset(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, \ std::vector*>& indices, \ DatasetViewT merged_dataset, \ + std::vector const& offsets, \ const cuvs::neighbors::cagra::merge_params& merge_params, \ cuvs::neighbors::filtering::base_filter const& row_filter); diff --git a/cpp/src/neighbors/cagra_merge_inst.cu.in b/cpp/src/neighbors/cagra_merge_inst.cu.in index da95e279b6..d20fc94f0c 100644 --- a/cpp/src/neighbors/cagra_merge_inst.cu.in +++ b/cpp/src/neighbors/cagra_merge_inst.cu.in @@ -38,18 +38,21 @@ CUVS_EXPORT void merge_instantiation_keepalive() using padded_index_t = cuvs::neighbors::cagra::index; using standard_index_t = cuvs::neighbors::cagra::index; - using filter_t = cuvs::neighbors::filtering::base_filter const&; + using filter_t = cuvs::neighbors::filtering::base_filter const&; + using offsets_t = std::vector const&; (void)static_cast&, inst_device_padded_view_t, + offsets_t, filter_t)>( &cuvs::neighbors::cagra::merge); (void)static_cast&, inst_device_padded_view_t, + offsets_t, cuvs::neighbors::cagra::merge_params const&, filter_t)>( &cuvs::neighbors::cagra::merge); @@ -57,15 +60,23 @@ CUVS_EXPORT void merge_instantiation_keepalive() cuvs::neighbors::cagra::index_params const&, std::vector&, inst_device_standard_view_t, + offsets_t, filter_t)>( &cuvs::neighbors::cagra::merge); (void)static_cast&, inst_device_standard_view_t, + offsets_t, cuvs::neighbors::cagra::merge_params const&, filter_t)>( &cuvs::neighbors::cagra::merge); + (void)static_cast (*)( + raft::resources const&, std::vector const&, filter_t)>( + &cuvs::neighbors::cagra::merged_dataset_offsets); + (void)static_cast (*)( + raft::resources const&, std::vector const&, filter_t)>( + &cuvs::neighbors::cagra::merged_dataset_offsets); } template CUVS_EXPORT void merge_instantiation_keepalive(); diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index 5069016b76..78714923a1 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -29,6 +29,7 @@ #include +#include #include #include #include @@ -38,41 +39,97 @@ namespace cuvs::neighbors::cagra::detail { +/** Per-index write offsets into a merged dataset buffer, length `indices.size() + 1`, with the + * last entry equal to the final row count. Entry `i` is the row at which caller-concatenated data + * for `indices[i]` must start. For `row_filter = none_sample_filter`, offsets are just the + * cumulative sizes of `indices` -- callers can compute those directly and do not need this + * function. For a bitset `row_filter`, the per-index surviving row counts are not derivable from + * public APIs alone, so this function walks the filter's sorted surviving-row list (via + * `bitset_view::to_csr`) and locates each index's boundary in it. */ template -int64_t merged_dataset_size( +std::vector merged_dataset_offsets( raft::resources const& handle, std::vector*> const& indices, cuvs::neighbors::filtering::base_filter const& row_filter) { - int64_t merged_rows = 0; + std::vector unfiltered_offsets; + unfiltered_offsets.reserve(indices.size() + 1); + unfiltered_offsets.push_back(0); for (auto* index : indices) { RAFT_EXPECTS(index != nullptr, "Null pointer detected in 'indices'. Ensure all elements are valid before usage."); - merged_rows += static_cast(index->size()); + unfiltered_offsets.push_back(unfiltered_offsets.back() + static_cast(index->size())); } - if (row_filter.get_filter_type() == cuvs::neighbors::filtering::FilterType::Bitset) { - auto const& actual_filter = - dynamic_cast&>(row_filter); - return actual_filter.view().count(handle); + + if (row_filter.get_filter_type() == cuvs::neighbors::filtering::FilterType::None) { + return unfiltered_offsets; + } + RAFT_EXPECTS(row_filter.get_filter_type() == cuvs::neighbors::filtering::FilterType::Bitset, + "Only none and bitset filters are supported by cagra::merged_dataset_offsets"); + + auto const& actual_filter = + dynamic_cast&>(row_filter); + int64_t const final_rows = actual_filter.view().count(handle); + + auto surviving_rows = raft::make_device_csr_matrix( + handle, 1, static_cast(unfiltered_offsets.back())); + surviving_rows.initialize_sparsity(final_rows); + actual_filter.view().to_csr(handle, surviving_rows); + auto const csr_indices = surviving_rows.structure_view().get_indices(); + + std::vector surviving_rows_host(csr_indices.size()); + raft::copy(surviving_rows_host.data(), + csr_indices.data(), + csr_indices.size(), + raft::resource::get_cuda_stream(handle)); + raft::resource::sync_stream(handle); + + std::vector filtered_offsets; + filtered_offsets.reserve(unfiltered_offsets.size()); + for (int64_t boundary : unfiltered_offsets) { + filtered_offsets.push_back(static_cast( + std::lower_bound(surviving_rows_host.begin(), surviving_rows_host.end(), boundary) - + surviving_rows_host.begin())); + } + return filtered_offsets; +} + +/** Validate that `offsets` is a well-formed length-`indices.size() + 1` boundary vector ending at + * `final_rows`: starts at 0, non-decreasing, and the caller-supplied row counts stay in range. */ +inline void validate_merge_offsets(std::vector const& offsets, + std::size_t num_indices, + int64_t final_rows) +{ + RAFT_EXPECTS(offsets.size() == num_indices + 1, + "offsets must have indices.size() + 1 (%zu) entries, got %zu", + num_indices + 1, + offsets.size()); + RAFT_EXPECTS(offsets.front() == 0, "offsets[0] must be 0"); + for (std::size_t i = 0; i + 1 < offsets.size(); ++i) { + RAFT_EXPECTS(offsets[i] <= offsets[i + 1], "offsets must be non-decreasing"); } - RAFT_EXPECTS(row_filter.get_filter_type() == cuvs::neighbors::filtering::FilterType::None, - "Only none and bitset filters are supported inside cagra::merge"); - return merged_rows; + RAFT_EXPECTS(offsets.back() == final_rows, + "offsets.back() (%ld) must equal merged_dataset's row count (%ld)", + long(offsets.back()), + long(final_rows)); } +/** Build a fresh CAGRA graph over a caller-populated, already-concatenated (and, if applicable, + * already-filtered) merged dataset. The caller owns `merged_dataset`; this only merges the graph + * and rebinds a view of it, mirroring the `extend()` contract. */ template cuvs::neighbors::cagra::index merge_rebuild( raft::resources const& handle, const cagra::index_params& params, std::vector*>& indices, DatasetViewT merged_dataset, + std::vector const& offsets, const cuvs::neighbors::filtering::base_filter& row_filter) { using cagra_index_t = cuvs::neighbors::cagra::index; - int64_t merged_rows = 0; - uint32_t dim = 0; - int64_t stride = -1; + uint32_t dim = 0; + int64_t stride = -1; RAFT_EXPECTS(row_filter.get_filter_type() != cuvs::neighbors::filtering::FilterType::Bitmap, "Bitmap filter isn't supported inside cagra::merge"); @@ -98,21 +155,12 @@ cuvs::neighbors::cagra::index merge_rebuild( RAFT_EXPECTS(stride == static_cast(dataset.stride()), "Row stride of datasets in indices must be equal."); } - merged_rows += static_cast(index->size()); } else { RAFT_FAIL("cagra::merge only supports an uncompressed dense device dataset index"); } } - bool const bitset_filtered = - row_filter.get_filter_type() == cuvs::neighbors::filtering::FilterType::Bitset; - int64_t const final_rows = - merged_dataset_size(handle, indices, row_filter); - - RAFT_EXPECTS(merged_dataset.n_rows() == final_rows, - "merged_dataset rows (%ld) must equal the final merged row count (%ld)", - long(merged_dataset.n_rows()), - long(final_rows)); + validate_merge_offsets(offsets, indices.size(), static_cast(merged_dataset.n_rows())); RAFT_EXPECTS(merged_dataset.dim() == dim, "merged_dataset dimension (%u) must equal the input dimension (%u)", unsigned(merged_dataset.dim()), @@ -122,84 +170,9 @@ cuvs::neighbors::cagra::index merge_rebuild( unsigned(merged_dataset.stride()), long(stride)); - auto output_const_view = merged_dataset.view(); - auto output_view = raft::make_device_matrix_view( - const_cast(output_const_view.data_handle()), final_rows, stride); - - auto merge_dataset = [&](T* dst, std::size_t dst_ld) { - IdxT row_offset = 0; - for (cagra_index_t* index : indices) { - const T* src_ptr = nullptr; - std::size_t n_rows = 0; - auto const& v = index->dataset(); - if constexpr (cuvs::neighbors::is_dense_row_major_dataset_view_v>) { - src_ptr = v.view().data_handle(); - n_rows = static_cast(v.n_rows()); - } else { - RAFT_FAIL("cagra::merge: unexpected dataset type while copying rows"); - } - raft::copy_matrix(dst + static_cast(row_offset) * dst_ld, - dst_ld, - src_ptr, - static_cast(stride), - static_cast(dim), - n_rows, - raft::resource::get_cuda_stream(handle)); - - row_offset += IdxT(index->dataset().n_rows()); - } - }; - - cudaStream_t stream = raft::resource::get_cuda_stream(handle); - - if (bitset_filtered) { - auto staging = raft::make_device_mdarray( - handle, - raft::resource::get_large_workspace_resource_ref(handle), - raft::make_extents(merged_rows, stride)); - RAFT_CUDA_TRY(cudaMemsetAsync( - staging.data_handle(), 0, static_cast(staging.size()) * sizeof(T), stream)); - merge_dataset(staging.data_handle(), static_cast(stride)); - - auto actual_filter = - dynamic_cast&>(row_filter); - - auto indices_csr = raft::make_device_csr_matrix( - handle, 1, static_cast(merged_rows)); - indices_csr.initialize_sparsity(final_rows); - - actual_filter.view().to_csr(handle, indices_csr); - - auto csr_indices = indices_csr.structure_view().get_indices(); - auto indices_view = raft::make_device_vector_view( - csr_indices.data(), static_cast(csr_indices.size())); - - RAFT_CUDA_TRY(cudaMemsetAsync( - output_view.data_handle(), - 0, - static_cast(final_rows) * static_cast(stride) * sizeof(T), - stream)); - - raft::matrix::copy_rows( - handle, raft::make_const_mdspan(staging.view()), output_view, indices_view); - - auto index = ::cuvs::neighbors::cagra::detail::build_from_device_matrix( - handle, params, merged_dataset); - index = ::cuvs::neighbors::cagra::update_dataset(handle, std::move(index), merged_dataset); - RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); - return index; - } - - RAFT_CUDA_TRY(cudaMemsetAsync( - output_view.data_handle(), - 0, - static_cast(final_rows) * static_cast(stride) * sizeof(T), - stream)); - merge_dataset(output_view.data_handle(), static_cast(stride)); auto index = ::cuvs::neighbors::cagra::detail::build_from_device_matrix( handle, params, merged_dataset); index = ::cuvs::neighbors::cagra::update_dataset(handle, std::move(index), merged_dataset); - RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); return index; } @@ -219,6 +192,7 @@ auto preflight_fastener( cagra::index_params const& params, cagra::merge_params const& merge_params, std::vector*> const& indices, + std::vector const& offsets, cuvs::neighbors::filtering::base_filter const& row_filter) -> fastener_preflight_result { fastener_preflight_result result; @@ -326,6 +300,14 @@ auto preflight_fastener( result.offsets.push_back(static_cast(rows)); } + // Fastener never applies a row filter (rejected above), so the caller-supplied offsets must be + // exactly the unfiltered per-index cumulative sizes computed above -- merge_dataset_offsets() + // returns this same vector for an unfiltered merge, so a caller who used it will always match. + if (offsets != result.offsets) { + return reject( + "offsets must equal the cumulative unfiltered row counts of each input index for Fastener"); + } + if (result.dim <= 0 || result.dim > std::numeric_limits::max()) { return reject("dataset dimension must be positive and fit cuBLAS int dimensions"); } @@ -366,30 +348,10 @@ auto preflight_fastener( return result; } -/** Copy every input dataset into its row range of the caller-supplied merged dataset. Both sides - * carry a row pitch: the inputs share one stride (enforced by preflight) and the destination uses - * the merged dataset's own stride. */ -template -void copy_input_datasets( - raft::resources const& handle, - std::vector*> const& indices, - std::vector const& offsets, - int64_t dim, - int64_t destination_stride, - T* destination) -{ - for (std::size_t i = 0; i < indices.size(); ++i) { - auto const& source = indices[i]->dataset(); - raft::copy_matrix(destination + offsets[i] * destination_stride, - static_cast(destination_stride), - source.view().data_handle(), - static_cast(source.stride()), - static_cast(dim), - static_cast(source.n_rows()), - raft::resource::get_cuda_stream(handle)); - } -} - +/** Build a merged CAGRA graph via Fastener over a caller-populated, already-concatenated merged + * dataset (Fastener never applies a row filter, so `merged_dataset` always holds the full, + * unfiltered concatenation of every input in `indices` order). The caller owns `merged_dataset`; + * this only merges the graph and rebinds a view of it. */ template auto merge_fastener(raft::resources const& handle, cagra::index_params const& params, @@ -410,22 +372,6 @@ auto merge_fastener(raft::resources const& handle, long(preflight.dim)); auto const output_const_view = merged_dataset.view(); - auto* destination = const_cast(output_const_view.data_handle()); - { - raft::common::nvtx::range scope("cagra::merge/consolidate"); - // The copy below overwrites columns [0, dim), while the sorter computes L2 over [0, stride). - // Zero the remaining padded columns; when stride == dim, there are none to initialize. - if (stride > preflight.dim) { - RAFT_CUDA_TRY(cudaMemset2DAsync(destination + preflight.dim, - static_cast(stride) * sizeof(T), - 0, - static_cast(stride - preflight.dim) * sizeof(T), - static_cast(preflight.rows), - raft::resource::get_cuda_stream(handle))); - } - copy_input_datasets( - handle, indices, preflight.offsets, preflight.dim, stride, destination); - } // The scaffold and the sorter read the consolidated rows with this pitch; dim stays logical. auto dataset_view = raft::make_device_matrix_view( output_const_view.data_handle(), preflight.rows, stride); @@ -490,6 +436,7 @@ auto merge(raft::resources const& handle, cagra::index_params const& params, std::vector*>& indices, DatasetViewT merged_dataset, + std::vector const& offsets, cagra::merge_params const& merge_params, cuvs::neighbors::filtering::base_filter const& row_filter) -> cuvs::neighbors::cagra::index @@ -503,15 +450,15 @@ auto merge(raft::resources const& handle, "Unknown cagra::merge algorithm"); if (merge_params.algo == cagra::merge_algo::REBUILD) { return merge_rebuild( - handle, params, indices, merged_dataset, row_filter); + handle, params, indices, merged_dataset, offsets, row_filter); } - auto preflight = - preflight_fastener(handle, params, merge_params, indices, row_filter); + auto preflight = preflight_fastener( + handle, params, merge_params, indices, offsets, row_filter); if (!preflight.eligible) { if (merge_params.algo == cagra::merge_algo::AUTO) { return merge_rebuild( - handle, params, indices, merged_dataset, row_filter); + handle, params, indices, merged_dataset, offsets, row_filter); } RAFT_FAIL("FASTENER cagra::merge is unsupported: %s", preflight.reason.c_str()); } @@ -529,7 +476,7 @@ auto merge(raft::resources const& handle, RAFT_LOG_WARN("Fastener cagra::merge could not allocate (%s); falling back to rebuild", failure.what()); return merge_rebuild( - handle, params, indices, merged_dataset, row_filter); + handle, params, indices, merged_dataset, offsets, row_filter); } } @@ -543,13 +490,14 @@ auto merge(raft::resources const& handle, cagra::index_params const& params, std::vector*>& indices, DatasetViewT merged_dataset, + std::vector const& offsets, cuvs::neighbors::filtering::base_filter const& row_filter) -> cuvs::neighbors::cagra::index { // Fully qualified: an unqualified call also finds cuvs::neighbors::cagra::merge via ADL on the // index arguments, which is ambiguous with this overload. return cuvs::neighbors::cagra::detail::merge( - handle, params, indices, merged_dataset, cagra::merge_params{}, row_filter); + handle, params, indices, merged_dataset, offsets, cagra::merge_params{}, row_filter); } } // namespace cuvs::neighbors::cagra::detail diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index ebfbab759e..d1d4344d29 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -17,6 +17,8 @@ #include #include #include +#include +#include #include #include #include @@ -29,6 +31,7 @@ #include #include #include +#include #include #include @@ -1326,14 +1329,37 @@ class AnnCagraIndexFilteredMergeTest : public ::testing::TestWithParam( + auto offsets = + cuvs::neighbors::cagra::merged_dataset_offsets(handle_, indices, bitset_filter_obj); + int64_t const merged_rows = offsets.back(); + + // `merge()` no longer gathers the filtered rows itself: do what it used to do internally + // -- find the sorted list of rows the bitset keeps (via `to_csr`, same as + // `merged_dataset_offsets()`) and gather them from `database`. `index0`/`index1` were built + // from contiguous halves of `database` (see `database0_view`/`database1_view` above), so + // `database` in its original row order is exactly the unfiltered concatenation of + // `index0`'s and `index1`'s rows, in `indices` order. + auto surviving_rows_csr = raft::make_device_csr_matrix( + handle_, 1, static_cast(ps.n_rows)); + surviving_rows_csr.initialize_sparsity(merged_rows); + bitset_filter_obj.view().to_csr(handle_, surviving_rows_csr); + auto surviving_row_indices = surviving_rows_csr.structure_view().get_indices(); + + auto database_view = raft::make_device_matrix_view( + (const DataT*)database.data(), ps.n_rows, ps.dim); + auto gathered_matrix = + raft::make_device_matrix(handle_, merged_rows, ps.dim); + raft::matrix::copy_rows( handle_, - ps.n_rows - static_cast(test_cagra_sample_filter::offset), - static_cast(index0.dataset().stride())); - auto merged_dataset = cuvs::neighbors::device_padded_dataset( - std::move(merged_matrix), static_cast(ps.dim)); + database_view, + gathered_matrix.view(), + raft::make_device_vector_view(surviving_row_indices.data(), + merged_rows)); + + cuvs::neighbors::test::padded_device_matrix_for_cagra merged_dataset( + handle_, raft::make_const_mdspan(gathered_matrix.view())); auto merge_idx = cuvs::neighbors::cagra::merge( - handle_, index_params, indices, merged_dataset.as_dataset_view(), bitset_filter_obj); + handle_, index_params, indices, merged_dataset.view, offsets, bitset_filter_obj); auto search_queries_view = raft::make_device_matrix_view( search_queries.data(), ps.n_queries, ps.dim); @@ -1572,21 +1598,26 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { if (ps.merge_strategy == cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL) { // The merged index holds only a view, so merged_dataset must outlive it. - auto const merged_rows = - static_cast(index0.size()) + static_cast(index1.size()); - auto merged_matrix = raft::make_device_matrix( - handle_, merged_rows, static_cast(index0.dataset().stride())); - auto merged_dataset = cuvs::neighbors::device_padded_dataset( - std::move(merged_matrix), static_cast(ps.dim)); + std::vector offsets{ + 0, + static_cast(index0.size()), + static_cast(index0.size()) + static_cast(index1.size())}; + // index0/index1 were built from contiguous halves of `database` (database0_view / + // database1_view above), so `database` in its original row order is exactly the + // concatenation of index0's and index1's rows, in indices_to_merge order. + auto database_view = raft::make_device_matrix_view( + (const DataT*)database.data(), ps.n_rows, ps.dim); + cuvs::neighbors::test::padded_device_matrix_for_cagra merged_dataset( + handle_, database_view); auto merged_idx = ps.physical_merge_params.has_value() ? cagra::merge(handle_, index_params, indices_to_merge, - merged_dataset.as_dataset_view(), + merged_dataset.view, + offsets, *ps.physical_merge_params) - : cagra::merge( - handle_, index_params, indices_to_merge, merged_dataset.as_dataset_view()); + : cagra::merge(handle_, index_params, indices_to_merge, merged_dataset.view, offsets); cagra::search(handle_, search_params, merged_idx, diff --git a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu index 9dee454d8c..3df58958f8 100644 --- a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu +++ b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu @@ -95,6 +95,35 @@ auto make_merged_storage(raft::resources const& res, int64_t rows, int64_t dim) return padded_storage{std::move(matrix), view}; } +/** Fill `dest` (from `make_merged_storage`) with the concatenation of `indices`' dataset rows, in + * order -- what cagra::merge() used to do internally for an unfiltered merge. Returns the + * per-index offsets merge() now requires. */ +template +auto populate_merged_storage(raft::resources const& res, + std::vector*> const& indices, + padded_storage& dest) -> std::vector +{ + auto stream = raft::resource::get_cuda_stream(res); + int64_t const dest_stride = static_cast(dest.view.stride()); + std::vector offsets{0}; + for (auto* index : indices) { + auto src = index->dataset(); + int64_t const rows = src.n_rows(); + int64_t const dim = static_cast(src.dim()); + int64_t const stride = static_cast(src.stride()); + raft::copy_matrix(dest.matrix.data_handle() + offsets.back() * dest_stride, + static_cast(dest_stride), + src.view().data_handle(), + static_cast(stride), + static_cast(dim), + static_cast(rows), + stream); + offsets.push_back(offsets.back() + rows); + } + raft::resource::sync_stream(res); + return offsets; +} + template auto make_dataset(raft::resources const& res, int64_t rows, @@ -284,12 +313,12 @@ void run_explicit_fastener(cuvs::distance::DistanceType metric) fastener.leaf_degree = 4; std::vector*> indices{&index0, &index1}; - auto merged_storage = make_merged_storage(res, rows * 2, dim); - RAFT_CUDA_TRY(cudaMemsetAsync(merged_storage.matrix.data_handle(), - 0xff, - merged_storage.matrix.size() * sizeof(T), - raft::resource::get_cuda_stream(res))); - auto merged = merge(res, params, indices, merged_storage.view, fastener); + // merge() only rebuilds the graph now; the caller must populate the merged dataset buffer. + // `expected` is already the concatenation of index0's and index1's rows, so build the padded + // merged buffer directly from it. + auto merged_storage = make_padded(res, raft::make_const_mdspan(expected.view())); + std::vector offsets{0, rows, rows * 2}; + auto merged = merge(res, params, indices, merged_storage.view, offsets, fastener); expect_valid_graph(merged, rows * 2, degree); expect_dataset_order(res, merged, raft::make_const_mdspan(expected.view())); expect_zero_padding(res, merged); @@ -745,7 +774,8 @@ TEST(CagraMergeFastener, IdenticalInputsAlignedToLeafSizeStayConnected) std::vector*> indices{&index0, &index1}; auto merged_storage = make_merged_storage(res, rows, dim); - auto merged = merge(res, params, indices, merged_storage.view, fastener); + auto offsets = populate_merged_storage(res, indices, merged_storage); + auto merged = merge(res, params, indices, merged_storage.view, offsets, fastener); expect_valid_graph(merged, rows, degree); EXPECT_EQ(count_connected_components(merged), 1) << "the inputs were merged without any edge connecting them"; @@ -802,8 +832,9 @@ TEST(CagraMergeFastener, LeafGemmLimitsRejectOversizedWorkspaceDimensions) fastener.leaf_size = leaf_size; std::vector*> indices{&index0, &index1}; + std::vector offsets{0, rows, rows * 2}; auto result = detail::preflight_fastener( - res, params, fastener, indices, cuvs::neighbors::filtering::none_sample_filter{}); + res, params, fastener, indices, offsets, cuvs::neighbors::filtering::none_sample_filter{}); EXPECT_FALSE(result.eligible); EXPECT_EQ(result.reason, "dataset dimension exceeds the leaf GEMM workspace limit"); EXPECT_EQ(index0.dataset().n_rows(), rows); @@ -877,20 +908,27 @@ TEST(CagraMergeFastener, AssignmentGemmLimitsRejectDimensionsThatFitALeaf) fastener.leader_fraction = worst_case.leader_fraction; fastener.max_leaders = worst_case.max_leaders; fastener.leaf_size = worst_case.leaf_size; - auto result = detail::preflight_fastener( - res, params, fastener, indices, cuvs::neighbors::filtering::none_sample_filter{}); + std::vector preflight_offsets{0, rows_per_input, rows_per_input * 2}; + auto result = detail::preflight_fastener(res, + params, + fastener, + indices, + preflight_offsets, + cuvs::neighbors::filtering::none_sample_filter{}); EXPECT_FALSE(result.eligible); EXPECT_EQ(result.reason, "dataset dimension exceeds the assignment GEMM workspace limit"); // Explicit Fastener surfaces the rejection; AUTO takes the rebuild path and still produces an // index. Reaching rebuild at all is the regression this test guards. auto fastener_storage = make_merged_storage(res, merged_rows, dim); - EXPECT_ANY_THROW(merge(res, params, indices, fastener_storage.view, fastener)); + auto fastener_offsets = populate_merged_storage(res, indices, fastener_storage); + EXPECT_ANY_THROW(merge(res, params, indices, fastener_storage.view, fastener_offsets, fastener)); auto automatic = fastener; automatic.algo = merge_algo::AUTO; auto merged_storage = make_merged_storage(res, merged_rows, dim); - auto merged = merge(res, params, indices, merged_storage.view, automatic); + auto merged_offsets = populate_merged_storage(res, indices, merged_storage); + auto merged = merge(res, params, indices, merged_storage.view, merged_offsets, automatic); EXPECT_EQ(merged.size(), merged_rows); EXPECT_EQ(index0.dataset().n_rows(), rows_per_input); @@ -935,8 +973,9 @@ TEST(CagraMergeFastener, PreflightRejectsCandidateDegreeAboveSortLimit) merge_params fastener; fastener.algo = merge_algo::FASTENER; std::vector*> indices{&index0, &index1}; + std::vector offsets{0, rows, rows * 2}; return detail::preflight_fastener( - res, params, fastener, indices, cuvs::neighbors::filtering::none_sample_filter{}); + res, params, fastener, indices, offsets, cuvs::neighbors::filtering::none_sample_filter{}); }; // Exactly at the limit stays eligible; one degree past it is rejected for that specific reason. @@ -997,7 +1036,8 @@ TEST(CagraMergeFastener, BatchesWithinConfiguredWorkspace) std::vector*> indices{&index0, &index1}; auto merged_storage = make_merged_storage(res, rows * 2, dim); - auto merged = merge(res, params, indices, merged_storage.view, fastener); + auto offsets = populate_merged_storage(res, indices, merged_storage); + auto merged = merge(res, params, indices, merged_storage.view, offsets, fastener); expect_valid_graph(merged, rows * 2, degree); EXPECT_EQ(raft::resource::get_workspace_used_bytes(res), 0); } @@ -1053,8 +1093,11 @@ TEST(CagraMergeFastener, MixedDatasetOwnershipPreservesInputs) fastener.leaf_degree = 4; std::vector*> indices{&index0, &index1}; - auto merged_storage = make_merged_storage(res, rows * 2, dim); - auto merged = merge(res, params, indices, merged_storage.view, fastener); + // `expected` is already the concatenation of index0's and index1's rows; build the padded + // merged buffer directly from it instead of letting merge() copy it internally. + auto merged_storage = make_padded(res, raft::make_const_mdspan(expected.view())); + std::vector offsets{0, rows, rows * 2}; + auto merged = merge(res, params, indices, merged_storage.view, offsets, fastener); expect_valid_graph(merged, rows * 2, degree); expect_dataset_order(res, merged, raft::make_const_mdspan(expected.view())); EXPECT_EQ(merged.dataset().n_rows(), rows * 2); @@ -1164,9 +1207,17 @@ void run_fastener_merge_recall(size_t n_inputs, double min_recall) inputs.push_back(&part); } merge_params fastener; - fastener.algo = merge_algo::FASTENER; - auto merged_storage = make_merged_storage(res, rows, dim); - auto merged = merge(res, build_params, inputs, merged_storage.view, fastener); + fastener.algo = merge_algo::FASTENER; + + // Each part was built from a contiguous slice of `data`, so the concatenation of the parts in + // `inputs` order is exactly `data` itself; build the merged buffer directly from it. + cuvs::neighbors::test::padded_device_matrix_for_cagra merged_storage( + res, raft::make_const_mdspan(data.view())); + std::vector offsets{0}; + for (auto& part : parts) { + offsets.push_back(offsets.back() + static_cast(part.size())); + } + auto merged = merge(res, build_params, inputs, merged_storage.view, offsets, fastener); ASSERT_EQ(merged.size(), rows); auto found_indices = raft::make_device_matrix(res, n_query, k); @@ -1276,7 +1327,8 @@ TEST(CagraMergeFastener, MixedDegreesAndThreeLevelUint8OptionsProduceExactDegree std::vector*> indices{&index0, &index1}; auto merged_storage = make_merged_storage(res, rows * 2, dim); - auto merged = merge(res, params, indices, merged_storage.view, fastener); + auto offsets = populate_merged_storage(res, indices, merged_storage); + auto merged = merge(res, params, indices, merged_storage.view, offsets, fastener); expect_valid_graph(merged, rows * 2, 8); } @@ -1380,7 +1432,8 @@ TEST(CagraMergeFastener, MergeSupportsOutputDegreeBelowInputDegree) std::vector*> indices{&index0, &index1}; auto merged_storage = make_merged_storage(res, rows * 2, dim); - auto merged = merge(res, params, indices, merged_storage.view, fastener); + auto offsets = populate_merged_storage(res, indices, merged_storage); + auto merged = merge(res, params, indices, merged_storage.view, offsets, fastener); expect_valid_graph(merged, rows * 2, params.graph_degree); } @@ -1457,9 +1510,10 @@ TEST(CagraMergeFastener, InvalidManywayOptionsFailPreflightWithoutMutation) value.lower_fanout = 32; }); + std::vector offsets{0, rows, rows * 2}; for (auto const& candidate : invalid) { auto result = detail::preflight_fastener( - res, params, candidate, indices, cuvs::neighbors::filtering::none_sample_filter{}); + res, params, candidate, indices, offsets, cuvs::neighbors::filtering::none_sample_filter{}); EXPECT_FALSE(result.eligible) << result.reason; } EXPECT_EQ(owned.indices[0].dataset().n_rows(), rows); @@ -1490,10 +1544,11 @@ TEST(CagraMergeFastener, DispatchRejectsOrFallsBackBeforeMutation) auto throwaway_storage = make_merged_storage(res, rows * 2, dim); std::vector*> indices{&owned.indices[0], &owned.indices[1]}; + auto offsets = populate_merged_storage(res, indices, throwaway_storage); merge_params unsupported; unsupported.algo = merge_algo::FASTENER; unsupported.leaf_size = 512; - EXPECT_ANY_THROW(merge(res, params, indices, throwaway_storage.view, unsupported)); + EXPECT_ANY_THROW(merge(res, params, indices, throwaway_storage.view, offsets, unsupported)); EXPECT_EQ(owned.indices[0].dataset().n_rows(), rows); EXPECT_EQ(owned.indices[1].dataset().n_rows(), rows); } @@ -1505,10 +1560,12 @@ TEST(CagraMergeFastener, DispatchRejectsOrFallsBackBeforeMutation) auto throwaway_storage = make_merged_storage(res, rows * 2, dim); std::vector*> indices{&owned.indices[0], &owned.indices[1]}; + auto offsets = populate_merged_storage(res, indices, throwaway_storage); merge_params fastener; fastener.algo = merge_algo::FASTENER; fastener.leaf_size = 64; - EXPECT_ANY_THROW(merge(res, inner_product_params, indices, throwaway_storage.view, fastener)); + EXPECT_ANY_THROW( + merge(res, inner_product_params, indices, throwaway_storage.view, offsets, fastener)); EXPECT_EQ(owned.indices[0].dataset().n_rows(), rows); EXPECT_EQ(owned.indices[1].dataset().n_rows(), rows); } @@ -1517,10 +1574,11 @@ TEST(CagraMergeFastener, DispatchRejectsOrFallsBackBeforeMutation) auto throwaway_storage = make_merged_storage(res, rows * 2, dim); std::vector*> indices{&owned.indices[0], &owned.indices[1]}; + auto offsets = populate_merged_storage(res, indices, throwaway_storage); merge_params automatic; automatic.algo = merge_algo::AUTO; automatic.leaf_size = 512; - auto merged = merge(res, params, indices, throwaway_storage.view, automatic); + auto merged = merge(res, params, indices, throwaway_storage.view, offsets, automatic); EXPECT_EQ(merged.size(), rows * 2); EXPECT_EQ(owned.indices[0].dataset().n_rows(), rows); EXPECT_EQ(owned.indices[1].dataset().n_rows(), rows); @@ -1532,7 +1590,8 @@ TEST(CagraMergeFastener, DispatchRejectsOrFallsBackBeforeMutation) &owned.indices[1]}; merge_params rebuild{merge_algo::REBUILD}; auto merged_storage = make_merged_storage(res, rows * 2, dim); - auto merged = merge(res, params, indices, merged_storage.view, rebuild); + auto offsets = populate_merged_storage(res, indices, merged_storage); + auto merged = merge(res, params, indices, merged_storage.view, offsets, rebuild); EXPECT_EQ(merged.size(), rows * 2); EXPECT_EQ(owned.indices[0].dataset().n_rows(), rows); EXPECT_EQ(owned.indices[1].dataset().n_rows(), rows); diff --git a/go/cagra/cagra.go b/go/cagra/cagra.go index bb7103b111..2eac3f0439 100644 --- a/go/cagra/cagra.go +++ b/go/cagra/cagra.go @@ -128,6 +128,28 @@ func (view *PaddedDatasetView) datasetHandle() C.cuvsDataset_t { return view.view } +// PaddedDatasetCloser is a PaddedDatasetHandle that owns resources needing Close. +type PaddedDatasetCloser interface { + PaddedDatasetHandle + Close() error +} + +// MakePaddedDatasetAuto builds a PaddedDatasetCloser from a tensor, choosing the +// non-owning MakePaddedDatasetView when the tensor's row stride is already +// CAGRA-padded (MakePaddedDataset rejects already-aligned sources), and the +// owning MakePaddedDataset otherwise. Mirrors the branch BuildIndex uses. +func MakePaddedDatasetAuto[T any](Resources cuvs.Resource, dataset *cuvs.Tensor[T]) (PaddedDatasetCloser, error) { + if dataset == nil || dataset.C_tensor == nil { + return nil, errors.New("dataset is nil") + } + datasetTensor := (*C.DLManagedTensor)(unsafe.Pointer(dataset.C_tensor)) + var zero T + if isCagraPaddedTensor(datasetTensor, int(unsafe.Sizeof(zero))) { + return MakePaddedDatasetView(Resources, dataset) + } + return MakePaddedDataset(Resources, dataset) +} + // Destroys an owning padded dataset handle. func (dataset *PaddedDataset) Close() error { if dataset == nil || dataset.dataset == nil { @@ -295,6 +317,168 @@ func ExtendIndex(Resources cuvs.Resource, params *ExtendParams, extended_dataset return nil } +// buildAllowListFilter builds a cuvsFilter (and its backing bitset tensor, when +// allowList is non-nil) suitable for passing to a filtered C call. The caller +// must keep the returned tensor alive (and Close it) for the duration of the +// call, since the filter's addr points into it. +func buildAllowListFilter(Resources cuvs.Resource, allowList []uint32) (C.cuvsFilter, *cuvs.Tensor[uint32], error) { + if allowList == nil { + return C.cuvsFilter{ + _type: C.NO_FILTER, + addr: C.uintptr_t(0), + }, nil, nil + } + + bitset := createBitset(allowList) + allowListTensor, err := cuvs.NewVector[uint32](bitset) + if err != nil { + return C.cuvsFilter{}, nil, err + } + if _, err := allowListTensor.ToDevice(&Resources); err != nil { + allowListTensor.Close() + return C.cuvsFilter{}, nil, err + } + + filter := C.cuvsFilter{ + _type: C.BITSET, + addr: C.uintptr_t(uintptr(unsafe.Pointer(allowListTensor.C_tensor))), + } + return filter, &allowListTensor, nil +} + +// MergedDatasetOffsets computes, for each input index (in order), the row at +// which its (post-filter) rows must start within a caller-built merged dataset +// buffer, for use with MergeIndex/MergeIndexWithParams. The returned slice has +// len(indices)+1 entries; the last entry is the total merged row count. +// +// For an unfiltered merge (allowList == nil) this is just the cumulative sum +// of each index's row count, and calling this function is unnecessary. For a +// filtered merge, this must be called to determine how many rows survive the +// filter for each index. +func MergedDatasetOffsets(Resources cuvs.Resource, indices []*CagraIndex, allowList []uint32) ([]int64, error) { + if len(indices) == 0 { + return nil, errors.New("indices must not be empty") + } + + filter, filterTensor, err := buildAllowListFilter(Resources, allowList) + if err != nil { + return nil, err + } + if filterTensor != nil { + defer filterTensor.Close() + } + + cIndices := make([]C.cuvsCagraIndex_t, len(indices)) + for i, idx := range indices { + cIndices[i] = idx.index + } + + cOffsets := make([]C.int64_t, len(indices)+1) + + err = cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraMergedDatasetOffsets( + C.cuvsResources_t(Resources.Resource), + &cIndices[0], + C.size_t(len(indices)), + filter, + &cOffsets[0], + ))) + if err != nil { + return nil, err + } + + offsets := make([]int64, len(cOffsets)) + for i, v := range cOffsets { + offsets[i] = int64(v) + } + return offsets, nil +} + +// MergeIndex merges multiple CAGRA indices into output using AUTO merge +// parameters. +// +// # Arguments +// +// * `Resources` - Resources to use +// * `params` - Parameters for the output index +// * `indices` - Input indices to merge, in the order they were concatenated into mergedDataset +// * `mergedDataset` - Caller-owned padded dataset already containing the concatenation (and, +// if allowList is set, already-filtered rows) of every input index's dataset, in `indices` order +// * `offsets` - Per-index starting row within mergedDataset; len(indices)+1 entries, as returned +// by MergedDatasetOffsets (or the trivial cumulative sizes, if allowList is nil) +// * `allowList` - Row filter already applied by the caller while building mergedDataset, or nil +// * `index` - Output CagraIndex, must be created with CreateIndex before use +func MergeIndex(Resources cuvs.Resource, params *IndexParams, indices []*CagraIndex, mergedDataset PaddedDatasetHandle, offsets []int64, allowList []uint32, index *CagraIndex) error { + return mergeIndex(Resources, params, nil, indices, mergedDataset, offsets, allowList, index) +} + +// MergeIndexWithParams merges multiple CAGRA indices into output, using +// explicit mergeParams to control the merge algorithm. See MergeIndex for the +// remaining arguments. +func MergeIndexWithParams(Resources cuvs.Resource, params *IndexParams, mergeParams *MergeParams, indices []*CagraIndex, mergedDataset PaddedDatasetHandle, offsets []int64, allowList []uint32, index *CagraIndex) error { + return mergeIndex(Resources, params, mergeParams, indices, mergedDataset, offsets, allowList, index) +} + +func mergeIndex(Resources cuvs.Resource, params *IndexParams, mergeParams *MergeParams, indices []*CagraIndex, mergedDataset PaddedDatasetHandle, offsets []int64, allowList []uint32, index *CagraIndex) error { + if len(indices) == 0 { + return errors.New("indices must not be empty") + } + if mergedDataset == nil || mergedDataset.datasetHandle() == nil { + return errors.New("mergedDataset is nil") + } + if len(offsets) != len(indices)+1 { + return errors.New("offsets must have len(indices)+1 entries") + } + + filter, filterTensor, err := buildAllowListFilter(Resources, allowList) + if err != nil { + return err + } + if filterTensor != nil { + defer filterTensor.Close() + } + + cIndices := make([]C.cuvsCagraIndex_t, len(indices)) + for i, idx := range indices { + cIndices[i] = idx.index + } + + cOffsets := make([]C.int64_t, len(offsets)) + for i, v := range offsets { + cOffsets[i] = C.int64_t(v) + } + + if mergeParams == nil { + err = cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraMerge( + C.cuvsResources_t(Resources.Resource), + params.params, + &cIndices[0], + C.size_t(len(indices)), + filter, + mergedDataset.datasetHandle(), + &cOffsets[0], + index.index, + ))) + } else { + err = cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraMergeWithParams( + C.cuvsResources_t(Resources.Resource), + params.params, + mergeParams.params, + &cIndices[0], + C.size_t(len(indices)), + filter, + mergedDataset.datasetHandle(), + &cOffsets[0], + index.index, + ))) + } + if err != nil { + return err + } + + index.trained = true + return nil +} + // Destroys the Cagra Index func (index *CagraIndex) Close() error { err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraIndexDestroy(index.index))) diff --git a/go/cagra/merge_params.go b/go/cagra/merge_params.go new file mode 100644 index 0000000000..fbcbcfa792 --- /dev/null +++ b/go/cagra/merge_params.go @@ -0,0 +1,107 @@ +package cagra + +// #include +import "C" + +import ( + "errors" + + cuvs "github.com/nvidia/cuvs/go" +) + +// Parameters controlling how physical CAGRA indices are merged. +type MergeParams struct { + params C.cuvsCagraMergeParams_t +} + +// Algorithm used to merge physical CAGRA indices. +type MergeAlgo int + +const ( + MergeAuto MergeAlgo = iota + MergeFastener + MergeRebuild +) + +var cMergeAlgos = map[MergeAlgo]int{ + MergeAuto: C.CUVS_CAGRA_MERGE_AUTO, + MergeFastener: C.CUVS_CAGRA_MERGE_FASTENER, + MergeRebuild: C.CUVS_CAGRA_MERGE_REBUILD, +} + +// Creates a new MergeParams, populated with AUTO defaults. +func CreateMergeParams() (*MergeParams, error) { + var params C.cuvsCagraMergeParams_t + + err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraMergeParamsCreate(¶ms))) + if err != nil { + return nil, err + } + + MergeParams := &MergeParams{params: params} + + return MergeParams, nil +} + +// Algorithm used to merge the physical CAGRA indices. +func (p *MergeParams) SetAlgo(algo MergeAlgo) (*MergeParams, error) { + CMergeAlgo, exists := cMergeAlgos[algo] + + if !exists { + return nil, errors.New("cuvs: invalid merge algo") + } + p.params.algo = uint32(CMergeAlgo) + + return p, nil +} + +// Number of levels used by the merge algorithm. +func (p *MergeParams) SetLevels(levels uint32) (*MergeParams, error) { + p.params.levels = C.uint32_t(levels) + return p, nil +} + +// Fanout of the root level. +func (p *MergeParams) SetRootFanout(root_fanout uint32) (*MergeParams, error) { + p.params.root_fanout = C.uint32_t(root_fanout) + return p, nil +} + +// Fanout of the lower levels. +func (p *MergeParams) SetLowerFanout(lower_fanout uint32) (*MergeParams, error) { + p.params.lower_fanout = C.uint32_t(lower_fanout) + return p, nil +} + +// Fraction of points selected as leaders. +func (p *MergeParams) SetLeaderFraction(leader_fraction float64) (*MergeParams, error) { + p.params.leader_fraction = C.double(leader_fraction) + return p, nil +} + +// Maximum number of leaders. +func (p *MergeParams) SetMaxLeaders(max_leaders uint32) (*MergeParams, error) { + p.params.max_leaders = C.uint32_t(max_leaders) + return p, nil +} + +// Size of the leaf partitions. +func (p *MergeParams) SetLeafSize(leaf_size uint32) (*MergeParams, error) { + p.params.leaf_size = C.uint32_t(leaf_size) + return p, nil +} + +// Degree used within the leaf partitions. +func (p *MergeParams) SetLeafDegree(leaf_degree uint32) (*MergeParams, error) { + p.params.leaf_degree = C.uint32_t(leaf_degree) + return p, nil +} + +// Destroys MergeParams +func (p *MergeParams) Close() error { + err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraMergeParamsDestroy(p.params))) + if err != nil { + return err + } + return nil +} diff --git a/go/cagra/merge_test.go b/go/cagra/merge_test.go new file mode 100644 index 0000000000..8dbea71bac --- /dev/null +++ b/go/cagra/merge_test.go @@ -0,0 +1,277 @@ +package cagra + +import ( + "math/rand/v2" + "testing" + + cuvs "github.com/nvidia/cuvs/go" +) + +func TestCagraMerge(t *testing.T) { + const ( + nDataPoints1 = 256 + nDataPoints2 = 256 + nFeatures = 16 + nQueries = 4 + k = 4 + epsilon = 0.001 + ) + r := rand.New(rand.NewPCG(7, 0)) + + resource, _ := cuvs.NewResource(nil) + defer resource.Close() + + // Build two disjoint datasets. + dataset1 := make([][]float32, nDataPoints1) + for i := range dataset1 { + dataset1[i] = make([]float32, nFeatures) + for j := range dataset1[i] { + dataset1[i][j] = r.Float32() + } + } + dataset2 := make([][]float32, nDataPoints2) + for i := range dataset2 { + dataset2[i] = make([]float32, nFeatures) + for j := range dataset2[i] { + dataset2[i][j] = r.Float32() + } + } + + // The merged dataset the caller must build: dataset1 || dataset2, in the + // same order the indices will be passed to MergeIndex. + mergedRaw := make([][]float32, 0, nDataPoints1+nDataPoints2) + mergedRaw = append(mergedRaw, dataset1...) + mergedRaw = append(mergedRaw, dataset2...) + + tensor1, err := cuvs.NewTensor(dataset1) + if err != nil { + t.Fatalf("error creating dataset1 tensor: %v", err) + } + defer tensor1.Close() + tensor2, err := cuvs.NewTensor(dataset2) + if err != nil { + t.Fatalf("error creating dataset2 tensor: %v", err) + } + defer tensor2.Close() + + if _, err := tensor1.ToDevice(&resource); err != nil { + t.Fatalf("error moving dataset1 to device: %v", err) + } + if _, err := tensor2.ToDevice(&resource); err != nil { + t.Fatalf("error moving dataset2 to device: %v", err) + } + + indexParams, err := CreateIndexParams() + if err != nil { + t.Fatalf("error creating index params: %v", err) + } + defer indexParams.Close() + + index1, err := CreateIndex() + if err != nil { + t.Fatalf("error creating index1: %v", err) + } + defer index1.Close() + index2, err := CreateIndex() + if err != nil { + t.Fatalf("error creating index2: %v", err) + } + defer index2.Close() + + if err := BuildIndex(resource, indexParams, &tensor1, index1); err != nil { + t.Fatalf("error building index1: %v", err) + } + if err := BuildIndex(resource, indexParams, &tensor2, index2); err != nil { + t.Fatalf("error building index2: %v", err) + } + + if err := resource.Sync(); err != nil { + t.Fatalf("error syncing resource: %v", err) + } + + mergedTensor, err := cuvs.NewTensor(mergedRaw) + if err != nil { + t.Fatalf("error creating merged tensor: %v", err) + } + defer mergedTensor.Close() + if _, err := mergedTensor.ToDevice(&resource); err != nil { + t.Fatalf("error moving merged dataset to device: %v", err) + } + + mergedDataset, err := MakePaddedDatasetAuto(resource, &mergedTensor) + if err != nil { + t.Fatalf("error making padded merged dataset: %v", err) + } + defer mergedDataset.Close() + + // Unfiltered merge: offsets are just the cumulative row counts. + offsets := []int64{0, nDataPoints1, nDataPoints1 + nDataPoints2} + + mergedIndex, err := CreateIndex() + if err != nil { + t.Fatalf("error creating merged index: %v", err) + } + defer mergedIndex.Close() + + if err := MergeIndex(resource, indexParams, []*CagraIndex{index1, index2}, mergedDataset, offsets, nil, mergedIndex); err != nil { + t.Fatalf("error merging indices: %v", err) + } + + if err := resource.Sync(); err != nil { + t.Fatalf("error syncing resource: %v", err) + } + + // Query with points from both halves; each should find itself as the + // nearest neighbor in the merged index. + queries := make([][]float32, 0, 2*nQueries) + queries = append(queries, dataset1[:nQueries]...) + queries = append(queries, dataset2[:nQueries]...) + expected := []uint32{0, 1, 2, 3, nDataPoints1, nDataPoints1 + 1, nDataPoints1 + 2, nDataPoints1 + 3} + + queriesTensor, err := cuvs.NewTensor(queries) + if err != nil { + t.Fatalf("error creating queries tensor: %v", err) + } + defer queriesTensor.Close() + if _, err := queriesTensor.ToDevice(&resource); err != nil { + t.Fatalf("error moving queries to device: %v", err) + } + + neighbors, err := cuvs.NewTensorOnDevice[uint32](&resource, []int64{int64(len(queries)), int64(k)}) + if err != nil { + t.Fatalf("error creating neighbors tensor: %v", err) + } + defer neighbors.Close() + + distances, err := cuvs.NewTensorOnDevice[float32](&resource, []int64{int64(len(queries)), int64(k)}) + if err != nil { + t.Fatalf("error creating distances tensor: %v", err) + } + defer distances.Close() + + searchParams, err := CreateSearchParams() + if err != nil { + t.Fatalf("error creating search params: %v", err) + } + defer searchParams.Close() + + if err := SearchIndex(resource, searchParams, mergedIndex, &queriesTensor, &neighbors, &distances, nil); err != nil { + t.Fatalf("error searching merged index: %v", err) + } + + if _, err := neighbors.ToHost(&resource); err != nil { + t.Fatalf("error moving neighbors to host: %v", err) + } + if _, err := distances.ToHost(&resource); err != nil { + t.Fatalf("error moving distances to host: %v", err) + } + if err := resource.Sync(); err != nil { + t.Fatalf("error syncing resource: %v", err) + } + + neighborsSlice, err := neighbors.Slice() + if err != nil { + t.Fatalf("error getting neighbors slice: %v", err) + } + distancesSlice, err := distances.Slice() + if err != nil { + t.Fatalf("error getting distances slice: %v", err) + } + + for i := range neighborsSlice { + if neighborsSlice[i][0] != expected[i] { + t.Error("wrong neighbor, expected", expected[i], "got", neighborsSlice[i][0]) + } + if distancesSlice[i][0] >= epsilon || distancesSlice[i][0] <= -epsilon { + t.Error("distance should be close to 0, got", distancesSlice[i][0]) + } + } +} + +func TestMergedDatasetOffsetsUnfiltered(t *testing.T) { + const ( + nDataPoints1 = 32 + nDataPoints2 = 48 + nFeatures = 8 + ) + r := rand.New(rand.NewPCG(11, 0)) + + resource, _ := cuvs.NewResource(nil) + defer resource.Close() + + dataset1 := make([][]float32, nDataPoints1) + for i := range dataset1 { + dataset1[i] = make([]float32, nFeatures) + for j := range dataset1[i] { + dataset1[i][j] = r.Float32() + } + } + dataset2 := make([][]float32, nDataPoints2) + for i := range dataset2 { + dataset2[i] = make([]float32, nFeatures) + for j := range dataset2[i] { + dataset2[i][j] = r.Float32() + } + } + + tensor1, err := cuvs.NewTensor(dataset1) + if err != nil { + t.Fatalf("error creating dataset1 tensor: %v", err) + } + defer tensor1.Close() + tensor2, err := cuvs.NewTensor(dataset2) + if err != nil { + t.Fatalf("error creating dataset2 tensor: %v", err) + } + defer tensor2.Close() + + if _, err := tensor1.ToDevice(&resource); err != nil { + t.Fatalf("error moving dataset1 to device: %v", err) + } + if _, err := tensor2.ToDevice(&resource); err != nil { + t.Fatalf("error moving dataset2 to device: %v", err) + } + + indexParams, err := CreateIndexParams() + if err != nil { + t.Fatalf("error creating index params: %v", err) + } + defer indexParams.Close() + + index1, err := CreateIndex() + if err != nil { + t.Fatalf("error creating index1: %v", err) + } + defer index1.Close() + index2, err := CreateIndex() + if err != nil { + t.Fatalf("error creating index2: %v", err) + } + defer index2.Close() + + if err := BuildIndex(resource, indexParams, &tensor1, index1); err != nil { + t.Fatalf("error building index1: %v", err) + } + if err := BuildIndex(resource, indexParams, &tensor2, index2); err != nil { + t.Fatalf("error building index2: %v", err) + } + + if err := resource.Sync(); err != nil { + t.Fatalf("error syncing resource: %v", err) + } + + offsets, err := MergedDatasetOffsets(resource, []*CagraIndex{index1, index2}, nil) + if err != nil { + t.Fatalf("error computing merged dataset offsets: %v", err) + } + + expected := []int64{0, nDataPoints1, nDataPoints1 + nDataPoints2} + if len(offsets) != len(expected) { + t.Fatalf("expected %d offsets, got %d", len(expected), len(offsets)) + } + for i := range expected { + if offsets[i] != expected[i] { + t.Errorf("offsets[%d]: expected %d, got %d", i, expected[i], offsets[i]) + } + } +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index 96e31431f8..4adc5593c1 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java @@ -305,26 +305,124 @@ static Builder newBuilder(CuVSResources cuvsResources) { /** * Merges multiple CAGRA indexes into a single index using default merge parameters. * + *

The caller is responsible for concatenating every input index's dataset (in {@code + * indexes} order) into a single caller-owned padded dataset before calling this, and for + * providing {@code offsets}: entry {@code i} is the row at which {@code indexes[i]}'s rows + * start in {@code mergedDataset}, and the last entry ({@code offsets[indexes.length]}) must + * equal {@code mergedDataset}'s total row count. For example, with no filtering, {@code + * offsets} is simply the cumulative row counts of {@code indexes} in order. This mirrors the + * caller-owned-buffer contract used by {@link #updateDataset(PaddedDataset)}. Keep {@code + * mergedDataset} alive for as long as the returned index remains in use. + * * @param indexes Array of CAGRA indexes to merge + * @param mergedDataset Caller-owned padded dataset holding the concatenation of every input + * index's rows, in {@code indexes} order + * @param offsets Per-index starting row within {@code mergedDataset}. Array of {@code + * indexes.length + 1} entries; the last entry must equal {@code mergedDataset}'s + * row count * @return A new merged CAGRA index * @throws Throwable if an error occurs during the merge operation */ - static CagraIndex merge(CagraIndex[] indexes) throws Throwable { - return merge(indexes, null); + static CagraIndex merge(CagraIndex[] indexes, PaddedDataset mergedDataset, long[] offsets) + throws Throwable { + return merge(indexes, mergedDataset, offsets, null); } /** * Merges multiple CAGRA indexes into a single index with the specified merge parameters. * + *

See {@link #merge(CagraIndex[], PaddedDataset, long[])} for the {@code mergedDataset}/ + * {@code offsets} contract. + * * @param indexes Array of CAGRA indexes to merge + * @param mergedDataset Caller-owned padded dataset holding the concatenation of every input + * index's rows, in {@code indexes} order + * @param offsets Per-index starting row within {@code mergedDataset}. Array of {@code + * indexes.length + 1} entries; the last entry must equal {@code mergedDataset}'s + * row count * @param mergeParams Parameters to control the merge operation, or null to use defaults * @return A new merged CAGRA index * @throws Throwable if an error occurs during the merge operation */ - static CagraIndex merge(CagraIndex[] indexes, CagraIndexParams mergeParams) throws Throwable { + static CagraIndex merge( + CagraIndex[] indexes, + PaddedDataset mergedDataset, + long[] offsets, + CagraIndexParams mergeParams) + throws Throwable { + validateMergeArgs(indexes, offsets); + Objects.requireNonNull(mergedDataset); + if (!mergedDataset.isPresent()) { + throw new IllegalArgumentException("mergedDataset is uninitialized"); + } + return CuVSProvider.provider() + .mergeCagraIndexes(indexes, mergedDataset.nativeHandleAddress(), offsets, mergeParams); + } + + /** + * Merges multiple CAGRA indexes into a single index using default merge parameters, from a + * caller-owned padded dataset view over a buffer that is already padded to CAGRA's required + * row stride. + * + *

See {@link #merge(CagraIndex[], PaddedDataset, long[])} for the {@code mergedDataset}/ + * {@code offsets} contract. + * + * @param indexes Array of CAGRA indexes to merge + * @param mergedDataset Caller-owned padded dataset view holding the concatenation of every + * input index's rows, in {@code indexes} order + * @param offsets Per-index starting row within {@code mergedDataset}. Array of {@code + * indexes.length + 1} entries; the last entry must equal {@code mergedDataset}'s + * row count + * @return A new merged CAGRA index + * @throws Throwable if an error occurs during the merge operation + */ + static CagraIndex merge(CagraIndex[] indexes, PaddedDatasetView mergedDataset, long[] offsets) + throws Throwable { + return merge(indexes, mergedDataset, offsets, null); + } + + /** + * Merges multiple CAGRA indexes into a single index with the specified merge parameters, from a + * caller-owned padded dataset view over a buffer that is already padded to CAGRA's required + * row stride. + * + *

See {@link #merge(CagraIndex[], PaddedDataset, long[])} for the {@code mergedDataset}/ + * {@code offsets} contract. + * + * @param indexes Array of CAGRA indexes to merge + * @param mergedDataset Caller-owned padded dataset view holding the concatenation of every + * input index's rows, in {@code indexes} order + * @param offsets Per-index starting row within {@code mergedDataset}. Array of {@code + * indexes.length + 1} entries; the last entry must equal {@code mergedDataset}'s + * row count + * @param mergeParams Parameters to control the merge operation, or null to use defaults + * @return A new merged CAGRA index + * @throws Throwable if an error occurs during the merge operation + */ + static CagraIndex merge( + CagraIndex[] indexes, + PaddedDatasetView mergedDataset, + long[] offsets, + CagraIndexParams mergeParams) + throws Throwable { + validateMergeArgs(indexes, offsets); + Objects.requireNonNull(mergedDataset); + if (!mergedDataset.isPresent()) { + throw new IllegalArgumentException("mergedDataset is uninitialized"); + } + return CuVSProvider.provider() + .mergeCagraIndexes(indexes, mergedDataset.nativeHandleAddress(), offsets, mergeParams); + } + + private static void validateMergeArgs(CagraIndex[] indexes, long[] offsets) { if (indexes == null || indexes.length == 0) { throw new IllegalArgumentException("At least one index must be provided for merging"); } + Objects.requireNonNull(offsets); + if (offsets.length != indexes.length + 1) { + throw new IllegalArgumentException( + "offsets must have indexes.length + 1 entries, got " + offsets.length); + } CuVSResources resources = indexes[0].getCuVSResources(); for (int i = 1; i < indexes.length; i++) { @@ -332,8 +430,6 @@ static CagraIndex merge(CagraIndex[] indexes, CagraIndexParams mergeParams) thro throw new IllegalArgumentException("All indexes must use the same CuVSResources instance"); } } - - return CuVSProvider.provider().mergeCagraIndexes(indexes, mergeParams); } /** diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java index a148ee76ae..def5da4e18 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java @@ -166,27 +166,31 @@ TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) throws UnsupportedOperationException; /** - * Merges multiple CAGRA indexes into a single index. + * Merges multiple CAGRA indexes into a single index, using a caller-owned pre-concatenated + * padded dataset. * - * @param indexes Array of CAGRA indexes to merge - * @return A new merged CAGRA index - * @throws Throwable if an error occurs during the merge operation - */ - CagraIndex mergeCagraIndexes(CagraIndex[] indexes) throws Throwable; - - /** - * Merges multiple CAGRA indexes into a single index with the specified merge parameters. + *

See {@link CagraIndex#merge(CagraIndex[], CagraIndex.PaddedDataset, long[])} for the full + * {@code mergedDatasetHandleAddress}/{@code offsets} contract; this SPI method takes the raw + * native handle address so implementations don't need to depend on the concrete dataset + * wrapper type. * * @param indexes Array of CAGRA indexes to merge + * @param mergedDatasetHandleAddress native handle address of the caller-owned padded dataset + * (or padded dataset view) holding the concatenation of every + * input index's rows, in {@code indexes} order + * @param offsets Per-index starting row within the merged dataset. Array of {@code + * indexes.length + 1} entries; the last entry must equal the merged dataset's + * row count * @param mergeParams Parameters to control the merge operation, or null to use defaults * @return A new merged CAGRA index * @throws Throwable if an error occurs during the merge operation */ - default CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraIndexParams mergeParams) - throws Throwable { - // Default implementation falls back to the method without parameters - return mergeCagraIndexes(indexes); - } + CagraIndex mergeCagraIndexes( + CagraIndex[] indexes, + long mergedDatasetHandleAddress, + long[] offsets, + CagraIndexParams mergeParams) + throws Throwable; /** * Reports whether the rows of {@code dataset} already sit at the row stride CAGRA requires, which diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index 66b0b8cc47..32f85e1a5f 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -81,7 +81,11 @@ public TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) { } @Override - public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) { + public CagraIndex mergeCagraIndexes( + CagraIndex[] indexes, + long mergedDatasetHandleAddress, + long[] offsets, + CagraIndexParams mergeParams) { throw new UnsupportedOperationException(reasons); } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 9506b130cd..5904a08c4d 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -923,23 +923,29 @@ public static CagraIndex.Builder newBuilder(CuVSResources cuvsResources) { } /** - * Merges multiple CAGRA indexes into a single index. + * Merges multiple CAGRA indexes into a single index, using a caller-owned pre-concatenated + * padded dataset. * - * @param indexes Array of CAGRA indexes to merge - * @return A new merged CAGRA index - */ - public static CagraIndex merge(CagraIndex[] indexes) { - return merge(indexes, null); - } - - /** - * Merges multiple CAGRA indexes into a single index with specified merge parameters. + *

The dataset stays owned by the caller: this method never allocates or copies dataset rows, + * it only merges the graphs and binds the returned index to the native dataset handle at {@code + * mergedDatasetHandleAddress}. Keep the underlying dataset alive while the returned index is in + * use. * * @param indexes Array of CAGRA indexes to merge + * @param mergedDatasetHandleAddress native handle address of the caller-owned padded dataset (or + * padded dataset view) holding the concatenation of every + * input index's rows, in {@code indexes} order + * @param offsets Per-index starting row within the merged dataset. Array of {@code + * indexes.length + 1} entries; the last entry must equal the merged dataset's row + * count * @param mergeParams Parameters to control the merge operation, or null to use defaults * @return A new merged CAGRA index */ - public static CagraIndex merge(CagraIndex[] indexes, CagraIndexParams mergeParams) { + public static CagraIndex merge( + CagraIndex[] indexes, + long mergedDatasetHandleAddress, + long[] offsets, + CagraIndexParams mergeParams) { CuVSResources resources = indexes[0].getCuVSResources(); var mergedIndex = createCagraIndex(); @@ -961,30 +967,22 @@ public static CagraIndex merge(CagraIndex[] indexes, CagraIndexParams mergeParam cuvsFilter.type(mergeFilter, 0); // NO_FILTER cuvsFilter.addr(mergeFilter, 0); - MemorySegment mergedDatasetPtr = localArena.allocate(cuvsDataset_t); - checkCuVSError(cuvsDatasetCreate(mergedDatasetPtr), "cuvsDatasetCreate"); - MemorySegment mergedDataset = mergedDatasetPtr.get(cuvsDataset_t, 0); - AutoCloseable datasetOwner = new DatasetCloseDelegate(mergedDataset); - try { - checkCuVSError( - cuvsCagraMerge( - cuvsRes, - nativeMergeParams.handle(), - indexesSegment, - indexes.length, - mergeFilter, - mergedDataset, - mergedIndex), - "cuvsCagraMerge"); - return new CagraIndexImpl(new IndexReference(mergedIndex, null, datasetOwner), resources); - } catch (Throwable e) { - try { - datasetOwner.close(); - } catch (Exception closeError) { - e.addSuppressed(closeError); - } - throw e; - } + MemorySegment mergedDataset = MemorySegment.ofAddress(mergedDatasetHandleAddress); + MemorySegment offsetsSegment = buildMemorySegment(localArena, offsets); + + checkCuVSError( + cuvsCagraMerge( + cuvsRes, + nativeMergeParams.handle(), + indexesSegment, + indexes.length, + mergeFilter, + mergedDataset, + offsetsSegment, + mergedIndex), + "cuvsCagraMerge"); + // mergedDataset is caller-owned; the returned index does not take ownership of it. + return new CagraIndexImpl(new IndexReference(mergedIndex, null, null), resources); } } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 9cc4a5499c..12ecbe5994 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -302,19 +302,15 @@ public TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) { } @Override - public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) { + public CagraIndex mergeCagraIndexes( + CagraIndex[] indexes, + long mergedDatasetHandleAddress, + long[] offsets, + CagraIndexParams mergeParams) { if (indexes == null || indexes.length == 0) { throw new IllegalArgumentException("At least one index must be provided for merging"); } - return CagraIndexImpl.merge(indexes); - } - - @Override - public CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraIndexParams mergeParams) { - if (indexes == null || indexes.length == 0) { - throw new IllegalArgumentException("At least one index must be provided for merging"); - } - return CagraIndexImpl.merge(indexes, mergeParams); + return CagraIndexImpl.merge(indexes, mergedDatasetHandleAddress, offsets, mergeParams); } @Override diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java index 500ef451f3..c0ba585f28 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java @@ -950,15 +950,29 @@ public void testMergingIndexes() throws Throwable { // Host-built indexes are not mergeable. Dim=2 is not 16-byte aligned, so upload to device, // allocate owning padded copies, and attach them before merge. Keep them alive until the // inputs are closed. + // + // The native merge contract now requires the caller to hand in the already-concatenated + // (vector1 || vector2) dataset plus each index's starting row within it. + float[][] mergedVectors = { + {0.0f, 0.0f}, + {1.0f, 1.0f}, + {10.0f, 10.0f}, + {11.0f, 11.0f} + }; + long[] mergeOffsets = {0L, vector1.length, vector1.length + (long) vector2.length}; + try (var device1 = CuVSMatrix.ofArray(vector1).toDevice(resources); var device2 = CuVSMatrix.ofArray(vector2).toDevice(resources); var padded1 = index1.makePaddedDataset(device1); - var padded2 = index2.makePaddedDataset(device2)) { + var padded2 = index2.makePaddedDataset(device2); + var mergedDevice = CuVSMatrix.ofArray(mergedVectors).toDevice(resources); + var mergedDataset = index1.makePaddedDataset(mergedDevice)) { index1.updateDataset(padded1); index2.updateDataset(padded2); log.trace("Merging indexes..."); - CagraIndex mergedIndex = CagraIndex.merge(new CagraIndex[] {index1, index2}); + CagraIndex mergedIndex = + CagraIndex.merge(new CagraIndex[] {index1, index2}, mergedDataset, mergeOffsets); log.trace("Merge completed successfully"); // Pin SINGLE_CTA; AUTO may pick MULTI_CTA, which drops neighbors on this tiny dataset. @@ -1068,16 +1082,33 @@ public void testMergeStrategies() throws Throwable { // Host-built indexes are not mergeable. Dim=2 is not 16-byte aligned, so upload to device, // allocate owning padded copies, and attach them before merge. Keep them alive until the // inputs are closed. + // + // The native merge contract now requires the caller to hand in the already-concatenated + // (vector1 || vector2) dataset plus each index's starting row within it. + float[][] mergedVectors = { + {0.0f, 0.0f}, + {1.0f, 1.0f}, + {10.0f, 10.0f}, + {11.0f, 11.0f} + }; + long[] mergeOffsets = {0L, vector1.length, vector1.length + (long) vector2.length}; + try (var device1 = CuVSMatrix.ofArray(vector1).toDevice(resources); var device2 = CuVSMatrix.ofArray(vector2).toDevice(resources); var padded1 = index1.makePaddedDataset(device1); - var padded2 = index2.makePaddedDataset(device2)) { + var padded2 = index2.makePaddedDataset(device2); + var mergedDevice = CuVSMatrix.ofArray(mergedVectors).toDevice(resources); + var mergedDataset = index1.makePaddedDataset(mergedDevice)) { index1.updateDataset(padded1); index2.updateDataset(padded2); log.trace("Merging indexes with PHYSICAL strategy..."); try (CagraIndex physicalMergedIndex = - CagraIndex.merge(new CagraIndex[] {index1, index2}, outputIndexParams)) { + CagraIndex.merge( + new CagraIndex[] {index1, index2}, + mergedDataset, + mergeOffsets, + outputIndexParams)) { log.trace("Physical merge completed successfully"); // Pin SINGLE_CTA; AUTO may pick MULTI_CTA, which drops neighbors on this tiny dataset. diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index aead683dc4..731234723c 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -395,6 +395,13 @@ static int distFuncToOrd(VectorSimilarityFunction func) { * * This is currently (and intentionally) marked as unused and will be plugged in later. * + *

The native {@code CagraIndex.merge} contract requires the caller to hand it a single + * pre-concatenated dataset (in {@code cagraIndexes} order) plus the per-index row offsets + * within that dataset; the merge call itself never reads or copies vectors. We therefore have to + * walk each source reader's own (unmerged, un-filtered) vectors for this field, in the same + * order we collect its {@link CagraIndex}, to build both the concatenated dataset and the + * matching offsets array. + * * @param fieldInfo instance of the FieldInfo * @param mergeState instance of the MergeState * @throws IOException I/O Exceptions @@ -403,6 +410,9 @@ static int distFuncToOrd(VectorSimilarityFunction func) { private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throws IOException { try { List cagraIndexes = new ArrayList<>(); + List mergedVectors = new ArrayList<>(); + List offsets = new ArrayList<>(); + offsets.add(0L); // We need this count so that the merged segment's meta information has the vector count. int totalVectorCount = 0; for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) { @@ -411,10 +421,19 @@ private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throw if (knnReader != null) { if (knnReader instanceof CuVS2510GPUVectorsReader cvr) { if (cvr != null) { - totalVectorCount += cvr.getFieldEntries().get(fieldInfo.number).count(); CagraIndex cagraIndex = getCagraIndexFromReader(cvr, fieldInfo.name); if (cagraIndex != null) { cagraIndexes.add(cagraIndex); + int count = cvr.getFieldEntries().get(fieldInfo.number).count(); + totalVectorCount += count; + // Append this reader's own vectors for the field, in the same order the native + // index stores them, so the concatenated dataset lines up with `cagraIndexes`. + FloatVectorValues values = knnReader.getFloatVectorValues(fieldInfo.name); + KnnVectorValues.DocIndexIterator iter = values.iterator(); + for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { + mergedVectors.add(values.vectorValue(iter.index()).clone()); + } + offsets.add((long) mergedVectors.size()); } } } else { @@ -425,9 +444,27 @@ private void mergeCagraIndexes(FieldInfo fieldInfo, MergeState mergeState) throw } } assert cagraIndexes.size() > 1; - CagraIndex mergedIndex = - CagraIndex.merge(cagraIndexes.toArray(new CagraIndex[cagraIndexes.size()])); - writeMergedCagraIndex(fieldInfo, mergedIndex, totalVectorCount); + long[] offsetsArray = offsets.stream().mapToLong(Long::longValue).toArray(); + CuVSMatrix mergedMatrix = + Utils.createFloatMatrix( + mergedVectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); + CagraIndex[] indexesArray = cagraIndexes.toArray(new CagraIndex[cagraIndexes.size()]); + CagraIndex mergedIndex; + try (var deviceVectors = mergedMatrix.toDevice(getCuVSResourcesInstance())) { + // cuVS rejects makePaddedDataset for a device matrix whose rows already sit at the + // required stride, and asks for a view over that storage instead (see writeCagraIndex). + if (CagraIndex.isPaddedDataset(deviceVectors)) { + try (var mergedDatasetView = cagraIndexes.get(0).makePaddedDatasetView(deviceVectors)) { + mergedIndex = CagraIndex.merge(indexesArray, mergedDatasetView, offsetsArray); + writeMergedCagraIndex(fieldInfo, mergedIndex, totalVectorCount); + } + } else { + try (var mergedDataset = cagraIndexes.get(0).makePaddedDataset(deviceVectors)) { + mergedIndex = CagraIndex.merge(indexesArray, mergedDataset, offsetsArray); + writeMergedCagraIndex(fieldInfo, mergedIndex, totalVectorCount); + } + } + } info( infoStream, COMPONENT, diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java index 1c74649843..34edec73b9 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java @@ -81,8 +81,13 @@ public HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources) } @Override - public CagraIndex mergeCagraIndexes(CagraIndex[] arg0) throws Throwable { - return delegate.mergeCagraIndexes(arg0); + public CagraIndex mergeCagraIndexes( + CagraIndex[] indexes, + long mergedDatasetHandleAddress, + long[] offsets, + CagraIndexParams mergeParams) + throws Throwable { + return delegate.mergeCagraIndexes(indexes, mergedDatasetHandleAddress, offsets, mergeParams); } @Override diff --git a/python/cuvs/cuvs/neighbors/cagra/__init__.py b/python/cuvs/cuvs/neighbors/cagra/__init__.py index 60811a23eb..7e3d0cef17 100644 --- a/python/cuvs/cuvs/neighbors/cagra/__init__.py +++ b/python/cuvs/cuvs/neighbors/cagra/__init__.py @@ -9,11 +9,14 @@ ExtendParams, Index, IndexParams, + MergeParams, SearchParams, build, extend, from_graph, load, + merge, + merged_dataset_offsets, save, search, update_dataset, @@ -25,11 +28,14 @@ "ExtendParams", "Index", "IndexParams", + "MergeParams", "SearchParams", "build", "extend", "from_graph", "load", + "merge", + "merged_dataset_offsets", "save", "search", "update_dataset", diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd index 9e4dbdb6f3..581aefd421 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd @@ -187,6 +187,51 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: int64_t new_start_row, cuvsCagraIndex_t index) + ctypedef enum cuvsCagraMergeAlgo: + CUVS_CAGRA_MERGE_AUTO + CUVS_CAGRA_MERGE_FASTENER + CUVS_CAGRA_MERGE_REBUILD + + ctypedef struct cuvsCagraMergeParams: + cuvsCagraMergeAlgo algo + uint32_t levels + uint32_t root_fanout + uint32_t lower_fanout + double leader_fraction + uint32_t max_leaders + uint32_t leaf_size + uint32_t leaf_degree + + ctypedef cuvsCagraMergeParams* cuvsCagraMergeParams_t + + cuvsError_t cuvsCagraMergeParamsCreate(cuvsCagraMergeParams_t* params) + cuvsError_t cuvsCagraMergeParamsDestroy(cuvsCagraMergeParams_t params) + + cuvsError_t cuvsCagraMergedDatasetOffsets(cuvsResources_t res, + cuvsCagraIndex_t* indices, + size_t num_indices, + cuvsFilter filter, + int64_t* offsets) + + cuvsError_t cuvsCagraMerge(cuvsResources_t res, + cuvsCagraIndexParams_t params, + cuvsCagraIndex_t* indices, + size_t num_indices, + cuvsFilter filter, + cuvsDataset_t merged_dataset, + const int64_t* offsets, + cuvsCagraIndex_t output_index) + + cuvsError_t cuvsCagraMergeWithParams(cuvsResources_t res, + cuvsCagraIndexParams_t params, + cuvsCagraMergeParams_t merge_params, + cuvsCagraIndex_t* indices, + size_t num_indices, + cuvsFilter filter, + cuvsDataset_t merged_dataset, + const int64_t* offsets, + cuvsCagraIndex_t output_index) + cdef class Index: """ diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index d10895bcf8..25d103cc67 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -1152,3 +1152,267 @@ def extend(ExtendParams params, Index index, extended_dataset, new_start_row, _keep_dataset_alive(index, dataset_obj, source_array) return index + + +cdef class MergeParams: + """ Supplemental parameters controlling how physical CAGRA indices are + merged. + + Parameters + ---------- + algo : str, default = "auto" + String denoting the merge algorithm to use. Valid values for + algo: ["auto", "fastener", "rebuild"], where + + - auto will automatically select the merge algorithm + - fastener will stitch the input graphs together + - rebuild will build the output graph from scratch + levels : int, default = 0 + root_fanout : int, default = 0 + lower_fanout : int, default = 0 + leader_fraction : float, default = 0 + max_leaders : int, default = 0 + leaf_size : int, default = 0 + leaf_degree : int, default = 0 + """ + + cdef cuvsCagraMergeParams* params + + def __cinit__(self): + check_cuvs(cuvsCagraMergeParamsCreate(&self.params)) + + def __dealloc__(self): + if self.params != NULL: + check_cuvs(cuvsCagraMergeParamsDestroy(self.params)) + + def __init__(self, *, + algo="auto", + levels=None, + root_fanout=None, + lower_fanout=None, + leader_fraction=None, + max_leaders=None, + leaf_size=None, + leaf_degree=None): + if algo == "auto": + self.params.algo = cuvsCagraMergeAlgo.CUVS_CAGRA_MERGE_AUTO + elif algo == "fastener": + self.params.algo = cuvsCagraMergeAlgo.CUVS_CAGRA_MERGE_FASTENER + elif algo == "rebuild": + self.params.algo = cuvsCagraMergeAlgo.CUVS_CAGRA_MERGE_REBUILD + else: + raise ValueError(f"Unknown algo '{algo}'") + + if levels is not None: + self.params.levels = levels + if root_fanout is not None: + self.params.root_fanout = root_fanout + if lower_fanout is not None: + self.params.lower_fanout = lower_fanout + if leader_fraction is not None: + self.params.leader_fraction = leader_fraction + if max_leaders is not None: + self.params.max_leaders = max_leaders + if leaf_size is not None: + self.params.leaf_size = leaf_size + if leaf_degree is not None: + self.params.leaf_degree = leaf_degree + + def get_handle(self): + return self.params + + @property + def algo(self): + algo = self.params.algo + if algo == cuvsCagraMergeAlgo.CUVS_CAGRA_MERGE_AUTO: + return "auto" + elif algo == cuvsCagraMergeAlgo.CUVS_CAGRA_MERGE_FASTENER: + return "fastener" + elif algo == cuvsCagraMergeAlgo.CUVS_CAGRA_MERGE_REBUILD: + return "rebuild" + + @property + def levels(self): + return self.params.levels + + @property + def root_fanout(self): + return self.params.root_fanout + + @property + def lower_fanout(self): + return self.params.lower_fanout + + @property + def leader_fraction(self): + return self.params.leader_fraction + + @property + def max_leaders(self): + return self.params.max_leaders + + @property + def leaf_size(self): + return self.params.leaf_size + + @property + def leaf_degree(self): + return self.params.leaf_degree + + +cdef cuvsCagraIndex_t* _index_array(indices) except NULL: + """ Allocate and fill a C array of cuvsCagraIndex_t handles from a + Python sequence of `Index` objects. Caller must `free()` the result. """ + cdef size_t n = len(indices) + cdef cuvsCagraIndex_t* arr = malloc( + n * sizeof(cuvsCagraIndex_t)) + if arr == NULL: + raise MemoryError("Could not allocate index array") + cdef Index idx + cdef size_t i + for i in range(n): + idx = indices[i] + if not idx.trained: + free(arr) + raise ValueError("All input indices must be trained/built") + arr[i] = idx.index + return arr + + +@auto_sync_resources +def merged_dataset_offsets(indices, filter=None, resources=None): + """ + Compute the per-index starting row offsets required by ``merge``. + + Only needed when merging with a bitset ``filter``. For an unfiltered + merge, the offsets are simply the cumulative row counts of ``indices`` + and this function is not required. + + Parameters + ---------- + indices : list[Index] + Input CAGRA indices that will be passed to ``merge``. + filter : Optional cuvs.neighbors.filters.Prefilter + Filter that will be passed to ``merge``. Only bitset filters (or no + filter) are supported. (default None) + {resources_docstring} + + Returns + ------- + offsets : numpy.ndarray of dtype int64, shape (len(indices) + 1,) + """ + if filter is None: + filter = no_filter() + + cdef size_t num_indices = len(indices) + cdef cuvsCagraIndex_t* c_indices = _index_array(indices) + cdef cuvsResources_t res = resources.get_c_obj() + + offsets = np.empty(num_indices + 1, dtype=np.int64) + cdef int64_t[::1] offsets_view = offsets + + try: + with cuda_interruptible(): + check_cuvs(cuvsCagraMergedDatasetOffsets( + res, + c_indices, + num_indices, + filter.prefilter, + &offsets_view[0])) + finally: + free(c_indices) + + return offsets + + +@auto_sync_resources +def merge(IndexParams params, indices, merged_dataset, offsets, + merge_params=None, filter=None, resources=None): + """ + Merge multiple CAGRA indices into a single CAGRA index. + + The caller owns dataset concatenation. Build a single padded + ``merged_dataset`` containing the concatenation, in ``indices`` order, + of every input index's dataset (with ``filter`` already applied, if + any). ``offsets`` gives the row at which each input index's (post + filter) rows begin in ``merged_dataset``; see + ``merged_dataset_offsets`` for the bitset-filtered case. For an + unfiltered merge, ``offsets`` is just the cumulative row counts of + ``indices``. + + This function only merges the graph and rebinds the output index to + ``merged_dataset``; keep that dataset alive for the resulting index's + lifetime. + + Parameters + ---------- + params : IndexParams object + Parameters for the output (merged) index. + indices : list[Index] + Input CAGRA indices to merge. + merged_dataset : Dataset or array + Padded dataset already containing the concatenated (and, if + ``filter`` is set, already-filtered) rows of every input index. + offsets : array-like of int64, shape (len(indices) + 1,) + Per-index starting row within ``merged_dataset``. The last entry + must equal ``merged_dataset``'s row count. + merge_params : MergeParams, optional + Parameters controlling the merge algorithm. Defaults to AUTO. + filter : Optional cuvs.neighbors.filters.Prefilter + Filter already applied by the caller while building + ``merged_dataset``. (default None) + {resources_docstring} + + Returns + ------- + index: cuvs.cagra.Index + """ + cdef Dataset dataset_obj + source_array = None + if isinstance(merged_dataset, Dataset): + dataset_obj = merged_dataset + else: + source_array = merged_dataset + dataset_obj = make_device_padded_dataset(merged_dataset, resources=resources) + + cdef cuvsDataset_t merged_handle = _cagra_dataset_handle(dataset_obj) + if dataset_obj.layout != "padded": + raise TypeError("merged_dataset must have padded layout") + + if filter is None: + filter = no_filter() + + cdef size_t num_indices = len(indices) + offsets_arr = np.ascontiguousarray(offsets, dtype=np.int64) + if offsets_arr.shape[0] != num_indices + 1: + raise ValueError( + "offsets must have len(indices) + 1 entries") + cdef int64_t[::1] offsets_view = offsets_arr + + cdef cuvsCagraIndex_t* c_indices = _index_array(indices) + cdef Index out_idx = Index() + cdef cuvsResources_t res = resources.get_c_obj() + cdef cuvsCagraMergeParams_t merge_params_ptr = NULL + if merge_params is not None: + merge_params_ptr = (merge_params).params + + try: + with cuda_interruptible(): + check_cuvs(cuvsCagraMergeWithParams( + res, + params.params, + merge_params_ptr, + c_indices, + num_indices, + filter.prefilter, + merged_handle, + &offsets_view[0], + out_idx.index)) + finally: + free(c_indices) + + out_idx.trained = True + out_idx.active_index_type = np.dtype( + dl_data_type_to_numpy(out_idx.index.dtype)).name + _keep_dataset_alive(out_idx, dataset_obj, source_array) + return out_idx diff --git a/python/cuvs/cuvs/tests/test_cagra.py b/python/cuvs/cuvs/tests/test_cagra.py index 25893a16f7..6e65768a44 100644 --- a/python/cuvs/cuvs/tests/test_cagra.py +++ b/python/cuvs/cuvs/tests/test_cagra.py @@ -226,6 +226,77 @@ def test_cagra_build_from_dataset_handle( assert distances.shape == (n_queries, k) +@pytest.mark.parametrize("n_rows", [1024, 2048]) +@pytest.mark.parametrize("n_cols", [16, 32]) +def test_cagra_merge(n_rows, n_cols): + n_queries = 32 + k = 10 + dtype = np.float32 + + dataset = generate_data((n_rows, n_cols), dtype) + split = n_rows // 2 + dataset_1 = dataset[:split, :] + dataset_2 = dataset[split:, :] + + build_params = cagra.IndexParams(metric="sqeuclidean") + index_1 = cagra.build(build_params, device_ndarray(dataset_1)) + index_2 = cagra.build(build_params, device_ndarray(dataset_2)) + + # unfiltered merge: offsets are simply the cumulative row counts + offsets = np.array( + [0, dataset_1.shape[0], dataset_1.shape[0] + dataset_2.shape[0]], + dtype=np.int64, + ) + + merged_dataset = make_device_padded_dataset( + device_ndarray(np.concatenate((dataset_1, dataset_2), axis=0)) + ) + + merge_params = cagra.IndexParams(metric="sqeuclidean") + merged_index = cagra.merge( + merge_params, + [index_1, index_2], + merged_dataset, + offsets, + ) + + assert merged_index.trained + assert len(merged_index) == n_rows + + queries = generate_data((n_queries, n_cols), dtype) + queries_device = device_ndarray(queries) + + search_params = cagra.SearchParams() + dist_device, idx_device = cagra.search( + search_params, merged_index, queries_device, k + ) + + nn_skl = NearestNeighbors( + n_neighbors=k, algorithm="brute", metric="sqeuclidean" + ) + nn_skl.fit(dataset) + skl_idx = nn_skl.kneighbors(queries, return_distance=False) + + recall = calc_recall(idx_device.copy_to_host(), skl_idx) + assert recall > 0.7 + + # also exercise merged_dataset_offsets() for the unfiltered case; it + # should reproduce the same cumulative offsets computed above. + computed_offsets = cagra.merged_dataset_offsets([index_1, index_2]) + assert np.array_equal(computed_offsets, offsets) + + # merging with explicit MergeParams (rebuild algo) should also work + merged_index_rebuild = cagra.merge( + merge_params, + [index_1, index_2], + merged_dataset, + offsets, + merge_params=cagra.MergeParams(algo="rebuild"), + ) + assert merged_index_rebuild.trained + assert len(merged_index_rebuild) == n_rows + + @pytest.mark.parametrize("sparsity", [0.2, 0.5, 0.7, 1.0]) def test_filtered_cagra(sparsity): run_filtered_search_test(cagra, sparsity) diff --git a/rust/cuvs-sys/src/bindings.rs b/rust/cuvs-sys/src/bindings.rs index e723abaaea..57e3e379ab 100644 --- a/rust/cuvs-sys/src/bindings.rs +++ b/rust/cuvs-sys/src/bindings.rs @@ -5,25 +5,42 @@ use crate::{cudaDataType_t, cudaStream_t}; #[repr(u32)] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum DLDeviceType { + #[doc = " \\brief CPU device"] kDLCPU = 1, + #[doc = " \\brief CUDA GPU device"] kDLCUDA = 2, + #[doc = " \\brief Pinned CUDA CPU memory by cudaMallocHost"] kDLCUDAHost = 3, + #[doc = " \\brief OpenCL devices."] kDLOpenCL = 4, + #[doc = " \\brief Vulkan buffer for next generation graphics."] kDLVulkan = 7, + #[doc = " \\brief Metal for Apple GPU."] kDLMetal = 8, + #[doc = " \\brief Verilog simulator buffer"] kDLVPI = 9, + #[doc = " \\brief ROCm GPUs for AMD GPUs"] kDLROCM = 10, + #[doc = " \\brief Pinned ROCm CPU memory allocated by hipMallocHost"] kDLROCMHost = 11, + #[doc = " \\brief Reserved extension device type,\n used for quickly test extension device\n The semantics can differ depending on the implementation."] kDLExtDev = 12, + #[doc = " \\brief CUDA managed/unified memory allocated by cudaMallocManaged"] kDLCUDAManaged = 13, + #[doc = " \\brief Unified shared memory allocated on a oneAPI non-partititioned\n device. Call to oneAPI runtime is required to determine the device\n type, the USM allocation type and the sycl context it is bound to.\n"] kDLOneAPI = 14, + #[doc = " \\brief GPU support for next generation WebGPU standard."] kDLWebGPU = 15, + #[doc = " \\brief Qualcomm Hexagon DSP"] kDLHexagon = 16, } +#[doc = " \\brief A Device for Tensor and operator."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct DLDevice { + #[doc = " \\brief The device type used in the device."] pub device_type: DLDeviceType, + #[doc = " \\brief The device index.\n For vanilla CPU memory, pinned memory, or managed memory, this is set to 0."] pub device_id: i32, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -35,21 +52,33 @@ const _: () = { ["Offset of field: DLDevice::device_id"][::std::mem::offset_of!(DLDevice, device_id) - 4usize]; }; #[repr(u32)] +#[doc = " \\brief The type code options DLDataType."] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum DLDataTypeCode { + #[doc = " \\brief signed integer"] kDLInt = 0, + #[doc = " \\brief unsigned integer"] kDLUInt = 1, + #[doc = " \\brief IEEE floating point"] kDLFloat = 2, + #[doc = " \\brief Opaque handle type, reserved for testing purposes.\n Frameworks need to agree on the handle data type for the exchange to be well-defined."] kDLOpaqueHandle = 3, + #[doc = " \\brief bfloat16"] kDLBfloat = 4, + #[doc = " \\brief complex number\n (C/C++/Python layout: compact struct per complex number)"] kDLComplex = 5, + #[doc = " \\brief boolean"] kDLBool = 6, } +#[doc = " \\brief The data type the tensor can hold. The data type is assumed to follow the\n native endian-ness. An explicit error message should be raised when attempting to\n export an array with non-native endianness\n\n Examples\n - float: type_code = 2, bits = 32, lanes = 1\n - float4(vectorized 4 float): type_code = 2, bits = 32, lanes = 4\n - int8: type_code = 0, bits = 8, lanes = 1\n - std::complex: type_code = 5, bits = 64, lanes = 1\n - bool: type_code = 6, bits = 8, lanes = 1 (as per common array library convention, the underlying storage size of bool is 8 bits)"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct DLDataType { + #[doc = " \\brief Type code of base types.\n We keep it uint8_t instead of DLDataTypeCode for minimal memory\n footprint, but the value should be one of DLDataTypeCode enum values."] pub code: u8, + #[doc = " \\brief Number of bits, common choices are 8, 16, 32."] pub bits: u8, + #[doc = " \\brief Number of lanes in the type, used for vector types."] pub lanes: u16, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -60,15 +89,23 @@ const _: () = { ["Offset of field: DLDataType::bits"][::std::mem::offset_of!(DLDataType, bits) - 1usize]; ["Offset of field: DLDataType::lanes"][::std::mem::offset_of!(DLDataType, lanes) - 2usize]; }; +#[doc = " \\brief Plain C Tensor object, does not manage memory."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct DLTensor { + #[doc = " \\brief The data pointer points to the allocated data. This will be CUDA\n device pointer or cl_mem handle in OpenCL. It may be opaque on some device\n types. This pointer is always aligned to 256 bytes as in CUDA. The\n `byte_offset` field should be used to point to the beginning of the data.\n\n Note that as of Nov 2021, multiply libraries (CuPy, PyTorch, TensorFlow,\n TVM, perhaps others) do not adhere to this 256 byte aligment requirement\n on CPU/CUDA/ROCm, and always use `byte_offset=0`. This must be fixed\n (after which this note will be updated); at the moment it is recommended\n to not rely on the data pointer being correctly aligned.\n\n For given DLTensor, the size of memory required to store the contents of\n data is calculated as follows:\n\n \\code{.c}\n static inline size_t GetDataSize(const DLTensor* t) {\n size_t size = 1;\n for (tvm_index_t i = 0; i < t->ndim; ++i) {\n size *= t->shape[i];\n }\n size *= (t->dtype.bits * t->dtype.lanes + 7) / 8;\n return size;\n }\n \\endcode"] pub data: *mut ::std::os::raw::c_void, + #[doc = " \\brief The device of the tensor"] pub device: DLDevice, + #[doc = " \\brief Number of dimensions"] pub ndim: i32, + #[doc = " \\brief The data type of the pointer"] pub dtype: DLDataType, + #[doc = " \\brief The shape of the tensor"] pub shape: *mut i64, + #[doc = " \\brief strides of the tensor (in number of elements, not bytes)\n can be NULL, indicating tensor is compact and row-majored."] pub strides: *mut i64, + #[doc = " \\brief The offset in bytes to the beginning pointer to data"] pub byte_offset: u64, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -84,11 +121,15 @@ const _: () = { ["Offset of field: DLTensor::byte_offset"] [::std::mem::offset_of!(DLTensor, byte_offset) - 40usize]; }; +#[doc = " \\brief C Tensor object, manage memory of DLTensor. This data structure is\n intended to facilitate the borrowing of DLTensor by another framework. It is\n not meant to transfer the tensor. When the borrowing framework doesn't need\n the tensor, it should call the deleter to notify the host that the resource\n is no longer needed."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct DLManagedTensor { + #[doc = " \\brief DLTensor which is being memory managed"] pub dl_tensor: DLTensor, + #[doc = " \\brief the context of the original host framework of DLManagedTensor in\n which DLManagedTensor is used in the framework. It can also be NULL."] pub manager_ctx: *mut ::std::os::raw::c_void, + #[doc = " \\brief Destructor signature void (*)(void*) - this should be called\n to destruct manager_ctx which holds the DLManagedTensor. It can be NULL\n if there is no way for the caller to provide a reasonable destructor.\n The destructors deletes the argument self as well."] pub deleter: ::std::option::Option, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -103,6 +144,7 @@ const _: () = { [::std::mem::offset_of!(DLManagedTensor, deleter) - 56usize]; }; #[repr(u32)] +#[doc = " @defgroup error_c cuVS Error Messages\n @{\n/\n/**\n @brief An enum denoting error statuses for function calls\n"] #[must_use] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsError_t { @@ -110,12 +152,15 @@ pub enum cuvsError_t { CUVS_SUCCESS = 1, } unsafe extern "C" { + #[doc = " @brief Returns a string describing the last seen error on this thread, or\n NULL if the last function succeeded."] pub fn cuvsGetLastErrorText() -> *const ::std::os::raw::c_char; } unsafe extern "C" { + #[doc = " @brief Sets a string describing an error seen on the thread. Passing NULL\n clears any previously seen error message."] pub fn cuvsSetLastErrorText(error: *const ::std::os::raw::c_char); } #[repr(u32)] +#[doc = " @brief An enum denoting log levels\n"] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsLogLevel_t { CUVS_LOG_LEVEL_TRACE = 0, @@ -127,18 +172,23 @@ pub enum cuvsLogLevel_t { CUVS_LOG_LEVEL_OFF = 6, } unsafe extern "C" { + #[doc = " @brief Returns the current log level"] pub fn cuvsGetLogLevel() -> cuvsLogLevel_t; } unsafe extern "C" { + #[doc = " @brief Sets the log level"] pub fn cuvsSetLogLevel(arg1: cuvsLogLevel_t); } +#[doc = " @brief An opaque C handle for C++ type `raft::resources`\n"] pub type cuvsResources_t = usize; unsafe extern "C" { #[must_use] + #[doc = " @brief Create an Initialized opaque C handle for C++ type `raft::resources`\n\n @param[in] res cuvsResources_t opaque C handle\n @return cuvsError_t"] pub fn cuvsResourcesCreate(res: *mut cuvsResources_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Create an opaque C handle for C++ type `raft::resources` whose memory\n allocations are tracked and written as CSV samples from a background\n thread.\n\n The returned handle wraps all reachable memory resources (host, pinned,\n managed, device, workspace, large_workspace) with allocation-tracking\n adaptors and replaces the global host and device memory resources for the\n lifetime of the handle. It is otherwise indistinguishable from a handle\n created by ::cuvsResourcesCreate and can be used wherever a\n ::cuvsResources_t is accepted. The CSV reporter is stopped and the global\n memory resources are restored when the handle is destroyed via\n ::cuvsResourcesDestroy.\n\n @param[out] res cuvsResources_t opaque C handle\n @param[in] csv_path Path to the output CSV file\n (created/truncated). Must be a non-empty,\n null-terminated UTF-8 string.\n @param[in] sample_interval_ms Minimum time in milliseconds between\n successive CSV samples. Pass 10 to match the\n C++ default.\n @return cuvsError_t"] pub fn cuvsResourcesCreateWithMemoryTracking( res: *mut cuvsResources_t, csv_path: *const ::std::os::raw::c_char, @@ -147,22 +197,27 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Destroy and de-allocate opaque C handle for C++ type `raft::resources`\n\n @param[in] res cuvsResources_t opaque C handle\n @return cuvsError_t"] pub fn cuvsResourcesDestroy(res: cuvsResources_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Set cudaStream_t on cuvsResources_t to queue CUDA kernels on APIs\n that accept a cuvsResources_t handle\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] stream cudaStream_t stream to queue CUDA kernels\n @return cuvsError_t"] pub fn cuvsStreamSet(res: cuvsResources_t, stream: cudaStream_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the cudaStream_t from a cuvsResources_t\n\n @param[in] res cuvsResources_t opaque C handle\n @param[out] stream cudaStream_t stream to queue CUDA kernels\n @return cuvsError_t"] pub fn cuvsStreamGet(res: cuvsResources_t, stream: *mut cudaStream_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Syncs the current CUDA stream on the resources object\n\n @param[in] res cuvsResources_t opaque C handle\n @return cuvsError_t"] pub fn cuvsStreamSync(res: cuvsResources_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the id of the device associated with this cuvsResources_t\n\n @param[in] res cuvsResources_t opaque C handle\n @param[out] device_id int the id of the device associated with res\n @return cuvsError_t"] pub fn cuvsDeviceIdGet( res: cuvsResources_t, device_id: *mut ::std::os::raw::c_int, @@ -170,6 +225,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Configure the temporary workspace on this resources object as an uncapped pool, backed\n by the current device memory resource. After the initial reservation is allocated on\n first use, subsequent calls to cuvsRMMAlloc / cuvsRMMFree on the same resources handle\n hit the pool cache rather than calling cudaMallocAsync / cudaFreeAsync, reducing CUDA\n context lock contention under concurrent query threads. The pool grows without shrinking:\n freed allocations are returned to the pool rather than to the device, so the pool's\n high-water mark only increases until the resources object is destroyed.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] initial_size_bytes initial pool reservation in bytes; size to cover the\n steady-state working set to avoid growth after warmup\n @return cuvsError_t"] pub fn cuvsResourcesSetWorkspacePool( res: cuvsResources_t, initial_size_bytes: usize, @@ -177,10 +233,12 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Create an Initialized opaque C handle for C++ type `raft::device_resources_snmg`\n for multi-GPU operations\n\n @param[in] res cuvsResources_t opaque C handle\n @return cuvsError_t"] pub fn cuvsMultiGpuResourcesCreate(res: *mut cuvsResources_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Create an Initialized opaque C handle for C++ type `raft::device_resources_snmg`\n for multi-GPU operations with specific device IDs\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] device_ids DLManagedTensor* containing device IDs to use\n @return cuvsError_t"] pub fn cuvsMultiGpuResourcesCreateWithDeviceIds( res: *mut cuvsResources_t, device_ids: *mut DLManagedTensor, @@ -188,10 +246,12 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Destroy and de-allocate opaque C handle for C++ type `raft::device_resources_snmg`\n\n @param[in] res cuvsResources_t opaque C handle\n @return cuvsError_t"] pub fn cuvsMultiGpuResourcesDestroy(res: cuvsResources_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Set a memory pool on all devices managed by the multi-GPU resources\n\n @param[in] res cuvsResources_t opaque C handle for multi-GPU resources\n @param[in] percent_of_free_memory Percent of free memory to allocate for the pool\n @return cuvsError_t"] pub fn cuvsMultiGpuResourcesSetMemoryPool( res: cuvsResources_t, percent_of_free_memory: ::std::os::raw::c_int, @@ -199,6 +259,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Allocates device memory using RMM\n\n\n @param[in] res cuvsResources_t opaque C handle\n @param[out] ptr Pointer to allocated device memory\n @param[in] bytes Size in bytes to allocate\n @return cuvsError_t"] pub fn cuvsRMMAlloc( res: cuvsResources_t, ptr: *mut *mut ::std::os::raw::c_void, @@ -207,6 +268,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Deallocates device memory using RMM\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] ptr Pointer to allocated device memory to free\n @param[in] bytes Size in bytes to allocate\n @return cuvsError_t"] pub fn cuvsRMMFree( res: cuvsResources_t, ptr: *mut ::std::os::raw::c_void, @@ -215,6 +277,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Switches the working memory resource to use the RMM pool memory resource, which will\n bypass unnecessary synchronizations by allocating a chunk of device memory up front and carving\n that up for temporary memory allocations within algorithms. Be aware that this function will\n change the memory resource for the whole process and the new memory resource will be used until\n explicitly changed.\n\n @param[in] initial_pool_size_percent The initial pool size as a percentage of the total\n available memory\n @param[in] max_pool_size_percent The maximum pool size as a percentage of the total\n available memory\n @param[in] managed Whether to use a managed memory resource as upstream resource or not\n @return cuvsError_t"] pub fn cuvsRMMPoolMemoryResourceEnable( initial_pool_size_percent: ::std::os::raw::c_int, max_pool_size_percent: ::std::os::raw::c_int, @@ -223,26 +286,32 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Switches the working memory resource to use stream-ordered asynchronous allocation\n (cudaMallocAsync / cudaFreeAsync). Unlike the pool resource, this resource returns memory to\n the stream immediately without blocking the CPU, eliminating device-wide synchronization on\n deallocation. This is especially beneficial when multiple CAGRA searches run concurrently on\n separate CUDA streams, because the internal workspace allocations no longer serialize kernel\n launches. Be aware that this function will change the memory resource for the whole process\n and the new memory resource will be used until explicitly changed.\n\n @return cuvsError_t"] pub fn cuvsRMMAsyncMemoryResourceEnable() -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Resets the memory resource to use the default memory resource (cuda_memory_resource)\n @return cuvsError_t"] pub fn cuvsRMMMemoryResourceReset() -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Allocates pinned memory on the host using RMM\n @param[out] ptr Pointer to allocated host memory\n @param[in] bytes Size in bytes to allocate\n @return cuvsError_t"] pub fn cuvsRMMHostAlloc(ptr: *mut *mut ::std::os::raw::c_void, bytes: usize) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Deallocates pinned memory on the host using RMM\n @param[in] ptr Pointer to allocated host memory to free\n @param[in] bytes Size in bytes to deallocate\n @return cuvsError_t"] pub fn cuvsRMMHostFree(ptr: *mut ::std::os::raw::c_void, bytes: usize) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the version of the cuVS library\n @param[out] major Major version\n @param[out] minor Minor version\n @param[out] patch Patch version\n @return cuvsError_t"] pub fn cuvsVersionGet(major: *mut u16, minor: *mut u16, patch: *mut u16) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Copy a matrix\n\n This function copies a matrix from dst to src. This lets you copy a matrix\n from device memory to host memory (or vice versa), while accounting for\n differences in strides.\n\n Both src and dst must have the same shape and dtype, but can have different\n strides and device type. The memory for the output dst tensor must already be\n allocated and the tensor initialized.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] src Pointer to DLManagedTensor to copy\n @param[out] dst Pointer to DLManagedTensor to receive copy of data"] pub fn cuvsMatrixCopy( res: cuvsResources_t, src: *mut DLManagedTensor, @@ -251,6 +320,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Slices rows from a matrix\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] src Pointer to DLManagedTensor to copy\n @param[in] start First row index to include in the output\n @param[in] end Last row index to include in the output\n @param[out] dst Pointer to DLManagedTensor to receive slice from matrix"] pub fn cuvsMatrixSliceRows( res: cuvsResources_t, src: *mut DLManagedTensor, @@ -260,17 +330,20 @@ unsafe extern "C" { ) -> cuvsError_t; } #[repr(u32)] +#[doc = " @brief Generic dataset layout kind for C API dataset handles."] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsDatasetLayout_t { CUVS_DATASET_LAYOUT_STANDARD = 0, CUVS_DATASET_LAYOUT_PADDED = 1, } #[repr(u32)] +#[doc = " @brief Memory space holding a C API dataset handle's data."] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsDatasetMemType_t { CUVS_DATASET_MEM_TYPE_HOST = 0, CUVS_DATASET_MEM_TYPE_DEVICE = 1, } +#[doc = " @brief Dataset handle representing owning storage or a non-owning view.\n\n `addr` points to C++ dataset storage or view metadata managed by the C API. `mem_type`\n identifies the memory space, `layout` identifies the data layout, and `is_owning` indicates\n whether the handle owns its backing data."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsDataset { @@ -299,10 +372,12 @@ const _: () = { pub type cuvsDataset_t = *mut cuvsDataset; unsafe extern "C" { #[must_use] + #[doc = " @brief Create an empty owning dataset handle.\n\n The dataset storage, memory type, layout, and dtype are populated by the operation that fills\n this handle."] pub fn cuvsDatasetCreate(dataset: *mut cuvsDataset_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Create an owning padded dataset in the requested memory space.\n\n The source tensor may reside in host- or device-accessible memory. Its contents are copied into\n newly allocated padded storage in `target_mem_type`.\n\n @param[in] res cuVS resources\n @param[in] dataset source tensor\n @param[in] target_mem_type memory space in which to allocate the padded dataset\n @param[out] padded_dataset newly allocated owning padded dataset"] pub fn cuvsDatasetMakePadded( res: cuvsResources_t, dataset: *mut DLManagedTensor, @@ -312,6 +387,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Create a non-owning padded dataset view from a host- or device-resident tensor.\n\n Memory residency is inferred from the tensor."] pub fn cuvsDatasetMakePaddedView( res: cuvsResources_t, dataset: *mut DLManagedTensor, @@ -320,6 +396,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Create a non-owning standard dataset view from a host- or device-resident tensor.\n\n Memory residency is inferred from the tensor."] pub fn cuvsDatasetMakeStandardView( res: cuvsResources_t, dataset: *mut DLManagedTensor, @@ -328,10 +405,12 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Destroy a dataset handle created by a `cuvsDatasetMake*` function."] pub fn cuvsDatasetDestroy(dataset: cuvsDataset_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the memory residency of a dataset handle."] pub fn cuvsDatasetGetMemType( dataset: cuvsDataset_t, mem_type: *mut cuvsDatasetMemType_t, @@ -339,6 +418,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the layout of a dataset handle."] pub fn cuvsDatasetGetLayout( dataset: cuvsDataset_t, layout: *mut cuvsDatasetLayout_t, @@ -346,60 +426,102 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get whether a dataset handle owns its backing storage."] pub fn cuvsDatasetGetIsOwning(dataset: cuvsDataset_t, is_owning: *mut bool) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the element dtype of a dataset handle."] pub fn cuvsDatasetGetDtype(dataset: cuvsDataset_t, dtype: *mut DLDataType) -> cuvsError_t; } #[repr(u32)] +#[doc = " enum to tell how to compute distance"] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsDistanceType { + #[doc = " evaluate as dist_ij = sum(x_ik^2) + sum(y_ij)^2 - 2*sum(x_ik * y_jk)"] L2Expanded = 0, + #[doc = " same as above, but inside the epilogue, perform square root operation"] L2SqrtExpanded = 1, + #[doc = " cosine distance"] CosineExpanded = 2, + #[doc = " L1 distance"] L1 = 3, + #[doc = " evaluate as dist_ij += (x_ik - y-jk)^2"] L2Unexpanded = 4, + #[doc = " same as above, but inside the epilogue, perform square root operation"] L2SqrtUnexpanded = 5, + #[doc = " basic inner product"] InnerProduct = 6, + #[doc = " Chebyshev (Linf) distance"] Linf = 7, + #[doc = " Canberra distance"] Canberra = 8, + #[doc = " Generalized Minkowski distance"] LpUnexpanded = 9, + #[doc = " Correlation distance"] CorrelationExpanded = 10, + #[doc = " Jaccard distance"] JaccardExpanded = 11, + #[doc = " Hellinger distance"] HellingerExpanded = 12, + #[doc = " Haversine distance"] Haversine = 13, + #[doc = " Bray-Curtis distance"] BrayCurtis = 14, + #[doc = " Jensen-Shannon distance"] JensenShannon = 15, + #[doc = " Hamming distance"] HammingUnexpanded = 16, + #[doc = " KLDivergence"] KLDivergence = 17, + #[doc = " RusselRao"] RusselRaoExpanded = 18, + #[doc = " Dice-Sorensen distance"] DiceExpanded = 19, + #[doc = " Bitstring Hamming distance"] BitwiseHamming = 20, + #[doc = " Precomputed (special value)"] Precomputed = 100, } #[repr(u32)] +#[doc = " @defgroup kmeans_c_params k-means hyperparameters\n @{"] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsKMeansInitMethod { + #[doc = " Sample the centroids using the kmeans++ strategy"] KMeansPlusPlus = 0, + #[doc = " Sample the centroids uniformly at random"] Random = 1, + #[doc = " User provides the array of initial centroids"] Array = 2, } +#[doc = " @brief Hyper-parameters for the kmeans algorithm"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsKMeansParams { pub metric: cuvsDistanceType, + #[doc = " The number of clusters to form as well as the number of centroids to generate (default:8)."] pub n_clusters: ::std::os::raw::c_int, + #[doc = " Method for initialization, defaults to k-means++:\n - cuvsKMeansInitMethod::KMeansPlusPlus (k-means++): Use scalable k-means++ algorithm\n to select the initial cluster centers.\n - cuvsKMeansInitMethod::Random (random): Choose 'n_clusters' observations (rows) at\n random from the input data for the initial centroids.\n - cuvsKMeansInitMethod::Array (ndarray): Use 'centroids' as initial cluster centers."] pub init: cuvsKMeansInitMethod, + #[doc = " Maximum number of iterations of the k-means algorithm for a single run."] pub max_iter: ::std::os::raw::c_int, + #[doc = " Relative tolerance with regards to inertia to declare convergence."] pub tol: f64, + #[doc = " Number of instance k-means algorithm will be run with different seeds."] pub n_init: ::std::os::raw::c_int, + #[doc = " Oversampling factor for use in the k-means|| algorithm"] pub oversampling_factor: f64, + #[doc = " batch_samples and batch_centroids are used to tile 1NN computation which is\n useful to optimize/control the memory footprint\n Default tile is [batch_samples x n_clusters] i.e. when batch_centroids is 0\n then don't tile the centroids"] pub batch_samples: ::std::os::raw::c_int, + #[doc = " if 0 then batch_centroids = n_clusters"] pub batch_centroids: ::std::os::raw::c_int, + #[doc = " Whether to use hierarchical (balanced) kmeans or not"] pub hierarchical: bool, + #[doc = " For hierarchical k-means , defines the number of training iterations"] pub hierarchical_n_iters: ::std::os::raw::c_int, + #[doc = " Number of samples to process per GPU batch for the batched (host-data) API.\n When set to 0, defaults to n_samples (process all at once)."] pub device_buffer_samples: i64, + #[doc = " Number of samples to draw for KMeansPlusPlus initialization.\n When set to 0, uses heuristic min(3 * n_clusters, n_samples) for host data,\n or n_samples for device data."] pub init_size: i64, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -436,13 +558,16 @@ const _: () = { pub type cuvsKMeansParams_t = *mut cuvsKMeansParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate KMeans params, and populate with default values\n\n @param[in] params cuvsKMeansParams_t to allocate\n @return cuvsError_t"] pub fn cuvsKMeansParamsCreate(params: *mut cuvsKMeansParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate KMeans params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsKMeansParamsDestroy(params: cuvsKMeansParams_t) -> cuvsError_t; } #[repr(u32)] +#[doc = " @brief Type of k-means algorithm."] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsKMeansType { CUVS_KMEANS_TYPE_KMEANS = 0, @@ -450,6 +575,7 @@ pub enum cuvsKMeansType { } unsafe extern "C" { #[must_use] + #[doc = " @brief Find clusters with k-means algorithm.\n\n Initial centroids are chosen with k-means++ algorithm. Empty\n clusters are reinitialized by choosing new centroids with\n k-means++ algorithm.\n\n X may reside on either host (CPU) or device (GPU) memory.\n When X is on the host the data is buffered to the GPU in\n batches controlled by params->device_buffer_samples.\n\n @param[in] res opaque C handle\n @param[in] params Parameters for KMeans model.\n @param[in] X Training instances to cluster. The data must\n be in row-major format. May be on host or\n device memory.\n [dim = n_samples x n_features]\n @param[in] sample_weight Optional weights for each observation in X.\n Must be on the same memory space as X.\n [len = n_samples]\n @param[inout] centroids [in] When init is InitMethod::Array, use\n centroids as the initial cluster centers.\n [out] The generated centroids from the\n kmeans algorithm are stored at the address\n pointed by 'centroids'. Must be on device.\n [dim = n_clusters x n_features]\n @param[out] inertia Sum of squared distances of samples to their\n closest cluster center.\n @param[out] n_iter Number of iterations run."] pub fn cuvsKMeansFit( res: cuvsResources_t, params: cuvsKMeansParams_t, @@ -462,6 +588,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Predict the closest cluster each sample in X belongs to.\n\n @param[in] res opaque C handle\n @param[in] params Parameters for KMeans model.\n @param[in] X New data to predict.\n [dim = n_samples x n_features]\n @param[in] sample_weight Optional weights for each observation in X.\n [len = n_samples]\n @param[in] centroids Cluster centroids. The data must be in\n row-major format.\n [dim = n_clusters x n_features]\n @param[in] normalize_weight True if the weights should be normalized\n @param[out] labels Index of the cluster each sample in X\n belongs to.\n [len = n_samples]\n @param[out] inertia Sum of squared distances of samples to\n their closest cluster center."] pub fn cuvsKMeansPredict( res: cuvsResources_t, params: cuvsKMeansParams_t, @@ -475,6 +602,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Compute cluster cost\n\n @param[in] res opaque C handle\n @param[in] X Training instances to cluster. The data must\n be in row-major format.\n [dim = n_samples x n_features]\n @param[in] centroids Cluster centroids. The data must be in\n row-major format.\n [dim = n_clusters x n_features]\n @param[out] cost Resulting cluster cost\n"] pub fn cuvsKMeansClusterCost( res: cuvsResources_t, X: *mut DLManagedTensor, @@ -484,6 +612,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Compute pairwise distances for two matrices\n\n\n Usage example:\n @code{.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Assume a populated `DLManagedTensor` type here\n DLManagedTensor x;\n DLManagedTensor y;\n DLManagedTensor dist;\n\n cuvsPairwiseDistance(res, &x, &y, &dist, L2SqrtUnexpanded, 2.0);\n @endcode\n\n @param[in] res cuvs resources object for managing expensive resources\n @param[in] x first set of points (size n*k)\n @param[in] y second set of points (size m*k)\n @param[out] dist output distance matrix (size n*m)\n @param[in] metric distance to evaluate\n @param[in] metric_arg metric argument (used for Minkowski distance)"] pub fn cuvsPairwiseDistance( res: cuvsResources_t, x: *mut DLManagedTensor, @@ -494,32 +623,48 @@ unsafe extern "C" { ) -> cuvsError_t; } #[repr(u32)] +#[doc = " @defgroup ivf_pq_c_index_params IVF-PQ index build parameters\n @{\n/\n/**\n @brief A type for specifying how PQ codebooks are created\n"] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsIvfPqCodebookGen { CUVS_IVF_PQ_CODEBOOK_GEN_PER_SUBSPACE = 0, CUVS_IVF_PQ_CODEBOOK_GEN_PER_CLUSTER = 1, } #[repr(u32)] +#[doc = " @brief A type for specifying the memory layout of IVF-PQ list data\n"] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsIvfPqListLayout { CUVS_IVF_PQ_LIST_LAYOUT_FLAT = 0, CUVS_IVF_PQ_LIST_LAYOUT_INTERLEAVED = 1, } +#[doc = " @brief Supplemental parameters to build IVF-PQ Index\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsIvfPqIndexParams { + #[doc = " Distance type."] pub metric: cuvsDistanceType, + #[doc = " The argument used by some distance metrics."] pub metric_arg: f32, + #[doc = " Whether to add the dataset content to the index, i.e.:\n\n - `true` means the index is filled with the dataset vectors and ready to search after calling\n `build`.\n - `false` means `build` only trains the underlying model (e.g. quantizer or clustering), but\n the index is left empty; you'd need to call `extend` on the index afterwards to populate it."] pub add_data_on_build: bool, + #[doc = " The number of inverted lists (clusters)\n\n Hint: the number of vectors per cluster (`n_rows/n_lists`) should be approximately 1,000 to\n 10,000."] pub n_lists: u32, + #[doc = " The number of iterations searching for kmeans centers (index building)."] pub kmeans_n_iters: u32, + #[doc = " The fraction of data to use during iterative kmeans building."] pub kmeans_trainset_fraction: f64, + #[doc = " The bit length of the vector element after compression by PQ.\n\n Possible values: [4, 5, 6, 7, 8].\n\n Hint: the smaller the 'pq_bits', the smaller the index size and the better the search\n performance, but the lower the recall."] pub pq_bits: u32, + #[doc = " The dimensionality of the vector after compression by PQ. When zero, an optimal value is\n selected using a heuristic.\n\n NB: `pq_dim * pq_bits` must be a multiple of 8.\n\n Hint: a smaller 'pq_dim' results in a smaller index size and better search performance, but\n lower recall. If 'pq_bits' is 8, 'pq_dim' can be set to any number, but multiple of 8 are\n desirable for good performance. If 'pq_bits' is not 8, 'pq_dim' should be a multiple of 8.\n For good performance, it is desirable that 'pq_dim' is a multiple of 32. Ideally, 'pq_dim'\n should be also a divisor of the dataset dim."] pub pq_dim: u32, + #[doc = " How PQ codebooks are created."] pub codebook_kind: cuvsIvfPqCodebookGen, + #[doc = " Apply a random rotation matrix on the input data and queries even if `dim % pq_dim == 0`.\n\n Note: if `dim` is not multiple of `pq_dim`, a random rotation is always applied to the input\n data and queries to transform the working space from `dim` to `rot_dim`, which may be slightly\n larger than the original space and and is a multiple of `pq_dim` (`rot_dim % pq_dim == 0`).\n However, this transform is not necessary when `dim` is multiple of `pq_dim`\n (`dim == rot_dim`, hence no need in adding \"extra\" data columns / features).\n\n By default, if `dim == rot_dim`, the rotation transform is initialized with the identity\n matrix. When `force_random_rotation == true`, a random orthogonal transform matrix is generated\n regardless of the values of `dim` and `pq_dim`."] pub force_random_rotation: bool, + #[doc = " By default, the algorithm allocates more space than necessary for individual clusters\n (`list_data`). This allows to amortize the cost of memory allocation and reduce the number of\n data copies during repeated calls to `extend` (extending the database).\n\n The alternative is the conservative allocation behavior; when enabled, the algorithm always\n allocates the minimum amount of memory required to store the given number of records. Set this\n flag to `true` if you prefer to use as little GPU memory for the database as possible."] pub conservative_memory_allocation: bool, + #[doc = " The max number of data points to use per PQ code during PQ codebook training. Using more data\n points per PQ code may increase the quality of PQ codebook but may also increase the build\n time. The parameter is applied to both PQ codebook generation methods, i.e., PER_SUBSPACE and\n PER_CLUSTER. In both cases, we will use `pq_book_size * max_train_points_per_pq_code` training\n points to train each codebook."] pub max_train_points_per_pq_code: u32, + #[doc = " Memory layout of the IVF-PQ list data.\n\n - CUVS_IVF_PQ_LIST_LAYOUT_FLAT: Codes are stored contiguously, one vector's codes after another.\n - CUVS_IVF_PQ_LIST_LAYOUT_INTERLEAVED: Codes are interleaved for optimized search performance.\n This is the default and recommended for search workloads."] pub codes_layout: cuvsIvfPqListLayout, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -556,19 +701,28 @@ const _: () = { pub type cuvsIvfPqIndexParams_t = *mut cuvsIvfPqIndexParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate IVF-PQ Index params, and populate with default values\n\n @param[in] index_params cuvsIvfPqIndexParams_t to allocate\n @return cuvsError_t"] pub fn cuvsIvfPqIndexParamsCreate(index_params: *mut cuvsIvfPqIndexParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate IVF-PQ Index params\n\n @param[in] index_params\n @return cuvsError_t"] pub fn cuvsIvfPqIndexParamsDestroy(index_params: cuvsIvfPqIndexParams_t) -> cuvsError_t; } +#[doc = " @defgroup ivf_pq_c_search_params IVF-PQ index search parameters\n @{\n/\n/**\n @brief Supplemental parameters to search IVF-PQ index\n"] #[repr(C)] pub struct cuvsIvfPqSearchParams { + #[doc = " The number of clusters to search."] pub n_probes: u32, + #[doc = " Data type of look up table to be created dynamically at search time.\n\n Possible values: [CUDA_R_32F, CUDA_R_16F, CUDA_R_8U]\n\n The use of low-precision types reduces the amount of shared memory required at search time, so\n fast shared memory kernels can be used even for datasets with large dimansionality. Note that\n the recall is slightly degraded when low-precision type is selected."] pub lut_dtype: cudaDataType_t, + #[doc = " Storage data type for distance/similarity computed at search time.\n\n Possible values: [CUDA_R_16F, CUDA_R_32F]\n\n If the performance limiter at search time is device memory access, selecting FP16 will improve\n performance slightly."] pub internal_distance_dtype: cudaDataType_t, + #[doc = " The data type to use as the GEMM element type when searching the clusters to probe.\n\n Possible values: [CUDA_R_8I, CUDA_R_16F, CUDA_R_32F].\n\n - Legacy default: CUDA_R_32F (float)\n - Recommended for performance: CUDA_R_16F (half)\n - Experimental/low-precision: CUDA_R_8I (int8_t)\n (WARNING: int8_t variant degrades recall unless data is normalized and low-dimensional)"] pub coarse_search_dtype: cudaDataType_t, + #[doc = " Set the internal batch size to improve GPU utilization at the cost of larger memory footprint."] pub max_internal_batch_size: u32, + #[doc = " Preferred fraction of SM's unified memory / L1 cache to be used as shared memory.\n\n Possible values: [0.0 - 1.0] as a fraction of the `sharedMemPerMultiprocessor`.\n\n One wants to increase the carveout to make sure a good GPU occupancy for the main search\n kernel, but not to keep it too high to leave some memory to be used as L1 cache. Note, this\n value is interpreted only as a hint. Moreover, a GPU usually allows only a fixed set of cache\n configurations, so the provided value is rounded up to the nearest configuration. Refer to the\n NVIDIA tuning guide for the target GPU architecture.\n\n Note, this is a low-level tuning parameter that can have drastic negative effects on the search\n performance if tweaked incorrectly."] pub preferred_shmem_carveout: f64, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -592,12 +746,15 @@ const _: () = { pub type cuvsIvfPqSearchParams_t = *mut cuvsIvfPqSearchParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate IVF-PQ search params, and populate with default values\n\n @param[in] params cuvsIvfPqSearchParams_t to allocate\n @return cuvsError_t"] pub fn cuvsIvfPqSearchParamsCreate(params: *mut cuvsIvfPqSearchParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate IVF-PQ search params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsIvfPqSearchParamsDestroy(params: cuvsIvfPqSearchParams_t) -> cuvsError_t; } +#[doc = " @defgroup ivf_pq_c_index IVF-PQ index\n @{\n/\n/**\n @brief Struct to hold address of cuvs::neighbors::ivf_pq::index and its active trained dtype\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsIvfPqIndex { @@ -616,38 +773,47 @@ const _: () = { pub type cuvsIvfPqIndex_t = *mut cuvsIvfPqIndex; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate IVF-PQ index\n\n @param[in] index cuvsIvfPqIndex_t to allocate\n @return cuvsError_t"] pub fn cuvsIvfPqIndexCreate(index: *mut cuvsIvfPqIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate IVF-PQ index\n\n @param[in] index cuvsIvfPqIndex_t to de-allocate"] pub fn cuvsIvfPqIndexDestroy(index: cuvsIvfPqIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " Get the number of clusters/inverted lists"] pub fn cuvsIvfPqIndexGetNLists(index: cuvsIvfPqIndex_t, n_lists: *mut i64) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " Get the dimensionality"] pub fn cuvsIvfPqIndexGetDim(index: cuvsIvfPqIndex_t, dim: *mut i64) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " Get the size of the index"] pub fn cuvsIvfPqIndexGetSize(index: cuvsIvfPqIndex_t, size: *mut i64) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " Get the dimensionality of an encoded vector after compression by PQ."] pub fn cuvsIvfPqIndexGetPqDim(index: cuvsIvfPqIndex_t, pq_dim: *mut i64) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " Get the bit length of an encoded vector element after compression by PQ."] pub fn cuvsIvfPqIndexGetPqBits(index: cuvsIvfPqIndex_t, pq_bits: *mut i64) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " Get the Dimensionality of a subspace, i.e. the number of vector\n components mapped to a subspace"] pub fn cuvsIvfPqIndexGetPqLen(index: cuvsIvfPqIndex_t, pq_len: *mut i64) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the cluster centers corresponding to the lists in the original space\n\n @param[in] index cuvsIvfPqIndex_t Built Ivf-Pq index\n @param[out] centers Output tensor that will be populated with a non-owning view of the data\n @return cuvsError_t"] pub fn cuvsIvfPqIndexGetCenters( index: cuvsIvfPqIndex_t, centers: *mut DLManagedTensor, @@ -655,6 +821,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the padded cluster centers [n_lists, dim_ext]\n where dim_ext = round_up(dim + 1, 8)\n\n This returns the full padded centers as a contiguous array, suitable for\n use with cuvsIvfPqBuildPrecomputed.\n\n @param[in] index cuvsIvfPqIndex_t Built Ivf-Pq index\n @param[out] centers Output tensor that will be populated with a non-owning view of the data\n @return cuvsError_t"] pub fn cuvsIvfPqIndexGetCentersPadded( index: cuvsIvfPqIndex_t, centers: *mut DLManagedTensor, @@ -662,6 +829,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the PQ cluster centers\n\n - CUVS_IVF_PQ_CODEBOOK_GEN_PER_SUBSPACE: [pq_dim , pq_len, pq_book_size]\n - CUVS_IVF_PQ_CODEBOOK_GEN_PER_CLUSTER: [n_lists, pq_len, pq_book_size]\n\n @param[in] index cuvsIvfPqIndex_t Built Ivf-Pq index\n @param[out] pq_centers Output tensor that will be populated with a non-owning view of the data\n @return cuvsError_t"] pub fn cuvsIvfPqIndexGetPqCenters( index: cuvsIvfPqIndex_t, pq_centers: *mut DLManagedTensor, @@ -669,6 +837,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the rotated cluster centers [n_lists, rot_dim]\n where rot_dim = pq_len * pq_dim\n\n @param[in] index cuvsIvfPqIndex_t Built Ivf-Pq index\n @param[out] centers_rot Output tensor that will be populated with a non-owning view of the data\n @return cuvsError_t"] pub fn cuvsIvfPqIndexGetCentersRot( index: cuvsIvfPqIndex_t, centers_rot: *mut DLManagedTensor, @@ -676,6 +845,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the rotation matrix [rot_dim, dim]\n Transform matrix (original space -> rotated padded space)\n\n @param[in] index cuvsIvfPqIndex_t Built Ivf-Pq index\n @param[out] rotation_matrix Output tensor that will be populated with a non-owning view of the\n data\n @return cuvsError_t"] pub fn cuvsIvfPqIndexGetRotationMatrix( index: cuvsIvfPqIndex_t, rotation_matrix: *mut DLManagedTensor, @@ -683,6 +853,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the sizes of each list\n\n @param[in] index cuvsIvfPqIndex_t Built Ivf-Pq index\n @param[out] list_sizes Output tensor that will be populated with a non-owning view of the data\n @return cuvsError_t"] pub fn cuvsIvfPqIndexGetListSizes( index: cuvsIvfPqIndex_t, list_sizes: *mut DLManagedTensor, @@ -690,6 +861,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Unpack `n_rows` consecutive PQ encoded vectors of a single list (cluster) in the\n compressed index starting at given `offset`, not expanded to one code per byte. Each code in the\n output buffer occupies ceildiv(index.pq_dim() * index.pq_bits(), 8) bytes.\n\n @param[in] res raft resource\n @param[in] index cuvsIvfPqIndex_t Built Ivf-Pq index\n @param[out] out_codes\n the destination buffer [n_rows, ceildiv(index.pq_dim() * index.pq_bits(), 8)].\n The length `n_rows` defines how many records to unpack,\n offset + n_rows must be smaller than or equal to the list size.\n This DLManagedTensor must already point to allocated device memory\n @param[in] label\n The id of the list (cluster) to decode.\n @param[in] offset\n How many records in the list to skip."] pub fn cuvsIvfPqIndexUnpackContiguousListData( res: cuvsResources_t, index: cuvsIvfPqIndex_t, @@ -700,6 +872,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the indices of each vector in a ivf-pq list\n\n @param[in] index cuvsIvfPqIndex_t Built Ivf-Pq index\n @param[in] label\n The id of the list (cluster) to decode.\n @param[out] out_labels\n output tensor that will be populated with a non-owning view of the data\n @return cuvsError_t"] pub fn cuvsIvfPqIndexGetListIndices( index: cuvsIvfPqIndex_t, label: u32, @@ -708,6 +881,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup ivf_pq_c_index_build IVF-PQ index build\n @{\n/\n/**\n @brief Build a IVF-PQ index with a `DLManagedTensor` which has underlying\n `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`,\n or `kDLCPU`. Also, acceptable underlying types are:\n 1. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n 2. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 16`\n 3. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8`\n 4. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8`\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Assume a populated `DLManagedTensor` type here\n DLManagedTensor dataset;\n\n // Create default index params\n cuvsIvfPqIndexParams_t index_params;\n cuvsError_t params_create_status = cuvsIvfPqIndexParamsCreate(&index_params);\n\n // Create IVF-PQ index\n cuvsIvfPqIndex_t index;\n cuvsError_t index_create_status = cuvsIvfPqIndexCreate(&index);\n\n // Build the IVF-PQ Index\n cuvsError_t build_status = cuvsIvfPqBuild(res, index_params, &dataset, index);\n\n // de-allocate `index_params`, `index` and `res`\n cuvsError_t params_destroy_status = cuvsIvfPqIndexParamsDestroy(index_params);\n cuvsError_t index_destroy_status = cuvsIvfPqIndexDestroy(index);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params cuvsIvfPqIndexParams_t used to build IVF-PQ index\n @param[in] dataset DLManagedTensor* training dataset\n @param[out] index cuvsIvfPqIndex_t Newly built IVF-PQ index\n @return cuvsError_t"] pub fn cuvsIvfPqBuild( res: cuvsResources_t, params: cuvsIvfPqIndexParams_t, @@ -717,6 +891,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Build a view-type IVF-PQ index from device memory precomputed centroids and codebook.\n\n This function creates a non-owning index that stores a reference to the provided device data.\n All parameters must be provided with correct extents. The caller is responsible for ensuring\n the lifetime of the input data exceeds the lifetime of the returned index.\n\n The index_params must be consistent with the provided matrices. Specifically:\n - index_params.codebook_kind determines the expected shape of pq_centers\n - index_params.metric will be stored in the index\n - index_params.conservative_memory_allocation will be stored in the index\n The function will verify consistency between index_params, dim, and the matrix extents.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params cuvsIvfPqIndexParams_t used to configure the index (must be consistent with\n matrices)\n @param[in] dim dimensionality of the input data\n @param[in] pq_centers PQ codebook on device memory with required shape:\n - codebook_kind CUVS_IVF_PQ_CODEBOOK_GEN_PER_SUBSPACE: [pq_dim, pq_len, pq_book_size]\n - codebook_kind CUVS_IVF_PQ_CODEBOOK_GEN_PER_CLUSTER: [n_lists, pq_len, pq_book_size]\n @param[in] centers Cluster centers in the original space [n_lists, dim_ext]\n where dim_ext = round_up(dim + 1, 8)\n @param[in] centers_rot Rotated cluster centers [n_lists, rot_dim]\n where rot_dim = pq_len * pq_dim\n @param[in] rotation_matrix Transform matrix (original space -> rotated padded space) [rot_dim,\n dim]\n @param[out] index cuvsIvfPqIndex_t Newly built view-type IVF-PQ index\n @return cuvsError_t"] pub fn cuvsIvfPqBuildPrecomputed( res: cuvsResources_t, params: cuvsIvfPqIndexParams_t, @@ -730,6 +905,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup ivf_pq_c_index_search IVF-PQ index search\n @{\n/\n/**\n @brief Search a IVF-PQ index with a `DLManagedTensor` which has underlying\n `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`.\n It is also important to note that the IVF-PQ Index must have been built\n with the same type of `queries`, such that `index.dtype.code ==\n queries.dl_tensor.dtype.code` Types for input are:\n 1. `queries`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n or `kDLDataType.bits = 16`\n 2. `neighbors`: `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 32`\n 3. `distances`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Assume a populated `DLManagedTensor` type here\n DLManagedTensor dataset;\n DLManagedTensor queries;\n DLManagedTensor neighbors;\n\n // Create default search params\n cuvsIvfPqSearchParams_t search_params;\n cuvsError_t params_create_status = cuvsIvfPqSearchParamsCreate(&search_params);\n\n // Search the `index` built using `cuvsIvfPqBuild`\n cuvsError_t search_status = cuvsIvfPqSearch(res, search_params, index, &queries, &neighbors,\n &distances);\n\n // de-allocate `search_params` and `res`\n cuvsError_t params_destroy_status = cuvsIvfPqSearchParamsDestroy(search_params);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] search_params cuvsIvfPqSearchParams_t used to search IVF-PQ index\n @param[in] index cuvsIvfPqIndex which has been returned by `cuvsIvfPqBuild`\n @param[in] queries DLManagedTensor* queries dataset to search\n @param[out] neighbors DLManagedTensor* output `k` neighbors for queries\n @param[out] distances DLManagedTensor* output `k` distances for queries"] pub fn cuvsIvfPqSearch( res: cuvsResources_t, search_params: cuvsIvfPqSearchParams_t, @@ -741,6 +917,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup ivf_pq_c_index_serialize IVF-PQ C-API serialize functions\n @{\n/\n/**\n Save the index to file.\n\n Experimental, both the API and the serialization format are subject to change.\n\n @code{.cpp}\n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // create an index with `cuvsIvfPqBuild`\n cuvsIvfPqSerialize(res, \"/path/to/index\", index, true);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename the file name for saving the index\n @param[in] index IVF-PQ index"] pub fn cuvsIvfPqSerialize( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -749,6 +926,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " Load index from file.\n\n Experimental, both the API and the serialization format are subject to change.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename the name of the file that stores the index\n @param[out] index IVF-PQ index loaded disk"] pub fn cuvsIvfPqDeserialize( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -757,6 +935,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup ivf_pq_c_index_extend IVF-PQ index extend\n @{\n/\n/**\n @brief Extend the index with the new data.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] new_vectors DLManagedTensor* the new vectors to add to the index\n @param[in] new_indices DLManagedTensor* vector of new indices for the new vectors\n @param[inout] index IVF-PQ index to be extended\n @return cuvsError_t"] pub fn cuvsIvfPqExtend( res: cuvsResources_t, new_vectors: *mut DLManagedTensor, @@ -766,6 +945,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup ivf_pq_c_index_transform IVF-PQ index transform\n @{\n/\n/**\n @brief Transform the input data by applying pq-encoding\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] index IVF-PQ index\n @param[in] input_dataset DLManagedTensor* vectors to transform\n @param[out] output_labels DLManagedTensor* Vector of cluster labels for each vector in the input\n @param[out] output_dataset DLManagedTensor* input vectors after pq-encoding\n @return cuvsError_t"] pub fn cuvsIvfPqTransform( res: cuvsResources_t, index: cuvsIvfPqIndex_t, @@ -775,12 +955,14 @@ unsafe extern "C" { ) -> cuvsError_t; } #[repr(u32)] +#[doc = " @brief Dtype to use for distance computation\n - `NND_DIST_COMP_AUTO`: Automatically determine the best dtype for distance computation based on the dataset dimensions.\n - `NND_DIST_COMP_FP32`: Use fp32 distance computation for better precision at the cost of performance and memory usage.\n - `NND_DIST_COMP_FP16`: Use fp16 distance computation."] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsNNDescentDistCompDtype { NND_DIST_COMP_AUTO = 0, NND_DIST_COMP_FP32 = 1, NND_DIST_COMP_FP16 = 2, } +#[doc = " @defgroup nn_descent_c_index_params The nn-descent algorithm parameters.\n @{\n/\n/**\n @brief Parameters used to build an nn-descent index\n\n `metric`: The distance metric to use\n `metric_arg`: The argument used by distance metrics like Minkowskidistance\n `graph_degree`: For an input dataset of dimensions (N, D),\n determines the final dimensions of the all-neighbors knn graph\n which turns out to be of dimensions (N, graph_degree)\n `intermediate_graph_degree`: Internally, nn-descent builds an\n all-neighbors knn graph of dimensions (N, intermediate_graph_degree)\n before selecting the final `graph_degree` neighbors. It's recommended\n that `intermediate_graph_degree` >= 1.5 * graph_degree\n `max_iterations`: The number of iterations that nn-descent will refine\n the graph for. More iterations produce a better quality graph at cost of performance\n `termination_threshold`: The delta at which nn-descent will terminate its iterations\n `return_distances`: Boolean to decide whether to return distances array\n `dist_comp_dtype`: dtype to use for distance computation. Defaults to `NND_DIST_COMP_AUTO` which automatically determines the best dtype for distance computation based on the dataset dimensions. Use `NND_DIST_COMP_FP32` for better precision at the cost of performance and memory usage. This option is only valid when data type is fp32. Use `NND_DIST_COMP_FP16` for better performance and memory usage at the cost of precision."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsNNDescentIndexParams { @@ -819,15 +1001,18 @@ const _: () = { pub type cuvsNNDescentIndexParams_t = *mut cuvsNNDescentIndexParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate NN-Descent Index params, and populate with default values\n\n @param[in] index_params cuvsNNDescentIndexParams_t to allocate\n @return cuvsError_t"] pub fn cuvsNNDescentIndexParamsCreate( index_params: *mut cuvsNNDescentIndexParams_t, ) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate NN-Descent Index params\n\n @param[in] index_params\n @return cuvsError_t"] pub fn cuvsNNDescentIndexParamsDestroy(index_params: cuvsNNDescentIndexParams_t) -> cuvsError_t; } +#[doc = " @defgroup nn_descent_c_index NN-Descent index\n @{\n/\n/**\n @brief Struct to hold address of cuvs::neighbors::nn_descent::index and its active trained dtype\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsNNDescentIndex { @@ -846,14 +1031,17 @@ const _: () = { pub type cuvsNNDescentIndex_t = *mut cuvsNNDescentIndex; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate NN-Descent index\n\n @param[in] index cuvsNNDescentIndex_t to allocate\n @return cuvsError_t"] pub fn cuvsNNDescentIndexCreate(index: *mut cuvsNNDescentIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate NN-Descent index\n\n @param[in] index cuvsNNDescentIndex_t to de-allocate"] pub fn cuvsNNDescentIndexDestroy(index: cuvsNNDescentIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @defgroup nn_descent_c_index_build NN-Descent index build\n @{\n/\n/**\n @brief Build a NN-Descent index with a `DLManagedTensor` which has underlying\n `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`,\n or `kDLCPU`. Also, acceptable underlying types are:\n 1. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n 2. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 16`\n 3. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8`\n 4. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8`\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Assume a populated `DLManagedTensor` type here\n DLManagedTensor dataset;\n\n // Create default index params\n cuvsNNDescentIndexParams_t index_params;\n cuvsError_t params_create_status = cuvsNNDescentIndexParamsCreate(&index_params);\n\n // Create NN-Descent index\n cuvsNNDescentIndex_t index;\n cuvsError_t index_create_status = cuvsNNDescentIndexCreate(&index);\n\n // Build the NN-Descent Index\n cuvsError_t build_status = cuvsNNDescentBuild(res, index_params, &dataset, index);\n\n // de-allocate `index_params`, `index` and `res`\n cuvsError_t params_destroy_status = cuvsNNDescentIndexParamsDestroy(index_params);\n cuvsError_t index_destroy_status = cuvsNNDescentIndexDestroy(index);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] index_params cuvsNNDescentIndexParams_t used to build NN-Descent index\n @param[in] dataset DLManagedTensor* training dataset on host or device memory\n @param[inout] graph Optional preallocated graph on host memory to store output\n @param[out] index cuvsNNDescentIndex_t Newly built NN-Descent index\n @return cuvsError_t"] pub fn cuvsNNDescentBuild( res: cuvsResources_t, index_params: cuvsNNDescentIndexParams_t, @@ -864,6 +1052,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the KNN graph from a built NN-Descent index\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] index cuvsNNDescentIndex_t Built NN-Descent index\n @param[out] graph Preallocated graph on host memory to store output\n @return cuvsError_t"] pub fn cuvsNNDescentIndexGetGraph( res: cuvsResources_t, index: cuvsNNDescentIndex_t, @@ -872,6 +1061,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the distances from a build NN_Descent index\n\n This requires that the `return_distances` parameter was set when building the\n graph\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] index cuvsNNDescentIndex_t Built NN-Descent index\n @param[out] distances Preallocated memory to store the output distances tensor\n @return cuvsError_t"] pub fn cuvsNNDescentIndexGetDistances( res: cuvsResources_t, index: cuvsNNDescentIndex_t, @@ -879,20 +1069,31 @@ unsafe extern "C" { ) -> cuvsError_t; } #[repr(u32)] +#[doc = " @brief Graph build algorithm selection."] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsAllNeighborsAlgo { + #[doc = "< Use Brute Force for local kNN subgraphs"] CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE = 0, + #[doc = "< Use IVF-PQ for local kNN subgraphs (host dataset only)"] CUVS_ALL_NEIGHBORS_ALGO_IVF_PQ = 1, + #[doc = "< Use NN-Descent for local kNN subgraphs"] CUVS_ALL_NEIGHBORS_ALGO_NN_DESCENT = 2, } +#[doc = " @brief Parameters controlling SNMG all-neighbors build."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsAllNeighborsIndexParams { + #[doc = "< Local kNN graph build algorithm"] pub algo: cuvsAllNeighborsAlgo, + #[doc = "< Number of clusters each point is assigned to (must be < n_clusters)"] pub overlap_factor: usize, + #[doc = "< Number of clusters/batches to partition the dataset into (> overlap_factor)"] pub n_clusters: usize, + #[doc = "< Distance metric used for graph construction"] pub metric: cuvsDistanceType, + #[doc = "< Parameters for IVF-PQ algorithm (when algo ==\n< CUVS_ALL_NEIGHBORS_ALGO_IVF_PQ)"] pub ivf_pq_params: cuvsIvfPqIndexParams_t, + #[doc = "< Parameters for NN-Descent algorithm (when algo\n< == CUVS_ALL_NEIGHBORS_ALGO_NN_DESCENT)"] pub nn_descent_params: cuvsNNDescentIndexParams_t, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -917,18 +1118,21 @@ const _: () = { pub type cuvsAllNeighborsIndexParams_t = *mut cuvsAllNeighborsIndexParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Create a default all-neighbors index parameters struct.\n\n @param[out] index_params Pointer to allocated index_params struct\n\n @return cuvsError_t"] pub fn cuvsAllNeighborsIndexParamsCreate( index_params: *mut cuvsAllNeighborsIndexParams_t, ) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Destroy an all-neighbors index parameters struct.\n\n @param[in] index_params Index parameters struct to destroy\n\n @return cuvsError_t"] pub fn cuvsAllNeighborsIndexParamsDestroy( index_params: cuvsAllNeighborsIndexParams_t, ) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Build an all-neighbors k-NN graph automatically detecting host vs device dataset.\n\n @param[in] res Can be a SNMG multi-GPU resources (`cuvsResources_t`) or single-GPU\n resources\n @param[in] params Build parameters (see cuvsAllNeighborsIndexParams)\n @param[in] dataset 2D tensor [num_rows x dim] on host or device (auto-detected)\n @param[out] indices 2D tensor [num_rows x k] (int64), host or device\n @param[out] distances Optional 2D tensor [num_rows x k] (float32), host or device; can be\n NULL\n @param[out] core_distances Optional 1D tensor [num_rows] (float32), host or device; can be NULL\n @param[in] alpha Mutual-reachability scaling; used only when core_distances is provided\n\n The function automatically detects whether the dataset is host-resident or device-resident\n and calls the appropriate implementation. For host datasets, it partitions data into\n `n_clusters` clusters and assigns each row to `overlap_factor` nearest clusters. For device\n datasets, `n_clusters` must be 1 (no batching); `overlap_factor` is ignored.\n\n Output memory space: a host dataset supports host- or device-resident outputs; a device dataset\n requires device-resident outputs. All provided outputs must share the same memory space."] pub fn cuvsAllNeighborsBuild( res: cuvsResources_t, params: cuvsAllNeighborsIndexParams_t, @@ -940,12 +1144,14 @@ unsafe extern "C" { ) -> cuvsError_t; } #[repr(u32)] +#[doc = " @brief Enum to denote filter type."] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsFilterType { NO_FILTER = 0, BITSET = 1, BITMAP = 2, } +#[doc = " @brief Struct to hold address of cuvs::neighbors::prefilter and its type\n\n `addr` points to a filter object owned by the caller; the library performs no caching of the\n underlying bitset across search calls. Allocating and populating the device bitset may be more\n expensive than a single filtered search, so callers that issue repeated searches against the same\n filter (e.g. many queries over one index) should build the bitset once and reuse the same\n cuvsFilter across those calls rather than rebuild it per search. Reusing the bitset is essential\n for realizing the full throughput of filtered search."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsFilter { @@ -960,11 +1166,15 @@ const _: () = { ["Offset of field: cuvsFilter::type_"][::std::mem::offset_of!(cuvsFilter, type_) - 8usize]; }; #[repr(u32)] +#[doc = " @brief Strategy for merging indices."] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsMergeStrategy { + #[doc = "< Merge indices physically"] MERGE_STRATEGY_PHYSICAL = 0, + #[doc = "< Merge indices logically"] MERGE_STRATEGY_LOGICAL = 1, } +#[doc = " @defgroup bruteforce_c_index Bruteforce index\n @{\n/\n/**\n @brief Struct to hold address of cuvs::neighbors::brute_force::index and its active trained dtype\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsBruteForceIndex { @@ -983,14 +1193,17 @@ const _: () = { pub type cuvsBruteForceIndex_t = *mut cuvsBruteForceIndex; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate BRUTEFORCE index\n\n @param[in] index cuvsBruteForceIndex_t to allocate\n @return cuvsError_t"] pub fn cuvsBruteForceIndexCreate(index: *mut cuvsBruteForceIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate BRUTEFORCE index\n\n @param[in] index cuvsBruteForceIndex_t to de-allocate"] pub fn cuvsBruteForceIndexDestroy(index: cuvsBruteForceIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @defgroup bruteforce_c_index_build Bruteforce index build\n @{\n/\n/**\n @brief Build a BRUTEFORCE index with a `DLManagedTensor` which has underlying\n `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`,\n or `kDLCPU`. Also, acceptable underlying types are:\n 1. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n 2. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 16`\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Assume a populated `DLManagedTensor` type here\n DLManagedTensor dataset;\n\n // Create BRUTEFORCE index\n cuvsBruteForceIndex_t index;\n cuvsError_t index_create_status = cuvsBruteForceIndexCreate(&index);\n\n // Build the BRUTEFORCE Index\n cuvsError_t build_status = cuvsBruteForceBuild(res, &dataset_tensor, L2Expanded, 0.f, index);\n\n // de-allocate `index` and `res`\n cuvsError_t index_destroy_status = cuvsBruteForceIndexDestroy(index);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] dataset DLManagedTensor* training dataset\n @param[in] metric metric\n @param[in] metric_arg metric_arg\n @param[out] index cuvsBruteForceIndex_t Newly built BRUTEFORCE index\n @return cuvsError_t"] pub fn cuvsBruteForceBuild( res: cuvsResources_t, dataset: *mut DLManagedTensor, @@ -1001,6 +1214,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup bruteforce_c_index_search Bruteforce index search\n @{\n/\n/**\n @brief Search a BRUTEFORCE index with a `DLManagedTensor` which has underlying\n `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`.\n It is also important to note that the BRUTEFORCE index must have been built\n with the same type of `queries`, such that `index.dtype.code ==\n queries.dl_tensor.dtype.code` Types for input are:\n 1. `queries`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32` or\n `kDLDataType.bits = 16`\n 2. `neighbors`: `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 32`\n 3. `distances`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Assume a populated `DLManagedTensor` type here\n DLManagedTensor dataset;\n DLManagedTensor queries;\n DLManagedTensor neighbors;\n DLManagedTensor bitmap;\n\n cuvsFilter prefilter{(uintptr_t)&bitmap, BITMAP};\n\n // Search the `index` built using `cuvsBruteForceBuild`\n cuvsError_t search_status = cuvsBruteForceSearch(res, index, &queries, &neighbors, &distances,\n prefilter);\n\n // de-allocate `res`\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] index cuvsBruteForceIndex which has been returned by `cuvsBruteForceBuild`\n @param[in] queries DLManagedTensor* queries dataset to search\n @param[out] neighbors DLManagedTensor* output `k` neighbors for queries\n @param[out] distances DLManagedTensor* output `k` distances for queries\n @param[in] prefilter cuvsFilter input prefilter that can be used\nto filter queries and neighbors based on the given bitmap."] pub fn cuvsBruteForceSearch( res: cuvsResources_t, index: cuvsBruteForceIndex_t, @@ -1012,6 +1226,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup bruteforce_c_index_serialize BRUTEFORCE C-API serialize functions\n @{\n/\n/**\n Save the index to file.\n The serialization format can be subject to changes, therefore loading\n an index saved with a previous version of cuvs is not guaranteed\n to work.\n\n @code{.c}\n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // create an index with `cuvsBruteforceBuild`\n cuvsBruteForceSerialize(res, \"/path/to/index\", index);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename the file name for saving the index\n @param[in] index BRUTEFORCE index\n"] pub fn cuvsBruteForceSerialize( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -1020,6 +1235,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " Load index from file.\n The serialization format can be subject to changes, therefore loading\n an index saved with a previous version of cuvs is not guaranteed\n to work.\n\n @code{.c}\n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Deserialize an index previously built with `cuvsBruteforceBuild`\n cuvsBruteForceIndex_t index;\n cuvsBruteForceIndexCreate(&index);\n cuvsBruteForceDeserialize(res, \"/path/to/index\", index);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename the name of the file that stores the index\n @param[out] index BRUTEFORCE index loaded disk"] pub fn cuvsBruteForceDeserialize( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -1027,28 +1243,40 @@ unsafe extern "C" { ) -> cuvsError_t; } #[repr(u32)] +#[doc = " @brief Enum to denote which ANN algorithm is used to build CAGRA graph\n"] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsCagraGraphBuildAlgo { AUTO_SELECT = 0, IVF_PQ = 1, NN_DESCENT = 2, ITERATIVE_CAGRA_SEARCH = 3, + #[doc = " Experimental, use ACE (Augmented Core Extraction) to build the graph. ACE partitions the\n dataset into core and augmented partitions and builds a sub-index for each partition. This\n enables building indices for datasets too large to fit in GPU or host memory.\n See cuvsAceParams for more details about the ACE algorithm and its parameters."] ACE = 4, } #[repr(u32)] +#[doc = " @brief A strategy for selecting the graph build parameters based on similar HNSW index\n parameters.\n\n Define how cuvsCagraIndexParamsFromHnswParams should construct a graph to construct a graph\n that is to be converted to (used by) a CPU HNSW index."] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsCagraHnswHeuristicType { + #[doc = " Create a graph that is very similar to an HNSW graph in\n terms of the number of nodes and search performance. Since HNSW produces a variable-degree\n graph (2M being the max graph degree) and CAGRA produces a fixed-degree graph, there's always a\n difference in the performance of the two.\n\n This function attempts to produce such a graph that the QPS and recall of the two graphs being\n searched by HNSW are close for any search parameter combination. The CAGRA-produced graph tends\n to have a \"longer tail\" on the low recall side (that is being slightly faster and less\n precise).\n"] CUVS_CAGRA_HEURISTIC_SIMILAR_SEARCH_PERFORMANCE = 0, + #[doc = " Create a graph that has the same binary size as an HNSW graph with the given parameters\n (graph_degree = 2 * M) while trying to match the search performance as closely as possible.\n\n The reference HNSW index and the corresponding from-CAGRA generated HNSW index will NOT produce\n the same recalls and QPS for the same parameter ef. The graphs are different internally. For\n the same ef, the from-CAGRA index likely has a slightly higher recall and slightly lower QPS.\n However, the Recall-QPS curves should be similar (i.e. the points are just shifted along the\n curve)."] CUVS_CAGRA_HEURISTIC_SAME_GRAPH_FOOTPRINT = 1, } +#[doc = " Parameters for VPQ compression."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsCagraCompressionParams { + #[doc = " The bit length of the vector element after compression by PQ.\n\n Possible values: [4, 5, 6, 7, 8].\n\n Hint: the smaller the 'pq_bits', the smaller the index size and the better the search\n performance, but the lower the recall."] pub pq_bits: u32, + #[doc = " The dimensionality of the vector after compression by PQ.\n When zero, an optimal value is selected using a heuristic.\n\n TODO: at the moment `dim` must be a multiple `pq_dim`."] pub pq_dim: u32, + #[doc = " Vector Quantization (VQ) codebook size - number of \"coarse cluster centers\".\n When zero, an optimal value is selected using a heuristic."] pub vq_n_centers: u32, + #[doc = " The number of iterations searching for kmeans centers (both VQ & PQ phases)."] pub kmeans_n_iters: u32, + #[doc = " The fraction of data to use during iterative kmeans building (VQ phase).\n When zero, an optimal value is selected using a heuristic."] pub vq_kmeans_trainset_fraction: f64, + #[doc = " The fraction of data to use during iterative kmeans building (PQ phase).\n When zero, an optimal value is selected using a heuristic."] pub pq_kmeans_trainset_fraction: f64, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -1090,14 +1318,21 @@ const _: () = { [::std::mem::offset_of!(cuvsIvfPqParams, refinement_rate) - 16usize]; }; pub type cuvsIvfPqParams_t = *mut cuvsIvfPqParams; +#[doc = " Parameters for ACE (Augmented Core Extraction) graph build.\n ACE enables building indexes for datasets too large to fit in GPU memory by:\n 1. Partitioning the dataset in core (closest) and augmented (second-closest)\n partitions using balanced k-means.\n 2. Building sub-indexes for each partition independently\n 3. Concatenating sub-graphs into a final unified index"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsAceParams { + #[doc = " Number of partitions for ACE (Augmented Core Extraction) partitioned build.\n\n When set to 0 (default), the number of partitions is automatically derived\n based on available host and GPU memory to maximize partition size while\n ensuring the build fits in memory.\n\n Small values might improve recall but potentially degrade performance and\n increase memory usage. Partitions should not be too small to prevent issues\n in KNN graph construction. The partition size is on average 2 * (n_rows /\n npartitions) * dim * sizeof(T). 2 is because of the core and augmented\n vectors. Please account for imbalance in the partition sizes (up to 3x in\n our tests).\n\n If the specified number of partitions results in partitions that exceed\n available memory, the value will be automatically increased to fit memory\n constraints and a warning will be issued."] pub npartitions: usize, + #[doc = " The index quality for the ACE build.\n\n Bigger values increase the index quality. At some point, increasing this will no longer\n improve the quality."] pub ef_construction: usize, + #[doc = " Directory to store ACE build artifacts (e.g., KNN graph, optimized graph).\n\n Used when `use_disk` is true or when the graph does not fit in host and GPU\n memory. This should be the fastest disk in the system and hold enough space\n for twice the dataset, final graph, and label mapping."] pub build_dir: *const ::std::os::raw::c_char, + #[doc = " Whether to use disk-based storage for ACE build.\n\n When true, enables disk-based operations for memory-efficient graph construction."] pub use_disk: bool, + #[doc = " Maximum host memory to use for ACE build in GiB.\n\n When set to 0 (default), uses available host memory.\n When set to a positive value, limits host memory usage to the specified amount.\n Useful for testing or when running alongside other memory-intensive processes."] pub max_host_memory_gb: f64, + #[doc = " Maximum GPU memory to use for ACE build in GiB.\n\n When set to 0 (default), uses available GPU memory.\n When set to a positive value, limits GPU memory usage to the specified amount.\n Useful for testing or when running alongside other memory-intensive processes."] pub max_gpu_memory_gb: f64, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -1118,14 +1353,21 @@ const _: () = { [::std::mem::offset_of!(cuvsAceParams, max_gpu_memory_gb) - 40usize]; }; pub type cuvsAceParams_t = *mut cuvsAceParams; +#[doc = " @brief Supplemental parameters to build CAGRA Index\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsCagraIndexParams { + #[doc = " Distance type."] pub metric: cuvsDistanceType, + #[doc = " Degree of input graph for pruning."] pub intermediate_graph_degree: usize, + #[doc = " Degree of output graph."] pub graph_degree: usize, + #[doc = " ANN algorithm to build knn graph."] pub build_algo: cuvsCagraGraphBuildAlgo, + #[doc = " Number of Iterations to run if building with NN_DESCENT"] pub nn_descent_niter: usize, + #[doc = " Optional: specify graph build params based on build_algo\n - IVF_PQ: cuvsIvfPqParams_t\n - ACE: cuvsAceParams_t\n - Others: nullptr"] pub graph_build_params: *mut ::std::os::raw::c_void, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -1146,34 +1388,94 @@ const _: () = { [::std::mem::offset_of!(cuvsCagraIndexParams, graph_build_params) - 40usize]; }; pub type cuvsCagraIndexParams_t = *mut cuvsCagraIndexParams; +#[repr(u32)] +#[doc = " Algorithm used to merge physical CAGRA indices."] +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +pub enum cuvsCagraMergeAlgo { + CUVS_CAGRA_MERGE_AUTO = 0, + CUVS_CAGRA_MERGE_FASTENER = 1, + CUVS_CAGRA_MERGE_REBUILD = 2, +} +#[doc = " Parameters controlling how physical CAGRA indices are merged."] +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct cuvsCagraMergeParams { + pub algo: cuvsCagraMergeAlgo, + pub levels: u32, + pub root_fanout: u32, + pub lower_fanout: u32, + pub leader_fraction: f64, + pub max_leaders: u32, + pub leaf_size: u32, + pub leaf_degree: u32, +} +#[allow(clippy::unnecessary_operation, clippy::identity_op)] +const _: () = { + ["Size of cuvsCagraMergeParams"][::std::mem::size_of::() - 40usize]; + ["Alignment of cuvsCagraMergeParams"][::std::mem::align_of::() - 8usize]; + ["Offset of field: cuvsCagraMergeParams::algo"] + [::std::mem::offset_of!(cuvsCagraMergeParams, algo) - 0usize]; + ["Offset of field: cuvsCagraMergeParams::levels"] + [::std::mem::offset_of!(cuvsCagraMergeParams, levels) - 4usize]; + ["Offset of field: cuvsCagraMergeParams::root_fanout"] + [::std::mem::offset_of!(cuvsCagraMergeParams, root_fanout) - 8usize]; + ["Offset of field: cuvsCagraMergeParams::lower_fanout"] + [::std::mem::offset_of!(cuvsCagraMergeParams, lower_fanout) - 12usize]; + ["Offset of field: cuvsCagraMergeParams::leader_fraction"] + [::std::mem::offset_of!(cuvsCagraMergeParams, leader_fraction) - 16usize]; + ["Offset of field: cuvsCagraMergeParams::max_leaders"] + [::std::mem::offset_of!(cuvsCagraMergeParams, max_leaders) - 24usize]; + ["Offset of field: cuvsCagraMergeParams::leaf_size"] + [::std::mem::offset_of!(cuvsCagraMergeParams, leaf_size) - 28usize]; + ["Offset of field: cuvsCagraMergeParams::leaf_degree"] + [::std::mem::offset_of!(cuvsCagraMergeParams, leaf_degree) - 32usize]; +}; +pub type cuvsCagraMergeParams_t = *mut cuvsCagraMergeParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate CAGRA Index params, and populate with default values\n\n @param[in] params cuvsCagraIndexParams_t to allocate\n @return cuvsError_t"] pub fn cuvsCagraIndexParamsCreate(params: *mut cuvsCagraIndexParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate CAGRA Index params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsCagraIndexParamsDestroy(params: cuvsCagraIndexParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " Allocate CAGRA merge params and populate them with AUTO defaults."] + pub fn cuvsCagraMergeParamsCreate(params: *mut cuvsCagraMergeParams_t) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + #[doc = " De-allocate CAGRA merge params."] + pub fn cuvsCagraMergeParamsDestroy(params: cuvsCagraMergeParams_t) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + #[doc = " @brief Allocate CAGRA Compression params, and populate with default values\n\n @param[in] params cuvsCagraCompressionParams_t to allocate\n @return cuvsError_t"] pub fn cuvsCagraCompressionParamsCreate( params: *mut cuvsCagraCompressionParams_t, ) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate CAGRA Compression params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsCagraCompressionParamsDestroy(params: cuvsCagraCompressionParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate ACE params, and populate with default values\n\n @param[in] params cuvsAceParams_t to allocate\n @return cuvsError_t"] pub fn cuvsAceParamsCreate(params: *mut cuvsAceParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate ACE params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsAceParamsDestroy(params: cuvsAceParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Create CAGRA index parameters similar to an HNSW index\n\n This factory function creates CAGRA parameters that yield a graph compatible with\n an HNSW graph with the given parameters.\n\n @param[out] params The CAGRA index params to populate\n @param[in] n_rows Number of rows in the dataset\n @param[in] dim Number of dimensions in the dataset\n @param[in] M HNSW index parameter M\n @param[in] ef_construction HNSW index parameter ef_construction\n @param[in] heuristic Strategy for parameter selection\n @param[in] metric Distance metric to use\n @return cuvsError_t"] pub fn cuvsCagraIndexParamsFromHnswParams( params: cuvsCagraIndexParams_t, n_rows: i64, @@ -1186,6 +1488,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Create CAGRA index parameters heuristically tuned for a dataset\n\n This factory function selects the graph build algorithm and its parameters based on the shape of\n the dataset.\n\n @param[out] params The CAGRA index params to populate\n @param[in] n_rows Number of rows in the dataset\n @param[in] dim Number of dimensions in the dataset\n @param[in] graph_degree Degree of the output graph\n @param[in] metric Distance metric to use\n @param[in] build_quality Higher values increase build quality (and cost) up to a point\n @return cuvsError_t"] pub fn cuvsCagraIndexParamsFromDataset( params: cuvsCagraIndexParams_t, n_rows: i64, @@ -1195,9 +1498,11 @@ unsafe extern "C" { build_quality: usize, ) -> cuvsError_t; } +#[doc = " @defgroup cagra_c_extend_params C API for CUDA ANN Graph-based nearest neighbor search\n @{\n/\n/**\n @brief Supplemental parameters to extend CAGRA Index\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsCagraExtendParams { + #[doc = " The additional dataset is divided into chunks and added to the graph. This is the knob to\n adjust the tradeoff between the recall and operation throughput. Large chunk sizes can result\n in high throughput, but use more working memory (O(max_chunk_size*degree^2)). This can also\n degrade recall because no edges are added between the nodes in the same chunk. Auto select when\n 0."] pub max_chunk_size: u32, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -1211,45 +1516,70 @@ const _: () = { pub type cuvsCagraExtendParams_t = *mut cuvsCagraExtendParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate CAGRA Extend params, and populate with default values\n\n @param[in] params cuvsCagraExtendParams_t to allocate\n @return cuvsError_t"] pub fn cuvsCagraExtendParamsCreate(params: *mut cuvsCagraExtendParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate CAGRA Extend params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsCagraExtendParamsDestroy(params: cuvsCagraExtendParams_t) -> cuvsError_t; } #[repr(u32)] +#[doc = " @brief Enum to denote algorithm used to search CAGRA Index\n"] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsCagraSearchAlgo { + #[doc = " For large batch sizes."] SINGLE_CTA = 0, + #[doc = " For small batch sizes."] MULTI_CTA = 1, + #[doc = " For small batch sizes."] MULTI_KERNEL = 2, + #[doc = " For small batch sizes."] AUTO = 100, } #[repr(u32)] +#[doc = " @brief Enum to denote Hash Mode used while searching CAGRA index\n"] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsCagraHashMode { HASH = 0, SMALL = 1, AUTO_HASH = 100, } +#[doc = " @brief Supplemental parameters to search CAGRA index\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsCagraSearchParams { + #[doc = " Maximum number of queries to search at the same time (batch size). Auto select when 0."] pub max_queries: usize, + #[doc = " Number of intermediate search results retained during the search.\n\n This is the main knob to adjust trade off between accuracy and search speed.\n Higher values improve the search accuracy."] pub itopk_size: usize, + #[doc = " Upper limit of search iterations. Auto select when 0."] pub max_iterations: usize, + #[doc = " Which search implementation to use."] pub algo: cuvsCagraSearchAlgo, + #[doc = " Number of threads used to calculate a single distance. 4, 8, 16, or 32."] pub team_size: usize, + #[doc = " Number of graph nodes to select as the starting point for the search in each iteration. aka\n search width?"] pub search_width: usize, + #[doc = " Lower limit of search iterations."] pub min_iterations: usize, + #[doc = " Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0."] pub thread_block_size: usize, + #[doc = " Hashmap type. Auto selection when AUTO."] pub hashmap_mode: cuvsCagraHashMode, + #[doc = " Lower limit of hashmap bit length. More than 8."] pub hashmap_min_bitlen: usize, + #[doc = " Upper limit of hashmap fill rate. More than 0.1, less than 0.9."] pub hashmap_max_fill_rate: f32, + #[doc = " Number of iterations of initial random seed node selection. 1 or more."] pub num_random_samplings: u32, + #[doc = " Bit mask used for initial random seed node selection."] pub rand_xor_mask: u64, + #[doc = " Whether to use the persistent version of the kernel (only SINGLE_CTA is supported a.t.m.)"] pub persistent: bool, + #[doc = " Persistent kernel: time in seconds before the kernel stops if no requests received."] pub persistent_lifetime: f32, + #[doc = " Set the fraction of maximum grid size used by persistent kernel.\n Value 1.0 means the kernel grid size is maximum possible for the selected device.\n The value must be greater than 0.0 and not greater than 1.0.\n\n One may need to run other kernels alongside this persistent kernel. This parameter can\n be used to reduce the grid size of the persistent kernel to leave a few SMs idle.\n Note: running any other work on GPU alongside with the persistent kernel makes the setup\n fragile.\n - Running another kernel in another thread usually works, but no progress guaranteed\n - Any CUDA allocations block the context (this issue may be obscured by using pools)\n - Memory copies to not-pinned host memory may block the context\n\n Even when we know there are no other kernels working at the same time, setting\n kDeviceUsage to 1.0 surprisingly sometimes hurts performance. Proceed with care.\n If you suspect this is an issue, you can reduce this number to ~0.9 without a significant\n impact on the throughput."] pub persistent_device_usage: f32, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -1293,12 +1623,15 @@ const _: () = { pub type cuvsCagraSearchParams_t = *mut cuvsCagraSearchParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate CAGRA search params, and populate with default values\n\n @param[in] params cuvsCagraSearchParams_t to allocate\n @return cuvsError_t"] pub fn cuvsCagraSearchParamsCreate(params: *mut cuvsCagraSearchParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate CAGRA search params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsCagraSearchParamsDestroy(params: cuvsCagraSearchParams_t) -> cuvsError_t; } +#[doc = " @brief Struct holding the CAGRA index storage address and vector element dtype (DLPack-style)\n\n Matches the usual cuVS C index pattern (`addr` + `dtype`). \\p addr points at implementation-owned\n storage (not always a bare `cagra::index*`); free only via \\ref cuvsCagraIndexDestroy. \\p dtype\n describes index vector elements for queries and template dispatch."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsCagraIndex { @@ -1317,22 +1650,27 @@ const _: () = { pub type cuvsCagraIndex_t = *mut cuvsCagraIndex; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate CAGRA index\n\n @param[in] index cuvsCagraIndex_t to allocate\n @return cagraError_t"] pub fn cuvsCagraIndexCreate(index: *mut cuvsCagraIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate CAGRA index\n\n @param[in] index cuvsCagraIndex_t to de-allocate"] pub fn cuvsCagraIndexDestroy(index: cuvsCagraIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Get dimension of the CAGRA index\n\n @param[in] index CAGRA index\n @param[out] dim return dimension of the index\n @return cuvsError_t"] pub fn cuvsCagraIndexGetDims(index: cuvsCagraIndex_t, dim: *mut i64) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Get size of the CAGRA index\n\n @param[in] index CAGRA index\n @param[out] size return number of vectors in the index\n @return cuvsError_t"] pub fn cuvsCagraIndexGetSize(index: cuvsCagraIndex_t, size: *mut i64) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Get graph degree of the CAGRA index\n\n @param[in] index CAGRA index\n @param[out] graph_degree return graph degree\n @return cuvsError_t"] pub fn cuvsCagraIndexGetGraphDegree( index: cuvsCagraIndex_t, graph_degree: *mut i64, @@ -1340,6 +1678,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Returns a view of the CAGRA dataset\n\n This function returns a non-owning view of the CAGRA dataset.\n The output will be referencing device memory that is directly used\n in CAGRA, without copying the dataset at all. This means that the\n output is only valid as long as the CAGRA index is alive, and once\n cuvsCagraIndexDestroy is called on the cagra index - the returned\n dataset view will be invalid.\n\n Note that the DLManagedTensor dataset returned will have an associated\n 'deleter' function that must be called when the dataset is no longer\n needed. This will free up host memory that stores the shape of the\n dataset view.\n\n @param[in] index CAGRA index\n @param[out] dataset the dataset used in cagra\n @return cuvsError_t"] pub fn cuvsCagraIndexGetDataset( index: cuvsCagraIndex_t, dataset: *mut DLManagedTensor, @@ -1347,6 +1686,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Returns a view of the CAGRA graph\n\n This function returns a non-owning view of the CAGRA graph.\n The output will be referencing device memory that is directly used\n in CAGRA, without copying the graph at all. This means that the\n output is only valid as long as the CAGRA index is alive, and once\n cuvsCagraIndexDestroy is called on the cagra index - the returned\n graph view will be invalid.\n\n Note that the DLManagedTensor graph returned will have an associated\n 'deleter' function that must be called when the graph is no longer\n needed. This will free up host memory that stores the metadata for the\n graph view.\n\n @param[in] index CAGRA index\n @param[out] graph the output knn graph.\n @return cuvsError_t"] pub fn cuvsCagraIndexGetGraph( index: cuvsCagraIndex_t, graph: *mut DLManagedTensor, @@ -1354,6 +1694,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Update a CAGRA index with a device-padded dataset.\n\n This is the centralized dataset update operation for C callers. If \\p index\n is already device-padded, its dataset view is replaced in place. Otherwise,\n the index is converted and its opaque handle is rebound to a search-ready\n device-padded index. Caller retains ownership of\n \\p device_padded_dataset and must keep it alive while \\p index uses it.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] device_padded_dataset owning or non-owning device-padded dataset handle\n @param[inout] index CAGRA index handle\n @return cuvsError_t"] pub fn cuvsCagraUpdateDataset( res: cuvsResources_t, device_padded_dataset: cuvsDataset_t, @@ -1362,6 +1703,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Build a CAGRA index from a dataset handle. Acceptable underlying\n types are:\n 1. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n 2. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 16`\n 3. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8`\n 4. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8`\n\n The memory space and layout \\p dataset was constructed with select the C++ build overload.\n Build the handle with an owning factory or the matching dataset view factory\n (`cuvsDatasetMakePaddedView` / `cuvsDatasetMakeStandardView`).\n\n Note that a dataset residing in host memory produces a host-backed index, which\n must be made search-ready with `cuvsCagraUpdateDataset` (using a device-padded\n dataset) before calling `cuvsCagraSearch`.\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Assume a populated `DLManagedTensor` type here holding a device padded dataset\n DLManagedTensor dataset;\n\n // Wrap it in a non-owning dataset handle\n cuvsDataset_t dataset_view;\n cuvsError_t view_create_status = cuvsDatasetMakePaddedView(res, &dataset, &dataset_view);\n\n // Create default index params\n cuvsCagraIndexParams_t params;\n cuvsError_t params_create_status = cuvsCagraIndexParamsCreate(¶ms);\n\n // Create CAGRA index\n cuvsCagraIndex_t index;\n cuvsError_t index_create_status = cuvsCagraIndexCreate(&index);\n\n // Build the CAGRA Index\n cuvsError_t build_status = cuvsCagraBuild(res, params, dataset_view, index);\n\n // de-allocate `dataset_view`, `params`, `index` and `res`\n cuvsError_t view_destroy_status = cuvsDatasetDestroy(dataset_view);\n cuvsError_t params_destroy_status = cuvsCagraIndexParamsDestroy(params);\n cuvsError_t index_destroy_status = cuvsCagraIndexDestroy(index);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params cuvsCagraIndexParams_t used to build CAGRA index\n @param[in] dataset cuvsDataset_t training dataset or dataset view\n @param[inout] index cuvsCagraIndex_t Newly built CAGRA index. This index needs to be already\n created with cuvsCagraIndexCreate.\n @return cuvsError_t"] pub fn cuvsCagraBuild( res: cuvsResources_t, params: cuvsCagraIndexParams_t, @@ -1371,6 +1713,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Extend a CAGRA index with a caller-owned pre-concatenated padded dataset.\n\n The caller must build `extended_dataset` as `old || new` (size `n_old + n_new`) before calling.\n Rows `[0, new_start_row)` are the original vectors; rows `[new_start_row, n_rows)` are the\n additional vectors. `new_start_row` must equal the current index size. The library only extends\n the graph and rebinds the index to `extended_dataset`; keep that dataset alive for the index\n lifetime.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params cuvsCagraExtendParams_t used to extend CAGRA index\n @param[in] extended_dataset cuvsDataset_t caller-owned device-padded dataset of old || new\n @param[in] new_start_row row index where the additional vectors begin\n @param[in,out] index cuvsCagraIndex_t CAGRA index\n @return cuvsError_t"] pub fn cuvsCagraExtend( res: cuvsResources_t, params: cuvsCagraExtendParams_t, @@ -1381,6 +1724,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup cagra_c_index_search C API for CUDA ANN Graph-based nearest neighbor search\n @{\n/\n/**\n @brief Search a CAGRA index with a `DLManagedTensor` which has underlying\n `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`.\n It is also important to note that the CAGRA Index must have been built\n with the same type of `queries`, such that `index.dtype.code ==\n queries.dl_tensor.dtype.code` Types for input are:\n 1. `queries`:\n a. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n b. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 16`\n c. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8`\n d. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8`\n 2. `neighbors`: `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 32`\n or `kDLDataType.code == kDLInt` and `kDLDataType.bits = 64`\n 3. `distances`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Assume a populated `DLManagedTensor` type here\n DLManagedTensor dataset;\n DLManagedTensor queries;\n DLManagedTensor neighbors;\n\n // Create default search params\n cuvsCagraSearchParams_t params;\n cuvsError_t params_create_status = cuvsCagraSearchParamsCreate(¶ms);\n\n // Search the `index` built using `cuvsCagraBuild`\n cuvsError_t search_status = cuvsCagraSearch(res, params, index, &queries, &neighbors,\n &distances);\n\n // de-allocate `params` and `res`\n cuvsError_t params_destroy_status = cuvsCagraSearchParamsDestroy(params);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params cuvsCagraSearchParams_t used to search CAGRA index\n @param[in] index cuvsCagraIndex which has been returned by `cuvsCagraBuild`\n @param[in] queries DLManagedTensor* queries dataset to search\n @param[out] neighbors DLManagedTensor* output `k` neighbors for queries\n @param[out] distances DLManagedTensor* output `k` distances for queries\n @param[in] filter cuvsFilter input filter that can be used\nto filter queries and neighbors based on the given bitset."] pub fn cuvsCagraSearch( res: cuvsResources_t, params: cuvsCagraSearchParams_t, @@ -1393,6 +1737,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Search multiple CAGRA index partitions concurrently and return the global top-k per\n query.\n\n For each query row, the function searches all partitions in parallel into an internal\n intermediate buffer, applies per-partition distance post-processing, runs a batched top-k\n merge across partitions, and writes the final outputs to the caller-supplied device tensors.\n All work is submitted to the CUDA stream associated with @p res; use @c cuvsStreamSync to\n wait for completion.\n\n The index element type may be float32, float16, int8, or uint8. All partitions must share the\n same element type, and the queries must use that same type.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params search parameters (shared across partitions)\n @param[in] num_partitions number of index partitions\n @param[in] indices array of num_partitions cuvsCagraIndex_t pointers, all of the same\n element type\n @param[in] queries DLManagedTensor* (device, same dtype as the indices, [n_queries,\n dim]); the queries matrix is searched against every partition\n @param[out] partition_ids DLManagedTensor* (device, uint32, [n_queries, k]); which partition\n each returned neighbor came from\n @param[out] neighbors DLManagedTensor* (device, uint32 or int64, [n_queries, k]); ordinal\n in the corresponding partition's dataset\n @param[out] distances DLManagedTensor* (device, float32, [n_queries, k]); post-processed\n distance for each (query, neighbor)\n @param[in] filters array of `num_partitions` filters, one per partition (or NULL for a\n fully unfiltered search). `filters[i]` applies to partition `i`: use\n {.type=NO_FILTER, .addr=0} for no filter on that partition, or\n {.type=BITSET, .addr=ptr} where ptr is a uintptr_t-cast\n DLManagedTensor* holding that partition's own bitset (one bit per\n vector in that partition; standard 32-bit packing)."] pub fn cuvsCagraSearchMultiPartition( res: cuvsResources_t, params: cuvsCagraSearchParams_t, @@ -1407,6 +1752,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup cagra_c_index_serialize CAGRA C-API serialize functions\n @{\n/\n/**\n Save the CAGRA graph to file without its dataset.\n\n Experimental, both the API and the serialization format are subject to change.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename the file name for saving the graph\n @param[in] index CAGRA index"] pub fn cuvsCagraSerializeGraph( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -1415,6 +1761,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " Save the CAGRA graph and its attached dataset to file.\n\n The index stores a non-owning dataset view. The caller must keep the dataset backing that view\n alive while this function runs. Returns CUVS_ERROR without modifying the destination file if\n the index has no attached dataset.\n\n Experimental, both the API and the serialization format are subject to change.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename the file name for saving the graph and dataset\n @param[in] index CAGRA index with an attached host or device dataset"] pub fn cuvsCagraSerializeGraphAndDataset( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -1423,6 +1770,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " Save the CAGRA index to file in hnswlib format.\n NOTE: The saved index can only be read by the hnswlib wrapper in cuVS,\n as the serialization format is not compatible with the original hnswlib.\n\n Experimental, both the API and the serialization format are subject to change.\n\n @code{.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // create an index with `cuvsCagraBuild`\n cuvsCagraSerializeHnswlib(res, \"/path/to/index\", index);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename the file name for saving the index\n @param[in] index CAGRA index\n"] pub fn cuvsCagraSerializeToHnswlib( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -1431,6 +1779,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " Load the CAGRA graph from file without retaining a serialized dataset.\n\n This succeeds whether or not the file contains a dataset. Use cuvsCagraUpdateDataset to attach a\n caller-owned device-padded dataset view before searching the graph-only index.\n\n Experimental, both the API and the serialization format are subject to change.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename the name of the file that stores the index\n @param[inout] index pre-created CAGRA index populated on success and unchanged on failure"] pub fn cuvsCagraDeserializeGraph( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -1439,6 +1788,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " Load the CAGRA graph and dataset from file.\n\n The returned owning dataset preserves the serialized host/device memory type and\n standard/padded layout. The index stores a non-owning view into it, so the caller must keep the\n dataset alive while the index uses it and destroy it separately with cuvsDatasetDestroy. Only a\n device-padded result is immediately searchable through the C API; attach a caller-owned\n device-padded view with cuvsCagraUpdateDataset for any other kind. The output pointer\n must point to a null handle on entry; deserialization acts as a factory and transfers ownership\n of the allocated dataset handle on success. Returns CUVS_ERROR when the file has no dataset; the\n index and output handle are unchanged on failure.\n\n Experimental, both the API and the serialization format are subject to change.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename the name of the file that stores the graph and dataset\n @param[inout] index pre-created CAGRA index populated on success and unchanged on failure\n @param[out] out_dataset receives the allocated owning dataset handle; must point to null on entry"] pub fn cuvsCagraDeserializeGraphAndDataset( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -1448,6 +1798,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " Load index from a dataset and graph\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] metric cuvsDistanceType to use in the index\n @param[in] graph the knn graph to use, shape (size, graph_degree)\n @param[in] dataset the dataset to use, shape (size, dim)\n @param[inout] index cuvsCagraIndex_t CAGRA index populated with the graph and dataset.\n This index needs to be already created with\n cuvsCagraIndexCreate.\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Create CAGRA index\n cuvsCagraIndex_t index;\n cuvsError_t index_create_status = cuvsCagraIndexCreate(&index);\n\n // Assume a populated `DLManagedTensor` type here for the graph and dataset\n DLManagedTensor dataset;\n DLManagedTensor graph;\n\n cuvsDistanceType metric = L2Expanded;\n\n // Build the CAGRA Index from the graph/dataset\n cuvsError_t status = cuvsCagraIndexFromArgs(res, metric, &graph, &dataset, index);\n\n @endcode"] pub fn cuvsCagraIndexFromArgs( res: cuvsResources_t, metric: cuvsDistanceType, @@ -1458,6 +1809,18 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Compute per-index write offsets for a merged dataset buffer.\n\n `cuvsCagraMerge`/`cuvsCagraMergeWithParams` require the caller to have already concatenated\n every input index's dataset (in `indices` order, applying `filter` if any) into a single buffer\n and to know each index's starting row within it. For `filter.type == NO_FILTER`, those offsets\n are just the cumulative sizes of `indices` and this function is not needed. For `BITSET`, the\n number of surviving rows per index cannot be derived any other way, so call this first.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] indices Array of input cuvsCagraIndex_t handles that will be passed to merge\n @param[in] num_indices Number of input indices\n @param[in] filter Filter that will be passed to merge. Only `NO_FILTER` and `BITSET` supported.\n @param[out] offsets Caller-allocated array of `num_indices + 1` int64_t. Entry `i` is the row at\n which `indices[i]`'s surviving rows must start in the merged buffer; the\n last entry is the total row count of the merged buffer.\n @return cuvsError_t"] + pub fn cuvsCagraMergedDatasetOffsets( + res: cuvsResources_t, + indices: *mut cuvsCagraIndex_t, + num_indices: usize, + filter: cuvsFilter, + offsets: *mut i64, + ) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + #[doc = " @brief Merge multiple CAGRA indices into a single CAGRA index.\n\n The caller is responsible for concatenating every input index's dataset (applying `filter` if\n any) into a single `merged_dataset` buffer before calling this, and for computing `offsets`\n (see `cuvsCagraMergedDatasetOffsets`). This function only builds/merges the graph and rebinds\n the output index to `merged_dataset` -- it never allocates or copies dataset rows itself. This\n mirrors the `cuvsCagraExtend` contract.\n\n All input indices must have been built with the same data type (`index.dtype`) and\n have the same dimensionality (`index.dims`). The merged index uses the output\n parameters specified in `cuvsCagraIndexParams`. The merge algorithm is selected automatically.\n\n Input indices must have:\n - `index.dtype.code` and `index.dtype.bits` matching across all indices.\n - Supported data types for indices:\n a. `kDLFloat` with `bits = 32`\n b. `kDLFloat` with `bits = 16`\n c. `kDLInt` with `bits = 8`\n d. `kDLUInt` with `bits = 8`\n\n The resulting output index will have the same data type as the input indices.\n\n Example:\n @code{.c}\n #include \n #include \n\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n cuvsCagraIndex_t index1, index2, merged_index;\n cuvsCagraIndexCreate(&index1);\n cuvsCagraIndexCreate(&index2);\n cuvsCagraIndexCreate(&merged_index);\n\n // Assume index1 and index2 have device datasets and were built using cuvsCagraBuild.\n\n cuvsCagraIndexParams_t merge_params;\n cuvsError_t params_create_status = cuvsCagraIndexParamsCreate(&merge_params);\n\n // Build `merged_dataset` as the caller-owned concatenation of index1 || index2 (e.g. via\n // cuvsDatasetMakePadded over a device buffer you populated yourself).\n cuvsDataset_t merged_dataset = ...;\n int64_t offsets[3] = {0, index1_size, index1_size + index2_size};\n cuvsFilter filter = {.type = NO_FILTER, .addr = 0};\n\n cuvsError_t merge_status = cuvsCagraMerge(res, merge_params, (cuvsCagraIndex_t[]){index1,\n index2}, 2, filter, merged_dataset, offsets, merged_index);\n\n // Use merged_index for search operations\n\n cuvsCagraIndexDestroy(merged_index);\n cuvsDatasetDestroy(merged_dataset);\n cuvsError_t params_destroy_status = cuvsCagraIndexParamsDestroy(merge_params);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params cuvsCagraIndexParams_t parameters for the output index\n @param[in] indices Array of input cuvsCagraIndex_t handles to merge\n @param[in] num_indices Number of input indices\n @param[in] filter Filter, already applied by the caller while building `merged_dataset`\n @param[in] merged_dataset Caller-owned dataset handle already containing the concatenated (and,\n if `filter` is set, already-filtered) dataset, with the same layout as\n the input indices. Keep this dataset alive while using\n \\p output_index. A host-backed dataset must be updated with\n `cuvsCagraUpdateDataset` before device search.\n @param[in] offsets Per-index starting row within `merged_dataset`, as returned by\n `cuvsCagraMergedDatasetOffsets`. Array of `num_indices + 1` int64_t; the last\n entry must equal `merged_dataset`'s row count.\n @param[out] output_index Output handle that will store the merged index.\n Must be initialized using `cuvsCagraIndexCreate` before use."] pub fn cuvsCagraMerge( res: cuvsResources_t, params: cuvsCagraIndexParams_t, @@ -1465,19 +1828,44 @@ unsafe extern "C" { num_indices: usize, filter: cuvsFilter, merged_dataset: cuvsDataset_t, + offsets: *const i64, + output_index: cuvsCagraIndex_t, + ) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + #[doc = " @brief Merge multiple CAGRA indices with explicit merge parameters.\n\n See `cuvsCagraMerge` for the full `merged_dataset`/`offsets` contract.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params cuvsCagraIndexParams_t parameters for the output index\n @param[in] merge_params cuvsCagraMergeParams_t parameters controlling the merge algorithm, or\n NULL to use AUTO defaults\n @param[in] indices Array of input cuvsCagraIndex_t handles to merge\n @param[in] num_indices Number of input indices\n @param[in] filter Filter, already applied by the caller while building `merged_dataset`\n @param[in] merged_dataset Caller-owned dataset handle already containing the concatenated (and,\n if `filter` is set, already-filtered) dataset. Keep this dataset alive\n while using `output_index`. A host-backed dataset must be updated with\n `cuvsCagraUpdateDataset` before device search.\n @param[in] offsets Per-index starting row within `merged_dataset`, as returned by\n `cuvsCagraMergedDatasetOffsets`. Array of `num_indices + 1` int64_t.\n @param[out] output_index Output handle initialized with `cuvsCagraIndexCreate`"] + pub fn cuvsCagraMergeWithParams( + res: cuvsResources_t, + params: cuvsCagraIndexParams_t, + merge_params: cuvsCagraMergeParams_t, + indices: *mut cuvsCagraIndex_t, + num_indices: usize, + filter: cuvsFilter, + merged_dataset: cuvsDataset_t, + offsets: *const i64, output_index: cuvsCagraIndex_t, ) -> cuvsError_t; } +#[doc = " @defgroup ivf_flat_c_index_params IVF-Flat index build parameters\n @{\n/\n/**\n @brief Supplemental parameters to build IVF-Flat Index\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsIvfFlatIndexParams { + #[doc = " Distance type."] pub metric: cuvsDistanceType, + #[doc = " The argument used by some distance metrics."] pub metric_arg: f32, + #[doc = " Whether to add the dataset content to the index, i.e.:\n\n - `true` means the index is filled with the dataset vectors and ready to search after calling\n `build`.\n - `false` means `build` only trains the underlying model (e.g. quantizer or clustering), but\n the index is left empty; you'd need to call `extend` on the index afterwards to populate it."] pub add_data_on_build: bool, + #[doc = " The number of inverted lists (clusters)"] pub n_lists: u32, + #[doc = " The number of iterations searching for kmeans centers (index building)."] pub kmeans_n_iters: u32, + #[doc = " The fraction of data to use during iterative kmeans building."] pub kmeans_trainset_fraction: f64, + #[doc = " By default (adaptive_centers = false), the cluster centers are trained in `ivf_flat::build`,\n and never modified in `ivf_flat::extend`. As a result, you may need to retrain the index\n from scratch after invoking (`ivf_flat::extend`) a few times with new data, the distribution of\n which is no longer representative of the original training set.\n\n The alternative behavior (adaptive_centers = true) is to update the cluster centers for new\n data when it is added. In this case, `index.centers()` are always exactly the centroids of the\n data in the corresponding clusters. The drawback of this behavior is that the centroids depend\n on the order of adding new data (through the classification of the added data); that is,\n `index.centers()` \"drift\" together with the changing distribution of the newly added data."] pub adaptive_centers: bool, + #[doc = " By default, the algorithm allocates more space than necessary for individual clusters\n (`list_data`). This allows to amortize the cost of memory allocation and reduce the number of\n data copies during repeated calls to `extend` (extending the database).\n\n The alternative is the conservative allocation behavior; when enabled, the algorithm always\n allocates the minimum amount of memory required to store the given number of records. Set this\n flag to `true` if you prefer to use as little GPU memory for the database as possible."] pub conservative_memory_allocation: bool, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -1505,16 +1893,20 @@ const _: () = { pub type cuvsIvfFlatIndexParams_t = *mut cuvsIvfFlatIndexParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate IVF-Flat Index params, and populate with default values\n\n @param[in] index_params cuvsIvfFlatIndexParams_t to allocate\n @return cuvsError_t"] pub fn cuvsIvfFlatIndexParamsCreate(index_params: *mut cuvsIvfFlatIndexParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate IVF-Flat Index params\n\n @param[in] index_params\n @return cuvsError_t"] pub fn cuvsIvfFlatIndexParamsDestroy(index_params: cuvsIvfFlatIndexParams_t) -> cuvsError_t; } +#[doc = " @defgroup ivf_flat_c_search_params IVF-Flat index search parameters\n @{\n/\n/**\n @brief Supplemental parameters to search IVF-Flat index\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsIvfFlatSearchParams { + #[doc = " The number of clusters to search."] pub n_probes: u32, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -1528,12 +1920,15 @@ const _: () = { pub type cuvsIvfFlatSearchParams_t = *mut cuvsIvfFlatSearchParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate IVF-Flat search params, and populate with default values\n\n @param[in] params cuvsIvfFlatSearchParams_t to allocate\n @return cuvsError_t"] pub fn cuvsIvfFlatSearchParamsCreate(params: *mut cuvsIvfFlatSearchParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate IVF-Flat search params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsIvfFlatSearchParamsDestroy(params: cuvsIvfFlatSearchParams_t) -> cuvsError_t; } +#[doc = " @defgroup ivf_flat_c_index IVF-Flat index\n @{\n/\n/**\n @brief Struct to hold address of cuvs::neighbors::ivf_flat::index and its active trained dtype\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsIvfFlatIndex { @@ -1552,22 +1947,27 @@ const _: () = { pub type cuvsIvfFlatIndex_t = *mut cuvsIvfFlatIndex; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate IVF-Flat index\n\n @param[in] index cuvsIvfFlatIndex_t to allocate\n @return cuvsError_t"] pub fn cuvsIvfFlatIndexCreate(index: *mut cuvsIvfFlatIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate IVF-Flat index\n\n @param[in] index cuvsIvfFlatIndex_t to de-allocate"] pub fn cuvsIvfFlatIndexDestroy(index: cuvsIvfFlatIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the number of clusters/inverted lists in the index\n\n @param[in] index cuvsIvfFlatIndex_t Built IVF-Flat index\n @param[out] n_lists Pointer to store the number of lists\n @return cuvsError_t"] pub fn cuvsIvfFlatIndexGetNLists(index: cuvsIvfFlatIndex_t, n_lists: *mut i64) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the dimensionality of the indexed data\n\n @param[in] index cuvsIvfFlatIndex_t Built IVF-Flat index\n @param[out] dim Pointer to store the dimensionality\n @return cuvsError_t"] pub fn cuvsIvfFlatIndexGetDim(index: cuvsIvfFlatIndex_t, dim: *mut i64) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the cluster centers corresponding to the lists [n_lists, dim]\n\n @param[in] index cuvsIvfFlatIndex_t Built Ivf-Flat Index\n @param[out] centers Preallocated array on host or device memory to store output, [n_lists, dim]\n @return cuvsError_t"] pub fn cuvsIvfFlatIndexGetCenters( index: cuvsIvfFlatIndex_t, centers: *mut DLManagedTensor, @@ -1575,6 +1975,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup ivf_flat_c_index_build IVF-Flat index build\n @{\n/\n/**\n @brief Build a IVF-Flat index with a `DLManagedTensor` which has underlying\n `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`,\n or `kDLCPU`. Also, acceptable underlying types are:\n 1. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n 2. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8`\n 3. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8`\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Assume a populated `DLManagedTensor` type here\n DLManagedTensor dataset;\n\n // Create default index params\n cuvsIvfFlatIndexParams_t index_params;\n cuvsError_t params_create_status = cuvsIvfFlatIndexParamsCreate(&index_params);\n\n // Create IVF-Flat index\n cuvsIvfFlatIndex_t index;\n cuvsError_t index_create_status = cuvsIvfFlatIndexCreate(&index);\n\n // Build the IVF-Flat Index\n cuvsError_t build_status = cuvsIvfFlatBuild(res, index_params, &dataset, index);\n\n // de-allocate `index_params`, `index` and `res`\n cuvsError_t params_destroy_status = cuvsIvfFlatIndexParamsDestroy(index_params);\n cuvsError_t index_destroy_status = cuvsIvfFlatIndexDestroy(index);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] index_params cuvsIvfFlatIndexParams_t used to build IVF-Flat index\n @param[in] dataset DLManagedTensor* training dataset\n @param[out] index cuvsIvfFlatIndex_t Newly built IVF-Flat index\n @return cuvsError_t"] pub fn cuvsIvfFlatBuild( res: cuvsResources_t, index_params: cuvsIvfFlatIndexParams_t, @@ -1584,6 +1985,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup ivf_flat_c_index_search IVF-Flat index search\n @{\n/\n/**\n @brief Search a IVF-Flat index with a `DLManagedTensor` which has underlying\n `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`.\n It is also important to note that the IVF-Flat Index must have been built\n with the same type of `queries`, such that `index.dtype.code ==\n queries.dl_tensor.dtype.code` Types for input are:\n 1. `queries`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n 2. `neighbors`: `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 32`\n 3. `distances`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Assume a populated `DLManagedTensor` type here\n DLManagedTensor dataset;\n DLManagedTensor queries;\n DLManagedTensor neighbors;\n\n // Create default search params\n cuvsIvfFlatSearchParams_t search_params;\n cuvsError_t params_create_status = cuvsIvfFlatSearchParamsCreate(&search_params);\n\n // Search the `index` built using `ivfFlatBuild`\n cuvsError_t search_status = cuvsIvfFlatSearch(res, search_params, index, &queries, &neighbors,\n &distances);\n\n // de-allocate `search_params` and `res`\n cuvsError_t params_destroy_status = cuvsIvfFlatSearchParamsDestroy(search_params);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] search_params cuvsIvfFlatSearchParams_t used to search IVF-Flat index\n @param[in] index ivfFlatIndex which has been returned by `ivfFlatBuild`\n @param[in] queries DLManagedTensor* queries dataset to search\n @param[out] neighbors DLManagedTensor* output `k` neighbors for queries\n @param[out] distances DLManagedTensor* output `k` distances for queries\n @param[in] filter cuvsFilter input filter that can be used\nto filter queries and neighbors based on the given bitset."] pub fn cuvsIvfFlatSearch( res: cuvsResources_t, search_params: cuvsIvfFlatSearchParams_t, @@ -1596,6 +1998,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup ivf_flat_c_index_serialize IVF-Flat C-API serialize functions\n @{\n/\n/**\n Save the index to file.\n\n Experimental, both the API and the serialization format are subject to change.\n\n @code{.cpp}\n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // create an index with `cuvsIvfFlatBuild`\n cuvsIvfFlatSerialize(res, \"/path/to/index\", index, true);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename the file name for saving the index\n @param[in] index IVF-Flat index"] pub fn cuvsIvfFlatSerialize( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -1604,6 +2007,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " Load index from file.\n\n Experimental, both the API and the serialization format are subject to change.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename the name of the file that stores the index\n @param[out] index IVF-Flat index loaded disk"] pub fn cuvsIvfFlatDeserialize( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -1612,6 +2016,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup ivf_flat_c_index_extend IVF-Flat index extend\n @{\n/\n/**\n @brief Extend the index with the new data.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] new_vectors DLManagedTensor* the new vectors to add to the index\n @param[in] new_indices DLManagedTensor* vector of new indices for the new vectors\n @param[inout] index IVF-Flat index to be extended\n @return cuvsError_t"] pub fn cuvsIvfFlatExtend( res: cuvsResources_t, new_vectors: *mut DLManagedTensor, @@ -1619,15 +2024,23 @@ unsafe extern "C" { index: cuvsIvfFlatIndex_t, ) -> cuvsError_t; } +#[doc = " @defgroup ivf_sq_c_index_params IVF-SQ index build parameters\n @{\n/\n/**\n @brief Supplemental parameters to build IVF-SQ Index\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsIvfSqIndexParams { + #[doc = " Distance type."] pub metric: cuvsDistanceType, + #[doc = " The argument used by some distance metrics."] pub metric_arg: f32, + #[doc = " Whether to add the dataset content to the index, i.e.:\n\n - `true` means the index is filled with the dataset vectors and ready to search after calling\n `build`.\n - `false` means `build` only trains the underlying model (e.g. quantizer or clustering), but\n the index is left empty; you'd need to call `extend` on the index afterwards to populate it."] pub add_data_on_build: bool, + #[doc = " The number of inverted lists (clusters)"] pub n_lists: u32, + #[doc = " The number of iterations searching for kmeans centers (index building)."] pub kmeans_n_iters: u32, + #[doc = " The number of data vectors per cluster to use during iterative kmeans building.\n The index uses at most `n_lists * max_train_points_per_cluster` rows for training."] pub max_train_points_per_cluster: u32, + #[doc = " By default, the algorithm allocates more space than necessary for individual clusters\n (`list_data`). This allows to amortize the cost of memory allocation and reduce the number of\n data copies during repeated calls to `extend` (extending the database).\n\n The alternative is the conservative allocation behavior; when enabled, the algorithm always\n allocates the minimum amount of memory required to store the given number of records. Set this\n flag to `true` if you prefer to use as little GPU memory for the database as possible."] pub conservative_memory_allocation: bool, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -1652,15 +2065,19 @@ const _: () = { pub type cuvsIvfSqIndexParams_t = *mut cuvsIvfSqIndexParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate IVF-SQ Index params, and populate with default values\n\n @param[in] index_params cuvsIvfSqIndexParams_t to allocate\n @return cuvsError_t"] pub fn cuvsIvfSqIndexParamsCreate(index_params: *mut cuvsIvfSqIndexParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate IVF-SQ Index params\n\n @param[in] index_params\n @return cuvsError_t"] pub fn cuvsIvfSqIndexParamsDestroy(index_params: cuvsIvfSqIndexParams_t) -> cuvsError_t; } +#[doc = " @defgroup ivf_sq_c_search_params IVF-SQ index search parameters\n @{\n/\n/**\n @brief Supplemental parameters to search IVF-SQ index\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsIvfSqSearchParams { + #[doc = " The number of clusters to search."] pub n_probes: u32, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -1674,12 +2091,15 @@ const _: () = { pub type cuvsIvfSqSearchParams_t = *mut cuvsIvfSqSearchParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate IVF-SQ search params, and populate with default values\n\n @param[in] params cuvsIvfSqSearchParams_t to allocate\n @return cuvsError_t"] pub fn cuvsIvfSqSearchParamsCreate(params: *mut cuvsIvfSqSearchParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate IVF-SQ search params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsIvfSqSearchParamsDestroy(params: cuvsIvfSqSearchParams_t) -> cuvsError_t; } +#[doc = " @defgroup ivf_sq_c_index IVF-SQ index\n @{\n/\n/**\n @brief Struct to hold address of cuvs::neighbors::ivf_sq::index and its active trained dtype\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsIvfSqIndex { @@ -1698,26 +2118,32 @@ const _: () = { pub type cuvsIvfSqIndex_t = *mut cuvsIvfSqIndex; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate IVF-SQ index\n\n @param[in] index cuvsIvfSqIndex_t to allocate\n @return cuvsError_t"] pub fn cuvsIvfSqIndexCreate(index: *mut cuvsIvfSqIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate IVF-SQ index\n\n @param[in] index cuvsIvfSqIndex_t to de-allocate"] pub fn cuvsIvfSqIndexDestroy(index: cuvsIvfSqIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " Get the number of clusters/inverted lists"] pub fn cuvsIvfSqIndexGetNLists(index: cuvsIvfSqIndex_t, n_lists: *mut i64) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " Get the dimensionality of the data"] pub fn cuvsIvfSqIndexGetDim(index: cuvsIvfSqIndex_t, dim: *mut i64) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " Get the size of the index"] pub fn cuvsIvfSqIndexGetSize(index: cuvsIvfSqIndex_t, size: *mut i64) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the cluster centers corresponding to the lists [n_lists, dim]\n\n @param[in] index cuvsIvfSqIndex_t Built Ivf-SQ Index\n @param[out] centers Preallocated array on host or device memory to store output, [n_lists, dim]\n @return cuvsError_t"] pub fn cuvsIvfSqIndexGetCenters( index: cuvsIvfSqIndex_t, centers: *mut DLManagedTensor, @@ -1725,6 +2151,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup ivf_sq_c_index_build IVF-SQ index build\n @{\n/\n/**\n @brief Build an IVF-SQ index with a `DLManagedTensor` which has underlying\n `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`,\n or `kDLCPU`. Also, acceptable underlying types are:\n 1. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n 2. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 16`\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Assume a populated `DLManagedTensor` type here\n DLManagedTensor dataset;\n\n // Create default index params\n cuvsIvfSqIndexParams_t index_params;\n cuvsError_t params_create_status = cuvsIvfSqIndexParamsCreate(&index_params);\n\n // Create IVF-SQ index\n cuvsIvfSqIndex_t index;\n cuvsError_t index_create_status = cuvsIvfSqIndexCreate(&index);\n\n // Build the IVF-SQ Index\n cuvsError_t build_status = cuvsIvfSqBuild(res, index_params, &dataset, index);\n\n // de-allocate `index_params`, `index` and `res`\n cuvsError_t params_destroy_status = cuvsIvfSqIndexParamsDestroy(index_params);\n cuvsError_t index_destroy_status = cuvsIvfSqIndexDestroy(index);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] index_params cuvsIvfSqIndexParams_t used to build IVF-SQ index\n @param[in] dataset DLManagedTensor* training dataset\n @param[out] index cuvsIvfSqIndex_t Newly built IVF-SQ index\n @return cuvsError_t"] pub fn cuvsIvfSqBuild( res: cuvsResources_t, index_params: cuvsIvfSqIndexParams_t, @@ -1734,6 +2161,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup ivf_sq_c_index_search IVF-SQ index search\n @{\n/\n/**\n @brief Search an IVF-SQ index with a `DLManagedTensor` which has underlying\n `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`.\n Types for input are:\n 1. `queries`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32` or 16\n 2. `neighbors`: `kDLDataType.code == kDLInt` and `kDLDataType.bits = 64`\n 3. `distances`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Assume a populated `DLManagedTensor` type here\n DLManagedTensor queries;\n DLManagedTensor neighbors;\n DLManagedTensor distances;\n\n // Create default search params\n cuvsIvfSqSearchParams_t search_params;\n cuvsError_t params_create_status = cuvsIvfSqSearchParamsCreate(&search_params);\n\n // Search the `index` built using `cuvsIvfSqBuild`\n cuvsError_t search_status = cuvsIvfSqSearch(\n res, search_params, index, &queries, &neighbors, &distances, (cuvsFilter){});\n\n // de-allocate `search_params` and `res`\n cuvsError_t params_destroy_status = cuvsIvfSqSearchParamsDestroy(search_params);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] search_params cuvsIvfSqSearchParams_t used to search IVF-SQ index\n @param[in] index ivfSqIndex which has been returned by `cuvsIvfSqBuild`\n @param[in] queries DLManagedTensor* queries dataset to search\n @param[out] neighbors DLManagedTensor* output `k` neighbors for queries\n @param[out] distances DLManagedTensor* output `k` distances for queries\n @param[in] filter cuvsFilter input filter that can be used\n to filter queries and neighbors based on the given bitset."] pub fn cuvsIvfSqSearch( res: cuvsResources_t, search_params: cuvsIvfSqSearchParams_t, @@ -1746,6 +2174,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup ivf_sq_c_index_serialize IVF-SQ C-API serialize functions\n @{\n/\n/**\n Save the index to file.\n\n Experimental, both the API and the serialization format are subject to change.\n\n @code{.c}\n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // create an index with `cuvsIvfSqBuild`\n cuvsIvfSqSerialize(res, \"/path/to/index\", index);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename the file name for saving the index\n @param[in] index IVF-SQ index"] pub fn cuvsIvfSqSerialize( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -1754,6 +2183,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " Load index from file.\n\n Experimental, both the API and the serialization format are subject to change.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename the name of the file that stores the index\n @param[out] index IVF-SQ index loaded from disk"] pub fn cuvsIvfSqDeserialize( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -1762,6 +2192,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup ivf_sq_c_index_extend IVF-SQ index extend\n @{\n/\n/**\n @brief Extend the index with the new data.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] new_vectors DLManagedTensor* the new vectors to add to the index\n @param[in] new_indices DLManagedTensor* vector of new indices for the new vectors. If the index\n is empty, this can be NULL to imply a continuous range `[0...n_rows)`.\n @param[inout] index IVF-SQ index to be extended\n @return cuvsError_t"] pub fn cuvsIvfSqExtend( res: cuvsResources_t, new_vectors: *mut DLManagedTensor, @@ -1771,6 +2202,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup ann_refine_c Approximate Nearest Neighbors Refinement C-API\n @{\n/\n/**\n @brief Refine nearest neighbor search.\n\n Refinement is an operation that follows an approximate NN search. The approximate search has\n already selected n_candidates neighbor candidates for each query. We narrow it down to k\n neighbors. For each query, we calculate the exact distance between the query and its\n n_candidates neighbor candidate, and select the k nearest ones.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] dataset device matrix that stores the dataset [n_rows, dims]\n @param[in] queries device matrix of the queries [n_queris, dims]\n @param[in] candidates indices of candidate vectors [n_queries, n_candidates], where\n n_candidates >= k\n @param[in] metric distance metric to use. Euclidean (L2) is used by default\n @param[out] indices device matrix that stores the refined indices [n_queries, k]\n @param[out] distances device matrix that stores the refined distances [n_queries, k]"] pub fn cuvsRefine( res: cuvsResources_t, dataset: *mut DLManagedTensor, @@ -1782,12 +2214,14 @@ unsafe extern "C" { ) -> cuvsError_t; } #[repr(u32)] +#[doc = " @brief Enum to hold which ANN algorithm is being used in the tiered index"] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsTieredIndexANNAlgo { CUVS_TIERED_INDEX_ALGO_CAGRA = 0, CUVS_TIERED_INDEX_ALGO_IVF_FLAT = 1, CUVS_TIERED_INDEX_ALGO_IVF_PQ = 2, } +#[doc = " @defgroup tiered_index_c_index Tiered Index\n @{\n/\n/**\n @brief Struct to hold address of cuvs::neighbors::tiered_index::index and its active trained\n dtype\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsTieredIndex { @@ -1809,21 +2243,31 @@ const _: () = { pub type cuvsTieredIndex_t = *mut cuvsTieredIndex; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Tiered Index\n\n @param[in] index cuvsTieredIndex_t to allocate\n @return cuvsError_t"] pub fn cuvsTieredIndexCreate(index: *mut cuvsTieredIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Tiered index\n\n @param[in] index cuvsTieredIndex_t to de-allocate"] pub fn cuvsTieredIndexDestroy(index: cuvsTieredIndex_t) -> cuvsError_t; } +#[doc = " @defgroup tiered_c_index_params Tiered Index build parameters\n @{\n/\n/**\n @brief Supplemental parameters to build a TieredIndex"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsTieredIndexParams { + #[doc = " Distance type."] pub metric: cuvsDistanceType, + #[doc = " The type of ANN algorithm we are using"] pub algo: cuvsTieredIndexANNAlgo, + #[doc = " The minimum number of rows necessary in the index to create an\nann index"] pub min_ann_rows: i64, + #[doc = " Whether or not to create a new ann index on extend, if the number\nof rows in the incremental (bfknn) portion is above min_ann_rows"] pub create_ann_index_on_extend: bool, + #[doc = " Optional parameters for building a cagra index"] pub cagra_params: cuvsCagraIndexParams_t, + #[doc = " Optional parameters for building a ivf_flat index"] pub ivf_flat_params: cuvsIvfFlatIndexParams_t, + #[doc = " Optional parameters for building a ivf-pq index"] pub ivf_pq_params: cuvsIvfPqIndexParams_t, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -1849,14 +2293,17 @@ const _: () = { pub type cuvsTieredIndexParams_t = *mut cuvsTieredIndexParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Tiered Index Params and populate with default values\n\n @param[in] index_params cuvsTieredIndexParams_t to allocate\n @return cuvsError_t"] pub fn cuvsTieredIndexParamsCreate(index_params: *mut cuvsTieredIndexParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Tiered Index params\n\n @param[in] index_params\n @return cuvsError_t"] pub fn cuvsTieredIndexParamsDestroy(index_params: cuvsTieredIndexParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @defgroup tieredindex_c_index_build Tiered index build\n @{\n/\n/**\n @brief Build a TieredIndex index with a `DLManagedTensor` which has underlying\n `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`,\n or `kDLCPU`. Also, acceptable underlying types are:\n 1. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n 2. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 16`\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Assume a populated `DLManagedTensor` type here\n DLManagedTensor dataset;\n\n // Create TieredIndex index\n cuvsTieredIndex_t index;\n cuvsError_t index_create_status = cuvsTieredIndexCreate(&index);\n\n // Create default index params\n cuvsTieredIndexParams_t index_params;\n cuvsError_t params_create_status = cuvsTieredIndexParamsCreate(&index_params);\n\n // Build the TieredIndex Index\n cuvsError_t build_status = cuvsTieredIndexBuild(res, index_params, &dataset_tensor, index);\n\n // de-allocate `index` and `res`\n cuvsError_t index_destroy_status = cuvsTieredIndexDestroy(index);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] dataset DLManagedTensor* training dataset\n @param[in] index_params Index parameters to use when building the index\n @param[out] index cuvsTieredIndex_t Newly built TieredIndex index\n @return cuvsError_t"] pub fn cuvsTieredIndexBuild( res: cuvsResources_t, index_params: cuvsTieredIndexParams_t, @@ -1866,6 +2313,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup tieredindex_c_index_search Tiered index search\n @{\n/\n/**\n @brief Search a TieredIndex index with a `DLManagedTensor`\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Assume a populated `DLManagedTensor` type here\n DLManagedTensor dataset;\n DLManagedTensor queries;\n DLManagedTensor neighbors;\n DLManagedTensor bitmap;\n\n cuvsFilter prefilter{(uintptr_t)&bitmap, BITMAP};\n\n // Search the `index` built using `cuvsTieredIndexBuild`\n cuvsError_t search_status = cuvsTieredIndexSearch(res, index, &queries, &neighbors, &distances,\n prefilter);\n\n // de-allocate `res`\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] search_params params used to the ANN index, should be one of\n cuvsCagraSearchParams_t, cuvsIvfFlatSearchParams_t, cuvsIvfPqSearchParams_t\n depending on the type of the tiered index used\n @param[in] index cuvsTieredIndex which has been returned by `cuvsTieredIndexBuild`\n @param[in] queries DLManagedTensor* queries dataset to search\n @param[out] neighbors DLManagedTensor* output `k` neighbors for queries\n @param[out] distances DLManagedTensor* output `k` distances for queries\n @param[in] prefilter cuvsFilter input prefilter that can be used\nto filter queries and neighbors based on the given bitmap."] pub fn cuvsTieredIndexSearch( res: cuvsResources_t, search_params: *mut ::std::os::raw::c_void, @@ -1878,6 +2326,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @}\n/\n/**\n @defgroup tiered_c_index_extend Tiered index extend\n @{\n/\n/**\n @brief Extend the index with the new data.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] new_vectors DLManagedTensor* the new vectors to add to the index\n @param[inout] index Tiered index to be extended\n @return cuvsError_t"] pub fn cuvsTieredIndexExtend( res: cuvsResources_t, new_vectors: *mut DLManagedTensor, @@ -1886,6 +2335,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @defgroup tiered_c_index_merge Tiered index merge\n @{\n/\n/**\n @brief Merge multiple indices together into a single index\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] index_params Index parameters to use when merging\n @param[in] indices pointers to indices to merge together\n @param[in] num_indices the number of indices to merge\n @param[out] output_index the merged index\n @return cuvsError_t"] pub fn cuvsTieredIndexMerge( res: cuvsResources_t, index_params: cuvsTieredIndexParams_t, @@ -1894,17 +2344,27 @@ unsafe extern "C" { output_index: cuvsTieredIndex_t, ) -> cuvsError_t; } +#[doc = " @brief Supplemental parameters to build Vamana Index\n\n `graph_degree`: Maximum degree of graph; corresponds to the R parameter of\n Vamana algorithm in the literature.\n `visited_size`: Maximum number of visited nodes per search during Vamana algorithm.\n Loosely corresponds to the L parameter in the literature.\n `vamana_iters`: The number of times all vectors are inserted into the graph. If > 1,\n all vectors are re-inserted to improve graph quality.\n `max_fraction`: The maximum batch size is this fraction of the total dataset size. Larger\n gives faster build but lower graph quality.\n `alpha`: Used to determine how aggressive the pruning will be."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsVamanaIndexParams { + #[doc = " Distance type."] pub metric: cuvsDistanceType, + #[doc = " Maximum degree of output graph corresponds to the R parameter in the original Vamana\n literature."] pub graph_degree: u32, + #[doc = " Maximum number of visited nodes per search corresponds to the L parameter in the Vamana\n literature"] pub visited_size: u32, + #[doc = " Number of Vamana vector insertion iterations (each iteration inserts all vectors)."] pub vamana_iters: f32, + #[doc = " Alpha for pruning parameter"] pub alpha: f32, + #[doc = " Maximum fraction of dataset inserted per batch. *\n Larger max batch decreases graph quality, but improves speed"] pub max_fraction: f32, + #[doc = " Base of growth rate of batch sizes"] pub batch_base: f32, + #[doc = " Size of candidate queue structure - should be (2^x)-1"] pub queue_size: u32, + #[doc = " Max batchsize of reverse edge processing (reduces memory footprint)"] pub reverse_batchsize: u32, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -1934,12 +2394,15 @@ const _: () = { pub type cuvsVamanaIndexParams_t = *mut cuvsVamanaIndexParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Vamana Index params, and populate with default values\n\n @param[in] params cuvsVamanaIndexParams_t to allocate\n @return cuvsError_t"] pub fn cuvsVamanaIndexParamsCreate(params: *mut cuvsVamanaIndexParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Vamana Index params\n\n @param[in] params cuvsVamanaIndexParams_t to de-allocate\n @return cuvsError_t"] pub fn cuvsVamanaIndexParamsDestroy(params: cuvsVamanaIndexParams_t) -> cuvsError_t; } +#[doc = " @brief Struct to hold address of cuvs::neighbors::vamana::index and its active trained dtype\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsVamanaIndex { @@ -1958,14 +2421,17 @@ const _: () = { pub type cuvsVamanaIndex_t = *mut cuvsVamanaIndex; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Vamana index\n\n @param[in] index cuvsVamanaIndex_t to allocate\n @return cuvsError_t"] pub fn cuvsVamanaIndexCreate(index: *mut cuvsVamanaIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Vamana index\n\n @param[in] index cuvsVamanaIndex_t to de-allocate\n @return cuvsError_t"] pub fn cuvsVamanaIndexDestroy(index: cuvsVamanaIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the dimension of the index\n\n @param[in] index cuvsVamanaIndex_t to get dimension of\n @param[out] dim pointer to dimension to set\n @return cuvsError_t"] pub fn cuvsVamanaIndexGetDims( index: cuvsVamanaIndex_t, dim: *mut ::std::os::raw::c_int, @@ -1973,6 +2439,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Build Vamana index\n\n Build the index from the dataset for efficient DiskANN search.\n\n The build uses the Vamana insertion-based algorithm to create the graph. The algorithm\n starts with an empty graph and iteratively inserts batches of nodes. Each batch involves\n performing a greedy search for each vector to be inserted, and inserting it with edges to\n all nodes traversed during the search. Reverse edges are also inserted and robustPrune is applied\n to improve graph quality. The index_params struct controls the degree of the final graph.\n\n The following distance metrics are supported:\n - L2\n\n Usage example:\n @code{.c}\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsResourcesCreate(&res);\n\n // Assume a row-major dataset [n_rows, n_cols] is defined as `float* dataset`\n cuvsVamanaIndexParams_t index_params;\n cuvsVamanaIndexParamsCreate(&index_params);\n index_params->metric = L2Expanded; // set distance metric\n cuvsVamanaIndex_t index;\n cuvsVamanaIndexCreate(&index);\n cuvsVamanaBuild(res, index_params, dataset, index);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params cuvsVamanaIndexParams_t used to build Vamana index\n @param[in] dataset DLManagedTensor* training dataset\n @param[out] index cuvsVamanaIndex_t Vamana index\n @return cuvsError_t"] pub fn cuvsVamanaBuild( res: cuvsResources_t, params: cuvsVamanaIndexParams_t, @@ -1982,6 +2449,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Save Vamana index to file\n\n Matches the file format used by the DiskANN open-source repository, allowing cross-compatibility.\n\n Serialized Index is to be used by the DiskANN open-source repository for graph search.\n\n @code{.c}\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsResourcesCreate(&res);\n\n // create an index with `cuvsVamanaBuild`\n cuvsVamanaSerialize(res, \"/path/to/index\", index, true);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename the file prefix for where the index is saved\n @param[in] index cuvsVamanaIndex_t to serialize\n @param[in] include_dataset whether to include the dataset in the serialized index\n @return cuvsError_t"] pub fn cuvsVamanaSerialize( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -1990,19 +2458,26 @@ unsafe extern "C" { ) -> cuvsError_t; } #[repr(u32)] +#[doc = " @brief Hierarchy for HNSW index when converting from CAGRA index\n\n NOTE: When the value is `NONE`, the HNSW index is built as a base-layer-only index."] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsHnswHierarchy { NONE = 0, CPU = 1, GPU = 2, } +#[doc = " Parameters for ACE (Augmented Core Extraction) graph build for HNSW.\n ACE enables building indexes for datasets too large to fit in GPU memory by:\n 1. Partitioning the dataset in core and augmented partitions using balanced k-means\n 2. Building sub-indexes for each partition independently\n 3. Concatenating sub-graphs into a final unified index"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsHnswAceParams { + #[doc = " Number of partitions for ACE partitioned build.\n\n When set to 0 (default), the number of partitions is automatically derived\n based on available host and GPU memory to maximize partition size while\n ensuring the build fits in memory.\n\n Small values might improve recall but potentially degrade performance and\n increase memory usage. The partition size is on average 2 * (n_rows /\n npartitions) * dim * sizeof(T). 2 is because of the core and augmented\n vectors. Please account for imbalance in the partition sizes (up to 3x in\n our tests).\n\n If the specified number of partitions results in partitions that exceed\n available memory, the value will be automatically increased to fit memory\n constraints and a warning will be issued."] pub npartitions: usize, + #[doc = " Directory to store ACE build artifacts (e.g., KNN graph, optimized graph).\n Used when `use_disk` is true or when the graph does not fit in memory."] pub build_dir: *const ::std::os::raw::c_char, + #[doc = " Whether to use disk-based storage for ACE build.\n When true, enables disk-based operations for memory-efficient graph construction."] pub use_disk: bool, + #[doc = " Maximum host memory to use for ACE build in GiB.\n When set to 0 (default), uses available host memory.\n Useful for testing or when running alongside other memory-intensive processes."] pub max_host_memory_gb: f64, + #[doc = " Maximum GPU memory to use for ACE build in GiB.\n When set to 0 (default), uses available GPU memory.\n Useful for testing or when running alongside other memory-intensive processes."] pub max_gpu_memory_gb: f64, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -2023,20 +2498,27 @@ const _: () = { pub type cuvsHnswAceParams_t = *mut cuvsHnswAceParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate HNSW ACE params, and populate with default values\n\n @param[in] params cuvsHnswAceParams_t to allocate\n @return cuvsError_t"] pub fn cuvsHnswAceParamsCreate(params: *mut cuvsHnswAceParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate HNSW ACE params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsHnswAceParamsDestroy(params: cuvsHnswAceParams_t) -> cuvsError_t; } #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsHnswIndexParams { pub hierarchy: cuvsHnswHierarchy, + #[doc = " Size of the candidate list during hierarchy construction when hierarchy is `CPU`"] pub ef_construction: ::std::os::raw::c_int, + #[doc = " Number of host threads to use to construct hierarchy when hierarchy is `CPU` or `GPU`.\nWhen the value is 0, the number of threads is automatically determined to the\nmaximum number of threads available.\nNOTE: When hierarchy is `GPU`, while the majority of the work is done on the GPU,\ninitialization of the HNSW index itself and some other work\nis parallelized with the help of CPU threads."] pub num_threads: ::std::os::raw::c_int, + #[doc = " HNSW M parameter: number of bi-directional links per node (used when building with ACE).\n graph_degree = m * 2, intermediate_graph_degree = m * 3."] pub M: usize, + #[doc = " Distance type for the index."] pub metric: cuvsDistanceType, + #[doc = " Optional: specify ACE parameters for building HNSW index using ACE algorithm.\n Set to nullptr for default behavior (from_cagra conversion)."] pub ace_params: cuvsHnswAceParams_t, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -2059,12 +2541,15 @@ const _: () = { pub type cuvsHnswIndexParams_t = *mut cuvsHnswIndexParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate HNSW Index params, and populate with default values\n\n @param[in] params cuvsHnswIndexParams_t to allocate\n @return cuvsError_t"] pub fn cuvsHnswIndexParamsCreate(params: *mut cuvsHnswIndexParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate HNSW Index params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsHnswIndexParamsDestroy(params: cuvsHnswIndexParams_t) -> cuvsError_t; } +#[doc = " @brief Struct to hold address of cuvs::neighbors::Hnsw::index and its active trained dtype\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsHnswIndex { @@ -2082,15 +2567,19 @@ const _: () = { pub type cuvsHnswIndex_t = *mut cuvsHnswIndex; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate HNSW index\n\n @param[in] index cuvsHnswIndex_t to allocate\n @return HnswError_t"] pub fn cuvsHnswIndexCreate(index: *mut cuvsHnswIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate HNSW index\n\n @param[in] index cuvsHnswIndex_t to de-allocate"] pub fn cuvsHnswIndexDestroy(index: cuvsHnswIndex_t) -> cuvsError_t; } +#[doc = " @defgroup hnsw_c_extend_params Parameters for extending HNSW index\n @{"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsHnswExtendParams { + #[doc = " Number of CPU threads used to extend additional vectors"] pub num_threads: ::std::os::raw::c_int, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -2103,14 +2592,17 @@ const _: () = { pub type cuvsHnswExtendParams_t = *mut cuvsHnswExtendParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate HNSW extend params, and populate with default values\n\n @param[in] params cuvsHnswExtendParams_t to allocate\n @return cuvsError_t"] pub fn cuvsHnswExtendParamsCreate(params: *mut cuvsHnswExtendParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate HNSW extend params\n\n @param[in] params cuvsHnswExtendParams_t to de-allocate\n @return cuvsError_t"] pub fn cuvsHnswExtendParamsDestroy(params: cuvsHnswExtendParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Convert a CAGRA Index to an HNSW index.\n NOTE: When hierarchy is:\n 1. `NONE`: This method uses the filesystem to write the CAGRA index in\n `/tmp/.bin` before reading it as an hnswlib index, then deleting the temporary\n file. The returned index is immutable and can only be searched by the hnswlib wrapper in cuVS,\n as the format is not compatible with the original hnswlib.\n 2. `CPU`: The returned index is mutable and can be extended with additional vectors. The\n serialized index is also compatible with the original hnswlib library.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params cuvsHnswIndexParams_t used to load Hnsw index\n @param[in] cagra_index cuvsCagraIndex_t to convert to HNSW index\n @param[out] hnsw_index cuvsHnswIndex_t to return the HNSW index\n\n @return cuvsError_t\n\n @code{.c}\n #include \n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // create a CAGRA index with `cuvsCagraBuild`\n\n // Convert the CAGRA index to an HNSW index\n cuvsHnswIndex_t hnsw_index;\n cuvsHnswIndexCreate(&hnsw_index);\n cuvsHnswIndexParams_t hnsw_params;\n cuvsHnswIndexParamsCreate(&hnsw_params);\n cuvsHnswFromCagra(res, hnsw_params, cagra_index, hnsw_index);\n\n // de-allocate `hnsw_params`, `hnsw_index` and `res`\n cuvsError_t hnsw_params_destroy_status = cuvsHnswIndexParamsDestroy(hnsw_params);\n cuvsError_t hnsw_index_destroy_status = cuvsHnswIndexDestroy(hnsw_index);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode"] pub fn cuvsHnswFromCagra( res: cuvsResources_t, params: cuvsHnswIndexParams_t, @@ -2130,6 +2622,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Build an HNSW index using ACE (Augmented Core Extraction) algorithm.\n\n ACE enables building HNSW indexes for datasets too large to fit in GPU memory by:\n 1. Partitioning the dataset using balanced k-means into core and augmented partitions\n 2. Building sub-indexes for each partition independently\n 3. Concatenating sub-graphs into a final unified index\n\n NOTE: This function requires CUDA to be available at runtime.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params cuvsHnswIndexParams_t with ACE parameters configured\n @param[in] dataset DLManagedTensor* host dataset to build index from\n @param[out] index cuvsHnswIndex_t to return the built HNSW index\n\n @return cuvsError_t\n\n @code{.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsResourcesCreate(&res);\n\n // Create ACE parameters\n cuvsHnswAceParams_t ace_params;\n cuvsHnswAceParamsCreate(&ace_params);\n ace_params->npartitions = 4;\n ace_params->use_disk = true;\n ace_params->build_dir = \"/tmp/hnsw_ace_build\";\n\n // Create index parameters\n cuvsHnswIndexParams_t params;\n cuvsHnswIndexParamsCreate(¶ms);\n params->hierarchy = GPU;\n params->ace_params = ace_params;\n params->M = 32;\n params->ef_construction = 120;\n\n // Create HNSW index\n cuvsHnswIndex_t hnsw_index;\n cuvsHnswIndexCreate(&hnsw_index);\n\n // Assume dataset is a populated DLManagedTensor with host data\n DLManagedTensor dataset;\n\n // Build the index\n cuvsHnswBuild(res, params, &dataset, hnsw_index);\n\n // Clean up\n cuvsHnswAceParamsDestroy(ace_params);\n cuvsHnswIndexParamsDestroy(params);\n cuvsHnswIndexDestroy(hnsw_index);\n cuvsResourcesDestroy(res);\n @endcode"] pub fn cuvsHnswBuild( res: cuvsResources_t, params: cuvsHnswIndexParams_t, @@ -2139,6 +2632,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Add new vectors to an HNSW index\n NOTE: The HNSW index can only be extended when the hierarchy is `CPU`\n when converting from a CAGRA index.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params cuvsHnswExtendParams_t used to extend Hnsw index\n @param[in] additional_dataset DLManagedTensor* additional dataset to extend the index\n @param[inout] index cuvsHnswIndex_t to extend\n\n @return cuvsError_t\n\n @code{.c}\n #include \n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // create an index with `cuvsCagraBuild`\n\n // Convert the CAGRA index to an HNSW index\n cuvsHnswIndex_t hnsw_index;\n cuvsHnswIndexCreate(&hnsw_index);\n cuvsHnswIndexParams_t hnsw_params;\n cuvsHnswIndexParamsCreate(&hnsw_params);\n cuvsHnswFromCagra(res, hnsw_params, cagra_index, hnsw_index);\n\n // Extend the HNSW index with additional vectors\n DLManagedTensor additional_dataset;\n cuvsHnswExtendParams_t extend_params;\n cuvsHnswExtendParamsCreate(&extend_params);\n cuvsHnswExtend(res, extend_params, additional_dataset, hnsw_index);\n\n // de-allocate `hnsw_params`, `hnsw_index`, `extend_params` and `res`\n cuvsError_t hnsw_params_destroy_status = cuvsHnswIndexParamsDestroy(hnsw_params);\n cuvsError_t hnsw_index_destroy_status = cuvsHnswIndexDestroy(hnsw_index);\n cuvsError_t extend_params_destroy_status = cuvsHnswExtendParamsDestroy(extend_params);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode"] pub fn cuvsHnswExtend( res: cuvsResources_t, params: cuvsHnswExtendParams_t, @@ -2146,6 +2640,7 @@ unsafe extern "C" { index: cuvsHnswIndex_t, ) -> cuvsError_t; } +#[doc = " @defgroup hnsw_c_search_params C API for hnswlib wrapper search params\n @{"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsHnswSearchParams { @@ -2164,14 +2659,17 @@ const _: () = { pub type cuvsHnswSearchParams_t = *mut cuvsHnswSearchParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate HNSW search params, and populate with default values\n\n @param[in] params cuvsHnswSearchParams_t to allocate\n @return cuvsError_t"] pub fn cuvsHnswSearchParamsCreate(params: *mut cuvsHnswSearchParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate HNSW search params\n\n @param[in] params cuvsHnswSearchParams_t to de-allocate\n @return cuvsError_t"] pub fn cuvsHnswSearchParamsDestroy(params: cuvsHnswSearchParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @defgroup hnsw_c_index_search C API for CUDA ANN Graph-based nearest neighbor search\n @{\n/\n/**\n @brief Search a HNSW index with a `DLManagedTensor` which has underlying\n `DLDeviceType` equal to `kDLCPU`, `kDLCUDAHost`, or `kDLCUDAManaged`.\n It is also important to note that the HNSW Index must have been built\n with the same type of `queries`, such that `index.dtype.code ==\n queries.dl_tensor.dtype.code`\n Supported types for input are:\n 1. `queries`:\n a. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n b. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8`\n c. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8`\n 2. `neighbors`: `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 64`\n 3. `distances`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`\n NOTE: When hierarchy is `NONE`, the HNSW index can only be searched by the hnswlib wrapper in\n cuVS, as the format is not compatible with the original hnswlib.\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // Assume a populated `DLManagedTensor` type here\n DLManagedTensor dataset;\n DLManagedTensor queries;\n DLManagedTensor neighbors;\n\n // Create default search params\n cuvsHnswSearchParams_t params;\n cuvsError_t params_create_status = cuvsHnswSearchParamsCreate(¶ms);\n\n // Search the `index` built using `cuvsHnswFromCagra`\n cuvsError_t search_status = cuvsHnswSearch(res, params, index, &queries, &neighbors,\n &distances);\n\n // de-allocate `params` and `res`\n cuvsError_t params_destroy_status = cuvsHnswSearchParamsDestroy(params);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params cuvsHnswSearchParams_t used to search Hnsw index\n @param[in] index cuvsHnswIndex which has been returned by `cuvsHnswFromCagra`\n @param[in] queries DLManagedTensor* queries dataset to search\n @param[out] neighbors DLManagedTensor* output `k` neighbors for queries\n @param[out] distances DLManagedTensor* output `k` distances for queries"] pub fn cuvsHnswSearch( res: cuvsResources_t, params: cuvsHnswSearchParams_t, @@ -2183,6 +2681,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Serialize a CAGRA index to a file as an hnswlib index\n NOTE: When hierarchy is `NONE`, the saved hnswlib index is immutable and can only be read by\n the hnswlib wrapper in cuVS, as the serialization format is not compatible with the original\n hnswlib. However, when hierarchy is `CPU`, the saved hnswlib index is compatible with the\n original hnswlib library.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename the name of the file to save the index\n @param[in] index cuvsHnswIndex_t to serialize\n @return cuvsError_t\n\n @code{.c}\n #include \n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // create an index with `cuvsCagraBuild`\n\n // Convert the CAGRA index to an HNSW index\n cuvsHnswIndex_t hnsw_index;\n cuvsHnswIndexCreate(&hnsw_index);\n cuvsHnswIndexParams_t hnsw_params;\n cuvsHnswIndexParamsCreate(&hnsw_params);\n cuvsHnswFromCagra(res, hnsw_params, cagra_index, hnsw_index);\n\n // Serialize the HNSW index\n cuvsHnswSerialize(res, \"/path/to/index\", hnsw_index);\n\n // de-allocate `hnsw_params`, `hnsw_index` and `res`\n cuvsError_t hnsw_params_destroy_status = cuvsHnswIndexParamsDestroy(hnsw_params);\n cuvsError_t hnsw_index_destroy_status = cuvsHnswIndexDestroy(hnsw_index);\n cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);\n @endcode"] pub fn cuvsHnswSerialize( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -2191,6 +2690,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " Load hnswlib index from file which was serialized from a HNSW index.\n NOTE: When hierarchy is `NONE`, the loaded hnswlib index is immutable, and only be read by the\n hnswlib wrapper in cuVS, as the serialization format is not compatible with the original\n hnswlib. Experimental, both the API and the serialization format are subject to change.\n\n @code{.c}\n #include \n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsError_t res_create_status = cuvsResourcesCreate(&res);\n\n // create an index with `cuvsCagraBuild`\n cuvsCagraSerializeHnswlib(res, \"/path/to/index\", index);\n\n // Load the serialized CAGRA index from file as an hnswlib index\n // The index should have the same dtype as the one used to build CAGRA the index\n cuvsHnswIndex_t hnsw_index;\n cuvsHnswIndexCreate(&hnsw_index);\n cuvsHnsWIndexParams_t hnsw_params;\n cuvsHnswIndexParamsCreate(&hnsw_params);\n hnsw_params->hierarchy = NONE;\n hnsw_index->dtype = index->dtype;\n cuvsHnswDeserialize(res, hnsw_params, \"/path/to/index\", dim, metric hnsw_index);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params cuvsHnswIndexParams_t used to load Hnsw index\n @param[in] filename the name of the file that stores the index\n @param[in] dim the dimension of the vectors in the index\n @param[in] metric the distance metric used to build the index\n @param[out] index HNSW index loaded disk"] pub fn cuvsHnswDeserialize( res: cuvsResources_t, params: cuvsHnswIndexParams_t, @@ -2202,6 +2702,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Find clusters with single-node multi-GPU k-means using host data.\n\n X, sample_weight, and centroids must be host-accessible, row-major,\n C-contiguous DLPack tensors. X and centroids must have dtype float32 or\n float64, and sample_weight must match X when provided.\n\n @param[in] res cuvsMultiGpuResources_t opaque C handle\n created by cuvsMultiGpuResourcesCreate or\n cuvsMultiGpuResourcesCreateWithDeviceIds.\n @param[in] params Parameters for KMeans model.\n @param[in] X Host training instances to cluster.\n [dim = n_samples x n_features]\n @param[in] sample_weight Optional host weights for each observation in X.\n [len = n_samples]\n @param[inout] centroids Host centroids. When init is Array, used as the\n initial cluster centers. The final generated\n centroids are copied back to this tensor.\n [dim = n_clusters x n_features]\n @param[out] inertia Sum of squared distances of samples to their\n closest cluster center.\n @param[out] n_iter Number of iterations run."] pub fn cuvsMultiGpuKMeansFit( res: cuvsResources_t, params: cuvsKMeansParams_t, @@ -2213,27 +2714,39 @@ unsafe extern "C" { ) -> cuvsError_t; } #[repr(u32)] +#[doc = " @brief Distribution mode for multi-GPU indexes"] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsMultiGpuDistributionMode { + #[doc = " Index is replicated on each device, favors throughput"] CUVS_NEIGHBORS_MG_REPLICATED = 0, + #[doc = " Index is split on several devices, favors scaling"] CUVS_NEIGHBORS_MG_SHARDED = 1, } #[repr(u32)] +#[doc = " @brief Search mode when using a replicated index"] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsMultiGpuReplicatedSearchMode { + #[doc = " Search queries are split to maintain equal load on GPUs"] CUVS_NEIGHBORS_MG_LOAD_BALANCER = 0, + #[doc = " Each search query is processed by a single GPU in a round-robin fashion"] CUVS_NEIGHBORS_MG_ROUND_ROBIN = 1, } #[repr(u32)] +#[doc = " @brief Merge mode when using a sharded index"] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsMultiGpuShardedMergeMode { + #[doc = " Search batches are merged on the root rank"] CUVS_NEIGHBORS_MG_MERGE_ON_ROOT_RANK = 0, + #[doc = " Search batches are merged in a tree reduction fashion"] CUVS_NEIGHBORS_MG_TREE_MERGE = 1, } +#[doc = " @brief Multi-GPU parameters to build CAGRA Index\n\n This structure extends the base CAGRA index parameters with multi-GPU specific settings."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsMultiGpuCagraIndexParams { + #[doc = " Base CAGRA index parameters"] pub base_params: cuvsCagraIndexParams_t, + #[doc = " Distribution mode for multi-GPU setup"] pub mode: cuvsMultiGpuDistributionMode, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -2250,22 +2763,29 @@ const _: () = { pub type cuvsMultiGpuCagraIndexParams_t = *mut cuvsMultiGpuCagraIndexParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Multi-GPU CAGRA Index params, and populate with default values\n\n @param[in] index_params cuvsMultiGpuCagraIndexParams_t to allocate\n @return cuvsError_t"] pub fn cuvsMultiGpuCagraIndexParamsCreate( index_params: *mut cuvsMultiGpuCagraIndexParams_t, ) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Multi-GPU CAGRA Index params\n\n @param[in] index_params\n @return cuvsError_t"] pub fn cuvsMultiGpuCagraIndexParamsDestroy( index_params: cuvsMultiGpuCagraIndexParams_t, ) -> cuvsError_t; } +#[doc = " @brief Multi-GPU parameters to search CAGRA index\n\n This structure extends the base CAGRA search parameters with multi-GPU specific settings."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsMultiGpuCagraSearchParams { + #[doc = " Base CAGRA search parameters"] pub base_params: cuvsCagraSearchParams_t, + #[doc = " Replicated search mode"] pub search_mode: cuvsMultiGpuReplicatedSearchMode, + #[doc = " Sharded merge mode"] pub merge_mode: cuvsMultiGpuShardedMergeMode, + #[doc = " Number of rows per batch"] pub n_rows_per_batch: i64, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -2286,16 +2806,19 @@ const _: () = { pub type cuvsMultiGpuCagraSearchParams_t = *mut cuvsMultiGpuCagraSearchParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Multi-GPU CAGRA search params, and populate with default values\n\n @param[in] params cuvsMultiGpuCagraSearchParams_t to allocate\n @return cuvsError_t"] pub fn cuvsMultiGpuCagraSearchParamsCreate( params: *mut cuvsMultiGpuCagraSearchParams_t, ) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Multi-GPU CAGRA search params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsMultiGpuCagraSearchParamsDestroy( params: cuvsMultiGpuCagraSearchParams_t, ) -> cuvsError_t; } +#[doc = " @brief Struct to hold address of cuvs::neighbors::mg_index and its active trained\n dtype"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsMultiGpuCagraIndex { @@ -2315,14 +2838,17 @@ const _: () = { pub type cuvsMultiGpuCagraIndex_t = *mut cuvsMultiGpuCagraIndex; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Multi-GPU CAGRA index\n\n @param[in] index cuvsMultiGpuCagraIndex_t to allocate\n @return cuvsError_t"] pub fn cuvsMultiGpuCagraIndexCreate(index: *mut cuvsMultiGpuCagraIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Multi-GPU CAGRA index\n\n @param[in] index cuvsMultiGpuCagraIndex_t to de-allocate\n @return cuvsError_t"] pub fn cuvsMultiGpuCagraIndexDestroy(index: cuvsMultiGpuCagraIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Build a Multi-GPU CAGRA index\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params Multi-GPU CAGRA index parameters\n @param[in] dataset_tensor DLManagedTensor* training dataset\n @param[out] index Multi-GPU CAGRA index\n @return cuvsError_t"] pub fn cuvsMultiGpuCagraBuild( res: cuvsResources_t, params: cuvsMultiGpuCagraIndexParams_t, @@ -2332,6 +2858,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Update a Multi-GPU CAGRA index with a device-padded dataset.\n\n Standard indexes are converted to device-padded indexes. Existing device-padded indexes are\n updated in place with the same layout.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] device_padded_dataset caller-owned device-padded dataset view\n @param[in,out] index Multi-GPU CAGRA index\n @return cuvsError_t"] pub fn cuvsMultiGpuCagraUpdateDataset( res: cuvsResources_t, device_padded_dataset: cuvsDataset_t, @@ -2340,6 +2867,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Search a Multi-GPU CAGRA index\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params Multi-GPU CAGRA search parameters\n @param[in] index Multi-GPU CAGRA index\n @param[in] queries_tensor DLManagedTensor* queries dataset\n @param[out] neighbors_tensor DLManagedTensor* output neighbors\n @param[out] distances_tensor DLManagedTensor* output distances\n @return cuvsError_t"] pub fn cuvsMultiGpuCagraSearch( res: cuvsResources_t, params: cuvsMultiGpuCagraSearchParams_t, @@ -2351,6 +2879,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Extend a Multi-GPU CAGRA index\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in,out] index Multi-GPU CAGRA index to extend\n @param[in] new_vectors_tensor DLManagedTensor* new vectors to add\n @param[in] new_indices_tensor DLManagedTensor* new indices (optional, can be NULL)\n @return cuvsError_t"] pub fn cuvsMultiGpuCagraExtend( res: cuvsResources_t, index: cuvsMultiGpuCagraIndex_t, @@ -2360,6 +2889,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Serialize a Multi-GPU CAGRA index to file\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] index Multi-GPU CAGRA index to serialize\n @param[in] filename Path to the output file\n @return cuvsError_t"] pub fn cuvsMultiGpuCagraSerialize( res: cuvsResources_t, index: cuvsMultiGpuCagraIndex_t, @@ -2368,6 +2898,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Deserialize a Multi-GPU CAGRA index from file\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename Path to the input file\n @param[out] index Multi-GPU CAGRA index\n @return cuvsError_t"] pub fn cuvsMultiGpuCagraDeserialize( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -2376,16 +2907,20 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Distribute a local CAGRA index to create a Multi-GPU index\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename Path to the local index file\n @param[out] index Multi-GPU CAGRA index\n @return cuvsError_t"] pub fn cuvsMultiGpuCagraDistribute( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, index: cuvsMultiGpuCagraIndex_t, ) -> cuvsError_t; } +#[doc = " @brief Multi-GPU parameters to build IVF-Flat Index\n\n This structure extends the base IVF-Flat index parameters with multi-GPU specific settings."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsMultiGpuIvfFlatIndexParams { + #[doc = " Base IVF-Flat index parameters"] pub base_params: cuvsIvfFlatIndexParams_t, + #[doc = " Distribution mode for multi-GPU setup"] pub mode: cuvsMultiGpuDistributionMode, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -2402,22 +2937,29 @@ const _: () = { pub type cuvsMultiGpuIvfFlatIndexParams_t = *mut cuvsMultiGpuIvfFlatIndexParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Multi-GPU IVF-Flat Index params, and populate with default values\n\n @param[in] index_params cuvsMultiGpuIvfFlatIndexParams_t to allocate\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfFlatIndexParamsCreate( index_params: *mut cuvsMultiGpuIvfFlatIndexParams_t, ) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Multi-GPU IVF-Flat Index params\n\n @param[in] index_params\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfFlatIndexParamsDestroy( index_params: cuvsMultiGpuIvfFlatIndexParams_t, ) -> cuvsError_t; } +#[doc = " @brief Multi-GPU parameters to search IVF-Flat index\n\n This structure extends the base IVF-Flat search parameters with multi-GPU specific settings."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsMultiGpuIvfFlatSearchParams { + #[doc = " Base IVF-Flat search parameters"] pub base_params: cuvsIvfFlatSearchParams_t, + #[doc = " Replicated search mode"] pub search_mode: cuvsMultiGpuReplicatedSearchMode, + #[doc = " Sharded merge mode"] pub merge_mode: cuvsMultiGpuShardedMergeMode, + #[doc = " Number of rows per batch"] pub n_rows_per_batch: i64, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -2438,16 +2980,19 @@ const _: () = { pub type cuvsMultiGpuIvfFlatSearchParams_t = *mut cuvsMultiGpuIvfFlatSearchParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Multi-GPU IVF-Flat search params, and populate with default values\n\n @param[in] params cuvsMultiGpuIvfFlatSearchParams_t to allocate\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfFlatSearchParamsCreate( params: *mut cuvsMultiGpuIvfFlatSearchParams_t, ) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Multi-GPU IVF-Flat search params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfFlatSearchParamsDestroy( params: cuvsMultiGpuIvfFlatSearchParams_t, ) -> cuvsError_t; } +#[doc = " @brief Struct to hold address of cuvs::neighbors::mg_index and its active\n trained dtype"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsMultiGpuIvfFlatIndex { @@ -2468,14 +3013,17 @@ const _: () = { pub type cuvsMultiGpuIvfFlatIndex_t = *mut cuvsMultiGpuIvfFlatIndex; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Multi-GPU IVF-Flat index\n\n @param[in] index cuvsMultiGpuIvfFlatIndex_t to allocate\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfFlatIndexCreate(index: *mut cuvsMultiGpuIvfFlatIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Multi-GPU IVF-Flat index\n\n @param[in] index cuvsMultiGpuIvfFlatIndex_t to de-allocate\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfFlatIndexDestroy(index: cuvsMultiGpuIvfFlatIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Build a Multi-GPU IVF-Flat index\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params Multi-GPU IVF-Flat index parameters\n @param[in] dataset_tensor DLManagedTensor* training dataset\n @param[out] index Multi-GPU IVF-Flat index\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfFlatBuild( res: cuvsResources_t, params: cuvsMultiGpuIvfFlatIndexParams_t, @@ -2485,6 +3033,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Search a Multi-GPU IVF-Flat index\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params Multi-GPU IVF-Flat search parameters\n @param[in] index Multi-GPU IVF-Flat index\n @param[in] queries_tensor DLManagedTensor* queries dataset\n @param[out] neighbors_tensor DLManagedTensor* output neighbors\n @param[out] distances_tensor DLManagedTensor* output distances\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfFlatSearch( res: cuvsResources_t, params: cuvsMultiGpuIvfFlatSearchParams_t, @@ -2496,6 +3045,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Extend a Multi-GPU IVF-Flat index\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in,out] index Multi-GPU IVF-Flat index to extend\n @param[in] new_vectors_tensor DLManagedTensor* new vectors to add\n @param[in] new_indices_tensor DLManagedTensor* new indices (optional, can be NULL)\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfFlatExtend( res: cuvsResources_t, index: cuvsMultiGpuIvfFlatIndex_t, @@ -2505,6 +3055,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Serialize a Multi-GPU IVF-Flat index to file\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] index Multi-GPU IVF-Flat index to serialize\n @param[in] filename Path to the output file\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfFlatSerialize( res: cuvsResources_t, index: cuvsMultiGpuIvfFlatIndex_t, @@ -2513,6 +3064,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Deserialize a Multi-GPU IVF-Flat index from file\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename Path to the input file\n @param[out] index Multi-GPU IVF-Flat index\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfFlatDeserialize( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -2521,16 +3073,20 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Distribute a local IVF-Flat index to create a Multi-GPU index\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename Path to the local index file\n @param[out] index Multi-GPU IVF-Flat index\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfFlatDistribute( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, index: cuvsMultiGpuIvfFlatIndex_t, ) -> cuvsError_t; } +#[doc = " @brief Multi-GPU parameters to build IVF-PQ Index\n\n This structure extends the base IVF-PQ index parameters with multi-GPU specific settings."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsMultiGpuIvfPqIndexParams { + #[doc = " Base IVF-PQ index parameters"] pub base_params: cuvsIvfPqIndexParams_t, + #[doc = " Distribution mode for multi-GPU setup"] pub mode: cuvsMultiGpuDistributionMode, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -2547,22 +3103,29 @@ const _: () = { pub type cuvsMultiGpuIvfPqIndexParams_t = *mut cuvsMultiGpuIvfPqIndexParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Multi-GPU IVF-PQ Index params, and populate with default values\n\n @param[in] index_params cuvsMultiGpuIvfPqIndexParams_t to allocate\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfPqIndexParamsCreate( index_params: *mut cuvsMultiGpuIvfPqIndexParams_t, ) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Multi-GPU IVF-PQ Index params\n\n @param[in] index_params\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfPqIndexParamsDestroy( index_params: cuvsMultiGpuIvfPqIndexParams_t, ) -> cuvsError_t; } +#[doc = " @brief Multi-GPU parameters to search IVF-PQ index\n\n This structure extends the base IVF-PQ search parameters with multi-GPU specific settings."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsMultiGpuIvfPqSearchParams { + #[doc = " Base IVF-PQ search parameters"] pub base_params: cuvsIvfPqSearchParams_t, + #[doc = " Replicated search mode"] pub search_mode: cuvsMultiGpuReplicatedSearchMode, + #[doc = " Sharded merge mode"] pub merge_mode: cuvsMultiGpuShardedMergeMode, + #[doc = " Number of rows per batch"] pub n_rows_per_batch: i64, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -2583,16 +3146,19 @@ const _: () = { pub type cuvsMultiGpuIvfPqSearchParams_t = *mut cuvsMultiGpuIvfPqSearchParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Multi-GPU IVF-PQ search params, and populate with default values\n\n @param[in] params cuvsMultiGpuIvfPqSearchParams_t to allocate\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfPqSearchParamsCreate( params: *mut cuvsMultiGpuIvfPqSearchParams_t, ) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Multi-GPU IVF-PQ search params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfPqSearchParamsDestroy( params: cuvsMultiGpuIvfPqSearchParams_t, ) -> cuvsError_t; } +#[doc = " @brief Struct to hold address of cuvs::neighbors::mg_index and its active trained\n dtype"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsMultiGpuIvfPqIndex { @@ -2612,14 +3178,17 @@ const _: () = { pub type cuvsMultiGpuIvfPqIndex_t = *mut cuvsMultiGpuIvfPqIndex; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Multi-GPU IVF-PQ index\n\n @param[in] index cuvsMultiGpuIvfPqIndex_t to allocate\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfPqIndexCreate(index: *mut cuvsMultiGpuIvfPqIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Multi-GPU IVF-PQ index\n\n @param[in] index cuvsMultiGpuIvfPqIndex_t to de-allocate\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfPqIndexDestroy(index: cuvsMultiGpuIvfPqIndex_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Build a Multi-GPU IVF-PQ index\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params Multi-GPU IVF-PQ index parameters\n @param[in] dataset_tensor DLManagedTensor* training dataset\n @param[out] index Multi-GPU IVF-PQ index\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfPqBuild( res: cuvsResources_t, params: cuvsMultiGpuIvfPqIndexParams_t, @@ -2629,6 +3198,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Search a Multi-GPU IVF-PQ index\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params Multi-GPU IVF-PQ search parameters\n @param[in] index Multi-GPU IVF-PQ index\n @param[in] queries_tensor DLManagedTensor* queries dataset\n @param[out] neighbors_tensor DLManagedTensor* output neighbors\n @param[out] distances_tensor DLManagedTensor* output distances\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfPqSearch( res: cuvsResources_t, params: cuvsMultiGpuIvfPqSearchParams_t, @@ -2640,6 +3210,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Extend a Multi-GPU IVF-PQ index\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in,out] index Multi-GPU IVF-PQ index to extend\n @param[in] new_vectors_tensor DLManagedTensor* new vectors to add\n @param[in] new_indices_tensor DLManagedTensor* new indices (optional, can be NULL)\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfPqExtend( res: cuvsResources_t, index: cuvsMultiGpuIvfPqIndex_t, @@ -2649,6 +3220,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Serialize a Multi-GPU IVF-PQ index to file\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] index Multi-GPU IVF-PQ index to serialize\n @param[in] filename Path to the output file\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfPqSerialize( res: cuvsResources_t, index: cuvsMultiGpuIvfPqIndex_t, @@ -2657,6 +3229,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Deserialize a Multi-GPU IVF-PQ index from file\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename Path to the input file\n @param[out] index Multi-GPU IVF-PQ index\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfPqDeserialize( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -2665,6 +3238,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Distribute a local IVF-PQ index to create a Multi-GPU index\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] filename Path to the local index file\n @param[out] index Multi-GPU IVF-PQ index\n @return cuvsError_t"] pub fn cuvsMultiGpuIvfPqDistribute( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -2672,19 +3246,29 @@ unsafe extern "C" { ) -> cuvsError_t; } #[repr(u32)] +#[doc = " @brief Solver algorithm for PCA eigen decomposition."] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsPcaSolver { + #[doc = " Covariance + divide-and-conquer eigen decomposition"] CUVS_PCA_COV_EIG_DQ = 0, + #[doc = " Covariance + Jacobi eigen decomposition"] CUVS_PCA_COV_EIG_JACOBI = 1, } +#[doc = " @brief Parameters for PCA decomposition."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsPcaParams { + #[doc = " Number of principal components to keep."] pub n_components: ::std::os::raw::c_int, + #[doc = " If false, data passed to fit are overwritten and running fit(X).transform(X) will\n not yield the expected results; use fit_transform(X) instead."] pub copy: bool, + #[doc = " When true the component vectors are multiplied by the square root of n_samples and then\n divided by the singular values to ensure uncorrelated outputs with unit component-wise\n variances."] pub whiten: bool, + #[doc = " Solver algorithm to use."] pub algorithm: cuvsPcaSolver, + #[doc = " Tolerance for singular values (used by Jacobi solver)."] pub tol: f32, + #[doc = " Number of iterations for the power method (Jacobi solver)."] pub n_iterations: ::std::os::raw::c_int, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -2705,14 +3289,17 @@ const _: () = { pub type cuvsPcaParams_t = *mut cuvsPcaParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate PCA params and populate with default values.\n\n @param[out] params cuvsPcaParams_t to allocate\n @return cuvsError_t"] pub fn cuvsPcaParamsCreate(params: *mut cuvsPcaParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate PCA params.\n\n @param[in] params cuvsPcaParams_t to de-allocate\n @return cuvsError_t"] pub fn cuvsPcaParamsDestroy(params: cuvsPcaParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Perform PCA fit operation.\n\n Computes the principal components, explained variances, singular values, and column means\n from the input data.\n\n The layout of `input` (C-contiguous / row-major or F-contiguous / col-major) is detected\n from its DLPack strides; `components` must use the same layout as `input`.\n\n @code {.c}\n #include \n #include \n\n // Create cuvsResources_t\n cuvsResources_t res;\n cuvsResourcesCreate(&res);\n\n // Create PCA params\n cuvsPcaParams_t params;\n cuvsPcaParamsCreate(¶ms);\n params->n_components = 2;\n\n // Assume populated DLManagedTensor objects (float32, device memory)\n DLManagedTensor input; // [n_rows x n_cols] (C- or F-contiguous)\n DLManagedTensor components; // [n_components x n_cols] (same layout as input)\n DLManagedTensor explained_var; // [n_components]\n DLManagedTensor explained_var_ratio; // [n_components]\n DLManagedTensor singular_vals; // [n_components]\n DLManagedTensor mu; // [n_cols]\n DLManagedTensor noise_vars; // [1] (scalar)\n\n cuvsPcaFit(res, params, &input, &components, &explained_var,\n &explained_var_ratio, &singular_vals, &mu, &noise_vars, false);\n\n // Cleanup\n cuvsPcaParamsDestroy(params);\n cuvsResourcesDestroy(res);\n @endcode\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params PCA parameters\n @param[inout] input input data [n_rows x n_cols] (C- or F-contiguous, float32, device)\n @param[out] components principal components [n_components x n_cols] (same layout as input)\n @param[out] explained_var explained variances [n_components] (float32, device)\n @param[out] explained_var_ratio explained variance ratios [n_components] (float32, device)\n @param[out] singular_vals singular values [n_components] (float32, device)\n @param[out] mu column means [n_cols] (float32, device)\n @param[out] noise_vars noise variance [1] (float32, device)\n @param[in] flip_signs_based_on_U whether to determine signs by U (true) or V.T (false)\n @return cuvsError_t"] pub fn cuvsPcaFit( res: cuvsResources_t, params: cuvsPcaParams_t, @@ -2728,6 +3315,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Perform PCA fit and transform in a single operation.\n\n Computes the principal components and transforms the input data into the eigenspace.\n The layout of `input` (C- or F-contiguous) is detected from its DLPack strides; all\n other matrix tensors must use the same layout.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params PCA parameters\n @param[inout] input input data [n_rows x n_cols] (C- or F-contiguous, float32, device)\n @param[out] trans_input transformed data [n_rows x n_components] (same layout as input)\n @param[out] components principal components [n_components x n_cols] (same layout as input)\n @param[out] explained_var explained variances [n_components] (float32, device)\n @param[out] explained_var_ratio explained variance ratios [n_components] (float32, device)\n @param[out] singular_vals singular values [n_components] (float32, device)\n @param[out] mu column means [n_cols] (float32, device)\n @param[out] noise_vars noise variance [1] (float32, device)\n @param[in] flip_signs_based_on_U whether to determine signs by U (true) or V.T (false)\n @return cuvsError_t"] pub fn cuvsPcaFitTransform( res: cuvsResources_t, params: cuvsPcaParams_t, @@ -2744,6 +3332,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Perform PCA transform operation.\n\n Transforms the input data into the eigenspace using previously computed principal components.\n The layout of `input` (C- or F-contiguous) is detected from its DLPack strides; all other\n matrix tensors must use the same layout.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params PCA parameters\n @param[inout] input data to transform [n_rows x n_cols] (C- or F-contiguous, float32, device)\n @param[in] components principal components [n_components x n_cols] (same layout as input)\n @param[in] singular_vals singular values [n_components] (float32, device)\n @param[in] mu column means [n_cols] (float32, device)\n @param[out] trans_input transformed data [n_rows x n_components] (same layout as input)\n @return cuvsError_t"] pub fn cuvsPcaTransform( res: cuvsResources_t, params: cuvsPcaParams_t, @@ -2756,6 +3345,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Perform PCA inverse transform operation.\n\n Transforms data from the eigenspace back to the original space.\n The layout of `trans_input` (C- or F-contiguous) is detected from its DLPack strides;\n all other matrix tensors must use the same layout.\n\n @param[in] res cuvsResources_t opaque C handle\n @param[in] params PCA parameters\n @param[in] trans_input transformed data [n_rows x n_components] (C- or F-contiguous,\n float32, device)\n @param[in] components principal components [n_components x n_cols] (same layout as trans_input)\n @param[in] singular_vals singular values [n_components] (float32, device)\n @param[in] mu column means [n_cols] (float32, device)\n @param[out] output reconstructed data [n_rows x n_cols] (same layout as trans_input)\n @return cuvsError_t"] pub fn cuvsPcaInverseTransform( res: cuvsResources_t, params: cuvsPcaParams_t, @@ -2767,12 +3357,14 @@ unsafe extern "C" { ) -> cuvsError_t; } #[repr(u32)] +#[doc = " @defgroup preprocessing_c_binary C API for Binary Quantizer\n @{\n/\n/**\n @brief In the cuvsBinaryQuantizerTransform function, a bit is set if the corresponding element in\n the dataset vector is greater than the corresponding element in the threshold vector. The mean\n and sampling_median thresholds are calculated separately for each dimension.\n"] #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] pub enum cuvsBinaryQuantizerThreshold { ZERO = 0, MEAN = 1, SAMPLING_MEDIAN = 2, } +#[doc = " @brief Binary quantizer parameters."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsBinaryQuantizerParams { @@ -2793,13 +3385,16 @@ const _: () = { pub type cuvsBinaryQuantizerParams_t = *mut cuvsBinaryQuantizerParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Binary Quantizer params, and populate with default values\n\n @param[in] params cuvsBinaryQuantizerParams_t to allocate\n @return cuvsError_t"] pub fn cuvsBinaryQuantizerParamsCreate(params: *mut cuvsBinaryQuantizerParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Binary Quantizer params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsBinaryQuantizerParamsDestroy(params: cuvsBinaryQuantizerParams_t) -> cuvsError_t; } +#[doc = " @brief Defines and stores threshold for quantization upon training\n\n The quantization is performed by a linear mapping of an interval in the\n float data type to the full range of the quantized int type."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsBinaryQuantizer { @@ -2818,14 +3413,17 @@ const _: () = { pub type cuvsBinaryQuantizer_t = *mut cuvsBinaryQuantizer; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Binary Quantizer and populate with default values\n\n @param[in] quantizer cuvsBinaryQuantizer_t to allocate\n @return cuvsError_t"] pub fn cuvsBinaryQuantizerCreate(quantizer: *mut cuvsBinaryQuantizer_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Binary Quantizer\n\n @param[in] quantizer\n @return cuvsError_t"] pub fn cuvsBinaryQuantizerDestroy(quantizer: cuvsBinaryQuantizer_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Trains a binary quantizer to be used later for quantizing the dataset.\n\n @param[in] res raft resource\n @param[in] params configure binary quantizer, e.g. threshold\n @param[in] dataset a row-major host or device matrix\n @param[out] quantizer trained binary quantizer"] pub fn cuvsBinaryQuantizerTrain( res: cuvsResources_t, params: cuvsBinaryQuantizerParams_t, @@ -2835,6 +3433,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Applies binary quantization transform to the given dataset\n\n This applies binary quantization to a dataset, changing any positive\n values to a bitwise 1. This is useful for searching with the\n BitwiseHamming distance type.\n\n @param[in] res raft resource\n @param[in] dataset a row-major host or device matrix to transform\n @param[out] out a row-major host or device matrix to store transformed data"] pub fn cuvsBinaryQuantizerTransform( res: cuvsResources_t, dataset: *mut DLManagedTensor, @@ -2843,6 +3442,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Applies binary quantization transform to the given dataset\n\n This applies binary quantization to a dataset, changing any values that are larger than the\n threshold specified in the param to a bitwise 1. This is useful for searching with the\n BitwiseHamming distance type.\n\n @param[in] res raft resource\n @param[in] quantizer binary quantizer\n @param[in] dataset a row-major host or device matrix to transform\n @param[out] out a row-major host or device matrix to store transformed data"] pub fn cuvsBinaryQuantizerTransformWithParams( res: cuvsResources_t, quantizer: cuvsBinaryQuantizer_t, @@ -2850,17 +3450,27 @@ unsafe extern "C" { out: *mut DLManagedTensor, ) -> cuvsError_t; } +#[doc = " @defgroup preprocessing_c_pq C API for Product Quantizer\n @{\n/\n/**\n @brief Product quantizer parameters."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsProductQuantizerParams { + #[doc = " The bit length of the vector element after compression by PQ.\n\n Possible values: within [4, 16].\n\n Hint: the smaller the 'pq_bits', the smaller the index size and the better the search\n performance, but the lower the recall."] pub pq_bits: u32, + #[doc = " The dimensionality of the vector after compression by PQ.\n When zero, an optimal value is selected using a heuristic.\n\n TODO: at the moment `dim` must be a multiple `pq_dim`."] pub pq_dim: u32, + #[doc = " Whether to use subspaces for product quantization (PQ).\n When true, one PQ codebook is used for each subspace. Otherwise, a single\n PQ codebook is used."] pub use_subspaces: bool, + #[doc = " Whether to use Vector Quantization (KMeans) before product quantization (PQ).\n When true, VQ is used before PQ. When false, only product quantization is used."] pub use_vq: bool, + #[doc = " Vector Quantization (VQ) codebook size - number of \"coarse cluster centers\".\n When zero, an optimal value is selected using a heuristic.\n When one, only product quantization is used."] pub vq_n_centers: u32, + #[doc = " The number of iterations searching for kmeans centers (both VQ & PQ phases)."] pub kmeans_n_iters: u32, + #[doc = " The type of kmeans algorithm to use for PQ training."] pub pq_kmeans_type: cuvsKMeansType, + #[doc = " The max number of data points to use per PQ code during PQ codebook training. Using more data\n points per PQ code may increase the quality of PQ codebook but may also increase the build\n time. We will use `pq_n_centers * max_train_points_per_pq_code` training\n points to train each PQ codebook."] pub max_train_points_per_pq_code: u32, + #[doc = " The max number of data points to use per VQ cluster."] pub max_train_points_per_vq_cluster: u32, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] @@ -2895,14 +3505,17 @@ const _: () = { pub type cuvsProductQuantizerParams_t = *mut cuvsProductQuantizerParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Product Quantizer params, and populate with default values\n\n @param[in] params cuvsProductQuantizerParams_t to allocate\n @return cuvsError_t"] pub fn cuvsProductQuantizerParamsCreate( params: *mut cuvsProductQuantizerParams_t, ) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Product Quantizer params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsProductQuantizerParamsDestroy(params: cuvsProductQuantizerParams_t) -> cuvsError_t; } +#[doc = " @brief Defines and stores product quantizer upon training\n\n The quantization is performed by a linear mapping of an interval in the\n float data type to the full range of the quantized int type."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsProductQuantizer { @@ -2921,14 +3534,17 @@ const _: () = { pub type cuvsProductQuantizer_t = *mut cuvsProductQuantizer; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Product Quantizer\n\n @param[in] quantizer cuvsProductQuantizer_t to allocate\n @return cuvsError_t"] pub fn cuvsProductQuantizerCreate(quantizer: *mut cuvsProductQuantizer_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Product Quantizer\n\n @param[in] quantizer\n @return cuvsError_t"] pub fn cuvsProductQuantizerDestroy(quantizer: cuvsProductQuantizer_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Builds a product quantizer to be used later for quantizing the dataset.\n\n @param[in] res raft resource\n @param[in] params Parameters for product quantizer training\n @param[in] dataset a row-major host or device matrix\n @param[out] quantizer trained product quantizer"] pub fn cuvsProductQuantizerBuild( res: cuvsResources_t, params: cuvsProductQuantizerParams_t, @@ -2938,6 +3554,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Applies product quantization transform to the given dataset\n\n This applies product quantization to a dataset.\n\n @param[in] res raft resource\n @param[in] quantizer product quantizer\n @param[in] dataset a row-major host or device matrix to transform\n @param[out] codes_out a row-major device matrix to store transformed data\n @param[out] vq_labels a device vector to store VQ labels.\n Optional, can be NULL."] pub fn cuvsProductQuantizerTransform( res: cuvsResources_t, quantizer: cuvsProductQuantizer_t, @@ -2948,6 +3565,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Applies product quantization inverse transform to the given quantized codes\n\n This applies product quantization inverse transform to the given quantized codes.\n\n @param[in] res raft resource\n @param[in] quantizer product quantizer\n @param[in] pq_codes a row-major device matrix of quantized codes\n @param[out] out a row-major device matrix to store the original data\n @param[out] vq_labels a device vector containing the VQ labels when VQ is used.\n Optional, can be NULL."] pub fn cuvsProductQuantizerInverseTransform( res: cuvsResources_t, quantizer: cuvsProductQuantizer_t, @@ -2958,6 +3576,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the bit length of the vector element after compression by PQ.\n\n @param[in] quantizer product quantizer\n @param[out] pq_bits bit length of the vector element after compression by PQ"] pub fn cuvsProductQuantizerGetPqBits( quantizer: cuvsProductQuantizer_t, pq_bits: *mut u32, @@ -2965,6 +3584,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the dimensionality of the vector after compression by PQ.\n\n @param[in] quantizer product quantizer\n @param[out] pq_dim dimensionality of the vector after compression by PQ"] pub fn cuvsProductQuantizerGetPqDim( quantizer: cuvsProductQuantizer_t, pq_dim: *mut u32, @@ -2972,6 +3592,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the PQ codebook.\n\n @param[in] quantizer product quantizer\n @param[out] pq_codebook PQ codebook"] pub fn cuvsProductQuantizerGetPqCodebook( quantizer: cuvsProductQuantizer_t, pq_codebook: *mut DLManagedTensor, @@ -2979,6 +3600,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the VQ codebook.\n\n @param[in] quantizer product quantizer\n @param[out] vq_codebook VQ codebook"] pub fn cuvsProductQuantizerGetVqCodebook( quantizer: cuvsProductQuantizer_t, vq_codebook: *mut DLManagedTensor, @@ -2986,6 +3608,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get the encoded dimension of the quantized dataset.\n\n @param[in] quantizer product quantizer\n @param[out] encoded_dim encoded dimension of the quantized dataset"] pub fn cuvsProductQuantizerGetEncodedDim( quantizer: cuvsProductQuantizer_t, encoded_dim: *mut u32, @@ -2993,11 +3616,13 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Get whether VQ is used.\n\n @param[in] quantizer product quantizer\n @param[out] use_vq whether VQ is used"] pub fn cuvsProductQuantizerGetUseVq( quantizer: cuvsProductQuantizer_t, use_vq: *mut bool, ) -> cuvsError_t; } +#[doc = " @defgroup preprocessing_c_scalar C API for Scalar Quantizer\n @{\n/\n/**\n @brief Scalar quantizer parameters."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsScalarQuantizerParams { @@ -3015,13 +3640,16 @@ const _: () = { pub type cuvsScalarQuantizerParams_t = *mut cuvsScalarQuantizerParams; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Scalar Quantizer params, and populate with default values\n\n @param[in] params cuvsScalarQuantizerParams_t to allocate\n @return cuvsError_t"] pub fn cuvsScalarQuantizerParamsCreate(params: *mut cuvsScalarQuantizerParams_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Scalar Quantizer params\n\n @param[in] params\n @return cuvsError_t"] pub fn cuvsScalarQuantizerParamsDestroy(params: cuvsScalarQuantizerParams_t) -> cuvsError_t; } +#[doc = " @brief Defines and stores scalar for quantisation upon training\n\n The quantization is performed by a linear mapping of an interval in the\n float data type to the full range of the quantized int type."] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsScalarQuantizer { @@ -3040,14 +3668,17 @@ const _: () = { pub type cuvsScalarQuantizer_t = *mut cuvsScalarQuantizer; unsafe extern "C" { #[must_use] + #[doc = " @brief Allocate Scalar Quantizer and populate with default values\n\n @param[in] quantizer cuvsScalarQuantizer_t to allocate\n @return cuvsError_t"] pub fn cuvsScalarQuantizerCreate(quantizer: *mut cuvsScalarQuantizer_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief De-allocate Scalar Quantizer\n\n @param[in] quantizer\n @return cuvsError_t"] pub fn cuvsScalarQuantizerDestroy(quantizer: cuvsScalarQuantizer_t) -> cuvsError_t; } unsafe extern "C" { #[must_use] + #[doc = " @brief Trains a scalar quantizer to be used later for quantizing the dataset.\n\n @param[in] res raft resource\n @param[in] params configure scalar quantizer, e.g. quantile\n @param[in] dataset a row-major host or device matrix\n @param[out] quantizer trained scalar quantizer"] pub fn cuvsScalarQuantizerTrain( res: cuvsResources_t, params: cuvsScalarQuantizerParams_t, @@ -3057,6 +3688,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Applies quantization transform to given dataset\n\n @param[in] res raft resource\n @param[in] quantizer a scalar quantizer\n @param[in] dataset a row-major host or device matrix to transform\n @param[out] out a row-major host or device matrix to store transformed data"] pub fn cuvsScalarQuantizerTransform( res: cuvsResources_t, quantizer: cuvsScalarQuantizer_t, @@ -3066,6 +3698,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Perform inverse quantization step on previously quantized dataset\n\n Note that depending on the chosen data types train dataset the conversion is\n not lossless.\n\n @param[in] res raft resource\n @param[in] quantizer a scalar quantizer\n @param[in] dataset a row-major host or device matrix\n @param[out] out a row-major host or device matrix\n"] pub fn cuvsScalarQuantizerInverseTransform( res: cuvsResources_t, quantizer: cuvsScalarQuantizer_t, @@ -3075,6 +3708,7 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] + #[doc = " @brief Select the k smallest values from a flat device array of n candidates.\n\n Treats `in_val` as a matrix of shape [1, n] and selects the `k` smallest\n float values. `out_idx` receives the int64 column positions of the selected\n values in [0, n), so the caller can recover per-segment identity as:\n\n segment_index = out_idx[j] / segment_k\n position_in_segment = out_idx[j] % segment_k\n\n @param[in] res cuvsResources_t handle\n @param[in] in_val DLManagedTensor* shape [1, n], float32, device memory\n @param[out] out_val DLManagedTensor* shape [1, k], float32, device memory\n @param[out] out_idx DLManagedTensor* shape [1, k], int64, device memory\n @return cuvsError_t"] pub fn cuvsSelectK( res: cuvsResources_t, in_val: *mut DLManagedTensor, diff --git a/rust/cuvs/src/neighbors/cagra/index.rs b/rust/cuvs/src/neighbors/cagra/index.rs index a93c90ca4d..244e80357f 100644 --- a/rust/cuvs/src/neighbors/cagra/index.rs +++ b/rust/cuvs/src/neighbors/cagra/index.rs @@ -6,7 +6,7 @@ use std::marker::PhantomData; use std::path::Path; -use super::{CagraError, IndexParams, SearchParams}; +use super::{CagraError, IndexParams, MergeParams, SearchParams}; use crate::dataset::private::Sealed as _; use crate::dataset::{CuvsDataset, Dataset, DatasetKind, DatasetView}; use crate::dlpack::{AsDlTensor, AsDlTensorMut, DLTensorView, DLTensorViewMut}; @@ -124,6 +124,152 @@ impl<'d> Index<'d> { Ok(Index { handle, _dataset: PhantomData }) } + /// Merges multiple CAGRA indices into a new index backed by `merged_dataset`. + /// + /// The caller must have already concatenated every input index's dataset (in + /// `indices` order) into `merged_dataset`, and computed `offsets` as the + /// cumulative row counts of each input index: `offsets[i]` is the row at + /// which `indices[i]`'s rows start in `merged_dataset`, and + /// `offsets[indices.len()]` must equal `merged_dataset`'s total row count. + /// See [`merged_dataset_offsets`] for the bitset-filtered case. + /// + /// The returned [`Index`] borrows `merged_dataset` for `'a` and cannot + /// outlive it, mirroring [`Index::update_dataset`]. + pub fn merge<'a, D>( + res: &Resources, + params: &IndexParams, + indices: &[&Index<'_>], + merged_dataset: &'a D, + offsets: &[i64], + ) -> Result> + where + D: CuvsDataset + ?Sized, + { + Self::merge_impl(res, params, None, indices, None, merged_dataset, offsets) + } + + /// Merges multiple CAGRA indices, applying a row-level bitset `filter`. + /// + /// `merged_dataset` must already contain only the rows surviving `filter` + /// (in `indices` order); use [`merged_dataset_offsets`] to compute the + /// per-index row offsets within it. + pub fn merge_filtered<'a, D>( + res: &Resources, + params: &IndexParams, + indices: &[&Index<'_>], + filter: &Filter<'_, Bitset>, + merged_dataset: &'a D, + offsets: &[i64], + ) -> Result> + where + D: CuvsDataset + ?Sized, + { + Self::merge_impl(res, params, None, indices, Some(filter), merged_dataset, offsets) + } + + /// Merges multiple CAGRA indices using explicit [`MergeParams`]. + /// + /// See [`Index::merge`] for the `merged_dataset`/`offsets` contract. + pub fn merge_with_params<'a, D>( + res: &Resources, + params: &IndexParams, + merge_params: &MergeParams, + indices: &[&Index<'_>], + merged_dataset: &'a D, + offsets: &[i64], + ) -> Result> + where + D: CuvsDataset + ?Sized, + { + Self::merge_impl(res, params, Some(merge_params), indices, None, merged_dataset, offsets) + } + + /// Merges multiple CAGRA indices using explicit [`MergeParams`] and a + /// row-level bitset `filter`. + /// + /// See [`Index::merge`] and [`Index::merge_filtered`] for the + /// `merged_dataset`/`offsets` contract. + #[allow(clippy::too_many_arguments)] + pub fn merge_filtered_with_params<'a, D>( + res: &Resources, + params: &IndexParams, + merge_params: &MergeParams, + indices: &[&Index<'_>], + filter: &Filter<'_, Bitset>, + merged_dataset: &'a D, + offsets: &[i64], + ) -> Result> + where + D: CuvsDataset + ?Sized, + { + Self::merge_impl( + res, + params, + Some(merge_params), + indices, + Some(filter), + merged_dataset, + offsets, + ) + } + + #[allow(clippy::too_many_arguments)] + fn merge_impl<'a, D>( + res: &Resources, + params: &IndexParams, + merge_params: Option<&MergeParams>, + indices: &[&Index<'_>], + filter: Option<&Filter<'_, Bitset>>, + merged_dataset: &'a D, + offsets: &[i64], + ) -> Result> + where + D: CuvsDataset + ?Sized, + { + if offsets.len() != indices.len() + 1 { + return Err(CagraError::Validation(format!( + "offsets must have indices.len() + 1 ({}) entries, got {}", + indices.len() + 1, + offsets.len() + ))); + } + + let mut raw_indices: Vec = + indices.iter().map(|index| index.handle.raw()).collect(); + let merged_dataset = merged_dataset.raw_dataset_handle(); + let handle = IndexHandle::new()?; + + with_filter(filter, |c_filter| { + check_cuvs(unsafe { + match merge_params { + Some(merge_params) => ffi::cuvsCagraMergeWithParams( + res.handle(), + params.handle(), + merge_params.handle(), + raw_indices.as_mut_ptr(), + raw_indices.len(), + c_filter, + merged_dataset, + offsets.as_ptr(), + handle.raw(), + ), + None => ffi::cuvsCagraMerge( + res.handle(), + params.handle(), + raw_indices.as_mut_ptr(), + raw_indices.len(), + c_filter, + merged_dataset, + offsets.as_ptr(), + handle.raw(), + ), + } + }) + })?; + + Ok(Index { handle, _dataset: PhantomData }) + } + /// Searches the index for the `k` nearest neighbors of each query. /// /// `queries`, `neighbors`, and `distances` must reside in device memory and @@ -366,6 +512,43 @@ impl DeserializedIndex { } } +/// Computes per-index row offsets within a to-be-built merge buffer. +/// +/// `cuvsCagraMerge`/`Index::merge` require the caller to have already +/// concatenated every input index's dataset (in `indices` order, applying +/// `filter` if any) into a single buffer, and to know each index's starting +/// row within it. For an unfiltered merge those offsets are just the +/// cumulative row counts of `indices`, so this function is unnecessary. For a +/// bitset `filter`, the number of surviving rows per index cannot be derived +/// any other way, so call this first. +/// +/// Returns a `Vec` of `indices.len() + 1` entries: entry `i` is the row at +/// which `indices[i]`'s surviving rows must start in the merged buffer; the +/// last entry is the total row count of the merged buffer. +pub fn merged_dataset_offsets( + res: &Resources, + indices: &[&Index<'_>], + filter: Option<&Filter<'_, Bitset>>, +) -> Result> { + let mut raw_indices: Vec = + indices.iter().map(|index| index.handle.raw()).collect(); + let mut offsets = vec![0i64; indices.len() + 1]; + + with_filter(filter, |c_filter| { + check_cuvs(unsafe { + ffi::cuvsCagraMergedDatasetOffsets( + res.handle(), + raw_indices.as_mut_ptr(), + raw_indices.len(), + c_filter, + offsets.as_mut_ptr(), + ) + }) + })?; + + Ok(offsets) +} + fn search_impl( handle: &IndexHandle, res: &Resources, @@ -419,7 +602,7 @@ fn serialize_to_hnswlib_impl(handle: &IndexHandle, res: &Resources, filename: &P #[cfg(test)] mod tests { use super::*; - use crate::dataset::PaddedDataset; + use crate::dataset::{DatasetView, PaddedDataset}; use crate::neighbors::filters::{Bitset, Filter}; use crate::test_utils::DeviceTensor; use ndarray::s; @@ -794,4 +977,148 @@ mod tests { .expect_err("serialize should reject paths with interior NUL"); assert!(matches!(err, CagraError::InvalidPath(_)), "expected InvalidPath, got {err:?}"); } + + /// Build two indices, merge them over a caller-concatenated buffer, and + /// verify every row still finds itself as its own nearest neighbor. + #[test] + fn test_cagra_merge() { + let res = Resources::new().unwrap(); + let build_params = IndexParams::builder().build().unwrap(); + + let n1 = 128usize; + let n2 = 96usize; + let dataset_a = + ndarray::Array::::random((n1, N_FEATURES), Uniform::new(0., 1.0).unwrap()); + let dataset_b = + ndarray::Array::::random((n2, N_FEATURES), Uniform::new(0., 1.0).unwrap()); + + let device_a = DeviceTensor::from_host(&res, &dataset_a).unwrap(); + let device_b = DeviceTensor::from_host(&res, &dataset_b).unwrap(); + + let index_a = + Index::build(&res, &build_params, &device_a).expect("failed to build index_a"); + let index_b = + Index::build(&res, &build_params, &device_b).expect("failed to build index_b"); + + let merged_host = + ndarray::concatenate(ndarray::Axis(0), &[dataset_a.view(), dataset_b.view()]).unwrap(); + let merged_device = DeviceTensor::from_host(&res, &merged_host).unwrap(); + let merged_view = DatasetView::new(&res, &merged_device).unwrap(); + + let offsets: Vec = vec![0, n1 as i64, (n1 + n2) as i64]; + + let merged_index = + Index::merge(&res, &build_params, &[&index_a, &index_b], &merged_view, &offsets) + .expect("merge failed"); + + search_and_verify_self_neighbors(&res, &merged_index, &merged_host, 4, 10); + } + + /// Same as `test_cagra_merge`, but exercising `merge_with_params` with an + /// explicit `MergeParams` instance. + #[test] + fn test_cagra_merge_with_params() { + let res = Resources::new().unwrap(); + let build_params = IndexParams::builder().build().unwrap(); + let merge_params = MergeParams::builder().build().unwrap(); + + let n1 = 64usize; + let n2 = 64usize; + let dataset_a = + ndarray::Array::::random((n1, N_FEATURES), Uniform::new(0., 1.0).unwrap()); + let dataset_b = + ndarray::Array::::random((n2, N_FEATURES), Uniform::new(0., 1.0).unwrap()); + + let device_a = DeviceTensor::from_host(&res, &dataset_a).unwrap(); + let device_b = DeviceTensor::from_host(&res, &dataset_b).unwrap(); + + let index_a = + Index::build(&res, &build_params, &device_a).expect("failed to build index_a"); + let index_b = + Index::build(&res, &build_params, &device_b).expect("failed to build index_b"); + + let merged_host = + ndarray::concatenate(ndarray::Axis(0), &[dataset_a.view(), dataset_b.view()]).unwrap(); + let merged_device = DeviceTensor::from_host(&res, &merged_host).unwrap(); + let merged_view = DatasetView::new(&res, &merged_device).unwrap(); + + let offsets: Vec = vec![0, n1 as i64, (n1 + n2) as i64]; + + let merged_index = Index::merge_with_params( + &res, + &build_params, + &merge_params, + &[&index_a, &index_b], + &merged_view, + &offsets, + ) + .expect("merge_with_params failed"); + + search_and_verify_self_neighbors(&res, &merged_index, &merged_host, 4, 10); + } + + #[test] + fn merge_rejects_mismatched_offsets_length() { + let res = Resources::new().unwrap(); + let build_params = IndexParams::builder().build().unwrap(); + let dataset = ndarray::Array::::random( + (N_DATAPOINTS, N_FEATURES), + Uniform::new(0., 1.0).unwrap(), + ); + let device = DeviceTensor::from_host(&res, &dataset).unwrap(); + let index = Index::build(&res, &build_params, &device).unwrap(); + let view = DatasetView::new(&res, &device).unwrap(); + + // Only one index but two offsets entries provided instead of the required two. + let err = Index::merge(&res, &build_params, &[&index], &view, &[0]) + .expect_err("offsets.len() must equal indices.len() + 1"); + assert!(matches!(err, CagraError::Validation(_)), "unexpected error: {err:?}"); + } + + #[test] + fn merged_dataset_offsets_without_filter_is_cumulative_sizes() { + let res = Resources::new().unwrap(); + let build_params = IndexParams::builder().build().unwrap(); + + let n1 = 96usize; + let n2 = 32usize; + let dataset_a = + ndarray::Array::::random((n1, N_FEATURES), Uniform::new(0., 1.0).unwrap()); + let dataset_b = + ndarray::Array::::random((n2, N_FEATURES), Uniform::new(0., 1.0).unwrap()); + let device_a = DeviceTensor::from_host(&res, &dataset_a).unwrap(); + let device_b = DeviceTensor::from_host(&res, &dataset_b).unwrap(); + + let index_a = Index::build(&res, &build_params, &device_a).unwrap(); + let index_b = Index::build(&res, &build_params, &device_b).unwrap(); + + let offsets = merged_dataset_offsets(&res, &[&index_a, &index_b], None).unwrap(); + assert_eq!(offsets, vec![0, n1 as i64, (n1 + n2) as i64]); + } + + #[test] + fn merged_dataset_offsets_reflects_bitset_filter() { + let res = Resources::new().unwrap(); + let build_params = IndexParams::builder().build().unwrap(); + + let n_datapoints = 64; + let dataset = ndarray::Array::::random( + (n_datapoints, N_FEATURES), + Uniform::new(0., 1.0).unwrap(), + ); + let dataset_device = DeviceTensor::from_host(&res, &dataset).unwrap(); + let index = Index::build(&res, &build_params, &dataset_device).unwrap(); + + // Keep only the first half of the rows. + let n_words = n_datapoints.div_ceil(32); + let mut bitset_host = ndarray::Array::::zeros(ndarray::Ix1(n_words)); + for i in 0..n_datapoints / 2 { + bitset_host[i / 32] |= 1u32 << (i % 32); + } + let bitset = DeviceTensor::from_host(&res, &bitset_host).unwrap(); + let filter = Filter::::new(&bitset).unwrap(); + + let offsets = merged_dataset_offsets(&res, &[&index], Some(&filter)).unwrap(); + assert_eq!(offsets, vec![0, (n_datapoints / 2) as i64]); + } } diff --git a/rust/cuvs/src/neighbors/cagra/mod.rs b/rust/cuvs/src/neighbors/cagra/mod.rs index c8ab865e02..3711c8a6bc 100644 --- a/rust/cuvs/src/neighbors/cagra/mod.rs +++ b/rust/cuvs/src/neighbors/cagra/mod.rs @@ -22,8 +22,8 @@ mod params; pub use crate::dataset::{CuvsDataset, Dataset, DatasetKind, DatasetView, PaddedDataset}; pub use crate::neighbors::filters::{Bitset, Filter}; -pub use index::{DeserializedIndex, Index}; -pub use params::{IndexParams, SearchParams}; +pub use index::{DeserializedIndex, Index, merged_dataset_offsets}; +pub use params::{IndexParams, MergeParams, SearchParams}; use crate::dlpack::DLPackError; use crate::error::LibraryError; @@ -104,6 +104,38 @@ impl From for SearchAlgo { } } +/// Algorithm used to merge multiple CAGRA indices into one. +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +#[non_exhaustive] +pub enum MergeAlgo { + /// Automatically select the best merge algorithm. + Auto, + /// Fast hierarchical merge ("Fastener"). + Fastener, + /// Rebuild the graph from scratch over the merged dataset. + Rebuild, +} + +impl From for ffi::cuvsCagraMergeAlgo { + fn from(v: MergeAlgo) -> Self { + match v { + MergeAlgo::Auto => Self::CUVS_CAGRA_MERGE_AUTO, + MergeAlgo::Fastener => Self::CUVS_CAGRA_MERGE_FASTENER, + MergeAlgo::Rebuild => Self::CUVS_CAGRA_MERGE_REBUILD, + } + } +} + +impl From for MergeAlgo { + fn from(v: ffi::cuvsCagraMergeAlgo) -> Self { + match v { + ffi::cuvsCagraMergeAlgo::CUVS_CAGRA_MERGE_AUTO => Self::Auto, + ffi::cuvsCagraMergeAlgo::CUVS_CAGRA_MERGE_FASTENER => Self::Fastener, + ffi::cuvsCagraMergeAlgo::CUVS_CAGRA_MERGE_REBUILD => Self::Rebuild, + } + } +} + /// Hash-table mode used during search. #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] #[non_exhaustive] diff --git a/rust/cuvs/src/neighbors/cagra/params.rs b/rust/cuvs/src/neighbors/cagra/params.rs index 4fcc3d18af..fce71215b6 100644 --- a/rust/cuvs/src/neighbors/cagra/params.rs +++ b/rust/cuvs/src/neighbors/cagra/params.rs @@ -20,7 +20,7 @@ use bon::bon; use crate::distance::DistanceType; use crate::error::check_cuvs; -use super::{CagraError, GraphBuildAlgo, HashMode, SearchAlgo}; +use super::{CagraError, GraphBuildAlgo, HashMode, MergeAlgo, SearchAlgo}; #[derive(Debug)] enum RequestedGraphBuild { @@ -398,6 +398,93 @@ impl Drop for SearchParams { } } +// --------------------------------------------------------------------------- +// MergeParams +// --------------------------------------------------------------------------- + +/// Parameters controlling how physical CAGRA indices are merged. +/// +/// ```ignore +/// use cuvs::neighbors::cagra::{MergeAlgo, MergeParams}; +/// +/// let params = MergeParams::builder().algo(MergeAlgo::Fastener).build()?; +/// ``` +pub struct MergeParams { + handle: ffi::cuvsCagraMergeParams_t, +} + +#[bon] +impl MergeParams { + #[builder] + #[allow(clippy::too_many_arguments)] + pub fn new( + algo: Option, + levels: Option, + root_fanout: Option, + lower_fanout: Option, + leader_fraction: Option, + max_leaders: Option, + leaf_size: Option, + leaf_degree: Option, + ) -> Result { + let params = Self::create_handle()?; + + unsafe { + if let Some(v) = algo { + (*params.handle).algo = v.into(); + } + if let Some(v) = levels { + (*params.handle).levels = v; + } + if let Some(v) = root_fanout { + (*params.handle).root_fanout = v; + } + if let Some(v) = lower_fanout { + (*params.handle).lower_fanout = v; + } + if let Some(v) = leader_fraction { + (*params.handle).leader_fraction = v; + } + if let Some(v) = max_leaders { + (*params.handle).max_leaders = v; + } + if let Some(v) = leaf_size { + (*params.handle).leaf_size = v; + } + if let Some(v) = leaf_degree { + (*params.handle).leaf_degree = v; + } + } + + Ok(params) + } +} + +impl MergeParams { + /// Allocate parameters populated with the AUTO-algorithm defaults. + fn create_handle() -> Result { + let mut handle = ptr::null_mut(); + check_cuvs(unsafe { ffi::cuvsCagraMergeParamsCreate(&mut handle) })?; + Ok(Self { handle }) + } + + pub(super) fn handle(&self) -> ffi::cuvsCagraMergeParams_t { + self.handle + } +} + +impl fmt::Debug for MergeParams { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("MergeParams").field(unsafe { &*self.handle }).finish() + } +} + +impl Drop for MergeParams { + fn drop(&mut self) { + let _ = unsafe { ffi::cuvsCagraMergeParamsDestroy(self.handle) }; + } +} + // --------------------------------------------------------------------------- // Tests // --------------------------------------------------------------------------- @@ -537,6 +624,40 @@ mod tests { assert!(err.to_string().contains("512")); } + #[test] + fn merge_params_all_defaults() { + let params = MergeParams::builder().build().unwrap(); + unsafe { + assert_eq!((*params.handle).algo, ffi::cuvsCagraMergeAlgo::CUVS_CAGRA_MERGE_AUTO); + } + } + + #[test] + fn merge_params_with_values() { + let params = MergeParams::builder() + .algo(MergeAlgo::Fastener) + .levels(3) + .root_fanout(4) + .lower_fanout(2) + .leader_fraction(0.1) + .max_leaders(8) + .leaf_size(16) + .leaf_degree(32) + .build() + .unwrap(); + + unsafe { + assert_eq!((*params.handle).algo, ffi::cuvsCagraMergeAlgo::CUVS_CAGRA_MERGE_FASTENER); + assert_eq!((*params.handle).levels, 3); + assert_eq!((*params.handle).root_fanout, 4); + assert_eq!((*params.handle).lower_fanout, 2); + assert_eq!((*params.handle).leader_fraction, 0.1); + assert_eq!((*params.handle).max_leaders, 8); + assert_eq!((*params.handle).leaf_size, 16); + assert_eq!((*params.handle).leaf_degree, 32); + } + } + #[test] fn search_params_rejects_small_hash_with_multi_cta() { let err = SearchParams::builder()