Add support for checking more than one output from delegate in numerical comparator (pytorch#13722)

Summary:

We add
- num outputs produced by each instruction as separate info in etrecord
- in numerical debugger, we check how many outputs are generated and for each output we find corresponding debug handle and use that to find corresponding tensor in AOT

Reviewed By: Gasoonjia

Differential Revision: D80272038

Author: kimishpatel

devtools/etrecord/_etrecord.py

@@ -52,6 +52,7 @@ class ETRecordReservedFileNames(StrEnum):
     ET_DIALECT_GRAPH_MODULE = "et_dialect_graph_module"
     DEBUG_HANDLE_MAP_NAME = "debug_handle_map"
     DELEGATE_MAP_NAME = "delegate_map"
+    INSTRUCTION_ID_TO_NUM_OUTS_MAP_NAME = "instruction_id_to_num_outs_map"
     REFERENCE_OUTPUTS = "reference_outputs"
     REPRESENTATIVE_INPUTS = "representative_inputs"
 
@@ -67,6 +68,9 @@ def __init__(
         _delegate_map: Optional[
             Dict[str, Dict[int, Dict[str, Union[str, _DelegateDebugIdentifierMap]]]]
         ] = None,
+        _instruction_id_to_num_outs_map: Optional[
+            Dict[str, Dict[int, Union[int, List[int]]]]
+        ] = None,
         _reference_outputs: Optional[Dict[str, List[ProgramOutput]]] = None,
         _representative_inputs: Optional[List[ProgramInput]] = None,
     ):
@@ -92,6 +96,7 @@ def __init__(
         self.graph_map = graph_map
         self._debug_handle_map = _debug_handle_map
         self._delegate_map = _delegate_map
+        self._instruction_id_to_num_outs_map = _instruction_id_to_num_outs_map
         self._reference_outputs = _reference_outputs
         self._representative_inputs = _representative_inputs
 
@@ -172,6 +177,12 @@ def _save_metadata(self, etrecord_zip: ZipFile) -> None:
                 json.dumps(self._delegate_map),
             )
 
+        if self._instruction_id_to_num_outs_map is not None:
+            etrecord_zip.writestr(
+                ETRecordReservedFileNames.INSTRUCTION_ID_TO_NUM_OUTS_MAP_NAME,
+                json.dumps(self._instruction_id_to_num_outs_map),
+            )
+
         if self._reference_outputs is not None:
             etrecord_zip.writestr(
                 ETRecordReservedFileNames.REFERENCE_OUTPUTS,
@@ -284,6 +295,7 @@ def add_executorch_program(
         if (
             self._debug_handle_map is not None
             or self._delegate_map is not None
+            or self._instruction_id_to_num_outs_map is not None
             or self._reference_outputs is not None
             or self._representative_inputs is not None
         ):
@@ -293,13 +305,18 @@ def add_executorch_program(
             )
 
         # Process executorch program and extract data
-        debug_handle_map, delegate_map, reference_outputs, representative_inputs = (
-            _process_executorch_program(executorch_program)
-        )
+        (
+            debug_handle_map,
+            delegate_map,
+            instruction_id_to_num_outs_map,
+            reference_outputs,
+            representative_inputs,
+        ) = _process_executorch_program(executorch_program)
 
         # Set the extracted data
         self._debug_handle_map = debug_handle_map
         self._delegate_map = delegate_map
+        self._instruction_id_to_num_outs_map = instruction_id_to_num_outs_map
         self._reference_outputs = reference_outputs
         self._representative_inputs = representative_inputs
 
@@ -593,7 +610,9 @@ def _process_executorch_program(
     executorch_program: Union[
         ExecutorchProgram, ExecutorchProgramManager, BundledProgram
     ]
-) -> tuple[Optional[Dict], Optional[Dict], Optional[Dict], Optional[List]]:
+) -> tuple[
+    Optional[Dict], Optional[Dict], Optional[Dict], Optional[Dict], Optional[List]
+]:
     """Process executorch program and return debug maps and bundled program data."""
     if isinstance(executorch_program, BundledProgram):
         reference_outputs = _get_reference_outputs(executorch_program)
@@ -602,11 +621,30 @@ def _process_executorch_program(
         debug_handle_map = executorch_program.executorch_program.debug_handle_map
         # pyre-ignore[16]: Item `None` of `typing.Union[None, exir.program._program.ExecutorchProgram, exir.program._program.ExecutorchProgramManager]` has no attribute `debug_handle_map`
         delegate_map = executorch_program.executorch_program.delegate_map
-        return debug_handle_map, delegate_map, reference_outputs, representative_inputs
+        # pyre-ignore[16]: Item `None` of `typing.Union[None, exir.program._program.ExecutorchProgram, exir.program._program.ExecutorchProgramManager]` has no attribute `instruction_id_to_num_outs_map`
+        instruction_id_to_num_outs_map = (
+            executorch_program.executorch_program.instruction_id_to_num_outs_map
+        )
+        return (
+            debug_handle_map,
+            delegate_map,
+            instruction_id_to_num_outs_map,
+            reference_outputs,
+            representative_inputs,
+        )
     else:
         debug_handle_map = executorch_program.debug_handle_map
         delegate_map = executorch_program.delegate_map
-        return debug_handle_map, delegate_map, None, None
+        instruction_id_to_num_outs_map = (
+            executorch_program.instruction_id_to_num_outs_map
+        )
+        return (
+            debug_handle_map,
+            delegate_map,
+            instruction_id_to_num_outs_map,
+            None,
+            None,
+        )
 
 
 def parse_etrecord(etrecord_path: str) -> ETRecord:  # noqa: C901
@@ -640,6 +678,7 @@ def parse_etrecord(etrecord_path: str) -> ETRecord:  # noqa: C901
     graph_map: Dict[str, ExportedProgram] = {}
     debug_handle_map = None
     delegate_map = None
+    instruction_id_to_num_outs_map = None
     exported_program = None
     edge_dialect_program = None
     reference_outputs = None
@@ -659,6 +698,12 @@ def parse_etrecord(etrecord_path: str) -> ETRecord:  # noqa: C901
             delegate_map = json.loads(
                 etrecord_zip.read(ETRecordReservedFileNames.DELEGATE_MAP_NAME)
             )
+        elif entry == ETRecordReservedFileNames.INSTRUCTION_ID_TO_NUM_OUTS_MAP_NAME:
+            instruction_id_to_num_outs_map = json.loads(
+                etrecord_zip.read(
+                    ETRecordReservedFileNames.INSTRUCTION_ID_TO_NUM_OUTS_MAP_NAME
+                )
+            )
         elif entry == ETRecordReservedFileNames.ETRECORD_IDENTIFIER:
             continue
         elif entry == ETRecordReservedFileNames.EDGE_DIALECT_EXPORTED_PROGRAM:
@@ -724,6 +769,7 @@ def parse_etrecord(etrecord_path: str) -> ETRecord:  # noqa: C901
         graph_map=graph_map,
         _debug_handle_map=debug_handle_map,
         _delegate_map=delegate_map,
+        _instruction_id_to_num_outs_map=instruction_id_to_num_outs_map,
         _reference_outputs=reference_outputs,
         _representative_inputs=representative_inputs,
         export_graph_id=export_graph_id,

devtools/etrecord/tests/etrecord_test.py

@@ -219,6 +219,10 @@ def test_etrecord_generation(self):
                 etrecord._debug_handle_map,
                 json.loads(json.dumps(et_output.debug_handle_map)),
             )
+            self.assertEqual(
+                etrecord._instruction_id_to_num_outs_map,
+                json.loads(json.dumps(et_output.instruction_id_to_num_outs_map)),
+            )
 
     def test_etrecord_generation_with_bundled_program(self):
         (

devtools/inspector/_inspector.py

@@ -317,6 +317,8 @@ class Event:
         op_type: List of op types corresponding to the event.
         delegate_debug_identifier: Supplemental identifier used in combination with instruction id.
         debug_handles: Debug handles in the model graph to which this event is correlated.
+        num_outputs: Indicates the number of outputs generated by the node.
+                     Right now only used for call_delegate nodes that output more than one tensor.
         stack_trace: A dictionary mapping the name of each associated op to its stack trace.
         module_hierarchy: A dictionary mapping the name of each associated op to its module hierarchy.
         is_delegated_op: Whether or not the event was delegated.
@@ -337,6 +339,7 @@ class Event:
     op_types: List[str] = dataclasses.field(default_factory=list)
     delegate_debug_identifier: Optional[Union[int, str]] = None
     debug_handles: Optional[Union[int, Sequence[int]]] = None
+    num_outputs: int = 1
     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
@@ -928,6 +931,7 @@ def _gen_resolve_debug_handles(
         self,
         handle_map: Dict[str, List[int]],
         delegate_map: Optional[Dict[str, DelegateMetadata]] = None,
+        instruction_id_to_num_outs_map: Dict[int, int] = None,
     ):
         """
         Given mappings from instruction id to debug handles, populate the
@@ -945,6 +949,10 @@ def _gen_resolve_debug_handles(
             if (instruction_id := str(event._instruction_id)) not in handle_map:
                 continue
 
+            num_outputs = 1
+            if instruction_id_to_num_outs_map is not None:
+                num_outputs = instruction_id_to_num_outs_map.get(instruction_id, 1)
+            event.num_outputs = num_outputs
             # For non-delegated event, handles are found in handle_map
             if (delegate_debug_id := event.delegate_debug_identifier) is None:
                 event.debug_handles = handle_map[instruction_id]
@@ -1131,6 +1139,7 @@ def _consume_etrecord(self) -> None:
                     if self._etrecord._delegate_map is not None
                     else None
                 ),
+                self._etrecord._instruction_id_to_num_outs_map[FORWARD],
             )
 
         # (2) Event Metadata Association
@@ -1196,7 +1205,7 @@ 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, List[str]]]:
+    ) -> Tuple[Dict[DebugHandle, Tuple[Any, int]], 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.
@@ -1217,12 +1226,15 @@ def _get_runtime_intermediate_outputs_and_op_names(
                     debug_handle = (debug_handle,)
                 else:
                     debug_handle = tuple(debug_handle)
-                current_entry = debug_handle_to_output.get(debug_handle, (-1, None))
+                current_entry = debug_handle_to_output.get(
+                    debug_handle, (-1, None, event.num_outputs)
+                )
                 # When event has same debug_handle, only keep the one with the largest instruction id
                 if event._instruction_id > current_entry[0]:
                     debug_handle_to_output[debug_handle] = (
                         event._instruction_id,
                         event.debug_data,
+                        event.num_outputs,
                     )
                     # TODO: One debug handle can be associated with multiple op names
                     debug_handle_to_op_names[debug_handle] = [event.name]
@@ -1231,7 +1243,7 @@ def _get_runtime_intermediate_outputs_and_op_names(
             debug_handle_to_output
         )
         return {
-            k: v[1] for k, v in debug_handle_to_output.items()
+            k: (v[1], v[2]) for k, v in debug_handle_to_output.items()
         }, debug_handle_to_op_names
 
     def to_dataframe(