Skip to content

Commit 1916626

Browse files
denix56denix56
committed
Use sorted batch pointers for Triton masks
1 parent 6dabb04 commit 1916626

4 files changed

Lines changed: 963 additions & 0 deletions

File tree

‎setup.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,7 @@ def get_extensions():
110110

111111
install_requires = [
112112
'scipy',
113+
'triton==3.6.0',
113114
]
114115

115116
test_requires = [

‎test/test_triton.py‎

Lines changed: 276 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,276 @@
1+
import importlib.util
2+
import time
3+
4+
import pytest
5+
import torch
6+
7+
from torch_cluster import (
8+
fps,
9+
graclus_cluster,
10+
grid_cluster,
11+
knn,
12+
knn_graph,
13+
nearest,
14+
radius,
15+
radius_graph,
16+
random_walk,
17+
)
18+
19+
HAS_CUDA = torch.cuda.is_available()
20+
HAS_TRITON = importlib.util.find_spec('triton') is not None
21+
22+
if HAS_TRITON:
23+
from torch_cluster import (
24+
fps__triton,
25+
graclus_cluster__triton,
26+
grid_cluster__triton,
27+
knn__triton,
28+
knn_graph__triton,
29+
nearest__triton,
30+
radius__triton,
31+
radius_graph__triton,
32+
random_walk__triton,
33+
)
34+
35+
pytestmark = pytest.mark.skipif(
36+
not (HAS_CUDA and HAS_TRITON),
37+
reason='CUDA and Triton are required for Triton parity tests.',
38+
)
39+
40+
41+
def _sort_edge_index(edge_index: torch.Tensor, num_nodes: int) -> torch.Tensor:
42+
key = edge_index[0] * num_nodes + edge_index[1]
43+
perm = key.argsort()
44+
return edge_index[:, perm]
45+
46+
47+
def _benchmark(fn, warmup: int = 3, iters: int = 10) -> float:
48+
for _ in range(warmup):
49+
fn()
50+
torch.cuda.synchronize()
51+
start = time.perf_counter()
52+
for _ in range(iters):
53+
fn()
54+
torch.cuda.synchronize()
55+
return (time.perf_counter() - start) / iters
56+
57+
58+
def test_knn_triton_matches_cuda():
59+
torch.manual_seed(42)
60+
x = torch.randn(128, 16, device='cuda')
61+
y = torch.randn(64, 16, device='cuda')
62+
batch_x = torch.zeros(x.size(0), dtype=torch.long, device='cuda')
63+
batch_y = torch.zeros(y.size(0), dtype=torch.long, device='cuda')
64+
65+
out_cuda = knn(x, y, k=8, batch_x=batch_x, batch_y=batch_y)
66+
out_triton = knn__triton(x, y, k=8, batch_x=batch_x, batch_y=batch_y)
67+
assert torch.equal(
68+
_sort_edge_index(out_cuda, x.size(0)),
69+
_sort_edge_index(out_triton, x.size(0)),
70+
)
71+
72+
out_cuda = knn(x, y, k=8, batch_x=batch_x, batch_y=batch_y, cosine=True)
73+
out_triton = knn__triton(x,
74+
y,
75+
k=8,
76+
batch_x=batch_x,
77+
batch_y=batch_y,
78+
cosine=True)
79+
assert torch.equal(
80+
_sort_edge_index(out_cuda, x.size(0)),
81+
_sort_edge_index(out_triton, x.size(0)),
82+
)
83+
84+
85+
def test_knn_graph_triton_matches_cuda():
86+
torch.manual_seed(1)
87+
x = torch.randn(64, 8, device='cuda')
88+
batch = torch.zeros(x.size(0), dtype=torch.long, device='cuda')
89+
90+
out_cuda = knn_graph(x, k=4, batch=batch, loop=False)
91+
out_triton = knn_graph__triton(x, k=4, batch=batch, loop=False)
92+
assert torch.equal(
93+
_sort_edge_index(out_cuda, x.size(0)),
94+
_sort_edge_index(out_triton, x.size(0)),
95+
)
96+
97+
98+
def test_radius_triton_matches_cuda():
99+
torch.manual_seed(7)
100+
x = torch.randn(128, 3, device='cuda')
101+
y = torch.randn(64, 3, device='cuda')
102+
batch_x = torch.zeros(x.size(0), dtype=torch.long, device='cuda')
103+
batch_y = torch.zeros(y.size(0), dtype=torch.long, device='cuda')
104+
105+
out_cuda = radius(x,
106+
y,
107+
r=1.5,
108+
batch_x=batch_x,
109+
batch_y=batch_y,
110+
max_num_neighbors=x.size(0))
111+
out_triton = radius__triton(x,
112+
y,
113+
r=1.5,
114+
batch_x=batch_x,
115+
batch_y=batch_y,
116+
max_num_neighbors=x.size(0))
117+
assert torch.equal(
118+
_sort_edge_index(out_cuda, x.size(0)),
119+
_sort_edge_index(out_triton, x.size(0)),
120+
)
121+
122+
out_cuda = radius(x,
123+
y,
124+
r=1.5,
125+
batch_x=batch_x,
126+
batch_y=batch_y,
127+
max_num_neighbors=x.size(0),
128+
ignore_same_index=True)
129+
out_triton = radius__triton(x,
130+
y,
131+
r=1.5,
132+
batch_x=batch_x,
133+
batch_y=batch_y,
134+
max_num_neighbors=x.size(0),
135+
ignore_same_index=True)
136+
assert torch.equal(
137+
_sort_edge_index(out_cuda, x.size(0)),
138+
_sort_edge_index(out_triton, x.size(0)),
139+
)
140+
141+
142+
def test_radius_graph_triton_matches_cuda():
143+
torch.manual_seed(3)
144+
x = torch.randn(64, 4, device='cuda')
145+
batch = torch.zeros(x.size(0), dtype=torch.long, device='cuda')
146+
147+
out_cuda = radius_graph(x, r=2.0, batch=batch, loop=False)
148+
out_triton = radius_graph__triton(x, r=2.0, batch=batch, loop=False)
149+
assert torch.equal(
150+
_sort_edge_index(out_cuda, x.size(0)),
151+
_sort_edge_index(out_triton, x.size(0)),
152+
)
153+
154+
155+
def test_nearest_triton_matches_cuda():
156+
torch.manual_seed(123)
157+
x = torch.randn(128, 8, device='cuda')
158+
y = torch.randn(32, 8, device='cuda')
159+
batch_x = torch.zeros(x.size(0), dtype=torch.long, device='cuda')
160+
batch_y = torch.zeros(y.size(0), dtype=torch.long, device='cuda')
161+
162+
out_cuda = nearest(x, y, batch_x, batch_y)
163+
out_triton = nearest__triton(x, y, batch_x, batch_y)
164+
assert torch.equal(out_cuda, out_triton)
165+
166+
167+
def test_grid_cluster_triton_matches_cuda():
168+
torch.manual_seed(5)
169+
pos = torch.randn(128, 3, device='cuda')
170+
size = torch.tensor([0.5, 0.5, 0.5], device='cuda')
171+
172+
out_cuda = grid_cluster(pos, size)
173+
out_triton = grid_cluster__triton(pos, size)
174+
assert torch.equal(out_cuda, out_triton)
175+
176+
177+
def test_fps_triton_matches_cuda():
178+
torch.manual_seed(11)
179+
src = torch.randn(256, 3, device='cuda')
180+
batch = torch.zeros(src.size(0), dtype=torch.long, device='cuda')
181+
182+
out_cuda = fps(src, batch=batch, ratio=0.25, random_start=False)
183+
out_triton = fps__triton(src, batch=batch, ratio=0.25, random_start=False)
184+
assert torch.equal(out_cuda, out_triton)
185+
186+
187+
def test_graclus_triton_matches_cuda_on_empty_graph():
188+
row = torch.empty(0, dtype=torch.long, device='cuda')
189+
col = torch.empty(0, dtype=torch.long, device='cuda')
190+
out_cuda = graclus_cluster(row, col, num_nodes=4)
191+
out_triton = graclus_cluster__triton(row, col, num_nodes=4)
192+
assert torch.equal(out_cuda, out_triton)
193+
194+
195+
def test_random_walk_triton_matches_cuda_on_deterministic_graph():
196+
row = torch.tensor([0, 1, 2, 3], device='cuda')
197+
col = torch.tensor([1, 2, 3, 0], device='cuda')
198+
start = torch.tensor([0, 1, 2, 3], device='cuda')
199+
200+
node_cuda, edge_cuda = random_walk(row,
201+
col,
202+
start,
203+
walk_length=4,
204+
return_edge_indices=True)
205+
node_triton, edge_triton = random_walk__triton(row,
206+
col,
207+
start,
208+
walk_length=4,
209+
return_edge_indices=True)
210+
assert torch.equal(node_cuda, node_triton)
211+
assert torch.equal(edge_cuda, edge_triton)
212+
213+
214+
def test_triton_edge_cases():
215+
empty = torch.empty(0, 2, device='cuda')
216+
out_cuda = knn(empty, empty, k=2)
217+
out_triton = knn__triton(empty, empty, k=2)
218+
assert torch.equal(out_cuda, out_triton)
219+
220+
out_cuda = radius(empty, empty, r=1.0)
221+
out_triton = radius__triton(empty, empty, r=1.0)
222+
assert torch.equal(out_cuda, out_triton)
223+
224+
x = torch.tensor([1.0, 2.0, 3.0], device='cuda')
225+
y = torch.tensor([1.5, 2.5], device='cuda')
226+
out_cuda = knn(x, y, k=1)
227+
out_triton = knn__triton(x, y, k=1)
228+
assert torch.equal(out_cuda, out_triton)
229+
230+
231+
@pytest.mark.parametrize('num_x,num_y', [(256, 128), (1024, 512), (4096, 2048)])
232+
def test_triton_knn_performance(num_x, num_y):
233+
torch.manual_seed(99)
234+
x = torch.randn(num_x, 16, device='cuda')
235+
y = torch.randn(num_y, 16, device='cuda')
236+
batch_x = torch.zeros(num_x, dtype=torch.long, device='cuda')
237+
batch_y = torch.zeros(num_y, dtype=torch.long, device='cuda')
238+
239+
def cuda_fn():
240+
knn(x, y, k=16, batch_x=batch_x, batch_y=batch_y)
241+
242+
def triton_fn():
243+
knn__triton(x, y, k=16, batch_x=batch_x, batch_y=batch_y)
244+
245+
cuda_time = _benchmark(cuda_fn)
246+
triton_time = _benchmark(triton_fn)
247+
assert triton_time <= cuda_time * 20
248+
249+
250+
@pytest.mark.parametrize('num_x,num_y', [(256, 128), (1024, 512), (4096, 2048)])
251+
def test_triton_radius_performance(num_x, num_y):
252+
torch.manual_seed(199)
253+
x = torch.randn(num_x, 8, device='cuda')
254+
y = torch.randn(num_y, 8, device='cuda')
255+
batch_x = torch.zeros(num_x, dtype=torch.long, device='cuda')
256+
batch_y = torch.zeros(num_y, dtype=torch.long, device='cuda')
257+
258+
def cuda_fn():
259+
radius(x,
260+
y,
261+
r=0.5,
262+
batch_x=batch_x,
263+
batch_y=batch_y,
264+
max_num_neighbors=num_x)
265+
266+
def triton_fn():
267+
radius__triton(x,
268+
y,
269+
r=0.5,
270+
batch_x=batch_x,
271+
batch_y=batch_y,
272+
max_num_neighbors=num_x)
273+
274+
cuda_time = _benchmark(cuda_fn)
275+
triton_time = _benchmark(triton_fn)
276+
assert triton_time <= cuda_time * 20

‎torch_cluster/__init__.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import importlib
2+
import importlib.util
23
import os.path as osp
34

45
import torch
@@ -45,6 +46,22 @@
4546
from .rw import random_walk # noqa
4647
from .sampler import neighbor_sampler # noqa
4748

49+
_HAS_TRITON = importlib.util.find_spec('triton') is not None
50+
51+
if _HAS_TRITON:
52+
from .triton import ( # noqa
53+
fps__triton,
54+
graclus_cluster__triton,
55+
grid_cluster__triton,
56+
knn__triton,
57+
knn_graph__triton,
58+
nearest__triton,
59+
neighbor_sampler__triton,
60+
radius__triton,
61+
radius_graph__triton,
62+
random_walk__triton,
63+
)
64+
4865
__all__ = [
4966
'graclus_cluster',
5067
'grid_cluster',
@@ -58,3 +75,17 @@
5875
'neighbor_sampler',
5976
'__version__',
6077
]
78+
79+
if _HAS_TRITON:
80+
__all__ += [
81+
'fps__triton',
82+
'graclus_cluster__triton',
83+
'grid_cluster__triton',
84+
'knn__triton',
85+
'knn_graph__triton',
86+
'nearest__triton',
87+
'neighbor_sampler__triton',
88+
'radius__triton',
89+
'radius_graph__triton',
90+
'random_walk__triton',
91+
]

0 commit comments

Comments
 (0)