[ET-VK][conv2d] Fuse residual and activation epilogues into conv2d shaders

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

@@ -12,8 +12,6 @@
 
 #define VEC4_T ${texel_type(DTYPE)}
 
-#define op(X, A, B) ${OPERATOR}
-
 #include "indexing_utils.h"
 
 layout(std430) buffer;
@@ -22,17 +20,25 @@ ${layout_declare_tensor(0, "w", "t_out", DTYPE, "texture3d")}
 ${layout_declare_tensor(1, "r", "t_in", DTYPE, "texture3d")}
 ${layout_declare_tensor(2, "r", "t_kernel", DTYPE, "texture2d")}
 ${layout_declare_tensor(3, "r", "t_bias", DTYPE, "texture2d")}
-${layout_declare_ubo(4, "ivec3", "out_limits")}
-${layout_declare_ubo(5, "ivec4", "in_sizes")}
-${layout_declare_ubo(6, "ivec2", "kernel_size", "ivec2", "stride", "ivec2", "padding", "ivec2", "dilation")}
-${layout_declare_ubo(7, "ivec2", "overlay_region", "int", "in_group_size")}
-${layout_declare_ubo(8, "float", "out_min", "float", "out_max")}
+${layout_declare_tensor(4, "r", "t_residual", DTYPE, "texture3d")}
+${layout_declare_tensor(5, "r", "t_scale", DTYPE, "texture2d")}
+${layout_declare_ubo(6, "ivec3", "out_limits")}
+${layout_declare_ubo(7, "ivec4", "in_sizes")}
+${layout_declare_ubo(8, "ivec2", "kernel_size", "ivec2", "stride", "ivec2", "padding", "ivec2", "dilation")}
+${layout_declare_ubo(9, "ivec2", "overlay_region", "int", "in_group_size")}
+${layout_declare_ubo(10, "float", "activation_a", "float", "activation_b", "float", "residual_alpha")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
 #include "dispatch.glslh"
 
-${layout_declare_spec_const(C, "int", "ngroups", "1")}
+${layout_declare_spec_const(C, "bool", "has_activation", "false")}
+${layout_declare_spec_const(C, "int", "activation_op", "0")}
+${layout_declare_spec_const(C, "bool", "has_residual", "false")}
+${layout_declare_spec_const(C, "bool", "has_scale", "false")}
+${layout_declare_spec_const(C, "int", "residual_op", "0")}
+
+#include "conv2d_epilogue.glslh"
 
 /*
  * Computes a depthwise convolution. Each shader invocation calculates the
@@ -79,5 +85,7 @@ void main() {
     }
   }
 
-  imageStore(t_out, pos, op(sum, out_min, out_max));
+  const vec4 scale =
+      has_scale ? texelFetch(t_scale, ivec2(pos.z, 0), 0) : vec4(1);
+  imageStore(t_out, pos, apply_epilogue(vec4(sum), pos, scale));
 }

backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_dw.yaml

@@ -6,7 +6,6 @@
 
 conv2d_dw:
   parameter_names_with_default_values:
-    OPERATOR: X
     NDIM: 3
     DTYPE: float
   generate_variant_forall:
@@ -15,5 +14,3 @@ conv2d_dw:
       - VALUE: float
   shader_variants:
     - NAME: conv2d_dw
-    - NAME: conv2d_dw_clamp
-      OPERATOR: clamp(X, A, B)

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

@@ -20,8 +20,6 @@
 
 #define LOCAL_WG_SIZE 64
 
-#define op(X, A, B) ${OPERATOR}
-
 #include "indexing_utils.h"
 
 layout(std430) buffer;
@@ -30,6 +28,8 @@ ${layout_declare_tensor(0, "w", "t_out", DTYPE, "texture3d")}
 ${layout_declare_tensor(1, "r", "t_in", DTYPE, "texture3d")}
 ${layout_declare_tensor(2, "r", "t_kernel", DTYPE, "texture2d")}
 ${layout_declare_tensor(3, "r", "t_bias", DTYPE, "texture2d")}
+${layout_declare_tensor(4, "r", "t_residual", DTYPE, "texture3d")}
+${layout_declare_tensor(5, "r", "t_scale", DTYPE, "texture2d")}
 
 layout(push_constant) uniform restrict Block {
   ivec4 out_limits;
@@ -41,12 +41,21 @@ layout(push_constant) uniform restrict Block {
   ivec2 overlay_region;
   int in_group_size;
   int dummy_padding;
-  float out_min;
-  float out_max;
+  float activation_a;
+  float activation_b;
+  float residual_alpha;
 };
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+${layout_declare_spec_const(C, "bool", "has_activation", "false")}
+${layout_declare_spec_const(C, "int", "activation_op", "0")}
+${layout_declare_spec_const(C, "bool", "has_residual", "false")}
+${layout_declare_spec_const(C, "bool", "has_scale", "false")}
+${layout_declare_spec_const(C, "int", "residual_op", "0")}
+
+#include "conv2d_epilogue.glslh"
+
 /*
  * Computes a depthwise convolution. Each shader invocation calculates the
  * output at a single output location.
@@ -125,11 +134,16 @@ void main() {
   }
 
   const VEC4_T bias = texelFetch(t_bias, ivec2(pos.z, 0), 0);
+  const vec4 scale =
+      has_scale ? texelFetch(t_scale, ivec2(pos.z, 0), 0) : vec4(1);
   for (int y = 0; y < BATCH_SIZE_Y; y++) {
     for (int x = 0; x < BATCH_SIZE_X; x++) {
       const ivec3 out_pos = ivec3(pos.x + x, pos.y + y, pos.z);
       if (all(lessThan(out_pos.xy, out_limits.xy))) {
-        imageStore(t_out, out_pos, op(sum[y * BATCH_SIZE_X + x] + bias, out_min, out_max));
+        imageStore(
+            t_out,
+            out_pos,
+            apply_epilogue(vec4(sum[y * BATCH_SIZE_X + x] + bias), out_pos, scale));
       }
     }
   }

backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_dw_output_tile.yaml

@@ -6,7 +6,6 @@
 
 conv2d_dw_output_tile:
   parameter_names_with_default_values:
-    OPERATOR: X
     NDIM: 3
     DTYPE: float
     TILE_SIZE: 3
@@ -18,26 +17,12 @@ conv2d_dw_output_tile:
       - VALUE: float
   shader_variants:
     - NAME: conv2d_dw_output_tile_3x3
-    - NAME: conv2d_dw_output_tile_3x3_clamp
-      OPERATOR: clamp(X, A, B)
     - NAME: conv2d_dw_output_tile_5x5
       TILE_SIZE: 5
-    - NAME: conv2d_dw_output_tile_5x5_clamp
-      OPERATOR: clamp(X, A, B)
-      TILE_SIZE: 5
     - NAME: conv2d_dw_output_tile_3x3_b1x1
       BATCH_SIZE_X: 1
       BATCH_SIZE_Y: 1
-    - NAME: conv2d_dw_output_tile_3x3_b1x1_clamp
-      OPERATOR: clamp(X, A, B)
-      BATCH_SIZE_X: 1
-      BATCH_SIZE_Y: 1
     - NAME: conv2d_dw_output_tile_5x5_b1x1
       TILE_SIZE: 5
       BATCH_SIZE_X: 1
       BATCH_SIZE_Y: 1
-    - NAME: conv2d_dw_output_tile_5x5_b1x1_clamp
-      OPERATOR: clamp(X, A, B)
-      TILE_SIZE: 5
-      BATCH_SIZE_X: 1
-      BATCH_SIZE_Y: 1

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

@@ -14,8 +14,6 @@
 
 #define TILE_SIZE ${TILE_SIZE}
 
-#define op(X, A, B) ${OPERATOR}
-
 #include "indexing_utils.h"
 
 layout(std430) buffer;
@@ -24,6 +22,8 @@ ${layout_declare_tensor(0, "w", "t_out", DTYPE, "texture3d")}
 ${layout_declare_tensor(1, "r", "t_in", DTYPE, "texture3d")}
 ${layout_declare_tensor(2, "r", "t_kernel", DTYPE, "texture2d")}
 ${layout_declare_tensor(3, "r", "t_bias", DTYPE, "texture2d")}
+${layout_declare_tensor(4, "r", "t_residual", DTYPE, "texture3d")}
+${layout_declare_tensor(5, "r", "t_scale", DTYPE, "texture2d")}
 
 layout(push_constant) uniform restrict Block {
   ivec4 out_limits;
@@ -35,12 +35,21 @@ layout(push_constant) uniform restrict Block {
   ivec2 overlay_region;
   int in_group_size;
   int dummy_padding;
-  float out_min;
-  float out_max;
+  float activation_a;
+  float activation_b;
+  float residual_alpha;
 };
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+${layout_declare_spec_const(C, "bool", "has_activation", "false")}
+${layout_declare_spec_const(C, "int", "activation_op", "0")}
+${layout_declare_spec_const(C, "bool", "has_residual", "false")}
+${layout_declare_spec_const(C, "bool", "has_scale", "false")}
+${layout_declare_spec_const(C, "int", "residual_op", "0")}
+
+#include "conv2d_epilogue.glslh"
+
 /*
  * Computes a depthwise convolution. Each shader invocation calculates the
  * output at a single output location.
@@ -79,5 +88,7 @@ void main() {
     }
   }
 
-  imageStore(t_out, pos, op(sum, out_min, out_max));
+  const vec4 scale =
+      has_scale ? texelFetch(t_scale, ivec2(pos.z, 0), 0) : vec4(1);
+  imageStore(t_out, pos, apply_epilogue(vec4(sum), pos, scale));
 }

backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_dw_sned_output_tile.yaml

@@ -6,7 +6,6 @@
 
 conv2d_dw_sned_output_tile:
   parameter_names_with_default_values:
-    OPERATOR: X
     NDIM: 3
     DTYPE: float
     TILE_SIZE: 3
@@ -16,10 +15,5 @@ conv2d_dw_sned_output_tile:
       - VALUE: float
   shader_variants:
     - NAME: conv2d_dw_sned_output_tile_3x3
-    - NAME: conv2d_dw_sned_output_tile_3x3_clamp
-      OPERATOR: clamp(X, A, B)
     - NAME: conv2d_dw_sned_output_tile_5x5
       TILE_SIZE: 5
-    - NAME: conv2d_dw_sned_output_tile_5x5_clamp
-      OPERATOR: clamp(X, A, B)
-      TILE_SIZE: 5

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

@@ -0,0 +1,86 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+/*
+ * Element-wise ops applied to a convolution output texel before it is stored:
+ *
+ *   out = activation(residual_op(conv, residual * scale, alpha))
+ *
+ * residual_op is add, sub or mul, and activation is clamp, relu or leaky_relu.
+ *
+ * The including shader binds t_residual, a texture3d with the output's sizes,
+ * and declares the specialization constants has_residual, has_scale,
+ * residual_op, has_activation and activation_op, and the uniforms
+ * residual_alpha, activation_a and activation_b. Operands that are absent are
+ * bound to other tensors of the dispatch and never read.
+ */
+
+#ifndef CONV2D_EPILOGUE_GLSLH
+#define CONV2D_EPILOGUE_GLSLH
+
+#define IN_T vec4
+#define OUT_T vec4
+#define ALPHA_T float
+#define PARAM_T float
+#define BINARY_OP_VEC4
+#define BINARY_OP_IS_FLOAT
+#define UNARY_OP_IS_FLOAT
+
+#include "binary/binary_op.glslh"
+#include "unary/unary_op.glslh"
+
+#undef IN_T
+#undef OUT_T
+#undef ALPHA_T
+#undef PARAM_T
+#undef BINARY_OP_VEC4
+#undef BINARY_OP_IS_FLOAT
+#undef UNARY_OP_IS_FLOAT
+
+// apply_binary_op() and apply_unary_op() are not reused for binary size: the
+// 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) {
+  switch (residual_op) {
+    case BINARY_OP_SUB:
+      return v - residual_alpha * residual;
+    case BINARY_OP_MUL:
+      return v * residual;
+    default: // BINARY_OP_ADD
+      return v + residual_alpha * residual;
+  }
+}
+
+vec4 apply_activation(const vec4 v) {
+  switch (activation_op) {
+    case UNARY_OP_RELU:
+      return max(v, vec4(0));
+    case UNARY_OP_LEAKY_RELU:
+      return leaky_relu(v, activation_a);
+    default: // UNARY_OP_CLAMP
+      return clamp(v, activation_a, 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) {
+  if (has_residual) {
+    vec4 residual = texelFetch(t_residual, pos, 0);
+    if (has_scale) {
+      residual *= scale;
+    }
+    v = apply_residual_op(v, residual);
+  }
+  if (has_activation) {
+    v = apply_activation(v);
+  }
+  return v;
+}
+
+#endif // CONV2D_EPILOGUE_GLSLH

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

@@ -40,6 +40,8 @@ $if WEIGHT_STORAGE == "buffer":
 $if PIXEL_TILES_ON_X:
   #define PIXEL_TILES_ON_X
 ${layout_declare_tensor(B, "r", "t_bias", DTYPE, "texture2d")}
+${layout_declare_tensor(B, "r", "t_residual", DTYPE, "texture3d")}
+${layout_declare_tensor(B, "r", "t_scale", DTYPE, "texture2d")}
 
 ${layout_declare_ubo(B, "ivec4", "in_sizes")}
 ${layout_declare_ubo(B, "ivec4", "out_sizes")}
@@ -49,14 +51,21 @@ layout(push_constant) uniform restrict Block {
   int stride_w;
   int padding_h;
   int padding_w;
-  float out_min;
-  float out_max;
+  float activation_a;
+  float activation_b;
+  float residual_alpha;
 };
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
 ${layout_declare_spec_const(C, "int", "stride_1_padding_0", "0")}
-${layout_declare_spec_const(C, "int", "activation_type", "0")}
+${layout_declare_spec_const(C, "bool", "has_activation", "false")}
+${layout_declare_spec_const(C, "int", "activation_op", "0")}
+${layout_declare_spec_const(C, "bool", "has_residual", "false")}
+${layout_declare_spec_const(C, "bool", "has_scale", "false")}
+${layout_declare_spec_const(C, "int", "residual_op", "0")}
+
+#include "conv2d_epilogue.glslh"
 
 #include "gemm/tile_utils/linear_fp_input_tile.glslh"
 #include "gemm/tile_utils/linear_fp_packed_weight_tile_load.glslh"
@@ -140,18 +149,22 @@ void store_output_tile_with_checks(
     const InCoords coords,
     const int n4_start,
     const int N4) {
+  vec4 scale[TILE_N4];
+  [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) {
+    scale[n4] = has_scale && n4_start + n4 < N4
+        ? texelFetch(t_scale, ivec2(n4_start + n4, 0), 0)
+        : vec4(1);
+  }
   // Reuse the coordinates computed for this invocation rather than dividing by
   // W_out again per output texel.
   [[unroll]] for (int m = 0; m < TILE_M; ++m) {
     [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) {
       if (coords.in_bounds[m] && n4_start + n4 < N4) {
-        VEC4_T texel = out_tile.data[m][n4];
-        if (activation_type == 1) {
-          texel = max(texel, VEC4_T(0.0));
-        } else if (activation_type == 2) {
-          texel = clamp(texel, VEC4_T(out_min), VEC4_T(out_max));
-        }
-        imageStore(t_out, ivec3(coords.out_pos[m], n4_start + n4), vec4(texel));
+        const ivec3 pos = ivec3(coords.out_pos[m], n4_start + n4);
+        imageStore(
+            t_out,
+            pos,
+            apply_epilogue(vec4(out_tile.data[m][n4]), pos, scale[n4]));
       }
     }
   }