Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion c/include/cuvs/core/dataset.h
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@ extern "C" {
*/
typedef enum {
CUVS_DATASET_LAYOUT_STANDARD = 0,
CUVS_DATASET_LAYOUT_PADDED = 1
CUVS_DATASET_LAYOUT_PADDED = 1,
/** Device VPQ storage with f16 codebooks (CAGRA-Q search dataset). */
CUVS_DATASET_LAYOUT_VPQ_F16 = 2

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
CUVS_DATASET_LAYOUT_VPQ_F16 = 2
CUVS_DATASET_LAYOUT_PQ = 2

} cuvsDatasetLayout_t;

/**
Expand Down
42 changes: 32 additions & 10 deletions c/include/cuvs/neighbors/cagra.h
Original file line number Diff line number Diff line change
Expand Up @@ -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,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's drop the V from this and call it cuvsDatasetMakePq.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The 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?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The 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
*
Expand Down Expand Up @@ -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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The 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: device-padded or pq`. Also- there's no reason we should need to specify "owning device" here... that should be opaque to the user (they just create a dataset in the c layer and they pass it in and we worry about the view creation under the hood).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The 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:
static_assert(std::is_same_v<CODE_BOOK_T, half>, "Only CODE_BOOK_T = half is supported now");

* @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);

/**
Expand Down
397 changes: 296 additions & 101 deletions c/src/neighbors/cagra.cpp

Large diffs are not rendered by default.

107 changes: 107 additions & 0 deletions c/tests/neighbors/ann_cagra_c.cu
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
#include <string>
#include <type_traits>
#include <unistd.h>
#include <vector>

#include <cuda_runtime.h>
#include <gtest/gtest.h>
Expand Down Expand Up @@ -2008,3 +2009,109 @@ TEST(CagraC, SearchMultiPartitionMultiKernelRejected)
}
cuvsResourcesDestroy(res);
}

TEST(CagraC, BuildAttachVpqSearch)
{
// CAGRA-Q smoke test: dense build → MakeVpq → UpdateDataset(VPQ) → Search.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The 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);
}
65 changes: 65 additions & 0 deletions cpp/include/cuvs/neighbors/cagra.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The 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.
I am making the change to update_dataset to do the move instead of copy. We should be able to reuse update_dataset for this.

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.
*
Expand Down Expand Up @@ -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,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The 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 update_dataset. Since we're not rushing to get this into 26.08, we need to fix this.

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
Expand Down
38 changes: 21 additions & 17 deletions cpp/include/cuvs/preprocessing/quantize/pq.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please rename- pq_dataset. We don't need the vpq in the name.

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;
Expand Down
10 changes: 6 additions & 4 deletions cpp/src/preprocessing/quantize/pq.cu
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@

#include "./detail/pq.cuh"

#include <cuvs/core/cuda_fp16.hpp>
#include <cuvs/core/export.hpp>
#include <cuvs/preprocessing/quantize/pq.hpp>

#include <raft/matrix/copy.cuh>
Expand Down Expand Up @@ -101,28 +103,28 @@ auto vpq_train_from_device_rows(raft::resources const& res,

} // namespace detail

template cuvs::neighbors::device_vpq_dataset<half, int64_t>
template CUVS_EXPORT cuvs::neighbors::device_vpq_dataset<half, int64_t>
detail::vpq_train_from_device_rows<float>(raft::resources const&,
cuvs::neighbors::vpq_params const&,
float const*,
int64_t,
int64_t,
int64_t);
template cuvs::neighbors::device_vpq_dataset<half, int64_t>
template CUVS_EXPORT cuvs::neighbors::device_vpq_dataset<half, int64_t>
detail::vpq_train_from_device_rows<half>(raft::resources const&,
cuvs::neighbors::vpq_params const&,
half const*,
int64_t,
int64_t,
int64_t);
template cuvs::neighbors::device_vpq_dataset<half, int64_t>
template CUVS_EXPORT cuvs::neighbors::device_vpq_dataset<half, int64_t>
detail::vpq_train_from_device_rows<int8_t>(raft::resources const&,
cuvs::neighbors::vpq_params const&,
int8_t const*,
int64_t,
int64_t,
int64_t);
template cuvs::neighbors::device_vpq_dataset<half, int64_t>
template CUVS_EXPORT cuvs::neighbors::device_vpq_dataset<half, int64_t>
detail::vpq_train_from_device_rows<uint8_t>(raft::resources const&,
cuvs::neighbors::vpq_params const&,
uint8_t const*,
Expand Down
Loading
Loading