Skip to content

Commit 5af902d

Browse files
denix56denix56
committed
Add Triton support
1 parent e903e66 commit 5af902d

16 files changed

Lines changed: 702 additions & 527 deletions

File tree

‎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/knn.cpp‎

Lines changed: 3 additions & 3 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,7 @@ TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
5050
}
5151

5252
#ifdef WITH_CUDA
53-
torch::Tensor knn_cuda_wrap(torch::Tensor x, torch::Tensor y,
53+
inline torch::Tensor knn_cuda_wrap(torch::Tensor x, torch::Tensor y,
5454
std::optional<torch::Tensor> ptr_x,
5555
std::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
5656
int64_t num_workers) {

‎csrc/radius.cpp‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,10 +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);
4654
}
4755
TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
48-
m.impl("radius", &radius_cuda);
56+
m.impl("radius", &radius_cuda_wrap);
4957
}
5058
#endif

‎test/test_knn.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import importlib.util
12
from itertools import product
23

34
import pytest
@@ -6,6 +7,9 @@
67
from torch_cluster import knn, knn_graph
78
from torch_cluster.testing import devices, grad_dtypes, tensor, triton_wrap
89

10+
HAS_CUDA = torch.cuda.is_available()
11+
HAS_TRITON = importlib.util.find_spec('triton') is not None
12+
913

1014
def to_set(edge_index):
1115
return set([(i, j) for i, j in edge_index.t().tolist()])
@@ -73,6 +77,7 @@ def test_knn_graph(dtype, device, use_triton):
7377

7478
@pytest.mark.parametrize('dtype,device,use_triton', triton_wrap(product([torch.float], devices)))
7579
def test_knn_graph_large(dtype, device, use_triton):
80+
torch.manual_seed(29)
7681
x = torch.randn(1000, 3, dtype=dtype, device=device)
7782

7883
edge_index = knn_graph(x, k=5, flow='target_to_source', loop=True, use_triton=use_triton)
@@ -82,3 +87,33 @@ def test_knn_graph_large(dtype, device, use_triton):
8287
truth = set([(i, j) for i, ns in enumerate(col) for j in ns])
8388

8489
assert to_set(edge_index.cpu()) == truth
90+
91+
92+
@pytest.mark.skipif(not (HAS_CUDA and HAS_TRITON),
93+
reason='CUDA and Triton are required for Triton parity tests.')
94+
def test_knn_triton_matches_cuda():
95+
torch.manual_seed(42)
96+
x = torch.randn(128, 16, device='cuda')
97+
y = torch.randn(64, 16, device='cuda')
98+
batch_x = torch.zeros(x.size(0), dtype=torch.long, device='cuda')
99+
batch_y = torch.zeros(y.size(0), dtype=torch.long, device='cuda')
100+
101+
out_cuda = knn(x, y, k=8, batch_x=batch_x, batch_y=batch_y, use_triton=False)
102+
out_triton = knn(x, y, k=8, batch_x=batch_x, batch_y=batch_y, use_triton=True)
103+
assert to_set(out_cuda) == to_set(out_triton)
104+
105+
out_cuda = knn(x, y, k=8, batch_x=batch_x, batch_y=batch_y, cosine=True, use_triton=False)
106+
out_triton = knn(x, y, k=8, batch_x=batch_x, batch_y=batch_y, cosine=True, use_triton=True)
107+
assert to_set(out_cuda) == to_set(out_triton)
108+
109+
110+
@pytest.mark.skipif(not (HAS_CUDA and HAS_TRITON),
111+
reason='CUDA and Triton are required for Triton parity tests.')
112+
def test_knn_graph_triton_matches_cuda():
113+
torch.manual_seed(1)
114+
x = torch.randn(64, 8, device='cuda')
115+
batch = torch.zeros(x.size(0), dtype=torch.long, device='cuda')
116+
117+
out_cuda = knn_graph(x, k=4, batch=batch, loop=False, use_triton=False)
118+
out_triton = knn_graph(x, k=4, batch=batch, loop=False, use_triton=True)
119+
assert to_set(out_cuda) == to_set(out_triton)

‎test/test_nearest.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import importlib.util
12
from itertools import product
23

34
import pytest
@@ -6,6 +7,31 @@
67
from torch_cluster.testing import devices, grad_dtypes, tensor, triton_wrap
78

89

10+
@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 parity tests.',
13+
)
14+
def test_nearest_triton_matches_cuda():
15+
torch.manual_seed(123)
16+
x = torch.randn(128, 8, device='cuda')
17+
y = torch.randn(32, 8, device='cuda')
18+
batch_x = torch.zeros(x.size(0), dtype=torch.long, device='cuda')
19+
batch_y = torch.zeros(y.size(0), dtype=torch.long, device='cuda')
20+
21+
out_cuda = nearest(x, y, batch_x, batch_y, use_triton=False)
22+
out_triton = nearest(x, y, batch_x, batch_y, use_triton=True)
23+
assert torch.equal(out_cuda, out_triton)
24+
25+
batch_x = torch.zeros(x.size(0), dtype=torch.long, device='cuda')
26+
batch_y = torch.zeros(y.size(0), dtype=torch.long, device='cuda')
27+
batch_x[x.size(0) // 2:] = 1
28+
batch_y[y.size(0) // 2:] = 1
29+
30+
out_cuda = nearest(x, y, batch_x, batch_y, use_triton=False)
31+
out_triton = nearest(x, y, batch_x, batch_y, use_triton=True)
32+
assert torch.equal(out_cuda, out_triton)
33+
34+
935
@pytest.mark.parametrize('dtype,device,use_triton', triton_wrap(product(grad_dtypes, devices)))
1036
def test_nearest(dtype, device, use_triton):
1137
x = tensor([

0 commit comments

Comments
 (0)