AST-based graph variable renamer

Author: TelGome

graph_net/test_compiler_util.py

@@ -144,7 +144,7 @@ def print_with_log_prompt(key, value, log_prompt):
 def print_basic_config(args, hardware_name, compile_framework_version):
     model_path = os.path.normpath(args.model_path)
     model_name = get_model_name(model_path)
-    print_with_log_prompt("[Config] model:", model_name, args.log_prompt)
+    print_with_log_prompt("[Processing] model:", model_name, args.log_prompt)
 
     print_with_log_prompt("[Config] device:", args.device, args.log_prompt)
     print_with_log_prompt("[Config] hardware:", hardware_name, args.log_prompt)

graph_net/torch/graph_variable_renamer.py

@@ -1,17 +1,116 @@
 import os
-import torch
 import shutil
 import tempfile
-
-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
-from graph_net.tensor_meta import TensorMeta
+import ast
+import inspect
+import torch
 from pathlib import Path
-from graph_net.torch.utils import apply_templates
 from graph_net.imp_util import load_module
+from graph_net.tensor_meta import TensorMeta
 from graph_net.hash_util import get_sha256_hash
 
 
+class AstGraphRenamer(ast.NodeTransformer):
+    def __init__(self, rename_map, input_arg_names, weight_arg_names):
+        self.rename_map = rename_map
+        self.input_arg_names = set(input_arg_names)
+        self.weight_arg_names = set(weight_arg_names)
+        self.counters = {"tmp": 0}
+        self.in_forward = False
+
+    def visit_FunctionDef(self, node):
+        if node.name == "forward":
+            self.in_forward = True
+            new_args = []
+            for arg in node.args.args:
+                if arg.arg == "self":
+                    new_args.append(arg)
+                    continue
+
+                if arg.arg in self.rename_map:
+                    new_arg_name = self.rename_map[arg.arg]
+                    new_args.append(
+                        ast.arg(arg=new_arg_name, annotation=arg.annotation)
+                    )
+                else:
+                    new_args.append(arg)
+
+            node.args.args = new_args
+
+            new_body = []
+            for stmt in node.body:
+                stmt = self._remove_clear_stmt_of_args(stmt)
+                if stmt is None:
+                    continue
+
+                stmt = self.visit(stmt)
+                new_body.append(stmt)
+
+            node.body = new_body
+            self.in_forward = False
+        return node
+
+    def visit_Assign(self, node):
+        if not self.in_forward:
+            return node
+
+        for target in node.targets:
+            if isinstance(target, ast.Name):
+                old_name = target.id
+                if old_name not in self.rename_map:
+                    new_name = f"tmp_{self.counters['tmp']}"
+                    self.counters["tmp"] += 1
+                    self.rename_map[old_name] = new_name
+
+        self.generic_visit(node)
+        return node
+
+    def visit_Name(self, node):
+        if not self.in_forward:
+            return node
+        if node.id in self.rename_map:
+            return ast.Name(id=self.rename_map[node.id], ctx=node.ctx)
+        return node
+
+    def _remove_clear_stmt_of_args(self, stmt):
+        args_names = self.input_arg_names | self.weight_arg_names
+
+        def _need_remove(target):
+            return isinstance(target, ast.Name) and target.id in args_names
+
+        if (
+            isinstance(stmt, ast.Assign)
+            and isinstance(stmt.value, ast.Constant)
+            and stmt.value.value is None
+        ):
+            # remove stmt like w_0 = None
+            new_targets = [t for t in stmt.targets if not _need_remove(t)]
+            if not new_targets:
+                return None
+            stmt.targets = new_targets
+        elif isinstance(stmt, ast.Delete):
+            # remove stmt like del w_0
+            new_targets = []
+            for t in stmt.targets:
+                if isinstance(t, ast.Tuple):
+                    kept = [e for e in t.elts if not _need_remove(e)]
+                    if kept:
+                        new_targets.append(ast.Tuple(elts=kept, ctx=ast.Del()))
+                elif not _need_remove(t):
+                    new_targets.append(t)
+            if not new_targets:
+                return None
+            stmt.targets = new_targets
+        return stmt
+
+
+def load_class_from_file(file_path: str, class_name: str):
+    print(f"Load {class_name} from {file_path}")
+    module = load_module(file_path, "unnamed_graph_module")
+    model_class = getattr(module, class_name, None)
+    return model_class
+
+
 class GraphVariableRenamer:
     """
     Used by graph_net.model_path_handler
@@ -44,18 +143,13 @@ def _make_config(
         output_dir="./tmp/graph_variable_renamer_dir",
         filter_path=None,
         filter_config=None,
-        post_extract_process_path=None,
-        post_extract_process_class_name=None,
-        post_extract_process_config=None,
         data_input_predicator_class_name="DataInputPredicator",
         model_runnable_predicator_class_name="ModelRunner",
         data_input_predicator_config=None,
         model_runnable_predicator_config=None,
         model_path_prefix="",
         **kwargs,
     ):
-        if post_extract_process_config is None:
-            post_extract_process_config = {}
         if data_input_predicator_config is None:
             data_input_predicator_config = {}
         if model_runnable_predicator_config is None:
@@ -65,9 +159,6 @@ def _make_config(
             "output_dir": output_dir,
             "filter_path": filter_path,
             "filter_config": filter_config if filter_config is not None else {},
-            "post_extract_process_path": post_extract_process_path,
-            "post_extract_process_class_name": post_extract_process_class_name,
-            "post_extract_process_config": post_extract_process_config,
             "data_input_predicator_filepath": data_input_predicator_filepath,
             "data_input_predicator_class_name": data_input_predicator_class_name,
             "data_input_predicator_config": data_input_predicator_config,
@@ -89,133 +180,94 @@ def __call__(self, rel_model_path):
             return
 
         src_model_path = os.path.join(self.config["model_path_prefix"], rel_model_path)
-        module, inputs = get_torch_module_and_inputs(src_model_path)
-        gm = parse_sole_graph_module(module, inputs)
-        gm, rename_map = self.rename_graph_variables(gm, inputs, src_model_path)
-
         Path(dst_model_path).parent.mkdir(parents=True, exist_ok=True)
+        graph_module = load_class_from_file(
+            os.path.join(src_model_path, "model.py"), class_name="GraphModule"
+        )
+        input_arg_names, weight_arg_names = self._get_input_and_weight_arg_names(
+            graph_module, src_model_path
+        )
+
+        rename_map = {}
+        for idx, name in enumerate(input_arg_names):
+            rename_map[name] = f"in_{idx}"
+        for idx, name in enumerate(weight_arg_names):
+            rename_map[name] = f"w_{idx}"
+
         with tempfile.TemporaryDirectory(prefix="graph_variable_renamer_") as temp_dir:
             temp_model_path = os.path.join(temp_dir, os.path.basename(dst_model_path))
             shutil.copytree(src_model_path, temp_model_path, dirs_exist_ok=True)
-            self._update_model_py_file(gm, temp_model_path)
-            self._update_weight_meta_py_file(
-                src_model_path, temp_model_path, rename_map
+            self._update_model_py_file(
+                temp_model_path, rename_map, input_arg_names, weight_arg_names
             )
-            self._update_input_meta_py_file(src_model_path, temp_model_path, rename_map)
-            # print("Try to run renamed model...")
-            # self._try_run(temp_model_path)
-            shutil.copytree(temp_model_path, dst_model_path)
+            self._update_meta_file(temp_model_path, "weight_meta.py", rename_map)
+            self._update_meta_file(temp_model_path, "input_meta.py", rename_map)
+            print(f"Verifying {rel_model_path}...")
+            self._try_run(temp_model_path)
+            shutil.copytree(temp_model_path, dst_model_path, dirs_exist_ok=True)
 
-    def _try_run(self, model_path):
-        assert self.model_runnable_predicator(
-            model_path
-        ), f"{model_path} is not a runnable model"
+    def _get_input_and_weight_arg_names(self, graph_module, model_path):
+        input_arg_names = []
+        weight_arg_names = []
+        sig = inspect.signature(graph_module.forward)
+        for name, param in sig.parameters.items():
+            if name == "self":
+                continue
+            is_not_data_input = not self.data_input_predicator(model_path, name)
+            is_parameter_type = self._is_parameter_type(param.annotation)
+            if is_not_data_input or is_parameter_type:
+                weight_arg_names.append(name)
+            else:
+                input_arg_names.append(name)
+        return input_arg_names, weight_arg_names
 
-    def _update_model_py_file(self, graph_module, model_path):
-        py_code = apply_templates(graph_module.code)
-        (Path(model_path) / "model.py").write_text(py_code)
-        file_hash = get_sha256_hash(py_code)
-        (Path(model_path) / "graph_hash.txt").write_text(file_hash)
+    def _is_parameter_type(self, annotation):
+        return annotation is torch.nn.parameter.Parameter
 
-    def _update_weight_meta_py_file(self, src_model_path, dst_model_path, rename_map):
-        tensor_metas = TensorMeta.unserialize_from_py_file(
-            os.path.join(src_model_path, "weight_meta.py"),
-        )
-        for weight_meta in tensor_metas:
-            meta_name = self._find_name_in_rename_map(weight_meta.name, rename_map)
-            assert (
-                meta_name is not None
-            ), f"{weight_meta.name} is not found in rename_map"
-            if weight_meta.original_name is None:
-                weight_meta.original_name = weight_meta.name
-            weight_meta.name = rename_map[meta_name]
-
-        py_code = "\n\n".join(
-            [weight_meta.serialize_to_py_str() for weight_meta in tensor_metas]
-        )
-        (Path(dst_model_path) / "weight_meta.py").write_text(py_code)
+    def _update_model_py_file(
+        self, model_path, rename_map, input_arg_names, weight_arg_names
+    ):
+        model_file = Path(model_path) / "model.py"
+        with open(model_file, "r", encoding="utf-8") as f:
+            source = f.read()
 
-    def _update_input_meta_py_file(self, src_model_path, dst_model_path, rename_map):
-        tensor_metas = TensorMeta.unserialize_from_py_file(
-            os.path.join(src_model_path, "input_meta.py"),
-        )
-        for input_meta in tensor_metas:
-            meta_name = self._find_name_in_rename_map(input_meta.name, rename_map)
-            assert (
-                meta_name is not None
-            ), f"{input_meta.name} is not found in rename_map"
-            if input_meta.original_name is None:
-                input_meta.original_name = input_meta.name
-            input_meta.name = rename_map[meta_name]
-
-        py_code = "\n\n".join(
-            [input_meta.serialize_to_py_str() for input_meta in tensor_metas]
-        )
-        (Path(dst_model_path) / "input_meta.py").write_text(py_code)
-
-    def _find_name_in_rename_map(self, raw_name, rename_map):
-        if raw_name in rename_map:
-            return raw_name
-        # s1 -> s1_
-        elif (raw_name + "_") in rename_map:
-            return raw_name + "_"
+        tree = ast.parse(source)
+        for node in tree.body:
+            if isinstance(node, ast.ClassDef) and node.name == "GraphModule":
+                transformer = AstGraphRenamer(
+                    rename_map, input_arg_names, weight_arg_names
+                )
+                transformer.visit(node)
+                break
+
+        if hasattr(ast, "unparse"):
+            py_code = ast.unparse(tree)
         else:
-            return None
+            import astor
 
-    def _get_model(self, model_path):
-        py_module = load_module(os.path.join(model_path, "model.py"))
-        GraphModule = getattr(py_module, "GraphModule")
-        GraphModule.__graph_net_file_path__ = py_module.__graph_net_file_path__
-        return GraphModule()
+            py_code = astor.to_source(tree)
 
-    def rename_graph_variables(
-        self, gm: torch.fx.GraphModule, sample_inputs, model_path
-    ):
-        counters = {"in": 0, "w": 0, "tmp": 0}
-        rename_map = {}
-        # graph may not have input, only contain weights
-        arg_iter = iter(sample_inputs) if sample_inputs else iter([])
-        for node in gm.graph.nodes:
-            self._process_single_node(node, arg_iter, counters, model_path, rename_map)
-        gm.graph.lint()
-        gm.recompile()
-        return gm, rename_map
-
-    def _process_single_node(self, node, arg_iter, counters, model_path, rename_map):
-        if "original_name" not in node.meta:
-            node.meta["original_name"] = node.name
-        if node.op == "placeholder":
-            self._handle_placeholder(node, arg_iter, counters, model_path, rename_map)
-        elif node.op == "get_attr":
-            self._apply_rename(node, "w", counters, rename_map)
-        elif node.op != "output":
-            self._apply_rename(node, "tmp", counters, rename_map)
-        else:
-            # Do nothing
-            pass
-
-    def _handle_placeholder(self, node, arg_iter, counters, model_path, rename_map):
-        real_arg = next(arg_iter, None)
-        is_weight = self._is_weight_node(node, real_arg, model_path)
-        prefix = "w" if is_weight else "in"
-        self._apply_rename(node, prefix, counters, rename_map, update_target=True)
-
-    def _apply_rename(self, node, prefix, counters, rename_map, update_target=False):
-        old_name = node.name
-        new_name = f"{prefix}_{counters[prefix]}"
-        counters[prefix] += 1
-        node.name = new_name
-        if update_target:
-            node.target = new_name
-
-        rename_map[old_name] = new_name
-
-    def _is_weight_node(self, node, real_arg, model_path):
-        is_not_data_input = not self.data_input_predicator(model_path, node.name)
-        is_parameter_type = (
-            node.type is not None
-            and isinstance(node.type, type)
-            and issubclass(node.type, torch.nn.parameter.Parameter)
-        )
-        is_parameter_value = isinstance(real_arg, torch.nn.Parameter)
-        return is_not_data_input or is_parameter_type or is_parameter_value
+        model_file.write_text(py_code, encoding="utf-8")
+        file_hash = get_sha256_hash(py_code)
+        (Path(model_path) / "graph_hash.txt").write_text(file_hash)
+
+    def _update_meta_file(self, model_path, meta_filename, rename_map):
+        meta_file = Path(model_path) / meta_filename
+        tensor_metas = TensorMeta.unserialize_from_py_file(str(meta_file))
+        for meta in tensor_metas:
+            if meta.name in rename_map:
+                if meta.original_name is None:
+                    meta.original_name = meta.name
+                meta.name = rename_map[meta.name]
+            else:
+                print(
+                    f"[Warning] {meta.name} in {meta_filename} not found in rename_map, skipping."
+                )
+
+        py_code = "\n\n".join([meta.serialize_to_py_str() for meta in tensor_metas])
+        meta_file.write_text(py_code)
+
+    def _try_run(self, model_path):
+        assert self.model_runnable_predicator(
+            model_path
+        ), f"{model_path} is not a runnable model"