| 47 | |
| 48 | |
| 49 | class GraphModuleSerializer(export_serialize.GraphModuleSerializer): |
| 50 | def __init__( |
| 51 | self, |
| 52 | graph_signature: ep.ExportGraphSignature, |
| 53 | module_call_graph: List[ep.ModuleCallEntry], |
| 54 | ) -> None: |
| 55 | super().__init__(graph_signature, module_call_graph) |
| 56 | self.state_dict: Dict[str, torch.Tensor] = {} # TODO(T157676982) |
| 57 | |
| 58 | def serialize_operator( |
| 59 | self, |
| 60 | target: Union[ |
| 61 | str, |
| 62 | EdgeOpOverload, |
| 63 | BackendOpOverload, |
| 64 | torch._ops.OpOverload, |
| 65 | torch._ops.HigherOrderOperator, |
| 66 | ], |
| 67 | ) -> str: |
| 68 | if isinstance(target, str): |
| 69 | return target |
| 70 | elif target.__module__.startswith("executorch.exir.dialects.edge"): |
| 71 | # TODO(zhxchen17) Maybe provide a function name helper in FX. |
| 72 | # From torch.fx.node._get_qualified_name |
| 73 | module = target.__module__.replace( |
| 74 | "executorch.exir.dialects.edge._ops", |
| 75 | "executorch.exir.dialects.edge.ops", |
| 76 | ) |
| 77 | return f"{module}.{target.__name__}" |
| 78 | elif target.__module__.startswith("executorch.exir.dialects.backend"): |
| 79 | module = target.__module__.replace( |
| 80 | "executorch.exir.dialects.backend._ops", |
| 81 | "executorch.exir.dialects.backend.ops", |
| 82 | ) |
| 83 | return f"{module}.{target.__name__}" |
| 84 | |
| 85 | return super().serialize_operator(target) |
| 86 | |
| 87 | def handle_call_function(self, node: torch.fx.Node) -> None: |
| 88 | assert node.op == "call_function" |
| 89 | |
| 90 | if node.target is memory.alloc: |
| 91 | ex_node = schema.Node( |
| 92 | name=node.name, |
| 93 | target="memory.alloc", |
| 94 | inputs=self.serialize_alloc_inputs(node.args), |
| 95 | outputs=self.serialize_arbitrary_outputs(node), |
| 96 | metadata=self.serialize_metadata(node), |
| 97 | ) |
| 98 | self.graph_state.nodes.append(ex_node) |
| 99 | return |
| 100 | elif isinstance(node.target, EdgeOpOverload): |
| 101 | assert node.target._op is not None |
| 102 | ex_node = schema.Node( |
| 103 | name=node.name, |
| 104 | target=self.serialize_operator(node.target), |
| 105 | # pyre-ignore Undefined attribute [16]: Item `typing.Callable` of |
| 106 | # `typing.Union[typing.Callable[..., typing.Any], str]` has no attribute `_op`. |