Star844
1313
代码
代码
Issues1468
Pull Requests2319
流水线
项目讨论
Wiki
分析
Star844
1313
已合并
fix: npugraph_ex补充用户自定义backend和自定义pass缺失的接口 #30367
Xu Zhenhua创建于 2月2日
fix: npugraph_ex补充用户自定义backend和自定义pass缺失的接口 #30367
已合并
Xu Zhenhua创建于 2月2日
从xuzhenhua2025/pytorchv2.10.0合入到Ascend/pytorchv2.10.0
共 5 个文件变更+256-73
Mtest/dynamo/test_npu_backend.py+0-64
@@ -1,18 +1,10 @@
1# Owner(s): ["module: dynamo"]1# Owner(s): ["module: dynamo"]
2-import unittest
3-import dataclasses
4-from typing import List
5- 
6import torch2import torch
7from torch._dynamo.test_case import TestCase3from torch._dynamo.test_case import TestCase
8-from torch._dynamo.testing import same
9 4 
10import torch_npu5import torch_npu
11 6 
12 7 
13-DEVICE_NAME = torch_npu.npu.get_device_name(0)[:10]
14- 
15- 
16class TestNpuBackend(TestCase):8class TestNpuBackend(TestCase):
17 def test_optimize_npu(self):9 def test_optimize_npu(self):
18 func = torch.ops.aten.relu.default10 func = torch.ops.aten.relu.default
@@ -23,62 +15,6 @@ class TestNpuBackend(TestCase):
23 dynamo_result = dynamo_func(x)15 dynamo_result = dynamo_func(x)
24 self.assertEqual(eager_result, dynamo_result)16 self.assertEqual(eager_result, dynamo_result)
25 17 
26- @unittest.skipIf(DEVICE_NAME == 'Ascend910A', "capture is not supported on 910A, skip this ut.")
27- def test_npugraph_ex_backend(self):
28- class Model(torch.nn.Module):
29- def __init__(self):
30- super().__init__()
31- 
32- def forward(self, x, y):
33- return x + y
34- 
35- compiled_model = torch.compile(Model().npu(), backend="npugraph_ex", fullgraph=True, dynamic=False)
36- x = torch.ones(1, dtype=torch.int32).npu()
37- y = torch.ones(1, dtype=torch.int32).npu()
38- z = compiled_model(x, y)
39- self.assertEqual(z.item(), 2)
40- 
41- @unittest.skipIf(DEVICE_NAME == 'Ascend910A', "capture is not supported on 910A, skip this ut.")
42- def test_npugraph_ex_cache_compile(self):
43- @dataclasses.dataclass
44- class InputMeta:
45- data: torch.Tensor
46- is_prompt: bool
47- 
48- class Model(torch.nn.Module):
49- def __init__(self):
50- super().__init__()
51- self.linear1 = torch.nn.Linear(2, 1)
52- self.linear2 = torch.nn.Linear(2, 1)
53- for param in self.parameters():
54- torch.nn.init.ones_(param)
55- self.cached_prompt = torch.npu.npugraph_ex.inference.cache_compile(self.prompt)
56- self.cached_decode = torch.npu.npugraph_ex.inference.cache_compile(self.decode)
57- 
58- def forward(self, x: InputMeta, kv: List[torch.Tensor]):
59- if x.is_prompt:
60- return self.cached_prompt(x, kv)
61- return self.cached_decode(x, kv)
62- 
63- def _forward(self, x, kv):
64- return self.linear2(x.data) + self.linear2(kv[0])
65- 
66- def prompt(self, x, y):
67- return self._forward(x, y)
68- 
69- def decode(self, x, y):
70- return self._forward(x, y)
71- 
72- x = InputMeta(data=torch.ones(2, 2).npu(), is_prompt=True)
73- kv = [torch.ones(2, 2).npu()]
74- model = Model().npu()
75- res_prompt = model(x, kv)
76- x.is_prompt = False
77- res_decode = model(x, kv)
78- res = torch.empty(2, 1).npu().fill_(6.0)
79- self.assertTrue(same(res, res_prompt))
80- self.assertTrue(same(res, res_decode))
81- 
82 18 
83if __name__ == "__main__":19if __name__ == "__main__":
84 from torch._dynamo.test_case import run_tests20 from torch._dynamo.test_case import run_tests
Atest/dynamo/test_npugraph_ex.py+210-0
@@ -0,0 +1,210 @@
1+# Owner(s): ["module: dynamo"]
2+import dataclasses
3+import functools
4+import unittest
5+from typing import List
6+ 
7+import torch
8+from torch._dynamo.test_case import TestCase
9+from torch._dynamo.testing import same
10+from torch._functorch.aot_autograd import aot_module_simplified
11+from torch._inductor.pattern_matcher import Match
12+from torch._subclasses.fake_tensor import FakeTensorMode
13+ 
14+import torch_npu
15+ 
16+DEVICE_NAME = torch_npu.npu.get_device_name(0)[:10]
17+ 
18+ 
19+@unittest.skipIf(DEVICE_NAME == 'Ascend910A', "capture is not supported on 910A, skip this ut.")
20+class TestNpuGraphEx(TestCase):
21+ def test_backend(self):
22+ class Model(torch.nn.Module):
23+ def __init__(self):
24+ super().__init__()
25+ 
26+ def forward(self, x, y):
27+ return x + y
28+ 
29+ compiled_model = torch.compile(Model().npu(), backend="npugraph_ex", fullgraph=True, dynamic=False)
30+ x = torch.ones(1, dtype=torch.int32, device="npu")
31+ y = torch.ones(1, dtype=torch.int32, device="npu")
32+ z = compiled_model(x, y)
33+ self.assertEqual(z.item(), 2)
34+ 
35+ def test_compile_fx(self):
36+ class Model(torch.nn.Module):
37+ def __init__(self):
38+ super().__init__()
39+ 
40+ def forward(self, x, y):
41+ return x + y
42+ 
43+ def my_backend(gm: torch.fx.GraphModule, example_inputs):
44+ compiler = torch.npu.npugraph_ex.compile_fx()
45+ return aot_module_simplified(gm, example_inputs, fw_compiler=compiler)
46+ 
47+ compiled_model = torch.compile(Model().npu(), backend=my_backend, fullgraph=True, dynamic=False)
48+ x = torch.ones(1, dtype=torch.int32, device="npu")
49+ y = torch.ones(1, dtype=torch.int32, device="npu")
50+ z = compiled_model(x, y)
51+ self.assertEqual(z.item(), 2)
52+ 
53+ def test_cache_compile(self):
54+ @dataclasses.dataclass
55+ class InputMeta:
56+ data: torch.Tensor
57+ is_prompt: bool
58+ 
59+ class Model(torch.nn.Module):
60+ def __init__(self):
61+ super().__init__()
62+ self.linear1 = torch.nn.Linear(2, 1)
63+ self.linear2 = torch.nn.Linear(2, 1)
64+ for param in self.parameters():
65+ torch.nn.init.ones_(param)
66+ self.cached_prompt = torch.npu.npugraph_ex.inference.cache_compile(self.prompt)
67+ self.cached_decode = torch.npu.npugraph_ex.inference.cache_compile(self.decode)
68+ 
69+ def forward(self, x: InputMeta, kv: List[torch.Tensor]):
70+ if x.is_prompt:
71+ return self.cached_prompt(x, kv)
72+ return self.cached_decode(x, kv)
73+ 
74+ def _forward(self, x, kv):
75+ return self.linear2(x.data) + self.linear2(kv[0])
76+ 
77+ def prompt(self, x, y):
78+ return self._forward(x, y)
79+ 
80+ def decode(self, x, y):
81+ return self._forward(x, y)
82+ 
83+ x = InputMeta(data=torch.ones(2, 2).npu(), is_prompt=True)
84+ kv = [torch.ones(2, 2).npu()]
85+ model = Model().npu()
86+ res_prompt = model(x, kv)
87+ x.is_prompt = False
88+ res_decode = model(x, kv)
89+ res = torch.empty(2, 1).npu().fill_(6.0)
90+ self.assertTrue(same(res, res_prompt))
91+ self.assertTrue(same(res, res_decode))
92+ 
93+ def test_register_replacement(self):
94+ def search_fn(x1, x2, gamma):
95+ x_out = torch.add(x1, x2)
96+ y, _ = torch_npu.npu_rms_norm(x_out, gamma)
97+ return y, x_out
98+ 
99+ def replace_fn(x1, x2, gamma):
100+ y, _, x_out = torch_npu.npu_add_rms_norm(
101+ x1, x2, gamma
102+ )
103+ return y, x_out
104+ 
105+ def extra_check(match: Match):
106+ x1 = match.kwargs.get("x1")
107+ 
108+ if x1 is None:
109+ return False
110+ if not hasattr(x1, "meta") or "val" not in x1.meta:
111+ return False
112+ 
113+ a_shape = x1.meta["val"].shape
114+ return a_shape[-1] == 7168
115+ 
116+ fake_mode = FakeTensorMode()
117+ with fake_mode:
118+ input_tensor = functools.partial(torch.empty, (1, 1, 2), dtype=torch.float16, device="npu")
119+ kwargs_tensor = functools.partial(torch.empty, 2, dtype=torch.float16, device="npu")
120+ 
121+ torch.npu.npugraph_ex.register_replacement(
122+ search_fn=search_fn,
123+ replace_fn=replace_fn,
124+ example_inputs=(input_tensor(), input_tensor(), kwargs_tensor()),
125+ extra_check=extra_check
126+ )
127+ 
128+ class Model(torch.nn.Module):
129+ def __init__(self):
130+ super(Model, self).__init__()
131+ 
132+ def forward(self, data1, data2, gamma):
133+ x_out = torch.add(data1, data2)
134+ y, _ = torch_npu.npu_rms_norm(x_out, gamma)
135+ 
136+ abs_01 = torch.abs(y)
137+ sqrt_01 = torch.sqrt(x_out)
138+ return abs_01, sqrt_01
139+ 
140+ model = Model().npu()
141+ 
142+ x1 = torch.randn(1, 1, 7168, dtype=torch.float16, device='npu')
143+ x2 = torch.randn(1, 1, 7168, dtype=torch.float16, device='npu')
144+ gamma = torch.ones(7168, dtype=torch.float16, device='npu')
145+ 
146+ model_compile = torch.compile(model, backend="npugraph_ex", fullgraph=True, dynamic=False)
147+ res = model_compile(x1, x2, gamma)
148+ self.assertEqual(res[0].shape[2], 7168)
149+ self.assertEqual(res[1].shape[2], 7168)
150+ 
151+ def test_multi_stream(self):
152+ class Model(torch.nn.Module):
153+ def __init__(self):
154+ super().__init__()
155+ self.tagged_event1 = torch.npu.npugraph_ex.ops.npu_create_tagged_event(tag="66")
156+ self.tagged_event2 = torch.npu.npugraph_ex.ops.npu_create_tagged_event(tag="77")
157+ 
158+ def forward(self, in1, in2, in3, in4):
159+ add_result = torch.add(in1, in2)
160+ torch.npu.npugraph_ex.ops.npu_tagged_event_record(self.tagged_event1)
161+ with torch.npu.npugraph_ex.scope.npu_stream_switch('1'):
162+ torch.npu.npugraph_ex.ops.npu_tagged_event_wait(self.tagged_event1)
163+ mm_result = torch.mm(in3, in4)
164+ torch.npu.npugraph_ex.ops.npu_tagged_event_record(self.tagged_event2)
165+ B = in3 + in4
166+ torch.npu.npugraph_ex.ops.npu_record_tagged_stream(B, '1')
167+ mm1 = torch.mm(in3, in4)
168+ del B
169+ C = torch.ones(1000, 1000, dtype=torch.float16, device="npu")
170+ C.add_(2)
171+ with torch.npu.npugraph_ex.scope.npu_stream_switch('2'):
172+ torch.npu.npugraph_ex.ops.npu_tagged_event_wait(self.tagged_event2)
173+ add2 = torch.add(in3, in4)
174+ return add_result, mm_result, mm1, add2, C
175+ 
176+ model = Model().npu()
177+ model = torch.compile(model, backend="npugraph_ex", fullgraph=True, dynamic=False)
178+ in1 = torch.randn(1000, 1000, dtype=torch.float16, device="npu")
179+ in2 = torch.randn(1000, 1000, dtype=torch.float16, device="npu")
180+ in3 = torch.randn(1000, 1000, dtype=torch.float16, device="npu")
181+ in4 = torch.randn(1000, 1000, dtype=torch.float16, device="npu")
182+ res = model(in1, in2, in3, in4)
183+ self.assertEqual(res[0].shape[0], 1000)
184+ self.assertEqual(res[1].shape[0], 1000)
185+ self.assertEqual(res[2].shape[0], 1000)
186+ self.assertEqual(res[3].shape[0], 1000)
187+ self.assertEqual(res[4].shape[0], 1000)
188+ 
189+ def test_record_and_wait(self):
190+ def demo(x):
191+ with torch.npu.npugraph_ex.scope.npu_stream_switch('1'):
192+ mm = torch.mm(x, x)
193+ res = torch.abs(mm)
194+ record = torch.npu.npugraph_ex.ops.record()
195+ add = torch.add(res, 1)
196+ torch.npu.npugraph_ex.ops.wait([record])
197+ sub = torch.sub(x, mm)
198+ return add, sub
199+ 
200+ func = torch.compile(demo, backend="npugraph_ex", fullgraph=True, dynamic=False)
201+ input1 = torch.ones(2, 2).npu()
202+ res = func(input1)
203+ self.assertTrue(same(res[0], torch.empty(2, 2).fill_(2).npu()))
204+ self.assertTrue(same(res[1], torch.empty(2, 2).fill_(-1).npu()))
205+ 
206+ 
207+if __name__ == "__main__":
208+ from torch._dynamo.test_case import run_tests
209+ 
210+ run_tests()
Mtest/torch_npu_schema.json+9-3
@@ -2255,12 +2255,12 @@
2255 "torch_npu.npu.npugraph_ex.ops.npu_tagged_event_record": {2255 "torch_npu.npu.npugraph_ex.ops.npu_tagged_event_record": {
2256 "signature": "(event)"2256 "signature": "(event)"
2257 },2257 },
2258- "torch_npu.npu.npugraph_ex.ops.npu_print": {
2259- "signature": "(*args, summarize_size=3, tensor_detail=False)"
2260- },
2261 "torch_npu.npu.npugraph_ex.ops.npu_tagged_event_wait": {2258 "torch_npu.npu.npugraph_ex.ops.npu_tagged_event_wait": {
2262 "signature": "(event)"2259 "signature": "(event)"
2263 },2260 },
2261+ "torch_npu.npu.npugraph_ex.ops.record": {
2262+ "signature": ""
2263+ },
2264 "torch_npu.npu.npugraph_ex.ops.wait": {2264 "torch_npu.npu.npugraph_ex.ops.wait": {
2265 "signature": "(tensors: list)"2265 "signature": "(tensors: list)"
2266 },2266 },
@@ -2273,6 +2273,12 @@
2273 "torch_npu.npu.npugraph_ex.scope.limit_core_num": {2273 "torch_npu.npu.npugraph_ex.scope.limit_core_num": {
2274 "signature": "(op_aicore_num: int, op_vectorcore_num: int)"2274 "signature": "(op_aicore_num: int, op_vectorcore_num: int)"
2275 },2275 },
2276+ "torch_npu.npu.npugraph_ex.compile_fx": {
2277+ "signature": "(options: dict = None)"
2278+ },
2279+ "torch_npu.npu.npugraph_ex.register_replacement": {
2280+ "signature": "(search_fn, replace_fn, example_inputs, trace_fn=<function fwd_only>, extra_check=<function _return_true>, search_fn_pattern=None)"
2281+ },
2276 "torch_npu.distributed.run.parse_args": {2282 "torch_npu.distributed.run.parse_args": {
2277 "signature": "(args)"2283 "signature": "(args)"
2278 },2284 },
Mtorch_npu/npu/npugraph_ex/__init__.py+32-1
@@ -1,4 +1,35 @@
1+__all__ = ["compile_fx", "register_replacement"]
2+ 
3+from typing import Match
4+ 
5+try:
6+ from torch._inductor.pattern_matcher import fwd_only
7+except ImportError:
8+ from torch._inductor.pattern_matcher import inference_graph as fwd_only
9+ 
1from . import experimental10from . import experimental
2from . import inference11from . import inference
3from . import ops12from . import ops
4-from . import scope13+from . import scope
14+ 
15+ 
16+def compile_fx(options: dict = None):
17+ import torchair
18+ from torchair.configs import npugraphex_config
19+ 
20+ compiler_config = torchair.CompilerConfig()
21+ compiler_config.mode = "npugraph_ex"
22+ npugraphex_config._process_kwargs_options(compiler_config, {options: options})
23+ return torchair.get_compiler(npugraphex_config)
24+ 
25+ 
26+def _return_true(match: Match):
27+ return True
28+ 
29+ 
30+def register_replacement(search_fn, replace_fn, example_inputs, trace_fn=fwd_only, extra_check=_return_true,
31+ search_fn_pattern=None):
32+ import torchair
33+ return torchair.patterns.pattern_pass_manager.register_replacement(search_fn, replace_fn, example_inputs,
34+ trace_fn=trace_fn, extra_check=extra_check,
35+ search_fn_pattern=search_fn_pattern)
Mtorch_npu/npu/npugraph_ex/ops/__init__.py+5-5
@@ -1,11 +1,6 @@
1import torch1import torch
2 2 
3 3 
4-def npu_print(*args, summarize_size=3, tensor_detail=False):
5- from torch_npu.dynamo.torchair import ops
6- return ops.npu_print(*args, summarize_size=summarize_size, tensor_detail=tensor_detail)
7- 
8- 
9def npu_create_tagged_event(tag: str):4def npu_create_tagged_event(tag: str):
10 from torch_npu.dynamo.torchair import ops5 from torch_npu.dynamo.torchair import ops
11 return ops.npu_create_tagged_event(tag=tag)6 return ops.npu_create_tagged_event(tag=tag)
@@ -26,6 +21,11 @@ def npu_record_tagged_stream(input: torch.Tensor, tagged_stream: str):
26 return ops.npu_record_tagged_stream(input=input, tagged_stream=tagged_stream)21 return ops.npu_record_tagged_stream(input=input, tagged_stream=tagged_stream)
27 22 
28 23 
24+def record():
25+ from torch_npu.dynamo.torchair import ops
26+ return ops.record()
27+ 
28+ 
29def wait(tensors: list):29def wait(tensors: list):
30 from torch_npu.dynamo.torchair import ops30 from torch_npu.dynamo.torchair import ops
31 return ops.wait(tensors=tensors)31 return ops.wait(tensors=tensors)