Add an aten.mul.Scalar kernel (#25026)

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(),

backends/webgpu/runtime/WebGPUShaderRegistry.cpp

@@ -139,6 +139,7 @@
 #include <executorch/backends/webgpu/runtime/ops/unary/erf_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/exp_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/unary/mul_scalar_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/neg_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/pow_scalar_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/round_wgsl.h>
@@ -160,7 +161,7 @@
 namespace executorch::backends::webgpu {
 namespace {
 
-constexpr std::array<WebGPUShaderInfo, 142> kShaderRegistry = {{
+constexpr std::array<WebGPUShaderInfo, 143> kShaderRegistry = {{
     {
         "abs",
         kAbsWGSL,
@@ -714,6 +715,13 @@ constexpr std::array<WebGPUShaderInfo, 142> kShaderRegistry = {{
         kMmVec4WorkgroupSizeY,
         kMmVec4WorkgroupSizeZ,
     },
+    {
+        "mul_scalar",
+        kMulScalarWGSL,
+        kMulScalarWorkgroupSizeX,
+        kMulScalarWorkgroupSizeY,
+        kMulScalarWorkgroupSizeZ,
+    },
     {
         "native_layer_norm",
         kNativeLayerNormWGSL,

backends/webgpu/runtime/ops/unary/Activations.cpp

@@ -16,6 +16,7 @@
 #include <executorch/backends/webgpu/runtime/ops/unary/erf_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/exp_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/unary/mul_scalar_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/neg_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/pow_scalar_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/round_wgsl.h>
@@ -178,6 +179,24 @@ void pow_scalar_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       exponent);
 }
 
+void mul_scalar_impl(WebGPUGraph& graph, const std::vector<int>& 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

backends/webgpu/runtime/ops/unary/mul_scalar.wgsl

@@ -0,0 +1,21 @@
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+
+struct Params {
+  num_elements: u32,
+  scalar: f32,
+}
+@group(0) @binding(2) var<uniform> params: Params;
+
+override wg_size: u32 = 256u;
+
+@compute @workgroup_size(wg_size)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    if (idx >= params.num_elements) {
+        return;
+    }
+    output[idx] = input[idx] * params.scalar;
+}

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 <cstdint>
+
+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<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+
+struct Params {
+  num_elements: u32,
+  scalar: f32,
+}
+@group(0) @binding(2) var<uniform> params: Params;
+
+override wg_size: u32 = 256u;
+
+@compute @workgroup_size(wg_size)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    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

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 = [

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";

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.

backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -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 @@
     _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 @@
     # 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 @@
                 "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",

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)

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}),