41100e51
Jiaming Guo
add cpu kdtree ve...
|
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
|
// data should be stored in row-major
// x1,x2,x3,x4,x5......
// y1,y2,y3,y4,y5......
// ....................
// ....................
#ifndef KDTREE_H
#define KDTREE_H
|
a0f09cd5
Jiaming Guo
add #define stack...
|
10
|
#define stack_size 50
|
cc09e435
David Mayerich
added Jack's KD-T...
|
11
12
13
14
15
16
|
#include "device_launch_parameters.h"
#include <cuda.h>
#include <cuda_runtime_api.h>
#include "cuda_runtime.h"
#include <vector>
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
17
|
#include <cstring>
|
cc09e435
David Mayerich
added Jack's KD-T...
|
18
19
20
|
#include <float.h>
#include <iostream>
#include <algorithm>
|
268037bc
Jiaming Guo
made code more el...
|
21
|
#include <stim/cuda/cudatools/error.h>
|
87d2ffe6
Jiaming Guo
add aabb to the k...
|
22
|
#include <stim/visualization/aabbn.h>
|
cc09e435
David Mayerich
added Jack's KD-T...
|
23
24
25
|
namespace stim {
namespace kdtree {
|
268037bc
Jiaming Guo
made code more el...
|
26
|
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...
|
27
|
struct point {
|
268037bc
Jiaming Guo
made code more el...
|
28
|
T dim[D]; // create a structure to store every one input point
|
cc09e435
David Mayerich
added Jack's KD-T...
|
29
30
31
|
};
template<typename T>
|
268037bc
Jiaming Guo
made code more el...
|
32
|
class kdnode {
|
cc09e435
David Mayerich
added Jack's KD-T...
|
33
|
public:
|
268037bc
Jiaming Guo
made code more el...
|
34
35
|
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...
|
36
37
|
left = NULL;
right = NULL;
|
268037bc
Jiaming Guo
made code more el...
|
38
39
40
41
|
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...
|
42
|
}
|
268037bc
Jiaming Guo
made code more el...
|
43
44
45
46
47
48
|
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...
|
49
|
};
|
268037bc
Jiaming Guo
made code more el...
|
50
|
} // end of namespace kdtree
|
cc09e435
David Mayerich
added Jack's KD-T...
|
51
|
|
268037bc
Jiaming Guo
made code more el...
|
52
|
template <typename T, int D = 3> // set dimension of data to default 3
|
cc09e435
David Mayerich
added Jack's KD-T...
|
53
|
class cpu_kdtree {
|
cc09e435
David Mayerich
added Jack's KD-T...
|
54
|
protected:
|
268037bc
Jiaming Guo
made code more el...
|
55
|
int current_axis; // current judging axis
|
268037bc
Jiaming Guo
made code more el...
|
56
|
int n_id; // store the total number of nodes
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
57
58
|
std::vector < typename kdtree::point<T, D> > *tmp_points; // transfer or temperary points
std::vector < typename kdtree::point<T, D> > cpu_tmp_points; // for cpu searching
|
268037bc
Jiaming Guo
made code more el...
|
59
60
|
kdtree::kdnode<T> *root; // root node
static cpu_kdtree<T, D> *cur_tree_ptr;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
61
|
public:
|
268037bc
Jiaming Guo
made code more el...
|
62
|
cpu_kdtree() { // constructor for creating a cpu_kdtree
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
63
|
cur_tree_ptr = this; // create a class pointer points to the current class value
|
268037bc
Jiaming Guo
made code more el...
|
64
|
n_id = 0; // set total number of points to default 0
|
cc09e435
David Mayerich
added Jack's KD-T...
|
65
|
}
|
268037bc
Jiaming Guo
made code more el...
|
66
67
68
69
70
71
72
73
|
~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...
|
74
|
if (cur->left)
|
268037bc
Jiaming Guo
made code more el...
|
75
|
next_search_nodes.push_back(cur->left);
|
cc09e435
David Mayerich
added Jack's KD-T...
|
76
|
if (cur->right)
|
268037bc
Jiaming Guo
made code more el...
|
77
|
next_search_nodes.push_back(cur->right);
|
cc09e435
David Mayerich
added Jack's KD-T...
|
78
79
|
delete cur;
}
|
268037bc
Jiaming Guo
made code more el...
|
80
|
next_nodes = next_search_nodes;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
81
|
}
|
268037bc
Jiaming Guo
made code more el...
|
82
|
root = NULL;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
83
|
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
84
85
|
void cpu_create(std::vector < typename kdtree::point<T, D> > &reference_points, size_t max_levels) {
tmp_points = &reference_points;
|
268037bc
Jiaming Guo
made code more el...
|
86
87
88
89
90
91
|
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...
|
92
|
}
|
268037bc
Jiaming Guo
made code more el...
|
93
94
95
96
97
98
99
100
101
102
103
104
105
|
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++;
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
106
|
split(current_node, left, right); // split left and right and determine a node
|
268037bc
Jiaming Guo
made code more el...
|
107
|
std::vector <size_t> temp; // empty vecters of int
|
cc09e435
David Mayerich
added Jack's KD-T...
|
108
|
//temp.resize(current_node->indices.size());
|
268037bc
Jiaming Guo
made code more el...
|
109
|
current_node->indices.swap(temp); // clean up current node's indices
|
cc09e435
David Mayerich
added Jack's KD-T...
|
110
111
|
current_node->left = left;
current_node->right = right;
|
268037bc
Jiaming Guo
made code more el...
|
112
113
|
current_node->left_idx = left->idx;
current_node->right_idx = right->idx;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
114
|
if (right->indices.size())
|
268037bc
Jiaming Guo
made code more el...
|
115
116
117
|
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...
|
118
119
120
|
}
}
}
|
268037bc
Jiaming Guo
made code more el...
|
121
|
next_nodes = next_search_nodes; // go deeper within the tree
|
cc09e435
David Mayerich
added Jack's KD-T...
|
122
123
|
}
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
124
|
static bool sort_points(const size_t a, const size_t b) { // create functor for std::sort
|
e8c9a82b
David Mayerich
fixed GCC errors ...
|
125
|
std::vector < typename kdtree::point<T, D> > &pts = *cur_tree_ptr->tmp_points; // put cur_tree_ptr to current input points' pointer
|
268037bc
Jiaming Guo
made code more el...
|
126
|
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...
|
127
|
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
128
|
void split(kdtree::kdnode<T> *cur, kdtree::kdnode<T> *left, kdtree::kdnode<T> *right) {
|
e8c9a82b
David Mayerich
fixed GCC errors ...
|
129
|
std::vector < typename kdtree::point<T, D> > &pts = *tmp_points;
|
268037bc
Jiaming Guo
made code more el...
|
130
|
current_axis = cur->level % D; // indicate the judicative dimension or axis
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
131
|
std::sort(cur->indices.begin(), cur->indices.end(), sort_points); // using SortPoints as comparison function to sort the data
|
268037bc
Jiaming Guo
made code more el...
|
132
133
134
|
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...
|
135
|
right->parent = cur;
|
268037bc
Jiaming Guo
made code more el...
|
136
|
left->level = cur->level + 1; // level + 1
|
cc09e435
David Mayerich
added Jack's KD-T...
|
137
|
right->level = cur->level + 1;
|
268037bc
Jiaming Guo
made code more el...
|
138
139
140
|
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...
|
141
|
size_t idx = cur->indices[i];
|
268037bc
Jiaming Guo
made code more el...
|
142
|
if (pts[idx].dim[current_axis] < cur->split_value)
|
cc09e435
David Mayerich
added Jack's KD-T...
|
143
144
145
146
147
|
left->indices.push_back(idx);
else
right->indices.push_back(idx);
}
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
148
149
150
151
152
153
154
155
156
|
void create(T *h_reference_points, size_t reference_count, size_t max_levels) {
std::vector < typename 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 < D; i++)
reference_points[j].dim[i] = h_reference_points[j * D + i];
cpu_create(reference_points, max_levels);
cpu_tmp_points = *tmp_points;
}
int get_num_nodes() const { // get the total number of nodes
|
268037bc
Jiaming Guo
made code more el...
|
157
158
|
return n_id;
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
159
|
kdtree::kdnode<T>* get_root() const { // get the root node of tree
|
268037bc
Jiaming Guo
made code more el...
|
160
161
|
return root;
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
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
|
T cpu_distance(const kdtree::point<T, D> &a, const kdtree::point<T, D> &b) {
T distance = 0;
for (size_t i = 0; i < D; i++) {
T d = a.dim[i] - b.dim[i];
distance += d*d;
}
return distance;
}
void cpu_search_at_node(kdtree::kdnode<T> *cur, const kdtree::point<T, D> &query, size_t *index, T *distance, kdtree::kdnode<T> **node) {
T best_distance = FLT_MAX; // initialize the best distance to max of floating point
size_t best_index = 0;
std::vector < typename kdtree::point<T, D> > pts = cpu_tmp_points;
while (true) {
size_t split_axis = cur->level % D;
if (cur->left == NULL) { // risky but acceptable, same goes for right because left and right are in same pace
*node = cur; // pointer points to a pointer
for (size_t i = 0; i < cur->indices.size(); i++) {
size_t idx = cur->indices[i];
T d = cpu_distance(query, pts[idx]); // compute distances
/// if we want to compute k nearest neighbor, we can input the last resul
/// (last_best_dist < dist < best_dist) to select the next point until reaching to k
if (d < best_distance) {
best_distance = d;
best_index = idx; // record the nearest neighbor index
}
}
break; // find the target point then break the loop
}
else if (query.dim[split_axis] < cur->split_value) { // if it has son node, visit the next node on either left side or right side
cur = cur->left;
}
else {
cur = cur->right;
}
}
*index = best_index;
*distance = best_distance;
}
void cpu_search_at_node_range(kdtree::kdnode<T> *cur, const kdtree::point<T, D> &query, T range, size_t *index, T *distance) {
T best_distance = FLT_MAX; // initialize the best distance to max of floating point
size_t best_index = 0;
std::vector < typename kdtree::point<T, D> > pts = cpu_tmp_points;
std::vector < typename kdtree::kdnode<T>*> next_node;
next_node.push_back(cur);
while (next_node.size()) {
std::vector<typename kdtree::kdnode<T>*> next_search;
while (next_node.size()) {
cur = next_node.back();
next_node.pop_back();
size_t split_axis = cur->level % D;
if (cur->left == NULL) {
for (size_t i = 0; i < cur->indices.size(); i++) {
size_t idx = cur->indices[i];
T d = cpu_distance(query, pts[idx]);
if (d < best_distance) {
best_distance = d;
best_index = idx;
}
}
}
else {
T d = query.dim[split_axis] - cur->split_value; // computer distance along specific axis or dimension
/// there are three possibilities: on either left or right, and on both left and right
if (fabs(d) > range) { // absolute value of floating point to see if distance will be larger that best_dist
if (d < 0)
next_search.push_back(cur->left); // every left[split_axis] is less and equal to cur->split_value, so it is possible to find the nearest point in this region
else
next_search.push_back(cur->right);
}
else { // it is possible that nereast neighbor will appear on both left and right
next_search.push_back(cur->left);
next_search.push_back(cur->right);
}
}
}
next_node = next_search; // pop out at least one time
}
*index = best_index;
*distance = best_distance;
}
void cpu_search(T *h_query_points, size_t query_count, size_t *h_indices, T *h_distances) {
/// first convert the input query point into specific type
kdtree::point<T, D> query;
for (size_t j = 0; j < query_count; j++) {
for (size_t i = 0; i < D; i++)
query.dim[i] = h_query_points[j * D + i];
/// find the nearest node, this will be the upper bound for the next time searching
kdtree::kdnode<T> *best_node = NULL;
T best_distance = FLT_MAX;
size_t best_index = 0;
T radius = 0; // radius for range
cpu_search_at_node(root, query, &best_index, &best_distance, &best_node); // simple search to rougly determine a result for next search step
radius = sqrt(best_distance); // It is possible that nearest will appear in another region
/// find other possibilities
kdtree::kdnode<T> *cur = best_node;
while (cur->parent != NULL) { // every node that you pass will be possible to be the best node
/// go up
kdtree::kdnode<T> *parent = cur->parent; // travel back to every node that we pass through
size_t split_axis = (parent->level) % D;
/// search other nodes
size_t tmp_index;
T tmp_distance = FLT_MAX;
if (fabs(parent->split_value - query.dim[split_axis]) <= radius) {
/// search opposite node
if (parent->left != cur)
cpu_search_at_node_range(parent->left, query, radius, &tmp_index, &tmp_distance); // to see whether it is its mother node's left son node
else
cpu_search_at_node_range(parent->right, query, radius, &tmp_index, &tmp_distance);
}
if (tmp_distance < best_distance) {
best_distance = tmp_distance;
best_index = tmp_index;
}
cur = parent;
}
h_indices[j] = best_index;
h_distances[j] = best_distance;
}
}
|
cc09e435
David Mayerich
added Jack's KD-T...
|
282
283
284
|
}; //end class kdtree
template <typename T, int D>
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
285
|
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...
|
286
287
|
template <typename T>
|
268037bc
Jiaming Guo
made code more el...
|
288
289
|
struct cuda_kdnode {
int parent, left, right;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
290
|
T split_value;
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
291
292
|
size_t num_index; // number of indices it has
int index; // the beginning index
|
268037bc
Jiaming Guo
made code more el...
|
293
|
size_t level;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
294
295
296
|
};
template <typename T, int D>
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
297
298
|
__device__ T gpu_distance(kdtree::point<T, D> &a, kdtree::point<T, D> &b) {
T distance = 0;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
299
300
|
for (size_t i = 0; i < D; i++) {
|
268037bc
Jiaming Guo
made code more el...
|
301
|
T d = a.dim[i] - b.dim[i];
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
302
|
distance += d*d;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
303
|
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
304
|
return distance;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
305
306
|
}
template <typename T, int D>
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
307
|
__device__ void search_at_node(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) {
|
268037bc
Jiaming Guo
made code more el...
|
308
309
|
T best_distance = FLT_MAX;
size_t best_index = 0;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
310
|
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
311
|
while (true) { // break until reach the bottom
|
cc09e435
David Mayerich
added Jack's KD-T...
|
312
|
int split_axis = nodes[cur].level % D;
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
313
|
if (nodes[cur].left == -1) { // check whether it has left node or not
|
268037bc
Jiaming Guo
made code more el...
|
314
315
316
|
*d_node = cur;
for (int i = 0; i < nodes[cur].num_index; i++) {
size_t idx = indices[nodes[cur].index + i];
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
317
|
T dist = gpu_distance<T, D>(d_query_point, d_reference_points[idx]);
|
268037bc
Jiaming Guo
made code more el...
|
318
319
320
|
if (dist < best_distance) {
best_distance = dist;
best_index = idx;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
321
322
|
}
}
|
268037bc
Jiaming Guo
made code more el...
|
323
|
break;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
324
|
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
325
|
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...
|
326
327
328
329
330
331
|
cur = nodes[cur].left;
}
else {
cur = nodes[cur].right;
}
}
|
268037bc
Jiaming Guo
made code more el...
|
332
333
|
*d_distance = best_distance;
*d_index = best_index;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
334
335
|
}
template <typename T, int D>
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
336
|
__device__ void search_at_node_range(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) {
|
268037bc
Jiaming Guo
made code more el...
|
337
338
339
|
T best_distance = FLT_MAX;
size_t best_index = 0;
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
340
|
int next_nodes_pos = 0; // initialize pop out order index
|
a0f09cd5
Jiaming Guo
add #define stack...
|
341
|
next_nodes[id * stack_size + next_nodes_pos] = cur; // find data that belongs to the very specific thread
|
268037bc
Jiaming Guo
made code more el...
|
342
343
344
|
next_nodes_pos++;
while (next_nodes_pos) {
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
345
|
int next_search_nodes_pos = 0; // record push back order index
|
268037bc
Jiaming Guo
made code more el...
|
346
|
while (next_nodes_pos) {
|
a0f09cd5
Jiaming Guo
add #define stack...
|
347
|
cur = next_nodes[id * stack_size + next_nodes_pos - 1]; // pop out the last push in one and keep poping out
|
268037bc
Jiaming Guo
made code more el...
|
348
|
next_nodes_pos--;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
349
350
351
|
int split_axis = nodes[cur].level % D;
if (nodes[cur].left == -1) {
|
268037bc
Jiaming Guo
made code more el...
|
352
|
for (int i = 0; i < nodes[cur].num_index; i++) {
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
353
354
|
int idx = indices[nodes[cur].index + i]; // all indices are stored in one array, pick up from every node's beginning index
T d = gpu_distance<T>(d_query_point, d_reference_points[idx]);
|
268037bc
Jiaming Guo
made code more el...
|
355
356
357
|
if (d < best_distance) {
best_distance = d;
best_index = idx;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
358
359
360
361
|
}
}
}
else {
|
268037bc
Jiaming Guo
made code more el...
|
362
|
T d = d_query_point.dim[split_axis] - nodes[cur].split_value;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
363
364
|
if (fabs(d) > range) {
|
268037bc
Jiaming Guo
made code more el...
|
365
|
if (d < 0) {
|
a0f09cd5
Jiaming Guo
add #define stack...
|
366
|
next_search_nodes[id * stack_size + next_search_nodes_pos] = nodes[cur].left;
|
268037bc
Jiaming Guo
made code more el...
|
367
368
369
|
next_search_nodes_pos++;
}
else {
|
a0f09cd5
Jiaming Guo
add #define stack...
|
370
|
next_search_nodes[id * stack_size + next_search_nodes_pos] = nodes[cur].right;
|
268037bc
Jiaming Guo
made code more el...
|
371
372
|
next_search_nodes_pos++;
}
|
cc09e435
David Mayerich
added Jack's KD-T...
|
373
374
|
}
else {
|
a0f09cd5
Jiaming Guo
add #define stack...
|
375
|
next_search_nodes[id * stack_size + next_search_nodes_pos] = nodes[cur].right;
|
268037bc
Jiaming Guo
made code more el...
|
376
|
next_search_nodes_pos++;
|
a0f09cd5
Jiaming Guo
add #define stack...
|
377
|
next_search_nodes[id * stack_size + next_search_nodes_pos] = nodes[cur].left;
|
268037bc
Jiaming Guo
made code more el...
|
378
|
next_search_nodes_pos++;
|
a0f09cd5
Jiaming Guo
add #define stack...
|
379
|
if (next_search_nodes_pos > stack_size) {
|
268037bc
Jiaming Guo
made code more el...
|
380
381
382
|
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...
|
383
384
385
|
}
}
}
|
268037bc
Jiaming Guo
made code more el...
|
386
|
for (int i = 0; i < next_search_nodes_pos; i++)
|
a0f09cd5
Jiaming Guo
add #define stack...
|
387
|
next_nodes[id * stack_size + i] = next_search_nodes[id * stack_size + i];
|
268037bc
Jiaming Guo
made code more el...
|
388
|
next_nodes_pos = next_search_nodes_pos;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
389
|
}
|
268037bc
Jiaming Guo
made code more el...
|
390
391
|
*d_distance = best_distance;
*d_index = best_index;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
392
393
|
}
template <typename T, int D>
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
394
|
__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...
|
395
|
int best_node = 0;
|
268037bc
Jiaming Guo
made code more el...
|
396
397
|
T best_distance = FLT_MAX;
size_t best_index = 0;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
398
|
T radius = 0;
|
268037bc
Jiaming Guo
made code more el...
|
399
|
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
400
|
search_at_node<T, D>(nodes, indices, d_reference_points, 0, d_query_point, &best_index, &best_distance, &best_node);
|
268037bc
Jiaming Guo
made code more el...
|
401
|
radius = sqrt(best_distance); // get range
|
cc09e435
David Mayerich
added Jack's KD-T...
|
402
403
404
|
int cur = best_node;
while (nodes[cur].parent != -1) {
|
cc09e435
David Mayerich
added Jack's KD-T...
|
405
406
|
int parent = nodes[cur].parent;
int split_axis = nodes[parent].level % D;
|
268037bc
Jiaming Guo
made code more el...
|
407
|
|
cc09e435
David Mayerich
added Jack's KD-T...
|
408
409
|
T tmp_dist = FLT_MAX;
size_t tmp_idx;
|
268037bc
Jiaming Guo
made code more el...
|
410
|
if (fabs(nodes[parent].split_value - d_query_point.dim[split_axis]) <= radius) {
|
cc09e435
David Mayerich
added Jack's KD-T...
|
411
|
if (nodes[parent].left != cur)
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
412
|
search_at_node_range(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...
|
413
|
else
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
414
|
search_at_node_range(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...
|
415
|
}
|
268037bc
Jiaming Guo
made code more el...
|
416
417
418
|
if (tmp_dist < best_distance) {
best_distance = tmp_dist;
best_index = tmp_idx;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
419
420
421
|
}
cur = parent;
}
|
268037bc
Jiaming Guo
made code more el...
|
422
423
|
*d_distance = sqrt(best_distance);
*d_index = best_index;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
424
425
|
}
template <typename T, int D>
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
426
|
__global__ void search_batch(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) {
|
268037bc
Jiaming Guo
made code more el...
|
427
428
|
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...
|
429
|
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
430
|
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...
|
431
|
}
|
cc09e435
David Mayerich
added Jack's KD-T...
|
432
433
434
435
|
template <typename T, int D = 3>
class cuda_kdtree {
protected:
|
268037bc
Jiaming Guo
made code more el...
|
436
437
438
|
cuda_kdnode<T> *d_nodes;
size_t *d_index;
kdtree::point<T, D>* d_reference_points;
|
fcd2eb7c
David Mayerich
added support for...
|
439
|
size_t npts;
|
87d2ffe6
Jiaming Guo
add aabb to the k...
|
440
|
int num_nodes;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
441
442
|
public:
~cuda_kdtree() {
|
268037bc
Jiaming Guo
made code more el...
|
443
444
445
|
HANDLE_ERROR(cudaFree(d_nodes));
HANDLE_ERROR(cudaFree(d_index));
HANDLE_ERROR(cudaFree(d_reference_points));
|
cc09e435
David Mayerich
added Jack's KD-T...
|
446
|
}
|
fcd2eb7c
David Mayerich
added support for...
|
447
448
449
450
451
452
|
/// Create a KD-tree given a pointer to an array of reference points and the number of reference points
/// @param h_reference_points is a host array containing the reference points in (x0, y0, z0, ...., ) order
/// @param reference_count is the number of reference point in the array
/// @param max_levels is the deepest number of tree levels allowed
void create(T *h_reference_points, size_t reference_count, size_t max_levels = 3) {
|
268037bc
Jiaming Guo
made code more el...
|
453
454
455
|
if (max_levels > 10) {
std::cout<<"The max_tree_levels should be smaller!"<<std::endl;
exit(1);
|
87d2ffe6
Jiaming Guo
add aabb to the k...
|
456
|
}
|
fcd2eb7c
David Mayerich
added support for...
|
457
458
|
//bb.init(&h_reference_points[0]);
//aaboundingboxing<T, D>(bb, h_reference_points, reference_count);
|
87d2ffe6
Jiaming Guo
add aabb to the k...
|
459
|
|
feb0c9dd
David Mayerich
network processin...
|
460
|
std::vector < typename kdtree::point<T, D> > reference_points(reference_count); // restore the reference points in particular way
|
268037bc
Jiaming Guo
made code more el...
|
461
|
for (size_t j = 0; j < reference_count; j++)
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
462
463
|
for (size_t i = 0; i < D; i++)
reference_points[j].dim[i] = h_reference_points[j * D + i];
|
268037bc
Jiaming Guo
made code more el...
|
464
|
cpu_kdtree<T, D> tree; // creating a tree on cpu
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
465
466
|
tree.cpu_create(reference_points, max_levels); // building a tree on cpu
kdtree::kdnode<T> *d_root = tree.get_root();
|
87d2ffe6
Jiaming Guo
add aabb to the k...
|
467
|
num_nodes = tree.get_num_nodes();
|
fcd2eb7c
David Mayerich
added support for...
|
468
|
npts = reference_count; // also equals to reference_count
|
cc09e435
David Mayerich
added Jack's KD-T...
|
469
|
|
268037bc
Jiaming Guo
made code more el...
|
470
|
HANDLE_ERROR(cudaMalloc((void**)&d_nodes, sizeof(cuda_kdnode<T>) * num_nodes)); // copy data from host to device
|
fcd2eb7c
David Mayerich
added support for...
|
471
472
|
HANDLE_ERROR(cudaMalloc((void**)&d_index, sizeof(size_t) * npts));
HANDLE_ERROR(cudaMalloc((void**)&d_reference_points, sizeof(kdtree::point<T, D>) * npts));
|
cc09e435
David Mayerich
added Jack's KD-T...
|
473
|
|
e8c9a82b
David Mayerich
fixed GCC errors ...
|
474
|
std::vector < cuda_kdnode<T> > tmp_nodes(num_nodes);
|
fcd2eb7c
David Mayerich
added support for...
|
475
|
std::vector <size_t> indices(npts);
|
268037bc
Jiaming Guo
made code more el...
|
476
|
std::vector <kdtree::kdnode<T>*> next_nodes;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
477
|
size_t cur_pos = 0;
|
268037bc
Jiaming Guo
made code more el...
|
478
479
480
481
482
483
484
485
486
487
488
489
490
|
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...
|
491
492
493
494
|
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...
|
495
496
|
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...
|
497
498
|
}
else {
|
268037bc
Jiaming Guo
made code more el...
|
499
|
tmp_nodes[id].index = -1;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
500
501
502
|
}
if (cur->left)
|
268037bc
Jiaming Guo
made code more el...
|
503
|
next_search_nodes.push_back(cur->left);
|
cc09e435
David Mayerich
added Jack's KD-T...
|
504
505
|
if (cur->right)
|
268037bc
Jiaming Guo
made code more el...
|
506
|
next_search_nodes.push_back(cur->right);
|
cc09e435
David Mayerich
added Jack's KD-T...
|
507
|
}
|
268037bc
Jiaming Guo
made code more el...
|
508
|
next_nodes = next_search_nodes;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
509
|
}
|
268037bc
Jiaming Guo
made code more el...
|
510
511
512
|
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...
|
513
|
}
|
fcd2eb7c
David Mayerich
added support for...
|
514
515
516
517
518
519
|
/// Search the KD tree for nearest neighbors to a set of specified query points
/// @param h_query_points an array of query points in (x0, y0, z0, ...) order
/// @param query_count is the number of query points
/// @param indices are the indices to the nearest reference point for each query points
/// @param distances is an array containing the distance between each query point and the nearest reference point
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
520
|
void search(T *h_query_points, size_t query_count, size_t *indices, T *distances) {
|
e8c9a82b
David Mayerich
fixed GCC errors ...
|
521
|
std::vector < typename kdtree::point<T, D> > query_points(query_count);
|
268037bc
Jiaming Guo
made code more el...
|
522
|
for (size_t j = 0; j < query_count; j++)
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
523
524
|
for (size_t i = 0; i < D; i++)
query_points[j].dim[i] = h_query_points[j * D + i];
|
268037bc
Jiaming Guo
made code more el...
|
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
|
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()));
|
a0f09cd5
Jiaming Guo
add #define stack...
|
541
542
|
HANDLE_ERROR(cudaMalloc((void**)&next_nodes, threads * blocks * stack_size * sizeof(int))); // STACK size right now is 50, you can change it if you mean to
HANDLE_ERROR(cudaMalloc((void**)&next_search_nodes, threads * blocks * stack_size * sizeof(int)));
|
268037bc
Jiaming Guo
made code more el...
|
543
544
|
HANDLE_ERROR(cudaMemcpy(d_query_points, &query_points[0], sizeof(T) * query_points.size() * D, cudaMemcpyHostToDevice));
|
fcd2eb7c
David Mayerich
added support for...
|
545
|
search_batch<<<blocks, threads>>> (d_nodes, d_index, d_reference_points, d_query_points, query_points.size(), d_indices, d_distances, next_nodes, next_search_nodes, Judge);
|
268037bc
Jiaming Guo
made code more el...
|
546
547
548
|
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));
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
549
|
HANDLE_ERROR(cudaMemcpy(distances, d_distances, sizeof(T) * query_points.size(), cudaMemcpyDeviceToHost));
|
268037bc
Jiaming Guo
made code more el...
|
550
|
}
|
cc09e435
David Mayerich
added Jack's KD-T...
|
551
|
|
268037bc
Jiaming Guo
made code more el...
|
552
553
554
555
556
|
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...
|
557
|
}
|
fcd2eb7c
David Mayerich
added support for...
|
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
|
/// Return the number of points in the KD tree
size_t num_points() {
return npts;
}
stim::aabbn<T, D> getbox() {
size_t N = npts;
//std::vector < typename kdtree::point<T, D> > cpu_ref(npts); //allocate space on the CPU for the reference points
T* cpu_ref = (T*)malloc(N * D * sizeof(T)); //allocate space on the CPU for the reference points
HANDLE_ERROR(cudaMemcpy(cpu_ref, d_reference_points, N * D * sizeof(T), cudaMemcpyDeviceToHost)); //copy from GPU to CPU
stim::aabbn<T, D> bb(cpu_ref);
for (size_t i = 1; i < N; i++) { //for each reference point
//std::cout << "( " << cpu_ref[i * D + 0] << ", " << cpu_ref[i * D + 1] << ", " << cpu_ref[i * D + 2] << ")" << std::endl;
bb.insert(&cpu_ref[i * D]);
}
return bb;
}
//generate an implicit distance field for the KD-tree
void dist_field3(T* dist, size_t* dims, stim::aabbn<T, 3> bb) {
size_t N = 1; //number of query points that make up the distance field
for (size_t d = 0; d < 3; d++) N *= dims[d]; //calculate the total number of query points
//calculate the grid spatial parameters
T dx = 0;
if (dims[0] > 1) dx = bb.length(0) / dims[0];
T dy = 0;
if (dims[1] > 1) dy = bb.length(1) / dims[1];
T dz = 0;
if (dims[2] > 1) dz = bb.length(2) / dims[2];
T* Q = (T*)malloc(N * 3 * sizeof(T)); //allocate space for the query points
size_t i;
for (size_t z = 0; z < dims[2]; z++) { //for each query point (which is a point in the grid)
for (size_t y = 0; y < dims[1]; y++) {
for (size_t x = 0; x < dims[0]; x++) {
i = z * dims[1] * dims[0] + y * dims[0] + x;
Q[i * 3 + 0] = bb.low[0] + x * dx + dx / 2;
Q[i * 3 + 1] = bb.low[1] + y * dy + dy / 2;
Q[i * 3 + 2] = bb.low[2] + z * dz + dz / 2;
//std::cout << i<<" "<<Q[i * 3 + 0] << " " << Q[i * 3 + 1] << " " << Q[i * 3 + 2] << std::endl;
}
}
}
size_t* temp = (size_t*)malloc(N * sizeof(size_t)); //allocate space to store the indices (unused)
search(Q, N, temp, dist);
}
//generate an implicit distance field for the KD-tree
void dist_field3(T* dist, size_t* dims) {
stim::aabbn<T, D> bb = getbox(); //get a bounding box around the tree
dist_field3(dist, dims, bb);
}
|
cc09e435
David Mayerich
added Jack's KD-T...
|
615
616
617
|
};
} //end namespace stim
#endif
|