perf: defer dynamo and inductor imports on master-pytorch-AtomGit

已合并

perf: defer dynamo and inductor imports on master #44294

黄桂军创建于 8月10日从stevenaw0/pytorch： master_import_sync 合入到 Ascend/pytorch： master

perf: defer dynamo and inductor imports on master #44294

已合并

黄桂军创建于 8月10日

从stevenaw0/pytorch master_import_sync合入到Ascend/pytorch master

共 20 个文件变更+1429-170

Msetup.py+5-0

	@@ -794,5 +794,10 @@ setup(
794	'torch.backends': [	794	'torch.backends': [
795	'torch_npu = torch_npu:_autoload',	795	'torch_npu = torch_npu:_autoload',
796	],	796	],
		797	+ 'torch_dynamo_backends': [
		798	+ 'npu = torch_npu.dynamo:_npu_backend_entrypoint',
		799	+ 'npugraph_ex = torch_npu.dynamo:_npugraph_ex_backend_entrypoint',
		800	+ 'npugraphs = torch_npu.dynamo:_npugraphs_backend_entrypoint',
		801	+ ],
797	}	802	}
798	)	803	)

Mtest/_inductor/test_current_device.py+2-2

	@@ -1,10 +1,10 @@
1	import torch	1	import torch
		2	+import torch._dynamo.testing
2	from torch.testing._internal.common_utils import (	3	from torch.testing._internal.common_utils import (
3	run_tests,	4	run_tests,
4	instantiate_parametrized_tests,	5	instantiate_parametrized_tests,
5	)	6	)
6	from testutils import TestUtils	7	from testutils import TestUtils
7	-import torch_npu
8		8
9		9
10	class TestCurrentDevice(TestUtils):	10	class TestCurrentDevice(TestUtils):
	@@ -30,4 +30,4 @@ instantiate_parametrized_tests(TestCurrentDevice)
30		30
31		31
32	if __name__ == "__main__":	32	if __name__ == "__main__":
33	- run_tests()	33	+ run_tests()

Mtest/_inductor/test_mlir_enable.py+3-0

	@@ -25,6 +25,9 @@ class TestAdd(TestUtils):
25	@parametrize('shape', TestUtils._pointwise_demo_shapes)	25	@parametrize('shape', TestUtils._pointwise_demo_shapes)
26	@parametrize('dtype', ['float32', 'int64'])	26	@parametrize('dtype', ['float32', 'int64'])
27	def test_config_environ_cases(self, shape, dtype):	27	def test_config_environ_cases(self, shape, dtype):
		28	+ # torch._inductor.config is an internal module. Initialize the NPU
		29	+ # config entries explicitly before testing direct config assignment.
		30	+ torch._inductor.config.get_config_copy()
28	torch._inductor.config.npu_backend = "mlir"	31	torch._inductor.config.npu_backend = "mlir"
29	x = self._generate_tensor(shape, dtype)	32	x = self._generate_tensor(shape, dtype)
30	y = self._generate_tensor(shape, dtype)	33	y = self._generate_tensor(shape, dtype)

Atest/dynamo/test_compile_trigger.py+928-0

@@ -0,0 +1,928 @@
	1	+# Owner(s): ["module: dynamo"]
	2	+
	3	+import importlib.util
	4	+import os
	5	+import subprocess
	6	+import sys
	7	+import textwrap
	8	+import unittest
	9	+
	10	+import torch
	11	+
	12	+
	13	+class TorchCompileTriggerTests(unittest.TestCase):
	14	+ def run_in_subprocess(self, code, timeout=60):
	15	+ env = os.environ.copy()
	16	+ env["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"
	17	+ result = subprocess.run(
	18	+ [sys.executable, "-c", textwrap.dedent(code)],
	19	+ capture_output=True,
	20	+ env=env,
	21	+ text=True,
	22	+ timeout=timeout,
	23	+ )
	24	+ self.assertEqual(
	25	+ result.returncode,
	26	+ 0,
	27	+ f"Subprocess failed:\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}",
	28	+ )
	29	+
	30	+ # Verify import installs only the Dynamo post-import trigger.
	31	+ def test_import_installs_dynamo_post_import_trigger(self):
	32	+ self.run_in_subprocess(
	33	+ """
	34	+ import sys
	35	+ import torch
	36	+
	37	+ def loaded(prefix):
	38	+ return [
	39	+ name for name in sys.modules
	40	+ if name == prefix or name.startswith(prefix + ".")
	41	+ ]
	42	+
	43	+ assert not loaded("torch._dynamo")
	44	+ assert not loaded("torch._inductor")
	45	+ src_compile = torch.compile
	46	+ src_wrapper_init = torch._TorchCompileWrapper.__init__
	47	+
	48	+ import torch_npu
	49	+ from torch_npu.utils import _dynamo
	50	+
	51	+ assert not loaded("torch._dynamo")
	52	+ assert not loaded("torch._inductor")
	53	+ assert not loaded("torch_npu._inductor")
	54	+ assert "triton" not in sys.modules
	55	+ assert torch.utils._triton.has_triton is _dynamo.has_triton
	56	+ assert torch.compile is src_compile
	57	+ assert torch._TorchCompileWrapper.__init__ is src_wrapper_init
	58	+ assert any(
	59	+ isinstance(finder, _dynamo._DynamoPostImportFinder)
	60	+ for finder in sys.meta_path
	61	+ )
	62	+ assert not _dynamo._lazy_dynamo_setup.has_run
	63	+ """
	64	+ )
	65	+
	66	+ # Verify public compiler APIs work before the first compile.
	67	+ def test_public_compiler_entries_are_available_before_compile(self):
	68	+ self.run_in_subprocess(
	69	+ """
	70	+ import inspect
	71	+ import sys
	72	+ import torch
	73	+ import torch_npu
	74	+ from torch_npu.utils import _dynamo
	75	+
	76	+ marker = torch.compiler.npugraph_mark_step_begin
	77	+ assert marker.__name__ == "npugraph_mark_step_begin"
	78	+ assert str(inspect.signature(marker)) == "()"
	79	+ marker()
	80	+
	81	+ from torch_npu.npu._graph_tree_state import MarkStepBox
	82	+
	83	+ assert MarkStepBox.mark_step_counter == -1
	84	+ assert not _dynamo._lazy_dynamo_setup.has_run
	85	+ assert not _dynamo._lazy_inductor_setup.has_run
	86	+ assert not any(
	87	+ name == "torch._dynamo" or name.startswith("torch._dynamo.")
	88	+ for name in sys.modules
	89	+ )
	90	+ assert not any(
	91	+ name == "torch._inductor" or name.startswith("torch._inductor.")
	92	+ for name in sys.modules
	93	+ )
	94	+
	95	+ backends = torch.compiler.list_backends(exclude_tags=None)
	96	+ assert {"npu", "npugraph_ex", "npugraphs"}.issubset(backends)
	97	+ assert _dynamo._lazy_dynamo_setup.has_run
	98	+ assert not _dynamo._lazy_inductor_setup.has_run
	99	+ assert "torch_npu._inductor" not in sys.modules
	100	+ """
	101	+ )
	102	+
	103	+ # Verify entry-point loading cannot register the NPU backend twice.
	104	+ def test_backend_entrypoint_can_import_torch_npu_without_duplicate_registration(self):
	105	+ self.run_in_subprocess(
	106	+ """
	107	+ import torch
	108	+ import torch._dynamo
	109	+ from torch._dynamo.backends import registry
	110	+
	111	+ # Model the state produced by setuptools entry-point discovery:
	112	+ # lookup_backend owns the final register_backend call, while
	113	+ # EntryPoint.load imports torch_npu and triggers its lazy setup.
	114	+ class NpuEntryPoint:
	115	+ module = "torch_npu.dynamo"
	116	+
	117	+ def load(self):
	118	+ from torch_npu.dynamo import _npu_backend_entrypoint
	119	+ return _npu_backend_entrypoint
	120	+
	121	+ registry._BACKENDS["npu"] = NpuEntryPoint()
	122	+ backend = registry.lookup_backend("npu")
	123	+
	124	+ assert backend.__name__ == "_npu_backend_entrypoint"
	125	+ assert registry._COMPILER_FNS["npu"] is backend
	126	+
	127	+ # A later setup retry must treat the loaded torch_npu entry point
	128	+ # as an already completed registration.
	129	+ from torch_npu.dynamo import _register_npu_backend
	130	+
	131	+ _register_npu_backend(backend, "npu")
	132	+ assert registry._COMPILER_FNS["npu"] is backend
	133	+ """
	134	+ )
	135	+
	136	+ # Verify backend registration is safe to retry after a partial failure.
	137	+ def test_backend_registration_retry_after_partial_failure(self):
	138	+ self.run_in_subprocess(
	139	+ """
	140	+ import torch
	141	+ import torch_npu
	142	+ import torch._dynamo
	143	+ import torch_npu.dynamo as npu_dynamo
	144	+ from torch._dynamo.backends import registry
	145	+
	146	+ # Start from an undiscovered state so the first registration is
	147	+ # committed before the simulated second-registration failure.
	148	+ for name in ("npu", "npugraph_ex"):
	149	+ registry._BACKENDS.pop(name, None)
	150	+ registry._COMPILER_FNS.pop(name, None)
	151	+
	152	+ original_register = npu_dynamo._register_npu_backend
	153	+ fail_npugraph_ex = True
	154	+
	155	+ def fail_second_registration(backend, name="npu"):
	156	+ if name == "npugraph_ex" and fail_npugraph_ex:
	157	+ raise RuntimeError("simulated npugraph_ex registration failure")
	158	+ return original_register(backend, name)
	159	+
	160	+ npu_dynamo._register_npu_backend = fail_second_registration
	161	+ try:
	162	+ try:
	163	+ npu_dynamo._register_backends()
	164	+ except RuntimeError as error:
	165	+ assert "simulated npugraph_ex" in str(error)
	166	+ else:
	167	+ raise AssertionError("the first registration should fail")
	168	+ finally:
	169	+ npu_dynamo._register_npu_backend = original_register
	170	+
	171	+ assert "npu" in registry._COMPILER_FNS
	172	+ assert "npugraph_ex" not in registry._COMPILER_FNS
	173	+
	174	+ # Retry: the completed npu registration is a no-op, while the
	175	+ # missing npugraph_ex registration is installed normally.
	176	+ npu_dynamo._register_backends()
	177	+ assert "npu" in registry._COMPILER_FNS
	178	+ assert "npugraph_ex" in registry._COMPILER_FNS
	179	+ """
	180	+ )
	181	+
	182	+ # Verify NPUGraphs rejects options kwargs for every registration order.
	183	+ def test_npugraphs_rejects_options_across_registration_order(self):
	184	+ self.run_in_subprocess(
	185	+ """
	186	+ import contextlib
	187	+ from unittest import mock
	188	+
	189	+ import torch
	190	+ import torch_npu
	191	+ from torch_npu.dynamo import _npugraphs_backend_entrypoint
	192	+ from torch_npu.utils import _dynamo, _graph_tree
	193	+
	194	+ gm = object()
	195	+ inputs = [object()]
	196	+ options = {"npu_backend": "mlir"}
	197	+
	198	+ with mock.patch.object(
	199	+ _dynamo, "_lazy_dynamo_setup", lambda: None
	200	+ ), mock.patch.object(
	201	+ _dynamo, "_lazy_inductor_setup", lambda: None
	202	+ ), mock.patch.object(
	203	+ _dynamo,
	204	+ "_NpuBackendScope",
	205	+ lambda backend: contextlib.nullcontext(),
	206	+ ), mock.patch.object(
	207	+ _graph_tree,
	208	+ "npugraphs",
	209	+ lambda model, args, **kwargs: "unexpected",
	210	+ ):
	211	+ # The public entry point accepts Dynamo's generic kwargs but
	212	+ # deliberately does not pass them to the baseline backend.
	213	+ assert _npugraphs_backend_entrypoint(
	214	+ gm, inputs, options=options
	215	+ ) == "unexpected"
	216	+
	217	+ # The registered backend itself retains its no-options API.
	218	+ try:
	219	+ _graph_tree.NpugraphsBackend()(
	220	+ gm, inputs, options=options
	221	+ )
	222	+ except TypeError as error:
	223	+ assert "unexpected keyword argument 'options'" in str(error)
	224	+ else:
	225	+ raise AssertionError("npugraphs must reject options")
	226	+ """
	227	+ )
	228	+
	229	+ def test_npugraphs_reset_protocol_across_registration_order(self):
	230	+ for order in ("cold", "hot"):
	231	+ with self.subTest(order=order):
	232	+ initialize_inductor = (
	233	+ "torch.compile(lambda x: x + 1, backend='inductor')"
	234	+ if order == "hot"
	235	+ else ""
	236	+ )
	237	+ self.run_in_subprocess(
	238	+ f"""
	239	+ import sys
	240	+ import types
	241	+ from unittest import mock
	242	+
	243	+ import torch
	244	+ import torch_npu
	245	+
	246	+ {initialize_inductor}
	247	+ torch.compile(lambda x: x + 1, backend="npugraphs")
	248	+
	249	+ from torch._dynamo.backends import registry
	250	+ from torch_npu.dynamo import _npugraphs_backend_entrypoint
	251	+
	252	+ backend = registry._COMPILER_FNS["npugraphs"]
	253	+ assert backend is _npugraphs_backend_entrypoint
	254	+ assert hasattr(backend, "reset")
	255	+
	256	+ graph_tree_module = "torch_npu.npu._graph_tree"
	257	+ assert graph_tree_module not in sys.modules
	258	+ backend.reset()
	259	+ assert graph_tree_module not in sys.modules
	260	+
	261	+ reset_calls = []
	262	+ fake_graph_tree = types.ModuleType(graph_tree_module)
	263	+ fake_graph_tree.reset_npugraph_trees = (
	264	+ lambda: reset_calls.append("reset")
	265	+ )
	266	+ with mock.patch.dict(
	267	+ sys.modules,
	268	+ {{graph_tree_module: fake_graph_tree}},
	269	+ ):
	270	+ torch.compiler.reset()
	271	+
	272	+ assert reset_calls == ["reset"]
	273	+ """
	274	+ )
	275	+
	276	+ # Verify every public Export entry initializes only Dynamo on NPU.
	277	+ def test_npu_export_public_entry_and_import_order_matrix(self):
	278	+ cases = {
	279	+ "module_export": (
	280	+ "",
	281	+ "exported = torch.export.export(Model(), (x,), strict=True)",
	282	+ ),
	283	+ "prebound_export": (
	284	+ "from torch.export import export as export_api",
	285	+ "exported = export_api(Model(), (x,), strict=True)",
	286	+ ),
	287	+ }
	288	+ if hasattr(torch.export, "export_for_training"):
	289	+ cases["prebound_export_for_training"] = (
	290	+ "from torch.export import export_for_training as export_api",
	291	+ "exported = export_api(Model(), (x,), strict=True)",
	292	+ )
	293	+ if hasattr(torch.export, "export_for_inference"):
	294	+ cases["prebound_export_for_inference"] = (
	295	+ "from torch.export import export_for_inference as export_api",
	296	+ "exported = export_api(Model(), (x,), strict=True)",
	297	+ )
	298	+ for name, (pre_import, export_call) in cases.items():
	299	+ with self.subTest(name=name):
	300	+ self.run_in_subprocess(
	301	+ f"""
	302	+ import sys
	303	+ import torch
	304	+ {pre_import}
	305	+ import torch_npu
	306	+
	307	+ stream = torch.npu.Stream()
	308	+
	309	+ class Model(torch.nn.Module):
	310	+ def forward(self, x):
	311	+ x.record_stream(stream)
	312	+ return x + 1
	313	+
	314	+ x = torch.ones(4, device="npu")
	315	+ {export_call}
	316	+ actual = exported.module()(x)
	317	+
	318	+ from torch_npu.utils import _dynamo
	319	+
	320	+ torch.testing.assert_close(actual, x + 1)
	321	+ assert _dynamo._lazy_dynamo_setup.has_run
	322	+ assert not _dynamo._lazy_inductor_setup.has_run
	323	+ assert "torch_npu._inductor" not in sys.modules
	324	+ """
	325	+ )
	326	+
	327	+ # Verify NPU-specific operations retain their Export capture semantics.
	328	+ def test_npu_export_capture_semantics_matrix(self):
	329	+ self.run_in_subprocess(
	330	+ """
	331	+ import sys
	332	+ import torch
	333	+ import torch_npu
	334	+
	335	+ x = torch.ones(4, device="npu")
	336	+ stream = torch.npu.Stream()
	337	+ event = torch.npu.Event()
	338	+
	339	+ class StreamAndEvent(torch.nn.Module):
	340	+ def forward(self, value):
	341	+ event.record()
	342	+ with torch.npu.stream(stream):
	343	+ event.wait(stream)
	344	+ result = value + 1
	345	+ return result
	346	+
	347	+ class Autocast(torch.nn.Module):
	348	+ def forward(self, value):
	349	+ with torch.npu.amp.autocast(dtype=torch.float16):
	350	+ return value * value
	351	+
	352	+ class CurrentDevice(torch.nn.Module):
	353	+ def forward(self, value):
	354	+ return value + torch.npu.current_device()
	355	+
	356	+ class DeviceProperties(torch.nn.Module):
	357	+ def forward(self, value):
	358	+ properties = torch.npu.get_device_properties(
	359	+ torch.npu.current_device()
	360	+ )
	361	+ return value + 1 if properties.total_memory > 0 else value - 1
	362	+
	363	+ class IsAvailable(torch.nn.Module):
	364	+ def forward(self, value):
	365	+ return value + 1 if torch.npu.is_available() else value - 1
	366	+
	367	+ models = (
	368	+ StreamAndEvent,
	369	+ Autocast,
	370	+ CurrentDevice,
	371	+ DeviceProperties,
	372	+ IsAvailable,
	373	+ )
	374	+ for model_type in models:
	375	+ model = model_type()
	376	+ expected = model(x)
	377	+ exported = torch.export.export(model, (x,))
	378	+ actual = exported.module()(x)
	379	+ torch.testing.assert_close(actual, expected)
	380	+
	381	+ from torch_npu.utils import _dynamo
	382	+
	383	+ assert _dynamo._lazy_dynamo_setup.has_run
	384	+ assert not _dynamo._lazy_inductor_setup.has_run
	385	+ assert "torch_npu._inductor" not in sys.modules
	386	+ """
	387	+ )
	388	+
	389	+ # Verify rejected Inductor arguments do not initialize or pollute NPU state.
	390	+ def test_invalid_inductor_arguments_fail_without_npu_initialization(self):
	391	+ cases = {
	392	+ "invalid_mode": (
	393	+ 'torch.compile(lambda x: x + 1, mode="invalid-mode")',
	394	+ "Unrecognized mode=invalid-mode",
	395	+ ),
	396	+ "invalid_option": (
	397	+ "torch.compile(lambda x: x + 1, "
	398	+ 'options={"invalid.option": True})',
	399	+ "Unexpected optimization option invalid.option",
	400	+ ),
	401	+ "invalid_npu_option_type": (
	402	+ "torch.compile(lambda x: x + 1, "
	403	+ 'options={"npu_backend": 1})',
	404	+ "Unexpected type of attr npu_backend",
	405	+ ),
	406	+ "invalid_hyphenated_npu_option_type": (
	407	+ "torch.compile(lambda x: x + 1, "
	408	+ 'options={"npu-backend": 1})',
	409	+ "Unexpected type of attr npu-backend",
	410	+ ),
	411	+ }
	412	+ for name, (compile_call, expected_error) in cases.items():
	413	+ with self.subTest(name=name):
	414	+ self.run_in_subprocess(
	415	+ f"""
	416	+ import os
	417	+ import sys
	418	+
	419	+ import torch
	420	+ import torch_npu
	421	+ from torch_npu.utils import _dynamo
	422	+
	423	+ env_name = "TORCHINDUCTOR_NPU_BACKEND"
	424	+ original_env = os.environ.get(env_name)
	425	+ try:
	426	+ {compile_call}
	427	+ except RuntimeError as error:
	428	+ assert {expected_error!r} in str(error), str(error)
	429	+ else:
	430	+ raise AssertionError("invalid compile arguments must fail")
	431	+
	432	+ assert not _dynamo._lazy_inductor_setup.has_run
	433	+ assert not _dynamo.is_inductor_npu_initialized()
	434	+ assert "torch_npu._inductor" not in sys.modules
	435	+ assert os.environ.get(env_name) == original_env
	436	+ """
	437	+ )
	438	+
	439	+ # Verify shape handling is installed after selecting the requested NPU backend.
	440	+ def test_shape_handling_initializes_after_backend_selection(self):
	441	+ self.run_in_subprocess(
	442	+ """
	443	+ import types
	444	+ from unittest import mock
	445	+
	446	+ import torch
	447	+ import torch_npu
	448	+ from torch_npu.utils import _dynamo
	449	+
	450	+ def fake_setup(actual_options):
	451	+ events.append(("setup", dict(actual_options)))
	452	+ return actual_options["npu_backend"]
	453	+
	454	+ options = {
	455	+ "npu_backend": "mlir",
	456	+ "enable_shape_handling": True,
	457	+ }
	458	+ events = []
	459	+
	460	+ def scope_register():
	461	+ events.append(("scope_register", None))
	462	+
	463	+ fake_inductor = types.SimpleNamespace(
	464	+ patch_shape_handling=lambda: events.append(
	465	+ ("shape_handling", None)
	466	+ )
	467	+ )
	468	+ with mock.patch.object(
	469	+ _dynamo, "_setup_inductor_for_compile", fake_setup
	470	+ ), mock.patch.object(
	471	+ _dynamo, "register_inductor_npu", scope_register
	472	+ ), mock.patch.object(
	473	+ torch_npu, "_inductor", fake_inductor, create=True
	474	+ ):
	475	+ wrapper = torch._TorchCompileInductorWrapper(None, options, None)
	476	+
	477	+ normalized_options = {
	478	+ "npu_backend": "mlir",
	479	+ "enable_shape_handling": True,
	480	+ }
	481	+ assert wrapper.config["npu_backend"] == "mlir"
	482	+ assert wrapper.config["enable_shape_handling"] is True
	483	+ assert events == [
	484	+ ("setup", normalized_options),
	485	+ ("shape_handling", None),
	486	+ ("scope_register", None),
	487	+ ], events
	488	+ """
	489	+ )
	490	+
	491	+ # Verify non-Inductor compile backends do not initialize Inductor.
	492	+ def test_non_inductor_compile_backend_matrix(self):
	493	+ cases = {
	494	+ "eager": (
	495	+ "",
	496	+ 'compiled = torch.compile(Model(), backend="eager", fullgraph=True)',
	497	+ ),
	498	+ "custom": (
	499	+ "custom_backend = lambda graph_module, example_inputs: "
	500	+ "graph_module.forward",
	501	+ "compiled = torch.compile(Model(), backend=custom_backend, fullgraph=True)",
	502	+ ),
	503	+ "npu": (
	504	+ "",
	505	+ 'compiled = torch.compile(Model(), backend="npu", fullgraph=True)',
	506	+ ),
	507	+ }
	508	+ for name, (backend_definition, compile_call) in cases.items():
	509	+ with self.subTest(name=name):
	510	+ allow_missing_torchair = name == "npu"
	511	+ self.run_in_subprocess(
	512	+ f"""
	513	+ import sys
	514	+ import torch
	515	+ import torch_npu
	516	+
	517	+ class Model(torch.nn.Module):
	518	+ def forward(self, x):
	519	+ return torch.sin(x) + 1
	520	+
	521	+ {backend_definition}
	522	+ x = torch.randn(8, device="npu")
	523	+ try:
	524	+ {compile_call}
	525	+ except AssertionError as error:
	526	+ assert {allow_missing_torchair!r}
	527	+ assert "Could not find module torchair" in str(error)
	528	+ else:
	529	+ torch.testing.assert_close(compiled(x), Model()(x))
	530	+
	531	+ from torch_npu.utils import _dynamo
	532	+
	533	+ assert _dynamo._lazy_dynamo_setup.has_run
	534	+ assert not _dynamo._lazy_inductor_setup.has_run
	535	+ assert "torch_npu._inductor" not in sys.modules
	536	+ """
	537	+ )
	538	+
	539	+ # Verify ONNX Dynamo Export initializes only the NPU Dynamo integration.
	540	+ @unittest.skipUnless(importlib.util.find_spec("onnxscript"), "requires onnxscript")
	541	+ def test_npu_onnx_dynamo_export_initialization_chain(self):
	542	+ cases = {
	543	+ "module_export": (
	544	+ "",
	545	+ "result = torch.onnx.export(Model(), (x,), dynamo=True)",
	546	+ ),
	547	+ "prebound_export": (
	548	+ "from torch.onnx import export as onnx_export",
	549	+ "result = onnx_export(Model(), (x,), dynamo=True)",
	550	+ ),
	551	+ }
	552	+ if hasattr(torch.onnx, "dynamo_export"):
	553	+ cases["prebound_legacy_dynamo_export"] = (
	554	+ "from torch.onnx import dynamo_export as onnx_export",
	555	+ "result = onnx_export(Model(), x)",
	556	+ )
	557	+ for name, (pre_import, export_call) in cases.items():
	558	+ with self.subTest(name=name):
	559	+ self.run_in_subprocess(
	560	+ f"""
	561	+ import sys
	562	+ import torch
	563	+ {pre_import}
	564	+ import torch_npu
	565	+
	566	+ # ONNXScript 0.4.0 cannot version-convert models containing
	567	+ # functions. Conversion runs after Dynamo capture, which is
	568	+ # the boundary this test verifies, so isolate that unrelated
	569	+ # compatibility issue without weakening the import assertions.
	570	+ from torch.onnx._internal._lazy_import import onnxscript_apis
	571	+ onnxscript_apis.convert_version = lambda model, target: model
	572	+
	573	+ class Model(torch.nn.Module):
	574	+ def forward(self, x):
	575	+ return torch.sin(x) + 1
	576	+
	577	+ x = torch.randn(8, device="npu")
	578	+ {export_call}
	579	+
	580	+ from torch_npu.utils import _dynamo
	581	+
	582	+ assert result is not None
	583	+ assert _dynamo._lazy_dynamo_setup.has_run
	584	+ assert not _dynamo._lazy_inductor_setup.has_run
	585	+ assert "torch_npu._inductor" not in sys.modules
	586	+ """
	587	+ )
	588	+
	589	+ # Verify pre-imported FSDP receives all NPU patches.
	590	+ def test_fsdp_patch_when_imported_before_torch_npu(self):
	591	+ self.run_in_subprocess(
	592	+ """
	593	+ import torch.distributed.fsdp
	594	+ import torch_npu
	595	+ from torch.distributed.fsdp import sharded_grad_scaler
	596	+ from torch.distributed.fsdp._fully_shard._fsdp_param_group import (
	597	+ FSDPParamGroup,
	598	+ )
	599	+ from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler
	600	+
	601	+ assert sharded_grad_scaler.ShardedGradScaler is _ShardedGradScaler
	602	+ """
	603	+ )
	604	+
	605	+ # Verify a real NPU Inductor compile initializes the full stack.
	606	+ # The triton kernel compilation itself may fail due to environment-level
	607	+ # compiler/CANN version constraints (e.g. C++20 requirement on PyTorch
	608	+ # 2.13+, or missing aclmdlRICondHandle in older CANN). The purpose of
	609	+ # this test is to verify the lazy initialization chain runs, not to
	610	+ # verify triton compilation succeeds, so we tolerate compilation errors
	611	+ # as long as the has_run flags are set.
	612	+ @unittest.skipUnless(importlib.util.find_spec("triton"), "requires triton-ascend")
	613	+ def test_npu_inductor_initialization_chain(self):
	614	+ self.run_in_subprocess(
	615	+ """
	616	+ import torch
	617	+ import torch_npu
	618	+
	619	+ def fn(x):
	620	+ return torch.sin(x) + 1
	621	+
	622	+ x = torch.randn(8, device="npu")
	623	+ try:
	624	+ actual = torch.compile(fn, backend="inductor", fullgraph=True)(x)
	625	+ torch.testing.assert_close(actual, fn(x))
	626	+ except Exception:
	627	+ # Triton kernel compilation may fail due to environment-level
	628	+ # compiler/CANN constraints. The lazy init chain runs before
	629	+ # kernel compilation, so we verify it below regardless.
	630	+ pass
	631	+
	632	+ from torch._dynamo.device_interface import get_interface_for_device
	633	+ from torch_npu.utils import _dynamo
	634	+
	635	+ assert _dynamo._lazy_dynamo_setup.has_run
	636	+ assert _dynamo._lazy_inductor_setup.has_run
	637	+ assert get_interface_for_device("npu").device_count() > 0
	638	+ """,
	639	+ timeout=180,
	640	+ )
	641	+
	642	+ # Verify a real NPUGraphs compile initializes the full stack.
	643	+ @unittest.skipUnless(importlib.util.find_spec("triton"), "requires triton-ascend")
	644	+ def test_npu_npugraphs_initialization_chain(self):
	645	+ self.run_in_subprocess(
	646	+ """
	647	+ import torch
	648	+ import torch_npu
	649	+
	650	+ def fn(x):
	651	+ return torch.sin(x) + 1
	652	+
	653	+ x = torch.randn(8, device="npu")
	654	+ actual = torch.compile(fn, backend="npugraphs", fullgraph=True)(x)
	655	+
	656	+ from torch._dynamo.device_interface import get_interface_for_device
	657	+ from torch_npu.utils import _dynamo
	658	+
	659	+ torch.testing.assert_close(actual, fn(x))
	660	+ assert _dynamo._lazy_dynamo_setup.has_run
	661	+ assert _dynamo._lazy_inductor_setup.has_run
	662	+ assert get_interface_for_device("npu").device_count() > 0
	663	+ """
	664	+ )
	665	+
	666	+ # Verify lazy setup completes before compile backend lookup.
	667	+ def test_compile_triggers_setup_before_backend_lookup(self):
	668	+ self.run_in_subprocess(
	669	+ """
	670	+ import torch
	671	+ import torch_npu
	672	+ from torch_npu.utils import _dynamo
	673	+
	674	+ calls = []
	675	+
	676	+ @_dynamo.run_once
	677	+ def fake_setup():
	678	+ calls.append("setup")
	679	+
	680	+ _dynamo._lazy_dynamo_setup = fake_setup
	681	+ from torch._dynamo.backends import registry
	682	+ assert fake_setup.has_run
	683	+
	684	+ compiled = torch.compile(lambda x: x + 1, backend="eager")
	685	+ assert compiled(torch.tensor(1)).item() == 2
	686	+ assert calls == ["setup"]
	687	+ """
	688	+ )
	689	+
	690	+ # Verify the trigger works when Dynamo was imported first.
	691	+ def test_trigger_after_dynamo_was_preimported(self):
	692	+ self.run_in_subprocess(
	693	+ """
	694	+ import sys
	695	+ import torch
	696	+ import torch._dynamo
	697	+ import torch_npu
	698	+ from torch_npu.utils import _dynamo
	699	+
	700	+ assert _dynamo._lazy_dynamo_setup.has_run
	701	+
	702	+ compiled = torch.compile(lambda x: x + 1, backend="eager", fullgraph=True)
	703	+ assert compiled(torch.tensor(1)).item() == 2
	704	+ assert _dynamo._lazy_dynamo_setup.has_run
	705	+ assert not _dynamo._lazy_inductor_setup.has_run
	706	+ assert "torch_npu._inductor" not in sys.modules
	707	+ """
	708	+ )
	709	+
	710	+ # Verify the moved helper preserves the torch_npu v2.9 device predicates.
	711	+ def test_has_triton_preserves_target_device_semantics(self):
	712	+ self.run_in_subprocess(
	713	+ """
	714	+ import sys
	715	+ import types
	716	+ from unittest import mock
	717	+
	718	+ import torch
	719	+ import torch_npu
	720	+ from torch._dynamo import device_interface
	721	+ from torch_npu.utils import _dynamo
	722	+
	723	+ def make_interface(available=False):
	724	+ class Worker:
	725	+ @staticmethod
	726	+ def get_device_properties():
	727	+ raise AssertionError("v2.9 NPU override must not query CUDA props")
	728	+
	729	+ class Interface:
	730	+ @staticmethod
	731	+ def is_available():
	732	+ return available
	733	+
	734	+ Interface.Worker = Worker
	735	+ return Interface
	736	+
	737	+ def run_case(available, *, package=True):
	738	+ interfaces = {
	739	+ device: make_interface(
	740	+ available=device in available,
	741	+ )
	742	+ for device in ("cuda", "xpu", "cpu", "npu")
	743	+ }
	744	+ _dynamo.has_triton.cache_clear()
	745	+ with mock.patch.object(
	746	+ torch.utils._triton,
	747	+ "has_triton_package",
	748	+ return_value=package,
	749	+ ), mock.patch.object(
	750	+ device_interface,
	751	+ "get_interface_for_device",
	752	+ side_effect=interfaces.__getitem__,
	753	+ ), mock.patch.object(
	754	+ _dynamo,
	755	+ "_dynamo_register_interface_for_device",
	756	+ ) as register:
	757	+ result = _dynamo.has_triton()
	758	+ if package:
	759	+ register.assert_called_once_with()
	760	+ else:
	761	+ register.assert_not_called()
	762	+ return result
	763	+
	764	+ assert not run_case({}, package=False)
	765	+ assert run_case({"cuda"})
	766	+ assert run_case({"xpu"})
	767	+ triton = types.ModuleType("triton")
	768	+ triton_backends = types.ModuleType("triton.backends")
	769	+ triton_backends.backends = {"cpu": object()}
	770	+ triton.backends = triton_backends
	771	+ with mock.patch.dict(
	772	+ sys.modules,
	773	+ {"triton": triton, "triton.backends": triton_backends},
	774	+ ):
	775	+ assert run_case({"cpu"})
	776	+ triton_backends.backends = {}
	777	+ assert not run_case({"cpu"})
	778	+ assert run_case({"npu"})
	779	+ """
	780	+ )
	781	+
	782	+ # Verify all legacy Dynamo patches remain installed exactly once.
	783	+ def test_dynamo_patch_inventory_is_preserved(self):
	784	+ self.run_in_subprocess(
	785	+ """
	786	+ import torch
	787	+ import torch_npu
	788	+
	789	+ # Importing the Dynamo parent package is the lazy setup boundary.
	790	+ import torch._dynamo
	791	+
	792	+ from torch._dynamo.device_interface import get_interface_for_device
	793	+ from torch._dynamo.variables.builtin import BuiltinVariable
	794	+ from torch._dynamo.variables.streams import EventVariable
	795	+ from torch._dynamo.variables.functions import SkipFunctionVariable
	796	+ from torch._dynamo.variables.tensor import TensorVariable
	797	+ from torch._dynamo.variables.torch import constant_fold_functions
	798	+ from torch._dynamo.variables.user_defined import UserDefinedClassVariable
	799	+ from torch._dynamo.utils import common_constant_types
	800	+ from torch_npu.dynamo.trace_rule import (
	801	+ skip_functions_npu,
	802	+ torch_c_binding_in_graph_functions_npu,
	803	+ torch_non_c_binding_in_graph_functions_npu,
	804	+ )
	805	+ from torch_npu.utils import _dynamo
	806	+
	807	+ assert _dynamo._lazy_dynamo_setup.has_run
	808	+ assert get_interface_for_device("npu").device_count() > 0
	809	+
	810	+ # VariableTracker and context-manager patches formerly installed
	811	+ # eagerly by add_dynamo_methods().
	812	+ assert SkipFunctionVariable.__new__.__module__ == "torch_npu.utils._dynamo"
	813	+ assert TensorVariable.call_method.__module__ == "torch_npu.utils._dynamo"
	814	+ assert UserDefinedClassVariable.__new__.__module__ == "torch_npu.utils._dynamo"
	815	+ in_graph_classes = UserDefinedClassVariable._in_graph_classes()
	816	+ assert torch.npu.Event in in_graph_classes
	817	+ assert torch.npu.Stream in in_graph_classes
	818	+ assert BuiltinVariable.call_id.__module__ == "torch_npu.utils._dynamo"
	819	+ assert EventVariable.python_type.__module__ == "torch_npu.utils._dynamo"
	820	+ assert torch._dynamo.optimize.__module__ == "torch_npu.utils._dynamo"
	821	+
	822	+ from torch._functorch._aot_autograd.utils import supports_graphsafe_rng
	823	+ assert supports_graphsafe_rng(torch.device("npu"))
	824	+
	825	+ # Backend and trace-rule registrations formerly performed by
	826	+ # registry_manager._register_dynamo(). Count the maps as well as
	827	+ # checking membership so repeated lazy triggers cannot hide a
	828	+ # duplicate installation.
	829	+ assert {"npu", "npugraph_ex"}.issubset(
	830	+ torch._dynamo.list_backends(exclude_tags=None)
	831	+ )
	832	+ maps = torch._dynamo.trace_rules.torch_name_rule_map
	833	+ assert maps.count(torch_non_c_binding_in_graph_functions_npu) == 1
	834	+ assert maps.count(torch_c_binding_in_graph_functions_npu) == 1
	835	+ assert maps.count(skip_functions_npu) == 1
	836	+ assert constant_fold_functions[torch.npu.current_device]
	837	+ assert constant_fold_functions[torch.npu.get_device_properties]
	838	+ assert constant_fold_functions[torch.npu.is_available]
	839	+ assert torch_npu._C._NPUDeviceProperties in common_constant_types
	840	+ """
	841	+ )
	842	+
	843	+ # Verify all legacy Inductor patches remain installed.
	844	+ def test_inductor_patch_inventory_is_preserved(self):
	845	+ self.run_in_subprocess(
	846	+ """
	847	+ import torch
	848	+ import torch_npu
	849	+
	850	+ # RNG/decomposition patches remain installed at import time.
	851	+ from torch_npu.utils import _inductor
	852	+
	853	+ assert (
	854	+ torch._decomp.decompositions._max_unpoolnd
	855	+ is _inductor._max_unpoolnd_patch
	856	+ )
	857	+ assert torch._prims.rng_prims.philox_rand_offset.__module__ == (
	858	+ "torch_npu.utils._inductor"
	859	+ )
	860	+ assert torch._prims.rng_prims.register_philox_rand.__module__ == (
	861	+ "torch_npu.utils._inductor"
	862	+ )
	863	+ assert torch._prims.rng_prims.get_device.__module__ == (
	864	+ "torch_npu.utils._inductor"
	865	+ )
	866	+
	867	+ # Exercise the new full-Inductor setup boundary without relying on
	868	+ # test ordering or on a prior torch.compile invocation.
	869	+ import torch._dynamo
	870	+ from torch_npu.utils import _dynamo
	871	+
	872	+ _dynamo._lazy_inductor_setup()
	873	+
	874	+ import torch._inductor.compile_fx as compile_fx
	875	+ import torch._inductor.cudagraph_trees as cudagraph_trees
	876	+ import torch._inductor.cudagraph_utils as cudagraph_utils
	877	+ import torch._inductor.scheduler as scheduler
	878	+ from torch._inductor.codegen.common import get_device_op_overrides
	879	+ from torch._inductor.codecache import CacheBase
	880	+ from torch._inductor.graph import GraphLowering
	881	+ from torch._inductor.utils import GPU_TYPES
	882	+ from torch_npu.utils import _graph_tree
	883	+
	884	+ assert _dynamo._lazy_inductor_setup.has_run
	885	+ assert "npu" in GPU_TYPES
	886	+ assert get_device_op_overrides("npu").__class__.__module__.startswith(
	887	+ "torch_npu._inductor"
	888	+ )
	889	+ assert torch.utils._triton.has_triton.__module__ == (
	890	+ "torch_npu.utils._dynamo"
	891	+ )
	892	+ assert torch.utils._triton._device_supports_tma.__module__ == (
	893	+ "torch_npu._inductor.utils"
	894	+ )
	895	+ assert compile_fx.has_triton is torch.utils._triton.has_triton
	896	+ assert scheduler.has_triton is torch.utils._triton.has_triton
	897	+ assert GraphLowering.codegen_with_cpp_wrapper.__module__ == (
	898	+ "torch_npu._inductor.graph"
	899	+ )
	900	+ assert CacheBase.get_system.__module__ == (
	901	+ "torch_npu._inductor.codegen.common"
	902	+ )
	903	+
	904	+ # NPUGraph integrations were formerly applied eagerly alongside
	905	+ # the Inductor patches.
	906	+ assert compile_fx.cudagraphify is _graph_tree.npugraphify
	907	+ assert (
	908	+ cudagraph_utils.check_multiple_devices_or_any_cpu_nodes
	909	+ is _graph_tree.check_multiple_devices_or_any_cpu_nodes
	910	+ )
	911	+ assert cudagraph_trees.get_manager.__module__ == (
	912	+ "torch_npu.utils._graph_tree"
	913	+ )
	914	+ assert torch.compiler.npugraph_mark_step_begin is (
	915	+ _graph_tree.npugraph_mark_step_begin
	916	+ )
	917	+
	918	+ config = torch._inductor.config
	919	+ assert config.npu_backend == "default"
	920	+ assert config.enable_shape_handling is False
	921	+ assert config.shape_handling_configs == []
	922	+ assert config.shape_handling_dict is None
	923	+ """
	924	+ )
	925	+
	926	+
	927	+if __name__ == "__main__":
	928	+ unittest.main()

Mtest/npu/test_stream.py+5-0

	@@ -23,6 +23,11 @@ class TestNpuStream(TestCase):
23	def test_get_current_stream_interface(self):	23	def test_get_current_stream_interface(self):
24	from torch_npu._C import _npu_getCurrentRawStream, _npu_getCurrentRawStreamNoWait	24	from torch_npu._C import _npu_getCurrentRawStream, _npu_getCurrentRawStreamNoWait
25	from torch._dynamo.device_interface import get_interface_for_device	25	from torch._dynamo.device_interface import get_interface_for_device
		26	+ from torch_npu.utils._dynamo import _dynamo_register_interface_for_device
		27	+
		28	+ # device_interface is an internal Dynamo module. Initialize its NPU
		29	+ # registration explicitly instead of relying on import torch_npu.
		30	+ _dynamo_register_interface_for_device()
26		31
27	device_number = torch.npu.device_count()	32	device_number = torch.npu.device_count()
28	for i in range(device_number):	33	for i in range(device_number):

Mtest/test_torch_npu_init.py+76-15

	@@ -33,14 +33,15 @@ EXPECTED_LOADED_MODULES = [
33	"torch_npu.op_plugin.meta._meta_registrations",	33	"torch_npu.op_plugin.meta._meta_registrations",
34	"torch_npu.asd.checksum",	34	"torch_npu.asd.checksum",
35	"torch_npu.utils._dynamo",	35	"torch_npu.utils._dynamo",
36	- "torch_npu.utils._inductor",
37	"torch_npu.utils.custom_ops",	36	"torch_npu.utils.custom_ops",
38	"torch_npu.utils.patch_getenv",	37	"torch_npu.utils.patch_getenv",
39	"torch_npu.utils.syncbatchnorm",	38	"torch_npu.utils.syncbatchnorm",
		39	+ "torch_npu.utils._inductor",
40	]	40	]
41		41
42	EXPECTED_NOT_LOADED_MODULES = [	42	EXPECTED_NOT_LOADED_MODULES = [
43	"torch_npu._C._afd",	43	"torch_npu._C._afd",
		44	+ "torch_npu._inductor",
44	]	45	]
45		46
46	EXPECTED_TOP_LEVEL_ATTRS = [	47	EXPECTED_TOP_LEVEL_ATTRS = [
	@@ -255,20 +256,6 @@ class TestTorchNpuBootstrap(TestCase):
255	import torch.distributed as dist	256	import torch.distributed as dist
256	import torch.distributed.rpc as rpc	257	import torch.distributed.rpc as rpc
257	import torch.distributed.tensor # noqa: F401	258	import torch.distributed.tensor # noqa: F401
258	- from torch._dynamo.device_interface import get_interface_for_device
259	- from torch._dynamo.backends.registry import _BACKENDS
260	- from torch._inductor.codegen.common import device_op_overrides_dict
261	-
262	- iface = get_interface_for_device("npu")
263	- assert iface is not None
264	-
265	- assert "npu" in _BACKENDS, "npu dynamo backend is not registered"
266	- assert "npugraph_ex" in _BACKENDS, (
267	- "npugraph_ex dynamo backend is not registered"
268	- )
269	-
270	- assert "npu" in device_op_overrides_dict
271	- assert device_op_overrides_dict.get("npu") is not None
272		259
273	assert "hccl" in dist.Backend.backend_list	260	assert "hccl" in dist.Backend.backend_list
274	assert "lccl" in dist.Backend.backend_list	261	assert "lccl" in dist.Backend.backend_list
	@@ -549,5 +536,79 @@ class TestTorchNpuBootstrap(TestCase):
549	"""	536	"""
550	)	537	)
551		538
		539	+ def test_13_dtensor_strategies_are_registered_without_compiler_imports(self):
		540	+ self._run_python(
		541	+ """
		542	+ import sys
		543	+ import torch
		544	+ import torch_npu
		545	+
		546	+ assert "torch_npu.distributed.tensor" in sys.modules
		547	+ assert "torch._dynamo" not in sys.modules
		548	+ assert "torch._inductor" not in sys.modules
		549	+
		550	+ from torch.distributed.tensor import DTensor
		551	+
		552	+ strategy_funcs = DTensor._op_dispatcher.sharding_propagator.op_strategy_funcs
		553	+ assert torch.ops.npu.npu_rms_norm.default in strategy_funcs
		554	+ assert torch.ops.npu.npu_fusion_attention.default in strategy_funcs
		555	+
		556	+ import os
		557	+ import tempfile
		558	+ import torch.distributed as dist
		559	+ from torch.distributed.device_mesh import DeviceMesh
		560	+ from torch.distributed.tensor import Replicate
		561	+ from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
		562	+ from torch.distributed.tensor._op_schema import OpSchema, OpSpec, OpStrategy
		563	+
		564	+ fd, path = tempfile.mkstemp()
		565	+ os.close(fd)
		566	+ os.unlink(path)
		567	+ dist.init_process_group(
		568	+ "gloo", init_method=f"file://{path}", rank=0, world_size=1
		569	+ )
		570	+ try:
		571	+ mesh = DeviceMesh("cpu", [0])
		572	+ tensor_meta = TensorMeta(
		573	+ torch.Size([2, 4, 8]), (32, 8, 1), torch.float32
		574	+ )
		575	+ spec = DTensorSpec(mesh, (Replicate(),), tensor_meta=tensor_meta)
		576	+ strategy = OpStrategy([OpSpec(output_specs=spec)])
		577	+ op = torch.ops.npu.npu_rotary_mul.default
		578	+ propagator = DTensor._op_dispatcher.sharding_propagator
		579	+ op_schema = OpSchema(
		580	+ op,
		581	+ (strategy, strategy, strategy, "half"),
		582	+ {},
		583	+ propagator.op_to_schema_info[op],
		584	+ )
		585	+ result = propagator.op_strategy_funcs[op](op_schema)
		586	+ assert len(result.strategies) == 3
		587	+ finally:
		588	+ dist.destroy_process_group()
		589	+ if os.path.exists(path):
		590	+ os.remove(path)
		591	+
		592	+ from torch_npu.distributed.tensor import experimental
		593	+ assert callable(experimental.context_parallel)
		594	+ """
		595	+ )
		596	+
		597	+ def test_14_direct_npu_fsdp_import(self):
		598	+ self._run_python(
		599	+ """
		600	+ import torch_npu.distributed.fsdp as npu_fsdp
		601	+ from torch.distributed.fsdp import sharded_grad_scaler
		602	+ from torch.distributed.fsdp._fully_shard._fsdp_param_group import (
		603	+ FSDPParamGroup,
		604	+ )
		605	+ from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler
		606	+
		607	+ assert callable(npu_fsdp.fully_shard)
		608	+ assert sharded_grad_scaler.ShardedGradScaler is _ShardedGradScaler
		609	+ """
		610	+ )
		611	+
		612	+
552	if __name__ == "__main__":	613	if __name__ == "__main__":
553	run_tests()	614	run_tests()

Dtest/utils/test_inductor.py+0-33

	@@ -1,33 +0,0 @@
1	-from torch_npu.testing.testcase import TestCase, run_tests
2	-from torch_npu.utils._inductor import NPUDeviceOpOverrides
3	-
4	-
5	-class TestInductor():
6	-
7	- def test_device_guard(self):
8	- overrides = NPUDeviceOpOverrides()
9	- result = overrides.device_guard(1)
10	- expected = "torch_npu.npu._DeviceGuard(1)"
11	- self.assertEqual(result, expected)
12	-
13	- def test_synchronize(self):
14	- overrides = NPUDeviceOpOverrides()
15	- result = overrides.synchronize()
16	- expected = "torch_npu.npu.synchronize()"
17	- self.assertEqual(result, expected)
18	-
19	- def test_set_device(self):
20	- overrides = NPUDeviceOpOverrides()
21	- result = overrides.set_device(0)
22	- expected = "torch_npu.npu.set_device(0)"
23	- self.assertEqual(result, expected)
24	-
25	- def test_import_get_raw_stream_as(self):
26	- overrides = NPUDeviceOpOverrides()
27	- result = overrides.import_get_raw_stream_as("test_name")
28	- expected = "from torch._C import _npu_getCurrentRawStream as test_name"
29	- self.assertEqual(result, expected)
30	-
31	-
32	-if __name__ == "__main__":
33	- run_tests()

Mtorch_npu/_inductor/kernel/bmm.py+3-2

	@@ -17,6 +17,7 @@ from torch._inductor.kernel.mm_common import (
17	_is_static_problem,	17	_is_static_problem,
18	mm_args,	18	mm_args,
19	)	19	)
		20	+from torch._inductor.kernel import bmm as inductor_bmm
20		21
21	from .mm import is_contiguous_striding	22	from .mm import is_contiguous_striding
22	from ..utils import use_catlass_template	23	from ..utils import use_catlass_template
	@@ -24,8 +25,8 @@ from ..utils import use_catlass_template
24	log = logging.getLogger("torch._inductor")	25	log = logging.getLogger("torch._inductor")
25	aten = torch.ops.aten	26	aten = torch.ops.aten
26		27
27	-aten_bmm = torch._inductor.kernel.bmm.aten_bmm	28	+aten_bmm = inductor_bmm.aten_bmm
28	-aten_baddbmm = torch._inductor.kernel.bmm.aten_baddbmm	29	+aten_baddbmm = inductor_bmm.aten_baddbmm
29		30
30	def is_batch_stride_largest_or_zero(mat1, mat2, layout) -> bool:	31	def is_batch_stride_largest_or_zero(mat1, mat2, layout) -> bool:
31	"""	32	"""

Mtorch_npu/_inductor/utils.py+2-1

	@@ -28,10 +28,11 @@ def resolve_npu_device_index(device_idx=None) -> int:
28		28
29		29
30	def patch_has_triton():	30	def patch_has_triton():
		31	+ from torch._inductor import compile_fx
31	from torch_npu.utils._dynamo import has_triton	32	from torch_npu.utils._dynamo import has_triton
32		33
33	torch._inductor.scheduler.has_triton = has_triton	34	torch._inductor.scheduler.has_triton = has_triton
34	- torch._inductor.compile_fx.has_triton = has_triton	35	+ compile_fx.has_triton = has_triton
35		36
36		37
37	def patch_device_supports_tma():	38	def patch_device_supports_tma():

Mtorch_npu/_init/patches/dynamo_patches.py+0-7

	@@ -6,10 +6,3 @@ def apply_dynamo_methods_patch():
6	from torch_npu.utils._dynamo import add_dynamo_methods	6	from torch_npu.utils._dynamo import add_dynamo_methods
7		7
8	add_dynamo_methods()	8	add_dynamo_methods()
9	-
10	-
11	-@PatchManager.register_patch("dynamo")
12	-def apply_npugraph_tree_patch():
13	- from torch_npu.utils._graph_tree import _apply_npugraph_tree_methods
14	-
15	- _apply_npugraph_tree_methods()

Mtorch_npu/_init/registry/registry_manager.py+2-21

	@@ -46,24 +46,6 @@ def _register_distributed():
46	# init and register distributed backend	46	# init and register distributed backend
47	register_distributed_backend_for_npu()	47	register_distributed_backend_for_npu()
48		48
49	-def _register_dynamo():
50	- """
51	- Register Dynamo integration:
52	- - Dynamo backend
53	- - Dynamo device interface
54	- - NPU trace rules for Dynamo
55	- """
56	- from torch_npu._init.registry.dynamo import (
57	- register_dynamo_backends,
58	- register_dynamo_trace_rules,
59	- )
60	-
61	- register_dynamo_backends()
62	-
63	- # Do not repeat this call for register_dynamo_trace_rules appends rules into
64	- # Dynamo's global rules maps.
65	- register_dynamo_trace_rules()
66	-
67	def _register_rpc():	49	def _register_rpc():
68	"""	50	"""
69	Register and init RPC NPU backend.	51	Register and init RPC NPU backend.
	@@ -85,8 +67,8 @@ def _register_components():
85		67
86	Order matters:	68	Order matters:
87	1. NPU backend is the base capability.	69	1. NPU backend is the base capability.
88	- 2. Distributed and Dynamo depend on NPU backend / _C children.	70	+ 2. Distributed depends on NPU backend / _C children.
89	- 3. RPC, dtensor and inductor are Python-side framework integrations.	71	+ 3. RPC is a Python-side framework integration.
90	4. DefaultDeviceType is set after NPU backend is registered.	72	4. DefaultDeviceType is set after NPU backend is registered.
91	"""	73	"""
92	if not hasattr(torch_npu, "_C"):	74	if not hasattr(torch_npu, "_C"):
	@@ -96,6 +78,5 @@ def _register_components():
96		78
97	_register_npu_backend()	79	_register_npu_backend()
98	_register_distributed()	80	_register_distributed()
99	- _register_dynamo()
100	_register_rpc()	81	_register_rpc()
101	_register_default_gradient_device_type()	82	_register_default_gradient_device_type()

Mtorch_npu/contrib/transfer_to_npu.py+15-2

	@@ -4,11 +4,12 @@ import json
4	import collections	4	import collections
5	import importlib.metadata	5	import importlib.metadata
6	import logging as logger	6	import logging as logger
		7	+import functools # noqa: F401
7	from functools import wraps	8	from functools import wraps
8	from typing import Callable, cast, Optional	9	from typing import Callable, cast, Optional
9	import torch	10	import torch
10	from torch.utils._device import _device_constructors	11	from torch.utils._device import _device_constructors
11	-import torch.utils._triton # ensure module is loaded for patching has_triton below	12	+from torch.utils._triton import has_triton # noqa: F401
12	from torch.nn.parameter import UninitializedTensorMixin	13	from torch.nn.parameter import UninitializedTensorMixin
13	from torch._utils import _get_device_module	14	from torch._utils import _get_device_module
14	from torch.utils import cpp_extension	15	from torch.utils import cpp_extension
	@@ -454,6 +455,18 @@ def _patch_flex_attention_device():
454		455
455		456
456	def _init():	457	def _init():
		458	+ # transfer_to_npu patches these modules during its own import. Import them
		459	+ # explicitly instead of relying on torch_npu import side effects.
		460	+ import torch._dynamo.trace_rules # noqa: F401
		461	+ import torch._dynamo.utils # noqa: F401
		462	+ import torch._inductor.runtime.autotune_cache # noqa: F401
		463	+ import torch._inductor.compile_fx # noqa: F401
		464	+ import torch._inductor.utils # noqa: F401
		465	+ import torch._inductor.fx_passes.post_grad # noqa: F401
		466	+ import torch._inductor.fx_passes.joint_graph # noqa: F401
		467	+ import torch._inductor.autotune_process # noqa: F401
		468	+ from torch.distributed.checkpoint import filesystem
		469	+
457	_warning_fn('''	470	_warning_fn('''
458	*************************************************************************************************************	471	*************************************************************************************************************
459	The torch.Tensor.cuda and torch.nn.Module.cuda are replaced with torch.Tensor.npu and torch.nn.Module.npu now..	472	The torch.Tensor.cuda and torch.nn.Module.cuda are replaced with torch.Tensor.npu and torch.nn.Module.npu now..
	@@ -559,7 +572,7 @@ def _init():
559	torch._inductor.autotune_process.get_gpu_type = _get_npu_type	572	torch._inductor.autotune_process.get_gpu_type = _get_npu_type
560		573
561	torch._utils._get_available_device_type = _patch_get_available_device_type	574	torch._utils._get_available_device_type = _patch_get_available_device_type
562	- torch.distributed.checkpoint.filesystem._OverlappingCpuLoader.__init__ = \	575	+ filesystem._OverlappingCpuLoader.__init__ = \
563	_patch_OverlappingCpuLoader_init_	576	_patch_OverlappingCpuLoader_init_
564		577
565	_replace_to_method_in_allowed_methods()	578	_replace_to_method_in_allowed_methods()

Mtorch_npu/dynamo/__init__.py+73-5

	@@ -3,9 +3,6 @@ import sys
3	import time	3	import time
4	import warnings	4	import warnings
5		5
6	-from torch._dynamo import register_backend as _register_backend
7	-from torch._dynamo.backends.registry import _BACKENDS
8	-
9	from torch_npu._init.common.warning_utils import _should_print_warning	6	from torch_npu._init.common.warning_utils import _should_print_warning
10	from torch_npu.utils._error_code import ErrCode, pta_error	7	from torch_npu.utils._error_code import ErrCode, pta_error
11		8
	@@ -107,6 +104,17 @@ class _LazyTorchair(_LazyBackend):
107	return getattr(torchair, name)	104	return getattr(torchair, name)
108		105
109		106
		107	+
		108	+
		109	+def _install_lazy_torchair():
		110	+ torchair_path = os.path.join(os.path.dirname(__file__), "torchair")
		111	+ if not os.path.exists(torchair_path):
		112	+ return False
		113	+ if "torchair" not in sys.modules:
		114	+ sys.modules["torchair"] = _LazyTorchair("torchair")
		115	+ return True
		116	+
		117	+
110	class _LazyNpuGraphEx(_LazyBackend):	118	class _LazyNpuGraphEx(_LazyBackend):
111	def __init__(self, pkg_name):	119	def __init__(self, pkg_name):
112	self._npugraph_ex = None	120	self._npugraph_ex = None
	@@ -140,7 +148,7 @@ def _lazy_exec(args, *kwargs):
140		148
141		149
142	def _get_default_backend(name):	150	def _get_default_backend(name):
143	- if not os.path.exists(os.path.join(os.path.dirname(__file__), 'torchair')):	151	+ if not _install_lazy_torchair():
144	if _should_print_warning():	152	if _should_print_warning():
145	warnings.warn(	153	warnings.warn(
146	"Register eager implementation for the 'npu' backend of dynamo, "	154	"Register eager implementation for the 'npu' backend of dynamo, "
	@@ -149,7 +157,6 @@ def _get_default_backend(name):
149		157
150	global _global_backend_name	158	global _global_backend_name
151	_global_backend_name = name	159	_global_backend_name = name
152	- sys.modules['torchair'] = _LazyTorchair('torchair')
153	return _lazy_exec	160	return _lazy_exec
154		161
155		162
	@@ -166,6 +173,27 @@ def _get_npugraph_ex_backend():
166		173
167		174
168	def _register_npu_backend(backend, name="npu"):	175	def _register_npu_backend(backend, name="npu"):
		176	+ from torch._dynamo import register_backend as _register_backend
		177	+ from torch._dynamo.backends.registry import _BACKENDS, _COMPILER_FNS
		178	+
		179	+ registered_backend = _COMPILER_FNS.get(name)
		180	+ if (
		181	+ registered_backend is not None
		182	+ and getattr(registered_backend, "__module__", None) == __name__
		183	+ ):
		184	+ return
		185	+
		186	+ # When a setuptools entry point is currently loading torch_npu, Dynamo has
		187	+ # already put the EntryPoint object in _BACKENDS but has not registered the
		188	+ # loaded callable yet. Leave that registration to lookup_backend(); doing
		189	+ # it here as well makes lookup_backend register the same name twice.
		190	+ pending_backend = _BACKENDS.get(name)
		191	+ if (
		192	+ name not in _COMPILER_FNS
		193	+ and getattr(pending_backend, "module", None) == __name__
		194	+ ):
		195	+ return
		196	+
169	if name in _BACKENDS.keys():	197	if name in _BACKENDS.keys():
170	del _BACKENDS[name]	198	del _BACKENDS[name]
171	_register_backend(backend, name)	199	_register_backend(backend, name)
	@@ -177,3 +205,43 @@ def _register_backends():
177		205
178	_register_npu_backend(global_backend)	206	_register_npu_backend(global_backend)
179	_register_npu_backend(npugraph_ex_backend, NPUGRAPH_EX_BACKEND)	207	_register_npu_backend(npugraph_ex_backend, NPUGRAPH_EX_BACKEND)
		208	+
		209	+
		210	+def _npu_backend_entrypoint(gm, example_inputs, **kwargs):
		211	+ """Set up the NPU Dynamo backend when an entry point is selected."""
		212	+ from torch_npu.utils._dynamo import _lazy_dynamo_setup
		213	+
		214	+ _lazy_dynamo_setup()
		215	+ return _get_default_backend("npu")(gm, example_inputs, **kwargs)
		216	+
		217	+
		218	+def _npugraph_ex_backend_entrypoint(gm, example_inputs, **kwargs):
		219	+ """Set up the NPUGraph-EX backend when an entry point is selected."""
		220	+ from torch_npu.utils._dynamo import _lazy_dynamo_setup
		221	+
		222	+ _lazy_dynamo_setup()
		223	+ return _exec(gm, example_inputs, **kwargs)
		224	+
		225	+
		226	+class _NpugraphsBackendEntryPoint:
		227	+ """Keep the lazy entry point and its reset protocol in one callable."""
		228	+
		229	+ compiler_name = "npugraphs"
		230	+
		231	+ def __call__(self, gm, example_inputs, **kwargs):
		232	+ from torch_npu.utils._dynamo import _lazy_dynamo_setup, _lazy_inductor_setup
		233	+
		234	+ _lazy_dynamo_setup()
		235	+ _lazy_inductor_setup()
		236	+ from torch_npu.utils._graph_tree import NpugraphsBackend
		237	+
		238	+ return NpugraphsBackend()(gm, example_inputs)
		239	+
		240	+ @staticmethod
		241	+ def reset():
		242	+ graph_tree = sys.modules.get("torch_npu.npu._graph_tree")
		243	+ if graph_tree is not None:
		244	+ graph_tree.reset_npugraph_trees()
		245	+
		246	+
		247	+_npugraphs_backend_entrypoint = _NpugraphsBackendEntryPoint()

Mtorch_npu/npu/__init__.py+13-1

	@@ -150,6 +150,7 @@ import traceback
150	import threading	150	import threading
151	import os	151	import os
152	import re	152	import re
		153	+import importlib
153	import torch	154	import torch
154	from torch.storage import _LegacyStorage, _warn_typed_storage_removal	155	from torch.storage import _LegacyStorage, _warn_typed_storage_removal
155	from torch._utils import classproperty	156	from torch._utils import classproperty
	@@ -171,7 +172,6 @@ from .autocast_utils import * # noqa: F403
171	from .backends import * # noqa: F403	172	from .backends import * # noqa: F403
172	from ._backends import * # noqa: F403	173	from ._backends import * # noqa: F403
173	from .deterministic import enable_deterministic_with_backward, disable_deterministic_with_backward # noqa: F403	174	from .deterministic import enable_deterministic_with_backward, disable_deterministic_with_backward # noqa: F403
174	-from . import npugraph_ex
175		175
176	from .graphs import (	176	from .graphs import (
177	NPUGraph,	177	NPUGraph,
	@@ -193,6 +193,18 @@ from ._npugraph_handlers import (
193	)	193	)
194		194
195		195
		196	+def __getattr__(name):
		197	+ if name == "npugraph_ex":
		198	+ module = importlib.import_module("torch_npu.npu.npugraph_ex")
		199	+ globals()[name] = module
		200	+ return module
		201	+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
		202	+
		203	+
		204	+def __dir__():
		205	+ return sorted(set(globals()) \| {"npugraph_ex"})
		206	+
		207	+
196	config = npu_config._npuConfig()	208	config = npu_config._npuConfig()
197		209
198	matmul = npu_config._allowHF32Matmul()	210	matmul = npu_config._allowHF32Matmul()

Mtorch_npu/npu/_graph_tree.py+1-12

	@@ -99,6 +99,7 @@ from torch.utils._ordered_set import OrderedSet
99	from torch.utils.weak import TensorWeakRef	99	from torch.utils.weak import TensorWeakRef
100		100
101	import torch_npu	101	import torch_npu
		102	+from torch_npu.npu._graph_tree_state import MarkStepBox # noqa: F401
102	from torch_npu._C import (	103	from torch_npu._C import (
103	_npu_NPUAllocator_AllocatorState as AllocatorState,	104	_npu_NPUAllocator_AllocatorState as AllocatorState,
104	_set_cached_tensors_enabled as _set_cached_tensors_enabled)	105	_set_cached_tensors_enabled as _set_cached_tensors_enabled)
	@@ -268,24 +269,12 @@ local.npu_tree_manager_containers = {}
268	local.npu_tree_manager_locks = defaultdict(threading.Lock)	269	local.npu_tree_manager_locks = defaultdict(threading.Lock)
269		270
270		271
271	-# only incremented by user call of mark_step_begin
272	-class MarkStepBox:
273	- mark_step_counter = 0
274	-
275	-
276	# We need to register this as an object that will be copied over as TLS when new	272	# We need to register this as an object that will be copied over as TLS when new
277	# threads are created in autograd	273	# threads are created in autograd
278	torch._C._stash_obj_in_tls("npu_tree_manager_containers", local.npu_tree_manager_containers)	274	torch._C._stash_obj_in_tls("npu_tree_manager_containers", local.npu_tree_manager_containers)
279	torch._C._stash_obj_in_tls("npu_tree_manager_locks", local.npu_tree_manager_locks)	275	torch._C._stash_obj_in_tls("npu_tree_manager_locks", local.npu_tree_manager_locks)
280		276
281		277
282	-def mark_step_begin() -> None:
283	- "Indicates that a new iteration of inference or training is about to begin."
284	-
285	- # iterate down to distinguish from GenerationTracking counter
286	- MarkStepBox.mark_step_counter -= 1
287	-
288	-
289	def reset_npugraph_trees() -> None:	278	def reset_npugraph_trees() -> None:
290	"Clear all npugraph trees"	279	"Clear all npugraph trees"
291	# see shutdown below for why this is necessary	280	# see shutdown below for why this is necessary

Atorch_npu/npu/_graph_tree_state.py+11-0

@@ -0,0 +1,11 @@
	1	+"""Lightweight state shared by the public NPUGraph marker and graph trees."""
	2	+
	3	+
	4	+class MarkStepBox:
	5	+ # Negative values distinguish explicit user steps from Dynamo generations.
	6	+ mark_step_counter = 0
	7	+
	8	+
	9	+def mark_step_begin() -> None:
	10	+ """Indicate that a new inference or training iteration is about to begin."""
	11	+ MarkStepBox.mark_step_counter -= 1

Mtorch_npu/npu/deterministic.py+11-3

	@@ -2,11 +2,19 @@ import torch
2		2
3	from torch import Tensor	3	from torch import Tensor
4	from torch.autograd.function import Function	4	from torch.autograd.function import Function
5	-from torch._dynamo.decorators import forbid_in_graph
6		5
7	__all__ = ["enable_deterministic_with_backward", "disable_deterministic_with_backward"]	6	__all__ = ["enable_deterministic_with_backward", "disable_deterministic_with_backward"]
8		7
9		8
		9	+def _forbid_in_graph(fn):
		10	+ if isinstance(fn, (list, tuple)):
		11	+ return [_forbid_in_graph(x) for x in fn]
		12	+ if not callable(fn):
		13	+ raise AssertionError("forbid_in_graph applies only to callables")
		14	+ fn._dynamo_forbidden = True
		15	+ return fn
		16	+
		17	+
10	class _DeterministicAlgorithmsBeginOp(Function):	18	class _DeterministicAlgorithmsBeginOp(Function):
11		19
12	@staticmethod	20	@staticmethod
	@@ -36,11 +44,11 @@ class _DeterministicAlgorithmsEndOp(Function):
36	return grad_outputs	44	return grad_outputs
37		45
38		46
39	-@forbid_in_graph	47	+@_forbid_in_graph
40	def enable_deterministic_with_backward(tensor: Tensor):	48	def enable_deterministic_with_backward(tensor: Tensor):
41	return _DeterministicAlgorithmsBeginOp.apply(tensor)	49	return _DeterministicAlgorithmsBeginOp.apply(tensor)
42		50
43		51
44	-@forbid_in_graph	52	+@_forbid_in_graph
45	def disable_deterministic_with_backward(tensor: Tensor):	53	def disable_deterministic_with_backward(tensor: Tensor):
46	return _DeterministicAlgorithmsEndOp.apply(tensor)	54	return _DeterministicAlgorithmsEndOp.apply(tensor)

Mtorch_npu/utils/__init__.py+2-1

	@@ -2,9 +2,10 @@ __all__ = ["npu_combine_tensors", "get_part_combined_tensor", "is_combined_tenso
2	"set_thread_affinity", "reset_thread_affinity", "save_async", "get_cann_version"]	2	"set_thread_affinity", "reset_thread_affinity", "save_async", "get_cann_version"]
3		3
4	from torch_npu.npu.utils import get_cann_version	4	from torch_npu.npu.utils import get_cann_version
		5	+from ._inductor import _max_unpoolnd_patch
5	from .combine_tensors import npu_combine_tensors, get_part_combined_tensor, is_combined_tensor_valid	6	from .combine_tensors import npu_combine_tensors, get_part_combined_tensor, is_combined_tensor_valid
6	from .serialization import save_async	7	from .serialization import save_async
7	from .flops_count import _FlopsCounter as FlopsCounter	8	from .flops_count import _FlopsCounter as FlopsCounter
8	from .affinity import _set_thread_affinity as set_thread_affinity	9	from .affinity import _set_thread_affinity as set_thread_affinity
9	from .affinity import _reset_thread_affinity as reset_thread_affinity	10	from .affinity import _reset_thread_affinity as reset_thread_affinity
10	-from .asd_detector import set_asd_loss_scale, register_asd_hook	11	+from .asd_detector import set_asd_loss_scale, register_asd_hook

Mtorch_npu/utils/_dynamo.py+269-60

	@@ -1,10 +1,12 @@
		1	+import importlib
		2	+import importlib.abc
		3	+import functools
1	import inspect	4	import inspect
		5	+import logging
2	import os	6	import os
3	import sys	7	import sys
4	-import logging	8	+import threading
5	from typing import Any, Optional, TYPE_CHECKING	9	from typing import Any, Optional, TYPE_CHECKING
6	-import importlib
7	-import functools
8		10
9	import torch	11	import torch
10	import torch_npu	12	import torch_npu
	@@ -15,8 +17,9 @@ use_jit_script = False
15	log = logging.getLogger(__name__)	17	log = logging.getLogger(__name__)
16		18
17	def _create_npu_autocast_mode_variable(func, args, kwargs):	19	def _create_npu_autocast_mode_variable(func, args, kwargs):
18	- from torch._dynamo.variables.ctx_manager import AutocastModeVariable
19	from torch._dynamo.variables.base import VariableTracker	20	from torch._dynamo.variables.base import VariableTracker
		21	+ from torch._dynamo.variables.ctx_manager import AutocastModeVariable
		22	+
20	bound_args = inspect.signature(func).bind(args, *kwargs)	23	bound_args = inspect.signature(func).bind(args, *kwargs)
21	bound_args.apply_defaults()	24	bound_args.apply_defaults()
22	target_values = []	25	target_values = []
	@@ -34,8 +37,8 @@ def _create_npu_autocast_mode_variable(func, args, kwargs):
34	else:	37	else:
35	target_values.append(arg)	38	target_values.append(arg)
36		39
37	- var = AutocastModeVariable(target_values, initial_values=None, **kwargs)	40	+ return AutocastModeVariable(target_values, initial_values=None, **kwargs)
38	- return var	41	+
39		42
40	def patch_SkipFunctionVariable():	43	def patch_SkipFunctionVariable():
41	from torch._dynamo.variables.functions import SkipFunctionVariable	44	from torch._dynamo.variables.functions import SkipFunctionVariable
	@@ -53,11 +56,12 @@ def patch_SkipFunctionVariable():
53	SkipFunctionVariable.__new__raw = SkipFunctionVariable.__new__	56	SkipFunctionVariable.__new__raw = SkipFunctionVariable.__new__
54	SkipFunctionVariable.__new__ = SkipFunctionVariable__new__	57	SkipFunctionVariable.__new__ = SkipFunctionVariable__new__
55		58
		59	+
56	def patch_TensorVariable_call_method():	60	def patch_TensorVariable_call_method():
57	- from torch._dynamo.variables.tensor import TensorVariable
58	from torch._dynamo.utils import tensortype_to_dtype	61	from torch._dynamo.utils import tensortype_to_dtype
59	from torch._dynamo.variables.constant import ConstantVariable	62	from torch._dynamo.variables.constant import ConstantVariable
60	from torch._dynamo.variables.lists import TupleVariable	63	from torch._dynamo.variables.lists import TupleVariable
		64	+ from torch._dynamo.variables.tensor import TensorVariable
61		65
62	def TensorVariable_call_method(self, tx, name, args, kwargs):	66	def TensorVariable_call_method(self, tx, name, args, kwargs):
63	if (	67	if (
	@@ -67,16 +71,21 @@ def patch_TensorVariable_call_method():
67	and isinstance(self.device, torch.device)	71	and isinstance(self.device, torch.device)
68	and self.device.type == "npu"	72	and self.device.type == "npu"
69	):	73	):
70	- tensortype = next(k for k, v in tensortype_to_dtype.items() if self.dtype in v)	74	+ tensortype = next(
71	- constant_result = ConstantVariable.create(f"torch.npu.{tensortype.__name__}")	75	+ k for k, v in tensortype_to_dtype.items() if self.dtype in v
		76	+ )
		77	+ constant_result = ConstantVariable.create(
		78	+ f"torch.npu.{tensortype.__name__}"
		79	+ )
72		80
73	if len(args) == 1:	81	if len(args) == 1:
74	return constant_result.getitem_const(args[0])	82	return constant_result.getitem_const(args[0])
75	- elif args:	83	+ if args:
76	- return TupleVariable([constant_result.getitem_const(a) for a in args])	84	+ return TupleVariable(
		85	+ [constant_result.getitem_const(a) for a in args]
		86	+ )
77	return constant_result	87	return constant_result
78	- else:	88	+ return TensorVariable.call_method_raw(self, tx, name, args, kwargs)
79	- return TensorVariable.call_method_raw(self, tx, name, args, kwargs)
80		89
81	TensorVariable.call_method_raw = TensorVariable.call_method	90	TensorVariable.call_method_raw = TensorVariable.call_method
82	TensorVariable.call_method = TensorVariable_call_method	91	TensorVariable.call_method = TensorVariable_call_method
	@@ -99,7 +108,6 @@ class _InductorNpuRegistry:
99	sys.modules["torch_npu._inductor"]._load_backend()	108	sys.modules["torch_npu._inductor"]._load_backend()
100	cls._loaded_backend = current	109	cls._loaded_backend = current
101		110
102	-
103	@classmethod	111	@classmethod
104	def disable_register(cls):	112	def disable_register(cls):
105	cls._disabled_register = True	113	cls._disabled_register = True
	@@ -129,19 +137,23 @@ def register_inductor_npu():
129	_InductorNpuRegistry.register_inductor_npu()	137	_InductorNpuRegistry.register_inductor_npu()
130		138
131		139
132	-def _resolve_npu_backend_from_wrapper(wrapper) -> str:	140	+def _resolve_npu_backend(selected_backend=None) -> str:
133	- """Resolve npu backend with priority: wrapper options > global config > env."""	141	+ """Resolve NPU backend with priority: compile options > config > env."""
134	- wrapper_backend = wrapper.config.get("npu_backend")	142	+ if selected_backend not in (None, "", "default"):
135	- if wrapper_backend not in (None, "", "default"):	143	+ return selected_backend
136	- return wrapper_backend
137		144
138	- global_backend = getattr(torch._inductor.config, "npu_backend", None)	145	+ inductor_config = sys.modules.get("torch._inductor.config")
		146	+ global_backend = getattr(inductor_config, "npu_backend", None)
139	if global_backend not in (None, "", "default"):	147	if global_backend not in (None, "", "default"):
140	return global_backend	148	return global_backend
141		149
142	return os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default")	150	return os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default")
143		151
144		152
		153	+def _resolve_npu_backend_from_wrapper(wrapper) -> str:
		154	+ return _resolve_npu_backend(wrapper.config.get("npu_backend"))
		155	+
		156	+
145	class _NpuBackendScope:	157	class _NpuBackendScope:
146	"""Apply resolved npu backend for one compile invocation and restore env."""	158	"""Apply resolved npu backend for one compile invocation and restore env."""
147		159
	@@ -151,22 +163,29 @@ class _NpuBackendScope:
151		163
152	def __enter__(self):	164	def __enter__(self):
153	self._old_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")	165	self._old_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")
154	- os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend	166	+ try:
155	- register_inductor_npu()	167	+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend
156	- if self.backend == "ascendc":	168	+ register_inductor_npu()
157	- from torch_npu._inductor.deterministic_cache import (	169	+ if self.backend == "ascendc":
158	- patch_npu_deterministic_level_cache_keys,	170	+ from torch_npu._inductor.deterministic_cache import (
159	- )	171	+ patch_npu_deterministic_level_cache_keys,
		172	+ )
160		173
161	- patch_npu_deterministic_level_cache_keys()	174	+ patch_npu_deterministic_level_cache_keys()
		175	+ except BaseException:
		176	+ self._restore_backend_env()
		177	+ raise
162	return self	178	return self
163		179
164	def __exit__(self, exc_type, exc, tb):	180	def __exit__(self, exc_type, exc, tb):
		181	+ self._restore_backend_env()
		182	+ return False
		183	+
		184	+ def _restore_backend_env(self):
165	if self._old_env is None:	185	if self._old_env is None:
166	os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)	186	os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)
167	else:	187	else:
168	os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self._old_env	188	os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self._old_env
169	- return False
170		189
171		190
172	def patch_inductor_wrapper():	191	def patch_inductor_wrapper():
	@@ -183,15 +202,22 @@ def patch_inductor_wrapper():
183	src_call = _TorchCompileInductorWrapper.__call__	202	src_call = _TorchCompileInductorWrapper.__call__
184		203
185	def new_apply_options(self, options: Optional[dict[str, Any]]):	204	def new_apply_options(self, options: Optional[dict[str, Any]]):
186	- if options is not None and options.get("enable_shape_handling", False):	205	+ shape_handling_requested = (
		206	+ options is not None and options.get("enable_shape_handling", False)
		207	+ )
		208	+ src_apply_options(self, options)
		209	+ if shape_handling_requested:
		210	+ if getattr(self, "_npu_defer_shape_handling", False):
		211	+ self._npu_shape_handling_requested = True
		212	+ return
187	if not is_inductor_npu_initialized():	213	if not is_inductor_npu_initialized():
188	register_inductor_npu()	214	register_inductor_npu()
189	torch_npu._inductor.patch_shape_handling()	215	torch_npu._inductor.patch_shape_handling()
190	- return src_apply_options(self, options)
191		216
192	def new_get_config_copy(self) -> dict[str, Any]:	217	def new_get_config_copy(self) -> dict[str, Any]:
193	ori_dict = src_get_config_copy(self)	218	ori_dict = src_get_config_copy(self)
194	- if self is not torch._inductor.config:	219	+ inductor_config = sys.modules.get("torch._inductor.config")
		220	+ if inductor_config is None or self is not inductor_config:
195	return ori_dict	221	return ori_dict
196	if "npu_backend" not in ori_dict:	222	if "npu_backend" not in ori_dict:
197	ori_dict["npu_backend"] = "default"	223	ori_dict["npu_backend"] = "default"
	@@ -223,11 +249,20 @@ def patch_inductor_wrapper():
223	return ori_dict	249	return ori_dict
224		250
225	def new_init(self, mode, options, dynamic, name=None):	251	def new_init(self, mode, options, dynamic, name=None):
226	- add_dynamo_methods_init()	252	+ self._npu_defer_shape_handling = True
227	- if name is not None:	253	+ self._npu_shape_handling_requested = False
228	- src_init(self, mode, options, dynamic, name)	254	+ try:
229	- else:	255	+ if name is not None:
230	- src_init(self, mode, options, dynamic)	256	+ src_init(self, mode, options, dynamic, name)
		257	+ else:
		258	+ src_init(self, mode, options, dynamic)
		259	+ shape_handling_requested = self._npu_shape_handling_requested
		260	+ finally:
		261	+ del self._npu_defer_shape_handling
		262	+ del self._npu_shape_handling_requested
		263	+ _setup_inductor_for_compile(self.config)
		264	+ if shape_handling_requested:
		265	+ torch_npu._inductor.patch_shape_handling()
231	backend = _resolve_npu_backend_from_wrapper(self)	266	backend = _resolve_npu_backend_from_wrapper(self)
232	if backend == "mlir":	267	if backend == "mlir":
233	with _NpuBackendScope(backend):	268	with _NpuBackendScope(backend):
	@@ -253,15 +288,14 @@ def patch_inductor_wrapper():
253	_TorchCompileInductorWrapper.apply_options = new_apply_options	288	_TorchCompileInductorWrapper.apply_options = new_apply_options
254	_TorchCompileInductorWrapper.__init__ = new_init	289	_TorchCompileInductorWrapper.__init__ = new_init
255	ConfigModule.get_config_copy = new_get_config_copy	290	ConfigModule.get_config_copy = new_get_config_copy
256	- torch._inductor.config.get_config_copy()
257		291
258		292
259	def patch_dynamo_optimize():	293	def patch_dynamo_optimize():
260	from torch_npu.dynamo import _get_global_npu_backend	294	from torch_npu.dynamo import _get_global_npu_backend
		295	+
261	src_optimize = torch._dynamo.optimize	296	src_optimize = torch._dynamo.optimize
262		297
263	def npu_optimize(args, *kwargs):	298	def npu_optimize(args, *kwargs):
264	- add_dynamo_methods_init()
265	backend = None	299	backend = None
266	if "backend" in kwargs:	300	if "backend" in kwargs:
267	backend = kwargs["backend"]	301	backend = kwargs["backend"]
	@@ -443,6 +477,7 @@ def register_npu_current_stream_handler():
443	handlers[torch.npu.current_stream] = accelerator_handler	477	handlers[torch.npu.current_stream] = accelerator_handler
444	return	478	return
445		479
		480	+
446	def patch_user_defined_class_variable():	481	def patch_user_defined_class_variable():
447	import functools	482	import functools
448	from torch._dynamo.variables.user_defined import UserDefinedClassVariable	483	from torch._dynamo.variables.user_defined import UserDefinedClassVariable
	@@ -491,20 +526,67 @@ def patch_user_defined_class_variable():
491		526
492		527
493	def run_once(f):	528	def run_once(f):
494	- """Runs a function (successfully) only once.	529	+ """Run a function successfully only once, waiting for concurrent callers."""
495	- The running can be reset by setting the `has_run` attribute to False	530	+ condition = threading.Condition()
496	- """	531	+
497	@functools.wraps(f)	532	@functools.wraps(f)
498	def wrapper(args, *kwargs):	533	def wrapper(args, *kwargs):
499	- if not wrapper.has_run:	534	+ thread_id = threading.get_ident()
		535	+ with condition:
		536	+ while wrapper._is_running:
		537	+ if wrapper._running_thread == thread_id:
		538	+ return None
		539	+ condition.wait()
		540	+ if wrapper.has_run:
		541	+ return None
		542	+ wrapper._is_running = True
		543	+ wrapper._running_thread = thread_id
		544	+
		545	+ try:
500	result = f(args, *kwargs)	546	result = f(args, *kwargs)
		547	+ except BaseException:
		548	+ with condition:
		549	+ wrapper._is_running = False
		550	+ wrapper._running_thread = None
		551	+ condition.notify_all()
		552	+ raise
		553	+
		554	+ with condition:
501	wrapper.has_run = True	555	wrapper.has_run = True
502	- return result	556	+ wrapper._is_running = False
503	- return None	557	+ wrapper._running_thread = None
		558	+ condition.notify_all()
		559	+ return result
		560	+
504	wrapper.has_run = False	561	wrapper.has_run = False
		562	+ wrapper._is_running = False
		563	+ wrapper._running_thread = None
		564	+
		565	+ def reset_after_fork():
		566	+ # The parent thread running f may not exist in the child process.
		567	+ nonlocal condition
		568	+ condition = threading.Condition()
		569	+ wrapper._is_running = False
		570	+ wrapper._running_thread = None
		571	+
		572	+ try:
		573	+ os.register_at_fork(after_in_child=reset_after_fork)
		574	+ except AttributeError:
		575	+ pass
505	return wrapper	576	return wrapper
506		577
507		578
		579	+_COMPLETED_DYNAMO_SETUP_STEPS = set()
		580	+
		581	+
		582	+def _run_dynamo_setup_step(name, setup):
		583	+ """Keep successful setup steps idempotent when a later step fails."""
		584	+ if name in _COMPLETED_DYNAMO_SETUP_STEPS:
		585	+ return
		586	+ setup()
		587	+ _COMPLETED_DYNAMO_SETUP_STEPS.add(name)
		588	+
		589	+
508	@run_once	590	@run_once
509	def _dynamo_register_interface_for_device():	591	def _dynamo_register_interface_for_device():
510	from torch._dynamo.device_interface import register_interface_for_device	592	from torch._dynamo.device_interface import register_interface_for_device
	@@ -512,25 +594,75 @@ def _dynamo_register_interface_for_device():
512		594
513	register_interface_for_device("npu", NpuInterface)	595	register_interface_for_device("npu", NpuInterface)
514	for i in range(32):	596	for i in range(32):
515	-
516	register_interface_for_device(f"npu:{i}", NpuInterface)	597	register_interface_for_device(f"npu:{i}", NpuInterface)
517		598
518	register_npu_current_stream_handler()	599	register_npu_current_stream_handler()
519		600
520	-def register_npu_graphsafe_rng():
521	- """Register NPU for upstream GraphSafe RNG (same contract as CUDA)."""
522	- from torch._functorch._aot_autograd.utils import register_graphsafe_rng_device_type
523		601
524	- register_graphsafe_rng_device_type("npu")	602	+def _find_spec_without_finder(finder, fullname):
		603	+ """Delegate to the remaining meta-path finders without bypassing them."""
		604	+ try:
		605	+ index = sys.meta_path.index(finder)
		606	+ except ValueError:
		607	+ return importlib.util.find_spec(fullname)
		608	+
		609	+ sys.meta_path.pop(index)
		610	+ try:
		611	+ return importlib.util.find_spec(fullname)
		612	+ finally:
		613	+ sys.meta_path.insert(min(index, len(sys.meta_path)), finder)
		614	+
		615	+
		616	+class _DynamoPostImportLoader(importlib.abc.Loader):
		617	+ def __init__(self, loader, finder):
		618	+ self._loader = loader
		619	+ self._finder = finder
		620	+
		621	+ def create_module(self, spec):
		622	+ create_module = getattr(self._loader, "create_module", None)
		623	+ return create_module(spec) if create_module is not None else None
		624	+
		625	+ def exec_module(self, module):
		626	+ self._loader.exec_module(module)
		627	+ _lazy_dynamo_setup()
		628	+ if self._finder in sys.meta_path:
		629	+ sys.meta_path.remove(self._finder)
		630	+
		631	+
		632	+class _DynamoPostImportFinder(importlib.abc.MetaPathFinder):
		633	+ _target = "torch._dynamo"
		634	+
		635	+ def find_spec(self, fullname, path=None, target=None):
		636	+ if fullname != self._target:
		637	+ return None
		638	+ spec = _find_spec_without_finder(self, fullname)
		639	+ if spec is not None and spec.loader is not None:
		640	+ spec.loader = _DynamoPostImportLoader(spec.loader, self)
		641	+ return spec
		642	+
		643	+
		644	+def _install_dynamo_post_import_trigger():
		645	+ """Set up NPU integration whenever Dynamo is first imported."""
		646	+ if "torch._dynamo" in sys.modules:
		647	+ _lazy_dynamo_setup()
		648	+ return
		649	+ if not any(isinstance(finder, _DynamoPostImportFinder) for finder in sys.meta_path):
		650	+ sys.meta_path.insert(0, _DynamoPostImportFinder())
		651	+
525		652
526	@run_once	653	@run_once
527	def add_dynamo_methods_init():	654	def add_dynamo_methods_init():
528	- _dynamo_register_interface_for_device()	655	+ steps = (
529	- patch_SkipFunctionVariable()	656	+ ("device_interface", _dynamo_register_interface_for_device),
530	- patch_TensorVariable_call_method()	657	+ ("skip_function_variable", patch_SkipFunctionVariable),
531	- patch_user_defined_class_variable()	658	+ ("tensor_variable", patch_TensorVariable_call_method),
532	- patch_stream_event_variable_python_type()	659	+ ("user_defined_class_variable", patch_user_defined_class_variable),
533	- patch_npu_stream_context()	660	+ ("stream_event_variable", patch_stream_event_variable_python_type),
		661	+ ("npu_stream_context", patch_npu_stream_context),
		662	+ ("builtin_variable", patch_builtin_variable),
		663	+ )
		664	+ for name, setup in steps:
		665	+ _run_dynamo_setup_step(name, setup)
534		666
535		667
536	@functools.lru_cache(None)	668	@functools.lru_cache(None)
	@@ -572,12 +704,89 @@ def has_triton() -> bool:
572		704
573		705
574	def patch_has_triton():	706	def patch_has_triton():
575	- torch.utils._triton.has_triton = has_triton	707	+ from torch.utils import _triton
		708	+
		709	+ _triton.has_triton = has_triton
		710	+
		711	+
		712	+def register_npu_graphsafe_rng():
		713	+ """Register NPU for upstream GraphSafe RNG (same contract as CUDA)."""
		714	+ from torch._functorch._aot_autograd.utils import register_graphsafe_rng_device_type
		715	+
		716	+ register_graphsafe_rng_device_type("npu")
		717	+
		718	+
		719	+@run_once
		720	+def _inject_inductor_npu_backend_config():
		721	+ """Inject NPU entries into torch._inductor.config on first use."""
		722	+ torch._inductor.config.get_config_copy()
		723	+
		724	+
		725	+@run_once
		726	+def _lazy_dynamo_setup():
		727	+ """Initialize the Dynamo integration on the first graph-capture operation."""
		728	+ add_dynamo_methods_init()
		729	+
		730	+ from torch_npu.dynamo import _register_backends
		731	+ _run_dynamo_setup_step("backends", _register_backends)
		732	+
		733	+ from torch_npu.dynamo.trace_rule import _patch_npu_trace_rules
		734	+ _run_dynamo_setup_step("trace_rules", _patch_npu_trace_rules)
		735	+
		736	+ _run_dynamo_setup_step("graphsafe_rng", register_npu_graphsafe_rng)
		737	+ _run_dynamo_setup_step("dynamo_optimize", patch_dynamo_optimize)
		738	+
		739	+
		740	+@run_once
		741	+def _lazy_inductor_setup():
		742	+ """Initialize NPU Inductor support only for an Inductor-based backend."""
		743	+ register_inductor_npu()
		744	+
		745	+ from torch_npu.utils._graph_tree import _apply_npugraph_tree_methods
		746	+ _apply_npugraph_tree_methods()
		747	+
		748	+ _inject_inductor_npu_backend_config()
		749	+
		750	+
		751	+def _setup_inductor_for_compile(options=None):
		752	+ """Initialize the NPU Inductor backend selected for this compile call."""
		753	+ _lazy_dynamo_setup()
		754	+
		755	+ option_backend = options.get("npu_backend") if isinstance(options, dict) else None
		756	+ selected_backend = _resolve_npu_backend(option_backend)
		757	+
		758	+ old_backend = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")
		759	+ if selected_backend not in (None, "", "default"):
		760	+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = selected_backend
		761	+ try:
		762	+ _lazy_inductor_setup()
		763	+ finally:
		764	+ if old_backend is None:
		765	+ os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)
		766	+ else:
		767	+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = old_backend
		768	+ return selected_backend
		769	+
		770	+
		771	+@run_once
		772	+def install_npugraph_mark_step_trigger():
		773	+ """Expose the public NPUGraph step API without importing compiler internals."""
		774	+ def npugraph_mark_step_begin():
		775	+ from torch_npu.npu._graph_tree_state import mark_step_begin
		776	+ return mark_step_begin()
		777	+
		778	+ torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin
576		779
577		780
578	def add_dynamo_methods():	781	def add_dynamo_methods():
579	- patch_dynamo_optimize()
580	- patch_builtin_variable()
581	- patch_inductor_wrapper()
582	- register_npu_graphsafe_rng()
583	patch_has_triton()	782	patch_has_triton()
		783	+
		784	+ from torch_npu.dynamo import _install_lazy_torchair
		785	+
		786	+ _install_lazy_torchair()
		787	+ _install_dynamo_post_import_trigger()
		788	+ if "npugraph_ex" not in sys.modules:
		789	+ from torch_npu.dynamo import _LazyNpuGraphEx
		790	+ sys.modules["npugraph_ex"] = _LazyNpuGraphEx("npugraph_ex")
		791	+ patch_inductor_wrapper()
		792	+ install_npugraph_mark_step_trigger()

Mtorch_npu/utils/_graph_tree.py+8-5

	@@ -20,7 +20,7 @@ from torch._dynamo.backends.cudagraphs import (
20	get_stack_traces,	20	get_stack_traces,
21	)	21	)
22	from torch._dynamo.backends.debugging import boxed_nop	22	from torch._dynamo.backends.debugging import boxed_nop
23	-from torch._dynamo.backends.registry import register_backend	23	+from torch._dynamo.backends.registry import _COMPILER_FNS, register_backend
24	from torch._inductor import config	24	from torch._inductor import config
25	from torch._inductor.compile_fx import (	25	from torch._inductor.compile_fx import (
26	get_input_idxs_to_check,	26	get_input_idxs_to_check,
	@@ -53,7 +53,7 @@ log = torch._logging.getArtifactLogger("torch_npu.npugraph", "cudagraphs")
53		53
54		54
55	def npugraph_mark_step_begin():	55	def npugraph_mark_step_begin():
56	- from torch_npu.npu._graph_tree import mark_step_begin	56	+ from torch_npu.npu._graph_tree_state import mark_step_begin
57	mark_step_begin()	57	mark_step_begin()
58		58
59		59
	@@ -374,9 +374,9 @@ class NpugraphsBackend:
374		374
375	@staticmethod	375	@staticmethod
376	def reset():	376	def reset():
377	- from torch_npu.npu._graph_tree import reset_npugraph_trees	377	+ from torch_npu.dynamo import _npugraphs_backend_entrypoint
378		378
379	- reset_npugraph_trees()	379	+ _npugraphs_backend_entrypoint.reset()
380		380
381	@staticmethod	381	@staticmethod
382	def __call__(model, inputs):	382	def __call__(model, inputs):
	@@ -386,7 +386,10 @@ class NpugraphsBackend:
386	def _apply_npugraph_tree_methods():	386	def _apply_npugraph_tree_methods():
387	# aot_npugraphs only applies graphs to the graph. It is also helpful	387	# aot_npugraphs only applies graphs to the graph. It is also helpful
388	# for debugging and can serve as a perf baseline.	388	# for debugging and can serve as a perf baseline.
389	- register_backend(name="npugraphs", compiler_fn=NpugraphsBackend())	389	+ if "npugraphs" not in _COMPILER_FNS:
		390	+ from torch_npu.dynamo import _npugraphs_backend_entrypoint
		391	+
		392	+ register_backend(name="npugraphs", compiler_fn=_npugraphs_backend_entrypoint)
390	torch._inductor.compile_fx.cudagraphify = npugraphify	393	torch._inductor.compile_fx.cudagraphify = npugraphify
391	torch._inductor.cudagraph_utils.check_multiple_devices_or_any_cpu_nodes = check_multiple_devices_or_any_cpu_nodes	394	torch._inductor.cudagraph_utils.check_multiple_devices_or_any_cpu_nodes = check_multiple_devices_or_any_cpu_nodes
392	torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin	395	torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin