268037bc
Jiaming Guo
made code more el...
|
1
|
// right now the size of CUDA STACK is set to 1000, increase it if you mean to make deeper tree
|
cc09e435
David Mayerich
added Jack's KD-T...
|
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
|
// data should be stored in row-major
// x1,x2,x3,x4,x5......
// y1,y2,y3,y4,y5......
// ....................
// ....................
#ifndef KDTREE_H
#define KDTREE_H
#include "device_launch_parameters.h"
#include <cuda.h>
#include <cuda_runtime_api.h>
#include "cuda_runtime.h"
#include <vector>
#include <float.h>
#include <iostream>
#include <algorithm>
|
268037bc
Jiaming Guo
made code more el...
|
19
|
#include <stim/cuda/cudatools/error.h>
|
cc09e435
David Mayerich
added Jack's KD-T...
|
20
21
22
|
namespace stim {
namespace kdtree {
|
268037bc
Jiaming Guo
made code more el...
|
23
|
template<typename T, int D> // typename refers to float or double while D refers to dimension of points
|
cc09e435
David Mayerich
added Jack's KD-T...
|
24
|
struct point {
|
268037bc
Jiaming Guo
made code more el...
|
25
|
T dim[D]; // create a structure to store every one input point
|
cc09e435
David Mayerich
added Jack's KD-T...
|
26
27
28
|
};
template<typename T>
|
268037bc
Jiaming Guo
made code more el...
|
29
|
class kdnode {
|
cc09e435
David Mayerich
added Jack's KD-T...
|
30
|
public:
|
268037bc
Jiaming Guo
made code more el...
|
31
32
|
kdnode() { // constructor for initializing a kdnode
parent = NULL; // set every node's parent, left and right kdnode pointers to NULL
|
cc09e435
David Mayerich
added Jack's KD-T...
|
33
34
|
left = NULL;
right = NULL;
|
268037bc
Jiaming Guo
made code more el...
|
35
36
37
38
|
parent_idx = -1; // set parent node index to default -1
left_idx = -1;
right_idx = -1;
split_value = -1; // set split_value to default -1
|
cc09e435
David Mayerich
added Jack's KD-T...
|
39
|
}
|
268037bc
Jiaming Guo
made code more el...
|
40
41
42
43
44
45
|
int idx; // index of current node
int parent_idx, left_idx, right_idx; // index of parent, left and right nodes
kdnode *parent, *left, *right; // parent, left and right kdnodes
T split_value; // splitting value of current node
std::vector <size_t> indices; // it indicates the points' indices that current node has
size_t level; // tree level of current node
|
cc09e435
David Mayerich
added Jack's KD-T...
|
46
|
};
|
268037bc
Jiaming Guo
made code more el...
|
47
|
} // end of namespace kdtree
|
cc09e435
David Mayerich
added Jack's KD-T...
|
48
|
|
268037bc
Jiaming Guo
made code more el...
|
49
|
template <typename T, int D = 3> // set dimension of data to default 3
|
cc09e435
David Mayerich
added Jack's KD-T...
|
50
|
class cpu_kdtree {
|
cc09e435
David Mayerich
added Jack's KD-T...
|
51
|
protected:
|
268037bc
Jiaming Guo
made code more el...
|
52
53
54
55
56
57
|
int current_axis; // current judging axis
int cmps; // count how many time of comparisons (just for cpu-kdtree)
int n_id; // store the total number of nodes
std::vector <kdtree::point<T, D>> *tmp_points; // transfer or temp points
kdtree::kdnode<T> *root; // root node
static cpu_kdtree<T, D> *cur_tree_ptr;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
58
|
public:
|
268037bc
Jiaming Guo
made code more el...
|
59
60
61
|
cpu_kdtree() { // constructor for creating a cpu_kdtree
cur_tree_ptr = this; // create a class pointer points to the current class value
n_id = 0; // set total number of points to default 0
|
cc09e435
David Mayerich
added Jack's KD-T...
|
62
|
}
|
268037bc
Jiaming Guo
made code more el...
|
63
64
65
66
67
68
69
70
|
~cpu_kdtree() { // destructor of cpu_kdtree
std::vector <kdtree::kdnode<T>*> next_nodes;
next_nodes.push_back(root);
while (next_nodes.size()) {
std::vector <kdtree::kdnode<T>*> next_search_nodes;
while (next_nodes.size()) {
kdtree::kdnode<T> *cur = next_nodes.back();
next_nodes.pop_back();
|
cc09e435
David Mayerich
added Jack's KD-T...
|
71
|
if (cur->left)
|
268037bc
Jiaming Guo
made code more el...
|
72
|
next_search_nodes.push_back(cur->left);
|
cc09e435
David Mayerich
added Jack's KD-T...
|
73
|
if (cur->right)
|
268037bc
Jiaming Guo
made code more el...
|
74
|
next_search_nodes.push_back(cur->right);
|
cc09e435
David Mayerich
added Jack's KD-T...
|
75
76
|
delete cur;
}
|
268037bc
Jiaming Guo
made code more el...
|
77
|
next_nodes = next_search_nodes;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
78
|
}
|
268037bc
Jiaming Guo
made code more el...
|
79
|
root = NULL;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
80
|
}
|
268037bc
Jiaming Guo
made code more el...
|
81
82
83
84
85
86
87
88
|
void Create(std::vector <kdtree::point<T, D>> &reference_points, size_t max_levels) {
tmp_points = &reference_points;
root = new kdtree::kdnode<T>(); // initializing the root node
root->idx = n_id++; // the index of root is 0
root->level = 0; // tree level begins at 0
root->indices.resize(reference_points.size()); // get the number of points
for (size_t i = 0; i < reference_points.size(); i++) {
root->indices[i] = i; // set indices of input points
|
cc09e435
David Mayerich
added Jack's KD-T...
|
89
|
}
|
268037bc
Jiaming Guo
made code more el...
|
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
|
std::vector <kdtree::kdnode<T>*> next_nodes; // next nodes
next_nodes.push_back(root); // push back the root node
while (next_nodes.size()) {
std::vector <kdtree::kdnode<T>*> next_search_nodes; // next search nodes
while (next_nodes.size()) { // two same WHILE is because we need to make a new vector to store nodes for search
kdtree::kdnode<T> *current_node = next_nodes.back(); // handle node one by one (right first)
next_nodes.pop_back(); // pop out current node in order to store next round of nodes
if (current_node->level < max_levels) {
if (current_node->indices.size() > 1) { // split if the nonleaf node contains more than one point
kdtree::kdnode<T> *left = new kdtree::kdnode<T>();
kdtree::kdnode<T> *right = new kdtree::kdnode<T>();
left->idx = n_id++; // set the index of current node's left node
right->idx = n_id++;
Split(current_node, left, right); // split left and right and determine a node
std::vector <size_t> temp; // empty vecters of int
|
cc09e435
David Mayerich
added Jack's KD-T...
|
105
|
//temp.resize(current_node->indices.size());
|
268037bc
Jiaming Guo
made code more el...
|
106
|
current_node->indices.swap(temp); // clean up current node's indices
|
cc09e435
David Mayerich
added Jack's KD-T...
|
107
108
|
current_node->left = left;
current_node->right = right;
|
268037bc
Jiaming Guo
made code more el...
|
109
110
|
current_node->left_idx = left->idx;
current_node->right_idx = right->idx;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
111
|
if (right->indices.size())
|
268037bc
Jiaming Guo
made code more el...
|
112
113
114
|
next_search_nodes.push_back(right); // left pop out first
if (left->indices.size())
next_search_nodes.push_back(left);
|
cc09e435
David Mayerich
added Jack's KD-T...
|
115
116
117
|
}
}
}
|
268037bc
Jiaming Guo
made code more el...
|
118
|
next_nodes = next_search_nodes; // go deeper within the tree
|
cc09e435
David Mayerich
added Jack's KD-T...
|
119
120
|
}
}
|
268037bc
Jiaming Guo
made code more el...
|
121
122
123
|
static bool SortPoints(const size_t a, const size_t b) { // create functor for std::sort
std::vector <kdtree::point<T, D>> &pts = *cur_tree_ptr->tmp_points; // put cur_tree_ptr to current input points' pointer
return pts[a].dim[cur_tree_ptr->current_axis] < pts[b].dim[cur_tree_ptr->current_axis];
|
cc09e435
David Mayerich
added Jack's KD-T...
|
124
|
}
|
268037bc
Jiaming Guo
made code more el...
|
125
126
127
|
void Split(kdtree::kdnode<T> *cur, kdtree::kdnode<T> *left, kdtree::kdnode<T> *right) {
std::vector <kdtree::point<T, D>> &pts = *tmp_points;
current_axis = cur->level % D; // indicate the judicative dimension or axis
|
cc09e435
David Mayerich
added Jack's KD-T...
|
128
|
std::sort(cur->indices.begin(), cur->indices.end(), SortPoints); // using SortPoints as comparison function to sort the data
|
268037bc
Jiaming Guo
made code more el...
|
129
130
131
|
size_t mid_value = cur->indices[cur->indices.size() / 2]; // odd in the mid_value, even take the floor
cur->split_value = pts[mid_value].dim[current_axis]; // get the parent node
left->parent = cur; // set the parent of the next search nodes to current node
|
cc09e435
David Mayerich
added Jack's KD-T...
|
132
|
right->parent = cur;
|
268037bc
Jiaming Guo
made code more el...
|
133
|
left->level = cur->level + 1; // level + 1
|
cc09e435
David Mayerich
added Jack's KD-T...
|
134
|
right->level = cur->level + 1;
|
268037bc
Jiaming Guo
made code more el...
|
135
136
137
|
left->parent_idx = cur->idx; // set its parent node's index
right->parent_idx = cur->idx;
for (size_t i = 0; i < cur->indices.size(); i++) { // split into left and right half-space one by one
|
cc09e435
David Mayerich
added Jack's KD-T...
|
138
|
size_t idx = cur->indices[i];
|
268037bc
Jiaming Guo
made code more el...
|
139
|
if (pts[idx].dim[current_axis] < cur->split_value)
|
cc09e435
David Mayerich
added Jack's KD-T...
|
140
141
142
143
144
|
left->indices.push_back(idx);
else
right->indices.push_back(idx);
}
}
|
268037bc
Jiaming Guo
made code more el...
|
145
146
147
148
149
150
|
int GetNumNodes() const { // get the total number of nodes
return n_id;
}
kdtree::kdnode<T>* GetRoot() const { // get the root node of tree
return root;
}
|
cc09e435
David Mayerich
added Jack's KD-T...
|
151
152
153
|
}; //end class kdtree
template <typename T, int D>
|
268037bc
Jiaming Guo
made code more el...
|
154
|
cpu_kdtree<T, D>* cpu_kdtree<T, D>::cur_tree_ptr = NULL; // definition of cur_tree_ptr pointer points to the current class
|
cc09e435
David Mayerich
added Jack's KD-T...
|
155
156
|
template <typename T>
|
268037bc
Jiaming Guo
made code more el...
|
157
158
|
struct cuda_kdnode {
int parent, left, right;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
159
|
T split_value;
|
268037bc
Jiaming Guo
made code more el...
|
160
161
162
|
size_t num_index; // number of indices it has
int index; // the beginning index
size_t level;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
163
164
165
166
167
168
169
|
};
template <typename T, int D>
__device__ T Distance(kdtree::point<T, D> &a, kdtree::point<T, D> &b) {
T dist = 0;
for (size_t i = 0; i < D; i++) {
|
268037bc
Jiaming Guo
made code more el...
|
170
|
T d = a.dim[i] - b.dim[i];
|
cc09e435
David Mayerich
added Jack's KD-T...
|
171
172
173
174
175
|
dist += d*d;
}
return dist;
}
template <typename T, int D>
|
268037bc
Jiaming Guo
made code more el...
|
176
177
178
|
__device__ void SearchAtNode(cuda_kdnode<T> *nodes, size_t *indices, kdtree::point<T, D> *d_reference_points, int cur, kdtree::point<T, D> &d_query_point, size_t *d_index, T *d_distance, int *d_node) {
T best_distance = FLT_MAX;
size_t best_index = 0;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
179
|
|
268037bc
Jiaming Guo
made code more el...
|
180
|
while (true) { // break until reach the bottom
|
cc09e435
David Mayerich
added Jack's KD-T...
|
181
|
int split_axis = nodes[cur].level % D;
|
268037bc
Jiaming Guo
made code more el...
|
182
183
184
185
186
187
188
189
|
if (nodes[cur].left == -1) { // check whether it has left node or not
*d_node = cur;
for (int i = 0; i < nodes[cur].num_index; i++) {
size_t idx = indices[nodes[cur].index + i];
T dist = Distance<T, D>(d_query_point, d_reference_points[idx]);
if (dist < best_distance) {
best_distance = dist;
best_index = idx;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
190
191
|
}
}
|
268037bc
Jiaming Guo
made code more el...
|
192
|
break;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
193
|
}
|
268037bc
Jiaming Guo
made code more el...
|
194
|
else if (d_query_point.dim[split_axis] < nodes[cur].split_value) { // jump into specific son node
|
cc09e435
David Mayerich
added Jack's KD-T...
|
195
196
197
198
199
200
|
cur = nodes[cur].left;
}
else {
cur = nodes[cur].right;
}
}
|
268037bc
Jiaming Guo
made code more el...
|
201
202
|
*d_distance = best_distance;
*d_index = best_index;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
203
204
|
}
template <typename T, int D>
|
268037bc
Jiaming Guo
made code more el...
|
205
206
207
208
209
210
211
212
213
214
215
216
217
|
__device__ void SearchAtNodeRange(cuda_kdnode<T> *nodes, size_t *indices, kdtree::point<T, D> *d_reference_points, kdtree::point<T, D> &d_query_point, int cur, T range, size_t *d_index, T *d_distance, size_t id, int *next_nodes, int *next_search_nodes, int *Judge) {
T best_distance = FLT_MAX;
size_t best_index = 0;
int next_nodes_pos = 0; // initialize pop out order index
next_nodes[id * 1000 + next_nodes_pos] = cur; // find data that belongs to the very specific thread
next_nodes_pos++;
while (next_nodes_pos) {
int next_search_nodes_pos = 0; // record push back order index
while (next_nodes_pos) {
cur = next_nodes[id * 1000 + next_nodes_pos - 1]; // pop out the last push in one and keep poping out
next_nodes_pos--;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
218
219
220
|
int split_axis = nodes[cur].level % D;
if (nodes[cur].left == -1) {
|
268037bc
Jiaming Guo
made code more el...
|
221
222
223
224
225
226
|
for (int i = 0; i < nodes[cur].num_index; i++) {
int idx = indices[nodes[cur].index + i]; // all indices are stored in one array, pick up from every node's beginning index
T d = Distance<T>(d_query_point, d_reference_points[idx]);
if (d < best_distance) {
best_distance = d;
best_index = idx;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
227
228
229
230
|
}
}
}
else {
|
268037bc
Jiaming Guo
made code more el...
|
231
|
T d = d_query_point.dim[split_axis] - nodes[cur].split_value;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
232
233
|
if (fabs(d) > range) {
|
268037bc
Jiaming Guo
made code more el...
|
234
235
236
237
238
239
240
241
|
if (d < 0) {
next_search_nodes[id * 1000 + next_search_nodes_pos] = nodes[cur].left;
next_search_nodes_pos++;
}
else {
next_search_nodes[id * 1000 + next_search_nodes_pos] = nodes[cur].right;
next_search_nodes_pos++;
}
|
cc09e435
David Mayerich
added Jack's KD-T...
|
242
243
|
}
else {
|
268037bc
Jiaming Guo
made code more el...
|
244
245
246
247
248
249
250
251
|
next_search_nodes[id * 1000 + next_search_nodes_pos] = nodes[cur].right;
next_search_nodes_pos++;
next_search_nodes[id * 1000 + next_search_nodes_pos] = nodes[cur].left;
next_search_nodes_pos++;
if (next_search_nodes_pos > 1000) {
printf("Thread conflict might be caused by thread %d, so please try smaller input max_tree_levels\n", id);
(*Judge)++;
}
|
cc09e435
David Mayerich
added Jack's KD-T...
|
252
253
254
|
}
}
}
|
268037bc
Jiaming Guo
made code more el...
|
255
256
257
|
for (int i = 0; i < next_search_nodes_pos; i++)
next_nodes[id * 1000 + i] = next_search_nodes[id * 1000 + i];
next_nodes_pos = next_search_nodes_pos;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
258
|
}
|
268037bc
Jiaming Guo
made code more el...
|
259
260
|
*d_distance = best_distance;
*d_index = best_index;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
261
262
|
}
template <typename T, int D>
|
268037bc
Jiaming Guo
made code more el...
|
263
|
__device__ void Search(cuda_kdnode<T> *nodes, size_t *indices, kdtree::point<T, D> *d_reference_points, kdtree::point<T, D> &d_query_point, size_t *d_index, T *d_distance, size_t id, int *next_nodes, int *next_search_nodes, int *Judge) {
|
cc09e435
David Mayerich
added Jack's KD-T...
|
264
|
int best_node = 0;
|
268037bc
Jiaming Guo
made code more el...
|
265
266
|
T best_distance = FLT_MAX;
size_t best_index = 0;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
267
|
T radius = 0;
|
268037bc
Jiaming Guo
made code more el...
|
268
269
270
|
SearchAtNode<T, D>(nodes, indices, d_reference_points, 0, d_query_point, &best_index, &best_distance, &best_node);
radius = sqrt(best_distance); // get range
|
cc09e435
David Mayerich
added Jack's KD-T...
|
271
272
273
|
int cur = best_node;
while (nodes[cur].parent != -1) {
|
cc09e435
David Mayerich
added Jack's KD-T...
|
274
275
|
int parent = nodes[cur].parent;
int split_axis = nodes[parent].level % D;
|
268037bc
Jiaming Guo
made code more el...
|
276
|
|
cc09e435
David Mayerich
added Jack's KD-T...
|
277
278
|
T tmp_dist = FLT_MAX;
size_t tmp_idx;
|
268037bc
Jiaming Guo
made code more el...
|
279
|
if (fabs(nodes[parent].split_value - d_query_point.dim[split_axis]) <= radius) {
|
cc09e435
David Mayerich
added Jack's KD-T...
|
280
|
if (nodes[parent].left != cur)
|
268037bc
Jiaming Guo
made code more el...
|
281
|
SearchAtNodeRange(nodes, indices, d_reference_points, d_query_point, nodes[parent].left, radius, &tmp_idx, &tmp_dist, id, next_nodes, next_search_nodes, Judge);
|
cc09e435
David Mayerich
added Jack's KD-T...
|
282
|
else
|
268037bc
Jiaming Guo
made code more el...
|
283
|
SearchAtNodeRange(nodes, indices, d_reference_points, d_query_point, nodes[parent].right, radius, &tmp_idx, &tmp_dist, id, next_nodes, next_search_nodes, Judge);
|
cc09e435
David Mayerich
added Jack's KD-T...
|
284
|
}
|
268037bc
Jiaming Guo
made code more el...
|
285
286
287
|
if (tmp_dist < best_distance) {
best_distance = tmp_dist;
best_index = tmp_idx;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
288
289
290
|
}
cur = parent;
}
|
268037bc
Jiaming Guo
made code more el...
|
291
292
|
*d_distance = sqrt(best_distance);
*d_index = best_index;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
293
294
|
}
template <typename T, int D>
|
268037bc
Jiaming Guo
made code more el...
|
295
296
297
|
__global__ void SearchBatch(cuda_kdnode<T> *nodes, size_t *indices, kdtree::point<T, D> *d_reference_points, kdtree::point<T, D> *d_query_points, size_t d_query_count, size_t *d_indices, T *d_distances, int *next_nodes, int *next_search_nodes, int *Judge) {
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= d_query_count) return; // avoid segfault
|
cc09e435
David Mayerich
added Jack's KD-T...
|
298
|
|
268037bc
Jiaming Guo
made code more el...
|
299
|
Search<T, D>(nodes, indices, d_reference_points, d_query_points[idx], &d_indices[idx], &d_distances[idx], idx, next_nodes, next_search_nodes, Judge); // every query points are independent
|
cc09e435
David Mayerich
added Jack's KD-T...
|
300
301
302
303
304
|
}
template <typename T, int D = 3>
class cuda_kdtree {
protected:
|
268037bc
Jiaming Guo
made code more el...
|
305
306
307
308
|
cuda_kdnode<T> *d_nodes;
size_t *d_index;
kdtree::point<T, D>* d_reference_points;
size_t d_reference_count;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
309
310
|
public:
~cuda_kdtree() {
|
268037bc
Jiaming Guo
made code more el...
|
311
312
313
|
HANDLE_ERROR(cudaFree(d_nodes));
HANDLE_ERROR(cudaFree(d_index));
HANDLE_ERROR(cudaFree(d_reference_points));
|
cc09e435
David Mayerich
added Jack's KD-T...
|
314
|
}
|
268037bc
Jiaming Guo
made code more el...
|
315
316
317
318
319
320
321
322
323
324
325
326
|
void CreateKDTree(T *h_reference_points, size_t reference_count, size_t dim_count, size_t max_levels) {
if (max_levels > 10) {
std::cout<<"The max_tree_levels should be smaller!"<<std::endl;
exit(1);
}
std::vector <kdtree::point<T, D>> reference_points(reference_count); // restore the reference points in particular way
for (size_t j = 0; j < reference_count; j++)
for (size_t i = 0; i < dim_count; i++)
reference_points[j].dim[i] = h_reference_points[j * dim_count + i];
cpu_kdtree<T, D> tree; // creating a tree on cpu
tree.Create(reference_points, max_levels); // building a tree on cpu
kdtree::kdnode<T> *d_root = tree.GetRoot();
|
cc09e435
David Mayerich
added Jack's KD-T...
|
327
|
int num_nodes = tree.GetNumNodes();
|
268037bc
Jiaming Guo
made code more el...
|
328
|
d_reference_count = reference_points.size(); // also equals to reference_count
|
cc09e435
David Mayerich
added Jack's KD-T...
|
329
|
|
268037bc
Jiaming Guo
made code more el...
|
330
331
332
|
HANDLE_ERROR(cudaMalloc((void**)&d_nodes, sizeof(cuda_kdnode<T>) * num_nodes)); // copy data from host to device
HANDLE_ERROR(cudaMalloc((void**)&d_index, sizeof(size_t) * d_reference_count));
HANDLE_ERROR(cudaMalloc((void**)&d_reference_points, sizeof(kdtree::point<T, D>) * d_reference_count));
|
cc09e435
David Mayerich
added Jack's KD-T...
|
333
|
|
268037bc
Jiaming Guo
made code more el...
|
334
335
336
|
std::vector <cuda_kdnode<T>> tmp_nodes(num_nodes);
std::vector <size_t> indices(d_reference_count);
std::vector <kdtree::kdnode<T>*> next_nodes;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
337
|
size_t cur_pos = 0;
|
268037bc
Jiaming Guo
made code more el...
|
338
339
340
341
342
343
344
345
346
347
348
349
350
|
next_nodes.push_back(d_root);
while (next_nodes.size()) {
std::vector <typename kdtree::kdnode<T>*> next_search_nodes;
while (next_nodes.size()) {
kdtree::kdnode<T> *cur = next_nodes.back();
next_nodes.pop_back();
int id = cur->idx; // the nodes at same level are independent
tmp_nodes[id].level = cur->level;
tmp_nodes[id].parent = cur->parent_idx;
tmp_nodes[id].left = cur->left_idx;
tmp_nodes[id].right = cur->right_idx;
tmp_nodes[id].split_value = cur->split_value;
tmp_nodes[id].num_index = cur->indices.size(); // number of index
|
cc09e435
David Mayerich
added Jack's KD-T...
|
351
352
353
354
|
if (cur->indices.size()) {
for (size_t i = 0; i < cur->indices.size(); i++)
indices[cur_pos + i] = cur->indices[i];
|
268037bc
Jiaming Guo
made code more el...
|
355
356
|
tmp_nodes[id].index = (int)cur_pos; // beginning index of reference_points that every bottom node has
cur_pos += cur->indices.size(); // store indices continuously for every query_point
|
cc09e435
David Mayerich
added Jack's KD-T...
|
357
358
|
}
else {
|
268037bc
Jiaming Guo
made code more el...
|
359
|
tmp_nodes[id].index = -1;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
360
361
362
|
}
if (cur->left)
|
268037bc
Jiaming Guo
made code more el...
|
363
|
next_search_nodes.push_back(cur->left);
|
cc09e435
David Mayerich
added Jack's KD-T...
|
364
365
|
if (cur->right)
|
268037bc
Jiaming Guo
made code more el...
|
366
|
next_search_nodes.push_back(cur->right);
|
cc09e435
David Mayerich
added Jack's KD-T...
|
367
|
}
|
268037bc
Jiaming Guo
made code more el...
|
368
|
next_nodes = next_search_nodes;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
369
|
}
|
268037bc
Jiaming Guo
made code more el...
|
370
371
372
|
HANDLE_ERROR(cudaMemcpy(d_nodes, &tmp_nodes[0], sizeof(cuda_kdnode<T>) * tmp_nodes.size(), cudaMemcpyHostToDevice));
HANDLE_ERROR(cudaMemcpy(d_index, &indices[0], sizeof(size_t) * indices.size(), cudaMemcpyHostToDevice));
HANDLE_ERROR(cudaMemcpy(d_reference_points, &reference_points[0], sizeof(kdtree::point<T, D>) * reference_points.size(), cudaMemcpyHostToDevice));
|
cc09e435
David Mayerich
added Jack's KD-T...
|
373
|
}
|
268037bc
Jiaming Guo
made code more el...
|
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
|
void Search(T *h_query_points, size_t query_count, size_t dim_count, T *dists, size_t *indices) {
std::vector <kdtree::point<T, D>> query_points(query_count);
for (size_t j = 0; j < query_count; j++)
for (size_t i = 0; i < dim_count; i++)
query_points[j].dim[i] = h_query_points[j * dim_count + i];
unsigned int threads = (unsigned int)(query_points.size() > 1024 ? 1024 : query_points.size());
unsigned int blocks = (unsigned int)(query_points.size() / threads + (query_points.size() % threads ? 1 : 0));
kdtree::point<T, D> *d_query_points; // create a pointer pointing to query points on gpu
size_t *d_indices;
T *d_distances;
int *next_nodes; // create two STACK-like array
int *next_search_nodes;
int *Judge = NULL; // judge variable to see whether one thread is overwrite another thread's memory
HANDLE_ERROR(cudaMalloc((void**)&d_query_points, sizeof(T) * query_points.size() * D));
HANDLE_ERROR(cudaMalloc((void**)&d_indices, sizeof(size_t) * query_points.size()));
HANDLE_ERROR(cudaMalloc((void**)&d_distances, sizeof(T) * query_points.size()));
HANDLE_ERROR(cudaMalloc((void**)&next_nodes, threads * blocks * 1000 * sizeof(int))); // STACK size right now is 1000, you can change it if you mean to
HANDLE_ERROR(cudaMalloc((void**)&next_search_nodes, threads * blocks * 1000 * sizeof(int)));
HANDLE_ERROR(cudaMemcpy(d_query_points, &query_points[0], sizeof(T) * query_points.size() * D, cudaMemcpyHostToDevice));
SearchBatch<<<threads, blocks>>> (d_nodes, d_index, d_reference_points, d_query_points, query_points.size(), d_indices, d_distances, next_nodes, next_search_nodes, Judge);
if (Judge == NULL) { // do the following work if the thread works safely
HANDLE_ERROR(cudaMemcpy(indices, d_indices, sizeof(size_t) * query_points.size(), cudaMemcpyDeviceToHost));
HANDLE_ERROR(cudaMemcpy(dists, d_distances, sizeof(T) * query_points.size(), cudaMemcpyDeviceToHost));
}
|
cc09e435
David Mayerich
added Jack's KD-T...
|
405
|
|
268037bc
Jiaming Guo
made code more el...
|
406
407
408
409
410
|
HANDLE_ERROR(cudaFree(next_nodes));
HANDLE_ERROR(cudaFree(next_search_nodes));
HANDLE_ERROR(cudaFree(d_query_points));
HANDLE_ERROR(cudaFree(d_indices));
HANDLE_ERROR(cudaFree(d_distances));
|
cc09e435
David Mayerich
added Jack's KD-T...
|
411
412
413
414
|
}
};
} //end namespace stim
#endif
|