From 78490d50c7dd9c12f2c3d0f50d5a5e0b4b87b43a Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 14 Jul 2026 14:48:59 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 3 + .../webgpu/runtime/ops/dim_order/DimOrder.cpp | 31 ++++ .../runtime/ops/expand_copy/ExpandCopy.cpp | 139 ++++++++++++++ .../runtime/ops/expand_copy/expand_copy.wgsl | 29 +++ .../ops/expand_copy/expand_copy_wgsl.h | 53 ++++++ backends/webgpu/runtime/ops/fill/Fill.cpp | 171 ++++++++++++++++++ backends/webgpu/runtime/ops/fill/fill.wgsl | 20 ++ backends/webgpu/runtime/ops/fill/fill_wgsl.h | 44 +++++ 8 files changed, 490 insertions(+) create mode 100644 backends/webgpu/runtime/ops/dim_order/DimOrder.cpp create mode 100644 backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp create mode 100644 backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl create mode 100644 backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h create mode 100644 backends/webgpu/runtime/ops/fill/Fill.cpp create mode 100644 backends/webgpu/runtime/ops/fill/fill.wgsl create mode 100644 backends/webgpu/runtime/ops/fill/fill_wgsl.h diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index f92ed119b4e..4dbd6e77965 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -62,6 +62,9 @@ set(WEBGPU_SRCS runtime/ops/where/Where.cpp runtime/ops/compare/Compare.cpp runtime/ops/gather/Gather.cpp + runtime/ops/expand_copy/ExpandCopy.cpp + runtime/ops/fill/Fill.cpp + runtime/ops/dim_order/DimOrder.cpp runtime/ops/linear/Linear.cpp runtime/ops/embedding/Embedding.cpp runtime/ops/logical_not/LogicalNot.cpp diff --git a/backends/webgpu/runtime/ops/dim_order/DimOrder.cpp b/backends/webgpu/runtime/ops/dim_order/DimOrder.cpp new file mode 100644 index 00000000000..e57e0375a4f --- /dev/null +++ b/backends/webgpu/runtime/ops/dim_order/DimOrder.cpp @@ -0,0 +1,31 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +namespace { + +// _clone_dim_order = numel-preserving flat copy (shared DMA helper). +void clone_dim_order_impl(WebGPUGraph& graph, const std::vector& args) { + add_flat_copy(graph, args.at(0), args.at(args.size() - 1)); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP( + dim_order_ops._clone_dim_order.default, clone_dim_order_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp new file mode 100644 index 00000000000..507677c7bca --- /dev/null +++ b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp @@ -0,0 +1,139 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +namespace { + +// out coord -> in coord; size-1 in-dims broadcast via clamp (mirrors mul). +void expand_copy_impl(WebGPUGraph& graph, const std::vector& args) { + const int in_id = args.at(0); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("expand_copy: in/out arg is not a tensor"); + } + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + + TensorMeta out_meta; + TensorMeta in_meta; + fill_tensor_meta(out_tensor, &out_meta); + fill_tensor_meta_broadcast(in_tensor, out_meta.ndim, &in_meta); + if (out_tensor.nbytes != + static_cast(out_meta.numel) * sizeof(float) || + in_tensor.nbytes != static_cast(in_meta.numel) * sizeof(float)) { + throw std::runtime_error( + "expand_copy: non-fp32 operand (nbytes != numel*4)"); + } + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kExpandCopyWorkgroupSizeX); + uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, out_meta.numel, wg_size, "expand_copy"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer out_meta_buf = + utils::make_uniform(device, &out_meta, sizeof(TensorMeta)); + WGPUBuffer in_meta_buf = + utils::make_uniform(device, &in_meta, sizeof(TensorMeta)); + graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kExpandCopyWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[4] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + entries[3].binding = 3; + entries[3].visibility = WGPUShaderStage_Compute; + entries[3].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 4; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[4] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = out_meta_buf; + bg_entries[2].size = sizeof(TensorMeta); + bg_entries[3].binding = 3; + bg_entries[3].buffer = in_meta_buf; + bg_entries[3].size = sizeof(TensorMeta); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 4; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch({pipeline, bind_group, workgroup_count}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + wgpuBufferRelease(out_meta_buf); + wgpuBufferRelease(in_meta_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.expand_copy.default, expand_copy_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl b/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl new file mode 100644 index 00000000000..d6c65588978 --- /dev/null +++ b/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl @@ -0,0 +1,29 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(2) var out_meta: TensorMeta; +@group(0) @binding(3) var in_meta: TensorMeta; + +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 >= out_meta.numel) { + return; + } + var rem = idx; + var l: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d]; + rem = rem % out_meta.strides[d]; + l = l + min(coord, in_meta.sizes[d] - 1u) * in_meta.strides[d]; + } + output[idx] = input[l]; +} diff --git a/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h b/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h new file mode 100644 index 00000000000..56b3f708767 --- /dev/null +++ b/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h @@ -0,0 +1,53 @@ +/* + * 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 expand_copy.wgsl - DO NOT EDIT. +// wgsl-sha256: ad996b7fd6eca5c5773715af3a9a117da83ef522042eb0ee918a623096417815 +inline constexpr const char* kExpandCopyWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(2) var out_meta: TensorMeta; +@group(0) @binding(3) var in_meta: TensorMeta; + +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 >= out_meta.numel) { + return; + } + var rem = idx; + var l: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d]; + rem = rem % out_meta.strides[d]; + l = l + min(coord, in_meta.sizes[d] - 1u) * in_meta.strides[d]; + } + output[idx] = input[l]; +} +)"; + +inline constexpr uint32_t kExpandCopyWorkgroupSizeX = 64; +inline constexpr uint32_t kExpandCopyWorkgroupSizeY = 1; +inline constexpr uint32_t kExpandCopyWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/fill/Fill.cpp b/backends/webgpu/runtime/ops/fill/Fill.cpp new file mode 100644 index 00000000000..6c0441ff5ad --- /dev/null +++ b/backends/webgpu/runtime/ops/fill/Fill.cpp @@ -0,0 +1,171 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform buffer layout matching the WGSL Params struct (16-byte aligned). +struct FillParams { + uint32_t num_elements; + float fill_value; + uint32_t _pad[2]; +}; + +float read_scalar(WebGPUGraph& graph, int id, const char* op_name) { + if (graph.get_value_type(id) == WebGPUGraph::ValueType::Double) { + return static_cast(graph.get_double(id)); + } + if (graph.get_value_type(id) == WebGPUGraph::ValueType::Int) { + return static_cast(graph.get_int(id)); + } + throw std::runtime_error( + std::string(op_name) + ": fill value is not a scalar"); +} + +// Fills the (pre-allocated) output buffer with a constant scalar. +void add_fill( + WebGPUGraph& graph, + int out_id, + float fill_value, + const char* op_name) { + WGPUDevice device = graph.device(); + const auto& out_tensor = graph.get_tensor(out_id); + if (out_tensor.buffer == nullptr) { + throw std::runtime_error(std::string(op_name) + ": null output buffer"); + } + uint32_t num_elements = + static_cast(out_tensor.nbytes / sizeof(float)); + + uint32_t wg_size = utils::clamp_workgroup_size(device, kFillWorkgroupSizeX); + utils::WgCount workgroup_count = + utils::compute_2d_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); + + FillParams params = {}; + params.num_elements = num_elements; + params.fill_value = fill_value; + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(FillParams)); + graph.add_uniform_buffer_bytes(sizeof(FillParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kFillWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[2] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_Storage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 2; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[2] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = out_tensor.buffer; + bg_entries[0].size = out_tensor.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = uniform_buffer; + bg_entries[1].size = sizeof(FillParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 2; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, bind_group, workgroup_count.x, "", workgroup_count.y}); + + // Dynamic shapes: recompute num_elements/dispatch from the live output dims. + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + out_id, + [out_id, fill_value, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(out_id); + const uint64_t numel = utils::numel_of(d); + FillParams p = {}; + p.num_elements = static_cast(numel); + p.fill_value = fill_value; + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), static_cast(numel), wg_size, "fill(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(uniform_buffer); +} + +void full_impl(WebGPUGraph& graph, const std::vector& args) { + add_fill( + graph, args.at(args.size() - 1), read_scalar(graph, args.at(1), "full"), + "full"); +} + +void full_like_impl(WebGPUGraph& graph, const std::vector& args) { + add_fill( + graph, args.at(args.size() - 1), + read_scalar(graph, args.at(1), "full_like"), "full_like"); +} + +void scalar_tensor_impl(WebGPUGraph& graph, const std::vector& args) { + add_fill( + graph, args.at(args.size() - 1), + read_scalar(graph, args.at(0), "scalar_tensor"), "scalar_tensor"); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.full.default, full_impl); + WEBGPU_REGISTER_OP(aten.full_like.default, full_like_impl); + WEBGPU_REGISTER_OP(aten.scalar_tensor.default, scalar_tensor_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/fill/fill.wgsl b/backends/webgpu/runtime/ops/fill/fill.wgsl new file mode 100644 index 00000000000..684c350e8d6 --- /dev/null +++ b/backends/webgpu/runtime/ops/fill/fill.wgsl @@ -0,0 +1,20 @@ +@group(0) @binding(0) var output: array; + +struct Params { + num_elements: u32, + fill_value: f32, +} +@group(0) @binding(1) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = params.fill_value; +} diff --git a/backends/webgpu/runtime/ops/fill/fill_wgsl.h b/backends/webgpu/runtime/ops/fill/fill_wgsl.h new file mode 100644 index 00000000000..a083af10af3 --- /dev/null +++ b/backends/webgpu/runtime/ops/fill/fill_wgsl.h @@ -0,0 +1,44 @@ +/* + * 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 fill.wgsl - DO NOT EDIT. +// wgsl-sha256: 9acfd8079cf404b22856c157b339922911e15bb0e808aa76b5a7fb81a8505cf7 +inline constexpr const char* kFillWGSL = R"( +@group(0) @binding(0) var output: array; + +struct Params { + num_elements: u32, + fill_value: f32, +} +@group(0) @binding(1) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = params.fill_value; +} +)"; + +inline constexpr uint32_t kFillWorkgroupSizeX = 256; +inline constexpr uint32_t kFillWorkgroupSizeY = 1; +inline constexpr uint32_t kFillWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu