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

@@ -18,6 +18,7 @@ python_unittest(
         "//executorch/sdk/edir:base_schema",
         "//executorch/sdk/edir:et_schema",
         "//executorch/sdk/etdump:schema",
+        "//executorch/sdk/fb:et_schema",
     ],
 )
 
@@ -34,6 +35,7 @@ python_library(
         "//executorch/sdk/edir:base_schema",
         "//executorch/sdk/edir:et_schema",
         "//executorch/sdk/etdump:schema",
+        "//executorch/sdk/fb:et_schema",
     ],
 )
 
@@ -42,7 +44,7 @@ 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",
@@ -51,5 +53,6 @@ python_binary(
         "//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
@@ -19,12 +15,9 @@
 )
 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

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,