From ede3bea6895639cc8bdb4e784b2ff351395ec8e8 Mon Sep 17 00:00:00 2001 From: Utkarsh Bahuguna Date: Mon, 12 Oct 2026 00:17:57 +0530 Subject: [PATCH] Support multi-method ETRecords in Inspector --- devtools/inspector/_inspector.py | 6 ++- devtools/inspector/_inspector_utils.py | 17 +++++++- devtools/inspector/tests/inspector_test.py | 47 ++++++++++++++++++++++ 3 files changed, 66 insertions(+), 4 deletions(-) diff --git a/devtools/inspector/_inspector.py b/devtools/inspector/_inspector.py index 1d5a52715c9..eefb80d3f53 100644 --- a/devtools/inspector/_inspector.py +++ b/devtools/inspector/_inspector.py @@ -58,6 +58,7 @@ gen_etdump_object, gen_graphs_from_etrecord, get_aot_debug_handle_to_op_name_mapping, + get_forward_edge_dialect_program, inflate_runtime_output, is_debug_output, is_inference_output_equal, @@ -1197,6 +1198,7 @@ def _resolve_reference_graph( debug handle backpropagation fails for "exported_program". """ resolved_graph_name = reference_graph + edge_dialect_program = get_forward_edge_dialect_program(self._etrecord) # Determine the reference graph to use if reference_graph is None or reference_graph == "exported_program": @@ -1204,7 +1206,7 @@ def _resolve_reference_graph( if self._etrecord.exported_program and propagate_back_debug_handle( self._etrecord.exported_program, self._etrecord.export_graph_id, - self._etrecord.edge_dialect_program, + edge_dialect_program, disable_debug_handle_validation, ): resolved_graph_name = "exported_program" @@ -1221,7 +1223,7 @@ def _resolve_reference_graph( ) if resolved_graph_name == "edge_dialect_exported_program": - export_program = self._etrecord.edge_dialect_program + export_program = edge_dialect_program log.info( "Using 'edge_dialect_exported_program' (Edge dialect) as reference graph" ) diff --git a/devtools/inspector/_inspector_utils.py b/devtools/inspector/_inspector_utils.py index fa38b79a358..f1e5311c461 100644 --- a/devtools/inspector/_inspector_utils.py +++ b/devtools/inspector/_inspector_utils.py @@ -296,6 +296,18 @@ def is_debug_output(value: Value) -> bool: return value.output is not None and value.output.bool_val +def get_forward_edge_dialect_program(etrecord: ETRecord) -> Optional[ExportedProgram]: + """ + Return the edge dialect program of the forward method. For multi-method + models, ETRecord stores a dict of method name -> edge dialect program, and + the Inspector only correlates the forward method. + """ + edge_dialect_program = etrecord.edge_dialect_program + if isinstance(edge_dialect_program, dict): + return edge_dialect_program.get(FORWARD) + return edge_dialect_program + + def gen_graphs_from_etrecord( etrecord: ETRecord, enable_module_hierarchy: bool = False ) -> Mapping[str, OperatorGraph]: @@ -308,9 +320,10 @@ def gen_graphs_from_etrecord( ) for name, exported_program in etrecord.graph_map.items() } - if etrecord.edge_dialect_program is not None: + edge_dialect_program = get_forward_edge_dialect_program(etrecord) + if edge_dialect_program is not None: op_graph_map[EDGE_DIALECT_GRAPH_KEY] = FXOperatorGraph.gen_operator_graph( - etrecord.edge_dialect_program.graph_module, + edge_dialect_program.graph_module, enable_module_hierarchy=enable_module_hierarchy, ) diff --git a/devtools/inspector/tests/inspector_test.py b/devtools/inspector/tests/inspector_test.py index 4c59190650c..05767e23fe0 100644 --- a/devtools/inspector/tests/inspector_test.py +++ b/devtools/inspector/tests/inspector_test.py @@ -1293,6 +1293,53 @@ def element_compare(self, a, b) -> float: ) self.assertIn("Invalid runtime debug handle", str(context.exception)) + def test_inspector_with_multi_method_etrecord(self): + class AddOne(torch.nn.Module): + def forward(self, x): + return x + 1 + + class TimesTwo(torch.nn.Module): + def forward(self, x): + return x * 2 + + example_inputs = (torch.randn(2, 2),) + et_program_manager = to_edge( + { + "forward": export(AddOne(), example_inputs), + "times_two": export(TimesTwo(), example_inputs), + }, + generate_etrecord=True, + ).to_executorch() + + with tempfile.TemporaryDirectory() as tmp_dir: + etrecord_path = os.path.join(tmp_dir, "etrecord.bin") + et_program_manager.get_etrecord().save(etrecord_path) + + with patch.object( + _inspector, "gen_etdump_object", return_value=None + ), patch.object(EventBlock, "_gen_from_etdump", return_value=[]): + # Used to raise AttributeError: 'dict' object has no attribute 'graph_module' + inspector_instance = Inspector( + etdump_path=ETDUMP_PATH, + etrecord=etrecord_path, + ) + + # Multi-method ETRecords store the edge dialect programs as a dict + self.assertIsInstance(inspector_instance._etrecord.edge_dialect_program, dict) + + # The forward method's edge dialect program is used as the reference graph + graph_module, graph_name = inspector_instance._resolve_reference_graph( + "edge_dialect_exported_program" + ) + self.assertEqual(graph_name, "edge_dialect_exported_program") + targets = [ + str(node.target) + for node in graph_module.graph.nodes + if node.op == "call_function" + ] + self.assertTrue(any("add" in target for target in targets)) + self.assertFalse(any("mul" in target for target in targets)) + def test_calculate_numeric_gap_with_reference_graph_name(self): """Test calculate_numeric_gap with the reference_graph parameter using a custom graph from graph_map.""" # Create a context manager to patch functions called by Inspector.__init__