From 1e84adafff0322488ffe2e4032a3a0b78be2e0c2 Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Tue, 3 Nov 2026 06:24:32 -0800 Subject: [PATCH] [executorch][vulkan] Round f16 image stores to nearest even Pull Request resolved: https://github.com/pytorch/executorch/pull/24372 Vulkan lets the device choose how stores to rgba16f images round, and devices that round toward zero bias f16 outputs; shader float controls don't govern that conversion. `gen_vulkan_spv.py` now wraps `imageStore` on writable rgba16f images in a round-to-nearest-even helper, and six shaders that declared such images by hand, two of them test shaders, now use the codegen. On devices that round toward zero, stores of 65520 and above now give inf, as the CPU conversion does, instead of 65504. The four texture shaders that already call `round_to_half_rte` on their outputs (`binary_scalar_texture`, `full_texture`, `reduce`, `scalar_tensor`) now round twice, with identical results; removing those calls is a follow-up. `aten._to_copy` from f32 to f16 textures is a `vkCmdBlitImage`, which the codegen doesn't reach, so `test_to_copy` keeps its one-ulp tolerance. Authored with Claude Code. @bypass-github-export-checks @bypass-github-pytorch-ci-checks @bypass-github-executorch-ci-checks ghstack-source-id: 463657749 @exported-using-ghexport Differential Revision: [D128254886](https://our.internmc.facebook.com/intern/diff/D128254886/) --- backends/vulkan/runtime/gen_vulkan_spv.py | 37 +++++++++++ .../conv2d_dw_prepack_weights.glsl | 2 +- .../pack_weights/conv2d_prepack_weights.glsl | 2 +- .../conv_transpose2d_prepack_weights.glsl | 2 +- .../graph/ops/glsl/repeat_channel.glsl | 2 +- .../glsl/binary_op_nobroadcast__test.glsl | 4 +- .../vulkan/test/glsl/fill_texture__test.glsl | 2 +- .../test/spv_codegen/test_gen_vulkan_spv.py | 16 +++++ .../vulkan/test/vulkan_compute_api_test.cpp | 65 +++++++++++++++++++ 9 files changed, 124 insertions(+), 8 deletions(-) diff --git a/backends/vulkan/runtime/gen_vulkan_spv.py b/backends/vulkan/runtime/gen_vulkan_spv.py index 5701248266e..76999cb0995 100644 --- a/backends/vulkan/runtime/gen_vulkan_spv.py +++ b/backends/vulkan/runtime/gen_vulkan_spv.py @@ -355,6 +355,38 @@ def layout_declare_buffer( return out_str +# Vulkan leaves the rounding of a store to an rgba16f image to the device, and +# some devices round toward zero. Rounding to the nearest f16, ties to even, in +# the shader makes the store exact. Integer texels are stored as is, but stores +# to rgba32f images would be rounded too, so a shader may not write both. +ROUND_F16_IMAGE_STORES = """ +#ifdef WRITES_RGBA32F_IMAGE +#error "A shader cannot write both rgba16f and rgba32f images" +#endif +#ifndef WRITES_RGBA16F_IMAGE +#define WRITES_RGBA16F_IMAGE +vec4 round_to_f16(const vec4 v) { + const uvec4 u = floatBitsToUint(v); + // Keep 10 mantissa bits: add just under half an f16 ulp, plus one when the + // kept lowest bit is odd so ties go to even, then clear the 13 dropped bits. + vec4 r = uintBitsToFloat((u + 0xFFFu + ((u >> 13) & 1u)) & ~0x1FFFu); + // Below the smallest normal f16, round to a multiple of 2^-24. + r = mix(r, roundEven(v * 16777216.0) * 5.9604644775390625e-8, lessThan(abs(v), vec4(6.103515625e-5))); + r = mix(r, uintBitsToFloat((u & 0x80000000u) | 0x7F800000u), greaterThan(abs(r), vec4(65504.0))); + return mix(r, v, isnan(v)); +} +ivec4 round_to_f16(const ivec4 v) { return v; } +uvec4 round_to_f16(const uvec4 v) { return v; } +#define imageStore(image, pos, texel) imageStore(image, pos, round_to_f16(texel)) +#endif""" + +WRITES_RGBA32F_IMAGE = """ +#ifdef WRITES_RGBA16F_IMAGE +#error "A shader cannot write both rgba16f and rgba32f images" +#endif +#define WRITES_RGBA32F_IMAGE""" + + def layout_declare_image( slot: Union[int, List[int]], access_type: str, @@ -367,6 +399,11 @@ def layout_declare_image( image_type = TYPE_MAPPINGS["IMAGE_T"][image_ndim][dtype] ret_str = f"layout(set = 0, binding = {get_slot_val(slot)}, {image_format}) uniform {precision} restrict {get_access_qualifier(access_type)} {image_type} {var_name};" + if access_type.lower() != "r": + if image_format == "rgba16f": + ret_str += ROUND_F16_IMAGE_STORES + elif image_format == "rgba32f": + ret_str += WRITES_RGBA32F_IMAGE if isinstance(slot, list): slot[0] = slot[0] + 1 diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_dw_prepack_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_dw_prepack_weights.glsl index 4707bcf9f9c..4a16ee05d42 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_dw_prepack_weights.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_dw_prepack_weights.glsl @@ -21,7 +21,7 @@ ${define_required_extensions("buffer", BUF_DTYPE)} layout(std430) buffer; -layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[2][DTYPE]} image_out; +${layout_declare_tensor(0, "w", "image_out", DTYPE, "texture2d")} layout(set = 0, binding = 1) buffer PRECISION restrict readonly Buffer { BUF_T buffer_in[]; }; diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_prepack_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_prepack_weights.glsl index b5bec8fca0b..95c8cc3c4ee 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_prepack_weights.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_prepack_weights.glsl @@ -29,7 +29,7 @@ $if OUTPUT_BUFFER: f16vec4 buffer_out[]; }; $else: - layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[2][DTYPE]} image_out; + ${layout_declare_tensor(0, "w", "image_out", DTYPE, "texture2d")} layout(set = 0, binding = 1) buffer PRECISION restrict readonly Buffer { BUF_T buffer_in[]; }; diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/pack_weights/conv_transpose2d_prepack_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/pack_weights/conv_transpose2d_prepack_weights.glsl index 65b998a9c5b..4a47892375e 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/pack_weights/conv_transpose2d_prepack_weights.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/pack_weights/conv_transpose2d_prepack_weights.glsl @@ -21,7 +21,7 @@ ${define_required_extensions("buffer", BUF_DTYPE)} layout(std430) buffer; -layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[2][DTYPE]} image_out; +${layout_declare_tensor(0, "w", "image_out", DTYPE, "texture2d")} layout(set = 0, binding = 1) buffer PRECISION restrict readonly Buffer { BUF_T buffer_in[]; }; diff --git a/backends/vulkan/runtime/graph/ops/glsl/repeat_channel.glsl b/backends/vulkan/runtime/graph/ops/glsl/repeat_channel.glsl index 42c7f86aea8..96e9f7aac4c 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/repeat_channel.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/repeat_channel.glsl @@ -16,7 +16,7 @@ layout(std430) buffer; #include "indexing_utils.h" -layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} image_out; +${layout_declare_image(0, "w", "image_out", DTYPE, image_ndim=NDIM)} layout(set = 0, binding = 1) uniform PRECISION sampler3D image_in; layout(set = 0, binding = 2) uniform PRECISION restrict RepeatArgs { diff --git a/backends/vulkan/test/glsl/binary_op_nobroadcast__test.glsl b/backends/vulkan/test/glsl/binary_op_nobroadcast__test.glsl index 7f72ac58972..5fcd1513c4c 100644 --- a/backends/vulkan/test/glsl/binary_op_nobroadcast__test.glsl +++ b/backends/vulkan/test/glsl/binary_op_nobroadcast__test.glsl @@ -14,9 +14,7 @@ layout(std430) buffer; -// clang-format off -layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly image3D image_out; -// clang-format on +${layout_declare_image(0, "w", "image_out", DTYPE)} layout(set = 0, binding = 1) uniform PRECISION sampler3D image_in; layout(set = 0, binding = 2) uniform PRECISION sampler3D image_other; diff --git a/backends/vulkan/test/glsl/fill_texture__test.glsl b/backends/vulkan/test/glsl/fill_texture__test.glsl index 76c630de55e..637c12fbb51 100644 --- a/backends/vulkan/test/glsl/fill_texture__test.glsl +++ b/backends/vulkan/test/glsl/fill_texture__test.glsl @@ -12,7 +12,7 @@ layout(std430) buffer; -layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} uOutput; +${layout_declare_image(0, "w", "uOutput", DTYPE, image_ndim=NDIM)} layout(set = 0, binding = 1) uniform PRECISION restrict Block { ivec3 size; int fill; diff --git a/backends/vulkan/test/spv_codegen/test_gen_vulkan_spv.py b/backends/vulkan/test/spv_codegen/test_gen_vulkan_spv.py index aef464dbf38..e569d6fe27d 100644 --- a/backends/vulkan/test/spv_codegen/test_gen_vulkan_spv.py +++ b/backends/vulkan/test/spv_codegen/test_gen_vulkan_spv.py @@ -13,6 +13,22 @@ from runtime import gen_vulkan_spv +class TestF16ImageStores(unittest.TestCase): + def test_writable_rgba16f_images_round_stores(self) -> None: + for access in ("w", "rw"): + with self.subTest(access=access): + declaration = gen_vulkan_spv.layout_declare_image( + 0, access, "t_out", "half" + ) + self.assertIn("round_to_f16(texel)", declaration) + + def test_other_images_store_unchanged(self) -> None: + for access, dtype in (("r", "half"), ("w", "float"), ("w", "int32")): + with self.subTest(access=access, dtype=dtype): + declaration = gen_vulkan_spv.layout_declare_image(0, access, "t", dtype) + self.assertNotIn("round_to_f16", declaration) + + class TestNestedShaderSources(unittest.TestCase): def _write_fake_compiler( self, directory: Path, assertions: str = "" diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp index b087b0e4baa..7f4ef389040 100644 --- a/backends/vulkan/test/vulkan_compute_api_test.cpp +++ b/backends/vulkan/test/vulkan_compute_api_test.cpp @@ -1202,6 +1202,71 @@ TEST_F(VulkanComputeAPITest, spec_var_shader_test) { } } +TEST_F(VulkanComputeAPITest, half_texture_store_rounds_to_nearest_even) { + // Ties, their neighbors, f16 subnormals and values at and past the largest + // f16, where rounding toward zero differs from rounding to nearest even. + const std::vector values = { + 1.0f + 0x1p-11f, + 1.0f + 0x3p-11f, + -(1.0f + 0x3p-11f), + std::nextafter(1.0f + 0x1p-11f, 0.0f), + std::nextafter(1.0f + 0x1p-11f, 2.0f), + 0x1p-25f, + 0x3p-25f, + std::nextafter(0x1p-25f, 0.0f), + 0x1.ffcp-15f, + -0x1.0p-24f, + 65504.0f, + 65519.0f, + 65520.0f, + -65520.0f, + 1e6f, + -1e6f, + }; + + struct Params final { + utils::ivec3 size; + int32_t fill; + utils::vec4 values; + }; + + for (size_t i = 0; i < values.size(); i += 4) { + vTensor a( + context(), + {4, 1, 1}, + vkapi::kHalf, + utils::kTexture3D, + utils::kChannelsPacked); + Params block{ + {1, 1, 1}, + 0, + {values[i], values[i + 1], values[i + 2], values[i + 3]}, + }; + ParamsBuffer params(context(), block); + + vkapi::PipelineBarrier pipeline_barrier{}; + context()->submit_compute_job( + VK_KERNEL(fill_texture__test_half), + pipeline_barrier, + GlobalWorkGrid({1u, 1u, 1u}, kExplicitWorkGrid), + LocalWorkGroup(1u, 1u, 1u), + {}, + VK_NULL_HANDLE, + 0, + a.image( + pipeline_barrier, + vkapi::PipelineStage::COMPUTE, + vkapi::MemoryAccessType::WRITE), + params.buffer()); + + const std::vector out = extract_vtensor(a); + for (size_t j = 0; j < 4; ++j) { + EXPECT_EQ(out[j], float(executorch::aten::Half(values[i + j]))) + << "input " << values[i + j]; + } + } +} + TEST_F(VulkanComputeAPITest, update_params_between_submit) { context()->set_cmd(/*reusable = */ true); std::vector sizes = {4, 4, 2};