Skip to content

Commit 02cc6ae

Browse files
denix56denix56
committed
Replace jit.script with custom ops registration using TORCH_LIBRARY
1 parent 7b1f2db commit 02cc6ae

10 files changed

Lines changed: 117 additions & 60 deletions

File tree

‎csrc/cluster.h‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#pragma once
2+
#include <torch/library.h>
23

34
#include "extensions.h"
45

@@ -37,3 +38,17 @@ random_walk(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start,
3738

3839
CLUSTER_API torch::Tensor neighbor_sampler(torch::Tensor start, torch::Tensor rowptr,
3940
int64_t count, double factor);
41+
42+
TORCH_LIBRARY(torch_cluster, m) {
43+
m.def("fps(Tensor src, Tensor ptr, Tensor ratio, bool random_start = False) -> Tensor");
44+
m.def("graclus(Tensor rowptr, Tensor col, Tensor? weight = None) -> Tensor");
45+
m.def("grid(Tensor pos, Tensor size, Tensor? start = None, Tensor? end = None) -> Tensor");
46+
m.def("knn(Tensor a, Tensor b, Tensor? ptr_x, Tensor? ptr_y, int k, bool cosine = False, int num_workers = 1) -> Tensor");
47+
m.def("nearest(Tensor x, Tensor y, Tensor ptr_x, Tensor ptr_y) -> Tensor");
48+
m.def("radius(Tensor x, Tensor y, Tensor? ptr_x, Tensor? ptr_y, float r, int max_num_neighbors, int num_workers = 1, bool ignore_same_index = False) -> Tensor");
49+
m.def("random_walk(Tensor rowptr, Tensor col, Tensor start, int walk_length, float p = 1, float q = 1) -> (Tensor, Tensor)");
50+
m.def("neighbor_sampler(Tensor start, Tensor rowptr, int count, float factor) -> Tensor");
51+
m.def("cuda_version() -> int");
52+
}
53+
54+
}

‎csrc/fps.cpp‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,17 +20,21 @@ PyMODINIT_FUNC PyInit__fps_cpu(void) { return NULL; }
2020
#endif
2121
#endif
2222

23-
CLUSTER_API torch::Tensor fps(torch::Tensor src, torch::Tensor ptr,
24-
torch::Tensor ratio, bool random_start) {
25-
return fps_cpu(src, ptr, ratio, random_start);
26-
}
27-
28-
TORCH_LIBRARY(torch_cluster, m) {
29-
m.def("fps(Tensor src, Tensor ptr, Tensor ratio, bool random_start = False) -> Tensor");
23+
CLUSTER_API torch::Tensor fps(torch::Tensor src, torch::Tensor ptr, torch::Tensor ratio,
24+
bool random_start) {
25+
if (src.device().is_cuda()) {
26+
#ifdef WITH_CUDA
27+
return fps_cuda(src, ptr, ratio, random_start);
28+
#else
29+
AT_ERROR("Not compiled with CUDA support");
30+
#endif
31+
} else {
32+
return fps_cpu(src, ptr, ratio, random_start);
33+
}
3034
}
3135

3236
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
33-
m.impl("fps", &fps);
37+
m.impl("fps", &fps_cpu);
3438
}
3539

3640
#ifdef WITH_CUDA

‎csrc/graclus.cpp‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,16 +21,20 @@ PyMODINIT_FUNC PyInit__graclus_cpu(void) { return NULL; }
2121
#endif
2222

2323
CLUSTER_API torch::Tensor graclus(torch::Tensor rowptr, torch::Tensor col,
24-
std::optional<torch::Tensor> optional_weight) {
25-
return graclus_cpu(rowptr, col, optional_weight);
26-
}
27-
28-
TORCH_LIBRARY(torch_cluster, m) {
29-
m.def("graclus(Tensor rowptr, Tensor col, Tensor? weight = None) -> Tensor");
24+
std::optional<torch::Tensor> optional_weight) {
25+
if (rowptr.device().is_cuda()) {
26+
#ifdef WITH_CUDA
27+
return graclus_cuda(rowptr, col, optional_weight);
28+
#else
29+
AT_ERROR("Not compiled with CUDA support");
30+
#endif
31+
} else {
32+
return graclus_cpu(rowptr, col, optional_weight);
33+
}
3034
}
3135

3236
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
33-
m.impl("graclus", &graclus);
37+
m.impl("graclus", &graclus_cpu);
3438
}
3539

3640
#ifdef WITH_CUDA

‎csrc/grid.cpp‎

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -21,21 +21,25 @@ PyMODINIT_FUNC PyInit__grid_cpu(void) { return NULL; }
2121
#endif
2222

2323
CLUSTER_API torch::Tensor grid(torch::Tensor pos, torch::Tensor size,
24-
std::optional<torch::Tensor> optional_start,
25-
std::optional<torch::Tensor> optional_end) {
26-
return grid_cpu(pos, size, optional_start, optional_end);
27-
}
28-
29-
TORCH_LIBRARY(torch_cluster, m) {
30-
m.def("grid(Tensor pos, Tensor size, Tensor? start = None, Tensor? end = None) -> Tensor");
24+
std::optional<torch::Tensor> optional_start,
25+
std::optional<torch::Tensor> optional_end) {
26+
if (pos.device().is_cuda()) {
27+
#ifdef WITH_CUDA
28+
return grid_cuda(pos, size, optional_start, optional_end);
29+
#else
30+
AT_ERROR("Not compiled with CUDA support");
31+
#endif
32+
} else {
33+
return grid_cpu(pos, size, optional_start, optional_end);
34+
}
3135
}
3236

3337
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
34-
m.impl("grid", &grid);
38+
m.impl("grid", &grid_cpu);
3539
}
3640

3741
#ifdef WITH_CUDA
3842
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
39-
m.impl("grid", &grid_cuda);
43+
m.impl("grid", &grid_cuda);
4044
}
4145
#endif

‎csrc/knn.cpp‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,16 +24,29 @@ CLUSTER_API torch::Tensor knn(torch::Tensor x, torch::Tensor y,
2424
std::optional<torch::Tensor> ptr_x,
2525
std::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
2626
int64_t num_workers) {
27-
TORCH_CHECK(!cosine, "`cosine` argument not supported on CPU");
27+
if (x.device().is_cuda()) {
28+
#ifdef WITH_CUDA
29+
return knn_cuda(x, y, ptr_x, ptr_y, k, cosine);
30+
#else
31+
AT_ERROR("Not compiled with CUDA support");
32+
#endif
33+
} else {
34+
if (cosine)
35+
AT_ERROR("`cosine` argument not supported on CPU");
2836
return knn_cpu(x, y, ptr_x, ptr_y, k, num_workers);
37+
}
2938
}
3039

31-
TORCH_LIBRARY(torch_cluster, m) {
32-
m.def("knn(Tensor a, Tensor b, Tensor? ptr_x, Tensor? ptr_y, int k, bool cosine = False, int num_workers = 1) -> Tensor");
40+
torch::Tensor knn_cpu_wrap(torch::Tensor x, torch::Tensor y,
41+
std::optional<torch::Tensor> ptr_x,
42+
std::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
43+
int64_t num_workers) {
44+
TORCH_CHECK(!cosine, "`cosine` argument not supported on CPU");
45+
return knn_cpu(x, y, ptr_x, ptr_y, k, num_workers);
3346
}
3447

3548
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
36-
m.impl("knn", &knn);
49+
m.impl("knn", &knn_cpu_wrap);
3750
}
3851

3952
#ifdef WITH_CUDA

‎csrc/nearest.cpp‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,21 @@ PyMODINIT_FUNC PyInit__nearest_cpu(void) { return NULL; }
2020
#endif
2121
#endif
2222

23-
TORCH_LIBRARY(torch_cluster, m) {
24-
m.def("nearest(Tensor x, Tensor y, Tensor ptr_x, Tensor ptr_y) -> Tensor");
23+
CLUSTER_API torch::Tensor nearest(torch::Tensor x, torch::Tensor y, torch::Tensor ptr_x,
24+
torch::Tensor ptr_y) {
25+
if (x.device().is_cuda()) {
26+
#ifdef WITH_CUDA
27+
return nearest_cuda(x, y, ptr_x, ptr_y);
28+
#else
29+
AT_ERROR("Not compiled with CUDA support");
30+
#endif
31+
} else {
32+
AT_ERROR("No CPU version supported");
33+
}
2534
}
2635

2736
#ifdef WITH_CUDA
2837
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
29-
m.impl("nearest", &nearest_cuda);
38+
m.impl("nearest", &nearest_cuda);
3039
}
3140
#endif

‎csrc/radius.cpp‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -21,21 +21,23 @@ PyMODINIT_FUNC PyInit__radius_cpu(void) { return NULL; }
2121
#endif
2222

2323
CLUSTER_API torch::Tensor radius(torch::Tensor x, torch::Tensor y,
24-
std::optional<torch::Tensor> ptr_x,
25-
std::optional<torch::Tensor> ptr_y, double r,
26-
int64_t max_num_neighbors,
27-
int64_t num_workers,
28-
bool ignore_same_index) {
29-
return radius_cpu(x, y, ptr_x, ptr_y, r, max_num_neighbors, num_workers,
30-
ignore_same_index);
31-
}
32-
33-
TORCH_LIBRARY(torch_cluster, m) {
34-
m.def("radius(Tensor x, Tensor y, Tensor? ptr_x, Tensor? ptr_y, float r, int max_num_neighbors, int num_workers = 1, bool ignore_same_index = False) -> Tensor");
24+
std::optional<torch::Tensor> ptr_x,
25+
std::optional<torch::Tensor> ptr_y, double r,
26+
int64_t max_num_neighbors, int64_t num_workers,
27+
bool ignore_same_index) {
28+
if (x.device().is_cuda()) {
29+
#ifdef WITH_CUDA
30+
return radius_cuda(x, y, ptr_x, ptr_y, r, max_num_neighbors, ignore_same_index);
31+
#else
32+
AT_ERROR("Not compiled with CUDA support");
33+
#endif
34+
} else {
35+
return radius_cpu(x, y, ptr_x, ptr_y, r, max_num_neighbors, num_workers, ignore_same_index);
36+
}
3537
}
3638

3739
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
38-
m.impl("radius", &radius);
40+
m.impl("radius", &radius_cpu);
3941
}
4042

4143
#ifdef WITH_CUDA

‎csrc/rw.cpp‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,15 +23,19 @@ PyMODINIT_FUNC PyInit__rw_cpu(void) { return NULL; }
2323
CLUSTER_API std::tuple<torch::Tensor, torch::Tensor>
2424
random_walk(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start,
2525
int64_t walk_length, double p, double q) {
26-
return random_walk_cpu(rowptr, col, start, walk_length, p, q);
27-
}
28-
29-
TORCH_LIBRARY(torch_cluster, m) {
30-
m.def("random_walk(Tensor rowptr, Tensor col, Tensor start, int walk_length, float p = 1, float q = 1) -> (Tensor, Tensor)");
26+
if (rowptr.device().is_cuda()) {
27+
#ifdef WITH_CUDA
28+
return random_walk_cuda(rowptr, col, start, walk_length, p, q);
29+
#else
30+
AT_ERROR("Not compiled with CUDA support");
31+
#endif
32+
} else {
33+
return random_walk_cpu(rowptr, col, start, walk_length, p, q);
34+
}
3135
}
3236

3337
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
34-
m.impl("random_walk", &random_walk);
38+
m.impl("random_walk", &random_walk_cpu);
3539
}
3640

3741
#ifdef WITH_CUDA

‎csrc/sampler.cpp‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -16,16 +16,19 @@ PyMODINIT_FUNC PyInit__sampler_cpu(void) { return NULL; }
1616
#endif
1717
#endif
1818

19-
CLUSTER_API torch::Tensor neighbor_sampler(torch::Tensor start,
20-
torch::Tensor rowptr, int64_t count,
21-
double factor) {
22-
return neighbor_sampler_cpu(start, rowptr, count, factor);
23-
}
24-
25-
TORCH_LIBRARY(torch_cluster, m) {
26-
m.def("neighbor_sampler(Tensor start, Tensor rowptr, int count, float factor) -> Tensor");
19+
CLUSTER_API torch::Tensor neighbor_sampler(torch::Tensor start, torch::Tensor rowptr,
20+
int64_t count, double factor) {
21+
if (rowptr.device().is_cuda()) {
22+
#ifdef WITH_CUDA
23+
AT_ERROR("No CUDA version supported");
24+
#else
25+
AT_ERROR("Not compiled with CUDA support");
26+
#endif
27+
} else {
28+
return neighbor_sampler_cpu(start, rowptr, count, factor);
29+
}
2730
}
2831

2932
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
30-
m.impl("neighbor_sampler", &neighbor_sampler);
33+
m.impl("neighbor_sampler", &neighbor_sampler_cpu);
3134
}

‎csrc/version.cpp‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,6 @@ CLUSTER_API int64_t cuda_version() noexcept {
3838
}
3939
} // namespace cluster
4040

41-
TORCH_LIBRARY(torch_cluster, m) {
42-
m.def("cuda_version() -> int");
41+
TORCH_LIBRARY_IMPL(torch_cluster, m) {
4342
m.impl("cuda_version", [] { return cluster::cuda_version(); });
4443
}

0 commit comments

Comments
 (0)