geometry_op_impl.cc 8.74 KB
Newer Older
1
2
/*!
 *  Copyright (c) 2019 by Contributors
3
4
 * @file array/cpu/geometry_op_impl.cc
 * @brief Geometry operator CPU implementation
5
 */
6
#include <dgl/random.h>
7

8
#include <numeric>
9
#include <utility>
10
11
#include <vector>

12
#include "../geometry_op.h"
13
14
15
16
17
18

namespace dgl {
using runtime::NDArray;
namespace geometry {
namespace impl {

19
/*! @brief Knuth shuffle algorithm */
20
21
22
23
24
25
26
27
28
29
template <typename IdType>
void IndexShuffle(IdType *idxs, int64_t num_elems) {
  for (int64_t i = num_elems - 1; i > 0; --i) {
    int64_t j = dgl::RandomEngine::ThreadLocal()->RandInt(i);
    std::swap(idxs[i], idxs[j]);
  }
}
template void IndexShuffle<int32_t>(int32_t *idxs, int64_t num_elems);
template void IndexShuffle<int64_t>(int64_t *idxs, int64_t num_elems);

30
/*! @brief Groupwise index shuffle algorithm. This function will perform shuffle
31
32
33
 * in subarrays indicated by group index. The group index is similar to indptr
 * in CSRMatrix.
 *
34
35
36
37
 * @param group_idxs group index array.
 * @param idxs index array for shuffle.
 * @param num_groups_idxs length of group_idxs
 * @param num_elems length of idxs
38
39
 */
template <typename IdType>
40
41
42
void GroupIndexShuffle(
    const IdType *group_idxs, IdType *idxs, int64_t num_groups_idxs,
    int64_t num_elems) {
43
  if (num_groups_idxs < 2) return;  // empty idxs array
44
45
  CHECK_LE(group_idxs[num_groups_idxs - 1], num_elems)
      << "group_idxs out of range";
46
47
48
49
50
51
  for (int64_t i = 0; i < num_groups_idxs - 1; ++i) {
    auto subarray_len = group_idxs[i + 1] - group_idxs[i];
    IndexShuffle(idxs + group_idxs[i], subarray_len);
  }
}
template void GroupIndexShuffle<int32_t>(
52
53
    const int32_t *group_idxs, int32_t *idxs, int64_t num_groups_idxs,
    int64_t num_elems);
54
template void GroupIndexShuffle<int64_t>(
55
56
    const int64_t *group_idxs, int64_t *idxs, int64_t num_groups_idxs,
    int64_t num_elems);
57
58
59

template <typename IdType>
IdArray RandomPerm(int64_t num_nodes) {
60
61
62
  IdArray perm =
      aten::NewIdArray(num_nodes, DGLContext{kDGLCPU, 0}, sizeof(IdType) * 8);
  IdType *perm_data = static_cast<IdType *>(perm->data);
63
64
65
66
67
68
  std::iota(perm_data, perm_data + num_nodes, 0);
  IndexShuffle(perm_data, num_nodes);
  return perm;
}

template <typename IdType>
69
70
71
72
73
IdArray GroupRandomPerm(
    const IdType *group_idxs, int64_t num_group_idxs, int64_t num_nodes) {
  IdArray perm =
      aten::NewIdArray(num_nodes, DGLContext{kDGLCPU, 0}, sizeof(IdType) * 8);
  IdType *perm_data = static_cast<IdType *>(perm->data);
74
75
76
77
78
  std::iota(perm_data, perm_data + num_nodes, 0);
  GroupIndexShuffle(group_idxs, perm_data, num_group_idxs, num_nodes);
  return perm;
}

79
/*!
80
 * @brief Farthest Point Sampler without the need to compute all pairs of
81
82
83
84
 * distance.
 *
 * The input array has shape (N, d), where N is the number of points, and d is
 * the dimension. It consists of a (flatten) batch of point clouds.
85
 *
86
87
88
89
 * In each batch, the algorithm starts with the sample index specified by
 * ``start_idx``. Then for each point, we maintain the minimum to-sample
 * distance. Finally, we pick the point with the maximum such distance. This
 * process will be repeated for ``sample_points`` - 1 times.
90
 */
91
template <DGLDeviceType XPU, typename FloatType, typename IdType>
92
93
94
95
void FarthestPointSampler(
    NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
    IdArray start_idx, IdArray result) {
  const FloatType *array_data = static_cast<FloatType *>(array->data);
96
97
98
99
  const int64_t point_in_batch = array->shape[0] / batch_size;
  const int64_t dim = array->shape[1];

  // distance
100
  FloatType *dist_data = static_cast<FloatType *>(dist->data);
101
102

  // sample for each cloud in the batch
103
  IdType *start_idx_data = static_cast<IdType *>(start_idx->data);
104
105

  // return value
106
  IdType *ret_data = static_cast<IdType *>(result->data);
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

  int64_t array_start = 0, ret_start = 0;
  // loop for each point cloud sample in this batch
  for (auto b = 0; b < batch_size; b++) {
    // random init start sample
    int64_t sample_idx = (int64_t)start_idx_data[b];
    ret_data[ret_start] = (IdType)(sample_idx);

    // sample the rest `sample_points - 1` points
    for (auto i = 0; i < sample_points - 1; i++) {
      // re-init distance and the argmax
      int64_t dist_argmax = 0;
      FloatType dist_max = -1;

      // update the distance
      for (auto j = 0; j < point_in_batch; j++) {
        // compute the distance on dimensions
        FloatType one_dist = 0;
        for (auto d = 0; d < dim; d++) {
          FloatType tmp = array_data[(array_start + j) * dim + d] -
127
                          array_data[(array_start + sample_idx) * dim + d];
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
          one_dist += tmp * tmp;
        }

        // for each out-of-set point, keep its nearest to-the-set distance
        if (i == 0 || dist_data[j] > one_dist) {
          dist_data[j] = one_dist;
        }
        // look for the farthest sample
        if (dist_data[j] > dist_max) {
          dist_argmax = j;
          dist_max = dist_data[j];
        }
      }
      // sample the `dist_argmax`-th point
      sample_idx = dist_argmax;
      ret_data[ret_start + i + 1] = (IdType)(sample_idx);
    }

    array_start += point_in_batch;
    ret_start += sample_points;
  }
}
150
template void FarthestPointSampler<kDGLCPU, float, int32_t>(
151
152
    NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
    IdArray start_idx, IdArray result);
153
template void FarthestPointSampler<kDGLCPU, float, int64_t>(
154
155
    NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
    IdArray start_idx, IdArray result);
156
template void FarthestPointSampler<kDGLCPU, double, int32_t>(
157
158
    NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
    IdArray start_idx, IdArray result);
159
template void FarthestPointSampler<kDGLCPU, double, int64_t>(
160
161
    NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
    IdArray start_idx, IdArray result);
162

163
template <DGLDeviceType XPU, typename FloatType, typename IdType>
164
165
void WeightedNeighborMatching(
    const aten::CSRMatrix &csr, const NDArray weight, IdArray result) {
166
  const int64_t num_nodes = result->shape[0];
167
168
169
170
  const IdType *indptr_data = static_cast<IdType *>(csr.indptr->data);
  const IdType *indices_data = static_cast<IdType *>(csr.indices->data);
  IdType *result_data = static_cast<IdType *>(result->data);
  FloatType *weight_data = static_cast<FloatType *>(weight->data);
171
172
173

  // build node visiting order
  IdArray vis_order = RandomPerm<IdType>(num_nodes);
174
  IdType *vis_order_data = static_cast<IdType *>(vis_order->data);
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196

  for (int64_t n = 0; n < num_nodes; ++n) {
    auto u = vis_order_data[n];

    // if marked
    if (result_data[u] >= 0) continue;

    auto v_max = u;
    FloatType weight_max = 0;

    for (auto e = indptr_data[u]; e < indptr_data[u + 1]; ++e) {
      auto v = indices_data[e];
      if (result_data[v] >= 0) continue;
      if (weight_data[e] >= weight_max) {
        v_max = v;
        weight_max = weight_data[e];
      }
    }
    result_data[u] = std::min(u, v_max);
    result_data[v_max] = result_data[u];
  }
}
197
template void WeightedNeighborMatching<kDGLCPU, float, int32_t>(
198
    const aten::CSRMatrix &csr, const NDArray weight, IdArray result);
199
template void WeightedNeighborMatching<kDGLCPU, float, int64_t>(
200
    const aten::CSRMatrix &csr, const NDArray weight, IdArray result);
201
template void WeightedNeighborMatching<kDGLCPU, double, int32_t>(
202
    const aten::CSRMatrix &csr, const NDArray weight, IdArray result);
203
template void WeightedNeighborMatching<kDGLCPU, double, int64_t>(
204
205
    const aten::CSRMatrix &csr, const NDArray weight, IdArray result);

206
template <DGLDeviceType XPU, typename IdType>
207
208
void NeighborMatching(const aten::CSRMatrix &csr, IdArray result) {
  const int64_t num_nodes = result->shape[0];
209
210
211
  const IdType *indptr_data = static_cast<IdType *>(csr.indptr->data);
  const IdType *indices_data = static_cast<IdType *>(csr.indices->data);
  IdType *result_data = static_cast<IdType *>(result->data);
212
213
214

  // build vis order
  IdArray u_vis_order = RandomPerm<IdType>(num_nodes);
215
  IdType *u_vis_order_data = static_cast<IdType *>(u_vis_order->data);
216
  IdArray v_vis_order = GroupRandomPerm<IdType>(
217
218
      indptr_data, csr.indptr->shape[0], csr.indices->shape[0]);
  IdType *v_vis_order_data = static_cast<IdType *>(v_vis_order->data);
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236

  for (int64_t n = 0; n < num_nodes; ++n) {
    auto u = u_vis_order_data[n];

    // if marked
    if (result_data[u] >= 0) continue;

    result_data[u] = u;

    for (auto e = indptr_data[u]; e < indptr_data[u + 1]; ++e) {
      auto v = indices_data[v_vis_order_data[e]];
      if (result_data[v] >= 0) continue;
      result_data[u] = std::min(u, v);
      result_data[v] = result_data[u];
      break;
    }
  }
}
237
238
239
240
template void NeighborMatching<kDGLCPU, int32_t>(
    const aten::CSRMatrix &csr, IdArray result);
template void NeighborMatching<kDGLCPU, int64_t>(
    const aten::CSRMatrix &csr, IdArray result);
241

242
243
244
}  // namespace impl
}  // namespace geometry
}  // namespace dgl