[ExecuTorch][WebGPU] Test byte-packed BOOL tails

backends/webgpu/test/op_tests/cases.py

@@ -153,8 +153,13 @@
 )
 
 from executorch.backends.webgpu.test.ops.test_to_copy import (
+    bool_tail_input,
+    compare_to_copy_input_a,
+    compare_to_copy_input_b,
+    CompareToCopyBoolToFloatModule,
     to_copy_float_input,
     to_copy_int_input,
+    ToCopyBoolToFloatModule,
     ToCopyFloatToIntToFloatModule,
     ToCopyIntToFloatModule,
 )
@@ -191,6 +196,39 @@ def _add_factory(variant: str = "regular") -> torch.nn.Module:
     }[variant]()
 
 
+@register_op_test("to_copy_bool_to_float")
+def _to_copy_bool_to_float_suite() -> WebGPUTestSuite:
+    return WebGPUTestSuite(
+        module_factory=CompareToCopyBoolToFloatModule,
+        cases=[
+            Case(
+                inputs=(
+                    InputSpec((n,), gen=compare_to_copy_input_a),
+                    InputSpec((n,), gen=compare_to_copy_input_b),
+                ),
+                name=f"length_{n}",
+            )
+            for n in (1, 4, 5, 67)
+        ],
+        golden_dtype="float32",
+    )
+
+
+@register_op_test("to_copy_bool_input_to_float")
+def _to_copy_bool_input_to_float_suite() -> WebGPUTestSuite:
+    return WebGPUTestSuite(
+        module_factory=ToCopyBoolToFloatModule,
+        cases=[
+            Case(
+                inputs=(InputSpec((n,), gen=bool_tail_input),),
+                name=f"length_{n}",
+            )
+            for n in (1, 4, 5, 67)
+        ],
+        golden_dtype="float32",
+    )
+
+
 @register_op_test("add")
 def _add_suite() -> WebGPUTestSuite:
     # Same-shape numeric coverage only: broadcast adds stay export-smoke in
@@ -295,10 +333,7 @@ def _minimum_suite() -> WebGPUTestSuite:
 
 
 def _compare_suite(op: str) -> WebGPUTestSuite:
-    # Elementwise fp32 comparison -> bool (byte-exact golden). The two inputs use
-    # DIFFERENT discrete-range seeds so a!=b (real lt/gt mix) while colliding
-    # often (eq/le/ge ties); all shapes have numel % 4 == 0 (bool output packs 4
-    # bytes/word). Same-shape only (flat kernel; broadcast=smoke).
+    # Distinct inputs and tail shapes cover byte-exact packed BOOL output.
     def case(name, shape):
         return Case(
             name=name,
@@ -310,7 +345,14 @@ def case(name, shape):
 
     return WebGPUTestSuite(
         module_factory=lambda: CompareModule(op),
-        cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))],
+        cases=[
+            case("tail_1", (1,)),
+            case("tail_5", (5,)),
+            case("tail_67", (67,)),
+            case("2d", (4, 8)),
+            case("3d", (2, 3, 8)),
+            case("sq", (16, 16)),
+        ],
         golden_dtype="bool",
     )
 

backends/webgpu/test/op_tests/generate_op_tests.py

@@ -39,6 +39,8 @@ def _materialize(spec) -> torch.Tensor:
         shape, gen = spec, "randn"
     if callable(gen):
         _t = gen(shape)
+        if _t.dtype == torch.bool:
+            return _t
         return (
             _t.to(torch.int32) if not _t.is_floating_point() else _t.to(torch.float32)
         )
@@ -169,7 +171,10 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> list[d
     input_entries: list[dict] = []
     for i, t in enumerate(inputs):
         rel = f"{case_id}.in{i}.bin"
-        if t.dtype == torch.int32:
+        if t.dtype == torch.bool:
+            _write_int8(t.to(torch.int8), os.path.join(out_dir, rel))
+            in_dtype = "bool"
+        elif t.dtype == torch.int32:
             t.detach().cpu().numpy().astype("<i4").tofile(os.path.join(out_dir, rel))
             in_dtype = "int32"
         else:

backends/webgpu/test/op_tests/op_test_driver.cpp

@@ -16,6 +16,7 @@
 
 #include <gtest/gtest.h>
 
+#include <cstdint>
 #include <cstdio>
 #include <cstdlib>
 #include <exception>
@@ -61,7 +62,15 @@ class OpCase : public ::testing::Test {
       const size_t n = numel(in.shape);
       std::vector<executorch::aten::SizesType> sizes(
           in.shape.begin(), in.shape.end());
-      if (in.dtype == "int32") {
+      if (in.dtype == "bool") {
+        auto data = load_int8_bin(in.path, n);
+        ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path;
+        std::vector<uint8_t> raw(data.begin(), data.end());
+        tensors.push_back(make_tensor_ptr(
+            std::move(sizes),
+            std::move(raw),
+            executorch::aten::ScalarType::Bool));
+      } else if (in.dtype == "int32") {
         auto data = load_int32_bin(in.path, n);
         ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path;
         tensors.push_back(make_tensor_ptr(std::move(sizes), std::move(data)));
@@ -96,10 +105,11 @@ class OpCase : public ::testing::Test {
       auto golden = load_int8_bin(e_.golden.path, gn);
       ASSERT_FALSE(golden.empty())
           << "missing/short golden: " << e_.golden.path;
-      const bool* out_p = out_tensor.const_data_ptr<bool>();
+      ASSERT_EQ(out_tensor.scalar_type(), executorch::aten::ScalarType::Bool);
+      const uint8_t* out_p = out_tensor.const_data_ptr<uint8_t>();
       int mism = -1;
       for (size_t i = 0; i < gn; i++) {
-        if (static_cast<int8_t>(out_p[i]) != golden[i]) {
+        if (out_p[i] != static_cast<uint8_t>(golden[i])) {
           mism = static_cast<int>(i);
           break;
         }

backends/webgpu/test/ops/test_to_copy.py

@@ -46,6 +46,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x.to(torch.float32, copy=True)
 
 
+class ToCopyBoolToFloatModule(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.to(torch.float32)
+
+
+class ToCopyInt8ToFloatModule(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.to(torch.float32)
+
+
+class CompareToCopyBoolToFloatModule(torch.nn.Module):
+    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        return (a > b).to(torch.float32)
+
+
 def to_copy_int_input(shape: tuple[int, ...]) -> torch.Tensor:
     n = math.prod(shape)
     return (torch.arange(n, dtype=torch.int32) - n // 2).reshape(shape)
@@ -61,14 +76,32 @@ def to_copy_float_input(shape: tuple[int, ...]) -> torch.Tensor:
     return pattern.repeat(repeats)[:n].reshape(shape)
 
 
-def _lower(model: torch.nn.Module, x: torch.Tensor):
-    ep = torch.export.export(model.eval(), (x,))
+def bool_tail_input(shape: tuple[int, ...]) -> torch.Tensor:
+    n = math.prod(shape)
+    pattern = torch.tensor([True, False, True, True, False, False, True])
+    repeats = (n + pattern.numel() - 1) // pattern.numel()
+    return pattern.repeat(repeats)[:n].reshape(shape)
+
+
+def compare_to_copy_input_a(shape: tuple[int, ...]) -> torch.Tensor:
+    n = math.prod(shape)
+    pattern = torch.tensor([1.0, -1.0, 2.0, -2.0, 3.0, -3.0, 4.0])
+    repeats = (n + pattern.numel() - 1) // pattern.numel()
+    return pattern.repeat(repeats)[:n].reshape(shape)
+
+
+def compare_to_copy_input_b(shape: tuple[int, ...]) -> torch.Tensor:
+    return torch.zeros(shape, dtype=torch.float32)
+
+
+def _lower(model: torch.nn.Module, *inputs: torch.Tensor):
+    ep = torch.export.export(model.eval(), inputs)
     edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()])
     return ep, edge
 
 
-def _export(model: torch.nn.Module, x: torch.Tensor):
-    _, edge = _lower(model, x)
+def _export(model: torch.nn.Module, *inputs: torch.Tensor):
+    _, edge = _lower(model, *inputs)
     return edge.to_executorch()
 
 
@@ -142,3 +175,22 @@ def test_float_passthrough_delegates(self) -> None:
         self.assertTrue(
             _delegated(et), "Expected a VulkanBackend delegate (to_copy float->float)"
         )
+
+    def test_bool_to_float_delegates(self) -> None:
+        x = bool_tail_input((5,))
+        ep, edge = _lower(ToCopyBoolToFloatModule(), x)
+        self.assertEqual(_prepartition_cast_dtypes(ep), [torch.float32])
+        self.assertEqual(_delegated_cast_dtypes(edge), [torch.float32])
+        self.assertTrue(_delegated(edge.to_executorch()))
+
+    def test_compare_bool_to_float_delegates(self) -> None:
+        a = compare_to_copy_input_a((5,))
+        b = compare_to_copy_input_b((5,))
+        ep, edge = _lower(CompareToCopyBoolToFloatModule(), a, b)
+        self.assertEqual(_prepartition_cast_dtypes(ep), [torch.float32])
+        self.assertEqual(_delegated_cast_dtypes(edge), [torch.float32])
+        self.assertTrue(_delegated(edge.to_executorch()))
+
+    def test_int8_to_float_does_not_delegate(self) -> None:
+        x = torch.tensor([-2, 0, 3], dtype=torch.int8)
+        self.assertFalse(_delegated(_export(ToCopyInt8ToFloatModule(), x)))