Arm backend: Fix FP8 mismatch diagnostics (#23991)

Use FP32 for diagnostic statistics when either output has dtype
torch.float8_e4m3fn or torch.float8_e5m2, avoiding unsupported CPU
operations. Preserve the original tensor dtypes in saved artifacts.

Correct the dump callback argument order so reference and result tensors
are saved to the right files.

Change-Id: I854c7985ee0eb598af78f025dc520a1821b7f178
Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>

Author: YufengShi-dudu

backends/arm/test/misc/test_debug_feats.py

@@ -17,6 +17,7 @@
 from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.runner_utils import dbg_tosa_fb_to_json
+from executorch.backends.arm.test.tester.arm_tester import ArmTester
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineINT,
     TosaPipelineFP,
@@ -128,6 +129,59 @@ def test_tosa_FP_can_skip_ref_model_output_comparison(test_data: input_t1):
     assert validation_compare_stage.kwargs["inputs"] == test_data
 
 
+@common.parametrize(
+    "dtype_extension",
+    {
+        "e4m3": (torch.float8_e4m3fn, torch.float8_e4m3fn, "fp8e4m3"),
+        "e5m2": (torch.float8_e5m2, torch.float8_e5m2, "fp8e5m2"),
+        "e4m3_fp32": (torch.float8_e4m3fn, torch.float32, "fp8e4m3"),
+        "fp32_e4m3": (torch.float32, torch.float8_e4m3fn, "fp8e4m3"),
+        "e5m2_fp32": (torch.float8_e5m2, torch.float32, "fp8e5m2"),
+        "fp32_e5m2": (torch.float32, torch.float8_e5m2, "fp8e5m2"),
+    },
+)
+def test_fp8_mismatch_reports_and_saves_diagnostics(dtype_extension, tmp_path, caplog):
+    reference_dtype, result_dtype, extension = dtype_extension
+    reference = (
+        torch.tensor([1.0, 2.0, 3.0, 4.0]).reshape(1, 2, 1, 2).to(reference_dtype)
+    )
+    result = torch.tensor([1.0, 2.0, 4.0, 5.0]).reshape(1, 2, 1, 2).to(result_dtype)
+    tester = ArmTester(
+        torch.nn.Identity(),
+        (reference,),
+        common.get_vgf_compile_spec(
+            f"TOSA-1.0+FP+{extension}", custom_path=str(tmp_path)
+        ),
+    )
+
+    # The base comparator only widens inputs when the result is FP8.
+    def compare_in_fp32(reference, result, _):
+        ArmTester._assert_outputs_equal(
+            [result.float()], [reference.float()], atol=0, rtol=0
+        )
+
+    with pytest.raises(
+        AssertionError, match="Output 0 does not match reference output"
+    ):
+        tester._compare_outputs(
+            reference,
+            result,
+            atol=0,
+            rtol=0,
+            compare_callback=compare_in_fp32 if result_dtype == torch.float32 else None,
+        )
+
+    assert "ERROR DIFFERENCE" in caplog.text
+    assert "MEAN      MEDIAN    MAX       MIN" in caplog.text
+    for filename, expected in (
+        ("torch_tosa_output.pt", result),
+        ("torch_ref_output.pt", reference),
+    ):
+        saved = torch.load(tmp_path / filename, weights_only=True)
+        assert saved.dtype == expected.dtype
+        torch.testing.assert_close(saved.float(), expected.float(), atol=0, rtol=0)
+
+
 @common.parametrize("test_data", Linear.inputs)
 def test_artifact_tosa_FP(test_data: input_t1):
     model = Linear()

backends/arm/test/tester/analyze_output_utils.py

@@ -173,6 +173,11 @@ def print_error_diffs(  # noqa: C901
         result = _ensure_tensor(result_or_reference)
         reference_tensor = _ensure_tensor(reference)
 
+    fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2)
+    if result.dtype in fp8_dtypes or reference_tensor.dtype in fp8_dtypes:
+        result = result.to(torch.float32)
+        reference_tensor = reference_tensor.to(torch.float32)
+
     if result.shape != reference_tensor.shape:
         raise ValueError(
             f"Output needs to be of same shape: {result.shape} != {reference_tensor.shape}"
@@ -294,8 +299,8 @@ def print_error_diffs(  # noqa: C901
 
 def dump_error_output(
     tester: Any,
-    reference_output: TensorLike,
     stage_output: TensorLike,
+    reference_output: TensorLike,
     quantization_scale: float | None = None,
     atol: float = 1e-03,
     rtol: float = 1e-03,