Open3D (C++ API)  0.16.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 TIndex,
196 class OUTPUT_ALLOCATOR,
197 int METRIC,
198 bool IGNORE_QUERY_POINT,
199 bool RETURN_DISTANCES>
200void _FixedRadiusSearchCPU(int64_t* query_neighbors_row_splits,
201 size_t num_points,
202 const T* const points,
203 size_t num_queries,
204 const T* const queries,
205 const T radius,
206 const size_t points_row_splits_size,
207 const int64_t* const points_row_splits,
208 const size_t queries_row_splits_size,
209 const int64_t* const queries_row_splits,
210 const uint32_t* const hash_table_splits,
211 const size_t hash_table_cell_splits_size,
212 const uint32_t* const hash_table_cell_splits,
213 const uint32_t* const hash_table_index,
214 OUTPUT_ALLOCATOR& output_allocator) {
215 using namespace open3d::utility;
216
217// number of elements for vectorization
218#define VECSIZE 8
219 typedef MiniVec<T, 3> Vec3_t;
220 typedef Eigen::Array<T, VECSIZE, 1> Vec_t;
221 typedef Eigen::Array<TIndex, VECSIZE, 1> Veci_t;
222
223 typedef Eigen::Array<T, 3, 1> Pos_t;
224 typedef Eigen::Array<T, VECSIZE, 3> Poslist_t;
225 typedef Eigen::Array<bool, VECSIZE, 1> Result_t;
226
227 const int batch_size = points_row_splits_size - 1;
228
229 // return empty output arrays if there are no points
230 if (num_points == 0 || num_queries == 0) {
231 std::fill(query_neighbors_row_splits,
232 query_neighbors_row_splits + num_queries + 1, 0);
233 TIndex* indices_ptr;
234 output_allocator.AllocIndices(&indices_ptr, 0);
235
236 T* distances_ptr;
237 output_allocator.AllocDistances(&distances_ptr, 0);
238
239 return;
240 }
241
242 // use squared radius for L2 to avoid sqrt
243 const T threshold = (METRIC == L2 ? radius * radius : radius);
244
245 const T voxel_size = 2 * radius;
246 const T inv_voxel_size = 1 / voxel_size;
247
248 // counts the number of indices we have to return. This is the number of all
249 // neighbors we find.
250 size_t num_indices = 0;
251
252 // count the number of neighbors for all query points and update num_indices
253 // and populate query_neighbors_row_splits with the number of neighbors
254 // for each query point
255 for (int i = 0; i < batch_size; ++i) {
256 const size_t hash_table_size =
257 hash_table_splits[i + 1] - hash_table_splits[i];
258 const size_t first_cell_idx = hash_table_splits[i];
259 tbb::parallel_for(
260 tbb::blocked_range<size_t>(queries_row_splits[i],
261 queries_row_splits[i + 1]),
262 [&](const tbb::blocked_range<size_t>& r) {
263 size_t num_indices_local = 0;
264 for (size_t i = r.begin(); i != r.end(); ++i) {
265 size_t neighbors_count = 0;
266
267 Vec3_t pos(queries + i * 3);
268
269 std::set<size_t> bins_to_visit;
270
271 auto voxel_index =
272 ComputeVoxelIndex(pos, inv_voxel_size);
273 size_t hash =
274 SpatialHash(voxel_index) % hash_table_size;
275
276 bins_to_visit.insert(first_cell_idx + hash);
277
278 for (int dz = -1; dz <= 1; dz += 2)
279 for (int dy = -1; dy <= 1; dy += 2)
280 for (int dx = -1; dx <= 1; dx += 2) {
281 Vec3_t p =
282 pos + radius * Vec3_t(T(dx), T(dy),
283 T(dz));
284 voxel_index = ComputeVoxelIndex(
285 p, inv_voxel_size);
286 hash = SpatialHash(voxel_index) %
287 hash_table_size;
288 bins_to_visit.insert(first_cell_idx + hash);
289 }
290
291 Poslist_t xyz;
292 int vec_i = 0;
293
294 for (size_t bin : bins_to_visit) {
295 size_t begin_idx = hash_table_cell_splits[bin];
296 size_t end_idx = hash_table_cell_splits[bin + 1];
297
298 for (size_t j = begin_idx; j < end_idx; ++j) {
299 uint32_t idx = hash_table_index[j];
300 if (IGNORE_QUERY_POINT) {
301 if (points[idx * 3 + 0] == pos[0] &&
302 points[idx * 3 + 1] == pos[1] &&
303 points[idx * 3 + 2] == pos[2])
304 continue;
305 }
306 xyz(vec_i, 0) = points[idx * 3 + 0];
307 xyz(vec_i, 1) = points[idx * 3 + 1];
308 xyz(vec_i, 2) = points[idx * 3 + 2];
309 ++vec_i;
310 if (VECSIZE == vec_i) {
311 Pos_t pos_arr(pos[0], pos[1], pos[2]);
312 Vec_t dist = NeighborsDist<METRIC, Pos_t,
313 VECSIZE>(pos_arr,
314 xyz);
315 Result_t test_result = dist <= threshold;
316 neighbors_count += test_result.count();
317 vec_i = 0;
318 }
319 }
320 }
321 // process the tail
322 if (vec_i) {
323 Pos_t pos_arr(pos[0], pos[1], pos[2]);
324 Vec_t dist = NeighborsDist<METRIC, Pos_t, VECSIZE>(
325 pos_arr, xyz);
326 Result_t test_result = dist <= threshold;
327 for (int k = 0; k < vec_i; ++k) {
328 neighbors_count += int(test_result(k));
329 }
330 vec_i = 0;
331 }
332 num_indices_local += neighbors_count;
333 // note the +1
334 query_neighbors_row_splits[i + 1] = neighbors_count;
335 }
336
338 num_indices_local);
339 });
340 }
341
342 // Allocate output arrays
343 // output for the indices to the neighbors
344 TIndex* indices_ptr;
345 output_allocator.AllocIndices(&indices_ptr, num_indices);
346
347 // output for the distances
348 T* distances_ptr;
349 if (RETURN_DISTANCES)
350 output_allocator.AllocDistances(&distances_ptr, num_indices);
351 else
352 output_allocator.AllocDistances(&distances_ptr, 0);
353
354 query_neighbors_row_splits[0] = 0;
355 InclusivePrefixSum(query_neighbors_row_splits + 1,
356 query_neighbors_row_splits + num_queries + 1,
357 query_neighbors_row_splits + 1);
358
359 // now populate the indices_ptr and distances_ptr array
360 for (int i = 0; i < batch_size; ++i) {
361 const size_t hash_table_size =
362 hash_table_splits[i + 1] - hash_table_splits[i];
363 const size_t first_cell_idx = hash_table_splits[i];
364 tbb::parallel_for(
365 tbb::blocked_range<size_t>(queries_row_splits[i],
366 queries_row_splits[i + 1]),
367 [&](const tbb::blocked_range<size_t>& r) {
368 for (size_t i = r.begin(); i != r.end(); ++i) {
369 size_t neighbors_count = 0;
370
371 size_t indices_offset = query_neighbors_row_splits[i];
372
373 Vec3_t pos(queries[i * 3 + 0], queries[i * 3 + 1],
374 queries[i * 3 + 2]);
375
376 std::set<size_t> bins_to_visit;
377
378 auto voxel_index =
379 ComputeVoxelIndex(pos, inv_voxel_size);
380 size_t hash =
381 SpatialHash(voxel_index) % hash_table_size;
382
383 bins_to_visit.insert(first_cell_idx + hash);
384
385 for (int dz = -1; dz <= 1; dz += 2)
386 for (int dy = -1; dy <= 1; dy += 2)
387 for (int dx = -1; dx <= 1; dx += 2) {
388 Vec3_t p =
389 pos + radius * Vec3_t(T(dx), T(dy),
390 T(dz));
391 voxel_index = ComputeVoxelIndex(
392 p, inv_voxel_size);
393 hash = SpatialHash(voxel_index) %
394 hash_table_size;
395 bins_to_visit.insert(first_cell_idx + hash);
396 }
397
398 Poslist_t xyz;
399 Veci_t idx_vec;
400 int vec_i = 0;
401
402 for (size_t bin : bins_to_visit) {
403 size_t begin_idx = hash_table_cell_splits[bin];
404 size_t end_idx = hash_table_cell_splits[bin + 1];
405
406 for (size_t j = begin_idx; j < end_idx; ++j) {
407 int64_t idx = hash_table_index[j];
408 if (IGNORE_QUERY_POINT) {
409 if (points[idx * 3 + 0] == pos[0] &&
410 points[idx * 3 + 1] == pos[1] &&
411 points[idx * 3 + 2] == pos[2])
412 continue;
413 }
414 xyz(vec_i, 0) = points[idx * 3 + 0];
415 xyz(vec_i, 1) = points[idx * 3 + 1];
416 xyz(vec_i, 2) = points[idx * 3 + 2];
417 idx_vec(vec_i) = idx;
418 ++vec_i;
419 if (VECSIZE == vec_i) {
420 Pos_t pos_arr(pos[0], pos[1], pos[2]);
421 Vec_t dist = NeighborsDist<METRIC, Pos_t,
422 VECSIZE>(pos_arr,
423 xyz);
424 Result_t test_result = dist <= threshold;
425 for (int k = 0; k < vec_i; ++k) {
426 if (test_result(k)) {
427 indices_ptr[indices_offset +
428 neighbors_count] =
429 idx_vec[k];
430 if (RETURN_DISTANCES) {
431 distances_ptr[indices_offset +
432 neighbors_count] =
433 dist[k];
434 }
435 }
436 neighbors_count += int(test_result(k));
437 }
438 vec_i = 0;
439 }
440 }
441 }
442 // process the tail
443 if (vec_i) {
444 Pos_t pos_arr(pos[0], pos[1], pos[2]);
445 Vec_t dist = NeighborsDist<METRIC, Pos_t, VECSIZE>(
446 pos_arr, xyz);
447 Result_t test_result = dist <= threshold;
448 for (int k = 0; k < vec_i; ++k) {
449 if (test_result(k)) {
450 indices_ptr[indices_offset +
451 neighbors_count] = idx_vec[k];
452 if (RETURN_DISTANCES) {
453 distances_ptr[indices_offset +
454 neighbors_count] =
455 dist[k];
456 }
457 }
458 neighbors_count += int(test_result(k));
459 }
460 vec_i = 0;
461 }
462 }
463 });
464 }
465#undef VECSIZE
466}
467
468} // namespace
469
548template <class T, class TIndex, class OUTPUT_ALLOCATOR>
549void FixedRadiusSearchCPU(int64_t* query_neighbors_row_splits,
550 const size_t num_points,
551 const T* const points,
552 const size_t num_queries,
553 const T* const queries,
554 const T radius,
555 const size_t points_row_splits_size,
556 const int64_t* const points_row_splits,
557 const size_t queries_row_splits_size,
558 const int64_t* const queries_row_splits,
559 const uint32_t* const hash_table_splits,
560 const size_t hash_table_cell_splits_size,
561 const uint32_t* const hash_table_cell_splits,
562 const uint32_t* const hash_table_index,
563 const Metric metric,
564 const bool ignore_query_point,
565 const bool return_distances,
566 OUTPUT_ALLOCATOR& output_allocator) {
567 // Dispatch all template parameter combinations
568
569#define FN_PARAMETERS \
570 query_neighbors_row_splits, num_points, points, num_queries, queries, \
571 radius, points_row_splits_size, points_row_splits, \
572 queries_row_splits_size, queries_row_splits, hash_table_splits, \
573 hash_table_cell_splits_size, hash_table_cell_splits, \
574 hash_table_index, output_allocator
575
576#define CALL_TEMPLATE(METRIC, IGNORE_QUERY_POINT, RETURN_DISTANCES) \
577 if (METRIC == metric && IGNORE_QUERY_POINT == ignore_query_point && \
578 RETURN_DISTANCES == return_distances) \
579 _FixedRadiusSearchCPU<T, TIndex, OUTPUT_ALLOCATOR, METRIC, \
580 IGNORE_QUERY_POINT, RETURN_DISTANCES>( \
581 FN_PARAMETERS);
582
583#define CALL_TEMPLATE2(METRIC) \
584 CALL_TEMPLATE(METRIC, true, true) \
585 CALL_TEMPLATE(METRIC, true, false) \
586 CALL_TEMPLATE(METRIC, false, true) \
587 CALL_TEMPLATE(METRIC, false, false)
588
589#define CALL_TEMPLATE3 \
590 CALL_TEMPLATE2(L1) \
591 CALL_TEMPLATE2(L2) \
592 CALL_TEMPLATE2(Linf)
593
595
596#undef CALL_TEMPLATE
597#undef CALL_TEMPLATE2
598#undef CALL_TEMPLATE3
599#undef FN_PARAMETERS
600}
601
602} // namespace impl
603} // namespace nns
604} // namespace core
605} // 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:549
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
Definition: Dispatch.h:110
void InclusivePrefixSum(const Tin *first, const Tin *last, Tout *out)
Definition: ParallelScan.h:90
Definition: PinholeCameraIntrinsic.cpp:35
Definition: MiniVec.h:43