frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add infra to run CPython tests under Dynamo (#150787)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/150787
Approved by: https://github.com/zou3519
This commit is contained in:
Guilherme Leobas
2025-05-05 16:51:54 +00:00
committed by PyTorch MergeBot
parent 50fe1b2349
commit 7c96dd8f0c
13 changed files with 175 additions and 1608 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
torch/testing/_internal/common_utils.py
View File

@ -3157,6 +3157,13 @@ class TestCase(expecttest.TestCase):
def wrap_with_cuda_memory_check(self, method):
return self.wrap_method_with_policy(method, self.assertLeaksNoCudaTensors)
def _dynamo_test_key(self):
return f"{self.__class__.__name__}.{self._testMethodName}"
def compile_fn(self, fn, backend, nopython):
# Allows subclasses to control compilation
return torch._dynamo.optimize(backend, nopython=nopython)(fn)
def _run_custom(self, result=None):
using_unittest = isinstance(result, unittest.TestResult)
@ -3232,16 +3239,16 @@ class TestCase(expecttest.TestCase):
with unittest.mock.patch("torch._dynamo.config.suppress_errors", suppress_errors), maybe_disable_size_asserts:
if TEST_WITH_AOT_EAGER:
super_run = torch._dynamo.optimize("aot_eager_decomp_partition")(super_run)
super_run = self.compile_fn(super_run, "aot_eager_decomp_partition", nopython)
elif TEST_WITH_TORCHDYNAMO or TEST_WITH_TORCHINDUCTOR:
if TEST_WITH_TORCHINDUCTOR:
super_run = torch._dynamo.optimize("inductor")(super_run)
super_run = self.compile_fn(super_run, "inductor", nopython)
else:
# Assume eager-generated GraphModules will not error out.
# If we do, this is probably a Dynamo bug!
super_run = torch._dynamo.optimize("eager_noexcept", nopython=nopython)(super_run)
super_run = self.compile_fn(super_run, "eager_noexcept", nopython)
key = f"{self.__class__.__name__}.{self._testMethodName}"
key = self._dynamo_test_key()
def expect_failure(f, file_name):
@wraps(f)