mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
c0ed38e644
See https://github.com/pytorch/pytorch/pull/129751#issue-2380881501. Most changes are auto-generated by linter. You can review these PRs via: ```bash git diff --ignore-all-space --ignore-blank-lines HEAD~1 ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/129754 Approved by: https://github.com/ezyang
62 lines
1.7 KiB
Python
62 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)
|