Star843
1300
代码
代码
Issues1414
Pull Requests2257
流水线
项目讨论
Wiki
分析
Star843
1300
已合并
perf: defer dynamo/inductor loading on v2.9.0 (#2788) #43444
黄桂军创建于 20 天前
perf: defer dynamo/inductor loading on v2.9.0 (#2788) #43444
已合并
黄桂军创建于 20 天前
从stevenaw0/pytorchv2.9.0_import合入到Ascend/pytorchv2.9.0
共 22 个文件变更+1476-174
Msetup.py+5-0
@@ -786,5 +786,10 @@ setup(
786 'torch.backends': [786 'torch.backends': [
787 'torch_npu = torch_npu:_autoload',787 'torch_npu = torch_npu:_autoload',
788 ],788 ],
789+ 'torch_dynamo_backends': [
790+ 'npu = torch_npu.dynamo:_npu_backend_entrypoint',
791+ 'npugraph_ex = torch_npu.dynamo:_npugraph_ex_backend_entrypoint',
792+ 'npugraphs = torch_npu.dynamo:_npugraphs_backend_entrypoint',
793+ ],
789 }794 }
790)795)
Mtest/_inductor/test_current_device.py+2-1
@@ -1,4 +1,5 @@
1import torch1import torch
2+import torch._dynamo.testing
2from torch.testing._internal.common_utils import (3from torch.testing._internal.common_utils import (
3 run_tests,4 run_tests,
4 instantiate_parametrized_tests,5 instantiate_parametrized_tests,
@@ -30,4 +31,4 @@ instantiate_parametrized_tests(TestCurrentDevice)
30 31 
31 32 
32if __name__ == "__main__":33if __name__ == "__main__":
33- run_tests()34+ run_tests()
Mtest/_inductor/test_mlir_enable.py+4-1
@@ -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)
@@ -38,4 +41,4 @@ class TestAdd(TestUtils):
38instantiate_parametrized_tests(TestAdd)41instantiate_parametrized_tests(TestAdd)
39 42 
40if __name__ == "__main__":43if __name__ == "__main__":
41- run_tests()44+ run_tests()
Atest/dynamo/test_compile_trigger.py+989-0
@@ -0,0 +1,989 @@
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):
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=60,
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+ 
183+ # Verify a forked child discards inherited in-progress lazy setup state.
184+ def test_public_lazy_setup_recovers_after_fork(self):
185+ self.run_in_subprocess(
186+ """
187+ import os
188+ import signal
189+ import threading
190+ 
191+ import torch_npu
192+ import torch_npu.dynamo as npu_dynamo
193+ from torch_npu.utils import _dynamo
194+ 
195+ parent_pid = os.getpid()
196+ original_add = _dynamo.add_dynamo_methods_init
197+ original_get_backend = npu_dynamo._get_default_backend
198+ entered = threading.Event()
199+ release = threading.Event()
200+ 
201+ def block_parent_setup():
202+ if os.getpid() == parent_pid:
203+ entered.set()
204+ assert release.wait(timeout=10)
205+ return original_add()
206+ 
207+ _dynamo.add_dynamo_methods_init = block_parent_setup
208+ setup_thread = threading.Thread(target=_dynamo._lazy_dynamo_setup)
209+ setup_thread.start()
210+ assert entered.wait(timeout=5)
211+ 
212+ try:
213+ child_pid = os.fork()
214+ if child_pid == 0:
215+ def timeout(_signal, _frame):
216+ raise TimeoutError("lazy setup hung after fork")
217+ 
218+ signal.signal(signal.SIGALRM, timeout)
219+ signal.alarm(5)
220+ try:
221+ npu_dynamo._get_default_backend = (
222+ lambda name: npu_dynamo._eager_npu_backend
223+ )
224+ graph_module = lambda x: x
225+ result = npu_dynamo._npu_backend_entrypoint(
226+ graph_module, []
227+ )
228+ assert result is graph_module
229+ signal.alarm(0)
230+ except BaseException as error:
231+ os.write(
232+ 2,
233+ f"child lazy setup failed: {error}\\n".encode(),
234+ )
235+ os._exit(1)
236+ os._exit(0)
237+ 
238+ _, status = os.waitpid(child_pid, 0)
239+ finally:
240+ release.set()
241+ setup_thread.join(timeout=10)
242+ _dynamo.add_dynamo_methods_init = original_add
243+ npu_dynamo._get_default_backend = original_get_backend
244+ 
245+ assert os.WIFEXITED(status)
246+ assert os.waitstatus_to_exitcode(status) == 0
247+ assert not setup_thread.is_alive()
248+ """
249+ )
250+ 
251+ 
252+ # Verify NPUGraphs rejects unsupported options in every registration order.
253+ def test_npugraphs_rejects_options_across_registration_order(self):
254+ self.run_in_subprocess(
255+ """
256+ import contextlib
257+ from unittest import mock
258+ 
259+ import torch
260+ import torch_npu
261+ from torch_npu.dynamo import _npugraphs_backend_entrypoint
262+ from torch_npu.utils import _dynamo, _graph_tree
263+ 
264+ gm = object()
265+ inputs = [object()]
266+ options = {"npu_backend": "mlir"}
267+ 
268+ with mock.patch.object(
269+ _dynamo, "_lazy_dynamo_setup", lambda: None
270+ ), mock.patch.object(
271+ _dynamo, "_lazy_inductor_setup", lambda: None
272+ ), mock.patch.object(
273+ _dynamo,
274+ "_NpuBackendScope",
275+ lambda backend: contextlib.nullcontext(),
276+ ), mock.patch.object(
277+ _graph_tree,
278+ "npugraphs",
279+ lambda model, args, **kwargs: "unexpected",
280+ ):
281+ for backend in (
282+ _npugraphs_backend_entrypoint,
283+ _graph_tree.NpugraphsBackend(),
284+ ):
285+ try:
286+ backend(gm, inputs, options=options)
287+ except TypeError as error:
288+ assert "unexpected keyword argument 'options'" in str(error)
289+ else:
290+ raise AssertionError("npugraphs must reject options")
291+ 
292+ compiled = torch.compile(
293+ lambda value: value + 1,
294+ backend="npugraphs",
295+ options=options,
296+ )
297+ try:
298+ compiled(torch.ones(1))
299+ except Exception as error:
300+ assert "unexpected keyword argument 'options'" in str(error)
301+ else:
302+ raise AssertionError("public npugraphs must reject options")
303+ """
304+ )
305+ 
306+ # Verify public reset reaches NPUGraphs for every registration order.
307+ 
308+ def test_npugraphs_reset_protocol_across_registration_order(self):
309+ for order in ("cold", "hot"):
310+ with self.subTest(order=order):
311+ initialize_inductor = (
312+ "torch.compile(lambda x: x + 1, backend='inductor')"
313+ if order == "hot"
314+ else ""
315+ )
316+ self.run_in_subprocess(
317+ f"""
318+ import sys
319+ import types
320+ from unittest import mock
321+ 
322+ import torch
323+ import torch_npu
324+ 
325+ {initialize_inductor}
326+ torch.compile(lambda x: x + 1, backend="npugraphs")
327+ 
328+ from torch._dynamo.backends import registry
329+ from torch_npu.dynamo import _npugraphs_backend_entrypoint
330+ 
331+ backend = registry._COMPILER_FNS["npugraphs"]
332+ assert backend is _npugraphs_backend_entrypoint
333+ assert hasattr(backend, "reset")
334+ 
335+ graph_tree_module = "torch_npu.npu._graph_tree"
336+ assert graph_tree_module not in sys.modules
337+ backend.reset()
338+ assert graph_tree_module not in sys.modules
339+ 
340+ reset_calls = []
341+ fake_graph_tree = types.ModuleType(graph_tree_module)
342+ fake_graph_tree.reset_npugraph_trees = (
343+ lambda: reset_calls.append("reset")
344+ )
345+ with mock.patch.dict(
346+ sys.modules,
347+ {{graph_tree_module: fake_graph_tree}},
348+ ):
349+ torch.compiler.reset()
350+ 
351+ assert reset_calls == ["reset"]
352+ """
353+ )
354+ 
355+ 
356+ # Verify every public Export entry initializes only Dynamo on NPU.
357+ def test_npu_export_public_entry_and_import_order_matrix(self):
358+ cases = {
359+ "module_export": (
360+ "",
361+ "exported = torch.export.export(Model(), (x,), strict=True)",
362+ ),
363+ "prebound_export": (
364+ "from torch.export import export as export_api",
365+ "exported = export_api(Model(), (x,), strict=True)",
366+ ),
367+ "prebound_export_for_training": (
368+ "from torch.export import export_for_training as export_api",
369+ "exported = export_api(Model(), (x,), strict=True)",
370+ ),
371+ }
372+ if hasattr(torch.export, "export_for_inference"):
373+ cases["prebound_export_for_inference"] = (
374+ "from torch.export import export_for_inference as export_api",
375+ "exported = export_api(Model(), (x,), strict=True)",
376+ )
377+ for name, (pre_import, export_call) in cases.items():
378+ with self.subTest(name=name):
379+ self.run_in_subprocess(
380+ f"""
381+ import sys
382+ import torch
383+ {pre_import}
384+ import torch_npu
385+ 
386+ stream = torch.npu.Stream()
387+ 
388+ class Model(torch.nn.Module):
389+ def forward(self, x):
390+ x.record_stream(stream)
391+ return x + 1
392+ 
393+ x = torch.ones(4, device="npu")
394+ {export_call}
395+ actual = exported.module()(x)
396+ 
397+ from torch_npu.utils import _dynamo
398+ 
399+ torch.testing.assert_close(actual, x + 1)
400+ assert _dynamo._lazy_dynamo_setup.has_run
401+ assert not _dynamo._lazy_inductor_setup.has_run
402+ assert "torch_npu._inductor" not in sys.modules
403+ """
404+ )
405+ 
406+ # Verify NPU-specific operations retain their Export capture semantics.
407+ def test_npu_export_capture_semantics_matrix(self):
408+ self.run_in_subprocess(
409+ """
410+ import sys
411+ import torch
412+ import torch_npu
413+ 
414+ x = torch.ones(4, device="npu")
415+ stream = torch.npu.Stream()
416+ event = torch.npu.Event()
417+ 
418+ class StreamAndEvent(torch.nn.Module):
419+ def forward(self, value):
420+ event.record()
421+ with torch.npu.stream(stream):
422+ event.wait(stream)
423+ result = value + 1
424+ return result
425+ 
426+ class Autocast(torch.nn.Module):
427+ def forward(self, value):
428+ with torch.npu.amp.autocast(dtype=torch.float16):
429+ return value * value
430+ 
431+ class CurrentDevice(torch.nn.Module):
432+ def forward(self, value):
433+ return value + torch.npu.current_device()
434+ 
435+ class DeviceProperties(torch.nn.Module):
436+ def forward(self, value):
437+ properties = torch.npu.get_device_properties(
438+ torch.npu.current_device()
439+ )
440+ return value + 1 if properties.total_memory > 0 else value - 1
441+ 
442+ class IsAvailable(torch.nn.Module):
443+ def forward(self, value):
444+ return value + 1 if torch.npu.is_available() else value - 1
445+ 
446+ models = (
447+ StreamAndEvent,
448+ Autocast,
449+ CurrentDevice,
450+ DeviceProperties,
451+ IsAvailable,
452+ )
453+ for model_type in models:
454+ model = model_type()
455+ expected = model(x)
456+ exported = torch.export.export(model, (x,))
457+ actual = exported.module()(x)
458+ torch.testing.assert_close(actual, expected)
459+ 
460+ from torch_npu.utils import _dynamo
461+ 
462+ assert _dynamo._lazy_dynamo_setup.has_run
463+ assert not _dynamo._lazy_inductor_setup.has_run
464+ assert "torch_npu._inductor" not in sys.modules
465+ """
466+ )
467+ 
468+ # Verify rejected Inductor arguments do not initialize or pollute NPU state.
469+ def test_invalid_inductor_arguments_fail_without_npu_initialization(self):
470+ cases = {
471+ "invalid_mode": (
472+ 'torch.compile(lambda x: x + 1, mode="invalid-mode")',
473+ "Unrecognized mode=invalid-mode",
474+ ),
475+ "invalid_option": (
476+ "torch.compile(lambda x: x + 1, "
477+ 'options={"invalid.option": True})',
478+ "Unexpected optimization option invalid.option",
479+ ),
480+ }
481+ for name, (compile_call, expected_error) in cases.items():
482+ with self.subTest(name=name):
483+ self.run_in_subprocess(
484+ f"""
485+ import os
486+ import sys
487+ 
488+ import torch
489+ import torch_npu
490+ from torch_npu.utils import _dynamo
491+ 
492+ env_name = "TORCHINDUCTOR_NPU_BACKEND"
493+ original_env = os.environ.get(env_name)
494+ try:
495+ {compile_call}
496+ except RuntimeError as error:
497+ assert {expected_error!r} in str(error), str(error)
498+ else:
499+ raise AssertionError("invalid compile arguments must fail")
500+ 
501+ assert not _dynamo._lazy_inductor_setup.has_run
502+ assert not _dynamo.is_inductor_npu_initialized()
503+ assert "torch_npu._inductor" not in sys.modules
504+ assert os.environ.get(env_name) == original_env
505+ """
506+ )
507+ 
508+ # Verify scope entry failures restore the process environment.
509+ def test_npu_backend_scope_restores_env_after_entry_failure(self):
510+ self.run_in_subprocess(
511+ """
512+ import os
513+ 
514+ import torch_npu
515+ from torch_npu.utils import _dynamo
516+ 
517+ env_name = "TORCHINDUCTOR_NPU_BACKEND"
518+ original_env = os.environ.get(env_name)
519+ try:
520+ with _dynamo._NpuBackendScope(1):
521+ raise AssertionError("scope entry must fail")
522+ except TypeError as error:
523+ assert "str" in str(error)
524+ else:
525+ raise AssertionError("a non-string backend must fail")
526+ 
527+ assert os.environ.get(env_name) == original_env
528+ """
529+ )
530+ 
531+ # Verify shape handling is installed after selecting the requested NPU backend.
532+ def test_shape_handling_initializes_after_backend_selection(self):
533+ self.run_in_subprocess(
534+ """
535+ import types
536+ from unittest import mock
537+ 
538+ import torch
539+ import torch_npu
540+ from torch_npu.utils import _dynamo
541+ 
542+ def fake_setup(actual_options):
543+ actual_options = dict(actual_options)
544+ events.append(("setup", actual_options))
545+ return actual_options["npu_backend"]
546+ 
547+ options = {
548+ "npu_backend": "mlir",
549+ "enable_shape_handling": True,
550+ }
551+ events = []
552+ 
553+ def scope_register():
554+ events.append(("scope_register", None))
555+ 
556+ fake_inductor = types.SimpleNamespace(
557+ patch_shape_handling=lambda: events.append(
558+ ("shape_handling", None)
559+ )
560+ )
561+ with mock.patch.object(
562+ _dynamo, "_setup_inductor_for_compile", fake_setup
563+ ), mock.patch.object(
564+ _dynamo, "register_inductor_npu", scope_register
565+ ), mock.patch.object(
566+ torch_npu, "_inductor", fake_inductor, create=True
567+ ):
568+ wrapper = torch._TorchCompileInductorWrapper(None, options, None)
569+ 
570+ assert wrapper.config["npu_backend"] == "mlir"
571+ assert wrapper.config["enable_shape_handling"] is True
572+ assert events == [
573+ ("setup", options),
574+ ("shape_handling", None),
575+ ("scope_register", None),
576+ ], events
577+ """
578+ )
579+ 
580+ # Verify non-Inductor compile backends do not initialize Inductor.
581+ def test_non_inductor_compile_backend_matrix(self):
582+ cases = {
583+ "eager": (
584+ "",
585+ 'compiled = torch.compile(Model(), backend="eager", fullgraph=True)',
586+ ),
587+ "custom": (
588+ "custom_backend = lambda graph_module, example_inputs: "
589+ "graph_module.forward",
590+ "compiled = torch.compile(Model(), backend=custom_backend, fullgraph=True)",
591+ ),
592+ "npu": (
593+ "",
594+ 'compiled = torch.compile(Model(), backend="npu", fullgraph=True)',
595+ ),
596+ }
597+ for name, (backend_definition, compile_call) in cases.items():
598+ with self.subTest(name=name):
599+ allow_missing_torchair = name == "npu"
600+ self.run_in_subprocess(
601+ f"""
602+ import sys
603+ import torch
604+ import torch_npu
605+ 
606+ class Model(torch.nn.Module):
607+ def forward(self, x):
608+ return torch.sin(x) + 1
609+ 
610+ {backend_definition}
611+ x = torch.randn(8, device="npu")
612+ try:
613+ {compile_call}
614+ except AssertionError as error:
615+ assert {allow_missing_torchair!r}
616+ assert "Could not find module torchair" in str(error)
617+ else:
618+ torch.testing.assert_close(compiled(x), Model()(x))
619+ 
620+ from torch_npu.utils import _dynamo
621+ 
622+ assert _dynamo._lazy_dynamo_setup.has_run
623+ assert not _dynamo._lazy_inductor_setup.has_run
624+ assert "torch_npu._inductor" not in sys.modules
625+ """
626+ )
627+ 
628+ # Verify ONNX Dynamo Export initializes only the NPU Dynamo integration.
629+ @unittest.skipUnless(importlib.util.find_spec("onnxscript"), "requires onnxscript")
630+ def test_npu_onnx_dynamo_export_initialization_chain(self):
631+ cases = {
632+ "module_export": (
633+ "",
634+ "result = torch.onnx.export(Model(), (x,), dynamo=True)",
635+ ),
636+ "prebound_export": (
637+ "from torch.onnx import export as onnx_export",
638+ "result = onnx_export(Model(), (x,), dynamo=True)",
639+ ),
640+ }
641+ if hasattr(torch.onnx, "dynamo_export"):
642+ cases["prebound_legacy_dynamo_export"] = (
643+ "from torch.onnx import dynamo_export as onnx_export",
644+ "result = onnx_export(Model(), x)",
645+ )
646+ for name, (pre_import, export_call) in cases.items():
647+ with self.subTest(name=name):
648+ self.run_in_subprocess(
649+ f"""
650+ import sys
651+ import torch
652+ {pre_import}
653+ import torch_npu
654+ 
655+ # ONNXScript 0.4.0 cannot version-convert models containing
656+ # functions. Conversion runs after Dynamo capture, which is
657+ # the boundary this test verifies, so isolate that unrelated
658+ # compatibility issue without weakening the import assertions.
659+ from torch.onnx._internal._lazy_import import onnxscript_apis
660+ onnxscript_apis.convert_version = lambda model, target: model
661+ 
662+ class Model(torch.nn.Module):
663+ def forward(self, x):
664+ return torch.sin(x) + 1
665+ 
666+ x = torch.randn(8, device="npu")
667+ {export_call}
668+ 
669+ from torch_npu.utils import _dynamo
670+ 
671+ assert result is not None
672+ assert _dynamo._lazy_dynamo_setup.has_run
673+ assert not _dynamo._lazy_inductor_setup.has_run
674+ assert "torch_npu._inductor" not in sys.modules
675+ """
676+ )
677+ 
678+ # Verify a real NPU Inductor compile initializes the full stack.
679+ @unittest.skipUnless(importlib.util.find_spec("triton"), "requires triton-ascend")
680+ def test_npu_inductor_initialization_chain(self):
681+ self.run_in_subprocess(
682+ """
683+ import torch
684+ import torch_npu
685+ 
686+ def fn(x):
687+ return torch.sin(x) + 1
688+ 
689+ x = torch.randn(8, device="npu")
690+ actual = torch.compile(fn, backend="inductor", fullgraph=True)(x)
691+ 
692+ from torch._dynamo.device_interface import get_interface_for_device
693+ from torch_npu.utils import _dynamo
694+ 
695+ torch.testing.assert_close(actual, fn(x))
696+ assert _dynamo._lazy_dynamo_setup.has_run
697+ assert _dynamo._lazy_inductor_setup.has_run
698+ assert get_interface_for_device("npu").device_count() > 0
699+ """
700+ )
701+ 
702+ # Verify a real NPUGraphs compile initializes the full stack.
703+ @unittest.skipUnless(importlib.util.find_spec("triton"), "requires triton-ascend")
704+ def test_npu_npugraphs_initialization_chain(self):
705+ self.run_in_subprocess(
706+ """
707+ import torch
708+ import torch_npu
709+ 
710+ def fn(x):
711+ return torch.sin(x) + 1
712+ 
713+ x = torch.randn(8, device="npu")
714+ actual = torch.compile(fn, backend="npugraphs", fullgraph=True)(x)
715+ 
716+ from torch._dynamo.device_interface import get_interface_for_device
717+ from torch_npu.utils import _dynamo
718+ 
719+ torch.testing.assert_close(actual, fn(x))
720+ assert _dynamo._lazy_dynamo_setup.has_run
721+ assert _dynamo._lazy_inductor_setup.has_run
722+ assert get_interface_for_device("npu").device_count() > 0
723+ """
724+ )
725+ 
726+ # Verify lazy setup completes before compile backend lookup.
727+ def test_compile_triggers_setup_before_backend_lookup(self):
728+ self.run_in_subprocess(
729+ """
730+ import torch
731+ import torch_npu
732+ from torch_npu.utils import _dynamo
733+ 
734+ calls = []
735+ 
736+ @_dynamo.run_once
737+ def fake_setup():
738+ calls.append("setup")
739+ 
740+ _dynamo._lazy_dynamo_setup = fake_setup
741+ from torch._dynamo.backends import registry
742+ assert fake_setup.has_run
743+ 
744+ compiled = torch.compile(lambda x: x + 1, backend="eager")
745+ assert compiled(torch.tensor(1)).item() == 2
746+ assert calls == ["setup"]
747+ """
748+ )
749+ 
750+ # Verify the trigger works when Dynamo was imported first.
751+ def test_trigger_after_dynamo_was_preimported(self):
752+ self.run_in_subprocess(
753+ """
754+ import sys
755+ import torch
756+ import torch._dynamo
757+ import torch_npu
758+ from torch_npu.utils import _dynamo
759+ 
760+ assert _dynamo._lazy_dynamo_setup.has_run
761+ 
762+ compiled = torch.compile(lambda x: x + 1, backend="eager", fullgraph=True)
763+ assert compiled(torch.tensor(1)).item() == 2
764+ assert _dynamo._lazy_dynamo_setup.has_run
765+ assert not _dynamo._lazy_inductor_setup.has_run
766+ assert "torch_npu._inductor" not in sys.modules
767+ """
768+ )
769+ 
770+ # Verify the moved helper preserves the torch_npu v2.9 device predicates.
771+ def test_has_triton_preserves_target_device_semantics(self):
772+ self.run_in_subprocess(
773+ """
774+ import sys
775+ import types
776+ from unittest import mock
777+ 
778+ import torch
779+ import torch_npu
780+ from torch._dynamo import device_interface
781+ from torch_npu.utils import _dynamo
782+ 
783+ def make_interface(available=False):
784+ class Worker:
785+ @staticmethod
786+ def get_device_properties():
787+ raise AssertionError("v2.9 NPU override must not query CUDA props")
788+ 
789+ class Interface:
790+ @staticmethod
791+ def is_available():
792+ return available
793+ 
794+ Interface.Worker = Worker
795+ return Interface
796+ 
797+ def run_case(available, *, package=True):
798+ interfaces = {
799+ device: make_interface(
800+ available=device in available,
801+ )
802+ for device in ("cuda", "xpu", "cpu", "npu")
803+ }
804+ _dynamo.has_triton.cache_clear()
805+ with mock.patch.object(
806+ torch.utils._triton,
807+ "has_triton_package",
808+ return_value=package,
809+ ), mock.patch.object(
810+ device_interface,
811+ "get_interface_for_device",
812+ side_effect=interfaces.__getitem__,
813+ ), mock.patch.object(
814+ _dynamo,
815+ "_dynamo_register_interface_for_device",
816+ ) as register:
817+ result = _dynamo.has_triton()
818+ if package:
819+ register.assert_called_once_with()
820+ else:
821+ register.assert_not_called()
822+ return result
823+ 
824+ assert not run_case({}, package=False)
825+ assert run_case({"cuda"})
826+ assert run_case({"xpu"})
827+ triton = types.ModuleType("triton")
828+ triton_backends = types.ModuleType("triton.backends")
829+ triton_backends.backends = {"cpu": object()}
830+ triton.backends = triton_backends
831+ with mock.patch.dict(
832+ sys.modules,
833+ {"triton": triton, "triton.backends": triton_backends},
834+ ):
835+ assert run_case({"cpu"})
836+ triton_backends.backends = {}
837+ assert not run_case({"cpu"})
838+ assert run_case({"npu"})
839+ """
840+ )
841+ 
842+ # Verify all legacy Dynamo patches remain installed exactly once.
843+ def test_dynamo_patch_inventory_is_preserved(self):
844+ self.run_in_subprocess(
845+ """
846+ import torch
847+ import torch_npu
848+ 
849+ # Importing the Dynamo parent package is the lazy setup boundary.
850+ import torch._dynamo
851+ 
852+ from torch._dynamo.device_interface import get_interface_for_device
853+ from torch._dynamo.variables.builtin import BuiltinVariable
854+ from torch._dynamo.variables.builder import VariableBuilder
855+ from torch._dynamo.variables.ctx_manager import EventVariable
856+ from torch._dynamo.variables.functions import SkipFunctionVariable
857+ from torch._dynamo.variables.tensor import TensorVariable
858+ from torch._dynamo.variables.torch import constant_fold_functions
859+ from torch._dynamo.variables.user_defined import UserDefinedClassVariable
860+ from torch._dynamo.utils import common_constant_types
861+ from torch_npu.dynamo.trace_rule import (
862+ skip_functions_npu,
863+ torch_c_binding_in_graph_functions_npu,
864+ torch_non_c_binding_in_graph_functions_npu,
865+ )
866+ from torch_npu.utils import _dynamo
867+ 
868+ assert _dynamo._lazy_dynamo_setup.has_run
869+ assert get_interface_for_device("npu").device_count() > 0
870+ 
871+ # VariableTracker and context-manager patches formerly installed
872+ # eagerly by add_dynamo_methods().
873+ assert SkipFunctionVariable.__new__.__module__ == "torch_npu.utils._dynamo"
874+ assert TensorVariable.call_method.__module__ == "torch_npu.utils._dynamo"
875+ assert UserDefinedClassVariable.__new__.__module__ == "torch_npu.utils._dynamo"
876+ in_graph_classes = UserDefinedClassVariable._in_graph_classes()
877+ assert torch.npu.Event in in_graph_classes
878+ assert torch.npu.Stream in in_graph_classes
879+ assert torch.npu.fake_record_stream is _dynamo.fake_record_stream
880+ assert TensorVariable.method_record_stream.__module__ == "torch_npu.utils._dynamo"
881+ assert VariableBuilder._wrap.__module__ == "torch_npu.utils._dynamo"
882+ assert BuiltinVariable.call_id.__module__ == "torch_npu.utils._dynamo"
883+ assert EventVariable.python_type.__module__ == "torch_npu.utils._dynamo"
884+ assert torch._dynamo.optimize.__module__ == "torch_npu.utils._dynamo"
885+ 
886+ # Backend and trace-rule registrations formerly performed by
887+ # registry_manager._register_dynamo(). Count the maps as well as
888+ # checking membership so repeated lazy triggers cannot hide a
889+ # duplicate installation.
890+ assert {"npu", "npugraph_ex"}.issubset(
891+ torch._dynamo.list_backends(exclude_tags=None)
892+ )
893+ maps = torch._dynamo.trace_rules.torch_name_rule_map
894+ assert maps.count(torch_non_c_binding_in_graph_functions_npu) == 1
895+ assert maps.count(torch_c_binding_in_graph_functions_npu) == 1
896+ assert maps.count(skip_functions_npu) == 1
897+ assert constant_fold_functions[torch.npu.current_device]
898+ assert constant_fold_functions[torch.npu.get_device_properties]
899+ assert constant_fold_functions[torch.npu.is_available]
900+ assert torch_npu._C._NPUDeviceProperties in common_constant_types
901+ """
902+ )
903+ 
904+ # Verify all legacy Inductor patches remain installed.
905+ def test_inductor_patch_inventory_is_preserved(self):
906+ self.run_in_subprocess(
907+ """
908+ import torch
909+ import torch_npu
910+ 
911+ # RNG/decomposition patches remain installed at import time.
912+ from torch_npu.utils import _inductor
913+ 
914+ assert (
915+ torch._decomp.decompositions._max_unpoolnd
916+ is _inductor._max_unpoolnd_patch
917+ )
918+ assert torch._prims.rng_prims.philox_rand_offset.__module__ == (
919+ "torch_npu.utils._inductor"
920+ )
921+ assert torch._prims.rng_prims.register_philox_rand.__module__ == (
922+ "torch_npu.utils._inductor"
923+ )
924+ assert torch._prims.rng_prims.get_device.__module__ == (
925+ "torch_npu.utils._inductor"
926+ )
927+ 
928+ # Exercise the new full-Inductor setup boundary without relying on
929+ # test ordering or on a prior torch.compile invocation.
930+ import torch._dynamo
931+ from torch_npu.utils import _dynamo
932+ 
933+ _dynamo._lazy_inductor_setup()
934+ 
935+ import torch._inductor.compile_fx as compile_fx
936+ import torch._inductor.cudagraph_trees as cudagraph_trees
937+ import torch._inductor.cudagraph_utils as cudagraph_utils
938+ import torch._inductor.scheduler as scheduler
939+ from torch._inductor.codegen.common import get_device_op_overrides
940+ from torch._inductor.codecache import CacheBase
941+ from torch._inductor.graph import GraphLowering
942+ from torch._inductor.utils import GPU_TYPES
943+ from torch_npu.utils import _graph_tree
944+ 
945+ assert _dynamo._lazy_inductor_setup.has_run
946+ assert "npu" in GPU_TYPES
947+ assert get_device_op_overrides("npu").__class__.__module__.startswith(
948+ "torch_npu._inductor"
949+ )
950+ assert torch.utils._triton.has_triton.__module__ == (
951+ "torch_npu.utils._dynamo"
952+ )
953+ assert torch.utils._triton._device_supports_tma.__module__ == (
954+ "torch_npu._inductor.utils"
955+ )
956+ assert compile_fx.has_triton is torch.utils._triton.has_triton
957+ assert scheduler.has_triton is torch.utils._triton.has_triton
958+ assert GraphLowering.codegen_with_cpp_wrapper.__module__ == (
959+ "torch_npu._inductor.graph"
960+ )
961+ assert CacheBase.get_system.__module__ == (
962+ "torch_npu._inductor.codegen.common"
963+ )
964+ 
965+ # NPUGraph integrations were formerly applied eagerly alongside
966+ # the Inductor patches.
967+ assert compile_fx.cudagraphify is _graph_tree.npugraphify
968+ assert (
969+ cudagraph_utils.check_multiple_devices_or_any_cpu_nodes
970+ is _graph_tree.check_multiple_devices_or_any_cpu_nodes
971+ )
972+ assert cudagraph_trees.get_manager.__module__ == (
973+ "torch_npu.utils._graph_tree"
974+ )
975+ assert torch.compiler.npugraph_mark_step_begin is (
976+ _graph_tree.npugraph_mark_step_begin
977+ )
978+ 
979+ config = torch._inductor.config
980+ assert config.npu_backend == "default"
981+ assert config.enable_shape_handling is False
982+ assert config.shape_handling_configs == []
983+ assert config.shape_handling_dict is None
984+ """
985+ )
986+ 
987+ 
988+if __name__ == "__main__":
989+ 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_getCurrentRawStreamNoWait24 from torch_npu._C import _npu_getCurrentRawStream, _npu_getCurrentRawStreamNoWait
25 from torch._dynamo.device_interface import get_interface_for_device25 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+60-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 
42EXPECTED_NOT_LOADED_MODULES = [42EXPECTED_NOT_LOADED_MODULES = [
43 "torch_npu._C._afd",43 "torch_npu._C._afd",
44+ "torch_npu._inductor",
44]45]
45 46 
46EXPECTED_TOP_LEVEL_ATTRS = [47EXPECTED_TOP_LEVEL_ATTRS = [
@@ -255,20 +256,6 @@ class TestTorchNpuBootstrap(TestCase):
255 import torch.distributed as dist256 import torch.distributed as dist
256 import torch.distributed.rpc as rpc257 import torch.distributed.rpc as rpc
257 import torch.distributed.tensor # noqa: F401258 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_list260 assert "hccl" in dist.Backend.backend_list
274 assert "lccl" in dist.Backend.backend_list261 assert "lccl" in dist.Backend.backend_list
@@ -588,5 +575,63 @@ class TestTorchNpuBootstrap(TestCase):
588 """575 """
589 )576 )
590 577 
578+ def test_13_dtensor_strategies_are_registered_without_compiler_imports(self):
579+ self._run_python(
580+ """
581+ import sys
582+ import torch
583+ import torch_npu
584+ from torch.distributed.tensor import DTensor
585+ 
586+ strategy_funcs = DTensor._op_dispatcher.sharding_propagator.op_strategy_funcs
587+ assert torch.ops.npu.npu_rms_norm.default in strategy_funcs
588+ assert torch.ops.npu.npu_fusion_attention.default in strategy_funcs
589+ 
590+ assert "torch_npu.distributed.tensor" in sys.modules
591+ assert "torch_npu.distributed.tensor.experimental" in sys.modules
592+ assert "torch._dynamo" not in sys.modules
593+ assert "torch._inductor" not in sys.modules
594+ 
595+ import os
596+ import tempfile
597+ import torch.distributed as dist
598+ from torch.distributed.device_mesh import DeviceMesh
599+ from torch.distributed.tensor import Replicate
600+ from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
601+ from torch.distributed.tensor._op_schema import OpSchema, OpSpec, OpStrategy
602+ 
603+ fd, path = tempfile.mkstemp()
604+ os.close(fd)
605+ os.unlink(path)
606+ dist.init_process_group(
607+ "gloo", init_method=f"file://{path}", rank=0, world_size=1
608+ )
609+ try:
610+ mesh = DeviceMesh("cpu", [0])
611+ tensor_meta = TensorMeta(
612+ torch.Size([2, 4, 8]), (32, 8, 1), torch.float32
613+ )
614+ spec = DTensorSpec(mesh, (Replicate(),), tensor_meta=tensor_meta)
615+ strategy = OpStrategy([OpSpec(output_specs=spec)])
616+ op = torch.ops.npu.npu_rotary_mul.default
617+ propagator = DTensor._op_dispatcher.sharding_propagator
618+ op_schema = OpSchema(
619+ op,
620+ (strategy, strategy, strategy, "half"),
621+ {},
622+ propagator.op_to_schema_info[op],
623+ )
624+ result = propagator.op_strategy_funcs[op](op_schema)
625+ assert len(result.strategies) == 3
626+ finally:
627+ dist.destroy_process_group()
628+ if os.path.exists(path):
629+ os.remove(path)
630+ 
631+ from torch_npu.distributed.tensor import experimental
632+ assert callable(experimental.context_parallel)
633+ """
634+ )
635+ 
591if __name__ == "__main__":636if __name__ == "__main__":
592 run_tests()637 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/decomposition.py+1-0
@@ -32,6 +32,7 @@ def _register_triton_decompositions():
32 aten.nll_loss_forward,32 aten.nll_loss_forward,
33 aten.nll_loss_backward,33 aten.nll_loss_backward,
34 aten._log_softmax_backward_data,34 aten._log_softmax_backward_data,
35+ aten.erfc,
35 aten.gelu,36 aten.gelu,
36 aten.native_layer_norm,37 aten.native_layer_norm,
37 aten.slice_backward,38 aten.slice_backward,
Mtorch_npu/_inductor/kernel/bmm.py+3-2
@@ -23,6 +23,7 @@ from torch._inductor.kernel.mm_common import (
23 addmm_epilogue,23 addmm_epilogue,
24 mm_args,24 mm_args,
25)25)
26+from torch._inductor.kernel import bmm as inductor_bmm
26 27 
27from .mm import is_contiguous_striding28from .mm import is_contiguous_striding
28from ..utils import use_catlass_template, use_triton_template29from ..utils import use_catlass_template, use_triton_template
@@ -31,8 +32,8 @@ from ..utils import use_catlass_template, use_triton_template
31log = logging.getLogger("torch._inductor")32log = logging.getLogger("torch._inductor")
32aten = torch.ops.aten33aten = torch.ops.aten
33 34 
34-aten_bmm = torch._inductor.kernel.bmm.aten_bmm35+aten_bmm = inductor_bmm.aten_bmm
35-aten_baddbmm = torch._inductor.kernel.bmm.aten_baddbmm36+aten_baddbmm = inductor_bmm.aten_baddbmm
36 37 
37 38 
38def is_batch_stride_largest_or_zero(mat1, mat2, layout) -> bool:39def is_batch_stride_largest_or_zero(mat1, mat2, layout) -> bool:
Mtorch_npu/_inductor/utils.py+2-2
@@ -5,7 +5,6 @@ import logging
5import os5import os
6 6 
7import torch7import torch
8-import torch_npu
9 8 
10import torch._inductor.config as inductor_config9import torch._inductor.config as inductor_config
11log = logging.getLogger("torch._inductor")10log = logging.getLogger("torch._inductor")
@@ -28,10 +27,11 @@ def resolve_npu_device_index(device_idx=None) -> int:
28 27 
29 28 
30def patch_has_triton():29def patch_has_triton():
30+ from torch._inductor import compile_fx
31 from torch_npu.utils._dynamo import has_triton31 from torch_npu.utils._dynamo import has_triton
32 32 
33 torch._inductor.scheduler.has_triton = has_triton33 torch._inductor.scheduler.has_triton = has_triton
34- torch._inductor.compile_fx.has_triton = has_triton34+ compile_fx.has_triton = has_triton
35 35 
36 36 
37def patch_device_supports_tma():37def 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_methods6 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-22
@@ -49,25 +49,6 @@ def _register_distributed():
49 register_distributed_backend_for_npu()49 register_distributed_backend_for_npu()
50 50 
51 51 
52-def _register_dynamo():
53- """
54- Register Dynamo integration:
55- - Dynamo backend
56- - Dynamo device interface
57- - NPU trace rules for Dynamo
58- """
59- from torch_npu._init.registry.dynamo import (
60- register_dynamo_backends,
61- register_dynamo_trace_rules,
62- )
63- 
64- register_dynamo_backends()
65- 
66- # Do not repeat this call for register_dynamo_trace_rules appends rules into
67- # Dynamo's global rules maps.
68- register_dynamo_trace_rules()
69- 
70- 
71def _register_rpc():52def _register_rpc():
72 """53 """
73 Register and init RPC NPU backend.54 Register and init RPC NPU backend.
@@ -100,8 +81,8 @@ def _register_components():
100 81 
101 Order matters:82 Order matters:
102 1. NPU backend is the base capability.83 1. NPU backend is the base capability.
103- 2. Distributed and Dynamo depend on NPU backend / _C children.84+ 2. Distributed depends on the NPU backend / _C children.
104- 3. RPC, dtensor and inductor are Python-side framework integrations.85+ 3. RPC and inductor are Python-side framework integrations.
105 4. DefaultDeviceType is set after NPU backend is registered.86 4. DefaultDeviceType is set after NPU backend is registered.
106 """87 """
107 if not hasattr(torch_npu, "_C"):88 if not hasattr(torch_npu, "_C"):
@@ -111,7 +92,6 @@ def _register_components():
111 92 
112 _register_npu_backend()93 _register_npu_backend()
113 _register_distributed()94 _register_distributed()
114- _register_dynamo()
115 _register_rpc()95 _register_rpc()
116 _register_inductor()96 _register_inductor()
117 _register_default_gradient_device_type()97 _register_default_gradient_device_type()
Mtorch_npu/contrib/transfer_to_npu.py+13-1
@@ -436,6 +436,18 @@ def _compose_wrappers(*wrappers):
436 436 
437 437 
438def _init():438def _init():
439+ # transfer_to_npu patches these modules during its own import. Import them
440+ # explicitly instead of relying on torch_npu import side effects.
441+ import torch._dynamo.trace_rules # noqa: F401
442+ import torch._dynamo.utils # noqa: F401
443+ import torch._inductor.runtime.autotune_cache # noqa: F401
444+ import torch._inductor.compile_fx # noqa: F401
445+ import torch._inductor.utils # noqa: F401
446+ import torch._inductor.fx_passes.post_grad # noqa: F401
447+ import torch._inductor.fx_passes.joint_graph # noqa: F401
448+ import torch._inductor.autotune_process # noqa: F401
449+ from torch.distributed.checkpoint import filesystem
450+ 
439 _warning_fn('''451 _warning_fn('''
440 *************************************************************************************************************452 *************************************************************************************************************
441 The torch.Tensor.cuda and torch.nn.Module.cuda are replaced with torch.Tensor.npu and torch.nn.Module.npu now..453 The torch.Tensor.cuda and torch.nn.Module.cuda are replaced with torch.Tensor.npu and torch.nn.Module.npu now..
@@ -540,7 +552,7 @@ def _init():
540 setattr(torch._inductor.autotune_process, "get_gpu_type", _get_npu_type)552 setattr(torch._inductor.autotune_process, "get_gpu_type", _get_npu_type)
541 553 
542 setattr(torch._utils, '_get_available_device_type', _patch_get_available_device_type)554 setattr(torch._utils, '_get_available_device_type', _patch_get_available_device_type)
543- setattr(torch.distributed.checkpoint.filesystem._OverlappingCpuLoader, '__init__',555+ setattr(filesystem._OverlappingCpuLoader, '__init__',
544 _patch_OverlappingCpuLoader_init_)556 _patch_OverlappingCpuLoader_init_)
545 557 
546 _replace_to_method_in_allowed_methods()558 _replace_to_method_in_allowed_methods()
Mtorch_npu/dynamo/__init__.py+71-5
@@ -3,9 +3,6 @@ import sys
3import time3import time
4import warnings4import warnings
5 5 
6-from torch._dynamo import register_backend as _register_backend
7-from torch._dynamo.backends.registry import _BACKENDS
8- 
9from torch_npu._init.common.warning_utils import _should_print_warning6from torch_npu._init.common.warning_utils import _should_print_warning
10from torch_npu.utils._error_code import ErrCode, pta_error7from torch_npu.utils._error_code import ErrCode, pta_error
11 8 
@@ -107,6 +104,15 @@ class _LazyTorchair(_LazyBackend):
107 return getattr(torchair, name)104 return getattr(torchair, name)
108 105 
109 106 
107+def _install_lazy_torchair():
108+ torchair_path = os.path.join(os.path.dirname(__file__), "torchair")
109+ if not os.path.exists(torchair_path):
110+ return False
111+ if "torchair" not in sys.modules:
112+ sys.modules["torchair"] = _LazyTorchair("torchair")
113+ return True
114+ 
115+ 
110class _LazyNpuGraphEx(_LazyBackend):116class _LazyNpuGraphEx(_LazyBackend):
111 def __init__(self, pkg_name):117 def __init__(self, pkg_name):
112 self._npugraph_ex = None118 self._npugraph_ex = None
@@ -140,7 +146,7 @@ def _lazy_exec(*args, **kwargs):
140 146 
141 147 
142def _get_default_backend(name):148def _get_default_backend(name):
143- if not os.path.exists(os.path.join(os.path.dirname(__file__), 'torchair')):149+ if not _install_lazy_torchair():
144 if _should_print_warning():150 if _should_print_warning():
145 warnings.warn(151 warnings.warn(
146 "Register eager implementation for the 'npu' backend of dynamo, "152 "Register eager implementation for the 'npu' backend of dynamo, "
@@ -149,7 +155,6 @@ def _get_default_backend(name):
149 155 
150 global _global_backend_name156 global _global_backend_name
151 _global_backend_name = name157 _global_backend_name = name
152- sys.modules['torchair'] = _LazyTorchair('torchair')
153 return _lazy_exec158 return _lazy_exec
154 159 
155 160 
@@ -166,6 +171,27 @@ def _get_npugraph_ex_backend():
166 171 
167 172 
168def _register_npu_backend(backend, name="npu"):173def _register_npu_backend(backend, name="npu"):
174+ from torch._dynamo import register_backend as _register_backend
175+ from torch._dynamo.backends.registry import _BACKENDS, _COMPILER_FNS
176+ 
177+ registered_backend = _COMPILER_FNS.get(name)
178+ if (
179+ registered_backend is not None
180+ and getattr(registered_backend, "__module__", None) == __name__
181+ ):
182+ return
183+ 
184+ # When a setuptools entry point is currently loading torch_npu, Dynamo has
185+ # already put the EntryPoint object in _BACKENDS but has not registered the
186+ # loaded callable yet. Leave that registration to lookup_backend(); doing
187+ # it here as well makes lookup_backend register the same name twice.
188+ pending_backend = _BACKENDS.get(name)
189+ if (
190+ name not in _COMPILER_FNS
191+ and getattr(pending_backend, "module", None) == __name__
192+ ):
193+ return
194+ 
169 if name in _BACKENDS.keys():195 if name in _BACKENDS.keys():
170 del _BACKENDS[name]196 del _BACKENDS[name]
171 _register_backend(backend, name)197 _register_backend(backend, name)
@@ -177,3 +203,43 @@ def _register_backends():
177 203 
178 _register_npu_backend(global_backend)204 _register_npu_backend(global_backend)
179 _register_npu_backend(npugraph_ex_backend, NPUGRAPH_EX_BACKEND)205 _register_npu_backend(npugraph_ex_backend, NPUGRAPH_EX_BACKEND)
206+ 
207+ 
208+def _npu_backend_entrypoint(gm, example_inputs, **kwargs):
209+ """Set up the NPU Dynamo backend when an entry point is selected."""
210+ from torch_npu.utils._dynamo import _lazy_dynamo_setup
211+ 
212+ _lazy_dynamo_setup()
213+ return _get_default_backend("npu")(gm, example_inputs, **kwargs)
214+ 
215+ 
216+def _npugraph_ex_backend_entrypoint(gm, example_inputs, **kwargs):
217+ """Set up the NPUGraph-EX backend when an entry point is selected."""
218+ from torch_npu.utils._dynamo import _lazy_dynamo_setup
219+ 
220+ _lazy_dynamo_setup()
221+ return _exec(gm, example_inputs, **kwargs)
222+ 
223+ 
224+class _NpugraphsBackendEntryPoint:
225+ """Keep the lazy entry point and its reset protocol in one callable."""
226+ 
227+ compiler_name = "npugraphs"
228+ 
229+ def __call__(self, gm, example_inputs, **kwargs):
230+ from torch_npu.utils._dynamo import _lazy_dynamo_setup, _lazy_inductor_setup
231+ 
232+ _lazy_dynamo_setup()
233+ _lazy_inductor_setup()
234+ from torch_npu.utils._graph_tree import NpugraphsBackend
235+ 
236+ return NpugraphsBackend()(gm, example_inputs, **kwargs)
237+ 
238+ @staticmethod
239+ def reset():
240+ graph_tree = sys.modules.get("torch_npu.npu._graph_tree")
241+ if graph_tree is not None:
242+ graph_tree.reset_npugraph_trees()
243+ 
244+ 
245+_npugraphs_backend_entrypoint = _NpugraphsBackendEntryPoint()
Mtorch_npu/npu/__init__.py+13-2
@@ -151,6 +151,7 @@ import traceback
151import threading151import threading
152import os152import os
153import re153import re
154+import importlib
154import torch155import torch
155from torch.storage import _LegacyStorage, _warn_typed_storage_removal156from torch.storage import _LegacyStorage, _warn_typed_storage_removal
156from torch._utils import classproperty157from torch._utils import classproperty
@@ -172,8 +173,6 @@ from .autocast_utils import * # noqa: F403
172from .backends import * # noqa: F403173from .backends import * # noqa: F403
173from ._backends import * # noqa: F403174from ._backends import * # noqa: F403
174from .deterministic import enable_deterministic_with_backward, disable_deterministic_with_backward # noqa: F403175from .deterministic import enable_deterministic_with_backward, disable_deterministic_with_backward # noqa: F403
175-from . import npugraph_ex
176- 
177from .graphs import (176from .graphs import (
178 NPUGraph,177 NPUGraph,
179 graph,178 graph,
@@ -194,6 +193,18 @@ from ._npugraph_handlers import (
194)193)
195 194 
196 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+ 
197config = npu_config._npuConfig()208config = npu_config._npuConfig()
198 209 
199matmul = npu_config._allowHF32Matmul()210matmul = npu_config._allowHF32Matmul()
Mtorch_npu/npu/_graph_tree.py+1-12
@@ -102,6 +102,7 @@ from torch.utils._ordered_set import OrderedSet
102from torch.utils.weak import TensorWeakRef102from torch.utils.weak import TensorWeakRef
103 103 
104import torch_npu104import torch_npu
105+from torch_npu.npu._graph_tree_state import MarkStepBox, mark_step_begin
105from torch_npu.npu import graphs as _npu_graphs106from torch_npu.npu import graphs as _npu_graphs
106from torch_npu._C import (107from torch_npu._C import (
107 _npu_NPUAllocator_AllocatorState as AllocatorState,108 _npu_NPUAllocator_AllocatorState as AllocatorState,
@@ -288,24 +289,12 @@ local.npu_tree_manager_containers = {}
288local.npu_tree_manager_locks = defaultdict(threading.Lock)289local.npu_tree_manager_locks = defaultdict(threading.Lock)
289 290 
290 291 
291-# only incremented by user call of mark_step_begin
292-class MarkStepBox:
293- mark_step_counter = 0
294- 
295- 
296# We need to register this as an object that will be copied over as TLS when new292# We need to register this as an object that will be copied over as TLS when new
297# threads are created in autograd293# threads are created in autograd
298torch._C._stash_obj_in_tls("npu_tree_manager_containers", local.npu_tree_manager_containers)294torch._C._stash_obj_in_tls("npu_tree_manager_containers", local.npu_tree_manager_containers)
299torch._C._stash_obj_in_tls("npu_tree_manager_locks", local.npu_tree_manager_locks)295torch._C._stash_obj_in_tls("npu_tree_manager_locks", local.npu_tree_manager_locks)
300 296 
301 297 
302-def mark_step_begin() -> None:
303- "Indicates that a new iteration of inference or training is about to begin."
304- 
305- # iterate down to distinguish from GenerationTracking counter
306- MarkStepBox.mark_step_counter -= 1
307- 
308- 
309def reset_npugraph_trees() -> None:298def reset_npugraph_trees() -> None:
310 "Clear all npugraph trees"299 "Clear all npugraph trees"
311 # see shutdown below for why this is necessary300 # 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 
3from torch import Tensor3from torch import Tensor
4from torch.autograd.function import Function4from 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+ 
10class _DeterministicAlgorithmsBeginOp(Function):18class _DeterministicAlgorithmsBeginOp(Function):
11 19 
12 @staticmethod20 @staticmethod
@@ -36,11 +44,11 @@ class _DeterministicAlgorithmsEndOp(Function):
36 return grad_outputs44 return grad_outputs
37 45 
38 46 
39-@forbid_in_graph47+@_forbid_in_graph
40def enable_deterministic_with_backward(tensor: Tensor):48def enable_deterministic_with_backward(tensor: Tensor):
41 return _DeterministicAlgorithmsBeginOp.apply(tensor)49 return _DeterministicAlgorithmsBeginOp.apply(tensor)
42 50 
43 51 
44-@forbid_in_graph52+@_forbid_in_graph
45def disable_deterministic_with_backward(tensor: Tensor):53def disable_deterministic_with_backward(tensor: Tensor):
46 return _DeterministicAlgorithmsEndOp.apply(tensor)54 return _DeterministicAlgorithmsEndOp.apply(tensor)
Mtorch_npu/npu/npugraph_ex/__init__.py+1-1
@@ -31,4 +31,4 @@ def register_replacement(search_fn: SearchFn, replace_fn: ReplaceFn, example_inp
31 trace_fn=trace_fn, extra_check=extra_check,31 trace_fn=trace_fn, extra_check=extra_check,
32 search_fn_pattern=search_fn_pattern,32 search_fn_pattern=search_fn_pattern,
33 scalar_workaround=scalar_workaround,33 scalar_workaround=scalar_workaround,
34- skip_duplicates=skip_duplicates)34+ skip_duplicates=skip_duplicates)
Mtorch_npu/utils/_dynamo.py+273-61
@@ -1,9 +1,12 @@
1-import inspect
2-import os
3-import logging
4-import sys
5import importlib1import importlib
2+import importlib.abc
6import functools3import functools
4+import inspect
5+import logging
6+import os
7+import sys
8+import threading
9+ 
7import torch10import torch
8import torch_npu11import torch_npu
9from torch import _TorchCompileWrapper12from torch import _TorchCompileWrapper
@@ -13,10 +16,10 @@ use_jit_script = False
13log = logging.getLogger(__name__)16log = logging.getLogger(__name__)
14 17 
15 18 
16- 
17def _create_npu_autocast_mode_variable(func, args, kwargs):19def _create_npu_autocast_mode_variable(func, args, kwargs):
18- from torch._dynamo.variables.ctx_manager import AutocastModeVariable
19 from torch._dynamo.variables.base import VariableTracker20 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,12 +37,13 @@ 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 var41+ 
39 42 
40def patch_SkipFunctionVariable():43def patch_SkipFunctionVariable():
41 from torch._dynamo.variables.functions import SkipFunctionVariable44 from torch._dynamo.variables.functions import SkipFunctionVariable
42 from torch._dynamo.variables.torch import TorchInGraphFunctionVariable45 from torch._dynamo.variables.torch import TorchInGraphFunctionVariable
46+ 
43 def SkipFunctionVariable__new__(cls, value, reason=None, **kwargs):47 def SkipFunctionVariable__new__(cls, value, reason=None, **kwargs):
44 if value in [48 if value in [
45 torch.npu.stream,49 torch.npu.stream,
@@ -52,11 +56,12 @@ def patch_SkipFunctionVariable():
52 SkipFunctionVariable.__new__raw = SkipFunctionVariable.__new__56 SkipFunctionVariable.__new__raw = SkipFunctionVariable.__new__
53 SkipFunctionVariable.__new__ = SkipFunctionVariable__new__57 SkipFunctionVariable.__new__ = SkipFunctionVariable__new__
54 58 
59+ 
55def patch_TensorVariable_call_method():60def patch_TensorVariable_call_method():
56- from torch._dynamo.variables.tensor import TensorVariable
57 from torch._dynamo.utils import tensortype_to_dtype61 from torch._dynamo.utils import tensortype_to_dtype
58 from torch._dynamo.variables.constant import ConstantVariable62 from torch._dynamo.variables.constant import ConstantVariable
59 from torch._dynamo.variables.lists import TupleVariable63 from torch._dynamo.variables.lists import TupleVariable
64+ from torch._dynamo.variables.tensor import TensorVariable
60 65 
61 def TensorVariable_call_method(self, tx, name, args, kwargs):66 def TensorVariable_call_method(self, tx, name, args, kwargs):
62 if (67 if (
@@ -66,25 +71,30 @@ def patch_TensorVariable_call_method():
66 and isinstance(self.device, torch.device)71 and isinstance(self.device, torch.device)
67 and self.device.type == "npu"72 and self.device.type == "npu"
68 ):73 ):
69- tensortype = next(k for k, v in tensortype_to_dtype.items() if self.dtype in v)74+ tensortype = next(
70- 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+ )
71 80 
72 if len(args) == 1:81 if len(args) == 1:
73 return constant_result.getitem_const(args[0])82 return constant_result.getitem_const(args[0])
74- elif args:83+ if args:
75- 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+ )
76 return constant_result87 return constant_result
77- else:88+ return TensorVariable.call_method_raw(self, tx, name, args, kwargs)
78- return TensorVariable.call_method_raw(self, tx, name, args, kwargs)
79 89 
80 TensorVariable.call_method_raw = TensorVariable.call_method90 TensorVariable.call_method_raw = TensorVariable.call_method
81 TensorVariable.call_method = TensorVariable_call_method91 TensorVariable.call_method = TensorVariable_call_method
82 92 
83 93 
84- 
85class _InductorNpuRegistry:94class _InductorNpuRegistry:
86 _disabled_register = False95 _disabled_register = False
87 _loaded_backend = None96 _loaded_backend = None
97+ 
88 @classmethod98 @classmethod
89 def register_inductor_npu(cls):99 def register_inductor_npu(cls):
90 if cls._disabled_register:100 if cls._disabled_register:
@@ -128,19 +138,23 @@ def register_inductor_npu():
128 _InductorNpuRegistry.register_inductor_npu()138 _InductorNpuRegistry.register_inductor_npu()
129 139 
130 140 
131-def _resolve_npu_backend_from_wrapper(wrapper) -> str:141+def _resolve_npu_backend(selected_backend=None) -> str:
132- """Resolve npu backend with priority: wrapper options > global config > env."""142+ """Resolve NPU backend with priority: compile options > config > env."""
133- wrapper_backend = wrapper.config.get("npu_backend")143+ if selected_backend not in (None, "", "default"):
134- if wrapper_backend not in (None, "", "default"):144+ return selected_backend
135- return wrapper_backend
136 145 
137- global_backend = getattr(torch._inductor.config, "npu_backend", None)146+ inductor_config = sys.modules.get("torch._inductor.config")
147+ global_backend = getattr(inductor_config, "npu_backend", None)
138 if global_backend not in (None, "", "default"):148 if global_backend not in (None, "", "default"):
139 return global_backend149 return global_backend
140 150 
141 return os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default")151 return os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default")
142 152 
143 153 
154+def _resolve_npu_backend_from_wrapper(wrapper) -> str:
155+ return _resolve_npu_backend(wrapper.config.get("npu_backend"))
156+ 
157+ 
144class _NpuBackendScope:158class _NpuBackendScope:
145 """Apply resolved npu backend for one compile invocation and restore env."""159 """Apply resolved npu backend for one compile invocation and restore env."""
146 160 
@@ -150,32 +164,37 @@ class _NpuBackendScope:
150 164 
151 def __enter__(self):165 def __enter__(self):
152 self._old_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")166 self._old_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")
153- os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend167+ try:
154- register_inductor_npu()168+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend
155- if self.backend == "ascendc":169+ register_inductor_npu()
156- from torch_npu._inductor.deterministic_cache import (170+ if self.backend == "ascendc":
157- patch_npu_deterministic_level_cache_keys,171+ from torch_npu._inductor.deterministic_cache import (
158- )172+ patch_npu_deterministic_level_cache_keys,
173+ )
159 174 
160- patch_npu_deterministic_level_cache_keys()175+ patch_npu_deterministic_level_cache_keys()
176+ except BaseException:
177+ self._restore_backend_env()
178+ raise
161 return self179 return self
162 180 
163 def __exit__(self, exc_type, exc, tb):181 def __exit__(self, exc_type, exc, tb):
182+ self._restore_backend_env()
183+ return False
184+ 
185+ def _restore_backend_env(self):
164 if self._old_env is None:186 if self._old_env is None:
165 os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)187 os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)
166 else:188 else:
167 os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self._old_env189 os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self._old_env
168- return False
169- 
170- 
171 190 
172 191 
173def patch_dynamo_optimize():192def patch_dynamo_optimize():
174 from torch_npu.dynamo import _get_global_npu_backend193 from torch_npu.dynamo import _get_global_npu_backend
194+ 
175 src_optimize = torch._dynamo.optimize195 src_optimize = torch._dynamo.optimize
176 196 
177 def npu_optimize(*args, **kwargs):197 def npu_optimize(*args, **kwargs):
178- add_dynamo_methods_init()
179 backend = None198 backend = None
180 if "backend" in kwargs:199 if "backend" in kwargs:
181 backend = kwargs["backend"]200 backend = kwargs["backend"]
@@ -197,11 +216,14 @@ def patch_dynamo_optimize():
197 216 
198 217 
199def patch_user_defined_class_variable():218def patch_user_defined_class_variable():
200- import functools219+ from torch._dynamo.variables.torch import (
220+ TorchCtxManagerClassVariable,
221+ TorchInGraphFunctionVariable,
222+ )
201 from torch._dynamo.variables.user_defined import UserDefinedClassVariable223 from torch._dynamo.variables.user_defined import UserDefinedClassVariable
202- from torch._dynamo.variables.torch import TorchCtxManagerClassVariable224+ 
203- from torch._dynamo.variables.torch import TorchInGraphFunctionVariable
204 original_method = UserDefinedClassVariable._in_graph_classes225 original_method = UserDefinedClassVariable._in_graph_classes
226+ 
205 class NPUTorchCtxManagerClassVariable(TorchCtxManagerClassVariable):227 class NPUTorchCtxManagerClassVariable(TorchCtxManagerClassVariable):
206 def call_function(self, tx, args, kwargs):228 def call_function(self, tx, args, kwargs):
207 return _create_npu_autocast_mode_variable(self.value, args, kwargs)229 return _create_npu_autocast_mode_variable(self.value, args, kwargs)
@@ -222,7 +244,7 @@ def patch_user_defined_class_variable():
222 torch_npu.npu.amp.autocast_mode.autocast,244 torch_npu.npu.amp.autocast_mode.autocast,
223 ]:245 ]:
224 return NPUTorchCtxManagerClassVariable(value, **kwargs)246 return NPUTorchCtxManagerClassVariable(value, **kwargs)
225- elif value in [247+ if value in [
226 torch.npu.Stream,248 torch.npu.Stream,
227 torch_npu.npu.Stream,249 torch_npu.npu.Stream,
228 torch.npu.streams.Stream,250 torch.npu.streams.Stream,
@@ -245,6 +267,7 @@ def patch_user_defined_class_variable():
245 UserDefinedClassVariable.__new__raw = UserDefinedClassVariable.__new__267 UserDefinedClassVariable.__new__raw = UserDefinedClassVariable.__new__
246 UserDefinedClassVariable.__new__ = UserDefinedClassVariable__new__268 UserDefinedClassVariable.__new__ = UserDefinedClassVariable__new__
247 269 
270+ 
248def fake_record_stream(self, s):271def fake_record_stream(self, s):
249 """272 """
250 let dynamo trace Tensor.record_stream as this empty function,273 let dynamo trace Tensor.record_stream as this empty function,
@@ -322,21 +345,69 @@ def patch_event_variable_python_type():
322 if "python_type" not in torch._dynamo.variables.ctx_manager.EventVariable.__dict__:345 if "python_type" not in torch._dynamo.variables.ctx_manager.EventVariable.__dict__:
323 torch._dynamo.variables.ctx_manager.EventVariable.python_type = python_type346 torch._dynamo.variables.ctx_manager.EventVariable.python_type = python_type
324 347 
348+ 
325def run_once(f):349def run_once(f):
326- """Runs a function (successfully) only once.350+ """Run a function successfully only once, waiting for concurrent callers."""
327- The running can be reset by setting the `has_run` attribute to False351+ condition = threading.Condition()
328- """352+ 
329 @functools.wraps(f)353 @functools.wraps(f)
330 def wrapper(*args, **kwargs):354 def wrapper(*args, **kwargs):
331- if not wrapper.has_run:355+ thread_id = threading.get_ident()
356+ with condition:
357+ while wrapper._is_running:
358+ if wrapper._running_thread == thread_id:
359+ return None
360+ condition.wait()
361+ if wrapper.has_run:
362+ return None
363+ wrapper._is_running = True
364+ wrapper._running_thread = thread_id
365+ 
366+ try:
332 result = f(*args, **kwargs)367 result = f(*args, **kwargs)
368+ except BaseException:
369+ with condition:
370+ wrapper._is_running = False
371+ wrapper._running_thread = None
372+ condition.notify_all()
373+ raise
374+ 
375+ with condition:
333 wrapper.has_run = True376 wrapper.has_run = True
334- return result377+ wrapper._is_running = False
335- return None378+ wrapper._running_thread = None
379+ condition.notify_all()
380+ return result
381+ 
336 wrapper.has_run = False382 wrapper.has_run = False
383+ wrapper._is_running = False
384+ wrapper._running_thread = None
385+ 
386+ def reset_after_fork():
387+ # The parent thread running f may not exist in the child process.
388+ nonlocal condition
389+ condition = threading.Condition()
390+ wrapper._is_running = False
391+ wrapper._running_thread = None
392+ 
393+ try:
394+ os.register_at_fork(after_in_child=reset_after_fork)
395+ except AttributeError:
396+ pass
337 return wrapper397 return wrapper
338 398 
339 399 
400+_COMPLETED_DYNAMO_SETUP_STEPS = set()
401+ 
402+ 
403+def _run_dynamo_setup_step(name, setup):
404+ """Keep successful setup steps idempotent when a later step fails."""
405+ if name in _COMPLETED_DYNAMO_SETUP_STEPS:
406+ return
407+ setup()
408+ _COMPLETED_DYNAMO_SETUP_STEPS.add(name)
409+ 
410+ 
340@run_once411@run_once
341def _dynamo_register_interface_for_device():412def _dynamo_register_interface_for_device():
342 from torch._dynamo.device_interface import register_interface_for_device413 from torch._dynamo.device_interface import register_interface_for_device
@@ -344,17 +415,74 @@ def _dynamo_register_interface_for_device():
344 415 
345 register_interface_for_device("npu", NpuInterface)416 register_interface_for_device("npu", NpuInterface)
346 for i in range(32):417 for i in range(32):
347- 
348 register_interface_for_device(f"npu:{i}", NpuInterface)418 register_interface_for_device(f"npu:{i}", NpuInterface)
349 419 
420+ 
421+def _find_spec_without_finder(finder, fullname):
422+ """Delegate to the remaining meta-path finders without bypassing them."""
423+ try:
424+ index = sys.meta_path.index(finder)
425+ except ValueError:
426+ return importlib.util.find_spec(fullname)
427+ 
428+ sys.meta_path.pop(index)
429+ try:
430+ return importlib.util.find_spec(fullname)
431+ finally:
432+ sys.meta_path.insert(min(index, len(sys.meta_path)), finder)
433+ 
434+ 
435+class _DynamoPostImportLoader(importlib.abc.Loader):
436+ def __init__(self, loader, finder):
437+ self._loader = loader
438+ self._finder = finder
439+ 
440+ def create_module(self, spec):
441+ create_module = getattr(self._loader, "create_module", None)
442+ return create_module(spec) if create_module is not None else None
443+ 
444+ def exec_module(self, module):
445+ self._loader.exec_module(module)
446+ _lazy_dynamo_setup()
447+ if self._finder in sys.meta_path:
448+ sys.meta_path.remove(self._finder)
449+ 
450+ 
451+class _DynamoPostImportFinder(importlib.abc.MetaPathFinder):
452+ _target = "torch._dynamo"
453+ 
454+ def find_spec(self, fullname, path=None, target=None):
455+ if fullname != self._target:
456+ return None
457+ spec = _find_spec_without_finder(self, fullname)
458+ if spec is not None and spec.loader is not None:
459+ spec.loader = _DynamoPostImportLoader(spec.loader, self)
460+ return spec
461+ 
462+ 
463+def _install_dynamo_post_import_trigger():
464+ """Set up NPU integration whenever Dynamo is first imported."""
465+ if "torch._dynamo" in sys.modules:
466+ _lazy_dynamo_setup()
467+ return
468+ if not any(isinstance(finder, _DynamoPostImportFinder) for finder in sys.meta_path):
469+ sys.meta_path.insert(0, _DynamoPostImportFinder())
470+ 
471+ 
350@run_once472@run_once
351def add_dynamo_methods_init():473def add_dynamo_methods_init():
352- _dynamo_register_interface_for_device()474+ steps = (
353- patch_SkipFunctionVariable()475+ ("device_interface", _dynamo_register_interface_for_device),
354- patch_TensorVariable_call_method()476+ ("skip_function_variable", patch_SkipFunctionVariable),
355- patch_user_defined_class_variable()477+ ("tensor_variable", patch_TensorVariable_call_method),
356- patch_record_stream()478+ ("user_defined_class_variable", patch_user_defined_class_variable),
357- patch_event_variable_python_type()479+ ("record_stream", patch_record_stream),
480+ ("event_variable", patch_event_variable_python_type),
481+ ("variable_builder", patch_variable_builder),
482+ ("builtin_variable", patch_builtin_variable),
483+ )
484+ for name, setup in steps:
485+ _run_dynamo_setup_step(name, setup)
358 486 
359 487 
360def patch_inductor_wrapper():488def patch_inductor_wrapper():
@@ -369,15 +497,22 @@ def patch_inductor_wrapper():
369 src_call = _TorchCompileInductorWrapper.__call__497 src_call = _TorchCompileInductorWrapper.__call__
370 498 
371 def new_apply_options(self, options: Optional[dict[str, Any]]):499 def new_apply_options(self, options: Optional[dict[str, Any]]):
372- if options is not None and options.get("enable_shape_handling", False):500+ shape_handling_requested = (
501+ options is not None and options.get("enable_shape_handling", False)
502+ )
503+ src_apply_options(self, options)
504+ if shape_handling_requested:
505+ if getattr(self, "_npu_defer_shape_handling", False):
506+ self._npu_shape_handling_requested = True
507+ return
373 if not is_inductor_npu_initialized():508 if not is_inductor_npu_initialized():
374 register_inductor_npu()509 register_inductor_npu()
375 torch_npu._inductor.patch_shape_handling()510 torch_npu._inductor.patch_shape_handling()
376- src_apply_options(self, options)
377 511 
378 def new_get_config_copy(self) -> dict[str, Any]:512 def new_get_config_copy(self) -> dict[str, Any]:
379 ori_dict = src_get_config_copy(self)513 ori_dict = src_get_config_copy(self)
380- if self is not torch._inductor.config:514+ inductor_config = sys.modules.get("torch._inductor.config")
515+ if inductor_config is None or self is not inductor_config:
381 return ori_dict516 return ori_dict
382 NpuBackendType = Literal["default", "mlir", "dvm"]517 NpuBackendType = Literal["default", "mlir", "dvm"]
383 if "npu_backend" not in ori_dict:518 if "npu_backend" not in ori_dict:
@@ -406,8 +541,17 @@ def patch_inductor_wrapper():
406 return ori_dict541 return ori_dict
407 542 
408 def new_init(self, mode, options, dynamic):543 def new_init(self, mode, options, dynamic):
409- add_dynamo_methods_init()544+ self._npu_defer_shape_handling = True
410- src_init(self, mode, options, dynamic)545+ self._npu_shape_handling_requested = False
546+ try:
547+ src_init(self, mode, options, dynamic)
548+ shape_handling_requested = self._npu_shape_handling_requested
549+ finally:
550+ del self._npu_defer_shape_handling
551+ del self._npu_shape_handling_requested
552+ _setup_inductor_for_compile(self.config)
553+ if shape_handling_requested:
554+ torch_npu._inductor.patch_shape_handling()
411 backend = _resolve_npu_backend_from_wrapper(self)555 backend = _resolve_npu_backend_from_wrapper(self)
412 if backend=="mlir":556 if backend=="mlir":
413 with _NpuBackendScope(backend):557 with _NpuBackendScope(backend):
@@ -433,7 +577,6 @@ def patch_inductor_wrapper():
433 _TorchCompileInductorWrapper.apply_options = new_apply_options577 _TorchCompileInductorWrapper.apply_options = new_apply_options
434 _TorchCompileInductorWrapper.__init__ = new_init578 _TorchCompileInductorWrapper.__init__ = new_init
435 ConfigModule.get_config_copy = new_get_config_copy579 ConfigModule.get_config_copy = new_get_config_copy
436- torch._inductor.config.get_config_copy()
437 580 
438 581 
439@functools.lru_cache(None)582@functools.lru_cache(None)
@@ -475,12 +618,81 @@ def has_triton() -> bool:
475 618 
476 619 
477def patch_has_triton():620def patch_has_triton():
478- torch.utils._triton.has_triton = has_triton621+ from torch.utils import _triton
622+ 
623+ _triton.has_triton = has_triton
624+ 
625+ 
626+@run_once
627+def _inject_inductor_npu_backend_config():
628+ """Inject NPU entries into torch._inductor.config on first use."""
629+ torch._inductor.config.get_config_copy()
630+ 
631+ 
632+@run_once
633+def _lazy_dynamo_setup():
634+ """Initialize the Dynamo integration on the first graph-capture operation."""
635+ add_dynamo_methods_init()
636+ 
637+ from torch_npu.dynamo import _register_backends
638+ _run_dynamo_setup_step("backends", _register_backends)
639+ 
640+ from torch_npu.dynamo.trace_rule import _patch_npu_trace_rules
641+ _run_dynamo_setup_step("trace_rules", _patch_npu_trace_rules)
642+ 
643+ _run_dynamo_setup_step("dynamo_optimize", patch_dynamo_optimize)
644+ 
645+ 
646+@run_once
647+def _lazy_inductor_setup():
648+ """Initialize NPU Inductor support only for an Inductor-based backend."""
649+ register_inductor_npu()
650+ 
651+ from torch_npu.utils._graph_tree import _apply_npugraph_tree_methods
652+ _apply_npugraph_tree_methods()
653+ 
654+ _inject_inductor_npu_backend_config()
655+ 
656+ 
657+def _setup_inductor_for_compile(options=None):
658+ """Initialize the NPU Inductor backend selected for this compile call."""
659+ _lazy_dynamo_setup()
660+ 
661+ option_backend = options.get("npu_backend") if isinstance(options, dict) else None
662+ selected_backend = _resolve_npu_backend(option_backend)
663+ 
664+ old_backend = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")
665+ if selected_backend not in (None, "", "default"):
666+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = selected_backend
667+ try:
668+ _lazy_inductor_setup()
669+ finally:
670+ if old_backend is None:
671+ os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)
672+ else:
673+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = old_backend
674+ return selected_backend
675+ 
676+ 
677+@run_once
678+def install_npugraph_mark_step_trigger():
679+ """Expose the public NPUGraph step API without importing compiler internals."""
680+ def npugraph_mark_step_begin():
681+ from torch_npu.npu._graph_tree_state import mark_step_begin
682+ return mark_step_begin()
683+ 
684+ torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin
479 685 
480 686 
481def add_dynamo_methods():687def add_dynamo_methods():
482- patch_dynamo_optimize()
483- patch_variable_builder()
484- patch_builtin_variable()
485- patch_inductor_wrapper()
486 patch_has_triton()688 patch_has_triton()
689+ 
690+ from torch_npu.dynamo import _install_lazy_torchair
691+ 
692+ _install_lazy_torchair()
693+ _install_dynamo_post_import_trigger()
694+ if "npugraph_ex" not in sys.modules:
695+ from torch_npu.dynamo import _LazyNpuGraphEx
696+ sys.modules["npugraph_ex"] = _LazyNpuGraphEx("npugraph_ex")
697+ patch_inductor_wrapper()
698+ install_npugraph_mark_step_trigger()
Mtorch_npu/utils/_graph_tree.py+8-5
@@ -22,7 +22,7 @@ from torch._dynamo.backends.cudagraphs import (
22 get_stack_traces,22 get_stack_traces,
23)23)
24from torch._dynamo.backends.debugging import boxed_nop24from torch._dynamo.backends.debugging import boxed_nop
25-from torch._dynamo.backends.registry import register_backend25+from torch._dynamo.backends.registry import _COMPILER_FNS, register_backend
26from torch._inductor import config26from torch._inductor import config
27from torch._inductor.compile_fx import (27from torch._inductor.compile_fx import (
28 get_input_idxs_to_check,28 get_input_idxs_to_check,
@@ -58,7 +58,7 @@ log = torch._logging.getArtifactLogger("torch_npu.npugraph", "cudagraphs")
58 58 
59 59 
60def npugraph_mark_step_begin():60def npugraph_mark_step_begin():
61- from torch_npu.npu._graph_tree import mark_step_begin61+ from torch_npu.npu._graph_tree_state import mark_step_begin
62 mark_step_begin()62 mark_step_begin()
63 63 
64 64 
@@ -373,9 +373,9 @@ class NpugraphsBackend:
373 373 
374 @staticmethod374 @staticmethod
375 def reset():375 def reset():
376- from torch_npu.npu._graph_tree import reset_npugraph_trees376+ from torch_npu.dynamo import _npugraphs_backend_entrypoint
377 377 
378- reset_npugraph_trees()378+ _npugraphs_backend_entrypoint.reset()
379 379 
380 @staticmethod380 @staticmethod
381 def __call__(model, inputs):381 def __call__(model, inputs):
@@ -385,7 +385,10 @@ class NpugraphsBackend:
385def _apply_npugraph_tree_methods():385def _apply_npugraph_tree_methods():
386 # aot_npugraphs only applies graphs to the graph. It is also helpful386 # aot_npugraphs only applies graphs to the graph. It is also helpful
387 # for debugging and can serve as a perf baseline.387 # for debugging and can serve as a perf baseline.
388- register_backend(name="npugraphs", compiler_fn=NpugraphsBackend())388+ if "npugraphs" not in _COMPILER_FNS:
389+ from torch_npu.dynamo import _npugraphs_backend_entrypoint
390+ 
391+ register_backend(name="npugraphs", compiler_fn=_npugraphs_backend_entrypoint)
389 torch._inductor.compile_fx.cudagraphify = npugraphify392 torch._inductor.compile_fx.cudagraphify = npugraphify
390 torch._inductor.cudagraph_utils.check_multiple_devices_or_any_cpu_nodes = check_multiple_devices_or_any_cpu_nodes393 torch._inductor.cudagraph_utils.check_multiple_devices_or_any_cpu_nodes = check_multiple_devices_or_any_cpu_nodes
391 torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin394 torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin
Mtorch_npu/utils/_inductor.py+1-1
@@ -244,4 +244,4 @@ patch_register_run_and_save_rng_state_op()
244patch_register_run_with_rng_state_op()244patch_register_run_with_rng_state_op()
245patch_philox_rand_offset()245patch_philox_rand_offset()
246patch_register_philox_rand()246patch_register_philox_rand()
247-patch_rng_prims_device()247+patch_rng_prims_device()