| # Owner(s): ["module: dynamo"] |
| |
| import os |
| import unittest |
| from unittest.mock import patch |
| |
| import torch |
| import torch._dynamo |
| import torch._dynamo.test_case |
| import torch._functorch._aot_autograd |
| from torch._dynamo import config as dynamo_config |
| from torch._dynamo.utils import counters |
| from torch._functorch import config as functorch_config |
| from torch._functorch._aot_autograd.autograd_cache import ( |
| AOTAutogradCache, |
| autograd_cache_key, |
| BypassAOTAutogradCache, |
| ) |
| from torch._functorch._aot_autograd.schemas import AOTConfig |
| from torch._inductor import config as inductor_config |
| from torch._inductor.test_case import TestCase as InductorTestCase |
| from torch.testing._internal.common_cuda import SM80OrLater |
| from torch.testing._internal.common_device_type import largeTensorTest |
| from torch.testing._internal.common_utils import ( |
| instantiate_parametrized_tests, |
| parametrize, |
| skipIfWindows, |
| ) |
| from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU |
| |
| |
| @instantiate_parametrized_tests |
| class AOTAutogradCacheTests(InductorTestCase): |
| def setUp(self): |
| """ |
| Reset all counters and caches before each unit test |
| """ |
| super().setUp() |
| counters.clear() |
| self._clear_all_caches() |
| |
| def _clear_all_caches(self): |
| """ |
| Clear every cache, including AOTAutogradCache and FXCache |
| """ |
| torch._inductor.codecache.FxGraphCache.clear() |
| AOTAutogradCache.clear() |
| self._clear_dynamo_and_codecache() |
| |
| def _clear_dynamo_and_codecache(self): |
| """ |
| Clear unrelated caches, like dynamo and PyCodeCache |
| """ |
| torch._dynamo.reset() |
| for m in torch._inductor.codecache.PyCodeCache.cache.values(): |
| os.remove(m.__file__) |
| torch._inductor.codecache.PyCodeCache.cache_clear() |
| |
| @inductor_config.patch("fx_graph_remote_cache", False) |
| @inductor_config.patch("fx_graph_cache", True) |
| @functorch_config.patch({"enable_autograd_cache": True}) |
| def test_basic(self): |
| """ |
| Verify the interactions between FXGraphCache and AOTAutogradCache. |
| """ |
| |
| def fn(x, y): |
| return (x * 2, y @ y) |
| |
| a = torch.rand(25) |
| b = torch.rand(5, 5) |
| |
| compiled_fn = torch.compile(fn, backend="inductor") |
| |
| # A first call should miss in the cache. |
| self.assertEqual(fn(a, b), compiled_fn(a, b)) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1) |
| |
| # A second call should hit. (First reset so in-memory guards |
| # don't prevent compilation). |
| self._clear_dynamo_and_codecache() |
| self.assertEqual(fn(a, b), compiled_fn(a, b)) |
| |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1) |
| |
| @inductor_config.patch("fx_graph_remote_cache", False) |
| @inductor_config.patch("fx_graph_cache", True) |
| @functorch_config.patch({"enable_autograd_cache": True}) |
| @skipIfWindows( |
| msg="Known issue: Window can't delete loaded modules, so we can't clear module cache." |
| ) |
| def test_clear_fx_graph_cache(self): |
| """ |
| Verify the interactions between FXGraphCache and AOTAutogradCache. |
| """ |
| |
| def fn(x, y): |
| return (x * 2, y @ y) |
| |
| a = torch.rand(25) |
| b = torch.rand(5, 5) |
| |
| compiled_fn = torch.compile(fn, backend="inductor") |
| |
| # A first call should miss in the cache. |
| self.assertEqual(fn(a, b), compiled_fn(a, b)) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1) |
| |
| # Clear FX graph cache: second call should also be a miss |
| self._clear_dynamo_and_codecache() |
| torch._inductor.codecache.FxGraphCache.clear() |
| self.assertEqual(fn(a, b), compiled_fn(a, b)) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 2) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0) |
| # We save again into the cache |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 2) |
| |
| @inductor_config.patch("fx_graph_remote_cache", False) |
| @inductor_config.patch("fx_graph_cache", False) |
| @functorch_config.patch({"enable_autograd_cache": True}) |
| def test_fx_graph_cache_off(self): |
| """ |
| Should not use cache if FXGraphCache is not enabled |
| """ |
| |
| def fn(x, y): |
| return (x * 2, y @ y) |
| |
| a = torch.rand(25) |
| b = torch.rand(5, 5) |
| |
| compiled_fn = torch.compile(fn, backend="inductor") |
| |
| # A first call should miss in the cache. |
| self.assertEqual(fn(a, b), compiled_fn(a, b)) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_bypass"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 0) |
| |
| # Clear FX graph cache: second call should also be a miss |
| self._clear_dynamo_and_codecache() |
| |
| self.assertEqual(fn(a, b), compiled_fn(a, b)) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_bypass"], 2) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 0) |
| |
| @inductor_config.patch("fx_graph_remote_cache", False) |
| @inductor_config.patch("fx_graph_cache", True) |
| @functorch_config.patch({"enable_autograd_cache": True}) |
| @dynamo_config.patch("compiled_autograd", True) |
| def test_compiled_autograd_bypass(self): |
| def fn(a, b): |
| out = a.cos() + b |
| loss = out.sum() |
| ga, gb = torch.autograd.grad(loss, inputs=[a, b]) |
| |
| a = torch.randn(25, requires_grad=True) |
| b = torch.randn(25, requires_grad=True) |
| a2 = a.detach().clone().requires_grad_(True) |
| b2 = b.detach().clone().requires_grad_(True) |
| compiled_fn = torch.compile(fn, backend="inductor") |
| self.assertEqual(fn(a, b), compiled_fn(a2, b2)) |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_miss"], 1 |
| ) # from compiled forward |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_bypass"], 1 |
| ) # from compiled autograd |
| |
| @inductor_config.patch("fx_graph_remote_cache", False) |
| @inductor_config.patch("fx_graph_cache", True) |
| @functorch_config.patch({"enable_autograd_cache": True}) |
| @dynamo_config.patch("compiled_autograd", True) |
| def test_inference_graph_cache_hit_with_compiled_autograd_enabled(self): |
| def fn(a, b): |
| out = a.cos() + b |
| return out.sum() |
| |
| a = torch.randn(25) |
| b = torch.randn(25) |
| compiled_fn = torch.compile(fn, backend="inductor") |
| self.assertEqual(fn(a, b), compiled_fn(a, b)) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1) |
| |
| # Clear dynamo and run again. Should be a cache hit. |
| counters.clear() |
| self._clear_dynamo_and_codecache() |
| self.assertEqual(fn(a, b), compiled_fn(a, b)) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 0) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 0) |
| |
| @inductor_config.patch("fx_graph_remote_cache", False) |
| @inductor_config.patch({"fx_graph_cache": True}) |
| @functorch_config.patch({"enable_autograd_cache": True}) |
| def test_autograd_lazy_backward(self): |
| """ |
| Lazily compile the backward, and lazily save to cache |
| """ |
| |
| def fn(a, b): |
| return a.cos() + b |
| |
| a = torch.randn(25, requires_grad=True) |
| b = torch.randn(25, requires_grad=True) |
| a2 = a.detach().clone().requires_grad_(True) |
| b2 = b.detach().clone().requires_grad_(True) |
| compiled_fn = torch.compile(fn, backend="inductor") |
| self.assertEqual(fn(a, b), compiled_fn(a2, b2)) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 0) |
| |
| # Clear dynamo and run again. Should be a cache miss still, because backward hasn't run |
| self._clear_dynamo_and_codecache() |
| self.assertEqual(fn(a, b), compiled_fn(a2, b2)) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 2) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 0) |
| |
| # Now let's run the backward |
| fn(a, b).sum().backward() |
| compiled_fn(a2, b2).sum().backward() |
| self.assertEqual(a.grad, a2.grad) |
| self.assertEqual(b.grad, b2.grad) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1) |
| |
| # Clear dynamo and rerun everything, now there should be a cache hit |
| self._clear_dynamo_and_codecache() |
| a = torch.randn(25, requires_grad=True) |
| b = torch.randn(25, requires_grad=True) |
| a2 = a.detach().clone().requires_grad_(True) |
| b2 = b.detach().clone().requires_grad_(True) |
| self.assertEqual(fn(a, b), compiled_fn(a2, b2)) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 2) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1) |
| fn(a, b).sum().backward() |
| compiled_fn(a2, b2).sum().backward() |
| self.assertEqual(a.grad, a2.grad) |
| self.assertEqual(b.grad, b2.grad) |
| |
| @inductor_config.patch("fx_graph_remote_cache", False) |
| @inductor_config.patch("fx_graph_cache", True) |
| @functorch_config.patch({"enable_autograd_cache": True}) |
| def test_autograd_function(self): |
| """ |
| Tests autograd cache hits |
| """ |
| |
| def fn(a, b): |
| return a.sin() + b |
| |
| a = torch.randn(25, requires_grad=True) |
| b = torch.randn(25, requires_grad=True) |
| a2 = a.detach().clone().requires_grad_(True) |
| b2 = b.detach().clone().requires_grad_(True) |
| |
| compiled_fn = torch.compile(fn, backend="inductor") |
| |
| # A first call should miss in the cache. |
| self.assertEqual(fn(a, b), compiled_fn(a2, b2)) |
| fn(a, b).sum().backward() |
| compiled_fn(a2, b2).sum().backward() |
| self.assertEqual(a.grad, a2.grad) |
| self.assertEqual(b.grad, b2.grad) |
| |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1) |
| |
| # Reset all tensors |
| a = torch.randn(25, requires_grad=True) |
| b = torch.randn(25, requires_grad=True) |
| a2 = a.detach().clone().requires_grad_(True) |
| b2 = b.detach().clone().requires_grad_(True) |
| |
| # A second call should hit. (First reset so in-memory guards |
| # don't prevent compilation). |
| self._clear_dynamo_and_codecache() |
| self.assertEqual(fn(a, b), compiled_fn(a2, b2)) |
| fn(a, b).sum().backward() |
| compiled_fn(a2, b2).sum().backward() |
| self.assertEqual(a.grad, a2.grad) |
| self.assertEqual(b.grad, b2.grad) |
| |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1) |
| |
| @largeTensorTest("64GB", device=GPU_TYPE) |
| @parametrize("device", (GPU_TYPE,)) |
| @parametrize("dtype", (torch.float16, torch.bfloat16)) |
| @inductor_config.patch("fx_graph_cache", True) |
| @inductor_config.patch("fx_graph_remote_cache", False) |
| @functorch_config.patch({"enable_autograd_cache": True}) |
| def test_autograd_guard_single_entry(self, device, dtype): |
| """ |
| Test caching the same graph, but under conditions that introduce guards |
| for tensor sizes < int32. See test_codecache::TestFxGraphCache::test_cache_load_with_guards_int32_bounds. |
| |
| This test in particular tests the behavior of a single entry cache. If we ever make AOTAutogradCache |
| support multiple entries under the same key, this test should be updated. |
| """ |
| if device == GPU_TYPE and not HAS_GPU: |
| raise unittest.SkipTest(f"requires {GPU_TYPE}") |
| if device == "cuda" and dtype == torch.bfloat16 and not SM80OrLater: |
| raise unittest.SkipTest("requires CUDA SM80 or later") |
| |
| def fn(x, y): |
| return (x + x, y + y) |
| |
| def expect_miss(compiled_fn, a, b): |
| self._clear_dynamo_and_codecache() |
| counters.clear() |
| res = compiled_fn(a, b) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1) |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_guard_miss"], |
| 0, |
| ) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0) |
| return res |
| |
| def expect_hit(compiled_fn, a, b): |
| self._clear_dynamo_and_codecache() |
| counters.clear() |
| res = compiled_fn(a, b) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 0) |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_guard_miss"], |
| 0, |
| ) |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_hit"], |
| 1, |
| ) |
| return res |
| |
| def expect_guard_miss(compiled_fn, a, b): |
| self._clear_dynamo_and_codecache() |
| counters.clear() |
| res = compiled_fn(a, b) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1) |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_guard_miss"], |
| 1, |
| ) |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_hit"], |
| 0, |
| ) |
| return res |
| |
| compiled_fn = torch.compile(fn, dynamic=True) |
| |
| a_shape = (5, 6) |
| b_shape = (7, 8) |
| a = torch.rand(a_shape, device=device, dtype=dtype) |
| b = torch.rand(b_shape, device=device, dtype=dtype) |
| res1 = expect_miss(compiled_fn, a, b) |
| |
| # Same shape, should cache hit |
| a2 = a.detach().clone() |
| b2 = b.detach().clone() |
| |
| res2 = expect_hit(compiled_fn, a2, b2) |
| |
| self.assertEqual(res1, res2) |
| |
| # By changing the shape greatly, despite the same exact input |
| # graph, inductor should report a guard miss, leading |
| # to a cache miss on our end. |
| a_shape = (5, 6) |
| b_shape = (47000, 47001) |
| a3 = torch.rand(a_shape, device=device, dtype=dtype) |
| b3 = torch.rand(b_shape, device=device, dtype=dtype) |
| |
| expect_guard_miss(compiled_fn, a3, b3) |
| |
| # Wobble the shape a bit, but not enough |
| # to trigger a guard miss (since 6, 7 is still less than int32) |
| # Should result in a cache hit |
| a_shape = (6, 7) |
| b_shape = (47000, 47001) |
| a4 = torch.rand(a_shape, device=device, dtype=dtype) |
| b4 = torch.rand(b_shape, device=device, dtype=dtype) |
| expect_hit(compiled_fn, a4, b4) |
| |
| # Change the shape back to the original, |
| # FXGraphCache should hit because it stores |
| # multiple entries |
| a_shape = (5, 6) |
| b_shape = (7, 8) |
| a5 = torch.rand(a_shape, device=device, dtype=dtype) |
| b5 = torch.rand(b_shape, device=device, dtype=dtype) |
| expect_hit(compiled_fn, a5, b5) |
| |
| @largeTensorTest("64GB", device=GPU_TYPE) |
| @parametrize("device", (GPU_TYPE,)) |
| @parametrize("dtype", (torch.float16, torch.bfloat16)) |
| @parametrize("requires_grad", (True, False)) |
| @inductor_config.patch("fx_graph_cache", True) |
| @inductor_config.patch("fx_graph_remote_cache", False) |
| @functorch_config.patch({"enable_autograd_cache": True}) |
| def test_autograd_inductor_guards(self, device, dtype, requires_grad): |
| """ |
| Test caching the same graph, but under conditions that introduce guards |
| for tensor sizes < int32. |
| See test_codecache::TestFxGraphCache::test_cache_load_with_guards_int32_bounds. |
| """ |
| if device == GPU_TYPE and not HAS_GPU: |
| raise unittest.SkipTest(f"requires {GPU_TYPE}") |
| if device == "cuda" and dtype == torch.bfloat16 and not SM80OrLater: |
| raise unittest.SkipTest("requires CUDA SM80 or later") |
| |
| def fn(x, y): |
| return (x + x, y + y) |
| |
| compiled_fn = torch.compile(fn, dynamic=True) |
| |
| # Iterate over different shapes, varying whether the total |
| # size is below or above int32. For each combination, we expect |
| # different guards around whether the symbolic sizes do or do |
| # not exceed int32. |
| shapes = ( |
| ((5, 6), (7, 8)), |
| ((5, 6), (47000, 47001)), |
| ((47000, 47001), (5, 6)), |
| ) |
| expected_hits = expected_misses = expected_saves = 0 |
| expected_guard_misses = 0 |
| for a_shape, b_shape in shapes: |
| a = torch.rand( |
| a_shape, device=device, dtype=dtype, requires_grad=requires_grad |
| ) |
| b = torch.rand( |
| b_shape, device=device, dtype=dtype, requires_grad=requires_grad |
| ) |
| |
| # AVOID a dynamo reset here. We expect guards to have been |
| # added that will be violated with the new shape. We should |
| # see a recompilation (along with a cache miss). |
| res1 = compiled_fn(a, b) |
| # A first call should miss in the cache. |
| expected_misses += 1 |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_miss"], expected_misses |
| ) |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_guard_miss"], |
| expected_guard_misses, |
| ) |
| |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_hit"], expected_hits |
| ) |
| # Because dynamic shapes are enabled, we expect backwards to be compiled ahead of time |
| # So we should see a cache save here |
| expected_saves += 1 |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_saved"], expected_saves |
| ) |
| if requires_grad: |
| res1[0].sum().backward() |
| # No extra saves |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_saved"], expected_saves |
| ) |
| |
| a2 = a.detach().clone().requires_grad_(requires_grad) |
| b2 = b.detach().clone().requires_grad_(requires_grad) |
| # A second call should hit. (First reset so in-memory guards |
| # don't prevent compilation). |
| |
| # Now clear dynamo and we should see a cache hit |
| # This should populate guards to dynamo's cache, so that a subsequent run with a different |
| # shape will still trigger a second call to autograd_cache. |
| self._clear_dynamo_and_codecache() |
| res2 = compiled_fn(a2, b2) |
| expected_hits += 1 |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_miss"], expected_misses |
| ) |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_guard_miss"], |
| expected_guard_misses, |
| ) |
| # First compile is a regular cache miss, subsequent are guard misses |
| expected_guard_misses += 1 |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_hit"], expected_hits |
| ) |
| self.assertEqual( |
| counters["aot_autograd"]["autograd_cache_saved"], expected_saves |
| ) |
| self.assertEqual(res1, res2) |
| if requires_grad: |
| res2[0].sum().backward() |
| self.assertEqual(a.grad, a2.grad) |
| |
| @inductor_config.patch("fx_graph_cache", True) |
| @inductor_config.patch("fx_graph_remote_cache", False) |
| @functorch_config.patch({"enable_autograd_cache": True}) |
| def test_nn_module_with_params_global_constant(self): |
| class MyMod(torch.nn.Module): |
| CONSTANT = torch.tensor([[2, 2], [2, 2]]) |
| |
| def __init__(self) -> None: |
| super().__init__() |
| self.param = torch.nn.Parameter(torch.randn([2, 2])) |
| |
| def forward(self, x): |
| return x.sin() + self.param + MyMod.CONSTANT |
| |
| with torch.no_grad(): |
| compiled_fn = torch.compile(MyMod(), backend="inductor", fullgraph=True) |
| res1 = compiled_fn(torch.ones([2, 2])) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1) |
| |
| self._clear_dynamo_and_codecache() |
| res2 = compiled_fn(torch.ones([2, 2])) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1) |
| |
| self.assertEqual(res1, res2) |
| # Edit the "constant". We'll get a cache hit, |
| # but it should result in a different result when run |
| # because MyMod.CONSTANT is an input to the graph |
| MyMod.CONSTANT = torch.tensor([[3, 3], [3, 3]]) |
| self._clear_dynamo_and_codecache() |
| res3 = compiled_fn(torch.ones([2, 2])) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 2) |
| self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1) |
| self.assertNotEqual(res1, res3) |
| self.assertEqual(res1, res3.sub(torch.ones(2, 2))) |
| |
| |
| @inductor_config.patch("fx_graph_cache", True) |
| class AOTAutogradCachePicklerTests(torch._dynamo.test_case.TestCase): |
| @property |
| def device_type(self) -> str: |
| return "cuda" if torch.cuda.is_available() else "cpu" |
| |
| def default_config(self): |
| return AOTConfig( |
| fw_compiler=None, |
| bw_compiler=None, |
| inference_compiler=None, |
| partition_fn=None, |
| decompositions={}, |
| num_params_buffers=0, |
| aot_id=0, |
| keep_inference_input_mutations=False, |
| dynamic_shapes=True, |
| aot_autograd_arg_pos_to_source=None, |
| is_export=False, |
| no_tangents=False, |
| enable_log=False, |
| ) |
| |
| def _get_dynamo_output(self, fn, *args, **kwargs): |
| # Reset dynamo between runs |
| torch._dynamo.reset() |
| fx_graph = None |
| example_inputs = None |
| |
| def compiler(gm, inputs, **kwargs): |
| nonlocal fx_graph |
| nonlocal example_inputs |
| fx_graph = gm |
| example_inputs = inputs |
| return gm |
| |
| g = torch.compile(fn, backend=compiler, fullgraph=True) |
| result = g(*args, **kwargs) |
| return (result, fx_graph, example_inputs) |
| |
| def gen_cache_key(self, f, config, inputs=None): |
| if inputs is None: |
| inputs = [torch.ones(3)] |
| _, fx_g, example_inputs = self._get_dynamo_output(f, *inputs) |
| return autograd_cache_key(fx_g, example_inputs, config, {}) |
| |
| def test_basic_hash_key(self): |
| def fn(x): |
| return x.sin().cos() |
| |
| config = self.default_config() |
| # Check hash is stable on multiple runs |
| c1 = self.gen_cache_key(fn, config) |
| c2 = self.gen_cache_key(fn, config) |
| self.assertEqual(c1, c2) |
| |
| def test_identical_graphs_and_configs(self): |
| def fn(x): |
| return x.sin().cos() |
| |
| def fn2(x): |
| y = x.sin() |
| z = y.cos() |
| return z |
| |
| # Make the id different, but otherwise identical |
| config = self.default_config() |
| config2 = self.default_config() |
| config2.aot_id = 1 |
| |
| c1 = self.gen_cache_key(fn, config) |
| c2 = self.gen_cache_key(fn, config2) |
| self.assertEqual(c1, c2) |
| |
| def test_different_graphs(self): |
| def fn(x): |
| return x.cos().sin() |
| |
| def fn2(x): |
| return x.sin().cos() |
| |
| config = self.default_config() |
| c1 = self.gen_cache_key(fn, config) |
| c2 = self.gen_cache_key(fn2, config) |
| self.assertNotEqual(c1, c2) |
| |
| def test_different_configs(self): |
| def fn(x): |
| return x.cos().sin() |
| |
| config = self.default_config() |
| config2 = self.default_config() |
| config2.dynamic_shapes = False |
| c1 = self.gen_cache_key(fn, config) |
| c2 = self.gen_cache_key(fn, config2) |
| self.assertNotEqual(c1, c2) |
| |
| def test_different_inputs(self): |
| def fn(x): |
| return x.cos().sin() |
| |
| config = self.default_config() |
| c1 = self.gen_cache_key(fn, config, inputs=[torch.ones(3)]) |
| c2 = self.gen_cache_key(fn, config, inputs=[torch.ones(2)]) |
| self.assertNotEqual(c1, c2) |
| |
| def test_different_global_configs(self): |
| def fn(x): |
| return x.cos().sin() |
| |
| config = self.default_config() |
| |
| c1 = self.gen_cache_key(fn, config) |
| c2 = self.gen_cache_key(fn, config) |
| self.assertEqual(c1, c2) |
| |
| c1 = self.gen_cache_key(fn, config) |
| |
| # Change functorch config |
| with functorch_config.patch( |
| {"debug_assert": not functorch_config.debug_assert} |
| ): |
| c2 = self.gen_cache_key(fn, config) |
| |
| self.assertNotEqual(c1, c2) |
| |
| c1 = self.gen_cache_key(fn, config) |
| # Change inductor config |
| with inductor_config.patch({"debug": not inductor_config.debug}): |
| c2 = self.gen_cache_key(fn, config) |
| |
| self.assertNotEqual(c1, c2) |
| |
| c1 = self.gen_cache_key(fn, config) |
| # Change torch grad enabled |
| with torch.no_grad(): |
| c2 = self.gen_cache_key(fn, config) |
| self.assertNotEqual(c1, c2) |
| |
| def test_incompatible_function(self): |
| @torch._dynamo.allow_in_graph |
| class AllowInGraphFunc(torch.autograd.Function): |
| @staticmethod |
| def forward(_, x): |
| torch._dynamo.graph_break() |
| return x.sin() |
| |
| def fn(x): |
| return AllowInGraphFunc.apply(x) |
| |
| config = self.default_config() |
| self.assertRaises( |
| BypassAOTAutogradCache, lambda: self.gen_cache_key(fn, config) |
| ) |
| |
| def test_private_namespace(self): |
| # TODO: anyone who monkeypatches a **public** function into torch namespace with @allow_in_graph |
| # could still break our sanity check and cache something bad. But that's an edge case we'll take the risk on. |
| # Monkeypatch some random private function into torch, see that it fails |
| @torch._dynamo.allow_in_graph |
| def my_private_fun(x): |
| return x.sin() |
| |
| with patch("torch._my_priv", new=my_private_fun, create=True): |
| |
| def fn(x): |
| return torch._my_priv(x) |
| |
| config = self.default_config() |
| self.assertRaises( |
| BypassAOTAutogradCache, lambda: self.gen_cache_key(fn, config) |
| ) |
| |
| def test_private_builtin(self): |
| # _foreach_add is a private torch function, but |
| # it's also a builtin_function_or_method, so it should be allowed to be cached |
| # since dynamo allows it in the graph |
| def fn(x, b): |
| y = (x, x) |
| return torch._foreach_add(y, b) |
| |
| config = self.default_config() |
| r1 = self.gen_cache_key(fn, config, inputs=[torch.ones(3), 1]) |
| r2 = self.gen_cache_key(fn, config, inputs=[torch.ones(3), 2]) |
| self.assertNotEqual(r1, r2) |
| |
| def test_nn_module_with_params(self): |
| class MyMod(torch.nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.seq = torch.nn.Parameter(torch.ones((3, 3))) |
| |
| def forward(self, x): |
| return self.seq + x |
| |
| config = self.default_config() |
| # Different inputs and parameters, but all the same size |
| c1 = self.gen_cache_key(MyMod(), config, inputs=[torch.ones((3, 3))]) |
| c2 = self.gen_cache_key(MyMod(), config, inputs=[torch.ones((3, 3))]) |
| self.assertEqual(c1, c2) |
| |
| def test_normal_torch_function(self): |
| @torch._dynamo.allow_in_graph |
| def fn(x): |
| y = torch.sin(x) |
| z = torch.cos(x) |
| w = y + z |
| w.abs() |
| return w |
| |
| config = self.default_config() |
| self.gen_cache_key(fn, config) |
| |
| |
| if __name__ == "__main__": |
| from torch._dynamo.test_case import run_tests |
| |
| run_tests() |
