Open3D (C++ API)  0.17.0
Loading...
Searching...
No Matches
InvertNeighborsList.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 "open3d/core/Atomic.h"
14
15namespace open3d {
16namespace ml {
17namespace impl {
18
70template <class TIndex, class TAttr>
71void InvertNeighborsListCPU(const TIndex* const inp_neighbors_index,
72 const TAttr* const inp_neighbors_attributes,
73 const int num_attributes_per_neighbor,
74 const int64_t* const inp_neighbors_row_splits,
75 const size_t inp_num_queries,
76 TIndex* out_neighbors_index,
77 TAttr* out_neighbors_attributes,
78 const size_t index_size,
79 int64_t* out_neighbors_row_splits,
80 const size_t out_num_queries) {
81 using namespace open3d::utility;
82
83 std::vector<uint32_t> tmp_neighbors_count(out_num_queries + 1, 0);
84
85 // count how often an idx appears in inp_neighbors_index
86 tbb::parallel_for(tbb::blocked_range<size_t>(0, index_size),
87 [&](const tbb::blocked_range<size_t>& r) {
88 for (size_t i = r.begin(); i != r.end(); ++i) {
89 TIndex idx = inp_neighbors_index[i];
91 &tmp_neighbors_count[idx + 1], 1);
92 }
93 });
94
95 InclusivePrefixSum(&tmp_neighbors_count[0],
96 &tmp_neighbors_count[tmp_neighbors_count.size()],
97 out_neighbors_row_splits);
98
99 memset(tmp_neighbors_count.data(), 0,
100 sizeof(uint32_t) * tmp_neighbors_count.size());
101
102 // fill the new index vector
103 tbb::parallel_for(
104 tbb::blocked_range<size_t>(0, inp_num_queries),
105 [&](const tbb::blocked_range<size_t>& r) {
106 for (size_t i = r.begin(); i != r.end(); ++i) {
107 TIndex query_idx = i;
108
109 size_t begin_idx = inp_neighbors_row_splits[i];
110 size_t end_idx = inp_neighbors_row_splits[i + 1];
111 for (size_t j = begin_idx; j < end_idx; ++j) {
112 TIndex neighbor_idx = inp_neighbors_index[j];
113
114 size_t list_offset =
115 out_neighbors_row_splits[neighbor_idx];
116 size_t item_offset = core::AtomicFetchAddRelaxed(
117 &tmp_neighbors_count[neighbor_idx], 1);
118 out_neighbors_index[list_offset + item_offset] =
119 query_idx;
120
121 if (inp_neighbors_attributes) {
122 TAttr* attr_ptr =
123 out_neighbors_attributes +
124 num_attributes_per_neighbor *
125 (list_offset + item_offset);
126 for (int attr_i = 0;
127 attr_i < num_attributes_per_neighbor;
128 ++attr_i) {
129 attr_ptr[attr_i] = inp_neighbors_attributes
130 [num_attributes_per_neighbor * j +
131 attr_i];
132 }
133 }
134 }
135 }
136 });
137}
138
139} // namespace impl
140} // namespace ml
141} // namespace open3d
uint32_t AtomicFetchAddRelaxed(uint32_t *address, uint32_t val)
Definition Atomic.h:25
void InvertNeighborsListCPU(const TIndex *const inp_neighbors_index, const TAttr *const inp_neighbors_attributes, const int num_attributes_per_neighbor, const int64_t *const inp_neighbors_row_splits, const size_t inp_num_queries, TIndex *out_neighbors_index, TAttr *out_neighbors_attributes, const size_t index_size, int64_t *out_neighbors_row_splits, const size_t out_num_queries)
Definition InvertNeighborsList.h:71
Definition Dispatch.h:91
Definition PinholeCameraIntrinsic.cpp:16