Modifies the given graph module in-place to separate out the given nodes into a submodule. The given node_list should form a fully connected subgraph. Args: gm: The graph module that we want to partition node_list: A list of nodes that belong in the partition R
(
gm: torch.fx.GraphModule,
node_list: NodeList,
tag: str,
skip_legalize_graph: bool = False,
)
| 838 | |
| 839 | |
| 840 | def create_submodule_from_nodes( |
| 841 | gm: torch.fx.GraphModule, |
| 842 | node_list: NodeList, |
| 843 | tag: str, |
| 844 | skip_legalize_graph: bool = False, |
| 845 | ) -> Tuple[torch.fx.GraphModule, torch.fx.Node]: |
| 846 | """ |
| 847 | Modifies the given graph module in-place to separate out the given nodes |
| 848 | into a submodule. The given node_list should form a fully connected |
| 849 | subgraph. |
| 850 | |
| 851 | Args: |
| 852 | gm: The graph module that we want to partition |
| 853 | node_list: A list of nodes that belong in the partition |
| 854 | |
| 855 | Returns: |
| 856 | The submodule that has been partitioned, the call_module node in the |
| 857 | toplevel graph module calling the submodule |
| 858 | """ |
| 859 | sorted_nodes = topo_sort(node_list) |
| 860 | |
| 861 | submodule_name = "fused_" + tag |
| 862 | sub_gm, orig_inputs, orig_outputs = fuse_as_graphmodule( |
| 863 | gm, sorted_nodes, submodule_name |
| 864 | ) |
| 865 | |
| 866 | _fixup_output_node(sub_gm) |
| 867 | |
| 868 | gm = insert_subgm(gm, sub_gm, orig_inputs, orig_outputs) |
| 869 | submodule_node = None |
| 870 | for node in gm.graph.nodes: |
| 871 | if node.op == "call_module" and node.target == submodule_name: |
| 872 | submodule_node = node |
| 873 | |
| 874 | if submodule_node is None: |
| 875 | raise RuntimeError( |
| 876 | f"The submodule created with nodes {node_list} did not form \ |
| 877 | one fully contained subgraph. Check that these nodes form a \ |
| 878 | fully contained graph. Partitioned graph: {gm.graph}." |
| 879 | ) |
| 880 | |
| 881 | if len(orig_outputs) == 1 and isinstance(orig_outputs[0].meta["val"], FakeTensor): |
| 882 | # If the original output is a single tensor, it has been |
| 883 | # pytree.tree_flatten-ed to be a singleton list, so we want to replace |
| 884 | # all uses with a getitem call to the 0th index of the result |
| 885 | with gm.graph.inserting_after(submodule_node): |
| 886 | proxy_out = torch.fx.Proxy(submodule_node)[0].node # type: ignore[index] |
| 887 | submodule_node.replace_all_uses_with(proxy_out) |
| 888 | proxy_out.meta["val"] = submodule_node.meta["val"] |
| 889 | # Reset the args since it was overwritten in the previous line |
| 890 | proxy_out.args = (submodule_node, 0) |
| 891 | else: |
| 892 | # fuse_as_graphmodule will automatically propagate the metadata of the |
| 893 | # partition's last node to the getitem nodes that appear after the |
| 894 | # call_module node. However, in the case of delegation we do not want |
| 895 | # these getitem nodes to contain irrelevant previous metadata |
| 896 | # (ex. source_fn, # nn_module_stack) |
| 897 | for user_node in submodule_node.users: |