Skip to content

Commit b75a19e

Browse files
rusty1srusty1s
authored
Merge pull request #68 from liaopeiyuan/cpu_radius
C++ CPU for radius and knn
2 parents 32fa325 + 29f9716 commit b75a19e

22 files changed

+2863
-201
lines changed

‎.travis.yml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,8 @@ install:
5959
- conda install pytorch=${TORCH_VERSION} ${TOOLKIT} -c pytorch --yes
6060
- source script/torch.sh
6161
- pip install flake8 codecov
62-
- python setup.py install
62+
- pip install scipy==1.4.1
63+
- pip install .[test]
6364
script:
6465
- flake8 .
6566
- python setup.py test

‎CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
cmake_minimum_required(VERSION 3.0)
22
project(torchcluster)
33
set(CMAKE_CXX_STANDARD 14)
4-
set(TORCHCLUSTER_VERSION 1.5.4)
4+
set(TORCHCLUSTER_VERSION 1.5.5)
55

66
option(WITH_CUDA "Enable CUDA support" OFF)
77

‎csrc/cpu/knn_cpu.cpp‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
#include "knn_cpu.h"
2+
3+
#include "utils.h"
4+
#include "utils/neighbors.cpp"
5+
6+
torch::Tensor knn_cpu(torch::Tensor x, torch::Tensor y,
7+
torch::optional<torch::Tensor> ptr_x,
8+
torch::optional<torch::Tensor> ptr_y, int64_t k,
9+
int64_t num_workers) {
10+
11+
CHECK_CPU(x);
12+
CHECK_INPUT(x.dim() == 2);
13+
CHECK_CPU(y);
14+
CHECK_INPUT(y.dim() == 2);
15+
16+
if (ptr_x.has_value()) {
17+
CHECK_CPU(ptr_x.value());
18+
CHECK_INPUT(ptr_x.value().dim() == 1);
19+
}
20+
if (ptr_y.has_value()) {
21+
CHECK_CPU(ptr_y.value());
22+
CHECK_INPUT(ptr_y.value().dim() == 1);
23+
}
24+
25+
std::vector<size_t> *out_vec = new std::vector<size_t>();
26+
27+
AT_DISPATCH_ALL_TYPES(x.scalar_type(), "radius_cpu", [&] {
28+
auto x_data = x.data_ptr<scalar_t>();
29+
auto y_data = y.data_ptr<scalar_t>();
30+
auto x_vec = std::vector<scalar_t>(x_data, x_data + x.numel());
31+
auto y_vec = std::vector<scalar_t>(y_data, y_data + y.numel());
32+
33+
if (!ptr_x.has_value()) {
34+
nanoflann_neighbors<scalar_t>(y_vec, x_vec, out_vec, 0, x.size(-1), 0,
35+
num_workers, k, 0);
36+
} else {
37+
auto sx = (ptr_x.value().narrow(0, 1, ptr_x.value().numel() - 1) -
38+
ptr_x.value().narrow(0, 0, ptr_x.value().numel() - 1));
39+
auto sy = (ptr_y.value().narrow(0, 1, ptr_y.value().numel() - 1) -
40+
ptr_y.value().narrow(0, 0, ptr_y.value().numel() - 1));
41+
auto sx_data = sx.data_ptr<int64_t>();
42+
auto sy_data = sy.data_ptr<int64_t>();
43+
auto sx_vec = std::vector<long>(sx_data, sx_data + sx.numel());
44+
auto sy_vec = std::vector<long>(sy_data, sy_data + sy.numel());
45+
batch_nanoflann_neighbors<scalar_t>(y_vec, x_vec, sy_vec, sx_vec, out_vec,
46+
k, x.size(-1), 0, k, 0);
47+
}
48+
});
49+
50+
const int64_t size = out_vec->size() / 2;
51+
auto out = torch::from_blob(out_vec->data(), {size, 2},
52+
x.options().dtype(torch::kLong));
53+
return out.t().index_select(0, torch::tensor({1, 0}));
54+
}

‎csrc/cpu/knn_cpu.h‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
#pragma once
2+
3+
#include <torch/extension.h>
4+
5+
torch::Tensor knn_cpu(torch::Tensor x, torch::Tensor y,
6+
torch::optional<torch::Tensor> ptr_x,
7+
torch::optional<torch::Tensor> ptr_y, int64_t k,
8+
int64_t num_workers);

‎csrc/cpu/radius_cpu.cpp‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
#include "radius_cpu.h"
2+
3+
#include "utils.h"
4+
#include "utils/neighbors.cpp"
5+
6+
torch::Tensor radius_cpu(torch::Tensor x, torch::Tensor y,
7+
torch::optional<torch::Tensor> ptr_x,
8+
torch::optional<torch::Tensor> ptr_y, double r,
9+
int64_t max_num_neighbors, int64_t num_workers) {
10+
11+
CHECK_CPU(x);
12+
CHECK_INPUT(x.dim() == 2);
13+
CHECK_CPU(y);
14+
CHECK_INPUT(y.dim() == 2);
15+
16+
if (ptr_x.has_value()) {
17+
CHECK_CPU(ptr_x.value());
18+
CHECK_INPUT(ptr_x.value().dim() == 1);
19+
}
20+
if (ptr_y.has_value()) {
21+
CHECK_CPU(ptr_y.value());
22+
CHECK_INPUT(ptr_y.value().dim() == 1);
23+
}
24+
25+
std::vector<size_t> *out_vec = new std::vector<size_t>();
26+
27+
AT_DISPATCH_ALL_TYPES(x.scalar_type(), "radius_cpu", [&] {
28+
auto x_data = x.data_ptr<scalar_t>();
29+
auto y_data = y.data_ptr<scalar_t>();
30+
auto x_vec = std::vector<scalar_t>(x_data, x_data + x.numel());
31+
auto y_vec = std::vector<scalar_t>(y_data, y_data + y.numel());
32+
33+
if (!ptr_x.has_value()) {
34+
nanoflann_neighbors<scalar_t>(y_vec, x_vec, out_vec, r, x.size(-1),
35+
max_num_neighbors, num_workers, 0, 1);
36+
} else {
37+
auto sx = (ptr_x.value().narrow(0, 1, ptr_x.value().numel() - 1) -
38+
ptr_x.value().narrow(0, 0, ptr_x.value().numel() - 1));
39+
auto sy = (ptr_y.value().narrow(0, 1, ptr_y.value().numel() - 1) -
40+
ptr_y.value().narrow(0, 0, ptr_y.value().numel() - 1));
41+
auto sx_data = sx.data_ptr<int64_t>();
42+
auto sy_data = sy.data_ptr<int64_t>();
43+
auto sx_vec = std::vector<long>(sx_data, sx_data + sx.numel());
44+
auto sy_vec = std::vector<long>(sy_data, sy_data + sy.numel());
45+
batch_nanoflann_neighbors<scalar_t>(y_vec, x_vec, sy_vec, sx_vec, out_vec,
46+
r, x.size(-1), max_num_neighbors, 0,
47+
1);
48+
}
49+
});
50+
51+
const int64_t size = out_vec->size() / 2;
52+
auto out = torch::from_blob(out_vec->data(), {size, 2},
53+
x.options().dtype(torch::kLong));
54+
return out.t().index_select(0, torch::tensor({1, 0}));
55+
}

‎csrc/cpu/radius_cpu.h‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
#pragma once
2+
3+
#include <torch/extension.h>
4+
5+
torch::Tensor radius_cpu(torch::Tensor x, torch::Tensor y,
6+
torch::optional<torch::Tensor> ptr_x,
7+
torch::optional<torch::Tensor> ptr_y, double r,
8+
int64_t max_num_neighbors, int64_t num_workers);

‎csrc/cpu/utils/cloud.h‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
#pragma once
2+
3+
#include <ATen/ATen.h>
4+
#include <algorithm>
5+
#include <cmath>
6+
#include <iomanip>
7+
#include <iostream>
8+
#include <map>
9+
#include <numeric>
10+
#include <unordered_map>
11+
#include <vector>
12+
13+
#include <time.h>
14+
15+
template <typename scalar_t> struct PointCloud {
16+
std::vector<std::vector<scalar_t> *> pts;
17+
18+
void set(std::vector<scalar_t> new_pts, int dim) {
19+
20+
std::vector<std::vector<scalar_t> *> temp(new_pts.size() / dim);
21+
for (size_t i = 0; i < new_pts.size(); i++) {
22+
if (i % dim == 0) {
23+
std::vector<scalar_t> *point = new std::vector<scalar_t>(dim);
24+
25+
for (size_t j = 0; j < (size_t)dim; j++) {
26+
(*point)[j] = new_pts[i + j];
27+
}
28+
temp[i / dim] = point;
29+
}
30+
}
31+
32+
pts = temp;
33+
}
34+
void set_batch(std::vector<scalar_t> new_pts, size_t begin, long size,
35+
int dim) {
36+
std::vector<std::vector<scalar_t> *> temp(size);
37+
for (size_t i = 0; i < (size_t)size; i++) {
38+
std::vector<scalar_t> *point = new std::vector<scalar_t>(dim);
39+
for (size_t j = 0; j < (size_t)dim; j++) {
40+
(*point)[j] = new_pts[dim * (begin + i) + j];
41+
}
42+
43+
temp[i] = point;
44+
}
45+
pts = temp;
46+
}
47+
48+
// Must return the number of data points.
49+
inline size_t kdtree_get_point_count() const { return pts.size(); }
50+
51+
// Returns the dim'th component of the idx'th point in the class:
52+
inline scalar_t kdtree_get_pt(const size_t idx, const size_t dim) const {
53+
return (*pts[idx])[dim];
54+
}
55+
56+
// Optional bounding-box computation: return false to default to a standard
57+
// bbox computation loop.
58+
// Return true if the BBOX was already computed by the class and returned in
59+
// "bb" so it can be avoided to redo it again. Look at bb.size() to find out
60+
// the expected dimensionality (e.g. 2 or 3 for point clouds)
61+
template <class BBOX> bool kdtree_get_bbox(BBOX & /* bb */) const {
62+
return false;
63+
}
64+
};

0 commit comments

Comments
 (0)