Skip to content

[ExecuTorch][WebGPU] Op-tests for clamp + hardtanh#21301

Merged
JCNTH merged 77 commits into
gh/JCNTH/116/origfrom
gh/JCNTH/117/orig
Jul 24, 2026
Merged

[ExecuTorch][WebGPU] Op-tests for clamp + hardtanh#21301
JCNTH merged 77 commits into
gh/JCNTH/116/origfrom
gh/JCNTH/117/orig

Conversation

@pytorchbot

Copy link
Copy Markdown
Collaborator

This PR was created by the merge bot to help merge the original PR into the main branch.
ghstack PR number: #21157 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/117/base
ghstack PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/117/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/JCNTH/116/orig
Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/117/orig

@diff-train-skip-merge

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/)
@pytorch-bot

pytorch-bot Bot commented Jul 24, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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 meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 24, 2026
JCNTH added 24 commits July 24, 2026 01:05
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: #21179

Ports `aten.grid_sampler_2d.default` as a bilinear grid sample, the last Phase B (vision) float op. Mirrors Vulkan's config-specific implementation: bilinear interpolation, border padding, align_corners=true.

Key changes:
- `runtime/ops/grid_sampler_2d/{GridSampler2d.cpp, grid_sampler_2d.wgsl}` (+ generated `_wgsl.h`) — per NCHW output element: read the normalized `(gx, gy)` from `grid[N, out_h, out_w, 2]`, unnormalize `(g+1)*0.5*(in_size-1)` (align_corners), clamp to `[0, in_size-1]` (border), gather the 4 neighbors, `out = mix(mix(s00,s10,wx), mix(s01,s11,wx), wy)`. Registered `aten.grid_sampler_2d.default`.
- `CMakeLists.txt` — `runtime/ops/grid_sampler_2d/GridSampler2d.cpp` in `WEBGPU_SRCS`.

Mirrors Vulkan `glsl/grid_sampler_2d.glsl` (same coord/clamp/interp) + `impl/GridSampler2d.cpp` (arg order `[in, grid, interp, padding, align, out]`; fail-loud `VK_CHECK_COND` guards for interp==bilinear / padding==border / align_corners==true; grid a contiguous `[N,Hout,Wout,2]` buffer). Buffer NCHW re-derivation (Vulkan is texture). The grid is bound as an `array<f32>` read buffer. A dynamic-shape resize hook recomputes out=`[in.N, in.C, grid.Hout, grid.Wout]` + params + dispatch; it triggers on BOTH `in` and `grid` (the out shape depends on grid), mirroring how `mul` hooks both operands. Fail-loud guards: config, 4D in/out/grid, fp32 in/out + grid `== N*Hout*Wout*2`, out numel `<= u32`.
ghstack-source-id: 406366792
@exported-using-ghexport

Differential Revision: [D112257585](https://our.internmc.facebook.com/intern/diff/D112257585/)
Pull Request resolved: #21180

Adds the op-test suite for the `grid_sampler_2d` op ported in the diff below.

Key changes:
- `test/ops/test_grid_sampler_2d.py` — `GridSampler2dModule` (`F.grid_sample`, bilinear/border/align_corners=True) + `GridSampler2dTest` delegation smoke test. Both inputs are float, so the framework feeds x + grid directly.
- `test/op_tests/cases.py` — `_grid_sampler_2d_suite`: sq, wide_in, batch.
ghstack-source-id: 406366795
@exported-using-ghexport

Differential Revision: [D112257658](https://our.internmc.facebook.com/intern/diff/D112257658/)
Pull Request resolved: #21181

Ports the depthwise conv1d configuration of `aten.convolution.default`, a Phase B (audio) op. Each channel convolves with its own 1D filter.

Key changes:
- `runtime/ops/conv1d_dw/{Conv1dDW.cpp, conv1d_dw.wgsl}` (+ generated `_wgsl.h`) — per NCL output element `(n, c, l_out)`: `sum_k weight[c, k] * input[n, c, l_out*stride - padding + k*dilation]` (window-clipped) + `bias[c]`. Registered `aten.convolution.default`, dispatching only the depthwise conv1d config (`groups==C`, weight `[C,1,K]`, not transposed) and throwing on any other convolution config (fail-loud).
- `CMakeLists.txt` — `runtime/ops/conv1d_dw/Conv1dDW.cpp` in `WEBGPU_SRCS`.

Mirrors Vulkan `glsl/conv1d_dw.glsl` (same `l_in` map + fma loop + optional bias) + `impl/Convolution.cpp:755` (the `is_depthwise = groups==weight[0] && weight[1]==1` runtime dispatch to conv1d_dw). Vulkan dispatches conv1d_dw at runtime from `aten.convolution` (no AOT rewrite), and its `[K,C]` weight is a runtime `prepack_standard(kChannelsPacked)` step; the WebGPU buffer path skips that and reads the raw serialized `[C,1,K]` constant directly (`weight[c*K+k]`). A dynamic-shape resize hook recomputes `L_out = (L + 2p - dilation*(K-1) - 1)/stride + 1`. Fail-loud guards: config, 3D, stride>=1, fp32 in/out + weight `== C*K`.
ghstack-source-id: 406366796
@exported-using-ghexport

Differential Revision: [D112257664](https://our.internmc.facebook.com/intern/diff/D112257664/)
JCNTH added 25 commits July 24, 2026 01:05
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/)
@JCNTH
JCNTH merged commit 2eef94f into gh/JCNTH/116/orig Jul 24, 2026
26 of 34 checks passed
@JCNTH
JCNTH deleted the gh/JCNTH/117/orig branch July 24, 2026 08:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants