merge develop

Author: lixinqi

graph_net/test/naive_graph_decomposer_test.sh

@@ -5,7 +5,7 @@ MODEL_PATH_IN_SAMPLES=/timm/resnet18
 read -r -d '' json_str <<'EOF'
 {
     "output_dir": "/tmp/naive_decompose_workspace",
-    "split_positions": [0, 32],
+    "split_positions": [8, 32],
     "group_head_and_tail": true
 }
 EOF

graph_net/torch/decompose_util.py

@@ -72,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,
@@ -144,9 +150,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