已合并
修改npugraph_ex.compile_fx接口签名 #32080
mihudan创建于 3月19日
修改npugraph_ex.compile_fx接口签名 #32080
已合并
共 3 个文件变更+18-9
| @@ -41,14 +41,22 @@ class TestNpuGraphEx(TestCase): | |||
| 41 | def forward(self, x, y): | 41 | def forward(self, x, y): |
| 42 | return x + y | 42 | return x + y |
| 43 | 43 | ||
| 44 | + def custom_compiler(gm: torch.fx.GraphModule, example_inputs): | ||
| 45 | + compiled_graph = torch.npu.npugraph_ex.compile_fx(gm, example_inputs) | ||
| 46 | + return compiled_graph | ||
| 47 | + | ||
| 48 | + def custom_compiler_with_options(gm: torch.fx.GraphModule, example_inputs): | ||
| 49 | + test_kwargs = { | ||
| 50 | + "clone_input": False | ||
| 51 | + } | ||
| 52 | + compiled_graph = torch.npu.npugraph_ex.compile_fx(gm, example_inputs, test_kwargs) | ||
| 53 | + return compiled_graph | ||
| 54 | + | ||
| 44 | def my_backend(gm: torch.fx.GraphModule, example_inputs): | 55 | def my_backend(gm: torch.fx.GraphModule, example_inputs): |
| 45 | - compiler = torch.npu.npugraph_ex.compile_fx() | 56 | + return aot_module_simplified(gm, example_inputs, fw_compiler=custom_compiler) |
| 46 | - return aot_module_simplified(gm, example_inputs, fw_compiler=compiler) | ||
| 47 | 57 | ||
| 48 | def my_backend_with_options(gm: torch.fx.GraphModule, example_inputs): | 58 | def my_backend_with_options(gm: torch.fx.GraphModule, example_inputs): |
| 49 | - options = {"clone_input": False} | 59 | + return aot_module_simplified(gm, example_inputs, fw_compiler=custom_compiler_with_options) |
| 50 | - compiler = torch.npu.npugraph_ex.compile_fx(options=options) | ||
| 51 | - return aot_module_simplified(gm, example_inputs, fw_compiler=compiler) | ||
| 52 | 60 | ||
| 53 | model = Model().npu() | 61 | model = Model().npu() |
| 54 | compiled_model = torch.compile(model, backend=my_backend, fullgraph=True, dynamic=False) | 62 | compiled_model = torch.compile(model, backend=my_backend, fullgraph=True, dynamic=False) |
| @@ -2256,7 +2256,7 @@ | |||
| 2256 | "signature": "(op_aicore_num: int, op_vectorcore_num: int)" | 2256 | "signature": "(op_aicore_num: int, op_vectorcore_num: int)" |
| 2257 | }, | 2257 | }, |
| 2258 | "torch_npu.npu.npugraph_ex.compile_fx": { | 2258 | "torch_npu.npu.npugraph_ex.compile_fx": { |
| 2259 | - "signature": "(options: dict = None)" | 2259 | + "signature": "(gm, example_inputs=None, options=None)" |
| 2260 | }, | 2260 | }, |
| 2261 | "torch_npu.npu.npugraph_ex.register_replacement": { | 2261 | "torch_npu.npu.npugraph_ex.register_replacement": { |
| 2262 | "signature": "(search_fn: torch._inductor.pattern_matcher.SearchFn, replace_fn: torch._inductor.pattern_matcher.ReplaceFn, example_inputs: collections.abc.Iterable[typing.Any], trace_fn: torch._inductor.pattern_matcher.TraceFn = <function fwd_only>, extra_check: Callable[[Match], bool] = <function _return_true>, search_fn_pattern: Optional[torch._inductor.pattern_matcher.PatternExpr] = None, scalar_workaround: Optional[dict[str, Union[float, int]]] = None, skip_duplicates: bool = False)" | 2262 | "signature": "(search_fn: torch._inductor.pattern_matcher.SearchFn, replace_fn: torch._inductor.pattern_matcher.ReplaceFn, example_inputs: collections.abc.Iterable[typing.Any], trace_fn: torch._inductor.pattern_matcher.TraceFn = <function fwd_only>, extra_check: Callable[[Match], bool] = <function _return_true>, search_fn_pattern: Optional[torch._inductor.pattern_matcher.PatternExpr] = None, scalar_workaround: Optional[dict[str, Union[float, int]]] = None, skip_duplicates: bool = False)" |
| @@ -1,7 +1,7 @@ | |||
| 1 | __all__ = ["compile_fx", "register_replacement"] | 1 | __all__ = ["compile_fx", "register_replacement"] |
| 2 | 2 | ||
| 3 | from collections.abc import Collection, Generator, Iterable, Mapping, Sequence | 3 | from collections.abc import Collection, Generator, Iterable, Mapping, Sequence |
| 4 | -from typing import Any, Callable, NoReturn, Optional, Protocol, TypeVar, Union, Match | 4 | +from typing import Any, Callable, NoReturn, Optional, Protocol, TypeVar, Union, Match, List |
| 5 | 5 | ||
| 6 | try: | 6 | try: |
| 7 | from torch._inductor.pattern_matcher import fwd_only, SearchFn, ReplaceFn, TraceFn, PatternExpr | 7 | from torch._inductor.pattern_matcher import fwd_only, SearchFn, ReplaceFn, TraceFn, PatternExpr |
| @@ -12,14 +12,15 @@ from . import inference | |||
| 12 | from . import scope | 12 | from . import scope |
| 13 | 13 | ||
| 14 | 14 | ||
| 15 | -def compile_fx(options: dict = None): | 15 | +def compile_fx(gm, example_inputs=None, options=None): |
| 16 | import npugraph_ex | 16 | import npugraph_ex |
| 17 | from npugraph_ex.configs import npugraphex_config | 17 | from npugraph_ex.configs import npugraphex_config |
| 18 | 18 | ||
| 19 | compiler_config = npugraph_ex.CompilerConfig() | 19 | compiler_config = npugraph_ex.CompilerConfig() |
| 20 | compiler_config.mode = "npugraph_ex" | 20 | compiler_config.mode = "npugraph_ex" |
| 21 | npugraphex_config._process_kwargs_options(compiler_config, {"options": {} if options is None else options}) | 21 | npugraphex_config._process_kwargs_options(compiler_config, {"options": {} if options is None else options}) |
| 22 | - return npugraph_ex.get_compiler(compiler_config) | 22 | + compiler = npugraph_ex.get_compiler(compiler_config) |
| 23 | + return compiler(gm, example_inputs) | ||
| 23 | 24 | ||
| 24 | 25 | ||
| 25 | def _return_true(match: Match): | 26 | def _return_true(match: Match): |