Update

[ghstack-poisoned]

Author: JulianCloudNTH

backends/webgpu/test/op_tests/cases.py

@@ -30,6 +30,10 @@
     AddModule,
     AddSelfModule,
 )
+from executorch.backends.webgpu.test.ops.mul.test_mul import (
+    CONFIGS as _MUL_CONFIGS,
+    MulModule,
+)
 from executorch.backends.webgpu.test.ops.rms_norm.test_rms_norm import (
     _CASES,
     _linspace_weight,
@@ -106,3 +110,14 @@ def _rms_norm_suite() -> WebGPUTestSuite:
             )
         )
     return WebGPUTestSuite(module_factory=_rms_norm_factory, cases=cases)
+
+
+@register_op_test("mul")
+def _mul_suite() -> WebGPUTestSuite:
+    # Full numeric coverage incl. broadcast (binary_mul.wgsl over a TensorMeta UBO); fp64 golden.
+    return WebGPUTestSuite(
+        module_factory=lambda: MulModule(),
+        cases=[
+            Case(name=name, inputs=(sa, sb)) for name, (sa, sb) in _MUL_CONFIGS.items()
+        ],
+    )

backends/webgpu/test/ops/mul/__init__.py

@@ -0,0 +1,5 @@
+# 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.

backends/webgpu/test/ops/mul/test_mul.py

@@ -0,0 +1,74 @@
+# 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.mul.Tensor` (full broadcast) module + 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). `TestMul` is
+the export-delegation + eager-correctness smoke test. Configs span the same-shape
+fast path (SwiGLU), last-dim broadcast at LLM width, and a mixed-rank left-pad case.
+"""
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> (shape_a, shape_b). Output shape is the broadcast of the two.
+CONFIGS = {
+    "same": ((8, 32), (8, 32)),  # fast path (SwiGLU same-shape)
+    "bcast_lastdim": ((1, 1, 7, 896), (1, 1, 7, 1)),  # last-dim broadcast, LLM width
+    "mixedrank": ((4,), (3, 4)),  # right-aligned left-pad (in.ndim < out.ndim)
+}
+
+
+class MulModule(torch.nn.Module):
+    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        return a * b
+
+
+def _det_inputs(shape_a, shape_b):
+    """Deterministic fp32 inputs (fixed seed) for a config."""
+    g = torch.Generator().manual_seed(0)
+    a = torch.randn(*shape_a, generator=g, dtype=torch.float32)
+    b = torch.randn(*shape_b, generator=g, dtype=torch.float32)
+    return a, b
+
+
+def _export(a: torch.Tensor, b: torch.Tensor):
+    ep = torch.export.export(MulModule().eval(), (a, b))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegated(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+class TestMul(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, (sa, sb) in CONFIGS.items():
+            a, b = _det_inputs(sa, sb)
+            et = _export(a, b)
+            self.assertTrue(
+                _delegated(et), f"Expected a VulkanBackend delegate (mul {name})"
+            )
+
+    def test_golden_matches_eager(self) -> None:
+        for _, (sa, sb) in CONFIGS.items():
+            a, b = _det_inputs(sa, sb)
+            torch.testing.assert_close(MulModule()(a, b), a * b)
+
+
+if __name__ == "__main__":
+    unittest.main()

backends/webgpu/test/tester.py

@@ -21,6 +21,7 @@
 WEBGPU_SUPPORTED_OPS = [
     exir_ops.edge.aten.add.Tensor,
     exir_ops.edge.et_vk.rms_norm.default,
+    exir_ops.edge.aten.mul.Tensor,
 ]