Skip to content

[ExecuTorch][WebGPU] Op-tests for pixel_shuffle#21322

Merged
JCNTH merged 56 commits into
gh/JCNTH/137/origfrom
gh/JCNTH/138/orig
Jul 24, 2026
Merged

[ExecuTorch][WebGPU] Op-tests for pixel_shuffle#21322
JCNTH merged 56 commits into
gh/JCNTH/137/origfrom
gh/JCNTH/138/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: #21178 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/138/base
ghstack PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/138/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/JCNTH/137/orig
Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/138/orig

@diff-train-skip-merge

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/)
@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/21322

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:01
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/)
Pull Request resolved: #21182

Adds the op-test suite for the depthwise conv1d config of `aten.convolution` ported in the diff below.

Key changes:
- `test/ops/test_conv1d_dw.py` — `Conv1dDWModule` (depthwise `nn.Conv1d(groups=C)`, seeded) + `Conv1dDWTest` delegation smoke test.
- `test/op_tests/cases.py` — `_conv1d_dw_suite`: k3s1p1, k3s2p1, dil2, k5_nobias.
ghstack-source-id: 406366798
@exported-using-ghexport

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

Adds the pointwise conv1d configuration to the `aten.convolution.default` handler (the depthwise config landed in the diff below). Pointwise conv1d (K=1, groups=1) is a per-position matmul over channels.

Key changes:
- `runtime/ops/conv1d_dw/conv1d_pw.wgsl` (+ generated `conv1d_pw_wgsl.h`) — per NCL output `(n, oc, l)`: `sum_ic weight[oc, ic] * input[n, ic, l] + bias[oc]`.
- `runtime/ops/conv1d_dw/Conv1dDW.cpp` — `add_conv1d_pw_node` helper + a pointwise dispatch branch in `convolution_impl`, checked BEFORE the depthwise branch (mirrors Vulkan `Convolution.cpp` is_pointwise-first). The one `aten.convolution.default` registration now dispatches depthwise vs pointwise, like Vulkan's single `Convolution.cpp`.

Mirrors Vulkan `glsl/conv1d_pw.glsl` (pointwise = channel matmul) + `impl/Convolution.cpp:755` (`is_pointwise = weight[2]==1`). Raw `[OC,IC,1]` weight (`oc*IC+ic`), same no-AOT-prepack basis as the depthwise diff. A dynamic-shape resize hook recomputes the length. No-bias binds the weight buffer as an unread placeholder.

Constraints: the pointwise branch is gated to `stride==1 && padding==0` (the standard pointwise); a strided/padded/grouped K=1 conv falls through to a fail-loud throw. This is stricter than Vulkan's `is_pointwise = weight[2]==1` alone (Vulkan's matmul shader has no stride/padding logic and would mis-handle those); the WebGPU handler fails loud instead.
ghstack-source-id: 406366804
@exported-using-ghexport

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

Adds the op-test suite for the pointwise conv1d config of `aten.convolution` ported in the diff below.

Key changes:
- `test/ops/test_conv1d_pw.py` — `Conv1dPwModule` (pointwise `nn.Conv1d(K=1)`, seeded) + `Conv1dPwTest` delegation smoke test.
- `test/op_tests/cases.py` — `_conv1d_pw_suite`: ic4_oc6, square, oc2_nobias, ic8_oc8, batch2.
ghstack-source-id: 406366805
@exported-using-ghexport

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

Problem: The WebGPU delegate is missing the pair-interleaved rotary embedding that the Vulkan delegate exposes as `et_vk.apply_rotary_emb_interleaved` (used by EdgeTAM). Without it these graphs fall back to CPU or fail to delegate.

Solution: Port the Vulkan `apply_rotary_emb_interleaved.glsl` to WGSL + a C++ handler, mirroring its math exactly. The last dim is `[r, i]` pairs and `freqs` is matching `[cos, sin]` pairs of shape `[seq, width]`; each output element applies the standard 2x2 rotation from its pair partner. Single dispatch, 2D-folded flat index to lift the 65535 1D-dispatch cap.

Implementation: `RotaryEmbeddingInterleaved.cpp` registers `et_vk.apply_rotary_emb_interleaved.default` (args `[x, freqs, out]`), guards rank-3 `[B, N, C]` / even width / fp32 / `freqs == [seq, width]` / `numel <= u32` / non-null buffers, and installs a resize hook that recomputes `seq`/`numel` + dispatch counts for dynamic shapes. `apply_rotary_emb_interleaved.wgsl` reads the `[cos, sin]` pair at the even-aligned column and writes the rotated element.
ghstack-source-id: 406366806
@exported-using-ghexport

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

Problem: The new `apply_rotary_emb_interleaved` WebGPU op needs golden coverage against a trusted reference and an export-delegation smoke test, matching the sibling rope ops.

Solution: Add an op-test suite that goldens the WebGPU kernel against the CPU eager `et_vk.apply_rotary_emb_interleaved` op (its own dtype-agnostic reference), plus an export test asserting the op is absorbed into the delegate rather than left as a CPU fallback.

Implementation: `cases.py` registers `_rope_interleaved_suite` over bnc/batch/c4/c16 shapes at atol=rtol=1e-3 with the default fp64 dual-oracle (fp32-eager vs fp64-golden gate). `test_rope_interleaved.py` calls the custom op directly (EdgeTAM has no aten lowering / fusion pattern for it) and verifies both `VulkanBackend` delegation and that no `interleaved` node is left top-level.
ghstack-source-id: 406366810
@exported-using-ghexport

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

Problem: The WebGPU delegate has no per-tensor quantization — the C0 gate for the int8 activation path that the Vulkan delegate's q8ta quantized ops build on. Without it, no int8-input/output op can be served.

Solution: Port `quantized_decomposed.quantize_per_tensor.default` (fp32 -> int8), the backend's first non-fp32 storage-buffer op. The int8 output is packed 4 elements per 32-bit word and bound as `array<u32>` (the landed `q4gsw` idiom), since WebGPU always allocates buffers and ignores the serialized `memory_layout`.

Implementation: `QuantizePerTensor.cpp` registers only `.default` (mirrors Vulkan), reads `scale`/`zero_point` at fixed arg indices with `out = args.back()` (robust to the overload's arg count), and guards int8 output / `numel % 4 == 0` (the `array<u32>` binding) / fp32 input / non-null buffers, all fail-loud. `quantize_per_tensor.wgsl` computes `round(x * inv_scale) + zero_point`, clamps to `[-128, 127]`, and packs; `inv_scale` is the reciprocal taken in double then cast to f32, bit-matching torch's `round(input * (1.0 / scale))`. `WebGPUTensor` gains an `is_int8` flag (distinguishes int8 from uint8/bool, which share a 1-byte size) so the handler can reject a non-int8 output. Mirrors Vulkan `runtime/graph/ops/glsl/q8ta_quantize.glsl` + `impl/QuantizeDequantize.cpp`.
ghstack-source-id: 406366809
@exported-using-ghexport

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

Problem: The WebGPU delegate has no per-tensor dequantization — the int8 -> fp32 half of the C0 quant gate, needed at the fp32 boundaries of a quantized graph.

Solution: Port `quantized_decomposed.dequantize_per_tensor.default` (int8 -> fp32). The int8 input is read as `array<u32>` (4 elements per word), each byte sign-extended, matching the layout `quantize_per_tensor` writes.

Implementation: `DequantizePerTensor.cpp` registers only `.default`, takes `out = args.back()` (robust to the extra `out_dtype` kwarg this overload carries), and guards int8 input (`is_int8`) / `numel % 4 == 0` / fp32 output / non-null buffers, fail-loud. `dequantize_per_tensor.wgsl` unpacks each byte, sign-extends via `(b ^ 0x80) - 128`, and computes `(q - zero_point) * scale`. Mirrors Vulkan `runtime/graph/ops/glsl/q8ta_dequantize.glsl` + `impl/QuantizeDequantize.cpp`.
ghstack-source-id: 406366819
@exported-using-ghexport

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

Problem: The op-test framework compares outputs as fp32 only, so an int8-output op (`quantize_per_tensor`, and the wider q8ta family) cannot be goldened — a round-trip through `dequantize` folds to identity under ET's QDQ-cancel and tests nothing.

Solution: Add an int8 golden path, keyed off the output tensor's dtype, alongside the existing fp32 path (which is byte-identical for every existing op). This mirrors the earlier multi-output framework extension.

Implementation: `generate_op_tests.py` detects an int8 output, writes the golden as raw int8 bytes, records `"dtype": "int8"` in the manifest, and skips the fp32 cast + dual-oracle gate (fp32/float64 goldens are untouched). `driver_util` parses the golden `dtype` (default `"float32"`, so old manifests are unchanged) and adds `load_int8_bin`. `op_test_driver.cpp` branches on the dtype: int8 goldens compare byte-exact (a discrete grid — any deviation is a real bug), fp32 goldens keep the `within_tol` path.
ghstack-source-id: 406366826
@exported-using-ghexport

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

Problem: The new `quantize_per_tensor` + `dequantize_per_tensor` ops need golden coverage. ET folds `dequantize(quantize(x))` back to `x` before delegation, so a round-trip would test a passthrough — each op must be tested alone.

Solution: `QuantizeModule` emits an int8 output goldened byte-exact against the torch int8 (via the int8-golden harness path). `DequantizeConstModule` dequantizes a baked int8 constant, verifying the dequantize stage independently against torch (breaking any compensating-bug blind spot a round-trip would have).

Implementation: `cases.py` registers a `quantize_per_tensor` suite (basic / nonzero zero-point / small-scale + a `ties` case whose `x * inv_scale` lands on half-integers, pinning ties-to-even rounding) and a `dequantize_per_tensor` suite (baked int8 edges spanning `-128/127/0/-1`). `test_quant.py` carries the export-delegation smoke tests.
ghstack-source-id: 406366824
@exported-using-ghexport

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

Problem: The WebGPU delegate has no quantized elementwise add — the first C1 q8ta op, on the int8 activation path C0 established.

Solution: Port `et_vk.q8ta_add` (two int8 tensors + per-tensor qparams -> int8): dequantize both, compute `a + alpha*b`, requantize. Reuses the C0 int8 buffer path (`array<u32>` pack/unpack, `is_int8` guard, multiply-by-`inv_scale`).

Implementation: `Q8taAdd.cpp` registers `et_vk.q8ta_add.default` (`out = args.back()`), guards all three tensors int8 + equal `numel` + `numel % 4 == 0`, fail-loud; `q8ta_add.wgsl` unpacks 4 int8/word, dequantizes `scale*(q - zero_point)`, forms `a + alpha*b`, and requantizes `clamp(round(v * inv_output_scale) + zero_point, -128, 127)`. Mirrors Vulkan `q8ta_binary.glsl` + `common.glslh`, with one deliberate divergence: the Vulkan glsl buffer kernel drops `alpha` (computes `a + b`), so this port implements the full `a + alpha*b` to match the CPU reference for any `alpha` (surfaced in the handler comment).

Constraints: same-shape only (no broadcast — the equal-`numel` guard rejects it, matching the `add` op baseline); `supports_resize=False` (no resize hook), mirroring the partitioner.
ghstack-source-id: 406366820
@exported-using-ghexport

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

Problem: The new `q8ta_add` int8 op needs golden coverage; its int8 output uses the int8-golden harness, and the `alpha` term (which the Vulkan glsl buffer path drops) must be pinned.

Solution: A `q8ta_add` suite goldens the kernel byte-exact against the CPU eager op with both operands as baked int8 constants, plus an export-delegation smoke test.

Implementation: `cases.py` registers `basic`, `alpha=2.0` (fails byte-exact if the `a + alpha*b` term is dropped), and `nonzero_zp` cases over int8 inputs spanning the sign-extend + requant-clamp edges (-128/127); `golden_dtype="float32"` (int8 result is exactly representable). `test_q8ta_add.py` carries the delegation smoke test.
ghstack-source-id: 406366821
@exported-using-ghexport

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

Problem: The WebGPU delegate has no quantized relu — a C1 q8ta activation on the int8 path C0 established.

Solution: Port `et_vk.q8ta_relu` (int8 -> int8): dequantize, `max(x, 0)`, requantize. Reuses the C0 int8 buffer path (`array<u32>` pack/unpack, `is_int8` guard, multiply-by-`inv_scale`), like `q8ta_add`.

Implementation: `Q8taRelu.cpp` registers `et_vk.q8ta_relu.default` (`out = args.back()`), guards in/out int8 + equal `numel` + `numel % 4 == 0`, fail-loud; `q8ta_relu.wgsl` computes `clamp(round(max(input_scale*(q - input_zero_point), 0) * inv_output_scale) + output_zero_point, -128, 127)`, packed 4 int8/word. Mirrors Vulkan `q8ta_unary.glsl` (`OPERATOR = max(X, 0)`) + `common.glslh`. Unlike `q8ta_add`, the partitioner sets `supports_resize=True`, so a resize hook recomputes `numel` + dispatch + rewrites the UBO on dynamic shapes (re-applying the `numel % 4` guard), mirroring the landed elementwise hooks.
ghstack-source-id: 406366825
@exported-using-ghexport

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

Problem: The new `q8ta_relu` int8 op needs golden coverage; the `max(x, 0)` relu term must be pinned.

Solution: A `q8ta_relu` suite goldens the kernel byte-exact against the CPU eager op with the int8 input as a baked constant, plus an export-delegation smoke test.

Implementation: `cases.py` registers `basic`, `diff_qparams` (input scale != output scale), and `nonzero_zp` cases over int8 inputs spanning the sign-extend edges (-128/127); values whose dequantized magnitude is negative are relu-clamped to 0, so a dropped `max(x, 0)` fails byte-exact. `golden_dtype="float32"` (int8 result is exactly representable). `test_q8ta_relu.py` carries the delegation smoke test.
ghstack-source-id: 406366822
@exported-using-ghexport

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

Problem: The WebGPU delegate has no quantized pixel_shuffle — the last C1 q8ta elementwise/shape op on the int8 path C0 established.

Solution: Port `et_vk.q8ta_pixel_shuffle` (int8 `[N, C*r*r, H, W]` -> int8 `[N, C, H*r, W*r]`): gather via the pixel_shuffle inverse, then dequantize -> requantize per element. Reuses the landed fp32 `pixel_shuffle` gather index math + the C0/`q8ta_add` int8 pack/unpack + requant.

Implementation: `Q8taPixelShuffle.cpp` registers `et_vk.q8ta_pixel_shuffle.default` (`out = args.back()`); one thread per packed output word gathers 4 int8 elements from arbitrary input positions (`c_in = c_out*r*r + (h_out%r)*r + (w_out%r)`, `h_in = h_out/r`, `w_in = w_out/r`), each `clamp(round(input_scale*(q - input_zero_point) * output_inv_scale) + output_zero_point, -128, 127)`. Guards int8 in/out + equal `numel` + `numel % 4 == 0` + `in_c % (r*r) == 0`, fail-loud; a resize hook (`supports_resize=True`) recomputes shape/numel/dispatch + rewrites the UBO. Mirrors Vulkan `q8ta_pixel_shuffle.glsl`.

Constraints / divergences: (1) the schema's 4th arg is `output_inv_scale` (already `1/output_scale`) — passed straight, NOT re-inverted (unlike `q8ta_relu`/`q8ta_add`). (2) Vulkan uses a channels-packed `int8x4` layout; this port uses the backend's flat NCHW buffer (always-buffer design) with a re-derived gather — logically identical for a per-tensor-quantized rearrange (byte-exact vs the eager op). (3) Vulkan's same-qparams passthrough fast-path is omitted; the general dequant->requant is byte-identical there (`round((q-zp)*s*(1/s))+zp == q`).
ghstack-source-id: 406366828
@exported-using-ghexport

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

Problem: The new `q8ta_pixel_shuffle` int8 op needs golden coverage of the gather + dequant->requant across qparam regimes and upscale factors.

Solution: A `q8ta_pixel_shuffle` suite goldens the kernel byte-exact against the CPU eager op with the int8 input a baked constant, plus an export-delegation smoke test.

Implementation: `cases.py` registers `basic`/`c2` (same qparams = pure gather), `r3` (upscale_factor=3), `diff_qparams` (input scale != output scale = dequant->requant rescale), and `nonzero_zp` (zero-point shift). `golden_dtype="float32"` (int8 result is exactly representable). `test_q8ta_pixel_shuffle.py` carries the delegation smoke test and passes `1/output_scale` as the 4th arg (schema `output_inv_scale`).
ghstack-source-id: 406366823
@exported-using-ghexport

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

Problem: The WebGPU delegate has no quantized linear — the core C2 GEMM. A plain `nn.Linear` through XNNPACK static PT2E lowers to a delegated `quantize_per_tensor -> q8ta_linear -> dequantize_per_tensor` subgraph; the quantize/dequantize are the landed C0 ops, so `q8ta_linear` is the one missing piece to run a quantized linear end-to-end.

Solution: Port `et_vk.q8ta_linear` (int8 activation x int8 per-channel weight -> int8), plus `et_vk.q8ta_linear_gemv` (identical schema, the M==1 case) on the same handler. Register-tiled i32 GEMM: `acc = Σ_k (x_int8[m,k] - input_zero_point) * weight_int8[n,k]`, dequantize `acc * input_scale * weight_scales[n] + bias`, requantize `clamp(round(v * inv_output_scale) + output_zero_point, -128, 127)`, packed 4 int8 per word. Mirrors Vulkan `q8ta_linear.glsl` + the landed q4gsw register tiling.

Implementation: `Q8taLinear.cpp` registers `q8ta_linear.default` + `q8ta_linear_gemv.default`, `out = args.back()`, guards int8 x/weight/out + `N % 4 == 0` (output word alignment) + `M*K`/`N*K` `% 4 == 0` (array<u32> binding), fail-loud. `activation` is read via a new `WebGPUGraph` string accessor (`strings_`/`get_string`, mirroring `ints_`/`doubles_`) and throws unless `"none"` — a relu-fused variant would otherwise silently emit non-relu output.

Constraints / divergences from the Vulkan reference (both numerically equivalent, byte-exact-verified): (1) the `weight_sums` arg is unused — the input zero-point correction is folded per-element (`Σ(x-zp)·w == Σx·w - zp·Σw`) instead of Vulkan's precomputed `weight_sums`. (2) x/weight are read row-major `[M,K]`/`[N,K]` (WebGPU's always-buffer convention; C0 `quantize_per_tensor` emits row-major int8 and no prepack runs), not Vulkan's 4H4W-packed/block-transposed layout — consistent with the landed q4gsw. `activation` scoped to `"none"` (the XNNPACK-static default; relu-fused is fail-loud, a follow-up).
ghstack-source-id: 406366829
@exported-using-ghexport

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

Problem: The new `q8ta_linear` int8 GEMM needs golden coverage of the full quantized-linear subgraph across GEMM/GEMV and bias regimes.

Solution: `make_q8ta_linear_module` runs a plain `nn.Linear` through XNNPACK static PT2E (per-channel weight, static per-tensor activation) in-process and returns the converted module, so the op-test framework goldens the WebGPU output against the converted eager (fp32, the fake-quant reference) e2e through `quantize_per_tensor -> q8ta_linear -> dequantize_per_tensor` (C0 + the new op) — no external model, no hand-written golden.

Implementation: `cases.py` registers `basic` (M=4), `gemv` (M=1 -> q8ta_linear_gemv), `k32`, and `no_bias` over N%4==0 shapes at atol=rtol=1e-3 (an off-by-one int8 level is 1.0, far above tol). `test_q8ta_linear.py` carries the delegation smoke test (asserts the `VulkanBackend` delegate contains `q8ta_linear`).
ghstack-source-id: 406366827
@exported-using-ghexport

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

Problem: The WebGPU delegate has no quantized pointwise (1x1) convolution. A plain 1x1 `nn.Conv2d` through XNNPACK static PT2E lowers to a delegated `quantize_per_tensor -> q8ta_conv2d_pw -> dequantize_per_tensor` subgraph; the quantize/dequantize are the landed C0 ops, so `q8ta_conv2d_pw` is the one missing piece to run a quantized 1x1 conv end-to-end.

Solution: Port `et_vk.q8ta_conv2d_pw` (int8 activation x int8 per-channel weight -> int8). A 1x1 conv is a per-(H,W)-position channel dot: `acc = Σ_ic (x_int8[n,ic,h,w] - input_zero_point) * weight_int8[oc,ic]`, dequantize `acc * input_scale * weight_scales[oc] + bias`, requantize `clamp(round(v * inv_output_scale) + output_zero_point, -128, 127)`, packed 4 int8 per word. Mirrors Vulkan `q8ta_conv2d_pw` (`Q8taConv2dPW.cpp` + `q8ta_conv2d_pw.glsl`).

Implementation: `Q8taConv2dPw.cpp` registers `q8ta_conv2d_pw.default`, `out = args.back()`, one thread per output word = 4 consecutive `W` positions of a fixed `(n, oc, h)` (`W % 4 == 0` for output word alignment), 2D dispatch fold to lift the 65535 cap. Guards int8 x/weight/out + `[N,IC,H,W]` / `[OC,IC]` / `[N,OC,H,W]` ranks, and scopes to `stride=1` / `pad=0` / `dilation=1` / `groups=1` / `activation="none"` (all fail-loud). `Q8taConvPwParams` is a 48-byte (16-aligned) uniform.

Constraints / divergences from the Vulkan reference (numerically equivalent, golden-verified): (1) the `weight_sums` arg is unused — the input zero-point correction is folded per-element (`Σ(x-zp)·w == Σx·w - zp·Σw`). (2) x/weight/out are read row-major NCHW / `[OC,IC]` (WebGPU's always-buffer convention; C0 `quantize_per_tensor` emits row-major int8 and the WebGPU prepack is a passthrough), not Vulkan's int8x4-block-packed 4W4C layout — consistent with the landed `q8ta_linear`.
ghstack-source-id: 406366831
@exported-using-ghexport

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

Problem: The new `q8ta_conv2d_pw` int8 pointwise conv needs golden coverage of the full quantized-conv subgraph across channel / bias / batch regimes.

Solution: `make_q8ta_conv2d_pw_module` runs a plain 1x1 `nn.Conv2d` through XNNPACK static PT2E (per-channel weight, static per-tensor activation) in-process and returns the converted module, so the op-test framework goldens the WebGPU output against the converted eager (fp32, the fake-quant reference) e2e through `quantize_per_tensor -> q8ta_conv2d_pw -> dequantize_per_tensor` (C0 + the new op) — no external model, no hand-written golden.

Implementation: `cases.py` registers `basic`, `ic8`, `no_bias`, and `batch2` (N=2) over `W % 4 == 0` shapes at atol=rtol=1e-3 (an off-by-one int8 level is 1.0, far above tol); calibration yields non-zero and negative activation zero-points, so the per-element zero-point-subtraction path is exercised. `test_q8ta_conv2d_pw.py` carries the delegation smoke test (asserts the `VulkanBackend` delegate contains `q8ta_conv2d_pw`).
ghstack-source-id: 406366834
@exported-using-ghexport

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

Problem: The WebGPU delegate has no quantized depthwise convolution. A depthwise `nn.Conv2d` (groups == channels) through XNNPACK static PT2E lowers to a delegated `quantize_per_tensor -> q8ta_conv2d_dw -> dequantize_per_tensor` subgraph; the quantize/dequantize are the landed C0 ops, so `q8ta_conv2d_dw` is the one missing piece to run a quantized depthwise conv end-to-end.

Solution: Port `et_vk.q8ta_conv2d_dw` (int8 activation x int8 per-channel weight -> int8). Depthwise means groups == C == OC: each output channel c convolves only input channel c with its own Kh x Kw filter. Per output element: `acc = Σ_{kh,kw} (x_int8[n,c,ih,iw] - input_zero_point) * weight_int8[kh,kw,c]` with `ih = oh*stride_h - pad_h + kh*dil_h`, `iw = ow*stride_w - pad_w + kw*dil_w` (out-of-bounds taps skipped = zero padding); dequantize `acc * input_scale * weight_scales[c] + bias`, requantize `clamp(round(v * inv_output_scale) + output_zero_point, -128, 127)`, packed 4 int8 per word. Mirrors Vulkan `q8ta_conv2d_dw` (`Q8taConv2dDW.cpp` + `q8ta_conv2d_dw.glsl`) and the landed `conv1d_dw` windowing.

Implementation: `Q8taConv2dDw.cpp` registers `q8ta_conv2d_dw.default`, `out = args.back()`, one thread per output word = 4 consecutive `W_out` positions of a fixed `(n, c, oh)` (`W_out % 4 == 0` for output word alignment), 2D dispatch fold to lift the 65535 cap. Reads full conv geometry (`stride`/`padding`/`dilation` int-lists, `groups`) and guards `groups == C == OC`, `[N,C,H_in,W_in]` / `[Kh,Kw,OC]` / `[N,C,H_out,W_out]` ranks, int8 dtypes, and `activation == "none"` (all fail-loud). `Q8taConvDwParams` is an 80-byte (16-aligned) uniform.

Constraints / divergences from the Vulkan reference (numerically equivalent, golden-verified): (1) the `weight_sums` arg is unused — the input zero-point correction is folded per-element (`Σ(x-zp)·w == Σx·w - zp·Σw`), so it matches Vulkan's `weight_sums` compensation exactly, including at zero-padded taps. (2) weight arrives raw `[Kh,Kw,OC]` int8 (the AOT pattern's depthwise reshape `permute(2,3,1,0)`; the WebGPU prepack is a passthrough), not Vulkan's int8x4-block 4W4C texture layout. `activation` scoped to `"none"` (the XNNPACK-static default; relu-fused is fail-loud, a follow-up).
ghstack-source-id: 406366835
@exported-using-ghexport

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

Problem: The new `q8ta_conv2d_dw` int8 depthwise conv needs golden coverage of the full quantized-conv subgraph across kernel-geometry, bias, stride, dilation, and batch regimes.

Solution: `make_q8ta_conv2d_dw_module` runs a depthwise `nn.Conv2d` (groups == channels) through XNNPACK static PT2E (per-channel weight, static per-tensor activation) in-process and returns the converted module, so the op-test framework goldens the WebGPU output against the converted eager (fp32, the fake-quant reference) e2e through `quantize_per_tensor -> q8ta_conv2d_dw -> dequantize_per_tensor` (C0 + the new op) — no external model, no hand-written golden.

Implementation: `cases.py` registers `k3` (3x3, stride 1, pad 1), `no_bias`, `stride2`, `dil2` (dilation 2 — exercises the `dil_h`/`dil_w` window taps), and `batch2` (N=2 — exercises the batch decomposition) over `C % 4 == 0` / `W_out % 4 == 0` shapes at atol=rtol=1e-3. XNNPACK-static calibration yields non-zero and negative activation zero-points, so the per-element zero-point-subtraction / requant path is exercised. `test_q8ta_conv2d_dw.py` carries the delegation smoke test (asserts the `VulkanBackend` delegate contains `q8ta_conv2d_dw`).
ghstack-source-id: 406366844
@exported-using-ghexport

Differential Revision: [D112257617](https://our.internmc.facebook.com/intern/diff/D112257617/)
JCNTH added 25 commits July 24, 2026 01:01
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 ca5f79b into gh/JCNTH/137/orig Jul 24, 2026
40 of 41 checks passed
@JCNTH
JCNTH deleted the gh/JCNTH/138/orig branch July 24, 2026 08:02
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