FixedRadiusSearchImpl.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2023 www.open3d.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------
 
#pragma once
 
#include <tbb/parallel_for.h>
 
#include <set>
 
#include "open3d/core/Atomic.h"
#include "open3d/core/nns/NeighborSearchCommon.h"
#include "open3d/utility/Eigen.h"
#include "open3d/utility/Helper.h"
#include "open3d/utility/ParallelScan.h"
 
namespace open3d {
namespace core {
namespace nns {
namespace impl {
 
namespace {
 
template <class T>
void BuildSpatialHashTableCPU(const size_t num_points,
                              const T* const points,
                              const T radius,
                              const size_t points_row_splits_size,
                              const int64_t* points_row_splits,
                              const uint32_t* hash_table_splits,
                              const size_t hash_table_cell_splits_size,
                              uint32_t* hash_table_cell_splits,
                              uint32_t* hash_table_index) {
    using namespace open3d::utility;
    typedef MiniVec<T, 3> Vec3_t;
 
    const int batch_size = points_row_splits_size - 1;
    const T voxel_size = 2 * radius;
    const T inv_voxel_size = 1 / voxel_size;
 
    memset(&hash_table_cell_splits[0], 0,
           sizeof(uint32_t) * hash_table_cell_splits_size);
 
    // compute number of points that map to each hash
    for (int i = 0; i < batch_size; ++i) {
        const size_t hash_table_size =
                hash_table_splits[i + 1] - hash_table_splits[i];
        const size_t first_cell_idx = hash_table_splits[i];
        tbb::parallel_for(
                tbb::blocked_range<int64_t>(points_row_splits[i],
                                            points_row_splits[i + 1]),
                [&](const tbb::blocked_range<int64_t>& r) {
                    for (int64_t i = r.begin(); i != r.end(); ++i) {
                        Vec3_t pos(points + 3 * i);
 
                        auto voxel_index =
                                ComputeVoxelIndex(pos, inv_voxel_size);
                        size_t hash =
                                SpatialHash(voxel_index) % hash_table_size;
 
                        // note the +1 because we want the first
                        // element to be 0
                        core::AtomicFetchAddRelaxed(
                                &hash_table_cell_splits[first_cell_idx + hash +
                                                        1],
                                1);
                    }
                });
    }
    InclusivePrefixSum(&hash_table_cell_splits[0],
                       &hash_table_cell_splits[hash_table_cell_splits_size],
                       &hash_table_cell_splits[0]);
 
    std::vector<uint32_t> count_tmp(hash_table_cell_splits_size - 1, 0);
 
    // now compute the indices for hash_table_index
    for (int i = 0; i < batch_size; ++i) {
        const size_t hash_table_size =
                hash_table_splits[i + 1] - hash_table_splits[i];
        const size_t first_cell_idx = hash_table_splits[i];
        tbb::parallel_for(
                tbb::blocked_range<size_t>(points_row_splits[i],
                                           points_row_splits[i + 1]),
                [&](const tbb::blocked_range<size_t>& r) {
                    for (size_t i = r.begin(); i != r.end(); ++i) {
                        Vec3_t pos(points + 3 * i);
 
                        auto voxel_index =
                                ComputeVoxelIndex(pos, inv_voxel_size);
                        size_t hash =
                                SpatialHash(voxel_index) % hash_table_size;
 
                        hash_table_index
                                [hash_table_cell_splits[hash + first_cell_idx] +
                                 core::AtomicFetchAddRelaxed(
                                         &count_tmp[hash + first_cell_idx],
                                         1)] = i;
                    }
                });
    }
}
 
template <int METRIC, class TDerived, int VECSIZE>
Eigen::Array<typename TDerived::Scalar, VECSIZE, 1> NeighborsDist(
        const Eigen::ArrayBase<TDerived>& p,
        const Eigen::Array<typename TDerived::Scalar, VECSIZE, 3>& points) {
    typedef Eigen::Array<typename TDerived::Scalar, VECSIZE, 1> VecN_t;
    VecN_t dist;
 
    dist.setZero();
    if (METRIC == Linf) {
        dist = (points.rowwise() - p.transpose()).abs().rowwise().maxCoeff();
    } else if (METRIC == L1) {
        dist = (points.rowwise() - p.transpose()).abs().rowwise().sum();
    } else {
        dist = (points.rowwise() - p.transpose()).square().rowwise().sum();
    }
    return dist;
}
 
template <class T,
          class TIndex,
          class OUTPUT_ALLOCATOR,
          int METRIC,
          bool IGNORE_QUERY_POINT,
          bool RETURN_DISTANCES>
void _FixedRadiusSearchCPU(int64_t* query_neighbors_row_splits,
                           size_t num_points,
                           const T* const points,
                           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,
                           OUTPUT_ALLOCATOR& output_allocator) {
    using namespace open3d::utility;
 
// number of elements for vectorization
#define VECSIZE 8
    typedef MiniVec<T, 3> Vec3_t;
    typedef Eigen::Array<T, VECSIZE, 1> Vec_t;
    typedef Eigen::Array<TIndex, VECSIZE, 1> Veci_t;
 
    typedef Eigen::Array<T, 3, 1> Pos_t;
    typedef Eigen::Array<T, VECSIZE, 3> Poslist_t;
    typedef Eigen::Array<bool, VECSIZE, 1> Result_t;
 
    const int batch_size = points_row_splits_size - 1;
 
    // return empty output arrays if there are no points
    if (num_points == 0 || num_queries == 0) {
        std::fill(query_neighbors_row_splits,
                  query_neighbors_row_splits + num_queries + 1, 0);
        TIndex* indices_ptr;
        output_allocator.AllocIndices(&indices_ptr, 0);
 
        T* distances_ptr;
        output_allocator.AllocDistances(&distances_ptr, 0);
 
        return;
    }
 
    // use squared radius for L2 to avoid sqrt
    const T threshold = (METRIC == L2 ? radius * radius : radius);
 
    const T voxel_size = 2 * radius;
    const T inv_voxel_size = 1 / voxel_size;
 
    // counts the number of indices we have to return. This is the number of all
    // neighbors we find.
    size_t num_indices = 0;
 
    // count the number of neighbors for all query points and update num_indices
    // and populate query_neighbors_row_splits with the number of neighbors
    // for each query point
    for (int i = 0; i < batch_size; ++i) {
        const size_t hash_table_size =
                hash_table_splits[i + 1] - hash_table_splits[i];
        const size_t first_cell_idx = hash_table_splits[i];
        tbb::parallel_for(
                tbb::blocked_range<size_t>(queries_row_splits[i],
                                           queries_row_splits[i + 1]),
                [&](const tbb::blocked_range<size_t>& r) {
                    size_t num_indices_local = 0;
                    for (size_t i = r.begin(); i != r.end(); ++i) {
                        size_t neighbors_count = 0;
 
                        Vec3_t pos(queries + i * 3);
 
                        std::set<size_t> bins_to_visit;
 
                        auto voxel_index =
                                ComputeVoxelIndex(pos, inv_voxel_size);
                        size_t hash =
                                SpatialHash(voxel_index) % hash_table_size;
 
                        bins_to_visit.insert(first_cell_idx + hash);
 
                        for (int dz = -1; dz <= 1; dz += 2)
                            for (int dy = -1; dy <= 1; dy += 2)
                                for (int dx = -1; dx <= 1; dx += 2) {
                                    Vec3_t p =
                                            pos + radius * Vec3_t(T(dx), T(dy),
                                                                  T(dz));
                                    voxel_index = ComputeVoxelIndex(
                                            p, inv_voxel_size);
                                    hash = SpatialHash(voxel_index) %
                                           hash_table_size;
                                    bins_to_visit.insert(first_cell_idx + hash);
                                }
 
                        Poslist_t xyz;
                        int vec_i = 0;
 
                        for (size_t bin : bins_to_visit) {
                            size_t begin_idx = hash_table_cell_splits[bin];
                            size_t end_idx = hash_table_cell_splits[bin + 1];
 
                            for (size_t j = begin_idx; j < end_idx; ++j) {
                                uint32_t idx = hash_table_index[j];
                                if (IGNORE_QUERY_POINT) {
                                    if (points[idx * 3 + 0] == pos[0] &&
                                        points[idx * 3 + 1] == pos[1] &&
                                        points[idx * 3 + 2] == pos[2])
                                        continue;
                                }
                                xyz(vec_i, 0) = points[idx * 3 + 0];
                                xyz(vec_i, 1) = points[idx * 3 + 1];
                                xyz(vec_i, 2) = points[idx * 3 + 2];
                                ++vec_i;
                                if (VECSIZE == vec_i) {
                                    Pos_t pos_arr(pos[0], pos[1], pos[2]);
                                    Vec_t dist = NeighborsDist<METRIC, Pos_t,
                                                               VECSIZE>(pos_arr,
                                                                        xyz);
                                    Result_t test_result = dist <= threshold;
                                    neighbors_count += test_result.count();
                                    vec_i = 0;
                                }
                            }
                        }
                        // process the tail
                        if (vec_i) {
                            Pos_t pos_arr(pos[0], pos[1], pos[2]);
                            Vec_t dist = NeighborsDist<METRIC, Pos_t, VECSIZE>(
                                    pos_arr, xyz);
                            Result_t test_result = dist <= threshold;
                            for (int k = 0; k < vec_i; ++k) {
                                neighbors_count += int(test_result(k));
                            }
                            vec_i = 0;
                        }
                        num_indices_local += neighbors_count;
                        // note the +1
                        query_neighbors_row_splits[i + 1] = neighbors_count;
                    }
 
                    core::AtomicFetchAddRelaxed((uint64_t*)&num_indices,
                                                num_indices_local);
                });
    }
 
    // Allocate output arrays
    // output for the indices to the neighbors
    TIndex* indices_ptr;
    output_allocator.AllocIndices(&indices_ptr, num_indices);
 
    // output for the distances
    T* distances_ptr;
    if (RETURN_DISTANCES)
        output_allocator.AllocDistances(&distances_ptr, num_indices);
    else
        output_allocator.AllocDistances(&distances_ptr, 0);
 
    query_neighbors_row_splits[0] = 0;
    InclusivePrefixSum(query_neighbors_row_splits + 1,
                       query_neighbors_row_splits + num_queries + 1,
                       query_neighbors_row_splits + 1);
 
    // now populate the indices_ptr and distances_ptr array
    for (int i = 0; i < batch_size; ++i) {
        const size_t hash_table_size =
                hash_table_splits[i + 1] - hash_table_splits[i];
        const size_t first_cell_idx = hash_table_splits[i];
        tbb::parallel_for(
                tbb::blocked_range<size_t>(queries_row_splits[i],
                                           queries_row_splits[i + 1]),
                [&](const tbb::blocked_range<size_t>& r) {
                    for (size_t i = r.begin(); i != r.end(); ++i) {
                        size_t neighbors_count = 0;
 
                        size_t indices_offset = query_neighbors_row_splits[i];
 
                        Vec3_t pos(queries[i * 3 + 0], queries[i * 3 + 1],
                                   queries[i * 3 + 2]);
 
                        std::set<size_t> bins_to_visit;
 
                        auto voxel_index =
                                ComputeVoxelIndex(pos, inv_voxel_size);
                        size_t hash =
                                SpatialHash(voxel_index) % hash_table_size;
 
                        bins_to_visit.insert(first_cell_idx + hash);
 
                        for (int dz = -1; dz <= 1; dz += 2)
                            for (int dy = -1; dy <= 1; dy += 2)
                                for (int dx = -1; dx <= 1; dx += 2) {
                                    Vec3_t p =
                                            pos + radius * Vec3_t(T(dx), T(dy),
                                                                  T(dz));
                                    voxel_index = ComputeVoxelIndex(
                                            p, inv_voxel_size);
                                    hash = SpatialHash(voxel_index) %
                                           hash_table_size;
                                    bins_to_visit.insert(first_cell_idx + hash);
                                }
 
                        Poslist_t xyz;
                        Veci_t idx_vec;
                        int vec_i = 0;
 
                        for (size_t bin : bins_to_visit) {
                            size_t begin_idx = hash_table_cell_splits[bin];
                            size_t end_idx = hash_table_cell_splits[bin + 1];
 
                            for (size_t j = begin_idx; j < end_idx; ++j) {
                                int64_t idx = hash_table_index[j];
                                if (IGNORE_QUERY_POINT) {
                                    if (points[idx * 3 + 0] == pos[0] &&
                                        points[idx * 3 + 1] == pos[1] &&
                                        points[idx * 3 + 2] == pos[2])
                                        continue;
                                }
                                xyz(vec_i, 0) = points[idx * 3 + 0];
                                xyz(vec_i, 1) = points[idx * 3 + 1];
                                xyz(vec_i, 2) = points[idx * 3 + 2];
                                idx_vec(vec_i) = idx;
                                ++vec_i;
                                if (VECSIZE == vec_i) {
                                    Pos_t pos_arr(pos[0], pos[1], pos[2]);
                                    Vec_t dist = NeighborsDist<METRIC, Pos_t,
                                                               VECSIZE>(pos_arr,
                                                                        xyz);
                                    Result_t test_result = dist <= threshold;
                                    for (int k = 0; k < vec_i; ++k) {
                                        if (test_result(k)) {
                                            indices_ptr[indices_offset +
                                                        neighbors_count] =
                                                    idx_vec[k];
                                            if (RETURN_DISTANCES) {
                                                distances_ptr[indices_offset +
                                                              neighbors_count] =
                                                        dist[k];
                                            }
                                        }
                                        neighbors_count += int(test_result(k));
                                    }
                                    vec_i = 0;
                                }
                            }
                        }
                        // process the tail
                        if (vec_i) {
                            Pos_t pos_arr(pos[0], pos[1], pos[2]);
                            Vec_t dist = NeighborsDist<METRIC, Pos_t, VECSIZE>(
                                    pos_arr, xyz);
                            Result_t test_result = dist <= threshold;
                            for (int k = 0; k < vec_i; ++k) {
                                if (test_result(k)) {
                                    indices_ptr[indices_offset +
                                                neighbors_count] = idx_vec[k];
                                    if (RETURN_DISTANCES) {
                                        distances_ptr[indices_offset +
                                                      neighbors_count] =
                                                dist[k];
                                    }
                                }
                                neighbors_count += int(test_result(k));
                            }
                            vec_i = 0;
                        }
                    }
                });
    }
#undef VECSIZE
}
 
}  // namespace
 
template <class T, class TIndex, class OUTPUT_ALLOCATOR>
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) {
    // Dispatch all template parameter combinations
 
#define FN_PARAMETERS                                                       \
    query_neighbors_row_splits, num_points, points, num_queries, queries,   \
            radius, points_row_splits_size, points_row_splits,              \
            queries_row_splits_size, queries_row_splits, hash_table_splits, \
            hash_table_cell_splits_size, hash_table_cell_splits,            \
            hash_table_index, output_allocator
 
#define CALL_TEMPLATE(METRIC, IGNORE_QUERY_POINT, RETURN_DISTANCES)     \
    if (METRIC == metric && IGNORE_QUERY_POINT == ignore_query_point && \
        RETURN_DISTANCES == return_distances)                           \
        _FixedRadiusSearchCPU<T, TIndex, OUTPUT_ALLOCATOR, METRIC,      \
                              IGNORE_QUERY_POINT, RETURN_DISTANCES>(    \
                FN_PARAMETERS);
 
#define CALL_TEMPLATE2(METRIC)         \
    CALL_TEMPLATE(METRIC, true, true)  \
    CALL_TEMPLATE(METRIC, true, false) \
    CALL_TEMPLATE(METRIC, false, true) \
    CALL_TEMPLATE(METRIC, false, false)
 
#define CALL_TEMPLATE3 \
    CALL_TEMPLATE2(L1) \
    CALL_TEMPLATE2(L2) \
    CALL_TEMPLATE2(Linf)
 
    CALL_TEMPLATE3
 
#undef CALL_TEMPLATE
#undef CALL_TEMPLATE2
#undef CALL_TEMPLATE3
#undef FN_PARAMETERS
}
 
}  // namespace impl
}  // namespace nns
}  // namespace core
}  // namespace open3d