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

backends/webgpu/test/ops/test_compare.py

@@ -0,0 +1,106 @@
+# 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}.Scalar` export + golden for the WebGPU backend.
+
+Each scalar comparison lowers to a single `aten.<op>.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")
+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,
+}
+
+
+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
+        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()

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("<u1")
+    et = _lower(m, x).to_executorch()
+    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} bytes); "
+        f"input {input_path} ({x.numel()} floats)"
+    )
+
+
+if __name__ == "__main__":
+    unittest.main()

backends/webgpu/test/ops/test_where.py

@@ -0,0 +1,86 @@
+# 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.where.self` export + fp64 golden for the WebGPU backend.
+
+`where(cond, a, b) -> 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<u32>` 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()