fix spec_prop_pass (#7974)

Summary:

fix spec_prop_pass by just consulting node.meta["val"]

Differential Revision: D66996237

Author: ydwu4

exir/passes/spec_prop_pass.py

@@ -15,6 +15,7 @@
 from torch.export.exported_program import ExportGraphSignature
 from torch.fx.node import Node
 from torch.utils import _pytree as pytree
+from torch.fx.passes.infra.pass_base import PassResult
 
 
 # pyre-ignore
@@ -44,19 +45,9 @@ def _is_mutable_buffer(
                 if fqn in graph_signature.buffers_to_mutate.values():
                     return True
     return False
-
-
-class SpecPropPass(ExportPass):
-    def __init__(self) -> None:
-        super().__init__()
-
-    def on_attr(self, attr: ProxyValue) -> None:
-        attr.node.meta["spec"] = pytree.tree_map_only(
-            torch.Tensor,
-            make_spec,
-            attr.data,
-        )
-
+class SpecPropPass:
+    def __call__(self, gm: torch.fx.GraphModule) -> PassResult:
+        return spec_prop_pass(gm)
     def update_placeholder_tensor_specs(
         self,
         exported_program: torch.export.ExportedProgram,
@@ -83,71 +74,16 @@ def update_placeholder_tensor_specs(
             ):
                 spec.const = True
 
-    # pyre-ignore
-    def placeholder(self, name: str, arg, meta):
-        meta["spec"] = make_spec(arg)
-        return super().placeholder(name, arg, meta)
-
-    # pyre-ignore
-    def call_operator(self, op, args, kwargs, meta):
-        args_data, kwargs_data = pytree.tree_map_only(
-            ProxyValue, lambda x: x.data, (args, kwargs)
-        )
-        meta["spec"] = pytree.tree_map(make_spec, op(*args_data, **kwargs_data))
-        return super().call_operator(op, args, kwargs, meta)
-
-    # pyre-ignore
-    def call_getitem(self, value, key: int, meta):
-        meta["spec"] = value.node.meta["spec"][key]
-        return super().call_getitem(value, key, meta)
-
-    # pyre-ignore
-    def call_cond(self, pred, true_fn, false_fn, inputs, meta):
-        # true_fn/false_fn return tensors of the same shape, so we can pick
-        # either one here.
-        *_, true_out_node = true_fn.graph.nodes
-        meta["spec"] = pytree.tree_map(make_spec, true_out_node.meta["val"])
-        return super().call_cond(pred, true_fn, false_fn, inputs, meta)
-
-    def call_map(
-        self,
-        f: torch.fx.GraphModule,
-        mapped_args: List[ProxyValue],
-        operands: List[ProxyValue],
-        meta: NodeMetadata,
-    ) -> ProxyValue:
-        mapped_dim_size = [arg.data for arg in mapped_args][0].size(0)
-        *_, body_out_node = f.graph.nodes
-        body_out_node_fake_tensor = body_out_node.meta["val"]
-        map_fake_tensor = pytree.tree_map_only(
-            torch.Tensor,
-            lambda x: x.new_empty(mapped_dim_size, *x.shape),
-            body_out_node_fake_tensor,
-        )
-        meta["spec"] = pytree.tree_map(make_spec, map_fake_tensor)
-        return super().call_map(f, mapped_args, operands, meta)
-
-    # pyre-ignore
-    def call_delegate(self, lowered_module, args, kwargs, meta):
-        args_data, kwargs_data = pytree.tree_map_only(
-            ProxyValue, lambda x: x.data, (args, kwargs)
-        )
-        # If spec is missing, re-genenrate it with args data
-        if "spec" not in meta:
-            meta["spec"] = pytree.tree_map(
-                make_spec,
-                executorch_call_delegate(lowered_module, *args_data),
-            )
-        return super().call_delegate(lowered_module, args, kwargs, meta)
-
-    # pyre-ignore
-    def output(self, results, meta):
-        # pyre-ignore
-        def get_spec(x):
-            if isinstance(x, ProxyValue):
-                return x.node.meta["spec"]
-            else:
-                return make_spec(x)
-
-        meta["spec"] = pytree.tree_map(get_spec, results)
-        return super().output(results, meta)
+def spec_prop_pass(gm: torch.fx.GraphModule) -> PassResult:
+    # Update all the meta["val"]
+    pass_result = ExportPass()(gm)
+    assert pass_result is not None
+    gm = pass_result.graph_module
+    # set node.meta["spec"] based on meta["val"]
+    for module in gm.modules():
+        if isinstance(module, torch.fx.GraphModule):
+            for node in module.graph.nodes:
+                if node.op == "get_attr":
+                    continue
+                node.meta["spec"] = pytree.tree_map(lambda meta_val: make_spec(meta_val), node.meta["val"])
+    return pass_result