mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
134bc4fc34
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/129763 Approved by: https://github.com/jansel
36 lines
964 B
Python
36 lines
964 B
Python
# Owner(s): ["module: inductor"]
|
|
|
|
import functools
|
|
import logging
|
|
|
|
import torch
|
|
from torch._inductor.runtime.runtime_utils import do_bench
|
|
from torch._inductor.test_case import run_tests, TestCase
|
|
from torch._inductor.utils import do_bench_using_profiling
|
|
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class TestBench(TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
super().setUpClass()
|
|
x = torch.rand(1024, 10).cuda().half()
|
|
w = torch.rand(512, 10).cuda().half()
|
|
cls._bench_fn = functools.partial(torch.nn.functional.linear, x, w)
|
|
|
|
def test_do_bench(self):
|
|
res = do_bench(self._bench_fn)
|
|
log.warning("do_bench result: %s", res)
|
|
self.assertGreater(res, 0)
|
|
|
|
def test_do_bench_using_profiling(self):
|
|
res = do_bench_using_profiling(self._bench_fn)
|
|
log.warning("do_bench_using_profiling result: %s", res)
|
|
self.assertGreater(res, 0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_tests("cuda")
|