diff --git a/backends/webgpu/test/ops/test_compare.py b/backends/webgpu/test/ops/test_compare.py new file mode 100644 index 00000000000..172c5e75113 --- /dev/null +++ b/backends/webgpu/test/ops/test_compare.py @@ -0,0 +1,109 @@ +# 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.{eq,ne,le,ge,lt,gt}.Scalar` export + golden for the WebGPU backend. + +Each scalar comparison lowers to a single `aten..Scalar` node that the +kernel computes as `cmp(self[i], scalar)` and writes as a byte-packed bool. The +delegation test locks that every variant partitions to `VulkanBackend`; the +golden test locks the fp32 module output against the fp64 torch truth. The +deterministic ramp is exact in fp32 and straddles (and hits) the scalar, so both +precisions agree bit-for-bit and every op has mixed True/False cases; a +non-multiple-of-4 numel exercises the kernel tail (4 elems per u32 word). +""" + +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 + +SCALAR = 0.0 +OPS = ("eq", "ne", "le", "ge", "lt", "gt") +SHAPES = {"tail": (3, 5), "3d": (2, 4, 8)} + +_TORCH_OP = { + "eq": torch.eq, + "ne": torch.ne, + "le": torch.le, + "ge": torch.ge, + "lt": torch.lt, + "gt": torch.gt, +} + + +class CompareModule(torch.nn.Module): + def __init__(self, op: str, scalar: float) -> None: + super().__init__() + self.op = op + self.scalar = scalar + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.op == "eq": + return x == self.scalar + if self.op == "ne": + return x != self.scalar + if self.op == "le": + return x <= self.scalar + if self.op == "ge": + return x >= self.scalar + if self.op == "lt": + return x < self.scalar + return x > self.scalar + + +def _det_input(shape: tuple[int, ...]) -> torch.Tensor: + """Deterministic fp32 spanning [-0.5, 0.5] in 1/16 steps (exact in fp32, + hits and straddles 0.0).""" + n = 1 + for d in shape: + n *= d + flat = torch.arange(n, dtype=torch.float32) + return (((flat % 17) - 8) / 16.0).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 _delegated(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +class TestCompare(unittest.TestCase): + def test_export_delegates(self) -> None: + for op in OPS: + for name, shape in SHAPES.items(): + with self.subTest(op=op, shape=name): + x = _det_input(shape) + et = _export(CompareModule(op, SCALAR).eval(), x) + self.assertTrue( + _delegated(et), + f"Expected a VulkanBackend delegate ({op}.Scalar {name})", + ) + + def test_module_matches_fp64_golden(self) -> None: + for op in OPS: + for name, shape in SHAPES.items(): + with self.subTest(op=op, shape=name): + x = _det_input(shape) + got = CompareModule(op, SCALAR)(x) + golden = _TORCH_OP[op](x.double(), float(SCALAR)) + torch.testing.assert_close(got, golden) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/webgpu/test/ops/test_logical_not.py b/backends/webgpu/test/ops/test_logical_not.py new file mode 100644 index 00000000000..6e9c15eaa08 --- /dev/null +++ b/backends/webgpu/test/ops/test_logical_not.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.logical_not.default` export + golden for the WebGPU backend. + +Exports single-op logical_not graphs through VulkanPartitioner and writes a +torch-computed golden. logical_not lowers to a byte-packed bool kernel (1 byte / +elem, one u32 word / thread), so the golden is an EXACT bool match -- the native +test compares the raw uint8 output byte-for-byte. The bool input is produced by a +delegated `x >= threshold` compare; the golden uses the De Morgan inverse +`x < threshold` as an independent reference. A non-multiple-of-4 numel exercises +the shader's partial-final-word masking (the `i < num_elements` branch). +""" + +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. "tail3x7" has numel 21 (not a multiple of 4) => partial word. +CONFIGS = { + "vec1d": (16,), + "mat2d": (4, 8), + "tail3x7": (3, 7), + "cube3d": (2, 4, 8), +} + + +class LogicalNotModule(torch.nn.Module): + def __init__(self, threshold: float) -> None: + super().__init__() + self.threshold = threshold + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.logical_not(x >= self.threshold) + + +def _det_input(shape: tuple[int, ...]) -> torch.Tensor: + """Deterministic fp32 spanning negatives, zero, and positives so the bool + input to logical_not carries both True and False values.""" + g = torch.Generator().manual_seed(0) + return torch.randn(*shape, generator=g, dtype=torch.float32) + + +def _lower(m: torch.nn.Module, x: torch.Tensor): + ep = torch.export.export(m, (x,)) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + +def _delegates(edge) -> bool: + et = edge.to_executorch() + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +def _top_level_targets(edge) -> list[str]: + return [ + str(n.target) + for n in edge.exported_program().graph_module.graph.nodes + if n.op == "call_function" + ] + + +class TestLogicalNot(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, shape in CONFIGS.items(): + with self.subTest(config=name): + edge = _lower(LogicalNotModule(0.0).eval(), _det_input(shape)) + targets = _top_level_targets(edge) + self.assertFalse( + any("logical_not" in t for t in targets), + f"logical_not must be delegated, not portable ({name}): {targets}", + ) + self.assertTrue( + _delegates(edge), + f"Expected a VulkanBackend delegate (logical_not {name})", + ) + + def test_golden_matches_eager(self) -> None: + for name, shape in CONFIGS.items(): + with self.subTest(config=name): + x = _det_input(shape) + got = LogicalNotModule(0.0)(x) + golden = x < 0.0 # De Morgan inverse of logical_not(x >= 0) + self.assertEqual(got.dtype, torch.bool) + torch.testing.assert_close(got, golden) + + +def export_logical_not_model(pte_path: str, golden_path: str, input_path: str) -> None: + """Write logical_not .pte + torch golden (raw uint8) + raw LE fp32 input.""" + m = LogicalNotModule(0.0).eval() + x = _det_input(CONFIGS["mat2d"]) + golden = m(x).detach().numpy().astype(" cond ? a : b`, with cond a 1-byte bool and a/b fp32 +(broadcast across all three operands). The kernel reads cond byte-packed as +`array` and relinearizes each out coord onto every operand. Configs cover +the equal-shape path plus broadcasts that exercise the size-1 clamp on cond, a, +and b. The native binary has no ATen, so the golden is computed with torch here +and checked in etvk CI. +""" + +import unittest + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# name -> (cond_shape, a_shape, b_shape). Output is the broadcast of all three. +CONFIGS = { + "equal": ((4, 8), (4, 8), (4, 8)), + "broadcast": ((4, 1), (4, 8), (1, 8)), + "cond_row": ((8,), (4, 8), (4, 8)), +} + + +class WhereModule(torch.nn.Module): + def forward( + self, cond: torch.Tensor, a: torch.Tensor, b: torch.Tensor + ) -> torch.Tensor: + return torch.where(cond, a, b) + + +def _det_inputs(cond_shape, a_shape, b_shape): + """Deterministic (bool cond, fp32 a, fp32 b) for a config.""" + g = torch.Generator().manual_seed(0) + cond = torch.rand(cond_shape, generator=g) > 0.5 + a = torch.randn(*a_shape, generator=g, dtype=torch.float32) + b = torch.randn(*b_shape, generator=g, dtype=torch.float32) + return cond, a, b + + +def _fp64_golden(cond, a, b): + return torch.where(cond, a.double(), b.double()).to(torch.float32) + + +def _export(cond, a, b): + ep = torch.export.export(WhereModule().eval(), (cond, 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 TestWhere(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, (cs, as_, bs) in CONFIGS.items(): + with self.subTest(config=name): + cond, a, b = _det_inputs(cs, as_, bs) + et = _export(cond, a, b) + self.assertTrue( + _delegates(et), f"Expected a VulkanBackend delegate (where {name})" + ) + + def test_op_matches_fp64_golden(self) -> None: + for name, (cs, as_, bs) in CONFIGS.items(): + with self.subTest(config=name): + cond, a, b = _det_inputs(cs, as_, bs) + out = WhereModule()(cond, a, b) + torch.testing.assert_close(out, _fp64_golden(cond, a, b)) + + +if __name__ == "__main__": + unittest.main()