Remove ExirMetadata

Summary: exir_meta.in_spec/out_spec should now belong in the toplevel exported_program.call_spec.in_spec/out_spec

Reviewed By: ydwu4

Differential Revision: D47392208

fbshipit-source-id: d510b7d45565138dac08ad8463b6680083f8b95b

Author: angelayi

backends/backend_api.py

@@ -23,11 +23,7 @@
     get_lowered_module_name,
     LoweredBackendModule,
 )
-from executorch.exir.graph_module import (
-    attach_export_graph_metadata,
-    ExirMetadata,
-    get_control_flow_submodules,
-)
+from executorch.exir.graph_module import get_control_flow_submodules
 from executorch.exir.pass_base import ExportPass
 from torch._export.exported_program import ExportedProgram
 
@@ -166,12 +162,6 @@ def _partition_and_lower(
             {},
             [],
         )
-        meta = ExirMetadata(
-            in_spec=None,
-            out_spec=None,
-            update_spec=0,
-        )
-        attach_export_graph_metadata(submodule_program.graph_module, meta)
 
         lowered_submodule = to_backend(
             delegation_spec.backend_id,

backends/xnnpack/test/TARGETS

@@ -23,7 +23,6 @@ python_unittest(
         "//executorch/bundled_program:config",
         "//executorch/bundled_program:core",
         "//executorch/bundled_program/serialize:lib",
-        "//executorch/exir:graph_module",
         "//executorch/exir:lib",
         "//executorch/exir/passes:spec_prop_pass",
         "//executorch/exir/serialize:lib",
@@ -54,7 +53,6 @@ python_unittest(
         "//executorch/bundled_program:config",
         "//executorch/bundled_program:core",
         "//executorch/bundled_program/serialize:lib",
-        "//executorch/exir:graph_module",
         "//executorch/exir:lib",
         "//executorch/exir/dialects:lib",
         "//executorch/exir/passes:spec_prop_pass",
@@ -86,7 +84,6 @@ python_unittest(
         "//executorch/bundled_program:config",
         "//executorch/bundled_program:core",
         "//executorch/bundled_program/serialize:lib",
-        "//executorch/exir:graph_module",
         "//executorch/exir:lib",
         "//executorch/exir/passes:spec_prop_pass",
         "//executorch/exir/serialize:lib",

backends/xnnpack/test/test_xnnpack_utils.py

@@ -27,8 +27,6 @@
     serialize_from_bundled_program_to_flatbuffer,
 )
 
-from executorch.exir.graph_module import attach_export_graph_metadata, get_exir_meta
-
 from executorch.exir.passes.spec_prop_pass import SpecPropPass
 from executorch.exir.serialize import serialize_to_flatbuffer
 
@@ -306,14 +304,13 @@ def quantize_and_test_model_with_quantizer(
         )
         captured_program = exir.capture(module, example_inputs, config=capture_config)
         m = captured_program.graph_module
-        exir_meta = get_exir_meta(m)
 
         quantizer = QNNPackQuantizer()
         quantization_config = get_symmetric_quantization_config()
         quantizer.set_global(quantization_config)
         prepared = prepare_pt2e_quantizer(captured_program.graph_module, quantizer)
         converted = convert_pt2e(prepared)
-        attach_export_graph_metadata(converted, exir_meta)
+
         captured_program.graph_module = converted
         edge_program = captured_program.to_edge(get_xnnpack_edge_compile_config())
         delegated_module = self.lower_module_and_test_output(

exir/TARGETS

@@ -34,9 +34,6 @@ python_library(
         "graph_module.py",
     ],
     deps = [
-        ":common",
-        ":error",
-        ":graph",
         "//caffe2:torch",
     ],
 )
@@ -129,7 +126,6 @@ python_library(
         ":control_flow",  # @manual
         ":dynamic_shape",
         ":error",
-        ":graph_module",
         ":pass_base",  # @manual
         ":pass_manager",
         ":schema",
@@ -156,7 +152,6 @@ python_library(
         ":control_flow",
         ":delegate",
         ":error",
-        ":graph_module",
         ":memory",
         ":schema",
         ":tensor",
@@ -234,7 +229,6 @@ python_library(
     deps = [
         ":delegate",
         ":error",
-        ":graph_module",
         ":memory",
         "//caffe2:torch",
         "//caffe2/functorch:functorch_src",

exir/__init__.py

@@ -11,13 +11,6 @@
 from executorch.exir.dynamic_shape import DynamicMemoryPlanningMode
 from executorch.exir.emit import emit_program, EmitterOutput
 from executorch.exir.error import ExportError, ExportErrorType, InternalError
-from executorch.exir.graph_module import (
-    attach_export_graph_metadata,
-    EXIR_METADATA,
-    get_exir_meta,
-    make_export_graph_module,
-    reduce_graph_module,
-)
 from executorch.exir.pass_manager import PassManager, PassType
 from executorch.exir.passes import (
     aten_to_edge_passes,
@@ -119,35 +112,17 @@ class ExecutorchBackendConfig:
 
 
 # TODO(ycao): set up "__all__" to limit symbol exposure
-def _to_edge(expo_prog, config: EdgeCompileConfig) -> "ExirExportedProgram":
-    meta = get_exir_meta(expo_prog.graph_module)
+def _to_edge(ep, config: EdgeCompileConfig) -> "ExirExportedProgram":
     if config._check_ir_validity:
         try:
-            EXIRATenDialectVerifier()(expo_prog.graph_module)
+            EXIRATenDialectVerifier()(ep.graph_module)
         except ExportError:
             logging.info(
                 "If you'd like to disable IR validation checking, please set _check_ir_validity in EdgeCompileConfig, "
                 "like *.to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))."
             )
             raise
-    mutation = meta.mutation
-    num_mutated = len(mutation) if mutation is not None else 0
-    output_node = next(iter(reversed(expo_prog.graph_module.graph.nodes)))
-    assert output_node.op == "output"
-    output_node.args = (output_node.args[0][num_mutated:],)
-    expo_prog.graph_module.graph.eliminate_dead_code()
 
-    ep = ExirExportedProgram(
-        expo_prog.graph_module,
-        expo_prog.graph_module.graph,
-        ExportGraphSignature([], [], [], [], {}, {}, {}, None),
-        CallSpec(meta.in_spec, meta.out_spec),
-        {},
-        {},
-        [],
-        False,
-    )
-    attach_export_graph_metadata(ep.graph_module, meta)
     op_replace_pass = [OpReplacePass()] if config._use_edge_ops else []
     passes = aten_to_edge_passes.passes + op_replace_pass + config.passes
     new_ep = ep.transform(*passes)
@@ -198,7 +173,6 @@ def transform(self, *passes: PassType) -> "ExirExportedProgram":
             ep.equality_constraints,
             self.after_to_edge_passes,
         )
-
         transformed_ep.graph_module.meta.update(ep.graph_module.meta)
         transformed_ep.graph_module.meta.update(self.graph_module.meta)
         return transformed_ep
@@ -227,24 +201,6 @@ def _to_server(
         # return graph_module now.
         return res.graph_module
 
-    @property
-    def meta(self):
-        return self.graph_module.meta
-
-    @property
-    def in_spec(self):
-        meta = get_exir_meta(self.graph_module)
-        return meta.in_spec
-
-    @property
-    def out_spec(self):
-        meta = get_exir_meta(self.graph_module)
-        return meta.out_spec
-
-    @property
-    def graph(self):
-        return self.graph_module.graph
-
     @property
     def code(self):
         return self.graph_module.code
@@ -257,7 +213,6 @@ def to_executorch(
             raise RuntimeError("Must run to_edge before to_executorch.")
         config = config or ExecutorchBackendConfig()
         new_prog = self.transform(*edge_to_executorch_passes(config))
-        meta = get_exir_meta(new_prog.graph_module)
         executorch_prog = ExecutorchProgram(
             new_prog,
             emit_stacktrace=config.emit_stacktrace,
@@ -266,20 +221,10 @@ def to_executorch(
             constant_tensor_alignment=config.constant_tensor_alignment,
             delegate_alignment=config.delegate_alignment,
         )
-        attach_export_graph_metadata(executorch_prog.graph_module, meta)
-        # We only need to update the meta of the root graph module since it is
-        # reconstructed in ExecutorchProgram. The submodules are exactly the
-        # original submodules in new_prog.
         executorch_prog.graph_module.meta.update(new_prog.graph_module.meta)
         executorch_prog.graph_module.meta.update(self.graph_module.meta)
         return executorch_prog
 
-    def __reduce__(
-        self,
-    ) -> Tuple[Callable[..., "ExirExportedProgram"], Tuple[bytes]]:
-        _, (pickled_states,) = self.graph_module.__reduce__()
-        return (edge_gm_deserializer, (pickled_states,))
-
     def __deepcopy__(self, memo: Optional[Dict[int, Any]] = None) -> "ExportedProgram":
         gm = self.graph_module.__deepcopy__(memo)
         new_ep = ExirExportedProgram(
@@ -292,9 +237,6 @@ def __deepcopy__(self, memo: Optional[Dict[int, Any]] = None) -> "ExportedProgra
             copy.deepcopy(self.equality_constraints),
             self.after_to_edge_passes,
         )
-        attach_export_graph_metadata(
-            new_ep.graph_module, get_exir_meta(self.graph_module)
-        )
         new_ep.graph_module.meta.update(self.graph_module.meta)
         return new_ep
 
@@ -418,27 +360,6 @@ def get_multi_method_graph_module(self) -> "MultiMethodExirExportedProgram":
         return self._executorch_dialect_ir_program
 
 
-def edge_gm_deserializer(pickled_states: bytes) -> "ExirExportedProgram":
-    loaded_gm = reduce_graph_module(pickled_states)
-    # restore node.meta["val"], which is deleted before pickling
-    annotated_gm = aten_to_edge_passes(loaded_gm).graph_module
-    meta = get_exir_meta(annotated_gm)
-
-    ep = ExirExportedProgram(
-        annotated_gm.module,
-        annotated_gm.graph,
-        ExportGraphSignature([], [], [], [], {}, {}, {}, None),
-        CallSpec(meta.in_spec, meta.out_spec),
-        {},
-        {},
-        [],
-        after_to_edge_passes=True,
-    )
-    attach_export_graph_metadata(ep.graph_module, meta)
-    ep.graph_module.meta.update(annotated_gm.graph_module.meta)
-    return ep
-
-
 # This is to bootstrap the missing meta["val"] when 1. ph consists of scalar
 # 2. meta["val"] is not properly set in dispatch_trace.
 def _instantiate_missing_placeholder_val_with_real_inputs(gm, args):
@@ -468,7 +389,6 @@ def capture(
         )
 
     config = config or CaptureConfig()
-    mutation = None
     out_spec = None
     # TODO (zhxchen17) Always functionalize in a second pass no matter which path is taken.
     flat_args = tuple(pytree.tree_flatten(args)[0])
@@ -578,15 +498,8 @@ def convert_to_fake(x):
                     tracing_mode=tracing_mode,
                     _allow_non_fake_inputs=True,
                 )(*args)
-        module = make_export_graph_module(
-            graph_module,
-            graph_module.graph,
-        )
 
-        flatten_output(module)
-        meta = get_exir_meta(module)
-        meta.in_spec = in_spec
-        meta.out_spec = out_spec
+        flatten_output(graph_module)
 
     else:
         warnings.warn(
@@ -604,24 +517,22 @@ def convert_to_fake(x):
             raise InternalError(
                 "Using AOT mode is not supported for leagacy capture mode, please use pt2_mode=True instead."
             )
-        module = dispatch_trace(f, args)
-        # TODO(shunting) move this into ExportModuleState
-        meta = module.meta[EXIR_METADATA]
-    _instantiate_missing_placeholder_val_with_real_inputs(module, flat_args)
-    module._apply(torch.Tensor.contiguous)
-    meta = get_exir_meta(module)
-    meta.mutation = mutation  # pyre-ignore
+        graph_module = dispatch_trace(f, args)
+        in_spec, out_spec = graph_module.in_spec, graph_module.out_spec
+
+    _instantiate_missing_placeholder_val_with_real_inputs(graph_module, flat_args)
+    graph_module._apply(torch.Tensor.contiguous)
+
     ep = ExirExportedProgram(
-        module,
-        module.graph,
+        graph_module,
+        graph_module.graph,
         ExportGraphSignature([], [], [], [], {}, {}, {}, None),
-        CallSpec(meta.in_spec, meta.out_spec),
+        CallSpec(in_spec, out_spec),
         {},
         {},
         [],
         False,
     )
-    attach_export_graph_metadata(ep.graph_module, meta)
     return ep
 
 
@@ -720,18 +631,6 @@ def access_property_of_default_method(self, property_name: str):
         "to access property: {property_name}."""
         return getattr(default_program.graph_module, property_name)
 
-    @property
-    def meta(self):
-        return self.access_property_of_default_method("meta")
-
-    @property
-    def in_spec(self):
-        return self.meta[EXIR_METADATA].in_spec
-
-    @property
-    def out_spec(self):
-        return self.meta[EXIR_METADATA].out_spec
-
     @property
     def graph(self):
         return self.access_property_of_default_method("graph")

exir/delegate.py

@@ -7,11 +7,7 @@
 import torch.utils._pytree as pytree
 
 from executorch.backends.compile_spec_schema import CompileSpec
-from executorch.exir.graph_module import (
-    _get_submodule,
-    EXIR_METADATA,
-    make_export_graph_module,
-)
+from executorch.exir.graph_module import _get_submodule
 from executorch.exir.tracer import Value
 from torch._export.exported_program import ExportedProgram
 from torch._functorch.eager_transforms import (
@@ -281,22 +277,6 @@ def _fixup_output_node(gm: torch.fx.GraphModule) -> None:
             return
 
 
-def generate_in_spec_out_spec(gm: torch.fx.GraphModule) -> None:
-    output_nodes = []
-    for node in gm.graph.nodes:
-        if node.op == "output":
-            output_nodes = node.args[0]
-
-    all_placeholders = [node for node in gm.graph.nodes if node.op == "placeholder"]
-
-    (_, pytree_in) = tree_flatten(tuple(all_placeholders))
-    (_, pytree_out) = tree_flatten(tuple(output_nodes))
-
-    meta = gm.meta[EXIR_METADATA]
-    meta.in_spec = pytree_in
-    meta.out_spec = pytree_out
-
-
 def create_submodule_from_nodes(
     gm: torch.fx.GraphModule,
     node_list: NodeList,
@@ -319,14 +299,11 @@ def create_submodule_from_nodes(
     sorted_nodes = topo_sort(node_list)
 
     submodule_name = "fused_" + tag
-    initial_sub_gm, orig_inputs, orig_outputs = fuse_as_graphmodule(
+    sub_gm, orig_inputs, orig_outputs = fuse_as_graphmodule(
         gm, sorted_nodes, submodule_name
     )
 
-    _fixup_output_node(initial_sub_gm)
-    sub_gm = make_export_graph_module(
-        initial_sub_gm, initial_sub_gm.graph, submodule_name
-    )
+    _fixup_output_node(sub_gm)
 
     gm = insert_subgm(gm, sub_gm, orig_inputs, orig_outputs)
     if len(orig_outputs) == 1 and isinstance(orig_outputs[0].meta["val"], FakeTensor):
@@ -365,7 +342,6 @@ def create_submodule_from_nodes(
         submodule_node is not None
     ), f"No submodule was created with the nodes {node_list} in the graph {gm.graph}"
 
-    generate_in_spec_out_spec(sub_gm)
     return sub_gm, submodule_node