已合并
[UT] add and adapt official source testcases for dynamo module #8742
AtomGit-Bot创建于 2024年1月2日
[UT] add and adapt official source testcases for dynamo module #8742
已合并
AtomGit-Bot创建于 2024年1月2日
master-dev合入到master
8 个文件变更+980-0
Atest/dynamo/test_after_aot.py+83-0
@@ -0,0 +1,83 @@
1+# Owner(s): ["module: dynamo"]
2+ 
3+import io
4+import os
5+import shutil
6+import sys
7+import tempfile
8+import unittest
9+ 
10+import torch
11+import torch_npu
12+import torch._dynamo.test_case
13+from torch._dynamo.repro.after_aot import InputReader, InputWriter, save_graph_repro
14+from torch.fx.experimental.proxy_tensor import make_fx
15+from torch.testing._internal.common_utils import IS_FBCODE
16+from torch.utils._traceback import report_compile_source_on_error
17+ 
18+ 
19+def strip_trailing_whitespace(r):
20+ return "\n".join([l_str.rstrip() for l_str in r.split("\n")])
21+ 
22+ 
23+class TestAfterAot(torch._dynamo.test_case.TestCase):
24+ @unittest.skipIf(IS_FBCODE, "NotImplementedError")
25+ def test_save_graph_repro(self):
26+ # do for later: This triggers CUDA context initialization, even though
27+ # it is CPU only
28+ buf = io.StringIO()
29+ args = [torch.randn(4)]
30+ 
31+ def f(x):
32+ return (x * x,)
33+ 
34+ gm = make_fx(f)(*args)
35+ with tempfile.TemporaryDirectory() as d:
36+ save_graph_repro(buf, gm, args, "inductor_accuracy", save_dir=d)
37+ r = buf.getvalue()
38+ with report_compile_source_on_error():
39+ exec(r, {"__compile_source__": r})
40+ 
41+ shutil.rmtree(os.path.join(d, "storages"))
42+ 
43+ # Should still work even without the save dir
44+ with report_compile_source_on_error():
45+ exec(r, {"__compile_source__": r})
46+ 
47+ @unittest.skipIf(sys.byteorder != "little", "checksum depends on endianness")
48+ def test_dump_tensor(self):
49+ def test(tensor, expected):
50+ with tempfile.TemporaryDirectory() as d:
51+ writer = InputWriter(d, stable_hash=True)
52+ writer.tensor("x", tensor)
53+ self.assertExpectedInline("\n".join(writer._lines), expected, skip=1)
54+ reader = InputReader(d)
55+ env = {"reader": reader, "torch": torch}
56+ # do for later: assert no logs
57+ exec("\n".join(writer._lines), env)
58+ self.assertEqual(reader.args[0], tensor)
59+ 
60+ test(
61+ torch.zeros(3, 4),
62+ """\
63+buf0 = reader.storage('c17fd92682ca5b304ac71074b558dda9e8eb4d66', 48)
64+reader.tensor(buf0, (3, 4), is_leaf=True) # x""",
65+ )
66+ test(
67+ torch.ones(3, 4, dtype=torch.int32),
68+ """\
69+buf0 = reader.storage('7c221e2da0c58c700cc2996644dd13d042bd552e', 48, dtype_hint=torch.int32)
70+reader.tensor(buf0, (3, 4), dtype=torch.int32, is_leaf=True) # x""",
71+ )
72+ test(
73+ torch.empty((3, 4, 5, 6), memory_format=torch.channels_last).fill_(2),
74+ """\
75+buf0 = reader.storage('49ebab3961d6221e64c4c72b0aefd976bdd2afc4', 1440)
76+reader.tensor(buf0, (3, 4, 5, 6), (120, 1, 24, 4), is_leaf=True) # x""",
77+ )
78+ 
79+ 
80+if __name__ == "__main__":
81+ from torch._dynamo.test_case import run_tests
82+ 
83+ run_tests()
Atest/dynamo/test_base_output.py+97-0
@@ -0,0 +1,97 @@
1+# Owner(s): ["module: dynamo"]
2+import unittest.mock
3+ 
4+import torch
5+import torch_npu
6+import torch._dynamo.test_case
7+import torch._dynamo.testing
8+from torch._dynamo.testing import same
9+ 
10+try:
11+ from diffusers.models import unet_2d
12+except ImportError:
13+ unet_2d = None
14+ 
15+ 
16+def maybe_skip(fn):
17+ if unet_2d is None:
18+ return unittest.skip("requires diffusers")(fn)
19+ return fn
20+ 
21+ 
22+class TestBaseOutput(torch._dynamo.test_case.TestCase):
23+ @maybe_skip
24+ def test_create(self):
25+ def fn(a):
26+ tmp = unet_2d.UNet2DOutput(a + 1)
27+ return tmp
28+ 
29+ torch._dynamo.testing.standard_test(self, fn=fn, nargs=1, expected_ops=1)
30+ 
31+ @maybe_skip
32+ def test_assign(self):
33+ def fn(a):
34+ tmp = unet_2d.UNet2DOutput(a + 1)
35+ tmp.sample = a + 2
36+ return tmp
37+ 
38+ args = [torch.randn(10)]
39+ obj1 = fn(*args)
40+ 
41+ cnts = torch._dynamo.testing.CompileCounter()
42+ opt_fn = torch._dynamo.optimize_assert(cnts)(fn)
43+ obj2 = opt_fn(*args)
44+ self.assertTrue(same(obj1.sample, obj2.sample))
45+ self.assertEqual(cnts.frame_count, 1)
46+ self.assertEqual(cnts.op_count, 2)
47+ 
48+ def _common(self, fn, op_count):
49+ args = [
50+ unet_2d.UNet2DOutput(
51+ sample=torch.randn(10),
52+ )
53+ ]
54+ obj1 = fn(*args)
55+ cnts = torch._dynamo.testing.CompileCounter()
56+ opt_fn = torch._dynamo.optimize_assert(cnts)(fn)
57+ obj2 = opt_fn(*args)
58+ self.assertTrue(same(obj1, obj2))
59+ self.assertEqual(cnts.frame_count, 1)
60+ self.assertEqual(cnts.op_count, op_count)
61+ 
62+ @maybe_skip
63+ def test_getattr(self):
64+ def fn(obj: unet_2d.UNet2DOutput):
65+ x = obj.sample * 10
66+ return x
67+ 
68+ self._common(fn, 1)
69+ 
70+ @maybe_skip
71+ def test_getitem(self):
72+ def fn(obj: unet_2d.UNet2DOutput):
73+ x = obj["sample"] * 10
74+ return x
75+ 
76+ self._common(fn, 1)
77+ 
78+ @maybe_skip
79+ def test_tuple(self):
80+ def fn(obj: unet_2d.UNet2DOutput):
81+ a = obj.to_tuple()
82+ return a[0] * 10
83+ 
84+ self._common(fn, 1)
85+ 
86+ @maybe_skip
87+ def test_index(self):
88+ def fn(obj: unet_2d.UNet2DOutput):
89+ return obj[0] * 10
90+ 
91+ self._common(fn, 1)
92+ 
93+ 
94+if __name__ == "__main__":
95+ from torch._dynamo.test_case import run_tests
96+ 
97+ run_tests()
Atest/dynamo/test_compile.py+107-0
@@ -0,0 +1,107 @@
1+# Owner(s): ["module: dynamo"]
2+ 
3+import inspect
4+import os
5+import tempfile
6+import unittest
7+ 
8+import torch
9+import torch_npu
10+from torch._dynamo.testing import CompileCounter
11+ 
12+ 
13+class ToyModel(torch.nn.Module):
14+ def __init__(self):
15+ super().__init__()
16+ self.linear = torch.nn.Linear(10, 10)
17+ self.relu = torch.nn.ReLU()
18+ 
19+ def forward(self, x):
20+ return self.relu(self.linear(x))
21+ 
22+ 
23+class InPlaceCompilationTests(unittest.TestCase):
24+ def test_compilation(self):
25+ torch._dynamo.reset()
26+ model = ToyModel()
27+ cnt = CompileCounter()
28+ model.compile(backend=cnt)
29+ x = torch.randn(10, 10)
30+ model(x)
31+ self.assertEqual(cnt.frame_count, 1)
32+ 
33+ def test_overwrite_call_impl(self):
34+ torch._dynamo.reset()
35+ model = ToyModel()
36+ self.assertIsNone(model._compiled_call_impl)
37+ model.compile()
38+ self.assertIsNotNone(model._compiled_call_impl)
39+ 
40+ def test_save(self):
41+ torch._dynamo.reset()
42+ model = ToyModel()
43+ model.compile()
44+ model(torch.randn(1, 10))
45+ 
46+ with tempfile.TemporaryDirectory() as tmpdirname:
47+ torch.save(model, os.path.join(tmpdirname, "model.pt"))
48+ loaded_model = torch.load(os.path.join(tmpdirname, "model.pt"))
49+ loaded_model(torch.randn(1, 10))
50+ 
51+ def test_state_dict_save(self):
52+ torch._dynamo.reset()
53+ model = ToyModel()
54+ model.compile()
55+ model(torch.randn(1, 10))
56+ with tempfile.TemporaryDirectory() as tmpdirname:
57+ torch.save(model.state_dict(), os.path.join(tmpdirname, "model.pt"))
58+ loaded_model = ToyModel()
59+ loaded_model.load_state_dict(
60+ torch.load(os.path.join(tmpdirname, "model.pt"))
61+ )
62+ loaded_model(torch.randn(1, 10))
63+ 
64+ def test_jit_save(self):
65+ torch._dynamo.reset()
66+ model = ToyModel()
67+ model.compile()
68+ model(torch.randn(1, 10))
69+ scripted_model = torch.jit.script(model)
70+ with tempfile.TemporaryDirectory() as tmpdirname:
71+ torch.jit.save(scripted_model, os.path.join(tmpdirname, "model.pt"))
72+ loaded_model = torch.jit.load(os.path.join(tmpdirname, "model.pt"))
73+ loaded_model(torch.randn(1, 10))
74+ 
75+ 
76+# The private variants of the below functions are extensively tested
77+# So as long as the signatures match we're good
78+class PublicTorchCompilerTests(unittest.TestCase):
79+ def check_signature(self, public_fn_name, private_fn_name, private_namespace):
80+ public_fn = getattr(torch.compiler, public_fn_name)
81+ private_fn = getattr(private_namespace, private_fn_name)
82+ 
83+ public_sig = inspect.signature(public_fn)
84+ private_sig = inspect.signature(private_fn)
85+ 
86+ self.assertEqual(
87+ public_sig,
88+ private_sig,
89+ f"Signatures do not match for function {public_fn_name}() \n Public: {public_sig} \n Private: {private_sig}",
90+ )
91+ 
92+ def test_dynamo_signatures(self):
93+ function_names = [
94+ "reset",
95+ "allow_in_graph",
96+ "list_backends",
97+ "assume_constant_result",
98+ "disable",
99+ ]
100+ 
101+ for fn_name in function_names:
102+ self.check_signature(fn_name, fn_name, torch._dynamo)
103+ 
104+ 
105+if __name__ == '__main__':
106+ unittest.main()
107+
Atest/dynamo/test_config.py+350-0
@@ -0,0 +1,350 @@
1+# Owner(s): ["module: dynamo"]
2+ 
3+import torch
4+import torch_npu
5+import torch._dynamo.test_case
6+import torch._dynamo.testing
7+from torch._dynamo.utils import disable_cache_limit
8+ 
9+# NB: do NOT include this test class in test_dynamic_shapes.py
10+ 
11+ 
12+class ConfigTests(torch._dynamo.test_case.TestCase):
13+ @disable_cache_limit()
14+ def test_no_automatic_dynamic(self):
15+ def fn(a, b):
16+ return a - b * 10
17+ 
18+ torch._dynamo.reset()
19+ cnt_static = torch._dynamo.testing.CompileCounter()
20+ with torch._dynamo.config.patch(
21+ automatic_dynamic_shapes=False, assume_static_by_default=True
22+ ):
23+ opt_fn = torch._dynamo.optimize(cnt_static)(fn)
24+ for i in range(2, 12):
25+ opt_fn(torch.randn(i), torch.randn(i))
26+ self.assertEqual(cnt_static.frame_count, 10)
27+ 
28+ @disable_cache_limit()
29+ def test_automatic_dynamic(self):
30+ def fn(a, b):
31+ return a - b * 10
32+ 
33+ torch._dynamo.reset()
34+ cnt_dynamic = torch._dynamo.testing.CompileCounter()
35+ with torch._dynamo.config.patch(
36+ automatic_dynamic_shapes=True, assume_static_by_default=True
37+ ):
38+ opt_fn = torch._dynamo.optimize(cnt_dynamic)(fn)
39+ # NB: must not do 0, 1 as they specialized
40+ for i in range(2, 12):
41+ opt_fn(torch.randn(i), torch.randn(i))
42+ # two graphs now rather than 10
43+ self.assertEqual(cnt_dynamic.frame_count, 2)
44+ 
45+ @disable_cache_limit()
46+ def test_no_assume_static_by_default(self):
47+ def fn(a, b):
48+ return a - b * 10
49+ 
50+ torch._dynamo.reset()
51+ cnt_dynamic = torch._dynamo.testing.CompileCounter()
52+ with torch._dynamo.config.patch(
53+ automatic_dynamic_shapes=True, assume_static_by_default=False
54+ ):
55+ opt_fn = torch._dynamo.optimize(cnt_dynamic)(fn)
56+ # NB: must not do 0, 1 as they specialized
57+ for i in range(2, 12):
58+ opt_fn(torch.randn(i), torch.randn(i))
59+ # one graph now, as we didn't wait for recompile
60+ self.assertEqual(cnt_dynamic.frame_count, 1)
61+ 
62+ def test_config_compile_ignored(self):
63+ # Remove from this list if no longer relevant
64+ dynamo_guarded_config_ignorelist = {
65+ "log_file_name",
66+ "verbose",
67+ "verify_correctness", # will not affect model, will raise RuntimeError
68+ # (no silent change to compilation behaviour)
69+ "cache_size_limit",
70+ "accumulated_cache_size_limit",
71+ "replay_record_enabled",
72+ "cprofile", # only wraps _compile, not graph
73+ "repro_after",
74+ "repro_level",
75+ "repro_forward_only",
76+ "repro_tolerance",
77+ "same_two_models_use_fp64",
78+ "error_on_recompile", # safe because: will throw error
79+ "report_guard_failures",
80+ "base_dir", # used for minifying / logging
81+ "DEBUG_DIR_VAR_NAME",
82+ "debug_dir_root",
83+ }
84+ for k in dynamo_guarded_config_ignorelist:
85+ assert k in torch._dynamo.config._compile_ignored_keys
86+ 
87+ def test_config_hash(self):
88+ config = torch._dynamo.config
89+ starting_hash = config.get_hash()
90+ 
91+ with config.patch({"verbose": not config.verbose}):
92+ new_hash = config.get_hash()
93+ assert "verbose" in config._compile_ignored_keys
94+ assert new_hash == starting_hash
95+ 
96+ new_hash = config.get_hash()
97+ assert new_hash == starting_hash
98+ 
99+ with config.patch({"dead_code_elimination": not config.dead_code_elimination}):
100+ changed_hash = config.get_hash()
101+ assert "dead_code_elimination" not in config._compile_ignored_keys
102+ assert changed_hash != starting_hash
103+ 
104+ # Test nested patch
105+ with config.patch({"verbose": not config.verbose}):
106+ inner_changed_hash = config.get_hash()
107+ assert inner_changed_hash == changed_hash
108+ assert inner_changed_hash != starting_hash
109+ 
110+ newest_hash = config.get_hash()
111+ assert changed_hash != newest_hash
112+ assert newest_hash == starting_hash
113+ 
114+ @disable_cache_limit()
115+ def test_no_saved_config(self):
116+ def fn(a, b):
117+ return a - b * 10
118+ 
119+ torch._dynamo.reset()
120+ cnt_dynamic = torch._dynamo.testing.CompileCounter()
121+ with torch._dynamo.config.patch(
122+ automatic_dynamic_shapes=False, assume_static_by_default=True
123+ ):
124+ opt_fn_static_shape = torch._dynamo.optimize(
125+ cnt_dynamic, save_config=False
126+ )(fn)
127+ opt_fn_static_shape(torch.randn(2), torch.randn(2))
128+ opt_fn_static_shape(torch.randn(3), torch.randn(3))
129+ 
130+ self.assertEqual(cnt_dynamic.frame_count, 2)
131+ 
132+ with torch._dynamo.config.patch(
133+ automatic_dynamic_shapes=True, assume_static_by_default=False
134+ ):
135+ for i in range(2, 12):
136+ opt_fn_static_shape(
137+ torch.randn(i), torch.randn(i)
138+ ) # will be recompiled under new config
139+ 
140+ self.assertEqual(cnt_dynamic.frame_count, 3)
141+ 
142+ @disable_cache_limit()
143+ def test_no_saved_config_nested(self):
144+ def fn(a, b):
145+ return a - b * 10
146+ 
147+ torch._dynamo.reset()
148+ cnt_dynamic = torch._dynamo.testing.CompileCounter()
149+ cnt_dynamic_1 = torch._dynamo.testing.CompileCounter()
150+ with torch._dynamo.config.patch(
151+ automatic_dynamic_shapes=True, assume_static_by_default=False
152+ ):
153+ opt_fn_static_shape = torch._dynamo.optimize(cnt_dynamic, dynamic=False)(fn)
154+ 
155+ # Will trigger recompile as compiled as static
156+ opt_fn_static_shape(torch.randn(2), torch.randn(2))
157+ opt_fn_static_shape(torch.randn(3), torch.randn(3))
158+ 
159+ self.assertEqual(cnt_dynamic.frame_count, 2)
160+ 
161+ opt_fn_try_dynamic = torch._dynamo.optimize(
162+ cnt_dynamic_1, save_config=False
163+ )(opt_fn_static_shape)
164+ 
165+ for i in range(2, 6):
166+ opt_fn_try_dynamic(torch.randn(i), torch.randn(i))
167+ self.assertEqual(cnt_dynamic_1.frame_count, 1)
168+ 
169+ # Saved config = False will use whatever config is available
170+ with torch._dynamo.config.patch(
171+ automatic_dynamic_shapes=False, assume_static_by_default=True
172+ ):
173+ for i in range(6, 12):
174+ opt_fn_try_dynamic(torch.randn(i), torch.randn(i))
175+ self.assertEqual(cnt_dynamic_1.frame_count, 7)
176+ 
177+ @disable_cache_limit()
178+ def test_config_changed_from_guarded_config_1(self):
179+ def fn(a, b):
180+ return a - b * 10
181+ 
182+ torch._dynamo.reset()
183+ 
184+ cnt_dynamic = torch._dynamo.testing.CompileCounter()
185+ with torch._dynamo.config.patch(
186+ automatic_dynamic_shapes=False, assume_static_by_default=True
187+ ):
188+ opt_fn_static_shape = torch._dynamo.optimize(cnt_dynamic)(fn)
189+ res = opt_fn_static_shape(torch.randn(2), torch.randn(2))
190+ opt_fn_static_shape(torch.randn(3), torch.randn(3))
191+ 
192+ self.assertEqual(cnt_dynamic.frame_count, 2)
193+ 
194+ with torch._dynamo.config.patch(
195+ automatic_dynamic_shapes=True, assume_static_by_default=False
196+ ):
197+ for i in range(2, 12):
198+ # Only 4-11 will now be recompiled under old config
199+ # 2-3 have been already been compiled under old config
200+ # and hence will hit cache
201+ opt_fn_static_shape(torch.randn(i), torch.randn(i))
202+ 
203+ self.assertEqual(cnt_dynamic.frame_count, 10)
204+ 
205+ @disable_cache_limit()
206+ def test_config_changed_from_guarded_config_2(self):
207+ def fn(a, b):
208+ return a - b * 10
209+ 
210+ torch._dynamo.reset()
211+ 
212+ cnt_dynamic = torch._dynamo.testing.CompileCounter()
213+ with torch._dynamo.config.patch(
214+ automatic_dynamic_shapes=True, assume_static_by_default=False
215+ ):
216+ opt_fn_dynamic_shape = torch._dynamo.optimize(cnt_dynamic)(fn)
217+ opt_fn_dynamic_shape(torch.randn(2), torch.randn(2))
218+ opt_fn_dynamic_shape(torch.randn(3), torch.randn(3))
219+ 
220+ self.assertEqual(cnt_dynamic.frame_count, 1)
221+ 
222+ with torch._dynamo.config.patch(
223+ automatic_dynamic_shapes=False, assume_static_by_default=True
224+ ):
225+ for i in range(2, 12):
226+ opt_fn_dynamic_shape(
227+ torch.randn(i), torch.randn(i)
228+ ) # will not be recompiled due to automatic dynamic shapes
229+ 
230+ self.assertEqual(cnt_dynamic.frame_count, 1)
231+ 
232+ @disable_cache_limit()
233+ def test_nested_compile_outer_wins(self):
234+ def fn(a, b):
235+ return a - b * 10
236+ 
237+ torch._dynamo.reset()
238+ 
239+ cnt_dynamic = torch._dynamo.testing.CompileCounter()
240+ cnt_dynamic_1 = torch._dynamo.testing.CompileCounter()
241+ with torch._dynamo.config.patch(
242+ automatic_dynamic_shapes=False, assume_static_by_default=True
243+ ):
244+ opt_fn_static_shape = torch._dynamo.optimize(cnt_dynamic)(fn)
245+ opt_fn_static_shape(torch.randn(2), torch.randn(2))
246+ opt_fn_static_shape(torch.randn(3), torch.randn(3))
247+ 
248+ self.assertEqual(cnt_dynamic.frame_count, 2)
249+ 
250+ with torch._dynamo.config.patch(
251+ automatic_dynamic_shapes=True, assume_static_by_default=False
252+ ):
253+ opt_fn_dynamic = torch._dynamo.optimize(cnt_dynamic_1)(
254+ lambda x, y: opt_fn_static_shape(x, y)
255+ )
256+ for i in range(2, 12):
257+ opt_fn_dynamic(
258+ torch.randn(i), torch.randn(i)
259+ ) # will be recompiled under new config
260+ 
261+ self.assertEqual(cnt_dynamic.frame_count, 2)
262+ self.assertEqual(cnt_dynamic_1.frame_count, 1)
263+ 
264+ @disable_cache_limit()
265+ def test_nested_fn_does_not_inherit_outer_config(self):
266+ def g1(x):
267+ return x + 1
268+ 
269+ def g2(x):
270+ return x * 2
271+ 
272+ def f(x):
273+ x = g1(x)
274+ torch._dynamo.graph_break()
275+ return g2(x)
276+ 
277+ torch._dynamo.reset()
278+ 
279+ cnt_dynamic = torch._dynamo.testing.CompileCounter()
280+ cnt_dynamic_1 = torch._dynamo.testing.CompileCounter()
281+ 
282+ opt_fn_static_shape = torch._dynamo.optimize(cnt_dynamic, dynamic=False)(f)
283+ opt_fn_static_shape(torch.randn(2))
284+ opt_fn_static_shape(torch.randn(3))
285+ self.assertEqual(cnt_dynamic.frame_count, 4) # 2 compiles * 2 graphs
286+ 
287+ opt_fn_dynamic = torch._dynamo.optimize(cnt_dynamic_1, dynamic=True)(g2)
288+ 
289+ for i in range(2, 12):
290+ opt_fn_dynamic(
291+ torch.randn(i),
292+ ) # will be recompiled under new config
293+ 
294+ self.assertEqual(cnt_dynamic_1.frame_count, 1)
295+ 
296+ @disable_cache_limit()
297+ def test_multiple_compile_recompiles(self):
298+ cnt_dynamic = torch._dynamo.testing.CompileCounter()
299+ 
300+ def f(dynamic, compile_count):
301+ @torch._dynamo.optimize(cnt_dynamic, dynamic=dynamic)
302+ def g(x):
303+ return x + 1
304+ 
305+ for i in range(2, 12):
306+ g(torch.randn(i)) # will be recompiled under new config
307+ self.assertEqual(cnt_dynamic.frame_count, compile_count)
308+ cnt_dynamic.clear()
309+ 
310+ f(dynamic=True, compile_count=1) # first compile
311+ f(dynamic=False, compile_count=10) # recompile
312+ f(dynamic=True, compile_count=0) # reuse first compile product
313+ 
314+ def test_cache_size_limit(self):
315+ cnt = torch._dynamo.testing.CompileCounter()
316+ key = "_ConfigTests___test_cache_size_limit_key"
317+ try:
318+ torch._dynamo.config._allowed_keys.add(key)
319+ torch._dynamo.config._ConfigTests___test_cache_size_limit_key = -1
320+ with torch._dynamo.config.patch(
321+ {"cache_size_limit": 1, "accumulated_cache_size_limit": 10}
322+ ):
323+ 
324+ def g(x):
325+ return x + 1
326+ 
327+ for i in range(12):
328+ with torch._dynamo.config.patch(
329+ {key: i % 6}
330+ ): # same config doesn't recompile
331+ opt_g = torch._dynamo.optimize(cnt)(g)
332+ opt_g(torch.randn(1))
333+ self.assertEqual(cnt.frame_count, 6)
334+ 
335+ for i in range(6, 12):
336+ with torch._dynamo.config.patch({key: i}):
337+ opt_g = torch._dynamo.optimize(cnt)(g)
338+ opt_g(torch.randn(1))
339+ self.assertEqual(
340+ cnt.frame_count, 10
341+ ) # only recompile up to cache size limit
342+ finally:
343+ if key in torch._dynamo.config._allowed_keys:
344+ torch._dynamo.config._allowed_keys.remove(key)
345+ 
346+ 
347+if __name__ == "__main__":
348+ from torch._dynamo.test_case import run_tests
349+ 
350+ run_tests()
Atest/dynamo/test_interop.py+66-0
@@ -0,0 +1,66 @@
1+# Owner(s): ["module: dynamo"]
2+import torch
3+import torch_npu
4+import torch._dynamo.test_case
5+import torch._dynamo.testing
6+import torch.onnx.operators
7+from torch._dynamo.testing import same
8+ 
9+ 
10+def fn(a, b):
11+ return a + b * 0.67
12+ 
13+ 
14+class InteropTests(torch._dynamo.test_case.TestCase):
15+ def _common(self, fn):
16+ inputs = [torch.randn(10), torch.randn(10)]
17+ ref = fn(*inputs)
18+ opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
19+ res = opt_fn(*inputs)
20+ self.assertTrue(same(ref, res))
21+ 
22+ def test_fx_fn(self):
23+ fx_fn = torch.fx.symbolic_trace(fn)
24+ self._common(lambda a, b: fx_fn(a, b) + 1)
25+ 
26+ def test_script_fn(self):
27+ script_fn = torch.jit.script(fn)
28+ self._common(lambda a, b: script_fn(a, b) + 1)
29+ 
30+ def test_trace_fn(self):
31+ trace_fn = torch.jit.trace(fn, [torch.zeros(10), torch.zeros(10)])
32+ self._common(lambda a, b: trace_fn(a, b) + 1)
33+ 
34+ def test_vmap_in_graph(self):
35+ from functools import wraps
36+ 
37+ from torch._dynamo import allow_in_graph
38+ 
39+ def traceable(f):
40+ f = allow_in_graph(f)
41+ 
42+ @wraps(f)
43+ def wrapper(*args, **kwargs):
44+ return f(*args, **kwargs)
45+ 
46+ return wrapper
47+ 
48+ cnts = torch._dynamo.testing.CompileCounter()
49+ x = torch.randn(3, 5, 3)
50+ 
51+ def fn1(x):
52+ return torch.vmap(torch.Tensor.t)(x)
53+ 
54+ fn_opt = torch.compile(fn1, backend=cnts, fullgraph=True)
55+ fn_opt_traceable = torch.compile(traceable(fn1), backend=cnts, fullgraph=True)
56+ 
57+ self.assertEqual(fn1(x), fn_opt(x))
58+ self.assertEqual(cnts.frame_count, 1)
59+ self.assertEqual(fn_opt(x), fn_opt_traceable(x))
60+ self.assertEqual(cnts.frame_count, 2)
61+ 
62+ 
63+if __name__ == "__main__":
64+ from torch._dynamo.test_case import run_tests
65+ 
66+ run_tests()
Atest/dynamo/test_pre_dispatch.py+76-0
@@ -0,0 +1,76 @@
1+# Owner(s): ["module: dynamo"]
2+import torch
3+import torch_npu
4+import torch._dynamo
5+import torch._dynamo.test_case
6+ 
7+ 
8+class PreDispatchTests(torch._dynamo.test_case.TestCase):
9+ def test_no_grad_simple(self):
10+ def f(a):
11+ b = a.sin()
12+ with torch.no_grad():
13+ c = b.cos()
14+ return b * c.sin()
15+ 
16+ f_compiled = torch.compile(f, backend="pre_dispatch_eager")
17+ 
18+ a_ref = torch.randn(4, requires_grad=True)
19+ a_test = a_ref.clone().detach().requires_grad_(True)
20+ 
21+ out_ref = f(a_ref)
22+ out_test = f_compiled(a_test)
23+ self.assertEqual(out_ref, out_test)
24+ 
25+ out_ref.sum().backward()
26+ out_test.sum().backward()
27+ self.assertEqual(a_ref.grad, a_test.grad)
28+ 
29+ def test_enable_grad_and_no_grad(self):
30+ def f(a):
31+ b = a * 2
32+ with torch.no_grad():
33+ c = b * 3
34+ with torch.enable_grad():
35+ d = c * 4
36+ e = d * 5
37+ return b + c + d + e
38+ 
39+ f_compiled = torch.compile(f, backend="pre_dispatch_eager")
40+ 
41+ a_ref = torch.randn(4, requires_grad=True)
42+ a_test = a_ref.clone().detach().requires_grad_(True)
43+ 
44+ out_ref = f(a_ref)
45+ out_test = f_compiled(a_test)
46+ self.assertEqual(out_ref, out_test)
47+ 
48+ out_ref.sum().backward()
49+ out_test.sum().backward()
50+ self.assertEqual(a_ref.grad, a_test.grad)
51+ 
52+ def test_autocast_simple(self):
53+ def f(a):
54+ b = a * 2
55+ with torch.amp.autocast(device_type="cpu"):
56+ c = torch.matmul(b, b)
57+ return b + c
58+ 
59+ f_compiled = torch.compile(f, backend="pre_dispatch_eager")
60+ 
61+ a_ref = torch.randn(4, device="cpu", requires_grad=True)
62+ a_test = a_ref.clone().detach().requires_grad_(True)
63+ 
64+ out_ref = f(a_ref)
65+ out_test = f_compiled(a_test)
66+ self.assertEqual(out_ref, out_test)
67+ 
68+ out_ref.sum().backward()
69+ out_test.sum().backward()
70+ self.assertEqual(a_ref.grad, a_test.grad)
71+ 
72+ 
73+if __name__ == "__main__":
74+ from torch._dynamo.test_case import run_tests
75+ 
76+ run_tests()
Atest/dynamo/test_skip_non_tensor.py+193-0
@@ -0,0 +1,193 @@
1+# Owner(s): ["module: dynamo"]
2+from unittest.mock import patch
3+ 
4+import torch
5+import torch_npu
6+import torch._dynamo
7+import torch._dynamo.test_case
8+from torch._dynamo.testing import CompileCounter
9+ 
10+_variable = 0
11+_variable_2 = 0
12+ 
13+ 
14+def user_function():
15+ return torch._utils.is_compiling()
16+ 
17+ 
18+def user_generator():
19+ for _ in range(1):
20+ yield torch._utils.is_compiling()
21+ return
22+ 
23+ 
24+class MyModule(torch.nn.Module):
25+ def __init__(self, mode: int):
26+ super().__init__()
27+ self.mode = mode
28+ self.register_forward_pre_hook(self.pre_forward, with_kwargs=True)
29+ 
30+ def pre_forward(self, module, args, kwargs):
31+ if self.mode == 5:
32+ if user_function():
33+ global _variable
34+ _variable += 1
35+ return args, kwargs
36+ 
37+ def forward(self, x):
38+ global _variable, _variable_2
39+ 
40+ if self.mode == 1:
41+ if torch._utils.is_compiling():
42+ _variable += 1
43+ else:
44+ _variable_2 += 1
45+ elif self.mode == 2:
46+ if user_function():
47+ _variable += 1
48+ elif self.mode == 3:
49+ lambda_f = lambda: torch._utils.is_compiling() # noqa: E731
50+ if lambda_f():
51+ _variable += 1
52+ elif self.mode == 4:
53+ for cond in user_generator():
54+ if cond:
55+ _variable += 1
56+ elif self.mode == 5:
57+ x += 1
58+ elif self.mode == 6:
59+ if user_function():
60+ torch._dynamo.graph_break()
61+ _variable += 1
62+ return x
63+ 
64+ 
65+class SkipNonTensorTests(torch._dynamo.test_case.TestCase):
66+ def test_add_tensor1(self):
67+ def fn(a, b):
68+ return a + b
69+ 
70+ counter = CompileCounter()
71+ x = torch.randn(4)
72+ y = 5
73+ opt_fn = torch._dynamo.optimize_assert(counter)(fn)
74+ opt_fn(x, y)
75+ 
76+ assert counter.op_count == 1
77+ 
78+ def test_add_tensor2(self):
79+ def fn(a, b):
80+ return torch.add(a, b)
81+ 
82+ counter = CompileCounter()
83+ 
84+ x = torch.randn(4)
85+ y = 5
86+ opt_fn = torch._dynamo.optimize_assert(counter)(fn)
87+ opt_fn(x, y)
88+ 
89+ assert counter.op_count == 1
90+ 
91+ def test_add_tensor_list(self):
92+ def fn(lst):
93+ return lst[0] + lst[1]
94+ 
95+ counter = CompileCounter()
96+ x = torch.randn(4)
97+ y = 5
98+ opt_fn = torch._dynamo.optimize_assert(counter)(fn)
99+ opt_fn([x, y])
100+ 
101+ assert counter.op_count == 1
102+ 
103+ def test_add_tensor_dict(self):
104+ def fn(dt):
105+ return dt["a"] + dt["b"]
106+ 
107+ counter = CompileCounter()
108+ x = torch.randn(4)
109+ y = 5
110+ opt_fn = torch._dynamo.optimize_assert(counter)(fn)
111+ opt_fn({"a": x, "b": y})
112+ 
113+ assert counter.op_count == 1
114+ 
115+ def test_add_skip(self):
116+ def fn(a, b):
117+ return a + b
118+ 
119+ counter = CompileCounter()
120+ opt_fn = torch._dynamo.optimize_assert(counter)(fn)
121+ x = 4
122+ y = 5
123+ opt_fn(x, y)
124+ 
125+ assert counter.op_count == 0
126+ 
127+ @patch.object(torch._dynamo.config, "raise_on_ctx_manager_usage", False)
128+ def test_recursive_list(self):
129+ def fn(x):
130+ return x
131+ 
132+ counter = CompileCounter()
133+ 
134+ x = []
135+ x.append(x)
136+ with torch._dynamo.optimize_assert(counter):
137+ fn(x)
138+ 
139+ assert counter.op_count == 0
140+ 
141+ @patch.object(torch._dynamo.config, "raise_on_ctx_manager_usage", False)
142+ def test_custom_list(self):
143+ def fn(x):
144+ return x[0] + x[1]
145+ 
146+ counter = CompileCounter()
147+ 
148+ class Foo(list):
149+ def __iter__(self):
150+ raise Exception()
151+ 
152+ def __len__(self):
153+ raise Exception()
154+ 
155+ x = Foo()
156+ x.append(torch.randn(4))
157+ x.append(torch.randn(4))
158+ with torch._dynamo.optimize_assert(counter):
159+ fn(x)
160+ 
161+ assert counter.op_count == 0
162+ 
163+ def test_do_not_skip_side_effects(self):
164+ # see pytorch issue 110765
165+ 
166+ # By invoking torch._utils.is_compiling(),
167+ # there may be side-effects inconsistent with eager when
168+ # compiling. Thus we force dynamo to commit the graph,
169+ # even if it does not perform any tensor operation
170+ global _variable, _variable_2
171+ 
172+ for mode in range(1, 7):
173+ _variable = 0
174+ _variable_2 = 0
175+ 
176+ mod = MyModule(mode=mode)
177+ model = torch._dynamo.optimize(backend="eager", nopython=mode != 6)(mod)
178+ assert _variable == 0
179+ assert _variable_2 == 0
180+ 
181+ model(torch.tensor([1]))
182+ assert _variable == 1
183+ assert _variable_2 == 0
184+ 
185+ model(torch.tensor([1]))
186+ assert _variable == 2
187+ assert _variable_2 == 0
188+ 
189+ 
190+if __name__ == "__main__":
191+ from torch._dynamo.test_case import run_tests
192+ 
193+ run_tests()
Mtest/unsupported_test_cases/.pytorch-disabled-tests.json+8-0
@@ -1,4 +1,12 @@
1{1{
2+ "test_vmap_in_graph (__main__.InteropTests)": ["", [""]],
3+ "test_disable_for_custom_op (__main__.DecoratorTests)": ["", [""]],
4+ "test_graph_break (__main__.DecoratorTests)": ["", [""]],
5+ "test_nested_disable_decorator (__main__.DecoratorTests)": ["", [""]],
6+ "test_create (__main__.TestBaseOutput)": ["", [""]],
7+ "test_assign (__main__.TestBaseOutput)": ["", [""]],
8+ "test_nested_fn_does_not_inherit_outer_config (__main__.ConfigTests)": ["", [""]],
9+ "test_do_not_skip_side_effects (__main__.SkipNonTensorTests)": ["", [""]],
2 "test_HF_bert_model_output (__main__.TestModelOutput)": ["", [""]],10 "test_HF_bert_model_output (__main__.TestModelOutput)": ["", [""]],
3 "test_module_attribute_mutation_violation_negative_4 (__main__.MutationExportTests)": ["", [""]],11 "test_module_attribute_mutation_violation_negative_4 (__main__.MutationExportTests)": ["", [""]],
4 "test_cudnn_rnn (__main__.FakeTensorTest)": ["", [""]],12 "test_cudnn_rnn (__main__.FakeTensorTest)": ["", [""]],