forked from rapidsai/cuvs
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathnn_descent_float.cu
More file actions
88 lines (82 loc) · 6.41 KB
/
nn_descent_float.cu
File metadata and controls
88 lines (82 loc) · 6.41 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
/*
* Copyright (c) 2024-2025, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "./detail/nn_descent_gnnd.hpp"
#include "./detail/reachability_types.cuh"
#include "nn_descent.cuh"
#include <cuvs/neighbors/nn_descent.hpp>
namespace cuvs::neighbors::nn_descent {
#define CUVS_INST_NN_DESCENT_BUILD(T, IdxT) \
auto build(raft::resources const& handle, \
const cuvs::neighbors::nn_descent::index_params& params, \
raft::device_matrix_view<const T, int64_t, raft::row_major> dataset, \
std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph) \
-> cuvs::neighbors::nn_descent::index<IdxT> \
{ \
if (!graph.has_value()) { \
return cuvs::neighbors::nn_descent::build<T, IdxT>(handle, params, dataset); \
} else { \
std::optional<raft::device_matrix_view<float, int64_t, raft::row_major>> distances = \
std::nullopt; \
cuvs::neighbors::nn_descent::index<IdxT> idx{ \
handle, graph.value(), distances, params.metric}; \
cuvs::neighbors::nn_descent::build<T, IdxT>(handle, params, dataset, idx); \
return idx; \
}; \
} \
auto build(raft::resources const& handle, \
const cuvs::neighbors::nn_descent::index_params& params, \
raft::host_matrix_view<const T, int64_t, raft::row_major> dataset, \
std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph) \
-> cuvs::neighbors::nn_descent::index<IdxT> \
{ \
if (!graph.has_value()) { \
return cuvs::neighbors::nn_descent::build<T, IdxT>(handle, params, dataset); \
} else { \
std::optional<raft::device_matrix_view<float, int64_t, raft::row_major>> distances = \
std::nullopt; \
cuvs::neighbors::nn_descent::index<IdxT> idx{ \
handle, graph.value(), distances, params.metric}; \
cuvs::neighbors::nn_descent::build<T, IdxT>(handle, params, dataset, idx); \
return idx; \
} \
}; \
template class detail::GNND<const T, int>; \
\
template void detail::GNND<const T, int>::build< \
cuvs::neighbors::detail::reachability::ReachabilityPostProcess<int, T>>( \
const T* data, \
const int nrow, \
int* output_graph, \
bool return_distances, \
float* output_distances, \
cuvs::neighbors::detail::reachability::ReachabilityPostProcess<int, T> dist_epilogue); \
template void detail::GNND<const T, int>::local_join< \
cuvs::neighbors::detail::reachability::ReachabilityPostProcess<int, T>>( \
cudaStream_t stream, \
cuvs::neighbors::detail::reachability::ReachabilityPostProcess<int, T> dist_epilogue); \
\
template void detail::GNND<const T, int>::build<raft::identity_op>( \
const T* data, \
const int nrow, \
int* output_graph, \
bool return_distances, \
float* output_distances, \
raft::identity_op dist_epilogue); \
template void detail::GNND<const T, int>::local_join<raft::identity_op>( \
cudaStream_t stream, raft::identity_op dist_epilogue);
CUVS_INST_NN_DESCENT_BUILD(float, uint32_t);
#undef CUVS_INST_NN_DESCENT_BUILD
} // namespace cuvs::neighbors::nn_descent