[ExecuTorch][WebGPU] Op-tests for q8ta_conv2d_pw#21344
Merged
Merged
Conversation
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/)
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21344
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
Pull Request resolved: #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/)
Pull Request resolved: #21203 Problem: The WebGPU delegate has no quantized general (ungrouped) convolution. A standard `nn.Conv2d` (groups == 1) through XNNPACK static PT2E lowers to a delegated `quantize_per_tensor -> q8ta_conv2d -> dequantize_per_tensor` subgraph; the quantize/dequantize are the landed C0 ops, so `q8ta_conv2d` is the one missing piece to run a quantized conv end-to-end. Together with the staged `q8ta_conv2d_pw` (1x1) and `q8ta_conv2d_dw` (depthwise), this covers the int8 conv2d family. Solution: Port `et_vk.q8ta_conv2d` (int8 activation x int8 per-channel weight -> int8) as a direct windowed convolution with full input-channel reduction. Per output element: `acc = Σ_{ic,kh,kw} (x_int8[n,ic,ih,iw] - input_zero_point) * weight_int8[oc, (kh*Kw+kw)*IC + ic]` 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[oc] + bias`, requantize `clamp(round(v * inv_output_scale) + output_zero_point, -128, 127)`, packed 4 int8 per word. Mirrors Vulkan `q8ta_conv2d` (`Q8taConv2d.cpp` + `q8ta_conv2d.glsl`, the direct general-shader path). Vulkan's `q8ta_conv2d` dispatches im2col-vs-general; the WebGPU port implements the direct-windowed general path only. Implementation: `Q8taConv2d.cpp` registers `q8ta_conv2d.default`, `out = args.back()`, one thread per output word = 4 consecutive `W_out` positions of a fixed `(n, oc, oh)` (`W_out % 4 == 0` for output word alignment), 2D dispatch fold to lift the 65535 cap. Reads the weight row stride from `weight.dims[1]` to tolerate the AOT's align-width padding of `Kh*Kw*IC`, and reads scales/bias as `>= [OC]` since the AOT pads `OC` to a multiple of 4 (the shader reads only `[0, OC)`). Guards `[N,IC,H_in,W_in]` / `[OC,Kh*Kw*IC]` / `[N,OC,H_out,W_out]` ranks, int8 dtypes, `groups == 1`, and `activation == "none"` (all fail-loud). `Q8taConvParams` is a 96-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`), matching Vulkan's `weight_sums` compensation exactly, including at zero-padded taps. (2) weight arrives raw `[OC, Kh*Kw*IC]` int8 (the AOT pattern's im2col reshape `permute(0,2,3,1)`; the WebGPU prepack is a passthrough), not Vulkan's int8x4-block texture layout. `groups > 1` and `activation != "none"` are fail-loud (follow-ups); `im2col` is not ported (the direct general path covers all shapes). ghstack-source-id: 406366853 @exported-using-ghexport Differential Revision: [D112257589](https://our.internmc.facebook.com/intern/diff/D112257589/)
Pull Request resolved: #21204 Problem: The new `q8ta_conv2d` int8 general conv needs golden coverage across channel-count (incl. OC/IC not multiples of 4), kernel-geometry, stride, dilation, and batch regimes. Solution: `make_q8ta_conv2d_module` runs a standard `nn.Conv2d` (groups == 1) 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 -> dequantize_per_tensor` (C0 + the new op) — no external model, no hand-written golden. Implementation: `cases.py` registers `k3` (3x3), `oc_gt` (IC=4, OC=8), `no_bias`, `stride2`, `dil2` (dilation 2), `oc6` (OC=6 — exercises the AOT `OC`-to-mult-4 scales/bias padding), `ic3` (IC=3 — exercises the `Kh*Kw*IC` weight-row align-width padding, the RGB first-conv case), `asym` (3x5 kernel — exercises the separate `Kh`/`Kw` window math), and `batch2` (N=2 — exercises the batch decomposition) over `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 / requant path is exercised. `test_q8ta_conv2d.py` carries the delegation smoke test (asserts the `VulkanBackend` delegate contains `q8ta_conv2d`). ghstack-source-id: 406366849 @exported-using-ghexport Differential Revision: [D112257676](https://our.internmc.facebook.com/intern/diff/D112257676/)
…onv) Pull Request resolved: #21205 Problem: The WebGPU delegate has no quantized transposed convolution (deconvolution). An `nn.ConvTranspose2d` through XNNPACK static PT2E lowers to a delegated `quantize_per_tensor -> q8ta_conv2d_transposed -> dequantize_per_tensor` subgraph; the quantize/dequantize are the landed C0 ops, so `q8ta_conv2d_transposed` is the one missing piece to run a quantized deconv end-to-end. This completes the int8 conv2d family (pointwise, depthwise, general, transposed). Solution: Port `et_vk.q8ta_conv2d_transposed` (int8 activation x int8 per-channel weight -> int8) as a direct gather. Per output element: `acc = Σ_{ic,kh,kw} (x_int8[n,ic,ih,iw] - input_zero_point) * weight_int8[oc, (kh*Kw+kw)*IC + ic]` where `ih = (oh + pad_h - kh*dil_h) / stride_h` is taken only when the numerator is non-negative and divisible by `stride_h` and `ih < H_in` (`iw` analogously); 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_transposed` (`Q8taConv2dTransposed.cpp` + `q8ta_conv2d_transposed.glsl`); reuses the general `q8ta_conv2d` int8 scaffold + `[OC, Kh*Kw*IC]` weight layout, changing only the index relation to the transposed gather. Implementation: `Q8taConv2dTransposed.cpp` registers `q8ta_conv2d_transposed.default`, `out = args.back()`, one thread per output word = 4 consecutive `W_out` positions of a fixed `(n, oc, oh)` (`W_out % 4 == 0`), 2D dispatch fold. The transposed schema inserts `output_padding` at arg 12, so `dilation` is read from arg 13 and `groups` from arg 14 (`output_padding` itself is unused — the output H/W come from the serialized output dims). The weight row stride is read from `weight.dims[1]` to tolerate align-width padding of `Kh*Kw*IC`; scales/bias are read as `>= [OC]` (the AOT pads `OC` to a multiple of 4; the shader reads only `[0, OC)`). Guards `[N,IC,H_in,W_in]` / `[OC,Kh*Kw*IC]` / `[N,OC,H_out,W_out]` ranks, int8 dtypes, `groups == 1`, `dilation == 1`, `stride >= 1`, and `activation == "none"` (all fail-loud). `Q8taConvTParams` is a 96-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`), matching Vulkan's `weight_sums` compensation exactly (Vulkan fills out-of-bounds taps with the zero-point, which contribute `(zp-zp)·w = 0` = the skipped taps here). (2) weight arrives raw `[OC, Kh*Kw*IC]` int8 (the AOT reshape maps `weight[oc,(kh*Kw+kw)*IC+ic] = original[ic,oc,kh,kw]`; the WebGPU prepack is a passthrough), not Vulkan's int8x4-block layout. `dilation != 1` (matching Vulkan's own restriction), `groups > 1`, and `activation != "none"` are fail-loud follow-ups. ghstack-source-id: 406366852 @exported-using-ghexport Differential Revision: [D112257642](https://our.internmc.facebook.com/intern/diff/D112257642/)
Pull Request resolved: #21206 Problem: The new `q8ta_conv2d_transposed` int8 transposed conv needs golden coverage across stride, kernel-geometry (incl. asymmetric), channel-count (incl. OC/IC not multiples of 4), bias, and batch regimes. Solution: `make_q8ta_conv2d_transposed_module` runs an `nn.ConvTranspose2d` (groups == 1, dilation == 1) 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_transposed -> dequantize_per_tensor` (C0 + the new op) — no external model, no hand-written golden. Implementation: `cases.py` registers `s2` (2x2, stride 2), `no_bias`, `k3` (3x3, stride 2, pad 1, output_padding 1), `oc6` (OC=6 — AOT `OC`-to-mult-4 scales/bias padding), `ic3` (IC=3 — `Kh*Kw*IC` weight-row align-width padding), `batch2` (N=2 — batch decomposition), and `asym` (3x2 kernel — the separate `Kh`/`Kw` gather math) over `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 / requant path is exercised. `test_q8ta_conv2d_transposed.py` carries the delegation smoke test (asserts the `VulkanBackend` delegate contains `q8ta_conv2d_transposed`). ghstack-source-id: 406366847 @exported-using-ghexport Differential Revision: [D112257602](https://our.internmc.facebook.com/intern/diff/D112257602/)
Pull Request resolved: #21207 Adds `div.Tensor_mode` (floor rounding) as `floor(in1 / in2)` with NumPy broadcasting, mirroring `mul` + Vulkan. Only `rounding_mode='floor'` is supported (fail-loud otherwise). Key changes: - `runtime/ops/floor_divide/{BinaryOp.cpp, binary_floor_divide.wgsl}` (+ generated `_wgsl.h`) — `floor_divide_impl` (args `[in1, in2, rounding_mode, out]`); broadcast `TensorMeta` path + dual-operand resize hook (mirrors `mul`); rounding_mode + int guards. - `CMakeLists.txt` — `runtime/ops/floor_divide/BinaryOp.cpp` in `WEBGPU_SRCS`. ghstack-source-id: 406366843 @exported-using-ghexport Differential Revision: [D112257613](https://our.internmc.facebook.com/intern/diff/D112257613/)
Pull Request resolved: #21208 Problem: The new `floor_divide` (`aten.div.Tensor_mode`) op needs golden coverage — and, because it is discontinuous, a correct oracle. Solution: `FloorDivideModule` runs `torch.div(a, b, rounding_mode="floor")` through the partitioner; the `floor_divide` suite goldens the WebGPU output against a `golden_fn` computing `floor(a / b)` in fp32 — the exact Vulkan glsl formula and the kernel's formula. This is deliberately NOT torch's eager `div.Tensor_mode` output: torch's `div_floor` is an fmod-corrected algorithm that can differ from `floor(a/b)` by 1 at fp-boundary quotients, so eager is the wrong oracle for a Vulkan-faithful `floor(a/b)` port (the "goldens = Vulkan" rule). Implementation: `cases.py` registers `2d` and `3d` same-shape cases with the divisor bounded away from zero (`_unary_lin(0.5, 4.0)`) and a widened dividend (`_unary_lin(-8, 8)`) for varied floor results; each case sets `golden_fn=_floor_div_golden`. `test_floor_divide.py` holds `FloorDivideModule` + the oracle rationale. ghstack-source-id: 406366848 @exported-using-ghexport Differential Revision: [D112257640](https://our.internmc.facebook.com/intern/diff/D112257640/)
Pull Request resolved: #21209 Problem: The WebGPU delegate has no `aten.argmax.default` / `aten.argmin.default`, and — more fundamentally — no int64-OUTPUT support: the AOT downcasts the int64 index to an int32 GPU buffer, but the ExecuTorch program output is int64, and `copy_outputs` raw-copied the int64 EValue's bytes from the int32 staging buffer (a size mismatch). argmax would be the first int64-output op. argmax unlocks on-GPU decode sampling. Solution: Port `et_vk.argmax`/`argmin` (last-dim arg-reduction -> int64 index) sharing one handler, mirroring Vulkan `ArgReduce.cpp` (`arg_reduce_impl`, last-dim only, the `add_reduce_per_row_node` accumulator that tracks `{val, idx}`). The kernel reuses the landed `amax` last-dim reduction + index-tracking: a strict `>` (argmax) / `<` (argmin) scan keeps the FIRST extremum (= torch tie-break), writing the index as int32 (1 u32/row) to the int32 GPU buffer. Add the int32->int64 output-widening path: `copy_outputs` now maps each output's LIVE staging size (`cur_nbytes`) and, when the host EValue is 2x the int staging buffer, sign-extends int32->int64 into the EValue; matched-dtype outputs keep the unchanged raw-copy path. Implementation: `argmax/Reduce.cpp` registers `aten.argmax.default` + `aten.argmin.default` -> `arg_reduce_impl(graph, args, is_argmin)`; args `[in, dim, keepdim, out]`, `dim` scalar (last-dim guard), `out = args.back()`, guards fp32 input / int32 output / shape, resize hook. `arg_reduce.wgsl` is one row per thread. `WebGPUGraph::copy_outputs` maps `cur_nbytes` + the guarded `dst == 2*map && is_int` widen. `WebGPUBackend::execute` wraps `execute()` + `copy_outputs()` in try/catch so a defensive throw never crosses the backend boundary. Constraints: last-dim reduction only (mirrors Vulkan `normalized_dim == ndim-1`); fp32 input, int32-backed int64 index output. The `copy_outputs` change is byte-identical for every existing fp32/int8 output (`cur_nbytes == EValue nbytes` for matched dtypes) — verified by regression (floor_divide fp32, q8ta convs int8 all still pass). ghstack-source-id: 406366846 @exported-using-ghexport Differential Revision: [D112257656](https://our.internmc.facebook.com/intern/diff/D112257656/)
Pull Request resolved: #21210 Problem: The new argmax/argmin ops output int64 indices, but the op-test harness was float/int8-only (it cast any int output to fp32, so the driver would read index bytes as float — a garbage compare). Solution: Extend the shared harness with an int64-golden path (mirroring the landed int8 path): the generator writes an int64 golden for int32/int64 raw outputs (`_write_int64`, dtype "int64"), and the driver reads `const_data_ptr<int64_t>()` and exact-compares to the int64 golden (`load_int64_bin`). The argmax/argmin suites use `golden_dtype="float32"` so the golden's fp32 argmax matches the fp32 kernel exactly (an fp64 oracle could flip a near-tie index). Implementation: `generate_op_tests.py` adds `_write_int64` + the `is_int64` branch. `op_test_driver.cpp` adds the int64 exact-compare branch; `driver_util.{h,cpp}` add `load_int64_bin`. `cases.py` registers `argmax`/`argmin` with randn 2d/3d cases (max/min at an INTERIOR index — exercises the reduction walk) and a deliberate-tie case per op (`argmax_tie_gen`/`argmin_tie_gen` place a repeated extremum at idx 1 and 3 so the FIRST occurrence wins — this discriminates the strict-`>`/`<` tie-break from a `>=`/`<=` bug). `test_argmax.py` holds the modules + tie gens. ghstack-source-id: 406366845 @exported-using-ghexport Differential Revision: [D112257677](https://our.internmc.facebook.com/intern/diff/D112257677/)
Pull Request resolved: #21211 Problem: The WebGPU delegate has no `et_vk.linear_qcs4w` — a 4-bit channels-symmetric-weight linear (per-output-channel symmetric weight, no zero-point). It is reachable via the `VulkanQuantizer` weight-only 4-bit path (distinct from the XNNPACK static-PT2E path that produces the q8ta ops), but had no WebGPU handler. Solution: Port `et_vk.linear_qcs4w` (fp32 activation, int4 weight), mirroring the landed `linear_q4gsw` register-tiled buffer GEMM (`quantized_linear/QuantizedLinear.cpp` + `q4gsw_linear.wgsl`) simplified to per-channel scale. Vulkan ref: `impl/QuantizedLinearQCSNW.cpp` `linear_qcs4w` (`check_linear_qcsnw_args`: args `[mat1, qmat2=[N,K/2] 4-bit, scales=[N], out]`, symmetric per-output-channel, no group, no zero-point, no bias) + `glsl/linear_qcsnw_coop.glsl`. Before/After vs q4gsw: q4gsw scale is grouped `scales[(k/group_size)*padded_N + n]`; qcs4w scale is per-channel `scales[n]` (1D [N]) — no group_size, no padded_N, no bias. Implementation: `linear_qcs4w/QuantizedLinearQcs4w.cpp` registers `et_vk.linear_qcs4w` -> `qcs4w_linear_impl`, args `[in, weight, scales, out]` (out=args.back()). `qcs4w_linear.wgsl`: register-tiled (TM=TN=4) GEMM, `acc += in[m,k] * (unpack_int4(w) - 8) * scales[n]`, 2D-folded dispatch (lifts the 65535 cap). `Qcs4wParams` (16 bytes: M/N/K/K_packed) matches the WGSL Params. Guards fp32 in/out, `K_packed==ceil(K/2)`, `N*K_packed%4==0` (u32-packed), `scales>=N`, all fail-loud. Resize hook recomputes live M + dispatch. Constraints / divergences from the Vulkan reference: (1) buffer re-derivation of Vulkan's texture-based qcs4w GEMM (WebGPU always buffers). (2) **CRITICAL — nibble order: the qcs4w AOT packer (`_passes/fuse_quantized_ops.py`) stores `(even_col<<4)|odd_col`, the SWAP of q4gsw's `pack_4bit_weight_tensor` `(odd<<4)|even`** — so this kernel reads even-k from the HIGH nibble and odd-k from the LOW nibble, the reverse of `q4gsw_linear.wgsl` (verified against the packer, the Vulkan `linear_qcsnw_coop.glsl` unpack, and a byte-search of the served `.pte`). (3) `+8`-shifted int4 recovered as signed `[-8,7]` (same as q4gsw). (4) one tiled kernel only (q4gsw's GEMV/shmem perf variants are separate, Canary-gated follow-ups). Bias is out of the op (a Linear bias lowers to a separate `aten.add`). ghstack-source-id: 406366854 @exported-using-ghexport Differential Revision: [D112257623](https://our.internmc.facebook.com/intern/diff/D112257623/)
Pull Request resolved: #21212 Problem: The new `et_vk.linear_qcs4w` op needs golden coverage — and, being a quantized op reachable only through a specific quantizer, a test that actually produces the op and validates the fake-quant numerics. Solution: `make_qcs4w_linear_module` runs a plain `nn.Linear` through the `VulkanQuantizer` weight-only 4-bit path (`get_symmetric_quantization_config(is_dynamic=False, weight_bits=4)` -> prepare_pt2e -> calibrate -> convert_pt2e), and the op-test `module_factory` returns the CONVERTED module, so the WebGPU output is goldened against the converted eager (fp32 per-channel fake-quant reference). This exercises the int4 unpack + per-channel dequant end-to-end. `bias=False` keeps the golden focused on qcs4w (a Linear bias lowers to a separate `aten.add`). Implementation: `cases.py` registers `linear_qcs4w` with `basic` (4x32x16), `gemv` (M=1 decode shape), `k64` (2x64x8), `n32` (3x32x32) — all with K even (2 nibbles/byte) and `N*ceil(K/2) % 4 == 0` (u32-packed weight); `test_linear_qcs4w.py` holds the module + a delegation smoke test asserting `et_vk.linear_qcs4w` is absorbed into the VulkanBackend delegate. ghstack-source-id: 406366862 @exported-using-ghexport Differential Revision: [D112257668](https://our.internmc.facebook.com/intern/diff/D112257668/)
Pull Request resolved: #21213 Problem: The WebGPU delegate has no `et_vk.linear_q8ta_q8csw` — an int8-activation x int8-channelwise-weight linear with a FP32 output (the "output-not-requantized" sibling of the landed `q8ta_linear`, e.g. a final projection). Its kernel design was ready but it was thought unreachable; the reachability blocker is now solved. Solution: Port `et_vk.linear_q8ta_q8csw` (int8 x, int8 per-channel weight, fp32 out), mirroring the landed `q8ta_linear` (`q8ta_linear/Q8taLinear.cpp`) with the output requant removed. Before/After vs `q8ta_linear`: same i32 GEMM `Σ(x_int8 - input_zp) * w_int8` + dequant `* input_scale * weight_scales[n] + bias`, but `q8ta_linear` then requantizes to int8 (needs output_scale/zp, N%4==0 for output packing) whereas this writes fp32 directly (no output qparams in the schema, N unconstrained). Mirrors Vulkan `impl/QuantizedLinear.cpp` `linear_q8ta_q8csw`. Implementation: `linear_q8ta_q8csw/LinearQ8taQ8csw.cpp` registers `et_vk.linear_q8ta_q8csw` -> `linear_q8ta_q8csw_impl`. Schema `[x, input_scale, input_zp, weights, weight_sums, weight_scales, bias?, out]`: x/weights int8 (bound `array<u32>`, unpacked in-shader), `weight_sums` (arg 4) folded per-element (unused, matching q8ta_linear), `out=args.back()`. `linear_q8ta_q8csw.wgsl` register-tiled (TM=TN=4) i32 GEMM, 2D-folded dispatch; the TN=4 tile guards `n<N` (no N%4 assumption — fp32 output has no packing constraint). `Q8taQ8cswParams` (32 bytes) matches the WGSL Params. Guards int8 x/weight (numel%4==0 for the u32 reads), fp32 out, scales fp32 [N], all fail-loud. Constraints / divergences from the Vulkan reference: fp32 output only. A bias-less TERMINAL linear can mis-fuse to an int8 output the schema cannot compute (no output scale/zp) — the handler fail-louds on it (`output must be fp32`); the correct fp32-out form is what the fusion routes when the output is not re-quantized. i32 accumulator (inherited from q8ta_linear). Buffer-only, so no int8x4 texture layout. ghstack-source-id: 406366865 @exported-using-ghexport Differential Revision: [D112257666](https://our.internmc.facebook.com/intern/diff/D112257666/)
Pull Request resolved: #21214 Problem: The new `et_vk.linear_q8ta_q8csw` op needs golden coverage, and the recipe to REACH it (a quantized linear whose output stays fp32) is non-obvious. Solution: `make_linear_q8ta_q8csw_module` runs a plain `nn.Linear` through XNNPACK static PT2E with the activation config's `output_activation` nulled (`dataclasses.replace(get_symmetric_quantization_config(is_per_channel=True, is_dynamic=False), output_activation=None)`): the linear's INPUT is statically per-tensor quantized but its OUTPUT is left fp32, so the Vulkan fusion routes to `linear_q8ta_q8csw` (fp32 out) instead of `q8ta_linear` (int8 out). The op-test `module_factory` returns the CONVERTED module, so the WebGPU output is goldened against the converted eager (fp32 fake-quant reference); the served subgraph is `quantize_per_tensor` (landed) -> `linear_q8ta_q8csw`. Cases use `bias=True` (a bias-less terminal linear mis-fuses to an int8 output the op fail-louds on) and N a multiple of 4 (the AOT pads the quantized weight's N). Implementation: `cases.py` registers `linear_q8ta_q8csw` with `basic` (4x32x16), `gemv` (M=1), `k48` (2x48x8), `n32` (3x32x32); `test_linear_q8ta_q8csw.py` holds the module + a delegation smoke test asserting `et_vk.linear_q8ta_q8csw` is absorbed into the VulkanBackend delegate. ghstack-source-id: 406366863 @exported-using-ghexport Differential Revision: [D112257659](https://our.internmc.facebook.com/intern/diff/D112257659/)
Pull Request resolved: #21215 Problem: The WebGPU delegate has no `et_vk.grid_priors` — a detection anchor-grid op (generates per-cell coordinate shifts). It delegates through `VulkanPartitioner` (`op_registry.py` tags `et_vk.grid_priors.default`; reachable via a direct custom-op call) but had no WebGPU handler. Solution: Port `et_vk.grid_priors` (fp32), mirroring Vulkan `impl/GridPriors.cpp` + `glsl/grid_priors.glsl`. The output is `[H*W, 2]` computed purely from the input's H/W plus scalar `stride`/`offset` — the input's VALUES are not read. Per output row `r`: `out[r,0] = (r%W + offset) * stride` (x-shift), `out[r,1] = (r/W + offset) * stride` (y-shift), matching the Vulkan glsl (`pos.x==0 -> pos.y%width`, `pos.x==1 -> pos.y/width`) and the CPU eager (`custom_ops_lib.py` grid_priors: `stack((shift_x_per_col, shift_y_per_row))`). Implementation: `grid_priors/GridPriors.cpp` registers `et_vk.grid_priors.default` -> `grid_priors_impl`. Args `[in, stride(int), offset(float), out]`; only `in.dims[-2:]` (H, W) are used. `grid_priors.wgsl`: one thread per output element, `row=idx/2`, even idx -> `row%W`, odd -> `row/W`, `(coord+offset)*stride`; 2D-folded dispatch. `GridPriorsParams` (16 bytes: numel/width/stride/offset) matches the WGSL Params. Only the output (storage) + params (uniform) are bound — the input's data is intentionally unread. Resize hook recomputes numel/width + dispatch + out dims `[H*W, 2]`. Guards: in >= 2D, W>0, numel<=u32, fp32 out, all fail-loud. Constraints / divergences from the Vulkan reference: buffer re-derivation of Vulkan's texture (`imageStore`/`VEC4_T`) kernel — the flat-index buffer form is forced by the backend's buffer-only design. `stride` (int) is promoted to f32 for the multiply (exact for detection-scale strides; matches the glsl's int->float promotion). ghstack-source-id: 406366861 @exported-using-ghexport Differential Revision: [D112257635](https://our.internmc.facebook.com/intern/diff/D112257635/)
Pull Request resolved: #21216 Problem: The new `et_vk.grid_priors` op needs golden coverage — and, since the input's values are unused (only its H/W matter), a test that validates the computed grid + the column ordering across shapes. Solution: `GridPriorsModule` calls the custom op directly with baked `stride`/`offset`, so only the float tensor `x` is a runtime input. The op has a CPU eager impl (an independent meshgrid/stack reference), so the op-test framework goldens the WebGPU output against it directly (float32 oracle — the output is an exact computed grid). Non-square shapes (8x10, 5x7) make the two coordinate ranges differ, so a swapped `r%W`/`r/W` column mapping would diverge from the golden. Implementation: `cases.py` registers `grid_priors` with `s8` (8x10, stride=8, offset=0.5), `s16` (4x4, stride=16, offset=0), `offset0` (5x7, stride=4, offset=0); `test_grid_priors.py` holds the module + a delegation smoke test asserting `et_vk.grid_priors` is absorbed into the VulkanBackend delegate. ghstack-source-id: 406366867 @exported-using-ghexport Differential Revision: [D112257599](https://our.internmc.facebook.com/intern/diff/D112257599/)
Pull Request resolved: #21217 Problem: The WebGPU delegate has no `et_vk.conv_with_clamp` — a fused fp32 2D convolution + output clamp (e.g. a `Conv2d` followed by `relu6`). It delegates through `VulkanPartitioner` (`op_registry.py` tags `et_vk.conv_with_clamp.default`; `nn.Conv2d`+`F.relu6` fuses to it) but had no WebGPU handler. Solution: Port `et_vk.conv_with_clamp` (fp32, groups==1), mirroring the Vulkan `conv` handler (`impl/Convolution.cpp:837`) + the eager `clamp(convolution(...), output_min, output_max)` (`custom_ops_lib.py`). Direct windowed conv reusing the staged `q8ta_conv2d` windowing structure but fp32 (no int8 unpack/dequant/requant) + a clamp epilogue: per NCHW output cell `out[n,oc,oh,ow] = clamp(bias[oc] + Σ_{ic,kh,kw} weight[oc,ic,kh,kw] * input[n,ic, oh*sh-ph+kh*dh, ow*sw-pw+kw*dw], output_min, output_max)`, OOB taps skipped. Implementation: `conv_with_clamp/ConvWithClamp.cpp` registers `et_vk.conv_with_clamp.default` -> `conv_with_clamp_impl`. Args `[in, weight, bias?, stride, padding, dilation, transposed, output_padding, groups, output_min?, output_max?, out]`; stride/padding/dilation as int-pairs, `output_min`/`output_max` as `Scalar?` (Double/Int -> value, absent -> -/+inf so `clamp` is a passthrough). `conv_with_clamp.wgsl`: one thread per output element, RAW `[OC,IC,Kh,Kw]` weight (the fp32 buffer path gets the raw constant, like conv1d_dw/pw — not the q8ta im2col `[OC,Kh*Kw*IC]`), 2D-folded dispatch. `ConvWithClampParams` (80 bytes) matches the WGSL Params. Guards: in/weight/out 4D + fp32, `weight.dims[1]==IC`, numel<=u32, groups==1 + not-transposed, all fail-loud. Constraints / divergences from the Vulkan reference: buffer re-derivation of Vulkan's texture conv (`conv2d_prepack_weights.glsl`) — the naive one-thread-per-output form is forced by the buffer-only backend. Scoped to `groups==1` + not-transposed (fail-loud otherwise). Note: overlaps the general `convolution` op — this handles the clamped variant only. ghstack-source-id: 406366872 @exported-using-ghexport Differential Revision: [D112257644](https://our.internmc.facebook.com/intern/diff/D112257644/)
Pull Request resolved: #21218 Problem: The new `et_vk.conv_with_clamp` op (a first-of-its-kind general fp32 conv2d) needs golden coverage that exercises the H/W axis separation, not just square/symmetric configs. Solution: `ConvWithClampModule` is a plain `nn.Conv2d` + `F.relu6`, which the Vulkan fusion rewrites to `et_vk.conv_with_clamp` (fp32 conv + clamp[0,6]); the conv weight/bias are baked params and only `x` is a runtime input. Goldened vs the module's fp32 eager. The `asym` case is fully axis-asymmetric (input H=7≠W=9, kernel Kh=2≠Kw=3, stride 1≠2, padding 1≠0, dilation 2≠1 → output H_out=7≠W_out=4), so an H↔W index swap anywhere (unravel divisors, stride/pad/dilation axis, or Kh/Kw) diverges from the golden. Implementation: `cases.py` registers `conv_with_clamp` with `k3p1`/`stride2`/`dil2`/`no_bias`/`asym` (all groups==1); `test_conv_with_clamp.py` holds the module + a delegation smoke test asserting `et_vk.conv_with_clamp` is absorbed into the VulkanBackend delegate. ghstack-source-id: 406366873 @exported-using-ghexport Differential Revision: [D112257609](https://our.internmc.facebook.com/intern/diff/D112257609/)
…> bool) Pull Request resolved: #21219 Problem: The WebGPU delegate has no elementwise comparison ops (`aten.eq/lt/le/gt/ge.Tensor`). They delegate through `VulkanPartitioner` (all 5 tagged, verified) and output BOOL tensors, but the backend had no handler and no bool-output path. Solution: Port all 5 comparisons via one shared handler + WGSL kernel (op-code switch), mirroring Vulkan `impl/BinaryOp.cpp` + `glsl/binary_op_buffer.yaml`. Output is a 1-byte bool tensor (`vk_datatype_size(BOOL)=1`, `is_int=true`, `is_int8=false`); the kernel packs 4 bool bytes per `u32` word (reusing the int8 buffer-binding idiom), and `copy_outputs` raw-copies it (dst==map, no widen). Per element: `out = (a OP b) ? 1 : 0` with OP in {==, <, <=, >, >=}. Implementation: `compare/Compare.cpp` registers `aten.eq/lt/le/gt/ge.Tensor` -> `compare_impl(graph, args, op)` (op 0=eq..4=ge). `compare.wgsl`: one thread per output word (4 bool bytes), 2D-folded dispatch, op switch. `CompareParams` (16 bytes: num_elements, op) matches the WGSL Params. Guards: fp32 inputs, 1-byte bool output, same-shape (in/out numel equal), `numel%4==0` (bool packs 4/word AND gates the readback map) — all fail-loud. Resize hook recomputes numel/dispatch + re-applies the %4 guard + cross-checks both operands. Constraints / divergences from the Vulkan reference: (1) `numel%4==0` (bool-output packing; fail-loud otherwise, mirroring the q8ta N%4 output-pack constraint). (2) same-shape only (flat kernel; broadcast is export-smoke, like `add`/`minimum`). (3) `eq` is torch-EXACT `a==b` — a DELIBERATE deviation from Vulkan's float `abs(X-Y)<1e-5`: the delegate serves torch-exact model semantics and the op-test golden is torch eager; `lt/le/gt/ge` mirror Vulkan's exact comparators. ghstack-source-id: 406366878 @exported-using-ghexport Differential Revision: [D112257588](https://our.internmc.facebook.com/intern/diff/D112257588/)
…rness Pull Request resolved: #21220 Problem: The new comparison ops (`aten.eq/lt/le/gt/ge.Tensor`) output BOOL, which the op-test golden harness did not handle (only fp32/int8/int64). And the goldens must actually discriminate the comparators. Solution: Extend the golden harness with a `bool` branch (mirrors the int8/int64 branches): `generate_op_tests` writes the bool golden as 0/1 bytes (`raw.dtype==torch.bool -> _write_int8(raw.to(int8))`, dtype "bool") and `op_test_driver` reads the bool output via `const_data_ptr<bool>()` and byte-compares. The change is additive (bool is a distinct dtype; the branch precedes int8) — existing fp32/int8/int64 paths are untouched (regression 27/27). `CompareModule(op)` covers all 5 ops; the two inputs draw from DIFFERENT discrete-range seeds (`compare_gen_a`/`compare_gen_b`) so `a!=b` with frequent collisions — giving `eq/le/ge` genuine ties and `lt/gt` a real true/false mix, so an op-switch swap diverges from the golden. Implementation: `cases.py` registers `eq`/`lt`/`le`/`gt`/`ge` via a shared `_compare_suite` (2d/3d/sq shapes, all numel%4==0); `test_compare.py` holds `CompareModule` + the seeded gens + a delegation smoke test. ghstack-source-id: 406388747 @exported-using-ghexport Differential Revision: [D112257586](https://our.internmc.facebook.com/intern/diff/D112257586/)
Pull Request resolved: #21221 Adds the elementwise bool `aten.logical_and.default` to the WebGPU delegate. Vulkan maps this op to its `bitwise_and` handler (`BinaryOp.cpp:161`), valid for canonical 0/1 bool bytes; this port mirrors that. Both operands and the output are 1-byte bool tensors packed 4/word into `array<u32>`, so a word-wise AND is exactly a per-byte AND. Key changes: - `runtime/ops/logical_and/logical_and.wgsl` — one thread per packed word, `t_out[w] = t_a[w] & t_b[w]` (2D-folded dispatch). - `runtime/ops/logical_and/LogicalAnd.cpp` — handler binds out (rw) + a/b (ro) + params (uniform); fail-loud guards on bool dtype, null buffer, and same-shape; resize hook re-applies `num_words` + dispatch. - `CMakeLists.txt` — register the op source. Same-shape only (`numel % 4 == 0` for the packed bool bindings), matching the backend's other bool ops. Co-authored-with: Claude Code. ghstack-source-id: 406388750 @exported-using-ghexport Differential Revision: [D112257653](https://our.internmc.facebook.com/intern/diff/D112257653/)
Pull Request resolved: #21222 Adds the `aten.logical_and.default` op-test suite to the manifest-driven WebGPU op-test framework, stacked above the op diff. Key changes: - `test/ops/test_logical_and.py` — `LogicalAndModule` derives its two bool operands on-GPU from float inputs (`a > 0`, `b > 0` via the delegated `gt.Tensor` against a baked zero buffer), so the only runtime inputs are the two float tensors (the op-test framework is float-input-only). Distinct `a`/`b` seeds make the two masks differ (each ~50% True, independent -> AND ~25% True), a real mix that an OR mutant would fail. `LogicalAndTest` is the export-delegation smoke test. - `test/op_tests/cases.py` — registers `logical_and` (shapes 2d/3d/sq, all `numel % 4 == 0`), byte-exact bool golden. Co-authored-with: Claude Code. ghstack-source-id: 406388749 @exported-using-ghexport Differential Revision: [D112257583](https://our.internmc.facebook.com/intern/diff/D112257583/)
Pull Request resolved: #21223 Adds the two bool bitwise ops the Vulkan partitioner tags (`op_registry.py` marks `bitwise_and.Tensor`, `bitwise_not.default`, `logical_and.default` all `BOOL_T`). `bitwise_and` on bool is identical to `logical_and`, so it shares the one AND handler — mirroring Vulkan, which registers both `aten.bitwise_and.Tensor` and `aten.logical_and.default` to its `bitwise_and` handler (`BinaryOp.cpp:160-161`). No new kernel needed. `bitwise_not` is the bool NOT: Vulkan uses `1 - X` on uint8 (`unary_op.yaml`); for canonical 0/1 that equals a per-byte low-bit flip, so the packed-word kernel is `t_out[w] = t_a[w] ^ 0x01010101u` (one thread per word). Key changes: - `runtime/ops/logical_and/LogicalAnd.cpp` — register `aten.bitwise_and.Tensor` on the existing shared bool-AND handler. - `runtime/ops/bitwise_not/{BitwiseNot.cpp,bitwise_not.wgsl}` — new unary; out (rw) + a (ro) + params (uniform); fail-loud bool-dtype/null/same-shape guards; resize hook re-applies `num_words`. - `CMakeLists.txt` — register `BitwiseNot.cpp`. Bool only (`numel % 4 == 0` for the packed bindings), matching the backend's other bool ops. Co-authored-with: Claude Code. ghstack-source-id: 406388751 @exported-using-ghexport Differential Revision: [D112257647](https://our.internmc.facebook.com/intern/diff/D112257647/)
Pull Request resolved: #21224 Adds the `aten.bitwise_and.Tensor` / `aten.bitwise_not.default` (bool) op-test suites, stacked above the op diff. Key changes: - `test/ops/test_bitwise.py` — `BitwiseAndModule`/`BitwiseNotModule` derive their bool operands on-GPU from float inputs (`a > 0` via the delegated `gt.Tensor` against a baked zero buffer), so the only runtime inputs are float tensors (the op-test framework is float-input-only). `bitwise_and` uses distinct `a`/`b` seeds (AND ~25% True); `bitwise_not` inverts one ~50% mask. - `test/op_tests/cases.py` — registers `bitwise_and`/`bitwise_not` (shapes 2d/3d/sq, all `numel % 4 == 0`), byte-exact bool goldens. Co-authored-with: Claude Code. ghstack-source-id: 406388753 @exported-using-ghexport Differential Revision: [D112257674](https://our.internmc.facebook.com/intern/diff/D112257674/)
… shared-memory Pull Request resolved: #21225 Problem: The reduction ops (amax/amin/argmax/sum/mean), softmax/log_softmax, and native_group_norm's reduce pass used naive one-thread-per-row serial loops — far below Vulkan's optimization level, where backends/vulkan/runtime/graph/ops/glsl/reduce.glsl uses a cooperative shared-memory reduction. Solution: Rewrite each to Vulkan's cooperative design — one workgroup per reduction row/group, each thread reduces a strided slice into a workgroup shared array, thread 0 aggregates the partials. Dispatch changed to one workgroup per row. argmax preserves torch's first-index tie-break; softmax keeps numerical stability + the log variant. ghstack-source-id: 406388754 @exported-using-ghexport Differential Revision: [D112257649](https://our.internmc.facebook.com/intern/diff/D112257649/)
…r, _to_copy) + enable unary activation family Pull Request resolved: #21226 Adds the creation + cast/copy ops and enables the unary activation family in the op allowlist. Key changes: - `runtime/ops/full/{Full.cpp, full.wgsl}` — `full`/`full_like`/`zeros`/`zeros_like`/`ones`/`ones_like`/`scalar_tensor` (shared fp32 fill; constant-folded before the delegate in practice, so effectively register-only) + a fail-loud non-fp32 output guard. - `runtime/ops/to_copy/{ToCopy.cpp, to_copy.wgsl}` — `_to_copy` / `_to_dim_order_copy`: same-dtype raw copy, and int<->float CONVERT via `f32(bitcast<i32>(x))` / `bitcast<f32>(i32(x))`. The dtype-promotion pass inserts an int->float `_to_copy` before binary ops, so a byte-reinterpret would ship denormal/inf garbage; this converts. Mirrors Vulkan `ToCopy.cpp` (view_convert vs blit branch). - `test/tester.py` — enable the unary activation family. ghstack-source-id: 406388756 @exported-using-ghexport Differential Revision: [D112257603](https://our.internmc.facebook.com/intern/diff/D112257603/)
…shape/movement, pooling, conv/sampling) Pull Request resolved: #21227 Problem: These ops already had kernels in the granular tree but were absent from WEBGPU_SUPPORTED_OPS, so the op-test framework neither delegated nor golden-tested them. Solution: Add to the allowlist: minimum, div.Tensor_mode (floor_divide), logical_and, bitwise_and, bitwise_not, flip, repeat, index_select, avg_pool2d, pixel_shuffle, convolution (depthwise conv1d), conv_with_clamp, grid_sampler_2d, grid_priors. ghstack-source-id: 406388757 @exported-using-ghexport Differential Revision: [D112257605](https://our.internmc.facebook.com/intern/diff/D112257605/)
…+ qcs4w/q8ta_q8csw linear) Pull Request resolved: #21228 Problem: The int8 quantized ops (q8ta elementwise/conv/linear + qcs4w and q8ta_q8csw linear) already had kernels in the granular tree but were absent from WEBGPU_SUPPORTED_OPS, so they were neither delegated nor golden-tested. Solution: Add to the allowlist: et_vk.linear_qcs4w, et_vk.linear_q8ta_q8csw, et_vk.q8ta_add, et_vk.q8ta_relu, et_vk.q8ta_pixel_shuffle, et_vk.q8ta_linear (+ q8ta_linear_gemv), et_vk.q8ta_conv2d, et_vk.q8ta_conv2d_dw, et_vk.q8ta_conv2d_pw. ghstack-source-id: 406388761 @exported-using-ghexport Differential Revision: [D112257636](https://our.internmc.facebook.com/intern/diff/D112257636/)
…8da4w) Pull Request resolved: #21229 Problem: WebGPU lacked the dynamic-8bit-activation x 4-bit-group-weight linear (the 8da4w path real LLMs use for dynamic quant). Vulkan registers `et_vk.linear_dq8ca_q4gsw` + `torchao.choose_qparams_affine`; WebGPU delegated neither, so any `Int8DynamicActivationIntxWeightConfig` model fell back to CPU. Solution: Author the reachable pair. `torchao.choose_qparams_affine` computes per-row asymmetric int8 activation scale/zp; `et_vk.linear_dq8ca_q4gsw` folds that dynamic activation quant into the existing q4gsw 4-bit-group GEMM. Both mirror the Vulkan reference (`ChooseQParams.cpp`, `QuantizedLinear.cpp:760`); the activation-quant + weight-dequant math was CPU-de-risked exact against torchao before authoring. Impl: `choose_qparams_affine.wgsl` does a cooperative per-row min/max reduction (one workgroup per block of 4 rows so the 4 int8 zps pack into one u32 with no write race) then `calculate_scale_and_zero_point` mirroring Vulkan. `linear_dq8ca_q4gsw.wgsl` is the register-tiled q4gsw GEMM with `out[m,n] = s[m] * sum_k (xq[m,k]-z[m]) * dequant(w)`, reading the packed-int8 zp. int8 zp (`elem_size` 1) is bound word-aligned; `num_rows` must be <=4 or a multiple of 4 (int8 buffers alloc `max(nbytes,4)`) — arbitrary-M prefill is a documented follow-up. ghstack-source-id: 406388760 @exported-using-ghexport Differential Revision: [D112257680](https://our.internmc.facebook.com/intern/diff/D112257680/)
Pull Request resolved: #21230 Add `aten.bitwise_or.Tensor` + `aten.logical_or.default` to the WebGPU backend — the OR counterpart to the existing `bitwise_and`/`logical_and` family, closing the last bool-binary parity gap with Vulkan. Both overloads share one handler (`logical_or_op`); on canonical 0/1 bool bytes the word-wise `|` over the packed `u32` buffer equals a per-byte OR. Kernel-only: the Vulkan partitioner already tags both ops (`op_registry.py` `register_bool_binary_ops`), so no partitioner change. Key changes: - `runtime/ops/logical_or/` — `logical_or.wgsl` (`t_out = t_a | t_b`, bool packed 4/word) + `LogicalOr.cpp` (one handler; bool-only + `numel % 4` guards + dynamic-shape resize hook; registers both overloads). Mirrors `logical_and`. - `test/ops/test_logical_or.py` + `op_tests/cases.py` — delegation smoke + op-test golden for both ops (two ~50% masks -> OR ~75% True, which distinguishes OR from an AND mutant). - `CMakeLists.txt` — wire `logical_or/LogicalOr.cpp`. Co-authored-with: Claude Code. ghstack-source-id: 406388764 @exported-using-ghexport Differential Revision: [D112483463](https://our.internmc.facebook.com/intern/diff/D112483463/)
… (0x30) Pull Request resolved: #21231 Rank-5+ tensors threw `0x30` (`DelegateInvalidCompatibility`) at delegate init: the shared `TensorMeta` UBO capped rank at 4 (`kTensorMetaMaxNdim`), so `fill_tensor_meta`/`fill_tensor_meta_broadcast` hard-throw during `build()` and the WGSL mirror stored `sizes`/`strides` as a single `vec4<u32>`. This clears `test_permute_different_shapes[webgpu]`, `test_transpose_different_shapes[webgpu]` (transpose lowers to `permute_copy`), and the rank>4 model failures (`conformer`, `maxvit_t`). Raise the global cap to 8 and widen the uniform to two `vec4<u32>` blocks, keeping the C++ and WGSL layouts byte-identical: - C++ `sizes[8]`/`strides[8]`: `sizeof(TensorMeta)` 48->80, `strides` offset 32->48 (`ndim`/`numel`/`sizes` offsets unchanged). - WGSL `sizes: array<vec4<u32>, 2>` / `strides: array<vec4<u32>, 2>`: in the uniform address space the `vec4<u32>` element stride is 16, so two blocks are 32 contiguous bytes matching `uint32_t[8]`; component k sits at byte `16 + 4*k` on both sides, so the raw memcpy stays correct. Runtime reads `m.strides[d]` become `m.strides[d >> 2u][d & 3u]`. - 8 (not 5): a WGSL uniform array cannot pack 5 `u32` tightly, so rank-5 would allocate the same two `vec4`s; 8 is the natural 2-`vec4` boundary and covers all realistic ranks. Key changes: - `runtime/ops/TensorMeta.h` - cap 4->8; `static_assert(sizeof == 80)` and `offsetof(strides) == 48`; throw text. - 15 shaders widened + reindexed (regenerated their 16 `_wgsl.h`): `permute`, `flip`, `cat`, `binary_op` (`div`/`sub`), `binary_mul`, `binary_pow`, `binary_minimum`, `binary_floor_divide`, `where`, `select`, `slice`, `expand_copy`, `gather`, `index_select`, `repeat`. - `permute`/`flip` also widen their per-dim `Params` array (`perm`/`flip` `vec4<u32>` -> `array<vec4<u32>, 2>`; `sizeof(PermuteParams)`/`sizeof(FlipParams)` 16->32). - Rank-guard throw messages "exceeds 4" -> "exceeds 8" (the guards themselves auto-follow the constant): `div`/`mul`/`sub`/`where`/`binary`/`index_select`/`repeat`. - Applied identically to both the `xplat/` and `fbcode/` mirrors (byte-identical). Co-authored-with: Claude Code. ghstack-source-id: 406388763 @exported-using-ghexport Differential Revision: [D113319595](https://our.internmc.facebook.com/intern/diff/D113319595/)
…match) Pull Request resolved: #21232 Under dynamic shapes `cat` produced a wrong (scrambled + tail-garbage) output. `cat_impl` baked every shape-dependent quantity - each input `in_meta` (strides/numel), the shared `out_meta` (strides), the per-input `off_k`, and each dispatch's `workgroup_count` - from the MAX (upper-bound) build shape and registered NO resize hook, then released the UBOs. Under a smaller live shape the kernel then decoded coords with max strides, scattered with max out-strides, and looped over the max numel. This is the `[dynamic]` output-mismatch behind `test_inception_v3`, `test_squeezenet1_1`, `test_densenet161` (all channel-cat on spatial-dynamic feature maps; the concat dim is static so `off_k` was already correct - the defect is the stale spatial strides/numel). Fix (mirrors the WebGPUGraph SwiGLU/QKV resize templates and the shipped `mul` hook): - Keep the per-input `in_meta`/`params` UBOs and the shared `out_meta` UBO alive via `own_uniform_buffer` (previously released after build); collect their handles plus each `add_dispatch` index. - Register an idempotent `add_tensor_resize_hook` on every input id: from `cur_dims` recompute live out dims (`set_cur_dims(out_id, ...)` to cascade to consumers + fix the delegate-output shape), rebuild `out_meta` + each input's `in_meta`/`params` and `wgpuQueueWriteBuffer` them, and rewrite each dispatch's `workgroup_count_x` via `compute_1d_workgroup_count`. On a static graph `cur_dims == dims`, so the hook rewrites identical values (no behavior change). Applied identically to both the `xplat/` and `fbcode/` mirrors (byte-identical). Co-authored-with: Claude Code. ghstack-source-id: 406388769 @exported-using-ghexport Differential Revision: [D113319596](https://our.internmc.facebook.com/intern/diff/D113319596/)
…dels Pull Request resolved: #21233 `add.Tensor` was elementwise-only: `binary_add.wgsl` computed `output[idx] = input1[idx] + alpha * input2[idx]` over a single flat index up to the output numel, so on a broadcast it read the smaller operand out of bounds (WebGPU robustness clamps/zeros -> silently wrong, not `0x30`). This is the failure mode behind the models the WebGPU flow skips (`resnet50`, `vit_b_16`, `swin_v2_t`, `convnext_small`, `mobilenet_v3_small`, `shufflenet_v2_x1_0`) and the `bcast_first`/`bcast_second` op cases. Give `add.Tensor` NumPy broadcasting by mirroring the shipped `mul`: - `binary_add.wgsl` - copy `binary_mul`'s kernel (rank-8 `TensorMeta` layout from the rank-cap fix below it in the stack): keep the identical-shape elementwise fast path (common case stays bit-for-bit `input1[idx] + alpha * input2[idx]`), else relinearize out idx -> per-input coords (clamp size-1 dims) and add. `alpha` is a second pipeline-override constant (read once at build, never rewritten on resize), so no extra UBO. - `add/BinaryOp.cpp` - port `mul_impl`: 3 `TensorMeta` UBOs via `fill_tensor_meta_broadcast`, rank + fp32 guards, 6-entry bind group, `constantCount = 2` ({wg_size, alpha}), 2D dispatch, and `mul`'s resize hook verbatim (rebuild the 3 metas from `cur_dims`, `set_cur_dims`, rewrite UBOs + dispatch); `own_uniform_buffer` the 3 metas. Drops the old flat `AddParams` path. - `flows/webgpu.py` - un-skip the 6 broadcast-add models + the `bcast_first`/`bcast_second` op cases (kept the `float16`/`float64` dtype skips and `hardswish`/`lstm_batch_sizes`/`upsample_nearest2d`). This sits above the rank-cap fix, so the shader uses the rank-8 `array<vec4<u32>, 2>` `TensorMeta` layout. Applied identically to both `xplat/` and `fbcode/` mirrors (byte-identical). Co-authored-with: Claude Code. ghstack-source-id: 406388773 @exported-using-ghexport Differential Revision: [D113319599](https://our.internmc.facebook.com/intern/diff/D113319599/)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR was created by the merge bot to help merge the original PR into the main branch.
ghstack PR number: #21200 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/160/base
ghstack PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/160/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/JCNTH/159/orig
Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/160/orig
@diff-train-skip-merge