Resolve data flow mismatch between adjacent subgraphs in chain decomposition. (#424)

* init 'symbolic_dimension_reifier' field in graph_net.json

* refactor model_path_handler

* Resolve data flow mismatch between adjacent subgraphs in chain decomposition.

Author: lixinqi

graph_net/config/decomposition_error_tmp_torch_samples_list.txt

@@ -0,0 +1 @@
+samples/transformers-auto-model/dbmdz_electra-large-discriminator-finetuned-conll03-english

graph_net/test/naive_graph_decomposer_test.sh

@@ -11,6 +11,7 @@ config_json_str=$(cat <<EOF
     "handler_path": "$GRAPH_NET_ROOT/torch/naive_graph_decomposer.py",
     "handler_class_name": "NaiveDecomposerExtractor",
     "handler_config": {
+        "model_path_prefix": "$GRAPH_NET_ROOT/../",
         "output_dir": "/tmp/naive_decompose_workspace",
         "split_positions": [8, 16, 32],
         "chain_style": true,
@@ -21,4 +22,5 @@ EOF
 )
 CONFIG=$(echo $config_json_str | base64 -w 0)
 
-python3 -m graph_net.model_path_handler --model-path $GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES --handler-config=$CONFIG
+# python3 -m graph_net.model_path_handler --model-path $GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES --handler-config=$CONFIG
+python3 -m graph_net.model_path_handler --model-path-list $GRAPH_NET_ROOT/config/decomposition_error_tmp_torch_samples_list.txt --handler-config=$CONFIG

graph_net/torch/decompose_util.py

@@ -17,9 +17,6 @@ def convert_to_submodules_graph(
     chain_style=True: decompose gm into g0 * g1 * g2 * g3
     """
     gm = copy.deepcopy(gm)
-    num_placeholders = len(
-        [node for node in gm.graph.nodes if node.op == "placeholder"]
-    )
     submodules_body_nodes = [
         node
         for node in gm.graph.nodes
@@ -207,6 +204,32 @@ def _get_submodule_inputs_and_outputs(
     start_node_idx: int,
     end_node_idx: int,
     chain_style=False,
+):
+    if not chain_style:
+        (
+            minimal_input_nodes,
+            minimal_output_nodes,
+        ) = _get_minimal_submodule_inputs_and_outputs(
+            gm=gm, start_node_idx=start_node_idx, end_node_idx=end_node_idx
+        )
+        return minimal_input_nodes, minimal_output_nodes, []
+    else:
+        node_list = list(gm.graph.nodes)
+        input_nodes, _ = _get_minimal_submodule_inputs_and_outputs(
+            gm=gm, start_node_idx=start_node_idx, end_node_idx=len(node_list)
+        )
+        output_nodes, _ = _get_minimal_submodule_inputs_and_outputs(
+            gm=gm, start_node_idx=end_node_idx, end_node_idx=len(node_list)
+        )
+        identity_nodes_set = set(input_nodes) & set(output_nodes)
+        identity_nodes = [node for node in input_nodes if node in identity_nodes_set]
+        return input_nodes, output_nodes, identity_nodes
+
+
+def _get_minimal_submodule_inputs_and_outputs(
+    gm: torch.fx.GraphModule,
+    start_node_idx: int,
+    end_node_idx: int,
 ):
     count_ctx = NodeProducedOrConsumedCountCtx(
         defaultdict(int),
@@ -215,33 +238,34 @@ def _get_submodule_inputs_and_outputs(
     )
     node_list = list(gm.graph.nodes)
 
-    def _hashable(obj):
-        if isinstance(obj, slice):
-            return ("__slice__", obj.start, obj.stop, obj.step)
-        elif isinstance(obj, (list, tuple)):
-            return tuple(_hashable(x) for x in obj)
+    def get_args_node(arg):
+        if isinstance(arg, torch.fx.Node):
+            yield arg
+        elif isinstance(arg, (tuple, list)):
+            for x in arg:
+                yield from get_args_node(x)
+        elif isinstance(arg, slice):
+            yield arg.start
+            yield arg.stop
+            yield arg.step
         else:
-            return obj
+            assert isinstance(arg, (int, bool, float, str, type(None))), f"{type(arg)=}"
 
-    def get_related_node(node):
+    def get_args_node_and_self_node(node):
         for arg in node.args:
-            if isinstance(arg, tuple):
-                for x in arg:
-                    yield _hashable(x)
-            else:
-                yield _hashable(arg)
-        yield _hashable(node)
+            yield from get_args_node(arg)
+        yield node
 
     for node in node_list[0:start_node_idx]:
-        for related_node in get_related_node(node):
+        for related_node in get_args_node_and_self_node(node):
             count_ctx.node2before_input[related_node] += 1
 
     for node in node_list[start_node_idx:end_node_idx]:
-        for related_node in get_related_node(node):
+        for related_node in get_args_node_and_self_node(node):
             count_ctx.node2body[related_node] += 1
 
     for node in node_list[end_node_idx:]:
-        for related_node in get_related_node(node):
+        for related_node in get_args_node_and_self_node(node):
             count_ctx.node2after_output[related_node] += 1
 
     input_nodes = [
@@ -257,24 +281,4 @@ def get_related_node(node):
         if count_ctx.node2body[node] > 0
         if count_ctx.node2after_output[node] > 0
     ]
-    if not chain_style:
-        identity_nodes = []
-    else:
-        identity_nodes = [
-            node
-            for node in node_list
-            if count_ctx.node2before_input[node] > 0
-            if count_ctx.node2body[node] == 0
-            if count_ctx.node2after_output[node] > 0
-        ]
-        input_nodes_set = set(input_nodes)
-        input_nodes = [
-            *input_nodes,
-            *[node for node in identity_nodes if node not in input_nodes_set],
-        ]
-        output_nodes_set = set(output_nodes)
-        output_nodes = [
-            *output_nodes,
-            *[node for node in identity_nodes if node not in output_nodes_set],
-        ]
-    return input_nodes, output_nodes, identity_nodes
+    return input_nodes, output_nodes