Skip to content

Commit 7b1f2db

Browse files
denix56denix56
committed
Add Triton support
1 parent f709778 commit 7b1f2db

29 files changed

Lines changed: 713 additions & 1161 deletions

‎csrc/fps.cpp‎

Lines changed: 18 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
#ifdef WITH_PYTHON
22
#include <Python.h>
33
#endif
4-
#include <torch/script.h>
4+
#include <torch/torch.h>
5+
#include <torch/library.h>
56

67
#include "cpu/fps_cpu.h"
78

@@ -19,18 +20,21 @@ PyMODINIT_FUNC PyInit__fps_cpu(void) { return NULL; }
1920
#endif
2021
#endif
2122

22-
CLUSTER_API torch::Tensor fps(torch::Tensor src, torch::Tensor ptr, torch::Tensor ratio,
23-
bool random_start) {
24-
if (src.device().is_cuda()) {
25-
#ifdef WITH_CUDA
26-
return fps_cuda(src, ptr, ratio, random_start);
27-
#else
28-
AT_ERROR("Not compiled with CUDA support");
29-
#endif
30-
} else {
31-
return fps_cpu(src, ptr, ratio, random_start);
32-
}
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");
3330
}
3431

35-
static auto registry =
36-
torch::RegisterOperators().op("torch_cluster::fps", &fps);
32+
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
33+
m.impl("fps", &fps);
34+
}
35+
36+
#ifdef WITH_CUDA
37+
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
38+
m.impl("fps", &fps_cuda);
39+
}
40+
#endif

‎csrc/graclus.cpp‎

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
#ifdef WITH_PYTHON
22
#include <Python.h>
33
#endif
4-
#include <torch/script.h>
4+
#include <torch/torch.h>
5+
#include <torch/library.h>
56

67
#include "cpu/graclus_cpu.h"
78

@@ -20,17 +21,20 @@ PyMODINIT_FUNC PyInit__graclus_cpu(void) { return NULL; }
2021
#endif
2122

2223
CLUSTER_API torch::Tensor graclus(torch::Tensor rowptr, torch::Tensor col,
23-
std::optional<torch::Tensor> optional_weight) {
24-
if (rowptr.device().is_cuda()) {
25-
#ifdef WITH_CUDA
26-
return graclus_cuda(rowptr, col, optional_weight);
27-
#else
28-
AT_ERROR("Not compiled with CUDA support");
29-
#endif
30-
} else {
31-
return graclus_cpu(rowptr, col, optional_weight);
32-
}
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");
3330
}
3431

35-
static auto registry =
36-
torch::RegisterOperators().op("torch_cluster::graclus", &graclus);
32+
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
33+
m.impl("graclus", &graclus);
34+
}
35+
36+
#ifdef WITH_CUDA
37+
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
38+
m.impl("graclus", &graclus_cuda);
39+
}
40+
#endif

‎csrc/grid.cpp‎

Lines changed: 18 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
#ifdef WITH_PYTHON
22
#include <Python.h>
33
#endif
4-
#include <torch/script.h>
4+
#include <torch/torch.h>
5+
#include <torch/library.h>
56

67
#include "cpu/grid_cpu.h"
78

@@ -20,18 +21,21 @@ PyMODINIT_FUNC PyInit__grid_cpu(void) { return NULL; }
2021
#endif
2122

2223
CLUSTER_API torch::Tensor grid(torch::Tensor pos, torch::Tensor size,
23-
std::optional<torch::Tensor> optional_start,
24-
std::optional<torch::Tensor> optional_end) {
25-
if (pos.device().is_cuda()) {
26-
#ifdef WITH_CUDA
27-
return grid_cuda(pos, size, optional_start, optional_end);
28-
#else
29-
AT_ERROR("Not compiled with CUDA support");
30-
#endif
31-
} else {
32-
return grid_cpu(pos, size, optional_start, optional_end);
33-
}
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");
3431
}
3532

36-
static auto registry =
37-
torch::RegisterOperators().op("torch_cluster::grid", &grid);
33+
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
34+
m.impl("grid", &grid);
35+
}
36+
37+
#ifdef WITH_CUDA
38+
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
39+
m.impl("grid", &grid_cuda);
40+
}
41+
#endif

‎csrc/knn.cpp‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
#ifdef WITH_PYTHON
22
#include <Python.h>
33
#endif
4-
#include <torch/script.h>
4+
#include <torch/torch.h>
5+
#include <torch/library.h>
56

67
#include "cpu/knn_cpu.h"
78

@@ -23,18 +24,20 @@ CLUSTER_API torch::Tensor knn(torch::Tensor x, torch::Tensor y,
2324
std::optional<torch::Tensor> ptr_x,
2425
std::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
2526
int64_t num_workers) {
26-
if (x.device().is_cuda()) {
27-
#ifdef WITH_CUDA
28-
return knn_cuda(x, y, ptr_x, ptr_y, k, cosine);
29-
#else
30-
AT_ERROR("Not compiled with CUDA support");
31-
#endif
32-
} else {
33-
if (cosine)
34-
AT_ERROR("`cosine` argument not supported on CPU");
27+
TORCH_CHECK(!cosine, "`cosine` argument not supported on CPU");
3528
return knn_cpu(x, y, ptr_x, ptr_y, k, num_workers);
36-
}
3729
}
3830

39-
static auto registry =
40-
torch::RegisterOperators().op("torch_cluster::knn", &knn);
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");
33+
}
34+
35+
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
36+
m.impl("knn", &knn);
37+
}
38+
39+
#ifdef WITH_CUDA
40+
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
41+
m.impl("knn", &knn_cuda);
42+
}
43+
#endif

‎csrc/nearest.cpp‎

Lines changed: 9 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
#ifdef WITH_PYTHON
22
#include <Python.h>
33
#endif
4-
#include <torch/script.h>
4+
#include <torch/torch.h>
5+
#include <torch/library.h>
56

67
#include "extensions.h"
78

@@ -19,18 +20,12 @@ PyMODINIT_FUNC PyInit__nearest_cpu(void) { return NULL; }
1920
#endif
2021
#endif
2122

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

35-
static auto registry =
36-
torch::RegisterOperators().op("torch_cluster::nearest", &nearest);
27+
#ifdef WITH_CUDA
28+
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
29+
m.impl("nearest", &nearest_cuda);
30+
}
31+
#endif

‎csrc/radius.cpp‎

Lines changed: 22 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
#ifdef WITH_PYTHON
22
#include <Python.h>
33
#endif
4-
#include <torch/script.h>
4+
#include <torch/torch.h>
5+
#include <torch/library.h>
56

67
#include "cpu/radius_cpu.h"
78

@@ -20,20 +21,25 @@ PyMODINIT_FUNC PyInit__radius_cpu(void) { return NULL; }
2021
#endif
2122

2223
CLUSTER_API torch::Tensor radius(torch::Tensor x, torch::Tensor y,
23-
std::optional<torch::Tensor> ptr_x,
24-
std::optional<torch::Tensor> ptr_y, double r,
25-
int64_t max_num_neighbors, int64_t num_workers,
26-
bool ignore_same_index) {
27-
if (x.device().is_cuda()) {
28-
#ifdef WITH_CUDA
29-
return radius_cuda(x, y, ptr_x, ptr_y, r, max_num_neighbors, ignore_same_index);
30-
#else
31-
AT_ERROR("Not compiled with CUDA support");
32-
#endif
33-
} else {
34-
return radius_cpu(x, y, ptr_x, ptr_y, r, max_num_neighbors, num_workers, ignore_same_index);
35-
}
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");
3635
}
3736

38-
static auto registry =
39-
torch::RegisterOperators().op("torch_cluster::radius", &radius);
37+
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
38+
m.impl("radius", &radius);
39+
}
40+
41+
#ifdef WITH_CUDA
42+
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
43+
m.impl("radius", &radius_cuda);
44+
}
45+
#endif

‎csrc/rw.cpp‎

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
#ifdef WITH_PYTHON
22
#include <Python.h>
33
#endif
4-
#include <torch/script.h>
4+
#include <torch/torch.h>
5+
#include <torch/library.h>
56

67
#include "cpu/rw_cpu.h"
78

@@ -22,16 +23,19 @@ PyMODINIT_FUNC PyInit__rw_cpu(void) { return NULL; }
2223
CLUSTER_API std::tuple<torch::Tensor, torch::Tensor>
2324
random_walk(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start,
2425
int64_t walk_length, double p, double q) {
25-
if (rowptr.device().is_cuda()) {
26-
#ifdef WITH_CUDA
27-
return random_walk_cuda(rowptr, col, start, walk_length, p, q);
28-
#else
29-
AT_ERROR("Not compiled with CUDA support");
30-
#endif
31-
} else {
32-
return random_walk_cpu(rowptr, col, start, walk_length, p, q);
33-
}
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)");
3431
}
3532

36-
static auto registry =
37-
torch::RegisterOperators().op("torch_cluster::random_walk", &random_walk);
33+
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
34+
m.impl("random_walk", &random_walk);
35+
}
36+
37+
#ifdef WITH_CUDA
38+
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
39+
m.impl("random_walk", &random_walk_cuda);
40+
}
41+
#endif

‎csrc/sampler.cpp‎

Lines changed: 13 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
#ifdef WITH_PYTHON
22
#include <Python.h>
33
#endif
4-
#include <torch/script.h>
4+
#include <torch/torch.h>
5+
#include <torch/library.h>
56

67
#include "cpu/sampler_cpu.h"
78

@@ -15,18 +16,16 @@ PyMODINIT_FUNC PyInit__sampler_cpu(void) { return NULL; }
1516
#endif
1617
#endif
1718

18-
CLUSTER_API torch::Tensor neighbor_sampler(torch::Tensor start, torch::Tensor rowptr,
19-
int64_t count, double factor) {
20-
if (rowptr.device().is_cuda()) {
21-
#ifdef WITH_CUDA
22-
AT_ERROR("No CUDA version supported");
23-
#else
24-
AT_ERROR("Not compiled with CUDA support");
25-
#endif
26-
} else {
27-
return neighbor_sampler_cpu(start, rowptr, count, factor);
28-
}
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);
2923
}
3024

31-
static auto registry = torch::RegisterOperators().op(
32-
"torch_cluster::neighbor_sampler", &neighbor_sampler);
25+
TORCH_LIBRARY(torch_cluster, m) {
26+
m.def("neighbor_sampler(Tensor start, Tensor rowptr, int count, float factor) -> Tensor");
27+
}
28+
29+
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
30+
m.impl("neighbor_sampler", &neighbor_sampler);
31+
}

‎csrc/version.cpp‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,8 @@
33
#endif
44
#include "cluster.h"
55
#include "macros.h"
6-
#include <torch/script.h>
6+
#include <torch/torch.h>
7+
#include <torch/library.h>
78

89
#ifdef WITH_CUDA
910
#ifdef USE_ROCM
@@ -37,5 +38,7 @@ CLUSTER_API int64_t cuda_version() noexcept {
3738
}
3839
} // namespace cluster
3940

40-
static auto registry = torch::RegisterOperators().op(
41-
"torch_cluster::cuda_version", [] { return cluster::cuda_version(); });
41+
TORCH_LIBRARY(torch_cluster, m) {
42+
m.def("cuda_version() -> int");
43+
m.impl("cuda_version", [] { return cluster::cuda_version(); });
44+
}

‎setup.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
from torch.utils.cpp_extension import (CUDA_HOME, BuildExtension, CppExtension,
1212
CUDAExtension)
1313

14-
__version__ = '1.6.3'
14+
__version__ = '2.0.0'
1515
URL = 'https://github.com/rusty1s/pytorch_cluster'
1616

1717
WITH_CUDA = False
@@ -25,7 +25,6 @@
2525
suffices = ['cuda']
2626
if os.getenv('FORCE_ONLY_CPU', '0') == '1':
2727
suffices = ['cpu']
28-
2928
BUILD_DOCS = os.getenv('BUILD_DOCS', '0') == '1'
3029

3130

0 commit comments

Comments
 (0)