Updated AOT debug_handle to operator names mapping (#12366)

Summary:
Pull Request resolved: #12366

This PR updates the AOT debug handle to operator names mapping. Previously, each debug handle was mapped to a single operator name, but this update allows for multiple operator names to be associated with a single debug handle by storing them in a list.

Reviewed By: Gasoonjia

Differential Revision: D78118798

Author: Juntian Liu

devtools/inspector/_inspector.py

@@ -1159,7 +1159,7 @@ def _consume_etrecord(self) -> None:
 
     def _get_aot_intermediate_outputs_and_op_names(
         self,
-    ) -> Tuple[Dict[DebugHandle, Any], Dict[DebugHandle, str]]:
+    ) -> Tuple[Dict[DebugHandle, Any], Dict[DebugHandle, List[str]]]:
         """
         Capture intermediate outputs only if _representative_inputs are provided
         when using bundled program to create the etrecord
@@ -1180,13 +1180,13 @@ def _get_aot_intermediate_outputs_and_op_names(
     # TODO: Make it more extensible to further merge overlapping debug handles
     def _get_runtime_intermediate_outputs_and_op_names(
         self,
-    ) -> Tuple[Dict[DebugHandle, Any], Dict[DebugHandle, str]]:
+    ) -> Tuple[Dict[DebugHandle, Any], Dict[DebugHandle, List[str]]]:
         """
         Retrieve the runtime intermediate outputs(debug handles and intermediate values mappings)
         from the event blocks, along with the corresponding debug handles and op names mapping.
         """
         debug_handle_to_output = {}
-        debug_handle_to_op_name = {}
+        debug_handle_to_op_names = {}
         for event_block in self.event_blocks:
             for event in event_block.events:
                 # Skip OPERATOR_CALL events to avoid double-counting and exclude framework tax
@@ -1208,12 +1208,13 @@ def _get_runtime_intermediate_outputs_and_op_names(
                         event._instruction_id,
                         event.debug_data,
                     )
-                    debug_handle_to_op_name[debug_handle] = event.name
+                    # TODO: One debug handle can be associated with multiple op names
+                    debug_handle_to_op_names[debug_handle] = [event.name]
 
         merge_runtime_overlapping_debug_handles(debug_handle_to_output)
         return {
             k: v[1] for k, v in debug_handle_to_output.items()
-        }, debug_handle_to_op_name
+        }, debug_handle_to_op_names
 
     def to_dataframe(
         self,
@@ -1385,15 +1386,15 @@ def calculate_numeric_gap(self, distance: str = "MSE") -> pd.DataFrame:
             pd.DataFrame: A DataFrame listing corresponding operator outputs from
                           both stages and their computed numerical gaps.
         """
-        aot_intermediate_outputs, aot_debug_handle_to_op_name = (
+        aot_intermediate_outputs, aot_debug_handle_to_op_names = (
             self._get_aot_intermediate_outputs_and_op_names()
         )
-        if len(aot_intermediate_outputs) == 0 or len(aot_debug_handle_to_op_name) == 0:
+        if len(aot_intermediate_outputs) == 0 or len(aot_debug_handle_to_op_names) == 0:
             raise ValueError(
                 "Missing etrecord or missing representative inputs within etrecord, both of which are required for calculating numerical gap"
             )
         # The runtime_op_names will be used later to map runtime debug_handle to op_name
-        runtime_intermediate_outputs, runtime_debug_handle_to_op_name = (
+        runtime_intermediate_outputs, runtime_debug_handle_to_op_names = (
             self._get_runtime_intermediate_outputs_and_op_names()
         )
         mapping = map_runtime_aot_intermediate_outputs(
@@ -1419,11 +1420,11 @@ def calculate_numeric_gap(self, distance: str = "MSE") -> pd.DataFrame:
             rows.append(
                 {
                     "aot_ops": find_op_names(
-                        aot_debug_handle, aot_debug_handle_to_op_name
+                        aot_debug_handle, aot_debug_handle_to_op_names
                     ),
                     "aot_intermediate_output": aot_intermediate_output,
                     "runtime_ops": find_op_names(
-                        runtime_debug_handle, runtime_debug_handle_to_op_name
+                        runtime_debug_handle, runtime_debug_handle_to_op_names
                     ),
                     "runtime_intermediate_output": runtime_intermediate_output,
                     "gap": compare_intermediate_outputs(

devtools/inspector/_inspector_utils.py

@@ -814,13 +814,13 @@ def convert_to_float_tensor(input_data: Any) -> torch.Tensor:
 
 def get_aot_debug_handle_to_op_name_mapping(
     graph_module: torch.fx.GraphModule,
-) -> Dict[DebugHandle, str]:
+) -> Dict[DebugHandle, List[str]]:
     """
     Get a mapping from debug handle to operator name from the ETRecord edge_dialect_program's graph module.
     Parameters:
     graph_module (torch.fx.GraphModule): The graph module to get the mapping from.
     Returns:
-    Dict[DebugHandle, str]: A dictionary mapping debug handles to operator names.
+    Dict[DebugHandle, List[str]]: A dictionary mapping debug handles to operator names.
     """
     node_filters = [
         NodeFilter("debug_handle", "call_function", exclude_ops=["getitem"])
@@ -836,26 +836,29 @@ def get_aot_debug_handle_to_op_name_mapping(
                 if isinstance(debug_handle, int)
                 else tuple(debug_handle)
             )
-            debug_handle_to_op_name[key] = node.name
+            if key in debug_handle_to_op_name:
+                debug_handle_to_op_name[key].append(node.name)
+            else:
+                debug_handle_to_op_name[key] = [node.name]
     return debug_handle_to_op_name
 
 
 def find_op_names(
     target_debug_handle: DebugHandle,
-    debug_handle_to_op_name: Dict[DebugHandle, str],
+    debug_handle_to_op_names: Dict[DebugHandle, List[str]],
 ) -> List[str]:
     """
     Record the operator names only if their debug handles are part of the target debug handle.
-    The debug handles in `debug_handle_to_op_name` have undergone merging and remain unchanged,
+    The debug handles in `debug_handle_to_op_names` have undergone merging and remain unchanged,
     and this function identifies operations corresponding to these transformed handles.
     """
     dh_set = set(target_debug_handle)
     result = []
 
-    for key_tuple, op_name in debug_handle_to_op_name.items():
+    for key_tuple, op_name in debug_handle_to_op_names.items():
         # Check if key is a subset of the target_debug_handle
         if set(key_tuple).issubset(dh_set):
-            result.append(op_name)
+            result.extend(op_name)
 
     return result
 

devtools/inspector/tests/inspector_test.py

@@ -44,7 +44,7 @@
     TimeScale,
 )
 from executorch.devtools.inspector.tests.inspector_test_utils import (
-    check_if_debug_handle_to_op_name_match,
+    check_if_debug_handle_to_op_names_match,
     check_if_final_outputs_match,
     model_registry,
 )
@@ -522,17 +522,18 @@ def test_etrecord_populates_correct_aot_intermediate_outputs(self):
                     _representative_inputs=aten_model.example_inputs[0],
                 )
                 inspector_instance._etrecord = etrecord
-                aot_intermediate_outputs, aot_debug_handle_to_op_name = (
+                aot_intermediate_outputs, aot_debug_handle_to_op_names = (
                     inspector_instance._get_aot_intermediate_outputs_and_op_names()
                 )
                 self.assertTrue(
                     check_if_final_outputs_match(
                         "ConvLinearModel", aot_intermediate_outputs
                     )
                 )
+
                 self.assertTrue(
-                    check_if_debug_handle_to_op_name_match(
-                        "ConvLinearModel", aot_debug_handle_to_op_name
+                    check_if_debug_handle_to_op_names_match(
+                        "ConvLinearModel", aot_debug_handle_to_op_names
                     )
                 )
 
@@ -584,14 +585,14 @@ def test_get_runtime_intermediate_outputs_and_op_names(self):
             self.assertTrue(
                 torch.allclose(runtime_outputs[(4,)][0], torch.tensor([4.0, 5.0, 6.0]))
             )
-            self.assertEqual(op_names[(4,)], "op_3")
+            self.assertEqual(op_names[(4,)], ["op_3"])
 
             # Check that keys (5,) to (8,) are in the dictionary and have values of the correct size
             for key in range(5, 9):
                 self.assertIn((key,), runtime_outputs)
                 self.assertIn((key,), op_names)
                 self.assertEqual(runtime_outputs[(key,)][0].size(0), RAW_DATA_SIZE)
-                self.assertEqual(op_names[(key,)], f"op_{key-1}")
+                self.assertEqual(op_names[(key,)], [f"op_{key-1}"])
 
     def test_calculate_numeric_gap(self):
         # Create a context manager to patch functions called by Inspector.__init__

devtools/inspector/tests/inspector_test_utils.py

@@ -76,22 +76,22 @@ def get_expected_intermediate_outputs():
         }
 
     @staticmethod
-    def get_expected_debug_handle_to_op_name():
+    def get_expected_debug_handle_to_op_names():
         """
-        Returns the expected debug handle and op name mapping for this model for the given input.
+        Returns the expected debug handle and op names mapping for this model for the given input.
         """
         return {
-            (1,): "aten_convolution_default",
-            (2,): "aten_view_copy_default",
-            (3,): "aten_addmm_default",
-            (4,): "aten_add_tensor",
-            (5,): "aten_sub_tensor",
-            (6,): "aten_mul_tensor",
-            (7,): "aten_add_tensor_1",
-            (8,): "aten_div_tensor",
-            (9,): "aten_relu_default",
-            (10,): "aten_sigmoid_default",
-            (11,): "aten_split_with_sizes_copy_default",
+            (1,): ["aten_convolution_default"],
+            (2,): ["aten_view_copy_default"],
+            (3,): ["aten_permute_copy_default", "aten_addmm_default"],
+            (4,): ["aten_add_tensor"],
+            (5,): ["aten_sub_tensor"],
+            (6,): ["aten_mul_tensor"],
+            (7,): ["aten_add_tensor_1"],
+            (8,): ["aten_div_tensor"],
+            (9,): ["aten_relu_default"],
+            (10,): ["aten_sigmoid_default"],
+            (11,): ["aten_split_with_sizes_copy_default"],
         }
 
 
@@ -129,14 +129,14 @@ def check_if_final_outputs_match(model_name, actual_outputs_with_handles):
     return True
 
 
-def check_if_debug_handle_to_op_name_match(model_name, actual_debug_handle_to_op_name):
+def check_if_debug_handle_to_op_names_match(model_name, actual_debug_handle_to_op_name):
     """
     Checks if the actual op names match the expected op names for the specified model.
     Returns True if all match, otherwise returns False.
     """
     model_instance = model_registry[model_name]
     expected_debug_handle_to_op_name = (
-        model_instance.get_expected_debug_handle_to_op_name()
+        model_instance.get_expected_debug_handle_to_op_names()
     )
     if len(actual_debug_handle_to_op_name) != len(expected_debug_handle_to_op_name):
         return False

devtools/inspector/tests/inspector_utils_test.py

@@ -450,7 +450,7 @@ def test_get_aot_debug_handle_to_op_name_mapping_single_debug_handle(self):
         )
         node.meta["debug_handle"] = 1
         debug_handle_to_op_name = get_aot_debug_handle_to_op_name_mapping(graph_module)
-        expected_result = {(1,): "op1"}
+        expected_result = {(1,): ["op1"]}
         self.assertEqual(debug_handle_to_op_name, expected_result)
 
     def test_get_aot_debug_handle_to_op_name_mapping_multiple_debug_handles(self):
@@ -469,8 +469,8 @@ def test_get_aot_debug_handle_to_op_name_mapping_multiple_debug_handles(self):
             (
                 1,
                 2,
-            ): "op1",
-            (3,): "op2",
+            ): ["op1"],
+            (3,): ["op2"],
         }
         self.assertEqual(debug_handle_to_op_name, expected_result)
 
@@ -550,21 +550,43 @@ def test_node_op_type_mismatch(self):
 
     def test_find_op_names_empty_debug_handle(self):
         debug_handle = ()
-        debug_handle_to_op_name = {(1, 2): "op1", (3, 4): "op2"}
+        debug_handle_to_op_name = {(1, 2): ["op1"], (3, 4): ["op2"]}
         self.assertEqual(find_op_names(debug_handle, debug_handle_to_op_name), [])
 
     def test_find_op_names_no_matching_handles(self):
         debug_handle = (1, 2)
-        debug_handle_to_op_name = {(3, 4): "op1", (5, 6): "op2"}
+        debug_handle_to_op_name = {(3, 4): ["op1"], (5, 6): ["op2"]}
         self.assertEqual(find_op_names(debug_handle, debug_handle_to_op_name), [])
 
     def test_find_op_names_matching_handles(self):
         debug_handle = (1, 2, 3)
-        debug_handle_to_op_name = {(1, 2): "op1", (2, 3): "op2", (4, 5, 6): "op3"}
+        debug_handle_to_op_name = {(1, 2): ["op1"], (2, 3): ["op2"], (4, 5, 6): ["op3"]}
         self.assertEqual(
             find_op_names(debug_handle, debug_handle_to_op_name), ["op1", "op2"]
         )
 
+    def test_find_op_names_multiple_ops_single_handle(self):
+        """Test when a single debug handle maps to multiple operator names"""
+        debug_handle = (1, 2, 3)
+        debug_handle_to_op_name = {(1, 2): ["op1", "op2", "op3"], (4, 5): ["op4"]}
+        self.assertEqual(
+            find_op_names(debug_handle, debug_handle_to_op_name), ["op1", "op2", "op3"]
+        )
+
+    def test_find_op_names_mixed_single_and_multiple_ops(self):
+        """Test mix of handles with single and multiple operator names"""
+        debug_handle = (1, 2, 3, 4, 5)
+        debug_handle_to_op_name = {
+            (1, 2): ["op1"],
+            (3,): ["op2", "op3"],
+            (4,): ["op4"],
+            (5,): ["op5", "op6", "op7"],  # Multiple ops
+        }
+        self.assertEqual(
+            find_op_names(debug_handle, debug_handle_to_op_name),
+            ["op1", "op2", "op3", "op4", "op5", "op6", "op7"],
+        )
+
     def test_compare_intermediate_outputs_sequences(self):
         a = [1.0, 2.0, 3.0]
         b = [1.0, 2.5, 3.5]