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

Atest/_inductor/test_inductor_graph_partition.py+485-0

@@ -0,0 +1,485 @@
	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.utils", "perf_hints")
	280	+ with ctx():
	281	+ foo(torch.ones([10], device="npu"), torch.ones([20]))
	282	+
	283	+ FileCheck().check_count(
	284	+ "cudagraph partition due to non gpu ops",
	285	+ 1,
	286	+ exactly=True,
	287	+ ).check_count(
	288	+ "cudagraph partition into 2 partitions", 1, exactly=True
	289	+ ).run(log_stream.getvalue())
	290	+
	291	+ log_stream, ctx = logs_to_string("torch_npu.npugraph", "cudagraphs")
	292	+ with ctx():
	293	+ # trigger recording
	294	+ foo(torch.ones([10], device="npu"), torch.ones([20]))
	295	+ foo(torch.ones([10], device="npu"), torch.ones([20]))
	296	+
	297	+ FileCheck().check_count(
	298	+ "[NPUGRAPH-TREE][Node][Record] function=0, graph=0",
	299	+ 1,
	300	+ exactly=True,
	301	+ ).run(log_stream.getvalue())
	302	+
	303	+ # -----------------------------------------------------------------------
	304	+ # CPU scalar tests
	305	+ # -----------------------------------------------------------------------
	306	+
	307	+ @config.patch("graph_partition", True)
	308	+ def test_graph_partition_cpu_scalar_device_put(self):
	309	+ @torch.compile(mode="reduce-overhead")
	310	+ def foo(x):
	311	+ y = x.to("npu")
	312	+ z = y.to("cpu")
	313	+ return z
	314	+
	315	+ x = torch.tensor(1)
	316	+ for _ in range(3):
	317	+ foo(x)
	318	+
	319	+ self.assertEqual(x, torch.tensor(1, device="cpu"))
	320	+
	321	+ @config.patch("graph_partition", True)
	322	+ def test_graph_partition_forward_with_skipped_cudagraphed_backward(self):
	323	+ @torch.compile(mode="reduce-overhead")
	324	+ def foo(x):
	325	+ return x * x * x
	326	+
	327	+ for _ in range(3):
	328	+ inp = torch.rand([20, 20], device="npu", requires_grad=True)
	329	+ out = foo(inp)
	330	+
	331	+ with config.patch(always_complex_memory_overlap_TESTING_ONLY=True):
	332	+ back_inp = torch.empty_strided([20, 20], [0, 1], device="npu")
	333	+ out.backward(back_inp)
	334	+
	335	+ # we should not have npugraph'd the backwards
	336	+ new_id = self.get_manager().new_graph_id().id
	337	+ self.assertEqual(new_id, 1)
	338	+
	339	+ self.assertFalse(self.get_manager().running_forwards_with_pending_backwards)
	340	+
	341	+ @config.patch("graph_partition", True)
	342	+ def test_graph_partition_dynamic_shapes(self):
	343	+ def foo(x):
	344	+ return x + 1
	345	+
	346	+ compiled_foo = torch.compile(foo, mode="reduce-overhead", fullgraph=True)
	347	+
	348	+ for input_shape in range(1, 4):
	349	+ for _ in range(3):
	350	+ compiled_foo(torch.randn(input_shape, device="npu"))
	351	+
	352	+ # 3 npugraphs for 3 input shapes
	353	+ self.assertEqual(self.get_manager().new_graph_id().id, 3)
	354	+
	355	+ @config.patch("graph_partition", True)
	356	+ def test_graph_partition_condition_op(self):
	357	+ def f(p, b):
	358	+ def true_fn(x):
	359	+ return torch.cos(x)
	360	+
	361	+ def false_fn(x):
	362	+ return torch.sin(x)
	363	+
	364	+ return torch.cond(p, true_fn, false_fn, [b])
	365	+
	366	+ compiled_f = torch.compile(f)
	367	+
	368	+ # static shape
	369	+ p = torch.tensor([True], device="npu")
	370	+ a = torch.ones([2, 3], device="npu")
	371	+ eager_out = f(p, a)
	372	+ compiled_out = compiled_f(p, a)
	373	+ self.assertEqual(eager_out, compiled_out)
	374	+
	375	+ # dynamic shape with backed symint
	376	+ p = torch.tensor([True], device="npu")
	377	+ a = torch.ones([4, 5], device="npu")
	378	+ eager_out = f(p, a)
	379	+ compiled_out = compiled_f(p, a)
	380	+ self.assertEqual(eager_out, compiled_out)
	381	+
	382	+ @config.patch("graph_partition", True)
	383	+ def test_graph_partition_reorder_cpu_and_gpu(self):
	384	+ def f(x_npu, y_cpu, z_npu, weight_npu, weight_cpu):
	385	+ x_npu0 = x_npu + 1
	386	+ x_npu1 = x_npu0 @ weight_npu
	387	+ x_npu2 = 2 * (x_npu1 + x_npu)
	388	+
	389	+ y_cpu0 = y_cpu + 1
	390	+ y_cpu1 = y_cpu0 @ weight_cpu
	391	+
	392	+ z_npu0 = z_npu + 1
	393	+ z_npu1 = z_npu0 @ weight_npu
	394	+ z_npu2 = 2 * (z_npu1 + z_npu)
	395	+
	396	+ return x_npu2, y_cpu1, z_npu2
	397	+
	398	+ x_npu = torch.randn(3, 3, device="npu")
	399	+ y_cpu = torch.randn(3, 3, device="cpu")
	400	+ z_npu = torch.randn(3, 3, device="npu")
	401	+ weight_npu = torch.randn(3, 3, device="npu")
	402	+ weight_cpu = torch.randn(3, 3, device="cpu")
	403	+
	404	+ eager_out = f(x_npu, y_cpu, z_npu, weight_npu, weight_cpu)
	405	+
	406	+ compiled_f = torch.compile(f, mode="reduce-overhead")
	407	+ for _ in range(3):
	408	+ compiled_out = compiled_f(x_npu, y_cpu, z_npu, weight_npu, weight_cpu)
	409	+ self.assertEqual(eager_out, compiled_out)
	410	+
	411	+ # reorder merges ops on npu into 1 graph partition
	412	+ self.assertEqual(self.get_manager().new_graph_id().id, 1)
	413	+
	414	+ @config.patch(implicit_fallbacks=True)
	415	+ @config.patch("graph_partition", True)
	416	+ def test_graph_partition_custom_op(self):
	417	+ @torch.library.custom_op(
	418	+ "mylib::movement_npu",
	419	+ mutates_args=(),
	420	+ tags=(torch._C.Tag.cudagraph_unsafe,),
	421	+ )
	422	+ def movement(pic: torch.Tensor) -> torch.Tensor:
	423	+ img = pic.cpu()
	424	+ cropped_img = (img + 1) * 2
	425	+ return cropped_img.npu() / 255.0
	426	+
	427	+ @movement.register_fake
	428	+ def _(pic):
	429	+ return torch.empty_like(pic)
	430	+
	431	+ @torch.library.custom_op(
	432	+ "mylib::modify_npu",
	433	+ mutates_args=(),
	434	+ tags=(torch._C.Tag.cudagraph_unsafe,),
	435	+ )
	436	+ def modify(pic: torch.Tensor) -> torch.Tensor:
	437	+ pic1 = pic + 1
	438	+ pic1_cpu = (pic1.cpu() + 1) * 2
	439	+ return pic1_cpu.npu() + pic
	440	+
	441	+ @modify.register_fake
	442	+ def _(pic):
	443	+ return torch.empty_like(pic)
	444	+
	445	+ @torch.library.custom_op("mylib::transform_npu", mutates_args=())
	446	+ def transform(pic: torch.Tensor) -> torch.Tensor:
	447	+ return (pic + 1) * 2
	448	+
	449	+ @transform.register_fake
	450	+ def _(pic):
	451	+ return torch.empty_like(pic)
	452	+
	453	+ img = torch.randn(3, 64, 64, device="npu")
	454	+
	455	+ def f(img):
	456	+ x = (img + 10) * 2
	457	+ y = movement(x)
	458	+ z = y + 1
	459	+ u = transform(z)
	460	+ v = 2 * u + 1
	461	+ out = modify(v)
	462	+ return out + 1
	463	+
	464	+ compiled_f = torch.compile(f, fullgraph=True)
	465	+
	466	+ eager_out = f(img)
	467	+ compiled_out = compiled_f(img)
	468	+
	469	+ self.assertEqual(eager_out, compiled_out)
	470	+
	471	+ compiled_f = torch.compile(f, mode="reduce-overhead", fullgraph=True)
	472	+
	473	+ eager_out = f(img)
	474	+
	475	+ for _ in range(3):
	476	+ compiled_out = compiled_f(img)
	477	+ self.assertEqual(eager_out, compiled_out)
	478	+
	479	+ # splitting on 2 custom gives 3 npugraphs
	480	+ self.assertEqual(self.get_manager().new_graph_id().id, 3)
	481	+
	482	+if __name__ == "__main__":
	483	+ from torch._inductor.test_case import run_tests
	484	+
	485	+ run_tests()

Mtest/npu/test_graph_tree.py+1-1

	@@ -295,7 +295,7 @@ class TestTreeManagerIntegration(TestCase):
295	def test_reset_npugraph_trees(self):	295	def test_reset_npugraph_trees(self):
296	get_container(0) # Initialize a container	296	get_container(0) # Initialize a container
297	reset_npugraph_trees()	297	reset_npugraph_trees()
298	- container_dict = getattr(local, "tree_manager_containers", {})	298	+ container_dict = getattr(local, "npu_tree_manager_containers", {})
299	self.assertEqual(len(container_dict), 0)	299	self.assertEqual(len(container_dict), 0)
300		300
301		301

Mtorch_npu/__init__.py+6-6

	@@ -16,7 +16,7 @@ import torch_npu.utils.patch_getenv
16	from torch_npu._init.core.module_loader import _load_core_modules	16	from torch_npu._init.core.module_loader import _load_core_modules
17	from torch_npu._init.core.optional_features import _enable_optional_features	17	from torch_npu._init.core.optional_features import _enable_optional_features
18	from torch_npu._init.core.runtime_lifecycle import _initialize_runtime_lifecycle	18	from torch_npu._init.core.runtime_lifecycle import _initialize_runtime_lifecycle
19	-from torch_npu._init.patches.patch_manager import _apply_patches	19	+from torch_npu._init.patches.patch_manager import _apply_all_patches
20	from torch_npu._init.registry.registry_manager import _register_components	20	from torch_npu._init.registry.registry_manager import _register_components
21	from torch_npu.version import __version__ as __version__	21	from torch_npu.version import __version__ as __version__
22		22
	@@ -53,14 +53,14 @@ def _initialize():
53	_register_components()	53	_register_components()
54		54
55	# 4. apply patches	55	# 4. apply patches
56	- _apply_patches()	56	+ _apply_all_patches()
57		57
58	- # 5. optional runtime features	58	+ # 5. final extension barrier and shutdown hook
59	- _enable_optional_features()
60	-
61	- # 6. final extension barrier and shutdown hook
62	_initialize_runtime_lifecycle()	59	_initialize_runtime_lifecycle()
63		60
		61	+ # 6. optional runtime features
		62	+ _enable_optional_features()
		63	+
64		64
65	_initialize()	65	_initialize()
66		66

Mtorch_npu/_inductor/codegen/wrapper.py+95-59

	@@ -22,7 +22,81 @@ from torch_npu._inductor.codegen.triton import NPUIndexTritonKernel
22	from torch_npu._inductor.npu_triton_heuristics import PrecomputedGridNpu, user_autotune_npu	22	from torch_npu._inductor.npu_triton_heuristics import PrecomputedGridNpu, user_autotune_npu
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: Optional[str] = None,
		72	+ gpu: bool = True,
		73	+ cpp_definition: Optional[str] = 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,7 +108,7 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
34	partition_signatures: Optional[GraphPartitionSignature] = None,	108	partition_signatures: Optional[GraphPartitionSignature] = None,
35	):	109	):
36	if is_subgraph:	110	if is_subgraph:
37	- return SubgraphPythonWrapperCodegen(subgraph_name, parent_wrapper, partition_signatures)	111	+ return NPUSubgraphWrapperCodegen(subgraph_name, parent_wrapper, partition_signatures)
38	return NPUWrapperCodeGen()	112	return NPUWrapperCodeGen()
39		113
40	def write_header(self) -> None:	114	def write_header(self) -> None:
	@@ -66,34 +140,6 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
66	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")
67	)	141	)
68		142
69	- # generate numel expr for range_tree_node
70	- def generate_node_numel_expr(self, kernel_name: str, node, numel_expr):
71	- expr = f"{kernel_name}_{node.name}_numel"
72	- if (expr, V.graph) not in self.kernel_numel_expr:
73	- # declare expr once in each graph (scope)
74	- self.kernel_numel_expr.add((expr, V.graph))
75	- self.writeline(
76	- f"{self.declare}{expr} = {self.expr_printer(numel_expr)}{self.ending}"
77	- )
78	- else:
79	- self.writeline(f"{expr} = {self.expr_printer(numel_expr)}{self.ending}")
80	- # We can get symbolic expressions here, like s0*64
81	- # It is fine to have them here, but we need to handle them correctly as their own type
82	- # This is tricky to do, so we wrap in a custom type, distinct from scalars, but also from sympy*
83	- # scalars as well.
84	- # This is handled in `generate_args_decl` which has a correct comment of: TODO: only works for
85	- # constant now, need type info. I agree, this needs type info, and while this is not true type info
86	- # it suffices as a type hint for the purposes of producing the correct code for this type.
87	- return SymbolicCallArg(expr, numel_expr)
88	-
89	- # don't assert
90	- def codegen_input_size_asserts(self) -> None:
91	- pass
92	-
93	- def get_next_kernel_suffix(self) -> str:
94	- iter_val = copy.copy(self._names_iter)
95	- return f"{next(iter_val)}"
96	-
97	def add_benchmark_harness(self, output):	143	def add_benchmark_harness(self, output):
98	"""	144	"""
99	Override, add aot-inductor debug kernel support.	145	Override, add aot-inductor debug kernel support.
	@@ -271,32 +317,22 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
271	self.wrapper_call.writeline('static_kernel_compiler.__exit__(*exc_info)')	317	self.wrapper_call.writeline('static_kernel_compiler.__exit__(*exc_info)')
272	super().generate_return(output_refs)	318	super().generate_return(output_refs)
273		319
274	- def define_kernel(	320	+class NPUSubgraphWrapperCodegen(_NPUKernelCodegenMixin, SubgraphPythonWrapperCodegen):
275	- self,	321	+ """
276	- kernel_name: str,	322	+ Partition subgraph wrapper for NPU.
277	- kernel_body: str,	323	+
278	- metadata: Optional[str] = None,	324	+ Inherits NPU kernel codegen specializations (define_kernel, numel_expr,
279	- gpu: bool = True,	325	+ make_buffer_free, codegen_input_size_asserts) via _NPUKernelCodegenMixin,
280	- cpp_definition: Optional[str] = None,	326	+ so user Triton kernels inside a partition subgraph get the NPU-flavored
281	- ):	327	+ user_autotune_npu / FixedGridNpu rewrite instead of the upstream default.
282	- # 重写父类逻辑，将triton_heuristics.user_autotune替换为npu_triton_heuristics.user_autotune_npu，	328	+
283	- # 将PrecomputedGrid替换为PrecomputedGridNpu，以适配NPU设备，避免core dump错误。	329	+ Also overrides get_next_kernel_suffix to delegate to parent_wrapper,
284	- if "user_autotune" in kernel_body and "user_autotune_npu" not in kernel_body:	330	+ matching the upstream next_kernel_suffix strategy - otherwise the
285	- kernel_body = kernel_body.replace(	331	+ "peek" counter in the subgraph would diverge from the "consume" counter
286	- "triton_heuristics.user_autotune(",	332	+ (which upstream already delegates to parent), producing mismatched
287	- "npu_triton_heuristics.user_autotune_npu("	333	+ kernel names between the kernel body placeholders and the real registered
288	- )	334	+ function name.
289	- kernel_body = kernel_body.replace(	335	+ """
290	- "PrecomputedGrid",	336	+
291	- "PrecomputedGridNpu"	337	+ def get_next_kernel_suffix(self) -> str:
292	- )	338	+ return self.parent_wrapper.get_next_kernel_suffix()
293	- kernel_body = kernel_body.replace(
294	- "FixedGrid",
295	- "FixedGridNpu"
296	- )
297	- #import npu_triton_heuristicsd相关头文件
298	- kernel_body = kernel_body.replace(
299	- "'''\n",
300	- "'''\n" + NPUIndexTritonKernel.gen_triton_ext_imports() + "\n"
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/_init/core/module_loader.py+0-3

	@@ -60,9 +60,6 @@ def _initialize_c_extension_children(required_children):
60	_create_child_once(_C, "_logging", "_logging_init")	60	_create_child_once(_C, "_logging", "_logging_init")
61	_create_child_once(_C, "_flops_count", "_flops_count_init")	61	_create_child_once(_C, "_flops_count", "_flops_count_init")
62		62
63	- # Optional RPC child, only if built.
64	- _create_child_once(_C, "_distributed_rpc", "_rpc_npu_init")
65	-
66	_register_c_extension_submodules(_C)	63	_register_c_extension_submodules(_C)
67	missing = [name for name in required_children if not hasattr(_C, name)]	64	missing = [name for name in required_children if not hasattr(_C, name)]
68	if missing:	65	if missing:

Mtorch_npu/_init/patches/patch_manager.py+3-2

	@@ -1,6 +1,7 @@
1	import pkgutil	1	import pkgutil
2	from collections import defaultdict	2	from collections import defaultdict
3	from collections.abc import Callable	3	from collections.abc import Callable
		4	+from typing import List, Optional
4	from importlib import import_module	5	from importlib import import_module
5		6
6		7
	@@ -45,7 +46,7 @@ class PatchManager:
45		46
46	_applied_patch_count = defaultdict(int)	47	_applied_patch_count = defaultdict(int)
47	_builtin_patches_registered = False	48	_builtin_patches_registered = False
48	- _custom_full_patch_order: list[str] \| None = None	49	+ _custom_full_patch_order: Optional[List[str]] = None
49	_patch_groups = defaultdict(list)	50	_patch_groups = defaultdict(list)
50	_patch_modules: list[str] = []	51	_patch_modules: list[str] = []
51		52
	@@ -207,7 +208,7 @@ class PatchManager:
207	cls._applied_patch_count.clear()	208	cls._applied_patch_count.clear()
208		209
209		210
210	-def _apply_patches():	211	+def _apply_all_patches():
211	PatchManager._register_builtin_patches()	212	PatchManager._register_builtin_patches()
212		213
213	for group in PatchManager._resolve_patch_order():	214	for group in PatchManager._resolve_patch_order():

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/contrib/transfer_to_npu.py+6-2

	@@ -340,7 +340,9 @@ def _patch_cuda():
340	['cuda.amp.common', torch_npu.npu.amp.common],	340	['cuda.amp.common', torch_npu.npu.amp.common],
341	['cuda.amp.grad_scaler', torch_npu.npu.amp.grad_scaler]	341	['cuda.amp.grad_scaler', torch_npu.npu.amp.grad_scaler]
342	]	342	]
343	- torch_npu._apply_patches(patchs)	343	+
		344	+ from torch_npu._init.patches.monkey_patches import _apply_patches
		345	+ _apply_patches(patchs)
344		346
345		347
346	def _patch_profiler():	348	def _patch_profiler():
	@@ -352,7 +354,9 @@ def _patch_profiler():
352	['profiler.ProfilerActivity.CUDA', torch_npu.profiler.ProfilerActivity.NPU],	354	['profiler.ProfilerActivity.CUDA', torch_npu.profiler.ProfilerActivity.NPU],
353	['profiler.ProfilerActivity.CPU', torch_npu.profiler.ProfilerActivity.CPU]	355	['profiler.ProfilerActivity.CPU', torch_npu.profiler.ProfilerActivity.CPU]
354	]	356	]
355	- torch_npu._apply_patches(patchs)	357	+
		358	+ from torch_npu._init.patches.monkey_patches import _apply_patches
		359	+ _apply_patches(patchs)
356		360
357		361
358	def _warning_fn(msg, rank0=True):	362	def _warning_fn(msg, rank0=True):

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/distributed/rpc/backend_registry.py+6-9

	@@ -286,16 +286,13 @@ def _npu_tensorpipe_init_backend_handler(
286		286
287		287
288	def _rpc_backend_registry():	288	def _rpc_backend_registry():
289	- if not hasattr(torch_npu._C, "_distributed_rpc"):	289	+ if hasattr(torch_npu._C, "_rpc_npu_init"):
290	- raise RuntimeError(	290	+ torch_npu._C._rpc_npu_init()
291	- "torch_npu._C._distributed_rpc must be initialized before RPC backend registration"	291	+ rpc.backend_registry.register_backend(
		292	+ "NPU_TENSORPIPE",
		293	+ _npu_tensorpipe_construct_rpc_backend_options_handler,
		294	+ _npu_tensorpipe_init_backend_handler,
292	)	295	)
293		296
294	- rpc.backend_registry.register_backend(
295	- "NPU_TENSORPIPE",
296	- _npu_tensorpipe_construct_rpc_backend_options_handler,
297	- _npu_tensorpipe_init_backend_handler,
298	- )
299	-
300	import torch.distributed.rpc as _rpc_module	297	import torch.distributed.rpc as _rpc_module
301	_rpc_module.BackendType = rpc.backend_registry.BackendType	298	_rpc_module.BackendType = rpc.backend_registry.BackendType

Mtorch_npu/npu/_graph_tree.py+14-9

	@@ -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)
	@@ -256,8 +256,13 @@ class TreeManagerContainer:
256	local = threading.local()	256	local = threading.local()
257		257
258	# one tree manager per device	258	# one tree manager per device
259	-local.tree_manager_containers = {}	259	+# Use npu-specific TLS keys to avoid colliding with upstream
260	-local.tree_manager_locks = defaultdict(threading.Lock)	260	+# torch._inductor.cudagraph_trees, which stashes the same-named objects under
		261	+# the keys "tree_manager_containers" / "tree_manager_locks". Upstream's stash
		262	+# does not INCREF, so overwriting an entry triggers use-after-free on the
		263	+# previous owner. See cudagraph_trees.py lines 306-321.
		264	+local.npu_tree_manager_containers = {}
		265	+local.npu_tree_manager_locks = defaultdict(threading.Lock)
261		266
262		267
263	# only incremented by user call of mark_step_begin	268	# only incremented by user call of mark_step_begin
	@@ -267,8 +272,8 @@ class MarkStepBox:
267		272
268	# We need to register this as an object that will be copied over as TLS when new	273	# We need to register this as an object that will be copied over as TLS when new
269	# threads are created in autograd	274	# threads are created in autograd
270	-torch._C._stash_obj_in_tls("tree_manager_containers", local.tree_manager_containers)	275	+torch._C._stash_obj_in_tls("npu_tree_manager_containers", local.npu_tree_manager_containers)
271	-torch._C._stash_obj_in_tls("tree_manager_locks", local.tree_manager_locks)	276	+torch._C._stash_obj_in_tls("npu_tree_manager_locks", local.npu_tree_manager_locks)
272		277
273		278
274	def mark_step_begin() -> None:	279	def mark_step_begin() -> None:
	@@ -281,8 +286,8 @@ def mark_step_begin() -> None:
281	def reset_npugraph_trees() -> None:	286	def reset_npugraph_trees() -> None:
282	"Clear all npugraph trees"	287	"Clear all npugraph trees"
283	# see shutdown below for why this is necessary	288	# see shutdown below for why this is necessary
284	- container_dict = get_obj(local, "tree_manager_containers")	289	+ container_dict = get_obj(local, "npu_tree_manager_containers")
285	- locks_dict = get_obj(local, "tree_manager_locks")	290	+ locks_dict = get_obj(local, "npu_tree_manager_locks")
286	for device, lock in locks_dict.items():	291	for device, lock in locks_dict.items():
287	with lock:	292	with lock:
288	container = container_dict.get(device)	293	container = container_dict.get(device)
	@@ -307,8 +312,8 @@ def get_obj(thread_local: Any, attr_name: str) -> Any:
307		312
308		313
309	def get_container(device_index: int) -> TreeManagerContainer:	314	def get_container(device_index: int) -> TreeManagerContainer:
310	- container_dict = get_obj(local, "tree_manager_containers")	315	+ container_dict = get_obj(local, "npu_tree_manager_containers")
311	- lock = get_obj(local, "tree_manager_locks")[device_index]	316	+ lock = get_obj(local, "npu_tree_manager_locks")[device_index]
312		317
313	with lock:	318	with lock:
314	if device_index not in container_dict:	319	if device_index not in container_dict:

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