From 39ab85ab731710d84e33ec264dd12fded36ef561 Mon Sep 17 00:00:00 2001 From: Prashant Rawat Date: Thu, 27 Aug 2026 13:21:40 -0700 Subject: [PATCH] Run the output shape matcher for delegates Summary: The Inspector's shape/dtype output matcher was gated behind `num_outputs == 1`. Dropping that one term is the whole functional change: if ( - num_outputs == 1 - and len(runtime_intermediate_output) > 1 + len(runtime_intermediate_output) > 1 and isinstance(aot_intermediate_output, torch.Tensor) ): The matcher was written for multi-output aten ops (`native_layer_norm.out`, `native_dropout.out`), where the runtime logs several tensors but AOT captured only the primary one, and the gate confined it to exactly that. Delegates returning more than one tensor have the same problem for a different reason and were left on positional pairing. However, positional pairing is unsound here. Take a fused FFN block -- five nodes in one partition, the last two escaping it: handle node shape 10 linear_gate [1, 4, 1408] 11 linear_up [1, 4, 1408] 12 silu [1, 4, 1408] 13 mul [1, 4, 1408] escapes 14 linear_down [1, 4, 512] escapes The delegate carries the whole tuple `(10, 11, 12, 13, 14)` as its debug handle and returns two tensors, so `num_outputs` is 2 and the walk is: i=0 negative_index = -1 AOT handle[-1] = 14, linear_down [1, 4, 512] runtime outputs[-1] = mul [1, 4, 1408] mispaired i=1 negative_index = -2 AOT handle[-2] = 13, mul [1, 4, 1408] runtime outputs[-2] = linear_down [1, 4, 512] mispaired Both AOT picks are the right nodes. Only the runtime side is crossed, because the two sides index different lists: * AOT side: `_combine_aot_overlapped_intermediate_outputs` indexes the delegate's *debug-handle tuple* -- `last_int = runtime_debug_handle[negative_index]`. * runtime side: `_process_single_runtime_output` indexes the delegate's *output list* -- `negative_index = -1 * (output_index + 1)`. The handle tuple is ordered by the backend's node serialization; the output list by the partition's output spec. Nothing ties the two orders together, and nothing in the debug handles records which fused node produced which output. With a single output they trivially agree, which is why the gate hid this. With several they usually do not, and the result is either * a shape mismatch, raising `ValueError: Error computing SNR difference between tensors: The size of tensor a (N) must match the size of tensor b (M) at non-singleton dimension K`. `NumericalComparatorBase.compare` builds its rows in a plain loop, so one bad pair aborts the entire DataFrame rather than the offending row. * a coincidental shape match, and a silently wrong SNR. Letting the matcher run replaces the runtime half of each pair by content: exactly one runtime output is `[1, 4, 512]`, so the AOT `linear_down` row takes that one instead of `outputs[-1]`, and `mul` gets the other. Shape settles this example; where two outputs share a shape the matcher narrows by dtype, which is what separates the int8 and float32 tensors a quantized partition returns. Differential Revision: D117735521 --- devtools/inspector/_inspector_utils.py | 44 ++++-- .../inspector/tests/inspector_utils_test.py | 146 ++++++++++++++++++ devtools/inspector/tests/targets.bzl | 1 + 3 files changed, 175 insertions(+), 16 deletions(-) diff --git a/devtools/inspector/_inspector_utils.py b/devtools/inspector/_inspector_utils.py index 556987e4bbf..11bb2b7d949 100644 --- a/devtools/inspector/_inspector_utils.py +++ b/devtools/inspector/_inspector_utils.py @@ -791,13 +791,15 @@ def _map_sequence_aot_output( def _find_matching_runtime_output_by_shape_and_dtype( aot_intermediate_output: torch.Tensor, runtime_intermediate_output: Sequence, + fallback_index: int, ) -> Any: """ Find the runtime output that matches the AOT output shape and dtype. - Used for multi-output operations (like native_layer_norm.out, native_dropout.out). + Used for multi-output operations (like native_layer_norm.out, native_dropout.out) + and for delegates that return more than one tensor. Returns: - The matching runtime output, or runtime_intermediate_output[-1] as fallback. + The matching runtime output, or runtime_intermediate_output[fallback_index]. """ # Find all runtime outputs that match the AOT shape matching_indices = [] @@ -824,14 +826,13 @@ def _find_matching_runtime_output_by_shape_and_dtype( # Exactly one dtype match - use it (e.g., dropout case where mask is bool) return runtime_intermediate_output[dtype_matching_indices[0]] - # No unique match found, return the last element as fallback - return runtime_intermediate_output[-1] + # No unique match found, fall back to positional pairing + return runtime_intermediate_output[fallback_index] def _map_non_sequence_aot_output( aot_intermediate_output: Any, runtime_intermediate_output: Any, - num_outputs: int, negative_index: int, ) -> Any: """ @@ -843,21 +844,33 @@ def _map_non_sequence_aot_output( if not isinstance(runtime_intermediate_output, Sequence): return runtime_intermediate_output - # Use the last element of the runtime output as fallback if no match is found + # Positional pairing is the fallback if no shape/dtype match is found. aot_mapped_runtime_intermediate_output = runtime_intermediate_output[negative_index] - # delegate runtime call and AOT intermediate is not a sequence. - # For multi-output operations (like native_layer_norm.out, native_dropout.out), - # the runtime captures all outputs but AOT only captures the primary output. - # We need to find the runtime output that matches the AOT output shape and dtype. - if ( - num_outputs == 1 - and len(runtime_intermediate_output) > 1 - and isinstance(aot_intermediate_output, torch.Tensor) + # The runtime event carries every tensor the op or delegate returned, while + # the AOT side has one tensor per node. Two cases land here: + # + # one runtime output per AOT node: a multi-output aten op + # (native_layer_norm.out, native_dropout.out) whose extra outputs AOT + # never captured. + # several runtime outputs: a delegate returning more than one tensor. The + # caller walks them positionally, taking the i-th-from-last debug handle + # in the delegate's handle tuple for the i-th-from-last runtime output. + # Nothing in the debug handles records which fused node produced which + # output, so handle order and output order agree only by luck -- with one + # output it is trivially true, with several it usually is not. + # + # Both cases are answered the same way: pick the runtime output whose shape + # and dtype match the AOT tensor, and only fall back to position when that + # search cannot single one out. + if len(runtime_intermediate_output) > 1 and isinstance( + aot_intermediate_output, torch.Tensor ): aot_mapped_runtime_intermediate_output = ( _find_matching_runtime_output_by_shape_and_dtype( - aot_intermediate_output, runtime_intermediate_output + aot_intermediate_output, + runtime_intermediate_output, + fallback_index=negative_index, ) ) @@ -903,7 +916,6 @@ def _process_single_runtime_output( aot_mapped_runtime_intermediate_output = _map_non_sequence_aot_output( aot_intermediate_output, runtime_intermediate_output, - num_outputs, negative_index, ) diff --git a/devtools/inspector/tests/inspector_utils_test.py b/devtools/inspector/tests/inspector_utils_test.py index cbdc557f405..1c9eee9493c 100644 --- a/devtools/inspector/tests/inspector_utils_test.py +++ b/devtools/inspector/tests/inspector_utils_test.py @@ -6,6 +6,7 @@ # pyre-unsafe +import math import tempfile import unittest from typing import Dict, Tuple @@ -44,6 +45,7 @@ propagate_back_debug_handle, TimeScale, ) +from executorch.devtools.inspector.numerical_comparator import SNRComparator from executorch.exir import to_edge from executorch.exir.debug_handle_utils import DEBUG_HANDLE_KEY, UNSET_DEBUG_HANDLE @@ -423,6 +425,150 @@ def test_map_runtime_aot_intermediate_outputs_delegated(self): break self.assertTrue(found) + def _multi_output_delegate_fixture(self): + """A delegate whose output order does not match its debug handle order. + + A delegate is logged as one runtime event carrying every tensor it + returned, and the mapping walks those outputs positionally: the + i-th-from-last debug handle in the delegate's handle tuple is paired with + the i-th-from-last runtime output. Nothing in the debug handles records + which fused node produced which output, so the two orders line up only + when the delegate has a single output. Here they deliberately disagree, + which is the shape a real fused delegate has. + """ + aot_intermediate_outputs = { + (1,): torch.randn(1, 4, 512), + (2,): torch.randn(1, 4, 1408), + (3,): torch.randn(1, 4, 2048), + } + runtime_intermediate_outputs = { + (1, 2, 3): ( + [ + aot_intermediate_outputs[(2,)].clone(), + aot_intermediate_outputs[(3,)].clone(), + aot_intermediate_outputs[(1,)].clone(), + ], + 3, + ), + } + return aot_intermediate_outputs, runtime_intermediate_outputs + + def test_map_runtime_aot_intermediate_outputs_multi_output_delegate(self): + # Each AOT tensor must be paired with the runtime tensor of the same + # shape. Pairing by position instead hands the 512-wide AOT tensor the + # 1408-wide runtime one. + aot_intermediate_outputs, runtime_intermediate_outputs = ( + self._multi_output_delegate_fixture() + ) + + actual = map_runtime_aot_intermediate_outputs( + aot_intermediate_outputs, runtime_intermediate_outputs + ) + + self.assertEqual(len(actual), 3) + for (_, aot_output), (runtime_debug_handle, runtime_output) in actual.items(): + self.assertEqual(runtime_debug_handle, (1, 2, 3)) + self.assertEqual(aot_output.shape, runtime_output.shape) + self.assertTrue(torch.allclose(aot_output, runtime_output)) + + def test_snr_comparator_on_multi_output_delegate(self): + # The mispairing surfaced inside the comparator rather than the mapper: + # subtracting two differently shaped tensors raised + # Error computing SNR difference between tensors: The size of tensor a + # (512) must match the size of tensor b (1408) at non-singleton + # dimension 2 + # which aborted the whole comparison, not just the offending row. + aot_intermediate_outputs, runtime_intermediate_outputs = ( + self._multi_output_delegate_fixture() + ) + mapping = map_runtime_aot_intermediate_outputs( + aot_intermediate_outputs, runtime_intermediate_outputs + ) + + df = SNRComparator().compare(mapping, {}, {}) + + self.assertEqual(len(df), 3) + for gap in df["gap"]: + self.assertEqual(len(gap), 1) + self.assertFalse(math.isnan(gap[0])) + + def test_map_runtime_aot_intermediate_outputs_delegate_with_interior_nodes(self): + # A delegate that fused more nodes than it returns tensors, which is the + # usual shape once a whole block lands in one partition: + # + # 10 linear_gate -> [1, 4, 1408] + # 11 linear_up -> [1, 4, 1408] + # 12 silu -> [1, 4, 1408] + # 13 mul -> [1, 4, 1408] escapes the partition + # 14 linear_down -> [1, 4, 512] escapes the partition + # + # Handles 10-12 are interior. The two that escape are the last two in + # the handle tuple, so both positional AOT picks do land on real + # delegate outputs -- but the delegate's output list is ordered the + # other way round, because the partition's output spec is built + # independently of the backend's handle serialization order. Pairing by + # position therefore hands the 512-wide AOT tensor the 1408-wide runtime + # one and vice versa, and the comparator aborts on the shape mismatch. + aot_intermediate_outputs = { + (10,): torch.randn(1, 4, 1408), + (11,): torch.randn(1, 4, 1408), + (12,): torch.randn(1, 4, 1408), + (13,): torch.randn(1, 4, 1408), + (14,): torch.randn(1, 4, 512), + } + runtime_intermediate_outputs = { + (10, 11, 12, 13, 14): ( + [ + aot_intermediate_outputs[(14,)].clone(), + aot_intermediate_outputs[(13,)].clone(), + ], + 2, + ), + } + + actual = map_runtime_aot_intermediate_outputs( + aot_intermediate_outputs, runtime_intermediate_outputs + ) + + self.assertEqual(len(actual), 2) + for (_, aot_output), (_, runtime_output) in actual.items(): + self.assertEqual(aot_output.shape, runtime_output.shape) + self.assertTrue(torch.allclose(aot_output, runtime_output)) + + # The mismatch surfaces in the comparator, and one bad pair aborts the + # whole DataFrame rather than a single row. + df = SNRComparator().compare(actual, {}, {}) + self.assertEqual(len(df), 2) + for gap in df["gap"]: + self.assertEqual(len(gap), 1) + self.assertFalse(math.isnan(gap[0])) + + def test_map_runtime_aot_intermediate_outputs_multi_output_ambiguous_shapes(self): + # Shape and dtype cannot separate outputs that share both, so pairing + # stays positional there. Guards against the search falling back to the + # last runtime output, which would collapse every AOT tensor onto it. + aot_intermediate_outputs = { + (1,): torch.tensor([1.0, 2.0, 3.0]), + (2,): torch.tensor([4.0, 5.0, 6.0]), + } + runtime_intermediate_outputs = { + (1, 2): ( + [ + aot_intermediate_outputs[(1,)].clone(), + aot_intermediate_outputs[(2,)].clone(), + ], + 2, + ), + } + + actual = map_runtime_aot_intermediate_outputs( + aot_intermediate_outputs, runtime_intermediate_outputs + ) + + self.assertEqual(len(actual), 2) + for (_, aot_output), (_, runtime_output) in actual.items(): + self.assertTrue(torch.allclose(aot_output, runtime_output)) + def test_convert_input_to_tensor_convertible_inputs(self): # Scalar -> tensor actual_output1 = convert_to_float_tensor(5) diff --git a/devtools/inspector/tests/targets.bzl b/devtools/inspector/tests/targets.bzl index c08ac0595db..47df4727bee 100644 --- a/devtools/inspector/tests/targets.bzl +++ b/devtools/inspector/tests/targets.bzl @@ -48,6 +48,7 @@ def define_common_targets(is_fbcode = False): "//executorch/devtools/etdump:schema_flatcc", "//executorch/devtools/etrecord/tests:etrecord_test_library", "//executorch/devtools/inspector:inspector_utils", + "//executorch/devtools/inspector/numerical_comparator:lib", ], )