feat: [graph partition] aclgraph support graph partition-pytorch-AtomGit

已合并

feat: [graph partition] aclgraph support graph partition #35327

luochao60创建于 5月11日从luochao60/pytorch_9490： pta_support_graph_partition_20260414_v2.12.0 合入到 Ascend/pytorch： v2.12.0

feat: [graph partition] aclgraph support graph partition #35327

已合并

luochao60创建于 5月11日

从luochao60/pytorch_9490 pta_support_graph_partition_20260414_v2.12.0合入到Ascend/pytorch v2.12.0

共 9 个文件变更+678-59

Atest/_inductor/test_inductor_graph_partition.py+545-0

@@ -0,0 +1,545 @@
	1	+import os
	2	+
	3	+# Some ops (e.g. aten.matmul_backward) only get their NPU meta registration when
	4	+# compatible impl mode is enabled. This env var is read at torch_npu import time,
	5	+# so it must be set before importing torch_npu.
	6	+os.environ.setdefault("TORCH_NPU_USE_COMPATIBLE_IMPL", "1")
	7	+
	8	+import contextlib
	9	+import gc
	10	+import math
	11	+import re
	12	+import sys
	13	+import unittest
	14	+import warnings
	15	+import weakref
	16	+from io import StringIO
	17	+
	18	+import torch
	19	+import torch.nn as nn
	20	+import torch._dynamo.config as dynamo_config
	21	+from torch._inductor import config
	22	+from torch._inductor.compile_fx import compile_fx_inner
	23	+from torch._inductor.utils import run_and_get_code
	24	+from torch._inductor.test_case import TestCase as InductorTestCase
	25	+from torch.fx.experimental.proxy_tensor import make_fx
	26	+from torch.testing import FileCheck
	27	+from torch.testing._internal.common_utils import (
	28	+ instantiate_parametrized_tests,
	29	+ parametrize,
	30	+)
	31	+from torch.testing._internal.logging_utils import logs_to_string
	32	+from torch.utils._python_dispatch import TorchDispatchMode
	33	+
	34	+import torch_npu # noqa: F401
	35	+from torch_npu.npu._graph_tree import get_container
	36	+
	37	+TEST_NPU = torch.npu.is_available()
	38	+aten = torch.ops.aten
	39	+
	40	+# ---------------------------------------------------------------------------
	41	+# Helpers
	42	+# ---------------------------------------------------------------------------
	43	+
	44	+
	45	+def get_num_partitions(code):
	46	+ """Get the number of graph partitions from generated code."""
	47	+ code = "".join(code)
	48	+ found = re.search(r"partitions=\[(.*)\]", code)
	49	+ assert found is not None, "Could not find partitions in generated code"
	50	+ partitions = found.group(1)
	51	+ return len([p for p in partitions.split(",") if p])
	52	+
	53	+
	54	+class capture_stderr(list):
	55	+ """Replace sys.stderr with a temporary StringIO."""
	56	+
	57	+ def __enter__(self):
	58	+ self.sys_stderr = sys.stderr
	59	+ self.stringio = StringIO()
	60	+ sys.stderr = self.stringio
	61	+ return self
	62	+
	63	+ def __exit__(self, *args):
	64	+ self.append(str(self.stringio.getvalue()))
	65	+ del self.stringio
	66	+ sys.stderr = self.sys_stderr
	67	+
	68	+
	69	+# ---------------------------------------------------------------------------
	70	+# Base test class
	71	+# ---------------------------------------------------------------------------
	72	+
	73	+
	74	+class TestCase(InductorTestCase):
	75	+ device = "npu"
	76	+
	77	+ @classmethod
	78	+ def setUpClass(cls):
	79	+ super().setUpClass()
	80	+ cls._stack = contextlib.ExitStack()
	81	+ cls._stack.enter_context(
	82	+ config.patch(
	83	+ {
	84	+ "debug": True,
	85	+ "cpp.min_chunk_size": 1,
	86	+ "triton.autotune_pointwise": False,
	87	+ "implicit_fallbacks": False,
	88	+ }
	89	+ )
	90	+ )
	91	+
	92	+ @classmethod
	93	+ def tearDownClass(cls):
	94	+ cls._stack.close()
	95	+ super().tearDownClass()
	96	+
	97	+ def setUp(self):
	98	+ torch._dynamo.reset()
	99	+ super().setUp()
	100	+
	101	+ def tearDown(self):
	102	+ super().tearDown()
	103	+ torch._dynamo.reset()
	104	+
	105	+
	106	+# ===========================================================================
	107	+# Graph Partition Tests — Codegen correctness
	108	+# (ported from test_torchinductor.py, device-agnostic via self.device)
	109	+# ===========================================================================
	110	+
	111	+
	112	+@unittest.skipIf(not TEST_NPU, "requires NPU")
	113	+class TestGraphPartitionCodegen(TestCase):
	114	+ """Tests that graph partition generates correct code and produces correct
	115	+ results. These tests do NOT verify npugraph tree state."""
	116	+
	117	+ @config.patch("graph_partition", True)
	118	+ def test_graph_partition_refcount(self):
	119	+ # Trigger NPU backend registration: compile_fx_inner is a low-level
	120	+ # entry that bypasses dynamo's lazy loading of torch_npu._inductor,
	121	+ # so without this warmup get_wrapper_codegen_for_device('npu') returns
	122	+ # None and init_wrapper_code asserts "Device npu not supported".
	123	+ @torch.compile
	124	+ def _warmup(x):
	125	+ return x + 1
	126	+ _warmup(torch.ones(2, device=self.device))
	127	+
	128	+ contexts = [
	129	+ contextlib.nullcontext,
	130	+ lambda: config.patch({"triton.cudagraphs": True}),
	131	+ ]
	132	+
	133	+ for context in contexts:
	134	+ with context():
	135	+ inps = [
	136	+ torch.rand([5, 5]).to(self.device),
	137	+ torch.rand([5, 5]).to(self.device),
	138	+ ]
	139	+ inp_refs = [weakref.ref(inp) for inp in inps]
	140	+
	141	+ def fn(x, y):
	142	+ a = x + y
	143	+ return (a @ a,)
	144	+
	145	+ fn_fx = make_fx(fn)(inps[0], inps[1])
	146	+ fn_compiled = compile_fx_inner(fn_fx, inps)
	147	+
	148	+ matmul_seen = False
	149	+
	150	+ class TestRefMode(TorchDispatchMode):
	151	+ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
	152	+ kwargs = kwargs if kwargs else {}
	153	+
	154	+ nonlocal inps
	155	+ nonlocal inp_refs
	156	+ nonlocal matmul_seen
	157	+
	158	+ gc.collect()
	159	+ if func is aten.mm.out:
	160	+ matmul_seen = True
	161	+ assert len(inps) == 0
	162	+ assert inp_refs[0]() is None
	163	+ assert inp_refs[1]() is None
	164	+
	165	+ return func(args, *kwargs)
	166	+
	167	+ with TestRefMode():
	168	+ fn_compiled(inps)
	169	+
	170	+ # do an extra run to make sure we are deallocating on warmup and record
	171	+ inps.extend(
	172	+ [
	173	+ torch.rand([5, 5]).to(self.device),
	174	+ torch.rand([5, 5]).to(self.device),
	175	+ ]
	176	+ )
	177	+ inp_refs.extend([weakref.ref(inp) for inp in inps])
	178	+ matmul_seen = False
	179	+
	180	+ with TestRefMode():
	181	+ fn_compiled(inps)
	182	+
	183	+ assert len(inps) == 0
	184	+
	185	+class TestGraphPartitionNPU(TestCase):
	186	+ """Tests that graph partition works end-to-end with NPU graph trees.
	187	+ Many tests verify npugraph tree state (partition count, graph id, etc.).
	188	+ """
	189	+
	190	+ def setUp(self):
	191	+ super().setUp()
	192	+ self.graph_stack = contextlib.ExitStack()
	193	+ self.graph_stack.enter_context(
	194	+ config.patch(
	195	+ {
	196	+ "triton.cudagraphs": True,
	197	+ "triton.cudagraph_trees": True,
	198	+ }
	199	+ )
	200	+ )
	201	+ self.graph_stack.enter_context(
	202	+ dynamo_config.patch(automatic_dynamic_shapes=True)
	203	+ )
	204	+ self.device_idx = torch.rand([0], device="npu").device.index
	205	+ warnings.filterwarnings("ignore")
	206	+
	207	+ def tearDown(self):
	208	+ super().tearDown()
	209	+ torch._dynamo.reset()
	210	+ gc.collect()
	211	+ torch.npu.empty_cache()
	212	+ self.graph_stack.close()
	213	+ # NPU's TreeManagerContainer holds a strong reference to the tree manager;
	214	+ # explicitly clear it so each test sees a fresh manager state under
	215	+ # pytest single-process execution.
	216	+ from torch_npu.npu._graph_tree import reset_npugraph_trees
	217	+ reset_npugraph_trees()
	218	+ warnings.resetwarnings()
	219	+
	220	+ def get_manager(self, device_index=None):
	221	+ return get_container(
	222	+ device_index if device_index else self.device_idx
	223	+ ).tree_manager
	224	+
	225	+ # -----------------------------------------------------------------------
	226	+ # Basic partition tests
	227	+ # -----------------------------------------------------------------------
	228	+
	229	+ def test_graph_partition_simple(self):
	230	+ def f(x, y):
	231	+ x1 = x + 1
	232	+ y1 = y + 1
	233	+ y_cpu = y1.cpu() + 1
	234	+ z = x @ y
	235	+ return x1 + y1 + z + y_cpu.to("npu")
	236	+
	237	+ x, y = [torch.ones(2, 2, device="npu") for _ in range(2)]
	238	+ x_cloned, y_cloned = [tmp.clone() for tmp in [x, y]]
	239	+ eager_out = f(x, y)
	240	+
	241	+ f_compiled = torch.compile(f)
	242	+ compiled_out = f_compiled(x_cloned, y_cloned)
	243	+ self.assertEqual(eager_out, compiled_out)
	244	+
	245	+ _, code = run_and_get_code(f_compiled, x_cloned, y_cloned)
	246	+
	247	+ if not config.cpp_wrapper:
	248	+ FileCheck().check("def partition_0(args):").check(
	249	+ "recursively_apply_fns = runner.recursively_apply_fns"
	250	+ ).run(code[0])
	251	+
	252	+ @config.patch("graph_partition", True)
	253	+ def test_graph_partition_view_fallback(self):
	254	+ def f(x):
	255	+ y = x + 1
	256	+ z = torch.ops.aten.view.dtype(y, torch.float8_e4m3fn)
	257	+ z_cpu = z.cpu()
	258	+ u_npu = z_cpu.npu()
	259	+ return u_npu
	260	+
	261	+ compiled_f = torch.compile(f, mode="reduce-overhead")
	262	+
	263	+ for _ in range(3):
	264	+ x = torch.ones(2, dtype=torch.int32, device="npu")
	265	+ eager_out = f(x)
	266	+ compiled_out = compiled_f(x)
	267	+ # NPU aclnnIsClose does not support float8, compare via int8 view
	268	+ self.assertEqual(
	269	+ eager_out.view(torch.int8), compiled_out.view(torch.int8)
	270	+ )
	271	+
	272	+ @config.patch("graph_partition", True)
	273	+ def test_graph_partition_log_message(self):
	274	+ def foo(x, y):
	275	+ return (x + 1, y + 2)
	276	+
	277	+ foo = torch.compile(foo, mode="reduce-overhead")
	278	+
	279	+ log_stream, ctx = logs_to_string("torch._inductor.scheduler", "cudagraphs")
	280	+ with ctx():
	281	+ foo(torch.ones([10], device="npu"), torch.ones([20]))
	282	+
	283	+ FileCheck().check_count(
	284	+ "Created 2 graph partitions: 1 cudagraphable, 1 non-cudagraphable",
	285	+ 1,
	286	+ exactly=True,
	287	+ ).check_count("reason=cpu ops", 1, exactly=True).run(log_stream.getvalue())
	288	+
	289	+ log_stream, ctx = logs_to_string("torch_npu.npugraph", "cudagraphs")
	290	+ with ctx():
	291	+ # trigger recording
	292	+ foo(torch.ones([10], device="npu"), torch.ones([20]))
	293	+ foo(torch.ones([10], device="npu"), torch.ones([20]))
	294	+
	295	+ FileCheck().check_count(
	296	+ "[NPUGRAPH-TREE][Node][Record] function=0, graph=0",
	297	+ 1,
	298	+ exactly=True,
	299	+ ).run(log_stream.getvalue())
	300	+
	301	+ # -----------------------------------------------------------------------
	302	+ # CPU scalar tests
	303	+ # -----------------------------------------------------------------------
	304	+
	305	+ @config.patch("graph_partition", True)
	306	+ def test_graph_partition_cpu_scalar_device_put(self):
	307	+ @torch.compile(mode="reduce-overhead")
	308	+ def foo(x):
	309	+ y = x.to("npu")
	310	+ z = y.to("cpu")
	311	+ return z
	312	+
	313	+ x = torch.tensor(1)
	314	+ for _ in range(3):
	315	+ foo(x)
	316	+
	317	+ self.assertEqual(x, torch.tensor(1, device="cpu"))
	318	+
	319	+ @config.patch("graph_partition", True)
	320	+ def test_graph_partition_forward_with_skipped_cudagraphed_backward(self):
	321	+ @torch.compile(mode="reduce-overhead")
	322	+ def foo(x):
	323	+ return x * x * x
	324	+
	325	+ for _ in range(3):
	326	+ inp = torch.rand([20, 20], device="npu", requires_grad=True)
	327	+ out = foo(inp)
	328	+
	329	+ with config.patch(always_complex_memory_overlap_TESTING_ONLY=True):
	330	+ back_inp = torch.empty_strided([20, 20], [0, 1], device="npu")
	331	+ out.backward(back_inp)
	332	+
	333	+ # we should not have npugraph'd the backwards
	334	+ new_id = self.get_manager().new_graph_id().id
	335	+ self.assertEqual(new_id, 1)
	336	+
	337	+ self.assertFalse(self.get_manager().running_forwards_with_pending_backwards)
	338	+
	339	+ @config.patch("graph_partition", True)
	340	+ def test_graph_partition_dynamic_shapes(self):
	341	+ def foo(x):
	342	+ return x + 1
	343	+
	344	+ compiled_foo = torch.compile(foo, mode="reduce-overhead", fullgraph=True)
	345	+
	346	+ for input_shape in range(1, 4):
	347	+ for _ in range(3):
	348	+ compiled_foo(torch.randn(input_shape, device="npu"))
	349	+
	350	+ # 3 npugraphs for 3 input shapes
	351	+ self.assertEqual(self.get_manager().new_graph_id().id, 3)
	352	+
	353	+ @config.patch("graph_partition", True)
	354	+ def test_graph_partition_condition_op(self):
	355	+ def f(p, b):
	356	+ def true_fn(x):
	357	+ return torch.cos(x)
	358	+
	359	+ def false_fn(x):
	360	+ return torch.sin(x)
	361	+
	362	+ return torch.cond(p, true_fn, false_fn, [b])
	363	+
	364	+ compiled_f = torch.compile(f)
	365	+
	366	+ # static shape
	367	+ p = torch.tensor([True], device="npu")
	368	+ a = torch.ones([2, 3], device="npu")
	369	+ eager_out = f(p, a)
	370	+ compiled_out = compiled_f(p, a)
	371	+ self.assertEqual(eager_out, compiled_out)
	372	+
	373	+ # dynamic shape with backed symint
	374	+ p = torch.tensor([True], device="npu")
	375	+ a = torch.ones([4, 5], device="npu")
	376	+ eager_out = f(p, a)
	377	+ compiled_out = compiled_f(p, a)
	378	+ self.assertEqual(eager_out, compiled_out)
	379	+
	380	+ @config.patch("graph_partition", True)
	381	+ def test_graph_partition_reorder_cpu_and_gpu(self):
	382	+ def f(x_npu, y_cpu, z_npu, weight_npu, weight_cpu):
	383	+ x_npu0 = x_npu + 1
	384	+ x_npu1 = x_npu0 @ weight_npu
	385	+ x_npu2 = 2 * (x_npu1 + x_npu)
	386	+
	387	+ y_cpu0 = y_cpu + 1
	388	+ y_cpu1 = y_cpu0 @ weight_cpu
	389	+
	390	+ z_npu0 = z_npu + 1
	391	+ z_npu1 = z_npu0 @ weight_npu
	392	+ z_npu2 = 2 * (z_npu1 + z_npu)
	393	+
	394	+ return x_npu2, y_cpu1, z_npu2
	395	+
	396	+ x_npu = torch.randn(3, 3, device="npu")
	397	+ y_cpu = torch.randn(3, 3, device="cpu")
	398	+ z_npu = torch.randn(3, 3, device="npu")
	399	+ weight_npu = torch.randn(3, 3, device="npu")
	400	+ weight_cpu = torch.randn(3, 3, device="cpu")
	401	+
	402	+ eager_out = f(x_npu, y_cpu, z_npu, weight_npu, weight_cpu)
	403	+
	404	+ compiled_f = torch.compile(f, mode="reduce-overhead")
	405	+ for _ in range(3):
	406	+ compiled_out = compiled_f(x_npu, y_cpu, z_npu, weight_npu, weight_cpu)
	407	+ self.assertEqual(eager_out, compiled_out)
	408	+
	409	+ # reorder merges ops on npu into 1 graph partition
	410	+ self.assertEqual(self.get_manager().new_graph_id().id, 1)
	411	+
	412	+ @config.patch(implicit_fallbacks=True)
	413	+ @config.patch("graph_partition", True)
	414	+ def test_graph_partition_custom_op(self):
	415	+ @torch.library.custom_op(
	416	+ "mylib::movement_npu",
	417	+ mutates_args=(),
	418	+ tags=(torch._C.Tag.cudagraph_unsafe,),
	419	+ )
	420	+ def movement(pic: torch.Tensor) -> torch.Tensor:
	421	+ img = pic.cpu()
	422	+ cropped_img = (img + 1) * 2
	423	+ return cropped_img.npu() / 255.0
	424	+
	425	+ @movement.register_fake
	426	+ def _(pic):
	427	+ return torch.empty_like(pic)
	428	+
	429	+ @torch.library.custom_op(
	430	+ "mylib::modify_npu",
	431	+ mutates_args=(),
	432	+ tags=(torch._C.Tag.cudagraph_unsafe,),
	433	+ )
	434	+ def modify(pic: torch.Tensor) -> torch.Tensor:
	435	+ pic1 = pic + 1
	436	+ pic1_cpu = (pic1.cpu() + 1) * 2
	437	+ return pic1_cpu.npu() + pic
	438	+
	439	+ @modify.register_fake
	440	+ def _(pic):
	441	+ return torch.empty_like(pic)
	442	+
	443	+ @torch.library.custom_op("mylib::transform_npu", mutates_args=())
	444	+ def transform(pic: torch.Tensor) -> torch.Tensor:
	445	+ return (pic + 1) * 2
	446	+
	447	+ @transform.register_fake
	448	+ def _(pic):
	449	+ return torch.empty_like(pic)
	450	+
	451	+ img = torch.randn(3, 64, 64, device="npu")
	452	+
	453	+ def f(img):
	454	+ x = (img + 10) * 2
	455	+ y = movement(x)
	456	+ z = y + 1
	457	+ u = transform(z)
	458	+ v = 2 * u + 1
	459	+ out = modify(v)
	460	+ return out + 1
	461	+
	462	+ compiled_f = torch.compile(f, fullgraph=True)
	463	+
	464	+ eager_out = f(img)
	465	+ compiled_out = compiled_f(img)
	466	+
	467	+ self.assertEqual(eager_out, compiled_out)
	468	+
	469	+ compiled_f = torch.compile(f, mode="reduce-overhead", fullgraph=True)
	470	+
	471	+ eager_out = f(img)
	472	+
	473	+ for _ in range(3):
	474	+ compiled_out = compiled_f(img)
	475	+ self.assertEqual(eager_out, compiled_out)
	476	+
	477	+ # splitting on 2 custom gives 3 npugraphs
	478	+ self.assertEqual(self.get_manager().new_graph_id().id, 3)
	479	+
	480	+ @config.patch("graph_partition", True)
	481	+ @config.patch("triton.cudagraphs", True)
	482	+ def test_graph_partition_subgraph_wrapper_user_autotune(self):
	483	+ """
	484	+ Probe for missing NPUSubgraphWrapperCodegen.define_kernel user_autotune
	485	+ replacement.
	486	+
	487	+ Trigger path:
	488	+ - user-defined @triton.jit kernel captured by torch.compile
	489	+ - upstream wrapper.py:2703 define_user_defined_triton_kernel produces
	490	+ kernel_body decorated with `@triton_heuristics.user_autotune(...)`
	491	+ - NPU must replace it to `npu_triton_heuristics.user_autotune_npu`,
	492	+ PrecomputedGrid/FixedGrid -> *Npu, and inject gen_triton_ext_imports.
	493	+ - if partition subgraph uses bare SubgraphPythonWrapperCodegen, the
	494	+ replacement never runs -> kernel_body keeps upstream CUDA-path
	495	+ decorator -> may core dump or misbehave on NPU.
	496	+
	497	+ Construction: cpu op forces partition boundary; user triton kernel is
	498	+ invoked inside partition.
	499	+ """
	500	+ import triton
	501	+ import triton.language as tl
	502	+
	503	+ @triton.jit
	504	+ def _my_add1_kernel(x_ptr, out_ptr, n, BLOCK: tl.constexpr):
	505	+ pid = tl.program_id(0)
	506	+ offs = pid * BLOCK + tl.arange(0, BLOCK)
	507	+ mask = offs < n
	508	+ x = tl.load(x_ptr + offs, mask=mask)
	509	+ tl.store(out_ptr + offs, x + 1, mask=mask)
	510	+
	511	+ def f(x):
	512	+ cpu_val = torch.tensor(3)
	513	+ _ = cpu_val.to("npu") # partition boundary
	514	+ out = torch.empty_like(x)
	515	+ n = x.numel()
	516	+ grid = lambda meta: (triton.cdiv(n, meta["BLOCK"]),)
	517	+ _my_add1_kernel[grid](x, out, n, BLOCK=128)
	518	+ return out + 0.0 # 防止 out 被优化掉
	519	+
	520	+ compiled_f = torch.compile(f, mode="reduce-overhead")
	521	+ x = torch.randn(128, device="npu")
	522	+ _, code = run_and_get_code(compiled_f, x)
	523	+ full_code = "\n".join(code) if isinstance(code, list) else code
	524	+ # Generated partition subgraph must carry NPU define_kernel overrides.
	525	+ # Currently failing because partition uses bare SubgraphPythonWrapperCodegen.
	526	+ self.assertIn(
	527	+ "user_autotune_npu", full_code,
	528	+ "partition subgraph did not apply NPU define_kernel override "
	529	+ "(expected `npu_triton_heuristics.user_autotune_npu`)",
	530	+ )
	531	+ self.assertIn(
	532	+ "FixedGridNpu", full_code,
	533	+ "partition subgraph did not apply NPU FixedGrid -> FixedGridNpu rewrite",
	534	+ )
	535	+ out = compiled_f(x)
	536	+ self.assertEqual(out, x + 1)
	537	+
	538	+
	539	+instantiate_parametrized_tests(TestGraphPartitionNPU)
	540	+
	541	+
	542	+if __name__ == "__main__":
	543	+ from torch._inductor.test_case import run_tests
	544	+
	545	+ run_tests()

Mtorch_npu/_inductor/codegen/wrapper.py+96-54

	@@ -22,7 +22,81 @@ from torch_npu._inductor import config as npu_config
22	from torch_npu._inductor.codegen.triton import NPUIndexTritonKernel	22	from torch_npu._inductor.codegen.triton import NPUIndexTritonKernel
23		23
24		24
25	-class NPUWrapperCodeGen(PythonWrapperCodegen):	25	+class _NPUKernelCodegenMixin:
		26	+ """
		27	+ NPU-specific cross-cutting overrides that both the main wrapper and the
		28	+ partition subgraph wrapper must apply. The mixin is not meant to be
		29	+ instantiated on its own; it is mixed into NPUWrapperCodeGen and
		30	+ NPUSubgraphWrapperCodegen as a base class.
		31	+
		32	+ Via cooperative multiple inheritance (super()), a single implementation
		33	+ works for both the main graph wrapper and the subgraph wrapper. This
		34	+ avoids code duplication and prevents main-wrapper-only logic
		35	+ (AOT debug / aclnn initialization / whole-graph benchmark harness, etc.)
		36	+ from leaking into subgraphs.
		37	+ """
		38	+
		39	+ # generate numel expr for range_tree_node
		40	+ def generate_node_numel_expr(self, kernel_name: str, node, numel_expr):
		41	+ expr = f"{kernel_name}_{node.name}_numel"
		42	+ if (expr, V.graph) not in self.kernel_numel_expr:
		43	+ # declare expr once in each graph (scope)
		44	+ self.kernel_numel_expr.add((expr, V.graph))
		45	+ self.writeline(
		46	+ f"{self.declare}{expr} = {self.expr_printer(numel_expr)}{self.ending}"
		47	+ )
		48	+ else:
		49	+ self.writeline(f"{expr} = {self.expr_printer(numel_expr)}{self.ending}")
		50	+ # We can get symbolic expressions here, like s0*64
		51	+ # It is fine to have them here, but we need to handle them correctly as their own type
		52	+ # This is tricky to do, so we wrap in a custom type, distinct from scalars, but also from sympy*
		53	+ # scalars as well.
		54	+ # This is handled in `generate_args_decl` which has a correct comment of: TODO: only works for
		55	+ # constant now, need type info. I agree, this needs type info, and while this is not true type info
		56	+ # it suffices as a type hint for the purposes of producing the correct code for this type.
		57	+ return SymbolicCallArg(expr, numel_expr)
		58	+
		59	+ # don't assert
		60	+ def codegen_input_size_asserts(self) -> None:
		61	+ pass
		62	+
		63	+ def get_next_kernel_suffix(self) -> str:
		64	+ iter_val = copy.copy(self._names_iter)
		65	+ return f"{next(iter_val)}"
		66	+
		67	+ def define_kernel(
		68	+ self,
		69	+ kernel_name: str,
		70	+ kernel_body: str,
		71	+ metadata: str \| None = None,
		72	+ gpu: bool = True,
		73	+ cpp_definition: str \| None = None,
		74	+ ):
		75	+ # Override the parent logic: replace triton_heuristics.user_autotune with
		76	+ # npu_triton_heuristics.user_autotune_npu, and replace PrecomputedGrid with
		77	+ # PrecomputedGridNpu, to adapt to the NPU device and avoid core dump errors.
		78	+ if "user_autotune" in kernel_body and "user_autotune_npu" not in kernel_body:
		79	+ kernel_body = kernel_body.replace(
		80	+ "triton_heuristics.user_autotune(",
		81	+ "npu_triton_heuristics.user_autotune_npu("
		82	+ )
		83	+ kernel_body = kernel_body.replace(
		84	+ "PrecomputedGrid",
		85	+ "PrecomputedGridNpu"
		86	+ )
		87	+ kernel_body = kernel_body.replace(
		88	+ "FixedGrid",
		89	+ "FixedGridNpu"
		90	+ )
		91	+ # import headers related to npu_triton_heuristics
		92	+ kernel_body = kernel_body.replace(
		93	+ "'''\n",
		94	+ "'''\n" + NPUIndexTritonKernel.gen_triton_ext_imports() + "\n"
		95	+ )
		96	+ super().define_kernel(kernel_name, kernel_body, metadata, gpu, cpp_definition)
		97	+
		98	+
		99	+class NPUWrapperCodeGen(_NPUKernelCodegenMixin, PythonWrapperCodegen):
26	def __init__(self):	100	def __init__(self):
27	super().__init__()	101	super().__init__()
28		102
	@@ -34,9 +108,7 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
34	partition_signatures: GraphPartitionSignature \| None = None,	108	partition_signatures: GraphPartitionSignature \| None = None,
35	):	109	):
36	if is_subgraph:	110	if is_subgraph:
37	- return SubgraphPythonWrapperCodegen(	111	+ return NPUSubgraphWrapperCodegen(subgraph_name, parent_wrapper, partition_signatures)
38	- subgraph_name, parent_wrapper, partition_signatures
39	- )
40	return NPUWrapperCodeGen()	112	return NPUWrapperCodeGen()
41		113
42	def write_header(self) -> None:	114	def write_header(self) -> None:
	@@ -68,34 +140,6 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
68	V.graph.device_ops.import_get_raw_stream_as("get_raw_stream")	140	V.graph.device_ops.import_get_raw_stream_as("get_raw_stream")
69	)	141	)
70		142
71	- # generate numel expr for range_tree_node
72	- def generate_node_numel_expr(self, kernel_name: str, node, numel_expr):
73	- expr = f"{kernel_name}_{node.name}_numel"
74	- if (expr, V.graph) not in self.kernel_numel_expr:
75	- # declare expr once in each graph (scope)
76	- self.kernel_numel_expr.add((expr, V.graph))
77	- self.writeline(
78	- f"{self.declare}{expr} = {self.expr_printer(numel_expr)}{self.ending}"
79	- )
80	- else:
81	- self.writeline(f"{expr} = {self.expr_printer(numel_expr)}{self.ending}")
82	- # We can get symbolic expressions here, like s0*64
83	- # It is fine to have them here, but we need to handle them correctly as their own type
84	- # This is tricky to do, so we wrap in a custom type, distinct from scalars, but also from sympy*
85	- # scalars as well.
86	- # This is handled in `generate_args_decl` which has a correct comment of: TODO: only works for
87	- # constant now, need type info. I agree, this needs type info, and while this is not true type info
88	- # it suffices as a type hint for the purposes of producing the correct code for this type.
89	- return SymbolicCallArg(expr, numel_expr)
90	-
91	- # don't assert
92	- def codegen_input_size_asserts(self) -> None:
93	- pass
94	-
95	- def get_next_kernel_suffix(self) -> str:
96	- iter_val = copy.copy(self._names_iter)
97	- return f"{next(iter_val)}"
98	-
99	def add_benchmark_harness(self, output):	143	def add_benchmark_harness(self, output):
100	"""	144	"""
101	Override, add aot-inductor debug kernel support.	145	Override, add aot-inductor debug kernel support.
	@@ -278,25 +322,23 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
278	)	322	)
279	super().generate_return(output_refs)	323	super().generate_return(output_refs)
280		324
281	- def define_kernel(	325	+
282	- self,	326	+class NPUSubgraphWrapperCodegen(_NPUKernelCodegenMixin, SubgraphPythonWrapperCodegen):
283	- kernel_name: str,	327	+ """
284	- kernel_body: str,	328	+ Partition subgraph wrapper for NPU.
285	- metadata: str \| None = None,	329	+
286	- gpu: bool = True,	330	+ Inherits NPU kernel codegen specializations (define_kernel, numel_expr,
287	- cpp_definition: str \| None = None,	331	+ codegen_input_size_asserts) via _NPUKernelCodegenMixin,
288	- ):	332	+ so user Triton kernels inside a partition subgraph get the NPU-flavored
289	- # 重写父类逻辑，将triton_heuristics.user_autotune替换为npu_triton_heuristics.user_autotune_npu，	333	+ user_autotune_npu / FixedGridNpu rewrite instead of the upstream default.
290	- # 将PrecomputedGrid替换为PrecomputedGridNpu，以适配NPU设备，避免core dump错误。	334	+
291	- if "user_autotune" in kernel_body and "user_autotune_npu" not in kernel_body:	335	+ Also overrides get_next_kernel_suffix to delegate to parent_wrapper,
292	- kernel_body = kernel_body.replace(	336	+ matching the upstream next_kernel_suffix strategy - otherwise the
293	- "triton_heuristics.user_autotune(",	337	+ "peek" counter in the subgraph would diverge from the "consume" counter
294	- "npu_triton_heuristics.user_autotune_npu(",	338	+ (which upstream already delegates to parent), producing mismatched
295	- )	339	+ kernel names between the kernel body placeholders and the real registered
296	- kernel_body = kernel_body.replace("PrecomputedGrid", "PrecomputedGridNpu")	340	+ function name.
297	- kernel_body = kernel_body.replace("FixedGrid", "FixedGridNpu")	341	+ """
298	- # import npu_triton_heuristicsd相关头文件	342	+
299	- kernel_body = kernel_body.replace(	343	+ def get_next_kernel_suffix(self) -> str:
300	- "'''\n", "'''\n" + NPUIndexTritonKernel.gen_triton_ext_imports() + "\n"	344	+ return self.parent_wrapper.get_next_kernel_suffix()
301	- )
302	- super().define_kernel(kernel_name, kernel_body, metadata, gpu, cpp_definition)

Mtorch_npu/_inductor/lowering_op_list.py+3-0

	@@ -8,6 +8,7 @@ prims = torch.ops.prims
8		8
9	GENERATE_LIST = [	9	GENERATE_LIST = [
10	prims.iota,	10	prims.iota,
		11	+ prims.device_put,
11	aten.full,	12	aten.full,
12	aten.mul,	13	aten.mul,
13	aten.add,	14	aten.add,
	@@ -74,11 +75,13 @@ GENERATE_LIST = [
74	aten.isnan,	75	aten.isnan,
75	aten.bitwise_and,	76	aten.bitwise_and,
76	aten.squeeze,	77	aten.squeeze,
		78	+ aten.unbind,
77	aten.copy,	79	aten.copy,
78	aten.reciprocal,	80	aten.reciprocal,
79	aten._assert_scalar,	81	aten._assert_scalar,
80	triton_kernel_wrapper_mutation,	82	triton_kernel_wrapper_mutation,
81	torch.ops.higher_order.invoke_subgraph,	83	torch.ops.higher_order.invoke_subgraph,
		84	+ torch.ops.higher_order.cond,
82	torch.ops._inductor_test.realize,	85	torch.ops._inductor_test.realize,
83	torch.ops._inductor_test.realize.default,	86	torch.ops._inductor_test.realize.default,
84	]	87	]

Mtorch_npu/_logging/_internal.py+1-2

	@@ -41,9 +41,8 @@ def _add_logging_module():
41	torch._logging._internal.register_log("shmem", "torch_npu.symmetric_memory")	41	torch._logging._internal.register_log("shmem", "torch_npu.symmetric_memory")
42	torch._logging._internal.register_log("env", "torch_npu.env")	42	torch._logging._internal.register_log("env", "torch_npu.env")
43	torch._logging._internal.register_log("acl", "torch_npu.acl")	43	torch._logging._internal.register_log("acl", "torch_npu.acl")
44	- torch._logging._internal.register_log("aclgraph", "torch_npu.aclgraph")	44	+ torch._logging._internal.register_log("aclgraph", "torch_npu.npugraph")
45	torch._logging._internal.register_log("npugraph", "torch_npu.npugraph")	45	torch._logging._internal.register_log("npugraph", "torch_npu.npugraph")
46	- torch._logging._internal.register_log("cudagraphs", "torch_npu.npugraph")
47		46
48		47
49	def _update_log_state_from_env():	48	def _update_log_state_from_env():

Mtorch_npu/csrc/core/npu/NPUHooksInterface.cpp+11-0

	@@ -1,5 +1,6 @@
1	#include "torch_npu/csrc/core/npu/NPUHooksInterface.h"	1	#include "torch_npu/csrc/core/npu/NPUHooksInterface.h"
2	#include "torch_npu/csrc/core/npu/NPUFunctions.h"	2	#include "torch_npu/csrc/core/npu/NPUFunctions.h"
		3	+#include "torch_npu/csrc/core/npu/CachingHostAllocator.h"
3	#include "torch_npu/csrc/core/NPUStorageImpl.h"	4	#include "torch_npu/csrc/core/NPUStorageImpl.h"
4	#include "torch_npu/csrc/framework/FormatHelper.h"	5	#include "torch_npu/csrc/framework/FormatHelper.h"
5	#include "torch_npu/csrc/aten/common/ResizeNpu.h"	6	#include "torch_npu/csrc/aten/common/ResizeNpu.h"
	@@ -63,6 +64,16 @@ bool NPUHooksInterface::isAvailable() const
63	return c10_npu::device_count() > 0;	64	return c10_npu::device_count() > 0;
64	}	65	}
65		66
		67	+bool NPUHooksInterface::isPinnedPtr(const void* data) const
		68	+{
		69	+ return at_npu::native::CachingHostAllocator_isPinned(const_cast<void*>(data));
		70	+}
		71	+
		72	+c10::Allocator* NPUHooksInterface::getPinnedMemoryAllocator() const
		73	+{
		74	+ return at_npu::native::getPinnedMemoryAllocator();
		75	+}
		76	+
66	at::PrivateUse1HooksInterface* get_npu_hooks()	77	at::PrivateUse1HooksInterface* get_npu_hooks()
67	{	78	{
68	static at::PrivateUse1HooksInterface* npu_hooks;	79	static at::PrivateUse1HooksInterface* npu_hooks;

Mtorch_npu/csrc/core/npu/NPUHooksInterface.h+2-0

	@@ -18,6 +18,8 @@ struct TORCH_API NPUHooksInterface : public at::PrivateUse1HooksInterface {
18	bool hasPrimaryContext(c10::DeviceIndex device_index) const override;	18	bool hasPrimaryContext(c10::DeviceIndex device_index) const override;
19	void resizePrivateUse1Bytes(const c10::Storage &storage, size_t new_bytes) const;	19	void resizePrivateUse1Bytes(const c10::Storage &storage, size_t new_bytes) const;
20	bool isAvailable() const override;	20	bool isAvailable() const override;
		21	+ bool isPinnedPtr(const void* data) const override;
		22	+ c10::Allocator* getPinnedMemoryAllocator() const override;
21	};	23	};
22		24
23	struct TORCH_API NPUHooksArgs : public at::PrivateUse1HooksArgs {};	25	struct TORCH_API NPUHooksArgs : public at::PrivateUse1HooksArgs {};

Mtorch_npu/npu/_graph_tree.py+1-1

	@@ -115,7 +115,7 @@ StorageWeakRefPointer = int
115	StorageDataPtr = int	115	StorageDataPtr = int
116	NBytes = int	116	NBytes = int
117	S = TypeVar("S", bound="StorageWeakRefWrapper")	117	S = TypeVar("S", bound="StorageWeakRefWrapper")
118	-log = logging.getLogger("torch_npu.npugraph")	118	+log = torch._logging.getArtifactLogger("torch_npu.npugraph", "cudagraphs")
119		119
120		120
121	@dataclasses.dataclass(frozen=True)	121	@dataclasses.dataclass(frozen=True)

Mtorch_npu/npu/graphs.py+1-1

	@@ -61,7 +61,7 @@ from torch_npu._C import ( # noqa: F401
61	)	61	)
62		62
63		63
64	-log = logging.getLogger("torch_npu.npugraph")	64	+log = torch._logging.getArtifactLogger("torch_npu.npugraph", "cudagraphs")
65		65
66		66
67	def is_current_stream_capturing():	67	def is_current_stream_capturing():

Mtorch_npu/utils/_graph_tree.py+18-1

	@@ -54,7 +54,7 @@ from torch.multiprocessing.reductions import StorageWeakRef
54	import torch_npu.npu.aclnn	54	import torch_npu.npu.aclnn
55		55
56		56
57	-log = logging.getLogger("torch_npu.aclgraph")	57	+log = torch._logging.getArtifactLogger("torch_npu.npugraph", "cudagraphs")
58		58
59		59
60	def npugraph_mark_step_begin():	60	def npugraph_mark_step_begin():
	@@ -69,6 +69,12 @@ def check_multiple_devices_or_any_cpu_nodes(
69	if npu_config.npugraph_trees.disable_cpu_input_check:	69	if npu_config.npugraph_trees.disable_cpu_input_check:
70	device_node_mapping.pop(torch.device("cpu"), None)	70	device_node_mapping.pop(torch.device("cpu"), None)
71		71
		72	+ device_node_mapping.pop(torch.device("meta"), None)
		73	+
		74	+ from torch._inductor.utils import is_using_cudagraph_partition
		75	+ if is_using_cudagraph_partition():
		76	+ device_node_mapping.pop(torch.device("cpu"), None)
		77	+
72	cpu_node = device_node_mapping.get(torch.device("cpu"))	78	cpu_node = device_node_mapping.get(torch.device("cpu"))
73	if cpu_node:	79	if cpu_node:
74	msg = f"cpu device ({cpu_node.name})"	80	msg = f"cpu device ({cpu_node.name})"
	@@ -382,3 +388,14 @@ def _apply_npugraph_tree_methods():
382	torch._inductor.compile_fx.cudagraphify = npugraphify	388	torch._inductor.compile_fx.cudagraphify = npugraphify
383	torch._inductor.cudagraph_utils.check_multiple_devices_or_any_cpu_nodes = check_multiple_devices_or_any_cpu_nodes	389	torch._inductor.cudagraph_utils.check_multiple_devices_or_any_cpu_nodes = check_multiple_devices_or_any_cpu_nodes
384	torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin	390	torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin
		391	+
		392	+ # Bridge upstream callers of `torch._inductor.cudagraph_trees.get_manager`
		393	+ # to the NPU manager registry. The only upstream call sites are
		394	+ # `_inductor/output_code.py:maybe_handle_backward_generation` (used when
		395	+ # forward was cudagraph'd but backward is not, to drive the cudagraph
		396	+ # generation state machine) and `_dynamo/backends/cudagraphs.py`. NPU
		397	+ # registers its manager under `torch_npu.npu._graph_tree`, so without
		398	+ # this forward those upstream paths raise AttributeError or return None.
		399	+ import torch._inductor.cudagraph_trees as _upstream_cgt # noqa: F401
		400	+ from torch_npu.npu._graph_tree import get_manager as _npu_get_manager
		401	+ _upstream_cgt.get_manager = _npu_get_manager