Naive decompose minor fix (#344)

* support checking model redundancy

* revert change of vision_model_test

* reformat python code.

* reformat bert_model_test.py and utils.py

* minor fix

* fix failed check by comparing directories after os.path.realpath()

* fix bugs in check_validate.sh

* set dynamic=False in single_device_runner.py

* reset graph hash

* minor fix for naive_graph_decomposer

Author: lixinqi

graph_net/test/chain_naive_graph_decomposer_test.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+set -x
+
+# input model path
+MODEL_PATH_IN_SAMPLES=/timm/resnet18 
+# extract subgraph 0-8, 8-16
+read -r -d '' json_str <<'EOF'
+{
+    "output_dir": "/tmp/naive_decompose_workspace",
+    "split_positions": [2, 4],
+    "group_head_and_tail": false,
+    "chain_style": true
+}
+EOF
+CONFIG=$(echo $json_str | base64 -w 0) 
+
+mkdir -p /tmp/naive_decompose_workspace
+GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
+os.path.dirname(graph_net.__file__))")
+python3 -m graph_net.torch.single_device_runner --model-path $GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES --enable-extract True --extract-name resnet18 --dump-graph-hash-key --custom-extractor-path=$GRAPH_NET_ROOT/torch/naive_graph_decomposer.py --custom-extractor-config=$CONFIG

graph_net/test/naive_graph_decomposer_test.sh

@@ -1,13 +1,17 @@
 #!/bin/bash
+
 # input model path
 MODEL_PATH_IN_SAMPLES=/timm/resnet18 
-# output model path
-OUTPUT_DIR=/tmp/naive_decompose_workspace
+read -r -d '' json_str <<'EOF'
+{
+    "output_dir": "/tmp/naive_decompose_workspace",
+    "split_positions": [8, 32],
+    "group_head_and_tail": true
+}
+EOF
+CONFIG=$(echo $json_str | base64 -w 0) 
 
-mkdir -p $OUTPUT_DIR
-# extract subgraph 0-8, 8-16
-export GRAPH_NET_NAIVE_DECOMPOSER_SPLIT_POS=0,8,16
-export GRAPH_NET_EXTRACT_WORKSPACE=$OUTPUT_DIR 
+mkdir -p /tmp/naive_decompose_workspace
 GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
 os.path.dirname(graph_net.__file__))")
-python3 -m graph_net.torch.single_device_runner --model-path $GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES --enable-extract True --extract-name resnet18 --dump-graph-hash-key --custom-extractor-path=$GRAPH_NET_ROOT/torch/naive_graph_decomposer.py
+python3 -m graph_net.torch.single_device_runner --model-path $GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES --enable-extract True --extract-name resnet18 --dump-graph-hash-key --custom-extractor-path=$GRAPH_NET_ROOT/torch/naive_graph_decomposer.py --custom-extractor-config=$CONFIG

graph_net/torch/decompose_util.py

@@ -9,9 +9,13 @@ def convert_to_submodules_graph(
     original_gm: torch.fx.GraphModule,
     split_positions: list[int],
     submodule_hook=None,
-    submodule_name_prefix="extraced_submodule",
+    submodule_name_prefix="extracted_submodule",
+    chain_style=False,
     group_head_and_tail=True,
 ):
+    """
+    chain_style=True: decompose original_gm into g0 * g1 * g2 * g3
+    """
     original_gm = copy.deepcopy(original_gm)
     num_placeholders = len(
         [node for node in original_gm.graph.nodes if node.op == "placeholder"]
@@ -68,6 +72,12 @@ def get_end_node_idx(range_idx):
                 return i + 1
         raise NotImplementedError("Dead code.")
 
+    def print_submodule_call(prompt, gm):
+        submodule_call_stmts = [
+            stmt for stmt in gm.code.split("\n") if "self.extracted_submodule" in stmt
+        ]
+        print(f"{prompt} ", submodule_call_stmts)
+
     for range_idx in range(len(range_idx2submodule_body_nodes)):
         (
             submodule_input_nodes,
@@ -76,6 +86,7 @@ def get_end_node_idx(range_idx):
             original_gm=original_gm,
             start_node_idx=get_start_node_idx(range_idx),
             end_node_idx=get_end_node_idx(range_idx),
+            chain_style=chain_style,
         )
 
         def get_input_nodes(range_idx):
@@ -136,9 +147,14 @@ def get_output_nodes(range_idx):
         # Erase old nodes
         for node in reversed(get_body_nodes(range_idx)):
             original_gm.graph.erase_node(node)
+        # print_submodule_call("(fx) after Erase old nodes", original_gm)
+
+    # print_submodule_call("(fx) before recompile", original_gm)
 
     original_gm.recompile()
 
+    # print_submodule_call("(fx) after recompile", original_gm)
+
     return original_gm
 
 
@@ -147,7 +163,7 @@ def fold_range_to_submodule(
     start_node_idx: int,
     end_node_idx: int,
     submodule_hook=None,
-    submodule_name="extraced_submodule",
+    submodule_name="extracted_submodule",
     group_head_and_tail=True,
 ):
     return convert_to_submodules_graph(
@@ -170,6 +186,7 @@ def _get_submodule_inputs_and_outputs(
     original_gm: torch.fx.GraphModule,
     start_node_idx: int,
     end_node_idx: int,
+    chain_style=False,
 ):
     count_ctx = NodeProducedOrConsumedCountCtx(
         defaultdict(int),
@@ -179,7 +196,11 @@ def _get_submodule_inputs_and_outputs(
     node_list = list(original_gm.graph.nodes)
 
     def get_related_node(node):
-        yield from node.args
+        for arg in node.args:
+            if isinstance(arg, tuple):
+                yield from arg
+            else:
+                yield arg
         yield node
 
     for node in node_list[0:start_node_idx]:
@@ -194,19 +215,33 @@ def get_related_node(node):
         for related_node in get_related_node(node):
             count_ctx.node2after_output[related_node] += 1
 
-    input_nodes = [
-        node
-        for node in node_list
-        if count_ctx.node2before_input[node] > 0
-        if count_ctx.node2body[node] > 0
-    ]
-
-    output_nodes = [
-        node
-        for node in node_list
-        if not (count_ctx.node2before_input[node] > 0)
-        if count_ctx.node2body[node] > 0
-        if count_ctx.node2after_output[node] > 0
-    ]
+    if chain_style:
+        input_nodes = [
+            node
+            for node in node_list
+            if (count_ctx.node2before_input[node] > 0)
+            if (count_ctx.node2body[node] > 0 or count_ctx.node2after_output[node] > 0)
+        ]
+        input_nodes_set = set(input_nodes)
+        output_nodes = [
+            node
+            for node in node_list
+            if (count_ctx.node2before_input[node] > 0 or count_ctx.node2body[node] > 0)
+            if (count_ctx.node2after_output[node] > 0)
+        ]
+    else:
+        input_nodes = [
+            node
+            for node in node_list
+            if count_ctx.node2before_input[node] > 0
+            if count_ctx.node2body[node] > 0
+        ]
+        output_nodes = [
+            node
+            for node in node_list
+            if not (count_ctx.node2before_input[node] > 0)
+            if count_ctx.node2body[node] > 0
+            if count_ctx.node2after_output[node] > 0
+        ]
 
     return input_nodes, output_nodes

graph_net/torch/extractor.py

@@ -15,14 +15,23 @@
 
 class GraphExtractor:
     def __init__(
-        self, name, dynamic, mut_graph_codes=None, placeholder_auto_rename=False
+        self,
+        name,
+        dynamic,
+        mut_graph_codes=None,
+        placeholder_auto_rename=False,
+        workspace_path=None,
     ):
         self.subgraph_counter = 0
         self.name = name
         self.dynamic = dynamic
         self.mut_graph_codes = mut_graph_codes
         self.placeholder_auto_rename = placeholder_auto_rename
-        self.workspace_path = os.environ.get("GRAPH_NET_EXTRACT_WORKSPACE")
+        self.workspace_path = (
+            workspace_path
+            if workspace_path is not None
+            else os.environ.get("GRAPH_NET_EXTRACT_WORKSPACE")
+        )
         if not self.workspace_path:
             raise EnvironmentError(
                 "Environment variable 'GRAPH_NET_EXTRACT_WORKSPACE' is not set."
@@ -130,7 +139,8 @@ def extract(
     dynamic=True,
     mut_graph_codes=None,
     placeholder_auto_rename=False,
-    custom_extractor_path=None,
+    custom_extractor_path: str = None,
+    custom_extractor_config: str = None,
 ):
     """
     Extract computation graphs from PyTorch nn.Module.
@@ -200,19 +210,20 @@ def forward(self, s0 : torch.SymInt, L_x_ : torch.Tensor):
         >>>
     """
 
-    def get_graph_extractor_cls():
+    def get_graph_extractor_maker():
         if custom_extractor_path is None:
             return GraphExtractor
         import importlib.util as imp
 
         spec = imp.spec_from_file_location("graph_extractor", custom_extractor_path)
         graph_extractor = imp.module_from_spec(spec)
         spec.loader.exec_module(graph_extractor)
-        return graph_extractor.GraphExtractor
+        cls = graph_extractor.GraphExtractor
+        return lambda *args, **kwargs: cls(custom_extractor_config, *args, **kwargs)
 
     def wrapper(model: torch.nn.Module):
         assert isinstance(model, torch.nn.Module), f"{type(model)=}"
-        extractor = get_graph_extractor_cls()(
+        extractor = get_graph_extractor_maker()(
             name, dynamic, mut_graph_codes, placeholder_auto_rename
         )
         # return torch.compile(backend=extractor, dynamic=dynamic)

graph_net/torch/naive_graph_decomposer.py

@@ -1,6 +1,7 @@
 import os
 import torch
 import json
+import base64
 import shutil
 from typing import Union, Callable
 from graph_net.torch import utils
@@ -10,36 +11,62 @@
 
 class GraphExtractor:
     def __init__(
-        self, name, dynamic, mut_graph_codes=None, placeholder_auto_rename=False
+        self,
+        config_str: str,
+        name,
+        dynamic,
+        mut_graph_codes=None,
+        placeholder_auto_rename=False,
     ):
         self.subgraph_counter = 0
         self.name = name
         self.dynamic = dynamic
         self.mut_graph_codes = mut_graph_codes
         self.placeholder_auto_rename = placeholder_auto_rename
-        self.workspace_path = os.environ.get("GRAPH_NET_EXTRACT_WORKSPACE")
-        if not self.workspace_path:
-            raise EnvironmentError(
-                "Environment variable 'GRAPH_NET_EXTRACT_WORKSPACE' is not set."
-            )
-        split_pos_str = os.environ.get("GRAPH_NET_NAIVE_DECOMPOSER_SPLIT_POS")
-        if split_pos_str is None:
-            raise EnvironmentError(
-                "Environment variable 'GRAPH_NET_NAIVE_DECOMPOSER_SPLIT_POS' is not set."
-            )
-        self.split_positions = [int(pos) for pos in split_pos_str.split(",")]
+        self.config = self.make_config(**self.convert_to_dict(config_str))
+
+    def make_config(
+        self,
+        split_positions=(),
+        group_head_and_tail=False,
+        chain_style=False,
+        output_dir="./tmp/naive_decomposer_dir",
+    ):
+        for pos in split_positions:
+            assert isinstance(
+                pos, int
+            ), f"split_positions should be list of int, {split_positions=}"
+        return {
+            "split_positions": split_positions,
+            "group_head_and_tail": group_head_and_tail,
+            "chain_style": chain_style,
+            "output_dir": output_dir,
+        }
 
     def __call__(self, gm: torch.fx.GraphModule, sample_inputs):
-        return convert_to_submodules_graph(
+        config = {
+            k: v
+            for k, v in self.config.items()
+            if k in {"split_positions", "group_head_and_tail", "chain_style"}
+        }
+        rewrited_gm = convert_to_submodules_graph(
             gm,
-            split_positions=self.split_positions,
             submodule_hook=self.get_naive_decomposer_extractor,
-            group_head_and_tail=False,
+            **config,
         )
+        return rewrited_gm
 
     def get_naive_decomposer_extractor(self, submodule, seq_no):
         return NaiveDecomposerExtractor(self, submodule, seq_no)
 
+    def convert_to_dict(self, config_str):
+        if config_str is None:
+            return {}
+        config_str = base64.b64decode(config_str).decode("utf-8")
+        config = json.loads(config_str)
+        assert isinstance(config, dict), f"config should be a dict. {config_str=}"
+        return config
+
 
 class NaiveDecomposerExtractor(torch.nn.Module):
     def __init__(self, parent_graph_extractor, submodule, seq_no):
@@ -54,6 +81,7 @@ def __init__(self, parent_graph_extractor, submodule, seq_no):
             dynamic=False,
             mut_graph_codes=[],
             placeholder_auto_rename=parent_graph_extractor.placeholder_auto_rename,
+            workspace_path=self.parent_graph_extractor.config["output_dir"],
         )
 
     def forward(self, *args):

graph_net/torch/single_device_runner.py

@@ -64,6 +64,7 @@ def main(args):
                 name=args.extract_name,
                 dynamic=False,
                 custom_extractor_path=args.custom_extractor_path,
+                custom_extractor_config=args.custom_extractor_config,
                 **dump_graph_options,
             )
             model = extract(**kwargs)(model)
@@ -123,5 +124,12 @@ def main(args):
         default=None,
         help="Custom extractor python file path",
     )
+    parser.add_argument(
+        "--custom-extractor-config",
+        type=str,
+        required=False,
+        default=None,
+        help="Custom extractor configuration string",
+    )
     args = parser.parse_args()
     main(args=args)