Cleaning et_schema of InferenceRun for OSS (#565)

Summary:

These are concepts that aren't useful in the OSS world. Associated tests were also moved

- Moved out InferenceRun, OperatorGraphWithStats, ExportedETOperatorGraph for internal
- Changed FXOperatorGraph to extending OperatorGraph instead of OperatorGraphWithStats
- Tests for EDIR (to be renamed) moved to internal since they are heavily coupled with InferenceRun

Reviewed By: tarun292

Differential Revision: D49839120

Author: Jack-Khuu

sdk/edir/TARGETS

@@ -17,13 +17,8 @@ python_library(
         "et_schema.py",
     ],
     deps = [
-        "fbsource//third-party/pypi/numpy:numpy",
         ":base_schema",
         "//caffe2:torch",
         "//executorch/exir:lib",
-        "//executorch/exir:schema",
-        "//executorch/exir/_serialize:lib",
-        "//executorch/sdk/etdump:schema",
-        "//executorch/sdk/etdump:serialize",
     ],
 )

sdk/edir/et_schema.py

@@ -15,9 +15,9 @@ python_unittest(
         "//caffe2:torch",
         "//executorch/exir:lib",
         "//executorch/exir:lowered_backend_module",
-        "//executorch/sdk/edir:base_schema",
         "//executorch/sdk/edir:et_schema",
         "//executorch/sdk/etdump:schema",
+        "//executorch/sdk/fb:et_schema",
     ],
 )
 
@@ -31,9 +31,9 @@ python_library(
         "//caffe2:torch",
         "//executorch/exir:lib",
         "//executorch/exir:lowered_backend_module",
-        "//executorch/sdk/edir:base_schema",
         "//executorch/sdk/edir:et_schema",
         "//executorch/sdk/etdump:schema",
+        "//executorch/sdk/fb:et_schema",
     ],
 )
 
@@ -42,14 +42,14 @@ python_binary(
     srcs = [
         "exported_op_graph_test.py",
     ],
-    main_module = "executorch.sdk.edir.tests.exported_op_graph_test",
+    main_module = "executorch.sdk.fb.tests.exported_op_graph_test",
     deps = [
         "fbsource//third-party/pypi/parameterized:parameterized",
         "//caffe2:torch",
         "//executorch/exir:lib",
         "//executorch/exir:lowered_backend_module",
-        "//executorch/sdk/edir:base_schema",
         "//executorch/sdk/edir:et_schema",
         "//executorch/sdk/etdump:schema",
+        "//executorch/sdk/fb:et_schema",
     ],
 )

sdk/edir/tests/TARGETS renamed to sdk/fb/tests/TARGETS

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
 
 import copy
 import unittest
@@ -18,13 +14,9 @@
     ExirExportedProgram,
 )
 from executorch.exir.lowered_backend_module import LoweredBackendModule
-from executorch.sdk.edir.base_schema import OperatorNode
-from executorch.sdk.edir.et_schema import (
-    ExportedETOperatorGraph,
-    FXOperatorGraph,
-    InferenceRun,
-)
+from executorch.sdk.edir.et_schema import FXOperatorGraph
 from executorch.sdk.etdump.schema import ETDump, ProfileBlock, ProfileEvent, RunData
+from executorch.sdk.fb.et_schema import ExportedETOperatorGraph, InferenceRun
 
 from parameterized import parameterized
 from torch import Tensor
@@ -697,37 +689,36 @@ def test_gen_from_fx_graph(self, model_name: str, model: torch.nn.Module) -> Non
         op_graph = gen_fx_graph_file_contents(et_program.dump_graph_module())
         self.check_graph_equal(op_graph, model_name, "et_dialect")
 
-    # pyre-ignore
-    @parameterized.expand(MODELS)
-    def test_metadata_attaching(self, model_name: str, model: torch.nn.Module) -> None:
-        _, _, et_program = gen_graphs_from_model(model)
-        op_graph = FXOperatorGraph.gen_operator_graph(et_program.dump_graph_module())
-        inference_run = model.gen_inference_run()
-        op_graph.attach_metadata(inference_run)
-
-        def verify_metadata_containment(
-            graph: FXOperatorGraph, inference_run: InferenceRun
-        ) -> None:
-            validation_map = inference_run.node_metadata
-
-            for node in graph.elements:
-                # Recursively check subgraph nodes
-                if isinstance(node, FXOperatorGraph):
-                    verify_metadata_containment(node, inference_run)
-                # Check that each node contains the corresponding metadata fields
-                if isinstance(node, OperatorNode) and node.metadata is not None:
-                    metadata = node.metadata
-                    debug_handle = metadata.get("debug_handle")
-                    if debug_handle in validation_map:
-                        self.assertDictContainsSubset(
-                            validation_map[debug_handle], metadata
-                        )
-
-        # Check for run level metadata
-        if op_graph.metadata is not None:
-            self.assertDictContainsSubset(inference_run.run_metadata, op_graph.metadata)
-
-        verify_metadata_containment(op_graph, inference_run)
+    # @parameterized.expand(MODELS)
+    # def test_metadata_attaching(self, model_name: str, model: torch.nn.Module) -> None:
+    #     _, _, et_program = gen_graphs_from_model(model)
+    #     op_graph = FXOperatorGraph.gen_operator_graph(et_program.dump_graph_module())
+    #     inference_run = model.gen_inference_run()
+    #     op_graph.attach_metadata(inference_run)
+
+    #     def verify_metadata_containment(
+    #         graph: FXOperatorGraph, inference_run: InferenceRun
+    #     ) -> None:
+    #         validation_map = inference_run.node_metadata
+
+    #         for node in graph.elements:
+    #             # Recursively check subgraph nodes
+    #             if isinstance(node, FXOperatorGraph):
+    #                 verify_metadata_containment(node, inference_run)
+    #             # Check that each node contains the corresponding metadata fields
+    #             if isinstance(node, OperatorNode) and node.metadata is not None:
+    #                 metadata = node.metadata
+    #                 debug_handle = metadata.get("debug_handle")
+    #                 if debug_handle in validation_map:
+    #                     self.assertDictContainsSubset(
+    #                         validation_map[debug_handle], metadata
+    #                     )
+
+    #     # Check for run level metadata
+    #     if op_graph.metadata is not None:
+    #         self.assertDictContainsSubset(inference_run.run_metadata, op_graph.metadata)
+
+    #     verify_metadata_containment(op_graph, inference_run)
 
 
 class ExportedOpGraphTest(unittest.TestCase):
@@ -747,36 +738,35 @@ def test_gen_from_emitted_program(
             "Please run `//executorch/sdk/edir/tests:generate_fixtures` to regenerate the fixtures.",
         )
 
-    # pyre-ignore
-    @parameterized.expand(MODELS)
-    def test_metadata_attaching(self, model_name: str, model: torch.nn.Module) -> None:
-        op_graph = generate_op_graph(model, model.get_random_inputs())
-        inference_run = model.gen_inference_run()
-        op_graph.attach_metadata(inference_run)
-
-        def verify_metadata_containment(
-            graph: ExportedETOperatorGraph, inference_run: InferenceRun
-        ) -> None:
-            validation_map = inference_run.node_metadata
-
-            for node in graph.elements:
-                # Recursively check subgraph nodes
-                if isinstance(node, ExportedETOperatorGraph):
-                    verify_metadata_containment(node, inference_run)
-                # Check that each node contains the corresponding metadata fields
-                if isinstance(node, OperatorNode) and node.metadata is not None:
-                    metadata = node.metadata
-                    debug_handle = metadata.get("debug_handle")
-                    if debug_handle in validation_map:
-                        self.assertDictContainsSubset(
-                            validation_map[debug_handle], metadata
-                        )
-
-        # Check for run level metadata
-        if op_graph.metadata is not None:
-            self.assertDictContainsSubset(inference_run.run_metadata, op_graph.metadata)
-
-        verify_metadata_containment(op_graph, inference_run)
+    # @parameterized.expand(MODELS)
+    # def test_metadata_attaching(self, model_name: str, model: torch.nn.Module) -> None:
+    #     op_graph = generate_op_graph(model, model.get_random_inputs())
+    #     inference_run = model.gen_inference_run()
+    #     op_graph.attach_metadata(inference_run)
+
+    #     def verify_metadata_containment(
+    #         graph: ExportedETOperatorGraph, inference_run: InferenceRun
+    #     ) -> None:
+    #         validation_map = inference_run.node_metadata
+
+    #         for node in graph.elements:
+    #             # Recursively check subgraph nodes
+    #             if isinstance(node, ExportedETOperatorGraph):
+    #                 verify_metadata_containment(node, inference_run)
+    #             # Check that each node contains the corresponding metadata fields
+    #             if isinstance(node, OperatorNode) and node.metadata is not None:
+    #                 metadata = node.metadata
+    #                 debug_handle = metadata.get("debug_handle")
+    #                 if debug_handle in validation_map:
+    #                     self.assertDictContainsSubset(
+    #                         validation_map[debug_handle], metadata
+    #                     )
+
+    #     # Check for run level metadata
+    #     if op_graph.metadata is not None:
+    #         self.assertDictContainsSubset(inference_run.run_metadata, op_graph.metadata)
+
+    #     verify_metadata_containment(op_graph, inference_run)
 
 
 class InferenceRunTest(unittest.TestCase):

sdk/edir/tests/exported_op_graph_test.py renamed to sdk/fb/tests/exported_op_graph_test.py

@@ -15,7 +15,7 @@ python_library(
         ":inspector_utils",
         "//caffe2:torch",
         "//executorch/exir:lib",
-        "//executorch/sdk/edir:et_schema",
+        "//executorch/sdk/edir:base_schema",
         "//executorch/sdk/etdump:schema_flatcc",
     ],
 )
@@ -26,6 +26,7 @@ python_library(
         "_inspector_utils.py",
     ],
     deps = [
+        "//executorch/sdk/edir:base_schema",
         "//executorch/sdk/edir:et_schema",
         "//executorch/sdk/etdump:schema_flatcc",
         "//executorch/sdk/etdump:serialize",

sdk/edir/tests/fixtures/composite_delegate_module.txt renamed to sdk/fb/tests/fixtures/composite_delegate_module.txt

@@ -6,11 +6,9 @@
 
 from typing import Dict, Mapping, Optional
 
-from executorch.sdk.edir.et_schema import (
-    FXOperatorGraph,
-    OperatorGraphWithStats,
-    OperatorNode,
-)
+from executorch.sdk.edir.base_schema import OperatorNode
+
+from executorch.sdk.edir.et_schema import FXOperatorGraph, OperatorGraph
 from executorch.sdk.etdump.schema_flatcc import ETDumpFlatCC
 
 from executorch.sdk.etdump.serialize import deserialize_from_etdump_flatcc
@@ -21,7 +19,7 @@
 
 def gen_graphs_from_etrecord(
     etrecord: ETRecord,
-) -> Mapping[str, OperatorGraphWithStats]:
+) -> Mapping[str, OperatorGraph]:
     op_graph_map = {}
     if etrecord.graph_map is not None:
         op_graph_map = {
@@ -39,7 +37,7 @@ def gen_graphs_from_etrecord(
 # TODO: use anonymous function to avoid passing the dict around
 # and move this inside of the OperatorGraphWithStats class
 def create_debug_handle_to_op_node_mapping(
-    op_graph: OperatorGraphWithStats,
+    op_graph: OperatorGraph,
     debug_handle_to_op_node_map: Dict[int, OperatorNode],
 ) -> None:
     """
@@ -48,7 +46,7 @@ def create_debug_handle_to_op_node_mapping(
     """
     # Recursively searches through the metadata of nodes
     for element in op_graph.elements:
-        if isinstance(element, OperatorGraphWithStats):
+        if isinstance(element, OperatorGraph):
             create_debug_handle_to_op_node_mapping(element, debug_handle_to_op_node_map)
         if isinstance(element, OperatorNode) and element.metadata is not None:
             metadata = element.metadata

sdk/edir/tests/fixtures/composite_delegate_module_fx_graph_aten_dialect.txt renamed to sdk/fb/tests/fixtures/composite_delegate_module_fx_graph_aten_dialect.txt

@@ -25,7 +25,7 @@
 import torch
 from executorch.exir import ExportedProgram
 
-from executorch.sdk.edir.et_schema import OperatorGraphWithStats, OperatorNode
+from executorch.sdk.edir.base_schema import OperatorGraph, OperatorNode
 from executorch.sdk.etdump.schema_flatcc import ETDumpFlatCC, ProfileEvent
 from executorch.sdk.inspector._inspector_utils import (
     create_debug_handle_to_op_node_mapping,
@@ -392,9 +392,9 @@ def __init__(
         etdump = gen_etdump_object(etdump_path=etdump_path)
         self.event_blocks = EventBlock._gen_from_etdump(etdump, etdump_scale)
 
-        self._op_graph_dict: Mapping[
-            str, OperatorGraphWithStats
-        ] = gen_graphs_from_etrecord(etrecord=self._etrecord)
+        self._op_graph_dict: Mapping[str, OperatorGraph] = gen_graphs_from_etrecord(
+            etrecord=self._etrecord
+        )
 
         # Use the delegate map from etrecord, associate debug handles with each event
         for event_block in self.event_blocks:

sdk/edir/tests/fixtures/composite_delegate_module_fx_graph_edge_dialect.txt renamed to sdk/fb/tests/fixtures/composite_delegate_module_fx_graph_edge_dialect.txt

@@ -27,6 +27,7 @@ python_unittest(
     name = "inspector_utils_test",
     srcs = ["inspector_utils_test.py"],
     deps = [
+        "//executorch/sdk/edir:base_schema",
         "//executorch/sdk/edir:et_schema",
         "//executorch/sdk/etrecord:etrecord",
         "//executorch/sdk/etrecord/tests:etrecord_test_library",

sdk/edir/tests/fixtures/composite_delegate_module_fx_graph_et_dialect.txt renamed to sdk/fb/tests/fixtures/composite_delegate_module_fx_graph_et_dialect.txt

@@ -8,12 +8,9 @@
 import unittest
 from typing import Dict, Tuple
 
-from executorch.sdk.edir.et_schema import (
-    FXOperatorGraph,
-    OperatorGraphWithStats,
-    OperatorNode,
-    ValueNode,
-)
+from executorch.sdk.edir.base_schema import OperatorGraph, OperatorNode, ValueNode
+
+from executorch.sdk.edir.et_schema import FXOperatorGraph
 from executorch.sdk.etrecord import generate_etrecord, parse_etrecord
 
 from executorch.sdk.etrecord.tests.etrecord_test import TestETRecord
@@ -62,7 +59,7 @@ def test_create_debug_handle_to_op_node_mapping(self):
 
 
 def gen_mock_operator_graph_with_expected_map() -> Tuple[
-    OperatorGraphWithStats, Dict[int, OperatorNode]
+    OperatorGraph, Dict[int, OperatorNode]
 ]:
     # Make a mock OperatorGraph instance for testing
     node_input = ValueNode("input")
@@ -113,7 +110,7 @@ def gen_mock_operator_graph_with_expected_map() -> Tuple[
     mapping[444] = node_div
     node_output = ValueNode("output", [node_div])
     return (
-        OperatorGraphWithStats(
+        OperatorGraph(
             "mock_et_model",
             [
                 node_input,