Add helper to resolve debug handles of Event (#368)

Summary:
Pull Request resolved: #368

Add a linking function in the form of `_gen_resolve_debug_handles`

It takes the derived EventBlocks (profiling info, etc.) and links it to ETRecord ("Debug Symbols") via the provided debug_handles

 ---
Note: this function only resolves the link (updates the debug_handles field), but does not join the components themselves

Reviewed By: tarun292

Differential Revision: D49117148

fbshipit-source-id: f71a8a7d3b8dcb4e559b733946f0df171147309e

Author: Jack-Khuu

sdk/etdb/inspector.py

@@ -5,9 +5,20 @@
 # LICENSE file in the root directory of this source tree.
 
 import dataclasses
+import logging
 from collections import defaultdict, OrderedDict
 from dataclasses import dataclass
-from typing import Dict, List, Mapping, NewType, Optional, Tuple, Union
+from typing import (
+    Dict,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeAlias,
+    TypedDict,
+    Union,
+)
 
 import numpy as np
 import pandas as pd
@@ -20,6 +31,7 @@
 from executorch.sdk.etrecord import parse_etrecord
 from tabulate import tabulate
 
+log: logging.Logger = logging.getLogger(__name__)
 
 # Signature of a ProfileEvent
 @dataclass(frozen=True, order=True)
@@ -46,7 +58,19 @@ def _gen_from_event(event: ProfileEvent) -> "ProfileEventSignature":
 
 
 # Signature of a RunData as defined by its ProfileEvents
-RunSignature = NewType("RunSignature", Tuple[ProfileEventSignature])
+RunSignature: TypeAlias = Tuple[ProfileEventSignature]
+
+
+# Typing for mapping Event.delegate_debug_identifiers to debug_handle(s)
+DelegateIdentifierDebugHandleMap: TypeAlias = Union[
+    Mapping[int, Tuple[int, ...]], Mapping[str, Tuple[int, ...]]
+]
+
+# Typing for Dict containig delegate metadata
+DelegateMetadata = TypedDict(
+    "DelegateMetadata",
+    {"name": str, "delegate_map": DelegateIdentifierDebugHandleMap},
+)
 
 
 @dataclass
@@ -97,7 +121,7 @@ class Event:
     delegate_debug_identifier: Optional[Union[int, str]] = None
 
     # Debug Handles in the model graph to which this event is correlated
-    debug_handles: Optional[Union[int, List[int]]] = None
+    debug_handles: Optional[Union[int, Sequence[int]]] = None
 
     stack_trace: Dict[str, str] = dataclasses.field(default_factory=dict)
     module_hierarchy: Dict[str, Dict] = dataclasses.field(default_factory=dict)
@@ -203,7 +227,7 @@ def _gen_from_etdump(etdump: ETDumpFlatCC) -> List["EventBlock"]:
                     profile_events[signature] = profile_event
 
             # Create a RunSignature from the ProfileEventSignature found
-            run_signature: RunSignature = RunSignature(tuple(profile_events.keys()))
+            run_signature: RunSignature = tuple(profile_events.keys())
 
             # Update the Profile Run Groups, indexed on the RunSignature
             run_signature_events: OrderedDict[
@@ -224,6 +248,47 @@ def _gen_from_etdump(etdump: ETDumpFlatCC) -> List["EventBlock"]:
             for index, profile_events in enumerate(profile_run_groups.values())
         ]
 
+    def _gen_resolve_debug_handles(
+        self,
+        handle_map: Dict[int, List[int]],
+        delegate_map: Optional[Dict[int, DelegateMetadata]] = None,
+    ):
+        """
+        Given mappings from instruction id to debug handles, populate the
+        debug_handles field of all underlying events
+
+        If the event is delegated, index with the instruction_id and delegate_debug_identifier
+        to obtain the debug_handle via the delegate map
+        """
+        for event in self.events:
+            # Check for the instruction_id in handle map
+            if (
+                instruction_id := event.instruction_id
+            ) is None or instruction_id not in handle_map:
+                continue
+
+            # For non-delegated event, handles are found in handle_map
+            if (delegate_debug_id := event.delegate_debug_identifier) is None:
+                event.debug_handles = handle_map[instruction_id]
+                continue
+
+            # Check that the delegated event has a corresponding mapping
+            if (
+                delegate_map is None
+                or (delegate_metadata := delegate_map.get(instruction_id)) is None
+            ):
+                event.debug_handles = handle_map[instruction_id]
+                log.warning(
+                    f" No delegate mapping found for delegate with instruction id {event.instruction_id}"
+                )
+                continue
+
+            # For delegated events, handles are found via delegateMetadata
+            event.delegate_backend_name = delegate_metadata.get("name", "")
+            event.debug_handles = delegate_metadata.get("delegate_map", {}).get(
+                delegate_debug_id  # pyre-ignore
+            )
+
 
 class Inspector:
     """

sdk/etdb/tests/event_blocks_test.py

@@ -10,12 +10,13 @@
 
 import executorch.sdk.etdump.schema_flatcc as flatcc
 from executorch.sdk.etdb.inspector import (
+    DelegateMetadata,
     Event,
     EventBlock,
     PerfData,
     ProfileEventSignature,
 )
-from executorch.sdk.etdump.schema_flatcc import ETDumpFlatCC
+from executorch.sdk.etdump.schema_flatcc import ETDumpFlatCC, ProfileEvent
 
 
 class TestEventBlock(unittest.TestCase):
@@ -203,3 +204,87 @@ def _test_profile_event_generation(
 
         # Delegate with String Debug ID
         _test_profile_event_generation("delegate", 1, None, "identifier")
+
+    def test_gen_resolve_debug_handles(self) -> None:
+        """
+        Test that gen_resolve_debug_handles() correctly populates the EventBlock
+        """
+
+        def _gen_event_helper(events: List[ProfileEvent]) -> Event:
+            """
+            Helper function to generate an Event given a set of ProfileEvents
+            """
+            signature = ProfileEventSignature._gen_from_event(events[0])
+            return Event._gen_from_profile_events(signature, events)
+
+        # Create Test Data
+
+        # Non-Delegated
+        non_delegated_profile_events_1 = [
+            TestEventBlock._gen_sample_profile_event("non_del_1", 0, (0, 1)),
+            TestEventBlock._gen_sample_profile_event("non_del_1", 0, (0, 1)),
+        ]
+        non_delegated_profile_events_2 = [
+            TestEventBlock._gen_sample_profile_event("non_del_2", 1, (0, 1)),
+        ]
+        non_delegated_event_1 = _gen_event_helper(non_delegated_profile_events_1)
+        non_delegated_event_2 = _gen_event_helper(non_delegated_profile_events_2)
+
+        # Delegated
+        delegated_profile_events_1 = [
+            TestEventBlock._gen_sample_profile_event("del_1", 0, (0, 1), 10),
+            TestEventBlock._gen_sample_profile_event("del_1", 0, (0, 1), 10),
+            TestEventBlock._gen_sample_profile_event("del_1", 0, (0, 1), 10),
+        ]
+        delegated_profile_events_2 = [
+            TestEventBlock._gen_sample_profile_event("del_2", 2, (0, 1), 20),
+        ]
+        delegated_event_1 = _gen_event_helper(delegated_profile_events_1)
+        delegated_event_2 = _gen_event_helper(delegated_profile_events_2)
+
+        # Create Test EventBlock
+        event_block = EventBlock(
+            name="test_name_1",
+            events=[
+                non_delegated_event_1,
+                non_delegated_event_2,
+                delegated_event_1,
+                delegated_event_2,
+            ],
+        )
+
+        # Create Test Maps
+        handle_map = {0: [100], 1: [110], 2: [120]}
+        delegate_map = {
+            0: DelegateMetadata(
+                {
+                    "name": "delegate",
+                    "delegate_map": {10: (100, 1000)},
+                }
+            ),
+            2: DelegateMetadata(
+                {
+                    "name": "delegate_2",
+                    "delegate_map": {20: (200,)},
+                }
+            ),
+        }
+        event_block._gen_resolve_debug_handles(handle_map, delegate_map)
+
+        # Verify Results
+        for event in event_block.events:
+            # To satisfy type checker
+            assert event.instruction_id is not None
+            if (
+                delegate_debug_identifier := event.delegate_debug_identifier
+            ) is not None:
+                # Delegated
+                metadata = delegate_map[event.instruction_id]
+                self.assertEqual(event.delegate_backend_name, metadata["name"])
+                self.assertEqual(
+                    event.debug_handles,
+                    metadata["delegate_map"][delegate_debug_identifier],  # pyre-ignore
+                )
+            else:
+                # Non Delegated
+                self.assertEqual(event.debug_handles, handle_map[event.instruction_id])