[ET-VK][conv2d] Skip the store rounding of texels computed in f16

backends/vulkan/runtime/api/containers/StagingBuffer.cpp

@@ -31,6 +31,8 @@ inline uint32_t fp32_to_bits(float f) {
   return bits;
 }
 
+} // namespace
+
 /*
  * Convert a 16-bit floating-point number in IEEE half-precision format, in bit
  * representation, to a 32-bit floating-point number in IEEE single-precision
@@ -127,8 +129,6 @@ uint16_t float_to_half(float f) {
       (shl1_w > UINT32_C(0xFF000000) ? UINT16_C(0x7E00) : nonsign));
 }
 
-} // namespace
-
 StagingBuffer::StagingBuffer(
     Context* context_p,
     const vkapi::ScalarType dtype,

backends/vulkan/runtime/api/containers/StagingBuffer.h

@@ -19,6 +19,11 @@
 namespace vkcompute {
 namespace api {
 
+// IEEE half-precision conversions; float_to_half rounds to the nearest f16,
+// ties to even.
+float half_to_float(uint16_t h);
+uint16_t float_to_half(float f);
+
 vkapi::ScalarType get_staging_dtype(
     Context* context_p,
     vkapi::ScalarType dtype);

backends/vulkan/runtime/gen_vulkan_spv.py

@@ -357,8 +357,11 @@ def layout_declare_buffer(
 
 # 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.
+# the shader makes the store exact. Integer texels are stored as is, and so are
+# f16 texels in shaders that define HAS_FLOAT16_ACCUMULATION before their first
+# writable image declaration; elsewhere, f16 texels widen to vec4 and take the
+# rounding path. 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"
@@ -377,6 +380,9 @@ def layout_declare_buffer(
 }
 ivec4 round_to_f16(const ivec4 v) { return v; }
 uvec4 round_to_f16(const uvec4 v) { return v; }
+#ifdef HAS_FLOAT16_ACCUMULATION
+vec4 round_to_f16(const f16vec4 v) { return vec4(v); }
+#endif
 #define imageStore(image, pos, texel) imageStore(image, pos, round_to_f16(texel))
 #endif"""
 

backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_epilogue.glslh

@@ -12,6 +12,10 @@
  *   out = activation(residual_op(conv, residual * scale, alpha))
  *
  * residual_op is add, sub or mul, and activation is clamp, relu or leaky_relu.
+ * The ops compute in VEC4_T, the including shader's texel type, so that texels
+ * computed in f16 stay f16 and are stored without rounding. Where VEC4_T is
+ * f16vec4, the host rounds the uniforms to the nearest f16, so converting them
+ * is exact.
  *
  * The including shader binds t_residual, a texture3d with the output's sizes,
  * and declares the specialization constants has_residual, has_scale,
@@ -46,34 +50,34 @@
 // op is a specialization constant, so every case of their switches stays in the
 // SPIR-V, and the epilogue is inlined at each output store of every conv2d
 // shader.
-vec4 apply_residual_op(const vec4 v, const vec4 residual) {
+VEC4_T apply_residual_op(const VEC4_T v, const VEC4_T residual) {
   switch (residual_op) {
     case BINARY_OP_SUB:
-      return v - residual_alpha * residual;
+      return v - VEC4_T(residual_alpha) * residual;
     case BINARY_OP_MUL:
       return v * residual;
     default: // BINARY_OP_ADD
-      return v + residual_alpha * residual;
+      return v + VEC4_T(residual_alpha) * residual;
   }
 }
 
-vec4 apply_activation(const vec4 v) {
+VEC4_T apply_activation(const VEC4_T v) {
   switch (activation_op) {
     case UNARY_OP_RELU:
-      return max(v, vec4(0));
+      return max(v, VEC4_T(0));
     case UNARY_OP_LEAKY_RELU:
-      return leaky_relu(v, activation_a);
+      return mix(v * VEC4_T(activation_a), v, greaterThan(v, VEC4_T(0)));
     default: // UNARY_OP_CLAMP
-      return clamp(v, activation_a, activation_b);
+      return clamp(v, VEC4_T(activation_a), VEC4_T(activation_b));
   }
 }
 
 // `scale` is the residual scale of the output channels in `pos.z`.
-vec4 apply_epilogue(vec4 v, const ivec3 pos, const vec4 scale) {
+VEC4_T apply_epilogue(VEC4_T v, const ivec3 pos, const vec4 scale) {
   if (has_residual) {
-    vec4 residual = texelFetch(t_residual, pos, 0);
+    VEC4_T residual = VEC4_T(texelFetch(t_residual, pos, 0));
     if (has_scale) {
-      residual *= scale;
+      residual *= VEC4_T(scale);
     }
     v = apply_residual_op(v, residual);
   }

backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw_tiled.glsl

@@ -165,7 +165,7 @@ void store_output_tile_with_checks(
         imageStore(
             t_out,
             pos,
-            apply_epilogue(vec4(out_tile.data[m][n4]), pos, scale[n4]));
+            apply_epilogue(out_tile.data[m][n4], pos, scale[n4]));
       }
     }
   }

backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_tiled.glsl

@@ -230,11 +230,11 @@ void store_tile(ivec3 pos) {
     const ivec2 offset = ivec2(i % TILE_X, i / TILE_X);
     if (pos.x + offset.x < out_limits.x && pos.y + offset.y < out_limits.y) {
       const ivec3 p0 = pos + ivec3(offset, 0);
-      imageStore(t_out, p0, apply_epilogue(vec4(sums[i][0]), p0, scale[0]));
+      imageStore(t_out, p0, apply_epilogue(sums[i][0], p0, scale[0]));
 #if OUT_N4 == 2
       if (second_slice) {
         const ivec3 p1 = pos + ivec3(offset, 1);
-        imageStore(t_out, p1, apply_epilogue(vec4(sums[i][1]), p1, scale[1]));
+        imageStore(t_out, p1, apply_epilogue(sums[i][1], p1, scale[1]));
       }
 #endif
     }

backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/conv2d_gemm.glsl

@@ -192,14 +192,11 @@ void store_output_tile_with_checks(
       const int ow = m_local % W_out;
       const int oh = OH_OFFSET + m_local / W_out;
       if (m_local < M_TILE && oh < H_out && n4_start + n4 < N4) {
-        // t_out is a texture3d, whose imageStore takes vec4 regardless of
-        // DTYPE, so the accumulator (f16vec4 for the buffer/half path) is cast
-        // to vec4 for the epilogue and store.
         const ivec3 pos = ivec3(ow, oh, n4_start + n4);
         imageStore(
             t_out,
             pos,
-            apply_epilogue(vec4(out_tile.data[m][n4]), pos, scale[n4]));
+            apply_epilogue(out_tile.data[m][n4], pos, scale[n4]));
       }
     }
   }

backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2d.cpp

@@ -53,9 +53,19 @@ void check_conv2d_residual_on_resize(
 }
 
 Conv2dEpilogueParams create_conv2d_epilogue_params(
-    const Conv2dEpilogue& epilogue) {
+    const Conv2dEpilogue& epilogue,
+    const bool f16_epilogue) {
+  if (!f16_epilogue) {
+    return {
+        epilogue.activation_a, epilogue.activation_b, epilogue.residual_alpha};
+  }
+  const auto round_to_half = [](const float f) {
+    return api::half_to_float(api::float_to_half(f));
+  };
   return {
-      epilogue.activation_a, epilogue.activation_b, epilogue.residual_alpha};
+      round_to_half(epilogue.activation_a),
+      round_to_half(epilogue.activation_b),
+      round_to_half(epilogue.residual_alpha)};
 }
 
 void append_conv2d_epilogue_spec_vars(

backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2d.h

@@ -95,8 +95,12 @@ void check_conv2d_residual_on_resize(
     const ValueRef out,
     const std::vector<ValueRef>& resize_args);
 
+// With `f16_epilogue`, for shaders that compute the epilogue in f16, the params
+// are rounded to the nearest f16 so that the shaders convert them exactly. Some
+// GPUs round f32 to f16 conversions toward zero.
 Conv2dEpilogueParams create_conv2d_epilogue_params(
-    const Conv2dEpilogue& epilogue);
+    const Conv2dEpilogue& epilogue,
+    const bool f16_epilogue);
 
 // Appends the has_activation, activation_op, has_residual, has_scale and
 // residual_op specialization constants.

backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dDW.cpp

@@ -188,7 +188,7 @@ void add_conv2d_dw_node(
     const Conv2dEpilogueArgs& epilogue_args) {
   vkapi::ShaderInfo shader = VK_KERNEL_FROM_STR(kernel_name);
   const Conv2dEpilogueParams epilogue_params =
-      create_conv2d_epilogue_params(epilogue);
+      create_conv2d_epilogue_params(epilogue, false);
   vkapi::SpecVarList spec_vars;
   append_conv2d_epilogue_spec_vars(spec_vars, epilogue);
 

backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dDirect.cpp

@@ -472,8 +472,9 @@ void add_conv2d_direct_node(
     const bool has_bias,
     const bool force_masked,
     const bool fast_tiled) {
-  const Conv2dEpilogueParams epilogue_params =
-      create_conv2d_epilogue_params(epilogue);
+  const Conv2dEpilogueParams epilogue_params = create_conv2d_epilogue_params(
+      epilogue,
+      graph.dtype_of(out) == vkapi::kHalf && !graph.get_bool(transposed));
   vkapi::ParamsBindList param_buffers = {
       graph.logical_limits_ubo(out),
       graph.sizes_ubo(in),

backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dPW.cpp

@@ -272,8 +272,8 @@ void add_conv2d_pw_tiled_node(
 
   utils::ivec4 stride_padding{stride_h, stride_w, padding_h, padding_w};
 
-  const Conv2dEpilogueParams epilogue_params =
-      create_conv2d_epilogue_params(epilogue);
+  const Conv2dEpilogueParams epilogue_params = create_conv2d_epilogue_params(
+      epilogue, graph.dtype_of(out) == vkapi::kHalf);
 
   vkapi::SpecVarList spec_vars = {s1p0 ? 1 : 0};
   append_conv2d_epilogue_spec_vars(spec_vars, epilogue);

backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/im2col/Conv2dGemm.cpp

@@ -276,8 +276,10 @@ void add_conv2d_gemm_node(
   // from the refreshed out_sizes UBO (a baked plain-data push constant cannot
   // be updated on resize, but oh_offset / oh_tile never change with shape).
   const utils::ivec4 gemm_dims{K4_total, oh_offset, oh_tile, 0};
-  const Conv2dEpilogueParams epilogue_params =
-      create_conv2d_epilogue_params(epilogue);
+  const Conv2dEpilogueParams epilogue_params = create_conv2d_epilogue_params(
+      epilogue,
+      graph.dtype_of(out) == vkapi::kHalf &&
+          graph.is_buffer_storage(im2col_in));
   vkapi::SpecVarList spec_vars;
   append_conv2d_epilogue_spec_vars(spec_vars, epilogue);