Open3D (C++ API)  0.15.1
FixedRadiusSearchImpl.h
Go to the documentation of this file.
1// ----------------------------------------------------------------------------
2// - Open3D: www.open3d.org -
3// ----------------------------------------------------------------------------
4// The MIT License (MIT)
5//
6// Copyright (c) 2018-2021 www.open3d.org
7//
8// Permission is hereby granted, free of charge, to any person obtaining a copy
9// of this software and associated documentation files (the "Software"), to deal
10// in the Software without restriction, including without limitation the rights
11// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
12// copies of the Software, and to permit persons to whom the Software is
13// furnished to do so, subject to the following conditions:
14//
15// The above copyright notice and this permission notice shall be included in
16// all copies or substantial portions of the Software.
17//
18// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
23// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
24// IN THE SOFTWARE.
25// ----------------------------------------------------------------------------
26
27#pragma once
28
29#include <tbb/parallel_for.h>
30
31#include <set>
32
33#include "open3d/core/Atomic.h"
38
39namespace open3d {
40namespace core {
41namespace nns {
42namespace impl {
43
44namespace {
45
78template <class T>
79void BuildSpatialHashTableCPU(const size_t num_points,
80 const T* const points,
81 const T radius,
82 const size_t points_row_splits_size,
83 const int64_t* points_row_splits,
84 const uint32_t* hash_table_splits,
85 const size_t hash_table_cell_splits_size,
86 uint32_t* hash_table_cell_splits,
87 uint32_t* hash_table_index) {
88 using namespace open3d::utility;
89 typedef MiniVec<T, 3> Vec3_t;
90
91 const int batch_size = points_row_splits_size - 1;
92 const T voxel_size = 2 * radius;
93 const T inv_voxel_size = 1 / voxel_size;
94
95 memset(&hash_table_cell_splits[0], 0,
96 sizeof(uint32_t) * hash_table_cell_splits_size);
97
98 // compute number of points that map to each hash
99 for (int i = 0; i < batch_size; ++i) {
100 const size_t hash_table_size =
101 hash_table_splits[i + 1] - hash_table_splits[i];
102 const size_t first_cell_idx = hash_table_splits[i];
103 tbb::parallel_for(
104 tbb::blocked_range<int64_t>(points_row_splits[i],
105 points_row_splits[i + 1]),
106 [&](const tbb::blocked_range<int64_t>& r) {
107 for (int64_t i = r.begin(); i != r.end(); ++i) {
108 Vec3_t pos(points + 3 * i);
109
110 auto voxel_index =
111 ComputeVoxelIndex(pos, inv_voxel_size);
112 size_t hash =
113 SpatialHash(voxel_index) % hash_table_size;
114
115 // note the +1 because we want the first
116 // element to be 0
117 core::AtomicFetchAddRelaxed(
118 &hash_table_cell_splits[first_cell_idx + hash +
119 1],
120 1);
121 }
122 });
123 }
124 InclusivePrefixSum(&hash_table_cell_splits[0],
125 &hash_table_cell_splits[hash_table_cell_splits_size],
126 &hash_table_cell_splits[0]);
127
128 std::vector<uint32_t> count_tmp(hash_table_cell_splits_size - 1, 0);
129
130 // now compute the indices for hash_table_index
131 for (int i = 0; i < batch_size; ++i) {
132 const size_t hash_table_size =
133 hash_table_splits[i + 1] - hash_table_splits[i];
134 const size_t first_cell_idx = hash_table_splits[i];
135 tbb::parallel_for(
136 tbb::blocked_range<size_t>(points_row_splits[i],
137 points_row_splits[i + 1]),
138 [&](const tbb::blocked_range<size_t>& r) {
139 for (size_t i = r.begin(); i != r.end(); ++i) {
140 Vec3_t pos(points + 3 * i);
141
142 auto voxel_index =
143 ComputeVoxelIndex(pos, inv_voxel_size);
144 size_t hash =
145 SpatialHash(voxel_index) % hash_table_size;
146
147 hash_table_index
148 [hash_table_cell_splits[hash + first_cell_idx] +
149 core::AtomicFetchAddRelaxed(
150 &count_tmp[hash + first_cell_idx],
151 1)] = i;
152 }
153 });
154 }
155}
156
174template <int METRIC, class TDerived, int VECSIZE>
175Eigen::Array<typename TDerived::Scalar, VECSIZE, 1> NeighborsDist(
176 const Eigen::ArrayBase<TDerived>& p,
177 const Eigen::Array<typename TDerived::Scalar, VECSIZE, 3>& points) {
178 typedef Eigen::Array<typename TDerived::Scalar, VECSIZE, 1> VecN_t;
179 VecN_t dist;
180
181 dist.setZero();
182 if (METRIC == Linf) {
183 dist = (points.rowwise() - p.transpose()).abs().rowwise().maxCoeff();
184 } else if (METRIC == L1) {
185 dist = (points.rowwise() - p.transpose()).abs().rowwise().sum();
186 } else {
187 dist = (points.rowwise() - p.transpose()).square().rowwise().sum();
188 }
189 return dist;
190}
191
194template <class T,
195 class OUTPUT_ALLOCATOR,
196 int METRIC,
197 bool IGNORE_QUERY_POINT,
198 bool RETURN_DISTANCES>
199void _FixedRadiusSearchCPU(int64_t* query_neighbors_row_splits,
200 size_t num_points,
201 const T* const points,
202 size_t num_queries,
203 const T* const queries,
204 const T radius,
205 const size_t points_row_splits_size,
206 const int64_t* const points_row_splits,
207 const size_t queries_row_splits_size,
208 const int64_t* const queries_row_splits,
209 const uint32_t* const hash_table_splits,
210 const size_t hash_table_cell_splits_size,
211 const uint32_t* const hash_table_cell_splits,
212 const uint32_t* const hash_table_index,
213 OUTPUT_ALLOCATOR& output_allocator) {
214 using namespace open3d::utility;
215
216// number of elements for vectorization
217#define VECSIZE 8
218 typedef MiniVec<T, 3> Vec3_t;
219 typedef Eigen::Array<T, VECSIZE, 1> Vec_t;
220 typedef Eigen::Array<int32_t, VECSIZE, 1> Veci_t;
221
222 typedef Eigen::Array<T, 3, 1> Pos_t;
223 typedef Eigen::Array<T, VECSIZE, 3> Poslist_t;
224 typedef Eigen::Array<bool, VECSIZE, 1> Result_t;
225
226 const int batch_size = points_row_splits_size - 1;
227
228 // return empty output arrays if there are no points
229 if (num_points == 0 || num_queries == 0) {
230 std::fill(query_neighbors_row_splits,
231 query_neighbors_row_splits + num_queries + 1, 0);
232 int32_t* indices_ptr;
233 output_allocator.AllocIndices(&indices_ptr, 0);
234
235 T* distances_ptr;
236 output_allocator.AllocDistances(&distances_ptr, 0);
237
238 return;
239 }
240
241 // use squared radius for L2 to avoid sqrt
242 const T threshold = (METRIC == L2 ? radius * radius : radius);
243
244 const T voxel_size = 2 * radius;
245 const T inv_voxel_size = 1 / voxel_size;
246
247 // counts the number of indices we have to return. This is the number of all
248 // neighbors we find.
249 size_t num_indices = 0;
250
251 // count the number of neighbors for all query points and update num_indices
252 // and populate query_neighbors_row_splits with the number of neighbors
253 // for each query point
254 for (int i = 0; i < batch_size; ++i) {
255 const size_t hash_table_size =
256 hash_table_splits[i + 1] - hash_table_splits[i];
257 const size_t first_cell_idx = hash_table_splits[i];
258 tbb::parallel_for(
259 tbb::blocked_range<size_t>(queries_row_splits[i],
260 queries_row_splits[i + 1]),
261 [&](const tbb::blocked_range<size_t>& r) {
262 size_t num_indices_local = 0;
263 for (size_t i = r.begin(); i != r.end(); ++i) {
264 size_t neighbors_count = 0;
265
266 Vec3_t pos(queries + i * 3);
267
268 std::set<size_t> bins_to_visit;
269
270 auto voxel_index =
271 ComputeVoxelIndex(pos, inv_voxel_size);
272 size_t hash =
273 SpatialHash(voxel_index) % hash_table_size;
274
275 bins_to_visit.insert(first_cell_idx + hash);
276
277 for (int dz = -1; dz <= 1; dz += 2)
278 for (int dy = -1; dy <= 1; dy += 2)
279 for (int dx = -1; dx <= 1; dx += 2) {
280 Vec3_t p =
281 pos + radius * Vec3_t(T(dx), T(dy),
282 T(dz));
283 voxel_index = ComputeVoxelIndex(
284 p, inv_voxel_size);
285 hash = SpatialHash(voxel_index) %
286 hash_table_size;
287 bins_to_visit.insert(first_cell_idx + hash);
288 }
289
290 Poslist_t xyz;
291 int vec_i = 0;
292
293 for (size_t bin : bins_to_visit) {
294 size_t begin_idx = hash_table_cell_splits[bin];
295 size_t end_idx = hash_table_cell_splits[bin + 1];
296
297 for (size_t j = begin_idx; j < end_idx; ++j) {
298 uint32_t idx = hash_table_index[j];
299 if (IGNORE_QUERY_POINT) {
300 if (points[idx * 3 + 0] == pos[0] &&
301 points[idx * 3 + 1] == pos[1] &&
302 points[idx * 3 + 2] == pos[2])
303 continue;
304 }
305 xyz(vec_i, 0) = points[idx * 3 + 0];
306 xyz(vec_i, 1) = points[idx * 3 + 1];
307 xyz(vec_i, 2) = points[idx * 3 + 2];
308 ++vec_i;
309 if (VECSIZE == vec_i) {
310 Pos_t pos_arr(pos[0], pos[1], pos[2]);
311 Vec_t dist = NeighborsDist<METRIC, Pos_t,
312 VECSIZE>(pos_arr,
313 xyz);
314 Result_t test_result = dist <= threshold;
315 neighbors_count += test_result.count();
316 vec_i = 0;
317 }
318 }
319 }
320 // process the tail
321 if (vec_i) {
322 Pos_t pos_arr(pos[0], pos[1], pos[2]);
323 Vec_t dist = NeighborsDist<METRIC, Pos_t, VECSIZE>(
324 pos_arr, xyz);
325 Result_t test_result = dist <= threshold;
326 for (int k = 0; k < vec_i; ++k) {
327 neighbors_count += int(test_result(k));
328 }
329 vec_i = 0;
330 }
331 num_indices_local += neighbors_count;
332 // note the +1
333 query_neighbors_row_splits[i + 1] = neighbors_count;
334 }
335
337 num_indices_local);
338 });
339 }
340
341 // Allocate output arrays
342 // output for the indices to the neighbors
343 int32_t* indices_ptr;
344 output_allocator.AllocIndices(&indices_ptr, num_indices);
345
346 // output for the distances
347 T* distances_ptr;
348 if (RETURN_DISTANCES)
349 output_allocator.AllocDistances(&distances_ptr, num_indices);
350 else
351 output_allocator.AllocDistances(&distances_ptr, 0);
352
353 query_neighbors_row_splits[0] = 0;
354 InclusivePrefixSum(query_neighbors_row_splits + 1,
355 query_neighbors_row_splits + num_queries + 1,
356 query_neighbors_row_splits + 1);
357
358 // now populate the indices_ptr and distances_ptr array
359 for (int i = 0; i < batch_size; ++i) {
360 const size_t hash_table_size =
361 hash_table_splits[i + 1] - hash_table_splits[i];
362 const size_t first_cell_idx = hash_table_splits[i];
363 tbb::parallel_for(
364 tbb::blocked_range<size_t>(queries_row_splits[i],
365 queries_row_splits[i + 1]),
366 [&](const tbb::blocked_range<size_t>& r) {
367 for (size_t i = r.begin(); i != r.end(); ++i) {
368 size_t neighbors_count = 0;
369
370 size_t indices_offset = query_neighbors_row_splits[i];
371
372 Vec3_t pos(queries[i * 3 + 0], queries[i * 3 + 1],
373 queries[i * 3 + 2]);
374
375 std::set<size_t> bins_to_visit;
376
377 auto voxel_index =
378 ComputeVoxelIndex(pos, inv_voxel_size);
379 size_t hash =
380 SpatialHash(voxel_index) % hash_table_size;
381
382 bins_to_visit.insert(first_cell_idx + hash);
383
384 for (int dz = -1; dz <= 1; dz += 2)
385 for (int dy = -1; dy <= 1; dy += 2)
386 for (int dx = -1; dx <= 1; dx += 2) {
387 Vec3_t p =
388 pos + radius * Vec3_t(T(dx), T(dy),
389 T(dz));
390 voxel_index = ComputeVoxelIndex(
391 p, inv_voxel_size);
392 hash = SpatialHash(voxel_index) %
393 hash_table_size;
394 bins_to_visit.insert(first_cell_idx + hash);
395 }
396
397 Poslist_t xyz;
398 Veci_t idx_vec;
399 int vec_i = 0;
400
401 for (size_t bin : bins_to_visit) {
402 size_t begin_idx = hash_table_cell_splits[bin];
403 size_t end_idx = hash_table_cell_splits[bin + 1];
404
405 for (size_t j = begin_idx; j < end_idx; ++j) {
406 int64_t idx = hash_table_index[j];
407 if (IGNORE_QUERY_POINT) {
408 if (points[idx * 3 + 0] == pos[0] &&
409 points[idx * 3 + 1] == pos[1] &&
410 points[idx * 3 + 2] == pos[2])
411 continue;
412 }
413 xyz(vec_i, 0) = points[idx * 3 + 0];
414 xyz(vec_i, 1) = points[idx * 3 + 1];
415 xyz(vec_i, 2) = points[idx * 3 + 2];
416 idx_vec(vec_i) = idx;
417 ++vec_i;
418 if (VECSIZE == vec_i) {
419 Pos_t pos_arr(pos[0], pos[1], pos[2]);
420 Vec_t dist = NeighborsDist<METRIC, Pos_t,
421 VECSIZE>(pos_arr,
422 xyz);
423 Result_t test_result = dist <= threshold;
424 for (int k = 0; k < vec_i; ++k) {
425 if (test_result(k)) {
426 indices_ptr[indices_offset +
427 neighbors_count] =
428 idx_vec[k];
429 if (RETURN_DISTANCES) {
430 distances_ptr[indices_offset +
431 neighbors_count] =
432 dist[k];
433 }
434 }
435 neighbors_count += int(test_result(k));
436 }
437 vec_i = 0;
438 }
439 }
440 }
441 // process the tail
442 if (vec_i) {
443 Pos_t pos_arr(pos[0], pos[1], pos[2]);
444 Vec_t dist = NeighborsDist<METRIC, Pos_t, VECSIZE>(
445 pos_arr, xyz);
446 Result_t test_result = dist <= threshold;
447 for (int k = 0; k < vec_i; ++k) {
448 if (test_result(k)) {
449 indices_ptr[indices_offset +
450 neighbors_count] = idx_vec[k];
451 if (RETURN_DISTANCES) {
452 distances_ptr[indices_offset +
453 neighbors_count] =
454 dist[k];
455 }
456 }
457 neighbors_count += int(test_result(k));
458 }
459 vec_i = 0;
460 }
461 }
462 });
463 }
464#undef VECSIZE
465}
466
467} // namespace
468
547template <class T, class OUTPUT_ALLOCATOR>
548void FixedRadiusSearchCPU(int64_t* query_neighbors_row_splits,
549 const size_t num_points,
550 const T* const points,
551 const size_t num_queries,
552 const T* const queries,
553 const T radius,
554 const size_t points_row_splits_size,
555 const int64_t* const points_row_splits,
556 const size_t queries_row_splits_size,
557 const int64_t* const queries_row_splits,
558 const uint32_t* const hash_table_splits,
559 const size_t hash_table_cell_splits_size,
560 const uint32_t* const hash_table_cell_splits,
561 const uint32_t* const hash_table_index,
562 const Metric metric,
563 const bool ignore_query_point,
564 const bool return_distances,
565 OUTPUT_ALLOCATOR& output_allocator) {
566 // Dispatch all template parameter combinations
567
568#define FN_PARAMETERS \
569 query_neighbors_row_splits, num_points, points, num_queries, queries, \
570 radius, points_row_splits_size, points_row_splits, \
571 queries_row_splits_size, queries_row_splits, hash_table_splits, \
572 hash_table_cell_splits_size, hash_table_cell_splits, \
573 hash_table_index, output_allocator
574
575#define CALL_TEMPLATE(METRIC, IGNORE_QUERY_POINT, RETURN_DISTANCES) \
576 if (METRIC == metric && IGNORE_QUERY_POINT == ignore_query_point && \
577 RETURN_DISTANCES == return_distances) \
578 _FixedRadiusSearchCPU<T, OUTPUT_ALLOCATOR, METRIC, IGNORE_QUERY_POINT, \
579 RETURN_DISTANCES>(FN_PARAMETERS);
580
581#define CALL_TEMPLATE2(METRIC) \
582 CALL_TEMPLATE(METRIC, true, true) \
583 CALL_TEMPLATE(METRIC, true, false) \
584 CALL_TEMPLATE(METRIC, false, true) \
585 CALL_TEMPLATE(METRIC, false, false)
586
587#define CALL_TEMPLATE3 \
588 CALL_TEMPLATE2(L1) \
589 CALL_TEMPLATE2(L2) \
590 CALL_TEMPLATE2(Linf)
591
593
594#undef CALL_TEMPLATE
595#undef CALL_TEMPLATE2
596#undef CALL_TEMPLATE3
597#undef FN_PARAMETERS
598}
599
600} // namespace impl
601} // namespace nns
602} // namespace core
603} // namespace open3d
#define VECSIZE
#define CALL_TEMPLATE3
int points
Definition: FilePCD.cpp:73
void FixedRadiusSearchCPU(int64_t *query_neighbors_row_splits, const size_t num_points, const T *const points, const size_t num_queries, const T *const queries, const T radius, const size_t points_row_splits_size, const int64_t *const points_row_splits, const size_t queries_row_splits_size, const int64_t *const queries_row_splits, const uint32_t *const hash_table_splits, const size_t hash_table_cell_splits_size, const uint32_t *const hash_table_cell_splits, const uint32_t *const hash_table_index, const Metric metric, const bool ignore_query_point, const bool return_distances, OUTPUT_ALLOCATOR &output_allocator)
Definition: FixedRadiusSearchImpl.h:548
Metric
Supported metrics.
Definition: NeighborSearchCommon.h:38
@ Linf
Definition: NeighborSearchCommon.h:38
@ L1
Definition: NeighborSearchCommon.h:38
@ L2
Definition: NeighborSearchCommon.h:38
void BuildSpatialHashTableCPU(const Tensor &points, double radius, const Tensor &points_row_splits, const Tensor &hash_table_splits, Tensor &hash_table_index, Tensor &hash_table_cell_splits)
Definition: FixedRadiusSearchOps.cpp:40
uint32_t AtomicFetchAddRelaxed(uint32_t *address, uint32_t val)
Definition: Atomic.h:44
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 timeout_in_ms capture_handle capture_handle capture_handle image_handle temperature_c k4a_image_t image_handle uint8_t image_handle image_handle image_handle image_handle uint32_t
Definition: K4aPlugin.cpp:567
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 timeout_in_ms capture_handle capture_handle capture_handle image_handle temperature_c int
Definition: K4aPlugin.cpp:493
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 uint64_t
Definition: K4aPlugin.cpp:362
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:414
Definition: Dispatch.h:110
void InclusivePrefixSum(const Tin *first, const Tin *last, Tout *out)
Definition: ParallelScan.h:67
Definition: PinholeCameraIntrinsic.cpp:35
Definition: MiniVec.h:43