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
10
11
12
13
14
15
|
// 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>
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
16
|
#include <cstring>
|
cc09e435
David Mayerich
added Jack's KD-T...
|
17
18
19
|
#include <float.h>
#include <iostream>
#include <algorithm>
|
268037bc
Jiaming Guo
made code more el...
|
20
|
#include <stim/cuda/cudatools/error.h>
|
87d2ffe6
Jiaming Guo
add aabb to the k...
|
21
|
#include <stim/visualization/aabbn.h>
|
cc09e435
David Mayerich
added Jack's KD-T...
|
22
23
24
|
namespace stim {
namespace kdtree {
|
268037bc
Jiaming Guo
made code more el...
|
25
|
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...
|
26
|
struct point {
|
268037bc
Jiaming Guo
made code more el...
|
27
|
T dim[D]; // create a structure to store every one input point
|
cc09e435
David Mayerich
added Jack's KD-T...
|
28
29
30
|
};
template<typename T>
|
268037bc
Jiaming Guo
made code more el...
|
31
|
class kdnode {
|
cc09e435
David Mayerich
added Jack's KD-T...
|
32
|
public:
|
268037bc
Jiaming Guo
made code more el...
|
33
34
|
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...
|
35
36
|
left = NULL;
right = NULL;
|
268037bc
Jiaming Guo
made code more el...
|
37
38
39
40
|
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...
|
41
|
}
|
268037bc
Jiaming Guo
made code more el...
|
42
43
44
45
46
47
|
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...
|
48
|
};
|
268037bc
Jiaming Guo
made code more el...
|
49
|
} // end of namespace kdtree
|
cc09e435
David Mayerich
added Jack's KD-T...
|
50
|
|
268037bc
Jiaming Guo
made code more el...
|
51
|
template <typename T, int D = 3> // set dimension of data to default 3
|
cc09e435
David Mayerich
added Jack's KD-T...
|
52
|
class cpu_kdtree {
|
cc09e435
David Mayerich
added Jack's KD-T...
|
53
|
protected:
|
268037bc
Jiaming Guo
made code more el...
|
54
|
int current_axis; // current judging axis
|
268037bc
Jiaming Guo
made code more el...
|
55
|
int n_id; // store the total number of nodes
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
56
57
|
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...
|
58
59
|
kdtree::kdnode<T> *root; // root node
static cpu_kdtree<T, D> *cur_tree_ptr;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
60
|
public:
|
268037bc
Jiaming Guo
made code more el...
|
61
|
cpu_kdtree() { // constructor for creating a cpu_kdtree
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
62
|
cur_tree_ptr = this; // create a class pointer points to the current class value
|
268037bc
Jiaming Guo
made code more el...
|
63
|
n_id = 0; // set total number of points to default 0
|
cc09e435
David Mayerich
added Jack's KD-T...
|
64
|
}
|
268037bc
Jiaming Guo
made code more el...
|
65
66
67
68
69
70
71
72
|
~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...
|
73
|
if (cur->left)
|
268037bc
Jiaming Guo
made code more el...
|
74
|
next_search_nodes.push_back(cur->left);
|
cc09e435
David Mayerich
added Jack's KD-T...
|
75
|
if (cur->right)
|
268037bc
Jiaming Guo
made code more el...
|
76
|
next_search_nodes.push_back(cur->right);
|
cc09e435
David Mayerich
added Jack's KD-T...
|
77
78
|
delete cur;
}
|
268037bc
Jiaming Guo
made code more el...
|
79
|
next_nodes = next_search_nodes;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
80
|
}
|
268037bc
Jiaming Guo
made code more el...
|
81
|
root = NULL;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
82
|
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
83
84
|
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...
|
85
86
87
88
89
90
|
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...
|
91
|
}
|
268037bc
Jiaming Guo
made code more el...
|
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++;
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
105
|
split(current_node, left, right); // split left and right and determine a node
|
268037bc
Jiaming Guo
made code more el...
|
106
|
std::vector <size_t> temp; // empty vecters of int
|
cc09e435
David Mayerich
added Jack's KD-T...
|
107
|
//temp.resize(current_node->indices.size());
|
268037bc
Jiaming Guo
made code more el...
|
108
|
current_node->indices.swap(temp); // clean up current node's indices
|
cc09e435
David Mayerich
added Jack's KD-T...
|
109
110
|
current_node->left = left;
current_node->right = right;
|
268037bc
Jiaming Guo
made code more el...
|
111
112
|
current_node->left_idx = left->idx;
current_node->right_idx = right->idx;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
113
|
if (right->indices.size())
|
268037bc
Jiaming Guo
made code more el...
|
114
115
116
|
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...
|
117
118
119
|
}
}
}
|
268037bc
Jiaming Guo
made code more el...
|
120
|
next_nodes = next_search_nodes; // go deeper within the tree
|
cc09e435
David Mayerich
added Jack's KD-T...
|
121
122
|
}
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
123
|
static bool sort_points(const size_t a, const size_t b) { // create functor for std::sort
|
e8c9a82b
David Mayerich
fixed GCC errors ...
|
124
|
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...
|
125
|
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...
|
126
|
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
127
|
void split(kdtree::kdnode<T> *cur, kdtree::kdnode<T> *left, kdtree::kdnode<T> *right) {
|
e8c9a82b
David Mayerich
fixed GCC errors ...
|
128
|
std::vector < typename kdtree::point<T, D> > &pts = *tmp_points;
|
268037bc
Jiaming Guo
made code more el...
|
129
|
current_axis = cur->level % D; // indicate the judicative dimension or axis
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
130
|
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...
|
131
132
133
|
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...
|
134
|
right->parent = cur;
|
268037bc
Jiaming Guo
made code more el...
|
135
|
left->level = cur->level + 1; // level + 1
|
cc09e435
David Mayerich
added Jack's KD-T...
|
136
|
right->level = cur->level + 1;
|
268037bc
Jiaming Guo
made code more el...
|
137
138
139
|
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...
|
140
|
size_t idx = cur->indices[i];
|
268037bc
Jiaming Guo
made code more el...
|
141
|
if (pts[idx].dim[current_axis] < cur->split_value)
|
cc09e435
David Mayerich
added Jack's KD-T...
|
142
143
144
145
146
|
left->indices.push_back(idx);
else
right->indices.push_back(idx);
}
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
147
148
149
150
151
152
153
154
155
|
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...
|
156
157
|
return n_id;
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
158
|
kdtree::kdnode<T>* get_root() const { // get the root node of tree
|
268037bc
Jiaming Guo
made code more el...
|
159
160
|
return root;
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
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
|
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...
|
281
282
283
|
}; //end class kdtree
template <typename T, int D>
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
284
|
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...
|
285
286
|
template <typename T>
|
268037bc
Jiaming Guo
made code more el...
|
287
288
|
struct cuda_kdnode {
int parent, left, right;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
289
|
T split_value;
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
290
291
|
size_t num_index; // number of indices it has
int index; // the beginning index
|
268037bc
Jiaming Guo
made code more el...
|
292
|
size_t level;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
293
294
295
|
};
template <typename T, int D>
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
296
297
|
__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...
|
298
299
|
for (size_t i = 0; i < D; i++) {
|
268037bc
Jiaming Guo
made code more el...
|
300
|
T d = a.dim[i] - b.dim[i];
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
301
|
distance += d*d;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
302
|
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
303
|
return distance;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
304
305
|
}
template <typename T, int D>
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
306
|
__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...
|
307
308
|
T best_distance = FLT_MAX;
size_t best_index = 0;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
309
|
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
310
|
while (true) { // break until reach the bottom
|
cc09e435
David Mayerich
added Jack's KD-T...
|
311
|
int split_axis = nodes[cur].level % D;
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
312
|
if (nodes[cur].left == -1) { // check whether it has left node or not
|
268037bc
Jiaming Guo
made code more el...
|
313
314
315
|
*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...
|
316
|
T dist = gpu_distance<T, D>(d_query_point, d_reference_points[idx]);
|
268037bc
Jiaming Guo
made code more el...
|
317
318
319
|
if (dist < best_distance) {
best_distance = dist;
best_index = idx;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
320
321
|
}
}
|
268037bc
Jiaming Guo
made code more el...
|
322
|
break;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
323
|
}
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
324
|
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...
|
325
326
327
328
329
330
|
cur = nodes[cur].left;
}
else {
cur = nodes[cur].right;
}
}
|
268037bc
Jiaming Guo
made code more el...
|
331
332
|
*d_distance = best_distance;
*d_index = best_index;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
333
334
|
}
template <typename T, int D>
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
335
|
__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...
|
336
337
338
|
T best_distance = FLT_MAX;
size_t best_index = 0;
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
339
340
|
int next_nodes_pos = 0; // initialize pop out order index
next_nodes[id * 50 + next_nodes_pos] = cur; // find data that belongs to the very specific thread
|
268037bc
Jiaming Guo
made code more el...
|
341
342
343
|
next_nodes_pos++;
while (next_nodes_pos) {
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
344
|
int next_search_nodes_pos = 0; // record push back order index
|
268037bc
Jiaming Guo
made code more el...
|
345
|
while (next_nodes_pos) {
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
346
|
cur = next_nodes[id * 50 + next_nodes_pos - 1]; // pop out the last push in one and keep poping out
|
268037bc
Jiaming Guo
made code more el...
|
347
|
next_nodes_pos--;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
348
349
350
|
int split_axis = nodes[cur].level % D;
if (nodes[cur].left == -1) {
|
268037bc
Jiaming Guo
made code more el...
|
351
|
for (int i = 0; i < nodes[cur].num_index; i++) {
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
352
353
|
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...
|
354
355
356
|
if (d < best_distance) {
best_distance = d;
best_index = idx;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
357
358
359
360
|
}
}
}
else {
|
268037bc
Jiaming Guo
made code more el...
|
361
|
T d = d_query_point.dim[split_axis] - nodes[cur].split_value;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
362
363
|
if (fabs(d) > range) {
|
268037bc
Jiaming Guo
made code more el...
|
364
|
if (d < 0) {
|
fbbc07be
Jiaming Guo
replace ANN by kd...
|
365
|
next_search_nodes[id * 50 + next_search_nodes_pos] = nodes[cur].left;
|
268037bc
Jiaming Guo
made code more el...
|
366
367
368
|
next_search_nodes_pos++;
}
else {
|
fbbc07be
Jiaming Guo
replace ANN by kd...
|
369
|
next_search_nodes[id * 50 + next_search_nodes_pos] = nodes[cur].right;
|
268037bc
Jiaming Guo
made code more el...
|
370
371
|
next_search_nodes_pos++;
}
|
cc09e435
David Mayerich
added Jack's KD-T...
|
372
373
|
}
else {
|
fbbc07be
Jiaming Guo
replace ANN by kd...
|
374
|
next_search_nodes[id * 50 + next_search_nodes_pos] = nodes[cur].right;
|
268037bc
Jiaming Guo
made code more el...
|
375
|
next_search_nodes_pos++;
|
fbbc07be
Jiaming Guo
replace ANN by kd...
|
376
|
next_search_nodes[id * 50 + next_search_nodes_pos] = nodes[cur].left;
|
268037bc
Jiaming Guo
made code more el...
|
377
|
next_search_nodes_pos++;
|
fbbc07be
Jiaming Guo
replace ANN by kd...
|
378
|
if (next_search_nodes_pos > 50) {
|
268037bc
Jiaming Guo
made code more el...
|
379
380
381
|
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...
|
382
383
384
|
}
}
}
|
268037bc
Jiaming Guo
made code more el...
|
385
|
for (int i = 0; i < next_search_nodes_pos; i++)
|
fbbc07be
Jiaming Guo
replace ANN by kd...
|
386
|
next_nodes[id * 50 + i] = next_search_nodes[id * 50 + i];
|
268037bc
Jiaming Guo
made code more el...
|
387
|
next_nodes_pos = next_search_nodes_pos;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
388
|
}
|
268037bc
Jiaming Guo
made code more el...
|
389
390
|
*d_distance = best_distance;
*d_index = best_index;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
391
392
|
}
template <typename T, int D>
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
393
|
__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...
|
394
|
int best_node = 0;
|
268037bc
Jiaming Guo
made code more el...
|
395
396
|
T best_distance = FLT_MAX;
size_t best_index = 0;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
397
|
T radius = 0;
|
268037bc
Jiaming Guo
made code more el...
|
398
|
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
399
|
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...
|
400
|
radius = sqrt(best_distance); // get range
|
cc09e435
David Mayerich
added Jack's KD-T...
|
401
402
403
|
int cur = best_node;
while (nodes[cur].parent != -1) {
|
cc09e435
David Mayerich
added Jack's KD-T...
|
404
405
|
int parent = nodes[cur].parent;
int split_axis = nodes[parent].level % D;
|
268037bc
Jiaming Guo
made code more el...
|
406
|
|
cc09e435
David Mayerich
added Jack's KD-T...
|
407
408
|
T tmp_dist = FLT_MAX;
size_t tmp_idx;
|
268037bc
Jiaming Guo
made code more el...
|
409
|
if (fabs(nodes[parent].split_value - d_query_point.dim[split_axis]) <= radius) {
|
cc09e435
David Mayerich
added Jack's KD-T...
|
410
|
if (nodes[parent].left != cur)
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
411
|
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...
|
412
|
else
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
413
|
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...
|
414
|
}
|
268037bc
Jiaming Guo
made code more el...
|
415
416
417
|
if (tmp_dist < best_distance) {
best_distance = tmp_dist;
best_index = tmp_idx;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
418
419
420
|
}
cur = parent;
}
|
268037bc
Jiaming Guo
made code more el...
|
421
422
|
*d_distance = sqrt(best_distance);
*d_index = best_index;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
423
424
|
}
template <typename T, int D>
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
425
|
__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...
|
426
427
|
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...
|
428
|
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
429
|
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...
|
430
|
}
|
cc09e435
David Mayerich
added Jack's KD-T...
|
431
432
433
434
|
template <typename T, int D = 3>
class cuda_kdtree {
protected:
|
268037bc
Jiaming Guo
made code more el...
|
435
436
437
|
cuda_kdnode<T> *d_nodes;
size_t *d_index;
kdtree::point<T, D>* d_reference_points;
|
fcd2eb7c
David Mayerich
added support for...
|
438
|
size_t npts;
|
87d2ffe6
Jiaming Guo
add aabb to the k...
|
439
|
int num_nodes;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
440
441
|
public:
~cuda_kdtree() {
|
268037bc
Jiaming Guo
made code more el...
|
442
443
444
|
HANDLE_ERROR(cudaFree(d_nodes));
HANDLE_ERROR(cudaFree(d_index));
HANDLE_ERROR(cudaFree(d_reference_points));
|
cc09e435
David Mayerich
added Jack's KD-T...
|
445
|
}
|
fcd2eb7c
David Mayerich
added support for...
|
446
447
448
449
450
451
|
/// 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...
|
452
453
454
|
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...
|
455
|
}
|
fcd2eb7c
David Mayerich
added support for...
|
456
457
|
//bb.init(&h_reference_points[0]);
//aaboundingboxing<T, D>(bb, h_reference_points, reference_count);
|
87d2ffe6
Jiaming Guo
add aabb to the k...
|
458
|
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
459
|
std::vector <kdtree::point<T, D>> reference_points(reference_count); // restore the reference points in particular way
|
268037bc
Jiaming Guo
made code more el...
|
460
|
for (size_t j = 0; j < reference_count; j++)
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
461
462
|
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...
|
463
|
cpu_kdtree<T, D> tree; // creating a tree on cpu
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
464
465
|
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...
|
466
|
num_nodes = tree.get_num_nodes();
|
fcd2eb7c
David Mayerich
added support for...
|
467
|
npts = reference_count; // also equals to reference_count
|
cc09e435
David Mayerich
added Jack's KD-T...
|
468
|
|
268037bc
Jiaming Guo
made code more el...
|
469
|
HANDLE_ERROR(cudaMalloc((void**)&d_nodes, sizeof(cuda_kdnode<T>) * num_nodes)); // copy data from host to device
|
fcd2eb7c
David Mayerich
added support for...
|
470
471
|
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...
|
472
|
|
e8c9a82b
David Mayerich
fixed GCC errors ...
|
473
|
std::vector < cuda_kdnode<T> > tmp_nodes(num_nodes);
|
fcd2eb7c
David Mayerich
added support for...
|
474
|
std::vector <size_t> indices(npts);
|
268037bc
Jiaming Guo
made code more el...
|
475
|
std::vector <kdtree::kdnode<T>*> next_nodes;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
476
|
size_t cur_pos = 0;
|
268037bc
Jiaming Guo
made code more el...
|
477
478
479
480
481
482
483
484
485
486
487
488
489
|
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...
|
490
491
492
493
|
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...
|
494
495
|
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...
|
496
497
|
}
else {
|
268037bc
Jiaming Guo
made code more el...
|
498
|
tmp_nodes[id].index = -1;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
499
500
501
|
}
if (cur->left)
|
268037bc
Jiaming Guo
made code more el...
|
502
|
next_search_nodes.push_back(cur->left);
|
cc09e435
David Mayerich
added Jack's KD-T...
|
503
504
|
if (cur->right)
|
268037bc
Jiaming Guo
made code more el...
|
505
|
next_search_nodes.push_back(cur->right);
|
cc09e435
David Mayerich
added Jack's KD-T...
|
506
|
}
|
268037bc
Jiaming Guo
made code more el...
|
507
|
next_nodes = next_search_nodes;
|
cc09e435
David Mayerich
added Jack's KD-T...
|
508
|
}
|
268037bc
Jiaming Guo
made code more el...
|
509
510
511
|
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...
|
512
|
}
|
fcd2eb7c
David Mayerich
added support for...
|
513
514
515
516
517
518
|
/// 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...
|
519
|
void search(T *h_query_points, size_t query_count, size_t *indices, T *distances) {
|
e8c9a82b
David Mayerich
fixed GCC errors ...
|
520
|
std::vector < typename kdtree::point<T, D> > query_points(query_count);
|
268037bc
Jiaming Guo
made code more el...
|
521
|
for (size_t j = 0; j < query_count; j++)
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
522
523
|
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...
|
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
|
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()));
|
41100e51
Jiaming Guo
add cpu kdtree ve...
|
540
|
HANDLE_ERROR(cudaMalloc((void**)&next_nodes, threads * blocks * 50 * sizeof(int))); // STACK size right now is 50, you can change it if you mean to
|
fbbc07be
Jiaming Guo
replace ANN by kd...
|
541
|
HANDLE_ERROR(cudaMalloc((void**)&next_search_nodes, threads * blocks * 50 * sizeof(int)));
|
268037bc
Jiaming Guo
made code more el...
|
542
543
|
HANDLE_ERROR(cudaMemcpy(d_query_points, &query_points[0], sizeof(T) * query_points.size() * D, cudaMemcpyHostToDevice));
|
fcd2eb7c
David Mayerich
added support for...
|
544
|
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...
|
545
546
547
|
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...
|
548
|
HANDLE_ERROR(cudaMemcpy(distances, d_distances, sizeof(T) * query_points.size(), cudaMemcpyDeviceToHost));
|
268037bc
Jiaming Guo
made code more el...
|
549
|
}
|
cc09e435
David Mayerich
added Jack's KD-T...
|
550
|
|
268037bc
Jiaming Guo
made code more el...
|
551
552
553
554
555
|
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...
|
556
|
}
|
fcd2eb7c
David Mayerich
added support for...
|
557
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
|
/// 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...
|
614
615
616
|
};
} //end namespace stim
#endif
|