Arm backend: Add FP8 support to the VGF runtime (#24028)

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

Author: YufengShi-dudu

backends/arm/runtime/VGFSetup.cpp

@@ -72,6 +72,8 @@ enum class FormatScalarKind {
   Sint,
   Float,
   BFloat,
+  Float8E4M3,
+  Float8E5M2,
 };
 
 struct FormatInfo {
@@ -134,6 +136,8 @@ static uint32_t get_format_component_count(VkFormat format) {
     case VK_FORMAT_R16_SINT:
     case VK_FORMAT_R16_SFLOAT:
     case VK_FORMAT_R16_SFLOAT_FPENCODING_BFLOAT16_ARM:
+    case VK_FORMAT_R8_SFLOAT_FPENCODING_FLOAT8E4M3_ARM:
+    case VK_FORMAT_R8_SFLOAT_FPENCODING_FLOAT8E5M2_ARM:
     case VK_FORMAT_R32_UINT:
     case VK_FORMAT_R32_SINT:
     case VK_FORMAT_R32_SFLOAT:
@@ -190,6 +194,12 @@ static bool get_format_info(VkFormat format, FormatInfo* info) {
     case VK_FORMAT_R16_SFLOAT_FPENCODING_BFLOAT16_ARM:
       *info = FormatInfo{1, 2, FormatScalarKind::BFloat};
       return true;
+    case VK_FORMAT_R8_SFLOAT_FPENCODING_FLOAT8E4M3_ARM:
+      *info = FormatInfo{1, 1, FormatScalarKind::Float8E4M3};
+      return true;
+    case VK_FORMAT_R8_SFLOAT_FPENCODING_FLOAT8E5M2_ARM:
+      *info = FormatInfo{1, 1, FormatScalarKind::Float8E5M2};
+      return true;
     case VK_FORMAT_R32_UINT:
       *info = FormatInfo{1, 4, FormatScalarKind::Uint};
       return true;
@@ -4875,6 +4885,8 @@ static uint32_t get_format_size(VkFormat format) {
     case VK_FORMAT_R8_BOOL_ARM:
     case VK_FORMAT_R8_UINT:
     case VK_FORMAT_R8_SINT:
+    case VK_FORMAT_R8_SFLOAT_FPENCODING_FLOAT8E4M3_ARM:
+    case VK_FORMAT_R8_SFLOAT_FPENCODING_FLOAT8E5M2_ARM:
       return 1;
     case VK_FORMAT_R16_UINT:
     case VK_FORMAT_R16_SINT:

backends/arm/test/misc/test_runner_utils.py

@@ -10,6 +10,7 @@
 from typing import Any, cast
 
 import numpy as np
+import pytest
 import torch
 from executorch.backends.arm.test import runner_utils
 
@@ -321,3 +322,43 @@ def test_add_known_vkml_validation_filters_uses_platform_separator() -> None:
         *runner_utils.KNOWN_VKML_VALIDATION_VUIDS,
     ]
     assert filters.count(first_known_vuid) == 1
+
+
+@pytest.mark.parametrize(
+    "dtype", [torch.float8_e4m3fn, torch.float8_e5m2, torch.bfloat16]
+)
+@pytest.mark.parametrize(
+    "shape, order",
+    [
+        ((3, 4), (0, 1)),
+        ((1, 3, 4, 5), runner_utils.NHWC_ORDER),
+        ((1, 2, 3, 4, 5), runner_utils.NNHWC_ORDER),
+    ],
+)
+@pytest.mark.parametrize("symbolic", [False, True])
+@pytest.mark.parametrize("opaque", [False, True])
+def test_raw_output_preserves_dtype_bytes_and_shape(
+    monkeypatch, dtype, shape, order, symbolic, opaque
+):
+    storage_dtype = np.uint16 if dtype == torch.bfloat16 else np.uint8
+    array = np.arange(np.prod(shape), dtype=storage_dtype).reshape(
+        tuple(shape[index] for index in order)
+    )
+    inverse_order = tuple(order.index(index) for index in range(len(shape)))
+    expected = torch.from_numpy(array).view(dtype).permute(inverse_order).contiguous()
+    metadata_shape = (object(), *shape[1:]) if symbolic else shape
+    output_tensor = SimpleNamespace(
+        shape=metadata_shape, dtype=dtype, dim_order=lambda: order
+    )
+    monkeypatch.setattr(
+        runner_utils, "get_first_fake_tensor", lambda node: output_tensor
+    )
+    if opaque:
+        array = array.view(np.dtype(f"V{array.itemsize}"))
+
+    result = runner_utils.numpy_to_torch_tensor(array, None)
+
+    assert result.dtype == dtype
+    assert tuple(result.shape) == shape
+    assert result.is_contiguous()
+    torch.testing.assert_close(result.view(torch.uint8), expected.view(torch.uint8))

backends/arm/test/ops/test_matmul.py

@@ -479,6 +479,30 @@ def test_matmul_vgf_no_quant(test_case: test_case_t):
     pipeline.run()
 
 
+@common.parametrize(
+    "test_case",
+    test_suite_fp8,
+    xfails={
+        "double_input_rand_rand_3d_fp8e5m2": (
+            "PyTorch lacks the required FP8 CPU kernels for bmm.",
+            NotImplementedError,
+        ),
+    },
+)
+@common.SkipIfNoModelConverter
+def test_matmul_vgf_no_quant_fp8(test_case: test_case_t):
+    test_data = test_case()
+    pipeline = VgfPipeline[input_t](
+        test_data.module,
+        test_data.input_factory(),
+        aten_op_mm,
+        list(test_data.exir_ops),
+        quantize=False,
+        tosa_extensions=["fp8e4m3", "fp8e5m2"],
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_case", test_suite, xfails=vgf_quant_xfails)
 @common.SkipIfNoModelConverter
 def test_matmul_vgf_quant(test_case: test_case_t):

backends/arm/test/runner_utils.py

@@ -300,8 +300,14 @@ def is_concrete_shape(shape_like) -> bool:
         return all(isinstance(dim, numbers.Integral) for dim in shape_like)
 
     def to_torch_tensor() -> torch.Tensor:
-        if output_tensor.dtype == torch.bfloat16 or array.dtype.type is np.void:
-            return torch.frombuffer(array, dtype=output_tensor.dtype)
+        if (
+            output_tensor.dtype
+            in (torch.bfloat16, torch.float8_e4m3fn, torch.float8_e5m2)
+            or array.dtype.type is np.void
+        ):
+            return torch.frombuffer(array, dtype=output_tensor.dtype).reshape(
+                array.shape
+            )
         return torch.from_numpy(array)
 
     if dim_order == NHWC_ORDER: