diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl b/backends/webgpu/runtime/ops/to_copy/to_copy_convert.wgsl similarity index 62% rename from backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl rename to backends/webgpu/runtime/ops/to_copy/to_copy_convert.wgsl index 3eb6cb44595..f1113f0e14c 100644 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_convert.wgsl @@ -1,5 +1,5 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; +@group(0) @binding(0) var input: array<${IN_TYPE}>; +@group(0) @binding(1) var output: array<${OUT_TYPE}>; struct Params { num_elements: u32, @@ -14,5 +14,5 @@ fn main(@builtin(global_invocation_id) gid: vec3) { if (idx >= params.num_elements) { return; } - output[idx] = i32(input[idx]); + output[idx] = ${OUT_TYPE}(input[idx]); } diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_convert.yaml b/backends/webgpu/runtime/ops/to_copy/to_copy_convert.yaml new file mode 100644 index 00000000000..219b2d8edd3 --- /dev/null +++ b/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 diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h b/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h index e7e0391dd13..1a384c747d8 100644 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h +++ b/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 input: array; diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl deleted file mode 100644 index 87affe78290..00000000000 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl +++ /dev/null @@ -1,18 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; - if (idx >= params.num_elements) { - return; - } - output[idx] = f32(input[idx]); -} diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h index fd18700f17c..6fdf37ec2b7 100644 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h +++ b/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 input: array; diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 19438d6acbc..be46d35b520 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/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) diff --git a/backends/webgpu/test/ops/test_to_copy.py b/backends/webgpu/test/ops/test_to_copy.py index baf068a3438..1fa2375f248 100644 --- a/backends/webgpu/test/ops/test_to_copy.py +++ b/backends/webgpu/test/ops/test_to_copy.py @@ -13,12 +13,14 @@ lvp golden (yolo11n / Depth-Anything-V2). """ +import math import unittest import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.exir import to_edge_transform_and_lower +from executorch.exir.lowered_backend_module import get_lowered_submodules class ToCopyIntToFloatModule(torch.nn.Module): @@ -27,6 +29,16 @@ 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 @@ -34,11 +46,30 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: 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 +80,29 @@ 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]: + 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 +111,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) diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 831d2ce0747..b8c239483f2 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/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