1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
|
from benchmark_helper import time_with_torch_timer
import torch
import torch._dynamo
import torch._dynamo.config
import torch._inductor.config as config
@torch._dynamo.optimize("inductor", nopython=True)
def inductor_aten_bmm(a, b):
return torch.bmm(a, b)
@torch._dynamo.optimize("inductor", nopython=True)
def inductor_triton_bmm(a, b):
return torch.bmm(a, b)
def torch_bmm(a, b):
return torch.bmm(a, b)
def test_total_time(shapes):
print("shape; torch bmm; inductor aten bmm; inductor triton bmm")
for i in range(len(shapes)):
a_shape, b_shape = shapes[i]
print(a_shape, "x", b_shape, end="; ")
a = torch.randn(a_shape, device="cuda", dtype=torch.float16)
b = torch.randn(b_shape, device="cuda", dtype=a.dtype)
config.triton.use_bmm = False
inductor_aten_bmm(a, b)
config.triton.use_bmm = True
inductor_triton_bmm(a, b)
torch_ms = time_with_torch_timer(torch_bmm, (a, b)).mean * 1000
config.triton.use_bmm = False
ind_aten_ms = time_with_torch_timer(inductor_aten_bmm, (a, b)).mean * 1000
config.triton.use_bmm = True
ind_triton_ms = time_with_torch_timer(inductor_triton_bmm, (a, b)).mean * 1000
print(torch_ms, ind_aten_ms, ind_triton_ms, sep="; ")
if __name__ == "__main__":
shapes = [
# BERT (all)
([192, 128, 64], [192, 64, 128]),
([192, 128, 128], [192, 128, 64]),
# hf_GPT2 (all)
([12, 1024, 1024], [12, 1024, 64]),
([12, 1024, 64], [12, 64, 1024]),
# hf_Albert (all)
([12, 512, 64], [12, 64, 512]),
([12, 512, 512], [12, 512, 64]),
]
test_total_time(shapes)
|
