refactor parse_sole_graph_module

Author: lixinqi

graph_net/test/batch_init_input_tensor_constraints_test.sh

@@ -11,7 +11,7 @@ config_json_str=$(cat <<EOF
     "handler_path": "$GRAPH_NET_ROOT/constraint_util.py",
     "handler_class_name": "UpdateInputTensorConstraints",
     "handler_config": {
-        "resume": false,
+        "resume": true,
         "model_path_prefix": "$GRAPH_NET_ROOT/../",
         "data_input_predicator_filepath": "$GRAPH_NET_ROOT/torch/constraint_util.py",
         "data_input_predicator_class_name": "NaiveDataInputPredicator",
@@ -26,4 +26,4 @@ EOF
 )
 CONFIG=$(echo $config_json_str | base64 -w 0)
 
-python3 -m graph_net.model_path_handler --model-path-list $GRAPH_NET_ROOT/config/small_torch_samples_list.txt --handler-config=$CONFIG
+python3 -m graph_net.model_path_handler --model-path-list $GRAPH_NET_ROOT/config/torch_samples_list.txt --handler-config=$CONFIG

graph_net/test/shape_prop_batch_init_input_tensor_constraints_test.sh

@@ -11,7 +11,7 @@ config_json_str=$(cat <<EOF
     "handler_path": "$GRAPH_NET_ROOT/constraint_util.py",
     "handler_class_name": "UpdateInputTensorConstraints",
     "handler_config": {
-        "resume": false,
+        "resume": true,
         "model_path_prefix": "$GRAPH_NET_ROOT/../",
         "data_input_predicator_filepath": "$GRAPH_NET_ROOT/torch/constraint_util.py",
         "data_input_predicator_class_name": "NaiveDataInputPredicator",
@@ -26,4 +26,4 @@ EOF
 )
 CONFIG=$(echo $config_json_str | base64 -w 0)
 
-python3 -m graph_net.model_path_handler --model-path-list $GRAPH_NET_ROOT/config/small_torch_samples_list.txt --handler-config=$CONFIG
+python3 -m graph_net.model_path_handler --model-path-list $GRAPH_NET_ROOT/config/torch_samples_list.txt --handler-config=$CONFIG

graph_net/torch/fx_graph_parse_util.py

@@ -1,4 +1,5 @@
 import torch
+import inspect
 
 
 def parse_sole_graph_module(module, inputs):
@@ -14,12 +15,46 @@ def my_backend(gm, sample_inputs):
 
     torch.compile(module, backend=my_backend)(*inputs)
     assert traced_module is not None
-    assert all(id(a) == id(b) for a, b in zip(inputs, traced_sample_inputs))
     for node in traced_module.graph.nodes:
         if node.op != "placeholder":
             continue
         assert node.target[:2] == "L_" or node.target[:2] == "l_", f"{node.target=}"
         node.target = node.target[2:]
+        if node.target[0] == "l":
+            node.target = "L" + node.target[1:]
         assert node.name[:2] == "L_" or node.name[:2] == "l_", f"{node.name=}"
         node.name = node.name[2:]
+        if node.name[0] == "l":
+            node.name = "L" + node.name[1:]
+
+    def get_input_names_from_signature():
+        return inspect.signature(module.forward).parameters
+
+    def get_input_names_from_placeholder():
+        return [
+            node.name for node in traced_module.graph.nodes if node.op == "placeholder"
+        ]
+
+    def get_diff_input_names():
+        placeholder_names = set(get_input_names_from_placeholder())
+        return [
+            (i, name)
+            for i, name in enumerate(get_input_names_from_signature())
+            if name not in placeholder_names
+        ]
+
+    if len(inputs) == len(traced_sample_inputs) + 1:
+        diff_input_names = get_diff_input_names()
+        assert len(diff_input_names) == 1, f"{diff_input_names=}"
+        pos, name = diff_input_names[0]
+        for i, node in enumerate(traced_module.graph.nodes):
+            if i < pos:
+                assert node.op == "placeholder"
+            elif i == pos:
+                with traced_module.graph.inserting_before(node):
+                    traced_module.graph.placeholder(name)
+            else:
+                break
+        traced_module.recompile()
+    assert len(get_diff_input_names()) == 0, f"{get_diff_input_names()=}"
     return traced_module