【rfc】mlir_triton_lowering_file_refactoring-pytorch-AtomGit

Atest/_inductor/test_multi_backend_compile.py+43-0

@@ -0,0 +1,43 @@
	1	+# Owner(s): ["module: tests"]
	2	+import torch
	3	+from torch._inductor.utils import run_and_get_code
	4	+from torch.testing._internal.common_utils import (
	5	+ TestCase,
	6	+ instantiate_parametrized_tests,
	7	+ parametrize,
	8	+ run_tests,
	9	+)
	10	+
	11	+import torch_npu
	12	+
	13	+
	14	+class TestMultiBackendMixedCompile(TestCase):
	15	+ @parametrize("dtype", [torch.float32])
	16	+ def test_mixed_decorators_in_one_scope(self, dtype):
	17	+
	18	+ a = torch.randn(2, 2, dtype=dtype, device="npu")
	19	+ b = torch.randn(2, 2, dtype=dtype, device="npu")
	20	+ x = torch.randn(3, 4, dtype=dtype, device="npu")
	21	+ y = torch.randn(3, 4, dtype=dtype, device="npu")
	22	+
	23	+ @torch.compile()
	24	+ def op_add(x, y):
	25	+ return x + y
	26	+
	27	+ add_out, add_codes = run_and_get_code(op_add, x, y)
	28	+ self.assertEqual(x + y, add_out, atol=1e-3, rtol=1e-3)
	29	+ self.assertIn("triton", add_codes[0])
	30	+
	31	+ @torch.compile(options={"npu_backend": "mlir"})
	32	+ def op_sub(a, b):
	33	+ return a - b
	34	+
	35	+ sub_out, sub_codes = run_and_get_code(op_sub, a, b)
	36	+ self.assertEqual(a - b, sub_out, atol=1e-3, rtol=1e-3)
	37	+ self.assertIn("mlir", sub_codes[0])
	38	+
	39	+
	40	+instantiate_parametrized_tests(TestMultiBackendMixedCompile)
	41	+
	42	+if __name__ == "__main__":
	43	+ run_tests()

Mtorch_npu/_inductor/__init__.py+18-7

	@@ -1,9 +1,6 @@
1	import os	1	import os
2		2
3	# all backends need register npu/cpu/mps device_op_overrides	3	# all backends need register npu/cpu/mps device_op_overrides
4	-
5	-from torch._inductor.lowering import make_fallback as _ori_make_fallback
6	-# All backends need npu/cpu/mps device_op_overrides.
7	from .codegen.common import register_device_op_overrides_npu, patch_cache_base_get_system	4	from .codegen.common import register_device_op_overrides_npu, patch_cache_base_get_system
8	from .graph import patch_codegen_with_cpp_wrapper	5	from .graph import patch_codegen_with_cpp_wrapper
9	from .utils import patch_has_triton, patch_device_supports_tma, patch_is_gpu, get_current_raw_stream	6	from .utils import patch_has_triton, patch_device_supports_tma, patch_is_gpu, get_current_raw_stream
	@@ -26,13 +23,23 @@ def _load_mlir_backend():
26	from torch_mlir import ir	23	from torch_mlir import ir
27	except ImportError as e:	24	except ImportError as e:
28	raise ImportError("torch_mlir is not installed, install it first.") from e	25	raise ImportError("torch_mlir is not installed, install it first.") from e
29	- global _ori_make_fallback	26	+ from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin, torch_mlir_patch
30	- torch._inductor.lowering.make_fallback = _ori_make_fallback	27	+ from .lowering_patch import apply_mlir_inductor_patch
31	- from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin	28	+ from .ascend_npu_ir.ascend_npu_ir.npu.npu_inductor_plugin import (
32	- from .ascend_npu_ir.ascend_npu_ir.npu import torch_mlir_patch	29	+ register_mlir_codegen_backend,
		30	+ )
		31	+
		32	+ apply_mlir_inductor_patch()
		33	+ register_mlir_codegen_backend()
33		34
34	def _load_dvm_backend():	35	def _load_dvm_backend():
35	from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin	36	from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin
		37	+ from .lowering_patch import apply_mlir_inductor_patch
		38	+ from .ascend_npu_ir.ascend_npu_ir.npu.npu_inductor_plugin import (
		39	+ register_mlir_codegen_backend,
		40	+ )
		41	+ apply_mlir_inductor_patch()
		42	+ register_mlir_codegen_backend()
36	from .dvm import mlir_fusion	43	from .dvm import mlir_fusion
37		44
38	def _load_triton_backend():	45	def _load_triton_backend():
	@@ -155,6 +162,10 @@ _BACKEND_LOADERS = {
155		162
156		163
157	def _load_backend():	164	def _load_backend():
		165	+ from .lowering_patch import restore_inductor_baseline
		166	+
		167	+ # Reset Inductor globals before each backend switch (mlir <-> triton in one process).
		168	+ restore_inductor_baseline()
158		169
159	backend = _get_backend()	170	backend = _get_backend()
160	loader = _BACKEND_LOADERS.get(backend, _load_triton_backend)	171	loader = _BACKEND_LOADERS.get(backend, _load_triton_backend)

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/config.py+1-0

	@@ -318,4 +318,5 @@ decomps_to_exclude_npu = [
318	aten.reflection_pad2d,	318	aten.reflection_pad2d,
319	aten.grid_sampler_2d,	319	aten.grid_sampler_2d,
320	aten.grid_sampler_2d_backward,	320	aten.grid_sampler_2d_backward,
		321	+ aten.expm1,
321	]	322	]

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/meta_kernel.py+9-3

	@@ -38,7 +38,11 @@ from torch._inductor.scheduler import Scheduler
38		38
39	from torch.fx.experimental.proxy_tensor import make_fx	39	from torch.fx.experimental.proxy_tensor import make_fx
40	from torch._dynamo.device_interface import get_interface_for_device	40	from torch._dynamo.device_interface import get_interface_for_device
41	-from ...npu.inductor_patch.lowering import map_strings_to_operators	41	+from torch_npu._inductor.lowering_common import (
		42	+ MLIR_OPERATOR_MAPPING,
		43	+ map_strings_to_operators as _map_strings_to_operators,
		44	+ merge_fx_graphs,
		45	+)
42	from ...npu.utils import (	46	from ...npu.utils import (
43	MLIRProcessor,	47	MLIRProcessor,
44	parse_fx_example_inputs,	48	parse_fx_example_inputs,
	@@ -52,12 +56,14 @@ from ...npu.utils import (
52	view_to_reshape	56	view_to_reshape
53	)	57	)
54	from ... import config as anir_config	58	from ... import config as anir_config
55	-from ...npu.inductor_patch.lowering import merge_fx_graphs
56	-
57		59
58	id_iter = count()	60	id_iter = count()
59		61
60		62
		63	+def map_strings_to_operators(expr_str: str):
		64	+ return _map_strings_to_operators(expr_str, MLIR_OPERATOR_MAPPING)
		65	+
		66	+
61	class NpuTritonKernel(TritonKernel):	67	class NpuTritonKernel(TritonKernel):
62	def __init__(self,	68	def __init__(self,
63	tiling: Dict[str, sympy.Expr],	69	tiling: Dict[str, sympy.Expr],

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/mlir.py+8-1

	@@ -25,7 +25,10 @@ from torch._inductor.utils import (
25	)	25	)
26	from torch._inductor import config, ir, scheduler	26	from torch._inductor import config, ir, scheduler
27	from ... import config as anir_config	27	from ... import config as anir_config
28	-from ...npu.inductor_patch.lowering import map_strings_to_operators	28	+from torch_npu._inductor.lowering_common import (
		29	+ MLIR_OPERATOR_MAPPING,
		30	+ map_strings_to_operators as _map_strings_to_operators,
		31	+)
29	from ...npu.utils import (	32	from ...npu.utils import (
30	MLIRProcessor,	33	MLIRProcessor,
31	fold_expand,	34	fold_expand,
	@@ -42,6 +45,10 @@ from ...npu.utils import (
42	from ...npu.codegen.meta_kernel import NpuMetaKernel, NpuMetaScheduling	45	from ...npu.codegen.meta_kernel import NpuMetaKernel, NpuMetaScheduling
43		46
44		47
		48	+def map_strings_to_operators(expr_str: str):
		49	+ return _map_strings_to_operators(expr_str, MLIR_OPERATOR_MAPPING)
		50	+
		51	+
45	class NpuMlirKernel(NpuMetaKernel):	52	class NpuMlirKernel(NpuMetaKernel):
46	def build_gm_with_prim_cast(self, gm):	53	def build_gm_with_prim_cast(self, gm):
47	return npu_cast_to_prim_cast(gm)	54	return npu_cast_to_prim_cast(gm)

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/inductor_patch/__init__.py+0-39

	@@ -1,43 +1,4 @@
1	import os	1	import os
2	-import sys
3	-import importlib
4	-import inspect
5	import pkgutil	2	import pkgutil
6		3
7	__all__ = list(module for _, module, _ in pkgutil.iter_modules([os.path.dirname(__file__)]))	4	__all__ = list(module for _, module, _ in pkgutil.iter_modules([os.path.dirname(__file__)]))
8	-
9	-from . import ir
10	-from . import lowering as npu_lowering
11	-from torch._inductor import lowering
12	-
13	-
14	-def get_functions_from_module(module):
15	- functions = {}
16	- members = inspect.getmembers(module, inspect.isfunction)
17	-
18	- for name, func in members:
19	- if inspect.getmodule(func) == module:
20	- functions[name] = func
21	-
22	- return functions
23	-
24	-
25	-npu_functions = get_functions_from_module(npu_lowering)
26	-functions = get_functions_from_module(lowering)
27	-for name, _ in functions.items():
28	- if name in npu_functions:
29	- setattr(lowering, name, npu_functions[name])
30	-
31	-extra_lowerings = set(lowering.lowerings.keys()) - set(npu_lowering.lowerings.keys())
32	-npu_lowering.lowerings.update({k: lowering.lowerings[k] for k in extra_lowerings})
33	-lowering.lowerings = npu_lowering.lowerings
34	-lowering._maybe_layout_constraints = npu_lowering._maybe_layout_constraints
35	-lowering.fallbacks = npu_lowering.fallbacks
36	-lowering.needs_realized_inputs = npu_lowering.needs_realized_inputs
37	-lowering.foreach_ops = npu_lowering.foreach_ops
38	-lowering.inplace_foreach_ops = npu_lowering.inplace_foreach_ops
39	-lowering.inplaceable_foreach_ops = npu_lowering.inplaceable_foreach_ops
40	-
41	-from torch._inductor import graph
42	-
43	-importlib.reload(graph)

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/inductor_patch/lowering.py+30-236

	@@ -133,262 +133,56 @@ inplace_foreach_ops = OrderedSet[torch._ops.OpOverload]()
133	inplaceable_foreach_ops: dict[torch._ops.OpOverload, torch._ops.OpOverload] = {}	133	inplaceable_foreach_ops: dict[torch._ops.OpOverload, torch._ops.OpOverload] = {}
134	quantized_decomposed = torch.ops.quantized_decomposed	134	quantized_decomposed = torch.ops.quantized_decomposed
135		135
		136	+from torch_npu._inductor.lowering_common import (
		137	+ TracedGraph,
		138	+ MLIR_OPERATOR_MAPPING,
		139	+ create_fake_input,
		140	+ create_sym_inputs as _create_sym_inputs,
		141	+ fetch_graphs as _fetch_graphs,
		142	+ get_reduction_type_to_aten_fn,
		143	+ map_operators_to_strings as _map_operators_to_strings,
		144	+ map_strings_to_operators as _map_strings_to_operators,
		145	+ merge_fx_graphs,
		146	+ merge_traced_graphs as _merge_traced_graphs,
		147	+ process_ir_constant as _process_ir_constant,
		148	+ register_fn_to_aten_fn as _register_fn_to_aten_fn,
		149	+ register_to_aten as _register_to_aten,
		150	+ subtract_graph,
		151	+)
		152	+
136	fn_to_aten_fn = {}	153	fn_to_aten_fn = {}
137	node_id = itertools.count(0)	154	node_id = itertools.count(0)
138		155
139	def register_fn_to_aten_fn(fn, aten_fn=None):	156	def register_fn_to_aten_fn(fn, aten_fn=None):
140	- if fn not in fn_to_aten_fn:	157	+ return _register_fn_to_aten_fn(fn_to_aten_fn, fn, aten_fn)
141	- fn_to_aten_fn[fn] = aten_fn
142	- return fn
143		158
144	def register_to_aten(aten_fn=None):	159	def register_to_aten(aten_fn=None):
145	- def decorator(fn):	160	+ return _register_to_aten(fn_to_aten_fn, aten_fn)
146	- if fn not in fn_to_aten_fn:
147	- fn_to_aten_fn[fn] = aten_fn
148	- return fn
149	- return decorator
150		161
151	-reduction_type_to_aten_fn = {	162	+reduction_type_to_aten_fn = get_reduction_type_to_aten_fn()
152	- "sum": aten.sum,
153	- "prod": aten.prod,
154	- "xor_sum": prims.xor_sum,
155	- "any": aten.any,
156	- "max": aten.amax,
157	- "min": aten.amin,
158	- "argmax": aten.argmax,
159	- "argmin": aten.argmin
160	-}
161		163
162	-operator_to_string = {	164	+operator_to_string = MLIR_OPERATOR_MAPPING.operator_to_string
163	- '+': 'xvxa',	165	+string_to_operator = MLIR_OPERATOR_MAPPING.string_to_operator
164	- '-': 'xvxb',
165	- '*': 'xvxc',
166	- '/': 'xvxd',
167	- '(': 'xvxe',
168	- ')': 'xvxf',
169	- '.': 'xvxg',
170	- ',': 'xvxh',
171	-}
172	-
173	-string_to_operator = {v: k for k, v in operator_to_string.items()}
174		166
175	def map_operators_to_strings(expr_str: str):	167	def map_operators_to_strings(expr_str: str):
176	- expr_str = expr_str.replace(' ', '')	168	+ return _map_operators_to_strings(expr_str, MLIR_OPERATOR_MAPPING)
177	- for op, string in operator_to_string.items():
178	- expr_str = expr_str.replace(op, string)
179	- return '_uwu_' + expr_str
180		169
181	def map_strings_to_operators(expr_str: str):	170	def map_strings_to_operators(expr_str: str):
182	- for op, string in string_to_operator.items():	171	+ return _map_strings_to_operators(expr_str, MLIR_OPERATOR_MAPPING)
183	- expr_str = expr_str.replace(op, string)
184	- return expr_str[5:]
185	-
186	-
187	-class TracedGraph:
188	- def __init__(self):
189	- self.graph = torch.fx.Graph()
190	- self.last_node: Optional[torch.fx.Node] = None
191	- self.sym_nodes: Dict[str, torch.fx.Node] = {}
192	-
193	- def __str__(self):
194	- return str(self.graph)
195	-
196	- def get_placeholder_names(self):
197	- placeholder_names = set()
198	- for node in self.graph.nodes:
199	- if node.op == 'placeholder' and node.name not in self.sym_nodes:
200	- placeholder_names.add(node.name)
201	- return placeholder_names
202	-
203	- __repr__ = __str__
204	-
205	-
206	-
207	-def create_fake_input(size, stride, device, dtype):
208	- size = [V.graph.sizevars.shape_env.create_symintnode(s, hint=None) \
209	- if isinstance(s, Expr) and not isinstance(s, Integer) else s for s in size]
210	- stride = [V.graph.sizevars.shape_env.create_symintnode(s, hint=None) \
211	- if isinstance(s, Expr) and not isinstance(s, Integer) else s for s in stride]
212	- with V.graph.fake_mode:
213	- fake_input = torch.empty_strided(size, stride, device=device, dtype=dtype)
214	- return fake_input
215	-
216		172
217	def create_sym_inputs(traced_graph: TracedGraph, size: List[Expr]):	173	def create_sym_inputs(traced_graph: TracedGraph, size: List[Expr]):
218	- for s in size:	174	+ return _create_sym_inputs(traced_graph, size, MLIR_OPERATOR_MAPPING)
219	- if isinstance(s, (List, Tuple)):
220	- create_sym_inputs(traced_graph, s)
221	- continue
222	- if isinstance(s, Expr) and not isinstance(s, Integer):
223	- s_name = str(s)
224	- if not isinstance(s, Symbol):
225	- s_name = map_operators_to_strings(s_name)
226	- if s_name in traced_graph.sym_nodes:
227	- continue
228	- new_node = traced_graph.graph.placeholder(s_name)
229	- new_node.meta['val'] = V.graph.sizevars.shape_env.create_symintnode(s, hint=None)
230	- traced_graph.sym_nodes.update({s_name: new_node})
231	-
232		175
233	def process_ir_constant(inp: ExpandView) -> Union[TracedGraph, int, float]:	176	def process_ir_constant(inp: ExpandView) -> Union[TracedGraph, int, float]:
234	- skip = False	177	+ return _process_ir_constant(inp, MLIR_OPERATOR_MAPPING)
235	- if isinstance(inp.data, IndexingConstant):
236	- dtype = inp.data.dtype
237	- inp = inp.data.index
238	- # convert to original dtype.
239	- if dtype in [torch.float32, torch.float16, torch.bfloat16]:
240	- # sympy inputs
241	- if isinstance(inp, Expr) and not isinstance(inp, Integer):
242	- traced_graph = TracedGraph()
243	- create_sym_inputs(traced_graph, [inp])
244	- s_name = str(inp)
245	- if not isinstance(inp, Symbol):
246	- s_name = map_operators_to_strings(str(inp))
247	- traced_graph.last_node = traced_graph.sym_nodes[s_name]
248	- inp = traced_graph
249	- else:
250	- inp = float(inp)
251	- elif isinstance(inp.data, ir.Constant):
252	- dtype = inp.data.dtype
253	- inp = inp.data.value
254	- else:
255	- skip = True
256	- return inp, skip
257	-
258		178
259	def fetch_graphs(inputs: Optional[List[TensorBox]]):	179	def fetch_graphs(inputs: Optional[List[TensorBox]]):
260	- if isinstance(inputs, (TensorBox, ir.StorageBox, ir.View, sympy.Symbol, ir.Constant, ir.ReinterpretView)):	180	+ return _fetch_graphs(inputs, MLIR_OPERATOR_MAPPING, use_npu_meta=False)
261	- inputs = [inputs]
262	- input_graphs = []
263	- for inp in inputs:
264	- if isinstance(inp, List):
265	- input_graphs.append(fetch_graphs(inp))
266	- continue
267	- if not isinstance(inp, (TensorBox, ir.StorageBox, ir.View, ir.ReinterpretView, ir.PermuteView, ir.SliceView, ir.ExpandView)):
268	- input_graphs.append(inp)
269	- continue
270	- if isinstance(inp, ExpandView):
271	- inp, skip = process_ir_constant(inp)
272	- if not skip:
273	- input_graphs.append(inp)
274	- continue
275	- name = inp.get_name()
276	- traced_graph = inp.get_traced_graph()
277	- if traced_graph is not None:
278	- input_graphs.append(traced_graph)
279	- continue
280	- traced_graph = TracedGraph()
281	- device = inp.get_device()
282	- dtype = inp.get_dtype()
283	- size = inp.get_size()
284	- stride = inp.get_stride()
285	- new_node = traced_graph.graph.placeholder(name)
286	- fake_input = create_fake_input(size, stride, device, dtype)
287	- new_node.meta['val'] = fake_input
288	- traced_graph.last_node = new_node
289	- input_graphs.append(traced_graph)
290	- return input_graphs
291	-
292		181
293	def merge_traced_graphs(input_graphs: List[TracedGraph], origin_fn, node_name, **kwargs):	182	def merge_traced_graphs(input_graphs: List[TracedGraph], origin_fn, node_name, **kwargs):
294	- new_graph = TracedGraph()	183	+ return _merge_traced_graphs(
295	- exist_nodes: Dict[str, torch.fx.Node] = {}	184	+ input_graphs, origin_fn, node_name, MLIR_OPERATOR_MAPPING, **kwargs
296	- def merge_graph(input_graphs: List[TracedGraph]):	185	+ )
297	- for input_graph in input_graphs:
298	- if isinstance(input_graph, List):
299	- merge_graph(input_graph)
300	- continue
301	- if not isinstance(input_graph, TracedGraph):
302	- continue
303	- for node in input_graph.graph.nodes:
304	- if node.name in exist_nodes:
305	- continue
306	- # [wtd#21] Use dict.get to avoid KeyError when n.name is not in exist_nodes
307	- new_node = new_graph.graph.node_copy(node, lambda n: exist_nodes.get(n.name, n))
308	- exist_nodes[node.name] = new_node
309	- if node.name in input_graph.sym_nodes:
310	- new_graph.sym_nodes.update({node.name: new_node})
311	-
312	- def parse_args(input_graphs, exist_nodes):
313	- args = []
314	- for input_graph in input_graphs:
315	- if isinstance(input_graph, TracedGraph):
316	- args.append(exist_nodes[input_graph.last_node.name])
317	- elif isinstance(input_graph, (List, Tuple)):
318	- args.append(parse_args(input_graph, exist_nodes))
319	- else:
320	- if isinstance(input_graph, Expr) and not isinstance(input_graph, Integer):
321	- if not isinstance(input_graph, Symbol):
322	- input_graph = map_operators_to_strings(str(input_graph))
323	- args.append(new_graph.sym_nodes[str(input_graph)])
324	- else:
325	- args.append(input_graph)
326	- return args
327	-
328	- num_args = len(input_graphs)
329	-
330	- for k, v in kwargs.items():
331	- if isinstance(v, Expr) and not isinstance(v, Integer):
332	- traced_graph = TracedGraph()
333	- create_sym_inputs(traced_graph, [v])
334	- s_name = str(v)
335	- if not isinstance(v, Symbol):
336	- s_name = map_operators_to_strings(str(v))
337	- traced_graph.last_node = traced_graph.sym_nodes[s_name]
338	- kwargs[k] = traced_graph.sym_nodes[s_name]
339	- input_graphs.append(traced_graph)
340	- merge_graph(input_graphs)
341	- input_graphs = input_graphs[:num_args]
342	- # if inputs do not have any valid graphs, like full/iota
343	- create_sym_inputs(new_graph, input_graphs)
344	- args = parse_args(input_graphs, exist_nodes)
345	- with new_graph.graph.inserting_after(new_graph.last_node):
346	- new_node = new_graph.graph.call_function(origin_fn, args=tuple(args), kwargs=kwargs)
347	- new_node.name = node_name
348	- new_graph.last_node = new_node
349	- return new_graph
350	-
351	-def merge_fx_graphs(traced_graphs: List[TracedGraph]):
352	- new_graph = TracedGraph()
353	- exist_nodes: Dict[str, torch.fx.Node] = {}
354	- last_nodes = []
355	- def merge_graph(input_graphs: List[TracedGraph]):
356	- for input_graph in input_graphs:
357	- if isinstance(input_graph, List):
358	- merge_graph(input_graph)
359	- continue
360	- if not isinstance(input_graph, TracedGraph):
361	- continue
362	- for node in input_graph.graph.nodes:
363	- if node.name in exist_nodes:
364	- continue
365	- # [wtd#21] Use dict.get to avoid KeyError when n.name is not in exist_nodes
366	- new_node = new_graph.graph.node_copy(node, lambda n: exist_nodes.get(n.name, n))
367	- exist_nodes[node.name] = new_node
368	- last_nodes.append(exist_nodes[input_graph.last_node.name])
369	- merge_graph(traced_graphs)
370	- new_graph.last_node = last_nodes
371	- return new_graph
372	-
373	-def subtract_graph(graph1: TracedGraph, graph2: TracedGraph, node_name=None) -> Tuple[TracedGraph, torch.fx.Node]:
374	- new_graph = TracedGraph()
375	- last_node2 = graph2.last_node
376	- graph1_node_names = {node.name for node in graph1.graph.nodes}
377	- graph2_node_names = {node.name for node in graph2.graph.nodes}
378	- placeholder = None
379	- exist_nodes: Dict[str, torch.fx.Node] = {}
380	- if node_name not in graph1_node_names:
381	- placeholder = new_graph.graph.placeholder(last_node2.name if node_name is None else node_name)
382	- exist_nodes[last_node2.name] = placeholder
383	- for node in graph1.graph.nodes:
384	- if node.name in graph2_node_names and node.name not in graph1.sym_nodes:
385	- continue
386	- # [wtd#21] Use dict.get to avoid KeyError when n.name is not in exist_nodes
387	- new_node = new_graph.graph.node_copy(node, lambda n: exist_nodes.get(n.name, n))
388	- exist_nodes[node.name] = new_node
389	- new_graph.last_node = exist_nodes[graph1.last_node.name]
390	- new_graph.sym_nodes = graph1.sym_nodes
391	- return new_graph, placeholder
392		186
393		187
394	def cur_node_has_non_foreach_users():	188	def cur_node_has_non_foreach_users():

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/inductor_patch/scheduler.py+2-2

	@@ -3,7 +3,8 @@ import sympy
3	import collections	3	import collections
4	from typing import (	4	from typing import (
5	Union,	5	Union,
6	- Optional	6	+ Optional,
		7	+ Sequence,
7	)	8	)
8		9
9	from torch._inductor import ir	10	from torch._inductor import ir
	@@ -18,7 +19,6 @@ from torch._inductor.scheduler import (
18	MultiOutput,	19	MultiOutput,
19	MultiOutputLayout,	20	MultiOutputLayout,
20	OrderedSet,	21	OrderedSet,
21	- Sequence,
22	get_dtype_size,	22	get_dtype_size,
23	sympy_product,	23	sympy_product,
24	V,	24	V,

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/npu_inductor_plugin.py+14-15

	@@ -68,23 +68,27 @@ if anir_config.online_acc_comp:
68	aten = torch.ops.aten	68	aten = torch.ops.aten
69		69
70	## Override original dynamo device interface in torch_npu	70	## Override original dynamo device interface in torch_npu
71	-from torch_npu.utils._dynamo_device import NpuInterface	71	+
72	-if os.getenv('TORCHINDUCTOR_USE_AKG', '0') == '1':	72	+
73	- try:	73	+def register_mlir_codegen_backend() -> None:
74	- import akg	74	+ """Register MLIR scheduling/wrapper; call on each mlir/dvm backend switch."""
75	- import torch_mlir	75	+ if os.getenv('TORCHINDUCTOR_USE_AKG', '0') == '1':
76	- register_backend_for_device("npu", AkgScheduling, NpuMlirWrapperCodeGen)	76	+ try:
77	- except:	77	+ import akg
78	- logger.warning(f"akg not found, fallback to torch-mlir for compilation.")	78	+ import torch_mlir
		79	+ register_backend_for_device("npu", AkgScheduling, NpuMlirWrapperCodeGen)
		80	+ except ImportError:
		81	+ logger.warning("akg not found, fallback to torch-mlir for compilation.")
		82	+ register_backend_for_device("npu", NpuMlirScheduling, NpuMlirWrapperCodeGen)
		83	+ else:
79	register_backend_for_device("npu", NpuMlirScheduling, NpuMlirWrapperCodeGen)	84	register_backend_for_device("npu", NpuMlirScheduling, NpuMlirWrapperCodeGen)
80	-else:
81	- register_backend_for_device("npu", NpuMlirScheduling, NpuMlirWrapperCodeGen)
82		85
83	try:	86	try:
84	from torch_npu.npu import device_count	87	from torch_npu.npu import device_count
85	except:	88	except:
86	from torch_npu.npu.utils import device_count	89	from torch_npu.npu.utils import device_count
87	from torch._dynamo.device_interface import register_interface_for_device	90	from torch._dynamo.device_interface import register_interface_for_device
		91	+from torch_npu.utils._dynamo_device import NpuInterface
88		92
89	class NewNpuInterface(NpuInterface):	93	class NewNpuInterface(NpuInterface):
90		94
	@@ -171,11 +175,6 @@ def _patch_run_node(tracer, node, args, kwargs, nnmodule):
171	raise AssertionError(op)	175	raise AssertionError(op)
172		176
173		177
174	-def _register_npu_inductor_fallbacks_operation():
175	- from ..npu import inductor_patch
176	-
177	-
178	-_register_npu_inductor_fallbacks_operation()
179	disable_implicit_decomposition()	178	disable_implicit_decomposition()
180	torch._dynamo.utils.run_node = _patch_run_node	179	torch._dynamo.utils.run_node = _patch_run_node
181		180

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/npu_lowering.py+14-68

	@@ -1,11 +1,15 @@
1	-from functools import reduce
2	-
3	import torch._ops	1	import torch._ops
4	from torch._inductor import decomposition, lowering	2	from torch._inductor import decomposition, lowering
5	from torch._inductor.fx_passes.control_dependencies import ControlDeps	3	from torch._inductor.fx_passes.control_dependencies import ControlDeps
6	from torch._inductor.lowering import lowerings, make_fallback	4	from torch._inductor.lowering import lowerings, make_fallback
7	from torch.utils._ordered_set import OrderedSet	5	from torch.utils._ordered_set import OrderedSet
8		6
		7	+from torch_npu._inductor.lowering_common import (
		8	+ add_overload,
		9	+ fallback_ops_with_meta,
		10	+ resolve_op_from_name,
		11	+)
		12	+
9	from .. import config	13	from .. import config
10	from ..npu.utils import get_anir_mode, run_once	14	from ..npu.utils import get_anir_mode, run_once
11	from .utils import logger	15	from .utils import logger
	@@ -23,21 +27,10 @@ def _register_npu_inductor_fallbacks():
23	fallback_set_exclude = OrderedSet()	27	fallback_set_exclude = OrderedSet()
24	env_fallback_list = config.enable_full_lowering_fallback	28	env_fallback_list = config.enable_full_lowering_fallback
25		29
26	- def _resolve_op_from_name(op_name: str):
27	- try:
28	- obj = torch.ops
29	- for part in op_name.split("."):
30	- obj = getattr(obj, part)
31	- return obj
32	- except AttributeError:
33	- logger.warning(
34	- "[npu\|inductor\|lowering\|fallback] invalid identifier name: %s", op_name
35	- )
36	-
37	if env_fallback_list:	30	if env_fallback_list:
38	for op_name in env_fallback_list.split(","):	31	for op_name in env_fallback_list.split(","):
39	op_name = op_name.strip()	32	op_name = op_name.strip()
40	- op = _resolve_op_from_name(op_name)	33	+ op = resolve_op_from_name(op_name, logger)
41	if isinstance(op, torch._ops.OpOverloadPacket):	34	if isinstance(op, torch._ops.OpOverloadPacket):
42	fallback_set.add(op)	35	fallback_set.add(op)
43	fallback_set_exclude.add(op)	36	fallback_set_exclude.add(op)
	@@ -51,19 +44,8 @@ def _register_npu_inductor_fallbacks():
51	op_name,	44	op_name,
52	)	45	)
53		46
54	- for fn in config.GENERATE_LIST:	47	+ add_overload(config.GENERATE_LIST, gen_set)
55	- gen_set.add(fn)	48	+ add_overload(config.FALLBACK_LIST, fallback_set)
56	- if isinstance(fn, torch._ops.OpOverloadPacket):
57	- for overload in fn.overloads():
58	- other_fn = getattr(fn, overload)
59	- gen_set.add(other_fn)
60	-
61	- for fn in config.FALLBACK_LIST:
62	- fallback_set.add(fn)
63	- if isinstance(fn, torch._ops.OpOverloadPacket):
64	- for overload in fn.overloads():
65	- other_fn = getattr(fn, overload)
66	- fallback_set.add(other_fn)
67		49
68	def fallback_except_gen_set(gen_set):	50	def fallback_except_gen_set(gen_set):
69	for op in lowering.lowerings:	51	for op in lowering.lowerings:
	@@ -110,7 +92,11 @@ def _register_npu_inductor_fallbacks():
110	for op in ops_to_fallback:	92	for op in ops_to_fallback:
111	make_fallback(op)	93	make_fallback(op)
112		94
113	- _fallback_ops_with_meta()	95	+ fallback_ops_with_meta(
		96	+ lowerings,
		97	+ decomposition.decompositions,
		98	+ make_fallback,
		99	+ )
114		100
115	if config.fallback_to_aten_mode not in {"off", "include", "exclude", "all"}:	101	if config.fallback_to_aten_mode not in {"off", "include", "exclude", "all"}:
116	raise AssertionError(	102	raise AssertionError(
	@@ -129,43 +115,3 @@ def _register_npu_inductor_fallbacks():
129	fallback_via_fallback_set(fallback_set=fallback_set_exclude)	115	fallback_via_fallback_set(fallback_set=fallback_set_exclude)
130	elif config.fallback_to_aten_mode == "all":	116	elif config.fallback_to_aten_mode == "all":
131	enable_full_lowering_fallback()	117	enable_full_lowering_fallback()
132	-
133	-
134	-def get_nested_attr(obj, attr_path, default=None):
135	- try:
136	- return reduce(getattr, attr_path.split("."), obj)
137	- except AttributeError:
138	- return default
139	-
140	-
141	-def _fallback_ops_with_meta():
142	- """
143	- Fallback all ops that have a Meta implementation but are not yet in lowerings
144	- """
145	- all_ops = torch._C._dispatch_get_all_op_names()
146	-
147	- for op_name in all_ops:
148	- has_meta = torch._C._dispatch_has_kernel_for_dispatch_key(op_name, "Meta")
149	- has_comp = torch._C._dispatch_has_kernel_for_dispatch_key(
150	- op_name, "CompositeImplicitAutograd"
151	- )
152	-
153	- if not (has_meta or has_comp):
154	- continue
155	-
156	- namespace, name_with_overload = op_name.split("::", 1)
157	-
158	- if "." in name_with_overload:
159	- name, overload = name_with_overload.rsplit(".", 1)
160	- else:
161	- name, overload = name_with_overload, "default"
162	-
163	- normalized_path = f"{namespace}.{name}.{overload}"
164	- op_overload = get_nested_attr(torch.ops, normalized_path)
165	- if not isinstance(op_overload, torch._ops.OpOverload):
166	- continue
167	-
168	- if op_overload in lowerings or op_overload in decomposition.decompositions:
169	- continue
170	-
171	- make_fallback(op_overload)

Mtorch_npu/_inductor/decomposition.py+4-2

	@@ -3,7 +3,7 @@ from torch._inductor.decomposition import decompositions, pw_cast_for_opmath
3	from torch._inductor.decomposition import register_decomposition	3	from torch._inductor.decomposition import register_decomposition
4	from torch._prims_common.wrappers import out_wrapper	4	from torch._prims_common.wrappers import out_wrapper
5		5
6	-from .lowering import _init_set	6	+from .lowering_common import add_overload
7		7
8	aten = torch.ops.aten	8	aten = torch.ops.aten
9		9
	@@ -17,12 +17,14 @@ DECOMPOSITION_OVERLOAD_OP = [
17	aten.embedding_dense_backward,	17	aten.embedding_dense_backward,
18	aten.addmm,	18	aten.addmm,
19	aten.gelu,	19	aten.gelu,
		20	+ aten.expm1,
		21	+ aten.erfc
20	]	22	]
21		23
22		24
23	def _register_npu_inductor_decompositons():	25	def _register_npu_inductor_decompositons():
24	overload_op_set = set()	26	overload_op_set = set()
25	- _init_set(DECOMPOSITION_OVERLOAD_OP, overload_op_set)	27	+ add_overload(DECOMPOSITION_OVERLOAD_OP, overload_op_set)
26		28
27	for op in overload_op_set:	29	for op in overload_op_set:
28	if (op in decompositions):	30	if (op in decompositions):

Mtorch_npu/_inductor/lowering.py+18-74

	@@ -1,6 +1,5 @@
1	import os	1	import os
2	import sympy	2	import sympy
3	-from functools import reduce
4		3
5	import torch._ops	4	import torch._ops
6	from torch._inductor import ir	5	from torch._inductor import ir
	@@ -39,6 +38,11 @@ from torch._inductor.lowering import (
39	get_promoted_dtype,	38	get_promoted_dtype,
40	)	39	)
41	from .. import npu_dtype_cast, _npu_dtype_cast	40	from .. import npu_dtype_cast, _npu_dtype_cast
		41	+from .lowering_common import (
		42	+ add_overload,
		43	+ enable_full_lowering_fallback as enable_full_lowering_fallback_common,
		44	+ resolve_op_from_name,
		45	+)
42	from .config import log, enable_full_lowering_fallback	46	from .config import log, enable_full_lowering_fallback
43	from .lowering_op_list import GENERATE_LIST, GENERATE_LIST2, FALLBACK_LIST, LOWERING_OVERLOAD_OP	47	from .lowering_op_list import GENERATE_LIST, GENERATE_LIST2, FALLBACK_LIST, LOWERING_OVERLOAD_OP
44	from . import config as npu_config	48	from . import config as npu_config
	@@ -142,37 +146,17 @@ tr_c10d = torch.ops.tr_c10d
142	prims = torch.ops.prims	146	prims = torch.ops.prims
143		147
144		148
145	-def _init_set(input_list, output_set):
146	- for fn in input_list:
147	- output_set.add(fn)
148	- if isinstance(fn, torch._ops.OpOverloadPacket):
149	- for overload in fn.overloads():
150	- other_fn = getattr(fn, overload)
151	- output_set.add(other_fn)
152	-
153	-
154	-def _resolve_op_from_name(op_name: str):
155	- try:
156	- obj = torch.ops
157	- for part in op_name.split('.'):
158	- obj = getattr(obj, part)
159	- return obj
160	- except AttributeError:
161	- log.warning(f"[npu\|inductor\|lowering\|fallback] invalid identifier name: {op_name}")
162	- return None
163	-
164	-
165	def _register_npu_inductor_fallbacks():	149	def _register_npu_inductor_fallbacks():
166	gen_set = set()	150	gen_set = set()
167	- _init_set(GENERATE_LIST, gen_set)	151	+ add_overload(GENERATE_LIST, gen_set)
168	overload_op_set = set()	152	overload_op_set = set()
169	- _init_set(LOWERING_OVERLOAD_OP, overload_op_set)	153	+ add_overload(LOWERING_OVERLOAD_OP, overload_op_set)
170		154
171	env_fallback_list = enable_full_lowering_fallback	155	env_fallback_list = enable_full_lowering_fallback
172	if env_fallback_list:	156	if env_fallback_list:
173	for op_name in env_fallback_list.split(','):	157	for op_name in env_fallback_list.split(','):
174	op_name = op_name.strip()	158	op_name = op_name.strip()
175	- op = _resolve_op_from_name(op_name)	159	+ op = resolve_op_from_name(op_name, log)
176	if isinstance(op, (torch._ops.OpOverloadPacket, torch._ops.OpOverload, torch._ops.HigherOrderOperator)):	160	if isinstance(op, (torch._ops.OpOverloadPacket, torch._ops.OpOverload, torch._ops.HigherOrderOperator)):
177	FALLBACK_LIST.append(op)	161	FALLBACK_LIST.append(op)
178	log.info(f"[npu\|inductor\|lowering\|fallback] User specified fallback: {op_name}")	162	log.info(f"[npu\|inductor\|lowering\|fallback] User specified fallback: {op_name}")
	@@ -351,54 +335,14 @@ def _register_npu_inductor_fallbacks():
351	make_fallback(aten.nll_loss_forward)	335	make_fallback(aten.nll_loss_forward)
352		336
353		337
354	-def get_nested_attr(obj, attr_path, default=None):
355	- try:
356	- return reduce(getattr, attr_path.split('.'), obj)
357	- except AttributeError:
358	- return default
359	-
360	-
361	-def _fallback_ops_with_meta():
362	- """
363	- Fallback all ops that have a Meta implementation but are not yet in lowerings
364	- """
365	- all_ops = torch._C._dispatch_get_all_op_names()
366	-
367	- for op_name in all_ops:
368	- has_meta = torch._C._dispatch_has_kernel_for_dispatch_key(op_name, "Meta")
369	- has_comp = torch._C._dispatch_has_kernel_for_dispatch_key(op_name, "CompositeImplicitAutograd")
370	-
371	- if not (has_meta or has_comp):
372	- continue
373	-
374	- namespace, name_with_overload = op_name.split("::", 1)
375	-
376	- if "." in name_with_overload:
377	- name, overload = name_with_overload.rsplit(".", 1)
378	- else:
379	- name, overload = name_with_overload, "default"
380	-
381	- normalized_path = f"{namespace}.{name}.{overload}"
382	- op_overload = get_nested_attr(torch.ops, normalized_path)
383	- if not isinstance(op_overload, torch._ops.OpOverload):
384	- continue
385	-
386	- if op_overload in lowerings or op_overload in decompositions:
387	- continue
388	-
389	- make_fallback(op_overload)
390	- FALLBACK_LIST.append(op_overload)
391	-
392	-
393	def _enable_full_lowering_fallback():	338	def _enable_full_lowering_fallback():
394	- ops_to_fallback = list(filter(	339	+ enable_full_lowering_fallback_common(
395	- lambda op: op not in decompositions and	340	+ lowerings,
396	- isinstance(op, (torch._ops.OpOverloadPacket, torch._ops.OpOverload, torch._ops.HigherOrderOperator)) and	341	+ decompositions,
397	- op not in (torch._higher_order_ops.triton_kernel_wrap.TritonKernelWrapperMutation,	342	+ make_fallback,
398	- torch._higher_order_ops.triton_kernel_wrap.TritonKernelWrapperFunctional),	343	+ FALLBACK_LIST,
399	- lowerings	344	+ excluded_ops=(
400	- ))	345	+ torch._higher_order_ops.triton_kernel_wrap.TritonKernelWrapperMutation,
401	- for op in ops_to_fallback:	346	+ torch._higher_order_ops.triton_kernel_wrap.TritonKernelWrapperFunctional,
402	- make_fallback(op)	347	+ ),
403	- FALLBACK_LIST.append(op)	348	+ )
404	- _fallback_ops_with_meta()

Atorch_npu/_inductor/lowering_common.py+462-0

@@ -0,0 +1,462 @@
	1	+# Copyright (c) 2026, Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache-2.0 License (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# https://github.com/pytorch/pytorch/blob/main/LICENSE
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+
	15	+from __future__ import annotations
	16	+
	17	+import inspect
	18	+import sympy
	19	+from functools import reduce
	20	+from dataclasses import dataclass
	21	+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
	22	+
	23	+import torch
	24	+import torch._ops
	25	+import torch.fx
	26	+from sympy.core import Expr, Integer, Symbol
	27	+from sympy.core.numbers import Number as SympyNumber
	28	+from torch._inductor import ir
	29	+from torch._inductor.ir import ExpandView, IndexingConstant, TensorBox
	30	+from torch._inductor.virtualized import V
	31	+
	32	+LOWERING_REGISTRY_ATTRS: tuple[str, ...] = (
	33	+ "lowerings",
	34	+ "_maybe_layout_constraints",
	35	+ "fallbacks",
	36	+ "needs_realized_inputs",
	37	+ "foreach_ops",
	38	+ "inplace_foreach_ops",
	39	+ "inplaceable_foreach_ops",
	40	+)
	41	+
	42	+
	43	+def get_module_functions(module: Any) -> dict[str, Callable[..., Any]]:
	44	+ functions: dict[str, Callable[..., Any]] = {}
	45	+ for name, func in inspect.getmembers(module, inspect.isfunction):
	46	+ if inspect.getmodule(func) is module:
	47	+ functions[name] = func
	48	+ return functions
	49	+
	50	+aten = torch.ops.aten
	51	+prims = torch.ops.prims
	52	+
	53	+
	54	+def add_overload(input_list, output_set):
	55	+ for fn in input_list:
	56	+ output_set.add(fn)
	57	+ if isinstance(fn, torch._ops.OpOverloadPacket):
	58	+ for overload in fn.overloads():
	59	+ other_fn = getattr(fn, overload)
	60	+ output_set.add(other_fn)
	61	+
	62	+
	63	+def resolve_op_from_name(op_name: str, logger=None):
	64	+ try:
	65	+ obj = torch.ops
	66	+ for part in op_name.split('.'):
	67	+ obj = getattr(obj, part)
	68	+ return obj
	69	+ except AttributeError:
	70	+ if logger is not None:
	71	+ logger.warning(f"[npu\|inductor\|lowering\|fallback] invalid identifier name: {op_name}")
	72	+ return None
	73	+
	74	+
	75	+def get_nested_attr(obj, attr_path, default=None):
	76	+ try:
	77	+ return reduce(getattr, attr_path.split('.'), obj)
	78	+ except AttributeError:
	79	+ return default
	80	+
	81	+
	82	+def fallback_ops_with_meta(lowerings, decompositions, make_fallback, fallback_list=None):
	83	+ """
	84	+ Fallback all ops that have a Meta implementation but are not yet in lowerings.
	85	+ """
	86	+ all_ops = torch._C._dispatch_get_all_op_names()
	87	+
	88	+ for op_name in all_ops:
	89	+ has_meta = torch._C._dispatch_has_kernel_for_dispatch_key(op_name, "Meta")
	90	+ has_comp = torch._C._dispatch_has_kernel_for_dispatch_key(op_name, "CompositeImplicitAutograd")
	91	+
	92	+ if not (has_meta or has_comp):
	93	+ continue
	94	+
	95	+ namespace, name_with_overload = op_name.split("::", 1)
	96	+
	97	+ if "." in name_with_overload:
	98	+ name, overload = name_with_overload.rsplit(".", 1)
	99	+ else:
	100	+ name, overload = name_with_overload, "default"
	101	+
	102	+ normalized_path = f"{namespace}.{name}.{overload}"
	103	+ op_overload = get_nested_attr(torch.ops, normalized_path)
	104	+ if not isinstance(op_overload, torch._ops.OpOverload):
	105	+ continue
	106	+
	107	+ if op_overload in lowerings or op_overload in decompositions:
	108	+ continue
	109	+
	110	+ make_fallback(op_overload)
	111	+ if fallback_list is not None:
	112	+ fallback_list.append(op_overload)
	113	+
	114	+
	115	+def enable_full_lowering_fallback(
	116	+ lowerings,
	117	+ decompositions,
	118	+ make_fallback,
	119	+ fallback_list=None,
	120	+ excluded_ops=(),
	121	+):
	122	+ ops_to_fallback = list(filter(
	123	+ lambda op: op not in decompositions and
	124	+ isinstance(op, (torch._ops.OpOverloadPacket, torch._ops.OpOverload, torch._ops.HigherOrderOperator)) and
	125	+ op not in excluded_ops,
	126	+ lowerings
	127	+ ))
	128	+ for op in ops_to_fallback:
	129	+ make_fallback(op)
	130	+ if fallback_list is not None:
	131	+ fallback_list.append(op)
	132	+
	133	+ fallback_ops_with_meta(lowerings, decompositions, make_fallback, fallback_list)
	134	+
	135	+
	136	+class TracedGraph:
	137	+ def __init__(self):
	138	+ self.graph = torch.fx.Graph()
	139	+ self.last_node: Optional[torch.fx.Node] = None
	140	+ self.sym_nodes: Dict[str, torch.fx.Node] = {}
	141	+
	142	+ def __str__(self):
	143	+ return str(self.graph)
	144	+
	145	+ def get_placeholder_names(self):
	146	+ placeholder_names = set()
	147	+ for node in self.graph.nodes:
	148	+ if node.op == "placeholder" and node.name not in self.sym_nodes:
	149	+ placeholder_names.add(node.name)
	150	+ return placeholder_names
	151	+
	152	+ __repr__ = __str__
	153	+
	154	+
	155	+def create_fake_input(size, stride, device, dtype):
	156	+ size = [
	157	+ V.graph.sizevars.shape_env.create_symintnode(s, hint=None)
	158	+ if isinstance(s, Expr) and not isinstance(s, Integer)
	159	+ else s
	160	+ for s in size
	161	+ ]
	162	+ stride = [
	163	+ V.graph.sizevars.shape_env.create_symintnode(s, hint=None)
	164	+ if isinstance(s, Expr) and not isinstance(s, Integer)
	165	+ else s
	166	+ for s in stride
	167	+ ]
	168	+ with V.graph.fake_mode:
	169	+ fake_input = torch.empty_strided(size, stride, device=device, dtype=dtype)
	170	+ return fake_input
	171	+
	172	+
	173	+def get_reduction_type_to_aten_fn():
	174	+ return {
	175	+ "sum": aten.sum,
	176	+ "prod": aten.prod,
	177	+ "xor_sum": prims.xor_sum,
	178	+ "any": aten.any,
	179	+ "max": aten.amax,
	180	+ "min": aten.amin,
	181	+ "argmax": aten.argmax,
	182	+ "argmin": aten.argmin,
	183	+ }
	184	+
	185	+
	186	+def register_fn_to_aten_fn(registry: Dict[Callable, object], fn, aten_fn=None):
	187	+ if fn not in registry:
	188	+ registry[fn] = aten_fn
	189	+ return fn
	190	+
	191	+
	192	+def register_to_aten(registry: Dict[Callable, object], aten_fn=None):
	193	+ def decorator(fn):
	194	+ if fn not in registry:
	195	+ registry[fn] = aten_fn
	196	+ return fn
	197	+
	198	+ return decorator
	199	+
	200	+
	201	+@dataclass(frozen=True)
	202	+class OperatorMapping:
	203	+ """Sympy expr symbol encoding for FX placeholder names."""
	204	+
	205	+ operator_to_string: Dict[str, str]
	206	+ encoded_prefix: str
	207	+ decoded_strip_len: int
	208	+
	209	+ @property
	210	+ def string_to_operator(self) -> Dict[str, str]:
	211	+ return {v: k for k, v in self.operator_to_string.items()}
	212	+
	213	+
	214	+TRITON_OPERATOR_MAPPING = OperatorMapping(
	215	+ operator_to_string={
	216	+ "+": "a",
	217	+ "-": "sub",
	218	+ "*": "m",
	219	+ "/": "d",
	220	+ "(": "l",
	221	+ ")": "r",
	222	+ ".": "p",
	223	+ },
	224	+ encoded_prefix="_",
	225	+ decoded_strip_len=1,
	226	+)
	227	+
	228	+MLIR_OPERATOR_MAPPING = OperatorMapping(
	229	+ operator_to_string={
	230	+ "+": "xvxa",
	231	+ "-": "xvxb",
	232	+ "*": "xvxc",
	233	+ "/": "xvxd",
	234	+ "(": "xvxe",
	235	+ ")": "xvxf",
	236	+ ".": "xvxg",
	237	+ ",": "xvxh",
	238	+ },
	239	+ encoded_prefix="_uwu_",
	240	+ decoded_strip_len=5,
	241	+)
	242	+
	243	+
	244	+def map_operators_to_strings(expr_str: str, mapping: OperatorMapping) -> str:
	245	+ expr_str = expr_str.replace(" ", "")
	246	+ for op, string in mapping.operator_to_string.items():
	247	+ expr_str = expr_str.replace(op, string)
	248	+ return mapping.encoded_prefix + expr_str
	249	+
	250	+
	251	+def map_strings_to_operators(expr_str: str, mapping: OperatorMapping) -> str:
	252	+ for op, string in mapping.string_to_operator.items():
	253	+ expr_str = expr_str.replace(op, string)
	254	+ return expr_str[mapping.decoded_strip_len:]
	255	+
	256	+
	257	+def create_sym_inputs(
	258	+ traced_graph: TracedGraph,
	259	+ size: List[Expr],
	260	+ operator_mapping: OperatorMapping,
	261	+) -> None:
	262	+ for s in size:
	263	+ if isinstance(s, (List, Tuple)):
	264	+ create_sym_inputs(traced_graph, s, operator_mapping)
	265	+ continue
	266	+ if isinstance(s, Expr) and not isinstance(s, Integer):
	267	+ s_name = str(s)
	268	+ if not isinstance(s, Symbol):
	269	+ s_name = map_operators_to_strings(s_name, operator_mapping)
	270	+ if s_name in traced_graph.sym_nodes:
	271	+ continue
	272	+ new_node = traced_graph.graph.placeholder(s_name)
	273	+ new_node.meta["val"] = V.graph.sizevars.shape_env.create_symintnode(s, hint=None)
	274	+ traced_graph.sym_nodes.update({s_name: new_node})
	275	+
	276	+
	277	+def process_ir_constant(
	278	+ inp: ExpandView,
	279	+ operator_mapping: OperatorMapping,
	280	+) -> tuple[Any, bool]:
	281	+ skip = False
	282	+ if isinstance(inp.data, IndexingConstant):
	283	+ dtype = inp.data.dtype
	284	+ inp = inp.data.index
	285	+ if dtype in [torch.float32, torch.float16, torch.bfloat16]:
	286	+ if isinstance(inp, Expr) and not isinstance(inp, SympyNumber):
	287	+ traced_graph = TracedGraph()
	288	+ create_sym_inputs(traced_graph, [inp], operator_mapping)
	289	+ s_name = str(inp)
	290	+ if not isinstance(inp, Symbol):
	291	+ s_name = map_operators_to_strings(str(inp), operator_mapping)
	292	+ traced_graph.last_node = traced_graph.sym_nodes[s_name]
	293	+ inp = traced_graph
	294	+ else:
	295	+ inp = float(inp)
	296	+ elif isinstance(inp.data, ir.Constant):
	297	+ inp = inp.data.value
	298	+ else:
	299	+ skip = True
	300	+ return inp, skip
	301	+
	302	+
	303	+def fetch_graphs(
	304	+ inputs: Optional[List[TensorBox]],
	305	+ operator_mapping: OperatorMapping,
	306	+ *,
	307	+ use_npu_meta: bool = False,
	308	+):
	309	+ if isinstance(inputs, (TensorBox, ir.StorageBox, ir.View, sympy.Symbol, ir.Constant, ir.ReinterpretView)):
	310	+ inputs = [inputs]
	311	+ input_graphs = []
	312	+ for inp in inputs:
	313	+ if isinstance(inp, List):
	314	+ input_graphs.append(fetch_graphs(inp, operator_mapping, use_npu_meta=use_npu_meta))
	315	+ continue
	316	+ if not isinstance(
	317	+ inp,
	318	+ (
	319	+ TensorBox,
	320	+ ir.StorageBox,
	321	+ ir.View,
	322	+ ir.ReinterpretView,
	323	+ ir.PermuteView,
	324	+ ir.SliceView,
	325	+ ir.ExpandView,
	326	+ ),
	327	+ ):
	328	+ input_graphs.append(inp)
	329	+ continue
	330	+ if isinstance(inp, ExpandView):
	331	+ inp, skip = process_ir_constant(inp, operator_mapping)
	332	+ if not skip:
	333	+ input_graphs.append(inp)
	334	+ continue
	335	+ name = inp.get_name()
	336	+ traced_graph = inp.get_traced_graph()
	337	+ if traced_graph is not None:
	338	+ input_graphs.append(traced_graph)
	339	+ continue
	340	+ traced_graph = TracedGraph()
	341	+ device = inp.get_device()
	342	+ dtype = inp.get_dtype()
	343	+ size = inp.get_size()
	344	+ stride = inp.get_stride()
	345	+ new_node = traced_graph.graph.placeholder(name)
	346	+ fake_input = create_fake_input(size, stride, device, dtype)
	347	+ new_node.meta["val"] = fake_input.npu() if use_npu_meta else fake_input
	348	+ traced_graph.last_node = new_node
	349	+ input_graphs.append(traced_graph)
	350	+ return input_graphs
	351	+
	352	+
	353	+def merge_traced_graphs(
	354	+ input_graphs: List[TracedGraph],
	355	+ origin_fn,
	356	+ node_name,
	357	+ operator_mapping: OperatorMapping,
	358	+ **kwargs,
	359	+):
	360	+ new_graph = TracedGraph()
	361	+ exist_nodes: Dict[str, torch.fx.Node] = {}
	362	+
	363	+ def merge_graph(subgraphs: List[TracedGraph]) -> None:
	364	+ for input_graph in subgraphs:
	365	+ if isinstance(input_graph, List):
	366	+ merge_graph(input_graph)
	367	+ continue
	368	+ if not isinstance(input_graph, TracedGraph):
	369	+ continue
	370	+ for node in input_graph.graph.nodes:
	371	+ if node.name in exist_nodes:
	372	+ continue
	373	+ new_node = new_graph.graph.node_copy(node, lambda n: exist_nodes.get(n.name, n))
	374	+ exist_nodes[node.name] = new_node
	375	+ if node.name in input_graph.sym_nodes:
	376	+ new_graph.sym_nodes.update({node.name: new_node})
	377	+
	378	+ def parse_args(subgraphs, nodes):
	379	+ args = []
	380	+ for input_graph in subgraphs:
	381	+ if isinstance(input_graph, TracedGraph):
	382	+ args.append(nodes[input_graph.last_node.name])
	383	+ elif isinstance(input_graph, (List, Tuple)):
	384	+ args.append(parse_args(input_graph, nodes))
	385	+ else:
	386	+ if isinstance(input_graph, Expr) and not isinstance(input_graph, Integer):
	387	+ if not isinstance(input_graph, Symbol):
	388	+ input_graph = map_operators_to_strings(str(input_graph), operator_mapping)
	389	+ args.append(new_graph.sym_nodes[str(input_graph)])
	390	+ else:
	391	+ args.append(input_graph)
	392	+ return args
	393	+
	394	+ num_args = len(input_graphs)
	395	+
	396	+ for k, v in kwargs.items():
	397	+ if isinstance(v, Expr) and not isinstance(v, Integer):
	398	+ traced_graph = TracedGraph()
	399	+ create_sym_inputs(traced_graph, [v], operator_mapping)
	400	+ s_name = str(v)
	401	+ if not isinstance(v, Symbol):
	402	+ s_name = map_operators_to_strings(str(v), operator_mapping)
	403	+ traced_graph.last_node = traced_graph.sym_nodes[s_name]
	404	+ kwargs[k] = traced_graph.sym_nodes[s_name]
	405	+ input_graphs.append(traced_graph)
	406	+ merge_graph(input_graphs)
	407	+ input_graphs = input_graphs[:num_args]
	408	+ create_sym_inputs(new_graph, input_graphs, operator_mapping)
	409	+ args = parse_args(input_graphs, exist_nodes)
	410	+ with new_graph.graph.inserting_after(new_graph.last_node):
	411	+ new_node = new_graph.graph.call_function(origin_fn, args=tuple(args), kwargs=kwargs)
	412	+ new_node.name = node_name
	413	+ new_graph.last_node = new_node
	414	+ return new_graph
	415	+
	416	+
	417	+def merge_fx_graphs(traced_graphs: List[TracedGraph]):
	418	+ new_graph = TracedGraph()
	419	+ exist_nodes: Dict[str, torch.fx.Node] = {}
	420	+ last_nodes = []
	421	+
	422	+ def merge_graph(subgraphs: List[TracedGraph]) -> None:
	423	+ for input_graph in subgraphs:
	424	+ if isinstance(input_graph, List):
	425	+ merge_graph(input_graph)
	426	+ continue
	427	+ if not isinstance(input_graph, TracedGraph):
	428	+ continue
	429	+ for node in input_graph.graph.nodes:
	430	+ if node.name in exist_nodes:
	431	+ continue
	432	+ new_node = new_graph.graph.node_copy(node, lambda n: exist_nodes.get(n.name, n))
	433	+ exist_nodes[node.name] = new_node
	434	+ last_nodes.append(exist_nodes[input_graph.last_node.name])
	435	+
	436	+ merge_graph(traced_graphs)
	437	+ new_graph.last_node = last_nodes
	438	+ return new_graph
	439	+
	440	+
	441	+def subtract_graph(
	442	+ graph1: TracedGraph,
	443	+ graph2: TracedGraph,
	444	+ node_name=None,
	445	+) -> Tuple[TracedGraph, torch.fx.Node]:
	446	+ new_graph = TracedGraph()
	447	+ last_node2 = graph2.last_node
	448	+ graph1_node_names = {node.name for node in graph1.graph.nodes}
	449	+ graph2_node_names = {node.name for node in graph2.graph.nodes}
	450	+ placeholder = None
	451	+ exist_nodes: Dict[str, torch.fx.Node] = {}
	452	+ if node_name not in graph1_node_names:
	453	+ placeholder = new_graph.graph.placeholder(last_node2.name if node_name is None else node_name)
	454	+ exist_nodes[last_node2.name] = placeholder
	455	+ for node in graph1.graph.nodes:
	456	+ if node.name in graph2_node_names and node.name not in graph1.sym_nodes:
	457	+ continue
	458	+ new_node = new_graph.graph.node_copy(node, lambda n: exist_nodes.get(n.name, n))
	459	+ exist_nodes[node.name] = new_node
	460	+ new_graph.last_node = exist_nodes[graph1.last_node.name]
	461	+ new_graph.sym_nodes = graph1.sym_nodes
	462	+ return new_graph, placeholder

Mtorch_npu/_inductor/lowering_fx.py+31-236

	@@ -60,6 +60,22 @@ from torch.utils._sympy.functions import (
60	ModularIndexing,	60	ModularIndexing,
61	)	61	)
62	from .config import log	62	from .config import log
		63	+from .lowering_common import (
		64	+ TracedGraph,
		65	+ TRITON_OPERATOR_MAPPING,
		66	+ create_fake_input,
		67	+ create_sym_inputs as _create_sym_inputs,
		68	+ fetch_graphs as _fetch_graphs,
		69	+ get_reduction_type_to_aten_fn,
		70	+ map_operators_to_strings as _map_operators_to_strings,
		71	+ map_strings_to_operators as _map_strings_to_operators,
		72	+ merge_fx_graphs,
		73	+ merge_traced_graphs as _merge_traced_graphs,
		74	+ process_ir_constant as _process_ir_constant,
		75	+ register_fn_to_aten_fn as _register_fn_to_aten_fn,
		76	+ register_to_aten as _register_to_aten,
		77	+ subtract_graph,
		78	+)
63	from .lowering_op_list import GENERATE_LIST, GENERATE_LIST2, FALLBACK_LIST, LOWERING_OVERLOAD_OP	79	from .lowering_op_list import GENERATE_LIST, GENERATE_LIST2, FALLBACK_LIST, LOWERING_OVERLOAD_OP
64		80
65	aten = torch.ops.aten	81	aten = torch.ops.aten
	@@ -74,264 +90,43 @@ snodes_to_fx = {}
74		90
75		91
76	def register_fn_to_aten_fn(fn, aten_fn=None):	92	def register_fn_to_aten_fn(fn, aten_fn=None):
77	- if fn not in fn_to_aten_fn:	93	+ return _register_fn_to_aten_fn(fn_to_aten_fn, fn, aten_fn)
78	- fn_to_aten_fn[fn] = aten_fn
79	- return fn
80		94
81		95
82	def register_to_aten(aten_fn=None):	96	def register_to_aten(aten_fn=None):
83	- def decorator(fn):	97	+ return _register_to_aten(fn_to_aten_fn, aten_fn)
84	- if fn not in fn_to_aten_fn:
85	- fn_to_aten_fn[fn] = aten_fn
86	- return fn
87	-
88	- return decorator
89		98
90		99
91	-reduction_type_to_aten_fn = {	100	+reduction_type_to_aten_fn = get_reduction_type_to_aten_fn()
92	- "sum": aten.sum,
93	- "prod": aten.prod,
94	- "xor_sum": prims.xor_sum,
95	- "any": aten.any,
96	- "max": aten.amax,
97	- "min": aten.amin,
98	- "argmax": aten.argmax,
99	- "argmin": aten.argmin
100	-}
101		101
102	-operator_to_string = {	102	+operator_to_string = TRITON_OPERATOR_MAPPING.operator_to_string
103	- '+': 'a',	103	+string_to_operator = TRITON_OPERATOR_MAPPING.string_to_operator
104	- '-': 'sub',
105	- '*': 'm',
106	- '/': 'd',
107	- '(': 'l',
108	- ')': 'r',
109	- '.': 'p',
110	-}
111	-
112	-string_to_operator = {v: k for k, v in operator_to_string.items()}
113		104
114		105
115	def map_operators_to_strings(expr_str: str):	106	def map_operators_to_strings(expr_str: str):
116	- expr_str = expr_str.replace(' ', '')	107	+ return _map_operators_to_strings(expr_str, TRITON_OPERATOR_MAPPING)
117	- for op, string in operator_to_string.items():
118	- expr_str = expr_str.replace(op, string)
119	- return '_' + expr_str
120		108
121		109
122	def map_strings_to_operators(expr_str: str):	110	def map_strings_to_operators(expr_str: str):
123	- for op, string in string_to_operator.items():	111	+ return _map_strings_to_operators(expr_str, TRITON_OPERATOR_MAPPING)
124	- expr_str = expr_str.replace(op, string)
125	- return expr_str[1:]
126	-
127	-
128	-class TracedGraph:
129	- def __init__(self):
130	- self.graph = torch.fx.Graph()
131	- self.last_node: Optional[torch.fx.Node] = None
132	- self.sym_nodes: Dict[str, torch.fx.Node] = {}
133	-
134	- def __str__(self):
135	- return str(self.graph)
136	-
137	- def get_placeholder_names(self):
138	- placeholder_names = set()
139	- for node in self.graph.nodes:
140	- if node.op == 'placeholder' and node.name not in self.sym_nodes:
141	- placeholder_names.add(node.name)
142	- return placeholder_names
143	-
144	- __repr__ = __str__
145	-
146	-
147	-def create_fake_input(size, stride, device, dtype):
148	- size = [V.graph.sizevars.shape_env.create_symintnode(s, hint=None) \
149	- if isinstance(s, Expr) and not isinstance(s, Integer) else s for s in size]
150	- stride = [V.graph.sizevars.shape_env.create_symintnode(s, hint=None) \
151	- if isinstance(s, Expr) and not isinstance(s, Integer) else s for s in stride]
152	- with V.graph.fake_mode:
153	- fake_input = torch.empty_strided(size, stride, device=device, dtype=dtype)
154	- return fake_input
155		112
156		113
157	def create_sym_inputs(traced_graph: TracedGraph, size: List[Expr]):	114	def create_sym_inputs(traced_graph: TracedGraph, size: List[Expr]):
158	- for s in size:	115	+ return _create_sym_inputs(traced_graph, size, TRITON_OPERATOR_MAPPING)
159	- if isinstance(s, (List, Tuple)):
160	- create_sym_inputs(traced_graph, s)
161	- continue
162	- if isinstance(s, Expr) and not isinstance(s, Integer):
163	- s_name = str(s)
164	- if not isinstance(s, Symbol):
165	- s_name = map_operators_to_strings(s_name)
166	- if s_name in traced_graph.sym_nodes:
167	- continue
168	- new_node = traced_graph.graph.placeholder(s_name)
169	- new_node.meta['val'] = V.graph.sizevars.shape_env.create_symintnode(s, hint=None)
170	- traced_graph.sym_nodes.update({s_name: new_node})
171		116
172		117
173	def process_ir_constant(inp: ExpandView) -> Union[TracedGraph, int, float]:	118	def process_ir_constant(inp: ExpandView) -> Union[TracedGraph, int, float]:
174	- skip = False	119	+ return _process_ir_constant(inp, TRITON_OPERATOR_MAPPING)
175	- if isinstance(inp.data, IndexingConstant):
176	- dtype = inp.data.dtype
177	- inp = inp.data.index
178	- # convert to original dtype.
179	- if dtype in [torch.float32, torch.float16, torch.bfloat16]:
180	- # sympy inputs
181	- if isinstance(inp, Expr) and not isinstance(inp, sympy.core.numbers.Number):
182	- traced_graph = TracedGraph()
183	- create_sym_inputs(traced_graph, [inp])
184	- s_name = str(inp)
185	- if not isinstance(inp, Symbol):
186	- s_name = map_operators_to_strings(str(inp))
187	- traced_graph.last_node = traced_graph.sym_nodes[s_name]
188	- inp = traced_graph
189	- else:
190	- inp = float(inp)
191	- elif isinstance(inp.data, ir.Constant):
192	- dtype = inp.data.dtype
193	- inp = inp.data.value
194	- else:
195	- skip = True
196	- return inp, skip
197		120
198		121
199	def fetch_graphs(inputs: Optional[List[TensorBox]]):	122	def fetch_graphs(inputs: Optional[List[TensorBox]]):
200	- if isinstance(inputs, (TensorBox, ir.StorageBox, ir.View, sympy.Symbol, ir.Constant)):	123	+ return _fetch_graphs(inputs, TRITON_OPERATOR_MAPPING, use_npu_meta=True)
201	- inputs = [inputs]
202	- input_graphs = []
203	- for inp in inputs:
204	- if isinstance(inp, List):
205	- input_graphs.append(fetch_graphs(inp))
206	- continue
207	- if not isinstance(inp, (
208	- TensorBox, ir.StorageBox, ir.View, ir.ReinterpretView, ir.PermuteView, ir.SliceView, ir.ExpandView)):
209	- input_graphs.append(inp)
210	- continue
211	- if isinstance(inp, ExpandView):
212	- inp, skip = process_ir_constant(inp)
213	- if not skip:
214	- input_graphs.append(inp)
215	- continue
216	- name = inp.get_name()
217	- traced_graph = inp.get_traced_graph()
218	- if traced_graph is not None:
219	- input_graphs.append(traced_graph)
220	- continue
221	- traced_graph = TracedGraph()
222	- device = inp.get_device()
223	- dtype = inp.get_dtype()
224	- size = inp.get_size()
225	- stride = inp.get_stride()
226	- new_node = traced_graph.graph.placeholder(name)
227	- fake_input = create_fake_input(size, stride, device, dtype)
228	- new_node.meta['val'] = fake_input
229	- traced_graph.last_node = new_node
230	- input_graphs.append(traced_graph)
231	- return input_graphs
232		124
233		125
234	def merge_traced_graphs(input_graphs: List[TracedGraph], origin_fn, node_name, **kwargs):	126	def merge_traced_graphs(input_graphs: List[TracedGraph], origin_fn, node_name, **kwargs):
235	- new_graph = TracedGraph()	127	+ return _merge_traced_graphs(
236	- exist_nodes: Dict[str, torch.fx.Node] = {}	128	+ input_graphs, origin_fn, node_name, TRITON_OPERATOR_MAPPING, **kwargs
237	-	129	+ )
238	- def merge_graph(input_graphs: List[TracedGraph]):
239	- for input_graph in input_graphs:
240	- if isinstance(input_graph, List):
241	- merge_graph(input_graph)
242	- continue
243	- if not isinstance(input_graph, TracedGraph):
244	- continue
245	- for node in input_graph.graph.nodes:
246	- if node.name in exist_nodes:
247	- continue
248	- new_node = new_graph.graph.node_copy(node, lambda n: exist_nodes[n.name])
249	- exist_nodes[node.name] = new_node
250	- if node.name in input_graph.sym_nodes:
251	- new_graph.sym_nodes.update({node.name: new_node})
252	-
253	- def parse_args(input_graphs, exist_nodes):
254	- args = []
255	- for input_graph in input_graphs:
256	- if isinstance(input_graph, TracedGraph):
257	- args.append(exist_nodes[input_graph.last_node.name])
258	- elif isinstance(input_graph, (List, Tuple)):
259	- args.append(parse_args(input_graph, exist_nodes))
260	- else:
261	- if isinstance(input_graph, Expr) and not isinstance(input_graph, Integer):
262	- if not isinstance(input_graph, Symbol):
263	- input_graph = map_operators_to_strings(str(input_graph))
264	- args.append(new_graph.sym_nodes[str(input_graph)])
265	- else:
266	- args.append(input_graph)
267	- return args
268	-
269	- num_args = len(input_graphs)
270	-
271	- for k, v in kwargs.items():
272	- if isinstance(v, Expr) and not isinstance(v, Integer):
273	- traced_graph = TracedGraph()
274	- create_sym_inputs(traced_graph, [v])
275	- s_name = str(v)
276	- if not isinstance(v, Symbol):
277	- s_name = map_operators_to_strings(str(v))
278	- traced_graph.last_node = traced_graph.sym_nodes[s_name]
279	- kwargs[k] = traced_graph.sym_nodes[s_name]
280	- input_graphs.append(traced_graph)
281	- merge_graph(input_graphs)
282	- input_graphs = input_graphs[:num_args]
283	- # if inputs do not have any valid graphs, like full/iota
284	- create_sym_inputs(new_graph, input_graphs)
285	- args = parse_args(input_graphs, exist_nodes)
286	- with new_graph.graph.inserting_after(new_graph.last_node):
287	- new_node = new_graph.graph.call_function(origin_fn, args=tuple(args), kwargs=kwargs)
288	- new_node.name = node_name
289	- new_graph.last_node = new_node
290	- return new_graph
291	-
292	-
293	-def merge_fx_graphs(traced_graphs: List[TracedGraph]):
294	- new_graph = TracedGraph()
295	- exist_nodes: Dict[str, torch.fx.Node] = {}
296	- last_nodes = []
297	-
298	- def merge_graph(input_graphs: List[TracedGraph]):
299	- for input_graph in input_graphs:
300	- if isinstance(input_graph, List):
301	- merge_graph(input_graph)
302	- continue
303	- if not isinstance(input_graph, TracedGraph):
304	- continue
305	- for node in input_graph.graph.nodes:
306	- if node.name in exist_nodes:
307	- continue
308	- new_node = new_graph.graph.node_copy(node, lambda n: exist_nodes[n.name])
309	- exist_nodes[node.name] = new_node
310	- last_nodes.append(exist_nodes[input_graph.last_node.name])
311	-
312	- merge_graph(traced_graphs)
313	- new_graph.last_node = last_nodes
314	- return new_graph
315	-
316	-
317	-def subtract_graph(graph1: TracedGraph, graph2: TracedGraph, node_name=None) -> Tuple[TracedGraph, torch.fx.Node]:
318	- new_graph = TracedGraph()
319	- last_node2 = graph2.last_node
320	- graph1_node_names = {node.name for node in graph1.graph.nodes}
321	- graph2_node_names = {node.name for node in graph2.graph.nodes}
322	- placeholder = None
323	- exist_nodes: Dict[str, torch.fx.Node] = {}
324	- if node_name not in graph1_node_names:
325	- placeholder = new_graph.graph.placeholder(last_node2.name if node_name is None else node_name)
326	- exist_nodes[last_node2.name] = placeholder
327	- for node in graph1.graph.nodes:
328	- if node.name in graph2_node_names and node.name not in graph1.sym_nodes:
329	- continue
330	- new_node = new_graph.graph.node_copy(node, lambda n: exist_nodes[n.name])
331	- exist_nodes[node.name] = new_node
332	- new_graph.last_node = exist_nodes[graph1.last_node.name]
333	- new_graph.sym_nodes = graph1.sym_nodes
334	- return new_graph, placeholder
335		130
336		131
337	def get_last_node(gm: torch.fx.GraphModule):	132	def get_last_node(gm: torch.fx.GraphModule):
	@@ -1030,9 +825,9 @@ def _register_npu_inductor_fallbacks_fx(make_reduction):
1030	return TensorBox(SqueezeView.create(x.data))	825	return TensorBox(SqueezeView.create(x.data))
1031		826
1032	dim = (	827	dim = (
1033	- V.graph.sizevars.evaluate_static_shape(dim)	828	+ V.graph.sizevars.guard_int(dim)
1034	if isinstance(dim, (int, sympy.Expr))	829	if isinstance(dim, (int, sympy.Expr))
1035	- else tuple(V.graph.sizevars.evaluate_static_shape(d) for d in dim)	830	+ else tuple(V.graph.sizevars.guard_int(d) for d in dim)
1036	)	831	)
1037	dim = canonicalize_dims(len(x.get_size()), dim) # type: ignore[call-overload]	832	dim = canonicalize_dims(len(x.get_size()), dim) # type: ignore[call-overload]
1038	dims = set((dim,) if not isinstance(dim, tuple) else dim)	833	dims = set((dim,) if not isinstance(dim, tuple) else dim)

Atorch_npu/_inductor/lowering_patch.py+175-0

@@ -0,0 +1,175 @@
	1	+# Copyright (c) 2026, Huawei Technologies Co., Ltd
	2	+#
	3	+# Lowering snapshot / restore for multi-backend (Triton vs MLIR/DVM) switching.
	4	+
	5	+from __future__ import annotations
	6	+
	7	+import copy
	8	+import importlib
	9	+from dataclasses import dataclass, field
	10	+from typing import Any, Callable, Optional
	11	+
	12	+from .lowering_common import LOWERING_REGISTRY_ATTRS, get_module_functions
	13	+
	14	+_BASELINE: Optional["LoweringSnapshot"] = None
	15	+_INDUCTOR_ATTR_BASELINE = None
	16	+
	17	+
	18	+@dataclass
	19	+class LoweringSnapshot:
	20	+ """Pristine torch._inductor.lowering state captured before any NPU patch."""
	21	+
	22	+ functions: dict[str, Callable[..., Any]]
	23	+ lowerings_ref: dict[Any, Any]
	24	+ lowerings_copy: dict[Any, Any]
	25	+ registry_copies: dict[str, Any] = field(default_factory=dict)
	26	+ make_reduction: Any = None
	27	+
	28	+
	29	+def _get_inductor_lowering():
	30	+ from torch._inductor import lowering as inductor_lowering
	31	+
	32	+ return inductor_lowering
	33	+
	34	+
	35	+def _copy_registry_value(value: Any) -> Any:
	36	+ if hasattr(value, "copy"):
	37	+ try:
	38	+ return value.copy()
	39	+ except TypeError:
	40	+ pass
	41	+ if isinstance(value, dict):
	42	+ return dict(value)
	43	+ if isinstance(value, (set, list)):
	44	+ return type(value)(value)
	45	+ return copy.copy(value)
	46	+
	47	+
	48	+def _module_functions(module: Any) -> dict[str, Callable[..., Any]]:
	49	+ return get_module_functions(module)
	50	+
	51	+
	52	+def capture_lowering_baseline() -> LoweringSnapshot:
	53	+ """Capture PT lowering once; safe to call repeatedly."""
	54	+ global _BASELINE
	55	+ if _BASELINE is not None:
	56	+ return _BASELINE
	57	+
	58	+ lowering = _get_inductor_lowering()
	59	+ registry_copies = {
	60	+ attr: _copy_registry_value(getattr(lowering, attr))
	61	+ for attr in LOWERING_REGISTRY_ATTRS
	62	+ }
	63	+ _BASELINE = LoweringSnapshot(
	64	+ functions=_module_functions(lowering),
	65	+ lowerings_ref=lowering.lowerings,
	66	+ lowerings_copy=dict(lowering.lowerings),
	67	+ registry_copies=registry_copies,
	68	+ make_reduction=getattr(lowering, "make_reduction", None),
	69	+ )
	70	+ return _BASELINE
	71	+
	72	+
	73	+def restore_lowering_baseline() -> None:
	74	+ """Reset torch._inductor.lowering to the captured PT baseline."""
	75	+ baseline = capture_lowering_baseline()
	76	+ lowering = _get_inductor_lowering()
	77	+
	78	+ for name, func in baseline.functions.items():
	79	+ if hasattr(lowering, name):
	80	+ setattr(lowering, name, func)
	81	+
	82	+ if lowering.lowerings is not baseline.lowerings_ref:
	83	+ lowering.lowerings = baseline.lowerings_ref
	84	+ baseline.lowerings_ref.clear()
	85	+ baseline.lowerings_ref.update(baseline.lowerings_copy)
	86	+
	87	+ for attr in LOWERING_REGISTRY_ATTRS:
	88	+ target = getattr(lowering, attr)
	89	+ snapshot_value = baseline.registry_copies[attr]
	90	+ if hasattr(target, "clear") and hasattr(target, "update"):
	91	+ target.clear()
	92	+ target.update(snapshot_value)
	93	+ elif isinstance(target, dict):
	94	+ target.clear()
	95	+ target.update(snapshot_value)
	96	+ else:
	97	+ setattr(lowering, attr, _copy_registry_value(snapshot_value))
	98	+
	99	+ lowering.make_reduction = baseline.make_reduction
	100	+
	101	+
	102	+def _snapshot_inductor_attr(owner, name):
	103	+ return owner, name, hasattr(owner, name), getattr(owner, name, None)
	104	+
	105	+
	106	+def _get_inductor_attr_baseline():
	107	+ global _INDUCTOR_ATTR_BASELINE
	108	+ if _INDUCTOR_ATTR_BASELINE is not None:
	109	+ return _INDUCTOR_ATTR_BASELINE
	110	+
	111	+ from torch._inductor import scheduler as inductor_scheduler
	112	+
	113	+ Scheduler = inductor_scheduler.Scheduler
	114	+ _INDUCTOR_ATTR_BASELINE = (
	115	+ _snapshot_inductor_attr(Scheduler, "_codegen"),
	116	+ _snapshot_inductor_attr(Scheduler, "compute_ancestors"),
	117	+ _snapshot_inductor_attr(inductor_scheduler, "_prune_redundant_deps"),
	118	+ _snapshot_inductor_attr(Scheduler, "can_fuse_vertical"),
	119	+ _snapshot_inductor_attr(Scheduler, "_get_unmet_dep_nodes"),
	120	+ )
	121	+ return _INDUCTOR_ATTR_BASELINE
	122	+
	123	+
	124	+def restore_inductor_baseline() -> None:
	125	+ """Reset lowering and scheduler hooks before loading a new NPU backend."""
	126	+ attr_baseline = _get_inductor_attr_baseline()
	127	+ restore_lowering_baseline()
	128	+ for owner, name, exists, value in attr_baseline:
	129	+ if exists:
	130	+ setattr(owner, name, value)
	131	+ elif hasattr(owner, name):
	132	+ delattr(owner, name)
	133	+
	134	+def merge_missing_lowerings(
	135	+ target_lowerings: dict[Any, Any],
	136	+ source_lowerings: dict[Any, Any],
	137	+) -> None:
	138	+ extra_keys = set(source_lowerings.keys()) - set(target_lowerings.keys())
	139	+ if extra_keys:
	140	+ target_lowerings.update({k: source_lowerings[k] for k in extra_keys})
	141	+
	142	+
	143	+def apply_mlir_lowering_patch(npu_lowering_module: Any) -> None:
	144	+ """Replace torch._inductor.lowering with the MLIR/DVM fork."""
	145	+ from torch._inductor import graph, lowering as inductor_lowering
	146	+
	147	+ npu_functions = _module_functions(npu_lowering_module)
	148	+ inductor_functions = _module_functions(inductor_lowering)
	149	+ for name in inductor_functions:
	150	+ if name in npu_functions:
	151	+ setattr(inductor_lowering, name, npu_functions[name])
	152	+
	153	+ merge_missing_lowerings(
	154	+ npu_lowering_module.lowerings,
	155	+ inductor_lowering.lowerings,
	156	+ )
	157	+
	158	+ for attr in LOWERING_REGISTRY_ATTRS:
	159	+ setattr(inductor_lowering, attr, getattr(npu_lowering_module, attr))
	160	+
	161	+ importlib.reload(graph)
	162	+
	163	+
	164	+def apply_mlir_inductor_patch() -> None:
	165	+ """MLIR/DVM: patch lowering + scheduler (called from _load_backend)."""
	166	+ from .ascend_npu_ir.ascend_npu_ir.npu.inductor_patch import lowering as npu_lowering
	167	+ from .ascend_npu_ir.ascend_npu_ir.npu.inductor_patch.scheduler import (
	168	+ _patch_scheduler,
	169	+ )
	170	+
	171	+ # Ensure IR patches (TracedGraph hooks) are registered.
	172	+ import torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.inductor_patch.ir # noqa: F401
	173	+
	174	+ apply_mlir_lowering_patch(npu_lowering)
	OO O openLiBingCI6月1日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月1日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月1日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月1日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月1日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月1日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月2日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月2日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月2日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月2日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月2日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月2日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月2日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月2日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月9日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月9日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月9日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月9日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月9日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月9日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月9日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月9日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月9日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月9日此条代码评论区间+170至+174 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
	175	+ _patch_scheduler()
	OO O openLiBingCI6月1日此条代码评论区间+169至+175 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月1日此条代码评论区间+169至+175 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。