diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 1b6b0fa7085..0216350a443 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/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, diff --git a/backends/webgpu/test/op_tests/test_generator.py b/backends/webgpu/test/op_tests/test_generator.py index deae6ecd4a6..9782d6fc1dc 100644 --- a/backends/webgpu/test/op_tests/test_generator.py +++ b/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) diff --git a/backends/webgpu/test/op_tests/test_schema.py b/backends/webgpu/test/op_tests/test_schema.py index bcc03a40fd9..c5c14b71754 100644 --- a/backends/webgpu/test/op_tests/test_schema.py +++ b/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" diff --git a/backends/webgpu/test/ops/test_mul.py b/backends/webgpu/test/ops/test_mul.py index cca04bcdaa4..36d5429c25a 100644 --- a/backends/webgpu/test/ops/test_mul.py +++ b/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) } diff --git a/backends/webgpu/test/ops/test_sigmoid.py b/backends/webgpu/test/ops/test_sigmoid.py index 0ba8c435a9a..b3cc4b8a165 100644 --- a/backends/webgpu/test/ops/test_sigmoid.py +++ b/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}", + )