subrange_graph_extract

Author: leon062112

graph_net/torch/subrange_graph.py

@@ -0,0 +1,66 @@
+import torch
+import torch.nn as nn
+import torch.fx as fx
+from typing import Set, Dict, List
+import itertools
+
+import argparse
+
+def subrange_graph_extract(original_gm: fx.GraphModule, start_idx: int, end_idx: int):
+    nodes_list = list(original_gm.graph.nodes)
+    if start_idx < 0 or end_idx >= len(nodes_list) or start_idx > end_idx:
+        raise ValueError("invalid indices")
+    
+    selected_nodes = set(nodes_list[start_idx:end_idx + 1])
+    
+    new_graph = fx.Graph()
+    node_map = {}
+    
+    external_inputs = set()
+    for node in selected_nodes:
+        if hasattr(node, 'args') and hasattr(node, 'kwargs'):
+            for arg in itertools.chain(node.args, node.kwargs.values()):
+                if isinstance(arg, fx.Node) and arg not in selected_nodes:
+                    external_inputs.add(arg)
+
+    print(f"found external inputs: {[node.name for node in external_inputs]}")
+    
+    for ext_node in external_inputs:
+        placeholder_name = f"input_{ext_node.name}"
+        new_placeholder = new_graph.placeholder(placeholder_name)
+        node_map[ext_node] = new_placeholder
+        print(f" {ext_node.name} create placeholder: {placeholder_name}")
+    
+    for node in selected_nodes:
+        if node.op == "placeholder":
+            new_node = new_graph.placeholder(node.name)
+            node_map[node] = new_node
+            print(f"kepp original placeholder: {node.name}")
+    
+    print(f"\ncopy node (node_map length: {len(node_map)}):")
+    for node in nodes_list[start_idx:end_idx + 1]:
+        if node.op == "placeholder":
+            continue 
+
+        print(f"  copy node: {node.name}")
+        missing_deps = []
+        if hasattr(node, 'args'):
+            for arg in node.args:
+                if isinstance(arg, fx.Node) and arg not in node_map:
+                    missing_deps.append(arg.name)
+        if missing_deps:
+            print(f"    missing deps: {missing_deps}")
+        else:
+            print(f"    deps ok")
+        
+        new_node = new_graph.node_copy(node, lambda n: node_map.get(n))
+        node_map[node] = new_node
+        print(f"    copy as: {new_node.name}")
+    
+    last_node = nodes_list[end_idx]
+    if last_node.op != "output":
+        new_graph.output(node_map[last_node])
+    else:
+        new_graph.output(node_map.get(last_node.args[0], last_node.args[0]))
+    
+    return fx.GraphModule(dict(original_gm.named_modules()), new_graph)