Skip to content

Commit 4b30dc5

Browse files
denix56denix56
committed
Fix mismatches with CUDA, added radius search support
1 parent 2be1748 commit 4b30dc5

11 files changed

Lines changed: 1200 additions & 248 deletions

File tree

‎benchmarks/test_benchmark_knn.py‎

Lines changed: 181 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -11,16 +11,21 @@
1111

1212
pytestmark = pytest.mark.skipif(
1313
not (
14-
torch.cuda.is_available()
14+
torch.ops.torch_cluster.cuda_version() != -1
1515
and importlib.util.find_spec('triton') is not None
1616
),
1717
reason='CUDA and Triton are required for Triton benchmark tests.',
1818
)
1919

2020
KNN_SIZES = [
2121
(256, 128),
22+
(512, 256),
2223
(1024, 512),
24+
(2048, 1024),
2325
(4096, 2048),
26+
(8192, 4096),
27+
(8192, 8192),
28+
(8201, 4103),
2429
(255, 127),
2530
(256, 5),
2631
(1024, 5),
@@ -29,6 +34,8 @@
2934
]
3035
KNN_GROUPS = [1, 2, 4, 8, 16, 32]
3136

37+
FEATURES = [8, 64, 200]
38+
3239

3340
def to_set(edge_index):
3441
return set([(i, j) for i, j in edge_index.t().tolist()])
@@ -55,19 +62,29 @@ def _make_batch(
5562
)
5663

5764

58-
@pytest.mark.parametrize(
59-
'num_x,num_y,num_groups',
60-
(
61-
(*p[0], p[1])
62-
for p in product(KNN_SIZES, KNN_GROUPS)
65+
def _knn_param_grid():
66+
return (
67+
(*p[0], p[1], p[2])
68+
for p in product(KNN_SIZES, KNN_GROUPS, FEATURES)
6369
if p[1] <= min(p[0])
64-
),
70+
)
71+
72+
73+
@pytest.mark.parametrize(
74+
'num_x,num_y,num_groups,num_features',
75+
_knn_param_grid(),
6576
)
6677
@pytest.mark.benchmark(group="knn")
67-
def test_triton_knn_benchmark_cuda(benchmark, num_x, num_y, num_groups):
78+
def test_triton_knn_benchmark_cuda(
79+
benchmark,
80+
num_x,
81+
num_y,
82+
num_groups,
83+
num_features,
84+
):
6885
torch.manual_seed(99)
69-
x = torch.randn(num_x, 16, device='cuda')
70-
y = torch.randn(num_y, 16, device='cuda')
86+
x = torch.randn(num_x, num_features, device='cuda')
87+
y = torch.randn(num_y, num_features, device='cuda')
7188
groups = min(num_groups, x.size(0), y.size(0))
7289
batch_x = _make_batch(num_x, groups, x.device)
7390
batch_y = _make_batch(num_y, groups, y.device)
@@ -94,23 +111,60 @@ def cuda_fn():
94111

95112

96113
@pytest.mark.parametrize(
97-
'num_x,num_y,num_groups',
98-
(
99-
(*p[0], p[1])
100-
for p in product(KNN_SIZES, KNN_GROUPS)
101-
if p[1] <= min(p[0])
102-
),
114+
'num_x,num_y,num_groups,num_features',
115+
_knn_param_grid(),
103116
)
104-
@pytest.mark.benchmark(group="knn")
117+
@pytest.mark.benchmark(group="knn_cosine")
118+
def test_triton_knn_benchmark_cuda_cosine(
119+
benchmark,
120+
num_x,
121+
num_y,
122+
num_groups,
123+
num_features,
124+
):
125+
torch.manual_seed(99)
126+
x = torch.randn(num_x, num_features, device='cuda')
127+
y = torch.randn(num_y, num_features, device='cuda')
128+
groups = min(num_groups, x.size(0), y.size(0))
129+
batch_x = _make_batch(num_x, groups, x.device)
130+
batch_y = _make_batch(num_y, groups, y.device)
131+
132+
def cuda_fn():
133+
knn(
134+
x,
135+
y,
136+
k=16,
137+
batch_x=batch_x,
138+
batch_y=batch_y,
139+
cosine=True,
140+
use_triton=False,
141+
)
142+
143+
for _ in range(5):
144+
cuda_fn()
145+
torch.cuda.synchronize()
146+
147+
benchmark(cuda_fn)
148+
print(
149+
f"[knn][cuda] num_x={num_x} num_y={num_y} groups={groups}"
150+
)
151+
152+
153+
@pytest.mark.parametrize(
154+
'num_x,num_y,num_groups,num_features',
155+
_knn_param_grid(),
156+
)
157+
@pytest.mark.benchmark(group="knn_cosine")
105158
def test_triton_knn_benchmark_triton_cosine(
106159
benchmark,
107160
num_x,
108161
num_y,
109162
num_groups,
163+
num_features,
110164
):
111165
torch.manual_seed(99)
112-
x = torch.randn(num_x, 16, device='cuda')
113-
y = torch.randn(num_y, 16, device='cuda')
166+
x = torch.randn(num_x, num_features, device='cuda')
167+
y = torch.randn(num_y, num_features, device='cuda')
114168
groups = min(num_groups, x.size(0), y.size(0))
115169
batch_x = _make_batch(num_x, groups, x.device)
116170
batch_y = _make_batch(num_y, groups, y.device)
@@ -141,7 +195,11 @@ def triton_fn():
141195
if i == 0:
142196
out_cuda = cuda_fn()
143197
out_triton = triton_fn()
144-
assert to_set(out_cuda) == to_set(out_triton)
198+
for a, b in zip(
199+
sorted(list(to_set(out_cuda))),
200+
sorted(list(to_set(out_triton))),
201+
):
202+
assert a == b
145203
else:
146204
triton_fn()
147205
torch.cuda.synchronize()
@@ -153,18 +211,20 @@ def triton_fn():
153211

154212

155213
@pytest.mark.parametrize(
156-
'num_x,num_y,num_groups',
157-
(
158-
(*p[0], p[1])
159-
for p in product(KNN_SIZES, KNN_GROUPS)
160-
if p[1] <= min(p[0])
161-
),
214+
'num_x,num_y,num_groups,num_features',
215+
_knn_param_grid(),
162216
)
163217
@pytest.mark.benchmark(group="knn")
164-
def test_triton_knn_benchmark_triton(benchmark, num_x, num_y, num_groups):
218+
def test_triton_knn_benchmark_triton(
219+
benchmark,
220+
num_x,
221+
num_y,
222+
num_groups,
223+
num_features,
224+
):
165225
torch.manual_seed(99)
166-
x = torch.randn(num_x, 16, device='cuda')
167-
y = torch.randn(num_y, 16, device='cuda')
226+
x = torch.randn(num_x, num_features, device='cuda')
227+
y = torch.randn(num_y, num_features, device='cuda')
168228
groups = min(num_groups, x.size(0), y.size(0))
169229
batch_x = _make_batch(num_x, groups, x.device)
170230
batch_y = _make_batch(num_y, groups, y.device)
@@ -195,7 +255,11 @@ def triton_fn():
195255
if i == 0:
196256
out_cuda = cuda_fn()
197257
out_triton = triton_fn()
198-
assert to_set(out_cuda) == to_set(out_triton)
258+
for a, b in zip(
259+
sorted(list(to_set(out_cuda))),
260+
sorted(list(to_set(out_triton))),
261+
):
262+
assert a == b
199263
else:
200264
triton_fn()
201265
torch.cuda.synchronize()
@@ -206,7 +270,7 @@ def triton_fn():
206270
)
207271

208272

209-
@pytest.mark.parametrize('num_x', [256, 1024, 4096, 255])
273+
@pytest.mark.parametrize('num_x', [256, 1024, 4096, 8192, 255])
210274
@pytest.mark.parametrize('num_groups', [1, 2, 4, 6, 8, 16, 24, 32])
211275
@pytest.mark.benchmark(group="knn_graph")
212276
def test_triton_knn_graph_benchmark_cuda(benchmark, num_x, num_groups):
@@ -229,7 +293,7 @@ def cuda_fn():
229293
)
230294

231295

232-
@pytest.mark.parametrize('num_x', [256, 1024, 4096, 255])
296+
@pytest.mark.parametrize('num_x', [256, 1024, 4096, 8192, 255])
233297
@pytest.mark.parametrize('num_groups', [1, 2, 4, 6, 8, 16, 24, 32])
234298
@pytest.mark.benchmark(group="knn_graph")
235299
def test_triton_knn_graph_benchmark_triton(benchmark, num_x, num_groups):
@@ -261,12 +325,93 @@ def triton_fn():
261325
if i == 0:
262326
out_cuda = cuda_fn()
263327
out_triton = triton_fn()
264-
assert to_set(out_cuda) == to_set(out_triton)
328+
for a, b in zip(
329+
sorted(list(to_set(out_cuda))),
330+
sorted(list(to_set(out_triton))),
331+
):
332+
assert a == b
265333
else:
266334
triton_fn()
267335
torch.cuda.synchronize()
268336

269337
benchmark(triton_fn)
270-
print(
271-
f"[knn_graph][triton] num_x={num_x} groups={groups} k={k}"
272-
)
338+
print(f"[knn_graph][triton] num_x={num_x} groups={groups} k={k}")
339+
340+
341+
@pytest.mark.parametrize('num_x', [256, 1024, 4096, 8192, 255])
342+
@pytest.mark.parametrize('num_groups', [1, 2, 4, 6, 8, 16, 24, 32])
343+
@pytest.mark.benchmark(group="knn_graph_cosine")
344+
def test_triton_knn_graph_benchmark_cuda_cosine(benchmark, num_x, num_groups):
345+
torch.manual_seed(199)
346+
x = torch.randn(num_x, 8, device='cuda')
347+
groups = min(num_groups, x.size(0))
348+
batch = _make_batch(num_x, groups, x.device)
349+
k = min(16, max(1, num_x - 1))
350+
351+
def cuda_fn():
352+
knn_graph(
353+
x,
354+
k=k,
355+
batch=batch,
356+
loop=False,
357+
use_triton=False,
358+
cosine=True,
359+
)
360+
361+
for _ in range(5):
362+
cuda_fn()
363+
torch.cuda.synchronize()
364+
365+
benchmark(cuda_fn)
366+
print(f"[knn_graph][cuda] num_x={num_x} groups={groups} k={k}")
367+
368+
369+
@pytest.mark.parametrize('num_x', [256, 1024, 4096, 8192, 255])
370+
@pytest.mark.parametrize('num_groups', [1, 2, 4, 6, 8, 16, 24, 32])
371+
@pytest.mark.benchmark(group="knn_graph_cosine")
372+
def test_triton_knn_graph_benchmark_triton_cosine(
373+
benchmark,
374+
num_x,
375+
num_groups,
376+
):
377+
torch.manual_seed(199)
378+
x = torch.randn(num_x, 8, device='cuda')
379+
groups = min(num_groups, x.size(0))
380+
batch = _make_batch(num_x, groups, x.device)
381+
k = min(16, max(1, num_x - 1))
382+
383+
def cuda_fn():
384+
return knn_graph(
385+
x,
386+
k=k,
387+
batch=batch,
388+
loop=False,
389+
use_triton=False,
390+
cosine=True,
391+
)
392+
393+
def triton_fn():
394+
return knn_graph(
395+
x,
396+
k=k,
397+
batch=batch,
398+
loop=False,
399+
use_triton=True,
400+
cosine=True,
401+
)
402+
403+
for i in range(5):
404+
if i == 0:
405+
out_cuda = cuda_fn()
406+
out_triton = triton_fn()
407+
for a, b in zip(
408+
sorted(list(to_set(out_cuda))),
409+
sorted(list(to_set(out_triton))),
410+
):
411+
assert a == b
412+
else:
413+
triton_fn()
414+
torch.cuda.synchronize()
415+
416+
benchmark(triton_fn)
417+
print(f"[knn_graph][triton] num_x={num_x} groups={groups} k={k}")

0 commit comments

Comments
 (0)