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};