Add schema changes for delegate profiling in etdump_flatcc (#280)

Summary:
Pull Request resolved: #280

Make a few changes to the flatbuffer+python schema to enable delegate profiling

- Remove `EventType` - we can determine the event type by simply checking which of the fields in Event is null.
- Add `instruction_id`, `delegate_debug_id_int` and `delegate_debug_id_str` to `ProfileEvent` to replace `debug_handle`

Reviewed By: Jack-Khuu

Differential Revision: D49076428

fbshipit-source-id: 332c0bf84fb65f04ff6d5d76557d67b08c5f52b3

Author: Varun Puri

exir/_serialize/_flatbuffer.py

@@ -13,7 +13,7 @@
 import tempfile
 
 from dataclasses import dataclass
-from typing import Callable, Dict, Optional, Sequence
+from typing import Callable, Dict, List, Optional, Sequence
 
 
 def _is_valid_alignment(alignment: int) -> bool:
@@ -212,7 +212,12 @@ def _flatc_compile(output_dir: str, schema_path: str, json_path: str) -> None:
     )
 
 
-def _flatc_decompile(output_dir: str, schema_path: str, bin_path: str) -> None:
+def _flatc_decompile(
+    output_dir: str,
+    schema_path: str,
+    bin_path: str,
+    flatc_additional_args: Optional[List[str]] = None,
+) -> None:
     """Deserializes binary flatbuffer data to a JSON file.
 
     Args:
@@ -223,8 +228,10 @@ def _flatc_decompile(output_dir: str, schema_path: str, bin_path: str) -> None:
         bin_path: Path to the data to deserialize, as binary data compatible
             with the schema.
     """
+    flatc_additional_args = flatc_additional_args if flatc_additional_args else []
     _run_flatc(
-        [
+        flatc_additional_args
+        + [
             "--json",
             "--defaults-json",
             "--strict-json",

sdk/etdump/TARGETS

@@ -39,7 +39,6 @@ python_library(
         "fbsource//third-party/pypi/setuptools:setuptools",
         ":schema",
         ":schema_flatcc",
-        "//executorch/exir/_serialize:_bindings",
         "//executorch/exir/_serialize:lib",
     ],
 )

sdk/etdump/etdump_schema_flatcc.fbs

@@ -82,12 +82,20 @@ table ProfileEvent {
 
   // The chain to which this instruction belongs to. For now it will always be 0,
   // as chains are not used, but left in place in case they are used in the future.
-  chain_idx:int;
+  chain_id:int;
 
-  // If this event corresponds to an operator execution then this is the debug
-  // handle that was generated by the compiler that will help us map back this
-  // operator to the source code.
-  debug_handle:int;
+  // Runtime instruction id to which this event corresponds to.
+  instruction_id:int = -1;
+
+  // If this is a delegate event then these will be the corresponding delegate
+  // debug identifiers of the event which occurred in the delegate backend. This
+  // should be the same delegate debug identifier that was generated AOT.
+
+  // Integer based delegate debug identifier.
+  delegate_debug_id_int:int = -1;
+
+  // String based delegate debug identifier.
+  delegate_debug_id_str:string;
 
   // Time at which this event started. Could be in units of time or CPU cycles.
   start_time:ulong;

sdk/etdump/schema_flatcc.py

@@ -91,8 +91,10 @@ class PROFILE_EVENT_ENUM(Enum):
 @dataclass
 class ProfileEvent:
     name: str
-    chain_idx: int
-    debug_handle: int
+    chain_id: int
+    instruction_id: int
+    delegate_debug_id_int: int
+    delegate_debug_id_str: str
     start_time: int
     end_time: int
 
@@ -118,8 +120,9 @@ class Event:
 
 @dataclass
 class RunData:
-    allocators: List[Allocator]
-    events: List[Event]
+    name: str
+    allocators: Optional[List[Allocator]]
+    events: Optional[List[Event]]
 
 
 @dataclass

sdk/etdump/serialize.py

@@ -11,12 +11,11 @@
 import tempfile
 from typing import Union
 
-# pyre-ignore[21]: Could not find module `executorch.exir._serialize._bindings`.
-import executorch.exir._serialize._bindings as bindings  # @manual=//executorch/exir/_serialize:_bindings
-
 import pkg_resources
 
 from executorch.exir._serialize._dataclass import _DataclassEncoder, _json_to_dataclass
+
+from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile
 from executorch.sdk.etdump.schema import ETDump
 from executorch.sdk.etdump.schema_flatcc import ETDumpFlatCC
 
@@ -55,8 +54,7 @@ def _convert_to_flatbuffer(etdump_json: str) -> bytes:
         with open(json_path, "wb") as json_file:
             json_file.write(etdump_json.encode("ascii"))
 
-        # pyre-ignore
-        bindings.flatc_compile(d, schema_path, json_path)
+        _flatc_compile(d, schema_path, json_path)
         output_path = os.path.join(d, "{}.etdp".format(ETDUMP_SCHEMA_NAME))
         with open(output_path, "rb") as output_file:
             return output_file.read()
@@ -71,8 +69,7 @@ def _convert_from_flatbuffer(etdump_flatbuffer: bytes) -> bytes:
         bin_path = os.path.join(d, "schema.bin")
         with open(bin_path, "wb") as bin_file:
             bin_file.write(etdump_flatbuffer)
-        # pyre-ignore
-        bindings.flatc_decompile(d, schema_path, bin_path)
+        _flatc_decompile(d, schema_path, bin_path)
         output_path = os.path.join(d, "schema.json")
         with open(output_path, "rb") as output_file:
             return output_file.read()
@@ -126,14 +123,13 @@ def _convert_to_flatcc(etdump_json: str) -> bytes:
         with open(json_path, "wb") as json_file:
             json_file.write(etdump_json.encode("ascii"))
 
-        # pyre-ignore
-        bindings.flatc_compile(d, schema_path, json_path)
+        _flatc_compile(d, schema_path, json_path)
         output_path = os.path.join(d, "{}.etdp".format(ETDUMP_FLATCC_SCHEMA_NAME))
         with open(output_path, "rb") as output_file:
             return output_file.read()
 
 
-def _convert_from_flatcc(etdump_flatbuffer: bytes) -> bytes:
+def _convert_from_flatcc(etdump_flatbuffer: bytes, size_prefixed: bool = True) -> bytes:
     with tempfile.TemporaryDirectory() as d:
         _write_schema(d, ETDUMP_FLATCC_SCHEMA_NAME)
         _write_schema(d, SCALAR_TYPE_SCHEMA_NAME)
@@ -142,8 +138,10 @@ def _convert_from_flatcc(etdump_flatbuffer: bytes) -> bytes:
         bin_path = os.path.join(d, "schema.bin")
         with open(bin_path, "wb") as bin_file:
             bin_file.write(etdump_flatbuffer)
-        # pyre-ignore
-        bindings.flatc_decompile(d, schema_path, bin_path)
+        additional_args = []
+        if size_prefixed:
+            additional_args = ["--size-prefixed"]
+        _flatc_decompile(d, schema_path, bin_path, additional_args)
         output_path = os.path.join(d, "schema.json")
         with open(output_path, "rb") as output_file:
             return output_file.read()
@@ -163,7 +161,9 @@ def serialize_to_etdump_flatcc(
     return _convert_to_flatcc(_serialize_from_etdump_to_json(etdump))
 
 
-def deserialize_from_etdump_flatcc(data: bytes) -> ETDumpFlatCC:
+def deserialize_from_etdump_flatcc(
+    data: bytes, size_prefixed: bool = True
+) -> ETDumpFlatCC:
     """
     Given an etdump binary blob (constructed using the FlatCC schema) this function will deserialize
     it and return the FlatCC python object representation of etdump.
@@ -172,4 +172,6 @@ def deserialize_from_etdump_flatcc(data: bytes) -> ETDumpFlatCC:
     Returns:
         Deserialized ETDump python object.
     """
-    return _deserialize_from_json_to_etdump_flatcc(_convert_from_flatcc(data))
+    return _deserialize_from_json_to_etdump_flatcc(
+        _convert_from_flatcc(data, size_prefixed)
+    )

sdk/etdump/tests/serialize_test.py

@@ -100,6 +100,7 @@ def get_sample_etdump_flatcc() -> flatcc.ETDumpFlatCC:
         version=0,
         run_data=[
             flatcc.RunData(
+                name="test_block",
                 allocators=[
                     flatcc.Allocator(
                         name="test_allocator",
@@ -109,8 +110,23 @@ def get_sample_etdump_flatcc() -> flatcc.ETDumpFlatCC:
                     flatcc.Event(
                         profile_event=flatcc.ProfileEvent(
                             name="test_profile_event",
-                            chain_idx=1,
-                            debug_handle=1,
+                            chain_id=1,
+                            instruction_id=1,
+                            delegate_debug_id_str="",
+                            delegate_debug_id_int=-1,
+                            start_time=1001,
+                            end_time=2002,
+                        ),
+                        allocation_event=None,
+                        debug_event=None,
+                    ),
+                    flatcc.Event(
+                        profile_event=flatcc.ProfileEvent(
+                            name="test_profile_event_delegated",
+                            chain_id=1,
+                            instruction_id=1,
+                            delegate_debug_id_str="",
+                            delegate_debug_id_int=13,
                             start_time=1001,
                             end_time=2002,
                         ),
@@ -169,7 +185,9 @@ def test_serialize(self) -> None:
         program = get_sample_etdump_flatcc()
 
         flatcc_from_py = serialize_to_etdump_flatcc(program)
-        deserialized_obj = deserialize_from_etdump_flatcc(flatcc_from_py)
+        deserialized_obj = deserialize_from_etdump_flatcc(
+            flatcc_from_py, size_prefixed=False
+        )
         self.assertEqual(
             program,
             deserialized_obj,