-
Notifications
You must be signed in to change notification settings - Fork 217
Recover cagra-q path in C API and downstream language wrappers #2413
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
309c59b
6418b95
f2ed015
fb5cb26
645519c
114b72a
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -255,6 +255,24 @@ CUVS_EXPORT cuvsError_t cuvsCagraCompressionParamsCreate(cuvsCagraCompressionPar | |
| */ | ||
| CUVS_EXPORT cuvsError_t cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionParams_t params); | ||
|
|
||
| /** | ||
| * @brief Train an owning device VPQ (f16 codebook) dataset from a device-padded source. | ||
| * | ||
| * Used for CAGRA-Q: build a dense CAGRA index, train VPQ with this factory, then attach via | ||
| * `cuvsCagraUpdateDataset`. Caller owns the returned dataset and must keep it alive while any | ||
| * index uses it. Metric for subsequent search must remain `L2Expanded`. | ||
| * | ||
| * @param[in] res cuvs resources | ||
| * @param[in] source_dataset device-padded dataset (owning or view) | ||
| * @param[in] params VPQ compression params; NULL selects defaults | ||
| * @param[out] vpq_dataset newly allocated owning VPQ dataset handle | ||
| * @return cuvsError_t | ||
| */ | ||
| CUVS_EXPORT cuvsError_t cuvsDatasetMakeVpq(cuvsResources_t res, | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Let's drop the V from this and call it
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Just wanted to clarify my understanding of make_vpq_dataset() factory. There is a PQ step and a VQ step but the VQ step is optional and run only when the flag bool use_vq = true. Is that why we are renaming the make_vpq_dataset() factory to make_pq_dataset()? So we will have pq dataset, rabitq dataset, bbq dataset, and sq dataset but no standalone vpq dataset? We will treat vpq as a special condition of pq dataset?
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The VQ is not a proper VQ. It's misnamed. It's coarse centroid assignments. But we don't need to include that in the name. |
||
| cuvsDataset_t source_dataset, | ||
| cuvsCagraCompressionParams_t params, | ||
| cuvsDataset_t* vpq_dataset); | ||
|
|
||
| /** | ||
| * @brief Allocate ACE params, and populate with default values | ||
| * | ||
|
|
@@ -580,21 +598,25 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexGetDataset(cuvsCagraIndex_t index, DLManag | |
| CUVS_EXPORT cuvsError_t cuvsCagraIndexGetGraph(cuvsCagraIndex_t index, DLManagedTensor* graph); | ||
|
|
||
| /** | ||
| * @brief Update a CAGRA index with a device-padded dataset. | ||
| * @brief Update a CAGRA index with a device dataset (padded or VPQ). | ||
| * | ||
| * This is the centralized dataset update/attach operation for C callers. | ||
| * | ||
| * - Device-padded dataset: if \p index is already device-padded, its dataset view is replaced in | ||
| * place (same index object); otherwise the index is converted via attach and rebound. | ||
| * - Device VPQ_F16 dataset (from `cuvsDatasetMakeVpq`): if \p index is already VPQ-typed, its | ||
| * dataset view is replaced in place; otherwise the graph is copied into a new VPQ-typed index | ||
| * (CAGRA-Q). Search requires metric `L2Expanded`. The VPQ handle must be owning. | ||
| * | ||
| * This is the centralized dataset update operation for C callers. If \p index | ||
| * is already device-padded, its dataset view is replaced in place. Otherwise, | ||
| * the index is converted and its opaque handle is rebound to a search-ready | ||
| * device-padded index. Caller retains ownership of | ||
| * \p device_padded_dataset and must keep it alive while \p index uses it. | ||
| * Caller retains ownership of \p dataset and must keep it alive while \p index uses it. | ||
| * | ||
| * @param[in] res cuvsResources_t opaque C handle | ||
| * @param[in] device_padded_dataset owning or non-owning device-padded dataset handle | ||
| * @param[inout] index CAGRA index handle | ||
| * @param[in] res cuvsResources_t opaque C handle | ||
| * @param[in] dataset device-padded or owning device VPQ_F16 dataset handle | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Hmm. I don't know that i like this naming. VPQ_F16... that's cryptic... and why only F16? Is that really the only data type that's supported? Let's keep the naming contention here:
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. VPQ_F32 is not supported right now in the cagra-q search kernels. Source vectors can be f32/f16/int8/uint8 but the codebook math type stored for search is forced to half. compute_distance_vpq-impl.cuh has this line: |
||
| * @param[inout] index CAGRA index handle | ||
| * @return cuvsError_t | ||
| */ | ||
| CUVS_EXPORT cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res, | ||
| cuvsDataset_t device_padded_dataset, | ||
| cuvsDataset_t dataset, | ||
| cuvsCagraIndex_t index); | ||
|
|
||
| /** | ||
|
|
||
Large diffs are not rendered by default.
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -19,6 +19,7 @@ | |
| #include <string> | ||
| #include <type_traits> | ||
| #include <unistd.h> | ||
| #include <vector> | ||
|
|
||
| #include <cuda_runtime.h> | ||
| #include <gtest/gtest.h> | ||
|
|
@@ -2008,3 +2009,109 @@ TEST(CagraC, SearchMultiPartitionMultiKernelRejected) | |
| } | ||
| cuvsResourcesDestroy(res); | ||
| } | ||
|
|
||
| TEST(CagraC, BuildAttachVpqSearch) | ||
| { | ||
| // CAGRA-Q smoke test: dense build → MakeVpq → UpdateDataset(VPQ) → Search. | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thanks for adding the tests to C api and language wrappers. I don't see tests in cpp though. We should add that. |
||
| constexpr int64_t n_rows = 256; | ||
| constexpr int64_t dim = 32; | ||
| constexpr int64_t n_queries = 4; | ||
| constexpr int64_t k = 1; | ||
|
|
||
| cuvsResources_t res; | ||
| ASSERT_EQ(cuvsResourcesCreate(&res), CUVS_SUCCESS); | ||
| cudaStream_t stream; | ||
| ASSERT_EQ(cuvsStreamGet(res, &stream), CUVS_SUCCESS); | ||
|
|
||
| rmm::device_uvector<float> dataset_d(n_rows * dim, stream); | ||
| { | ||
| std::vector<float> host(n_rows * dim); | ||
| for (int64_t i = 0; i < n_rows * dim; ++i) { | ||
| host[i] = static_cast<float>((i % 17) + 1); | ||
| } | ||
| raft::copy(dataset_d.data(), host.data(), host.size(), stream); | ||
| } | ||
|
|
||
| // dim=32 float already matches CAGRA padded row width; MakePadded refuses a | ||
| // no-op device copy — wrap with MakePaddedView instead. | ||
| DLManagedTensor dataset_tensor{}; | ||
| dataset_tensor.dl_tensor.data = dataset_d.data(); | ||
| dataset_tensor.dl_tensor.device.device_type = kDLCUDA; | ||
| dataset_tensor.dl_tensor.ndim = 2; | ||
| dataset_tensor.dl_tensor.dtype = {kDLFloat, 32, 1}; | ||
| int64_t dataset_shape[2] = {n_rows, dim}; | ||
| dataset_tensor.dl_tensor.shape = dataset_shape; | ||
| dataset_tensor.dl_tensor.strides = nullptr; | ||
|
|
||
| cuvsDataset_t padded; | ||
| ASSERT_EQ(cuvsDatasetMakePaddedView(res, &dataset_tensor, &padded), CUVS_SUCCESS); | ||
|
|
||
| cuvsCagraIndexParams_t build_params; | ||
| ASSERT_EQ(cuvsCagraIndexParamsCreate(&build_params), CUVS_SUCCESS); | ||
| cuvsCagraIndex_t index; | ||
| ASSERT_EQ(cuvsCagraIndexCreate(&index), CUVS_SUCCESS); | ||
| ASSERT_EQ(cuvsCagraBuild(res, build_params, padded, index), CUVS_SUCCESS); | ||
|
|
||
| cuvsCagraCompressionParams_t compression; | ||
| ASSERT_EQ(cuvsCagraCompressionParamsCreate(&compression), CUVS_SUCCESS); | ||
| compression->pq_bits = 8; | ||
| compression->pq_dim = 8; | ||
|
|
||
| cuvsDataset_t vpq = nullptr; | ||
| ASSERT_EQ(cuvsDatasetMakeVpq(res, padded, compression, &vpq), CUVS_SUCCESS); | ||
| { | ||
| cuvsDatasetLayout_t layout; | ||
| ASSERT_EQ(cuvsDatasetGetLayout(vpq, &layout), CUVS_SUCCESS); | ||
| EXPECT_EQ(layout, CUVS_DATASET_LAYOUT_VPQ_F16); | ||
| bool owning = false; | ||
| ASSERT_EQ(cuvsDatasetGetIsOwning(vpq, &owning), CUVS_SUCCESS); | ||
| EXPECT_TRUE(owning); | ||
| } | ||
|
|
||
| ASSERT_EQ(cuvsCagraUpdateDataset(res, vpq, index), CUVS_SUCCESS); | ||
|
|
||
| rmm::device_uvector<float> queries_d(n_queries * dim, stream); | ||
| raft::copy(queries_d.data(), dataset_d.data(), n_queries * dim, stream); | ||
| DLManagedTensor queries_tensor{}; | ||
| queries_tensor.dl_tensor.data = queries_d.data(); | ||
| queries_tensor.dl_tensor.device.device_type = kDLCUDA; | ||
| queries_tensor.dl_tensor.ndim = 2; | ||
| queries_tensor.dl_tensor.dtype = {kDLFloat, 32, 1}; | ||
| int64_t queries_shape[2] = {n_queries, dim}; | ||
| queries_tensor.dl_tensor.shape = queries_shape; | ||
|
|
||
| rmm::device_uvector<uint32_t> neighbors_d(n_queries * k, stream); | ||
| DLManagedTensor neighbors_tensor{}; | ||
| neighbors_tensor.dl_tensor.data = neighbors_d.data(); | ||
| neighbors_tensor.dl_tensor.device.device_type = kDLCUDA; | ||
| neighbors_tensor.dl_tensor.ndim = 2; | ||
| neighbors_tensor.dl_tensor.dtype = {kDLUInt, 32, 1}; | ||
| int64_t neighbors_shape[2] = {n_queries, k}; | ||
| neighbors_tensor.dl_tensor.shape = neighbors_shape; | ||
|
|
||
| rmm::device_uvector<float> distances_d(n_queries * k, stream); | ||
| DLManagedTensor distances_tensor{}; | ||
| distances_tensor.dl_tensor.data = distances_d.data(); | ||
| distances_tensor.dl_tensor.device.device_type = kDLCUDA; | ||
| distances_tensor.dl_tensor.ndim = 2; | ||
| distances_tensor.dl_tensor.dtype = {kDLFloat, 32, 1}; | ||
| int64_t distances_shape[2] = {n_queries, k}; | ||
| distances_tensor.dl_tensor.shape = distances_shape; | ||
|
|
||
| cuvsFilter filter; | ||
| filter.type = NO_FILTER; | ||
| filter.addr = (uintptr_t)NULL; | ||
| cuvsCagraSearchParams_t search_params; | ||
| ASSERT_EQ(cuvsCagraSearchParamsCreate(&search_params), CUVS_SUCCESS); | ||
| ASSERT_EQ(cuvsCagraSearch( | ||
| res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter), | ||
| CUVS_SUCCESS); | ||
|
|
||
| cuvsCagraSearchParamsDestroy(search_params); | ||
| cuvsCagraCompressionParamsDestroy(compression); | ||
| cuvsDatasetDestroy(vpq); | ||
| cuvsCagraIndexDestroy(index); | ||
| cuvsCagraIndexParamsDestroy(build_params); | ||
| cuvsDatasetDestroy(padded); | ||
| cuvsResourcesDestroy(res); | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -4548,6 +4548,42 @@ auto convert_standard_to_padded_index( | |
| return out; | ||
| } | ||
|
|
||
| /** | ||
| * @brief Convert a dense (non-VPQ) CAGRA index into a VPQ f16 index and attach a VPQ dataset. | ||
| * | ||
| * Copies graph/source-indices from `src` into a new `vpq_f16_index` and attaches `vpq_dataset`. | ||
| * Caller owns `vpq_dataset` storage for the lifetime of the returned index. | ||
| */ | ||
| template <typename T, typename IdxT, typename IndexViewT> | ||
| requires cuvs::neighbors::ann_dataset_view<IndexViewT> | ||
| auto convert_dense_to_vpq_f16_index( | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This function seems unnecessary. We shouldn't be copying the cagra graph just to change the template type. |
||
| raft::resources const& res, | ||
| index<T, IdxT, IndexViewT> const& src, | ||
| cuvs::neighbors::device_vpq_dataset_view<half, int64_t> const& vpq_dataset) | ||
| -> vpq_f16_index<T, IdxT> | ||
| { | ||
| RAFT_EXPECTS(vpq_dataset.n_rows() == src.size(), | ||
| "VPQ dataset row count must match the index size"); | ||
|
|
||
| vpq_f16_index<T, IdxT> out(res, src.metric()); | ||
| if (src.graph().extent(0) > 0) { | ||
| using GraphIndexType = typename index<T, IdxT, IndexViewT>::graph_index_type; | ||
| auto graph_host = | ||
| raft::make_host_matrix<GraphIndexType, int64_t>(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())); | ||
| } | ||
| if (src.source_indices().has_value()) { | ||
| out.update_source_indices(res, src.source_indices().value()); | ||
| } | ||
| out.update_device_dataset_same_layout(res, vpq_dataset); | ||
| return out; | ||
| } | ||
|
|
||
| /** | ||
| * @brief Attach a device-padded dataset and return a search-ready padded-device index. | ||
| * | ||
|
|
@@ -4598,6 +4634,35 @@ auto attach_dataset( | |
| } | ||
| } | ||
|
|
||
| /** | ||
| * @brief Attach a device VPQ (f16 codebook) dataset and return a search-ready `vpq_f16_index`. | ||
| * | ||
| * Builds a new VPQ-typed index by copying the graph from `idx`. Caller owns `vpq_dataset` storage. | ||
| * | ||
| * If `idx` is already a `vpq_f16_index`, call `idx.update_device_dataset_same_layout(res, | ||
| * vpq_dataset)` directly. | ||
| */ | ||
| template <typename T, typename IdxT, typename IndexViewT> | ||
| requires cuvs::neighbors::ann_dataset_view<IndexViewT> | ||
| auto attach_dataset(raft::resources const& res, | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Please stop naming this attach_dataset. The only function that should be exposed to the outside world is |
||
| index<T, IdxT, IndexViewT> const& idx, | ||
| cuvs::neighbors::device_vpq_dataset_view<half, int64_t> const& vpq_dataset) | ||
| -> vpq_f16_index<T, IdxT> | ||
| { | ||
| if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v<IndexViewT>) { | ||
| RAFT_LOG_WARN( | ||
| "cagra::attach_dataset called with an already vpq_f16 index. " | ||
| "To avoid an unnecessary index copy, call " | ||
| "index.update_device_dataset_same_layout(res, vpq_dataset) " | ||
| "directly on the original index."); | ||
| RAFT_FAIL( | ||
| "cagra::attach_dataset: vpq_f16_index input is not supported in this overload. " | ||
| "Call index.update_device_dataset_same_layout(res, vpq_dataset) directly."); | ||
| } else { | ||
| return convert_dense_to_vpq_f16_index(res, idx, vpq_dataset); | ||
| } | ||
| } | ||
|
|
||
| } // namespace cagra | ||
| } // namespace neighbors | ||
| } // namespace CUVS_EXPORT cuvs | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -246,14 +246,17 @@ void inverse_transform( | |
|
|
||
| namespace detail { | ||
|
|
||
| // Must be CUVS_EXPORT: libcuvs_c (and header-inline make_device_vpq_dataset) resolve this | ||
| // across the shared-library boundary. Without default visibility, -fvisibility=hidden | ||
| // + --gc-sections drop the pq.cu instantiations from libcuvs.so. | ||
| template <typename T> | ||
| [[nodiscard]] cuvs::neighbors::device_vpq_dataset<half, int64_t> 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); | ||
| [[nodiscard]] CUVS_EXPORT cuvs::neighbors::device_vpq_dataset<half, int64_t> | ||
| 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 | ||
|
|
||
|
|
@@ -267,32 +270,33 @@ template <typename T> | |
| * | ||
| * 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_device_dataset_same_layout` with a non-owning view. | ||
| * used for graph build, keep the `device_vpq_dataset` alive, and attach it with | ||
| * `cagra::attach_dataset` (returns a `vpq_f16_index`). | ||
| * | ||
| * @code{.cpp} | ||
| * #include <cuvs/neighbors/cagra.hpp> | ||
| * #include <cuvs/preprocessing/quantize/pq.hpp> | ||
| * | ||
| * // `idx` is a `cagra::index<float, uint32_t>` with graph built on dense rows. | ||
| * // `idx` is a dense CAGRA index with graph built on padded rows. | ||
| * // `padded` is a `device_padded_dataset_view<float, int64_t>` 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_device_dataset_same_layout(res, vpq.as_dataset_view()); | ||
| * auto vpq = cuvs::preprocessing::quantize::pq::make_device_vpq_dataset(res, vpq_params, | ||
| * padded.view()); auto vpq_idx = cuvs::neighbors::cagra::attach_dataset(res, idx, | ||
| * vpq.as_dataset_view()); | ||
| * @endcode | ||
| */ | ||
| template <typename SrcT> | ||
| [[nodiscard]] auto make_vpq_dataset(raft::resources const& res, | ||
| cuvs::neighbors::vpq_params const& params, | ||
| SrcT const& src) | ||
| [[nodiscard]] auto make_device_vpq_dataset(raft::resources const& res, | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Please rename- |
||
| cuvs::neighbors::vpq_params const& params, | ||
| SrcT const& src) | ||
| -> cuvs::neighbors::device_vpq_dataset<half, int64_t> | ||
| { | ||
| using T = typename SrcT::value_type; | ||
| RAFT_EXPECTS(src.extent(0) > 0, "make_vpq_dataset: dataset is empty"); | ||
| RAFT_EXPECTS(src.extent(0) > 0, "make_device_vpq_dataset: dataset is empty"); | ||
| cudaPointerAttributes ptr_attrs; | ||
| RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src.data_handle())); | ||
| auto const* device_ptr = reinterpret_cast<T const*>(ptr_attrs.devicePointer); | ||
| RAFT_EXPECTS(device_ptr != nullptr, "make_vpq_dataset: source must be device-accessible."); | ||
| RAFT_EXPECTS(device_ptr != nullptr, "make_device_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; | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.