Populate Event attributes with op nodes metadata linked by debug handles (#401)

Summary:

The main logic change is in the __init__() of Inspector.

Differential Revision: D49326221

Author: Olivia-liu

sdk/etdb/TARGETS

@@ -41,7 +41,6 @@ python_library(
         "//executorch/exir:lib",
         "//executorch/sdk/edir:et_schema",
         "//executorch/sdk/etdump:schema_flatcc",
-        "//executorch/sdk/etrecord:etrecord",
     ],
 )
 
@@ -52,6 +51,8 @@ python_library(
     ],
     deps = [
         "//executorch/sdk/edir:et_schema",
+        "//executorch/sdk/etdump:schema_flatcc",
+        "//executorch/sdk/etdump:serialize",
         "//executorch/sdk/etrecord:etrecord",
     ],
 )

sdk/etdb/_inspector_utils.py

@@ -4,13 +4,19 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Mapping
+from typing import Dict, Mapping, Optional, Tuple
 
-from executorch.sdk.edir.et_schema import FXOperatorGraph, OperatorGraphWithStats
-from executorch.sdk.etrecord import ETRecord
+from executorch.sdk.edir.et_schema import (
+    FXOperatorGraph,
+    OperatorGraphWithStats,
+    OperatorNode,
+)
+from executorch.sdk.etdump.schema_flatcc import ETDumpFlatCC
+
+from executorch.sdk.etdump.serialize import deserialize_from_etdump_flatcc
+from executorch.sdk.etrecord import ETRecord, parse_etrecord
 
 
-# TODO: add a unittest for this function
 def gen_graphs_from_etrecord(
     etrecord: ETRecord,
 ) -> Mapping[str, OperatorGraphWithStats]:
@@ -20,3 +26,44 @@ def gen_graphs_from_etrecord(
         name: FXOperatorGraph.gen_operator_graph(exported_program.graph_module)
         for name, exported_program in etrecord.graph_map.items()
     }
+
+
+def create_debug_handle_to_op_node_mapping(
+    op_graph: OperatorGraphWithStats,
+    debug_handle_to_op_node_map: Dict[int, OperatorNode],
+) -> None:
+    """
+    Recursive function to traverse all the operator graph nodes of input op_graph and build a mapping
+    from each debug handle to the operator node that contains the debug handle in its metadata.
+    """
+    # Recursively searches through the metadata of nodes
+    for element in op_graph.elements:
+        if isinstance(element, OperatorGraphWithStats):
+            create_debug_handle_to_op_node_mapping(element, debug_handle_to_op_node_map)
+        if isinstance(element, OperatorNode) and element.metadata is not None:
+            metadata = element.metadata
+            debug_handle = metadata.get("debug_handle")
+            if debug_handle is not None:
+                existing_entry = debug_handle_to_op_node_map.get(debug_handle)
+                if existing_entry is not None:
+                    raise ValueError(
+                        f"Duplicated debug handle {str(debug_handle)} shared between {element.name} and {existing_entry.name}. "
+                        "No two op nodes of the same graph should have the same debug handle."
+                    )
+                debug_handle_to_op_node_map[debug_handle] = element
+
+
+def gen_etdump_and_etrecord_objects(
+    etdump_path: Optional[str] = None, etrecord_path: Optional[str] = None
+) -> Tuple[ETDumpFlatCC, ETRecord]:
+    # Gen event blocks from etdump
+    if etdump_path is None:
+        raise ValueError("Etdump_path must be specified.")
+    with open(etdump_path, "rb") as buff:
+        etdump = deserialize_from_etdump_flatcc(buff.read())
+        # Gen op graphs from etrecord
+        if etrecord_path is None:
+            raise ValueError("Etrecord_path must be specified.")
+        etrecord = parse_etrecord(etrecord_path=etrecord_path)
+
+        return (etdump, etrecord)

sdk/etdb/inspector.py

@@ -25,10 +25,14 @@
 import torch
 from executorch.exir import ExportedProgram
 
-from executorch.sdk.edir.et_schema import OperatorGraphWithStats
-from executorch.sdk.etdb._inspector_utils import gen_graphs_from_etrecord
+from executorch.sdk.edir.et_schema import OperatorGraphWithStats, OperatorNode
+from executorch.sdk.etdb._inspector_utils import (
+    create_debug_handle_to_op_node_mapping,
+    gen_etdump_and_etrecord_objects,
+    gen_graphs_from_etrecord,
+)
 from executorch.sdk.etdump.schema_flatcc import ETDumpFlatCC, ProfileEvent
-from executorch.sdk.etrecord import parse_etrecord
+
 from tabulate import tabulate
 
 log: logging.Logger = logging.getLogger(__name__)
@@ -112,7 +116,7 @@ class Event:
 
     name: str
     perf_data: PerfData
-    op_type: List[str] = dataclasses.field(default_factory=list)
+    op_types: List[str] = dataclasses.field(default_factory=list)
 
     # Instruction Id of the original profiling event
     instruction_id: Optional[int] = None
@@ -123,7 +127,7 @@ class Event:
     # Debug Handles in the model graph to which this event is correlated
     debug_handles: Optional[Union[int, Sequence[int]]] = None
 
-    stack_trace: Dict[str, str] = dataclasses.field(default_factory=dict)
+    stack_traces: Dict[str, str] = dataclasses.field(default_factory=dict)
     module_hierarchy: Dict[str, Dict] = dataclasses.field(default_factory=dict)
     is_delegated_op: Optional[bool] = None
     delegate_backend_name: Optional[str] = None
@@ -158,6 +162,33 @@ def _gen_from_profile_events(
             is_delegated_op=is_delegated_op,
         )
 
+    def _associate_with_op_graph_nodes(
+        self, debug_handle_to_op_node_map: Dict[int, OperatorNode]
+    ) -> None:
+        """
+        Helper function to populate the stack_traces, module_hierarchy and op_types attributes
+        based on the debug handles of this event
+        """
+        debug_handles = []
+        if self.debug_handles is None:
+            return
+
+        if isinstance(self.debug_handles, int):
+            debug_handles = [self.debug_handles]
+        elif isinstance(self.debug_handles, Sequence):
+            debug_handles = self.debug_handles
+
+        for handle in debug_handles:
+            node = debug_handle_to_op_node_map.get(handle)
+            if node is not None and node.metadata is not None:
+                self.stack_traces[node.name] = node.metadata.get("stack_trace")
+                if node.metadata:
+                    self.module_hierarchy[node.name] = node.metadata.get(
+                        "nn_module_stack"
+                    )
+                if node.op:
+                    self.op_types += [node.op]
+
 
 @dataclass
 class EventBlock:
@@ -186,11 +217,11 @@ def to_dataframe(self) -> pd.DataFrame:
             "min": [event.perf_data.min for event in self.events],
             "max": [event.perf_data.max for event in self.events],
             "median": [event.perf_data.median for event in self.events],
-            "op_type": [event.op_type for event in self.events],
+            "op_types": [event.op_types for event in self.events],
             "delegate_debug_identifier": [
                 event.delegate_debug_identifier for event in self.events
             ],
-            "stack_traces": [event.stack_trace for event in self.events],
+            "stack_traces": [event.stack_traces for event in self.events],
             "module_hierarchy": [event.module_hierarchy for event in self.events],
             "is_delegated_op": [event.is_delegated_op for event in self.events],
             "delegate_backend_name": [
@@ -290,6 +321,9 @@ def _gen_resolve_debug_handles(
             )
 
 
+EDGE_DIALECT_GRAPH_KEY = "edge_dialect_output/forward"
+
+
 class Inspector:
     """
     APIs for examining model architecture and performance stats
@@ -301,16 +335,33 @@ def __init__(
         """
         Create an inspector instance from the provided ETDump/ETRecord
         """
+        etdump, self._etrecord = gen_etdump_and_etrecord_objects(
+            etdump_path=etdump_path, etrecord_path=etrecord_path
+        )
+
+        self.event_blocks = EventBlock._gen_from_etdump(etdump)
 
-        # Gen op graphs from etrecord
-        if etrecord_path is not None:
-            self._etrecord = parse_etrecord(etrecord_path=etrecord_path)
-            self._op_graph_dict: Mapping[
-                str, OperatorGraphWithStats
-            ] = gen_graphs_from_etrecord(etrecord=self._etrecord)
+        self._op_graph_dict: Mapping[
+            str, OperatorGraphWithStats
+        ] = gen_graphs_from_etrecord(etrecord=self._etrecord)
 
-        self.event_blocks: List[EventBlock] = []
-        # TODO: create event blocks from etdump, and associate events with op graph nodes
+        # Use the delegate map from etrecord, associate debug handles with each event
+        for event_block in self.event_blocks:
+            event_block._gen_resolve_debug_handles(
+                self._etrecord._debug_handle_map.get("forward"),
+                self._etrecord._delegate_map.get("forward"),
+            )
+
+        # Traverse the edge dialect op graph to create mapping from debug_handle to op node
+        debug_handle_to_op_node_map = {}
+        create_debug_handle_to_op_node_mapping(
+            self._op_graph_dict[EDGE_DIALECT_GRAPH_KEY],
+            debug_handle_to_op_node_map,
+        )
+
+        for event_block in self.event_blocks:
+            for event in event_block.events:
+                event._associate_with_op_graph_nodes(debug_handle_to_op_node_map)
 
     def print_data_tabular(self) -> None:
         """
@@ -353,12 +404,13 @@ def write_tensorboard_artifact(self, path: str) -> None:
         # TODO: implement
         pass
 
-    # TODO: add a unittest for this function
-    def get_exported_program(self, graph: Optional[str]) -> ExportedProgram:
+    def get_exported_program(
+        self, graph: Optional[str] = EDGE_DIALECT_GRAPH_KEY
+    ) -> 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"
         """
-        if not graph:
-            return self._etrecord["edge_dialect_output/forward"]
-        else:
-            return self._etrecord.get(graph)
+        return self._etrecord.graph_map.get(graph)

sdk/etdb/tests/TARGETS

@@ -6,7 +6,11 @@ python_unittest(
     name = "inspector_test",
     srcs = ["inspector_test.py"],
     deps = [
+        "//executorch/exir:lib",
+        "//executorch/sdk/edir:et_schema",
         "//executorch/sdk/etdb:inspector",
+        "//executorch/sdk/etrecord:etrecord",
+        "//executorch/sdk/etrecord/tests:etrecord_test_library",
     ],
 )
 
@@ -18,3 +22,14 @@ python_unittest(
         "//executorch/sdk/etdump:schema_flatcc",
     ],
 )
+
+python_unittest(
+    name = "inspector_utils_test",
+    srcs = ["inspector_utils_test.py"],
+    deps = [
+        "//executorch/sdk/edir:et_schema",
+        "//executorch/sdk/etdb:inspector_utils",
+        "//executorch/sdk/etrecord:etrecord",
+        "//executorch/sdk/etrecord/tests:etrecord_test_library",
+    ],
+)