[ExecuTorch][WebGPU] Shared pipeline + dispatch-grid + vec4-align helpers (op-suite foundation)#21243
Merged
Conversation
…pers (op-suite foundation) Pull Request resolved: #20842 **Foundation for the vision/voice/seg WebGPU op suite: the shared helpers every op builds on.** Adds to `WebGPUUtils.h` the `make_compute_pipeline` bundle (`BindingSpec`, `ComputePipelineBundle`, `make_storage_zeros`), `check_vec4_aligned`, and the validation/dispatch-config helpers every handler reuses: `check_fp32` and `check_elementwise_fp32_io` (fp32 byte-size + elementwise input/output guards), `checked_u32` (the `uint64_t`->`uint32_t` dispatch-count guard), `make_wg_size_constant` (the 1-element override-constant sibling of `make_grid_constants`), and `compute_tile_grid_2d` (the tiled-GEMM 2D grid, device-limit-queried); a new `WebGPUDispatchMath.h` with the near-square 2D dispatch-grid fold (`div_up`, `numel`); a `WebGPUGraph::add_dispatch_2d` convenience for the 2D dispatch fill; shared `WebGPUGraph` hooks; and op-test-driver infra (+ `test_webgpu_utils`). No op is registered here — each op lands on top using these helpers. Co-authored-with: Claude Code. ghstack-source-id: 405949691 @exported-using-ghexport Differential Revision: [D110836666](https://our.internmc.facebook.com/intern/diff/D110836666/)
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21243
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
Pull Request resolved: #20843 **Adds the `gelu` activation to the WebGPU backend.** GELU is the feed-forward (MLP) activation in the vision and VLM transformer blocks this backend targets — the Florence-2 DaViT encoder, BART, and the Whisper/Voxtral stacks — so it is on the critical path for delegating those models without graph breaks. **Problem** — The WebGPU backend had no `aten.gelu.default` kernel, so any model whose FFN uses GELU could not be fully delegated. PyTorch's default is the exact (erf) formulation, and the target encoders use it; the tanh approximation must also be supported. **Solution** - Before: `aten.gelu.default` was unsupported — no kernel, forcing a partition break around every GELU. - After: a single `gelu.wgsl` computes GELU on the GPU. The exact (erf) vs tanh choice is resolved by compiling the matching entry point (`main_erf` / `main_tanh`) into the pipeline, so each dispatch runs only its own formula — no per-invocation `select()` and no double-eval of both branches. **Implementation** - 1D dispatch with one thread per 4 elements: a vec4 body + scalar-tail idiom loads up to 4 elements via `select()`-guarded scalar reads, computes GELU as one `vec4<f32>` op, and scatters back only the in-bounds lanes — so arbitrary FFN widths that are not a multiple of 4 are handled without a separate remainder pass. - The erf path uses the Abramowitz & Stegun 7.1.26 rational approximation (max abs err ~1.5e-7); the tanh path uses the clamped `0.5*x*(1+tanh(...))` form. - The `approximate` arg (args[1]: "none" selects erf, anything else selects tanh) picks the entry point at build time through the shared `utils::make_compute_pipeline`; workgroup size comes from `utils::clamp_workgroup_size`, the dispatch count from `utils::compute_1d_workgroup_count`, and the uniform from `utils::make_uniform`. - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp` (its `gelu` handler; args[1] is the `approximate` string). Note the WebGPU kernel additionally implements the exact/erf path, which the Vulkan handler does not currently support. **Constraints** — fp32 only (both operands must be 4-byte aligned, enforced in the handler); input and output must have identical byte size; 1D dispatch only (throws if the workgroup count would exceed the 65535 cap). Element count need not be a multiple of 4 (scalar tail); static shape (no resize hook). Co-authored-with: Claude Code. ghstack-source-id: 405949686 @exported-using-ghexport Differential Revision: [D110836674](https://our.internmc.facebook.com/intern/diff/D110836674/)
Pull Request resolved: #20844 Splits the `gelu` op tests into their own diff, stacked directly above the `gelu` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_gelu.py` and registers the `gelu` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949721 @exported-using-ghexport Differential Revision: [D111072717](https://our.internmc.facebook.com/intern/diff/D111072717/)
Pull Request resolved: #20845 **Adds `native_layer_norm` to the WebGPU backend.** LayerNorm is pervasive in the BART decoder and the Florence-2 DaViT vision encoder (and the Whisper/Voxtral hidden dims 768/1024/1280/3072), so it is required to delegate those transformer stacks end-to-end. **Problem** — The backend had no `aten.native_layer_norm.default` kernel. LayerNorm also needs a numerically robust mean/variance (a naive `E[x^2]-E[x]^2` cancels badly on large-mean/small-variance activations), and some graphs (the group_norm LN-reframe) pass `None` for weight and bias, so the affine step must be optional. **Solution** - Before: `aten.native_layer_norm.default` was unsupported. - After: `native_layer_norm.wgsl` computes per-row mean, rstd, and the normalized (optionally affine) output in a single kernel launch — one workgroup per row — writing all three outputs (`out`, `mean`, `rstd`) of the op's ValueList. **Implementation** - Single-pass, numerically robust mean+variance via Chan et al.'s parallel Welford: each of the 64 threads folds its strided `vec4<f32>` slice of the row into a running `(n, mean, M2)`, then a shared-memory tree reduction pairwise-merges the per-thread triples — no `E[x^2]-E[x]^2` cancellation. - `t_in`/`t_out`/`t_weight`/`t_bias` are viewed as `array<vec4<f32>>` over the row width for wide loads/stores; the second pass re-streams the row applying `(x-mean)*rstd` and the affine only when `has_affine == 1`. - Rows past the 65535 per-dimension ceiling are handled by a near-square 2D workgroup grid: `utils::compute_row_dispatch_grid` computes `count_x`/`count_y` and a `stride_x` override constant, and the shader decodes the flat row index as `wid.y * stride_x + wid.x`. - Weight/bias are optional: when either is `None` the handler binds dummy storage via `utils::make_optional_binding` and sets `has_affine == 0`; the pipeline is built with the shared `utils::make_compute_pipeline`, `epsilon` read via `utils::scalar_or`. - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp` (same arg layout and `[out, mean, rstd]` ValueList; the WebGPU kernel additionally supports the no-affine path, which the Vulkan handler rejects). **Constraints** — fp32 only (byte-size must equal `numel * sizeof(float)`); normalizes over the last dim only; the last dim (`row_width`) must be a multiple of 4 for the vec4 view (all in-scope hidden dims are); non-empty input; 2D dispatch grid for the row count. Co-authored-with: Claude Code. ghstack-source-id: 405949712 @exported-using-ghexport Differential Revision: [D110836681](https://our.internmc.facebook.com/intern/diff/D110836681/)
Pull Request resolved: #20846 Splits the `layer_norm` op tests into their own diff, stacked directly above the `layer_norm` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_layer_norm.py` and registers the `layer_norm` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949713 @exported-using-ghexport Differential Revision: [D111072718](https://our.internmc.facebook.com/intern/diff/D111072718/)
Pull Request resolved: #20848 Splits the `linear_fp32` op tests into their own diff, stacked directly above the `linear_fp32` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_linear_fp32.py` and registers the `linear_fp32` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949719 @exported-using-ghexport Differential Revision: [D111072710](https://our.internmc.facebook.com/intern/diff/D111072710/)
Pull Request resolved: #20849 **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. ghstack-source-id: 405949715 @exported-using-ghexport Differential Revision: [D110836669](https://our.internmc.facebook.com/intern/diff/D110836669/)
Pull Request resolved: #20850 Splits the `conv2d` op tests into their own diff, stacked directly above the `conv2d` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_conv2d.py` and registers the `conv2d` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949711 @exported-using-ghexport Differential Revision: [D111072711](https://our.internmc.facebook.com/intern/diff/D111072711/)
Pull Request resolved: #20851 **Adds a non-causal fused scaled-dot-product-attention kernel to the WebGPU backend**, enabling the attention blocks of vision encoders (Florence-2 DaViT / SigLIP, SAM2) and non-causal cross-attention (BART) to run end-to-end on GPU. **Problem** — the delegate had no kernel for `et_vk.sdpa.default`, the plain non-causal attention `softmax(q @ kᵀ * scale + attn_mask) @ v` that the et_vk source transform plugs into every vision-encoder attention block. Without it the attention layers broke the graph. This is distinct from the causal KV-cache `sdpa_with_kv_cache`: no cache, an optional additive mask, and it must handle asymmetric sequence lengths (`S_q != S_kv`, e.g. a Hiera pooled query or cross-attention). **Solution** - Before: no non-causal fused-attention kernel — vision-encoder attention blocks could not lower into the delegate. - After: `et_vk_sdpa_impl` chains three compute dispatches over `[B, H, S, D]` (DSHB) row-major tensors — `et_vk_sdpa_qk.wgsl` computes the scaled `Q·Kᵀ` scores plus an optional additive mask, the reused `sdpa_softmax.wgsl` does the row-wise softmax, and `et_vk_sdpa_av.wgsl` computes `softmax · V` into the output. **Implementation** - QK phase (`et_vk_sdpa_qk.wgsl`): one GPU thread per `(b, h, s)` row of the `[B, H, S_q, S_kv]` attention-weight buffer; `q`/`k` are bound as `array<vec4<f32>>` over `D`, and the thread loops over `c` (key positions) and `d4` (`D/4`) accumulating `dot(q4, k4)`, then multiplies by `scale` and adds `mask[...]` when `has_mask`. Row count `B*H*S_q` stays well under the 1D dispatch limit for any ViT. - Softmax phase: reuses the existing `sdpa_softmax.wgsl` (one workgroup per row, hardcoded `@workgroup_size(64,1,1)`) dispatched on a near-square 2D workgroup grid past the 65535 ceiling; the shader recovers the flat row index from `@builtin(num_workgroups)`, so no override constant is needed. - AV phase (`et_vk_sdpa_av.wgsl`): one thread per `(b, h, s, d4)` computing a `vec4<f32>` of four output elements; `v`/`out` are bound as `array<vec4<f32>>` over `D` and the thread contracts over `c` scalar (`S_kv` is not guaranteed `% 4 == 0`), accumulating `sm[...] * v4`. - Two fp32 scratch buffers (`attn`, `softmax`, each `B*H*S_q*S_kv`) are allocated via `graph.create_scratch_buffer`; `scale` defaults to `1/sqrt(D)` when the arg is `None` or takes the `Double` value; the three uniforms are compact structs (`QkParams`/`AvParams` 32 bytes, `SoftmaxParams` 16 bytes). - Uses the shared runtime helpers: `utils::make_compute_pipeline`, `utils::make_uniform`, `utils::check_vec4_aligned` (guards `D % 4 == 0`), `utils::clamp_workgroup_size` + `utils::compute_1d_workgroup_count` (QK / AV grids), `utils::compute_2d_workgroup_count` (softmax grid), and `utils::make_optional_binding` (a 4-byte dummy satisfies the mask binding when absent; the shader never reads it under `has_mask == 0`). - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/SDPA.cpp` `SDPAMode::FUSED` (general non-cache SDPA, `[B, H, S, D]` DSHB layout, optional additive `attn_mask`, optional `scale`, unpadded fp32 attention weights). **Constraints** — fp32 only (bails on `q`/`out` byte mismatch); non-causal (causality is expressed as a baked additive mask input, not a code path); `D % 4 == 0` for the vec4 QK/AV kernels (every model in scope uses `D=64` or `128`); `q` rank ≥ 3; `k.dims == v.dims`, `q`/`k`/`v` share `H` and `D` and all leading batch dims, and `out.dims == q.dims`; asymmetric `S_q != S_kv` is supported (reduces bit-identically to self-attention when equal); a supplied mask must be `[B, H, S_q, S_kv]` fp32. Co-authored-with: Claude Code. ghstack-source-id: 405949724 @exported-using-ghexport Differential Revision: [D110836679](https://our.internmc.facebook.com/intern/diff/D110836679/)
Pull Request resolved: #20852 Splits the `SDPA` op tests into their own diff, stacked directly above the `SDPA` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_et_vk_sdpa.py` and registers the `et_vk_sdpa` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949730 @exported-using-ghexport Differential Revision: [D111072726](https://our.internmc.facebook.com/intern/diff/D111072726/)
Pull Request resolved: #20854 Splits the `embedding` op tests into their own diff, stacked directly above the `embedding` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_embedding.py` and registers the `embedding` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949731 @exported-using-ghexport Differential Revision: [D111072723](https://our.internmc.facebook.com/intern/diff/D111072723/)
Pull Request resolved: #20855 **fp32 `aten.addmm.default` now runs on the WebGPU backend as a shared-memory-tiled GEMM, so the fused bias-plus-matmul that HuggingFace `Linear` lowers to executes on-device.** In the Florence-2 BART and DaViT graphs the HF `nn.Linear` layers lower to `aten.addmm` (not `aten.linear`), making addmm the dominant dense GEMM in those stacks. **Problem** — The WebGPU backend had no `aten.addmm.default` kernel, so any graph whose linears lowered through addmm delegated at export but failed at runtime load. As with plain matmul, a naive per-output kernel would re-stream both operands from global memory with no reuse. **Solution** — Before: `aten.addmm.default` had no runtime kernel (load-time failure). After: the handler binds `self` (bias), `mat1 [M,K]`, `mat2 [K,N]`, and the output, and dispatches a tiled GEMM computing `out = beta*self + alpha*(mat1 @ mat2)`. Each workgroup cooperatively stages 32x32 slabs of `mat1` and `mat2` into workgroup shared memory, `workgroupBarrier`-syncs, and accumulates across `ceil(K/32)` k-tiles so the coalesced tile load is amortized over the whole output tile. The `beta*self + alpha*acc` epilogue broadcasts `self` from either `[N]` (the common HF bias case) or a full `[M,N]`. **Implementation** — Shares the tiled-GEMM approach with `linear_fp32`: `TILE = 32`, `RPT = 4`, `@workgroup_size(8, 8, 1)`, two `var<workgroup>` `a_sub`/`b_sub` arrays, each thread accumulating a 4x4 register sub-tile over `ceil(K/32)` tiles. It re-derives the B-side read for `mat2`'s actual `[K,N]` layout: `read_b` reads `t_mat2[krow*N + col]` (N contiguous), unlike `linear`'s transposed `[N,K]` weight read. No vec4 variant is provided: with `mat2 [K,N]` the N dimension is contiguous, so a vec4-over-K view would require a strided gather on the mat2 side that erodes the benefit, and the cooperative shared-memory tile load already closes the coalescing gap. Epilogue selects `t_self[r*N + c]` when `self_2d`, else the broadcast `t_self[c]`, and writes `beta*self_val + alpha*acc[ir][ic]`. `beta`/`alpha` are read with the shared `utils::scalar_or` (Double/Int/Null-permissive, default `1.0`). Dispatch is a 2D grid `ceil(N/32) x ceil(M/32)`, throwing if either dimension exceeds 65535. Uses the shared `utils::make_compute_pipeline` and `utils::make_uniform` (32-byte `AddmmParams {M,N,K,self_2d,beta,alpha,_pad[2]}`). Validates non-null buffers, 2D `mat1`/`mat2`/`out`, `mat2` shape `== [K,N]`, and `self` broadcastable from `[N]` or `[M,N]`. **Constraints** — fp32 only (output byte-size check); `mat1`/`mat2`/`out` must be 2D with `mat2` exactly `[K,N]`; `self` must be `[N]` or `[M,N]`; scalar-only tiled path (no vec4); 2D dispatch capped at 65535 workgroups per dimension; fixed `@workgroup_size(8, 8, 1)`. Co-authored-with: Claude Code. ghstack-source-id: 405949737 @exported-using-ghexport Differential Revision: [D110836673](https://our.internmc.facebook.com/intern/diff/D110836673/)
Pull Request resolved: #20856 Splits the `addmm` op tests into their own diff, stacked directly above the `addmm` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_addmm.py` and registers the `addmm` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949733 @exported-using-ghexport Differential Revision: [D111072709](https://our.internmc.facebook.com/intern/diff/D111072709/)
Pull Request resolved: #20857 **Adds `aten.constant_pad_nd.default` to the WebGPU backend, unblocking the DaViT window-padding path in vision models.** **Problem** — the backend had no `constant_pad_nd` handler, so any graph that pads a tensor (e.g. DaViT's window partitioning) could not fully delegate to WebGPU and threw at runtime. **Solution** — a single gather/fill compute kernel: Before — no handler; `aten.constant_pad_nd.default` unsupported at runtime. After — one thread per output element gathers the source element when its coordinates land inside the input, otherwise writes the constant fill `value`. **Implementation**: - The handler right-aligns the (rank 1..4) dims into fixed `vec4<u32>` params (`out_dims`, `in_dims`, `left`); leading slots get extent 1 / pad 0, so the WGSL is rank-agnostic and always iterates 4 dims. - The `pad` `IntList` is reversed-dim (innermost-first `(left, right)` pairs); the handler expands it to per-dim `left`/`right`, then validates `out.dims[d] == in.dims[d] + left[d] + right[d]` before any buffer allocation (loud-fail, no leak-on-throw). - The kernel decodes each output element's 4D coords (last dim fastest), subtracts each dim's `left` pad as an unsigned wrapping subtract (a negative coord wraps to a huge value and is rejected by the `< in_dims` bound check); if all four coords are in-bounds it copies `inp[flat_in]`, else it writes `value` — a pure copy/fill, so bit-exact. - The fill `value` is read via `utils::scalar_or` (a `Scalar` may serialize as `Int` or `Double`), defaulting to `0`. - Adaptive 1D->2D dispatch via `utils::compute_dispatch_grid` (workgroup size clamped to the device max, up to 256, plus a 2D spill past the 65535 workgroup-count ceiling; the `stride_x` override lets the shader decode `i = gid.y*stride_x + gid.x`). - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/Pad.cpp` (same reversed-dim `(before, after)` pad convention and `constant_pad_nd` resize logic). **Constraints** — fp32 only (`nbytes == numel*4` guard); rank 1..4; `pad` must be even-length and no longer than the rank; the output element count must fit `u32` (`<= 2^32`). Co-authored-with: Claude Code. ghstack-source-id: 405949740 @exported-using-ghexport Differential Revision: [D110836672](https://our.internmc.facebook.com/intern/diff/D110836672/)
Pull Request resolved: #20858 Splits the `constant_pad_nd` op tests into their own diff, stacked directly above the `constant_pad_nd` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_constant_pad_nd.py` and registers the `constant_pad_nd` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949732 @exported-using-ghexport Differential Revision: [D111072712](https://our.internmc.facebook.com/intern/diff/D111072712/)
Pull Request resolved: #20859 **Adds `aten.upsample_nearest2d.vec` to the WebGPU backend, enabling the SAM2/SAM3 pixel-decoder / FPN nearest-neighbour upsample path.** **Problem** — nearest-neighbour 2D upsample (`F.interpolate(..., mode="nearest")`) had no WebGPU kernel, blocking the FPN 2x upsample chain in the SAM2/SAM3 pixel decoder. **Solution** — Before — no handler; `aten.upsample_nearest2d.vec` unsupported. After — one thread per output element `(n, c, oh, ow)` maps back to its nearest source pixel and copies it. **Implementation**: - The kernel computes the source index with ATen's legacy nearest formula `ih = floor(oh*IH/OH)`, `iw = floor(ow*IW/OW)` (integer division), then copies `inp[((n*C+c)*IH+ih)*IW+iw]` — a plain NCHW row-major gather, bit-exact. - `OH`/`OW` are taken from the output tensor's own dims (not re-derived from the `size`/`scale` args), and the handler validates that `N`/`C` match between input and output. - Adaptive 1D->2D dispatch via `utils::compute_dispatch_grid` (workgroup size clamped to the device max, up to 256, plus a 2D spill past the 65535 ceiling; `stride_x` decode). - Divergence from Vulkan (intentional): this does NOT mirror Vulkan `backends/vulkan/runtime/graph/ops/impl/Upsample.cpp`, which computes the source index with the half-pixel-center reciprocal-scale formula. Half-pixel-center is correct for bilinear but wrong for non-exact nearest; the ATen `nearest_neighbor_compute_source_index` (`UpSample.h`) `floor(oh*IH/OH)` form is the one that matches PyTorch's `mode="nearest"`. The two agree on exact-integer ratios (e.g. 2x) but diverge on non-integer ratios (e.g. 5->8), which the op-test's `non_2x_ratio` case pins down. **Constraints** — fp32 only; 4D NCHW input and output; `nearest` mode only; non-zero spatial dims; the output element count must fit `u32`. Co-authored-with: Claude Code. ghstack-source-id: 405949741 @exported-using-ghexport Differential Revision: [D110836668](https://our.internmc.facebook.com/intern/diff/D110836668/)
Pull Request resolved: #20860 Splits the `upsample_nearest2d` op tests into their own diff, stacked directly above the `upsample_nearest2d` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_upsample_nearest2d.py` and registers the `upsample_nearest2d` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949736 @exported-using-ghexport Differential Revision: [D111072707](https://our.internmc.facebook.com/intern/diff/D111072707/)
Pull Request resolved: #20861 **Adds `aten.max_pool2d_with_indices.default` to the WebGPU backend, enabling the SAM2 Hiera q_pool path (values, with optional indices).** **Problem** — max pooling (`F.max_pool2d`, which decomposes to `max_pool2d_with_indices`) had no WebGPU kernel, blocking the SAM2 Hiera `q_pool` downsampling stages. **Solution** — Before — no handler; the multi-output `max_pool2d_with_indices` unsupported. After — one thread per output element `(n, c, oh, ow)` gathers its pooling window and writes the max (and, when the graph requests them, the argmax indices). **Implementation**: - The kernel iterates the `kH x kW` window with general `stride`/`padding`/`dilation`, skips out-of-range (padding) cells, and tracks the running max; `best` initialises to `-3.4e38` (a large finite negative, because Dawn/Tint rejects the exact `-FLT_MAX` literal). - Argmax is ALWAYS tracked; a `write_indices` override (a compile-time spec constant) gates only the final `out_idx` store, so the values-only and with-indices paths run one shader and stay bit-identical on `out_vals`. Indices are the flat `ih*IW+iw` spatial-plane offset (matching `torch.nn.functional.max_pool2d(return_indices=True)`, not an absolute NCHW offset). - The out is a `ValueList` `[values, indices]`; when indices are not requested the handler binds a tiny dummy storage buffer to slot 3 (the shader never writes it) rather than allocating a real indices tensor — mirroring the `dummy_affine` pattern in `NativeLayerNorm.cpp`. When indices ARE requested it validates that the indices tensor is `int32` (4 bytes/elem) and shares the values shape, else throws (the kernel writes `i32`, so an `int64` indices tensor would mis-stride the write). - The handler parses `kernel_size`/`stride`/`padding`/`dilation` via `utils::parse_hw` (a single value broadcasts to both spatial dims; `stride` defaults to `kernel_size` when empty), derives `OH`/`OW` from the pooling formula, and validates them against the serialized values output (loud-fail on a `ceil_mode` / layout mismatch). - Adaptive 1D->2D dispatch via `utils::compute_dispatch_grid`. - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/Pool.cpp` (the `ValueList[values, indices]` shape, the `write_indices` spec constant, unconditional argmax tracking, and 32-bit indices despite ATen's int64 schema). **Constraints** — fp32 values; 4D NCHW; general `kernel`/`stride`/`padding`/`dilation`; `ceil_mode` unsupported (the output-shape validation assumes floor); the optional indices output must be `int32`; the output element count must fit `u32`. Co-authored-with: Claude Code. ghstack-source-id: 405949738 @exported-using-ghexport Differential Revision: [D110836680](https://our.internmc.facebook.com/intern/diff/D110836680/)
Pull Request resolved: #20862 Splits the `max_pool2d` op tests into their own diff, stacked directly above the `max_pool2d` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_max_pool2d.py` and registers the `max_pool2d` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949735 @exported-using-ghexport Differential Revision: [D111072724](https://our.internmc.facebook.com/intern/diff/D111072724/)
…s make_compute_pipeline) Pull Request resolved: #20863 **Adds `relu` to the WebGPU backend via a shared elementwise-unary handler.** ReLU is on the SAM2/SAM3 mask-decoder MLP path, so it is needed to delegate those decoders. **Problem** — The backend had no `aten.relu.default` kernel, and `sigmoid` (the only prior unary op) built its compute pipeline inline rather than through the shared helper — duplicating the bind-group/dispatch boilerplate that a second unary op would repeat. **Solution** - Before: `sigmoid` was implemented with a bespoke inline pipeline; there was no `relu`. - After: a single generic `add_unary_op` helper (in `runtime/ops/sigmoid/UnaryOp.cpp`) builds the input/output/params binding and dispatch for any elementwise-unary WGSL; `sigmoid_impl` and the new `relu_impl` are thin wrappers over it, so `sigmoid` now goes through the same `utils::make_compute_pipeline` path as `relu`. `relu.wgsl` is a one-element-per-thread `output[idx] = max(input[idx], 0.0)`. **Implementation** - `add_unary_op(graph, in, out, wgsl_source, wg_size_x, op_name)` centralizes: the fp32/4-byte-alignment and same-size guards, `utils::clamp_workgroup_size` + `utils::compute_1d_workgroup_count` for the 1D dispatch, the `wg_size` override constant, the uniform (`num_elements`) via `utils::make_uniform`, and the three-binding pipeline via `utils::make_compute_pipeline`. - Dynamic shapes are supported: a `graph.add_tensor_resize_hook` recomputes `num_elements`, rewrites the uniform via `wgpuQueueWriteBuffer`, and updates the dispatch's workgroup count for the live shape; the graph owns the uniform buffer so the hook can rewrite it. - Both ops self-register: `aten.sigmoid.default -> sigmoid_impl` and `aten.relu.default -> relu_impl`. - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp` (`add_unary_op_node`); Vulkan expresses `relu` as `clamp(0, inf)`, whereas this kernel uses a direct `max(x, 0.0)`. **Constraints** — fp32 only (both operands 4-byte aligned); input and output must have identical byte size (same-shape elementwise); 1D dispatch only (throws past the 65535 workgroup cap). Co-authored-with: Claude Code. ghstack-source-id: 405949742 @exported-using-ghexport Differential Revision: [D110836664](https://our.internmc.facebook.com/intern/diff/D110836664/)
Pull Request resolved: #20864 Splits the `relu` op tests into their own diff, stacked directly above the `relu` op — keeping an op and its tests in separate diffs (op below, tests above) per this backend's convention. Adds `test/ops/test_relu.py` and registers the `relu` `@register_op_test` suite in `test/op_tests/cases.py`. Co-authored-with: Claude Code. ghstack-source-id: 405949743 @exported-using-ghexport Differential Revision: [D111072725](https://our.internmc.facebook.com/intern/diff/D111072725/)
…y/alias glue Pull Request resolved: #20865 **Hardens `slice_copy` scalar-arg decoding against edge-dialect `Double`-serialized indices and adds the `view_copy` / `alias` / `clone` reshape pass-throughs needed for graph glue.** **Problem** — two graph-glue gaps: (1) the edge dialect sometimes serializes an integer `slice` index (`dim` / `start` / `end` / `step`) as a floating-point `Double` (e.g. a `0` start), which the `slice` handler rejected as unsupported; (2) contiguous reshape / aliasing ops (`view_copy`, `alias_copy`, `clone`, `_clone_dim_order`) had no handler, breaking otherwise-delegatable subgraphs. **Solution** — Before — a `Double`-typed slice index threw, and reshape/alias ops had no handler. After — `slice_copy` scalar reads accept an integral `Double` (truncating to the int index) and reject only a genuinely fractional one, while `SymInt` (dynamic start/end) and `Null` (default) still resolve as before; and `view_copy` / `alias_copy` / `clone` / `_clone_dim_order` all lower to a single contiguous flat copy (or an in-place no-op when input and output alias the same buffer). **Implementation**: - `read_scalar` (`dim` / `step`) and `read_index` (`start` / `end`) switch on the value type: `Int` (`INT64_MAX` -> default), `Double` -> truncated int iff it round-trips (`static_cast<int64_t>(d)` back to `d`) else throw `"non-integral ..."` (NaN and out-of-`int64`-range doubles are rejected before the cast, since casting them is UB), `Null` -> default; `read_index` additionally resolves a `SymInt` via `read_symint`. - The slice kernel is an index gather: `out_bufi -> in_bufi` by walking per-dim strides, with the sliced dim's input coord `= start + coord*step`; dynamic `start` / `end` / `SymInt` are handled by a resize hook that recomputes the live `out[dim]` length (ceiling division) and rewrites the meta/params uniforms plus the dispatch count (mirrors Vulkan `resize_slice_copy_node`). - `add_flat_copy` (shared by all the reshape/alias ops) type-checks both args are tensors, guards 4-byte alignment and equal `nbytes` (a view preserves `numel`, so this also prevents an OOB copy), then either skips the copy when `in.buffer == out.buffer` (aliased, already in place; `CopyBufferToBuffer` rejects `src == dst`) or emits a buffer-to-buffer copy; a resize hook keeps the live output shape and copy byte-count in sync under dynamic shapes. - `_clone_dim_order` ignores its `dim_order` arg (the AOT pass elides it via shape and dtype). - Mirrors Vulkan `backends/vulkan/runtime/graph/ops/impl/Slice.cpp` (`normalize_idx` / `INT64_MAX` default and the ceiling-division length) and `backends/vulkan/runtime/graph/ops/impl/View.cpp` (the `view_buffer` no-remap contiguous reshape). **Constraints** — fp32 (4-byte-aligned) operands; `slice` requires `step >= 1` and an in-range `dim`; a fractional `Double` index is a hard error, not truncated; `view` / `alias` / `clone` require equal input/output `numel` (contiguous reshape only, no layout remap). Co-authored-with: Claude Code. ghstack-source-id: 405949746 @exported-using-ghexport Differential Revision: [D110836670](https://our.internmc.facebook.com/intern/diff/D110836670/)
Pull Request resolved: #20866 Splits the `SliceDoubleStart` native golden test into its own diff, stacked directly above the `slice_copy` / `view_copy` glue op (op below, tests above). Adds the double-start slice regression case to `test/test_webgpu_native.cpp`, covering an edge-dialect-serialized Double-typed slice `start` argument. Co-authored-with: Claude Code. ghstack-source-id: 405949754 @exported-using-ghexport Differential Revision: [D111072708](https://our.internmc.facebook.com/intern/diff/D111072708/)
… for channel attention (15-30x faster) Pull Request resolved: #20871 **Problem:** the fused `et_vk.sdpa` QK kernel runs one thread per (b,h,s) row with vec4 loads — ideal for standard attention, but on channel attention (DaViT/Florence, where `S_q = head_dim ~= 32`) `num_rows = B*H*S_q` is tiny, so only a handful of workgroups run serially over a huge `S_kv*D`, starving the GPU (the (2,1,1)@103ms dispatch). **Solution:** add a per-entry QK kernel (one thread per (b,h,s,c) attention entry, 2D-folded) and host-route to it when `num_rows` is below an occupancy floor (4096); standard attention keeps the per-row + vec4 path unchanged. **Before:** `et_vk_sdpa_qk` (per-row, vec4) — the only QK kernel; channel-attn shapes are occupancy-starved. **After:** router picks `et_vk_sdpa_qk_entry` (per-entry, scalar, 2D-folded) for small `num_rows`, else the unchanged per-row kernel. **Implementation:** - New `et_vk_sdpa_qk_entry.wgsl` (+ generated header) — same bindings and `Params` as the per-row kernel, so it is a drop-in under `layout:"auto"`; writes a layout-identical `attn[B,H,S_q,S_kv]` (`attn[idx]`), so softmax/AV are unchanged and either branch is numerically correct — the floor is a pure perf knob. - `EtVkSdpa.cpp` selects the shader and a 2D dispatch (`compute_2d_workgroup_count`, mirroring the softmax grid) when routed, else the existing 1D per-row dispatch; the grid + dispatch-limit check is computed up front (throw before any buffer alloc -> no leak). - Mirrors the codebase's host shape-router precedents (`LinearFp32.cpp` `K%4` vec4 selection, `Sdpa.cpp` variant selection). **Constraints:** per-entry drops vec4, so it only wins when the per-row path is occupancy-starved (small `num_rows`); the 4096 floor is Canary-tuned. Co-authored-with: Claude Code. ghstack-source-id: 405949750 @exported-using-ghexport Differential Revision: [D110994975](https://our.internmc.facebook.com/intern/diff/D110994975/)
Pull Request resolved: #20872 Splits the channel-attention routing case out of the `et_vk.sdpa` per-entry-QK op diff into its own test diff, stacked directly above it (op below, tests above). Adds the `chattn_davit` case to the `et_vk_sdpa` suite in `test/op_tests/cases.py`, exercising the per-entry QK kernel path (num_rows below the per-row floor). Co-authored-with: Claude Code. ghstack-source-id: 405949753 @exported-using-ghexport Differential Revision: [D111072706](https://our.internmc.facebook.com/intern/diff/D111072706/)
…rough an im2col tiled GEMM (1.1-2.4x) Pull Request resolved: #20873 **Problem:** the direct conv2d kernel runs one thread per output element and re-reads the input receptive field from global memory for every output — zero cross-thread reuse. For the patch-embed stem (3-channel RGB) the vec4-over-IC path is inert (icpg=3 fails the `%4` gate), so it runs the scalar direct path with no reuse at all. **Solution:** route groups==1 non-transposed convs through an implicit-im2col tiled GEMM that reuses the linear tiled-GEMM skeleton — M=OC, N=B*OH*OW, K=IC*KH*KW; shared-memory 32x32 tiles + 4x4 register blocking; the input is im2col-sampled on the fly (out-of-range -> 0.0 implements padding). Grouped/depthwise/transpose stay on the direct/gather kernels. **Before:** every conv -> direct kernel (scalar, or vec4-over-IC when icpg%4==0), no input reuse. **After:** groups==1 -> `conv2d_gemm` (shared-mem tiling + register blocking, input-tile reuse across output positions); grouped/transpose -> unchanged. **Implementation:** - New `conv2d_gemm.wgsl` (+ generated header): forks `linear_fp32_tiled.wgsl` — `read_a` loads the weight `[OC, K]`, `read_b` im2col-samples the input (decodes n->(b,oh,ow), kk->(ic,kh,kw); ih=oh*sH-pH+kh*dH; bounds-check->0), bias per-row (OC), output written NCHW. Reuses the existing `ConvParams` uniform. - `Conv2d.cpp` branches on `groups==1`: GEMM via `compute_tile_grid_2d` + `add_dispatch_2d` (mirrors `LinearFp32.cpp`); else the existing direct dispatch. The grouped path is byte-identical; both grids are computed before any buffer alloc (throw-before-leak). Mirrors Vulkan's own `should_use_conv2d_im2col` groups==1 routing. **Constraints:** scalar GEMM (no vec4) — NCHW's channel stride isn't contiguous, so vec4-over-K would be a strided gather (no compute win on Apple's scalar ALU); ORT skips vec4 for NCHW too. Co-authored-with: Claude Code. ghstack-source-id: 405949751 @exported-using-ghexport Differential Revision: [D110995347](https://our.internmc.facebook.com/intern/diff/D110995347/)
Pull Request resolved: #20874 Splits the im2col-GEMM routing cases out of the `conv2d` im2col-GEMM op diff into their own test diff, stacked directly above it (op below, tests above). Adds the `grouped_vec4` and `gemm_batched` cases to the `conv2d` suite in `test/op_tests/cases.py`, covering `groups==1` im2col-GEMM routing versus the direct vec4 / scalar kernels. Co-authored-with: Claude Code. ghstack-source-id: 405949758 @exported-using-ghexport Differential Revision: [D111072713](https://our.internmc.facebook.com/intern/diff/D111072713/)
…lf RoPE runtime op (unblocks Qwen3) Pull Request resolved: #20875 **Problem:** the WebGPU runtime registers only `et_vk.apply_rotary_emb` (the interleaved/Meta RoPE convention). HuggingFace-derived models (Qwen3, etc.) export the rotate-half convention, which fuses under VulkanPartitioner into `et_vk.apply_rotary_emb_hf` — an op the runtime graph builder has no handler for, so `WebGPUGraph::build()` throws and the delegate is rejected at load with `DelegateInvalidCompatibility` (et_load error 48). The whole model then fails to load on WebGPU. **Solution:** add the `et_vk.apply_rotary_emb_hf` runtime kernel + handler as a rotate-half sibling of the interleaved op. **Before:** only `apply_rotary_emb` (interleaved) is registered; HF-RoPE models throw at load. **After:** both conventions are handled; HF-RoPE models (Qwen3) load and run. **Implementation:** - New `rotary_embedding_hf.wgsl`: one thread per (i, i+half_dim) pair (rotate-half pairing vs the interleaved even/odd), reading a full `[max_seq, rotary_dim]` freqs table indexed at row `start_pos + s`. Scalar, `wg_size` 64 — structural + optimization parity with the interleaved kernel (RoPE is ~1% of runtime; vec4 is neutral for this elementwise-class op on Apple's scalar ALU). - `RotaryEmbedding.cpp`: `apply_rotary_emb_hf_impl` mirrors the interleaved handler; it parses the extra `start_pos` arg as a build-time Int (baked) or a runtime SymInt (dynamic KV-cache decode) exactly as `Sdpa.cpp` handles `input_pos`, and registers a seq resize hook (xq/xk) plus a start_pos resize hook (dynamic decode). Full rotary only (`rotary_dim == head_dim`); partial-rotary passthrough throws (documented follow-up; Qwen3 uses full RoPE). Mirrors Vulkan `et_vk.apply_rotary_emb_hf` (`backends/vulkan/runtime/graph/ops/impl/RotaryEmbedding.cpp`). - Registers `et_vk.apply_rotary_emb_hf.default`. **Constraints:** full rotary only for now; scalar one-thread-per-pair, kept at parity with the interleaved sibling rather than vec4 (neutral for RoPE per the closed vec4 sweep). Co-authored-with: Claude Code. ghstack-source-id: 405949770 @exported-using-ghexport Differential Revision: [D111009173](https://our.internmc.facebook.com/intern/diff/D111009173/)
Pull Request resolved: #20876 Splits the `apply_rotary_emb_hf` op tests into their own diff, stacked directly above the op (op below, tests above). Adds `test/ops/test_rope_hf.py`, the per-op export test for the HuggingFace rotate-half RoPE runtime op. Co-authored-with: Claude Code. ghstack-source-id: 405949767 @exported-using-ghexport Differential Revision: [D111072714](https://our.internmc.facebook.com/intern/diff/D111072714/)
Pull Request resolved: #20986 **Tests for `aten.sub.Tensor` broadcast** Adds op-test coverage for `aten.sub.Tensor`, stacked directly on the sub op diff (op below, tests above). `test/ops/test_sub.py` provides `SubModule` + `CONFIGS` (same-shape, the middle/spatial broadcast `[N,C,H,W] - [N,C,1,1]`, and an alpha != 1 case) plus the export-delegation smoke test; `test/op_tests/cases.py` registers the matching numeric suite (fp64 torch golden on Dawn, mirroring `_mul_suite`), with `alpha` baked into the `.pte` as a construct constant. Co-authored-with: Claude Code. ghstack-source-id: 405949771 @exported-using-ghexport Differential Revision: [D112378930](https://our.internmc.facebook.com/intern/diff/D112378930/)
…ic convert Pull Request resolved: #20987 **Add `aten._to_copy.default` with int↔float numeric convert** **Problem:** The copy-family ops byte-copied across dtypes, so an int32 -> fp32 cast reinterpreted the raw bits — int32 `2` = `0x2` decodes as the fp32 denormal `2.8e-45` — producing wrong values (and div-by-~0 `inf` downstream). Separately, `aten._to_copy.default` was unregistered, so any delegate containing it failed to load. **Solution:** Add `add_to_copy_node`: same-dtype copies stay a flat byte copy, while int<->float copies run a numeric-convert compute shader (`f32(i32)` / `i32(f32)`). Register `aten._to_copy.default`, and route the dim-order copy ops (`dim_order_ops._clone_dim_order.default` / `._to_dim_order_copy.default`) through the same convert-aware path so an int<->float dim-order copy numeric-converts instead of byte-reinterpreting; `view_copy` / `clone` / `alias_copy` stay on the flat copy. Mirrors Vulkan `ToCopy.cpp` (BlitNode vs the view_convert path). **Implementation:** `runtime/ops/to_copy/{ToCopy.cpp,to_copy.h,to_copy_int_to_float.wgsl,to_copy_float_to_int.wgsl}` provide `add_to_copy_node` and register `aten._to_copy.default`; `runtime/ops/view_copy/ViewCopy.cpp` re-points the two dim-order copy ops at `add_to_copy_node`. One `WEBGPU_SRCS` entry. **Constraints:** 32-bit only (int64 constants are downcast to int32 by the Vulkan serializer); fails loud on any other element width. Co-authored-with: Claude Code. ghstack-source-id: 405949773 @exported-using-ghexport Differential Revision: [D112378932](https://our.internmc.facebook.com/intern/diff/D112378932/)
Pull Request resolved: #20988 **Tests for `aten._to_copy.default` int↔float convert** Adds `test/ops/test_to_copy.py`, stacked directly on the to_copy op diff (op below, tests above). Two export-delegation smoke tests (mirroring `test_view_copy.py`): int32 -> fp32 (input int `[1, 2, 3]`, the numeric-convert path) and fp32 -> fp32 (same-dtype flat copy, `copy=True` so the op is not elided). The int -> float value correctness — `[1, 2, 3]` -> `[1.0, 2.0, 3.0]`, NOT the bit-reinterpretation `0x1 -> 1.4e-45` — is checked by the lvp golden. Co-authored-with: Claude Code. ghstack-source-id: 405949769 @exported-using-ghexport Differential Revision: [D112378931](https://our.internmc.facebook.com/intern/diff/D112378931/)
Pull Request resolved: #20989 **Add `aten.leaky_relu.default` to the WebGPU backend** — the SRVGGNetCompact body activation in Real-ESRGAN x4plus super-resolution, so that model can fully delegate to the GPU. **Problem**: The WebGPU delegate had no `leaky_relu` handler, so a model using it could not produce a fully-delegated `.pte`. **Solution**: A scalar-parameter elementwise fp32 kernel computing `x >= 0 ? x : negative_slope * x`, with `negative_slope` carried in the uniform and a 2D-spill dispatch for tensors exceeding the 1D workgroup-count limit. **Implementation**: - `runtime/ops/leaky_relu/{LeakyRelu.cpp,leaky_relu.wgsl,leaky_relu_wgsl.h}` registering `aten.leaky_relu.default`; uses `utils::make_compute_pipeline` + `utils::compute_dispatch_grid`. - Mirrors the Vulkan `leaky_relu.default` delegate (scalar-in-uniform, like `pow.Tensor_Scalar`). - CMake `WEBGPU_SRCS` entry. **Constraints**: fp32-only — throws on non-fp32 or input/output size mismatch (fail-loud, never a silent zero output); no change to existing ops. ghstack-source-id: 405949766 @exported-using-ghexport Differential Revision: [D112417289](https://our.internmc.facebook.com/intern/diff/D112417289/)
Pull Request resolved: #20990 **Op-test suite for `aten.leaky_relu.default`** (stacked on the leaky_relu op diff). Adds the declarative op-test entry: `test/ops/test_leaky_relu.py` (`LeakyReluModule`) + a `@register_op_test("leaky_relu")` suite in `test/op_tests/cases.py`. The framework exports each case via `VulkanPartitioner`, computes the fp64 torch golden, and compares the on-GPU output at `atol=rtol=1e-3`. Cases: `default_slope` (4D `[1,16,8,8]`, slope 0.01) + `slope_0_2` (2D `[3,32]`, slope 0.2). The deterministic input spans negatives so the `negative_slope` branch is exercised. ghstack-source-id: 405949781 @exported-using-ghexport Differential Revision: [D112417280](https://our.internmc.facebook.com/intern/diff/D112417280/)
Pull Request resolved: #20991 **Add `aten.upsample_bilinear2d.vec` to the WebGPU backend** — the bilinear resize on the Depth-Anything / DPT reassemble+fusion head, which upsamples the ViT patch grid back to image resolution. **Problem**: The WebGPU delegate had only nearest-neighbor upsample, so depth-estimation models using bilinear resize could not fully delegate to the GPU. **Solution**: A 4D NCHW fp32 kernel where each output pixel bilinearly interpolates its four source neighbors. `align_corners` selects the source-index formula (matching ATen `area_pixel_compute_source_index`); output H/W come from the output tensor's own dims. **Implementation**: - `runtime/ops/upsample_bilinear2d/{UpsampleBilinear2d.cpp,upsample_bilinear2d.wgsl,upsample_bilinear2d_wgsl.h}` registering `aten.upsample_bilinear2d.vec`; uses `utils::make_compute_pipeline` + `utils::compute_dispatch_grid` + `utils::make_grid_constants`. - Mirrors the Vulkan `upsample_bilinear2d.vec` delegate. - CMake `WEBGPU_SRCS` entry. **Constraints**: fp32-only, 4D in/out with N/C preserved — throws on rank/shape/dtype mismatch (fail-loud, never a silent zero output); no change to existing ops. ghstack-source-id: 405949780 @exported-using-ghexport Differential Revision: [D112417281](https://our.internmc.facebook.com/intern/diff/D112417281/)
Pull Request resolved: #20992 **Op-test suite for `aten.upsample_bilinear2d.vec`** (stacked on the upsample_bilinear2d op diff). Adds `test/ops/test_upsample_bilinear2d.py` (`UpsampleBilinear2dModule`) + a `@register_op_test("upsample_bilinear2d")` suite in `test/op_tests/cases.py` (5 cases). Covers both `align_corners` branches and a non-integer ratio (5->8) that discriminates the two source-index formulas. ghstack-source-id: 405949782 @exported-using-ghexport Differential Revision: [D112417283](https://our.internmc.facebook.com/intern/diff/D112417283/)
Pull Request resolved: #20993 **Add `aten._native_batch_norm_legit_no_training.default` to the WebGPU backend** — inference batch norm on MODNet's decoder and CNN backbones. **Problem**: The WebGPU delegate had no batch-norm handler, so MODNet (background removal) and other CNN models could not fully delegate to the GPU. **Solution**: A 4D NCHW fp32 kernel applying the per-channel inference affine `y = (x - running_mean) / sqrt(running_var + eps) * weight + bias`. `weight`/`bias` are optional (affine=False → unit scale / zero shift), bound via `utils::make_optional_binding` with a dummy buffer when absent. **Implementation**: - `runtime/ops/batch_norm/{BatchNorm.cpp,batch_norm.wgsl,batch_norm_wgsl.h}` registering `aten._native_batch_norm_legit_no_training.default`; uses `utils::make_compute_pipeline` + `utils::make_optional_binding`. - Multi-output op: reads the `out` entry of the output ValueList (`save_mean`/`save_invstd` unused in inference). - Mirrors the Vulkan `_native_batch_norm_legit_no_training` delegate. - CMake `WEBGPU_SRCS` entry. **Constraints**: fp32-only, 4D in/out — throws on rank/shape/dtype mismatch (fail-loud); no change to existing ops. ghstack-source-id: 405949779 @exported-using-ghexport Differential Revision: [D112417288](https://our.internmc.facebook.com/intern/diff/D112417288/)
Pull Request resolved: #20994 **Op-test suite for `aten._native_batch_norm_legit_no_training.default`** (stacked on the batch_norm op diff). Adds `test/ops/test_batch_norm.py` (`BatchNorm2dModule` — `nn.BatchNorm2d.eval()` with deterministic running stats + affine) + a `@register_op_test("batch_norm")` suite in `test/op_tests/cases.py` (3 cases). Covers affine + non-affine (optional weight/bias) and an odd H*W; only the populated `out` ValueList entry is compared (out_index 0). ghstack-source-id: 405949784 @exported-using-ghexport Differential Revision: [D112417282](https://our.internmc.facebook.com/intern/diff/D112417282/)
Pull Request resolved: #20995 **Add `aten.split_with_sizes_copy.default` to the WebGPU backend** — the split on YOLO's Detect head (separating the concatenated box / objectness / class predictions). **Problem**: The WebGPU delegate had no `split_with_sizes_copy`, so YOLO object-detection could not fully delegate to the GPU. **Solution**: Split `self` along `dim` into N contiguous chunks. Each chunk is a step-1 slice from the running offset, reusing the `slice.wgsl` gather kernel — one dispatch per output. Outputs arrive as a serialized ValueList. Each chunk writes its own distinct output buffer and reads only the shared input, so there is no cross-dispatch read-after-write hazard. **Implementation**: - `runtime/ops/split_with_sizes/SplitWithSizes.cpp` registering `aten.split_with_sizes_copy.default`; reuses `slice_wgsl.h` (no new shader). Uses `utils::make_compute_pipeline` (auto-derived bind-group layout) rather than hand-rolling the layout / pipeline / bind group. - Mirrors the Vulkan `split_with_sizes_copy` delegate. - CMake `WEBGPU_SRCS` entry. **Constraints**: fp32-only; `dim` normalized + range-checked; `outputs == sizes` count enforced (fail-loud). Reuses the `slice` op's `slice_wgsl.h`, so this diff stacks above `slice` and must land after it. No change to existing ops. ghstack-source-id: 405949786 @exported-using-ghexport Differential Revision: [D112417284](https://our.internmc.facebook.com/intern/diff/D112417284/)
Pull Request resolved: #20996 **Op-test suite for `aten.split_with_sizes_copy.default`** (stacked on the split_with_sizes_copy op diff). Adds `test/ops/test_split_with_sizes_copy.py` (`SplitWithSizesModule` — `torch.split` by a size list) + a `@register_op_test("split_with_sizes_copy")` suite in `test/op_tests/cases.py` (3 cases: a 3-way channel split, a dim-0 split, and a last-dim split). Multi-output: the framework compares chunk 0 (out_index 0) while each case exercises all N per-chunk dispatches. `copy` is bit-exact, so the golden is float32. ghstack-source-id: 405949787 @exported-using-ghexport Differential Revision: [D112417279](https://our.internmc.facebook.com/intern/diff/D112417279/)
…ipeline Pull Request resolved: #20868 Route `add`, `mul`, `index`, `permute`, `select`, `slice`, `update_cache`, `rms_norm`, and `embedding_q4gsw` through the shared `utils::make_compute_pipeline` helper (which uses `layout:"auto"`), replacing each op's hand-written `WGPUBindGroupLayoutEntry[]` + pipeline-layout + bind-group boilerplate (~50-110 lines each) with a single helper call. The driver now derives the bind-group layout from the shader's statically-used bindings. Byte-behavior is preserved: identical binding indices/types/buffers/sizes, dispatch workgroup counts, resize hooks, and validations; override constants (`wg_size`) passed via the helper's `constants` param. Extends the Diff 1 layout:"auto" adoption to the trivial single-dispatch ops. ghstack-source-id: 405949795 @exported-using-ghexport Differential Revision: [D110836665](https://our.internmc.facebook.com/intern/diff/D110836665/)
JCNTH
self-requested a review
July 23, 2026 16:13
JCNTH
approved these changes
Jul 23, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR was created by the merge bot to help merge the original PR into the main branch.
ghstack PR number: #20842 by @JCNTH
^ Please use this as the source of truth for the PR details, comments, and reviews
ghstack PR base: https://github.com/pytorch/executorch/tree/gh/JCNTH/17/base
ghstack PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/17/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/main
Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/17/orig
@diff-train-skip-merge