[ET-VK][unary] Select unary ops with a specialization constant

backends/vulkan/op_registry.py

@@ -213,7 +213,7 @@ def register_copy_op():
 
 
 # =============================================================================
-# UnaryOp.cpp
+# unary/UnaryOp.cpp
 # =============================================================================
 
 
@@ -648,7 +648,7 @@ def register_q8ta_add():
 
 
 # =============================================================================
-# Q8taUnary.cpp
+# unary/q8ta/Q8taUnary.cpp
 # =============================================================================
 
 

backends/vulkan/runtime/graph/ops/glsl/q8ta_unary.glsl → backends/vulkan/runtime/graph/ops/glsl/unary/q8ta/q8ta_unary.glsl

@@ -12,7 +12,10 @@
 
 ${define_active_storage_type("buffer")}
 
-#define op(X) ${OPERATOR}
+#define IN_T vec4
+#define OUT_T vec4
+#define PARAM_T float
+#define UNARY_OP_IS_FLOAT
 
 layout(std430) buffer;
 
@@ -45,6 +48,9 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 ${layout_declare_spec_const(C, "int", "out_layout", "CONTIG_LAYOUT_INT")}
 ${layout_declare_spec_const(C, "int", "in_layout", "CONTIG_LAYOUT_INT")}
 ${layout_declare_spec_const(C, "int", "block_config", "0")}
+${layout_declare_spec_const(C, "int", "op_type", "0")}
+
+#include "unary_op.glslh"
 
 // Generate loading functions for input buffer
 define_load_int8x4_buffer_fns(t_in)
@@ -74,7 +80,8 @@ void main() {
     vec4 in_texel = unpack_and_dequantize(
         in_block[row], input_scale, input_zp);
 
-    vec4 out_texel = op(in_texel);
+    // Only ops without parameters are dispatched to this shader
+    vec4 out_texel = apply_unary_op(in_texel, 0.0, 0.0);
     out_block[row] = quantize_and_pack(out_texel, output_inv_scale, output_zp);
   }
 

backends/vulkan/runtime/graph/ops/glsl/q8ta_unary.yaml → backends/vulkan/runtime/graph/ops/glsl/unary/q8ta/q8ta_unary.yaml

@@ -4,9 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# The op is selected at dispatch through the op_type specialization constant.
 q8ta_unary:
-  parameter_names_with_default_values:
-    OPERATOR: X
+  parameter_names_with_default_values: {}
   shader_variants:
-    - NAME: q8ta_relu_buffer
-      OPERATOR: max(X, vec4(0.0))
+    - NAME: q8ta_unary_buffer

backends/vulkan/runtime/graph/ops/glsl/unary_op.glsl → backends/vulkan/runtime/graph/ops/glsl/unary/unary.glsl

@@ -15,7 +15,22 @@ ${define_required_extensions(STORAGE, DTYPE)}
 #define VEC4_T ${texel_load_type(DTYPE, STORAGE)}
 #define T ${buffer_scalar_type(DTYPE)}
 
-#define op(X, A, B) ${OPERATOR}
+// Buffer ops compute in the parameter type. Runtime bounds of int32 clamps stay
+// int so that large values remain exact.
+$if DYNAMIC_PARAMS:
+  #define PARAM_T ${texel_load_component_type(DTYPE, "texture3d")}
+$else:
+  #define PARAM_T float
+$if STORAGE == "buffer":
+  #define IN_T PARAM_T
+  #define OUT_T T
+$else:
+  #define IN_T VEC4_T
+  #define OUT_T VEC4_T
+$if DTYPE in ["half", "float"]:
+  #define UNARY_OP_IS_FLOAT
+$elif DTYPE == "uint8":
+  #define UNARY_OP_IS_BITWISE
 
 ${define_active_storage_type(STORAGE)}
 
@@ -44,14 +59,22 @@ $else:
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+${layout_declare_spec_const(C, "int", "op_type", "0")}
+
 #include "dispatch.glslh"
-#include "activations.h"
+#include "unary_op.glslh"
 
 $if DYNAMIC_PARAMS:
   float decode_bound(const uint value, const int is_int) {
     return is_int != 0 ? float(int(value)) : uintBitsToFloat(value);
   }
 
+  #define PARAM_A PARAM_T(decode_bound(minimum, bounds_are_int.x))
+  #define PARAM_B PARAM_T(decode_bound(maximum, bounds_are_int.y))
+$else:
+  #define PARAM_A minimum
+  #define PARAM_B maximum
+
 #ifdef USING_BUFFER
 
 void main() {
@@ -60,14 +83,7 @@ void main() {
     return;
   }
 
-$if DYNAMIC_PARAMS:
-  const T in_val = T(t_in[i]);
-  const T minimum_val = T(decode_bound(minimum, bounds_are_int.x));
-  const T maximum_val = T(decode_bound(maximum, bounds_are_int.y));
-  t_out[i] = T(op(in_val, minimum_val, maximum_val));
-$else:
-  const float in_val = float(t_in[i]);
-  t_out[i] = T(op(in_val, minimum, maximum));
+  t_out[i] = apply_unary_op(IN_T(t_in[i]), PARAM_A, PARAM_B);
 }
 
 #else
@@ -80,12 +96,7 @@ void main() {
   }
 
   VEC4_T in_texel = texelFetch(t_in, pos, 0);
-$if DYNAMIC_PARAMS:
-  const VEC4_T minimum_val = VEC4_T(decode_bound(minimum, bounds_are_int.x));
-  const VEC4_T maximum_val = VEC4_T(decode_bound(maximum, bounds_are_int.y));
-  imageStore(t_out, pos, op(in_texel, minimum_val, maximum_val));
-$else:
-  imageStore(t_out, pos, VEC4_T(op(in_texel, minimum, maximum)));
+  imageStore(t_out, pos, apply_unary_op(in_texel, PARAM_A, PARAM_B));
 }
 
 #endif

backends/vulkan/runtime/graph/ops/glsl/unary/unary.yaml

@@ -0,0 +1,23 @@
+# The op is selected at dispatch through the op_type specialization constant.
+unary:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: texture3d
+    DYNAMIC_PARAMS: false
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: half
+      - VALUE: float
+      - VALUE: int32
+    STORAGE:
+      - VALUE: texture3d
+      - VALUE: buffer
+  shader_variants:
+    # int32 implements clamp only; uint8 implements bitwise_not only.
+    - NAME: unary_op
+    - NAME: unary_op_uint8
+      DTYPE: uint8
+    # Reads its parameters from uniform buffers, for clamp bounds that are only
+    # known at runtime.
+    - NAME: unary_op_dynamic
+      DYNAMIC_PARAMS: true