From 0ce99c2d93bf4ab4ece361990c1507df39f7ae85 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 14 Jul 2026 14:49:04 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/test/ops/test_dim_order.py | 97 +++++++++++++++ backends/webgpu/test/ops/test_expand_copy.py | 122 +++++++++++++++++++ backends/webgpu/test/ops/test_fill.py | 93 ++++++++++++++ 3 files changed, 312 insertions(+) create mode 100644 backends/webgpu/test/ops/test_dim_order.py create mode 100644 backends/webgpu/test/ops/test_expand_copy.py create mode 100644 backends/webgpu/test/ops/test_fill.py diff --git a/backends/webgpu/test/ops/test_dim_order.py b/backends/webgpu/test/ops/test_dim_order.py new file mode 100644 index 00000000000..bb032802bcc --- /dev/null +++ b/backends/webgpu/test/ops/test_dim_order.py @@ -0,0 +1,97 @@ +# 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. + +"""`dim_order_ops._clone_dim_order.default` export delegation for the WebGPU backend. + +`torch.clone(x)` lowers to `dim_order_ops._clone_dim_order.default` (the edge +dialect's dim-order-aware clone). On the buffer-only WebGPU backend it is a +numel-preserving flat copy handled by the shared `add_flat_copy` DMA helper (no +WGSL). The partitioner tags it (single-node partitions are allowed), so a +`VulkanBackend` delegate is formed; `RemoveRedundantOpsTransform` then folds the +identity clone out of the delegate in preprocess, so it never reaches the native +runtime (hence no `.golden.bin` / native sweep — like `clone`). + +`test_export_delegates` locks the contract that the op is absorbed into the +delegate and never survives as a top-level portable (CPU-fallback) node; +`test_golden_matches_eager` locks the fp64 `x.clone()` reference. Configs cover +1D, 3D, and 4D contiguous inputs. +""" + +from __future__ import annotations + +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 +CONFIGS = { + "flat": (16,), + "3d": (2, 3, 4), + "4d": (1, 3, 4, 5), +} + + +class CloneDimOrderModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.clone(x) + + +def _det_input(shape): + g = torch.Generator().manual_seed(0) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(x: torch.Tensor): + ep = torch.export.export(CloneDimOrderModule().eval(), (x,)) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + +def _delegates(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +def _op_absent_from_toplevel(edge, op_substr: str) -> bool: + # Delegated/folded ops are absorbed; none may survive as a top-level node. + gm = edge.exported_program().graph_module + return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + + +class TestCloneDimOrder(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, shape in CONFIGS.items(): + with self.subTest(name=name): + edge = _lower(_det_input(shape)) + et = edge.to_executorch() + self.assertTrue( + _delegates(et), + f"Expected a VulkanBackend delegate (clone_dim_order {name})", + ) + self.assertTrue( + _op_absent_from_toplevel(edge, "_clone_dim_order"), + f"_clone_dim_order left as a top-level portable op for {name}", + ) + + def test_golden_matches_eager(self) -> None: + for name, shape in CONFIGS.items(): + with self.subTest(name=name): + x = _det_input(shape) + ref = x.to(torch.float64).clone() + torch.testing.assert_close( + CloneDimOrderModule()(x).to(torch.float64), ref + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/webgpu/test/ops/test_expand_copy.py b/backends/webgpu/test/ops/test_expand_copy.py new file mode 100644 index 00000000000..787230314f3 --- /dev/null +++ b/backends/webgpu/test/ops/test_expand_copy.py @@ -0,0 +1,122 @@ +# 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.expand_copy.default` export + golden for the WebGPU backend. + +Exports single-op expand-copy graphs through VulkanPartitioner and checks an fp64 +torch golden. expand_copy materializes a broadcasted view (size-1 input dims, and +rank-increasing leading dims) into the target shape via a pure gather -- no +arithmetic, so fp64 and fp32 agree exactly. Configs cover a broadcast leading +dim, a broadcast middle dim, and a rank increase (input rank < output rank) that +exercises the kernel's right-alignment of input dims into the output rank. +""" + +from __future__ import annotations + +import math +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, expanded shape). +CONFIGS = { + "broadcast_leading": ((1, 4), (3, 4)), + "broadcast_middle": ((2, 1, 5), (2, 4, 5)), + "rank_increase": ((4,), (3, 4)), +} + + +class ExpandCopyModule(torch.nn.Module): + def __init__(self, shape: tuple[int, ...]) -> None: + super().__init__() + self.shape = shape + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.expand(self.shape).clone() + + +def _det_input(shape: tuple[int, ...]) -> torch.Tensor: + """Deterministic fp32 ramp; a broadcast dim repeats it so the copy is visible.""" + return torch.arange(math.prod(shape), dtype=torch.float32).reshape(shape) + + +def _export(m: torch.nn.Module, x: torch.Tensor): + ep = torch.export.export(m, (x,)) + return to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + + +def _delegates(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +def _top_level_op_names(et) -> set[str]: + return { + op.name + for plan in et.executorch_program.execution_plan + for op in plan.operators + } + + +class TestExpandCopy(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, (in_shape, out_shape) in CONFIGS.items(): + with self.subTest(name=name): + x = _det_input(in_shape) + et = _export(ExpandCopyModule(out_shape).eval(), x) + self.assertTrue( + _delegates(et), + f"Expected a VulkanBackend delegate (expand_copy {name})", + ) + # Delegated => expand_copy absent from top-level portable ops. + self.assertFalse( + any("expand_copy" in n for n in _top_level_op_names(et)), + f"expand_copy leaked into top-level ops ({name})", + ) + + def test_golden_matches_torch(self) -> None: + for name, (in_shape, out_shape) in CONFIGS.items(): + with self.subTest(name=name): + x = _det_input(in_shape) + golden = x.double().expand(out_shape).clone() + got = ExpandCopyModule(out_shape)(x) + torch.testing.assert_close(got.double(), golden) + + +def export_expand_copy_model( + out_shape: tuple[int, ...], + in_shape: tuple[int, ...], + pte_path: str, + golden_path: str, + input_path: str, +) -> None: + """Write an expand_copy .pte + torch golden (raw LE fp32) + raw LE fp32 input.""" + m = ExpandCopyModule(out_shape).eval() + x = _det_input(in_shape) + golden = m(x).detach().numpy().astype(" (kind, shape, fill_value) +CONFIGS = { + "full_1d": ("full", (37,), 0.0), # non-multiple of 256: bounds-guard path + "full_2d": ("full", (4, 8), 3.0), + "full_4d": ("full", (2, 3, 4, 5), -1.5), + "full_like_2d": ("full_like", (16, 16), 7.25), +} + + +class FullModule(torch.nn.Module): + def __init__(self, val: float) -> None: + super().__init__() + self.val = val + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.full(x.shape, self.val) + + +class FullLikeModule(torch.nn.Module): + def __init__(self, val: float) -> None: + super().__init__() + self.val = val + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.full_like(x, self.val) + + +def _module(kind: str, val: float) -> torch.nn.Module: + return (FullModule(val) if kind == "full" else FullLikeModule(val)).eval() + + +def _export(m: torch.nn.Module, x: torch.Tensor): + ep = torch.export.export(m, (x,)) + return to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + + +def _delegates(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +class TestFill(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, (kind, shape, val) in CONFIGS.items(): + with self.subTest(name=name): + x = torch.zeros(shape, dtype=torch.float32) + et = _export(_module(kind, val), x) + self.assertTrue( + _delegates(et), f"Expected a VulkanBackend delegate (fill {name})" + ) + + def test_golden_matches_eager(self) -> None: + for name, (kind, shape, val) in CONFIGS.items(): + with self.subTest(name=name): + x = torch.zeros(shape, dtype=torch.float32) + golden = torch.full(shape, val, dtype=torch.float64) + torch.testing.assert_close(_module(kind, val)(x).double(), golden) + + +if __name__ == "__main__": + unittest.main()