Skip to content

Commit 0cef61e

Browse files
GPUtesterGPUtester
authored
Merge pull request #653 from rapidsai/branch-25.02
Forward-merge branch-25.02 into branch-25.04
2 parents 5ad560c + 45703bf commit 0cef61e

13 files changed

Lines changed: 768 additions & 3 deletions

File tree

‎cpp/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -691,6 +691,7 @@ target_compile_definitions(cuvs::cuvs INTERFACE $<$<BOOL:${CUVS_NVTX}>:NVTX_ENAB
691691
src/neighbors/ivf_pq_c.cpp
692692
src/neighbors/cagra_c.cpp
693693
$<$<BOOL:${BUILD_CAGRA_HNSWLIB}>:src/neighbors/hnsw_c.cpp>
694+
src/neighbors/nn_descent_c.cpp
694695
src/neighbors/refine/refine_c.cpp
695696
src/preprocessing/quantize/scalar_c.cpp
696697
src/distance/pairwise_distance_c.cpp

‎cpp/include/cuvs/neighbors/nn_descent.h‎

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,181 @@
1+
/*
2+
* Copyright (c) 2024, NVIDIA CORPORATION.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
#pragma once
18+
19+
#include <cuvs/core/c_api.h>
20+
#include <cuvs/distance/distance.h>
21+
#include <dlpack/dlpack.h>
22+
#include <stdbool.h>
23+
#include <stdint.h>
24+
25+
#ifdef __cplusplus
26+
extern "C" {
27+
#endif
28+
29+
/**
30+
* @defgroup nn_descent_c_index_params The nn-descent algorithm parameters.
31+
* @{
32+
*/
33+
/**
34+
* @brief Parameters used to build an nn-descent index
35+
*
36+
* `metric`: The distance metric to use
37+
* `metric_arg`: The argument used by distance metrics like Minkowskidistance
38+
* `graph_degree`: For an input dataset of dimensions (N, D),
39+
* determines the final dimensions of the all-neighbors knn graph
40+
* which turns out to be of dimensions (N, graph_degree)
41+
* `intermediate_graph_degree`: Internally, nn-descent builds an
42+
* all-neighbors knn graph of dimensions (N, intermediate_graph_degree)
43+
* before selecting the final `graph_degree` neighbors. It's recommended
44+
* that `intermediate_graph_degree` >= 1.5 * graph_degree
45+
* `max_iterations`: The number of iterations that nn-descent will refine
46+
* the graph for. More iterations produce a better quality graph at cost of performance
47+
* `termination_threshold`: The delta at which nn-descent will terminate its iterations
48+
*/
49+
struct cuvsNNDescentIndexParams {
50+
cuvsDistanceType metric;
51+
float metric_arg;
52+
size_t graph_degree;
53+
size_t intermediate_graph_degree;
54+
size_t max_iterations;
55+
float termination_threshold;
56+
bool return_distances;
57+
size_t n_clusters;
58+
};
59+
60+
typedef struct cuvsNNDescentIndexParams* cuvsNNDescentIndexParams_t;
61+
62+
/**
63+
* @brief Allocate NN-Descent Index params, and populate with default values
64+
*
65+
* @param[in] index_params cuvsNNDescentIndexParams_t to allocate
66+
* @return cuvsError_t
67+
*/
68+
cuvsError_t cuvsNNDescentIndexParamsCreate(cuvsNNDescentIndexParams_t* index_params);
69+
70+
/**
71+
* @brief De-allocate NN-Descent Index params
72+
*
73+
* @param[in] index_params
74+
* @return cuvsError_t
75+
*/
76+
cuvsError_t cuvsNNDescentIndexParamsDestroy(cuvsNNDescentIndexParams_t index_params);
77+
/**
78+
* @}
79+
*/
80+
81+
/**
82+
* @defgroup nn_descent_c_index NN-Descent index
83+
* @{
84+
*/
85+
/**
86+
* @brief Struct to hold address of cuvs::neighbors::nn_descent::index and its active trained dtype
87+
*
88+
*/
89+
typedef struct {
90+
uintptr_t addr;
91+
DLDataType dtype;
92+
} cuvsNNDescentIndex;
93+
94+
typedef cuvsNNDescentIndex* cuvsNNDescentIndex_t;
95+
96+
/**
97+
* @brief Allocate NN-Descent index
98+
*
99+
* @param[in] index cuvsNNDescentIndex_t to allocate
100+
* @return cuvsError_t
101+
*/
102+
cuvsError_t cuvsNNDescentIndexCreate(cuvsNNDescentIndex_t* index);
103+
104+
/**
105+
* @brief De-allocate NN-Descent index
106+
*
107+
* @param[in] index cuvsNNDescentIndex_t to de-allocate
108+
*/
109+
cuvsError_t cuvsNNDescentIndexDestroy(cuvsNNDescentIndex_t index);
110+
/**
111+
* @}
112+
*/
113+
114+
/**
115+
* @defgroup nn_descent_c_index_build NN-Descent index build
116+
* @{
117+
*/
118+
/**
119+
* @brief Build a NN-Descent index with a `DLManagedTensor` which has underlying
120+
* `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`,
121+
* or `kDLCPU`. Also, acceptable underlying types are:
122+
* 1. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`
123+
* 2. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 16`
124+
* 3. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8`
125+
* 4. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8`
126+
*
127+
* @code {.c}
128+
* #include <cuvs/core/c_api.h>
129+
* #include <cuvs/neighbors/nn_descent.h>
130+
*
131+
* // Create cuvsResources_t
132+
* cuvsResources_t res;
133+
* cuvsError_t res_create_status = cuvsResourcesCreate(&res);
134+
*
135+
* // Assume a populated `DLManagedTensor` type here
136+
* DLManagedTensor dataset;
137+
*
138+
* // Create default index params
139+
* cuvsNNDescentIndexParams_t index_params;
140+
* cuvsError_t params_create_status = cuvsNNDescentIndexParamsCreate(&index_params);
141+
*
142+
* // Create NN-Descent index
143+
* cuvsNNDescentIndex_t index;
144+
* cuvsError_t index_create_status = cuvsNNDescentIndexCreate(&index);
145+
*
146+
* // Build the NN-Descent Index
147+
* cuvsError_t build_status = cuvsNNDescentBuild(res, index_params, &dataset, index);
148+
*
149+
* // de-allocate `index_params`, `index` and `res`
150+
* cuvsError_t params_destroy_status = cuvsNNDescentIndexParamsDestroy(index_params);
151+
* cuvsError_t index_destroy_status = cuvsNNDescentIndexDestroy(index);
152+
* cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);
153+
* @endcode
154+
*
155+
* @param[in] res cuvsResources_t opaque C handle
156+
* @param[in] index_params cuvsNNDescentIndexParams_t used to build NN-Descent index
157+
* @param[in] dataset DLManagedTensor* training dataset on host or device memory
158+
* @param[inout] graph Optional preallocated graph on host memory to store output
159+
* @param[out] index cuvsNNDescentIndex_t Newly built NN-Descent index
160+
* @return cuvsError_t
161+
*/
162+
cuvsError_t cuvsNNDescentBuild(cuvsResources_t res,
163+
cuvsNNDescentIndexParams_t index_params,
164+
DLManagedTensor* dataset,
165+
DLManagedTensor* graph,
166+
cuvsNNDescentIndex_t index);
167+
/**
168+
* @}
169+
*/
170+
171+
/**
172+
* @brief Get the KNN graph from a built NN-Descent index
173+
*
174+
* @param[in] index cuvsNNDescentIndex_t Built NN-Descent index
175+
* @param[inout] graph Optional preallocated graph on host memory to store output
176+
* @return cuvsError_t
177+
*/
178+
cuvsError_t cuvsNNDescentIndexGetGraph(cuvsNNDescentIndex_t index, DLManagedTensor* graph);
179+
#ifdef __cplusplus
180+
}
181+
#endif

‎cpp/src/neighbors/nn_descent_c.cpp‎

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
/*
2+
* Copyright (c) 2025, NVIDIA CORPORATION.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
#include <cstdint>
18+
#include <dlpack/dlpack.h>
19+
20+
#include <raft/core/copy.hpp>
21+
#include <raft/core/error.hpp>
22+
#include <raft/core/mdspan_types.hpp>
23+
#include <raft/core/resource/cuda_stream.hpp>
24+
#include <raft/core/resources.hpp>
25+
#include <raft/core/serialize.hpp>
26+
27+
#include <cuvs/core/c_api.h>
28+
#include <cuvs/core/exceptions.hpp>
29+
#include <cuvs/core/interop.hpp>
30+
#include <cuvs/neighbors/nn_descent.h>
31+
#include <cuvs/neighbors/nn_descent.hpp>
32+
33+
#include <fstream>
34+
35+
namespace {
36+
37+
template <typename T, typename IdxT = uint32_t>
38+
void* _build(cuvsResources_t res,
39+
cuvsNNDescentIndexParams params,
40+
DLManagedTensor* dataset_tensor,
41+
DLManagedTensor* graph_tensor)
42+
{
43+
auto res_ptr = reinterpret_cast<raft::resources*>(res);
44+
auto dataset = dataset_tensor->dl_tensor;
45+
46+
auto build_params = cuvs::neighbors::nn_descent::index_params();
47+
build_params.metric = static_cast<cuvs::distance::DistanceType>((int)params.metric),
48+
build_params.metric_arg = params.metric_arg;
49+
build_params.graph_degree = params.graph_degree;
50+
build_params.intermediate_graph_degree = params.intermediate_graph_degree;
51+
build_params.max_iterations = params.max_iterations;
52+
build_params.termination_threshold = params.termination_threshold;
53+
build_params.return_distances = params.return_distances;
54+
build_params.n_clusters = params.n_clusters;
55+
56+
using graph_type = raft::host_matrix_view<IdxT, int64_t, raft::row_major>;
57+
std::optional<graph_type> graph;
58+
if (graph_tensor != NULL) { graph = cuvs::core::from_dlpack<graph_type>(graph_tensor); }
59+
60+
if (cuvs::core::is_dlpack_device_compatible(dataset)) {
61+
using dataset_type = raft::device_matrix_view<T const, int64_t, raft::row_major>;
62+
auto dataset = cuvs::core::from_dlpack<dataset_type>(dataset_tensor);
63+
auto index = cuvs::neighbors::nn_descent::build(*res_ptr, build_params, dataset, graph);
64+
return new cuvs::neighbors::nn_descent::index<IdxT>(std::move(index));
65+
} else if (cuvs::core::is_dlpack_host_compatible(dataset)) {
66+
using dataset_type = raft::host_matrix_view<T const, int64_t, raft::row_major>;
67+
auto dataset = cuvs::core::from_dlpack<dataset_type>(dataset_tensor);
68+
auto index = cuvs::neighbors::nn_descent::build(*res_ptr, build_params, dataset, graph);
69+
return new cuvs::neighbors::nn_descent::index<IdxT>(std::move(index));
70+
} else {
71+
RAFT_FAIL("dataset must be accessible on host or device memory");
72+
}
73+
}
74+
} // namespace
75+
76+
extern "C" cuvsError_t cuvsNNDescentIndexCreate(cuvsNNDescentIndex_t* index)
77+
{
78+
return cuvs::core::translate_exceptions([=] { *index = new cuvsNNDescentIndex{}; });
79+
}
80+
81+
extern "C" cuvsError_t cuvsNNDescentIndexDestroy(cuvsNNDescentIndex_t index_c_ptr)
82+
{
83+
return cuvs::core::translate_exceptions([=] {
84+
auto index = *index_c_ptr;
85+
if ((index.dtype.code == kDLUInt) && (index.dtype.bits == 32)) {
86+
auto index_ptr = reinterpret_cast<cuvs::neighbors::nn_descent::index<uint32_t>*>(index.addr);
87+
delete index_ptr;
88+
} else {
89+
RAFT_FAIL(
90+
"Unsupported nn-descent index dtype: %d and bits: %d", index.dtype.code, index.dtype.bits);
91+
}
92+
delete index_c_ptr;
93+
});
94+
}
95+
96+
extern "C" cuvsError_t cuvsNNDescentBuild(cuvsResources_t res,
97+
cuvsNNDescentIndexParams_t params,
98+
DLManagedTensor* dataset_tensor,
99+
DLManagedTensor* graph_tensor,
100+
cuvsNNDescentIndex_t index)
101+
{
102+
return cuvs::core::translate_exceptions([=] {
103+
index->dtype.code = kDLUInt;
104+
index->dtype.bits = 32;
105+
106+
auto dtype = dataset_tensor->dl_tensor.dtype;
107+
108+
if ((dtype.code == kDLFloat) && (dtype.bits == 32)) {
109+
index->addr = reinterpret_cast<uintptr_t>(
110+
_build<float, uint32_t>(res, *params, dataset_tensor, graph_tensor));
111+
} else if ((dtype.code == kDLFloat) && (dtype.bits == 16)) {
112+
index->addr = reinterpret_cast<uintptr_t>(
113+
_build<half, uint32_t>(res, *params, dataset_tensor, graph_tensor));
114+
} else if ((dtype.code == kDLInt) && (dtype.bits == 8)) {
115+
index->addr = reinterpret_cast<uintptr_t>(
116+
_build<int8_t, uint32_t>(res, *params, dataset_tensor, graph_tensor));
117+
} else if ((dtype.code == kDLUInt) && (dtype.bits == 8)) {
118+
index->addr = reinterpret_cast<uintptr_t>(
119+
_build<uint8_t, uint32_t>(res, *params, dataset_tensor, graph_tensor));
120+
} else {
121+
RAFT_FAIL("Unsupported nn-descent dataset dtype: %d and bits: %d", dtype.code, dtype.bits);
122+
}
123+
});
124+
}
125+
126+
extern "C" cuvsError_t cuvsNNDescentIndexParamsCreate(cuvsNNDescentIndexParams_t* params)
127+
{
128+
return cuvs::core::translate_exceptions([=] {
129+
// get defaults from cpp parameters struct
130+
cuvs::neighbors::nn_descent::index_params cpp_params;
131+
132+
*params = new cuvsNNDescentIndexParams{
133+
.metric = cpp_params.metric,
134+
.metric_arg = cpp_params.metric_arg,
135+
.graph_degree = cpp_params.graph_degree,
136+
.intermediate_graph_degree = cpp_params.intermediate_graph_degree,
137+
.max_iterations = cpp_params.max_iterations,
138+
.termination_threshold = cpp_params.termination_threshold,
139+
.return_distances = cpp_params.return_distances,
140+
.n_clusters = cpp_params.n_clusters};
141+
});
142+
}
143+
144+
extern "C" cuvsError_t cuvsNNDescentIndexParamsDestroy(cuvsNNDescentIndexParams_t params)
145+
{
146+
return cuvs::core::translate_exceptions([=] { delete params; });
147+
}
148+
149+
extern "C" cuvsError_t cuvsNNDescentIndexGetGraph(cuvsNNDescentIndex_t index,
150+
DLManagedTensor* graph)
151+
{
152+
return cuvs::core::translate_exceptions([=] {
153+
auto dtype = index->dtype;
154+
if ((dtype.code == kDLUInt) && (dtype.bits == 32)) {
155+
auto index_ptr = reinterpret_cast<cuvs::neighbors::nn_descent::index<uint32_t>*>(index->addr);
156+
using output_mdspan_type = raft::host_matrix_view<uint32_t, int64_t, raft::row_major>;
157+
auto dst = cuvs::core::from_dlpack<output_mdspan_type>(graph);
158+
auto src = index_ptr->graph();
159+
160+
RAFT_EXPECTS(src.extent(0) == dst.extent(0), "Output graph has incorrect number of rows");
161+
RAFT_EXPECTS(src.extent(1) == dst.extent(1), "Output graph has incorrect number of cols");
162+
std::copy(src.data_handle(), src.data_handle() + dst.size(), dst.data_handle());
163+
} else {
164+
RAFT_FAIL("Unsupported nn-descent index dtype: %d and bits: %d", dtype.code, dtype.bits);
165+
}
166+
});
167+
}

‎python/cuvs/cuvs/distance/__init__.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,6 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
from .distance import DISTANCE_TYPES, pairwise_distance
15+
from .distance import DISTANCE_NAMES, DISTANCE_TYPES, pairwise_distance
1616

17-
__all__ = ["DISTANCE_TYPES", "pairwise_distance"]
17+
__all__ = ["DISTANCE_NAMES", "DISTANCE_TYPES", "pairwise_distance"]

‎python/cuvs/cuvs/distance/distance.pyx‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@ DISTANCE_TYPES = {
4848
"dice": cuvsDistanceType.DiceExpanded,
4949
}
5050

51+
DISTANCE_NAMES = {v: k for k, v in DISTANCE_TYPES.items()}
52+
5153
SUPPORTED_DISTANCES = ["euclidean", "l1", "cityblock", "l2", "inner_product",
5254
"chebyshev", "minkowski", "canberra", "kl_divergence",
5355
"correlation", "russellrao", "hellinger", "lp",

‎python/cuvs/cuvs/neighbors/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ add_subdirectory(hnsw)
1818
add_subdirectory(ivf_flat)
1919
add_subdirectory(ivf_pq)
2020
add_subdirectory(filters)
21+
add_subdirectory(nn_descent)
2122

2223
# Set the list of Cython files to build
2324
set(cython_sources refine.pyx)

0 commit comments

Comments
 (0)