Skip to content

Fix #1870: fuse_as_graphmodule input ordering #82

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
laifenxiawucha wants to merge 1 commit into from
Closed

Fix #1870: fuse_as_graphmodule input ordering #82

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 27 additions & 7 deletions torch/fx/passes/utils/fuser_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,33 @@ def fuse_as_graphmodule(
] = {} # mapping of nodes from old graph to placeholder in new graph
node_map: dict[Node, Node] = {} # mapping of nodes from old graph to new graph

def create_placeholder(input_node: Node) -> None:
placeholder_node = subgraph.placeholder(
input_node.name, type_expr=input_node.type
)
placeholder_node.meta = copy.copy(input_node.meta)
node_to_placeholder[input_node] = placeholder_node

external_inputs: list[Node] = []
external_inputs_set: set[Node] = set()
for node in nodes:
for input_node in node.all_input_nodes:
if (
input_node not in partition_lookup_table
and input_node not in external_inputs_set
):
external_inputs.append(input_node)
external_inputs_set.add(input_node)

if external_inputs and all(
input_node.op == "placeholder" for input_node in external_inputs
):
graph_order = {node: i for i, node in enumerate(gm.graph.nodes)}
external_inputs.sort(key=lambda node: graph_order[node])

for input_node in external_inputs:
create_placeholder(input_node)

# handles inputs through graph.node_copy's arg_transform functions
def remap_inputs(x: Node) -> Node:
if x.op == "get_attr":
Expand All @@ -167,13 +194,6 @@ def remap_inputs(x: Node) -> Node:
# the node should have been copied already, as we are copying graph in the topological order
return node_map[x]

if x not in node_to_placeholder:
# x is not in subgraph, create a new placeholder for subgraph
placeholder_node = subgraph.placeholder(x.name, type_expr=x.type)
# copy all meta fields, even if some fields might be irrelevant for the placeholder node
placeholder_node.meta = copy.copy(x.meta)
node_to_placeholder[x] = placeholder_node

return node_to_placeholder[x]

# copy nodes in topological order
Expand Down