Open3D (C++ API)  0.17.0
Loading...
Searching...
No Matches
FixedRadiusSearchImpl.h
Go to the documentation of this file.
1// ----------------------------------------------------------------------------
2// - Open3D: www.open3d.org -
3// ----------------------------------------------------------------------------
4// Copyright (c) 2018-2023 www.open3d.org
5// SPDX-License-Identifier: MIT
6// ----------------------------------------------------------------------------
7
8#pragma once
9
10#include <tbb/parallel_for.h>
11
12#include <set>
13
14#include "open3d/core/Atomic.h"
19
20namespace open3d {
21namespace core {
22namespace nns {
23namespace impl {
24
25namespace {
26
59template <class T>
60void BuildSpatialHashTableCPU(const size_t num_points,
61 const T* const points,
62 const T radius,
63 const size_t points_row_splits_size,
64 const int64_t* points_row_splits,
65 const uint32_t* hash_table_splits,
66 const size_t hash_table_cell_splits_size,
67 uint32_t* hash_table_cell_splits,
68 uint32_t* hash_table_index) {
69 using namespace open3d::utility;
70 typedef MiniVec<T, 3> Vec3_t;
71
72 const int batch_size = points_row_splits_size - 1;
73 const T voxel_size = 2 * radius;
74 const T inv_voxel_size = 1 / voxel_size;
75
76 memset(&hash_table_cell_splits[0], 0,
77 sizeof(uint32_t) * hash_table_cell_splits_size);
78
79 // compute number of points that map to each hash
80 for (int i = 0; i < batch_size; ++i) {
81 const size_t hash_table_size =
82 hash_table_splits[i + 1] - hash_table_splits[i];
83 const size_t first_cell_idx = hash_table_splits[i];
84 tbb::parallel_for(
85 tbb::blocked_range<int64_t>(points_row_splits[i],
86 points_row_splits[i + 1]),
87 [&](const tbb::blocked_range<int64_t>& r) {
88 for (int64_t i = r.begin(); i != r.end(); ++i) {
89 Vec3_t pos(points + 3 * i);
90
91 auto voxel_index =
92 ComputeVoxelIndex(pos, inv_voxel_size);
93 size_t hash =
94 SpatialHash(voxel_index) % hash_table_size;
95
96 // note the +1 because we want the first
97 // element to be 0
99 &hash_table_cell_splits[first_cell_idx + hash +
100 1],
101 1);
102 }
103 });
104 }
105 InclusivePrefixSum(&hash_table_cell_splits[0],
106 &hash_table_cell_splits[hash_table_cell_splits_size],
107 &hash_table_cell_splits[0]);
108
109 std::vector<uint32_t> count_tmp(hash_table_cell_splits_size - 1, 0);
110
111 // now compute the indices for hash_table_index
112 for (int i = 0; i < batch_size; ++i) {
113 const size_t hash_table_size =
114 hash_table_splits[i + 1] - hash_table_splits[i];
115 const size_t first_cell_idx = hash_table_splits[i];
116 tbb::parallel_for(
117 tbb::blocked_range<size_t>(points_row_splits[i],
118 points_row_splits[i + 1]),
119 [&](const tbb::blocked_range<size_t>& r) {
120 for (size_t i = r.begin(); i != r.end(); ++i) {
121 Vec3_t pos(points + 3 * i);
122
123 auto voxel_index =
124 ComputeVoxelIndex(pos, inv_voxel_size);
125 size_t hash =
126 SpatialHash(voxel_index) % hash_table_size;
127
128 hash_table_index
129 [hash_table_cell_splits[hash + first_cell_idx] +
131 &count_tmp[hash + first_cell_idx],
132 1)] = i;
133 }
134 });
135 }
136}
137
155template <int METRIC, class TDerived, int VECSIZE>
156Eigen::Array<typename TDerived::Scalar, VECSIZE, 1> NeighborsDist(
157 const Eigen::ArrayBase<TDerived>& p,
158 const Eigen::Array<typename TDerived::Scalar, VECSIZE, 3>& points) {
159 typedef Eigen::Array<typename TDerived::Scalar, VECSIZE, 1> VecN_t;
160 VecN_t dist;
161
162 dist.setZero();
163 if (METRIC == Linf) {
164 dist = (points.rowwise() - p.transpose()).abs().rowwise().maxCoeff();
165 } else if (METRIC == L1) {
166 dist = (points.rowwise() - p.transpose()).abs().rowwise().sum();
167 } else {
168 dist = (points.rowwise() - p.transpose()).square().rowwise().sum();
169 }
170 return dist;
171}
172
175template <class T,
176 class TIndex,
177 class OUTPUT_ALLOCATOR,
178 int METRIC,
179 bool IGNORE_QUERY_POINT,
180 bool RETURN_DISTANCES>
181void _FixedRadiusSearchCPU(int64_t* query_neighbors_row_splits,
182 size_t num_points,
183 const T* const points,
184 size_t num_queries,
185 const T* const queries,
186 const T radius,
187 const size_t points_row_splits_size,
188 const int64_t* const points_row_splits,
189 const size_t queries_row_splits_size,
190 const int64_t* const queries_row_splits,
191 const uint32_t* const hash_table_splits,
192 const size_t hash_table_cell_splits_size,
193 const uint32_t* const hash_table_cell_splits,
194 const uint32_t* const hash_table_index,
195 OUTPUT_ALLOCATOR& output_allocator) {
196 using namespace open3d::utility;
197
198// number of elements for vectorization
199#define VECSIZE 8
200 typedef MiniVec<T, 3> Vec3_t;
201 typedef Eigen::Array<T, VECSIZE, 1> Vec_t;
202 typedef Eigen::Array<TIndex, VECSIZE, 1> Veci_t;
203
204 typedef Eigen::Array<T, 3, 1> Pos_t;
205 typedef Eigen::Array<T, VECSIZE, 3> Poslist_t;
206 typedef Eigen::Array<bool, VECSIZE, 1> Result_t;
207
208 const int batch_size = points_row_splits_size - 1;
209
210 // return empty output arrays if there are no points
211 if (num_points == 0 || num_queries == 0) {
212 std::fill(query_neighbors_row_splits,
213 query_neighbors_row_splits + num_queries + 1, 0);
214 TIndex* indices_ptr;
215 output_allocator.AllocIndices(&indices_ptr, 0);
216
217 T* distances_ptr;
218 output_allocator.AllocDistances(&distances_ptr, 0);
219
220 return;
221 }
222
223 // use squared radius for L2 to avoid sqrt
224 const T threshold = (METRIC == L2 ? radius * radius : radius);
225
226 const T voxel_size = 2 * radius;
227 const T inv_voxel_size = 1 / voxel_size;
228
229 // counts the number of indices we have to return. This is the number of all
230 // neighbors we find.
231 size_t num_indices = 0;
232
233 // count the number of neighbors for all query points and update num_indices
234 // and populate query_neighbors_row_splits with the number of neighbors
235 // for each query point
236 for (int i = 0; i < batch_size; ++i) {
237 const size_t hash_table_size =
238 hash_table_splits[i + 1] - hash_table_splits[i];
239 const size_t first_cell_idx = hash_table_splits[i];
240 tbb::parallel_for(
241 tbb::blocked_range<size_t>(queries_row_splits[i],
242 queries_row_splits[i + 1]),
243 [&](const tbb::blocked_range<size_t>& r) {
244 size_t num_indices_local = 0;
245 for (size_t i = r.begin(); i != r.end(); ++i) {
246 size_t neighbors_count = 0;
247
248 Vec3_t pos(queries + i * 3);
249
250 std::set<size_t> bins_to_visit;
251
252 auto voxel_index =
253 ComputeVoxelIndex(pos, inv_voxel_size);
254 size_t hash =
255 SpatialHash(voxel_index) % hash_table_size;
256
257 bins_to_visit.insert(first_cell_idx + hash);
258
259 for (int dz = -1; dz <= 1; dz += 2)
260 for (int dy = -1; dy <= 1; dy += 2)
261 for (int dx = -1; dx <= 1; dx += 2) {
262 Vec3_t p =
263 pos + radius * Vec3_t(T(dx), T(dy),
264 T(dz));
265 voxel_index = ComputeVoxelIndex(
266 p, inv_voxel_size);
267 hash = SpatialHash(voxel_index) %
268 hash_table_size;
269 bins_to_visit.insert(first_cell_idx + hash);
270 }
271
272 Poslist_t xyz;
273 int vec_i = 0;
274
275 for (size_t bin : bins_to_visit) {
276 size_t begin_idx = hash_table_cell_splits[bin];
277 size_t end_idx = hash_table_cell_splits[bin + 1];
278
279 for (size_t j = begin_idx; j < end_idx; ++j) {
280 uint32_t idx = hash_table_index[j];
281 if (IGNORE_QUERY_POINT) {
282 if (points[idx * 3 + 0] == pos[0] &&
283 points[idx * 3 + 1] == pos[1] &&
284 points[idx * 3 + 2] == pos[2])
285 continue;
286 }
287 xyz(vec_i, 0) = points[idx * 3 + 0];
288 xyz(vec_i, 1) = points[idx * 3 + 1];
289 xyz(vec_i, 2) = points[idx * 3 + 2];
290 ++vec_i;
291 if (VECSIZE == vec_i) {
292 Pos_t pos_arr(pos[0], pos[1], pos[2]);
293 Vec_t dist = NeighborsDist<METRIC, Pos_t,
294 VECSIZE>(pos_arr,
295 xyz);
296 Result_t test_result = dist <= threshold;
297 neighbors_count += test_result.count();
298 vec_i = 0;
299 }
300 }
301 }
302 // process the tail
303 if (vec_i) {
304 Pos_t pos_arr(pos[0], pos[1], pos[2]);
305 Vec_t dist = NeighborsDist<METRIC, Pos_t, VECSIZE>(
306 pos_arr, xyz);
307 Result_t test_result = dist <= threshold;
308 for (int k = 0; k < vec_i; ++k) {
309 neighbors_count += int(test_result(k));
310 }
311 vec_i = 0;
312 }
313 num_indices_local += neighbors_count;
314 // note the +1
315 query_neighbors_row_splits[i + 1] = neighbors_count;
316 }
317
318 core::AtomicFetchAddRelaxed((uint64_t*)&num_indices,
319 num_indices_local);
320 });
321 }
322
323 // Allocate output arrays
324 // output for the indices to the neighbors
325 TIndex* indices_ptr;
326 output_allocator.AllocIndices(&indices_ptr, num_indices);
327
328 // output for the distances
329 T* distances_ptr;
330 if (RETURN_DISTANCES)
331 output_allocator.AllocDistances(&distances_ptr, num_indices);
332 else
333 output_allocator.AllocDistances(&distances_ptr, 0);
334
335 query_neighbors_row_splits[0] = 0;
336 InclusivePrefixSum(query_neighbors_row_splits + 1,
337 query_neighbors_row_splits + num_queries + 1,
338 query_neighbors_row_splits + 1);
339
340 // now populate the indices_ptr and distances_ptr array
341 for (int i = 0; i < batch_size; ++i) {
342 const size_t hash_table_size =
343 hash_table_splits[i + 1] - hash_table_splits[i];
344 const size_t first_cell_idx = hash_table_splits[i];
345 tbb::parallel_for(
346 tbb::blocked_range<size_t>(queries_row_splits[i],
347 queries_row_splits[i + 1]),
348 [&](const tbb::blocked_range<size_t>& r) {
349 for (size_t i = r.begin(); i != r.end(); ++i) {
350 size_t neighbors_count = 0;
351
352 size_t indices_offset = query_neighbors_row_splits[i];
353
354 Vec3_t pos(queries[i * 3 + 0], queries[i * 3 + 1],
355 queries[i * 3 + 2]);
356
357 std::set<size_t> bins_to_visit;
358
359 auto voxel_index =
360 ComputeVoxelIndex(pos, inv_voxel_size);
361 size_t hash =
362 SpatialHash(voxel_index) % hash_table_size;
363
364 bins_to_visit.insert(first_cell_idx + hash);
365
366 for (int dz = -1; dz <= 1; dz += 2)
367 for (int dy = -1; dy <= 1; dy += 2)
368 for (int dx = -1; dx <= 1; dx += 2) {
369 Vec3_t p =
370 pos + radius * Vec3_t(T(dx), T(dy),
371 T(dz));
372 voxel_index = ComputeVoxelIndex(
373 p, inv_voxel_size);
374 hash = SpatialHash(voxel_index) %
375 hash_table_size;
376 bins_to_visit.insert(first_cell_idx + hash);
377 }
378
379 Poslist_t xyz;
380 Veci_t idx_vec;
381 int vec_i = 0;
382
383 for (size_t bin : bins_to_visit) {
384 size_t begin_idx = hash_table_cell_splits[bin];
385 size_t end_idx = hash_table_cell_splits[bin + 1];
386
387 for (size_t j = begin_idx; j < end_idx; ++j) {
388 int64_t idx = hash_table_index[j];
389 if (IGNORE_QUERY_POINT) {
390 if (points[idx * 3 + 0] == pos[0] &&
391 points[idx * 3 + 1] == pos[1] &&
392 points[idx * 3 + 2] == pos[2])
393 continue;
394 }
395 xyz(vec_i, 0) = points[idx * 3 + 0];
396 xyz(vec_i, 1) = points[idx * 3 + 1];
397 xyz(vec_i, 2) = points[idx * 3 + 2];
398 idx_vec(vec_i) = idx;
399 ++vec_i;
400 if (VECSIZE == vec_i) {
401 Pos_t pos_arr(pos[0], pos[1], pos[2]);
402 Vec_t dist = NeighborsDist<METRIC, Pos_t,
403 VECSIZE>(pos_arr,
404 xyz);
405 Result_t test_result = dist <= threshold;
406 for (int k = 0; k < vec_i; ++k) {
407 if (test_result(k)) {
408 indices_ptr[indices_offset +
409 neighbors_count] =
410 idx_vec[k];
411 if (RETURN_DISTANCES) {
412 distances_ptr[indices_offset +
413 neighbors_count] =
414 dist[k];
415 }
416 }
417 neighbors_count += int(test_result(k));
418 }
419 vec_i = 0;
420 }
421 }
422 }
423 // process the tail
424 if (vec_i) {
425 Pos_t pos_arr(pos[0], pos[1], pos[2]);
426 Vec_t dist = NeighborsDist<METRIC, Pos_t, VECSIZE>(
427 pos_arr, xyz);
428 Result_t test_result = dist <= threshold;
429 for (int k = 0; k < vec_i; ++k) {
430 if (test_result(k)) {
431 indices_ptr[indices_offset +
432 neighbors_count] = idx_vec[k];
433 if (RETURN_DISTANCES) {
434 distances_ptr[indices_offset +
435 neighbors_count] =
436 dist[k];
437 }
438 }
439 neighbors_count += int(test_result(k));
440 }
441 vec_i = 0;
442 }
443 }
444 });
445 }
446#undef VECSIZE
447}
448
449} // namespace
450
529template <class T, class TIndex, class OUTPUT_ALLOCATOR>
530void FixedRadiusSearchCPU(int64_t* query_neighbors_row_splits,
531 const size_t num_points,
532 const T* const points,
533 const size_t num_queries,
534 const T* const queries,
535 const T radius,
536 const size_t points_row_splits_size,
537 const int64_t* const points_row_splits,
538 const size_t queries_row_splits_size,
539 const int64_t* const queries_row_splits,
540 const uint32_t* const hash_table_splits,
541 const size_t hash_table_cell_splits_size,
542 const uint32_t* const hash_table_cell_splits,
543 const uint32_t* const hash_table_index,
544 const Metric metric,
545 const bool ignore_query_point,
546 const bool return_distances,
547 OUTPUT_ALLOCATOR& output_allocator) {
548 // Dispatch all template parameter combinations
549
550#define FN_PARAMETERS \
551 query_neighbors_row_splits, num_points, points, num_queries, queries, \
552 radius, points_row_splits_size, points_row_splits, \
553 queries_row_splits_size, queries_row_splits, hash_table_splits, \
554 hash_table_cell_splits_size, hash_table_cell_splits, \
555 hash_table_index, output_allocator
556
557#define CALL_TEMPLATE(METRIC, IGNORE_QUERY_POINT, RETURN_DISTANCES) \
558 if (METRIC == metric && IGNORE_QUERY_POINT == ignore_query_point && \
559 RETURN_DISTANCES == return_distances) \
560 _FixedRadiusSearchCPU<T, TIndex, OUTPUT_ALLOCATOR, METRIC, \
561 IGNORE_QUERY_POINT, RETURN_DISTANCES>( \
562 FN_PARAMETERS);
563
564#define CALL_TEMPLATE2(METRIC) \
565 CALL_TEMPLATE(METRIC, true, true) \
566 CALL_TEMPLATE(METRIC, true, false) \
567 CALL_TEMPLATE(METRIC, false, true) \
568 CALL_TEMPLATE(METRIC, false, false)
569
570#define CALL_TEMPLATE3 \
571 CALL_TEMPLATE2(L1) \
572 CALL_TEMPLATE2(L2) \
573 CALL_TEMPLATE2(Linf)
574
576
577#undef CALL_TEMPLATE
578#undef CALL_TEMPLATE2
579#undef CALL_TEMPLATE3
580#undef FN_PARAMETERS
581}
582
583} // namespace impl
584} // namespace nns
585} // namespace core
586} // namespace open3d
#define VECSIZE
#define CALL_TEMPLATE3
int points
Definition FilePCD.cpp:54
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:530
HOST_DEVICE utility::MiniVec< int, 3 > ComputeVoxelIndex(const TVecf &pos, const typename TVecf::Scalar_t &inv_voxel_size)
Definition NeighborSearchCommon.h:42
Metric
Supported metrics.
Definition NeighborSearchCommon.h:19
@ Linf
Definition NeighborSearchCommon.h:19
@ L1
Definition NeighborSearchCommon.h:19
@ L2
Definition NeighborSearchCommon.h:19
HOST_DEVICE size_t SpatialHash(int x, int y, int z)
Spatial hashing function for integer coordinates.
Definition NeighborSearchCommon.h:28
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:21
uint32_t AtomicFetchAddRelaxed(uint32_t *address, uint32_t val)
Definition Atomic.h:25
Definition Dispatch.h:91
Definition PinholeCameraIntrinsic.cpp:16
Definition MiniVec.h:24