Merge intermediate output with overlapped debug_handles (#11529)

Summary:
Pull Request resolved: #11529

This PR added the function merge_overlapping_debug_handles to merge the overlapping debug handles in the intermediate outputs mapping, ensuring that only the last intermediate output is retained for each overlap.

It will be called later when do the mapping between aot_intermediate_outputs and runtime_intermediate_outputs as the first step to do pre-processing.

Reviewed By: larryliu0820

Differential Revision: D76304528

Author: Juntian Liu

devtools/inspector/_inspector_utils.py

@@ -9,7 +9,7 @@
 import math
 import sys
 from enum import Enum
-from typing import Dict, IO, List, Mapping, Optional, Tuple, TypeAlias, Union
+from typing import Any, Dict, IO, List, Mapping, Optional, Tuple, TypeAlias, Union
 
 import executorch.devtools.etdump.schema_flatcc as flatcc
 
@@ -483,3 +483,32 @@ def compare_results(
                 print("\n")
 
     return results
+
+
+def merge_overlapping_debug_handles(intermediate_outputs: Dict[Tuple[int, ...], Any]):
+    """
+    Merge overlapping debug handles int a single key
+    """
+    if not intermediate_outputs:
+        return
+    # Extract and normalize into (start, end, val)
+    intervals = [(min(key), max(key), val) for key, val in intermediate_outputs.items()]
+    intervals.sort(key=lambda x: x[0])
+
+    # Merge overlapping debug_hanldes, picking the last value
+    merged_intermediate_outputs = []
+    cur_start, cur_end, cur_val = intervals[0]
+    for start, end, val in intervals[1:]:
+        if start <= cur_end:  # Overlaps
+            if end > cur_end:  # Extend if this one goes further
+                cur_end, cur_val = end, val
+
+        else:
+            merged_intermediate_outputs.append((cur_start, cur_end, cur_val))
+            cur_start, cur_end, cur_val = start, end, val
+    merged_intermediate_outputs.append((cur_start, cur_end, cur_val))
+
+    # Clear original one and populate with merged keys (value will point to the same object)
+    intermediate_outputs.clear()
+    for start, end, val in merged_intermediate_outputs:
+        intermediate_outputs[tuple(range(start, end + 1))] = val

devtools/inspector/tests/inspector_utils_test.py

@@ -34,6 +34,7 @@
     find_populated_event,
     gen_graphs_from_etrecord,
     is_inference_output_equal,
+    merge_overlapping_debug_handles,
     TimeScale,
 )
 
@@ -217,6 +218,26 @@ def test_compare_results_uint8(self):
         self.assertGreater(calculate_snr([a], [b])[0], 30.0)
         self.assertAlmostEqual(calculate_cosine_similarity([a], [b])[0], 1.0)
 
+    def test_merge_overlapping_debug_handles(self):
+        big_tensor = torch.rand(100, 100)
+        intermediate_outputs = {
+            (1, 2, 3): "val1",
+            (2, 3, 4, 5): "val2",
+            (6, 7, 8): "val3",
+            (10, 11): "val4",
+            (11, 12): big_tensor,
+        }
+        # basic merge behavior
+        merge_overlapping_debug_handles(intermediate_outputs)
+        expected_intermediate_outputs = {
+            (1, 2, 3, 4, 5): "val2",
+            (6, 7, 8): "val3",
+            (10, 11, 12): big_tensor,
+        }
+
+        self.assertEqual(intermediate_outputs, expected_intermediate_outputs)
+        self.assertIs(expected_intermediate_outputs[(10, 11, 12)], big_tensor)
+
 
 def gen_mock_operator_graph_with_expected_map() -> (
     Tuple[OperatorGraph, Dict[int, OperatorNode]]