codebase / stimlib

Commit 87d2ffe658ed5974e68f3000b6c2f633dd727a55

Authored by Jiaming Guo
1 parent c17d9f44

add aabb to the kdtree::create

Showing 1 changed file with 14 additions and 3 deletions   Show diff stats
Inline Side-by-side
stim/structures/kdtree.cuh
  Show/Hide comments   View file @87d2ffe
... ... @@ -18,6 +18,7 @@
18 18 #include <iostream>
19 19 #include <algorithm>
20 20 #include <stim/cuda/cudatools/error.h>
  21 +#include <stim/visualization/aabbn.h>
21 22  
22 23 namespace stim {
23 24 namespace kdtree {
... ... @@ -427,6 +428,12 @@ namespace stim {
427 428  
428 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
429 430 }
  431 + template <typename T, int D>
  432 + __host__ void aaboundingboxing(stim::aabbn<T, D> &bb, T *reference_points, size_t reference_count) {
  433 + for(size_t i = 0; i < reference_count; i++) {
  434 + bb.insert(&reference_points[i]);
  435 + }
  436 + }
430 437  
431 438 template <typename T, int D = 3>
432 439 class cuda_kdtree {
... ... @@ -435,17 +442,21 @@ namespace stim {
435 442 size_t *d_index;
436 443 kdtree::point<T, D>* d_reference_points;
437 444 size_t d_reference_count;
  445 + int num_nodes;
438 446 public:
439 447 ~cuda_kdtree() {
440 448 HANDLE_ERROR(cudaFree(d_nodes));
441 449 HANDLE_ERROR(cudaFree(d_index));
442 450 HANDLE_ERROR(cudaFree(d_reference_points));
443 451 }
444   - void create(T *h_reference_points, size_t reference_count, size_t max_levels) {
  452 + void create(T *h_reference_points, size_t reference_count, size_t max_levels, stim::aabbn<T, D> &bb) {
445 453 if (max_levels > 10) {
446 454 std::cout<<"The max_tree_levels should be smaller!"<<std::endl;
447 455 exit(1);
448   - }
  456 + }
  457 + bb.init(&h_reference_points[0]);
  458 + aaboundingboxing<T, D>(bb, h_reference_points, reference_count);
  459 +
449 460 std::vector <kdtree::point<T, D>> reference_points(reference_count); // restore the reference points in particular way
450 461 for (size_t j = 0; j < reference_count; j++)
451 462 for (size_t i = 0; i < D; i++)
... ... @@ -453,7 +464,7 @@ namespace stim {
453 464 cpu_kdtree<T, D> tree; // creating a tree on cpu
454 465 tree.cpu_create(reference_points, max_levels); // building a tree on cpu
455 466 kdtree::kdnode<T> *d_root = tree.get_root();
456   - int num_nodes = tree.get_num_nodes();
  467 + num_nodes = tree.get_num_nodes();
457 468 d_reference_count = reference_count; // also equals to reference_count
458 469  
459 470 HANDLE_ERROR(cudaMalloc((void**)&d_nodes, sizeof(cuda_kdnode<T>) * num_nodes)); // copy data from host to device
... ...