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

Pull Request resolved: #24807

## Ulterior Motive

Closing the DSR FP16 gap to LiteRT. DSR has 83 skip connections of the form `conv(...) + x * scale`. ET-VK runs each as a separate `mul` and `add` (168 dispatches, 347 ms per inference on Galaxy S25 against LiteRT's 148 ms). This commit adds runtime support for fusing them into the preceding convolution; the next commit fuses them during lowering.

## Rationale

**What**: Every conv2d shader (direct, pointwise, depthwise, im2col GEMM and transposed) gains residual and scale bindings and applies an epilogue before storing:

```
out = activation(residual_op(conv + bias, residual * scale, alpha))
```

Specialization constants select each stage. A new `et_vk.conv2d.with_epilogue` op exposes it.

**Why**: One dispatch replaces conv, `mul` and `add`, and the conv output and scaled residual are never written to memory. Because every route applies the epilogue itself, the op never falls back to separate element-wise dispatches.

## Details

- **Epilogue** (`conv2d_epilogue.glslh`): calls `apply_binary_op()` for the residual op and `apply_unary_op()` for the activation.
  - The C++ side accepts add, sub and mul residual ops and clamp, relu and leaky_relu activations.
  - Inlining the `apply_unary_op()` switch changes Adreno's code generation for the 2x2 tiled direct shaders, even with the activation specialized away: AOC instruction counts rise 6–25% (`conv2d_2x2x8_border_float`: 532 → 660) with unchanged register counts.
  - Against a switch limited to the three supported activations, large FP16 3x3 and 5x5 direct convolutions are 2.1–2.9% slower on Pixel 8 Pro, and small 3x3 direct convolutions 4.4–5.7% slower on Galaxy S23. Calling `apply_unary_op()` avoids a second activation dispatcher.
- **Absent operands**: The residual slot binds the convolution input and the scale slot binds the packed bias, which have the slots' storage types. `has_residual` / `has_scale` gate the reads.
- **Scale**: A constant `[1, C, 1, 1]` or `[C, 1, 1]` residual scale. Each route prepacks it like the bias (texture2d, using the transposed weight layout for transposed convolutions) with `prepack_conv2d_residual_scale()`, next to the bias prepack, before building the bindings (`Conv2dEpilogueArgs`).
- **Validation**: `conv2d_epilogue_from_args()` throws at graph build when:
  - the residual is not texture3d or does not match the output's sizes, dtype and layout;
  - the scale is not a constant, is given without a residual, does not have the weight's dtype, or is not `[1, C, 1, 1]` or `[C, 1, 1]`. Other shapes with C values, such as `[C]`, broadcast along another dim in PyTorch;
  - the residual op or activation is unsupported.
  conv2d nodes pass the residual as their last resize argument, and their resize functions check that it still matches the resized output's sizes.
- **Clamp**
  - **Before**: Each shader family selected its clamp differently: a `clamp_out` spec constant and an `out_min` / `out_max` UBO (direct); an `activation_type` 0/1/2 spec constant (pointwise and im2col GEMM); separately generated `_clamp` shader variants (depthwise and transposed).
  - **After**: All use `has_activation` + `activation_op` with `activation_a` / `activation_b`. The `_clamp` variants are gone. `conv_with_clamp` builds a clamp epilogue.
- **Bindings and constants**
  - `conv2d_tiled`, `conv2d_dw` and `conv_transpose2d`: residual at 4, scale at 5, UBOs move from 4–8 to 6–10. UBO 10 holds `activation_a`, `activation_b` and `residual_alpha`.
  - `conv2d_pw_tiled`, `conv2d_gemm` and the depthwise output-tile shaders: residual and scale follow the bias; the clamp push constants become `activation_a`, `activation_b` and `residual_alpha`.
  - The epilogue specialization constants `has_activation`, `activation_op`, `has_residual`, `has_scale` and `residual_op` are contiguous in every shader, and `append_conv2d_epilogue_spec_vars()` appends all five. In `conv2d_tiled`, `has_bias` moves from ID 7 to 6, the group channel counts from 8–9 to 7–8, and `has_activation` (formerly `clamp_out`, 6) takes 9, ahead of the other epilogue constants at 10–13.
  - The unused `ngroups` (`conv2d_dw`) and `packed_dim` (`conv_transpose2d`) spec constants are removed. The transposed convolution no longer needs the `has_activation` resize argument that picked its `_clamp` variant.
- **Entry point**: `conv2d_with_epilogue_impl()` replaces `conv2d_impl()` and implements `aten.convolution`, `conv_with_clamp` and `conv2d.with_epilogue` for 4D inputs.
- **Depthwise tile override**: `conv2d_dw_impl` takes an optional `b1x1_override`. The benchmark's depthwise `b1x1` / `b4x2` routes call it instead of carrying a copy of the depthwise dispatch. The `b1x1` route now also selects the 1x1 tile for 5x5 kernels, as the Mali default already does.
- **Benchmark**
  - Adds a `conv2d_with_epilogue` reference and a transposed convolution reference; the conv2d reference previously rejected transposed convolutions. The epilogue reference rejects the scale shapes and the scale without a residual that the runtime rejects.
  - `et_vk.conv2d.with_epilogue` accepts the `direct`, `im2col`, `b1x1` and `b4x2` routes and reports throughput.
  - Adds 204 correctness cases. 200 cover 5 epilogues × FP32/FP16 on 20 shape and route combinations:
    - one shape per route: direct, pointwise, im2col, the three depthwise shaders, transposed;
    - partial output texels or tiles on the direct, pointwise, im2col, depthwise tile and transposed routes;
    - a strided and dilated direct conv and a grouped direct conv.
  - The other 4 run a pointwise, direct, im2col and depthwise epilogue through a full → shrunk → full resize.
  - Adds the `operator_dynamic_resize` case check. It builds an operator case, which must expand to one case, once, then for each step resizes the staged inputs, checks the output sizes and validates the output. The q8ta streaming-shrink check now uses it instead of its own copy.
  - The depthwise `b1x1` / `b4x2` routes report batch > 1 as unsupported (SKIP), since the depthwise shaders support batch 1 only.
  - The route inspector now recognizes the tiled direct shader names.
  - `conv2d_gemm_sustained.json` binds the new GEMM slots.

Usage:

```cpp
// Parses and validates the epilogue arguments of a conv2d op
const Conv2dEpilogue epilogue = conv2d_epilogue_from_args(graph, args);
conv2d_direct_impl(graph, in, weight, bias, stride, padding, dilation,
                   transposed, output_padding, groups, out, epilogue);

// Inside a route
const ValueRef packed_scale =
    prepack_conv2d_residual_scale(graph, epilogue, weight, transposed);
const Conv2dEpilogueArgs epilogue_args(epilogue, in, packed_bias, packed_scale);
// Bind epilogue_args.residual and epilogue_args.scale, append the epilogue's
// specialization constants, and pass epilogue.residual as the last resize arg.
append_conv2d_epilogue_spec_vars(spec_vars, epilogue);
```
ghstack-source-id: 469036183
@exported-using-ghexport

Differential Revision: [D130216287](https://our.internmc.facebook.com/intern/diff/D130216287/)

Author: SS-JIA

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