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

backends/webgpu/test/ops/test_div.py

@@ -0,0 +1,94 @@
+# 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.div.Tensor` (broadcast) export + fp64 golden for the WebGPU backend.
+
+Exports single-op divide graphs through VulkanPartitioner and asserts they
+delegate to the Vulkan backend (div is absent from the top-level portable ops),
+then locks the golden math against an fp64 torch reference (`a / b` with PyTorch
+broadcasting). Configs span the same-shape fast path, a last-dim broadcast at
+LLM width, and a mixed-rank left-pad case. Divisors are bounded away from zero
+so the fp32-vs-fp64 comparison stays well-conditioned.
+"""
+
+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 -> (shape_a, shape_b). Output shape is the broadcast of the two.
+CONFIGS = {
+    "same": ((8, 32), (8, 32)),  # fast path (same-shape elementwise)
+    "bcast_lastdim": ((1, 1, 7, 896), (1, 1, 7, 1)),  # last-dim broadcast, LLM
+    "mixedrank": ((4,), (3, 4)),  # right-aligned left-pad (in.ndim < out.ndim)
+}
+
+
+class DivModule(torch.nn.Module):
+    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        return a / b
+
+
+def _det_inputs(shape_a, shape_b):
+    """Deterministic fp32 inputs (fixed seed); divisor bounded away from zero."""
+    g = torch.Generator().manual_seed(0)
+    a = torch.randn(*shape_a, generator=g, dtype=torch.float32)
+    b = torch.randn(*shape_b, generator=g, dtype=torch.float32).abs() + 0.5
+    return a, b
+
+
+def _export(a: torch.Tensor, b: torch.Tensor):
+    ep = torch.export.export(DivModule().eval(), (a, b))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegated(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 TestDiv(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, (sa, sb) in CONFIGS.items():
+            with self.subTest(name=name):
+                a, b = _det_inputs(sa, sb)
+                et = _export(a, b)
+                self.assertTrue(
+                    _delegated(et), f"Expected a VulkanBackend delegate (div {name})"
+                )
+                self.assertFalse(
+                    any("div" in n for n in _top_level_op_names(et)),
+                    f"div should be delegated, not a top-level portable op (div {name})",
+                )
+
+    def test_op_matches_fp64_golden(self) -> None:
+        for name, (sa, sb) in CONFIGS.items():
+            with self.subTest(name=name):
+                a, b = _det_inputs(sa, sb)
+                got = DivModule()(a, b)
+                golden = (a.double() / b.double()).to(torch.float32)
+                torch.testing.assert_close(got, golden, atol=5e-4, rtol=1e-3)
+
+
+if __name__ == "__main__":
+    unittest.main()

backends/webgpu/test/ops/test_sub.py

@@ -0,0 +1,114 @@
+# 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.sub.Tensor` export + golden for the WebGPU backend.
+
+Exports single-op subtraction graphs through VulkanPartitioner (the WebGPU runtime
+consumes the Vulkan VK00 delegate directly) and checks an fp64 torch golden
+(`out = in1 - alpha * in2`). The native/etvk numeric oracle compares the GPU kernel
+against this same reference; this suite locks delegation + the reference math.
+Configs span same-shape 2D/3D, trailing- and leading-dim broadcast, and alpha != 1.
+"""
+
+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 -> (shape_a, shape_b); b broadcasts into a.
+CONFIGS = {
+    "2d": ((4, 4), (4, 4)),
+    "3d": ((2, 3, 4), (2, 3, 4)),
+    "bcast_last": ((4, 4), (4, 1)),
+    "bcast_row": ((4, 4), (1, 4)),
+}
+
+
+class SubModule(torch.nn.Module):
+    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        return torch.sub(a, b)
+
+
+class SubAlphaModule(torch.nn.Module):
+    def __init__(self, alpha: float) -> None:
+        super().__init__()
+        self.alpha = alpha
+
+    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        return torch.sub(a, b, alpha=self.alpha)
+
+
+def _det_inputs(shape_a, shape_b):
+    """Deterministic fp32 inputs (fixed seed) for a config."""
+    g = torch.Generator().manual_seed(0)
+    a = torch.randn(*shape_a, generator=g, dtype=torch.float32)
+    b = torch.randn(*shape_b, generator=g, dtype=torch.float32)
+    return a, b
+
+
+def _export(m: torch.nn.Module, a: torch.Tensor, b: torch.Tensor):
+    ep = torch.export.export(m.eval(), (a, b))
+    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 TestSub(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        # Delegation => no aten.sub.Tensor left in the top-level portable graph.
+        for name, (sa, sb) in CONFIGS.items():
+            with self.subTest(name=name):
+                a, b = _det_inputs(sa, sb)
+                et = _export(SubModule(), a, b)
+                self.assertTrue(
+                    _delegates(et),
+                    f"Expected a VulkanBackend delegate (sub {name})",
+                )
+
+    def test_golden_matches_fp64(self) -> None:
+        for name, (sa, sb) in CONFIGS.items():
+            with self.subTest(name=name):
+                a, b = _det_inputs(sa, sb)
+                ref = (a.double() - b.double()).to(torch.float32)
+                torch.testing.assert_close(SubModule()(a, b), ref)
+
+    def test_golden_matches_fp64_alpha(self) -> None:
+        # Locks the handler's out = in1 - alpha * in2 path (alpha != 1).
+        alpha = 2.5
+        a, b = _det_inputs((4, 4), (4, 4))
+        ref = (a.double() - alpha * b.double()).to(torch.float32)
+        torch.testing.assert_close(SubAlphaModule(alpha)(a, b), ref)
+
+
+def export_sub_model(pte_path: str, golden_path: str, input_path: str) -> None:
+    """Write sub(a, b) .pte + fp64-computed torch golden + raw LE fp32 inputs (in1, in2)."""
+    a, b = _det_inputs((1024, 1024), (1024, 1024))
+    golden = (a.double() - b.double()).to(torch.float32).numpy().astype("<f4")
+    et = _export(SubModule(), a, b)
+    with open(pte_path, "wb") as f:
+        f.write(et.buffer)
+    golden.tofile(golden_path)
+    with open(input_path, "wb") as f:
+        a.numpy().astype("<f4").tofile(f)
+        b.numpy().astype("<f4").tofile(f)
+    print(f"Exported {pte_path}; golden {golden_path} ({golden.size} floats)")
+
+
+if __name__ == "__main__":
+    unittest.main()