[fix]import_all_patch-pytorch-AtomGit

Mtest/_inductor/test_lazy_register.py+0-6

	@@ -8,13 +8,7 @@ import torch_npu
8		8
9	@skipIf(torch_npu.utils._dynamo.is_inductor_npu_initialized(), reason="Inductor npu has initialized")	9	@skipIf(torch_npu.utils._dynamo.is_inductor_npu_initialized(), reason="Inductor npu has initialized")
10	class TestLazyRegister(TestUtils):	10	class TestLazyRegister(TestUtils):
11	- def test_compile_but_not_invoked(self):
12		11
13	- def run(x, y):
14	- return x + y
15	-
16	- run = torch.compile(run)
17	- self.assertFalse(torch_npu.utils._dynamo.is_inductor_npu_initialized())
18		12
19	def test_disable_register_inductor_npu(self):	13	def test_disable_register_inductor_npu(self):
20	torch_npu.utils._dynamo.disable_register_inductor_npu()	14	torch_npu.utils._dynamo.disable_register_inductor_npu()

Mtest/_inductor/test_mlir_enable.py+0-1

	@@ -3,7 +3,6 @@ from torch.testing._internal.common_utils import run_tests, parametrize, instant
3	from testutils import TestUtils	3	from testutils import TestUtils
4	from torch._inductor.utils import run_and_get_code	4	from torch._inductor.utils import run_and_get_code
5	import torch_npu	5	import torch_npu
6	-import torch_npu._inductor
7		6
8		7
9	class TestAdd(TestUtils):	8	class TestAdd(TestUtils):

Mtorch_npu/_inductor/__init__.py+63-27

	@@ -1,6 +1,5 @@
1	import os	1	import os
2		2
3	-
4	ORG_AUTOLOAD = os.getenv("TORCH_DEVICE_BACKEND_AUTOLOAD", "1")	3	ORG_AUTOLOAD = os.getenv("TORCH_DEVICE_BACKEND_AUTOLOAD", "1")
5	os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"	4	os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"
6	from torch._inductor.async_compile import AsyncCompile	5	from torch._inductor.async_compile import AsyncCompile
	@@ -10,25 +9,51 @@ AsyncCompile.warm_pool()
10	os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = ORG_AUTOLOAD	9	os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = ORG_AUTOLOAD
11		10
12	# all backends need register npu/cpu/mps device_op_overrides	11	# all backends need register npu/cpu/mps device_op_overrides
		12	+from .codecache import patch_cache_base_get_system
		13	+
		14	+# All backends need npu/cpu/mps device_op_overrides.
13	from .codegen.common import register_device_op_overrides_npu	15	from .codegen.common import register_device_op_overrides_npu
14	-	16	+from .graph import patch_codegen_with_cpp_wrapper, patch_count_bytes, patch_run_node
15	-	17	+from .shape_handling import NPUShapeHandling, patch_shape_handling
		18	+from .utils import patch_has_triton, patch_has_triton_tma, patch_is_gpu
		19	+from .autotune_process import patch_tuning_process, patch_tuning_process_pool
		20	+from .codegen.cpp_utils import patch_device_to_aten
16	register_device_op_overrides_npu()	21	register_device_op_overrides_npu()
17		22
18	-if os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default") == "mlir":	23	+patch_has_triton()
		24	+patch_has_triton_tma()
		25	+patch_is_gpu()
		26	+patch_cache_base_get_system()
		27	+patch_codegen_with_cpp_wrapper()
		28	+patch_count_bytes()
		29	+patch_run_node()
		30	+patch_tuning_process()
		31	+patch_tuning_process_pool()
		32	+patch_device_to_aten()
		33	+def _get_backend() -> str:
		34	+ return os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default")
		35	+
		36	+if _get_backend() == "mlir":
		37	+ import torch
		38	+ import torch_npu
19	try:	39	try:
20	import torch_mlir	40	import torch_mlir
21	from torch_mlir import ir	41	from torch_mlir import ir
22	except ImportError as e:	42	except ImportError as e:
23	raise ImportError("torch_mlir is not installed, install it first.") from e	43	raise ImportError("torch_mlir is not installed, install it first.") from e
24	from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin, torch_mlir_patch	44	from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin, torch_mlir_patch
		45	+ device_id = torch_npu.npu.current_device()
		46	+ torch_npu._C._recovery_all_npu_stream(device_id)
25		47
26	-elif os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default") == "dvm":	48	+elif _get_backend() == "dvm":
27	from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin	49	from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin
28	from .dvm import mlir_fusion	50	from .dvm import mlir_fusion
		51	+
		52	+
29	else:	53	else:
30	import os	54	import os
31	-	55	+ import logging
		56	+ log = logging.getLogger(__name__)
32	import torch	57	import torch
33	from torch._dynamo.device_interface import get_interface_for_device	58	from torch._dynamo.device_interface import get_interface_for_device
34	from torch._inductor import lowering as inductor_lowering	59	from torch._inductor import lowering as inductor_lowering
	@@ -40,10 +65,8 @@ else:
40		65
41	from . import codegen, config as npu_config	66	from . import codegen, config as npu_config
42	from .async_compile import patch_async_compile	67	from .async_compile import patch_async_compile
43	- from .autotune_process import patch_tuning_process, patch_tuning_process_pool	68	+ from .codecache import patch_aot_code_compiler_compile
44	- from .codecache import patch_aot_code_compiler_compile, patch_cache_base_get_system
45	from .codegen._sizevars import patch_simplify	69	from .codegen._sizevars import patch_simplify
46	- from .codegen.cpp_utils import patch_device_to_aten
47	from .codegen.ir import patch_indexing, patch_loop_body	70	from .codegen.ir import patch_indexing, patch_loop_body
48	from .codegen.triton import (	71	from .codegen.triton import (
49	patch_gen_common_triton_ext_imports,	72	patch_gen_common_triton_ext_imports,
	@@ -66,9 +89,8 @@ else:
66	post_grad_custom_pass_fuc,	89	post_grad_custom_pass_fuc,
67	pre_grad_custom_pass_fuc,	90	pre_grad_custom_pass_fuc,
68	)	91	)
69	- from .fx_passes.joint_graph import patch_constant_fold_uniform_value
70	from .fx_passes.pattern_match.npu_fusion_attention_graph import register_fa_pass	92	from .fx_passes.pattern_match.npu_fusion_attention_graph import register_fa_pass
71	- from .graph import patch_codegen_with_cpp_wrapper, patch_count_bytes, patch_run_node	93	+ from .fx_passes.joint_graph import patch_constant_fold_uniform_value
72	from .ir import patch_fallback_kernel_codegen, patch_num_splits	94	from .ir import patch_fallback_kernel_codegen, patch_num_splits
73	from .kernel import (	95	from .kernel import (
74	_register_npu_inductor_addmm,	96	_register_npu_inductor_addmm,
	@@ -87,12 +109,8 @@ else:
87	from .scheduler import patch_scheduler	109	from .scheduler import patch_scheduler
88	from .select_algorithm import patch_algorithm_selector	110	from .select_algorithm import patch_algorithm_selector
89	from .shape_handling import NPUShapeHandling, patch_shape_handling	111	from .shape_handling import NPUShapeHandling, patch_shape_handling
90	- from .utils import (	112	+ from .utils import patch_get_first_incompatible_cudagraph_node
91	- patch_get_first_incompatible_cudagraph_node,	113	+
92	- patch_has_triton,
93	- patch_has_triton_tma,
94	- patch_is_gpu,
95	- )
96		114
97	flex_attention._validate_device = _validate_device	115	flex_attention._validate_device = _validate_device
98		116
	@@ -111,9 +129,7 @@ else:
111		129
112	inductor_lowering.make_reduction = make_reduction	130	inductor_lowering.make_reduction = make_reduction
113		131
114	- patch_codegen_with_cpp_wrapper()
115	patch_get_cpp_torch_device_options()	132	patch_get_cpp_torch_device_options()
116	- patch_device_to_aten()
117	patch_constant_fold_uniform_value()	133	patch_constant_fold_uniform_value()
118	patch_fallback_kernel_codegen()	134	patch_fallback_kernel_codegen()
119	patch_aot_code_compiler_compile()	135	patch_aot_code_compiler_compile()
	@@ -143,8 +159,6 @@ else:
143		159
144	patch_pattern_mm_plus_mm()	160	patch_pattern_mm_plus_mm()
145	patch_algorithm_selector()	161	patch_algorithm_selector()
146	- patch_tuning_process()
147	- patch_tuning_process_pool()
148	patch_async_compile()	162	patch_async_compile()
149	patch_scheduler()	163	patch_scheduler()
150	patch_gen_common_triton_ext_imports()	164	patch_gen_common_triton_ext_imports()
	@@ -189,12 +203,6 @@ else:
189	_replace_precompile()	203	_replace_precompile()
190		204
191	register_fa_pass()	205	register_fa_pass()
192	- patch_cache_base_get_system()
193	- patch_count_bytes()
194	- patch_run_node()
195	- patch_is_gpu()
196	- patch_has_triton()
197	- patch_has_triton_tma()
198	patch_get_first_incompatible_cudagraph_node()	206	patch_get_first_incompatible_cudagraph_node()
199	patch_get_optimization_cflags()	207	patch_get_optimization_cflags()
200	patch_extract_read_writes()	208	patch_extract_read_writes()
	@@ -253,3 +261,31 @@ else:
253	)	261	)
254		262
255	add_additional_op()	263	add_additional_op()
		264	+ torch._inductor.config.comprehensive_padding = False
		265	+
		266	+ compile_threads = int(
		267	+ os.environ.get("TORCHINDUCTOR_COMPILE_THREADS") or "1"
		268	+ )
		269	+ os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = str(compile_threads)
		270	+ torch._inductor.config.compile_threads = compile_threads
		271	+
		272	+ _fasta_autotune = os.environ.get("FASTAUTOTUNE", "0") == "1"
		273	+ _fasta_autotune_method = os.getenv("AUTOTUNE_METHOD", "Expert")
		274	+ if _fasta_autotune:
		275	+ if os.environ.get("ENABLE_PRINT_UB_BITS", "0") == "0":
		276	+ log.warnings(
		277	+ "Please set ENABLE_PRINT_UB_BITS to 1. Fasta autotune need to know real ub usage."
		278	+ )
		279	+ os.environ["ENABLE_PRINT_UB_BITS"] = "1"
		280	+
		281	+ if (
		282	+ _fasta_autotune_method == "SampleStack"
		283	+ and torch._inductor.config.compile_threads != 1
		284	+ ):
		285	+ log.warnings(
		286	+ "fasta SampleStack method is not temporarily compatible with multi-process compile, "
		287	+ "fasta_autotune set TORCHINDUCTOR_COMPILE_THREADS "
		288	+ f"from {torch._inductor.config.compile_threads} to 1."
		289	+ )
		290	+ os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
		291	+ torch._inductor.config.compile_threads = 1

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/npu_inductor_plugin.py+0-3

	@@ -65,9 +65,6 @@ from .. import config as anir_config
65	from . import npu_patch_deprecated	65	from . import npu_patch_deprecated
66	from .npu_meta import npu_patch_meta	66	from .npu_meta import npu_patch_meta
67		67
68	-_triton.has_triton = lambda: False
69	-_triton.has_triton_package = lambda: False
70	-
71	# Fix Error: Exit earlier than child process.	68	# Fix Error: Exit earlier than child process.
72	atexit.register(shutdown_compile_workers)	69	atexit.register(shutdown_compile_workers)
73		70

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/npu_patch_deprecated.py+1-9

	@@ -26,15 +26,7 @@ def py_contiguous(x, memory_format=torch.contiguous_format):
26	return x.clone(memory_format=memory_format)	26	return x.clone(memory_format=memory_format)
27		27
28		28
29	-@staticmethod
30	-@functools.lru_cache(None)
31	-def _patch_get_system() -> Dict[str, Any]:
32	- system = {}
33	- system["hash"] = hashlib.sha256(
34	- json.dumps(system, sort_keys=True).encode("utf-8")
35	- ).hexdigest()
36		29
37	- return system
38		30
39	def _patch_add_ephemeral_timeout_for_all_pgs(timeout: timedelta) -> None:	31	def _patch_add_ephemeral_timeout_for_all_pgs(timeout: timedelta) -> None:
40	"""	32	"""
	@@ -60,7 +52,7 @@ def _patch_add_ephemeral_timeout_for_all_pgs(timeout: timedelta) -> None:
60	if torch.device("npu") in devices:	52	if torch.device("npu") in devices:
61	backend = pg._get_backend(torch.device("npu"))	53	backend = pg._get_backend(torch.device("npu"))
62		54
63	-CacheBase.get_system = _patch_get_system	55	+
64	distributed_c10d._add_ephemeral_timeout_for_all_pgs = _patch_add_ephemeral_timeout_for_all_pgs	56	distributed_c10d._add_ephemeral_timeout_for_all_pgs = _patch_add_ephemeral_timeout_for_all_pgs
65		57
66	if get_anir_mode() == 'O0':	58	if get_anir_mode() == 'O0':

Mtorch_npu/utils/_dynamo.py+13-68

	@@ -116,15 +116,18 @@ def TensorVariable_call_method(self, tx, name, args, kwargs):
116		116
117	class _InductorNpuRegistry:	117	class _InductorNpuRegistry:
118	_disabled_register = False	118	_disabled_register = False
119	- _has_inited = False	119	+ _loaded_backend: str \| None = None
120		120
121	@classmethod	121	@classmethod
122	def register_inductor_npu(cls):	122	def register_inductor_npu(cls):
123	- if cls.has_initialized() or cls._disabled_register:	123	+ if cls._disabled_register:
124	return	124	return
125	- import torch_npu._inductor # noqa:F401
126		125
127	- cls._has_inited = True	126	+ current = os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default")
		127	+ if cls._loaded_backend != current:
		128	+ import torch_npu._inductor # noqa:F401
		129	+
		130	+ cls._loaded_backend = current
128		131
129	@classmethod	132	@classmethod
130	def disable_register(cls):	133	def disable_register(cls):
	@@ -136,12 +139,7 @@ class _InductorNpuRegistry:
136		139
137	@classmethod	140	@classmethod
138	def has_initialized(cls):	141	def has_initialized(cls):
139	- if cls._has_inited:	142	+ return cls._loaded_backend is not None
140	- return True
141	- # Maybe initialized by call `import torch_npu._inductor` manually.
142	- if "torch_npu._inductor" in sys.modules:
143	- cls._has_inited = True
144	- return cls._has_inited
145		143
146		144
147	def is_inductor_npu_initialized():	145	def is_inductor_npu_initialized():
	@@ -166,14 +164,10 @@ def patch_inductor_wrapper():
166	from torch import _TorchCompileInductorWrapper	164	from torch import _TorchCompileInductorWrapper
167	from torch.utils._config_module import _ConfigEntry, Config, ConfigModule	165	from torch.utils._config_module import _ConfigEntry, Config, ConfigModule
168		166
169	- src_call = _TorchCompileInductorWrapper.__call__
170	src_apply_options = _TorchCompileInductorWrapper.apply_options	167	src_apply_options = _TorchCompileInductorWrapper.apply_options
171	src_init = _TorchCompileInductorWrapper.__init__	168	src_init = _TorchCompileInductorWrapper.__init__
172	src_get_config_copy = ConfigModule.get_config_copy	169	src_get_config_copy = ConfigModule.get_config_copy
173		170
174	- def new_call(self, model_, inputs_):
175	- register_inductor_npu()
176	- return src_call(self, model_, inputs_)
177		171
178	def new_apply_options(self, options: Optional[dict[str, Any]]):	172	def new_apply_options(self, options: Optional[dict[str, Any]]):
179	if options is not None and options.get("enable_shape_handling", False):	173	if options is not None and options.get("enable_shape_handling", False):
	@@ -217,66 +211,17 @@ def patch_inductor_wrapper():
217	or torch._inductor.config.npu_backend == "mlir"	211	or torch._inductor.config.npu_backend == "mlir"
218	):	212	):
219	os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "mlir"	213	os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "mlir"
220	- device_id = torch_npu.npu.current_device()	214	+
221	- torch_npu._C._recovery_all_npu_stream(device_id)
222	- try:
223	- import torch_mlir # noqa:F401
224	- except ImportError as e:
225	- raise ImportError(
226	- "torch_mlir is not installed, install it first."
227	- ) from e
228	- importlib.import_module(
229	- "torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.npu_inductor_plugin"
230	- )
231	- importlib.import_module(
232	- "torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.torch_mlir_patch"
233	- )
234		215
235	elif (	216	elif (
236	self.config.get("npu_backend") == "dvm"	217	self.config.get("npu_backend") == "dvm"
237	or torch._inductor.config.npu_backend == "dvm"	218	or torch._inductor.config.npu_backend == "dvm"
238	):	219	):
239	os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "dvm"	220	os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "dvm"
240	- importlib.import_module(	221	+
241	- "torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.npu_inductor_plugin"	222	+ register_inductor_npu()
242	- )	223	+
243	- importlib.import_module("torch_npu._inductor.dvm.mlir_fusion")	224	+
244	-
245	- elif (
246	- self.config.get("npu_backend") == "default"
247	- or torch._inductor.config.npu_backend == "default"
248	- ):
249	- os.environ["TORCHINDUCTOR_COMPREHENSIVE_PADDING"] = "0"
250	- torch._inductor.config.comprehensive_padding = False
251	-
252	- compile_threads = int(
253	- os.environ.get("TORCHINDUCTOR_COMPILE_THREADS") or "1"
254	- )
255	- os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = str(compile_threads)
256	- torch._inductor.config.compile_threads = compile_threads
257	-
258	- _fasta_autotune = os.environ.get("FASTAUTOTUNE", "0") == "1"
259	- _fasta_autotune_method = os.getenv("AUTOTUNE_METHOD", "Expert")
260	- if _fasta_autotune:
261	- if os.environ.get("ENABLE_PRINT_UB_BITS", "0") == "0":
262	- log.warnings(
263	- "Please set ENABLE_PRINT_UB_BITS to 1. Fasta autotune need to know real ub usage."
264	- )
265	- os.environ["ENABLE_PRINT_UB_BITS"] = "1"
266	-
267	- if (
268	- _fasta_autotune_method == "SampleStack"
269	- and torch._inductor.config.compile_threads != 1
270	- ):
271	- log.warnings(
272	- "fasta SampleStack method is not temporarily compatible with multi-process compile, "
273	- "fasta_autotune set TORCHINDUCTOR_COMPILE_THREADS "
274	- f"from {torch._inductor.config.compile_threads} to 1."
275	- )
276	- os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
277	- torch._inductor.config.compile_threads = 1
278	-
279	- _TorchCompileInductorWrapper.__call__ = new_call
280	_TorchCompileInductorWrapper.apply_options = new_apply_options	225	_TorchCompileInductorWrapper.apply_options = new_apply_options
281	_TorchCompileInductorWrapper.__init__ = new_init	226	_TorchCompileInductorWrapper.__init__ = new_init
282	ConfigModule.get_config_copy = new_get_config_copy	227	ConfigModule.get_config_copy = new_get_config_copy