[ExecuTorch][WebGPU] Generate typed to_copy shader variants

backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl → backends/webgpu/runtime/ops/to_copy/to_copy_convert.wgsl

@@ -1,5 +1,5 @@
-@group(0) @binding(0) var<storage, read> input: array<f32>;
-@group(0) @binding(1) var<storage, read_write> output: array<i32>;
+@group(0) @binding(0) var<storage, read> input: array<${IN_TYPE}>;
+@group(0) @binding(1) var<storage, read_write> output: array<${OUT_TYPE}>;
 
 struct Params {
   num_elements: u32,
@@ -14,5 +14,5 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     if (idx >= params.num_elements) {
         return;
     }
-    output[idx] = i32(input[idx]);
+    output[idx] = ${OUT_TYPE}(input[idx]);
 }

backends/webgpu/runtime/ops/to_copy/to_copy_convert.yaml

@@ -0,0 +1,15 @@
+# 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.
+
+to_copy_convert:
+  parameter_names_with_default_values:
+    IN_TYPE: f32
+    OUT_TYPE: i32
+  shader_variants:
+    - NAME: to_copy_float_to_int
+    - NAME: to_copy_int_to_float
+      IN_TYPE: i32
+      OUT_TYPE: f32

backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h

@@ -12,7 +12,7 @@
 
 namespace executorch::backends::webgpu {
 
-// @generated from to_copy_float_to_int.wgsl - DO NOT EDIT.
+// @generated from to_copy_convert.wgsl - DO NOT EDIT.
 // wgsl-sha256: c331e00e3171eecbe6317ac9df0a5f9cd6d25da26a9a587250f1cc6086dc3c8f
 inline constexpr const char* kToCopyFloatToIntWGSL = R"(
 @group(0) @binding(0) var<storage, read> input: array<f32>;

backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h

@@ -12,7 +12,7 @@
 
 namespace executorch::backends::webgpu {
 
-// @generated from to_copy_int_to_float.wgsl - DO NOT EDIT.
+// @generated from to_copy_convert.wgsl - DO NOT EDIT.
 // wgsl-sha256: e18dd733a3838f83eded4977a2a2b21119099c8409b234f12474fae5acc9b195
 inline constexpr const char* kToCopyIntToFloatWGSL = R"(
 @group(0) @binding(0) var<storage, read> input: array<i32>;

backends/webgpu/test/op_tests/cases.py

@@ -145,6 +145,13 @@
     SqueezeModule,
 )
 
+from executorch.backends.webgpu.test.ops.test_to_copy import (
+    to_copy_float_input,
+    to_copy_int_input,
+    ToCopyFloatToIntToFloatModule,
+    ToCopyIntToFloatModule,
+)
+
 from executorch.backends.webgpu.test.ops.test_unary_activations import (
     _lin as _unary_lin,
     CLAMP_CONFIGS,
@@ -845,6 +852,38 @@ def _view_copy_suite() -> WebGPUTestSuite:
     return _fn_config_suite(ViewModule, _VIEW_CONFIGS)
 
 
+def _to_copy_factory(variant: str) -> torch.nn.Module:
+    return {
+        "int_to_float": ToCopyIntToFloatModule,
+        "float_roundtrip": ToCopyFloatToIntToFloatModule,
+    }[variant]()
+
+
+@register_op_test("to_copy")
+def _to_copy_suite() -> WebGPUTestSuite:
+    cases = []
+    for n in (63, 64, 65, 257):
+        cases.extend(
+            [
+                Case(
+                    name=f"int_to_float_{n}",
+                    construct={"variant": "int_to_float"},
+                    inputs=(InputSpec(shape=(n,), gen=to_copy_int_input),),
+                ),
+                Case(
+                    name=f"float_roundtrip_{n}",
+                    construct={"variant": "float_roundtrip"},
+                    inputs=(InputSpec(shape=(n,), gen=to_copy_float_input),),
+                ),
+            ]
+        )
+    return WebGPUTestSuite(
+        module_factory=_to_copy_factory,
+        cases=cases,
+        golden_dtype="float32",
+    )
+
+
 @register_op_test("select")
 def _select_suite() -> WebGPUTestSuite:
     return _fn_config_suite(SelectModule, _SELECT_CONFIGS)

backends/webgpu/test/ops/test_to_copy.py

@@ -13,6 +13,7 @@
 lvp golden (yolo11n / Depth-Anything-V2).
 """
 
+import math
 import unittest
 
 import torch
@@ -27,18 +28,47 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x.to(torch.float32)
 
 
+class ToCopyFloatToIntModule(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.to(torch.int32)
+
+
+class ToCopyFloatToIntToFloatModule(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.to(torch.int32).to(torch.float32)
+
+
 class ToCopyFloatModule(torch.nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # Same-dtype copy (flat byte-copy path); copy=True keeps the op from
         # being elided as a no-op.
         return x.to(torch.float32, copy=True)
 
 
-def _export(model: torch.nn.Module, x: torch.Tensor):
+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)
+
+
+def to_copy_float_input(shape: tuple[int, ...]) -> torch.Tensor:
+    n = math.prod(shape)
+    pattern = torch.tensor(
+        [-8.75, -3.0, -1.5, -0.25, 0.0, 0.25, 1.5, 3.0, 8.75],
+        dtype=torch.float32,
+    )
+    repeats = (n + pattern.numel() - 1) // pattern.numel()
+    return pattern.repeat(repeats)[:n].reshape(shape)
+
+
+def _lower(model: torch.nn.Module, x: torch.Tensor):
     ep = torch.export.export(model.eval(), (x,))
-    return to_edge_transform_and_lower(
-        ep, partitioner=[VulkanPartitioner()]
-    ).to_executorch()
+    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)
+    return edge.to_executorch()
 
 
 def _delegated(et) -> bool:
@@ -49,6 +79,31 @@ def _delegated(et) -> bool:
     )
 
 
+def _prepartition_cast_dtypes(ep) -> list[torch.dtype]:
+    return [
+        node.args[1]
+        for node in ep.graph_module.graph.nodes
+        if node.op == "call_function" and node.target == torch.ops.aten.to.dtype
+    ]
+
+
+def _delegated_cast_dtypes(edge) -> list[torch.dtype]:
+    from executorch.exir.lowered_backend_module import get_lowered_submodules
+
+    graph_module = edge.exported_program().graph_module
+    if any(
+        "_to_dim_order_copy" in str(getattr(node, "target", ""))
+        for node in graph_module.graph.nodes
+    ):
+        return []
+    return [
+        node.kwargs["dtype"]
+        for _, lowered, _ in get_lowered_submodules(graph_module)
+        for node in lowered.original_module.graph_module.graph.nodes
+        if "_to_dim_order_copy" in str(getattr(node, "target", ""))
+    ]
+
+
 class ToCopyTest(unittest.TestCase):
     def test_int_to_float_delegates(self) -> None:
         x = torch.tensor([1, 2, 3], dtype=torch.int32)
@@ -57,6 +112,31 @@ def test_int_to_float_delegates(self) -> None:
             _delegated(et), "Expected a VulkanBackend delegate (to_copy int->float)"
         )
 
+    def test_float_to_int_delegates(self) -> None:
+        x = torch.tensor([-3.75, -1.0, 0.0, 1.9, 63.0], dtype=torch.float32)
+        et = _export(ToCopyFloatToIntModule(), x)
+        self.assertTrue(
+            _delegated(et), "Expected a VulkanBackend delegate (to_copy float->int)"
+        )
+
+    def test_roundtrip_keeps_both_casts_in_delegate(self) -> None:
+        x = torch.tensor([-3.75, -1.0, 0.0, 1.9, 63.0], dtype=torch.float32)
+        ep, edge = _lower(ToCopyFloatToIntToFloatModule(), x)
+        expected = [torch.int32, torch.float32]
+        self.assertEqual(_prepartition_cast_dtypes(ep), expected)
+        self.assertEqual(_delegated_cast_dtypes(edge), expected)
+        self.assertTrue(_delegated(edge.to_executorch()))
+
+        for module, one_direction_input in (
+            (ToCopyFloatToIntModule(), x),
+            (
+                ToCopyIntToFloatModule(),
+                torch.tensor([-3, -1, 0, 1, 63], dtype=torch.int32),
+            ),
+        ):
+            _, one_direction_edge = _lower(module, one_direction_input)
+            self.assertNotEqual(_delegated_cast_dtypes(one_direction_edge), expected)
+
     def test_float_passthrough_delegates(self) -> None:
         x = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
         et = _export(ToCopyFloatModule(), x)

backends/webgpu/test/test_wgsl_codegen.py

@@ -213,7 +213,7 @@ def test_generated_output_manifest_digest(self) -> None:
         self.assertEqual(len(outputs), 134)
         self.assertEqual(
             digest.hexdigest(),
-            "19a0baf9345bec02fe2091a0e6320b81d966724bedc33feb0779c6a824925972",
+            "ab15be30e7cfa2cb2f6fa7743d3b9f03535f5cc0b88d64d9b6bad8f777efda25",
         )
 
     def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None:
@@ -907,6 +907,54 @@ def test_rms_norm_template_roundtrip_byte_identical(self) -> None:
                 got, want, f"{header_name} not reproduced from rms_norm.wgsl template"
             )
 
+    def test_to_copy_convert_template_roundtrip_byte_identical(self) -> None:
+        to_copy_dir = g.BACKEND_ROOT / "runtime/ops/to_copy"
+        template_path = to_copy_dir / "to_copy_convert.wgsl"
+        spec = g.parse_template_spec(template_path.with_suffix(".yaml"))
+        variants = {params["NAME"]: params for params in spec[template_path.stem]}
+        expected = {
+            "to_copy_float_to_int": (
+                "f32",
+                "i32",
+                "c331e00e3171eecbe6317ac9df0a5f9cd6d25da26a9a587250f1cc6086dc3c8f",
+            ),
+            "to_copy_int_to_float": (
+                "i32",
+                "f32",
+                "e18dd733a3838f83eded4977a2a2b21119099c8409b234f12474fae5acc9b195",
+            ),
+        }
+        self.assertEqual(set(variants), set(expected))
+        template = template_path.read_text()
+
+        for name, (in_type, out_type, expected_hash) in expected.items():
+            params = variants[name]
+            self.assertEqual(
+                (params["IN_TYPE"], params["OUT_TYPE"]), (in_type, out_type)
+            )
+            expanded = g.preprocess(template, {**g.WGSL_HELPERS, **params})
+            self.assertEqual(g.wgsl_sha256(expanded), expected_hash)
+
+            header = (to_copy_dir / f"{name}_wgsl.h").read_text()
+            body = header.split('R"(', 1)[1].split(')";', 1)[0][1:]
+            self.assertEqual(body, expanded)
+            self.assertEqual(g.embedded_sha256(header), expected_hash)
+            self.assertEqual(g.parse_workgroup_size(body), (64, 1, 1))
+
+        entries = {entry.name: entry for entry in g.registry_entries()}
+        self.assertEqual(
+            entries["to_copy_float_to_int"].include,
+            "runtime/ops/to_copy/to_copy_float_to_int_wgsl.h",
+        )
+        self.assertEqual(
+            entries["to_copy_int_to_float"].include,
+            "runtime/ops/to_copy/to_copy_int_to_float_wgsl.h",
+        )
+        self.assertEqual(
+            hashlib.sha256(g.registry_path().read_bytes()).hexdigest(),
+            "ce1777820bffe77e7cdda312f86a3fd41a090f8f282a6add66666446d06b1608",
+        )
+
     def test_rms_norm_half_variant_is_type_correct(self) -> None:
         # A DTYPE=half expansion must emit compilable WGSL: `enable f16;`, an f32
         # accumulator, loads widened to f32 for the reduction, and the store