From 3b55bdd11c82f511a5fc7c7be8fe685c5ede9573 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 14 Jul 2026 14:48:22 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/test/ops/test_div.py | 96 ++++++++++++++++++++++ backends/webgpu/test/ops/test_sub.py | 116 +++++++++++++++++++++++++++ 2 files changed, 212 insertions(+) create mode 100644 backends/webgpu/test/ops/test_div.py create mode 100644 backends/webgpu/test/ops/test_sub.py diff --git a/backends/webgpu/test/ops/test_div.py b/backends/webgpu/test/ops/test_div.py new file mode 100644 index 00000000000..912c094e8d2 --- /dev/null +++ b/backends/webgpu/test/ops/test_div.py @@ -0,0 +1,96 @@ +# 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() diff --git a/backends/webgpu/test/ops/test_sub.py b/backends/webgpu/test/ops/test_sub.py new file mode 100644 index 00000000000..5f711fece85 --- /dev/null +++ b/backends/webgpu/test/ops/test_sub.py @@ -0,0 +1,116 @@ +# 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("