Mapping AOT debug_handles to op names

Summary: This PR adds a function to map AOT debug handles to operator names in the Export graph. It will be used later to enhance how numerical  discrepancy results are shown, making it easier for users to understand.

Differential Revision: D77244175

Author: Juntian Liu

devtools/inspector/_inspector.py

@@ -52,6 +52,7 @@
     FORWARD,
     gen_etdump_object,
     gen_graphs_from_etrecord,
+    get_aot_debug_handle_to_op_name_mapping,
     inflate_runtime_output,
     is_debug_output,
     is_inference_output_equal,
@@ -1084,6 +1085,7 @@ def __init__(
         self._reference_outputs: Dict[str, List[ProgramOutput]] = {}
         self._enable_module_hierarchy = enable_module_hierarchy
         self._aot_intermediate_outputs: Optional[Dict[Tuple[int, ...], Any]] = None
+        self._aot_debug_handles_to_op_names: Optional[Dict[Tuple[int, ...], str]] = None
         self._consume_etrecord()
 
     def _consume_etrecord(self) -> None:
@@ -1150,6 +1152,9 @@ def _consume_etrecord(self) -> None:
             return
         export_program = self._etrecord.edge_dialect_program
         graph_module = export_program.module()
+        self._aot_debug_handles_to_op_names = get_aot_debug_handle_to_op_name_mapping(
+            graph_module
+        )
         capturer = IntermediateOutputCapturer(graph_module)
         self._aot_intermediate_outputs = capturer.run_and_capture(
             self._etrecord._representative_inputs

devtools/inspector/_inspector_utils.py

@@ -734,3 +734,27 @@ def convert_to_float_tensor(input_data: Any) -> torch.Tensor:
     if torch.isnan(input_tensor).any():
         input_tensor = torch.nan_to_num(input_tensor)
     return input_tensor
+
+
+def get_aot_debug_handle_to_op_name_mapping(
+    graph_module: torch.fx.GraphModule,
+) -> Dict[Tuple[int, ...], 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[Tuple[int, ...], str]: A dictionary mapping debug handles to operator names.
+    """
+    debug_handle_to_op_name = {}
+    for node in graph_module.graph.nodes:
+        if "debug_handle" in node.meta:
+            debug_handle = node.meta["debug_handle"]
+            # Convert the debug handle to a tuple to use as a dictionary key
+            key = (
+                (debug_handle,)
+                if isinstance(debug_handle, int)
+                else tuple(debug_handle)
+            )
+            debug_handle_to_op_name[key] = node.name
+    return debug_handle_to_op_name

devtools/inspector/tests/inspector_utils_test.py

@@ -34,6 +34,7 @@
     EDGE_DIALECT_GRAPH_KEY,
     find_populated_event,
     gen_graphs_from_etrecord,
+    get_aot_debug_handle_to_op_name_mapping,
     is_inference_output_equal,
     map_runtime_aot_intermediate_outputs,
     merge_overlapping_debug_handles,
@@ -364,6 +365,45 @@ class X:
         msg = str(cm.exception)
         self.assertIn("Cannot convert value of type", msg)
 
+    def test_get_aot_debug_handle_to_op_name_mapping_single_debug_handle(self):
+        # Create a simple graph module with one node
+        graph_module = torch.fx.GraphModule({}, torch.fx.Graph())
+        node = graph_module.graph.create_node(
+            "call_function", target=torch.mul, args=(), kwargs={}, name="op1"
+        )
+        node.meta["debug_handle"] = 1
+        debug_handle_to_op_name = get_aot_debug_handle_to_op_name_mapping(graph_module)
+        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):
+        # Create a simple graph module with two nodes
+        graph_module = torch.fx.GraphModule({}, torch.fx.Graph())
+        node1 = graph_module.graph.create_node(
+            "call_function", target=torch.mul, args=(), kwargs={}, name="op1"
+        )
+        node1.meta["debug_handle"] = (1, 2)
+        node2 = graph_module.graph.create_node(
+            "call_function", target=torch.mul, args=(), kwargs={}, name="op2"
+        )
+        node2.meta["debug_handle"] = 3
+        debug_handle_to_op_name = get_aot_debug_handle_to_op_name_mapping(graph_module)
+        expected_result = {
+            (
+                1,
+                2,
+            ): "op1",
+            (3,): "op2",
+        }
+        self.assertEqual(debug_handle_to_op_name, expected_result)
+
+    def test_get_aot_debug_handle_to_op_name_mapping_no_debug_handles(self):
+        # Create a simple graph module with no nodes
+        graph_module = torch.fx.GraphModule({}, torch.fx.Graph())
+        debug_handle_to_op_name = get_aot_debug_handle_to_op_name_mapping(graph_module)
+        expected_result = {}
+        self.assertEqual(debug_handle_to_op_name, expected_result)
+
 
 def gen_mock_operator_graph_with_expected_map() -> (
     Tuple[OperatorGraph, Dict[int, OperatorNode]]