Convert Flatcc ETDump Profile Entries into Underlying EventBlocks

sdk/etdb/TARGETS

@@ -40,6 +40,7 @@ python_library(
         "//caffe2:torch",
         "//executorch/exir:lib",
         "//executorch/sdk/edir:et_schema",
+        "//executorch/sdk/etdump:schema_flatcc",
         "//executorch/sdk/etrecord:etrecord",
     ],
 )

sdk/etdb/inspector.py

@@ -5,8 +5,9 @@
 # LICENSE file in the root directory of this source tree.
 
 import dataclasses
+from collections import defaultdict, OrderedDict
 from dataclasses import dataclass
-from typing import Dict, List, Mapping, Optional, Union
+from typing import Dict, List, Mapping, NewType, Optional, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -15,10 +16,39 @@
 
 from executorch.sdk.edir.et_schema import OperatorGraphWithStats
 from executorch.sdk.etdb._inspector_utils import gen_graphs_from_etrecord
+from executorch.sdk.etdump.schema_flatcc import ETDumpFlatCC, ProfileEvent
 from executorch.sdk.etrecord import parse_etrecord
 from tabulate import tabulate
 
 
+# Signature of a ProfileEvent
+@dataclass(frozen=True, order=True)
+class ProfileEventSignature:
+    name: str
+    instruction_id: Optional[int]
+    delegate_id: Optional[int] = None
+    delegate_id_str: Optional[str] = None
+
+    @staticmethod
+    def _gen_from_event(event: ProfileEvent) -> "ProfileEventSignature":
+        """
+        Given a ProfileEvent, extract the fields into a signature
+
+        ProfileEvents from ETDump default to "" and -1 when the field is not populated
+        The Signature will convert these back to the intended None value
+        """
+        return ProfileEventSignature(
+            event.name,
+            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,
+        )
+
+
+# Signature of a RunData as defined by its ProfileEvents
+RunSignature = NewType("RunSignature", Tuple[ProfileEventSignature])
+
+
 @dataclass
 class PerfData:
     def __init__(self, raw: List[float]):
@@ -75,6 +105,35 @@ class Event:
     delegate_backend_name: Optional[str] = None
     debug_data: List[torch.Tensor] = dataclasses.field(default_factory=list)
 
+    @staticmethod
+    def _gen_from_profile_events(
+        signature: ProfileEventSignature, events: List[ProfileEvent]
+    ) -> "Event":
+        """
+        Given a ProfileEventSignature and a list of ProfileEvents with that signature,
+        return an Event object matching the ProfileEventSignature, with perf_data
+        populated from the list of ProfileEvents
+        """
+        delegate_debug_identifier = (
+            signature.delegate_id or signature.delegate_id_str or None
+        )
+
+        # Use the delegate identifier as the event name if delegated
+        is_delegated_op = delegate_debug_identifier is not None
+        name = signature.name if not is_delegated_op else str(delegate_debug_identifier)
+
+        perf_data = PerfData(
+            [float(event.end_time - event.start_time) / 1000 for event in events]
+        )
+
+        return Event(
+            name=name,
+            perf_data=perf_data,
+            instruction_id=signature.instruction_id,
+            delegate_debug_identifier=delegate_debug_identifier,
+            is_delegated_op=is_delegated_op,
+        )
+
 
 @dataclass
 class EventBlock:
@@ -118,6 +177,53 @@ def to_dataframe(self) -> pd.DataFrame:
         df = pd.DataFrame(data)
         return df
 
+    @staticmethod
+    def _gen_from_etdump(etdump: ETDumpFlatCC) -> List["EventBlock"]:
+        """
+        Given an etdump, generate a list of EventBlocks corresponding to the
+        contents
+        """
+
+        # Group all the RunData by the set of profile events
+        profile_run_groups: Mapping[
+            RunSignature,
+            OrderedDict[ProfileEventSignature, List[ProfileEvent]],
+        ] = defaultdict(OrderedDict)
+        for run in etdump.run_data:
+            if (run_events := run.events) is None:
+                continue
+
+            # Identify all the ProfileEventSignatures
+            profile_events: OrderedDict[
+                ProfileEventSignature, ProfileEvent
+            ] = OrderedDict()
+            for event in run_events:
+                if (profile_event := event.profile_event) is not None:
+                    signature = ProfileEventSignature._gen_from_event(profile_event)
+                    profile_events[signature] = profile_event
+
+            # Create a RunSignature from the ProfileEventSignature found
+            run_signature: RunSignature = RunSignature(tuple(profile_events.keys()))
+
+            # Update the Profile Run Groups, indexed on the RunSignature
+            run_signature_events: OrderedDict[
+                ProfileEventSignature, List[ProfileEvent]
+            ] = profile_run_groups[run_signature]
+            for event_signature, event in profile_events.items():
+                run_signature_events.setdefault(event_signature, []).append(event)
+
+        # Create EventBlocks from the Profile Run Groups
+        return [
+            EventBlock(
+                name=str(index),
+                events=[
+                    Event._gen_from_profile_events(signature, event)
+                    for signature, event in profile_events.items()
+                ],
+            )
+            for index, profile_events in enumerate(profile_run_groups.values())
+        ]
+
 
 class Inspector:
     """

sdk/etdb/tests/TARGETS

@@ -9,3 +9,12 @@ python_unittest(
         "//executorch/sdk/etdb:inspector",
     ],
 )
+
+python_unittest(
+    name = "event_blocks_test",
+    srcs = ["event_blocks_test.py"],
+    deps = [
+        "//executorch/sdk/etdb:inspector",
+        "//executorch/sdk/etdump:schema_flatcc",
+    ],
+)

sdk/etdb/tests/event_blocks_test.py

@@ -0,0 +1,205 @@
+# 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.
+
+# pyre-strict
+import unittest
+from typing import List, Optional, Tuple, Union
+
+import executorch.sdk.etdump.schema_flatcc as flatcc
+from executorch.sdk.etdb.inspector import (
+    Event,
+    EventBlock,
+    PerfData,
+    ProfileEventSignature,
+)
+from executorch.sdk.etdump.schema_flatcc import ETDumpFlatCC
+
+
+class TestEventBlock(unittest.TestCase):
+
+    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Test Helpers ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+    @staticmethod
+    def _gen_sample_profile_event(
+        name: str,
+        instruction_id: int,
+        time: Tuple[int, int],
+        delegate_debug_id: Optional[Union[int, str]] = None,
+    ) -> flatcc.ProfileEvent:
+        """
+        Helper for generating test ProfileEvents
+
+        Notably:
+        - the timestamp is specified as a tuple of two separate integers
+        - delegate_debug_id takes either the str or int representation
+        - chain_idx is auto-populated to 0
+        """
+        delegate_debug_id_int = (
+            delegate_debug_id if isinstance(delegate_debug_id, int) else -1
+        )
+        delegate_debug_id_str = (
+            delegate_debug_id if isinstance(delegate_debug_id, str) else ""
+        )
+        return flatcc.ProfileEvent(
+            name,
+            0,
+            instruction_id,
+            delegate_debug_id_int,
+            delegate_debug_id_str,
+            start_time=time[0],
+            end_time=time[1],
+        )
+
+    @staticmethod
+    def _get_sample_etdump_flatcc() -> flatcc.ETDumpFlatCC:
+        """
+        Helper for getting a sample ETDumpFlatCC object with 3 RunData:
+        - run_data_1 has a signature with just profile_1
+        - run_data_2 has the same signature with run_data_1, but differnt times
+        - run_data_3 has a signature with both (profile_1, profile_2)
+        """
+        profile_event_1 = TestEventBlock._gen_sample_profile_event(
+            name="profile_1", instruction_id=1, time=(0, 1), delegate_debug_id=100
+        )
+        run_data_1 = flatcc.RunData(
+            name="run_data_1",
+            allocators=[],
+            events=[
+                flatcc.Event(
+                    allocation_event=None,
+                    debug_event=None,
+                    profile_event=profile_event_1,
+                )
+            ],
+        )
+        profile_event_2 = TestEventBlock._gen_sample_profile_event(
+            name="profile_1", instruction_id=1, time=(2, 4), delegate_debug_id=100
+        )
+        run_data_2 = flatcc.RunData(
+            name="run_data_2",
+            allocators=[],
+            events=[
+                flatcc.Event(
+                    allocation_event=None,
+                    debug_event=None,
+                    profile_event=profile_event_2,
+                )
+            ],
+        )
+
+        profile_event_3 = TestEventBlock._gen_sample_profile_event(
+            name="profile_1", instruction_id=1, time=(5, 6), delegate_debug_id=100
+        )
+        profile_event_4 = TestEventBlock._gen_sample_profile_event(
+            name="profile_2", instruction_id=2, time=(7, 8), delegate_debug_id=100
+        )
+        run_data_3 = flatcc.RunData(
+            name="run_data_3",
+            allocators=[],
+            events=[
+                flatcc.Event(
+                    allocation_event=None,
+                    debug_event=None,
+                    profile_event=profile_event_3,
+                ),
+                flatcc.Event(
+                    allocation_event=None,
+                    debug_event=None,
+                    profile_event=profile_event_4,
+                ),
+            ],
+        )
+
+        return ETDumpFlatCC(version=0, run_data=[run_data_1, run_data_2, run_data_3])
+
+    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Tests ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+    def test_gen_from_etdump(self) -> None:
+        """
+        Test "e2e" generation of EventBlocks given an ETDump
+            - Generated via EventBock.gen_from_etdump
+
+        Specifically it tests for external correctness:
+        - Correct number of EventBlocks
+        - Correct number of Events and Raw Data values (iterations)
+        """
+        etdump: ETDumpFlatCC = TestEventBlock._get_sample_etdump_flatcc()
+        blocks: List[EventBlock] = EventBlock._gen_from_etdump(etdump)
+
+        self.assertEqual(len(blocks), 2, f"Expected 2 runs, got {len(blocks)}")
+
+        # One EventBlock should have 1 event with 2 iterations
+        # The other EventBlock should have 2 events with 1 iterations
+        run_counts = {
+            (len(block.events), len(block.events[0].perf_data.raw)) for block in blocks
+        }
+        self.assertSetEqual(run_counts, {(1, 2), (2, 1)})
+
+    def test_inspector_event_generation(self) -> None:
+        """
+        Test Inspector.Event derivation from various ProfileEvent cases
+        - Non Delegated
+        - Delegate with Int Debug ID
+        - Delegate with String Debug ID
+        """
+
+        def _test_profile_event_generation(
+            name: str,
+            instruction_id: int,
+            delegate_debug_id_int: Optional[int] = None,
+            delegate_debug_id_str: Optional[str] = None,
+        ) -> None:
+            """
+            Helper function for testing that the provided ProfileEvent fields are
+            properly translated to Inspector.ProfileEventSignature and Inspector.Event
+            """
+            delegate_debug_id = delegate_debug_id_int or delegate_debug_id_str
+            profile_event: flatcc.ProfileEvent = (
+                TestEventBlock._gen_sample_profile_event(
+                    name,
+                    instruction_id,
+                    (0, 1),
+                    delegate_debug_id,
+                )
+            )
+
+            # Test Signature Generation
+            signature = ProfileEventSignature._gen_from_event(profile_event)
+            expected_signature = ProfileEventSignature(
+                name, instruction_id, delegate_debug_id_int, delegate_debug_id_str
+            )
+            self.assertEqual(signature, expected_signature)
+
+            # Test Event Generation
+            durations = [10, 20, 30]
+            profile_events: List[flatcc.ProfileEvent] = [
+                TestEventBlock._gen_sample_profile_event(
+                    name,
+                    instruction_id,
+                    (0, 10),
+                    delegate_debug_id,
+                )
+                for time in durations
+            ]
+            event = Event._gen_from_profile_events(signature, profile_events)
+
+            is_delegated = delegate_debug_id is not None
+            expected_event = Event(
+                str(delegate_debug_id) if is_delegated else name,
+                PerfData([float(duration) / 1000 for duration in durations]),
+                instruction_id=signature.instruction_id,
+                delegate_debug_identifier=delegate_debug_id,
+                is_delegated_op=is_delegated,
+            )
+            self.assertEqual(event, expected_event)
+
+        # Non Delegated
+        _test_profile_event_generation("non-delegate", 1)
+
+        # Delegate with Int Debug ID
+        _test_profile_event_generation("delegate", 1, 100)
+
+        # Delegate with String Debug ID
+        _test_profile_event_generation("delegate", 1, None, "identifier")