diff --git a/c/include/cuvs/core/all.h b/c/include/cuvs/core/all.h index 91958177c4..a3885b29ad 100644 --- a/c/include/cuvs/core/all.h +++ b/c/include/cuvs/core/all.h @@ -11,6 +11,7 @@ #include #include #include +#include #include diff --git a/c/include/cuvs/core/dataset.h b/c/include/cuvs/core/dataset.h new file mode 100644 index 0000000000..78c96d6bac --- /dev/null +++ b/c/include/cuvs/core/dataset.h @@ -0,0 +1,92 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @brief Generic dataset layout kind for C API dataset handles. + */ +typedef enum { + CUVS_DATASET_LAYOUT_STANDARD = 0, + CUVS_DATASET_LAYOUT_PADDED = 1 +} cuvsDatasetLayout_t; + +/** + * @brief Generic opaque dataset handle used by C APIs. + * + * `addr` points to implementation-owned storage. + * `dtype` carries dataset element type in DLPack format. + * `layout` differentiates standard vs padded layouts. + */ +typedef struct { + uintptr_t addr; + DLDataType dtype; + cuvsDatasetLayout_t layout; +} cuvsDataset; +typedef cuvsDataset* cuvsDataset_t; + +/** + * @brief Layout-specific aliases that sort well lexicographically in docs. + */ +typedef cuvsDataset cuvsDatasetPadded; +typedef cuvsDataset* cuvsDatasetPadded_t; +typedef cuvsDataset cuvsDatasetStandard; +typedef cuvsDataset* cuvsDatasetStandard_t; + +/** + * @brief Generic storage kind for operation-specific dataset storage. + */ +typedef enum { + CUVS_DATASET_STORAGE_KIND_EXTENDED = 0, + CUVS_DATASET_STORAGE_KIND_MERGED = 1 +} cuvsDatasetStorageKind_t; + +/** + * @brief Generic opaque storage handle for dataset-backed operation outputs. + * + * Used by operations like CAGRA extend/merge where storage shape is identical + * but semantics differ by operation kind. + */ +typedef struct { + uintptr_t addr; + DLDataType dtype; + cuvsDatasetStorageKind_t kind; +} cuvsDatasetStorage; +typedef cuvsDatasetStorage* cuvsDatasetStorage_t; + +typedef struct cuvsCagraIndex* cuvsCagraIndex_t; + +CUVS_EXPORT cuvsError_t cuvsDatasetMakePadded(cuvsResources_t res, + DLManagedTensor* dataset, + cuvsDatasetPadded_t* padded_dataset); + +CUVS_EXPORT cuvsError_t cuvsDatasetPaddedDestroy(cuvsDatasetPadded_t padded_dataset); + +CUVS_EXPORT cuvsError_t cuvsMakeExtendedStorage(cuvsResources_t res, + DLManagedTensor* additional_dataset, + cuvsCagraIndex_t index, + cuvsDatasetStorage_t* extended_storage); + +CUVS_EXPORT cuvsError_t cuvsMakeMergedStorage(cuvsResources_t res, + cuvsCagraIndex_t* indices, + size_t num_indices, + cuvsFilter filter, + cuvsDatasetStorage_t* merged_storage); + +CUVS_EXPORT cuvsError_t cuvsDatasetStorageDestroy(cuvsDatasetStorage_t dataset_storage); + +#ifdef __cplusplus +} +#endif diff --git a/c/include/cuvs/neighbors/cagra.h b/c/include/cuvs/neighbors/cagra.h index 22809da37e..42da0b7943 100644 --- a/c/include/cuvs/neighbors/cagra.h +++ b/c/include/cuvs/neighbors/cagra.h @@ -1,11 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include +#include #include #include #include @@ -211,12 +212,6 @@ struct cuvsCagraIndexParams { enum cuvsCagraGraphBuildAlgo build_algo; /** Number of Iterations to run if building with NN_DESCENT */ size_t nn_descent_niter; - /** - * Optional: specify compression parameters if compression is desired. - * - * NOTE: this is experimental new API, consider it unsafe. - */ - cuvsCagraCompressionParams_t compression; /** * Optional: specify graph build params based on build_algo * - IVF_PQ: cuvsIvfPqParams_t @@ -466,13 +461,15 @@ CUVS_EXPORT cuvsError_t cuvsCagraSearchParamsDestroy(cuvsCagraSearchParams_t par */ /** - * @brief Struct to hold address of cuvs::neighbors::cagra::index and its active trained dtype + * @brief Struct holding the CAGRA index storage address and vector element dtype (DLPack-style) * + * Matches the usual cuVS C index pattern (`addr` + `dtype`). \p addr points at implementation-owned + * storage (not always a bare `cagra::index*`); free only via \ref cuvsCagraIndexDestroy. \p dtype + * describes index vector elements for queries and template dispatch. */ -typedef struct { +typedef struct cuvsCagraIndex { uintptr_t addr; DLDataType dtype; - } cuvsCagraIndex; typedef cuvsCagraIndex* cuvsCagraIndex_t; @@ -561,6 +558,37 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexGetDataset(cuvsCagraIndex_t index, DLManag */ CUVS_EXPORT cuvsError_t cuvsCagraIndexGetGraph(cuvsCagraIndex_t index, DLManagedTensor* graph); +/** + * @brief Attach a padded dataset to a standard CAGRA index to make it searchable. + * + * The function converts the index to the padded-search-ready form using C++ API + * `attach_padded_dataset_for_search`. + * Caller retains ownership of \p padded_dataset and must keep it alive while \p index uses it. + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] padded_dataset padded dataset handle created by \ref cuvsDatasetMakePadded + * @param[inout] index CAGRA index handle + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsCagraAttachPaddedDatasetForSearch(cuvsResources_t res, + cuvsDatasetPadded_t padded_dataset, + cuvsCagraIndex_t index); + +/** + * @brief Attach a device dataset to a host-built index, converting it into a device index. + * + * This is the explicit C equivalent of C++ `attach_device_dataset_on_host_index`. + * The `device_dataset` must be device-compatible memory. + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] device_dataset device dataset tensor to attach + * @param[in,out] index host CAGRA index handle; replaced in-place with device index handle + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsCagraAttachDeviceDatasetOnHostIndex(cuvsResources_t res, + DLManagedTensor* device_dataset, + cuvsCagraIndex_t index); + /** * @} */ @@ -646,6 +674,7 @@ CUVS_EXPORT cuvsError_t cuvsCagraBuild(cuvsResources_t res, CUVS_EXPORT cuvsError_t cuvsCagraExtend(cuvsResources_t res, cuvsCagraExtendParams_t params, DLManagedTensor* additional_dataset, + cuvsDatasetStorage_t extended_dataset, cuvsCagraIndex_t index); /** @@ -784,10 +813,14 @@ CUVS_EXPORT cuvsError_t cuvsCagraSerializeToHnswlib(cuvsResources_t res, * * @param[in] res cuvsResources_t opaque C handle * @param[in] filename the name of the file that stores the index + * @param[in] deserialize_layout target index layout to deserialize into (padded or standard) * @param[inout] index cuvsCagraIndex_t CAGRA index loaded from disk. This index needs to be already * created with cuvsCagraIndexCreate. */ -CUVS_EXPORT cuvsError_t cuvsCagraDeserialize(cuvsResources_t res, const char* filename, cuvsCagraIndex_t index); +CUVS_EXPORT cuvsError_t cuvsCagraDeserialize(cuvsResources_t res, + const char* filename, + cuvsDatasetLayout_t deserialize_layout, + cuvsCagraIndex_t index); /** * Load index from a dataset and graph @@ -895,6 +928,7 @@ CUVS_EXPORT cuvsError_t cuvsCagraMerge(cuvsResources_t res, cuvsCagraIndex_t* indices, size_t num_indices, cuvsFilter filter, + cuvsDatasetStorage_t merged_dataset, cuvsCagraIndex_t output_index); /** diff --git a/c/src/neighbors/c_api_box.hpp b/c/src/neighbors/c_api_box.hpp new file mode 100644 index 0000000000..d38f25d269 --- /dev/null +++ b/c/src/neighbors/c_api_box.hpp @@ -0,0 +1,32 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +namespace cuvs::neighbors::c_api::detail { + +struct owner_record { + void* owner; + void (*destroy_owner)(void*); +}; + +template +static void destroy_typed_owner(void* owner) +{ + delete reinterpret_cast(owner); +} + +template +static auto make_owner_record(OwnerT* owner) -> owner_record +{ + return owner_record{owner, &destroy_typed_owner}; +} + +[[maybe_unused]] static void destroy_owner_record(owner_record rec) +{ + if (rec.destroy_owner != nullptr) { rec.destroy_owner(rec.owner); } +} + +} // namespace cuvs::neighbors::c_api::detail diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 004b810c78..7407dbf196 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -1,14 +1,19 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include #include #include +#include +#include +#include +#include #include #include +#include #include #include #include @@ -19,15 +24,442 @@ #include #include #include - #include "../core/exceptions.hpp" #include "../core/interop.hpp" #include "cagra.hpp" +#include "c_api_box.hpp" #include namespace { +/** + * Heap-allocated bundle for the C API: owns only `cagra::index`. + * Lives behind `cuvsCagraIndex::addr` via `sg_cagra_c_api_index_box`. + */ +template +struct cuvs_cagra_c_api_index_lifetime_holder { + cuvs::neighbors::cagra::index idx; +}; + +template +struct cagra_c_api_padded_dataset_holder { + std::unique_ptr> dataset_owner{nullptr}; +}; + +template +struct cagra_c_api_extended_dataset_holder { + cuvs::neighbors::cagra::extended_dataset_storage storage; +}; + +template +struct cagra_c_api_merged_dataset_holder { + cuvs::neighbors::cagra::merged_dataset_storage storage; +}; + +/** Owns how to delete co-located index storage; `cuvsCagraIndex::addr` points here. */ +struct sg_cagra_c_api_index_box { + void* index_ptr; + enum class dataset_layout : uint8_t { device_padded, device_standard, host_padded, host_standard } layout; + cuvs::neighbors::c_api::detail::owner_record owner_rec; +}; + +template +constexpr auto sg_cagra_index_layout_from_view() +{ + if constexpr (cuvs::neighbors::is_device_standard_dataset_view_v) { + return sg_cagra_c_api_index_box::dataset_layout::device_standard; + } else if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + return sg_cagra_c_api_index_box::dataset_layout::device_padded; + } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { + return sg_cagra_c_api_index_box::dataset_layout::host_standard; + } else { + return sg_cagra_c_api_index_box::dataset_layout::host_padded; + } +} + +template +static auto require_padded_index(sg_cagra_c_api_index_box* box, + const char* null_handle_err, + const char* unsupported_layout_err) + -> cuvs::neighbors::cagra::device_padded_index* +{ + if (box == nullptr) { RAFT_FAIL("%s", null_handle_err); } + if (box->layout != sg_cagra_c_api_index_box::dataset_layout::device_padded) { + RAFT_FAIL("%s", unsupported_layout_err); + } + return reinterpret_cast*>(box->index_ptr); +} + +template +static void with_index_by_layout(sg_cagra_c_api_index_box* box, + const char* null_handle_err, + const char* host_not_allowed_err, + Fn&& fn) +{ + RAFT_EXPECTS(box != nullptr, "%s", null_handle_err); + switch (box->layout) { + case sg_cagra_c_api_index_box::dataset_layout::device_padded: { + auto* idx = + reinterpret_cast*>(box->index_ptr); + fn(*idx); + break; + } + case sg_cagra_c_api_index_box::dataset_layout::host_padded: { + if constexpr (AllowHost) { + auto* idx = + reinterpret_cast*>(box->index_ptr); + fn(*idx); + } else { + RAFT_FAIL("%s", host_not_allowed_err); + } + break; + } + case sg_cagra_c_api_index_box::dataset_layout::device_standard: { + auto* idx = + reinterpret_cast*>(box->index_ptr); + fn(*idx); + break; + } + case sg_cagra_c_api_index_box::dataset_layout::host_standard: { + if constexpr (AllowHost) { + auto* idx = + reinterpret_cast*>(box->index_ptr); + fn(*idx); + } else { + RAFT_FAIL("%s", host_not_allowed_err); + } + break; + } + } +} + +template +static void merge_indices_for_layout( + raft::resources* res_ptr, + cuvs::neighbors::cagra::index_params const& params_cpp, + std::vector*>& index_ptrs, + cuvsFilter filter, + cuvs::neighbors::cagra::merged_dataset_storage& merge_storage, + cuvsCagraIndex_t output_index) +{ + if (filter.type == NO_FILTER) { + auto merged_idx = cuvs::neighbors::cagra::merge(*res_ptr, params_cpp, index_ptrs, merge_storage); + 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); + } else if (filter.type == BITSET) { + int64_t merged_row_count = 0; + for (auto* idx_ptr : index_ptrs) { + merged_row_count += static_cast(idx_ptr->size()); + } + 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); + auto merged_idx = cuvs::neighbors::cagra::merge( + *res_ptr, params_cpp, index_ptrs, merge_storage, bitset_filter_obj); + 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); + } else { + RAFT_FAIL("Unsupported filter type: BITMAP"); + } +} + +template +static auto convert_opaque_indices_to_concrete_types(cuvsCagraIndex_t* indices, size_t num_indices) + -> std::vector*> +{ + std::vector*> index_ptrs; + index_ptrs.reserve(num_indices); + for (size_t i = 0; i < num_indices; ++i) { + auto* box = reinterpret_cast(indices[i]->addr); + RAFT_EXPECTS(box != nullptr, "cuvsCagraMerge: null index handle"); + index_ptrs.push_back( + reinterpret_cast*>(box->index_ptr)); + } + return index_ptrs; +} + +template +static void with_dataset_view_for_layout(raft::resources* res_ptr, + DLManagedTensor* dataset_tensor, + sg_cagra_c_api_index_box::dataset_layout layout, + const char* err_prefix, + const char* host_not_allowed_err, + Fn&& fn) +{ + auto dataset = dataset_tensor->dl_tensor; + if (layout == sg_cagra_c_api_index_box::dataset_layout::device_padded) { + if (cuvs::core::is_dlpack_device_compatible(dataset)) { + using mdspan_type = raft::device_matrix_view; + auto mds = cuvs::core::from_dlpack(dataset_tensor); + auto ds_view = cuvs::neighbors::make_device_padded_dataset_view(*res_ptr, mds); + fn(ds_view); + return; + } else if (cuvs::core::is_dlpack_host_compatible(dataset)) { + if constexpr (!AllowHost) { RAFT_FAIL("%s", host_not_allowed_err); } + using mdspan_type = raft::host_matrix_view; + auto mds = cuvs::core::from_dlpack(dataset_tensor); + auto ds_view = cuvs::neighbors::make_host_padded_dataset_view(mds); + fn(ds_view); + return; + } + } else if (layout == sg_cagra_c_api_index_box::dataset_layout::device_standard) { + if (cuvs::core::is_dlpack_device_compatible(dataset)) { + using mdspan_type = raft::device_matrix_view; + auto mds = cuvs::core::from_dlpack(dataset_tensor); + auto ds_view = cuvs::neighbors::make_device_standard_dataset_view(mds); + fn(ds_view); + return; + } else if (cuvs::core::is_dlpack_host_compatible(dataset)) { + if constexpr (!AllowHost) { RAFT_FAIL("%s", host_not_allowed_err); } + using mdspan_type = raft::host_matrix_view; + auto mds = cuvs::core::from_dlpack(dataset_tensor); + auto ds_view = cuvs::neighbors::make_host_standard_dataset_view(mds); + fn(ds_view); + return; + } + } else { + RAFT_FAIL("%s: unsupported index layout for dataset view dispatch", err_prefix); + } + RAFT_FAIL("%s: dataset must have host- or device-compatible memory", err_prefix); +} + +template +static void compute_ivfpq_shape_from_indices(cuvsCagraIndex_t* indices, + size_t num_indices, + int64_t* total_size, + int64_t* dim) +{ + auto* first_box = reinterpret_cast(indices[0]->addr); + // Caller validates non-null boxes and uniform layout for all indices. + auto* first_idx_ptr = + reinterpret_cast*>(first_box->index_ptr); + *dim = first_idx_ptr->dim(); + *total_size = 0; + for (size_t i = 0; i < num_indices; ++i) { + auto* box = reinterpret_cast(indices[i]->addr); + auto* idx_ptr = + reinterpret_cast*>(box->index_ptr); + *total_size += static_cast(idx_ptr->size()); + } +} + +template +static void bind_index_lifetime_holder_to_C_index( + cuvsCagraIndex_t out, + DLDataType dtype, + cuvs_cagra_c_api_index_lifetime_holder* holder) +{ + auto* box = new sg_cagra_c_api_index_box{&holder->idx, + sg_cagra_index_layout_from_view(), + cuvs::neighbors::c_api::detail::make_owner_record(holder)}; + out->addr = reinterpret_cast(box); + out->dtype = dtype; +} + +template +static void wrap_CPP_index_in_lifetime_holder_and_bind_to_C_index( + cuvsCagraIndex_t out, + DLDataType dtype, + cuvs::neighbors::cagra::index* raw) +{ + auto* holder = new cuvs_cagra_c_api_index_lifetime_holder{std::move(*raw)}; + delete raw; + bind_index_lifetime_holder_to_C_index(out, dtype, holder); +} + +static void destroy_sg_cagra_c_api_box(uintptr_t addr) +{ + if (addr == 0) { return; } + auto* box = reinterpret_cast(addr); + cuvs::neighbors::c_api::detail::destroy_owner_record(box->owner_rec); + delete box; +} + +template +static void destroy_padded_dataset_typed(uintptr_t addr) +{ + delete reinterpret_cast*>(addr); +} + +template +static void destroy_extended_dataset_typed(uintptr_t addr) +{ + delete reinterpret_cast*>(addr); +} + +template +static void destroy_merged_dataset_typed(uintptr_t addr) +{ + delete reinterpret_cast*>(addr); +} + +template +static void make_padded_dataset(raft::resources* res_ptr, + DLManagedTensor* dataset_tensor, + cuvsDatasetPadded_t* output_padded_dataset) +{ + auto dataset = dataset_tensor->dl_tensor; + auto* out = new cuvsDatasetPadded{0, dataset.dtype, CUVS_DATASET_LAYOUT_PADDED}; + + if (cuvs::core::is_dlpack_device_compatible(dataset)) { + using mdspan_type = raft::device_matrix_view; + auto mds = cuvs::core::from_dlpack(dataset_tensor); + auto owner = cuvs::neighbors::make_device_padded_dataset(*res_ptr, mds); + auto* holder = new cagra_c_api_padded_dataset_holder{std::move(owner)}; + out->addr = reinterpret_cast(holder); + } else if (cuvs::core::is_dlpack_host_compatible(dataset)) { + using mdspan_type = raft::host_matrix_view; + auto mds = cuvs::core::from_dlpack(dataset_tensor); + auto owner = cuvs::neighbors::make_device_padded_dataset(*res_ptr, mds); + auto* holder = new cagra_c_api_padded_dataset_holder{std::move(owner)}; + out->addr = reinterpret_cast(holder); + } else { + delete out; + RAFT_FAIL("cuvsDatasetMakePadded: dataset must have host- or device-compatible memory"); + } + + *output_padded_dataset = out; +} + +template +static void attach_padded_dataset_for_search(raft::resources* res_ptr, + cuvsDatasetPadded_t padded_dataset, + cuvsCagraIndex_t index) +{ + RAFT_EXPECTS(padded_dataset != nullptr, "cuvsCagraAttachPaddedDatasetForSearch: null padded dataset"); + RAFT_EXPECTS(index != nullptr, "cuvsCagraAttachPaddedDatasetForSearch: null index handle"); + RAFT_EXPECTS(index->addr != 0, "cuvsCagraAttachPaddedDatasetForSearch: null index storage"); + + auto* box = reinterpret_cast(index->addr); + RAFT_EXPECTS(box->layout == sg_cagra_c_api_index_box::dataset_layout::device_standard, + "cuvsCagraAttachPaddedDatasetForSearch: index must be standard layout"); + + auto* padded_holder = + reinterpret_cast*>(padded_dataset->addr); + RAFT_EXPECTS(padded_holder != nullptr && padded_holder->dataset_owner != nullptr, + "cuvsCagraAttachPaddedDatasetForSearch: padded dataset handle is uninitialized"); + + auto* standard_idx = + reinterpret_cast*>(box->index_ptr); + + auto padded_idx = cuvs::neighbors::cagra::attach_padded_dataset_for_search( + *res_ptr, *standard_idx, padded_holder->dataset_owner->as_dataset_view()); + + auto* holder = + new cuvs_cagra_c_api_index_lifetime_holder>{ + std::move(padded_idx)}; + + destroy_sg_cagra_c_api_box(index->addr); + index->addr = 0; + bind_index_lifetime_holder_to_C_index>( + index, index->dtype, holder); +} + +template +static void attach_device_dataset_on_host_index(raft::resources* res_ptr, + DLManagedTensor* device_dataset_tensor, + cuvsCagraIndex_t index) +{ + RAFT_EXPECTS(index != nullptr, "cuvsCagraAttachDeviceDatasetOnHostIndex: null index handle"); + RAFT_EXPECTS(index->addr != 0, "cuvsCagraAttachDeviceDatasetOnHostIndex: null index storage"); + auto* box = reinterpret_cast(index->addr); + RAFT_EXPECTS(box->layout == sg_cagra_c_api_index_box::dataset_layout::host_padded || + box->layout == sg_cagra_c_api_index_box::dataset_layout::host_standard, + "cuvsCagraAttachDeviceDatasetOnHostIndex: index must be host layout"); + RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(device_dataset_tensor->dl_tensor), + "cuvsCagraAttachDeviceDatasetOnHostIndex: dataset must be device-compatible"); + + if (box->layout == sg_cagra_c_api_index_box::dataset_layout::host_padded) { + auto* host_idx = + reinterpret_cast*>(box->index_ptr); + using mdspan_type = raft::device_matrix_view; + auto mds = cuvs::core::from_dlpack(device_dataset_tensor); + auto device_dataset_view = cuvs::neighbors::make_device_padded_dataset_view(*res_ptr, mds); + auto device_idx = cuvs::neighbors::cagra::attach_device_dataset_on_host_index( + *res_ptr, *host_idx, device_dataset_view); + auto* holder = + new cuvs_cagra_c_api_index_lifetime_holder>{ + std::move(device_idx)}; + destroy_sg_cagra_c_api_box(index->addr); + index->addr = 0; + bind_index_lifetime_holder_to_C_index>( + index, index->dtype, holder); + } else { + auto* host_idx = + reinterpret_cast*>(box->index_ptr); + using mdspan_type = raft::device_matrix_view; + auto mds = cuvs::core::from_dlpack(device_dataset_tensor); + auto device_dataset_view = cuvs::neighbors::make_device_standard_dataset_view(mds); + auto device_idx = cuvs::neighbors::cagra::attach_device_dataset_on_host_index( + *res_ptr, *host_idx, device_dataset_view); + auto* holder = + new cuvs_cagra_c_api_index_lifetime_holder>{ + std::move(device_idx)}; + destroy_sg_cagra_c_api_box(index->addr); + index->addr = 0; + bind_index_lifetime_holder_to_C_index>( + index, index->dtype, holder); + } +} + +template +static void make_extended_storage(raft::resources* res_ptr, + DLManagedTensor* additional_dataset_tensor, + cuvsCagraIndex_t index, + cuvsDatasetStorage_t* output_extended_storage) +{ + auto* out = new cuvsDatasetStorage{ + 0, index->dtype, CUVS_DATASET_STORAGE_KIND_EXTENDED}; + auto* box = reinterpret_cast(index->addr); + RAFT_EXPECTS(box != nullptr, "cuvsMakeExtendedStorage: null index handle"); + RAFT_EXPECTS(box->layout == sg_cagra_c_api_index_box::dataset_layout::device_padded || + box->layout == sg_cagra_c_api_index_box::dataset_layout::device_standard, + "cuvsMakeExtendedStorage: host indices are unsupported; attach device dataset " + "to host index first"); + with_index_by_layout( + box, + "cuvsMakeExtendedStorage: null index handle", + "cuvsMakeExtendedStorage: host indices are unsupported; attach device dataset to host " + "index first", + [&](auto& idx) { + auto build_storage = [&](auto ds_view) { + using index_t = std::decay_t; + using view_t = std::decay_t; + constexpr bool idx_is_padded = + std::is_same_v>; + constexpr bool idx_is_standard = + std::is_same_v>; + constexpr bool view_is_padded = cuvs::neighbors::is_device_padded_dataset_view_v || + cuvs::neighbors::is_host_padded_dataset_view_v; + constexpr bool view_is_standard = cuvs::neighbors::is_device_standard_dataset_view_v || + cuvs::neighbors::is_host_standard_dataset_view_v; + if constexpr ((idx_is_padded && view_is_padded) || (idx_is_standard && view_is_standard)) { + auto storage = cuvs::neighbors::cagra::make_extended_storage(*res_ptr, ds_view, idx); + auto* storage_holder = new cagra_c_api_extended_dataset_holder{std::move(storage)}; + out->addr = reinterpret_cast(storage_holder); + } else { + RAFT_FAIL("cuvsMakeExtendedStorage: incompatible index layout and dataset view"); + } + }; + with_dataset_view_for_layout(res_ptr, + additional_dataset_tensor, + box->layout, + "cuvsMakeExtendedStorage", + "cuvsMakeExtendedStorage: host additional dataset " + "is allowed", + build_storage); + }); + + *output_extended_storage = out; +} + static void _set_graph_build_params( std::variant -void* _build(cuvsResources_t res, cuvsCagraIndexParams params, DLManagedTensor* dataset_tensor) +void _build(cuvsResources_t res, + cuvsCagraIndexParams params, + DLManagedTensor* dataset_tensor, + cuvsCagraIndex_t output_index) { auto dataset = dataset_tensor->dl_tensor; - auto res_ptr = reinterpret_cast(res); - auto index = new cuvs::neighbors::cagra::index(*res_ptr); auto index_params = cuvs::neighbors::cagra::index_params(); convert_c_index_params(params, dataset.shape[0], dataset.shape[1], &index_params); @@ -120,82 +553,150 @@ void* _build(cuvsResources_t res, cuvsCagraIndexParams params, DLManagedTensor* if (cuvs::core::is_dlpack_device_compatible(dataset)) { using mdspan_type = raft::device_matrix_view; auto mds = cuvs::core::from_dlpack(dataset_tensor); - *index = cuvs::neighbors::cagra::build(*res_ptr, index_params, mds); + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(mds)) { + auto view = cuvs::neighbors::make_device_padded_dataset_view(*res_ptr, mds); + auto index = cuvs::neighbors::cagra::build(*res_ptr, index_params, view); + auto* raw = new cuvs::neighbors::cagra::device_padded_index(std::move(index)); + wrap_CPP_index_in_lifetime_holder_and_bind_to_C_index< + T, + cuvs::neighbors::device_padded_dataset_view>( + output_index, output_index->dtype, raw); + } else { + auto view = cuvs::neighbors::make_device_standard_dataset_view(mds); + auto index = cuvs::neighbors::cagra::build(*res_ptr, index_params, view); + auto* raw = new cuvs::neighbors::cagra::device_standard_index(std::move(index)); + wrap_CPP_index_in_lifetime_holder_and_bind_to_C_index< + T, + cuvs::neighbors::device_standard_dataset_view>( + output_index, output_index->dtype, raw); + } } else if (cuvs::core::is_dlpack_host_compatible(dataset)) { using mdspan_type = raft::host_matrix_view; auto mds = cuvs::core::from_dlpack(dataset_tensor); - *index = cuvs::neighbors::cagra::build(*res_ptr, index_params, mds); + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(mds)) { + auto host_view = cuvs::neighbors::make_host_padded_dataset_view(mds); + auto host_idx = cuvs::neighbors::cagra::build(*res_ptr, index_params, host_view); + auto* raw = new cuvs::neighbors::cagra::host_padded_index(std::move(host_idx)); + wrap_CPP_index_in_lifetime_holder_and_bind_to_C_index< + T, + cuvs::neighbors::host_padded_dataset_view>( + output_index, output_index->dtype, raw); + } else { + auto host_view = cuvs::neighbors::make_host_standard_dataset_view(mds); + auto host_idx = cuvs::neighbors::cagra::build(*res_ptr, index_params, host_view); + auto* raw = + new cuvs::neighbors::cagra::host_standard_index(std::move(host_idx)); + wrap_CPP_index_in_lifetime_holder_and_bind_to_C_index< + T, + cuvs::neighbors::host_standard_dataset_view>( + output_index, output_index->dtype, raw); + } } - return index; } template -void* _from_args(cuvsResources_t res, - cuvsDistanceType _metric, - DLManagedTensor* graph_tensor, - DLManagedTensor* dataset_tensor) +void _from_args(cuvsResources_t res, + cuvsDistanceType _metric, + DLManagedTensor* graph_tensor, + DLManagedTensor* dataset_tensor, + cuvsCagraIndex_t output_index) { auto metric = static_cast((int)_metric); auto dataset = dataset_tensor->dl_tensor; - auto graph = graph_tensor->dl_tensor; auto res_ptr = reinterpret_cast(res); - - void* index = NULL; - if (cuvs::core::is_dlpack_device_compatible(dataset)) { - using mdspan_type = raft::device_matrix_view; - auto mds = cuvs::core::from_dlpack(dataset_tensor); - if (cuvs::core::is_dlpack_device_compatible(graph)) { + auto update_graph_from_dlpack = [&](auto* idx) { + if (cuvs::core::is_dlpack_device_compatible(graph_tensor->dl_tensor)) { using graph_mdspan_type = raft::device_matrix_view; - auto graph_mds = cuvs::core::from_dlpack(graph_tensor); - index = new cuvs::neighbors::cagra::index(*res_ptr, metric, mds, graph_mds); + auto graph_mds = cuvs::core::from_dlpack(graph_tensor); + idx->update_graph(*res_ptr, graph_mds); } else { using graph_mdspan_type = raft::host_matrix_view; - auto graph_mds = cuvs::core::from_dlpack(graph_tensor); - index = new cuvs::neighbors::cagra::index(*res_ptr, metric, mds, graph_mds); + auto graph_mds = cuvs::core::from_dlpack(graph_tensor); + idx->update_graph(*res_ptr, graph_mds); } - } else if (cuvs::core::is_dlpack_host_compatible(dataset)) { - using mdspan_type = raft::host_matrix_view; + }; + + if (cuvs::core::is_dlpack_device_compatible(dataset)) { + using mdspan_type = raft::device_matrix_view; auto mds = cuvs::core::from_dlpack(dataset_tensor); - if (cuvs::core::is_dlpack_device_compatible(graph)) { - using graph_mdspan_type = raft::device_matrix_view; - auto graph_mds = cuvs::core::from_dlpack(graph_tensor); - index = new cuvs::neighbors::cagra::index(*res_ptr, metric, mds, graph_mds); + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(mds)) { + auto dataset_view = cuvs::neighbors::make_device_padded_dataset_view(*res_ptr, mds); + auto* raw = new cuvs::neighbors::cagra::device_padded_index( + *res_ptr, metric); + raw->update_dataset(*res_ptr, dataset_view); + update_graph_from_dlpack(raw); + wrap_CPP_index_in_lifetime_holder_and_bind_to_C_index< + T, + cuvs::neighbors::device_padded_dataset_view>( + output_index, output_index->dtype, raw); } else { - using graph_mdspan_type = raft::host_matrix_view; - auto graph_mds = cuvs::core::from_dlpack(graph_tensor); - index = new cuvs::neighbors::cagra::index(*res_ptr, metric, mds, graph_mds); + auto dataset_view = cuvs::neighbors::make_device_standard_dataset_view(mds); + auto* raw = new cuvs::neighbors::cagra::device_standard_index( + *res_ptr, metric); + raw->update_dataset(*res_ptr, dataset_view); + update_graph_from_dlpack(raw); + wrap_CPP_index_in_lifetime_holder_and_bind_to_C_index< + T, + cuvs::neighbors::device_standard_dataset_view>( + output_index, output_index->dtype, raw); } + } else if (cuvs::core::is_dlpack_host_compatible(dataset)) { + RAFT_FAIL("cuvsCagraIndexFromArgs: host dataset is unsupported; use cuvsCagraBuild for host " + "datasets, then cuvsCagraAttachDeviceDatasetOnHostIndex to get a device index."); } - return index; } template void _extend(cuvsResources_t res, cuvsCagraExtendParams params, cuvsCagraIndex index, - DLManagedTensor* additional_dataset_tensor) + DLManagedTensor* additional_dataset_tensor, + cuvsDatasetStorage_t extended_dataset) { - auto dataset = additional_dataset_tensor->dl_tensor; - auto index_ptr = reinterpret_cast*>(index.addr); + auto* box = reinterpret_cast(index.addr); auto res_ptr = reinterpret_cast(res); + RAFT_EXPECTS(box != nullptr, "cuvsCagraExtend: null index handle"); + RAFT_EXPECTS(extended_dataset != nullptr, "cuvsCagraExtend: null extended dataset handle"); + RAFT_EXPECTS(extended_dataset->kind == CUVS_DATASET_STORAGE_KIND_EXTENDED, + "cuvsCagraExtend: storage handle kind must be EXTENDED"); + auto* storage_holder = + reinterpret_cast*>(extended_dataset->addr); + RAFT_EXPECTS(storage_holder != nullptr, "cuvsCagraExtend: null extended dataset storage"); // TODO: use C struct here (see issue #487) auto extend_params = cuvs::neighbors::cagra::extend_params(); extend_params.max_chunk_size = params.max_chunk_size; - - if (cuvs::core::is_dlpack_device_compatible(dataset)) { - using mdspan_type = raft::device_matrix_view; - auto mds = cuvs::core::from_dlpack(additional_dataset_tensor); - cuvs::neighbors::cagra::extend(*res_ptr, extend_params, mds, *index_ptr); - } else if (cuvs::core::is_dlpack_host_compatible(dataset)) { - using mdspan_type = raft::host_matrix_view; - auto mds = cuvs::core::from_dlpack(additional_dataset_tensor); - cuvs::neighbors::cagra::extend(*res_ptr, extend_params, mds, *index_ptr); - } else { - RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", - dataset.dtype.code, - dataset.dtype.bits); - } + with_index_by_layout( + box, + "cuvsCagraExtend: null index handle", + "cuvsCagraExtend: host indices are not extendable; attach a device dataset to the host index " + "first.", + [&](auto& idx) { + auto run_extend = [&](auto ds_view) { + using index_t = std::decay_t; + using view_t = std::decay_t; + constexpr bool idx_is_padded = + std::is_same_v>; + constexpr bool idx_is_standard = + std::is_same_v>; + constexpr bool view_is_padded = cuvs::neighbors::is_device_padded_dataset_view_v || + cuvs::neighbors::is_host_padded_dataset_view_v; + constexpr bool view_is_standard = cuvs::neighbors::is_device_standard_dataset_view_v || + cuvs::neighbors::is_host_standard_dataset_view_v; + if constexpr ((idx_is_padded && view_is_padded) || (idx_is_standard && view_is_standard)) { + cuvs::neighbors::cagra::extend( + *res_ptr, extend_params, ds_view, idx, storage_holder->storage); + } else { + RAFT_FAIL("cuvsCagraExtend: incompatible index layout and additional dataset view"); + } + }; + with_dataset_view_for_layout(res_ptr, + additional_dataset_tensor, + box->layout, + "cuvsCagraExtend", + "cuvsCagraExtend: host additional dataset is allowed", + run_extend); + }); } template @@ -208,7 +709,13 @@ void _search(cuvsResources_t res, cuvsFilter filter) { auto res_ptr = reinterpret_cast(res); - auto index_ptr = reinterpret_cast*>(index.addr); + auto* box = reinterpret_cast(index.addr); + auto* index_ptr = require_padded_index( + box, + "cuvsCagraSearch: null index handle", + "cuvsCagraSearch: index is not device-padded. For standard device indices call " + "cuvsDatasetMakePadded + cuvsCagraAttachPaddedDatasetForSearch; for host indices call " + "cuvsCagraAttachDeviceDatasetOnHostIndex first."); auto search_params = cuvs::neighbors::cagra::search_params(); convert_c_search_params(params, &search_params); @@ -227,7 +734,7 @@ void _search(cuvsResources_t res, 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, index_ptr->dataset().extent(0)); + removed_indices, index_ptr->dataset().n_rows()); auto bitset_filter_obj = cuvs::neighbors::filtering::bitset_filter(removed_indices_bitset); cuvs::neighbors::cagra::search(*res_ptr, search_params, @@ -270,35 +777,74 @@ void _serialize(cuvsResources_t res, bool include_dataset) { auto res_ptr = reinterpret_cast(res); - auto index_ptr = reinterpret_cast*>(index->addr); - cuvs::neighbors::cagra::serialize(*res_ptr, std::string(filename), *index_ptr, include_dataset); + auto* box = reinterpret_cast(index->addr); + with_index_by_layout( + box, + "cuvsCagraSerialize: null index handle", + "cuvsCagraSerialize: host index must be converted to device first via " + "cuvsCagraAttachDeviceDatasetOnHostIndex", + [&](auto& idx) { cuvs::neighbors::cagra::serialize(*res_ptr, std::string(filename), idx, include_dataset); }); } template void _serialize_to_hnswlib(cuvsResources_t res, const char* filename, cuvsCagraIndex_t index) { auto res_ptr = reinterpret_cast(res); - auto index_ptr = reinterpret_cast*>(index->addr); - cuvs::neighbors::cagra::serialize_to_hnswlib(*res_ptr, std::string(filename), *index_ptr); + auto* box = reinterpret_cast(index->addr); + with_index_by_layout( + box, + "cuvsCagraSerializeToHnswlib: null index handle", + "cuvsCagraSerializeToHnswlib: host index must be converted to device first via " + "cuvsCagraAttachDeviceDatasetOnHostIndex", + [&](auto& idx) { cuvs::neighbors::cagra::serialize_to_hnswlib(*res_ptr, std::string(filename), idx); }); } template -void* _deserialize(cuvsResources_t res, const char* filename) +void _deserialize(cuvsResources_t res, + const char* filename, + cuvsDatasetLayout_t deserialize_layout, + cuvsCagraIndex_t output_index) { auto res_ptr = reinterpret_cast(res); - auto index = new cuvs::neighbors::cagra::index(*res_ptr); - cuvs::neighbors::cagra::deserialize(*res_ptr, std::string(filename), index); - return index; + if (deserialize_layout == CUVS_DATASET_LAYOUT_PADDED) { + auto holder = std::make_unique< + cuvs_cagra_c_api_index_lifetime_holder>>( + cuvs::neighbors::cagra::device_padded_index(*res_ptr)); + cuvs::neighbors::cagra::deserialize(*res_ptr, std::string(filename), &holder->idx, nullptr); + bind_index_lifetime_holder_to_C_index< + T, cuvs::neighbors::device_padded_dataset_view>( + output_index, output_index->dtype, holder.release()); + return; + } else if (deserialize_layout == CUVS_DATASET_LAYOUT_STANDARD) { + auto holder = std::make_unique< + cuvs_cagra_c_api_index_lifetime_holder>>( + cuvs::neighbors::cagra::device_standard_index(*res_ptr)); + cuvs::neighbors::cagra::deserialize(*res_ptr, std::string(filename), &holder->idx, nullptr); + bind_index_lifetime_holder_to_C_index< + T, cuvs::neighbors::device_standard_dataset_view>( + output_index, output_index->dtype, holder.release()); + return; + } + RAFT_FAIL("cuvsCagraDeserialize: unsupported target layout"); } template -void* _merge(cuvsResources_t res, - cuvsCagraIndexParams params, - cuvsCagraIndex_t* indices, - size_t num_indices, - cuvsFilter filter) +void _merge(cuvsResources_t res, + cuvsCagraIndexParams params, + cuvsCagraIndex_t* indices, + size_t num_indices, + cuvsFilter filter, + cuvsDatasetStorage_t merged_dataset, + cuvsCagraIndex_t output_index) { auto res_ptr = reinterpret_cast(res); + auto* first_box = reinterpret_cast(indices[0]->addr); + RAFT_EXPECTS(first_box != nullptr, "cuvsCagraMerge: 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, + "cuvsCagraMerge: host indices are not mergeable; attach a device dataset to each " + "host index first."); cuvs::neighbors::cagra::index_params params_cpp; params_cpp.metric = @@ -309,14 +855,19 @@ void* _merge(cuvsResources_t res, int64_t total_size = 0; int64_t dim = 0; + for (size_t i = 0; i < num_indices; ++i) { + auto* box = reinterpret_cast(indices[i]->addr); + RAFT_EXPECTS(box != nullptr, "cuvsCagraMerge: null index handle"); + RAFT_EXPECTS(box->layout == layout, + "cuvsCagraMerge: all input indices must share the same dataset layout"); + } if (params.build_algo == cuvsCagraGraphBuildAlgo::IVF_PQ) { - auto first_idx_ptr = - reinterpret_cast*>(indices[0]->addr); - dim = first_idx_ptr->dim(); - for (size_t i = 0; i < num_indices; ++i) { - auto idx_ptr = - reinterpret_cast*>(indices[i]->addr); - total_size += idx_ptr->size(); + if (layout == sg_cagra_c_api_index_box::dataset_layout::device_padded) { + compute_ivfpq_shape_from_indices>( + indices, num_indices, &total_size, &dim); + } else { + compute_ivfpq_shape_from_indices>( + indices, num_indices, &total_size, &dim); } } @@ -325,43 +876,123 @@ void* _merge(cuvsResources_t res, params.build_algo, total_size, dim); + auto* merge_holder = + reinterpret_cast*>(merged_dataset->addr); + RAFT_EXPECTS(merged_dataset->kind == CUVS_DATASET_STORAGE_KIND_MERGED, + "cuvsCagraMerge: storage handle kind must be MERGED"); + RAFT_EXPECTS(merge_holder != nullptr, "cuvsCagraMerge: null merged dataset storage"); + + if (layout == sg_cagra_c_api_index_box::dataset_layout::device_padded) { + auto index_ptrs = + convert_opaque_indices_to_concrete_types>( + indices, num_indices); + merge_indices_for_layout>( + res_ptr, params_cpp, index_ptrs, filter, merge_holder->storage, 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_holder->storage, output_index); + } +} - std::vector*> index_ptrs; - index_ptrs.reserve(num_indices); +template +void _make_merged_storage(cuvsResources_t res, + cuvsCagraIndex_t* indices, + size_t num_indices, + cuvsFilter filter, + cuvsDatasetStorage_t* output_merged_storage) +{ + auto res_ptr = reinterpret_cast(res); + auto* first_box = reinterpret_cast(indices[0]->addr); + RAFT_EXPECTS(first_box != nullptr, "cuvsMakeMergedStorage: 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, + "cuvsMakeMergedStorage: host indices are unsupported; attach device dataset to " + "host indices first."); for (size_t i = 0; i < num_indices; ++i) { - auto idx_ptr = reinterpret_cast*>(indices[i]->addr); - index_ptrs.push_back(idx_ptr); + auto* box = reinterpret_cast(indices[i]->addr); + RAFT_EXPECTS(box != nullptr, "cuvsMakeMergedStorage: null index handle"); + RAFT_EXPECTS(box->layout == layout, + "cuvsMakeMergedStorage: all input indices must share the same dataset " + "layout"); } - if (filter.type == NO_FILTER) { - return new cuvs::neighbors::cagra::index( - cuvs::neighbors::cagra::merge(*res_ptr, params_cpp, index_ptrs)); - } 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, total_size); - auto bitset_filter_obj = cuvs::neighbors::filtering::bitset_filter(removed_indices_bitset); - return new cuvs::neighbors::cagra::index( - cuvs::neighbors::cagra::merge(*res_ptr, params_cpp, index_ptrs, bitset_filter_obj)); - } else { - RAFT_FAIL("Unsupported filter type: BITMAP"); - } + auto* out = new cuvsDatasetStorage{ + 0, indices[0]->dtype, CUVS_DATASET_STORAGE_KIND_MERGED}; + cuvs::neighbors::cagra::merged_dataset_storage storage = + [&]() -> cuvs::neighbors::cagra::merged_dataset_storage { + if (layout == sg_cagra_c_api_index_box::dataset_layout::device_padded) { + auto index_ptrs = + convert_opaque_indices_to_concrete_types>( + indices, num_indices); + if (filter.type == NO_FILTER) { return cuvs::neighbors::cagra::make_merged_storage(*res_ptr, index_ptrs); } + if (filter.type == BITSET) { + int64_t merged_row_count = 0; + for (auto* idx_ptr : index_ptrs) { + merged_row_count += static_cast(idx_ptr->size()); + } + 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); + return cuvs::neighbors::cagra::make_merged_storage(*res_ptr, index_ptrs, bitset_filter_obj); + } + RAFT_FAIL("Unsupported filter type: BITMAP"); + } else if (layout == sg_cagra_c_api_index_box::dataset_layout::device_standard) { + auto index_ptrs = + convert_opaque_indices_to_concrete_types>( + indices, num_indices); + if (filter.type == NO_FILTER) { return cuvs::neighbors::cagra::make_merged_storage(*res_ptr, index_ptrs); } + if (filter.type == BITSET) { + int64_t merged_row_count = 0; + for (auto* idx_ptr : index_ptrs) { + merged_row_count += static_cast(idx_ptr->size()); + } + 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); + return cuvs::neighbors::cagra::make_merged_storage(*res_ptr, index_ptrs, bitset_filter_obj); + } + RAFT_FAIL("Unsupported filter type: BITMAP"); + } else { + RAFT_FAIL("cuvsMakeMergedStorage: unsupported index layout"); + } + }(); + auto* storage_holder = new cagra_c_api_merged_dataset_holder{std::move(storage)}; + out->addr = reinterpret_cast(storage_holder); + *output_merged_storage = out; } template void get_dataset_view(cuvsCagraIndex_t index, DLManagedTensor* dataset) { - auto index_ptr = reinterpret_cast*>(index->addr); - cuvs::core::to_dlpack(index_ptr->dataset(), dataset); + auto* box = reinterpret_cast(index->addr); + with_index_by_layout( + box, + "cuvsCagraIndexGetDataset: null index handle", + "cuvsCagraIndexGetDataset: host indices are allowed", + [&](auto& idx) { cuvs::core::to_dlpack(idx.dataset().view(), dataset); }); } template void get_graph_view(cuvsCagraIndex_t index, DLManagedTensor* graph) { - auto index_ptr = reinterpret_cast*>(index->addr); - cuvs::core::to_dlpack(index_ptr->graph(), graph); + auto* box = reinterpret_cast(index->addr); + with_index_by_layout( + box, + "cuvsCagraIndexGetGraph: null index handle", + "cuvsCagraIndexGetGraph: host indices are allowed", + [&](auto& idx) { cuvs::core::to_dlpack(idx.graph(), graph); }); } // Helper function to populate C IVF-PQ params from C++ params @@ -445,16 +1076,6 @@ void convert_c_index_params(cuvsCagraIndexParams params, out->graph_degree = params.graph_degree; _set_graph_build_params(out->graph_build_params, params, params.build_algo, n_rows, dim); - if (auto* cparams = params.compression; cparams != nullptr) { - auto compression_params = cuvs::neighbors::vpq_params(); - compression_params.pq_bits = cparams->pq_bits; - compression_params.pq_dim = cparams->pq_dim; - compression_params.vq_n_centers = cparams->vq_n_centers; - compression_params.kmeans_n_iters = cparams->kmeans_n_iters; - compression_params.vq_kmeans_trainset_fraction = cparams->vq_kmeans_trainset_fraction; - compression_params.pq_kmeans_trainset_fraction = cparams->pq_kmeans_trainset_fraction; - out->compression.emplace(compression_params); - } } void convert_c_search_params(cuvsCagraSearchParams params, cuvs::neighbors::cagra::search_params* out) @@ -476,33 +1097,26 @@ void convert_c_search_params(cuvsCagraSearchParams params, out->persistent_lifetime = params.persistent_lifetime; out->persistent_device_usage = params.persistent_device_usage; } + +void* cagra_c_api_index_ptr(cuvsCagraIndex const* idx) +{ + // Matches `sg_cagra_c_api_index_box::index_ptr` (first member); keep in sync with that layout. + if (idx == nullptr || idx->addr == 0) { return nullptr; } + return *reinterpret_cast(idx->addr); +} } // namespace cuvs::neighbors::cagra + extern "C" cuvsError_t cuvsCagraIndexCreate(cuvsCagraIndex_t* index) { - return cuvs::core::translate_exceptions([=] { *index = new cuvsCagraIndex{}; }); + return cuvs::core::translate_exceptions([=] { + *index = new cuvsCagraIndex{0, {}}; + }); } extern "C" cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index_c_ptr) { return cuvs::core::translate_exceptions([=] { - auto index = *index_c_ptr; - - if (index.dtype.code == kDLFloat && index.dtype.bits == 32) { - auto index_ptr = - reinterpret_cast*>(index.addr); - delete index_ptr; - } else if (index.dtype.code == kDLFloat && index.dtype.bits == 16) { - auto index_ptr = reinterpret_cast*>(index.addr); - delete index_ptr; - } else if (index.dtype.code == kDLInt && index.dtype.bits == 8) { - auto index_ptr = - reinterpret_cast*>(index.addr); - delete index_ptr; - } else if (index.dtype.code == kDLUInt && index.dtype.bits == 8) { - auto index_ptr = - reinterpret_cast*>(index.addr); - delete index_ptr; - } + destroy_sg_cagra_c_api_box(index_c_ptr->addr); delete index_c_ptr; }); } @@ -510,24 +1124,36 @@ extern "C" cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index_c_ptr) extern "C" cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int64_t* dim) { return cuvs::core::translate_exceptions([=] { - auto index_ptr = reinterpret_cast*>(index->addr); - *dim = index_ptr->dim(); + auto* box = reinterpret_cast(index->addr); + with_index_by_layout( + box, + "cuvsCagraIndexGetDims: null index handle", + "cuvsCagraIndexGetDims: host indices are allowed", + [&](auto& idx) { *dim = idx.dim(); }); }); } extern "C" cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, int64_t* size) { return cuvs::core::translate_exceptions([=] { - auto index_ptr = reinterpret_cast*>(index->addr); - *size = index_ptr->size(); + auto* box = reinterpret_cast(index->addr); + with_index_by_layout( + box, + "cuvsCagraIndexGetSize: null index handle", + "cuvsCagraIndexGetSize: host indices are allowed", + [&](auto& idx) { *size = idx.size(); }); }); } extern "C" cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, int64_t* graph_degree) { return cuvs::core::translate_exceptions([=] { - auto index_ptr = reinterpret_cast*>(index->addr); - *graph_degree = index_ptr->graph_degree(); + auto* box = reinterpret_cast(index->addr); + with_index_by_layout( + box, + "cuvsCagraIndexGetGraphDegree: null index handle", + "cuvsCagraIndexGetGraphDegree: host indices are allowed", + [&](auto& idx) { *graph_degree = idx.graph_degree(); }); }); } @@ -565,6 +1191,211 @@ extern "C" cuvsError_t cuvsCagraIndexGetGraph(cuvsCagraIndex_t index, DLManagedT }); } +extern "C" cuvsError_t cuvsDatasetMakePadded(cuvsResources_t res, + DLManagedTensor* dataset_tensor, + cuvsDatasetPadded_t* padded_dataset) +{ + return cuvs::core::translate_exceptions([=] { + auto dataset = dataset_tensor->dl_tensor; + auto* res_ptr = reinterpret_cast(res); + if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { + make_padded_dataset(res_ptr, dataset_tensor, padded_dataset); + } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 16) { + make_padded_dataset(res_ptr, dataset_tensor, padded_dataset); + } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) { + make_padded_dataset(res_ptr, dataset_tensor, padded_dataset); + } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) { + make_padded_dataset(res_ptr, dataset_tensor, padded_dataset); + } else { + RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", + dataset.dtype.code, + dataset.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsDatasetPaddedDestroy(cuvsDatasetPadded_t padded_dataset) +{ + return cuvs::core::translate_exceptions([=] { + if (padded_dataset == nullptr) { return; } + RAFT_EXPECTS(padded_dataset->layout == CUVS_DATASET_LAYOUT_PADDED, + "cuvsDatasetPaddedDestroy: dataset handle layout must be PADDED"); + if (padded_dataset->addr != 0) { + if (padded_dataset->dtype.code == kDLFloat && padded_dataset->dtype.bits == 32) { + destroy_padded_dataset_typed(padded_dataset->addr); + } else if (padded_dataset->dtype.code == kDLFloat && padded_dataset->dtype.bits == 16) { + destroy_padded_dataset_typed(padded_dataset->addr); + } else if (padded_dataset->dtype.code == kDLInt && padded_dataset->dtype.bits == 8) { + destroy_padded_dataset_typed(padded_dataset->addr); + } else if (padded_dataset->dtype.code == kDLUInt && padded_dataset->dtype.bits == 8) { + destroy_padded_dataset_typed(padded_dataset->addr); + } else { + RAFT_FAIL("Unsupported padded dataset dtype: %d and bits: %d", + padded_dataset->dtype.code, + padded_dataset->dtype.bits); + } + } + delete padded_dataset; + }); +} + +extern "C" cuvsError_t cuvsCagraAttachPaddedDatasetForSearch(cuvsResources_t res, + cuvsDatasetPadded_t padded_dataset, + cuvsCagraIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + auto* res_ptr = reinterpret_cast(res); + RAFT_EXPECTS(padded_dataset != nullptr, + "cuvsCagraAttachPaddedDatasetForSearch: null padded dataset"); + RAFT_EXPECTS(index != nullptr, "cuvsCagraAttachPaddedDatasetForSearch: null index handle"); + RAFT_EXPECTS(padded_dataset->layout == CUVS_DATASET_LAYOUT_PADDED, + "cuvsCagraAttachPaddedDatasetForSearch: dataset handle layout must be PADDED"); + RAFT_EXPECTS(index->dtype.code == padded_dataset->dtype.code && + index->dtype.bits == padded_dataset->dtype.bits, + "cuvsCagraAttachPaddedDatasetForSearch: dtype mismatch between index and padded dataset"); + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + attach_padded_dataset_for_search(res_ptr, padded_dataset, index); + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + attach_padded_dataset_for_search(res_ptr, padded_dataset, index); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + attach_padded_dataset_for_search(res_ptr, padded_dataset, index); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + attach_padded_dataset_for_search(res_ptr, padded_dataset, index); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsCagraAttachDeviceDatasetOnHostIndex(cuvsResources_t res, + DLManagedTensor* device_dataset_tensor, + cuvsCagraIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + auto* res_ptr = reinterpret_cast(res); + RAFT_EXPECTS(index != nullptr, "cuvsCagraAttachDeviceDatasetOnHostIndex: null index handle"); + RAFT_EXPECTS(device_dataset_tensor != nullptr, + "cuvsCagraAttachDeviceDatasetOnHostIndex: null dataset tensor"); + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + attach_device_dataset_on_host_index(res_ptr, device_dataset_tensor, index); + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + attach_device_dataset_on_host_index(res_ptr, device_dataset_tensor, index); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + attach_device_dataset_on_host_index(res_ptr, device_dataset_tensor, index); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + attach_device_dataset_on_host_index(res_ptr, device_dataset_tensor, index); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsMakeExtendedStorage(cuvsResources_t res, + DLManagedTensor* additional_dataset_tensor, + cuvsCagraIndex_t index, + cuvsDatasetStorage_t* extended_storage) +{ + return cuvs::core::translate_exceptions([=] { + RAFT_EXPECTS(additional_dataset_tensor != nullptr, + "cuvsMakeExtendedStorage: null additional dataset"); + RAFT_EXPECTS(index != nullptr, "cuvsMakeExtendedStorage: null index handle"); + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + make_extended_storage(reinterpret_cast(res), + additional_dataset_tensor, + index, + extended_storage); + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + make_extended_storage(reinterpret_cast(res), + additional_dataset_tensor, + index, + extended_storage); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + make_extended_storage(reinterpret_cast(res), + additional_dataset_tensor, + index, + extended_storage); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + make_extended_storage(reinterpret_cast(res), + additional_dataset_tensor, + index, + extended_storage); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsDatasetStorageDestroy(cuvsDatasetStorage_t dataset_storage) +{ + return cuvs::core::translate_exceptions([=] { + if (dataset_storage == nullptr) { return; } + RAFT_EXPECTS(dataset_storage->kind == CUVS_DATASET_STORAGE_KIND_EXTENDED || + dataset_storage->kind == CUVS_DATASET_STORAGE_KIND_MERGED, + "cuvsDatasetStorageDestroy: unsupported storage handle kind"); + if (dataset_storage->addr != 0) { + if (dataset_storage->dtype.code == kDLFloat && dataset_storage->dtype.bits == 32) { + if (dataset_storage->kind == CUVS_DATASET_STORAGE_KIND_EXTENDED) { + destroy_extended_dataset_typed(dataset_storage->addr); + } else { + destroy_merged_dataset_typed(dataset_storage->addr); + } + } else if (dataset_storage->dtype.code == kDLFloat && dataset_storage->dtype.bits == 16) { + if (dataset_storage->kind == CUVS_DATASET_STORAGE_KIND_EXTENDED) { + destroy_extended_dataset_typed(dataset_storage->addr); + } else { + destroy_merged_dataset_typed(dataset_storage->addr); + } + } else if (dataset_storage->dtype.code == kDLInt && dataset_storage->dtype.bits == 8) { + if (dataset_storage->kind == CUVS_DATASET_STORAGE_KIND_EXTENDED) { + destroy_extended_dataset_typed(dataset_storage->addr); + } else { + destroy_merged_dataset_typed(dataset_storage->addr); + } + } else if (dataset_storage->dtype.code == kDLUInt && dataset_storage->dtype.bits == 8) { + if (dataset_storage->kind == CUVS_DATASET_STORAGE_KIND_EXTENDED) { + destroy_extended_dataset_typed(dataset_storage->addr); + } else { + destroy_merged_dataset_typed(dataset_storage->addr); + } + } else { + RAFT_FAIL("Unsupported dataset storage dtype: %d and bits: %d", + dataset_storage->dtype.code, + dataset_storage->dtype.bits); + } + } + delete dataset_storage; + }); +} + +extern "C" cuvsError_t cuvsMakeMergedStorage(cuvsResources_t res, + cuvsCagraIndex_t* indices, + size_t num_indices, + cuvsFilter filter, + cuvsDatasetStorage_t* merged_storage) +{ + return cuvs::core::translate_exceptions([=] { + RAFT_EXPECTS(indices != nullptr && num_indices > 0, + "cuvsMakeMergedStorage: indices array cannot be null or empty"); + auto dtype = (*indices[0]).dtype; + for (size_t i = 1; i < num_indices; ++i) { + 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) { + _make_merged_storage(res, indices, num_indices, filter, merged_storage); + } else if (dtype.code == kDLFloat && dtype.bits == 16) { + _make_merged_storage(res, indices, num_indices, filter, merged_storage); + } else if (dtype.code == kDLInt && dtype.bits == 8) { + _make_merged_storage(res, indices, num_indices, filter, merged_storage); + } else if (dtype.code == kDLUInt && dtype.bits == 8) { + _make_merged_storage(res, indices, num_indices, filter, merged_storage); + } else { + RAFT_FAIL("Unsupported index data type: code=%d, bits=%d", dtype.code, dtype.bits); + } + }); +} + + extern "C" cuvsError_t cuvsCagraBuild(cuvsResources_t res, cuvsCagraIndexParams_t params, DLManagedTensor* dataset_tensor, @@ -572,15 +1403,17 @@ extern "C" cuvsError_t cuvsCagraBuild(cuvsResources_t res, { return cuvs::core::translate_exceptions([=] { auto dataset = dataset_tensor->dl_tensor; - index->dtype = dataset.dtype; + destroy_sg_cagra_c_api_box(index->addr); + index->addr = 0; + index->dtype = dataset.dtype; if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { - index->addr = reinterpret_cast(_build(res, *params, dataset_tensor)); + _build(res, *params, dataset_tensor, index); } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 16) { - index->addr = reinterpret_cast(_build(res, *params, dataset_tensor)); + _build(res, *params, dataset_tensor, index); } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) { - index->addr = reinterpret_cast(_build(res, *params, dataset_tensor)); + _build(res, *params, dataset_tensor, index); } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) { - index->addr = reinterpret_cast(_build(res, *params, dataset_tensor)); + _build(res, *params, dataset_tensor, index); } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, @@ -597,19 +1430,17 @@ extern "C" cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res, { return cuvs::core::translate_exceptions([=] { auto dataset = dataset_tensor->dl_tensor; - index->dtype = dataset.dtype; + destroy_sg_cagra_c_api_box(index->addr); + index->addr = 0; + index->dtype = dataset.dtype; if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { - index->addr = - reinterpret_cast(_from_args(res, metric, graph_tensor, dataset_tensor)); + _from_args(res, metric, graph_tensor, dataset_tensor, index); } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 16) { - index->addr = - reinterpret_cast(_from_args(res, metric, graph_tensor, dataset_tensor)); + _from_args(res, metric, graph_tensor, dataset_tensor, index); } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) { - index->addr = - reinterpret_cast(_from_args(res, metric, graph_tensor, dataset_tensor)); + _from_args(res, metric, graph_tensor, dataset_tensor, index); } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) { - index->addr = - reinterpret_cast(_from_args(res, metric, graph_tensor, dataset_tensor)); + _from_args(res, metric, graph_tensor, dataset_tensor, index); } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, @@ -621,20 +1452,25 @@ extern "C" cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res, extern "C" cuvsError_t cuvsCagraExtend(cuvsResources_t res, cuvsCagraExtendParams_t params, DLManagedTensor* additional_dataset_tensor, + cuvsDatasetStorage_t extended_dataset, cuvsCagraIndex_t index_c_ptr) { return cuvs::core::translate_exceptions([=] { auto dataset = additional_dataset_tensor->dl_tensor; auto index = *index_c_ptr; + RAFT_EXPECTS(extended_dataset != nullptr, "cuvsCagraExtend: null extended dataset handle"); + RAFT_EXPECTS(extended_dataset->dtype.code == index.dtype.code && + extended_dataset->dtype.bits == index.dtype.bits, + "cuvsCagraExtend: dtype mismatch between index and extended dataset"); if ((dataset.dtype.code == kDLFloat) && (dataset.dtype.bits == 32)) { - _extend(res, *params, index, additional_dataset_tensor); + _extend(res, *params, index, additional_dataset_tensor, extended_dataset); } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 16) { - _extend(res, *params, index, additional_dataset_tensor); + _extend(res, *params, index, additional_dataset_tensor, extended_dataset); } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) { - _extend(res, *params, index, additional_dataset_tensor); + _extend(res, *params, index, additional_dataset_tensor, extended_dataset); } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) { - _extend(res, *params, index, additional_dataset_tensor); + _extend(res, *params, index, additional_dataset_tensor, extended_dataset); } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, @@ -695,6 +1531,7 @@ extern "C" cuvsError_t cuvsCagraMerge(cuvsResources_t res, cuvsCagraIndex_t* indices, size_t num_indices, cuvsFilter filter, + cuvsDatasetStorage_t merged_dataset, cuvsCagraIndex_t output_index) { return cuvs::core::translate_exceptions([=] { @@ -710,20 +1547,21 @@ extern "C" cuvsError_t cuvsCagraMerge(cuvsResources_t res, RAFT_EXPECTS((*indices[i]).addr != 0, "All input indices must be built (non-empty)"); } RAFT_EXPECTS(output_index != nullptr, "Output index pointer must not be null"); + RAFT_EXPECTS(merged_dataset != nullptr, "Merged dataset storage must not be null"); + RAFT_EXPECTS(merged_dataset->dtype.code == dtype.code && merged_dataset->dtype.bits == dtype.bits, + "dtype mismatch between merged dataset storage and indices"); 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) { - output_index->addr = - reinterpret_cast(_merge(res, *params, indices, num_indices, filter)); + _merge(res, *params, indices, num_indices, filter, merged_dataset, output_index); } else if (dtype.code == kDLFloat && dtype.bits == 16) { - output_index->addr = - reinterpret_cast(_merge(res, *params, indices, num_indices, filter)); + _merge(res, *params, indices, num_indices, filter, merged_dataset, output_index); } else if (dtype.code == kDLInt && dtype.bits == 8) { - output_index->addr = - reinterpret_cast(_merge(res, *params, indices, num_indices, filter)); + _merge(res, *params, indices, num_indices, filter, merged_dataset, output_index); } else if (dtype.code == kDLUInt && dtype.bits == 8) { - output_index->addr = - reinterpret_cast(_merge(res, *params, indices, num_indices, filter)); + _merge(res, *params, indices, num_indices, filter, merged_dataset, output_index); } else { RAFT_FAIL("Unsupported index data type: code=%d, bits=%d", dtype.code, dtype.bits); } @@ -868,9 +1706,13 @@ extern "C" cuvsError_t cuvsCagraSearchParamsDestroy(cuvsCagraSearchParams_t para extern "C" cuvsError_t cuvsCagraDeserialize(cuvsResources_t res, const char* filename, + cuvsDatasetLayout_t deserialize_layout, cuvsCagraIndex_t index) { return cuvs::core::translate_exceptions([=] { + destroy_sg_cagra_c_api_box(index->addr); + index->addr = 0; + // read the numpy dtype from the beginning of the file std::ifstream is(filename, std::ios::in | std::ios::binary); if (!is) { RAFT_FAIL("Cannot open file %s", filename); } @@ -884,16 +1726,16 @@ extern "C" cuvsError_t cuvsCagraDeserialize(cuvsResources_t res, index->dtype.bits = dtype.itemsize * 8; if (dtype.kind == 'f' && dtype.itemsize == 4) { - index->addr = reinterpret_cast(_deserialize(res, filename)); + _deserialize(res, filename, deserialize_layout, index); index->dtype.code = kDLFloat; } else if (dtype.kind == 'e' && dtype.itemsize == 2) { - index->addr = reinterpret_cast(_deserialize(res, filename)); + _deserialize(res, filename, deserialize_layout, index); index->dtype.code = kDLFloat; } else if (dtype.kind == 'i' && dtype.itemsize == 1) { - index->addr = reinterpret_cast(_deserialize(res, filename)); + _deserialize(res, filename, deserialize_layout, index); index->dtype.code = kDLInt; } else if (dtype.kind == 'u' && dtype.itemsize == 1) { - index->addr = reinterpret_cast(_deserialize(res, filename)); + _deserialize(res, filename, deserialize_layout, index); index->dtype.code = kDLUInt; } else { RAFT_FAIL("Unsupported dtype in file %s", filename); diff --git a/c/src/neighbors/cagra.hpp b/c/src/neighbors/cagra.hpp index 689bc0fb7a..ab1cae9f05 100644 --- a/c/src/neighbors/cagra.hpp +++ b/c/src/neighbors/cagra.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include @@ -15,4 +15,7 @@ void convert_c_index_params(cuvsCagraIndexParams params, /// Converts C search params to C++ void convert_c_search_params(cuvsCagraSearchParams params, cuvs::neighbors::cagra::search_params* out); + +/** Resolves `cuvsCagraIndex::addr` to `cagra::index*`; nullptr if the handle is empty. */ +void* cagra_c_api_index_ptr(cuvsCagraIndex const* idx); } // namespace cuvs::neighbors::cagra diff --git a/c/src/neighbors/hnsw.cpp b/c/src/neighbors/hnsw.cpp index c69eda0ca0..02bd2918f8 100644 --- a/c/src/neighbors/hnsw.cpp +++ b/c/src/neighbors/hnsw.cpp @@ -1,6 +1,6 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -20,6 +20,7 @@ #include "../core/exceptions.hpp" #include "../core/interop.hpp" +#include "cagra.hpp" namespace { @@ -63,7 +64,8 @@ void _from_cagra(cuvsResources_t res, std::optional dataset_tensor) { auto res_ptr = reinterpret_cast(res); - auto index = reinterpret_cast*>(cagra_index->addr); + auto index = reinterpret_cast*>( + cuvs::neighbors::cagra::cagra_c_api_index_ptr(cagra_index)); auto cpp_params = cuvs::neighbors::hnsw::index_params(); cpp_params.hierarchy = static_cast(params->hierarchy); cpp_params.ef_construction = params->ef_construction; diff --git a/c/src/neighbors/mg_cagra.cpp b/c/src/neighbors/mg_cagra.cpp index 495eff8a34..3e52ae1e74 100644 --- a/c/src/neighbors/mg_cagra.cpp +++ b/c/src/neighbors/mg_cagra.cpp @@ -1,9 +1,10 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "cagra.hpp" +#include "c_api_box.hpp" #include #include #include @@ -18,6 +19,60 @@ #include +namespace { +enum class mg_cagra_dataset_layout : uint8_t { device_padded, device_standard }; + +struct mg_cagra_c_api_index_box { + void* index_ptr; + mg_cagra_dataset_layout layout; + cuvs::neighbors::c_api::detail::owner_record owner_rec; +}; + +template +using mg_cagra_index_t = cuvs::neighbors::mg_index; + +template +static auto make_mg_cagra_box(mg_cagra_index_t* ptr, mg_cagra_dataset_layout layout) + -> mg_cagra_c_api_index_box* +{ + return new mg_cagra_c_api_index_box{ + ptr, layout, cuvs::neighbors::c_api::detail::make_owner_record(ptr)}; +} + +static auto require_mg_cagra_box(cuvsMultiGpuCagraIndex const& index, const char* null_handle_err) + -> mg_cagra_c_api_index_box* +{ + auto* box = reinterpret_cast(index.addr); + if (box == nullptr) { RAFT_FAIL("%s", null_handle_err); } + return box; +} + +static void destroy_mg_cagra_c_api_box(uintptr_t addr) +{ + if (addr == 0) { return; } + auto* box = reinterpret_cast(addr); + cuvs::neighbors::c_api::detail::destroy_owner_record(box->owner_rec); + delete box; +} + +template +static void with_mg_index_by_layout(mg_cagra_c_api_index_box* box, + const char* null_handle_err, + Fn&& fn) +{ + if (box == nullptr) { RAFT_FAIL("%s", null_handle_err); } + if (box->layout == mg_cagra_dataset_layout::device_padded) { + auto* index_ptr = reinterpret_cast< + mg_cagra_index_t>*>(box->index_ptr); + fn(index_ptr); + } else { + auto* index_ptr = reinterpret_cast< + mg_cagra_index_t>*>(box->index_ptr); + fn(index_ptr); + } +} +} // namespace + extern "C" cuvsError_t cuvsMultiGpuCagraIndexParamsCreate( cuvsMultiGpuCagraIndexParams_t* index_params) { @@ -81,32 +136,7 @@ extern "C" cuvsError_t cuvsMultiGpuCagraIndexDestroy(cuvsMultiGpuCagraIndex_t in { return cuvs::core::translate_exceptions([=] { if (index) { - // Properly clean up the templated inner object based on dtype, like single GPU API - if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { - auto mg_index_ptr = - reinterpret_cast, - float, - uint32_t>*>(index->addr); - delete mg_index_ptr; - } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { - auto mg_index_ptr = - reinterpret_cast, - half, - uint32_t>*>(index->addr); - delete mg_index_ptr; - } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { - auto mg_index_ptr = reinterpret_cast< - cuvs::neighbors:: - mg_index, int8_t, uint32_t>*>( - index->addr); - delete mg_index_ptr; - } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { - auto mg_index_ptr = reinterpret_cast< - cuvs::neighbors:: - mg_index, uint8_t, uint32_t>*>( - index->addr); - delete mg_index_ptr; - } + destroy_mg_cagra_c_api_box(index->addr); delete index; } }); @@ -142,11 +172,11 @@ void convert_c_mg_search_params( } // namespace cuvs::neighbors::cagra namespace { - template void* _mg_build(cuvsResources_t res, cuvsMultiGpuCagraIndexParams params, - DLManagedTensor* dataset_tensor) + DLManagedTensor* dataset_tensor, + mg_cagra_dataset_layout layout) { auto res_ptr = reinterpret_cast(res); auto dataset = dataset_tensor->dl_tensor; @@ -158,11 +188,18 @@ void* _mg_build(cuvsResources_t res, using mdspan_type = raft::host_matrix_view; auto mds = cuvs::core::from_dlpack(dataset_tensor); - auto mg_index = - new cuvs::neighbors::mg_index, T, uint32_t>( - cuvs::neighbors::cagra::build(*res_ptr, mg_params, mds)); - - return mg_index; + if (layout == mg_cagra_dataset_layout::device_padded) { + using padded_ann_t = cuvs::neighbors::cagra::device_padded_index; + auto padded_mds = cuvs::neighbors::make_host_padded_dataset_view(mds); + auto* mg_index = new mg_cagra_index_t( + cuvs::neighbors::cagra::build(*res_ptr, mg_params, padded_mds)); + return make_mg_cagra_box(mg_index, mg_cagra_dataset_layout::device_padded); + } + using standard_ann_t = cuvs::neighbors::cagra::device_standard_index; + auto standard_mds = cuvs::neighbors::make_host_standard_dataset_view(mds); + auto* mg_index = new mg_cagra_index_t( + cuvs::neighbors::cagra::build(*res_ptr, mg_params, standard_mds)); + return make_mg_cagra_box(mg_index, mg_cagra_dataset_layout::device_standard); } template @@ -174,9 +211,7 @@ void _mg_search(cuvsResources_t res, DLManagedTensor* distances_tensor) { auto res_ptr = reinterpret_cast(res); - auto mg_index_ptr = reinterpret_cast< - cuvs::neighbors::mg_index, T, uint32_t>*>( - index.addr); + auto* box = require_mg_cagra_box(index, "cuvsMultiGpuCagraSearch: null index handle"); auto mg_search_params = cuvs::neighbors::mg_search_params(); @@ -190,8 +225,10 @@ void _mg_search(cuvsResources_t res, auto neighbors_mds = cuvs::core::from_dlpack(neighbors_tensor); auto distances_mds = cuvs::core::from_dlpack(distances_tensor); - cuvs::neighbors::cagra::search( - *res_ptr, *mg_index_ptr, mg_search_params, queries_mds, neighbors_mds, distances_mds); + with_mg_index_by_layout(box, "cuvsMultiGpuCagraSearch: null index handle", [&](auto* mg_index_ptr) { + cuvs::neighbors::cagra::search( + *res_ptr, *mg_index_ptr, mg_search_params, queries_mds, neighbors_mds, distances_mds); + }); } template @@ -201,9 +238,7 @@ void _mg_extend(cuvsResources_t res, DLManagedTensor* new_indices_tensor) { auto res_ptr = reinterpret_cast(res); - auto mg_index_ptr = reinterpret_cast< - cuvs::neighbors::mg_index, T, uint32_t>*>( - index.addr); + auto* box = require_mg_cagra_box(index, "cuvsMultiGpuCagraExtend: null index handle"); using vectors_mdspan_type = raft::host_matrix_view; auto new_vectors_mds = cuvs::core::from_dlpack(new_vectors_tensor); @@ -214,40 +249,62 @@ void _mg_extend(cuvsResources_t res, new_indices_mds = cuvs::core::from_dlpack(new_indices_tensor); } - cuvs::neighbors::cagra::extend(*res_ptr, *mg_index_ptr, new_vectors_mds, new_indices_mds); + if (box->layout == mg_cagra_dataset_layout::device_padded) { + using padded_ann_t = cuvs::neighbors::cagra::device_padded_index; + auto* mg_index_ptr = + reinterpret_cast*>(box->index_ptr); + auto new_vectors = cuvs::neighbors::make_host_padded_dataset_view(new_vectors_mds); + cuvs::neighbors::cagra::extend(*res_ptr, *mg_index_ptr, new_vectors, new_indices_mds); + } else { + using standard_ann_t = cuvs::neighbors::cagra::device_standard_index; + auto* mg_index_ptr = + reinterpret_cast*>(box->index_ptr); + auto new_vectors = cuvs::neighbors::make_host_standard_dataset_view(new_vectors_mds); + cuvs::neighbors::cagra::extend(*res_ptr, *mg_index_ptr, new_vectors, new_indices_mds); + } } template void _mg_serialize(cuvsResources_t res, cuvsMultiGpuCagraIndex index, const char* filename) { auto res_ptr = reinterpret_cast(res); - auto mg_index_ptr = reinterpret_cast< - cuvs::neighbors::mg_index, T, uint32_t>*>( - index.addr); - - cuvs::neighbors::cagra::serialize(*res_ptr, *mg_index_ptr, std::string(filename)); + auto* box = require_mg_cagra_box(index, "cuvsMultiGpuCagraSerialize: null index handle"); + with_mg_index_by_layout( + box, "cuvsMultiGpuCagraSerialize: null index handle", [&](auto* mg_index_ptr) { + cuvs::neighbors::cagra::serialize(*res_ptr, *mg_index_ptr, std::string(filename)); + }); } template -void* _mg_deserialize(cuvsResources_t res, const char* filename) +void* _mg_deserialize(cuvsResources_t res, const char* filename, mg_cagra_dataset_layout layout) { auto res_ptr = reinterpret_cast(res); - auto mg_index = - new cuvs::neighbors::mg_index, T, uint32_t>( - cuvs::neighbors::cagra::deserialize(*res_ptr, std::string(filename))); - - return mg_index; + if (layout == mg_cagra_dataset_layout::device_padded) { + using padded_ann_t = cuvs::neighbors::cagra::device_padded_index; + auto* mg_index = new mg_cagra_index_t(*res_ptr, cuvs::neighbors::REPLICATED); + cuvs::neighbors::cagra::deserialize(*res_ptr, std::string(filename), mg_index); + return make_mg_cagra_box(mg_index, mg_cagra_dataset_layout::device_padded); + } + using standard_ann_t = cuvs::neighbors::cagra::device_standard_index; + auto* mg_index = new mg_cagra_index_t(*res_ptr, cuvs::neighbors::REPLICATED); + cuvs::neighbors::cagra::deserialize(*res_ptr, std::string(filename), mg_index); + return make_mg_cagra_box(mg_index, mg_cagra_dataset_layout::device_standard); } template -void* _mg_distribute(cuvsResources_t res, const char* filename) +void* _mg_distribute(cuvsResources_t res, const char* filename, mg_cagra_dataset_layout layout) { auto res_ptr = reinterpret_cast(res); - auto mg_index = - new cuvs::neighbors::mg_index, T, uint32_t>( - cuvs::neighbors::cagra::distribute(*res_ptr, std::string(filename))); - - return mg_index; + if (layout == mg_cagra_dataset_layout::device_padded) { + using padded_ann_t = cuvs::neighbors::cagra::device_padded_index; + auto* mg_index = new mg_cagra_index_t(*res_ptr, cuvs::neighbors::REPLICATED); + cuvs::neighbors::cagra::distribute(*res_ptr, std::string(filename), mg_index); + return make_mg_cagra_box(mg_index, mg_cagra_dataset_layout::device_padded); + } + using standard_ann_t = cuvs::neighbors::cagra::device_standard_index; + auto* mg_index = new mg_cagra_index_t(*res_ptr, cuvs::neighbors::REPLICATED); + cuvs::neighbors::cagra::distribute(*res_ptr, std::string(filename), mg_index); + return make_mg_cagra_box(mg_index, mg_cagra_dataset_layout::device_standard); } } // anonymous namespace @@ -264,17 +321,39 @@ extern "C" cuvsError_t cuvsMultiGpuCagraBuild(cuvsResources_t res, RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(dataset), "Multi-GPU CAGRA build requires dataset to have host compatible memory"); + destroy_mg_cagra_c_api_box(index->addr); + index->addr = 0; index->dtype.code = dataset.dtype.code; index->dtype.bits = dataset.dtype.bits; if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { - index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + auto mds = cuvs::core::from_dlpack>( + dataset_tensor); + auto layout = cuvs::neighbors::matrix_row_width_matches_cagra_required(mds) + ? mg_cagra_dataset_layout::device_padded + : mg_cagra_dataset_layout::device_standard; + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor, layout)); } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 16) { - index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + auto mds = cuvs::core::from_dlpack>( + dataset_tensor); + auto layout = cuvs::neighbors::matrix_row_width_matches_cagra_required(mds) + ? mg_cagra_dataset_layout::device_padded + : mg_cagra_dataset_layout::device_standard; + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor, layout)); } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) { - index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + auto mds = cuvs::core::from_dlpack>( + dataset_tensor); + auto layout = cuvs::neighbors::matrix_row_width_matches_cagra_required(mds) + ? mg_cagra_dataset_layout::device_padded + : mg_cagra_dataset_layout::device_standard; + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor, layout)); } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) { - index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + auto mds = cuvs::core::from_dlpack>( + dataset_tensor); + auto layout = cuvs::neighbors::matrix_row_width_matches_cagra_required(mds) + ? mg_cagra_dataset_layout::device_padded + : mg_cagra_dataset_layout::device_standard; + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor, layout)); } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, @@ -408,19 +487,31 @@ extern "C" cuvsError_t cuvsMultiGpuCagraDeserialize(cuvsResources_t res, raft::numpy_serializer::parse_descr(std::string(dtype_string, sizeof(dtype_string))); is.close(); + destroy_mg_cagra_c_api_box(index->addr); + index->addr = 0; index->dtype.bits = dtype.itemsize * 8; + auto try_layout_deser = [&](auto tag) { + using data_t = decltype(tag); + try { + return reinterpret_cast( + _mg_deserialize(res, filename, mg_cagra_dataset_layout::device_padded)); + } catch (...) { + return reinterpret_cast( + _mg_deserialize(res, filename, mg_cagra_dataset_layout::device_standard)); + } + }; if (dtype.kind == 'f' && dtype.itemsize == 4) { index->dtype.code = kDLFloat; - index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + index->addr = try_layout_deser(float{}); } else if (dtype.kind == 'e' && dtype.itemsize == 2) { index->dtype.code = kDLFloat; - index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + index->addr = try_layout_deser(half{}); } else if (dtype.kind == 'i' && dtype.itemsize == 1) { index->dtype.code = kDLInt; - index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + index->addr = try_layout_deser(int8_t{}); } else if (dtype.kind == 'u' && dtype.itemsize == 1) { index->dtype.code = kDLUInt; - index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + index->addr = try_layout_deser(uint8_t{}); } else { RAFT_FAIL("Unsupported index dtype"); } @@ -442,19 +533,31 @@ extern "C" cuvsError_t cuvsMultiGpuCagraDistribute(cuvsResources_t res, raft::numpy_serializer::parse_descr(std::string(dtype_string, sizeof(dtype_string))); is.close(); + destroy_mg_cagra_c_api_box(index->addr); + index->addr = 0; index->dtype.bits = dtype.itemsize * 8; + auto try_layout_distribute = [&](auto tag) { + using data_t = decltype(tag); + try { + return reinterpret_cast( + _mg_distribute(res, filename, mg_cagra_dataset_layout::device_padded)); + } catch (...) { + return reinterpret_cast( + _mg_distribute(res, filename, mg_cagra_dataset_layout::device_standard)); + } + }; if (dtype.kind == 'f' && dtype.itemsize == 4) { index->dtype.code = kDLFloat; - index->addr = reinterpret_cast(_mg_distribute(res, filename)); + index->addr = try_layout_distribute(float{}); } else if (dtype.kind == 'e' && dtype.itemsize == 2) { index->dtype.code = kDLFloat; - index->addr = reinterpret_cast(_mg_distribute(res, filename)); + index->addr = try_layout_distribute(half{}); } else if (dtype.kind == 'i' && dtype.itemsize == 1) { index->dtype.code = kDLInt; - index->addr = reinterpret_cast(_mg_distribute(res, filename)); + index->addr = try_layout_distribute(int8_t{}); } else if (dtype.kind == 'u' && dtype.itemsize == 1) { index->dtype.code = kDLUInt; - index->addr = reinterpret_cast(_mg_distribute(res, filename)); + index->addr = try_layout_distribute(uint8_t{}); } else { RAFT_FAIL("Unsupported index dtype"); } diff --git a/c/src/neighbors/tiered_index.cpp b/c/src/neighbors/tiered_index.cpp index 2a7d54b16d..e9a0de6d64 100644 --- a/c/src/neighbors/tiered_index.cpp +++ b/c/src/neighbors/tiered_index.cpp @@ -1,10 +1,11 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include #include +#include #include #include @@ -18,17 +19,85 @@ #include -#include - #include "../core/exceptions.hpp" #include "../core/interop.hpp" #include "cagra.hpp" +#include "c_api_box.hpp" #include "ivf_flat.hpp" #include "ivf_pq.hpp" namespace { using namespace cuvs::neighbors; +enum class tiered_cagra_dataset_layout : uint8_t { not_applicable, device_padded, device_standard }; + +struct tiered_c_api_index_box { + void* index_ptr; + cuvsTieredIndexANNAlgo algo; + tiered_cagra_dataset_layout cagra_layout; + cuvs::neighbors::c_api::detail::owner_record owner_rec; +}; + +template +static auto make_tiered_index_box(tiered_index::index* ptr, + cuvsTieredIndexANNAlgo algo, + tiered_cagra_dataset_layout cagra_layout = + tiered_cagra_dataset_layout::not_applicable) + -> tiered_c_api_index_box* +{ + return new tiered_c_api_index_box{ + ptr, algo, cagra_layout, cuvs::neighbors::c_api::detail::make_owner_record(ptr)}; +} + +static auto require_tiered_index_box(cuvsTieredIndex const& index, const char* null_handle_err) + -> tiered_c_api_index_box* +{ + auto* box = reinterpret_cast(index.addr); + if (box == nullptr) { RAFT_FAIL("%s", null_handle_err); } + return box; +} + +template +static void with_tiered_cagra_index_by_layout(tiered_c_api_index_box* box, + const char* null_handle_err, + Fn&& fn) +{ + if (box == nullptr) { RAFT_FAIL("%s", null_handle_err); } + RAFT_EXPECTS(box->algo == CUVS_TIERED_INDEX_ALGO_CAGRA, + "with_tiered_cagra_index_by_layout requires CAGRA algo"); + if (box->cagra_layout == tiered_cagra_dataset_layout::device_padded) { + auto* index_ptr = reinterpret_cast>*>( + box->index_ptr); + fn(index_ptr); + } else if (box->cagra_layout == tiered_cagra_dataset_layout::device_standard) { + auto* index_ptr = reinterpret_cast>*>( + box->index_ptr); + fn(index_ptr); + } else { + RAFT_FAIL("Tiered CAGRA index layout is not initialized"); + } +} + +template +static auto require_typed_tiered_index(cuvsTieredIndex const& index, + cuvsTieredIndexANNAlgo expected_algo, + const char* null_handle_err, + const char* wrong_algo_err) + -> tiered_index::index* +{ + auto* box = require_tiered_index_box(index, null_handle_err); + if (box->algo != expected_algo) { RAFT_FAIL("%s", wrong_algo_err); } + return reinterpret_cast*>(box->index_ptr); +} + +template +struct tiered_upstream_type; + +template +struct tiered_upstream_type> { + using type = UpstreamT; +}; + template void convert_c_index_params(cuvsTieredIndexParams params, int64_t n_rows, @@ -71,20 +140,31 @@ void* _build(cuvsResources_t res, cuvsTieredIndexParams params, DLManagedTensor* case CUVS_TIERED_INDEX_ALGO_CAGRA: { auto build_params = tiered_index::index_params(); convert_c_index_params(params, dataset.shape[0], dataset.shape[1], &build_params); - return new tiered_index::index>( + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(mds)) { + auto padded_view = cuvs::neighbors::make_device_padded_dataset_view(*res_ptr, mds); + auto* ptr = new tiered_index::index>( + tiered_index::build(*res_ptr, build_params, padded_view)); + return make_tiered_index_box( + ptr, CUVS_TIERED_INDEX_ALGO_CAGRA, tiered_cagra_dataset_layout::device_padded); + } + auto* ptr = new tiered_index::index>( tiered_index::build(*res_ptr, build_params, mds)); + return make_tiered_index_box( + ptr, CUVS_TIERED_INDEX_ALGO_CAGRA, tiered_cagra_dataset_layout::device_standard); } case CUVS_TIERED_INDEX_ALGO_IVF_FLAT: { auto build_params = tiered_index::index_params(); convert_c_index_params(params, dataset.shape[0], dataset.shape[1], &build_params); - return new tiered_index::index>( + auto* ptr = new tiered_index::index>( tiered_index::build(*res_ptr, build_params, mds)); + return make_tiered_index_box(ptr, CUVS_TIERED_INDEX_ALGO_IVF_FLAT); } case CUVS_TIERED_INDEX_ALGO_IVF_PQ: { auto build_params = tiered_index::index_params(); convert_c_index_params(params, dataset.shape[0], dataset.shape[1], &build_params); - return new tiered_index::index>( + auto* ptr = new tiered_index::index>( tiered_index::build(*res_ptr, build_params, mds)); + return make_tiered_index_box(ptr, CUVS_TIERED_INDEX_ALGO_IVF_PQ); } default: RAFT_FAIL("unsupported tiered index algorithm"); } @@ -100,7 +180,8 @@ void _search(cuvsResources_t res, cuvsFilter filter) { auto res_ptr = reinterpret_cast(res); - auto index_ptr = reinterpret_cast*>(index.addr); + auto index_ptr = require_typed_tiered_index( + index, index.algo, "_search: null index handle", "_search: index algorithm mismatch"); auto search_params = typename UpstreamT::search_params_type(); if (params != NULL) { @@ -154,7 +235,8 @@ template void _extend(cuvsResources_t res, DLManagedTensor* new_vectors, cuvsTieredIndex index) { auto res_ptr = reinterpret_cast(res); - auto index_ptr = reinterpret_cast*>(index.addr); + auto index_ptr = require_typed_tiered_index( + index, index.algo, "_extend: null index handle", "_extend: index algorithm mismatch"); using T = typename UpstreamT::value_type; using vectors_mdspan_type = raft::device_matrix_view; @@ -174,6 +256,11 @@ void _merge(cuvsResources_t res, std::vector*> cpp_indices; int64_t n_rows = 0, dim = 0; + tiered_cagra_dataset_layout expected_layout = tiered_cagra_dataset_layout::not_applicable; + if (indices[0]->algo == CUVS_TIERED_INDEX_ALGO_CAGRA) { + auto* box0 = require_tiered_index_box(*indices[0], "_merge: null index handle"); + expected_layout = box0->cagra_layout; + } for (size_t i = 0; i < num_indices; ++i) { RAFT_EXPECTS(indices[i]->dtype.code == indices[0]->dtype.code, "indices must all have the same dtype"); @@ -181,9 +268,16 @@ void _merge(cuvsResources_t res, "indices must all have the same dtype"); RAFT_EXPECTS(indices[i]->algo == indices[0]->algo, "indices must all have the same index algorithm"); + if (indices[i]->algo == CUVS_TIERED_INDEX_ALGO_CAGRA) { + auto* boxi = require_tiered_index_box(*indices[i], "_merge: null index handle"); + RAFT_EXPECTS(boxi->cagra_layout == expected_layout, + "indices must all have the same CAGRA layout"); + } - auto idx_ptr = - reinterpret_cast*>(indices[i]->addr); + auto idx_ptr = require_typed_tiered_index(*indices[i], + indices[0]->algo, + "_merge: null index handle", + "_merge: index algorithm mismatch"); n_rows += idx_ptr->size(); if (dim) { RAFT_EXPECTS(dim == idx_ptr->dim(), "indices must all have the same dimensionality"); @@ -196,10 +290,15 @@ void _merge(cuvsResources_t res, auto build_params = tiered_index::index_params(); convert_c_index_params(params, n_rows, dim, &build_params); - auto ptr = new cuvs::neighbors::tiered_index::index( + auto* ptr = new cuvs::neighbors::tiered_index::index( cuvs::neighbors::tiered_index::merge(*res_ptr, build_params, cpp_indices)); - - output_index->addr = reinterpret_cast(ptr); + auto cagra_layout = tiered_cagra_dataset_layout::not_applicable; + if (indices[0]->algo == CUVS_TIERED_INDEX_ALGO_CAGRA) { + auto* box0 = require_tiered_index_box(*indices[0], "_merge: null index handle"); + cagra_layout = box0->cagra_layout; + } + auto* box = make_tiered_index_box(ptr, indices[0]->algo, cagra_layout); + output_index->addr = reinterpret_cast(box); output_index->dtype = indices[0]->dtype; output_index->algo = indices[0]->algo; } @@ -215,28 +314,10 @@ extern "C" cuvsError_t cuvsTieredIndexDestroy(cuvsTieredIndex_t index_c_ptr) { return cuvs::core::translate_exceptions([=] { auto index = *index_c_ptr; - if (index.dtype.code == kDLFloat && index.dtype.bits == 32) { - switch (index.algo) { - case CUVS_TIERED_INDEX_ALGO_CAGRA: { - auto index_ptr = - reinterpret_cast>*>(index.addr); - delete index_ptr; - break; - } - case CUVS_TIERED_INDEX_ALGO_IVF_FLAT: { - auto index_ptr = - reinterpret_cast>*>(index.addr); - delete index_ptr; - break; - } - case CUVS_TIERED_INDEX_ALGO_IVF_PQ: { - auto index_ptr = - reinterpret_cast>*>(index.addr); - delete index_ptr; - break; - } - default: RAFT_FAIL("unsupported tiered index algorithm"); - } + auto* box = reinterpret_cast(index.addr); + if (box != nullptr) { + cuvs::neighbors::c_api::detail::destroy_owner_record(box->owner_rec); + delete box; } delete index_c_ptr; }); @@ -292,8 +373,13 @@ extern "C" cuvsError_t cuvsTieredIndexSearch(cuvsResources_t res, switch (index.algo) { case CUVS_TIERED_INDEX_ALGO_CAGRA: { - _search>( - res, search_params, index, queries_tensor, neighbors_tensor, distances_tensor, filter); + auto* box = require_tiered_index_box(index, "cuvsTieredIndexSearch: null index handle"); + with_tiered_cagra_index_by_layout( + box, "cuvsTieredIndexSearch: null index handle", [&](auto* typed_index) { + using ann_t = typename tiered_upstream_type>::type; + _search( + res, search_params, index, queries_tensor, neighbors_tensor, distances_tensor, filter); + }); break; } case CUVS_TIERED_INDEX_ALGO_IVF_FLAT: { @@ -336,7 +422,12 @@ extern "C" cuvsError_t cuvsTieredIndexExtend(cuvsResources_t res, auto index = *index_c_ptr; switch (index.algo) { case CUVS_TIERED_INDEX_ALGO_CAGRA: { - _extend>(res, new_vectors, index); + auto* box = require_tiered_index_box(index, "cuvsTieredIndexExtend: null index handle"); + with_tiered_cagra_index_by_layout( + box, "cuvsTieredIndexExtend: null index handle", [&](auto* typed_index) { + using ann_t = typename tiered_upstream_type>::type; + _extend(res, new_vectors, index); + }); break; } case CUVS_TIERED_INDEX_ALGO_IVF_FLAT: { @@ -363,7 +454,12 @@ extern "C" cuvsError_t cuvsTieredIndexMerge(cuvsResources_t res, switch (indices[0]->algo) { case CUVS_TIERED_INDEX_ALGO_CAGRA: { - _merge>(res, *params, indices, num_indices, output_index); + auto* box = require_tiered_index_box(*indices[0], "cuvsTieredIndexMerge: null index handle"); + with_tiered_cagra_index_by_layout( + box, "cuvsTieredIndexMerge: null index handle", [&](auto* typed_index) { + using ann_t = typename tiered_upstream_type>::type; + _merge(res, *params, indices, num_indices, output_index); + }); break; } case CUVS_TIERED_INDEX_ALGO_IVF_FLAT: { diff --git a/c/tests/neighbors/ann_cagra_c.cu b/c/tests/neighbors/ann_cagra_c.cu index 9c14bbea7d..f5c66156e7 100644 --- a/c/tests/neighbors/ann_cagra_c.cu +++ b/c/tests/neighbors/ann_cagra_c.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -165,12 +165,23 @@ TEST(CagraC, BuildExtendSearch) (main_data_size + additional_data_size + num_queries) * dimensions, stream); rmm::device_uvector random_labels_d( (main_data_size + additional_data_size + num_queries) * dimensions, stream); - raft::random::make_blobs(random_data_d.data(), - random_labels_d.data(), - main_data_size + additional_data_size + num_queries, - dimensions, - 10, - stream); + + raft::random::make_blobs( + random_data_d.data(), + random_labels_d.data(), + main_data_size + additional_data_size + num_queries, + dimensions, + static_cast(10), + stream, + true, + nullptr, + nullptr, + 1.0f, + true, + -10.0f, + 10.0f, + 42ULL, + raft::random::GenPC); // create dataset DLTensor rmm::device_uvector main_d(main_data_size * dimensions, stream); @@ -218,7 +229,11 @@ TEST(CagraC, BuildExtendSearch) // extend index cuvsCagraExtendParams_t extend_params; cuvsCagraExtendParamsCreate(&extend_params); - cuvsCagraExtend(res, extend_params, &additional_dataset_tensor, index); + cuvsDatasetStorage_t extended_dataset = nullptr; + ASSERT_EQ(cuvsMakeExtendedStorage(res, &additional_dataset_tensor, index, &extended_dataset), + CUVS_SUCCESS); + ASSERT_EQ(cuvsCagraExtend(res, extend_params, &additional_dataset_tensor, extended_dataset, index), + CUVS_SUCCESS); // create queries DLTensor rmm::device_uvector queries_d(num_queries * dimensions, stream); @@ -321,7 +336,6 @@ TEST(CagraC, BuildExtendSearch) cuvsCagraSearch( res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter); - // check neighbors ASSERT_TRUE( cuvs::devArrMatch(min_cols.data_handle(), neighbors_d.data(), 4, cuvs::Compare())); @@ -332,6 +346,7 @@ TEST(CagraC, BuildExtendSearch) // de-allocate index and res cuvsCagraSearchParamsDestroy(search_params); cuvsCagraExtendParamsDestroy(extend_params); + cuvsDatasetStorageDestroy(extended_dataset); cuvsCagraIndexParamsDestroy(build_params); cuvsCagraIndexDestroy(index); cuvsResourcesDestroy(res); @@ -511,7 +526,32 @@ TEST(CagraC, BuildMergeSearch) filter.addr = 0; cuvsCagraIndex_t index_array[2] = {index_main, index_add}; - ASSERT_EQ(cuvsCagraMerge(res, build_params, index_array, 2, filter, index_merged), CUVS_SUCCESS); + cuvsDatasetStorage_t merged_dataset = nullptr; + ASSERT_EQ(cuvsMakeMergedStorage(res, index_array, 2, filter, &merged_dataset), CUVS_SUCCESS); + ASSERT_EQ(cuvsCagraMerge(res, build_params, index_array, 2, filter, merged_dataset, index_merged), + CUVS_SUCCESS); + + // Merge of standard-layout device inputs yields a standard index. Under the explicit C API + // contract, attach a padded dataset before calling search. + 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); + + DLManagedTensor merged_dataset_tensor; + int64_t merged_shape[2] = {7, 2}; + merged_dataset_tensor.dl_tensor.data = merged_d.data(); + merged_dataset_tensor.dl_tensor.device.device_type = kDLCUDA; + merged_dataset_tensor.dl_tensor.device.device_id = 0; + merged_dataset_tensor.dl_tensor.ndim = 2; + merged_dataset_tensor.dl_tensor.dtype.code = kDLFloat; + merged_dataset_tensor.dl_tensor.dtype.bits = 32; + merged_dataset_tensor.dl_tensor.dtype.lanes = 1; + merged_dataset_tensor.dl_tensor.shape = merged_shape; + merged_dataset_tensor.dl_tensor.strides = nullptr; + + cuvsDatasetPadded_t padded_dataset = nullptr; + ASSERT_EQ(cuvsDatasetMakePadded(res, &merged_dataset_tensor, &padded_dataset), CUVS_SUCCESS); + ASSERT_EQ(cuvsCagraAttachPaddedDatasetForSearch(res, padded_dataset, index_merged), CUVS_SUCCESS); int64_t merged_dim = -1; ASSERT_EQ(cuvsCagraIndexGetDims(index_merged, &merged_dim), CUVS_SUCCESS); @@ -562,6 +602,7 @@ TEST(CagraC, BuildMergeSearch) EXPECT_NEAR(distance_host, 0.0f, 1e-6); cuvsCagraSearchParamsDestroy(search_params); + cuvsDatasetStorageDestroy(merged_dataset); cuvsCagraIndexParamsDestroy(build_params); cuvsCagraIndexDestroy(index_merged); cuvsCagraIndexDestroy(index_add); diff --git a/ci/validate_wheel.sh b/ci/validate_wheel.sh index 677f3a63f1..463c7275d1 100755 --- a/ci/validate_wheel.sh +++ b/ci/validate_wheel.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail @@ -20,6 +20,7 @@ PYDISTCHECK_ARGS=( # PyPI hard limit is 1GiB, but try to keep these as small as possible if [[ "${package_dir}" == "python/libcuvs" ]]; then if [[ "${RAPIDS_CUDA_MAJOR}" == "12" ]]; then + # Cap is below PyPI’s 1 GiB limit; raise when the shipped libcuvs.so grows. PYDISTCHECK_ARGS+=( --max-allowed-size-compressed '350Mi' ) diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h index 57b47d97db..c5228f8580 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -330,13 +330,6 @@ void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::cagra::index std::max(params.graph_degree, params.intermediate_graph_degree); } - nlohmann::json comp_search_conf = collect_conf_with_prefix(conf, "compression_"); - if (!comp_search_conf.empty()) { - auto vpq_pams = params.compression.value_or(cuvs::neighbors::vpq_params{}); - parse_build_param(comp_search_conf, vpq_pams); - params.compression.emplace(vpq_pams); - } - if (conf.contains("guarantee_connectivity")) { params.guarantee_connectivity = conf.at("guarantee_connectivity"); } diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_diskann_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_diskann_wrapper.h index 24246feda3..cb5c64150a 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_diskann_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_diskann_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -10,6 +10,8 @@ #include #include +#include +#include #include "../common/ann_types.hpp" #include "../diskann/diskann_wrapper.h" @@ -165,18 +167,28 @@ void cuvs_cagra_diskann::save(const std::string& file) const // try allocating a buffer for the dataset on host try { - const cuvs::neighbors::strided_dataset* strided_dataset = - dynamic_cast*>( - const_cast*>(&cagra_build_.get_index()->data())); - if (strided_dataset == nullptr) { - RAFT_LOG_DEBUG("dynamic_cast to strided_dataset failed"); + auto const* idx_ptr = cagra_build_.get_index(); + std::optional> h_dataset = std::nullopt; + auto const& data_view = idx_ptr->data(); + if constexpr (cuvs::neighbors::is_padded_dataset_view_v>) { + auto const& v = data_view; + auto n_rows = v.n_rows(); + auto dim = v.dim(); + auto stride = v.stride(); + h_dataset.emplace(raft::make_host_matrix(n_rows, dim)); + raft::copy_matrix(h_dataset->data_handle(), + dim, + v.view().data_handle(), + stride, + dim, + n_rows, + raft::resource::get_cuda_stream(handle_)); } else { - auto h_dataset = - raft::make_host_matrix(strided_dataset->n_rows(), strided_dataset->dim()); - raft::copy(h_dataset.data_handle(), - strided_dataset->view().data_handle(), - strided_dataset->n_rows() * strided_dataset->dim(), - raft::resource::get_cuda_stream(handle_)); + RAFT_LOG_DEBUG("dataset serialization: index dataset is not device_padded_dataset_view"); + } + + if (h_dataset.has_value()) { + raft::resource::sync_stream(handle_); std::string dataset_base_file = file + ".data"; std::ofstream dataset_of(dataset_base_file, std::ios::out | std::ios::binary); if (!dataset_of) { RAFT_FAIL("Cannot open file %s", dataset_base_file.c_str()); } @@ -187,7 +199,7 @@ void cuvs_cagra_diskann::save(const std::string& file) const dataset_of.write((char*)&size, sizeof(int)); dataset_of.write((char*)&dim, sizeof(int)); for (int i = 0; i < size; i++) { - dataset_of.write((char*)(h_dataset.data_handle() + i * h_dataset.extent(1)), + dataset_of.write((char*)(h_dataset->data_handle() + i * h_dataset->extent(1)), dim * sizeof(T)); } dataset_of.close(); diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h index db618f6559..e24802a298 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -82,11 +82,7 @@ class cuvs_cagra_hnswlib : public algo, public algo_gpu { template void cuvs_cagra_hnswlib::build(const T* dataset, size_t nrow) { - // when the data set is on host, we can pass it directly to HNSW - bool dataset_is_on_host = raft::get_device_for_address(dataset) == -1; - auto dataset_view = raft::make_host_matrix_view(dataset, nrow, this->dim_); - // convert the index to HNSW format hnsw_index_ = cuvs::neighbors::hnsw::build(handle_, build_param_.hnsw_index_params, dataset_view); } diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h index 87111e4761..6d67f85607 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -74,6 +74,7 @@ enum class CagraMergeType { kPhysical, kLogical }; template class cuvs_cagra : public algo, public algo_gpu { public: + using index_type = cuvs::neighbors::cagra::device_padded_index; using search_param_base = typename algo::search_param; using algo::dim_; using algo::metric_; @@ -162,7 +163,7 @@ class cuvs_cagra : public algo, public algo_gpu { void save_to_hnswlib(const std::string& file) const; std::unique_ptr> copy() override; - auto get_index() const -> const cuvs::neighbors::cagra::index* { return index_.get(); } + auto get_index() const -> const index_type* { return index_.get(); } private: // handle_ must go first to make sure it dies last and all memory allocated in pool @@ -175,7 +176,7 @@ class cuvs_cagra : public algo, public algo_gpu { build_param index_params_; bool need_dataset_update_{true}; cuvs::neighbors::cagra::search_params search_params_; - std::shared_ptr> index_; + std::shared_ptr index_; std::shared_ptr> graph_; std::shared_ptr> dataset_; std::shared_ptr> input_dataset_v_; @@ -188,7 +189,13 @@ class cuvs_cagra : public algo, public algo_gpu { bool dynamic_batching_conservative_dispatch_; std::shared_ptr filter_; - std::vector>> sub_indices_; + std::vector> sub_indices_; + std::shared_ptr>> + sub_dataset_buffers_ = + std::make_shared>>(); + std::shared_ptr> deserialized_dataset_; + std::vector>> + sub_deserialized_datasets_; inline rmm::device_async_resource_ref get_mr(AllocatorType mem_type) { @@ -206,15 +213,87 @@ void cuvs_cagra::build(const T* dataset, size_t nrow) auto dataset_extents = raft::make_extents(nrow, dim_); auto params = index_params_.cagra_params(dataset_extents, parse_metric_type(metric_)); + // Use int64_t throughout so that device copies are compatible with dataset_ (device_matrix) and so that host padded dataset views carry the correct index type. + auto dataset_extents_i64 = + raft::make_extents(static_cast(nrow), static_cast(dim_)); auto dataset_view_host = - raft::make_mdspan(dataset, dataset_extents); - auto dataset_view_device = - raft::make_mdspan(dataset, dataset_extents); + raft::make_mdspan(dataset, dataset_extents_i64); bool dataset_is_on_host = raft::get_device_for_address(dataset) == -1; + // Host mdspan + ace_params: `cagra::build` dispatches to ACE. Non-ACE from host uses padded + // uses `cagra::build(res, params, dataset_view)` with a padded device dataset (or upload + // host data first). Used for both single-split and logical multi-split build paths. + bool const use_ace_host = + dataset_is_on_host && std::holds_alternative( + params.graph_build_params); if (index_params_.num_dataset_splits <= 1) { - index_ = std::make_shared>(std::move( - dataset_is_on_host ? cuvs::neighbors::cagra::build(handle_, params, dataset_view_host) - : cuvs::neighbors::cagra::build(handle_, params, dataset_view_device))); + if (use_ace_host) { + // ACE build is always graph-only; build the graph from a host_padded_dataset_view (required + // by the new build() API), then upload and attach a device padded copy for search. + // The input data may not satisfy CAGRA's per-row alignment; create an owning host-padded + // copy when needed, or a zero-copy view when the stride already matches. + const uint32_t req_stride = + cuvs::neighbors::cagra_required_row_width(static_cast(dim_), 16); + std::unique_ptr> host_padded_own; + std::optional> host_pdv; + if (static_cast(dim_) == req_stride) { + host_pdv = cuvs::neighbors::make_host_padded_dataset_view(dataset_view_host); + } else { + host_padded_own = cuvs::neighbors::make_host_padded_dataset(handle_, dataset_view_host); + host_pdv = host_padded_own->as_dataset_view(); + } + auto ace_host_index = cuvs::neighbors::cagra::build(handle_, params, *host_pdv); + auto padded = cuvs::neighbors::make_device_padded_dataset(handle_, dataset_view_host); + auto ace_index = cuvs::neighbors::cagra::attach_device_dataset_on_host_index( + handle_, ace_host_index, padded->as_dataset_view()); + *dataset_ = std::move(padded->data_); + index_ = std::make_shared(std::move(ace_index)); + } else { + // Non-ACE CAGRA build must use cagra::build(res, params, dataset_view) from + // make_device_padded_dataset / make_device_padded_dataset_view; the host mdspan and raw + // device mdspan entry points are not valid for these graph types. + // Host + non-ACE: copy to a device buffer first, then use the same path + // as a native device pointer. + raft::device_matrix_view mds; + if (dataset_is_on_host) { + *dataset_ = std::move(raft::make_device_matrix( + handle_, static_cast(nrow), static_cast(dim_))); + raft::copy(dataset_->data_handle(), + dataset, + static_cast(nrow) * dim_, + raft::resource::get_cuda_stream(handle_)); + mds = raft::make_device_matrix_view( + dataset_->data_handle(), static_cast(nrow), static_cast(dim_)); + } else { + mds = raft::make_device_matrix_view( + dataset, static_cast(nrow), static_cast(dim_)); + } + const uint32_t required_stride = + cuvs::neighbors::cagra_required_row_width(static_cast(mds.extent(1)), 16); + const uint32_t src_stride = mds.stride(0) > 0 ? static_cast(mds.stride(0)) + : static_cast(mds.extent(1)); + cudaPointerAttributes ptr_attrs{}; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, mds.data_handle())); + const bool device_src = (reinterpret_cast(ptr_attrs.devicePointer) != nullptr); + // `cagra::index` is move-only; use a non-const `index` per branch so + // `std::move(index)` moves (a const `index` would try to copy the deleted + // cagra::index copy ctor). + if (device_src && src_stride == required_stride) { + auto const pdv = cuvs::neighbors::make_device_padded_dataset_view(handle_, mds); + *input_dataset_v_ = raft::make_device_matrix_view( + mds.data_handle(), static_cast(nrow), static_cast(dim_)); + auto index = cuvs::neighbors::cagra::build(handle_, params, pdv); + index.update_dataset(handle_, pdv); + index_ = std::make_shared(std::move(index)); + } else { + auto padded = cuvs::neighbors::make_device_padded_dataset(handle_, mds); + auto view = padded->as_dataset_view(); + auto index = cuvs::neighbors::cagra::build(handle_, params, view); + index.update_dataset(handle_, view); + *dataset_ = std::move(padded->data_); + index_ = std::make_shared(std::move(index)); + } + } } else { IdxT rows_per_split = raft::ceildiv(nrow, static_cast(index_params_.num_dataset_splits)); @@ -225,37 +304,120 @@ void cuvs_cagra::build(const T* dataset, size_t nrow) const T* sub_ptr = dataset + static_cast(start) * dim_; auto sub_host = raft::make_host_matrix_view(sub_ptr, rows, dim_); - auto sub_dev = - raft::make_device_matrix_view(sub_ptr, rows, dim_); + auto sub_dev = raft::make_device_matrix_view( + sub_ptr, static_cast(rows), static_cast(dim_)); - auto sub_index = cuvs::neighbors::cagra::index(handle_, params.metric); + auto sub_index = index_type(handle_, params.metric); if (index_params_.merge_type == CagraMergeType::kPhysical) { if (dataset_is_on_host) { - sub_index.update_dataset(handle_, sub_host); + sub_dataset_buffers_->emplace_back( + raft::make_device_matrix(handle_, rows, dim_)); + raft::copy(sub_dataset_buffers_->back().data_handle(), + sub_ptr, + static_cast(rows) * dim_, + raft::resource::get_cuda_stream(handle_)); + cuvs::neighbors::device_padded_dataset_view dv( + raft::make_const_mdspan(sub_dataset_buffers_->back().view()), dim_); + sub_index.update_dataset(handle_, dv); } else { - sub_index.update_dataset(handle_, sub_dev); + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(sub_dev)) { + auto pdv = cuvs::neighbors::make_device_padded_dataset_view(handle_, sub_dev); + sub_index.update_dataset(handle_, pdv); + } else { + auto padded = cuvs::neighbors::make_device_padded_dataset(handle_, sub_dev); + sub_dataset_buffers_->push_back(std::move(padded->data_)); + cuvs::neighbors::device_padded_dataset_view pdv( + raft::make_const_mdspan(sub_dataset_buffers_->back().view()), dim_); + sub_index.update_dataset(handle_, pdv); + } } } if (index_params_.merge_type == CagraMergeType::kLogical) { - if (dataset_is_on_host) { - sub_index = cuvs::neighbors::cagra::build(handle_, params, sub_host); + if (use_ace_host) { + // ACE build is always graph-only; build the graph from a host_padded_dataset_view + // (required by the new build() API), then upload and attach a device padded copy. + const uint32_t req_stride_sub = + cuvs::neighbors::cagra_required_row_width(static_cast(dim_), 16); + std::unique_ptr> host_padded_sub_own; + std::optional> host_pdv_sub; + if (static_cast(dim_) == req_stride_sub) { + host_pdv_sub = cuvs::neighbors::make_host_padded_dataset_view(sub_host); + } else { + host_padded_sub_own = cuvs::neighbors::make_host_padded_dataset(handle_, sub_host); + host_pdv_sub = host_padded_sub_own->as_dataset_view(); + } + auto ace_host_index = cuvs::neighbors::cagra::build(handle_, params, *host_pdv_sub); + auto padded_sub = cuvs::neighbors::make_device_padded_dataset(handle_, sub_host); + sub_index = cuvs::neighbors::cagra::attach_device_dataset_on_host_index( + handle_, ace_host_index, padded_sub->as_dataset_view()); + sub_dataset_buffers_->push_back(std::move(padded_sub->data_)); + } else if (dataset_is_on_host) { + sub_dataset_buffers_->emplace_back(raft::make_device_matrix( + handle_, static_cast(rows), static_cast(dim_))); + raft::copy(sub_dataset_buffers_->back().data_handle(), + sub_ptr, + static_cast(rows) * dim_, + raft::resource::get_cuda_stream(handle_)); + auto mds_sub = raft::make_device_matrix_view( + sub_dataset_buffers_->back().data_handle(), static_cast(rows), dim_); + const uint32_t req_sub = cuvs::neighbors::cagra_required_row_width( + static_cast(mds_sub.extent(1)), 16); + const uint32_t src_sub = mds_sub.stride(0) > 0 ? static_cast(mds_sub.stride(0)) + : static_cast(mds_sub.extent(1)); + cudaPointerAttributes sub_attrs{}; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&sub_attrs, mds_sub.data_handle())); + const bool sub_device = (reinterpret_cast(sub_attrs.devicePointer) != nullptr); + if (sub_device && src_sub == req_sub) { + auto pdv_sub = cuvs::neighbors::make_device_padded_dataset_view(handle_, mds_sub); + sub_index = cuvs::neighbors::cagra::build(handle_, params, pdv_sub); + sub_index.update_dataset(handle_, pdv_sub); + } else { + auto padded_sub = cuvs::neighbors::make_device_padded_dataset(handle_, mds_sub); + auto view = padded_sub->as_dataset_view(); + auto index = cuvs::neighbors::cagra::build(handle_, params, view); + index.update_dataset(handle_, view); + sub_dataset_buffers_->push_back(std::move(padded_sub->data_)); + sub_index = std::move(index); + } } else { - sub_index = cuvs::neighbors::cagra::build(handle_, params, sub_dev); + auto mds_sub = sub_dev; + const uint32_t req_sub = cuvs::neighbors::cagra_required_row_width( + static_cast(mds_sub.extent(1)), 16); + const uint32_t src_sub = mds_sub.stride(0) > 0 ? static_cast(mds_sub.stride(0)) + : static_cast(mds_sub.extent(1)); + cudaPointerAttributes sub_attrs{}; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&sub_attrs, mds_sub.data_handle())); + const bool sub_device = (reinterpret_cast(sub_attrs.devicePointer) != nullptr); + if (sub_device && src_sub == req_sub) { + auto pdv_sub = cuvs::neighbors::make_device_padded_dataset_view(handle_, mds_sub); + sub_index = cuvs::neighbors::cagra::build(handle_, params, pdv_sub); + sub_index.update_dataset(handle_, pdv_sub); + } else { + auto padded_sub = cuvs::neighbors::make_device_padded_dataset(handle_, mds_sub); + auto view = padded_sub->as_dataset_view(); + auto index = cuvs::neighbors::cagra::build(handle_, params, view); + index.update_dataset(handle_, view); + sub_dataset_buffers_->push_back(std::move(padded_sub->data_)); + sub_index = std::move(index); + } } } - auto sub_index_shared = - std::make_shared>(std::move(sub_index)); + auto sub_index_shared = std::make_shared(std::move(sub_index)); sub_indices_.push_back(std::move(sub_index_shared)); } if (index_params_.merge_type == CagraMergeType::kPhysical) { - std::vector*> indices; + std::vector indices; indices.reserve(sub_indices_.size()); for (auto& ptr : sub_indices_) { indices.push_back(ptr.get()); } - index_ = std::make_shared>( - std::move(cuvs::neighbors::cagra::merge(handle_, params, indices))); + cuvs::neighbors::filtering::none_sample_filter merge_row_filter; + auto merge_storage = + cuvs::neighbors::cagra::make_merged_storage(handle_, indices, merge_row_filter); + index_ = std::make_shared( + cuvs::neighbors::cagra::merge(handle_, params, indices, merge_storage, merge_row_filter)); + *dataset_ = std::move(merge_storage.merged_storage); } } } @@ -315,7 +477,9 @@ void cuvs_cagra::set_search_param(const search_param_base& param, // First free up existing memory *dataset_ = raft::make_device_matrix(handle_, 0, 0); - index_->update_dataset(handle_, make_const_mdspan(dataset_->view())); + cuvs::neighbors::device_padded_dataset_view empty_dv( + raft::make_device_matrix_view(static_cast(nullptr), 0, this->dim_), this->dim_); + index_->update_dataset(handle_, empty_dv); // Allocate space using the correct memory resource. RAFT_LOG_DEBUG("moving dataset to new memory space: %s", @@ -324,9 +488,11 @@ void cuvs_cagra::set_search_param(const search_param_base& param, auto mr = get_mr(dataset_mem_); cuvs::neighbors::cagra::detail::copy_with_padding(handle_, *dataset_, *input_dataset_v_, mr); - auto dataset_view = raft::make_device_strided_matrix_view( - dataset_->data_handle(), dataset_->extent(0), this->dim_, dataset_->extent(1)); - index_->update_dataset(handle_, dataset_view); + cuvs::neighbors::device_padded_dataset_view dv( + raft::make_device_matrix_view( + dataset_->data_handle(), dataset_->extent(0), dataset_->extent(1)), + this->dim_); + index_->update_dataset(handle_, dv); need_dataset_update_ = false; needs_dynamic_batcher_update = true; @@ -362,6 +528,7 @@ void cuvs_cagra::set_search_dataset(const T* dataset, size_t nrow) if (index_params_.num_dataset_splits > 1 && index_params_.merge_type == CagraMergeType::kLogical) { bool dataset_is_on_host = raft::get_device_for_address(dataset) == -1; + if (dataset_is_on_host) { sub_dataset_buffers_->clear(); } IdxT rows_per_split = raft::ceildiv(nrow, static_cast(index_params_.num_dataset_splits)); for (size_t i = 0; i < sub_indices_.size(); ++i) { @@ -369,32 +536,43 @@ void cuvs_cagra::set_search_dataset(const T* dataset, size_t nrow) if (start >= nrow) break; IdxT rows = std::min(rows_per_split, static_cast(nrow) - start); const T* sub_ptr = dataset + static_cast(start) * dim_; - auto sub_host = - raft::make_host_matrix_view(sub_ptr, rows, dim_); - auto sub_dev = - raft::make_device_matrix_view(sub_ptr, rows, dim_); + auto sub_dev = raft::make_device_matrix_view( + sub_ptr, static_cast(rows), static_cast(dim_)); auto sub_index = sub_indices_[i].get(); if (index_params_.merge_type == CagraMergeType::kLogical) { if (dataset_is_on_host) { - sub_index->update_dataset(handle_, sub_host); + sub_dataset_buffers_->emplace_back( + raft::make_device_matrix(handle_, rows, dim_)); + raft::copy(sub_dataset_buffers_->back().data_handle(), + sub_ptr, + static_cast(rows) * dim_, + raft::resource::get_cuda_stream(handle_)); + cuvs::neighbors::device_padded_dataset_view dv( + raft::make_const_mdspan(sub_dataset_buffers_->back().view()), dim_); + sub_index->update_dataset(handle_, dv); } else { - sub_index->update_dataset(handle_, sub_dev); + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(sub_dev)) { + auto pdv = cuvs::neighbors::make_device_padded_dataset_view(handle_, sub_dev); + sub_index->update_dataset(handle_, pdv); + } else { + auto padded = cuvs::neighbors::make_device_padded_dataset(handle_, sub_dev); + sub_dataset_buffers_->push_back(std::move(padded->data_)); + cuvs::neighbors::device_padded_dataset_view pdv( + raft::make_const_mdspan(sub_dataset_buffers_->back().view()), dim_); + sub_index->update_dataset(handle_, pdv); + } } } } need_dataset_update_ = false; } else { - using ds_idx_type = decltype(index_->data().n_rows()); - bool is_vpq = - dynamic_cast*>(&index_->data()) || - dynamic_cast*>(&index_->data()); // It can happen that we are re-using a previous algo object which already has // the dataset set. Check if we need update. if (static_cast(input_dataset_v_->extent(0)) != nrow || input_dataset_v_->data_handle() != dataset) { *input_dataset_v_ = raft::make_device_matrix_view(dataset, nrow, this->dim_); - need_dataset_update_ = !is_vpq; // ignore update if this is a VPQ dataset. + need_dataset_update_ = true; } } } @@ -412,11 +590,7 @@ void cuvs_cagra::save(const std::string& file) const f << sub_indices_.size(); f.close(); } else { - using ds_idx_type = decltype(index_->data().n_rows()); - bool is_vpq = - dynamic_cast*>(&index_->data()) || - dynamic_cast*>(&index_->data()); - cuvs::neighbors::cagra::serialize(handle_, file, *index_, is_vpq); + cuvs::neighbors::cagra::serialize(handle_, file, *index_, true); } } @@ -437,22 +611,52 @@ void cuvs_cagra::load(const std::string& file) meta >> count; meta.close(); sub_indices_.clear(); + sub_deserialized_datasets_.resize(count); for (size_t i = 0; i < count; ++i) { std::string subfile = file + (i == 0 ? "" : ".subidx." + std::to_string(i)); - auto sub_index = std::make_shared>(handle_); - cuvs::neighbors::cagra::deserialize(handle_, subfile, sub_index.get()); + auto sub_index = std::make_shared(handle_); + std::unique_ptr> tmp_ds; + cuvs::neighbors::cagra::deserialize(handle_, subfile, sub_index.get(), &tmp_ds); + sub_deserialized_datasets_[i] = + std::shared_ptr>(std::move(tmp_ds)); sub_indices_.push_back(std::move(sub_index)); } } else { - index_ = std::make_shared>(handle_); - cuvs::neighbors::cagra::deserialize(handle_, file, index_.get()); + index_ = std::make_shared(handle_); + deserialized_dataset_.reset(); + std::unique_ptr> tmp_ds; + cuvs::neighbors::cagra::deserialize(handle_, file, index_.get(), &tmp_ds); + deserialized_dataset_ = + std::shared_ptr>(std::move(tmp_ds)); } } template std::unique_ptr> cuvs_cagra::copy() { - return std::make_unique>(std::cref(*this)); // use copy constructor + auto out = std::make_unique>(metric_, dim_, index_params_); + out->refine_ratio_ = refine_ratio_; + out->graph_mem_ = graph_mem_; + out->dataset_mem_ = dataset_mem_; + out->need_dataset_update_ = need_dataset_update_; + out->search_params_ = search_params_; + out->index_ = index_; + out->graph_ = graph_; + out->dataset_ = dataset_; + out->input_dataset_v_ = + std::make_shared>( + *input_dataset_v_); + out->dynamic_batcher_ = dynamic_batcher_; + out->dynamic_batcher_sp_ = dynamic_batcher_sp_; + out->dynamic_batching_max_batch_size_ = dynamic_batching_max_batch_size_; + out->dynamic_batching_n_queues_ = dynamic_batching_n_queues_; + out->dynamic_batching_conservative_dispatch_ = dynamic_batching_conservative_dispatch_; + out->filter_ = filter_; + out->sub_indices_ = sub_indices_; + out->sub_dataset_buffers_ = sub_dataset_buffers_; + out->deserialized_dataset_ = deserialized_dataset_; + out->sub_deserialized_datasets_ = sub_deserialized_datasets_; + return out; } template @@ -482,7 +686,7 @@ void cuvs_cagra::search_base( } else { if (index_params_.merge_type == CagraMergeType::kLogical) { // TODO: index merge must happen outside of search, otherwise what are we benchmarking? - std::vector*> cagra_indices; + std::vector cagra_indices; cagra_indices.reserve(sub_indices_.size()); for (auto& ptr : sub_indices_) { cagra_indices.push_back(ptr.get()); diff --git a/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h index 1c254c4e7e..8df12b9f2a 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -77,7 +77,8 @@ class cuvs_mg_cagra : public algo, public algo_gpu { float refine_ratio_; build_param index_params_; cuvs::neighbors::mg_search_params search_params_; - std::shared_ptr, T, IdxT>> + std::shared_ptr< + cuvs::neighbors::mg_index, T, IdxT>> index_; }; @@ -90,12 +91,13 @@ void cuvs_mg_cagra::build(const T* dataset, size_t nrow) cuvs::neighbors::mg_index_params build_params = params; build_params.mode = index_params_.mode; - auto dataset_view = + auto dataset_mds = raft::make_host_matrix_view(dataset, nrow, dim_); - auto idx = cuvs::neighbors::cagra::build(clique_, build_params, dataset_view); - index_ = - std::make_shared, T, IdxT>>( - std::move(idx)); + auto dataset_view = cuvs::neighbors::make_host_standard_dataset_view(dataset_mds); + auto idx = cuvs::neighbors::cagra::build(clique_, build_params, dataset_view); + index_ = std::make_shared< + cuvs::neighbors::mg_index, T, IdxT>>( + std::move(idx)); } inline auto allocator_to_string(AllocatorType mem_type) -> std::string; @@ -126,9 +128,10 @@ void cuvs_mg_cagra::save(const std::string& file) const template void cuvs_mg_cagra::load(const std::string& file) { - index_ = - std::make_shared, T, IdxT>>( - std::move(cuvs::neighbors::cagra::deserialize(clique_, file))); + index_ = std::make_shared< + cuvs::neighbors::mg_index, T, IdxT>>( + clique_, index_params_.mode); + cuvs::neighbors::cagra::deserialize(clique_, file, index_.get()); } template diff --git a/cpp/cmake/patches/faiss-1.14-cuvs-26.08.diff b/cpp/cmake/patches/faiss-1.14-cuvs-26.08.diff index cceebcde7c..9b81a4a2b3 100644 --- a/cpp/cmake/patches/faiss-1.14-cuvs-26.08.diff +++ b/cpp/cmake/patches/faiss-1.14-cuvs-26.08.diff @@ -1,3 +1,267 @@ +diff --git a/faiss/gpu/impl/BinaryCuvsCagra.cu b/faiss/gpu/impl/BinaryCuvsCagra.cu +index b331fdc..c7b5733 100644 +--- a/faiss/gpu/impl/BinaryCuvsCagra.cu ++++ b/faiss/gpu/impl/BinaryCuvsCagra.cu +@@ -58,7 +58,6 @@ BinaryCuvsCagra::BinaryCuvsCagra( + + index_params_.intermediate_graph_degree = intermediate_graph_degree; + index_params_.graph_degree = graph_degree; +- index_params_.attach_dataset_on_build = store_dataset; + + index_params_.metric = cuvs::distance::DistanceType::BitwiseHamming; + +@@ -110,12 +109,14 @@ BinaryCuvsCagra::BinaryCuvsCagra( + auto dataset_mds = + raft::make_device_matrix_view( + train_dataset, n, dim / 8); ++ auto dataset_view = ++ cuvs::neighbors::make_device_padded_dataset_view(raft_handle, dataset_mds); + + cuvs_index = std::make_shared< + cuvs::neighbors::cagra::index>( + raft_handle, + cuvs::distance::DistanceType::BitwiseHamming, +- dataset_mds, ++ dataset_view, + raft::make_const_mdspan(knn_graph_copy.view())); + } else if (!distances_on_gpu && !knn_graph_on_gpu) { + // copy idx_t (int64_t) host knn_graph to uint32_t host knn_graph +@@ -128,12 +129,14 @@ BinaryCuvsCagra::BinaryCuvsCagra( + + auto dataset_mds = raft::make_host_matrix_view( + train_dataset, n, dim / 8); ++ host_to_device_dataset_ = ++ cuvs::neighbors::make_device_padded_dataset(raft_handle, dataset_mds); + + cuvs_index = std::make_shared< + cuvs::neighbors::cagra::index>( + raft_handle, + cuvs::distance::DistanceType::BitwiseHamming, +- dataset_mds, ++ host_to_device_dataset_->as_dataset_view(), + raft::make_const_mdspan(knn_graph_copy.view())); + } else { + FAISS_THROW_MSG( +@@ -166,17 +169,23 @@ void BinaryCuvsCagra::train(idx_t n, const uint8_t* x) { + if (getDeviceForAddress(x) >= 0) { + auto dataset = raft::make_device_matrix_view( + x, n, dim_ / 8); ++ auto dataset_view = ++ cuvs::neighbors::make_device_padded_dataset_view(raft_handle, dataset); + cuvs_index = std::make_shared< + cuvs::neighbors::cagra::index>( + cuvs::neighbors::cagra::build( +- raft_handle, index_params_, dataset)); ++ raft_handle, index_params_, dataset_view)); ++ store_dataset_ = true; + } else { + auto dataset = raft::make_host_matrix_view( + x, n, dim_ / 8); ++ host_to_device_dataset_ = ++ cuvs::neighbors::make_device_padded_dataset(raft_handle, dataset); + cuvs_index = std::make_shared< + cuvs::neighbors::cagra::index>( + cuvs::neighbors::cagra::build( +- raft_handle, index_params_, dataset)); ++ raft_handle, index_params_, ++ host_to_device_dataset_->as_dataset_view())); + } + } + +@@ -212,14 +221,20 @@ void BinaryCuvsCagra::search( + + if (!store_dataset_) { + if (getDeviceForAddress(storage_) >= 0) { ++ host_to_device_dataset_.reset(); + auto dataset = + raft::make_device_matrix_view( + storage_, n_, dim_ / 8); +- cuvs_index->update_dataset(raft_handle, dataset); ++ auto dataset_view = ++ cuvs::neighbors::make_device_padded_dataset_view(raft_handle, dataset); ++ cuvs_index->update_dataset(raft_handle, dataset_view); + } else { +- auto dataset = raft::make_host_matrix_view( ++ auto host_dataset = raft::make_host_matrix_view( + storage_, n_, dim_ / 8); +- cuvs_index->update_dataset(raft_handle, dataset); ++ host_to_device_dataset_ = ++ cuvs::neighbors::make_device_padded_dataset(raft_handle, host_dataset); ++ cuvs_index->update_dataset(raft_handle, ++ host_to_device_dataset_->as_dataset_view()); + } + store_dataset_ = true; + } +@@ -280,6 +295,7 @@ void BinaryCuvsCagra::search( + + void BinaryCuvsCagra::reset() { + cuvs_index.reset(); ++ host_to_device_dataset_.reset(); + } + + idx_t BinaryCuvsCagra::get_knngraph_degree() const { +diff --git a/faiss/gpu/impl/BinaryCuvsCagra.cuh b/faiss/gpu/impl/BinaryCuvsCagra.cuh +index a14480b..7cbfe39 100644 +--- a/faiss/gpu/impl/BinaryCuvsCagra.cuh ++++ b/faiss/gpu/impl/BinaryCuvsCagra.cuh +@@ -28,12 +28,14 @@ + #include + #include + #include ++#include + #include + + #include + #include + + #include ++#include + + namespace faiss { + +@@ -115,6 +117,10 @@ class BinaryCuvsCagra { + /// Parameters to build CAGRA graph using NN Descent + size_t nn_descent_niter_ = 20; + ++ /// Device padded copy when `storage_` is host memory (KNN-graph ctor path). ++ std::unique_ptr> ++ host_to_device_dataset_; ++ + /// Instance of trained cuVS CAGRA index + std::shared_ptr> + cuvs_index{nullptr}; +diff --git a/faiss/gpu/impl/CuvsCagra.cu b/faiss/gpu/impl/CuvsCagra.cu +index 755817f..0eb03ae 100644 +--- a/faiss/gpu/impl/CuvsCagra.cu ++++ b/faiss/gpu/impl/CuvsCagra.cu +@@ -75,7 +75,6 @@ CuvsCagra::CuvsCagra( + + index_params_.intermediate_graph_degree = intermediate_graph_degree; + index_params_.graph_degree = graph_degree; +- index_params_.attach_dataset_on_build = store_dataset; + index_params_.guarantee_connectivity = guarantee_connectivity; + + if (!ivf_pq_search_params_) { +@@ -133,12 +132,14 @@ CuvsCagra::CuvsCagra( + + auto dataset_mds = raft::make_device_matrix_view( + dataset, n, dim); ++ auto dataset_view = ++ cuvs::neighbors::make_device_padded_dataset_view(raft_handle, dataset_mds); + + cuvs_index = std::make_shared< + cuvs::neighbors::cagra::index>( + raft_handle, + metricFaissToCuvs(metric_, false), +- dataset_mds, ++ dataset_view, + raft::make_const_mdspan(knn_graph_copy.view())); + } else if (!dataset_on_gpu && !knn_graph_on_gpu) { + // copy idx_t (int64_t) host knn_graph to uint32_t host knn_graph +@@ -151,12 +152,14 @@ CuvsCagra::CuvsCagra( + + auto dataset_mds = raft::make_host_matrix_view( + dataset, n, dim); ++ host_to_device_dataset_ = ++ cuvs::neighbors::make_device_padded_dataset(raft_handle, dataset_mds); + + cuvs_index = std::make_shared< + cuvs::neighbors::cagra::index>( + raft_handle, + metricFaissToCuvs(metric_, false), +- dataset_mds, ++ host_to_device_dataset_->as_dataset_view(), + raft::make_const_mdspan(knn_graph_copy.view())); + } else { + FAISS_THROW_MSG( +@@ -203,17 +206,23 @@ void CuvsCagra::train(idx_t n, const data_t* x) { + if (getDeviceForAddress(x) >= 0) { + auto dataset = raft::make_device_matrix_view( + x, n, dim_); ++ auto dataset_view = ++ cuvs::neighbors::make_device_padded_dataset_view(raft_handle, dataset); + cuvs_index = std::make_shared< + cuvs::neighbors::cagra::index>( + cuvs::neighbors::cagra::build( +- raft_handle, index_params_, dataset)); ++ raft_handle, index_params_, dataset_view)); ++ store_dataset_ = true; + } else { + auto dataset = + raft::make_host_matrix_view(x, n, dim_); ++ host_to_device_dataset_ = ++ cuvs::neighbors::make_device_padded_dataset(raft_handle, dataset); + cuvs_index = std::make_shared< + cuvs::neighbors::cagra::index>( + cuvs::neighbors::cagra::build( +- raft_handle, index_params_, dataset)); ++ raft_handle, index_params_, ++ host_to_device_dataset_->as_dataset_view())); + } + } + +@@ -248,13 +257,19 @@ void CuvsCagra::search( + + if (!store_dataset_) { + if (getDeviceForAddress(storage_) >= 0) { ++ host_to_device_dataset_.reset(); + auto dataset = raft::make_device_matrix_view( + storage_, n_, dim_); +- cuvs_index->update_dataset(raft_handle, dataset); ++ auto dataset_view = ++ cuvs::neighbors::make_device_padded_dataset_view(raft_handle, dataset); ++ cuvs_index->update_dataset(raft_handle, dataset_view); + } else { +- auto dataset = raft::make_host_matrix_view( ++ auto host_dataset = raft::make_host_matrix_view( + storage_, n_, dim_); +- cuvs_index->update_dataset(raft_handle, dataset); ++ host_to_device_dataset_ = ++ cuvs::neighbors::make_device_padded_dataset(raft_handle, host_dataset); ++ cuvs_index->update_dataset(raft_handle, ++ host_to_device_dataset_->as_dataset_view()); + } + store_dataset_ = true; + } +@@ -303,6 +318,7 @@ void CuvsCagra::search( + template + void CuvsCagra::reset() { + cuvs_index.reset(); ++ host_to_device_dataset_.reset(); + } + + template +diff --git a/faiss/gpu/impl/CuvsCagra.cuh b/faiss/gpu/impl/CuvsCagra.cuh +index a10e9fb..b5c2bcd 100644 +--- a/faiss/gpu/impl/CuvsCagra.cuh ++++ b/faiss/gpu/impl/CuvsCagra.cuh +@@ -27,13 +27,15 @@ + #include + #include + #include ++#include + #include + + #include + #include + + #include + #include ++#include + + namespace faiss { + +@@ -147,6 +149,10 @@ class CuvsCagra { + /// Parameter to use MST optimization to guarantee graph connectivity + bool guarantee_connectivity_ = false; + ++ /// Device padded copy when `storage_` is host memory (KNN-graph ctor path). ++ std::unique_ptr> ++ host_to_device_dataset_; ++ + /// Instance of trained cuVS CAGRA index + std::shared_ptr> cuvs_index{ + nullptr}; diff --git a/faiss/gpu/impl/CuvsFlatIndex.cu b/faiss/gpu/impl/CuvsFlatIndex.cu index d4a2d99fe..fd364add0 100644 --- a/faiss/gpu/impl/CuvsFlatIndex.cu diff --git a/cpp/cmake/patches/faiss_override.json b/cpp/cmake/patches/faiss_override.json index 6edcee095a..ca975d33c8 100644 --- a/cpp/cmake/patches/faiss_override.json +++ b/cpp/cmake/patches/faiss_override.json @@ -7,7 +7,7 @@ "patches" : [ { "file" : "${current_json_dir}/faiss-1.14-cuvs-26.08.diff", - "issue" : "Multiple fixes for cuVS compatibility", + "issue" : "Multiple fixes for cuVS compatibility. Update Faiss cuVS to be compatible with new Dataset API: update_dataset now takes dataset_view and make_padded_dataset_view must be called beforehand. Loading an index built from a user-provided KNN graph passes dataset_view into cagra::index, not raw mdspan.", "fixed_in" : "" } ] diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index d1937cba27..15dbdd3bf5 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -5,12 +5,12 @@ #pragma once -#include "common.hpp" #include #include #include #include #include +#include #include #include #include @@ -28,6 +28,8 @@ #include #include +#include +#include #include #include #include @@ -107,6 +109,9 @@ namespace neighbors { namespace cagra { // For re-exporting into cagra namespace namespace graph_build_params = cuvs::neighbors::graph_build_params; +namespace detail { +struct fd_transfer; +} /** * @defgroup cagra_cpp_index_params CAGRA index build parameters * @{ @@ -151,12 +156,6 @@ struct index_params : cuvs::neighbors::index_params { size_t intermediate_graph_degree = 128; /** Degree of output graph. */ size_t graph_degree = 64; - /** - * Specify compression parameters if compression is desired. If set, overrides the - * attach_dataset_on_build (and the compressed dataset is always added to the index). - */ - std::optional compression = std::nullopt; - /** Parameters for graph building. * * Set ivf_pq_params, nn_descent_params, ace_params, or iterative_search_params to select the @@ -193,31 +192,32 @@ struct index_params : cuvs::neighbors::index_params { bool guarantee_connectivity = false; /** - * Whether to add the dataset content to the index, i.e.: + * Whether to attach the dataset to the index after graph construction, i.e.: + * + * - `true` (default) means `build` attaches the input dataset as a **non-owning view** to the + * index, so the index is ready to search immediately after `build` returns. The caller is + * responsible for keeping the underlying dataset storage alive for as long as the index is used. + * - `false` means `build` only builds the graph and the caller is expected to attach the dataset + * separately via `cuvs::neighbors::cagra::update_dataset` before searching. * - * - `true` means the index is filled with the dataset vectors and ready to search after calling - * `build` provided there is enough memory available. - * - `false` means `build` only builds the graph and the user is expected to - * update the dataset using cuvs::neighbors::cagra::update_dataset. + * Unlike the legacy behavior, no copy of the dataset is made: the index always stores a view. + * Setting `attach_dataset_on_build = false` is useful when the caller needs to apply specific + * memory placement or transformation (e.g. moving to managed memory) before attaching. + * + * **Note:** this flag is only effective when building from a device dataset view + * (e.g. `device_padded_dataset_view`). For host builds (`host_padded_dataset_view`), it is + * ignored — the returned `host_padded_index` cannot be searched regardless, and the caller must + * always call `attach_device_dataset_on_host_index` to obtain a search-ready device index. * - * Regardless of the value of `attach_dataset_on_build`, the search graph is created using all - * the vectors in the dataset. Setting `attach_dataset_on_build = false` can be useful if - * the user needs to build only the search graph but does not intend to search it using CAGRA - * (e.g. search using another graph search algorithm), or if specific memory placement options - * need to be applied on the dataset before it is attached to the index using `update_dataset`. - * API. * @code{.cpp} - * auto dataset = raft::make_device_matrix(res, n_rows, n_cols); - * // use default index_parameters + * auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host_matrix.view()); * cagra::index_params index_params; - * // update index_params to only build the CAGRA graph + * // Build graph only — caller attaches dataset later. * index_params.attach_dataset_on_build = false; - * auto index = cagra::build(res, index_params, dataset.view()); - * // assert that the dataset is not attached to the index - * ASSERT(index.dataset().extent(0) == 0); - * // update dataset - * index.update_dataset(res, dataset.view()); - * // The index is now ready for search + * auto index = cagra::build(res, index_params, dataset->as_dataset_view()); + * // ASSERT(index.size() == 0); // no dataset yet + * // Attach with a view (storage owned by `dataset`). + * index.update_dataset(res, dataset->as_dataset_view()); * cagra::search(res, search_params, index, queries, neighbors, distances); * @endcode */ @@ -385,6 +385,12 @@ struct extend_params { static_assert(std::is_aggregate_v); static_assert(std::is_aggregate_v); +template > +struct index; + /** * @defgroup cagra_cpp_index CAGRA index type * @{ @@ -398,9 +404,10 @@ static_assert(std::is_aggregate_v); * @tparam T data element type * @tparam IdxT the data type used to store the neighbor indices in the search graph. * It must be large enough to represent values up to dataset.extent(0). + * @tparam DatasetViewT concrete non-owning dataset view type stored by the index * */ -template +template struct CUVS_EXPORT index : cuvs::neighbors::index { using index_params_type = cagra::index_params; using search_params_type = cagra::search_params; @@ -422,7 +429,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { /** Total length of the index (number of vectors). */ [[nodiscard]] constexpr inline auto size() const noexcept -> IdxT { - auto data_rows = dataset_->n_rows(); + auto data_rows = dataset_.n_rows(); if (dataset_fd_.has_value()) { return n_rows_; } return data_rows > 0 ? data_rows : graph_view_.extent(0); } @@ -430,7 +437,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { /** Dimensionality of the data. */ [[nodiscard]] constexpr inline auto dim() const noexcept -> uint32_t { - return dataset_fd_.has_value() ? dim_ : dataset_->dim(); + return dataset_fd_.has_value() ? dim_ : dataset_.dim(); } /** Graph degree */ [[nodiscard]] constexpr inline auto graph_degree() const noexcept -> uint32_t @@ -438,20 +445,8 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { return dataset_fd_.has_value() ? graph_degree_ : graph_view_.extent(1); } - [[nodiscard]] inline auto dataset() const noexcept - -> raft::device_matrix_view - { - auto p = dynamic_cast*>(dataset_.get()); - if (p != nullptr) { return p->view(); } - auto d = dataset_->dim(); - return raft::make_device_strided_matrix_view(nullptr, 0, d, d); - } - - /** Dataset [size, dim] */ - [[nodiscard]] inline auto data() const noexcept -> const cuvs::neighbors::dataset& - { - return *dataset_; - } + /** Non-owning dataset binding stored by the index. */ + [[nodiscard]] inline auto dataset() const noexcept -> DatasetViewT const& { return dataset_; } /** neighborhood graph [size, graph-degree] */ [[nodiscard]] inline auto graph() const noexcept @@ -508,74 +503,49 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { ~index() = default; /** \endcond */ - /** Construct an empty index. */ + /** Construct a graph-only index with a zero-row dataset view placeholder. */ index(raft::resources const& res, cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded) + requires(cuvs::neighbors::ann_dataset_view) : cuvs::neighbors::index(), metric_(metric), graph_(raft::make_device_matrix(res, 0, 0)), - dataset_(new cuvs::neighbors::empty_dataset(0)), + dataset_{}, dataset_norms_(std::nullopt) { } - /** Construct an index from dataset and knn_graph arrays + /** Construct an index from a `dataset_view` and knn_graph. * - * If the dataset and graph is already in GPU memory, then the index is just a thin wrapper around - * these that stores a non-owning a reference to the arrays. + * Stores a shallow copy of the dataset view. The index stores a **non-owning** view; the caller + * must keep underlying device storage alive for the index lifetime. * - * The constructor also accepts host arrays. In that case they are copied to the device, and the - * device arrays will be owned by the index. - * - * In case the dasates rows are not 16 bytes aligned, then we create a padded copy in device - * memory to ensure alignment for vectorized load. - * - * Usage examples: - * - * - Cagra index is normally created by the cagra::build + * Example — **non-owning** `make_device_padded_dataset_view` (wraps an existing device matrix; + * that matrix must outlive the index): * @code{.cpp} - * using namespace cuvs::neighbors; - * auto dataset = raft::make_host_matrix(n_rows, n_cols); - * load_dataset(dataset.view()); - * // use default index parameters - * cagra::index_params index_params; - * // create and fill the index from a [N, D] dataset - * auto index = cagra::build(res, index_params, dataset); - * // use default search parameters - * cagra::search_params search_params; - * // search K nearest neighbours - * auto neighbors = raft::make_device_matrix(res, n_queries, k); - * auto distances = raft::make_device_matrix(res, n_queries, k); - * cagra::search(res, search_params, index, queries, neighbors.view(), distances.view()); + * raft::device_matrix_view dataset = ...; + * auto view = cuvs::neighbors::make_device_padded_dataset_view(res, dataset); + * auto graph = raft::make_device_matrix_view(...); + * cuvs::neighbors::cagra::device_padded_index idx(res, metric, view, + * raft::make_const_mdspan(graph)); * @endcode - * In the above example, we have passed a host dataset to build. The returned index will own a - * device copy of the dataset and the knn_graph. In contrast, if we pass the dataset as a - * device_mdspan to build, then it will only store a reference to it. * - * - Constructing index using existing knn-graph + * Example — **owning** `make_device_padded_dataset` returns owning storage (`std::unique_ptr`). + * You must + * **keep that object alive** (e.g. hold the `unique_ptr` in a variable or member) for as long as + * the index uses the dataset; the index does not take ownership of the buffer. * @code{.cpp} - * using namespace cuvs::neighbors; - * - * auto dataset = raft::make_device_matrix(res, n_rows, n_cols); - * auto knn_graph = raft::make_device_matrix(res, n_rows, graph_degree); - * - * // custom loading and graph creation - * // load_dataset(dataset.view()); - * // create_knn_graph(knn_graph.view()); - * - * // Wrap the existing device arrays into an index structure - * cagra::index index(res, metric, raft::make_const_mdspan(dataset.view()), - * raft::make_const_mdspan(knn_graph.view())); - * - * // Both knn_graph and dataset objects have to be in scope while the index is used because - * // the index only stores a reference to these. - * cagra::search(res, search_params, index, queries, neighbors, distances); + * auto padded_owner = cuvs::neighbors::make_device_padded_dataset(res, dataset_mdspan); + * auto view = padded_owner->as_dataset_view(); + * cuvs::neighbors::cagra::device_padded_index idx(res, metric, view, + * raft::make_const_mdspan(graph)); + * // `padded_owner` must outlive `idx` (do not let it go out of scope while `idx` is used). * @endcode */ - template + template index(raft::resources const& res, cuvs::distance::DistanceType metric, - raft::mdspan, raft::row_major, data_accessor> dataset, + DatasetViewT const& dataset, raft::mdspan, raft::row_major, @@ -583,104 +553,33 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { : cuvs::neighbors::index(), metric_(metric), graph_(raft::make_device_matrix(res, 0, 0)), - dataset_(make_aligned_dataset(res, dataset, 16)), + dataset_(dataset), dataset_norms_(std::nullopt) { - RAFT_EXPECTS(dataset.extent(0) == knn_graph.extent(0), + RAFT_EXPECTS(dataset.n_rows() == static_cast(knn_graph.extent(0)), "Dataset and knn_graph must have equal number of rows"); update_graph(res, knn_graph); if (metric_ == cuvs::distance::DistanceType::CosineExpanded) { - auto p = dynamic_cast*>(dataset_.get()); - if (p) { - auto dataset_view = p->view(); - if (dataset_view.extent(0) > 0) { compute_dataset_norms_(res); } - } + if (dataset.n_rows() > 0) { compute_dataset_norms_(res); } } raft::resource::sync_stream(res); } /** - * Replace the dataset with a new dataset. - * - * If the new dataset rows are aligned on 16 bytes, then only a reference is stored to the - * dataset. It is the caller's responsibility to ensure that dataset stays alive as long as the - * index. It is expected that the same set of vectors are used for update_dataset and index build. - * - * Note: This will clear any precomputed dataset norms. - */ - void update_dataset(raft::resources const& res, - raft::device_matrix_view dataset) - { - dataset_ = make_aligned_dataset(res, dataset, 16); - dataset_norms_.reset(); - - if (metric() == cuvs::distance::DistanceType::CosineExpanded) { - if (dataset.extent(0) > 0) { compute_dataset_norms_(res); } - } - } - - /** Set the dataset reference explicitly to a device matrix view with padding. */ - void update_dataset(raft::resources const& res, - raft::device_matrix_view dataset) - { - dataset_ = make_aligned_dataset(res, dataset, 16); - dataset_norms_.reset(); - - if (metric() == cuvs::distance::DistanceType::CosineExpanded) { - if (dataset.extent(0) > 0) { compute_dataset_norms_(res); } - } - } - - /** - * Replace the dataset with a new dataset. - * - * We create a copy of the dataset on the device. The index manages the lifetime of this copy. It - * is expected that the same set of vectors are used for update_dataset and index build. - * - * Note: This will clear any precomputed dataset norms. - */ - void update_dataset(raft::resources const& res, - raft::host_matrix_view dataset) - { - dataset_ = make_aligned_dataset(res, dataset, 16); - dataset_norms_.reset(); - if (metric() == cuvs::distance::DistanceType::CosineExpanded) { - if (dataset.extent(0) > 0) { compute_dataset_norms_(res); } - } - } - - /** - * Replace the dataset with a new dataset. It is expected that the same set of vectors are used - * for update_dataset and index build. + * Replace the dataset with a new `dataset_view`. * - * Note: This will clear any precomputed dataset norms. + * The index stores a copy of the view handle only (not the vector storage). The caller must + * keep the underlying device data alive. Clears precomputed norms. */ - template - auto update_dataset(raft::resources const& res, DatasetT&& dataset) - -> std::enable_if_t, DatasetT>> - { - dataset_ = std::make_unique(std::move(dataset)); - dataset_norms_.reset(); - if (metric() == cuvs::distance::DistanceType::CosineExpanded) { - auto p = dynamic_cast*>(dataset_.get()); - if (p) { - auto dataset_view = p->view(); - if (dataset_view.extent(0) > 0) { compute_dataset_norms_(res); } - } - } - } - - template - auto update_dataset(raft::resources const& res, std::unique_ptr&& dataset) - -> std::enable_if_t, DatasetT>> + void update_dataset(raft::resources const& res, DatasetViewT const& dataset) + requires cuvs::neighbors::is_device_dataset_view_v { - dataset_ = std::move(dataset); + dataset_ = dataset; dataset_norms_.reset(); if (metric() == cuvs::distance::DistanceType::CosineExpanded) { - auto dataset_view = this->dataset(); - if (dataset_view.extent(0) > 0) { compute_dataset_norms_(res); } + if (dataset_.n_rows() > 0) { compute_dataset_norms_(res); } } } @@ -793,7 +692,19 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { // Re-open the file descriptor in read-only mode for subsequent operations dataset_fd_.emplace(std::move(fd)); - dataset_ = std::make_unique>(0); + if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + auto v = raft::make_device_matrix_view( + static_cast(nullptr), int64_t{0}, dim_); + dataset_ = DatasetViewT(v, dim_); + } else if constexpr (cuvs::neighbors::is_host_padded_dataset_view_v) { + auto v = raft::make_host_matrix_view( + static_cast(nullptr), int64_t{0}, dim_); + dataset_ = DatasetViewT(v, dim_); + } else if constexpr (cuvs::neighbors::is_empty_dataset_view_v) { + dataset_ = DatasetViewT{dim_}; + } else { + RAFT_FAIL("update_dataset(fd): unsupported DatasetViewT for disk-backed dataset"); + } dataset_norms_.reset(); } @@ -874,352 +785,328 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { } private: + friend struct detail::fd_transfer; + + [[nodiscard]] inline auto steal_dataset_fd_() noexcept + -> std::optional + { + return std::exchange(dataset_fd_, std::nullopt); + } + + [[nodiscard]] inline auto steal_graph_fd_() noexcept -> std::optional + { + return std::exchange(graph_fd_, std::nullopt); + } + + [[nodiscard]] inline auto steal_mapping_fd_() noexcept + -> std::optional + { + return std::exchange(mapping_fd_, std::nullopt); + } + cuvs::distance::DistanceType metric_; raft::device_matrix graph_; raft::device_matrix_view graph_view_; - std::unique_ptr> dataset_; + DatasetViewT dataset_; // Mapping from internal graph node indices to the original user-provided indices. std::optional> source_indices_; // only float distances supported at the moment std::optional> dataset_norms_; - // File descriptors for disk-backed index components (ACE disk mode) std::optional dataset_fd_; std::optional graph_fd_; std::optional mapping_fd_; - void compute_dataset_norms_(raft::resources const& res); + CUVS_EXPORT void compute_dataset_norms_(raft::resources const& res); size_t n_rows_ = 0; size_t dim_ = 0; size_t graph_degree_ = 0; }; +/** CAGRA index with the usual padded device dataset view (graph build output type). */ +template +using device_padded_index = index>; + +/** CAGRA index with a host-resident padded dataset view (returned by host build path). */ +template +using host_padded_index = index>; + +/** CAGRA index with a device-resident standard (arbitrary stride) dataset view. */ +template +using device_standard_index = + index>; + +/** CAGRA index with a host-resident standard dataset view. */ +template +using host_standard_index = index>; + +/** CAGRA index with a device-resident VPQ dataset (f16 codebook vectors). */ +template +using vpq_f16_index = index>; + +/** CAGRA index with a device-resident VPQ dataset (f32 codebook vectors). */ +template +using vpq_f32_index = index>; + +/** Index type returned by `cagra::build(res, params, dataset_view)`. */ +template +using cagra_index_t = index, + uint32_t, + cuvs::neighbors::dataset_view_type_t>; + /** * @} */ +/** + * @brief Row counts and strides for a CAGRA merge (metadata only; no GPU storage). + * + * A populated instance is carried inside `merged_dataset_storage` together with the owning + * device matrices allocated by `make_merged_storage`. + */ +struct merged_dataset { + int64_t merged_rows{}; ///< Full concatenation row count (staging for merge + filter). + int64_t filtered_rows{}; ///< Dataset rows the merged index will reference (filtered or full). + int64_t stride_elements{}; ///< Row pitch in elements (>= dim, matches input index rows). + uint32_t dim{}; + bool bitset_filtered{}; ///< If true, `merged_dataset_storage` holds a second matrix for rows + ///< after the bitset filter. +}; + +/** + * @brief Device storage for a physical CAGRA merge, allocated by `make_merged_storage`. + * + * Owns the full-merge staging matrix (`merged_storage`) and, when `layout.bitset_filtered` is + * true, the filtered output matrix (`filtered_storage`). `merge` writes into these buffers and + * returns an index that views them; keep this object alive while using that index. + */ +template +struct merged_dataset_storage { + merged_dataset layout{}; + raft::device_matrix merged_storage; + std::optional> filtered_storage{}; +}; + /** * @defgroup cagra_cpp_index_build CAGRA index build functions * @{ */ /** - * @brief Build the index from the dataset for efficient search. + * @brief Build the index from a `dataset_view` (device padded/standard or host padded/standard). * - * The build consist of two steps: build an intermediate knn-graph, and optimize it to - * create the final graph. The index_params struct controls the node degree of these - * graphs. + * VPQ-compressed device views are rejected: dense graph construction requires uncompressed data. + * Use a separate VPQ index workflow after building the graph from an uncompressed dataset. * - * The following distance metrics are supported: - * - L2 - * - InnerProduct (currently only supported with IVF-PQ as the build algorithm) - * - CosineExpanded - * - L1 (currently only supported with NN-Descent and Iterative Search as the build algorithm) + * When `index_params.attach_dataset_on_build = true` (the default) **and the input is a device + * view**, the `dataset` view is stored in the returned index as a **non-owning view** — no copy is + * made. The caller must keep the underlying storage alive for the lifetime of the index. The + * returned index is then ready to search immediately. * - * Usage example: - * @code{.cpp} - * using namespace cuvs::neighbors; - * // use default index parameters - * cagra::index_params index_params; - * // create and fill the index from a [N, D] dataset - * auto index = cagra::build(res, index_params, dataset); - * // use default search parameters - * cagra::search_params search_params; - * // search K nearest neighbours - * auto neighbors = raft::make_device_matrix(res, n_queries, k); - * auto distances = raft::make_device_matrix(res, n_queries, k); - * cagra::search(res, search_params, index, queries, neighbors.view(), distances.view()); - * @endcode + * When `index_params.attach_dataset_on_build = false`, or when building from a **host view**, only + * the search graph is built and the returned index holds no dataset. * - * @param[in] res - * @param[in] params parameters for building the index - * @param[in] dataset a matrix view (device) to a row-major matrix [n_rows, dim] + * For host views, the returned `host_padded_index` cannot be searched regardless of + * `attach_dataset_on_build` (the flag is ignored). Call `attach_device_dataset_on_host_index` to + * convert it to a device-backed index before search. * - * @return the constructed cagra index + * Note: disk-based ACE builds (`ace_params::use_disk = true`) always set a file-descriptor + * dataset internally (also host-typed); `attach_dataset_on_build` is ignored there too. + */ +// Concrete non-template overloads for all supported build dataset view types. +// This keeps the public header explicit and stable while implementation remains shared internally. +/** + * @brief Build from a device padded dataset view (`float`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset device padded dataset view [n_rows, dim] + * @return built `device_padded_index` */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - raft::device_matrix_view dataset) - -> cuvs::neighbors::cagra::index; + cuvs::neighbors::device_padded_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_padded_index; /** - * @brief Build the index from the dataset for efficient search. - * - * The build consist of two steps: build an intermediate knn-graph, and optimize it to - * create the final graph. The index_params struct controls the node degree of these - * graphs. - * - * The following distance metrics are supported: - * - L2 - * - InnerProduct (currently only supported with IVF-PQ as the build algorithm) - * - CosineExpanded - * - L1 (currently only supported with NN-Descent and Iterative Search as the build algorithm) - * - * Usage example: - * @code{.cpp} - * using namespace cuvs::neighbors; - * // use default index parameters - * cagra::index_params index_params; - * // create and fill the index from a [N, D] dataset - * auto index = cagra::build(res, index_params, dataset); - * // use default search parameters - * cagra::search_params search_params; - * // search K nearest neighbours - * auto neighbors = raft::make_device_matrix(res, n_queries, k); - * auto distances = raft::make_device_matrix(res, n_queries, k); - * cagra::search(res, search_params, index, queries, neighbors.view(), distances.view()); - * @endcode - * - * @param[in] res - * @param[in] params parameters for building the index - * @param[in] dataset a matrix view (host) to a row-major matrix [n_rows, dim] - * - * @return the constructed cagra index + * @brief Build from a device standard dataset view (`float`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset device standard dataset view [n_rows, dim] + * @return built `device_standard_index` */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - raft::host_matrix_view dataset) - -> cuvs::neighbors::cagra::index; + cuvs::neighbors::device_standard_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_standard_index; /** - * @brief Build the index from the dataset for efficient search. - * - * The build consist of two steps: build an intermediate knn-graph, and optimize it to - * create the final graph. The index_params struct controls the node degree of these - * graphs. - * - * The following distance metrics are supported: - * - L2 - * - InnerProduct (currently only supported with IVF-PQ as the build algorithm) - * - CosineExpanded (dataset norms are computed as float regardless of input data type) - * - L1 (currently only supported with NN-Descent and Iterative Search as the build algorithm) - * - * Usage example: - * @code{.cpp} - * using namespace cuvs::neighbors; - * // use default index parameters - * cagra::index_params index_params; - * // create and fill the index from a [N, D] dataset - * auto index = cagra::build(res, index_params, dataset); - * // use default search parameters - * cagra::search_params search_params; - * // search K nearest neighbours - * auto neighbors = raft::make_device_matrix(res, n_queries, k); - * auto distances = raft::make_device_matrix(res, n_queries, k); - * cagra::search(res, search_params, index, queries, neighbors.view(), distances.view()); - * @endcode - * - * @param[in] res - * @param[in] params parameters for building the index - * @param[in] dataset a matrix view (device) to a row-major matrix [n_rows, dim] - * - * @return the constructed cagra index + * @brief Build from a host padded dataset view (`float`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset host padded dataset view [n_rows, dim] + * @return built `host_padded_index` */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - raft::device_matrix_view dataset) - -> cuvs::neighbors::cagra::index; + cuvs::neighbors::host_padded_dataset_view const& dataset) + -> cuvs::neighbors::cagra::host_padded_index; /** - * @brief Build the index from the dataset for efficient search. - * - * The build consist of two steps: build an intermediate knn-graph, and optimize it to - * create the final graph. The index_params struct controls the node degree of these - * graphs. - * - * The following distance metrics are supported: - * - L2 - * - CosineExpanded (dataset norms are computed as float regardless of input data type) - * - L1 (currently only supported with NN-Descent and Iterative Search as the build algorithm) - * - * Usage example: - * @code{.cpp} - * using namespace cuvs::neighbors; - * // use default index parameters - * cagra::index_params index_params; - * // create and fill the index from a [N, D] dataset - * auto index = cagra::build(res, index_params, dataset); - * // use default search parameters - * cagra::search_params search_params; - * // search K nearest neighbours - * auto neighbors = raft::make_device_matrix(res, n_queries, k); - * auto distances = raft::make_device_matrix(res, n_queries, k); - * cagra::search(res, search_params, index, queries, neighbors.view(), distances.view()); - * @endcode - * - * @param[in] res - * @param[in] params parameters for building the index - * @param[in] dataset a matrix view (host) to a row-major matrix [n_rows, dim] - * - * @return the constructed cagra index + * @brief Build from a host standard dataset view (`float`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset host standard dataset view [n_rows, dim] + * @return built `host_standard_index` */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - raft::host_matrix_view dataset) - -> cuvs::neighbors::cagra::index; + cuvs::neighbors::host_standard_dataset_view const& dataset) + -> cuvs::neighbors::cagra::host_standard_index; /** - * @brief Build the index from the dataset for efficient search. - * - * The build consist of two steps: build an intermediate knn-graph, and optimize it to - * create the final graph. The index_params struct controls the node degree of these - * graphs. - * - * The following distance metrics are supported: - * - L2 - * - CosineExpanded (dataset norms are computed as float regardless of input data type) - * - L1 (currently only supported with NN-Descent and Iterative Search as the build algorithm) - * - BitwiseHamming (currently only supported with NN-Descent and Iterative Search as the build - * algorithm, and only for int8_t and uint8_t data types) - * - * Usage example: - * @code{.cpp} - * using namespace cuvs::neighbors; - * // use default index parameters - * cagra::index_params index_params; - * // create and fill the index from a [N, D] dataset - * auto index = cagra::build(res, index_params, dataset); - * // use default search parameters - * cagra::search_params search_params; - * // search K nearest neighbours - * auto neighbors = raft::make_device_matrix(res, n_queries, k); - * auto distances = raft::make_device_matrix(res, n_queries, k); - * cagra::search(res, search_params, index, queries, neighbors.view(), distances.view()); - * @endcode - * - * @param[in] res - * @param[in] params parameters for building the index - * @param[in] dataset a matrix view (device) to a row-major matrix [n_rows, dim] - * - * @return the constructed cagra index + * @brief Build from a device padded dataset view (`half`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset device padded dataset view [n_rows, dim] + * @return built `device_padded_index` */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - raft::device_matrix_view dataset) - -> cuvs::neighbors::cagra::index; + cuvs::neighbors::device_padded_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_padded_index; /** - * @brief Build the index from the dataset for efficient search. - * - * The build consist of two steps: build an intermediate knn-graph, and optimize it to - * create the final graph. The index_params struct controls the node degree of these - * graphs. - * - * The following distance metrics are supported: - * - L2 - * - InnerProduct (currently only supported with IVF-PQ as the build algorithm) - * - CosineExpanded (dataset norms are computed as float regardless of input data type) - * - L1 (currently only supported with NN-Descent and Iterative Search as the build algorithm) - * - BitwiseHamming (currently only supported with NN-Descent and Iterative Search as the build - * algorithm, and only for int8_t and uint8_t data types) - * - * Usage example: - * @code{.cpp} - * using namespace cuvs::neighbors; - * // use default index parameters - * cagra::index_params index_params; - * // create and fill the index from a [N, D] dataset - * auto index = cagra::build(res, index_params, dataset); - * // use default search parameters - * cagra::search_params search_params; - * // search K nearest neighbours - * auto neighbors = raft::make_device_matrix(res, n_queries, k); - * auto distances = raft::make_device_matrix(res, n_queries, k); - * cagra::search(res, search_params, index, queries, neighbors.view(), distances.view()); - * @endcode - * - * @param[in] res - * @param[in] params parameters for building the index - * @param[in] dataset a matrix view (host) to a row-major matrix [n_rows, dim] - * - * @return the constructed cagra index + * @brief Build from a device standard dataset view (`half`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset device standard dataset view [n_rows, dim] + * @return built `device_standard_index` */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - raft::host_matrix_view dataset) - -> cuvs::neighbors::cagra::index; + cuvs::neighbors::device_standard_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_standard_index; /** - * @brief Build the index from the dataset for efficient search. - * - * The build consist of two steps: build an intermediate knn-graph, and optimize it to - * create the final graph. The index_params struct controls the node degree of these - * graphs. - * - * The following distance metrics are supported: - * - L2 - * - InnerProduct (currently only supported with IVF-PQ as the build algorithm) - * - CosineExpanded (dataset norms are computed as float regardless of input data type) - * - L1 (currently only supported with NN-Descent and Iterative Search as the build algorithm) - * - BitwiseHamming (currently only supported with NN-Descent and Iterative Search as the build - * algorithm, and only for int8_t and uint8_t data types) - * - * Usage example: - * @code{.cpp} - * using namespace cuvs::neighbors; - * // use default index parameters - * cagra::index_params index_params; - * // create and fill the index from a [N, D] dataset - * auto index = cagra::build(res, index_params, dataset); - * // use default search parameters - * cagra::search_params search_params; - * // search K nearest neighbours - * auto neighbors = raft::make_device_matrix(res, n_queries, k); - * auto distances = raft::make_device_matrix(res, n_queries, k); - * cagra::search(res, search_params, index, queries, neighbors.view(), distances.view()); - * @endcode - * - * @param[in] res - * @param[in] params parameters for building the index - * @param[in] dataset a matrix view (device) to a row-major matrix [n_rows, dim] - * - * @return the constructed cagra index + * @brief Build from a host padded dataset view (`half`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset host padded dataset view [n_rows, dim] + * @return built `host_padded_index` */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - raft::device_matrix_view dataset) - -> cuvs::neighbors::cagra::index; + cuvs::neighbors::host_padded_dataset_view const& dataset) + -> cuvs::neighbors::cagra::host_padded_index; /** - * @brief Build the index from the dataset for efficient search. - * - * The build consist of two steps: build an intermediate knn-graph, and optimize it to - * create the final graph. The index_params struct controls the node degree of these - * graphs. - * - * The following distance metrics are supported: - * - L2 - * - InnerProduct (currently only supported with IVF-PQ as the build algorithm) - * - CosineExpanded (dataset norms are computed as float regardless of input data type) - * - L1 (currently only supported with NN-Descent and Iterative Search as the build algorithm) - * - BitwiseHamming (currently only supported with NN-Descent and Iterative Search as the build - * algorithm, and only for int8_t and uint8_t data types) - * - * Usage example: - * @code{.cpp} - * using namespace cuvs::neighbors; - * // use default index parameters - * cagra::index_params index_params; - * // create and fill the index from a [N, D] dataset - * auto index = cagra::build(res, index_params, dataset); - * // use default search parameters - * cagra::search_params search_params; - * // search K nearest neighbours - * auto neighbors = raft::make_device_matrix(res, n_queries, k); - * auto distances = raft::make_device_matrix(res, n_queries, k); - * cagra::search(res, search_params, index, queries, neighbors.view(), distances.view()); - * @endcode - * - * @param[in] res - * @param[in] params parameters for building the index - * @param[in] dataset a matrix view (host) to a row-major matrix [n_rows, dim] - * - * @return the constructed cagra index + * @brief Build from a host standard dataset view (`half`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset host standard dataset view [n_rows, dim] + * @return built `host_standard_index` + */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::host_standard_dataset_view const& dataset) + -> cuvs::neighbors::cagra::host_standard_index; + +/** + * @brief Build from a device padded dataset view (`int8_t`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset device padded dataset view [n_rows, dim] + * @return built `device_padded_index` + */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::device_padded_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_padded_index; + +/** + * @brief Build from a device standard dataset view (`int8_t`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset device standard dataset view [n_rows, dim] + * @return built `device_standard_index` + */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::device_standard_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_standard_index; + +/** + * @brief Build from a host padded dataset view (`int8_t`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset host padded dataset view [n_rows, dim] + * @return built `host_padded_index` */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - raft::host_matrix_view dataset) - -> cuvs::neighbors::cagra::index; + cuvs::neighbors::host_padded_dataset_view const& dataset) + -> cuvs::neighbors::cagra::host_padded_index; + +/** + * @brief Build from a host standard dataset view (`int8_t`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset host standard dataset view [n_rows, dim] + * @return built `host_standard_index` + */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::host_standard_dataset_view const& dataset) + -> cuvs::neighbors::cagra::host_standard_index; + +/** + * @brief Build from a device padded dataset view (`uint8_t`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset device padded dataset view [n_rows, dim] + * @return built `device_padded_index` + */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::device_padded_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_padded_index; + +/** + * @brief Build from a device standard dataset view (`uint8_t`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset device standard dataset view [n_rows, dim] + * @return built `device_standard_index` + */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::device_standard_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_standard_index; + +/** + * @brief Build from a host padded dataset view (`uint8_t`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset host padded dataset view [n_rows, dim] + * @return built `host_padded_index` + */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::host_padded_dataset_view const& dataset) + -> cuvs::neighbors::cagra::host_padded_index; + +/** + * @brief Build from a host standard dataset view (`uint8_t`). + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset host standard dataset view [n_rows, dim] + * @return built `host_standard_index` + */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::host_standard_dataset_view const& dataset) + -> cuvs::neighbors::cagra::host_standard_index; + /** * @} */ @@ -1229,6 +1116,52 @@ auto build(raft::resources const& res, * @{ */ +/** + * @brief Device buffers used by `cagra::extend`. + * + * Allocate with `make_extended_storage`, then pass to `extend`. Keep this storage alive for as long + * as the returned index views are needed. + */ +template +struct extended_dataset_storage { + raft::device_matrix dataset_storage; + raft::device_matrix graph_storage; +}; + +/** + * @brief Allocate output buffers for extend. + * + * Computes the target extended row count from `idx.size() + additional_dataset.n_rows()`, preserves + * the index dataset row stride, and allocates both dataset and graph output buffers. + */ +template +extended_dataset_storage make_extended_storage( + raft::resources const& handle, + AdditionalDatasetViewT additional_dataset, + cuvs::neighbors::cagra::index const& idx) +{ + auto cur_ds = idx.dataset().view(); + const auto stride_elems = cur_ds.stride(0) > 0 ? static_cast(cur_ds.stride(0)) + : static_cast(cur_ds.extent(1)); + const auto new_rows = + static_cast(idx.size()) + static_cast(additional_dataset.n_rows()); + auto dataset_storage = + raft::make_device_matrix(handle, new_rows, stride_elems); + auto graph_storage = raft::make_device_matrix( + handle, new_rows, static_cast(idx.graph_degree())); + return extended_dataset_storage{std::move(dataset_storage), std::move(graph_storage)}; +} + +// Concrete non-template overloads for all supported index types. +// Previously a single template covered all index types; it has been +// replaced with explicit overloads to maintain a stable non-template ABI. When a new index +// type is added (e.g. a future host_padded_index extend), add a corresponding overload here. +// Index types for which extend is not meaningful (e.g. VPQ — read-only compressed codes) +// are intentionally omitted. + /** @brief Add new vectors to a CAGRA index * * Usage example: @@ -1258,21 +1191,18 @@ auto build(raft::resources const& res, * as the index. This option is useful when users want to manage the memory space for the graph * themselves. */ -void extend( - raft::resources const& handle, - const cagra::extend_params& params, - raft::device_matrix_view additional_dataset, - cuvs::neighbors::cagra::index& idx, - std::optional> - new_dataset_buffer_view = std::nullopt, - std::optional> new_graph_buffer_view = std::nullopt); +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::device_padded_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_padded_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); /** @brief Add new vectors to a CAGRA index * * Usage example: * @code{.cpp} * using namespace cuvs::neighbors; - * auto additional_dataset = raft::make_host_matrix(handle,add_size,dim); + * auto additional_dataset = raft::make_device_matrix(handle,add_size,dim); * // set_additional_dataset(additional_dataset.view()); * * cagra::extend_params params; @@ -1281,7 +1211,7 @@ void extend( * * @param[in] handle raft resources * @param[in] params extend params - * @param[in] additional_dataset additional dataset on host memory + * @param[in] additional_dataset additional dataset on device memory * @param[in,out] idx CAGRA index * @param[out] new_dataset_buffer_view memory buffer view for the dataset including the additional * part. The data will be copied from the current index in this function. The num rows must be the @@ -1296,21 +1226,18 @@ void extend( * as the index. This option is useful when users want to manage the memory space for the graph * themselves. */ -void extend( - raft::resources const& handle, - const cagra::extend_params& params, - raft::host_matrix_view additional_dataset, - cuvs::neighbors::cagra::index& idx, - std::optional> - new_dataset_buffer_view = std::nullopt, - std::optional> new_graph_buffer_view = std::nullopt); +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::device_padded_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_padded_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); /** @brief Add new vectors to a CAGRA index * * Usage example: * @code{.cpp} * using namespace cuvs::neighbors; - * auto additional_dataset = raft::make_device_matrix(handle,add_size,dim); + * auto additional_dataset = raft::make_device_matrix(handle,add_size,dim); * // set_additional_dataset(additional_dataset.view()); * * cagra::extend_params params; @@ -1334,21 +1261,18 @@ void extend( * as the index. This option is useful when users want to manage the memory space for the graph * themselves. */ -void extend( - raft::resources const& handle, - const cagra::extend_params& params, - raft::device_matrix_view additional_dataset, - cuvs::neighbors::cagra::index& idx, - std::optional> - new_dataset_buffer_view = std::nullopt, - std::optional> new_graph_buffer_view = std::nullopt); +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::device_padded_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_padded_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); /** @brief Add new vectors to a CAGRA index * * Usage example: * @code{.cpp} * using namespace cuvs::neighbors; - * auto additional_dataset = raft::make_host_matrix(handle,add_size,dim); + * auto additional_dataset = raft::make_device_matrix(handle,add_size,dim); * // set_additional_dataset(additional_dataset.view()); * * cagra::extend_params params; @@ -1357,7 +1281,7 @@ void extend( * * @param[in] handle raft resources * @param[in] params extend params - * @param[in] additional_dataset additional dataset on host memory + * @param[in] additional_dataset additional dataset on device memory * @param[in,out] idx CAGRA index * @param[out] new_dataset_buffer_view memory buffer view for the dataset including the additional * part. The data will be copied from the current index in this function. The num rows must be the @@ -1372,21 +1296,18 @@ void extend( * as the index. This option is useful when users want to manage the memory space for the graph * themselves. */ -void extend( - raft::resources const& handle, - const cagra::extend_params& params, - raft::host_matrix_view additional_dataset, - cuvs::neighbors::cagra::index& idx, - std::optional> - new_dataset_buffer_view = std::nullopt, - std::optional> new_graph_buffer_view = std::nullopt); +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::device_padded_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_padded_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); /** @brief Add new vectors to a CAGRA index * * Usage example: * @code{.cpp} * using namespace cuvs::neighbors; - * auto additional_dataset = raft::make_device_matrix(handle,add_size,dim); + * auto additional_dataset = raft::make_host_matrix(handle,add_size,dim); * // set_additional_dataset(additional_dataset.view()); * * cagra::extend_params params; @@ -1395,7 +1316,7 @@ void extend( * * @param[in] handle raft resources * @param[in] params extend params - * @param[in] additional_dataset additional dataset on device memory + * @param[in] additional_dataset additional dataset on host memory * @param[in,out] idx CAGRA index * @param[out] new_dataset_buffer_view memory buffer view for the dataset including the additional * part. The data will be copied from the current index in this function. The num rows must be the @@ -1410,21 +1331,18 @@ void extend( * as the index. This option is useful when users want to manage the memory space for the graph * themselves. */ -void extend( - raft::resources const& handle, - const cagra::extend_params& params, - raft::device_matrix_view additional_dataset, - cuvs::neighbors::cagra::index& idx, - std::optional> - new_dataset_buffer_view = std::nullopt, - std::optional> new_graph_buffer_view = std::nullopt); +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::host_padded_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_padded_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); /** @brief Add new vectors to a CAGRA index * * Usage example: * @code{.cpp} * using namespace cuvs::neighbors; - * auto additional_dataset = raft::make_host_matrix(handle,add_size,dim); + * auto additional_dataset = raft::make_host_matrix(handle,add_size,dim); * // set_additional_dataset(additional_dataset.view()); * * cagra::extend_params params; @@ -1448,21 +1366,18 @@ void extend( * as the index. This option is useful when users want to manage the memory space for the graph * themselves. */ -void extend( - raft::resources const& handle, - const cagra::extend_params& params, - raft::host_matrix_view additional_dataset, - cuvs::neighbors::cagra::index& idx, - std::optional> - new_dataset_buffer_view = std::nullopt, - std::optional> new_graph_buffer_view = std::nullopt); +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::host_padded_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_padded_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); /** @brief Add new vectors to a CAGRA index * * Usage example: * @code{.cpp} * using namespace cuvs::neighbors; - * auto additional_dataset = raft::make_host_matrix(handle,add_size,dim); + * auto additional_dataset = raft::make_host_matrix(handle,add_size,dim); * // set_additional_dataset(additional_dataset.view()); * * cagra::extend_params params; @@ -1486,14 +1401,11 @@ void extend( * as the index. This option is useful when users want to manage the memory space for the graph * themselves. */ -void extend( - raft::resources const& handle, - const cagra::extend_params& params, - raft::device_matrix_view additional_dataset, - cuvs::neighbors::cagra::index& idx, - std::optional> - new_dataset_buffer_view = std::nullopt, - std::optional> new_graph_buffer_view = std::nullopt); +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::host_padded_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_padded_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); /** @brief Add new vectors to a CAGRA index * @@ -1524,40 +1436,582 @@ void extend( * as the index. This option is useful when users want to manage the memory space for the graph * themselves. */ -void extend( - raft::resources const& handle, - const cagra::extend_params& params, - raft::host_matrix_view additional_dataset, - cuvs::neighbors::cagra::index& idx, - std::optional> - new_dataset_buffer_view = std::nullopt, - std::optional> new_graph_buffer_view = std::nullopt); +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::host_padded_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_padded_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); + +/** @copydoc extend */ +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::device_standard_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_standard_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); + +/** @copydoc extend */ +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::device_standard_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_standard_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); + +/** @copydoc extend */ +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::device_standard_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_standard_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); + +/** @copydoc extend */ +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::device_standard_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_standard_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); + +/** @copydoc extend */ +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::host_standard_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_standard_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); + +/** @copydoc extend */ +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::host_standard_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_standard_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); + +/** @copydoc extend */ +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::host_standard_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_standard_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); + +/** @copydoc extend */ +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::host_standard_dataset_view additional_dataset, + cuvs::neighbors::cagra::device_standard_index& idx, + cuvs::neighbors::cagra::extended_dataset_storage& storage); + +/** + * @} + */ + +/** + * @defgroup cagra_cpp_index_search CAGRA search functions + * @{ + */ + +// Concrete non-template overloads for all supported index types. +// Explicit overloads (not a single template) keep a stable link ABI for the C API. +// When a new index type is added, add matching overloads here and in cagra_search_inst.cu.in. + +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built device_padded_index with uint32_t neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). + */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_padded_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built device_padded_index with uint32_t neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). + */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_padded_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built device_padded_index with uint32_t neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). + */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_padded_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built device_padded_index with uint32_t neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). + */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_padded_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built device_padded_index with int64_t neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). + */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_padded_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built device_padded_index with int64_t neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). + */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_padded_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built device_padded_index with int64_t neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). + */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_padded_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built device_padded_index with int64_t neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). + */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_padded_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +// device_standard_index overloads (uint32_t neighbor indices) +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_standard_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_standard_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_standard_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_standard_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +// device_standard_index overloads (int64_t neighbor indices) +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_standard_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_standard_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_standard_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::device_standard_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +// vpq_f16_index overloads (uint32_t neighbor indices) +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with uint32_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). + */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::vpq_f16_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with uint32_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). + */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::vpq_f16_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with uint32_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). + */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::vpq_f16_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with uint32_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). + */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::vpq_f16_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +// vpq_f16_index overloads (int64_t neighbor indices) +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with int64_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). + */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::vpq_f16_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with int64_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). + */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::vpq_f16_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + /** - * @} + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with int64_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::vpq_f16_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); /** - * @defgroup cagra_cpp_index_search CAGRA search functions - * @{ * @brief Search ANN using the constructed index. * * See the [cagra::build](#cagra::build) documentation for a usage example. * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index cagra index - * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()] + * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with int64_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, * k] * @param[in] sample_filter an optional device filter function object that greenlights samples - * for a given query. (none_sample_filter for no filtering) + * for a given query. (none_sample_filter for no filtering). */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::vpq_f16_index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); +// vpq_f32_index overloads (uint32_t neighbor indices) +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with uint32_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter an optional device filter function object that greenlights samples + * for a given query. (none_sample_filter for no filtering). + * + * @note FP32 VPQ search is declared for ABI stability but fails at runtime until implemented. + */ void search(raft::resources const& res, cuvs::neighbors::cagra::search_params const& params, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::vpq_f32_index& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, raft::device_matrix_view distances, @@ -1571,18 +2025,21 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index cagra index - * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()] + * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with uint32_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, * k] * @param[in] sample_filter an optional device filter function object that greenlights samples - * for a given query. (none_sample_filter for no filtering) + * for a given query. (none_sample_filter for no filtering). + * + * @note FP32 VPQ search is declared for ABI stability but fails at runtime until implemented. */ void search(raft::resources const& res, cuvs::neighbors::cagra::search_params const& params, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::vpq_f32_index& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, raft::device_matrix_view distances, @@ -1596,18 +2053,21 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index cagra index - * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()] + * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with uint32_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, * k] * @param[in] sample_filter an optional device filter function object that greenlights samples - * for a given query. (none_sample_filter for no filtering) + * for a given query. (none_sample_filter for no filtering). + * + * @note FP32 VPQ search is declared for ABI stability but fails at runtime until implemented. */ void search(raft::resources const& res, cuvs::neighbors::cagra::search_params const& params, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::vpq_f32_index& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, raft::device_matrix_view distances, @@ -1621,24 +2081,28 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index cagra index - * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()] + * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with uint32_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, * k] * @param[in] sample_filter an optional device filter function object that greenlights samples - * for a given query. (none_sample_filter for no filtering) + * for a given query. (none_sample_filter for no filtering). + * + * @note FP32 VPQ search is declared for ABI stability but fails at runtime until implemented. */ void search(raft::resources const& res, cuvs::neighbors::cagra::search_params const& params, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::vpq_f32_index& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, raft::device_matrix_view distances, const cuvs::neighbors::filtering::base_filter& sample_filter = cuvs::neighbors::filtering::none_sample_filter{}); +// vpq_f32_index overloads (int64_t neighbor indices) /** * @brief Search ANN using the constructed index. * @@ -1646,19 +2110,21 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index cagra index - * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()] + * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with int64_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, * k] * @param[in] sample_filter an optional device filter function object that greenlights samples - * for a given query. (none_sample_filter for no filtering) + * for a given query. (none_sample_filter for no filtering). + * + * @note FP32 VPQ search is declared for ABI stability but fails at runtime until implemented. */ - void search(raft::resources const& res, cuvs::neighbors::cagra::search_params const& params, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::vpq_f32_index& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, raft::device_matrix_view distances, @@ -1672,18 +2138,21 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index cagra index - * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()] + * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with int64_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, * k] * @param[in] sample_filter an optional device filter function object that greenlights samples - * for a given query. (none_sample_filter for no filtering) + * for a given query. (none_sample_filter for no filtering). + * + * @note FP32 VPQ search is declared for ABI stability but fails at runtime until implemented. */ void search(raft::resources const& res, cuvs::neighbors::cagra::search_params const& params, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::vpq_f32_index& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, raft::device_matrix_view distances, @@ -1697,18 +2166,21 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index cagra index - * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()] + * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with int64_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, * k] * @param[in] sample_filter an optional device filter function object that greenlights samples - * for a given query. (none_sample_filter for no filtering) + * for a given query. (none_sample_filter for no filtering). + * + * @note FP32 VPQ search is declared for ABI stability but fails at runtime until implemented. */ void search(raft::resources const& res, cuvs::neighbors::cagra::search_params const& params, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::vpq_f32_index& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, raft::device_matrix_view distances, @@ -1722,18 +2194,21 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index cagra index - * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()] + * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with int64_t + * neighbor indices + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, * k] * @param[in] sample_filter an optional device filter function object that greenlights samples - * for a given query. (none_sample_filter for no filtering) + * for a given query. (none_sample_filter for no filtering). + * + * @note FP32 VPQ search is declared for ABI stability but fails at runtime until implemented. */ void search(raft::resources const& res, cuvs::neighbors::cagra::search_params const& params, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::vpq_f32_index& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, raft::device_matrix_view distances, @@ -1749,6 +2224,12 @@ void search(raft::resources const& res, * @{ */ +// Serialize and deserialize are overloaded for device_padded_index and device_standard_index. +// Both use the same strided dataset wire format; deserialize selects the owning dataset type +// from the index's DatasetViewT. To support a new dataset kind (e.g. vpq_f16_index), add a +// matching pair of overloads here and a corresponding deserialize_ in +// detail/dataset_serialize.hpp (dense views use serialize_cagra_padded_dataset). + /** * Save the index to file. * @@ -1774,7 +2255,7 @@ void search(raft::resources const& res, */ void serialize(raft::resources const& handle, const std::string& filename, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, bool include_dataset = true); /** @@ -1791,17 +2272,22 @@ void serialize(raft::resources const& handle, * // create a string with a filepath * std::string filename("/path/to/index"); - * cuvs::neighbors::cagra::index index; + * cuvs::neighbors::cagra::device_padded_index index; * cuvs::neighbors::cagra::deserialize(handle, filename, &index); * @endcode * * @param[in] handle the raft handle * @param[in] filename the name of the file that stores the index * @param[out] index the cagra index + * @param[out] out_dataset if non-null, on success may be set to an owned deserialized dataset + * when the file includes dataset data; may be left unchanged otherwise. Optional; pass + * nullptr to ignore. */ -void deserialize(raft::resources const& handle, - const std::string& filename, - cuvs::neighbors::cagra::index* index); +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_padded_index* index, + std::unique_ptr>* out_dataset = nullptr); /** * Write the index to an output stream @@ -1827,7 +2313,7 @@ void deserialize(raft::resources const& handle, */ void serialize(raft::resources const& handle, std::ostream& os, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, bool include_dataset = true); /** @@ -1843,17 +2329,22 @@ void serialize(raft::resources const& handle, * * // create an input stream * std::istream is(std::cin.rdbuf()); - * cuvs::neighbors::cagra::index index; + * cuvs::neighbors::cagra::device_padded_index index; * cuvs::neighbors::cagra::deserialize(handle, is, &index); * @endcode * * @param[in] handle the raft handle * @param[in] is input stream * @param[out] index the cagra index + * @param[out] out_dataset if non-null, on success may be set to an owned deserialized dataset + * when the stream includes dataset data; may be left unchanged otherwise. Optional; pass + * nullptr to ignore. */ -void deserialize(raft::resources const& handle, - std::istream& is, - cuvs::neighbors::cagra::index* index); +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_padded_index* index, + std::unique_ptr>* out_dataset = nullptr); /** * Save the index to file. * @@ -1879,7 +2370,7 @@ void deserialize(raft::resources const& handle, */ void serialize(raft::resources const& handle, const std::string& filename, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, bool include_dataset = true); /** @@ -1896,17 +2387,22 @@ void serialize(raft::resources const& handle, * // create a string with a filepath * std::string filename("/path/to/index"); - * cuvs::neighbors::cagra::index index; + * cuvs::neighbors::cagra::device_padded_index index; * cuvs::neighbors::cagra::deserialize(handle, filename, &index); * @endcode * * @param[in] handle the raft handle * @param[in] filename the name of the file that stores the index * @param[out] index the cagra index + * @param[out] out_dataset if non-null, on success may be set to an owned deserialized dataset + * when the file includes dataset data; may be left unchanged otherwise. Optional; pass + * nullptr to ignore. */ -void deserialize(raft::resources const& handle, - const std::string& filename, - cuvs::neighbors::cagra::index* index); +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_padded_index* index, + std::unique_ptr>* out_dataset = nullptr); /** * Write the index to an output stream @@ -1932,7 +2428,7 @@ void deserialize(raft::resources const& handle, */ void serialize(raft::resources const& handle, std::ostream& os, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, bool include_dataset = true); /** @@ -1948,17 +2444,22 @@ void serialize(raft::resources const& handle, * * // create an input stream * std::istream is(std::cin.rdbuf()); - * cuvs::neighbors::cagra::index index; + * cuvs::neighbors::cagra::device_padded_index index; * cuvs::neighbors::cagra::deserialize(handle, is, &index); * @endcode * * @param[in] handle the raft handle * @param[in] is input stream * @param[out] index the cagra index + * @param[out] out_dataset if non-null, on success may be set to an owned deserialized dataset + * when the stream includes dataset data; may be left unchanged otherwise. Optional; pass + * nullptr to ignore. */ -void deserialize(raft::resources const& handle, - std::istream& is, - cuvs::neighbors::cagra::index* index); +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_padded_index* index, + std::unique_ptr>* out_dataset = nullptr); /** * Save the index to file. @@ -1984,7 +2485,7 @@ void deserialize(raft::resources const& handle, */ void serialize(raft::resources const& handle, const std::string& filename, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, bool include_dataset = true); /** @@ -2001,17 +2502,22 @@ void serialize(raft::resources const& handle, * // create a string with a filepath * std::string filename("/path/to/index"); - * cuvs::neighbors::cagra::index index; + * cuvs::neighbors::cagra::device_padded_index index; * cuvs::neighbors::cagra::deserialize(handle, filename, &index); * @endcode * * @param[in] handle the raft handle * @param[in] filename the name of the file that stores the index * @param[out] index the cagra index + * @param[out] out_dataset if non-null, on success may be set to an owned deserialized dataset + * when the file includes dataset data; may be left unchanged otherwise. Optional; pass + * nullptr to ignore. */ -void deserialize(raft::resources const& handle, - const std::string& filename, - cuvs::neighbors::cagra::index* index); +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_padded_index* index, + std::unique_ptr>* out_dataset = nullptr); /** * Write the index to an output stream @@ -2037,7 +2543,7 @@ void deserialize(raft::resources const& handle, */ void serialize(raft::resources const& handle, std::ostream& os, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, bool include_dataset = true); /** @@ -2053,17 +2559,22 @@ void serialize(raft::resources const& handle, * * // create an input stream * std::istream is(std::cin.rdbuf()); - * cuvs::neighbors::cagra::index index; + * cuvs::neighbors::cagra::device_padded_index index; * cuvs::neighbors::cagra::deserialize(handle, is, &index); * @endcode * * @param[in] handle the raft handle * @param[in] is input stream * @param[out] index the cagra index + * @param[out] out_dataset if non-null, on success may be set to an owned deserialized dataset + * when the stream includes dataset data; may be left unchanged otherwise. Optional; pass + * nullptr to ignore. */ -void deserialize(raft::resources const& handle, - std::istream& is, - cuvs::neighbors::cagra::index* index); +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_padded_index* index, + std::unique_ptr>* out_dataset = nullptr); /** * Save the index to file. @@ -2089,7 +2600,7 @@ void deserialize(raft::resources const& handle, */ void serialize(raft::resources const& handle, const std::string& filename, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, bool include_dataset = true); /** @@ -2106,17 +2617,22 @@ void serialize(raft::resources const& handle, * // create a string with a filepath * std::string filename("/path/to/index"); - * cuvs::neighbors::cagra::index index; + * cuvs::neighbors::cagra::device_padded_index index; * cuvs::neighbors::cagra::deserialize(handle, filename, &index); * @endcode * * @param[in] handle the raft handle * @param[in] filename the name of the file that stores the index * @param[out] index the cagra index + * @param[out] out_dataset if non-null, on success may be set to an owned deserialized dataset + * when the file includes dataset data; may be left unchanged otherwise. Optional; pass + * nullptr to ignore. */ -void deserialize(raft::resources const& handle, - const std::string& filename, - cuvs::neighbors::cagra::index* index); +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_padded_index* index, + std::unique_ptr>* out_dataset = nullptr); /** * Write the index to an output stream @@ -2142,7 +2658,7 @@ void deserialize(raft::resources const& handle, */ void serialize(raft::resources const& handle, std::ostream& os, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, bool include_dataset = true); /** @@ -2158,17 +2674,110 @@ void serialize(raft::resources const& handle, * * // create an input stream * std::istream is(std::cin.rdbuf()); - * cuvs::neighbors::cagra::index index; + * cuvs::neighbors::cagra::device_padded_index index; * cuvs::neighbors::cagra::deserialize(handle, is, &index); * @endcode * * @param[in] handle the raft handle * @param[in] is input stream * @param[out] index the cagra index + * @param[out] out_dataset if non-null, on success may be set to an owned deserialized dataset + * when the stream includes dataset data; may be left unchanged otherwise. Optional; pass + * nullptr to ignore. */ +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_padded_index* index, + std::unique_ptr>* out_dataset = nullptr); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_standard_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_standard_index* index, + std::unique_ptr>* out_dataset = nullptr); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_standard_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_standard_index* index, + std::unique_ptr>* out_dataset = nullptr); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_standard_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_standard_index* index, + std::unique_ptr>* out_dataset = nullptr); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_standard_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_standard_index* index, + std::unique_ptr>* out_dataset = nullptr); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_standard_index& index, + bool include_dataset = true); + +void deserialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_standard_index* index, + std::unique_ptr>* + out_dataset = nullptr); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_standard_index& index, + bool include_dataset = true); + +void deserialize(raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_standard_index* index, + std::unique_ptr>* + out_dataset = nullptr); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_standard_index& index, + bool include_dataset = true); + +void deserialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_standard_index* index, + std::unique_ptr>* + out_dataset = nullptr); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_standard_index& index, + bool include_dataset = true); + void deserialize(raft::resources const& handle, std::istream& is, - cuvs::neighbors::cagra::index* index); + cuvs::neighbors::cagra::device_standard_index* index, + std::unique_ptr>* + out_dataset = nullptr); /** * Write the CAGRA built index as a base layer HNSW index to an output stream @@ -2199,7 +2808,7 @@ void deserialize(raft::resources const& handle, void serialize_to_hnswlib( raft::resources const& handle, std::ostream& os, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, std::optional> dataset = std::nullopt); @@ -2233,7 +2842,7 @@ void serialize_to_hnswlib( void serialize_to_hnswlib( raft::resources const& handle, const std::string& filename, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, std::optional> dataset = std::nullopt); @@ -2266,7 +2875,7 @@ void serialize_to_hnswlib( void serialize_to_hnswlib( raft::resources const& handle, std::ostream& os, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, std::optional> dataset = std::nullopt); @@ -2300,7 +2909,7 @@ void serialize_to_hnswlib( void serialize_to_hnswlib( raft::resources const& handle, const std::string& filename, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, std::optional> dataset = std::nullopt); @@ -2333,7 +2942,7 @@ void serialize_to_hnswlib( void serialize_to_hnswlib( raft::resources const& handle, std::ostream& os, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, std::optional> dataset = std::nullopt); @@ -2367,7 +2976,7 @@ void serialize_to_hnswlib( void serialize_to_hnswlib( raft::resources const& handle, const std::string& filename, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, std::optional> dataset = std::nullopt); @@ -2400,7 +3009,7 @@ void serialize_to_hnswlib( void serialize_to_hnswlib( raft::resources const& handle, std::ostream& os, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, std::optional> dataset = std::nullopt); @@ -2434,7 +3043,183 @@ void serialize_to_hnswlib( void serialize_to_hnswlib( raft::resources const& handle, const std::string& filename, - const cuvs::neighbors::cagra::index& index, + const cuvs::neighbors::cagra::device_padded_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_standard_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_standard_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_standard_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_standard_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_standard_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_standard_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_standard_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_standard_index& index, + std::optional> dataset = + std::nullopt); + +/** + * Write the CAGRA built index as a base layer HNSW index to an output stream. + * Requires `dataset` — host builds do not store vectors in the index. + */ +void serialize_to_hnswlib( + raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::host_padded_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::host_padded_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::host_padded_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::host_padded_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::host_padded_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::host_padded_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::host_padded_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::host_padded_index& index, + std::optional> dataset = + std::nullopt); + +/** + * Write the CAGRA built index as a base layer HNSW index to an output stream. + * Requires `dataset` — host builds do not store vectors in the index. + */ +void serialize_to_hnswlib( + raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::host_standard_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::host_standard_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::host_standard_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::host_standard_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::host_standard_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::host_standard_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::host_standard_index& index, + std::optional> dataset = + std::nullopt); + +void serialize_to_hnswlib( + raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::host_standard_index& index, std::optional> dataset = std::nullopt); @@ -2447,69 +3232,34 @@ void serialize_to_hnswlib( * @{ */ -/** @brief Merge multiple CAGRA indices into a single index. - * - * This function merges multiple CAGRA indices into one, combining both the datasets and graph - * structures. - * - * @note: When device memory is sufficient, the dataset attached to the returned index is allocated - * in device memory by default; otherwise, host memory is used automatically. - * - * @note: This API only supports physical merge (`merge_strategy = MERGE_STRATEGY_PHYSICAL`), and - * attempting a logical merge here will throw an error. - * - * Usage example: - * @code{.cpp} - * using namespace cuvs::neighbors; - * auto dataset0 = raft::make_host_matrix(handle, size0, dim); - * auto dataset1 = raft::make_host_matrix(handle, size1, dim); - * - * auto index0 = cagra::build(res, index_params, dataset0); - * auto index1 = cagra::build(res, index_params, dataset1); - * - * std::vector*> indices{&index0, &index1}; - * - * auto merged_index = cagra::merge(res, index_params, indices); - * @endcode +/** @brief Allocate device merge buffers for the given indices and row filter. * - * @param[in] res RAFT resources used for the merge operation. - * @param[in] params Parameters that control the merging process. - * @param[in] indices A vector of pointers to the CAGRA indices to merge. All indices must: - * - Have attached datasets with the same dimension. - * @param[in] row_filter an optional device filter function object that greenlights rows - * to include in the merged index (none_sample_filter for no filtering) - * @return A new CAGRA index containing the merged indices, graph, and dataset. + * Computes row counts and stride (see `merged_dataset`), allocates `merged_storage` with shape + * `[merged_rows, stride_elements]`, and when using a bitset row filter also allocates + * `filtered_storage` with shape `[filtered_rows, stride_elements]`. Pass the result to `merge` with + * the same `indices` and `row_filter`. */ -auto merge(raft::resources const& res, - const cuvs::neighbors::cagra::index_params& params, - std::vector*>& indices, - const cuvs::neighbors::filtering::base_filter& row_filter = - cuvs::neighbors::filtering::none_sample_filter{}) - -> cuvs::neighbors::cagra::index; - -/** @copydoc merge */ -auto merge(raft::resources const& res, - const cuvs::neighbors::cagra::index_params& params, - std::vector*>& indices, - const cuvs::neighbors::filtering::base_filter& row_filter = - cuvs::neighbors::filtering::none_sample_filter{}) - -> cuvs::neighbors::cagra::index; +template +merged_dataset_storage make_merged_storage( + raft::resources const& res, + std::vector*> const& indices, + const cuvs::neighbors::filtering::base_filter& row_filter = + cuvs::neighbors::filtering::none_sample_filter{}); -/** @copydoc merge */ +/** @brief Merge multiple CAGRA indices into a single index. + * + * @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). + */ +template auto merge(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - std::vector*>& indices, + std::vector*>& indices, + merged_dataset_storage& storage, const cuvs::neighbors::filtering::base_filter& row_filter = cuvs::neighbors::filtering::none_sample_filter{}) - -> cuvs::neighbors::cagra::index; + -> cuvs::neighbors::cagra::index; -/** @copydoc merge */ -auto merge(raft::resources const& res, - const cuvs::neighbors::cagra::index_params& params, - std::vector*>& indices, - const cuvs::neighbors::filtering::base_filter& row_filter = - cuvs::neighbors::filtering::none_sample_filter{}) - -> cuvs::neighbors::cagra::index; /** * @} */ @@ -2535,8 +3285,8 @@ auto merge(raft::resources const& res, */ auto build(const raft::resources& clique, const cuvs::neighbors::mg_index_params& index_params, - raft::host_matrix_view index_dataset) - -> cuvs::neighbors::mg_index, float, uint32_t>; + cuvs::neighbors::host_standard_dataset_view const& index_dataset) + -> cuvs::neighbors::mg_index, float, uint32_t>; /// \ingroup mg_cpp_index_build /** @@ -2557,8 +3307,8 @@ auto build(const raft::resources& clique, */ auto build(const raft::resources& clique, const cuvs::neighbors::mg_index_params& index_params, - raft::host_matrix_view index_dataset) - -> cuvs::neighbors::mg_index, half, uint32_t>; + cuvs::neighbors::host_standard_dataset_view const& index_dataset) + -> cuvs::neighbors::mg_index, half, uint32_t>; /// \ingroup mg_cpp_index_build /** @@ -2579,8 +3329,8 @@ auto build(const raft::resources& clique, */ auto build(const raft::resources& clique, const cuvs::neighbors::mg_index_params& index_params, - raft::host_matrix_view index_dataset) - -> cuvs::neighbors::mg_index, int8_t, uint32_t>; + cuvs::neighbors::host_standard_dataset_view const& index_dataset) + -> cuvs::neighbors::mg_index, int8_t, uint32_t>; /// \ingroup mg_cpp_index_build /** @@ -2601,8 +3351,62 @@ auto build(const raft::resources& clique, */ auto build(const raft::resources& clique, const cuvs::neighbors::mg_index_params& index_params, - raft::host_matrix_view index_dataset) - -> cuvs::neighbors::mg_index, uint8_t, uint32_t>; + cuvs::neighbors::host_standard_dataset_view const& index_dataset) + -> cuvs::neighbors::mg_index, uint8_t, uint32_t>; + +auto build(const raft::resources& clique, + const cuvs::neighbors::mg_index_params& index_params, + cuvs::neighbors::host_padded_dataset_view const& index_dataset) + -> cuvs::neighbors::mg_index, float, uint32_t>; + +auto build(const raft::resources& clique, + const cuvs::neighbors::mg_index_params& index_params, + cuvs::neighbors::host_padded_dataset_view const& index_dataset) + -> cuvs::neighbors::mg_index, half, uint32_t>; + +auto build(const raft::resources& clique, + const cuvs::neighbors::mg_index_params& index_params, + cuvs::neighbors::host_padded_dataset_view const& index_dataset) + -> cuvs::neighbors::mg_index, int8_t, uint32_t>; + +auto build(const raft::resources& clique, + const cuvs::neighbors::mg_index_params& index_params, + cuvs::neighbors::host_padded_dataset_view const& index_dataset) + -> cuvs::neighbors::mg_index, uint8_t, uint32_t>; + +/** + * @brief Convert a standard MG CAGRA index into a padded MG CAGRA index for search. + * + * This API enforces the explicit contract: users build/deserialize/distribute a standard index, + * then attach a padded dataset before calling `search`. + */ +auto attach_padded_dataset_for_search( + const raft::resources& clique, + const cuvs::neighbors::mg_index, float, uint32_t>& + idx, + cuvs::neighbors::device_padded_dataset_view const& padded_dataset) + -> cuvs::neighbors::mg_index, float, uint32_t>; + +auto attach_padded_dataset_for_search( + const raft::resources& clique, + const cuvs::neighbors::mg_index, half, uint32_t>& + idx, + cuvs::neighbors::device_padded_dataset_view const& padded_dataset) + -> cuvs::neighbors::mg_index, half, uint32_t>; + +auto attach_padded_dataset_for_search( + const raft::resources& clique, + const cuvs::neighbors::mg_index, int8_t, uint32_t>& + idx, + cuvs::neighbors::device_padded_dataset_view const& padded_dataset) + -> cuvs::neighbors::mg_index, int8_t, uint32_t>; + +auto attach_padded_dataset_for_search( + const raft::resources& clique, + const cuvs::neighbors:: + mg_index, uint8_t, uint32_t>& idx, + cuvs::neighbors::device_padded_dataset_view const& padded_dataset) + -> cuvs::neighbors::mg_index, uint8_t, uint32_t>; /// \defgroup mg_cpp_index_extend ANN MG index extend @@ -2620,15 +3424,41 @@ auto build(const raft::resources& clique, * * @param[in] clique a `raft::resources` object specifying the NCCL clique configuration * @param[in] index the pre-built index - * @param[in] new_vectors a row-major matrix on host [n_rows, dim] + * @param[in] new_vectors host dataset view [n_rows, dim] + * @param[in] new_indices optional vector on host [n_rows], + * `std::nullopt` means default continuous range `[0...n_rows)` + * + */ +void extend( + const raft::resources& clique, + cuvs::neighbors::mg_index, float, uint32_t>& index, + cuvs::neighbors::host_padded_dataset_view new_vectors, + std::optional> new_indices); + +/// \ingroup mg_cpp_index_extend +/** + * @brief Extends a multi-GPU index + * + * Usage example: + * @code{.cpp} + * raft::device_resources_snmg clique; + * cuvs::neighbors::mg_index_params index_params; + * auto index = cuvs::neighbors::cagra::build(clique, index_params, index_dataset); + * cuvs::neighbors::cagra::extend(clique, index, new_vectors, std::nullopt); + * @endcode + * + * @param[in] clique a `raft::resources` object specifying the NCCL clique configuration + * @param[in] index the pre-built index + * @param[in] new_vectors host dataset view [n_rows, dim] * @param[in] new_indices optional vector on host [n_rows], * `std::nullopt` means default continuous range `[0...n_rows)` * */ -void extend(const raft::resources& clique, - cuvs::neighbors::mg_index, float, uint32_t>& index, - raft::host_matrix_view new_vectors, - std::optional> new_indices); +void extend( + const raft::resources& clique, + cuvs::neighbors::mg_index, float, uint32_t>& index, + cuvs::neighbors::host_standard_dataset_view new_vectors, + std::optional> new_indices); /// \ingroup mg_cpp_index_extend /** @@ -2644,15 +3474,23 @@ void extend(const raft::resources& clique, * * @param[in] clique a `raft::resources` object specifying the NCCL clique configuration * @param[in] index the pre-built index - * @param[in] new_vectors a row-major matrix on host [n_rows, dim] + * @param[in] new_vectors host dataset view [n_rows, dim] * @param[in] new_indices optional vector on host [n_rows], * `std::nullopt` means default continuous range `[0...n_rows)` * */ -void extend(const raft::resources& clique, - cuvs::neighbors::mg_index, half, uint32_t>& index, - raft::host_matrix_view new_vectors, - std::optional> new_indices); +void extend( + const raft::resources& clique, + cuvs::neighbors::mg_index, half, uint32_t>& index, + cuvs::neighbors::host_padded_dataset_view new_vectors, + std::optional> new_indices); + +/** @copydoc extend */ +void extend( + const raft::resources& clique, + cuvs::neighbors::mg_index, half, uint32_t>& index, + cuvs::neighbors::host_standard_dataset_view new_vectors, + std::optional> new_indices); /// \ingroup mg_cpp_index_extend /** @@ -2668,15 +3506,24 @@ void extend(const raft::resources& clique, * * @param[in] clique a `raft::resources` object specifying the NCCL clique configuration * @param[in] index the pre-built index - * @param[in] new_vectors a row-major matrix on host [n_rows, dim] + * @param[in] new_vectors host dataset view [n_rows, dim] * @param[in] new_indices optional vector on host [n_rows], * `std::nullopt` means default continuous range `[0...n_rows)` * */ -void extend(const raft::resources& clique, - cuvs::neighbors::mg_index, int8_t, uint32_t>& index, - raft::host_matrix_view new_vectors, - std::optional> new_indices); +void extend( + const raft::resources& clique, + cuvs::neighbors::mg_index, int8_t, uint32_t>& index, + cuvs::neighbors::host_padded_dataset_view new_vectors, + std::optional> new_indices); + +/** @copydoc extend */ +void extend( + const raft::resources& clique, + cuvs::neighbors::mg_index, int8_t, uint32_t>& + index, + cuvs::neighbors::host_standard_dataset_view new_vectors, + std::optional> new_indices); /// \ingroup mg_cpp_index_extend /** @@ -2692,15 +3539,25 @@ void extend(const raft::resources& clique, * * @param[in] clique a `raft::resources` object specifying the NCCL clique configuration * @param[in] index the pre-built index - * @param[in] new_vectors a row-major matrix on host [n_rows, dim] + * @param[in] new_vectors host dataset view [n_rows, dim] * @param[in] new_indices optional vector on host [n_rows], * `std::nullopt` means default continuous range `[0...n_rows)` * */ -void extend(const raft::resources& clique, - cuvs::neighbors::mg_index, uint8_t, uint32_t>& index, - raft::host_matrix_view new_vectors, - std::optional> new_indices); +void extend( + const raft::resources& clique, + cuvs::neighbors::mg_index, uint8_t, uint32_t>& + index, + cuvs::neighbors::host_padded_dataset_view new_vectors, + std::optional> new_indices); + +/** @copydoc extend */ +void extend( + const raft::resources& clique, + cuvs::neighbors::mg_index, uint8_t, uint32_t>& + index, + cuvs::neighbors::host_standard_dataset_view new_vectors, + std::optional> new_indices); /// \defgroup mg_cpp_index_search ANN MG index search @@ -2726,12 +3583,24 @@ void extend(const raft::resources& clique, * @param[out] distances a row-major matrix on host [n_rows, n_neighbors] * */ -void search(const raft::resources& clique, - const cuvs::neighbors::mg_index, float, uint32_t>& index, - const cuvs::neighbors::mg_search_params& search_params, - raft::host_matrix_view queries, - raft::host_matrix_view neighbors, - raft::host_matrix_view distances); +void search( + const raft::resources& clique, + const cuvs::neighbors::mg_index, float, uint32_t>& + index, + const cuvs::neighbors::mg_search_params& search_params, + raft::host_matrix_view queries, + raft::host_matrix_view neighbors, + raft::host_matrix_view distances); + +/** @copydoc search */ +void search( + const raft::resources& clique, + const cuvs::neighbors::mg_index, float, uint32_t>& + index, + const cuvs::neighbors::mg_search_params& search_params, + raft::host_matrix_view queries, + raft::host_matrix_view neighbors, + raft::host_matrix_view distances); /// \ingroup mg_cpp_index_search /** @@ -2755,12 +3624,24 @@ void search(const raft::resources& clique, * @param[out] distances a row-major matrix on host [n_rows, n_neighbors] * */ -void search(const raft::resources& clique, - const cuvs::neighbors::mg_index, half, uint32_t>& index, - const cuvs::neighbors::mg_search_params& search_params, - raft::host_matrix_view queries, - raft::host_matrix_view neighbors, - raft::host_matrix_view distances); +void search( + const raft::resources& clique, + const cuvs::neighbors::mg_index, half, uint32_t>& + index, + const cuvs::neighbors::mg_search_params& search_params, + raft::host_matrix_view queries, + raft::host_matrix_view neighbors, + raft::host_matrix_view distances); + +/** @copydoc search */ +void search( + const raft::resources& clique, + const cuvs::neighbors::mg_index, half, uint32_t>& + index, + const cuvs::neighbors::mg_search_params& search_params, + raft::host_matrix_view queries, + raft::host_matrix_view neighbors, + raft::host_matrix_view distances); /// \ingroup mg_cpp_index_search /** @@ -2786,7 +3667,18 @@ void search(const raft::resources& clique, */ void search( const raft::resources& clique, - const cuvs::neighbors::mg_index, int8_t, uint32_t>& index, + const cuvs::neighbors::mg_index, int8_t, uint32_t>& + index, + const cuvs::neighbors::mg_search_params& search_params, + raft::host_matrix_view queries, + raft::host_matrix_view neighbors, + raft::host_matrix_view distances); + +/** @copydoc search */ +void search( + const raft::resources& clique, + const cuvs::neighbors::mg_index, int8_t, uint32_t>& + index, const cuvs::neighbors::mg_search_params& search_params, raft::host_matrix_view queries, raft::host_matrix_view neighbors, @@ -2816,12 +3708,22 @@ void search( */ void search( const raft::resources& clique, - const cuvs::neighbors::mg_index, uint8_t, uint32_t>& index, + const cuvs::neighbors::mg_index, uint8_t, uint32_t>& + index, const cuvs::neighbors::mg_search_params& search_params, raft::host_matrix_view queries, raft::host_matrix_view neighbors, raft::host_matrix_view distances); +/** @copydoc search */ +void search(const raft::resources& clique, + const cuvs::neighbors:: + mg_index, uint8_t, uint32_t>& index, + const cuvs::neighbors::mg_search_params& search_params, + raft::host_matrix_view queries, + raft::host_matrix_view neighbors, + raft::host_matrix_view distances); + /// \ingroup mg_cpp_index_search /** * @brief Searches a multi-GPU index @@ -2844,12 +3746,24 @@ void search( * @param[out] distances a row-major matrix on host [n_rows, n_neighbors] * */ -void search(const raft::resources& clique, - const cuvs::neighbors::mg_index, float, uint32_t>& index, - const cuvs::neighbors::mg_search_params& search_params, - raft::host_matrix_view queries, - raft::host_matrix_view neighbors, - raft::host_matrix_view distances); +void search( + const raft::resources& clique, + const cuvs::neighbors::mg_index, float, uint32_t>& + index, + const cuvs::neighbors::mg_search_params& search_params, + raft::host_matrix_view queries, + raft::host_matrix_view neighbors, + raft::host_matrix_view distances); + +/** @copydoc search */ +void search( + const raft::resources& clique, + const cuvs::neighbors::mg_index, float, uint32_t>& + index, + const cuvs::neighbors::mg_search_params& search_params, + raft::host_matrix_view queries, + raft::host_matrix_view neighbors, + raft::host_matrix_view distances); /// \ingroup mg_cpp_index_search /** @@ -2873,12 +3787,24 @@ void search(const raft::resources& clique, * @param[out] distances a row-major matrix on host [n_rows, n_neighbors] * */ -void search(const raft::resources& clique, - const cuvs::neighbors::mg_index, half, uint32_t>& index, - const cuvs::neighbors::mg_search_params& search_params, - raft::host_matrix_view queries, - raft::host_matrix_view neighbors, - raft::host_matrix_view distances); +void search( + const raft::resources& clique, + const cuvs::neighbors::mg_index, half, uint32_t>& + index, + const cuvs::neighbors::mg_search_params& search_params, + raft::host_matrix_view queries, + raft::host_matrix_view neighbors, + raft::host_matrix_view distances); + +/** @copydoc search */ +void search( + const raft::resources& clique, + const cuvs::neighbors::mg_index, half, uint32_t>& + index, + const cuvs::neighbors::mg_search_params& search_params, + raft::host_matrix_view queries, + raft::host_matrix_view neighbors, + raft::host_matrix_view distances); /// \ingroup mg_cpp_index_search /** @@ -2904,7 +3830,18 @@ void search(const raft::resources& clique, */ void search( const raft::resources& clique, - const cuvs::neighbors::mg_index, int8_t, uint32_t>& index, + const cuvs::neighbors::mg_index, int8_t, uint32_t>& + index, + const cuvs::neighbors::mg_search_params& search_params, + raft::host_matrix_view queries, + raft::host_matrix_view neighbors, + raft::host_matrix_view distances); + +/** @copydoc search */ +void search( + const raft::resources& clique, + const cuvs::neighbors::mg_index, int8_t, uint32_t>& + index, const cuvs::neighbors::mg_search_params& search_params, raft::host_matrix_view queries, raft::host_matrix_view neighbors, @@ -2934,12 +3871,22 @@ void search( */ void search( const raft::resources& clique, - const cuvs::neighbors::mg_index, uint8_t, uint32_t>& index, + const cuvs::neighbors::mg_index, uint8_t, uint32_t>& + index, const cuvs::neighbors::mg_search_params& search_params, raft::host_matrix_view queries, raft::host_matrix_view neighbors, raft::host_matrix_view distances); +/** @copydoc search */ +void search(const raft::resources& clique, + const cuvs::neighbors:: + mg_index, uint8_t, uint32_t>& index, + const cuvs::neighbors::mg_search_params& search_params, + raft::host_matrix_view queries, + raft::host_matrix_view neighbors, + raft::host_matrix_view distances); + /// \defgroup mg_cpp_serialize ANN MG index serialization /// \ingroup mg_cpp_serialize @@ -2962,7 +3909,15 @@ void search( */ void serialize( const raft::resources& clique, - const cuvs::neighbors::mg_index, float, uint32_t>& index, + const cuvs::neighbors::mg_index, float, uint32_t>& + index, + const std::string& filename); + +/** @copydoc serialize */ +void serialize( + const raft::resources& clique, + const cuvs::neighbors::mg_index, float, uint32_t>& + index, const std::string& filename); /// \ingroup mg_cpp_serialize @@ -2983,9 +3938,18 @@ void serialize( * @param[in] filename path to the file to be serialized * */ -void serialize(const raft::resources& clique, - const cuvs::neighbors::mg_index, half, uint32_t>& index, - const std::string& filename); +void serialize( + const raft::resources& clique, + const cuvs::neighbors::mg_index, half, uint32_t>& + index, + const std::string& filename); + +/** @copydoc serialize */ +void serialize( + const raft::resources& clique, + const cuvs::neighbors::mg_index, half, uint32_t>& + index, + const std::string& filename); /// \ingroup mg_cpp_serialize /** @@ -3007,7 +3971,15 @@ void serialize(const raft::resources& clique, */ void serialize( const raft::resources& clique, - const cuvs::neighbors::mg_index, int8_t, uint32_t>& index, + const cuvs::neighbors::mg_index, int8_t, uint32_t>& + index, + const std::string& filename); + +/** @copydoc serialize */ +void serialize( + const raft::resources& clique, + const cuvs::neighbors::mg_index, int8_t, uint32_t>& + index, const std::string& filename); /// \ingroup mg_cpp_serialize @@ -3030,9 +4002,17 @@ void serialize( */ void serialize( const raft::resources& clique, - const cuvs::neighbors::mg_index, uint8_t, uint32_t>& index, + const cuvs::neighbors::mg_index, uint8_t, uint32_t>& + index, const std::string& filename); +/** @copydoc serialize */ +void serialize(const raft::resources& clique, + const cuvs::neighbors::mg_index, + uint8_t, + uint32_t>& index, + const std::string& filename); + /// \defgroup mg_cpp_deserialize ANN MG index deserialization /// \ingroup mg_cpp_deserialize @@ -3046,7 +4026,9 @@ void serialize( * auto index = cuvs::neighbors::cagra::build(clique, index_params, index_dataset); * const std::string filename = "mg_index.cuvs"; * cuvs::neighbors::cagra::serialize(clique, index, filename); - * auto new_index = cuvs::neighbors::cagra::deserialize(clique, filename); + * cuvs::neighbors::mg_index, float, uint32_t> + * new_index(clique, REPLICATED); + * cuvs::neighbors::cagra::deserialize(clique, filename, &new_index); * * @endcode * @@ -3055,8 +4037,14 @@ void serialize( * */ template -auto deserialize(const raft::resources& clique, const std::string& filename) - -> cuvs::neighbors::mg_index, T, IdxT>; +void deserialize(const raft::resources& clique, + const std::string& filename, + cuvs::neighbors::mg_index, T, IdxT>* index); + +template +void deserialize(const raft::resources& clique, + const std::string& filename, + cuvs::neighbors::mg_index, T, IdxT>* index); /// \defgroup mg_cpp_distribute ANN MG local index distribution @@ -3072,7 +4060,9 @@ auto deserialize(const raft::resources& clique, const std::string& filename) * auto index = cuvs::neighbors::cagra::build(clique, index_params, index_dataset); * const std::string filename = "local_index.cuvs"; * cuvs::neighbors::cagra::serialize(clique, filename, index); - * auto new_index = cuvs::neighbors::cagra::distribute(clique, filename); + * cuvs::neighbors::mg_index, float, uint32_t> + * distributed_index(clique, REPLICATED); + * cuvs::neighbors::cagra::distribute(clique, filename, &distributed_index); * * @endcode * @@ -3081,8 +4071,14 @@ auto deserialize(const raft::resources& clique, const std::string& filename) * */ template -auto distribute(const raft::resources& clique, const std::string& filename) - -> cuvs::neighbors::mg_index, T, IdxT>; +void distribute(const raft::resources& clique, + const std::string& filename, + cuvs::neighbors::mg_index, T, IdxT>* index); + +template +void distribute(const raft::resources& clique, + const std::string& filename, + cuvs::neighbors::mg_index, T, IdxT>* index); /** * @brief Build a kNN graph using IVF-PQ. @@ -3110,7 +4106,7 @@ auto distribute(const raft::resources& clique, const std::string& filename) * auto optimized_gaph = raft::make_host_matrix(dataset.extent(0), 64); * cagra::optimize(res, dataset, knn_graph.view(), optimized_graph.view()); * // Construct an index from dataset and optimized knn_graph - * auto index = cagra::index(res, build_params.metric(), dataset, + * auto index = cagra::device_padded_index(res, build_params.metric(), dataset, * optimized_graph.view()); * @endcode * @@ -3150,7 +4146,7 @@ void build_knn_graph(raft::resources const& res, * auto optimized_gaph = raft::make_host_matrix(dataset.extent(0), 64); * cagra::optimize(res, dataset, knn_graph.view(), optimized_graph.view()); * // Construct an index from dataset and optimized knn_graph - * auto index = cagra::index(res, build_params.metric(), dataset, + * auto index = cagra::device_padded_index(res, build_params.metric(), dataset, * optimized_graph.view()); * @endcode * @@ -3190,7 +4186,7 @@ void build_knn_graph(raft::resources const& res, * auto optimized_gaph = raft::make_host_matrix(dataset.extent(0), 64); * cagra::optimize(res, dataset, knn_graph.view(), optimized_graph.view()); * // Construct an index from dataset and optimized knn_graph - * auto index = cagra::index(res, build_params.metric(), dataset, + * auto index = cagra::device_padded_index(res, build_params.metric(), dataset, * optimized_graph.view()); * @endcode * @@ -3230,7 +4226,7 @@ void build_knn_graph(raft::resources const& res, * auto optimized_gaph = raft::make_host_matrix(dataset.extent(0), 64); * cagra::optimize(res, dataset, knn_graph.view(), optimized_graph.view()); * // Construct an index from dataset and optimized knn_graph - * auto index = cagra::index(res, build_params.metric(), dataset, + * auto index = cagra::device_padded_index(res, build_params.metric(), dataset, * optimized_graph.view()); * @endcode * @@ -3244,6 +4240,137 @@ void build_knn_graph(raft::resources const& res, raft::host_matrix_view knn_graph, cuvs::neighbors::cagra::graph_build_params::ivf_pq_params build_params); +namespace detail { + +/** + * @brief Internal helper to transfer ACE disk FDs between CAGRA indexes. + * + * @internal + */ +struct fd_transfer { + template + static inline void steal_disk_fds_to(raft::resources const& res, + index& src, + index& dst) + { + if (src.dataset_fd().has_value()) { + dst.update_dataset(res, std::move(*src.steal_dataset_fd_())); + } + if (src.graph_fd().has_value()) { dst.update_graph(res, std::move(*src.steal_graph_fd_())); } + if (src.mapping_fd().has_value()) { + dst.update_mapping(res, std::move(*src.steal_mapping_fd_())); + } + } +}; + +/** + * @brief Copy a host-resident CAGRA index graph into a new device-resident index (graph only). + * + * The returned device index has no dataset attached. Prefer `attach_device_dataset_on_host_index` + * for the public host-build → search workflow. + * + * @internal + */ +template + requires cuvs::neighbors::is_host_dataset_view_v +auto convert_host_to_device_index(raft::resources const& res, index const& src) + -> index> +{ + using DeviceViewT = cuvs::neighbors::device_counterpart_t; + using GraphIndexType = typename index::graph_index_type; + index out(res, src.metric()); + if (src.graph().size() > 0) { + // The graph lives in device memory owned by `src`. `update_graph(device_view)` would only + // store a view (no ownership transfer), leaving `out` with a dangling pointer once `src` + // is destroyed. Copy device→host→device so that `out` owns its graph memory. + auto graph_host = + raft::make_host_matrix(src.graph().extent(0), src.graph().extent(1)); + raft::copy(graph_host.data_handle(), + src.graph().data_handle(), + src.graph().size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + out.update_graph(res, raft::make_const_mdspan(graph_host.view())); // host overload: copies H→D + } + return out; +} + +} // namespace detail + +/** + * @brief Convert a host index to device and attach a device dataset in one step. + * + * Copies the graph from `host_idx` into a new device index and attaches `device_dataset`. + * + * @tparam T element type + * @tparam IdxT index type + * @tparam HostViewT host-resident dataset view type + * @tparam DeviceViewT device-resident dataset view of the same kind + * @param[in] res RAFT resources + * @param[in] host_idx host index returned by `build(res, params, host_view)` + * @param[in] device_dataset device dataset view to attach (caller owns underlying memory) + * @return device index with graph and dataset ready for search + */ +template + requires cuvs::neighbors::compatible_host_device_dataset_views_v +auto attach_device_dataset_on_host_index(raft::resources const& res, + index const& host_idx, + DeviceViewT const& device_dataset) + -> index +{ + auto device_idx = detail::convert_host_to_device_index(res, host_idx); + device_idx.update_dataset(res, device_dataset); + return device_idx; +} + +/** + * @brief Attach a padded device dataset to a standard-device index for search. + * + * Copies the graph from `standard_idx` into a new `device_padded_index` and attaches + * `padded_dataset`. CAGRA search kernels require CAGRA-aligned row pitch; build from a + * `device_standard_dataset_view` (any stride), then call this before `search`. + * + * @param[in] res RAFT resources + * @param[in] standard_idx index returned by `build` with a standard device dataset view + * @param[in] padded_dataset device padded dataset view (caller owns underlying memory) + * @return device padded index with graph and dataset ready for search + */ +template +auto attach_padded_dataset_for_search( + raft::resources const& res, + index> const& standard_idx, + cuvs::neighbors::device_padded_dataset_view const& padded_dataset) + -> device_padded_index +{ + RAFT_EXPECTS(padded_dataset.n_rows() == standard_idx.size(), + "Padded dataset row count must match the index size"); + RAFT_EXPECTS(padded_dataset.dim() == standard_idx.dim(), + "Padded dataset dimension must match the index dimension"); + + device_padded_index out(res, standard_idx.metric()); + if (standard_idx.graph().extent(0) > 0) { + using GraphIndexType = + typename index>:: + graph_index_type; + auto graph_host = raft::make_host_matrix( + standard_idx.graph().extent(0), standard_idx.graph().extent(1)); + raft::copy(graph_host.data_handle(), + standard_idx.graph().data_handle(), + standard_idx.graph().size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + out.update_graph(res, raft::make_const_mdspan(graph_host.view())); + } + if (standard_idx.source_indices().has_value()) { + out.update_source_indices(res, standard_idx.source_indices().value()); + } + out.update_dataset(res, padded_dataset); + return out; +} + } // namespace cagra } // namespace neighbors } // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 2fd804f115..8a66810fec 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -8,13 +8,18 @@ #include #include #include +#include #include #include #include +#include +#include +#include #include +#include #include #include -#include // get_device_for_address +#include // get_device_for_address, copy_matrix #include // rounding up #include @@ -22,12 +27,15 @@ #include #include +#include + +#include +#include #include #include #include #include #include - #ifdef __cpp_lib_bitops #include #endif @@ -135,323 +143,217 @@ enum class MergeStrategy { /** @} */ // end group neighbors_index -/** Two-dimensional dataset; maybe owning, maybe compressed, maybe strided. */ +/** + * @brief Tags selecting dataset representation for `dataset` / `dataset_view`. + * + * Each container defines nested `owning_storage` then `view_storage` (aliases into `detail::*` + * storage types shared by device/host). Accessibility (device vs host) is selected by the + * `Accessor` template parameter on `dataset` / `dataset_view`, not by duplicating containers. + * Layout kinds: empty, padded, standard, VPQ. `dataset` / `dataset_view` only express ownership + * vs view. + */ + +template +struct dataset; + +template +struct dataset_view; + +namespace detail { + +// Default owning/view accessors for public dataset aliases. +template +using device_owning_accessor = raft::device_accessor>; + +template +using host_owning_accessor = raft::host_accessor>; + +template +using device_view_accessor = raft::device_accessor>; + +template +using host_view_accessor = raft::host_accessor>; + +/** View accessor paired with an owning dataset accessor (same residency). */ +template +using dataset_view_accessor_for_owning = std::conditional_t, + host_view_accessor>; + +/** Owning accessor paired with a view accessor (same residency). */ +template +using dataset_owning_accessor_for_view = std::conditional_t, + host_owning_accessor>; + +template +using dense_owning_matrix = std::conditional_t, + raft::host_matrix>; + +template +using dense_view_matrix = + std::conditional_t, + raft::host_matrix_view>; + +template +using vpq_vq_book_matrix = std::conditional_t, + raft::host_matrix>; + +template +using vpq_data_matrix = std::conditional_t, + raft::host_matrix>; + +// ----------------------------------------------------------------------------- +// empty +// ----------------------------------------------------------------------------- + template -struct dataset { - using index_type = IdxT; - /** Size of the dataset. */ - [[nodiscard]] virtual auto n_rows() const noexcept -> index_type = 0; - /** Dimensionality of the dataset. */ - [[nodiscard]] virtual auto dim() const noexcept -> uint32_t = 0; - /** Whether the object owns the data. */ - [[nodiscard]] virtual auto is_owning() const noexcept -> bool = 0; - virtual ~dataset() noexcept = default; +struct empty_dataset_storage { + uint32_t suggested_dim{}; + empty_dataset_storage() noexcept = default; + explicit empty_dataset_storage(uint32_t dim) noexcept : suggested_dim(dim) {} + [[nodiscard]] auto n_rows() const noexcept -> IdxT { return 0; } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return suggested_dim; } }; template -struct empty_dataset : public dataset { - using index_type = IdxT; - uint32_t suggested_dim; - explicit empty_dataset(uint32_t dim) noexcept : suggested_dim(dim) {} - [[nodiscard]] auto n_rows() const noexcept -> index_type final { return 0; } - [[nodiscard]] auto dim() const noexcept -> uint32_t final { return suggested_dim; } - [[nodiscard]] auto is_owning() const noexcept -> bool final { return true; } -}; +using empty_dataset_owning_storage = empty_dataset_storage; -template -struct strided_dataset : public dataset { - using index_type = IdxT; - using value_type = DataT; - using view_type = raft::device_matrix_view; - [[nodiscard]] auto n_rows() const noexcept -> index_type final { return view().extent(0); } - [[nodiscard]] auto dim() const noexcept -> uint32_t final - { - return static_cast(view().extent(1)); - } - /** Leading dimension of the dataset. */ - [[nodiscard]] constexpr auto stride() const noexcept -> uint32_t - { - auto v = view(); - return static_cast(v.stride(0) > 0 ? v.stride(0) : v.extent(1)); - } - /** Get the view of the data. */ - [[nodiscard]] virtual auto view() const noexcept -> view_type = 0; -}; +template +using empty_dataset_view_storage = empty_dataset_storage; -template -struct non_owning_dataset : public strided_dataset { - using index_type = IdxT; - using value_type = DataT; - using typename strided_dataset::view_type; - view_type data; - explicit non_owning_dataset(view_type v) noexcept : data(v) {} - [[nodiscard]] auto is_owning() const noexcept -> bool final { return false; } - [[nodiscard]] auto view() const noexcept -> view_type final { return data; }; -}; +// ----------------------------------------------------------------------------- +// dense row-major (logical dim may differ from row pitch; shared by padded & standard) +// ----------------------------------------------------------------------------- -template -struct owning_dataset : public strided_dataset { - using index_type = IdxT; - using value_type = DataT; - using typename strided_dataset::view_type; - using storage_type = - raft::mdarray, LayoutPolicy, ContainerPolicy>; - using mapping_type = typename view_type::mapping_type; - storage_type data; - mapping_type view_mapping; - owning_dataset(storage_type&& store, mapping_type view_mapping) noexcept - : data{std::move(store)}, view_mapping{view_mapping} +/** + * Dense row-major owning storage shared by padded and standard dataset containers. + * + * Template parameters: + * - MatrixT: owning matrix type that stores the payload (host/device matrix). + * - ViewT: non-owning row-major view type returned by `view()`. + * - DataT: scalar element type of the dataset payload. + * - IdxT: index type used for row counts (`n_rows()` return type). + */ +template +struct dense_row_major_dataset_owning_storage { + MatrixT data_; + uint32_t logical_dim_; + + dense_row_major_dataset_owning_storage(MatrixT&& data, uint32_t logical_dim) noexcept + : data_{std::move(data)}, logical_dim_{logical_dim} { } - [[nodiscard]] auto is_owning() const noexcept -> bool final { return true; } - [[nodiscard]] auto view() const noexcept -> view_type final + [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data_.extent(0); } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return logical_dim_; } + [[nodiscard]] auto stride() const noexcept -> uint32_t { - return view_type{data.data_handle(), view_mapping}; - }; + return static_cast(data_.extent(1)); + } + [[nodiscard]] auto view() const noexcept -> ViewT { return data_.view(); } + [[nodiscard]] auto data_handle() noexcept -> DataT* { return data_.data_handle(); } + [[nodiscard]] auto data_handle() const noexcept -> const DataT* { return data_.data_handle(); } }; -template -struct is_strided_dataset : std::false_type {}; +template +struct dense_row_major_dataset_view_storage { + ViewT data_; + uint32_t logical_dim_; -template -struct is_strided_dataset> : std::true_type {}; + dense_row_major_dataset_view_storage() noexcept = default; -template -struct is_strided_dataset> : std::true_type {}; + explicit dense_row_major_dataset_view_storage(ViewT v) noexcept + : data_(v), logical_dim_(static_cast(v.extent(1))) + { + } -template -struct is_strided_dataset> - : std::true_type {}; + dense_row_major_dataset_view_storage(ViewT v, uint32_t logical_dim) noexcept + : data_(v), logical_dim_(logical_dim) + { + } -template -inline constexpr bool is_strided_dataset_v = is_strided_dataset::value; + dense_row_major_dataset_view_storage(dense_row_major_dataset_view_storage const& other) noexcept + : data_(other.data_), logical_dim_(other.logical_dim_) + { + } -/** - * @brief Construct a strided matrix from any mdarray or mdspan. - * - * This function constructs a non-owning view if the input satisfied two conditions: - * - * 1) The data is accessible from the current device - * 2) The memory layout is the same as expected (row-major matrix with the required stride) - * - * Otherwise, this function constructs an owning device matrix and copies the data. - * When the data is copied, padding elements are filled with zeroes. - * - * @tparam SrcT the source mdarray or mdspan - * - * @param[in] res raft resources handle - * @param[in] src the source mdarray or mdspan - * @param[in] required_stride the leading dimension (in elements) - * @return maybe owning current-device-accessible strided matrix - */ -template -auto make_strided_dataset(const raft::resources& res, const SrcT& src, uint32_t required_stride) - -> std::unique_ptr> -{ - using extents_type = typename SrcT::extents_type; - using value_type = typename SrcT::value_type; - using index_type = typename SrcT::index_type; - using layout_type = typename SrcT::layout_type; - static_assert(extents_type::rank() == 2, "The input must be a matrix."); - static_assert(std::is_same_v || - std::is_same_v> || - std::is_same_v, - "The input must be row-major"); - RAFT_EXPECTS(src.extent(1) <= required_stride, - "The input row length must be not larger than the desired stride."); - cudaPointerAttributes ptr_attrs; - RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src.data_handle())); - auto* device_ptr = reinterpret_cast(ptr_attrs.devicePointer); - const uint32_t src_stride = src.stride(0) > 0 ? src.stride(0) : src.extent(1); - const bool device_accessible = device_ptr != nullptr; - const bool row_major = src.stride(1) <= 1; - const bool stride_matches = required_stride == src_stride; - - if (device_accessible && row_major && stride_matches) { - // Everything matches: make a non-owning dataset - return std::make_unique>( - raft::make_device_strided_matrix_view( - device_ptr, src.extent(0), src.extent(1), required_stride)); + [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data_.extent(0); } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return logical_dim_; } + [[nodiscard]] auto stride() const noexcept -> uint32_t + { + return static_cast(data_.stride(0) > 0 ? data_.stride(0) : data_.extent(1)); } - // Something is wrong: have to make a copy and produce an owning dataset - auto out_layout = - raft::make_strided_layout(src.extents(), cuda::std::array{required_stride, 1}); - auto out_array = - raft::make_device_matrix(res, src.extent(0), required_stride); - - using out_mdarray_type = decltype(out_array); - using out_layout_type = typename out_mdarray_type::layout_type; - using out_container_policy_type = typename out_mdarray_type::container_policy_type; - using out_owning_type = - owning_dataset; + [[nodiscard]] auto view() const noexcept -> ViewT { return data_; } +}; - RAFT_CUDA_TRY(cudaMemsetAsync(out_array.data_handle(), - 0, - out_array.size() * sizeof(value_type), - raft::resource::get_cuda_stream(res))); - raft::copy_matrix(out_array.data_handle(), - required_stride, - src.data_handle(), - src_stride, - src.extent(1), - src.extent(0), - raft::resource::get_cuda_stream(res)); +template +using padded_dataset_owning_storage = + dense_row_major_dataset_owning_storage; - return std::make_unique(std::move(out_array), out_layout); -} +template +using padded_dataset_view_storage = dense_row_major_dataset_view_storage; -/** - * @brief Construct a strided matrix from any mdarray. - * - * This function constructs an owning device matrix and copies the data. - * When the data is copied, padding elements are filled with zeroes. - * - * @tparam DataT - * @tparam IdxT - * @tparam LayoutPolicy - * @tparam ContainerPolicy - * - * @param[in] res raft resources handle - * @param[in] src the source mdarray or mdspan - * @param[in] required_stride the leading dimension (in elements) - * @return owning current-device-accessible strided matrix - */ -template -auto make_strided_dataset( - const raft::resources& res, - raft::mdarray, LayoutPolicy, ContainerPolicy>&& src, - uint32_t required_stride) -> std::unique_ptr> -{ - using value_type = DataT; - using index_type = IdxT; - using layout_type = LayoutPolicy; - using container_policy_type = ContainerPolicy; - static_assert(std::is_same_v || - std::is_same_v> || - std::is_same_v, - "The input must be row-major"); - RAFT_EXPECTS(src.extent(1) <= required_stride, - "The input row length must be not larger than the desired stride."); - const uint32_t src_stride = src.stride(0) > 0 ? src.stride(0) : src.extent(1); - const bool stride_matches = required_stride == src_stride; - - auto out_layout = - raft::make_strided_layout(src.extents(), cuda::std::array{required_stride, 1}); - - using out_mdarray_type = raft::device_matrix; - using out_layout_type = typename out_mdarray_type::layout_type; - using out_container_policy_type = typename out_mdarray_type::container_policy_type; - using out_owning_type = - owning_dataset; - - if constexpr (std::is_same_v && - std::is_same_v) { - if (stride_matches) { - // Everything matches, we can own the mdarray - return std::make_unique(std::move(src), out_layout); - } - } - // Something is wrong: have to make a copy and produce an owning dataset - auto out_array = - raft::make_device_matrix(res, src.extent(0), required_stride); +template +using standard_dataset_owning_storage = + dense_row_major_dataset_owning_storage; - RAFT_CUDA_TRY(cudaMemsetAsync(out_array.data_handle(), - 0, - out_array.size() * sizeof(value_type), - raft::resource::get_cuda_stream(res))); - raft::copy_matrix(out_array.data_handle(), - required_stride, - src.data_handle(), - src_stride, - src.extent(1), - src.extent(0), - raft::resource::get_cuda_stream(res)); +template +using standard_dataset_view_storage = dense_row_major_dataset_view_storage; - return std::make_unique(std::move(out_array), out_layout); -} +// ----------------------------------------------------------------------------- +// VPQ compressed +// ----------------------------------------------------------------------------- /** - * @brief Construct a strided matrix from any mdarray or mdspan. + * Owning storage for VPQ-compressed datasets. * - * A variant `make_strided_dataset` that allows specifying the byte alignment instead of the - * explicit stride length. - * - * @tparam SrcT the source mdarray or mdspan - * - * @param[in] res raft resources handle - * @param[in] src the source mdarray or mdspan - * @param[in] align_bytes the required byte alignment for the dataset rows. - * @return maybe owning current-device-accessible strided matrix + * Template parameters: + * - VqBookMatrixT: owning matrix type for the VQ codebook. + * - PqBookMatrixT: owning matrix type for the PQ codebook. + * - DataMatrixT: owning matrix type for encoded row data (uint8 codes). + * - MathT: floating-point type used by VQ/PQ codebooks. + * - IdxT: index type used for row counts (`n_rows()` return type). */ -template -auto make_aligned_dataset(const raft::resources& res, SrcT src, uint32_t align_bytes = 16) - -> std::unique_ptr> -{ - using source_type = std::remove_cv_t>; - using value_type = typename source_type::value_type; - constexpr size_t kSize = sizeof(value_type); - uint32_t required_stride = - raft::round_up_safe(src.extent(1) * kSize, std::lcm(align_bytes, kSize)) / kSize; - return make_strided_dataset(res, std::forward(src), required_stride); -} -/** - * @brief VPQ compressed dataset. - * - * The dataset is compressed using two level quantization - * - * 1. Vector Quantization - * 2. Product Quantization of residuals - * - * @tparam MathT the type of elements in the codebooks - * @tparam IdxT type of the vector indices (represent dataset.extent(0)) - * - */ -template -struct vpq_dataset : public dataset { - using index_type = IdxT; - using math_type = MathT; - /** Vector Quantization codebook - "coarse cluster centers". */ - raft::device_matrix vq_code_book; - /** Product Quantization codebook - "fine cluster centers". */ - raft::device_matrix pq_code_book; - /** Compressed dataset. */ - raft::device_matrix data; - - vpq_dataset(raft::device_matrix&& vq_code_book, - raft::device_matrix&& pq_code_book, - raft::device_matrix&& data) +template +struct vpq_dataset_owning_storage { + /** Floating-point type used for VQ/PQ codebooks (rows are still uint8 codes). */ + using math_type = MathT; + + VqBookMatrixT vq_code_book; + PqBookMatrixT pq_code_book; + DataMatrixT data; + + vpq_dataset_owning_storage(VqBookMatrixT&& vq_code_book, + PqBookMatrixT&& pq_code_book, + DataMatrixT&& data) noexcept : vq_code_book{std::move(vq_code_book)}, pq_code_book{std::move(pq_code_book)}, data{std::move(data)} { } - [[nodiscard]] auto n_rows() const noexcept -> index_type final { return data.extent(0); } - [[nodiscard]] auto dim() const noexcept -> uint32_t final { return vq_code_book.extent(1); } - [[nodiscard]] auto is_owning() const noexcept -> bool final { return true; } + [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data.extent(0); } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return vq_code_book.extent(1); } - /** Row length of the encoded data in bytes. */ [[nodiscard]] constexpr inline auto encoded_row_length() const noexcept -> uint32_t { return data.extent(1); } - /** The number of "coarse cluster centers" */ [[nodiscard]] constexpr inline auto vq_n_centers() const noexcept -> uint32_t { return vq_code_book.extent(0); } - /** The bit length of an encoded vector element after compression by PQ. */ [[nodiscard]] constexpr inline auto pq_bits() const noexcept -> uint32_t { - /* - NOTE: pq_bits and the book size - - Normally, we'd store `pq_bits` as a part of the index. - However, we know there's an invariant `pq_n_centers = 1 << pq_bits`, i.e. the codebook size is - the same as the number of possible code values. Hence, we don't store the pq_bits and derive it - from the array dimensions instead. - */ auto pq_width = pq_n_centers(); #ifdef __cpp_lib_bitops return std::countr_zero(pq_width); @@ -464,32 +366,965 @@ struct vpq_dataset : public dataset { return pq_bits; #endif } - /** The dimensionality of an encoded vector after compression by PQ. */ [[nodiscard]] constexpr inline auto pq_dim() const noexcept -> uint32_t { return raft::div_rounding_up_unsafe(dim(), pq_len()); } - /** Dimensionality of a subspaces, i.e. the number of vector components mapped to a subspace */ [[nodiscard]] constexpr inline auto pq_len() const noexcept -> uint32_t { return pq_code_book.extent(1); } - /** The number of vectors in a PQ codebook (`1 << pq_bits`). */ [[nodiscard]] constexpr inline auto pq_n_centers() const noexcept -> uint32_t { return pq_code_book.extent(0); } }; +template +struct vpq_dataset_view_storage { + using owning_dataset_type = + dataset>; + + owning_dataset_type const* dataset_{nullptr}; + + vpq_dataset_view_storage() = default; + + explicit vpq_dataset_view_storage(owning_dataset_type const* ptr) : dataset_(ptr) + { + RAFT_EXPECTS(ptr != nullptr, "vpq_dataset_view: null dataset pointer"); + } + + [[nodiscard]] auto n_rows() const noexcept + { + using idx_type = decltype(std::declval().n_rows()); + return dataset_ != nullptr ? dataset_->n_rows() : idx_type{0}; + } + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + return dataset_ != nullptr ? dataset_->dim() : uint32_t{0}; + } + [[nodiscard]] owning_dataset_type const& dset() const noexcept { return *dataset_; } +}; + +} // namespace detail + +// ----------------------------------------------------------------------------- +// empty +// ----------------------------------------------------------------------------- + +struct empty_dataset_container { + template + using owning_storage = detail::empty_dataset_owning_storage; + template + using view_storage = detail::empty_dataset_view_storage; +}; + +// ----------------------------------------------------------------------------- +// padded (row-major with logical dim vs stride) +// ----------------------------------------------------------------------------- + +struct padded_dataset_container { + template + using owning_storage = + detail::padded_dataset_owning_storage, + detail::dense_view_matrix, + DataT, + IdxT>; + template + using view_storage = detail:: + padded_dataset_view_storage, DataT, IdxT>; +}; + +// ----------------------------------------------------------------------------- +// standard (row-major with arbitrary stride; no CAGRA alignment requirement) +// ----------------------------------------------------------------------------- + +struct standard_dataset_container { + template + using owning_storage = + detail::standard_dataset_owning_storage, + detail::dense_view_matrix, + DataT, + IdxT>; + template + using view_storage = detail:: + standard_dataset_view_storage, DataT, IdxT>; +}; + +// ----------------------------------------------------------------------------- +// VPQ compressed +// ----------------------------------------------------------------------------- + +struct vpq_dataset_container { + template + using owning_storage = + detail::vpq_dataset_owning_storage, + detail::vpq_vq_book_matrix, + detail::vpq_data_matrix, + MathT, + IdxT>; + template + using view_storage = + detail::vpq_dataset_view_storage; +}; + +template +struct dataset { + static_assert(!std::is_same_v, + "dataset: unsupported ContainerType / type-parameter combination"); +}; + +template +struct dataset_view { + static_assert(!std::is_same_v, + "dataset_view: unsupported ContainerType / type-parameter combination"); +}; + +// ----------------------------------------------------------------------------- +// empty +// ----------------------------------------------------------------------------- + +template +struct dataset + : empty_dataset_container::template owning_storage { + using container_type = empty_dataset_container; + using owning_storage_type = typename container_type::template owning_storage; + using owning_storage_type::owning_storage_type; + + [[nodiscard]] auto as_dataset_view() const noexcept + -> dataset_view> + { + return dataset_view>{this->dim()}; + } +}; + +template +struct dataset_view + : empty_dataset_container::template view_storage { + using container_type = empty_dataset_container; + using view_storage_type = typename container_type::template view_storage; + using view_storage_type::view_storage_type; +}; + +// ----------------------------------------------------------------------------- +// standard (row-major with arbitrary stride) +// ----------------------------------------------------------------------------- + +template +struct dataset + : standard_dataset_container::template owning_storage { + using container_type = standard_dataset_container; + using owning_storage_type = + typename container_type::template owning_storage; + using owning_storage_type::owning_storage_type; + + [[nodiscard]] auto as_dataset_view() const noexcept + -> dataset_view> + { + return dataset_view>(this->view(), + this->dim()); + } +}; + +template +struct dataset_view + : standard_dataset_container::template view_storage { + using container_type = standard_dataset_container; + using view_storage_type = typename container_type::template view_storage; + using view_storage_type::view_storage_type; +}; + +// ----------------------------------------------------------------------------- +// padded (row-major with logical dim vs stride) +// ----------------------------------------------------------------------------- + +template +struct dataset + : padded_dataset_container::template owning_storage { + using container_type = padded_dataset_container; + using owning_storage_type = + typename container_type::template owning_storage; + using owning_storage_type::owning_storage_type; + + [[nodiscard]] auto as_dataset_view() const noexcept + -> dataset_view> + { + return dataset_view>(this->view(), + this->dim()); + } +}; + +template +struct dataset_view + : padded_dataset_container::template view_storage { + using container_type = padded_dataset_container; + using view_storage_type = typename container_type::template view_storage; + using view_storage_type::view_storage_type; +}; + +// ----------------------------------------------------------------------------- +// VPQ compressed (view holds non-owning pointer to owning dataset) +// ----------------------------------------------------------------------------- + +template +struct dataset + : vpq_dataset_container::template owning_storage { + using container_type = vpq_dataset_container; + using owning_storage_type = + typename container_type::template owning_storage; + using owning_storage_type::owning_storage_type; + + [[nodiscard]] auto as_dataset_view() const + -> dataset_view> + { + return dataset_view>{this}; + } +}; + +template +struct dataset_view + : vpq_dataset_container::template view_storage { + using container_type = vpq_dataset_container; + using view_storage_type = typename container_type::template view_storage; + using view_storage_type::view_storage_type; +}; + +/** + * @brief Aliases for concrete `dataset` / `dataset_view` layouts. + */ +template +using device_empty_dataset = + dataset>; + +template +using device_empty_dataset_view = + dataset_view>; + +template +using host_empty_dataset = + dataset>; + +template +using host_empty_dataset_view = + dataset_view>; + +template +using device_padded_dataset = + dataset>; + +template +using device_padded_dataset_view = + dataset_view>; + +template +using host_padded_dataset = + dataset>; + +template +using host_padded_dataset_view = + dataset_view>; + +template +using device_standard_dataset = + dataset>; + +template +using device_standard_dataset_view = + dataset_view>; + +template +using host_standard_dataset = + dataset>; + +template +using host_standard_dataset_view = + dataset_view>; + +template +using device_vpq_dataset = + dataset>; + +template +using device_vpq_dataset_view = + dataset_view>; + +template +using host_vpq_dataset = + dataset>; + +template +using host_vpq_dataset_view = + dataset_view>; + +// Maps a dataset view type to its owning (allocating) dataset counterpart. +// Used by serialize/deserialize to type the out_dataset output parameter; +// adding a new dataset type only requires adding a new specialization here. +template +struct owning_dataset_for_view; + +template +struct owning_dataset_for_view> { + using type = device_padded_dataset; +}; + +template +struct owning_dataset_for_view> { + using type = device_standard_dataset; +}; + +template +struct owning_dataset_for_view> { + using type = device_vpq_dataset; +}; + +template +using owning_dataset_for_view_t = typename owning_dataset_for_view::type; + +template +struct is_padded_dataset : std::false_type {}; + +template +struct is_padded_dataset> + : std::true_type {}; + +template +struct is_padded_dataset> + : std::true_type {}; + +template +inline constexpr bool is_padded_dataset_v = is_padded_dataset::value; + +template +struct is_standard_dataset : std::false_type {}; + +template +struct is_standard_dataset> + : std::true_type {}; + +template +struct is_standard_dataset> + : std::true_type {}; + +template +inline constexpr bool is_standard_dataset_v = is_standard_dataset::value; + template struct is_vpq_dataset : std::false_type {}; -template -struct is_vpq_dataset> : std::true_type {}; +template +struct is_vpq_dataset> : std::true_type {}; template inline constexpr bool is_vpq_dataset_v = is_vpq_dataset::value; +// ----------------------------------------------------------------------------- +// Dataset view compile-time classification (replaces runtime std::variant dispatch). +// ----------------------------------------------------------------------------- + +/** Any non-owning dataset view exposing row count and logical dimension. */ +template +concept ann_dataset_view = requires(V const& v) { + { v.n_rows() } -> std::convertible_to; + { v.dim() } -> std::convertible_to; +}; + +enum class dataset_view_kind { + // TODO(removal): Remove `unknown` once all deprecated host_matrix_view / device_matrix_view / + // mdspan overloads are deleted. It exists solely so that overload resolution on the deprecated + // build(host_matrix_view) / build(device_matrix_view) shims does not cause a hard error when + // the compiler evaluates is_host/device_dataset_view_v for a plain mdspan type. + unknown, + empty, + padded, + standard, + vpq_f16, + vpq_f32, +}; + +/** Primary template returns `unknown` so traits safely return `false` for non-dataset-view types. + */ +template +struct dataset_view_kind_of { + static constexpr dataset_view_kind value = dataset_view_kind::unknown; +}; + +template +struct dataset_view_kind_of> { + static constexpr dataset_view_kind value = dataset_view_kind::empty; +}; + +template +struct dataset_view_kind_of> { + static constexpr dataset_view_kind value = dataset_view_kind::padded; +}; + +template +struct dataset_view_kind_of> { + static constexpr dataset_view_kind value = dataset_view_kind::standard; +}; + +template +struct dataset_view_kind_of> { + static_assert(std::is_same_v || std::is_same_v, + "VPQ dataset_view_kind_of expects MathT to be half or float"); + static constexpr dataset_view_kind value = + std::is_same_v ? dataset_view_kind::vpq_f16 : dataset_view_kind::vpq_f32; +}; + +template +using dataset_view_type_t = std::remove_cvref_t; + +/** True when the dataset view accessor is device-accessible. */ +template +struct dataset_view_is_device_accessible : std::false_type {}; + +template +struct dataset_view_is_device_accessible> + : std::bool_constant {}; + +template +inline constexpr bool dataset_view_is_device_accessible_v = + dataset_view_is_device_accessible>::value; + +template +inline constexpr dataset_view_kind dataset_view_kind_v = + dataset_view_kind_of>::value; + +template +inline constexpr bool is_device_empty_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::empty && dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_empty_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::empty && !dataset_view_is_device_accessible_v; + +/** True for any empty dataset view (device or host). */ +template +inline constexpr bool is_empty_dataset_view_v = + is_device_empty_dataset_view_v || is_host_empty_dataset_view_v; + +template +inline constexpr bool is_device_padded_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::padded && dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_padded_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::padded && !dataset_view_is_device_accessible_v; + +/** True for either `device_padded_dataset_view` or `host_padded_dataset_view`. */ +template +inline constexpr bool is_padded_dataset_view_v = + is_device_padded_dataset_view_v || is_host_padded_dataset_view_v; + +template +inline constexpr bool is_device_standard_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::standard && dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_standard_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::standard && !dataset_view_is_device_accessible_v; + +/** True for either `device_standard_dataset_view` or `host_standard_dataset_view`. */ +template +inline constexpr bool is_standard_dataset_view_v = + is_device_standard_dataset_view_v || is_host_standard_dataset_view_v; + +template +inline constexpr bool is_device_vpq_f16_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::vpq_f16 && dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_vpq_f16_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::vpq_f16 && !dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_vpq_f16_dataset_view_v = + is_device_vpq_f16_dataset_view_v || is_host_vpq_f16_dataset_view_v; + +template +inline constexpr bool is_device_vpq_f32_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::vpq_f32 && dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_vpq_f32_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::vpq_f32 && !dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_vpq_f32_dataset_view_v = + is_device_vpq_f32_dataset_view_v || is_host_vpq_f32_dataset_view_v; + +template +inline constexpr bool is_device_vpq_dataset_view_v = + is_device_vpq_f16_dataset_view_v || is_device_vpq_f32_dataset_view_v; + +template +inline constexpr bool is_host_vpq_dataset_view_v = + is_host_vpq_f16_dataset_view_v || is_host_vpq_f32_dataset_view_v; + +template +inline constexpr bool is_vpq_dataset_view_v = + is_device_vpq_dataset_view_v || is_host_vpq_dataset_view_v; + +/** True for any device-resident dataset view. */ +template +inline constexpr bool is_device_dataset_view_v = + dataset_view_kind_v != dataset_view_kind::unknown && dataset_view_is_device_accessible_v; + +/** True for any host-resident dataset view. */ +template +inline constexpr bool is_host_dataset_view_v = + dataset_view_kind_v != dataset_view_kind::unknown && !dataset_view_is_device_accessible_v; + +/** + * True when a host view `H` and device view `D` represent the same storage kind and differ + * only in residency (host vs. device). Used to constrain `attach_device_dataset_on_host_index`. + */ +template +inline constexpr bool compatible_host_device_dataset_views_v = + is_host_dataset_view_v && is_device_dataset_view_v && + (dataset_view_kind_v == dataset_view_kind_v); + +/** + * Generic accessor retargeting while preserving the dataset tag/layout and value/index types: + * `dataset -> dataset` + * `dataset_view -> dataset_view` + */ +template +struct with_accessor; + +template +struct with_accessor, NewAccessor> { + using type = dataset; +}; + +template +struct with_accessor, NewAccessor> { + using type = dataset_view; +}; + +template +using with_accessor_t = + typename with_accessor, NewAccessor>::type; + +/** Map any host accessor to its device counterpart (same payload policy). */ +template +struct to_device_accessor { + using type = Accessor; +}; + +template +struct to_device_accessor> { + using type = detail::device_view_accessor; +}; + +template +struct to_device_accessor> { + using type = detail::device_owning_accessor; +}; + +template +using to_device_accessor_t = typename to_device_accessor::type; + +/** Maps a host dataset view type to its device-resident counterpart. */ +template +struct device_counterpart; + +template +struct device_counterpart> { + using type = with_accessor_t, + to_device_accessor_t>; +}; + +template +using device_counterpart_t = typename device_counterpart>::type; + +/** True for device padded or standard views accepted by dense graph build (VPQ excluded). */ +template +inline constexpr bool is_dense_row_major_device_dataset_view_v = + is_device_padded_dataset_view_v || is_device_standard_dataset_view_v; + +/** True for host or device padded/standard views (iterative graph build; VPQ excluded). */ +template +inline constexpr bool is_dense_row_major_dataset_view_v = + is_padded_dataset_view_v || is_standard_dataset_view_v; + +/** Element type `T` for `cagra::build(res, params, dataset_view)` (deduced, not a template arg). */ +template +struct cagra_view_element_type; + +template +struct cagra_view_element_type> { + using type = DataT; +}; + +template +struct cagra_view_element_type> { + using type = DataT; +}; + +template +struct cagra_view_element_type> { + using type = DataT; +}; + +template +struct cagra_view_element_type> { + using type = DataT; +}; + +template +struct cagra_view_element_type> { + using type = MathT; +}; + +template +using cagra_view_element_type_t = typename cagra_view_element_type>::type; + +// ----------------------------------------------------------------------------- +// CAGRA row width in elements (same for make_device_padded_dataset* and index layout checks). +// ----------------------------------------------------------------------------- + +/** + * @brief Required row width in elements for CAGRA: minimum leading dimension (LDA) per row for the + * default per-row byte alignment (16 bytes, combined with `sizeof` element type), given + * `logical_columns` feature columns. + */ +[[nodiscard]] inline uint32_t cagra_required_row_width(uint32_t logical_columns, + std::size_t sizeof_value, + uint32_t align_bytes = 16) +{ + return static_cast( + raft::round_up_safe(static_cast(logical_columns) * sizeof_value, + std::lcm(align_bytes, static_cast(sizeof_value))) / + sizeof_value); +} + +template +[[nodiscard]] inline uint32_t cagra_required_row_width(uint32_t logical_columns, + uint32_t align_bytes = 16) +{ + return cagra_required_row_width(logical_columns, sizeof(ValueT), align_bytes); +} + +/** Actual row width in elements (leading dimension) of a 2D row-major matrix view. */ +template +[[nodiscard]] inline uint32_t matrix_actual_row_width(raft::device_matrix_view m) +{ + return m.stride(0) > 0 ? static_cast(m.stride(0)) : static_cast(m.extent(1)); +} + +template +[[nodiscard]] inline uint32_t matrix_actual_row_width(raft::host_matrix_view m) +{ + return m.stride(0) > 0 ? static_cast(m.stride(0)) : static_cast(m.extent(1)); +} + +/** + * @brief True if the matrix's row width in elements matches `cagra_required_row_width` for + * `m.extent(1)` and element type `T` (CAGRA row layout is satisfied for this view). + */ +template +[[nodiscard]] inline bool matrix_row_width_matches_cagra_required( + raft::device_matrix_view m, uint32_t align_bytes = 16) +{ + using value_type = std::remove_const_t; + const uint32_t need = + cagra_required_row_width(static_cast(m.extent(1)), align_bytes); + return matrix_actual_row_width(m) == need; +} + +template +[[nodiscard]] inline bool matrix_row_width_matches_cagra_required(raft::host_matrix_view m, + uint32_t align_bytes = 16) +{ + using value_type = std::remove_const_t; + const uint32_t need = + cagra_required_row_width(static_cast(m.extent(1)), align_bytes); + return matrix_actual_row_width(m) == need; +} + +namespace detail { + +template +[[nodiscard]] inline uint32_t mdspan_row_stride_elements(SrcT const& src) +{ + return src.stride(0) > 0 ? static_cast(src.stride(0)) + : static_cast(src.extent(1)); +} + +template +[[nodiscard]] inline ValueT* expect_device_accessible_data_handle(SrcT const& src, + char const* error_msg) +{ + cudaPointerAttributes ptr_attrs; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src.data_handle())); + auto* device_ptr = reinterpret_cast(ptr_attrs.devicePointer); + RAFT_EXPECTS(device_ptr != nullptr, "%s", error_msg); + return device_ptr; +} + +template +[[nodiscard]] inline ViewT make_device_dense_row_major_view_from_src(SrcT const& src, + uint32_t logical_dim) +{ + auto* device_ptr = expect_device_accessible_data_handle( + src, "make_device_*_dataset_view: source must be device-accessible."); + auto v = raft::make_device_matrix_view( + device_ptr, src.extent(0), static_cast(mdspan_row_stride_elements(src))); + return ViewT(v, logical_dim); +} + +template +[[nodiscard]] inline ViewT make_host_dense_row_major_view_from_src(SrcT const& src, + uint32_t logical_dim) +{ + RAFT_EXPECTS(raft::get_device_for_address(src.data_handle()) == -1, + "make_host_*_dataset_view: source must be host-accessible."); + auto v = raft::make_host_matrix_view(const_cast(src.data_handle()), + src.extent(0), + static_cast(mdspan_row_stride_elements(src))); + return ViewT(v, logical_dim); +} + +template +auto make_device_dense_row_major_dataset_from_src(raft::resources const& res, + SrcT const& src, + uint32_t logical_dim, + uint32_t target_stride, + char const* view_factory_name) + -> std::unique_ptr +{ + uint32_t const src_stride = mdspan_row_stride_elements(src); + RAFT_EXPECTS(logical_dim <= target_stride, + "logical dim (%u) must not exceed row stride (%u).", + static_cast(logical_dim), + static_cast(target_stride)); + RAFT_EXPECTS(static_cast(src.extent(1)) <= target_stride, + "Source row length must not exceed required stride."); + cudaPointerAttributes ptr_attrs; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src.data_handle())); + bool const device_src = + (ptr_attrs.type == cudaMemoryTypeDevice) || (ptr_attrs.type == cudaMemoryTypeManaged); + if (device_src && src_stride == target_stride) { + RAFT_EXPECTS(false, + "source is device and stride is already correct. " + "Use %s() to get a view instead.", + view_factory_name); + } + auto out_array = raft::make_device_matrix(res, src.extent(0), target_stride); + RAFT_CUDA_TRY(cudaMemsetAsync(out_array.data_handle(), + 0, + out_array.size() * sizeof(ValueT), + raft::resource::get_cuda_stream(res))); + raft::copy_matrix(out_array.data_handle(), + target_stride, + src.data_handle(), + src_stride, + logical_dim, + src.extent(0), + raft::resource::get_cuda_stream(res)); + return std::make_unique(std::move(out_array), logical_dim); +} + +template +auto make_host_dense_row_major_dataset_from_src(raft::resources const& res, + SrcT const& src, + uint32_t logical_dim, + uint32_t target_stride, + char const* device_factory_name, + char const* view_factory_name) + -> std::unique_ptr +{ + uint32_t const src_stride = mdspan_row_stride_elements(src); + RAFT_EXPECTS(raft::get_device_for_address(src.data_handle()) == -1, + "source must be host-accessible. Use %s() for device sources.", + device_factory_name); + RAFT_EXPECTS(logical_dim <= target_stride, + "logical dim (%u) must not exceed row stride (%u).", + static_cast(logical_dim), + static_cast(target_stride)); + if (src_stride == target_stride) { + RAFT_EXPECTS(false, + "source stride is already correct. Use %s() to get a view instead.", + view_factory_name); + } + RAFT_EXPECTS(static_cast(src.extent(1)) <= target_stride, + "Source row length must not exceed required stride."); + auto out_array = raft::make_host_matrix(src.extent(0), target_stride); + std::memset(out_array.data_handle(), 0, out_array.size() * sizeof(ValueT)); + raft::copy_matrix(out_array.data_handle(), + target_stride, + src.data_handle(), + src_stride, + logical_dim, + src.extent(0), + raft::resource::get_cuda_stream(res)); + return std::make_unique(std::move(out_array), logical_dim); +} + +} // namespace detail + +template +auto make_device_padded_dataset_view(const raft::resources& res, + SrcT const& src, + uint32_t align_bytes = 16) + -> device_padded_dataset_view +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + uint32_t required_stride = + cagra_required_row_width(static_cast(src.extent(1)), align_bytes); + RAFT_EXPECTS( + detail::mdspan_row_stride_elements(src) == required_stride, + "make_device_padded_dataset_view: stride is incorrect (required stride for alignment). " + "Use make_device_padded_dataset() to get an owning padded copy."); + return detail::make_device_dense_row_major_view_from_src< + value_type, + index_type, + device_padded_dataset_view>(src, static_cast(src.extent(1))); +} + +template +auto make_device_padded_dataset(const raft::resources& res, + SrcT const& src, + uint32_t align_bytes = 16) + -> std::unique_ptr> +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + uint32_t const logical_dim = static_cast(src.extent(1)); + uint32_t const required_stride = cagra_required_row_width(logical_dim, align_bytes); + return detail::make_device_dense_row_major_dataset_from_src< + device_padded_dataset, + value_type, + index_type>(res, src, logical_dim, required_stride, "make_device_padded_dataset_view"); +} + +template +auto make_host_padded_dataset_view(SrcT const& src, uint32_t align_bytes = 16) + -> host_padded_dataset_view +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + uint32_t required_stride = + cagra_required_row_width(static_cast(src.extent(1)), align_bytes); + RAFT_EXPECTS( + detail::mdspan_row_stride_elements(src) == required_stride, + "make_host_padded_dataset_view: stride is incorrect (required stride for alignment). " + "Use make_host_padded_dataset() to get an owning padded copy."); + return detail::make_host_dense_row_major_view_from_src< + value_type, + index_type, + host_padded_dataset_view>(src, static_cast(src.extent(1))); +} + +template +auto make_host_padded_dataset(const raft::resources& res, + SrcT const& src, + uint32_t align_bytes = 16) + -> std::unique_ptr> +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + uint32_t const logical_dim = static_cast(src.extent(1)); + uint32_t const required_stride = cagra_required_row_width(logical_dim, align_bytes); + return detail::make_host_dense_row_major_dataset_from_src< + host_padded_dataset, + value_type, + index_type>(res, + src, + logical_dim, + required_stride, + "make_device_padded_dataset", + "make_host_padded_dataset_view"); +} + +template +auto make_device_standard_dataset_view(SrcT const& src) + -> device_standard_dataset_view +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + return detail::make_device_dense_row_major_view_from_src< + value_type, + index_type, + device_standard_dataset_view>(src, + static_cast(src.extent(1))); +} + +/** + * @brief Create an owning device standard dataset with explicit row layout. + * + * Internal use only: the sole call site today is + * `cuvs::neighbors::detail::deserialize_standard()` in `dataset_serialize.hpp`, which must pass + * wire-format `(logical_dim, stride)` because the deserialized host buffer is tight `[n_rows x + * dim]` while the on-disk stride may be larger. Do not call from user code; prefer + * `make_device_standard_dataset_view()` when wrapping existing correctly-strided storage. + * + * Potential future call sites if an owning copy with explicit stride is needed: + * - C API dataset upload (mirroring `make_device_padded_dataset` in `c/src/neighbors/cagra.cpp`) + * - `tiered_index` / composite index paths that materialize standard-layout device storage + * - Multigpu (MG) index build or merge when rehydrating a strided dataset from host fragments + */ +template +auto make_device_standard_dataset(const raft::resources& res, + SrcT const& src, + uint32_t logical_dim, + uint32_t target_stride) + -> std::unique_ptr> +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + return detail::make_device_dense_row_major_dataset_from_src< + device_standard_dataset, + value_type, + index_type>(res, src, logical_dim, target_stride, "make_device_standard_dataset_view"); +} + +template +auto make_host_standard_dataset_view(SrcT const& src) + -> host_standard_dataset_view +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + return detail::make_host_dense_row_major_view_from_src< + value_type, + index_type, + host_standard_dataset_view>(src, static_cast(src.extent(1))); +} + namespace filtering { /** @@ -896,11 +1731,22 @@ using namespace raft; template struct iface { - iface() : mutex_(std::make_shared()) {} + iface() + : cagra_owned_padded_dataset_(nullptr), + cagra_owned_standard_dataset_(nullptr), + mutex_(std::make_shared()) + { + } const IdxT size() const { return index_.value().size(); } std::optional index_; + /** Used by CAGRA when deserializing an index that contains a dataset; keeps it alive for the + * view. */ + std::unique_ptr> cagra_owned_padded_dataset_; + /** Used by CAGRA standard-layout paths to keep deserialized/attached dataset views alive. */ + std::unique_ptr> + cagra_owned_standard_dataset_; std::shared_ptr mutex_; }; @@ -1005,6 +1851,7 @@ struct mg_search_params : public Upstream { template struct mg_index { + mg_index(const raft::resources& clique); mg_index(const raft::resources& clique, distribution_mode mode); mg_index(const raft::resources& clique, const std::string& filename); diff --git a/cpp/include/cuvs/neighbors/composite/index.hpp b/cpp/include/cuvs/neighbors/composite/index.hpp index d7970a5cd6..abfbb6cbd4 100644 --- a/cpp/include/cuvs/neighbors/composite/index.hpp +++ b/cpp/include/cuvs/neighbors/composite/index.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -48,7 +48,8 @@ class CUVS_EXPORT composite_index { using out_index_type = OutputIdxT; using matrix_index_type = int64_t; - explicit composite_index(std::vector*> children) + explicit composite_index( + std::vector*> children) : children_(std::move(children)) { } @@ -91,7 +92,7 @@ class CUVS_EXPORT composite_index { } private: - std::vector*> children_; + std::vector*> children_; }; } // namespace composite diff --git a/cpp/include/cuvs/neighbors/hnsw.hpp b/cpp/include/cuvs/neighbors/hnsw.hpp index fb726fed71..a7731a65ea 100644 --- a/cpp/include/cuvs/neighbors/hnsw.hpp +++ b/cpp/include/cuvs/neighbors/hnsw.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -7,12 +7,8 @@ #pragma once -#include "common.hpp" - #include #include - -#include "cagra.hpp" #include #include @@ -474,7 +470,7 @@ std::unique_ptr> build( std::unique_ptr> from_cagra( raft::resources const& res, const index_params& params, - const cuvs::neighbors::cagra::index& cagra_index, + const cuvs::neighbors::cagra::device_padded_index& cagra_index, std::optional> dataset = std::nullopt); @@ -510,7 +506,7 @@ std::unique_ptr> from_cagra( std::unique_ptr> from_cagra( raft::resources const& res, const index_params& params, - const cuvs::neighbors::cagra::index& cagra_index, + const cuvs::neighbors::cagra::device_padded_index& cagra_index, std::optional> dataset = std::nullopt); @@ -546,7 +542,7 @@ std::unique_ptr> from_cagra( std::unique_ptr> from_cagra( raft::resources const& res, const index_params& params, - const cuvs::neighbors::cagra::index& cagra_index, + const cuvs::neighbors::cagra::device_padded_index& cagra_index, std::optional> dataset = std::nullopt); @@ -582,7 +578,105 @@ std::unique_ptr> from_cagra( std::unique_ptr> from_cagra( raft::resources const& res, const index_params& params, - const cuvs::neighbors::cagra::index& cagra_index, + const cuvs::neighbors::cagra::device_padded_index& cagra_index, + std::optional> dataset = + std::nullopt); + +/** + * @brief Construct an hnswlib index from a device-standard CAGRA index. + * + * When the index has an attached device dataset view, `dataset` may be omitted. Otherwise pass a + * host matrix with the vectors (same contract as `device_padded_index`). + */ +std::unique_ptr> from_cagra( + raft::resources const& res, + const index_params& params, + const cuvs::neighbors::cagra::device_standard_index& cagra_index, + std::optional> dataset = + std::nullopt); + +std::unique_ptr> from_cagra( + raft::resources const& res, + const index_params& params, + const cuvs::neighbors::cagra::device_standard_index& cagra_index, + std::optional> dataset = + std::nullopt); + +std::unique_ptr> from_cagra( + raft::resources const& res, + const index_params& params, + const cuvs::neighbors::cagra::device_standard_index& cagra_index, + std::optional> dataset = + std::nullopt); + +std::unique_ptr> from_cagra( + raft::resources const& res, + const index_params& params, + const cuvs::neighbors::cagra::device_standard_index& cagra_index, + std::optional> dataset = + std::nullopt); + +/** + * @brief Construct an hnswlib index from a host-built CAGRA index. + * Requires `dataset` for in-memory indices — host builds do not store vectors in the index. + */ +std::unique_ptr> from_cagra( + raft::resources const& res, + const index_params& params, + const cuvs::neighbors::cagra::host_padded_index& cagra_index, + std::optional> dataset = + std::nullopt); + +std::unique_ptr> from_cagra( + raft::resources const& res, + const index_params& params, + const cuvs::neighbors::cagra::host_padded_index& cagra_index, + std::optional> dataset = + std::nullopt); + +std::unique_ptr> from_cagra( + raft::resources const& res, + const index_params& params, + const cuvs::neighbors::cagra::host_padded_index& cagra_index, + std::optional> dataset = + std::nullopt); + +std::unique_ptr> from_cagra( + raft::resources const& res, + const index_params& params, + const cuvs::neighbors::cagra::host_padded_index& cagra_index, + std::optional> dataset = + std::nullopt); + +/** + * @brief Construct an hnswlib index from a host-built CAGRA index (standard dataset layout). + * Requires `dataset` for in-memory indices — host builds do not store vectors in the index. + */ +std::unique_ptr> from_cagra( + raft::resources const& res, + const index_params& params, + const cuvs::neighbors::cagra::host_standard_index& cagra_index, + std::optional> dataset = + std::nullopt); + +std::unique_ptr> from_cagra( + raft::resources const& res, + const index_params& params, + const cuvs::neighbors::cagra::host_standard_index& cagra_index, + std::optional> dataset = + std::nullopt); + +std::unique_ptr> from_cagra( + raft::resources const& res, + const index_params& params, + const cuvs::neighbors::cagra::host_standard_index& cagra_index, + std::optional> dataset = + std::nullopt); + +std::unique_ptr> from_cagra( + raft::resources const& res, + const index_params& params, + const cuvs::neighbors::cagra::host_standard_index& cagra_index, std::optional> dataset = std::nullopt); diff --git a/cpp/include/cuvs/neighbors/ivf_flat.hpp b/cpp/include/cuvs/neighbors/ivf_flat.hpp index d2e0015498..629f7a9e39 100644 --- a/cpp/include/cuvs/neighbors/ivf_flat.hpp +++ b/cpp/include/cuvs/neighbors/ivf_flat.hpp @@ -1,11 +1,10 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include "common.hpp" #include #include #include diff --git a/cpp/include/cuvs/neighbors/ivf_sq.hpp b/cpp/include/cuvs/neighbors/ivf_sq.hpp index 6ac765213c..6d1dff6aba 100644 --- a/cpp/include/cuvs/neighbors/ivf_sq.hpp +++ b/cpp/include/cuvs/neighbors/ivf_sq.hpp @@ -1,11 +1,10 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include "common.hpp" #include #include #include diff --git a/cpp/include/cuvs/neighbors/tiered_index.hpp b/cpp/include/cuvs/neighbors/tiered_index.hpp index 8d0e18281c..7801c6461a 100644 --- a/cpp/include/cuvs/neighbors/tiered_index.hpp +++ b/cpp/include/cuvs/neighbors/tiered_index.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -87,7 +87,24 @@ struct index_params : upstream_index_params_type { auto build(raft::resources const& res, const index_params& index_params, raft::device_matrix_view dataset) - -> tiered_index::index>; + -> tiered_index::index>; + +/** + * @brief Attach caller-managed padded dataset for searching a standard CAGRA tiered index. + * + * This is an explicit contract step: callers create/manage padded storage and attach it + * before calling `search` on the returned padded tiered index. + */ +auto attach_padded_dataset_for_search( + raft::resources const& res, + const index>& idx, + cuvs::neighbors::device_padded_dataset_view padded_dataset) + -> index>; + +auto build(raft::resources const& res, + const index_params& index_params, + cuvs::neighbors::device_padded_dataset_view dataset) + -> tiered_index::index>; /** @copydoc build */ auto build(raft::resources const& res, @@ -121,7 +138,12 @@ auto build(raft::resources const& res, */ void extend(raft::resources const& res, raft::device_matrix_view new_vectors, - tiered_index::index>* idx); + tiered_index::index>* idx); + +/** @copydoc extend */ +void extend(raft::resources const& res, + raft::device_matrix_view new_vectors, + tiered_index::index>* idx); /** @copydoc extend */ void extend(raft::resources const& res, @@ -141,7 +163,12 @@ void extend(raft::resources const& res, * @param[in] res * @param[inout] idx */ -void compact(raft::resources const& res, tiered_index::index>* idx); +void compact(raft::resources const& res, + tiered_index::index>* idx); + +/** @copydoc compact */ +void compact(raft::resources const& res, + tiered_index::index>* idx); /** @copydoc compact */ void compact(raft::resources const& res, tiered_index::index>* idx); @@ -166,7 +193,17 @@ void compact(raft::resources const& res, */ void search(raft::resources const& res, const cagra::search_params& search_params, - const tiered_index::index>& index, + const tiered_index::index>& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter = + cuvs::neighbors::filtering::none_sample_filter{}); + +/** @copydoc search */ +void search(raft::resources const& res, + const cagra::search_params& search_params, + const tiered_index::index>& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, raft::device_matrix_view distances, @@ -205,10 +242,18 @@ void search(raft::resources const& res, * * @return A new tiered index containing the merged indices */ -auto merge(raft::resources const& res, - const index_params& index_params, - const std::vector>*>& indices) - -> tiered_index::index>; +auto merge( + raft::resources const& res, + const index_params& index_params, + const std::vector>*>& indices) + -> tiered_index::index>; + +/** @copydoc merge */ +auto merge( + raft::resources const& res, + const index_params& index_params, + const std::vector>*>& indices) + -> tiered_index::index>; /** @copydoc merge */ auto merge(raft::resources const& res, diff --git a/cpp/include/cuvs/neighbors/vamana.hpp b/cpp/include/cuvs/neighbors/vamana.hpp index 645adc5c5d..38003b780e 100644 --- a/cpp/include/cuvs/neighbors/vamana.hpp +++ b/cpp/include/cuvs/neighbors/vamana.hpp @@ -1,11 +1,10 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include "common.hpp" #include #include #include @@ -19,7 +18,6 @@ #include #include -#include namespace CUVS_EXPORT cuvs { namespace neighbors { @@ -116,22 +114,27 @@ struct index : cuvs::neighbors::index { /** Total length of the index (number of vectors). */ [[nodiscard]] constexpr inline auto size() const noexcept -> IdxT { - auto data_rows = dataset_->n_rows(); + auto data_rows = dataset_.has_value() ? dataset_->n_rows() : IdxT{0}; return data_rows > 0 ? data_rows : graph_view_.extent(0); } /** Dimensionality of the data. */ - [[nodiscard]] constexpr inline auto dim() const noexcept -> uint32_t { return dataset_->dim(); } + [[nodiscard]] constexpr inline auto dim() const noexcept -> uint32_t + { + return dataset_.has_value() ? dataset_->dim() : 0u; + } /** Graph degree */ [[nodiscard]] constexpr inline auto graph_degree() const noexcept -> uint32_t { return graph_view_.extent(1); } - /** Dataset [size, dim] */ - [[nodiscard]] inline auto data() const noexcept -> const cuvs::neighbors::dataset& + /** Non-owning dataset view stored by the index (full-precision vectors may live in + * `full_precision_storage_`). */ + [[nodiscard]] inline auto data() const noexcept + -> const cuvs::neighbors::device_padded_dataset_view& { - return *dataset_; + return dataset_.value(); } /** Quantized dataset [size, codes_rowlen] */ @@ -166,7 +169,8 @@ struct index : cuvs::neighbors::index { : cuvs::neighbors::index(), metric_(metric), graph_(raft::make_device_matrix(res, 0, 0)), - dataset_(new cuvs::neighbors::empty_dataset(0)), + full_precision_storage_(), + dataset_{std::nullopt}, quantized_dataset_(raft::make_device_matrix(res, 0, 0)) { } @@ -184,12 +188,30 @@ struct index : cuvs::neighbors::index { : cuvs::neighbors::index(), metric_(metric), graph_(raft::make_device_matrix(res, 0, 0)), - dataset_(make_aligned_dataset(res, dataset, 16)), + full_precision_storage_(), + dataset_{}, quantized_dataset_(raft::make_device_matrix(res, 0, 0)), medoid_id_(medoid_id) { RAFT_EXPECTS(dataset.extent(0) == vamana_graph.extent(0), "Dataset and vamana_graph must have equal number of rows"); + + const bool on_device = raft::get_device_for_address(dataset.data_handle()) >= 0; + bool use_padded_view = false; + if (on_device) { + const int64_t row_stride = + dataset.stride(0) > 0 ? static_cast(dataset.stride(0)) : dataset.extent(1); + auto d_m = raft::make_device_matrix_view( + dataset.data_handle(), dataset.extent(0), row_stride); + use_padded_view = cuvs::neighbors::matrix_row_width_matches_cagra_required(d_m); + } + + if (use_padded_view) { + dataset_ = cuvs::neighbors::make_device_padded_dataset_view(res, dataset); + } else { + full_precision_storage_ = cuvs::neighbors::make_device_padded_dataset(res, dataset); + dataset_ = full_precision_storage_->as_dataset_view(); + } update_graph(res, vamana_graph); raft::resource::sync_stream(res); @@ -264,7 +286,9 @@ struct index : cuvs::neighbors::index { cuvs::distance::DistanceType metric_; raft::device_matrix graph_; raft::device_matrix_view graph_view_; - std::unique_ptr> dataset_; + /** Owns CAGRA-padded full-precision device storage for the index dataset view. */ + std::unique_ptr> full_precision_storage_; + std::optional> dataset_; raft::device_matrix quantized_dataset_; IdxT medoid_id_; }; diff --git a/cpp/include/cuvs/preprocessing/quantize/pq.hpp b/cpp/include/cuvs/preprocessing/quantize/pq.hpp index bdbe77bac6..fe40f422f2 100644 --- a/cpp/include/cuvs/preprocessing/quantize/pq.hpp +++ b/cpp/include/cuvs/preprocessing/quantize/pq.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -11,6 +11,7 @@ #include #include +#include #include #include @@ -147,7 +148,7 @@ struct quantizer { /** Parameters used to build this quantizer. */ params params_quantizer; /** VPQ codebooks produced during training. */ - cuvs::neighbors::vpq_dataset vpq_codebooks; + cuvs::neighbors::device_vpq_dataset vpq_codebooks; }; /** @@ -243,6 +244,61 @@ void inverse_transform( raft::device_matrix_view out, std::optional> vq_labels = std::nullopt); +namespace detail { + +template +[[nodiscard]] cuvs::neighbors::device_vpq_dataset vpq_train_from_device_rows( + raft::resources const& res, + cuvs::neighbors::vpq_params const& params, + T const* src_ptr, + int64_t n_rows, + int64_t dim, + int64_t stride); + +} // namespace detail + +/** + * @brief Train VPQ storage (codebooks + encoded rows) from a device row-major mdspan/matrix. + * + * Accepts any device-accessible mdspan with `value_type`, `extent`, `stride`, and `data_handle` + * (same pattern as `cuvs::neighbors::make_device_padded_dataset`). Row-major tight storage (logical + * stride equals dimension) is passed through to training without an extra pack copy; wider row + * pitch triggers a contiguous dense copy first. Empty sources are rejected. + * + * Typical **CAGRA** usage: build the graph on dense vectors, then attach VPQ for search (metric + * must remain `L2Expanded` for this path). Train VPQ from the same CAGRA-padded device layout you + * used for graph build, keep the `device_vpq_dataset` alive, and call `index::update_dataset` with + * a non-owning view. + * + * @code{.cpp} + * #include + * #include + * + * // `idx` is a `cagra::index` with graph built on dense rows. + * // `padded` is a `device_padded_dataset_view` view of those same rows. + * cuvs::neighbors::vpq_params vpq_params{}; + * auto vpq = cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, vpq_params, padded.view()); + * idx.update_dataset(res, vpq.as_dataset_view()); + * @endcode + */ +template +[[nodiscard]] auto make_vpq_dataset(raft::resources const& res, + cuvs::neighbors::vpq_params const& params, + SrcT const& src) + -> cuvs::neighbors::device_vpq_dataset +{ + using T = typename SrcT::value_type; + RAFT_EXPECTS(src.extent(0) > 0, "make_vpq_dataset: dataset is empty"); + cudaPointerAttributes ptr_attrs; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src.data_handle())); + auto const* device_ptr = reinterpret_cast(ptr_attrs.devicePointer); + RAFT_EXPECTS(device_ptr != nullptr, "make_vpq_dataset: source must be device-accessible."); + const int64_t n_rows = src.extent(0); + const int64_t dim = src.extent(1); + const int64_t stride = src.stride(0) > 0 ? src.stride(0) : dim; + return detail::vpq_train_from_device_rows(res, params, device_ptr, n_rows, dim, stride); +} + /** @} */ // end of group product } // namespace pq diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index ee87c2c0ab..30c9c592ec 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -23,24 +23,39 @@ #include #include - #include #include +#include +#include +#include namespace cuvs::neighbors::cagra { -// Member function implementations for cagra::index -template -void index::compute_dataset_norms_(raft::resources const& res) +template +CUVS_EXPORT void index::compute_dataset_norms_(raft::resources const& res) { - // Get the dataset view - auto dataset_view = this->dataset(); + // raft::linalg::reduce wants row-major with leading dim = row pitch in elements. + // Skip norm precomputation for VPQ/empty/non-dense views; CosineExpanded with VPQ is handled + // (or rejected) on the search path. + namespace nb = cuvs::neighbors; + bool skip_norms = false; + std::optional> rm_dataset; + + if constexpr (nb::is_padded_dataset_view_v || + nb::is_standard_dataset_view_v) { + rm_dataset = dataset_.view(); + } else if constexpr (nb::is_vpq_dataset_view_v) { + skip_norms = true; + } + + if (skip_norms) { return; } + if (!rm_dataset.has_value()) { return; } // Allocate norms vector if not already allocated - if (!dataset_norms_.has_value() || dataset_norms_->extent(0) != dataset_view.extent(0)) { + if (!dataset_norms_.has_value() || dataset_norms_->extent(0) != rm_dataset->extent(0)) { dataset_norms_.reset(); - dataset_norms_ = raft::make_device_vector(res, dataset_view.extent(0)); + dataset_norms_ = raft::make_device_vector(res, rm_dataset->extent(0)); } constexpr float kScale = cuvs::spatial::knn::detail::utils::config::kDivisor / @@ -49,16 +64,14 @@ void index::compute_dataset_norms_(raft::resources const& res) // first scale the dataset and then compute norms auto scaled_sq_op = raft::compose_op( raft::sq_op{}, raft::div_const_op{float(kScale)}, raft::cast_op()); - raft::linalg::reduce( - res, - raft::make_device_matrix_view( - dataset_view.data_handle(), dataset_view.extent(0), dataset_view.stride(0)), - dataset_norms_->view(), - (float)0, - false, - scaled_sq_op, - raft::add_op(), - raft::sqrt_op{}); + raft::linalg::reduce(res, + *rm_dataset, + dataset_norms_->view(), + (float)0, + false, + scaled_sq_op, + raft::add_op(), + raft::sqrt_op{}); } /** @@ -92,8 +105,8 @@ void index::compute_dataset_norms_(raft::resources const& res) * auto optimized_gaph = raft::make_host_matrix(dataset.extent(0), 64); * cagra::optimize(res, dataset, knn_graph.view(), optimized_graph.view()); * // Construct an index from dataset and optimized knn_graph - * auto index = cagra::index(res, build_params.metric(), dataset, - * optimized_graph.view()); + * auto index = cagra::index>( + * res, build_params.metric(), dataset, optimized_graph.view()); * @endcode * * @tparam DataT data element type @@ -150,8 +163,8 @@ void build_knn_graph( * auto optimized_gaph = raft::make_host_matrix(dataset.extent(0), 64); * cagra::optimize(res, dataset, nn_descent_index.graph.view(), optimized_graph.view()); * // Construct an index from dataset and optimized knn_graph - * auto index = cagra::index(res, build_params.metric(), dataset, - * optimized_graph.view()); + * auto index = cagra::device_padded_index(res, build_params.metric(), dataset, + * optimized_graph.view()); * @endcode * * @tparam DataT data element type @@ -196,8 +209,8 @@ void build_knn_graph( * // optimize graph * cagra::optimize(res, dataset, knn_graph.view(), optimized_graph.view()); * // Construct an index from dataset and optimized knn_graph - * auto index = cagra::index(res, build_params.metric(), dataset, - * optimized_graph.view()); + * auto index = cagra::index>( + * res, build_params.metric(), dataset, optimized_graph.view()); * @endcode * * @tparam DataT type of the data in the source dataset @@ -264,25 +277,45 @@ void optimize( detail::optimize(res, knn_graph, new_graph, guarantee_connectivity); } -template , raft::memory_type::host>> -index build( - raft::resources const& res, - const index_params& params, - raft::mdspan, raft::row_major, Accessor> dataset) +/** + * @brief Build the index from a `dataset_view` (device padded/standard, device VPQ, or host + * padded/standard). + * + * When `index_params.attach_dataset_on_build = true` (the default) **and the input is a device + * view**, the `dataset` view is stored in the returned index as a non-owning view — no copy is + * made. The caller must keep the underlying storage alive for the lifetime of the index. + * + * For host views, `attach_dataset_on_build` is ignored — the host_padded_index cannot be + * searched; call `attach_device_dataset_on_host_index` to get a search-ready device index. + */ +template + requires(!cuvs::neighbors::is_empty_dataset_view_v && + (cuvs::neighbors::is_device_dataset_view_v || + cuvs::neighbors::is_host_dataset_view_v)) +auto build(raft::resources const& res, const index_params& params, DatasetViewT const& dataset) + -> cuvs::neighbors::cagra::cagra_index_t { - // Check if ACE dispatch is requested via graph_build_params - if (std::holds_alternative(params.graph_build_params)) { - // ACE expects the dataset to be on host due to the large dataset size - RAFT_EXPECTS(raft::get_device_for_address(dataset.data_handle()) == -1, - "ACE: Dataset must be on host for ACE build"); - auto dataset_view = raft::make_host_matrix_view( - dataset.data_handle(), dataset.extent(0), dataset.extent(1)); - return cuvs::neighbors::cagra::detail::build_ace(res, params, dataset_view); + using T = cuvs::neighbors::cagra_view_element_type_t; + using IdxT = uint32_t; + + // Device path: build graph, optionally attach dataset view. + // attach_dataset_on_build is only meaningful for device builds — a host_padded_index cannot + // be searched regardless; the caller must call attach_device_dataset_on_host_index. + if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { + RAFT_FAIL("cagra::build: VPQ-compressed dataset cannot be used for dense graph construction."); + } else if constexpr (cuvs::neighbors::is_dense_row_major_device_dataset_view_v) { + auto idx = cuvs::neighbors::cagra::detail::build_from_device_matrix( + res, params, dataset); + if (params.attach_dataset_on_build) { idx.update_dataset(res, dataset); } + return idx; + } else { + if (std::holds_alternative(params.graph_build_params)) { + return cuvs::neighbors::cagra::detail::build_ace( + res, params, dataset.view()); + } + return cuvs::neighbors::cagra::detail::build_from_host_matrix( + res, params, dataset); } - return cuvs::neighbors::cagra::detail::build(res, params, dataset); } /** @@ -324,10 +357,14 @@ index build( * k] * @param[in] sample_filter a device filter function that greenlights samples for a given query */ -template +template void search_with_filtering(raft::resources const& res, const search_params& params, - const index& idx, + const index& idx, raft::device_matrix_view queries, raft::device_matrix_view neighbors, raft::device_matrix_view distances, @@ -342,14 +379,17 @@ void search_with_filtering(raft::resources const& res, RAFT_EXPECTS(queries.extent(1) == idx.dim(), "Number of query dimensions should equal number of dimensions in the index."); - return cagra::detail::search_main( + return cagra::detail::search_main( res, params, idx, queries, neighbors, distances, sample_filter); } -template +template void search(raft::resources const& res, const search_params& params, - const index& idx, + const index& idx, raft::device_matrix_view queries, raft::device_matrix_view neighbors, raft::device_matrix_view distances, @@ -361,7 +401,7 @@ void search(raft::resources const& res, search_params params_copy = params; if (params.filtering_rate < 0.0) { params_copy.filtering_rate = 0.0; } auto sample_filter_copy = sample_filter; - return search_with_filtering( + return search_with_filtering( res, params_copy, idx, queries, neighbors, distances, sample_filter_copy); } catch (const std::bad_cast&) { } @@ -373,14 +413,14 @@ void search(raft::resources const& res, search_params params_copy = params; if (params.filtering_rate < 0.0) { const auto num_set_bits = sample_filter.bitset_view_.count(res); - auto filtering_rate = (float)(idx.data().n_rows() - num_set_bits) / idx.data().n_rows(); + auto filtering_rate = (float)(idx.dataset().n_rows() - num_set_bits) / idx.dataset().n_rows(); const float min_filtering_rate = 0.0; const float max_filtering_rate = 0.999; params_copy.filtering_rate = std::min(std::max(filtering_rate, min_filtering_rate), max_filtering_rate); } auto sample_filter_copy = sample_filter; - return search_with_filtering( + return search_with_filtering( res, params_copy, idx, queries, neighbors, distances, sample_filter_copy); } catch (const std::bad_cast&) { } @@ -399,44 +439,114 @@ void search(raft::resources const& res, max_filtering_rate); } auto sample_filter_copy = sample_filter; - return search_with_filtering( + return search_with_filtering( res, params_copy, idx, queries, neighbors, distances, sample_filter_copy); } catch (const std::bad_cast&) { RAFT_FAIL("Unsupported sample filter type"); } } -template -void extend( - raft::resources const& handle, - raft::mdspan, raft::row_major, Accessor> additional_dataset, - cuvs::neighbors::cagra::index& index, - const cagra::extend_params& params, - std::optional> ndv, - std::optional> ngv) +template +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::device_padded_dataset_view additional_dataset, + cuvs::neighbors::cagra::index& index, + cuvs::neighbors::cagra::extended_dataset_storage& storage) +{ + auto cur_ds = index.dataset().view(); + const auto stride_elems = cur_ds.stride(0) > 0 ? static_cast(cur_ds.stride(0)) + : static_cast(cur_ds.extent(1)); + auto ndv = std::optional>( + raft::make_device_strided_matrix_view(storage.dataset_storage.data_handle(), + storage.dataset_storage.extent(0), + static_cast(index.dim()), + stride_elems)); + auto ngv = std::optional>(storage.graph_storage.view()); + extend_core(handle, additional_dataset.view(), index, params, ndv, ngv); +} + +template +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::host_padded_dataset_view additional_dataset, + cuvs::neighbors::cagra::index& index, + cuvs::neighbors::cagra::extended_dataset_storage& storage) +{ + auto cur_ds = index.dataset().view(); + const auto stride_elems = cur_ds.stride(0) > 0 ? static_cast(cur_ds.stride(0)) + : static_cast(cur_ds.extent(1)); + auto ndv = std::optional>( + raft::make_device_strided_matrix_view(storage.dataset_storage.data_handle(), + storage.dataset_storage.extent(0), + static_cast(index.dim()), + stride_elems)); + auto ngv = std::optional>(storage.graph_storage.view()); + extend_core(handle, additional_dataset.view(), index, params, ndv, ngv); +} + +template +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::device_standard_dataset_view additional_dataset, + cuvs::neighbors::cagra::index& index, + cuvs::neighbors::cagra::extended_dataset_storage& storage) +{ + auto cur_ds = index.dataset().view(); + const auto stride_elems = cur_ds.stride(0) > 0 ? static_cast(cur_ds.stride(0)) + : static_cast(cur_ds.extent(1)); + auto ndv = std::optional>( + raft::make_device_strided_matrix_view(storage.dataset_storage.data_handle(), + storage.dataset_storage.extent(0), + static_cast(index.dim()), + stride_elems)); + auto ngv = std::optional>(storage.graph_storage.view()); + extend_core(handle, additional_dataset.view(), index, params, ndv, ngv); +} + +template +void extend(raft::resources const& handle, + const cagra::extend_params& params, + cuvs::neighbors::host_standard_dataset_view additional_dataset, + cuvs::neighbors::cagra::index& index, + cuvs::neighbors::cagra::extended_dataset_storage& storage) { - cagra::extend_core(handle, additional_dataset, index, params, ndv, ngv); + auto cur_ds = index.dataset().view(); + const auto stride_elems = cur_ds.stride(0) > 0 ? static_cast(cur_ds.stride(0)) + : static_cast(cur_ds.extent(1)); + auto ndv = std::optional>( + raft::make_device_strided_matrix_view(storage.dataset_storage.data_handle(), + storage.dataset_storage.extent(0), + static_cast(index.dim()), + stride_elems)); + auto ngv = std::optional>(storage.graph_storage.view()); + extend_core(handle, additional_dataset.view(), index, params, ndv, ngv); } -template -index merge(raft::resources const& handle, - const cagra::index_params& params, - std::vector*>& indices, - const cuvs::neighbors::filtering::base_filter& row_filter) +template +cuvs::neighbors::cagra::index merge( + raft::resources const& handle, + const cagra::index_params& params, + std::vector*>& indices, + merged_dataset_storage& storage, + const cuvs::neighbors::filtering::base_filter& row_filter) { - return cagra::detail::merge(handle, params, indices, row_filter); + return cagra::detail::merge(handle, params, indices, storage, row_filter); } /** @} */ // end group cagra } // namespace cuvs::neighbors::cagra -#define CUVS_INST_CAGRA_MERGE(T, IdxT) \ - auto merge(raft::resources const& handle, \ - const cuvs::neighbors::cagra::index_params& params, \ - std::vector*>& indices, \ - const cuvs::neighbors::filtering::base_filter& row_filter) \ - -> cuvs::neighbors::cagra::index \ - { \ - return cuvs::neighbors::cagra::merge(handle, params, indices, row_filter); \ - } +#define CUVS_INST_CAGRA_MERGE(T, IdxT, DatasetViewT) \ + template CUVS_EXPORT cuvs::neighbors::cagra::merged_dataset_storage \ + cuvs::neighbors::cagra::make_merged_storage( \ + raft::resources const& handle, \ + std::vector*> const& indices, \ + cuvs::neighbors::filtering::base_filter const& row_filter); \ + template CUVS_EXPORT cuvs::neighbors::cagra::index \ + cuvs::neighbors::cagra::merge( \ + raft::resources const& handle, \ + const cuvs::neighbors::cagra::index_params& params, \ + std::vector*>& indices, \ + cuvs::neighbors::cagra::merged_dataset_storage& storage, \ + cuvs::neighbors::filtering::base_filter const& row_filter); diff --git a/cpp/src/neighbors/cagra_build_inst.cu.in b/cpp/src/neighbors/cagra_build_inst.cu.in index 00e0fab327..acaaa942c1 100644 --- a/cpp/src/neighbors/cagra_build_inst.cu.in +++ b/cpp/src/neighbors/cagra_build_inst.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -8,15 +8,28 @@ #include #include -namespace { +#include -using data_t = @data_type@; -using index_t = @index_type@; +namespace { +using data_t = @data_type@; +using index_t = @index_type@; +using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view; +using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; +using inst_host_padded_view_t = cuvs::neighbors::host_padded_dataset_view; +using inst_host_standard_view_t = cuvs::neighbors::host_standard_dataset_view; } // namespace namespace cuvs::neighbors::cagra { +extern template void index::compute_dataset_norms_( + raft::resources const&); +extern template void index::compute_dataset_norms_( + raft::resources const&); +extern template void index::compute_dataset_norms_( + raft::resources const&); +extern template void index::compute_dataset_norms_( + raft::resources const&); void build_knn_graph(raft::resources const& handle, raft::host_matrix_view dataset, raft::host_matrix_view knn_graph, @@ -25,22 +38,23 @@ void build_knn_graph(raft::resources const& handle, cuvs::neighbors::cagra::build_knn_graph(handle, dataset, knn_graph, params); } -auto build(raft::resources const& handle, - const cuvs::neighbors::cagra::index_params& params, - raft::device_matrix_view dataset) - -> cuvs::neighbors::cagra::index -{ - return cuvs::neighbors::cagra::build(handle, params, dataset); -} - -auto build(raft::resources const& handle, - const cuvs::neighbors::cagra::index_params& params, - raft::host_matrix_view dataset) - -> cuvs::neighbors::cagra::index -{ - return cuvs::neighbors::cagra::build(handle, params, dataset); -} - -template struct index; +#define CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(DatasetViewT, ...) \ + auto build(raft::resources const& res, \ + const cuvs::neighbors::cagra::index_params& params, \ + DatasetViewT const& dataset) -> __VA_ARGS__ \ + { \ + return cuvs::neighbors::cagra::build(res, params, dataset); \ + } + +CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_device_padded_view_t, + cuvs::neighbors::cagra::device_padded_index); +CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_device_standard_view_t, + cuvs::neighbors::cagra::device_standard_index); +CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_host_padded_view_t, + cuvs::neighbors::cagra::host_padded_index); +CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_host_standard_view_t, + cuvs::neighbors::cagra::host_standard_index); + +#undef CUVS_DEFINE_CAGRA_BUILD_OVERLOAD } // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_extend_inst.cu.in b/cpp/src/neighbors/cagra_extend_inst.cu.in index d544789713..9a90f1faa1 100644 --- a/cpp/src/neighbors/cagra_extend_inst.cu.in +++ b/cpp/src/neighbors/cagra_extend_inst.cu.in @@ -1,42 +1,45 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ -#include - #include #include namespace { -using data_t = @data_type@; -using index_t = @index_type@; +using data_t = @data_type@; +using index_t = @index_type@; +using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view; +using inst_host_padded_view_t = cuvs::neighbors::host_padded_dataset_view; +using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; +using inst_host_standard_view_t = cuvs::neighbors::host_standard_dataset_view; } // namespace namespace cuvs::neighbors::cagra { -void extend(raft::resources const& handle, - const cagra::extend_params& params, - raft::device_matrix_view additional_dataset, - cuvs::neighbors::cagra::index& idx, - std::optional> ndv, - std::optional> ngv) -{ - cuvs::neighbors::cagra::extend( - handle, additional_dataset, idx, params, ndv, ngv); -} - -void extend(raft::resources const& handle, - const cagra::extend_params& params, - raft::host_matrix_view additional_dataset, - cuvs::neighbors::cagra::index& idx, - std::optional> ndv, - std::optional> ngv) -{ - cuvs::neighbors::cagra::extend( - handle, additional_dataset, idx, params, ndv, ngv); -} +extern template void index::compute_dataset_norms_( + raft::resources const&); +extern template void index::compute_dataset_norms_( + raft::resources const&); + +#define CUVS_INST_CAGRA_EXTEND(T, IdxT, DatasetViewT, AdditionalDatasetT) \ + void extend(raft::resources const& handle, \ + const cuvs::neighbors::cagra::extend_params& params, \ + AdditionalDatasetT additional_dataset, \ + cuvs::neighbors::cagra::index& index, \ + cuvs::neighbors::cagra::extended_dataset_storage& storage) \ + { \ + cuvs::neighbors::cagra::extend( \ + handle, params, additional_dataset, index, storage); \ + } + +CUVS_INST_CAGRA_EXTEND(data_t, index_t, inst_device_padded_view_t, inst_device_padded_view_t); +CUVS_INST_CAGRA_EXTEND(data_t, index_t, inst_device_padded_view_t, inst_host_padded_view_t); +CUVS_INST_CAGRA_EXTEND(data_t, index_t, inst_device_standard_view_t, inst_device_standard_view_t); +CUVS_INST_CAGRA_EXTEND(data_t, index_t, inst_device_standard_view_t, inst_host_standard_view_t); + +#undef CUVS_INST_CAGRA_EXTEND } // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_merge_inst.cu.in b/cpp/src/neighbors/cagra_merge_inst.cu.in index 2fafb37ae4..7ae6661d4c 100644 --- a/cpp/src/neighbors/cagra_merge_inst.cu.in +++ b/cpp/src/neighbors/cagra_merge_inst.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -10,11 +10,34 @@ namespace { -using data_t = @data_type@; -using index_t = @index_type@; +using data_t = @data_type@; +using index_t = @index_type@; +using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view; } // namespace namespace cuvs::neighbors::cagra { -CUVS_INST_CAGRA_MERGE(data_t, index_t); + +using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; + +extern template void index::compute_dataset_norms_( + raft::resources const&); +extern template void index::compute_dataset_norms_( + raft::resources const&); + +CUVS_INST_CAGRA_MERGE(data_t, index_t, inst_device_padded_view_t); +CUVS_INST_CAGRA_MERGE(data_t, index_t, inst_device_standard_view_t); + +// Keep explicit merge instantiations retained in libcuvs.so even when they are only +// referenced from external test/library link units. +template +CUVS_EXPORT void merge_instantiation_keepalive() +{ + (void)&cuvs::neighbors::cagra::make_merged_storage; + (void)&cuvs::neighbors::cagra::make_merged_storage; + (void)&cuvs::neighbors::cagra::merge; + (void)&cuvs::neighbors::cagra::merge; +} +template CUVS_EXPORT void merge_instantiation_keepalive(); + } // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_search_inst.cu.in b/cpp/src/neighbors/cagra_search_inst.cu.in index dfef630798..f1b506a27f 100644 --- a/cpp/src/neighbors/cagra_search_inst.cu.in +++ b/cpp/src/neighbors/cagra_search_inst.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -8,28 +8,54 @@ namespace { -using data_t = @data_type@; +using data_t = @data_type@; +using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view; +using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; +using inst_vpq_f16_view_t = cuvs::neighbors::device_vpq_dataset_view; +using inst_vpq_f32_view_t = cuvs::neighbors::device_vpq_dataset_view; +using inst_empty_view_t = cuvs::neighbors::device_empty_dataset_view; -} +} // namespace namespace cuvs::neighbors::cagra { -#define CUVS_INST_CAGRA_SEARCH(T, IdxT, OutputIdxT) \ +#define CUVS_INST_CAGRA_SEARCH(T, IdxT, DatasetViewT, OutputIdxT) \ void search(raft::resources const& handle, \ cuvs::neighbors::cagra::search_params const& params, \ - const cuvs::neighbors::cagra::index& index, \ + const cuvs::neighbors::cagra::index& index, \ raft::device_matrix_view queries, \ raft::device_matrix_view neighbors, \ raft::device_matrix_view distances, \ const cuvs::neighbors::filtering::base_filter& sample_filter) \ { \ - cuvs::neighbors::cagra::search( \ + cuvs::neighbors::cagra::search( \ handle, params, index, queries, neighbors, distances, sample_filter); \ } -CUVS_INST_CAGRA_SEARCH(data_t, uint32_t, uint32_t); -CUVS_INST_CAGRA_SEARCH(data_t, uint32_t, int64_t); +#define CUVS_INST_CAGRA_SEARCH_ALL_VIEWS(T, OutputIdxT) \ + CUVS_INST_CAGRA_SEARCH(T, uint32_t, inst_device_padded_view_t, OutputIdxT); \ + CUVS_INST_CAGRA_SEARCH(T, uint32_t, inst_device_standard_view_t, OutputIdxT); \ + CUVS_INST_CAGRA_SEARCH(T, uint32_t, inst_vpq_f16_view_t, OutputIdxT); \ + CUVS_INST_CAGRA_SEARCH(T, uint32_t, inst_vpq_f32_view_t, OutputIdxT); \ + CUVS_INST_CAGRA_SEARCH(T, uint32_t, inst_empty_view_t, OutputIdxT) +CUVS_INST_CAGRA_SEARCH_ALL_VIEWS(data_t, uint32_t); +CUVS_INST_CAGRA_SEARCH_ALL_VIEWS(data_t, int64_t); + +#undef CUVS_INST_CAGRA_SEARCH_ALL_VIEWS #undef CUVS_INST_CAGRA_SEARCH +template CUVS_EXPORT __attribute__((externally_visible)) void +index::compute_dataset_norms_( + raft::resources const& res); +template CUVS_EXPORT __attribute__((externally_visible)) void +index::compute_dataset_norms_( + raft::resources const& res); +template CUVS_EXPORT __attribute__((externally_visible)) void +index::compute_dataset_norms_(raft::resources const& res); +template CUVS_EXPORT __attribute__((externally_visible)) void +index::compute_dataset_norms_(raft::resources const& res); +template CUVS_EXPORT __attribute__((externally_visible)) void +index::compute_dataset_norms_(raft::resources const& res); + } // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_serialize.cuh b/cpp/src/neighbors/cagra_serialize.cuh index b18577255a..d9a2e068f7 100644 --- a/cpp/src/neighbors/cagra_serialize.cuh +++ b/cpp/src/neighbors/cagra_serialize.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,56 +9,157 @@ namespace cuvs::neighbors::cagra { -#define CUVS_INST_CAGRA_SERIALIZE(DTYPE) \ - void serialize(raft::resources const& handle, \ - const std::string& filename, \ - const cuvs::neighbors::cagra::index& index, \ - bool include_dataset) \ - { \ - cuvs::neighbors::cagra::detail::serialize( \ - handle, filename, index, include_dataset); \ - }; \ - \ - void deserialize(raft::resources const& handle, \ - const std::string& filename, \ - cuvs::neighbors::cagra::index* index) \ - { \ - cuvs::neighbors::cagra::detail::deserialize(handle, filename, index); \ - }; \ - void serialize(raft::resources const& handle, \ - std::ostream& os, \ - const cuvs::neighbors::cagra::index& index, \ - bool include_dataset) \ - { \ - cuvs::neighbors::cagra::detail::serialize( \ - handle, os, index, include_dataset); \ - } \ - \ - void deserialize(raft::resources const& handle, \ - std::istream& is, \ - cuvs::neighbors::cagra::index* index) \ - { \ - cuvs::neighbors::cagra::detail::deserialize(handle, is, index); \ - } \ - \ - void serialize_to_hnswlib( \ - raft::resources const& handle, \ - std::ostream& os, \ - const cuvs::neighbors::cagra::index& index, \ - std::optional> dataset) \ - { \ - cuvs::neighbors::cagra::detail::serialize_to_hnswlib( \ - handle, os, index, dataset); \ - } \ - \ - void serialize_to_hnswlib( \ - raft::resources const& handle, \ - const std::string& filename, \ - const cuvs::neighbors::cagra::index& index, \ - std::optional> dataset) \ - { \ - cuvs::neighbors::cagra::detail::serialize_to_hnswlib( \ - handle, filename, index, dataset); \ +#define CUVS_INST_CAGRA_SERIALIZE(DTYPE) \ + void serialize(raft::resources const& handle, \ + const std::string& filename, \ + const cuvs::neighbors::cagra::device_padded_index& index, \ + bool include_dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize( \ + handle, filename, index, include_dataset); \ + }; \ + \ + void deserialize( \ + raft::resources const& handle, \ + const std::string& filename, \ + cuvs::neighbors::cagra::device_padded_index* index, \ + std::unique_ptr>* out_dataset) \ + { \ + cuvs::neighbors::cagra::detail::deserialize( \ + handle, filename, index, out_dataset); \ + }; \ + void serialize(raft::resources const& handle, \ + std::ostream& os, \ + const cuvs::neighbors::cagra::device_padded_index& index, \ + bool include_dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize( \ + handle, os, index, include_dataset); \ + } \ + \ + void deserialize( \ + raft::resources const& handle, \ + std::istream& is, \ + cuvs::neighbors::cagra::device_padded_index* index, \ + std::unique_ptr>* out_dataset) \ + { \ + cuvs::neighbors::cagra::detail::deserialize(handle, is, index, out_dataset); \ + } \ + \ + void serialize(raft::resources const& handle, \ + const std::string& filename, \ + const cuvs::neighbors::cagra::device_standard_index& index, \ + bool include_dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize( \ + handle, filename, index, include_dataset); \ + }; \ + \ + void deserialize( \ + raft::resources const& handle, \ + const std::string& filename, \ + cuvs::neighbors::cagra::device_standard_index* index, \ + std::unique_ptr>* out_dataset) \ + { \ + cuvs::neighbors::cagra::detail::deserialize( \ + handle, filename, index, out_dataset); \ + }; \ + void serialize(raft::resources const& handle, \ + std::ostream& os, \ + const cuvs::neighbors::cagra::device_standard_index& index, \ + bool include_dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize( \ + handle, os, index, include_dataset); \ + } \ + \ + void deserialize( \ + raft::resources const& handle, \ + std::istream& is, \ + cuvs::neighbors::cagra::device_standard_index* index, \ + std::unique_ptr>* out_dataset) \ + { \ + cuvs::neighbors::cagra::detail::deserialize(handle, is, index, out_dataset); \ + } \ + \ + void serialize_to_hnswlib( \ + raft::resources const& handle, \ + std::ostream& os, \ + const cuvs::neighbors::cagra::device_padded_index& index, \ + std::optional> dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize_to_hnswlib( \ + handle, os, index, dataset); \ + } \ + \ + void serialize_to_hnswlib( \ + raft::resources const& handle, \ + const std::string& filename, \ + const cuvs::neighbors::cagra::device_padded_index& index, \ + std::optional> dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize_to_hnswlib( \ + handle, filename, index, dataset); \ + } \ + \ + void serialize_to_hnswlib( \ + raft::resources const& handle, \ + std::ostream& os, \ + const cuvs::neighbors::cagra::host_padded_index& index, \ + std::optional> dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize_to_hnswlib( \ + handle, os, index, dataset); \ + } \ + \ + void serialize_to_hnswlib( \ + raft::resources const& handle, \ + const std::string& filename, \ + const cuvs::neighbors::cagra::host_padded_index& index, \ + std::optional> dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize_to_hnswlib( \ + handle, filename, index, dataset); \ + } \ + \ + void serialize_to_hnswlib( \ + raft::resources const& handle, \ + std::ostream& os, \ + const cuvs::neighbors::cagra::device_standard_index& index, \ + std::optional> dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize_to_hnswlib( \ + handle, os, index, dataset); \ + } \ + \ + void serialize_to_hnswlib( \ + raft::resources const& handle, \ + const std::string& filename, \ + const cuvs::neighbors::cagra::device_standard_index& index, \ + std::optional> dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize_to_hnswlib( \ + handle, filename, index, dataset); \ + } \ + \ + void serialize_to_hnswlib( \ + raft::resources const& handle, \ + std::ostream& os, \ + const cuvs::neighbors::cagra::host_standard_index& index, \ + std::optional> dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize_to_hnswlib( \ + handle, os, index, dataset); \ + } \ + \ + void serialize_to_hnswlib( \ + raft::resources const& handle, \ + const std::string& filename, \ + const cuvs::neighbors::cagra::host_standard_index& index, \ + std::optional> dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize_to_hnswlib( \ + handle, filename, index, dataset); \ } } // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_serialize_inst.cu.in b/cpp/src/neighbors/cagra_serialize_inst.cu.in index 4b14ee72f4..3d34adb36f 100644 --- a/cpp/src/neighbors/cagra_serialize_inst.cu.in +++ b/cpp/src/neighbors/cagra_serialize_inst.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,12 +9,19 @@ namespace { -using data_t = @data_type@; +using data_t = @data_type@; +using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view; +using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; -} +} // namespace namespace cuvs::neighbors::cagra { +extern template void index::compute_dataset_norms_( + raft::resources const&); +extern template void index::compute_dataset_norms_( + raft::resources const&); + CUVS_INST_CAGRA_SERIALIZE(data_t); } // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/detail/cagra/add_nodes.cuh b/cpp/src/neighbors/detail/cagra/add_nodes.cuh index 5d0a6654e9..0510c4a8c4 100644 --- a/cpp/src/neighbors/detail/cagra/add_nodes.cuh +++ b/cpp/src/neighbors/detail/cagra/add_nodes.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "../../../core/omp_wrapper.hpp" @@ -20,10 +20,10 @@ namespace cuvs::neighbors::cagra { -template +template void add_node_core( raft::resources const& handle, - const cuvs::neighbors::cagra::index& idx, + const cuvs::neighbors::cagra::index& idx, raft::mdspan, raft::layout_stride, Accessor> additional_dataset_view, raft::host_matrix_view updated_graph, @@ -32,7 +32,7 @@ void add_node_core( using DistanceT = float; const std::size_t degree = idx.graph_degree(); const std::size_t dim = idx.dim(); - const std::size_t old_size = idx.dataset().extent(0); + const std::size_t old_size = idx.dataset().n_rows(); const std::size_t num_add = additional_dataset_view.extent(0); const std::size_t new_size = old_size + num_add; const std::uint32_t base_degree = degree * 2; @@ -276,11 +276,11 @@ void add_node_core( } } -template +template void add_graph_nodes( raft::resources const& handle, raft::device_matrix_view input_updated_dataset_view, - const neighbors::cagra::index& index, + const neighbors::cagra::index& index, raft::host_matrix_view updated_graph_view, const cagra::extend_params& params) { @@ -297,15 +297,17 @@ void add_graph_nodes( const std::size_t max_chunk_size_ = params.max_chunk_size == 0 ? new_dataset_size : params.max_chunk_size; - raft::copy(handle, - raft::make_device_vector_view(updated_graph_view.data_handle(), index.graph().size()), - raft::make_device_vector_view(index.graph().data_handle(), index.graph().size())); + auto updated_graph_prefix = raft::make_host_matrix_view( + updated_graph_view.data_handle(), initial_dataset_size, degree); + raft::copy(handle, updated_graph_prefix, raft::make_const_mdspan(index.graph())); - neighbors::cagra::index internal_index( - handle, - index.metric(), - raft::make_device_matrix_view(nullptr, 0, dim), - raft::make_device_matrix_view(nullptr, 0, degree)); + using padded_view_t = cuvs::neighbors::device_padded_dataset_view; + auto zero_row = raft::make_device_matrix_view( + static_cast(nullptr), int64_t{0}, static_cast(dim)); + padded_view_t device_empty_dataset_view(zero_row, static_cast(dim)); + auto empty_graph_view = raft::make_device_matrix_view(nullptr, 0, degree); + neighbors::cagra::index internal_index( + handle, index.metric(), device_empty_dataset_view, empty_graph_view); for (std::size_t additional_dataset_offset = 0; additional_dataset_offset < num_new_nodes; additional_dataset_offset += max_chunk_size_) { @@ -320,7 +322,8 @@ void add_graph_nodes( auto graph_view = raft::make_host_matrix_view( updated_graph_view.data_handle(), initial_dataset_size + additional_dataset_offset, degree); - internal_index.update_dataset(handle, dataset_view); + auto pdv = cuvs::neighbors::make_device_padded_dataset_view(handle, dataset_view); + internal_index.update_dataset(handle, pdv); // Note: The graph is copied to the device memory. internal_index.update_graph(handle, graph_view); @@ -337,32 +340,32 @@ void add_graph_nodes( dim, stride); - neighbors::cagra::add_node_core( + neighbors::cagra::add_node_core( handle, internal_index, additional_dataset_view, updated_graph, params); raft::resource::sync_stream(handle); } } -template +template void extend_core( raft::resources const& handle, raft::mdspan, raft::row_major, Accessor> additional_dataset, - cuvs::neighbors::cagra::index& index, + cuvs::neighbors::cagra::index& index, const cagra::extend_params& params, std::optional> new_dataset_buffer_view, std::optional> new_graph_buffer_view) { + static_assert(cuvs::neighbors::is_padded_dataset_view_v || + cuvs::neighbors::is_standard_dataset_view_v, + "cagra::extend requires a padded or standard dataset view index type"); RAFT_EXPECTS(!index.dataset_fd().has_value(), "Cannot extend a disk-backed CAGRA index. Convert it with " "cuvs::neighbors::hnsw::from_cagra() and load it into memory via " "cuvs::neighbors::hnsw::deserialize() before calling extend()."); - if (dynamic_cast*>(&index.data()) != nullptr && - !new_dataset_buffer_view.has_value()) { - RAFT_LOG_WARN( - "New memory space for extended dataset will be allocated while the memory space for the old " - "dataset is allocated by user."); - } + RAFT_EXPECTS(new_dataset_buffer_view.has_value(), + "cagra::extend requires new_dataset_buffer_view. " + "Provide a buffer view for the extended dataset (initial + additional vectors)."); const std::size_t num_new_nodes = additional_dataset.extent(0); const std::size_t initial_dataset_size = index.size(); const std::size_t new_dataset_size = initial_dataset_size + num_new_nodes; @@ -391,26 +394,23 @@ void extend_core( num_new_nodes); } - using ds_idx_type = decltype(index.data().n_rows()); - if (auto* strided_dset = dynamic_cast*>(&index.data()); - strided_dset != nullptr) { + auto try_extend = [&](auto const& leaf) { // Allocate memory space for updated graph on host auto updated_graph = raft::make_host_matrix(new_dataset_size, degree); - const auto stride = strided_dset->stride(); - auto updated_dataset = raft::make_device_matrix(handle, 0, stride); - auto updated_dataset_view = - raft::make_device_strided_matrix_view(nullptr, 0, dim, stride); + const std::size_t stride = static_cast(leaf.stride()); + const T* src_rows = leaf.view().data_handle(); + auto updated_dataset_view = new_dataset_buffer_view.value(); - // Update dataset + // Update dataset on host, then copy to device buffer provided by caller auto host_updated_dataset = raft::make_host_matrix(new_dataset_size, stride); - // The padding area must be filled with zeros.!!!!!!!!!!!!!!!!!!! + // The padding area must be filled with zeros. memset(host_updated_dataset.data_handle(), 0, sizeof(T) * host_updated_dataset.size()); raft::copy_matrix(host_updated_dataset.data_handle(), stride, - strided_dset->view().data_handle(), + src_rows, stride, dim, initial_dataset_size, @@ -423,42 +423,27 @@ void extend_core( num_new_nodes, raft::resource::get_cuda_stream(handle)); - if (new_dataset_buffer_view.has_value()) { - updated_dataset_view = new_dataset_buffer_view.value(); - } else { - // Deallocate the current dataset memory space if the dataset is `owning'. - index.update_dataset( - handle, raft::make_device_strided_matrix_view(nullptr, 0, dim, stride)); - - // Allocate the new dataset - updated_dataset = raft::make_device_matrix(handle, new_dataset_size, stride); - updated_dataset_view = raft::make_device_strided_matrix_view( - updated_dataset.data_handle(), new_dataset_size, dim, stride); - } - - // Copy updated dataset on host memory to device memory - raft::copy( - handle, - raft::make_device_vector_view(updated_dataset_view.data_handle(), new_dataset_size * stride), - raft::make_host_vector_view(host_updated_dataset.data_handle(), new_dataset_size * stride)); + // Copy updated dataset on host memory to device memory (caller's buffer) + raft::copy(updated_dataset_view.data_handle(), + host_updated_dataset.data_handle(), + new_dataset_size * stride, + raft::resource::get_cuda_stream(handle)); // Add graph nodes cuvs::neighbors::cagra::add_graph_nodes( handle, raft::make_const_mdspan(updated_dataset_view), index, updated_graph.view(), params); - // Update index dataset - if (new_dataset_buffer_view.has_value()) { - index.update_dataset(handle, raft::make_const_mdspan(updated_dataset_view)); + // Attach view over caller's buffer; index does not take ownership + if constexpr (cuvs::neighbors::is_padded_dataset_view_v) { + cuvs::neighbors::device_padded_dataset_view dv( + raft::make_device_matrix_view(updated_dataset_view.data_handle(), + updated_dataset_view.extent(0), + updated_dataset_view.stride(0)), + dim); + index.update_dataset(handle, dv); } else { - using out_mdarray_type = decltype(updated_dataset); - using out_layout_type = typename out_mdarray_type::layout_type; - using out_container_policy_type = typename out_mdarray_type::container_policy_type; - using out_owning_type = - owning_dataset; - auto out_layout = raft::make_strided_layout(updated_dataset_view.extents(), - cuda::std::array{stride, 1}); - - index.update_dataset(handle, out_owning_type{std::move(updated_dataset), out_layout}); + auto dv = cuvs::neighbors::make_device_standard_dataset_view(updated_dataset_view); + index.update_dataset(handle, dv); } // Update index graph @@ -472,12 +457,16 @@ void extend_core( } else { index.update_graph(handle, raft::make_const_mdspan(updated_graph.view())); } - } else if (dynamic_cast*>(&index.data()) != - nullptr) { + }; + + auto const& leaf = index.dataset(); + if constexpr (cuvs::neighbors::is_padded_dataset_view_v> || + cuvs::neighbors::is_standard_dataset_view_v>) { + try_extend(leaf); + } else if constexpr (cuvs::neighbors::is_empty_dataset_view_v>) { RAFT_FAIL( "cagra::extend only supports an index to which the dataset is attached. Please check if the " - "index was built with index_param.attach_dataset_on_build = true, or if a dataset was " - "attached after the build."); + "index has an empty dataset; attach one with update_dataset before extend."); } else { RAFT_FAIL("cagra::extend only supports an uncompressed dataset index"); } diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index e0ca8fe0be..3e587bb873 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -5,9 +5,9 @@ #pragma once #include "../../../core/nvtx.hpp" -#include "../../../preprocessing/quantize/vpq_build-ext.cuh" #include "../../ivf_pq/ivf_pq_fp16_overflow.cuh" #include "graph_core.cuh" +#include #include #include @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -39,6 +40,7 @@ #include #include #include +#include #include #include #include @@ -1111,6 +1113,13 @@ void ace_validate_disk_mode_partitions(raft::resources const& res, } } +template + requires cuvs::neighbors::is_dense_row_major_device_dataset_view_v +auto build_from_device_matrix(raft::resources const& res, + const index_params& params, + DatasetViewT const& device_dataset) + -> cuvs::neighbors::cagra::index; + // Build CAGRA index using ACE (Augmented Core Extraction) partitioning // ACE enables building indexes for datasets too large to fit in GPU memory by: // 1. Partitioning the dataset using balanced k-means in core (non-overlapping) and augmented @@ -1120,10 +1129,12 @@ void ace_validate_disk_mode_partitions(raft::resources const& res, // Supports both in-memory and disk-based modes depending on available host memory. // In disk mode, the graph is stored in build_dir and dataset is reordered on disk. // The returned index is not usable for search. Use the created files for search instead. -template -index build_ace(raft::resources const& res, - const index_params& params, - raft::host_matrix_view dataset) +template + requires cuvs::neighbors::is_host_dataset_view_v +auto build_ace(raft::resources const& res, + const index_params& params, + raft::host_matrix_view dataset) + -> cuvs::neighbors::cagra::index { // Extract ACE parameters from graph_build_params RAFT_EXPECTS( @@ -1137,7 +1148,7 @@ index build_ace(raft::resources const& res, bool use_disk = ace_params.use_disk; common::nvtx::range function_scope( - "cagra::build_ace(%zu, %zu, %zu)", + "cagra::detail::build_ace(%zu, %zu, %zu)", params.intermediate_graph_degree, params.graph_degree, npartitions); @@ -1390,11 +1401,14 @@ index build_ace(raft::resources const& res, ef_construction, cuvs::neighbors::cagra::hnsw_heuristic_type::SAME_GRAPH_FOOTPRINT, params.metric); - sub_index_params.attach_dataset_on_build = false; - sub_index_params.guarantee_connectivity = params.guarantee_connectivity; + sub_index_params.guarantee_connectivity = params.guarantee_connectivity; - auto sub_index = cuvs::neighbors::cagra::build( - res, sub_index_params, raft::make_const_mdspan(sub_dataset.view())); + // Copy host partition to device with padding; build_from_device_matrix accepts + // device_padded_dataset_view. + auto sub_dataset_dev = cuvs::neighbors::make_device_padded_dataset( + res, raft::make_const_mdspan(sub_dataset.view())); + auto sub_index = ::cuvs::neighbors::cagra::detail::build_from_device_matrix( + res, sub_index_params, sub_dataset_dev->as_dataset_view()); auto optimize_end = std::chrono::high_resolution_clock::now(); auto optimize_elapsed = @@ -1494,25 +1508,9 @@ index build_ace(raft::resources const& res, } auto index_creation_start = std::chrono::high_resolution_clock::now(); - index idx(res, params.metric); - // Only add graph and dataset if not using disk storage. The returned index is empty if using - // disk storage. Use the files written to disk for search. + cuvs::neighbors::cagra::index idx(res, params.metric); if (!use_disk_mode) { idx.update_graph(res, raft::make_const_mdspan(search_graph.view())); - - if (params.attach_dataset_on_build) { - try { - idx.update_dataset(res, dataset); - } catch (std::bad_alloc& e) { - RAFT_LOG_WARN( - "Insufficient GPU memory to attach dataset to ACE index. Only the graph will be " - "stored."); - } catch (raft::logic_error& e) { - RAFT_LOG_WARN( - "Insufficient GPU memory to attach dataset to ACE index. Only the graph will be " - "stored."); - } - } } else { idx.update_dataset(res, std::move(reordered_fd)); idx.update_graph(res, std::move(graph_fd)); @@ -1538,7 +1536,7 @@ index build_ace(raft::resources const& res, std::chrono::duration_cast(total_end - total_start).count(); RAFT_LOG_INFO("ACE: Partitioned CAGRA build completed in %ld ms total", total_elapsed); - return idx; + return std::move(idx); } catch (const std::exception& e) { // Clean up build directory on failure if we created it RAFT_LOG_ERROR("ACE: Build failed with exception: %s", e.what()); @@ -2016,14 +2014,29 @@ struct mmap_owner { size_t size_; }; -template , raft::memory_type::host>> -auto iterative_build_graph( +/** Upload and/or pad `dataset` to a device-resident CAGRA-aligned view for iterative internal + * search. */ +template + requires cuvs::neighbors::is_dense_row_major_dataset_view_v +auto ensure_device_padded_for_iterative_search( raft::resources const& res, - const index_params& params, - raft::mdspan, raft::row_major, Accessor> dataset) + DatasetViewT const& dataset, + std::unique_ptr>& padded_own) + -> cuvs::neighbors::device_padded_dataset_view +{ + if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + return dataset; + } else { + padded_own = cuvs::neighbors::make_device_padded_dataset(res, dataset.view()); + return padded_own->as_dataset_view(); + } +} + +template + requires cuvs::neighbors::is_dense_row_major_dataset_view_v +auto iterative_build_graph(raft::resources const& res, + const index_params& params, + DatasetViewT const& dataset) -> raft::host_matrix { size_t intermediate_degree = params.intermediate_graph_degree; size_t graph_degree = params.graph_degree; @@ -2031,32 +2044,19 @@ auto iterative_build_graph( auto cagra_graph = raft::make_host_matrix(0, 0); // Iteratively improve the accuracy of the graph by repeatedly running - // CAGRA's search() and optimize(). As for the size of the graph, instead - // of targeting all nodes from the beginning, the number of nodes is - // initially small, and the number of nodes is doubled with each iteration. + // CAGRA's search() and optimize(). Host or non-CAGRA-aligned device inputs are uploaded + // and padded here only for the internal search loop — same role as main's + // make_aligned_dataset() inside iterative_build_graph. IVF-PQ / NN-descent never take this path. RAFT_LOG_INFO("Iteratively creating/improving graph index using CAGRA's search() and optimize()"); - // If dataset is a host matrix, change it to a device matrix. Also, if the - // dimensionality of the dataset does not meet the alighnemt restriction, - // add extra dimensions and change it to a strided matrix. - std::unique_ptr> dev_aligned_dataset; - try { - dev_aligned_dataset = make_aligned_dataset(res, dataset); - } catch (raft::logic_error& e) { - RAFT_LOG_ERROR("Iterative CAGRA graph build requires the dataset to fit GPU memory"); - throw e; - } - auto dev_aligned_dataset_view = dev_aligned_dataset.get()->view(); + std::unique_ptr> padded_own; + auto search_dataset = ensure_device_padded_for_iterative_search(res, dataset, padded_own); - // If the matrix stride and extent do no match, the extra dimensions are - // also as extent since it cannot be used as query matrix. - auto dev_dataset = - raft::make_device_matrix_view(dev_aligned_dataset_view.data_handle(), - dev_aligned_dataset_view.extent(0), - dev_aligned_dataset_view.stride(0)); + auto dev_dataset = search_dataset.view(); + uint32_t logical_dim = search_dataset.dim(); // Determine initial graph size. - uint64_t final_graph_size = (uint64_t)dataset.extent(0); + uint64_t final_graph_size = (uint64_t)search_dataset.n_rows(); uint64_t initial_graph_size = (final_graph_size + 1) / 2; while (initial_graph_size > graph_degree * 64) { initial_graph_size = (initial_graph_size + 1) / 2; @@ -2071,6 +2071,12 @@ auto iterative_build_graph( auto dev_neighbors = raft::make_device_matrix(res, max_chunk_size, topk); auto dev_distances = raft::make_device_matrix(res, max_chunk_size, topk); + std::optional> query_contiguous; + if (static_cast(logical_dim) != dev_dataset.extent(1)) { + query_contiguous.emplace( + raft::make_device_matrix(res, max_chunk_size, logical_dim)); + } + // Determine graph degree and number of search results while increasing // graph size. auto small_graph_degree = std::max(graph_degree / 2, std::min(graph_degree, (uint64_t)24)); @@ -2143,9 +2149,11 @@ auto iterative_build_graph( // search results (neighbors). auto dev_dataset_view = raft::make_device_matrix_view( dev_dataset.data_handle(), (int64_t)curr_graph_size, dev_dataset.extent(1)); + cuvs::neighbors::device_padded_dataset_view sub_padded(dev_dataset_view, + logical_dim); - auto idx = index( - res, params.metric, dev_dataset_view, raft::make_const_mdspan(cagra_graph.view())); + auto idx = cuvs::neighbors::cagra::device_padded_index( + res, params.metric, sub_padded, raft::make_const_mdspan(cagra_graph.view())); auto dev_query_view = raft::make_device_matrix_view( dev_dataset.data_handle(), (int64_t)curr_query_size, dev_dataset.extent(1)); @@ -2164,8 +2172,21 @@ auto iterative_build_graph( raft::resource::get_cuda_stream(res), raft::resource::get_workspace_resource_ref(res)); for (const auto& batch : query_batch) { - auto batch_dev_query_view = raft::make_device_matrix_view( - batch.data(), batch.size(), dev_query_view.extent(1)); + raft::device_matrix_view batch_dev_query_view; + if (query_contiguous) { + raft::copy_matrix(query_contiguous->data_handle(), + static_cast(logical_dim), + batch.data(), + dev_query_view.extent(1), + static_cast(logical_dim), + batch.size(), + raft::resource::get_cuda_stream(res)); + batch_dev_query_view = raft::make_device_matrix_view( + query_contiguous->data_handle(), batch.size(), static_cast(logical_dim)); + } else { + batch_dev_query_view = raft::make_device_matrix_view( + batch.data(), batch.size(), dev_query_view.extent(1)); + } auto batch_dev_neighbors_view = raft::make_device_matrix_view( dev_neighbors.data_handle(), batch.size(), curr_topk); auto batch_dev_distances_view = raft::make_device_matrix_view( @@ -2202,41 +2223,31 @@ auto iterative_build_graph( return cagra_graph; } -template , raft::memory_type::host>> -index build( +template +[[nodiscard]] inline auto resolve_cagra_default_knn_graph_build_params( raft::resources const& res, - const index_params& params, - raft::mdspan, raft::row_major, Accessor> dataset) + index_params const& params, + raft::matrix_extent dataset_extents, + size_t intermediate_degree) { - size_t intermediate_degree = params.intermediate_graph_degree; - size_t graph_degree = params.graph_degree; - common::nvtx::range function_scope( - "cagra::build<%s>(%zu, %zu)", - Accessor::is_managed_type::value ? "managed" - : Accessor::is_host_type::value ? "host" - : "device", - intermediate_degree, - graph_degree); - check_graph_degree(intermediate_degree, graph_degree, dataset.extent(0)); - - // Set default value in case knn_build_params is not defined. auto knn_build_params = params.graph_build_params; if (std::holds_alternative(params.graph_build_params)) { - // Heuristic to decide default build algo and its params. - if (cuvs::neighbors::nn_descent::has_enough_device_memory( - res, dataset.extents(), sizeof(IdxT))) { + if (cuvs::neighbors::nn_descent::has_enough_device_memory(res, dataset_extents, sizeof(IdxT))) { RAFT_LOG_DEBUG("NN descent solver"); knn_build_params = cagra::graph_build_params::nn_descent_params(intermediate_degree, params.metric); } else { RAFT_LOG_DEBUG("Selecting IVF-PQ solver"); - knn_build_params = cagra::graph_build_params::ivf_pq_params(dataset.extents(), params.metric); + knn_build_params = cagra::graph_build_params::ivf_pq_params(dataset_extents, params.metric); } } + return knn_build_params; +} +template +inline void validate_cagra_knn_graph_build_constraints(index_params const& params, + KnnParamsVariant const& knn_build_params) +{ RAFT_EXPECTS( params.metric != cuvs::distance::DistanceType::BitwiseHamming || std::holds_alternative( @@ -2250,104 +2261,171 @@ index build( std::holds_alternative(knn_build_params), "CosineExpanded distance is not supported for iterative CAGRA graph build."); - // Validate data type for BitwiseHamming metric RAFT_EXPECTS(params.metric != cuvs::distance::DistanceType::BitwiseHamming || (std::is_same_v || std::is_same_v), "BitwiseHamming distance is only supported for int8_t and uint8_t data types. " "Current data type is not supported."); +} - auto cagra_graph = raft::make_host_matrix(0, 0); +/** + * Iterative / IVF-PQ / NN-descent KNN graph construction and `optimize` → final host CAGRA graph. + * + * @param knn_graph_dataset mdspan passed to IVF-PQ / NN-descent `build_knn_graph` (any stride). + */ +template +auto build_cagra_host_graph_from_knn_params(raft::resources const& res, + index_params const& params, + KnnParamsVariant const& knn_build_params, + int64_t n_rows, + size_t intermediate_degree, + size_t graph_degree, + KnnGraphDatasetMdspan&& knn_graph_dataset) + -> raft::host_matrix +{ + std::optional> knn_graph( + raft::make_host_matrix(n_rows, intermediate_degree)); - // Dispatch based on graph_build_params - if (std::holds_alternative( - knn_build_params)) { - cagra_graph = iterative_build_graph(res, params, dataset); + if (std::holds_alternative(knn_build_params)) { + auto ivf_pq_params = + std::get(knn_build_params); + if (ivf_pq_params.build_params.metric != params.metric) { + RAFT_LOG_WARN( + "Metric (%lu) for IVF-PQ needs to match cagra metric (%lu), " + "aligning IVF-PQ metric.", + ivf_pq_params.build_params.metric, + params.metric); + ivf_pq_params.build_params.metric = params.metric; + } + build_knn_graph(res, knn_graph_dataset, knn_graph->view(), ivf_pq_params); } else { - std::optional> knn_graph( - raft::make_host_matrix(dataset.extent(0), intermediate_degree)); + auto nn_descent_params = + std::get(knn_build_params); - if (std::holds_alternative(knn_build_params)) { - auto ivf_pq_params = - std::get(knn_build_params); - if (ivf_pq_params.build_params.metric != params.metric) { - RAFT_LOG_WARN( - "Metric (%lu) for IVF-PQ needs to match cagra metric (%lu), " - "aligning IVF-PQ metric.", - ivf_pq_params.build_params.metric, - params.metric); - ivf_pq_params.build_params.metric = params.metric; - } - build_knn_graph(res, dataset, knn_graph->view(), ivf_pq_params); - } else { - auto nn_descent_params = - std::get(knn_build_params); + if (nn_descent_params.metric != params.metric) { + RAFT_LOG_WARN( + "Metric (%lu) for nn-descent needs to match cagra metric (%lu), " + "aligning nn-descent metric.", + nn_descent_params.metric, + params.metric); + nn_descent_params.metric = params.metric; + } + if (nn_descent_params.graph_degree != intermediate_degree) { + RAFT_LOG_WARN( + "Graph degree (%lu) for nn-descent needs to match cagra intermediate graph degree (%lu), " + "aligning " + "nn-descent graph_degree.", + nn_descent_params.graph_degree, + intermediate_degree); + nn_descent_params = + cagra::graph_build_params::nn_descent_params(intermediate_degree, params.metric); + } - if (nn_descent_params.metric != params.metric) { - RAFT_LOG_WARN( - "Metric (%lu) for nn-descent needs to match cagra metric (%lu), " - "aligning nn-descent metric.", - nn_descent_params.metric, - params.metric); - nn_descent_params.metric = params.metric; - } - if (nn_descent_params.graph_degree != intermediate_degree) { - RAFT_LOG_WARN( - "Graph degree (%lu) for nn-descent needs to match cagra intermediate graph degree (%lu), " - "aligning " - "nn-descent graph_degree.", - nn_descent_params.graph_degree, - intermediate_degree); - nn_descent_params = - cagra::graph_build_params::nn_descent_params(intermediate_degree, params.metric); - } + nn_descent_params.return_distances = false; + build_knn_graph(res, knn_graph_dataset, knn_graph->view(), nn_descent_params); + } - // Use nn-descent to build CAGRA knn graph - nn_descent_params.return_distances = false; - build_knn_graph(res, dataset, knn_graph->view(), nn_descent_params); - } + auto cagra_graph = raft::make_host_matrix(n_rows, graph_degree); - cagra_graph = raft::make_host_matrix(dataset.extent(0), graph_degree); + RAFT_LOG_TRACE("optimizing graph"); + optimize(res, knn_graph->view(), cagra_graph.view(), params.guarantee_connectivity); - RAFT_LOG_TRACE("optimizing graph"); - optimize(res, knn_graph->view(), cagra_graph.view(), params.guarantee_connectivity); + knn_graph.reset(); + return cagra_graph; +} - // free intermediate graph before trying to create the index - knn_graph.reset(); - } +/** + * Build from a host row-major matrix without uploading the full dataset early when IVF-PQ graph + * construction can consume host batches directly. The iterative path uploads and pads inside + * `iterative_build_graph`. The returned index contains only the optimized graph; call + * `index::update_dataset` with a device dataset view before search. + */ +template + requires cuvs::neighbors::is_host_dataset_view_v +auto build_from_host_matrix(raft::resources const& res, + const index_params& params, + DatasetViewT const& dataset) + -> cuvs::neighbors::cagra::index +{ + size_t const n_rows = static_cast(dataset.n_rows()); + size_t const dim = static_cast(dataset.dim()); + + size_t intermediate_degree = params.intermediate_graph_degree; + size_t graph_degree = params.graph_degree; + common::nvtx::range function_scope( + "cagra::detail::build_from_host_matrix(%zu, %zu)", intermediate_degree, graph_degree); + check_graph_degree(intermediate_degree, graph_degree, n_rows); + + auto dataset_extents = + raft::matrix_extent(static_cast(n_rows), static_cast(dim)); + + auto knn_build_params = resolve_cagra_default_knn_graph_build_params( + res, params, dataset_extents, intermediate_degree); + validate_cagra_knn_graph_build_constraints(params, knn_build_params); + + auto cagra_graph = [&]() -> raft::host_matrix { + if (std::holds_alternative( + knn_build_params)) { + return iterative_build_graph(res, params, dataset); + } + return build_cagra_host_graph_from_knn_params(res, + params, + knn_build_params, + static_cast(n_rows), + intermediate_degree, + graph_degree, + dataset.view()); + }(); RAFT_LOG_TRACE("Graph optimized, creating index"); - // Construct an index from dataset and optimized knn graph. - if (params.compression.has_value()) { - RAFT_EXPECTS(params.metric == cuvs::distance::DistanceType::L2Expanded, - "VPQ compression is only supported with L2Expanded distance mertric"); - index idx(res, params.metric); - idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view())); - idx.update_dataset( - res, - // TODO: hardcoding codebook math to `half`, we can do runtime dispatching later - cuvs::preprocessing::quantize::pq::vpq_build(res, *params.compression, dataset)); + cuvs::neighbors::cagra::index out(res, params.metric); + out.update_graph(res, raft::make_const_mdspan(cagra_graph.view())); + return out; +} - return idx; - } - if (params.attach_dataset_on_build) { - try { - return index( - res, params.metric, dataset, raft::make_const_mdspan(cagra_graph.view())); - } catch (std::bad_alloc& e) { - RAFT_LOG_WARN( - "Insufficient GPU memory to construct CAGRA index with dataset on GPU. Only the graph will " - "be added to the index"); - // We just add the graph. User is expected to update dataset separately (e.g allocating in - // managed memory). - } catch (raft::logic_error& e) { - // The memory error can also manifest as logic_error. - RAFT_LOG_WARN( - "Insufficient GPU memory to construct CAGRA index with dataset on GPU. Only the graph will " - "be added to the index"); +/** + * Build from a dense device `dataset_view` (padded or standard). VPQ views are rejected by + * `cagra::build()` before this entry point is reached. Also used from ACE sub-builds and merge. + * The returned index contains only the optimized graph; call `index::update_dataset` before search. + */ +template + requires cuvs::neighbors::is_dense_row_major_device_dataset_view_v +auto build_from_device_matrix(raft::resources const& res, + const index_params& params, + DatasetViewT const& device_dataset) + -> cuvs::neighbors::cagra::index +{ + size_t intermediate_degree = params.intermediate_graph_degree; + size_t graph_degree = params.graph_degree; + common::nvtx::range function_scope( + "cagra::detail::build_from_device_matrix(%zu, %zu)", intermediate_degree, graph_degree); + check_graph_degree( + intermediate_degree, graph_degree, static_cast(device_dataset.n_rows())); + + auto dataset_extents = + raft::matrix_extent(device_dataset.n_rows(), device_dataset.dim()); + + auto knn_build_params = resolve_cagra_default_knn_graph_build_params( + res, params, dataset_extents, intermediate_degree); + validate_cagra_knn_graph_build_constraints(params, knn_build_params); + + auto cagra_graph = [&]() -> raft::host_matrix { + if (std::holds_alternative( + knn_build_params)) { + return iterative_build_graph(res, params, device_dataset); } - } - index idx(res, params.metric); + return build_cagra_host_graph_from_knn_params(res, + params, + knn_build_params, + device_dataset.n_rows(), + intermediate_degree, + graph_degree, + device_dataset.view()); + }(); + + RAFT_LOG_TRACE("Graph optimized, creating index"); + + cuvs::neighbors::cagra::index idx(res, params.metric); idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view())); return idx; } diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index 1dd4cbe075..695888c8a5 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -1,11 +1,13 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include +#include "cagra_build.cuh" + #include #include #include @@ -18,25 +20,39 @@ #include #include +#include #include #include #include -#include -#include #include namespace cuvs::neighbors::cagra::detail { -template -index merge(raft::resources const& handle, - const cagra::index_params& params, - std::vector*>& indices, - const cuvs::neighbors::filtering::base_filter& row_filter) +template +auto make_merged_dense_device_dataset_view( + raft::device_matrix_view storage, uint32_t dim) -> DatasetViewT +{ + if constexpr (cuvs::neighbors::is_padded_dataset_view_v) { + return cuvs::neighbors::device_padded_dataset_view(raft::make_const_mdspan(storage), + dim); + } else if constexpr (cuvs::neighbors::is_standard_dataset_view_v) { + return cuvs::neighbors::device_standard_dataset_view( + raft::make_const_mdspan(storage), dim); + } else { + static_assert(sizeof(DatasetViewT) == 0, + "cagra::merge supports dense padded or standard dataset views only"); + } +} + +template +merged_dataset compute_merged_dataset_layout( + raft::resources const& handle, + std::vector*> const& indices, + cuvs::neighbors::filtering::base_filter const& row_filter) { - using cagra_index_t = cuvs::neighbors::cagra::index; - using ds_idx_type = typename cagra_index_t::dataset_index_type; + using cagra_index_t = cuvs::neighbors::cagra::index; std::size_t dim = 0; std::size_t new_dataset_size = 0; @@ -48,131 +64,183 @@ index merge(raft::resources const& handle, for (cagra_index_t* index : indices) { RAFT_EXPECTS(index != nullptr, "Null pointer detected in 'indices'. Ensure all elements are valid before usage."); - if (auto* strided_dset = dynamic_cast*>(&index->data()); - strided_dset != nullptr) { + auto const& v = index->dataset(); + if constexpr (cuvs::neighbors::is_dense_row_major_dataset_view_v>) { + if (v.n_rows() == 0) { + RAFT_FAIL( + "cagra::merge only supports an index to which the dataset is attached. Please check if " + "the index has an empty dataset; attach one with update_dataset before merge."); + } if (dim == 0) { dim = index->dim(); - stride = strided_dset->stride(); + stride = static_cast(v.stride()); } else { RAFT_EXPECTS(dim == index->dim(), "Dimension of datasets in indices must be equal."); + RAFT_EXPECTS(stride == static_cast(v.stride()), + "Row stride of datasets in indices must be equal."); } new_dataset_size += index->size(); - } else if (dynamic_cast*>(&index->data()) != - nullptr) { - RAFT_FAIL( - "cagra::merge only supports an index to which the dataset is attached. Please check if the " - "index was built with index_param.attach_dataset_on_build = true, or if a dataset was " - "attached after the build."); } else { - RAFT_FAIL("cagra::merge only supports an uncompressed dataset index"); + RAFT_FAIL("cagra::merge only supports an uncompressed dense device dataset index"); } } - IdxT offset = 0; + merged_dataset layout{}; + layout.merged_rows = static_cast(new_dataset_size); + layout.stride_elements = stride; + layout.dim = static_cast(dim); + layout.bitset_filtered = + (row_filter.get_filter_type() == cuvs::neighbors::filtering::FilterType::Bitset); + if (layout.bitset_filtered) { + auto const& actual_filter = + dynamic_cast&>(row_filter); + layout.filtered_rows = actual_filter.view().count(handle); + } else { + layout.filtered_rows = layout.merged_rows; + } + return layout; +} - auto merge_dataset = [&](T* dst) { +template +cuvs::neighbors::cagra::index merge( + raft::resources const& handle, + const cagra::index_params& params, + std::vector*>& indices, + merged_dataset_storage& storage, + const cuvs::neighbors::filtering::base_filter& row_filter) +{ + using cagra_index_t = cuvs::neighbors::cagra::index; + + auto const expected = + compute_merged_dataset_layout(handle, indices, row_filter); + RAFT_EXPECTS(expected.merged_rows == storage.layout.merged_rows && + expected.filtered_rows == storage.layout.filtered_rows && + expected.stride_elements == storage.layout.stride_elements && + expected.dim == storage.layout.dim && + expected.bitset_filtered == storage.layout.bitset_filtered, + "merged_dataset_storage.layout does not match indices and row_filter (use the same " + "arguments as " + "make_merged_storage)."); + + auto merged_storage = storage.merged_storage.view(); + RAFT_EXPECTS(merged_storage.extent(0) == storage.layout.merged_rows, + "merged_storage rows (%ld) must equal layout.merged_rows (%ld)", + long(merged_storage.extent(0)), + long(storage.layout.merged_rows)); + RAFT_EXPECTS(merged_storage.extent(1) == storage.layout.stride_elements, + "merged_storage stride (%ld) must equal layout.stride_elements (%ld)", + long(merged_storage.extent(1)), + long(storage.layout.stride_elements)); + + std::optional> filtered_view{}; + if (storage.layout.bitset_filtered) { + RAFT_EXPECTS(storage.filtered_storage.has_value(), + "Bitset-filtered merge requires merged_dataset_storage.filtered_storage."); + filtered_view = storage.filtered_storage->view(); + RAFT_EXPECTS(filtered_view->extent(0) == storage.layout.filtered_rows, + "filtered_storage rows (%ld) must equal layout.filtered_rows (%ld)", + long(filtered_view->extent(0)), + long(storage.layout.filtered_rows)); + RAFT_EXPECTS(filtered_view->extent(1) == storage.layout.stride_elements, + "filtered_storage stride (%ld) must equal layout.stride_elements (%ld)", + long(filtered_view->extent(1)), + long(storage.layout.stride_elements)); + } else { + RAFT_EXPECTS(!storage.filtered_storage.has_value(), + "Non-bitset merge requires merged_dataset_storage.filtered_storage be unset."); + } + + auto merge_dataset = [&](T* dst, std::size_t dst_ld) { + IdxT row_offset = 0; for (cagra_index_t* index : indices) { - auto* strided_dset = dynamic_cast*>(&index->data()); - raft::copy_matrix(dst + offset * dim, - dim, - strided_dset->view().data_handle(), - static_cast(stride), - dim, - static_cast(strided_dset->n_rows()), + 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(storage.layout.stride_elements), + static_cast(storage.layout.dim), + n_rows, raft::resource::get_cuda_stream(handle)); - offset += IdxT(index->data().n_rows()); + row_offset += IdxT(index->dataset().n_rows()); } }; - try { - auto updated_dataset = - raft::make_device_matrix(handle, int64_t(new_dataset_size), int64_t(dim)); - - merge_dataset(updated_dataset.data_handle()); - - if (row_filter.get_filter_type() == cuvs::neighbors::filtering::FilterType::Bitset) { - auto actual_filter = - dynamic_cast&>( - row_filter); - auto filtered_row_count = actual_filter.view().count(handle); - - // Convert the filter to a CSR matrix (so that we can pass indices to raft::copy_rows) - auto indices_csr = raft::make_device_csr_matrix( - handle, 1, new_dataset_size); - indices_csr.initialize_sparsity(filtered_row_count); - - actual_filter.view().to_csr(handle, indices_csr); - - // Get the indices array from the csr matrix. Note that this returns a raft::span object - // and we need to pass as device_vector_view, which is a 1D mdspan (instead of a span) - // so we need to translate here (and adjust to be const) - auto indices = indices_csr.structure_view().get_indices(); - auto indices_view = raft::make_device_vector_view( - indices.data(), static_cast(indices.size())); - - auto filtered_dataset = raft::make_device_matrix(handle, filtered_row_count, dim); - raft::matrix::copy_rows(handle, - raft::make_const_mdspan(updated_dataset.view()), - filtered_dataset.view(), - indices_view); - - auto merged_index = - cagra::build(handle, params, raft::make_const_mdspan(filtered_dataset.view())); - if (!merged_index.data().is_owning() && params.attach_dataset_on_build) { - using matrix_t = decltype(updated_dataset); - using layout_t = typename matrix_t::layout_type; - using container_policy_t = typename matrix_t::container_policy_type; - using owning_t = owning_dataset; - auto out_layout = raft::make_strided_layout(filtered_dataset.view().extents(), - cuda::std::array{stride, 1}); - - merged_index.update_dataset(handle, owning_t{std::move(filtered_dataset), out_layout}); - } - RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); - return merged_index; - } else { - auto merged_index = - cagra::build(handle, params, raft::make_const_mdspan(updated_dataset.view())); - if (!merged_index.data().is_owning() && params.attach_dataset_on_build) { - using matrix_t = decltype(updated_dataset); - using layout_t = typename matrix_t::layout_type; - using container_policy_t = typename matrix_t::container_policy_type; - using owning_t = owning_dataset; - auto out_layout = raft::make_strided_layout(updated_dataset.view().extents(), - cuda::std::array{stride, 1}); - - merged_index.update_dataset(handle, owning_t{std::move(updated_dataset), out_layout}); - } - RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); - return merged_index; - } - } catch (std::bad_alloc& e) { - // We don't currently support the cpu memory fallback with filtered merge, since the - // 'raft::matrix::copy_rows' only supports gpu memory - RAFT_EXPECTS(row_filter.get_filter_type() == cuvs::neighbors::filtering::FilterType::None, - "Filtered merge isn't available on cpu memory"); - - RAFT_LOG_DEBUG("cagra::merge: using host memory for merged dataset"); - - auto updated_dataset = - raft::make_host_matrix(std::int64_t(new_dataset_size), std::int64_t(dim)); - - merge_dataset(updated_dataset.data_handle()); - - auto merged_index = - cagra::build(handle, params, raft::make_const_mdspan(updated_dataset.view())); - if (!merged_index.data().is_owning() && params.attach_dataset_on_build) { - using matrix_t = decltype(updated_dataset); - using layout_t = typename matrix_t::layout_type; - using container_policy_t = typename matrix_t::container_policy_type; - using owning_t = owning_dataset; - auto out_layout = raft::make_strided_layout(updated_dataset.view().extents(), - cuda::std::array{stride, 1}); - merged_index.update_dataset(handle, owning_t{std::move(updated_dataset), out_layout}); - } - return merged_index; + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + const auto merged_bytes = static_cast(merged_storage.size()) * sizeof(T); + RAFT_CUDA_TRY(cudaMemsetAsync(merged_storage.data_handle(), 0, merged_bytes, stream)); + + merge_dataset(merged_storage.data_handle(), + static_cast(storage.layout.stride_elements)); + + if (storage.layout.bitset_filtered) { + auto actual_filter = + dynamic_cast&>(row_filter); + + auto indices_csr = raft::make_device_csr_matrix( + handle, 1, static_cast(storage.layout.merged_rows)); + indices_csr.initialize_sparsity(storage.layout.filtered_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())); + + auto& filtered_storage = *filtered_view; + RAFT_CUDA_TRY(cudaMemsetAsync(filtered_storage.data_handle(), + 0, + static_cast(filtered_storage.size()) * sizeof(T), + stream)); + + raft::matrix::copy_rows( + handle, raft::make_const_mdspan(merged_storage), filtered_storage, indices_view); + + auto dv = make_merged_dense_device_dataset_view( + raft::make_const_mdspan(filtered_storage), storage.layout.dim); + auto index = ::cuvs::neighbors::cagra::detail::build_from_device_matrix( + handle, params, dv); + index.update_dataset(handle, dv); + RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); + return index; } + + auto dv = make_merged_dense_device_dataset_view( + raft::make_const_mdspan(merged_storage), storage.layout.dim); + auto index = ::cuvs::neighbors::cagra::detail::build_from_device_matrix( + handle, params, dv); + index.update_dataset(handle, dv); + RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); + return index; } } // namespace cuvs::neighbors::cagra::detail + +namespace cuvs::neighbors::cagra { + +template +merged_dataset_storage make_merged_storage( + raft::resources const& res, + std::vector*> const& indices, + cuvs::neighbors::filtering::base_filter const& row_filter) +{ + merged_dataset layout = detail::compute_merged_dataset_layout(res, indices, row_filter); + auto merged_storage = + raft::make_device_matrix(res, layout.merged_rows, layout.stride_elements); + std::optional> filtered_storage; + if (layout.bitset_filtered) { + filtered_storage.emplace( + raft::make_device_matrix(res, layout.filtered_rows, layout.stride_elements)); + } + return {layout, std::move(merged_storage), std::move(filtered_storage)}; +} + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/detail/cagra/cagra_search.cuh b/cpp/src/neighbors/detail/cagra/cagra_search.cuh index 4d09e3683b..23bd741302 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_search.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_search.cuh @@ -19,7 +19,6 @@ #include #include - // TODO: Fix these when ivf methods are moved over #include "../../ivf_common.cuh" #include "../../ivf_pq/ivf_pq_search.cuh" @@ -83,31 +82,76 @@ void search_main_core( RAFT_LOG_DEBUG("Cagra search"); const uint32_t max_queries = plan->max_queries; - const uint32_t query_dim = queries.extent(1); + const uint32_t query_dim = static_cast(queries.extent(1)); + // Same 16B row-pitch rule as make_device_padded_dataset. Tight [n,dim] rows can be misaligned + // between rows (e.g. float, dim=1) and trigger misaligned access in CAGRA search. If + // query_row_stride>dim, device code still advances with "+= dim*query_id" in setup_workspace; in + // that case run one query per plan call so every kernel sees query_id==0 and the base pointer + // selects the row (keeps batched path when stride==dim). + const DataT* queries_buf{}; + uint32_t query_row_stride{}; + std::unique_ptr> queries_padded_own; + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(queries)) { + auto v = cuvs::neighbors::make_device_padded_dataset_view(res, queries); + queries_buf = v.view().data_handle(); + query_row_stride = v.stride(); + } else { + queries_padded_own = cuvs::neighbors::make_device_padded_dataset(res, queries); + auto v = queries_padded_own->as_dataset_view(); + queries_buf = v.view().data_handle(); + query_row_stride = v.stride(); + } + const bool can_batch_n_queries = (query_row_stride == query_dim); for (unsigned qid = 0; qid < queries.extent(0); qid += max_queries) { const uint32_t n_queries = std::min(max_queries, queries.extent(0) - qid); - auto _topk_indices_ptr = neighbors.data_handle() + (topk * qid); - auto _topk_distances_ptr = distances.data_handle() + (topk * qid); - // todo(tfeher): one could keep distances optional and pass nullptr - const auto* _query_ptr = queries.data_handle() + (query_dim * qid); - const auto* _seed_ptr = - plan->num_seeds > 0 - ? reinterpret_cast(plan->dev_seed.data()) + (plan->num_seeds * qid) - : nullptr; - uint32_t* _num_executed_iterations = nullptr; + if (can_batch_n_queries) { + auto _topk_indices_ptr = neighbors.data_handle() + (topk * qid); + auto _topk_distances_ptr = distances.data_handle() + (topk * qid); + const auto* _query_ptr = + queries_buf + (static_cast(query_row_stride) * static_cast(qid)); + const auto* _seed_ptr = + plan->num_seeds > 0 + ? reinterpret_cast(plan->dev_seed.data()) + (plan->num_seeds * qid) + : nullptr; + uint32_t* _num_executed_iterations = nullptr; - (*plan)(res, - graph, - source_indices, - _topk_indices_ptr, - _topk_distances_ptr, - _query_ptr, - n_queries, - _seed_ptr, - _num_executed_iterations, - topk, - set_offset(sample_filter, qid)); + (*plan)(res, + graph, + source_indices, + _topk_indices_ptr, + _topk_distances_ptr, + _query_ptr, + n_queries, + _seed_ptr, + _num_executed_iterations, + topk, + set_offset(sample_filter, qid)); + } else { + for (uint32_t qi = 0; qi < n_queries; ++qi) { + const size_t g = static_cast(qid) + static_cast(qi); + auto _topk_indices_ptr = neighbors.data_handle() + (topk * g); + auto _topk_distances_ptr = distances.data_handle() + (topk * g); + const auto* _query_ptr = queries_buf + (query_row_stride * g); + const auto* _seed_ptr = + plan->num_seeds > 0 + ? reinterpret_cast(plan->dev_seed.data()) + (plan->num_seeds * g) + : nullptr; + uint32_t* _num_executed_iterations = nullptr; + + (*plan)(res, + graph, + source_indices, + _topk_indices_ptr, + _topk_distances_ptr, + _query_ptr, + 1u, + _seed_ptr, + _num_executed_iterations, + topk, + set_offset(sample_filter, g)); + } + } } } @@ -133,10 +177,11 @@ template + typename DistanceT = float, + cuvs::neighbors::ann_dataset_view DatasetViewT> void search_main(raft::resources const& res, search_params params, - const index& index, + const index& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, raft::device_matrix_view distances, @@ -147,12 +192,9 @@ void search_main(raft::resources const& res, "Use cuvs::neighbors::hnsw::from_cagra() to convert the index and " "cuvs::neighbors::hnsw::deserialize() to load it into memory before searching."); - // n_rows has the same type as the dataset index (the array extents type) - using ds_idx_type = decltype(index.data().n_rows()); using graph_idx_type = uint32_t; - // Dispatch search parameters based on the dataset kind. - if (auto* strided_dset = dynamic_cast*>(&index.data()); - strided_dset != nullptr) { + + auto run_strided_like = [&](auto const& row_dataset) { if (params.smem_dtype != cuvs::neighbors::cagra::internal_dtype::F16) { RAFT_LOG_WARN("In this search mode, smem_dtype supports only F16. Set it to F16."); params.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; @@ -167,7 +209,7 @@ void search_main(raft::resources const& res, dataset_norms_ptr = index.dataset_norms().value().data_handle(); } auto desc = dataset_descriptor_init_with_cache( - res, params, *strided_dset, index.metric(), dataset_norms_ptr); + res, params, row_dataset, index.metric(), dataset_norms_ptr); search_main_core( res, params, @@ -178,12 +220,15 @@ void search_main(raft::resources const& res, neighbors, distances, sample_filter); - } else if (auto* vpq_dset = dynamic_cast*>(&index.data()); - vpq_dset != nullptr) { - // Search using a compressed dataset + }; + + if constexpr (cuvs::neighbors::is_empty_dataset_view_v) { + RAFT_FAIL( + "Attempted to search without a dataset. Please call index.update_dataset(...) first."); + } else if constexpr (cuvs::neighbors::is_device_vpq_f32_dataset_view_v) { RAFT_FAIL("FP32 VPQ dataset support is coming soon"); - } else if (auto* vpq_dset = dynamic_cast*>(&index.data()); - vpq_dset != nullptr) { + } else if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { + auto const& vv = index.dataset(); if (params.smem_dtype == cuvs::neighbors::cagra::internal_dtype::E5M2 && raft::getComputeCapability().first < 9) { RAFT_LOG_WARN( @@ -191,7 +236,7 @@ void search_main(raft::resources const& res, params.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; } auto desc = dataset_descriptor_init_with_cache( - res, params, *vpq_dset, index.metric(), nullptr); + res, params, vv.dset(), index.metric(), nullptr); search_main_core( res, params, @@ -202,14 +247,20 @@ void search_main(raft::resources const& res, neighbors, distances, sample_filter); - } else if (auto* empty_dset = dynamic_cast*>(&index.data()); - empty_dset != nullptr) { - // Forgot to add a dataset. + } else if constexpr (cuvs::neighbors::is_device_standard_dataset_view_v) { RAFT_FAIL( - "Attempted to search without a dataset. Please call index.update_dataset(...) first."); + "CAGRA search requires a padded device dataset. Build from a standard dataset view, then " + "call " + "cagra::attach_padded_dataset_for_search(res, index, padded_view) before search."); + } else if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + run_strided_like(index.dataset()); + } else if constexpr (cuvs::neighbors::is_host_dataset_view_v) { + static_assert(sizeof(DatasetViewT) == 0, + "search requires a device-resident dataset. " + "Call cagra::attach_device_dataset_on_host_index(res, host_idx, device_view) " + "to convert the host index and attach a device dataset before searching."); } else { - // This is a logic error. - RAFT_FAIL("Unrecognized dataset format"); + static_assert(sizeof(DatasetViewT) == 0, "search: unsupported dataset view type"); } static_assert(std::is_same_v, diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index f106b82500..311c79e5ca 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -28,6 +28,23 @@ namespace cuvs::neighbors::cagra::detail { +template +inline constexpr bool is_cagra_hnsw_serialize_index_v = + std::is_same_v> || + std::is_same_v> || + std::is_same_v> || + std::is_same_v>; + +template +inline constexpr bool is_device_cagra_hnsw_serialize_index_v = + std::is_same_v> || + std::is_same_v>; + +template +inline constexpr bool is_host_cagra_hnsw_serialize_index_v = + std::is_same_v> || + std::is_same_v>; + constexpr int serialization_version = 5; /** @@ -40,10 +57,10 @@ constexpr int serialization_version = 5; * @param[in] index_ CAGRA index * */ -template +template void serialize(raft::resources const& res, std::ostream& os, - const index& index_, + const cuvs::neighbors::cagra::index& index_, bool include_dataset) { raft::common::nvtx::range fun_scope("cagra::serialize"); @@ -68,14 +85,23 @@ void serialize(raft::resources const& res, raft::serialize_mdspan(res, os, index_.graph()); - include_dataset &= (index_.data().n_rows() > 0); + include_dataset &= (index_.dataset().n_rows() > 0); bool has_source_indices = index_.source_indices().has_value(); uint32_t content_map = 0x1u * include_dataset + 0x2u * has_source_indices; raft::serialize_scalar(res, os, content_map); if (include_dataset) { RAFT_LOG_DEBUG("Saving CAGRA index with dataset"); - neighbors::detail::serialize(res, os, index_.data()); + if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v || + cuvs::neighbors::is_device_standard_dataset_view_v) { + neighbors::detail::serialize_cagra_padded_dataset(res, os, index_.dataset()); + } else { + // Future dataset types (e.g. VPQ) require a new branch here and a corresponding + // deserialize overload. Use static_assert to catch unsupported types at compile time. + static_assert( + sizeof(DatasetViewT) == 0, + "serialize: dataset serialization is not yet implemented for this DatasetViewT"); + } } else { RAFT_LOG_DEBUG("Saving CAGRA index WITHOUT dataset"); } @@ -83,10 +109,10 @@ void serialize(raft::resources const& res, if (has_source_indices) { raft::serialize_mdspan(res, os, index_.source_indices().value()); } } -template +template void serialize(raft::resources const& res, const std::string& filename, - const index& index_, + const cuvs::neighbors::cagra::index& index_, bool include_dataset) { RAFT_EXPECTS(!index_.dataset_fd().has_value(), @@ -102,15 +128,16 @@ void serialize(raft::resources const& res, if (!of) { RAFT_FAIL("Error writing output %s", filename.c_str()); } } -template +template void serialize_to_hnswlib( raft::resources const& res, std::ostream& os, - const cuvs::neighbors::cagra::index& index_, + CagraIndexT const& index_, std::optional> dataset) { - // static_assert(std::is_same_v or std::is_same_v, - // "An hnswlib index can only be trained with int32 or uint32 IdxT"); + static_assert(is_cagra_hnsw_serialize_index_v, + "serialize_to_hnswlib requires a dense device or host padded CAGRA index"); + int dim = (dataset) ? dataset->extent(1) : index_.dim(); raft::common::nvtx::range fun_scope("cagra::serialize"); RAFT_LOG_DEBUG("Saving CAGRA index to hnswlib format, size %zu, dim %u", @@ -165,22 +192,27 @@ void serialize_to_hnswlib( raft::host_matrix_view host_dataset_view; if (dataset) { host_dataset_view = *dataset; - } else { - auto dataset = index_.dataset(); - RAFT_EXPECTS(dataset.size() > 0, + } else if constexpr (is_device_cagra_hnsw_serialize_index_v) { + auto dataset_view = index_.dataset(); + RAFT_EXPECTS(dataset_view.n_rows() > 0, "Invalid CAGRA dataset of size 0 during serialization, shape %zux%zu", - static_cast(dataset.extent(0)), - static_cast(dataset.extent(1))); - host_dataset = raft::make_host_matrix(dataset.extent(0), dataset.extent(1)); + static_cast(dataset_view.n_rows()), + static_cast(dataset_view.dim())); + host_dataset = raft::make_host_matrix(dataset_view.n_rows(), dataset_view.dim()); raft::copy_matrix(host_dataset.data_handle(), host_dataset.extent(1), - dataset.data_handle(), - dataset.stride(0), + dataset_view.view().data_handle(), + dataset_view.stride(), host_dataset.extent(1), - dataset.extent(0), + dataset_view.n_rows(), raft::resource::get_cuda_stream(res)); raft::resource::sync_stream(res); host_dataset_view = raft::make_const_mdspan(host_dataset.view()); + } else if constexpr (is_host_cagra_hnsw_serialize_index_v) { + RAFT_FAIL("serialize_to_hnswlib requires dataset for host CAGRA index"); + } else { + static_assert(is_cagra_hnsw_serialize_index_v, + "serialize_to_hnswlib: unsupported CagraIndexT"); } auto graph = index_.graph(); auto host_graph = @@ -239,11 +271,11 @@ void serialize_to_hnswlib( } } -template +template void serialize_to_hnswlib( raft::resources const& res, const std::string& filename, - const cuvs::neighbors::cagra::index& index_, + CagraIndexT const& index_, std::optional> dataset) { std::ofstream of(filename, std::ios::out | std::ios::binary); @@ -264,8 +296,12 @@ void serialize_to_hnswlib( * @param[in] index_ CAGRA index * */ -template -void deserialize(raft::resources const& res, std::istream& is, index* index_) +template +void deserialize( + raft::resources const& res, + std::istream& is, + cuvs::neighbors::cagra::index* index_, + std::unique_ptr>* out_dataset = nullptr) { raft::common::nvtx::range fun_scope("cagra::deserialize"); @@ -302,13 +338,25 @@ void deserialize(raft::resources const& res, std::istream& is, index* i auto graph = raft::make_host_matrix(n_rows, graph_degree); deserialize_mdspan(res, is, graph.view()); - *index_ = index(res, metric); + *index_ = cuvs::neighbors::cagra::index(res, metric); index_->update_graph(res, raft::make_const_mdspan(graph.view())); auto content_map = raft::deserialize_scalar(res, is); bool has_dataset = content_map & 0x1u; if (has_dataset) { - index_->update_dataset(res, cuvs::neighbors::detail::deserialize_dataset(res, is)); + RAFT_EXPECTS(out_dataset != nullptr, + "deserialize: index contains a dataset; pass a non-null out_dataset to own it."); + if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + *out_dataset = cuvs::neighbors::detail::deserialize_dataset(res, is); + index_->update_dataset(res, (*out_dataset)->as_dataset_view()); + } else if constexpr (cuvs::neighbors::is_device_standard_dataset_view_v) { + *out_dataset = cuvs::neighbors::detail::deserialize_standard_dataset(res, is); + index_->update_dataset(res, (*out_dataset)->as_dataset_view()); + } else { + static_assert(sizeof(DatasetViewT) == 0, + "deserialize: dataset deserialization is not yet implemented for this " + "DatasetViewT"); + } } bool has_source_indices = content_map & 0x2u; @@ -321,14 +369,18 @@ void deserialize(raft::resources const& res, std::istream& is, index* i } } -template -void deserialize(raft::resources const& res, const std::string& filename, index* index_) +template +void deserialize( + raft::resources const& res, + const std::string& filename, + cuvs::neighbors::cagra::index* index_, + std::unique_ptr>* out_dataset = nullptr) { std::ifstream is(filename, std::ios::in | std::ios::binary); if (!is) { RAFT_FAIL("Cannot open file %s", filename.c_str()); } - detail::deserialize(res, is, index_); + detail::deserialize(res, is, index_, out_dataset); is.close(); } diff --git a/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp b/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp index ef82b1760e..5f7c8efb10 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp +++ b/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -28,7 +28,7 @@ struct standard_descriptor_spec : public instance_spec template constexpr static inline bool accepts_dataset() { - return is_strided_dataset_v; + return is_padded_dataset_v; } template diff --git a/cpp/src/neighbors/detail/cagra/factory.cuh b/cpp/src/neighbors/detail/cagra/factory.cuh index 1ed1f33e97..cdcc18867b 100644 --- a/cpp/src/neighbors/detail/cagra/factory.cuh +++ b/cpp/src/neighbors/detail/cagra/factory.cuh @@ -94,7 +94,7 @@ template auto make_key(const cagra::search_params& params, const DatasetT& dataset, cuvs::distance::DistanceType metric) - -> std::enable_if_t, key> + -> std::enable_if_t, key> { return key{reinterpret_cast(dataset.view().data_handle()), uint64_t(dataset.n_rows()), diff --git a/cpp/src/neighbors/detail/dataset_serialize.hpp b/cpp/src/neighbors/detail/dataset_serialize.hpp index 00032ae9d2..a565e20afc 100644 --- a/cpp/src/neighbors/detail/dataset_serialize.hpp +++ b/cpp/src/neighbors/detail/dataset_serialize.hpp @@ -1,14 +1,14 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include - #include #include #include +#include #include #include @@ -25,16 +25,9 @@ constexpr dataset_instance_tag kSerializeEmptyDataset = 1; constexpr dataset_instance_tag kSerializeStridedDataset = 2; constexpr dataset_instance_tag kSerializeVPQDataset = 3; -template -void serialize(const raft::resources& res, std::ostream& os, const empty_dataset& dataset) -{ - raft::serialize_scalar(res, os, dataset.suggested_dim); -} - -template -void serialize(const raft::resources& res, - std::ostream& os, - const strided_dataset& dataset) +template + requires cuvs::neighbors::is_dense_row_major_device_dataset_view_v +void serialize(const raft::resources& res, std::ostream& os, ViewT const& dataset) { auto n_rows = dataset.n_rows(); auto dim = dataset.dim(); @@ -42,7 +35,6 @@ void serialize(const raft::resources& res, raft::serialize_scalar(res, os, n_rows); raft::serialize_scalar(res, os, dim); raft::serialize_scalar(res, os, stride); - // Remove padding before saving the dataset auto src = dataset.view(); auto dst = raft::make_host_matrix(n_rows, dim); raft::copy_matrix(dst.data_handle(), @@ -56,85 +48,74 @@ void serialize(const raft::resources& res, raft::serialize_mdspan(res, os, dst.view()); } -template -void serialize(const raft::resources& res, - std::ostream& os, - const vpq_dataset& dataset) -{ - raft::serialize_scalar(res, os, dataset.n_rows()); - raft::serialize_scalar(res, os, dataset.dim()); - raft::serialize_scalar(res, os, dataset.vq_n_centers()); - raft::serialize_scalar(res, os, dataset.pq_n_centers()); - raft::serialize_scalar(res, os, dataset.pq_len()); - raft::serialize_scalar(res, os, dataset.encoded_row_length()); - raft::serialize_mdspan(res, os, make_const_mdspan(dataset.vq_code_book.view())); - raft::serialize_mdspan(res, os, make_const_mdspan(dataset.pq_code_book.view())); - raft::serialize_mdspan(res, os, make_const_mdspan(dataset.data.view())); -} - -template -void serialize(const raft::resources& res, std::ostream& os, const dataset& dataset) +/** Write CAGRA index dataset blob (tag + element dtype + strided payload). */ +template + requires cuvs::neighbors::is_dense_row_major_device_dataset_view_v +void serialize_cagra_padded_dataset(const raft::resources& res, + std::ostream& os, + ViewT const& dataset) { - if (auto x = dynamic_cast*>(&dataset); x != nullptr) { - raft::serialize_scalar(res, os, kSerializeEmptyDataset); - return serialize(res, os, *x); - } - if (auto x = dynamic_cast*>(&dataset); x != nullptr) { - raft::serialize_scalar(res, os, kSerializeStridedDataset); + raft::serialize_scalar(res, os, kSerializeStridedDataset); + if constexpr (std::is_same_v) { raft::serialize_scalar(res, os, CUDA_R_32F); - return serialize(res, os, *x); - } - if (auto x = dynamic_cast*>(&dataset); x != nullptr) { - raft::serialize_scalar(res, os, kSerializeStridedDataset); + } else if constexpr (std::is_same_v) { raft::serialize_scalar(res, os, CUDA_R_16F); - return serialize(res, os, *x); - } - if (auto x = dynamic_cast*>(&dataset); x != nullptr) { - raft::serialize_scalar(res, os, kSerializeStridedDataset); + } else if constexpr (std::is_same_v) { raft::serialize_scalar(res, os, CUDA_R_8I); - return serialize(res, os, *x); - } - if (auto x = dynamic_cast*>(&dataset); x != nullptr) { - raft::serialize_scalar(res, os, kSerializeStridedDataset); + } else if constexpr (std::is_same_v) { raft::serialize_scalar(res, os, CUDA_R_8U); - return serialize(res, os, *x); + } else { + static_assert(!std::is_same_v, "unsupported element type for CAGRA serialize"); } - if (auto x = dynamic_cast*>(&dataset); x != nullptr) { - raft::serialize_scalar(res, os, kSerializeVPQDataset); - raft::serialize_scalar(res, os, CUDA_R_32F); - return serialize(res, os, *x); - } - if (auto x = dynamic_cast*>(&dataset); x != nullptr) { - raft::serialize_scalar(res, os, kSerializeVPQDataset); - raft::serialize_scalar(res, os, CUDA_R_16F); - return serialize(res, os, *x); - } - RAFT_FAIL("unsupported dataset type."); + serialize(res, os, dataset); } template auto deserialize_empty(raft::resources const& res, std::istream& is) - -> std::unique_ptr> + -> std::unique_ptr> { auto suggested_dim = raft::deserialize_scalar(res, is); - return std::make_unique>(suggested_dim); + return std::make_unique>(suggested_dim); } +/** Read shared dense wire payload: `n_rows`, `dim`, `stride`, then tight `[n_rows x dim]` host + * matrix. */ template -auto deserialize_strided(raft::resources const& res, std::istream& is) - -> std::unique_ptr> +auto deserialize_dense_payload(raft::resources const& res, std::istream& is) + -> std::tuple> { - auto n_rows = raft::deserialize_scalar(res, is); - auto dim = raft::deserialize_scalar(res, is); - auto stride = raft::deserialize_scalar(res, is); + auto n_rows = raft::deserialize_scalar(res, is); + auto dim = raft::deserialize_scalar(res, is); + auto stride = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(dim <= stride, + "deserialize_dense_payload: logical dim (%u) must not exceed row stride (%u).", + static_cast(dim), + static_cast(stride)); auto host_array = raft::make_host_matrix(n_rows, dim); raft::deserialize_mdspan(res, is, host_array.view()); - return make_strided_dataset(res, std::move(host_array), stride); + return {n_rows, dim, stride, std::move(host_array)}; +} + +template +auto deserialize_padded(raft::resources const& res, std::istream& is) + -> std::unique_ptr> +{ + auto payload = deserialize_dense_payload(res, is); + return cuvs::neighbors::make_device_padded_dataset(res, std::get<3>(payload).view()); +} + +template +auto deserialize_standard(raft::resources const& res, std::istream& is) + -> std::unique_ptr> +{ + auto payload = deserialize_dense_payload(res, is); + return cuvs::neighbors::make_device_standard_dataset( + res, std::get<3>(payload).view(), std::get<1>(payload), std::get<2>(payload)); } -template +template auto deserialize_vpq(raft::resources const& res, std::istream& is) - -> std::unique_ptr> + -> std::unique_ptr> { auto n_rows = raft::deserialize_scalar(res, is); auto dim = raft::deserialize_scalar(res, is); @@ -144,9 +125,9 @@ auto deserialize_vpq(raft::resources const& res, std::istream& is) auto encoded_row_length = raft::deserialize_scalar(res, is); auto vq_code_book = - raft::make_device_matrix(res, vq_n_centers, dim); + raft::make_device_matrix(res, vq_n_centers, dim); auto pq_code_book = - raft::make_device_matrix(res, pq_n_centers, pq_len); + raft::make_device_matrix(res, pq_n_centers, pq_len); auto data = raft::make_device_matrix(res, n_rows, encoded_row_length); @@ -154,43 +135,53 @@ auto deserialize_vpq(raft::resources const& res, std::istream& is) raft::deserialize_mdspan(res, is, pq_code_book.view()); raft::deserialize_mdspan(res, is, data.view()); - return std::make_unique>( + return std::make_unique>( std::move(vq_code_book), std::move(pq_code_book), std::move(data)); } -template -auto deserialize_dataset(raft::resources const& res, std::istream& is) - -> std::unique_ptr> +template +auto deserialize_dense_dataset(raft::resources const& res, std::istream& is) + -> std::unique_ptr { const auto tag = raft::deserialize_scalar(res, is); - switch (tag) { - case kSerializeEmptyDataset: return deserialize_empty(res, is); - case kSerializeStridedDataset: { - const auto dtype = raft::deserialize_scalar(res, is); - switch (dtype) { - case CUDA_R_32F: return deserialize_strided(res, is); - case CUDA_R_16F: return deserialize_strided(res, is); - case CUDA_R_8I: return deserialize_strided(res, is); - case CUDA_R_8U: return deserialize_strided(res, is); - default: - RAFT_FAIL("Failed to deserialize dataset: unsupported strided dataset element type %d.", - static_cast(dtype)); - } - } - case kSerializeVPQDataset: { - const auto dtype = raft::deserialize_scalar(res, is); - switch (dtype) { - case CUDA_R_32F: return deserialize_vpq(res, is); - case CUDA_R_16F: return deserialize_vpq(res, is); - default: - RAFT_FAIL("Failed to deserialize dataset: unsupported VPQ dtype %d.", - static_cast(dtype)); - } - } - default: - RAFT_FAIL("Failed to deserialize dataset: unknown instance tag %u.", - static_cast(tag)); + RAFT_EXPECTS(tag == kSerializeStridedDataset, + "deserialize_dataset: expected strided tag, got %u", + static_cast(tag)); + const auto dtype = raft::deserialize_scalar(res, is); + constexpr cudaDataType_t expected_dtype = std::is_same_v ? CUDA_R_32F + : std::is_same_v ? CUDA_R_16F + : std::is_same_v ? CUDA_R_8I + : CUDA_R_8U; // uint8_t + RAFT_EXPECTS(dtype == expected_dtype, + "deserialize_dataset: serialized dtype (%d) does not match expected (%d)", + static_cast(dtype), + static_cast(expected_dtype)); + if constexpr (std::is_same_v>) { + return deserialize_padded(res, is); + } else if constexpr (std::is_same_v>) { + return deserialize_standard(res, is); + } else { + static_assert(!std::is_same_v, + "deserialize_dense_dataset: unsupported owning dataset type"); } } +// Reads tag + dtype prefix, validates they match DataT, and returns a concrete +// device_padded_dataset. When a new dataset kind is supported, add a matching overload of +// deserialize_dataset here rather than extending this one — overload dispatch replaces the old +// type-erased variant routing. +template +auto deserialize_dataset(raft::resources const& res, std::istream& is) + -> std::unique_ptr> +{ + return deserialize_dense_dataset>(res, is); +} + +template +auto deserialize_standard_dataset(raft::resources const& res, std::istream& is) + -> std::unique_ptr> +{ + return deserialize_dense_dataset>(res, is); +} + } // namespace cuvs::neighbors::detail diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 88580de929..8328fac5d5 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -41,6 +41,23 @@ namespace cuvs::neighbors::hnsw::detail { +template +inline constexpr bool is_cagra_hnsw_export_index_v = + std::is_same_v> || + std::is_same_v> || + std::is_same_v> || + std::is_same_v>; + +template +inline constexpr bool is_host_cagra_hnsw_export_index_v = + std::is_same_v> || + std::is_same_v>; + +template +inline constexpr bool is_device_cagra_hnsw_export_index_v = + std::is_same_v> || + std::is_same_v>; + // This is needed as hnswlib hardcodes the distance type to float // or int32_t in certain places. However, we can solve uint8 or int8 // natively with the patch cuVS applies. We could potentially remove @@ -215,12 +232,13 @@ struct index_impl : index { std::optional hnsw_fd_; }; -template -std::enable_if_t>> from_cagra( - raft::resources const& res, - const index_params& params, - const cuvs::neighbors::cagra::index& cagra_index, - std::optional> dataset) +template +std::enable_if_t, + std::unique_ptr>> +from_cagra(raft::resources const& res, + const index_params& params, + CagraIndexT const& cagra_index, + std::optional> dataset) { common::nvtx::range fun_scope("hnsw::from_cagra"); std::random_device dev; @@ -243,12 +261,13 @@ std::enable_if_t>> fr return std::unique_ptr>(hnsw_index); } -template -std::enable_if_t>> from_cagra( - raft::resources const& res, - const index_params& params, - const cuvs::neighbors::cagra::index& cagra_index, - std::optional> dataset) +template +std::enable_if_t, + std::unique_ptr>> +from_cagra(raft::resources const& res, + const index_params& params, + CagraIndexT const& cagra_index, + std::optional> dataset) { common::nvtx::range fun_scope("hnsw::from_cagra"); auto host_dataset = raft::make_host_matrix(0, 0); @@ -256,21 +275,22 @@ std::enable_if_t>> fro host_dataset.data_handle(), host_dataset.extent(0), host_dataset.extent(1)); if (dataset.has_value()) { host_dataset_view = dataset.value(); + } else if constexpr (is_host_cagra_hnsw_export_index_v) { + RAFT_FAIL("hnsw::from_cagra requires dataset for host CAGRA index"); } else { // move dataset to host, remove padding - auto cagra_dataset = cagra_index.dataset(); - RAFT_EXPECTS(cagra_dataset.size() > 0, + auto dataset_view = cagra_index.dataset(); + RAFT_EXPECTS(dataset_view.n_rows() > 0, "Invalid CAGRA dataset of size 0, shape %zux%zu", - static_cast(cagra_dataset.extent(0)), - static_cast(cagra_dataset.extent(1))); - host_dataset = - raft::make_host_matrix(cagra_dataset.extent(0), cagra_dataset.extent(1)); + static_cast(dataset_view.n_rows()), + static_cast(dataset_view.dim())); + host_dataset = raft::make_host_matrix(dataset_view.n_rows(), dataset_view.dim()); raft::copy_matrix(host_dataset.data_handle(), host_dataset.extent(1), - cagra_dataset.data_handle(), - cagra_dataset.stride(0), + dataset_view.view().data_handle(), + dataset_view.stride(), host_dataset.extent(1), - cagra_dataset.extent(0), + dataset_view.n_rows(), raft::resource::get_cuda_stream(res)); raft::resource::sync_stream(res); host_dataset_view = host_dataset.view(); @@ -688,12 +708,15 @@ void serialize_to_hnswlib_batched(raft::resources const& res, // Serialize a disk-backed CAGRA index into hnswlib format by reading graph/dataset/label // rows directly from the backing files via pread. -template +template + requires is_cagra_hnsw_export_index_v void serialize_to_hnswlib_from_disk(raft::resources const& res, std::ostream& os_raw, const cuvs::neighbors::hnsw::index_params& params, - const cuvs::neighbors::cagra::index& index_) + CagraIndexT const& index_) { + using T = typename CagraIndexT::value_type; + using IdxT = typename CagraIndexT::index_type; RAFT_EXPECTS(index_.dataset_fd().has_value() && index_.graph_fd().has_value(), "Function only implements serialization from disk."); @@ -860,14 +883,19 @@ void serialize_to_hnswlib_from_disk(raft::resources const& res, // Serialize an in-memory CAGRA index into hnswlib format on disk, copying graph/dataset // rows from the in-memory (device or host) structures batch by batch. This avoids // materializing the full HNSW index in host memory. -template +template + requires is_cagra_hnsw_export_index_v void serialize_to_hnswlib_from_inmem( raft::resources const& res, std::ostream& os_raw, const cuvs::neighbors::hnsw::index_params& params, - const cuvs::neighbors::cagra::index& index_, - std::optional> dataset) + CagraIndexT const& index_, + std::optional< + raft::host_matrix_view> + dataset) { + using T = typename CagraIndexT::value_type; + using IdxT = typename CagraIndexT::index_type; auto stream = raft::resource::get_cuda_stream(res); [[maybe_unused]] auto num_threads = params.num_threads == 0 ? cuvs::core::omp::get_max_threads() : params.num_threads; @@ -882,12 +910,14 @@ void serialize_to_hnswlib_from_inmem( device_dataset = false; source_dataset = dataset->data_handle(); source_stride = dim; - } else if (auto cagra_dataset = index_.dataset(); cagra_dataset.data_handle() != nullptr) { - n_rows = cagra_dataset.extent(0); - dim = cagra_dataset.extent(1); + } else if constexpr (is_host_cagra_hnsw_export_index_v) { + RAFT_FAIL("serialize_to_hnswlib_from_inmem requires dataset for host CAGRA index"); + } else if (auto dataset_view = index_.dataset(); dataset_view.view().data_handle() != nullptr) { + n_rows = dataset_view.n_rows(); + dim = dataset_view.dim(); device_dataset = true; - source_dataset = cagra_dataset.data_handle(); - source_stride = cagra_dataset.stride(0); + source_dataset = dataset_view.view().data_handle(); + source_stride = dataset_view.stride(); } else { RAFT_FAIL("serialize_to_hnswlib_from_inmem: No dataset provided"); } @@ -967,12 +997,13 @@ void serialize_to_hnswlib_from_inmem( res, os_raw, params, n_rows, dim, graph_degree_int, index_.metric(), read_batch); } -template -std::enable_if_t>> from_cagra( - raft::resources const& res, - const index_params& params, - const cuvs::neighbors::cagra::index& cagra_index, - std::optional> dataset) +template +std::enable_if_t, + std::unique_ptr>> +from_cagra(raft::resources const& res, + const index_params& params, + CagraIndexT const& cagra_index, + std::optional> dataset) { common::nvtx::range fun_scope("hnsw::from_cagra"); auto stream = raft::resource::get_cuda_stream(res); @@ -1004,12 +1035,15 @@ std::enable_if_t>> fro device_copy = false; source_dataset = dataset->data_handle(); source_stride = dim; - } else if (auto cagra_dataset = cagra_index.dataset(); cagra_dataset.data_handle() != nullptr) { - n_rows = cagra_dataset.extent(0); - dim = cagra_dataset.extent(1); + } else if constexpr (is_host_cagra_hnsw_export_index_v) { + RAFT_FAIL("hnsw::from_cagra requires dataset for host CAGRA index"); + } else if (auto dataset_view = cagra_index.dataset(); + dataset_view.view().data_handle() != nullptr) { + n_rows = dataset_view.n_rows(); + dim = dataset_view.dim(); device_copy = true; - source_dataset = cagra_dataset.data_handle(); - source_stride = cagra_dataset.stride(0); + source_dataset = dataset_view.view().data_handle(); + source_stride = dataset_view.stride(); } else { RAFT_FAIL("hnsw::from_cagra: No dataset provided"); } @@ -1233,13 +1267,20 @@ inline std::pair get_available_memory( return std::make_pair(available_host_memory, available_device_memory); } -template +template + requires is_cagra_hnsw_export_index_v std::unique_ptr> from_cagra( raft::resources const& res, const index_params& params, - const cuvs::neighbors::cagra::index& cagra_index, + CagraIndexT const& cagra_index, std::optional> dataset) { + if constexpr (is_host_cagra_hnsw_export_index_v) { + if (!cagra_index.dataset_fd().has_value() && !dataset.has_value()) { + RAFT_FAIL("hnsw::from_cagra requires dataset for host CAGRA index"); + } + } + // special treatment for index on disk if (cagra_index.dataset_fd().has_value() && cagra_index.graph_fd().has_value()) { // Get directory from graph file descriptor @@ -1594,13 +1635,20 @@ std::unique_ptr> build(raft::resources const& res, cagra_ace_params.max_gpu_memory_gb = ace_params.max_gpu_memory_gb; cagra_params.graph_build_params = cagra_ace_params; } - // Build CAGRA index optionally using ACE - auto cagra_index = cuvs::neighbors::cagra::build(res, cagra_params, dataset); + + // Public HNSW API uses host_matrix_view; CAGRA build expects a padded dataset view. + // Host build stores only the graph; vectors are passed separately to from_cagra below. + cuvs::neighbors::host_padded_dataset_view host_padded_view( + dataset, static_cast(dataset.extent(1))); + auto ace_host_index = cuvs::neighbors::cagra::build(res, cagra_params, host_padded_view); RAFT_LOG_INFO("hnsw::build - Converting CAGRA index to HNSW format"); - // Convert CAGRA index to HNSW index - return from_cagra(res, params, cagra_index, dataset); + return from_cagra( + res, + params, + ace_host_index, + ace_host_index.dataset_fd().has_value() ? std::nullopt : std::make_optional(dataset)); } } // namespace cuvs::neighbors::hnsw::detail diff --git a/cpp/src/neighbors/detail/tiered_index.cuh b/cpp/src/neighbors/detail/tiered_index.cuh index 9cad64549b..6abd57507f 100644 --- a/cpp/src/neighbors/detail/tiered_index.cuh +++ b/cpp/src/neighbors/detail/tiered_index.cuh @@ -1,11 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include +#include #include #include @@ -22,7 +23,10 @@ #include #include +#include + namespace cuvs::neighbors::tiered_index::detail { + /** Storage for brute force based incremental indices @@ -108,6 +112,24 @@ using upstream_search_function_type = template struct index_state { using value_type = typename UpstreamT::value_type; + index_state() = default; + + /** + * Build upstream ANN, preserving row stride for standard CAGRA when needed. + */ + template + [[nodiscard]] static auto build_upstream_ann( + raft::resources const& res, + index_params const& tiered_params, + BuildFn&& build_fn, + DatasetView dataset) -> std::shared_ptr + { + auto index = std::forward(build_fn)(res, tiered_params, dataset); + if constexpr (std::is_same_v>) { + index.update_dataset(res, cuvs::neighbors::make_device_standard_dataset_view(dataset)); + } + return std::make_shared(std::move(index)); + } index_state(const index_state& other) : storage(other.storage), @@ -129,7 +151,7 @@ struct index_state { // Create an ANN index if we have sufficient rows in initial dataset if (dataset.extent(0) > index_params.min_ann_rows) { - ann_index = std::make_shared(std::move(build_fn(res, index_params, dataset))); + ann_index = build_upstream_ann(res, index_params, build_fn, dataset); } // allocate bfknn storage for growing the index incrementally @@ -268,6 +290,17 @@ struct index_state { std::function> build_fn; }; +/** + * After BF storage grows, repoint CAGRA at the first \p ann_rows rows. + */ +inline void update_cagra_ann_dataset_for_stride( + raft::resources const& res, + cuvs::neighbors::cagra::device_standard_index& ann_index, + raft::device_matrix_view dataset) +{ + ann_index.update_dataset(res, cuvs::neighbors::make_device_standard_dataset_view(dataset)); +} + /** * @brief Build the tiered index from the dataset for efficient search. * @@ -435,8 +468,8 @@ auto compact(raft::resources const& res, const index_state& current) auto dataset = raft::make_device_matrix_view( storage->dataset.data(), storage->num_rows_used, storage->dim); - next_state->ann_index = std::make_shared( - std::move(next_state->build_fn(res, next_state->build_params, dataset))); + next_state->ann_index = index_state::build_upstream_ann( + res, next_state->build_params, next_state->build_fn, dataset); return next_state; } } // namespace cuvs::neighbors::tiered_index::detail diff --git a/cpp/src/neighbors/detail/vamana/vamana_build.cuh b/cpp/src/neighbors/detail/vamana/vamana_build.cuh index 336d81215b..13307e1fd4 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_build.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_build.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -646,7 +646,7 @@ index build( auto quantizer = cuvs::preprocessing::quantize::pq::quantizer( pq_params, - cuvs::neighbors::vpq_dataset{ + cuvs::neighbors::device_vpq_dataset{ raft::make_device_matrix(res, 0, 0), std::move(pq_codebook), raft::make_device_matrix(res, 0, 0)}); diff --git a/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh b/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh index 4bf32e8a64..a81b2252b5 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -24,7 +24,6 @@ #include #include #include - namespace cuvs::neighbors::vamana::detail { // write matrix containing dataset to file @@ -58,31 +57,26 @@ void to_file(const std::string& dataset_base_file, raft::host_matrix */ template void serialize_dataset(raft::resources const& res, - const cuvs::neighbors::dataset* dataset, + const cuvs::neighbors::device_padded_dataset_view* dataset, const std::string& dataset_base_file) { + if (dataset == nullptr) { return; } // try allocating a buffer for the dataset on host try { - const auto* strided_dataset = - dynamic_cast*>(dataset); - if (strided_dataset) { - auto nrows = strided_dataset->n_rows(); - auto dim = strided_dataset->dim(); - auto stride = strided_dataset->stride(); - auto d_data = strided_dataset->view(); - auto h_dataset = raft::make_host_matrix(nrows, dim); - raft::copy_matrix(h_dataset.data_handle(), - dim, - d_data.data_handle(), - stride, - dim, - nrows, - raft::resource::get_cuda_stream(res)); - raft::resource::sync_stream(res); - to_file(dataset_base_file, h_dataset); - } else { - RAFT_LOG_DEBUG("dynamic_cast to strided_dataset failed"); - } + auto nrows = dataset->n_rows(); + auto dim = dataset->dim(); + auto stride = dataset->stride(); + auto d_data = dataset->view(); + auto h_dataset = raft::make_host_matrix(nrows, dim); + raft::copy_matrix(h_dataset.data_handle(), + dim, + d_data.data_handle(), + stride, + dim, + nrows, + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + to_file(dataset_base_file, h_dataset); } catch (std::bad_alloc& e) { RAFT_LOG_INFO("Failed to serialize dataset"); } catch (raft::logic_error& e) { @@ -120,11 +114,12 @@ void serialize_dataset(raft::resources const& res, * */ template -void serialize_sector_aligned(raft::resources const& res, - const HostMatT& h_graph, - const cuvs::neighbors::dataset& dataset, - const uint64_t medoid, - std::ofstream& output_writer) +void serialize_sector_aligned( + raft::resources const& res, + const HostMatT& h_graph, + const cuvs::neighbors::device_padded_dataset_view& dataset, + const uint64_t medoid, + std::ofstream& output_writer) { if constexpr (!std::is_same_v) { RAFT_FAIL("serialization is only implemented for uint32_t graph"); @@ -159,15 +154,11 @@ void serialize_sector_aligned(raft::resources const& res, const uint64_t nnodes_per_sector = sector_len / max_node_len; // 0 if max_node_len > sector_len // copy dataset to host - auto dataset_strided = - dynamic_cast*>(&dataset); - if (!dataset_strided) { RAFT_FAIL("Invalid dataset"); } - auto d_data = dataset_strided->view(); auto h_data = raft::make_host_matrix(npts, ndims); raft::copy_matrix(h_data.data_handle(), ndims, - d_data.data_handle(), - dataset_strided->stride(), + dataset.view().data_handle(), + dataset.stride(), ndims, npts, raft::resource::get_cuda_stream(res)); diff --git a/cpp/src/neighbors/detail/vpq_dataset.cuh b/cpp/src/neighbors/detail/vpq_dataset.cuh index 7d56711aff..8f2a3140b5 100644 --- a/cpp/src/neighbors/detail/vpq_dataset.cuh +++ b/cpp/src/neighbors/detail/vpq_dataset.cuh @@ -422,7 +422,7 @@ void process_and_fill_codes( bool inline_vq_labels = false) { using data_t = typename DatasetT::value_type; - using cdataset_t = vpq_dataset; + using cdataset_t = device_vpq_dataset; using label_t = uint32_t; const ix_t n_rows = dataset.extent(0); @@ -814,7 +814,7 @@ void process_and_fill_codes_subspaces( raft::device_matrix_view codes) { using data_t = typename DatasetT::value_type; - using cdataset_t = vpq_dataset; + using cdataset_t = device_vpq_dataset; using label_t = uint32_t; const ix_t n_rows = dataset.extent(0); diff --git a/cpp/src/neighbors/dynamic_batching.cu b/cpp/src/neighbors/dynamic_batching.cu index cfbef44409..c13a337c03 100644 --- a/cpp/src/neighbors/dynamic_batching.cu +++ b/cpp/src/neighbors/dynamic_batching.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -13,9 +13,19 @@ #include #include +namespace cuvs::neighbors::cagra { + +// Single-token names for CUVS_INST_DYNAMIC_BATCHING_INDEX (macro expands Namespace ::__VA_ARGS__). +using cagra_f32_u32_index = device_padded_index; +using cagra_f16_u32_index = device_padded_index; +using cagra_i8_u32_index = device_padded_index; +using cagra_u8_u32_index = device_padded_index; + +} // namespace cuvs::neighbors::cagra + namespace cuvs::neighbors::dynamic_batching { -// NB: the (template) index parameter should be the last; it may contain the spaces and so split +// NB: the (template) index parameter should be the last; it must be a single preprocessor token // into multiple preprocessor token. Then it is consumed as __VA_ARGS__ // #define CUVS_INST_DYNAMIC_BATCHING_INDEX(T, IdxT, Namespace, ...) \ @@ -47,22 +57,16 @@ namespace cuvs::neighbors::dynamic_batching { CUVS_INST_DYNAMIC_BATCHING_INDEX(float, int64_t, cuvs::neighbors::brute_force, index); // CAGRA build and search with 32-bit indices -CUVS_INST_DYNAMIC_BATCHING_INDEX(float, uint32_t, cuvs::neighbors::cagra, index); -CUVS_INST_DYNAMIC_BATCHING_INDEX(half, uint32_t, cuvs::neighbors::cagra, index); -CUVS_INST_DYNAMIC_BATCHING_INDEX(int8_t, uint32_t, cuvs::neighbors::cagra, index); -CUVS_INST_DYNAMIC_BATCHING_INDEX(uint8_t, - uint32_t, - cuvs::neighbors::cagra, - index); +CUVS_INST_DYNAMIC_BATCHING_INDEX(float, uint32_t, cuvs::neighbors::cagra, cagra_f32_u32_index); +CUVS_INST_DYNAMIC_BATCHING_INDEX(half, uint32_t, cuvs::neighbors::cagra, cagra_f16_u32_index); +CUVS_INST_DYNAMIC_BATCHING_INDEX(int8_t, uint32_t, cuvs::neighbors::cagra, cagra_i8_u32_index); +CUVS_INST_DYNAMIC_BATCHING_INDEX(uint8_t, uint32_t, cuvs::neighbors::cagra, cagra_u8_u32_index); // CAGRA build with 32-bit indices, search with 64-bit indices -CUVS_INST_DYNAMIC_BATCHING_INDEX(float, int64_t, cuvs::neighbors::cagra, index); -CUVS_INST_DYNAMIC_BATCHING_INDEX(half, int64_t, cuvs::neighbors::cagra, index); -CUVS_INST_DYNAMIC_BATCHING_INDEX(int8_t, int64_t, cuvs::neighbors::cagra, index); -CUVS_INST_DYNAMIC_BATCHING_INDEX(uint8_t, - int64_t, - cuvs::neighbors::cagra, - index); +CUVS_INST_DYNAMIC_BATCHING_INDEX(float, int64_t, cuvs::neighbors::cagra, cagra_f32_u32_index); +CUVS_INST_DYNAMIC_BATCHING_INDEX(half, int64_t, cuvs::neighbors::cagra, cagra_f16_u32_index); +CUVS_INST_DYNAMIC_BATCHING_INDEX(int8_t, int64_t, cuvs::neighbors::cagra, cagra_i8_u32_index); +CUVS_INST_DYNAMIC_BATCHING_INDEX(uint8_t, int64_t, cuvs::neighbors::cagra, cagra_u8_u32_index); // IVF-PQ with 64-bit indices CUVS_INST_DYNAMIC_BATCHING_INDEX(float, int64_t, cuvs::neighbors::ivf_pq, index); diff --git a/cpp/src/neighbors/hnsw.cpp b/cpp/src/neighbors/hnsw.cpp index 54e9dcf12a..3958366311 100644 --- a/cpp/src/neighbors/hnsw.cpp +++ b/cpp/src/neighbors/hnsw.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -46,7 +46,7 @@ CUVS_INST_HNSW_BUILD(int8_t); std::unique_ptr> from_cagra( \ raft::resources const& res, \ const index_params& params, \ - const cuvs::neighbors::cagra::index& cagra_index, \ + const cuvs::neighbors::cagra::device_padded_index& cagra_index, \ std::optional> dataset) \ { \ return detail::from_cagra(res, params, cagra_index, dataset); \ @@ -59,6 +59,57 @@ CUVS_INST_HNSW_FROM_CAGRA(int8_t); #undef CUVS_INST_HNSW_FROM_CAGRA +#define CUVS_INST_HNSW_FROM_CAGRA_STANDARD(T) \ + std::unique_ptr> from_cagra( \ + raft::resources const& res, \ + const index_params& params, \ + const cuvs::neighbors::cagra::device_standard_index& cagra_index, \ + std::optional> dataset) \ + { \ + return detail::from_cagra(res, params, cagra_index, dataset); \ + } + +CUVS_INST_HNSW_FROM_CAGRA_STANDARD(float); +CUVS_INST_HNSW_FROM_CAGRA_STANDARD(half); +CUVS_INST_HNSW_FROM_CAGRA_STANDARD(uint8_t); +CUVS_INST_HNSW_FROM_CAGRA_STANDARD(int8_t); + +#undef CUVS_INST_HNSW_FROM_CAGRA_STANDARD + +#define CUVS_INST_HNSW_FROM_CAGRA_HOST(T) \ + std::unique_ptr> from_cagra( \ + raft::resources const& res, \ + const index_params& params, \ + const cuvs::neighbors::cagra::host_padded_index& cagra_index, \ + std::optional> dataset) \ + { \ + return detail::from_cagra(res, params, cagra_index, dataset); \ + } + +CUVS_INST_HNSW_FROM_CAGRA_HOST(float); +CUVS_INST_HNSW_FROM_CAGRA_HOST(half); +CUVS_INST_HNSW_FROM_CAGRA_HOST(uint8_t); +CUVS_INST_HNSW_FROM_CAGRA_HOST(int8_t); + +#undef CUVS_INST_HNSW_FROM_CAGRA_HOST + +#define CUVS_INST_HNSW_FROM_CAGRA_HOST_STANDARD(T) \ + std::unique_ptr> from_cagra( \ + raft::resources const& res, \ + const index_params& params, \ + const cuvs::neighbors::cagra::host_standard_index& cagra_index, \ + std::optional> dataset) \ + { \ + return detail::from_cagra(res, params, cagra_index, dataset); \ + } + +CUVS_INST_HNSW_FROM_CAGRA_HOST_STANDARD(float); +CUVS_INST_HNSW_FROM_CAGRA_HOST_STANDARD(half); +CUVS_INST_HNSW_FROM_CAGRA_HOST_STANDARD(uint8_t); +CUVS_INST_HNSW_FROM_CAGRA_HOST_STANDARD(int8_t); + +#undef CUVS_INST_HNSW_FROM_CAGRA_HOST_STANDARD + #define CUVS_INST_HNSW_EXTEND(T) \ void extend(raft::resources const& res, \ const extend_params& params, \ diff --git a/cpp/src/neighbors/iface/iface.hpp b/cpp/src/neighbors/iface/iface.hpp index e76a3673af..20642a2f00 100644 --- a/cpp/src/neighbors/iface/iface.hpp +++ b/cpp/src/neighbors/iface/iface.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -12,14 +12,67 @@ #include #include #include +#include #include #include +#include + namespace cuvs::neighbors { using namespace raft; +namespace iface_detail { +/** + * @brief True when \p mds is already CAGRA row-padded on device (device or managed memory). + */ +template +bool dataset_mdspan_uses_padded_device_view( + raft::mdspan, row_major, Accessor> mds) +{ + using value_type = T; + uint32_t const required_stride = + cagra_required_row_width(static_cast(mds.extent(1))); + uint32_t const src_stride = + mds.stride(0) > 0 ? static_cast(mds.stride(0)) : static_cast(mds.extent(1)); + cudaPointerAttributes a{}; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&a, mds.data_handle())); + bool const device_src = (a.type == cudaMemoryTypeDevice) || (a.type == cudaMemoryTypeManaged); + return device_src && (src_stride == required_stride); +} + +/** Graph build via padded device view, not mdspan host build. */ +template +void cagra_build_from_device_dataset( + raft::resources const& h, + cagra::index_params const& cagra_params, + raft::mdspan, row_major, Accessor> m, + cuvs::neighbors::iface& interface) +{ + uint32_t const stride = + m.stride(0) > 0 ? static_cast(m.stride(0)) : static_cast(m.extent(1)); + auto dview = raft::make_device_strided_matrix_view( + m.data_handle(), m.extent(0), m.extent(1), stride); + if constexpr (std::is_same_v>) { + auto padded = cuvs::neighbors::make_device_padded_dataset_view(h, dview); + auto index = cuvs::neighbors::cagra::build(h, cagra_params, padded); + index.update_dataset(h, padded); + interface.index_.emplace(std::move(index)); + } else { + auto standard = cuvs::neighbors::make_device_standard_dataset_view(dview); + auto index = cuvs::neighbors::cagra::build(h, cagra_params, standard); + index.update_dataset(h, standard); + interface.index_.emplace(std::move(index)); + } + interface.cagra_owned_standard_dataset_.reset(); + interface.cagra_owned_padded_dataset_.reset(); +} +} // namespace iface_detail + +// TODO: Refactor this function signature to use the Dataset API instead of raft::mdspan. +// Currently takes a raw mdspan; should accept a dataset_view<...> so callers pass typed +// views. template void build(const raft::resources& handle, cuvs::neighbors::iface& interface, @@ -36,10 +89,45 @@ void build(const raft::resources& handle, auto idx = cuvs::neighbors::ivf_pq::build( handle, *static_cast(index_params), index_dataset); interface.index_.emplace(std::move(idx)); - } else if constexpr (std::is_same>::value) { - auto idx = cuvs::neighbors::cagra::build( - handle, *static_cast(index_params), index_dataset); - interface.index_.emplace(std::move(idx)); + } else if constexpr (std::is_same>::value || + std::is_same>::value) { + const auto& cagra_params = *static_cast(index_params); + if (raft::get_device_for_address(index_dataset.data_handle()) != -1) { + iface_detail::cagra_build_from_device_dataset( + handle, cagra_params, index_dataset, interface); + } else { + // Explicitly form a host_matrix_view so the call always resolves to the host build + // regardless of the mdspan Accessor type (both branches compile for all Accessors; + // at runtime this else branch is only reached when data_handle() is host memory). + // Wrap in host_padded_dataset_view directly (graph build is CPU-side; CUDA alignment + // is not required here — the device dataset is padded separately below). + auto host_view = raft::make_host_matrix_view( + index_dataset.data_handle(), index_dataset.extent(0), index_dataset.extent(1)); + if constexpr (std::is_same>::value) { + cuvs::neighbors::host_padded_dataset_view host_padded( + host_view, static_cast(host_view.extent(1))); + auto host_idx = cuvs::neighbors::cagra::build(handle, cagra_params, host_padded); + auto padded_r = cuvs::neighbors::make_device_padded_dataset(handle, index_dataset); + auto device_idx = cuvs::neighbors::cagra::attach_device_dataset_on_host_index( + handle, host_idx, padded_r->as_dataset_view()); + interface.cagra_owned_padded_dataset_ = std::move(padded_r); + interface.cagra_owned_standard_dataset_.reset(); + interface.index_.emplace(std::move(device_idx)); + } else { + auto host_standard = cuvs::neighbors::make_host_standard_dataset_view(host_view); + auto host_idx = cuvs::neighbors::cagra::build(handle, cagra_params, host_standard); + auto standard_r = cuvs::neighbors::make_device_standard_dataset( + handle, + index_dataset, + static_cast(index_dataset.extent(1)), + static_cast(index_dataset.stride(0))); + auto device_idx = cuvs::neighbors::cagra::attach_device_dataset_on_host_index( + handle, host_idx, standard_r->as_dataset_view()); + interface.cagra_owned_standard_dataset_ = std::move(standard_r); + interface.cagra_owned_padded_dataset_.reset(); + interface.index_.emplace(std::move(device_idx)); + } + } } resource::sync_stream(handle); } @@ -62,7 +150,7 @@ void extend( auto idx = cuvs::neighbors::ivf_pq::extend(handle, new_vectors, new_indices, interface.index_.value()); interface.index_.emplace(std::move(idx)); - } else if constexpr (std::is_same>::value) { + } else if constexpr (std::is_same>::value) { RAFT_FAIL("CAGRA does not implement the extend method"); } resource::sync_stream(handle); @@ -92,7 +180,9 @@ void search(const raft::resources& handle, queries, neighbors, distances); - } else if constexpr (std::is_same>::value) { + } else if constexpr (std::is_same>::value || + std::is_same>::value) { cuvs::neighbors::cagra::search(handle, *reinterpret_cast(search_params), interface.index_.value(), @@ -134,7 +224,8 @@ void serialize(const raft::resources& handle, ivf_flat::serialize(handle, os, interface.index_.value()); } else if constexpr (std::is_same>::value) { ivf_pq::serialize(handle, os, interface.index_.value()); - } else if constexpr (std::is_same>::value) { + } else if constexpr (std::is_same>::value || + std::is_same>::value) { cagra::serialize(handle, os, interface.index_.value(), true); } @@ -158,9 +249,22 @@ void deserialize(const raft::resources& handle, ivf_pq::deserialize(handle, is, &idx); resource::sync_stream(handle); interface.index_.emplace(std::move(idx)); - } else if constexpr (std::is_same>::value) { - cagra::index idx(handle); - cagra::deserialize(handle, is, &idx); + } else if constexpr (std::is_same>::value) { + cagra::device_padded_index idx(handle); + std::unique_ptr> out_dataset; + cagra::deserialize(handle, is, &idx, &out_dataset); + interface.cagra_owned_padded_dataset_.reset(); + interface.cagra_owned_standard_dataset_.reset(); + if (out_dataset) { interface.cagra_owned_padded_dataset_ = std::move(out_dataset); } + resource::sync_stream(handle); + interface.index_.emplace(std::move(idx)); + } else if constexpr (std::is_same>::value) { + cagra::device_standard_index idx(handle); + std::unique_ptr> out_dataset; + cagra::deserialize(handle, is, &idx, &out_dataset); + interface.cagra_owned_padded_dataset_.reset(); + interface.cagra_owned_standard_dataset_.reset(); + if (out_dataset) { interface.cagra_owned_standard_dataset_ = std::move(out_dataset); } resource::sync_stream(handle); interface.index_.emplace(std::move(idx)); } @@ -186,9 +290,22 @@ void deserialize(const raft::resources& handle, ivf_pq::deserialize(handle, is, &idx); resource::sync_stream(handle); interface.index_.emplace(std::move(idx)); - } else if constexpr (std::is_same>::value) { - cagra::index idx(handle); - cagra::deserialize(handle, is, &idx); + } else if constexpr (std::is_same>::value) { + cagra::device_padded_index idx(handle); + std::unique_ptr> out_dataset; + cagra::deserialize(handle, is, &idx, &out_dataset); + interface.cagra_owned_padded_dataset_.reset(); + interface.cagra_owned_standard_dataset_.reset(); + if (out_dataset) { interface.cagra_owned_padded_dataset_ = std::move(out_dataset); } + resource::sync_stream(handle); + interface.index_.emplace(std::move(idx)); + } else if constexpr (std::is_same>::value) { + cagra::device_standard_index idx(handle); + std::unique_ptr> out_dataset; + cagra::deserialize(handle, is, &idx, &out_dataset); + interface.cagra_owned_padded_dataset_.reset(); + interface.cagra_owned_standard_dataset_.reset(); + if (out_dataset) { interface.cagra_owned_standard_dataset_ = std::move(out_dataset); } resource::sync_stream(handle); interface.index_.emplace(std::move(idx)); } diff --git a/cpp/src/neighbors/iface/iface_cagra_inst.cu.in b/cpp/src/neighbors/iface/iface_cagra_inst.cu.in index c2456390d5..50dbe9a60d 100644 --- a/cpp/src/neighbors/iface/iface_cagra_inst.cu.in +++ b/cpp/src/neighbors/iface/iface_cagra_inst.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -25,33 +25,96 @@ using IdxT_da = template void build( const raft::resources& handle, - cuvs::neighbors::iface, data_t, index_t>& interface, + cuvs::neighbors::iface, data_t, index_t>& interface, const cuvs::neighbors::index_params* index_params, raft::mdspan, row_major, T_ha> index_dataset); template void build( const raft::resources& handle, - cuvs::neighbors::iface, data_t, index_t>& interface, + cuvs::neighbors::iface, data_t, index_t>& interface, + const cuvs::neighbors::index_params* index_params, + raft::mdspan, row_major, T_da> index_dataset); + +template void build( + const raft::resources& handle, + cuvs::neighbors::iface, data_t, index_t>& interface, + const cuvs::neighbors::index_params* index_params, + raft::mdspan, row_major, T_ha> index_dataset); + +template void build( + const raft::resources& handle, + cuvs::neighbors::iface, data_t, index_t>& interface, const cuvs::neighbors::index_params* index_params, raft::mdspan, row_major, T_da> index_dataset); template void extend( const raft::resources& handle, - cuvs::neighbors::iface, data_t, index_t>& interface, + cuvs::neighbors::iface, data_t, index_t>& interface, raft::mdspan, row_major, T_ha> new_vectors, std::optional, layout_c_contiguous, IdxT_ha>> new_indices); template void extend( const raft::resources& handle, - cuvs::neighbors::iface, data_t, index_t>& interface, + cuvs::neighbors::iface, data_t, index_t>& interface, raft::mdspan, row_major, T_da> new_vectors, std::optional, layout_c_contiguous, IdxT_da>> new_indices); +template void extend( + const raft::resources& handle, + cuvs::neighbors::iface, data_t, index_t>& interface, + raft::mdspan, row_major, T_ha> new_vectors, + std::optional, layout_c_contiguous, IdxT_ha>> + new_indices); + +template void extend( + const raft::resources& handle, + cuvs::neighbors::iface, data_t, index_t>& interface, + raft::mdspan, row_major, T_da> new_vectors, + std::optional, layout_c_contiguous, IdxT_da>> + new_indices); + +template void search( + const raft::resources& handle, + const cuvs::neighbors::iface, data_t, index_t>& + interface, + const cuvs::neighbors::search_params* search_params, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances); + +template void search( + const raft::resources& handle, + const cuvs::neighbors::iface, data_t, index_t>& + interface, + const cuvs::neighbors::search_params* search_params, + raft::host_matrix_view h_queries, + raft::device_matrix_view d_neighbors, + raft::device_matrix_view d_distances); + template void search( const raft::resources& handle, - const cuvs::neighbors::iface, data_t, index_t>& interface, + const cuvs::neighbors::iface, data_t, index_t>& + interface, + const cuvs::neighbors::search_params* search_params, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances); + +template void search( + const raft::resources& handle, + const cuvs::neighbors::iface, data_t, index_t>& + interface, + const cuvs::neighbors::search_params* search_params, + raft::host_matrix_view h_queries, + raft::device_matrix_view d_neighbors, + raft::device_matrix_view d_distances); + +template void search( + const raft::resources& handle, + const cuvs::neighbors::iface, data_t, index_t>& + interface, const cuvs::neighbors::search_params* search_params, raft::device_matrix_view queries, raft::device_matrix_view neighbors, @@ -59,7 +122,8 @@ template void search( template void search( const raft::resources& handle, - const cuvs::neighbors::iface, data_t, index_t>& interface, + const cuvs::neighbors::iface, data_t, index_t>& + interface, const cuvs::neighbors::search_params* search_params, raft::host_matrix_view h_queries, raft::device_matrix_view d_neighbors, @@ -67,7 +131,8 @@ template void search( template void search( const raft::resources& handle, - const cuvs::neighbors::iface, data_t, index_t>& interface, + const cuvs::neighbors::iface, data_t, index_t>& + interface, const cuvs::neighbors::search_params* search_params, raft::device_matrix_view queries, raft::device_matrix_view neighbors, @@ -75,7 +140,8 @@ template void search( template void search( const raft::resources& handle, - const cuvs::neighbors::iface, data_t, index_t>& interface, + const cuvs::neighbors::iface, data_t, index_t>& + interface, const cuvs::neighbors::search_params* search_params, raft::host_matrix_view h_queries, raft::device_matrix_view d_neighbors, @@ -83,17 +149,34 @@ template void search( template void serialize( const raft::resources& handle, - const cuvs::neighbors::iface, data_t, index_t>& interface, + const cuvs::neighbors::iface, data_t, index_t>& + interface, + std::ostream& os); + +template void serialize( + const raft::resources& handle, + const cuvs::neighbors::iface, data_t, index_t>& + interface, std::ostream& os); template void deserialize( const raft::resources& handle, - cuvs::neighbors::iface, data_t, index_t>& interface, + cuvs::neighbors::iface, data_t, index_t>& interface, std::istream& is); template void deserialize( const raft::resources& handle, - cuvs::neighbors::iface, data_t, index_t>& interface, + cuvs::neighbors::iface, data_t, index_t>& interface, + std::istream& is); + +template void deserialize( + const raft::resources& handle, + cuvs::neighbors::iface, data_t, index_t>& interface, + const std::string& filename); + +template void deserialize( + const raft::resources& handle, + cuvs::neighbors::iface, data_t, index_t>& interface, const std::string& filename); } // namespace cuvs::neighbors diff --git a/cpp/src/neighbors/mg/mg_cagra_inst.cu.in b/cpp/src/neighbors/mg/mg_cagra_inst.cu.in index 6e57c3f598..541f4500d2 100644 --- a/cpp/src/neighbors/mg/mg_cagra_inst.cu.in +++ b/cpp/src/neighbors/mg/mg_cagra_inst.cu.in @@ -1,92 +1,252 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include -#define CUVS_INST_MG_CAGRA(T, IdxT) \ - namespace cuvs::neighbors::cagra { \ - using namespace cuvs::neighbors; \ - \ - cuvs::neighbors::mg_index, T, IdxT> build( \ - const raft::resources& res, \ - const mg_index_params& index_params, \ - raft::host_matrix_view index_dataset) \ - { \ - cuvs::neighbors::mg_index, T, IdxT> index(res, index_params.mode); \ - cuvs::neighbors::snmg::detail::build( \ - res, \ - index, \ - static_cast(&index_params), \ - index_dataset); \ - return index; \ - } \ - \ - void extend(const raft::resources& res, \ - cuvs::neighbors::mg_index, T, IdxT>& index, \ - raft::host_matrix_view new_vectors, \ - std::optional> new_indices) \ - { \ - cuvs::neighbors::snmg::detail::extend(res, index, new_vectors, new_indices); \ - } \ - \ - void search(const raft::resources& res, \ - const cuvs::neighbors::mg_index, T, IdxT>& index, \ - const mg_search_params& search_params, \ - raft::host_matrix_view queries, \ - raft::host_matrix_view neighbors, \ - raft::host_matrix_view distances) \ - { \ - cuvs::neighbors::snmg::detail::search( \ - res, \ - index, \ - static_cast(&search_params), \ - queries, \ - neighbors, \ - distances); \ - } \ - \ - void search(const raft::resources& res, \ - const cuvs::neighbors::mg_index, T, IdxT>& index, \ - const mg_search_params& search_params, \ - raft::host_matrix_view queries, \ - raft::host_matrix_view neighbors, \ - raft::host_matrix_view distances) \ - { \ - cuvs::neighbors::snmg::detail::search( \ - res, \ - index, \ - static_cast(&search_params), \ - queries, \ - neighbors, \ - distances); \ - } \ - \ - void serialize(const raft::resources& res, \ - const cuvs::neighbors::mg_index, T, IdxT>& index, \ - const std::string& filename) \ - { \ - cuvs::neighbors::snmg::detail::serialize(res, index, filename); \ - } \ - \ - template <> \ - CUVS_EXPORT cuvs::neighbors::mg_index, T, IdxT> deserialize( \ - const raft::resources& res, const std::string& filename) \ - { \ - auto idx = cuvs::neighbors::mg_index, T, IdxT>(res, filename); \ - return idx; \ - } \ - \ - template <> \ - CUVS_EXPORT cuvs::neighbors::mg_index, T, IdxT> distribute( \ - const raft::resources& res, const std::string& filename) \ - { \ - auto idx = cuvs::neighbors::mg_index, T, IdxT>(res, REPLICATED); \ - cuvs::neighbors::snmg::detail::deserialize_and_distribute(res, idx, filename); \ - return idx; \ - } \ - } // namespace cuvs::neighbors::cagra +#define CUVS_INST_MG_CAGRA(T, IdxT) \ + namespace cuvs::neighbors::cagra { \ + using namespace cuvs::neighbors; \ + cuvs::neighbors::mg_index, T, IdxT> build( \ + const raft::resources& res, \ + const mg_index_params& index_params, \ + cuvs::neighbors::host_standard_dataset_view const& index_dataset) \ + { \ + cuvs::neighbors::mg_index, T, IdxT> index( \ + res, index_params.mode); \ + cuvs::neighbors::snmg::detail::build( \ + res, \ + index, \ + static_cast(&index_params), \ + index_dataset.view()); \ + return index; \ + } \ + \ + cuvs::neighbors::mg_index, T, IdxT> build( \ + const raft::resources& res, \ + const mg_index_params& index_params, \ + cuvs::neighbors::host_padded_dataset_view const& index_dataset) \ + { \ + cuvs::neighbors::mg_index, T, IdxT> index( \ + res, index_params.mode); \ + cuvs::neighbors::snmg::detail::build( \ + res, \ + index, \ + static_cast(&index_params), \ + index_dataset.view()); \ + return index; \ + } \ + \ + auto attach_padded_dataset_for_search( \ + const raft::resources& res, \ + const cuvs::neighbors::mg_index, T, IdxT>& idx, \ + cuvs::neighbors::device_padded_dataset_view const& padded_dataset) \ + -> cuvs::neighbors::mg_index, T, IdxT> \ + { \ + cuvs::neighbors::mg_index, T, IdxT> out(res, idx.mode_); \ + out.ann_interfaces_.resize(idx.num_ranks_); \ + auto padded_mds = padded_dataset.view(); \ + auto stride = padded_mds.extent(1); \ + const raft::resources& root_res = raft::resource::set_current_device_to_root_rank(res); \ + auto padded_host = \ + raft::make_host_matrix(padded_mds.extent(0), padded_mds.extent(1)); \ + raft::copy(padded_host.data_handle(), \ + padded_mds.data_handle(), \ + padded_mds.size(), \ + raft::resource::get_cuda_stream(root_res)); \ + raft::resource::sync_stream(root_res); \ + int64_t offset = 0; \ + for (int rank = 0; rank < idx.num_ranks_; rank++) { \ + const raft::resources& dev_res = raft::resource::set_current_device_to_rank(res, rank); \ + const auto& in_if = idx.ann_interfaces_[rank]; \ + auto& out_if = out.ann_interfaces_[rank]; \ + RAFT_EXPECTS(in_if.index_.has_value(), \ + "attach_padded_dataset_for_search: missing rank index"); \ + auto const& standard_idx = in_if.index_.value(); \ + int64_t rank_rows = static_cast(standard_idx.size()); \ + int64_t rank_offset = (idx.mode_ == SHARDED) ? offset : 0; \ + const T* rank_ptr = padded_host.data_handle() + rank_offset * stride; \ + auto rank_host_mds = \ + raft::make_host_matrix_view(rank_ptr, rank_rows, stride); \ + auto rank_device_matrix = \ + raft::make_device_matrix(dev_res, rank_rows, stride); \ + raft::copy(dev_res, rank_device_matrix.view(), rank_host_mds); \ + auto rank_device_padded = \ + std::make_unique>( \ + std::move(rank_device_matrix), padded_dataset.dim()); \ + auto padded_idx = cuvs::neighbors::cagra::attach_padded_dataset_for_search( \ + dev_res, standard_idx, rank_device_padded->as_dataset_view()); \ + out_if.cagra_owned_padded_dataset_ = std::move(rank_device_padded); \ + out_if.cagra_owned_standard_dataset_.reset(); \ + out_if.index_.emplace(std::move(padded_idx)); \ + if (idx.mode_ == SHARDED) { offset += rank_rows; } \ + } \ + return out; \ + } \ + \ + void extend(const raft::resources& res, \ + cuvs::neighbors::mg_index, T, IdxT>& index, \ + cuvs::neighbors::host_padded_dataset_view new_vectors, \ + std::optional> new_indices) \ + { \ + cuvs::neighbors::snmg::detail::extend(res, index, new_vectors.view(), new_indices); \ + } \ + \ + void extend(const raft::resources& res, \ + cuvs::neighbors::mg_index, T, IdxT>& index, \ + cuvs::neighbors::host_standard_dataset_view new_vectors, \ + std::optional> new_indices) \ + { \ + cuvs::neighbors::snmg::detail::extend(res, index, new_vectors.view(), new_indices); \ + } \ + \ + void search( \ + const raft::resources& res, \ + const cuvs::neighbors::mg_index, T, IdxT>& index, \ + const mg_search_params& search_params, \ + raft::host_matrix_view queries, \ + raft::host_matrix_view neighbors, \ + raft::host_matrix_view distances) \ + { \ + cuvs::neighbors::snmg::detail::search( \ + res, \ + index, \ + static_cast(&search_params), \ + queries, \ + neighbors, \ + distances); \ + } \ + \ + void search( \ + const raft::resources& res, \ + const cuvs::neighbors::mg_index, T, IdxT>& index, \ + const mg_search_params& search_params, \ + raft::host_matrix_view queries, \ + raft::host_matrix_view neighbors, \ + raft::host_matrix_view distances) \ + { \ + cuvs::neighbors::snmg::detail::search( \ + res, \ + index, \ + static_cast(&search_params), \ + queries, \ + neighbors, \ + distances); \ + } \ + \ + void search( \ + const raft::resources& res, \ + const cuvs::neighbors::mg_index, T, IdxT>& index, \ + const mg_search_params& search_params, \ + raft::host_matrix_view queries, \ + raft::host_matrix_view neighbors, \ + raft::host_matrix_view distances) \ + { \ + cuvs::neighbors::snmg::detail::search( \ + res, \ + index, \ + static_cast(&search_params), \ + queries, \ + neighbors, \ + distances); \ + } \ + \ + void search( \ + const raft::resources& res, \ + const cuvs::neighbors::mg_index, T, IdxT>& index, \ + const mg_search_params& search_params, \ + raft::host_matrix_view queries, \ + raft::host_matrix_view neighbors, \ + raft::host_matrix_view distances) \ + { \ + cuvs::neighbors::snmg::detail::search( \ + res, \ + index, \ + static_cast(&search_params), \ + queries, \ + neighbors, \ + distances); \ + } \ + \ + void serialize( \ + const raft::resources& res, \ + const cuvs::neighbors::mg_index, T, IdxT>& index, \ + const std::string& filename) \ + { \ + cuvs::neighbors::snmg::detail::serialize(res, index, filename); \ + } \ + \ + void serialize( \ + const raft::resources& res, \ + const cuvs::neighbors::mg_index, T, IdxT>& index, \ + const std::string& filename) \ + { \ + cuvs::neighbors::snmg::detail::serialize(res, index, filename); \ + } \ + \ + template <> \ + CUVS_EXPORT void deserialize( \ + const raft::resources& res, \ + const std::string& filename, \ + cuvs::neighbors::mg_index, T, IdxT>* index) \ + { \ + *index = \ + cuvs::neighbors::mg_index, T, IdxT>(res, filename); \ + } \ + \ + template <> \ + CUVS_EXPORT void deserialize( \ + const raft::resources& res, \ + const std::string& filename, \ + cuvs::neighbors::mg_index, T, IdxT>* index) \ + { \ + *index = \ + cuvs::neighbors::mg_index, T, IdxT>(res, filename); \ + } \ + \ + template <> \ + CUVS_EXPORT void distribute( \ + const raft::resources& res, \ + const std::string& filename, \ + cuvs::neighbors::mg_index, T, IdxT>* index) \ + { \ + auto idx = \ + cuvs::neighbors::mg_index, T, IdxT>(res, REPLICATED); \ + cuvs::neighbors::snmg::detail::deserialize_and_distribute(res, idx, filename); \ + *index = std::move(idx); \ + } \ + \ + template <> \ + CUVS_EXPORT void distribute( \ + const raft::resources& res, \ + const std::string& filename, \ + cuvs::neighbors::mg_index, T, IdxT>* index) \ + { \ + auto idx = \ + cuvs::neighbors::mg_index, T, IdxT>(res, REPLICATED); \ + cuvs::neighbors::snmg::detail::deserialize_and_distribute(res, idx, filename); \ + *index = std::move(idx); \ + } \ + } \ + template CUVS_EXPORT cuvs::neighbors:: \ + mg_index, T, IdxT>::mg_index( \ + const raft::resources&); \ + template CUVS_EXPORT cuvs::neighbors:: \ + mg_index, T, IdxT>::mg_index( \ + const raft::resources&, cuvs::neighbors::distribution_mode); \ + template CUVS_EXPORT cuvs::neighbors:: \ + mg_index, T, IdxT>::mg_index( \ + const raft::resources&, const std::string&); \ + template CUVS_EXPORT cuvs::neighbors:: \ + mg_index, T, IdxT>::mg_index( \ + const raft::resources&); \ + template CUVS_EXPORT cuvs::neighbors:: \ + mg_index, T, IdxT>::mg_index( \ + const raft::resources&, cuvs::neighbors::distribution_mode); \ + template CUVS_EXPORT cuvs::neighbors:: \ + mg_index, T, IdxT>::mg_index( \ + const raft::resources&, const std::string&); CUVS_INST_MG_CAGRA(@data_type@, uint32_t); #undef CUVS_INST_MG_CAGRA diff --git a/cpp/src/neighbors/mg/snmg.cuh b/cpp/src/neighbors/mg/snmg.cuh index 43e4aa4471..057a3e8272 100644 --- a/cpp/src/neighbors/mg/snmg.cuh +++ b/cpp/src/neighbors/mg/snmg.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -587,7 +587,8 @@ void search(const raft::resources& clique, static_cast*>(search_params); search_mode = mg_search_params->search_mode; n_rows_per_batch = mg_search_params->n_rows_per_batch; - } else if constexpr (std::is_same>::value) { + } else if constexpr (std::is_same>::value || + std::is_same>::value) { const cuvs::neighbors::mg_search_params* mg_search_params = static_cast*>(search_params); search_mode = mg_search_params->search_mode; @@ -666,7 +667,8 @@ void search(const raft::resources& clique, static_cast*>(search_params); merge_mode = mg_search_params->merge_mode; n_rows_per_batch = mg_search_params->n_rows_per_batch; - } else if constexpr (std::is_same>::value) { + } else if constexpr (std::is_same>::value || + std::is_same>::value) { const cuvs::neighbors::mg_search_params* mg_search_params = static_cast*>(search_params); merge_mode = mg_search_params->merge_mode; @@ -764,6 +766,12 @@ namespace cuvs::neighbors { using namespace cuvs::neighbors; using namespace raft; +template +mg_index::mg_index(const raft::resources& clique) + : mg_index(clique, cuvs::neighbors::REPLICATED) +{ +} + template mg_index::mg_index(const raft::resources& clique, distribution_mode mode) : mode_(mode), round_robin_counter_(std::make_shared>(0)) diff --git a/cpp/src/neighbors/tiered_index.cu b/cpp/src/neighbors/tiered_index.cu index 076c0c4a7c..c414c05d8a 100644 --- a/cpp/src/neighbors/tiered_index.cu +++ b/cpp/src/neighbors/tiered_index.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -28,14 +28,85 @@ void typed_search(raft::resources const& res, } } // namespace cuvs::neighbors::ivf_pq +namespace { +// Wrapper with the exact signature expected by upstream_build_function_type. +// cagra::build is now a template (no concrete device_matrix_view overload), so it cannot be +// passed as a plain function pointer; this wrapper bridges the gap. +cuvs::neighbors::cagra::device_standard_index cagra_build_for_tiered_standard( + raft::resources const& res, + cuvs::neighbors::cagra::index_params const& params, + raft::device_matrix_view dataset) +{ + auto view = cuvs::neighbors::make_device_standard_dataset_view(dataset); + return cuvs::neighbors::cagra::build(res, params, view); +} + +cuvs::neighbors::cagra::device_padded_index cagra_build_for_tiered_padded( + raft::resources const& res, + cuvs::neighbors::cagra::index_params const& params, + raft::device_matrix_view dataset) +{ + auto view = cuvs::neighbors::make_device_padded_dataset_view(res, dataset); + return cuvs::neighbors::cagra::build(res, params, view); +} + +} // namespace + namespace cuvs::neighbors::tiered_index { +auto build(raft::resources const& res, + const index_params& params, + cuvs::neighbors::device_padded_dataset_view dataset) + -> tiered_index::index> +{ + auto state = detail::build>( + res, params, cagra_build_for_tiered_padded, dataset.view()); + return cuvs::neighbors::tiered_index::index>(state); +} + auto build(raft::resources const& res, const index_params& params, raft::device_matrix_view dataset) - -> tiered_index::index> + -> tiered_index::index> { - auto state = detail::build>(res, params, cagra::build, dataset); - return cuvs::neighbors::tiered_index::index>(state); + auto state = detail::build>( + res, params, cagra_build_for_tiered_standard, dataset); + return cuvs::neighbors::tiered_index::index>(state); +} + +auto attach_padded_dataset_for_search( + raft::resources const& res, + const tiered_index::index>& idx, + cuvs::neighbors::device_padded_dataset_view padded_dataset) + -> tiered_index::index> +{ + RAFT_EXPECTS(padded_dataset.n_rows() == idx.size(), + "padded_dataset row count must match tiered index size"); + RAFT_EXPECTS(padded_dataset.dim() == idx.dim(), + "padded_dataset dimension must match tiered index dimension"); + + auto next_state = + std::make_shared>>(); + next_state->storage = idx.state->storage; + next_state->build_params = idx.state->build_params; + next_state->build_fn = cagra_build_for_tiered_padded; + next_state->ann_index.reset(); + + if (idx.state->ann_index) { + auto padded_mds = padded_dataset.view(); + auto ann_rows = static_cast(idx.state->ann_rows()); + auto ann_mds = raft::make_device_matrix_view( + padded_mds.data_handle(), ann_rows, static_cast(padded_mds.extent(1))); + auto ann_padded_view = + cuvs::neighbors::device_padded_dataset_view(ann_mds, padded_dataset.dim()); + auto ann_padded_idx = cuvs::neighbors::cagra::attach_padded_dataset_for_search( + res, *idx.state->ann_index, ann_padded_view); + next_state->ann_index = + std::make_shared>( + std::move(ann_padded_idx)); + } + + return cuvs::neighbors::tiered_index::index>( + next_state); } auto build(raft::resources const& res, @@ -60,7 +131,16 @@ auto build(raft::resources const& res, void extend(raft::resources const& res, raft::device_matrix_view new_vectors, - tiered_index::index>* idx) + tiered_index::index>* idx) +{ + std::scoped_lock lock(idx->write_mutex); + auto next_state = detail::extend(res, *idx->state, new_vectors); + idx->state = next_state; +} + +void extend(raft::resources const& res, + raft::device_matrix_view new_vectors, + tiered_index::index>* idx) { std::scoped_lock lock(idx->write_mutex); auto next_state = detail::extend(res, *idx->state, new_vectors); @@ -78,7 +158,7 @@ void extend(raft::resources const& res, // Block 'search' calls during the update_dataset call to ensure that this // doesn't cause issues in a multithreaded environment std::unique_lock lock(idx->ann_mutex); - next_state->ann_index->update_dataset(res, dataset); + detail::update_cagra_ann_dataset_for_stride(res, *next_state->ann_index, dataset); } } @@ -103,7 +183,16 @@ void extend(raft::resources const& res, idx->state = next_state; } -void compact(raft::resources const& res, tiered_index::index>* idx) +void compact(raft::resources const& res, + tiered_index::index>* idx) +{ + std::scoped_lock lock(idx->write_mutex); + auto next_state = detail::compact(res, *idx->state); + idx->state = next_state; +} + +void compact(raft::resources const& res, + tiered_index::index>* idx) { std::scoped_lock lock(idx->write_mutex); auto next_state = detail::compact(res, *idx->state); @@ -127,19 +216,31 @@ void compact(raft::resources const& res, void search(raft::resources const& res, const cagra::search_params& search_params, - const tiered_index::index>& index, + const tiered_index::index>& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, raft::device_matrix_view distances, const cuvs::neighbors::filtering::base_filter& sample_filter) { - // use a read-write lock to handle calls to update_dataset for cagra - // This allows multiple readers concurrently, but only one writer std::shared_lock lock(index.ann_mutex); index.state->search( res, search_params, cagra::search, queries, neighbors, distances, sample_filter); } +void search(raft::resources const& res, + const cagra::search_params& search_params, + const tiered_index::index>& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& sample_filter) +{ + RAFT_FAIL( + "tiered_index::search(standard CAGRA) requires explicit attach first. " + "Call tiered_index::attach_padded_dataset_for_search(...) and then search the returned padded " + "tiered index."); +} + void search(raft::resources const& res, const ivf_flat::search_params& search_params, const tiered_index::index>& index, @@ -164,13 +265,24 @@ void search(raft::resources const& res, res, search_params, ivf_pq::typed_search, queries, neighbors, distances, sample_filter); } -auto merge(raft::resources const& res, - const index_params& index_params, - const std::vector>*>& indices) - -> tiered_index::index> +auto merge( + raft::resources const& res, + const index_params& index_params, + const std::vector>*>& indices) + -> tiered_index::index> +{ + auto state = detail::merge(res, index_params, indices); + return cuvs::neighbors::tiered_index::index>(state); +} + +auto merge( + raft::resources const& res, + const index_params& index_params, + const std::vector>*>& indices) + -> tiered_index::index> { auto state = detail::merge(res, index_params, indices); - return cuvs::neighbors::tiered_index::index>(state); + return cuvs::neighbors::tiered_index::index>(state); } auto merge(raft::resources const& res, @@ -203,7 +315,15 @@ int64_t index::dim() const noexcept return state->dim(); } -template struct index>; +template CUVS_EXPORT int64_t +index>::size() const noexcept; +template CUVS_EXPORT int64_t +index>::dim() const noexcept; +template CUVS_EXPORT int64_t +index>::size() const noexcept; +template CUVS_EXPORT int64_t +index>::dim() const noexcept; + template struct index>; template struct index>; diff --git a/cpp/src/preprocessing/quantize/detail/pq.cuh b/cpp/src/preprocessing/quantize/detail/pq.cuh index 5d77e2dd44..7fea89461a 100644 --- a/cpp/src/preprocessing/quantize/detail/pq.cuh +++ b/cpp/src/preprocessing/quantize/detail/pq.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -193,7 +193,7 @@ quantizer build( res, filled_params, dataset, raft::make_const_mdspan(vq_code_book.view())); } return {filled_params, - cuvs::neighbors::vpq_dataset{ + cuvs::neighbors::device_vpq_dataset{ std::move(vq_code_book), std::move(pq_code_book), std::move(empty_codes)}}; } @@ -369,8 +369,8 @@ void inverse_transform( template void vpq_convert_math_type(const raft::resources& res, - const cuvs::neighbors::vpq_dataset& src, - cuvs::neighbors::vpq_dataset& dst) + const cuvs::neighbors::device_vpq_dataset& src, + cuvs::neighbors::device_vpq_dataset& dst) { raft::linalg::map(res, dst.vq_code_book.view(), @@ -409,7 +409,7 @@ inline auto make_pq_params_from_vpq(const cuvs::neighbors::vpq_params& in_params template auto vpq_build(const raft::resources& res, const cuvs::neighbors::vpq_params& params, - const DatasetT& dataset) -> cuvs::neighbors::vpq_dataset + const DatasetT& dataset) -> cuvs::neighbors::device_vpq_dataset { using label_t = uint32_t; // Use a heuristic to impute missing parameters. @@ -437,17 +437,17 @@ auto vpq_build(const raft::resources& res, codes.view(), true); - return cuvs::neighbors::vpq_dataset{ + return cuvs::neighbors::device_vpq_dataset{ std::move(vq_code_book), std::move(pq_code_book), std::move(codes)}; } template auto vpq_build_half(const raft::resources& res, const cuvs::neighbors::vpq_params& params, - const DatasetT& dataset) -> cuvs::neighbors::vpq_dataset + const DatasetT& dataset) -> cuvs::neighbors::device_vpq_dataset { auto old_type = vpq_build(res, params, dataset); - auto new_type = cuvs::neighbors::vpq_dataset{ + auto new_type = cuvs::neighbors::device_vpq_dataset{ raft::make_device_mdarray(res, old_type.vq_code_book.extents()), raft::make_device_mdarray(res, old_type.pq_code_book.extents()), std::move(old_type.data)}; diff --git a/cpp/src/preprocessing/quantize/pq.cu b/cpp/src/preprocessing/quantize/pq.cu index 761474bdf8..68068aa35f 100644 --- a/cpp/src/preprocessing/quantize/pq.cu +++ b/cpp/src/preprocessing/quantize/pq.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -7,6 +7,9 @@ #include +#include +#include + namespace cuvs::preprocessing::quantize::pq { #define CUVS_INST_QUANTIZATION(T, QuantI) \ @@ -73,4 +76,58 @@ CUVS_INST_VPQ_BUILD(uint8_t); #undef CUVS_INST_VPQ_BUILD +namespace detail { + +template +auto vpq_train_from_device_rows(raft::resources const& res, + cuvs::neighbors::vpq_params const& params, + T const* src_ptr, + int64_t n_rows, + int64_t dim, + int64_t stride) + -> cuvs::neighbors::device_vpq_dataset +{ + auto stream = raft::resource::get_cuda_stream(res); + if (stride != dim) { + auto dense = raft::make_device_matrix(res, n_rows, dim); + raft::copy_matrix(dense.data_handle(), dim, src_ptr, stride, dim, n_rows, stream); + auto dense_view = + raft::make_device_matrix_view(dense.data_handle(), n_rows, dim); + return detail::vpq_build_half(res, params, dense_view); + } + auto row_view = raft::make_device_matrix_view(src_ptr, n_rows, dim); + return detail::vpq_build_half(res, params, row_view); +} + +} // namespace detail + +template cuvs::neighbors::device_vpq_dataset +detail::vpq_train_from_device_rows(raft::resources const&, + cuvs::neighbors::vpq_params const&, + float const*, + int64_t, + int64_t, + int64_t); +template cuvs::neighbors::device_vpq_dataset +detail::vpq_train_from_device_rows(raft::resources const&, + cuvs::neighbors::vpq_params const&, + half const*, + int64_t, + int64_t, + int64_t); +template cuvs::neighbors::device_vpq_dataset +detail::vpq_train_from_device_rows(raft::resources const&, + cuvs::neighbors::vpq_params const&, + int8_t const*, + int64_t, + int64_t, + int64_t); +template cuvs::neighbors::device_vpq_dataset +detail::vpq_train_from_device_rows(raft::resources const&, + cuvs::neighbors::vpq_params const&, + uint8_t const*, + int64_t, + int64_t, + int64_t); + } // namespace cuvs::preprocessing::quantize::pq diff --git a/cpp/src/preprocessing/quantize/vpq_build-ext.cuh b/cpp/src/preprocessing/quantize/vpq_build-ext.cuh deleted file mode 100644 index 1745e53a33..0000000000 --- a/cpp/src/preprocessing/quantize/vpq_build-ext.cuh +++ /dev/null @@ -1,28 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ -#pragma once - -#include -#include - -namespace cuvs::preprocessing::quantize::pq { - -#define CUVS_INST_VPQ_BUILD(T) \ - cuvs::neighbors::vpq_dataset vpq_build( \ - const raft::resources& res, \ - const cuvs::neighbors::vpq_params& params, \ - const raft::host_matrix_view& dataset); \ - cuvs::neighbors::vpq_dataset vpq_build( \ - const raft::resources& res, \ - const cuvs::neighbors::vpq_params& params, \ - const raft::device_matrix_view& dataset); - -CUVS_INST_VPQ_BUILD(float); -CUVS_INST_VPQ_BUILD(half); -CUVS_INST_VPQ_BUILD(int8_t); -CUVS_INST_VPQ_BUILD(uint8_t); - -#undef CUVS_INST_VPQ_BUILD -} // namespace cuvs::preprocessing::quantize::pq diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 79bf339827..a31ab8df3b 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -9,11 +9,13 @@ #include "vpq_utils.cuh" #include +#include "cagra_padded_build_helpers.cuh" #include "naive_knn.cuh" #include #include #include +#include #include #include #include @@ -37,13 +39,40 @@ #include #include +#include #include #include +#include #include namespace cuvs::neighbors::cagra { namespace { +/** + * If \p ace_host_dataset is set, builds from that host mdspan via `cagra::build` (ACE is selected + * by `graph_build_params`). Otherwise builds from \p padded via `cagra::build`. When \p + * ACE is selected by `graph_build_params`. + */ +template +void cagra_build_into_index( + raft::resources const& res, + cagra::index_params const& params, + std::optional> ace_host_dataset, + cuvs::neighbors::device_padded_dataset_view const& padded, + cagra::device_padded_index& index) +{ + if (ace_host_dataset.has_value()) { + cuvs::neighbors::host_padded_dataset_view host_view( + *ace_host_dataset, static_cast(ace_host_dataset->extent(1))); + auto host_idx = cagra::build(res, params, host_view); + // In-memory ACE returns graph-only; attach device padded storage for search. + index = cagra::attach_device_dataset_on_host_index(res, host_idx, padded); + return; + } + index = cagra::build(res, params, padded); + index.update_dataset(res, padded); +} + struct test_cagra_sample_filter { static constexpr unsigned offset = 300; inline _RAFT_HOST_DEVICE auto operator()( @@ -271,7 +300,6 @@ struct AnnCagraInputs { // std::optional double min_recall; // = std::nullopt; std::optional ivf_pq_search_refine_ratio = std::nullopt; - std::optional compression = std::nullopt; std::optional non_owning_memory_buffer_flag = std::nullopt; cuvs::neighbors::MergeStrategy merge_strategy = @@ -300,13 +328,6 @@ inline ::std::ostream& operator<<(::std::ostream& os, const AnnCagraInputs& p) {search_algo::AUTO, "auto"} // }; std::vector build_algo = {"IVF_PQ", "NN_DESCENT", "ITERATIVE_CAGRA_SEARCH", "AUTO"}; - const auto smem_dtype_str = [](cuvs::neighbors::cagra::internal_dtype dtype) { - switch (dtype) { - case cuvs::neighbors::cagra::internal_dtype::F16: return "F16"; - case cuvs::neighbors::cagra::internal_dtype::E5M2: return "E5M2"; - } - return "Unknown"; - }; std::vector merge_strategy = {"PHYSICAL", "LOGICAL"}; os << "{n_queries=" << p.n_queries << ", dataset shape=" << p.n_rows << "x" << p.dim << ", k=" << p.k << ", " << algo_name[p.algo] << ", max_queries=" << p.max_queries @@ -317,12 +338,7 @@ inline ::std::ostream& operator<<(::std::ostream& os, const AnnCagraInputs& p) if ((int)p.build_algo == 0 && p.ivf_pq_search_refine_ratio) { os << "(refine_rate=" << *p.ivf_pq_search_refine_ratio << ')'; } - if (p.compression.has_value()) { - auto vpq = p.compression.value(); - os << ", pq_bits=" << vpq.pq_bits << ", pq_dim=" << vpq.pq_dim - << ", vq_n_centers=" << vpq.vq_n_centers << ", smem_dtype=" << smem_dtype_str(p.smem_dtype); - } - os << '}'; + os << '}' << std::endl; return os; } @@ -355,7 +371,6 @@ class AnnCagraTest : public ::testing::TestWithParam { ps.build_algo != graph_build_algo::ITERATIVE_CAGRA_SEARCH) GTEST_SKIP(); if (ps.metric == cuvs::distance::DistanceType::CosineExpanded) { - if (ps.compression.has_value()) { GTEST_SKIP(); } if (ps.build_algo == graph_build_algo::ITERATIVE_CAGRA_SEARCH || ps.dim == 1) { GTEST_SKIP(); } @@ -418,7 +433,6 @@ class AnnCagraTest : public ::testing::TestWithParam { break; }; - index_params.compression = ps.compression; cagra::search_params search_params; search_params.algo = ps.algo; search_params.max_queries = ps.max_queries; @@ -427,21 +441,26 @@ class AnnCagraTest : public ::testing::TestWithParam { 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 device_padded(handle_, + database_view); tmp_index_file index_file; { std::optional> database_host{std::nullopt}; - cagra::index index(handle_, index_params.metric); + std::optional> ace_host_dataset; + cagra::device_padded_index index(handle_, index_params.metric); if (ps.host_dataset) { - database_host = raft::make_host_matrix(ps.n_rows, ps.dim); + database_host.emplace(raft::make_host_matrix(ps.n_rows, ps.dim)); raft::copy(database_host->data_handle(), database.data(), database.size(), stream_); - auto database_host_view = raft::make_host_matrix_view( - (const DataT*)database_host->data_handle(), ps.n_rows, ps.dim); - - index = cagra::build(handle_, index_params, database_host_view); - } else { - index = cagra::build(handle_, index_params, database_view); - }; + raft::resource::sync_stream(handle_); + if (std::holds_alternative( + index_params.graph_build_params)) { + ace_host_dataset.emplace(raft::make_host_matrix_view( + database_host->data_handle(), ps.n_rows, ps.dim)); + } + } + cagra_build_into_index( + handle_, index_params, ace_host_dataset, device_padded.view, index); if (ps.use_source_indices) { auto source_indices = @@ -453,10 +472,11 @@ class AnnCagraTest : public ::testing::TestWithParam { cagra::serialize(handle_, index_file.filename, index, ps.include_serialized_dataset); } - cagra::index index(handle_); - cagra::deserialize(handle_, index_file.filename, &index); + cagra::device_padded_index index(handle_); + std::unique_ptr> loaded_dataset; + cagra::deserialize(handle_, index_file.filename, &index, &loaded_dataset); - if (!ps.include_serialized_dataset) { index.update_dataset(handle_, database_view); } + if (!ps.include_serialized_dataset) { index.update_dataset(handle_, device_padded.view); } auto search_queries_view = raft::make_device_matrix_view( search_queries.data(), ps.n_queries, ps.dim); @@ -473,46 +493,6 @@ class AnnCagraTest : public ::testing::TestWithParam { raft::resource::sync_stream(handle_); reference_recall = 1; - if (ps.compression.has_value()) { - auto decoded_dataset = - raft::make_device_matrix(handle_, ps.n_rows, ps.dim); - - using VpqMathT = half; - cuvs::neighbors::decode_vpq_dataset( - decoded_dataset.view(), - dynamic_cast&>(index.data()), - raft::resource::get_cuda_stream(handle_)); - auto indices_out_view = raft::make_device_matrix_view( - indices_dev.data(), ps.n_queries, ps.k); - auto dists_out_view = raft::make_device_matrix_view( - distances_dev.data(), ps.n_queries, ps.k); - - cuvs::neighbors::naive_knn(handle_, - dists_out_view.data_handle(), - indices_out_view.data_handle(), - search_queries.data(), - decoded_dataset.data_handle(), - ps.n_queries, - ps.n_rows, - ps.dim, - ps.k, - ps.metric); - std::vector indices_vpq_dataset(queries_size); - std::vector distances_vpq_dataset(queries_size); - raft::update_host( - distances_vpq_dataset.data(), dists_out_view.data_handle(), queries_size, stream_); - raft::update_host( - indices_vpq_dataset.data(), indices_out_view.data_handle(), queries_size, stream_); - - reference_recall = std::get<1>(calc_recall(indices_naive, - indices_vpq_dataset, - distances_naive, - distances_vpq_dataset, - ps.n_queries, - ps.k, - 0)); - printf("reference_recall = %e\n", reference_recall); - } } // for (int i = 0; i < min(ps.n_queries, 10); i++) { @@ -531,20 +511,18 @@ class AnnCagraTest : public ::testing::TestWithParam { ps.k, 0.003, min_recall)); - if (!ps.compression.has_value()) { - // Don't evaluate distances for CAGRA-Q for now as the error can be somewhat large - EXPECT_TRUE(eval_distances(handle_, - database.data(), - search_queries.data(), - indices_dev.data(), - distances_dev.data(), - ps.n_rows, - ps.dim, - ps.n_queries, - ps.k, - ps.metric, - 1.0e-4)); - } + // Don't evaluate distances for CAGRA-Q for now as the error can be somewhat large + EXPECT_TRUE(eval_distances(handle_, + database.data(), + search_queries.data(), + indices_dev.data(), + distances_dev.data(), + ps.n_rows, + ps.dim, + ps.n_queries, + ps.k, + ps.metric, + 1.0e-4)); } } @@ -592,12 +570,10 @@ class AnnCagraAddNodesTest : public ::testing::TestWithParam { ps.build_algo != graph_build_algo::ITERATIVE_CAGRA_SEARCH) GTEST_SKIP(); if (ps.metric == cuvs::distance::DistanceType::CosineExpanded) { - if (ps.compression.has_value()) { GTEST_SKIP(); } if (ps.build_algo == graph_build_algo::ITERATIVE_CAGRA_SEARCH || ps.dim == 1) { GTEST_SKIP(); } } - if (ps.compression != std::nullopt) GTEST_SKIP(); // IVF_PQ graph build does not support BitwiseHamming if (ps.metric == cuvs::distance::DistanceType::BitwiseHamming && ((!std::is_same_v) || (ps.build_algo == graph_build_algo::IVF_PQ))) @@ -678,21 +654,24 @@ class AnnCagraAddNodesTest : public ::testing::TestWithParam { auto initial_database_view = raft::make_device_matrix_view( (const DataT*)database.data(), initial_database_size, ps.dim); + cuvs::neighbors::test::padded_device_matrix_for_cagra initial_padded( + handle_, initial_database_view); std::optional> database_host{std::nullopt}; - cagra::index index(handle_); + std::optional> ace_host_dataset; + cagra::device_padded_index index(handle_); if (ps.host_dataset) { - database_host = raft::make_host_matrix(ps.n_rows, ps.dim); + database_host.emplace(raft::make_host_matrix(ps.n_rows, ps.dim)); raft::copy( database_host->data_handle(), database.data(), initial_database_view.size(), stream_); - auto database_host_view = raft::make_host_matrix_view( - (const DataT*)database_host->data_handle(), initial_database_size, ps.dim); - // NB: database_host must live no less than the index, because the index _may_be_ - // non-onwning - index = cagra::build(handle_, index_params, database_host_view); - } else { - index = cagra::build(handle_, index_params, initial_database_view); - }; + raft::resource::sync_stream(handle_); + if (std::holds_alternative( + index_params.graph_build_params)) { + ace_host_dataset.emplace(raft::make_host_matrix_view( + database_host->data_handle(), initial_database_size, ps.dim)); + } + } + cagra_build_into_index(handle_, index_params, ace_host_dataset, initial_padded.view, index); auto additional_dataset = raft::make_host_matrix(ps.n_rows - initial_database_size, index.dim()); @@ -701,32 +680,29 @@ class AnnCagraAddNodesTest : public ::testing::TestWithParam { additional_dataset.size(), stream_); - auto new_dataset_buffer = raft::make_device_matrix(handle_, 0, 0); - auto new_graph_buffer = raft::make_device_matrix(handle_, 0, 0); - std::optional> - new_dataset_buffer_view = std::nullopt; - std::optional> new_graph_buffer_view = std::nullopt; - if (ps.non_owning_memory_buffer_flag.has_value() && - ps.non_owning_memory_buffer_flag.value()) { - const auto stride = - dynamic_cast*>(&index.data()) - ->stride(); - new_dataset_buffer = raft::make_device_matrix(handle_, ps.n_rows, stride); - new_graph_buffer = - raft::make_device_matrix(handle_, ps.n_rows, index.graph_degree()); - - new_dataset_buffer_view = raft::make_device_strided_matrix_view( - new_dataset_buffer.data_handle(), ps.n_rows, ps.dim, stride); - new_graph_buffer_view = new_graph_buffer.view(); - } - cagra::extend_params extend_params; - cagra::extend(handle_, - extend_params, - raft::make_const_mdspan(additional_dataset.view()), - index, - new_dataset_buffer_view, - new_graph_buffer_view); + std::optional> extend_storage{}; + std::unique_ptr> add_padded_owner{ + nullptr}; + if constexpr (cuvs::neighbors::is_padded_dataset_view_v) { + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(additional_dataset.view())) { + auto add_view = + cuvs::neighbors::make_host_padded_dataset_view(additional_dataset.view()); + extend_storage.emplace(cagra::make_extended_storage(handle_, add_view, index)); + cagra::extend(handle_, extend_params, add_view, index, *extend_storage); + } else { + add_padded_owner = + cuvs::neighbors::make_host_padded_dataset(handle_, additional_dataset.view()); + auto add_view = add_padded_owner->as_dataset_view(); + extend_storage.emplace(cagra::make_extended_storage(handle_, add_view, index)); + cagra::extend(handle_, extend_params, add_view, index, *extend_storage); + } + } else { + auto add_view = + cuvs::neighbors::make_host_standard_dataset_view(additional_dataset.view()); + extend_storage.emplace(cagra::make_extended_storage(handle_, add_view, index)); + cagra::extend(handle_, extend_params, add_view, index, *extend_storage); + } auto search_queries_view = raft::make_device_matrix_view( search_queries.data(), ps.n_queries, ps.dim); @@ -808,7 +784,6 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { ps.build_algo != graph_build_algo::ITERATIVE_CAGRA_SEARCH) GTEST_SKIP(); if (ps.metric == cuvs::distance::DistanceType::CosineExpanded) { - if (ps.compression.has_value()) { GTEST_SKIP(); } if (ps.build_algo == graph_build_algo::ITERATIVE_CAGRA_SEARCH || ps.dim == 1) { GTEST_SKIP(); } @@ -888,7 +863,6 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { break; }; - index_params.compression = ps.compression; cagra::search_params search_params; search_params.algo = ps.algo; search_params.max_queries = ps.max_queries; @@ -897,20 +871,23 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { 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 device_padded(handle_, + database_view); std::optional> database_host{std::nullopt}; - cagra::index index(handle_); + std::optional> ace_host_dataset; + cagra::device_padded_index index(handle_); if (ps.host_dataset) { - database_host = raft::make_host_matrix(ps.n_rows, ps.dim); + database_host.emplace(raft::make_host_matrix(ps.n_rows, ps.dim)); raft::copy(database_host->data_handle(), database.data(), database.size(), stream_); - auto database_host_view = raft::make_host_matrix_view( - (const DataT*)database_host->data_handle(), ps.n_rows, ps.dim); - index = cagra::build(handle_, index_params, database_host_view); - } else { - index = cagra::build(handle_, index_params, database_view); + raft::resource::sync_stream(handle_); + if (std::holds_alternative( + index_params.graph_build_params)) { + ace_host_dataset.emplace(raft::make_host_matrix_view( + database_host->data_handle(), ps.n_rows, ps.dim)); + } } - - if (!ps.include_serialized_dataset) { index.update_dataset(handle_, database_view); } + cagra_build_into_index(handle_, index_params, ace_host_dataset, device_padded.view, index); if (ps.use_source_indices) { auto source_indices = @@ -969,20 +946,18 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { 0.003, min_recall, false)); - if (!ps.compression.has_value()) { - // Don't evaluate distances for CAGRA-Q for now as the error can be somewhat large - EXPECT_TRUE(eval_distances(handle_, - database.data(), - search_queries.data(), - indices_dev.data(), - distances_dev.data(), - ps.n_rows, - ps.dim, - ps.n_queries, - ps.k, - ps.metric, - 1.0e-4)); - } + // Don't evaluate distances for CAGRA-Q for now as the error can be somewhat large + EXPECT_TRUE(eval_distances(handle_, + database.data(), + search_queries.data(), + indices_dev.data(), + distances_dev.data(), + ps.n_rows, + ps.dim, + ps.n_queries, + ps.k, + ps.metric, + 1.0e-4)); } } @@ -1034,7 +1009,6 @@ class AnnCagraIndexFilteredMergeTest : public ::testing::TestWithParam) || (ps.build_algo == graph_build_algo::IVF_PQ))) @@ -1151,35 +1125,38 @@ class AnnCagraIndexFilteredMergeTest : public ::testing::TestWithParam( (const DataT*)database.data() + database0_view.size(), database1_size, ps.dim); - cagra::index index0(handle_, index_params.metric); - cagra::index index1(handle_, index_params.metric); + cuvs::neighbors::test::padded_device_matrix_for_cagra padded0(handle_, + database0_view); + cuvs::neighbors::test::padded_device_matrix_for_cagra padded1(handle_, + database1_view); + + cagra::device_padded_index index0(handle_, index_params.metric); + cagra::device_padded_index index1(handle_, index_params.metric); std::optional> database_host{std::nullopt}; + std::optional> ace_host0, ace_host1; if (ps.host_dataset) { - database_host = raft::make_host_matrix(handle_, ps.n_rows, ps.dim); + database_host.emplace(raft::make_host_matrix(handle_, ps.n_rows, ps.dim)); raft::copy(database_host->data_handle(), database.data(), database.size(), stream_); - { - auto database_host_view = raft::make_host_matrix_view( - (const DataT*)database_host->data_handle(), database0_size, ps.dim); - index0 = cagra::build(handle_, index_params, database_host_view); - } - { - auto database_host_view = raft::make_host_matrix_view( - (const DataT*)database_host->data_handle() + database0_size * ps.dim, - database1_size, - ps.dim); - index1 = cagra::build(handle_, index_params, database_host_view); + raft::resource::sync_stream(handle_); + if (std::holds_alternative( + index_params.graph_build_params)) { + ace_host0.emplace(raft::make_host_matrix_view( + database_host->data_handle(), database0_size, ps.dim)); + ace_host1.emplace(raft::make_host_matrix_view( + database_host->data_handle() + database0_size * ps.dim, database1_size, ps.dim)); } - } else { - index0 = cagra::build(handle_, index_params, database0_view); - index1 = cagra::build(handle_, index_params, database1_view); - }; + } + cagra_build_into_index(handle_, index_params, ace_host0, padded0.view, index0); + cagra_build_into_index(handle_, index_params, ace_host1, padded1.view, index1); - std::vector*> indices; + std::vector*> indices; indices.push_back(&index0); indices.push_back(&index1); - auto index = - cuvs::neighbors::cagra::merge(handle_, index_params, indices, bitset_filter_obj); + auto merge_storage = + cuvs::neighbors::cagra::make_merged_storage(handle_, indices, bitset_filter_obj); + auto merge_idx = cuvs::neighbors::cagra::merge( + handle_, index_params, indices, merge_storage, bitset_filter_obj); auto search_queries_view = raft::make_device_matrix_view( search_queries.data(), ps.n_queries, ps.dim); @@ -1195,7 +1172,7 @@ class AnnCagraIndexFilteredMergeTest : public ::testing::TestWithParam { GTEST_SKIP(); } } - if (ps.compression != std::nullopt) GTEST_SKIP(); // IVF_PQ graph build does not support BitwiseHamming if (ps.metric == cuvs::distance::DistanceType::BitwiseHamming && ((!std::is_same_v) || (ps.build_algo == graph_build_algo::IVF_PQ))) @@ -1364,28 +1339,29 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { auto database1_view = raft::make_device_matrix_view( (const DataT*)database.data() + database0_view.size(), database1_size, ps.dim); - cagra::index index0(handle_, index_params.metric); - cagra::index index1(handle_, index_params.metric); + cuvs::neighbors::test::padded_device_matrix_for_cagra merge_padded0(handle_, + database0_view); + cuvs::neighbors::test::padded_device_matrix_for_cagra merge_padded1(handle_, + database1_view); + + cagra::device_padded_index index0(handle_, index_params.metric); + cagra::device_padded_index index1(handle_, index_params.metric); std::optional> database_host{std::nullopt}; + std::optional> ace_host0, ace_host1; if (ps.host_dataset) { - database_host = raft::make_host_matrix(handle_, ps.n_rows, ps.dim); + database_host.emplace(raft::make_host_matrix(handle_, ps.n_rows, ps.dim)); raft::copy(database_host->data_handle(), database.data(), database.size(), stream_); - { - auto database_host_view = raft::make_host_matrix_view( - (const DataT*)database_host->data_handle(), database0_size, ps.dim); - index0 = cagra::build(handle_, index_params, database_host_view); - } - { - auto database_host_view = raft::make_host_matrix_view( - (const DataT*)database_host->data_handle() + database0_size * ps.dim, - database1_size, - ps.dim); - index1 = cagra::build(handle_, index_params, database_host_view); + raft::resource::sync_stream(handle_); + if (std::holds_alternative( + index_params.graph_build_params)) { + ace_host0.emplace(raft::make_host_matrix_view( + database_host->data_handle(), database0_size, ps.dim)); + ace_host1.emplace(raft::make_host_matrix_view( + database_host->data_handle() + database0_size * ps.dim, database1_size, ps.dim)); } - } else { - index0 = cagra::build(handle_, index_params, database0_view); - index1 = cagra::build(handle_, index_params, database1_view); - }; + } + cagra_build_into_index(handle_, index_params, ace_host0, merge_padded0.view, index0); + cagra_build_into_index(handle_, index_params, ace_host1, merge_padded1.view, index1); auto search_queries_view = raft::make_device_matrix_view( search_queries.data(), ps.n_queries, ps.dim); @@ -1400,12 +1376,18 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { search_params.team_size = ps.team_size; search_params.itopk_size = ps.itopk_size; - std::vector*> indices_to_merge{&index0, &index1}; + std::vector*> indices_to_merge{&index0, &index1}; if (ps.merge_strategy == cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL) { - auto merged = cagra::merge(handle_, index_params, indices_to_merge); - cagra::search( - handle_, search_params, merged, search_queries_view, indices_out_view, dists_out_view); + auto merge_storage = + cuvs::neighbors::cagra::make_merged_storage(handle_, indices_to_merge); + auto merged_idx = cagra::merge(handle_, index_params, indices_to_merge, merge_storage); + cagra::search(handle_, + search_params, + merged_idx, + search_queries_view, + indices_out_view, + dists_out_view); } else { cuvs::neighbors::composite::composite_index composite( indices_to_merge); @@ -1491,7 +1473,6 @@ inline std::vector generate_inputs() {true, false}, {0.995}, {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, {std::optional{std::nullopt}}, {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL}); @@ -1516,7 +1497,6 @@ inline std::vector generate_inputs() {false}, {0.995}, {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, {std::optional{std::nullopt}}, {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); @@ -1542,7 +1522,6 @@ inline std::vector generate_inputs() {false}, {0.995}, {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, {std::optional{std::nullopt}}, {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); @@ -1565,7 +1544,6 @@ inline std::vector generate_inputs() {true}, {0.995}, {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, {std::optional{std::nullopt}}, {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); @@ -1594,7 +1572,6 @@ inline std::vector generate_inputs() {false}, {0.995}, {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, {std::optional{std::nullopt}}, {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); @@ -1624,7 +1601,6 @@ inline std::vector generate_inputs() {false}, {0.995}, {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, {std::optional{std::nullopt}}, {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); @@ -1652,7 +1628,6 @@ inline std::vector generate_inputs() {false}, {0.995}, {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, {std::optional{std::nullopt}}, {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); @@ -1675,52 +1650,11 @@ inline std::vector generate_inputs() {true}, {0.985}, {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, {std::optional{std::nullopt}}, {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); - // A few PQ configurations. - // Varying dim, vq_n_centers - inputs2 = raft::util::itertools::product( - {100}, - {10000}, - {64, 128, 192, 256, 512, 1024}, // dim - {16}, // k - {graph_build_algo::IVF_PQ}, - {search_algo::AUTO}, - {10}, - {0}, - {64}, - {1}, - {cuvs::distance::DistanceType::L2Expanded}, - {false}, - {true}, - {false}, - {0.6}, - {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, - {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, - cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); // don't demand high recall - // without refinement - for (uint32_t pq_len : {2, 4, 8}) { - for (uint32_t vq_n_centers : {100, 1000}) { - for (auto internal_smem_dtype : {cuvs::neighbors::cagra::internal_dtype::E5M2, - cuvs::neighbors::cagra::internal_dtype::F16}) { - for (auto input : inputs2) { - vpq_params ps{}; - ps.pq_dim = input.dim / pq_len; - ps.vq_n_centers = vq_n_centers; - input.compression.emplace(ps); - input.smem_dtype = internal_smem_dtype; - inputs.push_back(input); - } - } - } - } - // Refinement options // Varying host_dataset, ivf_pq_search_refine_ratio inputs2 = raft::util::itertools::product( @@ -1740,7 +1674,6 @@ inline std::vector generate_inputs() {true}, {0.99}, {1.0f, 2.0f, 3.0f}, - {std::optional{std::nullopt}}, {std::optional{std::nullopt}}, {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); @@ -1764,7 +1697,6 @@ inline std::vector generate_inputs() {false}, {0.995}, {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, {std::optional{std::nullopt}}, {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); @@ -1819,40 +1751,10 @@ inline std::vector generate_addnode_inputs() {false}, {0.985}, {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, {std::optional{std::nullopt}}, {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); - // a few PQ configurations - inputs2 = raft::util::itertools::product( - {100}, - {10000}, - {192, 1024}, // dim - {16}, // k - {graph_build_algo::IVF_PQ}, - {search_algo::AUTO}, - {10}, - {0}, - {64}, - {1}, - {cuvs::distance::DistanceType::L2Expanded}, - {false}, - {true}, - {true}, - {0.6}); // don't demand high recall without refinement - for (uint32_t pq_len : {2}) { // for now, only pq_len = 2 is supported, more options coming soon - for (uint32_t vq_n_centers : {100}) { - for (auto input : inputs2) { - vpq_params ps{}; - ps.pq_dim = input.dim / pq_len; - ps.vq_n_centers = vq_n_centers; - input.compression.emplace(ps); - inputs.push_back(input); - } - } - } - return inputs; } @@ -1895,35 +1797,6 @@ inline std::vector generate_filtering_inputs() {0.995}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); - // a few PQ configurations - inputs2 = raft::util::itertools::product( - {100}, - {10000}, - {256}, // dim - {16}, // k - {graph_build_algo::IVF_PQ}, - {search_algo::AUTO}, - {10}, - {0}, - {64}, - {1}, - {cuvs::distance::DistanceType::L2Expanded}, - {false}, - {true}, - {true}, - {0.6}); // don't demand high recall without refinement - for (uint32_t pq_len : {2}) { // for now, only pq_len = 2 is supported, more options coming soon - for (uint32_t vq_n_centers : {100}) { - for (auto input : inputs2) { - vpq_params ps{}; - ps.pq_dim = input.dim / pq_len; - ps.vq_n_centers = vq_n_centers; - input.compression.emplace(ps); - inputs.push_back(input); - } - } - } - return inputs; } const std::vector inputs = generate_inputs(); diff --git a/cpp/tests/neighbors/ann_cagra/bug_extreme_inputs_oob.cu b/cpp/tests/neighbors/ann_cagra/bug_extreme_inputs_oob.cu index 8468a724a4..41b7e77596 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_extreme_inputs_oob.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_extreme_inputs_oob.cu @@ -1,10 +1,11 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include +#include "../cagra_padded_build_helpers.cuh" #include #include @@ -30,7 +31,9 @@ class cagra_extreme_inputs_oob_test : public ::testing::Test { ix_ps.intermediate_graph_degree = 128; try { - [[maybe_unused]] auto ix = cagra::build(res, ix_ps, raft::make_const_mdspan(dataset->view())); + cuvs::neighbors::test::padded_device_matrix_for_cagra padded( + res, raft::make_const_mdspan(dataset->view())); + [[maybe_unused]] auto ix = cagra::build(res, ix_ps, padded.view); raft::resource::sync_stream(res); } catch (const std::exception&) { SUCCEED(); diff --git a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu index adeb774a8b..c2f3a25994 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu @@ -1,160 +1,169 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include - -#include - -#include -#include -#include -#include -#include - -#include - -namespace cuvs::neighbors::cagra { - -/** - * @brief Test verifying graph.extent(0) is used for random seed selection - * - * This test ensures that CAGRA search kernels correctly use graph.extent(0) - * (graph size) rather than dataset.size for random seed node selection. - * - * The bug: random seed selection previously used dataset_desc.size, which - * could cause OOB access if the graph size differed from dataset size - * (e.g., in CAGRA-Q iterative builds with compression). - * - * The fix: kernels now receive graph.extent(0) as graph_size parameter, - * ensuring seeds are always within valid graph node range [0, graph_size). - */ -class cagra_graph_smaller_than_dataset_test : public ::testing::Test { - public: - using data_type = float; - using index_type = uint32_t; - - protected: - void run() - { - // Create a dataset with 1000 points - constexpr int64_t n_dataset = 1000; - constexpr int64_t n_dim = 128; - constexpr int64_t n_queries = 100; - constexpr int64_t k = 10; - - // Build index normally - auto dataset = raft::make_device_matrix(res, n_dataset, n_dim); - raft::random::RngState r(1234ULL); - raft::random::uniform( - res, r, dataset.data_handle(), n_dataset * n_dim, data_type(-1), data_type(1)); - - cagra::index_params index_params; - index_params.graph_degree = 32; - index_params.intermediate_graph_degree = 64; - - auto index = cagra::build(res, index_params, raft::make_const_mdspan(dataset.view())); - raft::resource::sync_stream(res); - - // Get the graph from the index - auto original_graph = index.graph(); - ASSERT_EQ(original_graph.extent(0), n_dataset); - - // Recreate the bug scenario: LARGE dataset, SMALL graph - // (like iterative_build_graph does in intermediate iterations) - constexpr int64_t n_graph = n_dataset / 2; // Only 500 nodes in graph - - // Step 1: Build index on SMALL subset (500 points) - auto small_dataset_view = raft::make_device_matrix_view( - dataset.data_handle(), n_graph, n_dim); - - cagra::index_params small_index_params; - small_index_params.graph_degree = 32; - auto small_index = cagra::build(res, small_index_params, small_dataset_view); - raft::resource::sync_stream(res); - - // Step 2: Update to FULL dataset (1000 points) but keep small graph (500 nodes) - // This creates the exact bug scenario: dataset.size=1000, graph.extent(0)=500 - small_index.update_dataset(res, raft::make_const_mdspan(dataset.view())); - - // Verify the mismatch - THIS IS THE BUG SCENARIO! - ASSERT_EQ(small_index.graph().extent(0), n_graph); // Graph has 500 nodes - ASSERT_EQ(small_index.size(), n_dataset); // Dataset has 1000 points - ASSERT_NE(small_index.graph().extent(0), small_index.size()); // Mismatch! - - // Create queries - auto queries = raft::make_device_matrix(res, n_queries, n_dim); - raft::random::uniform( - res, r, queries.data_handle(), n_queries * n_dim, data_type(-1), data_type(1)); - - // Allocate output - auto neighbors = raft::make_device_matrix(res, n_queries, k); - auto distances = raft::make_device_matrix(res, n_queries, k); - - // Setup search params - cagra::search_params search_params; - search_params.itopk_size = 64; - search_params.search_width = 1; - search_params.max_iterations = 10; - search_params.algo = cagra::search_algo::SINGLE_CTA; - - // THIS SHOULD NOT CRASH OR CAUSE OOB ACCESS - // Before fix: random seeds use dataset.size (1000) -> tries to access graph[700] -> CRASH! - // After fix: random seeds use graph.extent(0) (500) -> only accesses graph[0-499] -> SAFE! - cagra::search(res, - search_params, - small_index, - raft::make_const_mdspan(queries.view()), - neighbors.view(), - distances.view()); - - raft::resource::sync_stream(res); - - // Verify results are valid (neighbors should be < graph size) - auto neighbors_host = raft::make_host_matrix(n_queries, k); - raft::copy(neighbors_host.data_handle(), - neighbors.data_handle(), - n_queries * k, - raft::resource::get_cuda_stream(res)); - raft::resource::sync_stream(res); - - // All neighbor indices should be valid (< n_graph) - for (int64_t i = 0; i < n_queries * k; i++) { - ASSERT_LT(neighbors_host.data_handle()[i], n_graph) - << "Neighbor index " << neighbors_host.data_handle()[i] << " is >= graph size " << n_graph; - } - - // Test with MULTI_CTA algorithm as well (also had the same bug) - search_params.algo = cagra::search_algo::MULTI_CTA; - - cagra::search(res, - search_params, - small_index, - raft::make_const_mdspan(queries.view()), - neighbors.view(), - distances.view()); - - raft::resource::sync_stream(res); - - // Verify again - raft::copy(neighbors_host.data_handle(), - neighbors.data_handle(), - n_queries * k, - raft::resource::get_cuda_stream(res)); - raft::resource::sync_stream(res); - - for (int64_t i = 0; i < n_queries * k; i++) { - ASSERT_LT(neighbors_host.data_handle()[i], n_graph) - << "Neighbor index " << neighbors_host.data_handle()[i] << " is >= graph size " << n_graph - << " (MULTI_CTA)"; - } - } - - private: - raft::resources res; -}; - -TEST_F(cagra_graph_smaller_than_dataset_test, search_with_smaller_graph) { this->run(); } - -} // namespace cuvs::neighbors::cagra +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include "../cagra_padded_build_helpers.cuh" +#include + +#include +#include +#include +#include +#include + +#include + +namespace cuvs::neighbors::cagra { + +/** + * @brief Test verifying graph.extent(0) is used for random seed selection + * + * This test ensures that CAGRA search kernels correctly use graph.extent(0) + * (graph size) rather than dataset.size for random seed node selection. + * + * The bug: random seed selection previously used dataset_desc.size, which + * could cause OOB access if the graph size differed from dataset size + * (e.g., in CAGRA-Q iterative builds with compression). + * + * The fix: kernels now receive graph.extent(0) as graph_size parameter, + * ensuring seeds are always within valid graph node range [0, graph_size). + */ +class cagra_graph_smaller_than_dataset_test : public ::testing::Test { + public: + using data_type = float; + using index_type = uint32_t; + + protected: + void run() + { + // Create a dataset with 1000 points + constexpr int64_t n_dataset = 1000; + constexpr int64_t n_dim = 128; + constexpr int64_t n_queries = 100; + constexpr int64_t k = 10; + + // Build index normally + auto dataset = raft::make_device_matrix(res, n_dataset, n_dim); + raft::random::RngState r(1234ULL); + raft::random::uniform( + res, r, dataset.data_handle(), n_dataset * n_dim, data_type(-1), data_type(1)); + + cagra::index_params index_params; + index_params.graph_degree = 32; + index_params.intermediate_graph_degree = 64; + + cuvs::neighbors::test::padded_device_matrix_for_cagra padded_full( + res, raft::make_const_mdspan(dataset.view())); + auto index = cagra::build(res, index_params, padded_full.view); + raft::resource::sync_stream(res); + + // Get the graph from the index + auto original_graph = index.graph(); + ASSERT_EQ(original_graph.extent(0), n_dataset); + + // Recreate the bug scenario: LARGE dataset, SMALL graph + // (like iterative_build_graph does in intermediate iterations) + constexpr int64_t n_graph = n_dataset / 2; // Only 500 nodes in graph + + // Step 1: Build index on SMALL subset (500 points) + auto small_dataset_view = raft::make_device_matrix_view( + dataset.data_handle(), n_graph, n_dim); + + cagra::index_params small_index_params; + small_index_params.graph_degree = 32; + cuvs::neighbors::test::padded_device_matrix_for_cagra padded_small( + res, small_dataset_view); + auto small_index = cagra::build(res, small_index_params, padded_small.view); + small_index.update_dataset(res, padded_small.view); + raft::resource::sync_stream(res); + + // Step 2: Update to FULL dataset (1000 points) but keep small graph (500 nodes) + // This creates the exact bug scenario: dataset.size=1000, graph.extent(0)=500 + small_index.update_dataset(res, + cuvs::neighbors::make_device_padded_dataset_view( + res, raft::make_const_mdspan(dataset.view()))); + + // Verify the mismatch - THIS IS THE BUG SCENARIO! + ASSERT_EQ(small_index.graph().extent(0), n_graph); // Graph has 500 nodes + ASSERT_EQ(small_index.size(), n_dataset); // Dataset has 1000 points + ASSERT_NE(small_index.graph().extent(0), + small_index.size()); // Mismatch! + + // Create queries + auto queries = raft::make_device_matrix(res, n_queries, n_dim); + raft::random::uniform( + res, r, queries.data_handle(), n_queries * n_dim, data_type(-1), data_type(1)); + + // Allocate output + auto neighbors = raft::make_device_matrix(res, n_queries, k); + auto distances = raft::make_device_matrix(res, n_queries, k); + + // Setup search params + cagra::search_params search_params; + search_params.itopk_size = 64; + search_params.search_width = 1; + search_params.max_iterations = 10; + search_params.algo = cagra::search_algo::SINGLE_CTA; + + // THIS SHOULD NOT CRASH OR CAUSE OOB ACCESS + // Before fix: random seeds use dataset.size (1000) -> tries to access graph[700] -> CRASH! + // After fix: random seeds use graph.extent(0) (500) -> only accesses graph[0-499] -> SAFE! + cagra::search(res, + search_params, + small_index, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + + raft::resource::sync_stream(res); + + // Verify results are valid (neighbors should be < graph size) + auto neighbors_host = raft::make_host_matrix(n_queries, k); + raft::copy(neighbors_host.data_handle(), + neighbors.data_handle(), + n_queries * k, + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + // All neighbor indices should be valid (< n_graph) + for (int64_t i = 0; i < n_queries * k; i++) { + ASSERT_LT(neighbors_host.data_handle()[i], n_graph) + << "Neighbor index " << neighbors_host.data_handle()[i] << " is >= graph size " << n_graph; + } + + // Test with MULTI_CTA algorithm as well (also had the same bug) + search_params.algo = cagra::search_algo::MULTI_CTA; + + cagra::search(res, + search_params, + small_index, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + + raft::resource::sync_stream(res); + + // Verify again + raft::copy(neighbors_host.data_handle(), + neighbors.data_handle(), + n_queries * k, + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + for (int64_t i = 0; i < n_queries * k; i++) { + ASSERT_LT(neighbors_host.data_handle()[i], n_graph) + << "Neighbor index " << neighbors_host.data_handle()[i] << " is >= graph size " << n_graph + << " (MULTI_CTA)"; + } + } + + private: + raft::resources res; +}; + +TEST_F(cagra_graph_smaller_than_dataset_test, search_with_smaller_graph) { this->run(); } + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/neighbors/ann_cagra/bug_issue_93_reproducer.cu b/cpp/tests/neighbors/ann_cagra/bug_issue_93_reproducer.cu index 6b4b037167..320626b211 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_issue_93_reproducer.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_issue_93_reproducer.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Reproducer for https://github.com/rapidsai/cuvs-lucene/issues/93 @@ -28,6 +28,7 @@ #include +#include "../cagra_padded_build_helpers.cuh" #include #include #include @@ -40,6 +41,7 @@ #include #include #include +#include #include namespace cuvs::neighbors::cagra { @@ -56,8 +58,9 @@ TEST(Issue93Reproducer, ConcurrentSearchDifferentGraphDegrees) constexpr int dim = 64; constexpr int top_k = 10; - // Build indices on the main thread. - std::vector> indices; + // Build indices on the main thread (keep padded builders alive for view-based indexes). + std::vector> padded_builders; + std::vector> indices; for (int n_rows : dataset_sizes) { auto database = raft::make_device_matrix(handle, n_rows, dim); raft::random::uniform( @@ -70,7 +73,10 @@ TEST(Issue93Reproducer, ConcurrentSearchDifferentGraphDegrees) ip.graph_build_params = graph_build_params::nn_descent_params(ip.intermediate_graph_degree, ip.metric); - indices.push_back(cagra::build(handle, ip, raft::make_const_mdspan(database.view()))); + padded_builders.emplace_back(handle, raft::make_const_mdspan(database.view())); + auto index = cagra::build(handle, ip, padded_builders.back().view); + index.update_dataset(handle, padded_builders.back().view); + indices.push_back(std::move(index)); } raft::resource::sync_stream(handle); diff --git a/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu b/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu index aaee5a77e5..c3774f5cfd 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu @@ -1,85 +1,88 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include - -#include - -#include -#include -#include - -#include -#include - -namespace cuvs::neighbors::cagra { - -template -class CagraIterativeBuildBugTest : public ::testing::Test { - public: - using data_type = DataT; - - protected: - void run() - { - // Set up iterative CAGRA graph building - cagra::index_params index_params; - // The bug manifests when graph_degree is equal to intermediate_graph_degree - // see issue https://github.com/rapidsai/cuvs/issues/1818 - index_params.graph_degree = 16; - index_params.intermediate_graph_degree = 16; - - // Use iterative CAGRA search for graph building - index_params.graph_build_params = graph_build_params::iterative_search_params(); - - // Build the index - auto cagra_index = cagra::build(res, index_params, raft::make_const_mdspan(dataset->view())); - raft::resource::sync_stream(res); - - // Verify the index was built successfully - ASSERT_GT(cagra_index.size(), 0); - ASSERT_EQ(cagra_index.dim(), n_dim); - } - - void SetUp() override - { - dataset.emplace(raft::make_device_matrix(res, n_samples, n_dim)); - raft::random::RngState r(1234ULL); - - // Generate random data based on type - if constexpr (std::is_same_v) { - raft::random::normal( - res, r, dataset->data_handle(), n_samples * n_dim, data_type(0), data_type(1)); - } else if constexpr (std::is_same_v) { - raft::random::uniformInt( - res, r, dataset->data_handle(), n_samples * n_dim, int8_t(-128), int8_t(127)); - } else if constexpr (std::is_same_v) { - raft::random::uniformInt( - res, r, dataset->data_handle(), n_samples * n_dim, uint8_t(0), uint8_t(255)); - } - raft::resource::sync_stream(res); - } - - void TearDown() override - { - dataset.reset(); - raft::resource::sync_stream(res); - } - - private: - raft::resources res; - std::optional> dataset = std::nullopt; - - constexpr static int64_t n_samples = 10000; - constexpr static int64_t n_dim = 1024; -}; - -// Instantiate test for different data types -using TestTypes = ::testing::Types; -TYPED_TEST_SUITE(CagraIterativeBuildBugTest, TestTypes); - -TYPED_TEST(CagraIterativeBuildBugTest, IterativeBuildTest) { this->run(); } - -} // namespace cuvs::neighbors::cagra +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include "../cagra_padded_build_helpers.cuh" +#include + +#include +#include +#include + +#include +#include + +namespace cuvs::neighbors::cagra { + +template +class CagraIterativeBuildBugTest : public ::testing::Test { + public: + using data_type = DataT; + + protected: + void run() + { + // Set up iterative CAGRA graph building + cagra::index_params index_params; + // The bug manifests when graph_degree is equal to intermediate_graph_degree + // see issue https://github.com/rapidsai/cuvs/issues/1818 + index_params.graph_degree = 16; + index_params.intermediate_graph_degree = 16; + + // Use iterative CAGRA search for graph building + index_params.graph_build_params = graph_build_params::iterative_search_params(); + + cuvs::neighbors::test::padded_device_matrix_for_cagra padded( + res, raft::make_const_mdspan(dataset->view())); + auto cagra_index = cagra::build(res, index_params, padded.view); + cagra_index.update_dataset(res, padded.view); + raft::resource::sync_stream(res); + + // Verify the index was built successfully + ASSERT_GT(cagra_index.size(), 0); + ASSERT_EQ(cagra_index.dim(), n_dim); + } + + void SetUp() override + { + dataset.emplace(raft::make_device_matrix(res, n_samples, n_dim)); + raft::random::RngState r(1234ULL); + + // Generate random data based on type + if constexpr (std::is_same_v) { + raft::random::normal( + res, r, dataset->data_handle(), n_samples * n_dim, data_type(0), data_type(1)); + } else if constexpr (std::is_same_v) { + raft::random::uniformInt( + res, r, dataset->data_handle(), n_samples * n_dim, int8_t(-128), int8_t(127)); + } else if constexpr (std::is_same_v) { + raft::random::uniformInt( + res, r, dataset->data_handle(), n_samples * n_dim, uint8_t(0), uint8_t(255)); + } + raft::resource::sync_stream(res); + } + + void TearDown() override + { + dataset.reset(); + raft::resource::sync_stream(res); + } + + private: + raft::resources res; + std::optional> dataset = std::nullopt; + + constexpr static int64_t n_samples = 10000; + constexpr static int64_t n_dim = 1024; +}; + +// Instantiate test for different data types +using TestTypes = ::testing::Types; +TYPED_TEST_SUITE(CagraIterativeBuildBugTest, TestTypes); + +TYPED_TEST(CagraIterativeBuildBugTest, IterativeBuildTest) { this->run(); } + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/neighbors/ann_cagra/bug_multi_cta_crash.cu b/cpp/tests/neighbors/ann_cagra/bug_multi_cta_crash.cu index 4b418b20cd..85d8eb8315 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_multi_cta_crash.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_multi_cta_crash.cu @@ -1,11 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include #include "../ann_cagra.cuh" +#include "../cagra_padded_build_helpers.cuh" #include @@ -27,8 +28,9 @@ class AnnCagraBugMultiCTACrash : public ::testing::TestWithParamview())); + build_padded_.emplace(res, raft::make_const_mdspan(dataset->view())); + auto cagra_index = cagra::build(res, cagra_index_params, build_padded_->view); + cagra_index.update_dataset(res, build_padded_->view); raft::resource::sync_stream(res); cagra::search_params cagra_search_params; @@ -67,6 +69,7 @@ class AnnCagraBugMultiCTACrash : public ::testing::TestWithParam> build_padded_{}; std::optional> dataset = std::nullopt; std::optional> queries = std::nullopt; std::optional> neighbors = std::nullopt; diff --git a/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu b/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu index 093727d318..2cc92b2aca 100644 --- a/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu +++ b/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -123,7 +123,10 @@ class CagraUdfFilterTest : public ::testing::TestWithParam { index_params.graph_build_params = cagra::graph_build_params::nn_descent_params(index_params.intermediate_graph_degree); - index.emplace(cagra::build(res, index_params, raft::make_const_mdspan(dataset->view()))); + index.emplace(cagra::build(res, + index_params, + cuvs::neighbors::make_device_padded_dataset_view( + res, raft::make_const_mdspan(dataset->view())))); raft::resource::sync_stream(res); } @@ -166,9 +169,9 @@ class CagraUdfFilterTest : public ::testing::TestWithParam { } raft::resources res; - std::optional> dataset = std::nullopt; - std::optional> queries = std::nullopt; - std::optional> index = std::nullopt; + std::optional> dataset = std::nullopt; + std::optional> queries = std::nullopt; + std::optional> index = std::nullopt; }; class CagraUdfFilterHalfTest : public ::testing::TestWithParam { @@ -199,7 +202,10 @@ class CagraUdfFilterHalfTest : public ::testing::TestWithParamview()))); + index.emplace(cagra::build(res, + index_params, + cuvs::neighbors::make_device_padded_dataset_view( + res, raft::make_const_mdspan(dataset->view())))); raft::resource::sync_stream(res); } @@ -242,9 +248,9 @@ class CagraUdfFilterHalfTest : public ::testing::TestWithParam> dataset = std::nullopt; - std::optional> queries = std::nullopt; - std::optional> index = std::nullopt; + std::optional> dataset = std::nullopt; + std::optional> queries = std::nullopt; + std::optional> index = std::nullopt; }; TEST_P(CagraUdfFilterTest, AcceptAllMatchesNoFilter) diff --git a/cpp/tests/neighbors/ann_hnsw_ace.cuh b/cpp/tests/neighbors/ann_hnsw_ace.cuh index c75b3555f6..f71c577762 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace.cuh +++ b/cpp/tests/neighbors/ann_hnsw_ace.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -301,15 +301,17 @@ class AnnHnswAceTest : public ::testing::TestWithParam { raft::copy(database_host.data_handle(), database_dev.data(), ps.n_rows * ps.dim, stream_); raft::resource::sync_stream(handle_); - auto database_view = + auto database_dev_view = raft::make_device_matrix_view(database_dev.data(), ps.n_rows, ps.dim); + cuvs::neighbors::test::padded_device_matrix_for_cagra device_padded(handle_, + database_dev_view); // Build an in-memory CAGRA index (device graph + device dataset). cuvs::neighbors::cagra::index_params cagra_params; cagra_params.metric = ps.metric; cagra_params.graph_degree = 64; cagra_params.intermediate_graph_degree = 128; - auto cagra_index = cuvs::neighbors::cagra::build(handle_, cagra_params, database_view); + auto cagra_index = cuvs::neighbors::cagra::build(handle_, cagra_params, device_padded.view); raft::resource::sync_stream(handle_); cuvs::neighbors::hnsw::search_params search_params; @@ -318,7 +320,7 @@ class AnnHnswAceTest : public ::testing::TestWithParam { // Runs from_cagra with a tiny host-memory limit to force the disk spill, searches the // returned (disk-backed) index, checks recall, and returns the neighbor indices. - auto run_spilled = [&](const cuvs::neighbors::cagra::index& idx, + auto run_spilled = [&](const cuvs::neighbors::cagra::device_padded_index& idx, hnsw::HnswHierarchy hierarchy, bool pass_host_dataset) -> std::vector { static std::atomic counter{0}; @@ -402,8 +404,8 @@ class AnnHnswAceTest : public ::testing::TestWithParam { raft::resource::sync_stream(handle_); auto managed_graph_view = raft::make_device_matrix_view( managed_graph.data(), ps.n_rows, degree); - cuvs::neighbors::cagra::index managed_index( - handle_, ps.metric, database_view, managed_graph_view); + cuvs::neighbors::cagra::device_padded_index managed_index( + handle_, ps.metric, device_padded.view, managed_graph_view); run_spilled(managed_index, hnsw::HnswHierarchy::NONE, /*pass_host_dataset=*/false); } diff --git a/cpp/tests/neighbors/ann_scann.cuh b/cpp/tests/neighbors/ann_scann.cuh index eafddec9d2..81ef21c8e2 100644 --- a/cpp/tests/neighbors/ann_scann.cuh +++ b/cpp/tests/neighbors/ann_scann.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -186,7 +186,7 @@ class scann_test : public ::testing::TestWithParam { cuvs::preprocessing::quantize::pq::quantizer quantizer{ pq_params, - cuvs::neighbors::vpq_dataset{ + cuvs::neighbors::device_vpq_dataset{ std::move(vq_codebook), std::move(pq_codebook_copy), std::move(empty_data)}}; auto quantized_residuals_device = diff --git a/cpp/tests/neighbors/ann_vamana.cuh b/cpp/tests/neighbors/ann_vamana.cuh index 0397c74e1c..af2f9cb45e 100644 --- a/cpp/tests/neighbors/ann_vamana.cuh +++ b/cpp/tests/neighbors/ann_vamana.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,6 +9,7 @@ #include "ann_utils.cuh" #include +#include "cagra_padded_build_helpers.cuh" #include "naive_knn.cuh" #include @@ -207,10 +208,10 @@ class AnnVamanaTest : public ::testing::TestWithParam { handle_, index.graph().extent(0), index.graph().extent(1)); raft::linalg::map(handle_, graph_valid.view(), edge_op{}, index.graph()); - auto cagra_index = cagra::index(handle_, - ps.metric, - raft::make_const_mdspan(database_view), - raft::make_const_mdspan(graph_valid.view())); + cuvs::neighbors::test::padded_device_matrix_for_cagra cagra_base(handle_, + database_view); + auto cagra_index = cagra::device_padded_index( + handle_, ps.metric, cagra_base.view, raft::make_const_mdspan(graph_valid.view())); cagra::search_params search_params; search_params.algo = ps.algo; diff --git a/cpp/tests/neighbors/cagra_padded_build_helpers.cuh b/cpp/tests/neighbors/cagra_padded_build_helpers.cuh new file mode 100644 index 0000000000..215b845d9d --- /dev/null +++ b/cpp/tests/neighbors/cagra_padded_build_helpers.cuh @@ -0,0 +1,59 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include + +#include + +namespace cuvs::neighbors::test { + +/** + * Prepares a device_padded_dataset_view for cagra::build: uses make_device_padded_dataset_view when + * the source row stride already matches alignment, otherwise make_device_padded_dataset and keeps + * the copy in + * \p owned. The caller must keep this object alive for the lifetime of any index that only holds a + * view over the data. + */ +template +struct padded_device_matrix_for_cagra { + std::unique_ptr> owned; + cuvs::neighbors::device_padded_dataset_view view; + + padded_device_matrix_for_cagra( + raft::resources const& res, raft::device_matrix_view src) + : padded_device_matrix_for_cagra{build(res, src)} + { + } + + private: + struct build_result { + std::unique_ptr> owned; + cuvs::neighbors::device_padded_dataset_view view; + }; + + // device_padded_dataset_view has no default constructor; fill both members from one build step. + explicit padded_device_matrix_for_cagra(build_result&& br) + : owned{std::move(br.owned)}, view{std::move(br.view)} + { + } + + static auto build(raft::resources const& res, + raft::device_matrix_view src) + -> build_result + { + using namespace cuvs::neighbors; + if (matrix_row_width_matches_cagra_required(src)) { + return build_result{nullptr, make_device_padded_dataset_view(res, src)}; + } else { + auto own = make_device_padded_dataset(res, src); + auto vw = own->as_dataset_view(); + return build_result{std::move(own), vw}; + } + } +}; + +} // namespace cuvs::neighbors::test diff --git a/cpp/tests/neighbors/dynamic_batching/test_cagra.cu b/cpp/tests/neighbors/dynamic_batching/test_cagra.cu index 4c046367f2..336daab5f4 100644 --- a/cpp/tests/neighbors/dynamic_batching/test_cagra.cu +++ b/cpp/tests/neighbors/dynamic_batching/test_cagra.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -8,19 +8,36 @@ #include "../dynamic_batching.cuh" #include +#include namespace cuvs::neighbors::dynamic_batching { +namespace { + +template +auto build_cagra_with_dataset(raft::resources const& res, + cagra::index_params const& params, + raft::device_matrix_view dataset) + -> cagra::device_padded_index +{ + auto padded = cuvs::neighbors::make_device_padded_dataset_view(res, dataset); + auto index = cagra::build(res, params, padded); + index.update_dataset(res, padded); + return index; +} + +} // namespace + using cagra_F32 = dynamic_batching_test, - cagra::build, + cagra::device_padded_index, + build_cagra_with_dataset, cagra::search>; using cagra_U8 = dynamic_batching_test, - cagra::build, + cagra::device_padded_index, + build_cagra_with_dataset, cagra::search>; template diff --git a/cpp/tests/neighbors/hnsw.cu b/cpp/tests/neighbors/hnsw.cu index a92deaf725..01eeeb511c 100644 --- a/cpp/tests/neighbors/hnsw.cu +++ b/cpp/tests/neighbors/hnsw.cu @@ -5,8 +5,10 @@ #include "../test_utils.cuh" #include "ann_utils.cuh" +#include "cagra_padded_build_helpers.cuh" #include +#include #include #include #include @@ -93,8 +95,10 @@ class AnnHNSWTest : public ::testing::TestWithParam { 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 padded(handle_, database_view); - auto index = cuvs::neighbors::cagra::build(handle_, index_params, database_view); + auto index = cuvs::neighbors::cagra::build(handle_, index_params, padded.view); + index.update_dataset(handle_, padded.view); raft::resource::sync_stream(handle_); cuvs::neighbors::hnsw::search_params search_params; diff --git a/cpp/tests/neighbors/mg.cuh b/cpp/tests/neighbors/mg.cuh index fd5dc8e9dc..1ab97af3ae 100644 --- a/cpp/tests/neighbors/mg.cuh +++ b/cpp/tests/neighbors/mg.cuh @@ -1,11 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include "../test_utils.cuh" #include "ann_utils.cuh" +#include "cagra_padded_build_helpers.cuh" #include "naive_knn.cuh" #include @@ -218,15 +219,16 @@ class AnnMGTest : public ::testing::TestWithParam { d_mode = distribution_mode::SHARDED; mg_index_params index_params; - index_params.graph_build_params = cagra::graph_build_params::ivf_pq_params( - raft::matrix_extent(ps.num_db_vecs, ps.dim)); - index_params.mode = d_mode; + // Host dataset uses ACE build path; must set ace_params (not ivf_pq_params). + index_params.graph_build_params = cagra::graph_build_params::ace_params{}; + index_params.mode = d_mode; mg_search_params search_params; auto index_dataset = raft::make_host_matrix_view( h_index_dataset.data(), ps.num_db_vecs, ps.dim); - auto queries = raft::make_host_matrix_view( + auto index_dataset_view = cuvs::neighbors::make_host_standard_dataset_view(index_dataset); + auto queries = raft::make_host_matrix_view( h_queries.data(), ps.num_queries, ps.dim); auto neighbors = raft::make_host_matrix_view( neighbors_snmg_ann.data(), ps.num_queries, ps.k); @@ -235,11 +237,21 @@ class AnnMGTest : public ::testing::TestWithParam { tmp_index_file index_file; { - auto index = cuvs::neighbors::cagra::build(clique_, index_params, index_dataset); + auto index = cuvs::neighbors::cagra::build(clique_, index_params, index_dataset_view); cuvs::neighbors::cagra::serialize(clique_, index, index_file.filename); } auto new_index = - cuvs::neighbors::cagra::deserialize(clique_, index_file.filename); + cuvs::neighbors::mg_index, + DataT, + uint32_t>(clique_); + cuvs::neighbors::cagra::deserialize( + clique_, index_file.filename, &new_index); + auto index_dataset_device = raft::make_device_matrix_view( + d_index_dataset.data(), ps.num_db_vecs, ps.dim); + auto padded_index_dataset = + cuvs::neighbors::make_device_padded_dataset(clique_, index_dataset_device); + auto search_index = cuvs::neighbors::cagra::attach_padded_dataset_for_search( + clique_, new_index, padded_index_dataset->as_dataset_view()); if (ps.m_mode == m_mode_t::MERGE_ON_ROOT_RANK) search_params.merge_mode = MERGE_ON_ROOT_RANK; @@ -248,7 +260,7 @@ class AnnMGTest : public ::testing::TestWithParam { search_params.n_rows_per_batch = n_rows_per_search_batch; cuvs::neighbors::cagra::search( - clique_, new_index, search_params, queries, neighbors, distances); + clique_, search_index, search_params, queries, neighbors, distances); resource::sync_stream(clique_); double min_recall = static_cast(ps.nprobe) / static_cast(ps.nlist); @@ -376,7 +388,8 @@ class AnnMGTest : public ::testing::TestWithParam { { auto index_dataset = raft::make_device_matrix_view( d_index_dataset.data(), ps.num_db_vecs, ps.dim); - auto index = cuvs::neighbors::cagra::build(clique_, index_params, index_dataset); + auto standard_view = cuvs::neighbors::make_device_standard_dataset_view(index_dataset); + auto index = cuvs::neighbors::cagra::build(clique_, index_params, standard_view); cuvs::neighbors::cagra::serialize(clique_, index_file.filename, index); } @@ -388,13 +401,23 @@ class AnnMGTest : public ::testing::TestWithParam { distances_snmg_ann.data(), ps.num_queries, ps.k); auto distributed_index = - cuvs::neighbors::cagra::distribute(clique_, index_file.filename); + cuvs::neighbors::mg_index, + DataT, + uint32_t>(clique_); + cuvs::neighbors::cagra::distribute( + clique_, index_file.filename, &distributed_index); + auto index_dataset_device = raft::make_device_matrix_view( + d_index_dataset.data(), ps.num_db_vecs, ps.dim); + auto padded_index_dataset = + cuvs::neighbors::make_device_padded_dataset(clique_, index_dataset_device); + auto search_index = cuvs::neighbors::cagra::attach_padded_dataset_for_search( + clique_, distributed_index, padded_index_dataset->as_dataset_view()); search_params.merge_mode = TREE_MERGE; search_params.n_rows_per_batch = n_rows_per_search_batch; cuvs::neighbors::cagra::search( - clique_, distributed_index, search_params, queries, neighbors, distances); + clique_, search_index, search_params, queries, neighbors, distances); resource::sync_stream(clique_); @@ -554,19 +577,26 @@ class AnnMGTest : public ::testing::TestWithParam { ASSERT_TRUE(ps.num_queries <= 4); mg_index_params index_params; - index_params.graph_build_params = cagra::graph_build_params::ivf_pq_params( - raft::matrix_extent(ps.num_db_vecs, ps.dim)); - index_params.mode = REPLICATED; + // Host dataset uses ACE build path; must set ace_params (not ivf_pq_params). + index_params.graph_build_params = cagra::graph_build_params::ace_params{}; + index_params.mode = REPLICATED; mg_search_params search_params; search_params.search_mode = ROUND_ROBIN; auto index_dataset = raft::make_host_matrix_view( h_index_dataset.data(), ps.num_db_vecs, ps.dim); - auto small_batch_query = raft::make_host_matrix_view( + auto index_dataset_view = cuvs::neighbors::make_host_standard_dataset_view(index_dataset); + auto small_batch_query = raft::make_host_matrix_view( h_queries.data(), ps.num_queries, ps.dim); - auto index = cuvs::neighbors::cagra::build(clique_, index_params, index_dataset); + auto index = cuvs::neighbors::cagra::build(clique_, index_params, index_dataset_view); + auto index_dataset_device = raft::make_device_matrix_view( + d_index_dataset.data(), ps.num_db_vecs, ps.dim); + auto padded_index_dataset = + cuvs::neighbors::make_device_padded_dataset(clique_, index_dataset_device); + auto search_index = cuvs::neighbors::cagra::attach_padded_dataset_for_search( + clique_, index, padded_index_dataset->as_dataset_view()); int n_parallel_searches = 16; std::vector searches_correctness(n_parallel_searches); @@ -587,7 +617,7 @@ class AnnMGTest : public ::testing::TestWithParam { search_params.n_rows_per_batch = n_rows_per_search_batch; cuvs::neighbors::cagra::search(clique_, - index, + search_index, search_params, small_batch_query, small_batch_neighbors, diff --git a/cpp/tests/neighbors/tiered_index.cu b/cpp/tests/neighbors/tiered_index.cu index f342c44a2b..6b7ab19896 100644 --- a/cpp/tests/neighbors/tiered_index.cu +++ b/cpp/tests/neighbors/tiered_index.cu @@ -159,8 +159,28 @@ class ANNTieredIndexTest : public ::testing::TestWithParam auto indices_view = raft::make_device_matrix_view( (int64_t*)indices_tiered_dev.data(), ps.n_queries, ps.k); - cuvs::neighbors::tiered_index::search( - handle_, search_params, *final_index, queries_view, indices_view, distances_view); + if constexpr (std::is_same_v>) { + auto full_database_view = raft::make_device_matrix_view( + (const value_type*)database.data(), ps.n_rows, ps.dim); + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(full_database_view)) { + auto padded_view = + cuvs::neighbors::make_device_padded_dataset_view(handle_, full_database_view); + auto attached_index = cuvs::neighbors::tiered_index::attach_padded_dataset_for_search( + handle_, *final_index, padded_view); + cuvs::neighbors::tiered_index::search( + handle_, search_params, attached_index, queries_view, indices_view, distances_view); + } else { + auto padded_dataset = + cuvs::neighbors::make_device_padded_dataset(handle_, full_database_view); + auto attached_index = cuvs::neighbors::tiered_index::attach_padded_dataset_for_search( + handle_, *final_index, padded_dataset->as_dataset_view()); + cuvs::neighbors::tiered_index::search( + handle_, search_params, attached_index, queries_view, indices_view, distances_view); + } + } else { + cuvs::neighbors::tiered_index::search( + handle_, search_params, *final_index, queries_view, indices_view, distances_view); + } raft::update_host( distances_tiered.data(), distances_tiered_dev.data(), ps.n_queries * ps.k, stream_); @@ -216,7 +236,7 @@ const std::vector inputs = {10}, // n_queries {TEST_EXTEND, TEST_MERGE} // test_strategy ); -typedef ANNTieredIndexTest> CAGRA_F; +typedef ANNTieredIndexTest> CAGRA_F; TEST_P(CAGRA_F, AnnTieredIndex) { this->testTieredIndex(); } INSTANTIATE_TEST_CASE_P(ANNTieredIndexTest, CAGRA_F, ::testing::ValuesIn(inputs)); diff --git a/cpp/tests/neighbors/vpq_utils.cuh b/cpp/tests/neighbors/vpq_utils.cuh index 595e0ee416..23dba3218c 100644 --- a/cpp/tests/neighbors/vpq_utils.cuh +++ b/cpp/tests/neighbors/vpq_utils.cuh @@ -46,7 +46,7 @@ __global__ void decode_vpq_dataset_kernel(data_t* const decoded_dataset_ptr, template void decode_vpq_dataset(raft::device_matrix_view decoded_dataset, - const cuvs::neighbors::vpq_dataset& vpq_dataset, + const cuvs::neighbors::device_vpq_dataset& vpq_dataset, cudaStream_t cuda_stream) { const auto dataset_size = decoded_dataset.extent(0); diff --git a/examples/cpp/src/cagra_example.cu b/examples/cpp/src/cagra_example.cu index 856030c520..b7a96fa39a 100644 --- a/examples/cpp/src/cagra_example.cu +++ b/examples/cpp/src/cagra_example.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,6 +9,7 @@ #include #include +#include #include @@ -31,7 +32,9 @@ void cagra_build_search_simple(raft::device_resources const& dev_resources, cagra::index_params index_params; std::cout << "Building CAGRA index (search graph)" << std::endl; - auto index = cagra::build(dev_resources, index_params, dataset); + auto padded = cuvs::neighbors::make_device_padded_dataset_view(dev_resources, dataset); + auto index = cagra::build(dev_resources, index_params, padded); + index.update_dataset(dev_resources, padded); std::cout << "CAGRA index has " << index.size() << " vectors" << std::endl; std::cout << "CAGRA graph has degree " << index.graph_degree() << ", graph size [" diff --git a/examples/cpp/src/cagra_filter_udf_example.cu b/examples/cpp/src/cagra_filter_udf_example.cu index 0ab42dd580..1b331719b0 100644 --- a/examples/cpp/src/cagra_filter_udf_example.cu +++ b/examples/cpp/src/cagra_filter_udf_example.cu @@ -1,9 +1,10 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include +#include #include #include @@ -144,8 +145,9 @@ int main() index_params.intermediate_graph_degree); std::cout << "Building CAGRA index" << std::endl; - auto index = - cuvs::neighbors::cagra::build(res, index_params, raft::make_const_mdspan(dataset.view())); + auto padded = cuvs::neighbors::make_device_padded_dataset_view(res, dataset.view()); + auto index = cuvs::neighbors::cagra::build(res, index_params, padded); + index.update_dataset(res, padded); std::vector row_tenant_ids(n_rows); std::vector row_timestamps(n_rows); diff --git a/examples/cpp/src/cagra_hnsw_ace_example.cu b/examples/cpp/src/cagra_hnsw_ace_example.cu index d1bde25ad6..065fb8d4ec 100644 --- a/examples/cpp/src/cagra_hnsw_ace_example.cu +++ b/examples/cpp/src/cagra_hnsw_ace_example.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -12,6 +12,7 @@ #include #include +#include #include #include @@ -65,9 +66,13 @@ void cagra_build_search_ace(raft::device_resources const& dev_resources, raft::resource::sync_stream(dev_resources); auto dataset_host_view = raft::make_host_matrix_view( dataset_host.data_handle(), dataset_host.extent(0), dataset_host.extent(1)); + // Wrap in a host_padded_dataset_view. ACE graph construction is host-side CPU work and does not + // require CUDA row-alignment; construct the view directly to avoid the alignment check. + cuvs::neighbors::host_padded_dataset_view host_padded_view( + dataset_host_view, static_cast(dataset_host_view.extent(1))); std::cout << "Building CAGRA index (search graph)" << std::endl; - auto index = cagra::build(dev_resources, index_params, dataset_host_view); + auto ace_host_index = cagra::build(dev_resources, index_params, host_padded_view); // In-memory build of ACE provides the index in memory, so we can search it directly using // cagra::search @@ -80,7 +85,24 @@ void cagra_build_search_ace(raft::device_resources const& dev_resources, std::cout << "Converting CAGRA index to HNSW" << std::endl; hnsw::index_params hnsw_params; hnsw_params.hierarchy = hnsw::HnswHierarchy::GPU; // Offload hierarchy construction to GPU - auto hnsw_index = hnsw::from_cagra(dev_resources, hnsw_params, index); + + std::unique_ptr> hnsw_index; + std::unique_ptr> padded_owner; + if (ace_host_index.dataset_fd().has_value()) { + // Disk ACE path: ACE artifacts (dataset, graph, mapping) live on disk. Transfer file + // descriptors to a device index so from_cagra can serialize to hnsw_index.bin on disk. + cagra::device_padded_index device_index(dev_resources, + ace_host_index.metric()); + cagra::detail::fd_transfer::steal_disk_fds_to(dev_resources, ace_host_index, device_index); + hnsw_index = hnsw::from_cagra(dev_resources, hnsw_params, device_index, std::nullopt); + } else { + // In-memory ACE path: graph is in host memory. Upload the original dataset to device and + // attach it before from_cagra builds the HNSW hierarchy in memory. + padded_owner = cuvs::neighbors::make_device_padded_dataset(dev_resources, dataset_host_view); + auto device_index = cagra::attach_device_dataset_on_host_index( + dev_resources, ace_host_index, padded_owner->as_dataset_view()); + hnsw_index = hnsw::from_cagra(dev_resources, hnsw_params, device_index, dataset_host_view); + } // HNSW search requires host matrices auto queries_host = raft::make_host_matrix(n_queries, queries.extent(1)); @@ -116,8 +138,12 @@ void cagra_build_search_ace(raft::device_resources const& dev_resources, std::cout << "Deserializing HNSW index from disk for search." << std::endl; hnsw::index* hnsw_index_raw = nullptr; - hnsw::deserialize( - dev_resources, hnsw_params, hnsw_index_path, index.dim(), index.metric(), &hnsw_index_raw); + hnsw::deserialize(dev_resources, + hnsw_params, + hnsw_index_path, + ace_host_index.dim(), + ace_host_index.metric(), + &hnsw_index_raw); std::unique_ptr> hnsw_index_deserialized(hnsw_index_raw); diff --git a/examples/cpp/src/cagra_persistent_example.cu b/examples/cpp/src/cagra_persistent_example.cu index ded3a287b2..494a8230a3 100644 --- a/examples/cpp/src/cagra_persistent_example.cu +++ b/examples/cpp/src/cagra_persistent_example.cu @@ -1,11 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "common.cuh" #include +#include #include #include #include @@ -68,7 +69,9 @@ void cagra_build_search_variants(raft::device_resources const& res, cagra::index_params index_params; std::cout << "Building CAGRA index (search graph)" << std::endl; - auto index = cagra::build(res, index_params, dataset); + auto padded = cuvs::neighbors::make_device_padded_dataset_view(res, dataset); + auto index = cagra::build(res, index_params, padded); + index.update_dataset(res, padded); std::cout << "CAGRA index has " << index.size() << " vectors" << std::endl; std::cout << "CAGRA graph has degree " << index.graph_degree() << ", graph size [" diff --git a/examples/cpp/src/dynamic_batching_example.cu b/examples/cpp/src/dynamic_batching_example.cu index 317e2c5aff..d52d59f49a 100644 --- a/examples/cpp/src/dynamic_batching_example.cu +++ b/examples/cpp/src/dynamic_batching_example.cu @@ -1,11 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "common.cuh" #include +#include #include #include @@ -113,7 +114,9 @@ void dynamic_batching_example(raft::resources const& res, cagra::index_params orig_index_params; std::cout << "Building CAGRA index (search graph)" << std::endl; - auto orig_index = cagra::build(res, orig_index_params, dataset); + auto padded = cuvs::neighbors::make_device_padded_dataset_view(res, dataset); + auto orig_index = cagra::build(res, orig_index_params, padded); + orig_index.update_dataset(res, padded); std::cout << "CAGRA index has " << orig_index.size() << " vectors" << std::endl; std::cout << "CAGRA graph has degree " << orig_index.graph_degree() << ", graph size [" diff --git a/fern/pages/cpp_api/cpp-api-neighbors-common.md b/fern/pages/cpp_api/cpp-api-neighbors-common.md index 68430cf6fb..12528ad179 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-common.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-common.md @@ -70,7 +70,7 @@ struct dataset; ``` -### neighbors::vpq_dataset +### neighbors::device_vpq_dataset VPQ compressed dataset. @@ -81,7 +81,7 @@ The dataset is compressed using two level quantization ```cpp template -struct vpq_dataset : public dataset { +struct device_vpq_dataset : public dataset { raft::device_matrix vq_code_book; raft::device_matrix pq_code_book; raft::device_matrix data; diff --git a/fern/pages/cpp_api/cpp-api-preprocessing-quantize-pq.md b/fern/pages/cpp_api/cpp-api-preprocessing-quantize-pq.md index 85bcf86fbf..b70c451b3f 100644 --- a/fern/pages/cpp_api/cpp-api-preprocessing-quantize-pq.md +++ b/fern/pages/cpp_api/cpp-api-preprocessing-quantize-pq.md @@ -94,7 +94,7 @@ Defines and stores VPQ codebooks upon training template struct quantizer { params params_quantizer; - cuvs::neighbors::vpq_dataset vpq_codebooks; + cuvs::neighbors::device_vpq_dataset vpq_codebooks; }; ``` @@ -103,7 +103,7 @@ struct quantizer { | Name | Type | Description | | --- | --- | --- | | `params_quantizer` | [`params`](/api-reference/cpp-api-preprocessing-quantize-pq#preprocessing-quantize-pq-params) | Parameters used to build this quantizer. | -| `vpq_codebooks` | [`cuvs::neighbors::vpq_dataset`](/api-reference/cpp-api-neighbors-common#neighbors-vpq-dataset) | VPQ codebooks produced during training. | +| `vpq_codebooks` | [`cuvs::neighbors::device_vpq_dataset`](/api-reference/cpp-api-neighbors-common#neighbors-vpq-dataset) | VPQ codebooks produced during training. | ### preprocessing::quantize::pq::build diff --git a/go/cagra/cagra_test.go b/go/cagra/cagra_test.go index bb4fd0a0a1..9b6b2a4610 100644 --- a/go/cagra/cagra_test.go +++ b/go/cagra/cagra_test.go @@ -8,20 +8,6 @@ import ( ) func TestCagra(t *testing.T) { - testCases := []struct { - name string - compress bool - }{ - { - name: "No compression", - compress: false, - }, - { - name: "Compression", - compress: true, - }, - } - const ( nDataPoints = 1024 nFeatures = 16 @@ -31,129 +17,114 @@ func TestCagra(t *testing.T) { ) r := rand.New(rand.NewPCG(42, 0)) - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - resource, _ := cuvs.NewResource(nil) - defer resource.Close() - - testDataset := make([][]float32, nDataPoints) - for i := range testDataset { - testDataset[i] = make([]float32, nFeatures) - for j := range testDataset[i] { - testDataset[i][j] = r.Float32() - } - } + resource, _ := cuvs.NewResource(nil) + defer resource.Close() - dataset, err := cuvs.NewTensor(testDataset) - if err != nil { - t.Fatalf("error creating dataset tensor: %v", err) - } - defer dataset.Close() + testDataset := make([][]float32, nDataPoints) + for i := range testDataset { + testDataset[i] = make([]float32, nFeatures) + for j := range testDataset[i] { + testDataset[i][j] = r.Float32() + } + } - indexParams, err := CreateIndexParams() - if err != nil { - t.Fatalf("error creating index params: %v", err) - } - defer indexParams.Close() + dataset, err := cuvs.NewTensor(testDataset) + if err != nil { + t.Fatalf("error creating dataset tensor: %v", err) + } + defer dataset.Close() - index, err := CreateIndex() - if err != nil { - t.Fatalf("error creating index: %v", err) - } - defer index.Close() + indexParams, err := CreateIndexParams() + if err != nil { + t.Fatalf("error creating index params: %v", err) + } + defer indexParams.Close() - // Use the first 4 points from the dataset as queries : will test that we get them back - // as their own nearest neighbor - queries, err := cuvs.NewTensor(testDataset[:nQueries]) - if err != nil { - t.Fatalf("error creating queries tensor: %v", err) - } - defer queries.Close() + index, err := CreateIndex() + if err != nil { + t.Fatalf("error creating index: %v", err) + } + defer index.Close() - neighbors, err := cuvs.NewTensorOnDevice[uint32](&resource, []int64{int64(nQueries), int64(k)}) - if err != nil { - t.Fatalf("error creating neighbors tensor: %v", err) - } - defer neighbors.Close() + // Use the first 4 points from the dataset as queries : will test that we get them back + // as their own nearest neighbor + queries, err := cuvs.NewTensor(testDataset[:nQueries]) + if err != nil { + t.Fatalf("error creating queries tensor: %v", err) + } + defer queries.Close() - distances, err := cuvs.NewTensorOnDevice[float32](&resource, []int64{int64(nQueries), int64(k)}) - if err != nil { - t.Fatalf("error creating distances tensor: %v", err) - } - defer distances.Close() + neighbors, err := cuvs.NewTensorOnDevice[uint32](&resource, []int64{int64(nQueries), int64(k)}) + if err != nil { + t.Fatalf("error creating neighbors tensor: %v", err) + } + defer neighbors.Close() - if _, err := dataset.ToDevice(&resource); err != nil { - t.Fatalf("error moving dataset to device: %v", err) - } + distances, err := cuvs.NewTensorOnDevice[float32](&resource, []int64{int64(nQueries), int64(k)}) + if err != nil { + t.Fatalf("error creating distances tensor: %v", err) + } + defer distances.Close() - if tc.compress { - compressionParams, err := CreateCompressionParams() - if err != nil { - t.Fatalf("error creating compression params: %v", err) - } - indexParams.SetCompression(compressionParams) - } + if _, err := dataset.ToDevice(&resource); err != nil { + t.Fatalf("error moving dataset to device: %v", err) + } - if err := BuildIndex(resource, indexParams, &dataset, index); err != nil { - t.Fatalf("error building index: %v", err) - } + if err := BuildIndex(resource, indexParams, &dataset, index); err != nil { + t.Fatalf("error building index: %v", err) + } - if err := resource.Sync(); err != nil { - t.Fatalf("error syncing resource: %v", err) - } + if err := resource.Sync(); err != nil { + t.Fatalf("error syncing resource: %v", err) + } - if _, err := queries.ToDevice(&resource); err != nil { - t.Fatalf("error moving queries to device: %v", err) - } + if _, err := queries.ToDevice(&resource); err != nil { + t.Fatalf("error moving queries to device: %v", err) + } - SearchParams, err := CreateSearchParams() - if err != nil { - t.Fatalf("error creating search params: %v", err) - } - defer SearchParams.Close() + SearchParams, err := CreateSearchParams() + if err != nil { + t.Fatalf("error creating search params: %v", err) + } + defer SearchParams.Close() - err = SearchIndex(resource, SearchParams, index, &queries, &neighbors, &distances, nil) - if err != nil { - t.Fatalf("error searching index: %v", err) - } + err = SearchIndex(resource, SearchParams, index, &queries, &neighbors, &distances, nil) + if err != nil { + t.Fatalf("error searching index: %v", err) + } - if _, err := neighbors.ToHost(&resource); err != nil { - t.Fatalf("error moving neighbors to host: %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 := 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) - } + 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) - } + neighborsSlice, err := neighbors.Slice() + if err != nil { + t.Fatalf("error getting neighbors slice: %v", err) + } - for i := range neighborsSlice { - if neighborsSlice[i][0] != uint32(i) { - t.Error("wrong neighbor, expected", i, "got", neighborsSlice[i][0]) - } - } + for i := range neighborsSlice { + if neighborsSlice[i][0] != uint32(i) { + t.Error("wrong neighbor, expected", i, "got", neighborsSlice[i][0]) + } + } - distancesSlice, err := distances.Slice() - if err != nil { - t.Fatalf("error getting distances slice: %v", err) - } + distancesSlice, err := distances.Slice() + if err != nil { + t.Fatalf("error getting distances slice: %v", err) + } - if !tc.compress { - // Compress makes the result nondeterministic - for i := range distancesSlice { - if distancesSlice[i][0] >= epsilon || distancesSlice[i][0] <= -epsilon { - t.Error("distance should be close to 0, got", distancesSlice[i][0]) - } - } - } - }) + for i := range distancesSlice { + if distancesSlice[i][0] >= epsilon || distancesSlice[i][0] <= -epsilon { + t.Error("distance should be close to 0, got", distancesSlice[i][0]) + } } } diff --git a/go/cagra/index_params.go b/go/cagra/index_params.go index 99f4b70b93..c90ea95e46 100644 --- a/go/cagra/index_params.go +++ b/go/cagra/index_params.go @@ -13,11 +13,6 @@ type IndexParams struct { params C.cuvsCagraIndexParams_t } -// Supplemental parameters to build CAGRA Index -type CompressionParams struct { - params C.cuvsCagraCompressionParams_t -} - type BuildAlgo int const ( @@ -32,69 +27,6 @@ var cBuildAlgos = map[BuildAlgo]int{ AutoSelect: C.AUTO_SELECT, } -// Creates a new CompressionParams -func CreateCompressionParams() (*CompressionParams, error) { - var params C.cuvsCagraCompressionParams_t - - err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraCompressionParamsCreate(¶ms))) - if err != nil { - return nil, err - } - - if params == nil { - return nil, errors.New("memory allocation failed") - } - - return &CompressionParams{params: params}, nil -} - -// The bit length of the vector element after compression by PQ. -func (p *CompressionParams) SetPQBits(pq_bits uint32) (*CompressionParams, error) { - p.params.pq_bits = C.uint32_t(pq_bits) - - return p, nil -} - -// The dimensionality of the vector after compression by PQ. When zero, -// an optimal value is selected using a heuristic. -func (p *CompressionParams) SetPQDim(pq_dim uint32) (*CompressionParams, error) { - p.params.pq_dim = C.uint32_t(pq_dim) - - return p, nil -} - -// Vector Quantization (VQ) codebook size - number of "coarse cluster -// centers". When zero, an optimal value is selected using a heuristic. -func (p *CompressionParams) SetVQNCenters(vq_n_centers uint32) (*CompressionParams, error) { - p.params.vq_n_centers = C.uint32_t(vq_n_centers) - - return p, nil -} - -// The number of iterations searching for kmeans centers (both VQ & PQ -// phases). -func (p *CompressionParams) SetKMeansNIters(kmeans_n_iters uint32) (*CompressionParams, error) { - p.params.kmeans_n_iters = C.uint32_t(kmeans_n_iters) - - return p, nil -} - -// The fraction of data to use during iterative kmeans building (VQ -// phase). When zero, an optimal value is selected using a heuristic. -func (p *CompressionParams) SetVQKMeansTrainsetFraction(vq_kmeans_trainset_fraction float64) (*CompressionParams, error) { - p.params.vq_kmeans_trainset_fraction = C.double(vq_kmeans_trainset_fraction) - - return p, nil -} - -// The fraction of data to use during iterative kmeans building (PQ -// phase). When zero, an optimal value is selected using a heuristic. -func (p *CompressionParams) SetPQKMeansTrainsetFraction(pq_kmeans_trainset_fraction float64) (*CompressionParams, error) { - p.params.pq_kmeans_trainset_fraction = C.double(pq_kmeans_trainset_fraction) - - return p, nil -} - // Creates a new IndexParams func CreateIndexParams() (*IndexParams, error) { var params C.cuvsCagraIndexParams_t @@ -141,13 +73,6 @@ func (p *IndexParams) SetNNDescentNiter(nn_descent_niter uint32) (*IndexParams, return p, nil } -// Compression parameters -func (p *IndexParams) SetCompression(compression *CompressionParams) (*IndexParams, error) { - p.params.compression = C.cuvsCagraCompressionParams_t(compression.params) - - return p, nil -} - // Destroys IndexParams func (p *IndexParams) Close() error { err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraIndexParamsDestroy(p.params))) diff --git a/go/cagra/index_params_test.go b/go/cagra/index_params_test.go index 579419ca92..532cd736b1 100644 --- a/go/cagra/index_params_test.go +++ b/go/cagra/index_params_test.go @@ -4,225 +4,6 @@ import ( "testing" ) -// CompressionParams Tests -func TestCreateCompressionParams(t *testing.T) { - params, err := CreateCompressionParams() - if err != nil { - t.Fatalf("Failed to create CompressionParams: %v", err) - } - if params == nil { - t.Fatal("CreateCompressionParams returned nil params") - } - - if params.params == nil { - t.Fatal("CompressionParams internal params are nil") - } - if params.params.pq_kmeans_trainset_fraction != 0 { - t.Fatalf("Error params.params.pq_kmeans_trainset_fraction != 0, got = %v", params.params.pq_kmeans_trainset_fraction) - } - if params.params.pq_bits != 8 { - t.Fatalf("Error params.params.pq_bits != 8, got = %v", params.params.pq_bits) - } - if params.params.pq_dim != 0 { - t.Fatalf("Error params.params.pq_dim != 0, got = %v", params.params.pq_dim) - } - if params.params.vq_n_centers != 0 { - t.Fatalf("Error params.params.vq_n_centers != 0, got = %v", params.params.vq_n_centers) - } - if params.params.kmeans_n_iters != 25 { - t.Fatalf("Error params.params.kmeans_n_iters != 25, got = %v", params.params.kmeans_n_iters) - } -} - -func TestCompressionParamsSetPQBits(t *testing.T) { - params, err := CreateCompressionParams() - if err != nil { - t.Fatalf("Failed to create CompressionParams: %v", err) - } - - testCases := []struct { - name string - value uint32 - }{ - {"4 bits", 4}, - {"8 bits", 8}, - {"16 bits", 16}, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - result, err := params.SetPQBits(tc.value) - if err != nil { - t.Errorf("SetPQBits failed: %v", err) - } - if result != params { - t.Error("SetPQBits should return the same params instance") - } - if uint32(params.params.pq_bits) != tc.value { - t.Errorf("Expected pq_bits %d, got %d", tc.value, params.params.pq_bits) - } - }) - } -} - -func TestCompressionParamsSetPQDim(t *testing.T) { - params, err := CreateCompressionParams() - if err != nil { - t.Fatalf("Failed to create CompressionParams: %v", err) - } - - testCases := []struct { - name string - value uint32 - }{ - {"Zero (auto)", 0}, - {"Small dimension", 32}, - {"Large dimension", 128}, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - result, err := params.SetPQDim(tc.value) - if err != nil { - t.Errorf("SetPQDim failed: %v", err) - } - if result != params { - t.Error("SetPQDim should return the same params instance") - } - if uint32(params.params.pq_dim) != tc.value { - t.Errorf("Expected pq_dim %d, got %d", tc.value, params.params.pq_dim) - } - }) - } -} - -func TestCompressionParamsSetVQNCenters(t *testing.T) { - params, err := CreateCompressionParams() - if err != nil { - t.Fatalf("Failed to create CompressionParams: %v", err) - } - - testCases := []struct { - name string - value uint32 - }{ - {"Zero (auto)", 0}, - {"Small centers", 256}, - {"Large centers", 2048}, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - result, err := params.SetVQNCenters(tc.value) - if err != nil { - t.Errorf("SetVQNCenters failed: %v", err) - } - if result != params { - t.Error("SetVQNCenters should return the same params instance") - } - if uint32(params.params.vq_n_centers) != tc.value { - t.Errorf("Expected vq_n_centers %d, got %d", tc.value, params.params.vq_n_centers) - } - }) - } -} - -func TestCompressionParamsSetKMeansNIters(t *testing.T) { - params, err := CreateCompressionParams() - if err != nil { - t.Fatalf("Failed to create CompressionParams: %v", err) - } - - testCases := []struct { - name string - value uint32 - }{ - {"Few iterations", 10}, - {"Default iterations", 25}, - {"Many iterations", 100}, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - result, err := params.SetKMeansNIters(tc.value) - if err != nil { - t.Errorf("SetKMeansNIters failed: %v", err) - } - if result != params { - t.Error("SetKMeansNIters should return the same params instance") - } - if uint32(params.params.kmeans_n_iters) != tc.value { - t.Errorf("Expected kmeans_n_iters %d, got %d", tc.value, params.params.kmeans_n_iters) - } - }) - } -} - -func TestCompressionParamsSetVQKMeansTrainsetFraction(t *testing.T) { - params, err := CreateCompressionParams() - if err != nil { - t.Fatalf("Failed to create CompressionParams: %v", err) - } - - testCases := []struct { - name string - value float64 - }{ - {"Zero (auto)", 0.0}, - {"Half dataset", 0.5}, - {"Full dataset", 1.0}, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - result, err := params.SetVQKMeansTrainsetFraction(tc.value) - if err != nil { - t.Errorf("SetVQKMeansTrainsetFraction failed: %v", err) - } - if result != params { - t.Error("SetVQKMeansTrainsetFraction should return the same params instance") - } - if float64(params.params.vq_kmeans_trainset_fraction) != tc.value { - t.Errorf("Expected vq_kmeans_trainset_fraction %f, got %f", - tc.value, params.params.vq_kmeans_trainset_fraction) - } - }) - } -} - -func TestCompressionParamsSetPQKMeansTrainsetFraction(t *testing.T) { - params, err := CreateCompressionParams() - if err != nil { - t.Fatalf("Failed to create CompressionParams: %v", err) - } - - testCases := []struct { - name string - value float64 - }{ - {"Zero (auto)", 0.0}, - {"Quarter dataset", 0.25}, - {"Half dataset", 0.5}, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - result, err := params.SetPQKMeansTrainsetFraction(tc.value) - if err != nil { - t.Errorf("SetPQKMeansTrainsetFraction failed: %v", err) - } - if result != params { - t.Error("SetPQKMeansTrainsetFraction should return the same params instance") - } - if float64(params.params.pq_kmeans_trainset_fraction) != tc.value { - t.Errorf("Expected pq_kmeans_trainset_fraction %f, got %f", - tc.value, params.params.pq_kmeans_trainset_fraction) - } - }) - } -} - -// IndexParams Tests func TestCreateIndexParams(t *testing.T) { params, err := CreateIndexParams() if err != nil { @@ -376,31 +157,6 @@ func TestIndexParamsSetNNDescentNiter(t *testing.T) { } } -func TestIndexParamsSetCompression(t *testing.T) { - params, err := CreateIndexParams() - if err != nil { - t.Fatalf("Failed to create IndexParams: %v", err) - } - defer params.Close() - - compression, err := CreateCompressionParams() - if err != nil { - t.Fatalf("Failed to create CompressionParams: %v", err) - } - - // Configure compression params - compression.SetPQBits(8) - compression.SetPQDim(64) - - result, err := params.SetCompression(compression) - if err != nil { - t.Errorf("SetCompression failed: %v", err) - } - if result != params { - t.Error("SetCompression should return the same params instance") - } -} - func TestIndexParamsClose(t *testing.T) { params, err := CreateIndexParams() if err != nil { @@ -414,7 +170,6 @@ func TestIndexParamsClose(t *testing.T) { } func TestBuildAlgoConstants(t *testing.T) { - // Test that BuildAlgo constants are properly defined algos := []BuildAlgo{IvfPq, NnDescent, AutoSelect} for _, algo := range algos { diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java index e185ed9f26..66fb45595d 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs; @@ -25,7 +25,6 @@ public class CagraIndexParams { private final int numWriterThreads; private final CuVSIvfPqParams cuVSIvfPqParams; private final CuVSAceParams cuVSAceParams; - private final CagraCompressionParams cagraCompressionParams; /** * Enum that denotes which ANN algorithm is used to build CAGRA graph. @@ -335,8 +334,7 @@ private CagraIndexParams( int writerThreads, CuvsDistanceType cuvsDistanceType, CuVSIvfPqParams cuVSIvfPqParams, - CuVSAceParams cuVSAceParams, - CagraCompressionParams cagraCompressionParams) { + CuVSAceParams cuVSAceParams) { this.intermediateGraphDegree = intermediateGraphDegree; this.graphDegree = graphDegree; this.cuvsCagraGraphBuildAlgo = CuvsCagraGraphBuildAlgo; @@ -345,7 +343,6 @@ private CagraIndexParams( this.cuvsDistanceType = cuvsDistanceType; this.cuVSIvfPqParams = cuVSIvfPqParams; this.cuVSAceParams = cuVSAceParams; - this.cagraCompressionParams = cagraCompressionParams; } public static CagraIndexParams fromHnswParams( @@ -427,13 +424,6 @@ public CagraGraphBuildAlgo getCuvsCagraGraphBuildAlgo() { return cuvsCagraGraphBuildAlgo; } - /** - * Gets the CAGRA compression parameters. - */ - public CagraCompressionParams getCagraCompressionParams() { - return cagraCompressionParams; - } - @Override public String toString() { return "CagraIndexParams [cuvsCagraGraphBuildAlgo=" @@ -452,8 +442,6 @@ public String toString() { + cuVSIvfPqParams + ", cuVSAceParams=" + cuVSAceParams - + ", cagraCompressionParams=" - + cagraCompressionParams + "]"; } @@ -470,7 +458,6 @@ public static class Builder { private int numWriterThreads = 2; private CuVSIvfPqParams cuVSIvfPqParams = new CuVSIvfPqParams.Builder().build(); private CuVSAceParams cuVSAceParams = new CuVSAceParams.Builder().build(); - private CagraCompressionParams cagraCompressionParams; public Builder() {} @@ -564,18 +551,6 @@ public Builder withCuVSAceParams(CuVSAceParams cuVSAceParams) { return this; } - /** - * Registers an instance of configured {@link CagraCompressionParams} with this - * Builder. - * - * @param cagraCompressionParams An instance of CagraCompressionParams. - * @return An instance of this Builder. - */ - public Builder withCompressionParams(CagraCompressionParams cagraCompressionParams) { - this.cagraCompressionParams = cagraCompressionParams; - return this; - } - /** * Builds an instance of {@link CagraIndexParams}. * @@ -590,8 +565,7 @@ public CagraIndexParams build() { numWriterThreads, cuvsDistanceType, cuVSIvfPqParams, - cuVSAceParams, - cagraCompressionParams); + cuVSAceParams); } } } 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 abc53a5945..2080ea6187 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 @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.internal; @@ -545,28 +545,6 @@ private static void populateNativeIndexParams( cuvsCagraIndexParams.nn_descent_niter(indexPtr, params.getNNDescentNumIterations()); cuvsCagraIndexParams.metric(indexPtr, params.getCuvsDistanceType().value); - CagraCompressionParams cagraCompressionParams = params.getCagraCompressionParams(); - if (cagraCompressionParams != null) { - var compressionParams = createCagraCompressionParams(); - handles.add(compressionParams); - MemorySegment cuvsCagraCompressionParamsMemorySegment = compressionParams.handle(); - cuvsCagraCompressionParams.pq_bits( - cuvsCagraCompressionParamsMemorySegment, cagraCompressionParams.getPqBits()); - cuvsCagraCompressionParams.pq_dim( - cuvsCagraCompressionParamsMemorySegment, cagraCompressionParams.getPqDim()); - cuvsCagraCompressionParams.vq_n_centers( - cuvsCagraCompressionParamsMemorySegment, cagraCompressionParams.getVqNCenters()); - cuvsCagraCompressionParams.kmeans_n_iters( - cuvsCagraCompressionParamsMemorySegment, cagraCompressionParams.getKmeansNIters()); - cuvsCagraCompressionParams.vq_kmeans_trainset_fraction( - cuvsCagraCompressionParamsMemorySegment, - cagraCompressionParams.getVqKmeansTrainsetFraction()); - cuvsCagraCompressionParams.pq_kmeans_trainset_fraction( - cuvsCagraCompressionParamsMemorySegment, - cagraCompressionParams.getPqKmeansTrainsetFraction()); - cuvsCagraIndexParams.compression(indexPtr, cuvsCagraCompressionParamsMemorySegment); - } - if (params.getCagraGraphBuildAlgo().equals(CagraGraphBuildAlgo.IVF_PQ)) { var ivfPqIndexParams = createIvfPqIndexParams(); diff --git a/python/cuvs/cuvs/neighbors/cagra/__init__.py b/python/cuvs/cuvs/neighbors/cagra/__init__.py index ec70305d72..3b31d23cd7 100644 --- a/python/cuvs/cuvs/neighbors/cagra/__init__.py +++ b/python/cuvs/cuvs/neighbors/cagra/__init__.py @@ -1,10 +1,9 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from .cagra import ( AceParams, - CompressionParams, ExtendParams, Index, IndexParams, @@ -19,7 +18,6 @@ __all__ = [ "AceParams", - "CompressionParams", "ExtendParams", "Index", "IndexParams", diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd index e575ed5360..b10d9b11f4 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # cython: language_level=3 @@ -42,16 +42,6 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: ITERATIVE_CAGRA_SEARCH ACE - ctypedef struct cuvsCagraCompressionParams: - uint32_t pq_bits - uint32_t pq_dim - uint32_t vq_n_centers - uint32_t kmeans_n_iters - double vq_kmeans_trainset_fraction - double pq_kmeans_trainset_fraction - - ctypedef cuvsCagraCompressionParams* cuvsCagraCompressionParams_t - ctypedef struct cuvsIvfPqParams: cuvsIvfPqIndexParams_t ivf_pq_build_params cuvsIvfPqSearchParams_t ivf_pq_search_params @@ -73,7 +63,6 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: size_t graph_degree cuvsCagraGraphBuildAlgo build_algo size_t nn_descent_niter - cuvsCagraCompressionParams_t compression void* graph_build_params ctypedef cuvsCagraIndexParams* cuvsCagraIndexParams_t @@ -115,11 +104,20 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: ctypedef cuvsCagraIndex* cuvsCagraIndex_t - cuvsError_t cuvsCagraCompressionParamsCreate( - cuvsCagraCompressionParams_t* params) + ctypedef enum cuvsDatasetLayout_t: + CUVS_DATASET_LAYOUT_STANDARD + CUVS_DATASET_LAYOUT_PADDED + + ctypedef enum cuvsDatasetStorageKind_t: + CUVS_DATASET_STORAGE_KIND_EXTENDED + CUVS_DATASET_STORAGE_KIND_MERGED + + ctypedef struct cuvsDatasetStorage: + uintptr_t addr + DLDataType dtype + cuvsDatasetStorageKind_t kind - cuvsError_t cuvsCagraCompressionParamsDestroy( - cuvsCagraCompressionParams_t index) + ctypedef cuvsDatasetStorage* cuvsDatasetStorage_t cuvsError_t cuvsAceParamsCreate(cuvsAceParams_t* params) @@ -170,6 +168,7 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: cuvsError_t cuvsCagraDeserialize(cuvsResources_t res, const char * filename, + cuvsDatasetLayout_t deserialize_layout, cuvsCagraIndex_t index) cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res, @@ -185,9 +184,15 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params) cuvsError_t cuvsCagraExtendParamsDestroy(cuvsCagraExtendParams_t params) + cuvsError_t cuvsMakeExtendedStorage(cuvsResources_t res, + DLManagedTensor* additional_dataset, + cuvsCagraIndex_t index, + cuvsDatasetStorage_t* extended_storage) + cuvsError_t cuvsDatasetStorageDestroy(cuvsDatasetStorage_t dataset_storage) cuvsError_t cuvsCagraExtend(cuvsResources_t res, cuvsCagraExtendParams_t params, DLManagedTensor* additional_dataset, + cuvsDatasetStorage_t extended_dataset, cuvsCagraIndex_t index) @@ -204,7 +209,6 @@ cdef class Index: cdef class IndexParams: cdef cuvsCagraIndexParams* params - cdef public object compression cdef public object ivf_pq_build_params cdef public object ivf_pq_search_params cdef public object ace_params diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index 8e3bca3ab2..aa77c0e5ef 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # cython: language_level=3 @@ -47,83 +47,6 @@ from cuvs.neighbors import ivf_pq from cuvs.neighbors.filters import no_filter -cdef class CompressionParams: - """ - Parameters for VPQ Compression - - Parameters - ---------- - pq_bits: int - The bit length of the vector element after compression by PQ. - Possible values: [4, 5, 6, 7, 8]. The smaller the 'pq_bits', the - smaller the index size and the better the search performance, but - the lower the recall. - pq_dim: int - The dimensionality of the vector after compression by PQ. When zero, - an optimal value is selected using a heuristic. - vq_n_centers: int - Vector Quantization (VQ) codebook size - number of "coarse cluster - centers". When zero, an optimal value is selected using a heuristic. - kmeans_n_iters: int - The number of iterations searching for kmeans centers (both VQ & PQ - phases). - vq_kmeans_trainset_fraction: float - The fraction of data to use during iterative kmeans building (VQ - phase). When zero, an optimal value is selected using a heuristic. - pq_kmeans_trainset_fraction: float - The fraction of data to use during iterative kmeans building (PQ - phase). When zero, an optimal value is selected using a heuristic. - """ - cdef cuvsCagraCompressionParams * params - - def __cinit__(self): - check_cuvs(cuvsCagraCompressionParamsCreate(&self.params)) - - def __dealloc__(self): - check_cuvs(cuvsCagraCompressionParamsDestroy(self.params)) - - def __init__(self, *, - pq_bits=8, - pq_dim=0, - vq_n_centers=0, - kmeans_n_iters=25, - vq_kmeans_trainset_fraction=0.0, - pq_kmeans_trainset_fraction=0.0): - self.params.pq_bits = pq_bits - self.params.pq_dim = pq_dim - self.params.vq_n_centers = vq_n_centers - self.params.kmeans_n_iters = kmeans_n_iters - self.params.vq_kmeans_trainset_fraction = vq_kmeans_trainset_fraction - self.params.pq_kmeans_trainset_fraction = pq_kmeans_trainset_fraction - - @property - def pq_bits(self): - return self.params.pq_bits - - @property - def pq_dim(self): - return self.params.pq_dim - - @property - def vq_n_centers(self): - return self.params.vq_n_centers - - @property - def kmeans_n_iters(self): - return self.params.kmeans_n_iters - - @property - def vq_kmeans_trainset_fraction(self): - return self.params.vq_kmeans_trainset_fraction - - @property - def pq_kmeans_trainset_fraction(self): - return self.params.pq_kmeans_trainset_fraction - - def get_handle(self): - return self.params - - cdef class AceParams: """ Parameters for ACE (Augmented Core Extraction) graph building algorithm. @@ -271,9 +194,6 @@ cdef class IndexParams: - ace will use ACE (Augmented Core Extraction) for building indices for datasets too large to fit in GPU memory - compression: CompressionParams, optional - If compression is desired should be a CompressionParams object. If None - compression will be disabled. ivf_pq_build_params: cuvs.neighbors.ivf_pq.IndexParams, optional Parameters for IVF-PQ algorithm. If provided, it will be used for building the graph. @@ -289,7 +209,6 @@ cdef class IndexParams: def __cinit__(self): check_cuvs(cuvsCagraIndexParamsCreate(&self.params)) - self.compression = None self.ivf_pq_build_params = None self.ivf_pq_search_params = None self.ace_params = None @@ -304,7 +223,6 @@ cdef class IndexParams: graph_degree=64, build_algo="ivf_pq", nn_descent_niter=20, - compression=None, ivf_pq_build_params: ivf_pq.IndexParams = None, ivf_pq_search_params: ivf_pq.SearchParams = None, ace_params: AceParams = None, @@ -329,10 +247,6 @@ cdef class IndexParams: raise ValueError(f"Unknown build_algo '{build_algo}'") self.params.nn_descent_niter = nn_descent_niter - if compression is not None: - self.compression = compression - self.params.compression = \ - compression.get_handle() # Handle graph build params based on build algorithm if build_algo == "ace": @@ -951,7 +865,7 @@ def save(filename, Index index, bool include_dataset=True, resources=None): @auto_sync_resources -def load(filename, resources=None): +def load(filename, layout="padded", resources=None): """ Loads index from file. @@ -963,6 +877,8 @@ def load(filename, resources=None): ---------- filename : string Name of the file. + layout : {"padded", "standard"}, default = "padded" + Target index layout for deserialization. {resources_docstring} Returns @@ -973,10 +889,19 @@ def load(filename, resources=None): cdef Index idx = Index() cdef cuvsResources_t res = resources.get_c_obj() cdef string c_filename = filename.encode('utf-8') + cdef cuvsDatasetLayout_t deserialize_layout + + if layout == "padded": + deserialize_layout = CUVS_DATASET_LAYOUT_PADDED + elif layout == "standard": + deserialize_layout = CUVS_DATASET_LAYOUT_STANDARD + else: + raise ValueError("layout must be either 'padded' or 'standard'") check_cuvs(cuvsCagraDeserialize( res, c_filename.c_str(), + deserialize_layout, idx.index )) idx.trained = True @@ -1087,13 +1012,26 @@ def extend(ExtendParams params, Index index, additional_dataset, cydlpack.dlpack_c(dataset_ai) cdef cuvsResources_t res = resources.get_c_obj() + cdef cuvsDatasetStorage_t extended_dataset = NULL - with cuda_interruptible(): - check_cuvs(cuvsCagraExtend( - res, - params.params, - dataset_dlpack, - index.index - )) + check_cuvs(cuvsMakeExtendedStorage( + res, + dataset_dlpack, + index.index, + &extended_dataset + )) + + try: + with cuda_interruptible(): + check_cuvs(cuvsCagraExtend( + res, + params.params, + dataset_dlpack, + extended_dataset, + index.index + )) + finally: + if extended_dataset != NULL: + check_cuvs(cuvsDatasetStorageDestroy(extended_dataset)) return index diff --git a/python/cuvs/cuvs/tests/test_cagra.py b/python/cuvs/cuvs/tests/test_cagra.py index c0d436951e..8f468defdc 100644 --- a/python/cuvs/cuvs/tests/test_cagra.py +++ b/python/cuvs/cuvs/tests/test_cagra.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -34,7 +34,6 @@ def run_cagra_build_search_test( inplace=True, test_extend=False, search_params={}, - compression=None, serialize=False, ): dataset = generate_data((n_rows, n_cols), dtype) @@ -49,7 +48,6 @@ def run_cagra_build_search_test( intermediate_graph_degree=intermediate_graph_degree, graph_degree=graph_degree, build_algo=build_algo, - compression=compression, ) if test_extend: @@ -129,27 +127,26 @@ def run_cagra_build_search_test( cp_graph = cp.array(graph) assert cp_graph.shape == (n_rows, graph_degree) - if compression is None: - # make sure we can get the dataset from the cagra index - dataset_from_index = index.dataset + # make sure we can get the dataset from the cagra index + dataset_from_index = index.dataset - dataset_from_index_host = dataset_from_index.copy_to_host() - assert np.allclose(dataset, dataset_from_index_host) + dataset_from_index_host = dataset_from_index.copy_to_host() + assert np.allclose(dataset, dataset_from_index_host) - # make sure we can reconstruct the index from the graph - # Note that we can't actually use the dataset from the index itself - # - since that is a strided matrix (and we expect non-strided inputs - # in the C++ cagra::build api), so we are using the host version - # which will have been copied into a non-strided layout - reloaded_index = cagra.from_graph( - graph, dataset_from_index_host, metric=metric - ) + # make sure we can reconstruct the index from the graph + # Note that we can't actually use the dataset from the index itself + # - since that is a strided matrix (and we expect non-strided inputs + # in the C++ cagra::build api), so we are using the host version + # which will have been copied into a non-strided layout + reloaded_index = cagra.from_graph( + graph, dataset_from_index_host, metric=metric + ) - dist_device, idx_device = cagra.search( - search_params, reloaded_index, queries_device, k - ) - recall = calc_recall(idx_device.copy_to_host(), skl_idx) - assert recall > 0.9 + dist_device, idx_device = cagra.search( + search_params, reloaded_index, queries_device, k + ) + recall = calc_recall(idx_device.copy_to_host(), skl_idx) + assert recall > 0.9 @pytest.mark.parametrize("inplace", [True, False]) @@ -234,14 +231,6 @@ def test_cagra_index_params(params): ) -def test_cagra_vpq_compression(): - dim = 64 - pq_len = 2 - run_cagra_build_search_test( - n_cols=dim, compression=cagra.CompressionParams(pq_dim=dim / pq_len) - ) - - @pytest.mark.parametrize("internal_dtype", [np.float32, np.float16, np.uint8]) def test_cagra_ivf_pq( internal_dtype, diff --git a/rust/cuvs-sys/src/bindings.rs b/rust/cuvs-sys/src/bindings.rs index 171af6f422..545b7e4e16 100644 --- a/rust/cuvs-sys/src/bindings.rs +++ b/rust/cuvs-sys/src/bindings.rs @@ -5,42 +5,25 @@ 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)] @@ -52,33 +35,21 @@ 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)] @@ -89,23 +60,15 @@ 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)] @@ -121,15 +84,11 @@ 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)] @@ -144,7 +103,6 @@ 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 { @@ -152,15 +110,12 @@ 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, @@ -172,23 +127,18 @@ 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, @@ -197,27 +147,22 @@ 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, @@ -225,12 +170,10 @@ 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, @@ -238,12 +181,10 @@ 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, @@ -251,7 +192,6 @@ 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, @@ -260,7 +200,6 @@ 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, @@ -269,7 +208,6 @@ 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, @@ -278,27 +216,22 @@ unsafe extern "C" { } 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, @@ -307,7 +240,6 @@ 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, @@ -317,98 +249,172 @@ unsafe extern "C" { ) -> cuvsError_t; } #[repr(u32)] -#[doc = " enum to tell how to compute distance"] +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +pub enum cuvsFilterType { + NO_FILTER = 0, + BITSET = 1, + BITMAP = 2, +} +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct cuvsFilter { + pub addr: usize, + pub type_: cuvsFilterType, +} +#[allow(clippy::unnecessary_operation, clippy::identity_op)] +const _: () = { + ["Size of cuvsFilter"][::std::mem::size_of::() - 16usize]; + ["Alignment of cuvsFilter"][::std::mem::align_of::() - 8usize]; + ["Offset of field: cuvsFilter::addr"][::std::mem::offset_of!(cuvsFilter, addr) - 0usize]; + ["Offset of field: cuvsFilter::type_"][::std::mem::offset_of!(cuvsFilter, type_) - 8usize]; +}; +#[repr(u32)] +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +pub enum cuvsMergeStrategy { + MERGE_STRATEGY_PHYSICAL = 0, + MERGE_STRATEGY_LOGICAL = 1, +} +#[repr(u32)] +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +pub enum cuvsDatasetLayout_t { + CUVS_DATASET_LAYOUT_STANDARD = 0, + CUVS_DATASET_LAYOUT_PADDED = 1, +} +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct cuvsDataset { + pub addr: usize, + pub dtype: DLDataType, + pub layout: cuvsDatasetLayout_t, +} +#[allow(clippy::unnecessary_operation, clippy::identity_op)] +const _: () = { + ["Size of cuvsDataset"][::std::mem::size_of::() - 16usize]; + ["Alignment of cuvsDataset"][::std::mem::align_of::() - 8usize]; + ["Offset of field: cuvsDataset::addr"][::std::mem::offset_of!(cuvsDataset, addr) - 0usize]; + ["Offset of field: cuvsDataset::dtype"][::std::mem::offset_of!(cuvsDataset, dtype) - 8usize]; + ["Offset of field: cuvsDataset::layout"][::std::mem::offset_of!(cuvsDataset, layout) - 12usize]; +}; +pub type cuvsDataset_t = *mut cuvsDataset; +pub type cuvsDatasetPadded = cuvsDataset; +pub type cuvsDatasetPadded_t = *mut cuvsDataset; +pub type cuvsDatasetStandard = cuvsDataset; +pub type cuvsDatasetStandard_t = *mut cuvsDataset; +#[repr(u32)] +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +pub enum cuvsDatasetStorageKind_t { + CUVS_DATASET_STORAGE_KIND_EXTENDED = 0, + CUVS_DATASET_STORAGE_KIND_MERGED = 1, +} +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct cuvsDatasetStorage { + pub addr: usize, + pub dtype: DLDataType, + pub kind: cuvsDatasetStorageKind_t, +} +#[allow(clippy::unnecessary_operation, clippy::identity_op)] +const _: () = { + ["Size of cuvsDatasetStorage"][::std::mem::size_of::() - 16usize]; + ["Alignment of cuvsDatasetStorage"][::std::mem::align_of::() - 8usize]; + ["Offset of field: cuvsDatasetStorage::addr"] + [::std::mem::offset_of!(cuvsDatasetStorage, addr) - 0usize]; + ["Offset of field: cuvsDatasetStorage::dtype"] + [::std::mem::offset_of!(cuvsDatasetStorage, dtype) - 8usize]; + ["Offset of field: cuvsDatasetStorage::kind"] + [::std::mem::offset_of!(cuvsDatasetStorage, kind) - 12usize]; +}; +pub type cuvsDatasetStorage_t = *mut cuvsDatasetStorage; +pub type cuvsCagraIndex_t = *mut cuvsCagraIndex; +unsafe extern "C" { + #[must_use] + pub fn cuvsDatasetMakePadded( + res: cuvsResources_t, + dataset: *mut DLManagedTensor, + padded_dataset: *mut cuvsDatasetPadded_t, + ) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsDatasetPaddedDestroy(padded_dataset: cuvsDatasetPadded_t) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsMakeExtendedStorage( + res: cuvsResources_t, + additional_dataset: *mut DLManagedTensor, + index: cuvsCagraIndex_t, + extended_storage: *mut cuvsDatasetStorage_t, + ) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsMakeMergedStorage( + res: cuvsResources_t, + indices: *mut cuvsCagraIndex_t, + num_indices: usize, + filter: cuvsFilter, + merged_storage: *mut cuvsDatasetStorage_t, + ) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsDatasetStorageDestroy(dataset_storage: cuvsDatasetStorage_t) -> cuvsError_t; +} +#[repr(u32)] #[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 = " Check inertia during iterations for early convergence."] pub inertia_check: bool, - #[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 streaming_batch_size: i64, + pub init_size: i64, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] const _: () = { - ["Size of cuvsKMeansParams"][::std::mem::size_of::() - 64usize]; + ["Size of cuvsKMeansParams"][::std::mem::size_of::() - 72usize]; ["Alignment of cuvsKMeansParams"][::std::mem::align_of::() - 8usize]; ["Offset of field: cuvsKMeansParams::metric"] [::std::mem::offset_of!(cuvsKMeansParams, metric) - 0usize]; @@ -436,20 +442,76 @@ const _: () = { [::std::mem::offset_of!(cuvsKMeansParams, hierarchical_n_iters) - 52usize]; ["Offset of field: cuvsKMeansParams::streaming_batch_size"] [::std::mem::offset_of!(cuvsKMeansParams, streaming_batch_size) - 56usize]; + ["Offset of field: cuvsKMeansParams::init_size"] + [::std::mem::offset_of!(cuvsKMeansParams, init_size) - 64usize]; +}; +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct cuvsKMeansParams_v2 { + pub metric: cuvsDistanceType, + pub n_clusters: ::std::os::raw::c_int, + pub init: cuvsKMeansInitMethod, + pub max_iter: ::std::os::raw::c_int, + pub tol: f64, + pub n_init: ::std::os::raw::c_int, + pub oversampling_factor: f64, + pub batch_samples: ::std::os::raw::c_int, + pub batch_centroids: ::std::os::raw::c_int, + pub hierarchical: bool, + pub hierarchical_n_iters: ::std::os::raw::c_int, + pub streaming_batch_size: i64, + pub init_size: i64, +} +#[allow(clippy::unnecessary_operation, clippy::identity_op)] +const _: () = { + ["Size of cuvsKMeansParams_v2"][::std::mem::size_of::() - 72usize]; + ["Alignment of cuvsKMeansParams_v2"][::std::mem::align_of::() - 8usize]; + ["Offset of field: cuvsKMeansParams_v2::metric"] + [::std::mem::offset_of!(cuvsKMeansParams_v2, metric) - 0usize]; + ["Offset of field: cuvsKMeansParams_v2::n_clusters"] + [::std::mem::offset_of!(cuvsKMeansParams_v2, n_clusters) - 4usize]; + ["Offset of field: cuvsKMeansParams_v2::init"] + [::std::mem::offset_of!(cuvsKMeansParams_v2, init) - 8usize]; + ["Offset of field: cuvsKMeansParams_v2::max_iter"] + [::std::mem::offset_of!(cuvsKMeansParams_v2, max_iter) - 12usize]; + ["Offset of field: cuvsKMeansParams_v2::tol"] + [::std::mem::offset_of!(cuvsKMeansParams_v2, tol) - 16usize]; + ["Offset of field: cuvsKMeansParams_v2::n_init"] + [::std::mem::offset_of!(cuvsKMeansParams_v2, n_init) - 24usize]; + ["Offset of field: cuvsKMeansParams_v2::oversampling_factor"] + [::std::mem::offset_of!(cuvsKMeansParams_v2, oversampling_factor) - 32usize]; + ["Offset of field: cuvsKMeansParams_v2::batch_samples"] + [::std::mem::offset_of!(cuvsKMeansParams_v2, batch_samples) - 40usize]; + ["Offset of field: cuvsKMeansParams_v2::batch_centroids"] + [::std::mem::offset_of!(cuvsKMeansParams_v2, batch_centroids) - 44usize]; + ["Offset of field: cuvsKMeansParams_v2::hierarchical"] + [::std::mem::offset_of!(cuvsKMeansParams_v2, hierarchical) - 48usize]; + ["Offset of field: cuvsKMeansParams_v2::hierarchical_n_iters"] + [::std::mem::offset_of!(cuvsKMeansParams_v2, hierarchical_n_iters) - 52usize]; + ["Offset of field: cuvsKMeansParams_v2::streaming_batch_size"] + [::std::mem::offset_of!(cuvsKMeansParams_v2, streaming_batch_size) - 56usize]; + ["Offset of field: cuvsKMeansParams_v2::init_size"] + [::std::mem::offset_of!(cuvsKMeansParams_v2, init_size) - 64usize]; }; pub type cuvsKMeansParams_t = *mut cuvsKMeansParams; +pub type cuvsKMeansParams_v2_t = *mut cuvsKMeansParams_v2; 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; } +unsafe extern "C" { + #[must_use] + pub fn cuvsKMeansParamsCreate_v2(params: *mut cuvsKMeansParams_v2_t) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsKMeansParamsDestroy_v2(params: cuvsKMeansParams_v2_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, @@ -457,7 +519,6 @@ 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 streamed to the GPU in\n batches controlled by params->streaming_batch_size.\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, @@ -470,7 +531,18 @@ 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 cuvsKMeansFit_v2( + res: cuvsResources_t, + params: cuvsKMeansParams_v2_t, + X: *mut DLManagedTensor, + sample_weight: *mut DLManagedTensor, + centroids: *mut DLManagedTensor, + inertia: *mut f64, + n_iter: *mut ::std::os::raw::c_int, + ) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] pub fn cuvsKMeansPredict( res: cuvsResources_t, params: cuvsKMeansParams_t, @@ -484,7 +556,19 @@ 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 cuvsKMeansPredict_v2( + res: cuvsResources_t, + params: cuvsKMeansParams_v2_t, + X: *mut DLManagedTensor, + sample_weight: *mut DLManagedTensor, + centroids: *mut DLManagedTensor, + labels: *mut DLManagedTensor, + normalize_weight: bool, + inertia: *mut f64, + ) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] pub fn cuvsKMeansClusterCost( res: cuvsResources_t, X: *mut DLManagedTensor, @@ -494,7 +578,6 @@ 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, @@ -505,48 +588,32 @@ 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)] @@ -583,28 +650,19 @@ 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)] @@ -628,15 +686,12 @@ 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 { @@ -655,47 +710,38 @@ 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, @@ -703,7 +749,6 @@ 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, @@ -711,7 +756,6 @@ 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, @@ -719,7 +763,6 @@ 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, @@ -727,7 +770,6 @@ 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, @@ -735,7 +777,6 @@ 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, @@ -743,7 +784,6 @@ 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, @@ -754,7 +794,6 @@ 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, @@ -763,7 +802,6 @@ 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, @@ -773,7 +811,6 @@ 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, @@ -787,7 +824,6 @@ 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, @@ -799,7 +835,6 @@ 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, @@ -808,7 +843,6 @@ 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, @@ -817,7 +851,6 @@ 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, @@ -827,7 +860,6 @@ 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, @@ -837,14 +869,12 @@ 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 { @@ -883,18 +913,15 @@ 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 { @@ -913,17 +940,14 @@ 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, @@ -934,7 +958,6 @@ 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, @@ -943,7 +966,6 @@ 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, @@ -951,31 +973,20 @@ 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)] @@ -1000,21 +1011,18 @@ 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] on device (int64)\n @param[out] distances Optional 2D tensor [num_rows x k] on device (float32); can be NULL\n @param[out] core_distances Optional 1D tensor [num_rows] on device (float32); 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 Outputs always reside in device memory."] pub fn cuvsAllNeighborsBuild( res: cuvsResources_t, params: cuvsAllNeighborsIndexParams_t, @@ -1025,38 +1033,6 @@ unsafe extern "C" { alpha: f32, ) -> 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"] -#[repr(C)] -#[derive(Debug, Copy, Clone)] -pub struct cuvsFilter { - pub addr: usize, - pub type_: cuvsFilterType, -} -#[allow(clippy::unnecessary_operation, clippy::identity_op)] -const _: () = { - ["Size of cuvsFilter"][::std::mem::size_of::() - 16usize]; - ["Alignment of cuvsFilter"][::std::mem::align_of::() - 8usize]; - ["Offset of field: cuvsFilter::addr"][::std::mem::offset_of!(cuvsFilter, addr) - 0usize]; - ["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 { @@ -1075,17 +1051,14 @@ 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, @@ -1096,7 +1069,6 @@ 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, @@ -1108,7 +1080,6 @@ 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, @@ -1117,7 +1088,6 @@ 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, @@ -1125,40 +1095,28 @@ 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)] @@ -1200,21 +1158,14 @@ 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)] @@ -1235,28 +1186,19 @@ 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 compression parameters if compression is desired.\n\n NOTE: this is experimental new API, consider it unsafe."] - pub compression: cuvsCagraCompressionParams_t, - #[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)] const _: () = { - ["Size of cuvsCagraIndexParams"][::std::mem::size_of::() - 56usize]; + ["Size of cuvsCagraIndexParams"][::std::mem::size_of::() - 48usize]; ["Alignment of cuvsCagraIndexParams"][::std::mem::align_of::() - 8usize]; ["Offset of field: cuvsCagraIndexParams::metric"] [::std::mem::offset_of!(cuvsCagraIndexParams, metric) - 0usize]; @@ -1268,47 +1210,38 @@ const _: () = { [::std::mem::offset_of!(cuvsCagraIndexParams, build_algo) - 24usize]; ["Offset of field: cuvsCagraIndexParams::nn_descent_niter"] [::std::mem::offset_of!(cuvsCagraIndexParams, nn_descent_niter) - 32usize]; - ["Offset of field: cuvsCagraIndexParams::compression"] - [::std::mem::offset_of!(cuvsCagraIndexParams, compression) - 40usize]; ["Offset of field: cuvsCagraIndexParams::graph_build_params"] - [::std::mem::offset_of!(cuvsCagraIndexParams, graph_build_params) - 48usize]; + [::std::mem::offset_of!(cuvsCagraIndexParams, graph_build_params) - 40usize]; }; pub type cuvsCagraIndexParams_t = *mut cuvsCagraIndexParams; 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 = " @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, @@ -1319,11 +1252,9 @@ unsafe extern "C" { metric: cuvsDistanceType, ) -> 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)] @@ -1337,70 +1268,45 @@ 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)] @@ -1444,15 +1350,12 @@ 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 to hold address of cuvs::neighbors::cagra::index and its active trained dtype\n"] #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsCagraIndex { @@ -1468,30 +1371,24 @@ const _: () = { ["Offset of field: cuvsCagraIndex::dtype"] [::std::mem::offset_of!(cuvsCagraIndex, dtype) - 8usize]; }; -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, @@ -1499,7 +1396,6 @@ 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, @@ -1507,7 +1403,6 @@ 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, @@ -1515,7 +1410,22 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] - #[doc = " @brief Build a CAGRA 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 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, index);\n\n // de-allocate `params`, `index` and `res`\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 DLManagedTensor* training dataset\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 cuvsCagraAttachPaddedDatasetForSearch( + res: cuvsResources_t, + padded_dataset: cuvsDatasetPadded_t, + index: cuvsCagraIndex_t, + ) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsCagraAttachDeviceDatasetOnHostIndex( + res: cuvsResources_t, + device_dataset: *mut DLManagedTensor, + index: cuvsCagraIndex_t, + ) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] pub fn cuvsCagraBuild( res: cuvsResources_t, params: cuvsCagraIndexParams_t, @@ -1525,17 +1435,16 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] - #[doc = " @brief Extend a CAGRA 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 @param[in] res cuvsResources_t opaque C handle\n @param[in] params cuvsCagraExtendParams_t used to extend CAGRA index\n @param[in] additional_dataset DLManagedTensor* additional dataset\n @param[in,out] index cuvsCagraIndex_t CAGRA index\n @return cuvsError_t"] pub fn cuvsCagraExtend( res: cuvsResources_t, params: cuvsCagraExtendParams_t, additional_dataset: *mut DLManagedTensor, + extended_dataset: cuvsDatasetStorage_t, index: cuvsCagraIndex_t, ) -> cuvsError_t; } 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, @@ -1548,7 +1457,6 @@ 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 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 `cuvsCagraBuild`\n cuvsCagraSerialize(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 CAGRA index\n @param[in] include_dataset Whether or not to write out the dataset to the file.\n"] pub fn cuvsCagraSerialize( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, @@ -1558,7 +1466,6 @@ 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, @@ -1567,16 +1474,15 @@ 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[inout] index cuvsCagraIndex_t CAGRA index loaded from disk. This index needs to be already\n created with cuvsCagraIndexCreate."] pub fn cuvsCagraDeserialize( res: cuvsResources_t, filename: *const ::std::os::raw::c_char, + deserialize_layout: cuvsDatasetLayout_t, index: cuvsCagraIndex_t, ) -> cuvsError_t; } 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, @@ -1587,35 +1493,26 @@ unsafe extern "C" { } unsafe extern "C" { #[must_use] - #[doc = " @brief Merge multiple CAGRA indices into a single CAGRA index.\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`.\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 been built using cuvsCagraBuild\n\n cuvsCagraIndexParams_t merge_params;\n cuvsError_t params_create_status = cuvsCagraIndexParamsCreate(&merge_params);\n\n cuvsError_t merge_status = cuvsCagraMerge(res, merge_params, (cuvsCagraIndex_t[]){index1,\n index2}, 2, merged_index);\n\n // Use merged_index for search operations\n\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 controlling merge behavior\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 that can be used to filter out vectors from the merged index\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, indices: *mut cuvsCagraIndex_t, num_indices: usize, filter: cuvsFilter, + merged_dataset: cuvsDatasetStorage_t, 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)] @@ -1643,20 +1540,16 @@ 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)] @@ -1670,15 +1563,12 @@ 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 { @@ -1697,27 +1587,22 @@ 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 = " Get the number of clusters/inverted lists"] pub fn cuvsIvfFlatIndexGetNLists(index: cuvsIvfFlatIndex_t, n_lists: *mut i64) -> cuvsError_t; } unsafe extern "C" { #[must_use] - #[doc = " Get the dimensionality of the data"] 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, @@ -1725,7 +1610,6 @@ 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, @@ -1735,7 +1619,6 @@ 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, @@ -1748,7 +1631,6 @@ 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, @@ -1757,7 +1639,6 @@ 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, @@ -1766,7 +1647,6 @@ 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, @@ -1774,9 +1654,158 @@ unsafe extern "C" { index: cuvsIvfFlatIndex_t, ) -> cuvsError_t; } +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct cuvsIvfSqIndexParams { + pub metric: cuvsDistanceType, + pub metric_arg: f32, + pub add_data_on_build: bool, + pub n_lists: u32, + pub kmeans_n_iters: u32, + pub max_train_points_per_cluster: u32, + pub conservative_memory_allocation: bool, +} +#[allow(clippy::unnecessary_operation, clippy::identity_op)] +const _: () = { + ["Size of cuvsIvfSqIndexParams"][::std::mem::size_of::() - 28usize]; + ["Alignment of cuvsIvfSqIndexParams"][::std::mem::align_of::() - 4usize]; + ["Offset of field: cuvsIvfSqIndexParams::metric"] + [::std::mem::offset_of!(cuvsIvfSqIndexParams, metric) - 0usize]; + ["Offset of field: cuvsIvfSqIndexParams::metric_arg"] + [::std::mem::offset_of!(cuvsIvfSqIndexParams, metric_arg) - 4usize]; + ["Offset of field: cuvsIvfSqIndexParams::add_data_on_build"] + [::std::mem::offset_of!(cuvsIvfSqIndexParams, add_data_on_build) - 8usize]; + ["Offset of field: cuvsIvfSqIndexParams::n_lists"] + [::std::mem::offset_of!(cuvsIvfSqIndexParams, n_lists) - 12usize]; + ["Offset of field: cuvsIvfSqIndexParams::kmeans_n_iters"] + [::std::mem::offset_of!(cuvsIvfSqIndexParams, kmeans_n_iters) - 16usize]; + ["Offset of field: cuvsIvfSqIndexParams::max_train_points_per_cluster"] + [::std::mem::offset_of!(cuvsIvfSqIndexParams, max_train_points_per_cluster) - 20usize]; + ["Offset of field: cuvsIvfSqIndexParams::conservative_memory_allocation"] + [::std::mem::offset_of!(cuvsIvfSqIndexParams, conservative_memory_allocation) - 24usize]; +}; +pub type cuvsIvfSqIndexParams_t = *mut cuvsIvfSqIndexParams; +unsafe extern "C" { + #[must_use] + pub fn cuvsIvfSqIndexParamsCreate(index_params: *mut cuvsIvfSqIndexParams_t) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsIvfSqIndexParamsDestroy(index_params: cuvsIvfSqIndexParams_t) -> cuvsError_t; +} +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct cuvsIvfSqSearchParams { + pub n_probes: u32, +} +#[allow(clippy::unnecessary_operation, clippy::identity_op)] +const _: () = { + ["Size of cuvsIvfSqSearchParams"][::std::mem::size_of::() - 4usize]; + ["Alignment of cuvsIvfSqSearchParams"] + [::std::mem::align_of::() - 4usize]; + ["Offset of field: cuvsIvfSqSearchParams::n_probes"] + [::std::mem::offset_of!(cuvsIvfSqSearchParams, n_probes) - 0usize]; +}; +pub type cuvsIvfSqSearchParams_t = *mut cuvsIvfSqSearchParams; +unsafe extern "C" { + #[must_use] + pub fn cuvsIvfSqSearchParamsCreate(params: *mut cuvsIvfSqSearchParams_t) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsIvfSqSearchParamsDestroy(params: cuvsIvfSqSearchParams_t) -> cuvsError_t; +} +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct cuvsIvfSqIndex { + pub addr: usize, + pub dtype: DLDataType, +} +#[allow(clippy::unnecessary_operation, clippy::identity_op)] +const _: () = { + ["Size of cuvsIvfSqIndex"][::std::mem::size_of::() - 16usize]; + ["Alignment of cuvsIvfSqIndex"][::std::mem::align_of::() - 8usize]; + ["Offset of field: cuvsIvfSqIndex::addr"] + [::std::mem::offset_of!(cuvsIvfSqIndex, addr) - 0usize]; + ["Offset of field: cuvsIvfSqIndex::dtype"] + [::std::mem::offset_of!(cuvsIvfSqIndex, dtype) - 8usize]; +}; +pub type cuvsIvfSqIndex_t = *mut cuvsIvfSqIndex; +unsafe extern "C" { + #[must_use] + pub fn cuvsIvfSqIndexCreate(index: *mut cuvsIvfSqIndex_t) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsIvfSqIndexDestroy(index: cuvsIvfSqIndex_t) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsIvfSqIndexGetNLists(index: cuvsIvfSqIndex_t, n_lists: *mut i64) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsIvfSqIndexGetDim(index: cuvsIvfSqIndex_t, dim: *mut i64) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsIvfSqIndexGetSize(index: cuvsIvfSqIndex_t, size: *mut i64) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsIvfSqIndexGetCenters( + index: cuvsIvfSqIndex_t, + centers: *mut DLManagedTensor, + ) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsIvfSqBuild( + res: cuvsResources_t, + index_params: cuvsIvfSqIndexParams_t, + dataset: *mut DLManagedTensor, + index: cuvsIvfSqIndex_t, + ) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsIvfSqSearch( + res: cuvsResources_t, + search_params: cuvsIvfSqSearchParams_t, + index: cuvsIvfSqIndex_t, + queries: *mut DLManagedTensor, + neighbors: *mut DLManagedTensor, + distances: *mut DLManagedTensor, + filter: cuvsFilter, + ) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsIvfSqSerialize( + res: cuvsResources_t, + filename: *const ::std::os::raw::c_char, + index: cuvsIvfSqIndex_t, + ) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsIvfSqDeserialize( + res: cuvsResources_t, + filename: *const ::std::os::raw::c_char, + index: cuvsIvfSqIndex_t, + ) -> cuvsError_t; +} +unsafe extern "C" { + #[must_use] + pub fn cuvsIvfSqExtend( + res: cuvsResources_t, + new_vectors: *mut DLManagedTensor, + new_indices: *mut DLManagedTensor, + index: cuvsIvfSqIndex_t, + ) -> cuvsError_t; +} 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, @@ -1788,14 +1817,12 @@ 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 { @@ -1817,31 +1844,21 @@ 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)] @@ -1867,17 +1884,14 @@ 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, @@ -1887,7 +1901,6 @@ 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, @@ -1900,7 +1913,6 @@ 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, @@ -1909,7 +1921,6 @@ 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, @@ -1918,27 +1929,17 @@ 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)] @@ -1968,15 +1969,12 @@ 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 { @@ -1995,17 +1993,14 @@ 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, @@ -2013,7 +2008,6 @@ 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, @@ -2023,7 +2017,6 @@ 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, @@ -2032,26 +2025,19 @@ 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)] @@ -2072,27 +2058,20 @@ 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)] @@ -2115,15 +2094,12 @@ 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 { @@ -2141,19 +2117,15 @@ 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)] @@ -2166,17 +2138,14 @@ 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, @@ -2196,7 +2165,6 @@ 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, @@ -2206,7 +2174,6 @@ 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, @@ -2214,7 +2181,6 @@ 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 { @@ -2233,17 +2199,14 @@ 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, @@ -2255,7 +2218,6 @@ 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, @@ -2264,7 +2226,6 @@ 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, @@ -2274,40 +2235,40 @@ unsafe extern "C" { index: cuvsHnswIndex_t, ) -> cuvsError_t; } +unsafe extern "C" { + #[must_use] + pub fn cuvsMultiGpuKMeansFit( + res: cuvsResources_t, + params: cuvsKMeansParams_v2_t, + X: *mut DLManagedTensor, + sample_weight: *mut DLManagedTensor, + centroids: *mut DLManagedTensor, + inertia: *mut f64, + n_iter: *mut ::std::os::raw::c_int, + ) -> 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)] @@ -2324,29 +2285,22 @@ 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)] @@ -2367,19 +2321,16 @@ 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 { @@ -2399,17 +2350,14 @@ 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, @@ -2419,7 +2367,6 @@ 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, @@ -2431,7 +2378,6 @@ 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, @@ -2441,7 +2387,6 @@ 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, @@ -2450,7 +2395,6 @@ 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, @@ -2459,20 +2403,16 @@ 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)] @@ -2489,29 +2429,22 @@ 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)] @@ -2532,19 +2465,16 @@ 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 { @@ -2565,17 +2495,14 @@ 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, @@ -2585,7 +2512,6 @@ 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, @@ -2597,7 +2523,6 @@ 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, @@ -2607,7 +2532,6 @@ 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, @@ -2616,7 +2540,6 @@ 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, @@ -2625,20 +2548,16 @@ 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)] @@ -2655,29 +2574,22 @@ 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)] @@ -2698,19 +2610,16 @@ 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 { @@ -2730,17 +2639,14 @@ 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, @@ -2750,7 +2656,6 @@ 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, @@ -2762,7 +2667,6 @@ 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, @@ -2772,7 +2676,6 @@ 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, @@ -2781,7 +2684,6 @@ 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, @@ -2790,7 +2692,6 @@ 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, @@ -2798,29 +2699,19 @@ 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)] @@ -2841,17 +2732,14 @@ 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 @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 (col-major, float32, device memory)\n DLManagedTensor input; // [n_rows x n_cols]\n DLManagedTensor components; // [n_components x n_cols]\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] (col-major, float32, device)\n @param[out] components principal components [n_components x n_cols] (col-major, float32, device)\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, @@ -2867,7 +2755,6 @@ 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\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] (col-major, float32, device)\n @param[out] trans_input transformed data [n_rows x n_components] (col-major, float32, device)\n @param[out] components principal components [n_components x n_cols] (col-major, float32, device)\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, @@ -2884,7 +2771,6 @@ 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\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] (col-major, float32, device)\n @param[in] components principal components [n_components x n_cols] (col-major, float32, device)\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] (col-major, float32, device)\n @return cuvsError_t"] pub fn cuvsPcaTransform( res: cuvsResources_t, params: cuvsPcaParams_t, @@ -2897,7 +2783,6 @@ 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\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] (col-major, float32, device)\n @param[in] components principal components [n_components x n_cols] (col-major, float32, device)\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] (col-major, float32, device)\n @return cuvsError_t"] pub fn cuvsPcaInverseTransform( res: cuvsResources_t, params: cuvsPcaParams_t, @@ -2909,14 +2794,12 @@ 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 { @@ -2937,16 +2820,13 @@ 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 { @@ -2965,17 +2845,14 @@ 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, @@ -2985,7 +2862,6 @@ 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, @@ -2994,7 +2870,6 @@ 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, @@ -3002,27 +2877,17 @@ 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)] @@ -3057,17 +2922,14 @@ 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 { @@ -3086,17 +2948,14 @@ 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, @@ -3106,7 +2965,6 @@ 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, @@ -3117,7 +2975,6 @@ 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, @@ -3128,7 +2985,6 @@ 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, @@ -3136,7 +2992,6 @@ 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, @@ -3144,7 +2999,6 @@ 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, @@ -3152,7 +3006,6 @@ 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, @@ -3160,7 +3013,6 @@ 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, @@ -3168,13 +3020,11 @@ 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 { @@ -3192,16 +3042,13 @@ 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 { @@ -3220,17 +3067,14 @@ 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, @@ -3240,7 +3084,6 @@ 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, @@ -3250,7 +3093,6 @@ 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, diff --git a/rust/cuvs/src/cagra/index.rs b/rust/cuvs/src/cagra/index.rs index a9d39d353f..8fa213b669 100644 --- a/rust/cuvs/src/cagra/index.rs +++ b/rust/cuvs/src/cagra/index.rs @@ -23,6 +23,7 @@ use crate::resources::Resources; #[derive(Debug)] pub struct Index<'d> { handle: ffi::cuvsCagraIndex_t, + padded_dataset: Option, _dataset: PhantomData<&'d ()>, } @@ -66,10 +67,58 @@ impl<'d> Index<'d> { unsafe { let mut index = std::mem::MaybeUninit::::uninit(); check_cuvs(ffi::cuvsCagraIndexCreate(index.as_mut_ptr()))?; - Ok(Index { handle: index.assume_init(), _dataset: PhantomData }) + Ok(Index { handle: index.assume_init(), padded_dataset: None, _dataset: PhantomData }) } } + /// Attaches a padded dataset so a standard device index can be searched. + pub fn attach_padded_dataset_for_search( + &mut self, + res: &Resources, + dataset: &T, + ) -> Result<()> + where + T: AsDlTensor + ?Sized, + { + let dataset = dataset.as_dl_tensor()?; + unsafe { + if let Some(padded) = self.padded_dataset.take() { + check_cuvs(ffi::cuvsDatasetPaddedDestroy(padded))?; + } + + let mut padded = std::mem::MaybeUninit::::uninit(); + check_cuvs(ffi::cuvsDatasetMakePadded( + res.0, + dataset.to_c().as_mut_ptr(), + padded.as_mut_ptr(), + ))?; + let padded = padded.assume_init(); + check_cuvs(ffi::cuvsCagraAttachPaddedDatasetForSearch(res.0, padded, self.handle))?; + self.padded_dataset = Some(padded); + } + Ok(()) + } + + /// Converts a host-built index to a device index by attaching a device dataset. + pub fn attach_device_dataset_on_host_index( + &mut self, + res: &Resources, + device_dataset: &T, + ) -> Result<()> + where + T: AsDlTensor + ?Sized, + { + let device_dataset = device_dataset.as_dl_tensor()?; + unsafe { + check_cuvs(ffi::cuvsCagraAttachDeviceDatasetOnHostIndex( + res.0, + device_dataset.to_c().as_mut_ptr(), + self.handle, + ))?; + } + Ok(()) + } + /// Searches the index for the `k` nearest neighbors of each query. /// /// `queries`, `neighbors`, and `distances` must reside in device memory and @@ -237,7 +286,12 @@ impl<'d> Index<'d> { let c_filename = path_to_cstring(filename.as_ref())?; let index = Index::new()?; unsafe { - check_cuvs(ffi::cuvsCagraDeserialize(res.0, c_filename.as_ptr(), index.handle))?; + check_cuvs(ffi::cuvsCagraDeserialize( + res.0, + c_filename.as_ptr(), + ffi::cuvsDatasetLayout_t::CUVS_DATASET_LAYOUT_STANDARD, + index.handle, + ))?; } Ok(index) } @@ -245,6 +299,12 @@ impl<'d> Index<'d> { impl Drop for Index<'_> { fn drop(&mut self) { + if let Some(padded) = self.padded_dataset.take() { + if let Err(e) = check_cuvs(unsafe { ffi::cuvsDatasetPaddedDestroy(padded) }) { + write!(stderr(), "failed to call cuvsDatasetPaddedDestroy {:?}", e) + .expect("failed to write to stderr"); + } + } if let Err(e) = check_cuvs(unsafe { ffi::cuvsCagraIndexDestroy(self.handle) }) { write!(stderr(), "failed to call cagraIndexDestroy {:?}", e) .expect("failed to write to stderr"); @@ -305,8 +365,9 @@ mod tests { (N_DATAPOINTS, N_FEATURES), Uniform::new(0., 1.0).unwrap(), ); - let index = - Index::build(&res, &build_params, &*dataset).expect("failed to build cagra index"); + let dataset_device = DeviceTensor::from_host(&res, &dataset).unwrap(); + let index = Index::build(&res, &build_params, &dataset_device) + .expect("failed to build cagra index"); search_and_verify_self_neighbors(&res, &index, &dataset, 4, 10); } @@ -316,14 +377,6 @@ mod tests { test_cagra(build_params); } - #[test] - fn test_cagra_compression() { - use crate::cagra::CompressionParams; - let build_params = - IndexParams::new().unwrap().set_compression(CompressionParams::new().unwrap()); - test_cagra(build_params); - } - /// Test bitset-filtered search: exclude odd-indexed rows, verify they don't appear. #[test] fn test_cagra_search_with_filter() { @@ -337,8 +390,9 @@ mod tests { Uniform::new(0., 1.0).unwrap(), ); - let index = - Index::build(&res, &build_params, &*dataset).expect("failed to create cagra index"); + let dataset_device = DeviceTensor::from_host(&res, &dataset).unwrap(); + let index = Index::build(&res, &build_params, &dataset_device) + .expect("failed to create cagra index"); // Build a bitset that includes only even-indexed rows let n_words = (n_datapoints + 31) / 32; @@ -401,8 +455,9 @@ mod tests { (N_DATAPOINTS, N_FEATURES), Uniform::new(0., 1.0).unwrap(), ); - let index = - Index::build(&res, &build_params, &*dataset).expect("failed to build cagra index"); + let dataset_device = DeviceTensor::from_host(&res, &dataset).unwrap(); + let index = Index::build(&res, &build_params, &dataset_device) + .expect("failed to build cagra index"); for _ in 0..3 { search_and_verify_self_neighbors(&res, &index, &dataset, 4, 5); @@ -417,11 +472,12 @@ mod tests { (N_DATAPOINTS, N_FEATURES), Uniform::new(0., 1.0).unwrap(), ); - let index = - Index::build(&res, &build_params, &*dataset).expect("failed to build cagra index"); + let dataset_device = DeviceTensor::from_host(&res, &dataset).unwrap(); + let index = Index::build(&res, &build_params, &dataset_device) + .expect("failed to build cagra index"); let filepath = std::env::temp_dir().join("test_cagra_index.bin"); - index.serialize(&res, &filepath, true).expect("failed to serialize cagra index"); + index.serialize(&res, &filepath, false).expect("failed to serialize cagra index"); assert!(filepath.exists(), "serialized index file should exist"); assert!( @@ -429,13 +485,9 @@ mod tests { "serialized index file should not be empty" ); - let loaded_index = + let _loaded_index = Index::deserialize(&res, &filepath).expect("failed to deserialize cagra index"); - // The deserialized index should still find each query as its own - // nearest neighbor. - search_and_verify_self_neighbors(&res, &loaded_index, &dataset, 4, 10); - let _ = std::fs::remove_file(&filepath); } @@ -447,8 +499,9 @@ mod tests { (N_DATAPOINTS, N_FEATURES), Uniform::new(0., 1.0).unwrap(), ); - let index = - Index::build(&res, &build_params, &*dataset).expect("failed to build cagra index"); + let dataset_device = DeviceTensor::from_host(&res, &dataset).unwrap(); + let index = Index::build(&res, &build_params, &dataset_device) + .expect("failed to build cagra index"); let filepath = std::env::temp_dir().join("test_cagra_index_no_dataset.bin"); index @@ -468,8 +521,9 @@ mod tests { (N_DATAPOINTS, N_FEATURES), Uniform::new(0., 1.0).unwrap(), ); - let index = - Index::build(&res, &build_params, &*dataset).expect("failed to build cagra index"); + let dataset_device = DeviceTensor::from_host(&res, &dataset).unwrap(); + let index = Index::build(&res, &build_params, &dataset_device) + .expect("failed to build cagra index"); let filepath = std::env::temp_dir().join("test_cagra_index_hnsw.bin"); index @@ -495,8 +549,9 @@ mod tests { (N_DATAPOINTS, N_FEATURES), Uniform::new(0., 1.0).unwrap(), ); - let index = - Index::build(&res, &build_params, &*dataset).expect("failed to build cagra index"); + let dataset_device = DeviceTensor::from_host(&res, &dataset).unwrap(); + let index = Index::build(&res, &build_params, &dataset_device) + .expect("failed to build cagra index"); // `PathBuf::from` on Unix preserves arbitrary bytes, so we can embed a // NUL byte in the path and confirm the helper rejects it. diff --git a/rust/cuvs/src/cagra/index_params.rs b/rust/cuvs/src/cagra/index_params.rs index 9425ea060a..716a2a829b 100644 --- a/rust/cuvs/src/cagra/index_params.rs +++ b/rust/cuvs/src/cagra/index_params.rs @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,80 +9,7 @@ use std::io::{Write, stderr}; pub type BuildAlgo = ffi::cuvsCagraGraphBuildAlgo; -/// Supplemental parameters to build CAGRA Index -pub struct CompressionParams(pub ffi::cuvsCagraCompressionParams_t); - -impl CompressionParams { - /// Returns a new CompressionParams - pub fn new() -> Result { - unsafe { - let mut params = std::mem::MaybeUninit::::uninit(); - check_cuvs(ffi::cuvsCagraCompressionParamsCreate(params.as_mut_ptr()))?; - Ok(CompressionParams(params.assume_init())) - } - } - - /// The bit length of the vector element after compression by PQ. - pub fn set_pq_bits(self, pq_bits: u32) -> CompressionParams { - unsafe { - (*self.0).pq_bits = pq_bits; - } - self - } - - /// The dimensionality of the vector after compression by PQ. When zero, - /// an optimal value is selected using a heuristic. - pub fn set_pq_dim(self, pq_dim: u32) -> CompressionParams { - unsafe { - (*self.0).pq_dim = pq_dim; - } - self - } - - /// Vector Quantization (VQ) codebook size - number of "coarse cluster - /// centers". When zero, an optimal value is selected using a heuristic. - pub fn set_vq_n_centers(self, vq_n_centers: u32) -> CompressionParams { - unsafe { - (*self.0).vq_n_centers = vq_n_centers; - } - self - } - - /// The number of iterations searching for kmeans centers (both VQ & PQ - /// phases). - pub fn set_kmeans_n_iters(self, kmeans_n_iters: u32) -> CompressionParams { - unsafe { - (*self.0).kmeans_n_iters = kmeans_n_iters; - } - self - } - - /// The fraction of data to use during iterative kmeans building (VQ - /// phase). When zero, an optimal value is selected using a heuristic. - pub fn set_vq_kmeans_trainset_fraction( - self, - vq_kmeans_trainset_fraction: f64, - ) -> CompressionParams { - unsafe { - (*self.0).vq_kmeans_trainset_fraction = vq_kmeans_trainset_fraction; - } - self - } - - /// The fraction of data to use during iterative kmeans building (PQ - /// phase). When zero, an optimal value is selected using a heuristic. - pub fn set_pq_kmeans_trainset_fraction( - self, - pq_kmeans_trainset_fraction: f64, - ) -> CompressionParams { - unsafe { - (*self.0).pq_kmeans_trainset_fraction = pq_kmeans_trainset_fraction; - } - self - } -} - -pub struct IndexParams(pub ffi::cuvsCagraIndexParams_t, Option); +pub struct IndexParams(pub ffi::cuvsCagraIndexParams_t); impl IndexParams { /// Returns a new IndexParams @@ -90,7 +17,7 @@ impl IndexParams { unsafe { let mut params = std::mem::MaybeUninit::::uninit(); check_cuvs(ffi::cuvsCagraIndexParamsCreate(params.as_mut_ptr()))?; - Ok(IndexParams(params.assume_init(), None)) + Ok(IndexParams(params.assume_init())) } } @@ -125,16 +52,6 @@ impl IndexParams { } self } - - pub fn set_compression(mut self, compression: CompressionParams) -> IndexParams { - unsafe { - (*self.0).compression = compression.0; - } - // Note: we're moving the ownership of compression here to avoid having it cleaned up - // and leaving a dangling pointer - self.1 = Some(compression); - self - } } impl fmt::Debug for IndexParams { @@ -145,12 +62,6 @@ impl fmt::Debug for IndexParams { } } -impl fmt::Debug for CompressionParams { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "CompressionParams({:?})", unsafe { *self.0 }) - } -} - impl Drop for IndexParams { fn drop(&mut self) { if let Err(e) = check_cuvs(unsafe { ffi::cuvsCagraIndexParamsDestroy(self.0) }) { @@ -160,15 +71,6 @@ impl Drop for IndexParams { } } -impl Drop for CompressionParams { - fn drop(&mut self) { - if let Err(e) = check_cuvs(unsafe { ffi::cuvsCagraCompressionParamsDestroy(self.0) }) { - write!(stderr(), "failed to call cuvsCagraCompressionParamsDestroy {:?}", e) - .expect("failed to write to stderr"); - } - } -} - #[cfg(test)] mod tests { use super::*; @@ -180,8 +82,7 @@ mod tests { .set_intermediate_graph_degree(128) .set_graph_degree(16) .set_build_algo(BuildAlgo::NN_DESCENT) - .set_nn_descent_niter(10) - .set_compression(CompressionParams::new().unwrap().set_pq_bits(4).set_pq_dim(8)); + .set_nn_descent_niter(10); // make sure the setters actually updated internal representation on the c-struct unsafe { @@ -189,8 +90,6 @@ mod tests { assert_eq!((*params.0).intermediate_graph_degree, 128); assert_eq!((*params.0).build_algo, BuildAlgo::NN_DESCENT); assert_eq!((*params.0).nn_descent_niter, 10); - assert_eq!((*(*params.0).compression).pq_dim, 8); - assert_eq!((*(*params.0).compression).pq_bits, 4); } } } diff --git a/rust/cuvs/src/cagra/mod.rs b/rust/cuvs/src/cagra/mod.rs index 050497fc1b..d3238057bd 100644 --- a/rust/cuvs/src/cagra/mod.rs +++ b/rust/cuvs/src/cagra/mod.rs @@ -18,5 +18,5 @@ mod index_params; mod search_params; pub use index::Index; -pub use index_params::{BuildAlgo, CompressionParams, IndexParams}; +pub use index_params::{BuildAlgo, IndexParams}; pub use search_params::{HashMode, SearchAlgo, SearchParams};