From a219bb55a759170d40f6ff7f97df9ad8657b38d2 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 10 Jul 2026 11:25:58 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/runtime/ops/relu/relu.wgsl | 20 +++ backends/webgpu/runtime/ops/relu/relu_wgsl.h | 44 ++++++ .../webgpu/runtime/ops/sigmoid/UnaryOp.cpp | 135 ++++++------------ .../webgpu/runtime/ops/sigmoid/sigmoid.wgsl | 10 +- .../webgpu/runtime/ops/sigmoid/sigmoid_wgsl.h | 14 +- 5 files changed, 120 insertions(+), 103 deletions(-) create mode 100644 backends/webgpu/runtime/ops/relu/relu.wgsl create mode 100644 backends/webgpu/runtime/ops/relu/relu_wgsl.h diff --git a/backends/webgpu/runtime/ops/relu/relu.wgsl b/backends/webgpu/runtime/ops/relu/relu.wgsl new file mode 100644 index 00000000000..0d9885421f9 --- /dev/null +++ b/backends/webgpu/runtime/ops/relu/relu.wgsl @@ -0,0 +1,20 @@ +@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 = 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] = max(input[idx], 0.0); +} diff --git a/backends/webgpu/runtime/ops/relu/relu_wgsl.h b/backends/webgpu/runtime/ops/relu/relu_wgsl.h new file mode 100644 index 00000000000..ae74266ce8a --- /dev/null +++ b/backends/webgpu/runtime/ops/relu/relu_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 relu.wgsl - DO NOT EDIT. +// wgsl-sha256: 2f7fac19d55cb7e55749f2bc1856278c1f2afaf0bc7ff8e663fb9e14b4188199 +inline constexpr const char* kReluWGSL = 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 = 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] = max(input[idx], 0.0); +} +)"; + +inline constexpr uint32_t kReluWorkgroupSizeX = 256; +inline constexpr uint32_t kReluWorkgroupSizeY = 1; +inline constexpr uint32_t kReluWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp b/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp index 7c8b8c0e9ce..fc228853d93 100644 --- a/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp +++ b/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include @@ -39,31 +40,20 @@ void add_unary_op( 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"); - } - // 4-byte (fp32) alignment guard on both operands; also the dtype guard. - if (in_tensor.nbytes % sizeof(float) != 0 || - out_tensor.nbytes % sizeof(float) != 0) { - throw std::runtime_error( - std::string(op_name) + ": operand not 4-byte aligned"); - } - if (in_tensor.nbytes != out_tensor.nbytes) { - throw std::runtime_error( - std::string(op_name) + ": input/output size mismatch"); - } + utils::check_elementwise_fp32_io(in_tensor, out_tensor, op_name); uint32_t num_elements = static_cast(out_tensor.nbytes / sizeof(float)); + // Adaptive 1D->2D dispatch: wg=clamp(device,256) + 2D-spill past the 65535 + // per-dim ceiling. The shader decodes idx via num_workgroups.x, so the live + // count_x sets the stride at runtime (resize-safe, no override to re-bake). 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); + 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); + WGPUConstantEntry wg_size_constant = utils::make_wg_size_constant(wg_size); UnaryParams params = {}; params.num_elements = num_elements; @@ -72,71 +62,29 @@ void add_unary_op( utils::make_uniform(device, ¶ms, sizeof(UnaryParams)); graph.add_uniform_buffer_bytes(sizeof(UnaryParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {wgsl_source, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group layout: input (read storage) + output (storage) + params. - WGPUBindGroupLayoutEntry entries[3] = {}; - - 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; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - 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[3] = {}; - - 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 = uniform_buffer; - bg_entries[2].size = sizeof(UnaryParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - 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}); + // input (read storage) + output (storage) + params. + 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(UnaryParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_idx = graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, workgroup_count.x, workgroup_count.y); // Dynamic shapes: recompute num_elements/dispatch for the live shape. WGPUBuffer params_buf = uniform_buffer; @@ -149,18 +97,12 @@ void add_unary_op( UnaryParams 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, - "unary(resize)"); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), static_cast(numel), wg_size, "unary(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }); - // Release intermediates (pipeline + bind_group are kept by dispatch). - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } @@ -176,10 +118,17 @@ void sigmoid_impl(WebGPUGraph& graph, const std::vector& args) { "sigmoid"); } +void relu_impl(WebGPUGraph& graph, const std::vector& args) { + // aten.relu.default args: [in, out] + add_unary_op( + graph, args.at(0), args.at(1), kReluWGSL, kReluWorkgroupSizeX, "relu"); +} + } // namespace WEBGPU_REGISTER_OPERATORS { WEBGPU_REGISTER_OP(aten.sigmoid.default, sigmoid_impl); + WEBGPU_REGISTER_OP(aten.relu.default, relu_impl); } } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sigmoid/sigmoid.wgsl b/backends/webgpu/runtime/ops/sigmoid/sigmoid.wgsl index 09b3e5457b3..a57dd46f150 100644 --- a/backends/webgpu/runtime/ops/sigmoid/sigmoid.wgsl +++ b/backends/webgpu/runtime/ops/sigmoid/sigmoid.wgsl @@ -6,11 +6,13 @@ struct Params { } @group(0) @binding(2) var params: Params; -override wg_size: u32 = 64u; +override wg_size: u32 = 256; -@compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +@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; } diff --git a/backends/webgpu/runtime/ops/sigmoid/sigmoid_wgsl.h b/backends/webgpu/runtime/ops/sigmoid/sigmoid_wgsl.h index 48e6efc607a..fc87b981d34 100644 --- a/backends/webgpu/runtime/ops/sigmoid/sigmoid_wgsl.h +++ b/backends/webgpu/runtime/ops/sigmoid/sigmoid_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from sigmoid.wgsl - DO NOT EDIT. -// wgsl-sha256: 70395dbb107b8b95ae13c0a6fb12a8415c561c645da0347294c92904314ae84c +// wgsl-sha256: 79a1554e9d957e10c0f4379b5e0240191743952198e9ed9c0342402bacd356de inline constexpr const char* kSigmoidWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -23,11 +23,13 @@ struct Params { } @group(0) @binding(2) var params: Params; -override wg_size: u32 = 64u; +override wg_size: u32 = 256; -@compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +@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; } @@ -35,7 +37,7 @@ fn main(@builtin(global_invocation_id) gid: vec3) { } )"; -inline constexpr uint32_t kSigmoidWorkgroupSizeX = 64; +inline constexpr uint32_t kSigmoidWorkgroupSizeX = 256; inline constexpr uint32_t kSigmoidWorkgroupSizeY = 1; inline constexpr uint32_t kSigmoidWorkgroupSizeZ = 1;