Skip to content

Commit ee711ea

Browse files
denix56denix56
authored
Add Triton support
* Compute split-N tile counts via heuristics * Add Triton support * Add Triton support
1 parent 3b9b3d8 commit ee711ea

23 files changed

Lines changed: 776 additions & 266 deletions

‎benchmarks/test_benchmark_knn.py‎

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
import importlib.util
2+
from itertools import product
3+
4+
import pytest
5+
import torch
6+
import torch_cluster as tc
7+
8+
knn = tc.knn
9+
knn_graph = tc.knn_graph
10+
11+
12+
pytestmark = pytest.mark.skipif(
13+
not (torch.cuda.is_available() and importlib.util.find_spec('triton') is not None),
14+
reason='CUDA and Triton are required for Triton benchmark tests.',
15+
)
16+
17+
18+
def to_set(edge_index):
19+
return set([(i, j) for i, j in edge_index.t().tolist()])
20+
21+
22+
def _make_batch(num_nodes: int, num_groups: int,
23+
device: torch.device) -> torch.Tensor:
24+
groups = max(1, min(num_groups, num_nodes))
25+
counts = torch.full((groups, ), num_nodes // groups, device=device,
26+
dtype=torch.long)
27+
remainder = num_nodes % groups
28+
if remainder:
29+
counts[:remainder] += 1
30+
return torch.repeat_interleave(torch.arange(groups, device=device),
31+
counts)
32+
33+
34+
@pytest.mark.parametrize('num_x,num_y,num_groups',
35+
((*p[0], p[1]) for p in product([(256, 128), (1024, 512), (4096, 2048), (255, 127),
36+
(256, 5), (1024, 5), (4096, 5), (255, 5)],
37+
[1, 2, 4, 8, 16, 32]) if p[1] <= min(p[0])))
38+
@pytest.mark.benchmark(group="knn")
39+
def test_triton_knn_benchmark_cuda(benchmark, num_x, num_y, num_groups):
40+
torch.manual_seed(99)
41+
x = torch.randn(num_x, 16, device='cuda')
42+
y = torch.randn(num_y, 16, device='cuda')
43+
groups = min(num_groups, x.size(0), y.size(0))
44+
batch_x = _make_batch(num_x, groups, x.device)
45+
batch_y = _make_batch(num_y, groups, y.device)
46+
47+
def cuda_fn():
48+
knn(x, y, k=16, batch_x=batch_x, batch_y=batch_y, cosine=False, use_triton=False)
49+
50+
for _ in range(5):
51+
cuda_fn()
52+
torch.cuda.synchronize()
53+
54+
benchmark(cuda_fn)
55+
print(f"[knn][cuda] num_x={num_x} num_y={num_y} groups={groups}")
56+
57+
58+
@pytest.mark.parametrize('num_x,num_y,num_groups',
59+
((*p[0], p[1]) for p in product([(256, 128), (1024, 512), (4096, 2048), (255, 127),
60+
(256, 5), (1024, 5), (4096, 5), (255, 5)],
61+
[1, 2, 4, 8, 16, 32]) if p[1] <= min(p[0])))
62+
@pytest.mark.benchmark(group="knn")
63+
def test_triton_knn_benchmark_triton_cosine(benchmark, num_x, num_y, num_groups):
64+
torch.manual_seed(99)
65+
x = torch.randn(num_x, 16, device='cuda')
66+
y = torch.randn(num_y, 16, device='cuda')
67+
groups = min(num_groups, x.size(0), y.size(0))
68+
batch_x = _make_batch(num_x, groups, x.device)
69+
batch_y = _make_batch(num_y, groups, y.device)
70+
71+
def cuda_fn():
72+
return knn(x, y, k=16, batch_x=batch_x, batch_y=batch_y, cosine=True, use_triton=False)
73+
74+
def triton_fn():
75+
return knn(x, y, k=16, batch_x=batch_x, batch_y=batch_y, cosine=True, use_triton=True)
76+
77+
for i in range(5):
78+
if i == 0:
79+
out_cuda = cuda_fn()
80+
out_triton = triton_fn()
81+
assert to_set(out_cuda) == to_set(out_triton)
82+
else:
83+
triton_fn()
84+
torch.cuda.synchronize()
85+
86+
benchmark(triton_fn)
87+
print(f"[knn][triton] num_x={num_x} num_y={num_y} groups={groups}")
88+
89+
90+
@pytest.mark.parametrize('num_x,num_y,num_groups',
91+
((*p[0], p[1]) for p in product([(256, 128), (1024, 512), (4096, 2048), (255, 127),
92+
(256, 5), (1024, 5), (4096, 5), (255, 5)],
93+
[1, 2, 4, 8, 16, 32]) if p[1] <= min(p[0])))
94+
@pytest.mark.benchmark(group="knn")
95+
def test_triton_knn_benchmark_triton(benchmark, num_x, num_y, num_groups):
96+
torch.manual_seed(99)
97+
x = torch.randn(num_x, 16, device='cuda')
98+
y = torch.randn(num_y, 16, device='cuda')
99+
groups = min(num_groups, x.size(0), y.size(0))
100+
batch_x = _make_batch(num_x, groups, x.device)
101+
batch_y = _make_batch(num_y, groups, y.device)
102+
103+
def cuda_fn():
104+
return knn(x, y, k=16, batch_x=batch_x, batch_y=batch_y, cosine=False, use_triton=False)
105+
106+
def triton_fn():
107+
return knn(x, y, k=16, batch_x=batch_x, batch_y=batch_y, cosine=False, use_triton=True)
108+
109+
for i in range(5):
110+
if i == 0:
111+
out_cuda = cuda_fn()
112+
out_triton = triton_fn()
113+
assert to_set(out_cuda) == to_set(out_triton)
114+
else:
115+
triton_fn()
116+
torch.cuda.synchronize()
117+
118+
benchmark(triton_fn)
119+
print(f"[knn][triton] num_x={num_x} num_y={num_y} groups={groups}")
120+
121+
122+
@pytest.mark.parametrize('num_x', [256, 1024, 4096, 255])
123+
@pytest.mark.parametrize('num_groups', [1, 2, 4, 6, 8, 16, 24, 32])
124+
@pytest.mark.benchmark(group="knn_graph")
125+
def test_triton_knn_graph_benchmark_cuda(benchmark, num_x, num_groups):
126+
torch.manual_seed(199)
127+
x = torch.randn(num_x, 8, device='cuda')
128+
groups = min(num_groups, x.size(0))
129+
batch = _make_batch(num_x, groups, x.device)
130+
k = min(16, max(1, num_x - 1))
131+
132+
def cuda_fn():
133+
knn_graph(x, k=k, batch=batch, loop=False, use_triton=False)
134+
135+
for _ in range(5):
136+
cuda_fn()
137+
torch.cuda.synchronize()
138+
139+
benchmark(cuda_fn)
140+
print(f"[knn_graph][cuda] num_x={num_x} groups={groups} k={k}")
141+
142+
143+
@pytest.mark.parametrize('num_x', [256, 1024, 4096, 255])
144+
@pytest.mark.parametrize('num_groups', [1, 2, 4, 6, 8, 16, 24, 32])
145+
@pytest.mark.benchmark(group="knn_graph")
146+
def test_triton_knn_graph_benchmark_triton(benchmark, num_x, num_groups):
147+
torch.manual_seed(199)
148+
x = torch.randn(num_x, 8, device='cuda')
149+
groups = min(num_groups, x.size(0))
150+
batch = _make_batch(num_x, groups, x.device)
151+
k = min(16, max(1, num_x - 1))
152+
153+
def cuda_fn():
154+
return knn_graph(x, k=k, batch=batch, loop=False, use_triton=False)
155+
156+
def triton_fn():
157+
return knn_graph(x, k=k, batch=batch, loop=False, use_triton=True)
158+
159+
for i in range(5):
160+
if i == 0:
161+
out_cuda = cuda_fn()
162+
out_triton = triton_fn()
163+
assert to_set(out_cuda) == to_set(out_triton)
164+
else:
165+
triton_fn()
166+
torch.cuda.synchronize()
167+
168+
benchmark(triton_fn)
169+
print(f"[knn_graph][triton] num_x={num_x} groups={groups} k={k}")

‎benchmarks/test_benchmark_nearest.py‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
import importlib.util
2+
from itertools import product
3+
4+
import pytest
5+
import torch
6+
import torch_cluster as tc
7+
8+
nearest = tc.nearest
9+
10+
pytestmark = pytest.mark.skipif(
11+
not (torch.cuda.is_available() and importlib.util.find_spec('triton') is not None),
12+
reason='CUDA and Triton are required for Triton benchmark tests.',
13+
)
14+
15+
16+
def _make_batch(num_nodes: int, num_groups: int,
17+
device: torch.device) -> torch.Tensor:
18+
groups = max(1, min(num_groups, num_nodes))
19+
counts = torch.full((groups, ), num_nodes // groups, device=device,
20+
dtype=torch.long)
21+
remainder = num_nodes % groups
22+
if remainder:
23+
counts[:remainder] += 1
24+
return torch.repeat_interleave(torch.arange(groups, device=device),
25+
counts)
26+
27+
28+
@pytest.mark.parametrize('num_x,num_y,num_groups',
29+
((*p[0], p[1]) for p in product([(256, 128), (1024, 512), (4096, 2048), (255, 127),
30+
(256, 5), (1024, 5), (4096, 5), (255, 5)],
31+
[1, 2, 4, 8, 16, 32]) if p[1] <= min(p[0])))
32+
@pytest.mark.benchmark(group="nearest")
33+
def test_triton_nearest_benchmark_cuda(benchmark, num_x, num_y, num_groups):
34+
torch.manual_seed(123)
35+
x = torch.randn(num_x, 16, device='cuda')
36+
y = torch.randn(num_y, 16, device='cuda')
37+
groups = min(num_groups, x.size(0), y.size(0))
38+
batch_x = _make_batch(num_x, groups, x.device)
39+
batch_y = _make_batch(num_y, groups, y.device)
40+
41+
def cuda_fn():
42+
nearest(x, y, batch_x, batch_y, use_triton=False)
43+
44+
for _ in range(5):
45+
cuda_fn()
46+
torch.cuda.synchronize()
47+
48+
benchmark(cuda_fn)
49+
print(f"[nearest][cuda] num_x={num_x} num_y={num_y} groups={groups}")
50+
51+
52+
@pytest.mark.parametrize('num_x,num_y,num_groups',
53+
((*p[0], p[1]) for p in product([(256, 128), (1024, 512), (4096, 2048), (255, 127),
54+
(256, 5), (1024, 5), (4096, 5), (255, 5)],
55+
[1, 2, 4, 8, 16, 32]) if p[1] <= min(p[0])))
56+
@pytest.mark.benchmark(group="nearest")
57+
def test_triton_nearest_benchmark_triton(benchmark, num_x, num_y, num_groups):
58+
torch.manual_seed(123)
59+
x = torch.randn(num_x, 16, device='cuda')
60+
y = torch.randn(num_y, 16, device='cuda')
61+
groups = min(num_groups, x.size(0), y.size(0))
62+
batch_x = _make_batch(num_x, groups, x.device)
63+
batch_y = _make_batch(num_y, groups, y.device)
64+
65+
def cuda_fn():
66+
return nearest(x, y, batch_x, batch_y, use_triton=False)
67+
68+
def triton_fn():
69+
return nearest(x, y, batch_x, batch_y, use_triton=True)
70+
71+
for i in range(5):
72+
if i == 0:
73+
out_cuda = cuda_fn()
74+
out_triton = triton_fn()
75+
assert torch.equal(out_cuda, out_triton)
76+
else:
77+
triton_fn()
78+
torch.cuda.synchronize()
79+
80+
benchmark(triton_fn)
81+
print(f"[nearest][triton] num_x={num_x} num_y={num_y} groups={groups}")

‎csrc/cluster.h‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,5 +50,3 @@ TORCH_LIBRARY(torch_cluster, m) {
5050
m.def("neighbor_sampler(Tensor start, Tensor rowptr, int count, float factor) -> Tensor");
5151
m.def("cuda_version() -> int");
5252
}
53-
54-
}

‎csrc/fps.cpp‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,4 +41,7 @@ TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
4141
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
4242
m.impl("fps", &fps_cuda);
4343
}
44+
TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
45+
m.impl("fps", &fps_cuda);
46+
}
4447
#endif

‎csrc/graclus.cpp‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,4 +41,7 @@ TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
4141
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
4242
m.impl("graclus", &graclus_cuda);
4343
}
44+
TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
45+
m.impl("graclus", &graclus_cuda);
46+
}
4447
#endif

‎csrc/grid.cpp‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,4 +42,7 @@ TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
4242
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
4343
m.impl("grid", &grid_cuda);
4444
}
45+
TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
46+
m.impl("grid", &grid_cuda);
47+
}
4548
#endif

‎csrc/knn.cpp‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -37,10 +37,10 @@ CLUSTER_API torch::Tensor knn(torch::Tensor x, torch::Tensor y,
3737
}
3838
}
3939

40-
torch::Tensor knn_cpu_wrap(torch::Tensor x, torch::Tensor y,
40+
inline torch::Tensor knn_cpu_wrap(torch::Tensor x, torch::Tensor y,
4141
std::optional<torch::Tensor> ptr_x,
4242
std::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
43-
int64_t num_workers) {
43+
int64_t num_workers = 1) {
4444
TORCH_CHECK(!cosine, "`cosine` argument not supported on CPU");
4545
return knn_cpu(x, y, ptr_x, ptr_y, k, num_workers);
4646
}
@@ -50,7 +50,16 @@ TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
5050
}
5151

5252
#ifdef WITH_CUDA
53-
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
54-
m.impl("knn", &knn_cuda);
55-
}
53+
inline torch::Tensor knn_cuda_wrap(torch::Tensor x, torch::Tensor y,
54+
std::optional<torch::Tensor> ptr_x,
55+
std::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
56+
int64_t num_workers) {
57+
return knn_cuda(x, y, ptr_x, ptr_y, k, cosine);
58+
}
59+
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
60+
m.impl("knn", &knn_cuda_wrap);
61+
}
62+
TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
63+
m.impl("knn", &knn_cuda_wrap);
64+
}
5665
#endif

‎csrc/nearest.cpp‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,4 +37,7 @@ CLUSTER_API torch::Tensor nearest(torch::Tensor x, torch::Tensor y, torch::Tenso
3737
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
3838
m.impl("nearest", &nearest_cuda);
3939
}
40+
TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
41+
m.impl("nearest", &nearest_cuda);
42+
}
4043
#endif

‎csrc/radius.cpp‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,18 @@ TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
4141
}
4242

4343
#ifdef WITH_CUDA
44+
inline torch::Tensor radius_cuda_wrap(torch::Tensor x, torch::Tensor y,
45+
std::optional<torch::Tensor> ptr_x,
46+
std::optional<torch::Tensor> ptr_y, double r,
47+
int64_t max_num_neighbors, int64_t num_workers,
48+
bool ignore_same_index) {
49+
return radius_cuda(x, y, ptr_x, ptr_y, r, max_num_neighbors, ignore_same_index);
50+
}
51+
4452
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
45-
m.impl("radius", &radius_cuda);
53+
m.impl("radius", &radius_cuda_wrap);
54+
}
55+
TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
56+
m.impl("radius", &radius_cuda_wrap);
4657
}
4758
#endif

‎csrc/rw.cpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,4 +42,8 @@ TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
4242
TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
4343
m.impl("random_walk", &random_walk_cuda);
4444
}
45+
46+
TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
47+
m.impl("random_walk", &random_walk_cuda);
48+
}
4549
#endif

0 commit comments

Comments
 (0)