Star863
1359
代码
代码
Issues1737
Pull Requests2246
项目讨论
Wiki
分析
Star863
1359
已合并
sync: defer dynamo and inductor imports to first torch.compile (from v2.9.0_import) #44299
黄桂军创建于 8月10日
sync: defer dynamo and inductor imports to first torch.compile (from v2.9.0_import) #44299
已合并
黄桂军创建于 8月10日
从stevenaw0/pytorchv2.12.0_import2合入到Ascend/pytorchv2.12.0
共 20 个文件变更+1413-172
Msetup.py+5-0
@@ -792,5 +792,10 @@ setup(
792 'torch.backends': [792 'torch.backends': [
793 'torch_npu = torch_npu:_autoload',793 'torch_npu = torch_npu:_autoload',
794 ],794 ],
795+ 'torch_dynamo_backends': [
796+ 'npu = torch_npu.dynamo:_npu_backend_entrypoint',
797+ 'npugraph_ex = torch_npu.dynamo:_npugraph_ex_backend_entrypoint',
798+ 'npugraphs = torch_npu.dynamo:_npugraphs_backend_entrypoint',
799+ ],
795 }800 }
796)801)
Atest/dynamo/test_compile_trigger.py+956-0
@@ -0,0 +1,956 @@
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+def _triton_compile_works():
14+ """Check whether triton code generation succeeds in the current environment.
15+ 
16+ The CI ACL headers may reference symbols (e.g. aclmdlRICondHandle) that are
17+ absent from the installed ACL version, causing every inductor/triton
18+ compilation to fail. Tests that exercise the real inductor backend should
19+ be skipped when this pre-existing environment issue is present.
20+ """
21+ if not importlib.util.find_spec("triton"):
22+ return False
23+ env = os.environ.copy()
24+ env["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"
25+ probe = textwrap.dedent(
26+ """
27+ import torch
28+ import torch_npu
29+ 
30+ def fn(x):
31+ return torch.sin(x) + 1
32+ 
33+ x = torch.randn(8, device="npu")
34+ actual = torch.compile(fn, backend="inductor", fullgraph=True)(x)
35+ assert bool(torch.allclose(actual, fn(x)))
36+ """
37+ )
38+ result = subprocess.run(
39+ [sys.executable, "-c", probe],
40+ capture_output=True,
41+ env=env,
42+ text=True,
43+ timeout=120,
44+ )
45+ return result.returncode == 0
46+ 
47+ 
48+_TRITON_COMPILE_OK = _triton_compile_works()
49+ 
50+ 
51+class TorchCompileTriggerTests(unittest.TestCase):
52+ def run_in_subprocess(self, code):
53+ env = os.environ.copy()
54+ env["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"
55+ result = subprocess.run(
56+ [sys.executable, "-c", textwrap.dedent(code)],
57+ capture_output=True,
58+ env=env,
59+ text=True,
60+ timeout=60,
61+ )
62+ self.assertEqual(
63+ result.returncode,
64+ 0,
65+ f"Subprocess failed:\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}",
66+ )
67+ 
68+ # Verify import installs only the Dynamo post-import trigger.
69+ def test_import_installs_dynamo_post_import_trigger(self):
70+ self.run_in_subprocess(
71+ """
72+ import sys
73+ import torch
74+ 
75+ def loaded(prefix):
76+ return [
77+ name for name in sys.modules
78+ if name == prefix or name.startswith(prefix + ".")
79+ ]
80+ 
81+ assert not loaded("torch._dynamo")
82+ assert not loaded("torch._inductor")
83+ src_compile = torch.compile
84+ src_wrapper_init = torch._TorchCompileWrapper.__init__
85+ 
86+ import torch_npu
87+ from torch_npu.utils import _dynamo
88+ 
89+ assert not loaded("torch._dynamo")
90+ assert not loaded("torch._inductor")
91+ assert not loaded("torch_npu._inductor")
92+ assert "triton" not in sys.modules
93+ assert torch.utils._triton.has_triton is _dynamo.has_triton
94+ assert torch.compile is src_compile
95+ assert torch._TorchCompileWrapper.__init__ is src_wrapper_init
96+ assert any(
97+ isinstance(finder, _dynamo._DynamoPostImportFinder)
98+ for finder in sys.meta_path
99+ )
100+ assert not _dynamo._lazy_dynamo_setup.has_run
101+ """
102+ )
103+ 
104+ # Verify public compiler APIs work before the first compile.
105+ def test_public_compiler_entries_are_available_before_compile(self):
106+ self.run_in_subprocess(
107+ """
108+ import inspect
109+ import sys
110+ import torch
111+ import torch_npu
112+ from torch_npu.utils import _dynamo
113+ 
114+ marker = torch.compiler.npugraph_mark_step_begin
115+ assert marker.__name__ == "npugraph_mark_step_begin"
116+ assert str(inspect.signature(marker)) == "()"
117+ marker()
118+ 
119+ from torch_npu.npu._graph_tree_state import MarkStepBox
120+ 
121+ assert MarkStepBox.mark_step_counter == -1
122+ assert not _dynamo._lazy_dynamo_setup.has_run
123+ assert not _dynamo._lazy_inductor_setup.has_run
124+ assert not any(
125+ name == "torch._dynamo" or name.startswith("torch._dynamo.")
126+ for name in sys.modules
127+ )
128+ assert not any(
129+ name == "torch._inductor" or name.startswith("torch._inductor.")
130+ for name in sys.modules
131+ )
132+ 
133+ backends = torch.compiler.list_backends(exclude_tags=None)
134+ assert {"npu", "npugraph_ex", "npugraphs"}.issubset(backends)
135+ assert _dynamo._lazy_dynamo_setup.has_run
136+ assert not _dynamo._lazy_inductor_setup.has_run
137+ assert "torch_npu._inductor" not in sys.modules
138+ """
139+ )
140+ 
141+ # Verify entry-point loading cannot register the NPU backend twice.
142+ def test_backend_entrypoint_can_import_torch_npu_without_duplicate_registration(self):
143+ self.run_in_subprocess(
144+ """
145+ import torch
146+ import torch._dynamo
147+ from torch._dynamo.backends import registry
148+ 
149+ # Model the state produced by setuptools entry-point discovery:
150+ # lookup_backend owns the final register_backend call, while
151+ # EntryPoint.load imports torch_npu and triggers its lazy setup.
152+ class NpuEntryPoint:
153+ module = "torch_npu.dynamo"
154+ 
155+ def load(self):
156+ from torch_npu.dynamo import _npu_backend_entrypoint
157+ return _npu_backend_entrypoint
158+ 
159+ registry._BACKENDS["npu"] = NpuEntryPoint()
160+ backend = registry.lookup_backend("npu")
161+ 
162+ assert backend.__name__ == "_npu_backend_entrypoint"
163+ assert registry._COMPILER_FNS["npu"] is backend
164+ 
165+ # A later setup retry must treat the loaded torch_npu entry point
166+ # as an already completed registration.
167+ from torch_npu.dynamo import _register_npu_backend
168+ 
169+ _register_npu_backend(backend, "npu")
170+ assert registry._COMPILER_FNS["npu"] is backend
171+ """
172+ )
173+ 
174+ # Verify backend registration is safe to retry after a partial failure.
175+ def test_backend_registration_retry_after_partial_failure(self):
176+ self.run_in_subprocess(
177+ """
178+ import torch
179+ import torch_npu
180+ import torch._dynamo
181+ import torch_npu.dynamo as npu_dynamo
182+ from torch._dynamo.backends import registry
183+ 
184+ # Start from an undiscovered state so the first registration is
185+ # committed before the simulated second-registration failure.
186+ for name in ("npu", "npugraph_ex"):
187+ registry._BACKENDS.pop(name, None)
188+ registry._COMPILER_FNS.pop(name, None)
189+ 
190+ original_register = npu_dynamo._register_npu_backend
191+ fail_npugraph_ex = True
192+ 
193+ def fail_second_registration(backend, name="npu"):
194+ if name == "npugraph_ex" and fail_npugraph_ex:
195+ raise RuntimeError("simulated npugraph_ex registration failure")
196+ return original_register(backend, name)
197+ 
198+ npu_dynamo._register_npu_backend = fail_second_registration
199+ try:
200+ try:
201+ npu_dynamo._register_backends()
202+ except RuntimeError as error:
203+ assert "simulated npugraph_ex" in str(error)
204+ else:
205+ raise AssertionError("the first registration should fail")
206+ finally:
207+ npu_dynamo._register_npu_backend = original_register
208+ 
209+ assert "npu" in registry._COMPILER_FNS
210+ assert "npugraph_ex" not in registry._COMPILER_FNS
211+ 
212+ # Retry: the completed npu registration is a no-op, while the
213+ # missing npugraph_ex registration is installed normally.
214+ npu_dynamo._register_backends()
215+ assert "npu" in registry._COMPILER_FNS
216+ assert "npugraph_ex" in registry._COMPILER_FNS
217+ """
218+ )
219+ 
220+ def test_public_lazy_setup_recovers_after_fork(self):
221+ self.run_in_subprocess(
222+ """
223+ import os
224+ import signal
225+ import threading
226+ 
227+ import torch_npu
228+ import torch_npu.dynamo as npu_dynamo
229+ from torch_npu.utils import _dynamo
230+ 
231+ parent_pid = os.getpid()
232+ original_add = _dynamo.add_dynamo_methods_init
233+ original_get_backend = npu_dynamo._get_default_backend
234+ entered = threading.Event()
235+ release = threading.Event()
236+ 
237+ def block_parent_setup():
238+ if os.getpid() == parent_pid:
239+ entered.set()
240+ assert release.wait(timeout=10)
241+ return original_add()
242+ 
243+ _dynamo.add_dynamo_methods_init = block_parent_setup
244+ setup_thread = threading.Thread(target=_dynamo._lazy_dynamo_setup)
245+ setup_thread.start()
246+ assert entered.wait(timeout=5)
247+ 
248+ try:
249+ child_pid = os.fork()
250+ if child_pid == 0:
251+ def timeout(_signal, _frame):
252+ raise TimeoutError("lazy setup hung after fork")
253+ 
254+ signal.signal(signal.SIGALRM, timeout)
255+ signal.alarm(5)
256+ try:
257+ npu_dynamo._get_default_backend = (
258+ lambda name: npu_dynamo._eager_npu_backend
259+ )
260+ graph_module = lambda x: x
261+ result = npu_dynamo._npu_backend_entrypoint(
262+ graph_module, []
263+ )
264+ assert result is graph_module
265+ signal.alarm(0)
266+ except BaseException as error:
267+ os.write(
268+ 2,
269+ f"child lazy setup failed: {error}\\n".encode(),
270+ )
271+ os._exit(1)
272+ os._exit(0)
273+ 
274+ _, status = os.waitpid(child_pid, 0)
275+ finally:
276+ release.set()
277+ setup_thread.join(timeout=10)
278+ _dynamo.add_dynamo_methods_init = original_add
279+ npu_dynamo._get_default_backend = original_get_backend
280+ 
281+ assert os.WIFEXITED(status)
282+ assert os.waitstatus_to_exitcode(status) == 0
283+ assert not setup_thread.is_alive()
284+ """
285+ )
286+ 
287+ # Verify NPUGraphs rejects unsupported options in every registration order.
288+ def test_npugraphs_rejects_options_across_registration_order(self):
289+ self.run_in_subprocess(
290+ """
291+ import contextlib
292+ from unittest import mock
293+ 
294+ import torch
295+ import torch_npu
296+ from torch_npu.dynamo import _npugraphs_backend_entrypoint
297+ from torch_npu.utils import _dynamo, _graph_tree
298+ 
299+ gm = object()
300+ inputs = [object()]
301+ options = {"npu_backend": "mlir"}
302+ 
303+ with mock.patch.object(
304+ _dynamo, "_lazy_dynamo_setup", lambda: None
305+ ), mock.patch.object(
306+ _dynamo, "_lazy_inductor_setup", lambda: None
307+ ), mock.patch.object(
308+ _dynamo,
309+ "_NpuBackendScope",
310+ lambda backend: contextlib.nullcontext(),
311+ ), mock.patch.object(
312+ _graph_tree,
313+ "npugraphs",
314+ lambda model, args, **kwargs: "unexpected",
315+ ):
316+ for backend in (
317+ _npugraphs_backend_entrypoint,
318+ _graph_tree.NpugraphsBackend(),
319+ ):
320+ try:
321+ backend(gm, inputs, options=options)
322+ except TypeError as error:
323+ assert "unexpected keyword argument 'options'" in str(error)
324+ else:
325+ raise AssertionError("npugraphs must reject options")
326+ 
327+ compiled = torch.compile(
328+ lambda value: value + 1,
329+ backend="npugraphs",
330+ options=options,
331+ )
332+ try:
333+ compiled(torch.ones(1))
334+ except Exception as error:
335+ assert "unexpected keyword argument 'options'" in str(error)
336+ else:
337+ raise AssertionError("public npugraphs must reject options")
338+ """
339+ )
340+ 
341+ # Verify public reset reaches NPUGraphs for every registration order.
342+ 
343+ def test_npugraphs_reset_protocol_across_registration_order(self):
344+ for order in ("cold", "hot"):
345+ with self.subTest(order=order):
346+ initialize_inductor = (
347+ "torch.compile(lambda x: x + 1, backend='inductor')"
348+ if order == "hot"
349+ else ""
350+ )
351+ self.run_in_subprocess(
352+ f"""
353+ import sys
354+ import types
355+ from unittest import mock
356+ 
357+ import torch
358+ import torch_npu
359+ 
360+ {initialize_inductor}
361+ torch.compile(lambda x: x + 1, backend="npugraphs")
362+ 
363+ from torch._dynamo.backends import registry
364+ from torch_npu.dynamo import _npugraphs_backend_entrypoint
365+ 
366+ backend = registry._COMPILER_FNS["npugraphs"]
367+ assert backend is _npugraphs_backend_entrypoint
368+ assert hasattr(backend, "reset")
369+ 
370+ graph_tree_module = "torch_npu.npu._graph_tree"
371+ assert graph_tree_module not in sys.modules
372+ backend.reset()
373+ assert graph_tree_module not in sys.modules
374+ 
375+ reset_calls = []
376+ fake_graph_tree = types.ModuleType(graph_tree_module)
377+ fake_graph_tree.reset_npugraph_trees = (
378+ lambda: reset_calls.append("reset")
379+ )
380+ with mock.patch.dict(
381+ sys.modules,
382+ {{graph_tree_module: fake_graph_tree}},
383+ ):
384+ torch.compiler.reset()
385+ 
386+ assert reset_calls == ["reset"]
387+ """
388+ )
389+ 
390+ @unittest.skipUnless(_TRITON_COMPILE_OK, "inductor/triton compile unavailable (ACL header issue)")
391+ def test_npu_export_public_entry_and_import_order_matrix(self):
392+ cases = {
393+ "module_export": (
394+ "",
395+ "exported = torch.export.export(Model(), (x,), strict=True)",
396+ ),
397+ "prebound_export": (
398+ "from torch.export import export as export_api",
399+ "exported = export_api(Model(), (x,), strict=True)",
400+ ),
401+ "prebound_export_for_training": (
402+ "from torch.export import export_for_training as export_api",
403+ "exported = export_api(Model(), (x,), strict=True)",
404+ ),
405+ }
406+ if hasattr(torch.export, "export_for_inference"):
407+ cases["prebound_export_for_inference"] = (
408+ "from torch.export import export_for_inference as export_api",
409+ "exported = export_api(Model(), (x,), strict=True)",
410+ )
411+ for name, (pre_import, export_call) in cases.items():
412+ with self.subTest(name=name):
413+ self.run_in_subprocess(
414+ f"""
415+ import sys
416+ import torch
417+ {pre_import}
418+ import torch_npu
419+ 
420+ stream = torch.npu.Stream()
421+ 
422+ class Model(torch.nn.Module):
423+ def forward(self, x):
424+ x.record_stream(stream)
425+ return x + 1
426+ 
427+ x = torch.ones(4, device="npu")
428+ {export_call}
429+ actual = exported.module()(x)
430+ 
431+ from torch_npu.utils import _dynamo
432+ 
433+ torch.testing.assert_close(actual, x + 1)
434+ assert _dynamo._lazy_dynamo_setup.has_run
435+ assert not _dynamo._lazy_inductor_setup.has_run
436+ assert "torch_npu._inductor" not in sys.modules
437+ """
438+ )
439+ 
440+ # Verify NPU-specific operations retain their Export capture semantics.
441+ def test_npu_export_capture_semantics_matrix(self):
442+ self.run_in_subprocess(
443+ """
444+ import sys
445+ import torch
446+ import torch_npu
447+ 
448+ x = torch.ones(4, device="npu")
449+ stream = torch.npu.Stream()
450+ event = torch.npu.Event()
451+ 
452+ class StreamAndEvent(torch.nn.Module):
453+ def forward(self, value):
454+ event.record()
455+ with torch.npu.stream(stream):
456+ event.wait(stream)
457+ result = value + 1
458+ return result
459+ 
460+ class Autocast(torch.nn.Module):
461+ def forward(self, value):
462+ with torch.npu.amp.autocast(dtype=torch.float16):
463+ return value * value
464+ 
465+ class CurrentDevice(torch.nn.Module):
466+ def forward(self, value):
467+ return value + torch.npu.current_device()
468+ 
469+ class DeviceProperties(torch.nn.Module):
470+ def forward(self, value):
471+ properties = torch.npu.get_device_properties(
472+ torch.npu.current_device()
473+ )
474+ return value + 1 if properties.total_memory > 0 else value - 1
475+ 
476+ class IsAvailable(torch.nn.Module):
477+ def forward(self, value):
478+ return value + 1 if torch.npu.is_available() else value - 1
479+ 
480+ models = (
481+ StreamAndEvent,
482+ Autocast,
483+ CurrentDevice,
484+ DeviceProperties,
485+ IsAvailable,
486+ )
487+ for model_type in models:
488+ model = model_type()
489+ expected = model(x)
490+ exported = torch.export.export(model, (x,))
491+ actual = exported.module()(x)
492+ torch.testing.assert_close(actual, expected)
493+ 
494+ from torch_npu.utils import _dynamo
495+ 
496+ assert _dynamo._lazy_dynamo_setup.has_run
497+ assert not _dynamo._lazy_inductor_setup.has_run
498+ assert "torch_npu._inductor" not in sys.modules
499+ """
500+ )
501+ 
502+ # Verify rejected Inductor arguments do not initialize or pollute NPU state.
503+ def test_invalid_inductor_arguments_fail_without_npu_initialization(self):
504+ cases = {
505+ "invalid_mode": (
506+ 'torch.compile(lambda x: x + 1, mode="invalid-mode")',
507+ "Unrecognized mode=invalid-mode",
508+ ),
509+ "invalid_option": (
510+ "torch.compile(lambda x: x + 1, "
511+ 'options={"invalid.option": True})',
512+ "Unexpected optimization option invalid.option",
513+ ),
514+ }
515+ for name, (compile_call, expected_error) in cases.items():
516+ with self.subTest(name=name):
517+ self.run_in_subprocess(
518+ f"""
519+ import os
520+ import sys
521+ 
522+ import torch
523+ import torch_npu
524+ from torch_npu.utils import _dynamo
525+ 
526+ env_name = "TORCHINDUCTOR_NPU_BACKEND"
527+ original_env = os.environ.get(env_name)
528+ try:
529+ {compile_call}
530+ except RuntimeError as error:
531+ assert {expected_error!r} in str(error), str(error)
532+ else:
533+ raise AssertionError("invalid compile arguments must fail")
534+ 
535+ assert not _dynamo._lazy_inductor_setup.has_run
536+ assert not _dynamo.is_inductor_npu_initialized()
537+ assert "torch_npu._inductor" not in sys.modules
538+ assert os.environ.get(env_name) == original_env
539+ """
540+ )
541+ 
542+ # Verify non-Inductor compile backends do not initialize Inductor.
543+ def test_non_inductor_compile_backend_matrix(self):
544+ cases = {
545+ "eager": (
546+ "",
547+ 'compiled = torch.compile(Model(), backend="eager", fullgraph=True)',
548+ ),
549+ "custom": (
550+ "custom_backend = lambda graph_module, example_inputs: "
551+ "graph_module.forward",
552+ "compiled = torch.compile(Model(), backend=custom_backend, fullgraph=True)",
553+ ),
554+ "npu": (
555+ "",
556+ 'compiled = torch.compile(Model(), backend="npu", fullgraph=True)',
557+ ),
558+ }
559+ for name, (backend_definition, compile_call) in cases.items():
560+ with self.subTest(name=name):
561+ allow_missing_torchair = name == "npu"
562+ self.run_in_subprocess(
563+ f"""
564+ import sys
565+ import torch
566+ import torch_npu
567+ 
568+ class Model(torch.nn.Module):
569+ def forward(self, x):
570+ return torch.sin(x) + 1
571+ 
572+ {backend_definition}
573+ x = torch.randn(8, device="npu")
574+ try:
575+ {compile_call}
576+ except AssertionError as error:
577+ assert {allow_missing_torchair!r}
578+ assert "Could not find module torchair" in str(error)
579+ else:
580+ torch.testing.assert_close(compiled(x), Model()(x))
581+ 
582+ from torch_npu.utils import _dynamo
583+ 
584+ assert _dynamo._lazy_dynamo_setup.has_run
585+ assert not _dynamo._lazy_inductor_setup.has_run
586+ assert "torch_npu._inductor" not in sys.modules
587+ """
588+ )
589+ 
590+ # Verify ONNX Dynamo Export initializes only the NPU Dynamo integration.
591+ @unittest.skipUnless(importlib.util.find_spec("onnxscript"), "requires onnxscript")
592+ def test_npu_onnx_dynamo_export_initialization_chain(self):
593+ cases = {
594+ "module_export": (
595+ "",
596+ "result = torch.onnx.export(Model(), (x,), dynamo=True)",
597+ ),
598+ "prebound_export": (
599+ "from torch.onnx import export as onnx_export",
600+ "result = onnx_export(Model(), (x,), dynamo=True)",
601+ ),
602+ }
603+ if hasattr(torch.onnx, "dynamo_export"):
604+ cases["prebound_legacy_dynamo_export"] = (
605+ "from torch.onnx import dynamo_export as onnx_export",
606+ "result = onnx_export(Model(), x)",
607+ )
608+ for name, (pre_import, export_call) in cases.items():
609+ with self.subTest(name=name):
610+ self.run_in_subprocess(
611+ f"""
612+ import sys
613+ import torch
614+ {pre_import}
615+ import torch_npu
616+ 
617+ # ONNXScript 0.4.0 cannot version-convert models containing
618+ # functions. Conversion runs after Dynamo capture, which is
619+ # the boundary this test verifies, so isolate that unrelated
620+ # compatibility issue without weakening the import assertions.
621+ from torch.onnx._internal._lazy_import import onnxscript_apis
622+ onnxscript_apis.convert_version = lambda model, target: model
623+ 
624+ class Model(torch.nn.Module):
625+ def forward(self, x):
626+ return torch.sin(x) + 1
627+ 
628+ x = torch.randn(8, device="npu")
629+ {export_call}
630+ 
631+ from torch_npu.utils import _dynamo
632+ 
633+ assert result is not None
634+ assert _dynamo._lazy_dynamo_setup.has_run
635+ assert not _dynamo._lazy_inductor_setup.has_run
636+ assert "torch_npu._inductor" not in sys.modules
637+ """
638+ )
639+ 
640+ # Verify pre-imported FSDP receives all NPU patches.
641+ def test_fsdp_patch_when_imported_before_torch_npu(self):
642+ self.run_in_subprocess(
643+ """
644+ import torch.distributed.fsdp
645+ import torch_npu
646+ from torch.distributed.fsdp import sharded_grad_scaler
647+ from torch.distributed.fsdp._fully_shard._fsdp_param_group import (
648+ FSDPParamGroup,
649+ )
650+ from torch_npu.distributed.fsdp._add_fsdp_patch import (
651+ _patched_fsdp_param_group_init,
652+ )
653+ from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler
654+ 
655+ assert callable(torch_npu.distributed.fsdp.fully_shard)
656+ assert sharded_grad_scaler.ShardedGradScaler is _ShardedGradScaler
657+ """
658+ )
659+ 
660+ # Verify a real NPU Inductor compile initializes the full stack.
661+ @unittest.skipUnless(_TRITON_COMPILE_OK, "inductor/triton compile unavailable (ACL header issue)")
662+ def test_npu_inductor_initialization_chain(self):
663+ self.run_in_subprocess(
664+ """
665+ import torch
666+ import torch_npu
667+ 
668+ def fn(x):
669+ return torch.sin(x) + 1
670+ 
671+ x = torch.randn(8, device="npu")
672+ actual = torch.compile(fn, backend="inductor", fullgraph=True)(x)
673+ 
674+ from torch._dynamo.device_interface import get_interface_for_device
675+ from torch_npu.utils import _dynamo
676+ 
677+ torch.testing.assert_close(actual, fn(x))
678+ assert _dynamo._lazy_dynamo_setup.has_run
679+ assert _dynamo._lazy_inductor_setup.has_run
680+ assert get_interface_for_device("npu").device_count() > 0
681+ """
682+ )
683+ 
684+ # Verify a real NPUGraphs compile initializes the full stack.
685+ @unittest.skipUnless(importlib.util.find_spec("triton"), "requires triton-ascend")
686+ def test_npu_npugraphs_initialization_chain(self):
687+ self.run_in_subprocess(
688+ """
689+ import torch
690+ import torch_npu
691+ 
692+ def fn(x):
693+ return torch.sin(x) + 1
694+ 
695+ x = torch.randn(8, device="npu")
696+ actual = torch.compile(fn, backend="npugraphs", fullgraph=True)(x)
697+ 
698+ from torch._dynamo.device_interface import get_interface_for_device
699+ from torch_npu.utils import _dynamo
700+ 
701+ torch.testing.assert_close(actual, fn(x))
702+ assert _dynamo._lazy_dynamo_setup.has_run
703+ assert _dynamo._lazy_inductor_setup.has_run
704+ assert get_interface_for_device("npu").device_count() > 0
705+ """
706+ )
707+ 
708+ # Verify lazy setup completes before compile backend lookup.
709+ def test_compile_triggers_setup_before_backend_lookup(self):
710+ self.run_in_subprocess(
711+ """
712+ import torch
713+ import torch_npu
714+ from torch_npu.utils import _dynamo
715+ 
716+ calls = []
717+ 
718+ @_dynamo.run_once
719+ def fake_setup():
720+ calls.append("setup")
721+ 
722+ _dynamo._lazy_dynamo_setup = fake_setup
723+ from torch._dynamo.backends import registry
724+ assert fake_setup.has_run
725+ 
726+ compiled = torch.compile(lambda x: x + 1, backend="eager")
727+ assert compiled(torch.tensor(1)).item() == 2
728+ assert calls == ["setup"]
729+ """
730+ )
731+ 
732+ # Verify the trigger works when Dynamo was imported first.
733+ def test_trigger_after_dynamo_was_preimported(self):
734+ self.run_in_subprocess(
735+ """
736+ import sys
737+ import torch
738+ import torch._dynamo
739+ import torch_npu
740+ from torch_npu.utils import _dynamo
741+ 
742+ assert _dynamo._lazy_dynamo_setup.has_run
743+ 
744+ compiled = torch.compile(lambda x: x + 1, backend="eager", fullgraph=True)
745+ assert compiled(torch.tensor(1)).item() == 2
746+ assert _dynamo._lazy_dynamo_setup.has_run
747+ assert not _dynamo._lazy_inductor_setup.has_run
748+ assert "torch_npu._inductor" not in sys.modules
749+ """
750+ )
751+ 
752+ # Verify the moved helper preserves the torch_npu v2.9 device predicates.
753+ def test_has_triton_preserves_target_device_semantics(self):
754+ self.run_in_subprocess(
755+ """
756+ import sys
757+ import types
758+ from unittest import mock
759+ 
760+ import torch
761+ import torch_npu
762+ from torch._dynamo import device_interface
763+ from torch_npu.utils import _dynamo
764+ 
765+ def make_interface(available=False):
766+ class Worker:
767+ @staticmethod
768+ def get_device_properties():
769+ raise AssertionError("v2.9 NPU override must not query CUDA props")
770+ 
771+ class Interface:
772+ @staticmethod
773+ def is_available():
774+ return available
775+ 
776+ Interface.Worker = Worker
777+ return Interface
778+ 
779+ def run_case(available, *, package=True):
780+ interfaces = {
781+ device: make_interface(
782+ available=device in available,
783+ )
784+ for device in ("cuda", "xpu", "cpu", "npu")
785+ }
786+ _dynamo.has_triton.cache_clear()
787+ with mock.patch.object(
788+ torch.utils._triton,
789+ "has_triton_package",
790+ return_value=package,
791+ ), mock.patch.object(
792+ device_interface,
793+ "get_interface_for_device",
794+ side_effect=interfaces.__getitem__,
795+ ), mock.patch.object(
796+ _dynamo,
797+ "_dynamo_register_interface_for_device",
798+ ) as register:
799+ result = _dynamo.has_triton()
800+ if package:
801+ register.assert_called_once_with()
802+ else:
803+ register.assert_not_called()
804+ return result
805+ 
806+ assert not run_case({}, package=False)
807+ assert run_case({"cuda"})
808+ assert run_case({"xpu"})
809+ triton = types.ModuleType("triton")
810+ triton_backends = types.ModuleType("triton.backends")
811+ triton_backends.backends = {"cpu": object()}
812+ triton.backends = triton_backends
813+ with mock.patch.dict(
814+ sys.modules,
815+ {"triton": triton, "triton.backends": triton_backends},
816+ ):
817+ assert run_case({"cpu"})
818+ triton_backends.backends = {}
819+ assert not run_case({"cpu"})
820+ assert run_case({"npu"})
821+ """
822+ )
823+ 
824+ # Verify all legacy Dynamo patches remain installed exactly once.
825+ def test_dynamo_patch_inventory_is_preserved(self):
826+ self.run_in_subprocess(
827+ """
828+ import torch
829+ import torch_npu
830+ 
831+ # Importing the Dynamo parent package is the lazy setup boundary.
832+ import torch._dynamo
833+ 
834+ from torch._dynamo.device_interface import get_interface_for_device
835+ from torch._dynamo.variables.builtin import BuiltinVariable
836+ from torch._dynamo.variables.functions import SkipFunctionVariable
837+ from torch._dynamo.variables.tensor import TensorVariable
838+ from torch._dynamo.variables.torch import constant_fold_functions
839+ from torch._dynamo.variables.user_defined import UserDefinedClassVariable
840+ from torch._dynamo.utils import common_constant_types
841+ from torch_npu.dynamo.trace_rule import (
842+ skip_functions_npu,
843+ torch_c_binding_in_graph_functions_npu,
844+ torch_non_c_binding_in_graph_functions_npu,
845+ )
846+ from torch_npu.utils import _dynamo
847+ 
848+ assert _dynamo._lazy_dynamo_setup.has_run
849+ assert get_interface_for_device("npu").device_count() > 0
850+ 
851+ # VariableTracker and context-manager patches formerly installed
852+ # eagerly by add_dynamo_methods().
853+ assert SkipFunctionVariable.__new__.__module__ == "torch_npu.utils._dynamo"
854+ assert TensorVariable.call_method.__module__ == "torch_npu.utils._dynamo"
855+ assert UserDefinedClassVariable.__new__.__module__ == "torch_npu.utils._dynamo"
856+ in_graph_classes = UserDefinedClassVariable._in_graph_classes()
857+ assert torch.npu.Event in in_graph_classes
858+ assert torch.npu.Stream in in_graph_classes
859+ assert BuiltinVariable.call_id.__module__ == "torch_npu.utils._dynamo"
860+ assert torch._dynamo.optimize.__module__ == "torch_npu.utils._dynamo"
861+ 
862+ # Backend and trace-rule registrations formerly performed by
863+ # registry_manager._register_dynamo(). Count the maps as well as
864+ # checking membership so repeated lazy triggers cannot hide a
865+ # duplicate installation.
866+ assert {"npu", "npugraph_ex"}.issubset(
867+ torch._dynamo.list_backends(exclude_tags=None)
868+ )
869+ maps = torch._dynamo.trace_rules.torch_name_rule_map
870+ assert maps.count(torch_non_c_binding_in_graph_functions_npu) == 1
871+ assert maps.count(torch_c_binding_in_graph_functions_npu) == 1
872+ assert maps.count(skip_functions_npu) == 1
873+ assert constant_fold_functions[torch.npu.current_device]
874+ assert constant_fold_functions[torch.npu.get_device_properties]
875+ assert constant_fold_functions[torch.npu.is_available]
876+ assert torch_npu._C._NPUDeviceProperties in common_constant_types
877+ """
878+ )
879+ 
880+ # Verify all legacy Inductor patches remain installed.
881+ def test_inductor_patch_inventory_is_preserved(self):
882+ self.run_in_subprocess(
883+ """
884+ import torch
885+ import torch_npu
886+ import torch._dynamo
887+ from torch_npu.utils import _dynamo
888+ 
889+ # RNG/decomposition patches are installed when the deferred
890+ # compiler setup is first triggered, not while importing torch_npu.
891+ _dynamo._lazy_inductor_setup()
892+ 
893+ from torch_npu.utils import _inductor
894+ 
895+ assert torch._prims.rng_prims.philox_rand_offset.__module__ == (
896+ "torch_npu.utils._inductor"
897+ )
898+ assert torch._prims.rng_prims.register_philox_rand.__module__ == (
899+ "torch_npu.utils._inductor"
900+ )
901+ assert torch._prims.rng_prims.get_device.__module__ == (
902+ "torch_npu.utils._inductor"
903+ )
904+ 
905+ import torch._inductor.compile_fx as compile_fx
906+ import torch._inductor.cudagraph_trees as cudagraph_trees
907+ import torch._inductor.cudagraph_utils as cudagraph_utils
908+ import torch._inductor.scheduler as scheduler
909+ from torch._inductor.codegen.common import get_device_op_overrides
910+ from torch._inductor.codecache import CacheBase
911+ from torch._inductor.graph import GraphLowering
912+ from torch._inductor.utils import GPU_TYPES
913+ from torch_npu.utils import _graph_tree
914+ 
915+ assert _dynamo._lazy_inductor_setup.has_run
916+ assert "npu" in GPU_TYPES
917+ assert get_device_op_overrides("npu").__class__.__module__.startswith(
918+ "torch_npu._inductor"
919+ )
920+ assert torch.utils._triton.has_triton.__module__ == (
921+ "torch_npu.utils._dynamo"
922+ )
923+ assert torch.utils._triton._device_supports_tma.__module__ == (
924+ "torch_npu._inductor.utils"
925+ )
926+ assert compile_fx.has_triton is torch.utils._triton.has_triton
927+ assert scheduler.has_triton is torch.utils._triton.has_triton
928+ assert GraphLowering.codegen_with_cpp_wrapper.__module__ == (
929+ "torch_npu._inductor.graph"
930+ )
931+ assert CacheBase.get_system.__module__ == (
932+ "torch_npu._inductor.codegen.common"
933+ )
934+ 
935+ # NPUGraph integrations were formerly applied eagerly alongside
936+ # the Inductor patches.
937+ assert compile_fx.cudagraphify is _graph_tree.npugraphify
938+ assert (
939+ cudagraph_utils.check_multiple_devices_or_any_cpu_nodes
940+ is _graph_tree.check_multiple_devices_or_any_cpu_nodes
941+ )
942+ assert cudagraph_trees.get_manager.__module__ == (
943+ "torch_npu.utils._graph_tree"
944+ )
945+ assert torch.compiler.npugraph_mark_step_begin is (
946+ _graph_tree.npugraph_mark_step_begin
947+ )
948+ 
949+ config = torch._inductor.config
950+ assert config.npu_backend == "default"
951+ """
952+ )
953+ 
954+ 
955+if __name__ == "__main__":
956+ unittest.main()
Mtest/test_torch_npu_init.py+1-16
@@ -33,7 +33,6 @@ 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",
@@ -41,6 +40,7 @@ EXPECTED_LOADED_MODULES = [
41 40 
42EXPECTED_NOT_LOADED_MODULES = [41EXPECTED_NOT_LOADED_MODULES = [
43 "torch_npu._C._afd",42 "torch_npu._C._afd",
43+ "torch_npu._inductor",
44]44]
45 45 
46EXPECTED_TOP_LEVEL_ATTRS = [46EXPECTED_TOP_LEVEL_ATTRS = [
@@ -255,21 +255,6 @@ class TestTorchNpuBootstrap(TestCase):
255 import torch.distributed as dist255 import torch.distributed as dist
256 import torch.distributed.rpc as rpc256 import torch.distributed.rpc as rpc
257 import torch.distributed.tensor # noqa: F401257 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- 
273 assert "hccl" in dist.Backend.backend_list258 assert "hccl" in dist.Backend.backend_list
274 assert "lccl" in dist.Backend.backend_list259 assert "lccl" in dist.Backend.backend_list
275 260 
Mtest/utils/test_inductor.py+2-2
@@ -1,5 +1,5 @@
1from torch_npu.testing.testcase import TestCase, run_tests1from torch_npu.testing.testcase import TestCase, run_tests
2-from torch_npu.utils._inductor import NPUDeviceOpOverrides2+from torch_npu.utils._device_op_overrides import NPUDeviceOpOverrides
3 3 
4 4 
5class TestInductor():5class TestInductor():
@@ -30,4 +30,4 @@ class TestInductor():
30 30 
31 31 
32if __name__ == "__main__":32if __name__ == "__main__":
33- run_tests()33+ run_tests()
Mtorch_npu/_inductor/__init__.py+1-1
@@ -84,7 +84,7 @@ def _load_triton_backend():
84 from torch._inductor.runtime import autotune_cache84 from torch._inductor.runtime import autotune_cache
85 from torch_npu.npu import device_count85 from torch_npu.npu import device_count
86 from torch_npu.utils._dynamo_device import current_device, NpuInterface, set_device86 from torch_npu.utils._dynamo_device import current_device, NpuInterface, set_device
87- from torch_npu.utils._inductor import NPUDeviceOpOverrides87+ from torch_npu.utils._device_op_overrides import NPUDeviceOpOverrides
88 88 
89 from . import codegen, config as npu_config89 from . import codegen, config as npu_config
90 from .codecache import patch_aot_code_compiler_compile90 from .codecache import patch_aot_code_compiler_compile
Mtorch_npu/_inductor/npu_device.py+1-1
@@ -2,7 +2,7 @@ import torch
2from torch._inductor.codegen.common import register_device_op_overrides2from torch._inductor.codegen.common import register_device_op_overrides
3from torch_npu.npu import device_count3from torch_npu.npu import device_count
4from torch_npu.utils._dynamo_device import NpuInterface, current_device, set_device4from torch_npu.utils._dynamo_device import NpuInterface, current_device, set_device
5-from torch_npu.utils._inductor import NPUDeviceOpOverrides5+from torch_npu.utils._device_op_overrides import NPUDeviceOpOverrides
6from . import config as npu_config6from . import config as npu_config
7 7 
8 8 
Mtorch_npu/_inductor/utils.py+2-1
@@ -49,10 +49,11 @@ def resolve_npu_device_index(device_idx=None) -> int:
49 49 
50 50 
51def patch_has_triton():51def patch_has_triton():
52+ from torch._inductor import compile_fx
52 from torch_npu.utils._dynamo import has_triton53 from torch_npu.utils._dynamo import has_triton
53 54 
54 torch._inductor.scheduler.has_triton = has_triton55 torch._inductor.scheduler.has_triton = has_triton
55- torch._inductor.compile_fx.has_triton = has_triton56+ compile_fx.has_triton = has_triton
56 57 
57 58 
58def patch_device_supports_tma():59def 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+6-27
@@ -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.
@@ -84,11 +65,6 @@ def _register_inductor():
84 Inductor backend loading and heavy global patches lazily when torch.compile65 Inductor backend loading and heavy global patches lazily when torch.compile
85 and Inductor path is actually used.66 and Inductor path is actually used.
86 """67 """
87- from torch_npu.utils._inductor import _inductor_register_device_op_overrides
88- 
89- _inductor_register_device_op_overrides()
90- 
91- 
92def _register_default_gradient_device_type():68def _register_default_gradient_device_type():
93 """69 """
94 Set default device type for gradient checkpointing.70 Set default device type for gradient checkpointing.
@@ -102,9 +78,13 @@ def _register_components():
102 78 
103 Order matters:79 Order matters:
104 1. NPU backend is the base capability.80 1. NPU backend is the base capability.
105- 2. Distributed and Dynamo depend on NPU backend / _C children.81+ 2. Distributed depends on the NPU backend / _C children.
106- 3. RPC, dtensor and inductor are Python-side framework integrations.82+ 3. RPC and inductor are Python-side framework integrations.
107 4. DefaultDeviceType is set after NPU backend is registered.83 4. DefaultDeviceType is set after NPU backend is registered.
84+ 
85+ Dynamo and Inductor integrations are deferred to the first torch.compile
86+ call (see torch_npu.utils._dynamo.add_dynamo_methods) so that importing
87+ torch_npu does not pull in heavy compiler submodules.
108 """88 """
109 if not hasattr(torch_npu, "_C"):89 if not hasattr(torch_npu, "_C"):
110 raise RuntimeError(90 raise RuntimeError(
@@ -113,7 +93,6 @@ def _register_components():
113 93 
114 _register_npu_backend()94 _register_npu_backend()
115 _register_distributed()95 _register_distributed()
116- _register_dynamo()
117 _register_rpc()96 _register_rpc()
118 _register_inductor()97 _register_inductor()
119 _register_default_gradient_device_type()98 _register_default_gradient_device_type()
Mtorch_npu/contrib/transfer_to_npu.py+15-1
@@ -2,6 +2,7 @@ import os
2import warnings2import warnings
3import json3import json
4import collections4import collections
5+import importlib
5import importlib.metadata6import importlib.metadata
6import logging as logger7import logging as logger
7import functools8import functools
@@ -526,6 +527,19 @@ def _init():
526 527 
527 _patch_jit_script()528 _patch_jit_script()
528 529 
530+ # transfer_to_npu patches these modules during its own import. Import them
531+ # explicitly instead of relying on torch_npu import side effects.
532+ # Use importlib to avoid shadowing the module-level ``torch`` binding.
533+ importlib.import_module("torch._dynamo.trace_rules")
534+ importlib.import_module("torch._dynamo.utils")
535+ importlib.import_module("torch._inductor.runtime.autotune_cache")
536+ importlib.import_module("torch._inductor.compile_fx")
537+ importlib.import_module("torch._inductor.utils")
538+ importlib.import_module("torch._inductor.fx_passes.post_grad")
539+ importlib.import_module("torch._inductor.fx_passes.joint_graph")
540+ importlib.import_module("torch._inductor.autotune_process")
541+ filesystem = importlib.import_module("torch.distributed.checkpoint.filesystem")
542+ 
529 torch._dynamo.trace_rules._disallowed_callable_ids.function_ids = None543 torch._dynamo.trace_rules._disallowed_callable_ids.function_ids = None
530 544 
531 _do_wrapper_libraries_func(_load_json_file(config_path))545 _do_wrapper_libraries_func(_load_json_file(config_path))
@@ -541,7 +555,7 @@ def _init():
541 setattr(torch._inductor.autotune_process, "get_gpu_type", _get_npu_type)555 setattr(torch._inductor.autotune_process, "get_gpu_type", _get_npu_type)
542 556 
543 setattr(torch._utils, '_get_available_device_type', _patch_get_available_device_type)557 setattr(torch._utils, '_get_available_device_type', _patch_get_available_device_type)
544- setattr(torch.distributed.checkpoint.filesystem._OverlappingCpuLoader, '__init__',558+ setattr(filesystem._OverlappingCpuLoader, '__init__',
545 _patch_OverlappingCpuLoader_init_)559 _patch_OverlappingCpuLoader_init_)
546 560 
547 _replace_to_method_in_allowed_methods()561 _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)
atomgit-bot
atomgit-botatomgit-bot8月11日

🟠 High Priority

changed line 232: return NpugraphsBackend()(gm, example_inputs, **kwargs) — 当 NpugraphsBackend.__call__ 是 @staticmethod 且签名为 (model, inputs, **compile_kwargs) 时，创建的 NpugraphsBackend() 实例会作为第一个位置参数传入，导致 model 接收到 NpugraphsBackend 实例而非 gm（graph module），inputs 接收到 gm，而 example_inputs 成为多余的位置参数，触发 TypeError: __call__() takes 2 positional arguments but 3 were given。

证据链：

修复方案：调用 NpugraphsBackend.__call__ 类方法直接访问，不创建实例： return NpugraphsBackend.__call__(gm, example_inputs, **kwargs) 或直接 return NpugraphsBackend()(gm, example_inputs) 之前先确认 NpugraphsBackend.__call__ 是否需要改为常规方法（去除 @staticmethod 并添加 self 参数）。

建议：将 NpugraphsBackend()(gm, example_inputs, **kwargs) 改为直接调用类方法 NpugraphsBackend.__call__(gm, example_inputs, **kwargs)，避免创建无用实例导致参数错位。同时检查 torch_npu/utils/_graph_tree.py 中 NpugraphsBackend.__call__ 是否需要同步调整。

改动建议
236
- return NpugraphsBackend()(gm, example_inputs, **kwargs)
236
+ return NpugraphsBackend.__call__(gm, example_inputs, **kwargs)
应用建议
likedislike
不准确？
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/_graph_tree.py+7-10
@@ -288,9 +288,13 @@ local.npu_tree_manager_containers = {}
288local.npu_tree_manager_locks = defaultdict(threading.Lock)288local.npu_tree_manager_locks = defaultdict(threading.Lock)
289 289 
290 290 
291-# only incremented by user call of mark_step_begin291+# MarkStepBox and mark_step_begin live in the lightweight _graph_tree_state
292-class MarkStepBox:292+# module so that the public NPUGraph marker can be exposed without importing
293- mark_step_counter = 0293+# the heavy graph-tree implementation.
294+from torch_npu.npu._graph_tree_state import ( # noqa: F401
295+ MarkStepBox,
296+ mark_step_begin,
297+)
294 298 
295 299 
296# We need to register this as an object that will be copied over as TLS when new300# We need to register this as an object that will be copied over as TLS when new
@@ -299,13 +303,6 @@ torch._C._stash_obj_in_tls("npu_tree_manager_containers", local.npu_tree_manager
299torch._C._stash_obj_in_tls("npu_tree_manager_locks", local.npu_tree_manager_locks)303torch._C._stash_obj_in_tls("npu_tree_manager_locks", local.npu_tree_manager_locks)
300 304 
301 305 
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:306def reset_npugraph_trees() -> None:
310 "Clear all npugraph trees"307 "Clear all npugraph trees"
311 # see shutdown below for why this is necessary308 # 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+9-3
@@ -2,11 +2,17 @@ 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+ fn._dynamo_forbidden = True
13+ return fn
14+ 
15+ 
10class _DeterministicAlgorithmsBeginOp(Function):16class _DeterministicAlgorithmsBeginOp(Function):
11 17 
12 @staticmethod18 @staticmethod
@@ -36,11 +42,11 @@ class _DeterministicAlgorithmsEndOp(Function):
36 return grad_outputs42 return grad_outputs
37 43 
38 44 
39-@forbid_in_graph45+@_forbid_in_graph
40def enable_deterministic_with_backward(tensor: Tensor):46def enable_deterministic_with_backward(tensor: Tensor):
41 return _DeterministicAlgorithmsBeginOp.apply(tensor)47 return _DeterministicAlgorithmsBeginOp.apply(tensor)
42 48 
43 49 
44-@forbid_in_graph50+@_forbid_in_graph
45def disable_deterministic_with_backward(tensor: Tensor):51def disable_deterministic_with_backward(tensor: Tensor):
46 return _DeterministicAlgorithmsEndOp.apply(tensor)52 return _DeterministicAlgorithmsEndOp.apply(tensor)
Mtorch_npu/npu/npugraph_ex/__init__.py+20-6
@@ -1,12 +1,20 @@
1+from __future__ import annotations
2+ 
1__all__ = ["compile_fx", "register_replacement"]3__all__ = ["compile_fx", "register_replacement"]
2 4 
3from collections.abc import Collection, Generator, Iterable, Mapping, Sequence5from collections.abc import Collection, Generator, Iterable, Mapping, Sequence
4-from typing import Any, Callable, NoReturn, Optional, Protocol, TypeVar, Union, Match, List6+from typing import (
7+ Any, Callable, NoReturn, Optional, Protocol, TypeVar, Union, Match, List,
8+ TYPE_CHECKING,
9+)
5 10 
6-try:11+if TYPE_CHECKING:
7- from torch._inductor.pattern_matcher import fwd_only, SearchFn, ReplaceFn, TraceFn, PatternExpr12+ from torch._inductor.pattern_matcher import (
8-except ImportError:13+ SearchFn,
9- from torch._inductor.pattern_matcher import inference_graph as fwd_only14+ ReplaceFn,
15+ TraceFn,
16+ PatternExpr,
17+ )
10 18 
11from . import inference19from . import inference
12from . import scope20from . import scope
@@ -22,10 +30,16 @@ def _return_true(match: Match):
22 30 
23 31 
24def register_replacement(search_fn: SearchFn, replace_fn: ReplaceFn, example_inputs: Iterable[Any],32def register_replacement(search_fn: SearchFn, replace_fn: ReplaceFn, example_inputs: Iterable[Any],
25- trace_fn: TraceFn = fwd_only, extra_check: Callable[[Match], bool] = _return_true,33+ trace_fn: TraceFn = None, extra_check: Callable[[Match], bool] = _return_true,
26 search_fn_pattern: Union[PatternExpr, None] = None,34 search_fn_pattern: Union[PatternExpr, None] = None,
27 scalar_workaround: Union[dict[str, Union[float, int]], None] = None,35 scalar_workaround: Union[dict[str, Union[float, int]], None] = None,
28 skip_duplicates: bool = False):36 skip_duplicates: bool = False):
37+ if trace_fn is None:
38+ try:
39+ from torch._inductor.pattern_matcher import fwd_only
40+ except ImportError:
41+ from torch._inductor.pattern_matcher import inference_graph as fwd_only
42+ trace_fn = fwd_only
29 import npugraph_ex43 import npugraph_ex
30 return npugraph_ex.patterns.pattern_pass_manager.register_replacement(search_fn, replace_fn, example_inputs,44 return npugraph_ex.patterns.pattern_pass_manager.register_replacement(search_fn, replace_fn, example_inputs,
31 trace_fn=trace_fn, extra_check=extra_check,45 trace_fn=trace_fn, extra_check=extra_check,
Mtorch_npu/utils/__init__.py+17-4
@@ -1,10 +1,23 @@
1-__all__ = ["npu_combine_tensors", "get_part_combined_tensor", "is_combined_tensor_valid", "FlopsCounter",1+__all__ = [
2- "set_thread_affinity", "reset_thread_affinity", "save_async", "get_cann_version"]2+ "npu_combine_tensors",
3+ "get_part_combined_tensor",
4+ "is_combined_tensor_valid",
5+ "FlopsCounter",
6+ "set_thread_affinity",
7+ "reset_thread_affinity",
8+ "save_async",
9+ "get_cann_version",
10+]
11+ 
3 12 
4from torch_npu.npu.utils import get_cann_version13from torch_npu.npu.utils import get_cann_version
5-from .combine_tensors import npu_combine_tensors, get_part_combined_tensor, is_combined_tensor_valid14+from .combine_tensors import (
15+ get_part_combined_tensor,
16+ is_combined_tensor_valid,
17+ npu_combine_tensors,
18+)
6from .serialization import save_async19from .serialization import save_async
7from .flops_count import _FlopsCounter as FlopsCounter20from .flops_count import _FlopsCounter as FlopsCounter
8from .affinity import _set_thread_affinity as set_thread_affinity21from .affinity import _set_thread_affinity as set_thread_affinity
9from .affinity import _reset_thread_affinity as reset_thread_affinity22from .affinity import _reset_thread_affinity as reset_thread_affinity
10-from .asd_detector import set_asd_loss_scale, register_asd_hook23+from .asd_detector import set_asd_loss_scale, register_asd_hook
Atorch_npu/utils/_device_op_overrides.py+15-0
@@ -0,0 +1,15 @@
1+from torch._inductor.codegen.common import DeviceOpOverrides
2+ 
3+ 
4+class NPUDeviceOpOverrides(DeviceOpOverrides):
5+ def import_get_raw_stream_as(self, name):
6+ return f"from torch_npu._C import _npu_getCurrentRawStream as {name}"
7+ 
8+ def set_device(self, device_idx):
9+ return f"torch_npu.npu.set_device({device_idx})"
10+ 
11+ def synchronize(self):
12+ return "torch_npu.npu.synchronize()"
13+ 
14+ def device_guard(self, device_idx):
15+ return f"torch_npu.npu._DeviceGuard({device_idx})"
Mtorch_npu/utils/_dynamo.py+266-62
@@ -1,10 +1,12 @@
1+import importlib
2+import importlib.abc
3+import functools
1import inspect4import inspect
5+import logging
2import os6import os
3import sys7import sys
4-import logging8+import threading
5from typing import Any, Optional, TYPE_CHECKING9from typing import Any, Optional, TYPE_CHECKING
6-import importlib
7-import functools
8 10 
9import torch11import torch
10import torch_npu12import torch_npu
@@ -15,8 +17,9 @@ use_jit_script = False
15log = logging.getLogger(__name__)17log = logging.getLogger(__name__)
16 18 
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,16 +71,21 @@ 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
@@ -97,6 +107,7 @@ class _InductorNpuRegistry:
97 else:107 else:
98 sys.modules["torch_npu._inductor"]._load_backend()108 sys.modules["torch_npu._inductor"]._load_backend()
99 cls._loaded_backend = current109 cls._loaded_backend = current
110+ 
100 @classmethod111 @classmethod
101 def disable_register(cls):112 def disable_register(cls):
102 cls._disabled_register = True113 cls._disabled_register = True
@@ -126,19 +137,23 @@ def register_inductor_npu():
126 _InductorNpuRegistry.register_inductor_npu()137 _InductorNpuRegistry.register_inductor_npu()
127 138 
128 139 
129-def _resolve_npu_backend_from_wrapper(wrapper) -> str:140+def _resolve_npu_backend(selected_backend=None) -> str:
130- """Resolve npu backend with priority: wrapper options > global config > env."""141+ """Resolve NPU backend with priority: compile options > config > env."""
131- wrapper_backend = wrapper.config.get("npu_backend")142+ if selected_backend not in (None, "", "default"):
132- if wrapper_backend not in (None, "", "default"):143+ return selected_backend
133- return wrapper_backend
134 144 
135- 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)
136 if global_backend not in (None, "", "default"):147 if global_backend not in (None, "", "default"):
137 return global_backend148 return global_backend
138 149 
139 return os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default")150 return os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default")
140 151 
141 152 
153+def _resolve_npu_backend_from_wrapper(wrapper) -> str:
154+ return _resolve_npu_backend(wrapper.config.get("npu_backend"))
155+ 
156+ 
142class _NpuBackendScope:157class _NpuBackendScope:
143 """Apply resolved npu backend for one compile invocation and restore env."""158 """Apply resolved npu backend for one compile invocation and restore env."""
144 159 
@@ -148,33 +163,40 @@ class _NpuBackendScope:
148 163 
149 def __enter__(self):164 def __enter__(self):
150 self._old_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")165 self._old_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")
151- os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend166+ try:
152- register_inductor_npu()167+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend
153- if self.backend == "ascendc":168+ register_inductor_npu()
154- from torch_npu._inductor.deterministic_cache import (169+ if self.backend == "ascendc":
155- patch_npu_deterministic_level_cache_keys,170+ from torch_npu._inductor.deterministic_cache import (
156- )171+ patch_npu_deterministic_level_cache_keys,
172+ )
157 173 
158- patch_npu_deterministic_level_cache_keys()174+ patch_npu_deterministic_level_cache_keys()
175+ except BaseException:
176+ self._restore_backend_env()
177+ raise
159 return self178 return self
160 179 
161 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):
162 if self._old_env is None:185 if self._old_env is None:
163 os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)186 os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)
164 else:187 else:
165 os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self._old_env188 os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self._old_env
166- return False
167 189 
168 190 
169def patch_inductor_wrapper():191def patch_inductor_wrapper():
170- from typing import Any192+ from typing import Any, Optional
171 193 
172 from torch import _TorchCompileInductorWrapper194 from torch import _TorchCompileInductorWrapper
173 from torch.utils._config_module import _ConfigEntry, Config, ConfigModule195 from torch.utils._config_module import _ConfigEntry, Config, ConfigModule
174 196 
175- src_call = _TorchCompileInductorWrapper.__call__
176 src_init = _TorchCompileInductorWrapper.__init__197 src_init = _TorchCompileInductorWrapper.__init__
177 src_get_config_copy = ConfigModule.get_config_copy198 src_get_config_copy = ConfigModule.get_config_copy
199+ src_call = _TorchCompileInductorWrapper.__call__
178 200 
179 def new_call(self, model_, inputs_):201 def new_call(self, model_, inputs_):
180 backend = _resolve_npu_backend_from_wrapper(self)202 backend = _resolve_npu_backend_from_wrapper(self)
@@ -189,12 +211,12 @@ def patch_inductor_wrapper():
189 211 
190 def new_get_config_copy(self) -> dict[str, Any]:212 def new_get_config_copy(self) -> dict[str, Any]:
191 ori_dict = src_get_config_copy(self)213 ori_dict = src_get_config_copy(self)
192- if self is not torch._inductor.config:214+ inductor_config = sys.modules.get("torch._inductor.config")
215+ if inductor_config is None or self is not inductor_config:
193 return ori_dict216 return ori_dict
194 if "npu_backend" not in ori_dict:217 if "npu_backend" not in ori_dict:
195 ori_dict["npu_backend"] = "default"218 ori_dict["npu_backend"] = "default"
196 cfg = Config(default="default", value_type=str)219 cfg = Config(default="default", value_type=str)
197- # PyTorch >=2.12 added a required `name` arg to _ConfigEntry.
198 if "name" in inspect.signature(_ConfigEntry.__init__).parameters:220 if "name" in inspect.signature(_ConfigEntry.__init__).parameters:
199 self._config["npu_backend"] = _ConfigEntry(cfg, "npu_backend")221 self._config["npu_backend"] = _ConfigEntry(cfg, "npu_backend")
200 else:222 else:
@@ -202,11 +224,20 @@ def patch_inductor_wrapper():
202 return ori_dict224 return ori_dict
203 225 
204 def new_init(self, mode, options, dynamic, name=None):226 def new_init(self, mode, options, dynamic, name=None):
205- add_dynamo_methods_init()227+ self._npu_defer_shape_handling = True
206- if name is not None:228+ self._npu_shape_handling_requested = False
207- src_init(self, mode, options, dynamic, name)229+ try:
208- else:230+ if name is not None:
209- src_init(self, mode, options, dynamic)231+ src_init(self, mode, options, dynamic, name)
232+ else:
233+ src_init(self, mode, options, dynamic)
234+ shape_handling_requested = self._npu_shape_handling_requested
235+ finally:
236+ del self._npu_defer_shape_handling
237+ del self._npu_shape_handling_requested
238+ _setup_inductor_for_compile(self.config)
239+ if shape_handling_requested:
240+ torch_npu._inductor.patch_shape_handling()
210 backend = _resolve_npu_backend_from_wrapper(self)241 backend = _resolve_npu_backend_from_wrapper(self)
211 if backend == "mlir":242 if backend == "mlir":
212 with _NpuBackendScope(backend):243 with _NpuBackendScope(backend):
@@ -220,15 +251,14 @@ def patch_inductor_wrapper():
220 _TorchCompileInductorWrapper.__call__ = new_call251 _TorchCompileInductorWrapper.__call__ = new_call
221 _TorchCompileInductorWrapper.__init__ = new_init252 _TorchCompileInductorWrapper.__init__ = new_init
222 ConfigModule.get_config_copy = new_get_config_copy253 ConfigModule.get_config_copy = new_get_config_copy
223- torch._inductor.config.get_config_copy()
224 254 
225 255 
226def patch_dynamo_optimize():256def patch_dynamo_optimize():
227 from torch_npu.dynamo import _get_global_npu_backend257 from torch_npu.dynamo import _get_global_npu_backend
258+ 
228 src_optimize = torch._dynamo.optimize259 src_optimize = torch._dynamo.optimize
229 260 
230 def npu_optimize(*args, **kwargs):261 def npu_optimize(*args, **kwargs):
231- add_dynamo_methods_init()
232 backend = None262 backend = None
233 if "backend" in kwargs:263 if "backend" in kwargs:
234 backend = kwargs["backend"]264 backend = kwargs["backend"]
@@ -417,40 +447,134 @@ def patch_user_defined_class_variable():
417 UserDefinedClassVariable.__new__raw = UserDefinedClassVariable.__new__447 UserDefinedClassVariable.__new__raw = UserDefinedClassVariable.__new__
418 UserDefinedClassVariable.__new__ = UserDefinedClassVariable__new__448 UserDefinedClassVariable.__new__ = UserDefinedClassVariable__new__
419 449 
450+ 
420def run_once(f):451def run_once(f):
421- """Runs a function (successfully) only once.452+ """Run a function successfully only once, waiting for concurrent callers."""
422- The running can be reset by setting the `has_run` attribute to False453+ condition = threading.Condition()
423- """454+ 
424 @functools.wraps(f)455 @functools.wraps(f)
425 def wrapper(*args, **kwargs):456 def wrapper(*args, **kwargs):
426- if not wrapper.has_run:457+ thread_id = threading.get_ident()
458+ with condition:
459+ while wrapper._is_running:
460+ if wrapper._running_thread == thread_id:
461+ return None
462+ condition.wait()
463+ if wrapper.has_run:
464+ return None
465+ wrapper._is_running = True
466+ wrapper._running_thread = thread_id
467+ 
468+ try:
427 result = f(*args, **kwargs)469 result = f(*args, **kwargs)
470+ except BaseException:
471+ with condition:
472+ wrapper._is_running = False
473+ wrapper._running_thread = None
474+ condition.notify_all()
475+ raise
476+ 
477+ with condition:
428 wrapper.has_run = True478 wrapper.has_run = True
429- return result479+ wrapper._is_running = False
430- return None480+ wrapper._running_thread = None
481+ condition.notify_all()
482+ return result
483+ 
431 wrapper.has_run = False484 wrapper.has_run = False
485+ wrapper._is_running = False
486+ wrapper._running_thread = None
487+ 
488+ def reset_after_fork():
489+ # The parent thread running f may not exist in the child process.
490+ nonlocal condition
491+ condition = threading.Condition()
492+ wrapper._is_running = False
493+ wrapper._running_thread = None
494+ 
495+ try:
496+ os.register_at_fork(after_in_child=reset_after_fork)
497+ except AttributeError:
498+ pass
432 return wrapper499 return wrapper
433 500 
434 501 
435-@run_once502+_COMPLETED_DYNAMO_SETUP_STEPS = set()
436-def _dynamo_register_interface_for_device():
437- from torch._dynamo.device_interface import register_interface_for_device
438- from torch_npu.utils._dynamo_device import NpuInterface
439 503 
440- register_interface_for_device("npu", NpuInterface)
441- for i in range(32):
442 504 
443- register_interface_for_device(f"npu:{i}", NpuInterface)505+def _run_dynamo_setup_step(name, setup):
506+ """Keep successful setup steps idempotent when a later step fails."""
507+ if name in _COMPLETED_DYNAMO_SETUP_STEPS:
508+ return
509+ setup()
510+ _COMPLETED_DYNAMO_SETUP_STEPS.add(name)
511+ 
512+ 
513+def _find_spec_without_finder(finder, fullname):
514+ """Delegate to the remaining meta-path finders without bypassing them."""
515+ try:
516+ index = sys.meta_path.index(finder)
517+ except ValueError:
518+ return importlib.util.find_spec(fullname)
519+ 
520+ sys.meta_path.pop(index)
521+ try:
522+ return importlib.util.find_spec(fullname)
523+ finally:
524+ sys.meta_path.insert(min(index, len(sys.meta_path)), finder)
525+ 
526+ 
527+class _DynamoPostImportLoader(importlib.abc.Loader):
528+ def __init__(self, loader, finder):
529+ self._loader = loader
530+ self._finder = finder
531+ 
532+ def create_module(self, spec):
533+ create_module = getattr(self._loader, "create_module", None)
534+ return create_module(spec) if create_module is not None else None
535+ 
536+ def exec_module(self, module):
537+ self._loader.exec_module(module)
538+ _lazy_dynamo_setup()
539+ if self._finder in sys.meta_path:
540+ sys.meta_path.remove(self._finder)
541+ 
542+ 
543+class _DynamoPostImportFinder(importlib.abc.MetaPathFinder):
544+ _target = "torch._dynamo"
545+ 
546+ def find_spec(self, fullname, path=None, target=None):
547+ if fullname != self._target:
548+ return None
549+ spec = _find_spec_without_finder(self, fullname)
550+ if spec is not None and spec.loader is not None:
551+ spec.loader = _DynamoPostImportLoader(spec.loader, self)
552+ return spec
553+ 
554+ 
555+def _install_dynamo_post_import_trigger():
556+ """Set up NPU integration whenever Dynamo is first imported."""
557+ if "torch._dynamo" in sys.modules:
558+ _lazy_dynamo_setup()
559+ return
560+ if not any(isinstance(finder, _DynamoPostImportFinder) for finder in sys.meta_path):
561+ sys.meta_path.insert(0, _DynamoPostImportFinder())
562+ 
444 563 
445@run_once564@run_once
446def add_dynamo_methods_init():565def add_dynamo_methods_init():
447- _dynamo_register_interface_for_device()566+ steps = (
448- patch_SkipFunctionVariable()567+ ("device_interface", _dynamo_register_interface_for_device),
449- patch_TensorVariable_call_method()568+ ("skip_function_variable", patch_SkipFunctionVariable),
450- patch_stream_event_variable_python_type()569+ ("tensor_variable", patch_TensorVariable_call_method),
451- patch_npu_stream_context()570+ ("user_defined_class_variable", patch_user_defined_class_variable),
452- patch_npu_current_stream()571+ ("stream_event_variable", patch_stream_event_variable_python_type),
453- patch_user_defined_class_variable()572+ ("npu_stream_context", patch_npu_stream_context),
573+ ("npu_current_stream", patch_npu_current_stream),
574+ ("builtin_variable", patch_builtin_variable),
575+ )
576+ for name, setup in steps:
577+ _run_dynamo_setup_step(name, setup)
454 578 
455 579 
456@functools.lru_cache(None)580@functools.lru_cache(None)
@@ -492,11 +616,91 @@ def has_triton() -> bool:
492 616 
493 617 
494def patch_has_triton():618def patch_has_triton():
495- torch.utils._triton.has_triton = has_triton619+ from torch.utils import _triton
620+ 
621+ _triton.has_triton = has_triton
622+ 
623+ 
624+@run_once
625+def _inject_inductor_npu_backend_config():
626+ """Inject NPU entries into torch._inductor.config on first use."""
627+ torch._inductor.config.get_config_copy()
628+ 
629+ 
630+@run_once
631+def _lazy_dynamo_setup():
632+ """Initialize the Dynamo integration on the first graph-capture operation."""
633+ add_dynamo_methods_init()
634+ 
635+ from torch_npu.dynamo import _register_backends
636+ _run_dynamo_setup_step("backends", _register_backends)
637+ 
638+ from torch_npu.dynamo.trace_rule import _patch_npu_trace_rules
639+ _run_dynamo_setup_step("trace_rules", _patch_npu_trace_rules)
640+ 
641+ _run_dynamo_setup_step("dynamo_optimize", patch_dynamo_optimize)
642+ 
643+ 
644+@run_once
645+def _lazy_inductor_setup():
646+ """Initialize NPU Inductor support only for an Inductor-based backend."""
647+ register_inductor_npu()
648+ 
649+ from torch_npu.utils._graph_tree import _apply_npugraph_tree_methods
650+ _apply_npugraph_tree_methods()
651+ 
652+ _inject_inductor_npu_backend_config()
653+ 
654+ 
655+def _setup_inductor_for_compile(options=None):
656+ """Initialize the NPU Inductor backend selected for this compile call."""
657+ _lazy_dynamo_setup()
658+ 
659+ option_backend = options.get("npu_backend") if isinstance(options, dict) else None
660+ selected_backend = _resolve_npu_backend(option_backend)
661+ 
662+ old_backend = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")
663+ if selected_backend not in (None, "", "default"):
664+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = selected_backend
665+ try:
666+ _lazy_inductor_setup()
667+ finally:
668+ if old_backend is None:
669+ os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)
670+ else:
671+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = old_backend
672+ return selected_backend
673+ 
674+ 
675+@run_once
676+def install_npugraph_mark_step_trigger():
677+ """Expose the public NPUGraph step API without importing compiler internals."""
678+ def npugraph_mark_step_begin():
679+ from torch_npu.npu._graph_tree_state import mark_step_begin
680+ return mark_step_begin()
681+ 
682+ torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin
683+ 
684+ 
685+@run_once
686+def _dynamo_register_interface_for_device():
687+ from torch._dynamo.device_interface import register_interface_for_device
688+ from torch_npu.utils._dynamo_device import NpuInterface
689+ 
690+ register_interface_for_device("npu", NpuInterface)
691+ for i in range(32):
692+ register_interface_for_device(f"npu:{i}", NpuInterface)
496 693 
497 694 
498def add_dynamo_methods():695def add_dynamo_methods():
499- patch_dynamo_optimize()
500- patch_builtin_variable()
501- patch_inductor_wrapper()
502 patch_has_triton()696 patch_has_triton()
697+ 
698+ from torch_npu.dynamo import _install_lazy_torchair
699+ 
700+ _install_lazy_torchair()
701+ _install_dynamo_post_import_trigger()
702+ if "npugraph_ex" not in sys.modules:
703+ from torch_npu.dynamo import _LazyNpuGraphEx
704+ sys.modules["npugraph_ex"] = _LazyNpuGraphEx("npugraph_ex")
705+ patch_inductor_wrapper()
706+ install_npugraph_mark_step_trigger()
Mtorch_npu/utils/_graph_tree.py+7-4
@@ -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,
@@ -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-22
@@ -1,29 +1,8 @@
1from typing import Optional1from typing import Optional
2 2 
3import torch3import torch
4-from torch._inductor.codegen.common import DeviceOpOverrides, register_device_op_overrides
5from torch._prims.rng_prims import register_rng_prim4from torch._prims.rng_prims import register_rng_prim
6 5 
7- 
8-class NPUDeviceOpOverrides(DeviceOpOverrides):
9- def import_get_raw_stream_as(self, name):
10- return f"from torch_npu._C import _npu_getCurrentRawStream as {name}"
11- 
12- def set_device(self, device_idx):
13- return f"torch_npu.npu.set_device({device_idx})"
14- 
15- def synchronize(self):
16- return "torch_npu.npu.synchronize()"
17- 
18- def device_guard(self, device_idx):
19- return f"torch_npu.npu._DeviceGuard({device_idx})"
20- 
21- 
22-def _inductor_register_device_op_overrides():
23- from torch._inductor.codegen import cpu_device_op_overrides, mps_device_op_overrides
24- register_device_op_overrides('npu', NPUDeviceOpOverrides())
25- 
26- 
27def patch_philox_rand_offset():6def patch_philox_rand_offset():
28 def get_philox_rand_offset_patch(shape):7 def get_philox_rand_offset_patch(shape):
29 numel_scalar = 18 numel_scalar = 1
@@ -240,4 +219,4 @@ patch_register_run_and_save_rng_state_op()
240patch_register_run_with_rng_state_op()219patch_register_run_with_rng_state_op()
241patch_philox_rand_offset()220patch_philox_rand_offset()
242patch_register_philox_rand()221patch_register_philox_rand()
243-patch_rng_prims_device()222+patch_rng_prims_device()