29 #include <tbb/parallel_for.h> 77 template <
class TReal,
class TIndex>
81 const size_t points_row_splits_size,
82 const int64_t* points_row_splits,
83 const TIndex* hash_table_splits,
84 const size_t hash_table_cell_splits_size,
85 TIndex* hash_table_cell_splits,
86 TIndex* hash_table_index) {
91 const int batch_size = points_row_splits_size - 1;
92 const TReal voxel_size = 2 * radius;
93 const TReal inv_voxel_size = 1 / voxel_size;
95 memset(&hash_table_cell_splits[0], 0,
96 sizeof(TIndex) * hash_table_cell_splits_size);
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];
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);
117 &hash_table_cell_splits[first_cell_idx + hash +
124 &hash_table_cell_splits[hash_table_cell_splits_size],
125 &hash_table_cell_splits[0]);
127 std::vector<uint32_t> count_tmp(hash_table_cell_splits_size - 1, 0);
130 for (
int i = 0; i < batch_size; ++i) {
131 const size_t hash_table_size =
132 hash_table_splits[i + 1] - hash_table_splits[i];
133 const size_t first_cell_idx = hash_table_splits[i];
135 tbb::blocked_range<size_t>(points_row_splits[i],
136 points_row_splits[i + 1]),
137 [&](
const tbb::blocked_range<size_t>& r) {
138 for (
size_t i = r.begin(); i != r.end(); ++i) {
139 Vec3_t pos(points + 3 * i);
147 [hash_table_cell_splits[hash + first_cell_idx] +
149 &count_tmp[hash + first_cell_idx],
173 template <
int METRIC,
class TDerived,
int VECSIZE>
174 Eigen::Array<typename TDerived::Scalar, VECSIZE, 1> NeighborsDist(
175 const Eigen::ArrayBase<TDerived>& p,
176 const Eigen::Array<typename TDerived::Scalar, VECSIZE, 3>& points) {
177 typedef Eigen::Array<typename TDerived::Scalar, VECSIZE, 1> VecN_t;
181 if (METRIC ==
Linf) {
182 dist = (points.rowwise() - p.transpose()).abs().rowwise().maxCoeff();
183 }
else if (METRIC ==
L1) {
184 dist = (points.rowwise() - p.transpose()).abs().rowwise().sum();
186 dist = (points.rowwise() - p.transpose()).square().rowwise().sum();
194 class OUTPUT_ALLOCATOR,
196 bool IGNORE_QUERY_POINT,
197 bool RETURN_DISTANCES>
198 void _FixedRadiusSearchCPU(int64_t* query_neighbors_row_splits,
200 const T*
const points,
202 const T*
const queries,
204 const size_t points_row_splits_size,
205 const int64_t*
const points_row_splits,
206 const size_t queries_row_splits_size,
207 const int64_t*
const queries_row_splits,
208 const uint32_t*
const hash_table_splits,
209 const size_t hash_table_cell_splits_size,
210 const uint32_t*
const hash_table_cell_splits,
211 const uint32_t*
const hash_table_index,
212 OUTPUT_ALLOCATOR& output_allocator) {
216 const int VECSIZE = 8;
218 typedef Eigen::Array<T, VECSIZE, 1> Vec_t;
219 typedef Eigen::Array<int32_t, VECSIZE, 1> Veci_t;
221 const int batch_size = points_row_splits_size - 1;
224 if (num_points == 0 || num_queries == 0) {
225 std::fill(query_neighbors_row_splits,
226 query_neighbors_row_splits + num_queries + 1, 0);
228 output_allocator.AllocIndices(&indices_ptr, 0);
231 output_allocator.AllocDistances(&distances_ptr, 0);
237 const T threshold = (METRIC ==
L2 ? radius * radius : radius);
239 const T voxel_size = 2 * radius;
240 const T inv_voxel_size = 1 / voxel_size;
244 size_t num_indices = 0;
249 for (
int i = 0; i < batch_size; ++i) {
250 const size_t hash_table_size =
251 hash_table_splits[i + 1] - hash_table_splits[i];
252 const size_t first_cell_idx = hash_table_splits[i];
254 tbb::blocked_range<size_t>(queries_row_splits[i],
255 queries_row_splits[i + 1]),
256 [&](
const tbb::blocked_range<size_t>& r) {
257 size_t num_indices_local = 0;
258 for (
size_t i = r.begin(); i != r.end(); ++i) {
259 size_t neighbors_count = 0;
261 Vec3_t pos(queries + i * 3);
263 std::set<size_t> bins_to_visit;
270 bins_to_visit.insert(first_cell_idx + hash);
272 for (
int dz = -1; dz <= 1; dz += 2)
273 for (
int dy = -1; dy <= 1; dy += 2)
274 for (
int dx = -1; dx <= 1; dx += 2) {
276 pos + radius * Vec3_t(T(dx), T(dy),
282 bins_to_visit.insert(first_cell_idx + hash);
285 Eigen::Array<T, VECSIZE, 3> xyz;
288 for (
size_t bin : bins_to_visit) {
289 size_t begin_idx = hash_table_cell_splits[bin];
290 size_t end_idx = hash_table_cell_splits[bin + 1];
292 for (
size_t j = begin_idx; j < end_idx; ++j) {
294 if (IGNORE_QUERY_POINT) {
295 if (points[idx * 3 + 0] == pos[0] &&
296 points[idx * 3 + 1] == pos[1] &&
297 points[idx * 3 + 2] == pos[2])
300 xyz(vec_i, 0) = points[idx * 3 + 0];
301 xyz(vec_i, 1) = points[idx * 3 + 1];
302 xyz(vec_i, 2) = points[idx * 3 + 2];
304 if (VECSIZE == vec_i) {
305 Eigen::Array<T, 3, 1> pos_arr(
306 pos[0], pos[1], pos[2]);
308 NeighborsDist<METRIC>(pos_arr, xyz);
309 Eigen::Array<bool, VECSIZE, 1> test_result =
311 neighbors_count += test_result.count();
318 Eigen::Array<T, 3, 1> pos_arr(pos[0], pos[1],
320 Eigen::Array<T, VECSIZE, 1> dist =
321 NeighborsDist<METRIC>(pos_arr, xyz);
322 Eigen::Array<bool, VECSIZE, 1> test_result =
324 for (
int k = 0; k < vec_i; ++k) {
325 neighbors_count +=
int(test_result(k));
329 num_indices_local += neighbors_count;
331 query_neighbors_row_splits[i + 1] = neighbors_count;
342 output_allocator.AllocIndices(&indices_ptr, num_indices);
346 if (RETURN_DISTANCES)
347 output_allocator.AllocDistances(&distances_ptr, num_indices);
349 output_allocator.AllocDistances(&distances_ptr, 0);
351 query_neighbors_row_splits[0] = 0;
353 query_neighbors_row_splits + num_queries + 1,
354 query_neighbors_row_splits + 1);
357 for (
int i = 0; i < batch_size; ++i) {
358 const size_t hash_table_size =
359 hash_table_splits[i + 1] - hash_table_splits[i];
360 const size_t first_cell_idx = hash_table_splits[i];
362 tbb::blocked_range<size_t>(queries_row_splits[i],
363 queries_row_splits[i + 1]),
364 [&](
const tbb::blocked_range<size_t>& r) {
365 for (
size_t i = r.begin(); i != r.end(); ++i) {
366 size_t neighbors_count = 0;
368 size_t indices_offset = query_neighbors_row_splits[i];
370 Vec3_t pos(queries[i * 3 + 0], queries[i * 3 + 1],
373 std::set<size_t> bins_to_visit;
380 bins_to_visit.insert(first_cell_idx + hash);
382 for (
int dz = -1; dz <= 1; dz += 2)
383 for (
int dy = -1; dy <= 1; dy += 2)
384 for (
int dx = -1; dx <= 1; dx += 2) {
386 pos + radius * Vec3_t(T(dx), T(dy),
392 bins_to_visit.insert(first_cell_idx + hash);
395 Eigen::Array<T, VECSIZE, 3> xyz;
399 for (
size_t bin : bins_to_visit) {
400 size_t begin_idx = hash_table_cell_splits[bin];
401 size_t end_idx = hash_table_cell_splits[bin + 1];
403 for (
size_t j = begin_idx; j < end_idx; ++j) {
405 if (IGNORE_QUERY_POINT) {
406 if (points[idx * 3 + 0] == pos[0] &&
407 points[idx * 3 + 1] == pos[1] &&
408 points[idx * 3 + 2] == pos[2])
411 xyz(vec_i, 0) = points[idx * 3 + 0];
412 xyz(vec_i, 1) = points[idx * 3 + 1];
413 xyz(vec_i, 2) = points[idx * 3 + 2];
414 idx_vec(vec_i) = idx;
416 if (VECSIZE == vec_i) {
417 Eigen::Array<T, 3, 1> pos_arr(
418 pos[0], pos[1], pos[2]);
419 Eigen::Array<T, VECSIZE, 1> dist =
420 NeighborsDist<METRIC>(pos_arr, xyz);
421 Eigen::Array<bool, VECSIZE, 1> test_result =
423 for (
int k = 0; k < vec_i; ++k) {
424 if (test_result(k)) {
425 indices_ptr[indices_offset +
428 if (RETURN_DISTANCES) {
429 distances_ptr[indices_offset +
434 neighbors_count +=
int(test_result(k));
442 Eigen::Array<T, 3, 1> pos_arr(pos[0], pos[1],
444 Eigen::Array<T, VECSIZE, 1> dist =
445 NeighborsDist<METRIC>(pos_arr, xyz);
446 Eigen::Array<bool, VECSIZE, 1> test_result =
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 +
458 neighbors_count +=
int(test_result(k));
547 template <
class T,
class OUTPUT_ALLOCATOR>
549 const size_t num_points,
550 const T*
const points,
551 const size_t num_queries,
552 const T*
const queries,
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,
563 const bool ignore_query_point,
564 const bool return_distances,
565 OUTPUT_ALLOCATOR& output_allocator) {
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 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); 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) 587 #define CALL_TEMPLATE3 \ 595 #undef CALL_TEMPLATE2 596 #undef CALL_TEMPLATE3
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: FixedRadiusSearch.h:548
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:557
void InclusivePrefixSum(const Tin *first, const Tin *last, Tout *out)
Definition: ParallelScan.h:67
void BuildSpatialHashTableCPU(const torch::Tensor &points, double radius, const torch::Tensor &points_row_splits, const std::vector< uint32_t > &hash_table_splits, torch::Tensor &hash_table_index, torch::Tensor &hash_table_cell_splits)
Definition: BuildSpatialHashTableOpKernel.cpp:33
Metric
Supported metrics.
Definition: NeighborSearchCommon.h:38
HOST_DEVICE utility::MiniVec< int, 3 > ComputeVoxelIndex(const TVecf &pos, const typename TVecf::Scalar_t &inv_voxel_size)
Definition: NeighborSearchCommon.h:61
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
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 uint64_t
Definition: K4aPlugin.cpp:352
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 timeout_in_ms capture_handle capture_handle capture_handle image_handle temperature_c int
Definition: K4aPlugin.cpp:479
int points
Definition: FilePCD.cpp:73
Definition: NeighborSearchCommon.h:38
Definition: PinholeCameraIntrinsic.cpp:35
HOST_DEVICE size_t SpatialHash(int x, int y, int z)
Spatial hashing function for integer coordinates.
Definition: NeighborSearchCommon.h:47
uint32_t AtomicFetchAddRelaxed(uint32_t *address, uint32_t val)
Definition: Atomic.h:40
Definition: Console.cpp:51