Skip to content

[ExecuTorch][WebGPU] Add optimized conv2d op#20849

Merged
JCNTH merged 6 commits into
gh/JCNTH/24/basefrom
gh/JCNTH/24/head
Jul 23, 2026
Merged

[ExecuTorch][WebGPU] Add optimized conv2d op#20849
JCNTH merged 6 commits into
gh/JCNTH/24/basefrom
gh/JCNTH/24/head

Conversation

@JCNTH

@JCNTH JCNTH commented Jul 10, 2026

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

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.

Problematen.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.
@exported-using-ghexport

Differential Revision: D110836669

Differential Revision: D110836669

[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Jul 10, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

❌ 68 New Failures

As of commit a395a9a with merge base 430b73d (image):

NEW FAILURES - The following jobs have failed:

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

@github-actions

Copy link
Copy Markdown

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.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 10, 2026
This was referenced Jul 10, 2026
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* __->__ #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D110836672](https://our.internmc.facebook.com/intern/diff/D110836672)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* __->__ #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



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.
@exported-using-ghexport

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

Differential Revision:
[D111072712](https://our.internmc.facebook.com/intern/diff/D111072712)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* __->__ #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D110836668](https://our.internmc.facebook.com/intern/diff/D110836668)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* __->__ #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



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.
@exported-using-ghexport

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

Differential Revision:
[D111072707](https://our.internmc.facebook.com/intern/diff/D111072707)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* __->__ #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D110836680](https://our.internmc.facebook.com/intern/diff/D110836680)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* __->__ #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



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.
@exported-using-ghexport

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

Differential Revision:
[D111072724](https://our.internmc.facebook.com/intern/diff/D111072724)
JCNTH added a commit that referenced this pull request Jul 23, 2026
…s make_compute_pipeline) (#20863)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* __->__ #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D110836664](https://our.internmc.facebook.com/intern/diff/D110836664)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* __->__ #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



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.
@exported-using-ghexport

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

Differential Revision:
[D111072725](https://our.internmc.facebook.com/intern/diff/D111072725)
JCNTH added a commit that referenced this pull request Jul 23, 2026
…y/alias glue (#20865)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* __->__ #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D110836670](https://our.internmc.facebook.com/intern/diff/D110836670)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* __->__ #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



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.
@exported-using-ghexport

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

Differential Revision:
[D111072708](https://our.internmc.facebook.com/intern/diff/D111072708)
JCNTH added a commit that referenced this pull request Jul 23, 2026
… for channel attention (15-30x faster) (#20871)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* __->__ #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D110994975](https://our.internmc.facebook.com/intern/diff/D110994975)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* __->__ #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



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.
@exported-using-ghexport

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

Differential Revision:
[D111072706](https://our.internmc.facebook.com/intern/diff/D111072706)
JCNTH added a commit that referenced this pull request Jul 23, 2026
…rough an im2col tiled GEMM (1.1-2.4x) (#20873)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* __->__ #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D110995347](https://our.internmc.facebook.com/intern/diff/D110995347)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* __->__ #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



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.
@exported-using-ghexport

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

Differential Revision:
[D111072713](https://our.internmc.facebook.com/intern/diff/D111072713)
JCNTH added a commit that referenced this pull request Jul 23, 2026
…lf RoPE runtime op (unblocks Qwen3) (#20875)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* __->__ #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D111009173](https://our.internmc.facebook.com/intern/diff/D111009173)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* __->__ #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



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.
@exported-using-ghexport

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

Differential Revision:
[D111072714](https://our.internmc.facebook.com/intern/diff/D111072714)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* __->__ #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D112378930](https://our.internmc.facebook.com/intern/diff/D112378930)
JCNTH added a commit that referenced this pull request Jul 23, 2026
…ic convert (#20987)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* __->__ #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D112378932](https://our.internmc.facebook.com/intern/diff/D112378932)
JCNTH added a commit that referenced this pull request Jul 23, 2026
…ert (#20988)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* __->__ #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D112378931](https://our.internmc.facebook.com/intern/diff/D112378931)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* __->__ #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D112417289](https://our.internmc.facebook.com/intern/diff/D112417289)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* __->__ #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D112417280](https://our.internmc.facebook.com/intern/diff/D112417280)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* __->__ #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D112417281](https://our.internmc.facebook.com/intern/diff/D112417281)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* #20993
* __->__ #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D112417283](https://our.internmc.facebook.com/intern/diff/D112417283)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* #20994
* __->__ #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D112417288](https://our.internmc.facebook.com/intern/diff/D112417288)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* #20995
* __->__ #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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).
@exported-using-ghexport

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

Differential Revision:
[D112417282](https://our.internmc.facebook.com/intern/diff/D112417282)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* #20996
* __->__ #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D112417284](https://our.internmc.facebook.com/intern/diff/D112417284)
JCNTH added a commit that referenced this pull request Jul 23, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* #20868
* __->__ #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



**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.
@exported-using-ghexport

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

Differential Revision:
[D112417279](https://our.internmc.facebook.com/intern/diff/D112417279)
JCNTH added a commit that referenced this pull request Jul 23, 2026
…ipeline (#20868)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* __->__ #20868
* #20996
* #20995
* #20994
* #20993
* #20992
* #20991
* #20990
* #20989
* #20988
* #20987
* #20986
* #20876
* #20875
* #20874
* #20873
* #20872
* #20871
* #20866
* #20865
* #20864
* #20863
* #20862
* #20861
* #20860
* #20859
* #20858
* #20857
* #20856
* #20855
* #20854
* #20852
* #20851
* #20850
* #20849
* #20848
* #20846
* #20845
* #20844
* #20843
* #20842



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.
@exported-using-ghexport

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

Differential Revision:
[D110836665](https://our.internmc.facebook.com/intern/diff/D110836665)
JCNTH added a commit that referenced this pull request Jul 23, 2026
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/)
JCNTH added a commit that referenced this pull request Jul 23, 2026
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/)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants