29 #include <tbb/parallel_for.h> 44 template <
class T,
class OUTPUT_ALLOCATOR,
int METRIC>
45 void _KnnSearchCPU(int64_t* query_neighbors_row_splits,
49 const T*
const queries,
51 bool ignore_query_point,
52 bool return_distances,
53 OUTPUT_ALLOCATOR& output_allocator) {
57 if (num_points == 0 || num_queries == 0) {
58 std::fill(query_neighbors_row_splits,
59 query_neighbors_row_splits + num_queries + 1, 0);
61 output_allocator.AllocIndices(&indices_ptr, 0);
64 output_allocator.AllocDistances(&distances_ptr, 0);
69 Adaptor<T> adaptor(num_points, points);
72 typename SelectNanoflannAdaptor<METRIC, T>::Adaptor_t, Adaptor<T>,
76 KDTree_t index(3, adaptor);
83 std::vector<Pair> pairs;
85 auto points_equal = [](
const T*
const p1,
const T*
const p2) {
86 return p1[0] == p2[0] && p1[1] == p2[1] && p1[2] == p2[2];
89 auto distance_fn = [](
const T*
const p1,
const T*
const p2) {
94 return dx * dx + dy * dy + dz * dz;
97 return std::abs(dx) + std::abs(dy) + std::abs(dz);
101 std::mutex pairs_mutex;
102 std::vector<uint32_t> neighbors_count(num_queries, 0);
106 tbb::blocked_range<size_t>(0, num_queries),
107 [&](
const tbb::blocked_range<size_t>& r) {
108 std::vector<Pair> pairs_private;
110 std::vector<size_t> result_indices(k);
111 std::vector<T> result_distances(k);
112 for (
size_t i = r.begin(); i != r.end(); ++i) {
113 size_t num_valid = index.knnSearch(&queries[i * 3], k,
114 result_indices.data(),
115 result_distances.data());
117 int num_neighbors = 0;
118 for (
size_t valid_i = 0; valid_i < num_valid; ++valid_i) {
119 size_t idx = result_indices[valid_i];
120 if (ignore_query_point &&
121 points_equal(&queries[i * 3], &points[idx * 3])) {
127 neighbors_count[i] = num_neighbors;
130 std::lock_guard<std::mutex> lock(pairs_mutex);
131 pairs.insert(pairs.end(), pairs_private.begin(),
132 pairs_private.end());
136 query_neighbors_row_splits[0] = 0;
138 &neighbors_count[neighbors_count.size()],
139 query_neighbors_row_splits + 1);
141 int32_t* neighbors_indices_ptr;
142 output_allocator.AllocIndices(&neighbors_indices_ptr, pairs.size());
144 if (return_distances)
145 output_allocator.AllocDistances(&distances_ptr, pairs.size());
147 output_allocator.AllocDistances(&distances_ptr, 0);
149 std::fill(neighbors_count.begin(), neighbors_count.end(), 0);
152 tbb::parallel_for(tbb::blocked_range<size_t>(0, pairs.size()),
153 [&](
const tbb::blocked_range<size_t>& r) {
154 for (
size_t i = r.begin(); i != r.end(); ++i) {
155 Pair pair = pairs[i];
158 query_neighbors_row_splits[pair.i] +
160 &neighbors_count[pair.i], 1);
161 neighbors_indices_ptr[idx] = pair.j;
163 if (return_distances) {
164 T dist = distance_fn(&points[pair.j * 3],
165 &queries[pair.i * 3]);
166 distances_ptr[idx] = dist;
221 template <
class T,
class OUTPUT_ALLOCATOR>
224 const T*
const points,
226 const T*
const queries,
229 const bool ignore_query_point,
230 const bool return_distances,
231 OUTPUT_ALLOCATOR& output_allocator) {
232 #define FN_PARAMETERS \ 233 query_neighbors_row_splits, num_points, points, num_queries, queries, k, \ 234 ignore_query_point, return_distances, output_allocator 236 #define CALL_TEMPLATE(METRIC) \ 237 if (METRIC == metric) \ 238 _KnnSearchCPU<T, OUTPUT_ALLOCATOR, METRIC>(FN_PARAMETERS); 240 #define CALL_TEMPLATE2 \ 247 #undef CALL_TEMPLATE2
void InclusivePrefixSum(const Tin *first, const Tin *last, Tout *out)
Definition: ParallelScan.h:67
Metric
Supported metrics.
Definition: NeighborSearchCommon.h:38
const char const char value recording_handle imu_sample recording_handle uint8_t size_t data_size k4a_record_configuration_t config target_format k4a_capture_t capture_handle k4a_imu_sample_t imu_sample playback_handle k4a_logging_message_cb_t void min_level device_handle k4a_imu_sample_t int32_t
Definition: K4aPlugin.cpp:398
int points
Definition: FilePCD.cpp:73
Definition: NanoFlannIndex.h:45
Definition: NeighborSearchCommon.h:38
Definition: PinholeCameraIntrinsic.cpp:35
uint32_t AtomicFetchAddRelaxed(uint32_t *address, uint32_t val)
Definition: Atomic.h:40
Definition: Console.cpp:51
void KnnSearchCPU(int64_t *query_neighbors_row_splits, size_t num_points, const T *const points, size_t num_queries, const T *const queries, const int k, const Metric metric, const bool ignore_query_point, const bool return_distances, OUTPUT_ALLOCATOR &output_allocator)
Definition: KnnSearch.h:222