From 65da4d5cc42229648d224a2120eabf20938638c9 Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Thu, 12 Nov 2026 22:34:18 -0800 Subject: [PATCH] Add an aten.mul.Scalar kernel (#25026) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/25026 The WebGPU partitioner delegated aten.mul.Scalar (from the Vulkan op registry), but the WebGPU runtime had no kernel, so programs with SDPA's math decomposition failed to load. This adds a one-dispatch fp32 mul.Scalar with a resize hook, and the partitioner keeps non-fp32 mul.Scalar on CPU. Use case: unmodified Hugging Face ESM-2 and BERT-style models (SDPA without a fused kernel), including dynamic sequence length. Reviewed By: SS-JIA Differential Revision: D129036504 --- backends/webgpu/partitioner/partitioner.py | 5 +++ .../webgpu/runtime/WebGPUShaderRegistry.cpp | 10 ++++- .../webgpu/runtime/ops/unary/Activations.cpp | 20 +++++++++ .../webgpu/runtime/ops/unary/mul_scalar.wgsl | 21 +++++++++ .../runtime/ops/unary/mul_scalar_wgsl.h | 45 +++++++++++++++++++ backends/webgpu/test/BUCK | 1 + .../webgpu/test/native/test_dynamic_shape.cpp | 11 ++++- backends/webgpu/test/op_tests/cases.py | 23 ++++++++++ .../test_dynamic_shape_export.py | 19 ++++++++ backends/webgpu/test/ops/test_mul.py | 36 ++++++++++++++- .../webgpu/test/test_webgpu_partitioner.py | 17 +++++++ 11 files changed, 205 insertions(+), 3 deletions(-) create mode 100644 backends/webgpu/runtime/ops/unary/mul_scalar.wgsl create mode 100644 backends/webgpu/runtime/ops/unary/mul_scalar_wgsl.h diff --git a/backends/webgpu/partitioner/partitioner.py b/backends/webgpu/partitioner/partitioner.py index 30aa2d1d630..3f77f59cd70 100644 --- a/backends/webgpu/partitioner/partitioner.py +++ b/backends/webgpu/partitioner/partitioner.py @@ -115,6 +115,11 @@ def _fp32_unary_features() -> OpFeatures: _convolution_check, dtypes={torch.float32}, ), + exir_ops.edge.aten.mul.Scalar: _webgpu_features( + exir_ops.edge.aten.mul.Scalar, + lambda node, vulkan_check: vulkan_check(node), + dtypes={torch.float32}, + ), exir_ops.edge.aten.elu.default: _fp32_unary_features(), exir_ops.edge.aten.erf.default: _fp32_unary_features(), exir_ops.edge.aten.silu.default: _fp32_unary_features(), diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index 17b55f9f49a..540f2c50e50 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -139,6 +139,7 @@ #include #include #include +#include #include #include #include @@ -160,7 +161,7 @@ namespace executorch::backends::webgpu { namespace { -constexpr std::array kShaderRegistry = {{ +constexpr std::array kShaderRegistry = {{ { "abs", kAbsWGSL, @@ -714,6 +715,13 @@ constexpr std::array kShaderRegistry = {{ kMmVec4WorkgroupSizeY, kMmVec4WorkgroupSizeZ, }, + { + "mul_scalar", + kMulScalarWGSL, + kMulScalarWorkgroupSizeX, + kMulScalarWorkgroupSizeY, + kMulScalarWorkgroupSizeZ, + }, { "native_layer_norm", kNativeLayerNormWGSL, diff --git a/backends/webgpu/runtime/ops/unary/Activations.cpp b/backends/webgpu/runtime/ops/unary/Activations.cpp index f41aa577e17..0cdd1772df6 100644 --- a/backends/webgpu/runtime/ops/unary/Activations.cpp +++ b/backends/webgpu/runtime/ops/unary/Activations.cpp @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -178,6 +179,24 @@ void pow_scalar_impl(WebGPUGraph& graph, const std::vector& args) { exponent); } +void mul_scalar_impl(WebGPUGraph& graph, const std::vector& args) { + // aten.mul.Scalar args: [in, other, out]; other = min slot. + const auto vt = graph.get_value_type(args.at(1)); + if (vt != WebGPUGraph::ValueType::Int && + vt != WebGPUGraph::ValueType::Double) { + throw std::runtime_error( + "WebGPU mul.Scalar: other must be an Int or Double"); + } + add_unary_op( + graph, + args.at(0), + args.at(2), + kMulScalarWGSL, + kMulScalarWorkgroupSizeX, + "mul_scalar", + utils::scalar_or(graph, args.at(1), 0.0f)); +} + } // namespace WEBGPU_REGISTER_OPERATORS { @@ -197,6 +216,7 @@ WEBGPU_REGISTER_OPERATORS { 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); + WEBGPU_REGISTER_OP(aten.mul.Scalar, mul_scalar_impl); } } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/unary/mul_scalar.wgsl b/backends/webgpu/runtime/ops/unary/mul_scalar.wgsl new file mode 100644 index 00000000000..e7144ee5db1 --- /dev/null +++ b/backends/webgpu/runtime/ops/unary/mul_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, + scalar: 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] = input[idx] * params.scalar; +} diff --git a/backends/webgpu/runtime/ops/unary/mul_scalar_wgsl.h b/backends/webgpu/runtime/ops/unary/mul_scalar_wgsl.h new file mode 100644 index 00000000000..c790f7e2b80 --- /dev/null +++ b/backends/webgpu/runtime/ops/unary/mul_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 mul_scalar.wgsl - DO NOT EDIT. +// wgsl-sha256: 31c39096da58d01129a1c09c499c1a39aa0b06e61b47b5ddb9fba2e66d38968e +inline constexpr const char* kMulScalarWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, + scalar: 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] = input[idx] * params.scalar; +} +)"; + +inline constexpr uint32_t kMulScalarWorkgroupSizeX = 256; +inline constexpr uint32_t kMulScalarWorkgroupSizeY = 1; +inline constexpr uint32_t kMulScalarWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/BUCK b/backends/webgpu/test/BUCK index 0567f19fdd7..c2f72f32644 100644 --- a/backends/webgpu/test/BUCK +++ b/backends/webgpu/test/BUCK @@ -40,6 +40,7 @@ fbcode_target( "ops/test_conv_transpose1d.py", "ops/test_elu_silu.py", "ops/test_erf.py", + "ops/test_mul.py", "test_webgpu_partitioner.py", ], deps = [ diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index a4ab5664d89..e6bd8246c7e 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -21,7 +21,7 @@ // O ONE dyn_conv1d graph reused across live input lengths // P ONE dyn_gelu graph reused above -> at -> above the old 1D dispatch cap // Q dyn_layer_norm(_rows) R dyn_mean_keepdim S dyn_elu_silu (reused) -// S2 dyn_erf_gelu (reused) +// S2 dyn_erf_gelu (reused) S3 dyn_mul_scalar (reused) // .pte + goldens from test/ops/dynamic_shape/test_dynamic_shape_export.py. // // Artifacts dir: $WEBGPU_DYNAMIC_SHAPE_DIR, else argv[1], else @@ -1971,6 +1971,15 @@ TEST(DynamicShape, ErfGeluReusedGraph) { } } +// S3: aten.mul.Scalar(x, 0.125) on ONE loaded graph, shrinking and growing S. +TEST(DynamicShape, MulScalarReusedGraph) { + Module m(g_dir + "/dyn_mul_scalar.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_mul_scalar.pte"; + for (int s : {128, 7, 1, 125, 7}) { + check_s(m, "dyn_mul_scalar", s); + } +} + TEST(DynamicShape, SwiGluReusedGraph) { Module module(g_dir + "/dyn_swiglu.pte"); ASSERT_EQ(module.load_forward(), Error::Ok) << "load dyn_swiglu.pte"; diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index a61c15636d0..f0a503daad5 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -111,6 +111,8 @@ from executorch.backends.webgpu.test.ops.test_mul import ( CONFIGS as _MUL_CONFIGS, MulModule, + MulScalarModule, + SCALAR_CONFIGS as _MUL_SCALAR_CONFIGS, ) from executorch.backends.webgpu.test.ops.test_permute import ( CONFIGS as _PERMUTE_CONFIGS, @@ -330,6 +332,27 @@ def _mul_suite() -> WebGPUTestSuite: ) +@register_op_test("mul_scalar") +def _mul_scalar_suite() -> WebGPUTestSuite: + return WebGPUTestSuite( + module_factory=MulScalarModule, + cases=[ + Case(name=name, construct={"scalar": scalar}, inputs=((4, 257),)) + for name, scalar in _MUL_SCALAR_CONFIGS.items() + ] + + [ + Case(name="rank4", construct={"scalar": 0.125}, inputs=((S1, S2, S2, M2),)), + Case( + name="param_input", + construct={"scalar": -1.5, "param_shape": (4, 257)}, + inputs=((4, 257),), + ), + ], + atol=1e-6, + rtol=1e-6, + ) + + def _fn_config_suite(module_cls, configs) -> WebGPUTestSuite: """Builder for ops whose per-case spec is a (shape, fn) pair (view/select/slice). The fn is a `construct` kwarg baked into the .pte module, never a serialized input. diff --git a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py index e4641c4de05..663b2507a31 100644 --- a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py +++ b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py @@ -23,6 +23,7 @@ from executorch.backends.webgpu.partitioner import WebGPUPartitioner from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dModule from executorch.backends.webgpu.test.ops.test_erf import ErfGeluModule +from executorch.backends.webgpu.test.ops.test_mul import MulScalarModule from executorch.backends.webgpu.test.ops.test_gelu import GeluModule from executorch.exir import to_edge_transform_and_lower from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes @@ -453,6 +454,18 @@ def export_dynamic_erf_gelu_cases(out_dir: str, s_dim) -> None: _write_goldens(erf_gelu, "dyn_erf_gelu", out_dir, [MAXS, 125, 7, 1]) +def export_dynamic_mul_scalar_cases(out_dir: str, s_dim) -> None: + mul_scalar = MulScalarModule(0.125) + _export( + mul_scalar, + (_ramp((1, 1, MAXS, HIDDEN)),), + {"x": {2: s_dim}}, + os.path.join(out_dir, "dyn_mul_scalar.pte"), + partitioner_cls=WebGPUPartitioner, + ) + _write_goldens(mul_scalar, "dyn_mul_scalar", out_dir, [MAXS, 125, 7, 1]) + + def export_dynamic_shape_cases(out_dir: str) -> None: """Write the dynamic + static .pte's and per-S goldens for the native test.""" os.makedirs(out_dir, exist_ok=True) @@ -571,6 +584,9 @@ def export_dynamic_shape_cases(out_dir: str) -> None: # 2h.4) The hand-written exact GELU of ESM-2 / BERT (erf) with a DYNAMIC S. export_dynamic_erf_gelu_cases(out_dir, s_dim) + # 2h.5) aten.mul.Scalar (from SDPA's math decomposition) with a DYNAMIC S. + export_dynamic_mul_scalar_cases(out_dir, s_dim) + # 2i) Dynamic SwiGLU graph patterns. The Llama-width fixture forces a 2D # fused dispatch at M=512; compact variants isolate commutative matching and # graph/intermediate ownership guards around same-input q4 projections. @@ -1788,6 +1804,9 @@ def test_export_dynamic_rms(self) -> None: "dyn_erf_gelu.pte", "dyn_erf_gelu.S125.input.bin", "dyn_erf_gelu.S7.golden.bin", + "dyn_mul_scalar.pte", + "dyn_mul_scalar.S125.input.bin", + "dyn_mul_scalar.S7.golden.bin", "dyn_linear_bk64.pte", "dyn_linear_bk64.S512.input.bin", "dyn_linear_bk64.S512.golden.bin", diff --git a/backends/webgpu/test/ops/test_mul.py b/backends/webgpu/test/ops/test_mul.py index 36d5429c25a..5c6337acdb9 100644 --- a/backends/webgpu/test/ops/test_mul.py +++ b/backends/webgpu/test/ops/test_mul.py @@ -4,7 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""`aten.mul.Tensor` (full broadcast) module + configs for the WebGPU op-test framework. +"""`aten.mul.Tensor` (full broadcast) and `aten.mul.Scalar` modules + configs for the +WebGPU op-test framework. `MulModule` + `CONFIGS` are imported by `cases.py` to drive the declarative op-test suite (export via VulkanPartitioner + fp64 torch golden, run on Dawn). `MulTest` is @@ -35,6 +36,39 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return a * b +# name -> scalar for `aten.mul.Scalar`, which SDPA's math decomposition emits +# with 1.0 (and which survives to Edge, unlike `x * 2.0`). +SCALAR_CONFIGS = { + "one": 1.0, + "negative": -2.5, + "int": 3, + "zero": 0.0, +} + + +class MulScalarModule(torch.nn.Module): + """`aten.mul.Scalar`; with `param_shape`, on a parameter added to the input.""" + + def __init__(self, scalar: float = 1.0, param_shape=None) -> None: + super().__init__() + self.scalar = scalar + self.param = ( + torch.nn.Parameter( + torch.linspace(-4.0, 4.0, param_shape[0] * param_shape[1]).reshape( + param_shape + ), + requires_grad=False, + ) + if param_shape is not None + else None + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.param is not None: + return torch.ops.aten.mul.Scalar(self.param, self.scalar) + x + return torch.ops.aten.mul.Scalar(x, self.scalar) + + def _det_inputs(shape_a, shape_b): """Deterministic fp32 inputs (fixed seed) for a config.""" g = torch.Generator().manual_seed(0) diff --git a/backends/webgpu/test/test_webgpu_partitioner.py b/backends/webgpu/test/test_webgpu_partitioner.py index 50d453fb98b..f41ba37b878 100644 --- a/backends/webgpu/test/test_webgpu_partitioner.py +++ b/backends/webgpu/test/test_webgpu_partitioner.py @@ -29,6 +29,7 @@ SiluMulModule, ) from executorch.backends.webgpu.test.ops.test_erf import ErfGeluModule, ErfModule +from executorch.backends.webgpu.test.ops.test_mul import MulScalarModule from executorch.backends.webgpu.test.tester import Partition, ToEdgeTransformAndLower from executorch.exir import ( ExecutorchProgramManager, @@ -372,6 +373,22 @@ def test_erf_is_delegated(self) -> None: self.assertEqual(non_lowered, []) self.assertIn(exir_ops.edge.aten.erf.default, delegated) + def test_mul_scalar_is_delegated_in_fp32_only(self) -> None: + mul_scalar = exir_ops.edge.aten.mul.Scalar + x = torch.linspace(-4.0, 4.0, 64).reshape(4, 16) + for name, (inputs, delegated) in { + "fp32": ((x,), True), + "fp16": ((x.half(),), False), + }.items(): + with self.subTest(case=name): + non_lowered, in_delegate = self._lower_split( + WebGPUPartitioner(), MulScalarModule(0.5), inputs + ) + self.assertEqual(mul_scalar in in_delegate, delegated) + self.assertEqual( + mul_scalar in self._targets(non_lowered), not delegated + ) + def test_dynamic_erf_is_delegated(self) -> None: non_lowered, delegated = self._lower_split( WebGPUPartitioner({"require_dynamic_shapes": True}),