[SS-JIA]

Stack from ghstack (oldest at bottom):

• [ET-VK][unary] Compute vec4 leaky_relu with a single mix #24813
• [ET-VK][conv2d] Fuse residual adds and activations into conv2d during lowering #24808
• -> [ET-VK][conv2d] Fuse residual and activation epilogues into conv2d shaders #24807
• [ET-VK][tools] Write etvk_ops_bench JSON files compactly #24806

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:

// 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);

Differential Revision: D130216287

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24807

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit de6a311 with merge base e5ffebe ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.