From 1022e76712bb475ed6797a07b85e2fa517602e82 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 16 Jul 2026 14:48:50 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 1 + .../webgpu/runtime/ops/to_copy/ToCopy.cpp | 163 ++++++++++++++++++ backends/webgpu/runtime/ops/to_copy/to_copy.h | 18 ++ .../ops/to_copy/to_copy_float_to_int.wgsl | 18 ++ .../ops/to_copy/to_copy_float_to_int_wgsl.h | 42 +++++ .../ops/to_copy/to_copy_int_to_float.wgsl | 18 ++ .../ops/to_copy/to_copy_int_to_float_wgsl.h | 42 +++++ .../webgpu/runtime/ops/view_copy/ViewCopy.cpp | 13 +- 8 files changed, 313 insertions(+), 2 deletions(-) create mode 100644 backends/webgpu/runtime/ops/to_copy/ToCopy.cpp create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy.h create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 54558a8972b..89f4ef62742 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -43,6 +43,7 @@ set(WEBGPU_SRCS runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp runtime/ops/rope/RotaryEmbedding.cpp runtime/ops/prepack/Prepack.cpp + runtime/ops/to_copy/ToCopy.cpp runtime/ops/view_copy/ViewCopy.cpp runtime/ops/select/Select.cpp runtime/ops/sigmoid/UnaryOp.cpp diff --git a/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp new file mode 100644 index 00000000000..05f7ff6fe24 --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp @@ -0,0 +1,163 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform buffer layout matching the WGSL Params struct; 16-byte aligned. +struct ConvertParams { + uint32_t num_elements; + uint32_t _pad[3]; +}; + +// Elementwise int32<->fp32 convert; mirrors Vulkan add_view_copy_convert_node. +void add_convert_op( + WebGPUGraph& graph, + int in_id, + int out_id, + const char* wgsl_source, + uint32_t wg_size_x, + const char* op_name) { + WGPUDevice device = graph.device(); + + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error(std::string(op_name) + ": null buffer binding"); + } + + // 32-bit only: int64 consts are downcast to int32 by the Vulkan serializer. + if (in_tensor.elem_size != 4 || out_tensor.elem_size != 4) { + throw std::runtime_error( + std::string(op_name) + ": only 32-bit int<->float convert supported"); + } + if (in_tensor.nbytes != out_tensor.nbytes) { + throw std::runtime_error(std::string(op_name) + ": numel mismatch"); + } + + uint32_t num_elements = static_cast(out_tensor.nbytes / 4); + + uint32_t wg_size = utils::clamp_workgroup_size(device, wg_size_x); + uint32_t workgroup_count = + utils::compute_1d_workgroup_count(device, num_elements, wg_size, op_name); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + ConvertParams params = {}; + params.num_elements = num_elements; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(ConvertParams)); + graph.add_uniform_buffer_bytes(sizeof(ConvertParams)); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + wgsl_source, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(ConvertParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_idx = + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + + // Dynamic shapes: recompute num_elements/dispatch for the live shape. + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + const uint64_t numel = utils::numel_of(d); + g.set_cur_dims(out_id, d); + ConvertParams p = {}; + p.num_elements = static_cast(numel); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = + utils::compute_1d_workgroup_count( + g.device(), + static_cast(numel), + wg_size, + "to_copy(resize)"); + }); + + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); +} + +void to_copy_impl(WebGPUGraph& graph, const std::vector& args) { + // aten._to_copy.default args: [self, ...kwargs, out]; out = last value id. + add_to_copy_node(graph, args.at(0), args.at(args.size() - 1)); +} + +} // namespace + +void add_to_copy_node(WebGPUGraph& graph, int in_id, int out_id) { + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + + // Same is_int+width = flat byte copy; unique dtype key in the 32-bit domain. + if (in_tensor.is_int == out_tensor.is_int && + in_tensor.elem_size == out_tensor.elem_size) { + add_flat_copy(graph, in_id, out_id); + return; + } + + // int<->float = numeric convert (mirrors Vulkan add_view_copy_convert_node). + if (in_tensor.is_int && !out_tensor.is_int) { + add_convert_op( + graph, + in_id, + out_id, + kToCopyIntToFloatWGSL, + kToCopyIntToFloatWorkgroupSizeX, + "to_copy_int_to_float"); + } else { + add_convert_op( + graph, + in_id, + out_id, + kToCopyFloatToIntWGSL, + kToCopyFloatToIntWorkgroupSizeX, + "to_copy_float_to_int"); + } +} + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten._to_copy.default, to_copy_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy.h b/backends/webgpu/runtime/ops/to_copy/to_copy.h new file mode 100644 index 00000000000..714e9a5224e --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy.h @@ -0,0 +1,18 @@ +/* + * 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. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// Copy: flat byte for same dtype, numeric convert for int<->float (Vulkan). +void add_to_copy_node(WebGPUGraph& graph, int in_id, int out_id); + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl b/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl new file mode 100644 index 00000000000..3eb6cb44595 --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl @@ -0,0 +1,18 @@ +@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] = i32(input[idx]); +} 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 new file mode 100644 index 00000000000..e7e0391dd13 --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h @@ -0,0 +1,42 @@ +/* + * 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. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from to_copy_float_to_int.wgsl - DO NOT EDIT. +// wgsl-sha256: c331e00e3171eecbe6317ac9df0a5f9cd6d25da26a9a587250f1cc6086dc3c8f +inline constexpr const char* kToCopyFloatToIntWGSL = R"( +@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] = i32(input[idx]); +} +)"; + +inline constexpr uint32_t kToCopyFloatToIntWorkgroupSizeX = 64; +inline constexpr uint32_t kToCopyFloatToIntWorkgroupSizeY = 1; +inline constexpr uint32_t kToCopyFloatToIntWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu 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 new file mode 100644 index 00000000000..87affe78290 --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl @@ -0,0 +1,18 @@ +@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 new file mode 100644 index 00000000000..fd18700f17c --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h @@ -0,0 +1,42 @@ +/* + * 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. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from to_copy_int_to_float.wgsl - DO NOT EDIT. +// wgsl-sha256: e18dd733a3838f83eded4977a2a2b21119099c8409b234f12474fae5acc9b195 +inline constexpr const char* kToCopyIntToFloatWGSL = R"( +@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]); +} +)"; + +inline constexpr uint32_t kToCopyIntToFloatWorkgroupSizeX = 64; +inline constexpr uint32_t kToCopyIntToFloatWorkgroupSizeY = 1; +inline constexpr uint32_t kToCopyIntToFloatWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp b/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp index 7b3b44eb86f..5c75b2d9084 100644 --- a/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp +++ b/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include @@ -98,14 +99,22 @@ void clone_impl(WebGPUGraph& graph, const std::vector& args) { add_flat_copy(graph, args.at(0), args.at(args.size() - 1)); } +// dim_order copies may change dtype -> convert-aware add_to_copy_node. +void clone_dim_order_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [self, non_blocking?, dim_order?, out]; out = last value id. + add_to_copy_node(graph, args.at(0), args.at(args.size() - 1)); +} + } // namespace WEBGPU_REGISTER_OPERATORS { WEBGPU_REGISTER_OP(aten.view_copy.default, view_copy_impl); WEBGPU_REGISTER_OP(aten.clone.default, clone_impl); WEBGPU_REGISTER_OP(aten.alias_copy.default, clone_impl); - // _clone_dim_order ignores dim_order (AOT pass elides via shape+dtype). - WEBGPU_REGISTER_OP(dim_order_ops._clone_dim_order.default, clone_impl); + WEBGPU_REGISTER_OP( + dim_order_ops._clone_dim_order.default, clone_dim_order_impl); + WEBGPU_REGISTER_OP( + dim_order_ops._to_dim_order_copy.default, clone_dim_order_impl); } } // namespace executorch::backends::webgpu