fix spec_prop_pass

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
@@ -45,109 +46,41 @@ def _is_mutable_buffer(
                     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,
-        )
-
-    def update_placeholder_tensor_specs(
-        self,
-        exported_program: torch.export.ExportedProgram,
-        graph_module: torch.fx.GraphModule,
-    ) -> None:
-        """
-        Update the tensor specs for all placeholder nodes such that
-        placeholders that are parameters are marked as constant.
-        """
-        for node in graph_module.graph.nodes:
-            if node.op != "placeholder":
-                continue
-            if "spec" not in node.meta:
-                raise RuntimeError(f"Placeholder node {node} missing meta['spec']")
-            spec = node.meta["spec"]
-            if isinstance(node.target, str) and (
-                node.target in exported_program.graph_signature.inputs_to_parameters
-                or (
-                    node.target in exported_program.graph_signature.inputs_to_buffers
-                    and not _is_mutable_buffer(node, exported_program.graph_signature)
-                )
-                or node.target
-                in exported_program.graph_signature.inputs_to_lifted_tensor_constants
-            ):
-                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),
+def SpecPropPass(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
+
+def update_placeholder_tensor_specs(
+    exported_program: torch.export.ExportedProgram,
+    graph_module: torch.fx.GraphModule,
+) -> None:
+    """
+    Update the tensor specs for all placeholder nodes such that
+    placeholders that are parameters are marked as constant.
+    """
+    for node in graph_module.graph.nodes:
+        if node.op != "placeholder":
+            continue
+        if "spec" not in node.meta:
+            raise RuntimeError(f"Placeholder node {node} missing meta['spec']")
+        spec = node.meta["spec"]
+        if isinstance(node.target, str) and (
+            node.target in exported_program.graph_signature.inputs_to_parameters
+            or (
+                node.target in exported_program.graph_signature.inputs_to_buffers
+                and not _is_mutable_buffer(node, exported_program.graph_signature)
             )
-        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)
+            or node.target
+            in exported_program.graph_signature.inputs_to_lifted_tensor_constants
+        ):
+            spec.const = True

exir/program/_program.py

@@ -54,7 +54,7 @@
 from executorch.exir.passes.replace_view_copy_with_view_pass import (
     ReplaceViewCopyWithViewPass,
 )
-from executorch.exir.passes.spec_prop_pass import SpecPropPass
+from executorch.exir.passes.spec_prop_pass import SpecPropPass, update_placeholder_tensor_specs
 from executorch.exir.passes.weights_to_outputs_pass import weights_to_outputs_pass
 from executorch.exir.print_program import pretty_print, print_program
 from executorch.exir.schema import Program
@@ -734,7 +734,7 @@ def edge_to_executorch_passes(
     """
     passes: List[PassType] = [
         *config.passes,
-        SpecPropPass(),
+        SpecPropPass,
         # ExecuTorch backend ops are unable to handle unbacked symints. So after
         # this pass, passes cannot be Interpreter-based, because it will fail if
         # there exists an unbacked symint operation.
@@ -1390,7 +1390,7 @@ def to_executorch(
                 new_gm_res = p(new_gm)
                 assert new_gm_res is not None
                 new_gm = new_gm_res.graph_module
-                if isinstance(p, SpecPropPass):
+                if p is SpecPropPass:
                     # Note that this is a hacky way to get around the fact that
                     # placeholder nodes corresponding to the parameters of the graph module
                     # shall not participate in memory planning. It increases runtime memory
@@ -1401,7 +1401,7 @@ def to_executorch(
                     # Working with GraphModule does not provide all the information contained
                     # in the ExportedProgram
                     # TODO(who?)
-                    p.update_placeholder_tensor_specs(program, new_gm)
+                    update_placeholder_tensor_specs(program, new_gm)
 
             # Extract constants if the config says too.
             if config.external_constants: