mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
7763c83af6
The `usort` config in `pyproject.toml` has no effect due to a typo. Fixing the typo make `usort` do more and generate the changes in the PR. Except `pyproject.toml`, all changes are generated by `lintrunner -a --take UFMT --all-files`. Pull Request resolved: https://github.com/pytorch/pytorch/pull/127126 Approved by: https://github.com/kit1980 ghstack dependencies: #127122, #127123, #127124, #127125
63 lines
1.7 KiB
Python
63 lines
1.7 KiB
Python
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)
|