kdtree.cuh
18.4 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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
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
405
406
407
408
409
410
411
412
413
414
// right now the size of CUDA STACK is set to 1000, increase it if you mean to make deeper tree
// 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>
#include <stim/cuda/cudatools/error.h>
namespace stim {
namespace kdtree {
template<typename T, int D> // typename refers to float or double while D refers to dimension of points
struct point {
T dim[D]; // create a structure to store every one input point
};
template<typename T>
class kdnode {
public:
kdnode() { // constructor for initializing a kdnode
parent = NULL; // set every node's parent, left and right kdnode pointers to NULL
left = NULL;
right = NULL;
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
}
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
};
} // end of namespace kdtree
template <typename T, int D = 3> // set dimension of data to default 3
class cpu_kdtree {
protected:
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;
public:
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
}
~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();
if (cur->left)
next_search_nodes.push_back(cur->left);
if (cur->right)
next_search_nodes.push_back(cur->right);
delete cur;
}
next_nodes = next_search_nodes;
}
root = NULL;
}
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
}
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
//temp.resize(current_node->indices.size());
current_node->indices.swap(temp); // clean up current node's indices
current_node->left = left;
current_node->right = right;
current_node->left_idx = left->idx;
current_node->right_idx = right->idx;
if (right->indices.size())
next_search_nodes.push_back(right); // left pop out first
if (left->indices.size())
next_search_nodes.push_back(left);
}
}
}
next_nodes = next_search_nodes; // go deeper within the tree
}
}
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];
}
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
std::sort(cur->indices.begin(), cur->indices.end(), SortPoints); // using SortPoints as comparison function to sort the data
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
right->parent = cur;
left->level = cur->level + 1; // level + 1
right->level = cur->level + 1;
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
size_t idx = cur->indices[i];
if (pts[idx].dim[current_axis] < cur->split_value)
left->indices.push_back(idx);
else
right->indices.push_back(idx);
}
}
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;
}
}; //end class kdtree
template <typename T, int D>
cpu_kdtree<T, D>* cpu_kdtree<T, D>::cur_tree_ptr = NULL; // definition of cur_tree_ptr pointer points to the current class
template <typename T>
struct cuda_kdnode {
int parent, left, right;
T split_value;
size_t num_index; // number of indices it has
int index; // the beginning index
size_t level;
};
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++) {
T d = a.dim[i] - b.dim[i];
dist += d*d;
}
return dist;
}
template <typename T, int D>
__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;
while (true) { // break until reach the bottom
int split_axis = nodes[cur].level % D;
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;
}
}
break;
}
else if (d_query_point.dim[split_axis] < nodes[cur].split_value) { // jump into specific son node
cur = nodes[cur].left;
}
else {
cur = nodes[cur].right;
}
}
*d_distance = best_distance;
*d_index = best_index;
}
template <typename T, int D>
__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--;
int split_axis = nodes[cur].level % D;
if (nodes[cur].left == -1) {
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;
}
}
}
else {
T d = d_query_point.dim[split_axis] - nodes[cur].split_value;
if (fabs(d) > range) {
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++;
}
}
else {
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)++;
}
}
}
}
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;
}
*d_distance = best_distance;
*d_index = best_index;
}
template <typename T, int D>
__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) {
int best_node = 0;
T best_distance = FLT_MAX;
size_t best_index = 0;
T radius = 0;
SearchAtNode<T, D>(nodes, indices, d_reference_points, 0, d_query_point, &best_index, &best_distance, &best_node);
radius = sqrt(best_distance); // get range
int cur = best_node;
while (nodes[cur].parent != -1) {
int parent = nodes[cur].parent;
int split_axis = nodes[parent].level % D;
T tmp_dist = FLT_MAX;
size_t tmp_idx;
if (fabs(nodes[parent].split_value - d_query_point.dim[split_axis]) <= radius) {
if (nodes[parent].left != cur)
SearchAtNodeRange(nodes, indices, d_reference_points, d_query_point, nodes[parent].left, radius, &tmp_idx, &tmp_dist, id, next_nodes, next_search_nodes, Judge);
else
SearchAtNodeRange(nodes, indices, d_reference_points, d_query_point, nodes[parent].right, radius, &tmp_idx, &tmp_dist, id, next_nodes, next_search_nodes, Judge);
}
if (tmp_dist < best_distance) {
best_distance = tmp_dist;
best_index = tmp_idx;
}
cur = parent;
}
*d_distance = sqrt(best_distance);
*d_index = best_index;
}
template <typename T, int D>
__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
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
}
template <typename T, int D = 3>
class cuda_kdtree {
protected:
cuda_kdnode<T> *d_nodes;
size_t *d_index;
kdtree::point<T, D>* d_reference_points;
size_t d_reference_count;
public:
~cuda_kdtree() {
HANDLE_ERROR(cudaFree(d_nodes));
HANDLE_ERROR(cudaFree(d_index));
HANDLE_ERROR(cudaFree(d_reference_points));
}
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();
int num_nodes = tree.GetNumNodes();
d_reference_count = reference_points.size(); // also equals to reference_count
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));
std::vector <cuda_kdnode<T>> tmp_nodes(num_nodes);
std::vector <size_t> indices(d_reference_count);
std::vector <kdtree::kdnode<T>*> next_nodes;
size_t cur_pos = 0;
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
if (cur->indices.size()) {
for (size_t i = 0; i < cur->indices.size(); i++)
indices[cur_pos + i] = cur->indices[i];
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
}
else {
tmp_nodes[id].index = -1;
}
if (cur->left)
next_search_nodes.push_back(cur->left);
if (cur->right)
next_search_nodes.push_back(cur->right);
}
next_nodes = next_search_nodes;
}
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));
}
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));
}
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));
}
};
} //end namespace stim
#endif