[ExecuTorch][WebGPU] Add relu op (shared unary handler; sigmoid adopts make_compute_pipeline)#20863
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20863
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: ❌ 2 New FailuresAs of commit 1790b25 with merge base 430b73d ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
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)
…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)
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)
… 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)
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)
…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)
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)
…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)
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)
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)
…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)
…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)
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)
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)
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)
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)
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)
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)
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)
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)
…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)
…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/)
…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/)
…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/)
…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/)
…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/)
Stack from ghstack (oldest at bottom):
Adds
reluto 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.defaultkernel, andsigmoid(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
sigmoidwas implemented with a bespoke inline pipeline; there was norelu.add_unary_ophelper (inruntime/ops/sigmoid/UnaryOp.cpp) builds the input/output/params binding and dispatch for any elementwise-unary WGSL;sigmoid_impland the newrelu_implare thin wrappers over it, sosigmoidnow goes through the sameutils::make_compute_pipelinepath asrelu.relu.wgslis a one-element-per-threadoutput[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_countfor the 1D dispatch, thewg_sizeoverride constant, the uniform (num_elements) viautils::make_uniform, and the three-binding pipeline viautils::make_compute_pipeline.graph.add_tensor_resize_hookrecomputesnum_elements, rewrites the uniform viawgpuQueueWriteBuffer, and updates the dispatch's workgroup count for the live shape; the graph owns the uniform buffer so the hook can rewrite it.aten.sigmoid.default -> sigmoid_implandaten.relu.default -> relu_impl.backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp(add_unary_op_node); Vulkan expressesreluasclamp(0, inf), whereas this kernel uses a directmax(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
Differential Revision: D110836664