(
exported_program: ExportedProgram,
edge_graph_module: torch.fx.GraphModule,
)
| 51 | |
| 52 | |
| 53 | def generate_node_to_external_map( |
| 54 | exported_program: ExportedProgram, |
| 55 | edge_graph_module: torch.fx.GraphModule, |
| 56 | ) -> Dict[torch.fx.Node, ExternalMeta]: |
| 57 | node_to_external_map = {} |
| 58 | for node in edge_graph_module.graph.nodes: |
| 59 | # The order in which we visit the placeholder node is same as the *args |
| 60 | # order for the forward(*args) signature for this gm. Using the order of |
| 61 | # the nodes as external_id to extract the right arg from *args at runtime |
| 62 | # |
| 63 | # Removing parameters/buffers since they will disappear from the signature |
| 64 | # at runtime |
| 65 | if node.op == "placeholder" and not is_param_node(exported_program, node): |
| 66 | node_to_external_map[node] = ExternalMeta( |
| 67 | external_id=len(node_to_external_map), |
| 68 | io_type=XNN_VALUE_FLAG_EXTERNAL_INPUT, |
| 69 | ) |
| 70 | for node in edge_graph_module.graph.nodes: |
| 71 | if node.op == "output": |
| 72 | for output_nodes in node.args: |
| 73 | for output_node in output_nodes: |
| 74 | if output_node in node_to_external_map: |
| 75 | raise RuntimeError( |
| 76 | f"Output node '{output_node}' is already in the inputs. " |
| 77 | "This is likely due to pass through arguments, which are not supported in XNNPACK Delegate." |
| 78 | ) |
| 79 | node_to_external_map[output_node] = ExternalMeta( |
| 80 | external_id=len(node_to_external_map), |
| 81 | io_type=XNN_VALUE_FLAG_EXTERNAL_OUTPUT, |
| 82 | ) |
| 83 | return node_to_external_map |
| 84 | |
| 85 | |
| 86 | def assert_default_dim_order(edge_graph_module: torch.fx.GraphModule) -> None: |
no test coverage detected