[ExecuTorch][WebGPU] Port 10 unary activations to the WebGPU backend#21298
Merged
Conversation
Pull Request resolved: #21154 Ports the Vulkan delegate's no-param unary activations — `abs`, `exp`, `sqrt`, `rsqrt`, `sin`, `cos`, `tanh`, `round`, `neg`, `hardswish` — built on the shared `add_unary_op` helper. Each is a per-op `.wgsl` (2D-fold dispatch, `wg_size=256`) whose body mirrors the Vulkan formula in `backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml` + `activations.h` verbatim, registered through `runtime/ops/unary/Activations.cpp`. Key changes: - `runtime/ops/unary/{abs,exp,sqrt,rsqrt,sin,cos,tanh,round,neg,hardswish}.wgsl` (+ generated `_wgsl.h`) — `output[idx] = <vulkan-math>`; notably `rsqrt` = `1.0 / sqrt(x)`, `tanh` = `tanh(clamp(x, -15.0, 15.0))`, `hardswish` = the piecewise `select`. - `runtime/ops/unary/Activations.cpp` — registers `aten.<op>.default` for the 10 ops, each calling the shared `add_unary_op`. - `CMakeLists.txt` — `runtime/ops/unary/Activations.cpp` in `WEBGPU_SRCS`. `tan` and `hardsigmoid` are intentionally excluded: the Vulkan partitioner (`backends/vulkan/op_registry.py`) does not tag them, so they cannot be delegated to this backend without a separate partitioner change. ghstack-source-id: 406360832 @exported-using-ghexport Differential Revision: [D112257637](https://our.internmc.facebook.com/intern/diff/D112257637/)
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21298
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
Pull Request resolved: #21155 Adds the declarative op-test suites for the unary activations ported in the diff below. Each op gets a `torch` reference plus a deterministic, range-bounded input so the fp64 golden is well-defined — positive for `sqrt`/`rsqrt`, spanning the ±3 knees for `hardswish`, and reaching the ±15 clamp for `tanh` — driving the manifest-based Dawn golden. Key changes: - `test/ops/test_unary_activations.py` — `UnaryModule` + `UNARY_G1` (op name -> torch fn + linspace input gen). - `test/op_tests/cases.py` — registers the 10 suites (a `mat` + a `rank3` case each) via a factory loop, mirroring the existing per-op registration convention. ghstack-source-id: 406360831 @exported-using-ghexport Differential Revision: [D112257657](https://our.internmc.facebook.com/intern/diff/D112257657/)
Pull Request resolved: #21156 Ports the Vulkan delegate's clamp-family activations. `clamp` and `hardtanh` share one `clamp(x, min, max)` shader (mirroring how Vulkan reuses its clamp shader for both), built on the shared unary helper. Bounds come from the op's scalar args, with `None -> +/-inf` matching Vulkan `get_val_or_inf` (`backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp:72`). Key changes: - `runtime/ops/unary/clamp.wgsl` (+ generated `_wgsl.h`) — `output[idx] = clamp(input[idx], params.minimum, params.maximum)`; the `minimum`/`maximum` member names mirror Vulkan `unary_op.glsl`. - `runtime/ops/unary/Activations.cpp` — `clamp_impl` / `hardtanh_impl` (args `[in, min, max, out]`) + `get_val_or_inf` (reads an Int/Double scalar, substitutes `+/-inf` for `None`, throws on any other type). - `runtime/ops/unary/UnaryOp.cpp` — the shared helper now rejects int operands (fail-loud): the fp32-only backend would otherwise reinterpret int bits as f32, and `clamp` is the one unary the Vulkan partitioner (`op_registry.py`) tags for integer tensors. `relu6` is excluded: the Vulkan partitioner does not tag it, so it cannot be delegated without a separate partitioner change. ghstack-source-id: 406360830 @exported-using-ghexport Differential Revision: [D112257672](https://our.internmc.facebook.com/intern/diff/D112257672/)
Pull Request resolved: #21157 Adds the op-test suites for the clamp-family ops ported in the diff below. `clamp` covers both-bounds and a `min=None` case (exercising the `None -> -inf` substitution); `hardtanh` covers the default `(-1, 1)` bounds and a wider `(-2, 2)`. Key changes: - `test/ops/test_unary_activations.py` — `ClampModule`, `HardtanhModule`, `CLAMP_CONFIGS`, `HARDTANH_CONFIGS`. - `test/op_tests/cases.py` — `_clamp_suite` + `_hardtanh_suite`. ghstack-source-id: 406360836 @exported-using-ghexport Differential Revision: [D112257631](https://our.internmc.facebook.com/intern/diff/D112257631/)
Pull Request resolved: #21158 Ports the Vulkan delegate's `minimum` with NumPy broadcasting: `out = min(in1, in2)`, mirroring the sibling `mul` op + Vulkan `binary_op_buffer` (3 per-tensor `TensorMeta` uniforms; size-1 input dims clamped and the flat output index relinearized; equal-shape fast path). Key changes: - `runtime/ops/minimum/{BinaryOp.cpp, binary_minimum.wgsl}` (+ generated `_wgsl.h`) — `minimum_impl` (args `[in1, in2, out]`); broadcast `TensorMeta` path + dual-operand resize hook (mirrors `mul`); fail-loud int guard (the fp32-only backend would otherwise reinterpret int bits as f32). - `CMakeLists.txt` — `runtime/ops/minimum/BinaryOp.cpp` in `WEBGPU_SRCS`. Note: WGSL `min` and `torch.minimum` differ on NaN, matching the Vulkan reference; the golden uses finite inputs. ghstack-source-id: 406360835 @exported-using-ghexport Differential Revision: [D112257607](https://our.internmc.facebook.com/intern/diff/D112257607/)
Pull Request resolved: #21159 Adds the minimum op-test suite (2d + 3d same-shape) for the op ported in the diff below. Key changes: - `test/ops/test_minimum.py` — `MinimumModule` (`torch.minimum`). - `test/op_tests/cases.py` — `_minimum_suite`. ghstack-source-id: 406360846 @exported-using-ghexport Differential Revision: [D112257590](https://our.internmc.facebook.com/intern/diff/D112257590/)
Pull Request resolved: #21160 Ports `pow.Tensor_Tensor` with NumPy broadcasting, mirroring the sibling `mul` op + Vulkan `binary_op_buffer` (per-tensor `TensorMeta` uniforms; size-1 input dims clamped and the flat output index relinearized; equal-shape fast path). Key changes: - `runtime/ops/pow/{BinaryOp.cpp, binary_pow.wgsl}` (+ generated `_wgsl.h`) — `pow_impl` (args `[in1, in2, out]`); broadcast `TensorMeta` path + dual-operand resize hook; fail-loud int guard (fp32-only backend). Note: WGSL `pow(x, y)` is undefined for `x < 0`, exact parity with the Vulkan GLSL `pow`. Co-authored-with: Claude Code. ghstack-source-id: 406360857 @exported-using-ghexport Differential Revision: [D112257673](https://our.internmc.facebook.com/intern/diff/D112257673/)
Pull Request resolved: #21161 Adds the op-test suite for `pow.Tensor_Tensor` (ported in the diff below). Uses a strictly-positive base so `pow(neg, frac)=NaN` is never exercised. Key changes: - `test/ops/test_pow.py` + `test/op_tests/cases.py` — `PowModule` + `_pow_suite` (2d + 3d equal-shape, `atol=rtol=1e-3` vs the `torch` golden). Co-authored-with: Claude Code. ghstack-source-id: 406360852 @exported-using-ghexport Differential Revision: [D112257639](https://our.internmc.facebook.com/intern/diff/D112257639/)
Pull Request resolved: #21162 Ports `pow.Tensor_Scalar` (tensor ^ scalar-exponent) on the shared unary helper: the exponent rides the `min` uniform slot, exactly like `leaky_relu`'s slope. `output[idx] = pow(input[idx], exponent)`. Key changes: - `runtime/ops/unary/pow_scalar.wgsl` (+ generated `_wgsl.h`) — `pow(input[idx], params.minimum)`. - `runtime/ops/unary/Activations.cpp` — `pow_scalar_impl` (args `[in, exponent, out]`; exponent read via the shared `get_val_or_inf` into the min slot), registered `aten.pow.Tensor_Scalar`. Inherits the shared helper's fp32 int-guard + resize hook (the exponent is re-applied on resize). Note: WGSL `pow` returns NaN for a negative base at any exponent (matching the reference GPU `pow`); the suite uses a positive base. ghstack-source-id: 406360862 @exported-using-ghexport Differential Revision: [D112257604](https://our.internmc.facebook.com/intern/diff/D112257604/)
Pull Request resolved: #21163 Adds the pow_scalar op-test suite (square + sqrt exponents, positive base) for the op ported in the diff below. Key changes: - `test/ops/test_unary_activations.py` — `PowScalarModule` + `POW_SCALAR_CONFIGS`. - `test/op_tests/cases.py` — `_pow_scalar_suite`. ghstack-source-id: 406360849 @exported-using-ghexport Differential Revision: [D112257612](https://our.internmc.facebook.com/intern/diff/D112257612/)
Pull Request resolved: #21164 Ports the Vulkan delegate's `amax`/`amin` as last-dim per-row reductions — the buffer-storage path Vulkan itself restricts reductions to (`impl/Reduce.cpp` `add_reduce_per_row_node`; its texture path throws on buffer storage). One thread per output row serially reduces over the contiguous last dim. Key changes: - `runtime/ops/amax/{Reduce.cpp, amax.wgsl}` + `runtime/ops/amin/{Reduce.cpp, amin.wgsl}` (+ generated `_wgsl.h`) — `out[row] = reduce(input[row*W .. +W])` where `W` = last-dim extent; args `[in, dim, keepdim, out]`; fp32 int-guard; keepdim-aware resize hook; registered `aten.amax.default` / `aten.amin.default`. - `CMakeLists.txt` — both in `WEBGPU_SRCS`. Constraint: only a single last-dim reduction (`dim == -1` or `ndim-1`) is supported on buffer storage; any other reduce dim throws (mirrors Vulkan's buffer-per-row gate). The accumulator seeds `input[base]` (not 0), so all-negative rows reduce correctly. ghstack-source-id: 406360850 @exported-using-ghexport Differential Revision: [D112257619](https://our.internmc.facebook.com/intern/diff/D112257619/)
Pull Request resolved: #21165 Adds op-test suites for the last-dim reductions ported in the diff below (keepdim + nodim, 2d + 3d). Key changes: - `test/ops/test_reduce.py` — `AmaxModule` / `AminModule` (reduce `dim=-1`). - `test/op_tests/cases.py` — `_reduce_suite` + `_amax_suite` / `_amin_suite`. ghstack-source-id: 406360856 @exported-using-ghexport Differential Revision: [D112257621](https://our.internmc.facebook.com/intern/diff/D112257621/)
Pull Request resolved: #21166 Ports `aten.flip.default` as a per-element index-transform, mirroring the landed `permute` op (same dual-`TensorMeta` + 5-binding bind group). flip is pure data movement, so it reverses the coord along each flagged dim and copies — no fp compute. Key changes: - `runtime/ops/flip/{Flip.cpp, flip.wgsl}` (+ generated `flip_wgsl.h`) — per out flat idx: unravel via `out_meta.strides`, `in_coord = select(coord, sizes[d]-1-coord, flip[d]==1)`, ravel via `in_meta.strides`, `output[out] = input[in]`; `Params { flip: vec4<u32> }` bitmap built from the normalized `dims` arg; registered `aten.flip.default`. - `CMakeLists.txt` — `runtime/ops/flip/Flip.cpp` in `WEBGPU_SRCS`. Matches Vulkan `impl/Flip.cpp` (`create_whcn_bitmap` normalizes/flags dims; reversal `size-1-coord`). No int dtype-guard (unlike clamp/minimum which do fp arithmetic): flip is bit-preserving movement, so any 4-byte dtype is correct; the `nbytes != numel*4` guard rejects int64/fp16. 2D-folded dispatch lifts the 65535 cap. `out` is the last arg (this backend's convention). ghstack-source-id: 406360854 @exported-using-ghexport Differential Revision: [D112257643](https://our.internmc.facebook.com/intern/diff/D112257643/)
Pull Request resolved: #21167 Adds the op-test suite for the `flip` op ported in the diff below. Key changes: - `test/ops/test_flip.py` — `FlipModule` (reverses `dims`) + `FlipTest` export-delegation smoke test (mirrors `PermuteTest`). - `test/op_tests/cases.py` — `_flip_suite`: last `[-1]`/2d, dim0 `[0]`/2d, both_3d `[0,2]`/3d, mid_3d `[1]`/3d, multi_4d `[1,3]`/4d. ghstack-source-id: 406360858 @exported-using-ghexport Differential Revision: [D112257622](https://our.internmc.facebook.com/intern/diff/D112257622/)
Pull Request resolved: #21168 Ports `aten.repeat.default` as a per-element tiling gather, mirroring the landed `permute`/`flip` index-transform (dual-`TensorMeta`, 2D-fold dispatch). repeat is pure data movement, so it maps each output element to its source via a per-dim modulo and copies — no fp compute. Key changes: - `runtime/ops/repeat/{Repeat.cpp, repeat.wgsl}` (+ generated `repeat_wgsl.h`) — per out flat idx: unravel via `out_meta.strides`, `in_coord = out_coord % in_meta.sizes[in_d]`, ravel via `in_meta.strides`, `output[out] = input[in]`; registered `aten.repeat.default`. - `CMakeLists.txt` — `runtime/ops/repeat/Repeat.cpp` in `WEBGPU_SRCS`. Mirrors Vulkan `glsl/repeat_buffer.glsl` (`in_tidx[d] = out_tidx[d] % size_at(in_meta, d)`). Vulkan's innermost-first `BufferMetadata` auto-aligns a shorter input; WebGPU's NCHW `TensorMeta` makes the right-alignment explicit via `offset = out_ndim - in_ndim` (leading `d < offset` dims are pure prepends, contributing 0). No `repeats`-arg read needed: the repeated out dims are baked in by export, so the tiling is recoverable from shapes alone. No int dtype-guard (bit-preserving movement, like `flip`); the `nbytes != numel*4` guard rejects int64/fp16. ghstack-source-id: 406360860 @exported-using-ghexport Differential Revision: [D112257679](https://our.internmc.facebook.com/intern/diff/D112257679/)
Pull Request resolved: #21169 Adds the op-test suite for the `repeat` op ported in the diff below. Key changes: - `test/ops/test_repeat.py` — `RepeatModule` (`x.repeat(*repeats)`) + `RepeatTest` export-delegation smoke test (mirrors `PermuteTest`/`FlipTest`). - `test/op_tests/cases.py` — `_repeat_suite`: tile_1d, tile_2d, prepend_3d `[1,3,2]`, prepend_ext `[2,3,1]`, tile_3d `[2,1,2]`. ghstack-source-id: 406360847 @exported-using-ghexport Differential Revision: [D112257616](https://our.internmc.facebook.com/intern/diff/D112257616/)
Pull Request resolved: #21170 Extends the op-test generator to golden-verify EVERY output of a multi-output op (e.g. `split_with_sizes_copy` → N tensors), not just `output[0]`. The C++ driver already selected `outs[output_index]` per manifest entry; only the generator hardcoded a single `output_index = 0`. Key changes: - `test/op_tests/generate_op_tests.py` — `generate_case` now returns one manifest entry PER output tensor: it iterates the module's tuple return, writes a golden per output (`{case}_out{i}.golden.bin`, `output_index = i`), and runs the dual-oracle gate per output. `generate()` uses `extend`. Single-output ops emit exactly one entry, byte-identical to before (no `_out` suffix, `output_index = 0`). - `test/op_tests/test_generator.py` — `test_generate_case_writes_artifacts` reads the returned list (asserts one entry for the single-output `add`). ghstack-source-id: 406360853 @exported-using-ghexport Differential Revision: [D112257608](https://our.internmc.facebook.com/intern/diff/D112257608/)
Pull Request resolved: #21171 Ports `aten.index_select.default` as a gather along a dim, mirroring the landed `flip`/`split` index-transform (dual-`TensorMeta`, 2D-fold dispatch) + the landed `index/Index.cpp` int-index binding. index_select is pure data movement, so it remaps the dim coord through an int index and copies — no fp compute. Key changes: - `runtime/ops/index_select/{IndexSelect.cpp, index_select.wgsl}` (+ generated `index_select_wgsl.h`) — per out flat idx: unravel via `out_meta.strides`, `in_coord = out_coord` with the dim coord remapped by `index[out_coord[dim]]`, ravel via `in_meta.strides`, `output[out] = input[in]`; the index tensor is bound as `array<i32>` (the int32 `downcast_64_bit` of the int64 index, same as `Index.cpp`); `Params { info: vec4<u32> }` carries `dim`. Registered `aten.index_select.default`. - `CMakeLists.txt` — `runtime/ops/index_select/IndexSelect.cpp` in `WEBGPU_SRCS`. Buffer re-derivation of the gather (Vulkan `IndexSelect.cpp` is texture/channel-packed only — `op_registry.py` tags it `CHANNELS_PACKED_TEXTURE`; this backend is buffer-only, so the NCHW index math is derived and CPU-de-risked, cf. `split.wgsl`). Index-value bounds are trusted/unclamped, matching Vulkan `index_select.glsl`; a negative index is a memory-safe WGSL OOB read (no trap), same safety class as Vulkan's clamped `texelFetch`. No int dtype-guard on self/out (bit-preserving); `nbytes` guards reject non-fp32 self/out and a non-int32 index. ghstack-source-id: 406360861 @exported-using-ghexport Differential Revision: [D112257660](https://our.internmc.facebook.com/intern/diff/D112257660/)
Pull Request resolved: #21172 Adds the op-test suite for the `index_select` op ported in the diff below. Key changes: - `test/ops/test_index_select.py` — `IndexSelectModule` gathers via a baked int index buffer, so the only runtime input is the float tensor (the op-test framework is float-only); `torch.index_select(x, dim, self.index)`. + `IndexSelectTest` delegation smoke test (mirrors `SplitTest`). - `test/op_tests/cases.py` — `_index_select_suite`: dim0_1d, dim0_2d, dim1_2d (repeated index), dim1_3d (middle dim), dim2_3d. ghstack-source-id: 406360855 @exported-using-ghexport Differential Revision: [D112257648](https://our.internmc.facebook.com/intern/diff/D112257648/)
Pull Request resolved: #21173 Ports `aten.native_group_norm.default` as a 2-pass, 3-output group normalization, the second Phase B (vision) op. Pass 1 reduces per group to `mean`/`rstd`; pass 2 applies the per-channel affine. Key changes: - `runtime/ops/native_group_norm/{GroupNorm.cpp, group_norm_reduce.wgsl, group_norm.wgsl}` (+ generated `_wgsl.h`) — reduce: one thread per `(n, group)` serially scans that group's contiguous NCHW block (`base = n*C*HxW + g*(C/G)*HxW`, `group_size = (C/G)*HxW` elements), writing `mean = s/count`, `rstd = inverseSqrt(ss/count - mean^2 + eps)`; normalize: one thread per element, `g = c / (C/G)`, `out = (x - mean[n*G+g]) * rstd[n*G+g] * weight[c] + bias[c]`. A file-local `add_gn_dispatch` helper builds each of the two dispatches from an explicit binding list sharing one params UBO. Registered `aten.native_group_norm.default`. - `CMakeLists.txt` — `runtime/ops/native_group_norm/GroupNorm.cpp` in `WEBGPU_SRCS`. Mirrors Vulkan `impl/GroupNorm.cpp` (2 nodes: reduce -> mean/rstd, then normalize) + `glsl/group_norm_reduce_texture.glsl` (`mean=sum/count`, `variance=sumsq/count - mean^2`, `rstd=1/sqrt(var+eps)`). Buffer NCHW re-derivation (Vulkan is texture/channels-packed), so a serial per-group reduce (a group's channels x HxW are contiguous in NCHW) replaces the cooperative texture reduction. All 3 outputs are real. The inter-pass RAW ordering (normalize reads the reduce's mean/rstd) is guaranteed because `execute()` runs one compute pass per dispatch (`WebGPUGraph.cpp`; proven by `test_dispatch_order.cpp` to 1M elems). A dynamic-shape resize hook recomputes params + both dispatch counts + propagates cur_dims for out/mean/rstd. Fail-loud guards: 4D, `C % group == 0`, fp32, weight/bias len `== C`, mean/rstd size `== N*group`. ghstack-source-id: 406360859 @exported-using-ghexport Differential Revision: [D112257596](https://our.internmc.facebook.com/intern/diff/D112257596/)
Pull Request resolved: #21174 Adds the op-test suite for the `native_group_norm` op ported in the diff below, using the multi-output golden support to verify all 3 outputs. Key changes: - `test/ops/test_group_norm.py` — `GroupNormModule` returns the full `torch.native_group_norm` (out, mean, rstd) tuple, so the multi-output golden path checks both the reduce (mean/rstd) and normalize (out) passes. + `GroupNormTest` delegation smoke test. - `test/op_tests/cases.py` — `_group_norm_suite`: c4_g2, c6_g3, c4_g1. ghstack-source-id: 406360848 @exported-using-ghexport Differential Revision: [D112257626](https://our.internmc.facebook.com/intern/diff/D112257626/)
Pull Request resolved: #21175 Ports `aten.avg_pool2d.default` as a windowed spatial average, a Phase B (vision) op. Each output cell averages its KxK input window. Key changes: - `runtime/ops/avg_pool2d/{AvgPool2d.cpp, avg_pool2d.wgsl}` (+ generated `avg_pool2d_wgsl.h`) — per NCHW output element: `ipos = out*stride - padding`, window `[max(0,ipos), min(ipos+kernel, in))`, `acc = sum(window)`, `out = acc / divisor`. The divisor mirrors the three ATen cases: `divisor_override` if given, else `count_include_pad` counts the padded window minus any overhang past the padded edge, else the actual number of summed cells. Registered `aten.avg_pool2d.default`. - `CMakeLists.txt` — `runtime/ops/avg_pool2d/AvgPool2d.cpp` in `WEBGPU_SRCS`. Mirrors Vulkan `glsl/avg_pool2d.glsl` (`ipos`/`start`/`end`/`sum`/`empty`-overhang divisor) + `impl/Pool.cpp` (arg order, `divisor_override` read only when Int, `count_include_pad` bool). Buffer NCHW re-derivation (Vulkan is texture/channels-packed). `read_pair` normalizes each IntList param (empty stride -> kernel, len-1 broadcast, len-2). A dynamic-shape resize hook recomputes the pooled output extent via `pool_out_dim` (floor, or ceil + last-window-drop for `ceil_mode`) + params + dispatch + cur_dims. Fail-loud guards: 4D, fp32. ghstack-source-id: 406360851 @exported-using-ghexport Differential Revision: [D112257592](https://our.internmc.facebook.com/intern/diff/D112257592/)
Pull Request resolved: #21176 Adds the op-test suite for the `avg_pool2d` op ported in the diff below. Key changes: - `test/ops/test_avg_pool2d.py` — `AvgPool2dModule` (`F.avg_pool2d`, parametrized by kernel/stride/padding/count_include_pad/ceil_mode/divisor) + `AvgPool2dTest` delegation smoke test. - `test/op_tests/cases.py` — `_avg_pool2d_suite`: basic, pad_cip, pad_nocip, asym, divisor, ceil_cip, ceil_nocip. ghstack-source-id: 406366786 @exported-using-ghexport Differential Revision: [D113493800](https://our.internmc.facebook.com/intern/diff/D113493800/)
Pull Request resolved: #21177 Ports `aten.pixel_shuffle.default` as a channel->space rearrange, a Phase B (vision) op. Pure data movement: each output element gathers from one input element. Key changes: - `runtime/ops/pixel_shuffle/{PixelShuffle.cpp, pixel_shuffle.wgsl}` (+ generated `_wgsl.h`) — for output `(n, c_out, h_out, w_out)`: `h_in = h_out / r`, `w_in = w_out / r`, `c_in = c_out*r*r + (h_out%r)*r + (w_out%r)`, then `output[out] = input[in]`. Leading dims collapse into a batch, so any rank >= 3 works. Registered `aten.pixel_shuffle.default`. - `CMakeLists.txt` — `runtime/ops/pixel_shuffle/PixelShuffle.cpp` in `WEBGPU_SRCS`. Mirrors Vulkan `glsl/pixel_shuffle_buffer.glsl` (same `c_in`/`h_in`/`w_in` map) + `impl/PixelShuffle.cpp` (arg order `[in, upscale_factor, out]`, `r >= 1`, `in_channels % (r*r) == 0`, `resize_pixel_shuffle_node`). No int dtype-guard (bit-preserving movement); `nbytes % 4` guard. A dynamic-shape resize hook recomputes the out shape (`in` with the last 3 dims rescaled by r) + params + dispatch + cur_dims. Fail-loud guards: rank >= 3, r >= 1, `in_channels % (r*r) == 0` (mirrors Vulkan VK_CHECK_COND), fp32. ghstack-source-id: 406366788 @exported-using-ghexport Differential Revision: [D112257663](https://our.internmc.facebook.com/intern/diff/D112257663/)
Pull Request resolved: #21178 Adds the op-test suite for the `pixel_shuffle` op ported in the diff below. Key changes: - `test/ops/test_pixel_shuffle.py` — `PixelShuffleModule` (`torch.pixel_shuffle`) + `PixelShuffleTest` delegation smoke test. - `test/op_tests/cases.py` — `_pixel_shuffle_suite`: r2, r2_batch, r3, r2_3d. ghstack-source-id: 406366790 @exported-using-ghexport Differential Revision: [D112257661](https://our.internmc.facebook.com/intern/diff/D112257661/)
Pull Request resolved: #21209 Problem: The WebGPU delegate has no `aten.argmax.default` / `aten.argmin.default`, and — more fundamentally — no int64-OUTPUT support: the AOT downcasts the int64 index to an int32 GPU buffer, but the ExecuTorch program output is int64, and `copy_outputs` raw-copied the int64 EValue's bytes from the int32 staging buffer (a size mismatch). argmax would be the first int64-output op. argmax unlocks on-GPU decode sampling. Solution: Port `et_vk.argmax`/`argmin` (last-dim arg-reduction -> int64 index) sharing one handler, mirroring Vulkan `ArgReduce.cpp` (`arg_reduce_impl`, last-dim only, the `add_reduce_per_row_node` accumulator that tracks `{val, idx}`). The kernel reuses the landed `amax` last-dim reduction + index-tracking: a strict `>` (argmax) / `<` (argmin) scan keeps the FIRST extremum (= torch tie-break), writing the index as int32 (1 u32/row) to the int32 GPU buffer. Add the int32->int64 output-widening path: `copy_outputs` now maps each output's LIVE staging size (`cur_nbytes`) and, when the host EValue is 2x the int staging buffer, sign-extends int32->int64 into the EValue; matched-dtype outputs keep the unchanged raw-copy path. Implementation: `argmax/Reduce.cpp` registers `aten.argmax.default` + `aten.argmin.default` -> `arg_reduce_impl(graph, args, is_argmin)`; args `[in, dim, keepdim, out]`, `dim` scalar (last-dim guard), `out = args.back()`, guards fp32 input / int32 output / shape, resize hook. `arg_reduce.wgsl` is one row per thread. `WebGPUGraph::copy_outputs` maps `cur_nbytes` + the guarded `dst == 2*map && is_int` widen. `WebGPUBackend::execute` wraps `execute()` + `copy_outputs()` in try/catch so a defensive throw never crosses the backend boundary. Constraints: last-dim reduction only (mirrors Vulkan `normalized_dim == ndim-1`); fp32 input, int32-backed int64 index output. The `copy_outputs` change is byte-identical for every existing fp32/int8 output (`cur_nbytes == EValue nbytes` for matched dtypes) — verified by regression (floor_divide fp32, q8ta convs int8 all still pass). ghstack-source-id: 406366846 @exported-using-ghexport Differential Revision: [D112257656](https://our.internmc.facebook.com/intern/diff/D112257656/)
Pull Request resolved: #21210 Problem: The new argmax/argmin ops output int64 indices, but the op-test harness was float/int8-only (it cast any int output to fp32, so the driver would read index bytes as float — a garbage compare). Solution: Extend the shared harness with an int64-golden path (mirroring the landed int8 path): the generator writes an int64 golden for int32/int64 raw outputs (`_write_int64`, dtype "int64"), and the driver reads `const_data_ptr<int64_t>()` and exact-compares to the int64 golden (`load_int64_bin`). The argmax/argmin suites use `golden_dtype="float32"` so the golden's fp32 argmax matches the fp32 kernel exactly (an fp64 oracle could flip a near-tie index). Implementation: `generate_op_tests.py` adds `_write_int64` + the `is_int64` branch. `op_test_driver.cpp` adds the int64 exact-compare branch; `driver_util.{h,cpp}` add `load_int64_bin`. `cases.py` registers `argmax`/`argmin` with randn 2d/3d cases (max/min at an INTERIOR index — exercises the reduction walk) and a deliberate-tie case per op (`argmax_tie_gen`/`argmin_tie_gen` place a repeated extremum at idx 1 and 3 so the FIRST occurrence wins — this discriminates the strict-`>`/`<` tie-break from a `>=`/`<=` bug). `test_argmax.py` holds the modules + tie gens. ghstack-source-id: 406366845 @exported-using-ghexport Differential Revision: [D112257677](https://our.internmc.facebook.com/intern/diff/D112257677/)
Pull Request resolved: #21211 Problem: The WebGPU delegate has no `et_vk.linear_qcs4w` — a 4-bit channels-symmetric-weight linear (per-output-channel symmetric weight, no zero-point). It is reachable via the `VulkanQuantizer` weight-only 4-bit path (distinct from the XNNPACK static-PT2E path that produces the q8ta ops), but had no WebGPU handler. Solution: Port `et_vk.linear_qcs4w` (fp32 activation, int4 weight), mirroring the landed `linear_q4gsw` register-tiled buffer GEMM (`quantized_linear/QuantizedLinear.cpp` + `q4gsw_linear.wgsl`) simplified to per-channel scale. Vulkan ref: `impl/QuantizedLinearQCSNW.cpp` `linear_qcs4w` (`check_linear_qcsnw_args`: args `[mat1, qmat2=[N,K/2] 4-bit, scales=[N], out]`, symmetric per-output-channel, no group, no zero-point, no bias) + `glsl/linear_qcsnw_coop.glsl`. Before/After vs q4gsw: q4gsw scale is grouped `scales[(k/group_size)*padded_N + n]`; qcs4w scale is per-channel `scales[n]` (1D [N]) — no group_size, no padded_N, no bias. Implementation: `linear_qcs4w/QuantizedLinearQcs4w.cpp` registers `et_vk.linear_qcs4w` -> `qcs4w_linear_impl`, args `[in, weight, scales, out]` (out=args.back()). `qcs4w_linear.wgsl`: register-tiled (TM=TN=4) GEMM, `acc += in[m,k] * (unpack_int4(w) - 8) * scales[n]`, 2D-folded dispatch (lifts the 65535 cap). `Qcs4wParams` (16 bytes: M/N/K/K_packed) matches the WGSL Params. Guards fp32 in/out, `K_packed==ceil(K/2)`, `N*K_packed%4==0` (u32-packed), `scales>=N`, all fail-loud. Resize hook recomputes live M + dispatch. Constraints / divergences from the Vulkan reference: (1) buffer re-derivation of Vulkan's texture-based qcs4w GEMM (WebGPU always buffers). (2) **CRITICAL — nibble order: the qcs4w AOT packer (`_passes/fuse_quantized_ops.py`) stores `(even_col<<4)|odd_col`, the SWAP of q4gsw's `pack_4bit_weight_tensor` `(odd<<4)|even`** — so this kernel reads even-k from the HIGH nibble and odd-k from the LOW nibble, the reverse of `q4gsw_linear.wgsl` (verified against the packer, the Vulkan `linear_qcsnw_coop.glsl` unpack, and a byte-search of the served `.pte`). (3) `+8`-shifted int4 recovered as signed `[-8,7]` (same as q4gsw). (4) one tiled kernel only (q4gsw's GEMV/shmem perf variants are separate, Canary-gated follow-ups). Bias is out of the op (a Linear bias lowers to a separate `aten.add`). ghstack-source-id: 406366854 @exported-using-ghexport Differential Revision: [D112257623](https://our.internmc.facebook.com/intern/diff/D112257623/)
Pull Request resolved: #21212 Problem: The new `et_vk.linear_qcs4w` op needs golden coverage — and, being a quantized op reachable only through a specific quantizer, a test that actually produces the op and validates the fake-quant numerics. Solution: `make_qcs4w_linear_module` runs a plain `nn.Linear` through the `VulkanQuantizer` weight-only 4-bit path (`get_symmetric_quantization_config(is_dynamic=False, weight_bits=4)` -> prepare_pt2e -> calibrate -> convert_pt2e), and the op-test `module_factory` returns the CONVERTED module, so the WebGPU output is goldened against the converted eager (fp32 per-channel fake-quant reference). This exercises the int4 unpack + per-channel dequant end-to-end. `bias=False` keeps the golden focused on qcs4w (a Linear bias lowers to a separate `aten.add`). Implementation: `cases.py` registers `linear_qcs4w` with `basic` (4x32x16), `gemv` (M=1 decode shape), `k64` (2x64x8), `n32` (3x32x32) — all with K even (2 nibbles/byte) and `N*ceil(K/2) % 4 == 0` (u32-packed weight); `test_linear_qcs4w.py` holds the module + a delegation smoke test asserting `et_vk.linear_qcs4w` is absorbed into the VulkanBackend delegate. ghstack-source-id: 406366862 @exported-using-ghexport Differential Revision: [D112257668](https://our.internmc.facebook.com/intern/diff/D112257668/)
Pull Request resolved: #21213 Problem: The WebGPU delegate has no `et_vk.linear_q8ta_q8csw` — an int8-activation x int8-channelwise-weight linear with a FP32 output (the "output-not-requantized" sibling of the landed `q8ta_linear`, e.g. a final projection). Its kernel design was ready but it was thought unreachable; the reachability blocker is now solved. Solution: Port `et_vk.linear_q8ta_q8csw` (int8 x, int8 per-channel weight, fp32 out), mirroring the landed `q8ta_linear` (`q8ta_linear/Q8taLinear.cpp`) with the output requant removed. Before/After vs `q8ta_linear`: same i32 GEMM `Σ(x_int8 - input_zp) * w_int8` + dequant `* input_scale * weight_scales[n] + bias`, but `q8ta_linear` then requantizes to int8 (needs output_scale/zp, N%4==0 for output packing) whereas this writes fp32 directly (no output qparams in the schema, N unconstrained). Mirrors Vulkan `impl/QuantizedLinear.cpp` `linear_q8ta_q8csw`. Implementation: `linear_q8ta_q8csw/LinearQ8taQ8csw.cpp` registers `et_vk.linear_q8ta_q8csw` -> `linear_q8ta_q8csw_impl`. Schema `[x, input_scale, input_zp, weights, weight_sums, weight_scales, bias?, out]`: x/weights int8 (bound `array<u32>`, unpacked in-shader), `weight_sums` (arg 4) folded per-element (unused, matching q8ta_linear), `out=args.back()`. `linear_q8ta_q8csw.wgsl` register-tiled (TM=TN=4) i32 GEMM, 2D-folded dispatch; the TN=4 tile guards `n<N` (no N%4 assumption — fp32 output has no packing constraint). `Q8taQ8cswParams` (32 bytes) matches the WGSL Params. Guards int8 x/weight (numel%4==0 for the u32 reads), fp32 out, scales fp32 [N], all fail-loud. Constraints / divergences from the Vulkan reference: fp32 output only. A bias-less TERMINAL linear can mis-fuse to an int8 output the schema cannot compute (no output scale/zp) — the handler fail-louds on it (`output must be fp32`); the correct fp32-out form is what the fusion routes when the output is not re-quantized. i32 accumulator (inherited from q8ta_linear). Buffer-only, so no int8x4 texture layout. ghstack-source-id: 406366865 @exported-using-ghexport Differential Revision: [D112257666](https://our.internmc.facebook.com/intern/diff/D112257666/)
Pull Request resolved: #21214 Problem: The new `et_vk.linear_q8ta_q8csw` op needs golden coverage, and the recipe to REACH it (a quantized linear whose output stays fp32) is non-obvious. Solution: `make_linear_q8ta_q8csw_module` runs a plain `nn.Linear` through XNNPACK static PT2E with the activation config's `output_activation` nulled (`dataclasses.replace(get_symmetric_quantization_config(is_per_channel=True, is_dynamic=False), output_activation=None)`): the linear's INPUT is statically per-tensor quantized but its OUTPUT is left fp32, so the Vulkan fusion routes to `linear_q8ta_q8csw` (fp32 out) instead of `q8ta_linear` (int8 out). The op-test `module_factory` returns the CONVERTED module, so the WebGPU output is goldened against the converted eager (fp32 fake-quant reference); the served subgraph is `quantize_per_tensor` (landed) -> `linear_q8ta_q8csw`. Cases use `bias=True` (a bias-less terminal linear mis-fuses to an int8 output the op fail-louds on) and N a multiple of 4 (the AOT pads the quantized weight's N). Implementation: `cases.py` registers `linear_q8ta_q8csw` with `basic` (4x32x16), `gemv` (M=1), `k48` (2x48x8), `n32` (3x32x32); `test_linear_q8ta_q8csw.py` holds the module + a delegation smoke test asserting `et_vk.linear_q8ta_q8csw` is absorbed into the VulkanBackend delegate. ghstack-source-id: 406366863 @exported-using-ghexport Differential Revision: [D112257659](https://our.internmc.facebook.com/intern/diff/D112257659/)
Pull Request resolved: #21215 Problem: The WebGPU delegate has no `et_vk.grid_priors` — a detection anchor-grid op (generates per-cell coordinate shifts). It delegates through `VulkanPartitioner` (`op_registry.py` tags `et_vk.grid_priors.default`; reachable via a direct custom-op call) but had no WebGPU handler. Solution: Port `et_vk.grid_priors` (fp32), mirroring Vulkan `impl/GridPriors.cpp` + `glsl/grid_priors.glsl`. The output is `[H*W, 2]` computed purely from the input's H/W plus scalar `stride`/`offset` — the input's VALUES are not read. Per output row `r`: `out[r,0] = (r%W + offset) * stride` (x-shift), `out[r,1] = (r/W + offset) * stride` (y-shift), matching the Vulkan glsl (`pos.x==0 -> pos.y%width`, `pos.x==1 -> pos.y/width`) and the CPU eager (`custom_ops_lib.py` grid_priors: `stack((shift_x_per_col, shift_y_per_row))`). Implementation: `grid_priors/GridPriors.cpp` registers `et_vk.grid_priors.default` -> `grid_priors_impl`. Args `[in, stride(int), offset(float), out]`; only `in.dims[-2:]` (H, W) are used. `grid_priors.wgsl`: one thread per output element, `row=idx/2`, even idx -> `row%W`, odd -> `row/W`, `(coord+offset)*stride`; 2D-folded dispatch. `GridPriorsParams` (16 bytes: numel/width/stride/offset) matches the WGSL Params. Only the output (storage) + params (uniform) are bound — the input's data is intentionally unread. Resize hook recomputes numel/width + dispatch + out dims `[H*W, 2]`. Guards: in >= 2D, W>0, numel<=u32, fp32 out, all fail-loud. Constraints / divergences from the Vulkan reference: buffer re-derivation of Vulkan's texture (`imageStore`/`VEC4_T`) kernel — the flat-index buffer form is forced by the backend's buffer-only design. `stride` (int) is promoted to f32 for the multiply (exact for detection-scale strides; matches the glsl's int->float promotion). ghstack-source-id: 406366861 @exported-using-ghexport Differential Revision: [D112257635](https://our.internmc.facebook.com/intern/diff/D112257635/)
Pull Request resolved: #21216 Problem: The new `et_vk.grid_priors` op needs golden coverage — and, since the input's values are unused (only its H/W matter), a test that validates the computed grid + the column ordering across shapes. Solution: `GridPriorsModule` calls the custom op directly with baked `stride`/`offset`, so only the float tensor `x` is a runtime input. The op has a CPU eager impl (an independent meshgrid/stack reference), so the op-test framework goldens the WebGPU output against it directly (float32 oracle — the output is an exact computed grid). Non-square shapes (8x10, 5x7) make the two coordinate ranges differ, so a swapped `r%W`/`r/W` column mapping would diverge from the golden. Implementation: `cases.py` registers `grid_priors` with `s8` (8x10, stride=8, offset=0.5), `s16` (4x4, stride=16, offset=0), `offset0` (5x7, stride=4, offset=0); `test_grid_priors.py` holds the module + a delegation smoke test asserting `et_vk.grid_priors` is absorbed into the VulkanBackend delegate. ghstack-source-id: 406366867 @exported-using-ghexport Differential Revision: [D112257599](https://our.internmc.facebook.com/intern/diff/D112257599/)
Pull Request resolved: #21217 Problem: The WebGPU delegate has no `et_vk.conv_with_clamp` — a fused fp32 2D convolution + output clamp (e.g. a `Conv2d` followed by `relu6`). It delegates through `VulkanPartitioner` (`op_registry.py` tags `et_vk.conv_with_clamp.default`; `nn.Conv2d`+`F.relu6` fuses to it) but had no WebGPU handler. Solution: Port `et_vk.conv_with_clamp` (fp32, groups==1), mirroring the Vulkan `conv` handler (`impl/Convolution.cpp:837`) + the eager `clamp(convolution(...), output_min, output_max)` (`custom_ops_lib.py`). Direct windowed conv reusing the staged `q8ta_conv2d` windowing structure but fp32 (no int8 unpack/dequant/requant) + a clamp epilogue: per NCHW output cell `out[n,oc,oh,ow] = clamp(bias[oc] + Σ_{ic,kh,kw} weight[oc,ic,kh,kw] * input[n,ic, oh*sh-ph+kh*dh, ow*sw-pw+kw*dw], output_min, output_max)`, OOB taps skipped. Implementation: `conv_with_clamp/ConvWithClamp.cpp` registers `et_vk.conv_with_clamp.default` -> `conv_with_clamp_impl`. Args `[in, weight, bias?, stride, padding, dilation, transposed, output_padding, groups, output_min?, output_max?, out]`; stride/padding/dilation as int-pairs, `output_min`/`output_max` as `Scalar?` (Double/Int -> value, absent -> -/+inf so `clamp` is a passthrough). `conv_with_clamp.wgsl`: one thread per output element, RAW `[OC,IC,Kh,Kw]` weight (the fp32 buffer path gets the raw constant, like conv1d_dw/pw — not the q8ta im2col `[OC,Kh*Kw*IC]`), 2D-folded dispatch. `ConvWithClampParams` (80 bytes) matches the WGSL Params. Guards: in/weight/out 4D + fp32, `weight.dims[1]==IC`, numel<=u32, groups==1 + not-transposed, all fail-loud. Constraints / divergences from the Vulkan reference: buffer re-derivation of Vulkan's texture conv (`conv2d_prepack_weights.glsl`) — the naive one-thread-per-output form is forced by the buffer-only backend. Scoped to `groups==1` + not-transposed (fail-loud otherwise). Note: overlaps the general `convolution` op — this handles the clamped variant only. ghstack-source-id: 406366872 @exported-using-ghexport Differential Revision: [D112257644](https://our.internmc.facebook.com/intern/diff/D112257644/)
Pull Request resolved: #21218 Problem: The new `et_vk.conv_with_clamp` op (a first-of-its-kind general fp32 conv2d) needs golden coverage that exercises the H/W axis separation, not just square/symmetric configs. Solution: `ConvWithClampModule` is a plain `nn.Conv2d` + `F.relu6`, which the Vulkan fusion rewrites to `et_vk.conv_with_clamp` (fp32 conv + clamp[0,6]); the conv weight/bias are baked params and only `x` is a runtime input. Goldened vs the module's fp32 eager. The `asym` case is fully axis-asymmetric (input H=7≠W=9, kernel Kh=2≠Kw=3, stride 1≠2, padding 1≠0, dilation 2≠1 → output H_out=7≠W_out=4), so an H↔W index swap anywhere (unravel divisors, stride/pad/dilation axis, or Kh/Kw) diverges from the golden. Implementation: `cases.py` registers `conv_with_clamp` with `k3p1`/`stride2`/`dil2`/`no_bias`/`asym` (all groups==1); `test_conv_with_clamp.py` holds the module + a delegation smoke test asserting `et_vk.conv_with_clamp` is absorbed into the VulkanBackend delegate. ghstack-source-id: 406366873 @exported-using-ghexport Differential Revision: [D112257609](https://our.internmc.facebook.com/intern/diff/D112257609/)
…> bool) Pull Request resolved: #21219 Problem: The WebGPU delegate has no elementwise comparison ops (`aten.eq/lt/le/gt/ge.Tensor`). They delegate through `VulkanPartitioner` (all 5 tagged, verified) and output BOOL tensors, but the backend had no handler and no bool-output path. Solution: Port all 5 comparisons via one shared handler + WGSL kernel (op-code switch), mirroring Vulkan `impl/BinaryOp.cpp` + `glsl/binary_op_buffer.yaml`. Output is a 1-byte bool tensor (`vk_datatype_size(BOOL)=1`, `is_int=true`, `is_int8=false`); the kernel packs 4 bool bytes per `u32` word (reusing the int8 buffer-binding idiom), and `copy_outputs` raw-copies it (dst==map, no widen). Per element: `out = (a OP b) ? 1 : 0` with OP in {==, <, <=, >, >=}. Implementation: `compare/Compare.cpp` registers `aten.eq/lt/le/gt/ge.Tensor` -> `compare_impl(graph, args, op)` (op 0=eq..4=ge). `compare.wgsl`: one thread per output word (4 bool bytes), 2D-folded dispatch, op switch. `CompareParams` (16 bytes: num_elements, op) matches the WGSL Params. Guards: fp32 inputs, 1-byte bool output, same-shape (in/out numel equal), `numel%4==0` (bool packs 4/word AND gates the readback map) — all fail-loud. Resize hook recomputes numel/dispatch + re-applies the %4 guard + cross-checks both operands. Constraints / divergences from the Vulkan reference: (1) `numel%4==0` (bool-output packing; fail-loud otherwise, mirroring the q8ta N%4 output-pack constraint). (2) same-shape only (flat kernel; broadcast is export-smoke, like `add`/`minimum`). (3) `eq` is torch-EXACT `a==b` — a DELIBERATE deviation from Vulkan's float `abs(X-Y)<1e-5`: the delegate serves torch-exact model semantics and the op-test golden is torch eager; `lt/le/gt/ge` mirror Vulkan's exact comparators. ghstack-source-id: 406366878 @exported-using-ghexport Differential Revision: [D112257588](https://our.internmc.facebook.com/intern/diff/D112257588/)
…rness Pull Request resolved: #21220 Problem: The new comparison ops (`aten.eq/lt/le/gt/ge.Tensor`) output BOOL, which the op-test golden harness did not handle (only fp32/int8/int64). And the goldens must actually discriminate the comparators. Solution: Extend the golden harness with a `bool` branch (mirrors the int8/int64 branches): `generate_op_tests` writes the bool golden as 0/1 bytes (`raw.dtype==torch.bool -> _write_int8(raw.to(int8))`, dtype "bool") and `op_test_driver` reads the bool output via `const_data_ptr<bool>()` and byte-compares. The change is additive (bool is a distinct dtype; the branch precedes int8) — existing fp32/int8/int64 paths are untouched (regression 27/27). `CompareModule(op)` covers all 5 ops; the two inputs draw from DIFFERENT discrete-range seeds (`compare_gen_a`/`compare_gen_b`) so `a!=b` with frequent collisions — giving `eq/le/ge` genuine ties and `lt/gt` a real true/false mix, so an op-switch swap diverges from the golden. Implementation: `cases.py` registers `eq`/`lt`/`le`/`gt`/`ge` via a shared `_compare_suite` (2d/3d/sq shapes, all numel%4==0); `test_compare.py` holds `CompareModule` + the seeded gens + a delegation smoke test. ghstack-source-id: 406388747 @exported-using-ghexport Differential Revision: [D112257586](https://our.internmc.facebook.com/intern/diff/D112257586/)
Pull Request resolved: #21221 Adds the elementwise bool `aten.logical_and.default` to the WebGPU delegate. Vulkan maps this op to its `bitwise_and` handler (`BinaryOp.cpp:161`), valid for canonical 0/1 bool bytes; this port mirrors that. Both operands and the output are 1-byte bool tensors packed 4/word into `array<u32>`, so a word-wise AND is exactly a per-byte AND. Key changes: - `runtime/ops/logical_and/logical_and.wgsl` — one thread per packed word, `t_out[w] = t_a[w] & t_b[w]` (2D-folded dispatch). - `runtime/ops/logical_and/LogicalAnd.cpp` — handler binds out (rw) + a/b (ro) + params (uniform); fail-loud guards on bool dtype, null buffer, and same-shape; resize hook re-applies `num_words` + dispatch. - `CMakeLists.txt` — register the op source. Same-shape only (`numel % 4 == 0` for the packed bool bindings), matching the backend's other bool ops. Co-authored-with: Claude Code. ghstack-source-id: 406388750 @exported-using-ghexport Differential Revision: [D112257653](https://our.internmc.facebook.com/intern/diff/D112257653/)
Pull Request resolved: #21222 Adds the `aten.logical_and.default` op-test suite to the manifest-driven WebGPU op-test framework, stacked above the op diff. Key changes: - `test/ops/test_logical_and.py` — `LogicalAndModule` derives its two bool operands on-GPU from float inputs (`a > 0`, `b > 0` via the delegated `gt.Tensor` against a baked zero buffer), so the only runtime inputs are the two float tensors (the op-test framework is float-input-only). Distinct `a`/`b` seeds make the two masks differ (each ~50% True, independent -> AND ~25% True), a real mix that an OR mutant would fail. `LogicalAndTest` is the export-delegation smoke test. - `test/op_tests/cases.py` — registers `logical_and` (shapes 2d/3d/sq, all `numel % 4 == 0`), byte-exact bool golden. Co-authored-with: Claude Code. ghstack-source-id: 406388749 @exported-using-ghexport Differential Revision: [D112257583](https://our.internmc.facebook.com/intern/diff/D112257583/)
Pull Request resolved: #21223 Adds the two bool bitwise ops the Vulkan partitioner tags (`op_registry.py` marks `bitwise_and.Tensor`, `bitwise_not.default`, `logical_and.default` all `BOOL_T`). `bitwise_and` on bool is identical to `logical_and`, so it shares the one AND handler — mirroring Vulkan, which registers both `aten.bitwise_and.Tensor` and `aten.logical_and.default` to its `bitwise_and` handler (`BinaryOp.cpp:160-161`). No new kernel needed. `bitwise_not` is the bool NOT: Vulkan uses `1 - X` on uint8 (`unary_op.yaml`); for canonical 0/1 that equals a per-byte low-bit flip, so the packed-word kernel is `t_out[w] = t_a[w] ^ 0x01010101u` (one thread per word). Key changes: - `runtime/ops/logical_and/LogicalAnd.cpp` — register `aten.bitwise_and.Tensor` on the existing shared bool-AND handler. - `runtime/ops/bitwise_not/{BitwiseNot.cpp,bitwise_not.wgsl}` — new unary; out (rw) + a (ro) + params (uniform); fail-loud bool-dtype/null/same-shape guards; resize hook re-applies `num_words`. - `CMakeLists.txt` — register `BitwiseNot.cpp`. Bool only (`numel % 4 == 0` for the packed bindings), matching the backend's other bool ops. Co-authored-with: Claude Code. ghstack-source-id: 406388751 @exported-using-ghexport Differential Revision: [D112257647](https://our.internmc.facebook.com/intern/diff/D112257647/)
Pull Request resolved: #21224 Adds the `aten.bitwise_and.Tensor` / `aten.bitwise_not.default` (bool) op-test suites, stacked above the op diff. Key changes: - `test/ops/test_bitwise.py` — `BitwiseAndModule`/`BitwiseNotModule` derive their bool operands on-GPU from float inputs (`a > 0` via the delegated `gt.Tensor` against a baked zero buffer), so the only runtime inputs are float tensors (the op-test framework is float-input-only). `bitwise_and` uses distinct `a`/`b` seeds (AND ~25% True); `bitwise_not` inverts one ~50% mask. - `test/op_tests/cases.py` — registers `bitwise_and`/`bitwise_not` (shapes 2d/3d/sq, all `numel % 4 == 0`), byte-exact bool goldens. Co-authored-with: Claude Code. ghstack-source-id: 406388753 @exported-using-ghexport Differential Revision: [D112257674](https://our.internmc.facebook.com/intern/diff/D112257674/)
… shared-memory Pull Request resolved: #21225 Problem: The reduction ops (amax/amin/argmax/sum/mean), softmax/log_softmax, and native_group_norm's reduce pass used naive one-thread-per-row serial loops — far below Vulkan's optimization level, where backends/vulkan/runtime/graph/ops/glsl/reduce.glsl uses a cooperative shared-memory reduction. Solution: Rewrite each to Vulkan's cooperative design — one workgroup per reduction row/group, each thread reduces a strided slice into a workgroup shared array, thread 0 aggregates the partials. Dispatch changed to one workgroup per row. argmax preserves torch's first-index tie-break; softmax keeps numerical stability + the log variant. ghstack-source-id: 406388754 @exported-using-ghexport Differential Revision: [D112257649](https://our.internmc.facebook.com/intern/diff/D112257649/)
…r, _to_copy) + enable unary activation family Pull Request resolved: #21226 Adds the creation + cast/copy ops and enables the unary activation family in the op allowlist. Key changes: - `runtime/ops/full/{Full.cpp, full.wgsl}` — `full`/`full_like`/`zeros`/`zeros_like`/`ones`/`ones_like`/`scalar_tensor` (shared fp32 fill; constant-folded before the delegate in practice, so effectively register-only) + a fail-loud non-fp32 output guard. - `runtime/ops/to_copy/{ToCopy.cpp, to_copy.wgsl}` — `_to_copy` / `_to_dim_order_copy`: same-dtype raw copy, and int<->float CONVERT via `f32(bitcast<i32>(x))` / `bitcast<f32>(i32(x))`. The dtype-promotion pass inserts an int->float `_to_copy` before binary ops, so a byte-reinterpret would ship denormal/inf garbage; this converts. Mirrors Vulkan `ToCopy.cpp` (view_convert vs blit branch). - `test/tester.py` — enable the unary activation family. ghstack-source-id: 406388756 @exported-using-ghexport Differential Revision: [D112257603](https://our.internmc.facebook.com/intern/diff/D112257603/)
…shape/movement, pooling, conv/sampling) Pull Request resolved: #21227 Problem: These ops already had kernels in the granular tree but were absent from WEBGPU_SUPPORTED_OPS, so the op-test framework neither delegated nor golden-tested them. Solution: Add to the allowlist: minimum, div.Tensor_mode (floor_divide), logical_and, bitwise_and, bitwise_not, flip, repeat, index_select, avg_pool2d, pixel_shuffle, convolution (depthwise conv1d), conv_with_clamp, grid_sampler_2d, grid_priors. ghstack-source-id: 406388757 @exported-using-ghexport Differential Revision: [D112257605](https://our.internmc.facebook.com/intern/diff/D112257605/)
…+ qcs4w/q8ta_q8csw linear) Pull Request resolved: #21228 Problem: The int8 quantized ops (q8ta elementwise/conv/linear + qcs4w and q8ta_q8csw linear) already had kernels in the granular tree but were absent from WEBGPU_SUPPORTED_OPS, so they were neither delegated nor golden-tested. Solution: Add to the allowlist: et_vk.linear_qcs4w, et_vk.linear_q8ta_q8csw, et_vk.q8ta_add, et_vk.q8ta_relu, et_vk.q8ta_pixel_shuffle, et_vk.q8ta_linear (+ q8ta_linear_gemv), et_vk.q8ta_conv2d, et_vk.q8ta_conv2d_dw, et_vk.q8ta_conv2d_pw. ghstack-source-id: 406388761 @exported-using-ghexport Differential Revision: [D112257636](https://our.internmc.facebook.com/intern/diff/D112257636/)
…8da4w) Pull Request resolved: #21229 Problem: WebGPU lacked the dynamic-8bit-activation x 4-bit-group-weight linear (the 8da4w path real LLMs use for dynamic quant). Vulkan registers `et_vk.linear_dq8ca_q4gsw` + `torchao.choose_qparams_affine`; WebGPU delegated neither, so any `Int8DynamicActivationIntxWeightConfig` model fell back to CPU. Solution: Author the reachable pair. `torchao.choose_qparams_affine` computes per-row asymmetric int8 activation scale/zp; `et_vk.linear_dq8ca_q4gsw` folds that dynamic activation quant into the existing q4gsw 4-bit-group GEMM. Both mirror the Vulkan reference (`ChooseQParams.cpp`, `QuantizedLinear.cpp:760`); the activation-quant + weight-dequant math was CPU-de-risked exact against torchao before authoring. Impl: `choose_qparams_affine.wgsl` does a cooperative per-row min/max reduction (one workgroup per block of 4 rows so the 4 int8 zps pack into one u32 with no write race) then `calculate_scale_and_zero_point` mirroring Vulkan. `linear_dq8ca_q4gsw.wgsl` is the register-tiled q4gsw GEMM with `out[m,n] = s[m] * sum_k (xq[m,k]-z[m]) * dequant(w)`, reading the packed-int8 zp. int8 zp (`elem_size` 1) is bound word-aligned; `num_rows` must be <=4 or a multiple of 4 (int8 buffers alloc `max(nbytes,4)`) — arbitrary-M prefill is a documented follow-up. ghstack-source-id: 406388760 @exported-using-ghexport Differential Revision: [D112257680](https://our.internmc.facebook.com/intern/diff/D112257680/)
Pull Request resolved: #21230 Add `aten.bitwise_or.Tensor` + `aten.logical_or.default` to the WebGPU backend — the OR counterpart to the existing `bitwise_and`/`logical_and` family, closing the last bool-binary parity gap with Vulkan. Both overloads share one handler (`logical_or_op`); on canonical 0/1 bool bytes the word-wise `|` over the packed `u32` buffer equals a per-byte OR. Kernel-only: the Vulkan partitioner already tags both ops (`op_registry.py` `register_bool_binary_ops`), so no partitioner change. Key changes: - `runtime/ops/logical_or/` — `logical_or.wgsl` (`t_out = t_a | t_b`, bool packed 4/word) + `LogicalOr.cpp` (one handler; bool-only + `numel % 4` guards + dynamic-shape resize hook; registers both overloads). Mirrors `logical_and`. - `test/ops/test_logical_or.py` + `op_tests/cases.py` — delegation smoke + op-test golden for both ops (two ~50% masks -> OR ~75% True, which distinguishes OR from an AND mutant). - `CMakeLists.txt` — wire `logical_or/LogicalOr.cpp`. Co-authored-with: Claude Code. ghstack-source-id: 406388764 @exported-using-ghexport Differential Revision: [D112483463](https://our.internmc.facebook.com/intern/diff/D112483463/)
… (0x30) Pull Request resolved: #21231 Rank-5+ tensors threw `0x30` (`DelegateInvalidCompatibility`) at delegate init: the shared `TensorMeta` UBO capped rank at 4 (`kTensorMetaMaxNdim`), so `fill_tensor_meta`/`fill_tensor_meta_broadcast` hard-throw during `build()` and the WGSL mirror stored `sizes`/`strides` as a single `vec4<u32>`. This clears `test_permute_different_shapes[webgpu]`, `test_transpose_different_shapes[webgpu]` (transpose lowers to `permute_copy`), and the rank>4 model failures (`conformer`, `maxvit_t`). Raise the global cap to 8 and widen the uniform to two `vec4<u32>` blocks, keeping the C++ and WGSL layouts byte-identical: - C++ `sizes[8]`/`strides[8]`: `sizeof(TensorMeta)` 48->80, `strides` offset 32->48 (`ndim`/`numel`/`sizes` offsets unchanged). - WGSL `sizes: array<vec4<u32>, 2>` / `strides: array<vec4<u32>, 2>`: in the uniform address space the `vec4<u32>` element stride is 16, so two blocks are 32 contiguous bytes matching `uint32_t[8]`; component k sits at byte `16 + 4*k` on both sides, so the raw memcpy stays correct. Runtime reads `m.strides[d]` become `m.strides[d >> 2u][d & 3u]`. - 8 (not 5): a WGSL uniform array cannot pack 5 `u32` tightly, so rank-5 would allocate the same two `vec4`s; 8 is the natural 2-`vec4` boundary and covers all realistic ranks. Key changes: - `runtime/ops/TensorMeta.h` - cap 4->8; `static_assert(sizeof == 80)` and `offsetof(strides) == 48`; throw text. - 15 shaders widened + reindexed (regenerated their 16 `_wgsl.h`): `permute`, `flip`, `cat`, `binary_op` (`div`/`sub`), `binary_mul`, `binary_pow`, `binary_minimum`, `binary_floor_divide`, `where`, `select`, `slice`, `expand_copy`, `gather`, `index_select`, `repeat`. - `permute`/`flip` also widen their per-dim `Params` array (`perm`/`flip` `vec4<u32>` -> `array<vec4<u32>, 2>`; `sizeof(PermuteParams)`/`sizeof(FlipParams)` 16->32). - Rank-guard throw messages "exceeds 4" -> "exceeds 8" (the guards themselves auto-follow the constant): `div`/`mul`/`sub`/`where`/`binary`/`index_select`/`repeat`. - Applied identically to both the `xplat/` and `fbcode/` mirrors (byte-identical). Co-authored-with: Claude Code. ghstack-source-id: 406388763 @exported-using-ghexport Differential Revision: [D113319595](https://our.internmc.facebook.com/intern/diff/D113319595/)
…match) Pull Request resolved: #21232 Under dynamic shapes `cat` produced a wrong (scrambled + tail-garbage) output. `cat_impl` baked every shape-dependent quantity - each input `in_meta` (strides/numel), the shared `out_meta` (strides), the per-input `off_k`, and each dispatch's `workgroup_count` - from the MAX (upper-bound) build shape and registered NO resize hook, then released the UBOs. Under a smaller live shape the kernel then decoded coords with max strides, scattered with max out-strides, and looped over the max numel. This is the `[dynamic]` output-mismatch behind `test_inception_v3`, `test_squeezenet1_1`, `test_densenet161` (all channel-cat on spatial-dynamic feature maps; the concat dim is static so `off_k` was already correct - the defect is the stale spatial strides/numel). Fix (mirrors the WebGPUGraph SwiGLU/QKV resize templates and the shipped `mul` hook): - Keep the per-input `in_meta`/`params` UBOs and the shared `out_meta` UBO alive via `own_uniform_buffer` (previously released after build); collect their handles plus each `add_dispatch` index. - Register an idempotent `add_tensor_resize_hook` on every input id: from `cur_dims` recompute live out dims (`set_cur_dims(out_id, ...)` to cascade to consumers + fix the delegate-output shape), rebuild `out_meta` + each input's `in_meta`/`params` and `wgpuQueueWriteBuffer` them, and rewrite each dispatch's `workgroup_count_x` via `compute_1d_workgroup_count`. On a static graph `cur_dims == dims`, so the hook rewrites identical values (no behavior change). Applied identically to both the `xplat/` and `fbcode/` mirrors (byte-identical). Co-authored-with: Claude Code. ghstack-source-id: 406388769 @exported-using-ghexport Differential Revision: [D113319596](https://our.internmc.facebook.com/intern/diff/D113319596/)
…dels Pull Request resolved: #21233 `add.Tensor` was elementwise-only: `binary_add.wgsl` computed `output[idx] = input1[idx] + alpha * input2[idx]` over a single flat index up to the output numel, so on a broadcast it read the smaller operand out of bounds (WebGPU robustness clamps/zeros -> silently wrong, not `0x30`). This is the failure mode behind the models the WebGPU flow skips (`resnet50`, `vit_b_16`, `swin_v2_t`, `convnext_small`, `mobilenet_v3_small`, `shufflenet_v2_x1_0`) and the `bcast_first`/`bcast_second` op cases. Give `add.Tensor` NumPy broadcasting by mirroring the shipped `mul`: - `binary_add.wgsl` - copy `binary_mul`'s kernel (rank-8 `TensorMeta` layout from the rank-cap fix below it in the stack): keep the identical-shape elementwise fast path (common case stays bit-for-bit `input1[idx] + alpha * input2[idx]`), else relinearize out idx -> per-input coords (clamp size-1 dims) and add. `alpha` is a second pipeline-override constant (read once at build, never rewritten on resize), so no extra UBO. - `add/BinaryOp.cpp` - port `mul_impl`: 3 `TensorMeta` UBOs via `fill_tensor_meta_broadcast`, rank + fp32 guards, 6-entry bind group, `constantCount = 2` ({wg_size, alpha}), 2D dispatch, and `mul`'s resize hook verbatim (rebuild the 3 metas from `cur_dims`, `set_cur_dims`, rewrite UBOs + dispatch); `own_uniform_buffer` the 3 metas. Drops the old flat `AddParams` path. - `flows/webgpu.py` - un-skip the 6 broadcast-add models + the `bcast_first`/`bcast_second` op cases (kept the `float16`/`float64` dtype skips and `hardswish`/`lstm_batch_sizes`/`upsample_nearest2d`). This sits above the rank-cap fix, so the shader uses the rank-8 `array<vec4<u32>, 2>` `TensorMeta` layout. Applied identically to both `xplat/` and `fbcode/` mirrors (byte-identical). Co-authored-with: Claude Code. ghstack-source-id: 406388773 @exported-using-ghexport Differential Revision: [D113319599](https://our.internmc.facebook.com/intern/diff/D113319599/)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR was created by the merge bot to help merge the original PR into the main branch.
ghstack PR number: #21154 by @JCNTH
^ Please use this as the source of truth for the PR details, comments, and reviews
ghstack PR base: https://github.com/pytorch/executorch/tree/gh/JCNTH/114/base
ghstack PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/114/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/JCNTH/113/orig
Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/114/orig
@diff-train-skip-merge