| 51 | |
| 52 | |
| 53 | class SpecPropPass(ExportPass): |
| 54 | def __init__(self) -> None: |
| 55 | super().__init__() |
| 56 | |
| 57 | def __call__(self, graph_module: torch.fx.GraphModule) -> PassResult: |
| 58 | # Re-trace metadata to ensure it's up to date. |
| 59 | res = ExportPass()(graph_module) |
| 60 | assert res is not None |
| 61 | gm = res.graph_module |
| 62 | |
| 63 | def get_spec(x): |
| 64 | if hasattr(x, "meta"): |
| 65 | return x.meta.get("spec", None) |
| 66 | else: |
| 67 | return None |
| 68 | |
| 69 | for module in gm.modules(): |
| 70 | if isinstance(module, torch.fx.GraphModule): |
| 71 | for node in module.graph.nodes: |
| 72 | meta_val = node.meta.get("val", None) |
| 73 | if node.op == "output": |
| 74 | node.meta["spec"] = pytree.tree_map(get_spec, node.args[0]) |
| 75 | elif node.op == "call_function" and node.target == operator.getitem: |
| 76 | value_spec = pytree.tree_map(get_spec, node.args[0]) |
| 77 | node.meta["spec"] = value_spec[node.args[1]] |
| 78 | elif ( |
| 79 | node.op == "call_function" |
| 80 | and node.target == executorch_call_delegate |
| 81 | ): |
| 82 | # Note: We currently rely on delegate node specs not being regenerated, |
| 83 | # as the spec is set somewhat manually when adding the call delegate node. |
| 84 | # If we regenerate, it can change and break lowering (it becomes a tuple?). |
| 85 | # Ideally, we should figure out how to make the spec regeneration not break |
| 86 | # things. |
| 87 | # |
| 88 | # We do need to regenerate non-call-delegate node specs, as this pass is called |
| 89 | # multiple times in some lowering paths (backends can and do call it). |
| 90 | if "spec" not in node.meta: |
| 91 | node.meta["spec"] = pytree.tree_map(make_spec, meta_val) |
| 92 | else: |
| 93 | node.meta["spec"] = pytree.tree_map(make_spec, meta_val) |
| 94 | return res |
| 95 | |
| 96 | def call(self, graph_module: torch.fx.GraphModule) -> PassResult: |
| 97 | return self(graph_module) |
| 98 | |
| 99 | def update_placeholder_tensor_specs( |
| 100 | self, |
| 101 | exported_program: torch.export.ExportedProgram, |
| 102 | graph_module: torch.fx.GraphModule, |
| 103 | ) -> None: |
| 104 | """ |
| 105 | Update the tensor specs for all placeholder nodes such that |
| 106 | placeholders that are parameters are marked as constant. |
| 107 | """ |
| 108 | for node in graph_module.graph.nodes: |
| 109 | if node.op != "placeholder": |
| 110 | continue |
no outgoing calls