Skip to content

Commit 3b9b3d8

Browse files
denix56denix56
authored
Merge pull request #3 from denix56/triton
Add Triton-based variants of KNN and Nearest search algorithms
2 parents 6dabb04 + 02cc6ae commit 3b9b3d8

27 files changed

Lines changed: 685 additions & 111 deletions

‎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: 11 additions & 3 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

@@ -32,5 +33,12 @@ CLUSTER_API torch::Tensor fps(torch::Tensor src, torch::Tensor ptr, torch::Tenso
3233
}
3334
}
3435

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

‎csrc/graclus.cpp‎

Lines changed: 11 additions & 3 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

@@ -32,5 +33,12 @@ CLUSTER_API torch::Tensor graclus(torch::Tensor rowptr, torch::Tensor col,
3233
}
3334
}
3435

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

‎csrc/grid.cpp‎

Lines changed: 11 additions & 3 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

@@ -33,5 +34,12 @@ CLUSTER_API torch::Tensor grid(torch::Tensor pos, torch::Tensor size,
3334
}
3435
}
3536

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

‎csrc/knn.cpp‎

Lines changed: 19 additions & 3 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

@@ -36,5 +37,20 @@ CLUSTER_API torch::Tensor knn(torch::Tensor x, torch::Tensor y,
3637
}
3738
}
3839

39-
static auto registry =
40-
torch::RegisterOperators().op("torch_cluster::knn", &knn);
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);
46+
}
47+
48+
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
49+
m.impl("knn", &knn_cpu_wrap);
50+
}
51+
52+
#ifdef WITH_CUDA
53+
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
54+
m.impl("knn", &knn_cuda);
55+
}
56+
#endif

‎csrc/nearest.cpp‎

Lines changed: 7 additions & 3 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

@@ -32,5 +33,8 @@ CLUSTER_API torch::Tensor nearest(torch::Tensor x, torch::Tensor y, torch::Tenso
3233
}
3334
}
3435

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

‎csrc/radius.cpp‎

Lines changed: 11 additions & 3 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

@@ -35,5 +36,12 @@ CLUSTER_API torch::Tensor radius(torch::Tensor x, torch::Tensor y,
3536
}
3637
}
3738

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

‎csrc/rw.cpp‎

Lines changed: 11 additions & 3 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

@@ -33,5 +34,12 @@ random_walk(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start,
3334
}
3435
}
3536

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

‎csrc/sampler.cpp‎

Lines changed: 5 additions & 3 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

@@ -28,5 +29,6 @@ CLUSTER_API torch::Tensor neighbor_sampler(torch::Tensor start, torch::Tensor ro
2829
}
2930
}
3031

31-
static auto registry = torch::RegisterOperators().op(
32-
"torch_cluster::neighbor_sampler", &neighbor_sampler);
32+
TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
33+
m.impl("neighbor_sampler", &neighbor_sampler_cpu);
34+
}

‎csrc/version.cpp‎

Lines changed: 5 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,6 @@ 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_IMPL(torch_cluster, m) {
42+
m.impl("cuda_version", [] { return cluster::cuda_version(); });
43+
}

0 commit comments

Comments
 (0)