Integrate the IntermediateOutputCapturer into Inspector

Differential Revision: D75828351

Pull Request resolved: #11384

Author: Juntian777

devtools/inspector/TARGETS

@@ -18,6 +18,7 @@ python_library(
         "//executorch/devtools/etdump:schema_flatcc",
         "//executorch/devtools/etrecord:etrecord",
         "//executorch/exir:lib",
+        "//executorch/devtools/inspector:intermediate_output_capturer",
     ],
 )
 

devtools/inspector/_inspector.py

@@ -59,6 +59,9 @@
     TimeScale,
     verify_debug_data_equivalence,
 )
+from executorch.devtools.inspector._intermediate_output_capturer import (
+    IntermediateOutputCapturer,
+)
 from executorch.exir import ExportedProgram
 
 
@@ -1074,6 +1077,7 @@ def __init__(
         # Key str is method name; value is list of ProgramOutputs because of list of test cases
         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._consume_etrecord()
 
     def _consume_etrecord(self) -> None:
@@ -1134,6 +1138,16 @@ def _consume_etrecord(self) -> None:
                     event_block.reference_output = self._reference_outputs[FORWARD][
                         index
                     ]
+        # Capture intermediate outputs only if _representative_inputs are provided
+        # when using bundled program to create the etrecord
+        if self._etrecord._representative_inputs is None:
+            return
+        export_program = self._etrecord.edge_dialect_program
+        graph_module = export_program.module()
+        capturer = IntermediateOutputCapturer(graph_module)
+        self._aot_intermediate_outputs = capturer.run_and_capture(
+            self._etrecord._representative_inputs
+        )
 
     def to_dataframe(
         self,

devtools/inspector/tests/inspector_test.py

@@ -6,6 +6,7 @@
 
 # pyre-unsafe
 
+import copy
 import random
 import statistics
 import tempfile
@@ -16,9 +17,13 @@
 
 from unittest.mock import patch
 
+import torch
+import torch.fx
+
 from executorch.devtools import generate_etrecord, parse_etrecord
 from executorch.devtools.debug_format.et_schema import OperatorNode
 from executorch.devtools.etdump.schema_flatcc import ProfileEvent
+from executorch.devtools.etrecord._etrecord import ETRecord
 from executorch.devtools.etrecord.tests.etrecord_test import TestETRecord
 
 from executorch.devtools.inspector import (
@@ -36,8 +41,17 @@
     ProfileEventSignature,
     TimeScale,
 )
-
-from executorch.exir import ExportedProgram
+from executorch.devtools.inspector.tests.inspector_test_utils import (
+    check_if_final_outputs_match,
+    model_registry,
+)
+from executorch.exir import (
+    EdgeCompileConfig,
+    EdgeProgramManager,
+    ExecutorchProgramManager,
+    to_edge,
+)
+from torch.export import export, ExportedProgram
 
 
 OP_TYPE = "aten::add"
@@ -452,6 +466,75 @@ def test_populate_debugging_related_fields_passes_for_consistent_events(self):
                 events=events,
             )
 
+    def test_no_capture_when_representative_inputs_are_none(self):
+        # Create a context manager to patch functions called by Inspector.__init__
+        with patch.object(
+            _inspector, "parse_etrecord", return_value=None
+        ), patch.object(
+            _inspector, "gen_etdump_object", return_value=None
+        ), patch.object(
+            EventBlock, "_gen_from_etdump"
+        ), patch.object(
+            _inspector, "gen_graphs_from_etrecord"
+        ):
+            # Call the constructor of Inspector
+            inspector_instance = Inspector(
+                etdump_path=ETDUMP_PATH,
+                etrecord=ETRECORD_PATH,
+            )
+            self.assertIsNone(inspector_instance._aot_intermediate_outputs)
+
+    def test_consume_etrecord_populates_correct_aot_intermediate_outputs(self):
+        with tempfile.NamedTemporaryFile(suffix=".bin") as tmp_file:
+            etrecord_path = tmp_file.name
+            mod = model_registry["ConvLinearModel"]()
+            input_tensor = torch.tensor(
+                [[[[1.0, 2.0], [3.0, 4.0]]]], requires_grad=True
+            )
+            aten_model: ExportedProgram = export(mod, (input_tensor,), strict=True)
+            edge_program_manager: EdgeProgramManager = to_edge(
+                aten_model, compile_config=EdgeCompileConfig(_check_ir_validity=True)
+            )
+            edge_program_manager_copy = copy.deepcopy(edge_program_manager)
+            et_program_manager: ExecutorchProgramManager = (
+                edge_program_manager.to_executorch()
+            )
+            # Generate ETRecord
+            generate_etrecord(
+                etrecord_path, edge_program_manager_copy, et_program_manager
+            )
+            original_consume_etrecord = Inspector._consume_etrecord
+            with patch.object(
+                Inspector, "_consume_etrecord", return_value=None
+            ), patch.object(
+                _inspector, "gen_etdump_object", return_value=None
+            ), patch.object(
+                EventBlock, "_gen_from_etdump"
+            ), patch.object(
+                _inspector, "gen_graphs_from_etrecord"
+            ):
+                # Call the constructor of Inspector
+                inspector_instance = Inspector(
+                    etdump_path=ETDUMP_PATH,
+                    etrecord=etrecord_path,
+                )
+                etrecord = ETRecord(
+                    edge_dialect_program=inspector_instance._etrecord.edge_dialect_program,
+                    graph_map=inspector_instance._etrecord.graph_map,
+                    _debug_handle_map=inspector_instance._etrecord._debug_handle_map,
+                    _delegate_map=inspector_instance._etrecord._delegate_map,
+                    _reference_outputs=inspector_instance._etrecord._reference_outputs,
+                    _representative_inputs=aten_model.example_inputs[0],
+                )
+                inspector_instance._etrecord = etrecord
+                Inspector._consume_etrecord = original_consume_etrecord
+                inspector_instance._consume_etrecord()
+                self.assertTrue(
+                    check_if_final_outputs_match(
+                        "ConvLinearModel", inspector_instance._aot_intermediate_outputs
+                    )
+                )
+
     def _gen_random_float_list(self) -> List[float]:
         return [random.uniform(0, 10) for _ in range(RAW_DATA_SIZE)]