-
Notifications
You must be signed in to change notification settings - Fork 7
/
Copy pathtest_kdtree.cpp
97 lines (78 loc) · 4.1 KB
/
test_kdtree.cpp
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
/**
*
* @file test_kdtree.cpp
* @author Thomas Grandits ([email protected])
* @brief Short test to see if the KD-Tree implementation works on the C++ side.
* This could also be used as a reference, if a simple C++ implementation is
* needed without using Tensorflow.
* @version 0.1
* @date 2020-11-16
*
* @copyright Copyright (c) 2020
*
*/
#include <random>
#include <iostream>
#include "kdtree.hpp"
#include "kdtree_g.hpp"
const int nr_points = 10000;
const int dims = 3;
const int levels = 11;
const int nr_query_points = 2000;
const point_i_knn_t knn = 20;
typedef double test_T;
template <typename T, dim_t dims>
std::tuple<T*, point_i_t*, T*> copyData(const std::vector<T>& result_dists, const std::vector<point_i_t>& result_idx,
const std::vector<std::array<T, dims>>& points_query);
template <typename T>
std::tuple<T*, point_i_t*> copyDataBackToHost(const T* result_dists, const point_i_t* result_idx, const size_t nr_query, const uint32_t nr_nns_searches);
template <typename T>
void freeGPUArray(T* arr);
int main()
{
std::random_device rd; //Will be used to obtain a seed for the random number engine
std::mt19937 gen(0); //Standard mersenne_twister_engine seeded with rd()
std::normal_distribution<test_T> distrib(0, 5);
std::vector<test_T> point_cloud;
for (int n=0; n<nr_points*dims; ++n)
//Use `distrib` to transform the random unsigned int generated by gen into an int in [1, 6]
point_cloud.push_back(distrib(gen));
std::cout << "Generated points..." << std::endl;
auto partition_info = createKDTree<test_T, dims>(point_cloud.data(), nr_points, levels);
assert(partition_info.levels == levels);
//fillKDTreeParents<test_T, dims>(main_partition);
std::vector<test_T> query_points;
for (int n=0; n<nr_query_points*dims; ++n)
//Use `distrib` to transform the random unsigned int generated by gen into an int in [1, 6]
query_points.push_back(distrib(gen));
std::vector<test_T> result_dists(nr_query_points*knn);
std::vector<point_i_t> result_idx(nr_query_points*knn);
KDTreeKNNSearch<test_T, test_T, dims>(partition_info,
nr_query_points, reinterpret_cast<std::array<test_T, dims>*>(query_points.data()), result_dists.data(), result_idx.data(), knn);
//GPU Test
PartitionInfoDevice<test_T, dims>* partition_info_d = copyPartitionToGPU<test_T, dims>(partition_info);
const auto tmp = copyData<test_T, dims>(result_dists, result_idx, std::vector<std::array<test_T, dims>>(reinterpret_cast<std::array<test_T, dims>*>(query_points.data()),
reinterpret_cast<std::array<test_T, dims>*>(query_points.data()) + query_points.size() / dims));
test_T* result_dists_d = std::get<0>(tmp);
point_i_t* result_idx_d = std::get<1>(tmp);
test_T* query_points_d = std::get<2>(tmp);
KDTreeKNNGPUSearch<test_T, test_T, dims>(partition_info_d, nr_query_points,
reinterpret_cast<std::array<test_T, dims>*>(query_points_d), result_dists_d, result_idx_d, knn);
auto result_gpu = copyDataBackToHost(result_dists_d, result_idx_d, nr_query_points, knn);
const auto result_dists_gpu = std::get<0>(result_gpu);
const auto result_idx_gpu = std::get<1>(result_gpu);
/*for (size_t i = 0; i < result_dists.size(); i++)
{
std::cout << result_dists[i] << ", " << result_dists_gpu[i] << " / " << result_idx[i] << ", " << result_idx_gpu[i] << "\n";
}*/
if(!std::equal(result_idx_gpu, result_idx_gpu + nr_query_points*knn, result_idx.begin()))
throw std::runtime_error("Error when comparing KD-Tree solution against the reference");
if(!std::equal(result_dists_gpu, result_dists_gpu + nr_query_points*knn, result_dists.begin(), [](const double lhs, const double rhs){return std::abs<double>(lhs - rhs) < 1e-7;}))
throw std::runtime_error("Error when comparing KD-Tree solution against the reference");
//Clean up
freePartitionFromGPU(partition_info_d);
freeGPUArray(result_dists_d);
freeGPUArray(result_idx_d);
std::cout << "Success" << std::endl;
return 0;
}