Skip to content

Commit 231b9f8

Browse files
SS-JIASS-JIA
committed
[executorch][arm] Add bfloat16 support to VGF backend
Re-land bf16 support for the Arm VGF backend, restoring the two GitHub PRs that were reverted (#20299 and #20660). bf16 support needs the `VK_FORMAT_R16_SFLOAT_FPENCODING_BFLOAT16_ARM` `VkFormat`, which was added in Vulkan headers 1.4.350. The internal headers were pinned to 1.4.341, so the original diff-train import of #20299 (D110192953) failed to compile internally and both PRs were reverted in OSS. The Vulkan 1.4.350 header import and the `//third-party/khronos` selector wiring are the diffs below this one in the stack; this diff is the backend code that depends on them. Changes: - #20299: handle bf16 in `VGFSetup.cpp` (new `FormatScalarKind::BFloat` plus `VK_FORMAT_R16_SFLOAT_FPENCODING_BFLOAT16_ARM` cases in `get_format_component_count`, `get_format_info`, and `get_format_size`), the VGF test infrastructure (`test_pipeline.py`, `runner_utils.py`, `test_matmul.py`), and the `vulkan_utils.sh` SDK bump for OSS parity. - #20660: extend bf16 test coverage across the VGF op and model tests. All source and test changes are mirrored across the `fbcode/` and `xplat/` copies of the Arm backend. ghstack-source-id: 4625b1c ghstack-comment-id: 4909629034 Pull-Request: #20776
1 parent dcc04ef commit 231b9f8

59 files changed

Lines changed: 449 additions & 102 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎backends/arm/runtime/VGFSetup.cpp‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ enum class FormatScalarKind {
6666
Uint,
6767
Sint,
6868
Float,
69+
BFloat,
6970
};
7071

7172
struct FormatInfo {
@@ -157,6 +158,7 @@ static uint32_t get_format_component_count(VkFormat format) {
157158
case VK_FORMAT_R16_UINT:
158159
case VK_FORMAT_R16_SINT:
159160
case VK_FORMAT_R16_SFLOAT:
161+
case VK_FORMAT_R16_SFLOAT_FPENCODING_BFLOAT16_ARM:
160162
case VK_FORMAT_R32_UINT:
161163
case VK_FORMAT_R32_SINT:
162164
case VK_FORMAT_R32_SFLOAT:
@@ -209,6 +211,9 @@ static bool get_format_info(VkFormat format, FormatInfo* info) {
209211
case VK_FORMAT_R16_SFLOAT:
210212
*info = FormatInfo{1, 2, FormatScalarKind::Float};
211213
return true;
214+
case VK_FORMAT_R16_SFLOAT_FPENCODING_BFLOAT16_ARM:
215+
*info = FormatInfo{1, 2, FormatScalarKind::BFloat};
216+
return true;
212217
case VK_FORMAT_R32_UINT:
213218
*info = FormatInfo{1, 4, FormatScalarKind::Uint};
214219
return true;
@@ -3696,6 +3701,7 @@ static uint32_t get_format_size(VkFormat format) {
36963701
case VK_FORMAT_R16_UINT:
36973702
case VK_FORMAT_R16_SINT:
36983703
case VK_FORMAT_R16_SFLOAT:
3704+
case VK_FORMAT_R16_SFLOAT_FPENCODING_BFLOAT16_ARM:
36993705
case VK_FORMAT_R8G8_UINT:
37003706
case VK_FORMAT_R8G8_SINT:
37013707
return 2;

‎backends/arm/scripts/vulkan_utils.sh‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -26,21 +26,21 @@ vulkan_sdk_arch="${ARCH}"
2626
# macOS and Linux x86_64 use the official LunarG SDK tarballs. Linux ARM64
2727
# uses a separately repackaged mirror of the same SDK version.
2828
if [[ "${os_name}" == "Darwin" ]]; then
29-
vulkan_sdk_version="1.4.341.1"
29+
vulkan_sdk_version="1.4.350.0"
3030
vulkan_sdk_arch="macOS"
3131
vulkan_sdk_url="https://sdk.lunarg.com/sdk/download/${vulkan_sdk_version}/mac/vulkansdk-macos-${vulkan_sdk_version}.zip"
32-
vulkan_sdk_sha256="632cbe96c8ed6ed00c6ce25e3a7738c466134f76586e1c51f1419410d7f9042e"
32+
vulkan_sdk_sha256="7acc181b8fd9b4781bf51ed086222ec95d22004b85b3d0a6683a7e48ca5a1679"
3333
elif [[ "${os_name}" == "Linux" ]] && [[ "${ARCH}" == "x86_64" ]]; then
34-
vulkan_sdk_version="1.4.341.1"
34+
vulkan_sdk_version="1.4.350.0"
3535
vulkan_sdk_url="https://sdk.lunarg.com/sdk/download/${vulkan_sdk_version}/linux/vulkansdk-linux-x86_64-${vulkan_sdk_version}.tar.xz"
36-
vulkan_sdk_sha256="3bf0f762afb6c79bc6a9d9fb5998745ccff928800a29619b501ed9de7fd9789b"
36+
vulkan_sdk_sha256="b65f068ab36263559da49d7cacd7e7b9df23824ca8b68ccc522a2b06f5725df2"
3737
elif [[ "${os_name}" == "Linux" ]] && ([[ "${ARCH}" == "aarch64" ]] || [[ "${ARCH}" == "arm64" ]]); then
38-
vulkan_sdk_version="1.4.341.1"
38+
vulkan_sdk_version="1.4.350.0"
3939
if [[ "${vulkan_sdk_arch}" == "arm64" ]]; then
4040
vulkan_sdk_arch="aarch64"
4141
fi
4242
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"
43-
vulkan_sdk_sha256="345312aee2c835e128b30653278593f899a659a7ba287c571cafb22acb708b8f"
43+
vulkan_sdk_sha256="9e403d444219bb7c17e9231b580d704453e2afa30a1c2fdd568d1776dc68790b"
4444
else
4545
log_step "vulkan" "Error: only macOS and Linux are supported (detected ${os_name}); architecture must be x86-64 or aarch64/arm64"
4646
exit 1

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -474,6 +474,8 @@ class Qwen3VLTestCase:
474474
"qk_norm": TOSA_FP_TEST_CASES["qk_norm"],
475475
}
476476

477+
VGF_NO_QUANT_BF16_TEST_CASES: dict[str, Qwen3VLTestCase] = TOSA_BF16_TEST_CASES
478+
477479

478480
TOSA_MXFP8_TEST_CASES: dict[str, Qwen3VLTestCase] = {
479481
"text_attention": Qwen3VLTestCase(
@@ -602,3 +604,25 @@ def test_qwen3_vl_vgf_no_quant(test_case: Qwen3VLTestCase):
602604
transform_passes=list(test_case.transform_passes),
603605
)
604606
pipeline.run()
607+
608+
609+
@common.SkipIfNoModelConverter
610+
@common.parametrize(
611+
"test_case",
612+
VGF_NO_QUANT_BF16_TEST_CASES,
613+
)
614+
def test_qwen3_vl_vgf_no_quant_bf16(test_case: Qwen3VLTestCase):
615+
model, inputs = test_case.model_cls.prepare_model_and_inputs()
616+
model, inputs = _to_bfloat16(model, inputs)
617+
with torch.no_grad():
618+
pipeline = VgfPipeline[input_t](
619+
model,
620+
inputs,
621+
aten_op=[],
622+
exir_op=[],
623+
quantize=False,
624+
transform_passes=list(test_case.transform_passes),
625+
atol=5e-2,
626+
rtol=5e-2,
627+
)
628+
pipeline.run()

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

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -180,14 +180,16 @@ def _dtype_cases() -> list:
180180
]
181181

182182

183-
def _vgf_dtype_cases() -> list:
183+
def _vgf_no_quant_dtype_cases() -> list:
184+
return [
185+
pytest.param(torch.float32, id="fp32"),
186+
pytest.param(torch.bfloat16, id="bf16"),
187+
]
188+
189+
190+
def _vgf_quant_dtype_cases() -> list:
184191
return [
185192
pytest.param(torch.float32, id="fp32"),
186-
pytest.param(
187-
torch.bfloat16,
188-
marks=pytest.mark.xfail(reason="BF16 runtime support not ready for VGF."),
189-
id="bf16",
190-
),
191193
]
192194

193195

@@ -233,12 +235,14 @@ def test_phi3_tosa_INT_layers(name, module_factory, input_factory):
233235

234236

235237
@common.SkipIfNoModelConverter
236-
@pytest.mark.parametrize("dtype", _vgf_dtype_cases())
238+
@pytest.mark.parametrize("dtype", _vgf_no_quant_dtype_cases())
237239
@pytest.mark.parametrize("name,module_factory,input_factory", _module_cases())
238240
def test_phi3_vgf_no_quant_layers(name, module_factory, input_factory, dtype):
239241
config = _phi3_config()
240242
module = module_factory(config).to(dtype)
241243
inputs = input_factory(config, dtype)
244+
atol = 1e-02 if dtype == torch.bfloat16 else 1e-03
245+
rtol = 1e-02 if dtype == torch.bfloat16 else 1e-03
242246

243247
pipeline = VgfPipeline[input_t1 if len(inputs) == 1 else input_t2](
244248
module,
@@ -250,12 +254,14 @@ def test_phi3_vgf_no_quant_layers(name, module_factory, input_factory, dtype):
250254
InsertInt32CastsAfterInt64PlaceholdersPass(),
251255
],
252256
quantize=False,
257+
atol=atol,
258+
rtol=rtol,
253259
)
254260
pipeline.run()
255261

256262

257263
@common.SkipIfNoModelConverter
258-
@pytest.mark.parametrize("dtype", _vgf_dtype_cases())
264+
@pytest.mark.parametrize("dtype", _vgf_quant_dtype_cases())
259265
@pytest.mark.parametrize("name,module_factory,input_factory", _module_cases())
260266
def test_phi3_vgf_quant_layers(name, module_factory, input_factory, dtype):
261267
config = _phi3_config()

‎backends/arm/test/models/test_mobilenet_v2_arm.py‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626
MobileNet_V2_Weights,
2727
)
2828

29-
3029
mv2 = models.mobilenetv2.mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT)
3130
mv2 = mv2.eval()
3231
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
@@ -191,6 +190,26 @@ def test_mv2_vgf_no_quant():
191190
pipeline.run()
192191

193192

193+
@common.SkipIfNoModelConverter
194+
def test_mv2_vgf_no_quant_bf16():
195+
bf16_model = models.mobilenetv2.mobilenet_v2(
196+
weights=MobileNet_V2_Weights.DEFAULT
197+
).eval()
198+
bf16_model = bf16_model.to(torch.bfloat16)
199+
bf16_input = normalize(torch.rand((1, 3, 224, 224))).to(torch.bfloat16)
200+
pipeline = VgfPipeline[input_t](
201+
bf16_model,
202+
(bf16_input,),
203+
aten_op=[],
204+
exir_op=[],
205+
use_to_edge_transform_and_lower=True,
206+
quantize=False,
207+
atol=2.5e-01,
208+
rtol=2.5e-01,
209+
)
210+
pipeline.run()
211+
212+
194213
def test_mv2_tosa_INT_FP_partial_quant():
195214
pipeline = TosaPipelineINT[input_t](
196215
mv2,

‎backends/arm/test/models/test_resnet18.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
EthosU85PipelineINT,
1515
TosaPipelineFP,
1616
TosaPipelineINT,
17+
VgfPipeline,
1718
)
1819

1920
from torchvision import transforms # type: ignore[import-untyped]
@@ -65,6 +66,22 @@ def test_resnet_18_tosa_FP_bf16():
6566
pipeline.run()
6667

6768

69+
@common.SkipIfNoModelConverter
70+
def test_resnet_18_vgf_no_quant_bf16():
71+
bf16_model = resnet18(weights=ResNet18_Weights).eval()
72+
bf16_model = bf16_model.to(torch.bfloat16)
73+
bf16_input = normalize(torch.rand((1, 3, 224, 224)) * 2 - 1).to(torch.bfloat16)
74+
pipeline = VgfPipeline[input_t](
75+
bf16_model,
76+
(bf16_input,),
77+
aten_op=[],
78+
quantize=False,
79+
atol=6e-01,
80+
rtol=6e-01,
81+
)
82+
pipeline.run()
83+
84+
6885
@common.parametrize("per_channel_quantization", quant_test_data)
6986
def test_resnet_18_tosa_INT(per_channel_quantization):
7087
pipeline = TosaPipelineINT[input_t](

‎backends/arm/test/ops/test_abs.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,10 @@ def test_abs_u85_INT(test_data: torch.Tensor):
8888
pipeline.run()
8989

9090

91-
@common.parametrize("test_data", Abs.test_parameters | Abs.test_parameters_fp16)
91+
@common.parametrize(
92+
"test_data",
93+
Abs.test_parameters | Abs.test_parameters_bf16 | Abs.test_parameters_fp16,
94+
)
9295
@common.SkipIfNoModelConverter
9396
def test_abs_vgf_no_quant(test_data: input_t1):
9497
pipeline = VgfPipeline[input_t1](

‎backends/arm/test/ops/test_add.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -263,7 +263,9 @@ def test_add_tensor_u85_INT_2(test_data: input_t2):
263263
pipeline.run()
264264

265265

266-
@common.parametrize("test_data", Add.test_data | Add.test_data_fp16)
266+
@common.parametrize(
267+
"test_data", Add.test_data | Add.test_data_fp16 | Add.test_data_bf16
268+
)
267269
@common.SkipIfNoModelConverter
268270
def test_add_tensor_vgf_no_quant(test_data: input_t1):
269271
pipeline = VgfPipeline[input_t1](

‎backends/arm/test/ops/test_alias_copy.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,10 @@ def test_alias_u85_INT(test_data: input_t1):
9898
).run()
9999

100100

101-
@common.parametrize("test_data", AliasCopy.test_data | AliasCopy.test_data_fp16)
101+
@common.parametrize(
102+
"test_data",
103+
AliasCopy.test_data | AliasCopy.test_data_bf16 | AliasCopy.test_data_fp16,
104+
)
102105
@common.SkipIfNoModelConverter
103106
def test_alias_vgf_no_quant(test_data: input_t1):
104107
pipeline = VgfPipeline[input_t1](

‎backends/arm/test/ops/test_amax.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,9 @@ def test_max_dim_tosa_FP_not_delegated():
212212
pipeline.run()
213213

214214

215-
@common.parametrize("test_data", Amax.test_data | Amax.test_data_fp16)
215+
@common.parametrize(
216+
"test_data", Amax.test_data | Amax.test_data_fp16 | Amax.test_data_bf16
217+
)
216218
@common.SkipIfNoModelConverter
217219
def test_amax_vgf_no_quant(test_data: Amax.input_t):
218220
data, dim, keep_dims = test_data()
@@ -240,7 +242,9 @@ def test_amax_vgf_quant(test_data: Amax.input_t):
240242
pipeline.run()
241243

242244

243-
@common.parametrize("test_data", Max.test_data | Max.test_data_fp16)
245+
@common.parametrize(
246+
"test_data", Max.test_data | Max.test_data_fp16 | Max.test_data_bf16
247+
)
244248
@common.SkipIfNoModelConverter
245249
def test_max_dim_vgf_no_quant_to_amax(test_data: Max.input_t):
246250
data, dim = test_data()

0 commit comments

Comments
 (0)