Skip to content

Commit e51f1a9

Browse files
iamorlandoJCNTH
andauthored
Add WebGPU mul and sigmoid test coverage (#20504)
### Summary Adds WebGPU test coverage for existing `aten.mul.Tensor` and `aten.sigmoid.default` support. This updates the WebGPU Python op tests to cover: - additional `mul` broadcast shapes, including leading-dimension and mixed 4D broadcast cases - wider sigmoid saturation inputs - chained sigmoid delegation - op-test registry/schema coverage for `mul` and `sigmoid` - generator delegation checks so registered `mul` and `sigmoid` cases must lower to the WebGPU delegate path No WebGPU runtime, shader, CMake, README, or TODO files are changed in this PR. ### Test plan - `git diff --check` - `.venv-pyproject/bin/python -m py_compile backends/webgpu/test/op_tests/cases.py backends/webgpu/test/op_tests/test_generator.py backends/webgpu/test/op_tests/test_schema.py backends/webgpu/test/ops/test_mul.py backends/webgpu/test/ops/test_sigmoid.py` - `PATH=/Users/orlando/Documents/dev/executorch/.venv-pyproject/bin:$PATH .venv-pyproject/bin/lintrunner --data-path /tmp/executorch_lintrunner_pyproject backends/webgpu/test/op_tests/cases.py backends/webgpu/test/op_tests/test_generator.py backends/webgpu/test/op_tests/test_schema.py backends/webgpu/test/ops/test_mul.py backends/webgpu/test/ops/test_sigmoid.py` - `PYTHONPYCACHEPREFIX=/tmp/executorch_pyproject_pycache .venv-pyproject/bin/python -m pytest backends/webgpu/test/ops/test_mul.py backends/webgpu/test/ops/test_sigmoid.py backends/webgpu/test/op_tests/test_schema.py backends/webgpu/test/op_tests/test_generator.py -vv` Result: `13 passed in 120.77s` --------- Co-authored-by: Julian Ng-Thow-Hing <107437036+JCNTH@users.noreply.github.com>
1 parent 96c83ed commit e51f1a9

5 files changed

Lines changed: 80 additions & 22 deletions

File tree

backends/webgpu/test/op_tests/cases.py

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,9 @@
5454
)
5555
from executorch.backends.webgpu.test.ops.test_sigmoid import (
5656
_det_input as _sigmoid_det_input,
57+
_wide_det_input as _sigmoid_wide_det_input,
5758
N as _SIGMOID_N,
59+
SigmoidChainedModule,
5860
SigmoidModule,
5961
)
6062

@@ -188,11 +190,22 @@ def _sigmoid_full_range(_shape) -> torch.Tensor:
188190
return _sigmoid_det_input()
189191

190192

193+
def _sigmoid_wide_range(_shape) -> torch.Tensor:
194+
return _sigmoid_wide_det_input()
195+
196+
197+
def _sigmoid_factory(variant: str = "regular") -> torch.nn.Module:
198+
return {
199+
"regular": SigmoidModule,
200+
"chained": SigmoidChainedModule,
201+
}[variant]()
202+
203+
191204
@register_op_test("sigmoid")
192205
def _sigmoid_suite() -> WebGPUTestSuite:
193206
# sigmoid has no CONFIGS table; cover unary shapes directly (tol 1e-4).
194207
return WebGPUTestSuite(
195-
module_factory=lambda: SigmoidModule(),
208+
module_factory=_sigmoid_factory,
196209
cases=[
197210
Case(name="vec", inputs=((M1,),)),
198211
Case(name="mat", inputs=((M1, M2),)),
@@ -203,6 +216,15 @@ def _sigmoid_suite() -> WebGPUTestSuite:
203216
name="saturation",
204217
inputs=(InputSpec(shape=(_SIGMOID_N,), gen=_sigmoid_full_range),),
205218
),
219+
Case(
220+
name="wide_saturation",
221+
inputs=(InputSpec(shape=(_SIGMOID_N,), gen=_sigmoid_wide_range),),
222+
),
223+
Case(
224+
name="chained",
225+
construct={"variant": "chained"},
226+
inputs=(InputSpec(shape=(_SIGMOID_N,), gen=_sigmoid_full_range),),
227+
),
206228
],
207229
atol=1e-4,
208230
rtol=1e-4,

backends/webgpu/test/op_tests/test_generator.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ def test_generate_manifest(tmp_path):
7777
def test_every_case_delegates():
7878
# Contract: every registered case must lower to a VulkanBackend delegate. An op that
7979
# silently CPU-falls-back would otherwise produce a misleading golden-equals-golden pass.
80-
for op in ("add", "rms_norm"):
80+
for op in ("add", "mul", "sigmoid", "rms_norm"):
8181
suite = op_test_registry[op]
8282
for case in suite.cases:
8383
_module, _inputs, prog = g.export_case(suite, case)

backends/webgpu/test/op_tests/test_schema.py

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,11 +36,13 @@ def _dummy():
3636
assert suite.atol == 1e-3 and isinstance(XL, int)
3737

3838

39-
def test_add_rms_norm_registered():
39+
def test_add_mul_sigmoid_rms_norm_registered():
4040
from executorch.backends.webgpu.test.op_tests import cases # registers
4141

42-
assert {"add", "rms_norm"} <= set(op_test_registry)
43-
assert len(op_test_registry["add"].cases) >= 3 # regular/self/scalar/chained
42+
assert {"add", "mul", "sigmoid", "rms_norm"} <= set(op_test_registry)
43+
assert len(op_test_registry["add"].cases) >= 3 # regular/self/chained
44+
assert len(op_test_registry["mul"].cases) >= len(cases._MUL_CONFIGS)
45+
assert len(op_test_registry["sigmoid"].cases) >= 7 # rank/saturation/chained
4446
# Exact parity, no hardcoded literal (real _CASES == 15; import so it can't drift):
4547
assert len(op_test_registry["rms_norm"].cases) == len(cases.RMS_NORM_CASES)
4648
# weight is a construction param, NOT a forward input:
@@ -74,7 +76,9 @@ def test_heavy_forces_not_required():
7476
def test_golden_dtype_default():
7577
from executorch.backends.webgpu.test.op_tests import cases # noqa: F401 registers
7678

77-
# fp64 oracle is the default; the two landed compute ops keep it. (Per-op golden_dtype
79+
# fp64 oracle is the default; these compute ops keep it. (Per-op golden_dtype
7880
# for gather/copy ops is asserted in each op's own tests-diff.)
7981
assert op_test_registry["add"].golden_dtype == "float64"
82+
assert op_test_registry["mul"].golden_dtype == "float64"
83+
assert op_test_registry["sigmoid"].golden_dtype == "float64"
8084
assert op_test_registry["rms_norm"].golden_dtype == "float64"

backends/webgpu/test/ops/test_mul.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,8 @@
99
`MulModule` + `CONFIGS` are imported by `cases.py` to drive the declarative op-test
1010
suite (export via VulkanPartitioner + fp64 torch golden, run on Dawn). `MulTest` is
1111
the export-delegation smoke test. Configs span the same-shape
12-
fast path (SwiGLU), last-dim broadcast at LLM width, and a mixed-rank left-pad case.
12+
fast path (SwiGLU), last-dim broadcast at LLM width, same-rank leading-dim
13+
broadcast, mixed 4D broadcast, and a mixed-rank left-pad case.
1314
"""
1415

1516
import unittest
@@ -23,6 +24,8 @@
2324
CONFIGS = {
2425
"same": ((8, 32), (8, 32)), # fast path (SwiGLU same-shape)
2526
"bcast_lastdim": ((1, 1, 7, 896), (1, 1, 7, 1)), # last-dim broadcast, LLM width
27+
"bcast_firstdim": ((4, 4), (1, 4)), # same-rank leading-dim broadcast
28+
"bcast_4d_mixed": ((3, 5, 7, 11), (1, 5, 1, 11)),
2629
"mixedrank": ((4,), (3, 4)), # right-aligned left-pad (in.ndim < out.ndim)
2730
}
2831

backends/webgpu/test/ops/test_sigmoid.py

Lines changed: 44 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,11 @@
66

77
"""`aten.sigmoid.default` module + input for the WebGPU op-test framework.
88
9-
`SigmoidModule`, `N`, and `_det_input` are imported by `cases.py` to drive the
10-
declarative op-test suite. `SigmoidTest` is the export-delegation
11-
smoke test. Sigmoid is on the Llama critical path (`F.silu` -> `sigmoid` + `mul`); the
12-
deterministic input spans the saturation tails.
9+
`SigmoidModule`, `SigmoidChainedModule`, `N`, `_det_input`, and
10+
`_wide_det_input` are imported by `cases.py` to drive the declarative op-test
11+
suite. `SigmoidTest` is the export-delegation smoke test. Sigmoid is on the
12+
Llama critical path (`F.silu` -> `sigmoid` + `mul`); the deterministic inputs
13+
span the saturation tails.
1314
"""
1415

1516
import unittest
@@ -28,24 +29,52 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
2829
return torch.sigmoid(x)
2930

3031

32+
class SigmoidChainedModule(torch.nn.Module):
33+
def forward(self, x: torch.Tensor) -> torch.Tensor:
34+
return torch.sigmoid(torch.sigmoid(x))
35+
36+
3137
def _det_input() -> torch.Tensor:
3238
"""Deterministic fp32 input spanning negatives, zero, and large magnitudes."""
3339
return torch.linspace(-12.0, 12.0, N, dtype=torch.float32)
3440

3541

36-
def _export(m: torch.nn.Module, x: torch.Tensor):
42+
def _wide_det_input() -> torch.Tensor:
43+
"""Deterministic fp32 input covering stronger saturation tails."""
44+
return torch.linspace(-20.0, 20.0, N, dtype=torch.float32)
45+
46+
47+
def _lower(m: torch.nn.Module, x: torch.Tensor):
3748
ep = torch.export.export(m, (x,))
38-
return to_edge_transform_and_lower(
39-
ep, partitioner=[VulkanPartitioner()]
40-
).to_executorch()
49+
return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()])
50+
51+
52+
def _delegated(et) -> bool:
53+
return any(
54+
d.id == "VulkanBackend"
55+
for plan in et.executorch_program.execution_plan
56+
for d in plan.delegates
57+
)
58+
59+
60+
def _op_delegated(edge, op_substr: str) -> bool:
61+
gm = edge.exported_program().graph_module
62+
return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes)
4163

4264

4365
class SigmoidTest(unittest.TestCase):
4466
def test_export_delegates(self) -> None:
45-
et = _export(SigmoidModule().eval(), _det_input())
46-
found = any(
47-
d.id == "VulkanBackend"
48-
for plan in et.executorch_program.execution_plan
49-
for d in plan.delegates
50-
)
51-
self.assertTrue(found, "Expected a VulkanBackend delegate (sigmoid)")
67+
for name, module, x in (
68+
("regular", SigmoidModule().eval(), _det_input()),
69+
("chained", SigmoidChainedModule().eval(), _det_input()),
70+
):
71+
with self.subTest(name=name):
72+
edge = _lower(module, x)
73+
et = edge.to_executorch()
74+
self.assertTrue(
75+
_delegated(et), f"Expected a VulkanBackend delegate ({name})"
76+
)
77+
self.assertTrue(
78+
_op_delegated(edge, "sigmoid.default"),
79+
f"sigmoid.default not delegated (fell back to CPU) for {name}",
80+
)

0 commit comments

Comments
 (0)