Add delegate time scale converter to Inspector

Differential Revision: D55298701

Author: tarun292

sdk/inspector/_inspector.py

@@ -312,6 +312,9 @@ class Event:
     _instruction_id: Optional[int] = None
 
     _delegate_metadata_parser: Optional[Callable[[List[str]], Dict[str, Any]]] = None
+    _delegate_time_scale_converter: Optional[
+        Callable[[Union[int, str], Union[int, float]], Union[int, float]]
+    ] = None
 
     @cached_property
     def delegate_debug_metadatas(self) -> Union[List[str], Dict[str, Any]]:
@@ -391,6 +394,9 @@ def _gen_from_inference_events(
         delegate_metadata_parser: Optional[
             Callable[[List[str]], Dict[str, Any]]
         ] = None,
+        delegate_time_scale_converter: Optional[
+            Callable[[Union[int, float]], Union[int, float]]
+        ] = None,
     ) -> "Event":
         """
         Given an EventSignature and a list of Events with that signature,
@@ -411,6 +417,7 @@ def _gen_from_inference_events(
             name="",
             _instruction_id=signature.instruction_id,
             _delegate_metadata_parser=delegate_metadata_parser,
+            _delegate_time_scale_converter=delegate_time_scale_converter,
         )
 
         # Populate fields from profile events
@@ -476,14 +483,31 @@ def _populate_profiling_related_fields(
                         f"Expected exactly one profile event per InstructionEvent when generating Inspector Event, but got {len(profile_events)}"
                     )
 
+                profile_event = profile_events[0]
+
                 # Scale factor should only be applied to non-delegated ops
-                scale_factor_updated = 1 if ret_event.is_delegated_op else scale_factor
+                if (
+                    ret_event.is_delegated_op
+                    and ret_event._delegate_time_scale_converter is not None
+                ):
+                    scaled_time = ret_event._delegate_time_scale_converter(
+                        ret_event.name,
+                        profile_event.end_time,
+                        # pyre-ignore
+                    ) - ret_event._delegate_time_scale_converter(
+                        ret_event.name, profile_event.start_time
+                    )
+                elif not ret_event.is_delegated_op:
+                    scaled_time = (
+                        float(profile_event.end_time - profile_event.start_time)
+                        / scale_factor
+                    )
+                else:
+                    scaled_time = float(
+                        profile_event.end_time - profile_event.start_time
+                    )
 
-                profile_event = profile_events[0]
-                data.append(
-                    float(profile_event.end_time - profile_event.start_time)
-                    / scale_factor_updated
-                )
+                data.append(scaled_time)
                 delegate_debug_metadatas.append(
                     profile_event.delegate_debug_metadata
                     if profile_event.delegate_debug_metadata
@@ -646,6 +670,9 @@ def _gen_from_etdump(
         delegate_metadata_parser: Optional[
             Callable[[List[str]], Dict[str, Any]]
         ] = None,
+        delegate_time_scale_converter: Optional[
+            Callable[[Union[int, float]], Union[int, float]]
+        ] = None,
     ) -> List["EventBlock"]:
         """
         Given an etdump, generate a list of EventBlocks corresponding to the
@@ -743,6 +770,7 @@ class GroupedRunInstances:
                     scale_factor,
                     output_buffer,
                     delegate_metadata_parser,
+                    delegate_time_scale_converter,
                 )
                 for signature, instruction_events in run_group.items()
             ]
@@ -875,6 +903,9 @@ def __init__(
         delegate_metadata_parser: Optional[
             Callable[[List[str]], Dict[str, Any]]
         ] = None,
+        delegate_time_scale_converter: Optional[
+            Callable[[Union[int, float]], Union[int, float]]
+        ] = None,
         enable_module_hierarchy: bool = False,
     ) -> None:
         r"""
@@ -930,6 +961,7 @@ def __init__(
             self._target_time_scale,
             output_buffer,
             delegate_metadata_parser=delegate_metadata_parser,
+            delegate_time_scale_converter=delegate_time_scale_converter,
         )
 
         # Connect ETRecord to EventBlocks

sdk/inspector/tests/TARGETS

@@ -9,6 +9,7 @@ python_unittest(
         "//executorch/exir:lib",
         "//executorch/sdk:lib",
         "//executorch/sdk/debug_format:et_schema",
+        "//executorch/sdk/etdump:schema_flatcc",
         "//executorch/sdk/etrecord/tests:etrecord_test_library",
         "//executorch/sdk/inspector:inspector",
         "//executorch/sdk/inspector:lib",

sdk/inspector/tests/inspector_test.py

@@ -17,9 +17,15 @@
 from executorch.exir import ExportedProgram
 from executorch.sdk import generate_etrecord, parse_etrecord
 from executorch.sdk.debug_format.et_schema import OperatorNode
+from executorch.sdk.etdump.schema_flatcc import ProfileEvent
 from executorch.sdk.etrecord.tests.etrecord_test import TestETRecord
 
 from executorch.sdk.inspector import _inspector, Event, EventBlock, Inspector, PerfData
+from executorch.sdk.inspector._inspector import (
+    InstructionEvent,
+    InstructionEventSignature,
+    ProfileEventSignature,
+)
 
 
 OP_TYPE = "aten::add"
@@ -183,6 +189,49 @@ def test_inspector_associate_with_op_graph_nodes_multiple_debug_handles(self):
         expected_ops = ["op_0", "op_1"]
         self.assertEqual(event_with_multiple_debug_handles.op_types, expected_ops)
 
+    def test_inspector_delegate_time_scale_converter(self):
+        def time_scale_converter(event_name, time):
+            return time / 10
+
+        event = Event(
+            name="",
+            _delegate_metadata_parser=None,
+            _delegate_time_scale_converter=None,
+        )
+        event_signature = ProfileEventSignature(
+            name="",
+            instruction_id=0,
+            delegate_id_str="test_event",
+        )
+        instruction_events = [
+            InstructionEvent(
+                signature=InstructionEventSignature(0, 0),
+                profile_events=[
+                    ProfileEvent(
+                        name="test_event",
+                        chain_index=0,
+                        instruction_id=0,
+                        delegate_debug_id_int=None,
+                        delegate_debug_id_str="test_event_delegated",
+                        start_time=100,
+                        end_time=200,
+                        delegate_debug_metadata=None,
+                    )
+                ],
+            )
+        ]
+        Event._populate_profiling_related_fields(
+            event, event_signature, instruction_events, 1
+        )
+        # Value of the perf data before scaling is done.
+        self.assertEqual(event.perf_data.raw[0], 100)
+        event._delegate_time_scale_converter = time_scale_converter
+        Event._populate_profiling_related_fields(
+            event, event_signature, instruction_events, 1
+        )
+        # Value of the perf data after scaling is done. 200/10 - 100/10.
+        self.assertEqual(event.perf_data.raw[0], 10)
+
     def test_inspector_get_exported_program(self):
         # Create a context manager to patch functions called by Inspector.__init__
         with patch.object(