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