test: sync dynamo compile cache tests to v2.12.0-pytorch-AtomGit

已合并

test: sync dynamo compile cache tests to v2.12.0 #37711

wangzili121创建于 6月5日从wangzili121/pytorch： v2.12.0 合入到 Ascend/pytorch： v2.12.0

test: sync dynamo compile cache tests to v2.12.0 #37711

已合并

wangzili121创建于 6月5日

从wangzili121/pytorch v2.12.0合入到Ascend/pytorch v2.12.0

共 4 个文件变更+734-9

Mtest/_inductor/test_codecache.py+136-9

	@@ -1,13 +1,139 @@
1	-import pytest	1	+# lintrunner: skip PYFMT
		2	+# Owner(s): ["module: inductor"]
		3	+"""Module for inductor codecache tests."""
		4	+
		5	+import os
		6	+import shutil
		7	+import unittest
		8	+
2	import torch	9	import torch
3	-from torch.testing._internal.common_utils import run_tests	10	+import torch_npu # noqa: F401
4	-from torch._inductor.codecache import CacheBase	11	+import torch_npu._inductor # noqa: F401
5	-from testutils import TestUtils	12	+from torch._dynamo.package import DynamoCache
		OO O openLiBingCI6月5日此条代码评论区间+6至+12 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月5日此条代码评论区间+6至+12 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
6	-import torch_npu	13	+from torch._dynamo.precompile_context import PrecompileContext
		OO O openLiBingCI6月5日此条代码评论区间+7至+13 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月5日此条代码评论区间+7至+13 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
7	-import torch_npu._inductor	14	+from torch._dynamo.utils import counters
		15	+from torch._functorch._aot_autograd.autograd_cache import AOTAutogradCache
		16	+from torch._inductor import config
		17	+from torch._inductor.codecache import CacheBase, PyCodeCache
		18	+from torch._inductor.runtime.runtime_utils import cache_dir
		19	+from torch._inductor.test_case import run_tests, TestCase
		20	+from torch._inductor.utils import clear_caches, fresh_cache
		21	+from torch.compiler._cache import CacheArtifactManager
		22	+from torch.testing._internal.common_utils import (
		23	+ instantiate_parametrized_tests,
		24	+ parametrize,
		25	+)
		26	+from torch.testing._internal.inductor_utils import requires_triton
8		27
9		28
10	-class TestCodeCache(TestUtils):	29	+DEVICES = ("npu", "cpu")
		30	+
		31	+
		32	+@instantiate_parametrized_tests
		33	+class TestCachingPrecompileCodeCache(TestCase):
		34	+ def setUp(self):
		35	+ super().setUp()
		36	+ counters.clear()
		37	+ DynamoCache.clear()
		38	+ PrecompileContext.clear()
		39	+ AOTAutogradCache.clear()
		40	+ CacheArtifactManager.clear()
		41	+ torch._dynamo.reset()
		42	+
		43	+ def reset(self):
		44	+ AOTAutogradCache.clear()
		45	+ DynamoCache.clear()
		46	+ PrecompileContext.clear()
		47	+ PyCodeCache.cache_clear(purge=True)
		48	+ torch._dynamo.reset()
		49	+ clear_caches()
		50	+
		51	+ @requires_triton()
		52	+ @config.patch(
		53	+ {
		54	+ "fx_graph_cache": True,
		55	+ "fx_graph_remote_cache": False,
		56	+ "autotune_local_cache": True,
		57	+ }
		58	+ )
		59	+ @torch._dynamo.config.patch({"caching_precompile": True})
		60	+ @parametrize("dynamic", (False, True))
		61	+ @parametrize("device", DEVICES)
		62	+ @parametrize("dtype", (torch.float32, torch.bfloat16))
		63	+ def test_cache_hot_load_caching_precompile(self, device, dtype, dynamic):
		64	+ if device == "npu" and not torch.npu.is_available():
		65	+ raise unittest.SkipTest("Requires NPU")
		66	+
		67	+ def fn(x, y):
		68	+ return x.sin() @ y
		69	+
		70	+ a = torch.rand(100, 100, dtype=dtype, device=device, requires_grad=True)
		71	+ b = torch.rand(100, 100, dtype=dtype, device=device, requires_grad=True)
		72	+
		73	+ # Record artifacts.
		74	+ with fresh_cache():
		75	+ compiled_fn = torch.compile(fn, dynamic=dynamic)
		76	+
		77	+ eager_result = fn(a, b)
		78	+ compiled_result = compiled_fn(a, b)
		79	+ compiled_result.sum().backward()
		80	+ self.assertEqual(eager_result, compiled_result)
		81	+ self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1)
		82	+ self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0)
		83	+ self.assertEqual(counters["dynamo_cache"]["dynamo_cache_miss"], 1)
		84	+ self.assertEqual(counters["dynamo_cache"]["dynamo_cache_hit"], 0)
		85	+
		86	+ artifacts = torch.compiler.save_cache_artifacts()
		87	+ self.assertEqual(artifacts is not None, True)
		88	+
		89	+ artifact_bytes, cache_info = artifacts
		90	+
		91	+ autotune_expect = 1 if device == "npu" else 0
		92	+ self.assertEqual(len(cache_info.inductor_artifacts), 2)
		93	+ self.assertEqual(len(cache_info.autotune_artifacts), autotune_expect)
		94	+ self.assertEqual(len(cache_info.aot_autograd_artifacts), 1)
		95	+ self.assertEqual(len(cache_info.pgo_artifacts), 0)
		96	+ self.assertEqual(len(cache_info.precompile_artifacts), 1)
		97	+
		98	+ self.reset()
		99	+ shutil.rmtree(os.path.join(cache_dir(), "triton"), ignore_errors=True)
		100	+
		101	+ # Without loading artifacts, a new compile should not hit dynamo cache.
		102	+ with fresh_cache():
		103	+ eager_result = fn(a, b)
		104	+ compiled_fn = torch.compile(fn, dynamic=dynamic)
		105	+ compiled_result = compiled_fn(a, b)
		106	+ compiled_result.sum().backward()
		107	+ self.assertEqual(eager_result, compiled_result)
		108	+ self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 2)
		109	+ self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0)
		110	+ self.assertEqual(counters["dynamo_cache"]["dynamo_cache_miss"], 2)
		111	+ self.assertEqual(counters["dynamo_cache"]["dynamo_cache_hit"], 0)
		112	+
		113	+ self.reset()
		114	+ shutil.rmtree(os.path.join(cache_dir(), "triton"), ignore_errors=True)
		115	+
		116	+ # After loading artifacts, recompilation is forbidden and dynamo cache hits.
		117	+ with fresh_cache(), torch.compiler.set_stance("fail_on_recompile"):
		118	+ cache_info = torch.compiler.load_cache_artifacts(artifact_bytes)
		119	+ self.assertEqual(len(cache_info.inductor_artifacts), 2)
		120	+ self.assertEqual(len(cache_info.autotune_artifacts), autotune_expect)
		121	+ self.assertEqual(len(cache_info.aot_autograd_artifacts), 1)
		122	+ self.assertEqual(len(cache_info.pgo_artifacts), 0)
		123	+ self.assertEqual(len(cache_info.precompile_artifacts), 1)
		124	+
		125	+ compiled_fn = torch.compile(fn, dynamic=dynamic)
		126	+ eager_result = fn(a, b)
		127	+ compiled_result = compiled_fn(a, b)
		128	+ compiled_result.sum().backward()
		129	+ self.assertEqual(eager_result, compiled_result)
		130	+ self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 2)
		131	+ self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0)
		132	+ self.assertEqual(counters["dynamo_cache"]["dynamo_cache_miss"], 2)
		133	+ self.assertEqual(counters["dynamo_cache"]["dynamo_cache_hit"], 1)
		134	+
		135	+
		136	+class TestCodeCache(TestCase):
11	def test_codecache(self):	137	def test_codecache(self):
12	device_properties = torch_npu.npu.get_device_properties(	138	device_properties = torch_npu.npu.get_device_properties(
13	torch_npu.npu.current_device()	139	torch_npu.npu.current_device()
	@@ -17,11 +143,12 @@ class TestCodeCache(TestUtils):
17	self.assertEqual(system1["device"]["name"], device_properties.name)	143	self.assertEqual(system1["device"]["name"], device_properties.name)
18	self.assertEqual(system1["version"]["cann"], torch.version.cann)	144	self.assertEqual(system1["version"]["cann"], torch.version.cann)
19		145
20	- from torch_npu.contrib import transfer_to_npu	146	+ from torch_npu.contrib import transfer_to_npu # noqa: F401
		147	+
21	system2 = CacheBase.get_system()	148	system2 = CacheBase.get_system()
22	self.assertEqual(system2["device"]["name"], device_properties.name)	149	self.assertEqual(system2["device"]["name"], device_properties.name)
23	self.assertEqual(system2["version"]["cann"], torch.version.cann)	150	self.assertEqual(system2["version"]["cann"], torch.version.cann)
24		151
25		152
26	if __name__ == "__main__":	153	if __name__ == "__main__":
27	- run_tests()	154	+ run_tests()
		OO O openLiBingCI6月5日此条代码评论区间+133至+154 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月5日此条代码评论区间+133至+154 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。

Atest/dynamo/test_guard_serialization.py+167-0

@@ -0,0 +1,167 @@
	1	+# lintrunner: skip PYFMT
	2	+# Owner(s): ["module: dynamo"]
	3	+"""Module for dynamo guard serialization tests."""
	4	+
	5	+import dataclasses
	6	+import sys
	7	+import types
	8	+
	9	+import torch
	10	+import torch_npu # noqa: F401
	11	+import torch_npu._inductor # noqa: F401
	12	+import torch._dynamo.package
	OO O openLiBingCI6月5日此条代码评论区间+7至+12 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月5日此条代码评论区间+7至+12 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
	13	+import torch._inductor.test_case
	OO O openLiBingCI6月5日此条代码评论区间+9至+13 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月5日此条代码评论区间+9至+13 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
	14	+from torch._dynamo.bytecode_transformation import transform_code_object
	15	+from torch._dynamo.guards import CheckFunctionManager, CompileId
	16	+from torch._dynamo.symbolic_convert import (
	17	+ ExceptionStack,
	18	+ InstructionTranslator,
	19	+ SpeculationLog,
	20	+)
	21	+from torch._dynamo.utils import dynamo_timed, get_metrics_context
	22	+from torch._guards import compile_context, CompileContext, tracing
	23	+
	24	+
	25	+@dataclasses.dataclass
	26	+class _FrameState:
	27	+ f_locals: dict
	28	+ f_globals: dict
	29	+ f_code: types.CodeType
	30	+ f_builtins: dict
	31	+
	32	+
	33	+class TestGuardSerializationBase(torch._inductor.test_case.TestCase):
	34	+ def _tracefunc(self, frame, event, arg):
	35	+ if event != "call":
	36	+ return
	37	+ if self._frame_state is not None:
	38	+ return
	39	+
	40	+ self._frame_state = _FrameState(
	41	+ f_locals=dict(frame.f_locals),
	42	+ f_globals=frame.f_globals,
	43	+ f_code=frame.f_code,
	44	+ f_builtins=frame.f_builtins,
	45	+ )
	46	+
	47	+ def _test_serialization(self, guard_type, fn, *args):
	48	+ torch._dynamo.reset()
	49	+ self._frame_state = None
	50	+ sys.settrace(self._tracefunc)
	51	+ try:
	52	+ fn(*args)
	53	+ finally:
	54	+ sys.settrace(None)
	55	+
	56	+ self.assertEqual(self._frame_state is not None, True)
	57	+
	58	+ def guard_filter_fn(guards):
	59	+ ret = [
	60	+ g.guard_type == guard_type or guard_type in g.derived_guard_types
	61	+ for g in guards
	62	+ ]
	63	+ self.assertEqual(any(ret), True)
	64	+ return ret
	65	+
	66	+ ref_gm = None
	67	+ loaded_gm = None
	68	+
	69	+ def transform(instructions: list, code_options: dict[str, object]):
	70	+ nonlocal ref_gm
	71	+ nonlocal loaded_gm
	72	+
	73	+ torch._dynamo.convert_frame.initial_global_state = (
	74	+ torch._C._dynamo.guards.GlobalStateGuard()
	75	+ )
	76	+ tracer = InstructionTranslator(
	77	+ instructions,
	78	+ self._frame_state.f_code,
	79	+ self._frame_state.f_locals,
	80	+ self._frame_state.f_globals,
	81	+ self._frame_state.f_builtins,
	82	+ fn.__closure__ or (),
	83	+ torch.overrides._get_current_function_mode_stack(),
	84	+ code_options,
	85	+ torch._dynamo.lookup_backend("eager"),
	86	+ one_graph=False,
	87	+ export=False,
	88	+ export_constraints=None,
	89	+ frame_state=None,
	90	+ speculation_log=SpeculationLog(),
	91	+ exn_vt_stack=ExceptionStack(),
	92	+ distributed_state=None,
	93	+ package=None,
	94	+ )
	95	+ with (
	96	+ compile_context(
	97	+ CompileContext(CompileId(frame_id=0, frame_compile_id=0))
	98	+ ),
	99	+ tracing(tracer.output.tracing_context),
	100	+ tracer.set_current_tx(),
	101	+ get_metrics_context(),
	102	+ dynamo_timed(""),
	103	+ ):
	104	+ tracer.run()
	105	+
	106	+ ref_gm = CheckFunctionManager(
	107	+ self._frame_state.f_code,
	108	+ tracer.output,
	109	+ guard_filter_fn=guard_filter_fn,
	110	+ ).guard_manager
	111	+
	112	+ check_fn_manager = CheckFunctionManager(
	113	+ self._frame_state.f_code,
	114	+ tracer.output,
	115	+ guard_filter_fn=guard_filter_fn,
	116	+ save_guards=True,
	117	+ )
	118	+ guards_state = check_fn_manager.guards_state
	119	+ self.assertEqual(guards_state is not None, True)
	120	+ guards_state = torch._dynamo.package.load_guards_state(guards_state)
	121	+
	122	+ loaded_gm = torch._dynamo.package.load_guard_manager(
	123	+ guards_state,
	124	+ self._frame_state.f_code,
	125	+ self._frame_state.f_globals,
	126	+ )
	127	+
	128	+ try:
	129	+ transform_code_object(self._frame_state.f_code, transform)
	130	+ finally:
	131	+ torch._dynamo.convert_frame.initial_global_state = None
	132	+ self._frame_state = None
	133	+
	134	+ self.assertEqual(ref_gm is not None, True)
	135	+ self.assertEqual(loaded_gm is not None, True)
	136	+ return ref_gm, loaded_gm
	137	+
	138	+ def _test_check_fn(self, ref, loaded, inputs, expected):
	139	+ self.assertEqual(isinstance(inputs, dict), True)
	140	+ self.assertEqual(ref.check(inputs), expected)
	141	+ self.assertEqual(ref.check(inputs), loaded.check(inputs))
	142	+
	143	+
	144	+@torch._dynamo.config.patch({"strict_precompile": True})
	145	+class TestCachingPrecompileGuardSerialization(TestGuardSerializationBase):
	146	+ @torch._dynamo.config.patch(caching_precompile=True)
	147	+ def test_id_match_with_config(self):
	148	+ def fn(x):
	149	+ return x + id(x)
	150	+
	151	+ ref, loaded = self._test_serialization("ID_MATCH", fn, torch.randn(3))
	152	+ self._test_check_fn(ref, loaded, {"x": torch.randn(3)}, True)
	153	+
	154	+ def fn(x):
	155	+ # torch.no_grad() installs a CLASS_MATCH guard.
	156	+ with torch.no_grad():
	157	+ y = x * 2
	158	+ return y
	159	+
	160	+ ref, loaded = self._test_serialization("CLASS_MATCH", fn, torch.randn(3))
	161	+ self._test_check_fn(ref, loaded, {"x": torch.randn(3)}, True)
	162	+
	163	+
	164	+if __name__ == "__main__":
	165	+ from torch._dynamo.test_case import run_tests
	166	+
	167	+ run_tests()

Atest/dynamo/test_package.py+301-0

@@ -0,0 +1,301 @@
	1	+# lintrunner: skip PYFMT
	2	+# Owner(s): ["module: dynamo"]
	3	+"""Module for dynamo package tests."""
	4	+
	5	+import sys
	6	+import unittest
	7	+
	8	+import torch
	9	+import torch_npu # noqa: F401
	10	+import torch_npu._inductor # noqa: F401
	11	+import torch._inductor.test_case
	OO O openLiBingCI6月5日此条代码评论区间+6至+11 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月5日此条代码评论区间+6至+11 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
	12	+from torch._dynamo.package import DynamoCache
	OO O openLiBingCI6月5日此条代码评论区间+8至+12 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月5日此条代码评论区间+8至+12 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
	13	+from torch._dynamo.precompile_context import PrecompileContext
	14	+from torch._dynamo.testing import reduce_to_scalar_loss
	15	+from torch._functorch import config as functorch_config
	16	+from torch.testing._internal.common_utils import (
	17	+ instantiate_parametrized_tests,
	18	+ parametrize,
	19	+)
	20	+
	21	+
	22	+DEVICES = ("cpu", "npu")
	23	+
	24	+
	25	+def compute_loss_helper(x):
	26	+ return reduce_to_scalar_loss(x)
	27	+
	28	+
	29	+@functorch_config.patch("bundled_autograd_cache", True)
	30	+@torch._dynamo.config.patch({"strict_precompile": True})
	31	+@instantiate_parametrized_tests
	32	+class TestCachingPrecompilePackage(torch._inductor.test_case.TestCase):
	33	+ def setUp(self):
	34	+ super().setUp()
	35	+ torch._dynamo.reset()
	36	+ torch._dynamo.utils.counters.clear()
	37	+ DynamoCache.clear()
	38	+ PrecompileContext.clear()
	39	+
	40	+ def _save_and_reload(self, expected_backends, expected_dynamo):
	41	+ debug_info = PrecompileContext.save_to_dynamo_cache()
	42	+ self.assertEqual(len(debug_info["dynamo"]), expected_dynamo)
	43	+ self.assertEqual(len(debug_info["backends"]), expected_backends)
	44	+ torch._dynamo.reset()
	45	+ PrecompileContext.clear()
	46	+
	47	+ @staticmethod
	48	+ def _check_device(device):
	49	+ if device == "npu" and not torch.npu.is_available():
	50	+ raise unittest.SkipTest("Requires NPU")
	51	+
	52	+ @parametrize("device", DEVICES)
	53	+ @torch._dynamo.config.patch(caching_precompile=True)
	54	+ def test_automatic_dynamo_serialize(self, device):
	55	+ self._check_device(device)
	56	+
	57	+ def fn(x):
	58	+ return x.sin() + x.cos()
	59	+
	60	+ def fn2(x):
	61	+ return x.cos() + x
	62	+
	63	+ arg1 = torch.randn(3, 2, device=device)
	64	+ arg2 = torch.randn(5, 2, device=device)
	65	+ expected = [fn(arg1), fn2(arg2)]
	66	+ compiled_fn1 = torch.compile(fn)
	67	+ compiled_fn2 = torch.compile(fn2)
	68	+ result = [compiled_fn1(arg1), compiled_fn2(arg2)]
	69	+ self.assertEqual(expected, result)
	70	+ DynamoCache.clear()
	71	+ total_frames = torch._dynamo.convert_frame.FRAME_COUNTER
	72	+
	73	+ self._save_and_reload(expected_backends=2, expected_dynamo=2)
	74	+
	75	+ compiled_fn1 = torch.compile(fn)
	76	+ compiled_fn2 = torch.compile(fn2)
	77	+ with torch.compiler.set_stance("fail_on_recompile"):
	78	+ result1 = compiled_fn1(arg1)
	79	+ result2 = compiled_fn2(arg2)
	80	+ self.assertEqual(expected, [result1, result2])
	81	+ self.assertEqual(torch._dynamo.convert_frame.FRAME_COUNTER, total_frames)
	82	+
	83	+ @parametrize("device", DEVICES)
	84	+ @torch._dynamo.config.patch(caching_precompile=True)
	85	+ def test_automatic_dynamo_recompiles(self, device):
	86	+ self._check_device(device)
	87	+
	88	+ def fn(x):
	89	+ return x.sin() + x.cos()
	90	+
	91	+ arg1 = torch.randn(3, 2, device=device)
	92	+ arg2 = torch.randn(5, 2, device=device)
	93	+ compiled_fn = torch.compile(fn)
	94	+ expected1 = compiled_fn(arg1)
	95	+ expected2 = compiled_fn(arg2)
	96	+ total_frames = torch._dynamo.convert_frame.FRAME_COUNTER
	97	+
	98	+ self._save_and_reload(expected_backends=2, expected_dynamo=1)
	99	+
	100	+ compiled_fn = torch.compile(fn)
	101	+ with torch.compiler.set_stance("fail_on_recompile"):
	102	+ result1 = compiled_fn(arg1)
	103	+ result2 = compiled_fn(arg2)
	104	+ arg3 = torch.randn(7, 2, device=device)
	105	+ compiled_fn(arg3)
	106	+ self.assertEqual(result1, expected1)
	107	+ self.assertEqual(result2, expected2)
	108	+ self.assertEqual(torch._dynamo.convert_frame.FRAME_COUNTER, total_frames)
	109	+
	110	+ @parametrize("device", ("cpu",))
	111	+ @torch._dynamo.config.patch(caching_precompile=True)
	112	+ def test_automatic_dynamo_graph_breaks(self, device):
	113	+ self._check_device(device)
	114	+
	115	+ def fn(x, l, r):
	116	+ if l > r:
	117	+ return x.sum()
	118	+ mid = (l + r) // 2
	119	+ if x.sum() == mid:
	120	+ return x.sum()
	121	+ elif x.sum() < mid:
	122	+ return fn(x, l, mid)
	123	+ else:
	124	+ return fn(x, mid + 1, r)
	125	+
	126	+ def guard_filter_fn(guards):
	127	+ return [
	128	+ guard.guard_type not in ("CLOSURE_MATCH", "FUNCTION_MATCH")
	129	+ for guard in guards
	130	+ ]
	131	+
	132	+ compiled_fn = torch._dynamo.optimize(
	133	+ backend="inductor", guard_filter_fn=guard_filter_fn
	134	+ )(fn)
	135	+ n = 10
	136	+ args_list = [(torch.tensor(x, device=device), 0, n - 1) for x in range(n)]
	137	+ for args in args_list:
	138	+ compiled_fn(*args)
	139	+
	140	+ total_frames = torch._dynamo.convert_frame.FRAME_COUNTER
	141	+ self._save_and_reload(expected_backends=9, expected_dynamo=1)
	142	+
	143	+ compiled_fn = torch._dynamo.optimize(
	144	+ backend="inductor", guard_filter_fn=guard_filter_fn
	145	+ )(fn)
	146	+ with torch.compiler.set_stance("fail_on_recompile"):
	147	+ for args in args_list:
	148	+ self.assertEqual(compiled_fn(*args), args[0].sum())
	149	+ self.assertEqual(torch._dynamo.convert_frame.FRAME_COUNTER, total_frames)
	150	+
	151	+ @parametrize("device", DEVICES)
	152	+ @torch._dynamo.config.patch(caching_precompile=True)
	153	+ def test_automatic_dynamo_lazy_backward(self, device):
	154	+ self._check_device(device)
	155	+
	156	+ def fn(x):
	157	+ return x.sin() + x.cos()
	158	+
	159	+ arg1 = torch.randn(3, 2, device=device, requires_grad=True)
	160	+ arg2 = arg1.clone().detach_().requires_grad_(True)
	161	+
	162	+ compiled_fn = torch.compile(fn)
	163	+ expected1 = compiled_fn(arg1)
	164	+ expected1.sum().backward()
	165	+ total_frames = torch._dynamo.convert_frame.FRAME_COUNTER
	166	+
	167	+ self._save_and_reload(expected_backends=1, expected_dynamo=1)
	168	+
	169	+ compiled_fn = torch.compile(fn)
	170	+ with torch.compiler.set_stance("fail_on_recompile"):
	171	+ expected2 = compiled_fn(arg2)
	172	+ expected2.sum().backward()
	173	+
	174	+ self.assertEqual(torch._dynamo.convert_frame.FRAME_COUNTER, total_frames)
	175	+
	176	+ @parametrize("device", DEVICES)
	177	+ @torch._dynamo.config.patch(caching_precompile=True)
	178	+ def test_graph_break_partial_backend(self, device):
	179	+ self._check_device(device)
	180	+
	181	+ def fn(x):
	182	+ y = x.sin()
	183	+ torch._dynamo.graph_break()
	184	+ return x.sin() + y
	185	+
	186	+ arg1 = torch.randn(3, 2, device=device, requires_grad=True)
	187	+ arg2 = arg1.clone().detach_().requires_grad_(True)
	188	+ compiled_fn = torch.compile(fn)
	189	+ expected1 = compiled_fn(arg1)
	190	+ expected1.sum().backward()
	191	+ total_frames = torch._dynamo.convert_frame.FRAME_COUNTER
	192	+
	193	+ dynamo_entry = next(iter(PrecompileContext._dynamo_cache_entries.values()))
	194	+ for code in dynamo_entry.codes:
	195	+ module = sys.modules[code.python_module]
	196	+ if code.install_to_global:
	197	+ for fn_name in code.function_names:
	198	+ module.__dict__.pop(fn_name)
	199	+ for fn_name in code.function_names:
	200	+ if "resume" in fn_name:
	201	+ self.assertEqual(len(code.backend_ids), 1)
	202	+ backend = code.backend_ids[0]
	203	+ del PrecompileContext._backend_artifacts_by_key[backend]
	204	+
	205	+ self._save_and_reload(expected_backends=1, expected_dynamo=1)
	206	+
	207	+ compiled_fn = torch.compile(fn)
	208	+ expected2 = compiled_fn(arg2)
	209	+ expected2.sum().backward()
	210	+ self.assertEqual(expected1, expected2)
	211	+ self.assertEqual(torch._dynamo.convert_frame.FRAME_COUNTER, total_frames + 1)
	212	+
	213	+ @parametrize("device", DEVICES)
	214	+ @torch._dynamo.config.patch(caching_precompile=True)
	215	+ def test_call_function_from_resume(self, device):
	216	+ self._check_device(device)
	217	+ mod = torch.nn.Linear(2, 3, device=device)
	218	+
	219	+ def foo(x, mod):
	220	+ pred = mod(x)
	221	+ compute_loss_helper(pred).backward()
	222	+ return None
	223	+
	224	+ args = (torch.randn(3, 2, device=device), mod)
	225	+ compiled_fn = torch.compile(foo)
	226	+ compiled_fn(*args)
	227	+ total_frames = torch._dynamo.convert_frame.FRAME_COUNTER
	228	+
	229	+ self._save_and_reload(expected_backends=1, expected_dynamo=1)
	230	+
	231	+ compiled_fn = torch.compile(foo)
	232	+ with torch.compiler.set_stance("fail_on_recompile"):
	233	+ compiled_fn(*args)
	234	+
	235	+ self.assertEqual(torch._dynamo.convert_frame.FRAME_COUNTER, total_frames)
	236	+
	237	+ @parametrize("device", DEVICES)
	238	+ @torch._dynamo.config.patch(caching_precompile=True)
	239	+ def test_code_with_generator(self, device):
	240	+ self._check_device(device)
	241	+
	242	+ def foo(set_of_x):
	243	+ if not all(isinstance(s, torch.Tensor) for s in set_of_x):
	244	+ raise TypeError(
	245	+ f"Expected all elements of set_of_x to be tensors, got {set_of_x}"
	246	+ )
	247	+
	248	+ return torch.cat(set_of_x, dim=0)
	249	+
	250	+ args = ([torch.randn(3, 2, device=device) for _ in range(3)],)
	251	+ compiled_fn = torch.compile(foo)
	252	+ compiled_fn(*args)
	253	+ self._save_and_reload(expected_backends=1, expected_dynamo=1)
	254	+
	255	+ @parametrize("device", DEVICES)
	256	+ @torch._dynamo.config.patch(caching_precompile=True)
	257	+ def test_automatic_dynamo_graph_breaks_from_print_model_as_fn(self, device):
	258	+ self._check_device(device)
	259	+
	260	+ def guard_filter_fn(guards):
	261	+ return [
	262	+ guard.guard_type not in ("CLOSURE_MATCH", "FUNCTION_MATCH")
	263	+ for guard in guards
	264	+ ]
	265	+
	266	+ class TempNN(torch.nn.Module):
	267	+ def forward(self, x):
	268	+ x = torch.nn.functional.relu(x)
	269	+ x *= x
	270	+ x /= 2
	271	+ print(x.sum().item())
	272	+ x += 1
	273	+ return x
	274	+
	275	+ x = torch.rand(10, device=device)
	276	+ model = TempNN()
	277	+ model(x)
	278	+ compiled_fn = torch.compile(
	279	+ model,
	280	+ backend="inductor",
	281	+ options=dict(guard_filter_fn=guard_filter_fn),
	282	+ )
	283	+
	284	+ compiled_fn(x)
	285	+ total_frames = torch._dynamo.convert_frame.FRAME_COUNTER
	286	+ self._save_and_reload(expected_backends=2, expected_dynamo=1)
	287	+
	288	+ del compiled_fn
	289	+
	290	+ with torch.compiler.set_stance("fail_on_recompile"):
	291	+ compiled_fn = torch.compile(
	292	+ model, backend="inductor", options=dict(guard_filter_fn=guard_filter_fn)
	293	+ )
	294	+ compiled_fn(x)
	295	+ self.assertEqual(torch._dynamo.convert_frame.FRAME_COUNTER, total_frames)
	296	+
	297	+
	298	+if __name__ == "__main__":
	299	+ from torch._dynamo.test_case import run_tests
	300	+
	301	+ run_tests()

Atest/dynamo/test_precompile_context.py+130-0

@@ -0,0 +1,130 @@
	1	+# lintrunner: skip PYFMT
	2	+# Owner(s): ["module: dynamo"]
	3	+"""Module for dynamo precompile context tests."""
	4	+
	5	+import unittest
	6	+
	7	+import torch
	8	+import torch_npu # noqa: F401
	9	+import torch_npu._inductor # noqa: F401
	10	+import torch._dynamo
	OO O openLiBingCI6月5日此条代码评论区间+5至+10 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月5日此条代码评论区间+5至+10 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
	11	+import torch._dynamo.test_case
	OO O openLiBingCI6月5日此条代码评论区间+7至+11 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月5日此条代码评论区间+7至+11 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
	12	+import torch._functorch
	13	+from torch._dynamo.precompile_context import BackendCacheArtifact, PrecompileContext
	14	+from torch._functorch import config as functorch_config
	15	+from torch._functorch._aot_autograd.autograd_cache import (
	16	+ BundledAOTAutogradCacheArtifact,
	17	+)
	18	+from torch._inductor.test_case import TestCase as InductorTestCase
	19	+from torch.testing._internal.inductor_utils import requires_triton
	20	+
	21	+
	22	+DEVICE = "npu"
	23	+
	24	+
	25	+@functorch_config.patch({"enable_autograd_cache": True})
	26	+@torch._dynamo.config.patch(
	27	+ {"caching_precompile": True}
	28	+) # Requires bundledaotautograd cache for now
	29	+class PrecompileContextTests(InductorTestCase):
	30	+ def setUp(self):
	31	+ """
	32	+ Reset all counters and caches before each unit test
	33	+ """
	34	+ super().setUp()
	35	+ # Clear PrecompileContext cache artifacts
	36	+ PrecompileContext.clear()
	37	+
	38	+ @staticmethod
	39	+ def _check_device():
	40	+ if not torch.npu.is_available():
	41	+ raise unittest.SkipTest("Requires NPU")
	42	+
	43	+ @requires_triton()
	44	+ def test_basic(self):
	45	+ """
	46	+ Test that after torch.compile, PrecompileContext._new_cache_artifacts length is 1
	47	+ """
	48	+ self._check_device()
	49	+
	50	+ def simple_function(x):
	51	+ return x.sin() + x.cos()
	52	+
	53	+ compiled_fn = torch.compile(simple_function)
	54	+
	55	+ # Run the compiled function
	56	+ x = torch.randn(10, device=DEVICE, requires_grad=True)
	57	+ result = compiled_fn(x)
	58	+ result.sum().backward()
	59	+ self.assertEqual(len(PrecompileContext._dynamo_cache_entries), 1)
	60	+ self.assertEqual(len(PrecompileContext._backend_artifacts_by_key), 1)
	61	+ cache_entries, _ = PrecompileContext.create_cache_entries()
	62	+ self.assertEqual(len(cache_entries), 1)
	63	+
	64	+ @requires_triton()
	65	+ def test_serialize_by_key(self):
	66	+ self._check_device()
	67	+
	68	+ def simple_function(x):
	69	+ return x.sin() + x.cos()
	70	+
	71	+ compiled_fn = torch.compile(simple_function)
	72	+
	73	+ # Run the compiled function
	74	+ x = torch.randn(10, device=DEVICE, requires_grad=True)
	75	+ result = compiled_fn(x)
	76	+ result.sum().backward()
	77	+ self.assertEqual(len(PrecompileContext._dynamo_cache_entries), 1)
	78	+ self.assertEqual(len(PrecompileContext._backend_artifacts_by_key), 1)
	79	+ for key in PrecompileContext._backend_artifacts_by_key:
	80	+ result = PrecompileContext.serialize_artifact_by_key(key)
	81	+ self.assertEqual(isinstance(result, BackendCacheArtifact), True)
	82	+ self.assertEqual(result.key, key)
	83	+
	84	+ # This should still work
	85	+ result, _ = PrecompileContext.create_cache_entries()
	86	+ self.assertEqual(len(result), 1)
	87	+
	88	+ @requires_triton()
	89	+ def test_editable(self):
	90	+ """
	91	+ Test that after torch.compile, PrecompileContext._new_cache_artifacts length is 1
	92	+ """
	93	+ self._check_device()
	94	+
	95	+ def simple_function(x):
	96	+ return x.sin() + x.cos()
	97	+
	98	+ compiled_fn = torch.compile(simple_function)
	99	+
	100	+ # Run the compiled function
	101	+ x = torch.randn(10, device=DEVICE, requires_grad=True)
	102	+ result = compiled_fn(x)
	103	+ result.sum().backward()
	104	+ self.assertEqual(len(PrecompileContext._dynamo_cache_entries), 1)
	105	+ self.assertEqual(len(PrecompileContext._backend_artifacts_by_key), 1)
	106	+ # Find the key for the artifact of type "precompile_aot_autograd"
	107	+ key = next(iter(PrecompileContext._backend_artifacts_by_key))
	108	+
	109	+ def edit_fn(x):
	110	+ x._my_private_field = 42
	111	+ return x
	112	+
	113	+ PrecompileContext.edit_artifact(key, edit_fn)
	114	+
	115	+ result = PrecompileContext.serialize_artifact_by_key(key)
	116	+ self.assertEqual(isinstance(result, BundledAOTAutogradCacheArtifact), True)
	117	+ self.assertEqual(result.key, key)
	118	+
	119	+ result, _ = PrecompileContext.create_cache_entries()
	120	+ self.assertEqual(len(result), 1)
	121	+ aot_autograd_artifacts = next(iter(result.values())).backends
	122	+ self.assertEqual(len(aot_autograd_artifacts), 1)
	123	+ entry = next(iter(aot_autograd_artifacts.values())).content
	124	+ self.assertEqual(entry._my_private_field, 42)
	125	+
	126	+
	127	+if __name__ == "__main__":
	128	+ from torch._dynamo.test_case import run_tests
	129	+
	130	+ run_tests()