kdtree.cuh
16.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
415
// change CUDA_STACK together with max_tree_levels in trial and error manner
// 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>
/// using API called HADDLE_ERROR
static void HandleError(cudaError_t err, const char *file, int line) {
if (err != cudaSuccess) {
std::cout<<cudaGetErrorString(err)<<" in"<< file <<" at line "<<line<<std::endl;
}
}
#define HANDLE_ERROR(err) (HandleError(err, __FILE__, __LINE__))
#define CUDA_STACK 2 // implementation "stacks" on CUDA as to do store the nodes information
namespace stim {
namespace kdtree {
template<typename T, int D>
struct point {
T coords[D]; // if we use size to measure a vector<point>, it will show the number of point structures
};
template<typename T>
class KDNode {
public:
KDNode() { // initialization
parent = NULL;
left = NULL;
right = NULL;
split_value = -1;
_parent = -1;
_left = -1;
_right = -1;
}
int id; // id for current node
size_t level;
KDNode *parent, *left, *right;
int _parent, _left, _right; // id for parent node
T split_value; // node value
std::vector <size_t> indices; // indices that indicate the data that current tree has
};
}
template <typename T, int D = 3>
class cpu_kdtree {
protected:
std::vector <kdtree::point<T, D>> *m_pts;
kdtree::KDNode<T> *m_root; // current node
int m_current_axis;
size_t m_levels;
int m_cmps; // count how many comparisons are to made in the tree for one query
int m_id; // level + 1
static cpu_kdtree<T, D> *myself;
public:
cpu_kdtree() { // initialization
myself = this;
m_id = 0; // id = level + 1, level -> axis index while id -> node identifier
}
~cpu_kdtree() { // destructor for deleting what was created by kdtree()
std::vector <kdtree::KDNode<T>*> next_node;
next_node.push_back(m_root);
while (next_node.size()) {
std::vector <kdtree::KDNode<T>*> next_search;
while (next_node.size()) {
kdtree::KDNode<T> *cur = next_node.back();
next_node.pop_back();
if (cur->left)
next_search.push_back(cur->left);
if (cur->right)
next_search.push_back(cur->right);
delete cur;
}
next_node = next_search;
}
m_root = NULL;
}
void Create(std::vector <kdtree::point<T, D>> &pts, size_t max_levels) {
m_pts = &pts; // create a pointer point to the input data
m_levels = max_levels; // stores max tree levels
m_root = new kdtree::KDNode<T>(); // using KDNode() to initialize an ancestor node
m_root->id = m_id++; // id is 1 while level is 0 at the very beginning
m_root->level = 0; // to begin with level 0
m_root->indices.resize(pts.size()); // initialize the size of whole indices
for (size_t i = 0; i < pts.size(); i++) {
m_root->indices[i] = i; // like what we did on Keys in GPU-BF part
}
std::vector <kdtree::KDNode<T>*> next_node; // next node
next_node.push_back(m_root); // new node
while (next_node.size()) {
std::vector <kdtree::KDNode<T>*> next_search;
while (next_node.size()) { // two same WHILE is because we need to make a new vector for searching
kdtree::KDNode<T> *current_node = next_node.back(); // pointer point to current node (right first)
next_node.pop_back(); // pop out current node in order to store next node
if (current_node->level < max_levels) { // max_levels should be reasonably small compared with numbers of data
if (current_node->indices.size() > 1) {
kdtree::KDNode<T> *left = new kdtree::KDNode<T>();
kdtree::KDNode<T> *right = new kdtree::KDNode<T>();
left->id = m_id++; // risky guessing but OK for large amount of data since max_level is small
right->id = m_id++; // risky guessing but OK for large amount of data since max_level is small
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 tree's indices
current_node->left = left;
current_node->right = right;
current_node->_left = left->id; // indicates it has left son node and gets its id
current_node->_right = right->id; // indicates it has right son node and gets its id
if (left->indices.size())
next_search.push_back(left); // right first then left according to stack(first in last out), it can be done in parallel for left and right are independent
if (right->indices.size())
next_search.push_back(right);
}
}
}
next_node = next_search;
}
}
static bool SortPoints(const size_t a, const size_t b) {
std::vector <kdtree::point<T, D>> &pts = *myself->m_pts;
return pts[a].coords[myself->m_current_axis] < pts[b].coords[myself->m_current_axis];
}
void Split(kdtree::KDNode<T> *cur, kdtree::KDNode<T> *left, kdtree::KDNode<T> *right) {
/// assume both two sides are created and sure it was
std::vector <kdtree::point<T, D>> &pts = *m_pts;
m_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 = cur->indices[cur->indices.size() / 2]; // odd in the mid, even take the floor
cur->split_value = pts[mid].coords[m_current_axis]; // get the mother node
left->parent = cur; // set the parent to current node for the next nodes
right->parent = cur;
left->level = cur->level + 1;
right->level = cur->level + 1;
left->_parent = cur->id; // indicates it has mother node and gets its id
right->_parent = cur->id; // indicates it has mother node and gets its id
for (size_t i = 0; i < cur->indices.size(); i++) { // split into left and right area one by one
size_t idx = cur->indices[i];
if (pts[idx].coords[m_current_axis] < cur->split_value)
left->indices.push_back(idx);
else
right->indices.push_back(idx);
}
}
int GetNumNodes() const { return m_id; }
kdtree::KDNode<T>* GetRoot() const { return m_root; }
}; //end class kdtree
template <typename T, int D>
cpu_kdtree<T, D>* cpu_kdtree<T, D>::myself = NULL; // definition of myself pointer points to the specific class
template <typename T>
struct CUDA_KDNode {
size_t level;
int parent, left, right; // indicates id of
T split_value;
size_t num_indices; // number of indices it has
int indices; // the beginning
};
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.coords[i] - b.coords[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> *pts, int cur, kdtree::point<T, D> &Query, size_t *ret_index, T *ret_dist, int *ret_node) {
/// finds the first possibility
size_t best_idx = 0;
T best_dist = FLT_MAX;
while (true) {
int split_axis = nodes[cur].level % D;
if (nodes[cur].left == -1) { // if it doesn't have left son node
*ret_node = cur;
for (int i = 0; i < nodes[cur].num_indices; i++) {
size_t idx = indices[nodes[cur].indices + i];
T dist = Distance<T, D>(Query, pts[idx]);
if (dist < best_dist) {
best_dist = dist;
best_idx = idx;
}
}
break;
}
else if (Query.coords[split_axis] < nodes[cur].split_value) {
cur = nodes[cur].left;
}
else {
cur = nodes[cur].right;
}
}
*ret_index = best_idx;
*ret_dist = best_dist;
}
template <typename T, int D>
__device__ void SearchAtNodeRange(CUDA_KDNode<T> *nodes, size_t *indices, kdtree::point<T, D> *pts, kdtree::point<T, D> &Query, int cur, T range, size_t *ret_index, T *ret_dist) {
/// search throught all the nodes that are within one range
size_t best_idx = 0;
T best_dist = FLT_MAX;
/// using fixed stack, increase it when need
int next_node[CUDA_STACK]; // should be larger than 1
int next_node_pos = 0; // initialize pop out order index
next_node[next_node_pos++] = cur; // equals to next_node[next_node_pos] = cur, next_node_pos++
while (next_node_pos) {
int next_search[CUDA_STACK]; // for store next nodes
int next_search_pos = 0; // record push back order index
while (next_node_pos) {
cur = next_node[next_node_pos - 1]; // pop out the last in one and keep poping out
next_node_pos--;
int split_axis = nodes[cur].level % D;
if (nodes[cur].left == -1) {
for (int i = 0; i < nodes[cur].num_indices; i++) {
int idx = indices[nodes[cur].indices + i]; // all indices are stored in one array, pick up from every node's beginning index
T d = Distance<T>(Query, pts[idx]);
if (d < best_dist) {
best_dist = d;
best_idx = idx;
}
}
}
else {
T d = Query.coords[split_axis] - nodes[cur].split_value;
if (fabs(d) > range) {
if (d < 0)
next_search[next_search_pos++] = nodes[cur].left;
else
next_search[next_search_pos++] = nodes[cur].right;
}
else {
next_search[next_search_pos++] = nodes[cur].left;
next_search[next_search_pos++] = nodes[cur].right;
}
}
}
for (int i = 0; i < next_search_pos; i++)
next_node[i] = next_search[i];
next_node_pos = next_search_pos; // operation that really resemble STACK, namely first in last out
}
*ret_index = best_idx;
*ret_dist = best_dist;
}
template <typename T, int D>
__device__ void Search(CUDA_KDNode<T> *nodes, size_t *indices, kdtree::point<T, D> *pts, kdtree::point<T, D> &Query, size_t *ret_index, T *ret_dist) {
/// find first nearest node
int best_node = 0;
size_t best_idx = 0;
T best_dist = FLT_MAX;
T radius = 0;
SearchAtNode<T, D>(nodes, indices, pts, 0, Query, &best_idx, &best_dist, &best_node);
radius = sqrt(best_dist);
/// find other possibilities
int cur = best_node;
while (nodes[cur].parent != -1) {
/// go up
int parent = nodes[cur].parent;
int split_axis = nodes[parent].level % D;
/// search other node
T tmp_dist = FLT_MAX;
size_t tmp_idx;
if (fabs(nodes[parent].split_value - Query.coords[split_axis]) <= radius) {
/// search opposite node
if (nodes[parent].left != cur)
SearchAtNodeRange(nodes, indices, pts, Query, nodes[parent].left, radius, &tmp_idx, &tmp_dist);
else
SearchAtNodeRange(nodes, indices, pts, Query, nodes[parent].right, radius, &tmp_idx, &tmp_dist);
}
if (tmp_dist < best_dist) {
best_dist = tmp_dist;
best_idx = tmp_idx;
}
cur = parent;
}
*ret_index = best_idx;
*ret_dist = sqrt(best_dist);
}
template <typename T, int D>
__global__ void SearchBatch(CUDA_KDNode<T> *nodes, size_t *indices, kdtree::point<T, D> *pts, size_t num_pts, kdtree::point<T, D> *Query, size_t num_Query, size_t *ret_index, T *ret_dist) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= num_Query) return;
Search<T, D>(nodes, indices, pts, Query[idx], &ret_index[idx], &ret_dist[idx]); // every query points are independent
}
template <typename T, int D = 3>
class cuda_kdtree {
protected:
CUDA_KDNode<T> *m_gpu_nodes; // store nodes
size_t *m_gpu_indices;
kdtree::point<T, D>* m_gpu_points;
size_t m_num_points;
public:
~cuda_kdtree() {
HANDLE_ERROR(cudaFree(m_gpu_nodes));
HANDLE_ERROR(cudaFree(m_gpu_indices));
HANDLE_ERROR(cudaFree(m_gpu_points));
}
void CreateKDTree(T *ReferencePoints, size_t ReferenceCount, size_t ColCount, size_t max_levels) {
std::vector < kdtree::point<T, D> > pts(ReferenceCount); // create specific struct of reference data
for (size_t j = 0; j < ReferenceCount; j++)
for (size_t i = 0; i < ColCount; i++)
pts[j].coords[i] = ReferencePoints[j * ColCount + i];
cpu_kdtree<T, D> tree; // initialize a tree
tree.Create(pts, max_levels); // create KD-Tree on host
kdtree::KDNode<T> *root = tree.GetRoot();
int num_nodes = tree.GetNumNodes();
/// create the same on CPU
m_num_points = pts.size(); // number of points for creating tree = reference_count in the case
HANDLE_ERROR(cudaMalloc((void**)&m_gpu_nodes, sizeof(CUDA_KDNode<T>) * num_nodes)); // private variables for kdtree
HANDLE_ERROR(cudaMalloc((void**)&m_gpu_indices, sizeof(size_t) * m_num_points));
HANDLE_ERROR(cudaMalloc((void**)&m_gpu_points, sizeof(kdtree::point<T, D>) * m_num_points));
std::vector <CUDA_KDNode<T>> cpu_nodes(num_nodes); // from left to right, id of nodes
std::vector <size_t> indices(m_num_points);
std::vector < kdtree::KDNode<T>* > next_node;
size_t cur_pos = 0;
next_node.push_back(root);
while (next_node.size()) {
std::vector <typename kdtree::KDNode<T>* > next_search;
while (next_node.size()) {
kdtree::KDNode<T> *cur = next_node.back();
next_node.pop_back();
int id = cur->id; // the nodes at same level are independent
cpu_nodes[id].level = cur->level;
cpu_nodes[id].parent = cur->_parent;
cpu_nodes[id].left = cur->_left;
cpu_nodes[id].right = cur->_right;
cpu_nodes[id].split_value = cur->split_value;
cpu_nodes[id].num_indices = 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];
cpu_nodes[id].indices = (int)cur_pos; // beginning index that every bottom node has
cur_pos += cur->indices.size(); // store indices continuously
}
else {
cpu_nodes[id].indices = -1;
}
if (cur->left)
next_search.push_back(cur->left);
if (cur->right)
next_search.push_back(cur->right);
}
next_node = next_search;
}
HANDLE_ERROR(cudaMemcpy(m_gpu_nodes, &cpu_nodes[0], sizeof(CUDA_KDNode<T>) * cpu_nodes.size(), cudaMemcpyHostToDevice));
HANDLE_ERROR(cudaMemcpy(m_gpu_indices, &indices[0], sizeof(size_t) * indices.size(), cudaMemcpyHostToDevice));
HANDLE_ERROR(cudaMemcpy(m_gpu_points, &pts[0], sizeof(kdtree::point<T, D>) * pts.size(), cudaMemcpyHostToDevice));
}
void Search(T *QueryPoints, size_t QueryCount, size_t ColCount, T *dists, size_t *indices) {
std::vector < kdtree::point<T, D> > query(QueryCount);
for (size_t j = 0; j < QueryCount; j++)
for (size_t i = 0; i < ColCount; i++)
query[j].coords[i] = QueryPoints[j * ColCount + i];
unsigned int threads = (unsigned int)(query.size() > 1024 ? 1024 : query.size());
//unsigned int blocks = (unsigned int)ceil(query.size() / threads);
unsigned int blocks = (unsigned int)(query.size() / threads + (query.size() % threads ? 1 : 0));
kdtree::point<T, D> *gpu_Query;
size_t *gpu_ret_indices;
T *gpu_ret_dist;
HANDLE_ERROR(cudaMalloc((void**)&gpu_Query, sizeof(T) * query.size() * D));
HANDLE_ERROR(cudaMalloc((void**)&gpu_ret_indices, sizeof(size_t) * query.size()));
HANDLE_ERROR(cudaMalloc((void**)&gpu_ret_dist, sizeof(T) * query.size()));
HANDLE_ERROR(cudaMemcpy(gpu_Query, &query[0], sizeof(T) * query.size() * D, cudaMemcpyHostToDevice));
SearchBatch << <threads, blocks >> > (m_gpu_nodes, m_gpu_indices, m_gpu_points, m_num_points, gpu_Query, query.size(), gpu_ret_indices, gpu_ret_dist);
HANDLE_ERROR(cudaMemcpy(indices, gpu_ret_indices, sizeof(size_t) * query.size(), cudaMemcpyDeviceToHost));
HANDLE_ERROR(cudaMemcpy(dists, gpu_ret_dist, sizeof(T) * query.size(), cudaMemcpyDeviceToHost));
HANDLE_ERROR(cudaFree(gpu_Query));
HANDLE_ERROR(cudaFree(gpu_ret_indices));
HANDLE_ERROR(cudaFree(gpu_ret_dist));
}
};
} //end namespace stim
#endif