Skip to content

Commit 29e0b41

Browse files
denix56denix56
committed
Fix flake8 errors
remove triton from dependencies
1 parent 56c55ce commit 29e0b41

14 files changed

Lines changed: 617 additions & 200 deletions

File tree

‎benchmarks/test_benchmark_knn.py‎

Lines changed: 141 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -10,31 +10,59 @@
1010

1111

1212
pytestmark = pytest.mark.skipif(
13-
not (torch.cuda.is_available() and importlib.util.find_spec('triton') is not None),
13+
not (
14+
torch.cuda.is_available()
15+
and importlib.util.find_spec('triton') is not None
16+
),
1417
reason='CUDA and Triton are required for Triton benchmark tests.',
1518
)
1619

20+
KNN_SIZES = [
21+
(256, 128),
22+
(1024, 512),
23+
(4096, 2048),
24+
(255, 127),
25+
(256, 5),
26+
(1024, 5),
27+
(4096, 5),
28+
(255, 5),
29+
]
30+
KNN_GROUPS = [1, 2, 4, 8, 16, 32]
31+
1732

1833
def to_set(edge_index):
1934
return set([(i, j) for i, j in edge_index.t().tolist()])
2035

2136

22-
def _make_batch(num_nodes: int, num_groups: int,
23-
device: torch.device) -> torch.Tensor:
37+
def _make_batch(
38+
num_nodes: int,
39+
num_groups: int,
40+
device: torch.device,
41+
) -> torch.Tensor:
2442
groups = max(1, min(num_groups, num_nodes))
25-
counts = torch.full((groups, ), num_nodes // groups, device=device,
26-
dtype=torch.long)
43+
counts = torch.full(
44+
(groups,),
45+
num_nodes // groups,
46+
device=device,
47+
dtype=torch.long,
48+
)
2749
remainder = num_nodes % groups
2850
if remainder:
2951
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])))
52+
return torch.repeat_interleave(
53+
torch.arange(groups, device=device),
54+
counts,
55+
)
56+
57+
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)
63+
if p[1] <= min(p[0])
64+
),
65+
)
3866
@pytest.mark.benchmark(group="knn")
3967
def test_triton_knn_benchmark_cuda(benchmark, num_x, num_y, num_groups):
4068
torch.manual_seed(99)
@@ -45,22 +73,41 @@ def test_triton_knn_benchmark_cuda(benchmark, num_x, num_y, num_groups):
4573
batch_y = _make_batch(num_y, groups, y.device)
4674

4775
def cuda_fn():
48-
knn(x, y, k=16, batch_x=batch_x, batch_y=batch_y, cosine=False, use_triton=False)
76+
knn(
77+
x,
78+
y,
79+
k=16,
80+
batch_x=batch_x,
81+
batch_y=batch_y,
82+
cosine=False,
83+
use_triton=False,
84+
)
4985

5086
for _ in range(5):
5187
cuda_fn()
5288
torch.cuda.synchronize()
5389

5490
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])))
91+
print(
92+
f"[knn][cuda] num_x={num_x} num_y={num_y} groups={groups}"
93+
)
94+
95+
96+
@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+
),
103+
)
62104
@pytest.mark.benchmark(group="knn")
63-
def test_triton_knn_benchmark_triton_cosine(benchmark, num_x, num_y, num_groups):
105+
def test_triton_knn_benchmark_triton_cosine(
106+
benchmark,
107+
num_x,
108+
num_y,
109+
num_groups,
110+
):
64111
torch.manual_seed(99)
65112
x = torch.randn(num_x, 16, device='cuda')
66113
y = torch.randn(num_y, 16, device='cuda')
@@ -69,10 +116,26 @@ def test_triton_knn_benchmark_triton_cosine(benchmark, num_x, num_y, num_groups)
69116
batch_y = _make_batch(num_y, groups, y.device)
70117

71118
def cuda_fn():
72-
return knn(x, y, k=16, batch_x=batch_x, batch_y=batch_y, cosine=True, use_triton=False)
119+
return knn(
120+
x,
121+
y,
122+
k=16,
123+
batch_x=batch_x,
124+
batch_y=batch_y,
125+
cosine=True,
126+
use_triton=False,
127+
)
73128

74129
def triton_fn():
75-
return knn(x, y, k=16, batch_x=batch_x, batch_y=batch_y, cosine=True, use_triton=True)
130+
return knn(
131+
x,
132+
y,
133+
k=16,
134+
batch_x=batch_x,
135+
batch_y=batch_y,
136+
cosine=True,
137+
use_triton=True,
138+
)
76139

77140
for i in range(5):
78141
if i == 0:
@@ -84,13 +147,19 @@ def triton_fn():
84147
torch.cuda.synchronize()
85148

86149
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])))
150+
print(
151+
f"[knn][triton] num_x={num_x} num_y={num_y} groups={groups}"
152+
)
153+
154+
155+
@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+
),
162+
)
94163
@pytest.mark.benchmark(group="knn")
95164
def test_triton_knn_benchmark_triton(benchmark, num_x, num_y, num_groups):
96165
torch.manual_seed(99)
@@ -101,10 +170,26 @@ def test_triton_knn_benchmark_triton(benchmark, num_x, num_y, num_groups):
101170
batch_y = _make_batch(num_y, groups, y.device)
102171

103172
def cuda_fn():
104-
return knn(x, y, k=16, batch_x=batch_x, batch_y=batch_y, cosine=False, use_triton=False)
173+
return knn(
174+
x,
175+
y,
176+
k=16,
177+
batch_x=batch_x,
178+
batch_y=batch_y,
179+
cosine=False,
180+
use_triton=False,
181+
)
105182

106183
def triton_fn():
107-
return knn(x, y, k=16, batch_x=batch_x, batch_y=batch_y, cosine=False, use_triton=True)
184+
return knn(
185+
x,
186+
y,
187+
k=16,
188+
batch_x=batch_x,
189+
batch_y=batch_y,
190+
cosine=False,
191+
use_triton=True,
192+
)
108193

109194
for i in range(5):
110195
if i == 0:
@@ -116,7 +201,9 @@ def triton_fn():
116201
torch.cuda.synchronize()
117202

118203
benchmark(triton_fn)
119-
print(f"[knn][triton] num_x={num_x} num_y={num_y} groups={groups}")
204+
print(
205+
f"[knn][triton] num_x={num_x} num_y={num_y} groups={groups}"
206+
)
120207

121208

122209
@pytest.mark.parametrize('num_x', [256, 1024, 4096, 255])
@@ -137,7 +224,9 @@ def cuda_fn():
137224
torch.cuda.synchronize()
138225

139226
benchmark(cuda_fn)
140-
print(f"[knn_graph][cuda] num_x={num_x} groups={groups} k={k}")
227+
print(
228+
f"[knn_graph][cuda] num_x={num_x} groups={groups} k={k}"
229+
)
141230

142231

143232
@pytest.mark.parametrize('num_x', [256, 1024, 4096, 255])
@@ -151,10 +240,22 @@ def test_triton_knn_graph_benchmark_triton(benchmark, num_x, num_groups):
151240
k = min(16, max(1, num_x - 1))
152241

153242
def cuda_fn():
154-
return knn_graph(x, k=k, batch=batch, loop=False, use_triton=False)
243+
return knn_graph(
244+
x,
245+
k=k,
246+
batch=batch,
247+
loop=False,
248+
use_triton=False,
249+
)
155250

156251
def triton_fn():
157-
return knn_graph(x, k=k, batch=batch, loop=False, use_triton=True)
252+
return knn_graph(
253+
x,
254+
k=k,
255+
batch=batch,
256+
loop=False,
257+
use_triton=True,
258+
)
158259

159260
for i in range(5):
160261
if i == 0:
@@ -166,4 +267,6 @@ def triton_fn():
166267
torch.cuda.synchronize()
167268

168269
benchmark(triton_fn)
169-
print(f"[knn_graph][triton] num_x={num_x} groups={groups} k={k}")
270+
print(
271+
f"[knn_graph][triton] num_x={num_x} groups={groups} k={k}"
272+
)

‎benchmarks/test_benchmark_nearest.py‎

Lines changed: 58 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -8,27 +8,55 @@
88
nearest = tc.nearest
99

1010
pytestmark = pytest.mark.skipif(
11-
not (torch.cuda.is_available() and importlib.util.find_spec('triton') is not None),
11+
not (
12+
torch.cuda.is_available()
13+
and importlib.util.find_spec('triton') is not None
14+
),
1215
reason='CUDA and Triton are required for Triton benchmark tests.',
1316
)
1417

15-
16-
def _make_batch(num_nodes: int, num_groups: int,
17-
device: torch.device) -> torch.Tensor:
18+
NEAREST_SIZES = [
19+
(256, 128),
20+
(1024, 512),
21+
(4096, 2048),
22+
(255, 127),
23+
(256, 5),
24+
(1024, 5),
25+
(4096, 5),
26+
(255, 5),
27+
]
28+
NEAREST_GROUPS = [1, 2, 4, 8, 16, 32]
29+
30+
31+
def _make_batch(
32+
num_nodes: int,
33+
num_groups: int,
34+
device: torch.device,
35+
) -> torch.Tensor:
1836
groups = max(1, min(num_groups, num_nodes))
19-
counts = torch.full((groups, ), num_nodes // groups, device=device,
20-
dtype=torch.long)
37+
counts = torch.full(
38+
(groups,),
39+
num_nodes // groups,
40+
device=device,
41+
dtype=torch.long,
42+
)
2143
remainder = num_nodes % groups
2244
if remainder:
2345
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])))
46+
return torch.repeat_interleave(
47+
torch.arange(groups, device=device),
48+
counts,
49+
)
50+
51+
52+
@pytest.mark.parametrize(
53+
'num_x,num_y,num_groups',
54+
(
55+
(*p[0], p[1])
56+
for p in product(NEAREST_SIZES, NEAREST_GROUPS)
57+
if p[1] <= min(p[0])
58+
),
59+
)
3260
@pytest.mark.benchmark(group="nearest")
3361
def test_triton_nearest_benchmark_cuda(benchmark, num_x, num_y, num_groups):
3462
torch.manual_seed(123)
@@ -46,13 +74,19 @@ def cuda_fn():
4674
torch.cuda.synchronize()
4775

4876
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])))
77+
print(
78+
f"[nearest][cuda] num_x={num_x} num_y={num_y} groups={groups}"
79+
)
80+
81+
82+
@pytest.mark.parametrize(
83+
'num_x,num_y,num_groups',
84+
(
85+
(*p[0], p[1])
86+
for p in product(NEAREST_SIZES, NEAREST_GROUPS)
87+
if p[1] <= min(p[0])
88+
),
89+
)
5690
@pytest.mark.benchmark(group="nearest")
5791
def test_triton_nearest_benchmark_triton(benchmark, num_x, num_y, num_groups):
5892
torch.manual_seed(123)
@@ -78,4 +112,6 @@ def triton_fn():
78112
torch.cuda.synchronize()
79113

80114
benchmark(triton_fn)
81-
print(f"[nearest][triton] num_x={num_x} num_y={num_y} groups={groups}")
115+
print(
116+
f"[nearest][triton] num_x={num_x} num_y={num_y} groups={groups}"
117+
)

‎setup.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,6 @@ def get_extensions():
109109

110110
install_requires = [
111111
'scipy',
112-
'triton==3.6.0',
113112
]
114113

115114
test_requires = [

0 commit comments

Comments
 (0)