Make edge dialect program a first class citizen (#508)

Summary:
Pull Request resolved: #508

As building out the Inspector APIs, we realized that we want to make edge dialect program a first class citizen of ETRecord, and also require executorch program when constructing an ETRecord (used to be optional).

Reviewed By: tarun292

Differential Revision: D49436040

fbshipit-source-id: 59986e75d3da141502066d152925119998426d4e

Author: Olivia-liu

docs/website/docs/sdk/01_generating_etrecord.md

@@ -13,7 +13,8 @@ There are two important API's users must be aware of when dealing with ETrecord:
 ```python
 generate_etrecord(
     etrecord_path: str,
-    program: Optional[Union[ExecutorchProgram, MultiMethodExecutorchProgram]] = None,
+    edge_dialect_program: ExirExportedProgram
+    executorch_program: Union[ExecutorchProgram, MultiMethodExecutorchProgram],
     export_modules: Optional[
         Dict[
             str, Union[MultiMethodExirExportedProgram, ExirExportedProgram]
@@ -23,13 +24,15 @@ generate_etrecord(
 ```
 
 Generates an ETRecord from the given objects and saves it to the given path.
-The objects that will be serialized to an ETRecord are all the graph modules present in the export_modules dict and also the graph module present in the program object, which is the closest graph module representation of what is eventually run on the device.
+The objects that will be serialized to an ETRecord are all the graph modules present in the export_modules dict, the graph module present in the edge dialect program object,
+and also the graph module present in the executorch program object, which is the closest graph module representation of what is eventually run on the device.
 
 In addition to all the graph modules we also serialize the program buffer which the users can provide to the ExecuTorch runtime to run the model.
 
 #### Parameters:
 - `etrecord_path` : Path to where the ETRecord file will be saved to.
-- `program`: ExecutorchProgram or MultiMethodExecutorchProgram for this model returned by the call to to_executorch()
+- `edge_dialect_program`: ExirExportedProgram for this model returned by the call to to_edge()
+- `executorch_program`: ExecutorchProgram or MultiMethodExecutorchProgram for this model returned by the call to to_executorch()
 - `export_modules`: Dictionary of graph modules with the key being the user provided name and the value is the corresponding exported module. The exported graph modules can be either the output of capture() or to_edge().
 
 #### Returns:

sdk/etdb/_inspector_utils.py

@@ -16,16 +16,24 @@
 from executorch.sdk.etdump.serialize import deserialize_from_etdump_flatcc
 from executorch.sdk.etrecord import ETRecord, parse_etrecord
 
+EDGE_DIALECT_GRAPH_KEY = "edge_dialect_graph_module"
+
 
 def gen_graphs_from_etrecord(
     etrecord: ETRecord,
 ) -> Mapping[str, OperatorGraphWithStats]:
-    if etrecord.graph_map is None:
-        return {}
-    return {
-        name: FXOperatorGraph.gen_operator_graph(exported_program.graph_module)
-        for name, exported_program in etrecord.graph_map.items()
-    }
+    op_graph_map = {}
+    if etrecord.graph_map is not None:
+        op_graph_map = {
+            name: FXOperatorGraph.gen_operator_graph(exported_program.graph_module)
+            for name, exported_program in etrecord.graph_map.items()
+        }
+    if etrecord.edge_dialect_program is not None:
+        op_graph_map[EDGE_DIALECT_GRAPH_KEY] = FXOperatorGraph.gen_operator_graph(
+            etrecord.edge_dialect_program.graph_module
+        )
+
+    return op_graph_map
 
 
 # TODO: use anonymous function to avoid passing the dict around

sdk/etdb/inspector.py

@@ -28,6 +28,7 @@
 from executorch.sdk.edir.et_schema import OperatorGraphWithStats, OperatorNode
 from executorch.sdk.etdb._inspector_utils import (
     create_debug_handle_to_op_node_mapping,
+    EDGE_DIALECT_GRAPH_KEY,
     gen_etdump_object,
     gen_etrecord_object,
     gen_graphs_from_etrecord,
@@ -71,7 +72,7 @@ def _gen_from_event(event: ProfileEvent) -> "ProfileEventSignature":
         The Signature will convert these back to the intended None value
         """
         return ProfileEventSignature(
-            event.name,
+            event.name or "",
             event.instruction_id if event.instruction_id != -1 else None,
             event.delegate_debug_id_int if event.delegate_debug_id_int != -1 else None,
             event.delegate_debug_id_str if event.delegate_debug_id_str != "" else None,
@@ -346,9 +347,6 @@ def _gen_resolve_debug_handles(
                 )
 
 
-EDGE_DIALECT_GRAPH_KEY = "edge_dialect_output/forward"
-
-
 class Inspector:
     """
     APIs for examining model architecture and performance stats.
@@ -455,13 +453,13 @@ def write_tensorboard_artifact(self, path: str) -> None:
         # TODO: implement
         pass
 
-    def get_exported_program(
-        self, graph: Optional[str] = EDGE_DIALECT_GRAPH_KEY
-    ) -> ExportedProgram:
+    def get_exported_program(self, graph: Optional[str] = None) -> ExportedProgram:
         """
         Access helper for ETRecord, defaults to returning Edge Dialect Program
 
         Args:
-            graph: Name of the graph to access, defaults to "edge_dialect_output/forward"
+            graph: Name of the graph to access. If None, returns the Edge Dialect Program.
         """
+        if graph is None:
+            return self._etrecord.edge_dialect_program
         return self._etrecord.graph_map.get(graph)

sdk/etdb/tests/inspector_test.py

@@ -197,6 +197,8 @@ def test_inspector_get_exported_program(self):
             EventBlock, "_gen_from_etdump"
         ), patch.object(
             inspector, "gen_graphs_from_etrecord"
+        ), patch.object(
+            inspector, "create_debug_handle_to_op_node_mapping"
         ):
             # Call the constructor of Inspector
             inspector_instance = Inspector(
@@ -209,10 +211,10 @@ def test_inspector_get_exported_program(self):
             with tempfile.TemporaryDirectory() as tmpdirname:
                 generate_etrecord(
                     tmpdirname + "/etrecord.bin",
+                    edge_output,
                     et_output,
                     {
                         "aten_dialect_output": captured_output,
-                        "edge_dialect_output": edge_output,
                     },
                 )
 

sdk/etdb/tests/inspector_utils_test.py

@@ -16,6 +16,7 @@
 )
 from executorch.sdk.etdb._inspector_utils import (
     create_debug_handle_to_op_node_mapping,
+    EDGE_DIALECT_GRAPH_KEY,
     gen_graphs_from_etrecord,
 )
 from executorch.sdk.etrecord import generate_etrecord, parse_etrecord
@@ -29,10 +30,10 @@ def test_gen_graphs_from_etrecord(self):
         with tempfile.TemporaryDirectory() as tmpdirname:
             generate_etrecord(
                 tmpdirname + "/etrecord.bin",
+                edge_output,
                 et_output,
                 {
                     "aten_dialect_output": captured_output,
-                    "edge_dialect_output": edge_output,
                 },
             )
 
@@ -42,17 +43,15 @@ def test_gen_graphs_from_etrecord(self):
 
             self.assertTrue("aten_dialect_output/forward" in graphs)
             self.assertTrue("et_dialect_graph_module/forward" in graphs)
-            self.assertTrue("edge_dialect_output/forward" in graphs)
+            self.assertTrue(EDGE_DIALECT_GRAPH_KEY in graphs)
 
             self.assertTrue(
                 isinstance(graphs["aten_dialect_output/forward"], FXOperatorGraph)
             )
             self.assertTrue(
                 isinstance(graphs["et_dialect_graph_module/forward"], FXOperatorGraph)
             )
-            self.assertTrue(
-                isinstance(graphs["edge_dialect_output/forward"], FXOperatorGraph)
-            )
+            self.assertTrue(isinstance(graphs[EDGE_DIALECT_GRAPH_KEY], FXOperatorGraph))
 
     def test_create_debug_handle_to_op_node_mapping(self):
         debug_handle_to_op_node_map = {}

sdk/etrecord/_etrecord.py

@@ -24,13 +24,15 @@
 class ETRecordReservedFileNames(str, Enum):
     ETRECORD_IDENTIFIER = "ETRECORD_V0"
     PROGRAM_BUFFER = "program_buffer"
+    EDGE_DIALECT_EXPORTED_PROGRAM = "edge_dialect_exported_program"
     ET_DIALECT_GRAPH_MODULE = "et_dialect_graph_module"
     DEBUG_HANDLE_MAP_NAME = "debug_handle_map"
     DELEGATE_MAP_NAME = "delegate_map"
 
 
 @dataclass
 class ETRecord:
+    edge_dialect_program: Optional[ExportedProgram] = None
     graph_map: Optional[Dict[str, ExportedProgram]] = None
     program_buffer: Optional[bytes] = None
     _debug_handle_map: Optional[Dict[int, Union[int, List[int]]]] = None
@@ -76,7 +78,7 @@ def _handle_export_module(
         raise RuntimeError(f"Unsupported graph module type. {type(export_module)}")
 
 
-def _handle_program(
+def _handle_executorch_program(
     etrecord_zip: ZipFile,
     program: Union[ExecutorchProgram, MultiMethodExecutorchProgram],
 ) -> None:
@@ -111,9 +113,25 @@ def _handle_program(
         )
 
 
+def _handle_edge_dialect_exported_program(
+    etrecord_zip: ZipFile, edge_dialect_exported_program: ExportedProgram
+) -> None:
+    serialized_ep, serialized_state_dict = serialize(edge_dialect_exported_program)
+
+    etrecord_zip.writestr(
+        ETRecordReservedFileNames.EDGE_DIALECT_EXPORTED_PROGRAM,
+        serialized_ep,
+    )
+    etrecord_zip.writestr(
+        f"{ETRecordReservedFileNames.EDGE_DIALECT_EXPORTED_PROGRAM}_state_dict",
+        serialized_state_dict,
+    )
+
+
 def generate_etrecord(
     etrecord_path: str,
-    program: Optional[Union[ExecutorchProgram, MultiMethodExecutorchProgram]] = None,
+    edge_dialect_program: ExirExportedProgram,
+    executorch_program: Union[ExecutorchProgram, MultiMethodExecutorchProgram],
     export_modules: Optional[
         Dict[
             str,
@@ -132,7 +150,8 @@ def generate_etrecord(
 
     Args:
         etrecord_path: Path to where the ETRecord file will be saved to.
-        program: ExecutorchProgram or MultiMethodExecutorchProgram for this model returned by the
+        edge_dialect_program: ExirExportedProgram for this model returned by the call to to_edge()
+        executorch_program: ExecutorchProgram or MultiMethodExecutorchProgram for this model returned by the
             call to to_executorch()
         export_modules: Dictionary of graph modules with the key being the user provided name and the
             value is the corresponding exported module. The exported graph modules can be either the
@@ -158,18 +177,22 @@ def generate_etrecord(
                 )
             _handle_export_module(etrecord_zip, export_module, module_name)
 
-    if program is not None:
-        _handle_program(etrecord_zip, program)
+    _handle_executorch_program(etrecord_zip, executorch_program)
 
-        etrecord_zip.writestr(
-            ETRecordReservedFileNames.DEBUG_HANDLE_MAP_NAME,
-            json.dumps(program.debug_handle_map),
-        )
+    _handle_edge_dialect_exported_program(
+        etrecord_zip,
+        edge_dialect_program.exported_program,
+    )
 
-        etrecord_zip.writestr(
-            ETRecordReservedFileNames.DELEGATE_MAP_NAME,
-            json.dumps(program.delegate_map),
-        )
+    etrecord_zip.writestr(
+        ETRecordReservedFileNames.DEBUG_HANDLE_MAP_NAME,
+        json.dumps(executorch_program.debug_handle_map),
+    )
+
+    etrecord_zip.writestr(
+        ETRecordReservedFileNames.DELEGATE_MAP_NAME,
+        json.dumps(executorch_program.delegate_map),
+    )
 
 
 def parse_etrecord(etrecord_path: str) -> ETRecord:
@@ -202,6 +225,7 @@ def parse_etrecord(etrecord_path: str) -> ETRecord:
     debug_handle_map = None
     delegate_map = None
     program_buffer = None
+    edge_dialect_program = None
 
     serialized_exported_program_files = set()
     serialized_state_dict_files = set()
@@ -218,6 +242,13 @@ def parse_etrecord(etrecord_path: str) -> ETRecord:
             program_buffer = etrecord_zip.read(ETRecordReservedFileNames.PROGRAM_BUFFER)
         elif entry == ETRecordReservedFileNames.ETRECORD_IDENTIFIER:
             continue
+        elif entry == ETRecordReservedFileNames.EDGE_DIALECT_EXPORTED_PROGRAM:
+            edge_dialect_program = deserialize(
+                etrecord_zip.read(
+                    ETRecordReservedFileNames.EDGE_DIALECT_EXPORTED_PROGRAM
+                ),
+                etrecord_zip.read(f"{entry}_state_dict"),
+            )
         else:
             if entry.endswith("state_dict"):
                 serialized_state_dict_files.add(entry)
@@ -235,6 +266,7 @@ def parse_etrecord(etrecord_path: str) -> ETRecord:
         )
 
     return ETRecord(
+        edge_dialect_program=edge_dialect_program,
         graph_map=graph_map,
         program_buffer=program_buffer,
         _debug_handle_map=debug_handle_map,

sdk/etrecord/tests/etrecord_test.py

@@ -51,10 +51,10 @@ def test_etrecord_generation(self):
         with tempfile.TemporaryDirectory() as tmpdirname:
             generate_etrecord(
                 tmpdirname + "/etrecord.bin",
+                edge_output,
                 et_output,
                 {
                     "aten_dialect_output": captured_output,
-                    "edge_dialect_output": edge_output,
                 },
             )
 
@@ -64,7 +64,7 @@ def test_etrecord_generation(self):
                 captured_output.exported_program.graph_module,
             )
             self.check_graph_closeness(
-                etrecord.graph_map["edge_dialect_output/forward"],
+                etrecord.edge_dialect_program,
                 edge_output.exported_program.graph_module,
             )
             self.check_graph_closeness(
@@ -82,7 +82,10 @@ def test_etrecord_invalid_input(self):
         with tempfile.TemporaryDirectory() as tmpdirname:
             with self.assertRaises(RuntimeError):
                 generate_etrecord(
-                    tmpdirname + "/etrecord.bin", {"fail_test_case": et_output}
+                    tmpdirname + "/etrecord.bin",
+                    edge_output,
+                    et_output,
+                    {"fail_test_case": et_output},
                 )
 
     def test_etrecord_reserved_name(self):
@@ -92,5 +95,7 @@ def test_etrecord_reserved_name(self):
                 with self.assertRaises(RuntimeError):
                     generate_etrecord(
                         tmpdirname + "/etrecord.bin",
+                        edge_output,
+                        et_output,
                         {reserved_name: captured_output.exported_program.graph_module},
                     )