Skip to content

Commit a5f8d43

Browse files
GPUtesterGPUtester
authored
Merge pull request #1729 from rapidsai/release/26.02
Forward-merge release/26.02 into main
2 parents 7da9d10 + d1b11d5 commit a5f8d43

37 files changed

Lines changed: 3004 additions & 213 deletions

File tree

‎c/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@ add_library(
100100
src/neighbors/tiered_index.cpp
101101
src/neighbors/all_neighbors.cpp
102102
src/preprocessing/quantize/binary.cpp
103+
src/preprocessing/quantize/pq.cpp
103104
src/preprocessing/quantize/scalar.cpp
104105
src/distance/pairwise_distance.cpp
105106
)

‎c/include/cuvs/cluster/kmeans.h‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

@@ -121,6 +121,11 @@ cuvsError_t cuvsKMeansParamsCreate(cuvsKMeansParams_t* params);
121121
*/
122122
cuvsError_t cuvsKMeansParamsDestroy(cuvsKMeansParams_t params);
123123

124+
/**
125+
* @brief Type of k-means algorithm.
126+
*/
127+
typedef enum { CUVS_KMEANS_TYPE_KMEANS = 0, CUVS_KMEANS_TYPE_KMEANS_BALANCED = 1 } cuvsKMeansType;
128+
124129
/**
125130
* @}
126131
*/

‎c/include/cuvs/core/all.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

@@ -39,4 +39,5 @@
3939
#endif
4040

4141
#include <cuvs/preprocessing/quantize/binary.h>
42+
#include <cuvs/preprocessing/quantize/pq.h>
4243
#include <cuvs/preprocessing/quantize/scalar.h>

‎c/include/cuvs/preprocessing/quantize/binary.h‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

@@ -13,6 +13,10 @@
1313
extern "C" {
1414
#endif
1515

16+
/**
17+
* @defgroup preprocessing_c_binary C API for Binary Quantizer
18+
* @{
19+
*/
1620
/**
1721
* @brief In the cuvsBinaryQuantizerTransform function, a bit is set if the corresponding element in
1822
* the dataset vector is greater than the corresponding element in the threshold vector. The mean
@@ -132,6 +136,10 @@ cuvsError_t cuvsBinaryQuantizerTransformWithParams(cuvsResources_t res,
132136
DLManagedTensor* dataset,
133137
DLManagedTensor* out);
134138

139+
/**
140+
* @}
141+
*/
142+
135143
#ifdef __cplusplus
136144
}
137145
#endif

‎c/include/cuvs/preprocessing/quantize/pq.h‎

Lines changed: 227 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,227 @@
1+
/*
2+
* SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
#pragma once
7+
8+
#include <cuvs/cluster/kmeans.h>
9+
#include <cuvs/core/c_api.h>
10+
#include <dlpack/dlpack.h>
11+
#include <stdint.h>
12+
13+
#ifdef __cplusplus
14+
extern "C" {
15+
#endif
16+
17+
/**
18+
* @defgroup preprocessing_c_pq C API for Product Quantizer
19+
* @{
20+
*/
21+
/**
22+
* @brief Product quantizer parameters.
23+
*/
24+
struct cuvsProductQuantizerParams {
25+
/**
26+
* The bit length of the vector element after compression by PQ.
27+
*
28+
* Possible values: within [4, 16].
29+
*
30+
* Hint: the smaller the 'pq_bits', the smaller the index size and the better the search
31+
* performance, but the lower the recall.
32+
*/
33+
uint32_t pq_bits;
34+
/**
35+
* The dimensionality of the vector after compression by PQ.
36+
* When zero, an optimal value is selected using a heuristic.
37+
*
38+
* TODO: at the moment `dim` must be a multiple `pq_dim`.
39+
*/
40+
uint32_t pq_dim;
41+
/**
42+
* Whether to use subspaces for product quantization (PQ).
43+
* When true, one PQ codebook is used for each subspace. Otherwise, a single
44+
* PQ codebook is used.
45+
*/
46+
bool use_subspaces;
47+
/**
48+
* Whether to use Vector Quantization (KMeans) before product quantization (PQ).
49+
* When true, VQ is used before PQ. When false, only product quantization is used.
50+
*/
51+
bool use_vq;
52+
/**
53+
* Vector Quantization (VQ) codebook size - number of "coarse cluster centers".
54+
* When zero, an optimal value is selected using a heuristic.
55+
* When one, only product quantization is used.
56+
*/
57+
uint32_t vq_n_centers;
58+
/** The number of iterations searching for kmeans centers (both VQ & PQ phases). */
59+
uint32_t kmeans_n_iters;
60+
/**
61+
* The type of kmeans algorithm to use for PQ training.
62+
*/
63+
cuvsKMeansType pq_kmeans_type;
64+
/**
65+
* The max number of data points to use per PQ code during PQ codebook training. Using more data
66+
* points per PQ code may increase the quality of PQ codebook but may also increase the build
67+
* time. We will use `pq_n_centers * max_train_points_per_pq_code` training
68+
* points to train each PQ codebook.
69+
*/
70+
uint32_t max_train_points_per_pq_code;
71+
/**
72+
* The max number of data points to use per VQ cluster.
73+
*/
74+
uint32_t max_train_points_per_vq_cluster;
75+
};
76+
77+
typedef struct cuvsProductQuantizerParams* cuvsProductQuantizerParams_t;
78+
79+
/**
80+
* @brief Allocate Product Quantizer params, and populate with default values
81+
*
82+
* @param[in] params cuvsProductQuantizerParams_t to allocate
83+
* @return cuvsError_t
84+
*/
85+
cuvsError_t cuvsProductQuantizerParamsCreate(cuvsProductQuantizerParams_t* params);
86+
87+
/**
88+
* @brief De-allocate Product Quantizer params
89+
*
90+
* @param[in] params
91+
* @return cuvsError_t
92+
*/
93+
cuvsError_t cuvsProductQuantizerParamsDestroy(cuvsProductQuantizerParams_t params);
94+
95+
/**
96+
* @brief Defines and stores product quantizer upon training
97+
*
98+
* The quantization is performed by a linear mapping of an interval in the
99+
* float data type to the full range of the quantized int type.
100+
*/
101+
typedef struct {
102+
uintptr_t addr;
103+
DLDataType dtype;
104+
} cuvsProductQuantizer;
105+
106+
typedef cuvsProductQuantizer* cuvsProductQuantizer_t;
107+
108+
/**
109+
* @brief Allocate Product Quantizer
110+
*
111+
* @param[in] quantizer cuvsProductQuantizer_t to allocate
112+
* @return cuvsError_t
113+
*/
114+
cuvsError_t cuvsProductQuantizerCreate(cuvsProductQuantizer_t* quantizer);
115+
116+
/**
117+
* @brief De-allocate Product Quantizer
118+
*
119+
* @param[in] quantizer
120+
* @return cuvsError_t
121+
*/
122+
cuvsError_t cuvsProductQuantizerDestroy(cuvsProductQuantizer_t quantizer);
123+
124+
/**
125+
* @brief Builds a product quantizer to be used later for quantizing the dataset.
126+
*
127+
* @param[in] res raft resource
128+
* @param[in] params Parameters for product quantizer training
129+
* @param[in] dataset a row-major host or device matrix
130+
* @param[out] quantizer trained product quantizer
131+
*/
132+
cuvsError_t cuvsProductQuantizerBuild(cuvsResources_t res,
133+
cuvsProductQuantizerParams_t params,
134+
DLManagedTensor* dataset,
135+
cuvsProductQuantizer_t quantizer);
136+
137+
/**
138+
* @brief Applies product quantization transform to the given dataset
139+
*
140+
* This applies product quantization to a dataset.
141+
*
142+
* @param[in] res raft resource
143+
* @param[in] quantizer product quantizer
144+
* @param[in] dataset a row-major host or device matrix to transform
145+
* @param[out] codes_out a row-major device matrix to store transformed data
146+
* @param[out] vq_labels a device vector to store VQ labels.
147+
* Optional, can be NULL.
148+
*/
149+
cuvsError_t cuvsProductQuantizerTransform(cuvsResources_t res,
150+
cuvsProductQuantizer_t quantizer,
151+
DLManagedTensor* dataset,
152+
DLManagedTensor* codes_out,
153+
DLManagedTensor* vq_labels);
154+
155+
/**
156+
* @brief Applies product quantization inverse transform to the given quantized codes
157+
*
158+
* This applies product quantization inverse transform to the given quantized codes.
159+
*
160+
* @param[in] res raft resource
161+
* @param[in] quantizer product quantizer
162+
* @param[in] pq_codes a row-major device matrix of quantized codes
163+
* @param[out] out a row-major device matrix to store the original data
164+
* @param[out] vq_labels a device vector containing the VQ labels when VQ is used.
165+
* Optional, can be NULL.
166+
*/
167+
cuvsError_t cuvsProductQuantizerInverseTransform(cuvsResources_t res,
168+
cuvsProductQuantizer_t quantizer,
169+
DLManagedTensor* pq_codes,
170+
DLManagedTensor* out,
171+
DLManagedTensor* vq_labels);
172+
173+
/**
174+
* @brief Get the bit length of the vector element after compression by PQ.
175+
*
176+
* @param[in] quantizer product quantizer
177+
* @param[out] pq_bits bit length of the vector element after compression by PQ
178+
*/
179+
cuvsError_t cuvsProductQuantizerGetPqBits(cuvsProductQuantizer_t quantizer, uint32_t* pq_bits);
180+
181+
/**
182+
* @brief Get the dimensionality of the vector after compression by PQ.
183+
*
184+
* @param[in] quantizer product quantizer
185+
* @param[out] pq_dim dimensionality of the vector after compression by PQ
186+
*/
187+
cuvsError_t cuvsProductQuantizerGetPqDim(cuvsProductQuantizer_t quantizer, uint32_t* pq_dim);
188+
189+
/**
190+
* @brief Get the PQ codebook.
191+
*
192+
* @param[in] quantizer product quantizer
193+
* @param[out] pq_codebook PQ codebook
194+
*/
195+
cuvsError_t cuvsProductQuantizerGetPqCodebook(cuvsProductQuantizer_t quantizer,
196+
DLManagedTensor* pq_codebook);
197+
198+
/**
199+
* @brief Get the VQ codebook.
200+
*
201+
* @param[in] quantizer product quantizer
202+
* @param[out] vq_codebook VQ codebook
203+
*/
204+
cuvsError_t cuvsProductQuantizerGetVqCodebook(cuvsProductQuantizer_t quantizer,
205+
DLManagedTensor* vq_codebook);
206+
/**
207+
* @brief Get the encoded dimension of the quantized dataset.
208+
*
209+
* @param[in] quantizer product quantizer
210+
* @param[out] encoded_dim encoded dimension of the quantized dataset
211+
*/
212+
cuvsError_t cuvsProductQuantizerGetEncodedDim(cuvsProductQuantizer_t quantizer,
213+
uint32_t* encoded_dim);
214+
215+
/**
216+
* @brief Get whether VQ is used.
217+
*
218+
* @param[in] quantizer product quantizer
219+
* @param[out] use_vq whether VQ is used
220+
*/
221+
cuvsError_t cuvsProductQuantizerGetUseVq(cuvsProductQuantizer_t quantizer, bool* use_vq);
222+
/**
223+
* @}
224+
*/
225+
#ifdef __cplusplus
226+
}
227+
#endif

‎c/include/cuvs/preprocessing/quantize/scalar.h‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

@@ -13,6 +13,10 @@
1313
extern "C" {
1414
#endif
1515

16+
/**
17+
* @defgroup preprocessing_c_scalar C API for Scalar Quantizer
18+
* @{
19+
*/
1620
/**
1721
* @brief Scalar quantizer parameters.
1822
*/
@@ -114,6 +118,10 @@ cuvsError_t cuvsScalarQuantizerInverseTransform(cuvsResources_t res,
114118
DLManagedTensor* dataset,
115119
DLManagedTensor* out);
116120

121+
/**
122+
* @}
123+
*/
124+
117125
#ifdef __cplusplus
118126
}
119127
#endif

0 commit comments

Comments
 (0)