Given `gm` and some graph module which is called with target name `inline_mod_name`, this helper will inline all of the nodes from that called graph module into `gm`.
(gm: torch.fx.GraphModule, inline_mod_name: str)
| 75 | |
| 76 | |
| 77 | def _inline_module(gm: torch.fx.GraphModule, inline_mod_name: str): |
| 78 | """ |
| 79 | Given `gm` and some graph module which is called with target name `inline_mod_name`, |
| 80 | this helper will inline all of the nodes from that called graph module into `gm`. |
| 81 | """ |
| 82 | # Fetch the inner graph module that we want to inline inside `gm`. |
| 83 | inline_mod = dict(gm.named_modules())[inline_mod_name] |
| 84 | assert isinstance(inline_mod, torch.fx.GraphModule) |
| 85 | call_mod_node_to_replace = None |
| 86 | for node in gm.graph.nodes: |
| 87 | if node.op == "call_module" and node.target == inline_mod_name: |
| 88 | call_mod_node_to_replace = node |
| 89 | break |
| 90 | assert call_mod_node_to_replace is not None |
| 91 | |
| 92 | # Now actually do the swap. Note that we have to keep track of new nodes that are |
| 93 | # copied into `gm` -- we do this via replacement_mapping. |
| 94 | call_mod_args = call_mod_node_to_replace.args |
| 95 | replacement_mapping: Dict[torch.fx.Node, torch.fx.Node] = {} |
| 96 | ph_count = 0 |
| 97 | |
| 98 | def replacement_fn(node): |
| 99 | new_node = replacement_mapping[node] |
| 100 | new_node.meta = node.meta.copy() |
| 101 | return new_node |
| 102 | |
| 103 | for inline_node in inline_mod.graph.nodes: |
| 104 | if inline_node.op == "placeholder": |
| 105 | replacement_mapping[inline_node] = call_mod_args[ph_count] |
| 106 | ph_count += 1 |
| 107 | continue |
| 108 | |
| 109 | if inline_node.op == "output": |
| 110 | outputs = inline_node.args[0] |
| 111 | output_replacements = map_arg(outputs, replacement_fn) |
| 112 | call_mod_node_to_replace.replace_all_uses_with(output_replacements) |
| 113 | continue |
| 114 | |
| 115 | with gm.graph.inserting_before(call_mod_node_to_replace): |
| 116 | new_node = gm.graph.node_copy(inline_node, replacement_fn) |
| 117 | replacement_mapping[inline_node] = new_node |
| 118 | |
| 119 | gm.graph.eliminate_dead_code() |
| 120 | |
| 121 | |
| 122 | def get_unique_attr_name_in_module(mod_traced: torch.fx.GraphModule, name: str) -> str: |
no test coverage detected
searching dependent graphs…