Merge intermediate output with overlapped debug_handles

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

Differential Revision: D76304528

Author: Juntian777

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 Dict, IO, List, Mapping, Optional, Tuple, TypeAlias, Union, Any
 
 import executorch.devtools.etdump.schema_flatcc as flatcc
 
@@ -483,3 +483,30 @@ 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
+    """
+
+    # 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

@@ -35,6 +35,7 @@
     gen_graphs_from_etrecord,
     is_inference_output_equal,
     TimeScale,
+    merge_overlapping_debug_handles,
 )
 
 
@@ -217,6 +218,25 @@ 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]]