diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 30bd86004f7..61afeb32dd7 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -49,6 +49,16 @@ N as _SIGMOID_N, SigmoidModule, ) + +from executorch.backends.webgpu.test.ops.squeeze.test_squeeze import ( + CONFIGS as _SQUEEZE_CONFIGS, + SqueezeModule, +) + +from executorch.backends.webgpu.test.ops.unsqueeze.test_unsqueeze import ( + CONFIGS as _UNSQUEEZE_CONFIGS, + UnsqueezeModule, +) from executorch.backends.webgpu.test.ops.view_copy.test_view_copy import ( CONFIGS as _VIEW_CONFIGS, ViewModule, @@ -184,3 +194,29 @@ def _sigmoid_suite() -> WebGPUTestSuite: atol=1e-4, rtol=1e-4, ) + + +@register_op_test("squeeze") +def _squeeze_suite() -> WebGPUTestSuite: + # CONFIGS: name -> (shape, dim) where dim is an int or a tuple. + return WebGPUTestSuite( + module_factory=lambda dim: SqueezeModule(dim), + cases=[ + Case(name=n, construct={"dim": dim}, inputs=(shape,)) + for n, (shape, dim) in _SQUEEZE_CONFIGS.items() + ], + golden_dtype="float32", # reshape copies values; fp64 bit-identical + ) + + +@register_op_test("unsqueeze") +def _unsqueeze_suite() -> WebGPUTestSuite: + # CONFIGS: name -> (shape, dim). + return WebGPUTestSuite( + module_factory=lambda dim: UnsqueezeModule(dim), + cases=[ + Case(name=n, construct={"dim": dim}, inputs=(shape,)) + for n, (shape, dim) in _UNSQUEEZE_CONFIGS.items() + ], + golden_dtype="float32", # reshape copies values; fp64 bit-identical + ) diff --git a/backends/webgpu/test/ops/squeeze/__init__.py b/backends/webgpu/test/ops/squeeze/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/webgpu/test/ops/squeeze/__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. diff --git a/backends/webgpu/test/ops/squeeze/test_squeeze.py b/backends/webgpu/test/ops/squeeze/test_squeeze.py new file mode 100644 index 00000000000..ad129744ffb --- /dev/null +++ b/backends/webgpu/test/ops/squeeze/test_squeeze.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.squeeze_copy.dims` module + configs for the WebGPU op-test framework. + +`SqueezeModule` + `CONFIGS` are imported by `cases.py` to drive the declarative +op-test suite. `SqueezeTest` is the export-delegation smoke +test. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# name -> (input_shape, squeeze_dim) +CONFIGS = { + "dim0": ((1, 3, 4), 0), + "mid": ((2, 1, 4), 1), + "multi": ((1, 3, 1, 4), (0, 2)), +} + + +class SqueezeModule(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.squeeze(x, self.dim) + + +def _det_input(shape): + g = torch.Generator().manual_seed(0) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(dim, x: torch.Tensor): + ep = torch.export.export(SqueezeModule(dim).eval(), (x,)) + 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: + # op must be absorbed into the delegate, not left as a CPU-fallback node. + gm = edge.exported_program().graph_module + return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + + +class SqueezeTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, (shape, dim) in CONFIGS.items(): + edge = _lower(dim, _det_input(shape)) + et = edge.to_executorch() + self.assertTrue( + _delegated(et), + f"Expected a VulkanBackend delegate (squeeze {name})", + ) + self.assertTrue( + _op_delegated(edge, "squeeze_copy"), + f"squeeze_copy not delegated (fell back to CPU) for {name}", + ) diff --git a/backends/webgpu/test/ops/unsqueeze/__init__.py b/backends/webgpu/test/ops/unsqueeze/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/webgpu/test/ops/unsqueeze/__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. diff --git a/backends/webgpu/test/ops/unsqueeze/test_unsqueeze.py b/backends/webgpu/test/ops/unsqueeze/test_unsqueeze.py new file mode 100644 index 00000000000..dada77df00e --- /dev/null +++ b/backends/webgpu/test/ops/unsqueeze/test_unsqueeze.py @@ -0,0 +1,73 @@ +# 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.unsqueeze_copy.default` module + configs for the WebGPU op-test framework. + +`UnsqueezeModule` + `CONFIGS` are imported by `cases.py` to drive the declarative +op-test suite. `UnsqueezeTest` is the export-delegation smoke +test. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# name -> (input_shape, unsqueeze_dim) +CONFIGS = { + "front": ((3, 4), 0), + "mid": ((2, 4), 1), + "last": ((3, 4), 2), +} + + +class UnsqueezeModule(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.unsqueeze(x, self.dim) + + +def _det_input(shape): + g = torch.Generator().manual_seed(0) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(dim, x: torch.Tensor): + ep = torch.export.export(UnsqueezeModule(dim).eval(), (x,)) + 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: + # op must be absorbed into the delegate, not left as a top-level CPU-fallback node. + gm = edge.exported_program().graph_module + return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + + +class UnsqueezeTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, (shape, dim) in CONFIGS.items(): + edge = _lower(dim, _det_input(shape)) + et = edge.to_executorch() + self.assertTrue( + _delegated(et), f"Expected a VulkanBackend delegate (unsqueeze {name})" + ) + self.assertTrue( + _op_delegated(edge, "unsqueeze_copy"), + f"unsqueeze_copy not delegated (fell back to CPU) for {name}", + ) diff --git a/backends/webgpu/test/tester.py b/backends/webgpu/test/tester.py index e5dd510d49b..53a745a16df 100644 --- a/backends/webgpu/test/tester.py +++ b/backends/webgpu/test/tester.py @@ -25,6 +25,8 @@ exir_ops.edge.aten.view_copy.default, exir_ops.edge.aten.select_copy.int, exir_ops.edge.aten.sigmoid.default, + exir_ops.edge.aten.squeeze_copy.dims, + exir_ops.edge.aten.unsqueeze_copy.default, ]