From 5e2bd00601c811b39d82c20b17326ae619071048 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:13 -0700 Subject: [PATCH 01/75] [ExecuTorch][WebGPU] Op-test for minimum MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21159 Adds the minimum op-test suite (2d + 3d same-shape) for the op ported in the diff below. Key changes: - `test/ops/test_minimum.py` — `MinimumModule` (`torch.minimum`). - `test/op_tests/cases.py` — `_minimum_suite`. ghstack-source-id: 406360846 @exported-using-ghexport Differential Revision: [D112257590](https://our.internmc.facebook.com/intern/diff/D112257590/) --- backends/webgpu/test/op_tests/cases.py | 13 +++++++++++++ backends/webgpu/test/ops/test_minimum.py | 18 ++++++++++++++++++ 2 files changed, 31 insertions(+) create mode 100644 backends/webgpu/test/ops/test_minimum.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 1780f65ec11..d60e8a9568b 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -34,6 +34,7 @@ CatModule, CONFIGS as _CAT_CONFIGS, ) +from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule from executorch.backends.webgpu.test.ops.test_mul import ( CONFIGS as _MUL_CONFIGS, MulModule, @@ -184,6 +185,18 @@ def _fn_config_suite(module_cls, configs) -> WebGPUTestSuite: ) +@register_op_test("minimum") +def _minimum_suite() -> WebGPUTestSuite: + # Same-shape numeric coverage (flat binary kernel; broadcast stays smoke). + return WebGPUTestSuite( + module_factory=lambda: MinimumModule(), + cases=[ + Case(name="2d", inputs=((M1, M2), (M1, M2))), + Case(name="3d", inputs=((S, S1, S2), (S, S1, S2))), + ], + ) + + @register_op_test("view_copy") def _view_copy_suite() -> WebGPUTestSuite: return _fn_config_suite(ViewModule, _VIEW_CONFIGS) diff --git a/backends/webgpu/test/ops/test_minimum.py b/backends/webgpu/test/ops/test_minimum.py new file mode 100644 index 00000000000..f8f8088fef5 --- /dev/null +++ b/backends/webgpu/test/ops/test_minimum.py @@ -0,0 +1,18 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.minimum.default` module for the WebGPU op-test framework. + +`MinimumModule` is imported by `cases.py`. minimum is a same-shape elementwise +binary op mirroring the landed `add`/`mul` pattern (flat 2D-dispatch kernel). +""" + +import torch + + +class MinimumModule(torch.nn.Module): + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return torch.minimum(a, b) From 0b9c32d066ffbb19872e34c3256ce08db6b88688 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:13 -0700 Subject: [PATCH 02/75] [ExecuTorch][WebGPU] Port pow (aten.pow.Tensor_Tensor) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21160 Ports `pow.Tensor_Tensor` with NumPy broadcasting, mirroring the sibling `mul` op + Vulkan `binary_op_buffer` (per-tensor `TensorMeta` uniforms; size-1 input dims clamped and the flat output index relinearized; equal-shape fast path). Key changes: - `runtime/ops/pow/{BinaryOp.cpp, binary_pow.wgsl}` (+ generated `_wgsl.h`) — `pow_impl` (args `[in1, in2, out]`); broadcast `TensorMeta` path + dual-operand resize hook; fail-loud int guard (fp32-only backend). Note: WGSL `pow(x, y)` is undefined for `x < 0`, exact parity with the Vulkan GLSL `pow`. Co-authored-with: Claude Code. ghstack-source-id: 406360857 @exported-using-ghexport Differential Revision: [D112257673](https://our.internmc.facebook.com/intern/diff/D112257673/) --- backends/webgpu/runtime/ops/pow/BinaryOp.cpp | 38 ++++++++++ .../webgpu/runtime/ops/pow/binary_pow.wgsl | 51 +++++++++++++ .../webgpu/runtime/ops/pow/binary_pow_wgsl.h | 75 +++++++++++++++++++ 3 files changed, 164 insertions(+) create mode 100644 backends/webgpu/runtime/ops/pow/BinaryOp.cpp create mode 100644 backends/webgpu/runtime/ops/pow/binary_pow.wgsl create mode 100644 backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h diff --git a/backends/webgpu/runtime/ops/pow/BinaryOp.cpp b/backends/webgpu/runtime/ops/pow/BinaryOp.cpp new file mode 100644 index 00000000000..d2f923880bf --- /dev/null +++ b/backends/webgpu/runtime/ops/pow/BinaryOp.cpp @@ -0,0 +1,38 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include + +#include + +namespace executorch::backends::webgpu { + +namespace { + +// aten.pow.Tensor_Tensor -> pow(in1, in2), with NumPy broadcasting (mirrors mul +// + Vulkan). WGSL pow(x,y) is undefined for x<0 (exact Vulkan-GLSL parity). +void pow_impl(WebGPUGraph& graph, const std::vector& args) { + add_binary_broadcast_op( + graph, + args.at(0), + args.at(1), + args.at(2), + kBinaryPowWGSL, + kBinaryPowWorkgroupSizeX, + "pow"); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.pow.Tensor_Tensor, pow_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/pow/binary_pow.wgsl b/backends/webgpu/runtime/ops/pow/binary_pow.wgsl new file mode 100644 index 00000000000..894c32a3dc2 --- /dev/null +++ b/backends/webgpu/runtime/ops/pow/binary_pow.wgsl @@ -0,0 +1,51 @@ +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in1_meta: TensorMeta; +@group(0) @binding(5) var in2_meta: TensorMeta; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= out_meta.numel) { + return; + } + + // Fast path: every input dim matches the output dim -> elementwise. + var same = true; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + if (in1_meta.sizes[d] != out_meta.sizes[d] || + in2_meta.sizes[d] != out_meta.sizes[d]) { + same = false; + } + } + if (same) { + output[idx] = pow(input1[idx], input2[idx]); + return; + } + + // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. + var rem = idx; + var l1: u32 = 0u; + var l2: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d]; + rem = rem % out_meta.strides[d]; + l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d]; + l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d]; + } + output[idx] = pow(input1[l1], input2[l2]); +} diff --git a/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h b/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h new file mode 100644 index 00000000000..6488e360518 --- /dev/null +++ b/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h @@ -0,0 +1,75 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from binary_pow.wgsl - DO NOT EDIT. +// wgsl-sha256: ffc116e845c861fb2a43f0c3ca86f150f8e1e95f9d3dde73a3886cd2a1ebff93 +inline constexpr const char* kBinaryPowWGSL = R"( +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in1_meta: TensorMeta; +@group(0) @binding(5) var in2_meta: TensorMeta; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= out_meta.numel) { + return; + } + + // Fast path: every input dim matches the output dim -> elementwise. + var same = true; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + if (in1_meta.sizes[d] != out_meta.sizes[d] || + in2_meta.sizes[d] != out_meta.sizes[d]) { + same = false; + } + } + if (same) { + output[idx] = pow(input1[idx], input2[idx]); + return; + } + + // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. + var rem = idx; + var l1: u32 = 0u; + var l2: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d]; + rem = rem % out_meta.strides[d]; + l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d]; + l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d]; + } + output[idx] = pow(input1[l1], input2[l2]); +} +)"; + +inline constexpr uint32_t kBinaryPowWorkgroupSizeX = 64; +inline constexpr uint32_t kBinaryPowWorkgroupSizeY = 1; +inline constexpr uint32_t kBinaryPowWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 03f5a3023eca6245ead657f6889144606294e189 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:14 -0700 Subject: [PATCH 03/75] [ExecuTorch][WebGPU] Op-tests for pow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21161 Adds the op-test suite for `pow.Tensor_Tensor` (ported in the diff below). Uses a strictly-positive base so `pow(neg, frac)=NaN` is never exercised. Key changes: - `test/ops/test_pow.py` + `test/op_tests/cases.py` — `PowModule` + `_pow_suite` (2d + 3d equal-shape, `atol=rtol=1e-3` vs the `torch` golden). Co-authored-with: Claude Code. ghstack-source-id: 406360852 @exported-using-ghexport Differential Revision: [D112257639](https://our.internmc.facebook.com/intern/diff/D112257639/) --- backends/webgpu/test/op_tests/cases.py | 25 +++++++++++++++++++++++++ backends/webgpu/test/ops/test_pow.py | 18 ++++++++++++++++++ 2 files changed, 43 insertions(+) create mode 100644 backends/webgpu/test/ops/test_pow.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index d60e8a9568b..1f4b2fcf94c 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -35,6 +35,7 @@ CONFIGS as _CAT_CONFIGS, ) from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule +from executorch.backends.webgpu.test.ops.test_pow import PowModule from executorch.backends.webgpu.test.ops.test_mul import ( CONFIGS as _MUL_CONFIGS, MulModule, @@ -197,6 +198,30 @@ def _minimum_suite() -> WebGPUTestSuite: ) +@register_op_test("pow") +def _pow_suite() -> WebGPUTestSuite: + # Positive base avoids pow(neg, frac)=NaN; exponent spans negative+positive. + return WebGPUTestSuite( + module_factory=lambda: PowModule(), + cases=[ + Case( + name="2d", + inputs=( + InputSpec(shape=(M1, M2), gen=_unary_lin(0.1, 3.0)), + InputSpec(shape=(M1, M2), gen=_unary_lin(-2.0, 3.0)), + ), + ), + Case( + name="3d", + inputs=( + InputSpec(shape=(S, S1, S2), gen=_unary_lin(0.1, 3.0)), + InputSpec(shape=(S, S1, S2), gen=_unary_lin(-2.0, 3.0)), + ), + ), + ], + ) + + @register_op_test("view_copy") def _view_copy_suite() -> WebGPUTestSuite: return _fn_config_suite(ViewModule, _VIEW_CONFIGS) diff --git a/backends/webgpu/test/ops/test_pow.py b/backends/webgpu/test/ops/test_pow.py new file mode 100644 index 00000000000..93596dc5d9f --- /dev/null +++ b/backends/webgpu/test/ops/test_pow.py @@ -0,0 +1,18 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.pow.Tensor_Tensor` module for the WebGPU op-test framework. + +`PowModule` is imported by `cases.py`. Same-shape elementwise `pow(a, b)`; the +suite uses a POSITIVE base so `pow(neg, frac)` (NaN) is never exercised. +""" + +import torch + + +class PowModule(torch.nn.Module): + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return torch.pow(a, b) From 47543b40540b59ad4fe9989ce6f4e2bc14d10187 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:14 -0700 Subject: [PATCH 04/75] [ExecuTorch][WebGPU] Port pow.Tensor_Scalar MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21162 Ports `pow.Tensor_Scalar` (tensor ^ scalar-exponent) on the shared unary helper: the exponent rides the `min` uniform slot, exactly like `leaky_relu`'s slope. `output[idx] = pow(input[idx], exponent)`. Key changes: - `runtime/ops/unary/pow_scalar.wgsl` (+ generated `_wgsl.h`) — `pow(input[idx], params.minimum)`. - `runtime/ops/unary/Activations.cpp` — `pow_scalar_impl` (args `[in, exponent, out]`; exponent read via the shared `get_val_or_inf` into the min slot), registered `aten.pow.Tensor_Scalar`. Inherits the shared helper's fp32 int-guard + resize hook (the exponent is re-applied on resize). Note: WGSL `pow` returns NaN for a negative base at any exponent (matching the reference GPU `pow`); the suite uses a positive base. ghstack-source-id: 406360862 @exported-using-ghexport Differential Revision: [D112257604](https://our.internmc.facebook.com/intern/diff/D112257604/) --- .../webgpu/runtime/ops/unary/Activations.cpp | 15 +++++++ .../webgpu/runtime/ops/unary/pow_scalar.wgsl | 21 +++++++++ .../runtime/ops/unary/pow_scalar_wgsl.h | 45 +++++++++++++++++++ 3 files changed, 81 insertions(+) create mode 100644 backends/webgpu/runtime/ops/unary/pow_scalar.wgsl create mode 100644 backends/webgpu/runtime/ops/unary/pow_scalar_wgsl.h diff --git a/backends/webgpu/runtime/ops/unary/Activations.cpp b/backends/webgpu/runtime/ops/unary/Activations.cpp index f56daf84d82..53b05e8b61b 100644 --- a/backends/webgpu/runtime/ops/unary/Activations.cpp +++ b/backends/webgpu/runtime/ops/unary/Activations.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -129,6 +130,19 @@ void hardtanh_impl(WebGPUGraph& graph, const std::vector& args) { hi); } +void pow_scalar_impl(WebGPUGraph& graph, const std::vector& args) { + // aten.pow.Tensor_Scalar args: [in, exponent, out]; exponent = min slot. + const float exponent = get_val_or_inf(graph, args.at(1), /*is_max=*/false); + add_unary_op( + graph, + args.at(0), + args.at(2), + kPowScalarWGSL, + kPowScalarWorkgroupSizeX, + "pow_scalar", + exponent); +} + } // namespace WEBGPU_REGISTER_OPERATORS { @@ -144,6 +158,7 @@ WEBGPU_REGISTER_OPERATORS { WEBGPU_REGISTER_OP(aten.hardswish.default, hardswish_impl); WEBGPU_REGISTER_OP(aten.clamp.default, clamp_impl); WEBGPU_REGISTER_OP(aten.hardtanh.default, hardtanh_impl); + WEBGPU_REGISTER_OP(aten.pow.Tensor_Scalar, pow_scalar_impl); } } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/unary/pow_scalar.wgsl b/backends/webgpu/runtime/ops/unary/pow_scalar.wgsl new file mode 100644 index 00000000000..db97cdffafc --- /dev/null +++ b/backends/webgpu/runtime/ops/unary/pow_scalar.wgsl @@ -0,0 +1,21 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, + minimum: f32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = pow(input[idx], params.minimum); +} diff --git a/backends/webgpu/runtime/ops/unary/pow_scalar_wgsl.h b/backends/webgpu/runtime/ops/unary/pow_scalar_wgsl.h new file mode 100644 index 00000000000..e296ae59a5e --- /dev/null +++ b/backends/webgpu/runtime/ops/unary/pow_scalar_wgsl.h @@ -0,0 +1,45 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from pow_scalar.wgsl - DO NOT EDIT. +// wgsl-sha256: a9176c6a6b0da421cd3649e396390cd10859c1256f6d38fce6bc730d3e3788e6 +inline constexpr const char* kPowScalarWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, + minimum: f32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = pow(input[idx], params.minimum); +} +)"; + +inline constexpr uint32_t kPowScalarWorkgroupSizeX = 256; +inline constexpr uint32_t kPowScalarWorkgroupSizeY = 1; +inline constexpr uint32_t kPowScalarWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 9d4cfae3da6b9e5d62af3e05a1840a8792643261 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:14 -0700 Subject: [PATCH 05/75] [ExecuTorch][WebGPU] Op-test for pow.Tensor_Scalar MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21163 Adds the pow_scalar op-test suite (square + sqrt exponents, positive base) for the op ported in the diff below. Key changes: - `test/ops/test_unary_activations.py` — `PowScalarModule` + `POW_SCALAR_CONFIGS`. - `test/op_tests/cases.py` — `_pow_scalar_suite`. ghstack-source-id: 406360849 @exported-using-ghexport Differential Revision: [D112257612](https://our.internmc.facebook.com/intern/diff/D112257612/) --- backends/webgpu/test/op_tests/cases.py | 18 ++++++++++++++++++ .../webgpu/test/ops/test_unary_activations.py | 18 ++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 1f4b2fcf94c..686e988d0a7 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -78,6 +78,8 @@ ClampModule, HARDTANH_CONFIGS, HardtanhModule, + POW_SCALAR_CONFIGS, + PowScalarModule, UNARY_G1, UnaryModule, ) @@ -1042,3 +1044,19 @@ def _hardtanh_suite() -> WebGPUTestSuite: for n, (lo, hi) in HARDTANH_CONFIGS.items() ], ) + + +@register_op_test("pow_scalar") +def _pow_scalar_suite() -> WebGPUTestSuite: + # Positive base: WGSL pow(neg base)=NaN for any exponent; exponent baked. + return WebGPUTestSuite( + module_factory=lambda exponent: PowScalarModule(exponent), + cases=[ + Case( + name=n, + construct={"exponent": e}, + inputs=(InputSpec(shape=(M1, M2), gen=_unary_lin(0.1, 4.0)),), + ) + for n, e in POW_SCALAR_CONFIGS.items() + ], + ) diff --git a/backends/webgpu/test/ops/test_unary_activations.py b/backends/webgpu/test/ops/test_unary_activations.py index bd722ebadfa..6b9509a042f 100644 --- a/backends/webgpu/test/ops/test_unary_activations.py +++ b/backends/webgpu/test/ops/test_unary_activations.py @@ -62,6 +62,24 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: } +class PowScalarModule(torch.nn.Module): + """aten.pow.Tensor_Scalar with a baked exponent (exponent → the min slot).""" + + def __init__(self, exponent) -> None: + super().__init__() + self.exponent = exponent + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.pow(x, self.exponent) + + +# name -> exponent construct kwarg; the suite uses a positive base to avoid NaN. +POW_SCALAR_CONFIGS = { + "square": 2.0, + "sqrt": 0.5, +} + + def _lin(lo: float, hi: float): """Deterministic linspace input of the requested shape over [lo, hi].""" From 6a59ec9074d26e78e6727f4de3554d9f0339f9df Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:15 -0700 Subject: [PATCH 06/75] [ExecuTorch][WebGPU] Port amax + amin (last-dim reductions) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21164 Ports the Vulkan delegate's `amax`/`amin` as last-dim per-row reductions — the buffer-storage path Vulkan itself restricts reductions to (`impl/Reduce.cpp` `add_reduce_per_row_node`; its texture path throws on buffer storage). One thread per output row serially reduces over the contiguous last dim. Key changes: - `runtime/ops/amax/{Reduce.cpp, amax.wgsl}` + `runtime/ops/amin/{Reduce.cpp, amin.wgsl}` (+ generated `_wgsl.h`) — `out[row] = reduce(input[row*W .. +W])` where `W` = last-dim extent; args `[in, dim, keepdim, out]`; fp32 int-guard; keepdim-aware resize hook; registered `aten.amax.default` / `aten.amin.default`. - `CMakeLists.txt` — both in `WEBGPU_SRCS`. Constraint: only a single last-dim reduction (`dim == -1` or `ndim-1`) is supported on buffer storage; any other reduce dim throws (mirrors Vulkan's buffer-per-row gate). The accumulator seeds `input[base]` (not 0), so all-negative rows reduce correctly. ghstack-source-id: 406360850 @exported-using-ghexport Differential Revision: [D112257619](https://our.internmc.facebook.com/intern/diff/D112257619/) --- backends/webgpu/runtime/ops/amax/Reduce.cpp | 190 +++++++++++++++++++ backends/webgpu/runtime/ops/amax/amax.wgsl | 26 +++ backends/webgpu/runtime/ops/amax/amax_wgsl.h | 50 +++++ backends/webgpu/runtime/ops/amin/Reduce.cpp | 190 +++++++++++++++++++ backends/webgpu/runtime/ops/amin/amin.wgsl | 26 +++ backends/webgpu/runtime/ops/amin/amin_wgsl.h | 50 +++++ 6 files changed, 532 insertions(+) create mode 100644 backends/webgpu/runtime/ops/amax/Reduce.cpp create mode 100644 backends/webgpu/runtime/ops/amax/amax.wgsl create mode 100644 backends/webgpu/runtime/ops/amax/amax_wgsl.h create mode 100644 backends/webgpu/runtime/ops/amin/Reduce.cpp create mode 100644 backends/webgpu/runtime/ops/amin/amin.wgsl create mode 100644 backends/webgpu/runtime/ops/amin/amin_wgsl.h diff --git a/backends/webgpu/runtime/ops/amax/Reduce.cpp b/backends/webgpu/runtime/ops/amax/Reduce.cpp new file mode 100644 index 00000000000..702aa21c218 --- /dev/null +++ b/backends/webgpu/runtime/ops/amax/Reduce.cpp @@ -0,0 +1,190 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform layout matching the WGSL Params struct; 16-byte aligned. +struct AmaxParams { + uint32_t num_rows; + uint32_t reduce_size; + uint32_t _pad[2]; +}; + +// Last-dim reduction; mirrors Vulkan add_reduce_per_row_node. +void amax_impl(WebGPUGraph& graph, const std::vector& args) { + // aten.amax.default args: [in, dim, keepdim, out] + const int in_id = args.at(0); + const int out_id = args.at(3); + + WGPUDevice device = graph.device(); + + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("amax: null buffer binding"); + } + if (in_tensor.is_int || out_tensor.is_int) { + throw std::runtime_error("amax: int dtype unsupported"); + } + + const std::vector& dims = graph.get_int_list(args.at(1)); + const int64_t ndim = static_cast(in_tensor.dims.size()); + if (dims.size() != 1 || (dims[0] != -1 && dims[0] != ndim - 1)) { + throw std::runtime_error("amax: only last-dim reduction is supported"); + } + + const uint32_t reduce_size = static_cast(in_tensor.dims.back()); + const uint32_t num_rows = + static_cast(out_tensor.nbytes / sizeof(float)); + if (reduce_size == 0u || + in_tensor.nbytes / sizeof(float) != + static_cast(num_rows) * reduce_size) { + throw std::runtime_error("amax: shape mismatch (num_rows * reduce_size)"); + } + + uint32_t wg_size = utils::clamp_workgroup_size(device, kAmaxWorkgroupSizeX); + utils::WgCount workgroup_count = + utils::compute_2d_workgroup_count(device, num_rows, wg_size, "amax"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + AmaxParams params = {}; + params.num_rows = num_rows; + params.reduce_size = reduce_size; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(AmaxParams)); + graph.add_uniform_buffer_bytes(sizeof(AmaxParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kAmaxWGSL, WGPU_STRLEN}; + + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // Bind group layout: input (read storage) + output (storage) + params. + WGPUBindGroupLayoutEntry entries[3] = {}; + + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 3; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[3] = {}; + + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + + bg_entries[2].binding = 2; + bg_entries[2].buffer = uniform_buffer; + bg_entries[2].size = sizeof(AmaxParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 3; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, bind_group, workgroup_count.x, "amax", workgroup_count.y}); + + // Dynamic shapes: recompute reduce_size (last dim) + num_rows + dispatch. + const bool keepdim = graph.get_bool(args.at(2)); + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, keepdim, wg_size, dispatch_idx, params_buf]( + WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + const uint32_t rsize = static_cast(d.back()); + if (rsize == 0u) { + throw std::runtime_error("amax(resize): zero reduce dim"); + } + const uint64_t total = utils::numel_of(d); + const uint32_t rows = static_cast(total / rsize); + std::vector od = d; + if (keepdim) { + od.back() = 1; + } else { + od.pop_back(); + } + g.set_cur_dims(out_id, od); + AmaxParams p = {}; + p.num_rows = rows; + p.reduce_size = rsize; + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), rows, wg_size, "amax"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }); + + // Release intermediates (pipeline + bind_group are kept by dispatch). + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.amax.default, amax_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/amax/amax.wgsl b/backends/webgpu/runtime/ops/amax/amax.wgsl new file mode 100644 index 00000000000..2ff33e0c20f --- /dev/null +++ b/backends/webgpu/runtime/ops/amax/amax.wgsl @@ -0,0 +1,26 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_rows: u32, + reduce_size: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let row = gid.x + gid.y * (num_workgroups.x * wg_size); + if (row >= params.num_rows) { + return; + } + let base = row * params.reduce_size; + var acc = input[base]; + for (var j = 1u; j < params.reduce_size; j = j + 1u) { + acc = max(acc, input[base + j]); + } + output[row] = acc; +} diff --git a/backends/webgpu/runtime/ops/amax/amax_wgsl.h b/backends/webgpu/runtime/ops/amax/amax_wgsl.h new file mode 100644 index 00000000000..1f82d7507cd --- /dev/null +++ b/backends/webgpu/runtime/ops/amax/amax_wgsl.h @@ -0,0 +1,50 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from amax.wgsl - DO NOT EDIT. +// wgsl-sha256: aae058ed0c432ac0cb54ea894e03e12a246d0d1c743d9ebb995b9773029e4652 +inline constexpr const char* kAmaxWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_rows: u32, + reduce_size: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let row = gid.x + gid.y * (num_workgroups.x * wg_size); + if (row >= params.num_rows) { + return; + } + let base = row * params.reduce_size; + var acc = input[base]; + for (var j = 1u; j < params.reduce_size; j = j + 1u) { + acc = max(acc, input[base + j]); + } + output[row] = acc; +} +)"; + +inline constexpr uint32_t kAmaxWorkgroupSizeX = 256; +inline constexpr uint32_t kAmaxWorkgroupSizeY = 1; +inline constexpr uint32_t kAmaxWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/amin/Reduce.cpp b/backends/webgpu/runtime/ops/amin/Reduce.cpp new file mode 100644 index 00000000000..35c7017436d --- /dev/null +++ b/backends/webgpu/runtime/ops/amin/Reduce.cpp @@ -0,0 +1,190 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform layout matching the WGSL Params struct; 16-byte aligned. +struct AminParams { + uint32_t num_rows; + uint32_t reduce_size; + uint32_t _pad[2]; +}; + +// Last-dim reduction; mirrors Vulkan add_reduce_per_row_node. +void amin_impl(WebGPUGraph& graph, const std::vector& args) { + // aten.amin.default args: [in, dim, keepdim, out] + const int in_id = args.at(0); + const int out_id = args.at(3); + + WGPUDevice device = graph.device(); + + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("amin: null buffer binding"); + } + if (in_tensor.is_int || out_tensor.is_int) { + throw std::runtime_error("amin: int dtype unsupported"); + } + + const std::vector& dims = graph.get_int_list(args.at(1)); + const int64_t ndim = static_cast(in_tensor.dims.size()); + if (dims.size() != 1 || (dims[0] != -1 && dims[0] != ndim - 1)) { + throw std::runtime_error("amin: only last-dim reduction is supported"); + } + + const uint32_t reduce_size = static_cast(in_tensor.dims.back()); + const uint32_t num_rows = + static_cast(out_tensor.nbytes / sizeof(float)); + if (reduce_size == 0u || + in_tensor.nbytes / sizeof(float) != + static_cast(num_rows) * reduce_size) { + throw std::runtime_error("amin: shape mismatch (num_rows * reduce_size)"); + } + + uint32_t wg_size = utils::clamp_workgroup_size(device, kAminWorkgroupSizeX); + utils::WgCount workgroup_count = + utils::compute_2d_workgroup_count(device, num_rows, wg_size, "amin"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + AminParams params = {}; + params.num_rows = num_rows; + params.reduce_size = reduce_size; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(AminParams)); + graph.add_uniform_buffer_bytes(sizeof(AminParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kAminWGSL, WGPU_STRLEN}; + + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // Bind group layout: input (read storage) + output (storage) + params. + WGPUBindGroupLayoutEntry entries[3] = {}; + + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 3; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[3] = {}; + + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + + bg_entries[2].binding = 2; + bg_entries[2].buffer = uniform_buffer; + bg_entries[2].size = sizeof(AminParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 3; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, bind_group, workgroup_count.x, "amin", workgroup_count.y}); + + // Dynamic shapes: recompute reduce_size (last dim) + num_rows + dispatch. + const bool keepdim = graph.get_bool(args.at(2)); + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, keepdim, wg_size, dispatch_idx, params_buf]( + WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + const uint32_t rsize = static_cast(d.back()); + if (rsize == 0u) { + throw std::runtime_error("amin(resize): zero reduce dim"); + } + const uint64_t total = utils::numel_of(d); + const uint32_t rows = static_cast(total / rsize); + std::vector od = d; + if (keepdim) { + od.back() = 1; + } else { + od.pop_back(); + } + g.set_cur_dims(out_id, od); + AminParams p = {}; + p.num_rows = rows; + p.reduce_size = rsize; + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), rows, wg_size, "amin"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }); + + // Release intermediates (pipeline + bind_group are kept by dispatch). + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.amin.default, amin_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/amin/amin.wgsl b/backends/webgpu/runtime/ops/amin/amin.wgsl new file mode 100644 index 00000000000..d8f4346d231 --- /dev/null +++ b/backends/webgpu/runtime/ops/amin/amin.wgsl @@ -0,0 +1,26 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_rows: u32, + reduce_size: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let row = gid.x + gid.y * (num_workgroups.x * wg_size); + if (row >= params.num_rows) { + return; + } + let base = row * params.reduce_size; + var acc = input[base]; + for (var j = 1u; j < params.reduce_size; j = j + 1u) { + acc = min(acc, input[base + j]); + } + output[row] = acc; +} diff --git a/backends/webgpu/runtime/ops/amin/amin_wgsl.h b/backends/webgpu/runtime/ops/amin/amin_wgsl.h new file mode 100644 index 00000000000..53f94e2f0e8 --- /dev/null +++ b/backends/webgpu/runtime/ops/amin/amin_wgsl.h @@ -0,0 +1,50 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from amin.wgsl - DO NOT EDIT. +// wgsl-sha256: 974a28fd80f089c8a52bf54d73f3bd03c195b2f9d7904bb4c31e6545813e0459 +inline constexpr const char* kAminWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_rows: u32, + reduce_size: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let row = gid.x + gid.y * (num_workgroups.x * wg_size); + if (row >= params.num_rows) { + return; + } + let base = row * params.reduce_size; + var acc = input[base]; + for (var j = 1u; j < params.reduce_size; j = j + 1u) { + acc = min(acc, input[base + j]); + } + output[row] = acc; +} +)"; + +inline constexpr uint32_t kAminWorkgroupSizeX = 256; +inline constexpr uint32_t kAminWorkgroupSizeY = 1; +inline constexpr uint32_t kAminWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From ee2cb9828902df3b5bd0b2378bca67a3307a3af0 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:15 -0700 Subject: [PATCH 07/75] [ExecuTorch][WebGPU] Op-tests for amax + amin MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21165 Adds op-test suites for the last-dim reductions ported in the diff below (keepdim + nodim, 2d + 3d). Key changes: - `test/ops/test_reduce.py` — `AmaxModule` / `AminModule` (reduce `dim=-1`). - `test/op_tests/cases.py` — `_reduce_suite` + `_amax_suite` / `_amin_suite`. ghstack-source-id: 406360856 @exported-using-ghexport Differential Revision: [D112257621](https://our.internmc.facebook.com/intern/diff/D112257621/) --- backends/webgpu/test/op_tests/cases.py | 34 ++++++++++++++++++++++++- backends/webgpu/test/ops/test_reduce.py | 22 ++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 686e988d0a7..5814c5bb1bf 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -35,7 +35,6 @@ CONFIGS as _CAT_CONFIGS, ) from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule -from executorch.backends.webgpu.test.ops.test_pow import PowModule from executorch.backends.webgpu.test.ops.test_mul import ( CONFIGS as _MUL_CONFIGS, MulModule, @@ -44,6 +43,8 @@ CONFIGS as _PERMUTE_CONFIGS, PermuteModule, ) +from executorch.backends.webgpu.test.ops.test_pow import PowModule +from executorch.backends.webgpu.test.ops.test_reduce import AmaxModule, AminModule from executorch.backends.webgpu.test.ops.test_rms_norm import ( _CASES, _linspace_weight, @@ -224,6 +225,37 @@ def _pow_suite() -> WebGPUTestSuite: ) +def _reduce_suite(module_cls) -> WebGPUTestSuite: + # Last-dim reduction; both keepdim variants over a 2d and a 3d shape. + return WebGPUTestSuite( + module_factory=lambda keepdim: module_cls(keepdim), + cases=[ + Case(name="keepdim_2d", construct={"keepdim": True}, inputs=((M1, M2),)), + Case(name="nodim_2d", construct={"keepdim": False}, inputs=((M1, M2),)), + Case( + name="keepdim_3d", + construct={"keepdim": True}, + inputs=((S, S1, S2),), + ), + Case( + name="nodim_3d", + construct={"keepdim": False}, + inputs=((S, S1, S2),), + ), + ], + ) + + +@register_op_test("amax") +def _amax_suite() -> WebGPUTestSuite: + return _reduce_suite(AmaxModule) + + +@register_op_test("amin") +def _amin_suite() -> WebGPUTestSuite: + return _reduce_suite(AminModule) + + @register_op_test("view_copy") def _view_copy_suite() -> WebGPUTestSuite: return _fn_config_suite(ViewModule, _VIEW_CONFIGS) diff --git a/backends/webgpu/test/ops/test_reduce.py b/backends/webgpu/test/ops/test_reduce.py index 39c9dc48506..c7b91d8156c 100644 --- a/backends/webgpu/test/ops/test_reduce.py +++ b/backends/webgpu/test/ops/test_reduce.py @@ -11,6 +11,10 @@ outer/r/inner decomposition: `dim=-1` gives inner=1 (unit-stride reduction), a middle dim gives inner>1 (the non-unit-stride path); `keepdim` toggles whether the reduced dim survives in the output shape. + +`AmaxModule`/`AminModule` (below) are imported by `cases.py` for the amax/amin +op-test suites. The WebGPU backend supports only the last-dim (per-row) reduction +on buffer storage, so those reduce over `dim=-1`, mirroring Vulkan's per-row path. """ from __future__ import annotations @@ -123,5 +127,23 @@ def export_reduce_model( print(f"Exported {pte_path}; golden {golden_path}; input {input_path}") +class AmaxModule(torch.nn.Module): + def __init__(self, keepdim: bool) -> None: + super().__init__() + self.keepdim = keepdim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.amax(x, dim=-1, keepdim=self.keepdim) + + +class AminModule(torch.nn.Module): + def __init__(self, keepdim: bool) -> None: + super().__init__() + self.keepdim = keepdim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.amin(x, dim=-1, keepdim=self.keepdim) + + if __name__ == "__main__": unittest.main() From 8bc5b9efdc984eca14c9289970dae6f60fb4d21e Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:16 -0700 Subject: [PATCH 08/75] [ExecuTorch][WebGPU] Port flip MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21166 Ports `aten.flip.default` as a per-element index-transform, mirroring the landed `permute` op (same dual-`TensorMeta` + 5-binding bind group). flip is pure data movement, so it reverses the coord along each flagged dim and copies — no fp compute. Key changes: - `runtime/ops/flip/{Flip.cpp, flip.wgsl}` (+ generated `flip_wgsl.h`) — per out flat idx: unravel via `out_meta.strides`, `in_coord = select(coord, sizes[d]-1-coord, flip[d]==1)`, ravel via `in_meta.strides`, `output[out] = input[in]`; `Params { flip: vec4 }` bitmap built from the normalized `dims` arg; registered `aten.flip.default`. - `CMakeLists.txt` — `runtime/ops/flip/Flip.cpp` in `WEBGPU_SRCS`. Matches Vulkan `impl/Flip.cpp` (`create_whcn_bitmap` normalizes/flags dims; reversal `size-1-coord`). No int dtype-guard (unlike clamp/minimum which do fp arithmetic): flip is bit-preserving movement, so any 4-byte dtype is correct; the `nbytes != numel*4` guard rejects int64/fp16. 2D-folded dispatch lifts the 65535 cap. `out` is the last arg (this backend's convention). ghstack-source-id: 406360854 @exported-using-ghexport Differential Revision: [D112257643](https://our.internmc.facebook.com/intern/diff/D112257643/) --- backends/webgpu/runtime/ops/flip/Flip.cpp | 197 +++++++++++++++++++ backends/webgpu/runtime/ops/flip/flip.wgsl | 41 ++++ backends/webgpu/runtime/ops/flip/flip_wgsl.h | 65 ++++++ 3 files changed, 303 insertions(+) create mode 100644 backends/webgpu/runtime/ops/flip/Flip.cpp create mode 100644 backends/webgpu/runtime/ops/flip/flip.wgsl create mode 100644 backends/webgpu/runtime/ops/flip/flip_wgsl.h diff --git a/backends/webgpu/runtime/ops/flip/Flip.cpp b/backends/webgpu/runtime/ops/flip/Flip.cpp new file mode 100644 index 00000000000..81943ac3151 --- /dev/null +++ b/backends/webgpu/runtime/ops/flip/Flip.cpp @@ -0,0 +1,197 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct FlipParams { + uint32_t flip[kTensorMetaMaxNdim]; +}; +static_assert( + sizeof(FlipParams) == 16, + "FlipParams must match the WGSL Params vec4 (16 bytes)"); + +// flip: reverse coords along dims (Vulkan Flip.cpp, NCHW; any 4-byte dtype). +void flip_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [self, dims, out]; dims = dims to reverse; out is the last value-id. + const int in_id = args.at(0); + const int dims_id = args.at(1); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("flip: in/out arg is not a tensor"); + } + if (graph.get_value_type(dims_id) != WebGPUGraph::ValueType::IntList) { + throw std::runtime_error("flip: dims arg is not an IntList"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + const int ndim = static_cast(in_tensor.dims.size()); + if (ndim > static_cast(kTensorMetaMaxNdim)) { + throw std::runtime_error("flip: tensor rank exceeds 4 (MAX_NDIM)"); + } + // flip preserves shape: the output rank + dims must equal the input's, else + // the kernel unravels out coords against a shape it wasn't reversed for. + if (static_cast(out_tensor.dims.size()) != ndim) { + throw std::runtime_error("flip: output rank != input rank"); + } + for (int i = 0; i < ndim; i++) { + if (out_tensor.dims[i] != in_tensor.dims[i]) { + throw std::runtime_error("flip: output shape != input shape"); + } + } + + // Build the flip bitmap: 1 for each (normalized) dim to reverse. + const std::vector& dims = graph.get_int_list(dims_id); + uint32_t flip[kTensorMetaMaxNdim] = {}; + bool seen[kTensorMetaMaxNdim] = {}; + for (int64_t dv : dims) { + int64_t d = dv < 0 ? dv + ndim : dv; + if (d < 0 || d >= ndim) { + throw std::runtime_error("flip: dim out of range"); + } + if (seen[static_cast(d)]) { + throw std::runtime_error("flip: duplicate dim"); + } + seen[static_cast(d)] = true; + flip[static_cast(d)] = 1u; + } + + TensorMeta out_meta; + TensorMeta in_meta; + fill_tensor_meta(out_tensor, &out_meta); + fill_tensor_meta(in_tensor, &in_meta); + if (out_tensor.nbytes != + static_cast(out_meta.numel) * sizeof(float) || + in_tensor.nbytes != static_cast(in_meta.numel) * sizeof(float)) { + throw std::runtime_error("flip: non-4-byte operand (nbytes != numel * 4)"); + } + + FlipParams params = {}; + std::memcpy(params.flip, flip, sizeof(flip)); + + uint32_t wg_size = utils::clamp_workgroup_size(device, kFlipWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, out_meta.numel, wg_size, "flip"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer out_meta_buf = + utils::make_uniform(device, &out_meta, sizeof(TensorMeta)); + WGPUBuffer in_meta_buf = + utils::make_uniform(device, &in_meta, sizeof(TensorMeta)); + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(FlipParams)); + graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta) + sizeof(FlipParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kFlipWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // Bind group: in, out (rw), out_meta, in_meta, params (3 uniforms). + WGPUBindGroupLayoutEntry entries[5] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + entries[3].binding = 3; + entries[3].visibility = WGPUShaderStage_Compute; + entries[3].buffer.type = WGPUBufferBindingType_Uniform; + entries[4].binding = 4; + entries[4].visibility = WGPUShaderStage_Compute; + entries[4].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 5; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[5] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = out_meta_buf; + bg_entries[2].size = sizeof(TensorMeta); + bg_entries[3].binding = 3; + bg_entries[3].buffer = in_meta_buf; + bg_entries[3].size = sizeof(TensorMeta); + bg_entries[4].binding = 4; + bg_entries[4].buffer = params_buf; + bg_entries[4].size = sizeof(FlipParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 5; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, bind_group, workgroup_count.x, "flip", workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Drop our refs; the bind group keeps the uniforms alive until release. + wgpuBufferRelease(out_meta_buf); + wgpuBufferRelease(in_meta_buf); + wgpuBufferRelease(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.flip.default, flip_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/flip/flip.wgsl b/backends/webgpu/runtime/ops/flip/flip.wgsl new file mode 100644 index 00000000000..b66c9cd6751 --- /dev/null +++ b/backends/webgpu/runtime/ops/flip/flip.wgsl @@ -0,0 +1,41 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(2) var out_meta: TensorMeta; +@group(0) @binding(3) var in_meta: TensorMeta; + +struct Params { + flip: vec4, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (out_bufi >= out_meta.numel) { + return; + } + + // Reverse the coord along each flagged dim (Vulkan flip, NCHW buffer). + var rem = out_bufi; + var in_bufi: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d]; + rem = rem % out_meta.strides[d]; + let flipped = in_meta.sizes[d] - 1u - coord; + let in_coord = select(coord, flipped, params.flip[d] == 1u); + in_bufi = in_bufi + in_coord * in_meta.strides[d]; + } + output[out_bufi] = input[in_bufi]; +} diff --git a/backends/webgpu/runtime/ops/flip/flip_wgsl.h b/backends/webgpu/runtime/ops/flip/flip_wgsl.h new file mode 100644 index 00000000000..138fa66aef5 --- /dev/null +++ b/backends/webgpu/runtime/ops/flip/flip_wgsl.h @@ -0,0 +1,65 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from flip.wgsl - DO NOT EDIT. +// wgsl-sha256: e85b3a7f753ec183e83f55f53cc764b5ffa748b4e229d81b978dabe58f576e80 +inline constexpr const char* kFlipWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(2) var out_meta: TensorMeta; +@group(0) @binding(3) var in_meta: TensorMeta; + +struct Params { + flip: vec4, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (out_bufi >= out_meta.numel) { + return; + } + + // Reverse the coord along each flagged dim (Vulkan flip, NCHW buffer). + var rem = out_bufi; + var in_bufi: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d]; + rem = rem % out_meta.strides[d]; + let flipped = in_meta.sizes[d] - 1u - coord; + let in_coord = select(coord, flipped, params.flip[d] == 1u); + in_bufi = in_bufi + in_coord * in_meta.strides[d]; + } + output[out_bufi] = input[in_bufi]; +} +)"; + +inline constexpr uint32_t kFlipWorkgroupSizeX = 64; +inline constexpr uint32_t kFlipWorkgroupSizeY = 1; +inline constexpr uint32_t kFlipWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From a1ad12ef471c4c6febf21a09a5bf8a1678df55e7 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:16 -0700 Subject: [PATCH 09/75] [ExecuTorch][WebGPU] Op-tests for flip MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21167 Adds the op-test suite for the `flip` op ported in the diff below. Key changes: - `test/ops/test_flip.py` — `FlipModule` (reverses `dims`) + `FlipTest` export-delegation smoke test (mirrors `PermuteTest`). - `test/op_tests/cases.py` — `_flip_suite`: last `[-1]`/2d, dim0 `[0]`/2d, both_3d `[0,2]`/3d, mid_3d `[1]`/3d, multi_4d `[1,3]`/4d. ghstack-source-id: 406360858 @exported-using-ghexport Differential Revision: [D112257622](https://our.internmc.facebook.com/intern/diff/D112257622/) --- backends/webgpu/test/op_tests/cases.py | 21 +++++++ backends/webgpu/test/ops/test_flip.py | 76 ++++++++++++++++++++++++++ 2 files changed, 97 insertions(+) create mode 100644 backends/webgpu/test/ops/test_flip.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 5814c5bb1bf..369f1f61e95 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -34,6 +34,7 @@ CatModule, CONFIGS as _CAT_CONFIGS, ) +from executorch.backends.webgpu.test.ops.test_flip import FlipModule from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule from executorch.backends.webgpu.test.ops.test_mul import ( CONFIGS as _MUL_CONFIGS, @@ -256,6 +257,26 @@ def _amin_suite() -> WebGPUTestSuite: return _reduce_suite(AminModule) +@register_op_test("flip") +def _flip_suite() -> WebGPUTestSuite: + # Reverse various dims; pure data movement -> float32 oracle. + return WebGPUTestSuite( + module_factory=lambda dims: FlipModule(dims), + cases=[ + Case(name="last", construct={"dims": [-1]}, inputs=((M1, M2),)), + Case(name="dim0", construct={"dims": [0]}, inputs=((M1, M2),)), + Case(name="both_3d", construct={"dims": [0, 2]}, inputs=((S, S1, S2),)), + Case(name="mid_3d", construct={"dims": [1]}, inputs=((S, S1, S2),)), + Case( + name="multi_4d", + construct={"dims": [1, 3]}, + inputs=((XS, S, S1, S2),), + ), + ], + golden_dtype="float32", + ) + + @register_op_test("view_copy") def _view_copy_suite() -> WebGPUTestSuite: return _fn_config_suite(ViewModule, _VIEW_CONFIGS) diff --git a/backends/webgpu/test/ops/test_flip.py b/backends/webgpu/test/ops/test_flip.py new file mode 100644 index 00000000000..902c99052f1 --- /dev/null +++ b/backends/webgpu/test/ops/test_flip.py @@ -0,0 +1,76 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.flip.default` module + configs for the WebGPU op-test framework. + +`FlipModule` reverses the given dims. flip is pure data movement (bit-identical), +so the suite uses the float32 oracle. `FlipTest` is the export-delegation smoke +test. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# name -> (input_shape, dims) +CONFIGS = { + "mid_3d": ((2, 3, 4), [1]), + "multi_4d": ((2, 3, 4, 5), [1, 3]), + "last_2d": ((3, 5), [-1]), + "all_3d": ((2, 3, 4), [0, 1, 2]), +} + + +class FlipModule(torch.nn.Module): + def __init__(self, dims) -> None: + super().__init__() + self.dims = dims + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.flip(x, self.dims) + + +def _det_input(shape): + g = torch.Generator().manual_seed(0) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(dims, x: torch.Tensor): + ep = torch.export.export(FlipModule(dims).eval(), (x,)) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + +def _delegated(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +def _op_delegated(edge, op_substr: str) -> bool: + # op must be absorbed into the delegate, not left as a top-level CPU-fallback node. + gm = edge.exported_program().graph_module + return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + + +class FlipTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, (shape, dims) in CONFIGS.items(): + with self.subTest(name=name): + edge = _lower(dims, _det_input(shape)) + et = edge.to_executorch() + self.assertTrue( + _delegated(et), + f"Expected a VulkanBackend delegate (flip {name})", + ) + self.assertTrue( + _op_delegated(edge, "flip"), + f"flip not delegated (fell back to CPU) for {name}", + ) From 5c70b0ce6d667f5935e94ac893c780d2e41f0e1f Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:16 -0700 Subject: [PATCH 10/75] [ExecuTorch][WebGPU] Port repeat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21168 Ports `aten.repeat.default` as a per-element tiling gather, mirroring the landed `permute`/`flip` index-transform (dual-`TensorMeta`, 2D-fold dispatch). repeat is pure data movement, so it maps each output element to its source via a per-dim modulo and copies — no fp compute. Key changes: - `runtime/ops/repeat/{Repeat.cpp, repeat.wgsl}` (+ generated `repeat_wgsl.h`) — per out flat idx: unravel via `out_meta.strides`, `in_coord = out_coord % in_meta.sizes[in_d]`, ravel via `in_meta.strides`, `output[out] = input[in]`; registered `aten.repeat.default`. - `CMakeLists.txt` — `runtime/ops/repeat/Repeat.cpp` in `WEBGPU_SRCS`. Mirrors Vulkan `glsl/repeat_buffer.glsl` (`in_tidx[d] = out_tidx[d] % size_at(in_meta, d)`). Vulkan's innermost-first `BufferMetadata` auto-aligns a shorter input; WebGPU's NCHW `TensorMeta` makes the right-alignment explicit via `offset = out_ndim - in_ndim` (leading `d < offset` dims are pure prepends, contributing 0). No `repeats`-arg read needed: the repeated out dims are baked in by export, so the tiling is recoverable from shapes alone. No int dtype-guard (bit-preserving movement, like `flip`); the `nbytes != numel*4` guard rejects int64/fp16. ghstack-source-id: 406360860 @exported-using-ghexport Differential Revision: [D112257679](https://our.internmc.facebook.com/intern/diff/D112257679/) --- backends/webgpu/runtime/ops/repeat/Repeat.cpp | 148 ++++++++++++++++++ .../webgpu/runtime/ops/repeat/repeat.wgsl | 40 +++++ .../webgpu/runtime/ops/repeat/repeat_wgsl.h | 64 ++++++++ 3 files changed, 252 insertions(+) create mode 100644 backends/webgpu/runtime/ops/repeat/Repeat.cpp create mode 100644 backends/webgpu/runtime/ops/repeat/repeat.wgsl create mode 100644 backends/webgpu/runtime/ops/repeat/repeat_wgsl.h diff --git a/backends/webgpu/runtime/ops/repeat/Repeat.cpp b/backends/webgpu/runtime/ops/repeat/Repeat.cpp new file mode 100644 index 00000000000..04553b0a67d --- /dev/null +++ b/backends/webgpu/runtime/ops/repeat/Repeat.cpp @@ -0,0 +1,148 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +// repeat: tile input along each dim (Vulkan Repeat.cpp, NCHW; 4-byte dtype). +void repeat_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [self, repeats, out]; out dims baked in by export (gather % size). + const int in_id = args.at(0); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("repeat: in/out arg is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (out_tensor.dims.size() > kTensorMetaMaxNdim || + in_tensor.dims.size() > kTensorMetaMaxNdim) { + throw std::runtime_error("repeat: tensor rank exceeds 4 (MAX_NDIM)"); + } + + TensorMeta out_meta; + TensorMeta in_meta; + fill_tensor_meta(out_tensor, &out_meta); + fill_tensor_meta(in_tensor, &in_meta); + if (out_tensor.nbytes != + static_cast(out_meta.numel) * sizeof(float) || + in_tensor.nbytes != static_cast(in_meta.numel) * sizeof(float)) { + throw std::runtime_error( + "repeat: non-4-byte operand (nbytes != numel * 4)"); + } + + uint32_t wg_size = utils::clamp_workgroup_size(device, kRepeatWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, out_meta.numel, wg_size, "repeat"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer out_meta_buf = + utils::make_uniform(device, &out_meta, sizeof(TensorMeta)); + WGPUBuffer in_meta_buf = + utils::make_uniform(device, &in_meta, sizeof(TensorMeta)); + graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kRepeatWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // Bind group: in, out (rw), out_meta, in_meta (2 uniforms). + WGPUBindGroupLayoutEntry entries[4] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + entries[3].binding = 3; + entries[3].visibility = WGPUShaderStage_Compute; + entries[3].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 4; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[4] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = out_meta_buf; + bg_entries[2].size = sizeof(TensorMeta); + bg_entries[3].binding = 3; + bg_entries[3].buffer = in_meta_buf; + bg_entries[3].size = sizeof(TensorMeta); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 4; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, bind_group, workgroup_count.x, "repeat", workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Drop our refs; the bind group keeps the uniforms alive until release. + wgpuBufferRelease(out_meta_buf); + wgpuBufferRelease(in_meta_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.repeat.default, repeat_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/repeat/repeat.wgsl b/backends/webgpu/runtime/ops/repeat/repeat.wgsl new file mode 100644 index 00000000000..8e0c2c4e676 --- /dev/null +++ b/backends/webgpu/runtime/ops/repeat/repeat.wgsl @@ -0,0 +1,40 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(2) var out_meta: TensorMeta; +@group(0) @binding(3) var in_meta: TensorMeta; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (out_bufi >= out_meta.numel) { + return; + } + + // Tile: gather in_coord[d] = out_coord[d] % in_size (Vulkan repeat_buffer). + // in is right-aligned into out; the leading offset dims are pure repeats. + let offset = out_meta.ndim - in_meta.ndim; + var rem = out_bufi; + var in_bufi: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let out_coord = rem / out_meta.strides[d]; + rem = rem % out_meta.strides[d]; + if (d >= offset) { + let in_d = d - offset; + let in_coord = out_coord % in_meta.sizes[in_d]; + in_bufi = in_bufi + in_coord * in_meta.strides[in_d]; + } + } + output[out_bufi] = input[in_bufi]; +} diff --git a/backends/webgpu/runtime/ops/repeat/repeat_wgsl.h b/backends/webgpu/runtime/ops/repeat/repeat_wgsl.h new file mode 100644 index 00000000000..b08d2d3d3c4 --- /dev/null +++ b/backends/webgpu/runtime/ops/repeat/repeat_wgsl.h @@ -0,0 +1,64 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from repeat.wgsl - DO NOT EDIT. +// wgsl-sha256: ae77e3d6846cd185e96b0fcae4b4bcee2096e659df1a7ba9d0129beb24d754c7 +inline constexpr const char* kRepeatWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(2) var out_meta: TensorMeta; +@group(0) @binding(3) var in_meta: TensorMeta; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (out_bufi >= out_meta.numel) { + return; + } + + // Tile: gather in_coord[d] = out_coord[d] % in_size (Vulkan repeat_buffer). + // in is right-aligned into out; the leading offset dims are pure repeats. + let offset = out_meta.ndim - in_meta.ndim; + var rem = out_bufi; + var in_bufi: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let out_coord = rem / out_meta.strides[d]; + rem = rem % out_meta.strides[d]; + if (d >= offset) { + let in_d = d - offset; + let in_coord = out_coord % in_meta.sizes[in_d]; + in_bufi = in_bufi + in_coord * in_meta.strides[in_d]; + } + } + output[out_bufi] = input[in_bufi]; +} +)"; + +inline constexpr uint32_t kRepeatWorkgroupSizeX = 64; +inline constexpr uint32_t kRepeatWorkgroupSizeY = 1; +inline constexpr uint32_t kRepeatWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 4174a3209fd6a209338c928f8a230d775e26951e Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:17 -0700 Subject: [PATCH 11/75] [ExecuTorch][WebGPU] Op-tests for repeat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21169 Adds the op-test suite for the `repeat` op ported in the diff below. Key changes: - `test/ops/test_repeat.py` — `RepeatModule` (`x.repeat(*repeats)`) + `RepeatTest` export-delegation smoke test (mirrors `PermuteTest`/`FlipTest`). - `test/op_tests/cases.py` — `_repeat_suite`: tile_1d, tile_2d, prepend_3d `[1,3,2]`, prepend_ext `[2,3,1]`, tile_3d `[2,1,2]`. ghstack-source-id: 406360847 @exported-using-ghexport Differential Revision: [D112257616](https://our.internmc.facebook.com/intern/diff/D112257616/) --- backends/webgpu/test/op_tests/cases.py | 29 ++++++++++ backends/webgpu/test/ops/test_repeat.py | 77 +++++++++++++++++++++++++ 2 files changed, 106 insertions(+) create mode 100644 backends/webgpu/test/ops/test_repeat.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 369f1f61e95..64663de7257 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -35,6 +35,7 @@ CONFIGS as _CAT_CONFIGS, ) from executorch.backends.webgpu.test.ops.test_flip import FlipModule +from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule from executorch.backends.webgpu.test.ops.test_mul import ( CONFIGS as _MUL_CONFIGS, @@ -277,6 +278,34 @@ def _flip_suite() -> WebGPUTestSuite: ) +@register_op_test("repeat") +def _repeat_suite() -> WebGPUTestSuite: + # Tile along dims; pure data movement -> float32 oracle. + return WebGPUTestSuite( + module_factory=lambda repeats: RepeatModule(repeats), + cases=[ + Case(name="tile_1d", construct={"repeats": [2]}, inputs=((XS,),)), + Case(name="tile_2d", construct={"repeats": [2, 2]}, inputs=((XS, S),)), + Case( + name="prepend_3d", + construct={"repeats": [1, 3, 2]}, + inputs=((XS, S),), + ), + Case( + name="prepend_ext", + construct={"repeats": [2, 3, 1]}, + inputs=((XS, S),), + ), + Case( + name="tile_3d", + construct={"repeats": [2, 1, 2]}, + inputs=((XS, S, S1),), + ), + ], + golden_dtype="float32", + ) + + @register_op_test("view_copy") def _view_copy_suite() -> WebGPUTestSuite: return _fn_config_suite(ViewModule, _VIEW_CONFIGS) diff --git a/backends/webgpu/test/ops/test_repeat.py b/backends/webgpu/test/ops/test_repeat.py new file mode 100644 index 00000000000..0e21c7d3d98 --- /dev/null +++ b/backends/webgpu/test/ops/test_repeat.py @@ -0,0 +1,77 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.repeat.default` module + configs for the WebGPU op-test framework. + +`RepeatModule` tiles the input along each dim. repeat is pure data movement +(bit-identical), so the suite uses the float32 oracle. `RepeatTest` is the +export-delegation smoke test. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# name -> (input_shape, repeats) +CONFIGS = { + "tile_1d": ((3,), (2,)), + "tile_2d": ((2, 3), (2, 2)), + "prepend_3d": ((2, 3), (1, 3, 2)), + "prepend_ext": ((2, 3), (2, 3, 1)), + "tile_3d": ((2, 3, 4), (2, 1, 2)), +} + + +class RepeatModule(torch.nn.Module): + def __init__(self, repeats) -> None: + super().__init__() + self.repeats = repeats + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.repeat(*self.repeats) + + +def _det_input(shape): + g = torch.Generator().manual_seed(0) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(repeats, x: torch.Tensor): + ep = torch.export.export(RepeatModule(repeats).eval(), (x,)) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + +def _delegated(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +def _op_delegated(edge, op_substr: str) -> bool: + # op must be absorbed into the delegate, not left as a top-level CPU-fallback node. + gm = edge.exported_program().graph_module + return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + + +class RepeatTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, (shape, repeats) in CONFIGS.items(): + with self.subTest(name=name): + edge = _lower(repeats, _det_input(shape)) + et = edge.to_executorch() + self.assertTrue( + _delegated(et), + f"Expected a VulkanBackend delegate (repeat {name})", + ) + self.assertTrue( + _op_delegated(edge, "repeat"), + f"repeat not delegated (fell back to CPU) for {name}", + ) From fd152d4650bce5f1df5b4e0874764b028d9e6491 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:17 -0700 Subject: [PATCH 12/75] [ExecuTorch][WebGPU] Op-test framework: multi-output goldens MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21170 Extends the op-test generator to golden-verify EVERY output of a multi-output op (e.g. `split_with_sizes_copy` → N tensors), not just `output[0]`. The C++ driver already selected `outs[output_index]` per manifest entry; only the generator hardcoded a single `output_index = 0`. Key changes: - `test/op_tests/generate_op_tests.py` — `generate_case` now returns one manifest entry PER output tensor: it iterates the module's tuple return, writes a golden per output (`{case}_out{i}.golden.bin`, `output_index = i`), and runs the dual-oracle gate per output. `generate()` uses `extend`. Single-output ops emit exactly one entry, byte-identical to before (no `_out` suffix, `output_index = 0`). - `test/op_tests/test_generator.py` — `test_generate_case_writes_artifacts` reads the returned list (asserts one entry for the single-output `add`). ghstack-source-id: 406360853 @exported-using-ghexport Differential Revision: [D112257608](https://our.internmc.facebook.com/intern/diff/D112257608/) --- .../webgpu/test/op_tests/generate_op_tests.py | 72 ++++++++++--------- .../webgpu/test/op_tests/test_generator.py | 5 +- 2 files changed, 44 insertions(+), 33 deletions(-) diff --git a/backends/webgpu/test/op_tests/generate_op_tests.py b/backends/webgpu/test/op_tests/generate_op_tests.py index d0b972ec74e..a4ff7aaa6e6 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -78,8 +78,10 @@ def _write_fp32(t: torch.Tensor, path: str) -> None: t.detach().contiguous().cpu().numpy().astype(" dict: - """Export one case, write its .pte + input/golden .bin, return its manifest entry.""" +def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> list[dict]: + """Export one case; write its .pte + input/golden .bin(s). Returns one manifest + entry per output tensor (multi-output ops emit N; single-output ops emit one, + byte-identical to the pre-multi-output form).""" module, inputs, prog = export_case(suite, case) if not _has_vulkan_delegate(prog): msg = ( @@ -90,11 +92,9 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> dict: raise RuntimeError(msg) print(f"WARNING: {msg}") golden_dtype = getattr(suite, "golden_dtype", "float64") - out_index = 0 with torch.no_grad(): # fp32 eager is the dual-oracle reference; compute it BEFORE any .double(). eager = module(*inputs) - eager_t = eager[out_index] if isinstance(eager, (tuple, list)) else eager if case.golden_fn is not None: golden = case.golden_fn(module, inputs) elif golden_dtype == "float64": @@ -105,16 +105,10 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> dict: ) else: golden = eager # gather/copy: fp64 is bit-identical, skip it - golden_t = golden[out_index] if isinstance(golden, (tuple, list)) else golden - out_t = golden_t.to(torch.float32) + eager_outs = list(eager) if isinstance(eager, (tuple, list)) else [eager] + golden_outs = list(golden) if isinstance(golden, (tuple, list)) else [golden] atol = case.atol if case.atol is not None else suite.atol rtol = case.rtol if case.rtol is not None else suite.rtol - # Dual-oracle gate: the fp64 golden must match the fp32 eager within tol — proves - # the oracle isn't itself buggy. Skipped for float32 suites. - if golden_dtype == "float64" and case.golden_fn is None: - torch.testing.assert_close( - eager_t.to(torch.float32), out_t, atol=atol, rtol=rtol - ) case_id = f"{op}__{case.name or 'case'}" pte_rel = f"{case_id}.pte" @@ -132,25 +126,39 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> dict: in_dtype = "float32" input_entries.append({"path": rel, "shape": list(t.shape), "dtype": in_dtype}) - golden_rel = f"{case_id}.golden.bin" - _write_fp32(out_t, os.path.join(out_dir, golden_rel)) - - return { - "op": op, - "case": case.name or "case", - "pte": pte_rel, - "inputs": input_entries, - "golden": { - "path": golden_rel, - "shape": list(out_t.shape), - "dtype": "float32", - "output_index": out_index, - }, - "atol": atol, - "rtol": rtol, - "required": case.required, - "heavy": case.heavy, - } + entries: list[dict] = [] + n_out = len(golden_outs) + for out_index in range(n_out): + out_t = golden_outs[out_index].to(torch.float32) + # Dual-oracle gate: the fp64 golden must match the fp32 eager within tol — + # proves the oracle isn't itself buggy. Skipped for float32 suites. + if golden_dtype == "float64" and case.golden_fn is None: + torch.testing.assert_close( + eager_outs[out_index].to(torch.float32), out_t, atol=atol, rtol=rtol + ) + # Single-output ops keep the original name/golden; multi-output ops suffix. + suffix = f"_out{out_index}" if n_out > 1 else "" + golden_rel = f"{case_id}{suffix}.golden.bin" + _write_fp32(out_t, os.path.join(out_dir, golden_rel)) + entries.append( + { + "op": op, + "case": f"{case.name or 'case'}{suffix}", + "pte": pte_rel, + "inputs": input_entries, + "golden": { + "path": golden_rel, + "shape": list(out_t.shape), + "dtype": "float32", + "output_index": out_index, + }, + "atol": atol, + "rtol": rtol, + "required": case.required, + "heavy": case.heavy, + } + ) + return entries def generate( @@ -177,7 +185,7 @@ def generate( if case.heavy and not heavy_enabled: # Export-gated: omitted from the default manifest (set WEBGPU_TEST_HEAVY). continue - entries.append(generate_case(op, suite, case, out_dir)) + entries.extend(generate_case(op, suite, case, out_dir)) with open(os.path.join(out_dir, "manifest.json"), "w") as f: json.dump(entries, f, indent=2) return entries diff --git a/backends/webgpu/test/op_tests/test_generator.py b/backends/webgpu/test/op_tests/test_generator.py index 9782d6fc1dc..fabf79171c0 100644 --- a/backends/webgpu/test/op_tests/test_generator.py +++ b/backends/webgpu/test/op_tests/test_generator.py @@ -28,7 +28,10 @@ def test_export_case_has_delegate(): def test_generate_case_writes_artifacts(tmp_path): suite, case = _add_regular_case() - entry = g.generate_case("add", suite, case, str(tmp_path)) + # generate_case returns one entry per output; add is single-output. + entries = g.generate_case("add", suite, case, str(tmp_path)) + assert len(entries) == 1 + entry = entries[0] # .pte + 2 input .bin + golden .bin all exist assert (tmp_path / entry["pte"]).exists() assert len(entry["inputs"]) == 2 From 22855e956197302694acbfa2027521d62ad4ecfe Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:18 -0700 Subject: [PATCH 13/75] [ExecuTorch][WebGPU] Port index_select MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21171 Ports `aten.index_select.default` as a gather along a dim, mirroring the landed `flip`/`split` index-transform (dual-`TensorMeta`, 2D-fold dispatch) + the landed `index/Index.cpp` int-index binding. index_select is pure data movement, so it remaps the dim coord through an int index and copies — no fp compute. Key changes: - `runtime/ops/index_select/{IndexSelect.cpp, index_select.wgsl}` (+ generated `index_select_wgsl.h`) — per out flat idx: unravel via `out_meta.strides`, `in_coord = out_coord` with the dim coord remapped by `index[out_coord[dim]]`, ravel via `in_meta.strides`, `output[out] = input[in]`; the index tensor is bound as `array` (the int32 `downcast_64_bit` of the int64 index, same as `Index.cpp`); `Params { info: vec4 }` carries `dim`. Registered `aten.index_select.default`. - `CMakeLists.txt` — `runtime/ops/index_select/IndexSelect.cpp` in `WEBGPU_SRCS`. Buffer re-derivation of the gather (Vulkan `IndexSelect.cpp` is texture/channel-packed only — `op_registry.py` tags it `CHANNELS_PACKED_TEXTURE`; this backend is buffer-only, so the NCHW index math is derived and CPU-de-risked, cf. `split.wgsl`). Index-value bounds are trusted/unclamped, matching Vulkan `index_select.glsl`; a negative index is a memory-safe WGSL OOB read (no trap), same safety class as Vulkan's clamped `texelFetch`. No int dtype-guard on self/out (bit-preserving); `nbytes` guards reject non-fp32 self/out and a non-int32 index. ghstack-source-id: 406360861 @exported-using-ghexport Differential Revision: [D112257660](https://our.internmc.facebook.com/intern/diff/D112257660/) --- .../runtime/ops/index_select/IndexSelect.cpp | 213 ++++++++++++++++++ .../ops/index_select/index_select.wgsl | 45 ++++ .../ops/index_select/index_select_wgsl.h | 69 ++++++ 3 files changed, 327 insertions(+) create mode 100644 backends/webgpu/runtime/ops/index_select/IndexSelect.cpp create mode 100644 backends/webgpu/runtime/ops/index_select/index_select.wgsl create mode 100644 backends/webgpu/runtime/ops/index_select/index_select_wgsl.h diff --git a/backends/webgpu/runtime/ops/index_select/IndexSelect.cpp b/backends/webgpu/runtime/ops/index_select/IndexSelect.cpp new file mode 100644 index 00000000000..dda307029d5 --- /dev/null +++ b/backends/webgpu/runtime/ops/index_select/IndexSelect.cpp @@ -0,0 +1,213 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct IndexSelectParams { + uint32_t info[4]; // info[0] = dim +}; +static_assert( + sizeof(IndexSelectParams) == 16, + "IndexSelectParams must match the WGSL Params vec4 (16 bytes)"); + +// index_select: gather rows along dim via an int index (Vulkan IndexSelect). +void index_select_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [self, dim, index, out]; index is an int32 tensor (downcast_64_bit). + const int self_id = args.at(0); + const int dim_id = args.at(1); + const int index_id = args.at(2); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(self_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(index_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("index_select: self/index/out is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& self_tensor = graph.get_tensor(self_id); + const auto& index_tensor = graph.get_tensor(index_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (self_tensor.buffer == nullptr || index_tensor.buffer == nullptr || + out_tensor.buffer == nullptr) { + throw std::runtime_error("index_select: null buffer binding"); + } + + const int64_t ndim = static_cast(self_tensor.dims.size()); + if (ndim > static_cast(kTensorMetaMaxNdim)) { + throw std::runtime_error("index_select: tensor rank exceeds 4 (MAX_NDIM)"); + } + + if (graph.get_value_type(dim_id) != WebGPUGraph::ValueType::Int) { + throw std::runtime_error("index_select: dim arg is not a static Int"); + } + int64_t dim = graph.get_int(dim_id); + if (dim < 0) { + dim += ndim; + } + if (dim < 0 || dim >= ndim) { + throw std::runtime_error("index_select: dim out of range"); + } + + TensorMeta out_meta; + TensorMeta in_meta; + fill_tensor_meta(out_tensor, &out_meta); + fill_tensor_meta(self_tensor, &in_meta); + if (out_tensor.nbytes != + static_cast(out_meta.numel) * sizeof(float) || + self_tensor.nbytes % sizeof(float) != 0) { + throw std::runtime_error("index_select: non-fp32 self/out (nbytes % 4)"); + } + // Index is the int32 downcast of the int64 index (downcast_64_bit). + if (index_tensor.nbytes % sizeof(int32_t) != 0) { + throw std::runtime_error("index_select: index buffer is not int32"); + } + // The index gathers out.dims[dim] rows (one per index element), so it must + // hold at least that many entries. + uint64_t index_numel = 1; + for (int64_t d : index_tensor.dims) { + index_numel *= static_cast(d); + } + if (index_numel < static_cast(out_tensor.dims.at(dim))) { + throw std::runtime_error("index_select: index numel < out.dims[dim]"); + } + + IndexSelectParams params = {}; + params.info[0] = static_cast(dim); + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kIndexSelectWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, out_meta.numel, wg_size, "index_select"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer out_meta_buf = + utils::make_uniform(device, &out_meta, sizeof(TensorMeta)); + WGPUBuffer in_meta_buf = + utils::make_uniform(device, &in_meta, sizeof(TensorMeta)); + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(IndexSelectParams)); + graph.add_uniform_buffer_bytes( + 2 * sizeof(TensorMeta) + sizeof(IndexSelectParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kIndexSelectWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // in, out (rw), index (read i32), out_meta, in_meta, params (3 uniforms). + WGPUBindGroupLayoutEntry entries[6] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[3].binding = 3; + entries[3].visibility = WGPUShaderStage_Compute; + entries[3].buffer.type = WGPUBufferBindingType_Uniform; + entries[4].binding = 4; + entries[4].visibility = WGPUShaderStage_Compute; + entries[4].buffer.type = WGPUBufferBindingType_Uniform; + entries[5].binding = 5; + entries[5].visibility = WGPUShaderStage_Compute; + entries[5].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 6; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[6] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = self_tensor.buffer; + bg_entries[0].size = self_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = index_tensor.buffer; + bg_entries[2].size = index_tensor.nbytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = out_meta_buf; + bg_entries[3].size = sizeof(TensorMeta); + bg_entries[4].binding = 4; + bg_entries[4].buffer = in_meta_buf; + bg_entries[4].size = sizeof(TensorMeta); + bg_entries[5].binding = 5; + bg_entries[5].buffer = params_buf; + bg_entries[5].size = sizeof(IndexSelectParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 6; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + // Static shapes only: index_select registers no resize hook, so the output + // extent (out.dims[dim] == index numel) is fixed at build time. + graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "index_select", + workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Drop our refs; the bind group keeps the uniforms alive until release. + wgpuBufferRelease(out_meta_buf); + wgpuBufferRelease(in_meta_buf); + wgpuBufferRelease(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.index_select.default, index_select_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/index_select/index_select.wgsl b/backends/webgpu/runtime/ops/index_select/index_select.wgsl new file mode 100644 index 00000000000..79b977e7f7c --- /dev/null +++ b/backends/webgpu/runtime/ops/index_select/index_select.wgsl @@ -0,0 +1,45 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var index: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in_meta: TensorMeta; + +struct Params { + info: vec4, +} +@group(0) @binding(5) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (out_bufi >= out_meta.numel) { + return; + } + + // Gather: in_coord = out_coord, but the dim coord is remapped by index[]. + let dim = params.info.x; + var rem = out_bufi; + var in_bufi: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d]; + rem = rem % out_meta.strides[d]; + var in_coord = coord; + if (d == dim) { + in_coord = u32(index[coord]); + } + in_bufi = in_bufi + in_coord * in_meta.strides[d]; + } + output[out_bufi] = input[in_bufi]; +} diff --git a/backends/webgpu/runtime/ops/index_select/index_select_wgsl.h b/backends/webgpu/runtime/ops/index_select/index_select_wgsl.h new file mode 100644 index 00000000000..6e83be3883b --- /dev/null +++ b/backends/webgpu/runtime/ops/index_select/index_select_wgsl.h @@ -0,0 +1,69 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from index_select.wgsl - DO NOT EDIT. +// wgsl-sha256: 63d3a49ddd72a67b05bb3d03a05ba375e4e3b096a47ec67406d05c3ddc9eb241 +inline constexpr const char* kIndexSelectWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var index: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in_meta: TensorMeta; + +struct Params { + info: vec4, +} +@group(0) @binding(5) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (out_bufi >= out_meta.numel) { + return; + } + + // Gather: in_coord = out_coord, but the dim coord is remapped by index[]. + let dim = params.info.x; + var rem = out_bufi; + var in_bufi: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d]; + rem = rem % out_meta.strides[d]; + var in_coord = coord; + if (d == dim) { + in_coord = u32(index[coord]); + } + in_bufi = in_bufi + in_coord * in_meta.strides[d]; + } + output[out_bufi] = input[in_bufi]; +} +)"; + +inline constexpr uint32_t kIndexSelectWorkgroupSizeX = 64; +inline constexpr uint32_t kIndexSelectWorkgroupSizeY = 1; +inline constexpr uint32_t kIndexSelectWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 8ecda196207ee0dba3d5da90ee61f8d5626a4ba0 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:18 -0700 Subject: [PATCH 14/75] [ExecuTorch][WebGPU] Op-tests for index_select MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21172 Adds the op-test suite for the `index_select` op ported in the diff below. Key changes: - `test/ops/test_index_select.py` — `IndexSelectModule` gathers via a baked int index buffer, so the only runtime input is the float tensor (the op-test framework is float-only); `torch.index_select(x, dim, self.index)`. + `IndexSelectTest` delegation smoke test (mirrors `SplitTest`). - `test/op_tests/cases.py` — `_index_select_suite`: dim0_1d, dim0_2d, dim1_2d (repeated index), dim1_3d (middle dim), dim2_3d. ghstack-source-id: 406360855 @exported-using-ghexport Differential Revision: [D112257648](https://our.internmc.facebook.com/intern/diff/D112257648/) --- backends/webgpu/test/op_tests/cases.py | 38 +++++++++ backends/webgpu/test/ops/test_index_select.py | 79 +++++++++++++++++++ 2 files changed, 117 insertions(+) create mode 100644 backends/webgpu/test/ops/test_index_select.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 64663de7257..2bb04164c6c 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -36,6 +36,7 @@ ) from executorch.backends.webgpu.test.ops.test_flip import FlipModule from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule +from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule from executorch.backends.webgpu.test.ops.test_mul import ( CONFIGS as _MUL_CONFIGS, @@ -306,6 +307,43 @@ def _repeat_suite() -> WebGPUTestSuite: ) +@register_op_test("index_select") +def _index_select_suite() -> WebGPUTestSuite: + # Gather rows along dim via a baked int index; float32 oracle. Only the float + # input is a runtime tensor (the index is a graph constant). + return WebGPUTestSuite( + module_factory=lambda dim, index: IndexSelectModule(dim, index), + cases=[ + Case( + name="dim0_1d", + construct={"dim": 0, "index": [0, 2, 4, 1]}, + inputs=((S,),), + ), + Case( + name="dim0_2d", + construct={"dim": 0, "index": [3, 0, 1]}, + inputs=((XS + 1, S1),), + ), + Case( + name="dim1_2d", + construct={"dim": 1, "index": [5, 2, 0, 2]}, + inputs=((XS + 1, S1),), + ), + Case( + name="dim1_3d", + construct={"dim": 1, "index": [4, 0, 2]}, + inputs=((XS, S, S1),), + ), + Case( + name="dim2_3d", + construct={"dim": 2, "index": [6, 1, 4]}, + inputs=((XS, S, S1),), + ), + ], + golden_dtype="float32", + ) + + @register_op_test("view_copy") def _view_copy_suite() -> WebGPUTestSuite: return _fn_config_suite(ViewModule, _VIEW_CONFIGS) diff --git a/backends/webgpu/test/ops/test_index_select.py b/backends/webgpu/test/ops/test_index_select.py new file mode 100644 index 00000000000..128f76274b5 --- /dev/null +++ b/backends/webgpu/test/ops/test_index_select.py @@ -0,0 +1,79 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.index_select.default` module + configs for the WebGPU op-test framework. + +`IndexSelectModule` gathers rows along a dim via a baked int index buffer, so the +only runtime input is the float tensor (the op-test framework is float-only). +index_select is pure data movement (bit-identical) -> float32 oracle. `IndexSelectTest` +is the export-delegation smoke test. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# name -> (input_shape, dim, index) +CONFIGS = { + "dim0_1d": ((5,), 0, [0, 2, 4, 1]), + "dim0_2d": ((4, 7), 0, [3, 0, 1]), + "dim1_2d": ((4, 7), 1, [5, 2, 0, 2]), + "dim1_3d": ((3, 5, 7), 1, [4, 0, 2]), + "dim2_3d": ((3, 5, 7), 2, [6, 1, 4]), +} + + +class IndexSelectModule(torch.nn.Module): + def __init__(self, dim, index) -> None: + super().__init__() + self.dim = dim + self.register_buffer("index", torch.tensor(index, dtype=torch.int64)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.index_select(x, self.dim, self.index) + + +def _det_input(shape): + g = torch.Generator().manual_seed(0) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(dim, index, x: torch.Tensor): + ep = torch.export.export(IndexSelectModule(dim, index).eval(), (x,)) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + +def _delegated(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +def _op_delegated(edge, op_substr: str) -> bool: + # op must be absorbed into the delegate, not left as a top-level CPU-fallback node. + gm = edge.exported_program().graph_module + return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + + +class IndexSelectTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, (shape, dim, index) in CONFIGS.items(): + with self.subTest(name=name): + edge = _lower(dim, index, _det_input(shape)) + et = edge.to_executorch() + self.assertTrue( + _delegated(et), + f"Expected a VulkanBackend delegate (index_select {name})", + ) + self.assertTrue( + _op_delegated(edge, "index_select"), + f"index_select not delegated (fell back to CPU) for {name}", + ) From ac29beda01ec435dbaca835f39febafaa9d68784 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:18 -0700 Subject: [PATCH 15/75] [ExecuTorch][WebGPU] Port native_group_norm MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21173 Ports `aten.native_group_norm.default` as a 2-pass, 3-output group normalization, the second Phase B (vision) op. Pass 1 reduces per group to `mean`/`rstd`; pass 2 applies the per-channel affine. Key changes: - `runtime/ops/native_group_norm/{GroupNorm.cpp, group_norm_reduce.wgsl, group_norm.wgsl}` (+ generated `_wgsl.h`) — reduce: one thread per `(n, group)` serially scans that group's contiguous NCHW block (`base = n*C*HxW + g*(C/G)*HxW`, `group_size = (C/G)*HxW` elements), writing `mean = s/count`, `rstd = inverseSqrt(ss/count - mean^2 + eps)`; normalize: one thread per element, `g = c / (C/G)`, `out = (x - mean[n*G+g]) * rstd[n*G+g] * weight[c] + bias[c]`. A file-local `add_gn_dispatch` helper builds each of the two dispatches from an explicit binding list sharing one params UBO. Registered `aten.native_group_norm.default`. - `CMakeLists.txt` — `runtime/ops/native_group_norm/GroupNorm.cpp` in `WEBGPU_SRCS`. Mirrors Vulkan `impl/GroupNorm.cpp` (2 nodes: reduce -> mean/rstd, then normalize) + `glsl/group_norm_reduce_texture.glsl` (`mean=sum/count`, `variance=sumsq/count - mean^2`, `rstd=1/sqrt(var+eps)`). Buffer NCHW re-derivation (Vulkan is texture/channels-packed), so a serial per-group reduce (a group's channels x HxW are contiguous in NCHW) replaces the cooperative texture reduction. All 3 outputs are real. The inter-pass RAW ordering (normalize reads the reduce's mean/rstd) is guaranteed because `execute()` runs one compute pass per dispatch (`WebGPUGraph.cpp`; proven by `test_dispatch_order.cpp` to 1M elems). A dynamic-shape resize hook recomputes params + both dispatch counts + propagates cur_dims for out/mean/rstd. Fail-loud guards: 4D, `C % group == 0`, fp32, weight/bias len `== C`, mean/rstd size `== N*group`. ghstack-source-id: 406360859 @exported-using-ghexport Differential Revision: [D112257596](https://our.internmc.facebook.com/intern/diff/D112257596/) --- .../ops/native_group_norm/GroupNorm.cpp | 296 ++++++++++++++++++ .../ops/native_group_norm/group_norm.wgsl | 38 +++ .../native_group_norm/group_norm_reduce.wgsl | 47 +++ .../group_norm_reduce_wgsl.h | 71 +++++ .../ops/native_group_norm/group_norm_wgsl.h | 62 ++++ 5 files changed, 514 insertions(+) create mode 100644 backends/webgpu/runtime/ops/native_group_norm/GroupNorm.cpp create mode 100644 backends/webgpu/runtime/ops/native_group_norm/group_norm.wgsl create mode 100644 backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce.wgsl create mode 100644 backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce_wgsl.h create mode 100644 backends/webgpu/runtime/ops/native_group_norm/group_norm_wgsl.h diff --git a/backends/webgpu/runtime/ops/native_group_norm/GroupNorm.cpp b/backends/webgpu/runtime/ops/native_group_norm/GroupNorm.cpp new file mode 100644 index 00000000000..71ec4a22a43 --- /dev/null +++ b/backends/webgpu/runtime/ops/native_group_norm/GroupNorm.cpp @@ -0,0 +1,296 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct GroupNormParams { + uint32_t n_channels; + uint32_t hxw; + uint32_t num_groups; + uint32_t chans_per_group; + uint32_t numel; + uint32_t mean_numel; + uint32_t group_size; + float eps; +}; +static_assert( + sizeof(GroupNormParams) == 32, + "GroupNormParams must match the WGSL Params struct (32 bytes)"); + +struct GnBinding { + WGPUBuffer buffer; + uint64_t size; + WGPUBufferBindingType type; +}; + +// Build one compute dispatch from a binding list (last = uniform). +size_t add_gn_dispatch( + WebGPUGraph& graph, + WGPUDevice device, + const char* wgsl_code, + uint32_t wg_size, + const std::vector& binds, + utils::WgCount wgc, + const char* label) { + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {wgsl_code, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + std::vector entries(binds.size()); + for (size_t i = 0; i < binds.size(); i++) { + entries[i].binding = static_cast(i); + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = binds[i].type; + } + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = entries.size(); + bgl_desc.entries = entries.data(); + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + std::vector bg_entries(binds.size()); + for (size_t i = 0; i < binds.size(); i++) { + bg_entries[i].binding = static_cast(i); + bg_entries[i].buffer = binds[i].buffer; + bg_entries[i].size = binds[i].size; + } + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = bg_entries.size(); + bg_desc.entries = bg_entries.data(); + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t idx = + graph.add_dispatch({pipeline, bind_group, wgc.x, label, wgc.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + return idx; +} + +// native_group_norm: per-group reduce + per-channel norm (Vulkan GroupNorm). +void native_group_norm_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [in, weight, bias, N, C, HxW, group, eps, out_tuple]. + const int in_id = args.at(0); + const int weight_id = args.at(1); + const int bias_id = args.at(2); + const int group_id = args.at(6); + const int eps_id = args.at(7); + const int out_list_id = args.at(8); + + if (graph.get_value_type(out_list_id) != WebGPUGraph::ValueType::ValueList) { + throw std::runtime_error("group_norm: out arg is not a value list"); + } + const std::vector& out_ids = graph.get_value_list(out_list_id); + if (out_ids.size() != 3) { + throw std::runtime_error("group_norm: expected 3 outputs (out/mean/rstd)"); + } + const int out_id = out_ids.at(0); + const int mean_id = out_ids.at(1); + const int rstd_id = out_ids.at(2); + + const int tensor_ids[6] = { + in_id, weight_id, bias_id, out_id, mean_id, rstd_id}; + for (int id : tensor_ids) { + if (graph.get_value_type(id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("group_norm: in/weight/bias/out not a tensor"); + } + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + const auto& weight_tensor = graph.get_tensor(weight_id); + const auto& bias_tensor = graph.get_tensor(bias_id); + const auto& mean_tensor = graph.get_tensor(mean_id); + const auto& rstd_tensor = graph.get_tensor(rstd_id); + + if (in_tensor.dims.size() != 4) { + throw std::runtime_error("group_norm: only 4D (NCHW) input is supported"); + } + const uint32_t n_batch = static_cast(in_tensor.dims.at(0)); + const uint32_t n_channels = static_cast(in_tensor.dims.at(1)); + const uint32_t hxw = + static_cast(in_tensor.dims.at(2) * in_tensor.dims.at(3)); + const int64_t group = graph.get_int(group_id); + if (group <= 0 || n_channels % static_cast(group) != 0) { + throw std::runtime_error("group_norm: C not divisible by group"); + } + const uint32_t num_groups = static_cast(group); + const uint32_t chans_per_group = n_channels / num_groups; + + const uint64_t numel = static_cast(n_batch) * n_channels * hxw; + const uint32_t mean_numel = n_batch * num_groups; + if (in_tensor.nbytes != numel * sizeof(float) || + out_tensor.nbytes != numel * sizeof(float)) { + throw std::runtime_error("group_norm: fp32-only (byte-size mismatch)"); + } + const size_t chan_bytes = static_cast(n_channels) * sizeof(float); + if (weight_tensor.nbytes != chan_bytes || bias_tensor.nbytes != chan_bytes) { + throw std::runtime_error("group_norm: weight/bias length != num_channels"); + } + const size_t mean_bytes = static_cast(mean_numel) * sizeof(float); + if (mean_tensor.nbytes != mean_bytes || rstd_tensor.nbytes != mean_bytes) { + throw std::runtime_error("group_norm: mean/rstd size != N * group"); + } + + float eps = std::numeric_limits::epsilon(); + if (graph.get_value_type(eps_id) == WebGPUGraph::ValueType::Double) { + eps = static_cast(graph.get_double(eps_id)); + } else if (graph.get_value_type(eps_id) == WebGPUGraph::ValueType::Int) { + eps = static_cast(graph.get_int(eps_id)); + } + + GroupNormParams params = {}; + params.n_channels = n_channels; + params.hxw = hxw; + params.num_groups = num_groups; + params.chans_per_group = chans_per_group; + params.numel = static_cast(numel); + params.mean_numel = mean_numel; + params.group_size = chans_per_group * hxw; + params.eps = eps; + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kGroupNormWorkgroupSizeX); + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(GroupNormParams)); + graph.add_uniform_buffer_bytes(sizeof(GroupNormParams)); + + const WGPUBufferBindingType ro = WGPUBufferBindingType_ReadOnlyStorage; + const WGPUBufferBindingType rw = WGPUBufferBindingType_Storage; + const WGPUBufferBindingType uni = WGPUBufferBindingType_Uniform; + + // Pass 1: reduce -> mean/rstd (one thread per (n, group)). + utils::WgCount reduce_wgc = utils::compute_2d_workgroup_count( + device, mean_numel, wg_size, "gn_reduce"); + const size_t reduce_idx = add_gn_dispatch( + graph, + device, + kGroupNormReduceWGSL, + wg_size, + {{in_tensor.buffer, in_tensor.nbytes, ro}, + {mean_tensor.buffer, mean_tensor.nbytes, rw}, + {rstd_tensor.buffer, rstd_tensor.nbytes, rw}, + {params_buf, sizeof(GroupNormParams), uni}}, + reduce_wgc, + "group_norm_reduce"); + + // Pass 2: normalize (execute() runs one pass/dispatch, so reduce precedes). + utils::WgCount norm_wgc = utils::compute_2d_workgroup_count( + device, static_cast(numel), wg_size, "gn_norm"); + const size_t norm_idx = add_gn_dispatch( + graph, + device, + kGroupNormWGSL, + wg_size, + {{in_tensor.buffer, in_tensor.nbytes, ro}, + {out_tensor.buffer, out_tensor.nbytes, rw}, + {weight_tensor.buffer, weight_tensor.nbytes, ro}, + {bias_tensor.buffer, bias_tensor.nbytes, ro}, + {mean_tensor.buffer, mean_tensor.nbytes, ro}, + {rstd_tensor.buffer, rstd_tensor.nbytes, ro}, + {params_buf, sizeof(GroupNormParams), uni}}, + norm_wgc, + "group_norm"); + + // Dynamic shapes: recompute params + both dispatch counts from the live dims. + WGPUBuffer p_buf = params_buf; + graph.add_tensor_resize_hook( + in_id, + [in_id, + out_id, + mean_id, + rstd_id, + num_groups, + eps, + wg_size, + reduce_idx, + norm_idx, + p_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + if (d.size() != 4) { + throw std::runtime_error("group_norm(resize): input is not 4D"); + } + const uint32_t nb = static_cast(d[0]); + const uint32_t c = static_cast(d[1]); + const uint32_t hw = static_cast(d[2] * d[3]); + const uint32_t dpg = c / num_groups; + GroupNormParams p = {}; + p.n_channels = c; + p.hxw = hw; + p.num_groups = num_groups; + p.chans_per_group = dpg; + p.numel = nb * c * hw; + p.mean_numel = nb * num_groups; + p.group_size = dpg * hw; + p.eps = eps; + wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p)); + const utils::WgCount rwgc = utils::compute_2d_workgroup_count( + g.device(), p.mean_numel, wg_size, "gn_reduce(resize)"); + g.dispatch_at(reduce_idx).workgroup_count_x = rwgc.x; + g.dispatch_at(reduce_idx).workgroup_count_y = rwgc.y; + const utils::WgCount nwgc = utils::compute_2d_workgroup_count( + g.device(), p.numel, wg_size, "gn_norm(resize)"); + g.dispatch_at(norm_idx).workgroup_count_x = nwgc.x; + g.dispatch_at(norm_idx).workgroup_count_y = nwgc.y; + g.set_cur_dims(out_id, d); + const std::vector mr = { + d[0], static_cast(num_groups)}; + g.set_cur_dims(mean_id, mr); + g.set_cur_dims(rstd_id, mr); + }); + + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.native_group_norm.default, native_group_norm_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/native_group_norm/group_norm.wgsl b/backends/webgpu/runtime/ops/native_group_norm/group_norm.wgsl new file mode 100644 index 00000000000..22b7a09eb51 --- /dev/null +++ b/backends/webgpu/runtime/ops/native_group_norm/group_norm.wgsl @@ -0,0 +1,38 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var mean: array; +@group(0) @binding(5) var rstd: array; + +struct Params { + n_channels: u32, + hxw: u32, + num_groups: u32, + chans_per_group: u32, + numel: u32, + mean_numel: u32, + group_size: u32, + eps: f32, +} +@group(0) @binding(6) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + + // NCHW channel -> its group; apply mean/rstd + per-channel affine. + let n = idx / (params.n_channels * params.hxw); + let c = (idx / params.hxw) % params.n_channels; + let g = c / params.chans_per_group; + let mg = n * params.num_groups + g; + output[idx] = (input[idx] - mean[mg]) * rstd[mg] * weight[c] + bias[c]; +} diff --git a/backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce.wgsl b/backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce.wgsl new file mode 100644 index 00000000000..6b0569dc213 --- /dev/null +++ b/backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce.wgsl @@ -0,0 +1,47 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var mean: array; +@group(0) @binding(2) var rstd: array; + +struct Params { + n_channels: u32, + hxw: u32, + num_groups: u32, + chans_per_group: u32, + numel: u32, + mean_numel: u32, + group_size: u32, + eps: f32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per (n, group) -> mean/rstd of shape [N, G]. + let mg = gid.x + gid.y * (num_workgroups.x * wg_size); + if (mg >= params.mean_numel) { + return; + } + + // A group's D channels x HxW form a contiguous NCHW block. + let n = mg / params.num_groups; + let g = mg % params.num_groups; + let base = n * params.n_channels * params.hxw + + g * params.chans_per_group * params.hxw; + + var s = 0.0; + var ss = 0.0; + for (var i = 0u; i < params.group_size; i = i + 1u) { + let v = input[base + i]; + s = s + v; + ss = ss + v * v; + } + let count = f32(params.group_size); + let m = s / count; + let variance = ss / count - m * m; + mean[mg] = m; + rstd[mg] = inverseSqrt(variance + params.eps); +} diff --git a/backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce_wgsl.h b/backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce_wgsl.h new file mode 100644 index 00000000000..77704092a20 --- /dev/null +++ b/backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce_wgsl.h @@ -0,0 +1,71 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from group_norm_reduce.wgsl - DO NOT EDIT. +// wgsl-sha256: 63e9e15a934f19ce0649323ae3820f1cd86deb9aeb0dbb4ba1de9ab4d7d74688 +inline constexpr const char* kGroupNormReduceWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var mean: array; +@group(0) @binding(2) var rstd: array; + +struct Params { + n_channels: u32, + hxw: u32, + num_groups: u32, + chans_per_group: u32, + numel: u32, + mean_numel: u32, + group_size: u32, + eps: f32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per (n, group) -> mean/rstd of shape [N, G]. + let mg = gid.x + gid.y * (num_workgroups.x * wg_size); + if (mg >= params.mean_numel) { + return; + } + + // A group's D channels x HxW form a contiguous NCHW block. + let n = mg / params.num_groups; + let g = mg % params.num_groups; + let base = n * params.n_channels * params.hxw + + g * params.chans_per_group * params.hxw; + + var s = 0.0; + var ss = 0.0; + for (var i = 0u; i < params.group_size; i = i + 1u) { + let v = input[base + i]; + s = s + v; + ss = ss + v * v; + } + let count = f32(params.group_size); + let m = s / count; + let variance = ss / count - m * m; + mean[mg] = m; + rstd[mg] = inverseSqrt(variance + params.eps); +} +)"; + +inline constexpr uint32_t kGroupNormReduceWorkgroupSizeX = 64; +inline constexpr uint32_t kGroupNormReduceWorkgroupSizeY = 1; +inline constexpr uint32_t kGroupNormReduceWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/native_group_norm/group_norm_wgsl.h b/backends/webgpu/runtime/ops/native_group_norm/group_norm_wgsl.h new file mode 100644 index 00000000000..1a963387379 --- /dev/null +++ b/backends/webgpu/runtime/ops/native_group_norm/group_norm_wgsl.h @@ -0,0 +1,62 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from group_norm.wgsl - DO NOT EDIT. +// wgsl-sha256: 911ba488b46495c887637a36d813cb81c639295881873c921c5939bdb1c397aa +inline constexpr const char* kGroupNormWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var mean: array; +@group(0) @binding(5) var rstd: array; + +struct Params { + n_channels: u32, + hxw: u32, + num_groups: u32, + chans_per_group: u32, + numel: u32, + mean_numel: u32, + group_size: u32, + eps: f32, +} +@group(0) @binding(6) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + + // NCHW channel -> its group; apply mean/rstd + per-channel affine. + let n = idx / (params.n_channels * params.hxw); + let c = (idx / params.hxw) % params.n_channels; + let g = c / params.chans_per_group; + let mg = n * params.num_groups + g; + output[idx] = (input[idx] - mean[mg]) * rstd[mg] * weight[c] + bias[c]; +} +)"; + +inline constexpr uint32_t kGroupNormWorkgroupSizeX = 64; +inline constexpr uint32_t kGroupNormWorkgroupSizeY = 1; +inline constexpr uint32_t kGroupNormWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From bdb95b93bdfd452e0f3805dbea4ae6fe7720817d Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:19 -0700 Subject: [PATCH 16/75] [ExecuTorch][WebGPU] Op-tests for native_group_norm MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21174 Adds the op-test suite for the `native_group_norm` op ported in the diff below, using the multi-output golden support to verify all 3 outputs. Key changes: - `test/ops/test_group_norm.py` — `GroupNormModule` returns the full `torch.native_group_norm` (out, mean, rstd) tuple, so the multi-output golden path checks both the reduce (mean/rstd) and normalize (out) passes. + `GroupNormTest` delegation smoke test. - `test/op_tests/cases.py` — `_group_norm_suite`: c4_g2, c6_g3, c4_g1. ghstack-source-id: 406360848 @exported-using-ghexport Differential Revision: [D112257626](https://our.internmc.facebook.com/intern/diff/D112257626/) --- backends/webgpu/test/op_tests/cases.py | 31 ++++++++ backends/webgpu/test/ops/test_group_norm.py | 83 +++++++++++++++++++++ 2 files changed, 114 insertions(+) create mode 100644 backends/webgpu/test/ops/test_group_norm.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 2bb04164c6c..3bd9a0358d3 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -36,6 +36,7 @@ ) from executorch.backends.webgpu.test.ops.test_flip import FlipModule from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule +from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule from executorch.backends.webgpu.test.ops.test_mul import ( @@ -307,6 +308,36 @@ def _repeat_suite() -> WebGPUTestSuite: ) +@register_op_test("native_group_norm") +def _group_norm_suite() -> WebGPUTestSuite: + # 2-pass norm returning (out, mean, rstd); the multi-output golden verifies + # both the reduce (mean/rstd) and normalize (out) passes. float64 oracle. + return WebGPUTestSuite( + module_factory=lambda num_channels, num_groups: GroupNormModule( + num_channels, num_groups + ), + cases=[ + Case( + name="c4_g2", + construct={"num_channels": 4, "num_groups": 2}, + inputs=((2, 4, 3, 5),), + ), + Case( + name="c6_g3", + construct={"num_channels": 6, "num_groups": 3}, + inputs=((1, 6, 2, 2),), + ), + Case( + name="c4_g1", + construct={"num_channels": 4, "num_groups": 1}, + inputs=((1, 4, 3, 3),), + ), + ], + atol=1e-3, + rtol=1e-3, + ) + + @register_op_test("index_select") def _index_select_suite() -> WebGPUTestSuite: # Gather rows along dim via a baked int index; float32 oracle. Only the float diff --git a/backends/webgpu/test/ops/test_group_norm.py b/backends/webgpu/test/ops/test_group_norm.py new file mode 100644 index 00000000000..2b51209047e --- /dev/null +++ b/backends/webgpu/test/ops/test_group_norm.py @@ -0,0 +1,83 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.native_group_norm.default` module + configs for the op-test framework. + +`GroupNormModule` returns the full (out, mean, rstd) tuple via +`torch.native_group_norm`, so the multi-output golden path verifies BOTH the +reduce pass (mean/rstd) and the normalize pass (out). `GroupNormTest` is the +export-delegation smoke test. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# name -> num_channels, num_groups, input_shape +CONFIGS = { + "c4_g2": (4, 2, (2, 4, 3, 5)), + "c6_g3": (6, 3, (1, 6, 2, 2)), + "c4_g1": (4, 1, (1, 4, 3, 3)), +} + + +class GroupNormModule(torch.nn.Module): + def __init__(self, num_channels, num_groups) -> None: + super().__init__() + self.num_groups = num_groups + g = torch.Generator().manual_seed(0) + self.register_buffer("w", torch.randn(num_channels, generator=g)) + self.register_buffer("b", torch.randn(num_channels, generator=g)) + self.eps = 1e-5 + + def forward(self, x: torch.Tensor): + n, c, h, w = x.shape + return torch.native_group_norm( + x, self.w, self.b, n, c, h * w, self.num_groups, self.eps + ) + + +def _det_input(shape): + g = torch.Generator().manual_seed(1) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(num_channels, num_groups, x: torch.Tensor): + ep = torch.export.export(GroupNormModule(num_channels, num_groups).eval(), (x,)) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + +def _delegated(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +def _op_delegated(edge, op_substr: str) -> bool: + # op must be absorbed into the delegate, not left as a top-level CPU-fallback node. + gm = edge.exported_program().graph_module + return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + + +class GroupNormTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, (num_channels, num_groups, shape) in CONFIGS.items(): + with self.subTest(name=name): + edge = _lower(num_channels, num_groups, _det_input(shape)) + et = edge.to_executorch() + self.assertTrue( + _delegated(et), + f"Expected a VulkanBackend delegate (group_norm {name})", + ) + self.assertTrue( + _op_delegated(edge, "native_group_norm"), + f"group_norm not delegated (fell back to CPU) for {name}", + ) From e237fd13439285a0acd32e7dfbc06583215bc094 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:19 -0700 Subject: [PATCH 17/75] [ExecuTorch][WebGPU] Port avg_pool2d MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21175 Ports `aten.avg_pool2d.default` as a windowed spatial average, a Phase B (vision) op. Each output cell averages its KxK input window. Key changes: - `runtime/ops/avg_pool2d/{AvgPool2d.cpp, avg_pool2d.wgsl}` (+ generated `avg_pool2d_wgsl.h`) — per NCHW output element: `ipos = out*stride - padding`, window `[max(0,ipos), min(ipos+kernel, in))`, `acc = sum(window)`, `out = acc / divisor`. The divisor mirrors the three ATen cases: `divisor_override` if given, else `count_include_pad` counts the padded window minus any overhang past the padded edge, else the actual number of summed cells. Registered `aten.avg_pool2d.default`. - `CMakeLists.txt` — `runtime/ops/avg_pool2d/AvgPool2d.cpp` in `WEBGPU_SRCS`. Mirrors Vulkan `glsl/avg_pool2d.glsl` (`ipos`/`start`/`end`/`sum`/`empty`-overhang divisor) + `impl/Pool.cpp` (arg order, `divisor_override` read only when Int, `count_include_pad` bool). Buffer NCHW re-derivation (Vulkan is texture/channels-packed). `read_pair` normalizes each IntList param (empty stride -> kernel, len-1 broadcast, len-2). A dynamic-shape resize hook recomputes the pooled output extent via `pool_out_dim` (floor, or ceil + last-window-drop for `ceil_mode`) + params + dispatch + cur_dims. Fail-loud guards: 4D, fp32. ghstack-source-id: 406360851 @exported-using-ghexport Differential Revision: [D112257592](https://our.internmc.facebook.com/intern/diff/D112257592/) --- .../runtime/ops/avg_pool2d/AvgPool2d.cpp | 296 ++++++++++++++++++ .../runtime/ops/avg_pool2d/avg_pool2d.wgsl | 74 +++++ .../runtime/ops/avg_pool2d/avg_pool2d_wgsl.h | 98 ++++++ 3 files changed, 468 insertions(+) create mode 100644 backends/webgpu/runtime/ops/avg_pool2d/AvgPool2d.cpp create mode 100644 backends/webgpu/runtime/ops/avg_pool2d/avg_pool2d.wgsl create mode 100644 backends/webgpu/runtime/ops/avg_pool2d/avg_pool2d_wgsl.h diff --git a/backends/webgpu/runtime/ops/avg_pool2d/AvgPool2d.cpp b/backends/webgpu/runtime/ops/avg_pool2d/AvgPool2d.cpp new file mode 100644 index 00000000000..a172bd71ec0 --- /dev/null +++ b/backends/webgpu/runtime/ops/avg_pool2d/AvgPool2d.cpp @@ -0,0 +1,296 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct AvgPoolParams { + uint32_t kh; + uint32_t kw; + uint32_t sh; + uint32_t sw; + uint32_t ph; + uint32_t pw; + uint32_t in_h; + uint32_t in_w; + uint32_t out_h; + uint32_t out_w; + uint32_t channels; + uint32_t numel; + int32_t divisor_override; + uint32_t count_include_pad; + uint32_t has_divisor_override; + uint32_t pad1; +}; +static_assert(sizeof(AvgPoolParams) == 64, "AvgPoolParams must be 64 bytes"); + +// Pooled output extent (in+2p-k)/s+1; ceil_mode: ceil + drop pad-only window. +uint32_t +pool_out_dim(int64_t in, int64_t k, int64_t s, int64_t p, bool ceil_mode) { + const int64_t num = in + 2 * p - k; + int64_t o = ceil_mode ? (num + s - 1) / s + 1 : num / s + 1; + if (ceil_mode && (o - 1) * s >= in + p) { + o -= 1; + } + if (o < 0) { + o = 0; + } + return static_cast(o); +} + +// Read an IntList[2]: empty->fallback (stride->kernel), len1->broadcast, len2. +void read_pair( + const std::vector& v, + uint32_t fallback_h, + uint32_t fallback_w, + uint32_t* h, + uint32_t* w) { + if (v.empty()) { + *h = fallback_h; + *w = fallback_w; + } else if (v.size() == 1) { + *h = static_cast(v[0]); + *w = static_cast(v[0]); + } else { + *h = static_cast(v[0]); + *w = static_cast(v[1]); + } +} + +// avg_pool2d: average over a KxK window per output cell (Vulkan Pool.cpp). +void avg_pool2d_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [in, kernel, stride, padding, ceil_mode, cip, divisor, out]. + const int in_id = args.at(0); + const int kernel_id = args.at(1); + const int stride_id = args.at(2); + const int padding_id = args.at(3); + const int ceil_id = args.at(4); + const int cip_id = args.at(5); + const int divisor_id = args.at(6); + const int out_id = args.at(7); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("avg_pool2d: in/out arg is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.dims.size() != 4 || out_tensor.dims.size() != 4) { + throw std::runtime_error("avg_pool2d: only 4D (NCHW) tensors supported"); + } + + uint32_t kh, kw, sh, sw, ph, pw; + read_pair(graph.get_int_list(kernel_id), 1, 1, &kh, &kw); + read_pair(graph.get_int_list(stride_id), kh, kw, &sh, &sw); + read_pair(graph.get_int_list(padding_id), 0, 0, &ph, &pw); + + int32_t divisor_override = 0; + uint32_t has_divisor_override = 0u; + if (graph.get_value_type(divisor_id) == WebGPUGraph::ValueType::Int) { + divisor_override = static_cast(graph.get_int(divisor_id)); + has_divisor_override = 1u; + } + const uint32_t count_include_pad = graph.get_bool(cip_id) ? 1u : 0u; + + const uint32_t channels = static_cast(in_tensor.dims.at(1)); + const uint32_t in_h = static_cast(in_tensor.dims.at(2)); + const uint32_t in_w = static_cast(in_tensor.dims.at(3)); + const uint32_t out_h = static_cast(out_tensor.dims.at(2)); + const uint32_t out_w = static_cast(out_tensor.dims.at(3)); + + uint64_t out_numel = 1; + for (int64_t d : out_tensor.dims) { + out_numel *= static_cast(d); + } + if (in_tensor.nbytes % sizeof(float) != 0 || + out_tensor.nbytes != out_numel * sizeof(float)) { + throw std::runtime_error("avg_pool2d: fp32-only (byte-size mismatch)"); + } + if (out_numel > UINT32_MAX) { + throw std::runtime_error("avg_pool2d: output numel exceeds u32"); + } + + AvgPoolParams params = {}; + params.kh = kh; + params.kw = kw; + params.sh = sh; + params.sw = sw; + params.ph = ph; + params.pw = pw; + params.in_h = in_h; + params.in_w = in_w; + params.out_h = out_h; + params.out_w = out_w; + params.channels = channels; + params.numel = static_cast(out_numel); + params.divisor_override = divisor_override; + params.count_include_pad = count_include_pad; + params.has_divisor_override = has_divisor_override; + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kAvgPool2dWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, static_cast(out_numel), wg_size, "avg_pool2d"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(AvgPoolParams)); + graph.add_uniform_buffer_bytes(sizeof(AvgPoolParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kAvgPool2dWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[3] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 3; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[3] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = params_buf; + bg_entries[2].size = sizeof(AvgPoolParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 3; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "avg_pool2d", + workgroup_count.y}); + + // Dynamic shapes: recompute the pooled output extent + params + dispatch. + const bool ceil_mode = graph.get_bool(ceil_id); + WGPUBuffer p_buf = params_buf; + graph.add_tensor_resize_hook( + in_id, + [in_id, + out_id, + kh, + kw, + sh, + sw, + ph, + pw, + divisor_override, + has_divisor_override, + count_include_pad, + ceil_mode, + wg_size, + dispatch_idx, + p_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + if (d.size() != 4) { + throw std::runtime_error("avg_pool2d(resize): input is not 4D"); + } + AvgPoolParams p = {}; + p.kh = kh; + p.kw = kw; + p.sh = sh; + p.sw = sw; + p.ph = ph; + p.pw = pw; + p.in_h = static_cast(d[2]); + p.in_w = static_cast(d[3]); + p.out_h = pool_out_dim(d[2], kh, sh, ph, ceil_mode); + p.out_w = pool_out_dim(d[3], kw, sw, pw, ceil_mode); + p.channels = static_cast(d[1]); + const uint64_t out_numel = + static_cast(d[0]) * d[1] * p.out_h * p.out_w; + if (out_numel > UINT32_MAX) { + throw std::runtime_error( + "avg_pool2d(resize): output numel exceeds u32"); + } + p.numel = static_cast(out_numel); + p.divisor_override = divisor_override; + p.has_divisor_override = has_divisor_override; + p.count_include_pad = count_include_pad; + wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), p.numel, wg_size, "avg_pool2d(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + const std::vector out_d = { + d[0], + d[1], + static_cast(p.out_h), + static_cast(p.out_w)}; + g.set_cur_dims(out_id, out_d); + }); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.avg_pool2d.default, avg_pool2d_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/avg_pool2d/avg_pool2d.wgsl b/backends/webgpu/runtime/ops/avg_pool2d/avg_pool2d.wgsl new file mode 100644 index 00000000000..77c55984333 --- /dev/null +++ b/backends/webgpu/runtime/ops/avg_pool2d/avg_pool2d.wgsl @@ -0,0 +1,74 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + kh: u32, + kw: u32, + sh: u32, + sw: u32, + ph: u32, + pw: u32, + in_h: u32, + in_w: u32, + out_h: u32, + out_w: u32, + channels: u32, + numel: u32, + divisor_override: i32, + count_include_pad: u32, + has_divisor_override: u32, + pad1: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let oi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (oi >= params.numel) { + return; + } + + // Unravel NCHW out index; average the clipped input window (Vulkan glsl). + let ow = oi % params.out_w; + let oh = (oi / params.out_w) % params.out_h; + let c = (oi / (params.out_w * params.out_h)) % params.channels; + let n = oi / (params.out_w * params.out_h * params.channels); + + let iph = i32(oh) * i32(params.sh) - i32(params.ph); + let ipw = i32(ow) * i32(params.sw) - i32(params.pw); + let sh0 = max(0, iph); + let eh = min(iph + i32(params.kh), i32(params.in_h)); + let sw0 = max(0, ipw); + let ew = min(ipw + i32(params.kw), i32(params.in_w)); + + let cbase = (n * params.channels + c) * params.in_h * params.in_w; + var acc = 0.0; + for (var ih = sh0; ih < eh; ih = ih + 1) { + for (var iw = sw0; iw < ew; iw = iw + 1) { + acc = acc + input[cbase + u32(ih) * params.in_w + u32(iw)]; + } + } + + var divv: i32; + if (params.has_divisor_override != 0u) { + divv = params.divisor_override; + } else if (params.count_include_pad != 0u) { + // Cells the window extends past the padded input's right/bottom edge. + let beh = iph + i32(params.kh) - i32(params.ph) - i32(params.in_h); + let bew = ipw + i32(params.kw) - i32(params.pw) - i32(params.in_w); + divv = (i32(params.kh) - max(beh, 0)) * (i32(params.kw) - max(bew, 0)); + } else { + divv = (eh - sh0) * (ew - sw0); + } + // Empty window (fully in padding): no cells summed, so avoid 0/0. + if (params.has_divisor_override == 0u && divv <= 0) { + output[oi] = 0.0; + return; + } + output[oi] = acc / f32(divv); +} diff --git a/backends/webgpu/runtime/ops/avg_pool2d/avg_pool2d_wgsl.h b/backends/webgpu/runtime/ops/avg_pool2d/avg_pool2d_wgsl.h new file mode 100644 index 00000000000..47ce9537481 --- /dev/null +++ b/backends/webgpu/runtime/ops/avg_pool2d/avg_pool2d_wgsl.h @@ -0,0 +1,98 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from avg_pool2d.wgsl - DO NOT EDIT. +// wgsl-sha256: 39d22da5d8dd6975bd3d98450ec2db052fde77e89805557cde0d86d9da2530e3 +inline constexpr const char* kAvgPool2dWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + kh: u32, + kw: u32, + sh: u32, + sw: u32, + ph: u32, + pw: u32, + in_h: u32, + in_w: u32, + out_h: u32, + out_w: u32, + channels: u32, + numel: u32, + divisor_override: i32, + count_include_pad: u32, + has_divisor_override: u32, + pad1: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let oi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (oi >= params.numel) { + return; + } + + // Unravel NCHW out index; average the clipped input window (Vulkan glsl). + let ow = oi % params.out_w; + let oh = (oi / params.out_w) % params.out_h; + let c = (oi / (params.out_w * params.out_h)) % params.channels; + let n = oi / (params.out_w * params.out_h * params.channels); + + let iph = i32(oh) * i32(params.sh) - i32(params.ph); + let ipw = i32(ow) * i32(params.sw) - i32(params.pw); + let sh0 = max(0, iph); + let eh = min(iph + i32(params.kh), i32(params.in_h)); + let sw0 = max(0, ipw); + let ew = min(ipw + i32(params.kw), i32(params.in_w)); + + let cbase = (n * params.channels + c) * params.in_h * params.in_w; + var acc = 0.0; + for (var ih = sh0; ih < eh; ih = ih + 1) { + for (var iw = sw0; iw < ew; iw = iw + 1) { + acc = acc + input[cbase + u32(ih) * params.in_w + u32(iw)]; + } + } + + var divv: i32; + if (params.has_divisor_override != 0u) { + divv = params.divisor_override; + } else if (params.count_include_pad != 0u) { + // Cells the window extends past the padded input's right/bottom edge. + let beh = iph + i32(params.kh) - i32(params.ph) - i32(params.in_h); + let bew = ipw + i32(params.kw) - i32(params.pw) - i32(params.in_w); + divv = (i32(params.kh) - max(beh, 0)) * (i32(params.kw) - max(bew, 0)); + } else { + divv = (eh - sh0) * (ew - sw0); + } + // Empty window (fully in padding): no cells summed, so avoid 0/0. + if (params.has_divisor_override == 0u && divv <= 0) { + output[oi] = 0.0; + return; + } + output[oi] = acc / f32(divv); +} +)"; + +inline constexpr uint32_t kAvgPool2dWorkgroupSizeX = 64; +inline constexpr uint32_t kAvgPool2dWorkgroupSizeY = 1; +inline constexpr uint32_t kAvgPool2dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 3742aab8be880b358ab2adbe6cc43ea5e5ba9493 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:20 -0700 Subject: [PATCH 18/75] [ExecuTorch][WebGPU] Op-tests for avg_pool2d MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21176 Adds the op-test suite for the `avg_pool2d` op ported in the diff below. Key changes: - `test/ops/test_avg_pool2d.py` — `AvgPool2dModule` (`F.avg_pool2d`, parametrized by kernel/stride/padding/count_include_pad/ceil_mode/divisor) + `AvgPool2dTest` delegation smoke test. - `test/op_tests/cases.py` — `_avg_pool2d_suite`: basic, pad_cip, pad_nocip, asym, divisor, ceil_cip, ceil_nocip. ghstack-source-id: 406366786 @exported-using-ghexport Differential Revision: [D113493800](https://our.internmc.facebook.com/intern/diff/D113493800/) --- backends/webgpu/test/op_tests/cases.py | 45 ++++++++++ backends/webgpu/test/ops/test_avg_pool2d.py | 98 +++++++++++++++++++++ 2 files changed, 143 insertions(+) create mode 100644 backends/webgpu/test/ops/test_avg_pool2d.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 3bd9a0358d3..e22150eda6b 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -36,6 +36,7 @@ ) from executorch.backends.webgpu.test.ops.test_flip import FlipModule from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule +from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule @@ -308,6 +309,50 @@ def _repeat_suite() -> WebGPUTestSuite: ) +@register_op_test("avg_pool2d") +def _avg_pool2d_suite() -> WebGPUTestSuite: + # Windowed spatial average; real fp math -> float64 oracle at 1e-3. + def mk(kernel, stride, padding, count_include_pad, ceil_mode, divisor): + return AvgPool2dModule( + kernel, stride, padding, count_include_pad, ceil_mode, divisor + ) + + def case(name, k, s, p, cip, ceil_mode, divisor, shape): + return Case( + name=name, + construct={ + "kernel": k, + "stride": s, + "padding": p, + "count_include_pad": cip, + "ceil_mode": ceil_mode, + "divisor": divisor, + }, + inputs=(shape,), + ) + + return WebGPUTestSuite( + module_factory=mk, + cases=[ + case("basic", [2, 2], [2, 2], [0, 0], True, False, None, (1, 2, 4, 4)), + case("pad_cip", [3, 3], [2, 2], [1, 1], True, False, None, (1, 2, 5, 5)), + case( + "pad_nocip", [3, 3], [2, 2], [1, 1], False, False, None, (1, 2, 5, 5) + ), + case("asym", [3, 2], [2, 3], [1, 1], True, False, None, (2, 3, 5, 7)), + case("divisor", [2, 2], [2, 2], [0, 0], True, False, 3, (1, 1, 4, 4)), + # ceil_mode: last window overhangs -> exercises the overhang divisor + # branch (beh/bew > 0) + the ceil output-size (3x3 vs floor 2x2). + case("ceil_cip", [2, 2], [2, 2], [0, 0], True, True, None, (1, 1, 5, 5)), + case( + "ceil_nocip", [3, 3], [2, 2], [0, 0], False, True, None, (1, 2, 5, 5) + ), + ], + atol=1e-3, + rtol=1e-3, + ) + + @register_op_test("native_group_norm") def _group_norm_suite() -> WebGPUTestSuite: # 2-pass norm returning (out, mean, rstd); the multi-output golden verifies diff --git a/backends/webgpu/test/ops/test_avg_pool2d.py b/backends/webgpu/test/ops/test_avg_pool2d.py new file mode 100644 index 00000000000..7dc047d4ec2 --- /dev/null +++ b/backends/webgpu/test/ops/test_avg_pool2d.py @@ -0,0 +1,98 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.avg_pool2d.default` module + configs for the WebGPU op-test framework. + +`AvgPool2dModule` averages over a KxK window (`F.avg_pool2d`). `AvgPool2dTest` +is the export-delegation smoke test. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# name -> input_shape, kernel, stride, padding, count_include_pad, ceil_mode, divisor +CONFIGS = { + "basic": ((1, 2, 4, 4), [2, 2], [2, 2], [0, 0], True, False, None), + "pad_cip": ((1, 2, 5, 5), [3, 3], [2, 2], [1, 1], True, False, None), + "pad_nocip": ((1, 2, 5, 5), [3, 3], [2, 2], [1, 1], False, False, None), + "asym": ((2, 3, 5, 7), [3, 2], [2, 3], [1, 1], True, False, None), + "divisor": ((1, 1, 4, 4), [2, 2], [2, 2], [0, 0], True, False, 3), + # ceil_mode: the last window overhangs the input -> exercises the overhang + # divisor branch (beh/bew > 0) + the ceil output-size (3x3 vs floor 2x2). + "ceil_cip": ((1, 1, 5, 5), [2, 2], [2, 2], [0, 0], True, True, None), + "ceil_nocip": ((1, 2, 5, 5), [3, 3], [2, 2], [0, 0], False, True, None), +} + + +class AvgPool2dModule(torch.nn.Module): + def __init__( + self, kernel, stride, padding, count_include_pad, ceil_mode, divisor + ) -> None: + super().__init__() + self.kernel = kernel + self.stride = stride + self.padding = padding + self.count_include_pad = count_include_pad + self.ceil_mode = ceil_mode + self.divisor = divisor + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.avg_pool2d( + x, + self.kernel, + self.stride, + self.padding, + ceil_mode=self.ceil_mode, + count_include_pad=self.count_include_pad, + divisor_override=self.divisor, + ) + + +def _det_input(shape): + g = torch.Generator().manual_seed(1) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(cfg, x: torch.Tensor): + _, kernel, stride, padding, cip, ceil_mode, divisor = cfg + ep = torch.export.export( + AvgPool2dModule(kernel, stride, padding, cip, ceil_mode, divisor).eval(), (x,) + ) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + +def _delegated(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +def _op_delegated(edge, op_substr: str) -> bool: + # op must be absorbed into the delegate, not left as a top-level CPU-fallback node. + gm = edge.exported_program().graph_module + return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + + +class AvgPool2dTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, cfg in CONFIGS.items(): + with self.subTest(name=name): + edge = _lower(cfg, _det_input(cfg[0])) + et = edge.to_executorch() + self.assertTrue( + _delegated(et), + f"Expected a VulkanBackend delegate (avg_pool2d {name})", + ) + self.assertTrue( + _op_delegated(edge, "avg_pool2d.default"), + f"avg_pool2d not delegated (fell back to CPU) for {name}", + ) From 5956e3aaa038cabdff395232d9f932d72308b43f Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:20 -0700 Subject: [PATCH 19/75] [ExecuTorch][WebGPU] Port pixel_shuffle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21177 Ports `aten.pixel_shuffle.default` as a channel->space rearrange, a Phase B (vision) op. Pure data movement: each output element gathers from one input element. Key changes: - `runtime/ops/pixel_shuffle/{PixelShuffle.cpp, pixel_shuffle.wgsl}` (+ generated `_wgsl.h`) — for output `(n, c_out, h_out, w_out)`: `h_in = h_out / r`, `w_in = w_out / r`, `c_in = c_out*r*r + (h_out%r)*r + (w_out%r)`, then `output[out] = input[in]`. Leading dims collapse into a batch, so any rank >= 3 works. Registered `aten.pixel_shuffle.default`. - `CMakeLists.txt` — `runtime/ops/pixel_shuffle/PixelShuffle.cpp` in `WEBGPU_SRCS`. Mirrors Vulkan `glsl/pixel_shuffle_buffer.glsl` (same `c_in`/`h_in`/`w_in` map) + `impl/PixelShuffle.cpp` (arg order `[in, upscale_factor, out]`, `r >= 1`, `in_channels % (r*r) == 0`, `resize_pixel_shuffle_node`). No int dtype-guard (bit-preserving movement); `nbytes % 4` guard. A dynamic-shape resize hook recomputes the out shape (`in` with the last 3 dims rescaled by r) + params + dispatch + cur_dims. Fail-loud guards: rank >= 3, r >= 1, `in_channels % (r*r) == 0` (mirrors Vulkan VK_CHECK_COND), fp32. ghstack-source-id: 406366788 @exported-using-ghexport Differential Revision: [D112257663](https://our.internmc.facebook.com/intern/diff/D112257663/) --- .../ops/pixel_shuffle/PixelShuffle.cpp | 218 ++++++++++++++++++ .../ops/pixel_shuffle/pixel_shuffle.wgsl | 42 ++++ .../ops/pixel_shuffle/pixel_shuffle_wgsl.h | 66 ++++++ 3 files changed, 326 insertions(+) create mode 100644 backends/webgpu/runtime/ops/pixel_shuffle/PixelShuffle.cpp create mode 100644 backends/webgpu/runtime/ops/pixel_shuffle/pixel_shuffle.wgsl create mode 100644 backends/webgpu/runtime/ops/pixel_shuffle/pixel_shuffle_wgsl.h diff --git a/backends/webgpu/runtime/ops/pixel_shuffle/PixelShuffle.cpp b/backends/webgpu/runtime/ops/pixel_shuffle/PixelShuffle.cpp new file mode 100644 index 00000000000..c8f7e399c99 --- /dev/null +++ b/backends/webgpu/runtime/ops/pixel_shuffle/PixelShuffle.cpp @@ -0,0 +1,218 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct PixelShuffleParams { + uint32_t r; + uint32_t out_c; + uint32_t out_h; + uint32_t out_w; + uint32_t in_c; + uint32_t in_h; + uint32_t in_w; + uint32_t numel; +}; +static_assert( + sizeof(PixelShuffleParams) == 32, + "PixelShuffleParams must match the WGSL Params struct (32 bytes)"); + +// pixel_shuffle: (N,C*r*r,H,W)->(N,C,H*r,W*r) rearrange (Vulkan glsl). +void pixel_shuffle_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [self, upscale_factor, out]. + const int in_id = args.at(0); + const int r_id = args.at(1); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("pixel_shuffle: in/out arg is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + const size_t ndim = in_tensor.dims.size(); + if (ndim < 3 || out_tensor.dims.size() != ndim) { + throw std::runtime_error("pixel_shuffle: expected matching rank >= 3"); + } + + const int64_t r = graph.get_int(r_id); + if (r < 1) { + throw std::runtime_error("pixel_shuffle: upscale_factor must be >= 1"); + } + + // C/H/W are the last 3 dims; any leading dims collapse into the batch. + const uint32_t in_c = static_cast(in_tensor.dims.at(ndim - 3)); + const uint32_t in_h = static_cast(in_tensor.dims.at(ndim - 2)); + const uint32_t in_w = static_cast(in_tensor.dims.at(ndim - 1)); + const uint32_t out_c = static_cast(out_tensor.dims.at(ndim - 3)); + const uint32_t out_h = static_cast(out_tensor.dims.at(ndim - 2)); + const uint32_t out_w = static_cast(out_tensor.dims.at(ndim - 1)); + // Mirror Vulkan VK_CHECK_COND(in_sizes[ndim-3] % (r*r) == 0). + if (in_c % (static_cast(r) * static_cast(r)) != 0) { + throw std::runtime_error("pixel_shuffle: in channels not divisible by r*r"); + } + + uint64_t out_numel = 1; + for (int64_t d : out_tensor.dims) { + out_numel *= static_cast(d); + } + if (in_tensor.nbytes % sizeof(float) != 0 || + out_tensor.nbytes != out_numel * sizeof(float)) { + throw std::runtime_error("pixel_shuffle: non-4-byte operand (nbytes % 4)"); + } + + PixelShuffleParams params = {}; + params.r = static_cast(r); + params.out_c = out_c; + params.out_h = out_h; + params.out_w = out_w; + params.in_c = in_c; + params.in_h = in_h; + params.in_w = in_w; + params.numel = static_cast(out_numel); + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kPixelShuffleWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, static_cast(out_numel), wg_size, "pixel_shuffle"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(PixelShuffleParams)); + graph.add_uniform_buffer_bytes(sizeof(PixelShuffleParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kPixelShuffleWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[3] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 3; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[3] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = params_buf; + bg_entries[2].size = sizeof(PixelShuffleParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 3; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "pixel_shuffle", + workgroup_count.y}); + + // Dynamic shapes: out = in last-3 rescaled by r; recompute params+dispatch. + const uint32_t r_u = static_cast(r); + WGPUBuffer p_buf = params_buf; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, r_u, wg_size, dispatch_idx, p_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + const size_t nd = d.size(); + if (nd < 3) { + throw std::runtime_error("pixel_shuffle(resize): rank < 3"); + } + PixelShuffleParams p = {}; + p.r = r_u; + p.in_c = static_cast(d[nd - 3]); + p.in_h = static_cast(d[nd - 2]); + p.in_w = static_cast(d[nd - 1]); + p.out_c = p.in_c / (r_u * r_u); + p.out_h = p.in_h * r_u; + p.out_w = p.in_w * r_u; + std::vector out_d(d); + out_d[nd - 3] = p.out_c; + out_d[nd - 2] = p.out_h; + out_d[nd - 1] = p.out_w; + uint64_t n = 1; + for (int64_t v : out_d) { + n *= static_cast(v); + } + p.numel = static_cast(n); + wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), p.numel, wg_size, "pixel_shuffle(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + g.set_cur_dims(out_id, out_d); + }); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.pixel_shuffle.default, pixel_shuffle_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/pixel_shuffle/pixel_shuffle.wgsl b/backends/webgpu/runtime/ops/pixel_shuffle/pixel_shuffle.wgsl new file mode 100644 index 00000000000..dd59e65b4dd --- /dev/null +++ b/backends/webgpu/runtime/ops/pixel_shuffle/pixel_shuffle.wgsl @@ -0,0 +1,42 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + r: u32, + out_c: u32, + out_h: u32, + out_w: u32, + in_c: u32, + in_h: u32, + in_w: u32, + numel: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let oi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (oi >= params.numel) { + return; + } + + // (N, C*r*r, H, W) -> (N, C, H*r, W*r); leading dims collapse into b. + let w_out = oi % params.out_w; + let h_out = (oi / params.out_w) % params.out_h; + let c_out = (oi / (params.out_w * params.out_h)) % params.out_c; + let b = oi / (params.out_w * params.out_h * params.out_c); + + let w_in = w_out / params.r; + let h_in = h_out / params.r; + let c_in = c_out * params.r * params.r + + (h_out % params.r) * params.r + (w_out % params.r); + + let in_bufi = + ((b * params.in_c + c_in) * params.in_h + h_in) * params.in_w + w_in; + output[oi] = input[in_bufi]; +} diff --git a/backends/webgpu/runtime/ops/pixel_shuffle/pixel_shuffle_wgsl.h b/backends/webgpu/runtime/ops/pixel_shuffle/pixel_shuffle_wgsl.h new file mode 100644 index 00000000000..322727aa90f --- /dev/null +++ b/backends/webgpu/runtime/ops/pixel_shuffle/pixel_shuffle_wgsl.h @@ -0,0 +1,66 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from pixel_shuffle.wgsl - DO NOT EDIT. +// wgsl-sha256: fc5271241b55091b4989aafec77221ed93115d69b402f46615713339765d91ee +inline constexpr const char* kPixelShuffleWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + r: u32, + out_c: u32, + out_h: u32, + out_w: u32, + in_c: u32, + in_h: u32, + in_w: u32, + numel: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let oi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (oi >= params.numel) { + return; + } + + // (N, C*r*r, H, W) -> (N, C, H*r, W*r); leading dims collapse into b. + let w_out = oi % params.out_w; + let h_out = (oi / params.out_w) % params.out_h; + let c_out = (oi / (params.out_w * params.out_h)) % params.out_c; + let b = oi / (params.out_w * params.out_h * params.out_c); + + let w_in = w_out / params.r; + let h_in = h_out / params.r; + let c_in = c_out * params.r * params.r + + (h_out % params.r) * params.r + (w_out % params.r); + + let in_bufi = + ((b * params.in_c + c_in) * params.in_h + h_in) * params.in_w + w_in; + output[oi] = input[in_bufi]; +} +)"; + +inline constexpr uint32_t kPixelShuffleWorkgroupSizeX = 64; +inline constexpr uint32_t kPixelShuffleWorkgroupSizeY = 1; +inline constexpr uint32_t kPixelShuffleWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 6c6aa21d267322441865697a1b8a30d1bd75e82a Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:21 -0700 Subject: [PATCH 20/75] [ExecuTorch][WebGPU] Op-tests for pixel_shuffle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 16 ++++ .../webgpu/test/ops/test_pixel_shuffle.py | 76 +++++++++++++++++++ 2 files changed, 92 insertions(+) create mode 100644 backends/webgpu/test/ops/test_pixel_shuffle.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index e22150eda6b..e89aa0e222e 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -37,6 +37,7 @@ from executorch.backends.webgpu.test.ops.test_flip import FlipModule from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule +from executorch.backends.webgpu.test.ops.test_pixel_shuffle import PixelShuffleModule from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule @@ -309,6 +310,21 @@ def _repeat_suite() -> WebGPUTestSuite: ) +@register_op_test("pixel_shuffle") +def _pixel_shuffle_suite() -> WebGPUTestSuite: + # Channel->space rearrange; pure data movement -> float32 oracle. + return WebGPUTestSuite( + module_factory=lambda r: PixelShuffleModule(r), + cases=[ + Case(name="r2", construct={"r": 2}, inputs=((1, 8, 2, 3),)), + Case(name="r2_batch", construct={"r": 2}, inputs=((2, 4, 3, 3),)), + Case(name="r3", construct={"r": 3}, inputs=((1, 9, 2, 2),)), + Case(name="r2_3d", construct={"r": 2}, inputs=((4, 2, 2),)), + ], + golden_dtype="float32", + ) + + @register_op_test("avg_pool2d") def _avg_pool2d_suite() -> WebGPUTestSuite: # Windowed spatial average; real fp math -> float64 oracle at 1e-3. diff --git a/backends/webgpu/test/ops/test_pixel_shuffle.py b/backends/webgpu/test/ops/test_pixel_shuffle.py new file mode 100644 index 00000000000..83377e6ca33 --- /dev/null +++ b/backends/webgpu/test/ops/test_pixel_shuffle.py @@ -0,0 +1,76 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.pixel_shuffle.default` module + configs for the WebGPU op-test framework. + +`PixelShuffleModule` rearranges (N, C*r*r, H, W) -> (N, C, H*r, W*r). pixel_shuffle +is pure data movement (bit-identical), so the suite uses the float32 oracle. +`PixelShuffleTest` is the export-delegation smoke test. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# name -> input_shape, upscale_factor +CONFIGS = { + "r2": ((1, 8, 2, 3), 2), + "r2_batch": ((2, 4, 3, 3), 2), + "r3": ((1, 9, 2, 2), 3), + "r2_3d": ((4, 2, 2), 2), +} + + +class PixelShuffleModule(torch.nn.Module): + def __init__(self, r) -> None: + super().__init__() + self.r = r + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.pixel_shuffle(x, self.r) + + +def _det_input(shape): + g = torch.Generator().manual_seed(1) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(r, x: torch.Tensor): + ep = torch.export.export(PixelShuffleModule(r).eval(), (x,)) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + +def _delegated(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +def _op_delegated(edge, op_substr: str) -> bool: + # op must be absorbed into the delegate, not left as a top-level CPU-fallback node. + gm = edge.exported_program().graph_module + return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + + +class PixelShuffleTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, (shape, r) in CONFIGS.items(): + with self.subTest(name=name): + edge = _lower(r, _det_input(shape)) + et = edge.to_executorch() + self.assertTrue( + _delegated(et), + f"Expected a VulkanBackend delegate (pixel_shuffle {name})", + ) + self.assertTrue( + _op_delegated(edge, "pixel_shuffle"), + f"pixel_shuffle not delegated (fell back to CPU) for {name}", + ) From 1e15a2ea10fb07c41d863a7a8cf32057294851d5 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:21 -0700 Subject: [PATCH 21/75] [ExecuTorch][WebGPU] Port grid_sampler_2d MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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` 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/) --- .../ops/grid_sampler_2d/GridSampler2d.cpp | 262 ++++++++++++++++++ .../ops/grid_sampler_2d/grid_sampler_2d.wgsl | 62 +++++ .../grid_sampler_2d/grid_sampler_2d_wgsl.h | 86 ++++++ 3 files changed, 410 insertions(+) create mode 100644 backends/webgpu/runtime/ops/grid_sampler_2d/GridSampler2d.cpp create mode 100644 backends/webgpu/runtime/ops/grid_sampler_2d/grid_sampler_2d.wgsl create mode 100644 backends/webgpu/runtime/ops/grid_sampler_2d/grid_sampler_2d_wgsl.h diff --git a/backends/webgpu/runtime/ops/grid_sampler_2d/GridSampler2d.cpp b/backends/webgpu/runtime/ops/grid_sampler_2d/GridSampler2d.cpp new file mode 100644 index 00000000000..27f88f7203d --- /dev/null +++ b/backends/webgpu/runtime/ops/grid_sampler_2d/GridSampler2d.cpp @@ -0,0 +1,262 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct GridSamplerParams { + uint32_t in_h; + uint32_t in_w; + uint32_t out_h; + uint32_t out_w; + uint32_t channels; + uint32_t numel; + uint32_t pad0; + uint32_t pad1; +}; +static_assert( + sizeof(GridSamplerParams) == 32, + "GridSamplerParams must match the WGSL Params struct (32 bytes)"); + +// grid_sampler_2d: bilinear+border+align_corners sample (Vulkan glsl). +void grid_sampler_2d_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [in, grid, interpolation_mode, padding_mode, align_corners, out]. + const int in_id = args.at(0); + const int grid_id = args.at(1); + const int interp_id = args.at(2); + const int padding_id = args.at(3); + const int align_id = args.at(4); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(grid_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("grid_sampler_2d: in/grid/out is not a tensor"); + } + + // Mirror Vulkan's config guards (bilinear=0, border=1, align_corners=true). + if (graph.get_int(interp_id) != 0) { + throw std::runtime_error("grid_sampler_2d: only bilinear is supported"); + } + if (graph.get_int(padding_id) != 1) { + throw std::runtime_error("grid_sampler_2d: only border padding supported"); + } + if (!graph.get_bool(align_id)) { + throw std::runtime_error("grid_sampler_2d: requires align_corners=true"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& grid_tensor = graph.get_tensor(grid_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.dims.size() != 4 || out_tensor.dims.size() != 4 || + grid_tensor.dims.size() != 4) { + throw std::runtime_error("grid_sampler_2d: in/out/grid must be 4D"); + } + + const uint32_t channels = static_cast(in_tensor.dims.at(1)); + const uint32_t in_h = static_cast(in_tensor.dims.at(2)); + const uint32_t in_w = static_cast(in_tensor.dims.at(3)); + const uint32_t out_h = static_cast(out_tensor.dims.at(2)); + const uint32_t out_w = static_cast(out_tensor.dims.at(3)); + if (in_h < 1u || in_w < 1u) { + throw std::runtime_error("grid_sampler_2d: input H/W must be >= 1"); + } + + uint64_t in_numel = 1; + for (int64_t d : in_tensor.dims) { + in_numel *= static_cast(d); + } + uint64_t out_numel = 1; + for (int64_t d : out_tensor.dims) { + out_numel *= static_cast(d); + } + // grid is [N, out_h, out_w, 2] fp32; validate its own shape (batch matches + // out, trailing coord pair is 2) and derive numel from the grid's dims. + if (grid_tensor.dims.at(0) != out_tensor.dims.at(0)) { + throw std::runtime_error("grid_sampler_2d: grid/out batch mismatch"); + } + if (grid_tensor.dims.at(3) != 2) { + throw std::runtime_error("grid_sampler_2d: grid last dim must be 2"); + } + if (static_cast(grid_tensor.dims.at(1)) != out_h || + static_cast(grid_tensor.dims.at(2)) != out_w) { + throw std::runtime_error("grid_sampler_2d: grid H/W must equal out H/W"); + } + uint64_t grid_numel = 1; + for (int64_t d : grid_tensor.dims) { + grid_numel *= static_cast(d); + } + if (in_tensor.nbytes != in_numel * sizeof(float) || + out_tensor.nbytes != out_numel * sizeof(float) || + grid_tensor.nbytes != grid_numel * sizeof(float)) { + throw std::runtime_error("grid_sampler_2d: fp32-only (byte-size mismatch)"); + } + if (out_numel > UINT32_MAX) { + throw std::runtime_error("grid_sampler_2d: output numel exceeds u32"); + } + + GridSamplerParams params = {}; + params.in_h = in_h; + params.in_w = in_w; + params.out_h = out_h; + params.out_w = out_w; + params.channels = channels; + params.numel = static_cast(out_numel); + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kGridSampler2dWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, static_cast(out_numel), wg_size, "grid_sampler_2d"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(GridSamplerParams)); + graph.add_uniform_buffer_bytes(sizeof(GridSamplerParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kGridSampler2dWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[4] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[3].binding = 3; + entries[3].visibility = WGPUShaderStage_Compute; + entries[3].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 4; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[4] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = grid_tensor.buffer; + bg_entries[2].size = grid_tensor.nbytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = params_buf; + bg_entries[3].size = sizeof(GridSamplerParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 4; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "grid_sampler_2d", + workgroup_count.y}); + + // Dynamic shapes: out shape depends on grid, so trigger on BOTH in and grid. + WGPUBuffer p_buf = params_buf; + auto gs_resize = + [in_id, grid_id, out_id, wg_size, dispatch_idx, p_buf](WebGPUGraph& g) { + const auto& id = g.cur_dims(in_id); + const auto& gd = g.cur_dims(grid_id); + if (id.size() != 4 || gd.size() != 4) { + throw std::runtime_error("grid_sampler_2d(resize): in/grid not 4D"); + } + if (gd[0] != id[0]) { + throw std::runtime_error( + "grid_sampler_2d(resize): grid/out batch mismatch"); + } + if (gd[3] != 2) { + throw std::runtime_error( + "grid_sampler_2d(resize): grid last dim must be 2"); + } + if (id[2] < 1 || id[3] < 1) { + throw std::runtime_error( + "grid_sampler_2d(resize): input H/W must be >= 1"); + } + GridSamplerParams p = {}; + p.in_h = static_cast(id[2]); + p.in_w = static_cast(id[3]); + p.out_h = static_cast(gd[1]); + p.out_w = static_cast(gd[2]); + p.channels = static_cast(id[1]); + p.numel = static_cast( + static_cast(id[0]) * id[1] * p.out_h * p.out_w); + wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), p.numel, wg_size, "grid_sampler_2d(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + const std::vector out_d = { + id[0], + id[1], + static_cast(p.out_h), + static_cast(p.out_w)}; + g.set_cur_dims(out_id, out_d); + }; + graph.add_tensor_resize_hook(in_id, gs_resize); + graph.add_tensor_resize_hook(grid_id, gs_resize); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.grid_sampler_2d.default, grid_sampler_2d_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/grid_sampler_2d/grid_sampler_2d.wgsl b/backends/webgpu/runtime/ops/grid_sampler_2d/grid_sampler_2d.wgsl new file mode 100644 index 00000000000..5ab2a55d2be --- /dev/null +++ b/backends/webgpu/runtime/ops/grid_sampler_2d/grid_sampler_2d.wgsl @@ -0,0 +1,62 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var grid: array; + +struct Params { + in_h: u32, + in_w: u32, + out_h: u32, + out_w: u32, + channels: u32, + numel: u32, + pad0: u32, + pad1: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let oi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (oi >= params.numel) { + return; + } + + // grid[N, out_h, out_w, 2] gives the normalized (x, y) sample coord; Vulkan + // grid_sampler_2d config: bilinear, border padding, align_corners=true. + let ow = oi % params.out_w; + let oh = (oi / params.out_w) % params.out_h; + let c = (oi / (params.out_w * params.out_h)) % params.channels; + let n = oi / (params.out_w * params.out_h * params.channels); + + let gb = ((n * params.out_h + oh) * params.out_w + ow) * 2u; + let gx_norm = grid[gb]; + let gy_norm = grid[gb + 1u]; + + let maxx = f32(params.in_w - 1u); + let maxy = f32(params.in_h - 1u); + let gx = clamp((gx_norm + 1.0) * 0.5 * maxx, 0.0, maxx); + let gy = clamp((gy_norm + 1.0) * 0.5 * maxy, 0.0, maxy); + + let mxi = i32(params.in_w) - 1; + let myi = i32(params.in_h) - 1; + let lx = i32(floor(gx)); + let ly = i32(floor(gy)); + let ux = clamp(lx + 1, 0, mxi); + let uy = clamp(ly + 1, 0, myi); + let wx = gx - f32(lx); + let wy = gy - f32(ly); + + let base = (n * params.channels + c) * params.in_h * params.in_w; + let s00 = input[base + u32(ly) * params.in_w + u32(lx)]; + let s10 = input[base + u32(ly) * params.in_w + u32(ux)]; + let s01 = input[base + u32(uy) * params.in_w + u32(lx)]; + let s11 = input[base + u32(uy) * params.in_w + u32(ux)]; + let top = s00 + wx * (s10 - s00); + let bot = s01 + wx * (s11 - s01); + output[oi] = top + wy * (bot - top); +} diff --git a/backends/webgpu/runtime/ops/grid_sampler_2d/grid_sampler_2d_wgsl.h b/backends/webgpu/runtime/ops/grid_sampler_2d/grid_sampler_2d_wgsl.h new file mode 100644 index 00000000000..78ef4165343 --- /dev/null +++ b/backends/webgpu/runtime/ops/grid_sampler_2d/grid_sampler_2d_wgsl.h @@ -0,0 +1,86 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from grid_sampler_2d.wgsl - DO NOT EDIT. +// wgsl-sha256: c69f408f77ce89d1ca8363e1ca7e2064be910c99d902a762bcbfcf1a2b059588 +inline constexpr const char* kGridSampler2dWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var grid: array; + +struct Params { + in_h: u32, + in_w: u32, + out_h: u32, + out_w: u32, + channels: u32, + numel: u32, + pad0: u32, + pad1: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let oi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (oi >= params.numel) { + return; + } + + // grid[N, out_h, out_w, 2] gives the normalized (x, y) sample coord; Vulkan + // grid_sampler_2d config: bilinear, border padding, align_corners=true. + let ow = oi % params.out_w; + let oh = (oi / params.out_w) % params.out_h; + let c = (oi / (params.out_w * params.out_h)) % params.channels; + let n = oi / (params.out_w * params.out_h * params.channels); + + let gb = ((n * params.out_h + oh) * params.out_w + ow) * 2u; + let gx_norm = grid[gb]; + let gy_norm = grid[gb + 1u]; + + let maxx = f32(params.in_w - 1u); + let maxy = f32(params.in_h - 1u); + let gx = clamp((gx_norm + 1.0) * 0.5 * maxx, 0.0, maxx); + let gy = clamp((gy_norm + 1.0) * 0.5 * maxy, 0.0, maxy); + + let mxi = i32(params.in_w) - 1; + let myi = i32(params.in_h) - 1; + let lx = i32(floor(gx)); + let ly = i32(floor(gy)); + let ux = clamp(lx + 1, 0, mxi); + let uy = clamp(ly + 1, 0, myi); + let wx = gx - f32(lx); + let wy = gy - f32(ly); + + let base = (n * params.channels + c) * params.in_h * params.in_w; + let s00 = input[base + u32(ly) * params.in_w + u32(lx)]; + let s10 = input[base + u32(ly) * params.in_w + u32(ux)]; + let s01 = input[base + u32(uy) * params.in_w + u32(lx)]; + let s11 = input[base + u32(uy) * params.in_w + u32(ux)]; + let top = s00 + wx * (s10 - s00); + let bot = s01 + wx * (s11 - s01); + output[oi] = top + wy * (bot - top); +} +)"; + +inline constexpr uint32_t kGridSampler2dWorkgroupSizeX = 64; +inline constexpr uint32_t kGridSampler2dWorkgroupSizeY = 1; +inline constexpr uint32_t kGridSampler2dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From b6037ec9a4e83693282250419a45de2eae7dfaac Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:22 -0700 Subject: [PATCH 22/75] [ExecuTorch][WebGPU] Op-tests for grid_sampler_2d MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 22 ++++++ .../webgpu/test/ops/test_grid_sampler_2d.py | 76 +++++++++++++++++++ 2 files changed, 98 insertions(+) create mode 100644 backends/webgpu/test/ops/test_grid_sampler_2d.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index e89aa0e222e..bd294b3a7a6 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -37,6 +37,9 @@ from executorch.backends.webgpu.test.ops.test_flip import FlipModule from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule +from executorch.backends.webgpu.test.ops.test_grid_sampler_2d import ( + GridSampler2dModule, +) from executorch.backends.webgpu.test.ops.test_pixel_shuffle import PixelShuffleModule from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule @@ -310,6 +313,25 @@ def _repeat_suite() -> WebGPUTestSuite: ) +@register_op_test("grid_sampler_2d") +def _grid_sampler_2d_suite() -> WebGPUTestSuite: + # Bilinear grid sample (border, align_corners); real fp math -> fp64 oracle. + return WebGPUTestSuite( + module_factory=lambda: GridSampler2dModule(), + cases=[ + Case(name="sq", construct={}, inputs=((1, 2, 4, 4), (1, 3, 3, 2))), + Case( + name="wide_in", + construct={}, + inputs=((1, 1, 3, 5), (1, 4, 4, 2)), + ), + Case(name="batch", construct={}, inputs=((2, 3, 4, 4), (2, 2, 6, 2))), + ], + atol=1e-3, + rtol=1e-3, + ) + + @register_op_test("pixel_shuffle") def _pixel_shuffle_suite() -> WebGPUTestSuite: # Channel->space rearrange; pure data movement -> float32 oracle. diff --git a/backends/webgpu/test/ops/test_grid_sampler_2d.py b/backends/webgpu/test/ops/test_grid_sampler_2d.py new file mode 100644 index 00000000000..069eee690e0 --- /dev/null +++ b/backends/webgpu/test/ops/test_grid_sampler_2d.py @@ -0,0 +1,76 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.grid_sampler_2d.default` module + configs for the WebGPU op-test framework. + +`GridSampler2dModule` samples the input at grid coords (bilinear, border padding, +align_corners=True — the only config the backend supports). Both inputs are +float, so the op-test framework feeds both directly. `GridSampler2dTest` is the +export-delegation smoke test. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# name -> (input_shape, grid_shape) +CONFIGS = { + "sq": ((1, 2, 4, 4), (1, 3, 3, 2)), + "wide_in": ((1, 1, 3, 5), (1, 4, 4, 2)), + "batch": ((2, 3, 4, 4), (2, 2, 6, 2)), +} + + +class GridSampler2dModule(torch.nn.Module): + def forward(self, x: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.grid_sample( + x, grid, mode="bilinear", padding_mode="border", align_corners=True + ) + + +def _det(shape, seed): + g = torch.Generator().manual_seed(seed) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(in_shape, grid_shape): + x = _det(in_shape, 1) + grid = _det(grid_shape, 2) + ep = torch.export.export(GridSampler2dModule().eval(), (x, grid)) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + +def _delegated(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +def _op_delegated(edge, op_substr: str) -> bool: + # op must be absorbed into the delegate, not left as a top-level CPU-fallback node. + gm = edge.exported_program().graph_module + return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + + +class GridSampler2dTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, (in_shape, grid_shape) in CONFIGS.items(): + with self.subTest(name=name): + edge = _lower(in_shape, grid_shape) + et = edge.to_executorch() + self.assertTrue( + _delegated(et), + f"Expected a VulkanBackend delegate (grid_sampler {name})", + ) + self.assertTrue( + _op_delegated(edge, "grid_sampler"), + f"grid_sampler not delegated (fell back to CPU) for {name}", + ) From 384cb6e246623d6116d78572d290fdc8baed3166 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:22 -0700 Subject: [PATCH 23/75] [ExecuTorch][WebGPU] Port aten.convolution (depthwise conv1d) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp | 304 ++++++++++++++++++ .../runtime/ops/conv1d_dw/conv1d_dw.wgsl | 53 +++ .../runtime/ops/conv1d_dw/conv1d_dw_wgsl.h | 77 +++++ 3 files changed, 434 insertions(+) create mode 100644 backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp create mode 100644 backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw.wgsl create mode 100644 backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw_wgsl.h diff --git a/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp new file mode 100644 index 00000000000..84cb8cc5b7d --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp @@ -0,0 +1,304 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct Conv1dDwParams { + uint32_t kernel_size; + uint32_t stride; + uint32_t padding; + uint32_t dilation; + uint32_t channels; + uint32_t in_len; + uint32_t out_len; + uint32_t numel; + uint32_t has_bias; + uint32_t pad0; + float output_min; + float output_max; +}; +static_assert( + sizeof(Conv1dDwParams) == 48, + "Conv1dDwParams must match the WGSL Params struct (48 bytes)"); + +// Convolved output length: (L + 2p - dilation*(K-1) - 1) / stride + 1 (floor). +uint32_t conv1d_out_len( + int64_t in_len, + int64_t k, + int64_t stride, + int64_t padding, + int64_t dilation) { + return static_cast( + (in_len + 2 * padding - dilation * (k - 1) - 1) / stride + 1); +} + +int64_t first_int(const std::vector& v) { + return v.empty() ? 0 : v[0]; +} + +// depthwise-conv1d (groups==C); mirrors Vulkan conv1d_dw (Convolution.cpp:755). +void convolution_impl(WebGPUGraph& graph, const std::vector& args) { + // args mirror Vulkan conv1d_dw; bias (arg 2) may be Null; out=args.back(). + const int in_id = args.at(0); + const int weight_id = args.at(1); + const int bias_id = args.at(2); + const int stride_id = args.at(3); + const int padding_id = args.at(4); + const int dilation_id = args.at(5); + const int transposed_id = args.at(6); + const int groups_id = args.at(8); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(weight_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("convolution: in/weight/out is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& weight_tensor = graph.get_tensor(weight_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.dims.size() != 3 || out_tensor.dims.size() != 3 || + weight_tensor.dims.size() != 3) { + throw std::runtime_error("convolution: only conv1d (3D) is supported"); + } + + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + + const uint32_t channels = static_cast(in_tensor.dims.at(1)); + const uint32_t in_len = static_cast(in_tensor.dims.at(2)); + const uint32_t out_len = static_cast(out_tensor.dims.at(2)); + const uint32_t kernel_size = static_cast(weight_tensor.dims.at(2)); + + // Only the depthwise config (groups==C, weight [C,1,K], not transposed). + const bool transposed = graph.get_bool(transposed_id); + const int64_t groups = graph.get_int(groups_id); + if (transposed || groups != static_cast(channels) || + weight_tensor.dims.at(0) != static_cast(channels) || + weight_tensor.dims.at(1) != 1) { + throw std::runtime_error( + "convolution: only depthwise conv1d (groups==C, weight [C,1,K])"); + } + + const int64_t stride_i = first_int(graph.get_int_list(stride_id)); + const int64_t padding_i = first_int(graph.get_int_list(padding_id)); + const int64_t dilation_i = first_int(graph.get_int_list(dilation_id)); + if (stride_i < 1) { + throw std::runtime_error("convolution: stride must be >= 1"); + } + if (padding_i < 0) { + throw std::runtime_error("convolution: padding must be >= 0"); + } + if (dilation_i < 1) { + throw std::runtime_error("convolution: dilation must be >= 1"); + } + const uint32_t stride = static_cast(stride_i); + const uint32_t padding = static_cast(padding_i); + const uint32_t dilation = static_cast(dilation_i); + + uint64_t out_numel = 1; + for (int64_t d : out_tensor.dims) { + out_numel *= static_cast(d); + } + if (in_tensor.nbytes % sizeof(float) != 0 || + out_tensor.nbytes != out_numel * sizeof(float) || + weight_tensor.nbytes != + static_cast(channels) * kernel_size * sizeof(float)) { + throw std::runtime_error("convolution: fp32-only (byte-size mismatch)"); + } + if (out_numel > UINT32_MAX) { + throw std::runtime_error("convolution: output numel exceeds u32"); + } + + // aten.convolution has no fused clamp (that is et_vk.conv_with_clamp). + const float output_min = std::numeric_limits::lowest(); + const float output_max = std::numeric_limits::max(); + + Conv1dDwParams params = {}; + params.kernel_size = kernel_size; + params.stride = stride; + params.padding = padding; + params.dilation = dilation; + params.channels = channels; + params.in_len = in_len; + params.out_len = out_len; + params.numel = static_cast(out_numel); + params.has_bias = has_bias ? 1u : 0u; + params.output_min = output_min; + params.output_max = output_max; + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kConv1dDwWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, static_cast(out_numel), wg_size, "conv1d_dw"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(Conv1dDwParams)); + graph.add_uniform_buffer_bytes(sizeof(Conv1dDwParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kConv1dDwWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[5] = {}; + for (int i = 0; i < 4; i++) { + entries[i].binding = static_cast(i); + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = (i == 1) ? WGPUBufferBindingType_Storage + : WGPUBufferBindingType_ReadOnlyStorage; + } + entries[4].binding = 4; + entries[4].visibility = WGPUShaderStage_Compute; + entries[4].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 5; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + // No-bias binds the weight buffer as an unread placeholder (has_bias gates). + WGPUBuffer bias_buf = + has_bias ? graph.get_tensor(bias_id).buffer : weight_tensor.buffer; + uint64_t bias_sz = + has_bias ? graph.get_tensor(bias_id).nbytes : weight_tensor.nbytes; + + WGPUBindGroupEntry bg_entries[5] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = weight_tensor.buffer; + bg_entries[2].size = weight_tensor.nbytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = bias_buf; + bg_entries[3].size = bias_sz; + bg_entries[4].binding = 4; + bg_entries[4].buffer = params_buf; + bg_entries[4].size = sizeof(Conv1dDwParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 5; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "conv1d_dw", + workgroup_count.y}); + + // Dynamic shapes: recompute out_len + params + dispatch. + WGPUBuffer p_buf = params_buf; + graph.add_tensor_resize_hook( + in_id, + [in_id, + out_id, + channels, + kernel_size, + stride, + padding, + dilation, + has_bias, + output_min, + output_max, + wg_size, + dispatch_idx, + p_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + if (d.size() != 3) { + throw std::runtime_error("conv1d_dw(resize): input is not 3D"); + } + if (d[1] != static_cast(channels)) { + throw std::runtime_error("conv1d_dw(resize): channel count changed"); + } + Conv1dDwParams p = {}; + p.kernel_size = kernel_size; + p.stride = stride; + p.padding = padding; + p.dilation = dilation; + p.channels = static_cast(d[1]); + p.in_len = static_cast(d[2]); + p.out_len = + conv1d_out_len(d[2], kernel_size, stride, padding, dilation); + const uint64_t out_numel = + static_cast(d[0]) * d[1] * p.out_len; + if (out_numel > UINT32_MAX) { + throw std::runtime_error( + "conv1d_dw(resize): output numel exceeds u32"); + } + p.numel = static_cast(out_numel); + p.has_bias = has_bias ? 1u : 0u; + p.output_min = output_min; + p.output_max = output_max; + wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), p.numel, wg_size, "conv1d_dw(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + const std::vector out_d = { + d[0], d[1], static_cast(p.out_len)}; + g.set_cur_dims(out_id, out_d); + }); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.convolution.default, convolution_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw.wgsl b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw.wgsl new file mode 100644 index 00000000000..04b02af85b7 --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw.wgsl @@ -0,0 +1,53 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + kernel_size: u32, + stride: u32, + padding: u32, + dilation: u32, + channels: u32, + in_len: u32, + out_len: u32, + numel: u32, + has_bias: u32, + pad0: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let oi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (oi >= params.numel) { + return; + } + + // Depthwise: each channel convolves with its own [K] filter (Vulkan glsl). + let l_out = oi % params.out_len; + let c = (oi / params.out_len) % params.channels; + let n = oi / (params.out_len * params.channels); + + let w_base = c * params.kernel_size; + let in_base = (n * params.channels + c) * params.in_len; + var s = 0.0; + for (var k: u32 = 0u; k < params.kernel_size; k = k + 1u) { + let l_in = i32(l_out) * i32(params.stride) - i32(params.padding) + + i32(k) * i32(params.dilation); + if (l_in >= 0 && l_in < i32(params.in_len)) { + s = s + weight[w_base + k] * input[in_base + u32(l_in)]; + } + } + if (params.has_bias != 0u) { + s = s + bias[c]; + } + output[oi] = clamp(s, params.output_min, params.output_max); +} diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw_wgsl.h b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw_wgsl.h new file mode 100644 index 00000000000..9a3b36ebdff --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw_wgsl.h @@ -0,0 +1,77 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from conv1d_dw.wgsl - DO NOT EDIT. +// wgsl-sha256: b803534172bc99cd9092a44998e3a18436bee720c9a3c9cdb9f04fdb538fafca +inline constexpr const char* kConv1dDwWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + kernel_size: u32, + stride: u32, + padding: u32, + dilation: u32, + channels: u32, + in_len: u32, + out_len: u32, + numel: u32, + has_bias: u32, + pad0: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let oi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (oi >= params.numel) { + return; + } + + // Depthwise: each channel convolves with its own [K] filter (Vulkan glsl). + let l_out = oi % params.out_len; + let c = (oi / params.out_len) % params.channels; + let n = oi / (params.out_len * params.channels); + + let w_base = c * params.kernel_size; + let in_base = (n * params.channels + c) * params.in_len; + var s = 0.0; + for (var k: u32 = 0u; k < params.kernel_size; k = k + 1u) { + let l_in = i32(l_out) * i32(params.stride) - i32(params.padding) + + i32(k) * i32(params.dilation); + if (l_in >= 0 && l_in < i32(params.in_len)) { + s = s + weight[w_base + k] * input[in_base + u32(l_in)]; + } + } + if (params.has_bias != 0u) { + s = s + bias[c]; + } + output[oi] = clamp(s, params.output_min, params.output_max); +} +)"; + +inline constexpr uint32_t kConv1dDwWorkgroupSizeX = 64; +inline constexpr uint32_t kConv1dDwWorkgroupSizeY = 1; +inline constexpr uint32_t kConv1dDwWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 31d98a1b40ddab832f430cc36c66c42724a7c9eb Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:23 -0700 Subject: [PATCH 24/75] [ExecuTorch][WebGPU] Op-tests for conv1d_dw (depthwise conv1d) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 34 +++++++ backends/webgpu/test/ops/test_conv1d_dw.py | 102 +++++++++++++++++++++ 2 files changed, 136 insertions(+) create mode 100644 backends/webgpu/test/ops/test_conv1d_dw.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index bd294b3a7a6..1c805ffe0c2 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -37,6 +37,7 @@ from executorch.backends.webgpu.test.ops.test_flip import FlipModule from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule +from executorch.backends.webgpu.test.ops.test_conv1d_dw import Conv1dDWModule from executorch.backends.webgpu.test.ops.test_grid_sampler_2d import ( GridSampler2dModule, ) @@ -313,6 +314,39 @@ def _repeat_suite() -> WebGPUTestSuite: ) +@register_op_test("conv1d_dw") +def _conv1d_dw_suite() -> WebGPUTestSuite: + # Depthwise 1D conv (aten.convolution, depthwise config); fp64 oracle. + def mk(C, kernel, stride, padding, dilation, bias): + return Conv1dDWModule(C, kernel, stride, padding, dilation, bias) + + def case(name, C, L, kernel, stride, padding, dilation, bias): + return Case( + name=name, + construct={ + "C": C, + "kernel": kernel, + "stride": stride, + "padding": padding, + "dilation": dilation, + "bias": bias, + }, + inputs=((1, C, L),), + ) + + return WebGPUTestSuite( + module_factory=mk, + cases=[ + case("k3s1p1", 4, 8, 3, 1, 1, 1, True), + case("k3s2p1", 4, 8, 3, 2, 1, 1, True), + case("dil2", 3, 10, 3, 1, 2, 2, True), + case("k5_nobias", 5, 7, 5, 1, 0, 1, False), + ], + atol=1e-3, + rtol=1e-3, + ) + + @register_op_test("grid_sampler_2d") def _grid_sampler_2d_suite() -> WebGPUTestSuite: # Bilinear grid sample (border, align_corners); real fp math -> fp64 oracle. diff --git a/backends/webgpu/test/ops/test_conv1d_dw.py b/backends/webgpu/test/ops/test_conv1d_dw.py new file mode 100644 index 00000000000..ad4d1c1a481 --- /dev/null +++ b/backends/webgpu/test/ops/test_conv1d_dw.py @@ -0,0 +1,102 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Depthwise conv1d module + configs for the WebGPU op-test framework. + +`Conv1dDWModule` is a depthwise `nn.Conv1d` (groups=C); it serializes as +`aten.convolution.default`, and the handler runs the depthwise-conv1d kernel. +`Conv1dDWTest` is the export-delegation smoke test. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# name -> C, L, kernel, stride, padding, dilation, bias +CONFIGS = { + "k3s1p1": (4, 8, 3, 1, 1, 1, True), + "k3s2p1": (4, 8, 3, 2, 1, 1, True), + "dil2": (3, 10, 3, 1, 2, 2, True), + "k5_nobias": (5, 7, 5, 1, 0, 1, False), +} + + +class Conv1dDWModule(torch.nn.Module): + def __init__(self, C, kernel, stride, padding, dilation, bias) -> None: + super().__init__() + g = torch.Generator().manual_seed(0) + self.conv = torch.nn.Conv1d( + C, + C, + kernel, + stride=stride, + padding=padding, + dilation=dilation, + groups=C, + bias=bias, + ) + with torch.no_grad(): + self.conv.weight.normal_(generator=g) + if bias: + self.conv.bias.normal_(generator=g) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + +def _det_input(shape): + g = torch.Generator().manual_seed(1) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(cfg): + C, L, kernel, stride, padding, dilation, bias = cfg + m = Conv1dDWModule(C, kernel, stride, padding, dilation, bias).eval() + ep = torch.export.export(m, (_det_input((1, C, L)),)) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + +def _delegated(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +def _op_delegated(edge, op_substr: str) -> bool: + # The op must be absorbed into a delegate: absent from the top-level graph AND + # present inside a lowered submodule reached by an executorch_call_delegate node + # (a bare absence check also passes for an empty graph or a renamed op). + from executorch.exir.lowered_backend_module import get_lowered_submodules + + gm = edge.exported_program().graph_module + if any(op_substr in str(getattr(n, "target", "")) for n in gm.graph.nodes): + return False + return any( + op_substr in str(getattr(dn, "target", "")) + for _, lowered, _ in get_lowered_submodules(gm) + for dn in lowered.original_module.graph_module.graph.nodes + ) + + +class Conv1dDWTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, cfg in CONFIGS.items(): + with self.subTest(name=name): + edge = _lower(cfg) + et = edge.to_executorch() + self.assertTrue( + _delegated(et), + f"Expected a VulkanBackend delegate (conv1d_dw {name})", + ) + self.assertTrue( + _op_delegated(edge, "convolution"), + f"conv1d not delegated (fell back to CPU) for {name}", + ) From 991142237b15510029fd6d26b9061207203332ff Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:23 -0700 Subject: [PATCH 25/75] [ExecuTorch][WebGPU] Port conv1d pointwise (aten.convolution) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp | 206 +++++++++++++++++- .../webgpu/runtime/ops/conv1d_dw/conv1d_dw.h | 23 ++ .../runtime/ops/conv1d_dw/conv1d_pw.wgsl | 45 ++++ .../runtime/ops/conv1d_dw/conv1d_pw_wgsl.h | 69 ++++++ .../runtime/ops/et_vk_conv2d/Conv2d.cpp | 8 + 5 files changed, 347 insertions(+), 4 deletions(-) create mode 100644 backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw.h create mode 100644 backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw.wgsl create mode 100644 backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw_wgsl.h diff --git a/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp index 84cb8cc5b7d..f876b095a22 100644 --- a/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp +++ b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp @@ -9,7 +9,9 @@ #include #include #include +#include #include +#include #include @@ -55,6 +57,191 @@ int64_t first_int(const std::vector& v) { return v.empty() ? 0 : v[0]; } +struct Conv1dPwParams { + uint32_t in_channels; + uint32_t out_channels; + uint32_t length; + uint32_t numel; + uint32_t has_bias; + uint32_t pad0; + uint32_t pad1; + uint32_t pad2; +}; +static_assert( + sizeof(Conv1dPwParams) == 32, + "Conv1dPwParams must match the WGSL Params struct (32 bytes)"); + +// Pointwise conv1d (K=1, groups=1): a per-position matmul over channels. +void add_conv1d_pw_node( + WebGPUGraph& graph, + int in_id, + int weight_id, + int bias_id, + int out_id) { + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& weight_tensor = graph.get_tensor(weight_id); + const auto& out_tensor = graph.get_tensor(out_id); + + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + const uint32_t in_channels = static_cast(in_tensor.dims.at(1)); + const uint32_t out_channels = static_cast(out_tensor.dims.at(1)); + const uint32_t length = static_cast(out_tensor.dims.at(2)); + if (in_tensor.dims.at(2) != out_tensor.dims.at(2)) { + throw std::runtime_error("conv1d_pw: in/out length (dim 2) mismatch"); + } + + uint64_t out_numel = 1; + for (int64_t d : out_tensor.dims) { + out_numel *= static_cast(d); + } + if (in_tensor.nbytes % sizeof(float) != 0 || + out_tensor.nbytes != out_numel * sizeof(float) || + weight_tensor.nbytes != + static_cast(out_channels) * in_channels * sizeof(float)) { + throw std::runtime_error("conv1d_pw: fp32-only (byte-size mismatch)"); + } + if (out_numel > UINT32_MAX) { + throw std::runtime_error("conv1d_pw: output numel exceeds u32"); + } + + Conv1dPwParams params = {}; + params.in_channels = in_channels; + params.out_channels = out_channels; + params.length = length; + params.numel = static_cast(out_numel); + params.has_bias = has_bias ? 1u : 0u; + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kConv1dPwWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, static_cast(out_numel), wg_size, "conv1d_pw"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(Conv1dPwParams)); + graph.add_uniform_buffer_bytes(sizeof(Conv1dPwParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kConv1dPwWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[5] = {}; + for (int i = 0; i < 4; i++) { + entries[i].binding = static_cast(i); + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = (i == 1) ? WGPUBufferBindingType_Storage + : WGPUBufferBindingType_ReadOnlyStorage; + } + entries[4].binding = 4; + entries[4].visibility = WGPUShaderStage_Compute; + entries[4].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 5; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBuffer bias_buf = + has_bias ? graph.get_tensor(bias_id).buffer : weight_tensor.buffer; + uint64_t bias_sz = + has_bias ? graph.get_tensor(bias_id).nbytes : weight_tensor.nbytes; + + WGPUBindGroupEntry bg_entries[5] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = weight_tensor.buffer; + bg_entries[2].size = weight_tensor.nbytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = bias_buf; + bg_entries[3].size = bias_sz; + bg_entries[4].binding = 4; + bg_entries[4].buffer = params_buf; + bg_entries[4].size = sizeof(Conv1dPwParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 5; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "conv1d_pw", + workgroup_count.y}); + + // Dynamic shapes: only the length varies; recompute params + dispatch. + WGPUBuffer p_buf = params_buf; + graph.add_tensor_resize_hook( + in_id, + [in_id, + out_id, + in_channels, + out_channels, + has_bias, + wg_size, + dispatch_idx, + p_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + if (d.size() != 3) { + throw std::runtime_error("conv1d_pw(resize): input is not 3D"); + } + if (d[1] != static_cast(in_channels)) { + throw std::runtime_error( + "conv1d_pw(resize): in channel count changed"); + } + Conv1dPwParams p = {}; + p.in_channels = static_cast(d[1]); + p.out_channels = out_channels; + p.length = static_cast(d[2]); + p.numel = static_cast( + static_cast(d[0]) * out_channels * d[2]); + p.has_bias = has_bias ? 1u : 0u; + wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), p.numel, wg_size, "conv1d_pw(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + const std::vector out_d = {d[0], out_channels, d[2]}; + g.set_cur_dims(out_id, out_d); + }); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(params_buf); +} + // depthwise-conv1d (groups==C); mirrors Vulkan conv1d_dw (Convolution.cpp:755). void convolution_impl(WebGPUGraph& graph, const std::vector& args) { // args mirror Vulkan conv1d_dw; bias (arg 2) may be Null; out=args.back(). @@ -91,14 +278,23 @@ void convolution_impl(WebGPUGraph& graph, const std::vector& args) { const uint32_t out_len = static_cast(out_tensor.dims.at(2)); const uint32_t kernel_size = static_cast(weight_tensor.dims.at(2)); - // Only the depthwise config (groups==C, weight [C,1,K], not transposed). const bool transposed = graph.get_bool(transposed_id); const int64_t groups = graph.get_int(groups_id); + + // Pointwise (K=1, groups=1): a matmul over channels; stride-1 / no-pad only. + if (!transposed && groups == 1 && weight_tensor.dims.at(2) == 1 && + first_int(graph.get_int_list(stride_id)) == 1 && + first_int(graph.get_int_list(padding_id)) == 0) { + add_conv1d_pw_node(graph, in_id, weight_id, bias_id, out_id); + return; + } + + // Otherwise only the depthwise config (groups==C, weight [C,1,K]). if (transposed || groups != static_cast(channels) || weight_tensor.dims.at(0) != static_cast(channels) || weight_tensor.dims.at(1) != 1) { throw std::runtime_error( - "convolution: only depthwise conv1d (groups==C, weight [C,1,K])"); + "convolution: only depthwise or pointwise conv1d supported"); } const int64_t stride_i = first_int(graph.get_int_list(stride_id)); @@ -297,8 +493,10 @@ void convolution_impl(WebGPUGraph& graph, const std::vector& args) { } // namespace -WEBGPU_REGISTER_OPERATORS { - WEBGPU_REGISTER_OP(aten.convolution.default, convolution_impl); +// Single aten.convolution.default registration lives in Conv2d.cpp, which +// dispatches on input rank; this is the 3D (conv1d) entry it calls. +void conv1d_dispatch(WebGPUGraph& graph, const std::vector& args) { + convolution_impl(graph, args); } } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw.h b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw.h new file mode 100644 index 00000000000..bcf7a0c8109 --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_dw.h @@ -0,0 +1,23 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +#include + +namespace executorch::backends::webgpu { + +// conv1d (3D) entry: pointwise (K=1, groups=1) or depthwise (groups==C). Called +// by the single aten.convolution.default handler in Conv2d.cpp, which +// dispatches on input rank (3D -> conv1d, 4D -> conv2d) so one registration +// serves both. +void conv1d_dispatch(WebGPUGraph& graph, const std::vector& args); + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw.wgsl b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw.wgsl new file mode 100644 index 00000000000..66260ec78ad --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw.wgsl @@ -0,0 +1,45 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + in_channels: u32, + out_channels: u32, + length: u32, + numel: u32, + has_bias: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let oi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (oi >= params.numel) { + return; + } + + // Pointwise (K=1): out[n,oc,l] = sum_ic weight[oc,ic] * input[n,ic,l]. + let l = oi % params.length; + let oc = (oi / params.length) % params.out_channels; + let n = oi / (params.length * params.out_channels); + + let w_base = oc * params.in_channels; + let in_base = n * params.in_channels * params.length + l; + var s = 0.0; + for (var ic: u32 = 0u; ic < params.in_channels; ic = ic + 1u) { + s = s + weight[w_base + ic] * input[in_base + ic * params.length]; + } + if (params.has_bias != 0u) { + s = s + bias[oc]; + } + output[oi] = s; +} diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw_wgsl.h b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw_wgsl.h new file mode 100644 index 00000000000..854a7160e95 --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_pw_wgsl.h @@ -0,0 +1,69 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from conv1d_pw.wgsl - DO NOT EDIT. +// wgsl-sha256: d47893e191276cf99ee6fd30002ff70f4a0656bcc8e8f90eaca95d9916b54448 +inline constexpr const char* kConv1dPwWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + in_channels: u32, + out_channels: u32, + length: u32, + numel: u32, + has_bias: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let oi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (oi >= params.numel) { + return; + } + + // Pointwise (K=1): out[n,oc,l] = sum_ic weight[oc,ic] * input[n,ic,l]. + let l = oi % params.length; + let oc = (oi / params.length) % params.out_channels; + let n = oi / (params.length * params.out_channels); + + let w_base = oc * params.in_channels; + let in_base = n * params.in_channels * params.length + l; + var s = 0.0; + for (var ic: u32 = 0u; ic < params.in_channels; ic = ic + 1u) { + s = s + weight[w_base + ic] * input[in_base + ic * params.length]; + } + if (params.has_bias != 0u) { + s = s + bias[oc]; + } + output[oi] = s; +} +)"; + +inline constexpr uint32_t kConv1dPwWorkgroupSizeX = 64; +inline constexpr uint32_t kConv1dPwWorkgroupSizeY = 1; +inline constexpr uint32_t kConv1dPwWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp b/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp index 5f127f6d5f2..3c9d49903bd 100644 --- a/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp +++ b/backends/webgpu/runtime/ops/et_vk_conv2d/Conv2d.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -236,6 +237,13 @@ void conv2d_impl(WebGPUGraph& graph, const std::vector& args) { const int groups_id = args.at(8); const int out_id = args.at(9); + // aten.convolution.default covers conv1d (3D) and conv2d (4D); the registry + // is first-wins, so one handler must serve both ranks. Route 3D to conv1d. + if (graph.get_tensor(in_id).dims.size() == 3) { + conv1d_dispatch(graph, args); + return; + } + WGPUDevice device = graph.device(); if (graph.get_bool(transposed_id)) { From 3948270e1dca6a03964b868baadcb4fe9fd50372 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:23 -0700 Subject: [PATCH 26/75] [ExecuTorch][WebGPU] Op-tests for conv1d pointwise MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 28 ++++++++ backends/webgpu/test/ops/test_conv1d_pw.py | 84 ++++++++++++++++++++++ 2 files changed, 112 insertions(+) create mode 100644 backends/webgpu/test/ops/test_conv1d_pw.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 1c805ffe0c2..e907ef576dc 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -38,6 +38,7 @@ from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule from executorch.backends.webgpu.test.ops.test_conv1d_dw import Conv1dDWModule +from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dPwModule from executorch.backends.webgpu.test.ops.test_grid_sampler_2d import ( GridSampler2dModule, ) @@ -314,6 +315,33 @@ def _repeat_suite() -> WebGPUTestSuite: ) +@register_op_test("conv1d_pw") +def _conv1d_pw_suite() -> WebGPUTestSuite: + # Pointwise 1D conv (aten.convolution, K=1 matmul); fp64 oracle. + def mk(in_channels, out_channels, bias): + return Conv1dPwModule(in_channels, out_channels, bias) + + def case(name, N, ic, oc, L, bias): + return Case( + name=name, + construct={"in_channels": ic, "out_channels": oc, "bias": bias}, + inputs=((N, ic, L),), + ) + + return WebGPUTestSuite( + module_factory=mk, + cases=[ + case("ic4_oc6", 1, 4, 6, 5, True), + case("square", 1, 3, 3, 7, True), + case("oc2_nobias", 1, 5, 2, 4, False), + case("ic8_oc8", 1, 8, 8, 3, True), + case("batch2", 2, 3, 4, 5, True), + ], + atol=1e-3, + rtol=1e-3, + ) + + @register_op_test("conv1d_dw") def _conv1d_dw_suite() -> WebGPUTestSuite: # Depthwise 1D conv (aten.convolution, depthwise config); fp64 oracle. diff --git a/backends/webgpu/test/ops/test_conv1d_pw.py b/backends/webgpu/test/ops/test_conv1d_pw.py new file mode 100644 index 00000000000..a0de8988f81 --- /dev/null +++ b/backends/webgpu/test/ops/test_conv1d_pw.py @@ -0,0 +1,84 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Pointwise conv1d module + configs for the WebGPU op-test framework. + +`Conv1dPwModule` is a pointwise `nn.Conv1d` (K=1, groups=1); it serializes as +`aten.convolution.default`, and the handler runs the pointwise-conv1d (matmul) +kernel. `Conv1dPwTest` is the export-delegation smoke test. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# name -> batch, in_channels, out_channels, L, bias +CONFIGS = { + "ic4_oc6": (1, 4, 6, 5, True), + "square": (1, 3, 3, 7, True), + "oc2_nobias": (1, 5, 2, 4, False), + "ic8_oc8": (1, 8, 8, 3, True), + "batch2": (2, 3, 4, 5, True), +} + + +class Conv1dPwModule(torch.nn.Module): + def __init__(self, in_channels, out_channels, bias) -> None: + super().__init__() + g = torch.Generator().manual_seed(0) + self.conv = torch.nn.Conv1d(in_channels, out_channels, 1, bias=bias) + with torch.no_grad(): + self.conv.weight.normal_(generator=g) + if bias: + self.conv.bias.normal_(generator=g) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + +def _det_input(shape): + g = torch.Generator().manual_seed(1) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(cfg): + n, ic, oc, L, bias = cfg + m = Conv1dPwModule(ic, oc, bias).eval() + ep = torch.export.export(m, (_det_input((n, ic, L)),)) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + +def _delegated(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +def _op_delegated(edge, op_substr: str) -> bool: + # op must be absorbed into the delegate, not left as a top-level CPU-fallback node. + gm = edge.exported_program().graph_module + return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + + +class Conv1dPwTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, cfg in CONFIGS.items(): + with self.subTest(name=name): + edge = _lower(cfg) + et = edge.to_executorch() + self.assertTrue( + _delegated(et), + f"Expected a VulkanBackend delegate (conv1d_pw {name})", + ) + self.assertTrue( + _op_delegated(edge, "convolution"), + f"conv1d not delegated (fell back to CPU) for {name}", + ) From e88dd79fa2c70df60163b7a28c4dfe7ae3014383 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:24 -0700 Subject: [PATCH 27/75] [ExecuTorch][WebGPU] Add apply_rotary_emb_interleaved op Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../ops/rope/RotaryEmbeddingInterleaved.cpp | 218 ++++++++++++++++++ .../rope/apply_rotary_emb_interleaved.wgsl | 37 +++ .../rope/apply_rotary_emb_interleaved_wgsl.h | 61 +++++ 3 files changed, 316 insertions(+) create mode 100644 backends/webgpu/runtime/ops/rope/RotaryEmbeddingInterleaved.cpp create mode 100644 backends/webgpu/runtime/ops/rope/apply_rotary_emb_interleaved.wgsl create mode 100644 backends/webgpu/runtime/ops/rope/apply_rotary_emb_interleaved_wgsl.h diff --git a/backends/webgpu/runtime/ops/rope/RotaryEmbeddingInterleaved.cpp b/backends/webgpu/runtime/ops/rope/RotaryEmbeddingInterleaved.cpp new file mode 100644 index 00000000000..73fb112e972 --- /dev/null +++ b/backends/webgpu/runtime/ops/rope/RotaryEmbeddingInterleaved.cpp @@ -0,0 +1,218 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct InterleavedParams { + uint32_t seq; + uint32_t width; + uint32_t numel; + uint32_t pad0; +}; +static_assert( + sizeof(InterleavedParams) == 16, + "InterleavedParams must match the WGSL Params struct (16 bytes)"); + +// Pair-interleaved rope; mirrors Vulkan apply_rotary_emb_interleaved.glsl. +void apply_rotary_emb_interleaved_impl( + WebGPUGraph& graph, + const std::vector& args) { + // args: [x, freqs, out]. + const int in_id = args.at(0); + const int freqs_id = args.at(1); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(freqs_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("rope_interleaved: in/freqs/out is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& freqs_tensor = graph.get_tensor(freqs_id); + const auto& out_tensor = graph.get_tensor(out_id); + // Require rank 3 [B,N,C]; a 4D input would mis-index freqs (mirrors Vulkan). + if (in_tensor.dims.size() != 3 || out_tensor.dims.size() != 3) { + throw std::runtime_error("rope_interleaved: in/out must be 3D [B, N, C]"); + } + if (in_tensor.buffer == nullptr || freqs_tensor.buffer == nullptr || + out_tensor.buffer == nullptr) { + throw std::runtime_error("rope_interleaved: null buffer binding"); + } + + const uint32_t width = static_cast(in_tensor.dims.back()); + const uint32_t seq = + static_cast(in_tensor.dims[in_tensor.dims.size() - 2]); + if (width == 0 || width % 2 != 0) { + throw std::runtime_error( + "rope_interleaved: last dim must be a multiple of 2"); + } + + uint64_t numel = 1; + for (int64_t d : in_tensor.dims) { + numel *= static_cast(d); + } + const uint64_t freqs_numel = utils::numel_of(freqs_tensor.dims); + if (in_tensor.nbytes != numel * sizeof(float) || + out_tensor.nbytes != numel * sizeof(float) || + freqs_numel != static_cast(seq) * width || + freqs_tensor.nbytes != freqs_numel * sizeof(float)) { + throw std::runtime_error( + "rope_interleaved: fp32 byte mismatch or freqs != [seq, width]"); + } + if (numel > UINT32_MAX) { + throw std::runtime_error("rope_interleaved: numel exceeds u32"); + } + + InterleavedParams params = {}; + params.seq = seq; + params.width = width; + params.numel = static_cast(numel); + + uint32_t wg_size = utils::clamp_workgroup_size( + device, kApplyRotaryEmbInterleavedWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, static_cast(numel), wg_size, "rope_interleaved"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(InterleavedParams)); + graph.add_uniform_buffer_bytes(sizeof(InterleavedParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kApplyRotaryEmbInterleavedWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[4] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[3].binding = 3; + entries[3].visibility = WGPUShaderStage_Compute; + entries[3].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 4; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[4] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = freqs_tensor.buffer; + bg_entries[2].size = freqs_tensor.nbytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = params_buf; + bg_entries[3].size = sizeof(InterleavedParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 4; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "rope_interleaved", + workgroup_count.y}); + + // Dynamic shapes: recompute seq/numel + dispatch (freqs stay max-allocated). + WGPUBuffer p_buf = params_buf; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, width, wg_size, dispatch_idx, p_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + if (d.size() != 3) { + throw std::runtime_error("rope_interleaved(resize): rank must be 3"); + } + // width is baked into the params + freqs allocation; only seq may vary. + if (d.back() != static_cast(width)) { + throw std::runtime_error( + "rope_interleaved(resize): last dim (width) changed"); + } + const uint64_t numel = utils::numel_of(d); + if (numel > UINT32_MAX) { + throw std::runtime_error( + "rope_interleaved(resize): numel exceeds u32"); + } + InterleavedParams p = {}; + p.seq = static_cast(d[d.size() - 2]); + p.width = width; + p.numel = static_cast(numel); + wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), p.numel, wg_size, "rope_interleaved(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + g.set_cur_dims(out_id, d); + }); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP( + et_vk.apply_rotary_emb_interleaved.default, + apply_rotary_emb_interleaved_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/rope/apply_rotary_emb_interleaved.wgsl b/backends/webgpu/runtime/ops/rope/apply_rotary_emb_interleaved.wgsl new file mode 100644 index 00000000000..3c255de328a --- /dev/null +++ b/backends/webgpu/runtime/ops/rope/apply_rotary_emb_interleaved.wgsl @@ -0,0 +1,37 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var freqs: array; + +struct Params { + seq: u32, + width: u32, + numel: u32, + pad0: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + + // Pair-interleaved rope (Vulkan glsl): last dim [r,i], freqs [cos,sin]. + let c = idx % params.width; + let n = (idx / params.width) % params.seq; + let ce = c & ~1u; + let base = n * params.width + ce; + let cos_v = freqs[base]; + let sin_v = freqs[base + 1u]; + if ((c & 1u) == 0u) { + output[idx] = input[idx] * cos_v - input[idx + 1u] * sin_v; + } else { + output[idx] = input[idx - 1u] * sin_v + input[idx] * cos_v; + } +} diff --git a/backends/webgpu/runtime/ops/rope/apply_rotary_emb_interleaved_wgsl.h b/backends/webgpu/runtime/ops/rope/apply_rotary_emb_interleaved_wgsl.h new file mode 100644 index 00000000000..f0f6d1f5920 --- /dev/null +++ b/backends/webgpu/runtime/ops/rope/apply_rotary_emb_interleaved_wgsl.h @@ -0,0 +1,61 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from apply_rotary_emb_interleaved.wgsl - DO NOT EDIT. +// wgsl-sha256: bcd9b98f9bade3c2cb56ef3c1b60012aed0bcf8171e66ca9f0e2dc1ad4c03cf8 +inline constexpr const char* kApplyRotaryEmbInterleavedWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var freqs: array; + +struct Params { + seq: u32, + width: u32, + numel: u32, + pad0: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + + // Pair-interleaved rope (Vulkan glsl): last dim [r,i], freqs [cos,sin]. + let c = idx % params.width; + let n = (idx / params.width) % params.seq; + let ce = c & ~1u; + let base = n * params.width + ce; + let cos_v = freqs[base]; + let sin_v = freqs[base + 1u]; + if ((c & 1u) == 0u) { + output[idx] = input[idx] * cos_v - input[idx + 1u] * sin_v; + } else { + output[idx] = input[idx - 1u] * sin_v + input[idx] * cos_v; + } +} +)"; + +inline constexpr uint32_t kApplyRotaryEmbInterleavedWorkgroupSizeX = 64; +inline constexpr uint32_t kApplyRotaryEmbInterleavedWorkgroupSizeY = 1; +inline constexpr uint32_t kApplyRotaryEmbInterleavedWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 9f95bf03f1a80f7aeec6dbbbdf5ea06739c3d760 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:24 -0700 Subject: [PATCH 28/75] [ExecuTorch][WebGPU] Test apply_rotary_emb_interleaved Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 22 ++++++ .../webgpu/test/ops/test_rope_interleaved.py | 77 +++++++++++++++++++ 2 files changed, 99 insertions(+) create mode 100644 backends/webgpu/test/ops/test_rope_interleaved.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index e907ef576dc..ffb3673f994 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -42,6 +42,9 @@ from executorch.backends.webgpu.test.ops.test_grid_sampler_2d import ( GridSampler2dModule, ) +from executorch.backends.webgpu.test.ops.test_rope_interleaved import ( + RopeInterleavedModule, +) from executorch.backends.webgpu.test.ops.test_pixel_shuffle import PixelShuffleModule from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule @@ -375,6 +378,25 @@ def case(name, C, L, kernel, stride, padding, dilation, bias): ) +@register_op_test("apply_rotary_emb_interleaved") +def _rope_interleaved_suite() -> WebGPUTestSuite: + # Pair-interleaved rope (direct custom-op call); CPU eager golden. + def case(name, in_shape, freqs_shape): + return Case(name=name, construct={}, inputs=(in_shape, freqs_shape)) + + return WebGPUTestSuite( + module_factory=lambda: RopeInterleavedModule(), + cases=[ + case("bnc", (1, 4, 8), (4, 8)), + case("batch", (2, 3, 8), (3, 8)), + case("c4", (1, 5, 4), (5, 4)), + case("c16", (1, 2, 16), (2, 16)), + ], + atol=1e-3, + rtol=1e-3, + ) + + @register_op_test("grid_sampler_2d") def _grid_sampler_2d_suite() -> WebGPUTestSuite: # Bilinear grid sample (border, align_corners); real fp math -> fp64 oracle. diff --git a/backends/webgpu/test/ops/test_rope_interleaved.py b/backends/webgpu/test/ops/test_rope_interleaved.py new file mode 100644 index 00000000000..efe8bd86f96 --- /dev/null +++ b/backends/webgpu/test/ops/test_rope_interleaved.py @@ -0,0 +1,77 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`et_vk.apply_rotary_emb_interleaved` module + configs for the op-test framework. + +`RopeInterleavedModule` calls the custom op directly (EdgeTAM has no aten +lowering / fusion pattern for it); both x and the [cos,sin] freqs are fp32 +runtime inputs and the op has a CPU eager impl, so the framework goldens it +directly. `RopeInterleavedTest` is the export-delegation smoke test. +""" + +import unittest + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# name -> input_shape (B, N, C), freqs_shape (N, C) +CONFIGS = { + "bnc": ((1, 4, 8), (4, 8)), + "batch": ((2, 3, 8), (3, 8)), + "c4": ((1, 5, 4), (5, 4)), + "c16": ((1, 2, 16), (2, 16)), +} + + +class RopeInterleavedModule(torch.nn.Module): + def forward(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + return torch.ops.et_vk.apply_rotary_emb_interleaved(x, freqs) + + +def _det(shape, seed): + g = torch.Generator().manual_seed(seed) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(in_shape, freqs_shape): + x = _det(in_shape, 1) + freqs = _det(freqs_shape, 2) + ep = torch.export.export(RopeInterleavedModule().eval(), (x, freqs)) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + +def _delegated(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +def _op_delegated(edge, op_substr: str) -> bool: + # op must be absorbed into the delegate, not left as a top-level CPU-fallback node. + gm = edge.exported_program().graph_module + return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + + +class RopeInterleavedTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, (in_shape, freqs_shape) in CONFIGS.items(): + with self.subTest(name=name): + edge = _lower(in_shape, freqs_shape) + et = edge.to_executorch() + self.assertTrue( + _delegated(et), + f"Expected a VulkanBackend delegate (rope_interleaved {name})", + ) + self.assertTrue( + _op_delegated(edge, "interleaved"), + f"rope_interleaved not delegated (CPU fallback) for {name}", + ) From d4ee755810e7a060978da12cbca3bc32b78b4325 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:25 -0700 Subject: [PATCH 29/75] [ExecuTorch][WebGPU] Add quantize_per_tensor op (int8 buffer path) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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` (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` 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/) --- backends/webgpu/runtime/WebGPUGraph.cpp | 1 + backends/webgpu/runtime/WebGPUGraph.h | 2 + .../ops/quantize/QuantizePerTensor.cpp | 174 ++++++++++++++++++ .../ops/quantize/quantize_per_tensor.wgsl | 31 ++++ .../ops/quantize/quantize_per_tensor_wgsl.h | 55 ++++++ 5 files changed, 263 insertions(+) create mode 100644 backends/webgpu/runtime/ops/quantize/QuantizePerTensor.cpp create mode 100644 backends/webgpu/runtime/ops/quantize/quantize_per_tensor.wgsl create mode 100644 backends/webgpu/runtime/ops/quantize/quantize_per_tensor_wgsl.h diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index abffa25e969..dc0662b1102 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -508,6 +508,7 @@ void WebGPUGraph::build( } tensor.elem_size = vk_datatype_size(vk_tensor->datatype()); tensor.is_int = vk_datatype_is_int(vk_tensor->datatype()); + tensor.is_int8 = vk_tensor->datatype() == vkgraph::VkDataType::INT8; tensor.nbytes = numel * tensor.elem_size; // Live dims start == max (serialized upper bound); resize_input shrinks // them per call. Static graphs keep cur == max forever. diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index c193de9fe35..daa083e992c 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -34,6 +34,8 @@ struct WebGPUTensor { // Serialized (GPU-side) element type, used to narrow wider host inputs. size_t elem_size = 0; bool is_int = false; + // Exactly int8 (not uint8/bool), so int8-only ops can guard their dtype. + bool is_int8 = false; }; // Host-side view of one graph input, passed to copy_inputs. diff --git a/backends/webgpu/runtime/ops/quantize/QuantizePerTensor.cpp b/backends/webgpu/runtime/ops/quantize/QuantizePerTensor.cpp new file mode 100644 index 00000000000..43405bd4962 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantize/QuantizePerTensor.cpp @@ -0,0 +1,174 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct QuantParams { + float inv_scale; + int32_t zero_point; + uint32_t numel; + uint32_t pad0; +}; +static_assert( + sizeof(QuantParams) == 16, + "QuantParams must match the WGSL Params struct (16 bytes)"); + +// fp32->int8; mirrors Vulkan q8ta_quantize.glsl (round(x * inv_scale) + zp). +void quantize_per_tensor_impl( + WebGPUGraph& graph, + const std::vector& args) { + // args: [x, scale, zp, qmin, qmax, dtype, out]; out is always args.back(). + const int in_id = args.at(0); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("quantize_per_tensor: in/out is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("quantize_per_tensor: null buffer binding"); + } + + const double scale = graph.get_double(args.at(1)); + const int zero_point = graph.get_int(args.at(2)); + + uint64_t numel = 1; + for (int64_t d : in_tensor.dims) { + numel *= static_cast(d); + } + // int8 buffer is bound as array; require a whole number of 4-elem words. + if (numel == 0 || numel % 4 != 0) { + throw std::runtime_error( + "quantize_per_tensor: numel must be a nonzero " + "multiple of 4"); + } + if (numel > UINT32_MAX) { + throw std::runtime_error("quantize_per_tensor: numel exceeds u32"); + } + if (in_tensor.nbytes != numel * sizeof(float)) { + throw std::runtime_error("quantize_per_tensor: input is not fp32"); + } + // int8 output (not uint8/bool): the kernel hardcodes the [-128,127] clamp. + if (!out_tensor.is_int8 || out_tensor.nbytes != numel) { + throw std::runtime_error("quantize_per_tensor: output is not int8"); + } + + QuantParams params = {}; + // Reciprocal in double then cast, matching torch's f32(1.0 / f64(scale)). + params.inv_scale = static_cast(1.0 / scale); + params.zero_point = static_cast(zero_point); + params.numel = static_cast(numel); + + const uint32_t num_words = static_cast(numel / 4); + uint32_t wg_size = + utils::clamp_workgroup_size(device, kQuantizePerTensorWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, num_words, wg_size, "quantize_per_tensor"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(QuantParams)); + graph.add_uniform_buffer_bytes(sizeof(QuantParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kQuantizePerTensorWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[3] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 3; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[3] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = params_buf; + bg_entries[2].size = sizeof(QuantParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 3; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "quantize_per_tensor", + workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP( + quantized_decomposed.quantize_per_tensor.default, + quantize_per_tensor_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/quantize/quantize_per_tensor.wgsl b/backends/webgpu/runtime/ops/quantize/quantize_per_tensor.wgsl new file mode 100644 index 00000000000..ae502dd20ea --- /dev/null +++ b/backends/webgpu/runtime/ops/quantize/quantize_per_tensor.wgsl @@ -0,0 +1,31 @@ +@group(0) @binding(0) var t_in: array; +@group(0) @binding(1) var t_out: array; + +struct Params { + inv_scale: f32, + zero_point: i32, + numel: u32, + pad0: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per packed u32 word (4 int8 elems); 2D-fold lifts the 65535 cap. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (widx >= params.numel / 4u) { + return; + } + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + // Multiply by inv_scale, matching torch round(x * (1/scale)). + var q = i32(round(t_in[widx * 4u + j] * params.inv_scale)) + params.zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (j * 8u)); + } + t_out[widx] = packed; +} diff --git a/backends/webgpu/runtime/ops/quantize/quantize_per_tensor_wgsl.h b/backends/webgpu/runtime/ops/quantize/quantize_per_tensor_wgsl.h new file mode 100644 index 00000000000..be07f347d8b --- /dev/null +++ b/backends/webgpu/runtime/ops/quantize/quantize_per_tensor_wgsl.h @@ -0,0 +1,55 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from quantize_per_tensor.wgsl - DO NOT EDIT. +// wgsl-sha256: 45111988635560213805aafd36e5235f4c0e75fba4c94c27c85cc981e06c0473 +inline constexpr const char* kQuantizePerTensorWGSL = R"( +@group(0) @binding(0) var t_in: array; +@group(0) @binding(1) var t_out: array; + +struct Params { + inv_scale: f32, + zero_point: i32, + numel: u32, + pad0: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per packed u32 word (4 int8 elems); 2D-fold lifts the 65535 cap. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (widx >= params.numel / 4u) { + return; + } + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + // Multiply by inv_scale, matching torch round(x * (1/scale)). + var q = i32(round(t_in[widx * 4u + j] * params.inv_scale)) + params.zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (j * 8u)); + } + t_out[widx] = packed; +} +)"; + +inline constexpr uint32_t kQuantizePerTensorWorkgroupSizeX = 64; +inline constexpr uint32_t kQuantizePerTensorWorkgroupSizeY = 1; +inline constexpr uint32_t kQuantizePerTensorWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From c45b869c1b814dc1a09224cdb9992efca3a4e379 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:25 -0700 Subject: [PATCH 30/75] [ExecuTorch][WebGPU] Add dequantize_per_tensor op (int8 buffer path) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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` (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/) --- .../ops/dequantize/DequantizePerTensor.cpp | 172 ++++++++++++++++++ .../ops/dequantize/dequantize_per_tensor.wgsl | 29 +++ .../dequantize/dequantize_per_tensor_wgsl.h | 53 ++++++ 3 files changed, 254 insertions(+) create mode 100644 backends/webgpu/runtime/ops/dequantize/DequantizePerTensor.cpp create mode 100644 backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor.wgsl create mode 100644 backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor_wgsl.h diff --git a/backends/webgpu/runtime/ops/dequantize/DequantizePerTensor.cpp b/backends/webgpu/runtime/ops/dequantize/DequantizePerTensor.cpp new file mode 100644 index 00000000000..83f9b62ce5a --- /dev/null +++ b/backends/webgpu/runtime/ops/dequantize/DequantizePerTensor.cpp @@ -0,0 +1,172 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct DequantParams { + float scale; + int32_t zero_point; + uint32_t numel; + uint32_t pad0; +}; +static_assert( + sizeof(DequantParams) == 16, + "DequantParams must match the WGSL Params struct (16 bytes)"); + +// int8->fp32; mirrors Vulkan q8ta_dequantize.glsl ((q-zp)*scale). +void dequantize_per_tensor_impl( + WebGPUGraph& graph, + const std::vector& args) { + // args: [q, scale, zp, ...]; out is always args.back() (skips out_dtype). + const int in_id = args.at(0); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("dequantize_per_tensor: in/out is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("dequantize_per_tensor: null buffer binding"); + } + + const double scale = graph.get_double(args.at(1)); + const int zero_point = graph.get_int(args.at(2)); + + uint64_t numel = 1; + for (int64_t d : out_tensor.dims) { + numel *= static_cast(d); + } + if (numel == 0 || numel % 4 != 0) { + throw std::runtime_error( + "dequantize_per_tensor: numel must be a nonzero " + "multiple of 4"); + } + if (numel > UINT32_MAX) { + throw std::runtime_error("dequantize_per_tensor: numel exceeds u32"); + } + // int8 input (not uint8/bool): the kernel sign-extends assuming int8. + if (!in_tensor.is_int8 || in_tensor.nbytes != numel) { + throw std::runtime_error("dequantize_per_tensor: input is not int8"); + } + if (out_tensor.nbytes != numel * sizeof(float)) { + throw std::runtime_error("dequantize_per_tensor: output is not fp32"); + } + + DequantParams params = {}; + params.scale = static_cast(scale); + params.zero_point = static_cast(zero_point); + params.numel = static_cast(numel); + + const uint32_t num_words = static_cast(numel / 4); + uint32_t wg_size = + utils::clamp_workgroup_size(device, kDequantizePerTensorWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, num_words, wg_size, "dequantize_per_tensor"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(DequantParams)); + graph.add_uniform_buffer_bytes(sizeof(DequantParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kDequantizePerTensorWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[3] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 3; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[3] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = params_buf; + bg_entries[2].size = sizeof(DequantParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 3; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "dequantize_per_tensor", + workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP( + quantized_decomposed.dequantize_per_tensor.default, + dequantize_per_tensor_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor.wgsl b/backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor.wgsl new file mode 100644 index 00000000000..ab5619dfa13 --- /dev/null +++ b/backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor.wgsl @@ -0,0 +1,29 @@ +@group(0) @binding(0) var t_in: array; +@group(0) @binding(1) var t_out: array; + +struct Params { + scale: f32, + zero_point: i32, + numel: u32, + pad0: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per packed u32 word (4 int8 elems); 2D-fold lifts the 65535 cap. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (widx >= params.numel / 4u) { + return; + } + let word = t_in[widx]; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let b = (word >> (j * 8u)) & 0xFFu; + let s = i32(b ^ 0x80u) - 128; + t_out[widx * 4u + j] = f32(s - params.zero_point) * params.scale; + } +} diff --git a/backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor_wgsl.h b/backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor_wgsl.h new file mode 100644 index 00000000000..a5e2c900fde --- /dev/null +++ b/backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor_wgsl.h @@ -0,0 +1,53 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from dequantize_per_tensor.wgsl - DO NOT EDIT. +// wgsl-sha256: b9e24b5f3b57f6eb842838399985ab2de20513319387fa3341624d310687f90e +inline constexpr const char* kDequantizePerTensorWGSL = R"( +@group(0) @binding(0) var t_in: array; +@group(0) @binding(1) var t_out: array; + +struct Params { + scale: f32, + zero_point: i32, + numel: u32, + pad0: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per packed u32 word (4 int8 elems); 2D-fold lifts the 65535 cap. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (widx >= params.numel / 4u) { + return; + } + let word = t_in[widx]; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let b = (word >> (j * 8u)) & 0xFFu; + let s = i32(b ^ 0x80u) - 128; + t_out[widx * 4u + j] = f32(s - params.zero_point) * params.scale; + } +} +)"; + +inline constexpr uint32_t kDequantizePerTensorWorkgroupSizeX = 64; +inline constexpr uint32_t kDequantizePerTensorWorkgroupSizeY = 1; +inline constexpr uint32_t kDequantizePerTensorWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From e16a2278942b46ee7d8bf3ac4721315ee5fbca88 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:26 -0700 Subject: [PATCH 31/75] [ExecuTorch][WebGPU] Op-test harness: int8-output golden support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/driver_util.cpp | 15 +++++++ backends/webgpu/test/op_tests/driver_util.h | 4 ++ .../webgpu/test/op_tests/generate_op_tests.py | 36 +++++++++++---- .../webgpu/test/op_tests/op_test_driver.cpp | 45 +++++++++++++------ 4 files changed, 78 insertions(+), 22 deletions(-) diff --git a/backends/webgpu/test/op_tests/driver_util.cpp b/backends/webgpu/test/op_tests/driver_util.cpp index 24003a2c3fb..75d398fdacf 100644 --- a/backends/webgpu/test/op_tests/driver_util.cpp +++ b/backends/webgpu/test/op_tests/driver_util.cpp @@ -56,6 +56,7 @@ std::vector parse_manifest(const std::string& manifest_path) { m.golden.path = join(base, g.at("path").get()); m.golden.shape = g.at("shape").get>(); m.golden.output_index = g.value("output_index", 0); + m.golden.dtype = g.value("dtype", std::string("float32")); m.atol = e.value("atol", 1e-3f); m.rtol = e.value("rtol", 1e-3f); m.required = e.value("required", true); @@ -93,6 +94,20 @@ std::vector load_fp32_bin(const std::string& path, size_t numel) { return g; } +std::vector load_int8_bin(const std::string& path, size_t numel) { + FILE* f = std::fopen(path.c_str(), "rb"); + if (!f) { + return {}; + } + std::vector g(numel); + const size_t n = std::fread(g.data(), sizeof(int8_t), numel, f); + std::fclose(f); + if (n != numel) { + return {}; + } + return g; +} + bool within_tol( const float* out, const float* golden, diff --git a/backends/webgpu/test/op_tests/driver_util.h b/backends/webgpu/test/op_tests/driver_util.h index 4c638c4f57d..c171c114309 100644 --- a/backends/webgpu/test/op_tests/driver_util.h +++ b/backends/webgpu/test/op_tests/driver_util.h @@ -24,6 +24,7 @@ struct GoldenRef { std::string path; std::vector shape; int output_index = 0; + std::string dtype = "float32"; // "float32" or "int8" (int8-output ops) }; struct ManifestEntry { @@ -45,6 +46,9 @@ std::vector parse_manifest(const std::string& manifest_path); std::vector load_fp32_bin(const std::string& path, size_t numel); std::vector load_int32_bin(const std::string& path, size_t numel); +/// Load raw int8 (one byte per element); empty on size/IO mismatch. +std::vector load_int8_bin(const std::string& path, size_t numel); + /// Element OK if abs_err <= atol OR rel_err <= rtol (rel floored at /// |golden|=1e-6). Sets the reported maxima; true iff all elements pass. bool within_tol( diff --git a/backends/webgpu/test/op_tests/generate_op_tests.py b/backends/webgpu/test/op_tests/generate_op_tests.py index a4ff7aaa6e6..862d5eaa98d 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -78,6 +78,10 @@ def _write_fp32(t: torch.Tensor, path: str) -> None: t.detach().contiguous().cpu().numpy().astype(" None: + t.detach().contiguous().cpu().numpy().astype(" list[dict]: """Export one case; write its .pte + input/golden .bin(s). Returns one manifest entry per output tensor (multi-output ops emit N; single-output ops emit one, @@ -129,17 +133,31 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> list[d entries: list[dict] = [] n_out = len(golden_outs) for out_index in range(n_out): - out_t = golden_outs[out_index].to(torch.float32) - # Dual-oracle gate: the fp64 golden must match the fp32 eager within tol — - # proves the oracle isn't itself buggy. Skipped for float32 suites. - if golden_dtype == "float64" and case.golden_fn is None: - torch.testing.assert_close( - eager_outs[out_index].to(torch.float32), out_t, atol=atol, rtol=rtol - ) + raw = golden_outs[out_index] + is_int8 = raw.dtype == torch.int8 + if is_int8: + # int8-output ops (quantize): byte-exact golden, no fp32 cast/oracle. + out_t = raw + out_dtype = "int8" + else: + out_t = raw.to(torch.float32) + out_dtype = "float32" + # Dual-oracle gate: the fp64 golden must match the fp32 eager within + # tol — proves the oracle isn't itself buggy. Skipped for float32. + if golden_dtype == "float64" and case.golden_fn is None: + torch.testing.assert_close( + eager_outs[out_index].to(torch.float32), + out_t, + atol=atol, + rtol=rtol, + ) # Single-output ops keep the original name/golden; multi-output ops suffix. suffix = f"_out{out_index}" if n_out > 1 else "" golden_rel = f"{case_id}{suffix}.golden.bin" - _write_fp32(out_t, os.path.join(out_dir, golden_rel)) + if is_int8: + _write_int8(out_t, os.path.join(out_dir, golden_rel)) + else: + _write_fp32(out_t, os.path.join(out_dir, golden_rel)) entries.append( { "op": op, @@ -149,7 +167,7 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> list[d "golden": { "path": golden_rel, "shape": list(out_t.shape), - "dtype": "float32", + "dtype": out_dtype, "output_index": out_index, }, "atol": atol, diff --git a/backends/webgpu/test/op_tests/op_test_driver.cpp b/backends/webgpu/test/op_tests/op_test_driver.cpp index 561e2a9ce57..ba5995ceb81 100644 --- a/backends/webgpu/test/op_tests/op_test_driver.cpp +++ b/backends/webgpu/test/op_tests/op_test_driver.cpp @@ -91,19 +91,38 @@ class OpCase : public ::testing::Test { out_tensor.sizes().begin(), out_tensor.sizes().end()); EXPECT_EQ(out_shape, e_.golden.shape) << "output shape != golden shape (numel matched but dims differ)"; - auto golden = load_fp32_bin(e_.golden.path, gn); - ASSERT_FALSE(golden.empty()) << "missing/short golden: " << e_.golden.path; - - float max_abs = 0.0f, max_rel = 0.0f; - EXPECT_TRUE(within_tol( - out_tensor.const_data_ptr(), - golden.data(), - static_cast(gn), - e_.atol, - e_.rtol, - &max_abs, - &max_rel)) - << "max_abs=" << max_abs << " max_rel=" << max_rel; + if (e_.golden.dtype == "int8") { + // int8-output ops (quantize) compare byte-exact: a discrete grid, so any + // deviation is a real bug, not tolerance. + auto golden = load_int8_bin(e_.golden.path, gn); + ASSERT_FALSE(golden.empty()) + << "missing/short golden: " << e_.golden.path; + const int8_t* out_p = out_tensor.const_data_ptr(); + int mism = -1; + for (size_t i = 0; i < gn; i++) { + if (out_p[i] != golden[i]) { + mism = static_cast(i); + break; + } + } + EXPECT_EQ(mism, -1) << "int8 mismatch at index " << mism + << ": out=" << (mism >= 0 ? int(out_p[mism]) : 0) + << " golden=" << (mism >= 0 ? int(golden[mism]) : 0); + } else { + auto golden = load_fp32_bin(e_.golden.path, gn); + ASSERT_FALSE(golden.empty()) + << "missing/short golden: " << e_.golden.path; + float max_abs = 0.0f, max_rel = 0.0f; + EXPECT_TRUE(within_tol( + out_tensor.const_data_ptr(), + golden.data(), + static_cast(gn), + e_.atol, + e_.rtol, + &max_abs, + &max_rel)) + << "max_abs=" << max_abs << " max_rel=" << max_rel; + } } private: From 4d0b3d993bf8aef38a266d38757afab3b6e2c6ba Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:26 -0700 Subject: [PATCH 32/75] [ExecuTorch][WebGPU] Op-tests for quantize/dequantize_per_tensor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 57 +++++++++++++++++++++ backends/webgpu/test/ops/test_quant.py | 70 ++++++++++++++++++++++++++ 2 files changed, 127 insertions(+) create mode 100644 backends/webgpu/test/ops/test_quant.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index ffb3673f994..b8c46d0f443 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -45,6 +45,10 @@ from executorch.backends.webgpu.test.ops.test_rope_interleaved import ( RopeInterleavedModule, ) +from executorch.backends.webgpu.test.ops.test_quant import ( + DequantizeConstModule, + QuantizeModule, +) from executorch.backends.webgpu.test.ops.test_pixel_shuffle import PixelShuffleModule from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule @@ -1378,3 +1382,56 @@ def _pow_scalar_suite() -> WebGPUTestSuite: for n, e in POW_SCALAR_CONFIGS.items() ], ) + + +# int8 stress: values at the sign-extend edges (-128/127/0/-1), numel % 4 == 0. +_DEQUANT_EDGE = [-128, -127, -1, 0, 1, 63, 126, 127, -64, -50, 50, 100, 7, -8, 42, -42] + + +@register_op_test("quantize_per_tensor") +def _quantize_suite() -> WebGPUTestSuite: + # Lone quantize (int8 output), byte-exact vs torch int8 -- a round-trip folds + # to identity (ET cancels dq(q(x))). golden_dtype float32: golden = int8 eager. + def case(name, shape, scale, zp): + return Case( + name=name, + construct={"scale": scale, "zero_point": zp}, + inputs=(shape,), + ) + + # scale 0.05 -> inv_scale 20; x = k*0.025 lands on half-integer k*0.5 (ties), + # where WGSL round() and torch (both ties-to-even) must still agree byte-exact. + def _ties(shape): + assert torch.Size(shape).numel() == 16, "ties input requires numel == 16" + return (torch.arange(-8, 8, dtype=torch.float32) * 0.025).reshape(shape) + + return WebGPUTestSuite( + module_factory=lambda scale, zero_point: QuantizeModule(scale, zero_point), + cases=[ + case("basic", (4, 8), 0.05, 0), + case("zp_nonzero", (2, 2, 4), 0.1, 5), + case("small_scale", (16,), 0.02, -13), + Case( + name="ties", + construct={"scale": 0.05, "zero_point": 0}, + inputs=(InputSpec(shape=(16,), gen=_ties),), + ), + ], + golden_dtype="float32", + ) + + +@register_op_test("dequantize_per_tensor") +def _dequant_const_suite() -> WebGPUTestSuite: + # Independent GPU dequantize check: baked int8 const vs torch dequant (breaks + # the round-trip's compensating-bug blind spot). golden = CPU eager. + return WebGPUTestSuite( + module_factory=lambda scale, zero_point: DequantizeConstModule( + scale, zero_point, _DEQUANT_EDGE + ), + cases=[ + Case(name="edges", construct={"scale": 0.05, "zero_point": 0}, inputs=()), + Case(name="zp", construct={"scale": 0.1, "zero_point": 7}, inputs=()), + ], + golden_dtype="float32", + ) diff --git a/backends/webgpu/test/ops/test_quant.py b/backends/webgpu/test/ops/test_quant.py new file mode 100644 index 00000000000..b3b3acf1e8f --- /dev/null +++ b/backends/webgpu/test/ops/test_quant.py @@ -0,0 +1,70 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`quantize_per_tensor` + `dequantize_per_tensor` modules + configs. + +Each op is tested ALONE, not as a round-trip: ET folds `dequantize(quantize(x))` +back to `x` before delegation (verified 2026-07-08), so a round-trip would test a +passthrough, not the kernels. `QuantizeModule` emits an int8 output compared +byte-exact to the torch int8 (needs the int8-golden harness path). +`DequantizeConstModule` dequantizes a BAKED int8 constant so the dequantize stage +is verified independently against torch. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +_QMIN, _QMAX = -128, 127 + + +class QuantizeModule(torch.nn.Module): + def __init__(self, scale: float = 0.05, zero_point: int = 0): + super().__init__() + self.scale = scale + self.zero_point = zero_point + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.quantized_decomposed.quantize_per_tensor( + x, self.scale, self.zero_point, _QMIN, _QMAX, torch.int8 + ) + + +class DequantizeConstModule(torch.nn.Module): + def __init__(self, scale: float = 0.05, zero_point: int = 0, q_values=None): + super().__init__() + self.scale = scale + self.zero_point = zero_point + vals = q_values if q_values is not None else list(range(-8, 8)) + self.register_buffer("q", torch.tensor(vals, dtype=torch.int8)) + + def forward(self) -> torch.Tensor: + return torch.ops.quantized_decomposed.dequantize_per_tensor( + self.q, self.scale, self.zero_point, _QMIN, _QMAX, torch.int8 + ) + + +def _delegates(module, inputs) -> bool: + ep = torch.export.export(module, inputs) + et = to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +class QuantTest(unittest.TestCase): + def test_quantize_delegates(self) -> None: + self.assertTrue(_delegates(QuantizeModule(), (torch.randn(4, 8),))) + + def test_dequant_const_delegates(self) -> None: + self.assertTrue(_delegates(DequantizeConstModule(), ())) From fdfd9b04b64d26e065862c9e19f68c50bf41b81f Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:26 -0700 Subject: [PATCH 33/75] [ExecuTorch][WebGPU] Add q8ta_add op (int8 elementwise add) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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` 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/) --- .../webgpu/runtime/ops/q8ta_add/Q8taAdd.cpp | 186 ++++++++++++++++++ .../webgpu/runtime/ops/q8ta_add/q8ta_add.wgsl | 44 +++++ .../runtime/ops/q8ta_add/q8ta_add_wgsl.h | 68 +++++++ 3 files changed, 298 insertions(+) create mode 100644 backends/webgpu/runtime/ops/q8ta_add/Q8taAdd.cpp create mode 100644 backends/webgpu/runtime/ops/q8ta_add/q8ta_add.wgsl create mode 100644 backends/webgpu/runtime/ops/q8ta_add/q8ta_add_wgsl.h diff --git a/backends/webgpu/runtime/ops/q8ta_add/Q8taAdd.cpp b/backends/webgpu/runtime/ops/q8ta_add/Q8taAdd.cpp new file mode 100644 index 00000000000..b2123cc0e46 --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_add/Q8taAdd.cpp @@ -0,0 +1,186 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct Q8taAddParams { + float inv_output_scale; + float a_scale; + float b_scale; + float alpha; + int32_t a_zero_point; + int32_t b_zero_point; + int32_t output_zero_point; + uint32_t numel; +}; +static_assert( + sizeof(Q8taAddParams) == 32, + "Q8taAddParams must match the WGSL Params struct (32 bytes)"); + +// int8 a+alpha*b then requant (CPU-golden; Vulkan glsl buffer drops alpha). +void q8ta_add_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [a, b, a_scale, a_zp, b_scale, b_zp, out_scale, out_zp, alpha, out]. + const int a_id = args.at(0); + const int b_id = args.at(1); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(a_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(b_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("q8ta_add: a/b/out is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& a_tensor = graph.get_tensor(a_id); + const auto& b_tensor = graph.get_tensor(b_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (a_tensor.buffer == nullptr || b_tensor.buffer == nullptr || + out_tensor.buffer == nullptr) { + throw std::runtime_error("q8ta_add: null buffer binding"); + } + + const double a_scale = graph.get_double(args.at(2)); + const int a_zero_point = graph.get_int(args.at(3)); + const double b_scale = graph.get_double(args.at(4)); + const int b_zero_point = graph.get_int(args.at(5)); + const double output_scale = graph.get_double(args.at(6)); + const int output_zero_point = graph.get_int(args.at(7)); + const double alpha = graph.get_double(args.at(8)); + + uint64_t numel = 1; + for (int64_t d : out_tensor.dims) { + numel *= static_cast(d); + } + if (numel == 0 || numel % 4 != 0) { + throw std::runtime_error("q8ta_add: numel must be a nonzero multiple of 4"); + } + if (numel > UINT32_MAX) { + throw std::runtime_error("q8ta_add: numel exceeds u32"); + } + // All three tensors are int8 (kernel clamps to [-128,127]). + if (!a_tensor.is_int8 || !b_tensor.is_int8 || !out_tensor.is_int8 || + a_tensor.nbytes != numel || b_tensor.nbytes != numel || + out_tensor.nbytes != numel) { + throw std::runtime_error("q8ta_add: a/b/out must be int8 of equal numel"); + } + + Q8taAddParams params = {}; + // Reciprocal in double then cast, matching torch's f32(1.0 / f64(scale)). + params.inv_output_scale = static_cast(1.0 / output_scale); + params.a_scale = static_cast(a_scale); + params.b_scale = static_cast(b_scale); + params.alpha = static_cast(alpha); + params.a_zero_point = static_cast(a_zero_point); + params.b_zero_point = static_cast(b_zero_point); + params.output_zero_point = static_cast(output_zero_point); + params.numel = static_cast(numel); + + const uint32_t num_words = static_cast(numel / 4); + uint32_t wg_size = + utils::clamp_workgroup_size(device, kQ8taAddWorkgroupSizeX); + utils::WgCount workgroup_count = + utils::compute_2d_workgroup_count(device, num_words, wg_size, "q8ta_add"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(Q8taAddParams)); + graph.add_uniform_buffer_bytes(sizeof(Q8taAddParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kQ8taAddWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[4] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Storage; + entries[3].binding = 3; + entries[3].visibility = WGPUShaderStage_Compute; + entries[3].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 4; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[4] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = a_tensor.buffer; + bg_entries[0].size = a_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = b_tensor.buffer; + bg_entries[1].size = b_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = out_tensor.buffer; + bg_entries[2].size = out_tensor.nbytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = params_buf; + bg_entries[3].size = sizeof(Q8taAddParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 4; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, bind_group, workgroup_count.x, "q8ta_add", workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.q8ta_add.default, q8ta_add_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/q8ta_add/q8ta_add.wgsl b/backends/webgpu/runtime/ops/q8ta_add/q8ta_add.wgsl new file mode 100644 index 00000000000..77cde5593ca --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_add/q8ta_add.wgsl @@ -0,0 +1,44 @@ +@group(0) @binding(0) var t_a: array; +@group(0) @binding(1) var t_b: array; +@group(0) @binding(2) var t_out: array; + +struct Params { + inv_output_scale: f32, + a_scale: f32, + b_scale: f32, + alpha: f32, + a_zero_point: i32, + b_zero_point: i32, + output_zero_point: i32, + numel: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +fn unpack_i8(word: u32, j: u32) -> i32 { + return i32(((word >> (j * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per packed u32 word (4 int8 elems); 2D-fold lifts the 65535 cap. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (widx >= params.numel / 4u) { + return; + } + let wa = t_a[widx]; + let wb = t_b[widx]; + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let da = params.a_scale * f32(unpack_i8(wa, j) - params.a_zero_point); + let db = params.b_scale * f32(unpack_i8(wb, j) - params.b_zero_point); + let deq = da + params.alpha * db; + var q = i32(round(deq * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (j * 8u)); + } + t_out[widx] = packed; +} diff --git a/backends/webgpu/runtime/ops/q8ta_add/q8ta_add_wgsl.h b/backends/webgpu/runtime/ops/q8ta_add/q8ta_add_wgsl.h new file mode 100644 index 00000000000..f1997add910 --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_add/q8ta_add_wgsl.h @@ -0,0 +1,68 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from q8ta_add.wgsl - DO NOT EDIT. +// wgsl-sha256: 67ce678dbe548b3b29a264dcd2b3070af903e17306ba23cc7ca14dbf629e72c6 +inline constexpr const char* kQ8taAddWGSL = R"( +@group(0) @binding(0) var t_a: array; +@group(0) @binding(1) var t_b: array; +@group(0) @binding(2) var t_out: array; + +struct Params { + inv_output_scale: f32, + a_scale: f32, + b_scale: f32, + alpha: f32, + a_zero_point: i32, + b_zero_point: i32, + output_zero_point: i32, + numel: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +fn unpack_i8(word: u32, j: u32) -> i32 { + return i32(((word >> (j * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per packed u32 word (4 int8 elems); 2D-fold lifts the 65535 cap. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (widx >= params.numel / 4u) { + return; + } + let wa = t_a[widx]; + let wb = t_b[widx]; + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let da = params.a_scale * f32(unpack_i8(wa, j) - params.a_zero_point); + let db = params.b_scale * f32(unpack_i8(wb, j) - params.b_zero_point); + let deq = da + params.alpha * db; + var q = i32(round(deq * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (j * 8u)); + } + t_out[widx] = packed; +} +)"; + +inline constexpr uint32_t kQ8taAddWorkgroupSizeX = 64; +inline constexpr uint32_t kQ8taAddWorkgroupSizeY = 1; +inline constexpr uint32_t kQ8taAddWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 9ec26cbb27fdca10d8e72c5360acb2fd627cf319 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:27 -0700 Subject: [PATCH 34/75] [ExecuTorch][WebGPU] Op-tests for q8ta_add Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 24 ++++++++ backends/webgpu/test/ops/test_q8ta_add.py | 70 +++++++++++++++++++++++ 2 files changed, 94 insertions(+) create mode 100644 backends/webgpu/test/ops/test_q8ta_add.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index b8c46d0f443..eaf1e44f61e 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -49,6 +49,7 @@ DequantizeConstModule, QuantizeModule, ) +from executorch.backends.webgpu.test.ops.test_q8ta_add import Q8taAddModule from executorch.backends.webgpu.test.ops.test_pixel_shuffle import PixelShuffleModule from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule @@ -1421,6 +1422,29 @@ def _ties(shape): ) +# span the int8 sign-extend + requant-clamp edges (-128/127); numel % 4 == 0. +_Q8_A = [-128, 127, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7] +_Q8_B = [-128, 127, 0, 1, -1, 50, -50, 100, -100, 3, -3, 42, -42, 9, -9, 63] + + +@register_op_test("q8ta_add") +def _q8ta_add_suite() -> WebGPUTestSuite: + # int8 add (baked int8 consts), byte-exact vs CPU eager. `alpha` case pins the + # a + alpha*b term (the Vulkan glsl buffer path drops alpha). + def case(name, **kw): + return Case(name=name, construct={"a_vals": _Q8_A, "b_vals": _Q8_B, **kw}) + + return WebGPUTestSuite( + module_factory=lambda **kw: Q8taAddModule(**kw), + cases=[ + case("basic"), + case("alpha", alpha=2.0), + case("nonzero_zp", a_zp=5, b_zp=-4, out_zp=7), + ], + golden_dtype="float32", + ) + + @register_op_test("dequantize_per_tensor") def _dequant_const_suite() -> WebGPUTestSuite: # Independent GPU dequantize check: baked int8 const vs torch dequant (breaks diff --git a/backends/webgpu/test/ops/test_q8ta_add.py b/backends/webgpu/test/ops/test_q8ta_add.py new file mode 100644 index 00000000000..13630e9a035 --- /dev/null +++ b/backends/webgpu/test/ops/test_q8ta_add.py @@ -0,0 +1,70 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`et_vk.q8ta_add` module + configs: int8 elementwise add. + +Both int8 operands are baked constants (only the scalar qparams vary), so the op +is exercised alone with a byte-exact int8 golden vs the CPU eager op. The kernel +implements `a + alpha*b` (the CPU reference semantics); an `alpha != 1` case pins +the alpha term, which the Vulkan glsl buffer path drops. +""" + +import unittest + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 registers et_vk ops + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + + +class Q8taAddModule(torch.nn.Module): + def __init__( + self, + a_vals, + b_vals, + a_scale=0.05, + a_zp=0, + b_scale=0.1, + b_zp=3, + out_scale=0.08, + out_zp=-2, + alpha=1.0, + ): + super().__init__() + self.register_buffer("a", torch.tensor(a_vals, dtype=torch.int8)) + self.register_buffer("b", torch.tensor(b_vals, dtype=torch.int8)) + self.a_scale, self.a_zp = a_scale, a_zp + self.b_scale, self.b_zp = b_scale, b_zp + self.out_scale, self.out_zp, self.alpha = out_scale, out_zp, alpha + + def forward(self) -> torch.Tensor: + return torch.ops.et_vk.q8ta_add( + self.a, + self.b, + self.a_scale, + self.a_zp, + self.b_scale, + self.b_zp, + self.out_scale, + self.out_zp, + self.alpha, + ) + + +class Q8taAddTest(unittest.TestCase): + def test_delegates(self) -> None: + m = Q8taAddModule(list(range(-8, 8)), list(range(7, -9, -1))) + et = to_edge_transform_and_lower( + torch.export.export(m, ()), partitioner=[VulkanPartitioner()] + ).to_executorch() + delegate_ids = [ + d.id + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ] + self.assertIn("VulkanBackend", delegate_ids) From 2687d124a8787b580f0fdde8b010d61f87e883e7 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:27 -0700 Subject: [PATCH 35/75] [ExecuTorch][WebGPU] Add q8ta_relu op (int8 relu) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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` 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/) --- .../webgpu/runtime/ops/q8ta_relu/Q8taRelu.cpp | 200 ++++++++++++++++++ .../runtime/ops/q8ta_relu/q8ta_relu.wgsl | 42 ++++ .../runtime/ops/q8ta_relu/q8ta_relu_wgsl.h | 66 ++++++ 3 files changed, 308 insertions(+) create mode 100644 backends/webgpu/runtime/ops/q8ta_relu/Q8taRelu.cpp create mode 100644 backends/webgpu/runtime/ops/q8ta_relu/q8ta_relu.wgsl create mode 100644 backends/webgpu/runtime/ops/q8ta_relu/q8ta_relu_wgsl.h diff --git a/backends/webgpu/runtime/ops/q8ta_relu/Q8taRelu.cpp b/backends/webgpu/runtime/ops/q8ta_relu/Q8taRelu.cpp new file mode 100644 index 00000000000..a6cafe467a6 --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_relu/Q8taRelu.cpp @@ -0,0 +1,200 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct Q8taReluParams { + float inv_output_scale; + float input_scale; + int32_t input_zero_point; + int32_t output_zero_point; + uint32_t numel; + uint32_t pad0; + uint32_t pad1; + uint32_t pad2; +}; +static_assert( + sizeof(Q8taReluParams) == 32, + "Q8taReluParams must match the WGSL Params struct (32 bytes)"); + +// int8 relu: dequant, max(x, 0), requant; mirrors Vulkan q8ta relu glsl. +void q8ta_relu_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [x, in_scale, in_zp, out_scale, out_zp, out]; out = args.back(). + const int in_id = args.at(0); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("q8ta_relu: in/out is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("q8ta_relu: null buffer binding"); + } + + const double input_scale = graph.get_double(args.at(1)); + const int input_zero_point = graph.get_int(args.at(2)); + const double output_scale = graph.get_double(args.at(3)); + const int output_zero_point = graph.get_int(args.at(4)); + + uint64_t numel = 1; + for (int64_t d : out_tensor.dims) { + numel *= static_cast(d); + } + if (numel == 0 || numel % 4 != 0) { + throw std::runtime_error( + "q8ta_relu: numel must be a nonzero multiple of 4"); + } + if (numel > UINT32_MAX) { + throw std::runtime_error("q8ta_relu: numel exceeds u32"); + } + // in/out int8 (kernel clamps to [-128,127]) of equal numel. + if (!in_tensor.is_int8 || !out_tensor.is_int8 || in_tensor.nbytes != numel || + out_tensor.nbytes != numel) { + throw std::runtime_error("q8ta_relu: in/out must be int8 of equal numel"); + } + + Q8taReluParams params = {}; + // Reciprocal in double then cast, matching torch's f32(1.0 / f64(scale)). + params.inv_output_scale = static_cast(1.0 / output_scale); + params.input_scale = static_cast(input_scale); + params.input_zero_point = static_cast(input_zero_point); + params.output_zero_point = static_cast(output_zero_point); + params.numel = static_cast(numel); + + const uint32_t num_words = static_cast(numel / 4); + uint32_t wg_size = + utils::clamp_workgroup_size(device, kQ8taReluWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, num_words, wg_size, "q8ta_relu"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(Q8taReluParams)); + graph.add_uniform_buffer_bytes(sizeof(Q8taReluParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kQ8taReluWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[3] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 3; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[3] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = params_buf; + bg_entries[2].size = sizeof(Q8taReluParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 3; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "q8ta_relu", + workgroup_count.y}); + + // Dynamic shapes (supports_resize): recompute numel + dispatch, rewrite UBO. + Q8taReluParams base = params; + WGPUBuffer p_buf = params_buf; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, base, wg_size, dispatch_idx, p_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + uint64_t n = utils::numel_of(d); + if (n == 0 || n % 4 != 0 || n > UINT32_MAX) { + throw std::runtime_error( + "q8ta_relu(resize): numel must be a u32 " + "nonzero multiple of 4"); + } + Q8taReluParams p = base; + p.numel = static_cast(n); + wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), + static_cast(n / 4), + wg_size, + "q8ta_relu(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + g.set_cur_dims(out_id, d); + }); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.q8ta_relu.default, q8ta_relu_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/q8ta_relu/q8ta_relu.wgsl b/backends/webgpu/runtime/ops/q8ta_relu/q8ta_relu.wgsl new file mode 100644 index 00000000000..ec5664d761e --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_relu/q8ta_relu.wgsl @@ -0,0 +1,42 @@ +@group(0) @binding(0) var t_in: array; +@group(0) @binding(1) var t_out: array; + +struct Params { + inv_output_scale: f32, + input_scale: f32, + input_zero_point: i32, + output_zero_point: i32, + numel: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +fn unpack_i8(word: u32, j: u32) -> i32 { + return i32(((word >> (j * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per packed u32 word (4 int8 elems); 2D-fold lifts the 65535 cap. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (widx >= params.numel / 4u) { + return; + } + let word = t_in[widx]; + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let deq = + params.input_scale * f32(unpack_i8(word, j) - params.input_zero_point); + let r = max(deq, 0.0); + var q = i32(round(r * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (j * 8u)); + } + t_out[widx] = packed; +} diff --git a/backends/webgpu/runtime/ops/q8ta_relu/q8ta_relu_wgsl.h b/backends/webgpu/runtime/ops/q8ta_relu/q8ta_relu_wgsl.h new file mode 100644 index 00000000000..27d8c7999ea --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_relu/q8ta_relu_wgsl.h @@ -0,0 +1,66 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from q8ta_relu.wgsl - DO NOT EDIT. +// wgsl-sha256: c9e7d458d6cae8a1bf2ae4dd575f8260ba139a60074c044ea7921860061e5a89 +inline constexpr const char* kQ8taReluWGSL = R"( +@group(0) @binding(0) var t_in: array; +@group(0) @binding(1) var t_out: array; + +struct Params { + inv_output_scale: f32, + input_scale: f32, + input_zero_point: i32, + output_zero_point: i32, + numel: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +fn unpack_i8(word: u32, j: u32) -> i32 { + return i32(((word >> (j * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per packed u32 word (4 int8 elems); 2D-fold lifts the 65535 cap. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (widx >= params.numel / 4u) { + return; + } + let word = t_in[widx]; + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let deq = + params.input_scale * f32(unpack_i8(word, j) - params.input_zero_point); + let r = max(deq, 0.0); + var q = i32(round(r * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (j * 8u)); + } + t_out[widx] = packed; +} +)"; + +inline constexpr uint32_t kQ8taReluWorkgroupSizeX = 64; +inline constexpr uint32_t kQ8taReluWorkgroupSizeY = 1; +inline constexpr uint32_t kQ8taReluWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 619702dce6f5db2b96e0179a74beea73c40e4d79 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:28 -0700 Subject: [PATCH 36/75] [ExecuTorch][WebGPU] Op-tests for q8ta_relu Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 19 +++++++ backends/webgpu/test/ops/test_q8ta_relu.py | 60 ++++++++++++++++++++++ 2 files changed, 79 insertions(+) create mode 100644 backends/webgpu/test/ops/test_q8ta_relu.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index eaf1e44f61e..043764aabaa 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -50,6 +50,7 @@ QuantizeModule, ) from executorch.backends.webgpu.test.ops.test_q8ta_add import Q8taAddModule +from executorch.backends.webgpu.test.ops.test_q8ta_relu import Q8taReluModule from executorch.backends.webgpu.test.ops.test_pixel_shuffle import PixelShuffleModule from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule @@ -1445,6 +1446,24 @@ def case(name, **kw): ) +@register_op_test("q8ta_relu") +def _q8ta_relu_suite() -> WebGPUTestSuite: + # int8 relu (baked int8 const), byte-exact vs CPU eager. Inputs whose dequant + # is negative are relu-clamped to 0 (pins max(x,0)); edges -128/127 covered. + def case(name, **kw): + return Case(name=name, construct={"x_vals": _Q8_A, **kw}) + + return WebGPUTestSuite( + module_factory=lambda **kw: Q8taReluModule(**kw), + cases=[ + case("basic"), + case("diff_qparams", output_scale=0.08, output_zp=5), + case("nonzero_zp", input_zp=10, output_zp=-20), + ], + golden_dtype="float32", + ) + + @register_op_test("dequantize_per_tensor") def _dequant_const_suite() -> WebGPUTestSuite: # Independent GPU dequantize check: baked int8 const vs torch dequant (breaks diff --git a/backends/webgpu/test/ops/test_q8ta_relu.py b/backends/webgpu/test/ops/test_q8ta_relu.py new file mode 100644 index 00000000000..b369c518930 --- /dev/null +++ b/backends/webgpu/test/ops/test_q8ta_relu.py @@ -0,0 +1,60 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`et_vk.q8ta_relu` module + configs: int8 relu. + +The int8 input is a baked constant, so the op is exercised alone with a byte-exact +int8 golden vs the CPU eager op. Inputs whose dequantized value is negative are +clamped to 0 by the relu, pinning the `max(x, 0)` term. +""" + +import unittest + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 registers et_vk ops + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + + +class Q8taReluModule(torch.nn.Module): + def __init__( + self, + x_vals, + input_scale=0.05, + input_zp=-3, + output_scale=0.05, + output_zp=-3, + ): + super().__init__() + self.register_buffer("x", torch.tensor(x_vals, dtype=torch.int8)) + self.input_scale, self.input_zp = input_scale, input_zp + self.output_scale, self.output_zp = output_scale, output_zp + + def forward(self) -> torch.Tensor: + return torch.ops.et_vk.q8ta_relu( + self.x, + self.input_scale, + self.input_zp, + self.output_scale, + self.output_zp, + ) + + +class Q8taReluTest(unittest.TestCase): + def test_delegates(self) -> None: + m = Q8taReluModule(list(range(-8, 8))) + et = to_edge_transform_and_lower( + torch.export.export(m, ()), partitioner=[VulkanPartitioner()] + ).to_executorch() + self.assertTrue( + any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + ) From f3a262fdb901ad04dc3bbd69d94bb6dd435616d1 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:28 -0700 Subject: [PATCH 37/75] [ExecuTorch][WebGPU] Add q8ta_pixel_shuffle op (int8 pixel_shuffle) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../q8ta_pixel_shuffle/Q8taPixelShuffle.cpp | 256 ++++++++++++++++++ .../q8ta_pixel_shuffle.wgsl | 57 ++++ .../q8ta_pixel_shuffle_wgsl.h | 81 ++++++ 3 files changed, 394 insertions(+) create mode 100644 backends/webgpu/runtime/ops/q8ta_pixel_shuffle/Q8taPixelShuffle.cpp create mode 100644 backends/webgpu/runtime/ops/q8ta_pixel_shuffle/q8ta_pixel_shuffle.wgsl create mode 100644 backends/webgpu/runtime/ops/q8ta_pixel_shuffle/q8ta_pixel_shuffle_wgsl.h diff --git a/backends/webgpu/runtime/ops/q8ta_pixel_shuffle/Q8taPixelShuffle.cpp b/backends/webgpu/runtime/ops/q8ta_pixel_shuffle/Q8taPixelShuffle.cpp new file mode 100644 index 00000000000..7fa56895856 --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_pixel_shuffle/Q8taPixelShuffle.cpp @@ -0,0 +1,256 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct Q8taPixelShuffleParams { + uint32_t r; + uint32_t out_c; + uint32_t out_h; + uint32_t out_w; + uint32_t in_c; + uint32_t in_h; + uint32_t in_w; + uint32_t numel; + float input_scale; + float inv_output_scale; + int32_t input_zero_point; + int32_t output_zero_point; +}; +static_assert( + sizeof(Q8taPixelShuffleParams) == 48, + "Q8taPixelShuffleParams must match the WGSL Params struct (48 bytes)"); + +// int8 pixel_shuffle gather + requant; mirrors Vulkan q8ta_pixel_shuffle.glsl. +void q8ta_pixel_shuffle_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [x, in_scale, in_zp, out_inv_scale, out_zp, upscale_factor, out]. + const int in_id = args.at(0); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("q8ta_pixel_shuffle: in/out is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + const size_t ndim = in_tensor.dims.size(); + if (ndim < 3 || out_tensor.dims.size() != ndim) { + throw std::runtime_error("q8ta_pixel_shuffle: matching rank >= 3 expected"); + } + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("q8ta_pixel_shuffle: null buffer binding"); + } + + const double input_scale = graph.get_double(args.at(1)); + const int input_zero_point = graph.get_int(args.at(2)); + // 4th arg is already 1/output_scale (output_inv_scale); not re-inverted. + const double output_inv_scale = graph.get_double(args.at(3)); + const int output_zero_point = graph.get_int(args.at(4)); + const int64_t r = graph.get_int(args.at(5)); + if (r < 1) { + throw std::runtime_error("q8ta_pixel_shuffle: upscale_factor must be >= 1"); + } + + const uint32_t in_c = static_cast(in_tensor.dims.at(ndim - 3)); + const uint32_t in_h = static_cast(in_tensor.dims.at(ndim - 2)); + const uint32_t in_w = static_cast(in_tensor.dims.at(ndim - 1)); + const uint32_t out_c = static_cast(out_tensor.dims.at(ndim - 3)); + const uint32_t out_h = static_cast(out_tensor.dims.at(ndim - 2)); + const uint32_t out_w = static_cast(out_tensor.dims.at(ndim - 1)); + const uint32_t rr = static_cast(r) * static_cast(r); + if (in_c % rr != 0) { + throw std::runtime_error("q8ta_pixel_shuffle: in channels not div by r*r"); + } + // The shader derives in c/h/w from the OUT dims, so an out shape inconsistent + // with pixel_shuffle(in, r) (but same total numel) would read OOB. Mirrors + // Vulkan Q8taPixelShuffle.cpp VK_CHECK_COND(in == out*r*r / out*r). + if (out_c != in_c / rr || out_h != in_h * static_cast(r) || + out_w != in_w * static_cast(r)) { + throw std::runtime_error( + "q8ta_pixel_shuffle: out dims != pixel_shuffle(in, r)"); + } + + uint64_t numel = 1; + for (int64_t d : out_tensor.dims) { + numel *= static_cast(d); + } + if (numel == 0 || numel % 4 != 0) { + throw std::runtime_error( + "q8ta_pixel_shuffle: numel must be a nonzero multiple of 4"); + } + if (numel > UINT32_MAX) { + throw std::runtime_error("q8ta_pixel_shuffle: numel exceeds u32"); + } + // in/out int8 (kernel clamps to [-128,127]) of equal numel. + if (!in_tensor.is_int8 || !out_tensor.is_int8 || in_tensor.nbytes != numel || + out_tensor.nbytes != numel) { + throw std::runtime_error( + "q8ta_pixel_shuffle: in/out must be int8 of equal numel"); + } + + Q8taPixelShuffleParams params = {}; + params.r = static_cast(r); + params.out_c = out_c; + params.out_h = out_h; + params.out_w = out_w; + params.in_c = in_c; + params.in_h = in_h; + params.in_w = in_w; + params.numel = static_cast(numel); + params.input_scale = static_cast(input_scale); + params.inv_output_scale = static_cast(output_inv_scale); + params.input_zero_point = static_cast(input_zero_point); + params.output_zero_point = static_cast(output_zero_point); + + const uint32_t num_words = static_cast(numel / 4); + uint32_t wg_size = + utils::clamp_workgroup_size(device, kQ8taPixelShuffleWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, num_words, wg_size, "q8ta_pixel_shuffle"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(Q8taPixelShuffleParams)); + graph.add_uniform_buffer_bytes(sizeof(Q8taPixelShuffleParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kQ8taPixelShuffleWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[3] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 3; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[3] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = params_buf; + bg_entries[2].size = sizeof(Q8taPixelShuffleParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 3; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "q8ta_pixel_shuffle", + workgroup_count.y}); + + // Dynamic shapes (supports_resize): recompute shape/numel/dispatch + UBO. + Q8taPixelShuffleParams base = params; + const uint32_t r_u = static_cast(r); + WGPUBuffer p_buf = params_buf; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, base, r_u, wg_size, dispatch_idx, p_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + const size_t nd = d.size(); + if (nd < 3) { + throw std::runtime_error("q8ta_pixel_shuffle(resize): rank < 3"); + } + Q8taPixelShuffleParams p = base; + p.in_c = static_cast(d[nd - 3]); + p.in_h = static_cast(d[nd - 2]); + p.in_w = static_cast(d[nd - 1]); + p.out_c = p.in_c / (r_u * r_u); + p.out_h = p.in_h * r_u; + p.out_w = p.in_w * r_u; + std::vector out_d(d); + out_d[nd - 3] = p.out_c; + out_d[nd - 2] = p.out_h; + out_d[nd - 1] = p.out_w; + uint64_t n = 1; + for (int64_t v : out_d) { + n *= static_cast(v); + } + if (n == 0 || n % 4 != 0 || n > UINT32_MAX) { + throw std::runtime_error( + "q8ta_pixel_shuffle(resize): numel must be a u32 multiple of 4"); + } + p.numel = static_cast(n); + wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), + static_cast(n / 4), + wg_size, + "q8ta_pixel_shuffle(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + g.set_cur_dims(out_id, out_d); + }); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.q8ta_pixel_shuffle.default, q8ta_pixel_shuffle_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/q8ta_pixel_shuffle/q8ta_pixel_shuffle.wgsl b/backends/webgpu/runtime/ops/q8ta_pixel_shuffle/q8ta_pixel_shuffle.wgsl new file mode 100644 index 00000000000..ba8c4d74d07 --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_pixel_shuffle/q8ta_pixel_shuffle.wgsl @@ -0,0 +1,57 @@ +@group(0) @binding(0) var t_in: array; +@group(0) @binding(1) var t_out: array; + +struct Params { + r: u32, + out_c: u32, + out_h: u32, + out_w: u32, + in_c: u32, + in_h: u32, + in_w: u32, + numel: u32, + input_scale: f32, + inv_output_scale: f32, + input_zero_point: i32, + output_zero_point: i32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +fn unpack_i8(word: u32, byte_idx: u32) -> i32 { + return i32(((word >> (byte_idx * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per packed u32 word (4 int8 outputs); 2D-fold lifts 65535. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (widx >= params.numel / 4u) { + return; + } + var packed: u32 = 0u; + for (var k: u32 = 0u; k < 4u; k = k + 1u) { + // (N, C*r*r, H, W) -> (N, C, H*r, W*r) gather; leading dims collapse to b. + let oi = widx * 4u + k; + let w_out = oi % params.out_w; + let h_out = (oi / params.out_w) % params.out_h; + let c_out = (oi / (params.out_w * params.out_h)) % params.out_c; + let b = oi / (params.out_w * params.out_h * params.out_c); + let w_in = w_out / params.r; + let h_in = h_out / params.r; + let c_in = c_out * params.r * params.r + + (h_out % params.r) * params.r + (w_out % params.r); + let in_flat = + ((b * params.in_c + c_in) * params.in_h + h_in) * params.in_w + w_in; + let q_in = unpack_i8(t_in[in_flat / 4u], in_flat % 4u); + let deq = params.input_scale * f32(q_in - params.input_zero_point); + var q = + i32(round(deq * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (k * 8u)); + } + t_out[widx] = packed; +} diff --git a/backends/webgpu/runtime/ops/q8ta_pixel_shuffle/q8ta_pixel_shuffle_wgsl.h b/backends/webgpu/runtime/ops/q8ta_pixel_shuffle/q8ta_pixel_shuffle_wgsl.h new file mode 100644 index 00000000000..13403cba3ad --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_pixel_shuffle/q8ta_pixel_shuffle_wgsl.h @@ -0,0 +1,81 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from q8ta_pixel_shuffle.wgsl - DO NOT EDIT. +// wgsl-sha256: d9f89aace1cb13a5ee6be263d36dd0636423b873678f6c14502aa9ba6807604a +inline constexpr const char* kQ8taPixelShuffleWGSL = R"( +@group(0) @binding(0) var t_in: array; +@group(0) @binding(1) var t_out: array; + +struct Params { + r: u32, + out_c: u32, + out_h: u32, + out_w: u32, + in_c: u32, + in_h: u32, + in_w: u32, + numel: u32, + input_scale: f32, + inv_output_scale: f32, + input_zero_point: i32, + output_zero_point: i32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +fn unpack_i8(word: u32, byte_idx: u32) -> i32 { + return i32(((word >> (byte_idx * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per packed u32 word (4 int8 outputs); 2D-fold lifts 65535. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (widx >= params.numel / 4u) { + return; + } + var packed: u32 = 0u; + for (var k: u32 = 0u; k < 4u; k = k + 1u) { + // (N, C*r*r, H, W) -> (N, C, H*r, W*r) gather; leading dims collapse to b. + let oi = widx * 4u + k; + let w_out = oi % params.out_w; + let h_out = (oi / params.out_w) % params.out_h; + let c_out = (oi / (params.out_w * params.out_h)) % params.out_c; + let b = oi / (params.out_w * params.out_h * params.out_c); + let w_in = w_out / params.r; + let h_in = h_out / params.r; + let c_in = c_out * params.r * params.r + + (h_out % params.r) * params.r + (w_out % params.r); + let in_flat = + ((b * params.in_c + c_in) * params.in_h + h_in) * params.in_w + w_in; + let q_in = unpack_i8(t_in[in_flat / 4u], in_flat % 4u); + let deq = params.input_scale * f32(q_in - params.input_zero_point); + var q = + i32(round(deq * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (k * 8u)); + } + t_out[widx] = packed; +} +)"; + +inline constexpr uint32_t kQ8taPixelShuffleWorkgroupSizeX = 64; +inline constexpr uint32_t kQ8taPixelShuffleWorkgroupSizeY = 1; +inline constexpr uint32_t kQ8taPixelShuffleWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From beb6881597f4f0b76fa8861b5754a0edd48fe2c5 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:28 -0700 Subject: [PATCH 38/75] [ExecuTorch][WebGPU] Op-tests for q8ta_pixel_shuffle Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 25 ++++++++ .../test/ops/test_q8ta_pixel_shuffle.py | 64 +++++++++++++++++++ 2 files changed, 89 insertions(+) create mode 100644 backends/webgpu/test/ops/test_q8ta_pixel_shuffle.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 043764aabaa..6d6f8c3ee35 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -51,6 +51,9 @@ ) from executorch.backends.webgpu.test.ops.test_q8ta_add import Q8taAddModule from executorch.backends.webgpu.test.ops.test_q8ta_relu import Q8taReluModule +from executorch.backends.webgpu.test.ops.test_q8ta_pixel_shuffle import ( + Q8taPixelShuffleModule, +) from executorch.backends.webgpu.test.ops.test_pixel_shuffle import PixelShuffleModule from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule @@ -1464,6 +1467,28 @@ def case(name, **kw): ) +@register_op_test("q8ta_pixel_shuffle") +def _q8ta_pixel_shuffle_suite() -> WebGPUTestSuite: + # int8 pixel_shuffle (baked int8 const), byte-exact vs CPU eager. Same-qparams + # is a pure gather; diff_qparams exercises the dequant->requant rescale. + def case(name, n_ch, h, w, **kw): + n = n_ch * h * w # [1, n_ch, h, w], n_ch = C*r*r + vals = [((i % 251) - 125) for i in range(n)] # spread across int8 range + return Case(name=name, construct={"x_vals": vals, "shape": (1, n_ch, h, w), **kw}) + + return WebGPUTestSuite( + module_factory=lambda **kw: Q8taPixelShuffleModule(**kw), + cases=[ + case("basic", 4, 2, 2), # r=2, C=1 -> [1,1,4,4] + case("c2", 8, 2, 2), # r=2, C=2 -> [1,2,4,4] + case("r3", 9, 2, 2, upscale_factor=3), # r=3, C=1 -> [1,1,6,6] + case("diff_qparams", 4, 3, 3, output_scale=0.08, output_zp=5), + case("nonzero_zp", 4, 2, 2, input_zp=10, output_zp=-20), + ], + golden_dtype="float32", + ) + + @register_op_test("dequantize_per_tensor") def _dequant_const_suite() -> WebGPUTestSuite: # Independent GPU dequantize check: baked int8 const vs torch dequant (breaks diff --git a/backends/webgpu/test/ops/test_q8ta_pixel_shuffle.py b/backends/webgpu/test/ops/test_q8ta_pixel_shuffle.py new file mode 100644 index 00000000000..41ac851d846 --- /dev/null +++ b/backends/webgpu/test/ops/test_q8ta_pixel_shuffle.py @@ -0,0 +1,64 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`et_vk.q8ta_pixel_shuffle` module + configs: int8 pixel_shuffle. + +The int8 input `[N, C*r*r, H, W]` is a baked constant, so the op is exercised +alone with a byte-exact int8 golden vs the CPU eager op (gather then dequant -> +requant). Note the schema's 4th arg is `output_inv_scale` (already 1/scale). +""" + +import unittest + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 registers et_vk ops + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + + +class Q8taPixelShuffleModule(torch.nn.Module): + def __init__( + self, + x_vals, + shape, + input_scale=0.05, + input_zp=-3, + output_scale=0.05, + output_zp=-3, + upscale_factor=2, + ): + super().__init__() + self.register_buffer("x", torch.tensor(x_vals, dtype=torch.int8).reshape(shape)) + self.input_scale, self.input_zp = input_scale, input_zp + self.output_inv_scale = 1.0 / output_scale + self.output_zp, self.upscale_factor = output_zp, upscale_factor + + def forward(self) -> torch.Tensor: + return torch.ops.et_vk.q8ta_pixel_shuffle( + self.x, + self.input_scale, + self.input_zp, + self.output_inv_scale, + self.output_zp, + self.upscale_factor, + ) + + +class Q8taPixelShuffleTest(unittest.TestCase): + def test_delegates(self) -> None: + m = Q8taPixelShuffleModule(list(range(-8, 8)), (1, 4, 2, 2)) + et = to_edge_transform_and_lower( + torch.export.export(m, ()), partitioner=[VulkanPartitioner()] + ).to_executorch() + self.assertTrue( + any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + ) From 740330045ada549a61aa7c7cc989faae9770da56 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:29 -0700 Subject: [PATCH 39/75] [ExecuTorch][WebGPU] Add q8ta_linear op (int8 quantized linear) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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 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/) --- .../runtime/ops/q8ta_linear/Q8taLinear.cpp | 235 ++++++++++++++++++ .../runtime/ops/q8ta_linear/q8ta_linear.wgsl | 89 +++++++ .../ops/q8ta_linear/q8ta_linear_wgsl.h | 113 +++++++++ 3 files changed, 437 insertions(+) create mode 100644 backends/webgpu/runtime/ops/q8ta_linear/Q8taLinear.cpp create mode 100644 backends/webgpu/runtime/ops/q8ta_linear/q8ta_linear.wgsl create mode 100644 backends/webgpu/runtime/ops/q8ta_linear/q8ta_linear_wgsl.h diff --git a/backends/webgpu/runtime/ops/q8ta_linear/Q8taLinear.cpp b/backends/webgpu/runtime/ops/q8ta_linear/Q8taLinear.cpp new file mode 100644 index 00000000000..fdf2eb38542 --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_linear/Q8taLinear.cpp @@ -0,0 +1,235 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct Q8taLinearParams { + uint32_t M; + uint32_t N; + uint32_t K; + int32_t input_zero_point; + int32_t output_zero_point; + float input_scale; + float inv_output_scale; + uint32_t has_bias; +}; +static_assert( + sizeof(Q8taLinearParams) == 32, + "Q8taLinearParams must match the WGSL Params struct (32 bytes)"); + +// int8-act x int8-weight linear->int8; mirrors Vulkan q8ta_linear (act=none). +void q8ta_linear_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [x,in_scale,in_zp,w,w_sums,w_scales,out_scale,out_zp,bias,act,out]. + const int in_id = args.at(0); + const int weight_id = args.at(3); + const int scales_id = args.at(5); + const int bias_id = args.at(8); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(weight_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(scales_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("q8ta_linear: in/weight/scales/out not tensor"); + } + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + // Only activation="none" is supported; relu-fused would compute wrong output. + const int act_id = args.at(args.size() - 2); + if (graph.get_value_type(act_id) != WebGPUGraph::ValueType::String || + graph.get_string(act_id) != "none") { + throw std::runtime_error("q8ta_linear: only activation='none' supported"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& weight_tensor = graph.get_tensor(weight_id); + const auto& scales_tensor = graph.get_tensor(scales_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || weight_tensor.buffer == nullptr || + scales_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("q8ta_linear: null buffer binding"); + } + if (weight_tensor.dims.size() != 2) { + throw std::runtime_error("q8ta_linear: weight must be 2D [N, K]"); + } + + const double input_scale = graph.get_double(args.at(1)); + const int input_zero_point = graph.get_int(args.at(2)); + const double output_scale = graph.get_double(args.at(6)); + const int output_zero_point = graph.get_int(args.at(7)); + + const uint64_t N = static_cast(weight_tensor.dims.at(0)); + const uint64_t K = static_cast(weight_tensor.dims.at(1)); + if (K == 0 || in_tensor.dims.empty() || + static_cast(in_tensor.dims.back()) != K) { + throw std::runtime_error("q8ta_linear: input last dim must equal K"); + } + uint64_t in_numel = 1; + for (int64_t d : in_tensor.dims) { + in_numel *= static_cast(d); + } + const uint64_t M = in_numel / K; + if (M == 0 || N == 0 || out_tensor.dims.empty() || + static_cast(out_tensor.dims.back()) != N) { + throw std::runtime_error("q8ta_linear: bad M/N shape"); + } + if (M > UINT32_MAX || N > UINT32_MAX || K > UINT32_MAX) { + throw std::runtime_error("q8ta_linear: dim exceeds u32"); + } + // N % 4 == 0 makes each row's TN=4 outputs land in one aligned int8 word. + if (N % 4 != 0) { + throw std::runtime_error("q8ta_linear: N must be a multiple of 4"); + } + // int8 x / weight / out bound as array: their numels must be %4 == 0. + if (!in_tensor.is_int8 || in_tensor.nbytes != M * K || (M * K) % 4 != 0 || + !weight_tensor.is_int8 || weight_tensor.nbytes != N * K || + (N * K) % 4 != 0 || !out_tensor.is_int8 || out_tensor.nbytes != M * N) { + throw std::runtime_error("q8ta_linear: int8 x/weight/out size mismatch"); + } + if (scales_tensor.nbytes != N * sizeof(float)) { + throw std::runtime_error("q8ta_linear: weight_scales must be fp32 [N]"); + } + if (has_bias) { + const auto& b = graph.get_tensor(bias_id); + if (b.buffer == nullptr || b.nbytes != N * sizeof(float)) { + throw std::runtime_error("q8ta_linear: bias must be fp32 [N]"); + } + } + + Q8taLinearParams params = {}; + params.M = static_cast(M); + params.N = static_cast(N); + params.K = static_cast(K); + params.input_zero_point = static_cast(input_zero_point); + params.output_zero_point = static_cast(output_zero_point); + params.input_scale = static_cast(input_scale); + // Reciprocal in double then cast, matching torch's f32(1.0 / f64(scale)). + params.inv_output_scale = static_cast(1.0 / output_scale); + params.has_bias = has_bias ? 1u : 0u; + + const uint64_t n_tiles_u64 = ((M + 3) / 4) * (N / 4); + if (n_tiles_u64 > UINT32_MAX) { + throw std::runtime_error("q8ta_linear: dispatch tile count exceeds u32"); + } + const uint32_t n_tiles = static_cast(n_tiles_u64); + uint32_t wg_size = + utils::clamp_workgroup_size(device, kQ8taLinearWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, n_tiles, wg_size, "q8ta_linear"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(Q8taLinearParams)); + graph.add_uniform_buffer_bytes(sizeof(Q8taLinearParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kQ8taLinearWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[6] = {}; + for (int i = 0; i < 6; i++) { + entries[i].binding = static_cast(i); + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = (i == 0) ? WGPUBufferBindingType_Storage + : (i == 5) ? WGPUBufferBindingType_Uniform + : WGPUBufferBindingType_ReadOnlyStorage; + } + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 6; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + // No-bias: bind scales as an unread placeholder (has_bias gates the read). + WGPUBuffer bias_buf = + has_bias ? graph.get_tensor(bias_id).buffer : scales_tensor.buffer; + const uint64_t bias_size = + has_bias ? graph.get_tensor(bias_id).nbytes : scales_tensor.nbytes; + + WGPUBindGroupEntry bg[6] = {}; + bg[0].binding = 0; + bg[0].buffer = out_tensor.buffer; + bg[0].size = out_tensor.nbytes; + bg[1].binding = 1; + bg[1].buffer = in_tensor.buffer; + bg[1].size = in_tensor.nbytes; + bg[2].binding = 2; + bg[2].buffer = weight_tensor.buffer; + bg[2].size = weight_tensor.nbytes; + bg[3].binding = 3; + bg[3].buffer = scales_tensor.buffer; + bg[3].size = scales_tensor.nbytes; + bg[4].binding = 4; + bg[4].buffer = bias_buf; + bg[4].size = bias_size; + bg[5].binding = 5; + bg[5].buffer = params_buf; + bg[5].size = sizeof(Q8taLinearParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 6; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "q8ta_linear", + workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.q8ta_linear.default, q8ta_linear_impl); + // GEMV (M==1) variant: identical schema/semantics, handled by the same GEMM. + WEBGPU_REGISTER_OP(et_vk.q8ta_linear_gemv.default, q8ta_linear_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/q8ta_linear/q8ta_linear.wgsl b/backends/webgpu/runtime/ops/q8ta_linear/q8ta_linear.wgsl new file mode 100644 index 00000000000..b4448f7f44a --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_linear/q8ta_linear.wgsl @@ -0,0 +1,89 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + M: u32, + N: u32, + K: u32, + input_zero_point: i32, + output_zero_point: i32, + input_scale: f32, + inv_output_scale: f32, + has_bias: u32, +} +@group(0) @binding(5) var params: Params; + +override wg_size: u32 = 64u; + +const TM: u32 = 4u; +const TN: u32 = 4u; +const TILE_ELEMS: u32 = TM * TN; + +fn unpack_i8(buf_idx: u32, arr_word: u32) -> i32 { + // arr_word is the value at t[buf_idx>>2]; extract the (buf_idx&3) byte. + return i32(((arr_word >> ((buf_idx & 3u) * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded tile index (lifts the 65535 1D-dispatch cap for large M*N). + let tile = gid.x + gid.y * (num_workgroups.x * wg_size); + let nct = (params.N + TN - 1u) / TN; + if (tile >= ((params.M + TM - 1u) / TM) * nct) { + return; + } + let m0 = (tile / nct) * TM; + let n0 = (tile % nct) * TN; + + var acc: array; + for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) { + acc[i] = 0; + } + + var k: u32 = 0u; + loop { + if (k >= params.K) { + break; + } + var xq: array; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m_eff = min(m0 + ml, params.M - 1u); + let bi = m_eff * params.K + k; + xq[ml] = unpack_i8(bi, t_x[bi >> 2u]) - params.input_zero_point; + } + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n_eff = min(n0 + nl, params.N - 1u); + let bi = n_eff * params.K + k; + let w = unpack_i8(bi, t_weight[bi >> 2u]); + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + acc[ml * TN + nl] = acc[ml * TN + nl] + xq[ml] * w; + } + } + k = k + 1u; + } + + // Requantize to int8; N%4==0 (host) so each row's 4 cols pack into one word. + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m = m0 + ml; + if (m >= params.M) { + continue; + } + var packed: u32 = 0u; + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n = n0 + nl; + var v = f32(acc[ml * TN + nl]) * params.input_scale * t_scales[n]; + if (params.has_bias != 0u) { + v = v + t_bias[n]; + } + var q = i32(round(v * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (nl * 8u)); + } + t_out[(m * params.N + n0) >> 2u] = packed; + } +} diff --git a/backends/webgpu/runtime/ops/q8ta_linear/q8ta_linear_wgsl.h b/backends/webgpu/runtime/ops/q8ta_linear/q8ta_linear_wgsl.h new file mode 100644 index 00000000000..4a5cd9bb40c --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_linear/q8ta_linear_wgsl.h @@ -0,0 +1,113 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from q8ta_linear.wgsl - DO NOT EDIT. +// wgsl-sha256: b08067158b832e6e25f55024f9aa3a4eded697f1d77b8d638fc96b1ecdbd5512 +inline constexpr const char* kQ8taLinearWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + M: u32, + N: u32, + K: u32, + input_zero_point: i32, + output_zero_point: i32, + input_scale: f32, + inv_output_scale: f32, + has_bias: u32, +} +@group(0) @binding(5) var params: Params; + +override wg_size: u32 = 64u; + +const TM: u32 = 4u; +const TN: u32 = 4u; +const TILE_ELEMS: u32 = TM * TN; + +fn unpack_i8(buf_idx: u32, arr_word: u32) -> i32 { + // arr_word is the value at t[buf_idx>>2]; extract the (buf_idx&3) byte. + return i32(((arr_word >> ((buf_idx & 3u) * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded tile index (lifts the 65535 1D-dispatch cap for large M*N). + let tile = gid.x + gid.y * (num_workgroups.x * wg_size); + let nct = (params.N + TN - 1u) / TN; + if (tile >= ((params.M + TM - 1u) / TM) * nct) { + return; + } + let m0 = (tile / nct) * TM; + let n0 = (tile % nct) * TN; + + var acc: array; + for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) { + acc[i] = 0; + } + + var k: u32 = 0u; + loop { + if (k >= params.K) { + break; + } + var xq: array; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m_eff = min(m0 + ml, params.M - 1u); + let bi = m_eff * params.K + k; + xq[ml] = unpack_i8(bi, t_x[bi >> 2u]) - params.input_zero_point; + } + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n_eff = min(n0 + nl, params.N - 1u); + let bi = n_eff * params.K + k; + let w = unpack_i8(bi, t_weight[bi >> 2u]); + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + acc[ml * TN + nl] = acc[ml * TN + nl] + xq[ml] * w; + } + } + k = k + 1u; + } + + // Requantize to int8; N%4==0 (host) so each row's 4 cols pack into one word. + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m = m0 + ml; + if (m >= params.M) { + continue; + } + var packed: u32 = 0u; + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n = n0 + nl; + var v = f32(acc[ml * TN + nl]) * params.input_scale * t_scales[n]; + if (params.has_bias != 0u) { + v = v + t_bias[n]; + } + var q = i32(round(v * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (nl * 8u)); + } + t_out[(m * params.N + n0) >> 2u] = packed; + } +} +)"; + +inline constexpr uint32_t kQ8taLinearWorkgroupSizeX = 64; +inline constexpr uint32_t kQ8taLinearWorkgroupSizeY = 1; +inline constexpr uint32_t kQ8taLinearWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From b1482d7be04faf5f757030f066d1ef9bf2db96ff Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:29 -0700 Subject: [PATCH 40/75] [ExecuTorch][WebGPU] Op-tests for q8ta_linear MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 26 +++++++++ backends/webgpu/test/ops/test_q8ta_linear.py | 61 ++++++++++++++++++++ 2 files changed, 87 insertions(+) create mode 100644 backends/webgpu/test/ops/test_q8ta_linear.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 6d6f8c3ee35..7dd03c265aa 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -54,6 +54,9 @@ from executorch.backends.webgpu.test.ops.test_q8ta_pixel_shuffle import ( Q8taPixelShuffleModule, ) +from executorch.backends.webgpu.test.ops.test_q8ta_linear import ( + make_q8ta_linear_module, +) from executorch.backends.webgpu.test.ops.test_pixel_shuffle import PixelShuffleModule from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule @@ -1489,6 +1492,29 @@ def case(name, n_ch, h, w, **kw): ) +@register_op_test("q8ta_linear") +def _q8ta_linear_suite() -> WebGPUTestSuite: + # XNNPACK-static-quantized nn.Linear -> delegated quantize->q8ta_linear-> + # dequantize (C0 + the new op). Golden = converted eager (fp32 fake-quant). + def case(name, m, k, n, **kw): + return Case( + name=name, construct={"k": k, "n": n, "m": m, **kw}, inputs=((m, k),) + ) + + return WebGPUTestSuite( + module_factory=lambda **kw: make_q8ta_linear_module(**kw), + cases=[ + case("basic", 4, 16, 8), + case("gemv", 1, 16, 8), # M==1 + case("k32", 8, 32, 4), + case("no_bias", 4, 16, 8, bias=False), + ], + golden_dtype="float32", + atol=1e-3, + rtol=1e-3, + ) + + @register_op_test("dequantize_per_tensor") def _dequant_const_suite() -> WebGPUTestSuite: # Independent GPU dequantize check: baked int8 const vs torch dequant (breaks diff --git a/backends/webgpu/test/ops/test_q8ta_linear.py b/backends/webgpu/test/ops/test_q8ta_linear.py new file mode 100644 index 00000000000..fd41543d98f --- /dev/null +++ b/backends/webgpu/test/ops/test_q8ta_linear.py @@ -0,0 +1,61 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`et_vk.q8ta_linear` module + configs: int8-activation x int8-weight linear. + +The op is produced by running a plain `nn.Linear` through XNNPACK static PT2E +(per-channel weight, static per-tensor activation): the converted module lowers +to a delegated `quantize_per_tensor -> q8ta_linear -> dequantize_per_tensor` +subgraph (the quantize/dequantize are the landed C0 ops). The `module_factory` +returns the CONVERTED module, so the op-test framework goldens the WebGPU output +against the converted eager (fp32, the fake-quant reference), exercising all three +ops end-to-end. `activation` is scoped to "none" (the XNNPACK-static default). +""" + +import unittest + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 registers et_vk ops + +import torch +import torch.nn as nn + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( + get_symmetric_quantization_config, + XNNPACKQuantizer, +) +from executorch.exir import to_edge_transform_and_lower +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + +def make_q8ta_linear_module(k, n, m, bias=True, seed=0): + torch.manual_seed(seed) + lin = nn.Linear(k, n, bias=bias).eval() + ex = (torch.randn(m, k),) + q = XNNPACKQuantizer().set_global( + get_symmetric_quantization_config(is_per_channel=True, is_dynamic=False) + ) + prepared = prepare_pt2e(torch.export.export(lin, ex).module(), q) + prepared(*ex) # calibrate + return convert_pt2e(prepared) + + +class Q8taLinearTest(unittest.TestCase): + def test_delegates(self) -> None: + m = make_q8ta_linear_module(16, 8, 4) + et = to_edge_transform_and_lower( + torch.export.export(m, (torch.randn(4, 16),)), + partitioner=[VulkanPartitioner()], + ).to_executorch() + self.assertTrue( + any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + ) + # The delegate must contain the new op (+ the landed C0 quant ops). + self.assertIn(b"q8ta_linear", et.buffer) From a429dd3c9bdea4bc49153195bf100a847aca8143 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:30 -0700 Subject: [PATCH 41/75] [ExecuTorch][WebGPU] Add q8ta_conv2d_pw op (int8 pointwise conv) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../ops/q8ta_conv2d_pw/Q8taConv2dPw.cpp | 244 ++++++++++++++++++ .../ops/q8ta_conv2d_pw/q8ta_conv2d_pw.wgsl | 78 ++++++ .../ops/q8ta_conv2d_pw/q8ta_conv2d_pw_wgsl.h | 102 ++++++++ 3 files changed, 424 insertions(+) create mode 100644 backends/webgpu/runtime/ops/q8ta_conv2d_pw/Q8taConv2dPw.cpp create mode 100644 backends/webgpu/runtime/ops/q8ta_conv2d_pw/q8ta_conv2d_pw.wgsl create mode 100644 backends/webgpu/runtime/ops/q8ta_conv2d_pw/q8ta_conv2d_pw_wgsl.h diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d_pw/Q8taConv2dPw.cpp b/backends/webgpu/runtime/ops/q8ta_conv2d_pw/Q8taConv2dPw.cpp new file mode 100644 index 00000000000..5dbf91558ab --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_conv2d_pw/Q8taConv2dPw.cpp @@ -0,0 +1,244 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct Q8taConvPwParams { + uint32_t N; + uint32_t OC; + uint32_t IC; + uint32_t H; + uint32_t W; + int32_t input_zero_point; + int32_t output_zero_point; + float input_scale; + float inv_output_scale; + uint32_t has_bias; + uint32_t pad0; + uint32_t pad1; +}; +static_assert( + sizeof(Q8taConvPwParams) == 48, + "Q8taConvPwParams must match the WGSL Params struct (48 bytes)"); + +bool is_unit_pair(const std::vector& v, int64_t a, int64_t b) { + return v.size() == 2 && v[0] == a && v[1] == b; +} + +// int8 1x1 conv; per-position channel dot; mirrors Vulkan q8ta_conv2d_pw. +void q8ta_conv2d_pw_impl(WebGPUGraph& graph, const std::vector& args) { + const int in_id = args.at(0); + const int weight_id = args.at(3); + const int scales_id = args.at(5); + const int bias_id = args.at(8); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(weight_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(scales_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("q8ta_conv2d_pw: in/weight/scales/out not tensor"); + } + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + + // Pointwise-standard config only; anything else is fail-loud. + if (!is_unit_pair(graph.get_int_list(args.at(10)), 1, 1) || + !is_unit_pair(graph.get_int_list(args.at(11)), 0, 0) || + !is_unit_pair(graph.get_int_list(args.at(12)), 1, 1) || + graph.get_int(args.at(13)) != 1) { + throw std::runtime_error("q8ta_conv2d_pw: only stride1/pad0/groups1"); + } + const int act_id = args.at(args.size() - 2); + if (graph.get_value_type(act_id) != WebGPUGraph::ValueType::String || + graph.get_string(act_id) != "none") { + throw std::runtime_error( + "q8ta_conv2d_pw: only activation='none' supported"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& weight_tensor = graph.get_tensor(weight_id); + const auto& scales_tensor = graph.get_tensor(scales_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || weight_tensor.buffer == nullptr || + scales_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("q8ta_conv2d_pw: null buffer binding"); + } + if (in_tensor.dims.size() != 4 || out_tensor.dims.size() != 4 || + weight_tensor.dims.size() != 2) { + throw std::runtime_error("q8ta_conv2d_pw: in/out must be 4D, weight 2D"); + } + + const double input_scale = graph.get_double(args.at(1)); + const int input_zero_point = graph.get_int(args.at(2)); + const double output_scale = graph.get_double(args.at(6)); + const int output_zero_point = graph.get_int(args.at(7)); + + const uint64_t N = static_cast(in_tensor.dims.at(0)); + const uint64_t IC = static_cast(in_tensor.dims.at(1)); + const uint64_t H = static_cast(in_tensor.dims.at(2)); + const uint64_t W = static_cast(in_tensor.dims.at(3)); + const uint64_t OC = static_cast(weight_tensor.dims.at(0)); + if (static_cast(weight_tensor.dims.at(1)) != IC) { + throw std::runtime_error("q8ta_conv2d_pw: weight must be [OC, IC]"); + } + const uint64_t out_numel = N * OC * H * W; + if (out_numel == 0) { + throw std::runtime_error("q8ta_conv2d_pw: output is empty"); + } + if (W % 4 != 0) { + throw std::runtime_error("q8ta_conv2d_pw: W must be a multiple of 4"); + } + if (N * OC * H * W > UINT32_MAX || N * IC * H * W > UINT32_MAX) { + throw std::runtime_error("q8ta_conv2d_pw: numel exceeds u32"); + } + if (!in_tensor.is_int8 || in_tensor.nbytes != N * IC * H * W || + (N * IC * H * W) % 4 != 0 || !weight_tensor.is_int8 || + weight_tensor.nbytes != OC * IC || (OC * IC) % 4 != 0 || + !out_tensor.is_int8 || out_tensor.nbytes != out_numel) { + throw std::runtime_error( + "q8ta_conv2d_pw: int8 in/weight/out size mismatch"); + } + if (scales_tensor.nbytes != OC * sizeof(float)) { + throw std::runtime_error("q8ta_conv2d_pw: weight_scales must be fp32 [OC]"); + } + if (has_bias) { + const auto& b = graph.get_tensor(bias_id); + if (b.buffer == nullptr || b.nbytes != OC * sizeof(float)) { + throw std::runtime_error("q8ta_conv2d_pw: bias must be fp32 [OC]"); + } + } + + Q8taConvPwParams params = {}; + params.N = static_cast(N); + params.OC = static_cast(OC); + params.IC = static_cast(IC); + params.H = static_cast(H); + params.W = static_cast(W); + params.input_zero_point = static_cast(input_zero_point); + params.output_zero_point = static_cast(output_zero_point); + params.input_scale = static_cast(input_scale); + // Reciprocal in double then cast, matching torch's f32(1.0 / f64(scale)). + params.inv_output_scale = static_cast(1.0 / output_scale); + params.has_bias = has_bias ? 1u : 0u; + + const uint32_t num_words = static_cast(out_numel / 4); + uint32_t wg_size = + utils::clamp_workgroup_size(device, kQ8taConv2dPwWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, num_words, wg_size, "q8ta_conv2d_pw"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(Q8taConvPwParams)); + graph.add_uniform_buffer_bytes(sizeof(Q8taConvPwParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kQ8taConv2dPwWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[6] = {}; + for (int i = 0; i < 6; i++) { + entries[i].binding = static_cast(i); + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = (i == 0) ? WGPUBufferBindingType_Storage + : (i == 5) ? WGPUBufferBindingType_Uniform + : WGPUBufferBindingType_ReadOnlyStorage; + } + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 6; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + // No-bias: bind scales as an unread placeholder (has_bias gates the read). + WGPUBuffer bias_buf = + has_bias ? graph.get_tensor(bias_id).buffer : scales_tensor.buffer; + const uint64_t bias_size = + has_bias ? graph.get_tensor(bias_id).nbytes : scales_tensor.nbytes; + + WGPUBindGroupEntry bg[6] = {}; + bg[0].binding = 0; + bg[0].buffer = out_tensor.buffer; + bg[0].size = out_tensor.nbytes; + bg[1].binding = 1; + bg[1].buffer = in_tensor.buffer; + bg[1].size = in_tensor.nbytes; + bg[2].binding = 2; + bg[2].buffer = weight_tensor.buffer; + bg[2].size = weight_tensor.nbytes; + bg[3].binding = 3; + bg[3].buffer = scales_tensor.buffer; + bg[3].size = scales_tensor.nbytes; + bg[4].binding = 4; + bg[4].buffer = bias_buf; + bg[4].size = bias_size; + bg[5].binding = 5; + bg[5].buffer = params_buf; + bg[5].size = sizeof(Q8taConvPwParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 6; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "q8ta_conv2d_pw", + workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.q8ta_conv2d_pw.default, q8ta_conv2d_pw_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d_pw/q8ta_conv2d_pw.wgsl b/backends/webgpu/runtime/ops/q8ta_conv2d_pw/q8ta_conv2d_pw.wgsl new file mode 100644 index 00000000000..e1f885984fb --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_conv2d_pw/q8ta_conv2d_pw.wgsl @@ -0,0 +1,78 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + N: u32, + OC: u32, + IC: u32, + H: u32, + W: u32, + input_zero_point: i32, + output_zero_point: i32, + input_scale: f32, + inv_output_scale: f32, + has_bias: u32, + pad0: u32, + pad1: u32, +} +@group(0) @binding(5) var params: Params; + +override wg_size: u32 = 64u; + +fn unpack_i8(bi: u32, word: u32) -> i32 { + return i32(((word >> ((bi & 3u) * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per output word = 4 W-positions of fixed (n,oc,h); W%4==0. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + let words = (params.N * params.OC * params.H * params.W) / 4u; + if (widx >= words) { + return; + } + let flat0 = widx * 4u; + let w0 = flat0 % params.W; + var r = flat0 / params.W; + let h = r % params.H; + r = r / params.H; + let oc = r % params.OC; + let n = r / params.OC; + + var acc: array; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + acc[j] = 0; + } + // 1x1 conv = per-position channel dot over input channels. + for (var ic: u32 = 0u; ic < params.IC; ic = ic + 1u) { + let wbi = oc * params.IC + ic; + let wv = unpack_i8(wbi, t_weight[wbi >> 2u]); + let base = ((n * params.IC + ic) * params.H + h) * params.W + w0; + // W%4==0 and w0%4==0 => base%4==0, so all 4 j share one input word. + let xword = t_x[base >> 2u]; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let xbi = base + j; + acc[j] = acc[j] + + (unpack_i8(xbi, xword) - params.input_zero_point) * wv; + } + } + + let wscale = t_scales[oc]; + let bias = t_bias[oc]; + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + var v = f32(acc[j]) * params.input_scale * wscale; + if (params.has_bias != 0u) { + v = v + bias; + } + var q = i32(round(v * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (j * 8u)); + } + t_out[widx] = packed; +} diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d_pw/q8ta_conv2d_pw_wgsl.h b/backends/webgpu/runtime/ops/q8ta_conv2d_pw/q8ta_conv2d_pw_wgsl.h new file mode 100644 index 00000000000..c2d7976ee3b --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_conv2d_pw/q8ta_conv2d_pw_wgsl.h @@ -0,0 +1,102 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from q8ta_conv2d_pw.wgsl - DO NOT EDIT. +// wgsl-sha256: b4d6ddfd2e0101bbff9a668895014c8822aa7d92147a982c1e6fe2c3bff847d0 +inline constexpr const char* kQ8taConv2dPwWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + N: u32, + OC: u32, + IC: u32, + H: u32, + W: u32, + input_zero_point: i32, + output_zero_point: i32, + input_scale: f32, + inv_output_scale: f32, + has_bias: u32, + pad0: u32, + pad1: u32, +} +@group(0) @binding(5) var params: Params; + +override wg_size: u32 = 64u; + +fn unpack_i8(bi: u32, word: u32) -> i32 { + return i32(((word >> ((bi & 3u) * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per output word = 4 W-positions of fixed (n,oc,h); W%4==0. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + let words = (params.N * params.OC * params.H * params.W) / 4u; + if (widx >= words) { + return; + } + let flat0 = widx * 4u; + let w0 = flat0 % params.W; + var r = flat0 / params.W; + let h = r % params.H; + r = r / params.H; + let oc = r % params.OC; + let n = r / params.OC; + + var acc: array; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + acc[j] = 0; + } + // 1x1 conv = per-position channel dot over input channels. + for (var ic: u32 = 0u; ic < params.IC; ic = ic + 1u) { + let wbi = oc * params.IC + ic; + let wv = unpack_i8(wbi, t_weight[wbi >> 2u]); + let base = ((n * params.IC + ic) * params.H + h) * params.W + w0; + // W%4==0 and w0%4==0 => base%4==0, so all 4 j share one input word. + let xword = t_x[base >> 2u]; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let xbi = base + j; + acc[j] = acc[j] + + (unpack_i8(xbi, xword) - params.input_zero_point) * wv; + } + } + + let wscale = t_scales[oc]; + let bias = t_bias[oc]; + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + var v = f32(acc[j]) * params.input_scale * wscale; + if (params.has_bias != 0u) { + v = v + bias; + } + var q = i32(round(v * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (j * 8u)); + } + t_out[widx] = packed; +} +)"; + +inline constexpr uint32_t kQ8taConv2dPwWorkgroupSizeX = 64; +inline constexpr uint32_t kQ8taConv2dPwWorkgroupSizeY = 1; +inline constexpr uint32_t kQ8taConv2dPwWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 49951665d72dae8a2b3ab5af75dc5c0bdd3dc965 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:30 -0700 Subject: [PATCH 42/75] [ExecuTorch][WebGPU] Op-tests for q8ta_conv2d_pw MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 27 +++++++++ .../webgpu/test/ops/test_q8ta_conv2d_pw.py | 58 +++++++++++++++++++ 2 files changed, 85 insertions(+) create mode 100644 backends/webgpu/test/ops/test_q8ta_conv2d_pw.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 7dd03c265aa..25d1574e1a4 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -57,6 +57,9 @@ from executorch.backends.webgpu.test.ops.test_q8ta_linear import ( make_q8ta_linear_module, ) +from executorch.backends.webgpu.test.ops.test_q8ta_conv2d_pw import ( + make_q8ta_conv2d_pw_module, +) from executorch.backends.webgpu.test.ops.test_pixel_shuffle import PixelShuffleModule from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule @@ -1515,6 +1518,30 @@ def case(name, m, k, n, **kw): ) +@register_op_test("q8ta_conv2d_pw") +def _q8ta_conv2d_pw_suite() -> WebGPUTestSuite: + # Golden = XNNPACK-static-PT2E converted eager (fp32); W%4==0 for output pack. + def case(name, ic, oc, h, w, n=1, **kw): + return Case( + name=name, + construct={"ic": ic, "oc": oc, "h": h, "w": w, "n": n, **kw}, + inputs=((n, ic, h, w),), + ) + + return WebGPUTestSuite( + module_factory=lambda **kw: make_q8ta_conv2d_pw_module(**kw), + cases=[ + case("basic", 4, 8, 6, 8), + case("ic8", 8, 4, 4, 4), + case("no_bias", 4, 8, 6, 8, bias=False), + case("batch2", 4, 8, 6, 8, n=2), + ], + golden_dtype="float32", + atol=1e-3, + rtol=1e-3, + ) + + @register_op_test("dequantize_per_tensor") def _dequant_const_suite() -> WebGPUTestSuite: # Independent GPU dequantize check: baked int8 const vs torch dequant (breaks diff --git a/backends/webgpu/test/ops/test_q8ta_conv2d_pw.py b/backends/webgpu/test/ops/test_q8ta_conv2d_pw.py new file mode 100644 index 00000000000..838142deebe --- /dev/null +++ b/backends/webgpu/test/ops/test_q8ta_conv2d_pw.py @@ -0,0 +1,58 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`et_vk.q8ta_conv2d_pw` module + configs: int8 pointwise (1x1) conv. + +Produced by running a plain 1x1 `nn.Conv2d` through XNNPACK static PT2E: the +converted module lowers to a delegated `quantize_per_tensor -> q8ta_conv2d_pw -> +dequantize_per_tensor` subgraph (quantize/dequantize = the landed C0 ops). The +`module_factory` returns the CONVERTED module, so the op-test framework goldens +the WebGPU output against the converted eager (fp32 fake-quant), e2e. +""" + +import unittest + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 registers et_vk ops + +import torch +import torch.nn as nn + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( + get_symmetric_quantization_config, + XNNPACKQuantizer, +) +from executorch.exir import to_edge_transform_and_lower +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + +def make_q8ta_conv2d_pw_module(ic, oc, h, w, n=1, bias=True, seed=0): + torch.manual_seed(seed) + conv = nn.Conv2d(ic, oc, 1, bias=bias).eval() + ex = (torch.randn(n, ic, h, w),) + q = XNNPACKQuantizer().set_global( + get_symmetric_quantization_config(is_per_channel=True, is_dynamic=False) + ) + prepared = prepare_pt2e(torch.export.export(conv, ex).module(), q) + prepared(*ex) # calibrate + return convert_pt2e(prepared) + + +class Q8taConv2dPwTest(unittest.TestCase): + def test_delegates(self) -> None: + m = make_q8ta_conv2d_pw_module(4, 8, 6, 8) + et = to_edge_transform_and_lower( + torch.export.export(m, (torch.randn(1, 4, 6, 8),)), + partitioner=[VulkanPartitioner()], + ).to_executorch() + self.assertTrue( + any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + ) + self.assertIn(b"q8ta_conv2d_pw", et.buffer) From 73ae18e6aac94438fc14bdd203474ad722f17b0a Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:30 -0700 Subject: [PATCH 43/75] [ExecuTorch][WebGPU] Add q8ta_conv2d_dw op (int8 depthwise conv) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../ops/q8ta_conv2d_dw/Q8taConv2dDw.cpp | 278 ++++++++++++++++++ .../ops/q8ta_conv2d_dw/q8ta_conv2d_dw.wgsl | 94 ++++++ .../ops/q8ta_conv2d_dw/q8ta_conv2d_dw_wgsl.h | 118 ++++++++ 3 files changed, 490 insertions(+) create mode 100644 backends/webgpu/runtime/ops/q8ta_conv2d_dw/Q8taConv2dDw.cpp create mode 100644 backends/webgpu/runtime/ops/q8ta_conv2d_dw/q8ta_conv2d_dw.wgsl create mode 100644 backends/webgpu/runtime/ops/q8ta_conv2d_dw/q8ta_conv2d_dw_wgsl.h diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d_dw/Q8taConv2dDw.cpp b/backends/webgpu/runtime/ops/q8ta_conv2d_dw/Q8taConv2dDw.cpp new file mode 100644 index 00000000000..e0b0a5e548c --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_conv2d_dw/Q8taConv2dDw.cpp @@ -0,0 +1,278 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct Q8taConvDwParams { + uint32_t N; + uint32_t C; + uint32_t H_in; + uint32_t W_in; + uint32_t H_out; + uint32_t W_out; + uint32_t Kh; + uint32_t Kw; + uint32_t stride_h; + uint32_t stride_w; + uint32_t pad_h; + uint32_t pad_w; + uint32_t dil_h; + uint32_t dil_w; + int32_t input_zero_point; + int32_t output_zero_point; + float input_scale; + float inv_output_scale; + uint32_t has_bias; + uint32_t pad0; +}; +static_assert( + sizeof(Q8taConvDwParams) == 80, + "Q8taConvDwParams must match the WGSL Params struct (80 bytes)"); + +std::pair pair_or_throw( + const std::vector& v, + const char* msg) { + if (v.size() != 2) { + throw std::runtime_error(msg); + } + return {v.at(0), v.at(1)}; +} + +// int8 depthwise conv; per-channel windowed dot; mirrors Vulkan q8ta_conv2d_dw. +void q8ta_conv2d_dw_impl(WebGPUGraph& graph, const std::vector& args) { + const int in_id = args.at(0); + const int weight_id = args.at(3); + const int scales_id = args.at(5); + const int bias_id = args.at(8); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(weight_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(scales_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("q8ta_conv2d_dw: in/weight/scales/out not tensor"); + } + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + + const int act_id = args.at(args.size() - 2); + if (graph.get_value_type(act_id) != WebGPUGraph::ValueType::String || + graph.get_string(act_id) != "none") { + throw std::runtime_error( + "q8ta_conv2d_dw: only activation='none' supported"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& weight_tensor = graph.get_tensor(weight_id); + const auto& scales_tensor = graph.get_tensor(scales_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || weight_tensor.buffer == nullptr || + scales_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("q8ta_conv2d_dw: null buffer binding"); + } + if (in_tensor.dims.size() != 4 || out_tensor.dims.size() != 4 || + weight_tensor.dims.size() != 3) { + throw std::runtime_error("q8ta_conv2d_dw: in/out must be 4D, weight 3D"); + } + + const double input_scale = graph.get_double(args.at(1)); + const int input_zero_point = graph.get_int(args.at(2)); + const double output_scale = graph.get_double(args.at(6)); + const int output_zero_point = graph.get_int(args.at(7)); + + const auto [stride_h, stride_w] = + pair_or_throw(graph.get_int_list(args.at(10)), "q8ta_conv2d_dw: stride"); + const auto [pad_h, pad_w] = + pair_or_throw(graph.get_int_list(args.at(11)), "q8ta_conv2d_dw: padding"); + const auto [dil_h, dil_w] = pair_or_throw( + graph.get_int_list(args.at(12)), "q8ta_conv2d_dw: dilation"); + const int groups = graph.get_int(args.at(13)); + + const uint64_t N = static_cast(in_tensor.dims.at(0)); + const uint64_t C = static_cast(in_tensor.dims.at(1)); + const uint64_t H_in = static_cast(in_tensor.dims.at(2)); + const uint64_t W_in = static_cast(in_tensor.dims.at(3)); + const uint64_t Kh = static_cast(weight_tensor.dims.at(0)); + const uint64_t Kw = static_cast(weight_tensor.dims.at(1)); + const uint64_t OC = static_cast(weight_tensor.dims.at(2)); + const uint64_t H_out = static_cast(out_tensor.dims.at(2)); + const uint64_t W_out = static_cast(out_tensor.dims.at(3)); + + if (OC != C || static_cast(groups) != C) { + throw std::runtime_error( + "q8ta_conv2d_dw: depthwise requires groups==C==OC"); + } + if (static_cast(out_tensor.dims.at(1)) != C || + static_cast(out_tensor.dims.at(0)) != N) { + throw std::runtime_error( + "q8ta_conv2d_dw: output must be [N, C, H_out, W_out]"); + } + const uint64_t out_numel = N * C * H_out * W_out; + if (out_numel == 0 || W_out % 4 != 0) { + throw std::runtime_error( + "q8ta_conv2d_dw: W_out must be a nonzero multiple of 4"); + } + if (out_numel > UINT32_MAX || N * C * H_in * W_in > UINT32_MAX) { + throw std::runtime_error("q8ta_conv2d_dw: numel exceeds u32"); + } + const uint64_t in_numel = N * C * H_in * W_in; + const uint64_t weight_numel = Kh * Kw * OC; + if (!in_tensor.is_int8 || in_tensor.nbytes != in_numel || in_numel % 4 != 0 || + !weight_tensor.is_int8 || weight_tensor.nbytes != weight_numel || + weight_numel % 4 != 0 || !out_tensor.is_int8 || + out_tensor.nbytes != out_numel) { + throw std::runtime_error( + "q8ta_conv2d_dw: int8 in/weight/out size mismatch"); + } + if (scales_tensor.nbytes != OC * sizeof(float)) { + throw std::runtime_error("q8ta_conv2d_dw: weight_scales must be fp32 [OC]"); + } + if (has_bias) { + const auto& b = graph.get_tensor(bias_id); + if (b.buffer == nullptr || b.nbytes != OC * sizeof(float)) { + throw std::runtime_error("q8ta_conv2d_dw: bias must be fp32 [OC]"); + } + } + + Q8taConvDwParams params = {}; + params.N = static_cast(N); + params.C = static_cast(C); + params.H_in = static_cast(H_in); + params.W_in = static_cast(W_in); + params.H_out = static_cast(H_out); + params.W_out = static_cast(W_out); + params.Kh = static_cast(Kh); + params.Kw = static_cast(Kw); + params.stride_h = static_cast(stride_h); + params.stride_w = static_cast(stride_w); + params.pad_h = static_cast(pad_h); + params.pad_w = static_cast(pad_w); + params.dil_h = static_cast(dil_h); + params.dil_w = static_cast(dil_w); + params.input_zero_point = static_cast(input_zero_point); + params.output_zero_point = static_cast(output_zero_point); + params.input_scale = static_cast(input_scale); + // Reciprocal in double then cast, matching torch's f32(1.0 / f64(scale)). + params.inv_output_scale = static_cast(1.0 / output_scale); + params.has_bias = has_bias ? 1u : 0u; + + const uint32_t num_words = static_cast(out_numel / 4); + uint32_t wg_size = + utils::clamp_workgroup_size(device, kQ8taConv2dDwWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, num_words, wg_size, "q8ta_conv2d_dw"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(Q8taConvDwParams)); + graph.add_uniform_buffer_bytes(sizeof(Q8taConvDwParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kQ8taConv2dDwWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[6] = {}; + for (int i = 0; i < 6; i++) { + entries[i].binding = static_cast(i); + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = (i == 0) ? WGPUBufferBindingType_Storage + : (i == 5) ? WGPUBufferBindingType_Uniform + : WGPUBufferBindingType_ReadOnlyStorage; + } + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 6; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + // No-bias: bind scales as an unread placeholder (has_bias gates the read). + WGPUBuffer bias_buf = + has_bias ? graph.get_tensor(bias_id).buffer : scales_tensor.buffer; + const uint64_t bias_size = + has_bias ? graph.get_tensor(bias_id).nbytes : scales_tensor.nbytes; + + WGPUBindGroupEntry bg[6] = {}; + bg[0].binding = 0; + bg[0].buffer = out_tensor.buffer; + bg[0].size = out_tensor.nbytes; + bg[1].binding = 1; + bg[1].buffer = in_tensor.buffer; + bg[1].size = in_tensor.nbytes; + bg[2].binding = 2; + bg[2].buffer = weight_tensor.buffer; + bg[2].size = weight_tensor.nbytes; + bg[3].binding = 3; + bg[3].buffer = scales_tensor.buffer; + bg[3].size = scales_tensor.nbytes; + bg[4].binding = 4; + bg[4].buffer = bias_buf; + bg[4].size = bias_size; + bg[5].binding = 5; + bg[5].buffer = params_buf; + bg[5].size = sizeof(Q8taConvDwParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 6; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "q8ta_conv2d_dw", + workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.q8ta_conv2d_dw.default, q8ta_conv2d_dw_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d_dw/q8ta_conv2d_dw.wgsl b/backends/webgpu/runtime/ops/q8ta_conv2d_dw/q8ta_conv2d_dw.wgsl new file mode 100644 index 00000000000..4d44dd7cdd9 --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_conv2d_dw/q8ta_conv2d_dw.wgsl @@ -0,0 +1,94 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + N: u32, + C: u32, + H_in: u32, + W_in: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + input_zero_point: i32, + output_zero_point: i32, + input_scale: f32, + inv_output_scale: f32, + has_bias: u32, + pad0: u32, +} +@group(0) @binding(5) var params: Params; + +override wg_size: u32 = 64u; + +fn unpack_i8(bi: u32, word: u32) -> i32 { + return i32(((word >> ((bi & 3u) * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per output word = 4 W-positions of fixed (n,c,oh); W_out%4==0. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + let words = (params.N * params.C * params.H_out * params.W_out) / 4u; + if (widx >= words) { + return; + } + let flat0 = widx * 4u; + let ow0 = flat0 % params.W_out; + var r = flat0 / params.W_out; + let oh = r % params.H_out; + r = r / params.H_out; + let c = r % params.C; + let n = r / params.C; + + var acc: array; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + acc[j] = 0; + } + // Depthwise per-channel [Kh,Kw] window; weight laid out [Kh,Kw,OC]. + for (var kh: u32 = 0u; kh < params.Kh; kh = kh + 1u) { + let ih = i32(oh) * i32(params.stride_h) - i32(params.pad_h) + + i32(kh) * i32(params.dil_h); + if (ih < 0 || ih >= i32(params.H_in)) { + continue; + } + for (var kw: u32 = 0u; kw < params.Kw; kw = kw + 1u) { + let wbi = (kh * params.Kw + kw) * params.C + c; + let wv = unpack_i8(wbi, t_weight[wbi >> 2u]); + let in_row = (n * params.C + c) * params.H_in + u32(ih); + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let iw = i32(ow0 + j) * i32(params.stride_w) - i32(params.pad_w) + + i32(kw) * i32(params.dil_w); + if (iw < 0 || iw >= i32(params.W_in)) { + continue; + } + let xbi = in_row * params.W_in + u32(iw); + acc[j] = acc[j] + + (unpack_i8(xbi, t_x[xbi >> 2u]) - params.input_zero_point) * wv; + } + } + } + + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + var v = f32(acc[j]) * params.input_scale * t_scales[c]; + if (params.has_bias != 0u) { + v = v + t_bias[c]; + } + var q = i32(round(v * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (j * 8u)); + } + t_out[widx] = packed; +} diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d_dw/q8ta_conv2d_dw_wgsl.h b/backends/webgpu/runtime/ops/q8ta_conv2d_dw/q8ta_conv2d_dw_wgsl.h new file mode 100644 index 00000000000..ba7779e0681 --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_conv2d_dw/q8ta_conv2d_dw_wgsl.h @@ -0,0 +1,118 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from q8ta_conv2d_dw.wgsl - DO NOT EDIT. +// wgsl-sha256: daabc65120a540a43e0199d8c0a61cf510db056854013693cea39869c78359f3 +inline constexpr const char* kQ8taConv2dDwWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + N: u32, + C: u32, + H_in: u32, + W_in: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + input_zero_point: i32, + output_zero_point: i32, + input_scale: f32, + inv_output_scale: f32, + has_bias: u32, + pad0: u32, +} +@group(0) @binding(5) var params: Params; + +override wg_size: u32 = 64u; + +fn unpack_i8(bi: u32, word: u32) -> i32 { + return i32(((word >> ((bi & 3u) * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per output word = 4 W-positions of fixed (n,c,oh); W_out%4==0. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + let words = (params.N * params.C * params.H_out * params.W_out) / 4u; + if (widx >= words) { + return; + } + let flat0 = widx * 4u; + let ow0 = flat0 % params.W_out; + var r = flat0 / params.W_out; + let oh = r % params.H_out; + r = r / params.H_out; + let c = r % params.C; + let n = r / params.C; + + var acc: array; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + acc[j] = 0; + } + // Depthwise per-channel [Kh,Kw] window; weight laid out [Kh,Kw,OC]. + for (var kh: u32 = 0u; kh < params.Kh; kh = kh + 1u) { + let ih = i32(oh) * i32(params.stride_h) - i32(params.pad_h) + + i32(kh) * i32(params.dil_h); + if (ih < 0 || ih >= i32(params.H_in)) { + continue; + } + for (var kw: u32 = 0u; kw < params.Kw; kw = kw + 1u) { + let wbi = (kh * params.Kw + kw) * params.C + c; + let wv = unpack_i8(wbi, t_weight[wbi >> 2u]); + let in_row = (n * params.C + c) * params.H_in + u32(ih); + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let iw = i32(ow0 + j) * i32(params.stride_w) - i32(params.pad_w) + + i32(kw) * i32(params.dil_w); + if (iw < 0 || iw >= i32(params.W_in)) { + continue; + } + let xbi = in_row * params.W_in + u32(iw); + acc[j] = acc[j] + + (unpack_i8(xbi, t_x[xbi >> 2u]) - params.input_zero_point) * wv; + } + } + } + + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + var v = f32(acc[j]) * params.input_scale * t_scales[c]; + if (params.has_bias != 0u) { + v = v + t_bias[c]; + } + var q = i32(round(v * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (j * 8u)); + } + t_out[widx] = packed; +} +)"; + +inline constexpr uint32_t kQ8taConv2dDwWorkgroupSizeX = 64; +inline constexpr uint32_t kQ8taConv2dDwWorkgroupSizeY = 1; +inline constexpr uint32_t kQ8taConv2dDwWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 1892701f96b30d67a42e1cdbf7d10dc3ab63e2b5 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:31 -0700 Subject: [PATCH 44/75] [ExecuTorch][WebGPU] Op-tests for q8ta_conv2d_dw MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 29 +++++++++ .../webgpu/test/ops/test_q8ta_conv2d_dw.py | 63 +++++++++++++++++++ 2 files changed, 92 insertions(+) create mode 100644 backends/webgpu/test/ops/test_q8ta_conv2d_dw.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 25d1574e1a4..ff32ca8d870 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -60,6 +60,9 @@ from executorch.backends.webgpu.test.ops.test_q8ta_conv2d_pw import ( make_q8ta_conv2d_pw_module, ) +from executorch.backends.webgpu.test.ops.test_q8ta_conv2d_dw import ( + make_q8ta_conv2d_dw_module, +) from executorch.backends.webgpu.test.ops.test_pixel_shuffle import PixelShuffleModule from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule @@ -1542,6 +1545,32 @@ def case(name, ic, oc, h, w, n=1, **kw): ) +@register_op_test("q8ta_conv2d_dw") +def _q8ta_conv2d_dw_suite() -> WebGPUTestSuite: + # Golden = XNNPACK-static-PT2E converted eager (fp32); depthwise, C%4==0, + # W_out%4==0 for output packing. + def case(name, c, k, h, w, n=1, **kw): + return Case( + name=name, + construct={"c": c, "k": k, "h": h, "w": w, "n": n, **kw}, + inputs=((n, c, h, w),), + ) + + return WebGPUTestSuite( + module_factory=lambda **kw: make_q8ta_conv2d_dw_module(**kw), + cases=[ + case("k3", 8, 3, 8, 8), + case("no_bias", 8, 3, 8, 8, bias=False), + case("stride2", 8, 3, 8, 8, stride=2), + case("dil2", 8, 3, 8, 8, dilation=2, padding=2), + case("batch2", 8, 3, 8, 8, n=2), + ], + golden_dtype="float32", + atol=1e-3, + rtol=1e-3, + ) + + @register_op_test("dequantize_per_tensor") def _dequant_const_suite() -> WebGPUTestSuite: # Independent GPU dequantize check: baked int8 const vs torch dequant (breaks diff --git a/backends/webgpu/test/ops/test_q8ta_conv2d_dw.py b/backends/webgpu/test/ops/test_q8ta_conv2d_dw.py new file mode 100644 index 00000000000..f4f2f7d78ab --- /dev/null +++ b/backends/webgpu/test/ops/test_q8ta_conv2d_dw.py @@ -0,0 +1,63 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`et_vk.q8ta_conv2d_dw` module + configs: int8 depthwise conv. + +A depthwise `nn.Conv2d` (groups == channels) through XNNPACK static PT2E lowers +to a delegated `quantize_per_tensor -> q8ta_conv2d_dw -> dequantize_per_tensor` +subgraph (quantize/dequantize = the landed C0 ops). The `module_factory` returns +the CONVERTED module, so the op-test framework goldens the WebGPU output against +the converted eager (fp32 fake-quant), e2e. +""" + +import unittest + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 registers et_vk ops + +import torch +import torch.nn as nn + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( + get_symmetric_quantization_config, + XNNPACKQuantizer, +) +from executorch.exir import to_edge_transform_and_lower +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + +def make_q8ta_conv2d_dw_module( + c, k, h, w, stride=1, padding=1, dilation=1, bias=True, n=1, seed=0 +): + torch.manual_seed(seed) + conv = nn.Conv2d( + c, c, k, stride=stride, padding=padding, dilation=dilation, groups=c, bias=bias + ) + conv = conv.eval() + ex = (torch.randn(n, c, h, w),) + q = XNNPACKQuantizer().set_global( + get_symmetric_quantization_config(is_per_channel=True, is_dynamic=False) + ) + prepared = prepare_pt2e(torch.export.export(conv, ex).module(), q) + prepared(*ex) # calibrate + return convert_pt2e(prepared) + + +class Q8taConv2dDwTest(unittest.TestCase): + def test_delegates(self) -> None: + m = make_q8ta_conv2d_dw_module(8, 3, 8, 8) + et = to_edge_transform_and_lower( + torch.export.export(m, (torch.randn(1, 8, 8, 8),)), + partitioner=[VulkanPartitioner()], + ).to_executorch() + self.assertTrue( + any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + ) + self.assertIn(b"q8ta_conv2d_dw", et.buffer) From fce43d01be976641d2f2057a643fa88bae59d663 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:31 -0700 Subject: [PATCH 45/75] [ExecuTorch][WebGPU] Add q8ta_conv2d op (int8 general conv) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../runtime/ops/q8ta_conv2d/Q8taConv2d.cpp | 290 ++++++++++++++++++ .../runtime/ops/q8ta_conv2d/q8ta_conv2d.wgsl | 104 +++++++ .../ops/q8ta_conv2d/q8ta_conv2d_wgsl.h | 128 ++++++++ 3 files changed, 522 insertions(+) create mode 100644 backends/webgpu/runtime/ops/q8ta_conv2d/Q8taConv2d.cpp create mode 100644 backends/webgpu/runtime/ops/q8ta_conv2d/q8ta_conv2d.wgsl create mode 100644 backends/webgpu/runtime/ops/q8ta_conv2d/q8ta_conv2d_wgsl.h diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d/Q8taConv2d.cpp b/backends/webgpu/runtime/ops/q8ta_conv2d/Q8taConv2d.cpp new file mode 100644 index 00000000000..a6d88e159ef --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_conv2d/Q8taConv2d.cpp @@ -0,0 +1,290 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct Q8taConvParams { + uint32_t N; + uint32_t IC; + uint32_t H_in; + uint32_t W_in; + uint32_t OC; + uint32_t H_out; + uint32_t W_out; + uint32_t Kh; + uint32_t Kw; + uint32_t stride_h; + uint32_t stride_w; + uint32_t pad_h; + uint32_t pad_w; + uint32_t dil_h; + uint32_t dil_w; + uint32_t weight_row_stride; + int32_t input_zero_point; + int32_t output_zero_point; + float input_scale; + float inv_output_scale; + uint32_t has_bias; + uint32_t pad0; + uint32_t pad1; + uint32_t pad2; +}; +static_assert( + sizeof(Q8taConvParams) == 96, + "Q8taConvParams must match the WGSL Params struct (96 bytes)"); + +std::pair pair_or_throw( + const std::vector& v, + const char* msg) { + if (v.size() != 2) { + throw std::runtime_error(msg); + } + return {v.at(0), v.at(1)}; +} + +// int8 general conv (groups==1); full-IC windowed dot; mirrors Vulkan. +void q8ta_conv2d_impl(WebGPUGraph& graph, const std::vector& args) { + const int in_id = args.at(0); + const int weight_id = args.at(3); + const int scales_id = args.at(5); + const int bias_id = args.at(8); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(weight_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(scales_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("q8ta_conv2d: in/weight/scales/out not tensor"); + } + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + + const int act_id = args.at(args.size() - 2); + if (graph.get_value_type(act_id) != WebGPUGraph::ValueType::String || + graph.get_string(act_id) != "none") { + throw std::runtime_error("q8ta_conv2d: only activation='none' supported"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& weight_tensor = graph.get_tensor(weight_id); + const auto& scales_tensor = graph.get_tensor(scales_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || weight_tensor.buffer == nullptr || + scales_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("q8ta_conv2d: null buffer binding"); + } + if (in_tensor.dims.size() != 4 || out_tensor.dims.size() != 4 || + weight_tensor.dims.size() != 2) { + throw std::runtime_error("q8ta_conv2d: in/out must be 4D, weight 2D"); + } + + const double input_scale = graph.get_double(args.at(1)); + const int input_zero_point = graph.get_int(args.at(2)); + const double output_scale = graph.get_double(args.at(6)); + const int output_zero_point = graph.get_int(args.at(7)); + + const auto [kernel_h, kernel_w] = + pair_or_throw(graph.get_int_list(args.at(9)), "q8ta_conv2d: kernel_size"); + const auto [stride_h, stride_w] = + pair_or_throw(graph.get_int_list(args.at(10)), "q8ta_conv2d: stride"); + const auto [pad_h, pad_w] = + pair_or_throw(graph.get_int_list(args.at(11)), "q8ta_conv2d: padding"); + const auto [dil_h, dil_w] = + pair_or_throw(graph.get_int_list(args.at(12)), "q8ta_conv2d: dilation"); + const int groups = graph.get_int(args.at(13)); + if (groups != 1) { + throw std::runtime_error("q8ta_conv2d: only groups==1 supported"); + } + + const uint64_t N = static_cast(in_tensor.dims.at(0)); + const uint64_t IC = static_cast(in_tensor.dims.at(1)); + const uint64_t H_in = static_cast(in_tensor.dims.at(2)); + const uint64_t W_in = static_cast(in_tensor.dims.at(3)); + const uint64_t OC = static_cast(weight_tensor.dims.at(0)); + const uint64_t weight_row_stride = + static_cast(weight_tensor.dims.at(1)); + const uint64_t Kh = static_cast(kernel_h); + const uint64_t Kw = static_cast(kernel_w); + const uint64_t H_out = static_cast(out_tensor.dims.at(2)); + const uint64_t W_out = static_cast(out_tensor.dims.at(3)); + + if (static_cast(out_tensor.dims.at(0)) != N || + static_cast(out_tensor.dims.at(1)) != OC) { + throw std::runtime_error( + "q8ta_conv2d: output must be [N, OC, H_out, W_out]"); + } + if (weight_row_stride < Kh * Kw * IC) { + throw std::runtime_error("q8ta_conv2d: weight row stride < Kh*Kw*IC"); + } + const uint64_t out_numel = N * OC * H_out * W_out; + if (out_numel == 0 || W_out % 4 != 0) { + throw std::runtime_error( + "q8ta_conv2d: W_out must be a nonzero multiple of 4"); + } + const uint64_t in_numel = N * IC * H_in * W_in; + const uint64_t weight_numel = OC * weight_row_stride; + if (out_numel > UINT32_MAX || in_numel > UINT32_MAX || + weight_numel > UINT32_MAX) { + throw std::runtime_error("q8ta_conv2d: numel exceeds u32"); + } + if (!in_tensor.is_int8 || in_tensor.nbytes != in_numel || in_numel % 4 != 0 || + !weight_tensor.is_int8 || weight_tensor.nbytes != weight_numel || + weight_numel % 4 != 0 || !out_tensor.is_int8 || + out_tensor.nbytes != out_numel) { + throw std::runtime_error("q8ta_conv2d: int8 in/weight/out size mismatch"); + } + // scales/bias fp32 [OC]; AOT pads OC to mult-4; shader reads [0,OC) only. + if (scales_tensor.nbytes < OC * sizeof(float)) { + throw std::runtime_error("q8ta_conv2d: weight_scales must be fp32 [OC]"); + } + if (has_bias) { + const auto& b = graph.get_tensor(bias_id); + if (b.buffer == nullptr || b.nbytes < OC * sizeof(float)) { + throw std::runtime_error("q8ta_conv2d: bias must be fp32 [OC]"); + } + } + + Q8taConvParams params = {}; + params.N = static_cast(N); + params.IC = static_cast(IC); + params.H_in = static_cast(H_in); + params.W_in = static_cast(W_in); + params.OC = static_cast(OC); + params.H_out = static_cast(H_out); + params.W_out = static_cast(W_out); + params.Kh = static_cast(Kh); + params.Kw = static_cast(Kw); + params.stride_h = static_cast(stride_h); + params.stride_w = static_cast(stride_w); + params.pad_h = static_cast(pad_h); + params.pad_w = static_cast(pad_w); + params.dil_h = static_cast(dil_h); + params.dil_w = static_cast(dil_w); + params.weight_row_stride = static_cast(weight_row_stride); + params.input_zero_point = static_cast(input_zero_point); + params.output_zero_point = static_cast(output_zero_point); + params.input_scale = static_cast(input_scale); + // Reciprocal in double then cast, matching torch's f32(1.0 / f64(scale)). + params.inv_output_scale = static_cast(1.0 / output_scale); + params.has_bias = has_bias ? 1u : 0u; + + const uint32_t num_words = static_cast(out_numel / 4); + uint32_t wg_size = + utils::clamp_workgroup_size(device, kQ8taConv2dWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, num_words, wg_size, "q8ta_conv2d"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(Q8taConvParams)); + graph.add_uniform_buffer_bytes(sizeof(Q8taConvParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kQ8taConv2dWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[6] = {}; + for (int i = 0; i < 6; i++) { + entries[i].binding = static_cast(i); + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = (i == 0) ? WGPUBufferBindingType_Storage + : (i == 5) ? WGPUBufferBindingType_Uniform + : WGPUBufferBindingType_ReadOnlyStorage; + } + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 6; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + // No-bias: bind scales as an unread placeholder (has_bias gates the read). + WGPUBuffer bias_buf = + has_bias ? graph.get_tensor(bias_id).buffer : scales_tensor.buffer; + const uint64_t bias_size = + has_bias ? graph.get_tensor(bias_id).nbytes : scales_tensor.nbytes; + + WGPUBindGroupEntry bg[6] = {}; + bg[0].binding = 0; + bg[0].buffer = out_tensor.buffer; + bg[0].size = out_tensor.nbytes; + bg[1].binding = 1; + bg[1].buffer = in_tensor.buffer; + bg[1].size = in_tensor.nbytes; + bg[2].binding = 2; + bg[2].buffer = weight_tensor.buffer; + bg[2].size = weight_tensor.nbytes; + bg[3].binding = 3; + bg[3].buffer = scales_tensor.buffer; + bg[3].size = scales_tensor.nbytes; + bg[4].binding = 4; + bg[4].buffer = bias_buf; + bg[4].size = bias_size; + bg[5].binding = 5; + bg[5].buffer = params_buf; + bg[5].size = sizeof(Q8taConvParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 6; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "q8ta_conv2d", + workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.q8ta_conv2d.default, q8ta_conv2d_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d/q8ta_conv2d.wgsl b/backends/webgpu/runtime/ops/q8ta_conv2d/q8ta_conv2d.wgsl new file mode 100644 index 00000000000..b40660d80aa --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_conv2d/q8ta_conv2d.wgsl @@ -0,0 +1,104 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + weight_row_stride: u32, + input_zero_point: i32, + output_zero_point: i32, + input_scale: f32, + inv_output_scale: f32, + has_bias: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(5) var params: Params; + +override wg_size: u32 = 64u; + +fn unpack_i8(bi: u32, word: u32) -> i32 { + return i32(((word >> ((bi & 3u) * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per output word = 4 W-positions of fixed (n,oc,oh); W_out%4==0. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + let words = (params.N * params.OC * params.H_out * params.W_out) / 4u; + if (widx >= words) { + return; + } + let flat0 = widx * 4u; + let ow0 = flat0 % params.W_out; + var r = flat0 / params.W_out; + let oh = r % params.H_out; + r = r / params.H_out; + let oc = r % params.OC; + let n = r / params.OC; + + // acc is i32: each tap |(x - zp) * w| <= 255 * 128 = 32640, so the sum + // stays exact while Kh*Kw*IC <= 65793 (2^31 / 32640) and can overflow + // beyond that. + var acc: array; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + acc[j] = 0; + } + // General conv: full IC reduction; weight [OC,Kh*Kw*IC]. + let w_row = oc * params.weight_row_stride; + for (var ic: u32 = 0u; ic < params.IC; ic = ic + 1u) { + for (var kh: u32 = 0u; kh < params.Kh; kh = kh + 1u) { + let ih = i32(oh) * i32(params.stride_h) - i32(params.pad_h) + + i32(kh) * i32(params.dil_h); + if (ih < 0 || ih >= i32(params.H_in)) { + continue; + } + let in_row = ((n * params.IC + ic) * params.H_in + u32(ih)) * params.W_in; + for (var kw: u32 = 0u; kw < params.Kw; kw = kw + 1u) { + let wbi = w_row + (kh * params.Kw + kw) * params.IC + ic; + let wv = unpack_i8(wbi, t_weight[wbi >> 2u]); + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let iw = i32(ow0 + j) * i32(params.stride_w) - i32(params.pad_w) + + i32(kw) * i32(params.dil_w); + if (iw < 0 || iw >= i32(params.W_in)) { + continue; + } + let xbi = in_row + u32(iw); + acc[j] = acc[j] + + (unpack_i8(xbi, t_x[xbi >> 2u]) - params.input_zero_point) * wv; + } + } + } + } + + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + var v = f32(acc[j]) * params.input_scale * t_scales[oc]; + if (params.has_bias != 0u) { + v = v + t_bias[oc]; + } + var q = i32(round(v * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (j * 8u)); + } + t_out[widx] = packed; +} diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d/q8ta_conv2d_wgsl.h b/backends/webgpu/runtime/ops/q8ta_conv2d/q8ta_conv2d_wgsl.h new file mode 100644 index 00000000000..bdb685fe608 --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_conv2d/q8ta_conv2d_wgsl.h @@ -0,0 +1,128 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from q8ta_conv2d.wgsl - DO NOT EDIT. +// wgsl-sha256: f93c4400b1e84039aa29ddbc6470cc0bea2a3e73370ce64700a2d63a42a462e7 +inline constexpr const char* kQ8taConv2dWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + weight_row_stride: u32, + input_zero_point: i32, + output_zero_point: i32, + input_scale: f32, + inv_output_scale: f32, + has_bias: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(5) var params: Params; + +override wg_size: u32 = 64u; + +fn unpack_i8(bi: u32, word: u32) -> i32 { + return i32(((word >> ((bi & 3u) * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per output word = 4 W-positions of fixed (n,oc,oh); W_out%4==0. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + let words = (params.N * params.OC * params.H_out * params.W_out) / 4u; + if (widx >= words) { + return; + } + let flat0 = widx * 4u; + let ow0 = flat0 % params.W_out; + var r = flat0 / params.W_out; + let oh = r % params.H_out; + r = r / params.H_out; + let oc = r % params.OC; + let n = r / params.OC; + + // acc is i32: each tap |(x - zp) * w| <= 255 * 128 = 32640, so the sum + // stays exact while Kh*Kw*IC <= 65793 (2^31 / 32640) and can overflow + // beyond that. + var acc: array; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + acc[j] = 0; + } + // General conv: full IC reduction; weight [OC,Kh*Kw*IC]. + let w_row = oc * params.weight_row_stride; + for (var ic: u32 = 0u; ic < params.IC; ic = ic + 1u) { + for (var kh: u32 = 0u; kh < params.Kh; kh = kh + 1u) { + let ih = i32(oh) * i32(params.stride_h) - i32(params.pad_h) + + i32(kh) * i32(params.dil_h); + if (ih < 0 || ih >= i32(params.H_in)) { + continue; + } + let in_row = ((n * params.IC + ic) * params.H_in + u32(ih)) * params.W_in; + for (var kw: u32 = 0u; kw < params.Kw; kw = kw + 1u) { + let wbi = w_row + (kh * params.Kw + kw) * params.IC + ic; + let wv = unpack_i8(wbi, t_weight[wbi >> 2u]); + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let iw = i32(ow0 + j) * i32(params.stride_w) - i32(params.pad_w) + + i32(kw) * i32(params.dil_w); + if (iw < 0 || iw >= i32(params.W_in)) { + continue; + } + let xbi = in_row + u32(iw); + acc[j] = acc[j] + + (unpack_i8(xbi, t_x[xbi >> 2u]) - params.input_zero_point) * wv; + } + } + } + } + + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + var v = f32(acc[j]) * params.input_scale * t_scales[oc]; + if (params.has_bias != 0u) { + v = v + t_bias[oc]; + } + var q = i32(round(v * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (j * 8u)); + } + t_out[widx] = packed; +} +)"; + +inline constexpr uint32_t kQ8taConv2dWorkgroupSizeX = 64; +inline constexpr uint32_t kQ8taConv2dWorkgroupSizeY = 1; +inline constexpr uint32_t kQ8taConv2dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 9440253bd5d617004b95840a93f847ac7100da4c Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:32 -0700 Subject: [PATCH 46/75] [ExecuTorch][WebGPU] Op-tests for q8ta_conv2d MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 33 +++++++++++ backends/webgpu/test/ops/test_q8ta_conv2d.py | 62 ++++++++++++++++++++ 2 files changed, 95 insertions(+) create mode 100644 backends/webgpu/test/ops/test_q8ta_conv2d.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index ff32ca8d870..1f7e1c4d5e2 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -63,6 +63,9 @@ from executorch.backends.webgpu.test.ops.test_q8ta_conv2d_dw import ( make_q8ta_conv2d_dw_module, ) +from executorch.backends.webgpu.test.ops.test_q8ta_conv2d import ( + make_q8ta_conv2d_module, +) from executorch.backends.webgpu.test.ops.test_pixel_shuffle import PixelShuffleModule from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule @@ -1571,6 +1574,36 @@ def case(name, c, k, h, w, n=1, **kw): ) +@register_op_test("q8ta_conv2d") +def _q8ta_conv2d_suite() -> WebGPUTestSuite: + # Golden = XNNPACK-static-PT2E converted eager (fp32); general (groups==1) + # conv, full IC reduction. W_out%4==0 for output packing. + def case(name, ic, oc, k, h, w, n=1, **kw): + return Case( + name=name, + construct={"ic": ic, "oc": oc, "k": k, "h": h, "w": w, "n": n, **kw}, + inputs=((n, ic, h, w),), + ) + + return WebGPUTestSuite( + module_factory=lambda **kw: make_q8ta_conv2d_module(**kw), + cases=[ + case("k3", 8, 8, 3, 8, 8), + case("oc_gt", 4, 8, 3, 8, 8), + case("no_bias", 8, 8, 3, 8, 8, bias=False), + case("stride2", 8, 8, 3, 8, 8, stride=2), + case("dil2", 8, 8, 3, 8, 8, dilation=2, padding=2), + case("oc6", 8, 6, 3, 8, 8), + case("ic3", 3, 8, 3, 8, 8), + case("asym", 8, 8, (3, 5), 8, 8, padding=(1, 2)), + case("batch2", 8, 8, 3, 8, 8, n=2), + ], + golden_dtype="float32", + atol=1e-3, + rtol=1e-3, + ) + + @register_op_test("dequantize_per_tensor") def _dequant_const_suite() -> WebGPUTestSuite: # Independent GPU dequantize check: baked int8 const vs torch dequant (breaks diff --git a/backends/webgpu/test/ops/test_q8ta_conv2d.py b/backends/webgpu/test/ops/test_q8ta_conv2d.py new file mode 100644 index 00000000000..5490f27d82a --- /dev/null +++ b/backends/webgpu/test/ops/test_q8ta_conv2d.py @@ -0,0 +1,62 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`et_vk.q8ta_conv2d` module + configs: int8 general (ungrouped) conv. + +A standard `nn.Conv2d` (groups == 1) through XNNPACK static PT2E lowers to a +delegated `quantize_per_tensor -> q8ta_conv2d -> dequantize_per_tensor` subgraph +(quantize/dequantize = the landed C0 ops). The `module_factory` returns the +CONVERTED module, so the op-test framework goldens the WebGPU output against the +converted eager (fp32 fake-quant), e2e. +""" + +import unittest + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 registers et_vk ops + +import torch +import torch.nn as nn + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( + get_symmetric_quantization_config, + XNNPACKQuantizer, +) +from executorch.exir import to_edge_transform_and_lower +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + +def make_q8ta_conv2d_module( + ic, oc, k, h, w, stride=1, padding=1, dilation=1, bias=True, n=1, seed=0 +): + torch.manual_seed(seed) + conv = nn.Conv2d( + ic, oc, k, stride=stride, padding=padding, dilation=dilation, bias=bias + ).eval() + ex = (torch.randn(n, ic, h, w),) + q = XNNPACKQuantizer().set_global( + get_symmetric_quantization_config(is_per_channel=True, is_dynamic=False) + ) + prepared = prepare_pt2e(torch.export.export(conv, ex).module(), q) + prepared(*ex) # calibrate + return convert_pt2e(prepared) + + +class Q8taConv2dTest(unittest.TestCase): + def test_delegates(self) -> None: + m = make_q8ta_conv2d_module(8, 8, 3, 8, 8) + et = to_edge_transform_and_lower( + torch.export.export(m, (torch.randn(1, 8, 8, 8),)), + partitioner=[VulkanPartitioner()], + ).to_executorch() + self.assertTrue( + any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + ) + self.assertIn(b"q8ta_conv2d", et.buffer) From b741ef3910b45e94cd6e5b39ad5cdfcc537bb7ab Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:32 -0700 Subject: [PATCH 47/75] [ExecuTorch][WebGPU] Add q8ta_conv2d_transposed op (int8 transposed conv) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../Q8taConv2dTransposed.cpp | 309 ++++++++++++++++++ .../q8ta_conv2d_transposed.wgsl | 108 ++++++ .../q8ta_conv2d_transposed_wgsl.h | 132 ++++++++ 3 files changed, 549 insertions(+) create mode 100644 backends/webgpu/runtime/ops/q8ta_conv2d_transposed/Q8taConv2dTransposed.cpp create mode 100644 backends/webgpu/runtime/ops/q8ta_conv2d_transposed/q8ta_conv2d_transposed.wgsl create mode 100644 backends/webgpu/runtime/ops/q8ta_conv2d_transposed/q8ta_conv2d_transposed_wgsl.h diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d_transposed/Q8taConv2dTransposed.cpp b/backends/webgpu/runtime/ops/q8ta_conv2d_transposed/Q8taConv2dTransposed.cpp new file mode 100644 index 00000000000..0990b49bb9d --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_conv2d_transposed/Q8taConv2dTransposed.cpp @@ -0,0 +1,309 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct Q8taConvTParams { + uint32_t N; + uint32_t IC; + uint32_t H_in; + uint32_t W_in; + uint32_t OC; + uint32_t H_out; + uint32_t W_out; + uint32_t Kh; + uint32_t Kw; + uint32_t stride_h; + uint32_t stride_w; + uint32_t pad_h; + uint32_t pad_w; + uint32_t dil_h; + uint32_t dil_w; + uint32_t weight_row_stride; + int32_t input_zero_point; + int32_t output_zero_point; + float input_scale; + float inv_output_scale; + uint32_t has_bias; + uint32_t pad0; + uint32_t pad1; + uint32_t pad2; +}; +static_assert( + sizeof(Q8taConvTParams) == 96, + "Q8taConvTParams must match the WGSL Params struct (96 bytes)"); + +std::pair pair_or_throw( + const std::vector& v, + const char* msg) { + if (v.size() != 2) { + throw std::runtime_error(msg); + } + return {v.at(0), v.at(1)}; +} + +// int8 transposed conv (groups==1, dilation==1); gather form; mirrors Vulkan. +void q8ta_conv2d_transposed_impl( + WebGPUGraph& graph, + const std::vector& args) { + const int in_id = args.at(0); + const int weight_id = args.at(3); + const int scales_id = args.at(5); + const int bias_id = args.at(8); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(weight_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(scales_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error( + "q8ta_conv2d_transposed: in/weight/scales/out not tensor"); + } + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + + const int act_id = args.at(args.size() - 2); + if (graph.get_value_type(act_id) != WebGPUGraph::ValueType::String || + graph.get_string(act_id) != "none") { + throw std::runtime_error( + "q8ta_conv2d_transposed: only activation='none' supported"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& weight_tensor = graph.get_tensor(weight_id); + const auto& scales_tensor = graph.get_tensor(scales_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || weight_tensor.buffer == nullptr || + scales_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("q8ta_conv2d_transposed: null buffer binding"); + } + if (in_tensor.dims.size() != 4 || out_tensor.dims.size() != 4 || + weight_tensor.dims.size() != 2) { + throw std::runtime_error( + "q8ta_conv2d_transposed: in/out must be 4D, weight 2D"); + } + + const double input_scale = graph.get_double(args.at(1)); + const int input_zero_point = graph.get_int(args.at(2)); + const double output_scale = graph.get_double(args.at(6)); + const int output_zero_point = graph.get_int(args.at(7)); + + // Transposed schema inserts output_padding at 12, shifting dilation to 13. + const auto [kernel_h, kernel_w] = pair_or_throw( + graph.get_int_list(args.at(9)), "q8ta_conv2d_transposed: kernel_size"); + const auto [stride_h, stride_w] = pair_or_throw( + graph.get_int_list(args.at(10)), "q8ta_conv2d_transposed: stride"); + const auto [pad_h, pad_w] = pair_or_throw( + graph.get_int_list(args.at(11)), "q8ta_conv2d_transposed: padding"); + const auto [dil_h, dil_w] = pair_or_throw( + graph.get_int_list(args.at(13)), "q8ta_conv2d_transposed: dilation"); + const int groups = graph.get_int(args.at(14)); + if (groups != 1) { + throw std::runtime_error( + "q8ta_conv2d_transposed: only groups==1 supported"); + } + if (dil_h != 1 || dil_w != 1) { + throw std::runtime_error( + "q8ta_conv2d_transposed: only dilation==1 supported"); + } + if (stride_h < 1 || stride_w < 1) { + throw std::runtime_error("q8ta_conv2d_transposed: stride must be >= 1"); + } + + const uint64_t N = static_cast(in_tensor.dims.at(0)); + const uint64_t IC = static_cast(in_tensor.dims.at(1)); + const uint64_t H_in = static_cast(in_tensor.dims.at(2)); + const uint64_t W_in = static_cast(in_tensor.dims.at(3)); + const uint64_t OC = static_cast(weight_tensor.dims.at(0)); + const uint64_t weight_row_stride = + static_cast(weight_tensor.dims.at(1)); + const uint64_t Kh = static_cast(kernel_h); + const uint64_t Kw = static_cast(kernel_w); + const uint64_t H_out = static_cast(out_tensor.dims.at(2)); + const uint64_t W_out = static_cast(out_tensor.dims.at(3)); + + if (static_cast(out_tensor.dims.at(0)) != N || + static_cast(out_tensor.dims.at(1)) != OC) { + throw std::runtime_error( + "q8ta_conv2d_transposed: output must be [N, OC, H_out, W_out]"); + } + if (weight_row_stride < Kh * Kw * IC) { + throw std::runtime_error( + "q8ta_conv2d_transposed: weight row stride < Kh*Kw*IC"); + } + const uint64_t out_numel = N * OC * H_out * W_out; + if (out_numel == 0 || W_out % 4 != 0) { + throw std::runtime_error( + "q8ta_conv2d_transposed: W_out must be a nonzero multiple of 4"); + } + const uint64_t in_numel = N * IC * H_in * W_in; + const uint64_t weight_numel = OC * weight_row_stride; + if (out_numel > UINT32_MAX || in_numel > UINT32_MAX || + weight_numel > UINT32_MAX) { + throw std::runtime_error("q8ta_conv2d_transposed: numel exceeds u32"); + } + if (!in_tensor.is_int8 || in_tensor.nbytes != in_numel || in_numel % 4 != 0 || + !weight_tensor.is_int8 || weight_tensor.nbytes != weight_numel || + weight_numel % 4 != 0 || !out_tensor.is_int8 || + out_tensor.nbytes != out_numel) { + throw std::runtime_error( + "q8ta_conv2d_transposed: int8 in/weight/out size mismatch"); + } + // scales/bias fp32 [OC]; AOT pads OC to mult-4; shader reads [0,OC) only. + if (scales_tensor.nbytes < OC * sizeof(float)) { + throw std::runtime_error( + "q8ta_conv2d_transposed: weight_scales must be fp32 [OC]"); + } + if (has_bias) { + const auto& b = graph.get_tensor(bias_id); + if (b.buffer == nullptr || b.nbytes < OC * sizeof(float)) { + throw std::runtime_error( + "q8ta_conv2d_transposed: bias must be fp32 [OC]"); + } + } + + Q8taConvTParams params = {}; + params.N = static_cast(N); + params.IC = static_cast(IC); + params.H_in = static_cast(H_in); + params.W_in = static_cast(W_in); + params.OC = static_cast(OC); + params.H_out = static_cast(H_out); + params.W_out = static_cast(W_out); + params.Kh = static_cast(Kh); + params.Kw = static_cast(Kw); + params.stride_h = static_cast(stride_h); + params.stride_w = static_cast(stride_w); + params.pad_h = static_cast(pad_h); + params.pad_w = static_cast(pad_w); + params.dil_h = static_cast(dil_h); + params.dil_w = static_cast(dil_w); + params.weight_row_stride = static_cast(weight_row_stride); + params.input_zero_point = static_cast(input_zero_point); + params.output_zero_point = static_cast(output_zero_point); + params.input_scale = static_cast(input_scale); + // Reciprocal in double then cast, matching torch's f32(1.0 / f64(scale)). + params.inv_output_scale = static_cast(1.0 / output_scale); + params.has_bias = has_bias ? 1u : 0u; + + const uint32_t num_words = static_cast(out_numel / 4); + uint32_t wg_size = + utils::clamp_workgroup_size(device, kQ8taConv2dTransposedWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, num_words, wg_size, "q8ta_conv2d_transposed"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(Q8taConvTParams)); + graph.add_uniform_buffer_bytes(sizeof(Q8taConvTParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kQ8taConv2dTransposedWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[6] = {}; + for (int i = 0; i < 6; i++) { + entries[i].binding = static_cast(i); + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = (i == 0) ? WGPUBufferBindingType_Storage + : (i == 5) ? WGPUBufferBindingType_Uniform + : WGPUBufferBindingType_ReadOnlyStorage; + } + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 6; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + // No-bias: bind scales as an unread placeholder (has_bias gates the read). + WGPUBuffer bias_buf = + has_bias ? graph.get_tensor(bias_id).buffer : scales_tensor.buffer; + const uint64_t bias_size = + has_bias ? graph.get_tensor(bias_id).nbytes : scales_tensor.nbytes; + + WGPUBindGroupEntry bg[6] = {}; + bg[0].binding = 0; + bg[0].buffer = out_tensor.buffer; + bg[0].size = out_tensor.nbytes; + bg[1].binding = 1; + bg[1].buffer = in_tensor.buffer; + bg[1].size = in_tensor.nbytes; + bg[2].binding = 2; + bg[2].buffer = weight_tensor.buffer; + bg[2].size = weight_tensor.nbytes; + bg[3].binding = 3; + bg[3].buffer = scales_tensor.buffer; + bg[3].size = scales_tensor.nbytes; + bg[4].binding = 4; + bg[4].buffer = bias_buf; + bg[4].size = bias_size; + bg[5].binding = 5; + bg[5].buffer = params_buf; + bg[5].size = sizeof(Q8taConvTParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 6; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "q8ta_conv2d_transposed", + workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP( + et_vk.q8ta_conv2d_transposed.default, q8ta_conv2d_transposed_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d_transposed/q8ta_conv2d_transposed.wgsl b/backends/webgpu/runtime/ops/q8ta_conv2d_transposed/q8ta_conv2d_transposed.wgsl new file mode 100644 index 00000000000..6c7a959441a --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_conv2d_transposed/q8ta_conv2d_transposed.wgsl @@ -0,0 +1,108 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + weight_row_stride: u32, + input_zero_point: i32, + output_zero_point: i32, + input_scale: f32, + inv_output_scale: f32, + has_bias: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(5) var params: Params; + +override wg_size: u32 = 64u; + +fn unpack_i8(bi: u32, word: u32) -> i32 { + return i32(((word >> ((bi & 3u) * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per output word = 4 W-positions of fixed (n,oc,oh); W_out%4==0. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + let words = (params.N * params.OC * params.H_out * params.W_out) / 4u; + if (widx >= words) { + return; + } + let flat0 = widx * 4u; + let ow0 = flat0 % params.W_out; + var r = flat0 / params.W_out; + let oh = r % params.H_out; + r = r / params.H_out; + let oc = r % params.OC; + let n = r / params.OC; + + var acc: array; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + acc[j] = 0; + } + // Transposed conv gather: ih=(oh+pad-kh*dil)/stride (must divide), full IC. + let w_row = oc * params.weight_row_stride; + for (var ic: u32 = 0u; ic < params.IC; ic = ic + 1u) { + for (var kh: u32 = 0u; kh < params.Kh; kh = kh + 1u) { + let ih_num = i32(oh) + i32(params.pad_h) - i32(kh) * i32(params.dil_h); + if (ih_num < 0 || (ih_num % i32(params.stride_h)) != 0) { + continue; + } + let ih = ih_num / i32(params.stride_h); + if (ih >= i32(params.H_in)) { + continue; + } + let in_row = ((n * params.IC + ic) * params.H_in + u32(ih)) * params.W_in; + for (var kw: u32 = 0u; kw < params.Kw; kw = kw + 1u) { + let wbi = w_row + (kh * params.Kw + kw) * params.IC + ic; + let wv = unpack_i8(wbi, t_weight[wbi >> 2u]); + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let iw_num = + i32(ow0 + j) + i32(params.pad_w) - i32(kw) * i32(params.dil_w); + if (iw_num < 0 || (iw_num % i32(params.stride_w)) != 0) { + continue; + } + let iw = iw_num / i32(params.stride_w); + if (iw >= i32(params.W_in)) { + continue; + } + let xbi = in_row + u32(iw); + acc[j] = acc[j] + + (unpack_i8(xbi, t_x[xbi >> 2u]) - params.input_zero_point) * wv; + } + } + } + } + + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + var v = f32(acc[j]) * params.input_scale * t_scales[oc]; + if (params.has_bias != 0u) { + v = v + t_bias[oc]; + } + var q = i32(round(v * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (j * 8u)); + } + t_out[widx] = packed; +} diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d_transposed/q8ta_conv2d_transposed_wgsl.h b/backends/webgpu/runtime/ops/q8ta_conv2d_transposed/q8ta_conv2d_transposed_wgsl.h new file mode 100644 index 00000000000..79d463f9ddd --- /dev/null +++ b/backends/webgpu/runtime/ops/q8ta_conv2d_transposed/q8ta_conv2d_transposed_wgsl.h @@ -0,0 +1,132 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from q8ta_conv2d_transposed.wgsl - DO NOT EDIT. +// wgsl-sha256: 426daa41bd5d7202d9260462c0ef4bb4c1737f0708f13331b22c4ba1f7885c1b +inline constexpr const char* kQ8taConv2dTransposedWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + weight_row_stride: u32, + input_zero_point: i32, + output_zero_point: i32, + input_scale: f32, + inv_output_scale: f32, + has_bias: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(5) var params: Params; + +override wg_size: u32 = 64u; + +fn unpack_i8(bi: u32, word: u32) -> i32 { + return i32(((word >> ((bi & 3u) * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per output word = 4 W-positions of fixed (n,oc,oh); W_out%4==0. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + let words = (params.N * params.OC * params.H_out * params.W_out) / 4u; + if (widx >= words) { + return; + } + let flat0 = widx * 4u; + let ow0 = flat0 % params.W_out; + var r = flat0 / params.W_out; + let oh = r % params.H_out; + r = r / params.H_out; + let oc = r % params.OC; + let n = r / params.OC; + + var acc: array; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + acc[j] = 0; + } + // Transposed conv gather: ih=(oh+pad-kh*dil)/stride (must divide), full IC. + let w_row = oc * params.weight_row_stride; + for (var ic: u32 = 0u; ic < params.IC; ic = ic + 1u) { + for (var kh: u32 = 0u; kh < params.Kh; kh = kh + 1u) { + let ih_num = i32(oh) + i32(params.pad_h) - i32(kh) * i32(params.dil_h); + if (ih_num < 0 || (ih_num % i32(params.stride_h)) != 0) { + continue; + } + let ih = ih_num / i32(params.stride_h); + if (ih >= i32(params.H_in)) { + continue; + } + let in_row = ((n * params.IC + ic) * params.H_in + u32(ih)) * params.W_in; + for (var kw: u32 = 0u; kw < params.Kw; kw = kw + 1u) { + let wbi = w_row + (kh * params.Kw + kw) * params.IC + ic; + let wv = unpack_i8(wbi, t_weight[wbi >> 2u]); + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let iw_num = + i32(ow0 + j) + i32(params.pad_w) - i32(kw) * i32(params.dil_w); + if (iw_num < 0 || (iw_num % i32(params.stride_w)) != 0) { + continue; + } + let iw = iw_num / i32(params.stride_w); + if (iw >= i32(params.W_in)) { + continue; + } + let xbi = in_row + u32(iw); + acc[j] = acc[j] + + (unpack_i8(xbi, t_x[xbi >> 2u]) - params.input_zero_point) * wv; + } + } + } + } + + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + var v = f32(acc[j]) * params.input_scale * t_scales[oc]; + if (params.has_bias != 0u) { + v = v + t_bias[oc]; + } + var q = i32(round(v * params.inv_output_scale)) + params.output_zero_point; + q = clamp(q, -128, 127); + packed = packed | ((bitcast(q) & 0xFFu) << (j * 8u)); + } + t_out[widx] = packed; +} +)"; + +inline constexpr uint32_t kQ8taConv2dTransposedWorkgroupSizeX = 64; +inline constexpr uint32_t kQ8taConv2dTransposedWorkgroupSizeY = 1; +inline constexpr uint32_t kQ8taConv2dTransposedWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From f2c3740d7fc27425f73ff872f7ff6fd3c3b988f7 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:32 -0700 Subject: [PATCH 48/75] [ExecuTorch][WebGPU] Op-tests for q8ta_conv2d_transposed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 31 +++++++++ .../test/ops/test_q8ta_conv2d_transposed.py | 68 +++++++++++++++++++ 2 files changed, 99 insertions(+) create mode 100644 backends/webgpu/test/ops/test_q8ta_conv2d_transposed.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 1f7e1c4d5e2..343193cd8ab 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -66,6 +66,9 @@ from executorch.backends.webgpu.test.ops.test_q8ta_conv2d import ( make_q8ta_conv2d_module, ) +from executorch.backends.webgpu.test.ops.test_q8ta_conv2d_transposed import ( + make_q8ta_conv2d_transposed_module, +) from executorch.backends.webgpu.test.ops.test_pixel_shuffle import PixelShuffleModule from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule @@ -1604,6 +1607,34 @@ def case(name, ic, oc, k, h, w, n=1, **kw): ) +@register_op_test("q8ta_conv2d_transposed") +def _q8ta_conv2d_transposed_suite() -> WebGPUTestSuite: + # Golden = XNNPACK-static-PT2E converted eager (fp32); transposed conv + # (groups==1, dilation==1). W_out%4==0 for output packing. + def case(name, ic, oc, k, h, w, n=1, **kw): + return Case( + name=name, + construct={"ic": ic, "oc": oc, "k": k, "h": h, "w": w, "n": n, **kw}, + inputs=((n, ic, h, w),), + ) + + return WebGPUTestSuite( + module_factory=lambda **kw: make_q8ta_conv2d_transposed_module(**kw), + cases=[ + case("s2", 8, 8, 2, 8, 8, stride=2), + case("no_bias", 8, 8, 2, 8, 8, stride=2, bias=False), + case("k3", 8, 8, 3, 8, 8, stride=2, padding=1, output_padding=1), + case("oc6", 8, 6, 2, 8, 8, stride=2), + case("ic3", 3, 8, 3, 8, 8, stride=2, padding=1, output_padding=1), + case("batch2", 8, 8, 2, 8, 8, stride=2, n=2), + case("asym", 8, 8, (3, 2), 8, 8, stride=2), + ], + golden_dtype="float32", + atol=1e-3, + rtol=1e-3, + ) + + @register_op_test("dequantize_per_tensor") def _dequant_const_suite() -> WebGPUTestSuite: # Independent GPU dequantize check: baked int8 const vs torch dequant (breaks diff --git a/backends/webgpu/test/ops/test_q8ta_conv2d_transposed.py b/backends/webgpu/test/ops/test_q8ta_conv2d_transposed.py new file mode 100644 index 00000000000..bc582a33608 --- /dev/null +++ b/backends/webgpu/test/ops/test_q8ta_conv2d_transposed.py @@ -0,0 +1,68 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`et_vk.q8ta_conv2d_transposed` module + configs: int8 transposed conv. + +An `nn.ConvTranspose2d` (groups == 1, dilation == 1) through XNNPACK static PT2E +lowers to a delegated `quantize_per_tensor -> q8ta_conv2d_transposed -> +dequantize_per_tensor` subgraph (quantize/dequantize = the landed C0 ops). The +`module_factory` returns the CONVERTED module, so the op-test framework goldens +the WebGPU output against the converted eager (fp32 fake-quant), e2e. +""" + +import unittest + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 registers et_vk ops + +import torch +import torch.nn as nn + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( + get_symmetric_quantization_config, + XNNPACKQuantizer, +) +from executorch.exir import to_edge_transform_and_lower +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + +def make_q8ta_conv2d_transposed_module( + ic, oc, k, h, w, stride=2, padding=0, output_padding=0, bias=True, n=1, seed=0 +): + torch.manual_seed(seed) + conv = nn.ConvTranspose2d( + ic, + oc, + k, + stride=stride, + padding=padding, + output_padding=output_padding, + bias=bias, + ).eval() + ex = (torch.randn(n, ic, h, w),) + q = XNNPACKQuantizer().set_global( + get_symmetric_quantization_config(is_per_channel=True, is_dynamic=False) + ) + prepared = prepare_pt2e(torch.export.export(conv, ex).module(), q) + prepared(*ex) # calibrate + return convert_pt2e(prepared) + + +class Q8taConv2dTransposedTest(unittest.TestCase): + def test_delegates(self) -> None: + m = make_q8ta_conv2d_transposed_module(8, 8, 2, 8, 8, stride=2) + et = to_edge_transform_and_lower( + torch.export.export(m, (torch.randn(1, 8, 8, 8),)), + partitioner=[VulkanPartitioner()], + ).to_executorch() + self.assertTrue( + any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + ) + self.assertIn(b"q8ta_conv2d_transposed", et.buffer) From 9c9cad74b7610dc28877c7777c2f7725e9cbed6c Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:33 -0700 Subject: [PATCH 49/75] [ExecuTorch][WebGPU] Add floor_divide op (aten.div.Tensor_mode) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../runtime/ops/floor_divide/BinaryOp.cpp | 46 ++++++++++++ .../ops/floor_divide/binary_floor_divide.wgsl | 51 +++++++++++++ .../floor_divide/binary_floor_divide_wgsl.h | 75 +++++++++++++++++++ 3 files changed, 172 insertions(+) create mode 100644 backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp create mode 100644 backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl create mode 100644 backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h diff --git a/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp b/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp new file mode 100644 index 00000000000..a35a7587311 --- /dev/null +++ b/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp @@ -0,0 +1,46 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +// aten.div.Tensor_mode -> floor(a/b), with NumPy broadcasting (mirrors mul + +// Vulkan). Only rounding_mode='floor' is supported. +void floor_divide_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [in1, in2, rounding_mode(str), out] + const int mode_id = args.at(2); + if (graph.get_value_type(mode_id) != WebGPUGraph::ValueType::String || + graph.get_string(mode_id) != "floor") { + throw std::runtime_error("floor_divide: only rounding_mode='floor'"); + } + add_binary_broadcast_op( + graph, + args.at(0), + args.at(1), + args.at(args.size() - 1), + kBinaryFloorDivideWGSL, + kBinaryFloorDivideWorkgroupSizeX, + "floor_divide"); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.div.Tensor_mode, floor_divide_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl b/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl new file mode 100644 index 00000000000..5191a220238 --- /dev/null +++ b/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl @@ -0,0 +1,51 @@ +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in1_meta: TensorMeta; +@group(0) @binding(5) var in2_meta: TensorMeta; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= out_meta.numel) { + return; + } + + // Fast path: every input dim matches the output dim -> elementwise. + var same = true; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + if (in1_meta.sizes[d] != out_meta.sizes[d] || + in2_meta.sizes[d] != out_meta.sizes[d]) { + same = false; + } + } + if (same) { + output[idx] = floor(input1[idx] / input2[idx]); + return; + } + + // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. + var rem = idx; + var l1: u32 = 0u; + var l2: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d]; + rem = rem % out_meta.strides[d]; + l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d]; + l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d]; + } + output[idx] = floor(input1[l1] / input2[l2]); +} diff --git a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h b/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h new file mode 100644 index 00000000000..63bd6d1f48b --- /dev/null +++ b/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h @@ -0,0 +1,75 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from binary_floor_divide.wgsl - DO NOT EDIT. +// wgsl-sha256: bbc1c3a39dd17c88e7df4f2adbf3b9cfe424693aed00deab268a03b2e0cdc958 +inline constexpr const char* kBinaryFloorDivideWGSL = R"( +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in1_meta: TensorMeta; +@group(0) @binding(5) var in2_meta: TensorMeta; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= out_meta.numel) { + return; + } + + // Fast path: every input dim matches the output dim -> elementwise. + var same = true; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + if (in1_meta.sizes[d] != out_meta.sizes[d] || + in2_meta.sizes[d] != out_meta.sizes[d]) { + same = false; + } + } + if (same) { + output[idx] = floor(input1[idx] / input2[idx]); + return; + } + + // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. + var rem = idx; + var l1: u32 = 0u; + var l2: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d]; + rem = rem % out_meta.strides[d]; + l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d]; + l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d]; + } + output[idx] = floor(input1[l1] / input2[l2]); +} +)"; + +inline constexpr uint32_t kBinaryFloorDivideWorkgroupSizeX = 64; +inline constexpr uint32_t kBinaryFloorDivideWorkgroupSizeY = 1; +inline constexpr uint32_t kBinaryFloorDivideWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 2fe2b6bec33b78b7e7d5a2307e65ba6280d677c0 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:33 -0700 Subject: [PATCH 50/75] [ExecuTorch][WebGPU] Op-tests for floor_divide MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 34 +++++++++++++++++++ backends/webgpu/test/ops/test_floor_divide.py | 21 ++++++++++++ 2 files changed, 55 insertions(+) create mode 100644 backends/webgpu/test/ops/test_floor_divide.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 343193cd8ab..94792fd0f5c 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -34,6 +34,7 @@ CatModule, CONFIGS as _CAT_CONFIGS, ) +from executorch.backends.webgpu.test.ops.test_floor_divide import FloorDivideModule from executorch.backends.webgpu.test.ops.test_flip import FlipModule from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule @@ -263,6 +264,39 @@ def _pow_suite() -> WebGPUTestSuite: ) +def _floor_div_golden(module, inputs): + # Vulkan-faithful oracle: floor(a/b) in fp32 (Vulkan glsl OPERATOR floor(X/Y)), + # NOT torch's fmod-corrected div_floor which can differ by 1 at fp boundaries. + return torch.floor(inputs[0] / inputs[1]) + + +@register_op_test("floor_divide") +def _floor_divide_suite() -> WebGPUTestSuite: + # aten.div.Tensor_mode; divisor bounded away from zero. golden_fn = the + # Vulkan floor(a/b) formula (same formula as the fp32 kernel). + return WebGPUTestSuite( + module_factory=lambda: FloorDivideModule(), + cases=[ + Case( + name="2d", + inputs=( + InputSpec(shape=(M1, M2), gen=_unary_lin(-8.0, 8.0)), + InputSpec(shape=(M1, M2), gen=_unary_lin(0.5, 4.0)), + ), + golden_fn=_floor_div_golden, + ), + Case( + name="3d", + inputs=( + InputSpec(shape=(S, S1, S2), gen=_unary_lin(-8.0, 8.0)), + InputSpec(shape=(S, S1, S2), gen=_unary_lin(0.5, 4.0)), + ), + golden_fn=_floor_div_golden, + ), + ], + ) + + def _reduce_suite(module_cls) -> WebGPUTestSuite: # Last-dim reduction; both keepdim variants over a 2d and a 3d shape. return WebGPUTestSuite( diff --git a/backends/webgpu/test/ops/test_floor_divide.py b/backends/webgpu/test/ops/test_floor_divide.py new file mode 100644 index 00000000000..6020480ff2c --- /dev/null +++ b/backends/webgpu/test/ops/test_floor_divide.py @@ -0,0 +1,21 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.div.Tensor_mode` module for the WebGPU op-test framework. + +`FloorDivideModule` is imported by `cases.py`. Same-shape elementwise +`div(a, b, rounding_mode="floor")`. The kernel computes `floor(a/b)` mirroring +the Vulkan `floor_divide` glsl (`floor(X/Y)`); this differs from torch's own +fmod-corrected `div_floor` at rare fp boundaries, so the suite goldens against a +`floor(a/b)` `golden_fn` (Vulkan-faithful), not this module's eager output. +""" + +import torch + + +class FloorDivideModule(torch.nn.Module): + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return torch.div(a, b, rounding_mode="floor") From 2c24bd696b35ad4febde03fe2b14951ae8e59a35 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:34 -0700 Subject: [PATCH 51/75] [ExecuTorch][WebGPU] Add argmax/argmin ops + int64-output path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/runtime/WebGPUBackend.cpp | 26 ++- backends/webgpu/runtime/WebGPUGraph.cpp | 68 ++++-- backends/webgpu/runtime/ops/argmax/Reduce.cpp | 215 ++++++++++++++++++ .../webgpu/runtime/ops/argmax/arg_reduce.wgsl | 41 ++++ .../runtime/ops/argmax/arg_reduce_wgsl.h | 65 ++++++ 5 files changed, 388 insertions(+), 27 deletions(-) create mode 100644 backends/webgpu/runtime/ops/argmax/Reduce.cpp create mode 100644 backends/webgpu/runtime/ops/argmax/arg_reduce.wgsl create mode 100644 backends/webgpu/runtime/ops/argmax/arg_reduce_wgsl.h diff --git a/backends/webgpu/runtime/WebGPUBackend.cpp b/backends/webgpu/runtime/WebGPUBackend.cpp index e1497bc4ed8..e7645514c9b 100644 --- a/backends/webgpu/runtime/WebGPUBackend.cpp +++ b/backends/webgpu/runtime/WebGPUBackend.cpp @@ -164,18 +164,22 @@ Error WebGPUBackend::execute( return Error::Internal; } - // Execute the compute graph - graph->execute(); - - // Copy outputs from GPU staging buffers to EValue tensor data pointers - std::vector> outputs; - outputs.reserve(num_outputs); - for (size_t i = 0; i < num_outputs; i++) { - const size_t arg_idx = num_inputs + i; - auto& tensor = args[arg_idx]->toTensor(); - outputs.emplace_back(tensor.mutable_data_ptr(), tensor.nbytes()); + // Execute + read back; fail loud as a runtime Error so a throw never crosses + // the backend boundary. + try { + graph->execute(); + std::vector> outputs; + outputs.reserve(num_outputs); + for (size_t i = 0; i < num_outputs; i++) { + const size_t arg_idx = num_inputs + i; + auto& tensor = args[arg_idx]->toTensor(); + outputs.emplace_back(tensor.mutable_data_ptr(), tensor.nbytes()); + } + graph->copy_outputs(outputs); + } catch (const std::exception& e) { + ET_LOG(Error, "WebGPU execute / output copy failed: %s", e.what()); + return Error::Internal; } - graph->copy_outputs(outputs); return Error::Ok; } diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index dc0662b1102..e3f4b1e48d0 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -1903,9 +1903,12 @@ void WebGPUGraph::copy_outputs(std::vector>& outputs) { std::vector cb_data(count); std::vector map_futures(count, WGPUFuture{}); + // Map each output's LIVE staging size (an int64 output is int32-backed). + std::vector map_nbytes(count, 0); for (size_t i = 0; i < count; i++) { - if (outputs[i].second == 0) { + map_nbytes[i] = tensors_[output_ids_[i]].cur_nbytes; + if (map_nbytes[i] == 0) { cb_data[i].status = WGPUMapAsyncStatus_Success; continue; } @@ -1917,29 +1920,62 @@ void WebGPUGraph::copy_outputs(std::vector>& outputs) { output_staging_buffers_[i], WGPUMapMode_Read, 0, - outputs[i].second, + map_nbytes[i], cb_info); } - for (size_t i = 0; i < count; i++) { - if (outputs[i].second != 0 && - webgpu_wait(instance_, map_futures[i]) != WGPUWaitStatus_Success) { - throw std::runtime_error("WebGPU: WaitAny failed for output map"); - } - } + // Tracks which output buffers are currently mapped so a mid-loop throw can + // release them before propagating (no dangling mapped buffers). + std::vector is_mapped(count, false); - for (size_t i = 0; i < count; i++) { - if (outputs[i].second == 0) { - continue; + try { + for (size_t i = 0; i < count; i++) { + if (map_nbytes[i] == 0) { + continue; + } + if (webgpu_wait(instance_, map_futures[i]) != WGPUWaitStatus_Success) { + throw std::runtime_error("WebGPU: WaitAny failed for output map"); + } + if (cb_data[i].status == WGPUMapAsyncStatus_Success) { + is_mapped[i] = true; + } } - if (cb_data[i].status == WGPUMapAsyncStatus_Success) { + + for (size_t i = 0; i < count; i++) { + if (map_nbytes[i] == 0) { + continue; + } + if (cb_data[i].status != WGPUMapAsyncStatus_Success) { + throw std::runtime_error("WebGPU buffer map failed for output"); + } const void* mapped = wgpuBufferGetConstMappedRange( - output_staging_buffers_[i], 0, outputs[i].second); - std::memcpy(outputs[i].first, mapped, outputs[i].second); + output_staging_buffers_[i], 0, map_nbytes[i]); + const size_t dst_nbytes = outputs[i].second; + if (dst_nbytes == map_nbytes[i]) { + std::memcpy(outputs[i].first, mapped, map_nbytes[i]); + } else if ( + dst_nbytes == 2 * map_nbytes[i] && tensors_[output_ids_[i]].is_int && + tensors_[output_ids_[i]].elem_size == 4) { + // int64 host output backed by an int32 GPU buffer: widen (sign-extend). + const int32_t* src = static_cast(mapped); + int64_t* dst = static_cast(outputs[i].first); + const size_t n = map_nbytes[i] / sizeof(int32_t); + for (size_t k = 0; k < n; k++) { + dst[k] = static_cast(src[k]); + } + } else { + throw std::runtime_error("WebGPU: output buffer size mismatch"); + } wgpuBufferUnmap(output_staging_buffers_[i]); - } else { - throw std::runtime_error("WebGPU buffer map failed for output"); + is_mapped[i] = false; + } + } catch (...) { + for (size_t j = 0; j < count; j++) { + if (is_mapped[j]) { + wgpuBufferUnmap(output_staging_buffers_[j]); + } } + throw; } } diff --git a/backends/webgpu/runtime/ops/argmax/Reduce.cpp b/backends/webgpu/runtime/ops/argmax/Reduce.cpp new file mode 100644 index 00000000000..636cae19b42 --- /dev/null +++ b/backends/webgpu/runtime/ops/argmax/Reduce.cpp @@ -0,0 +1,215 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform layout matching the WGSL Params struct; 16-byte aligned. +struct ArgReduceParams { + uint32_t num_rows; + uint32_t reduce_size; + uint32_t is_argmin; + uint32_t _pad; +}; + +// Last-dim argmax/argmin -> int32 index; mirrors Vulkan arg_reduce_impl. +void arg_reduce_impl( + WebGPUGraph& graph, + const std::vector& args, + uint32_t is_argmin) { + // aten.argmax/argmin.default args: [in, dim, keepdim, out] + const int in_id = args.at(0); + const int out_id = args.at(args.size() - 1); + + WGPUDevice device = graph.device(); + + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("arg_reduce: null buffer binding"); + } + // fp32 input; int index output (the AOT downcasts the int64 index to int32). + if (in_tensor.is_int || in_tensor.elem_size != 4) { + throw std::runtime_error("arg_reduce: input must be fp32"); + } + if (!out_tensor.is_int || out_tensor.elem_size != 4) { + throw std::runtime_error("arg_reduce: output must be int32 index"); + } + + const int64_t ndim = static_cast(in_tensor.dims.size()); + const int64_t dim = graph.get_int(args.at(1)); + if (dim != -1 && dim != ndim - 1) { + throw std::runtime_error( + "arg_reduce: only last-dim reduction is supported"); + } + + const uint32_t reduce_size = static_cast(in_tensor.dims.back()); + const uint32_t num_rows = + static_cast(out_tensor.nbytes / sizeof(int32_t)); + if (reduce_size == 0u || + in_tensor.nbytes / sizeof(float) != + static_cast(num_rows) * reduce_size) { + throw std::runtime_error("arg_reduce: shape mismatch (rows * reduce_size)"); + } + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kArgReduceWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, num_rows, wg_size, "arg_reduce"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + ArgReduceParams params = {}; + params.num_rows = num_rows; + params.reduce_size = reduce_size; + params.is_argmin = is_argmin; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(ArgReduceParams)); + graph.add_uniform_buffer_bytes(sizeof(ArgReduceParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kArgReduceWGSL, WGPU_STRLEN}; + + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // Bind group layout: input (read storage) + output (storage) + params. + WGPUBindGroupLayoutEntry entries[3] = {}; + + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 3; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[3] = {}; + + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + + bg_entries[2].binding = 2; + bg_entries[2].buffer = uniform_buffer; + bg_entries[2].size = sizeof(ArgReduceParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 3; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "arg_reduce", + workgroup_count.y}); + + // Dynamic shapes: recompute reduce_size (last dim) + num_rows + dispatch. + const bool keepdim = graph.get_bool(args.at(2)); + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, keepdim, is_argmin, wg_size, dispatch_idx, params_buf]( + WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + const uint32_t rsize = static_cast(d.back()); + if (rsize == 0u) { + throw std::runtime_error("arg_reduce(resize): zero reduce dim"); + } + const uint64_t total = utils::numel_of(d); + const uint32_t rows = static_cast(total / rsize); + std::vector od = d; + if (keepdim) { + od.back() = 1; + } else { + od.pop_back(); + } + g.set_cur_dims(out_id, od); + ArgReduceParams p = {}; + p.num_rows = rows; + p.reduce_size = rsize; + p.is_argmin = is_argmin; + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), rows, wg_size, "arg_reduce"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }); + + // Release intermediates (pipeline + bind_group are kept by dispatch). + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); +} + +void argmax_impl(WebGPUGraph& graph, const std::vector& args) { + arg_reduce_impl(graph, args, 0u); +} + +void argmin_impl(WebGPUGraph& graph, const std::vector& args) { + arg_reduce_impl(graph, args, 1u); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.argmax.default, argmax_impl); + WEBGPU_REGISTER_OP(aten.argmin.default, argmin_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/argmax/arg_reduce.wgsl b/backends/webgpu/runtime/ops/argmax/arg_reduce.wgsl new file mode 100644 index 00000000000..ef77e9bffd8 --- /dev/null +++ b/backends/webgpu/runtime/ops/argmax/arg_reduce.wgsl @@ -0,0 +1,41 @@ +@group(0) @binding(0) var t_in: array; +@group(0) @binding(1) var t_out: array; + +struct Params { + num_rows: u32, + reduce_size: u32, + is_argmin: u32, + pad0: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let row = gid.x + gid.y * (num_workgroups.x * wg_size); + if (row >= params.num_rows) { + return; + } + let base = row * params.reduce_size; + // Strict compare keeps the FIRST extremum = torch argmax/argmin tie-break. + var best = t_in[base]; + var best_idx: u32 = 0u; + for (var k: u32 = 1u; k < params.reduce_size; k = k + 1u) { + let v = t_in[base + k]; + if (params.is_argmin != 0u) { + if (v < best) { + best = v; + best_idx = k; + } + } else { + if (v > best) { + best = v; + best_idx = k; + } + } + } + t_out[row] = best_idx; +} diff --git a/backends/webgpu/runtime/ops/argmax/arg_reduce_wgsl.h b/backends/webgpu/runtime/ops/argmax/arg_reduce_wgsl.h new file mode 100644 index 00000000000..3ec5403acbe --- /dev/null +++ b/backends/webgpu/runtime/ops/argmax/arg_reduce_wgsl.h @@ -0,0 +1,65 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from arg_reduce.wgsl - DO NOT EDIT. +// wgsl-sha256: 0ee2f2b47cb29c1b26339314f6ede857efb0f558e53c6433bb3aae984952e1dd +inline constexpr const char* kArgReduceWGSL = R"( +@group(0) @binding(0) var t_in: array; +@group(0) @binding(1) var t_out: array; + +struct Params { + num_rows: u32, + reduce_size: u32, + is_argmin: u32, + pad0: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let row = gid.x + gid.y * (num_workgroups.x * wg_size); + if (row >= params.num_rows) { + return; + } + let base = row * params.reduce_size; + // Strict compare keeps the FIRST extremum = torch argmax/argmin tie-break. + var best = t_in[base]; + var best_idx: u32 = 0u; + for (var k: u32 = 1u; k < params.reduce_size; k = k + 1u) { + let v = t_in[base + k]; + if (params.is_argmin != 0u) { + if (v < best) { + best = v; + best_idx = k; + } + } else { + if (v > best) { + best = v; + best_idx = k; + } + } + } + t_out[row] = best_idx; +} +)"; + +inline constexpr uint32_t kArgReduceWorkgroupSizeX = 64; +inline constexpr uint32_t kArgReduceWorkgroupSizeY = 1; +inline constexpr uint32_t kArgReduceWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From bca821012e96375beea91c87130130c21da046f7 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:34 -0700 Subject: [PATCH 52/75] [ExecuTorch][WebGPU] Op-tests for argmax/argmin + int64-golden harness MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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()` 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/) --- backends/webgpu/test/op_tests/cases.py | 35 ++++++++++++++ backends/webgpu/test/op_tests/driver_util.cpp | 14 ++++++ backends/webgpu/test/op_tests/driver_util.h | 5 +- .../webgpu/test/op_tests/generate_op_tests.py | 11 +++++ .../webgpu/test/op_tests/op_test_driver.cpp | 17 +++++++ backends/webgpu/test/ops/test_argmax.py | 46 +++++++++++++++++++ 6 files changed, 127 insertions(+), 1 deletion(-) create mode 100644 backends/webgpu/test/ops/test_argmax.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 94792fd0f5c..8b0446955b2 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -35,6 +35,12 @@ CONFIGS as _CAT_CONFIGS, ) from executorch.backends.webgpu.test.ops.test_floor_divide import FloorDivideModule +from executorch.backends.webgpu.test.ops.test_argmax import ( + argmax_tie_gen, + argmin_tie_gen, + ArgmaxModule, + ArgminModule, +) from executorch.backends.webgpu.test.ops.test_flip import FlipModule from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule @@ -1683,3 +1689,32 @@ def _dequant_const_suite() -> WebGPUTestSuite: ], golden_dtype="float32", ) + + +@register_op_test("argmax") +def _argmax_suite() -> WebGPUTestSuite: + # Last-dim argmax -> int64 index. randn puts the max at an INTERIOR index + # (exercises the walk); the tie case pins the strict-`>` first-occurrence. + return WebGPUTestSuite( + module_factory=lambda: ArgmaxModule(), + cases=[ + Case(name="2d", inputs=((M1, M2),)), + Case(name="3d", inputs=((S, S1, S2),)), + Case(name="tie", inputs=(InputSpec(shape=(3, 6), gen=argmax_tie_gen),)), + ], + golden_dtype="float32", + ) + + +@register_op_test("argmin") +def _argmin_suite() -> WebGPUTestSuite: + # Last-dim argmin -> int64 index; randn interior min + strict-`<` tie case. + return WebGPUTestSuite( + module_factory=lambda: ArgminModule(), + cases=[ + Case(name="2d", inputs=((M1, M2),)), + Case(name="3d", inputs=((S, S1, S2),)), + Case(name="tie", inputs=(InputSpec(shape=(3, 6), gen=argmin_tie_gen),)), + ], + golden_dtype="float32", + ) diff --git a/backends/webgpu/test/op_tests/driver_util.cpp b/backends/webgpu/test/op_tests/driver_util.cpp index 75d398fdacf..455f61c9532 100644 --- a/backends/webgpu/test/op_tests/driver_util.cpp +++ b/backends/webgpu/test/op_tests/driver_util.cpp @@ -108,6 +108,20 @@ std::vector load_int8_bin(const std::string& path, size_t numel) { return g; } +std::vector load_int64_bin(const std::string& path, size_t numel) { + FILE* f = std::fopen(path.c_str(), "rb"); + if (!f) { + return {}; + } + std::vector g(numel); + const size_t n = std::fread(g.data(), sizeof(int64_t), numel, f); + std::fclose(f); + if (n != numel) { + return {}; + } + return g; +} + bool within_tol( const float* out, const float* golden, diff --git a/backends/webgpu/test/op_tests/driver_util.h b/backends/webgpu/test/op_tests/driver_util.h index c171c114309..96998cec49e 100644 --- a/backends/webgpu/test/op_tests/driver_util.h +++ b/backends/webgpu/test/op_tests/driver_util.h @@ -24,7 +24,7 @@ struct GoldenRef { std::string path; std::vector shape; int output_index = 0; - std::string dtype = "float32"; // "float32" or "int8" (int8-output ops) + std::string dtype = "float32"; // "float32" | "int8" | "int64" (argmax index) }; struct ManifestEntry { @@ -49,6 +49,9 @@ std::vector load_int32_bin(const std::string& path, size_t numel); /// Load raw int8 (one byte per element); empty on size/IO mismatch. std::vector load_int8_bin(const std::string& path, size_t numel); +/// Load raw little-endian int64 (8B/elem); empty on size/IO mismatch. +std::vector load_int64_bin(const std::string& path, size_t numel); + /// Element OK if abs_err <= atol OR rel_err <= rtol (rel floored at /// |golden|=1e-6). Sets the reported maxima; true iff all elements pass. bool within_tol( diff --git a/backends/webgpu/test/op_tests/generate_op_tests.py b/backends/webgpu/test/op_tests/generate_op_tests.py index 862d5eaa98d..1907f20bbe1 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -82,6 +82,10 @@ def _write_int8(t: torch.Tensor, path: str) -> None: t.detach().contiguous().cpu().numpy().astype(" None: + t.detach().contiguous().cpu().numpy().astype(" list[dict]: """Export one case; write its .pte + input/golden .bin(s). Returns one manifest entry per output tensor (multi-output ops emit N; single-output ops emit one, @@ -135,10 +139,15 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> list[d for out_index in range(n_out): raw = golden_outs[out_index] is_int8 = raw.dtype == torch.int8 + is_int64 = raw.dtype in (torch.int64, torch.int32) if is_int8: # int8-output ops (quantize): byte-exact golden, no fp32 cast/oracle. out_t = raw out_dtype = "int8" + elif is_int64: + # int64-index ops (argmax/argmin): exact golden, no fp32 cast/oracle. + out_t = raw.to(torch.int64) + out_dtype = "int64" else: out_t = raw.to(torch.float32) out_dtype = "float32" @@ -156,6 +165,8 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> list[d golden_rel = f"{case_id}{suffix}.golden.bin" if is_int8: _write_int8(out_t, os.path.join(out_dir, golden_rel)) + elif is_int64: + _write_int64(out_t, os.path.join(out_dir, golden_rel)) else: _write_fp32(out_t, os.path.join(out_dir, golden_rel)) entries.append( diff --git a/backends/webgpu/test/op_tests/op_test_driver.cpp b/backends/webgpu/test/op_tests/op_test_driver.cpp index ba5995ceb81..c29daf6e101 100644 --- a/backends/webgpu/test/op_tests/op_test_driver.cpp +++ b/backends/webgpu/test/op_tests/op_test_driver.cpp @@ -108,6 +108,23 @@ class OpCase : public ::testing::Test { EXPECT_EQ(mism, -1) << "int8 mismatch at index " << mism << ": out=" << (mism >= 0 ? int(out_p[mism]) : 0) << " golden=" << (mism >= 0 ? int(golden[mism]) : 0); + } else if (e_.golden.dtype == "int64") { + // int64-index ops (argmax/argmin) compare exact: indices are discrete. + auto golden = load_int64_bin(e_.golden.path, gn); + ASSERT_FALSE(golden.empty()) + << "missing/short golden: " << e_.golden.path; + const int64_t* out_p = out_tensor.const_data_ptr(); + int mism = -1; + for (size_t i = 0; i < gn; i++) { + if (out_p[i] != golden[i]) { + mism = static_cast(i); + break; + } + } + EXPECT_EQ(mism, -1) << "int64 mismatch at index " << mism << ": out=" + << (mism >= 0 ? (long long)out_p[mism] : 0) + << " golden=" + << (mism >= 0 ? (long long)golden[mism] : 0); } else { auto golden = load_fp32_bin(e_.golden.path, gn); ASSERT_FALSE(golden.empty()) diff --git a/backends/webgpu/test/ops/test_argmax.py b/backends/webgpu/test/ops/test_argmax.py new file mode 100644 index 00000000000..73aed92ca23 --- /dev/null +++ b/backends/webgpu/test/ops/test_argmax.py @@ -0,0 +1,46 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.argmax.default` / `aten.argmin.default` modules for the op-test framework. + +Last-dim arg-reduction -> int64 index. The kernel writes an int32 index (the AOT +downcasts the int64 output), which `copy_outputs` widens to the int64 program +output. `golden_dtype="float32"` so the golden's fp32 argmax matches the kernel's +fp32 argmax exactly (avoids an fp32-vs-fp64 tie-break discriminator flip). +""" + +import torch + + +class ArgmaxModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.argmax(x, dim=-1) + + +class ArgminModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.argmin(x, dim=-1) + + +def argmax_tie_gen(shape): + # Repeated max in the last dim (idx 1 and 3); the FIRST wins (index 1) — + # discriminates the strict-`>` tie-break from a `>=` (last-wins) bug. + assert shape[-1] >= 4, "tie generator writes idx 1 and 3; last dim must be >= 4" + base = torch.arange(shape[-1], dtype=torch.float32) * 0.01 + t = base.expand(shape).contiguous() + t[..., 1] = 10.0 + t[..., 3] = 10.0 + return t + + +def argmin_tie_gen(shape): + # Repeated min in the last dim (idx 1 and 3); the FIRST wins (index 1). + assert shape[-1] >= 4, "tie generator writes idx 1 and 3; last dim must be >= 4" + base = 5.0 + torch.arange(shape[-1], dtype=torch.float32) * 0.01 + t = base.expand(shape).contiguous() + t[..., 1] = -10.0 + t[..., 3] = -10.0 + return t From dbf93ca57eca67a47bc688f62fbd6edc64c570d3 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:35 -0700 Subject: [PATCH 53/75] [ExecuTorch][WebGPU] Add linear_qcs4w op (et_vk.linear_qcs4w) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../ops/linear_qcs4w/QuantizedLinearQcs4w.cpp | 254 ++++++++++++++++++ .../ops/linear_qcs4w/qcs4w_linear.wgsl | 87 ++++++ .../ops/linear_qcs4w/qcs4w_linear_wgsl.h | 111 ++++++++ 3 files changed, 452 insertions(+) create mode 100644 backends/webgpu/runtime/ops/linear_qcs4w/QuantizedLinearQcs4w.cpp create mode 100644 backends/webgpu/runtime/ops/linear_qcs4w/qcs4w_linear.wgsl create mode 100644 backends/webgpu/runtime/ops/linear_qcs4w/qcs4w_linear_wgsl.h diff --git a/backends/webgpu/runtime/ops/linear_qcs4w/QuantizedLinearQcs4w.cpp b/backends/webgpu/runtime/ops/linear_qcs4w/QuantizedLinearQcs4w.cpp new file mode 100644 index 00000000000..d75c2617072 --- /dev/null +++ b/backends/webgpu/runtime/ops/linear_qcs4w/QuantizedLinearQcs4w.cpp @@ -0,0 +1,254 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform layout matching the WGSL Params struct; 16-byte aligned. +struct Qcs4wParams { + uint32_t M; + uint32_t N; + uint32_t K; + uint32_t K_packed; +}; +static_assert(sizeof(Qcs4wParams) == 16, "Qcs4wParams must be 16 bytes"); + +// Register-tile dims; MUST match TM/TN in qcs4w_linear.wgsl. +constexpr int64_t kQcs4wTileM = 4; +constexpr int64_t kQcs4wTileN = 4; + +utils::WgCount compute_qcs4w_workgroup_count( + WGPUDevice device, + uint32_t m, + uint32_t n, + uint32_t wg_size) { + const int64_t total_tiles = utils::div_up(m, kQcs4wTileM) * + utils::div_up(n, kQcs4wTileN); + if (total_tiles > static_cast(UINT32_MAX)) { + throw std::runtime_error("WebGPU linear_qcs4w: tile count exceeds u32"); + } + return utils::compute_2d_workgroup_count( + device, static_cast(total_tiles), wg_size, "linear_qcs4w"); +} + +// linear_qcs4w: 4-bit channels-symmetric weight, per-channel scale (no bias). +void qcs4w_linear_impl(WebGPUGraph& graph, const std::vector& args) { + const int in_id = args.at(0); + const int weight_id = args.at(1); + const int scales_id = args.at(2); + const int out_id = args.at(args.size() - 1); + + WGPUDevice device = graph.device(); + + const auto& in = graph.get_tensor(in_id); + const auto& weight = graph.get_tensor(weight_id); + const auto& scales = graph.get_tensor(scales_id); + const auto& out = graph.get_tensor(out_id); + if (in.buffer == nullptr || weight.buffer == nullptr || + scales.buffer == nullptr || out.buffer == nullptr) { + throw std::runtime_error("WebGPU linear_qcs4w: null buffer binding"); + } + if (in.dims.empty() || weight.dims.size() < 2 || scales.dims.empty()) { + throw std::runtime_error("WebGPU linear_qcs4w: malformed input dims"); + } + + const uint32_t K = static_cast(in.dims.back()); + if (K == 0) { + throw std::runtime_error("WebGPU linear_qcs4w: K == 0"); + } + uint64_t in_numel = 1; + for (int64_t d : in.dims) { + in_numel *= static_cast(d); + } + if (in_numel % K != 0) { + throw std::runtime_error( + "WebGPU linear_qcs4w: input numel not a multiple of K"); + } + const uint32_t M = static_cast(in_numel / K); + const uint32_t N = static_cast(weight.dims[0]); + const uint32_t K_packed = static_cast(weight.dims[1]); + if (M == 0 || N == 0) { + throw std::runtime_error("WebGPU linear_qcs4w: M or N == 0"); + } + // int4 packing is 2 nibbles/byte, so K_packed must be ceil(K/2) (guards OOB). + if (K_packed != (K + 1) / 2) { + throw std::runtime_error("WebGPU linear_qcs4w: K_packed must be ceil(K/2)"); + } + // Weight is read as array; a non-multiple-of-4 byte count over-reads. + if ((static_cast(N) * K_packed) % 4u != 0u) { + throw std::runtime_error( + "WebGPU linear_qcs4w: N*K_packed must be a multiple of 4 (u32-packed)"); + } + + // fp32-only byte-size guards; scales is per-output-channel (1D, N entries). + const uint64_t weight_numel = + static_cast(N) * static_cast(K_packed); + if (in.nbytes != in_numel * sizeof(float) || + out.nbytes != static_cast(M) * N * sizeof(float) || + scales.nbytes < static_cast(N) * sizeof(float) || + weight.nbytes != weight_numel) { + throw std::runtime_error( + "WebGPU linear_qcs4w: fp32-only (byte-size mismatch)"); + } + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kQcs4wLinearWorkgroupSizeX); + const utils::WgCount workgroup_count = + compute_qcs4w_workgroup_count(device, M, N, wg_size); + + Qcs4wParams params = {}; + params.M = M; + params.N = N; + params.K = K; + params.K_packed = K_packed; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(Qcs4wParams)); + graph.add_uniform_buffer_bytes(sizeof(Qcs4wParams)); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kQcs4wLinearWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // Bind group layout: out (rw) + in/weight/scales (ro storage) + uniform. + WGPUBindGroupLayoutEntry entries[5] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_Storage; + for (uint32_t i = 1; i <= 3; i++) { + entries[i].binding = i; + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + } + entries[4].binding = 4; + entries[4].visibility = WGPUShaderStage_Compute; + entries[4].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 5; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[5] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = out.buffer; + bg_entries[0].size = out.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = in.buffer; + bg_entries[1].size = in.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = weight.buffer; + bg_entries[2].size = weight.nbytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = scales.buffer; + bg_entries[3].size = scales.nbytes; + bg_entries[4].binding = 4; + bg_entries[4].buffer = uniform_buffer; + bg_entries[4].size = sizeof(Qcs4wParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 5; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "linear_qcs4w", + workgroup_count.y}); + + // Dynamic shapes: recompute dispatch + params.M for the live M. + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, M, K, N, K_packed, wg_size, dispatch_idx, uniform_buffer]( + WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + if (d.empty()) { + throw std::runtime_error( + "WebGPU linear_qcs4w(resize): empty input dims"); + } + if (static_cast(d.back()) != K) { + throw std::runtime_error( + "WebGPU linear_qcs4w(resize): last dim must equal K"); + } + const uint64_t numel = utils::numel_of(d); + if (numel % static_cast(K) != 0u) { + throw std::runtime_error( + "WebGPU linear_qcs4w(resize): live numel not a multiple of K"); + } + const uint32_t m = + static_cast(numel / static_cast(K)); + if (m == 0u || m > M) { + throw std::runtime_error( + "WebGPU linear_qcs4w(resize): live M is 0 or exceeds build max"); + } + const utils::WgCount wgc = + compute_qcs4w_workgroup_count(g.device(), m, N, wg_size); + Qcs4wParams p = {}; + p.M = m; + p.N = N; + p.K = K; + p.K_packed = K_packed; + wgpuQueueWriteBuffer(g.queue(), uniform_buffer, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + std::vector od(d.begin(), d.end()); + od.back() = static_cast(N); + g.set_cur_dims(out_id, od); + }); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.linear_qcs4w.default, qcs4w_linear_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/linear_qcs4w/qcs4w_linear.wgsl b/backends/webgpu/runtime/ops/linear_qcs4w/qcs4w_linear.wgsl new file mode 100644 index 00000000000..2183bfbb672 --- /dev/null +++ b/backends/webgpu/runtime/ops/linear_qcs4w/qcs4w_linear.wgsl @@ -0,0 +1,87 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_input: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; + +struct Params { + M: u32, + N: u32, + K: u32, + K_packed: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +// Register-tiled GEMM: dequant weight once per (n,k), reused across TM rows. +const TM: u32 = 4u; +const TN: u32 = 4u; +const TILE_ELEMS: u32 = TM * TN; // acc size; kept in sync with TM/TN + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let nrt = (params.M + TM - 1u) / TM; + let nct = (params.N + TN - 1u) / TN; + let tiles = nrt * nct; + // 2D-folded flat tile index (lifts the 65535 1D-dispatch cap for large M/N). + let tile = gid.x + gid.y * (num_workgroups.x * wg_size); + // M==0 or N==0 -> tiles==0 -> every thread returns, so the M-1u/N-1u clamps + // below never underflow (the host also rejects M==0/N==0). + if (tile >= tiles) { + return; + } + let row_tile = tile / nct; + let col_tile = tile % nct; + let m0 = row_tile * TM; + let n0 = col_tile * TN; + + var acc: array; + for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) { + acc[i] = 0.0; + } + + var k: u32 = 0u; + loop { + if (k >= params.K) { + break; + } + // Load the TM input values for column k once; reused across all TN columns. + var in_reg: array; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m_eff = min(m0 + ml, params.M - 1u); + in_reg[ml] = t_input[m_eff * params.K + k]; + } + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + // Clamp to last valid column; overhang result is never stored. + let n_eff = min(n0 + nl, params.N - 1u); + let byte_idx = n_eff * params.K_packed + (k >> 1u); + let word = t_weight[byte_idx >> 2u]; + let b = (word >> ((byte_idx & 3u) * 8u)) & 0xFFu; + var nib: u32; + // qcs4w prepack packs (even<<4)|odd (swapped vs q4gsw's pack_4bit). + if ((k & 1u) == 0u) { + nib = (b >> 4u) & 0x0Fu; // even k -> high nibble + } else { + nib = b & 0x0Fu; // odd k -> low nibble + } + let q = f32(i32(nib) - 8); // +8-shifted on pack; recover signed [-8,7] + let dq = q * t_scales[n_eff]; // per-channel symmetric scale + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + acc[ml * TN + nl] = acc[ml * TN + nl] + in_reg[ml] * dq; + } + } + k = k + 1u; + } + + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m = m0 + ml; + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n = n0 + nl; + if (m < params.M && n < params.N) { + t_out[m * params.N + n] = acc[ml * TN + nl]; + } + } + } +} diff --git a/backends/webgpu/runtime/ops/linear_qcs4w/qcs4w_linear_wgsl.h b/backends/webgpu/runtime/ops/linear_qcs4w/qcs4w_linear_wgsl.h new file mode 100644 index 00000000000..a56e01b96bd --- /dev/null +++ b/backends/webgpu/runtime/ops/linear_qcs4w/qcs4w_linear_wgsl.h @@ -0,0 +1,111 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from qcs4w_linear.wgsl - DO NOT EDIT. +// wgsl-sha256: 59a17d804c020c01322cbbd25d0984183fa7016186b77bfedd53733185f83fc6 +inline constexpr const char* kQcs4wLinearWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_input: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; + +struct Params { + M: u32, + N: u32, + K: u32, + K_packed: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +// Register-tiled GEMM: dequant weight once per (n,k), reused across TM rows. +const TM: u32 = 4u; +const TN: u32 = 4u; +const TILE_ELEMS: u32 = TM * TN; // acc size; kept in sync with TM/TN + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let nrt = (params.M + TM - 1u) / TM; + let nct = (params.N + TN - 1u) / TN; + let tiles = nrt * nct; + // 2D-folded flat tile index (lifts the 65535 1D-dispatch cap for large M/N). + let tile = gid.x + gid.y * (num_workgroups.x * wg_size); + // M==0 or N==0 -> tiles==0 -> every thread returns, so the M-1u/N-1u clamps + // below never underflow (the host also rejects M==0/N==0). + if (tile >= tiles) { + return; + } + let row_tile = tile / nct; + let col_tile = tile % nct; + let m0 = row_tile * TM; + let n0 = col_tile * TN; + + var acc: array; + for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) { + acc[i] = 0.0; + } + + var k: u32 = 0u; + loop { + if (k >= params.K) { + break; + } + // Load the TM input values for column k once; reused across all TN columns. + var in_reg: array; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m_eff = min(m0 + ml, params.M - 1u); + in_reg[ml] = t_input[m_eff * params.K + k]; + } + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + // Clamp to last valid column; overhang result is never stored. + let n_eff = min(n0 + nl, params.N - 1u); + let byte_idx = n_eff * params.K_packed + (k >> 1u); + let word = t_weight[byte_idx >> 2u]; + let b = (word >> ((byte_idx & 3u) * 8u)) & 0xFFu; + var nib: u32; + // qcs4w prepack packs (even<<4)|odd (swapped vs q4gsw's pack_4bit). + if ((k & 1u) == 0u) { + nib = (b >> 4u) & 0x0Fu; // even k -> high nibble + } else { + nib = b & 0x0Fu; // odd k -> low nibble + } + let q = f32(i32(nib) - 8); // +8-shifted on pack; recover signed [-8,7] + let dq = q * t_scales[n_eff]; // per-channel symmetric scale + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + acc[ml * TN + nl] = acc[ml * TN + nl] + in_reg[ml] * dq; + } + } + k = k + 1u; + } + + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m = m0 + ml; + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n = n0 + nl; + if (m < params.M && n < params.N) { + t_out[m * params.N + n] = acc[ml * TN + nl]; + } + } + } +} +)"; + +inline constexpr uint32_t kQcs4wLinearWorkgroupSizeX = 64; +inline constexpr uint32_t kQcs4wLinearWorkgroupSizeY = 1; +inline constexpr uint32_t kQcs4wLinearWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 345f96b8cc21ec8761e8a52631dcd27156bdad07 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:35 -0700 Subject: [PATCH 54/75] [ExecuTorch][WebGPU] Op-tests for linear_qcs4w MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 27 +++++++++ backends/webgpu/test/ops/test_linear_qcs4w.py | 60 +++++++++++++++++++ 2 files changed, 87 insertions(+) create mode 100644 backends/webgpu/test/ops/test_linear_qcs4w.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 8b0446955b2..b1b075ecd27 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -61,6 +61,9 @@ from executorch.backends.webgpu.test.ops.test_q8ta_pixel_shuffle import ( Q8taPixelShuffleModule, ) +from executorch.backends.webgpu.test.ops.test_linear_qcs4w import ( + make_qcs4w_linear_module, +) from executorch.backends.webgpu.test.ops.test_q8ta_linear import ( make_q8ta_linear_module, ) @@ -1544,6 +1547,30 @@ def case(name, n_ch, h, w, **kw): ) +@register_op_test("linear_qcs4w") +def _linear_qcs4w_suite() -> WebGPUTestSuite: + # VulkanQuantizer weight-only 4-bit nn.Linear -> delegated et_vk.linear_qcs4w. + # Golden = converted eager (fp32 per-channel fake-quant). K even (2 nibbles/ + # byte), N*ceil(K/2) % 4 == 0 (u32-packed weight). M==1 = decode/gemv shape. + def case(name, m, k, n, **kw): + return Case( + name=name, construct={"k": k, "n": n, "m": m, **kw}, inputs=((m, k),) + ) + + return WebGPUTestSuite( + module_factory=lambda **kw: make_qcs4w_linear_module(**kw), + cases=[ + case("basic", 4, 32, 16), + case("gemv", 1, 32, 16), # M==1 + case("k64", 2, 64, 8), + case("n32", 3, 32, 32), + ], + golden_dtype="float32", + atol=1e-3, + rtol=1e-3, + ) + + @register_op_test("q8ta_linear") def _q8ta_linear_suite() -> WebGPUTestSuite: # XNNPACK-static-quantized nn.Linear -> delegated quantize->q8ta_linear-> diff --git a/backends/webgpu/test/ops/test_linear_qcs4w.py b/backends/webgpu/test/ops/test_linear_qcs4w.py new file mode 100644 index 00000000000..c006660aec8 --- /dev/null +++ b/backends/webgpu/test/ops/test_linear_qcs4w.py @@ -0,0 +1,60 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`et_vk.linear_qcs4w` module + configs: 4-bit channels-symmetric weight linear. + +The op is produced by running a plain `nn.Linear` through the `VulkanQuantizer` +weight-only 4-bit path (per-output-channel symmetric weight, no activation +quant): the converted module lowers to a delegated `et_vk.linear_qcs4w`. The +`module_factory` returns the CONVERTED module, so the op-test framework goldens +the WebGPU output against the converted eager (fp32, the fake-quant reference). +Bias is out of the op's args (a `nn.Linear` bias would lower to a separate +`aten.add`); the cases use `bias=False` to keep the golden focused on qcs4w. +""" + +import unittest + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 registers et_vk ops + +import torch +import torch.nn as nn + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.vulkan.quantizer.vulkan_quantizer import ( + get_symmetric_quantization_config, + VulkanQuantizer, +) +from executorch.exir import to_edge_transform_and_lower +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + +def make_qcs4w_linear_module(k, n, m, bias=False, seed=0): + torch.manual_seed(seed) + lin = nn.Linear(k, n, bias=bias).eval() + ex = (torch.randn(m, k),) + q = VulkanQuantizer().set_global( + get_symmetric_quantization_config(is_dynamic=False, weight_bits=4) + ) + prepared = prepare_pt2e(torch.export.export(lin, ex).module(), q) + prepared(*ex) # calibrate + return convert_pt2e(prepared) + + +class LinearQcs4wTest(unittest.TestCase): + def test_delegates(self) -> None: + m = make_qcs4w_linear_module(32, 16, 4) + et = to_edge_transform_and_lower( + torch.export.export(m, (torch.randn(4, 32),)), + partitioner=[VulkanPartitioner()], + ).to_executorch() + self.assertTrue( + any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + ) + self.assertIn(b"linear_qcs4w", et.buffer) From 439fa69330f6faf0312d1b451ef0f1c5102fe9d1 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:35 -0700 Subject: [PATCH 55/75] [ExecuTorch][WebGPU] Add linear_q8ta_q8csw op (et_vk.linear_q8ta_q8csw) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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`, 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 +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct Q8taQ8cswParams { + uint32_t M; + uint32_t N; + uint32_t K; + int32_t input_zero_point; + float input_scale; + uint32_t has_bias; + uint32_t _pad[2]; +}; +static_assert( + sizeof(Q8taQ8cswParams) == 32, + "Q8taQ8cswParams must match the WGSL Params struct (32 bytes)"); + +// int8-act x int8-weight -> fp32 (no requant); Vulkan QuantizedLinear.cpp:699. +void linear_q8ta_q8csw_impl(WebGPUGraph& graph, const std::vector& args) { + // args mirror Vulkan; weight_sums (arg 4) is unused (zp subtracted inline). + const int in_id = args.at(0); + const int weight_id = args.at(3); + const int scales_id = args.at(5); + const int out_id = args.at(args.size() - 1); + const int bias_id = args.size() >= 8 ? args.at(6) : -1; + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(weight_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(scales_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error( + "linear_q8ta_q8csw: in/weight/scales/out not tensor"); + } + const bool has_bias = bias_id >= 0 && + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& weight_tensor = graph.get_tensor(weight_id); + const auto& scales_tensor = graph.get_tensor(scales_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || weight_tensor.buffer == nullptr || + scales_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("linear_q8ta_q8csw: null buffer binding"); + } + if (weight_tensor.dims.size() != 2) { + throw std::runtime_error("linear_q8ta_q8csw: weight must be 2D [N, K]"); + } + + const double input_scale = graph.get_double(args.at(1)); + const int input_zero_point = graph.get_int(args.at(2)); + + const uint64_t N = static_cast(weight_tensor.dims.at(0)); + const uint64_t K = static_cast(weight_tensor.dims.at(1)); + if (K == 0 || in_tensor.dims.empty() || + static_cast(in_tensor.dims.back()) != K) { + throw std::runtime_error("linear_q8ta_q8csw: input last dim must equal K"); + } + uint64_t in_numel = 1; + for (int64_t d : in_tensor.dims) { + in_numel *= static_cast(d); + } + const uint64_t M = in_numel / K; + if (M == 0 || N == 0 || static_cast(out_tensor.dims.back()) != N) { + throw std::runtime_error("linear_q8ta_q8csw: bad M/N shape"); + } + if (M > UINT32_MAX || N > UINT32_MAX || K > UINT32_MAX) { + throw std::runtime_error("linear_q8ta_q8csw: dim exceeds u32"); + } + if (M * K > UINT32_MAX || N * K > UINT32_MAX) { + throw std::runtime_error("linear_q8ta_q8csw: M*K or N*K exceeds u32"); + } + // int8 x/weight as array (numel%4==0); fp32 out, N free (TN=4). + if (!in_tensor.is_int8 || in_tensor.nbytes != M * K || (M * K) % 4 != 0 || + !weight_tensor.is_int8 || weight_tensor.nbytes != N * K || + (N * K) % 4 != 0) { + throw std::runtime_error("linear_q8ta_q8csw: int8 x/weight size mismatch"); + } + if (out_tensor.is_int || out_tensor.nbytes != M * N * sizeof(float)) { + throw std::runtime_error("linear_q8ta_q8csw: output must be fp32 [M, N]"); + } + if (scales_tensor.nbytes != N * sizeof(float)) { + throw std::runtime_error( + "linear_q8ta_q8csw: weight_scales must be fp32 [N]"); + } + if (has_bias) { + const auto& b = graph.get_tensor(bias_id); + if (b.buffer == nullptr || b.nbytes != N * sizeof(float)) { + throw std::runtime_error("linear_q8ta_q8csw: bias must be fp32 [N]"); + } + } + + Q8taQ8cswParams params = {}; + params.M = static_cast(M); + params.N = static_cast(N); + params.K = static_cast(K); + params.input_zero_point = static_cast(input_zero_point); + params.input_scale = static_cast(input_scale); + params.has_bias = has_bias ? 1u : 0u; + + const uint64_t n_tiles_u64 = ((M + 3) / 4) * ((N + 3) / 4); + if (n_tiles_u64 > UINT32_MAX) { + throw std::runtime_error( + "linear_q8ta_q8csw: dispatch tile count exceeds u32"); + } + const uint32_t n_tiles = static_cast(n_tiles_u64); + uint32_t wg_size = + utils::clamp_workgroup_size(device, kLinearQ8taQ8cswWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, n_tiles, wg_size, "linear_q8ta_q8csw"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(Q8taQ8cswParams)); + graph.add_uniform_buffer_bytes(sizeof(Q8taQ8cswParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kLinearQ8taQ8cswWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[6] = {}; + for (int i = 0; i < 6; i++) { + entries[i].binding = static_cast(i); + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = (i == 0) ? WGPUBufferBindingType_Storage + : (i == 5) ? WGPUBufferBindingType_Uniform + : WGPUBufferBindingType_ReadOnlyStorage; + } + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 6; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + // No-bias: bind scales as an unread placeholder (has_bias gates the read). + WGPUBuffer bias_buf = + has_bias ? graph.get_tensor(bias_id).buffer : scales_tensor.buffer; + const uint64_t bias_size = + has_bias ? graph.get_tensor(bias_id).nbytes : scales_tensor.nbytes; + + WGPUBindGroupEntry bg[6] = {}; + bg[0].binding = 0; + bg[0].buffer = out_tensor.buffer; + bg[0].size = out_tensor.nbytes; + bg[1].binding = 1; + bg[1].buffer = in_tensor.buffer; + bg[1].size = in_tensor.nbytes; + bg[2].binding = 2; + bg[2].buffer = weight_tensor.buffer; + bg[2].size = weight_tensor.nbytes; + bg[3].binding = 3; + bg[3].buffer = scales_tensor.buffer; + bg[3].size = scales_tensor.nbytes; + bg[4].binding = 4; + bg[4].buffer = bias_buf; + bg[4].size = bias_size; + bg[5].binding = 5; + bg[5].buffer = params_buf; + bg[5].size = sizeof(Q8taQ8cswParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 6; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "linear_q8ta_q8csw", + workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.linear_q8ta_q8csw.default, linear_q8ta_q8csw_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/linear_q8ta_q8csw/linear_q8ta_q8csw.wgsl b/backends/webgpu/runtime/ops/linear_q8ta_q8csw/linear_q8ta_q8csw.wgsl new file mode 100644 index 00000000000..1c2fe046a1f --- /dev/null +++ b/backends/webgpu/runtime/ops/linear_q8ta_q8csw/linear_q8ta_q8csw.wgsl @@ -0,0 +1,88 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + M: u32, + N: u32, + K: u32, + input_zero_point: i32, + input_scale: f32, + has_bias: u32, + pad0: u32, + pad1: u32, +} +@group(0) @binding(5) var params: Params; + +override wg_size: u32 = 64u; + +const TM: u32 = 4u; +const TN: u32 = 4u; +const TILE_ELEMS: u32 = TM * TN; + +fn unpack_i8(buf_idx: u32, arr_word: u32) -> i32 { + // arr_word is the value at t[buf_idx>>2]; extract the (buf_idx&3) byte. + return i32(((arr_word >> ((buf_idx & 3u) * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded tile index (lifts the 65535 1D-dispatch cap for large M*N). + let tile = gid.x + gid.y * (num_workgroups.x * wg_size); + let nct = (params.N + TN - 1u) / TN; + if (tile >= ((params.M + TM - 1u) / TM) * nct) { + return; + } + let m0 = (tile / nct) * TM; + let n0 = (tile % nct) * TN; + + var acc: array; + for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) { + acc[i] = 0; + } + + var k: u32 = 0u; + loop { + if (k >= params.K) { + break; + } + var xq: array; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m_eff = min(m0 + ml, params.M - 1u); + let bi = m_eff * params.K + k; + xq[ml] = unpack_i8(bi, t_x[bi >> 2u]) - params.input_zero_point; + } + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n_eff = min(n0 + nl, params.N - 1u); + let bi = n_eff * params.K + k; + let w = unpack_i8(bi, t_weight[bi >> 2u]); + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + acc[ml * TN + nl] = acc[ml * TN + nl] + xq[ml] * w; + } + } + k = k + 1u; + } + + // Dequantize to fp32 (no output requant); guard n= params.M) { + continue; + } + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n = n0 + nl; + if (n >= params.N) { + continue; + } + var v = f32(acc[ml * TN + nl]) * params.input_scale * t_scales[n]; + if (params.has_bias != 0u) { + v = v + t_bias[n]; + } + t_out[m * params.N + n] = v; + } + } +} diff --git a/backends/webgpu/runtime/ops/linear_q8ta_q8csw/linear_q8ta_q8csw_wgsl.h b/backends/webgpu/runtime/ops/linear_q8ta_q8csw/linear_q8ta_q8csw_wgsl.h new file mode 100644 index 00000000000..c956ae3e55c --- /dev/null +++ b/backends/webgpu/runtime/ops/linear_q8ta_q8csw/linear_q8ta_q8csw_wgsl.h @@ -0,0 +1,112 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from linear_q8ta_q8csw.wgsl - DO NOT EDIT. +// wgsl-sha256: a395e49d4581d54036637b43637b3ad6cb6bd0a90676fa4b607621ad083f5b20 +inline constexpr const char* kLinearQ8taQ8cswWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + M: u32, + N: u32, + K: u32, + input_zero_point: i32, + input_scale: f32, + has_bias: u32, + pad0: u32, + pad1: u32, +} +@group(0) @binding(5) var params: Params; + +override wg_size: u32 = 64u; + +const TM: u32 = 4u; +const TN: u32 = 4u; +const TILE_ELEMS: u32 = TM * TN; + +fn unpack_i8(buf_idx: u32, arr_word: u32) -> i32 { + // arr_word is the value at t[buf_idx>>2]; extract the (buf_idx&3) byte. + return i32(((arr_word >> ((buf_idx & 3u) * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded tile index (lifts the 65535 1D-dispatch cap for large M*N). + let tile = gid.x + gid.y * (num_workgroups.x * wg_size); + let nct = (params.N + TN - 1u) / TN; + if (tile >= ((params.M + TM - 1u) / TM) * nct) { + return; + } + let m0 = (tile / nct) * TM; + let n0 = (tile % nct) * TN; + + var acc: array; + for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) { + acc[i] = 0; + } + + var k: u32 = 0u; + loop { + if (k >= params.K) { + break; + } + var xq: array; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m_eff = min(m0 + ml, params.M - 1u); + let bi = m_eff * params.K + k; + xq[ml] = unpack_i8(bi, t_x[bi >> 2u]) - params.input_zero_point; + } + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n_eff = min(n0 + nl, params.N - 1u); + let bi = n_eff * params.K + k; + let w = unpack_i8(bi, t_weight[bi >> 2u]); + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + acc[ml * TN + nl] = acc[ml * TN + nl] + xq[ml] * w; + } + } + k = k + 1u; + } + + // Dequantize to fp32 (no output requant); guard n= params.M) { + continue; + } + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n = n0 + nl; + if (n >= params.N) { + continue; + } + var v = f32(acc[ml * TN + nl]) * params.input_scale * t_scales[n]; + if (params.has_bias != 0u) { + v = v + t_bias[n]; + } + t_out[m * params.N + n] = v; + } + } +} +)"; + +inline constexpr uint32_t kLinearQ8taQ8cswWorkgroupSizeX = 64; +inline constexpr uint32_t kLinearQ8taQ8cswWorkgroupSizeY = 1; +inline constexpr uint32_t kLinearQ8taQ8cswWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 0c3343169ebf7460b7e3152b5de3a6472642ff45 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:36 -0700 Subject: [PATCH 56/75] [ExecuTorch][WebGPU] Op-tests for linear_q8ta_q8csw Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 32 +++++++++ .../webgpu/test/ops/test_linear_q8ta_q8csw.py | 65 +++++++++++++++++++ 2 files changed, 97 insertions(+) create mode 100644 backends/webgpu/test/ops/test_linear_q8ta_q8csw.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index b1b075ecd27..4977b49e8a9 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -64,6 +64,9 @@ from executorch.backends.webgpu.test.ops.test_linear_qcs4w import ( make_qcs4w_linear_module, ) +from executorch.backends.webgpu.test.ops.test_linear_q8ta_q8csw import ( + make_linear_q8ta_q8csw_module, +) from executorch.backends.webgpu.test.ops.test_q8ta_linear import ( make_q8ta_linear_module, ) @@ -1571,6 +1574,35 @@ def case(name, m, k, n, **kw): ) +@register_op_test("linear_q8ta_q8csw") +def _linear_q8ta_q8csw_suite() -> WebGPUTestSuite: + # XNNPACK-static nn.Linear with output_activation=None -> delegated + # quantize_per_tensor -> linear_q8ta_q8csw (fp32 out). Golden = converted + # eager (fp32 fake-quant). All cases use bias=True: a bias-less TERMINAL + # linear mis-fuses to an int8 output the schema (no output scale/zp) cannot + # compute -> the handler fail-louds on it; bias keeps the output fp32. N is a + # multiple of 4 (the AOT pads the quantized weight's N to a mult of 4). + def case(name, m, k, n, **kw): + return Case( + name=name, + construct={"k": k, "n": n, "m": m, "bias": True, **kw}, + inputs=((m, k),), + ) + + return WebGPUTestSuite( + module_factory=lambda **kw: make_linear_q8ta_q8csw_module(**kw), + cases=[ + case("basic", 4, 32, 16), + case("gemv", 1, 32, 16), # M==1 + case("k48", 2, 48, 8), # different K, smaller N + case("n32", 3, 32, 32), # larger N + ], + golden_dtype="float32", + atol=1e-3, + rtol=1e-3, + ) + + @register_op_test("q8ta_linear") def _q8ta_linear_suite() -> WebGPUTestSuite: # XNNPACK-static-quantized nn.Linear -> delegated quantize->q8ta_linear-> diff --git a/backends/webgpu/test/ops/test_linear_q8ta_q8csw.py b/backends/webgpu/test/ops/test_linear_q8ta_q8csw.py new file mode 100644 index 00000000000..3ed853fbe91 --- /dev/null +++ b/backends/webgpu/test/ops/test_linear_q8ta_q8csw.py @@ -0,0 +1,65 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`et_vk.linear_q8ta_q8csw` module + configs: int8-act x int8-channelwise-weight +linear with FP32 output (no output requant). + +Reached by running a plain `nn.Linear` through XNNPACK static PT2E with the +activation config's `output_activation` nulled (`dataclasses.replace(..., 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 `module_factory` returns the CONVERTED +module, so the op-test framework goldens the WebGPU output against the converted +eager (fp32 fake-quant reference); the served subgraph is `quantize_per_tensor` +(landed C0) -> `linear_q8ta_q8csw`. +""" + +import dataclasses +import unittest + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 registers et_vk ops + +import torch +import torch.nn as nn + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( + get_symmetric_quantization_config, + XNNPACKQuantizer, +) +from executorch.exir import to_edge_transform_and_lower +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + +def make_linear_q8ta_q8csw_module(k, n, m, bias=True, seed=0): + torch.manual_seed(seed) + lin = nn.Linear(k, n, bias=bias).eval() + ex = (torch.randn(m, k),) + cfg = dataclasses.replace( + get_symmetric_quantization_config(is_per_channel=True, is_dynamic=False), + output_activation=None, + ) + q = XNNPACKQuantizer().set_global(cfg) + prepared = prepare_pt2e(torch.export.export(lin, ex).module(), q) + prepared(*ex) # calibrate + return convert_pt2e(prepared) + + +class LinearQ8taQ8cswTest(unittest.TestCase): + def test_delegates(self) -> None: + m = make_linear_q8ta_q8csw_module(32, 16, 4, bias=True) + et = to_edge_transform_and_lower( + torch.export.export(m, (torch.randn(4, 32),)), + partitioner=[VulkanPartitioner()], + ).to_executorch() + self.assertTrue( + any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + ) + self.assertIn(b"linear_q8ta_q8csw", et.buffer) From ea80ee2f15210f563386ad2096dfdeabfd3cbf3e Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:36 -0700 Subject: [PATCH 57/75] [ExecuTorch][WebGPU] Add grid_priors op (et_vk.grid_priors) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../runtime/ops/grid_priors/GridPriors.cpp | 189 ++++++++++++++++++ .../runtime/ops/grid_priors/grid_priors.wgsl | 31 +++ .../ops/grid_priors/grid_priors_wgsl.h | 55 +++++ 3 files changed, 275 insertions(+) create mode 100644 backends/webgpu/runtime/ops/grid_priors/GridPriors.cpp create mode 100644 backends/webgpu/runtime/ops/grid_priors/grid_priors.wgsl create mode 100644 backends/webgpu/runtime/ops/grid_priors/grid_priors_wgsl.h diff --git a/backends/webgpu/runtime/ops/grid_priors/GridPriors.cpp b/backends/webgpu/runtime/ops/grid_priors/GridPriors.cpp new file mode 100644 index 00000000000..03ab5927001 --- /dev/null +++ b/backends/webgpu/runtime/ops/grid_priors/GridPriors.cpp @@ -0,0 +1,189 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform layout matching the WGSL Params struct; 16-byte aligned. +struct GridPriorsParams { + uint32_t numel; + uint32_t width; + float stride; + float offset; +}; +static_assert(sizeof(GridPriorsParams) == 16, "GridPriorsParams must be 16 B"); + +// grid_priors: anchor grid-shifts [H*W,2] from H/W (Vulkan GridPriors.cpp). +void grid_priors_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [in, stride(int), offset(float), out]; only in's H/W are used. + const int in_id = args.at(0); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("grid_priors: in/out is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (out_tensor.buffer == nullptr) { + throw std::runtime_error("grid_priors: null output buffer"); + } + if (in_tensor.dims.size() < 2) { + throw std::runtime_error("grid_priors: input must be at least 2D"); + } + + const int64_t stride = graph.get_int(args.at(1)); + const double offset = graph.get_double(args.at(2)); + const uint64_t height = + static_cast(in_tensor.dims.at(in_tensor.dims.size() - 2)); + const uint64_t width = + static_cast(in_tensor.dims.at(in_tensor.dims.size() - 1)); + const uint64_t numel = height * width * 2u; + if (width == 0u || numel == 0u || numel > UINT32_MAX) { + throw std::runtime_error("grid_priors: bad H/W (zero or numel > u32)"); + } + // Output is fp32 [H*W, 2]. + if (out_tensor.is_int || out_tensor.nbytes != numel * sizeof(float)) { + throw std::runtime_error("grid_priors: output must be fp32 [H*W, 2]"); + } + + GridPriorsParams params = {}; + params.numel = static_cast(numel); + params.width = static_cast(width); + params.stride = static_cast(stride); + params.offset = static_cast(offset); + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kGridPriorsWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, static_cast(numel), wg_size, "grid_priors"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(GridPriorsParams)); + graph.add_uniform_buffer_bytes(sizeof(GridPriorsParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kGridPriorsWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // Bind group: output (rw storage) + params (uniform); input data is unread. + WGPUBindGroupLayoutEntry entries[2] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_Storage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 2; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg[2] = {}; + bg[0].binding = 0; + bg[0].buffer = out_tensor.buffer; + bg[0].size = out_tensor.nbytes; + bg[1].binding = 1; + bg[1].buffer = uniform_buffer; + bg[1].size = sizeof(GridPriorsParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 2; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "grid_priors", + workgroup_count.y}); + + // Dynamic shapes: recompute numel/width + dispatch + out dims [H*W, 2]. + WGPUBuffer params_buf = uniform_buffer; + const float stride_f = static_cast(stride); + const float offset_f = static_cast(offset); + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, stride_f, offset_f, wg_size, dispatch_idx, params_buf]( + WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + if (d.size() < 2) { + throw std::runtime_error("grid_priors(resize): input < 2D"); + } + const uint64_t h = static_cast(d.at(d.size() - 2)); + const uint64_t w = static_cast(d.at(d.size() - 1)); + const uint64_t n = h * w * 2u; + if (w == 0u || n == 0u || n > UINT32_MAX) { + throw std::runtime_error("grid_priors(resize): bad H/W"); + } + GridPriorsParams p = {}; + p.numel = static_cast(n); + p.width = static_cast(w); + p.stride = stride_f; + p.offset = offset_f; + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), static_cast(n), wg_size, "grid_priors"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + g.set_cur_dims( + out_id, {static_cast(h * w), static_cast(2)}); + }); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.grid_priors.default, grid_priors_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/grid_priors/grid_priors.wgsl b/backends/webgpu/runtime/ops/grid_priors/grid_priors.wgsl new file mode 100644 index 00000000000..9046c3e0e10 --- /dev/null +++ b/backends/webgpu/runtime/ops/grid_priors/grid_priors.wgsl @@ -0,0 +1,31 @@ +@group(0) @binding(0) var t_out: array; + +struct Params { + numel: u32, + width: u32, + stride: f32, + offset: f32, +} +@group(0) @binding(1) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + // out[r,0]=(r%W+offset)*stride; out[r,1]=(r/W+offset)*stride (r=row=idx/2). + let row = idx / 2u; + var coord: u32; + if ((idx & 1u) == 0u) { + coord = row % params.width; + } else { + coord = row / params.width; + } + t_out[idx] = (f32(coord) + params.offset) * params.stride; +} diff --git a/backends/webgpu/runtime/ops/grid_priors/grid_priors_wgsl.h b/backends/webgpu/runtime/ops/grid_priors/grid_priors_wgsl.h new file mode 100644 index 00000000000..b6ee41a1b5e --- /dev/null +++ b/backends/webgpu/runtime/ops/grid_priors/grid_priors_wgsl.h @@ -0,0 +1,55 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from grid_priors.wgsl - DO NOT EDIT. +// wgsl-sha256: d9ded332b2d5a5c9aae62a13ddbc011c47f67731c3221ea191f64bba0a887e42 +inline constexpr const char* kGridPriorsWGSL = R"( +@group(0) @binding(0) var t_out: array; + +struct Params { + numel: u32, + width: u32, + stride: f32, + offset: f32, +} +@group(0) @binding(1) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + // out[r,0]=(r%W+offset)*stride; out[r,1]=(r/W+offset)*stride (r=row=idx/2). + let row = idx / 2u; + var coord: u32; + if ((idx & 1u) == 0u) { + coord = row % params.width; + } else { + coord = row / params.width; + } + t_out[idx] = (f32(coord) + params.offset) * params.stride; +} +)"; + +inline constexpr uint32_t kGridPriorsWorkgroupSizeX = 64; +inline constexpr uint32_t kGridPriorsWorkgroupSizeY = 1; +inline constexpr uint32_t kGridPriorsWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 9e3273eece3012d95983de9875017d982cc4ab80 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:37 -0700 Subject: [PATCH 58/75] [ExecuTorch][WebGPU] Op-tests for grid_priors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 24 ++++++ backends/webgpu/test/ops/test_grid_priors.py | 80 ++++++++++++++++++++ 2 files changed, 104 insertions(+) create mode 100644 backends/webgpu/test/ops/test_grid_priors.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 4977b49e8a9..7a277d5f2ea 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -46,6 +46,7 @@ from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule from executorch.backends.webgpu.test.ops.test_conv1d_dw import Conv1dDWModule from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dPwModule +from executorch.backends.webgpu.test.ops.test_grid_priors import GridPriorsModule from executorch.backends.webgpu.test.ops.test_grid_sampler_2d import ( GridSampler2dModule, ) @@ -612,6 +613,29 @@ def _index_select_suite() -> WebGPUTestSuite: ) +@register_op_test("grid_priors") +def _grid_priors_suite() -> WebGPUTestSuite: + # Detection anchor-grid op: out [H*W, 2] from the input's H/W (values unused) + # + baked stride/offset. Pure computed output -> float32 oracle. `offset0` + # covers offset=0; the shapes span square + non-square H/W. + def case(name, shape, stride, offset): + return Case( + name=name, + construct={"stride": stride, "offset": offset}, + inputs=(shape,), + ) + + return WebGPUTestSuite( + module_factory=lambda stride, offset: GridPriorsModule(stride, offset), + cases=[ + case("s8", (1, 3, 8, 10), 8, 0.5), + case("s16", (1, 3, 4, 4), 16, 0.0), + case("offset0", (1, 3, 5, 7), 4, 0.0), + ], + golden_dtype="float32", + ) + + @register_op_test("view_copy") def _view_copy_suite() -> WebGPUTestSuite: return _fn_config_suite(ViewModule, _VIEW_CONFIGS) diff --git a/backends/webgpu/test/ops/test_grid_priors.py b/backends/webgpu/test/ops/test_grid_priors.py new file mode 100644 index 00000000000..a5e7597fe8e --- /dev/null +++ b/backends/webgpu/test/ops/test_grid_priors.py @@ -0,0 +1,80 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`et_vk.grid_priors` module + configs for the WebGPU op-test framework. + +`GridPriorsModule` calls the custom op directly (detection anchor-grid op, no +aten lowering); `stride`/`offset` are baked construct kwargs and only the float +tensor `x` is a runtime input (its VALUES are unused — only its H/W set the +output shape [H*W, 2]). The op has a CPU eager impl, so the framework goldens it +directly. `GridPriorsTest` is the export-delegation smoke test. +""" + +import unittest + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# name -> (input_shape, stride, offset) +CONFIGS = { + "s8": ((1, 3, 8, 10), 8, 0.5), + "s16": ((1, 3, 4, 4), 16, 0.0), + "offset0": ((1, 3, 5, 7), 4, 0.0), +} + + +class GridPriorsModule(torch.nn.Module): + def __init__(self, stride, offset) -> None: + super().__init__() + self.stride = stride + self.offset = offset + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.et_vk.grid_priors.default(x, self.stride, self.offset) + + +def _det(shape): + g = torch.Generator().manual_seed(0) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(shape, stride, offset): + ep = torch.export.export(GridPriorsModule(stride, offset).eval(), (_det(shape),)) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + +def _delegated(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +def _op_delegated(edge, op_substr: str) -> bool: + # op must be absorbed into the delegate, not left as a top-level CPU-fallback node. + gm = edge.exported_program().graph_module + return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + + +class GridPriorsTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, (shape, stride, offset) in CONFIGS.items(): + with self.subTest(name=name): + edge = _lower(shape, stride, offset) + et = edge.to_executorch() + self.assertTrue( + _delegated(et), + f"Expected a VulkanBackend delegate (grid_priors {name})", + ) + self.assertTrue( + _op_delegated(edge, "grid_priors"), + f"grid_priors not delegated (CPU fallback) for {name}", + ) From 81fabd2a9e13ed9f77c3caa77baedbb8ccc7c97f Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:37 -0700 Subject: [PATCH 59/75] [ExecuTorch][WebGPU] Add conv_with_clamp op (et_vk.conv_with_clamp) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../ops/conv_with_clamp/ConvWithClamp.cpp | 302 ++++++++++++++++++ .../ops/conv_with_clamp/conv_with_clamp.wgsl | 85 +++++ .../conv_with_clamp/conv_with_clamp_wgsl.h | 109 +++++++ 3 files changed, 496 insertions(+) create mode 100644 backends/webgpu/runtime/ops/conv_with_clamp/ConvWithClamp.cpp create mode 100644 backends/webgpu/runtime/ops/conv_with_clamp/conv_with_clamp.wgsl create mode 100644 backends/webgpu/runtime/ops/conv_with_clamp/conv_with_clamp_wgsl.h diff --git a/backends/webgpu/runtime/ops/conv_with_clamp/ConvWithClamp.cpp b/backends/webgpu/runtime/ops/conv_with_clamp/ConvWithClamp.cpp new file mode 100644 index 00000000000..42f989009f4 --- /dev/null +++ b/backends/webgpu/runtime/ops/conv_with_clamp/ConvWithClamp.cpp @@ -0,0 +1,302 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct ConvWithClampParams { + uint32_t N; + uint32_t IC; + uint32_t H_in; + uint32_t W_in; + uint32_t OC; + uint32_t H_out; + uint32_t W_out; + uint32_t Kh; + uint32_t Kw; + uint32_t stride_h; + uint32_t stride_w; + uint32_t pad_h; + uint32_t pad_w; + uint32_t dil_h; + uint32_t dil_w; + uint32_t has_bias; + uint32_t numel; + uint32_t groups; + uint32_t ic_per_group; + uint32_t pad0; + uint32_t pad1; + uint32_t pad2; + float output_min; + float output_max; +}; +static_assert( + sizeof(ConvWithClampParams) == 96, + "ConvWithClampParams must match the WGSL Params struct (96 bytes)"); + +std::pair pair_or_throw( + const std::vector& v, + const char* msg) { + if (v.size() != 2) { + throw std::runtime_error(msg); + } + return {v.at(0), v.at(1)}; +} + +// A Scalar? clamp bound: Double/Int -> value; Null (absent) -> the default. +float scalar_or(WebGPUGraph& graph, int id, float dflt) { + const auto t = graph.get_value_type(id); + if (t == WebGPUGraph::ValueType::Double) { + return static_cast(graph.get_double(id)); + } + if (t == WebGPUGraph::ValueType::Int) { + return static_cast(graph.get_int(id)); + } + if (t == WebGPUGraph::ValueType::Null) { + return dflt; + } + throw std::runtime_error("conv_with_clamp: unexpected clamp bound type"); +} + +// fp32 general conv2d (groups==1) + clamp; mirrors Vulkan conv_with_clamp. +void conv_with_clamp_impl(WebGPUGraph& graph, const std::vector& args) { + // args mirror Vulkan convolution + clamp bounds; out=args.back(). + const int in_id = args.at(0); + const int weight_id = args.at(1); + const int bias_id = args.at(2); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(weight_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("conv_with_clamp: in/weight/out not tensor"); + } + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& weight_tensor = graph.get_tensor(weight_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || weight_tensor.buffer == nullptr || + out_tensor.buffer == nullptr) { + throw std::runtime_error("conv_with_clamp: null buffer binding"); + } + if (in_tensor.dims.size() != 4 || out_tensor.dims.size() != 4 || + weight_tensor.dims.size() != 4) { + throw std::runtime_error("conv_with_clamp: in/out/weight must be 4D"); + } + if (in_tensor.is_int || weight_tensor.is_int || out_tensor.is_int) { + throw std::runtime_error("conv_with_clamp: fp32 only"); + } + + const auto [stride_h, stride_w] = + pair_or_throw(graph.get_int_list(args.at(3)), "conv_with_clamp: stride"); + const auto [pad_h, pad_w] = + pair_or_throw(graph.get_int_list(args.at(4)), "conv_with_clamp: padding"); + const auto [dil_h, dil_w] = pair_or_throw( + graph.get_int_list(args.at(5)), "conv_with_clamp: dilation"); + if (graph.get_value_type(args.at(6)) != WebGPUGraph::ValueType::Bool) { + throw std::runtime_error("conv_with_clamp: transposed must be bool"); + } + if (graph.get_bool(args.at(6))) { + throw std::runtime_error("conv_with_clamp: transposed unsupported"); + } + const int64_t groups = graph.get_int(args.at(8)); + if (groups < 1) { + throw std::runtime_error("conv_with_clamp: groups must be >= 1"); + } + const float kInf = std::numeric_limits::infinity(); + const float output_min = scalar_or(graph, args.at(9), -kInf); + const float output_max = scalar_or(graph, args.at(10), kInf); + + const uint64_t N = static_cast(in_tensor.dims.at(0)); + const uint64_t IC = static_cast(in_tensor.dims.at(1)); + const uint64_t H_in = static_cast(in_tensor.dims.at(2)); + const uint64_t W_in = static_cast(in_tensor.dims.at(3)); + const uint64_t OC = static_cast(weight_tensor.dims.at(0)); + const uint64_t Kh = static_cast(weight_tensor.dims.at(2)); + const uint64_t Kw = static_cast(weight_tensor.dims.at(3)); + const uint64_t H_out = static_cast(out_tensor.dims.at(2)); + const uint64_t W_out = static_cast(out_tensor.dims.at(3)); + // Declared output spatial dims must match the conv2d formula result. + if (stride_h <= 0 || stride_w <= 0) { + throw std::runtime_error("conv_with_clamp: stride must be positive"); + } + const int64_t h_eff = static_cast(H_in) + 2 * pad_h - + dil_h * (static_cast(Kh) - 1) - 1; + const int64_t w_eff = static_cast(W_in) + 2 * pad_w - + dil_w * (static_cast(Kw) - 1) - 1; + if (h_eff < 0 || w_eff < 0 || + static_cast(h_eff / stride_h + 1) != H_out || + static_cast(w_eff / stride_w + 1) != W_out) { + throw std::runtime_error( + "conv_with_clamp: output dims inconsistent with conv2d formula"); + } + const uint64_t numel = N * OC * H_out * W_out; + const uint64_t ug = static_cast(groups); + // Grouped weight is [OC, IC/groups, Kh, Kw]; ic_per_group = weight.dims[1]. + const uint64_t ic_per_group = + static_cast(weight_tensor.dims.at(1)); + if (IC == 0 || numel == 0 || numel > UINT32_MAX || IC % ug != 0 || + OC % ug != 0 || ic_per_group * ug != IC) { + throw std::runtime_error("conv_with_clamp: bad shape (IC/numel/groups)"); + } + if (out_tensor.nbytes != numel * sizeof(float) || + in_tensor.nbytes != N * IC * H_in * W_in * sizeof(float) || + weight_tensor.nbytes != OC * ic_per_group * Kh * Kw * sizeof(float)) { + throw std::runtime_error("conv_with_clamp: fp32 byte-size mismatch"); + } + if (has_bias) { + const auto& b = graph.get_tensor(bias_id); + if (b.buffer == nullptr || b.nbytes != OC * sizeof(float)) { + throw std::runtime_error("conv_with_clamp: bias must be fp32 [OC]"); + } + } + + ConvWithClampParams params = {}; + params.N = static_cast(N); + params.IC = static_cast(IC); + params.H_in = static_cast(H_in); + params.W_in = static_cast(W_in); + params.OC = static_cast(OC); + params.H_out = static_cast(H_out); + params.W_out = static_cast(W_out); + params.Kh = static_cast(Kh); + params.Kw = static_cast(Kw); + params.stride_h = static_cast(stride_h); + params.stride_w = static_cast(stride_w); + params.pad_h = static_cast(pad_h); + params.pad_w = static_cast(pad_w); + params.dil_h = static_cast(dil_h); + params.dil_w = static_cast(dil_w); + params.has_bias = has_bias ? 1u : 0u; + params.numel = static_cast(numel); + params.groups = static_cast(groups); + params.ic_per_group = static_cast(ic_per_group); + params.output_min = output_min; + params.output_max = output_max; + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kConvWithClampWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, static_cast(numel), wg_size, "conv_with_clamp"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(ConvWithClampParams)); + graph.add_uniform_buffer_bytes(sizeof(ConvWithClampParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kConvWithClampWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // out (rw) + in/weight/bias (ro storage) + params (uniform). + WGPUBindGroupLayoutEntry entries[5] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_Storage; + for (uint32_t i = 1; i <= 3; i++) { + entries[i].binding = i; + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + } + entries[4].binding = 4; + entries[4].visibility = WGPUShaderStage_Compute; + entries[4].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 5; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + // No-bias: bind weight as an unread placeholder (has_bias gates the read). + WGPUBuffer bias_buf = + has_bias ? graph.get_tensor(bias_id).buffer : weight_tensor.buffer; + const uint64_t bias_size = + has_bias ? graph.get_tensor(bias_id).nbytes : weight_tensor.nbytes; + + WGPUBindGroupEntry bg[5] = {}; + bg[0].binding = 0; + bg[0].buffer = out_tensor.buffer; + bg[0].size = out_tensor.nbytes; + bg[1].binding = 1; + bg[1].buffer = in_tensor.buffer; + bg[1].size = in_tensor.nbytes; + bg[2].binding = 2; + bg[2].buffer = weight_tensor.buffer; + bg[2].size = weight_tensor.nbytes; + bg[3].binding = 3; + bg[3].buffer = bias_buf; + bg[3].size = bias_size; + bg[4].binding = 4; + bg[4].buffer = params_buf; + bg[4].size = sizeof(ConvWithClampParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 5; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "conv_with_clamp", + workgroup_count.y}); + + // conv2d is static-shape-only: no tensor resize hook is registered, so the + // output spatial dims stay fixed at their build-time (serialized) values. + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.conv_with_clamp.default, conv_with_clamp_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/conv_with_clamp/conv_with_clamp.wgsl b/backends/webgpu/runtime/ops/conv_with_clamp/conv_with_clamp.wgsl new file mode 100644 index 00000000000..f5bff8699fa --- /dev/null +++ b/backends/webgpu/runtime/ops/conv_with_clamp/conv_with_clamp.wgsl @@ -0,0 +1,85 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + // Unravel the NCHW output index -> (n, oc, oh, ow). + let ow = idx % params.W_out; + var r = idx / params.W_out; + let oh = r % params.H_out; + r = r / params.H_out; + let oc = r % params.OC; + let n = r / params.OC; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = t_bias[oc]; + } + // Grouped conv: output channel oc belongs to group g and dots only that + // group's ic_per_group input channels. weight is [OC, IC/groups, Kh, Kw]; + // groups==1 (ic_per_group==IC, g==0) is the general dense case. + let oc_per_group = params.OC / params.groups; + let g = oc / oc_per_group; + let ic_base = g * params.ic_per_group; + for (var ic_local: u32 = 0u; ic_local < params.ic_per_group; ic_local = ic_local + 1u) { + let ic = ic_base + ic_local; + for (var kh: u32 = 0u; kh < params.Kh; kh = kh + 1u) { + let ih = i32(oh) * i32(params.stride_h) - i32(params.pad_h) + + i32(kh) * i32(params.dil_h); + if (ih < 0 || ih >= i32(params.H_in)) { + continue; + } + let in_row = ((n * params.IC + ic) * params.H_in + u32(ih)) * params.W_in; + let w_row = + ((oc * params.ic_per_group + ic_local) * params.Kh + kh) * params.Kw; + for (var kw: u32 = 0u; kw < params.Kw; kw = kw + 1u) { + let iw = i32(ow) * i32(params.stride_w) - i32(params.pad_w) + + i32(kw) * i32(params.dil_w); + if (iw < 0 || iw >= i32(params.W_in)) { + continue; + } + acc = acc + t_x[in_row + u32(iw)] * t_weight[w_row + kw]; + } + } + } + t_out[idx] = clamp(acc, params.output_min, params.output_max); +} diff --git a/backends/webgpu/runtime/ops/conv_with_clamp/conv_with_clamp_wgsl.h b/backends/webgpu/runtime/ops/conv_with_clamp/conv_with_clamp_wgsl.h new file mode 100644 index 00000000000..059e82c5873 --- /dev/null +++ b/backends/webgpu/runtime/ops/conv_with_clamp/conv_with_clamp_wgsl.h @@ -0,0 +1,109 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from conv_with_clamp.wgsl - DO NOT EDIT. +// wgsl-sha256: 262acfcd0b8ea1743d618c0d21dc140563da3ebf98c0399507f1005a16d64b5e +inline constexpr const char* kConvWithClampWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + // Unravel the NCHW output index -> (n, oc, oh, ow). + let ow = idx % params.W_out; + var r = idx / params.W_out; + let oh = r % params.H_out; + r = r / params.H_out; + let oc = r % params.OC; + let n = r / params.OC; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = t_bias[oc]; + } + // Grouped conv: output channel oc belongs to group g and dots only that + // group's ic_per_group input channels. weight is [OC, IC/groups, Kh, Kw]; + // groups==1 (ic_per_group==IC, g==0) is the general dense case. + let oc_per_group = params.OC / params.groups; + let g = oc / oc_per_group; + let ic_base = g * params.ic_per_group; + for (var ic_local: u32 = 0u; ic_local < params.ic_per_group; ic_local = ic_local + 1u) { + let ic = ic_base + ic_local; + for (var kh: u32 = 0u; kh < params.Kh; kh = kh + 1u) { + let ih = i32(oh) * i32(params.stride_h) - i32(params.pad_h) + + i32(kh) * i32(params.dil_h); + if (ih < 0 || ih >= i32(params.H_in)) { + continue; + } + let in_row = ((n * params.IC + ic) * params.H_in + u32(ih)) * params.W_in; + let w_row = + ((oc * params.ic_per_group + ic_local) * params.Kh + kh) * params.Kw; + for (var kw: u32 = 0u; kw < params.Kw; kw = kw + 1u) { + let iw = i32(ow) * i32(params.stride_w) - i32(params.pad_w) + + i32(kw) * i32(params.dil_w); + if (iw < 0 || iw >= i32(params.W_in)) { + continue; + } + acc = acc + t_x[in_row + u32(iw)] * t_weight[w_row + kw]; + } + } + } + t_out[idx] = clamp(acc, params.output_min, params.output_max); +} +)"; + +inline constexpr uint32_t kConvWithClampWorkgroupSizeX = 64; +inline constexpr uint32_t kConvWithClampWorkgroupSizeY = 1; +inline constexpr uint32_t kConvWithClampWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 6d245a789d92dcb64ab1a08299ba1ac09b37fc98 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:37 -0700 Subject: [PATCH 60/75] [ExecuTorch][WebGPU] Op-tests for conv_with_clamp MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 39 ++++++++ .../webgpu/test/ops/test_conv_with_clamp.py | 95 +++++++++++++++++++ 2 files changed, 134 insertions(+) create mode 100644 backends/webgpu/test/ops/test_conv_with_clamp.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 7a277d5f2ea..03dc628ffae 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -47,6 +47,9 @@ from executorch.backends.webgpu.test.ops.test_conv1d_dw import Conv1dDWModule from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dPwModule from executorch.backends.webgpu.test.ops.test_grid_priors import GridPriorsModule +from executorch.backends.webgpu.test.ops.test_conv_with_clamp import ( + ConvWithClampModule, +) from executorch.backends.webgpu.test.ops.test_grid_sampler_2d import ( GridSampler2dModule, ) @@ -613,6 +616,42 @@ def _index_select_suite() -> WebGPUTestSuite: ) +@register_op_test("conv_with_clamp") +def _conv_with_clamp_suite() -> WebGPUTestSuite: + # nn.Conv2d + F.relu6 -> delegated et_vk.conv_with_clamp (fp32 conv + clamp + # [0,6]). Golden = fp32 eager. Covers k3/stride/dilation/no_bias (groups==1). + def case(name, shape, ic, oc, k, stride, padding, dilation, bias=True): + return Case( + name=name, + construct={ + "ic": ic, + "oc": oc, + "k": k, + "stride": stride, + "padding": padding, + "dilation": dilation, + "bias": bias, + }, + inputs=(shape,), + ) + + return WebGPUTestSuite( + module_factory=lambda **kw: ConvWithClampModule(**kw), + cases=[ + case("k3p1", (1, 4, 8, 8), 4, 8, 3, 1, 1, 1), + case("stride2", (1, 3, 10, 10), 3, 6, 3, 2, 1, 1), + case("dil2", (2, 3, 9, 9), 3, 5, 3, 1, 2, 2), + case("no_bias", (1, 4, 8, 8), 4, 8, 3, 1, 1, 1, bias=False), + # Fully axis-asymmetric (Kh!=Kw, H!=W, sh!=sw, ph!=pw, dh!=dw) so an + # H<->W index swap would diverge from the golden. + case("asym", (1, 3, 7, 9), 3, 5, (2, 3), (1, 2), (1, 0), (2, 1)), + ], + golden_dtype="float32", + atol=1e-3, + rtol=1e-3, + ) + + @register_op_test("grid_priors") def _grid_priors_suite() -> WebGPUTestSuite: # Detection anchor-grid op: out [H*W, 2] from the input's H/W (values unused) diff --git a/backends/webgpu/test/ops/test_conv_with_clamp.py b/backends/webgpu/test/ops/test_conv_with_clamp.py new file mode 100644 index 00000000000..1ca4d9def6b --- /dev/null +++ b/backends/webgpu/test/ops/test_conv_with_clamp.py @@ -0,0 +1,95 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`et_vk.conv_with_clamp` module + configs for the WebGPU op-test framework. + +`ConvWithClampModule` is a plain `nn.Conv2d` followed by `F.relu6`, which the +Vulkan fusion rewrites to a delegated `et_vk.conv_with_clamp` (fp32 general conv ++ output clamp[0,6]). The conv weight/bias are baked params; only the float +tensor `x` is a runtime input. Goldened vs the module's fp32 eager. +`ConvWithClampTest` is the export-delegation smoke test. +""" + +import unittest + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 + +import torch +import torch.nn as nn + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# name -> (input_shape, ic, oc, k, stride, padding, dilation, bias) +CONFIGS = { + "k3p1": ((1, 4, 8, 8), 4, 8, 3, 1, 1, 1, True), + "stride2": ((1, 3, 10, 10), 3, 6, 3, 2, 1, 1, True), + "dil2": ((2, 3, 9, 9), 3, 5, 3, 1, 2, 2, True), + "no_bias": ((1, 4, 8, 8), 4, 8, 3, 1, 1, 1, False), + "asym": ((1, 3, 7, 9), 3, 5, (2, 3), (1, 2), (1, 0), (2, 1), True), +} + + +class ConvWithClampModule(nn.Module): + def __init__(self, ic, oc, k, stride, padding, dilation, bias) -> None: + super().__init__() + self.c = nn.Conv2d( + ic, oc, k, stride=stride, padding=padding, dilation=dilation, bias=bias + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.relu6(self.c(x)) + + +def _det(shape): + g = torch.Generator().manual_seed(0) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(shape, ic, oc, k, stride, padding, dilation, bias): + mod = ConvWithClampModule(ic, oc, k, stride, padding, dilation, bias).eval() + ep = torch.export.export(mod, (_det(shape),)) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + +def _delegated(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +def _op_delegated(edge, op_substr: str) -> bool: + # The op must be absorbed into a delegate: absent from the top-level graph AND + # present inside a lowered submodule reached by an executorch_call_delegate node + # (a bare absence check also passes for an empty graph or a renamed op). + from executorch.exir.lowered_backend_module import get_lowered_submodules + + gm = edge.exported_program().graph_module + if any(op_substr in str(getattr(n, "target", "")) for n in gm.graph.nodes): + return False + return any( + op_substr in str(getattr(dn, "target", "")) + for _, lowered, _ in get_lowered_submodules(gm) + for dn in lowered.original_module.graph_module.graph.nodes + ) + + +class ConvWithClampTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, (shape, ic, oc, k, s, p, d, b) in CONFIGS.items(): + with self.subTest(name=name): + edge = _lower(shape, ic, oc, k, s, p, d, b) + et = edge.to_executorch() + self.assertTrue( + _delegated(et), + f"Expected a VulkanBackend delegate (conv_with_clamp {name})", + ) + self.assertTrue( + _op_delegated(edge, "conv_with_clamp"), + f"conv_with_clamp not delegated (CPU fallback) for {name}", + ) From a4e332b58375bfb8587c7450b9597f76fc7267ce Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:38 -0700 Subject: [PATCH 61/75] [ExecuTorch][WebGPU] Add comparison ops (aten.eq/lt/le/gt/ge.Tensor -> bool) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../webgpu/runtime/ops/compare/Compare.cpp | 212 ++++++++++++++++++ .../webgpu/runtime/ops/compare/compare.wgsl | 43 ++++ .../webgpu/runtime/ops/compare/compare_wgsl.h | 67 ++++++ 3 files changed, 322 insertions(+) create mode 100644 backends/webgpu/runtime/ops/compare/Compare.cpp create mode 100644 backends/webgpu/runtime/ops/compare/compare.wgsl create mode 100644 backends/webgpu/runtime/ops/compare/compare_wgsl.h diff --git a/backends/webgpu/runtime/ops/compare/Compare.cpp b/backends/webgpu/runtime/ops/compare/Compare.cpp new file mode 100644 index 00000000000..8ecd35a44ff --- /dev/null +++ b/backends/webgpu/runtime/ops/compare/Compare.cpp @@ -0,0 +1,212 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform layout matching the WGSL Params struct; 16-byte aligned. +struct CompareParams { + uint32_t num_elements; + uint32_t op; + uint32_t _pad[2]; +}; +static_assert(sizeof(CompareParams) == 16, "CompareParams must be 16 bytes"); + +// Elementwise fp32 compare -> bool (op 0=eq..4=ge); eq is torch-exact a==b. +void compare_impl( + WebGPUGraph& graph, + const std::vector& args, + uint32_t op) { + // args: [in1, in2, out]; fp32 in, bool out; out=args.back(). + const int in1_id = args.at(0); + const int in2_id = args.at(1); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in1_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(in2_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("compare: in1/in2/out is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& in1_tensor = graph.get_tensor(in1_id); + const auto& in2_tensor = graph.get_tensor(in2_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in1_tensor.buffer == nullptr || in2_tensor.buffer == nullptr || + out_tensor.buffer == nullptr) { + throw std::runtime_error("compare: null buffer binding"); + } + // fp32 inputs; bool output (1-byte int-typed, NOT int8-quantized). + if (in1_tensor.is_int || in2_tensor.is_int || in1_tensor.elem_size != 4 || + in2_tensor.elem_size != 4) { + throw std::runtime_error("compare: fp32 inputs only"); + } + if (!out_tensor.is_int || out_tensor.elem_size != 1) { + throw std::runtime_error("compare: output must be a 1-byte bool tensor"); + } + const uint64_t numel = out_tensor.nbytes; + // out bool packed 4/word (array); numel%4==0 gates the readback map. + if (numel == 0u || numel % 4u != 0u || numel > UINT32_MAX) { + throw std::runtime_error("compare: numel must be a nonzero mult of 4"); + } + const uint64_t in_numel = in1_tensor.nbytes / sizeof(float); + if (in1_tensor.nbytes != in2_tensor.nbytes || in_numel != numel) { + throw std::runtime_error("compare: in/out numel mismatch (same-shape)"); + } + + CompareParams params = {}; + params.num_elements = static_cast(numel); + params.op = op; + + const uint32_t words = static_cast(numel / 4u); + uint32_t wg_size = + utils::clamp_workgroup_size(device, kCompareWorkgroupSizeX); + utils::WgCount workgroup_count = + utils::compute_2d_workgroup_count(device, words, wg_size, "compare"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(CompareParams)); + graph.add_uniform_buffer_bytes(sizeof(CompareParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kCompareWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // out (rw storage) + in1/in2 (ro storage) + params (uniform). + WGPUBindGroupLayoutEntry entries[4] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_Storage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[3].binding = 3; + entries[3].visibility = WGPUShaderStage_Compute; + entries[3].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 4; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg[4] = {}; + bg[0].binding = 0; + bg[0].buffer = out_tensor.buffer; + bg[0].size = out_tensor.nbytes; + bg[1].binding = 1; + bg[1].buffer = in1_tensor.buffer; + bg[1].size = in1_tensor.nbytes; + bg[2].binding = 2; + bg[2].buffer = in2_tensor.buffer; + bg[2].size = in2_tensor.nbytes; + bg[3].binding = 3; + bg[3].buffer = uniform_buffer; + bg[3].size = sizeof(CompareParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 4; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, bind_group, workgroup_count.x, "compare", workgroup_count.y}); + + // Dynamic shapes: recompute numel/dispatch; out follows in1 (same-shape). + WGPUBuffer params_buf = uniform_buffer; + auto resize = [in1_id, in2_id, out_id, op, wg_size, dispatch_idx, params_buf]( + WebGPUGraph& g) { + const auto& d = g.cur_dims(in1_id); + const uint64_t n = utils::numel_of(d); + if (n == 0u || n % 4u != 0u || n > UINT32_MAX || + utils::numel_of(g.cur_dims(in2_id)) != n) { + throw std::runtime_error("compare(resize): numel must be a mult of 4"); + } + g.set_cur_dims(out_id, d); + CompareParams p = {}; + p.num_elements = static_cast(n); + p.op = op; + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), static_cast(n / 4u), wg_size, "compare"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }; + graph.add_tensor_resize_hook(in1_id, resize); + graph.add_tensor_resize_hook(in2_id, resize); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(uniform_buffer); +} + +void eq_op(WebGPUGraph& graph, const std::vector& args) { + compare_impl(graph, args, 0u); +} +void lt_op(WebGPUGraph& graph, const std::vector& args) { + compare_impl(graph, args, 1u); +} +void le_op(WebGPUGraph& graph, const std::vector& args) { + compare_impl(graph, args, 2u); +} +void gt_op(WebGPUGraph& graph, const std::vector& args) { + compare_impl(graph, args, 3u); +} +void ge_op(WebGPUGraph& graph, const std::vector& args) { + compare_impl(graph, args, 4u); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.eq.Tensor, eq_op); + WEBGPU_REGISTER_OP(aten.lt.Tensor, lt_op); + WEBGPU_REGISTER_OP(aten.le.Tensor, le_op); + WEBGPU_REGISTER_OP(aten.gt.Tensor, gt_op); + WEBGPU_REGISTER_OP(aten.ge.Tensor, ge_op); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/compare/compare.wgsl b/backends/webgpu/runtime/ops/compare/compare.wgsl new file mode 100644 index 00000000000..a16568e5533 --- /dev/null +++ b/backends/webgpu/runtime/ops/compare/compare.wgsl @@ -0,0 +1,43 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var input1: array; +@group(0) @binding(2) var input2: array; + +struct Params { + num_elements: u32, + op: u32, + pad0: u32, + pad1: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per output word = 4 bool bytes; num_elements%4==0 (host). + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + let words = (params.num_elements + 3u) / 4u; + if (widx >= words) { + return; + } + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let i = widx * 4u + j; + let a = input1[i]; + let b = input2[i]; + var r: bool; + switch params.op { + case 0u: { r = a == b; } // eq + case 1u: { r = a < b; } // lt + case 2u: { r = a <= b; } // le + case 3u: { r = a > b; } // gt + default: { r = a >= b; } // ge + } + if (r) { + packed = packed | (1u << (j * 8u)); + } + } + t_out[widx] = packed; +} diff --git a/backends/webgpu/runtime/ops/compare/compare_wgsl.h b/backends/webgpu/runtime/ops/compare/compare_wgsl.h new file mode 100644 index 00000000000..672c99b62d8 --- /dev/null +++ b/backends/webgpu/runtime/ops/compare/compare_wgsl.h @@ -0,0 +1,67 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from compare.wgsl - DO NOT EDIT. +// wgsl-sha256: 241e7e6762b1eded07d28a3767936c970f509f6591e7cbf599d0b1eb61efb181 +inline constexpr const char* kCompareWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var input1: array; +@group(0) @binding(2) var input2: array; + +struct Params { + num_elements: u32, + op: u32, + pad0: u32, + pad1: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One thread per output word = 4 bool bytes; num_elements%4==0 (host). + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + let words = (params.num_elements + 3u) / 4u; + if (widx >= words) { + return; + } + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let i = widx * 4u + j; + let a = input1[i]; + let b = input2[i]; + var r: bool; + switch params.op { + case 0u: { r = a == b; } // eq + case 1u: { r = a < b; } // lt + case 2u: { r = a <= b; } // le + case 3u: { r = a > b; } // gt + default: { r = a >= b; } // ge + } + if (r) { + packed = packed | (1u << (j * 8u)); + } + } + t_out[widx] = packed; +} +)"; + +inline constexpr uint32_t kCompareWorkgroupSizeX = 64; +inline constexpr uint32_t kCompareWorkgroupSizeY = 1; +inline constexpr uint32_t kCompareWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From d77a35ff681326b082281eb5b489dfe6d7c6a8b3 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:38 -0700 Subject: [PATCH 62/75] [ExecuTorch][WebGPU] Op-tests for comparisons + bool-output golden harness MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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()` 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/) --- backends/webgpu/test/op_tests/cases.py | 51 ++++++++++ .../webgpu/test/op_tests/generate_op_tests.py | 80 ++++++++++------ .../webgpu/test/op_tests/op_test_driver.cpp | 18 +++- backends/webgpu/test/ops/test_compare.py | 96 +++++++++++++++---- 4 files changed, 198 insertions(+), 47 deletions(-) diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 03dc628ffae..083e9b588c1 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -34,6 +34,11 @@ CatModule, CONFIGS as _CAT_CONFIGS, ) +from executorch.backends.webgpu.test.ops.test_compare import ( + CompareModule, + compare_gen_a, + compare_gen_b, +) from executorch.backends.webgpu.test.ops.test_floor_divide import FloorDivideModule from executorch.backends.webgpu.test.ops.test_argmax import ( argmax_tie_gen, @@ -256,6 +261,52 @@ def _minimum_suite() -> WebGPUTestSuite: ) +def _compare_suite(op: str) -> WebGPUTestSuite: + # Elementwise fp32 comparison -> bool (byte-exact golden). The two inputs use + # DIFFERENT discrete-range seeds so a!=b (real lt/gt mix) while colliding + # often (eq/le/ge ties); all shapes have numel % 4 == 0 (bool output packs 4 + # bytes/word). Same-shape only (flat kernel; broadcast=smoke). + def case(name, shape): + return Case( + name=name, + inputs=( + InputSpec(shape=shape, gen=compare_gen_a), + InputSpec(shape=shape, gen=compare_gen_b), + ), + ) + + return WebGPUTestSuite( + module_factory=lambda: CompareModule(op), + cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + golden_dtype="bool", + ) + + +@register_op_test("eq") +def _eq_suite() -> WebGPUTestSuite: + return _compare_suite("eq") + + +@register_op_test("lt") +def _lt_suite() -> WebGPUTestSuite: + return _compare_suite("lt") + + +@register_op_test("le") +def _le_suite() -> WebGPUTestSuite: + return _compare_suite("le") + + +@register_op_test("gt") +def _gt_suite() -> WebGPUTestSuite: + return _compare_suite("gt") + + +@register_op_test("ge") +def _ge_suite() -> WebGPUTestSuite: + return _compare_suite("ge") + + @register_op_test("pow") def _pow_suite() -> WebGPUTestSuite: # Positive base avoids pow(neg, frac)=NaN; exponent spans negative+positive. diff --git a/backends/webgpu/test/op_tests/generate_op_tests.py b/backends/webgpu/test/op_tests/generate_op_tests.py index 1907f20bbe1..72f819f94ce 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -86,6 +86,49 @@ def _write_int64(t: torch.Tensor, path: str) -> None: t.detach().contiguous().cpu().numpy().astype(" tuple[torch.Tensor, str]: + if raw.dtype == torch.bool: + # bool-output ops (comparisons): byte-exact 0/1 golden written as int8. + out_t = raw.to(torch.int8) + out_dtype = "bool" + elif raw.dtype == torch.int8: + # int8-output ops (quantize): byte-exact golden, no fp32 cast/oracle. + out_t = raw + out_dtype = "int8" + elif raw.dtype in (torch.int64, torch.int32): + # int64-index ops (argmax/argmin): exact golden, no fp32 cast/oracle. + out_t = raw.to(torch.int64) + out_dtype = "int64" + else: + out_t = raw.to(torch.float32) + out_dtype = "float32" + # Dual-oracle gate: the fp64 golden must match the fp32 eager within + # tol — proves the oracle isn't itself buggy. Skipped for float32. + if golden_dtype == "float64" and not has_custom_golden: + torch.testing.assert_close( + eager.to(torch.float32), + out_t, + atol=atol, + rtol=rtol, + ) + + if out_dtype in ("bool", "int8"): + _write_int8(out_t, path) + elif out_dtype == "int64": + _write_int64(out_t, path) + else: + _write_fp32(out_t, path) + return out_t, out_dtype + + def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> list[dict]: """Export one case; write its .pte + input/golden .bin(s). Returns one manifest entry per output tensor (multi-output ops emit N; single-output ops emit one, @@ -138,37 +181,18 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> list[d n_out = len(golden_outs) for out_index in range(n_out): raw = golden_outs[out_index] - is_int8 = raw.dtype == torch.int8 - is_int64 = raw.dtype in (torch.int64, torch.int32) - if is_int8: - # int8-output ops (quantize): byte-exact golden, no fp32 cast/oracle. - out_t = raw - out_dtype = "int8" - elif is_int64: - # int64-index ops (argmax/argmin): exact golden, no fp32 cast/oracle. - out_t = raw.to(torch.int64) - out_dtype = "int64" - else: - out_t = raw.to(torch.float32) - out_dtype = "float32" - # Dual-oracle gate: the fp64 golden must match the fp32 eager within - # tol — proves the oracle isn't itself buggy. Skipped for float32. - if golden_dtype == "float64" and case.golden_fn is None: - torch.testing.assert_close( - eager_outs[out_index].to(torch.float32), - out_t, - atol=atol, - rtol=rtol, - ) # Single-output ops keep the original name/golden; multi-output ops suffix. suffix = f"_out{out_index}" if n_out > 1 else "" golden_rel = f"{case_id}{suffix}.golden.bin" - if is_int8: - _write_int8(out_t, os.path.join(out_dir, golden_rel)) - elif is_int64: - _write_int64(out_t, os.path.join(out_dir, golden_rel)) - else: - _write_fp32(out_t, os.path.join(out_dir, golden_rel)) + out_t, out_dtype = _write_golden_output( + raw, + eager_outs[out_index], + golden_dtype, + case.golden_fn is not None, + atol, + rtol, + os.path.join(out_dir, golden_rel), + ) entries.append( { "op": op, diff --git a/backends/webgpu/test/op_tests/op_test_driver.cpp b/backends/webgpu/test/op_tests/op_test_driver.cpp index c29daf6e101..79b72c1bd2d 100644 --- a/backends/webgpu/test/op_tests/op_test_driver.cpp +++ b/backends/webgpu/test/op_tests/op_test_driver.cpp @@ -91,7 +91,23 @@ class OpCase : public ::testing::Test { out_tensor.sizes().begin(), out_tensor.sizes().end()); EXPECT_EQ(out_shape, e_.golden.shape) << "output shape != golden shape (numel matched but dims differ)"; - if (e_.golden.dtype == "int8") { + if (e_.golden.dtype == "bool") { + // bool-output ops (comparisons) compare byte-exact 0/1 (golden as int8). + auto golden = load_int8_bin(e_.golden.path, gn); + ASSERT_FALSE(golden.empty()) + << "missing/short golden: " << e_.golden.path; + const bool* out_p = out_tensor.const_data_ptr(); + int mism = -1; + for (size_t i = 0; i < gn; i++) { + if (static_cast(out_p[i]) != golden[i]) { + mism = static_cast(i); + break; + } + } + EXPECT_EQ(mism, -1) << "bool mismatch at index " << mism + << ": out=" << (mism >= 0 ? int(out_p[mism]) : 0) + << " golden=" << (mism >= 0 ? int(golden[mism]) : 0); + } else if (e_.golden.dtype == "int8") { // int8-output ops (quantize) compare byte-exact: a discrete grid, so any // deviation is a real bug, not tolerance. auto golden = load_int8_bin(e_.golden.path, gn); diff --git a/backends/webgpu/test/ops/test_compare.py b/backends/webgpu/test/ops/test_compare.py index 172c5e75113..9d571b6f653 100644 --- a/backends/webgpu/test/ops/test_compare.py +++ b/backends/webgpu/test/ops/test_compare.py @@ -4,15 +4,19 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""`aten.{eq,ne,le,ge,lt,gt}.Scalar` export + golden for the WebGPU backend. - -Each scalar comparison lowers to a single `aten..Scalar` node that the -kernel computes as `cmp(self[i], scalar)` and writes as a byte-packed bool. The -delegation test locks that every variant partitions to `VulkanBackend`; the -golden test locks the fp32 module output against the fp64 torch truth. The -deterministic ramp is exact in fp32 and straddles (and hits) the scalar, so both -precisions agree bit-for-bit and every op has mixed True/False cases; a -non-multiple-of-4 numel exercises the kernel tail (4 elems per u32 word). +"""Comparison-op tests for the WebGPU backend. + +Two suites share this file: +- `aten.{eq,ne,le,ge,lt,gt}.Scalar` (`ScalarCompareModule` / `TestCompare`): each + scalar comparison lowers to a single `aten..Scalar` node computed as + `cmp(self[i], scalar)` and written as a byte-packed bool. The deterministic ramp + is exact in fp32 and straddles (and hits) the scalar, so fp32/fp64 agree + bit-for-bit and every op has mixed True/False cases; a non-multiple-of-4 numel + exercises the kernel tail (4 elems per u32 word). +- `aten.{eq,lt,le,gt,ge}.Tensor` (`CompareModule` / `CompareTest`, imported by + `cases.py`): elementwise fp32 tensor comparison -> bool. Inputs come from a small + discrete range so `eq`/`le`/`ge` see genuine ties and `lt`/`gt` are a real + true/false mix; numel is a multiple of 4 (bool output packs 4/word). """ from __future__ import annotations @@ -25,8 +29,8 @@ from executorch.exir import to_edge_transform_and_lower SCALAR = 0.0 -OPS = ("eq", "ne", "le", "ge", "lt", "gt") -SHAPES = {"tail": (3, 5), "3d": (2, 4, 8)} +SCALAR_OPS = ("eq", "ne", "le", "ge", "lt", "gt") +SCALAR_SHAPES = {"tail": (3, 5), "3d": (2, 4, 8)} _TORCH_OP = { "eq": torch.eq, @@ -38,7 +42,7 @@ } -class CompareModule(torch.nn.Module): +class ScalarCompareModule(torch.nn.Module): def __init__(self, op: str, scalar: float) -> None: super().__init__() self.op = op @@ -85,25 +89,81 @@ def _delegated(et) -> bool: class TestCompare(unittest.TestCase): def test_export_delegates(self) -> None: - for op in OPS: - for name, shape in SHAPES.items(): + for op in SCALAR_OPS: + for name, shape in SCALAR_SHAPES.items(): with self.subTest(op=op, shape=name): x = _det_input(shape) - et = _export(CompareModule(op, SCALAR).eval(), x) + et = _export(ScalarCompareModule(op, SCALAR).eval(), x) self.assertTrue( _delegated(et), f"Expected a VulkanBackend delegate ({op}.Scalar {name})", ) def test_module_matches_fp64_golden(self) -> None: - for op in OPS: - for name, shape in SHAPES.items(): + for op in SCALAR_OPS: + for name, shape in SCALAR_SHAPES.items(): with self.subTest(op=op, shape=name): x = _det_input(shape) - got = CompareModule(op, SCALAR)(x) + got = ScalarCompareModule(op, SCALAR)(x) golden = _TORCH_OP[op](x.double(), float(SCALAR)) torch.testing.assert_close(got, golden) +class CompareModule(torch.nn.Module): + def __init__(self, op) -> None: + super().__init__() + self.op = op + + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + if self.op == "eq": + return a == b + if self.op == "lt": + return a < b + if self.op == "le": + return a <= b + if self.op == "gt": + return a > b + return a >= b # ge + + +def _cmp_gen(seed): + # Small discrete range from a per-input seed: the two inputs DIFFER (a!=b) + # while still colliding often (so eq/le/ge see genuine ties and lt/gt are a + # real true/false mix). numel is a multiple of 4 (bool output packs 4/word). + def g(shape): + gen = torch.Generator().manual_seed(seed) + return torch.randint(-2, 3, shape, generator=gen).to(torch.float32) + + return g + + +compare_gen_a = _cmp_gen(0) +compare_gen_b = _cmp_gen(1) + + +OPS = ["eq", "lt", "le", "gt", "ge"] +# All shapes have numel % 4 == 0 (bool output is packed 4 bytes/word). +SHAPES = [(4, 8), (2, 3, 8), (16, 16)] + + +class CompareTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for op in OPS: + with self.subTest(op=op): + a = compare_gen_a((4, 8)) + b = compare_gen_b((4, 8)) + ep = torch.export.export(CompareModule(op).eval(), (a, b)) + edge = to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ) + et = edge.to_executorch() + deleg = any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + self.assertTrue(deleg, f"Expected VulkanBackend delegate ({op})") + + if __name__ == "__main__": unittest.main() From e7c421cab1706b4b246563822884edb4b6527d98 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:39 -0700 Subject: [PATCH 63/75] [ExecuTorch][WebGPU] Add logical_and op (aten.logical_and.default) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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`, 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/) --- .../runtime/ops/logical_and/LogicalAnd.cpp | 190 ++++++++++++++++++ .../runtime/ops/logical_and/logical_and.wgsl | 25 +++ .../ops/logical_and/logical_and_wgsl.h | 49 +++++ 3 files changed, 264 insertions(+) create mode 100644 backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp create mode 100644 backends/webgpu/runtime/ops/logical_and/logical_and.wgsl create mode 100644 backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h diff --git a/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp new file mode 100644 index 00000000000..9c398cee862 --- /dev/null +++ b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp @@ -0,0 +1,190 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform layout matching the WGSL Params struct; 16-byte aligned. +struct LogicalAndParams { + uint32_t num_words; + uint32_t _pad[3]; +}; +static_assert( + sizeof(LogicalAndParams) == 16, + "LogicalAndParams must be 16 bytes"); + +// out = a & b (canonical bool); mirrors Vulkan bitwise_and (BinaryOp.cpp:161). +void logical_and_op(WebGPUGraph& graph, const std::vector& args) { + const int a_id = args.at(0); + const int b_id = args.at(1); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(a_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(b_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("logical_and: a/b/out is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& a_tensor = graph.get_tensor(a_id); + const auto& b_tensor = graph.get_tensor(b_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (a_tensor.buffer == nullptr || b_tensor.buffer == nullptr || + out_tensor.buffer == nullptr) { + throw std::runtime_error("logical_and: null buffer binding"); + } + // a/b/out are all 1-byte bool tensors (int-typed, NOT int8-quantized). + if (!a_tensor.is_int || !b_tensor.is_int || !out_tensor.is_int || + a_tensor.elem_size != 1 || b_tensor.elem_size != 1 || + out_tensor.elem_size != 1) { + throw std::runtime_error( + "logical_and: a/b/out must be 1-byte bool tensors"); + } + const uint64_t numel = out_tensor.nbytes; + // bool packed 4/word (array); numel%4==0 gates the u32 binding. + if (numel == 0u || numel % 4u != 0u || numel > UINT32_MAX) { + throw std::runtime_error("logical_and: numel must be a nonzero mult of 4"); + } + if (a_tensor.nbytes != numel || b_tensor.nbytes != numel) { + throw std::runtime_error( + "logical_and: a/b/out numel mismatch (same-shape)"); + } + + LogicalAndParams params = {}; + params.num_words = static_cast(numel / 4u); + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kLogicalAndWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, params.num_words, wg_size, "logical_and"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(LogicalAndParams)); + graph.add_uniform_buffer_bytes(sizeof(LogicalAndParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kLogicalAndWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // out (rw storage) + a/b (ro storage) + params (uniform). + WGPUBindGroupLayoutEntry entries[4] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_Storage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[3].binding = 3; + entries[3].visibility = WGPUShaderStage_Compute; + entries[3].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 4; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg[4] = {}; + bg[0].binding = 0; + bg[0].buffer = out_tensor.buffer; + bg[0].size = out_tensor.nbytes; + bg[1].binding = 1; + bg[1].buffer = a_tensor.buffer; + bg[1].size = a_tensor.nbytes; + bg[2].binding = 2; + bg[2].buffer = b_tensor.buffer; + bg[2].size = b_tensor.nbytes; + bg[3].binding = 3; + bg[3].buffer = uniform_buffer; + bg[3].size = sizeof(LogicalAndParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 4; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "logical_and", + workgroup_count.y}); + + // Dynamic shapes: recompute num_words/dispatch; out follows a (same-shape). + WGPUBuffer params_buf = uniform_buffer; + auto resize = + [a_id, b_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(a_id); + const uint64_t n = utils::numel_of(d); + if (n == 0u || n % 4u != 0u || n > UINT32_MAX || + utils::numel_of(g.cur_dims(b_id)) != n) { + throw std::runtime_error( + "logical_and(resize): numel must be a mult of 4"); + } + g.set_cur_dims(out_id, d); + LogicalAndParams p = {}; + p.num_words = static_cast(n / 4u); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), p.num_words, wg_size, "logical_and"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }; + graph.add_tensor_resize_hook(a_id, resize); + graph.add_tensor_resize_hook(b_id, resize); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.logical_and.default, logical_and_op); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/logical_and/logical_and.wgsl b/backends/webgpu/runtime/ops/logical_and/logical_and.wgsl new file mode 100644 index 00000000000..9acb583f51c --- /dev/null +++ b/backends/webgpu/runtime/ops/logical_and/logical_and.wgsl @@ -0,0 +1,25 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_a: array; +@group(0) @binding(2) var t_b: array; + +struct Params { + num_words: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // bool packed 4/word; canonical 0/1 bytes -> word-wise AND == per-byte AND. + let w = gid.x + gid.y * (num_workgroups.x * wg_size); + if (w >= params.num_words) { + return; + } + t_out[w] = t_a[w] & t_b[w]; +} diff --git a/backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h b/backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h new file mode 100644 index 00000000000..6a21a77a687 --- /dev/null +++ b/backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h @@ -0,0 +1,49 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from logical_and.wgsl - DO NOT EDIT. +// wgsl-sha256: cf7c1d1dbba94e429120796c9c25a6717786cca03c08f3bd1e291d5627089c20 +inline constexpr const char* kLogicalAndWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_a: array; +@group(0) @binding(2) var t_b: array; + +struct Params { + num_words: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // bool packed 4/word; canonical 0/1 bytes -> word-wise AND == per-byte AND. + let w = gid.x + gid.y * (num_workgroups.x * wg_size); + if (w >= params.num_words) { + return; + } + t_out[w] = t_a[w] & t_b[w]; +} +)"; + +inline constexpr uint32_t kLogicalAndWorkgroupSizeX = 64; +inline constexpr uint32_t kLogicalAndWorkgroupSizeY = 1; +inline constexpr uint32_t kLogicalAndWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 594f64d6a6701b9d779e81001348bc952bf36758 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:39 -0700 Subject: [PATCH 64/75] [ExecuTorch][WebGPU] Op-tests for logical_and MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 28 ++++++++ backends/webgpu/test/ops/test_logical_and.py | 76 ++++++++++++++++++++ 2 files changed, 104 insertions(+) create mode 100644 backends/webgpu/test/ops/test_logical_and.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 083e9b588c1..33f44cdc929 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -39,6 +39,11 @@ compare_gen_a, compare_gen_b, ) +from executorch.backends.webgpu.test.ops.test_logical_and import ( + la_gen_a, + la_gen_b, + LogicalAndModule, +) from executorch.backends.webgpu.test.ops.test_floor_divide import FloorDivideModule from executorch.backends.webgpu.test.ops.test_argmax import ( argmax_tie_gen, @@ -307,6 +312,29 @@ def _ge_suite() -> WebGPUTestSuite: return _compare_suite("ge") +@register_op_test("logical_and") +def _logical_and_suite() -> WebGPUTestSuite: + # out = (a>0) && (b>0): two bool masks derived on-GPU from float inputs via + # gt.Tensor (baked zeros), AND'd -> bool. Distinct a/b seeds so the masks + # differ (AND ~25% True, a real mix an OR mutant fails); all shapes numel % + # 4 == 0 (bool packs 4/word). float32 oracle (byte-exact bool golden). + def case(name, shape): + return Case( + name=name, + construct={"shape": shape}, + inputs=( + InputSpec(shape=shape, gen=la_gen_a), + InputSpec(shape=shape, gen=la_gen_b), + ), + ) + + return WebGPUTestSuite( + module_factory=lambda shape: LogicalAndModule(shape), + cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + golden_dtype="float32", + ) + + @register_op_test("pow") def _pow_suite() -> WebGPUTestSuite: # Positive base avoids pow(neg, frac)=NaN; exponent spans negative+positive. diff --git a/backends/webgpu/test/ops/test_logical_and.py b/backends/webgpu/test/ops/test_logical_and.py new file mode 100644 index 00000000000..4c3939002d0 --- /dev/null +++ b/backends/webgpu/test/ops/test_logical_and.py @@ -0,0 +1,76 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.logical_and.default` module + configs for the WebGPU op-test framework. + +`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). `a`/`b` use distinct seeds so the two bool masks differ (each +~50% True, independent -> AND ~25% True), a real mix that a wrong op (e.g. OR) +would fail. Output is bool (byte-exact golden). `LogicalAndTest` is the +export-delegation smoke test. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + + +class LogicalAndModule(torch.nn.Module): + def __init__(self, shape) -> None: + super().__init__() + self.register_buffer("z", torch.zeros(shape)) + + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return torch.logical_and(a > self.z, b > self.z) + + +def _la_gen(seed): + # Distinct per-input seed so the two derived bool masks differ. + def g(shape): + gen = torch.Generator().manual_seed(seed) + return torch.randn(*shape, generator=gen, dtype=torch.float32) + + return g + + +la_gen_a = _la_gen(0) +la_gen_b = _la_gen(1) + + +# All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). +SHAPES = [(4, 8), (2, 3, 8), (16, 16)] + + +class LogicalAndTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for shape in SHAPES: + with self.subTest(shape=shape): + a = la_gen_a(shape) + b = la_gen_b(shape) + ep = torch.export.export(LogicalAndModule(shape).eval(), (a, b)) + edge = to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ) + et = edge.to_executorch() + deleg = any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + self.assertTrue(deleg, f"Expected VulkanBackend delegate ({shape})") + gm = edge.exported_program().graph_module + self.assertTrue( + all( + "logical_and" not in str(getattr(n, "target", "")) + for n in gm.graph.nodes + ), + f"logical_and fell back to CPU for {shape}", + ) From 74509538c1fc5d4b676508726752ea1e1d2784b2 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:40 -0700 Subject: [PATCH 65/75] [ExecuTorch][WebGPU] Add bitwise_and + bitwise_not ops (bool) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../runtime/ops/bitwise_not/BitwiseNot.cpp | 175 ++++++++++++++++++ .../runtime/ops/bitwise_not/bitwise_not.wgsl | 24 +++ .../ops/bitwise_not/bitwise_not_wgsl.h | 48 +++++ .../runtime/ops/logical_and/LogicalAnd.cpp | 3 +- 4 files changed, 249 insertions(+), 1 deletion(-) create mode 100644 backends/webgpu/runtime/ops/bitwise_not/BitwiseNot.cpp create mode 100644 backends/webgpu/runtime/ops/bitwise_not/bitwise_not.wgsl create mode 100644 backends/webgpu/runtime/ops/bitwise_not/bitwise_not_wgsl.h diff --git a/backends/webgpu/runtime/ops/bitwise_not/BitwiseNot.cpp b/backends/webgpu/runtime/ops/bitwise_not/BitwiseNot.cpp new file mode 100644 index 00000000000..22ab8daa654 --- /dev/null +++ b/backends/webgpu/runtime/ops/bitwise_not/BitwiseNot.cpp @@ -0,0 +1,175 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform layout matching the WGSL Params struct; 16-byte aligned. +struct BitwiseNotParams { + uint32_t num_words; + uint32_t _pad[3]; +}; +static_assert( + sizeof(BitwiseNotParams) == 16, + "BitwiseNotParams must be 16 bytes"); + +// out = ~a on 1-byte bools; mirrors Vulkan bitwise_not (1-X). args: [a, out]. +void bitwise_not_op(WebGPUGraph& graph, const std::vector& args) { + const int a_id = args.at(0); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(a_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("bitwise_not: a/out is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& a_tensor = graph.get_tensor(a_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (a_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("bitwise_not: null buffer binding"); + } + // a/out are 1-byte bool tensors (int-typed, NOT int8-quantized). + if (!a_tensor.is_int || !out_tensor.is_int || a_tensor.elem_size != 1 || + out_tensor.elem_size != 1) { + throw std::runtime_error("bitwise_not: a/out must be 1-byte bool tensors"); + } + const uint64_t numel = out_tensor.nbytes; + // bool packed 4/word (array); numel%4==0 gates the u32 binding. + if (numel == 0u || numel % 4u != 0u || numel > UINT32_MAX) { + throw std::runtime_error("bitwise_not: numel must be a nonzero mult of 4"); + } + if (a_tensor.nbytes != numel) { + throw std::runtime_error("bitwise_not: a/out numel mismatch (same-shape)"); + } + + BitwiseNotParams params = {}; + params.num_words = static_cast(numel / 4u); + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kBitwiseNotWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, params.num_words, wg_size, "bitwise_not"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(BitwiseNotParams)); + graph.add_uniform_buffer_bytes(sizeof(BitwiseNotParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kBitwiseNotWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // out (rw storage) + a (ro storage) + params (uniform). + WGPUBindGroupLayoutEntry entries[3] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_Storage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 3; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg[3] = {}; + bg[0].binding = 0; + bg[0].buffer = out_tensor.buffer; + bg[0].size = out_tensor.nbytes; + bg[1].binding = 1; + bg[1].buffer = a_tensor.buffer; + bg[1].size = a_tensor.nbytes; + bg[2].binding = 2; + bg[2].buffer = uniform_buffer; + bg[2].size = sizeof(BitwiseNotParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 3; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "bitwise_not", + workgroup_count.y}); + + // Dynamic shapes: recompute num_words/dispatch; out follows a (same-shape). + WGPUBuffer params_buf = uniform_buffer; + auto resize = + [a_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(a_id); + const uint64_t n = utils::numel_of(d); + if (n == 0u || n % 4u != 0u || n > UINT32_MAX) { + throw std::runtime_error( + "bitwise_not(resize): numel must be a mult of 4"); + } + g.set_cur_dims(out_id, d); + BitwiseNotParams p = {}; + p.num_words = static_cast(n / 4u); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), p.num_words, wg_size, "bitwise_not"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }; + graph.add_tensor_resize_hook(a_id, resize); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.bitwise_not.default, bitwise_not_op); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/bitwise_not/bitwise_not.wgsl b/backends/webgpu/runtime/ops/bitwise_not/bitwise_not.wgsl new file mode 100644 index 00000000000..3104761a42a --- /dev/null +++ b/backends/webgpu/runtime/ops/bitwise_not/bitwise_not.wgsl @@ -0,0 +1,24 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_a: array; + +struct Params { + num_words: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // bool NOT packed 4/word: per-byte 1-x == x^1; word-wise ^0x01010101. + let w = gid.x + gid.y * (num_workgroups.x * wg_size); + if (w >= params.num_words) { + return; + } + t_out[w] = t_a[w] ^ 0x01010101u; +} diff --git a/backends/webgpu/runtime/ops/bitwise_not/bitwise_not_wgsl.h b/backends/webgpu/runtime/ops/bitwise_not/bitwise_not_wgsl.h new file mode 100644 index 00000000000..082654bb311 --- /dev/null +++ b/backends/webgpu/runtime/ops/bitwise_not/bitwise_not_wgsl.h @@ -0,0 +1,48 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from bitwise_not.wgsl - DO NOT EDIT. +// wgsl-sha256: 06c95212af5f02c08978288f25552da0e2476c37f9421e941d059338ffa1aa9f +inline constexpr const char* kBitwiseNotWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_a: array; + +struct Params { + num_words: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // bool NOT packed 4/word: per-byte 1-x == x^1; word-wise ^0x01010101. + let w = gid.x + gid.y * (num_workgroups.x * wg_size); + if (w >= params.num_words) { + return; + } + t_out[w] = t_a[w] ^ 0x01010101u; +} +)"; + +inline constexpr uint32_t kBitwiseNotWorkgroupSizeX = 64; +inline constexpr uint32_t kBitwiseNotWorkgroupSizeY = 1; +inline constexpr uint32_t kBitwiseNotWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp index 9c398cee862..f18305de01b 100644 --- a/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp +++ b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp @@ -30,7 +30,7 @@ static_assert( sizeof(LogicalAndParams) == 16, "LogicalAndParams must be 16 bytes"); -// out = a & b (canonical bool); mirrors Vulkan bitwise_and (BinaryOp.cpp:161). +// out = a & b (bool AND); serves logical_and + bitwise_and (mirrors Vulkan). void logical_and_op(WebGPUGraph& graph, const std::vector& args) { const int a_id = args.at(0); const int b_id = args.at(1); @@ -185,6 +185,7 @@ void logical_and_op(WebGPUGraph& graph, const std::vector& args) { WEBGPU_REGISTER_OPERATORS { WEBGPU_REGISTER_OP(aten.logical_and.default, logical_and_op); + WEBGPU_REGISTER_OP(aten.bitwise_and.Tensor, logical_and_op); } } // namespace executorch::backends::webgpu From 7c43f341afe0964aa3951172ada5b9bf8515f62e Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:40 -0700 Subject: [PATCH 66/75] [ExecuTorch][WebGPU] Op-tests for bitwise_and + bitwise_not MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/op_tests/cases.py | 46 +++++++++++++ backends/webgpu/test/ops/test_bitwise.py | 87 ++++++++++++++++++++++++ 2 files changed, 133 insertions(+) create mode 100644 backends/webgpu/test/ops/test_bitwise.py diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 33f44cdc929..9cf0c82e8d8 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -44,6 +44,12 @@ la_gen_b, LogicalAndModule, ) +from executorch.backends.webgpu.test.ops.test_bitwise import ( + BitwiseAndModule, + BitwiseNotModule, + bw_gen_a, + bw_gen_b, +) from executorch.backends.webgpu.test.ops.test_floor_divide import FloorDivideModule from executorch.backends.webgpu.test.ops.test_argmax import ( argmax_tie_gen, @@ -335,6 +341,46 @@ def case(name, shape): ) +@register_op_test("bitwise_and") +def _bitwise_and_suite() -> WebGPUTestSuite: + # bool bitwise AND == logical_and for canonical 0/1 (shares the handler). + # Two masks derived on-GPU from float inputs via gt.Tensor (baked zeros), + # distinct a/b seeds (AND ~25% True); all shapes numel % 4 == 0. + def case(name, shape): + return Case( + name=name, + construct={"shape": shape}, + inputs=( + InputSpec(shape=shape, gen=bw_gen_a), + InputSpec(shape=shape, gen=bw_gen_b), + ), + ) + + return WebGPUTestSuite( + module_factory=lambda shape: BitwiseAndModule(shape), + cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + golden_dtype="float32", + ) + + +@register_op_test("bitwise_not") +def _bitwise_not_suite() -> WebGPUTestSuite: + # bool NOT (1-x): one mask derived on-GPU from a float input via gt.Tensor + # (baked zeros), inverted -> bool (~50% True); all shapes numel % 4 == 0. + def case(name, shape): + return Case( + name=name, + construct={"shape": shape}, + inputs=(InputSpec(shape=shape, gen=bw_gen_a),), + ) + + return WebGPUTestSuite( + module_factory=lambda shape: BitwiseNotModule(shape), + cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + golden_dtype="float32", + ) + + @register_op_test("pow") def _pow_suite() -> WebGPUTestSuite: # Positive base avoids pow(neg, frac)=NaN; exponent spans negative+positive. diff --git a/backends/webgpu/test/ops/test_bitwise.py b/backends/webgpu/test/ops/test_bitwise.py new file mode 100644 index 00000000000..b61a37d9990 --- /dev/null +++ b/backends/webgpu/test/ops/test_bitwise.py @@ -0,0 +1,87 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.bitwise_and.Tensor` / `aten.bitwise_not.default` (bool) modules + configs. + +Both are partitioner-tagged for BOOL inputs only, so the modules derive their +bool operands on-GPU from float inputs (`a > 0` via the delegated `gt.Tensor` +against a baked zero buffer) — the only runtime inputs are float tensors (the +op-test framework is float-input-only). `bitwise_and` on bool is identical to +`logical_and` (shares the handler); `bitwise_not` is the bool NOT (1-x). Output +is bool (byte-exact golden). `BitwiseTest` is the export-delegation smoke test. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + + +class BitwiseAndModule(torch.nn.Module): + def __init__(self, shape) -> None: + super().__init__() + self.register_buffer("z", torch.zeros(shape)) + + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return torch.bitwise_and(a > self.z, b > self.z) + + +class BitwiseNotModule(torch.nn.Module): + def __init__(self, shape) -> None: + super().__init__() + self.register_buffer("z", torch.zeros(shape)) + + def forward(self, a: torch.Tensor) -> torch.Tensor: + return torch.bitwise_not(a > self.z) + + +def _bw_gen(seed): + # Distinct per-input seed so derived bool masks differ (bitwise_and). + def g(shape): + gen = torch.Generator().manual_seed(seed) + return torch.randn(*shape, generator=gen, dtype=torch.float32) + + return g + + +bw_gen_a = _bw_gen(0) +bw_gen_b = _bw_gen(1) + + +# All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). +SHAPES = [(4, 8), (2, 3, 8), (16, 16)] + + +class BitwiseTest(unittest.TestCase): + def _assert_delegates(self, mod, inputs, op_name, shape) -> None: + ep = torch.export.export(mod.eval(), inputs) + edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + et = edge.to_executorch() + deleg = any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + self.assertTrue(deleg, f"Expected VulkanBackend delegate ({op_name} {shape})") + gm = edge.exported_program().graph_module + self.assertTrue( + all(op_name not in str(getattr(n, "target", "")) for n in gm.graph.nodes), + f"{op_name} fell back to CPU for {shape}", + ) + + def test_export_delegates(self) -> None: + for shape in SHAPES: + with self.subTest(shape=shape): + a = bw_gen_a(shape) + b = bw_gen_b(shape) + self._assert_delegates( + BitwiseAndModule(shape), (a, b), "bitwise_and", shape + ) + self._assert_delegates( + BitwiseNotModule(shape), (a,), "bitwise_not", shape + ) From 4b116d6fedc495a0cd80e7b7c2bef53cef4dd52b Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:41 -0700 Subject: [PATCH 67/75] [ExecuTorch][WebGPU] Optimize reduction-family kernels to cooperative shared-memory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/runtime/ops/amax/Reduce.cpp | 15 +++- backends/webgpu/runtime/ops/amax/amax.wgsl | 33 +++++++-- backends/webgpu/runtime/ops/amax/amax_wgsl.h | 35 +++++++-- backends/webgpu/runtime/ops/amin/Reduce.cpp | 15 +++- backends/webgpu/runtime/ops/amin/amin.wgsl | 33 +++++++-- backends/webgpu/runtime/ops/amin/amin_wgsl.h | 35 +++++++-- backends/webgpu/runtime/ops/argmax/Reduce.cpp | 14 +++- .../webgpu/runtime/ops/argmax/arg_reduce.wgsl | 48 +++++++++--- .../runtime/ops/argmax/arg_reduce_wgsl.h | 50 +++++++++---- .../ops/native_group_norm/GroupNorm.cpp | 30 ++++++-- .../native_group_norm/group_norm_reduce.wgsl | 39 +++++++--- .../group_norm_reduce_wgsl.h | 41 +++++++--- backends/webgpu/runtime/ops/reduce/Reduce.cpp | 74 +++++++++++++------ backends/webgpu/test/tester.py | 5 ++ 14 files changed, 368 insertions(+), 99 deletions(-) diff --git a/backends/webgpu/runtime/ops/amax/Reduce.cpp b/backends/webgpu/runtime/ops/amax/Reduce.cpp index 702aa21c218..4e76bdf369f 100644 --- a/backends/webgpu/runtime/ops/amax/Reduce.cpp +++ b/backends/webgpu/runtime/ops/amax/Reduce.cpp @@ -19,6 +19,10 @@ namespace executorch::backends::webgpu { +static_assert( + kAmaxWorkgroupSizeX <= 256, + "amax.wgsl partials[] is sized 256; wg_size must not exceed it"); + namespace { // Uniform layout matching the WGSL Params struct; 16-byte aligned. @@ -47,6 +51,9 @@ void amax_impl(WebGPUGraph& graph, const std::vector& args) { const std::vector& dims = graph.get_int_list(args.at(1)); const int64_t ndim = static_cast(in_tensor.dims.size()); + if (ndim <= 0) { + throw std::runtime_error("amax: input must have at least one dim"); + } if (dims.size() != 1 || (dims[0] != -1 && dims[0] != ndim - 1)) { throw std::runtime_error("amax: only last-dim reduction is supported"); } @@ -61,8 +68,9 @@ void amax_impl(WebGPUGraph& graph, const std::vector& args) { } uint32_t wg_size = utils::clamp_workgroup_size(device, kAmaxWorkgroupSizeX); + // One workgroup per row (cooperative reduction); grid = num_rows workgroups. utils::WgCount workgroup_count = - utils::compute_2d_workgroup_count(device, num_rows, wg_size, "amax"); + utils::compute_2d_workgroup_count(device, num_rows, 1, "amax"); WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; @@ -155,6 +163,9 @@ void amax_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error("amax(resize): zero reduce dim"); } const uint64_t total = utils::numel_of(d); + if (total % rsize != 0u) { + throw std::runtime_error("amax(resize): numel not divisible by dim"); + } const uint32_t rows = static_cast(total / rsize); std::vector od = d; if (keepdim) { @@ -168,7 +179,7 @@ void amax_impl(WebGPUGraph& graph, const std::vector& args) { p.reduce_size = rsize; wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); const utils::WgCount wgc = utils::compute_2d_workgroup_count( - g.device(), rows, wg_size, "amax"); + g.device(), rows, 1, "amax"); g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }); diff --git a/backends/webgpu/runtime/ops/amax/amax.wgsl b/backends/webgpu/runtime/ops/amax/amax.wgsl index 2ff33e0c20f..2f23b9c38fa 100644 --- a/backends/webgpu/runtime/ops/amax/amax.wgsl +++ b/backends/webgpu/runtime/ops/amax/amax.wgsl @@ -9,18 +9,41 @@ struct Params { override wg_size: u32 = 256u; +// Cooperative shared-memory reduction; mirrors Vulkan reduce.glsl (a group of +// threads co-operates per reduction row, partials aggregated in shared memory). +// Fixed upper bound (>= any clamped wg_size); only [0, wg_size) is used. +var partials: array; + @compute @workgroup_size(wg_size) fn main( - @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - let row = gid.x + gid.y * (num_workgroups.x * wg_size); + // One workgroup per reduction row; 2D-fold lifts the 65535 grid cap. + let row = wid.x + wid.y * num_workgroups.x; if (row >= params.num_rows) { return; } let base = row * params.reduce_size; + + // Each thread reduces a strided slice of the row into a partial. Seed with + // the row's first element (always valid; reduce_size >= 1) so threads that + // own no element contribute a real value, not an out-of-range identity. var acc = input[base]; - for (var j = 1u; j < params.reduce_size; j = j + 1u) { - acc = max(acc, input[base + j]); + var i = lid.x; + while (i < params.reduce_size) { + acc = max(acc, input[base + i]); + i = i + wg_size; + } + partials[lid.x] = acc; + workgroupBarrier(); + + // Thread 0 aggregates the wg_size partials (mirrors Vulkan's group aggregate). + if (lid.x == 0u) { + var m = partials[0]; + for (var t = 1u; t < wg_size; t = t + 1u) { + m = max(m, partials[t]); + } + output[row] = m; } - output[row] = acc; } diff --git a/backends/webgpu/runtime/ops/amax/amax_wgsl.h b/backends/webgpu/runtime/ops/amax/amax_wgsl.h index 1f82d7507cd..48ec8f20e27 100644 --- a/backends/webgpu/runtime/ops/amax/amax_wgsl.h +++ b/backends/webgpu/runtime/ops/amax/amax_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from amax.wgsl - DO NOT EDIT. -// wgsl-sha256: aae058ed0c432ac0cb54ea894e03e12a246d0d1c743d9ebb995b9773029e4652 +// wgsl-sha256: 35fc059d7c72caa17f9cb1128823ecfd8f75be4ce24b6cd4f9629a97b52f64c0 inline constexpr const char* kAmaxWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -26,20 +26,43 @@ struct Params { override wg_size: u32 = 256u; +// Cooperative shared-memory reduction; mirrors Vulkan reduce.glsl (a group of +// threads co-operates per reduction row, partials aggregated in shared memory). +// Fixed upper bound (>= any clamped wg_size); only [0, wg_size) is used. +var partials: array; + @compute @workgroup_size(wg_size) fn main( - @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - let row = gid.x + gid.y * (num_workgroups.x * wg_size); + // One workgroup per reduction row; 2D-fold lifts the 65535 grid cap. + let row = wid.x + wid.y * num_workgroups.x; if (row >= params.num_rows) { return; } let base = row * params.reduce_size; + + // Each thread reduces a strided slice of the row into a partial. Seed with + // the row's first element (always valid; reduce_size >= 1) so threads that + // own no element contribute a real value, not an out-of-range identity. var acc = input[base]; - for (var j = 1u; j < params.reduce_size; j = j + 1u) { - acc = max(acc, input[base + j]); + var i = lid.x; + while (i < params.reduce_size) { + acc = max(acc, input[base + i]); + i = i + wg_size; + } + partials[lid.x] = acc; + workgroupBarrier(); + + // Thread 0 aggregates the wg_size partials (mirrors Vulkan's group aggregate). + if (lid.x == 0u) { + var m = partials[0]; + for (var t = 1u; t < wg_size; t = t + 1u) { + m = max(m, partials[t]); + } + output[row] = m; } - output[row] = acc; } )"; diff --git a/backends/webgpu/runtime/ops/amin/Reduce.cpp b/backends/webgpu/runtime/ops/amin/Reduce.cpp index 35c7017436d..ade0ef5b583 100644 --- a/backends/webgpu/runtime/ops/amin/Reduce.cpp +++ b/backends/webgpu/runtime/ops/amin/Reduce.cpp @@ -19,6 +19,10 @@ namespace executorch::backends::webgpu { +static_assert( + kAminWorkgroupSizeX <= 256, + "amin.wgsl partials[] is sized 256; wg_size must not exceed it"); + namespace { // Uniform layout matching the WGSL Params struct; 16-byte aligned. @@ -47,6 +51,9 @@ void amin_impl(WebGPUGraph& graph, const std::vector& args) { const std::vector& dims = graph.get_int_list(args.at(1)); const int64_t ndim = static_cast(in_tensor.dims.size()); + if (ndim <= 0) { + throw std::runtime_error("amin: input must have at least one dim"); + } if (dims.size() != 1 || (dims[0] != -1 && dims[0] != ndim - 1)) { throw std::runtime_error("amin: only last-dim reduction is supported"); } @@ -61,8 +68,9 @@ void amin_impl(WebGPUGraph& graph, const std::vector& args) { } uint32_t wg_size = utils::clamp_workgroup_size(device, kAminWorkgroupSizeX); + // One workgroup per row (cooperative reduction); grid = num_rows workgroups. utils::WgCount workgroup_count = - utils::compute_2d_workgroup_count(device, num_rows, wg_size, "amin"); + utils::compute_2d_workgroup_count(device, num_rows, 1, "amin"); WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; @@ -155,6 +163,9 @@ void amin_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error("amin(resize): zero reduce dim"); } const uint64_t total = utils::numel_of(d); + if (total % rsize != 0u) { + throw std::runtime_error("amin(resize): numel not divisible by dim"); + } const uint32_t rows = static_cast(total / rsize); std::vector od = d; if (keepdim) { @@ -168,7 +179,7 @@ void amin_impl(WebGPUGraph& graph, const std::vector& args) { p.reduce_size = rsize; wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); const utils::WgCount wgc = utils::compute_2d_workgroup_count( - g.device(), rows, wg_size, "amin"); + g.device(), rows, 1, "amin"); g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }); diff --git a/backends/webgpu/runtime/ops/amin/amin.wgsl b/backends/webgpu/runtime/ops/amin/amin.wgsl index d8f4346d231..4778800ab3d 100644 --- a/backends/webgpu/runtime/ops/amin/amin.wgsl +++ b/backends/webgpu/runtime/ops/amin/amin.wgsl @@ -9,18 +9,41 @@ struct Params { override wg_size: u32 = 256u; +// Cooperative shared-memory reduction; mirrors Vulkan reduce.glsl (a group of +// threads co-operates per reduction row, partials aggregated in shared memory). +// Fixed upper bound (>= any clamped wg_size); only [0, wg_size) is used. +var partials: array; + @compute @workgroup_size(wg_size) fn main( - @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - let row = gid.x + gid.y * (num_workgroups.x * wg_size); + // One workgroup per reduction row; 2D-fold lifts the 65535 grid cap. + let row = wid.x + wid.y * num_workgroups.x; if (row >= params.num_rows) { return; } let base = row * params.reduce_size; + + // Each thread reduces a strided slice of the row into a partial. Seed with + // the row's first element (always valid; reduce_size >= 1) so threads that + // own no element contribute a real value, not an out-of-range identity. var acc = input[base]; - for (var j = 1u; j < params.reduce_size; j = j + 1u) { - acc = min(acc, input[base + j]); + var i = lid.x; + while (i < params.reduce_size) { + acc = min(acc, input[base + i]); + i = i + wg_size; + } + partials[lid.x] = acc; + workgroupBarrier(); + + // Thread 0 aggregates the wg_size partials (mirrors Vulkan's group aggregate). + if (lid.x == 0u) { + var m = partials[0]; + for (var t = 1u; t < wg_size; t = t + 1u) { + m = min(m, partials[t]); + } + output[row] = m; } - output[row] = acc; } diff --git a/backends/webgpu/runtime/ops/amin/amin_wgsl.h b/backends/webgpu/runtime/ops/amin/amin_wgsl.h index 53f94e2f0e8..40a97c67a63 100644 --- a/backends/webgpu/runtime/ops/amin/amin_wgsl.h +++ b/backends/webgpu/runtime/ops/amin/amin_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from amin.wgsl - DO NOT EDIT. -// wgsl-sha256: 974a28fd80f089c8a52bf54d73f3bd03c195b2f9d7904bb4c31e6545813e0459 +// wgsl-sha256: 8cb6035ae4d34eb2a6cc973d93d9847905722e967239c96033fccfe3a1943cb2 inline constexpr const char* kAminWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -26,20 +26,43 @@ struct Params { override wg_size: u32 = 256u; +// Cooperative shared-memory reduction; mirrors Vulkan reduce.glsl (a group of +// threads co-operates per reduction row, partials aggregated in shared memory). +// Fixed upper bound (>= any clamped wg_size); only [0, wg_size) is used. +var partials: array; + @compute @workgroup_size(wg_size) fn main( - @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - let row = gid.x + gid.y * (num_workgroups.x * wg_size); + // One workgroup per reduction row; 2D-fold lifts the 65535 grid cap. + let row = wid.x + wid.y * num_workgroups.x; if (row >= params.num_rows) { return; } let base = row * params.reduce_size; + + // Each thread reduces a strided slice of the row into a partial. Seed with + // the row's first element (always valid; reduce_size >= 1) so threads that + // own no element contribute a real value, not an out-of-range identity. var acc = input[base]; - for (var j = 1u; j < params.reduce_size; j = j + 1u) { - acc = min(acc, input[base + j]); + var i = lid.x; + while (i < params.reduce_size) { + acc = min(acc, input[base + i]); + i = i + wg_size; + } + partials[lid.x] = acc; + workgroupBarrier(); + + // Thread 0 aggregates the wg_size partials (mirrors Vulkan's group aggregate). + if (lid.x == 0u) { + var m = partials[0]; + for (var t = 1u; t < wg_size; t = t + 1u) { + m = min(m, partials[t]); + } + output[row] = m; } - output[row] = acc; } )"; diff --git a/backends/webgpu/runtime/ops/argmax/Reduce.cpp b/backends/webgpu/runtime/ops/argmax/Reduce.cpp index 636cae19b42..27318ba0ab9 100644 --- a/backends/webgpu/runtime/ops/argmax/Reduce.cpp +++ b/backends/webgpu/runtime/ops/argmax/Reduce.cpp @@ -29,6 +29,12 @@ struct ArgReduceParams { uint32_t _pad; }; +// The cooperative reduction uses var part_val/part_idx arrays of +// width 256; the workgroup size must not exceed that. +static_assert( + kArgReduceWorkgroupSizeX <= 256, + "arg_reduce workgroup size exceeds the 256-wide shared partials"); + // Last-dim argmax/argmin -> int32 index; mirrors Vulkan arg_reduce_impl. void arg_reduce_impl( WebGPUGraph& graph, @@ -54,6 +60,9 @@ void arg_reduce_impl( } const int64_t ndim = static_cast(in_tensor.dims.size()); + if (ndim <= 0) { + throw std::runtime_error("arg_reduce: input must have at least one dim"); + } const int64_t dim = graph.get_int(args.at(1)); if (dim != -1 && dim != ndim - 1) { throw std::runtime_error( @@ -71,8 +80,9 @@ void arg_reduce_impl( uint32_t wg_size = utils::clamp_workgroup_size(device, kArgReduceWorkgroupSizeX); + // One workgroup per row (cooperative reduction); grid = num_rows workgroups. utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( - device, num_rows, wg_size, "arg_reduce"); + device, num_rows, 1, "arg_reduce"); WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; @@ -184,7 +194,7 @@ void arg_reduce_impl( p.is_argmin = is_argmin; wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); const utils::WgCount wgc = utils::compute_2d_workgroup_count( - g.device(), rows, wg_size, "arg_reduce"); + g.device(), rows, 1, "arg_reduce"); g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }); diff --git a/backends/webgpu/runtime/ops/argmax/arg_reduce.wgsl b/backends/webgpu/runtime/ops/argmax/arg_reduce.wgsl index ef77e9bffd8..098d55b23c5 100644 --- a/backends/webgpu/runtime/ops/argmax/arg_reduce.wgsl +++ b/backends/webgpu/runtime/ops/argmax/arg_reduce.wgsl @@ -11,31 +11,55 @@ struct Params { override wg_size: u32 = 64u; +// Cooperative shared-memory arg-reduction; mirrors Vulkan reduce.glsl. Each +// thread scans a strided slice for its local extremum (strict compare -> first +// index within the slice), then thread 0 aggregates the partials, breaking ties +// by lowest index = torch argmax/argmin semantics. +var part_val: array; +var part_idx: array; + @compute @workgroup_size(wg_size, 1, 1) fn main( - @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - let row = gid.x + gid.y * (num_workgroups.x * wg_size); + // One workgroup per reduction row; 2D-fold lifts the 65535 grid cap. + let row = wid.x + wid.y * num_workgroups.x; if (row >= params.num_rows) { return; } let base = row * params.reduce_size; - // Strict compare keeps the FIRST extremum = torch argmax/argmin tie-break. + + // Seed with element 0 (real value, lowest index); threads owning no element + // keep it, which is harmless since index 0 is a valid candidate. var best = t_in[base]; var best_idx: u32 = 0u; - for (var k: u32 = 1u; k < params.reduce_size; k = k + 1u) { + var k = lid.x; + while (k < params.reduce_size) { let v = t_in[base + k]; if (params.is_argmin != 0u) { - if (v < best) { - best = v; - best_idx = k; - } + if (v < best) { best = v; best_idx = k; } } else { - if (v > best) { - best = v; - best_idx = k; + if (v > best) { best = v; best_idx = k; } + } + k = k + wg_size; + } + part_val[lid.x] = best; + part_idx[lid.x] = best_idx; + workgroupBarrier(); + + if (lid.x == 0u) { + var bv = part_val[0]; + var bi = part_idx[0]; + for (var t: u32 = 1u; t < wg_size; t = t + 1u) { + let v = part_val[t]; + let idx = part_idx[t]; + if (params.is_argmin != 0u) { + if (v < bv || (v == bv && idx < bi)) { bv = v; bi = idx; } + } else { + if (v > bv || (v == bv && idx < bi)) { bv = v; bi = idx; } } } + t_out[row] = bi; } - t_out[row] = best_idx; } diff --git a/backends/webgpu/runtime/ops/argmax/arg_reduce_wgsl.h b/backends/webgpu/runtime/ops/argmax/arg_reduce_wgsl.h index 3ec5403acbe..b98fa79945f 100644 --- a/backends/webgpu/runtime/ops/argmax/arg_reduce_wgsl.h +++ b/backends/webgpu/runtime/ops/argmax/arg_reduce_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from arg_reduce.wgsl - DO NOT EDIT. -// wgsl-sha256: 0ee2f2b47cb29c1b26339314f6ede857efb0f558e53c6433bb3aae984952e1dd +// wgsl-sha256: 5f9f5c8f1bd83164b9e93a9b014f558fe427c0405fc4d8ab683433d72a040bb1 inline constexpr const char* kArgReduceWGSL = R"( @group(0) @binding(0) var t_in: array; @group(0) @binding(1) var t_out: array; @@ -28,33 +28,57 @@ struct Params { override wg_size: u32 = 64u; +// Cooperative shared-memory arg-reduction; mirrors Vulkan reduce.glsl. Each +// thread scans a strided slice for its local extremum (strict compare -> first +// index within the slice), then thread 0 aggregates the partials, breaking ties +// by lowest index = torch argmax/argmin semantics. +var part_val: array; +var part_idx: array; + @compute @workgroup_size(wg_size, 1, 1) fn main( - @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - let row = gid.x + gid.y * (num_workgroups.x * wg_size); + // One workgroup per reduction row; 2D-fold lifts the 65535 grid cap. + let row = wid.x + wid.y * num_workgroups.x; if (row >= params.num_rows) { return; } let base = row * params.reduce_size; - // Strict compare keeps the FIRST extremum = torch argmax/argmin tie-break. + + // Seed with element 0 (real value, lowest index); threads owning no element + // keep it, which is harmless since index 0 is a valid candidate. var best = t_in[base]; var best_idx: u32 = 0u; - for (var k: u32 = 1u; k < params.reduce_size; k = k + 1u) { + var k = lid.x; + while (k < params.reduce_size) { let v = t_in[base + k]; if (params.is_argmin != 0u) { - if (v < best) { - best = v; - best_idx = k; - } + if (v < best) { best = v; best_idx = k; } } else { - if (v > best) { - best = v; - best_idx = k; + if (v > best) { best = v; best_idx = k; } + } + k = k + wg_size; + } + part_val[lid.x] = best; + part_idx[lid.x] = best_idx; + workgroupBarrier(); + + if (lid.x == 0u) { + var bv = part_val[0]; + var bi = part_idx[0]; + for (var t: u32 = 1u; t < wg_size; t = t + 1u) { + let v = part_val[t]; + let idx = part_idx[t]; + if (params.is_argmin != 0u) { + if (v < bv || (v == bv && idx < bi)) { bv = v; bi = idx; } + } else { + if (v > bv || (v == bv && idx < bi)) { bv = v; bi = idx; } } } + t_out[row] = bi; } - t_out[row] = best_idx; } )"; diff --git a/backends/webgpu/runtime/ops/native_group_norm/GroupNorm.cpp b/backends/webgpu/runtime/ops/native_group_norm/GroupNorm.cpp index 71ec4a22a43..19be6a89a11 100644 --- a/backends/webgpu/runtime/ops/native_group_norm/GroupNorm.cpp +++ b/backends/webgpu/runtime/ops/native_group_norm/GroupNorm.cpp @@ -37,6 +37,12 @@ static_assert( sizeof(GroupNormParams) == 32, "GroupNormParams must match the WGSL Params struct (32 bytes)"); +// The reduce shader's cooperative reduction uses var f32[256] arrays; +// the workgroup size (used for both passes) must not exceed that width. +static_assert( + kGroupNormWorkgroupSizeX <= 256, + "group_norm workgroup size exceeds the 256-wide shared reduction arrays"); + struct GnBinding { WGPUBuffer buffer; uint64_t size; @@ -176,11 +182,13 @@ void native_group_norm_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error("group_norm: mean/rstd size != N * group"); } - float eps = std::numeric_limits::epsilon(); + float eps; if (graph.get_value_type(eps_id) == WebGPUGraph::ValueType::Double) { eps = static_cast(graph.get_double(eps_id)); } else if (graph.get_value_type(eps_id) == WebGPUGraph::ValueType::Int) { eps = static_cast(graph.get_int(eps_id)); + } else { + throw std::runtime_error("group_norm: unexpected eps value type"); } GroupNormParams params = {}; @@ -203,9 +211,9 @@ void native_group_norm_impl(WebGPUGraph& graph, const std::vector& args) { const WGPUBufferBindingType rw = WGPUBufferBindingType_Storage; const WGPUBufferBindingType uni = WGPUBufferBindingType_Uniform; - // Pass 1: reduce -> mean/rstd (one thread per (n, group)). + // Pass 1: reduce -> mean/rstd (one workgroup per (n, group), cooperative). utils::WgCount reduce_wgc = utils::compute_2d_workgroup_count( - device, mean_numel, wg_size, "gn_reduce"); + device, mean_numel, 1, "gn_reduce"); const size_t reduce_idx = add_gn_dispatch( graph, device, @@ -257,19 +265,29 @@ void native_group_norm_impl(WebGPUGraph& graph, const std::vector& args) { const uint32_t nb = static_cast(d[0]); const uint32_t c = static_cast(d[1]); const uint32_t hw = static_cast(d[2] * d[3]); + if (c % num_groups != 0) { + throw std::runtime_error( + "group_norm(resize): C not divisible by group"); + } const uint32_t dpg = c / num_groups; + const uint64_t numel64 = static_cast(nb) * c * hw; + const uint64_t group_size64 = static_cast(dpg) * hw; + if (numel64 > std::numeric_limits::max() || + group_size64 > std::numeric_limits::max()) { + throw std::runtime_error("group_norm(resize): numel exceeds uint32"); + } GroupNormParams p = {}; p.n_channels = c; p.hxw = hw; p.num_groups = num_groups; p.chans_per_group = dpg; - p.numel = nb * c * hw; + p.numel = static_cast(numel64); p.mean_numel = nb * num_groups; - p.group_size = dpg * hw; + p.group_size = static_cast(group_size64); p.eps = eps; wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p)); const utils::WgCount rwgc = utils::compute_2d_workgroup_count( - g.device(), p.mean_numel, wg_size, "gn_reduce(resize)"); + g.device(), p.mean_numel, 1, "gn_reduce(resize)"); g.dispatch_at(reduce_idx).workgroup_count_x = rwgc.x; g.dispatch_at(reduce_idx).workgroup_count_y = rwgc.y; const utils::WgCount nwgc = utils::compute_2d_workgroup_count( diff --git a/backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce.wgsl b/backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce.wgsl index 6b0569dc213..69ad9b8bbc2 100644 --- a/backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce.wgsl +++ b/backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce.wgsl @@ -16,12 +16,19 @@ struct Params { override wg_size: u32 = 64u; +// Cooperative shared-memory reduction; mirrors Vulkan group_norm_reduce. Threads +// co-operate per (n, group) to accumulate sum and sum-of-squares. +// Fixed upper bound (>= any clamped wg_size); only [0, wg_size) is used. +var psum: array; +var pss: array; + @compute @workgroup_size(wg_size, 1, 1) fn main( - @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - // One thread per (n, group) -> mean/rstd of shape [N, G]. - let mg = gid.x + gid.y * (num_workgroups.x * wg_size); + // One workgroup per (n, group) -> mean/rstd of shape [N, G]. + let mg = wid.x + wid.y * num_workgroups.x; if (mg >= params.mean_numel) { return; } @@ -34,14 +41,28 @@ fn main( var s = 0.0; var ss = 0.0; - for (var i = 0u; i < params.group_size; i = i + 1u) { + var i = lid.x; + while (i < params.group_size) { let v = input[base + i]; s = s + v; ss = ss + v * v; + i = i + wg_size; + } + psum[lid.x] = s; + pss[lid.x] = ss; + workgroupBarrier(); + + if (lid.x == 0u) { + var ts = psum[0]; + var tss = pss[0]; + for (var t = 1u; t < wg_size; t = t + 1u) { + ts = ts + psum[t]; + tss = tss + pss[t]; + } + let count = f32(params.group_size); + let m = ts / count; + let variance = tss / count - m * m; + mean[mg] = m; + rstd[mg] = inverseSqrt(variance + params.eps); } - let count = f32(params.group_size); - let m = s / count; - let variance = ss / count - m * m; - mean[mg] = m; - rstd[mg] = inverseSqrt(variance + params.eps); } diff --git a/backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce_wgsl.h b/backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce_wgsl.h index 77704092a20..d2ecff1a210 100644 --- a/backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce_wgsl.h +++ b/backends/webgpu/runtime/ops/native_group_norm/group_norm_reduce_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from group_norm_reduce.wgsl - DO NOT EDIT. -// wgsl-sha256: 63e9e15a934f19ce0649323ae3820f1cd86deb9aeb0dbb4ba1de9ab4d7d74688 +// wgsl-sha256: 6765e6dc7eace87dffb1b42b08cbde9b49d60c96074025b708073104075da60a inline constexpr const char* kGroupNormReduceWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var mean: array; @@ -33,12 +33,19 @@ struct Params { override wg_size: u32 = 64u; +// Cooperative shared-memory reduction; mirrors Vulkan group_norm_reduce. Threads +// co-operate per (n, group) to accumulate sum and sum-of-squares. +// Fixed upper bound (>= any clamped wg_size); only [0, wg_size) is used. +var psum: array; +var pss: array; + @compute @workgroup_size(wg_size, 1, 1) fn main( - @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - // One thread per (n, group) -> mean/rstd of shape [N, G]. - let mg = gid.x + gid.y * (num_workgroups.x * wg_size); + // One workgroup per (n, group) -> mean/rstd of shape [N, G]. + let mg = wid.x + wid.y * num_workgroups.x; if (mg >= params.mean_numel) { return; } @@ -51,16 +58,30 @@ fn main( var s = 0.0; var ss = 0.0; - for (var i = 0u; i < params.group_size; i = i + 1u) { + var i = lid.x; + while (i < params.group_size) { let v = input[base + i]; s = s + v; ss = ss + v * v; + i = i + wg_size; + } + psum[lid.x] = s; + pss[lid.x] = ss; + workgroupBarrier(); + + if (lid.x == 0u) { + var ts = psum[0]; + var tss = pss[0]; + for (var t = 1u; t < wg_size; t = t + 1u) { + ts = ts + psum[t]; + tss = tss + pss[t]; + } + let count = f32(params.group_size); + let m = ts / count; + let variance = tss / count - m * m; + mean[mg] = m; + rstd[mg] = inverseSqrt(variance + params.eps); } - let count = f32(params.group_size); - let m = s / count; - let variance = ss / count - m * m; - mean[mg] = m; - rstd[mg] = inverseSqrt(variance + params.eps); } )"; diff --git a/backends/webgpu/runtime/ops/reduce/Reduce.cpp b/backends/webgpu/runtime/ops/reduce/Reduce.cpp index 7e185d38f82..8ee7f1f9cfd 100644 --- a/backends/webgpu/runtime/ops/reduce/Reduce.cpp +++ b/backends/webgpu/runtime/ops/reduce/Reduce.cpp @@ -13,6 +13,7 @@ #include +#include #include #include #include @@ -31,28 +32,53 @@ struct ReduceParams { }; static_assert(sizeof(ReduceParams) == 16, "ReduceParams must be 16 bytes"); -void decompose( +// Normalize + validate CONTIGUOUS reduced dims, then fold the tensor into +// [outer, r, inner] where r spans the reduced range. The kernel reduces r +// elements per (outer, inner) output — identical for 1 or N contiguous dims +// (e.g. global avg pool over [H,W] -> r = H*W). Non-contiguous multi-dim reduce +// is rejected (rare; would need a separate gather). +void decompose_dims( const std::vector& dims, - int64_t dim, + std::vector reduce_dims, uint32_t& outer, uint32_t& r, uint32_t& inner) { const int64_t ndim = static_cast(dims.size()); - if (dim < 0) { - dim += ndim; - } - if (ndim == 0 || dim < 0 || dim >= ndim) { + if (ndim == 0 || reduce_dims.empty()) { throw std::runtime_error("WebGPU reduce: dim out of range"); } - uint64_t o = 1, in = 1; - for (int64_t d = 0; d < dim; ++d) { + for (int64_t& d : reduce_dims) { + if (d < 0) { + d += ndim; + } + if (d < 0 || d >= ndim) { + throw std::runtime_error("WebGPU reduce: dim out of range"); + } + } + std::sort(reduce_dims.begin(), reduce_dims.end()); + for (size_t i = 1; i < reduce_dims.size(); ++i) { + if (reduce_dims[i] == reduce_dims[i - 1]) { + throw std::runtime_error("WebGPU reduce: duplicate reduced dim"); + } + if (reduce_dims[i] != reduce_dims[i - 1] + 1) { + throw std::runtime_error( + "WebGPU reduce: only contiguous reduced dims supported"); + } + } + const int64_t first_rd = reduce_dims.front(); + const int64_t last_rd = reduce_dims.back(); + uint64_t o = 1, rr = 1, in = 1; + for (int64_t d = 0; d < first_rd; ++d) { o *= static_cast(dims[d]); } - for (int64_t d = dim + 1; d < ndim; ++d) { + for (int64_t d = first_rd; d <= last_rd; ++d) { + rr *= static_cast(dims[d]); + } + for (int64_t d = last_rd + 1; d < ndim; ++d) { in *= static_cast(dims[d]); } outer = static_cast(o); - r = static_cast(dims[dim]); + r = static_cast(rr); inner = static_cast(in); } @@ -81,16 +107,15 @@ void reduce_impl( if (graph.get_value_type(dim_id) != WebGPUGraph::ValueType::IntList) { throw std::runtime_error("WebGPU reduce: dim arg is not an IntList"); } - const std::vector& reduce_dims = graph.get_int_list(dim_id); - // Single-dim reduction only for now; multi-dim is a tracked extension. - if (reduce_dims.size() != 1) { - throw std::runtime_error( - "WebGPU reduce: only single-dim reduction is supported"); + // Contiguous multi-dim reduction (e.g. global avg pool over [H,W]) folds into + // one [outer, r, inner]; single-dim is the r = one-dim case. + const std::vector reduce_dims = graph.get_int_list(dim_id); + if (reduce_dims.empty()) { + throw std::runtime_error("WebGPU reduce: empty dim list"); } - const int64_t dim = reduce_dims[0]; uint32_t outer = 0, r = 0, inner = 0; - decompose(in.dims, dim, outer, r, inner); + decompose_dims(in.dims, reduce_dims, outer, r, inner); if (outer == 0 || r == 0 || inner == 0) { throw std::runtime_error("WebGPU reduce: zero-sized reduction"); } @@ -165,7 +190,7 @@ void reduce_impl( in_id, [in_id, out_id, - dim, + reduce_dims, keepdim, is_mean_u, build_outputs, @@ -173,7 +198,8 @@ void reduce_impl( params_buf](WebGPUGraph& g) { const auto& d = g.cur_dims(in_id); uint32_t o = 0, rr = 0, n = 0; - decompose(std::vector(d.begin(), d.end()), dim, o, rr, n); + decompose_dims( + std::vector(d.begin(), d.end()), reduce_dims, o, rr, n); if (o == 0u || rr == 0u || n == 0u) { throw std::runtime_error("WebGPU reduce: live zero-sized reduction"); } @@ -197,10 +223,16 @@ void reduce_impl( g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; // Propagate reduced output dims for downstream resize hooks. int64_t nd = static_cast(d.size()); - int64_t rd = dim < 0 ? dim + nd : dim; + std::vector is_rd(static_cast(nd), false); + for (int64_t rd : reduce_dims) { + const int64_t n2 = rd < 0 ? rd + nd : rd; + if (n2 >= 0 && n2 < nd) { + is_rd[static_cast(n2)] = true; + } + } std::vector od; for (int64_t i = 0; i < nd; ++i) { - if (i == rd) { + if (is_rd[static_cast(i)]) { if (keepdim) { od.push_back(1); } diff --git a/backends/webgpu/test/tester.py b/backends/webgpu/test/tester.py index ed23f01ef8e..5bd13567570 100644 --- a/backends/webgpu/test/tester.py +++ b/backends/webgpu/test/tester.py @@ -30,6 +30,11 @@ exir_ops.edge.aten.slice_copy.Tensor, exir_ops.edge.aten.permute_copy.default, exir_ops.edge.aten.cat.default, + exir_ops.edge.aten.amax.default, + exir_ops.edge.aten.amin.default, + exir_ops.edge.aten.argmax.default, + exir_ops.edge.aten.argmin.default, + exir_ops.edge.aten.native_group_norm.default, ] From ddae6a72b9b402dd7d3d12f6185bb572d6b5bb38 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:41 -0700 Subject: [PATCH 68/75] [ExecuTorch][WebGPU] Add creation/cast ops (full family, scalar_tensor, _to_copy) + enable unary activation family MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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(x))` / `bitcast(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/) --- backends/webgpu/runtime/ops/fill/Fill.cpp | 36 +++++++++++++++++++++-- backends/webgpu/test/tester.py | 22 ++++++++++++++ 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/backends/webgpu/runtime/ops/fill/Fill.cpp b/backends/webgpu/runtime/ops/fill/Fill.cpp index 8da3b31577d..50a8eb32733 100644 --- a/backends/webgpu/runtime/ops/fill/Fill.cpp +++ b/backends/webgpu/runtime/ops/fill/Fill.cpp @@ -13,6 +13,7 @@ #include +#include #include #include @@ -49,8 +50,19 @@ void add_fill( if (out_tensor.buffer == nullptr) { throw std::runtime_error(std::string(op_name) + ": null output buffer"); } - uint32_t num_elements = - static_cast(out_tensor.nbytes / sizeof(float)); + // fp32-only: the shader writes an f32 fill_value and numel below assumes a + // 4-byte element; a bool/int output would be miscounted + bit-wrong -> + // reject. + if (out_tensor.is_int || out_tensor.elem_size != sizeof(float)) { + throw std::runtime_error( + std::string(op_name) + ": only fp32 output is supported"); + } + const uint64_t numel = out_tensor.nbytes / sizeof(float); + if (numel == 0 || numel > UINT32_MAX) { + throw std::runtime_error( + std::string(op_name) + ": output numel is zero or exceeds u32"); + } + uint32_t num_elements = static_cast(numel); uint32_t wg_size = utils::clamp_workgroup_size(device, kFillWorkgroupSizeX); utils::WgCount workgroup_count = @@ -126,6 +138,22 @@ void full_like_impl(WebGPUGraph& graph, const std::vector& args) { "full_like"); } +void zeros_impl(WebGPUGraph& graph, const std::vector& args) { + add_fill(graph, args.at(args.size() - 1), 0.0f, "zeros"); +} + +void zeros_like_impl(WebGPUGraph& graph, const std::vector& args) { + add_fill(graph, args.at(args.size() - 1), 0.0f, "zeros_like"); +} + +void ones_impl(WebGPUGraph& graph, const std::vector& args) { + add_fill(graph, args.at(args.size() - 1), 1.0f, "ones"); +} + +void ones_like_impl(WebGPUGraph& graph, const std::vector& args) { + add_fill(graph, args.at(args.size() - 1), 1.0f, "ones_like"); +} + void scalar_tensor_impl(WebGPUGraph& graph, const std::vector& args) { add_fill( graph, @@ -139,6 +167,10 @@ void scalar_tensor_impl(WebGPUGraph& graph, const std::vector& args) { WEBGPU_REGISTER_OPERATORS { WEBGPU_REGISTER_OP(aten.full.default, full_impl); WEBGPU_REGISTER_OP(aten.full_like.default, full_like_impl); + WEBGPU_REGISTER_OP(aten.zeros.default, zeros_impl); + WEBGPU_REGISTER_OP(aten.zeros_like.default, zeros_like_impl); + WEBGPU_REGISTER_OP(aten.ones.default, ones_impl); + WEBGPU_REGISTER_OP(aten.ones_like.default, ones_like_impl); WEBGPU_REGISTER_OP(aten.scalar_tensor.default, scalar_tensor_impl); } diff --git a/backends/webgpu/test/tester.py b/backends/webgpu/test/tester.py index 5bd13567570..74382148016 100644 --- a/backends/webgpu/test/tester.py +++ b/backends/webgpu/test/tester.py @@ -35,6 +35,28 @@ exir_ops.edge.aten.argmax.default, exir_ops.edge.aten.argmin.default, exir_ops.edge.aten.native_group_norm.default, + exir_ops.edge.aten.full.default, + exir_ops.edge.aten.full_like.default, + exir_ops.edge.aten.zeros.default, + exir_ops.edge.aten.zeros_like.default, + exir_ops.edge.aten.ones.default, + exir_ops.edge.aten.ones_like.default, + exir_ops.edge.aten.scalar_tensor.default, + exir_ops.edge.aten._to_copy.default, + exir_ops.edge.dim_order_ops._to_dim_order_copy.default, + exir_ops.edge.aten.abs.default, + exir_ops.edge.aten.exp.default, + exir_ops.edge.aten.sqrt.default, + exir_ops.edge.aten.rsqrt.default, + exir_ops.edge.aten.sin.default, + exir_ops.edge.aten.cos.default, + exir_ops.edge.aten.tanh.default, + exir_ops.edge.aten.round.default, + exir_ops.edge.aten.neg.default, + exir_ops.edge.aten.hardswish.default, + exir_ops.edge.aten.clamp.default, + exir_ops.edge.aten.hardtanh.default, + exir_ops.edge.aten.pow.Tensor_Scalar, ] From a5652abaa24513be1e45403ec4d4ba65deb1d9c4 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:42 -0700 Subject: [PATCH 69/75] [ExecuTorch][WebGPU] Enable + golden granular ops (elementwise/bool, shape/movement, pooling, conv/sampling) Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/runtime/ops/linear/Linear.cpp | 42 +++++++++++++++---- .../runtime/ops/linear/linear_tiled.wgsl | 9 +++- .../runtime/ops/linear/linear_tiled_wgsl.h | 11 +++-- .../runtime/ops/linear/linear_vec4.wgsl | 9 +++- .../runtime/ops/linear/linear_vec4_wgsl.h | 11 +++-- .../runtime/ops/max_pool2d/MaxPool2d.cpp | 23 +++++++++- backends/webgpu/test/tester.py | 14 +++++++ 7 files changed, 100 insertions(+), 19 deletions(-) diff --git a/backends/webgpu/runtime/ops/linear/Linear.cpp b/backends/webgpu/runtime/ops/linear/Linear.cpp index 7b0de24918d..97f0591d101 100644 --- a/backends/webgpu/runtime/ops/linear/Linear.cpp +++ b/backends/webgpu/runtime/ops/linear/Linear.cpp @@ -27,18 +27,21 @@ struct LinearParams { uint32_t M; uint32_t N; uint32_t K; - uint32_t pad_; + uint32_t has_bias; }; static_assert(sizeof(LinearParams) == 16, "LinearParams must be 16 bytes"); constexpr uint32_t kTile = 32u; -// aten.linear (no bias); shared-memory tiled GEMM. +// aten.linear (+ optional bias); shared-memory tiled GEMM. void linear_impl(WebGPUGraph& graph, const std::vector& args) { - // out is the last arg; bias (if present) is unused on this path. + // args: [input, weight, bias?, out]; out is last. bias (arg 2) is a tensor + // when present, Null/absent otherwise; add_bias[col] is fused in the shader. const int in_id = args.at(0); const int w_id = args.at(1); const int out_id = args.at(args.size() - 1); + const bool has_bias = args.size() >= 4 && + graph.get_value_type(args.at(2)) == WebGPUGraph::ValueType::Tensor; WGPUDevice device = graph.device(); const auto& in = graph.get_tensor(in_id); @@ -72,10 +75,26 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error("WebGPU linear: tile grid exceeds dispatch limit"); } + if (has_bias) { + const auto& b = graph.get_tensor(args.at(2)); + if (b.nbytes != static_cast(N) * sizeof(float)) { + throw std::runtime_error("WebGPU linear: bias must be [N] fp32"); + } + } + LinearParams params = {}; params.M = M; params.N = N; params.K = K; + params.has_bias = has_bias ? 1u : 0u; + + // Bias binding (binding 4); a 4-byte dummy satisfies it when None + // (WGSL-gated). + utils::OptionalBinding bias = utils::make_optional_binding( + device, + has_bias, + has_bias ? graph.get_tensor(args.at(2)).buffer : nullptr, + has_bias ? graph.get_tensor(args.at(2)).nbytes : 0); WGPUBufferDescriptor uniform_desc = {}; uniform_desc.size = sizeof(LinearParams); @@ -97,7 +116,7 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { shader_desc.nextInChain = &wgsl_desc.chain; WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - WGPUBindGroupLayoutEntry entries[4] = {}; + WGPUBindGroupLayoutEntry entries[5] = {}; entries[0].binding = 0; entries[0].visibility = WGPUShaderStage_Compute; entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; @@ -110,9 +129,12 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { entries[3].binding = 3; entries[3].visibility = WGPUShaderStage_Compute; entries[3].buffer.type = WGPUBufferBindingType_Uniform; + entries[4].binding = 4; + entries[4].visibility = WGPUShaderStage_Compute; + entries[4].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; + bgl_desc.entryCount = 5; bgl_desc.entries = entries; WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); @@ -130,7 +152,7 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { WGPUComputePipeline pipeline = wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - WGPUBindGroupEntry bg_entries[4] = {}; + WGPUBindGroupEntry bg_entries[5] = {}; bg_entries[0].binding = 0; bg_entries[0].buffer = in.buffer; bg_entries[0].size = in.nbytes; @@ -143,12 +165,18 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { bg_entries[3].binding = 3; bg_entries[3].buffer = uniform_buffer; bg_entries[3].size = sizeof(LinearParams); + bg_entries[4].binding = 4; + bg_entries[4].buffer = bias.buffer; + bg_entries[4].size = bias.nbytes; WGPUBindGroupDescriptor bg_desc = {}; bg_desc.layout = bgl; - bg_desc.entryCount = 4; + bg_desc.entryCount = 5; bg_desc.entries = bg_entries; WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + if (bias.owned_dummy != nullptr) { + wgpuBufferRelease(bias.owned_dummy); + } WebGPUDispatch dispatch; dispatch.pipeline = pipeline; diff --git a/backends/webgpu/runtime/ops/linear/linear_tiled.wgsl b/backends/webgpu/runtime/ops/linear/linear_tiled.wgsl index 2b94814b75e..be9354da6ba 100644 --- a/backends/webgpu/runtime/ops/linear/linear_tiled.wgsl +++ b/backends/webgpu/runtime/ops/linear/linear_tiled.wgsl @@ -2,13 +2,14 @@ struct Params { M: u32, N: u32, K: u32, - pad_: u32, + has_bias: u32, }; @group(0) @binding(0) var input: array; @group(0) @binding(1) var weight: array; @group(0) @binding(2) var out: array; @group(0) @binding(3) var params: Params; +@group(0) @binding(4) var bias: array; const TILE: u32 = 32u; const RPT: u32 = 4u; @@ -75,7 +76,11 @@ fn main( let r = tile_row0 + tile_row + ir; let c = tile_col0 + tile_col + ic; if (r < params.M && c < params.N) { - out[r * params.N + c] = acc[ir][ic]; + var v = acc[ir][ic]; + if (params.has_bias != 0u) { + v = v + bias[c]; + } + out[r * params.N + c] = v; } } } diff --git a/backends/webgpu/runtime/ops/linear/linear_tiled_wgsl.h b/backends/webgpu/runtime/ops/linear/linear_tiled_wgsl.h index 880479b2301..0264d4645cc 100644 --- a/backends/webgpu/runtime/ops/linear/linear_tiled_wgsl.h +++ b/backends/webgpu/runtime/ops/linear/linear_tiled_wgsl.h @@ -13,19 +13,20 @@ namespace executorch::backends::webgpu { // @generated from linear_tiled.wgsl - DO NOT EDIT. -// wgsl-sha256: bb054c7ab0937a24089df4944001666c902c3c4b31e9555ec916bf357c2104f7 +// wgsl-sha256: 723ff884cd63485561079eaab4c93d604b938372ce79733402f17494b25a8e91 inline constexpr const char* kLinearTiledWGSL = R"( struct Params { M: u32, N: u32, K: u32, - pad_: u32, + has_bias: u32, }; @group(0) @binding(0) var input: array; @group(0) @binding(1) var weight: array; @group(0) @binding(2) var out: array; @group(0) @binding(3) var params: Params; +@group(0) @binding(4) var bias: array; const TILE: u32 = 32u; const RPT: u32 = 4u; @@ -92,7 +93,11 @@ fn main( let r = tile_row0 + tile_row + ir; let c = tile_col0 + tile_col + ic; if (r < params.M && c < params.N) { - out[r * params.N + c] = acc[ir][ic]; + var v = acc[ir][ic]; + if (params.has_bias != 0u) { + v = v + bias[c]; + } + out[r * params.N + c] = v; } } } diff --git a/backends/webgpu/runtime/ops/linear/linear_vec4.wgsl b/backends/webgpu/runtime/ops/linear/linear_vec4.wgsl index d3cfaa97d24..a404f94d724 100644 --- a/backends/webgpu/runtime/ops/linear/linear_vec4.wgsl +++ b/backends/webgpu/runtime/ops/linear/linear_vec4.wgsl @@ -2,13 +2,14 @@ struct Params { M: u32, N: u32, K: u32, - pad_: u32, + has_bias: u32, }; @group(0) @binding(0) var input: array>; @group(0) @binding(1) var weight: array>; @group(0) @binding(2) var out: array; @group(0) @binding(3) var params: Params; +@group(0) @binding(4) var bias: array; const TILE: u32 = 32u; const TILE4: u32 = 8u; @@ -76,7 +77,11 @@ fn main( let r = row0 + tr + ir; let c = col0 + tc + ic; if (r < params.M && c < params.N) { - out[r * params.N + c] = acc[ir][ic]; + var v = acc[ir][ic]; + if (params.has_bias != 0u) { + v = v + bias[c]; + } + out[r * params.N + c] = v; } } } diff --git a/backends/webgpu/runtime/ops/linear/linear_vec4_wgsl.h b/backends/webgpu/runtime/ops/linear/linear_vec4_wgsl.h index 361735ea63a..057536f2a3a 100644 --- a/backends/webgpu/runtime/ops/linear/linear_vec4_wgsl.h +++ b/backends/webgpu/runtime/ops/linear/linear_vec4_wgsl.h @@ -13,19 +13,20 @@ namespace executorch::backends::webgpu { // @generated from linear_vec4.wgsl - DO NOT EDIT. -// wgsl-sha256: 2d456961a1145cf9858237deb59fd0b37f390dfc253894b240914d2cbe056b66 +// wgsl-sha256: 900bc546260898ede5281dae5cfced8a5712c2d011d0e0e2742bf19350796dc8 inline constexpr const char* kLinearVec4WGSL = R"( struct Params { M: u32, N: u32, K: u32, - pad_: u32, + has_bias: u32, }; @group(0) @binding(0) var input: array>; @group(0) @binding(1) var weight: array>; @group(0) @binding(2) var out: array; @group(0) @binding(3) var params: Params; +@group(0) @binding(4) var bias: array; const TILE: u32 = 32u; const TILE4: u32 = 8u; @@ -93,7 +94,11 @@ fn main( let r = row0 + tr + ir; let c = col0 + tc + ic; if (r < params.M && c < params.N) { - out[r * params.N + c] = acc[ir][ic]; + var v = acc[ir][ic]; + if (params.has_bias != 0u) { + v = v + bias[c]; + } + out[r * params.N + c] = v; } } } diff --git a/backends/webgpu/runtime/ops/max_pool2d/MaxPool2d.cpp b/backends/webgpu/runtime/ops/max_pool2d/MaxPool2d.cpp index 682eff61197..ea6bbfe35e7 100644 --- a/backends/webgpu/runtime/ops/max_pool2d/MaxPool2d.cpp +++ b/backends/webgpu/runtime/ops/max_pool2d/MaxPool2d.cpp @@ -101,6 +101,14 @@ void max_pool2d_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error("WebGPU max_pool2d: zero stride"); } + // ceil_mode (arg 5) rounds the output size up. The shader already skips + // out-of-bounds window positions, so only the host output size changes. + bool ceil_mode = false; + if (args.size() > 5 && + graph.get_value_type(args.at(5)) == WebGPUGraph::ValueType::Bool) { + ceil_mode = graph.get_bool(args.at(5)); + } + const int64_t oh_num = int64_t(IH) + 2 * int64_t(pH) - int64_t(dH) * (int64_t(kH) - 1) - 1; const int64_t ow_num = @@ -108,8 +116,19 @@ void max_pool2d_impl(WebGPUGraph& graph, const std::vector& args) { if (oh_num < 0 || ow_num < 0) { throw std::runtime_error("WebGPU max_pool2d: kernel larger than input"); } - const uint32_t OH = static_cast(oh_num / int64_t(sH)) + 1; - const uint32_t OW = static_cast(ow_num / int64_t(sW)) + 1; + const int64_t addH = ceil_mode ? int64_t(sH) - 1 : 0; + const int64_t addW = ceil_mode ? int64_t(sW) - 1 : 0; + uint32_t OH = static_cast((oh_num + addH) / int64_t(sH)) + 1; + uint32_t OW = static_cast((ow_num + addW) / int64_t(sW)) + 1; + // ceil_mode: drop a trailing window that begins entirely in the padding. + if (ceil_mode) { + if ((int64_t(OH) - 1) * int64_t(sH) >= int64_t(IH) + int64_t(pH)) { + OH--; + } + if ((int64_t(OW) - 1) * int64_t(sW) >= int64_t(IW) + int64_t(pW)) { + OW--; + } + } // Validate against the serialized values output [N, C, OH, OW] (loud-fail if // the arg interpretation is wrong, e.g. ceil_mode or a different layout). diff --git a/backends/webgpu/test/tester.py b/backends/webgpu/test/tester.py index 74382148016..24c6a480b81 100644 --- a/backends/webgpu/test/tester.py +++ b/backends/webgpu/test/tester.py @@ -57,6 +57,20 @@ exir_ops.edge.aten.clamp.default, exir_ops.edge.aten.hardtanh.default, exir_ops.edge.aten.pow.Tensor_Scalar, + exir_ops.edge.aten.minimum.default, + exir_ops.edge.aten.div.Tensor_mode, + exir_ops.edge.aten.logical_and.default, + exir_ops.edge.aten.bitwise_and.Tensor, + exir_ops.edge.aten.bitwise_not.default, + exir_ops.edge.aten.flip.default, + exir_ops.edge.aten.repeat.default, + exir_ops.edge.aten.index_select.default, + exir_ops.edge.aten.avg_pool2d.default, + exir_ops.edge.aten.pixel_shuffle.default, + exir_ops.edge.aten.convolution.default, + exir_ops.edge.et_vk.conv_with_clamp.default, + exir_ops.edge.aten.grid_sampler_2d.default, + exir_ops.edge.et_vk.grid_priors.default, ] From bfcf42f290c828225569a25caeea16977bee1170 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:42 -0700 Subject: [PATCH 70/75] [ExecuTorch][WebGPU] Enable + golden int8 quantized ops (q8ta family + qcs4w/q8ta_q8csw linear) Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/test/tester.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/backends/webgpu/test/tester.py b/backends/webgpu/test/tester.py index 24c6a480b81..bee5d13f057 100644 --- a/backends/webgpu/test/tester.py +++ b/backends/webgpu/test/tester.py @@ -71,6 +71,16 @@ exir_ops.edge.et_vk.conv_with_clamp.default, exir_ops.edge.aten.grid_sampler_2d.default, exir_ops.edge.et_vk.grid_priors.default, + exir_ops.edge.et_vk.linear_qcs4w.default, + exir_ops.edge.et_vk.linear_q8ta_q8csw.default, + exir_ops.edge.et_vk.q8ta_add.default, + exir_ops.edge.et_vk.q8ta_relu.default, + exir_ops.edge.et_vk.q8ta_pixel_shuffle.default, + exir_ops.edge.et_vk.q8ta_linear.default, + exir_ops.edge.et_vk.q8ta_linear_gemv.default, + exir_ops.edge.et_vk.q8ta_conv2d.default, + exir_ops.edge.et_vk.q8ta_conv2d_dw.default, + exir_ops.edge.et_vk.q8ta_conv2d_pw.default, ] From 9a47856c3fa31b78447edb489641a2a610318916 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:43 -0700 Subject: [PATCH 71/75] [ExecuTorch][WebGPU] Add linear_dq8ca_q4gsw + choose_qparams_affine (8da4w) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../ChooseQparamsAffine.cpp | 240 ++++++++++++++++ .../choose_qparams_affine.wgsl | 115 ++++++++ .../choose_qparams_affine_wgsl.h | 139 +++++++++ .../linear_dq8ca_q4gsw/LinearDq8caQ4gsw.cpp | 272 ++++++++++++++++++ .../linear_dq8ca_q4gsw.wgsl | 120 ++++++++ .../linear_dq8ca_q4gsw_wgsl.h | 144 ++++++++++ backends/webgpu/test/op_tests/cases.py | 28 ++ .../test/ops/test_linear_dq8ca_q4gsw.py | 38 +++ backends/webgpu/test/tester.py | 2 + 9 files changed, 1098 insertions(+) create mode 100644 backends/webgpu/runtime/ops/choose_qparams_affine/ChooseQparamsAffine.cpp create mode 100644 backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine.wgsl create mode 100644 backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine_wgsl.h create mode 100644 backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/LinearDq8caQ4gsw.cpp create mode 100644 backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/linear_dq8ca_q4gsw.wgsl create mode 100644 backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/linear_dq8ca_q4gsw_wgsl.h create mode 100644 backends/webgpu/test/ops/test_linear_dq8ca_q4gsw.py diff --git a/backends/webgpu/runtime/ops/choose_qparams_affine/ChooseQparamsAffine.cpp b/backends/webgpu/runtime/ops/choose_qparams_affine/ChooseQparamsAffine.cpp new file mode 100644 index 00000000000..d896cd13409 --- /dev/null +++ b/backends/webgpu/runtime/ops/choose_qparams_affine/ChooseQparamsAffine.cpp @@ -0,0 +1,240 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct ChooseQParamsParams { + uint32_t num_rows; + uint32_t reduce_size; + int32_t quant_min; + int32_t quant_max; +}; +static_assert( + sizeof(ChooseQParamsParams) == 16, + "ChooseQParamsParams must match the WGSL Params struct (16 bytes)"); + +// torchao.choose_qparams_affine args (mirrors Vulkan ChooseQParams.cpp:158): +// [input, mapping_type, block_size, target_dtype, quant_min, quant_max, eps, +// scale_dtype, zero_point_dtype, out_tuple(scale, zp)]. Routes to the per-row +// (last-dim) path: one workgroup per row computes asymmetric scale/zp. +void choose_qparams_affine_impl( + WebGPUGraph& graph, + const std::vector& args) { + const int in_id = args.at(0); + const int out_list_id = args.at(args.size() - 1); + + const std::vector& out_ids = graph.get_value_list(out_list_id); + if (out_ids.size() != 2) { + throw std::runtime_error( + "choose_qparams_affine: expected 2 outputs (scale, zp)"); + } + const int scale_id = out_ids.at(0); + const int zp_id = out_ids.at(1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(scale_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(zp_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("choose_qparams_affine: in/scale/zp not tensor"); + } + + WGPUDevice device = graph.device(); + const auto& in = graph.get_tensor(in_id); + const auto& scale_t = graph.get_tensor(scale_id); + const auto& zp_t = graph.get_tensor(zp_id); + if (in.buffer == nullptr || scale_t.buffer == nullptr || + zp_t.buffer == nullptr) { + throw std::runtime_error("choose_qparams_affine: null buffer binding"); + } + if (in.dims.empty()) { + throw std::runtime_error("choose_qparams_affine: input has no dims"); + } + + const uint64_t reduce_size = static_cast(in.dims.back()); + if (reduce_size == 0) { + throw std::runtime_error("choose_qparams_affine: last dim == 0"); + } + uint64_t in_numel = 1; + for (int64_t d : in.dims) { + in_numel *= static_cast(d); + } + const uint64_t num_rows = in_numel / reduce_size; + if (num_rows == 0 || num_rows > UINT32_MAX || reduce_size > UINT32_MAX) { + throw std::runtime_error("choose_qparams_affine: bad row/reduce shape"); + } + if (in.nbytes != in_numel * sizeof(float)) { + throw std::runtime_error("choose_qparams_affine: input must be fp32"); + } + // scale is fp32[num_rows]; zp is int8[num_rows] (bound as array). + if (scale_t.nbytes != num_rows * sizeof(float)) { + throw std::runtime_error("choose_qparams_affine: scale must be fp32[rows]"); + } + // zp is int8[rows] (elem_size 1), packed 4-per-u32 in the shader. int8 buffers + // are allocated max(nbytes, 4); M<=4 pads to one word, M%4==0 is word-exact. + // Other M (5,6,7,...) would overflow the M-byte buffer -> reject. + if (!zp_t.is_int8 || zp_t.nbytes != num_rows) { + throw std::runtime_error("choose_qparams_affine: zp must be int8[rows]"); + } + if (num_rows > 4 && num_rows % 4 != 0) { + throw std::runtime_error( + "choose_qparams_affine: num_rows must be <=4 or a multiple of 4"); + } + + // The kernel implements only the asymmetric, per-row (last-dim), int8 path; + // validate the schema args it assumes and fail loud rather than silently + // ignoring them (mirrors Vulkan, which consumes block_size + quant_min/max). + const int quant_min_id = args.at(4); + const int quant_max_id = args.at(5); + if (graph.get_value_type(quant_min_id) != WebGPUGraph::ValueType::Int || + graph.get_value_type(quant_max_id) != WebGPUGraph::ValueType::Int) { + throw std::runtime_error( + "choose_qparams_affine: quant_min/quant_max must be int scalars"); + } + const int64_t quant_min = graph.get_int(quant_min_id); + const int64_t quant_max = graph.get_int(quant_max_id); + if (quant_min != -128 || quant_max != 127) { + throw std::runtime_error( + "choose_qparams_affine: only the int8 range [-128, 127] is supported"); + } + // Per-row fast path: block_size must be [1, ..., 1, reduce_size]. + const int block_size_id = args.at(2); + if (graph.get_value_type(block_size_id) != WebGPUGraph::ValueType::IntList) { + throw std::runtime_error( + "choose_qparams_affine: block_size must be an int list"); + } + const std::vector& block_size = graph.get_int_list(block_size_id); + if (block_size.empty() || + block_size.back() != static_cast(reduce_size)) { + throw std::runtime_error( + "choose_qparams_affine: block_size must reduce the last dim"); + } + for (size_t d = 0; d + 1 < block_size.size(); ++d) { + if (block_size[d] != 1) { + throw std::runtime_error( + "choose_qparams_affine: only per-row (last-dim) blocks are supported"); + } + } + + ChooseQParamsParams params = {}; + params.num_rows = static_cast(num_rows); + params.reduce_size = static_cast(reduce_size); + params.quant_min = static_cast(quant_min); + params.quant_max = static_cast(quant_max); + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kChooseQparamsAffineWorkgroupSizeX); + // One workgroup per block of 4 rows (wg_size threads cooperate per row); the + // block packs its 4 int8 zps into one u32. 2D-fold lifts the 65535 grid cap. + const uint32_t num_blocks = static_cast((num_rows + 3) / 4); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, num_blocks, 1, "choose_qparams_affine"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(ChooseQParamsParams)); + graph.add_uniform_buffer_bytes(sizeof(ChooseQParamsParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kChooseQparamsAffineWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[4] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Storage; + entries[3].binding = 3; + entries[3].visibility = WGPUShaderStage_Compute; + entries[3].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 4; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg[4] = {}; + bg[0].binding = 0; + bg[0].buffer = in.buffer; + bg[0].size = in.nbytes; + bg[1].binding = 1; + bg[1].buffer = scale_t.buffer; + bg[1].size = scale_t.nbytes; + bg[2].binding = 2; + bg[2].buffer = zp_t.buffer; + // Bind word-aligned (buffer is >= max(nbytes,4); array needs a mult of 4). + bg[2].size = ((zp_t.nbytes + 3u) / 4u) * 4u; + bg[3].binding = 3; + bg[3].buffer = params_buf; + bg[3].size = sizeof(ChooseQParamsParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 4; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "choose_qparams_affine", + workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP( + torchao.choose_qparams_affine.default, + choose_qparams_affine_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine.wgsl b/backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine.wgsl new file mode 100644 index 00000000000..ff0c9002ada --- /dev/null +++ b/backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine.wgsl @@ -0,0 +1,115 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var scales_out: array; +@group(0) @binding(2) var zps_out: array; + +struct Params { + num_rows: u32, + reduce_size: u32, + quant_min: i32, + quant_max: i32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256u; + +// Cooperative per-row min/max (mirrors Vulkan choose_qparams_per_row.glsl). One +// workgroup handles a block of up to 4 rows so the 4 int8 zero-points pack into a +// single u32 word with no cross-workgroup write race (zp is int8, elem_size 1). +// Fixed upper bound >= any clamped wg_size; only [0, wg_size) is used. +var part_min: array; +var part_max: array; + +const SMALL_SCALE_THRESHOLD: f32 = 6.1e-5; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One workgroup per block of 4 rows; 2D-fold lifts the 65535 grid cap. + let block = wid.x + wid.y * num_workgroups.x; + let row0 = block * 4u; + if (row0 >= params.num_rows) { + return; + } + let lim = min(4u, params.num_rows - row0); + let qmin = f32(params.quant_min); + let qmax = f32(params.quant_max); + + var packed_zp: u32 = 0u; + for (var r: u32 = 0u; r < lim; r = r + 1u) { + let row = row0 + r; + let base = row * params.reduce_size; + + // Seed with the row's first element so idle threads contribute a real value. + var lmin = input[base]; + var lmax = input[base]; + var i = lid.x; + while (i < params.reduce_size) { + let v = input[base + i]; + lmin = min(lmin, v); + lmax = max(lmax, v); + i = i + wg_size; + } + part_min[lid.x] = lmin; + part_max[lid.x] = lmax; + workgroupBarrier(); + + if (lid.x == 0u) { + var mn = part_min[0]; + var mx = part_max[0]; + for (var t = 1u; t < wg_size; t = t + 1u) { + mn = min(mn, part_min[t]); + mx = max(mx, part_max[t]); + } + + // calculate_scale_and_zero_point (mirrors Vulkan): extend [min,max] to + // contain 0, asymmetric scale, error-selected + nudged integer zp. + mn = min(mn, 0.0); + mx = max(mx, 0.0); + var scale = (mx - mn) / (qmax - qmin); + if (scale == 0.0) { + scale = 0.1; + } + if (scale < SMALL_SCALE_THRESHOLD) { + let org_scale = scale; + scale = SMALL_SCALE_THRESHOLD; + if (mn == 0.0) { + mx = SMALL_SCALE_THRESHOLD * (qmax - qmin); + } else if (mx == 0.0) { + mn = -SMALL_SCALE_THRESHOLD * (qmax - qmin); + } else { + let amplifier = SMALL_SCALE_THRESHOLD / org_scale; + mn = mn * amplifier; + mx = mx * amplifier; + } + } + + let zp_from_min = qmin - mn / scale; + let zp_from_max = qmax - mx / scale; + let zp_from_min_error = abs(qmin) - abs(mn / scale); + let zp_from_max_error = abs(qmax) - abs(mx / scale); + var initial_zp = zp_from_max; + if (zp_from_min_error < zp_from_max_error) { + initial_zp = zp_from_min; + } + var nudged: i32; + if (initial_zp < qmin) { + nudged = params.quant_min; + } else if (initial_zp > qmax) { + nudged = params.quant_max; + } else { + nudged = i32(round(initial_zp)); + } + + scales_out[row] = scale; + packed_zp = packed_zp | ((u32(nudged) & 0xFFu) << (r * 8u)); + } + // Ensure part_min/part_max are fully consumed before the next row reuses them. + workgroupBarrier(); + } + + if (lid.x == 0u) { + zps_out[block] = packed_zp; + } +} diff --git a/backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine_wgsl.h b/backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine_wgsl.h new file mode 100644 index 00000000000..bf2b027b906 --- /dev/null +++ b/backends/webgpu/runtime/ops/choose_qparams_affine/choose_qparams_affine_wgsl.h @@ -0,0 +1,139 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from choose_qparams_affine.wgsl - DO NOT EDIT. +// wgsl-sha256: 45b55ec7c432d5fbd9a7c9716b569c9c654e7ea566b6530c4b78f2924daa116f +inline constexpr const char* kChooseQparamsAffineWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var scales_out: array; +@group(0) @binding(2) var zps_out: array; + +struct Params { + num_rows: u32, + reduce_size: u32, + quant_min: i32, + quant_max: i32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256u; + +// Cooperative per-row min/max (mirrors Vulkan choose_qparams_per_row.glsl). One +// workgroup handles a block of up to 4 rows so the 4 int8 zero-points pack into a +// single u32 word with no cross-workgroup write race (zp is int8, elem_size 1). +// Fixed upper bound >= any clamped wg_size; only [0, wg_size) is used. +var part_min: array; +var part_max: array; + +const SMALL_SCALE_THRESHOLD: f32 = 6.1e-5; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // One workgroup per block of 4 rows; 2D-fold lifts the 65535 grid cap. + let block = wid.x + wid.y * num_workgroups.x; + let row0 = block * 4u; + if (row0 >= params.num_rows) { + return; + } + let lim = min(4u, params.num_rows - row0); + let qmin = f32(params.quant_min); + let qmax = f32(params.quant_max); + + var packed_zp: u32 = 0u; + for (var r: u32 = 0u; r < lim; r = r + 1u) { + let row = row0 + r; + let base = row * params.reduce_size; + + // Seed with the row's first element so idle threads contribute a real value. + var lmin = input[base]; + var lmax = input[base]; + var i = lid.x; + while (i < params.reduce_size) { + let v = input[base + i]; + lmin = min(lmin, v); + lmax = max(lmax, v); + i = i + wg_size; + } + part_min[lid.x] = lmin; + part_max[lid.x] = lmax; + workgroupBarrier(); + + if (lid.x == 0u) { + var mn = part_min[0]; + var mx = part_max[0]; + for (var t = 1u; t < wg_size; t = t + 1u) { + mn = min(mn, part_min[t]); + mx = max(mx, part_max[t]); + } + + // calculate_scale_and_zero_point (mirrors Vulkan): extend [min,max] to + // contain 0, asymmetric scale, error-selected + nudged integer zp. + mn = min(mn, 0.0); + mx = max(mx, 0.0); + var scale = (mx - mn) / (qmax - qmin); + if (scale == 0.0) { + scale = 0.1; + } + if (scale < SMALL_SCALE_THRESHOLD) { + let org_scale = scale; + scale = SMALL_SCALE_THRESHOLD; + if (mn == 0.0) { + mx = SMALL_SCALE_THRESHOLD * (qmax - qmin); + } else if (mx == 0.0) { + mn = -SMALL_SCALE_THRESHOLD * (qmax - qmin); + } else { + let amplifier = SMALL_SCALE_THRESHOLD / org_scale; + mn = mn * amplifier; + mx = mx * amplifier; + } + } + + let zp_from_min = qmin - mn / scale; + let zp_from_max = qmax - mx / scale; + let zp_from_min_error = abs(qmin) - abs(mn / scale); + let zp_from_max_error = abs(qmax) - abs(mx / scale); + var initial_zp = zp_from_max; + if (zp_from_min_error < zp_from_max_error) { + initial_zp = zp_from_min; + } + var nudged: i32; + if (initial_zp < qmin) { + nudged = params.quant_min; + } else if (initial_zp > qmax) { + nudged = params.quant_max; + } else { + nudged = i32(round(initial_zp)); + } + + scales_out[row] = scale; + packed_zp = packed_zp | ((u32(nudged) & 0xFFu) << (r * 8u)); + } + // Ensure part_min/part_max are fully consumed before the next row reuses them. + workgroupBarrier(); + } + + if (lid.x == 0u) { + zps_out[block] = packed_zp; + } +} +)"; + +inline constexpr uint32_t kChooseQparamsAffineWorkgroupSizeX = 256; +inline constexpr uint32_t kChooseQparamsAffineWorkgroupSizeY = 1; +inline constexpr uint32_t kChooseQparamsAffineWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/LinearDq8caQ4gsw.cpp b/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/LinearDq8caQ4gsw.cpp new file mode 100644 index 00000000000..96f456923a0 --- /dev/null +++ b/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/LinearDq8caQ4gsw.cpp @@ -0,0 +1,272 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct Dq8caParams { + uint32_t M; + uint32_t N; + uint32_t K; + uint32_t K_packed; + uint32_t group_size; + uint32_t padded_N; + uint32_t has_bias; + uint32_t _pad; +}; +static_assert(sizeof(Dq8caParams) == 32, "Dq8caParams must be 32 bytes"); + +constexpr int64_t kTileM = 4; // MUST match TM in linear_dq8ca_q4gsw.wgsl +constexpr int64_t kTileN = 4; // MUST match TN + +// et_vk.linear_dq8ca_q4gsw args (mirrors Vulkan QuantizedLinear.cpp:760): +// [in, input_scale, input_zp, weight, weight_sums, weight_scales, group_size, +// bias, out]. Dynamic per-row int8 activation quant x 4-bit-group symmetric +// weight. weight_sums (arg 4) is a perf shortcut; this v1 recomputes the sum +// inline so it is intentionally unused. Static-shape only (no resize hook yet). +void linear_dq8ca_q4gsw_impl( + WebGPUGraph& graph, + const std::vector& args) { + const int in_id = args.at(0); + const int input_scale_id = args.at(1); + const int input_zp_id = args.at(2); + const int weight_id = args.at(3); + const int scales_id = args.at(5); + const int group_size_id = args.at(6); + const int bias_id = args.at(7); + const int out_id = args.at(8); + + WGPUDevice device = graph.device(); + const auto& in = graph.get_tensor(in_id); + const auto& input_scale = graph.get_tensor(input_scale_id); + const auto& input_zp = graph.get_tensor(input_zp_id); + const auto& weight = graph.get_tensor(weight_id); + const auto& scales = graph.get_tensor(scales_id); + const auto& out = graph.get_tensor(out_id); + + if (in.dims.empty() || weight.dims.size() < 2 || scales.dims.size() < 2) { + throw std::runtime_error("linear_dq8ca_q4gsw: malformed dims"); + } + if (in.buffer == nullptr || input_scale.buffer == nullptr || + input_zp.buffer == nullptr || weight.buffer == nullptr || + scales.buffer == nullptr || out.buffer == nullptr) { + throw std::runtime_error("linear_dq8ca_q4gsw: null buffer binding"); + } + + const uint32_t K = static_cast(in.dims.back()); + if (K == 0) { + throw std::runtime_error("linear_dq8ca_q4gsw: K == 0"); + } + uint64_t in_numel = 1; + for (int64_t d : in.dims) { + in_numel *= static_cast(d); + } + if (in_numel % K != 0) { + throw std::runtime_error("linear_dq8ca_q4gsw: input numel % K != 0"); + } + const uint32_t M = static_cast(in_numel / K); + const uint32_t N = static_cast(weight.dims[0]); + const uint32_t K_packed = static_cast(weight.dims[1]); + const uint32_t num_groups = static_cast(scales.dims[0]); + const uint32_t padded_N = static_cast(scales.dims[1]); + if (M == 0 || N == 0) { + throw std::runtime_error("linear_dq8ca_q4gsw: M or N == 0"); + } + if (K_packed != (K + 1) / 2) { + throw std::runtime_error("linear_dq8ca_q4gsw: K_packed must be ceil(K/2)"); + } + if ((static_cast(N) * K_packed) % 4u != 0u) { + throw std::runtime_error( + "linear_dq8ca_q4gsw: N*K_packed must be a multiple of 4 (u32-packed)"); + } + + int64_t group_size = 0; + if (graph.get_value_type(group_size_id) == WebGPUGraph::ValueType::Int) { + group_size = graph.get_int(group_size_id); + } + if (group_size <= 0) { + throw std::runtime_error("linear_dq8ca_q4gsw: group_size <= 0"); + } + const uint32_t gs = static_cast(group_size); + + // fp32-only byte guards; per-row scale (f32[M]) + zp (int8[M]). + if (in.nbytes != in_numel * sizeof(float) || + out.nbytes != static_cast(M) * N * sizeof(float) || + scales.nbytes != + static_cast(num_groups) * padded_N * sizeof(float) || + weight.nbytes != static_cast(N) * K_packed) { + throw std::runtime_error("linear_dq8ca_q4gsw: fp32/byte-size mismatch"); + } + // Per-row activation scale (fp32[M]) + zp (int8[M], packed 4/word in-shader). + if (input_scale.nbytes != static_cast(M) * sizeof(float) || + !input_zp.is_int8 || input_zp.nbytes != static_cast(M)) { + throw std::runtime_error( + "linear_dq8ca_q4gsw: input scale fp32[M] / zp int8[M] required"); + } + // int8 zp is bound word-aligned over a max(nbytes,4) buffer; M in {5,6,7,...} + // would bind past the buffer. Mirrors the choose_qparams_affine producer guard. + if (M > 4u && M % 4u != 0u) { + throw std::runtime_error( + "linear_dq8ca_q4gsw: num_rows must be <=4 or a multiple of 4"); + } + if (num_groups < (K + gs - 1u) / gs || padded_N < N) { + throw std::runtime_error("linear_dq8ca_q4gsw: scales dims too small"); + } + + uint32_t has_bias = 0; + WGPUBuffer bias_buffer = nullptr; + uint64_t bias_size = 4; + if (graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor) { + const auto& bias = graph.get_tensor(bias_id); + if (bias.buffer != nullptr && bias.nbytes >= N * sizeof(float)) { + has_bias = 1; + bias_buffer = bias.buffer; + bias_size = bias.nbytes; + } + } + if (bias_buffer == nullptr) { + bias_buffer = graph.create_scratch_buffer(4); + } + + Dq8caParams params = {}; + params.M = M; + params.N = N; + params.K = K; + params.K_packed = K_packed; + params.group_size = gs; + params.padded_N = padded_N; + params.has_bias = has_bias; + + const int64_t total_tiles = utils::div_up(M, kTileM) * + utils::div_up(N, kTileN); + if (total_tiles > static_cast(UINT32_MAX)) { + throw std::runtime_error("linear_dq8ca_q4gsw: tile count exceeds u32"); + } + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kLinearDq8caQ4gswWorkgroupSizeX); + const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, static_cast(total_tiles), wg_size, "linear_dq8ca_q4gsw"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(Dq8caParams)); + graph.add_uniform_buffer_bytes(sizeof(Dq8caParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kLinearDq8caQ4gswWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // 0 out(rw), 1 in, 2 input_scale, 3 input_zp, 4 weight, 5 scales, 6 bias (ro), + // 7 uniform. + WGPUBindGroupLayoutEntry entries[8] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_Storage; + for (uint32_t i = 1; i <= 6; i++) { + entries[i].binding = i; + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + } + entries[7].binding = 7; + entries[7].visibility = WGPUShaderStage_Compute; + entries[7].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 8; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg[8] = {}; + bg[0].binding = 0; + bg[0].buffer = out.buffer; + bg[0].size = out.nbytes; + bg[1].binding = 1; + bg[1].buffer = in.buffer; + bg[1].size = in.nbytes; + bg[2].binding = 2; + bg[2].buffer = input_scale.buffer; + bg[2].size = input_scale.nbytes; + bg[3].binding = 3; + bg[3].buffer = input_zp.buffer; + // int8 zp bound as array; round to a multiple of 4 (buffer is >=4 bytes). + bg[3].size = ((input_zp.nbytes + 3u) / 4u) * 4u; + bg[4].binding = 4; + bg[4].buffer = weight.buffer; + bg[4].size = weight.nbytes; + bg[5].binding = 5; + bg[5].buffer = scales.buffer; + bg[5].size = scales.nbytes; + bg[6].binding = 6; + bg[6].buffer = bias_buffer; + bg[6].size = bias_size; + bg[7].binding = 7; + bg[7].buffer = params_buf; + bg[7].size = sizeof(Dq8caParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 8; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "linear_dq8ca_q4gsw", + workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP( + et_vk.linear_dq8ca_q4gsw.default, + linear_dq8ca_q4gsw_impl); +} + +} // namespace executorch::backends::webgpu + diff --git a/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/linear_dq8ca_q4gsw.wgsl b/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/linear_dq8ca_q4gsw.wgsl new file mode 100644 index 00000000000..81cffed5b0b --- /dev/null +++ b/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/linear_dq8ca_q4gsw.wgsl @@ -0,0 +1,120 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_input: array; +@group(0) @binding(2) var t_input_scale: array; +@group(0) @binding(3) var t_input_zp: array; +@group(0) @binding(4) var t_weight: array; +@group(0) @binding(5) var t_scales: array; +@group(0) @binding(6) var t_bias: array; + +struct Params { + M: u32, + N: u32, + K: u32, + K_packed: u32, + group_size: u32, + padded_N: u32, + has_bias: u32, + _pad: u32, +} +@group(0) @binding(7) var params: Params; + +override wg_size: u32 = 64u; + +// Register-tiled GEMM (mirrors q4gsw_linear.wgsl) with folded dynamic per-row +// int8 activation quant: out[m,n] = s[m] * sum_k (xq[m,k]-z[m]) * dequant(w). +// s[m]/z[m] come from choose_qparams_affine (per-row asymmetric). Weight side is +// the SAME symmetric 4-bit-group packing as linear_q4gsw (nibble - 8). +const TM: u32 = 4u; +const TN: u32 = 4u; +const TILE_ELEMS: u32 = TM * TN; + +// int8 zero-point is packed 4-per-u32 (elem_size 1); extract the signed byte. +fn unpack_zp(idx: u32) -> i32 { + let word = t_input_zp[idx >> 2u]; + return i32(((word >> ((idx & 3u) * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +// xq = clamp(round(x/s) + z, -128, 127); returns (xq - z) as f32 (the value the +// dequant sum uses; s is applied once to the accumulator at the end). +fn quant_act_centered(x: f32, s: f32, z: i32) -> f32 { + let q = clamp(i32(round(x / s)) + z, -128, 127); + return f32(q - z); +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let nrt = (params.M + TM - 1u) / TM; + let nct = (params.N + TN - 1u) / TN; + let tiles = nrt * nct; + // 2D-folded flat tile index (lifts the 65535 1D-dispatch cap on prefill). + let tile = gid.x + gid.y * (num_workgroups.x * wg_size); + if (tile >= tiles) { + return; + } + let row_tile = tile / nct; + let col_tile = tile % nct; + let m0 = row_tile * TM; + let n0 = col_tile * TN; + + // Per-row activation scale/zp (clamp row index for the TM tile overhang). + var s_reg: array; + var z_reg: array; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m_eff = min(m0 + ml, params.M - 1u); + s_reg[ml] = t_input_scale[m_eff]; + z_reg[ml] = unpack_zp(m_eff); + } + + var acc: array; + for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) { + acc[i] = 0.0; + } + + var k: u32 = 0u; + loop { + if (k >= params.K) { + break; + } + // Quantize the TM activation values for column k once; reused across TN cols. + var in_reg: array; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m_eff = min(m0 + ml, params.M - 1u); + in_reg[ml] = + quant_act_centered(t_input[m_eff * params.K + k], s_reg[ml], z_reg[ml]); + } + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n_eff = min(n0 + nl, params.N - 1u); + let byte_idx = n_eff * params.K_packed + (k >> 1u); + let word = t_weight[byte_idx >> 2u]; + let b = (word >> ((byte_idx & 3u) * 8u)) & 0xFFu; + var nib: u32; + if ((k & 1u) == 0u) { + nib = b & 0x0Fu; // even k -> low nibble + } else { + nib = (b >> 4u) & 0x0Fu; // odd k -> high nibble + } + let q = f32(i32(nib) - 8); // +8-shifted on pack; recover signed [-8,7] + let dq = q * t_scales[(k / params.group_size) * params.padded_N + n_eff]; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + acc[ml * TN + nl] = acc[ml * TN + nl] + in_reg[ml] * dq; + } + } + k = k + 1u; + } + + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m = m0 + ml; + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n = n0 + nl; + if (m < params.M && n < params.N) { + var v = s_reg[ml] * acc[ml * TN + nl]; + if (params.has_bias != 0u) { + v = v + t_bias[n]; + } + t_out[m * params.N + n] = v; + } + } + } +} diff --git a/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/linear_dq8ca_q4gsw_wgsl.h b/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/linear_dq8ca_q4gsw_wgsl.h new file mode 100644 index 00000000000..dde1e5258b8 --- /dev/null +++ b/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/linear_dq8ca_q4gsw_wgsl.h @@ -0,0 +1,144 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from linear_dq8ca_q4gsw.wgsl - DO NOT EDIT. +// wgsl-sha256: a3ec7f77b0bbb629b57f5a91a6ab33d96fcce3c7af3e23db1b599926ed04a5df +inline constexpr const char* kLinearDq8caQ4gswWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_input: array; +@group(0) @binding(2) var t_input_scale: array; +@group(0) @binding(3) var t_input_zp: array; +@group(0) @binding(4) var t_weight: array; +@group(0) @binding(5) var t_scales: array; +@group(0) @binding(6) var t_bias: array; + +struct Params { + M: u32, + N: u32, + K: u32, + K_packed: u32, + group_size: u32, + padded_N: u32, + has_bias: u32, + _pad: u32, +} +@group(0) @binding(7) var params: Params; + +override wg_size: u32 = 64u; + +// Register-tiled GEMM (mirrors q4gsw_linear.wgsl) with folded dynamic per-row +// int8 activation quant: out[m,n] = s[m] * sum_k (xq[m,k]-z[m]) * dequant(w). +// s[m]/z[m] come from choose_qparams_affine (per-row asymmetric). Weight side is +// the SAME symmetric 4-bit-group packing as linear_q4gsw (nibble - 8). +const TM: u32 = 4u; +const TN: u32 = 4u; +const TILE_ELEMS: u32 = TM * TN; + +// int8 zero-point is packed 4-per-u32 (elem_size 1); extract the signed byte. +fn unpack_zp(idx: u32) -> i32 { + let word = t_input_zp[idx >> 2u]; + return i32(((word >> ((idx & 3u) * 8u)) & 0xFFu) ^ 0x80u) - 128; +} + +// xq = clamp(round(x/s) + z, -128, 127); returns (xq - z) as f32 (the value the +// dequant sum uses; s is applied once to the accumulator at the end). +fn quant_act_centered(x: f32, s: f32, z: i32) -> f32 { + let q = clamp(i32(round(x / s)) + z, -128, 127); + return f32(q - z); +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let nrt = (params.M + TM - 1u) / TM; + let nct = (params.N + TN - 1u) / TN; + let tiles = nrt * nct; + // 2D-folded flat tile index (lifts the 65535 1D-dispatch cap on prefill). + let tile = gid.x + gid.y * (num_workgroups.x * wg_size); + if (tile >= tiles) { + return; + } + let row_tile = tile / nct; + let col_tile = tile % nct; + let m0 = row_tile * TM; + let n0 = col_tile * TN; + + // Per-row activation scale/zp (clamp row index for the TM tile overhang). + var s_reg: array; + var z_reg: array; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m_eff = min(m0 + ml, params.M - 1u); + s_reg[ml] = t_input_scale[m_eff]; + z_reg[ml] = unpack_zp(m_eff); + } + + var acc: array; + for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) { + acc[i] = 0.0; + } + + var k: u32 = 0u; + loop { + if (k >= params.K) { + break; + } + // Quantize the TM activation values for column k once; reused across TN cols. + var in_reg: array; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m_eff = min(m0 + ml, params.M - 1u); + in_reg[ml] = + quant_act_centered(t_input[m_eff * params.K + k], s_reg[ml], z_reg[ml]); + } + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n_eff = min(n0 + nl, params.N - 1u); + let byte_idx = n_eff * params.K_packed + (k >> 1u); + let word = t_weight[byte_idx >> 2u]; + let b = (word >> ((byte_idx & 3u) * 8u)) & 0xFFu; + var nib: u32; + if ((k & 1u) == 0u) { + nib = b & 0x0Fu; // even k -> low nibble + } else { + nib = (b >> 4u) & 0x0Fu; // odd k -> high nibble + } + let q = f32(i32(nib) - 8); // +8-shifted on pack; recover signed [-8,7] + let dq = q * t_scales[(k / params.group_size) * params.padded_N + n_eff]; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + acc[ml * TN + nl] = acc[ml * TN + nl] + in_reg[ml] * dq; + } + } + k = k + 1u; + } + + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m = m0 + ml; + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n = n0 + nl; + if (m < params.M && n < params.N) { + var v = s_reg[ml] * acc[ml * TN + nl]; + if (params.has_bias != 0u) { + v = v + t_bias[n]; + } + t_out[m * params.N + n] = v; + } + } + } +} +)"; + +inline constexpr uint32_t kLinearDq8caQ4gswWorkgroupSizeX = 64; +inline constexpr uint32_t kLinearDq8caQ4gswWorkgroupSizeY = 1; +inline constexpr uint32_t kLinearDq8caQ4gswWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 9cf0c82e8d8..5d660286cbf 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -87,6 +87,9 @@ from executorch.backends.webgpu.test.ops.test_linear_q8ta_q8csw import ( make_linear_q8ta_q8csw_module, ) +from executorch.backends.webgpu.test.ops.test_linear_dq8ca_q4gsw import ( + make_linear_dq8ca_q4gsw_module, +) from executorch.backends.webgpu.test.ops.test_q8ta_linear import ( make_q8ta_linear_module, ) @@ -1791,6 +1794,31 @@ def case(name, m, k, n, **kw): ) +@register_op_test("linear_dq8ca_q4gsw") +def _linear_dq8ca_q4gsw_suite() -> WebGPUTestSuite: + # torchao Int8DynamicActivationIntxWeightConfig(int4, PerGroup) -> delegated + # choose_qparams_affine (per-row act scale/zp) -> linear_dq8ca_q4gsw (fp32 + # out). Golden = converted eager (fp32 fake-quant). q4gsw needs K % gs == 0, + # K % 8 == 0, N % 8 == 0. + def case(name, m, k, n, **kw): + return Case( + name=name, construct={"k": k, "n": n, "m": m, **kw}, inputs=((m, k),) + ) + + return WebGPUTestSuite( + module_factory=lambda **kw: make_linear_dq8ca_q4gsw_module(**kw), + cases=[ + case("basic", 4, 64, 16), + case("gemv", 1, 64, 16), # M==1 decode + case("k128", 2, 128, 8), # more groups, smaller N + case("n32", 3, 64, 32), # larger N + ], + golden_dtype="float32", + atol=1e-3, + rtol=1e-3, + ) + + @register_op_test("q8ta_linear") def _q8ta_linear_suite() -> WebGPUTestSuite: # XNNPACK-static-quantized nn.Linear -> delegated quantize->q8ta_linear-> diff --git a/backends/webgpu/test/ops/test_linear_dq8ca_q4gsw.py b/backends/webgpu/test/ops/test_linear_dq8ca_q4gsw.py new file mode 100644 index 00000000000..387e7864b29 --- /dev/null +++ b/backends/webgpu/test/ops/test_linear_dq8ca_q4gsw.py @@ -0,0 +1,38 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`et_vk.linear_dq8ca_q4gsw` module + config: dynamic per-row int8 activation +quant x 4-bit-group symmetric weight (the 8da4w path). + +Reached by running a plain `nn.Linear` through torchao's +`Int8DynamicActivationIntxWeightConfig(weight_dtype=int4, PerGroup(gs))`: the +Vulkan partitioner fuses the dynamic-quant + linear into `choose_qparams_affine` +(per-row activation scale/zp) feeding `et_vk.linear_dq8ca_q4gsw` (fp32 out). The +factory returns the CONVERTED module so the op-test framework goldens the WebGPU +output against the converted eager forward (fp32 fake-quant reference). q4gsw +requires K % group_size == 0, K % 8 == 0, N % 8 == 0. +""" + +import torch +import torch.nn as nn + +from torchao.quantization.granularity import PerGroup +from torchao.quantization.quant_api import ( + Int8DynamicActivationIntxWeightConfig, + quantize_, +) + + +def make_linear_dq8ca_q4gsw_module(k, n, m, group_size=32, bias=False, seed=0): + torch.manual_seed(seed) # fixes the weights the golden derives from + lin = nn.Linear(k, n, bias=bias).eval() + quantize_( + lin, + Int8DynamicActivationIntxWeightConfig( + weight_dtype=torch.int4, weight_granularity=PerGroup(group_size) + ), + ) + return lin diff --git a/backends/webgpu/test/tester.py b/backends/webgpu/test/tester.py index bee5d13f057..69d291899b7 100644 --- a/backends/webgpu/test/tester.py +++ b/backends/webgpu/test/tester.py @@ -81,6 +81,8 @@ exir_ops.edge.et_vk.q8ta_conv2d.default, exir_ops.edge.et_vk.q8ta_conv2d_dw.default, exir_ops.edge.et_vk.q8ta_conv2d_pw.default, + exir_ops.edge.et_vk.linear_dq8ca_q4gsw.default, + exir_ops.edge.torchao.choose_qparams_affine.default, ] From fd47129cc6b2a8fb32ee1d18cd0a62cf2f993477 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:43 -0700 Subject: [PATCH 72/75] [ExecuTorch][WebGPU] Add bitwise_or + logical_or ops (bool) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../runtime/ops/logical_or/LogicalOr.cpp | 187 ++++++++++++++++++ .../runtime/ops/logical_or/logical_or.wgsl | 25 +++ .../runtime/ops/logical_or/logical_or_wgsl.h | 49 +++++ backends/webgpu/test/op_tests/cases.py | 171 +++++++++------- backends/webgpu/test/ops/test_logical_or.py | 89 +++++++++ 5 files changed, 455 insertions(+), 66 deletions(-) create mode 100644 backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp create mode 100644 backends/webgpu/runtime/ops/logical_or/logical_or.wgsl create mode 100644 backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h create mode 100644 backends/webgpu/test/ops/test_logical_or.py diff --git a/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp b/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp new file mode 100644 index 00000000000..5fe67cc49d1 --- /dev/null +++ b/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp @@ -0,0 +1,187 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform layout matching the WGSL Params struct; 16-byte aligned. +struct LogicalOrParams { + uint32_t num_words; + uint32_t _pad[3]; +}; +static_assert(sizeof(LogicalOrParams) == 16, "LogicalOrParams must be 16 bytes"); + +// out = a | b (bool OR); serves logical_or + bitwise_or (mirrors Vulkan). +void logical_or_op(WebGPUGraph& graph, const std::vector& args) { + const int a_id = args.at(0); + const int b_id = args.at(1); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(a_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(b_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("logical_or: a/b/out is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& a_tensor = graph.get_tensor(a_id); + const auto& b_tensor = graph.get_tensor(b_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (a_tensor.buffer == nullptr || b_tensor.buffer == nullptr || + out_tensor.buffer == nullptr) { + throw std::runtime_error("logical_or: null buffer binding"); + } + // a/b/out are all 1-byte bool tensors (int-typed, NOT int8-quantized). + if (!a_tensor.is_int || !b_tensor.is_int || !out_tensor.is_int || + a_tensor.elem_size != 1 || b_tensor.elem_size != 1 || + out_tensor.elem_size != 1) { + throw std::runtime_error("logical_or: a/b/out must be 1-byte bool tensors"); + } + const uint64_t numel = out_tensor.nbytes; + // bool packed 4/word (array); numel%4==0 gates the u32 binding. + if (numel == 0u || numel % 4u != 0u || numel > UINT32_MAX) { + throw std::runtime_error("logical_or: numel must be a nonzero mult of 4"); + } + if (a_tensor.nbytes != numel || b_tensor.nbytes != numel) { + throw std::runtime_error("logical_or: a/b/out numel mismatch (same-shape)"); + } + + LogicalOrParams params = {}; + params.num_words = static_cast(numel / 4u); + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kLogicalOrWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, params.num_words, wg_size, "logical_or"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(LogicalOrParams)); + graph.add_uniform_buffer_bytes(sizeof(LogicalOrParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kLogicalOrWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // out (rw storage) + a/b (ro storage) + params (uniform). + WGPUBindGroupLayoutEntry entries[4] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_Storage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[3].binding = 3; + entries[3].visibility = WGPUShaderStage_Compute; + entries[3].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 4; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg[4] = {}; + bg[0].binding = 0; + bg[0].buffer = out_tensor.buffer; + bg[0].size = out_tensor.nbytes; + bg[1].binding = 1; + bg[1].buffer = a_tensor.buffer; + bg[1].size = a_tensor.nbytes; + bg[2].binding = 2; + bg[2].buffer = b_tensor.buffer; + bg[2].size = b_tensor.nbytes; + bg[3].binding = 3; + bg[3].buffer = uniform_buffer; + bg[3].size = sizeof(LogicalOrParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 4; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "logical_or", + workgroup_count.y}); + + // Dynamic shapes: recompute num_words/dispatch; out follows a (same-shape). + WGPUBuffer params_buf = uniform_buffer; + auto resize = + [a_id, b_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(a_id); + const uint64_t n = utils::numel_of(d); + if (n == 0u || n % 4u != 0u || n > UINT32_MAX || + utils::numel_of(g.cur_dims(b_id)) != n) { + throw std::runtime_error( + "logical_or(resize): numel must be a mult of 4"); + } + g.set_cur_dims(out_id, d); + LogicalOrParams p = {}; + p.num_words = static_cast(n / 4u); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), p.num_words, wg_size, "logical_or"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }; + graph.add_tensor_resize_hook(a_id, resize); + graph.add_tensor_resize_hook(b_id, resize); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.logical_or.default, logical_or_op); + WEBGPU_REGISTER_OP(aten.bitwise_or.Tensor, logical_or_op); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl b/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl new file mode 100644 index 00000000000..d7e6176ba32 --- /dev/null +++ b/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl @@ -0,0 +1,25 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_a: array; +@group(0) @binding(2) var t_b: array; + +struct Params { + num_words: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // bool packed 4/word; canonical 0/1 bytes -> word-wise OR == per-byte OR. + let w = gid.x + gid.y * (num_workgroups.x * wg_size); + if (w >= params.num_words) { + return; + } + t_out[w] = t_a[w] | t_b[w]; +} diff --git a/backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h b/backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h new file mode 100644 index 00000000000..d64898cb523 --- /dev/null +++ b/backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h @@ -0,0 +1,49 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from logical_or.wgsl - DO NOT EDIT. +// wgsl-sha256: 4ad19ee04e2c7b396b4669cf44f95133d658c3ec2e6f37d7b271bedc0e582ecf +inline constexpr const char* kLogicalOrWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_a: array; +@group(0) @binding(2) var t_b: array; + +struct Params { + num_words: u32, + pad0: u32, + pad1: u32, + pad2: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // bool packed 4/word; canonical 0/1 bytes -> word-wise OR == per-byte OR. + let w = gid.x + gid.y * (num_workgroups.x * wg_size); + if (w >= params.num_words) { + return; + } + t_out[w] = t_a[w] | t_b[w]; +} +)"; + +inline constexpr uint32_t kLogicalOrWorkgroupSizeX = 64; +inline constexpr uint32_t kLogicalOrWorkgroupSizeY = 1; +inline constexpr uint32_t kLogicalOrWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 5d660286cbf..26b82699e38 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -30,101 +30,99 @@ AddModule, AddSelfModule, ) -from executorch.backends.webgpu.test.ops.test_cat import ( - CatModule, - CONFIGS as _CAT_CONFIGS, -) -from executorch.backends.webgpu.test.ops.test_compare import ( - CompareModule, - compare_gen_a, - compare_gen_b, -) -from executorch.backends.webgpu.test.ops.test_logical_and import ( - la_gen_a, - la_gen_b, - LogicalAndModule, +from executorch.backends.webgpu.test.ops.test_argmax import ( + argmax_tie_gen, + ArgmaxModule, + argmin_tie_gen, + ArgminModule, ) +from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule from executorch.backends.webgpu.test.ops.test_bitwise import ( BitwiseAndModule, BitwiseNotModule, bw_gen_a, bw_gen_b, ) -from executorch.backends.webgpu.test.ops.test_floor_divide import FloorDivideModule -from executorch.backends.webgpu.test.ops.test_argmax import ( - argmax_tie_gen, - argmin_tie_gen, - ArgmaxModule, - ArgminModule, +from executorch.backends.webgpu.test.ops.test_cat import ( + CatModule, + CONFIGS as _CAT_CONFIGS, +) +from executorch.backends.webgpu.test.ops.test_compare import ( + compare_gen_a, + compare_gen_b, + CompareModule, ) -from executorch.backends.webgpu.test.ops.test_flip import FlipModule -from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule -from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule from executorch.backends.webgpu.test.ops.test_conv1d_dw import Conv1dDWModule from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dPwModule +from executorch.backends.webgpu.test.ops.test_conv_with_clamp import ConvWithClampModule +from executorch.backends.webgpu.test.ops.test_flip import FlipModule +from executorch.backends.webgpu.test.ops.test_floor_divide import FloorDivideModule from executorch.backends.webgpu.test.ops.test_grid_priors import GridPriorsModule -from executorch.backends.webgpu.test.ops.test_conv_with_clamp import ( - ConvWithClampModule, -) -from executorch.backends.webgpu.test.ops.test_grid_sampler_2d import ( - GridSampler2dModule, -) -from executorch.backends.webgpu.test.ops.test_rope_interleaved import ( - RopeInterleavedModule, -) -from executorch.backends.webgpu.test.ops.test_quant import ( - DequantizeConstModule, - QuantizeModule, +from executorch.backends.webgpu.test.ops.test_grid_sampler_2d import GridSampler2dModule +from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule +from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule +from executorch.backends.webgpu.test.ops.test_linear_dq8ca_q4gsw import ( + make_linear_dq8ca_q4gsw_module, ) -from executorch.backends.webgpu.test.ops.test_q8ta_add import Q8taAddModule -from executorch.backends.webgpu.test.ops.test_q8ta_relu import Q8taReluModule -from executorch.backends.webgpu.test.ops.test_q8ta_pixel_shuffle import ( - Q8taPixelShuffleModule, +from executorch.backends.webgpu.test.ops.test_linear_q8ta_q8csw import ( + make_linear_q8ta_q8csw_module, ) from executorch.backends.webgpu.test.ops.test_linear_qcs4w import ( make_qcs4w_linear_module, ) -from executorch.backends.webgpu.test.ops.test_linear_q8ta_q8csw import ( - make_linear_q8ta_q8csw_module, +from executorch.backends.webgpu.test.ops.test_logical_and import ( + la_gen_a, + la_gen_b, + LogicalAndModule, ) -from executorch.backends.webgpu.test.ops.test_linear_dq8ca_q4gsw import ( - make_linear_dq8ca_q4gsw_module, +from executorch.backends.webgpu.test.ops.test_logical_or import ( + BitwiseOrModule, + lo_gen_a, + lo_gen_b, + LogicalOrModule, ) -from executorch.backends.webgpu.test.ops.test_q8ta_linear import ( - make_q8ta_linear_module, +from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule +from executorch.backends.webgpu.test.ops.test_mul import ( + CONFIGS as _MUL_CONFIGS, + MulModule, ) -from executorch.backends.webgpu.test.ops.test_q8ta_conv2d_pw import ( - make_q8ta_conv2d_pw_module, +from executorch.backends.webgpu.test.ops.test_permute import ( + CONFIGS as _PERMUTE_CONFIGS, + PermuteModule, ) +from executorch.backends.webgpu.test.ops.test_pixel_shuffle import PixelShuffleModule +from executorch.backends.webgpu.test.ops.test_pow import PowModule +from executorch.backends.webgpu.test.ops.test_q8ta_add import Q8taAddModule +from executorch.backends.webgpu.test.ops.test_q8ta_conv2d import make_q8ta_conv2d_module from executorch.backends.webgpu.test.ops.test_q8ta_conv2d_dw import ( make_q8ta_conv2d_dw_module, ) -from executorch.backends.webgpu.test.ops.test_q8ta_conv2d import ( - make_q8ta_conv2d_module, +from executorch.backends.webgpu.test.ops.test_q8ta_conv2d_pw import ( + make_q8ta_conv2d_pw_module, ) from executorch.backends.webgpu.test.ops.test_q8ta_conv2d_transposed import ( make_q8ta_conv2d_transposed_module, ) -from executorch.backends.webgpu.test.ops.test_pixel_shuffle import PixelShuffleModule -from executorch.backends.webgpu.test.ops.test_group_norm import GroupNormModule -from executorch.backends.webgpu.test.ops.test_index_select import IndexSelectModule -from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule -from executorch.backends.webgpu.test.ops.test_mul import ( - CONFIGS as _MUL_CONFIGS, - MulModule, +from executorch.backends.webgpu.test.ops.test_q8ta_linear import make_q8ta_linear_module +from executorch.backends.webgpu.test.ops.test_q8ta_pixel_shuffle import ( + Q8taPixelShuffleModule, ) -from executorch.backends.webgpu.test.ops.test_permute import ( - CONFIGS as _PERMUTE_CONFIGS, - PermuteModule, +from executorch.backends.webgpu.test.ops.test_q8ta_relu import Q8taReluModule +from executorch.backends.webgpu.test.ops.test_quant import ( + DequantizeConstModule, + QuantizeModule, ) -from executorch.backends.webgpu.test.ops.test_pow import PowModule from executorch.backends.webgpu.test.ops.test_reduce import AmaxModule, AminModule +from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule from executorch.backends.webgpu.test.ops.test_rms_norm import ( _CASES, _linspace_weight, _ramp, RmsNormModule, ) +from executorch.backends.webgpu.test.ops.test_rope_interleaved import ( + RopeInterleavedModule, +) from executorch.backends.webgpu.test.ops.test_select import ( CONFIGS as _SELECT_CONFIGS, SelectModule, @@ -384,6 +382,49 @@ def case(name, shape): ) +@register_op_test("logical_or") +def _logical_or_suite() -> WebGPUTestSuite: + # out = (a>0) || (b>0): two bool masks derived on-GPU from float inputs via + # gt.Tensor (baked zeros), OR'd -> bool. Distinct a/b seeds (~50% each, + # independent -> OR ~75% True, a real mix an AND mutant fails); all shapes + # numel % 4 == 0. float32 oracle (byte-exact bool golden). + def case(name, shape): + return Case( + name=name, + construct={"shape": shape}, + inputs=( + InputSpec(shape=shape, gen=lo_gen_a), + InputSpec(shape=shape, gen=lo_gen_b), + ), + ) + + return WebGPUTestSuite( + module_factory=lambda shape: LogicalOrModule(shape), + cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + golden_dtype="float32", + ) + + +@register_op_test("bitwise_or") +def _bitwise_or_suite() -> WebGPUTestSuite: + # bool bitwise OR == logical_or for canonical 0/1 (shares the handler). + def case(name, shape): + return Case( + name=name, + construct={"shape": shape}, + inputs=( + InputSpec(shape=shape, gen=lo_gen_a), + InputSpec(shape=shape, gen=lo_gen_b), + ), + ) + + return WebGPUTestSuite( + module_factory=lambda shape: BitwiseOrModule(shape), + cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + golden_dtype="float32", + ) + + @register_op_test("pow") def _pow_suite() -> WebGPUTestSuite: # Positive base avoids pow(neg, frac)=NaN; exponent spans negative+positive. @@ -660,17 +701,13 @@ def case(name, k, s, p, cip, ceil_mode, divisor, shape): cases=[ case("basic", [2, 2], [2, 2], [0, 0], True, False, None, (1, 2, 4, 4)), case("pad_cip", [3, 3], [2, 2], [1, 1], True, False, None, (1, 2, 5, 5)), - case( - "pad_nocip", [3, 3], [2, 2], [1, 1], False, False, None, (1, 2, 5, 5) - ), + case("pad_nocip", [3, 3], [2, 2], [1, 1], False, False, None, (1, 2, 5, 5)), case("asym", [3, 2], [2, 3], [1, 1], True, False, None, (2, 3, 5, 7)), case("divisor", [2, 2], [2, 2], [0, 0], True, False, 3, (1, 1, 4, 4)), # ceil_mode: last window overhangs -> exercises the overhang divisor # branch (beh/bew > 0) + the ceil output-size (3x3 vs floor 2x2). case("ceil_cip", [2, 2], [2, 2], [0, 0], True, True, None, (1, 1, 5, 5)), - case( - "ceil_nocip", [3, 3], [2, 2], [0, 0], False, True, None, (1, 2, 5, 5) - ), + case("ceil_nocip", [3, 3], [2, 2], [0, 0], False, True, None, (1, 2, 5, 5)), ], atol=1e-3, rtol=1e-3, @@ -1726,7 +1763,9 @@ def _q8ta_pixel_shuffle_suite() -> WebGPUTestSuite: def case(name, n_ch, h, w, **kw): n = n_ch * h * w # [1, n_ch, h, w], n_ch = C*r*r vals = [((i % 251) - 125) for i in range(n)] # spread across int8 range - return Case(name=name, construct={"x_vals": vals, "shape": (1, n_ch, h, w), **kw}) + return Case( + name=name, construct={"x_vals": vals, "shape": (1, n_ch, h, w), **kw} + ) return WebGPUTestSuite( module_factory=lambda **kw: Q8taPixelShuffleModule(**kw), diff --git a/backends/webgpu/test/ops/test_logical_or.py b/backends/webgpu/test/ops/test_logical_or.py new file mode 100644 index 00000000000..d71c1493ea9 --- /dev/null +++ b/backends/webgpu/test/ops/test_logical_or.py @@ -0,0 +1,89 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""`aten.logical_or.default` / `aten.bitwise_or.Tensor` (bool) modules + configs. + +Mirrors the logical_and/bitwise_and tests: the modules derive their 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). `a`/`b` use distinct seeds +so the two bool masks differ (each ~50% True, independent -> OR ~75% True), a +real mix that a wrong op (e.g. AND) would fail. `bitwise_or` on bool is identical +to `logical_or` (shares the handler). Output is bool (byte-exact golden). +`LogicalOrTest` is the export-delegation smoke test. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + + +class LogicalOrModule(torch.nn.Module): + def __init__(self, shape) -> None: + super().__init__() + self.register_buffer("z", torch.zeros(shape)) + + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return torch.logical_or(a > self.z, b > self.z) + + +class BitwiseOrModule(torch.nn.Module): + def __init__(self, shape) -> None: + super().__init__() + self.register_buffer("z", torch.zeros(shape)) + + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return torch.bitwise_or(a > self.z, b > self.z) + + +def _lo_gen(seed): + # Distinct per-input seed so the two derived bool masks differ. + def g(shape): + gen = torch.Generator().manual_seed(seed) + return torch.randn(*shape, generator=gen, dtype=torch.float32) + + return g + + +lo_gen_a = _lo_gen(0) +lo_gen_b = _lo_gen(1) + + +# All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). +SHAPES = [(4, 8), (2, 3, 8), (16, 16)] + + +class LogicalOrTest(unittest.TestCase): + def _assert_delegates(self, mod, inputs, op_name, shape) -> None: + ep = torch.export.export(mod.eval(), inputs) + edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + et = edge.to_executorch() + deleg = any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + self.assertTrue(deleg, f"Expected VulkanBackend delegate ({op_name} {shape})") + gm = edge.exported_program().graph_module + self.assertTrue( + all(op_name not in str(getattr(n, "target", "")) for n in gm.graph.nodes), + f"{op_name} fell back to CPU for {shape}", + ) + + def test_export_delegates(self) -> None: + for shape in SHAPES: + with self.subTest(shape=shape): + a = lo_gen_a(shape) + b = lo_gen_b(shape) + self._assert_delegates( + LogicalOrModule(shape), (a, b), "logical_or", shape + ) + self._assert_delegates( + BitwiseOrModule(shape), (a, b), "bitwise_or", shape + ) From bb297a79c377c247a5976c25fb6574d2edece587 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:44 -0700 Subject: [PATCH 73/75] [ExecuTorch][WebGPU] Raise TensorMeta rank cap 4->8 to fix rank>4 ops (0x30) Pull Request resolved: https://github.com/pytorch/executorch/pull/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`. 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` 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, 2>` / `strides: array, 2>`: in the uniform address space the `vec4` 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` -> `array, 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/) --- backends/webgpu/runtime/ops/TensorMeta.h | 14 +++++++------- .../webgpu/runtime/ops/binary/BinaryOp.cpp | 2 +- .../runtime/ops/binary_op/binary_div_wgsl.h | 18 +++++++++--------- .../runtime/ops/binary_op/binary_op.wgsl | 16 ++++++++-------- .../runtime/ops/binary_op/binary_sub_wgsl.h | 18 +++++++++--------- backends/webgpu/runtime/ops/cat/cat.wgsl | 10 +++++----- backends/webgpu/runtime/ops/cat/cat_wgsl.h | 12 ++++++------ backends/webgpu/runtime/ops/div/BinaryOp.cpp | 2 +- .../runtime/ops/expand_copy/expand_copy.wgsl | 10 +++++----- .../runtime/ops/expand_copy/expand_copy_wgsl.h | 12 ++++++------ backends/webgpu/runtime/ops/flip/Flip.cpp | 6 +++--- backends/webgpu/runtime/ops/flip/flip.wgsl | 16 ++++++++-------- backends/webgpu/runtime/ops/flip/flip_wgsl.h | 18 +++++++++--------- .../ops/floor_divide/binary_floor_divide.wgsl | 16 ++++++++-------- .../floor_divide/binary_floor_divide_wgsl.h | 18 +++++++++--------- backends/webgpu/runtime/ops/gather/gather.wgsl | 10 +++++----- .../webgpu/runtime/ops/gather/gather_wgsl.h | 12 ++++++------ .../runtime/ops/index_select/IndexSelect.cpp | 2 +- .../runtime/ops/index_select/index_select.wgsl | 10 +++++----- .../ops/index_select/index_select_wgsl.h | 12 ++++++------ .../runtime/ops/minimum/binary_minimum.wgsl | 16 ++++++++-------- .../runtime/ops/minimum/binary_minimum_wgsl.h | 18 +++++++++--------- backends/webgpu/runtime/ops/mul/BinaryOp.cpp | 2 +- .../webgpu/runtime/ops/mul/binary_mul.wgsl | 16 ++++++++-------- .../webgpu/runtime/ops/mul/binary_mul_wgsl.h | 18 +++++++++--------- .../webgpu/runtime/ops/permute/Permute.cpp | 6 +++--- .../webgpu/runtime/ops/permute/permute.wgsl | 13 +++++++------ .../webgpu/runtime/ops/permute/permute_wgsl.h | 15 ++++++++------- .../webgpu/runtime/ops/pow/binary_pow.wgsl | 16 ++++++++-------- .../webgpu/runtime/ops/pow/binary_pow_wgsl.h | 18 +++++++++--------- backends/webgpu/runtime/ops/repeat/Repeat.cpp | 2 +- backends/webgpu/runtime/ops/repeat/repeat.wgsl | 12 ++++++------ .../webgpu/runtime/ops/repeat/repeat_wgsl.h | 14 +++++++------- backends/webgpu/runtime/ops/select/select.wgsl | 12 ++++++------ .../webgpu/runtime/ops/select/select_wgsl.h | 14 +++++++------- backends/webgpu/runtime/ops/slice/slice.wgsl | 10 +++++----- backends/webgpu/runtime/ops/slice/slice_wgsl.h | 12 ++++++------ backends/webgpu/runtime/ops/sub/BinaryOp.cpp | 2 +- backends/webgpu/runtime/ops/where/Where.cpp | 2 +- backends/webgpu/runtime/ops/where/where.wgsl | 14 +++++++------- backends/webgpu/runtime/ops/where/where_wgsl.h | 16 ++++++++-------- 41 files changed, 242 insertions(+), 240 deletions(-) diff --git a/backends/webgpu/runtime/ops/TensorMeta.h b/backends/webgpu/runtime/ops/TensorMeta.h index 3a552851e08..bb1a0aa06af 100644 --- a/backends/webgpu/runtime/ops/TensorMeta.h +++ b/backends/webgpu/runtime/ops/TensorMeta.h @@ -16,9 +16,9 @@ namespace executorch::backends::webgpu { -constexpr uint32_t kTensorMetaMaxNdim = 4; +constexpr uint32_t kTensorMetaMaxNdim = 8; -// Per-tensor metadata UBO; mirrors Vulkan BufferMetadata (4-dim NCHW, std140). +// Per-tensor metadata UBO; mirrors Vulkan BufferMetadata (8-dim NCHW, std140). struct TensorMeta { uint32_t ndim; uint32_t numel; @@ -28,19 +28,19 @@ struct TensorMeta { }; static_assert( - sizeof(TensorMeta) == 48, - "TensorMeta std140 layout must be 48 bytes to match the WGSL uniform"); + sizeof(TensorMeta) == 80, + "TensorMeta std140 layout must be 80 bytes to match the WGSL uniform"); // Lock the std140 field offsets the WGSL uniform reads, not just total size. static_assert(offsetof(TensorMeta, ndim) == 0); static_assert(offsetof(TensorMeta, numel) == 4); static_assert(offsetof(TensorMeta, sizes) == 16); -static_assert(offsetof(TensorMeta, strides) == 32); +static_assert(offsetof(TensorMeta, strides) == 48); // Fill TensorMeta from NCHW dims: contiguous strides, padded trailing slots. inline void fill_tensor_meta(const WebGPUTensor& t, TensorMeta* m) { const uint32_t ndim = static_cast(t.dims.size()); if (ndim > kTensorMetaMaxNdim) { - throw std::runtime_error("TensorMeta: tensor rank exceeds 4 (MAX_NDIM)"); + throw std::runtime_error("TensorMeta: tensor rank exceeds 8 (MAX_NDIM)"); } *m = {}; for (uint32_t d = 0; d < kTensorMetaMaxNdim; d++) { @@ -67,7 +67,7 @@ inline void fill_tensor_meta_broadcast( TensorMeta* m) { const uint32_t rank = static_cast(t.dims.size()); if (out_ndim > kTensorMetaMaxNdim) { - throw std::runtime_error("TensorMeta: out_ndim exceeds 4 (MAX_NDIM)"); + throw std::runtime_error("TensorMeta: out_ndim exceeds 8 (MAX_NDIM)"); } if (rank > out_ndim) { throw std::runtime_error("TensorMeta: operand rank exceeds out_ndim"); diff --git a/backends/webgpu/runtime/ops/binary/BinaryOp.cpp b/backends/webgpu/runtime/ops/binary/BinaryOp.cpp index dd136dfe2d1..84436071e27 100644 --- a/backends/webgpu/runtime/ops/binary/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/binary/BinaryOp.cpp @@ -48,7 +48,7 @@ void add_binary_broadcast_op( if (out_tensor.dims.size() > kTensorMetaMaxNdim || in1_tensor.dims.size() > kTensorMetaMaxNdim || in2_tensor.dims.size() > kTensorMetaMaxNdim) { - throw std::runtime_error(name + ": tensor rank exceeds 4 (MAX_NDIM)"); + throw std::runtime_error(name + ": tensor rank exceeds 8 (MAX_NDIM)"); } const uint32_t out_ndim = static_cast(out_tensor.dims.size()); diff --git a/backends/webgpu/runtime/ops/binary_op/binary_div_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_div_wgsl.h index 3724e0e4011..c4ce390561a 100644 --- a/backends/webgpu/runtime/ops/binary_op/binary_div_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_div_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from binary_op.wgsl - DO NOT EDIT. -// wgsl-sha256: 233876a29c95d830f88fb22f2e3638b3b91e078b8f957b3d2c199352bdf40f93 +// wgsl-sha256: e36b560fd623dd5337b9ae57acd8981c9c635b995d6021caf1331c182cd3f0cd inline constexpr const char* kBinaryDivWGSL = R"( @group(0) @binding(0) var input1: array; @group(0) @binding(1) var input2: array; @@ -22,8 +22,8 @@ inline constexpr const char* kBinaryDivWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(3) var out_meta: TensorMeta; @group(0) @binding(4) var in1_meta: TensorMeta; @@ -47,8 +47,8 @@ fn main( var same = true; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d] != out_meta.sizes[d] || - in2_meta.sizes[d] != out_meta.sizes[d]) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { same = false; } } @@ -61,10 +61,10 @@ fn main( var l1: u32 = 0u; var l2: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d]; - l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; } output[idx] = op(input1[l1], input2[l2]); } diff --git a/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl b/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl index 070cb9adf95..98076ed98b5 100644 --- a/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl +++ b/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl @@ -5,8 +5,8 @@ struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(3) var out_meta: TensorMeta; @group(0) @binding(4) var in1_meta: TensorMeta; @@ -32,8 +32,8 @@ fn main( var same = true; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d] != out_meta.sizes[d] || - in2_meta.sizes[d] != out_meta.sizes[d]) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { same = false; } } @@ -46,10 +46,10 @@ fn main( var l1: u32 = 0u; var l2: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d]; - l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; } output[idx] = op(input1[l1], input2[l2]); } diff --git a/backends/webgpu/runtime/ops/binary_op/binary_sub_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_sub_wgsl.h index caffc03e6dc..d29ca722ece 100644 --- a/backends/webgpu/runtime/ops/binary_op/binary_sub_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_sub_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from binary_op.wgsl - DO NOT EDIT. -// wgsl-sha256: 496b343cef6838c8316686916a1c8fd971afc7a9abfc9abca4eb28db60611371 +// wgsl-sha256: 63209ff70422a21fc340d9aadba0945bc259bba89bdf05db018a6507d01c7ae5 inline constexpr const char* kBinarySubWGSL = R"( @group(0) @binding(0) var input1: array; @group(0) @binding(1) var input2: array; @@ -22,8 +22,8 @@ inline constexpr const char* kBinarySubWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(3) var out_meta: TensorMeta; @group(0) @binding(4) var in1_meta: TensorMeta; @@ -48,8 +48,8 @@ fn main( var same = true; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d] != out_meta.sizes[d] || - in2_meta.sizes[d] != out_meta.sizes[d]) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { same = false; } } @@ -62,10 +62,10 @@ fn main( var l1: u32 = 0u; var l2: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d]; - l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; } output[idx] = op(input1[l1], input2[l2]); } diff --git a/backends/webgpu/runtime/ops/cat/cat.wgsl b/backends/webgpu/runtime/ops/cat/cat.wgsl index 3b1f4aaaa4d..aa90f3979c2 100644 --- a/backends/webgpu/runtime/ops/cat/cat.wgsl +++ b/backends/webgpu/runtime/ops/cat/cat.wgsl @@ -4,8 +4,8 @@ struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(2) var out_meta: TensorMeta; @group(0) @binding(3) var in_meta: TensorMeta; @@ -29,13 +29,13 @@ fn main(@builtin(global_invocation_id) gid: vec3) { var rem = in_bufi; var out_bufi: u32 = 0u; for (var d: u32 = 0u; d < in_meta.ndim; d = d + 1u) { - let coord = rem / in_meta.strides[d]; - rem = rem % in_meta.strides[d]; + let coord = rem / in_meta.strides[d >> 2u][d & 3u]; + rem = rem % in_meta.strides[d >> 2u][d & 3u]; var out_coord = coord; if (d == params.concat_dim) { out_coord = coord + params.off_k; } - out_bufi = out_bufi + out_coord * out_meta.strides[d]; + out_bufi = out_bufi + out_coord * out_meta.strides[d >> 2u][d & 3u]; } output[out_bufi] = input[in_bufi]; } diff --git a/backends/webgpu/runtime/ops/cat/cat_wgsl.h b/backends/webgpu/runtime/ops/cat/cat_wgsl.h index 94d7e2afdc8..26e5f80f8a5 100644 --- a/backends/webgpu/runtime/ops/cat/cat_wgsl.h +++ b/backends/webgpu/runtime/ops/cat/cat_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from cat.wgsl - DO NOT EDIT. -// wgsl-sha256: d1fcb4da7e32c6295b80d581c093b78d0a4b43a972fe2d5d9d94c4f9ae459f4c +// wgsl-sha256: 1a5f66607e2959c12d757989a3c1ae6fc6f2ede139bc9afb6e14dd3961a6fcb7 inline constexpr const char* kCatWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -21,8 +21,8 @@ inline constexpr const char* kCatWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(2) var out_meta: TensorMeta; @group(0) @binding(3) var in_meta: TensorMeta; @@ -46,13 +46,13 @@ fn main(@builtin(global_invocation_id) gid: vec3) { var rem = in_bufi; var out_bufi: u32 = 0u; for (var d: u32 = 0u; d < in_meta.ndim; d = d + 1u) { - let coord = rem / in_meta.strides[d]; - rem = rem % in_meta.strides[d]; + let coord = rem / in_meta.strides[d >> 2u][d & 3u]; + rem = rem % in_meta.strides[d >> 2u][d & 3u]; var out_coord = coord; if (d == params.concat_dim) { out_coord = coord + params.off_k; } - out_bufi = out_bufi + out_coord * out_meta.strides[d]; + out_bufi = out_bufi + out_coord * out_meta.strides[d >> 2u][d & 3u]; } output[out_bufi] = input[in_bufi]; } diff --git a/backends/webgpu/runtime/ops/div/BinaryOp.cpp b/backends/webgpu/runtime/ops/div/BinaryOp.cpp index 279b50e5ec5..4f2f893272e 100644 --- a/backends/webgpu/runtime/ops/div/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/div/BinaryOp.cpp @@ -38,7 +38,7 @@ void div_impl(WebGPUGraph& graph, const std::vector& args) { if (out_tensor.dims.size() > kTensorMetaMaxNdim || in1_tensor.dims.size() > kTensorMetaMaxNdim || in2_tensor.dims.size() > kTensorMetaMaxNdim) { - throw std::runtime_error("div: tensor rank exceeds 4 (MAX_NDIM)"); + throw std::runtime_error("div: tensor rank exceeds 8 (MAX_NDIM)"); } const uint32_t out_ndim = static_cast(out_tensor.dims.size()); diff --git a/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl b/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl index d6c65588978..053311a69f4 100644 --- a/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl +++ b/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl @@ -4,8 +4,8 @@ struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(2) var out_meta: TensorMeta; @group(0) @binding(3) var in_meta: TensorMeta; @@ -21,9 +21,9 @@ fn main(@builtin(global_invocation_id) gid: vec3) { var rem = idx; var l: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - l = l + min(coord, in_meta.sizes[d] - 1u) * in_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l = l + min(coord, in_meta.sizes[d >> 2u][d & 3u] - 1u) * in_meta.strides[d >> 2u][d & 3u]; } output[idx] = input[l]; } diff --git a/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h b/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h index 56b3f708767..83c5881604f 100644 --- a/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h +++ b/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from expand_copy.wgsl - DO NOT EDIT. -// wgsl-sha256: ad996b7fd6eca5c5773715af3a9a117da83ef522042eb0ee918a623096417815 +// wgsl-sha256: 99953670bea89e42bc9c689ab80addfd9a331442c8c8f1a5b0c39dbe11c19370 inline constexpr const char* kExpandCopyWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -21,8 +21,8 @@ inline constexpr const char* kExpandCopyWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(2) var out_meta: TensorMeta; @group(0) @binding(3) var in_meta: TensorMeta; @@ -38,9 +38,9 @@ fn main(@builtin(global_invocation_id) gid: vec3) { var rem = idx; var l: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - l = l + min(coord, in_meta.sizes[d] - 1u) * in_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l = l + min(coord, in_meta.sizes[d >> 2u][d & 3u] - 1u) * in_meta.strides[d >> 2u][d & 3u]; } output[idx] = input[l]; } diff --git a/backends/webgpu/runtime/ops/flip/Flip.cpp b/backends/webgpu/runtime/ops/flip/Flip.cpp index 81943ac3151..20461da5406 100644 --- a/backends/webgpu/runtime/ops/flip/Flip.cpp +++ b/backends/webgpu/runtime/ops/flip/Flip.cpp @@ -27,8 +27,8 @@ struct FlipParams { uint32_t flip[kTensorMetaMaxNdim]; }; static_assert( - sizeof(FlipParams) == 16, - "FlipParams must match the WGSL Params vec4 (16 bytes)"); + sizeof(FlipParams) == 32, + "FlipParams must match the WGSL Params array, 2> (32 bytes)"); // flip: reverse coords along dims (Vulkan Flip.cpp, NCHW; any 4-byte dtype). void flip_impl(WebGPUGraph& graph, const std::vector& args) { @@ -50,7 +50,7 @@ void flip_impl(WebGPUGraph& graph, const std::vector& args) { const auto& out_tensor = graph.get_tensor(out_id); const int ndim = static_cast(in_tensor.dims.size()); if (ndim > static_cast(kTensorMetaMaxNdim)) { - throw std::runtime_error("flip: tensor rank exceeds 4 (MAX_NDIM)"); + throw std::runtime_error("flip: tensor rank exceeds 8 (MAX_NDIM)"); } // flip preserves shape: the output rank + dims must equal the input's, else // the kernel unravels out coords against a shape it wasn't reversed for. diff --git a/backends/webgpu/runtime/ops/flip/flip.wgsl b/backends/webgpu/runtime/ops/flip/flip.wgsl index b66c9cd6751..8392424b3fb 100644 --- a/backends/webgpu/runtime/ops/flip/flip.wgsl +++ b/backends/webgpu/runtime/ops/flip/flip.wgsl @@ -4,14 +4,14 @@ struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(2) var out_meta: TensorMeta; @group(0) @binding(3) var in_meta: TensorMeta; struct Params { - flip: vec4, + flip: array, 2>, } @group(0) @binding(4) var params: Params; @@ -31,11 +31,11 @@ fn main( var rem = out_bufi; var in_bufi: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - let flipped = in_meta.sizes[d] - 1u - coord; - let in_coord = select(coord, flipped, params.flip[d] == 1u); - in_bufi = in_bufi + in_coord * in_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + let flipped = in_meta.sizes[d >> 2u][d & 3u] - 1u - coord; + let in_coord = select(coord, flipped, params.flip[d >> 2u][d & 3u] == 1u); + in_bufi = in_bufi + in_coord * in_meta.strides[d >> 2u][d & 3u]; } output[out_bufi] = input[in_bufi]; } diff --git a/backends/webgpu/runtime/ops/flip/flip_wgsl.h b/backends/webgpu/runtime/ops/flip/flip_wgsl.h index 138fa66aef5..d1fc0c92abc 100644 --- a/backends/webgpu/runtime/ops/flip/flip_wgsl.h +++ b/backends/webgpu/runtime/ops/flip/flip_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from flip.wgsl - DO NOT EDIT. -// wgsl-sha256: e85b3a7f753ec183e83f55f53cc764b5ffa748b4e229d81b978dabe58f576e80 +// wgsl-sha256: 0c3a6b73a4a4923df6b742a9dacc0d73e6b5b033e4b91835ec71107a06848552 inline constexpr const char* kFlipWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -21,14 +21,14 @@ inline constexpr const char* kFlipWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(2) var out_meta: TensorMeta; @group(0) @binding(3) var in_meta: TensorMeta; struct Params { - flip: vec4, + flip: array, 2>, } @group(0) @binding(4) var params: Params; @@ -48,11 +48,11 @@ fn main( var rem = out_bufi; var in_bufi: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - let flipped = in_meta.sizes[d] - 1u - coord; - let in_coord = select(coord, flipped, params.flip[d] == 1u); - in_bufi = in_bufi + in_coord * in_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + let flipped = in_meta.sizes[d >> 2u][d & 3u] - 1u - coord; + let in_coord = select(coord, flipped, params.flip[d >> 2u][d & 3u] == 1u); + in_bufi = in_bufi + in_coord * in_meta.strides[d >> 2u][d & 3u]; } output[out_bufi] = input[in_bufi]; } diff --git a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl b/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl index 5191a220238..3b4edd788b0 100644 --- a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl +++ b/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl @@ -5,8 +5,8 @@ struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(3) var out_meta: TensorMeta; @group(0) @binding(4) var in1_meta: TensorMeta; @@ -27,8 +27,8 @@ fn main( // Fast path: every input dim matches the output dim -> elementwise. var same = true; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d] != out_meta.sizes[d] || - in2_meta.sizes[d] != out_meta.sizes[d]) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { same = false; } } @@ -42,10 +42,10 @@ fn main( var l1: u32 = 0u; var l2: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d]; - l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; } output[idx] = floor(input1[l1] / input2[l2]); } diff --git a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h b/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h index 63bd6d1f48b..fe2315df30f 100644 --- a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h +++ b/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from binary_floor_divide.wgsl - DO NOT EDIT. -// wgsl-sha256: bbc1c3a39dd17c88e7df4f2adbf3b9cfe424693aed00deab268a03b2e0cdc958 +// wgsl-sha256: baf71d277da79389315a6b96b439e7f0a55842e8288283f2af121f84536b3af3 inline constexpr const char* kBinaryFloorDivideWGSL = R"( @group(0) @binding(0) var input1: array; @group(0) @binding(1) var input2: array; @@ -22,8 +22,8 @@ inline constexpr const char* kBinaryFloorDivideWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(3) var out_meta: TensorMeta; @group(0) @binding(4) var in1_meta: TensorMeta; @@ -44,8 +44,8 @@ fn main( // Fast path: every input dim matches the output dim -> elementwise. var same = true; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d] != out_meta.sizes[d] || - in2_meta.sizes[d] != out_meta.sizes[d]) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { same = false; } } @@ -59,10 +59,10 @@ fn main( var l1: u32 = 0u; var l2: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d]; - l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; } output[idx] = floor(input1[l1] / input2[l2]); } diff --git a/backends/webgpu/runtime/ops/gather/gather.wgsl b/backends/webgpu/runtime/ops/gather/gather.wgsl index fb62206be64..5fa428952b3 100644 --- a/backends/webgpu/runtime/ops/gather/gather.wgsl +++ b/backends/webgpu/runtime/ops/gather/gather.wgsl @@ -5,8 +5,8 @@ struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(3) var out_meta: TensorMeta; @group(0) @binding(4) var self_meta: TensorMeta; @@ -27,13 +27,13 @@ fn main(@builtin(global_invocation_id) gid: vec3) { var rem = o; var self_idx: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let c = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; + let c = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; var coord = c; if (d == params.dim) { coord = u32(indices[o]); } - self_idx = self_idx + coord * self_meta.strides[d]; + self_idx = self_idx + coord * self_meta.strides[d >> 2u][d & 3u]; } output[o] = self_[self_idx]; } diff --git a/backends/webgpu/runtime/ops/gather/gather_wgsl.h b/backends/webgpu/runtime/ops/gather/gather_wgsl.h index 5b6922e21ea..8e679d9f9c2 100644 --- a/backends/webgpu/runtime/ops/gather/gather_wgsl.h +++ b/backends/webgpu/runtime/ops/gather/gather_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from gather.wgsl - DO NOT EDIT. -// wgsl-sha256: cd19a9ed2753e2a97e96fb90f7a72e8f40d08feb6c84fbed7f2c52e71d77a03c +// wgsl-sha256: c0811711f4603f840241a9cdd458f7b941dcf32336ff84f2b6f1333b8d6f4cfb inline constexpr const char* kGatherWGSL = R"( @group(0) @binding(0) var self_: array; @group(0) @binding(1) var indices: array; @@ -22,8 +22,8 @@ inline constexpr const char* kGatherWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(3) var out_meta: TensorMeta; @group(0) @binding(4) var self_meta: TensorMeta; @@ -44,13 +44,13 @@ fn main(@builtin(global_invocation_id) gid: vec3) { var rem = o; var self_idx: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let c = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; + let c = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; var coord = c; if (d == params.dim) { coord = u32(indices[o]); } - self_idx = self_idx + coord * self_meta.strides[d]; + self_idx = self_idx + coord * self_meta.strides[d >> 2u][d & 3u]; } output[o] = self_[self_idx]; } diff --git a/backends/webgpu/runtime/ops/index_select/IndexSelect.cpp b/backends/webgpu/runtime/ops/index_select/IndexSelect.cpp index dda307029d5..4f9f9c540f1 100644 --- a/backends/webgpu/runtime/ops/index_select/IndexSelect.cpp +++ b/backends/webgpu/runtime/ops/index_select/IndexSelect.cpp @@ -54,7 +54,7 @@ void index_select_impl(WebGPUGraph& graph, const std::vector& args) { const int64_t ndim = static_cast(self_tensor.dims.size()); if (ndim > static_cast(kTensorMetaMaxNdim)) { - throw std::runtime_error("index_select: tensor rank exceeds 4 (MAX_NDIM)"); + throw std::runtime_error("index_select: tensor rank exceeds 8 (MAX_NDIM)"); } if (graph.get_value_type(dim_id) != WebGPUGraph::ValueType::Int) { diff --git a/backends/webgpu/runtime/ops/index_select/index_select.wgsl b/backends/webgpu/runtime/ops/index_select/index_select.wgsl index 79b977e7f7c..d1fc827031c 100644 --- a/backends/webgpu/runtime/ops/index_select/index_select.wgsl +++ b/backends/webgpu/runtime/ops/index_select/index_select.wgsl @@ -5,8 +5,8 @@ struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(3) var out_meta: TensorMeta; @group(0) @binding(4) var in_meta: TensorMeta; @@ -33,13 +33,13 @@ fn main( var rem = out_bufi; var in_bufi: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; var in_coord = coord; if (d == dim) { in_coord = u32(index[coord]); } - in_bufi = in_bufi + in_coord * in_meta.strides[d]; + in_bufi = in_bufi + in_coord * in_meta.strides[d >> 2u][d & 3u]; } output[out_bufi] = input[in_bufi]; } diff --git a/backends/webgpu/runtime/ops/index_select/index_select_wgsl.h b/backends/webgpu/runtime/ops/index_select/index_select_wgsl.h index 6e83be3883b..e8248cb82d3 100644 --- a/backends/webgpu/runtime/ops/index_select/index_select_wgsl.h +++ b/backends/webgpu/runtime/ops/index_select/index_select_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from index_select.wgsl - DO NOT EDIT. -// wgsl-sha256: 63d3a49ddd72a67b05bb3d03a05ba375e4e3b096a47ec67406d05c3ddc9eb241 +// wgsl-sha256: cd1d68712a409b98f3526a9bb4ff0d815c8c7e4c4696578315ea7ac1bbba0685 inline constexpr const char* kIndexSelectWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -22,8 +22,8 @@ inline constexpr const char* kIndexSelectWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(3) var out_meta: TensorMeta; @group(0) @binding(4) var in_meta: TensorMeta; @@ -50,13 +50,13 @@ fn main( var rem = out_bufi; var in_bufi: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; var in_coord = coord; if (d == dim) { in_coord = u32(index[coord]); } - in_bufi = in_bufi + in_coord * in_meta.strides[d]; + in_bufi = in_bufi + in_coord * in_meta.strides[d >> 2u][d & 3u]; } output[out_bufi] = input[in_bufi]; } diff --git a/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl b/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl index 4688d0647d3..e79cb2d2bcc 100644 --- a/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl +++ b/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl @@ -5,8 +5,8 @@ struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(3) var out_meta: TensorMeta; @group(0) @binding(4) var in1_meta: TensorMeta; @@ -27,8 +27,8 @@ fn main( // Fast path: every input dim matches the output dim -> elementwise. var same = true; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d] != out_meta.sizes[d] || - in2_meta.sizes[d] != out_meta.sizes[d]) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { same = false; } } @@ -42,10 +42,10 @@ fn main( var l1: u32 = 0u; var l2: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d]; - l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; } output[idx] = min(input1[l1], input2[l2]); } diff --git a/backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h b/backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h index 2bba2afc6be..88d9614ba59 100644 --- a/backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h +++ b/backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from binary_minimum.wgsl - DO NOT EDIT. -// wgsl-sha256: 4e4fa9a44e6ec41fb595bec7766c65e7e503d02a9222cd45e976b459aee14b4c +// wgsl-sha256: 929b7ba85936e3652baea9f4e5e7f049d232c7ae7a74814a536b4c2674897972 inline constexpr const char* kBinaryMinimumWGSL = R"( @group(0) @binding(0) var input1: array; @group(0) @binding(1) var input2: array; @@ -22,8 +22,8 @@ inline constexpr const char* kBinaryMinimumWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(3) var out_meta: TensorMeta; @group(0) @binding(4) var in1_meta: TensorMeta; @@ -44,8 +44,8 @@ fn main( // Fast path: every input dim matches the output dim -> elementwise. var same = true; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d] != out_meta.sizes[d] || - in2_meta.sizes[d] != out_meta.sizes[d]) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { same = false; } } @@ -59,10 +59,10 @@ fn main( var l1: u32 = 0u; var l2: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d]; - l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; } output[idx] = min(input1[l1], input2[l2]); } diff --git a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp index a14e583037a..ed81fb3216c 100644 --- a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp @@ -38,7 +38,7 @@ void mul_impl(WebGPUGraph& graph, const std::vector& args) { if (out_tensor.dims.size() > kTensorMetaMaxNdim || in1_tensor.dims.size() > kTensorMetaMaxNdim || in2_tensor.dims.size() > kTensorMetaMaxNdim) { - throw std::runtime_error("mul: tensor rank exceeds 4 (MAX_NDIM)"); + throw std::runtime_error("mul: tensor rank exceeds 8 (MAX_NDIM)"); } const uint32_t out_ndim = static_cast(out_tensor.dims.size()); diff --git a/backends/webgpu/runtime/ops/mul/binary_mul.wgsl b/backends/webgpu/runtime/ops/mul/binary_mul.wgsl index 66722968e44..f82a16e4b21 100644 --- a/backends/webgpu/runtime/ops/mul/binary_mul.wgsl +++ b/backends/webgpu/runtime/ops/mul/binary_mul.wgsl @@ -5,8 +5,8 @@ struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(3) var out_meta: TensorMeta; @group(0) @binding(4) var in1_meta: TensorMeta; @@ -27,8 +27,8 @@ fn main( // Fast path: every input dim matches the output dim -> elementwise. var same = true; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d] != out_meta.sizes[d] || - in2_meta.sizes[d] != out_meta.sizes[d]) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { same = false; } } @@ -42,10 +42,10 @@ fn main( var l1: u32 = 0u; var l2: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d]; - l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; } output[idx] = input1[l1] * input2[l2]; } diff --git a/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h b/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h index 1c700615c7f..c9f60dbd200 100644 --- a/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h +++ b/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from binary_mul.wgsl - DO NOT EDIT. -// wgsl-sha256: cca69c3428f37f293942637e23f664225dec81a56f184bcb63185b6629dd155e +// wgsl-sha256: d248c0f1856b57115a5001a47f4936caa564dd3b787c02ceba504a13ab987812 inline constexpr const char* kBinaryMulWGSL = R"( @group(0) @binding(0) var input1: array; @group(0) @binding(1) var input2: array; @@ -22,8 +22,8 @@ inline constexpr const char* kBinaryMulWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(3) var out_meta: TensorMeta; @group(0) @binding(4) var in1_meta: TensorMeta; @@ -44,8 +44,8 @@ fn main( // Fast path: every input dim matches the output dim -> elementwise. var same = true; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d] != out_meta.sizes[d] || - in2_meta.sizes[d] != out_meta.sizes[d]) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { same = false; } } @@ -59,10 +59,10 @@ fn main( var l1: u32 = 0u; var l2: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d]; - l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; } output[idx] = input1[l1] * input2[l2]; } diff --git a/backends/webgpu/runtime/ops/permute/Permute.cpp b/backends/webgpu/runtime/ops/permute/Permute.cpp index fb23fce7019..ce9e55d04fc 100644 --- a/backends/webgpu/runtime/ops/permute/Permute.cpp +++ b/backends/webgpu/runtime/ops/permute/Permute.cpp @@ -27,8 +27,8 @@ struct PermuteParams { uint32_t perm[kTensorMetaMaxNdim]; }; static_assert( - sizeof(PermuteParams) == 16, - "PermuteParams must match the WGSL Params vec4 (16 bytes)"); + sizeof(PermuteParams) == 32, + "PermuteParams must match the WGSL Params array, 2> (32 bytes)"); // permute: out coord d -> in coord perm[d] (Vulkan permute_buffer.glsl, NCHW). void permute_impl(WebGPUGraph& graph, const std::vector& args) { @@ -60,7 +60,7 @@ void permute_impl(WebGPUGraph& graph, const std::vector& args) { uint32_t perm[kTensorMetaMaxNdim]; bool seen[kTensorMetaMaxNdim] = {}; if (ndim > static_cast(kTensorMetaMaxNdim)) { - throw std::runtime_error("permute: tensor rank exceeds 4 (MAX_NDIM)"); + throw std::runtime_error("permute: tensor rank exceeds 8 (MAX_NDIM)"); } for (int d = 0; d < ndim; d++) { int64_t p = dims[d]; diff --git a/backends/webgpu/runtime/ops/permute/permute.wgsl b/backends/webgpu/runtime/ops/permute/permute.wgsl index e62fa5624d5..9746a4641d3 100644 --- a/backends/webgpu/runtime/ops/permute/permute.wgsl +++ b/backends/webgpu/runtime/ops/permute/permute.wgsl @@ -4,14 +4,14 @@ struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(2) var out_meta: TensorMeta; @group(0) @binding(3) var in_meta: TensorMeta; struct Params { - perm: vec4, + perm: array, 2>, } @group(0) @binding(4) var params: Params; @@ -31,9 +31,10 @@ fn main( var rem = out_bufi; var in_bufi: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - in_bufi = in_bufi + coord * in_meta.strides[params.perm[d]]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + let p = params.perm[d >> 2u][d & 3u]; + in_bufi = in_bufi + coord * in_meta.strides[p >> 2u][p & 3u]; } output[out_bufi] = input[in_bufi]; } diff --git a/backends/webgpu/runtime/ops/permute/permute_wgsl.h b/backends/webgpu/runtime/ops/permute/permute_wgsl.h index b3d4684d54b..4d9b0a4ad25 100644 --- a/backends/webgpu/runtime/ops/permute/permute_wgsl.h +++ b/backends/webgpu/runtime/ops/permute/permute_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from permute.wgsl - DO NOT EDIT. -// wgsl-sha256: 05884aeb14426c979ea037b066266d8cab11f4fed76ee21ee8778e7fc13ad84e +// wgsl-sha256: ec71705cb9feb46edd5398986cfe4f5a40028b8d5a603f0d2bf06d056a844c26 inline constexpr const char* kPermuteWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -21,14 +21,14 @@ inline constexpr const char* kPermuteWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(2) var out_meta: TensorMeta; @group(0) @binding(3) var in_meta: TensorMeta; struct Params { - perm: vec4, + perm: array, 2>, } @group(0) @binding(4) var params: Params; @@ -48,9 +48,10 @@ fn main( var rem = out_bufi; var in_bufi: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - in_bufi = in_bufi + coord * in_meta.strides[params.perm[d]]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + let p = params.perm[d >> 2u][d & 3u]; + in_bufi = in_bufi + coord * in_meta.strides[p >> 2u][p & 3u]; } output[out_bufi] = input[in_bufi]; } diff --git a/backends/webgpu/runtime/ops/pow/binary_pow.wgsl b/backends/webgpu/runtime/ops/pow/binary_pow.wgsl index 894c32a3dc2..2114ef87fee 100644 --- a/backends/webgpu/runtime/ops/pow/binary_pow.wgsl +++ b/backends/webgpu/runtime/ops/pow/binary_pow.wgsl @@ -5,8 +5,8 @@ struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(3) var out_meta: TensorMeta; @group(0) @binding(4) var in1_meta: TensorMeta; @@ -27,8 +27,8 @@ fn main( // Fast path: every input dim matches the output dim -> elementwise. var same = true; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d] != out_meta.sizes[d] || - in2_meta.sizes[d] != out_meta.sizes[d]) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { same = false; } } @@ -42,10 +42,10 @@ fn main( var l1: u32 = 0u; var l2: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d]; - l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; } output[idx] = pow(input1[l1], input2[l2]); } diff --git a/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h b/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h index 6488e360518..3532c091160 100644 --- a/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h +++ b/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from binary_pow.wgsl - DO NOT EDIT. -// wgsl-sha256: ffc116e845c861fb2a43f0c3ca86f150f8e1e95f9d3dde73a3886cd2a1ebff93 +// wgsl-sha256: a88c161bd3f43d21a72ebd8ca6f8611b6b9b854e3572a8e6b820602091bc464c inline constexpr const char* kBinaryPowWGSL = R"( @group(0) @binding(0) var input1: array; @group(0) @binding(1) var input2: array; @@ -22,8 +22,8 @@ inline constexpr const char* kBinaryPowWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(3) var out_meta: TensorMeta; @group(0) @binding(4) var in1_meta: TensorMeta; @@ -44,8 +44,8 @@ fn main( // Fast path: every input dim matches the output dim -> elementwise. var same = true; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d] != out_meta.sizes[d] || - in2_meta.sizes[d] != out_meta.sizes[d]) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { same = false; } } @@ -59,10 +59,10 @@ fn main( var l1: u32 = 0u; var l2: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d]; - l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; } output[idx] = pow(input1[l1], input2[l2]); } diff --git a/backends/webgpu/runtime/ops/repeat/Repeat.cpp b/backends/webgpu/runtime/ops/repeat/Repeat.cpp index 04553b0a67d..67aaae25f43 100644 --- a/backends/webgpu/runtime/ops/repeat/Repeat.cpp +++ b/backends/webgpu/runtime/ops/repeat/Repeat.cpp @@ -38,7 +38,7 @@ void repeat_impl(WebGPUGraph& graph, const std::vector& args) { const auto& out_tensor = graph.get_tensor(out_id); if (out_tensor.dims.size() > kTensorMetaMaxNdim || in_tensor.dims.size() > kTensorMetaMaxNdim) { - throw std::runtime_error("repeat: tensor rank exceeds 4 (MAX_NDIM)"); + throw std::runtime_error("repeat: tensor rank exceeds 8 (MAX_NDIM)"); } TensorMeta out_meta; diff --git a/backends/webgpu/runtime/ops/repeat/repeat.wgsl b/backends/webgpu/runtime/ops/repeat/repeat.wgsl index 8e0c2c4e676..2c0b30fe810 100644 --- a/backends/webgpu/runtime/ops/repeat/repeat.wgsl +++ b/backends/webgpu/runtime/ops/repeat/repeat.wgsl @@ -4,8 +4,8 @@ struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(2) var out_meta: TensorMeta; @group(0) @binding(3) var in_meta: TensorMeta; @@ -28,12 +28,12 @@ fn main( var rem = out_bufi; var in_bufi: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let out_coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; + let out_coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; if (d >= offset) { let in_d = d - offset; - let in_coord = out_coord % in_meta.sizes[in_d]; - in_bufi = in_bufi + in_coord * in_meta.strides[in_d]; + let in_coord = out_coord % in_meta.sizes[in_d >> 2u][in_d & 3u]; + in_bufi = in_bufi + in_coord * in_meta.strides[in_d >> 2u][in_d & 3u]; } } output[out_bufi] = input[in_bufi]; diff --git a/backends/webgpu/runtime/ops/repeat/repeat_wgsl.h b/backends/webgpu/runtime/ops/repeat/repeat_wgsl.h index b08d2d3d3c4..ec4b859180d 100644 --- a/backends/webgpu/runtime/ops/repeat/repeat_wgsl.h +++ b/backends/webgpu/runtime/ops/repeat/repeat_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from repeat.wgsl - DO NOT EDIT. -// wgsl-sha256: ae77e3d6846cd185e96b0fcae4b4bcee2096e659df1a7ba9d0129beb24d754c7 +// wgsl-sha256: 51b2068fa868a260ce789d685d23ff286947e30dad171a9d2a39a955916bf297 inline constexpr const char* kRepeatWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -21,8 +21,8 @@ inline constexpr const char* kRepeatWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(2) var out_meta: TensorMeta; @group(0) @binding(3) var in_meta: TensorMeta; @@ -45,12 +45,12 @@ fn main( var rem = out_bufi; var in_bufi: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let out_coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; + let out_coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; if (d >= offset) { let in_d = d - offset; - let in_coord = out_coord % in_meta.sizes[in_d]; - in_bufi = in_bufi + in_coord * in_meta.strides[in_d]; + let in_coord = out_coord % in_meta.sizes[in_d >> 2u][in_d & 3u]; + in_bufi = in_bufi + in_coord * in_meta.strides[in_d >> 2u][in_d & 3u]; } } output[out_bufi] = input[in_bufi]; diff --git a/backends/webgpu/runtime/ops/select/select.wgsl b/backends/webgpu/runtime/ops/select/select.wgsl index 84938b2c840..ea05cabd10f 100644 --- a/backends/webgpu/runtime/ops/select/select.wgsl +++ b/backends/webgpu/runtime/ops/select/select.wgsl @@ -4,8 +4,8 @@ struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(2) var out_meta: TensorMeta; @group(0) @binding(3) var in_meta: TensorMeta; @@ -27,15 +27,15 @@ fn main(@builtin(global_invocation_id) gid: vec3) { // Gather: out dim od -> in dim (od if od < dim else od+1); sel dim = index. var rem = out_bufi; - var in_bufi: u32 = params.index * in_meta.strides[params.dim]; + var in_bufi: u32 = params.index * in_meta.strides[params.dim >> 2u][params.dim & 3u]; for (var od: u32 = 0u; od < out_meta.ndim; od = od + 1u) { - let coord = rem / out_meta.strides[od]; - rem = rem % out_meta.strides[od]; + let coord = rem / out_meta.strides[od >> 2u][od & 3u]; + rem = rem % out_meta.strides[od >> 2u][od & 3u]; var id = od; if (od >= params.dim) { id = od + 1u; } - in_bufi = in_bufi + coord * in_meta.strides[id]; + in_bufi = in_bufi + coord * in_meta.strides[id >> 2u][id & 3u]; } output[out_bufi] = input[in_bufi]; } diff --git a/backends/webgpu/runtime/ops/select/select_wgsl.h b/backends/webgpu/runtime/ops/select/select_wgsl.h index e66edde240d..d2e3e947bb8 100644 --- a/backends/webgpu/runtime/ops/select/select_wgsl.h +++ b/backends/webgpu/runtime/ops/select/select_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from select.wgsl - DO NOT EDIT. -// wgsl-sha256: 200cf5a8190045aa0562e782f01c1cfaf9681f30f679f5112ccc3d347a0ed8df +// wgsl-sha256: a33af9c7ee49c5d12b226befbb08add1538367750ab7437b5c0b4f5fa0291af9 inline constexpr const char* kSelectWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -21,8 +21,8 @@ inline constexpr const char* kSelectWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(2) var out_meta: TensorMeta; @group(0) @binding(3) var in_meta: TensorMeta; @@ -44,15 +44,15 @@ fn main(@builtin(global_invocation_id) gid: vec3) { // Gather: out dim od -> in dim (od if od < dim else od+1); sel dim = index. var rem = out_bufi; - var in_bufi: u32 = params.index * in_meta.strides[params.dim]; + var in_bufi: u32 = params.index * in_meta.strides[params.dim >> 2u][params.dim & 3u]; for (var od: u32 = 0u; od < out_meta.ndim; od = od + 1u) { - let coord = rem / out_meta.strides[od]; - rem = rem % out_meta.strides[od]; + let coord = rem / out_meta.strides[od >> 2u][od & 3u]; + rem = rem % out_meta.strides[od >> 2u][od & 3u]; var id = od; if (od >= params.dim) { id = od + 1u; } - in_bufi = in_bufi + coord * in_meta.strides[id]; + in_bufi = in_bufi + coord * in_meta.strides[id >> 2u][id & 3u]; } output[out_bufi] = input[in_bufi]; } diff --git a/backends/webgpu/runtime/ops/slice/slice.wgsl b/backends/webgpu/runtime/ops/slice/slice.wgsl index 7ed718b2ca1..e6940ee2e4e 100644 --- a/backends/webgpu/runtime/ops/slice/slice.wgsl +++ b/backends/webgpu/runtime/ops/slice/slice.wgsl @@ -4,8 +4,8 @@ struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(2) var out_meta: TensorMeta; @group(0) @binding(3) var in_meta: TensorMeta; @@ -30,13 +30,13 @@ fn main(@builtin(global_invocation_id) gid: vec3) { var rem = out_bufi; var in_bufi: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; var in_coord = coord; if (d == params.dim) { in_coord = params.start + coord * params.step; } - in_bufi = in_bufi + in_coord * in_meta.strides[d]; + in_bufi = in_bufi + in_coord * in_meta.strides[d >> 2u][d & 3u]; } output[out_bufi] = input[in_bufi]; } diff --git a/backends/webgpu/runtime/ops/slice/slice_wgsl.h b/backends/webgpu/runtime/ops/slice/slice_wgsl.h index 52693f3665b..ff70bacd4a6 100644 --- a/backends/webgpu/runtime/ops/slice/slice_wgsl.h +++ b/backends/webgpu/runtime/ops/slice/slice_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from slice.wgsl - DO NOT EDIT. -// wgsl-sha256: 6639d985420d43a67de0847749918ab6216e0785399bdcae737d49b81c773526 +// wgsl-sha256: 6a895a8c321cd3ddaffc468d7843c9dea3eaeb9d3de0088a1a09419b3ccfd10b inline constexpr const char* kSliceWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -21,8 +21,8 @@ inline constexpr const char* kSliceWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(2) var out_meta: TensorMeta; @group(0) @binding(3) var in_meta: TensorMeta; @@ -47,13 +47,13 @@ fn main(@builtin(global_invocation_id) gid: vec3) { var rem = out_bufi; var in_bufi: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; var in_coord = coord; if (d == params.dim) { in_coord = params.start + coord * params.step; } - in_bufi = in_bufi + in_coord * in_meta.strides[d]; + in_bufi = in_bufi + in_coord * in_meta.strides[d >> 2u][d & 3u]; } output[out_bufi] = input[in_bufi]; } diff --git a/backends/webgpu/runtime/ops/sub/BinaryOp.cpp b/backends/webgpu/runtime/ops/sub/BinaryOp.cpp index 13b609cc1e7..7b96bf0c9ac 100644 --- a/backends/webgpu/runtime/ops/sub/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/sub/BinaryOp.cpp @@ -45,7 +45,7 @@ void sub_impl(WebGPUGraph& graph, const std::vector& args) { if (out_tensor.dims.size() > kTensorMetaMaxNdim || in1_tensor.dims.size() > kTensorMetaMaxNdim || in2_tensor.dims.size() > kTensorMetaMaxNdim) { - throw std::runtime_error("sub: tensor rank exceeds 4 (MAX_NDIM)"); + throw std::runtime_error("sub: tensor rank exceeds 8 (MAX_NDIM)"); } const uint32_t out_ndim = static_cast(out_tensor.dims.size()); diff --git a/backends/webgpu/runtime/ops/where/Where.cpp b/backends/webgpu/runtime/ops/where/Where.cpp index f8aba4a5d01..e70888308bf 100644 --- a/backends/webgpu/runtime/ops/where/Where.cpp +++ b/backends/webgpu/runtime/ops/where/Where.cpp @@ -45,7 +45,7 @@ void where_impl(WebGPUGraph& graph, const std::vector& args) { cond_tensor.dims.size() > kTensorMetaMaxNdim || a_tensor.dims.size() > kTensorMetaMaxNdim || b_tensor.dims.size() > kTensorMetaMaxNdim) { - throw std::runtime_error("where: tensor rank exceeds 4 (MAX_NDIM)"); + throw std::runtime_error("where: tensor rank exceeds 8 (MAX_NDIM)"); } const uint32_t out_ndim = static_cast(out_tensor.dims.size()); diff --git a/backends/webgpu/runtime/ops/where/where.wgsl b/backends/webgpu/runtime/ops/where/where.wgsl index eb32f314f83..8c11db74200 100644 --- a/backends/webgpu/runtime/ops/where/where.wgsl +++ b/backends/webgpu/runtime/ops/where/where.wgsl @@ -6,8 +6,8 @@ struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(4) var out_meta: TensorMeta; @group(0) @binding(5) var cond_meta: TensorMeta; @@ -34,11 +34,11 @@ fn main(@builtin(global_invocation_id) gid: vec3) { var la: u32 = 0u; var lb: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - lc = lc + min(coord, cond_meta.sizes[d] - 1u) * cond_meta.strides[d]; - la = la + min(coord, a_meta.sizes[d] - 1u) * a_meta.strides[d]; - lb = lb + min(coord, b_meta.sizes[d] - 1u) * b_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + lc = lc + min(coord, cond_meta.sizes[d >> 2u][d & 3u] - 1u) * cond_meta.strides[d >> 2u][d & 3u]; + la = la + min(coord, a_meta.sizes[d >> 2u][d & 3u] - 1u) * a_meta.strides[d >> 2u][d & 3u]; + lb = lb + min(coord, b_meta.sizes[d >> 2u][d & 3u] - 1u) * b_meta.strides[d >> 2u][d & 3u]; } if (cond_is_true(lc)) { diff --git a/backends/webgpu/runtime/ops/where/where_wgsl.h b/backends/webgpu/runtime/ops/where/where_wgsl.h index a8ef64fb169..01177210207 100644 --- a/backends/webgpu/runtime/ops/where/where_wgsl.h +++ b/backends/webgpu/runtime/ops/where/where_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from where.wgsl - DO NOT EDIT. -// wgsl-sha256: e02e6b9e7a073d2aac3d2c110ed5c382ab6663ee9656dc24f4b9c9fdd9e98a38 +// wgsl-sha256: 2c5c6491e95822f767920ad82b81b005808ad5550567b85fa98c1521f188dbfc inline constexpr const char* kWhereWGSL = R"( @group(0) @binding(0) var cond: array; @group(0) @binding(1) var input_a: array; @@ -23,8 +23,8 @@ inline constexpr const char* kWhereWGSL = R"( struct TensorMeta { ndim: u32, numel: u32, - sizes: vec4, - strides: vec4, + sizes: array, 2>, + strides: array, 2>, } @group(0) @binding(4) var out_meta: TensorMeta; @group(0) @binding(5) var cond_meta: TensorMeta; @@ -51,11 +51,11 @@ fn main(@builtin(global_invocation_id) gid: vec3) { var la: u32 = 0u; var lb: u32 = 0u; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d]; - rem = rem % out_meta.strides[d]; - lc = lc + min(coord, cond_meta.sizes[d] - 1u) * cond_meta.strides[d]; - la = la + min(coord, a_meta.sizes[d] - 1u) * a_meta.strides[d]; - lb = lb + min(coord, b_meta.sizes[d] - 1u) * b_meta.strides[d]; + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + lc = lc + min(coord, cond_meta.sizes[d >> 2u][d & 3u] - 1u) * cond_meta.strides[d >> 2u][d & 3u]; + la = la + min(coord, a_meta.sizes[d >> 2u][d & 3u] - 1u) * a_meta.strides[d >> 2u][d & 3u]; + lb = lb + min(coord, b_meta.sizes[d >> 2u][d & 3u] - 1u) * b_meta.strides[d >> 2u][d & 3u]; } if (cond_is_true(lc)) { From 6ec735a8b3e1e062d862bcb2170bec1b0ec6aee2 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:44 -0700 Subject: [PATCH 74/75] [ExecuTorch][WebGPU] Add dynamic-shape resize hook to cat (output mismatch) Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/webgpu/runtime/ops/cat/Cat.cpp | 76 +++++++++++++++++++++++-- 1 file changed, 70 insertions(+), 6 deletions(-) diff --git a/backends/webgpu/runtime/ops/cat/Cat.cpp b/backends/webgpu/runtime/ops/cat/Cat.cpp index 0cfb857745c..28c327ae795 100644 --- a/backends/webgpu/runtime/ops/cat/Cat.cpp +++ b/backends/webgpu/runtime/ops/cat/Cat.cpp @@ -149,6 +149,11 @@ void cat_impl(WebGPUGraph& graph, const std::vector& args) { WGPUPipelineLayout pipeline_layout = wgpuDeviceCreatePipelineLayout(device, &pl_desc); + // Collected for the dynamic-shape resize hook (rewrites these on resize). + std::vector in_meta_bufs(ids.size()); + std::vector params_bufs(ids.size()); + std::vector dispatch_idxs(ids.size()); + uint32_t off_k = 0; for (size_t k = 0; k < ids.size(); k++) { const auto& in_tensor = graph.get_tensor(ids[k]); @@ -195,18 +200,77 @@ void cat_impl(WebGPUGraph& graph, const std::vector& args) { bg_desc.entries = bg_entries; WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - graph.add_dispatch({pipeline, bind_group, wg_counts[k]}); - // Drop our refs; this input's bind group keeps its uniforms alive. - wgpuBufferRelease(in_meta_buf); - wgpuBufferRelease(params_buf); + in_meta_bufs[k] = in_meta_buf; + params_bufs[k] = params_buf; + dispatch_idxs[k] = graph.add_dispatch({pipeline, bind_group, wg_counts[k]}); + // Uniforms kept alive (owned below) so the resize hook can rewrite them. off_k += static_cast(in_tensor.dims[dim]); } wgpuShaderModuleRelease(shader); wgpuBindGroupLayoutRelease(bgl); wgpuPipelineLayoutRelease(pipeline_layout); - // Drop our ref to the shared out_meta; the bind groups keep it alive. - wgpuBufferRelease(out_meta_buf); + + // Dynamic shapes: cat bakes strides/numel/off_k/workgroup_count from the max + // (build) shape, so under a smaller live shape it would scatter with stale + // strides and loop over the stale numel. Recompute every input's in_meta + + // params and the shared out_meta from live dims, and rewrite each dispatch's + // workgroup count. Mirrors the WebGPUGraph SwiGLU/QKV resize templates; on a + // static graph cur_dims == dims, so the hook rewrites identical values. + auto cat_resize = [ids, + out_id, + dim, + wg_size, + out_meta_buf, + in_meta_bufs, + params_bufs, + dispatch_idxs](WebGPUGraph& g) { + const size_t cdim = static_cast(dim); + std::vector out_d = g.cur_dims(ids[0]); + int64_t concat_sum = 0; + for (size_t k = 0; k < ids.size(); k++) { + concat_sum += g.cur_dims(ids[k])[cdim]; + } + out_d[cdim] = concat_sum; + g.set_cur_dims(out_id, out_d); + + WebGPUTensor to; + to.dims = out_d; + TensorMeta out_meta; + fill_tensor_meta(to, &out_meta); + wgpuQueueWriteBuffer( + g.queue(), out_meta_buf, 0, &out_meta, sizeof(out_meta)); + + uint32_t off = 0; + for (size_t k = 0; k < ids.size(); k++) { + const std::vector& in_dims = g.cur_dims(ids[k]); + WebGPUTensor ti; + ti.dims = in_dims; + TensorMeta in_meta; + fill_tensor_meta(ti, &in_meta); + wgpuQueueWriteBuffer( + g.queue(), in_meta_bufs[k], 0, &in_meta, sizeof(in_meta)); + CatParams params = {}; + params.concat_dim = static_cast(dim); + params.off_k = off; + wgpuQueueWriteBuffer( + g.queue(), params_bufs[k], 0, ¶ms, sizeof(params)); + g.dispatch_at(dispatch_idxs[k]).workgroup_count_x = + utils::compute_1d_workgroup_count( + g.device(), in_meta.numel, wg_size, "cat(resize)"); + off += static_cast(in_dims[cdim]); + } + }; + for (size_t k = 0; k < ids.size(); k++) { + graph.add_tensor_resize_hook(ids[k], cat_resize); + } + + // Graph owns the uniforms so the resize hook can rewrite them; freed in dtor. + graph.own_uniform_buffer(out_meta_buf); + for (size_t k = 0; k < ids.size(); k++) { + graph.own_uniform_buffer(in_meta_bufs[k]); + graph.own_uniform_buffer(params_bufs[k]); + } } } // namespace From 21315bc40156725d96d452f05894208ec7d29ab0 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 23 Jul 2026 23:01:45 -0700 Subject: [PATCH 75/75] [ExecuTorch][WebGPU] Broadcast add.Tensor (mirror mul) + un-skip 6 models Pull Request resolved: https://github.com/pytorch/executorch/pull/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, 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/) --- backends/test/suite/flows/webgpu.py | 10 - backends/webgpu/runtime/WebGPUUtils.h | 57 ++++ .../webgpu/runtime/ops/adamw/AdamwStep.cpp | 84 ++---- backends/webgpu/runtime/ops/add/BinaryOp.cpp | 266 ++++++++--------- .../webgpu/runtime/ops/add/binary_add.wgsl | 47 ++- .../webgpu/runtime/ops/add/binary_add_wgsl.h | 49 +++- backends/webgpu/runtime/ops/amax/Reduce.cpp | 95 ++----- backends/webgpu/runtime/ops/amin/Reduce.cpp | 95 ++----- backends/webgpu/runtime/ops/argmax/Reduce.cpp | 93 ++---- .../runtime/ops/avg_pool2d/AvgPool2d.cpp | 77 ++--- .../webgpu/runtime/ops/binary/BinaryOp.cpp | 109 +++---- .../runtime/ops/bitwise_not/BitwiseNot.cpp | 80 ++---- backends/webgpu/runtime/ops/bmm/Bmm.cpp | 74 +---- backends/webgpu/runtime/ops/cat/Cat.cpp | 100 ++----- .../ChooseQparamsAffine.cpp | 96 ++----- .../webgpu/runtime/ops/compare/Compare.cpp | 92 ++---- .../webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp | 182 ++++-------- .../ops/conv_with_clamp/ConvWithClamp.cpp | 97 ++----- .../ops/dequantize/DequantizePerTensor.cpp | 77 ++--- .../runtime/ops/embedding/Embedding.cpp | 89 ++---- .../ops/embedding_q4gsw/EmbeddingQ4gsw.cpp | 92 ++---- backends/webgpu/runtime/ops/flip/Flip.cpp | 94 ++---- .../webgpu/runtime/ops/fused_ce/FusedCe.cpp | 168 +++-------- .../runtime/ops/grid_priors/GridPriors.cpp | 70 ++--- .../ops/grid_sampler_2d/GridSampler2d.cpp | 90 ++---- .../runtime/ops/index_select/IndexSelect.cpp | 104 ++----- backends/webgpu/runtime/ops/linear/Linear.cpp | 84 +----- .../linear_dq8ca_q4gsw/LinearDq8caQ4gsw.cpp | 134 +++------ .../ops/linear_q8ta_q8csw/LinearQ8taQ8csw.cpp | 96 ++----- .../ops/linear_qcs4w/QuantizedLinearQcs4w.cpp | 91 ++---- .../runtime/ops/logical_and/LogicalAnd.cpp | 90 ++---- .../runtime/ops/logical_or/LogicalOr.cpp | 129 +++------ backends/webgpu/runtime/ops/mm/Mm.cpp | 74 +---- backends/webgpu/runtime/ops/mul/BinaryOp.cpp | 117 ++------ .../ops/native_group_norm/GroupNorm.cpp | 65 +---- .../ops/pixel_shuffle/PixelShuffle.cpp | 80 ++---- .../webgpu/runtime/ops/q8ta_add/Q8taAdd.cpp | 89 ++---- .../runtime/ops/q8ta_conv2d/Q8taConv2d.cpp | 96 ++----- .../ops/q8ta_conv2d_dw/Q8taConv2dDw.cpp | 96 ++----- .../ops/q8ta_conv2d_pw/Q8taConv2dPw.cpp | 96 ++----- .../Q8taConv2dTransposed.cpp | 96 ++----- .../runtime/ops/q8ta_linear/Q8taLinear.cpp | 96 ++----- .../q8ta_pixel_shuffle/Q8taPixelShuffle.cpp | 80 ++---- .../webgpu/runtime/ops/q8ta_relu/Q8taRelu.cpp | 80 ++---- .../ops/quantize/QuantizePerTensor.cpp | 77 ++--- .../ops/quantized_linear/QuantizedLinear.cpp | 93 ++---- backends/webgpu/runtime/ops/repeat/Repeat.cpp | 86 ++---- .../webgpu/runtime/ops/rms_norm/RmsNorm.cpp | 108 ++----- .../runtime/ops/rope/RotaryEmbedding.cpp | 268 +++++++++--------- backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 80 ++---- .../ops/sdpa_fd_decode/SdpaFdDecode.cpp | 71 ++--- .../webgpu/runtime/ops/unary/pow_scalar.wgsl | 18 +- .../runtime/ops/unary/pow_scalar_wgsl.h | 20 +- backends/webgpu/test/op_tests/cases.py | 23 +- 54 files changed, 1570 insertions(+), 3450 deletions(-) diff --git a/backends/test/suite/flows/webgpu.py b/backends/test/suite/flows/webgpu.py index 43fb1f572d0..ee0c24f221a 100644 --- a/backends/test/suite/flows/webgpu.py +++ b/backends/test/suite/flows/webgpu.py @@ -16,19 +16,9 @@ def _create_webgpu_flow() -> TestFlow: skip_patterns=[ "float16", "float64", # Not supported in swiftshader - # WebGPU add is elementwise-only; broadcasting add.Tensor unsupported. - "bcast_first", - "bcast_second", "hardswish", "lstm_batch_sizes", "upsample_nearest2d", - # torchvision models with broadcasting adds; resnet50 covers wide. - "mobilenet_v3_small", - "shufflenet_v2_x1_0", - "resnet50", - "vit_b_16", - "swin_v2_t", - "convnext_small", ], ) diff --git a/backends/webgpu/runtime/WebGPUUtils.h b/backends/webgpu/runtime/WebGPUUtils.h index dfcb0c1b03a..06c7f312dcd 100644 --- a/backends/webgpu/runtime/WebGPUUtils.h +++ b/backends/webgpu/runtime/WebGPUUtils.h @@ -457,6 +457,63 @@ inline ComputePipelineBundle make_compute_pipeline( return bundle; } +// Builds another pipeline + bind group from resources owned by an earlier +// bundle. The binding indices and types must match the shared bind-group +// layout. This preserves shared-shader/layout multi-pipeline construction +// without transferring ownership of those resources to the returned bundle. +inline ComputePipelineBundle make_compute_pipeline( + WGPUDevice device, + const ComputePipelineBundle& shared_resources, + const std::vector& bindings, + const WGPUConstantEntry* constants = nullptr, + size_t constant_count = 0, + const char* entry_point = "main") { + if (shared_resources.shader == nullptr || + shared_resources.bind_group_layout == nullptr || + shared_resources.pipeline_layout == nullptr) { + throw std::runtime_error( + "make_compute_pipeline: shared resources are not available"); + } + + ComputePipelineBundle bundle; + + std::vector bind_entries(bindings.size()); + for (size_t i = 0; i < bindings.size(); i++) { + bind_entries[i] = {}; + bind_entries[i].binding = bindings[i].binding; + bind_entries[i].buffer = bindings[i].buffer; + bind_entries[i].size = bindings[i].size; + } + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = shared_resources.pipeline_layout; + pipeline_desc.compute.module = shared_resources.shader; + pipeline_desc.compute.entryPoint = {entry_point, WGPU_STRLEN}; + pipeline_desc.compute.constantCount = constant_count; + pipeline_desc.compute.constants = constants; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + if (pipeline == nullptr) { + throw std::runtime_error( + "make_compute_pipeline: compute pipeline creation failed"); + } + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = shared_resources.bind_group_layout; + bg_desc.entryCount = bind_entries.size(); + bg_desc.entries = bind_entries.data(); + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + if (bind_group == nullptr) { + wgpuComputePipelineRelease(pipeline); + throw std::runtime_error( + "make_compute_pipeline: bind group creation failed"); + } + + bundle.pipeline = pipeline; + bundle.bind_group = bind_group; + return bundle; +} + // The {wg_size, stride_x} override-constant pair every 2D-spill dispatch // builds from its DispatchGrid; was hand-rolled identically at 7 call sites. inline std::array make_grid_constants( diff --git a/backends/webgpu/runtime/ops/adamw/AdamwStep.cpp b/backends/webgpu/runtime/ops/adamw/AdamwStep.cpp index 4a78694d40b..52fc4c0d471 100644 --- a/backends/webgpu/runtime/ops/adamw/AdamwStep.cpp +++ b/backends/webgpu/runtime/ops/adamw/AdamwStep.cpp @@ -97,78 +97,26 @@ void adamw_step_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(params)); graph.add_uniform_buffer_bytes(sizeof(params)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kAdamwStepWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[5] = {}; - for (uint32_t i = 0; i <= 2; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = WGPUBufferBindingType_Storage; - } - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[5] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = param.buffer; - bg_entries[0].size = param.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = m.buffer; - bg_entries[1].size = m.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = v.buffer; - bg_entries[2].size = v.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = grad.buffer; - bg_entries[3].size = grad.nbytes; - bg_entries[4].binding = 4; - bg_entries[4].buffer = uniform_buffer; - bg_entries[4].size = sizeof(params); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - - graph.add_dispatch({pipeline, bind_group, workgroup_count, "adamw_step"}); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kAdamwStepWGSL, + { + {0, WGPUBufferBindingType_Storage, param.buffer, param.nbytes}, + {1, WGPUBufferBindingType_Storage, m.buffer, m.nbytes}, + {2, WGPUBufferBindingType_Storage, v.buffer, v.nbytes}, + {3, WGPUBufferBindingType_ReadOnlyStorage, grad.buffer, grad.nbytes}, + {4, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(params)}, + }, + &wg_size_constant, + 1); + + graph.add_dispatch( + {bundle.pipeline, bundle.bind_group, workgroup_count, "adamw_step"}); + graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/add/BinaryOp.cpp b/backends/webgpu/runtime/ops/add/BinaryOp.cpp index ca49e21f046..8a8f10302f9 100644 --- a/backends/webgpu/runtime/ops/add/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/add/BinaryOp.cpp @@ -9,12 +9,14 @@ #include #include #include +#include #include #include -#include -#include +#include +#include +#include namespace executorch { namespace backends { @@ -22,16 +24,8 @@ namespace webgpu { namespace { -// Uniform buffer layout matching the WGSL Params struct. -// Must be 16-byte aligned for WebGPU uniform buffer requirements. -struct AddParams { - uint32_t num_elements; - float alpha; - uint32_t _pad[2]; // pad to 16 bytes -}; - void add_impl(WebGPUGraph& graph, const std::vector& args) { - // aten.add.Tensor args: [in1, in2, alpha, out] + // aten.add.Tensor args: [in1, in2, alpha, out]. const int in1_id = args.at(0); const int in2_id = args.at(1); const int alpha_id = args.at(2); @@ -39,7 +33,8 @@ void add_impl(WebGPUGraph& graph, const std::vector& args) { WGPUDevice device = graph.device(); - // Get alpha value (defaults to 1.0 if not a scalar) + // alpha is read once at build and fixed as a pipeline-override constant; it + // never changes on resize (resize only rewrites shapes/strides/numel). float alpha = 1.0f; if (graph.get_value_type(alpha_id) == WebGPUGraph::ValueType::Int) { alpha = static_cast(graph.get_int(alpha_id)); @@ -47,158 +42,131 @@ void add_impl(WebGPUGraph& graph, const std::vector& args) { alpha = static_cast(graph.get_double(alpha_id)); } - const auto& out_tensor = graph.get_tensor(out_id); - uint32_t num_elements = - static_cast(out_tensor.nbytes / sizeof(float)); - - uint32_t wg_size = - utils::clamp_workgroup_size(device, kBinaryAddWorkgroupSizeX); - utils::WgCount workgroup_count = - utils::compute_2d_workgroup_count(device, num_elements, wg_size, "add"); - - WGPUConstantEntry wg_size_constant = {}; - wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = static_cast(wg_size); - - // Create uniform buffer for params - AddParams params = {}; - params.num_elements = num_elements; - params.alpha = alpha; - - WGPUBufferDescriptor uniform_desc = {}; - uniform_desc.size = sizeof(AddParams); - uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; - uniform_desc.mappedAtCreation = true; - WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc); - void* mapped = wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(AddParams)); - std::memcpy(mapped, ¶ms, sizeof(AddParams)); - wgpuBufferUnmap(uniform_buffer); - - graph.add_uniform_buffer_bytes(sizeof(AddParams)); - - // Create shader module from built-in WGSL source - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kBinaryAddWGSL, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Create bind group layout: 3 storage buffers + 1 uniform - WGPUBindGroupLayoutEntry entries[4] = {}; - - // input1 - storage buffer, read-only - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - // input2 - storage buffer, read-only - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - // output - storage buffer, read-write - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Storage; - - // params - uniform buffer - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - // Create pipeline layout - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - // Create compute pipeline - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - // Create bind group with actual buffers const auto& in1_tensor = graph.get_tensor(in1_id); const auto& in2_tensor = graph.get_tensor(in2_id); + const auto& out_tensor = graph.get_tensor(out_id); - WGPUBindGroupEntry bg_entries[4] = {}; - - bg_entries[0].binding = 0; - bg_entries[0].buffer = in1_tensor.buffer; - bg_entries[0].size = in1_tensor.nbytes; - - bg_entries[1].binding = 1; - bg_entries[1].buffer = in2_tensor.buffer; - bg_entries[1].size = in2_tensor.nbytes; - - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_tensor.buffer; - bg_entries[2].size = out_tensor.nbytes; - - bg_entries[3].binding = 3; - bg_entries[3].buffer = uniform_buffer; - bg_entries[3].size = sizeof(AddParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + // Rank guard (shared TensorMeta cap). + if (out_tensor.dims.size() > kTensorMetaMaxNdim || + in1_tensor.dims.size() > kTensorMetaMaxNdim || + in2_tensor.dims.size() > kTensorMetaMaxNdim) { + throw std::runtime_error("add: tensor rank exceeds 8 (MAX_NDIM)"); + } - graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "", workgroup_count.y}); - const size_t dispatch_idx = graph.num_dispatches() - 1; + const uint32_t out_ndim = static_cast(out_tensor.dims.size()); + + // 3 per-tensor meta uniforms (mirror mul); inputs broadcast-aligned. + TensorMeta out_meta; + TensorMeta in1_meta; + TensorMeta in2_meta; + fill_tensor_meta_broadcast(out_tensor, out_ndim, &out_meta); + fill_tensor_meta_broadcast(in1_tensor, out_ndim, &in1_meta); + fill_tensor_meta_broadcast(in2_tensor, out_ndim, &in2_meta); + + // fp32-only: nbytes must equal numel * 4 for every operand. + if (out_tensor.nbytes != + static_cast(out_meta.numel) * sizeof(float) || + in1_tensor.nbytes != + static_cast(in1_meta.numel) * sizeof(float) || + in2_tensor.nbytes != + static_cast(in2_meta.numel) * sizeof(float)) { + throw std::runtime_error("add: non-fp32 operand (nbytes != numel * 4)"); + } - // Dynamic shapes: recompute numel/dispatch; out follows the larger operand. - WGPUBuffer params_buf = uniform_buffer; + uint32_t wg_size = + utils::clamp_workgroup_size(device, kBinaryAddWorkgroupSizeX); + utils::WgCount workgroup_count = + utils::compute_2d_workgroup_count(device, out_meta.numel, wg_size, "add"); + + // Two pipeline-override constants: workgroup size + the fixed alpha. + WGPUConstantEntry constants[2] = {}; + constants[0].key = {"wg_size", WGPU_STRLEN}; + constants[0].value = static_cast(wg_size); + constants[1].key = {"alpha", WGPU_STRLEN}; + constants[1].value = static_cast(alpha); + + WGPUBuffer out_meta_buf = + utils::make_uniform(device, &out_meta, sizeof(TensorMeta)); + WGPUBuffer in1_meta_buf = + utils::make_uniform(device, &in1_meta, sizeof(TensorMeta)); + WGPUBuffer in2_meta_buf = + utils::make_uniform(device, &in2_meta, sizeof(TensorMeta)); + graph.add_uniform_buffer_bytes(3 * sizeof(TensorMeta)); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kBinaryAddWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in1_tensor.buffer, + in1_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in2_tensor.buffer, + in2_tensor.nbytes}, + {2, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {3, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, in1_meta_buf, sizeof(TensorMeta)}, + {5, WGPUBufferBindingType_Uniform, in2_meta_buf, sizeof(TensorMeta)}, + }, + constants, + 2); + + const size_t dispatch_idx = graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "add", + workgroup_count.y}); + + // Dynamic shapes: rebuild all 3 broadcast TensorMeta UBOs + dispatch. alpha + // is a pipeline constant captured at build, so it is not rewritten here. + WGPUBuffer o_buf = out_meta_buf, a_buf = in1_meta_buf, b_buf = in2_meta_buf; auto add_resize = - [in1_id, in2_id, out_id, alpha, wg_size, dispatch_idx, params_buf]( + [in1_id, in2_id, out_id, wg_size, dispatch_idx, o_buf, a_buf, b_buf]( WebGPUGraph& g) { - const auto& d1 = g.cur_dims(in1_id); - const auto& d2 = g.cur_dims(in2_id); - const uint64_t n1 = utils::numel_of(d1); - const uint64_t n2 = utils::numel_of(d2); - const uint64_t numel = n2 > n1 ? n2 : n1; - const uint64_t n_min = n2 > n1 ? n1 : n2; - // The flat add follows the larger operand and broadcasts the smaller; - // valid only when the smaller tiles evenly into it (rejects e.g. [4,1] - // vs [1,3], whose true [4,3] result this flat kernel cannot produce). - if (n_min == 0u || numel % n_min != 0u) { - throw std::runtime_error( - "add(resize): operands are not broadcast-compatible by numel"); + const auto& a = g.cur_dims(in1_id); + const auto& b = g.cur_dims(in2_id); + const size_t r = std::max(a.size(), b.size()); + std::vector out_d(r, 1); + for (size_t i = 0; i < r; i++) { + const int64_t av = (i + a.size() < r) ? 1 : a[i - (r - a.size())]; + const int64_t bv = (i + b.size() < r) ? 1 : b[i - (r - b.size())]; + if (av != bv && av != 1 && bv != 1) { + throw std::runtime_error( + "add(resize): operands are not broadcast-compatible"); + } + out_d[i] = av > bv ? av : bv; } - g.set_cur_dims(out_id, n2 > n1 ? d2 : d1); - AddParams p = {}; - p.num_elements = static_cast(numel); - p.alpha = alpha; - wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.set_cur_dims(out_id, out_d); + const uint32_t out_ndim = static_cast(r); + WebGPUTensor ta, tb, to; + ta.dims = a; + tb.dims = b; + to.dims = out_d; + TensorMeta om, am, bm; + fill_tensor_meta_broadcast(to, out_ndim, &om); + fill_tensor_meta_broadcast(ta, out_ndim, &am); + fill_tensor_meta_broadcast(tb, out_ndim, &bm); + wgpuQueueWriteBuffer(g.queue(), o_buf, 0, &om, sizeof(om)); + wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am)); + wgpuQueueWriteBuffer(g.queue(), b_buf, 0, &bm, sizeof(bm)); const utils::WgCount wgc = utils::compute_2d_workgroup_count( - g.device(), static_cast(numel), wg_size, "add(resize)"); + g.device(), om.numel, wg_size, "add(resize)"); g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }; graph.add_tensor_resize_hook(in1_id, add_resize); graph.add_tensor_resize_hook(in2_id, add_resize); - // Release intermediate objects (pipeline and bind_group are kept by dispatch) - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); - // Graph owns it so a resize hook can rewrite it; freed in the dtor. - graph.own_uniform_buffer(uniform_buffer); + // Graph owns them so a resize hook can rewrite them; freed in the dtor. + graph.own_uniform_buffer(out_meta_buf); + graph.own_uniform_buffer(in1_meta_buf); + graph.own_uniform_buffer(in2_meta_buf); } } // namespace diff --git a/backends/webgpu/runtime/ops/add/binary_add.wgsl b/backends/webgpu/runtime/ops/add/binary_add.wgsl index c1cb9d5ffd7..cd6c5291a4c 100644 --- a/backends/webgpu/runtime/ops/add/binary_add.wgsl +++ b/backends/webgpu/runtime/ops/add/binary_add.wgsl @@ -2,21 +2,52 @@ @group(0) @binding(1) var input2: array; @group(0) @binding(2) var output: array; -struct Params { - num_elements: u32, - alpha: f32, +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: array, 2>, + strides: array, 2>, } -@group(0) @binding(3) var params: Params; +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in1_meta: TensorMeta; +@group(0) @binding(5) var in2_meta: TensorMeta; -override wg_size: u32 = 256; +override wg_size: u32 = 256u; +// add.Tensor alpha; read once from the graph and fixed at build (never resized). +override alpha: f32 = 1.0; -@compute @workgroup_size(wg_size) +@compute @workgroup_size(wg_size, 1, 1) fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { + if (idx >= out_meta.numel) { return; } - output[idx] = input1[idx] + params.alpha * input2[idx]; + + // Fast path: every input dim matches the output dim -> elementwise. + var same = true; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { + same = false; + } + } + if (same) { + output[idx] = input1[idx] + alpha * input2[idx]; + return; + } + + // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. + var rem = idx; + var l1: u32 = 0u; + var l2: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; + } + output[idx] = input1[l1] + alpha * input2[l2]; } diff --git a/backends/webgpu/runtime/ops/add/binary_add_wgsl.h b/backends/webgpu/runtime/ops/add/binary_add_wgsl.h index 829407bb383..ed7ab2cebfe 100644 --- a/backends/webgpu/runtime/ops/add/binary_add_wgsl.h +++ b/backends/webgpu/runtime/ops/add/binary_add_wgsl.h @@ -13,29 +13,60 @@ namespace executorch::backends::webgpu { // @generated from binary_add.wgsl - DO NOT EDIT. -// wgsl-sha256: e66bd67465c2a0296e09668df54f87605a4c91015a615f3734cdd0f140a74477 +// wgsl-sha256: 4a66de0d15fc8f31de0aa97ad5c1dc7e431ec690e829eaf222e77c1f7aabf01d inline constexpr const char* kBinaryAddWGSL = R"( @group(0) @binding(0) var input1: array; @group(0) @binding(1) var input2: array; @group(0) @binding(2) var output: array; -struct Params { - num_elements: u32, - alpha: f32, +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: array, 2>, + strides: array, 2>, } -@group(0) @binding(3) var params: Params; +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in1_meta: TensorMeta; +@group(0) @binding(5) var in2_meta: TensorMeta; -override wg_size: u32 = 256; +override wg_size: u32 = 256u; +// add.Tensor alpha; read once from the graph and fixed at build (never resized). +override alpha: f32 = 1.0; -@compute @workgroup_size(wg_size) +@compute @workgroup_size(wg_size, 1, 1) fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { + if (idx >= out_meta.numel) { return; } - output[idx] = input1[idx] + params.alpha * input2[idx]; + + // Fast path: every input dim matches the output dim -> elementwise. + var same = true; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { + same = false; + } + } + if (same) { + output[idx] = input1[idx] + alpha * input2[idx]; + return; + } + + // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. + var rem = idx; + var l1: u32 = 0u; + var l2: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; + } + output[idx] = input1[l1] + alpha * input2[l2]; } )"; diff --git a/backends/webgpu/runtime/ops/amax/Reduce.cpp b/backends/webgpu/runtime/ops/amax/Reduce.cpp index 4e76bdf369f..14675d97e96 100644 --- a/backends/webgpu/runtime/ops/amax/Reduce.cpp +++ b/backends/webgpu/runtime/ops/amax/Reduce.cpp @@ -84,71 +84,32 @@ void amax_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(AmaxParams)); graph.add_uniform_buffer_bytes(sizeof(AmaxParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kAmaxWGSL, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group layout: input (read storage) + output (storage) + params. - WGPUBindGroupLayoutEntry entries[3] = {}; - - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[3] = {}; - - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - - bg_entries[2].binding = 2; - bg_entries[2].buffer = uniform_buffer; - bg_entries[2].size = sizeof(AmaxParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kAmaxWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(AmaxParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "amax", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "amax", + workgroup_count.y}); // Dynamic shapes: recompute reduce_size (last dim) + num_rows + dispatch. const bool keepdim = graph.get_bool(args.at(2)); @@ -178,16 +139,12 @@ void amax_impl(WebGPUGraph& graph, const std::vector& args) { p.num_rows = rows; p.reduce_size = rsize; wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); - const utils::WgCount wgc = utils::compute_2d_workgroup_count( - g.device(), rows, 1, "amax"); + const utils::WgCount wgc = + utils::compute_2d_workgroup_count(g.device(), rows, 1, "amax"); g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }); - // Release intermediates (pipeline + bind_group are kept by dispatch). - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/amin/Reduce.cpp b/backends/webgpu/runtime/ops/amin/Reduce.cpp index ade0ef5b583..fbe574fdf0b 100644 --- a/backends/webgpu/runtime/ops/amin/Reduce.cpp +++ b/backends/webgpu/runtime/ops/amin/Reduce.cpp @@ -84,71 +84,32 @@ void amin_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(AminParams)); graph.add_uniform_buffer_bytes(sizeof(AminParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kAminWGSL, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group layout: input (read storage) + output (storage) + params. - WGPUBindGroupLayoutEntry entries[3] = {}; - - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[3] = {}; - - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - - bg_entries[2].binding = 2; - bg_entries[2].buffer = uniform_buffer; - bg_entries[2].size = sizeof(AminParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kAminWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(AminParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "amin", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "amin", + workgroup_count.y}); // Dynamic shapes: recompute reduce_size (last dim) + num_rows + dispatch. const bool keepdim = graph.get_bool(args.at(2)); @@ -178,16 +139,12 @@ void amin_impl(WebGPUGraph& graph, const std::vector& args) { p.num_rows = rows; p.reduce_size = rsize; wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); - const utils::WgCount wgc = utils::compute_2d_workgroup_count( - g.device(), rows, 1, "amin"); + const utils::WgCount wgc = + utils::compute_2d_workgroup_count(g.device(), rows, 1, "amin"); g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }); - // Release intermediates (pipeline + bind_group are kept by dispatch). - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/argmax/Reduce.cpp b/backends/webgpu/runtime/ops/argmax/Reduce.cpp index 27318ba0ab9..7b5b714f1d2 100644 --- a/backends/webgpu/runtime/ops/argmax/Reduce.cpp +++ b/backends/webgpu/runtime/ops/argmax/Reduce.cpp @@ -81,8 +81,8 @@ void arg_reduce_impl( uint32_t wg_size = utils::clamp_workgroup_size(device, kArgReduceWorkgroupSizeX); // One workgroup per row (cooperative reduction); grid = num_rows workgroups. - utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( - device, num_rows, 1, "arg_reduce"); + utils::WgCount workgroup_count = + utils::compute_2d_workgroup_count(device, num_rows, 1, "arg_reduce"); WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; @@ -97,72 +97,29 @@ void arg_reduce_impl( utils::make_uniform(device, ¶ms, sizeof(ArgReduceParams)); graph.add_uniform_buffer_bytes(sizeof(ArgReduceParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kArgReduceWGSL, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group layout: input (read storage) + output (storage) + params. - WGPUBindGroupLayoutEntry entries[3] = {}; - - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[3] = {}; - - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - - bg_entries[2].binding = 2; - bg_entries[2].buffer = uniform_buffer; - bg_entries[2].size = sizeof(ArgReduceParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kArgReduceWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(ArgReduceParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "arg_reduce", workgroup_count.y}); @@ -199,10 +156,6 @@ void arg_reduce_impl( g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }); - // Release intermediates (pipeline + bind_group are kept by dispatch). - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/avg_pool2d/AvgPool2d.cpp b/backends/webgpu/runtime/ops/avg_pool2d/AvgPool2d.cpp index a172bd71ec0..24da24c15f5 100644 --- a/backends/webgpu/runtime/ops/avg_pool2d/AvgPool2d.cpp +++ b/backends/webgpu/runtime/ops/avg_pool2d/AvgPool2d.cpp @@ -159,64 +159,26 @@ void avg_pool2d_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(AvgPoolParams)); graph.add_uniform_buffer_bytes(sizeof(AvgPoolParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kAvgPool2dWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[3] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[3] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = params_buf; - bg_entries[2].size = sizeof(AvgPoolParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kAvgPool2dWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, WGPUBufferBindingType_Uniform, params_buf, sizeof(AvgPoolParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "avg_pool2d", workgroup_count.y}); @@ -280,9 +242,6 @@ void avg_pool2d_impl(WebGPUGraph& graph, const std::vector& args) { g.set_cur_dims(out_id, out_d); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/binary/BinaryOp.cpp b/backends/webgpu/runtime/ops/binary/BinaryOp.cpp index 84436071e27..1e21f6e01e9 100644 --- a/backends/webgpu/runtime/ops/binary/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/binary/BinaryOp.cpp @@ -72,8 +72,8 @@ void add_binary_broadcast_op( } uint32_t wg_size = utils::clamp_workgroup_size(device, wg_size_default); - utils::WgCount workgroup_count = - utils::compute_2d_workgroup_count(device, out_meta.numel, wg_size, op_name); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, out_meta.numel, wg_size, op_name); WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; @@ -87,77 +87,47 @@ void add_binary_broadcast_op( utils::make_uniform(device, &in2_meta, sizeof(TensorMeta)); graph.add_uniform_buffer_bytes(3 * sizeof(TensorMeta)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {wgsl_code, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group: in1, in2, out (rw), out_meta, in1_meta, in2_meta (3 uniforms). - WGPUBindGroupLayoutEntry entries[6] = {}; - for (int i = 0; i < 6; i++) { - entries[i].binding = static_cast(i); - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = (i == 2) ? WGPUBufferBindingType_Storage - : (i >= 3) ? WGPUBufferBindingType_Uniform - : WGPUBufferBindingType_ReadOnlyStorage; - } - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 6; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[6] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in1_tensor.buffer; - bg_entries[0].size = in1_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = in2_tensor.buffer; - bg_entries[1].size = in2_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_tensor.buffer; - bg_entries[2].size = out_tensor.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = out_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - bg_entries[4].binding = 4; - bg_entries[4].buffer = in1_meta_buf; - bg_entries[4].size = sizeof(TensorMeta); - bg_entries[5].binding = 5; - bg_entries[5].buffer = in2_meta_buf; - bg_entries[5].size = sizeof(TensorMeta); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 6; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + wgsl_code, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in1_tensor.buffer, + in1_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in2_tensor.buffer, + in2_tensor.nbytes}, + {2, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {3, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, in1_meta_buf, sizeof(TensorMeta)}, + {5, WGPUBufferBindingType_Uniform, in2_meta_buf, sizeof(TensorMeta)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, op_name, workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + op_name, + workgroup_count.y}); // Dynamic shapes: rebuild all 3 broadcast TensorMeta UBOs + dispatch. WGPUBuffer o_buf = out_meta_buf, a_buf = in1_meta_buf, b_buf = in2_meta_buf; - auto resize = [in1_id, in2_id, out_id, wg_size, dispatch_idx, o_buf, a_buf, - b_buf, name](WebGPUGraph& g) { + auto resize = [in1_id, + in2_id, + out_id, + wg_size, + dispatch_idx, + o_buf, + a_buf, + b_buf, + name](WebGPUGraph& g) { const auto& a = g.cur_dims(in1_id); const auto& b = g.cur_dims(in2_id); const size_t r = std::max(a.size(), b.size()); @@ -193,9 +163,6 @@ void add_binary_broadcast_op( graph.add_tensor_resize_hook(in1_id, resize); graph.add_tensor_resize_hook(in2_id, resize); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns them so a resize hook can rewrite them; freed in the dtor. graph.own_uniform_buffer(out_meta_buf); graph.own_uniform_buffer(in1_meta_buf); diff --git a/backends/webgpu/runtime/ops/bitwise_not/BitwiseNot.cpp b/backends/webgpu/runtime/ops/bitwise_not/BitwiseNot.cpp index 22ab8daa654..fc02ea01c55 100644 --- a/backends/webgpu/runtime/ops/bitwise_not/BitwiseNot.cpp +++ b/backends/webgpu/runtime/ops/bitwise_not/BitwiseNot.cpp @@ -76,65 +76,30 @@ void bitwise_not_op(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(BitwiseNotParams)); graph.add_uniform_buffer_bytes(sizeof(BitwiseNotParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kBitwiseNotWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // out (rw storage) + a (ro storage) + params (uniform). - WGPUBindGroupLayoutEntry entries[3] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg[3] = {}; - bg[0].binding = 0; - bg[0].buffer = out_tensor.buffer; - bg[0].size = out_tensor.nbytes; - bg[1].binding = 1; - bg[1].buffer = a_tensor.buffer; - bg[1].size = a_tensor.nbytes; - bg[2].binding = 2; - bg[2].buffer = uniform_buffer; - bg[2].size = sizeof(BitwiseNotParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kBitwiseNotWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + a_tensor.buffer, + a_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(BitwiseNotParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "bitwise_not", workgroup_count.y}); @@ -160,9 +125,6 @@ void bitwise_not_op(WebGPUGraph& graph, const std::vector& args) { }; graph.add_tensor_resize_hook(a_id, resize); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/bmm/Bmm.cpp b/backends/webgpu/runtime/ops/bmm/Bmm.cpp index 04c41a51c2e..e9bf9d17c4d 100644 --- a/backends/webgpu/runtime/ops/bmm/Bmm.cpp +++ b/backends/webgpu/runtime/ops/bmm/Bmm.cpp @@ -89,69 +89,20 @@ void bmm_impl(WebGPUGraph& graph, const std::vector& args) { // vec4 path when K and N are multiples of 4 (wider 16B loads); else scalar. const bool use_vec4 = (K % 4u == 0u) && (N % 4u == 0u); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {use_vec4 ? kBmmVec4WGSL : kBmmTiledWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[4] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Storage; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - // Tiled kernel has a fixed @workgroup_size(8, 8, 1) — no override constant. - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[4] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = a.buffer; - bg_entries[0].size = a.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = b.buffer; - bg_entries[1].size = b.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out.buffer; - bg_entries[2].size = out.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = uniform_buffer; - bg_entries[3].size = sizeof(BmmParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + use_vec4 ? kBmmVec4WGSL : kBmmTiledWGSL, + { + {0, WGPUBufferBindingType_ReadOnlyStorage, a.buffer, a.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, b.buffer, b.nbytes}, + {2, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {3, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(BmmParams)}, + }); WebGPUDispatch dispatch; - dispatch.pipeline = pipeline; - dispatch.bind_group = bind_group; + dispatch.pipeline = bundle.pipeline; + dispatch.bind_group = bundle.bind_group; dispatch.workgroup_count_x = dispatch_x; dispatch.workgroup_count_y = dispatch_y; dispatch.kernel_name = "bmm"; @@ -209,9 +160,6 @@ void bmm_impl(WebGPUGraph& graph, const std::vector& args) { static_cast(live_n)}); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/cat/Cat.cpp b/backends/webgpu/runtime/ops/cat/Cat.cpp index 28c327ae795..3c3c34d571c 100644 --- a/backends/webgpu/runtime/ops/cat/Cat.cpp +++ b/backends/webgpu/runtime/ops/cat/Cat.cpp @@ -15,7 +15,9 @@ #include #include +#include #include +#include #include namespace executorch::backends::webgpu { @@ -113,47 +115,14 @@ void cat_impl(WebGPUGraph& graph, const std::vector& args) { wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - // Shared shader/layout; fresh pipeline+bind group per input (no double-free). - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kCatWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - // Collected for the dynamic-shape resize hook (rewrites these on resize). std::vector in_meta_bufs(ids.size()); std::vector params_bufs(ids.size()); std::vector dispatch_idxs(ids.size()); + // The first bundle owns the shader/layout resources. Later calls reuse them + // while still returning one distinct pipeline + bind group per input. + std::optional shared_resources; uint32_t off_k = 0; for (size_t k = 0; k < ids.size(); k++) { const auto& in_tensor = graph.get_tensor(ids[k]); @@ -168,49 +137,36 @@ void cat_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(CatParams)); graph.add_uniform_buffer_bytes(sizeof(TensorMeta) + sizeof(CatParams)); - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[5] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_meta_buf; - bg_entries[2].size = sizeof(TensorMeta); - bg_entries[3].binding = 3; - bg_entries[3].buffer = in_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - bg_entries[4].binding = 4; - bg_entries[4].buffer = params_buf; - bg_entries[4].size = sizeof(CatParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + const std::vector bindings = { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {3, WGPUBufferBindingType_Uniform, in_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, params_buf, sizeof(CatParams)}, + }; + utils::ComputePipelineBundle bundle = shared_resources.has_value() + ? utils::make_compute_pipeline( + device, *shared_resources, bindings, &wg_size_constant, 1) + : utils::make_compute_pipeline( + device, kCatWGSL, bindings, &wg_size_constant, 1); in_meta_bufs[k] = in_meta_buf; params_bufs[k] = params_buf; - dispatch_idxs[k] = graph.add_dispatch({pipeline, bind_group, wg_counts[k]}); + dispatch_idxs[k] = + graph.add_dispatch({bundle.pipeline, bundle.bind_group, wg_counts[k]}); + if (!shared_resources.has_value()) { + shared_resources.emplace(std::move(bundle)); + } // Uniforms kept alive (owned below) so the resize hook can rewrite them. off_k += static_cast(in_tensor.dims[dim]); } - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); - // Dynamic shapes: cat bakes strides/numel/off_k/workgroup_count from the max // (build) shape, so under a smaller live shape it would scatter with stale // strides and loop over the stale numel. Recompute every input's in_meta + diff --git a/backends/webgpu/runtime/ops/choose_qparams_affine/ChooseQparamsAffine.cpp b/backends/webgpu/runtime/ops/choose_qparams_affine/ChooseQparamsAffine.cpp index d896cd13409..f593334cb2a 100644 --- a/backends/webgpu/runtime/ops/choose_qparams_affine/ChooseQparamsAffine.cpp +++ b/backends/webgpu/runtime/ops/choose_qparams_affine/ChooseQparamsAffine.cpp @@ -86,9 +86,9 @@ void choose_qparams_affine_impl( if (scale_t.nbytes != num_rows * sizeof(float)) { throw std::runtime_error("choose_qparams_affine: scale must be fp32[rows]"); } - // zp is int8[rows] (elem_size 1), packed 4-per-u32 in the shader. int8 buffers - // are allocated max(nbytes, 4); M<=4 pads to one word, M%4==0 is word-exact. - // Other M (5,6,7,...) would overflow the M-byte buffer -> reject. + // zp is int8[rows] (elem_size 1), packed 4-per-u32 in the shader. int8 + // buffers are allocated max(nbytes, 4); M<=4 pads to one word, M%4==0 is + // word-exact. Other M (5,6,7,...) would overflow the M-byte buffer -> reject. if (!zp_t.is_int8 || zp_t.nbytes != num_rows) { throw std::runtime_error("choose_qparams_affine: zp must be int8[rows]"); } @@ -154,78 +154,33 @@ void choose_qparams_affine_impl( utils::make_uniform(device, ¶ms, sizeof(ChooseQParamsParams)); graph.add_uniform_buffer_bytes(sizeof(ChooseQParamsParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kChooseQparamsAffineWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[4] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Storage; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg[4] = {}; - bg[0].binding = 0; - bg[0].buffer = in.buffer; - bg[0].size = in.nbytes; - bg[1].binding = 1; - bg[1].buffer = scale_t.buffer; - bg[1].size = scale_t.nbytes; - bg[2].binding = 2; - bg[2].buffer = zp_t.buffer; - // Bind word-aligned (buffer is >= max(nbytes,4); array needs a mult of 4). - bg[2].size = ((zp_t.nbytes + 3u) / 4u) * 4u; - bg[3].binding = 3; - bg[3].buffer = params_buf; - bg[3].size = sizeof(ChooseQParamsParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kChooseQparamsAffineWGSL, + { + {0, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {1, WGPUBufferBindingType_Storage, scale_t.buffer, scale_t.nbytes}, + {2, + WGPUBufferBindingType_Storage, + zp_t.buffer, + // Bind word-aligned (buffer is >= max(nbytes,4); array needs a + // mult of 4). + ((zp_t.nbytes + 3u) / 4u) * 4u}, + {3, + WGPUBufferBindingType_Uniform, + params_buf, + sizeof(ChooseQParamsParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "choose_qparams_affine", workgroup_count.y}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } @@ -233,8 +188,7 @@ void choose_qparams_affine_impl( WEBGPU_REGISTER_OPERATORS { WEBGPU_REGISTER_OP( - torchao.choose_qparams_affine.default, - choose_qparams_affine_impl); + torchao.choose_qparams_affine.default, choose_qparams_affine_impl); } } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/compare/Compare.cpp b/backends/webgpu/runtime/ops/compare/Compare.cpp index 8ecd35a44ff..99200cb2425 100644 --- a/backends/webgpu/runtime/ops/compare/Compare.cpp +++ b/backends/webgpu/runtime/ops/compare/Compare.cpp @@ -89,70 +89,37 @@ void compare_impl( utils::make_uniform(device, ¶ms, sizeof(CompareParams)); graph.add_uniform_buffer_bytes(sizeof(CompareParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kCompareWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // out (rw storage) + in1/in2 (ro storage) + params (uniform). - WGPUBindGroupLayoutEntry entries[4] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg[4] = {}; - bg[0].binding = 0; - bg[0].buffer = out_tensor.buffer; - bg[0].size = out_tensor.nbytes; - bg[1].binding = 1; - bg[1].buffer = in1_tensor.buffer; - bg[1].size = in1_tensor.nbytes; - bg[2].binding = 2; - bg[2].buffer = in2_tensor.buffer; - bg[2].size = in2_tensor.nbytes; - bg[3].binding = 3; - bg[3].buffer = uniform_buffer; - bg[3].size = sizeof(CompareParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kCompareWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in1_tensor.buffer, + in1_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + in2_tensor.buffer, + in2_tensor.nbytes}, + {3, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(CompareParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "compare", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "compare", + workgroup_count.y}); // Dynamic shapes: recompute numel/dispatch; out follows in1 (same-shape). WGPUBuffer params_buf = uniform_buffer; @@ -177,9 +144,6 @@ void compare_impl( graph.add_tensor_resize_hook(in1_id, resize); graph.add_tensor_resize_hook(in2_id, resize); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp index f876b095a22..ae86ccffba8 100644 --- a/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp +++ b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp @@ -126,75 +126,39 @@ void add_conv1d_pw_node( utils::make_uniform(device, ¶ms, sizeof(Conv1dPwParams)); graph.add_uniform_buffer_bytes(sizeof(Conv1dPwParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kConv1dPwWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[5] = {}; - for (int i = 0; i < 4; i++) { - entries[i].binding = static_cast(i); - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = (i == 1) ? WGPUBufferBindingType_Storage - : WGPUBufferBindingType_ReadOnlyStorage; - } - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - WGPUBuffer bias_buf = has_bias ? graph.get_tensor(bias_id).buffer : weight_tensor.buffer; uint64_t bias_sz = has_bias ? graph.get_tensor(bias_id).nbytes : weight_tensor.nbytes; - WGPUBindGroupEntry bg_entries[5] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = weight_tensor.buffer; - bg_entries[2].size = weight_tensor.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = bias_buf; - bg_entries[3].size = bias_sz; - bg_entries[4].binding = 4; - bg_entries[4].buffer = params_buf; - bg_entries[4].size = sizeof(Conv1dPwParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kConv1dPwWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight_tensor.buffer, + weight_tensor.nbytes}, + {3, WGPUBufferBindingType_ReadOnlyStorage, bias_buf, bias_sz}, + {4, + WGPUBufferBindingType_Uniform, + params_buf, + sizeof(Conv1dPwParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "conv1d_pw", workgroup_count.y}); @@ -235,9 +199,6 @@ void add_conv1d_pw_node( g.set_cur_dims(out_id, out_d); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(params_buf); } @@ -357,76 +318,40 @@ void convolution_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(Conv1dDwParams)); graph.add_uniform_buffer_bytes(sizeof(Conv1dDwParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kConv1dDwWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[5] = {}; - for (int i = 0; i < 4; i++) { - entries[i].binding = static_cast(i); - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = (i == 1) ? WGPUBufferBindingType_Storage - : WGPUBufferBindingType_ReadOnlyStorage; - } - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - // No-bias binds the weight buffer as an unread placeholder (has_bias gates). WGPUBuffer bias_buf = has_bias ? graph.get_tensor(bias_id).buffer : weight_tensor.buffer; uint64_t bias_sz = has_bias ? graph.get_tensor(bias_id).nbytes : weight_tensor.nbytes; - WGPUBindGroupEntry bg_entries[5] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = weight_tensor.buffer; - bg_entries[2].size = weight_tensor.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = bias_buf; - bg_entries[3].size = bias_sz; - bg_entries[4].binding = 4; - bg_entries[4].buffer = params_buf; - bg_entries[4].size = sizeof(Conv1dDwParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kConv1dDwWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight_tensor.buffer, + weight_tensor.nbytes}, + {3, WGPUBufferBindingType_ReadOnlyStorage, bias_buf, bias_sz}, + {4, + WGPUBufferBindingType_Uniform, + params_buf, + sizeof(Conv1dDwParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "conv1d_dw", workgroup_count.y}); @@ -484,9 +409,6 @@ void convolution_impl(WebGPUGraph& graph, const std::vector& args) { g.set_cur_dims(out_id, out_d); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/conv_with_clamp/ConvWithClamp.cpp b/backends/webgpu/runtime/ops/conv_with_clamp/ConvWithClamp.cpp index 42f989009f4..4a0a9760aeb 100644 --- a/backends/webgpu/runtime/ops/conv_with_clamp/ConvWithClamp.cpp +++ b/backends/webgpu/runtime/ops/conv_with_clamp/ConvWithClamp.cpp @@ -154,8 +154,7 @@ void conv_with_clamp_impl(WebGPUGraph& graph, const std::vector& args) { const uint64_t numel = N * OC * H_out * W_out; const uint64_t ug = static_cast(groups); // Grouped weight is [OC, IC/groups, Kh, Kw]; ic_per_group = weight.dims[1]. - const uint64_t ic_per_group = - static_cast(weight_tensor.dims.at(1)); + const uint64_t ic_per_group = static_cast(weight_tensor.dims.at(1)); if (IC == 0 || numel == 0 || numel > UINT32_MAX || IC % ug != 0 || OC % ug != 0 || ic_per_group * ug != IC) { throw std::runtime_error("conv_with_clamp: bad shape (IC/numel/groups)"); @@ -208,88 +207,46 @@ void conv_with_clamp_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(ConvWithClampParams)); graph.add_uniform_buffer_bytes(sizeof(ConvWithClampParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kConvWithClampWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // out (rw) + in/weight/bias (ro storage) + params (uniform). - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - for (uint32_t i = 1; i <= 3; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - } - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - // No-bias: bind weight as an unread placeholder (has_bias gates the read). WGPUBuffer bias_buf = has_bias ? graph.get_tensor(bias_id).buffer : weight_tensor.buffer; const uint64_t bias_size = has_bias ? graph.get_tensor(bias_id).nbytes : weight_tensor.nbytes; - WGPUBindGroupEntry bg[5] = {}; - bg[0].binding = 0; - bg[0].buffer = out_tensor.buffer; - bg[0].size = out_tensor.nbytes; - bg[1].binding = 1; - bg[1].buffer = in_tensor.buffer; - bg[1].size = in_tensor.nbytes; - bg[2].binding = 2; - bg[2].buffer = weight_tensor.buffer; - bg[2].size = weight_tensor.nbytes; - bg[3].binding = 3; - bg[3].buffer = bias_buf; - bg[3].size = bias_size; - bg[4].binding = 4; - bg[4].buffer = params_buf; - bg[4].size = sizeof(ConvWithClampParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kConvWithClampWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight_tensor.buffer, + weight_tensor.nbytes}, + {3, WGPUBufferBindingType_ReadOnlyStorage, bias_buf, bias_size}, + {4, + WGPUBufferBindingType_Uniform, + params_buf, + sizeof(ConvWithClampParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "conv_with_clamp", workgroup_count.y}); // conv2d is static-shape-only: no tensor resize hook is registered, so the // output spatial dims stay fixed at their build-time (serialized) values. - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/dequantize/DequantizePerTensor.cpp b/backends/webgpu/runtime/ops/dequantize/DequantizePerTensor.cpp index 83f9b62ce5a..e6a413b4b02 100644 --- a/backends/webgpu/runtime/ops/dequantize/DequantizePerTensor.cpp +++ b/backends/webgpu/runtime/ops/dequantize/DequantizePerTensor.cpp @@ -93,71 +93,30 @@ void dequantize_per_tensor_impl( utils::make_uniform(device, ¶ms, sizeof(DequantParams)); graph.add_uniform_buffer_bytes(sizeof(DequantParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kDequantizePerTensorWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[3] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[3] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = params_buf; - bg_entries[2].size = sizeof(DequantParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kDequantizePerTensorWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, WGPUBufferBindingType_Uniform, params_buf, sizeof(DequantParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "dequantize_per_tensor", workgroup_count.y}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/embedding/Embedding.cpp b/backends/webgpu/runtime/ops/embedding/Embedding.cpp index b05b1af7d9d..f4771782cb9 100644 --- a/backends/webgpu/runtime/ops/embedding/Embedding.cpp +++ b/backends/webgpu/runtime/ops/embedding/Embedding.cpp @@ -78,73 +78,34 @@ void embedding_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(EmbeddingParams)); graph.add_uniform_buffer_bytes(sizeof(EmbeddingParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kEmbeddingWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[4] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Storage; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[4] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = weight.buffer; - bg_entries[0].size = weight.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = indices.buffer; - bg_entries[1].size = indices.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out.buffer; - bg_entries[2].size = out.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = uniform_buffer; - bg_entries[3].size = sizeof(EmbeddingParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kEmbeddingWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + weight.buffer, + weight.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + indices.buffer, + indices.nbytes}, + {2, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {3, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(EmbeddingParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "", + workgroup_count.y}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); wgpuBufferRelease(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp index 6f1febf77d8..9ed4946eab6 100644 --- a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp +++ b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp @@ -197,76 +197,37 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { wgpuBufferUnmap(uniform_buffer); graph.add_uniform_buffer_bytes(sizeof(EmbeddingParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kEmbeddingQ4gswWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group layout: out (rw) + indices/weight/scales (ro storage) + uniform. - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - for (uint32_t i = 1; i <= 3; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - } - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[5] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = out.buffer; - bg_entries[0].size = out.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = indices.buffer; - bg_entries[1].size = indices.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = weight.buffer; - bg_entries[2].size = weight.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = scales.buffer; - bg_entries[3].size = scales.nbytes; - bg_entries[4].binding = 4; - bg_entries[4].buffer = uniform_buffer; - bg_entries[4].size = sizeof(EmbeddingParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kEmbeddingQ4gswWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + indices.buffer, + indices.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight.buffer, + weight.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + scales.buffer, + scales.nbytes}, + {4, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(EmbeddingParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count, "embedding_q4gsw"}); + {bundle.pipeline, bundle.bind_group, workgroup_count, "embedding_q4gsw"}); // Dynamic shapes: recompute counts/dispatch; out = indices + [embed_dim]. const uint32_t gs_u = static_cast(group_size); @@ -297,9 +258,6 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { params_buf); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/flip/Flip.cpp b/backends/webgpu/runtime/ops/flip/Flip.cpp index 20461da5406..54c59a29206 100644 --- a/backends/webgpu/runtime/ops/flip/Flip.cpp +++ b/backends/webgpu/runtime/ops/flip/Flip.cpp @@ -108,80 +108,32 @@ void flip_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(FlipParams)); graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta) + sizeof(FlipParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kFlipWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group: in, out (rw), out_meta, in_meta, params (3 uniforms). - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[5] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_meta_buf; - bg_entries[2].size = sizeof(TensorMeta); - bg_entries[3].binding = 3; - bg_entries[3].buffer = in_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - bg_entries[4].binding = 4; - bg_entries[4].buffer = params_buf; - bg_entries[4].size = sizeof(FlipParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kFlipWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {3, WGPUBufferBindingType_Uniform, in_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, params_buf, sizeof(FlipParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "flip", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "flip", + workgroup_count.y}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Drop our refs; the bind group keeps the uniforms alive until release. wgpuBufferRelease(out_meta_buf); wgpuBufferRelease(in_meta_buf); diff --git a/backends/webgpu/runtime/ops/fused_ce/FusedCe.cpp b/backends/webgpu/runtime/ops/fused_ce/FusedCe.cpp index e51052f978d..97a47346db0 100644 --- a/backends/webgpu/runtime/ops/fused_ce/FusedCe.cpp +++ b/backends/webgpu/runtime/ops/fused_ce/FusedCe.cpp @@ -40,15 +40,6 @@ struct ReduceParams { }; static_assert(sizeof(ReduceParams) == 16, "ReduceParams must be 16 bytes"); -WGPUShaderModule make_shader(WGPUDevice device, const char* wgsl) { - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {wgsl, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - return wgpuDeviceCreateShaderModule(device, &shader_desc); -} - WGPUBuffer create_uniform( WebGPUGraph& graph, WGPUDevice device, @@ -127,79 +118,38 @@ void fused_ce_impl(WebGPUGraph& graph, const std::vector& args) { utils::clamp_workgroup_size(device, kFusedCeWorkgroupSizeX); WGPUBuffer ce_uniform = create_uniform(graph, device, &ce_params, sizeof(ce_params)); - WGPUShaderModule ce_shader = make_shader(device, kFusedCeWGSL); - - WGPUBindGroupLayoutEntry ce_entries[5] = {}; - ce_entries[0].binding = 0; - ce_entries[0].visibility = WGPUShaderStage_Compute; - ce_entries[0].buffer.type = WGPUBufferBindingType_Storage; - ce_entries[1].binding = 1; - ce_entries[1].visibility = WGPUShaderStage_Compute; - ce_entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - ce_entries[2].binding = 2; - ce_entries[2].visibility = WGPUShaderStage_Compute; - ce_entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - ce_entries[3].binding = 3; - ce_entries[3].visibility = WGPUShaderStage_Compute; - ce_entries[3].buffer.type = WGPUBufferBindingType_Storage; - ce_entries[4].binding = 4; - ce_entries[4].visibility = WGPUShaderStage_Compute; - ce_entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor ce_bgl_desc = {}; - ce_bgl_desc.entryCount = 5; - ce_bgl_desc.entries = ce_entries; - WGPUBindGroupLayout ce_bgl = - wgpuDeviceCreateBindGroupLayout(device, &ce_bgl_desc); - - WGPUPipelineLayoutDescriptor ce_pl_desc = {}; - ce_pl_desc.bindGroupLayoutCount = 1; - ce_pl_desc.bindGroupLayouts = &ce_bgl; - WGPUPipelineLayout ce_pl = - wgpuDeviceCreatePipelineLayout(device, &ce_pl_desc); - WGPUConstantEntry ce_wg_const = {}; ce_wg_const.key = {"wg_size", WGPU_STRLEN}; ce_wg_const.value = static_cast(ce_wg); - WGPUComputePipelineDescriptor ce_pipe_desc = {}; - ce_pipe_desc.layout = ce_pl; - ce_pipe_desc.compute.module = ce_shader; - ce_pipe_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - ce_pipe_desc.compute.constantCount = 1; - ce_pipe_desc.compute.constants = &ce_wg_const; - WGPUComputePipeline ce_pipe = - wgpuDeviceCreateComputePipeline(device, &ce_pipe_desc); - - WGPUBindGroupEntry ce_bg[5] = {}; - ce_bg[0].binding = 0; - ce_bg[0].buffer = dlogits.buffer; - ce_bg[0].size = dlogits.nbytes; - ce_bg[1].binding = 1; - ce_bg[1].buffer = logits.buffer; - ce_bg[1].size = logits.nbytes; - ce_bg[2].binding = 2; - ce_bg[2].buffer = labels.buffer; - ce_bg[2].size = labels.nbytes; - ce_bg[3].binding = 3; - ce_bg[3].buffer = loss_partial; - ce_bg[3].size = n_rows * sizeof(float); - ce_bg[4].binding = 4; - ce_bg[4].buffer = ce_uniform; - ce_bg[4].size = sizeof(ce_params); - - WGPUBindGroupDescriptor ce_bg_desc = {}; - ce_bg_desc.layout = ce_bgl; - ce_bg_desc.entryCount = 5; - ce_bg_desc.entries = ce_bg; - WGPUBindGroup ce_bind_group = wgpuDeviceCreateBindGroup(device, &ce_bg_desc); + utils::ComputePipelineBundle ce_bundle = utils::make_compute_pipeline( + device, + kFusedCeWGSL, + { + {0, WGPUBufferBindingType_Storage, dlogits.buffer, dlogits.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + logits.buffer, + logits.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + labels.buffer, + labels.nbytes}, + {3, + WGPUBufferBindingType_Storage, + loss_partial, + n_rows * sizeof(float)}, + {4, WGPUBufferBindingType_Uniform, ce_uniform, sizeof(ce_params)}, + }, + &ce_wg_const, + 1); graph.add_dispatch( - {ce_pipe, ce_bind_group, static_cast(n_rows), "fused_ce"}); + {ce_bundle.pipeline, + ce_bundle.bind_group, + static_cast(n_rows), + "fused_ce"}); - wgpuShaderModuleRelease(ce_shader); - wgpuBindGroupLayoutRelease(ce_bgl); - wgpuPipelineLayoutRelease(ce_pl); wgpuBufferRelease(ce_uniform); // reduce loss_partial[N] -> loss[1] (reuses reduce.wgsl) @@ -214,65 +164,27 @@ void fused_ce_impl(WebGPUGraph& graph, const std::vector& args) { utils::compute_1d_workgroup_count(device, 1u, r_wg, "fused_ce_reduce"); WGPUBuffer r_uniform = create_uniform(graph, device, &r_params, sizeof(r_params)); - WGPUShaderModule r_shader = make_shader(device, kReduceWGSL); - - WGPUBindGroupLayoutEntry r_entries[3] = {}; - r_entries[0].binding = 0; - r_entries[0].visibility = WGPUShaderStage_Compute; - r_entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - r_entries[1].binding = 1; - r_entries[1].visibility = WGPUShaderStage_Compute; - r_entries[1].buffer.type = WGPUBufferBindingType_Storage; - r_entries[2].binding = 2; - r_entries[2].visibility = WGPUShaderStage_Compute; - r_entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor r_bgl_desc = {}; - r_bgl_desc.entryCount = 3; - r_bgl_desc.entries = r_entries; - WGPUBindGroupLayout r_bgl = - wgpuDeviceCreateBindGroupLayout(device, &r_bgl_desc); - - WGPUPipelineLayoutDescriptor r_pl_desc = {}; - r_pl_desc.bindGroupLayoutCount = 1; - r_pl_desc.bindGroupLayouts = &r_bgl; - WGPUPipelineLayout r_pl = wgpuDeviceCreatePipelineLayout(device, &r_pl_desc); - WGPUConstantEntry r_wg_const = {}; r_wg_const.key = {"wg_size", WGPU_STRLEN}; r_wg_const.value = static_cast(r_wg); - WGPUComputePipelineDescriptor r_pipe_desc = {}; - r_pipe_desc.layout = r_pl; - r_pipe_desc.compute.module = r_shader; - r_pipe_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - r_pipe_desc.compute.constantCount = 1; - r_pipe_desc.compute.constants = &r_wg_const; - WGPUComputePipeline r_pipe = - wgpuDeviceCreateComputePipeline(device, &r_pipe_desc); + utils::ComputePipelineBundle r_bundle = utils::make_compute_pipeline( + device, + kReduceWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + loss_partial, + n_rows * sizeof(float)}, + {1, WGPUBufferBindingType_Storage, loss.buffer, loss.nbytes}, + {2, WGPUBufferBindingType_Uniform, r_uniform, sizeof(r_params)}, + }, + &r_wg_const, + 1); - WGPUBindGroupEntry r_bg[3] = {}; - r_bg[0].binding = 0; - r_bg[0].buffer = loss_partial; - r_bg[0].size = n_rows * sizeof(float); - r_bg[1].binding = 1; - r_bg[1].buffer = loss.buffer; - r_bg[1].size = loss.nbytes; - r_bg[2].binding = 2; - r_bg[2].buffer = r_uniform; - r_bg[2].size = sizeof(r_params); - - WGPUBindGroupDescriptor r_bg_desc = {}; - r_bg_desc.layout = r_bgl; - r_bg_desc.entryCount = 3; - r_bg_desc.entries = r_bg; - WGPUBindGroup r_bind_group = wgpuDeviceCreateBindGroup(device, &r_bg_desc); - - graph.add_dispatch({r_pipe, r_bind_group, r_wgc, "fused_ce_reduce"}); + graph.add_dispatch( + {r_bundle.pipeline, r_bundle.bind_group, r_wgc, "fused_ce_reduce"}); - wgpuShaderModuleRelease(r_shader); - wgpuBindGroupLayoutRelease(r_bgl); - wgpuPipelineLayoutRelease(r_pl); wgpuBufferRelease(r_uniform); } diff --git a/backends/webgpu/runtime/ops/grid_priors/GridPriors.cpp b/backends/webgpu/runtime/ops/grid_priors/GridPriors.cpp index 03ab5927001..98457835d0d 100644 --- a/backends/webgpu/runtime/ops/grid_priors/GridPriors.cpp +++ b/backends/webgpu/runtime/ops/grid_priors/GridPriors.cpp @@ -85,59 +85,26 @@ void grid_priors_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(GridPriorsParams)); graph.add_uniform_buffer_bytes(sizeof(GridPriorsParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kGridPriorsWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // Bind group: output (rw storage) + params (uniform); input data is unread. - WGPUBindGroupLayoutEntry entries[2] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 2; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg[2] = {}; - bg[0].binding = 0; - bg[0].buffer = out_tensor.buffer; - bg[0].size = out_tensor.nbytes; - bg[1].binding = 1; - bg[1].buffer = uniform_buffer; - bg[1].size = sizeof(GridPriorsParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 2; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kGridPriorsWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(GridPriorsParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "grid_priors", workgroup_count.y}); @@ -174,9 +141,6 @@ void grid_priors_impl(WebGPUGraph& graph, const std::vector& args) { out_id, {static_cast(h * w), static_cast(2)}); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/grid_sampler_2d/GridSampler2d.cpp b/backends/webgpu/runtime/ops/grid_sampler_2d/GridSampler2d.cpp index 27f88f7203d..c5dcd84a45e 100644 --- a/backends/webgpu/runtime/ops/grid_sampler_2d/GridSampler2d.cpp +++ b/backends/webgpu/runtime/ops/grid_sampler_2d/GridSampler2d.cpp @@ -134,70 +134,33 @@ void grid_sampler_2d_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(GridSamplerParams)); graph.add_uniform_buffer_bytes(sizeof(GridSamplerParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kGridSampler2dWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[4] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[4] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = grid_tensor.buffer; - bg_entries[2].size = grid_tensor.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = params_buf; - bg_entries[3].size = sizeof(GridSamplerParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kGridSampler2dWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + grid_tensor.buffer, + grid_tensor.nbytes}, + {3, + WGPUBufferBindingType_Uniform, + params_buf, + sizeof(GridSamplerParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "grid_sampler_2d", workgroup_count.y}); @@ -246,9 +209,6 @@ void grid_sampler_2d_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_tensor_resize_hook(in_id, gs_resize); graph.add_tensor_resize_hook(grid_id, gs_resize); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/index_select/IndexSelect.cpp b/backends/webgpu/runtime/ops/index_select/IndexSelect.cpp index 4f9f9c540f1..542d5902696 100644 --- a/backends/webgpu/runtime/ops/index_select/IndexSelect.cpp +++ b/backends/webgpu/runtime/ops/index_select/IndexSelect.cpp @@ -112,92 +112,42 @@ void index_select_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_uniform_buffer_bytes( 2 * sizeof(TensorMeta) + sizeof(IndexSelectParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kIndexSelectWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // in, out (rw), index (read i32), out_meta, in_meta, params (3 uniforms). - WGPUBindGroupLayoutEntry entries[6] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - entries[5].binding = 5; - entries[5].visibility = WGPUShaderStage_Compute; - entries[5].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 6; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[6] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = self_tensor.buffer; - bg_entries[0].size = self_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = index_tensor.buffer; - bg_entries[2].size = index_tensor.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = out_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - bg_entries[4].binding = 4; - bg_entries[4].buffer = in_meta_buf; - bg_entries[4].size = sizeof(TensorMeta); - bg_entries[5].binding = 5; - bg_entries[5].buffer = params_buf; - bg_entries[5].size = sizeof(IndexSelectParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 6; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kIndexSelectWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + self_tensor.buffer, + self_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + index_tensor.buffer, + index_tensor.nbytes}, + {3, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, in_meta_buf, sizeof(TensorMeta)}, + {5, + WGPUBufferBindingType_Uniform, + params_buf, + sizeof(IndexSelectParams)}, + }, + &wg_size_constant, + 1); // Static shapes only: index_select registers no resize hook, so the output // extent (out.dims[dim] == index numel) is fixed at build time. graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "index_select", workgroup_count.y}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Drop our refs; the bind group keeps the uniforms alive until release. wgpuBufferRelease(out_meta_buf); wgpuBufferRelease(in_meta_buf); diff --git a/backends/webgpu/runtime/ops/linear/Linear.cpp b/backends/webgpu/runtime/ops/linear/Linear.cpp index 97f0591d101..45cc27a515d 100644 --- a/backends/webgpu/runtime/ops/linear/Linear.cpp +++ b/backends/webgpu/runtime/ops/linear/Linear.cpp @@ -109,78 +109,27 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { // vec4-over-K path when K%4==0 (scalar out, only K%4, not N%4 like mm). const bool use_vec4 = (K % 4u == 0u); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {use_vec4 ? kLinearVec4WGSL : kLinearTiledWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Storage; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - // Tiled kernel has a fixed @workgroup_size(8, 8, 1) — no override constant. - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[5] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in.buffer; - bg_entries[0].size = in.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = w.buffer; - bg_entries[1].size = w.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out.buffer; - bg_entries[2].size = out.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = uniform_buffer; - bg_entries[3].size = sizeof(LinearParams); - bg_entries[4].binding = 4; - bg_entries[4].buffer = bias.buffer; - bg_entries[4].size = bias.nbytes; - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + use_vec4 ? kLinearVec4WGSL : kLinearTiledWGSL, + { + {0, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, w.buffer, w.nbytes}, + {2, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {3, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(LinearParams)}, + {4, WGPUBufferBindingType_ReadOnlyStorage, bias.buffer, bias.nbytes}, + }); if (bias.owned_dummy != nullptr) { wgpuBufferRelease(bias.owned_dummy); } WebGPUDispatch dispatch; - dispatch.pipeline = pipeline; - dispatch.bind_group = bind_group; + dispatch.pipeline = bundle.pipeline; + dispatch.bind_group = bundle.bind_group; dispatch.workgroup_count_x = dispatch_x; dispatch.workgroup_count_y = dispatch_y; dispatch.kernel_name = "linear"; @@ -215,9 +164,6 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { out_id, {static_cast(m), static_cast(N)}); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/LinearDq8caQ4gsw.cpp b/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/LinearDq8caQ4gsw.cpp index 96f456923a0..3b113b85fc2 100644 --- a/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/LinearDq8caQ4gsw.cpp +++ b/backends/webgpu/runtime/ops/linear_dq8ca_q4gsw/LinearDq8caQ4gsw.cpp @@ -41,9 +41,7 @@ constexpr int64_t kTileN = 4; // MUST match TN // bias, out]. Dynamic per-row int8 activation quant x 4-bit-group symmetric // weight. weight_sums (arg 4) is a perf shortcut; this v1 recomputes the sum // inline so it is intentionally unused. Static-shape only (no resize hook yet). -void linear_dq8ca_q4gsw_impl( - WebGPUGraph& graph, - const std::vector& args) { +void linear_dq8ca_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { const int in_id = args.at(0); const int input_scale_id = args.at(1); const int input_zp_id = args.at(2); @@ -121,7 +119,8 @@ void linear_dq8ca_q4gsw_impl( "linear_dq8ca_q4gsw: input scale fp32[M] / zp int8[M] required"); } // int8 zp is bound word-aligned over a max(nbytes,4) buffer; M in {5,6,7,...} - // would bind past the buffer. Mirrors the choose_qparams_affine producer guard. + // would bind past the buffer. Mirrors the choose_qparams_affine producer + // guard. if (M > 4u && M % 4u != 0u) { throw std::runtime_error( "linear_dq8ca_q4gsw: num_rows must be <=4 or a multiple of 4"); @@ -154,15 +153,18 @@ void linear_dq8ca_q4gsw_impl( params.padded_N = padded_N; params.has_bias = has_bias; - const int64_t total_tiles = utils::div_up(M, kTileM) * - utils::div_up(N, kTileN); + const int64_t total_tiles = + utils::div_up(M, kTileM) * utils::div_up(N, kTileN); if (total_tiles > static_cast(UINT32_MAX)) { throw std::runtime_error("linear_dq8ca_q4gsw: tile count exceeds u32"); } const uint32_t wg_size = utils::clamp_workgroup_size(device, kLinearDq8caQ4gswWorkgroupSizeX); const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( - device, static_cast(total_tiles), wg_size, "linear_dq8ca_q4gsw"); + device, + static_cast(total_tiles), + wg_size, + "linear_dq8ca_q4gsw"); WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; @@ -172,101 +174,51 @@ void linear_dq8ca_q4gsw_impl( utils::make_uniform(device, ¶ms, sizeof(Dq8caParams)); graph.add_uniform_buffer_bytes(sizeof(Dq8caParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kLinearDq8caQ4gswWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // 0 out(rw), 1 in, 2 input_scale, 3 input_zp, 4 weight, 5 scales, 6 bias (ro), - // 7 uniform. - WGPUBindGroupLayoutEntry entries[8] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - for (uint32_t i = 1; i <= 6; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - } - entries[7].binding = 7; - entries[7].visibility = WGPUShaderStage_Compute; - entries[7].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 8; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg[8] = {}; - bg[0].binding = 0; - bg[0].buffer = out.buffer; - bg[0].size = out.nbytes; - bg[1].binding = 1; - bg[1].buffer = in.buffer; - bg[1].size = in.nbytes; - bg[2].binding = 2; - bg[2].buffer = input_scale.buffer; - bg[2].size = input_scale.nbytes; - bg[3].binding = 3; - bg[3].buffer = input_zp.buffer; - // int8 zp bound as array; round to a multiple of 4 (buffer is >=4 bytes). - bg[3].size = ((input_zp.nbytes + 3u) / 4u) * 4u; - bg[4].binding = 4; - bg[4].buffer = weight.buffer; - bg[4].size = weight.nbytes; - bg[5].binding = 5; - bg[5].buffer = scales.buffer; - bg[5].size = scales.nbytes; - bg[6].binding = 6; - bg[6].buffer = bias_buffer; - bg[6].size = bias_size; - bg[7].binding = 7; - bg[7].buffer = params_buf; - bg[7].size = sizeof(Dq8caParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 8; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + // 0 out(rw), 1 in, 2 input_scale, 3 input_zp, 4 weight, 5 scales, 6 bias + // (ro), 7 uniform. + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kLinearDq8caQ4gswWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + input_scale.buffer, + input_scale.nbytes}, + // int8 zp bound as array; round to a multiple of 4 (buffer is + // >=4 bytes). + {3, + WGPUBufferBindingType_ReadOnlyStorage, + input_zp.buffer, + ((input_zp.nbytes + 3u) / 4u) * 4u}, + {4, + WGPUBufferBindingType_ReadOnlyStorage, + weight.buffer, + weight.nbytes}, + {5, + WGPUBufferBindingType_ReadOnlyStorage, + scales.buffer, + scales.nbytes}, + {6, WGPUBufferBindingType_ReadOnlyStorage, bias_buffer, bias_size}, + {7, WGPUBufferBindingType_Uniform, params_buf, sizeof(Dq8caParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "linear_dq8ca_q4gsw", workgroup_count.y}); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } } // namespace WEBGPU_REGISTER_OPERATORS { - WEBGPU_REGISTER_OP( - et_vk.linear_dq8ca_q4gsw.default, - linear_dq8ca_q4gsw_impl); + WEBGPU_REGISTER_OP(et_vk.linear_dq8ca_q4gsw.default, linear_dq8ca_q4gsw_impl); } } // namespace executorch::backends::webgpu - diff --git a/backends/webgpu/runtime/ops/linear_q8ta_q8csw/LinearQ8taQ8csw.cpp b/backends/webgpu/runtime/ops/linear_q8ta_q8csw/LinearQ8taQ8csw.cpp index eedffc96572..60a2c19f895 100644 --- a/backends/webgpu/runtime/ops/linear_q8ta_q8csw/LinearQ8taQ8csw.cpp +++ b/backends/webgpu/runtime/ops/linear_q8ta_q8csw/LinearQ8taQ8csw.cpp @@ -136,83 +136,47 @@ void linear_q8ta_q8csw_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(Q8taQ8cswParams)); graph.add_uniform_buffer_bytes(sizeof(Q8taQ8cswParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kLinearQ8taQ8cswWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[6] = {}; - for (int i = 0; i < 6; i++) { - entries[i].binding = static_cast(i); - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = (i == 0) ? WGPUBufferBindingType_Storage - : (i == 5) ? WGPUBufferBindingType_Uniform - : WGPUBufferBindingType_ReadOnlyStorage; - } - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 6; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - // No-bias: bind scales as an unread placeholder (has_bias gates the read). WGPUBuffer bias_buf = has_bias ? graph.get_tensor(bias_id).buffer : scales_tensor.buffer; const uint64_t bias_size = has_bias ? graph.get_tensor(bias_id).nbytes : scales_tensor.nbytes; - WGPUBindGroupEntry bg[6] = {}; - bg[0].binding = 0; - bg[0].buffer = out_tensor.buffer; - bg[0].size = out_tensor.nbytes; - bg[1].binding = 1; - bg[1].buffer = in_tensor.buffer; - bg[1].size = in_tensor.nbytes; - bg[2].binding = 2; - bg[2].buffer = weight_tensor.buffer; - bg[2].size = weight_tensor.nbytes; - bg[3].binding = 3; - bg[3].buffer = scales_tensor.buffer; - bg[3].size = scales_tensor.nbytes; - bg[4].binding = 4; - bg[4].buffer = bias_buf; - bg[4].size = bias_size; - bg[5].binding = 5; - bg[5].buffer = params_buf; - bg[5].size = sizeof(Q8taQ8cswParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 6; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kLinearQ8taQ8cswWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight_tensor.buffer, + weight_tensor.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + scales_tensor.buffer, + scales_tensor.nbytes}, + {4, WGPUBufferBindingType_ReadOnlyStorage, bias_buf, bias_size}, + {5, + WGPUBufferBindingType_Uniform, + params_buf, + sizeof(Q8taQ8cswParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "linear_q8ta_q8csw", workgroup_count.y}); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/linear_qcs4w/QuantizedLinearQcs4w.cpp b/backends/webgpu/runtime/ops/linear_qcs4w/QuantizedLinearQcs4w.cpp index d75c2617072..858a4a2e8cb 100644 --- a/backends/webgpu/runtime/ops/linear_qcs4w/QuantizedLinearQcs4w.cpp +++ b/backends/webgpu/runtime/ops/linear_qcs4w/QuantizedLinearQcs4w.cpp @@ -127,73 +127,31 @@ void qcs4w_linear_impl(WebGPUGraph& graph, const std::vector& args) { wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kQcs4wLinearWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group layout: out (rw) + in/weight/scales (ro storage) + uniform. - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - for (uint32_t i = 1; i <= 3; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - } - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[5] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = out.buffer; - bg_entries[0].size = out.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = in.buffer; - bg_entries[1].size = in.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = weight.buffer; - bg_entries[2].size = weight.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = scales.buffer; - bg_entries[3].size = scales.nbytes; - bg_entries[4].binding = 4; - bg_entries[4].buffer = uniform_buffer; - bg_entries[4].size = sizeof(Qcs4wParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kQcs4wLinearWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight.buffer, + weight.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + scales.buffer, + scales.nbytes}, + {4, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(Qcs4wParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "linear_qcs4w", workgroup_count.y}); @@ -238,9 +196,6 @@ void qcs4w_linear_impl(WebGPUGraph& graph, const std::vector& args) { g.set_cur_dims(out_id, od); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp index f18305de01b..b6a7e12830f 100644 --- a/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp +++ b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp @@ -83,71 +83,34 @@ void logical_and_op(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(LogicalAndParams)); graph.add_uniform_buffer_bytes(sizeof(LogicalAndParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kLogicalAndWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // out (rw storage) + a/b (ro storage) + params (uniform). - WGPUBindGroupLayoutEntry entries[4] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg[4] = {}; - bg[0].binding = 0; - bg[0].buffer = out_tensor.buffer; - bg[0].size = out_tensor.nbytes; - bg[1].binding = 1; - bg[1].buffer = a_tensor.buffer; - bg[1].size = a_tensor.nbytes; - bg[2].binding = 2; - bg[2].buffer = b_tensor.buffer; - bg[2].size = b_tensor.nbytes; - bg[3].binding = 3; - bg[3].buffer = uniform_buffer; - bg[3].size = sizeof(LogicalAndParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kLogicalAndWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + a_tensor.buffer, + a_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + b_tensor.buffer, + b_tensor.nbytes}, + {3, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(LogicalAndParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "logical_and", workgroup_count.y}); @@ -175,9 +138,6 @@ void logical_and_op(WebGPUGraph& graph, const std::vector& args) { graph.add_tensor_resize_hook(a_id, resize); graph.add_tensor_resize_hook(b_id, resize); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp b/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp index 5fe67cc49d1..d26f6b486d2 100644 --- a/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp +++ b/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp @@ -26,7 +26,9 @@ struct LogicalOrParams { uint32_t num_words; uint32_t _pad[3]; }; -static_assert(sizeof(LogicalOrParams) == 16, "LogicalOrParams must be 16 bytes"); +static_assert( + sizeof(LogicalOrParams) == 16, + "LogicalOrParams must be 16 bytes"); // out = a | b (bool OR); serves logical_or + bitwise_or (mirrors Vulkan). void logical_or_op(WebGPUGraph& graph, const std::vector& args) { @@ -79,101 +81,60 @@ void logical_or_op(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(LogicalOrParams)); graph.add_uniform_buffer_bytes(sizeof(LogicalOrParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kLogicalOrWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // out (rw storage) + a/b (ro storage) + params (uniform). - WGPUBindGroupLayoutEntry entries[4] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg[4] = {}; - bg[0].binding = 0; - bg[0].buffer = out_tensor.buffer; - bg[0].size = out_tensor.nbytes; - bg[1].binding = 1; - bg[1].buffer = a_tensor.buffer; - bg[1].size = a_tensor.nbytes; - bg[2].binding = 2; - bg[2].buffer = b_tensor.buffer; - bg[2].size = b_tensor.nbytes; - bg[3].binding = 3; - bg[3].buffer = uniform_buffer; - bg[3].size = sizeof(LogicalOrParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kLogicalOrWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + a_tensor.buffer, + a_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + b_tensor.buffer, + b_tensor.nbytes}, + {3, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(LogicalOrParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "logical_or", workgroup_count.y}); // Dynamic shapes: recompute num_words/dispatch; out follows a (same-shape). WGPUBuffer params_buf = uniform_buffer; - auto resize = - [a_id, b_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { - const auto& d = g.cur_dims(a_id); - const uint64_t n = utils::numel_of(d); - if (n == 0u || n % 4u != 0u || n > UINT32_MAX || - utils::numel_of(g.cur_dims(b_id)) != n) { - throw std::runtime_error( - "logical_or(resize): numel must be a mult of 4"); - } - g.set_cur_dims(out_id, d); - LogicalOrParams p = {}; - p.num_words = static_cast(n / 4u); - wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); - const utils::WgCount wgc = utils::compute_2d_workgroup_count( - g.device(), p.num_words, wg_size, "logical_or"); - g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; - g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; - }; + auto resize = [a_id, b_id, out_id, wg_size, dispatch_idx, params_buf]( + WebGPUGraph& g) { + const auto& d = g.cur_dims(a_id); + const uint64_t n = utils::numel_of(d); + if (n == 0u || n % 4u != 0u || n > UINT32_MAX || + utils::numel_of(g.cur_dims(b_id)) != n) { + throw std::runtime_error("logical_or(resize): numel must be a mult of 4"); + } + g.set_cur_dims(out_id, d); + LogicalOrParams p = {}; + p.num_words = static_cast(n / 4u); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), p.num_words, wg_size, "logical_or"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }; graph.add_tensor_resize_hook(a_id, resize); graph.add_tensor_resize_hook(b_id, resize); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/mm/Mm.cpp b/backends/webgpu/runtime/ops/mm/Mm.cpp index b93c2aff247..a37a03d5da9 100644 --- a/backends/webgpu/runtime/ops/mm/Mm.cpp +++ b/backends/webgpu/runtime/ops/mm/Mm.cpp @@ -89,69 +89,20 @@ void mm_impl(WebGPUGraph& graph, const std::vector& args) { // vec4 path when K and N are multiples of 4 (wider 16B loads); else scalar. const bool use_vec4 = (K % 4u == 0u) && (N % 4u == 0u); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {use_vec4 ? kMmVec4WGSL : kMmTiledWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[4] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Storage; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - // Tiled kernel has a fixed @workgroup_size(8, 8, 1) — no override constant. - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[4] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = a.buffer; - bg_entries[0].size = a.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = b.buffer; - bg_entries[1].size = b.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out.buffer; - bg_entries[2].size = out.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = uniform_buffer; - bg_entries[3].size = sizeof(MmParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + use_vec4 ? kMmVec4WGSL : kMmTiledWGSL, + { + {0, WGPUBufferBindingType_ReadOnlyStorage, a.buffer, a.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, b.buffer, b.nbytes}, + {2, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {3, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(MmParams)}, + }); WebGPUDispatch dispatch; - dispatch.pipeline = pipeline; - dispatch.bind_group = bind_group; + dispatch.pipeline = bundle.pipeline; + dispatch.bind_group = bundle.bind_group; dispatch.workgroup_count_x = dispatch_x; dispatch.workgroup_count_y = dispatch_y; dispatch.kernel_name = "mm"; @@ -187,9 +138,6 @@ void mm_impl(WebGPUGraph& graph, const std::vector& args) { out_id, {static_cast(m), static_cast(N)}); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp index ed81fb3216c..fdb7984b9e1 100644 --- a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp @@ -78,95 +78,35 @@ void mul_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, &in2_meta, sizeof(TensorMeta)); graph.add_uniform_buffer_bytes(3 * sizeof(TensorMeta)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kBinaryMulWGSL, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group: in1, in2, out (rw), out_meta, in1_meta, in2_meta (3 uniforms). - WGPUBindGroupLayoutEntry entries[6] = {}; - - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Storage; - - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - entries[5].binding = 5; - entries[5].visibility = WGPUShaderStage_Compute; - entries[5].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 6; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[6] = {}; - - bg_entries[0].binding = 0; - bg_entries[0].buffer = in1_tensor.buffer; - bg_entries[0].size = in1_tensor.nbytes; - - bg_entries[1].binding = 1; - bg_entries[1].buffer = in2_tensor.buffer; - bg_entries[1].size = in2_tensor.nbytes; - - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_tensor.buffer; - bg_entries[2].size = out_tensor.nbytes; - - bg_entries[3].binding = 3; - bg_entries[3].buffer = out_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - - bg_entries[4].binding = 4; - bg_entries[4].buffer = in1_meta_buf; - bg_entries[4].size = sizeof(TensorMeta); - - bg_entries[5].binding = 5; - bg_entries[5].buffer = in2_meta_buf; - bg_entries[5].size = sizeof(TensorMeta); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 6; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kBinaryMulWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in1_tensor.buffer, + in1_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in2_tensor.buffer, + in2_tensor.nbytes}, + {2, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {3, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, in1_meta_buf, sizeof(TensorMeta)}, + {5, WGPUBufferBindingType_Uniform, in2_meta_buf, sizeof(TensorMeta)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "mul", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "mul", + workgroup_count.y}); // Dynamic shapes: rebuild all 3 broadcast TensorMeta UBOs + dispatch. WGPUBuffer o_buf = out_meta_buf, a_buf = in1_meta_buf, b_buf = in2_meta_buf; @@ -207,9 +147,6 @@ void mul_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_tensor_resize_hook(in1_id, mul_resize); graph.add_tensor_resize_hook(in2_id, mul_resize); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns them so a resize hook can rewrite them; freed in the dtor. graph.own_uniform_buffer(out_meta_buf); graph.own_uniform_buffer(in1_meta_buf); diff --git a/backends/webgpu/runtime/ops/native_group_norm/GroupNorm.cpp b/backends/webgpu/runtime/ops/native_group_norm/GroupNorm.cpp index 19be6a89a11..5a0d06514e3 100644 --- a/backends/webgpu/runtime/ops/native_group_norm/GroupNorm.cpp +++ b/backends/webgpu/runtime/ops/native_group_norm/GroupNorm.cpp @@ -37,8 +37,8 @@ static_assert( sizeof(GroupNormParams) == 32, "GroupNormParams must match the WGSL Params struct (32 bytes)"); -// The reduce shader's cooperative reduction uses var f32[256] arrays; -// the workgroup size (used for both passes) must not exceed that width. +// The reduce shader's cooperative reduction uses var f32[256] +// arrays; the workgroup size (used for both passes) must not exceed that width. static_assert( kGroupNormWorkgroupSizeX <= 256, "group_norm workgroup size exceeds the 256-wide shared reduction arrays"); @@ -62,57 +62,20 @@ size_t add_gn_dispatch( wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {wgsl_code, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - std::vector entries(binds.size()); - for (size_t i = 0; i < binds.size(); i++) { - entries[i].binding = static_cast(i); - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = binds[i].type; - } - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = entries.size(); - bgl_desc.entries = entries.data(); - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - std::vector bg_entries(binds.size()); + std::vector specs(binds.size()); for (size_t i = 0; i < binds.size(); i++) { - bg_entries[i].binding = static_cast(i); - bg_entries[i].buffer = binds[i].buffer; - bg_entries[i].size = binds[i].size; + specs[i] = { + static_cast(i), + binds[i].type, + binds[i].buffer, + binds[i].size}; } - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = bg_entries.size(); - bg_desc.entries = bg_entries.data(); - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - const size_t idx = - graph.add_dispatch({pipeline, bind_group, wgc.x, label, wgc.y}); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, wgsl_code, specs, &wg_size_constant, 1); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); + const size_t idx = graph.add_dispatch( + {bundle.pipeline, bundle.bind_group, wgc.x, label, wgc.y}); return idx; } @@ -212,8 +175,8 @@ void native_group_norm_impl(WebGPUGraph& graph, const std::vector& args) { const WGPUBufferBindingType uni = WGPUBufferBindingType_Uniform; // Pass 1: reduce -> mean/rstd (one workgroup per (n, group), cooperative). - utils::WgCount reduce_wgc = utils::compute_2d_workgroup_count( - device, mean_numel, 1, "gn_reduce"); + utils::WgCount reduce_wgc = + utils::compute_2d_workgroup_count(device, mean_numel, 1, "gn_reduce"); const size_t reduce_idx = add_gn_dispatch( graph, device, diff --git a/backends/webgpu/runtime/ops/pixel_shuffle/PixelShuffle.cpp b/backends/webgpu/runtime/ops/pixel_shuffle/PixelShuffle.cpp index c8f7e399c99..500a0fbbc40 100644 --- a/backends/webgpu/runtime/ops/pixel_shuffle/PixelShuffle.cpp +++ b/backends/webgpu/runtime/ops/pixel_shuffle/PixelShuffle.cpp @@ -104,64 +104,29 @@ void pixel_shuffle_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(PixelShuffleParams)); graph.add_uniform_buffer_bytes(sizeof(PixelShuffleParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kPixelShuffleWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[3] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[3] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = params_buf; - bg_entries[2].size = sizeof(PixelShuffleParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kPixelShuffleWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + params_buf, + sizeof(PixelShuffleParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "pixel_shuffle", workgroup_count.y}); @@ -202,9 +167,6 @@ void pixel_shuffle_impl(WebGPUGraph& graph, const std::vector& args) { g.set_cur_dims(out_id, out_d); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/q8ta_add/Q8taAdd.cpp b/backends/webgpu/runtime/ops/q8ta_add/Q8taAdd.cpp index b2123cc0e46..8a207c201fd 100644 --- a/backends/webgpu/runtime/ops/q8ta_add/Q8taAdd.cpp +++ b/backends/webgpu/runtime/ops/q8ta_add/Q8taAdd.cpp @@ -107,73 +107,34 @@ void q8ta_add_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(Q8taAddParams)); graph.add_uniform_buffer_bytes(sizeof(Q8taAddParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kQ8taAddWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[4] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Storage; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[4] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = a_tensor.buffer; - bg_entries[0].size = a_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = b_tensor.buffer; - bg_entries[1].size = b_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_tensor.buffer; - bg_entries[2].size = out_tensor.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = params_buf; - bg_entries[3].size = sizeof(Q8taAddParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kQ8taAddWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + a_tensor.buffer, + a_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + b_tensor.buffer, + b_tensor.nbytes}, + {2, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {3, WGPUBufferBindingType_Uniform, params_buf, sizeof(Q8taAddParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "q8ta_add", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "q8ta_add", + workgroup_count.y}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d/Q8taConv2d.cpp b/backends/webgpu/runtime/ops/q8ta_conv2d/Q8taConv2d.cpp index a6d88e159ef..d5f8eae26a0 100644 --- a/backends/webgpu/runtime/ops/q8ta_conv2d/Q8taConv2d.cpp +++ b/backends/webgpu/runtime/ops/q8ta_conv2d/Q8taConv2d.cpp @@ -201,83 +201,47 @@ void q8ta_conv2d_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(Q8taConvParams)); graph.add_uniform_buffer_bytes(sizeof(Q8taConvParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kQ8taConv2dWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[6] = {}; - for (int i = 0; i < 6; i++) { - entries[i].binding = static_cast(i); - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = (i == 0) ? WGPUBufferBindingType_Storage - : (i == 5) ? WGPUBufferBindingType_Uniform - : WGPUBufferBindingType_ReadOnlyStorage; - } - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 6; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - // No-bias: bind scales as an unread placeholder (has_bias gates the read). WGPUBuffer bias_buf = has_bias ? graph.get_tensor(bias_id).buffer : scales_tensor.buffer; const uint64_t bias_size = has_bias ? graph.get_tensor(bias_id).nbytes : scales_tensor.nbytes; - WGPUBindGroupEntry bg[6] = {}; - bg[0].binding = 0; - bg[0].buffer = out_tensor.buffer; - bg[0].size = out_tensor.nbytes; - bg[1].binding = 1; - bg[1].buffer = in_tensor.buffer; - bg[1].size = in_tensor.nbytes; - bg[2].binding = 2; - bg[2].buffer = weight_tensor.buffer; - bg[2].size = weight_tensor.nbytes; - bg[3].binding = 3; - bg[3].buffer = scales_tensor.buffer; - bg[3].size = scales_tensor.nbytes; - bg[4].binding = 4; - bg[4].buffer = bias_buf; - bg[4].size = bias_size; - bg[5].binding = 5; - bg[5].buffer = params_buf; - bg[5].size = sizeof(Q8taConvParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 6; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kQ8taConv2dWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight_tensor.buffer, + weight_tensor.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + scales_tensor.buffer, + scales_tensor.nbytes}, + {4, WGPUBufferBindingType_ReadOnlyStorage, bias_buf, bias_size}, + {5, + WGPUBufferBindingType_Uniform, + params_buf, + sizeof(Q8taConvParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "q8ta_conv2d", workgroup_count.y}); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d_dw/Q8taConv2dDw.cpp b/backends/webgpu/runtime/ops/q8ta_conv2d_dw/Q8taConv2dDw.cpp index e0b0a5e548c..f49ea88ca12 100644 --- a/backends/webgpu/runtime/ops/q8ta_conv2d_dw/Q8taConv2dDw.cpp +++ b/backends/webgpu/runtime/ops/q8ta_conv2d_dw/Q8taConv2dDw.cpp @@ -189,83 +189,47 @@ void q8ta_conv2d_dw_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(Q8taConvDwParams)); graph.add_uniform_buffer_bytes(sizeof(Q8taConvDwParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kQ8taConv2dDwWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[6] = {}; - for (int i = 0; i < 6; i++) { - entries[i].binding = static_cast(i); - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = (i == 0) ? WGPUBufferBindingType_Storage - : (i == 5) ? WGPUBufferBindingType_Uniform - : WGPUBufferBindingType_ReadOnlyStorage; - } - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 6; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - // No-bias: bind scales as an unread placeholder (has_bias gates the read). WGPUBuffer bias_buf = has_bias ? graph.get_tensor(bias_id).buffer : scales_tensor.buffer; const uint64_t bias_size = has_bias ? graph.get_tensor(bias_id).nbytes : scales_tensor.nbytes; - WGPUBindGroupEntry bg[6] = {}; - bg[0].binding = 0; - bg[0].buffer = out_tensor.buffer; - bg[0].size = out_tensor.nbytes; - bg[1].binding = 1; - bg[1].buffer = in_tensor.buffer; - bg[1].size = in_tensor.nbytes; - bg[2].binding = 2; - bg[2].buffer = weight_tensor.buffer; - bg[2].size = weight_tensor.nbytes; - bg[3].binding = 3; - bg[3].buffer = scales_tensor.buffer; - bg[3].size = scales_tensor.nbytes; - bg[4].binding = 4; - bg[4].buffer = bias_buf; - bg[4].size = bias_size; - bg[5].binding = 5; - bg[5].buffer = params_buf; - bg[5].size = sizeof(Q8taConvDwParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 6; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kQ8taConv2dDwWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight_tensor.buffer, + weight_tensor.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + scales_tensor.buffer, + scales_tensor.nbytes}, + {4, WGPUBufferBindingType_ReadOnlyStorage, bias_buf, bias_size}, + {5, + WGPUBufferBindingType_Uniform, + params_buf, + sizeof(Q8taConvDwParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "q8ta_conv2d_dw", workgroup_count.y}); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d_pw/Q8taConv2dPw.cpp b/backends/webgpu/runtime/ops/q8ta_conv2d_pw/Q8taConv2dPw.cpp index 5dbf91558ab..315ec860269 100644 --- a/backends/webgpu/runtime/ops/q8ta_conv2d_pw/Q8taConv2dPw.cpp +++ b/backends/webgpu/runtime/ops/q8ta_conv2d_pw/Q8taConv2dPw.cpp @@ -155,83 +155,47 @@ void q8ta_conv2d_pw_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(Q8taConvPwParams)); graph.add_uniform_buffer_bytes(sizeof(Q8taConvPwParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kQ8taConv2dPwWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[6] = {}; - for (int i = 0; i < 6; i++) { - entries[i].binding = static_cast(i); - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = (i == 0) ? WGPUBufferBindingType_Storage - : (i == 5) ? WGPUBufferBindingType_Uniform - : WGPUBufferBindingType_ReadOnlyStorage; - } - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 6; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - // No-bias: bind scales as an unread placeholder (has_bias gates the read). WGPUBuffer bias_buf = has_bias ? graph.get_tensor(bias_id).buffer : scales_tensor.buffer; const uint64_t bias_size = has_bias ? graph.get_tensor(bias_id).nbytes : scales_tensor.nbytes; - WGPUBindGroupEntry bg[6] = {}; - bg[0].binding = 0; - bg[0].buffer = out_tensor.buffer; - bg[0].size = out_tensor.nbytes; - bg[1].binding = 1; - bg[1].buffer = in_tensor.buffer; - bg[1].size = in_tensor.nbytes; - bg[2].binding = 2; - bg[2].buffer = weight_tensor.buffer; - bg[2].size = weight_tensor.nbytes; - bg[3].binding = 3; - bg[3].buffer = scales_tensor.buffer; - bg[3].size = scales_tensor.nbytes; - bg[4].binding = 4; - bg[4].buffer = bias_buf; - bg[4].size = bias_size; - bg[5].binding = 5; - bg[5].buffer = params_buf; - bg[5].size = sizeof(Q8taConvPwParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 6; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kQ8taConv2dPwWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight_tensor.buffer, + weight_tensor.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + scales_tensor.buffer, + scales_tensor.nbytes}, + {4, WGPUBufferBindingType_ReadOnlyStorage, bias_buf, bias_size}, + {5, + WGPUBufferBindingType_Uniform, + params_buf, + sizeof(Q8taConvPwParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "q8ta_conv2d_pw", workgroup_count.y}); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/q8ta_conv2d_transposed/Q8taConv2dTransposed.cpp b/backends/webgpu/runtime/ops/q8ta_conv2d_transposed/Q8taConv2dTransposed.cpp index 0990b49bb9d..1987d699edd 100644 --- a/backends/webgpu/runtime/ops/q8ta_conv2d_transposed/Q8taConv2dTransposed.cpp +++ b/backends/webgpu/runtime/ops/q8ta_conv2d_transposed/Q8taConv2dTransposed.cpp @@ -219,83 +219,47 @@ void q8ta_conv2d_transposed_impl( utils::make_uniform(device, ¶ms, sizeof(Q8taConvTParams)); graph.add_uniform_buffer_bytes(sizeof(Q8taConvTParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kQ8taConv2dTransposedWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[6] = {}; - for (int i = 0; i < 6; i++) { - entries[i].binding = static_cast(i); - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = (i == 0) ? WGPUBufferBindingType_Storage - : (i == 5) ? WGPUBufferBindingType_Uniform - : WGPUBufferBindingType_ReadOnlyStorage; - } - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 6; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - // No-bias: bind scales as an unread placeholder (has_bias gates the read). WGPUBuffer bias_buf = has_bias ? graph.get_tensor(bias_id).buffer : scales_tensor.buffer; const uint64_t bias_size = has_bias ? graph.get_tensor(bias_id).nbytes : scales_tensor.nbytes; - WGPUBindGroupEntry bg[6] = {}; - bg[0].binding = 0; - bg[0].buffer = out_tensor.buffer; - bg[0].size = out_tensor.nbytes; - bg[1].binding = 1; - bg[1].buffer = in_tensor.buffer; - bg[1].size = in_tensor.nbytes; - bg[2].binding = 2; - bg[2].buffer = weight_tensor.buffer; - bg[2].size = weight_tensor.nbytes; - bg[3].binding = 3; - bg[3].buffer = scales_tensor.buffer; - bg[3].size = scales_tensor.nbytes; - bg[4].binding = 4; - bg[4].buffer = bias_buf; - bg[4].size = bias_size; - bg[5].binding = 5; - bg[5].buffer = params_buf; - bg[5].size = sizeof(Q8taConvTParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 6; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kQ8taConv2dTransposedWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight_tensor.buffer, + weight_tensor.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + scales_tensor.buffer, + scales_tensor.nbytes}, + {4, WGPUBufferBindingType_ReadOnlyStorage, bias_buf, bias_size}, + {5, + WGPUBufferBindingType_Uniform, + params_buf, + sizeof(Q8taConvTParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "q8ta_conv2d_transposed", workgroup_count.y}); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/q8ta_linear/Q8taLinear.cpp b/backends/webgpu/runtime/ops/q8ta_linear/Q8taLinear.cpp index fdf2eb38542..90924d2cecb 100644 --- a/backends/webgpu/runtime/ops/q8ta_linear/Q8taLinear.cpp +++ b/backends/webgpu/runtime/ops/q8ta_linear/Q8taLinear.cpp @@ -144,83 +144,47 @@ void q8ta_linear_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(Q8taLinearParams)); graph.add_uniform_buffer_bytes(sizeof(Q8taLinearParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kQ8taLinearWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[6] = {}; - for (int i = 0; i < 6; i++) { - entries[i].binding = static_cast(i); - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = (i == 0) ? WGPUBufferBindingType_Storage - : (i == 5) ? WGPUBufferBindingType_Uniform - : WGPUBufferBindingType_ReadOnlyStorage; - } - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 6; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - // No-bias: bind scales as an unread placeholder (has_bias gates the read). WGPUBuffer bias_buf = has_bias ? graph.get_tensor(bias_id).buffer : scales_tensor.buffer; const uint64_t bias_size = has_bias ? graph.get_tensor(bias_id).nbytes : scales_tensor.nbytes; - WGPUBindGroupEntry bg[6] = {}; - bg[0].binding = 0; - bg[0].buffer = out_tensor.buffer; - bg[0].size = out_tensor.nbytes; - bg[1].binding = 1; - bg[1].buffer = in_tensor.buffer; - bg[1].size = in_tensor.nbytes; - bg[2].binding = 2; - bg[2].buffer = weight_tensor.buffer; - bg[2].size = weight_tensor.nbytes; - bg[3].binding = 3; - bg[3].buffer = scales_tensor.buffer; - bg[3].size = scales_tensor.nbytes; - bg[4].binding = 4; - bg[4].buffer = bias_buf; - bg[4].size = bias_size; - bg[5].binding = 5; - bg[5].buffer = params_buf; - bg[5].size = sizeof(Q8taLinearParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 6; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kQ8taLinearWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight_tensor.buffer, + weight_tensor.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + scales_tensor.buffer, + scales_tensor.nbytes}, + {4, WGPUBufferBindingType_ReadOnlyStorage, bias_buf, bias_size}, + {5, + WGPUBufferBindingType_Uniform, + params_buf, + sizeof(Q8taLinearParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "q8ta_linear", workgroup_count.y}); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/q8ta_pixel_shuffle/Q8taPixelShuffle.cpp b/backends/webgpu/runtime/ops/q8ta_pixel_shuffle/Q8taPixelShuffle.cpp index 7fa56895856..5ee585d9c6a 100644 --- a/backends/webgpu/runtime/ops/q8ta_pixel_shuffle/Q8taPixelShuffle.cpp +++ b/backends/webgpu/runtime/ops/q8ta_pixel_shuffle/Q8taPixelShuffle.cpp @@ -136,64 +136,29 @@ void q8ta_pixel_shuffle_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(Q8taPixelShuffleParams)); graph.add_uniform_buffer_bytes(sizeof(Q8taPixelShuffleParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kQ8taPixelShuffleWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[3] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[3] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = params_buf; - bg_entries[2].size = sizeof(Q8taPixelShuffleParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kQ8taPixelShuffleWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + params_buf, + sizeof(Q8taPixelShuffleParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "q8ta_pixel_shuffle", workgroup_count.y}); @@ -241,9 +206,6 @@ void q8ta_pixel_shuffle_impl(WebGPUGraph& graph, const std::vector& args) { g.set_cur_dims(out_id, out_d); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/q8ta_relu/Q8taRelu.cpp b/backends/webgpu/runtime/ops/q8ta_relu/Q8taRelu.cpp index a6cafe467a6..5456d5446e5 100644 --- a/backends/webgpu/runtime/ops/q8ta_relu/Q8taRelu.cpp +++ b/backends/webgpu/runtime/ops/q8ta_relu/Q8taRelu.cpp @@ -97,64 +97,29 @@ void q8ta_relu_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(Q8taReluParams)); graph.add_uniform_buffer_bytes(sizeof(Q8taReluParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kQ8taReluWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[3] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[3] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = params_buf; - bg_entries[2].size = sizeof(Q8taReluParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kQ8taReluWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + params_buf, + sizeof(Q8taReluParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "q8ta_relu", workgroup_count.y}); @@ -185,9 +150,6 @@ void q8ta_relu_impl(WebGPUGraph& graph, const std::vector& args) { g.set_cur_dims(out_id, d); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/quantize/QuantizePerTensor.cpp b/backends/webgpu/runtime/ops/quantize/QuantizePerTensor.cpp index 43405bd4962..c9b5b03ff16 100644 --- a/backends/webgpu/runtime/ops/quantize/QuantizePerTensor.cpp +++ b/backends/webgpu/runtime/ops/quantize/QuantizePerTensor.cpp @@ -95,71 +95,30 @@ void quantize_per_tensor_impl( utils::make_uniform(device, ¶ms, sizeof(QuantParams)); graph.add_uniform_buffer_bytes(sizeof(QuantParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kQuantizePerTensorWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[3] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[3] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = params_buf; - bg_entries[2].size = sizeof(QuantParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kQuantizePerTensorWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, WGPUBufferBindingType_Uniform, params_buf, sizeof(QuantParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count.x, "quantize_per_tensor", workgroup_count.y}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index a931e22b08d..5280430df45 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -335,38 +335,6 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { wgpuBufferUnmap(uniform_buffer); graph.add_uniform_buffer_bytes(sizeof(Q4gswParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {shader_src, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group layout: out (rw) + in/weight/scales/bias (ro storage) + uniform. - WGPUBindGroupLayoutEntry entries[6] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - for (uint32_t i = 1; i <= 4; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - } - entries[5].binding = 5; - entries[5].visibility = WGPUShaderStage_Compute; - entries[5].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 6; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - // GEMV/tiled wire an override wg_size; steel (256) + shmem (64) are fixed. const bool fixed_wg = use_steel || use_shmem_gemm; WGPUConstantEntry wg_size_constant = {}; @@ -374,43 +342,31 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { wg_size_constant.value = static_cast(use_gemv ? gemv_wg_size : wg_size); - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = fixed_wg ? 0u : 1u; - pipeline_desc.compute.constants = fixed_wg ? nullptr : &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[6] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = out.buffer; - bg_entries[0].size = out.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = in.buffer; - bg_entries[1].size = in.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = weight.buffer; - bg_entries[2].size = weight.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = scales.buffer; - bg_entries[3].size = scales.nbytes; - bg_entries[4].binding = 4; - bg_entries[4].buffer = bias_buffer; - bg_entries[4].size = bias_size; - bg_entries[5].binding = 5; - bg_entries[5].buffer = uniform_buffer; - bg_entries[5].size = sizeof(Q4gswParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 6; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + shader_src, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight.buffer, + weight.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + scales.buffer, + scales.nbytes}, + {4, WGPUBufferBindingType_ReadOnlyStorage, bias_buffer, bias_size}, + {5, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(Q4gswParams)}, + }, + fixed_wg ? nullptr : &wg_size_constant, + fixed_wg ? 0u : 1u); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count, "linear_q4gsw"}); + {bundle.pipeline, bundle.bind_group, workgroup_count, "linear_q4gsw"}); // Dynamic shapes: recompute dispatch + params.M for the live M. use_gemv and // use_shmem_gemm are captured (routing is fixed at build); the helper re-runs @@ -478,9 +434,6 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { g.set_cur_dims(out_id, od); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/repeat/Repeat.cpp b/backends/webgpu/runtime/ops/repeat/Repeat.cpp index 67aaae25f43..570e1bafcf0 100644 --- a/backends/webgpu/runtime/ops/repeat/Repeat.cpp +++ b/backends/webgpu/runtime/ops/repeat/Repeat.cpp @@ -66,74 +66,32 @@ void repeat_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, &in_meta, sizeof(TensorMeta)); graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kRepeatWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // Bind group: in, out (rw), out_meta, in_meta (2 uniforms). - WGPUBindGroupLayoutEntry entries[4] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[4] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_meta_buf; - bg_entries[2].size = sizeof(TensorMeta); - bg_entries[3].binding = 3; - bg_entries[3].buffer = in_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kRepeatWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {3, WGPUBufferBindingType_Uniform, in_meta_buf, sizeof(TensorMeta)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "repeat", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "repeat", + workgroup_count.y}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Drop our refs; the bind group keeps the uniforms alive until release. wgpuBufferRelease(out_meta_buf); wgpuBufferRelease(in_meta_buf); diff --git a/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp b/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp index d9f6b89e9f2..50b8828de93 100644 --- a/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp +++ b/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp @@ -134,49 +134,7 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector& args) { // group + dispatch. const bool use_vec4 = (row_width % 4u == 0u); - // Create shader module from built-in WGSL source - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {use_vec4 ? kRmsNormVec4WGSL : kRmsNormWGSL, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Create bind group layout: out (rw) + in/weight (ro storage) + params - WGPUBindGroupLayoutEntry entries[4] = {}; - - // t_out - storage buffer, read-write - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - - // t_in - storage buffer, read-only - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - // t_weight - storage buffer, read-only - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - // params - uniform buffer - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - // Create pipeline layout - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); + const char* shader_src = use_vec4 ? kRmsNormVec4WGSL : kRmsNormWGSL; // Runtime-overridable workgroup size (mirrors add op); clamp only reduces. // Pow2 required: the kernel halves the reduction stride (wg_size / 2u). @@ -186,43 +144,31 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector& args) { wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - // Create compute pipeline - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - // Create bind group with actual buffers const auto& out_tensor = graph.get_tensor(out_id); const auto& weight_tensor = graph.get_tensor(weight_id); - - WGPUBindGroupEntry bg_entries[4] = {}; - - bg_entries[0].binding = 0; - bg_entries[0].buffer = out_tensor.buffer; - bg_entries[0].size = out_tensor.nbytes; - - bg_entries[1].binding = 1; - bg_entries[1].buffer = in_tensor.buffer; - bg_entries[1].size = in_tensor.nbytes; - - bg_entries[2].binding = 2; - bg_entries[2].buffer = weight_tensor.buffer; - bg_entries[2].size = weight_tensor.nbytes; - - bg_entries[3].binding = 3; - bg_entries[3].buffer = uniform_buffer; - bg_entries[3].size = sizeof(RmsNormParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + shader_src, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight_tensor.buffer, + weight_tensor.nbytes}, + {3, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(RmsNormParams)}, + }, + &wg_size_constant, + 1); // One workgroup per row (kRmsNormWorkgroupSizeX threads cooperate per row) static_assert( @@ -232,7 +178,7 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector& args) { kRmsNormVec4WorkgroupSizeX == 64, "kRmsNormVec4WorkgroupSizeX must match override wg_size default (64)"); const size_t dispatch_idx = - graph.add_dispatch({pipeline, bind_group, num_rows}); + graph.add_dispatch({bundle.pipeline, bundle.bind_group, num_rows}); // Dynamic shapes: recompute num_rows + rewrite the UBO for the live input. WGPUBuffer params_buf = uniform_buffer; @@ -244,10 +190,6 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector& args) { g, in_id, out_id, row_width, epsilon, dispatch_idx, params_buf); }); - // Release intermediate objects (pipeline and bind_group are kept by dispatch) - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp index af1783496d2..012875b21e7 100644 --- a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp +++ b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp @@ -16,7 +16,9 @@ #include #include +#include #include +#include namespace executorch::backends::webgpu { @@ -46,8 +48,8 @@ struct RopeDispatch { RopeDispatch add_rope_dispatch( WebGPUGraph& graph, WGPUDevice device, - WGPUComputePipeline pipeline, - WGPUBindGroupLayout bgl, + std::optional& shared_resources, + uint32_t wg_size, const WebGPUTensor& x, const WebGPUTensor& out, const WebGPUTensor& freqs_cos, @@ -79,31 +81,37 @@ RopeDispatch add_rope_dispatch( wgpuBufferUnmap(uniform_buffer); graph.add_uniform_buffer_bytes(sizeof(RotaryParams)); - WGPUBindGroupEntry bg_entries[5] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = out.buffer; - bg_entries[0].size = out.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = x.buffer; - bg_entries[1].size = x.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = freqs_cos.buffer; - bg_entries[2].size = freqs_cos.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = freqs_sin.buffer; - bg_entries[3].size = freqs_sin.nbytes; - bg_entries[4].binding = 4; - bg_entries[4].buffer = uniform_buffer; - bg_entries[4].size = sizeof(RotaryParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + const std::vector bindings = { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, x.buffer, x.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + freqs_cos.buffer, + freqs_cos.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + freqs_sin.buffer, + freqs_sin.nbytes}, + {4, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(RotaryParams)}, + }; + utils::ComputePipelineBundle bundle = shared_resources.has_value() + ? utils::make_compute_pipeline( + device, *shared_resources, bindings, &wg_size_constant, 1) + : utils::make_compute_pipeline( + device, kRotaryEmbeddingWGSL, bindings, &wg_size_constant, 1); const size_t dispatch_index = graph.add_dispatch( - {pipeline, bind_group, workgroup_count, "apply_rotary_emb"}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count, + "apply_rotary_emb"}); + if (!shared_resources.has_value()) { + shared_resources.emplace(std::move(bundle)); + } // Graph owns it so a resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); @@ -256,59 +264,12 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { wg_size, "apply_rotary_emb"); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kRotaryEmbeddingWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group: out (rw) + in/freqs_cos/freqs_sin (ro) + uniform. - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - for (uint32_t i = 1; i <= 3; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - } - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUConstantEntry wg_size_constant = {}; - wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = static_cast(wg_size); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - // One pipeline per dispatch; a shared handle would double-free. - WGPUComputePipeline pipeline_q = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - WGPUComputePipeline pipeline_k = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - + std::optional shared_resources; RopeDispatch q_disp = add_rope_dispatch( graph, device, - pipeline_q, - bgl, + shared_resources, + wg_size, xq, xq_out, freqs_cos, @@ -320,8 +281,8 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { RopeDispatch k_disp = add_rope_dispatch( graph, device, - pipeline_k, - bgl, + shared_resources, + wg_size, xk, xk_out, freqs_cos, @@ -372,11 +333,6 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { }; graph.add_tensor_resize_hook(xq_id, rope_hook); graph.add_tensor_resize_hook(xk_id, rope_hook); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); - // pipeline_q/pipeline_k owned by their dispatches; graph dtor frees. } // HuggingFace rotate-half RoPE (Qwen3 etc.). Structural sibling of the @@ -401,7 +357,6 @@ static_assert(sizeof(RotaryHfParams) == 32, "RotaryHfParams must be 32 bytes"); RopeDispatch add_rope_hf_dispatch( WebGPUGraph& graph, - WGPUDevice device, uint32_t wg_size, const WebGPUTensor& x, const WebGPUTensor& out, @@ -426,15 +381,8 @@ RopeDispatch add_rope_hf_dispatch( params.rotary_dim = rotary_dim; params.start_pos = start_pos; - WGPUBufferDescriptor uniform_desc = {}; - uniform_desc.size = sizeof(RotaryHfParams); - uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; - uniform_desc.mappedAtCreation = true; - WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc); - void* mapped = - wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(RotaryHfParams)); - std::memcpy(mapped, ¶ms, sizeof(RotaryHfParams)); - wgpuBufferUnmap(uniform_buffer); + WGPUBuffer uniform_buffer = + utils::make_uniform(graph.device(), ¶ms, sizeof(RotaryHfParams)); graph.add_uniform_buffer_bytes(sizeof(RotaryHfParams)); WGPUConstantEntry wg_size_constant = {}; @@ -442,7 +390,7 @@ RopeDispatch add_rope_hf_dispatch( wg_size_constant.value = static_cast(wg_size); utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( - device, + graph.device(), kRotaryEmbeddingHfWGSL, { {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, @@ -544,33 +492,26 @@ void resize_rope_hf( g.set_cur_dims(xk_out_id, kd); } -// args: [xq, xk, freqs_cos, freqs_sin, start_pos, out_list(ValueList[xq_out, -// xk_out])]. freqs is the FULL [max_seq, rotary_dim] table (start_pos offsets -// into it), unlike the pre-sliced interleaved freqs. -void apply_rotary_emb_hf_impl( - WebGPUGraph& graph, - const std::vector& args) { - const int xq_id = args.at(0); - const int xk_id = args.at(1); - const int freqs_cos_id = args.at(2); - const int freqs_sin_id = args.at(3); - const int start_pos_id = args.at(4); - - const std::vector& out_list = graph.get_value_list(args.at(5)); - if (out_list.size() != 2) { - throw std::runtime_error( - "WebGPU apply_rotary_emb_hf: expected an output ValueList of size 2"); - } - - WGPUDevice device = graph.device(); - - const auto& xq = graph.get_tensor(xq_id); - const auto& xk = graph.get_tensor(xk_id); - const auto& freqs_cos = graph.get_tensor(freqs_cos_id); - const auto& freqs_sin = graph.get_tensor(freqs_sin_id); - const auto& xq_out = graph.get_tensor(out_list[0]); - const auto& xk_out = graph.get_tensor(out_list[1]); +// Validated HF-rope shape, derived from the input tensors. +struct RopeHfShape { + uint32_t head_dim; + uint32_t seq; + uint32_t n_heads_q; + uint32_t n_heads_k; + uint32_t rotary_dim; + uint32_t half_dim; + uint64_t xq_numel; + uint64_t xk_numel; +}; +// Derive and validate the HF-rope input shapes; throws on any malformed input. +RopeHfShape validate_rope_hf_inputs( + const WebGPUTensor& xq, + const WebGPUTensor& xk, + const WebGPUTensor& freqs_cos, + const WebGPUTensor& freqs_sin, + const WebGPUTensor& xq_out, + const WebGPUTensor& xk_out) { // Shape contract: xq/xk (B,S,n_heads,head_dim), freqs (max_seq, rotary_dim). if (xq.dims.size() < 3 || xk.dims.size() < 3 || freqs_cos.dims.size() < 2) { throw std::runtime_error("WebGPU apply_rotary_emb_hf: malformed dims"); @@ -613,26 +554,6 @@ void apply_rotary_emb_hf_impl( throw std::runtime_error("WebGPU apply_rotary_emb_hf: null buffer binding"); } - const uint32_t half_dim = rotary_dim / 2u; - - // start_pos: build-time Int (baked) OR runtime SymInt (dynamic decode); - // mirrors sdpa's input_pos handling. - int64_t start_pos = 0; - const auto start_pos_type = graph.get_value_type(start_pos_id); - const bool dynamic_pos = start_pos_type == WebGPUGraph::ValueType::SymInt; - if (dynamic_pos) { - start_pos = graph.read_symint(start_pos_id); // build placeholder (e.g. 0) - } else if (start_pos_type == WebGPUGraph::ValueType::Int) { - start_pos = graph.get_int(start_pos_id); - } else { - throw std::runtime_error( - "WebGPU apply_rotary_emb_hf: start_pos must be Int or SymInt"); - } - if (start_pos < 0) { - throw std::runtime_error( - "WebGPU apply_rotary_emb_hf: start_pos must be non-negative"); - } - // All tensors are fp32; output shapes equal their inputs. const uint64_t xq_numel = utils::numel_of(xq.dims); const uint64_t xk_numel = utils::numel_of(xk.dims); @@ -654,6 +575,73 @@ void apply_rotary_emb_hf_impl( "WebGPU apply_rotary_emb_hf: pair count exceeds uint32 dispatch range"); } + return { + head_dim, + seq, + n_heads_q, + n_heads_k, + rotary_dim, + rotary_dim / 2u, + xq_numel, + xk_numel}; +} + +// args: [xq, xk, freqs_cos, freqs_sin, start_pos, out_list(ValueList[xq_out, +// xk_out])]. freqs is the FULL [max_seq, rotary_dim] table (start_pos offsets +// into it), unlike the pre-sliced interleaved freqs. +void apply_rotary_emb_hf_impl( + WebGPUGraph& graph, + const std::vector& args) { + const int xq_id = args.at(0); + const int xk_id = args.at(1); + const int freqs_cos_id = args.at(2); + const int freqs_sin_id = args.at(3); + const int start_pos_id = args.at(4); + + const std::vector& out_list = graph.get_value_list(args.at(5)); + if (out_list.size() != 2) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: expected an output ValueList of size 2"); + } + + WGPUDevice device = graph.device(); + + const auto& xq = graph.get_tensor(xq_id); + const auto& xk = graph.get_tensor(xk_id); + const auto& freqs_cos = graph.get_tensor(freqs_cos_id); + const auto& freqs_sin = graph.get_tensor(freqs_sin_id); + const auto& xq_out = graph.get_tensor(out_list[0]); + const auto& xk_out = graph.get_tensor(out_list[1]); + + const RopeHfShape shp = + validate_rope_hf_inputs(xq, xk, freqs_cos, freqs_sin, xq_out, xk_out); + const uint32_t head_dim = shp.head_dim; + const uint32_t seq = shp.seq; + const uint32_t n_heads_q = shp.n_heads_q; + const uint32_t n_heads_k = shp.n_heads_k; + const uint32_t rotary_dim = shp.rotary_dim; + const uint32_t half_dim = shp.half_dim; + const uint64_t xq_numel = shp.xq_numel; + const uint64_t xk_numel = shp.xk_numel; + + // start_pos: build-time Int (baked) OR runtime SymInt (dynamic decode); + // mirrors sdpa's input_pos handling. + int64_t start_pos = 0; + const auto start_pos_type = graph.get_value_type(start_pos_id); + const bool dynamic_pos = start_pos_type == WebGPUGraph::ValueType::SymInt; + if (dynamic_pos) { + start_pos = graph.read_symint(start_pos_id); // build placeholder (e.g. 0) + } else if (start_pos_type == WebGPUGraph::ValueType::Int) { + start_pos = graph.get_int(start_pos_id); + } else { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: start_pos must be Int or SymInt"); + } + if (start_pos < 0) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: start_pos must be non-negative"); + } + const uint32_t wg_size = utils::clamp_workgroup_size(device, kRotaryEmbeddingHfWorkgroupSizeX); // Validate both dispatches before any GPU-object alloc (no leak on throw). @@ -670,7 +658,6 @@ void apply_rotary_emb_hf_impl( RopeDispatch q_disp = add_rope_hf_dispatch( graph, - device, wg_size, xq, xq_out, @@ -685,7 +672,6 @@ void apply_rotary_emb_hf_impl( xq_wgc); RopeDispatch k_disp = add_rope_hf_dispatch( graph, - device, wg_size, xk, xk_out, diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index f3345c89c77..96ec782c72e 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -24,6 +24,7 @@ #include #include #include +#include namespace executorch::backends::webgpu { @@ -152,84 +153,47 @@ void build_dispatch( const char* kernel_name = "") { WGPUDevice device = graph.device(); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {wgsl_source, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // Bind group layout: storage entries then the uniform. constexpr uint32_t kMaxEntries = 8; if (n_storage + 1 > kMaxEntries) { throw std::runtime_error("WebGPU sdpa: n_storage exceeds kMaxEntries"); } - WGPUBindGroupLayoutEntry bgl_entries[kMaxEntries] = {}; const uint32_t uniform_binding = n_storage; + std::vector bindings; + bindings.reserve(n_storage + 1u); for (uint32_t i = 0; i < n_storage; i++) { - bgl_entries[i].binding = i; - bgl_entries[i].visibility = WGPUShaderStage_Compute; - bgl_entries[i].buffer.type = (i == 0) - ? WGPUBufferBindingType_Storage - : WGPUBufferBindingType_ReadOnlyStorage; + bindings.push_back( + {i, + (i == 0) ? WGPUBufferBindingType_Storage + : WGPUBufferBindingType_ReadOnlyStorage, + storage_bindings[i].buffer, + storage_bindings[i].size}); } - bgl_entries[uniform_binding].binding = uniform_binding; - bgl_entries[uniform_binding].visibility = WGPUShaderStage_Compute; - bgl_entries[uniform_binding].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = n_storage + 1; - bgl_desc.entries = bgl_entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); + bindings.push_back( + {uniform_binding, + WGPUBufferBindingType_Uniform, + uniform_buffer, + uniform_size}); // All callers pass an override wg_size; a 0 would keep the shader default. WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - if (wg_size != 0) { - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - } - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[kMaxEntries] = {}; - for (uint32_t i = 0; i < n_storage; i++) { - bg_entries[i].binding = i; - bg_entries[i].buffer = storage_bindings[i].buffer; - bg_entries[i].size = storage_bindings[i].size; - } - bg_entries[uniform_binding].binding = uniform_binding; - bg_entries[uniform_binding].buffer = uniform_buffer; - bg_entries[uniform_binding].size = uniform_size; - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = n_storage + 1; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + wgsl_source, + bindings, + wg_size != 0 ? &wg_size_constant : nullptr, + wg_size != 0 ? 1u : 0u); graph.add_dispatch( - {pipeline, - bind_group, + {bundle.pipeline, + bundle.bind_group, workgroup_count_x, kernel_name, workgroup_count_y}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); if (retain_uniform) { // Graph owns it so a resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp index ffd3b24dce3..82ead71b597 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp @@ -20,6 +20,7 @@ #include #include #include +#include namespace executorch::backends::webgpu { @@ -72,70 +73,34 @@ void build_dispatch( const char* kernel_name) { WGPUDevice device = graph.device(); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {wgsl_source, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - constexpr uint32_t kMaxEntries = 8; if (n_storage + 1u > kMaxEntries) { throw std::runtime_error( "WebGPU sdpa FlashDecoding: bind group entry count exceeds kMaxEntries"); } - WGPUBindGroupLayoutEntry bgl_entries[kMaxEntries] = {}; const uint32_t uniform_binding = n_storage; + std::vector bindings; + bindings.reserve(n_storage + 1u); for (uint32_t i = 0; i < n_storage; i++) { - bgl_entries[i].binding = i; - bgl_entries[i].visibility = WGPUShaderStage_Compute; - bgl_entries[i].buffer.type = (i < n_rw) - ? WGPUBufferBindingType_Storage - : WGPUBufferBindingType_ReadOnlyStorage; - } - bgl_entries[uniform_binding].binding = uniform_binding; - bgl_entries[uniform_binding].visibility = WGPUShaderStage_Compute; - bgl_entries[uniform_binding].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = n_storage + 1; - bgl_desc.entries = bgl_entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[kMaxEntries] = {}; - for (uint32_t i = 0; i < n_storage; i++) { - bg_entries[i].binding = i; - bg_entries[i].buffer = storage_bindings[i].buffer; - bg_entries[i].size = storage_bindings[i].size; + bindings.push_back( + {i, + (i < n_rw) ? WGPUBufferBindingType_Storage + : WGPUBufferBindingType_ReadOnlyStorage, + storage_bindings[i].buffer, + storage_bindings[i].size}); } - bg_entries[uniform_binding].binding = uniform_binding; - bg_entries[uniform_binding].buffer = uniform_buffer; - bg_entries[uniform_binding].size = uniform_size; + bindings.push_back( + {uniform_binding, + WGPUBufferBindingType_Uniform, + uniform_buffer, + uniform_size}); - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = n_storage + 1; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = + utils::make_compute_pipeline(device, wgsl_source, bindings); - graph.add_dispatch({pipeline, bind_group, workgroup_count_x, kernel_name}); + graph.add_dispatch( + {bundle.pipeline, bundle.bind_group, workgroup_count_x, kernel_name}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); wgpuBufferRelease(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/unary/pow_scalar.wgsl b/backends/webgpu/runtime/ops/unary/pow_scalar.wgsl index db97cdffafc..701da7df73f 100644 --- a/backends/webgpu/runtime/ops/unary/pow_scalar.wgsl +++ b/backends/webgpu/runtime/ops/unary/pow_scalar.wgsl @@ -17,5 +17,21 @@ fn main( if (idx >= params.num_elements) { return; } - output[idx] = pow(input[idx], params.minimum); + // WGSL pow(x, y) = exp2(y*log2(x)) is undefined for x < 0. Fix the common + // integer-exponent case to match PyTorch (|x|^y, sign-preserved for odd y); + // the rare non-integer negative case keeps the prior WGSL behavior. + let x = input[idx]; + let e = params.minimum; + var r: f32; + if (x >= 0.0) { + r = pow(x, e); + } else if (fract(e) == 0.0) { + r = pow(abs(x), e); + if (fract(e * 0.5) != 0.0) { + r = -r; + } + } else { + r = pow(x, e); + } + output[idx] = r; } diff --git a/backends/webgpu/runtime/ops/unary/pow_scalar_wgsl.h b/backends/webgpu/runtime/ops/unary/pow_scalar_wgsl.h index e296ae59a5e..8707bc75c5b 100644 --- a/backends/webgpu/runtime/ops/unary/pow_scalar_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/pow_scalar_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from pow_scalar.wgsl - DO NOT EDIT. -// wgsl-sha256: a9176c6a6b0da421cd3649e396390cd10859c1256f6d38fce6bc730d3e3788e6 +// wgsl-sha256: b5eda5edff9e51b475e258ddb14ce5bf1cc2399fcfc02b42d306cdfd5f4bf97b inline constexpr const char* kPowScalarWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -34,7 +34,23 @@ fn main( if (idx >= params.num_elements) { return; } - output[idx] = pow(input[idx], params.minimum); + // WGSL pow(x, y) = exp2(y*log2(x)) is undefined for x < 0. Fix the common + // integer-exponent case to match PyTorch (|x|^y, sign-preserved for odd y); + // the rare non-integer negative case keeps the prior WGSL behavior. + let x = input[idx]; + let e = params.minimum; + var r: f32; + if (x >= 0.0) { + r = pow(x, e); + } else if (fract(e) == 0.0) { + r = pow(abs(x), e); + if (fract(e * 0.5) != 0.0) { + r = -r; + } + } else { + r = pow(x, e); + } + output[idx] = r; } )"; diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 26b82699e38..19438d6acbc 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -1217,7 +1217,10 @@ def qkv(b, h, sq, skv, d): ) -from executorch.backends.webgpu.test.ops.test_embedding import EmbeddingModule +from executorch.backends.webgpu.test.ops.test_embedding import ( + _det_weight as _emb_det_weight, + EmbeddingModule, +) def _emb_idx_small(_shape): @@ -1234,7 +1237,7 @@ def _embedding_suite() -> WebGPUTestSuite: # framework's int-input path. return WebGPUTestSuite( module_factory=lambda num_embeddings, embed_dim: EmbeddingModule( - num_embeddings, embed_dim + _emb_det_weight(num_embeddings, embed_dim) ), cases=[ Case( @@ -1605,10 +1608,9 @@ def _sub_suite() -> WebGPUTestSuite: # over a TensorMeta UBO); fp64 golden. Mirrors _mul_suite. alpha is a # construct kwarg baked into the .pte, never a serialized input. return WebGPUTestSuite( - module_factory=lambda alpha=1.0: SubModule(alpha), + module_factory=lambda: SubModule(), cases=[ - Case(name=name, construct={"alpha": alpha}, inputs=(sa, sb)) - for name, (sa, sb, alpha) in _SUB_CONFIGS.items() + Case(name=name, inputs=(sa, sb)) for name, (sa, sb) in _SUB_CONFIGS.items() ], ) @@ -1664,7 +1666,9 @@ def _hardtanh_suite() -> WebGPUTestSuite: @register_op_test("pow_scalar") def _pow_scalar_suite() -> WebGPUTestSuite: - # Positive base: WGSL pow(neg base)=NaN for any exponent; exponent baked. + # Positive-base configs (exponent baked). Plus a negative-base integer-exponent + # case: pow(x, 2) == x*x. WGSL pow(neg) is undefined, so the shader special-cases + # it (fixes F.normalize's q^2 in Swin-V2 cosine attention). return WebGPUTestSuite( module_factory=lambda exponent: PowScalarModule(exponent), cases=[ @@ -1674,6 +1678,13 @@ def _pow_scalar_suite() -> WebGPUTestSuite: inputs=(InputSpec(shape=(M1, M2), gen=_unary_lin(0.1, 4.0)),), ) for n, e in POW_SCALAR_CONFIGS.items() + ] + + [ + Case( + name="neg_base_sq", + construct={"exponent": 2.0}, + inputs=(InputSpec(shape=(M1, M2), gen=_unary_lin(-3.0, 3.0)),), + ) ], )