Add WebGPU mul and sigmoid test coverage

backends/webgpu/test/op_tests/cases.py

@@ -54,7 +54,9 @@
 )
 from executorch.backends.webgpu.test.ops.test_sigmoid import (
     _det_input as _sigmoid_det_input,
+    _wide_det_input as _sigmoid_wide_det_input,
     N as _SIGMOID_N,
+    SigmoidChainedModule,
     SigmoidModule,
 )
 
@@ -188,11 +190,22 @@ def _sigmoid_full_range(_shape) -> torch.Tensor:
     return _sigmoid_det_input()
 
 
+def _sigmoid_wide_range(_shape) -> torch.Tensor:
+    return _sigmoid_wide_det_input()
+
+
+def _sigmoid_factory(variant: str = "regular") -> torch.nn.Module:
+    return {
+        "regular": SigmoidModule,
+        "chained": SigmoidChainedModule,
+    }[variant]()
+
+
 @register_op_test("sigmoid")
 def _sigmoid_suite() -> WebGPUTestSuite:
     # sigmoid has no CONFIGS table; cover unary shapes directly (tol 1e-4).
     return WebGPUTestSuite(
-        module_factory=lambda: SigmoidModule(),
+        module_factory=_sigmoid_factory,
         cases=[
             Case(name="vec", inputs=((M1,),)),
             Case(name="mat", inputs=((M1, M2),)),
@@ -203,6 +216,15 @@ def _sigmoid_suite() -> WebGPUTestSuite:
                 name="saturation",
                 inputs=(InputSpec(shape=(_SIGMOID_N,), gen=_sigmoid_full_range),),
             ),
+            Case(
+                name="wide_saturation",
+                inputs=(InputSpec(shape=(_SIGMOID_N,), gen=_sigmoid_wide_range),),
+            ),
+            Case(
+                name="chained",
+                construct={"variant": "chained"},
+                inputs=(InputSpec(shape=(_SIGMOID_N,), gen=_sigmoid_full_range),),
+            ),
         ],
         atol=1e-4,
         rtol=1e-4,

backends/webgpu/test/op_tests/test_generator.py

@@ -77,7 +77,7 @@ def test_generate_manifest(tmp_path):
 def test_every_case_delegates():
     # Contract: every registered case must lower to a VulkanBackend delegate. An op that
     # silently CPU-falls-back would otherwise produce a misleading golden-equals-golden pass.
-    for op in ("add", "rms_norm"):
+    for op in ("add", "mul", "sigmoid", "rms_norm"):
         suite = op_test_registry[op]
         for case in suite.cases:
             _module, _inputs, prog = g.export_case(suite, case)

backends/webgpu/test/op_tests/test_schema.py

@@ -36,11 +36,13 @@ def _dummy():
     assert suite.atol == 1e-3 and isinstance(XL, int)
 
 
-def test_add_rms_norm_registered():
+def test_add_mul_sigmoid_rms_norm_registered():
     from executorch.backends.webgpu.test.op_tests import cases  # registers
 
-    assert {"add", "rms_norm"} <= set(op_test_registry)
-    assert len(op_test_registry["add"].cases) >= 3  # regular/self/scalar/chained
+    assert {"add", "mul", "sigmoid", "rms_norm"} <= set(op_test_registry)
+    assert len(op_test_registry["add"].cases) >= 3  # regular/self/chained
+    assert len(op_test_registry["mul"].cases) >= len(cases._MUL_CONFIGS)
+    assert len(op_test_registry["sigmoid"].cases) >= 7  # rank/saturation/chained
     # Exact parity, no hardcoded literal (real _CASES == 15; import so it can't drift):
     assert len(op_test_registry["rms_norm"].cases) == len(cases.RMS_NORM_CASES)
     # weight is a construction param, NOT a forward input:
@@ -74,7 +76,9 @@ def test_heavy_forces_not_required():
 def test_golden_dtype_default():
     from executorch.backends.webgpu.test.op_tests import cases  # noqa: F401  registers
 
-    # fp64 oracle is the default; the two landed compute ops keep it. (Per-op golden_dtype
+    # fp64 oracle is the default; these compute ops keep it. (Per-op golden_dtype
     # for gather/copy ops is asserted in each op's own tests-diff.)
     assert op_test_registry["add"].golden_dtype == "float64"
+    assert op_test_registry["mul"].golden_dtype == "float64"
+    assert op_test_registry["sigmoid"].golden_dtype == "float64"
     assert op_test_registry["rms_norm"].golden_dtype == "float64"

backends/webgpu/test/ops/test_mul.py

@@ -9,7 +9,8 @@
 `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
 the export-delegation smoke test. Configs span the same-shape
-fast path (SwiGLU), last-dim broadcast at LLM width, and a mixed-rank left-pad case.
+fast path (SwiGLU), last-dim broadcast at LLM width, same-rank leading-dim
+broadcast, mixed 4D broadcast, and a mixed-rank left-pad case.
 """
 
 import unittest
@@ -23,6 +24,8 @@
 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
+    "bcast_firstdim": ((4, 4), (1, 4)),  # same-rank leading-dim broadcast
+    "bcast_4d_mixed": ((3, 5, 7, 11), (1, 5, 1, 11)),
     "mixedrank": ((4,), (3, 4)),  # right-aligned left-pad (in.ndim < out.ndim)
 }
 

backends/webgpu/test/ops/test_sigmoid.py

@@ -6,10 +6,11 @@
 
 """`aten.sigmoid.default` module + input for the WebGPU op-test framework.
 
-`SigmoidModule`, `N`, and `_det_input` are imported by `cases.py` to drive the
-declarative op-test suite. `SigmoidTest` is the export-delegation
-smoke test. Sigmoid is on the Llama critical path (`F.silu` -> `sigmoid` + `mul`); the
-deterministic input spans the saturation tails.
+`SigmoidModule`, `SigmoidChainedModule`, `N`, `_det_input`, and
+`_wide_det_input` are imported by `cases.py` to drive the declarative op-test
+suite. `SigmoidTest` is the export-delegation smoke test. Sigmoid is on the
+Llama critical path (`F.silu` -> `sigmoid` + `mul`); the deterministic inputs
+span the saturation tails.
 """
 
 import unittest
@@ -28,24 +29,52 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return torch.sigmoid(x)
 
 
+class SigmoidChainedModule(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.sigmoid(torch.sigmoid(x))
+
+
 def _det_input() -> torch.Tensor:
     """Deterministic fp32 input spanning negatives, zero, and large magnitudes."""
     return torch.linspace(-12.0, 12.0, N, dtype=torch.float32)
 
 
-def _export(m: torch.nn.Module, x: torch.Tensor):
+def _wide_det_input() -> torch.Tensor:
+    """Deterministic fp32 input covering stronger saturation tails."""
+    return torch.linspace(-20.0, 20.0, N, dtype=torch.float32)
+
+
+def _lower(m: torch.nn.Module, x: torch.Tensor):
     ep = torch.export.export(m, (x,))
-    return to_edge_transform_and_lower(
-        ep, partitioner=[VulkanPartitioner()]
-    ).to_executorch()
+    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:
+    gm = edge.exported_program().graph_module
+    return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes)
 
 
 class SigmoidTest(unittest.TestCase):
     def test_export_delegates(self) -> None:
-        et = _export(SigmoidModule().eval(), _det_input())
-        found = any(
-            d.id == "VulkanBackend"
-            for plan in et.executorch_program.execution_plan
-            for d in plan.delegates
-        )
-        self.assertTrue(found, "Expected a VulkanBackend delegate (sigmoid)")
+        for name, module, x in (
+            ("regular", SigmoidModule().eval(), _det_input()),
+            ("chained", SigmoidChainedModule().eval(), _det_input()),
+        ):
+            with self.subTest(name=name):
+                edge = _lower(module, x)
+                et = edge.to_executorch()
+                self.assertTrue(
+                    _delegated(et), f"Expected a VulkanBackend delegate ({name})"
+                )
+                self.assertTrue(
+                    _op_delegated(edge, "sigmoid.default"),
+                    f"sigmoid.default not delegated (fell back to CPU) for {name}",
+                )