[ExecuTorch][WebGPU] Tests for the shape_copy WebGPU ops

JCNTH · web-flow · commit 568dec107b12 · 2026-07-21T16:23:01.000Z
Differential Revision: D111755126 Pull Request resolved: #20930
diff --git a/backends/webgpu/test/ops/test_dim_order.py b/backends/webgpu/test/ops/test_dim_order.py
@@ -0,0 +1,95 @@
+# 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
@@ -0,0 +1,120 @@
+# 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("<f4")
+    et = _export(m, x)
+    with open(pte_path, "wb") as f:
+        f.write(et.buffer)
+    golden.tofile(golden_path)
+    x.numpy().astype("<f4").tofile(input_path)
+    print(
+        f"Exported {pte_path}; golden {golden_path} ({golden.size} floats); "
+        f"input {input_path} ({x.numel()} floats)"
+    )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_fill.py b/backends/webgpu/test/ops/test_fill.py
@@ -0,0 +1,91 @@
+# 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.full` / `aten.full_like` export + fp64 golden for the WebGPU backend.
+
+Both lower to the `fill` handler, which writes a constant scalar into the
+pre-allocated output buffer. `full` and `full_like` are on the training-backward
+critical path (gradient seeds). To keep the op in the graph (a constant-shaped
+`full` would fold away), the fill size is taken from a live input: `full` uses
+`x.shape`, `full_like` uses `x`. Configs span 1D/2D/4D and a non-multiple-of-256
+shape that exercises the dispatch bounds guard.
+"""
+
+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 -> (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()
