Skip to content

Commit 9b371e4

Browse files
GPUtesterGPUtester
authored
Merge pull request #1579 from rapidsai/release/25.12
Forward-merge release/25.12 into main
2 parents 0136dd0 + ad1438f commit 9b371e4

File tree

3 files changed

+170
-1
lines changed

3 files changed

+170
-1
lines changed

‎c/src/core/c_api.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,7 @@ extern "C" cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_per
185185
extern "C" cuvsError_t cuvsRMMMemoryResourceReset()
186186
{
187187
return cuvs::core::translate_exceptions([=] {
188-
rmm::mr::set_current_device_resource(nullptr);
188+
rmm::mr::set_current_device_resource(rmm::mr::detail::initial_resource());
189189
pool_mr.reset();
190190
});
191191
}

‎examples/c/CMakeLists.txt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,3 +45,7 @@ target_link_libraries(IVF_FLAT_C_EXAMPLE PRIVATE cuvs::c_api $<TARGET_NAME_IF_EX
4545
add_executable(IVF_PQ_C_EXAMPLE src/ivf_pq_c_example.c)
4646
target_include_directories(IVF_PQ_C_EXAMPLE PUBLIC "$<BUILD_INTERFACE:${DLPACK_INCLUDE_DIR}>")
4747
target_link_libraries(IVF_PQ_C_EXAMPLE PRIVATE cuvs::c_api $<TARGET_NAME_IF_EXISTS:conda_env>)
48+
49+
add_executable(BRUTEFORCE_C_EXAMPLE src/bruteforce_c_example.c)
50+
target_include_directories(BRUTEFORCE_C_EXAMPLE PUBLIC "$<BUILD_INTERFACE:${DLPACK_INCLUDE_DIR}>")
51+
target_link_libraries(BRUTEFORCE_C_EXAMPLE PRIVATE cuvs::c_api $<TARGET_NAME_IF_EXISTS:conda_env>)

‎examples/c/src/bruteforce_c_example.c‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
/*
2+
* SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
#include <cuvs/neighbors/brute_force.h>
7+
#include <stdint.h>
8+
9+
static const char dataset[] = {0.74021935f,
10+
0.9209938f,
11+
0.03902049f,
12+
0.9689629f,
13+
0.92514056f,
14+
0.4463501f,
15+
0.6673192f,
16+
0.10993068f};
17+
18+
static const char queries[] = {0.48216683f,
19+
0.0428398f,
20+
0.5084142f,
21+
0.6545497f,
22+
0.51260436f,
23+
0.2643005f,
24+
0.05198065f,
25+
0.5789965f};
26+
27+
void index_and_search()
28+
{
29+
int64_t n_rows = 4;
30+
int64_t n_queries = 4;
31+
int64_t n_dim = 2;
32+
uint32_t n_neighbors = 2;
33+
34+
float* index_data;
35+
float* query_data;
36+
37+
long indexBytes = sizeof(float) * n_rows * n_dim;
38+
long queriesBytes = sizeof(float) * n_queries * n_dim;
39+
long neighborsBytes = sizeof(long) * n_queries * n_neighbors;
40+
long distanceBytes = sizeof(float) * n_queries * n_neighbors;
41+
42+
uint32_t* prefilter_data = NULL;
43+
enum cuvsFilterType prefilter_type = NO_FILTER;
44+
45+
float* distances_data;
46+
int64_t* neighbors_data;
47+
48+
// create cuvsResources_t
49+
cuvsResources_t res;
50+
cuvsResourcesCreate(&res);
51+
52+
cuvsRMMAlloc(res, (void**)&index_data, indexBytes);
53+
cuvsRMMAlloc(res, (void**)&query_data, queriesBytes);
54+
cuvsRMMAlloc(res, (void**)&distances_data, distanceBytes);
55+
cuvsRMMAlloc(res, (void**)&neighbors_data, neighborsBytes);
56+
57+
cudaMemcpy(index_data, dataset, indexBytes, cudaMemcpyHostToDevice);
58+
cudaMemcpy(query_data, queries, queriesBytes, cudaMemcpyHostToDevice);
59+
60+
// create dataset DLTensor
61+
DLManagedTensor dataset_tensor;
62+
dataset_tensor.dl_tensor.data = index_data;
63+
dataset_tensor.dl_tensor.device.device_type = kDLCUDA;
64+
dataset_tensor.dl_tensor.ndim = 2;
65+
dataset_tensor.dl_tensor.dtype.code = kDLFloat;
66+
dataset_tensor.dl_tensor.dtype.bits = 32;
67+
dataset_tensor.dl_tensor.dtype.lanes = 1;
68+
int64_t dataset_shape[2] = {n_rows, n_dim};
69+
dataset_tensor.dl_tensor.shape = dataset_shape;
70+
dataset_tensor.dl_tensor.strides = NULL;
71+
72+
// create index
73+
cuvsBruteForceIndex_t index;
74+
cuvsBruteForceIndexCreate(&index);
75+
76+
// build index
77+
cuvsBruteForceBuild(res, &dataset_tensor, 0, 0.0f, index);
78+
79+
// create queries DLTensor
80+
DLManagedTensor queries_tensor;
81+
queries_tensor.dl_tensor.data = (void*)query_data;
82+
queries_tensor.dl_tensor.device.device_type = kDLCUDA;
83+
queries_tensor.dl_tensor.ndim = 2;
84+
queries_tensor.dl_tensor.dtype.code = kDLFloat;
85+
queries_tensor.dl_tensor.dtype.bits = 32;
86+
queries_tensor.dl_tensor.dtype.lanes = 1;
87+
int64_t queries_shape[2] = {n_queries, n_dim};
88+
queries_tensor.dl_tensor.shape = queries_shape;
89+
queries_tensor.dl_tensor.strides = NULL;
90+
91+
// create neighbors DLTensor
92+
DLManagedTensor neighbors_tensor;
93+
neighbors_tensor.dl_tensor.data = (void*)neighbors_data;
94+
neighbors_tensor.dl_tensor.device.device_type = kDLCUDA;
95+
neighbors_tensor.dl_tensor.ndim = 2;
96+
neighbors_tensor.dl_tensor.dtype.code = kDLInt;
97+
neighbors_tensor.dl_tensor.dtype.bits = 64;
98+
neighbors_tensor.dl_tensor.dtype.lanes = 1;
99+
int64_t neighbors_shape[2] = {n_queries, n_neighbors};
100+
neighbors_tensor.dl_tensor.shape = neighbors_shape;
101+
neighbors_tensor.dl_tensor.strides = NULL;
102+
103+
// create distances DLTensor
104+
DLManagedTensor distances_tensor;
105+
distances_tensor.dl_tensor.data = (void*)distances_data;
106+
distances_tensor.dl_tensor.device.device_type = kDLCUDA;
107+
distances_tensor.dl_tensor.ndim = 2;
108+
distances_tensor.dl_tensor.dtype.code = kDLFloat;
109+
distances_tensor.dl_tensor.dtype.bits = 32;
110+
distances_tensor.dl_tensor.dtype.lanes = 1;
111+
int64_t distances_shape[2] = {n_queries, n_neighbors};
112+
distances_tensor.dl_tensor.shape = distances_shape;
113+
distances_tensor.dl_tensor.strides = NULL;
114+
115+
cuvsFilter prefilter;
116+
117+
DLManagedTensor prefilter_tensor;
118+
if (prefilter_data == NULL || prefilter_type == NO_FILTER) {
119+
prefilter.type = NO_FILTER;
120+
prefilter.addr = (uintptr_t)NULL;
121+
} else {
122+
prefilter_tensor.dl_tensor.data = (void*)prefilter_data;
123+
prefilter_tensor.dl_tensor.device.device_type = kDLCUDA;
124+
prefilter_tensor.dl_tensor.ndim = 1;
125+
prefilter_tensor.dl_tensor.dtype.code = kDLUInt;
126+
prefilter_tensor.dl_tensor.dtype.bits = 32;
127+
prefilter_tensor.dl_tensor.dtype.lanes = 1;
128+
129+
int64_t prefilter_bits_num = (prefilter_type == BITMAP) ? n_queries * n_rows : n_rows;
130+
int64_t prefilter_shape[1] = {(prefilter_bits_num + 31) / 32};
131+
132+
prefilter_tensor.dl_tensor.shape = prefilter_shape;
133+
prefilter_tensor.dl_tensor.strides = NULL;
134+
135+
prefilter.type = prefilter_type;
136+
prefilter.addr = (uintptr_t)&prefilter_tensor;
137+
}
138+
139+
// search index
140+
cuvsBruteForceSearch(
141+
res, index, &queries_tensor, &neighbors_tensor, &distances_tensor, prefilter);
142+
143+
// de-allocate index and res
144+
cuvsBruteForceIndexDestroy(index);
145+
146+
cuvsRMMFree(res, index_data, indexBytes);
147+
cuvsRMMFree(res, query_data, queriesBytes);
148+
cuvsRMMFree(res, distances_data, distanceBytes);
149+
cuvsRMMFree(res, neighbors_data, neighborsBytes);
150+
151+
cuvsResourcesDestroy(res);
152+
}
153+
154+
int main()
155+
{
156+
// Perform indexing and search with pooled resources
157+
cuvsRMMPoolMemoryResourceEnable(10, 60, false);
158+
index_and_search();
159+
160+
// Perform indexing and search with the default memory resources
161+
cuvsRMMMemoryResourceReset();
162+
index_and_search();
163+
164+
return 0;
165+
}

0 commit comments

Comments
 (0)