[ExecuTorch][WebGPU] Add optimized conv2d op

[JCNTH]

Stack from ghstack (oldest at bottom):

• [ExecuTorch][WebGPU] apply_rotary_emb_hf op tests #20876
• [ExecuTorch][WebGPU] et_vk.apply_rotary_emb_hf: HuggingFace rotate-half RoPE runtime op (unblocks Qwen3) #20875
• [ExecuTorch][WebGPU] et_vk conv2d im2col-GEMM tests #20874
• [ExecuTorch][WebGPU] et_vk conv2d: route standard (groups==1) conv through an im2col tiled GEMM (1.1-2.4x) #20873
• [ExecuTorch][WebGPU] et_vk.sdpa channel-attention tests #20872
• [ExecuTorch][WebGPU] et_vk.sdpa: shape-route QK to a per-entry kernel for channel attention (15-30x faster) #20871
• [ExecuTorch][WebGPU] Migrate quantized_linear to make_compute_pipeline #20870
• [ExecuTorch][WebGPU] Migrate multi-dispatch ops to make_compute_pipeline #20869
• [ExecuTorch][WebGPU] Migrate 9 ops to make_compute_pipeline (layout:"auto") #20868
• [ExecuTorch][WebGPU] Adopt WebGPU layout:"auto" in make_compute_pipeline #20867
• [ExecuTorch][WebGPU] slice_copy native golden test #20866
• [ExecuTorch][WebGPU] slice_copy Double-arg scalar handling + view_copy/alias glue #20865
• [ExecuTorch][WebGPU] relu op tests #20864
• [ExecuTorch][WebGPU] Add relu op (shared unary handler; sigmoid adopts make_compute_pipeline) #20863
• [ExecuTorch][WebGPU] max_pool2d op tests #20862
• [ExecuTorch][WebGPU] Add max_pool2d op #20861
• [ExecuTorch][WebGPU] upsample_nearest2d op tests #20860
• [ExecuTorch][WebGPU] Add upsample_nearest2d op #20859
• [ExecuTorch][WebGPU] constant_pad_nd op tests #20858
• [ExecuTorch][WebGPU] Add constant_pad_nd op #20857
• [ExecuTorch][WebGPU] addmm op tests #20856
• [ExecuTorch][WebGPU] Add optimized addmm op #20855
• [ExecuTorch][WebGPU] embedding op tests #20854
• [ExecuTorch][WebGPU] Add embedding op #20853
• [ExecuTorch][WebGPU] SDPA op tests #20852
• [ExecuTorch][WebGPU] Add optimized SDPA op #20851
• [ExecuTorch][WebGPU] conv2d op tests #20850
• -> [ExecuTorch][WebGPU] Add optimized conv2d op #20849
• [ExecuTorch][WebGPU] linear_fp32 op tests #20848
• [ExecuTorch][WebGPU] Add optimized linear_fp32 op #20847
• [ExecuTorch][WebGPU] layer_norm op tests #20846
• [ExecuTorch][WebGPU] Add optimized layer_norm op #20845
• [ExecuTorch][WebGPU] gelu op tests #20844
• [ExecuTorch][WebGPU] Add optimized gelu op #20843
• [ExecuTorch][WebGPU] Shared pipeline + dispatch-grid + vec4-align helpers (op-suite foundation) #20842

Adds a general fp32 2D convolution kernel — direct and transposed — to the WebGPU backend, enabling the patch-embed / conv-stem and spatial-downsample layers of vision encoders (Florence-2 DaViT, SigLIP) to run GPU-accelerated in the browser.

Problem — aten.convolution.default had no WebGPU kernel, so any model with a convolutional stem or downsample block (the DaViT patch-embed and its strided downsample convs) could not lower into the delegate and broke the graph. Both the direct and the transposed forms serialize to the same aten.convolution.default op (distinguished only by the transposed arg), so a single handler has to cover both.

Solution

• Before: no conv kernel — a convolution in the exported graph fell out of the delegate.
• After: one conv2d_impl handler (registered once for aten.convolution.default) routes to three WGSL kernels by shape: conv2d_vec4.wgsl (vec4 fast path over input channels), conv2d.wgsl (scalar general path), and conv_transpose2d.wgsl (gather-form transposed conv). The transposed arg folds the transpose path into the same registration (a second WEBGPU_REGISTER_OP(aten.convolution.default) would be silently dropped), and among the non-transposed kernels the host picks vec4 vs scalar from the input-channels-per-group count.

Implementation

• Direct conv is one GPU thread per (b, oc, oh, ow) output element, looping over input-channel-per-group then (kh, kw), with continue guards skipping out-of-bounds padded taps; conv2d.wgsl accumulates scalar input[...] * weight[...].
• The vec4 path is selected only when icpg % 4 == 0 (use_vec4 in the handler); because NCHW's channel dim has stride IH*IW (not memory-contiguous), conv2d_vec4.wgsl gathers four strided scalar loads into a vec4<f32> and uses dot(in4, w4) — a register-packing vec4, not a coalesced load — cutting the input-channel loop trip count 4x. A real RGB stem (icpg=3) correctly falls to the scalar path.
• conv_transpose2d.wgsl implements the scatter-inversion as a gather: for each tap (kh, kw) an input row contributes only when (oh + pH - kh*dH) is divisible by sH and in range; weight is read in the un-flipped torch transposed layout [IC, OC/groups, KH, KW].
• All shape/stride/pad/dilation/groups constants ride in an 80-byte ConvParams uniform (16-byte-aligned); the handler parses stride/padding/dilation int-lists via utils::parse_hw (broadcasting a single value to both spatial dims).
• Dispatch uses the shared adaptive-grid helper utils::compute_dispatch_grid (workgroup clamped to the device max, default 256) with a 2D spill past the 65535 per-dimension ceiling; the stride_x override constant lets the shader decode i = gid.y*stride_x + gid.x. Pipeline creation, the uniform upload, the grid override constants, and the optional-bias binding go through the shared utils::make_compute_pipeline, utils::make_uniform, utils::make_grid_constants, and utils::make_optional_binding helpers (a 4-byte dummy storage satisfies the bias binding when bias is None, gated in WGSL by has_bias).
• Mirrors Vulkan backends/vulkan/runtime/graph/ops/impl/Convolution.cpp (its Conv2dMethod specializations Depthwise / Pointwise / SlidingWindow / Transposed collapse here into the vec4 / scalar / transpose routing).

Constraints — fp32 only (bails if nbytes != numel * sizeof(float)); 4D input/weight/output in NCHW; groups must divide both IC and OC and weight.dims[1] must equal IC/groups; output element count must fit uint32 for the (up to 2D) dispatch; non-zero output_padding is rejected on the non-transposed path (and output_padding >= stride on the transposed path); the fused et_vk.conv_with_clamp variant is not handled and would error clearly at load.

Co-authored-with: Claude Code.

Differential Revision: D110836669

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit 1c5c7b2 with merge base aceeb40 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest-nxp-neutron / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.

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.