[executorch][arm] Add bfloat16 support to VGF backend

backends/arm/runtime/VGFSetup.cpp

@@ -66,6 +66,7 @@ enum class FormatScalarKind {
   Uint,
   Sint,
   Float,
+  BFloat,
 };
 
 struct FormatInfo {
@@ -157,6 +158,7 @@ static uint32_t get_format_component_count(VkFormat format) {
     case VK_FORMAT_R16_UINT:
     case VK_FORMAT_R16_SINT:
     case VK_FORMAT_R16_SFLOAT:
+    case VK_FORMAT_R16_SFLOAT_FPENCODING_BFLOAT16_ARM:
     case VK_FORMAT_R32_UINT:
     case VK_FORMAT_R32_SINT:
     case VK_FORMAT_R32_SFLOAT:
@@ -209,6 +211,9 @@ static bool get_format_info(VkFormat format, FormatInfo* info) {
     case VK_FORMAT_R16_SFLOAT:
       *info = FormatInfo{1, 2, FormatScalarKind::Float};
       return true;
+    case VK_FORMAT_R16_SFLOAT_FPENCODING_BFLOAT16_ARM:
+      *info = FormatInfo{1, 2, FormatScalarKind::BFloat};
+      return true;
     case VK_FORMAT_R32_UINT:
       *info = FormatInfo{1, 4, FormatScalarKind::Uint};
       return true;
@@ -3696,6 +3701,7 @@ static uint32_t get_format_size(VkFormat format) {
     case VK_FORMAT_R16_UINT:
     case VK_FORMAT_R16_SINT:
     case VK_FORMAT_R16_SFLOAT:
+    case VK_FORMAT_R16_SFLOAT_FPENCODING_BFLOAT16_ARM:
     case VK_FORMAT_R8G8_UINT:
     case VK_FORMAT_R8G8_SINT:
       return 2;

backends/arm/scripts/vulkan_utils.sh

@@ -26,21 +26,21 @@ vulkan_sdk_arch="${ARCH}"
 # macOS and Linux x86_64 use the official LunarG SDK tarballs. Linux ARM64
 # uses a separately repackaged mirror of the same SDK version.
 if [[ "${os_name}" == "Darwin" ]]; then
-    vulkan_sdk_version="1.4.341.1"
+    vulkan_sdk_version="1.4.350.0"
     vulkan_sdk_arch="macOS"
     vulkan_sdk_url="https://sdk.lunarg.com/sdk/download/${vulkan_sdk_version}/mac/vulkansdk-macos-${vulkan_sdk_version}.zip"
-    vulkan_sdk_sha256="632cbe96c8ed6ed00c6ce25e3a7738c466134f76586e1c51f1419410d7f9042e"
+    vulkan_sdk_sha256="7acc181b8fd9b4781bf51ed086222ec95d22004b85b3d0a6683a7e48ca5a1679"
 elif [[ "${os_name}" == "Linux" ]] && [[ "${ARCH}" == "x86_64" ]]; then
-    vulkan_sdk_version="1.4.341.1"
+    vulkan_sdk_version="1.4.350.0"
     vulkan_sdk_url="https://sdk.lunarg.com/sdk/download/${vulkan_sdk_version}/linux/vulkansdk-linux-x86_64-${vulkan_sdk_version}.tar.xz"
-    vulkan_sdk_sha256="3bf0f762afb6c79bc6a9d9fb5998745ccff928800a29619b501ed9de7fd9789b"
+    vulkan_sdk_sha256="b65f068ab36263559da49d7cacd7e7b9df23824ca8b68ccc522a2b06f5725df2"
 elif [[ "${os_name}" == "Linux" ]] && ([[ "${ARCH}" == "aarch64" ]] || [[ "${ARCH}" == "arm64" ]]); then
-    vulkan_sdk_version="1.4.341.1"
+    vulkan_sdk_version="1.4.350.0"
     if [[ "${vulkan_sdk_arch}" == "arm64" ]]; then
         vulkan_sdk_arch="aarch64"
     fi
     vulkan_sdk_url="https://github.com/jakoch/vulkan-sdk-arm/releases/download/${vulkan_sdk_version}/vulkansdk-ubuntu-22.04-arm-${vulkan_sdk_version}.tar.xz"
-    vulkan_sdk_sha256="345312aee2c835e128b30653278593f899a659a7ba287c571cafb22acb708b8f"
+    vulkan_sdk_sha256="9e403d444219bb7c17e9231b580d704453e2afa30a1c2fdd568d1776dc68790b"
 else
     log_step "vulkan" "Error: only macOS and Linux are supported (detected ${os_name}); architecture must be x86-64 or aarch64/arm64"
     exit 1

backends/arm/test/models/Qwen3_VL/test_qwen3_vl_layers.py

@@ -474,6 +474,8 @@ class Qwen3VLTestCase:
     "qk_norm": TOSA_FP_TEST_CASES["qk_norm"],
 }
 
+VGF_NO_QUANT_BF16_TEST_CASES: dict[str, Qwen3VLTestCase] = TOSA_BF16_TEST_CASES
+
 
 TOSA_MXFP8_TEST_CASES: dict[str, Qwen3VLTestCase] = {
     "text_attention": Qwen3VLTestCase(
@@ -602,3 +604,25 @@ def test_qwen3_vl_vgf_no_quant(test_case: Qwen3VLTestCase):
             transform_passes=list(test_case.transform_passes),
         )
         pipeline.run()
+
+
+@common.SkipIfNoModelConverter
+@common.parametrize(
+    "test_case",
+    VGF_NO_QUANT_BF16_TEST_CASES,
+)
+def test_qwen3_vl_vgf_no_quant_bf16(test_case: Qwen3VLTestCase):
+    model, inputs = test_case.model_cls.prepare_model_and_inputs()
+    model, inputs = _to_bfloat16(model, inputs)
+    with torch.no_grad():
+        pipeline = VgfPipeline[input_t](
+            model,
+            inputs,
+            aten_op=[],
+            exir_op=[],
+            quantize=False,
+            transform_passes=list(test_case.transform_passes),
+            atol=5e-2,
+            rtol=5e-2,
+        )
+        pipeline.run()

backends/arm/test/models/phi3/test_phi3_layers.py

@@ -180,14 +180,16 @@ def _dtype_cases() -> list:
     ]
 
 
-def _vgf_dtype_cases() -> list:
+def _vgf_no_quant_dtype_cases() -> list:
+    return [
+        pytest.param(torch.float32, id="fp32"),
+        pytest.param(torch.bfloat16, id="bf16"),
+    ]
+
+
+def _vgf_quant_dtype_cases() -> list:
     return [
         pytest.param(torch.float32, id="fp32"),
-        pytest.param(
-            torch.bfloat16,
-            marks=pytest.mark.xfail(reason="BF16 runtime support not ready for VGF."),
-            id="bf16",
-        ),
     ]
 
 
@@ -233,12 +235,14 @@ def test_phi3_tosa_INT_layers(name, module_factory, input_factory):
 
 
 @common.SkipIfNoModelConverter
-@pytest.mark.parametrize("dtype", _vgf_dtype_cases())
+@pytest.mark.parametrize("dtype", _vgf_no_quant_dtype_cases())
 @pytest.mark.parametrize("name,module_factory,input_factory", _module_cases())
 def test_phi3_vgf_no_quant_layers(name, module_factory, input_factory, dtype):
     config = _phi3_config()
     module = module_factory(config).to(dtype)
     inputs = input_factory(config, dtype)
+    atol = 1e-02 if dtype == torch.bfloat16 else 1e-03
+    rtol = 1e-02 if dtype == torch.bfloat16 else 1e-03
 
     pipeline = VgfPipeline[input_t1 if len(inputs) == 1 else input_t2](
         module,
@@ -250,12 +254,14 @@ def test_phi3_vgf_no_quant_layers(name, module_factory, input_factory, dtype):
             InsertInt32CastsAfterInt64PlaceholdersPass(),
         ],
         quantize=False,
+        atol=atol,
+        rtol=rtol,
     )
     pipeline.run()
 
 
 @common.SkipIfNoModelConverter
-@pytest.mark.parametrize("dtype", _vgf_dtype_cases())
+@pytest.mark.parametrize("dtype", _vgf_quant_dtype_cases())
 @pytest.mark.parametrize("name,module_factory,input_factory", _module_cases())
 def test_phi3_vgf_quant_layers(name, module_factory, input_factory, dtype):
     config = _phi3_config()

backends/arm/test/models/test_mobilenet_v2_arm.py

@@ -191,6 +191,26 @@ def test_mv2_vgf_no_quant():
     pipeline.run()
 
 
+@common.SkipIfNoModelConverter
+def test_mv2_vgf_no_quant_bf16():
+    bf16_model = models.mobilenetv2.mobilenet_v2(
+        weights=MobileNet_V2_Weights.DEFAULT
+    ).eval()
+    bf16_model = bf16_model.to(torch.bfloat16)
+    bf16_input = normalize(torch.rand((1, 3, 224, 224))).to(torch.bfloat16)
+    pipeline = VgfPipeline[input_t](
+        bf16_model,
+        (bf16_input,),
+        aten_op=[],
+        exir_op=[],
+        use_to_edge_transform_and_lower=True,
+        quantize=False,
+        atol=2.5e-01,
+        rtol=2.5e-01,
+    )
+    pipeline.run()
+
+
 def test_mv2_tosa_INT_FP_partial_quant():
     pipeline = TosaPipelineINT[input_t](
         mv2,

backends/arm/test/models/test_resnet18.py

@@ -14,6 +14,7 @@
     EthosU85PipelineINT,
     TosaPipelineFP,
     TosaPipelineINT,
+    VgfPipeline,
 )
 
 from torchvision import transforms  # type: ignore[import-untyped]
@@ -65,6 +66,22 @@ def test_resnet_18_tosa_FP_bf16():
     pipeline.run()
 
 
+@common.SkipIfNoModelConverter
+def test_resnet_18_vgf_no_quant_bf16():
+    bf16_model = resnet18(weights=ResNet18_Weights).eval()
+    bf16_model = bf16_model.to(torch.bfloat16)
+    bf16_input = normalize(torch.rand((1, 3, 224, 224)) * 2 - 1).to(torch.bfloat16)
+    pipeline = VgfPipeline[input_t](
+        bf16_model,
+        (bf16_input,),
+        aten_op=[],
+        quantize=False,
+        atol=6e-01,
+        rtol=6e-01,
+    )
+    pipeline.run()
+
+
 @common.parametrize("per_channel_quantization", quant_test_data)
 def test_resnet_18_tosa_INT(per_channel_quantization):
     pipeline = TosaPipelineINT[input_t](

backends/arm/test/ops/test_abs.py

@@ -88,7 +88,10 @@ def test_abs_u85_INT(test_data: torch.Tensor):
     pipeline.run()
 
 
-@common.parametrize("test_data", Abs.test_parameters | Abs.test_parameters_fp16)
+@common.parametrize(
+    "test_data",
+    Abs.test_parameters | Abs.test_parameters_bf16 | Abs.test_parameters_fp16,
+)
 @common.SkipIfNoModelConverter
 def test_abs_vgf_no_quant(test_data: input_t1):
     pipeline = VgfPipeline[input_t1](

backends/arm/test/ops/test_add.py

@@ -263,7 +263,9 @@ def test_add_tensor_u85_INT_2(test_data: input_t2):
     pipeline.run()
 
 
-@common.parametrize("test_data", Add.test_data | Add.test_data_fp16)
+@common.parametrize(
+    "test_data", Add.test_data | Add.test_data_fp16 | Add.test_data_bf16
+)
 @common.SkipIfNoModelConverter
 def test_add_tensor_vgf_no_quant(test_data: input_t1):
     pipeline = VgfPipeline[input_t1](

backends/arm/test/ops/test_alias_copy.py

@@ -98,7 +98,10 @@ def test_alias_u85_INT(test_data: input_t1):
     ).run()
 
 
-@common.parametrize("test_data", AliasCopy.test_data | AliasCopy.test_data_fp16)
+@common.parametrize(
+    "test_data",
+    AliasCopy.test_data | AliasCopy.test_data_bf16 | AliasCopy.test_data_fp16,
+)
 @common.SkipIfNoModelConverter
 def test_alias_vgf_no_quant(test_data: input_t1):
     pipeline = VgfPipeline[input_t1](

backends/arm/test/ops/test_amax.py

@@ -212,7 +212,9 @@ def test_max_dim_tosa_FP_not_delegated():
     pipeline.run()
 
 
-@common.parametrize("test_data", Amax.test_data | Amax.test_data_fp16)
+@common.parametrize(
+    "test_data", Amax.test_data | Amax.test_data_fp16 | Amax.test_data_bf16
+)
 @common.SkipIfNoModelConverter
 def test_amax_vgf_no_quant(test_data: Amax.input_t):
     data, dim, keep_dims = test_data()
@@ -240,7 +242,9 @@ def test_amax_vgf_quant(test_data: Amax.input_t):
     pipeline.run()
 
 
-@common.parametrize("test_data", Max.test_data | Max.test_data_fp16)
+@common.parametrize(
+    "test_data", Max.test_data | Max.test_data_fp16 | Max.test_data_bf16
+)
 @common.SkipIfNoModelConverter
 def test_max_dim_vgf_no_quant_to_amax(test_data: Max.input_t):
     data, dim = test_data()

backends/arm/test/ops/test_amin.py

@@ -224,7 +224,9 @@ def test_min_dim_tosa_FP_not_delegated():
     pipeline.run()
 
 
-@common.parametrize("test_data", Amin.test_data | Amin.test_data_fp16)
+@common.parametrize(
+    "test_data", Amin.test_data | Amin.test_data_fp16 | Amin.test_data_bf16
+)
 @common.SkipIfNoModelConverter
 def test_amin_vgf_no_quant(test_data: Amin.input_t):
     data, dim, keep_dims = test_data()
@@ -250,7 +252,9 @@ def test_amin_vgf_quant(test_data: Amin.input_t):
     pipeline.run()
 
 
-@common.parametrize("test_data", Min.test_data | Min.test_data_fp16)
+@common.parametrize(
+    "test_data", Min.test_data | Min.test_data_fp16 | Min.test_data_bf16
+)
 @common.SkipIfNoModelConverter
 def test_min_dim_vgf_no_quant_to_amin(test_data: Min.input_t):
     data, dim = test_data()

backends/arm/test/ops/test_avg_pool2d.py

@@ -286,16 +286,31 @@ def test_avg_pool2d_16a8w_u85_INT(test_module):
     pipeline.run()
 
 
-@common.parametrize("test_module", test_modules | test_modules_fp16)
+@common.parametrize(
+    "test_module",
+    test_modules | test_modules_bf16 | test_modules_fp16,
+    xfails={
+        "kernel_3x3_stride_1_pad_1_bf16": "'Unsupported BF16 PAD constant encoding' in emulation layer. MLCE-1887."
+    },
+)
 @common.SkipIfNoModelConverter
 def test_avg_pool2d_vgf_no_quant(test_module):
     model, input_tensor = test_module()
+    match input_tensor[0].dtype:
+        case torch.bfloat16:
+            atol = 5e-3
+            rtol = 5e-3
+        case _:
+            atol = 1e-3
+            rtol = 1e-3
     pipeline = VgfPipeline[input_t](
         model,
         input_tensor,
         aten_op,
         exir_op,
         quantize=False,
+        atol=atol,
+        rtol=rtol,
     )
     pipeline.run()
 

backends/arm/test/ops/test_cat.py

@@ -201,7 +201,10 @@ def test_cat_u85_INT(test_data: Tuple):
     pipeline.run()
 
 
-@common.parametrize("test_data", Cat.test_parameters | Cat.test_parameters_fp16)
+@common.parametrize(
+    "test_data",
+    Cat.test_parameters | Cat.test_parameters_bf16 | Cat.test_parameters_fp16,
+)
 @common.SkipIfNoModelConverter
 def test_cat_vgf_no_quant(test_data: Tuple):
     pipeline = VgfPipeline[input_t1](

backends/arm/test/ops/test_ceil.py

@@ -122,7 +122,7 @@ def test_ceil_u85_INT(test_data: input_t1):
     pipeline.run()
 
 
-@common.parametrize("test_data", test_data | test_data_fp16)
+@common.parametrize("test_data", test_data | test_data_bf16 | test_data_fp16)
 @common.SkipIfNoModelConverter
 def test_ceil_vgf_no_quant(test_data: input_t1):
     module, data = test_data()

backends/arm/test/ops/test_clamp.py

@@ -230,7 +230,9 @@ def test_clamp_u85_INT_16a8w(test_data):
     pipeline.run()
 
 
-@common.parametrize("test_data", test_data_suite | test_data_suite_fp16)
+@common.parametrize(
+    "test_data", test_data_suite | test_data_suite_bf16 | test_data_suite_fp16
+)
 @common.SkipIfNoModelConverter
 def test_clamp_vgf_no_quant(test_data):
     input_tensor, min_val, max_val = test_data()
@@ -546,7 +548,10 @@ def test_clamp_u85_INT_16a8w_tensor(test_data):
 
 
 @common.parametrize(
-    "test_data", test_data_suite_tensor_FP | test_data_suite_tensor_fp16
+    "test_data",
+    test_data_suite_tensor_FP
+    | test_data_suite_tensor_bf16
+    | test_data_suite_tensor_fp16,
 )
 @common.SkipIfNoModelConverter
 def test_clamp_vgf_no_quant_tensor(test_data):

backends/arm/test/ops/test_constant_pad_nd.py

@@ -344,7 +344,9 @@ def test_constant_pad_nd_tosa_INT_a16w8(test_data: Tuple):
     pipeline.run()
 
 
-@common.parametrize("test_data", test_data_suite | test_data_suite_fp16)
+@common.parametrize(
+    "test_data", test_data_suite | test_data_suite_bf16 | test_data_suite_fp16
+)
 @common.SkipIfNoModelConverter
 def test_constant_pad_nd_vgf_no_quant(test_data: Tuple):
     inp, padding, value, mode = test_data()