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) {
83 std::vector<uint32_t> tmp_neighbors_count(out_num_queries + 1, 0);
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);
95 InclusivePrefixSum(&tmp_neighbors_count[0],
96 &tmp_neighbors_count[tmp_neighbors_count.size()],
97 out_neighbors_row_splits);
99 memset(tmp_neighbors_count.data(), 0,
100 sizeof(uint32_t) * tmp_neighbors_count.size());
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;
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];
115 out_neighbors_row_splits[neighbor_idx];
117 &tmp_neighbors_count[neighbor_idx], 1);
118 out_neighbors_index[list_offset + item_offset] =
121 if (inp_neighbors_attributes) {
123 out_neighbors_attributes +
124 num_attributes_per_neighbor *
125 (list_offset + item_offset);
127 attr_i < num_attributes_per_neighbor;
129 attr_ptr[attr_i] = inp_neighbors_attributes
130 [num_attributes_per_neighbor * j +
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