From 130c79957bf104a24f8f82c4f43deef7ea0f174d Mon Sep 17 00:00:00 2001 From: Honglei-Qiu <1044497581@qq.com> Date: Fri, 9 Jan 2026 04:32:02 +0000 Subject: [PATCH 1/2] Implement graph_net/torch/op_names_extractor.py as a subclass of SamplePass --- .../test/typical_sequence_decomposer_test.sh | 2 +- graph_net/tools/typical_sequence_decompose.sh | 3 +- .../{ => sample_pass}/op_names_extractor.py | 97 +++++++++---------- 3 files changed, 51 insertions(+), 51 deletions(-) rename graph_net/torch/{ => sample_pass}/op_names_extractor.py (90%) diff --git a/graph_net/test/typical_sequence_decomposer_test.sh b/graph_net/test/typical_sequence_decomposer_test.sh index ca327ba31..c28ec1334 100644 --- a/graph_net/test/typical_sequence_decomposer_test.sh +++ b/graph_net/test/typical_sequence_decomposer_test.sh @@ -14,7 +14,7 @@ python3 -m graph_net.model_path_handler \ --model-path-list $model_list \ --handler-config=$(base64 -w 0 < List[str]: + extractor = TypicalSequenceExtractor() + model, inputs = get_torch_module_and_inputs(model_path) + compiled_model, _ = parse_sole_graph_module_without_varify(model, inputs) + extractor.extract_compiler(compiled_model, inputs) + ops_info = extractor.extract_node + + return [op["target_name"] for op in ops_info] + class TypicalSequenceExtractor: def __init__(self): @@ -52,46 +94,3 @@ def extract_compiler(self, gm: torch.fx.GraphModule, inputs: List[torch.Tensor]) operator = self._extract_operators_from_graph(gm, inputs) self.extract_node = operator return gm.forward - - - -class OpNamesExtractor: - def __init__(self, config=None): - if config is None: - config = {} - - self.config = self._make_config(**config) - - def _make_config( - self, model_path_prefix: str, output_dir: str, resume: bool = False - ): - return { - "model_path_prefix": model_path_prefix, - "resume": resume, - "output_dir": output_dir, - } - - def __call__(self, rel_model_path: str): - torch.cuda.empty_cache() - model_path = os.path.join(self.config["model_path_prefix"], rel_model_path) - output_path = self._get_output_path(rel_model_path) - if self.config["resume"] and output_path.exists(): - return - op_names = self._extract_ops(model_path) - output_path.write_text("\n".join(op_names)) - print(f"Save op-names to {str(output_path)}") - - def _get_output_path(self, rel_model_path: str): - output_path_dir = Path(self.config["output_dir"]) / rel_model_path - output_path_dir.mkdir(parents=True, exist_ok=True) - output_path = output_path_dir / "op_names.txt" - return output_path - - def _extract_ops(self, model_path: str) -> List[str]: - extractor = TypicalSequenceExtractor() - model, inputs = get_torch_module_and_inputs(model_path) - compiled_model, _ = parse_sole_graph_module_without_varify(model, inputs) - extractor.extract_compiler(compiled_model, inputs) - ops_info = extractor.extract_node - - return [op["target_name"] for op in ops_info] From a2c85e34ab75e2f341d48b947c06e1a6e32a1a2a Mon Sep 17 00:00:00 2001 From: Honglei-Qiu <1044497581@qq.com> Date: Fri, 9 Jan 2026 11:02:45 +0000 Subject: [PATCH 2/2] [Feature Enhancement] Add resumable processing support to OpNamesExtractor using ResumableSamplePassMixin --- .../torch/sample_pass/op_names_extractor.py | 23 ++++++++++++++----- 1 file changed, 17 insertions(+), 6 deletions(-) mode change 100644 => 100755 graph_net/torch/sample_pass/op_names_extractor.py diff --git a/graph_net/torch/sample_pass/op_names_extractor.py b/graph_net/torch/sample_pass/op_names_extractor.py old mode 100644 new mode 100755 index f9a5421a5..28608160e --- a/graph_net/torch/sample_pass/op_names_extractor.py +++ b/graph_net/torch/sample_pass/op_names_extractor.py @@ -4,12 +4,13 @@ import torch import torch.nn as nn +from graph_net.sample_pass.resumable_sample_pass_mixin import ResumableSamplePassMixin from graph_net.torch.fx_graph_module_util import get_torch_module_and_inputs from graph_net.torch.fx_graph_parse_util import parse_sole_graph_module_without_varify from graph_net.sample_pass.sample_pass import SamplePass -class OpNamesExtractor(SamplePass): +class OpNamesExtractor(SamplePass, ResumableSamplePassMixin): def __init__(self, config=None): super().__init__(config) if config is None: @@ -18,20 +19,30 @@ def __init__(self, config=None): self.config = self._make_config(**config) def declare_config( - self, model_path_prefix: str, output_dir: str, resume: bool = False + self, + model_path_prefix: str, + output_dir: str, + resume: bool = False, + limits_handled_models: int = None, ): pass - def __call__(self, rel_model_path: str): + def sample_handled(self, rel_model_path: str) -> bool: + return self.naive_sample_handled( + rel_model_path, search_file_name="op_names.txt" + ) + + def resume(self, rel_model_path: str): torch.cuda.empty_cache() model_path = os.path.join(self.config["model_path_prefix"], rel_model_path) - output_path = self._get_output_path(rel_model_path) - if self.config["resume"] and output_path.exists(): - return op_names = self._extract_ops(model_path) + output_path = self._get_output_path(rel_model_path) output_path.write_text("\n".join(op_names)) print(f"Save op-names to {str(output_path)}") + def __call__(self, rel_model_path: str): + self.resumable_handle_sample(rel_model_path) + def _make_config( self, model_path_prefix: str, output_dir: str, resume: bool = False ):